From 616346b423e739f53621216a90144e6402e16989 Mon Sep 17 00:00:00 2001
From: mdfaijul <md.faijul.amin@intel.com>
Date: Wed, 6 Sep 2023 17:56:01 -0700
Subject: [PATCH 001/478] Enable quantized matmul fusions.

---
 tensorflow/core/kernels/mkl/BUILD             |  20 +
 .../core/kernels/mkl/mkl_fused_ops_test.cc    |   2 +-
 .../core/kernels/mkl/mkl_kernel_util.cc       |   4 +
 tensorflow/core/kernels/mkl/mkl_kernel_util.h |  32 +
 .../core/kernels/mkl/mkl_matmul_op_fused.cc   | 752 ++++++++++++++++--
 .../core/kernels/mkl/mkl_matmul_ops_common.h  | 116 ++-
 tensorflow/core/kernels/mkl/mkl_qmatmul_op.cc |   2 +-
 .../core/kernels/mkl/mkl_qmatmul_op_test.cc   | 425 +++++++---
 .../mkl/onednn_fused_matmul_ops_test.cc       | 749 +++++++++++++++++
 tensorflow/core/ops/mkl_nn_ops.cc             |  40 +
 10 files changed, 1915 insertions(+), 227 deletions(-)
 create mode 100644 tensorflow/core/kernels/mkl/onednn_fused_matmul_ops_test.cc

diff --git a/tensorflow/core/kernels/mkl/BUILD b/tensorflow/core/kernels/mkl/BUILD
index 329139874b101f..4dd5370d281aa2 100644
--- a/tensorflow/core/kernels/mkl/BUILD
+++ b/tensorflow/core/kernels/mkl/BUILD
@@ -95,6 +95,7 @@ tf_mkl_kernel_library(
     hdrs = [
         "mkl_kernel_util.h",
         "mkl_matmul_ops_common.h",
+        "mkl_quantized_conv_ops.h",
     ],
     deps = [
         "//tensorflow/core:graph",
@@ -171,6 +172,7 @@ tf_cc_test_mkl(
     srcs = ["mkl_qmatmul_op_test.cc"],
     linkstatic = 1,  # Fixes dyld error on MacOS.
     deps = [
+        ":mkl_matmul_op",
         ":mkl_qmatmul_op",
         "//tensorflow/core:array_ops_op_lib",
         "//tensorflow/core:math_ops_op_lib",
@@ -554,3 +556,21 @@ tf_cc_test_mkl(
         "//tensorflow/core/kernels/mkl:mkl_softmax_op",
     ] + MKL_TEST_DEPS,
 )
+
+tf_cc_test_mkl(
+    name = "onednn_fused_matmul_ops_test",
+    size = "medium",
+    srcs = ["onednn_fused_matmul_ops_test.cc"],
+    linkstatic = 1,  # Fixes dyld error on MacOS.
+    deps = [
+        ":mkl_matmul_op",
+        ":mkl_kernel_util",
+        "//tensorflow/cc:cc_ops_internal",
+        "//tensorflow/core:direct_session",
+        "//tensorflow/core/kernels:matmul_op",
+        "//tensorflow/core/kernels:relu_op",
+        "//tensorflow/core/kernels:bias_op",
+        "//tensorflow/core/kernels:quantization_utils",
+        "@com_google_absl//absl/strings",
+    ] + MKL_TEST_DEPS,
+)
diff --git a/tensorflow/core/kernels/mkl/mkl_fused_ops_test.cc b/tensorflow/core/kernels/mkl/mkl_fused_ops_test.cc
index 857a25bbe55e7a..6d709794b6fb92 100644
--- a/tensorflow/core/kernels/mkl/mkl_fused_ops_test.cc
+++ b/tensorflow/core/kernels/mkl/mkl_fused_ops_test.cc
@@ -1098,7 +1098,7 @@ class MklFusedMatMulCacheTest : public OpsTestBase {
     // Bias vector.
     AddInputFromArray<float>(TensorShape({4}), {1, 2, 3, 4});
 
-    using KernelType = MklDnnMatMulOpBase<float, void, float>;
+    using KernelType = MklDnnMatMulOpBase<float, float, float>;
     // Before the first time kernel execution, weight should be empty
     EXPECT_TRUE(static_cast<KernelType*>(this->kernel_.get())
                     ->IsWeightCacheEmpty(this->context_.get()));
diff --git a/tensorflow/core/kernels/mkl/mkl_kernel_util.cc b/tensorflow/core/kernels/mkl/mkl_kernel_util.cc
index a4731ea7888337..02e5d0a578df9c 100644
--- a/tensorflow/core/kernels/mkl/mkl_kernel_util.cc
+++ b/tensorflow/core/kernels/mkl/mkl_kernel_util.cc
@@ -22,6 +22,7 @@ limitations under the License.
 #include "tensorflow/cc/ops/array_ops.h"
 #include "tensorflow/cc/ops/const_op.h"
 #include "tensorflow/core/graph/node_builder.h"
+#include "tensorflow/core/lib/core/errors.h"
 
 namespace tensorflow {
 
@@ -40,6 +41,9 @@ void MklTestingUtil::RunMklQuantizeOp(const Tensor& input,
   Node* max_node = test::graph::Constant(&*graph, Tensor(max), "max");
 
   Node* quantize_op;
+  string round_mode =
+      (mode == "SCALE") ? "HALF_TO_EVEN" : "HALF_AWAY_FROM_ZERO";
+
   TF_CHECK_OK(NodeBuilder("mkl_quantizeV2", "_MklQuantizeV2")
                   .Input(input_node)
                   .Input(min_node)
diff --git a/tensorflow/core/kernels/mkl/mkl_kernel_util.h b/tensorflow/core/kernels/mkl/mkl_kernel_util.h
index fb9df4d237c74b..a3015a7b09025c 100644
--- a/tensorflow/core/kernels/mkl/mkl_kernel_util.h
+++ b/tensorflow/core/kernels/mkl/mkl_kernel_util.h
@@ -49,6 +49,38 @@ class MklTestingUtil {
     *tensor_min = min();
     *tensor_max = max();
   }
+
+  // This utility function mimics Quantization of float/bfloat16 tensor with
+  // oneDNN backend QuantizeV2 operation. Since the op signature requires min
+  // and max values to be in float type, min_tensor and max_tensor should have
+  // their dtype set to DT_FLOAT.
+  template <typename T>
+  static Status GetQuantizationTensors(const Tensor& input, Tensor* output,
+                                     DataType out_type, const string mode,
+                                     Tensor* min_tensor, Tensor* max_tensor) {
+    if (min_tensor->dtype() != DT_FLOAT || max_tensor->dtype() != DT_FLOAT) {
+      return absl::UnimplementedError("Tensor must be float32.");
+    }
+    T min;
+    T max;
+    ComputeMinMax<T>(input, &min, &max);
+
+    float adjusted_min = static_cast<float>(min);
+    float adjusted_max = static_cast<float>(max);
+    if (mode == "SCALED") {
+      if (output->dtype() != DT_QINT8) {
+        return absl::UnimplementedError("Tensor must be QInt8 in SCALED mode.");
+      }
+      float range = std::max(std::abs(adjusted_min), std::abs(adjusted_max));
+      adjusted_min = -range;
+      adjusted_max = range;
+    }
+    RunMklQuantizeOp(input, adjusted_min, adjusted_max, out_type, mode, output);
+    min_tensor->flat<float>()(0) = adjusted_min;
+    max_tensor->flat<float>()(0) = adjusted_max;
+
+    return OkStatus();
+  }
 };
 
 #ifdef ENABLE_ONEDNN_V3
diff --git a/tensorflow/core/kernels/mkl/mkl_matmul_op_fused.cc b/tensorflow/core/kernels/mkl/mkl_matmul_op_fused.cc
index 2d0065a52e5b4a..7e5544f1c3e3b4 100644
--- a/tensorflow/core/kernels/mkl/mkl_matmul_op_fused.cc
+++ b/tensorflow/core/kernels/mkl/mkl_matmul_op_fused.cc
@@ -1,4 +1,4 @@
-/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
 
 Licensed under the Apache License, Version 2.0 (the "License");
 you may not use this file except in compliance with the License.
@@ -19,19 +19,29 @@ limitations under the License.
 // Multiplication (MatMul) with bias (BiasAdd) operations.
 #if defined(INTEL_MKL)
 
+#include <type_traits>
+
+#include "oneapi/dnnl/dnnl.hpp"
 #include "tensorflow/core/framework/register_types.h"
 #include "tensorflow/core/kernels/fill_functor.h"
 #include "tensorflow/core/kernels/mkl/mkl_matmul_ops_common.h"
+#include "tensorflow/core/kernels/mkl/mkl_quantized_conv_ops.h"
 #include "tensorflow/core/lib/core/errors.h"
+#include "tensorflow/core/lib/gtl/inlined_vector.h"
+#include "tensorflow/core/platform/errors.h"
 
 namespace tensorflow {
 
 // Fuse Operation
-template <typename Device, typename T, bool native_format = false>
-class MklFusedMatMulOp : public MklDnnMatMulOpBase<T, void, T> {
+template <typename Device, typename T1, typename T2, typename Tbias,
+          typename Toutput, typename U, bool native_format = false>
+class MklFusedMatMulOp : public MklDnnMatMulOpBase<T2, Tbias, Toutput> {
  public:
   explicit MklFusedMatMulOp(OpKernelConstruction* ctx)
-      : MklDnnMatMulOpBase<T, void, T>(ctx) {
+      : MklDnnMatMulOpBase<T2, Tbias, Toutput>(ctx) {
+    if (std::is_same<T2, qint8>::value) {
+      return;  // Quantized version will have own contstruction code.
+    }
     OP_REQUIRES_OK(ctx, ctx->GetAttr("fused_ops", &fused_ops_));
     OP_REQUIRES_OK(ctx, ctx->GetAttr("transpose_a", &transpose_a_));
     OP_REQUIRES_OK(ctx, ctx->GetAttr("transpose_b", &transpose_b_));
@@ -41,7 +51,6 @@ class MklFusedMatMulOp : public MklDnnMatMulOpBase<T, void, T> {
       OP_REQUIRES_OK(
           ctx, ctx->GetAttr("is_filter_const", &(this->is_weight_const_)));
     }
-
     OP_REQUIRES(ctx, fused_ops_.size() <= 2,
                 absl::InvalidArgumentError(
                     "MklFusedMatMul must have 2 post-arguments at most."));
@@ -54,7 +63,7 @@ class MklFusedMatMulOp : public MklDnnMatMulOpBase<T, void, T> {
         ctx, transpose_a_ == false,
         absl::InvalidArgumentError("In[0] of MklMatMul can't be transposed."));
     if (fused_ops_.size() == 2 && fused_ops_[1] == "LeakyRelu") {
-      OP_REQUIRES_OK(ctx, ctx->GetAttr("leakyrelu_alpha", &leakyrelu_alpha));
+      OP_REQUIRES_OK(ctx, ctx->GetAttr("leakyrelu_alpha", &leakyrelu_alpha_));
     }
   }
 
@@ -64,7 +73,7 @@ class MklFusedMatMulOp : public MklDnnMatMulOpBase<T, void, T> {
     const Tensor& weight_tensor = ctx->input(this->kInputIndexWeight);
     const Tensor& bias_tensor = MklGetInput(ctx, this->kInputIndexBias);
 
-    if (std::is_same<T, float>::value) {
+    if (std::is_same<T1, float>::value) {
       (void)SetFPMathMode();
     }
 
@@ -134,15 +143,16 @@ class MklFusedMatMulOp : public MklDnnMatMulOpBase<T, void, T> {
         memory::format_tag::nc, this->is_weight_const_);
     // Extend the basic parameters for data types and fusions.
     ExtendMklDnnMatMulFwdParams(ctx, matmul_params);
-    auto st = ExecuteSingleThreadedGemm(batch, channel, k, sizeof(T));
+    auto st = ExecuteSingleThreadedGemm(batch, channel, k, sizeof(T1));
     // Create the oneDNN wrapper over Eigen threadpool and set max threads
     // in oneDNN.
     Eigen::ThreadPoolInterface* eigen_interface =
         EigenThreadPoolFromTfContext(ctx);
     tsl::OneDnnThreadPool eigen_tp(eigen_interface, ThreadPoolUseCallerThread(),
                                    st ? 1 : -1);
-    MklDnnMatMulFwdPrimitive<T, T, T, T, T>* matmul_prim =
-        MklDnnMatMulFwdPrimitiveFactory<T, T, T, T, T>::Get(matmul_params, 0);
+    MklDnnMatMulFwdPrimitive<float, T1, T2, Tbias, Toutput>* matmul_prim =
+        MklDnnMatMulFwdPrimitiveFactory<float, T1, T2, Tbias, Toutput>::Get(
+            matmul_params, 0);
 
     // Allocate output tensor.
     Tensor* dst_tensor = nullptr;
@@ -158,17 +168,17 @@ class MklFusedMatMulOp : public MklDnnMatMulOpBase<T, void, T> {
     TensorShape output_tf_shape({batch, channel});
 
     if (fuse_add_) {
-      const Tensor& add_tensor = MklGetInput(ctx, kInputIndex_Add);
+      const Tensor& add_tensor = MklGetInput(ctx, input_idx_add_);
       MklDnnShape add_mkl_shape;
-      GetMklShape(ctx, kInputIndex_Add, &add_mkl_shape, native_format);
+      GetMklShape(ctx, input_idx_add_, &add_mkl_shape, native_format);
 
       // For native format, we need not to set metadata.
-      if (native_format && ctx->forward_input_to_output_with_shape(
-                               kInputIndex_Add, kOutputIndex_Dst,
-                               output_tf_shape, &dst_tensor)) {
+      if (native_format &&
+          ctx->forward_input_to_output_with_shape(
+              input_idx_add_, kOutputIndex_Dst, output_tf_shape, &dst_tensor)) {
         ;  // Need to do nothing for native format
       } else if (!native_format && ForwardMklTensorInToOutWithMklShape(
-                                       ctx, kInputIndex_Add, kOutputIndex_Dst,
+                                       ctx, input_idx_add_, kOutputIndex_Dst,
                                        &dst_tensor, output_mkl_shape, false)) {
         ;  // If it's not native format, need to forward and set meta first
       } else {
@@ -182,19 +192,20 @@ class MklFusedMatMulOp : public MklDnnMatMulOpBase<T, void, T> {
         auto add_md =
             add_mkl_shape.IsMklTensor()
                 ? add_mkl_shape.GetMklLayout()
-                : memory::desc(dst_dims, MklDnnType<T>(), output_format_tag);
+                : memory::desc(dst_dims, MklDnnType<U>(), output_format_tag);
         auto dst_md =
-            memory::desc(dst_dims, MklDnnType<T>(), output_format_tag);
+            memory::desc(dst_dims, MklDnnType<Toutput>(), output_format_tag);
 
         void* add_buf =
-            static_cast<void*>(const_cast<T*>(add_tensor.flat<T>().data()));
-        void* dst_buf = static_cast<void*>((dst_tensor)->flat<T>().data());
+            static_cast<void*>(const_cast<U*>(add_tensor.flat<U>().data()));
+        void* dst_buf =
+            static_cast<void*>((dst_tensor)->flat<Toutput>().data());
 
         if (native_format) {
           // We are simply deep copying the add_tensor to dst_tensor without
           // changing memory layout, hence using same memory descriptor.
           add_md = dst_md =
-              memory::desc({add_tensor.NumElements()}, MklDnnType<T>(),
+              memory::desc({add_tensor.NumElements()}, MklDnnType<U>(),
                            dnnl::memory::format_tag::x);
         }
 
@@ -218,31 +229,33 @@ class MklFusedMatMulOp : public MklDnnMatMulOpBase<T, void, T> {
 
     try {
       // Prepare the input and output for primitive.
-      T* src_data = const_cast<T*>(src_tensor.flat<T>().data());
-      T* weight_data = const_cast<T*>(weight_tensor.flat<T>().data());
-      T* bias_data = const_cast<T*>(bias_tensor.flat<T>().data());
-      T* dst_data = const_cast<T*>(dst_tensor->flat<T>().data());
+      T1* src_data = const_cast<T1*>(src_tensor.flat<T1>().data());
+      T2* weight_data = const_cast<T2*>(weight_tensor.flat<T2>().data());
+      void* bias_data = static_cast<void*>(
+          const_cast<Tbias*>(bias_tensor.flat<Tbias>().data()));
+      Toutput* dst_data =
+          const_cast<Toutput*>(dst_tensor->flat<Toutput>().data());
 
       // Reorder input if necessary.
-      MklDnnData<T> src_mkl(&(this->cpu_engine_));
-      MklDnnData<T> weight_mkl(&(this->cpu_engine_));
+      MklDnnData<T1> src_mkl(&(this->cpu_engine_));
+      MklDnnData<T2> weight_mkl(&(this->cpu_engine_));
 
       auto src_md = src_mkl_shape.IsMklTensor()
                         ? src_mkl_shape.GetMklLayout()
-                        : memory::desc(src_dims, MklDnnType<T>(), src_format);
+                        : memory::desc(src_dims, MklDnnType<T1>(), src_format);
 
       if (src_md != matmul_pd->src_desc()) {
         src_mkl.SetUsrMem(src_md, src_data);
         src_mkl.CheckReorderToOpMem(matmul_pd.get()->src_desc(),
                                     this->cpu_engine_, ctx);
-        src_data = reinterpret_cast<T*>(src_mkl.GetOpMem().get_data_handle());
+        src_data = static_cast<T1*>(src_mkl.GetOpMem().get_data_handle());
       }
 
       // Get cached data when weight is const.
       const memory::desc weight_md =
-          memory::desc(weight_dims, MklDnnType<T>(), weight_format);
+          memory::desc(weight_dims, MklDnnType<T2>(), weight_format);
       if (weight_md != matmul_pd->weights_desc()) {
-        T* cached_weight_data = nullptr;
+        T2* cached_weight_data = nullptr;
 
         if (this->is_weight_const_) {
           // TODO(intel-tf): When oneDNN major version changes to v4.x, weight
@@ -268,16 +281,22 @@ class MklFusedMatMulOp : public MklDnnMatMulOpBase<T, void, T> {
           weight_mkl.CheckReorderToOpMem(matmul_pd.get()->weights_desc(),
                                          this->cpu_engine_, ctx);
           weight_data =
-              reinterpret_cast<T*>(weight_mkl.GetOpMem().get_data_handle());
+              static_cast<T2*>(weight_mkl.GetOpMem().get_data_handle());
         }
       }
       std::shared_ptr<stream> cpu_stream;
-
       cpu_stream.reset(CreateStream(&eigen_tp, matmul_prim->GetEngine()));
 
       UserScratchPad<unsigned char> scratch_pad;
       scratch_pad.AllocateSPTensor(matmul_prim, ctx);
 
+      // Temporary tensor for scaled bias when op is quantized version.
+      Tensor temp_scaled_bias_tensor;
+      if (std::is_same<T2, qint8>::value) {
+        this->GetScaledBias(ctx, matmul_pd, bias_tensor,
+                            &temp_scaled_bias_tensor, &bias_data);
+      }
+
       // Execute fused matmul op.
       matmul_prim->Execute(src_data, weight_data, bias_data, dst_data,
                            matmul_params, scratch_pad.Get(), cpu_stream);
@@ -290,30 +309,31 @@ class MklFusedMatMulOp : public MklDnnMatMulOpBase<T, void, T> {
     }
   }
 
-  void ExtendMklDnnMatMulFwdParams(OpKernelContext* ctx,
-                                   MklDnnMatMulFwdParams& params) {
+  virtual void ExtendMklDnnMatMulFwdParams(OpKernelContext* ctx,
+                                           MklDnnMatMulFwdParams& params) {
+    // Create a string from data types of input, weight, bias, and output.
+    params.dtypes.append(typeid(T1).name());
+    params.dtypes.append(typeid(T2).name());
+    params.dtypes.append(typeid(Tbias).name());
+    params.dtypes.append(typeid(Toutput).name());
     if (fused_ops_.size() == 2) {
       string post_op = fused_ops_[1];
-
-      if (post_op == "Relu") {
-        params.post_op_params.push_back({"relu", {1.0, 0.0, 0.0}});
-      } else if (post_op == "Relu6") {
-        params.post_op_params.push_back({"relu6", {1.0, 6.0, 0.0}});
+      float scale = 1.0f;
+      float alpha = 0.0f;
+      float beta = 0.0f;
+      if (post_op == "Relu6") {
+        alpha = 6.0f;
+      } else if (post_op == "LeakyRelu") {
+        alpha = leakyrelu_alpha_;
       } else if (post_op == "Elu") {
-        params.post_op_params.push_back({"elu", {1.0, 1.0, 0.0}});
-      } else if (post_op == "GeluApproximate") {
-        params.post_op_params.push_back({"gelu_approximate", {1.0, 1.0, 0.0}});
-      } else if (post_op == "GeluExact") {
-        params.post_op_params.push_back({"gelu_exact", {1.0, 1.0, 0.0}});
-      } else if (post_op == "Tanh") {
-        params.post_op_params.push_back({"tanh", {1.0, 0.0, 0.0}});
+        alpha = 1.0f;
+      }
+      if (post_op == "Relu" || post_op == "Relu6" || post_op == "LeakyRelu" ||
+          post_op == "Elu" || post_op == "GeluApproximate" ||
+          post_op == "GeluExact" || post_op == "Tanh" || post_op == "Sigmoid") {
+        params.post_op_params.push_back({post_op, {scale, alpha, beta}});
       } else if (post_op == "Add") {
         params.post_op_params.push_back({"sum", {1.0}});
-      } else if (post_op == "LeakyRelu") {
-        params.post_op_params.push_back(
-            {"leakyrelu", {1.0, leakyrelu_alpha, 0.0}});
-      } else if (post_op == "Sigmoid") {
-        params.post_op_params.push_back({"logistic", {1.0, 0.0, 0.0}});
       } else {
         OP_REQUIRES_OK(ctx, absl::InvalidArgumentError(absl::StrCat(
                                 "Unsupported post-argument in MklFusedMatMul: ",
@@ -322,34 +342,630 @@ class MklFusedMatMulOp : public MklDnnMatMulOpBase<T, void, T> {
     }
   }
 
- private:
+ protected:
+  virtual void GetScaledBias(
+      OpKernelContext*,
+      std::shared_ptr<dnnl::inner_product_forward::primitive_desc>&,
+      const Tensor&, Tensor*, void**) {}
+
   bool fuse_add_ = false;
   bool transpose_a_;
   bool transpose_b_;
-  float leakyrelu_alpha = 0.2;
+  float leakyrelu_alpha_ = 0.2;
   std::vector<string> fused_ops_;
-  const int kInputIndex_Add = 3;
+  int input_idx_add_ = 3;
   const int kOutputIndex_Dst = 0;
-};  // namespace tensorflow
+#ifdef DNNL_AARCH64_USE_ACL
+  const int kWeightTensorHashLength = 1024;
+#endif
+};
+
+namespace {
+
+enum class FusedComputationType {
+  kUndefined,
+  kBiasAdd,
+  kBiasAdd_Dequantize,
+  kBiasAdd_Requantize,
+  kBiasAdd_Activation,
+  kBiasAdd_Activation_Dequantize,
+  kBiasAdd_Activation_Requantize,
+  kBiasAdd_Add,
+  kBiasAdd_Add_Dequantize,
+  kBiasAdd_Add_Requantize,
+};
+
+struct FusedComputationPattern {
+  FusedComputationType fused_computation;
+  std::vector<string> fused_ops;
+};
+
+}  // namespace
+
+// OneDNN uses post-ops to implement different kind of fusions. The category of
+// each individual post-op can be inferred from the fused_ops attribute. The
+// following enum is used to identify list of required post-ops.
+enum class PostOpKind { kActivation, kSum, kOutputScale, kLinear };
+
+template <typename Device, typename T1, typename T2, typename Tbias,
+          typename Toutput, typename U, bool native_format = true>
+class QuantizedFusedMatMulOp
+    : public MklFusedMatMulOp<Device, T1, T2, Tbias, Toutput, U,
+                              native_format> {
+ protected:
+  string input_quant_mode_;   // 0-th input
+  string output_quant_mode_;  // 0-th output
+  string activation_type_;    // Activation op type
+
+  // Initialize minmax tensor indices with default values for the most common
+  // cases.
+  int input_min_idx_ = 3;
+  int input_max_idx_ = 4;
+  int weight_min_idx_ = 5;
+  int weight_max_idx_ = 6;
+
+  struct PostOpInfo {
+    PostOpKind post_op_kind;
+    struct OperandInfo {
+      int idx = -1;  // Operand tensor index if needed by a post-op.
+      // Indices of min and max value tensors, if the operand is quantized.
+      gtl::InlinedVector<int, 4> min_max_indices;
+    } operand_info;
+    // Indices of output min and max value tensors. It is used when requantize
+    // is fused.
+    gtl::InlinedVector<int, 4> min_max_indices;
+  };
+
+  gtl::InlinedVector<PostOpInfo, 4> post_op_info_list_;
+
+  void Initialize(OpKernelConstruction* context) {
+    OP_REQUIRES_OK(context,
+                   context->GetAttr("transpose_a", &this->transpose_a_));
+    OP_REQUIRES_OK(context,
+                   context->GetAttr("transpose_b", &this->transpose_b_));
+    OP_REQUIRES_OK(context,
+                   context->GetAttr("input_quant_mode", &input_quant_mode_));
+    OP_REQUIRES_OK(context,
+                   context->GetAttr("output_quant_mode", &output_quant_mode_));
+    OP_REQUIRES_OK(
+        context, context->GetAttr("is_weight_const", &this->is_weight_const_));
+    OP_REQUIRES_OK(context,
+                   context->GetAttr("is_bias_const", &this->is_bias_const_));
+    if (context->HasAttr("leakyrelu_alpha")) {
+      OP_REQUIRES_OK(context, context->GetAttr("leakyrelu_alpha",
+                                               &this->leakyrelu_alpha_));
+    }
+
+    // Extract activation info and canonicalize activation types to
+    // common name "Activation" in the fused_ops attribute.
+    std::vector<string> fused_ops;
+    OP_REQUIRES_OK(context, context->GetAttr("fused_ops", &fused_ops));
+    for (auto it = fused_ops.begin(); it != fused_ops.end(); ++it) {
+      if (*it == "Relu" || *it == "Relu6" || *it == "Elu" ||
+          *it == "GeluApproximate" || *it == "GeluExact" || *it == "Tanh" ||
+          *it == "LeakyRelu" || *it == "Sigmoid") {
+        if (*it != "Relu") {
+          string last_fusion = fused_ops.back();
+          OP_REQUIRES(
+              context,
+              (last_fusion == "Dequantize" || last_fusion == "Requantize"),
+              absl::UnimplementedError(absl::StrCat(
+                  "Nonlinear activation except Relu can be ",
+                  "supported only with Dequantize or Requantize fusion.")));
+        }
+        activation_type_ = *it;
+        // Canonicalize all activation types into "Activation" for simplifying
+        // post ops construction.
+        *it = "Activation";
+      }
+    }
+
+    using FCT = FusedComputationType;
+
+    // TODO(intel-tf): Add more patterns when implemented.
+    std::vector<FusedComputationPattern> patterns{
+        {FCT::kBiasAdd, {"BiasAdd"}},
+        {FCT::kBiasAdd_Dequantize, {"BiasAdd", "Dequantize"}},
+        {FCT::kBiasAdd_Requantize, {"BiasAdd", "Requantize"}},
+        {FCT::kBiasAdd_Activation, {"BiasAdd", "Activation"}},
+        {FCT::kBiasAdd_Activation_Dequantize,
+         {"BiasAdd", "Activation", "Dequantize"}},
+        {FCT::kBiasAdd_Activation_Requantize,
+         {"BiasAdd", "Activation", "Requantize"}},
+        {FCT::kBiasAdd_Add_Dequantize, {"BiasAdd", "Add", "Dequantize"}},
+    };
+
+    FusedComputationType fused_computation = FusedComputationType::kUndefined;
+    for (const auto& pattern : patterns) {
+      if (fused_ops == pattern.fused_ops) {
+        fused_computation = pattern.fused_computation;
+        break;
+      }
+    }
+
+    // Configure oneDNN post ops
+    switch (fused_computation) {
+      case FCT::kBiasAdd:
+        // No post op is required.
+        OP_REQUIRES(context, (std::is_same<Toutput, qint32>::value),
+                    absl::UnimplementedError(absl::StrCat(
+                        "Qunatized fusion: [", absl::StrJoin(fused_ops, ","),
+                        "] needs output in qint32.")));
+        break;
+      case FCT::kBiasAdd_Dequantize:
+        post_op_info_list_ = {{PostOpKind::kOutputScale, {}, {}}};
+        break;
+      case FCT::kBiasAdd_Requantize:
+        post_op_info_list_ = {{PostOpKind::kOutputScale, {}, {}},
+                              {PostOpKind::kLinear, {}, {7, 8}}};
+        break;
+      case FCT::kBiasAdd_Activation:
+        OP_REQUIRES(context,
+                    (std::is_same<Toutput, qint32>::value &&
+                     activation_type_ == "Relu"),
+                    absl::UnimplementedError(absl::StrCat(
+                        "Qunatized fusion: [", absl::StrJoin(fused_ops, ","),
+                        "] needs output in qint32 and ",
+                        "activation supported is only Relu")));
+        post_op_info_list_ = {{PostOpKind::kActivation, {}, {}}};
+        break;
+      case FCT::kBiasAdd_Activation_Dequantize:
+        post_op_info_list_ = {{PostOpKind::kOutputScale, {}, {}},
+                              {PostOpKind::kActivation, {}, {}}};
+        break;
+      case FCT::kBiasAdd_Activation_Requantize:
+        post_op_info_list_ = {{PostOpKind::kOutputScale, {}, {}},
+                              {PostOpKind::kActivation, {}, {}},
+                              {PostOpKind::kLinear, {}, {7, 8}}};
+        break;
+      case FCT::kBiasAdd_Add_Dequantize: {
+        OP_REQUIRES(
+            context,
+            (std::is_same<U, float>::value || std::is_same<U, bfloat16>::value),
+            absl::UnimplementedError(
+                "Quantized addend tensor is not implemented yet."));
+        // Addend tensor precedes all minmax tensors. Shift the indices from
+        // default initilized values.
+        input_min_idx_ += 1;
+        input_max_idx_ += 1;
+        weight_min_idx_ += 1;
+        weight_max_idx_ += 1;
+        post_op_info_list_ = {{PostOpKind::kOutputScale, {}, {}},
+                              {PostOpKind::kSum, {3, {}}, {}}};
+      } break;
+      default:
+        OP_REQUIRES(context, false,
+                    absl::UnimplementedError(
+                        absl::StrCat("Fusion is not implemented: [",
+                                     absl::StrJoin(fused_ops, ","), "]")));
+    }
+  }
+
+ public:
+  explicit QuantizedFusedMatMulOp(OpKernelConstruction* context)
+      : MklFusedMatMulOp<Device, T1, T2, Tbias, Toutput, U, true>(context) {
+    Initialize(context);
+  }
+
+  void Compute(OpKernelContext* ctx) override {
+    MklFusedMatMulOp<Device, T1, T2, Tbias, Toutput, U, true>::Compute(ctx);
+    // Compute additional outputs
+    if (std::is_same<Toutput, qint8>::value ||
+        std::is_same<Toutput, quint8>::value ||
+        std::is_same<Toutput, qint32>::value) {
+      Tensor* min_output = nullptr;
+      Tensor* max_output = nullptr;
+
+      const float min_input = ctx->input(input_min_idx_).flat<float>()(0);
+      const float max_input = ctx->input(input_max_idx_).flat<float>()(0);
+      const Tensor& min_weight = ctx->input(weight_min_idx_);
+      const Tensor& max_weight = ctx->input(weight_max_idx_);
+      OP_REQUIRES(ctx, min_weight.shape() == max_weight.shape(),
+                  absl::InvalidArgumentError(
+                      "Shape of min-weight and max-weight must be same."));
+
+      if (std::is_same<Toutput, qint32>::value) {
+        TensorShape output_minmax_shape = min_weight.shape();
+        OP_REQUIRES_OK(
+            ctx, ctx->allocate_output(1, output_minmax_shape, &min_output));
+        OP_REQUIRES_OK(
+            ctx, ctx->allocate_output(2, output_minmax_shape, &max_output));
+        if (min_weight.dims() == 0) {
+          float min_output_value;
+          float max_output_value;
+          MklQuantizationRangeForMultiplication<T1, T2, qint32>(
+              min_input, max_input, min_weight.flat<float>()(0),
+              max_weight.flat<float>()(0), &min_output_value,
+              &max_output_value);
+          min_output->flat<float>()(0) = min_output_value;
+          max_output->flat<float>()(0) = max_output_value;
+        } else {
+          MklQuantizationRangeForMultiplication<T1, T2, qint32>(
+              min_input, max_input, min_weight, max_weight, &min_output,
+              &max_output);
+        }
+      } else {
+        // When output type is qint8 or quint8, the kernel is registered for
+        // Requantize fusion.
+        OP_REQUIRES_OK(ctx, ctx->allocate_output(1, {}, &min_output));
+        OP_REQUIRES_OK(ctx, ctx->allocate_output(2, {}, &max_output));
+        int output_min_idx = ctx->num_inputs() - 2;
+        int output_max_idx = ctx->num_inputs() - 1;
+        const float requested_min = ctx->input(output_min_idx).flat<float>()(0);
+        const float requested_max = ctx->input(output_max_idx).flat<float>()(0);
+        if (output_quant_mode_ == "SCALED") {
+          const float range_output =
+              std::max(std::abs(requested_min), std::abs(requested_max));
+          if (std::is_same<Toutput, qint8>::value) {
+            min_output->flat<float>()(0) = -range_output;
+            max_output->flat<float>()(0) = range_output;
+          } else {
+            min_output->flat<float>()(0) = 0;
+            max_output->flat<float>()(0) = range_output;
+          }
+        } else {
+          min_output->flat<float>()(0) = requested_min;
+          max_output->flat<float>()(0) = requested_max;
+        }
+      }
+    } else if (std::is_same<Toutput, float>::value ||
+               std::is_same<Toutput, bfloat16>::value) {
+      // Kernel is registered for Dequantization fusion. Nothing to do.
+    } else {
+      OP_REQUIRES_OK(ctx,
+                     absl::InvalidArgumentError("Unsupported output type."));
+    }
+  }
+
+  void ExtendMklDnnMatMulFwdParams(OpKernelContext* ctx,
+                                   MklDnnMatMulFwdParams& params) override {
+    // Create a string from data types of input, weight, bias, and output.
+    params.dtypes.append(typeid(T1).name());
+    params.dtypes.append(typeid(T2).name());
+    params.dtypes.append(typeid(Tbias).name());
+    params.dtypes.append(typeid(Toutput).name());
+
+    params.input_quant_mode = input_quant_mode_;
+
+    for (const auto& post_op_info : post_op_info_list_) {
+      auto post_op_kind = post_op_info.post_op_kind;
+      switch (post_op_kind) {
+        case PostOpKind::kOutputScale: {
+          if constexpr (std::is_same<Toutput, qint32>::value) {
+            // No scaling is required.
+            break;
+          }
+          const float min_input = ctx->input(input_min_idx_).flat<float>()(0);
+          const float max_input = ctx->input(input_max_idx_).flat<float>()(0);
+          const Tensor& min_weight_tensor = ctx->input(weight_min_idx_);
+          const Tensor& max_weight_tensor = ctx->input(weight_max_idx_);
+          const float* min_weight = min_weight_tensor.flat<float>().data();
+          const float* max_weight = max_weight_tensor.flat<float>().data();
+          const size_t num_weight_scales = min_weight_tensor.NumElements();
+
+          const float max_int8_input =
+              (std::is_same<T1, quint8>::value) ? 255.0f : 127.0f;
+          const float max_int8_weight =
+              (std::is_same<T2, quint8>::value) ? 255.0f : 127.0f;
+          const float range_input =
+              (input_quant_mode_ == "MIN_FIRST")
+                  ? max_input - min_input
+                  : std::max(std::abs(min_input), std::abs(max_input));
+
+          const float src_scale = range_input / max_int8_input;
+          std::vector<float> wei_scales(num_weight_scales);
+#ifndef ENABLE_ONEDNN_V3
+          std::vector<float> output_scales(num_weight_scales);
+#endif  // ENABLE_ONEDNN_V3
+          for (size_t i = 0; i < num_weight_scales; ++i) {
+            float range_weight =
+                std::max(std::abs(min_weight[i]), std::abs(max_weight[i]));
+            wei_scales[i] = range_weight / max_int8_weight;
+#ifndef ENABLE_ONEDNN_V3
+            output_scales[i] = src_scale * wei_scales[i];
+#endif  // ENABLE_ONEDNN_V3
+          }
+          FactoryKeyCreator src_partial_key;
+          src_partial_key.AddAsKey<float>(min_input);
+          src_partial_key.AddAsKey<float>(max_input);
+
+          FactoryKeyCreator wei_partial_key;
+          wei_partial_key.AddAsKey<const float*>(min_weight);
+          wei_partial_key.AddAsKey<const float*>(max_weight);
+#ifndef ENABLE_ONEDNN_V3
+          FactoryKeyCreator output_scales_partial_key;
+          output_scales_partial_key.AddAsKey(src_partial_key.GetKey());
+          output_scales_partial_key.AddAsKey(wei_partial_key.GetKey());
+          params.post_op_params.push_back({"output_scale", output_scales,
+                                           output_scales_partial_key.GetKey()});
+#else
+          params.post_op_params.push_back(
+              {"src_scale", {src_scale}, src_partial_key.GetKey()});
+          params.post_op_params.push_back(
+              {"wei_scale", wei_scales, wei_partial_key.GetKey()});
+#endif  // ENABLE_ONEDNN_V3
+        } break;
+
+        case PostOpKind::kActivation: {
+          float scale = 1.0f;
+          float alpha = 0.0f;
+          float beta = 0.0f;
+          if (activation_type_ == "LeakyRelu")
+            alpha = this->leakyrelu_alpha_;
+          else if (activation_type_ == "Relu6")
+            alpha = 6.0f;
+          else if (activation_type_ == "Elu")
+            alpha = 1.0f;
+          params.post_op_params.push_back(
+              {activation_type_, {scale, alpha, beta}});
+        } break;
+
+        case PostOpKind::kLinear: {
+          // Update output_scale for requantize fusion.
+          auto output_min_idx = post_op_info.min_max_indices[0];
+          auto output_max_idx = post_op_info.min_max_indices[1];
+          const float min_output =
+              ctx->input(output_min_idx).template flat<float>()(0);
+          const float max_output =
+              ctx->input(output_max_idx).template flat<float>()(0);
+          const float max_int8_output =
+              (std::is_same<Toutput, quint8>::value) ? 255.0f : 127.0f;
+          const float range_output =
+              (output_quant_mode_ == "MIN_FIRST")
+                  ? max_output - min_output
+                  : std::max(std::abs(min_output), std::abs(max_output));
+          float req_scale = max_int8_output / range_output;
+          float req_shift = 0.0f;
+          if (output_quant_mode_ == "MIN_FIRST") {
+            req_shift = -min_output * max_int8_output / range_output;
+          }
+          params.post_op_params.push_back(
+              {"linear", {1.0, req_scale, req_shift}});
+        } break;
+
+        case PostOpKind::kSum: {
+          this->fuse_add_ = true;
+          this->input_idx_add_ = post_op_info.operand_info.idx;
+          params.post_op_params.push_back({"sum", {1.0}});
+        } break;
+
+        default:
+          OP_REQUIRES_OK(
+              ctx, absl::InvalidArgumentError("Unsupported post-op-kind."));
+      }
+    }
+  }
+
+  void GetScaledBias(
+      OpKernelContext* ctx,
+      std::shared_ptr<dnnl::inner_product_forward::primitive_desc>& matmul_pd,
+      const Tensor& bias_tensor, Tensor* temp_scaled_bias_tensor,
+      void** bias_data) override {
+#ifdef ENABLE_ONEDNN_V3
+#define TSCALED_BIAS float
+#else
+#define TSCALED_BIAS Tbias
+#endif  // ENABLE_ONEDNN_V3
+
+#ifndef ENABLE_ONEDNN_V3
+    if (std::is_same<Tbias, qint32>::value) {
+      // Bias already has been scaled for quantized input and weight.
+#else
+    if ((std::is_same<Tbias, float>::value ||
+         std::is_same<Tbias, bfloat16>::value) &&
+        input_quant_mode_ == "SCALED") {
+#endif  // !ENABLE_ONEDNN_V3
+      return;
+    } else {
+      const float min_input = ctx->input(input_min_idx_).flat<float>()(0);
+      const float max_input = ctx->input(input_max_idx_).flat<float>()(0);
+      const Tensor& min_weight_tensor = ctx->input(weight_min_idx_);
+      const Tensor& max_weight_tensor = ctx->input(weight_max_idx_);
+      const float* min_weight = min_weight_tensor.flat<float>().data();
+      const float* max_weight = max_weight_tensor.flat<float>().data();
+      bool is_cached_bias_valid = false;
+      bool is_bias_cache_empty = this->IsBiasCacheEmpty();
+      if (!is_bias_cache_empty) {
+        this->GetCachedBias(min_input, max_input, bias_data);
+        is_cached_bias_valid = (*bias_data != nullptr);
+      }
+      if (!is_cached_bias_valid) {
+        void* input_bias_buf = static_cast<void*>(
+            const_cast<Tbias*>(bias_tensor.flat<Tbias>().data()));
+        auto scaled_bias_md = matmul_pd->bias_desc();
+        TensorShape scaled_bias_shape;
+        scaled_bias_shape.AddDim((scaled_bias_md.get_size() / sizeof(float)));
+        OP_REQUIRES_OK(ctx, ctx->allocate_temp(
+                                DataTypeToEnum<TSCALED_BIAS>::v(),
+                                scaled_bias_shape, temp_scaled_bias_tensor));
+        void* scaled_bias_buf = static_cast<void*>(
+            temp_scaled_bias_tensor->flat<TSCALED_BIAS>().data());
+
+        const float max_int8_input =
+            (std::is_same<T1, quint8>::value) ? 255.0f : 127.0f;
+        const float max_int8_weight =
+            (std::is_same<T2, quint8>::value) ? 255.0f : 127.0f;
+        const float range_input =
+            (input_quant_mode_ == "MIN_FIRST")
+                ? max_input - min_input
+                : std::max(std::abs(min_input), std::abs(max_input));
+        const size_t num_weight_scales = min_weight_tensor.NumElements();
+        std::vector<float> bias_scales(num_weight_scales, 1.0);
+        for (size_t i = 0; i < num_weight_scales; ++i) {
+          float range_weight =
+              std::max(std::abs(min_weight[i]), std::abs(max_weight[i]));
+          float scale_factor =
+              (max_int8_input * max_int8_weight) / (range_input * range_weight);
+          bias_scales[i] = scale_factor;
+        }
+        if (input_quant_mode_ == "MIN_FIRST") {
+          Tbias* input_bias = (Tbias*)input_bias_buf;
+          TSCALED_BIAS* adjusted_bias = (TSCALED_BIAS*)scaled_bias_buf;
+          float q_min_input = max_int8_input * min_input / range_input;
+          const Tensor& weight_tensor = ctx->input(1);
+          int stride_ic = 1;
+          int stride_oc = 1;
+          int k = 0;
+          int n = 0;
+          if (this->transpose_b_) {
+            k = weight_tensor.dim_size(1);
+            n = weight_tensor.dim_size(0);
+            stride_ic = 1;
+            stride_oc = k;
+          } else {
+            k = weight_tensor.dim_size(0);
+            n = weight_tensor.dim_size(1);
+            stride_ic = n;
+            stride_oc = 1;
+          }
+          T2* weight_buf = const_cast<T2*>(weight_tensor.flat<T2>().data());
+          std::vector<float> scales(n);
+          if (num_weight_scales == 1) {
+            // Weights are quantized per_tensor. Scales need to be expanded to
+            // number of output channels.
+            std::fill(scales.begin(), scales.end(), bias_scales[0]);
+          } else {
+            scales = bias_scales;
+          }
+          // TODO(intel-tf): Paralellize loop for large weights.
+          for (int j = 0; j < n; ++j) {
+            int sum = 0;
+            for (int i = 0; i < k; ++i) {
+              sum += weight_buf[i * stride_ic + j * stride_oc];
+            }
+#ifndef ENABLE_ONEDNN_V3
+            adjusted_bias[j] = static_cast<TSCALED_BIAS>(
+                (static_cast<float>(input_bias[j]) * scales[j]) +
+                (sum * q_min_input));
+#else
+            // TODO(intel-tf): Use zeropoint for quantized input tensor instead
+            // of manual adjustments.
+            if (std::is_same<Tbias, qint32>::value) {
+              // Starting with oneDNN v3.0, bias is expected to be dequantized
+              // to float32.
+              adjusted_bias[j] = static_cast<float>(input_bias[j]) / scales[j];
+            } else {
+              // Bias is float32 or bfloat16 but still needs to be compensated.
+              adjusted_bias[j] = static_cast<float>(input_bias[j]) +
+                                 ((sum * q_min_input) / scales[j]);
+            }
+#endif  // !ENABLE_ONEDNN_V3
+          }
+        } else {
+          memory::dims input_bias_dims =
+              memory::dims({bias_tensor.shape().dim_size(0)});
+          auto input_bias_md = dnnl::memory::desc(
+              input_bias_dims, MklDnnType<Tbias>(), memory::format_tag::x);
+          auto input_bias_mem =
+              dnnl::memory(input_bias_md, this->cpu_engine_, input_bias_buf);
+          auto scaled_bias_mem =
+              dnnl::memory(scaled_bias_md, this->cpu_engine_, scaled_bias_buf);
+          dnnl::primitive_attr bias_attr;
+#ifndef ENABLE_ONEDNN_V3
+          (num_weight_scales == 1)
+              ? bias_attr.set_output_scales(0, bias_scales)
+              : bias_attr.set_output_scales(1, bias_scales);
+#else
+          (num_weight_scales == 1) ? bias_attr.set_scales_mask(DNNL_ARG_SRC, 0)
+                                   : bias_attr.set_scales_mask(DNNL_ARG_SRC, 1);
+#endif  // !ENABLE_ONEDNN_V3
+          auto reorder_prim =
+              dnnl::reorder(input_bias_mem, scaled_bias_mem, bias_attr);
+          std::unordered_map<int, memory> reorder_net_args = {
+              {DNNL_ARG_FROM, input_bias_mem}, {DNNL_ARG_TO, scaled_bias_mem}};
+#ifdef ENABLE_ONEDNN_V3
+          auto scale_mem =
+              memory({{1}, MklDnnType<float>(), memory::format_tag::x},
+                     this->cpu_engine_, bias_scales.data());
+          reorder_net_args.insert(
+              {DNNL_ARG_ATTR_SCALES | DNNL_ARG_SRC, scale_mem});
+#endif  // ENABLE_ONEDNN_V3
+          reorder_prim.execute(dnnl::stream(this->cpu_engine_),
+                               reorder_net_args);
+        }
+
+        *bias_data = temp_scaled_bias_tensor->flat<float>().data();
+
+        if (is_bias_cache_empty) {
+          // Only try to cache the bias in the first iteration.
+          this->CacheBias(ctx, *temp_scaled_bias_tensor, min_input, max_input);
+        }
+      }
+    }
+  }
+
+  bool IsCachedBiasValid(float current_min_input,
+                         float current_max_input) override
+      TF_LOCKS_EXCLUDED(this->bias_cache_mutex_) {
+    tf_shared_lock lock(this->bias_cache_mutex_);
+    if (this->is_bias_const_ && this->is_weight_const_ &&
+        std::abs(current_min_input - this->saved_min_input_) < 1e-5 &&
+        std::abs(current_max_input - this->saved_max_input_) < 1e-5) {
+      return true;
+    }
+    return false;
+  }
+};
 
 // Register mkl kernels for supported operations and types.
-#define REGISTER_FUSEDMATMUL_MKL_SUPPORTED_KERNELS_TYPES(type)                \
-  REGISTER_KERNEL_BUILDER(                                                    \
-      Name("_MklFusedMatMul")                                                 \
-          .Device(DEVICE_CPU)                                                 \
-          .TypeConstraint<type>("T")                                          \
-          .Label(mkl_op_registry::kMklLayoutDependentOpLabel),                \
-      MklFusedMatMulOp<CPUDevice, type>);                                     \
-  REGISTER_KERNEL_BUILDER(Name("_MklNativeFusedMatMul")                       \
-                              .Device(DEVICE_CPU)                             \
-                              .TypeConstraint<type>("T")                      \
-                              .Label(mkl_op_registry::kMklNameChangeOpLabel), \
-                          MklFusedMatMulOp<CPUDevice, type, true>);
+#define REGISTER_FUSEDMATMUL_MKL_SUPPORTED_KERNELS_TYPES(type)    \
+  REGISTER_KERNEL_BUILDER(                                        \
+      Name("_MklFusedMatMul")                                     \
+          .Device(DEVICE_CPU)                                     \
+          .TypeConstraint<type>("T")                              \
+          .Label(mkl_op_registry::kMklLayoutDependentOpLabel),    \
+      MklFusedMatMulOp<CPUDevice, type, type, type, type, type>); \
+  REGISTER_KERNEL_BUILDER(                                        \
+      Name("_MklNativeFusedMatMul")                               \
+          .Device(DEVICE_CPU)                                     \
+          .TypeConstraint<type>("T")                              \
+          .Label(mkl_op_registry::kMklNameChangeOpLabel),         \
+      MklFusedMatMulOp<CPUDevice, type, type, type, type, type, true>);
 TF_CALL_float(REGISTER_FUSEDMATMUL_MKL_SUPPORTED_KERNELS_TYPES);
 TF_CALL_bfloat16(REGISTER_FUSEDMATMUL_MKL_SUPPORTED_KERNELS_TYPES);
 TF_CALL_half(REGISTER_FUSEDMATMUL_MKL_SUPPORTED_KERNELS_TYPES);
 #undef REGISTER_FUSEDMATMUL_MKL_SUPPORTED_KERNELS_TYPES
 
+#define REGISTER_QUANTIZED_MATMUL(input_type, weight_type, bias_type,       \
+                                  output_type, additional_type)             \
+  REGISTER_KERNEL_BUILDER(                                                  \
+      Name("_QuantizedMatMul")                                              \
+          .Device(DEVICE_CPU)                                               \
+          .TypeConstraint<input_type>("T1")                                 \
+          .TypeConstraint<weight_type>("T2")                                \
+          .TypeConstraint<bias_type>("Tbias")                               \
+          .TypeConstraint<output_type>("Tout")                              \
+          .TypeConstraint<additional_type>("U"),                            \
+      QuantizedFusedMatMulOp<CPUDevice, input_type, weight_type, bias_type, \
+                             output_type, additional_type, true>);
+
+#define REGISTER_ALL_OUTPUT_TYPES(input_type, weight_type, bias_type,     \
+                                  additional_type)                        \
+  REGISTER_QUANTIZED_MATMUL(input_type, weight_type, bias_type, qint8,    \
+                            additional_type)                              \
+  REGISTER_QUANTIZED_MATMUL(input_type, weight_type, bias_type, quint8,   \
+                            additional_type)                              \
+  REGISTER_QUANTIZED_MATMUL(input_type, weight_type, bias_type, qint32,   \
+                            additional_type)                              \
+  REGISTER_QUANTIZED_MATMUL(input_type, weight_type, bias_type, float,    \
+                            additional_type)                              \
+  REGISTER_QUANTIZED_MATMUL(input_type, weight_type, bias_type, bfloat16, \
+                            additional_type)
+
+#define REGISTER_ALL_BIAS_OUTPUT_TYPES(input_type, weight_type,              \
+                                       additional_type)                      \
+  REGISTER_ALL_OUTPUT_TYPES(input_type, weight_type, float, additional_type) \
+  REGISTER_ALL_OUTPUT_TYPES(input_type, weight_type, bfloat16,               \
+                            additional_type)                                 \
+  REGISTER_ALL_OUTPUT_TYPES(input_type, weight_type, qint32, additional_type)
+
+#define REGISTER_ALL_INPUT_BIAS_OUTPUT_TYPES(weight_type, additional_type) \
+  REGISTER_ALL_BIAS_OUTPUT_TYPES(qint8, weight_type, additional_type)      \
+  REGISTER_ALL_BIAS_OUTPUT_TYPES(quint8, weight_type, additional_type)
+
+REGISTER_ALL_INPUT_BIAS_OUTPUT_TYPES(qint8, float);
+REGISTER_ALL_INPUT_BIAS_OUTPUT_TYPES(qint8, bfloat16);
+
 }  // namespace tensorflow
 
 #endif  // INTEL_MKL
diff --git a/tensorflow/core/kernels/mkl/mkl_matmul_ops_common.h b/tensorflow/core/kernels/mkl/mkl_matmul_ops_common.h
index 922e6464663bb5..6159b1d047ea65 100644
--- a/tensorflow/core/kernels/mkl/mkl_matmul_ops_common.h
+++ b/tensorflow/core/kernels/mkl/mkl_matmul_ops_common.h
@@ -21,7 +21,6 @@ limitations under the License.
 #include <string>
 #include <vector>
 
-#include "unsupported/Eigen/CXX11/Tensor"  // from @eigen_archive
 #include "dnnl.hpp"
 #include "tensorflow/core/framework/op.h"
 #include "tensorflow/core/framework/op_kernel.h"
@@ -29,6 +28,7 @@ limitations under the License.
 #include "tensorflow/core/kernels/mkl/mkl_kernel_util.h"
 #include "tensorflow/core/util/mkl_util.h"
 #include "tensorflow/core/util/onednn_env_vars.h"
+#include "unsupported/Eigen/CXX11/Tensor"  // from @eigen_archive
 #if defined(DNNL_AARCH64_USE_ACL) && defined(ENABLE_ONEDNN_OPENMP)
 #include "tensorflow/core/platform/mutex.h"
 #endif
@@ -102,8 +102,10 @@ struct MklDnnMatMulFwdParams {
   struct PostOpParam {
     string name;
     std::vector<float> param;
+    string partial_key = string("");
   };
   std::vector<PostOpParam> post_op_params;
+  string input_quant_mode;
 
   MklDnnMatMulFwdParams(
       memory::dims src_dims, memory::dims weight_dims, memory::dims bias_dims,
@@ -244,7 +246,7 @@ class MklDnnMatMulFwdPrimitive : public MklPrimitive {
           dst_scale_mem(nullptr),
 #ifndef ENABLE_ONEDNN_V3
           fwd_desc(nullptr),
-#endif  // !ENABLE_ONEDNN_V3
+#endif  // ENABLE_ONEDNN_V3
           fwd_pd(nullptr),
           src_md(nullptr),
           weight_md(nullptr),
@@ -276,15 +278,26 @@ class MklDnnMatMulFwdPrimitive : public MklPrimitive {
                                            MklDnnType<Toutput>(),
                                            matmul_fwd_params.dst_format));
 
-    if (std::is_same<Tbias, qint32>::value) {
-      context_.bias_md.reset(new memory::desc({matmul_fwd_params.bias_dims},
-                                              MklDnnType<TSCALED_BIAS>(),
-                                              memory::format_tag::any));
+    memory::data_type bias_dt;
+#ifndef ENABLE_ONEDNN_V3
+    bias_dt = MklDnnType<Tbias>();
+#else
+    if (std::is_same<Tweight, qint8>::value) {
+      // For QuantizedMatMul, bias needs to be passed to oneDNN as float of
+      // bfloat16 (even if Tbias is qint32).
+      if (std::is_same<Tbias, bfloat16>::value &&
+          matmul_fwd_params.input_quant_mode == "SCALED") {
+        bias_dt = MklDnnType<bfloat16>();
+      } else {
+        bias_dt = MklDnnType<float>();
+      }
     } else {
-      context_.bias_md.reset(new memory::desc({matmul_fwd_params.bias_dims},
-                                              MklDnnType<Tbias>(),
-                                              memory::format_tag::any));
+      bias_dt = MklDnnType<Tbias>();
     }
+#endif  // !ENABLE_ONEDNN_V3
+    context_.bias_md.reset(new memory::desc({matmul_fwd_params.bias_dims},
+                                            bias_dt, memory::format_tag::any));
+
     // Create an inner-product.
 #ifndef ENABLE_ONEDNN_V3
     context_.fwd_desc.reset(new inner_product_forward::desc(
@@ -304,60 +317,68 @@ class MklDnnMatMulFwdPrimitive : public MklPrimitive {
     std::unordered_map<string, bool> is_scale_set;
     if (!post_op_params.empty()) {
       for (auto const& post_op_param : post_op_params) {
-        if (post_op_param.name == "relu" || post_op_param.name == "leakyrelu") {
+        if (post_op_param.name == "Relu" || post_op_param.name == "LeakyRelu") {
           DCHECK_EQ(post_op_param.param.size(), 3);
           float op_scale = post_op_param.param[0];
           float op_alpha = post_op_param.param[1];
           float op_beta = post_op_param.param[2];
           post_ops.APPEND_ELTWISE(op_scale, dnnl::algorithm::eltwise_relu,
                                   op_alpha, op_beta);
-        } else if (post_op_param.name == "relu6") {
+        } else if (post_op_param.name == "Relu6") {
           DCHECK_EQ(post_op_param.param.size(), 3);
           float op_scale = post_op_param.param[0];
           float op_alpha = post_op_param.param[1];
           float op_beta = post_op_param.param[2];
           post_ops.APPEND_ELTWISE_RELU6(op_scale, op_alpha, op_beta);
-        } else if (post_op_param.name == "elu") {
+        } else if (post_op_param.name == "Elu") {
           DCHECK_EQ(post_op_param.param.size(), 3);
           float op_scale = post_op_param.param[0];
           float op_alpha = post_op_param.param[1];
           float op_beta = post_op_param.param[2];
           post_ops.APPEND_ELTWISE(op_scale, dnnl::algorithm::eltwise_elu,
                                   op_alpha, op_beta);
-        } else if (post_op_param.name == "gelu_approximate") {
+        } else if (post_op_param.name == "GeluApproximate") {
           DCHECK_EQ(post_op_param.param.size(), 3);
           float op_scale = post_op_param.param[0];
           float op_alpha = post_op_param.param[1];
           float op_beta = post_op_param.param[2];
           post_ops.APPEND_ELTWISE(op_scale, dnnl::algorithm::eltwise_gelu_tanh,
                                   op_alpha, op_beta);
-        } else if (post_op_param.name == "gelu_exact") {
+        } else if (post_op_param.name == "GeluExact") {
           DCHECK_EQ(post_op_param.param.size(), 3);
           float op_scale = post_op_param.param[0];
           float op_alpha = post_op_param.param[1];
           float op_beta = post_op_param.param[2];
           post_ops.APPEND_ELTWISE(op_scale, dnnl::algorithm::eltwise_gelu_erf,
                                   op_alpha, op_beta);
-        } else if (post_op_param.name == "tanh") {
+        } else if (post_op_param.name == "Tanh") {
           DCHECK_EQ(post_op_param.param.size(), 3);
           float op_scale = post_op_param.param[0];
           float op_alpha = post_op_param.param[1];
           float op_beta = post_op_param.param[2];
           post_ops.APPEND_ELTWISE(op_scale, dnnl::algorithm::eltwise_tanh,
                                   op_alpha, op_beta);
-        } else if (post_op_param.name == "logistic") {
+        } else if (post_op_param.name == "Sigmoid") {
           DCHECK_EQ(post_op_param.param.size(), 3);
           float op_scale = post_op_param.param[0];
           float op_alpha = post_op_param.param[1];
           float op_beta = post_op_param.param[2];
           post_ops.APPEND_ELTWISE(op_scale, dnnl::algorithm::eltwise_logistic,
                                   op_alpha, op_beta);
+        } else if (post_op_param.name == "linear") {
+          DCHECK_EQ(post_op_param.param.size(), 3);
+          float op_scale = post_op_param.param[0];
+          float op_alpha = post_op_param.param[1];
+          float op_beta = post_op_param.param[2];
+          post_ops.APPEND_ELTWISE(op_scale, dnnl::algorithm::eltwise_linear,
+                                  op_alpha, op_beta);
 #ifndef ENABLE_ONEDNN_V3
         } else if (post_op_param.name == "output_scale") {
-          DCHECK_EQ(post_op_param.param.size(), 1);
-          std::vector<float> scales;
-          scales.push_back(post_op_param.param[0]);
-          post_ops_attr.set_output_scales(0, scales);
+          if (post_op_param.param.size() == 1) {
+            post_ops_attr.set_output_scales(0, post_op_param.param);
+          } else {
+            post_ops_attr.set_output_scales(2, post_op_param.param);
+          }
 #else
         } else if (post_op_param.name == "src_scale") {
           is_scale_set.insert({"src", true});
@@ -368,14 +389,18 @@ class MklDnnMatMulFwdPrimitive : public MklPrimitive {
               new memory(*context_.src_scale_md, cpu_engine_, DummyData));
         } else if (post_op_param.name == "wei_scale") {
           is_scale_set.insert({"wei", true});
-          post_ops_attr.set_scales_mask(DNNL_ARG_WEIGHTS, 0);
-          context_.wei_scale_md.reset(new memory::desc({1}, MklDnnType<float>(),
-                                                       memory::format_tag::x));
+          const int scale_size = post_op_param.param.size();
+          const int mask = scale_size == 1 ? 0 : 1;
+          post_ops_attr.set_scales_mask(DNNL_ARG_WEIGHTS, mask);
+          context_.wei_scale_md.reset(new memory::desc(
+              {scale_size}, MklDnnType<float>(), memory::format_tag::x));
           context_.wei_scale_mem.reset(
               new memory(*context_.wei_scale_md, cpu_engine_, DummyData));
         } else if (post_op_param.name == "dst_scale") {
           is_scale_set.insert({"dst", true});
-          post_ops_attr.set_scales_mask(DNNL_ARG_DST, 0);
+          const int scale_size = post_op_param.param.size();
+          const int mask = scale_size == 1 ? 0 : 1;
+          post_ops_attr.set_scales_mask(DNNL_ARG_DST, mask);
           context_.dst_scale_md.reset(new memory::desc({1}, MklDnnType<float>(),
                                                        memory::format_tag::x));
           context_.dst_scale_mem.reset(
@@ -387,13 +412,15 @@ class MklDnnMatMulFwdPrimitive : public MklPrimitive {
           post_ops.append_sum(op_scale);
 
         } else {
-          DCHECK((post_op_param.name == "relu") ||
-                 (post_op_param.name == "relu6") ||
-                 (post_op_param.name == "elu") ||
-                 (post_op_param.name == "tanh") ||
-                 (post_op_param.name == "logistic") ||
+          DCHECK((post_op_param.name == "Relu") ||
+                 (post_op_param.name == "Relu6") ||
+                 (post_op_param.name == "Elu") ||
+                 (post_op_param.name == "GeluApproximate") ||
+                 (post_op_param.name == "GeluExact") ||
+                 (post_op_param.name == "Tanh") ||
+                 (post_op_param.name == "Sigmoid") ||
                  (post_op_param.name == "sum") ||
-                 (post_op_param.name == "leakyrelu") || OUTPUT_SCALE_DCHECK);
+                 (post_op_param.name == "Leakyrelu") || OUTPUT_SCALE_DCHECK);
         }
       }
       post_ops_attr.set_post_ops(post_ops);
@@ -433,11 +460,15 @@ class MklDnnMatMulFwdPrimitive : public MklPrimitive {
         {DNNL_ARG_SCRATCHPAD, *context_.sp_mem},
         {DNNL_ARG_DST, *context_.dst_mem}};
 #ifdef ENABLE_ONEDNN_V3
-    if (is_scale_set["src"] && is_scale_set["wei"] && is_scale_set["dst"]) {
+    if (is_scale_set["src"]) {
       net_args.insert(
           {DNNL_ARG_ATTR_SCALES | DNNL_ARG_SRC, *context_.src_scale_mem});
+    }
+    if (is_scale_set["wei"]) {
       net_args.insert(
           {DNNL_ARG_ATTR_SCALES | DNNL_ARG_WEIGHTS, *context_.wei_scale_mem});
+    }
+    if (is_scale_set["dst"]) {
       net_args.insert(
           {DNNL_ARG_ATTR_SCALES | DNNL_ARG_DST, *context_.dst_scale_mem});
     }
@@ -510,12 +541,12 @@ class MklDnnMatMulFwdPrimitiveFactory : public MklPrimitiveFactory<T> {
 
     // Generate keys for post-ops
     for (auto const& post_op_param : mkldnn_matmul_fwd_dims.post_op_params) {
-      if (post_op_param.name == "relu" || post_op_param.name == "relu6" ||
-          post_op_param.name == "elu" || post_op_param.name == "tanh" ||
-          post_op_param.name == "logistic" ||
-          post_op_param.name == "leakyrelu" ||
-          post_op_param.name == "gelu_approximate" ||
-          post_op_param.name == "gelu_exact") {
+      if (post_op_param.name == "Relu" || post_op_param.name == "Relu6" ||
+          post_op_param.name == "Elu" || post_op_param.name == "Tanh" ||
+          post_op_param.name == "Sigmoid" ||
+          post_op_param.name == "LeakyRelu" ||
+          post_op_param.name == "GeluApproximate" ||
+          post_op_param.name == "GeluExact" || post_op_param.name == "linear") {
         DCHECK_EQ(post_op_param.param.size(), 3);
         key_creator.AddAsKey(post_op_param.name);
         key_creator.AddAsKey(post_op_param.param[0]);
@@ -532,9 +563,16 @@ class MklDnnMatMulFwdPrimitiveFactory : public MklPrimitiveFactory<T> {
                  post_op_param.name == "wei_scale" ||
                  post_op_param.name == "dst_scale") {
 #endif  // !ENABLE_ONEDNN_V3
-        DCHECK_EQ(post_op_param.param.size(), 1);
         key_creator.AddAsKey(post_op_param.name);
-        key_creator.AddAsKey(post_op_param.param[0]);
+        if (post_op_param.partial_key.empty()) {
+          DCHECK_GE(post_op_param.param.size(), 1);
+          // Old Quantized MatMul kernels do not create part of key beforehand
+          // as primitive caching-key-creation optimization.
+          key_creator.AddAsKey(post_op_param.param[0]);
+        } else {
+          // New Quantized MatMul kernels pre-create partial key.
+          key_creator.AddAsKey(post_op_param.partial_key);
+        }
       } else {
         return string("not_a_key");
       }
diff --git a/tensorflow/core/kernels/mkl/mkl_qmatmul_op.cc b/tensorflow/core/kernels/mkl/mkl_qmatmul_op.cc
index efb33375d1669d..92a21e1255c778 100644
--- a/tensorflow/core/kernels/mkl/mkl_qmatmul_op.cc
+++ b/tensorflow/core/kernels/mkl/mkl_qmatmul_op.cc
@@ -783,7 +783,7 @@ class MklDnnQuantizedMatMulReluOp
     MklDnnQuantizedMatMulOp<Device, quint8, qint8, Tbias, Toutput,
                             native_format>::ExtendMklDnnMatMulFwdParams(context,
                                                                         params);
-    params.post_op_params.push_back({"relu", {1.0, 0.0, 0.0}});
+    params.post_op_params.push_back({"Relu", {1.0, 0.0, 0.0}});
   }
 };
 
diff --git a/tensorflow/core/kernels/mkl/mkl_qmatmul_op_test.cc b/tensorflow/core/kernels/mkl/mkl_qmatmul_op_test.cc
index 22b56e19e3bb63..3d862e0b8e8fd5 100644
--- a/tensorflow/core/kernels/mkl/mkl_qmatmul_op_test.cc
+++ b/tensorflow/core/kernels/mkl/mkl_qmatmul_op_test.cc
@@ -35,23 +35,43 @@ limitations under the License.
 
 namespace tensorflow {
 
-class QuantizedMatMulTest : public OpsTestBase {};
+class QuantizedMatMulTest : public OpsTestBase,
+                            public ::testing::WithParamInterface<bool> {};
 
 // Two small matrices A of type uint8 and B of type int8  are multiplied
 // and the result is added with int32 bias
-TEST_F(QuantizedMatMulTest, Small_withBias) {
-  TF_ASSERT_OK(
-      NodeDefBuilder("quantized_mat_mul_op", "_MklQuantizedMatMulWithBias")
-          .Input(FakeInput(DT_QUINT8))
-          .Input(FakeInput(DT_QINT8))
-          .Input(FakeInput(DT_QINT32))
-          .Input(FakeInput(DT_FLOAT))
-          .Input(FakeInput(DT_FLOAT))
-          .Input(FakeInput(DT_FLOAT))
-          .Input(FakeInput(DT_FLOAT))
-          .Attr("Toutput", DataTypeToEnum<qint32>::v())
-          .Attr("_kernel", "QuantizedMklOp")
-          .Finalize(node_def()));
+TEST_P(QuantizedMatMulTest, Small_withBias) {
+  const bool is_old_api = GetParam();
+  if (is_old_api) {
+    TF_ASSERT_OK(
+        NodeDefBuilder("quantized_mat_mul_op", "_MklQuantizedMatMulWithBias")
+            .Input(FakeInput(DT_QUINT8))
+            .Input(FakeInput(DT_QINT8))
+            .Input(FakeInput(DT_QINT32))
+            .Input(FakeInput(DT_FLOAT))
+            .Input(FakeInput(DT_FLOAT))
+            .Input(FakeInput(DT_FLOAT))
+            .Input(FakeInput(DT_FLOAT))
+            .Attr("Toutput", DataTypeToEnum<qint32>::v())
+            .Attr("_kernel", "QuantizedMklOp")
+            .Finalize(node_def()));
+  } else {
+    TF_ASSERT_OK(
+        NodeDefBuilder("quantized_mat_mul_op", "_QuantizedMatMul")
+            .Attr("Thost_inputs", {DT_QUINT8, DT_QINT8, DT_QINT32, DT_FLOAT,
+                                   DT_FLOAT, DT_FLOAT, DT_FLOAT})
+            .Attr("Thost_outputs", {DT_QINT32, DT_FLOAT, DT_FLOAT})
+            .Attr("Tdevice_inputs", std::vector<DataType>())
+            .Attr("Tdevice_outputs", std::vector<DataType>())
+            .Attr("T1", DT_QUINT8)
+            .Attr("T2", DT_QINT8)
+            .Attr("Tbias", DT_QINT32)
+            .Attr("Tout", DT_QINT32)
+            .Attr("fused_ops", {"BiasAdd"})
+            .Input(FakeInput())
+            .Input(FakeInput())
+            .Finalize(node_def()));
+  }
   TF_ASSERT_OK(InitOp());
   // A matrix is:
   // |  1 |  2 |  3 |
@@ -91,19 +111,38 @@ TEST_F(QuantizedMatMulTest, Small_withBias) {
 
 // Two small matrices A of type uint8 and B of type int8  are multiplied
 // and the result is added with neg bias as well
-TEST_F(QuantizedMatMulTest, Small_withNegBias) {
-  TF_ASSERT_OK(
-      NodeDefBuilder("quantized_mat_mul_op", "_MklQuantizedMatMulWithBias")
-          .Input(FakeInput(DT_QUINT8))
-          .Input(FakeInput(DT_QINT8))
-          .Input(FakeInput(DT_QINT32))
-          .Input(FakeInput(DT_FLOAT))
-          .Input(FakeInput(DT_FLOAT))
-          .Input(FakeInput(DT_FLOAT))
-          .Input(FakeInput(DT_FLOAT))
-          .Attr("Toutput", DataTypeToEnum<qint32>::v())
-          .Attr("_kernel", "QuantizedMklOp")
-          .Finalize(node_def()));
+TEST_P(QuantizedMatMulTest, Small_withNegBias) {
+  const bool is_old_api = GetParam();
+  if (is_old_api) {
+    TF_ASSERT_OK(
+        NodeDefBuilder("quantized_mat_mul_op", "_MklQuantizedMatMulWithBias")
+            .Input(FakeInput(DT_QUINT8))
+            .Input(FakeInput(DT_QINT8))
+            .Input(FakeInput(DT_QINT32))
+            .Input(FakeInput(DT_FLOAT))
+            .Input(FakeInput(DT_FLOAT))
+            .Input(FakeInput(DT_FLOAT))
+            .Input(FakeInput(DT_FLOAT))
+            .Attr("Toutput", DataTypeToEnum<qint32>::v())
+            .Attr("_kernel", "QuantizedMklOp")
+            .Finalize(node_def()));
+  } else {
+    TF_ASSERT_OK(
+        NodeDefBuilder("quantized_mat_mul_op", "_QuantizedMatMul")
+            .Attr("Thost_inputs", {DT_QUINT8, DT_QINT8, DT_QINT32, DT_FLOAT,
+                                   DT_FLOAT, DT_FLOAT, DT_FLOAT})
+            .Attr("Thost_outputs", {DT_QINT32, DT_FLOAT, DT_FLOAT})
+            .Attr("Tdevice_inputs", std::vector<DataType>())
+            .Attr("Tdevice_outputs", std::vector<DataType>())
+            .Attr("T1", DT_QUINT8)
+            .Attr("T2", DT_QINT8)
+            .Attr("Tbias", DT_QINT32)
+            .Attr("Tout", DT_QINT32)
+            .Attr("fused_ops", {"BiasAdd"})
+            .Input(FakeInput())
+            .Input(FakeInput())
+            .Finalize(node_def()));
+  }
   TF_ASSERT_OK(InitOp());
   // A matrix is:
   // |  1 |  2 |  3 |
@@ -144,20 +183,40 @@ TEST_F(QuantizedMatMulTest, Small_withNegBias) {
 
 // Two small matrices A of type uint8 (converted from signed integer)
 // and B of type int8  are multiplied and the result is added with float bias
-TEST_F(QuantizedMatMulTest, Small_WithNegInp) {
-  TF_ASSERT_OK(
-      NodeDefBuilder("quantized_mat_mul_op", "_MklQuantizedMatMulWithBias")
-          .Input(FakeInput(DT_QUINT8))
-          .Input(FakeInput(DT_QINT8))
-          .Input(FakeInput(DT_FLOAT))
-          .Input(FakeInput(DT_FLOAT))
-          .Input(FakeInput(DT_FLOAT))
-          .Input(FakeInput(DT_FLOAT))
-          .Input(FakeInput(DT_FLOAT))
-          .Attr("Toutput", DataTypeToEnum<qint32>::v())
-          .Attr("input_quant_mode", "MIN_FIRST")
-          .Attr("_kernel", "QuantizedMklOp")
-          .Finalize(node_def()));
+TEST_P(QuantizedMatMulTest, Small_WithNegInp) {
+  const bool is_old_api = GetParam();
+  if (is_old_api) {
+    TF_ASSERT_OK(
+        NodeDefBuilder("quantized_mat_mul_op", "_MklQuantizedMatMulWithBias")
+            .Input(FakeInput(DT_QUINT8))
+            .Input(FakeInput(DT_QINT8))
+            .Input(FakeInput(DT_FLOAT))
+            .Input(FakeInput(DT_FLOAT))
+            .Input(FakeInput(DT_FLOAT))
+            .Input(FakeInput(DT_FLOAT))
+            .Input(FakeInput(DT_FLOAT))
+            .Attr("Toutput", DataTypeToEnum<qint32>::v())
+            .Attr("input_quant_mode", "MIN_FIRST")
+            .Attr("_kernel", "QuantizedMklOp")
+            .Finalize(node_def()));
+  } else {
+    TF_ASSERT_OK(
+        NodeDefBuilder("quantized_mat_mul_op", "_QuantizedMatMul")
+            .Attr("Thost_inputs", {DT_QUINT8, DT_QINT8, DT_FLOAT, DT_FLOAT,
+                                   DT_FLOAT, DT_FLOAT, DT_FLOAT})
+            .Attr("Thost_outputs", {DT_QINT32, DT_FLOAT, DT_FLOAT})
+            .Attr("Tdevice_inputs", std::vector<DataType>())
+            .Attr("Tdevice_outputs", std::vector<DataType>())
+            .Attr("T1", DT_QUINT8)
+            .Attr("T2", DT_QINT8)
+            .Attr("Tbias", DT_FLOAT)
+            .Attr("Tout", DT_QINT32)
+            .Attr("fused_ops", {"BiasAdd"})
+            .Attr("input_quant_mode", "MIN_FIRST")
+            .Input(FakeInput())
+            .Input(FakeInput())
+            .Finalize(node_def()));
+  }
   TF_ASSERT_OK(InitOp());
   // The A matrix is:
   // |  -1 |  -5 |  -9 |
@@ -213,21 +272,41 @@ TEST_F(QuantizedMatMulTest, Small_WithNegInp) {
 
 // Two small matrices A of type uint8 and B of type int8  are multiplied
 // and the result is added with int32 bias and Requantization fusion
-TEST_F(QuantizedMatMulTest, Small_withBiasAndReq) {
-  TF_ASSERT_OK(NodeDefBuilder("quantized_mat_mul_op",
-                              "_MklQuantizedMatMulWithBiasAndRequantize")
-                   .Input(FakeInput(DT_QUINT8))
-                   .Input(FakeInput(DT_QINT8))
-                   .Input(FakeInput(DT_QINT32))
-                   .Input(FakeInput(DT_FLOAT))
-                   .Input(FakeInput(DT_FLOAT))
-                   .Input(FakeInput(DT_FLOAT))
-                   .Input(FakeInput(DT_FLOAT))
-                   .Input(FakeInput(DT_FLOAT))
-                   .Input(FakeInput(DT_FLOAT))
-                   .Attr("Toutput", DataTypeToEnum<quint8>::v())
-                   .Attr("_kernel", "QuantizedMklOp")
-                   .Finalize(node_def()));
+TEST_P(QuantizedMatMulTest, Small_withBiasAndReq) {
+  const bool is_old_api = GetParam();
+  if (is_old_api) {
+    TF_ASSERT_OK(NodeDefBuilder("quantized_mat_mul_op",
+                                "_MklQuantizedMatMulWithBiasAndRequantize")
+                     .Input(FakeInput(DT_QUINT8))
+                     .Input(FakeInput(DT_QINT8))
+                     .Input(FakeInput(DT_QINT32))
+                     .Input(FakeInput(DT_FLOAT))
+                     .Input(FakeInput(DT_FLOAT))
+                     .Input(FakeInput(DT_FLOAT))
+                     .Input(FakeInput(DT_FLOAT))
+                     .Input(FakeInput(DT_FLOAT))
+                     .Input(FakeInput(DT_FLOAT))
+                     .Attr("Toutput", DataTypeToEnum<quint8>::v())
+                     .Attr("_kernel", "QuantizedMklOp")
+                     .Finalize(node_def()));
+  } else {
+    TF_ASSERT_OK(NodeDefBuilder("quantized_mat_mul_op", "_QuantizedMatMul")
+                     .Attr("Thost_inputs",
+                           {DT_QUINT8, DT_QINT8, DT_QINT32, DT_FLOAT, DT_FLOAT,
+                            DT_FLOAT, DT_FLOAT, DT_FLOAT, DT_FLOAT})
+                     .Attr("Thost_outputs", {DT_QUINT8, DT_FLOAT, DT_FLOAT})
+                     .Attr("Tdevice_inputs", std::vector<DataType>())
+                     .Attr("Tdevice_outputs", std::vector<DataType>())
+                     .Attr("T1", DT_QUINT8)
+                     .Attr("T2", DT_QINT8)
+                     .Attr("Tbias", DT_QINT32)
+                     .Attr("Tout", DT_QUINT8)
+                     .Attr("fused_ops", {"BiasAdd", "Requantize"})
+                     .Input(FakeInput())
+                     .Input(FakeInput())
+                     .Finalize(node_def()));
+  }
+
   TF_ASSERT_OK(InitOp());
   // A matrix is:
   // |  1 |  2 |  3 |
@@ -273,11 +352,25 @@ TEST_F(QuantizedMatMulTest, Small_withBiasAndReq) {
   // 178 * 1.00392 ~= 178.698 ~= 179
 
   Tensor expected(allocator(), DT_QUINT8, TensorShape({2, 4}));
+  if (is_old_api) {
 #ifdef ENABLE_ONEDNN_V3
-  test::FillValues<quint8>(&expected, {84, 60, 116, 52, 183, 168, 233, 178});
+    test::FillValues<quint8>(&expected, {84, 60, 116, 52, 183, 168, 233, 178});
 #else
-  test::FillValues<quint8>(&expected, {84, 60, 116, 52, 184, 169, 234, 179});
+    test::FillValues<quint8>(&expected, {84, 60, 116, 52, 184, 169, 234, 179});
 #endif  // ENABLE_ONEDNN_V3
+  } else {
+    // New api uses more numerical precision preserving equation. Old api scales
+    // up to 32-bit and then scales down from 32-bit to 8-bit. New api instead
+    // does a dequantization followed by a scaling to 8-bit.
+    // In this test,
+    //    input deq. scale =  ((255.0 * 127.0) / (255.0 * 127.0)) = 1.0
+    //    output req. scale = 255.0 / 255.0 = 1.0
+    //    combined scale = 1.0 * 1.0 = 1.0
+    // Note: new api scale value is 1.0, whereas the old api scale is 1.000392.
+    // Correct value is 1.0f. The closer it is to the correct value the better
+    // the formula is.
+    test::FillValues<quint8>(&expected, {84, 60, 116, 52, 183, 168, 233, 178});
+  }
 
   const Tensor& output = *GetOutput(0);
   test::ExpectTensorEqual<quint8>(expected, output);
@@ -285,21 +378,40 @@ TEST_F(QuantizedMatMulTest, Small_withBiasAndReq) {
 
 // Two small matrices A of type uint8 and B of type int8  are multiplied
 // and the result is added with int32 bias and Requantization fusion
-TEST_F(QuantizedMatMulTest, Small_withBiasAndDeq) {
-  TF_ASSERT_OK(NodeDefBuilder("quantized_mat_mul_op",
-                              "_MklQuantizedMatMulWithBiasAndDequantize")
-                   .Input(FakeInput(DT_QUINT8))
-                   .Input(FakeInput(DT_QINT8))
-                   .Input(FakeInput(DT_QINT32))
-                   .Input(FakeInput(DT_FLOAT))
-                   .Input(FakeInput(DT_FLOAT))
-                   .Input(FakeInput(DT_FLOAT))
-                   .Input(FakeInput(DT_FLOAT))
-                   .Input(FakeInput(DT_FLOAT))
-                   .Input(FakeInput(DT_FLOAT))
-                   .Attr("Toutput", DataTypeToEnum<float>::v())
-                   .Attr("_kernel", "QuantizedMklOp")
-                   .Finalize(node_def()));
+TEST_P(QuantizedMatMulTest, Small_withBiasAndDeq) {
+  const bool is_old_api = GetParam();
+  if (is_old_api) {
+    TF_ASSERT_OK(NodeDefBuilder("quantized_mat_mul_op",
+                                "_MklQuantizedMatMulWithBiasAndDequantize")
+                     .Input(FakeInput(DT_QUINT8))
+                     .Input(FakeInput(DT_QINT8))
+                     .Input(FakeInput(DT_QINT32))
+                     .Input(FakeInput(DT_FLOAT))
+                     .Input(FakeInput(DT_FLOAT))
+                     .Input(FakeInput(DT_FLOAT))
+                     .Input(FakeInput(DT_FLOAT))
+                     .Input(FakeInput(DT_FLOAT))
+                     .Input(FakeInput(DT_FLOAT))
+                     .Attr("Toutput", DataTypeToEnum<float>::v())
+                     .Attr("_kernel", "QuantizedMklOp")
+                     .Finalize(node_def()));
+  } else {
+    TF_ASSERT_OK(
+        NodeDefBuilder("quantized_mat_mul_op", "_QuantizedMatMul")
+            .Attr("Thost_inputs", {DT_QUINT8, DT_QINT8, DT_QINT32, DT_FLOAT,
+                                   DT_FLOAT, DT_FLOAT, DT_FLOAT})
+            .Attr("Thost_outputs", {DT_FLOAT})
+            .Attr("Tdevice_inputs", std::vector<DataType>())
+            .Attr("Tdevice_outputs", std::vector<DataType>())
+            .Attr("T1", DT_QUINT8)
+            .Attr("T2", DT_QINT8)
+            .Attr("Tbias", DT_QINT32)
+            .Attr("Tout", DT_FLOAT)
+            .Attr("fused_ops", {"BiasAdd", "Dequantize"})
+            .Input(FakeInput())
+            .Input(FakeInput())
+            .Finalize(node_def()));
+  }
   TF_ASSERT_OK(InitOp());
   // A matrix is:
   // |  1 |  2 |  3 |
@@ -316,9 +428,11 @@ TEST_F(QuantizedMatMulTest, Small_withBiasAndDeq) {
   AddInputFromArray<float>(TensorShape({}), {255.0f});
   AddInputFromArray<float>(TensorShape({}), {-127.0f});
   AddInputFromArray<float>(TensorShape({}), {127.0f});
-  AddInputFromArray<float>(TensorShape({}), {0});
-  AddInputFromArray<float>(TensorShape({}), {255.0f});
 
+  if (is_old_api) {
+    AddInputFromArray<float>(TensorShape({}), {0});
+    AddInputFromArray<float>(TensorShape({}), {255.0f});
+  }
   TF_ASSERT_OK(RunOpKernel());
   // Here are the results we expect, from hand calculations:
   // (1 * 7) + (2 * 11) + (3 * 15) = 74
@@ -353,19 +467,38 @@ TEST_F(QuantizedMatMulTest, Small_withBiasAndDeq) {
 
 // Two small matrices A of type uint8 and B of type int8  are multiplied
 // and the result is added with float bias and then performed relu on the result
-TEST_F(QuantizedMatMulTest, Small_withBiasAndRelu) {
-  TF_ASSERT_OK(NodeDefBuilder("quantized_mat_mul_op",
-                              "_MklQuantizedMatMulWithBiasAndRelu")
-                   .Input(FakeInput(DT_QUINT8))
-                   .Input(FakeInput(DT_QINT8))
-                   .Input(FakeInput(DT_FLOAT))
-                   .Input(FakeInput(DT_FLOAT))
-                   .Input(FakeInput(DT_FLOAT))
-                   .Input(FakeInput(DT_FLOAT))
-                   .Input(FakeInput(DT_FLOAT))
-                   .Attr("Toutput", DataTypeToEnum<qint32>::v())
-                   .Attr("_kernel", "QuantizedMklOp")
-                   .Finalize(node_def()));
+TEST_P(QuantizedMatMulTest, Small_withBiasAndRelu) {
+  const bool is_old_api = GetParam();
+  if (is_old_api) {
+    TF_ASSERT_OK(NodeDefBuilder("quantized_mat_mul_op",
+                                "_MklQuantizedMatMulWithBiasAndRelu")
+                     .Input(FakeInput(DT_QUINT8))
+                     .Input(FakeInput(DT_QINT8))
+                     .Input(FakeInput(DT_FLOAT))
+                     .Input(FakeInput(DT_FLOAT))
+                     .Input(FakeInput(DT_FLOAT))
+                     .Input(FakeInput(DT_FLOAT))
+                     .Input(FakeInput(DT_FLOAT))
+                     .Attr("Toutput", DataTypeToEnum<qint32>::v())
+                     .Attr("_kernel", "QuantizedMklOp")
+                     .Finalize(node_def()));
+  } else {
+    TF_ASSERT_OK(
+        NodeDefBuilder("quantized_mat_mul_op", "_QuantizedMatMul")
+            .Attr("Thost_inputs", {DT_QUINT8, DT_QINT8, DT_FLOAT, DT_FLOAT,
+                                   DT_FLOAT, DT_FLOAT, DT_FLOAT})
+            .Attr("Thost_outputs", {DT_QINT32, DT_FLOAT, DT_FLOAT})
+            .Attr("Tdevice_inputs", std::vector<DataType>())
+            .Attr("Tdevice_outputs", std::vector<DataType>())
+            .Attr("T1", DT_QUINT8)
+            .Attr("T2", DT_QINT8)
+            .Attr("Tbias", DT_FLOAT)
+            .Attr("Tout", DT_QINT32)
+            .Attr("fused_ops", {"BiasAdd", "Relu"})
+            .Input(FakeInput())
+            .Input(FakeInput())
+            .Finalize(node_def()));
+  }
   TF_ASSERT_OK(InitOp());
   // A matrix is:
   // |  1 |  2 |  3 |
@@ -408,21 +541,41 @@ TEST_F(QuantizedMatMulTest, Small_withBiasAndRelu) {
 
 // Simple test for Matrix multiplication with Bias, Relu and
 // Requantization fusion
-TEST_F(QuantizedMatMulTest, Small_withBiasAndReluAndReq) {
-  TF_ASSERT_OK(NodeDefBuilder("quantized_mat_mul_op",
-                              "_MklQuantizedMatMulWithBiasAndReluAndRequantize")
-                   .Input(FakeInput(DT_QUINT8))
-                   .Input(FakeInput(DT_QINT8))
-                   .Input(FakeInput(DT_QINT32))
-                   .Input(FakeInput(DT_FLOAT))
-                   .Input(FakeInput(DT_FLOAT))
-                   .Input(FakeInput(DT_FLOAT))
-                   .Input(FakeInput(DT_FLOAT))
-                   .Input(FakeInput(DT_FLOAT))
-                   .Input(FakeInput(DT_FLOAT))
-                   .Attr("Toutput", DataTypeToEnum<quint8>::v())
-                   .Attr("_kernel", "QuantizedMklOp")
-                   .Finalize(node_def()));
+TEST_P(QuantizedMatMulTest, Small_withBiasAndReluAndReq) {
+  const bool is_old_api = GetParam();
+  if (is_old_api) {
+    TF_ASSERT_OK(
+        NodeDefBuilder("quantized_mat_mul_op",
+                       "_MklQuantizedMatMulWithBiasAndReluAndRequantize")
+            .Input(FakeInput(DT_QUINT8))
+            .Input(FakeInput(DT_QINT8))
+            .Input(FakeInput(DT_QINT32))
+            .Input(FakeInput(DT_FLOAT))
+            .Input(FakeInput(DT_FLOAT))
+            .Input(FakeInput(DT_FLOAT))
+            .Input(FakeInput(DT_FLOAT))
+            .Input(FakeInput(DT_FLOAT))
+            .Input(FakeInput(DT_FLOAT))
+            .Attr("Toutput", DataTypeToEnum<quint8>::v())
+            .Attr("_kernel", "QuantizedMklOp")
+            .Finalize(node_def()));
+  } else {
+    TF_ASSERT_OK(NodeDefBuilder("quantized_mat_mul_op", "_QuantizedMatMul")
+                     .Attr("Thost_inputs",
+                           {DT_QUINT8, DT_QINT8, DT_QINT32, DT_FLOAT, DT_FLOAT,
+                            DT_FLOAT, DT_FLOAT, DT_FLOAT, DT_FLOAT})
+                     .Attr("Thost_outputs", {DT_QUINT8, DT_FLOAT, DT_FLOAT})
+                     .Attr("Tdevice_inputs", std::vector<DataType>())
+                     .Attr("Tdevice_outputs", std::vector<DataType>())
+                     .Attr("T1", DT_QUINT8)
+                     .Attr("T2", DT_QINT8)
+                     .Attr("Tbias", DT_QINT32)
+                     .Attr("Tout", DT_QUINT8)
+                     .Attr("fused_ops", {"BiasAdd", "Relu", "Requantize"})
+                     .Input(FakeInput())
+                     .Input(FakeInput())
+                     .Finalize(node_def()));
+  }
   TF_ASSERT_OK(InitOp());
   // A matrix is:
   // |  1 |  2 |  3 |
@@ -470,11 +623,25 @@ TEST_F(QuantizedMatMulTest, Small_withBiasAndReluAndReq) {
   // 178 * 1.00392 ~= 178.698 ~= 179
 
   Tensor expected(allocator(), DT_QUINT8, TensorShape({2, 4}));
+  if (is_old_api) {
 #ifdef ENABLE_ONEDNN_V3
-  test::FillValues<quint8>(&expected, {84, 60, 116, 52, 183, 168, 233, 178});
+    test::FillValues<quint8>(&expected, {84, 60, 116, 52, 183, 168, 233, 178});
 #else
-  test::FillValues<quint8>(&expected, {84, 60, 116, 52, 184, 169, 234, 179});
+    test::FillValues<quint8>(&expected, {84, 60, 116, 52, 184, 169, 234, 179});
 #endif  // ENABLE_ONEDNN_V3
+  } else {
+    // New api uses more numerical precision preserving equation. Old api scales
+    // up to 32-bit and then scales down from 32-bit to 8-bit. New api instead
+    // does a dequantization followed by a scaling to 8-bit.
+    // In this test,
+    //    input deq. scale =  ((255.0 * 127.0) / (255.0 * 127.0)) = 1.0
+    //    output req. scale = 255.0 / 255.0 = 1.0
+    //    combined scale = 1.0 * 1.0 = 1.0
+    // Note: new api scale value is 1.0, whereas the old api scale is 1.000392.
+    // Correct value is 1.0f. The closer it is to the correct value the better
+    // the formula is.
+    test::FillValues<quint8>(&expected, {84, 60, 116, 52, 183, 168, 233, 178});
+  }
 
   const Tensor& output = *GetOutput(0);
   test::ExpectTensorEqual<quint8>(expected, output);
@@ -484,19 +651,38 @@ TEST_F(QuantizedMatMulTest, Small_withBiasAndReluAndReq) {
 // and the result is added with int32 bias
 // For the first time B matrix will be reordered and cached which will be
 // used for subsequent runs
-TEST_F(QuantizedMatMulTest, Small_withWeightCached) {
-  TF_ASSERT_OK(
-      NodeDefBuilder("quantized_mat_mul_op", "_MklQuantizedMatMulWithBias")
-          .Input(FakeInput(DT_QUINT8))
-          .Input(FakeInput(DT_QINT8))
-          .Input(FakeInput(DT_QINT32))
-          .Input(FakeInput(DT_FLOAT))
-          .Input(FakeInput(DT_FLOAT))
-          .Input(FakeInput(DT_FLOAT))
-          .Input(FakeInput(DT_FLOAT))
-          .Attr("Toutput", DataTypeToEnum<qint32>::v())
-          .Attr("_kernel", "QuantizedMklOp")
-          .Finalize(node_def()));
+TEST_P(QuantizedMatMulTest, Small_withWeightCached) {
+  const bool is_old_api = GetParam();
+  if (is_old_api) {
+    TF_ASSERT_OK(
+        NodeDefBuilder("quantized_mat_mul_op", "_MklQuantizedMatMulWithBias")
+            .Input(FakeInput(DT_QUINT8))
+            .Input(FakeInput(DT_QINT8))
+            .Input(FakeInput(DT_QINT32))
+            .Input(FakeInput(DT_FLOAT))
+            .Input(FakeInput(DT_FLOAT))
+            .Input(FakeInput(DT_FLOAT))
+            .Input(FakeInput(DT_FLOAT))
+            .Attr("Toutput", DataTypeToEnum<qint32>::v())
+            .Attr("_kernel", "QuantizedMklOp")
+            .Finalize(node_def()));
+  } else {
+    TF_ASSERT_OK(
+        NodeDefBuilder("quantized_mat_mul_op", "_QuantizedMatMul")
+            .Attr("Thost_inputs", {DT_QUINT8, DT_QINT8, DT_QINT32, DT_FLOAT,
+                                   DT_FLOAT, DT_FLOAT, DT_FLOAT})
+            .Attr("Thost_outputs", {DT_QINT32, DT_FLOAT, DT_FLOAT})
+            .Attr("Tdevice_inputs", std::vector<DataType>())
+            .Attr("Tdevice_outputs", std::vector<DataType>())
+            .Attr("T1", DT_QUINT8)
+            .Attr("T2", DT_QINT8)
+            .Attr("Tbias", DT_QINT32)
+            .Attr("Tout", DT_QINT32)
+            .Attr("fused_ops", {"BiasAdd"})
+            .Input(FakeInput())
+            .Input(FakeInput())
+            .Finalize(node_def()));
+  }
   TF_ASSERT_OK(InitOp());
   // The tensor shape of (1,3) is selected to allow the oneDNN expected
   // weight format to be made as OI rather than IO for BS > 1
@@ -549,6 +735,9 @@ TEST_F(QuantizedMatMulTest, Small_withWeightCached) {
   test::ExpectTensorEqual<qint32>(expected, output_new);
 }
 
+INSTANTIATE_TEST_SUITE_P(All, QuantizedMatMulTest,
+                         ::testing::Values(true, false));
+
 }  // namespace tensorflow
 
 #endif  // INTEL_MKL
diff --git a/tensorflow/core/kernels/mkl/onednn_fused_matmul_ops_test.cc b/tensorflow/core/kernels/mkl/onednn_fused_matmul_ops_test.cc
new file mode 100644
index 00000000000000..af1d056ca7530a
--- /dev/null
+++ b/tensorflow/core/kernels/mkl/onednn_fused_matmul_ops_test.cc
@@ -0,0 +1,749 @@
+/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+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.
+==============================================================================*/
+#if defined(INTEL_MKL)
+
+#define EIGEN_USE_THREADS
+
+#include <algorithm>
+#include <limits>
+
+#include "absl/algorithm/container.h"
+#include "absl/strings/match.h"
+#include "gtest/gtest.h"
+#include "tensorflow/cc/ops/array_ops.h"
+#include "tensorflow/cc/ops/const_op.h"
+#include "tensorflow/cc/ops/math_ops.h"
+#include "tensorflow/cc/ops/nn_ops.h"
+#include "tensorflow/cc/ops/nn_ops_internal.h"
+#include "tensorflow/cc/ops/standard_ops.h"
+#include "tensorflow/core/framework/fake_input.h"
+#include "tensorflow/core/framework/node_def_builder.h"
+#include "tensorflow/core/framework/tensor.h"
+#include "tensorflow/core/framework/types.pb.h"
+#include "tensorflow/core/graph/node_builder.h"
+#include "tensorflow/core/kernels/mkl/mkl_kernel_util.h"
+#include "tensorflow/core/kernels/ops_testutil.h"
+#include "tensorflow/core/kernels/ops_util.h"
+#include "tensorflow/core/kernels/quantization_utils.h"
+#include "tensorflow/core/lib/core/status_test_util.h"
+#include "tensorflow/core/platform/env.h"
+#include "tensorflow/core/platform/test.h"
+#include "tensorflow/core/protobuf/rewriter_config.pb.h"
+#include "tensorflow/core/public/session.h"
+#include "tensorflow/core/util/util.h"
+#include "unsupported/Eigen/CXX11/Tensor"
+
+namespace tensorflow {
+
+// The test suite contains different categories of tests.
+//    (1) Realnumber (float/bfloat16): The output of _FusedMatMul should be
+//    close enough to the final output of the sequence of unfused operations.
+//    Only Gelu fusion is included here. All other fusion tests can be found in
+//    tensorflow/core/kernels/mkl/mkl_fused_ops_test.cc
+//
+//    (2) Quantized: Possible fusions are done in _QuantizedMatMul op. The
+//    output of
+//        quantize --> quantized_op --> dequantize, or
+//        quantize --> quantized_op --> requantize --> dequantize
+//    should be close (with a higher tolerance) to the final output of the
+//    sequence of unfused real number type operations. For the quantized
+//    scenario, it is assumed that the first matrix of MatMul op represents
+//    feature, while the second matrix represents weight parameters. The feature
+//    matrix can be quantized with MIN_FIRST (to QUINT8) or SCALED (to QINT8)
+//    mode and always quantized per-tensor. The weight can be quantized with
+//    per-tensor or per-channel, only with SCALED mode to QINT8.
+
+// T: float or bfloat16 used as tensor type of the MatMul and fusion operation.
+template <typename T>
+class FusedMatMulOpsTest : public OpsTestBase {
+ private:
+  float leakyrelu_alpha_ = 0.2f;
+
+ protected:
+  struct FusedOpsAndDims {
+    // List of fusions.
+    std::vector<string> fused_ops;
+    // Tensor dimension associated with the fusions. It is assumed here that
+    // each fusion requires no more than one addtional tensor. If some fusion
+    // does not require a tensor, e.g., Relu, the tensor dimensions will be {0}
+    // implying an an empty tensor.
+    std::vector<std::vector<int64_t>> fusion_dims;
+  };
+
+  struct FusedOpsAndTensors {
+    // List of fusions.
+    std::vector<string> fused_ops;
+    // Tensors associated with the fusions. It is assumed here that each fusion
+    // requires no more than one additional tensor. If some fusion does not
+    // require a tensor, e.g., Relu, the tensor will be an empty tensor.
+    std::vector<Tensor> fusion_tensors;
+  };
+
+  using GraphRunner =
+      std::function<void(const Tensor& x, const Tensor& y,
+                         const FusedOpsAndTensors& fused_ops_and_tensors,
+                         Tensor* result, bool transpose_x, bool transpose_y)>;
+
+  using QuantizedGraphRunner = std::function<void(
+      const Tensor& x, const Tensor& y,
+      const FusedOpsAndTensors& fused_ops_and_tensors, Tensor* result,
+      bool transpose_x, bool transpose_y, string input_quant_mode,
+      string output_quant_mode, bool is_bias_quantized, bool is_perchannel,
+      bool requantize, float output_min, float output_max)>;
+
+  bool HasQuantizationSupport() {
+    return TestCPUFeature(tensorflow::port::CPUFeature::AVX_VNNI_INT8) ||
+           TestCPUFeature(tensorflow::port::CPUFeature::AVX512_VNNI) ||
+           TestCPUFeature(port::CPUFeature::AMX_INT8);
+  }
+
+  // Runs a Tensorflow graph defined by the root scope, and fetches the result
+  // of 'fetch' node into the outputs. Optional `add_nodes` parameter
+  // allows to define nodes directly using NodeDefBuilder.
+  void RunAndFetch(const tensorflow::Scope& root,
+                   const std::vector<string>& fetch,
+                   std::vector<Tensor>* outputs,
+                   const std::vector<const NodeDef*> add_nodes = {}) {
+    tensorflow::GraphDef graph;
+    TF_ASSERT_OK(root.ToGraphDef(&graph));
+
+    for (const NodeDef* add_node : add_nodes) {
+      *graph.add_node() = *add_node;
+    }
+
+    // We really want to make sure that graph executed exactly as we passed it
+    // to the session, so we disable various optimizations.
+    tensorflow::SessionOptions session_options;
+
+    // Disable common runtime constant folding.
+    session_options.config.mutable_graph_options()
+        ->mutable_optimizer_options()
+        ->set_opt_level(OptimizerOptions::L0);
+
+    // Disable Grappler optimizations for tests.
+    tensorflow::RewriterConfig* cfg =
+        session_options.config.mutable_graph_options()
+            ->mutable_rewrite_options();
+    cfg->set_constant_folding(tensorflow::RewriterConfig::OFF);
+    cfg->set_layout_optimizer(tensorflow::RewriterConfig::OFF);
+    cfg->set_remapping(tensorflow::RewriterConfig::OFF);
+
+    std::unique_ptr<tensorflow::Session> session(
+        tensorflow::NewSession(session_options));
+
+    const string device = "/device:CPU:0";
+    for (NodeDef& mutable_node : *graph.mutable_node()) {
+      mutable_node.set_device(device);
+    }
+
+    TF_ASSERT_OK(session->Create(graph));
+    TF_ASSERT_OK(session->Run({}, fetch, {}, outputs));
+  }
+
+  Output ActivationOp(Scope& root, string op, Output x, string name) {
+    // TODO(intel-tf): Add GeluExact (Erf op based) when the Erf op is enabled
+    // for bfloat16. GeluExact with float32 precision test can be found in
+    //    tensorflow/python/grappler/remapper_test.py
+    if (op == "Relu") {
+      return ops::Relu(root.WithOpName(name), x);
+    } else if (op == "Relu6") {
+      return ops::Relu6(root.WithOpName(name), x);
+    } else if (op == "LeakyRelu") {
+      return ops::internal::LeakyRelu(
+          root.WithOpName(name), x,
+          ops::internal::LeakyRelu::Attrs().Alpha(this->leakyrelu_alpha_));
+    } else if (op == "Elu") {
+      return ops::Elu(root.WithOpName(name), x);
+    } else if (op == "Tanh") {
+      return ops::Tanh(root.WithOpName(name), x);
+    } else if (op == "Sigmoid") {
+      return ops::Sigmoid(root.WithOpName(name), x);
+    } else if (op == "GeluApproximate") {
+      Output three = ops::Const<T>(root.WithOpName("gelu_three"), 3.0f);
+      Output empirical =
+          ops::Const<T>(root.WithOpName("gelu_empirical"), 0.044715f);
+      Output square_root_two_over_pi = ops::Const<T>(
+          root.WithOpName("gelu_square_root_two_over_pi"), 0.7978845608028654f);
+      Output one = ops::Const<T>(root.WithOpName("gelu_one"), 1.0f);
+      Output half = ops::Const<T>(root.WithOpName("gelu_half"), 0.5f);
+      Output pow = ops::Pow(root.WithOpName("gelu_pow"), x, three);
+      Output mul1 = ops::Multiply(root.WithOpName("gelu_mul1"), empirical, pow);
+      Output add1 = ops::AddV2(root.WithOpName("gelu_add1"), x, mul1);
+      Output mul2 = ops::Multiply(root.WithOpName("gelu_mul2"),
+                                  square_root_two_over_pi, add1);
+      Output tanh = ops::Tanh(root.WithOpName("gelu_tanh"), mul2);
+      Output add3 = ops::AddV2(root.WithOpName("gelu_add3"), one, tanh);
+      Output mul3 = ops::Multiply(root.WithOpName("gelu_mul3"), half, x);
+      return ops::Multiply(root.WithOpName(name), mul3, add3);
+    } else {
+      EXPECT_TRUE(false) << absl::StrCat("The activation: ", op,
+                                         " is not supported in this test.");
+    }
+  }
+
+  void RunMatMulAndFusedOps(const Tensor& x, const Tensor& y,
+                            const FusedOpsAndTensors& fused_ops_and_tensors,
+                            Tensor* result, bool transpose_x,
+                            bool transpose_y) {
+    Scope root = tensorflow::Scope::NewRootScope();
+
+    Output x_input =
+        ops::Const(root.WithOpName("x_input"), Input::Initializer(x));
+    Output y_input =
+        ops::Const(root.WithOpName("y_input"), Input::Initializer(y));
+    Output last_output = ops::MatMul(
+        root.WithOpName("matmul"), x_input, y_input,
+        ops::MatMul::Attrs().TransposeA(transpose_x).TransposeB(transpose_y));
+    auto& fused_ops = fused_ops_and_tensors.fused_ops;
+    auto& fusion_tensors = fused_ops_and_tensors.fusion_tensors;
+    for (int i = 0; i < fused_ops.size(); ++i) {
+      const string& op = fused_ops[i];
+      if (op == "BiasAdd") {
+        Output arg = ops::Const(root.WithOpName(absl::StrCat("arg", i)),
+                                Input::Initializer(fusion_tensors[i]));
+        last_output = ops::BiasAdd(
+            root.WithOpName(absl::StrCat("bias_add_at_", i)), last_output, arg);
+      } else if (op == "Relu" || op == "Relu6" || op == "LeakyRelu" ||
+                 op == "Elu" || op == "Tanh" || op == "Sigmoid" ||
+                 op == "GeluApproximate") {
+        last_output =
+            ActivationOp(root, op, last_output, absl::StrCat(op, "_at_", i));
+      } else if (op == "Add") {
+        ASSERT_EQ(x.dtype(), fusion_tensors[i].dtype());
+        Output arg = ops::Const(root.WithOpName(absl::StrCat("arg", i)),
+                                Input::Initializer(fusion_tensors[i]));
+        last_output = ops::AddV2(root.WithOpName(absl::StrCat("add_at_", i)),
+                                 last_output, arg);
+      } else {
+        EXPECT_TRUE(false) << absl::StrCat("The fusion: [",
+                                           absl::StrJoin(fused_ops, ","),
+                                           "] is not supported in this test.");
+      }
+    }
+    std::vector<Tensor> outputs;
+    RunAndFetch(root, {last_output.name()}, &outputs);
+    *result = outputs[0];
+  }
+
+  void RunFusedMatMul(const Tensor& x, const Tensor& y,
+                      const FusedOpsAndTensors& fused_ops_and_tensors,
+                      Tensor* result, bool transpose_x, bool transpose_y) {
+    Scope root = tensorflow::Scope::NewRootScope();
+
+    DataType dtype = DataTypeToEnum<T>::v();
+
+    Output x_input =
+        ops::Const(root.WithOpName("x_input"), Input::Initializer(x));
+    Output y_input =
+        ops::Const(root.WithOpName("y_input"), Input::Initializer(y));
+    auto& fused_ops = fused_ops_and_tensors.fused_ops;
+    auto& fusion_tensors = fused_ops_and_tensors.fusion_tensors;
+    int num_fusion_inputs = 0;
+    bool has_leaky_relu = false;
+    std::vector<NodeDefBuilder::NodeOut> fusion_inputs;
+    for (int i = 0; i < fused_ops.size(); ++i) {
+      const string& op = fused_ops[i];
+      if (op == "BiasAdd") {
+        Output arg = ops::Const(root.WithOpName(absl::StrCat("arg", i)),
+                                Input::Initializer(fusion_tensors[i]));
+        fusion_inputs.push_back({arg.name(), 0, dtype});
+        num_fusion_inputs++;
+      } else if (op == "Add") {
+        ASSERT_EQ(x.dtype(), fusion_tensors[i].dtype());
+        Output arg = ops::Const(root.WithOpName(absl::StrCat("arg", i)),
+                                Input::Initializer(fusion_tensors[i]));
+        fusion_inputs.push_back({arg.name(), 0, dtype});
+        num_fusion_inputs++;
+      } else if (op == "LeakyRelu") {
+        has_leaky_relu = true;
+      } else {
+        bool is_supported = op == "Relu" || op == "Relu6" ||
+                            op == "LeakyRelu" || op == "Elu" || op == "Tanh" ||
+                            op == "Sigmoid" || op == "GeluApproximate";
+        EXPECT_TRUE(is_supported)
+            << absl::StrCat("The fusion: [", absl::StrJoin(fused_ops, ","),
+                            "] is not supported in this test.");
+      }
+    }
+    NodeDef fused_matmul;
+    std::vector<const NodeDef*> add_nodes;
+    TF_EXPECT_OK(NodeDefBuilder("fused_batch_matmul", "_MklNativeFusedMatMul")
+                     .Input({x_input.name(), 0, dtype})
+                     .Input({y_input.name(), 0, dtype})
+                     .Input(fusion_inputs)
+                     .Attr("transpose_a", transpose_x)
+                     .Attr("transpose_b", transpose_y)
+                     .Attr("num_args", num_fusion_inputs)
+                     .Attr("fused_ops", fused_ops)
+                     .Attr("leakyrelu_alpha",
+                           has_leaky_relu ? this->leakyrelu_alpha_ : 0.2f)
+                     .Attr("_kernel", "MklNameChangeOp")
+                     .Finalize(&fused_matmul));
+    add_nodes = {&fused_matmul};
+    std::vector<Tensor> outputs;
+    RunAndFetch(root, {fused_matmul.name()}, &outputs, add_nodes);
+    *result = outputs[0];
+  }
+
+  // Compute quantized tensor perchannel (aka axis) in SCALED mode for 2D
+  // tensor.
+  template <bool transpose = false>
+  void GetPerchannelQuantizationTensors(const Tensor& input, Tensor* output,
+                                        Tensor* min_tensor,
+                                        Tensor* max_tensor) {
+    ASSERT_EQ(input.dims(), 2);
+    ASSERT_EQ(output->dtype(), DT_QINT8);
+    constexpr int axis = transpose ? 0 : 1;
+    int num_channels = input.dim_size(axis);
+    ASSERT_EQ(min_tensor->NumElements(), num_channels);
+    ASSERT_EQ(max_tensor->NumElements(), num_channels);
+
+    auto eigen_input_tensor = input.matrix<T>().template cast<float>();
+    auto eigen_output_tensor = output->matrix<qint8>();
+    std::vector<float> scales(num_channels);
+    float* min_tensor_buf = min_tensor->flat<float>().data();
+    float* max_tensor_buf = max_tensor->flat<float>().data();
+    for (int i = 0; i < num_channels; ++i) {
+      auto input_slice = eigen_input_tensor.template chip<axis>(i);
+      auto output_slice = eigen_output_tensor.template chip<axis>(i);
+      Eigen::Tensor<float, 0, Eigen::RowMajor> min = input_slice.minimum();
+      Eigen::Tensor<float, 0, Eigen::RowMajor> max = input_slice.maximum();
+      float min_i = min();
+      float max_i = max();
+      float range = std::max(std::abs(min_i), std::abs(max_i));
+      min_tensor_buf[i] = -range;
+      max_tensor_buf[i] = range;
+      const float scale = 127.0f / range;
+      output_slice = (input_slice * scale).round().template cast<qint8>();
+    }
+  }
+
+  void RunQuantizedMatMul(const Tensor& x, const Tensor& y,
+                          const FusedOpsAndTensors& fused_ops_and_tensors,
+                          Tensor* result, bool transpose_x, bool transpose_y,
+                          string input_quant_mode, string output_quant_mode,
+                          bool is_bias_quantized, bool is_perchannel,
+                          bool requantize, float output_min, float output_max) {
+    // TODO(intel-tf): Extend test with quantized bias
+    ASSERT_EQ(is_bias_quantized, false);
+
+    DataType real_dtype = DataTypeToEnum<T>::v();
+    DataType qinput_dtype =
+        (input_quant_mode == "MIN_FIRST") ? DT_QUINT8 : DT_QINT8;
+    // Quantize x and y
+    Tensor x_qtensor(qinput_dtype, x.shape());
+    Tensor x_min_tensor(DT_FLOAT, TensorShape({}));
+    Tensor x_max_tensor(DT_FLOAT, TensorShape({}));
+    auto status = MklTestingUtil::GetQuantizationTensors<T>(
+        x, &x_qtensor, qinput_dtype, input_quant_mode, &x_min_tensor,
+        &x_max_tensor);
+    ASSERT_TRUE(status.ok());
+    Tensor y_qtensor(DT_QINT8, y.shape());
+    const int num_channels = transpose_y ? y.dim_size(0) : y.dim_size(1);
+    TensorShape minmax_shape =
+        is_perchannel ? TensorShape({num_channels}) : TensorShape({});
+    Tensor y_min_tensor(DT_FLOAT, minmax_shape);
+    Tensor y_max_tensor(DT_FLOAT, minmax_shape);
+    if (is_perchannel) {
+      if (transpose_y) {
+        GetPerchannelQuantizationTensors<true>(y, &y_qtensor, &y_min_tensor,
+                                               &y_max_tensor);
+      } else {
+        GetPerchannelQuantizationTensors<false>(y, &y_qtensor, &y_min_tensor,
+                                                &y_max_tensor);
+      }
+    } else {
+      auto status = MklTestingUtil::GetQuantizationTensors<T>(
+          y, &y_qtensor, DT_QINT8, "SCALED", &y_min_tensor, &y_max_tensor);
+      ASSERT_TRUE(status.ok());
+    }
+
+    Scope root = tensorflow::Scope::NewRootScope();
+
+    Output x_input =
+        ops::Const(root.WithOpName("x_input"), Input::Initializer(x_qtensor));
+    Output x_min =
+        ops::Const(root.WithOpName("x_min"), Input::Initializer(x_min_tensor));
+    Output x_max =
+        ops::Const(root.WithOpName("x_max"), Input::Initializer(x_max_tensor));
+    Output y_input =
+        ops::Const(root.WithOpName("y_input"), Input::Initializer(y_qtensor));
+    Output y_min =
+        ops::Const(root.WithOpName("y_min"), Input::Initializer(y_min_tensor));
+    Output y_max =
+        ops::Const(root.WithOpName("y_max"), Input::Initializer(y_max_tensor));
+    auto& fused_ops = fused_ops_and_tensors.fused_ops;
+    auto& fusion_tensors = fused_ops_and_tensors.fusion_tensors;
+    int num_fusion_inputs = 0;
+    std::vector<NodeDefBuilder::NodeOut> fusion_inputs;
+    bool has_leaky_relu = false;
+    for (int i = 0; i < fused_ops.size(); ++i) {
+      const string& op = fused_ops[i];
+      if (op == "BiasAdd") {
+        Output arg = ops::Const(root.WithOpName(absl::StrCat("arg", i)),
+                                Input::Initializer(fusion_tensors[i]));
+        fusion_inputs.push_back({arg.name(), 0, real_dtype});
+        num_fusion_inputs++;
+      } else if (op == "Add") {
+        ASSERT_EQ(real_dtype, fusion_tensors[i].dtype());
+        Output arg = ops::Const(root.WithOpName(absl::StrCat("arg", i)),
+                                Input::Initializer(fusion_tensors[i]));
+        fusion_inputs.push_back({arg.name(), 0, real_dtype});
+        num_fusion_inputs++;
+      } else if (op == "LeakyRelu") {
+        has_leaky_relu = true;
+      }
+    }
+    NodeDef fused_matmul;
+    std::vector<const NodeDef*> add_nodes;
+    std::vector<Tensor> outputs;
+    std::vector<NodeDefBuilder::NodeOut> inputs;
+    inputs.push_back({"x_input", 0, qinput_dtype});
+    inputs.push_back({"y_input", 0, DT_QINT8});
+    inputs.insert(std::end(inputs), std::begin(fusion_inputs),
+                  std::end(fusion_inputs));
+    inputs.push_back({"x_min", 0, DT_FLOAT});
+    inputs.push_back({"x_max", 0, DT_FLOAT});
+    inputs.push_back({"y_min", 0, DT_FLOAT});
+    inputs.push_back({"y_max", 0, DT_FLOAT});
+    std::vector<string> extended_fused_ops(fused_ops);
+    DataType out_dtype;
+    if (requantize) {
+      if (output_quant_mode == "SCALED") {
+        out_dtype = DT_QINT8;
+      } else {
+        out_dtype = DT_QUINT8;
+      }
+    } else {
+      out_dtype = real_dtype;
+    }
+    std::vector<DataType> output_dtypes;
+    if (requantize) {
+      Output out_min = ops::Const(root.WithOpName("output_min"), output_min);
+      Output out_max = ops::Const(root.WithOpName("output_max"), output_max);
+      inputs.push_back({"output_min", 0, DT_FLOAT});
+      inputs.push_back({"output_max", 0, DT_FLOAT});
+      extended_fused_ops.push_back("Requantize");
+      output_dtypes = {out_dtype, DT_FLOAT, DT_FLOAT};
+    } else {
+      extended_fused_ops.push_back("Dequantize");
+      output_dtypes = {out_dtype};
+    }
+
+    TF_EXPECT_OK(NodeDefBuilder("quantized_fused_matmul", "_QuantizedMatMul")
+                     .Attr("Tdevice_inputs", std::vector<DataType>())
+                     .Input(FakeInput())
+                     .Input(inputs)
+                     .Attr("Thost_outputs", output_dtypes)
+                     .Attr("Tdevice_outputs", std::vector<DataType>())
+                     .Attr("T1", qinput_dtype)
+                     .Attr("T2", DT_QINT8)
+                     .Attr("Tbias", real_dtype)
+                     .Attr("Tout", out_dtype)
+                     .Attr("U", real_dtype)
+                     .Attr("transpose_a", transpose_x)
+                     .Attr("transpose_b", transpose_y)
+                     .Attr("fused_ops", extended_fused_ops)
+                     .Attr("leakyrelu_alpha",
+                           has_leaky_relu ? this->leakyrelu_alpha_ : 0.2f)
+                     .Attr("input_quant_mode", input_quant_mode)
+                     .Attr("output_quant_mode", output_quant_mode)
+                     .Finalize(&fused_matmul));
+    if (requantize) {
+      NodeDef dequantize;
+      TF_EXPECT_OK(NodeDefBuilder("dequantize", "Dequantize")
+                       .Input({"quantized_fused_matmul", 0, out_dtype})
+                       .Input({"quantized_fused_matmul", 1, DT_FLOAT})
+                       .Input({"quantized_fused_matmul", 2, DT_FLOAT})
+                       .Attr("dtype", real_dtype)
+                       .Attr("mode", output_quant_mode)
+                       .Finalize(&dequantize));
+      add_nodes = {&fused_matmul, &dequantize};
+      RunAndFetch(root, {dequantize.name()}, &outputs, add_nodes);
+    } else {
+      add_nodes = {&fused_matmul};
+      RunAndFetch(root, {fused_matmul.name()}, &outputs, add_nodes);
+    }
+    *result = outputs[0];
+  }
+
+  template <typename FusedGraphRunner>
+  void VerifyTensorsNear(const std::vector<int64_t>& x_dims,
+                         const std::vector<int64_t>& y_dims,
+                         const FusedOpsAndDims& fused_ops_and_dims,
+                         const GraphRunner& run_default,
+                         const FusedGraphRunner& run_fused, bool transpose_x,
+                         bool transpose_y, const double atol = 1e-5,
+                         // The following arguments are used by quantized fusion
+                         string input_quant_mode = "SCALED",
+                         string output_quant_mode = "SCALED",
+                         bool is_bias_quantized = false,
+                         bool is_perchannel = false, bool requantize = false) {
+    srand(1234);
+    DataType dtype = DataTypeToEnum<T>::v();
+    TensorShape x_shape = TensorShape(x_dims);
+    TensorShape y_shape = TensorShape(y_dims);
+
+    Tensor x_tensor(dtype, x_shape);
+    x_tensor.flat<T>().setRandom();
+    x_tensor.flat<T>() -= x_tensor.flat<T>().constant(static_cast<T>(0.5));
+
+    Tensor y_tensor(dtype, y_shape);
+    y_tensor.flat<T>().setRandom();
+    y_tensor.flat<T>() -= y_tensor.flat<T>().constant(static_cast<T>(0.5));
+
+    FusedOpsAndTensors fused_ops_and_tensors;
+    fused_ops_and_tensors.fused_ops = fused_ops_and_dims.fused_ops;
+    const auto& fused_ops = fused_ops_and_tensors.fused_ops;   // Alias to field
+    const auto& fusion_dims = fused_ops_and_dims.fusion_dims;  // Alias to field
+    auto& fusion_tensors = fused_ops_and_tensors.fusion_tensors;
+    for (int i = 0; i < fused_ops.size(); ++i) {
+      TensorShape arg_shape = TensorShape(fusion_dims[i]);
+      Tensor arg_tensor(dtype, arg_shape);
+      arg_tensor.flat<T>().setRandom();
+      arg_tensor.flat<T>() -=
+          arg_tensor.flat<T>().constant(static_cast<T>(0.5));
+      fusion_tensors.push_back(arg_tensor);
+    }
+    Tensor default_result;
+    run_default(x_tensor, y_tensor, fused_ops_and_tensors, &default_result,
+                transpose_x, transpose_y);
+
+    Tensor fused_result;
+    if constexpr (std::is_same<FusedGraphRunner, QuantizedGraphRunner>::value) {
+      float output_min = 1.0;
+      float output_max = 1.0 + std::numeric_limits<float>::epsilon();
+      if (requantize) {
+        T min;
+        T max;
+        MklTestingUtil::ComputeMinMax<T>(default_result, &min, &max);
+        output_min = static_cast<float>(min);
+        output_max = static_cast<float>(max);
+      }
+      // Run quantized fusion
+      run_fused(x_tensor, y_tensor, fused_ops_and_tensors, &fused_result,
+                transpose_x, transpose_y, input_quant_mode, output_quant_mode,
+                is_bias_quantized, is_perchannel, requantize, output_min,
+                output_max);
+    } else {
+      // Run realnumber type fusion
+      run_fused(x_tensor, y_tensor, fused_ops_and_tensors, &fused_result,
+                transpose_x, transpose_y);
+    }
+    std::vector<std::pair<Tensor, Tensor>> tensor_pairs = {
+        {default_result, fused_result}};
+    for (auto& pair : tensor_pairs) {
+      const Tensor& expected = pair.first;
+      const Tensor& evaluated = pair.second;
+
+      ASSERT_EQ(expected.dtype(), evaluated.dtype());
+      ASSERT_EQ(expected.shape(), evaluated.shape());
+
+      test::ExpectClose(expected, evaluated, atol);
+    }
+  }
+
+  void GetFusionConfiguration(const std::vector<string>& fused_ops,
+                              const int row, const int col,
+                              FusedOpsAndDims* fused_ops_and_dims) {
+    if (fused_ops == std::vector<string>{"BiasAdd"}) {
+      *fused_ops_and_dims = {fused_ops, {std::vector<int64_t>{col}}};
+    } else if (fused_ops == std::vector<string>{"BiasAdd", "Relu"} ||
+               fused_ops == std::vector<string>{"BiasAdd", "Relu6"} ||
+               fused_ops == std::vector<string>{"BiasAdd", "LeakyRelu"} ||
+               fused_ops == std::vector<string>{"BiasAdd", "Elu"} ||
+               fused_ops == std::vector<string>{"BiasAdd", "Tanh"} ||
+               fused_ops == std::vector<string>{"BiasAdd", "Sigmoid"} ||
+               fused_ops == std::vector<string>{"BiasAdd", "GeluApproximate"}) {
+      *fused_ops_and_dims = {
+          fused_ops, {std::vector<int64_t>{col}, std::vector<int64_t>{0}}};
+    } else if (fused_ops == std::vector<string>{"BiasAdd", "Add"}) {
+      *fused_ops_and_dims = {
+          fused_ops,
+          {std::vector<int64_t>{col}, std::vector<int64_t>{row, col}}};
+    } else {
+      EXPECT_TRUE(false) << absl::StrCat("The fusion: [",
+                                         absl::StrJoin(fused_ops, ","),
+                                         "] is not supported in this test.");
+    }
+  }
+
+  void VerifyFusedMatMul(std::vector<string> fused_ops) {
+    const GraphRunner run_default =
+        [&](const Tensor& x, const Tensor& y,
+            const FusedOpsAndTensors& fused_ops_and_tensors, Tensor* result,
+            bool transpose_x, bool transpose_y) {
+          this->RunMatMulAndFusedOps(x, y, fused_ops_and_tensors, result,
+                                     transpose_x, transpose_y);
+        };
+
+    const GraphRunner run_fused =
+        [&](const Tensor& x, const Tensor& y,
+            const FusedOpsAndTensors& fused_ops_and_tensors, Tensor* result,
+            bool transpose_x, bool transpose_y) {
+          this->RunFusedMatMul(x, y, fused_ops_and_tensors, result, transpose_x,
+                               transpose_y);
+        };
+    const double atol = std::is_same<T, bfloat16>::value ? 1e-2 : 1e-5;
+    constexpr int M = 3;
+    constexpr int K = 4;
+    constexpr int N = 5;
+    bool transpose_x = false;  // OpKernel does not support transpose_x.
+    std::vector<int64_t> x_dims;
+    std::vector<int64_t> y_dims;
+    FusedOpsAndDims fused_ops_and_dims;
+    GetFusionConfiguration(fused_ops, M, N, &fused_ops_and_dims);
+    for (bool transpose_y : {false, true}) {
+      x_dims =
+          transpose_x ? std::vector<int64_t>{K, M} : std::vector<int64_t>{M, K};
+      y_dims =
+          transpose_y ? std::vector<int64_t>{N, K} : std::vector<int64_t>{K, N};
+      VerifyTensorsNear<GraphRunner>(x_dims, y_dims, fused_ops_and_dims,
+                                     run_default, run_fused, transpose_x,
+                                     transpose_y, atol);
+    }
+  }
+
+  // The following test runs with 32 configurations.
+  //    (1) input quantization mode : {"MIN_FIRST", "SCALED"}
+  //    (2) input quantization mode : {"MIN_FIRST", "SCALED"}
+  //    (3) weight quantization per_channel : {false, true}
+  //    (4) output is requantized or dequantized:
+  //        false: dequantized
+  //        true: requantized
+  //    (5) weight matrix is transposed : {false, true}
+  void VerifyQuantizedMatMul(std::vector<string> fused_ops) {
+    if (!HasQuantizationSupport()) {
+      GTEST_SKIP() << "oneDNN based Quantized ops are not enabled on this CPU.";
+    }
+    const GraphRunner run_default =
+        [&](const Tensor& x, const Tensor& y,
+            const FusedOpsAndTensors& fused_ops_and_tensors, Tensor* result,
+            bool transpose_x, bool transpose_y) {
+          this->RunMatMulAndFusedOps(x, y, fused_ops_and_tensors, result,
+                                     transpose_x, transpose_y);
+        };
+
+    const QuantizedGraphRunner run_quantized =
+        [&](const Tensor& x, const Tensor& y,
+            const FusedOpsAndTensors& fused_ops_and_tensors, Tensor* result,
+            bool transpose_x, bool transpose_y, string input_quant_mode,
+            string output_quant_mode, bool is_bias_quantized,
+            bool is_perchannel, bool requantize, float output_min,
+            float output_max) {
+          this->RunQuantizedMatMul(
+              x, y, fused_ops_and_tensors, result, transpose_x, transpose_y,
+              input_quant_mode, output_quant_mode, is_bias_quantized,
+              is_perchannel, requantize, output_min, output_max);
+        };
+
+    const double atol = 1e-2;
+    constexpr int M = 3;
+    constexpr int K = 4;
+    constexpr int N = 5;
+    bool transpose_x = false;  // OpKernel does not support transpose_x.
+    std::vector<int64_t> x_dims;
+    std::vector<int64_t> y_dims;
+    FusedOpsAndDims fused_ops_and_dims;
+    GetFusionConfiguration(fused_ops, M, N, &fused_ops_and_dims);
+    std::vector<bool> requantization_config;
+    if (fused_ops == std::vector<string>{"BiasAdd", "Add"}) {
+      // MatMul + BiasAdd + Add + Requantize fusion is not supported yet.
+      requantization_config = {false};
+    } else {
+      requantization_config = {false, true};
+    }
+    for (bool transpose_y : {false, true}) {
+      x_dims =
+          transpose_x ? std::vector<int64_t>{K, M} : std::vector<int64_t>{M, K};
+      y_dims =
+          transpose_y ? std::vector<int64_t>{N, K} : std::vector<int64_t>{K, N};
+      for (bool per_channel : {false, true}) {
+        for (string input_quant_mode : {"MIN_FIRST", "SCALED"}) {
+          for (string output_quant_mode : {"MIN_FIRST", "SCALED"}) {
+            for (bool requantize : requantization_config) {
+              VerifyTensorsNear<QuantizedGraphRunner>(
+                  x_dims, y_dims, fused_ops_and_dims, run_default,
+                  run_quantized, transpose_x, transpose_y, atol,
+                  input_quant_mode, output_quant_mode, false, per_channel,
+                  requantize);
+            }
+          }
+        }
+      }
+    }
+  }
+};
+
+TYPED_TEST_SUITE_P(FusedMatMulOpsTest);
+
+// Realnumber typed test.
+TYPED_TEST_P(FusedMatMulOpsTest, BiasAddGeluApproximate) {
+  this->VerifyFusedMatMul({"BiasAdd", "GeluApproximate"});
+}
+
+// The following tests are for quantized fusions.
+TYPED_TEST_P(FusedMatMulOpsTest, Quantized_BiasAdd) {
+  this->VerifyQuantizedMatMul({"BiasAdd"});
+}
+
+TYPED_TEST_P(FusedMatMulOpsTest, Quantized_BiasAddRelu) {
+  this->VerifyQuantizedMatMul({"BiasAdd", "Relu"});
+}
+
+TYPED_TEST_P(FusedMatMulOpsTest, Quantized_BiasAddRelu6) {
+  this->VerifyQuantizedMatMul({"BiasAdd", "Relu6"});
+}
+
+TYPED_TEST_P(FusedMatMulOpsTest, Quantized_BiasAddLeakyRelu) {
+  this->VerifyQuantizedMatMul({"BiasAdd", "LeakyRelu"});
+}
+
+TYPED_TEST_P(FusedMatMulOpsTest, Quantized_BiasAddElu) {
+  this->VerifyQuantizedMatMul({"BiasAdd", "Elu"});
+}
+
+TYPED_TEST_P(FusedMatMulOpsTest, Quantized_BiasAddTanh) {
+  this->VerifyQuantizedMatMul({"BiasAdd", "Tanh"});
+}
+
+TYPED_TEST_P(FusedMatMulOpsTest, Quantized_BiasAddSigmoid) {
+  this->VerifyQuantizedMatMul({"BiasAdd", "Sigmoid"});
+}
+
+TYPED_TEST_P(FusedMatMulOpsTest, Quantized_BiasAddGeluApproximate) {
+  this->VerifyQuantizedMatMul({"BiasAdd", "GeluApproximate"});
+}
+
+TYPED_TEST_P(FusedMatMulOpsTest, Quantized_BiasAddAdd) {
+  this->VerifyQuantizedMatMul({"BiasAdd", "Add"});
+}
+
+REGISTER_TYPED_TEST_SUITE_P(FusedMatMulOpsTest, BiasAddGeluApproximate,
+                            Quantized_BiasAdd, Quantized_BiasAddRelu,
+                            Quantized_BiasAddRelu6, Quantized_BiasAddLeakyRelu,
+                            Quantized_BiasAddElu, Quantized_BiasAddTanh,
+                            Quantized_BiasAddSigmoid,
+                            Quantized_BiasAddGeluApproximate,
+                            Quantized_BiasAddAdd);
+
+// TODO(intel-tf): Add bfloat16 to Types when PR#56613 is merged.
+using DataTypes = ::testing::Types<float>;
+
+INSTANTIATE_TYPED_TEST_SUITE_P(Test, FusedMatMulOpsTest, DataTypes);
+
+}  // namespace tensorflow
+
+#endif  // INTEL_MKL
diff --git a/tensorflow/core/ops/mkl_nn_ops.cc b/tensorflow/core/ops/mkl_nn_ops.cc
index 868e15113b406c..14a42f29ddecbd 100644
--- a/tensorflow/core/ops/mkl_nn_ops.cc
+++ b/tensorflow/core/ops/mkl_nn_ops.cc
@@ -1937,6 +1937,46 @@ operation.
 expected to invoke these operators.
 )doc");
 
+REGISTER_OP("_QuantizedMatMul")
+    // Variable number of inputs depending on fusion. The inputs contain
+    // quantized or real tensors. Some of the inputs carry min-max values for
+    // quantized tensors.
+    .Input("device_inputs: Tdevice_inputs")
+    .Input("host_inputs: Thost_inputs")
+    // Variable number of outputs depending on the main output type. For
+    // example, quantized output will need additional tensors to carry min-max
+    // values. If the output type is real tensor (e.g. Dequantize fusion), the
+    // op should produce only single output tensor.
+    .Output("device_outputs: Tdevice_outputs")
+    .Output("host_outputs: Thost_outputs")
+    .Attr("Tdevice_inputs: list(type) >= 0 = []")
+    .Attr("Thost_inputs: list(type) >= 0 = []")
+    .Attr("Tdevice_outputs: list(type) >= 0 = []")
+    .Attr("Thost_outputs: list(type) >= 0 = []")
+    // The following attributes T1, T2, U, and Tout are members of Tinputs
+    // and Toutputs, used here for type constraints in the templatized OpKernel
+    // registrations.
+    .Attr("T1: quantizedtype")  // 0-th input    
+    .Attr("T2: quantizedtype")  // 1st input
+    .Attr("Tbias: {bfloat16, float, quantizedtype} = DT_FLOAT")
+    // Additional inputs' type. Currently, restricting all to be of same type.
+    .Attr("U: {bfloat16, float, quantizedtype} = DT_FLOAT")
+    .Attr("Tout: {bfloat16, float, quantizedtype} = DT_FLOAT")  // 0-th output  
+    .Attr("transpose_a: bool = false")
+    .Attr("transpose_b: bool = false")
+    .Attr("is_weight_const: bool = true")
+    .Attr("is_bias_const: bool = true")
+    .Attr("fused_ops: list(string) = []")
+    // Attribute for quantization mode of all quantized input tensors.
+    // Currently restricting all operands using same quantization mode.
+    .Attr("input_quant_mode: {'MIN_FIRST', 'SCALED'} = 'SCALED'")
+    // Attribute for activation (0-th output) requnatization mode
+    .Attr("output_quant_mode: {'MIN_FIRST', 'SCALED'} = 'SCALED'")
+    // Attributes for the LeakyRelu ----------------------------------------- //
+    .Attr("leakyrelu_alpha: float = 0.2")
+    // ---------------------------------------------------------------------- //
+    .SetShapeFn(shape_inference::MatMulShape);
+
 }  // namespace tensorflow
 
 #endif  // INTEL_MKL

From cf2678f166d5c13b7face1cebd23ce237108bd5c Mon Sep 17 00:00:00 2001
From: mdfaijul <md.faijul.amin@intel.com>
Date: Tue, 30 Apr 2024 11:52:11 -0700
Subject: [PATCH 002/478] Fix round mode.

---
 tensorflow/core/kernels/mkl/mkl_kernel_util.cc | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorflow/core/kernels/mkl/mkl_kernel_util.cc b/tensorflow/core/kernels/mkl/mkl_kernel_util.cc
index 02e5d0a578df9c..504247ff6613ec 100644
--- a/tensorflow/core/kernels/mkl/mkl_kernel_util.cc
+++ b/tensorflow/core/kernels/mkl/mkl_kernel_util.cc
@@ -50,7 +50,7 @@ void MklTestingUtil::RunMklQuantizeOp(const Tensor& input,
                   .Input(max_node)
                   .Attr("T", type)
                   .Attr("mode", mode)
-                  .Attr("round_mode", "HALF_TO_EVEN")
+                  .Attr("round_mode", round_mode)
                   .Attr("_kernel", "QuantizedMklOp")
                   .Finalize(&*graph, &quantize_op));
 

From 803d7917db1890f9f1af6be55a3f6a0b086a0f75 Mon Sep 17 00:00:00 2001
From: "dependabot[bot]" <49699333+dependabot[bot]@users.noreply.github.com>
Date: Mon, 6 May 2024 19:21:12 +0000
Subject: [PATCH 003/478] Bump werkzeug from 3.0.1 to 3.0.3

Bumps [werkzeug](https://github.com/pallets/werkzeug) from 3.0.1 to 3.0.3.
- [Release notes](https://github.com/pallets/werkzeug/releases)
- [Changelog](https://github.com/pallets/werkzeug/blob/main/CHANGES.rst)
- [Commits](https://github.com/pallets/werkzeug/compare/3.0.1...3.0.3)

---
updated-dependencies:
- dependency-name: werkzeug
  dependency-type: direct:production
...

Signed-off-by: dependabot[bot] <support@github.com>
---
 requirements_lock_3_10.txt | 6 +++---
 requirements_lock_3_11.txt | 6 +++---
 requirements_lock_3_12.txt | 6 +++---
 requirements_lock_3_9.txt  | 6 +++---
 4 files changed, 12 insertions(+), 12 deletions(-)

diff --git a/requirements_lock_3_10.txt b/requirements_lock_3_10.txt
index 05dc3940487eef..f17468ddaafd0a 100644
--- a/requirements_lock_3_10.txt
+++ b/requirements_lock_3_10.txt
@@ -522,9 +522,9 @@ urllib3==2.2.0 \
     --hash=sha256:051d961ad0c62a94e50ecf1af379c3aba230c66c710493493560c0c223c49f20 \
     --hash=sha256:ce3711610ddce217e6d113a2732fafad960a03fd0318c91faa79481e35c11224
     # via requests
-werkzeug==3.0.1 \
-    --hash=sha256:507e811ecea72b18a404947aded4b3390e1db8f826b494d76550ef45bb3b1dcc \
-    --hash=sha256:90a285dc0e42ad56b34e696398b8122ee4c681833fb35b8334a095d82c56da10
+werkzeug==3.0.3 \
+    --hash=sha256:097e5bfda9f0aba8da6b8545146def481d06aa7d3266e7448e2cccf67dd8bd18 \
+    --hash=sha256:fc9645dc43e03e4d630d23143a04a7f947a9a3b5727cd535fdfe155a17cc48c8
     # via tb-nightly
 wheel==0.41.3 \
     --hash=sha256:488609bc63a29322326e05560731bf7bfea8e48ad646e1f5e40d366607de0942 \
diff --git a/requirements_lock_3_11.txt b/requirements_lock_3_11.txt
index 05dc3940487eef..f17468ddaafd0a 100644
--- a/requirements_lock_3_11.txt
+++ b/requirements_lock_3_11.txt
@@ -522,9 +522,9 @@ urllib3==2.2.0 \
     --hash=sha256:051d961ad0c62a94e50ecf1af379c3aba230c66c710493493560c0c223c49f20 \
     --hash=sha256:ce3711610ddce217e6d113a2732fafad960a03fd0318c91faa79481e35c11224
     # via requests
-werkzeug==3.0.1 \
-    --hash=sha256:507e811ecea72b18a404947aded4b3390e1db8f826b494d76550ef45bb3b1dcc \
-    --hash=sha256:90a285dc0e42ad56b34e696398b8122ee4c681833fb35b8334a095d82c56da10
+werkzeug==3.0.3 \
+    --hash=sha256:097e5bfda9f0aba8da6b8545146def481d06aa7d3266e7448e2cccf67dd8bd18 \
+    --hash=sha256:fc9645dc43e03e4d630d23143a04a7f947a9a3b5727cd535fdfe155a17cc48c8
     # via tb-nightly
 wheel==0.41.3 \
     --hash=sha256:488609bc63a29322326e05560731bf7bfea8e48ad646e1f5e40d366607de0942 \
diff --git a/requirements_lock_3_12.txt b/requirements_lock_3_12.txt
index 120ec6ebcd7c72..0d045ea1a0579c 100644
--- a/requirements_lock_3_12.txt
+++ b/requirements_lock_3_12.txt
@@ -530,9 +530,9 @@ urllib3==2.2.0 \
     --hash=sha256:051d961ad0c62a94e50ecf1af379c3aba230c66c710493493560c0c223c49f20 \
     --hash=sha256:ce3711610ddce217e6d113a2732fafad960a03fd0318c91faa79481e35c11224
     # via requests
-werkzeug==3.0.1 \
-    --hash=sha256:507e811ecea72b18a404947aded4b3390e1db8f826b494d76550ef45bb3b1dcc \
-    --hash=sha256:90a285dc0e42ad56b34e696398b8122ee4c681833fb35b8334a095d82c56da10
+werkzeug==3.0.3 \
+    --hash=sha256:097e5bfda9f0aba8da6b8545146def481d06aa7d3266e7448e2cccf67dd8bd18 \
+    --hash=sha256:fc9645dc43e03e4d630d23143a04a7f947a9a3b5727cd535fdfe155a17cc48c8
     # via tb-nightly
 wheel==0.41.3 \
     --hash=sha256:488609bc63a29322326e05560731bf7bfea8e48ad646e1f5e40d366607de0942 \
diff --git a/requirements_lock_3_9.txt b/requirements_lock_3_9.txt
index 36a55514cd788b..48c74173fe553f 100644
--- a/requirements_lock_3_9.txt
+++ b/requirements_lock_3_9.txt
@@ -526,9 +526,9 @@ urllib3==2.2.0 \
     --hash=sha256:051d961ad0c62a94e50ecf1af379c3aba230c66c710493493560c0c223c49f20 \
     --hash=sha256:ce3711610ddce217e6d113a2732fafad960a03fd0318c91faa79481e35c11224
     # via requests
-werkzeug==3.0.1 \
-    --hash=sha256:507e811ecea72b18a404947aded4b3390e1db8f826b494d76550ef45bb3b1dcc \
-    --hash=sha256:90a285dc0e42ad56b34e696398b8122ee4c681833fb35b8334a095d82c56da10
+werkzeug==3.0.3 \
+    --hash=sha256:097e5bfda9f0aba8da6b8545146def481d06aa7d3266e7448e2cccf67dd8bd18 \
+    --hash=sha256:fc9645dc43e03e4d630d23143a04a7f947a9a3b5727cd535fdfe155a17cc48c8
     # via tb-nightly
 wheel==0.41.3 \
     --hash=sha256:488609bc63a29322326e05560731bf7bfea8e48ad646e1f5e40d366607de0942 \

From bfe2973e107bd6aaa12c29a3204a9b59ed018440 Mon Sep 17 00:00:00 2001
From: mdfaijul <md.faijul.amin@intel.com>
Date: Sat, 11 May 2024 13:41:38 -0700
Subject: [PATCH 004/478] Address review comments.

---
 tensorflow/core/kernels/mkl/mkl_kernel_util.h             | 4 ++--
 tensorflow/core/kernels/mkl/mkl_matmul_op_fused.cc        | 8 ++++----
 tensorflow/core/kernels/mkl/mkl_matmul_ops_common.h       | 2 +-
 .../core/kernels/mkl/onednn_fused_matmul_ops_test.cc      | 5 ++---
 4 files changed, 9 insertions(+), 10 deletions(-)

diff --git a/tensorflow/core/kernels/mkl/mkl_kernel_util.h b/tensorflow/core/kernels/mkl/mkl_kernel_util.h
index a3015a7b09025c..da600fb001e038 100644
--- a/tensorflow/core/kernels/mkl/mkl_kernel_util.h
+++ b/tensorflow/core/kernels/mkl/mkl_kernel_util.h
@@ -56,8 +56,8 @@ class MklTestingUtil {
   // their dtype set to DT_FLOAT.
   template <typename T>
   static Status GetQuantizationTensors(const Tensor& input, Tensor* output,
-                                     DataType out_type, const string mode,
-                                     Tensor* min_tensor, Tensor* max_tensor) {
+                                       DataType out_type, const string mode,
+                                       Tensor* min_tensor, Tensor* max_tensor) {
     if (min_tensor->dtype() != DT_FLOAT || max_tensor->dtype() != DT_FLOAT) {
       return absl::UnimplementedError("Tensor must be float32.");
     }
diff --git a/tensorflow/core/kernels/mkl/mkl_matmul_op_fused.cc b/tensorflow/core/kernels/mkl/mkl_matmul_op_fused.cc
index 7e5544f1c3e3b4..a78f81d09db25f 100644
--- a/tensorflow/core/kernels/mkl/mkl_matmul_op_fused.cc
+++ b/tensorflow/core/kernels/mkl/mkl_matmul_op_fused.cc
@@ -27,7 +27,7 @@ limitations under the License.
 #include "tensorflow/core/kernels/mkl/mkl_matmul_ops_common.h"
 #include "tensorflow/core/kernels/mkl/mkl_quantized_conv_ops.h"
 #include "tensorflow/core/lib/core/errors.h"
-#include "tensorflow/core/lib/gtl/inlined_vector.h"
+#include "absl/container/inlined_vector.h"
 #include "tensorflow/core/platform/errors.h"
 
 namespace tensorflow {
@@ -409,14 +409,14 @@ class QuantizedFusedMatMulOp
     struct OperandInfo {
       int idx = -1;  // Operand tensor index if needed by a post-op.
       // Indices of min and max value tensors, if the operand is quantized.
-      gtl::InlinedVector<int, 4> min_max_indices;
+      absl::InlinedVector<int, 4> min_max_indices;
     } operand_info;
     // Indices of output min and max value tensors. It is used when requantize
     // is fused.
-    gtl::InlinedVector<int, 4> min_max_indices;
+    absl::InlinedVector<int, 4> min_max_indices;
   };
 
-  gtl::InlinedVector<PostOpInfo, 4> post_op_info_list_;
+  absl::InlinedVector<PostOpInfo, 4> post_op_info_list_;
 
   void Initialize(OpKernelConstruction* context) {
     OP_REQUIRES_OK(context,
diff --git a/tensorflow/core/kernels/mkl/mkl_matmul_ops_common.h b/tensorflow/core/kernels/mkl/mkl_matmul_ops_common.h
index 6159b1d047ea65..dc4dfcb3663048 100644
--- a/tensorflow/core/kernels/mkl/mkl_matmul_ops_common.h
+++ b/tensorflow/core/kernels/mkl/mkl_matmul_ops_common.h
@@ -102,7 +102,7 @@ struct MklDnnMatMulFwdParams {
   struct PostOpParam {
     string name;
     std::vector<float> param;
-    string partial_key = string("");
+    string partial_key;
   };
   std::vector<PostOpParam> post_op_params;
   string input_quant_mode;
diff --git a/tensorflow/core/kernels/mkl/onednn_fused_matmul_ops_test.cc b/tensorflow/core/kernels/mkl/onednn_fused_matmul_ops_test.cc
index af1d056ca7530a..c26535012b2aee 100644
--- a/tensorflow/core/kernels/mkl/onednn_fused_matmul_ops_test.cc
+++ b/tensorflow/core/kernels/mkl/onednn_fused_matmul_ops_test.cc
@@ -491,7 +491,6 @@ class FusedMatMulOpsTest : public OpsTestBase {
                          string output_quant_mode = "SCALED",
                          bool is_bias_quantized = false,
                          bool is_perchannel = false, bool requantize = false) {
-    srand(1234);
     DataType dtype = DataTypeToEnum<T>::v();
     TensorShape x_shape = TensorShape(x_dims);
     TensorShape y_shape = TensorShape(y_dims);
@@ -532,13 +531,13 @@ class FusedMatMulOpsTest : public OpsTestBase {
         output_min = static_cast<float>(min);
         output_max = static_cast<float>(max);
       }
-      // Run quantized fusion
+      // Run quantized fusion.
       run_fused(x_tensor, y_tensor, fused_ops_and_tensors, &fused_result,
                 transpose_x, transpose_y, input_quant_mode, output_quant_mode,
                 is_bias_quantized, is_perchannel, requantize, output_min,
                 output_max);
     } else {
-      // Run realnumber type fusion
+      // Run realnumber type fusion.
       run_fused(x_tensor, y_tensor, fused_ops_and_tensors, &fused_result,
                 transpose_x, transpose_y);
     }

From 48fee16c70295881b78fe1a394480764a74114b1 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Mon, 13 May 2024 02:02:12 -0700
Subject: [PATCH 005/478] Update GraphDef version to 1861.

PiperOrigin-RevId: 633133739
---
 tensorflow/core/public/version.h | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorflow/core/public/version.h b/tensorflow/core/public/version.h
index ad2911fae4c67f..aafe138460a6db 100644
--- a/tensorflow/core/public/version.h
+++ b/tensorflow/core/public/version.h
@@ -108,7 +108,7 @@ limitations under the License.
 
 #define TF_GRAPH_DEF_VERSION_MIN_PRODUCER 0
 #define TF_GRAPH_DEF_VERSION_MIN_CONSUMER 0
-#define TF_GRAPH_DEF_VERSION 1860  // Updated: 2024/5/12
+#define TF_GRAPH_DEF_VERSION 1861  // Updated: 2024/5/13
 
 // Checkpoint compatibility versions (the versions field in SavedSliceMeta).
 //

From c04d2ba725b3ed6dec7cab9162adbb42bcc29151 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Mon, 13 May 2024 02:02:12 -0700
Subject: [PATCH 006/478] compat: Update forward compatibility horizon to
 2024-05-13

PiperOrigin-RevId: 633133741
---
 tensorflow/python/compat/compat.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorflow/python/compat/compat.py b/tensorflow/python/compat/compat.py
index 21f795a3a55a41..2a9869afa5dfb4 100644
--- a/tensorflow/python/compat/compat.py
+++ b/tensorflow/python/compat/compat.py
@@ -29,7 +29,7 @@
 # This value changes every day with an automatic CL. It can be modified in code
 # via `forward_compatibility_horizon()` or with the environment variable
 # TF_FORWARD_COMPATIBILITY_DELTA_DAYS, which is added to the compatibility date.
-_FORWARD_COMPATIBILITY_HORIZON = datetime.date(2024, 5, 12)
+_FORWARD_COMPATIBILITY_HORIZON = datetime.date(2024, 5, 13)
 _FORWARD_COMPATIBILITY_DELTA_DAYS_VAR_NAME = "TF_FORWARD_COMPATIBILITY_DELTA_DAYS"
 _FORWARD_COMPATIBILITY_DATE_NUMBER = None
 

From 76a05ff1230d0c4100ebce395b992c05700fedca Mon Sep 17 00:00:00 2001
From: George Karpenkov <cheshire@google.com>
Date: Mon, 13 May 2024 02:19:29 -0700
Subject: [PATCH 007/478] [XLA:NFC] Support generating LLVM CPU output for
 hlo-opt

PiperOrigin-RevId: 633137810
---
 third_party/xla/xla/tools/hlo_opt/BUILD       |  7 +++++
 .../xla/xla/tools/hlo_opt/cpu_llvm.hlo        | 19 ++++++++++++++
 third_party/xla/xla/tools/hlo_opt/cpu_opt.cc  | 26 +++++++++++++++++++
 3 files changed, 52 insertions(+)
 create mode 100644 third_party/xla/xla/tools/hlo_opt/cpu_llvm.hlo

diff --git a/third_party/xla/xla/tools/hlo_opt/BUILD b/third_party/xla/xla/tools/hlo_opt/BUILD
index 06e2380b70a542..dd7f62731a27f9 100644
--- a/third_party/xla/xla/tools/hlo_opt/BUILD
+++ b/third_party/xla/xla/tools/hlo_opt/BUILD
@@ -100,10 +100,16 @@ cc_library(
     srcs = ["cpu_opt.cc"],
     deps = [
         ":opt_lib",
+        "//xla/hlo/ir:hlo",
         "//xla/service:cpu_plugin",
+        "//xla/service:executable",
         "//xla/service:hlo_graph_dumper",
+        "//xla/service/cpu:cpu_executable",
         "//xla/stream_executor/host:host_platform",
         "//xla/stream_executor/platform",
+        "@com_google_absl//absl/status:statusor",
+        "@com_google_absl//absl/strings",
+        "@local_tsl//tsl/platform:statusor",
     ],
     alwayslink = True,  # Initializer needs to run.
 )
@@ -149,6 +155,7 @@ lit_test_suite(
     srcs = enforce_glob(
         [
             "cpu_hlo.hlo",
+            "cpu_llvm.hlo",
             "gpu_hlo.hlo",
             "gpu_hlo_backend.hlo",
             "gpu_hlo_buffers.hlo",
diff --git a/third_party/xla/xla/tools/hlo_opt/cpu_llvm.hlo b/third_party/xla/xla/tools/hlo_opt/cpu_llvm.hlo
new file mode 100644
index 00000000000000..fbb033e53b07eb
--- /dev/null
+++ b/third_party/xla/xla/tools/hlo_opt/cpu_llvm.hlo
@@ -0,0 +1,19 @@
+// RUN: hlo-opt %s --platform=cpu --stage=llvm-before-optimizations --split-input-file | FileCheck --check-prefixes=CHECK %s
+
+HloModule m
+
+add {
+ a = s8[] parameter(0)
+ b = s8[] parameter(1)
+ ROOT out = s8[] add(a, b)
+}
+
+// CHECK: i8
+ENTRY e {
+ p1 = s8[1048576] parameter(0)
+ i = s8[] constant(0)
+ ROOT out = s8[] reduce(p1, i), dimensions={0}, to_apply=add
+}
+
+
+
diff --git a/third_party/xla/xla/tools/hlo_opt/cpu_opt.cc b/third_party/xla/xla/tools/hlo_opt/cpu_opt.cc
index 9e7d5c2b72ace7..fc5e540fc7907a 100644
--- a/third_party/xla/xla/tools/hlo_opt/cpu_opt.cc
+++ b/third_party/xla/xla/tools/hlo_opt/cpu_opt.cc
@@ -14,10 +14,19 @@ limitations under the License.
 ==============================================================================*/
 
 #include <memory>
+#include <optional>
+#include <set>
 #include <string>
+#include <utility>
 
+#include "absl/status/statusor.h"
+#include "absl/strings/string_view.h"
+#include "xla/hlo/ir/hlo_module.h"
+#include "xla/service/cpu/cpu_executable.h"
+#include "xla/service/executable.h"
 #include "xla/stream_executor/platform/initialize.h"
 #include "xla/tools/hlo_opt/opt_lib.h"
+#include "tsl/platform/statusor.h"
 
 namespace xla {
 
@@ -25,6 +34,23 @@ namespace {
 
 class CpuOptProvider : public OptProvider {
  public:
+  absl::StatusOr<std::optional<std::string>> GenerateStage(
+      std::unique_ptr<HloModule> module, absl::string_view s) override {
+    if (s == "llvm-before-optimizations") {
+      TF_ASSIGN_OR_RETURN(std::unique_ptr<Executable> executable,
+                          GetExecutable(std::move(module)));
+      return static_cast<cpu::CpuExecutable*>(executable.get())
+          ->ir_module_string();
+    }
+    return OptProvider::GenerateStage(std::move(module), s);
+  }
+
+  std::set<std::string> SupportedStages() override {
+    std::set<std::string> supported = OptProvider::SupportedStages();
+    supported.insert({"llvm-before-optimizations"});
+    return supported;
+  }
+
   std::string GetPlatformName() override { return "cpu"; }
 };
 

From f0e890dd698abda68b7204c68bf3d4e0fb1056af Mon Sep 17 00:00:00 2001
From: George Karpenkov <cheshire@google.com>
Date: Mon, 13 May 2024 02:28:48 -0700
Subject: [PATCH 008/478] [XLA] Change default hlo-bisect output format to hlo

PiperOrigin-RevId: 633139592
---
 third_party/xla/xla/tools/hlo_bisect/hlo_bisect.cc | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/third_party/xla/xla/tools/hlo_bisect/hlo_bisect.cc b/third_party/xla/xla/tools/hlo_bisect/hlo_bisect.cc
index 73b018323f34c3..588e4f52999172 100644
--- a/third_party/xla/xla/tools/hlo_bisect/hlo_bisect.cc
+++ b/third_party/xla/xla/tools/hlo_bisect/hlo_bisect.cc
@@ -50,7 +50,7 @@ struct BisectOptions {
   std::string input = "";
   std::string script = "";
   std::string dump_path = "/tmp/hlo_bisect";
-  std::string output_format = "pb";
+  std::string output_format = "hlo";
   bool all_computations = false;
   std::string test_platform = "CUDA";
   std::string reference_platform = "Interpreter";

From 9213f23cdf5f1366ed569c2626963c660ba136fd Mon Sep 17 00:00:00 2001
From: George Karpenkov <cheshire@google.com>
Date: Mon, 13 May 2024 02:57:51 -0700
Subject: [PATCH 009/478] [XLA:GPU] [NFC] Propagate Status from allocation
 failures properly

Do not attempt to wrap/rewrap Status, as that loses the attached stack trace, and causes further errors.

PiperOrigin-RevId: 633145137
---
 third_party/xla/xla/service/gpu/gpu_executable.cc | 12 ++++--------
 1 file changed, 4 insertions(+), 8 deletions(-)

diff --git a/third_party/xla/xla/service/gpu/gpu_executable.cc b/third_party/xla/xla/service/gpu/gpu_executable.cc
index d815974e98aec7..4c56c84cdadbe8 100644
--- a/third_party/xla/xla/service/gpu/gpu_executable.cc
+++ b/third_party/xla/xla/service/gpu/gpu_executable.cc
@@ -754,16 +754,12 @@ absl::StatusOr<se::DeviceMemoryBase> GpuExecutable::BufferForAllocation(
     const int64_t buffer_size = allocation.size();
     se::DeviceMemoryBase buffer_address;
     if (buffer_size > 0) {
-      absl::StatusOr<se::OwningDeviceMemory> buffer =
+      TF_ASSIGN_OR_RETURN(
+          se::OwningDeviceMemory buffer,
           memory_allocator->Allocate(device_ordinal, buffer_size,
                                      /*retry_on_failure=*/true,
-                                     /*memory_space=*/allocation.color());
-      if (!buffer.ok()) {
-        return ResourceExhausted("%s\n%s\n", buffer.status().message(),
-                                 buffer_assignment_->ToVerboseString(
-                                     debug_buffer_assignment_show_max_));
-      }
-      buffer_address = buffer->Release();
+                                     /*memory_space=*/allocation.color()));
+      buffer_address = buffer.Release();
     }
     return buffer_address;
   }

From befd651123afc501e5e8d7f15a4a409f8654adc4 Mon Sep 17 00:00:00 2001
From: Johannes Reifferscheid <jreiffers@google.com>
Date: Mon, 13 May 2024 03:08:38 -0700
Subject: [PATCH 010/478] Add reduction group support.

After this, all legacy fusions are supported by MLIR emitters. Currently, there
is still one failure in the JAX test suite, but it is not a reduction fusion.

PiperOrigin-RevId: 633147696
---
 third_party/xla/xla/service/gpu/fusions/BUILD |   5 +-
 .../xla/xla/service/gpu/fusions/fusions.cc    |  36 +-
 .../fusions/mlir/computation_partitioner.cc   |   2 +
 .../gpu/fusions/mlir/mlir_fusion_emitter.cc   |   6 +-
 .../xla/service/gpu/fusions/reduction_mlir.cc | 388 +++++++++---------
 .../xla/service/gpu/fusions/reduction_mlir.h  |  14 +-
 .../gpu/fusions/reduction_mlir_test.cc        |  31 ++
 7 files changed, 258 insertions(+), 224 deletions(-)

diff --git a/third_party/xla/xla/service/gpu/fusions/BUILD b/third_party/xla/xla/service/gpu/fusions/BUILD
index 77e69877f3b5e3..cb89eb885d4e59 100644
--- a/third_party/xla/xla/service/gpu/fusions/BUILD
+++ b/third_party/xla/xla/service/gpu/fusions/BUILD
@@ -826,9 +826,9 @@ cc_library(
     srcs = ["reduction_mlir.cc"],
     hdrs = ["reduction_mlir.h"],
     deps = [
+        ":fusion_emitter",
         ":reduction_base",
         "//xla:shape_util",
-        "//xla:status_macros",
         "//xla/hlo/ir:hlo",
         "//xla/service/gpu:hlo_fusion_analysis",
         "//xla/service/gpu:ir_emission_utils",
@@ -842,6 +842,7 @@ cc_library(
         "//xla/service/gpu/model:indexing_map",
         "@com_google_absl//absl/algorithm:container",
         "@com_google_absl//absl/container:flat_hash_map",
+        "@com_google_absl//absl/container:flat_hash_set",
         "@com_google_absl//absl/status",
         "@com_google_absl//absl/types:span",
         "@llvm-project//llvm:Support",
@@ -865,11 +866,9 @@ xla_cc_test(
         ":reduction_mlir",
         "//xla:error_spec",
         "//xla/service:gpu_plugin",
-        "//xla/tests:filecheck",
         "//xla/tests:xla_internal_test_main",
         "@com_google_googletest//:gtest",
         "@local_tsl//tsl/lib/core:status_test_util",
-        "@local_tsl//tsl/platform:statusor",
     ],
 )
 
diff --git a/third_party/xla/xla/service/gpu/fusions/fusions.cc b/third_party/xla/xla/service/gpu/fusions/fusions.cc
index a44477e0f9cc42..3567463575539f 100644
--- a/third_party/xla/xla/service/gpu/fusions/fusions.cc
+++ b/third_party/xla/xla/service/gpu/fusions/fusions.cc
@@ -141,24 +141,16 @@ absl::StatusOr<std::unique_ptr<FusionInterface>> GetFusionEmitter(
                          .GetModule()
                          ->config()
                          .debug_options();
-  auto check_mlir_emitters = [&](std::function<bool(const HloFusionAnalysis&)>
-                                     support_check) {
+  auto check_mlir_emitters = [&](bool check = true) {
     if (!opts.xla_gpu_enable_mlir_emitters()) {
       return false;
     }
-    if (!mlir_converter::IsHloConversionSupported(
-            analysis.fusion(),
-            fusion_info.analysis().device_info().gpu_compute_capability())) {
-      VLOG(5) << "Skipping MLIR emission because the fusion contains "
-                 "unsupported instructions.";
-      return false;
-    }
-    if (support_check && !support_check(analysis)) {
-      VLOG(5) << "Skipping MLIR emission because the fusion emitter does not "
-                 "support "
-                 "the fusion.";
-      return false;
-    }
+    CHECK(!check ||
+          mlir_converter::IsHloConversionSupported(
+              analysis.fusion(),
+              fusion_info.analysis().device_info().gpu_compute_capability()))
+        << "Unsupported fusion: "
+        << analysis.fusion_root(0).instruction().parent()->ToString();
 
     static int num_mlir_emitters = 0;
     if (is_emission_phase) {
@@ -191,14 +183,14 @@ absl::StatusOr<std::unique_ptr<FusionInterface>> GetFusionEmitter(
       return std::make_unique<CustomFusion>();
     }
     case HloFusionAnalysis::EmitterFusionKind::kInputSlices:
-      if (check_mlir_emitters(nullptr)) {
+      if (check_mlir_emitters()) {
         return std::make_unique<MlirInputSlicesFusion>(analysis);
       }
       return std::make_unique<InputSlicesFusion>(analysis);
     case HloFusionAnalysis::EmitterFusionKind::kLoop: {
       if (IsDynamicUpdateSliceFusion(analysis) &&
           fusion_info.CanEmitDynamicUpdateSliceInPlace()) {
-        if (check_mlir_emitters(nullptr)) {
+        if (check_mlir_emitters()) {
           return std::make_unique<MlirInPlaceDynamicUpdateSliceFusion>(
               analysis);
         }
@@ -209,30 +201,30 @@ absl::StatusOr<std::unique_ptr<FusionInterface>> GetFusionEmitter(
         return *std::move(copy_fusion);
       }
 
-      if (check_mlir_emitters(nullptr)) {
+      if (check_mlir_emitters()) {
         return std::make_unique<MlirLoopFusion>(analysis);
       }
       return std::make_unique<LoopFusion>(analysis);
     }
     case HloFusionAnalysis::EmitterFusionKind::kReduction:
-      if (check_mlir_emitters(MlirReductionFusion::IsSupported)) {
+      if (check_mlir_emitters()) {
         return std::make_unique<MlirReductionFusion>(analysis);
       }
       return std::make_unique<ReductionFusion>(analysis);
     case HloFusionAnalysis::EmitterFusionKind::kScatter: {
-      if (check_mlir_emitters(MlirScatterFusion::IsSupported)) {
+      if (check_mlir_emitters(false)) {
         return std::make_unique<MlirScatterFusion>(analysis);
       }
       return std::make_unique<ScatterFusion>(analysis);
     }
     case HloFusionAnalysis::EmitterFusionKind::kTranspose: {
-      if (check_mlir_emitters(nullptr)) {
+      if (check_mlir_emitters()) {
         return std::make_unique<MlirTransposeFusion>(analysis);
       }
       return std::make_unique<TransposeFusion>(analysis);
     }
     case HloFusionAnalysis::EmitterFusionKind::kConcatenate: {
-      if (check_mlir_emitters(nullptr)) {
+      if (check_mlir_emitters()) {
         return std::make_unique<MlirConcatenateFusion>(analysis);
       }
       return std::make_unique<ConcatenateFusion>(analysis);
diff --git a/third_party/xla/xla/service/gpu/fusions/mlir/computation_partitioner.cc b/third_party/xla/xla/service/gpu/fusions/mlir/computation_partitioner.cc
index dffbb918b2d960..97c3cba8b3892f 100644
--- a/third_party/xla/xla/service/gpu/fusions/mlir/computation_partitioner.cc
+++ b/third_party/xla/xla/service/gpu/fusions/mlir/computation_partitioner.cc
@@ -321,6 +321,7 @@ PartitionedComputation::PartitionedComputation(
 
 PartitionedComputation::Subgraph PartitionedComputation::Subgraph::ForEpilogue(
     const EpilogueSpecification& epilogue) {
+  if (epilogue.roots.empty()) return {};
   const auto* computation = epilogue.heroes.front()->parent();
   PartitionedComputation::Subgraph subgraph;
   subgraph.name = llvm_ir::SanitizeFunctionName(
@@ -407,6 +408,7 @@ PartitionedComputations::DeclareFunctions(mlir::ModuleOp module) const {
   auto create_funcs =
       [&](absl::Span<const PartitionedComputation::Subgraph> subgraphs) {
         for (const auto& subgraph : subgraphs) {
+          if (subgraph.roots.empty()) continue;
           auto func_op = CreateSubgraphMlirFunction(subgraph, builder);
           func_op->setAttr("llvm.linkage", mlir::LLVM::LinkageAttr::get(
                                                module->getContext(),
diff --git a/third_party/xla/xla/service/gpu/fusions/mlir/mlir_fusion_emitter.cc b/third_party/xla/xla/service/gpu/fusions/mlir/mlir_fusion_emitter.cc
index d0969dde6517cb..286b3e211df558 100644
--- a/third_party/xla/xla/service/gpu/fusions/mlir/mlir_fusion_emitter.cc
+++ b/third_party/xla/xla/service/gpu/fusions/mlir/mlir_fusion_emitter.cc
@@ -480,7 +480,7 @@ absl::Status MlirFusionEmitterBase::EmitMlir(
 
   // The epilogue functions replace the root tuple.
   auto* root = fusion.fused_instructions_computation()->root_instruction();
-  if (!epilogues.empty() && root->opcode() == HloOpcode::kTuple) {
+  if (root->opcode() == HloOpcode::kTuple && !epilogues.empty()) {
     subgraph_to_mlir_fn.extract(&computations.FindSubgraph(root))
         .mapped()
         .erase();
@@ -497,6 +497,7 @@ absl::Status MlirFusionEmitterBase::EmitMlir(
     }
   }
   for (const auto& epilogue : computations.epilogues()) {
+    if (epilogue.roots.empty()) continue;
     TF_RETURN_IF_ERROR(mlir_converter::SubgraphToMlirFunction(
         computations.FindPartitionedComputation(
             fusion.fused_instructions_computation()),
@@ -524,6 +525,9 @@ MlirFusionEmitterBase::EmitEpilogue(
         injected,
     ValueRange output_indices, mlir::ImplicitLocOpBuilder& builder) const {
   const auto& epilogue = computations.epilogues().at(epilogue_index);
+  if (epilogue.roots.empty()) {
+    return {};
+  }
   auto epilogue_fn = mlir::cast<mlir::func::FuncOp>(
       entry_fn->getParentOfType<mlir::ModuleOp>().lookupSymbol(epilogue.name));
   SmallVector<Value> operands = ValueRange(entry_fn.getArguments().take_front(
diff --git a/third_party/xla/xla/service/gpu/fusions/reduction_mlir.cc b/third_party/xla/xla/service/gpu/fusions/reduction_mlir.cc
index bd75a7e2ecbea5..8c1dcdaf3b7bf9 100644
--- a/third_party/xla/xla/service/gpu/fusions/reduction_mlir.cc
+++ b/third_party/xla/xla/service/gpu/fusions/reduction_mlir.cc
@@ -15,14 +15,14 @@ limitations under the License.
 #include "xla/service/gpu/fusions/reduction_mlir.h"
 
 #include <cstdint>
-#include <iterator>
+#include <utility>
 #include <vector>
 
 #include "absl/algorithm/container.h"
 #include "absl/container/flat_hash_map.h"
+#include "absl/container/flat_hash_set.h"
 #include "absl/status/status.h"
 #include "absl/types/span.h"
-#include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/SmallVector.h"
 #include "mlir/Dialect/Arith/IR/Arith.h"  // from @llvm-project
@@ -40,6 +40,7 @@ limitations under the License.
 #include "mlir/IR/ValueRange.h"  // from @llvm-project
 #include "xla/hlo/ir/hlo_instruction.h"
 #include "xla/hlo/ir/hlo_instructions.h"
+#include "xla/service/gpu/fusions/fusion_emitter.h"
 #include "xla/service/gpu/fusions/mlir/computation_partitioner.h"
 #include "xla/service/gpu/fusions/mlir/elemental_hlo_to_mlir.h"
 #include "xla/service/gpu/fusions/mlir/ir/xla_gpu_ops.h"
@@ -51,11 +52,11 @@ limitations under the License.
 #include "xla/service/gpu/model/indexing_map.h"
 #include "xla/service/gpu/reduction_utils.h"
 #include "xla/shape_util.h"
-#include "xla/status_macros.h"
 
 namespace xla {
 namespace gpu {
 
+namespace ma = mlir::arith;
 using llvm::SmallVector;
 using mlir::Value;
 using mlir::ValueRange;
@@ -75,7 +76,9 @@ struct MlirReductionFusion::EmitterState {
         fusion(fusion),
         computations(computations),
         call_target(call_target),
-        builder(entry_function.getLoc(), entry_function) {
+        builder(entry_function.getLoc(), entry_function),
+        computation(computations.FindPartitionedComputation(
+            fusion.fused_instructions_computation())) {
     int index = 0;
     for (auto root : owner.analysis().fusion_roots()) {
       fusion_result_index_starts[root] = index;
@@ -83,9 +86,11 @@ struct MlirReductionFusion::EmitterState {
     }
   }
 
-  // Uses the given indexing map to reduce a subset of the inputs in a single
-  // thread. The subset may be a single element.
-  HloValueMap EmitPerThreadReducedElements(const HloValueMap& inits);
+  // Reduces a subset of the inputs in a single thread. Also writes side outputs
+  // to the output tensors. The map contains the reduced values for reductions
+  // and the written tensors for side outputs.
+  HloValueMap EmitPerThreadReducedElements(int group_id,
+                                           const HloValueMap& inits);
 
   mlir::func::FuncOp GetReducer(const HloInstruction* hero) const {
     return call_target(hero->called_computations()[0]->root_instruction());
@@ -109,87 +114,107 @@ struct MlirReductionFusion::EmitterState {
   const PartitionedComputations& computations;
   const mlir_converter::CallTargetProvider& call_target;
   mlir::ImplicitLocOpBuilder builder;
+  const mlir_converter::PartitionedComputation& computation;
   absl::flat_hash_map<const HloInstruction*, int> fusion_result_index_starts;
+  SmallVector<Value> thread_and_block_ids;
 };
 
 MlirReductionFusion::MlirReductionFusion(const HloFusionAnalysis& analysis)
     : ReductionFusionBase(analysis) {
+  CHECK(reduction_info().IsRaceFree())
+      << "Non-race-free reductions should have been decomposed. Did "
+         "tree_reduction_rewriter run?";
+
+  const auto& groups = reduction_info().GetGroups();
+  int num_groups = groups.grouped_roots.size();
+  side_output_roots_.resize(num_groups);
+  reduction_heroes_.resize(num_groups);
+  reduction_roots_.resize(num_groups);
+
   absl::flat_hash_set<const HloInstruction*> seen_heroes;
-  const auto& is_reduction_root =
-      reduction_info().GetGroups().is_reduction_root;
-  first_reduction_root_index_ = std::distance(
-      is_reduction_root.begin(), absl::c_find(is_reduction_root, true));
-  for (auto [root, hero, is_reduction] :
+  for (auto [root, hero, is_reduction, group_id] :
        llvm::zip(analysis.fusion_roots(), analysis.fusion_heroes(),
-                 reduction_info().GetGroups().is_reduction_root)) {
-    (is_reduction ? reduction_roots_ : side_output_roots_).push_back(root);
-    if (is_reduction && seen_heroes.insert(hero).second) {
-      reduction_heroes_.push_back(hero);
+                 groups.is_reduction_root, groups.group_id_per_root)) {
+    if (is_reduction) {
+      if (seen_heroes.insert(hero).second) {
+        reduction_heroes_[group_id].push_back(hero);
+      }
+      reduction_roots_[group_id].push_back(root);
+    } else {
+      side_output_roots_[group_id].push_back(root);
     }
   }
 }
 
-bool MlirReductionFusion::IsSupported(const HloFusionAnalysis& analysis) {
-  auto info = ReductionInfo::Create(analysis, /*for_mlir=*/true);
-  return info.GetGroups().grouped_roots.size() == 1 && info.IsRaceFree();
-}
-
 std::vector<mlir_converter::EpilogueSpecification>
 MlirReductionFusion::GetEpilogues(const HloFusionInstruction& fusion,
                                   mlir::MLIRContext* mlir_context) const {
-  return {mlir_converter::EpilogueSpecification::FromOutputIndexing(
-      analysis(), reduction_heroes_, reduction_roots_, *this, mlir_context)};
+  std::vector<mlir_converter::EpilogueSpecification> epilogues;
+  epilogues.reserve(reduction_heroes_.size());
+  for (const auto& [heroes, roots] :
+       llvm::zip(reduction_heroes_, reduction_roots_)) {
+    epilogues.push_back(
+        mlir_converter::EpilogueSpecification::FromOutputIndexing(
+            analysis(), heroes, roots, *this, mlir_context));
+  }
+  return epilogues;
 }
 
 absl::Status MlirReductionFusion::EmitEntryFunction(
-    const mlir_converter::PartitionedComputations& computations,
+    const PartitionedComputations& computations,
     const mlir_converter::CallTargetProvider& call_targets,
     mlir::func::FuncOp entry_function,
     const HloFusionInstruction& fusion) const {
-  // Reduction groups will probably be implemented in a separate pass, since
-  // they share nothing by definition.
-  TF_RET_CHECK(reduction_info().GetGroups().grouped_roots.size() == 1)
-      << "Only one reduction group is supported.";
   EmitterState state{*this, entry_function, fusion, computations, call_targets};
-  state.builder.setInsertionPointToStart(entry_function.addEntryBlock());
-  return EmitReduction(state);
+  auto& b = state.builder;
+  b.setInsertionPointToStart(entry_function.addEntryBlock());
+  state.thread_and_block_ids = EmitThreadAndBlockIds(b);
+  if (reduction_heroes_.size() == 1) {
+    b.create<mlir::func::ReturnOp>(EmitReduction(0, state));
+    return absl::OkStatus();
+  }
+  SmallVector<int64_t> cases(reduction_heroes_.size() - 1);
+  absl::c_iota(cases, 1);  // `default` is region 0.
+  auto switch_op = b.create<mlir::scf::IndexSwitchOp>(
+      entry_function.getResultTypes(), EmitBlockId(b, 1), cases, cases.size());
+  b.create<mlir::func::ReturnOp>(switch_op.getResults());
+  for (auto [id, region] : llvm::enumerate(switch_op->getRegions())) {
+    b.setInsertionPointToStart(&region.emplaceBlock());
+    b.create<mlir::scf::YieldOp>(EmitReduction(id, state));
+  }
+  return absl::OkStatus();
 }
 
-absl::Status MlirReductionFusion::EmitReduction(EmitterState& state) const {
-  auto& builder = state.builder;
+llvm::SmallVector<Value> MlirReductionFusion::EmitReduction(
+    int group_id, EmitterState& state) const {
+  auto& b = state.builder;
   const auto& tiling = reduction_info().GetTiling();
+  const auto& threads_per_block = tiling.GetThreadsPerBlock();
+  auto* ctx = state.entry_function.getContext();
 
   // The number of warps working on one element in a row reduction.
-  int num_warps_row = tiling.GetThreadsPerBlock()
-                          [ReductionDimensions::kRowMinorReducedDimension] /
-                      WarpSize();
-  auto ctx = state.entry_function.getContext();
-
-  auto zero = builder.create<mlir::arith::ConstantIndexOp>(0);
-  auto lane_id = builder.create<mlir::gpu::LaneIdOp>();
-  auto is_first_lane = builder.create<mlir::arith::CmpIOp>(
-      mlir::arith::CmpIPredicate::eq, lane_id, zero);
-  auto thread_id = EmitThreadId(builder, 0);
-  auto block_id = EmitBlockId(builder, 0);
-  Value cst_true = builder.create<mlir::arith::ConstantOp>(
-      builder.getIntegerAttr(builder.getI1Type(), 1));
-
+  int num_warps_row =
+      threads_per_block[ReductionDimensions::kRowMinorReducedDimension] /
+      WarpSize();
+
+  auto zero = b.create<ma::ConstantIndexOp>(0);
+  auto lane_id = b.create<mlir::gpu::LaneIdOp>();
+  auto is_first_lane =
+      b.create<ma::CmpIOp>(ma::CmpIPredicate::eq, lane_id, zero);
+  auto thread_id = state.thread_and_block_ids[0];
+  Value cst_true = b.create<ma::ConstantOp>(b.getOneAttr(b.getI1Type()));
+
+  auto delinearized =
+      DelinearizeInBoundsIndex(mlir::getAffineDimExpr(0, ctx),
+                               threads_per_block, tiling.GetThreadStrides());
   auto thread_ids = mlir_converter::ApplyAffineMap(
-      mlir::AffineMap::get(
-          /*dimCount=*/1, /*symbolCount=*/0,
-          DelinearizeInBoundsIndex(mlir::getAffineDimExpr(0, ctx),
-                                   tiling.GetThreadsPerBlock(),
-                                   tiling.GetThreadStrides()),
-          ctx),
-      {thread_id}, {}, builder);
-  SmallVector<Value> thread_and_block_indices{thread_id, zero, zero,
-                                              block_id,  zero, zero};
-
-  auto warp_id = builder.create<mlir::arith::DivUIOp>(
+      mlir::AffineMap::get(1, 0, delinearized, ctx), {thread_id}, {}, b);
+
+  auto warp_id = b.create<ma::DivUIOp>(
       reduction_info().IsRowReduction()
           ? thread_ids[ReductionDimensions::kRowMinorReducedDimension]
           : thread_id,
-      builder.create<mlir::arith::ConstantIndexOp>(WarpSize()));
+      b.create<ma::ConstantIndexOp>(WarpSize()));
 
   std::vector<int64_t> shared_tile_size;
   SmallVector<Value> shared_write_indices;
@@ -204,64 +229,60 @@ absl::Status MlirReductionFusion::EmitReduction(EmitterState& state) const {
     auto kKept = ReductionDimensions::kRowKeptDimension;
     shared_tile_size = {tiling.GetThreadsPerBlock()[kKept], num_warps_row};
     shared_write_condition = is_first_lane;
-    shared_read_condition = builder.create<mlir::arith::CmpIOp>(
-        mlir::arith::CmpIPredicate::ult,
+    shared_read_condition = b.create<ma::CmpIOp>(
+        ma::CmpIPredicate::ult,
         thread_ids[ReductionDimensions::kRowMinorReducedDimension],
-        builder.create<mlir::arith::ConstantIndexOp>(num_warps_row));
+        b.create<ma::ConstantIndexOp>(num_warps_row));
     shared_write_indices = {thread_ids[kKept], warp_id};
     shared_read_indices = {thread_ids[kKept], lane_id};
   }
-  bool use_shared = !shared_tile_size.empty();
-
-  auto thread_has_output = mlir_converter::CheckConstraints(
-      *ComputeThreadIdToOutputIndexing(first_reduction_root_index_, ctx),
-      thread_and_block_indices, {}, builder);
-
-  HloValueMap inits;
-  llvm::SmallVector<Value> outputs =
-      mlir::ValueRange(state.entry_function.getArguments().drop_front(
-          state.fusion.fused_parameters().size()));
-  HloValueMap root_output_indices;
-  llvm::SmallVector<Value> epilogue_input_dims;
-  const auto& epilogue = state.computations.epilogues().front();
-  epilogue_input_dims = EmitThreadAndBlockIds(builder);
-  llvm::SmallVector<Value> epilogue_input_symbols(
-      epilogue.root_indexing.front().getNumSymbols(), zero);
-  for (auto [index, root] : llvm::enumerate(epilogue.roots)) {
-    root_output_indices[root] = mlir_converter::ApplyAffineMap(
-        epilogue.root_indexing[index], epilogue_input_dims,
-        epilogue_input_symbols, builder);
-  }
-
-  for (auto [index, hero] : llvm::enumerate(reduction_heroes_)) {
-    int arity = hero->operand_count() / 2;
-    const auto& computation =
-        state.computations.FindPartitionedComputation(hero->parent());
-    inits[hero] =
-        ProvideParameterRange(computation, hero, arity, arity, {},
-                              state.call_target, state.entry_function, builder);
-  }
 
   auto evaluate_epilogue = [&](const HloValueMap& results,
                                llvm::SmallVector<Value> outputs) {
-    auto epilogue_indices = epilogue_input_dims;
-    epilogue_indices.append(epilogue_input_symbols);
+    const auto& epilogue = state.computations.epilogues()[group_id];
+    if (epilogue.roots.empty()) return outputs;
+
+    llvm::SmallVector<Value> epilogue_input_symbols(
+        epilogue.root_indexing.front().getNumSymbols(), zero);
+    auto epilogue_input_indices = state.thread_and_block_ids;
+    epilogue_input_indices.append(epilogue_input_symbols);
     auto values =
-        EmitEpilogue(/*epilogue_index=*/0, state.computations,
-                     state.entry_function, results, epilogue_indices, builder);
-    const auto& epilogue = state.computations.epilogues().front();
-    for (auto root : epilogue.roots) {
+        EmitEpilogue(group_id, state.computations, state.entry_function,
+                     results, epilogue_input_indices, b);
+    int first_root_index = state.OutputIndex(epilogue.roots.front(), 0);
+    auto thread_has_output = mlir_converter::CheckConstraints(
+        *ComputeThreadIdToOutputIndexing(first_root_index, ctx),
+        state.thread_and_block_ids, {}, b);
+    for (auto [index, root] : llvm::enumerate(epilogue.roots)) {
+      auto output_indices = mlir_converter::ApplyAffineMap(
+          epilogue.root_indexing[index], state.thread_and_block_ids,
+          epilogue_input_symbols, b);
       for (auto [result_index, result] : llvm::enumerate(values.at(root))) {
         auto& output = outputs[state.OutputIndex(root, result_index)];
-        output = builder.create<PredicatedInsertOp>(
-            thread_has_output, result, output, root_output_indices[root]);
+        output = b.create<PredicatedInsertOp>(thread_has_output, result, output,
+                                              output_indices);
       }
     }
     return outputs;
   };
 
-  auto accumulated = state.EmitPerThreadReducedElements(inits);
-  for (auto root : side_output_roots_) {
+  HloValueMap inits;
+  const auto& reductions = reduction_heroes_[group_id];
+  for (auto* hero : reductions) {
+    int arity = hero->operand_count() / 2;
+    inits[hero] =
+        ProvideParameterRange(state.computation, hero, arity, arity, {},
+                              state.call_target, state.entry_function, b);
+  }
+  llvm::SmallVector<Value> outputs =
+      mlir::ValueRange(state.entry_function.getArguments().drop_front(
+          state.fusion.fused_parameters().size()));
+  for (auto* side_output : side_output_roots_[group_id]) {
+    inits[side_output].push_back(outputs[state.OutputIndex(side_output, 0)]);
+  }
+
+  auto accumulated = state.EmitPerThreadReducedElements(group_id, inits);
+  for (auto root : side_output_roots_[group_id]) {
     outputs[state.OutputIndex(root, 0)] = accumulated[root].front();
   }
 
@@ -269,40 +290,37 @@ absl::Status MlirReductionFusion::EmitReduction(EmitterState& state) const {
   // memory. In column reductions, the members of the warp process different
   // output elements, so we need to transpose first.
   if (reduction_info().IsRowReduction()) {
-    for (auto* hero : reduction_heroes_) {
-      auto reducer = state.GetReducer(hero);
+    for (auto* reduction : reductions) {
+      auto reducer = state.GetReducer(reduction);
       int max_dist = WarpSize() / 2 / reduction_info().GetRowsPerWarp();
-      accumulated[hero] =
-          builder.create<ShuffleReduceOp>(reducer, accumulated[hero], max_dist)
-              .getResults();
+      auto& values = accumulated[reduction];
+      values =
+          b.create<ShuffleReduceOp>(reducer, values, max_dist).getResults();
     }
   }
 
-  if (!use_shared) {
-    builder.create<mlir::func::ReturnOp>(
-        evaluate_epilogue(accumulated, std::move(outputs)));
-    return absl::OkStatus();
+  if (shared_tile_size.empty()) {
+    return evaluate_epilogue(accumulated, std::move(outputs));
   }
 
   SmallVector<Value> shared_tiles;
   // Write results to shared memory.
-  for (auto hero : reduction_heroes_) {
-    const auto& result = accumulated[hero];
+  for (auto* hero : reductions) {
     auto dest = state.AllocateSharedTiles(hero, shared_tile_size);
-    for (auto [value, output] : llvm::zip(result, dest)) {
-      shared_tiles.push_back(builder.create<PredicatedInsertOp>(
+    for (auto [value, output] : llvm::zip(accumulated[hero], dest)) {
+      shared_tiles.push_back(b.create<PredicatedInsertOp>(
           shared_write_condition, value, output, shared_write_indices));
     }
   }
 
   // Wait for the entire tile to be written.
   auto synced_tiles =
-      builder.create<SyncThreadsOp>(mlir::TypeRange(shared_tiles), shared_tiles)
+      b.create<SyncThreadsOp>(mlir::TypeRange(shared_tiles), shared_tiles)
           .getResults();
   auto write_outputs = [&](mlir::OpBuilder then_builder, mlir::Location loc) {
     mlir::ImplicitLocOpBuilder b(loc, then_builder);
     int tile_index = 0;
-    for (auto* hero : reduction_heroes_) {
+    for (auto* hero : reductions) {
       // Load from shared memory.
       SmallVector<Value> reduced;
       for (auto init : inits[hero]) {
@@ -312,112 +330,98 @@ absl::Status MlirReductionFusion::EmitReduction(EmitterState& state) const {
                                synced_tiles[tile_index++], shared_read_indices)
                               .getResult());
       }
-      accumulated[hero] = builder
-                              .create<ShuffleReduceOp>(state.GetReducer(hero),
-                                                       reduced, WarpSize() / 2)
-                              .getResults();
+      const auto& reducer = state.GetReducer(hero);
+      accumulated[hero] =
+          b.create<ShuffleReduceOp>(reducer, reduced, WarpSize() / 2)
+              .getResults();
     }
 
     b.create<mlir::scf::YieldOp>(loc, evaluate_epilogue(accumulated, outputs));
   };
 
-  auto warp_writes = reduction_info().IsRowReduction()
-                         ? builder.create<mlir::arith::CmpIOp>(
-                               mlir::arith::CmpIPredicate::eq, warp_id, zero)
-                         : cst_true;
-  auto written = builder.create<mlir::scf::IfOp>(
-      warp_writes, write_outputs, [&](mlir::OpBuilder b, mlir::Location loc) {
-        b.create<mlir::scf::YieldOp>(loc, outputs);
-      });
-  builder.create<mlir::func::ReturnOp>(written.getResults());
-
-  return absl::OkStatus();
+  auto warp_writes =
+      reduction_info().IsRowReduction()
+          ? b.create<ma::CmpIOp>(ma::CmpIPredicate::eq, warp_id, zero)
+          : cst_true;
+  auto yield_outputs = [&](mlir::OpBuilder else_builder, mlir::Location loc) {
+    else_builder.create<mlir::scf::YieldOp>(loc, outputs);
+  };
+  return b.create<mlir::scf::IfOp>(warp_writes, write_outputs, yield_outputs)
+      .getResults();
 }
 
 HloValueMap MlirReductionFusion::EmitterState::EmitPerThreadReducedElements(
-    const HloValueMap& inits) {
+    int group_id, const HloValueMap& inits) {
   const auto& tiling = owner.reduction_info().GetTiling();
   auto tile_indexing = GetIndexingMapForTiling(tiling, builder.getContext());
+  tile_indexing
+      .GetMutableDimensionBound(
+          KernelFusionInterface::kIndexingMapBlockIdxDims[1])
+      .upper = owner.reduction_heroes_.size();
 
   SmallVector<Value> iter_arg_inits;
-  ValueRange output_args = entry_function.getArguments().drop_front(
-      fusion.fused_parameters().size());
-  for (auto [is_reduction, hero] :
-       llvm::zip(owner.reduction_info().GetGroups().is_reduction_root,
-                 owner.analysis().fusion_heroes())) {
-    if (is_reduction) {
-      iter_arg_inits.append(inits.at(hero));
-    } else {
-      iter_arg_inits.push_back(output_args[OutputIndex(hero, 0)]);
-    }
+  const auto& side_outputs = owner.side_output_roots_[group_id];
+  const auto& reductions = owner.reduction_heroes_[group_id];
+  absl::flat_hash_map<const HloInstruction*, int> iter_arg_starts;
+  for (const auto& [hero, init] : inits) {
+    iter_arg_starts[hero] = iter_arg_inits.size();
+    iter_arg_inits.append(init);
   }
 
-  const auto& computation = computations.FindPartitionedComputation(
-      fusion.fused_instructions_computation());
-
   auto body_builder = [&](ValueRange iter_args, ValueRange dim_values,
                           ValueRange symbol_values) -> SmallVector<Value> {
     auto tile_indices = mlir_converter::ApplyAffineMap(
         tile_indexing.GetAffineMap(), dim_values, symbol_values, builder);
 
     llvm::SmallVector<Value> results(iter_args.size(), nullptr);
-    struct SideOutput {
-      llvm::SmallVector<Value> indices;
-      Value scalar;
-      int result_index;
-    };
-    llvm::SmallVector<SideOutput> side_outputs;
-    for (auto [is_reduction, hero, root] :
-         llvm::zip(owner.reduction_info().GetGroups().is_reduction_root,
-                   owner.analysis().fusion_heroes(),
-                   owner.analysis().fusion_roots())) {
-      const xla::Shape& input_shape =
+    auto get_input_indices = [&](auto* hero, bool is_reduction) {
+      const auto& input_shape =
           is_reduction ? hero->operand(0)->shape() : hero->shape();
-      auto input_indices = mlir_converter::ApplyAffineMap(
+      return mlir_converter::ApplyAffineMap(
           GetBitcastMap(tiling.GetXlaShape(), input_shape, builder.getContext())
               .GetAffineMap(),
           tile_indices, {}, builder);
-      int start = fusion_result_index_starts[root];
-      if (is_reduction) {
-        int num_outs = hero->operand_count() / 2;
-        auto values = ProvideParameterRange(
-            computations.FindPartitionedComputation(hero->parent()), hero, 0,
-            num_outs, input_indices, call_target, entry_function, builder);
-        SmallVector<Value> reduce_args = iter_args.slice(start, num_outs);
-        reduce_args.append(values);
-        absl::c_copy(builder.create<PureCallOp>(GetReducer(hero), reduce_args)
-                         .getResults(),
-                     results.begin() + start);
-      } else {
-        auto* root_tuple = fusion.fused_expression_root();
-        Value value = mlir_converter::ProvideParameter(
-            computation, root_tuple, root_tuple->operand_index(hero),
-            input_indices, call_target, entry_function, builder)[0];
-        // Tensor insertions turn into writes, so they have to happen in the
-        // end. This could be considered a bug in the lowering, but since we
-        // don't have bufferization, we need to handle it here.
-        side_outputs.push_back({std::move(input_indices), value, start});
-      }
+    };
+    for (auto* reduction : reductions) {
+      int arity = reduction->operand_count() / 2;
+      int start = iter_arg_starts[reduction];
+      SmallVector<Value> reduce_args = iter_args.slice(start, arity);
+      reduce_args.append(ProvideParameterRange(
+          computation, reduction, 0, arity, get_input_indices(reduction, true),
+          call_target, entry_function, builder));
+      const auto& reducer = GetReducer(reduction);
+      absl::c_copy(
+          builder.create<PureCallOp>(reducer, reduce_args).getResults(),
+          results.begin() + start);
+    }
+    struct SideOutput {
+      llvm::SmallVector<Value> indices;
+      Value scalar;
+    };
+    llvm::SmallVector<SideOutput> side_output_values;
+    for (auto* side_output : side_outputs) {
+      auto indices = get_input_indices(side_output, false);
+      auto* root_tuple = fusion.fused_expression_root();
+      Value value = mlir_converter::ProvideParameter(
+          computation, root_tuple, root_tuple->operand_index(side_output),
+          indices, call_target, entry_function, builder)[0];
+      side_output_values.push_back({std::move(indices), value});
     }
-    for (auto& side_output : side_outputs) {
-      int index = side_output.result_index;
-      results[index] = builder.create<mlir::tensor::InsertOp>(
-          side_output.scalar, iter_args[index], side_output.indices);
+    for (const auto& [side_output, values] :
+         llvm::zip(side_outputs, side_output_values)) {
+      int offset = iter_arg_starts[side_output];
+      results[offset] = builder.create<mlir::tensor::InsertOp>(
+          values.scalar, iter_args[offset], values.indices);
     }
     return results;
   };
 
-  auto results = owner.EmitThreadLoopNest(builder, iter_arg_inits,
-                                          tile_indexing, body_builder);
-  mlir::ValueRange result_range = results;
+  auto results_vector = owner.EmitThreadLoopNest(builder, iter_arg_inits,
+                                                 tile_indexing, body_builder);
+  mlir::ValueRange results = results_vector;
   HloValueMap results_per_hero;
-  for (auto [is_reduction, hero] :
-       llvm::zip(owner.reduction_info().GetGroups().is_reduction_root,
-                 owner.analysis().fusion_heroes())) {
-    int num_outs =
-        hero->shape().IsTuple() ? hero->shape().tuple_shapes_size() : 1;
-    results_per_hero[hero] = result_range.take_front(num_outs);
-    result_range = result_range.drop_front(num_outs);
+  for (const auto& [hero, init] : inits) {
+    results_per_hero[hero] = results.slice(iter_arg_starts[hero], init.size());
   }
   return results_per_hero;
 }
diff --git a/third_party/xla/xla/service/gpu/fusions/reduction_mlir.h b/third_party/xla/xla/service/gpu/fusions/reduction_mlir.h
index b17285a3bc68ae..12bc8124064245 100644
--- a/third_party/xla/xla/service/gpu/fusions/reduction_mlir.h
+++ b/third_party/xla/xla/service/gpu/fusions/reduction_mlir.h
@@ -54,13 +54,15 @@ class MlirReductionFusion
   struct EmitterState;
   friend struct EmitterState;
 
-  absl::Status EmitReduction(EmitterState& state) const;
+  llvm::SmallVector<mlir::Value> EmitReduction(int group_id,
+                                               EmitterState& state) const;
 
-  std::vector<const HloInstruction*> reduction_heroes_;
-  // The roots that have reduction heroes.
-  std::vector<const HloInstruction*> reduction_roots_;
-  std::vector<const HloInstruction*> side_output_roots_;
-  int first_reduction_root_index_;
+  // The reduction heroes for each reduction group.
+  std::vector<std::vector<const HloInstruction*>> reduction_heroes_;
+  // The roots that have reduction heroes for each reduction group.
+  std::vector<std::vector<const HloInstruction*>> reduction_roots_;
+  // The side output roots for each reduction group.
+  std::vector<std::vector<const HloInstruction*>> side_output_roots_;
 };
 
 }  // namespace gpu
diff --git a/third_party/xla/xla/service/gpu/fusions/reduction_mlir_test.cc b/third_party/xla/xla/service/gpu/fusions/reduction_mlir_test.cc
index 1fb1f2544bcda1..45b46267a2cfa7 100644
--- a/third_party/xla/xla/service/gpu/fusions/reduction_mlir_test.cc
+++ b/third_party/xla/xla/service/gpu/fusions/reduction_mlir_test.cc
@@ -153,6 +153,37 @@ TEST_F(ReductionTest, RowReduceMOFEpilogue) {
   EXPECT_TRUE(RunAndCompareNoHloPasses(kHloString, ErrorSpec{1e-3}));
 }
 
+TEST_F(ReductionTest, RowReduceMOFGroups) {
+  constexpr auto kHloString = R"(
+    %add_f32 {
+      %x = f32[] parameter(0)
+      %y = f32[] parameter(1)
+      ROOT %add = f32[] add(%x, %y)
+    }
+
+    %fused_computation {
+      %param0 = f32[1024] parameter(0)
+      %param1 = f32[1024] parameter(1)
+      %constant0 = f32[] constant(0)
+      %reduce1 = f32[] reduce(%param0, %constant0), dimensions={0}, to_apply=%add_f32
+      %reduce2 = f32[] reduce(%param1, %constant0), dimensions={0}, to_apply=%add_f32
+      ROOT %tuple = (f32[], f32[]) tuple(%reduce1, %reduce2)
+    }
+
+    ENTRY %cluster {
+      %param0 = f32[1024] parameter(0)
+      %param1 = f32[1024] parameter(1)
+      ROOT %fusion = (f32[], f32[])
+          fusion(%param0, %param1), kind=kInput, calls=%fused_computation
+    })";
+  TF_ASSERT_OK(EmitAndCheckIR(kHloString, R"(
+    // CHECK: scf.index_switch %block_id_y
+    // CHECK: case 1 {
+    // CHECK: default {
+  )"));
+  EXPECT_TRUE(RunAndCompareNoHloPasses(kHloString, ErrorSpec{1e-3}));
+}
+
 TEST_F(ReductionTest, ColumnReduction) {
   constexpr auto kHloString = R"(
     HloModule Test, is_scheduled=true

From 35f634a2cb446e6d15dd820ed22cf8b2aa3938d5 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Mon, 13 May 2024 03:18:13 -0700
Subject: [PATCH 011/478] Update ops-related pbtxt files.

PiperOrigin-RevId: 633149675
---
 tensorflow/core/ops/ops.pbtxt | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorflow/core/ops/ops.pbtxt b/tensorflow/core/ops/ops.pbtxt
index 2c64cab5e211e2..0556051577f5e4 100644
--- a/tensorflow/core/ops/ops.pbtxt
+++ b/tensorflow/core/ops/ops.pbtxt
@@ -1,4 +1,4 @@
-go/nodeserialize    
+go/debugstr  
 op {
   name: "Abort"
   attr {

From d6a65049b15ec86e6ef455f7268a9273cdbcd1fb Mon Sep 17 00:00:00 2001
From: Quentin Khan <qkhan@google.com>
Date: Mon, 13 May 2024 04:02:10 -0700
Subject: [PATCH 012/478] When no weight cache is provided to XNNPack, create
 one to share packed weights between operations.

PiperOrigin-RevId: 633158328
---
 tensorflow/lite/core/c/common.h               |   2 +
 tensorflow/lite/core/interpreter_builder.cc   |   3 +-
 tensorflow/lite/core/subgraph.cc              |   6 +-
 tensorflow/lite/core/subgraph.h               |  17 +-
 tensorflow/lite/delegates/xnnpack/BUILD       |  41 +
 .../lite/delegates/xnnpack/weight_cache.cc    | 491 ++++++++++++
 .../lite/delegates/xnnpack/weight_cache.h     | 306 ++++++++
 .../delegates/xnnpack/weight_cache_schema.fbs |  52 ++
 .../xnnpack/weight_cache_schema_generated.h   | 422 ++++++++++
 .../delegates/xnnpack/weight_cache_test.cc    | 725 ++++++++++++++++++
 .../delegates/xnnpack/xnnpack_delegate.cc     |  54 +-
 .../lite/delegates/xnnpack/xnnpack_delegate.h |   8 +
 tensorflow/lite/tflite_with_xnnpack.cc        |   4 +
 tensorflow/opensource_only.files              |   1 +
 14 files changed, 2124 insertions(+), 8 deletions(-)
 create mode 100644 tensorflow/lite/delegates/xnnpack/weight_cache.cc
 create mode 100644 tensorflow/lite/delegates/xnnpack/weight_cache.h
 create mode 100644 tensorflow/lite/delegates/xnnpack/weight_cache_schema.fbs
 create mode 100755 tensorflow/lite/delegates/xnnpack/weight_cache_schema_generated.h
 create mode 100644 tensorflow/lite/delegates/xnnpack/weight_cache_test.cc

diff --git a/tensorflow/lite/core/c/common.h b/tensorflow/lite/core/c/common.h
index ea54be9490ef01..96f19f12336bc4 100644
--- a/tensorflow/lite/core/c/common.h
+++ b/tensorflow/lite/core/c/common.h
@@ -472,6 +472,8 @@ typedef enum TfLiteCustomAllocationFlags {
   kTfLiteCustomAllocationFlagsSkipAlignCheck = 1,
 } TfLiteCustomAllocationFlags;
 
+enum { kTfLiteNoBufferIdentifier = SIZE_MAX };
+
 /// A tensor in the interpreter system which is a wrapper around a buffer of
 /// data including a dimensionality (or NULL if not currently defined).
 #ifndef TF_LITE_STATIC_MEMORY
diff --git a/tensorflow/lite/core/interpreter_builder.cc b/tensorflow/lite/core/interpreter_builder.cc
index 41e62cfd675340..d8c6d181ebdd1a 100644
--- a/tensorflow/lite/core/interpreter_builder.cc
+++ b/tensorflow/lite/core/interpreter_builder.cc
@@ -691,7 +691,8 @@ TfLiteStatus InterpreterBuilder::ParseTensors(
 
       if (subgraph->SetTensorParametersReadOnly(
               i, type, get_name(tensor), dims, quantization, buffer_ptr,
-              buffer_size, allocation_, sparsity) != kTfLiteOk) {
+              buffer_size, allocation_, sparsity,
+              /*buffer_identifier=*/tensor->buffer()) != kTfLiteOk) {
         TF_LITE_REPORT_ERROR(error_reporter_,
                              "Tensor %d is invalidly specified in schema.\n",
                              i);
diff --git a/tensorflow/lite/core/subgraph.cc b/tensorflow/lite/core/subgraph.cc
index 26ba2037342405..ce3622105b1ce5 100644
--- a/tensorflow/lite/core/subgraph.cc
+++ b/tensorflow/lite/core/subgraph.cc
@@ -1856,7 +1856,8 @@ TfLiteStatus Subgraph::GetNodeAndRegistration(
 TfLiteStatus Subgraph::SetTensorParametersReadOnly(
     int tensor_index, TfLiteType type, const char* name, const size_t ndims,
     const int* dims, TfLiteQuantization quantization, const char* buffer,
-    size_t bytes, const Allocation* allocation, TfLiteSparsity* sparsity) {
+    size_t bytes, const Allocation* allocation, TfLiteSparsity* sparsity,
+    const size_t buffer_identifier) {
   // Ensure quantization cleanup on failure.
   ScopedTfLiteQuantization scoped_quantization(&quantization);
   ScopedTfLiteSparsity scoped_sparsity(sparsity);
@@ -1904,6 +1905,9 @@ TfLiteStatus Subgraph::SetTensorParametersReadOnly(
     tensor.quantization = *scoped_quantization.release();
     tensor.sparsity = scoped_sparsity.release();
   }
+  if (buffer_identifier != kTfLiteNoBufferIdentifier) {
+    tensor_buffer_identifiers_[tensor_index] = buffer_identifier;
+  }
   return kTfLiteOk;
 }
 
diff --git a/tensorflow/lite/core/subgraph.h b/tensorflow/lite/core/subgraph.h
index 5940bfbb232ca3..281ac04adc2096 100644
--- a/tensorflow/lite/core/subgraph.h
+++ b/tensorflow/lite/core/subgraph.h
@@ -23,6 +23,7 @@ limitations under the License.
 #include <map>
 #include <memory>
 #include <string>
+#include <unordered_map>
 #include <unordered_set>
 #include <utility>
 #include <vector>
@@ -132,16 +133,18 @@ class Subgraph {
       int tensor_index, TfLiteType type, const char* name,
       const std::vector<int>& dims, TfLiteQuantization quantization,
       const char* buffer, size_t bytes, const Allocation* allocation = nullptr,
-      TfLiteSparsity* sparsity = nullptr) {
+      TfLiteSparsity* sparsity = nullptr,
+      size_t buffer_identifier = kTfLiteNoBufferIdentifier) {
     return SetTensorParametersReadOnly(tensor_index, type, name, dims.size(),
                                        dims.data(), quantization, buffer, bytes,
-                                       allocation, sparsity);
+                                       allocation, sparsity, buffer_identifier);
   }
   TfLiteStatus SetTensorParametersReadOnly(
       int tensor_index, TfLiteType type, const char* name, const size_t ndims,
       const int* dims, TfLiteQuantization quantization, const char* buffer,
       size_t bytes, const Allocation* allocation = nullptr,
-      TfLiteSparsity* sparsity = nullptr);
+      TfLiteSparsity* sparsity = nullptr,
+      size_t buffer_identifier = kTfLiteNoBufferIdentifier);
 
   // Set description of inputs/outputs/data/fptrs for node `node_index`.
   // This variant assumes an external buffer has been allocated of size
@@ -589,6 +592,10 @@ class Subgraph {
   // Returns true if the subgraph has been fully delegated.
   bool IsFullyDelegated() const;
 
+  const std::unordered_map<size_t, size_t>& GetTensorBufferIdentifiers() {
+    return tensor_buffer_identifiers_;
+  }
+
  private:
 #ifndef DOXYGEN_SKIP
   friend class tflite::impl::InterpreterBuilder;
@@ -1153,6 +1160,10 @@ class Subgraph {
   /// The allocator used for holding memory of the model. Note that this will
   /// be null if the client provides a tflite::Model directly.
   const Allocation* allocation_ = nullptr;
+
+  // Maps tensor constant buffers used in the subgraph to a model-wide
+  // identifiers.
+  std::unordered_map<size_t, size_t> tensor_buffer_identifiers_;
 };
 
 }  // namespace tflite
diff --git a/tensorflow/lite/delegates/xnnpack/BUILD b/tensorflow/lite/delegates/xnnpack/BUILD
index c4f748280d70ec..911a9358129a29 100644
--- a/tensorflow/lite/delegates/xnnpack/BUILD
+++ b/tensorflow/lite/delegates/xnnpack/BUILD
@@ -1,3 +1,4 @@
+load("@flatbuffers//:build_defs.bzl", "flatbuffer_cc_library")
 load("//tensorflow:tensorflow.default.bzl", "get_compatible_with_portable")
 load("//tensorflow/lite:build_def.bzl", "tflite_copts")
 load("//tensorflow/lite:special_rules.bzl", "internal_visibility_allowlist", "tflite_portable_test_suite_combined")
@@ -251,6 +252,7 @@ cc_library(
         ":tflite_with_xnnpack_qs8",
         ":tflite_with_xnnpack_qu8",
         ":tflite_with_xnnpack_transient_indirection_buffer",
+        ":weight_cache",
         "//tensorflow/lite:kernel_api",
         "//tensorflow/lite:minimal_logging",
         "//tensorflow/lite/c:c_api_types",
@@ -289,6 +291,7 @@ cc_library(
     linkstatic = True,
     deps = [
         ":quantization_util",
+        ":weight_cache",
         "//tensorflow/lite:kernel_api",
         "//tensorflow/lite:minimal_logging",
         "//tensorflow/lite/c:c_api_types",
@@ -323,6 +326,30 @@ cc_library(
     ],
 )
 
+flatbuffer_cc_library(
+    name = "weight_cache_schema",
+    srcs = ["weight_cache_schema.fbs"],
+    compatible_with = get_compatible_with_portable(),
+    flatc_args = [
+        "--gen-mutable",
+        "--gen-object-api",
+    ],
+)
+
+cc_library(
+    name = "weight_cache",
+    srcs = ["weight_cache.cc"],
+    hdrs = ["weight_cache.h"],
+    compatible_with = get_compatible_with_portable(),
+    deps = [
+        ":weight_cache_schema",
+        "//tensorflow/lite:minimal_logging",
+        "//tensorflow/lite/c:common",
+        "@XNNPACK",
+        "@flatbuffers//:runtime_cc",
+    ],
+)
+
 ################################ Tester classes ################################
 
 cc_library(
@@ -2828,4 +2855,18 @@ cc_test(
     ],
 )
 
+cc_test(
+    name = "weight_cache_test",
+    srcs = ["weight_cache_test.cc"],
+    deps = [
+        ":test_main",
+        ":weight_cache",
+        ":weight_cache_schema",
+        "//tensorflow/lite/c:common",
+        "@XNNPACK",
+        "@com_google_googletest//:gtest",
+        "@flatbuffers//:runtime_cc",
+    ],
+)
+
 tflite_portable_test_suite_combined(combine_conditions = {"deps": [":test_main"]})
diff --git a/tensorflow/lite/delegates/xnnpack/weight_cache.cc b/tensorflow/lite/delegates/xnnpack/weight_cache.cc
new file mode 100644
index 00000000000000..cb178662831839
--- /dev/null
+++ b/tensorflow/lite/delegates/xnnpack/weight_cache.cc
@@ -0,0 +1,491 @@
+/* Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+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.
+==============================================================================*/
+#include "tensorflow/lite/delegates/xnnpack/weight_cache.h"
+
+#include <fcntl.h>
+#include <sys/stat.h>
+
+#if defined(_MSC_VER)
+#include <io.h>
+#else
+#include <sys/mman.h>
+#include <unistd.h>
+#endif
+
+#include <algorithm>
+#include <cerrno>
+#include <cstddef>
+#include <cstdint>
+#include <cstdio>
+#include <cstdlib>
+#include <cstring>
+#include <memory>
+#include <string>
+#include <unordered_map>
+#include <utility>
+
+#include "xnnpack.h"  // from @XNNPACK
+#include "flatbuffers/base.h"  // from @flatbuffers
+#include "flatbuffers/flatbuffer_builder.h"  // from @flatbuffers
+#include "flatbuffers/verifier.h"  // from @flatbuffers
+#include "tensorflow/lite/c/common.h"
+#include "tensorflow/lite/delegates/xnnpack/weight_cache_schema_generated.h"
+#include "tensorflow/lite/logger.h"
+#include "tensorflow/lite/minimal_logging.h"
+
+#define XNNPACK_ABORT_CHECK(TEST, ...)                      \
+  if (!(TEST)) {                                            \
+    TFLITE_LOG_PROD(tflite::TFLITE_LOG_ERROR, __VA_ARGS__); \
+    std::abort();                                           \
+  }
+
+namespace tflite::xnnpack {
+
+namespace {
+constexpr size_t kMinAlignment = 64;
+
+template <class F>
+class ScopeGuard {
+ public:
+  explicit ScopeGuard(F&& callback) : callback_(std::forward<F>(callback)) {}
+  ~ScopeGuard() {
+    if (active_) {
+      callback_();
+    }
+  }
+
+  void Deactivate() { active_ = false; }
+
+ private:
+  F callback_;
+  bool active_ = true;
+};
+
+template <class F>
+ScopeGuard(F&&) -> ScopeGuard<F>;
+
+}  // namespace
+
+void swap(MMapHandle& a, MMapHandle& b) {
+  using std::swap;
+  swap(a.size_, b.size_);
+  swap(a.data_, b.data_);
+}
+
+MMapHandle::~MMapHandle() { UnMap(); }
+
+MMapHandle::MMapHandle(MMapHandle&& other) { swap(*this, other); }
+
+MMapHandle& MMapHandle::operator=(MMapHandle&& other) {
+  swap(*this, other);
+  return *this;
+}
+
+bool MMapHandle::Map(const char* path) {
+  this->UnMap();
+
+  const int fd = open(path, O_RDONLY);
+  if (fd == -1) {
+    TFLITE_LOG_PROD(tflite::TFLITE_LOG_ERROR,
+                    "Could not open file to mmap: %s (%s).", strerror(errno),
+                    path)
+    return false;
+  }
+
+  const ScopeGuard close_fd_on_return([&fd] {
+    if (fd >= 0) {
+      close(fd);
+    }
+  });
+
+  struct stat file_stats;
+  if (fstat(fd, &file_stats)) {
+    TFLITE_LOG_PROD(tflite::TFLITE_LOG_ERROR,
+                    "Could not access file stats to get size: %s (%s).",
+                    strerror(errno), path)
+    return false;
+  }
+
+  size_ = file_stats.st_size;
+#if defined(_MSC_VER)
+  data_ = new uint8_t[size_];
+  {
+    uint8_t* data_reader = data_;
+    size_t remaining_bytes = size_;
+    while (remaining_bytes > 0) {
+      const auto bytes = read(fd, data_reader, remaining_bytes);
+      if (bytes == -1) {
+        TFLITE_LOG_PROD(tflite::TFLITE_LOG_ERROR,
+                        "Could not read file ('%s'): %s.", path,
+                        strerror(errno))
+        UnMap();
+        return false;
+      }
+      remaining_bytes -= bytes;
+      data_reader += bytes;
+    }
+  }
+#else
+  data_ = static_cast<uint8_t*>(
+      mmap(/*addr=*/nullptr, size_, PROT_READ, MAP_SHARED, fd, /*offset=*/0));
+  if (data_ == MAP_FAILED) {
+    TFLITE_LOG_PROD(tflite::TFLITE_LOG_ERROR, "Could not mmap file: %s (%s).",
+                    strerror(errno), path)
+    data_ = nullptr;
+    size_ = 0;
+    return false;
+  }
+#endif
+
+  return true;
+}
+
+void MMapHandle::UnMap() {
+  if (data_) {
+#if defined(_MSC_VER)
+    delete[] data_;
+#else
+    munmap(data_, size_);
+#endif
+    data_ = nullptr;
+    size_ = 0;
+  }
+}
+
+void* WeightCacheBuilder::Reserve(size_t size) {
+  size_t offset = buffer_data_.size();
+  const size_t misalign = offset % kMinAlignment;
+  if (misalign) {
+    size += kMinAlignment - misalign;
+    offset += kMinAlignment - misalign;
+  }
+  buffer_data_.resize(buffer_data_.size() + size);
+  return buffer_data_.data() + offset;
+}
+
+bool WeightCacheBuilder::SpanIsWithinBuffer(const void* ptr,
+                                            uint64_t size) const {
+  const uintptr_t buf_begin = reinterpret_cast<uintptr_t>(buffer_data_.data());
+  const uintptr_t buf_end = buf_begin + buffer_data_.size();
+  const uintptr_t ptr_begin = reinterpret_cast<uintptr_t>(ptr);
+  const uintptr_t ptr_end = ptr_begin + size;
+  return ptr_begin >= buf_begin && ptr_begin <= buf_end &&
+         ptr_end >= buf_begin && ptr_end <= buf_end;
+}
+
+BufferLocation WeightCacheBuilder::Append(PackIdentifier pack_id,
+                                          const void* data, uint64_t size) {
+  const void* append_data = data;
+  if (!SpanIsWithinBuffer(data, size)) {
+    void* reserved_data = Reserve(size);
+    std::memcpy(reserved_data, data, size);
+    append_data = reserved_data;
+  }
+  BufferLocation loc{.offset = reinterpret_cast<uintptr_t>(append_data) -
+                               reinterpret_cast<uintptr_t>(buffer_data_.data()),
+                     .size = size};
+  schema_.buffers.push_back(std::make_unique<cache::schema::BufferT>(
+      cache::schema::BufferT{.packing_algorithm_id = pack_id.pack_algorithm_id,
+                             .weights_id = pack_id.weights_id,
+                             .bias_id = pack_id.bias_id,
+                             .offset = loc.offset,
+                             .size = loc.size}));
+  return loc;
+}
+
+bool WeightCacheBuilder::ShouldWrite() const { return !buffer_data_.empty(); }
+
+namespace {
+
+bool WriteData(const int fd, const uint8_t* data, size_t size,
+               const char* const file_path, const char* step_description) {
+  for (size_t bytes = 0; bytes < size;) {
+    const auto written_bytes = write(fd, data + bytes, size - bytes);
+    if (written_bytes == -1) {
+      TFLITE_LOG_PROD(tflite::TFLITE_LOG_ERROR,
+                      "Cache file write incomplete (%s). %s: %s", file_path,
+                      step_description, strerror(errno))
+    }
+    bytes += written_bytes;
+  }
+
+  return true;
+}
+
+}  // namespace
+
+bool WeightCacheBuilder::Write(const char* path) {
+  const int fd = open(path, O_WRONLY | O_CREAT | O_TRUNC, 0644);
+  if (fd == -1) {
+    TFLITE_LOG_PROD(tflite::TFLITE_LOG_ERROR,
+                    "Could not open cache file ('%s') for writing: %s", path,
+                    strerror(errno))
+    return false;
+  }
+
+  const ScopeGuard close_fd_on_return([&fd] {
+    if (fd >= 0) {
+      close(fd);
+    }
+  });
+
+  flatbuffers::FlatBufferBuilder builder;
+  // Add a fake size and the base offset to mutate them afterwards. Otherwise
+  // space for it won't be added to the flatbuffer.
+  schema_.flatbuffer_size = 1;
+  schema_.base_offset = 1;
+  FinishPackedWeightsBuffer(
+      builder, cache::schema::PackedWeights::Pack(builder, &schema_));
+
+  // Mutate the flatbuffer size and base offset fields.
+  auto* mutable_packed_weights =
+      cache::schema::GetMutablePackedWeights(builder.GetBufferPointer());
+  mutable_packed_weights->mutate_flatbuffer_size(builder.GetSize());
+  const size_t misalign = builder.GetSize() % kMinAlignment;
+  const size_t alignment_offset = misalign ? kMinAlignment - misalign : 0;
+  mutable_packed_weights->mutate_base_offset(builder.GetSize() +
+                                             alignment_offset);
+
+  // Write the flatbuffer which serves as a header to index the following data.
+  if (!WriteData(fd, builder.GetBufferPointer(), builder.GetSize(), path,
+                 "Header")) {
+    return false;
+  }
+  // Add some padding so that the cache file can be mmaped and the buffers
+  // stay aligned correctly.
+  const uint8_t fill[kMinAlignment] = {0};
+  if (!WriteData(fd, fill, alignment_offset, path, "Alignment padding")) {
+    return false;
+  }
+  // Write the actual buffer data.
+  if (!WriteData(fd, buffer_data_.data(), buffer_data_.size(), path,
+                 "Buffer data")) {
+    return false;
+  }
+  return true;
+}
+
+MMapWeightCacheProvider::MMapWeightCacheProvider(
+    MMapWeightCacheProvider&& other) {
+  *this = std::move(other);
+}
+
+MMapWeightCacheProvider& MMapWeightCacheProvider::operator=(
+    MMapWeightCacheProvider&& other) {
+  using std::swap;
+  swap(cache_provider_, other.cache_provider_);
+  // The contexts need to keep pointing to their owning object.
+  cache_provider_.context = this;
+  other.cache_provider_.context = &other;
+  swap(file_path_, other.file_path_);
+  swap(buffer_address_to_identifier_, other.buffer_address_to_identifier_);
+  swap(cache_key_to_offset_, other.cache_key_to_offset_);
+  swap(mmap_handle_, other.mmap_handle_);
+  swap(mmap_buffer_base_offset_, other.mmap_buffer_base_offset_);
+  swap(builder_, other.builder_);
+  return *this;
+}
+
+void MMapWeightCacheProvider::SetFilePath(const char* path) {
+  XNNPACK_ABORT_CHECK(
+      !IsFinalized(),
+      "Cannot change the path of a cache that has already been loaded.");
+  file_path_ = path;
+}
+
+bool MMapWeightCacheProvider::Load(const std::string& path) {
+  file_path_ = path;
+  if (mmap_handle_.Map(path.c_str())) {
+    return Load(std::move(mmap_handle_));
+  }
+  return false;
+}
+
+bool MMapWeightCacheProvider::Load(MMapHandle&& handle) {
+  swap(mmap_handle_, handle);
+  // Verifiy the flabuffer part of the file.
+  const size_t verifier_size =
+      std::min(mmap_handle_.size(),
+               static_cast<size_t>(FLATBUFFERS_MAX_BUFFER_SIZE - 1));
+  flatbuffers::Verifier verifier(mmap_handle_.data(), verifier_size);
+  if (!cache::schema::VerifyPackedWeightsBuffer(verifier)) {
+    TFLITE_LOG_PROD(tflite::TFLITE_LOG_ERROR,
+                    "Packed weights buffer validation failed.");
+    return false;
+  }
+
+  // Load flatbuffer.
+  const cache::schema::PackedWeights* packed_weights =
+      cache::schema::GetPackedWeights(mmap_handle_.data());
+  if (!packed_weights) {
+    TFLITE_LOG_PROD(tflite::TFLITE_LOG_ERROR,
+                    "Could not get packed weights from flatbuffer.");
+    return false;
+  }
+  mmap_buffer_base_offset_ = packed_weights->base_offset();
+  if (const auto buffers = packed_weights->buffers(); buffers) {
+    for (auto* buffer : *buffers) {
+      if (!buffer) {
+        TFLITE_LOG_PROD(tflite::TFLITE_LOG_ERROR,
+                        "Invalid buffer address in buffer list.");
+        return false;
+      }
+      cache_key_to_offset_.emplace(
+          PackIdentifier{.pack_algorithm_id = buffer->packing_algorithm_id(),
+                         .weights_id = buffer->weights_id(),
+                         .bias_id = buffer->bias_id()},
+          BufferLocation{.offset = buffer->offset(), .size = buffer->size()});
+    }
+  }
+  return true;
+}
+
+void MMapWeightCacheProvider::MapTensorIdentifiers(
+    const TfLiteTensor* tensors, const size_t size,
+    const std::unordered_map<size_t, size_t>& tensor_index_to_identifier) {
+  for (const auto [index, identifier] : tensor_index_to_identifier) {
+    XNNPACK_ABORT_CHECK(index < size,
+                        "Tensor index corresponds to a non existing tensor.");
+    buffer_address_to_identifier_[tensors[index].data.data] = identifier;
+  }
+}
+
+size_t MMapWeightCacheProvider::LookUp(
+    const xnn_weights_cache_look_up_key* cache_key) {
+  if (!cache_key) {
+    TFLITE_LOG_PROD(tflite::TFLITE_LOG_ERROR, "A null cache key was provided.");
+    return SIZE_MAX;
+  }
+  const PackIdentifier pack_id = BuildPackIdentifier(*cache_key);
+  if (auto offset_it = cache_key_to_offset_.find(pack_id);
+      offset_it != cache_key_to_offset_.end()) {
+    return offset_it->second.offset;
+  }
+  return SIZE_MAX;
+}
+
+void* MMapWeightCacheProvider::ReserveSpace(size_t size) {
+  XNNPACK_ABORT_CHECK(!IsFinalized(),
+                      "Cannot reserve space in a finalized cache.");
+  return builder_.Reserve(size);
+}
+
+size_t MMapWeightCacheProvider::LookUpOrInsert(
+    const xnn_weights_cache_look_up_key* cache_key, void* ptr, size_t size) {
+  XNNPACK_ABORT_CHECK(cache_key, "A null cache key was provided.");
+
+  const PackIdentifier pack_id = BuildPackIdentifier(*cache_key);
+  if (auto offset_it = cache_key_to_offset_.find(pack_id);
+      offset_it != cache_key_to_offset_.end()) {
+    return offset_it->second.offset;
+  }
+
+  XNNPACK_ABORT_CHECK(!IsFinalized(),
+                      "Cannot insert a buffer in a finalized cache.");
+
+  const BufferLocation location = builder_.Append(pack_id, ptr, size);
+  cache_key_to_offset_.emplace(pack_id, location);
+  return location.offset;
+}
+
+void* MMapWeightCacheProvider::OffsetToAddr(const size_t offset) {
+  // While the cache is being built, the buffer could grow and need to be
+  // reallocated so we cannot ensure pointer stability.
+  XNNPACK_ABORT_CHECK(
+      IsFinalized(),
+      "Cannot get the address of a buffer in a non finalized cache.");
+  return mmap_handle_.data() + mmap_buffer_base_offset_ + offset;
+}
+
+void MMapWeightCacheProvider::Reset() {
+  MMapWeightCacheProvider empty;
+  std::swap(*this, empty);
+}
+
+bool MMapWeightCacheProvider::Finalize() {
+  if (IsFinalized()) {
+    return true;
+  }
+  if (file_path_.empty()) {
+    TFLITE_LOG_PROD(tflite::TFLITE_LOG_ERROR,
+                    "File path wasn't set. Cannot finalize the cache.");
+    return false;
+  }
+  std::string file_path = file_path_;
+  if (!builder_.Write(file_path.c_str())) {
+    return false;
+  }
+  // The buffer mapping needs to be kept. We save it and restore it after the
+  // Reset.
+  std::unordered_map<const void*, uint64_t>
+      buffer_address_to_identifier_backup =
+          std::move(buffer_address_to_identifier_);
+  Reset();
+  buffer_address_to_identifier_ =
+      std::move(buffer_address_to_identifier_backup);
+  return Load(file_path);
+}
+
+bool MMapWeightCacheProvider::IsFinalized() const {
+  return mmap_handle_.IsMapped();
+}
+
+size_t MMapWeightCacheProvider::look_up(
+    void* context, const xnn_weights_cache_look_up_key* cache_key) {
+  return reinterpret_cast<MMapWeightCacheProvider*>(context)->LookUp(cache_key);
+}
+
+void* MMapWeightCacheProvider::reserve_space(void* context, size_t n) {
+  return reinterpret_cast<MMapWeightCacheProvider*>(context)->ReserveSpace(n);
+}
+
+size_t MMapWeightCacheProvider::look_up_or_insert(
+    void* context, const xnn_weights_cache_look_up_key* cache_key, void* ptr,
+    size_t size) {
+  return reinterpret_cast<MMapWeightCacheProvider*>(context)->LookUpOrInsert(
+      cache_key, ptr, size);
+}
+
+bool MMapWeightCacheProvider::is_finalized(void* context) {
+  return reinterpret_cast<MMapWeightCacheProvider*>(context)->IsFinalized();
+}
+
+void* MMapWeightCacheProvider::offset_to_addr(void* context, size_t offset) {
+  return reinterpret_cast<MMapWeightCacheProvider*>(context)->OffsetToAddr(
+      offset);
+}
+
+enum xnn_status MMapWeightCacheProvider::delete_cache(void* context) {
+  reinterpret_cast<MMapWeightCacheProvider*>(context)->Reset();
+  return xnn_status_success;
+}
+
+PackIdentifier MMapWeightCacheProvider::BuildPackIdentifier(
+    const xnn_weights_cache_look_up_key& key) {
+  const auto get_buffer_id = [&](const void* buffer) -> size_t {
+    if (buffer) {
+      const auto identifier_it = buffer_address_to_identifier_.find(buffer);
+      XNNPACK_ABORT_CHECK(identifier_it != buffer_address_to_identifier_.end(),
+                          "Unknown constant buffer passed to HashCacheKey.");
+      return identifier_it->second;
+    }
+    return PackIdentifier::kNoId;
+  };
+  return PackIdentifier{.pack_algorithm_id = key.seed,
+                        .weights_id = get_buffer_id(key.kernel),
+                        .bias_id = get_buffer_id(key.bias)};
+}
+
+}  // namespace tflite::xnnpack
diff --git a/tensorflow/lite/delegates/xnnpack/weight_cache.h b/tensorflow/lite/delegates/xnnpack/weight_cache.h
new file mode 100644
index 00000000000000..8942ea3beb3e9e
--- /dev/null
+++ b/tensorflow/lite/delegates/xnnpack/weight_cache.h
@@ -0,0 +1,306 @@
+/* Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+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.
+==============================================================================*/
+#ifndef TENSORFLOW_LITE_DELEGATES_XNNPACK_WEIGHT_CACHE_H_
+#define TENSORFLOW_LITE_DELEGATES_XNNPACK_WEIGHT_CACHE_H_
+
+#include <cstddef>
+#include <cstdint>
+#include <functional>
+#include <string>
+#include <unordered_map>
+#include <vector>
+
+#include "xnnpack.h"  // from @XNNPACK
+#include "tensorflow/lite/c/common.h"
+#include "tensorflow/lite/delegates/xnnpack/weight_cache_schema_generated.h"
+
+// WARNING: the interface in this file is still under experimentation and WILL
+// CHANGE. Do not rely on it.
+
+// TFLite doesn't use absl hashing utilities.
+
+namespace tflite {
+namespace xnnpack {
+
+struct PackIdentifier {
+  enum { kNoId = SIZE_MAX };
+  uint64_t pack_algorithm_id = kNoId;
+  uint64_t weights_id = kNoId;
+  uint64_t bias_id = kNoId;
+
+  friend bool operator==(const PackIdentifier& a, const PackIdentifier& b) {
+    return a.pack_algorithm_id == b.pack_algorithm_id &&
+           a.weights_id == b.weights_id && a.bias_id == b.bias_id;
+  }
+
+  struct Hash {
+    size_t operator()(const PackIdentifier& p) const {
+      std::hash<uint64_t> hasher;
+      return hasher(p.pack_algorithm_id) ^ hasher(p.weights_id) ^
+             hasher(p.bias_id);
+    }
+  };
+};
+
+struct BufferLocation {
+  uint64_t offset;
+  uint64_t size;
+};
+
+// Handles MMap allocations lifetime.
+//
+// When mapped, provides a view over the allocation for convenience.
+//
+// WARNING: the interface in this file is still under experimentation and WILL
+// CHANGE. Do not rely on it.
+class MMapHandle {
+ public:
+  using value_type = uint8_t;
+
+  MMapHandle() = default;
+  ~MMapHandle();
+  MMapHandle(const MMapHandle&) = delete;
+  MMapHandle& operator=(const MMapHandle&) = delete;
+  MMapHandle(MMapHandle&&);
+  MMapHandle& operator=(MMapHandle&&);
+
+  // Maps the file at the given path.
+  [[nodiscard /*Mapping a file can fail.*/]]
+  bool Map(const char* path);
+
+  // Unmaps an existing mapping.
+  void UnMap();
+
+  // Returns true if a mapping exists.
+  bool IsMapped() const { return data_ != nullptr; }
+
+  // Returns the mapping buffer.
+  uint8_t* data() { return data_; }
+
+  // Returns the mapping buffer.
+  const uint8_t* data() const { return data_; }
+
+  // Returns the mapping size in bytes.
+  size_t size() const { return size_; }
+
+  uint8_t* begin() { return data(); }
+
+  const uint8_t* begin() const { return data(); }
+
+  uint8_t* end() { return data() + size(); }
+
+  const uint8_t* end() const { return data() + size(); }
+
+  friend void swap(MMapHandle& a, MMapHandle& b);
+
+ private:
+  size_t size_ = 0;
+  uint8_t* data_ = nullptr;
+};
+
+// Provides storage to write the packed buffers to and saves those to disk.
+//
+// WARNING: the interface in this file is still under experimentation and WILL
+// CHANGE. Do not rely on it.
+class WeightCacheBuilder {
+ public:
+  // Reserves space in the data buffer for the required size in bytes and
+  // returns the address of that space.
+  //
+  // Sets `last_reserve` to the offset from `buffer_data_`'s start and `n`.
+  //
+  // A call to `Reserve` should alway be followed by a call to `Append`.
+  [[nodiscard /*The pointer to reserved space should be used.*/]]
+  void* Reserve(size_t size);
+
+  // Adds a buffer to the cache.
+  //
+  // The buffer space must have been reserved before using `Reserve`. If not, a
+  // new call to `Reserve` will be done and the data will be copied over.
+  [[nodiscard /*The location to the appended data should be saved.*/]]
+  BufferLocation Append(PackIdentifier pack_id, const void* data,
+                        uint64_t size);
+
+  // Checks whether this builder has data that needs to be written to disk.
+  bool ShouldWrite() const;
+
+  // Writes the flatbuffer to disk.
+  [[nodiscard /*Writing the weight cache can fail.*/]]
+  bool Write(const char* path);
+
+  // Helper for testing.
+  //
+  // WARNING: this exposes class implementation details for testing purposes and
+  // may be removed at any time.
+  const std::vector<uint8_t>& BufferData() const { return buffer_data_; }
+
+ private:
+  bool SpanIsWithinBuffer(const void* ptr, uint64_t size) const;
+
+  cache::schema::PackedWeightsT schema_;
+  std::vector<uint8_t> buffer_data_;
+};
+
+// Allows XNNPack to directly load packed weights from disk instead of having to
+// repack them every time.
+//
+// XNNPack kernels do not have knowledge of the TFLite context. The only thing
+// they can access is the buffers address. We rely on the fact that the address
+// provided by TFLite is unique in order to find out the buffer identifier.
+//
+// To use the cache you need to:
+//
+//  - Map the buffer addresses to their identifier with `MapTensorIdentifiers`
+//  - Load the cache file.
+//  - Finalize the cache before calling the run functions of XNNPack (setup and
+//    reshape are ok).
+class MMapWeightCacheProvider {
+ public:
+  MMapWeightCacheProvider() = default;
+  MMapWeightCacheProvider(const MMapWeightCacheProvider&) = delete;
+  MMapWeightCacheProvider& operator=(const MMapWeightCacheProvider&) = delete;
+  MMapWeightCacheProvider(MMapWeightCacheProvider&&);
+  MMapWeightCacheProvider& operator=(MMapWeightCacheProvider&&);
+
+  // Changes the file path to save the cache to.
+  //
+  // WARNING: Can only be called if the cache isn't finalized.
+  void SetFilePath(const char* file_path);
+
+  // Loads a flatbuffer following the layout in weight_cache_schema.fbs and set
+  // the file path.
+  [[nodiscard /*Loading a cache file may fail.*/]]
+  bool Load(const std::string& path);
+
+  // Loads an MMap allocation following the layout in weight_cache_schema.fbs.
+  [[nodiscard /*Loading cache data may fail.*/]]
+  bool Load(MMapHandle&& mmap_handle);
+
+  // Creates the tensor map.
+  void MapTensorIdentifiers(
+      const TfLiteTensor* tensors, size_t size,
+      const std::unordered_map<size_t, size_t>& tensor_index_to_identifier);
+
+  // Returns the offset of the buffer identified by `cache_key`.
+  //
+  // If the buffer isn't found, return SIZE_MAX.
+  [[nodiscard]]
+  size_t LookUp(const xnn_weights_cache_look_up_key* cache_key);
+
+  // Reserves space for a buffer of given size and returns a pointer to it.
+  //
+  // The buffer data should be filled and `LookUpOrInsert` should be immediately
+  // called.
+  [[nodiscard]]
+  void* ReserveSpace(size_t size);
+
+  // Returns the offset of the buffer identified by `cache_key`. If the lookup
+  // fails, inserts the span `[ptr, ptr+size)`.
+  //
+  // This should be called after ReserveSpace and `ptr` should be the result of
+  // that call with the given `size`.
+  //
+  // WARNING: The cache key cannot be null.
+  [[nodiscard]]
+  size_t LookUpOrInsert(const xnn_weights_cache_look_up_key* cache_key,
+                        void* ptr, size_t size);
+
+  // Gets the pointer to the buffer at the given offset.
+  //
+  // WARNING: This requires the buffer to be finalized.
+  // WARNING: This does not check the validity of the passed offset.
+  void* OffsetToAddr(size_t offset);
+
+  // Resets the weight cache provider as if it had been default constructed.
+  void Reset();
+
+  // Ensures that the cache is ready.
+  //
+  // If the cache file already exists, this is a no-op. Otherwise, this writes
+  // the file to disk and reloads it.
+  [[nodiscard /*Writing the cache file may fail.*/]]
+  bool Finalize();
+
+  // Checks whether the cache is ready to be used.
+  bool IsFinalized() const;
+
+  // Returns true if any weights have been added to the underlying builder.
+  bool IsBuilding() const { return !IsFinalized() && !file_path_.empty(); };
+
+  // Returns true if a file is mapped or a file path is set.
+  bool IsActive() const { return IsFinalized() || !file_path_.empty(); };
+
+  // Returns the cache provider expected by XNNPack.
+  xnn_weights_cache_provider& GetCacheProvider() { return cache_provider_; }
+
+  // C interface: `xnn_weights_cache_provider` callback.
+  static size_t look_up(void* context,
+                        const xnn_weights_cache_look_up_key* cache_key);
+
+  // C interface: `xnn_weights_cache_provider` callback.
+  static void* reserve_space(void* context, size_t n);
+
+  // C interface: `xnn_weights_cache_provider` callback.
+  static size_t look_up_or_insert(
+      void* context, const xnn_weights_cache_look_up_key* cache_key, void* ptr,
+      size_t size);
+
+  // C interface: `xnn_weights_cache_provider` callback.
+  static bool is_finalized(void* context);
+
+  // C interface: `xnn_weights_cache_provider` callback.
+  static void* offset_to_addr(void* context, size_t offset);
+
+  // C interface: `xnn_weights_cache_provider` callback.
+  static enum xnn_status delete_cache(void* context);
+
+ private:
+  // Hashes a cache key to lookup in `cache_key_to_identifier_`.
+  PackIdentifier BuildPackIdentifier(const xnn_weights_cache_look_up_key& key);
+
+  // Cache provider implementation for XNNPack.
+  xnn_weights_cache_provider cache_provider_{
+      .context = this,
+      .look_up = MMapWeightCacheProvider::look_up,
+      .reserve_space = MMapWeightCacheProvider::reserve_space,
+      .look_up_or_insert = MMapWeightCacheProvider::look_up_or_insert,
+      .is_finalized = MMapWeightCacheProvider::is_finalized,
+      .offset_to_addr = MMapWeightCacheProvider::offset_to_addr,
+      .delete_cache = MMapWeightCacheProvider::delete_cache};
+
+  // Path to the cache file.
+  std::string file_path_;
+
+  // Maps buffer addresses to buffer identifiers.
+  std::unordered_map<const void*, uint64_t> buffer_address_to_identifier_;
+
+  // Maps cache request hashes to the buffer identifier.
+  std::unordered_multimap<PackIdentifier, BufferLocation, PackIdentifier::Hash>
+      cache_key_to_offset_;
+
+  // MMap allocation handler.
+  MMapHandle mmap_handle_;
+
+  // The offset to the first buffer data in the MMap allocation.
+  size_t mmap_buffer_base_offset_;
+
+  // Used to build the cache.
+  WeightCacheBuilder builder_;
+};
+
+}  // namespace xnnpack
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_DELEGATES_XNNPACK_WEIGHT_CACHE_H_
diff --git a/tensorflow/lite/delegates/xnnpack/weight_cache_schema.fbs b/tensorflow/lite/delegates/xnnpack/weight_cache_schema.fbs
new file mode 100644
index 00000000000000..0658054f21c07e
--- /dev/null
+++ b/tensorflow/lite/delegates/xnnpack/weight_cache_schema.fbs
@@ -0,0 +1,52 @@
+// Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// 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.
+
+// This is a list of buffers with identifiers, to host the CPU-specific cache on disk.
+namespace tflite.xnnpack.cache.schema;
+
+// Schema version.
+file_identifier "V001";
+// File extension of written files.
+file_extension "xnn_weights";
+
+table Buffer {
+  // To uniquely identify a packed buffer we need to keep track of the packing
+  // algorithm and of the buffers that were used to generate it.
+  packing_algorithm_id: uint64;
+  weights_id: uint64;
+  bias_id: uint64;
+
+  /// The buffer data is appended after the flatbuffer to bypass 2GB file size
+  /// limitation. The offset is calculated relative to the base offset.
+  /// (i.e. beginning of the file + base_offset).
+  offset: uint64;
+
+  /// Size of the buffer in bytes.
+  size: uint64;
+}
+
+table PackedWeights {
+  /// A list of buffers.
+  buffers: [Buffer];
+
+  /// The serialized file is `flatbuffer_size` of bytes representing
+  /// `NamedBuffers` appended with a blob representing the buffer content.
+  flatbuffer_size: uint64;
+
+  /// Defines the base offset for the data appended to the file. That offset
+  /// may be needed to guarantee data alignment.
+  base_offset:uint64;
+}
+
+root_type PackedWeights;
diff --git a/tensorflow/lite/delegates/xnnpack/weight_cache_schema_generated.h b/tensorflow/lite/delegates/xnnpack/weight_cache_schema_generated.h
new file mode 100755
index 00000000000000..fa5d30a4cdae65
--- /dev/null
+++ b/tensorflow/lite/delegates/xnnpack/weight_cache_schema_generated.h
@@ -0,0 +1,422 @@
+/* Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+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.
+==============================================================================*/
+// automatically generated by the FlatBuffers compiler, do not modify
+
+
+#ifndef FLATBUFFERS_GENERATED_WEIGHTCACHESCHEMA_TFLITE_XNNPACK_CACHE_SCHEMA_H_
+#define FLATBUFFERS_GENERATED_WEIGHTCACHESCHEMA_TFLITE_XNNPACK_CACHE_SCHEMA_H_
+
+#include "flatbuffers/flatbuffers.h"
+
+// Ensure the included flatbuffers.h is the same version as when this file was
+// generated, otherwise it may not be compatible.
+static_assert(FLATBUFFERS_VERSION_MAJOR == 24 &&
+              FLATBUFFERS_VERSION_MINOR == 3 &&
+              FLATBUFFERS_VERSION_REVISION == 25,
+             "Non-compatible flatbuffers version included");
+
+namespace tflite {
+namespace xnnpack {
+namespace cache {
+namespace schema {
+
+struct Buffer;
+struct BufferBuilder;
+struct BufferT;
+
+struct PackedWeights;
+struct PackedWeightsBuilder;
+struct PackedWeightsT;
+
+struct BufferT : public ::flatbuffers::NativeTable {
+  typedef Buffer TableType;
+  uint64_t packing_algorithm_id = 0;
+  uint64_t weights_id = 0;
+  uint64_t bias_id = 0;
+  uint64_t offset = 0;
+  uint64_t size = 0;
+};
+
+struct Buffer FLATBUFFERS_FINAL_CLASS : private ::flatbuffers::Table {
+  typedef BufferT NativeTableType;
+  typedef BufferBuilder Builder;
+  enum FlatBuffersVTableOffset FLATBUFFERS_VTABLE_UNDERLYING_TYPE {
+    VT_PACKING_ALGORITHM_ID = 4,
+    VT_WEIGHTS_ID = 6,
+    VT_BIAS_ID = 8,
+    VT_OFFSET = 10,
+    VT_SIZE = 12
+  };
+  uint64_t packing_algorithm_id() const {
+    return GetField<uint64_t>(VT_PACKING_ALGORITHM_ID, 0);
+  }
+  bool mutate_packing_algorithm_id(uint64_t _packing_algorithm_id = 0) {
+    return SetField<uint64_t>(VT_PACKING_ALGORITHM_ID, _packing_algorithm_id, 0);
+  }
+  uint64_t weights_id() const {
+    return GetField<uint64_t>(VT_WEIGHTS_ID, 0);
+  }
+  bool mutate_weights_id(uint64_t _weights_id = 0) {
+    return SetField<uint64_t>(VT_WEIGHTS_ID, _weights_id, 0);
+  }
+  uint64_t bias_id() const {
+    return GetField<uint64_t>(VT_BIAS_ID, 0);
+  }
+  bool mutate_bias_id(uint64_t _bias_id = 0) {
+    return SetField<uint64_t>(VT_BIAS_ID, _bias_id, 0);
+  }
+  /// The buffer data is appended after the flatbuffer to bypass 2GB file size
+  /// limitation. The offset is calculated relative to the base offset.
+  /// (i.e. beginning of the file + base_offset).
+  uint64_t offset() const {
+    return GetField<uint64_t>(VT_OFFSET, 0);
+  }
+  bool mutate_offset(uint64_t _offset = 0) {
+    return SetField<uint64_t>(VT_OFFSET, _offset, 0);
+  }
+  /// Size of the buffer in bytes.
+  uint64_t size() const {
+    return GetField<uint64_t>(VT_SIZE, 0);
+  }
+  bool mutate_size(uint64_t _size = 0) {
+    return SetField<uint64_t>(VT_SIZE, _size, 0);
+  }
+  bool Verify(::flatbuffers::Verifier &verifier) const {
+    return VerifyTableStart(verifier) &&
+           VerifyField<uint64_t>(verifier, VT_PACKING_ALGORITHM_ID, 8) &&
+           VerifyField<uint64_t>(verifier, VT_WEIGHTS_ID, 8) &&
+           VerifyField<uint64_t>(verifier, VT_BIAS_ID, 8) &&
+           VerifyField<uint64_t>(verifier, VT_OFFSET, 8) &&
+           VerifyField<uint64_t>(verifier, VT_SIZE, 8) &&
+           verifier.EndTable();
+  }
+  BufferT *UnPack(const ::flatbuffers::resolver_function_t *_resolver = nullptr) const;
+  void UnPackTo(BufferT *_o, const ::flatbuffers::resolver_function_t *_resolver = nullptr) const;
+  static ::flatbuffers::Offset<Buffer> Pack(::flatbuffers::FlatBufferBuilder &_fbb, const BufferT* _o, const ::flatbuffers::rehasher_function_t *_rehasher = nullptr);
+};
+
+struct BufferBuilder {
+  typedef Buffer Table;
+  ::flatbuffers::FlatBufferBuilder &fbb_;
+  ::flatbuffers::uoffset_t start_;
+  void add_packing_algorithm_id(uint64_t packing_algorithm_id) {
+    fbb_.AddElement<uint64_t>(Buffer::VT_PACKING_ALGORITHM_ID, packing_algorithm_id, 0);
+  }
+  void add_weights_id(uint64_t weights_id) {
+    fbb_.AddElement<uint64_t>(Buffer::VT_WEIGHTS_ID, weights_id, 0);
+  }
+  void add_bias_id(uint64_t bias_id) {
+    fbb_.AddElement<uint64_t>(Buffer::VT_BIAS_ID, bias_id, 0);
+  }
+  void add_offset(uint64_t offset) {
+    fbb_.AddElement<uint64_t>(Buffer::VT_OFFSET, offset, 0);
+  }
+  void add_size(uint64_t size) {
+    fbb_.AddElement<uint64_t>(Buffer::VT_SIZE, size, 0);
+  }
+  explicit BufferBuilder(::flatbuffers::FlatBufferBuilder &_fbb)
+        : fbb_(_fbb) {
+    start_ = fbb_.StartTable();
+  }
+  ::flatbuffers::Offset<Buffer> Finish() {
+    const auto end = fbb_.EndTable(start_);
+    auto o = ::flatbuffers::Offset<Buffer>(end);
+    return o;
+  }
+};
+
+inline ::flatbuffers::Offset<Buffer> CreateBuffer(
+    ::flatbuffers::FlatBufferBuilder &_fbb,
+    uint64_t packing_algorithm_id = 0,
+    uint64_t weights_id = 0,
+    uint64_t bias_id = 0,
+    uint64_t offset = 0,
+    uint64_t size = 0) {
+  BufferBuilder builder_(_fbb);
+  builder_.add_size(size);
+  builder_.add_offset(offset);
+  builder_.add_bias_id(bias_id);
+  builder_.add_weights_id(weights_id);
+  builder_.add_packing_algorithm_id(packing_algorithm_id);
+  return builder_.Finish();
+}
+
+::flatbuffers::Offset<Buffer> CreateBuffer(::flatbuffers::FlatBufferBuilder &_fbb, const BufferT *_o, const ::flatbuffers::rehasher_function_t *_rehasher = nullptr);
+
+struct PackedWeightsT : public ::flatbuffers::NativeTable {
+  typedef PackedWeights TableType;
+  std::vector<std::unique_ptr<tflite::xnnpack::cache::schema::BufferT>> buffers{};
+  uint64_t flatbuffer_size = 0;
+  uint64_t base_offset = 0;
+  PackedWeightsT() = default;
+  PackedWeightsT(const PackedWeightsT &o);
+  PackedWeightsT(PackedWeightsT&&) FLATBUFFERS_NOEXCEPT = default;
+  PackedWeightsT &operator=(PackedWeightsT o) FLATBUFFERS_NOEXCEPT;
+};
+
+struct PackedWeights FLATBUFFERS_FINAL_CLASS : private ::flatbuffers::Table {
+  typedef PackedWeightsT NativeTableType;
+  typedef PackedWeightsBuilder Builder;
+  enum FlatBuffersVTableOffset FLATBUFFERS_VTABLE_UNDERLYING_TYPE {
+    VT_BUFFERS = 4,
+    VT_FLATBUFFER_SIZE = 6,
+    VT_BASE_OFFSET = 8
+  };
+  /// A list of buffers.
+  const ::flatbuffers::Vector<::flatbuffers::Offset<tflite::xnnpack::cache::schema::Buffer>> *buffers() const {
+    return GetPointer<const ::flatbuffers::Vector<::flatbuffers::Offset<tflite::xnnpack::cache::schema::Buffer>> *>(VT_BUFFERS);
+  }
+  ::flatbuffers::Vector<::flatbuffers::Offset<tflite::xnnpack::cache::schema::Buffer>> *mutable_buffers() {
+    return GetPointer<::flatbuffers::Vector<::flatbuffers::Offset<tflite::xnnpack::cache::schema::Buffer>> *>(VT_BUFFERS);
+  }
+  /// The serialized file is `flatbuffer_size` of bytes representing
+  /// `NamedBuffers` appended with a blob representing the buffer content.
+  uint64_t flatbuffer_size() const {
+    return GetField<uint64_t>(VT_FLATBUFFER_SIZE, 0);
+  }
+  bool mutate_flatbuffer_size(uint64_t _flatbuffer_size = 0) {
+    return SetField<uint64_t>(VT_FLATBUFFER_SIZE, _flatbuffer_size, 0);
+  }
+  /// Defines the base offset for the data appended to the file. That offset
+  /// may be needed to guarantee data alignment.
+  uint64_t base_offset() const {
+    return GetField<uint64_t>(VT_BASE_OFFSET, 0);
+  }
+  bool mutate_base_offset(uint64_t _base_offset = 0) {
+    return SetField<uint64_t>(VT_BASE_OFFSET, _base_offset, 0);
+  }
+  bool Verify(::flatbuffers::Verifier &verifier) const {
+    return VerifyTableStart(verifier) &&
+           VerifyOffset(verifier, VT_BUFFERS) &&
+           verifier.VerifyVector(buffers()) &&
+           verifier.VerifyVectorOfTables(buffers()) &&
+           VerifyField<uint64_t>(verifier, VT_FLATBUFFER_SIZE, 8) &&
+           VerifyField<uint64_t>(verifier, VT_BASE_OFFSET, 8) &&
+           verifier.EndTable();
+  }
+  PackedWeightsT *UnPack(const ::flatbuffers::resolver_function_t *_resolver = nullptr) const;
+  void UnPackTo(PackedWeightsT *_o, const ::flatbuffers::resolver_function_t *_resolver = nullptr) const;
+  static ::flatbuffers::Offset<PackedWeights> Pack(::flatbuffers::FlatBufferBuilder &_fbb, const PackedWeightsT* _o, const ::flatbuffers::rehasher_function_t *_rehasher = nullptr);
+};
+
+struct PackedWeightsBuilder {
+  typedef PackedWeights Table;
+  ::flatbuffers::FlatBufferBuilder &fbb_;
+  ::flatbuffers::uoffset_t start_;
+  void add_buffers(::flatbuffers::Offset<::flatbuffers::Vector<::flatbuffers::Offset<tflite::xnnpack::cache::schema::Buffer>>> buffers) {
+    fbb_.AddOffset(PackedWeights::VT_BUFFERS, buffers);
+  }
+  void add_flatbuffer_size(uint64_t flatbuffer_size) {
+    fbb_.AddElement<uint64_t>(PackedWeights::VT_FLATBUFFER_SIZE, flatbuffer_size, 0);
+  }
+  void add_base_offset(uint64_t base_offset) {
+    fbb_.AddElement<uint64_t>(PackedWeights::VT_BASE_OFFSET, base_offset, 0);
+  }
+  explicit PackedWeightsBuilder(::flatbuffers::FlatBufferBuilder &_fbb)
+        : fbb_(_fbb) {
+    start_ = fbb_.StartTable();
+  }
+  ::flatbuffers::Offset<PackedWeights> Finish() {
+    const auto end = fbb_.EndTable(start_);
+    auto o = ::flatbuffers::Offset<PackedWeights>(end);
+    return o;
+  }
+};
+
+inline ::flatbuffers::Offset<PackedWeights> CreatePackedWeights(
+    ::flatbuffers::FlatBufferBuilder &_fbb,
+    ::flatbuffers::Offset<::flatbuffers::Vector<::flatbuffers::Offset<tflite::xnnpack::cache::schema::Buffer>>> buffers = 0,
+    uint64_t flatbuffer_size = 0,
+    uint64_t base_offset = 0) {
+  PackedWeightsBuilder builder_(_fbb);
+  builder_.add_base_offset(base_offset);
+  builder_.add_flatbuffer_size(flatbuffer_size);
+  builder_.add_buffers(buffers);
+  return builder_.Finish();
+}
+
+inline ::flatbuffers::Offset<PackedWeights> CreatePackedWeightsDirect(
+    ::flatbuffers::FlatBufferBuilder &_fbb,
+    const std::vector<::flatbuffers::Offset<tflite::xnnpack::cache::schema::Buffer>> *buffers = nullptr,
+    uint64_t flatbuffer_size = 0,
+    uint64_t base_offset = 0) {
+  auto buffers__ = buffers ? _fbb.CreateVector<::flatbuffers::Offset<tflite::xnnpack::cache::schema::Buffer>>(*buffers) : 0;
+  return tflite::xnnpack::cache::schema::CreatePackedWeights(
+      _fbb,
+      buffers__,
+      flatbuffer_size,
+      base_offset);
+}
+
+::flatbuffers::Offset<PackedWeights> CreatePackedWeights(::flatbuffers::FlatBufferBuilder &_fbb, const PackedWeightsT *_o, const ::flatbuffers::rehasher_function_t *_rehasher = nullptr);
+
+inline BufferT *Buffer::UnPack(const ::flatbuffers::resolver_function_t *_resolver) const {
+  auto _o = std::unique_ptr<BufferT>(new BufferT());
+  UnPackTo(_o.get(), _resolver);
+  return _o.release();
+}
+
+inline void Buffer::UnPackTo(BufferT *_o, const ::flatbuffers::resolver_function_t *_resolver) const {
+  (void)_o;
+  (void)_resolver;
+  { auto _e = packing_algorithm_id(); _o->packing_algorithm_id = _e; }
+  { auto _e = weights_id(); _o->weights_id = _e; }
+  { auto _e = bias_id(); _o->bias_id = _e; }
+  { auto _e = offset(); _o->offset = _e; }
+  { auto _e = size(); _o->size = _e; }
+}
+
+inline ::flatbuffers::Offset<Buffer> Buffer::Pack(::flatbuffers::FlatBufferBuilder &_fbb, const BufferT* _o, const ::flatbuffers::rehasher_function_t *_rehasher) {
+  return CreateBuffer(_fbb, _o, _rehasher);
+}
+
+inline ::flatbuffers::Offset<Buffer> CreateBuffer(::flatbuffers::FlatBufferBuilder &_fbb, const BufferT *_o, const ::flatbuffers::rehasher_function_t *_rehasher) {
+  (void)_rehasher;
+  (void)_o;
+  struct _VectorArgs { ::flatbuffers::FlatBufferBuilder *__fbb; const BufferT* __o; const ::flatbuffers::rehasher_function_t *__rehasher; } _va = { &_fbb, _o, _rehasher}; (void)_va;
+  auto _packing_algorithm_id = _o->packing_algorithm_id;
+  auto _weights_id = _o->weights_id;
+  auto _bias_id = _o->bias_id;
+  auto _offset = _o->offset;
+  auto _size = _o->size;
+  return tflite::xnnpack::cache::schema::CreateBuffer(
+      _fbb,
+      _packing_algorithm_id,
+      _weights_id,
+      _bias_id,
+      _offset,
+      _size);
+}
+
+inline PackedWeightsT::PackedWeightsT(const PackedWeightsT &o)
+      : flatbuffer_size(o.flatbuffer_size),
+        base_offset(o.base_offset) {
+  buffers.reserve(o.buffers.size());
+  for (const auto &buffers_ : o.buffers) { buffers.emplace_back((buffers_) ? new tflite::xnnpack::cache::schema::BufferT(*buffers_) : nullptr); }
+}
+
+inline PackedWeightsT &PackedWeightsT::operator=(PackedWeightsT o) FLATBUFFERS_NOEXCEPT {
+  std::swap(buffers, o.buffers);
+  std::swap(flatbuffer_size, o.flatbuffer_size);
+  std::swap(base_offset, o.base_offset);
+  return *this;
+}
+
+inline PackedWeightsT *PackedWeights::UnPack(const ::flatbuffers::resolver_function_t *_resolver) const {
+  auto _o = std::unique_ptr<PackedWeightsT>(new PackedWeightsT());
+  UnPackTo(_o.get(), _resolver);
+  return _o.release();
+}
+
+inline void PackedWeights::UnPackTo(PackedWeightsT *_o, const ::flatbuffers::resolver_function_t *_resolver) const {
+  (void)_o;
+  (void)_resolver;
+  { auto _e = buffers(); if (_e) { _o->buffers.resize(_e->size()); for (::flatbuffers::uoffset_t _i = 0; _i < _e->size(); _i++) { if(_o->buffers[_i]) { _e->Get(_i)->UnPackTo(_o->buffers[_i].get(), _resolver); } else { _o->buffers[_i] = std::unique_ptr<tflite::xnnpack::cache::schema::BufferT>(_e->Get(_i)->UnPack(_resolver)); }; } } else { _o->buffers.resize(0); } }
+  { auto _e = flatbuffer_size(); _o->flatbuffer_size = _e; }
+  { auto _e = base_offset(); _o->base_offset = _e; }
+}
+
+inline ::flatbuffers::Offset<PackedWeights> PackedWeights::Pack(::flatbuffers::FlatBufferBuilder &_fbb, const PackedWeightsT* _o, const ::flatbuffers::rehasher_function_t *_rehasher) {
+  return CreatePackedWeights(_fbb, _o, _rehasher);
+}
+
+inline ::flatbuffers::Offset<PackedWeights> CreatePackedWeights(::flatbuffers::FlatBufferBuilder &_fbb, const PackedWeightsT *_o, const ::flatbuffers::rehasher_function_t *_rehasher) {
+  (void)_rehasher;
+  (void)_o;
+  struct _VectorArgs { ::flatbuffers::FlatBufferBuilder *__fbb; const PackedWeightsT* __o; const ::flatbuffers::rehasher_function_t *__rehasher; } _va = { &_fbb, _o, _rehasher}; (void)_va;
+  auto _buffers = _o->buffers.size() ? _fbb.CreateVector<::flatbuffers::Offset<tflite::xnnpack::cache::schema::Buffer>> (_o->buffers.size(), [](size_t i, _VectorArgs *__va) { return CreateBuffer(*__va->__fbb, __va->__o->buffers[i].get(), __va->__rehasher); }, &_va ) : 0;
+  auto _flatbuffer_size = _o->flatbuffer_size;
+  auto _base_offset = _o->base_offset;
+  return tflite::xnnpack::cache::schema::CreatePackedWeights(
+      _fbb,
+      _buffers,
+      _flatbuffer_size,
+      _base_offset);
+}
+
+inline const tflite::xnnpack::cache::schema::PackedWeights *GetPackedWeights(const void *buf) {
+  return ::flatbuffers::GetRoot<tflite::xnnpack::cache::schema::PackedWeights>(buf);
+}
+
+inline const tflite::xnnpack::cache::schema::PackedWeights *GetSizePrefixedPackedWeights(const void *buf) {
+  return ::flatbuffers::GetSizePrefixedRoot<tflite::xnnpack::cache::schema::PackedWeights>(buf);
+}
+
+inline PackedWeights *GetMutablePackedWeights(void *buf) {
+  return ::flatbuffers::GetMutableRoot<PackedWeights>(buf);
+}
+
+inline tflite::xnnpack::cache::schema::PackedWeights *GetMutableSizePrefixedPackedWeights(void *buf) {
+  return ::flatbuffers::GetMutableSizePrefixedRoot<tflite::xnnpack::cache::schema::PackedWeights>(buf);
+}
+
+inline const char *PackedWeightsIdentifier() {
+  return "V001";
+}
+
+inline bool PackedWeightsBufferHasIdentifier(const void *buf) {
+  return ::flatbuffers::BufferHasIdentifier(
+      buf, PackedWeightsIdentifier());
+}
+
+inline bool SizePrefixedPackedWeightsBufferHasIdentifier(const void *buf) {
+  return ::flatbuffers::BufferHasIdentifier(
+      buf, PackedWeightsIdentifier(), true);
+}
+
+inline bool VerifyPackedWeightsBuffer(
+    ::flatbuffers::Verifier &verifier) {
+  return verifier.VerifyBuffer<tflite::xnnpack::cache::schema::PackedWeights>(PackedWeightsIdentifier());
+}
+
+inline bool VerifySizePrefixedPackedWeightsBuffer(
+    ::flatbuffers::Verifier &verifier) {
+  return verifier.VerifySizePrefixedBuffer<tflite::xnnpack::cache::schema::PackedWeights>(PackedWeightsIdentifier());
+}
+
+inline const char *PackedWeightsExtension() {
+  return "xnn_weights";
+}
+
+inline void FinishPackedWeightsBuffer(
+    ::flatbuffers::FlatBufferBuilder &fbb,
+    ::flatbuffers::Offset<tflite::xnnpack::cache::schema::PackedWeights> root) {
+  fbb.Finish(root, PackedWeightsIdentifier());
+}
+
+inline void FinishSizePrefixedPackedWeightsBuffer(
+    ::flatbuffers::FlatBufferBuilder &fbb,
+    ::flatbuffers::Offset<tflite::xnnpack::cache::schema::PackedWeights> root) {
+  fbb.FinishSizePrefixed(root, PackedWeightsIdentifier());
+}
+
+inline std::unique_ptr<tflite::xnnpack::cache::schema::PackedWeightsT> UnPackPackedWeights(
+    const void *buf,
+    const ::flatbuffers::resolver_function_t *res = nullptr) {
+  return std::unique_ptr<tflite::xnnpack::cache::schema::PackedWeightsT>(GetPackedWeights(buf)->UnPack(res));
+}
+
+inline std::unique_ptr<tflite::xnnpack::cache::schema::PackedWeightsT> UnPackSizePrefixedPackedWeights(
+    const void *buf,
+    const ::flatbuffers::resolver_function_t *res = nullptr) {
+  return std::unique_ptr<tflite::xnnpack::cache::schema::PackedWeightsT>(GetSizePrefixedPackedWeights(buf)->UnPack(res));
+}
+
+}  // namespace schema
+}  // namespace cache
+}  // namespace xnnpack
+}  // namespace tflite
+
+#endif  // FLATBUFFERS_GENERATED_WEIGHTCACHESCHEMA_TFLITE_XNNPACK_CACHE_SCHEMA_H_
diff --git a/tensorflow/lite/delegates/xnnpack/weight_cache_test.cc b/tensorflow/lite/delegates/xnnpack/weight_cache_test.cc
new file mode 100644
index 00000000000000..a270c4f712918f
--- /dev/null
+++ b/tensorflow/lite/delegates/xnnpack/weight_cache_test.cc
@@ -0,0 +1,725 @@
+/* Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+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.
+==============================================================================*/
+#include "tensorflow/lite/delegates/xnnpack/weight_cache.h"
+
+#include <fcntl.h>
+
+#include <algorithm>
+#include <cassert>
+#include <cstdint>
+#include <cstdio>
+#include <cstring>
+#include <iterator>
+#include <map>
+#include <ostream>
+#include <string>
+#include <unordered_map>
+#include <utility>
+#include <vector>
+
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+#include "xnnpack.h"  // from @XNNPACK
+#include "flatbuffers/verifier.h"  // from @flatbuffers
+#include "tensorflow/lite/c/common.h"
+#include "tensorflow/lite/delegates/xnnpack/weight_cache_schema_generated.h"
+
+namespace tflite::xnnpack {
+
+std::ostream& operator<<(std::ostream& os, const PackIdentifier& p) {
+  return os << "PackIdentifier{pack_algo: " << p.pack_algorithm_id
+            << ", weights_id: " << p.weights_id << ", bias_id: " << p.bias_id
+            << "}";
+}
+
+namespace {
+
+using testing::ElementsAreArray;
+using testing::Ge;
+
+#ifndef XNN_TEST_WEIGHT_CACHE_TEMP_FILE_TEMPATE
+#define XNN_TEST_WEIGHT_CACHE_TEMP_FILE_TEMPATE \
+  "/tmp/weight_cache_test_file.XXXXXX"
+#endif
+constexpr const char kTempFileTemplate[] =
+    XNN_TEST_WEIGHT_CACHE_TEMP_FILE_TEMPATE;
+
+// Wraps a call to `mkstemp` to create temporary files.
+class TempFileDesc {
+ public:
+  static constexpr struct AutoClose {
+  } kAutoCLose;
+
+#if defined(_MSC_VER)
+  TempFileDesc() : fd_() {
+    char filename[L_tmpnam_s];
+    errno_t err = tmpnam_s(filename, L_tmpnam_s);
+    if (err) {
+      fprintf(stderr, "Could not create temporary filename.\n");
+      std::abort();
+    }
+    path_ = filename;
+    fd_ = open(path_.c_str(), O_CREAT | O_EXCL | O_RDWR, 0644);
+    if (fd_ < 0) {
+      fprintf(stderr, "Could not create temporary filename.\n");
+      std::abort();
+    }
+  }
+#else
+  TempFileDesc() : fd_(mkstemp(path_.data())) {
+    if (GetFd() < 0) {
+      perror("Could not create temporary file");
+    }
+  }
+#endif
+
+  explicit TempFileDesc(AutoClose) : TempFileDesc() { Close(); }
+
+  TempFileDesc(const TempFileDesc&) = delete;
+  TempFileDesc& operator=(const TempFileDesc&) = delete;
+
+  friend void swap(TempFileDesc& a, TempFileDesc& b) {
+    std::swap(a.path_, b.path_);
+    std::swap(a.fd_, b.fd_);
+  }
+
+  TempFileDesc(TempFileDesc&& other) { swap(*this, other); }
+  TempFileDesc& operator=(TempFileDesc&& other) {
+    swap(*this, other);
+    return *this;
+  }
+
+  ~TempFileDesc() { Close(); }
+
+  void Close() {
+    if (GetFd() >= 0) {
+      close(fd_);
+      fd_ = -1;
+    }
+  }
+
+  const std::string& GetPath() const { return path_; }
+
+  const char* GetCPath() const { return path_.c_str(); }
+
+  int GetFd() const { return fd_; }
+
+  bool IsOpen() const { return fd_ >= 0; }
+
+ private:
+  std::string path_ = kTempFileTemplate;
+  int fd_ = -1;
+};
+
+TEST(MMapHandleTest, DefaultConstructs) {
+  MMapHandle handle;
+  EXPECT_FALSE(handle.IsMapped());
+  EXPECT_EQ(handle.data(), nullptr);
+  EXPECT_EQ(handle.size(), 0);
+}
+
+TEST(MMapHandleTest, MapNonExitxingFileFails) {
+  // I hope this path doesn't exist...
+  const char* file_path = "sdbgfd";
+  MMapHandle handle;
+  EXPECT_FALSE(handle.Map(file_path));
+}
+
+TEST(MMapHandleTest, MapExistingFileWorks) {
+  using std::size;
+
+  const std::string payload = "This is some data in the file.";
+
+  TempFileDesc tmp_file;
+  ASSERT_TRUE(tmp_file.IsOpen());
+  write(tmp_file.GetFd(), payload.c_str(), size(payload));
+  tmp_file.Close();
+
+  MMapHandle handle;
+  ASSERT_TRUE(handle.Map(tmp_file.GetCPath()));
+  EXPECT_TRUE(handle.IsMapped());
+  EXPECT_NE(handle.data(), nullptr);
+  EXPECT_THAT(handle.size(), Ge(size(payload)));
+  EXPECT_THAT(handle, ElementsAreArray(payload));
+
+  handle.UnMap();
+  EXPECT_FALSE(handle.IsMapped());
+  EXPECT_EQ(handle.data(), nullptr);
+  EXPECT_EQ(handle.size(), 0);
+}
+
+TEST(MMapHandleTest, MoveConstructs) {
+  const std::string payload = "This is some data in the file.";
+
+  TempFileDesc tmp_file;
+  ASSERT_TRUE(tmp_file.IsOpen());
+  write(tmp_file.GetFd(), payload.c_str(), size(payload));
+  tmp_file.Close();
+
+  MMapHandle handle;
+  ASSERT_TRUE(handle.Map(tmp_file.GetCPath()));
+
+  MMapHandle handle2(std::move(handle));
+
+  // We are checking that the moved from handle has lost control over the data.
+  // NOLINTBEGIN(bugprone-use-after-move)
+  EXPECT_FALSE(handle.IsMapped());
+  EXPECT_EQ(handle.data(), nullptr);
+  EXPECT_EQ(handle.size(), 0);
+  // NOLINTEND(bugprone-use-after-move)
+
+  EXPECT_TRUE(handle2.IsMapped());
+  EXPECT_NE(handle2.data(), nullptr);
+  EXPECT_THAT(handle2.size(), Ge(size(payload)));
+  EXPECT_THAT(handle2, ElementsAreArray(payload));
+}
+
+TEST(WeightCacheBuilderTest, ReserveAppendWriteWorks) {
+  using std::size;
+
+  const std::string payload = "This is some data in the file.";
+  const PackIdentifier dummy_id{1, 2, 3};
+
+  WeightCacheBuilder builder;
+
+  const size_t payload_size = size(payload);
+  void* buffer = builder.Reserve(payload_size);
+  std::memcpy(buffer, payload.c_str(), payload_size);
+  auto loc = builder.Append(dummy_id, buffer, payload_size);
+
+  EXPECT_EQ(loc.size, payload_size);
+  EXPECT_EQ(builder.BufferData().size(), payload_size);
+  EXPECT_TRUE(builder.ShouldWrite());
+
+  TempFileDesc tmp_file;
+  ASSERT_TRUE(tmp_file.IsOpen());
+  tmp_file.Close();
+
+  ASSERT_TRUE(builder.Write(tmp_file.GetCPath()));
+
+  MMapHandle handle;
+  ASSERT_TRUE(handle.Map(tmp_file.GetCPath()));
+
+  const cache::schema::PackedWeights* const packed_weights =
+      cache::schema::GetPackedWeights(handle.data());
+  ASSERT_NE(packed_weights, nullptr);
+  EXPECT_LE(packed_weights->flatbuffer_size(), size(handle) - size(payload));
+  ASSERT_NE(packed_weights->buffers(), nullptr);
+  ASSERT_EQ(packed_weights->buffers()->size(), 1);
+  ASSERT_NE(packed_weights->buffers()->Get(0), nullptr);
+  ASSERT_EQ(packed_weights->buffers()->Get(0)->size(), size(payload));
+  EXPECT_EQ(packed_weights->buffers()->Get(0)->offset(), 0);
+  ASSERT_EQ(packed_weights->buffers()->Get(0)->packing_algorithm_id(),
+            dummy_id.pack_algorithm_id);
+  ASSERT_EQ(packed_weights->buffers()->Get(0)->weights_id(),
+            dummy_id.weights_id);
+  ASSERT_EQ(packed_weights->buffers()->Get(0)->bias_id(), dummy_id.bias_id);
+
+  flatbuffers::Verifier verifier(handle.data(), handle.size());
+  EXPECT_TRUE(cache::schema::VerifyPackedWeightsBuffer(verifier))
+      << packed_weights->flatbuffer_size() << " " << handle.size() << " "
+      << packed_weights->buffers()->size() << "\n"
+      << tmp_file.GetPath();
+}
+
+TEST(WeightCacheBuilderTest, AppendWithoutReserveWriteWorks) {
+  using std::size;
+
+  const std::string payload = "This is some data in the file.";
+  const PackIdentifier dummy_id{1, 2, 3};
+
+  WeightCacheBuilder builder;
+
+  const size_t payload_size = size(payload);
+  auto loc = builder.Append(dummy_id, payload.c_str(), payload_size);
+
+  EXPECT_EQ(loc.size, payload_size);
+  EXPECT_EQ(builder.BufferData().size(), payload_size);
+  EXPECT_TRUE(builder.ShouldWrite());
+
+  TempFileDesc tmp_file;
+  ASSERT_TRUE(tmp_file.IsOpen());
+  tmp_file.Close();
+
+  ASSERT_TRUE(builder.Write(tmp_file.GetCPath()));
+
+  MMapHandle handle;
+  ASSERT_TRUE(handle.Map(tmp_file.GetCPath()));
+
+  const cache::schema::PackedWeights* const packed_weights =
+      cache::schema::GetPackedWeights(handle.data());
+  ASSERT_NE(packed_weights, nullptr);
+  EXPECT_LE(packed_weights->flatbuffer_size(), size(handle) - size(payload));
+  ASSERT_NE(packed_weights->buffers(), nullptr);
+  ASSERT_EQ(packed_weights->buffers()->size(), 1);
+  ASSERT_NE(packed_weights->buffers()->Get(0), nullptr);
+  ASSERT_EQ(packed_weights->buffers()->Get(0)->size(), size(payload));
+  EXPECT_EQ(packed_weights->buffers()->Get(0)->offset(), 0);
+  ASSERT_EQ(packed_weights->buffers()->Get(0)->packing_algorithm_id(),
+            dummy_id.pack_algorithm_id);
+  ASSERT_EQ(packed_weights->buffers()->Get(0)->weights_id(),
+            dummy_id.weights_id);
+  ASSERT_EQ(packed_weights->buffers()->Get(0)->bias_id(), dummy_id.bias_id);
+
+  flatbuffers::Verifier verifier(handle.data(), handle.size());
+  EXPECT_TRUE(cache::schema::VerifyPackedWeightsBuffer(verifier))
+      << packed_weights->flatbuffer_size() << " " << handle.size() << " "
+      << packed_weights->buffers()->size() << "\n"
+      << tmp_file.GetPath();
+}
+
+TEST(WeightCacheBuilderTest, NonExistingPathFails) {
+  using std::size;
+
+  const std::string payload = "This is some data in the file.";
+  const PackIdentifier dummy_id{1, 2, 3};
+
+  WeightCacheBuilder builder;
+
+  const size_t payload_size = size(payload);
+  auto loc = builder.Append(dummy_id, payload.c_str(), payload_size);
+
+  EXPECT_EQ(loc.size, payload_size);
+  EXPECT_EQ(builder.BufferData().size(), payload_size);
+  EXPECT_TRUE(builder.ShouldWrite());
+
+  EXPECT_FALSE(builder.Write(""));
+  EXPECT_FALSE(builder.Write("/selktjdsljf"));
+}
+
+struct FakeContext {
+  // Adds a new tensor and it's backing buffer to the context.
+  //
+  // The tensor `data` will not be set until `FinalizeTensors` is called.
+  void AddTensor(int buffer_identifier, size_t size) {
+    buffers.emplace_back(size, buffer_identifier);
+    tensors.push_back({});
+    tensors.back().allocation_type = kTfLiteMmapRo;
+    tensor_buffer_identifiers[tensors.size() - 1] = buffer_identifier;
+  }
+
+  // Updates the tensor data mappings.
+  //
+  // This needs to be called every time the context `tensors` list is
+  // reallocated (mainly because of insertions).
+  void FinalizeTensors() {
+    for (size_t i = 0; i < tensors.size(); ++i) {
+      tensors[i].data.data = buffers[i].data();
+      tensors[i].bytes = buffers[i].size();
+    }
+  }
+
+  // Creates a look up key for the XNNPack weight provider C interface.
+  xnn_weights_cache_look_up_key LookUpKey(const uint32_t algorithm_seed,
+                                          const int weights_index) const {
+    return {.seed = algorithm_seed,
+            .kernel = buffers[weights_index].data(),
+            .bias = nullptr};
+  }
+
+  // Creates a look up key for the XNNPack weight provider C interface.
+  xnn_weights_cache_look_up_key LookUpKey(const uint32_t algorithm_seed,
+                                          const int weights_index,
+                                          const int bias_index) const {
+    return {.seed = algorithm_seed,
+            .kernel = buffers[weights_index].data(),
+            .bias = buffers[bias_index].data()};
+  }
+
+  // Helps creating fake packed data.
+  void AddTensorToPack(std::vector<uint8_t>& pack_buffer, int index) {
+    const std::vector<uint8_t>& buffer = buffers[index];
+    pack_buffer.resize(std::max(size(pack_buffer), size(buffer)));
+    for (size_t i = 0; i < size(buffer); ++i) {
+      pack_buffer[i] ^= buffer[i];
+    }
+  }
+
+  // Packs the referenced tensors into one buffer.
+  //
+  // Returns the pack id to retrieve the packed reference data from
+  // `packed_buffers`.
+  template <class... Ids>
+  PackIdentifier PackTensors(xnn_weights_cache_t weight_cache,
+                             const uint32_t algorithm_seed,
+                             const Ids... tensor_indices) {
+    // Create fake packed and save the result for later lookup tests.
+
+    PackIdentifier pack_id{algorithm_seed,
+                           tensor_buffer_identifiers[tensor_indices]...};
+    PackedBuffer& packed =
+        packed_buffers.emplace(pack_id, PackedBuffer{})->second;
+    (AddTensorToPack(packed.buffer, tensor_indices), ...);
+
+    // Add the packed buffer to the XNNPack cache. Normaly you would pack in
+    // place where the reserved space is.
+    xnn_weights_cache_look_up_key look_up_key =
+        LookUpKey(algorithm_seed, tensor_indices...);
+    packed.offset = weight_cache->look_up_or_insert(
+        weight_cache->context, &look_up_key, packed.buffer.data(),
+        packed.buffer.size());
+    return pack_id;
+  }
+
+  struct PackedBuffer {
+    size_t offset;
+    std::vector<uint8_t> buffer;
+  };
+
+  std::vector<TfLiteTensor> tensors;
+  std::vector<std::vector<uint8_t>> buffers;
+  std::unordered_multimap<PackIdentifier, PackedBuffer, PackIdentifier::Hash>
+      packed_buffers;
+  std::unordered_map<size_t, size_t> tensor_buffer_identifiers;
+};
+
+struct BuildMMapWeightCacheProviderTest : testing::Test {
+  enum { kAlgoSeed1, kAlgoSeed2, kAlgoSeed3 };
+  enum { kBufferId1, kBufferId2, kBufferId3, kBufferId4 };
+
+  void SetUp() override {
+    AddTensors();
+    EndSetup();
+  }
+
+  void AddTensors() {
+    ctx.AddTensor(/*buffer_identifier=*/kBufferId1, /*size=*/12);
+    ctx.AddTensor(/*buffer_identifier=*/kBufferId2, /*size=*/43);
+    ctx.AddTensor(/*buffer_identifier=*/kBufferId3, /*size=*/64);
+    ctx.AddTensor(/*buffer_identifier=*/kBufferId4, /*size=*/8);
+  }
+
+  void EndSetup() {
+    ctx.FinalizeTensors();
+    cache_provider.MapTensorIdentifiers(ctx.tensors.data(), ctx.tensors.size(),
+                                        ctx.tensor_buffer_identifiers);
+  }
+
+  FakeContext ctx;
+  MMapWeightCacheProvider cache_provider;
+};
+
+TEST_F(BuildMMapWeightCacheProviderTest, LookUpFailsIfKeyDoesntMatch) {
+  xnn_weights_cache_look_up_key look_up_key{};
+  EXPECT_EQ(cache_provider.LookUp(&look_up_key), SIZE_MAX);
+}
+
+TEST_F(BuildMMapWeightCacheProviderTest, LookUpSucceeds) {
+  enum { kWeightIndex, kBiasIndex };
+  const auto pack_id = ctx.PackTensors(&cache_provider.GetCacheProvider(),
+                                       kAlgoSeed1, kWeightIndex, kBiasIndex);
+  const xnn_weights_cache_look_up_key look_up_key =
+      ctx.LookUpKey(kAlgoSeed1, kWeightIndex, kBiasIndex);
+
+  EXPECT_EQ(cache_provider.LookUp(&look_up_key),
+            ctx.packed_buffers.find(pack_id)->second.offset);
+}
+
+TEST_F(BuildMMapWeightCacheProviderTest,
+       DifferentAlgoSeedsSameTensorsDontConflict) {
+  enum { kWeightIndex, kBiasIndex };
+  const auto pack_id_1 = ctx.PackTensors(&cache_provider.GetCacheProvider(),
+                                         kAlgoSeed1, kWeightIndex, kBiasIndex);
+  const auto pack_id_2 = ctx.PackTensors(&cache_provider.GetCacheProvider(),
+                                         kAlgoSeed2, kWeightIndex, kBiasIndex);
+
+  const xnn_weights_cache_look_up_key look_up_key_1 =
+      ctx.LookUpKey(kAlgoSeed1, kWeightIndex, kBiasIndex);
+  const xnn_weights_cache_look_up_key look_up_key_2 =
+      ctx.LookUpKey(kAlgoSeed2, kWeightIndex, kBiasIndex);
+
+  EXPECT_EQ(cache_provider.LookUp(&look_up_key_1),
+            ctx.packed_buffers.find(pack_id_1)->second.offset);
+  EXPECT_EQ(cache_provider.LookUp(&look_up_key_2),
+            ctx.packed_buffers.find(pack_id_2)->second.offset);
+  EXPECT_NE(cache_provider.LookUp(&look_up_key_1),
+            cache_provider.LookUp(&look_up_key_2));
+}
+
+TEST_F(BuildMMapWeightCacheProviderTest,
+       SameAlgoSeedDifferentTensorsDontConflict) {
+  enum { kWeightIndex1, kWeightIndex2, kBiasIndex1, kBiasIndex2 };
+  const auto pack_id_1 =
+      ctx.PackTensors(&cache_provider.GetCacheProvider(), kAlgoSeed1,
+                      kWeightIndex1, kBiasIndex1);
+  const auto pack_id_2 =
+      ctx.PackTensors(&cache_provider.GetCacheProvider(), kAlgoSeed1,
+                      kWeightIndex2, kBiasIndex1);
+  const auto pack_id_3 =
+      ctx.PackTensors(&cache_provider.GetCacheProvider(), kAlgoSeed1,
+                      kWeightIndex1, kBiasIndex2);
+  const auto pack_id_4 =
+      ctx.PackTensors(&cache_provider.GetCacheProvider(), kAlgoSeed1,
+                      kWeightIndex2, kBiasIndex2);
+
+  const xnn_weights_cache_look_up_key look_up_key_1 =
+      ctx.LookUpKey(kAlgoSeed1, kWeightIndex1, kBiasIndex1);
+  const xnn_weights_cache_look_up_key look_up_key_2 =
+      ctx.LookUpKey(kAlgoSeed1, kWeightIndex2, kBiasIndex1);
+  const xnn_weights_cache_look_up_key look_up_key_3 =
+      ctx.LookUpKey(kAlgoSeed1, kWeightIndex1, kBiasIndex2);
+  const xnn_weights_cache_look_up_key look_up_key_4 =
+      ctx.LookUpKey(kAlgoSeed1, kWeightIndex2, kBiasIndex2);
+
+  EXPECT_EQ(cache_provider.LookUp(&look_up_key_1),
+            ctx.packed_buffers.find(pack_id_1)->second.offset);
+  EXPECT_EQ(cache_provider.LookUp(&look_up_key_2),
+            ctx.packed_buffers.find(pack_id_2)->second.offset);
+  EXPECT_EQ(cache_provider.LookUp(&look_up_key_3),
+            ctx.packed_buffers.find(pack_id_3)->second.offset);
+  EXPECT_EQ(cache_provider.LookUp(&look_up_key_4),
+            ctx.packed_buffers.find(pack_id_4)->second.offset);
+  EXPECT_NE(cache_provider.LookUp(&look_up_key_1),
+            cache_provider.LookUp(&look_up_key_2));
+  EXPECT_NE(cache_provider.LookUp(&look_up_key_1),
+            cache_provider.LookUp(&look_up_key_3));
+  EXPECT_NE(cache_provider.LookUp(&look_up_key_1),
+            cache_provider.LookUp(&look_up_key_4))
+      << pack_id_1 << " " << pack_id_4;
+  EXPECT_NE(cache_provider.LookUp(&look_up_key_2),
+            cache_provider.LookUp(&look_up_key_3));
+  EXPECT_NE(cache_provider.LookUp(&look_up_key_2),
+            cache_provider.LookUp(&look_up_key_4));
+  EXPECT_NE(cache_provider.LookUp(&look_up_key_3),
+            cache_provider.LookUp(&look_up_key_4));
+}
+
+TEST_F(BuildMMapWeightCacheProviderTest, FinalizeWorks) {
+  enum { kWeightIndex1, kBiasIndex, kWeightIndex2 };
+  TempFileDesc tmp_file;
+
+  ctx.PackTensors(&cache_provider.GetCacheProvider(), kAlgoSeed1, kWeightIndex1,
+                  kBiasIndex);
+  ctx.PackTensors(&cache_provider.GetCacheProvider(), kAlgoSeed2,
+                  kWeightIndex2);
+
+  EXPECT_FALSE(cache_provider.Finalize());
+  cache_provider.SetFilePath(tmp_file.GetCPath());
+
+  EXPECT_TRUE(cache_provider.IsActive());
+  EXPECT_TRUE(cache_provider.IsBuilding());
+  ASSERT_TRUE(cache_provider.Finalize());
+
+  ASSERT_TRUE(cache_provider.IsFinalized());
+}
+
+struct LoadMMapWeightCacheProviderTest : BuildMMapWeightCacheProviderTest {
+  enum { kWeightIndex1, kBiasIndex, kWeightIndex2 };
+
+  void SetUp() override {
+    BuildMMapWeightCacheProviderTest::SetUp();
+    cache_provider.SetFilePath(tmp_file.GetCPath());
+
+    pack_id_1 = ctx.PackTensors(&cache_provider.GetCacheProvider(), kAlgoSeed1,
+                                kWeightIndex1, kBiasIndex);
+    pack_id_2 = ctx.PackTensors(&cache_provider.GetCacheProvider(), kAlgoSeed2,
+                                kWeightIndex2);
+
+    ASSERT_TRUE(cache_provider.Finalize());
+    ASSERT_TRUE(cache_provider.IsFinalized());
+  }
+
+  xnn_weights_cache_look_up_key LookUpKey1() const {
+    return ctx.LookUpKey(kAlgoSeed1, kWeightIndex1, kBiasIndex);
+  }
+
+  xnn_weights_cache_look_up_key LookUpKey2() const {
+    return ctx.LookUpKey(kAlgoSeed2, kWeightIndex2);
+  }
+
+  TempFileDesc tmp_file;
+  PackIdentifier pack_id_1;
+  PackIdentifier pack_id_2;
+};
+
+TEST_F(LoadMMapWeightCacheProviderTest, LookUpFailsIfKeyDoesntMatch) {
+  xnn_weights_cache_look_up_key look_up_key{};
+  EXPECT_EQ(cache_provider.LookUp(&look_up_key), SIZE_MAX);
+}
+
+template <class T>
+class LightSpan {
+ public:
+  using value_type = T;
+
+  LightSpan(const void* data, const size_t size)
+      : ptr_(reinterpret_cast<T*>(data)), size_(size) {}
+
+  const T* begin() const { return ptr_; }
+  const T* end() const { return ptr_ + size_; }
+
+ private:
+  T* ptr_;
+  size_t size_;
+};
+
+TEST_F(LoadMMapWeightCacheProviderTest, LookUpSucceeds) {
+  const auto& reference_1 = ctx.packed_buffers.find(pack_id_1)->second;
+  const auto& reference_2 = ctx.packed_buffers.find(pack_id_2)->second;
+
+  const xnn_weights_cache_look_up_key look_up_key_1 = LookUpKey1();
+  const xnn_weights_cache_look_up_key look_up_key_2 = LookUpKey2();
+
+  const uint64_t offset_1 = cache_provider.LookUp(&look_up_key_1);
+  const uint64_t offset_2 = cache_provider.LookUp(&look_up_key_2);
+
+  ASSERT_EQ(offset_1, reference_1.offset);
+  ASSERT_EQ(offset_2, reference_2.offset);
+
+  const void* const addr_1 = cache_provider.OffsetToAddr(offset_1);
+  const void* const addr_2 = cache_provider.OffsetToAddr(offset_2);
+
+  ASSERT_NE(addr_1, nullptr);
+  ASSERT_NE(addr_2, nullptr);
+
+  EXPECT_THAT(LightSpan<const uint8_t>(addr_1, reference_1.buffer.size()),
+              ElementsAreArray(reference_1.buffer));
+  EXPECT_THAT(LightSpan<const uint8_t>(addr_2, reference_2.buffer.size()),
+              ElementsAreArray(reference_2.buffer));
+}
+
+TEST(MMapWeightCacheProviderTest, XnnpackCApiJourney) {
+  using std::size;
+  TempFileDesc temp_fd(TempFileDesc::kAutoCLose);
+  const int32_t fake_packing_algo_seed = 0xBA0BAB;
+  const char packed_data_ref_1[] = {1, 2, 3, 4, 5, 6, 7, 8, 9};
+  const char packed_data_ref_2[] = {26, 32, 43, 59, 34, 65, 80, 101};
+  auto bytes = [](const auto& array) { return size(array) * sizeof(array[0]); };
+
+  constexpr int kBufferCount = 10;
+  // We are going to feed dummy packed data. We only need a valid pointer
+  // address to map to a buffer identifier.
+  char fake_buffer_pointer[kBufferCount] = {0};
+
+  {  // Build and reload scenario.
+    TfLiteTensor tensors[kBufferCount];
+    std::unordered_map<size_t, size_t> tensor_buffer_identifiers;
+    for (int i = 0; i < kBufferCount; ++i) {
+      tensors[0].data.data = (void*)(fake_buffer_pointer + i);
+      tensor_buffer_identifiers[i] = i + 1;
+    }
+
+    MMapWeightCacheProvider cache_provider;
+    cache_provider.SetFilePath(temp_fd.GetCPath());
+
+    xnn_weights_cache_t cache = &cache_provider.GetCacheProvider();
+    cache_provider.MapTensorIdentifiers(tensors, size(tensors),
+                                        tensor_buffer_identifiers);
+
+    const xnn_weights_cache_look_up_key look_up_key_1{
+        .seed = fake_packing_algo_seed,
+        .kernel = tensors[0].data.data,
+        .bias = tensors[1].data.data};
+
+    // Lookup non-packed tensor.
+    ASSERT_EQ(cache->look_up(cache, &look_up_key_1), SIZE_MAX);
+    // Reserve space, write data and add packed data.
+    void* const reserved_ptr =
+        cache->reserve_space(cache, bytes(packed_data_ref_1));
+    ASSERT_NE(reserved_ptr, nullptr);
+    std::memcpy(reserved_ptr, packed_data_ref_1, bytes(packed_data_ref_1));
+    const size_t build_offset_1 = cache->look_up_or_insert(
+        cache, &look_up_key_1, reserved_ptr, bytes(packed_data_ref_1));
+
+    // Check that a second insertion with the same key returns the same offset.
+    const size_t build_offset_redundant = cache->look_up_or_insert(
+        cache, &look_up_key_1, reserved_ptr, bytes(packed_data_ref_1));
+    EXPECT_EQ(build_offset_1, build_offset_redundant);
+
+    // Lookup newly packed tensor.
+    ASSERT_EQ(cache->look_up(cache, &look_up_key_1), build_offset_1);
+
+    // Add a tensor without reserving before.
+    const xnn_weights_cache_look_up_key look_up_key_2{
+        .seed = fake_packing_algo_seed,
+        .kernel = tensors[2].data.data,
+        .bias = tensors[3].data.data};
+    const size_t build_offset_2 = cache->look_up_or_insert(
+        cache, &look_up_key_2, (void*)packed_data_ref_2,
+        bytes(packed_data_ref_2));
+
+    // Save the cache to disk and reload.
+    ASSERT_TRUE(cache_provider.Finalize());
+
+    ASSERT_TRUE(cache->is_finalized(cache));
+
+    const size_t reload_offset_1 = cache->look_up(cache, &look_up_key_1);
+    ASSERT_EQ(reload_offset_1, build_offset_1);
+
+    const void* const loaded_packed_data_1 =
+        cache->offset_to_addr(cache, reload_offset_1);
+    ASSERT_NE(loaded_packed_data_1, nullptr);
+    EXPECT_THAT(
+        LightSpan<const char>(loaded_packed_data_1, size(packed_data_ref_1)),
+        ElementsAreArray(packed_data_ref_1));
+
+    const size_t reload_offset_2 = cache->look_up(cache, &look_up_key_2);
+    ASSERT_EQ(reload_offset_2, build_offset_2);
+
+    const void* const loaded_packed_data_2 =
+        cache->offset_to_addr(cache, reload_offset_2);
+    ASSERT_NE(loaded_packed_data_2, nullptr);
+    EXPECT_THAT(
+        LightSpan<const char>(loaded_packed_data_2, size(packed_data_ref_2)),
+        ElementsAreArray(packed_data_ref_2));
+  }
+
+  {  // Load existing cache scenario.
+    TfLiteTensor tensors[kBufferCount];
+    std::unordered_map<size_t, size_t> tensor_buffer_identifiers;
+    for (int i = 0; i < kBufferCount; ++i) {
+      tensors[0].data.data = (void*)(fake_buffer_pointer + i);
+      tensor_buffer_identifiers[i] = i + 1;
+    }
+
+    MMapWeightCacheProvider cache_provider;
+    ASSERT_TRUE(cache_provider.Load(temp_fd.GetCPath()));
+
+    xnn_weights_cache_t cache = &cache_provider.GetCacheProvider();
+    cache_provider.MapTensorIdentifiers(tensors, size(tensors),
+                                        tensor_buffer_identifiers);
+
+    const xnn_weights_cache_look_up_key look_up_key_1{
+        .seed = fake_packing_algo_seed,
+        .kernel = tensors[0].data.data,
+        .bias = tensors[1].data.data};
+
+    const xnn_weights_cache_look_up_key look_up_key_2{
+        .seed = fake_packing_algo_seed,
+        .kernel = tensors[2].data.data,
+        .bias = tensors[3].data.data};
+
+    ASSERT_TRUE(cache->is_finalized(cache));
+
+    const size_t offset_1 = cache->look_up(cache, &look_up_key_1);
+    const void* const loaded_packed_data_1 =
+        cache->offset_to_addr(cache, offset_1);
+    ASSERT_NE(loaded_packed_data_1, nullptr);
+    EXPECT_THAT(
+        LightSpan<const char>(loaded_packed_data_1, size(packed_data_ref_1)),
+        ElementsAreArray(packed_data_ref_1));
+
+    const size_t offset_2 = cache->look_up(cache, &look_up_key_2);
+    const void* const loaded_packed_data_2 =
+        cache->offset_to_addr(cache, offset_2);
+    ASSERT_NE(loaded_packed_data_2, nullptr);
+    EXPECT_THAT(
+        LightSpan<const char>(loaded_packed_data_2, size(packed_data_ref_2)),
+        ElementsAreArray(packed_data_ref_2));
+  }
+}
+
+}  // namespace
+}  // namespace tflite::xnnpack
diff --git a/tensorflow/lite/delegates/xnnpack/xnnpack_delegate.cc b/tensorflow/lite/delegates/xnnpack/xnnpack_delegate.cc
index 76cc6dba209ab9..1503fea039580f 100644
--- a/tensorflow/lite/delegates/xnnpack/xnnpack_delegate.cc
+++ b/tensorflow/lite/delegates/xnnpack/xnnpack_delegate.cc
@@ -38,7 +38,9 @@ limitations under the License.
 #include "tensorflow/lite/core/api/profiler.h"
 #include "tensorflow/lite/core/c/builtin_op_data.h"
 #include "tensorflow/lite/core/c/common.h"
+#include "tensorflow/lite/core/subgraph.h"
 #include "tensorflow/lite/delegates/xnnpack/quantization_util.h"
+#include "tensorflow/lite/delegates/xnnpack/weight_cache.h"
 #include "tensorflow/lite/kernels/cpu_backend_context.h"
 #include "tensorflow/lite/kernels/internal/compatibility.h"
 #include "tensorflow/lite/kernels/internal/tensor_ctypes.h"
@@ -487,6 +489,8 @@ class VariableHolder {
   std::map<uint32_t, const TfLiteTensor*> global_id_to_dims_and_type_;
 };
 
+class Subgraph;
+
 class Delegate {
   friend class Subgraph;
 
@@ -525,6 +529,29 @@ class Delegate {
         options != nullptr ? *options : TfLiteXNNPackDelegateOptionsDefault();
     delegate_.flags = GetXNNPackDelegateFlags();
     workspace_.reset(workspace);
+
+    // If no weight cache is provided, add one when requested.
+    if (!options_.weights_cache) {
+      if (options_.experimental_weight_cache_file_path) {
+        if (weight_cache_provider_.Load(
+                options_.experimental_weight_cache_file_path)) {
+          TFLITE_LOG_PROD(tflite::TFLITE_LOG_INFO,
+                          "XNNPack weight cache loaded from '%s'.",
+                          options_.experimental_weight_cache_file_path);
+        } else {
+          TFLITE_LOG_PROD(
+              tflite::TFLITE_LOG_INFO,
+              "XNNPack weight cache not found at '%s', building it.",
+              options_.experimental_weight_cache_file_path);
+        }
+        options_.weights_cache =
+            reinterpret_cast<TfLiteXNNPackDelegateWeightsCache*>(
+                weight_cache_provider_.GetCacheProvider().context);
+      } else {
+        TFLITE_LOG_PROD(tflite::TFLITE_LOG_INFO,
+                        "XNNPack weight cache not enabled.");
+      }
+    }
   }
 
   TfLiteIntArray* PrepareOpsToDelegate(TfLiteContext* context);
@@ -711,6 +738,10 @@ class Delegate {
   TfLiteXNNPackDelegateOptions options_{};
   VariableHolder variable_holder_;
   std::mutex workspace_mutex_;
+
+  // If no weight cache is provided and a cache is set in the delegate options,
+  // this will be used as a weight cache.
+  MMapWeightCacheProvider weight_cache_provider_;
 };
 
 class Subgraph {
@@ -781,6 +812,13 @@ class Subgraph {
   static Subgraph* Create(TfLiteContext* context,
                           const TfLiteDelegateParams* params,
                           Delegate& delegate) {
+    // Map tensors identifiers before packing anything.
+    if (delegate.weight_cache_provider_.IsActive()) {
+      delegate.weight_cache_provider_.MapTensorIdentifiers(
+          context->tensors, context->tensors_size,
+          reinterpret_cast<tflite::Subgraph*>(context->impl_)
+              ->GetTensorBufferIdentifiers());
+    }
     // Convert subgraph inputs and outputs to hash sets for faster lookup.
     const std::unordered_set<int> inputs(
         &params->input_tensors->data[0],
@@ -1121,6 +1159,18 @@ class Subgraph {
   TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node,
                        bool enable_subgraph_reshaping, Delegate* delegate) {
     std::lock_guard<std::mutex> lock(delegate->workspace_mutex_);
+
+    // The weights cache needs to be finalized only once. Prepare will be called
+    // for each partition after all the partitions have been created (therefore
+    // all the weights are known and have been packed).
+    if (delegate->weight_cache_provider_.IsActive()) {
+      if (!delegate->weight_cache_provider_.Finalize()) {
+        TF_LITE_KERNEL_LOG(context,
+                           "XNNPack delegate failed to finalize cache.");
+        return kTfLiteError;
+      }
+    }
+
     if (enable_subgraph_reshaping) {
       xnn_status status = xnn_status_invalid_state;
       for (int i = 0; i < inputs_.size(); ++i) {
@@ -1170,10 +1220,8 @@ class Subgraph {
           return kTfLiteError;
         }
       }
-      return kTfLiteOk;
-    } else {
-      return kTfLiteOk;
     }
+    return kTfLiteOk;
   }
 
   TfLiteStatus Invoke(TfLiteContext* context, bool enable_subgraph_reshaping,
diff --git a/tensorflow/lite/delegates/xnnpack/xnnpack_delegate.h b/tensorflow/lite/delegates/xnnpack/xnnpack_delegate.h
index aa11998dc0fc49..1d40c04db940ac 100644
--- a/tensorflow/lite/delegates/xnnpack/xnnpack_delegate.h
+++ b/tensorflow/lite/delegates/xnnpack/xnnpack_delegate.h
@@ -70,6 +70,10 @@ typedef struct {
   bool handle_variable_ops;
   // Enable adaptive optimization for AVX CPUs.
   bool experimental_adaptive_avx_optimization;
+  // Path to the weight cache to load if `weight_cache` is undefined.
+  //
+  // WARNING this is an experimental flag.
+  const char* experimental_weight_cache_file_path;
 } TfLiteXNNPackDelegateOptions;
 
 // Returns a structure with the default XNNPack delegate options.
@@ -111,11 +115,13 @@ TFL_CAPI_EXPORT void TfLiteXNNPackDelegateDelete(TfLiteDelegate* delegate);
 // reduce memory bandwidth.
 TFL_CAPI_EXPORT struct TfLiteXNNPackDelegateWeightsCache*
 TfLiteXNNPackDelegateWeightsCacheCreate();
+
 // Creates a new weights cache with a specified initial size that can be shared
 // with multiple delegate instances. The weights cache can hold up to size bytes
 // without growing.
 TFL_CAPI_EXPORT struct TfLiteXNNPackDelegateWeightsCache*
 TfLiteXNNPackDelegateWeightsCacheCreateWithSize(size_t size);
+
 // Soft-finalize a weights cache. Extra space will be left in the weights cache
 // to allow for cache "insertion" only if it is a cache hit. This has memory
 // overhead compared to TfLiteXNNPackDelegateWeightsCacheFinalizeHard. Use this
@@ -124,6 +130,7 @@ TfLiteXNNPackDelegateWeightsCacheCreateWithSize(size_t size);
 // Returns true on success, false on error.
 TFL_CAPI_EXPORT bool TfLiteXNNPackDelegateWeightsCacheFinalizeSoft(
     struct TfLiteXNNPackDelegateWeightsCache* cache);
+
 // Hard-finalize a weights cache, cache is effectively frozen and no more cache
 // operations are allowed. Memory is resized to smallest possible. Use this if
 // the number of interpreter instances using XNNPACK delegate can be fixed and
@@ -132,6 +139,7 @@ TFL_CAPI_EXPORT bool TfLiteXNNPackDelegateWeightsCacheFinalizeSoft(
 // Returns true on success, false on error.
 TFL_CAPI_EXPORT bool TfLiteXNNPackDelegateWeightsCacheFinalizeHard(
     struct TfLiteXNNPackDelegateWeightsCache* cache);
+
 // Destroys a weights cache created with
 // `TfLiteXNNPackDelegateWeightsCacheCreate` call.
 TFL_CAPI_EXPORT void TfLiteXNNPackDelegateWeightsCacheDelete(
diff --git a/tensorflow/lite/tflite_with_xnnpack.cc b/tensorflow/lite/tflite_with_xnnpack.cc
index 22e8617ec74e21..d443d404c21f05 100644
--- a/tensorflow/lite/tflite_with_xnnpack.cc
+++ b/tensorflow/lite/tflite_with_xnnpack.cc
@@ -23,6 +23,10 @@ namespace tflite {
 std::unique_ptr<TfLiteDelegate, void (*)(TfLiteDelegate*)>
 AcquireXNNPACKDelegate() {
   auto opts = TfLiteXNNPackDelegateOptionsDefault();
+#ifdef TFLITE_XNNPACK_DELEGATE_EXPERIMENTAL_WEIGHT_CACHE_FILE_PATH
+  opts.experimental_weight_cache_file_path =
+      TFLITE_XNNPACK_DELEGATE_EXPERIMENTAL_WEIGHT_CACHE_FILE_PATH;
+#endif
   return std::unique_ptr<TfLiteDelegate, void (*)(TfLiteDelegate*)>(
       TfLiteXNNPackDelegateCreate(&opts), TfLiteXNNPackDelegateDelete);
 }
diff --git a/tensorflow/opensource_only.files b/tensorflow/opensource_only.files
index b2645a331739e3..b83a48134d02e1 100644
--- a/tensorflow/opensource_only.files
+++ b/tensorflow/opensource_only.files
@@ -101,6 +101,7 @@ tf_staging/tensorflow/lite/delegates/utils/experimental/stable_delegate/BUILD:
 tf_staging/tensorflow/lite/delegates/utils/experimental/stable_delegate/delegate_loader.cc:
 tf_staging/tensorflow/lite/delegates/utils/experimental/stable_delegate/delegate_loader.h:
 tf_staging/tensorflow/lite/delegates/utils/experimental/stable_delegate/delegate_loader_test.cc:
+tf_staging/tensorflow/lite/delegates/xnnpack/weight_cache_schema_generated.h:
 tf_staging/tensorflow/lite/experimental/acceleration/configuration/configuration_generated.h:
 tf_staging/tensorflow/lite/experimental/acceleration/mini_benchmark/c/c_api.h:
 tf_staging/tensorflow/lite/experimental/acceleration/mini_benchmark/libjpeg.h:

From fd2a7a57c5276612faff2fa0bf4a908e48c73811 Mon Sep 17 00:00:00 2001
From: Johannes Reifferscheid <jreiffers@google.com>
Date: Mon, 13 May 2024 04:03:57 -0700
Subject: [PATCH 013/478] Reduce input sizes in ReductionTest.

The interpreter is very slow for some of these (up to ~45 seconds
in unoptimized builds). With this change, runtime goes down to at
most ~1 second per test case. At this point, sharding is no longer
helpful (since the fixed setup cost dominates), so we can disable
it.

PiperOrigin-RevId: 633158930
---
 third_party/xla/xla/service/gpu/fusions/BUILD |  1 -
 .../gpu/fusions/reduction_mlir_test.cc        | 32 +++++++++----------
 2 files changed, 16 insertions(+), 17 deletions(-)

diff --git a/third_party/xla/xla/service/gpu/fusions/BUILD b/third_party/xla/xla/service/gpu/fusions/BUILD
index cb89eb885d4e59..02752d4bfffe04 100644
--- a/third_party/xla/xla/service/gpu/fusions/BUILD
+++ b/third_party/xla/xla/service/gpu/fusions/BUILD
@@ -859,7 +859,6 @@ cc_library(
 xla_cc_test(
     name = "reduction_mlir_test",
     srcs = ["reduction_mlir_test.cc"],
-    shard_count = 11,
     use_gpu = True,
     deps = [
         ":mlir_emitter_test_base",
diff --git a/third_party/xla/xla/service/gpu/fusions/reduction_mlir_test.cc b/third_party/xla/xla/service/gpu/fusions/reduction_mlir_test.cc
index 45b46267a2cfa7..a34843735301ab 100644
--- a/third_party/xla/xla/service/gpu/fusions/reduction_mlir_test.cc
+++ b/third_party/xla/xla/service/gpu/fusions/reduction_mlir_test.cc
@@ -41,26 +41,26 @@ TEST_F(ReductionTest, VariadicRowReduce) {
       ROOT t = (f32[], f32[]) tuple(add.0, add.1)
     }
     fused_computation {
-      param_0 = f32[5,200,2048] parameter(0)
-      param_1 = f32[5,200,2048] parameter(1)
+      param_0 = f32[2, 3, 2048] parameter(0)
+      param_1 = f32[2, 3, 2048] parameter(1)
       param_2 = f32[] parameter(2)
-      ROOT d.1 = (f32[5,200], f32[5,200])
+      ROOT d.1 = (f32[2, 3], f32[2, 3])
         reduce(param_0, param_1, param_2, param_2), dimensions={2}, to_apply=Add
     }
     ENTRY main {
-      a = f32[5, 200, 2048] parameter(0)
-      b = f32[5, 200, 2048] parameter(1)
+      a = f32[2, 3, 2048] parameter(0)
+      b = f32[2, 3, 2048] parameter(1)
       c = f32[] constant(0)
-      ROOT fusion = (f32[5,200], f32[5,200]) fusion(a, b, c),
+      ROOT fusion = (f32[2, 3], f32[2, 3]) fusion(a, b, c),
         kind=kInput, calls=fused_computation
     })";
   TF_ASSERT_OK(EmitAndCheckIR(kHloString, R"(
 // CHECK:      @fused_computation
-// CHECK-SAME:   %[[ARG0:.*]]: tensor<5x200x2048xf32> {xla.slice_index = 0
-// CHECK-SAME:   %[[ARG1:.*]]: tensor<5x200x2048xf32> {xla.slice_index = 1
+// CHECK-SAME:   %[[ARG0:.*]]: tensor<2x3x2048xf32> {xla.slice_index = 0
+// CHECK-SAME:   %[[ARG1:.*]]: tensor<2x3x2048xf32> {xla.slice_index = 1
 // CHECK-SAME:   %[[INIT_TENSOR:.*]]: tensor<f32> {xla.slice_index = 2
-// CHECK-SAME:   %[[OUT0:.*]]: tensor<5x200xf32> {xla.slice_index = 3
-// CHECK-SAME:   %[[OUT1:.*]]: tensor<5x200xf32> {xla.slice_index = 4
+// CHECK-SAME:   %[[OUT0:.*]]: tensor<2x3xf32> {xla.slice_index = 3
+// CHECK-SAME:   %[[OUT1:.*]]: tensor<2x3xf32> {xla.slice_index = 4
 // CHECK:        %[[INIT:.*]] = xla_gpu.pure_call @fused_computation_param_2
 // CHECK:        %[[PER_THREAD:.*]]:2 = scf.for
 // CHECK-SAME:       iter_args(%[[A:.*]] = %[[INIT]], %[[B:.*]] = %[[INIT]])
@@ -124,7 +124,7 @@ TEST_F(ReductionTest, RowReduceMOFEpilogue) {
       ROOT mul = f32[] multiply(lhs, rhs)
     }
     fused_computation {
-      param_0 = f32[8,2048] parameter(0)
+      param_0 = f32[8,1024] parameter(0)
       param_1 = f32[] parameter(1)
       reduce1 = f32[8] reduce(param_0, param_1), dimensions={1}, to_apply=Add
       reduce2 = f32[8] reduce(param_0, param_1), dimensions={1}, to_apply=Mul
@@ -134,7 +134,7 @@ TEST_F(ReductionTest, RowReduceMOFEpilogue) {
       ROOT tuple = (f32[8], f32[8], f32[8]) tuple(log, neg, abs)
     }
     ENTRY main {
-      a = f32[8,2048] parameter(0)
+      a = f32[8,1024] parameter(0)
       c = f32[] constant(0)
       ROOT fusion = (f32[8], f32[8], f32[8]) fusion(a, c), kind=kInput,
         calls=fused_computation
@@ -194,14 +194,14 @@ TEST_F(ReductionTest, ColumnReduction) {
       ROOT add = f32[] add(lhs, rhs)
     }
     fused_computation {
-      param_0 = f32[123,2051,321] parameter(0)
+      param_0 = f32[13,1051,321] parameter(0)
       param_1 = f32[] parameter(1)
-      ROOT reduce = f32[123,321] reduce(param_0, param_1), dimensions={1}, to_apply=Add
+      ROOT reduce = f32[13,321] reduce(param_0, param_1), dimensions={1}, to_apply=Add
     }
     ENTRY main {
-      a = f32[123,2051,321] parameter(0)
+      a = f32[13,1051,321] parameter(0)
       c = f32[] constant(0)
-      ROOT fusion = f32[123,321] fusion(a, c), kind=kInput, calls=fused_computation
+      ROOT fusion = f32[13,321] fusion(a, c), kind=kInput, calls=fused_computation
     })";
   TF_ASSERT_OK(EmitAndCheckIR(kHloString, R"(
     // CHECK: xla_gpu.pure_call @Add_add

From c70ea272d66d38af38d4ae424ed226bcadc05abd Mon Sep 17 00:00:00 2001
From: Adrian Kuegel <akuegel@google.com>
Date: Mon, 13 May 2024 04:11:08 -0700
Subject: [PATCH 014/478] Disable cache invalidation logic for mlir emitters.

We have logic to estimate the impact of certain ops on IR size. This problem is
fixed with the MLIR emitters, so we don't need it for them.

PiperOrigin-RevId: 633160564
---
 .../service/gpu/model/gpu_hlo_cost_analysis.cc  | 17 ++++++++++++-----
 1 file changed, 12 insertions(+), 5 deletions(-)

diff --git a/third_party/xla/xla/service/gpu/model/gpu_hlo_cost_analysis.cc b/third_party/xla/xla/service/gpu/model/gpu_hlo_cost_analysis.cc
index c1ec02bc5267cf..4bbe929730c661 100644
--- a/third_party/xla/xla/service/gpu/model/gpu_hlo_cost_analysis.cc
+++ b/third_party/xla/xla/service/gpu/model/gpu_hlo_cost_analysis.cc
@@ -218,12 +218,19 @@ bool GpuHloCostAnalysis::ProducerConsumerMergedTooLarge(
                      IrBasicBlockSplitCount(consumer);
   VLOG(5) << "Basic block split counts: " << IrBasicBlockSplitCount(producer)
           << ", " << IrBasicBlockSplitCount(consumer) << " -> " << n_splits;
-  if (n_splits > kMaxBasicBlockSplitsPerFusion) {
-    return true;
-  }
   int64_t merged_ir_size =
-      (IrSize(producer) * producer_replication + IrSize(consumer)) *
-      (1 << n_splits);
+      (IrSize(producer) * producer_replication + IrSize(consumer));
+  // The MLIR emitters don't have the problem with cache invalidation, so we
+  // don't need to evaluate basic block split counts.
+  if (!producer.GetModule()
+           ->config()
+           .debug_options()
+           .xla_gpu_enable_mlir_emitters()) {
+    if (n_splits > kMaxBasicBlockSplitsPerFusion) {
+      return true;
+    }
+    merged_ir_size *= (1 << n_splits);
+  }
   VLOG(5) << "IR sizes: " << IrSize(producer) << ", " << IrSize(consumer)
           << " -> " << merged_ir_size;
   return merged_ir_size > kMaxIRSize;

From 399d09503d7d0a7584b994fe544eb1fdc76faa37 Mon Sep 17 00:00:00 2001
From: Quentin Khan <qkhan@google.com>
Date: Mon, 13 May 2024 04:51:36 -0700
Subject: [PATCH 015/478] Fix weight cache test typo in XNNPack delegate.

PiperOrigin-RevId: 633168188
---
 tensorflow/lite/delegates/xnnpack/weight_cache_test.cc | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/tensorflow/lite/delegates/xnnpack/weight_cache_test.cc b/tensorflow/lite/delegates/xnnpack/weight_cache_test.cc
index a270c4f712918f..4cf49e2144cbe4 100644
--- a/tensorflow/lite/delegates/xnnpack/weight_cache_test.cc
+++ b/tensorflow/lite/delegates/xnnpack/weight_cache_test.cc
@@ -607,8 +607,8 @@ TEST(MMapWeightCacheProviderTest, XnnpackCApiJourney) {
     TfLiteTensor tensors[kBufferCount];
     std::unordered_map<size_t, size_t> tensor_buffer_identifiers;
     for (int i = 0; i < kBufferCount; ++i) {
-      tensors[0].data.data = (void*)(fake_buffer_pointer + i);
-      tensor_buffer_identifiers[i] = i + 1;
+      tensors[i].data.data = (void*)(fake_buffer_pointer + i);
+      tensor_buffer_identifiers[i] = i;
     }
 
     MMapWeightCacheProvider cache_provider;
@@ -680,8 +680,8 @@ TEST(MMapWeightCacheProviderTest, XnnpackCApiJourney) {
     TfLiteTensor tensors[kBufferCount];
     std::unordered_map<size_t, size_t> tensor_buffer_identifiers;
     for (int i = 0; i < kBufferCount; ++i) {
-      tensors[0].data.data = (void*)(fake_buffer_pointer + i);
-      tensor_buffer_identifiers[i] = i + 1;
+      tensors[i].data.data = (void*)(fake_buffer_pointer + i);
+      tensor_buffer_identifiers[i] = i;
     }
 
     MMapWeightCacheProvider cache_provider;

From 6487006a061ed706e713c5e025c46cf19b9d42c7 Mon Sep 17 00:00:00 2001
From: Johannes Reifferscheid <jreiffers@google.com>
Date: Mon, 13 May 2024 05:19:07 -0700
Subject: [PATCH 016/478] Enable VLOGging of pass pipelines in
 mlir_fusion_emitter.

PiperOrigin-RevId: 633174085
---
 .../xla/service/gpu/fusions/mlir/mlir_fusion_emitter.cc    | 7 +++++++
 1 file changed, 7 insertions(+)

diff --git a/third_party/xla/xla/service/gpu/fusions/mlir/mlir_fusion_emitter.cc b/third_party/xla/xla/service/gpu/fusions/mlir/mlir_fusion_emitter.cc
index 286b3e211df558..41108b774751b9 100644
--- a/third_party/xla/xla/service/gpu/fusions/mlir/mlir_fusion_emitter.cc
+++ b/third_party/xla/xla/service/gpu/fusions/mlir/mlir_fusion_emitter.cc
@@ -287,6 +287,9 @@ MlirFusionEmitterBase::CreateLLVMModule(
                                     buffer_assignment));
 
   mlir::PassManager pm(&mlir_context);
+  if (VLOG_IS_ON(5)) {
+    pm.enableIRPrinting();
+  }
   pm.addPass(mlir::createCanonicalizerPass());
   pm.addPass(mlir::createInlinerPass());
   pm.addPass(mlir::createCanonicalizerPass());
@@ -430,6 +433,10 @@ MlirFusionEmitterBase::CreateMLIRModule(
 
   // Run a minimal simplification pipeline.
   mlir::PassManager pm(&context);
+  if (VLOG_IS_ON(5)) {
+    context.disableMultithreading();
+    pm.enableIRPrinting();
+  }
   pm.addPass(CreateSimplifyArithPass());
   pm.addPass(mlir::createCanonicalizerPass());
   pm.addPass(mlir::createCSEPass());

From 33fe3156406ace60347447a4d5fdb3d8a2b29b94 Mon Sep 17 00:00:00 2001
From: Quentin Khan <qkhan@google.com>
Date: Mon, 13 May 2024 05:22:08 -0700
Subject: [PATCH 017/478] Do not deinitialize XNNPack in
 `TfLiteXNNPackDelegateWeightsCacheDelete`.

PiperOrigin-RevId: 633174598
---
 tensorflow/lite/delegates/xnnpack/xnnpack_delegate.cc | 1 -
 1 file changed, 1 deletion(-)

diff --git a/tensorflow/lite/delegates/xnnpack/xnnpack_delegate.cc b/tensorflow/lite/delegates/xnnpack/xnnpack_delegate.cc
index 1503fea039580f..6ed0696d8d9e5b 100644
--- a/tensorflow/lite/delegates/xnnpack/xnnpack_delegate.cc
+++ b/tensorflow/lite/delegates/xnnpack/xnnpack_delegate.cc
@@ -7846,7 +7846,6 @@ void TfLiteXNNPackDelegateWeightsCacheDelete(
   }
   auto weights_cache = reinterpret_cast<xnn_weights_cache_t>(cache);
   xnn_delete_weights_cache(weights_cache);
-  xnn_deinitialize();
 }
 
 TfLiteXNNPackDelegateOptions TfLiteXNNPackDelegateOptionsDefault() {

From ea106b7963a5ec91cf3eab9108f058d695289600 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Mon, 13 May 2024 05:25:26 -0700
Subject: [PATCH 018/478] Update ops-related pbtxt files.

PiperOrigin-RevId: 633175341
---
 tensorflow/core/ops/ops.pbtxt | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorflow/core/ops/ops.pbtxt b/tensorflow/core/ops/ops.pbtxt
index 0556051577f5e4..53c362ed0ab819 100644
--- a/tensorflow/core/ops/ops.pbtxt
+++ b/tensorflow/core/ops/ops.pbtxt
@@ -1,4 +1,4 @@
-go/debugstr  
+go/debugstr    
 op {
   name: "Abort"
   attr {

From 676078f7101ea1174eef7646a98ab563bfe3f213 Mon Sep 17 00:00:00 2001
From: Greg Olechwierowicz <olechwierowicz@google.com>
Date: Mon, 13 May 2024 07:05:18 -0700
Subject: [PATCH 019/478] [XLA:GPU] Add flag for fully unrolling while loops

PiperOrigin-RevId: 633196820
---
 third_party/xla/xla/debug_options_flags.cc    |  2 +
 .../xla/xla/service/gpu/gpu_compiler.cc       | 53 +++++++++++--------
 third_party/xla/xla/xla.proto                 | 14 ++++-
 3 files changed, 46 insertions(+), 23 deletions(-)

diff --git a/third_party/xla/xla/debug_options_flags.cc b/third_party/xla/xla/debug_options_flags.cc
index 3b31ed9bc3f0b7..3eb1a7a2fd4292 100644
--- a/third_party/xla/xla/debug_options_flags.cc
+++ b/third_party/xla/xla/debug_options_flags.cc
@@ -211,6 +211,8 @@ DebugOptions DefaultDebugOptionsIgnoringFlags() {
   opts.set_xla_gpu_cublas_fallback(true);
   opts.set_xla_gpu_cudnn_gemm_fusion_level(0);
   opts.set_xla_gpu_enable_while_loop_double_buffering(false);
+  opts.set_xla_gpu_enable_while_loop_unrolling(
+      DebugOptions::WHILE_LOOP_UNROLLING_NO_UNROLL);
   opts.set_xla_gpu_ensure_minor_dot_contraction_dims(false);
   opts.set_xla_gpu_filter_kernels_spilling_registers_on_autotuning(true);
   opts.set_xla_gpu_llvm_verification_level(0);
diff --git a/third_party/xla/xla/service/gpu/gpu_compiler.cc b/third_party/xla/xla/service/gpu/gpu_compiler.cc
index 9a861ec2cb0c71..2d33d67d2d6064 100644
--- a/third_party/xla/xla/service/gpu/gpu_compiler.cc
+++ b/third_party/xla/xla/service/gpu/gpu_compiler.cc
@@ -1066,36 +1066,27 @@ absl::Status RunPostFusionPasses(
     HloModule* hlo_module,
     std::function<absl::Status(HloPassPipeline*, const DebugOptions&)>
         add_custom_kernel_replacement_passes) {
+  const DebugOptions& opts = hlo_module->config().debug_options();
+
   HloPassPipeline pipeline("post-fusion optimization");
   pipeline.AddPass<RenameFusions>();
   pipeline.AddPass<AllGatherCombiner>(
-      hlo_module->config()
-          .debug_options()
-          .xla_gpu_all_gather_combine_threshold_bytes(),
+      opts.xla_gpu_all_gather_combine_threshold_bytes(),
       /*combine_threshold_count=*/256,
-      hlo_module->config()
-          .debug_options()
-          .xla_gpu_enable_all_gather_combine_by_dim());
+      opts.xla_gpu_enable_all_gather_combine_by_dim());
   pipeline.AddPass<AllReduceCombiner>(
-      hlo_module->config()
-          .debug_options()
-          .xla_gpu_all_reduce_combine_threshold_bytes(),
+      opts.xla_gpu_all_reduce_combine_threshold_bytes(),
       /*combine_threshold_count=*/256);
   pipeline.AddPass<ReduceScatterCombiner>(
-      hlo_module->config()
-          .debug_options()
-          .xla_gpu_reduce_scatter_combine_threshold_bytes(),
+      opts.xla_gpu_reduce_scatter_combine_threshold_bytes(),
       /*combine_threshold_count=*/256,
-      hlo_module->config()
-          .debug_options()
-          .xla_gpu_enable_reduce_scatter_combine_by_dim());
+      opts.xla_gpu_enable_reduce_scatter_combine_by_dim());
 
-  if (hlo_module->config().debug_options().xla_gpu_all_reduce_contiguous()) {
+  if (opts.xla_gpu_all_reduce_contiguous()) {
     pipeline.AddPass<AllReduceContiguous>();
   }
 
-  TF_RETURN_IF_ERROR(add_custom_kernel_replacement_passes(
-      &pipeline, hlo_module->config().debug_options()));
+  TF_RETURN_IF_ERROR(add_custom_kernel_replacement_passes(&pipeline, opts));
 
   int32_t blueconnect_num_devices_per_host =
       hlo_module->config()
@@ -1105,10 +1096,28 @@ absl::Status RunPostFusionPasses(
     pipeline.AddPass<AllReduceBlueConnect>(blueconnect_num_devices_per_host);
   }
 
-  if (hlo_module->config()
-          .debug_options()
-          .xla_gpu_enable_while_loop_double_buffering()) {
-    pipeline.AddPass<LoopDoubleBufferTransformer>();
+  std::optional<LoopDoubleBufferTransformer::UnrollStrategy> unroll_strategy =
+      std::nullopt;
+  // Support old flag.
+  if (opts.xla_gpu_enable_while_loop_double_buffering()) {
+    unroll_strategy =
+        LoopDoubleBufferTransformer::UnrollStrategy::kDoubleBuffer;
+  }
+  // Support new flag setting style, override the old one.
+  if (opts.xla_gpu_enable_while_loop_unrolling() ==
+      DebugOptions::WHILE_LOOP_UNROLLING_DOUBLE_BUFFER) {
+    unroll_strategy =
+        LoopDoubleBufferTransformer::UnrollStrategy::kDoubleBuffer;
+  }
+  if (opts.xla_gpu_enable_while_loop_unrolling() ==
+      DebugOptions::WHILE_LOOP_UNROLLING_FULL_UNROLL) {
+    LOG_IF(WARNING, unroll_strategy != std::nullopt)
+        << "Overriding double buffering set via "
+           "`xla_gpu_enable_while_loop_double_buffering` flag.";
+    unroll_strategy = LoopDoubleBufferTransformer::UnrollStrategy::kFullUnroll;
+  }
+  if (unroll_strategy != std::nullopt) {
+    pipeline.AddPass<LoopDoubleBufferTransformer>(*unroll_strategy);
     pipeline.AddPass<TupleSimplifier>();
     pipeline.AddPass<HloDCE>();
   }
diff --git a/third_party/xla/xla/xla.proto b/third_party/xla/xla/xla.proto
index e7055359768113..c49d4a509a2618 100644
--- a/third_party/xla/xla/xla.proto
+++ b/third_party/xla/xla/xla.proto
@@ -668,6 +668,18 @@ message DebugOptions {
   // Enable double buffering for loops.
   bool xla_gpu_enable_while_loop_double_buffering = 248;
 
+  enum WhileLoopUnrolling {
+    WHILE_LOOP_UNROLLING_NO_UNROLL = 0;
+    // Has the same effect as setting
+    // `xla_gpu_enable_while_loop_double_buffering`.
+    WHILE_LOOP_UNROLLING_DOUBLE_BUFFER = 1;
+    // Enables full loop unrolling using the same strategy as `DOUBLE_BUFFER`.
+    WHILE_LOOP_UNROLLING_FULL_UNROLL = 2;
+  }
+
+  // Determine the while loop unrolling scheme.
+  WhileLoopUnrolling xla_gpu_enable_while_loop_unrolling = 294;
+
   // Change the layout of the second triton dot operand to be column major.
   // Only works for (bf16 x bf16) -> bf16.
   bool xla_gpu_ensure_minor_dot_contraction_dims = 249;
@@ -767,7 +779,7 @@ message DebugOptions {
   // Base length to rewrite the reduce window to, no rewrite if set to 0.
   int64 xla_reduce_window_rewrite_base_length = 293;
 
-  // Next id: 294
+  // Next id: 295
 
   // Extra options to pass to the compilation backend (e.g. LLVM); specific
   // interpretation of these values is left to the backend.

From b010326f4d779f1540343faf838ee14480c34588 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Mon, 13 May 2024 07:17:38 -0700
Subject: [PATCH 020/478] Update ops-related pbtxt files.

PiperOrigin-RevId: 633199826
---
 tensorflow/core/ops/ops.pbtxt | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorflow/core/ops/ops.pbtxt b/tensorflow/core/ops/ops.pbtxt
index 53c362ed0ab819..68d78d0329a9ac 100644
--- a/tensorflow/core/ops/ops.pbtxt
+++ b/tensorflow/core/ops/ops.pbtxt
@@ -1,4 +1,4 @@
-go/debugstr    
+go/nodeserialize   
 op {
   name: "Abort"
   attr {

From bd8a712cbdfa268baf2c91bb165d0d9849576389 Mon Sep 17 00:00:00 2001
From: Johannes Reifferscheid <jreiffers@google.com>
Date: Mon, 13 May 2024 07:30:48 -0700
Subject: [PATCH 021/478] Reshape multi-dimensional constants to 1d.

The LLVM lowering doesn't support arbitrary shapes.

PiperOrigin-RevId: 633203497
---
 .../xla/xla/service/gpu/fusions/mlir/lower_tensors.cc      | 7 +++++++
 .../xla/service/gpu/fusions/mlir/tests/lower_tensors.mlir  | 4 ++--
 2 files changed, 9 insertions(+), 2 deletions(-)

diff --git a/third_party/xla/xla/service/gpu/fusions/mlir/lower_tensors.cc b/third_party/xla/xla/service/gpu/fusions/mlir/lower_tensors.cc
index a529e6d86d00e3..5677b45e5342bf 100644
--- a/third_party/xla/xla/service/gpu/fusions/mlir/lower_tensors.cc
+++ b/third_party/xla/xla/service/gpu/fusions/mlir/lower_tensors.cc
@@ -334,6 +334,13 @@ mlir::LLVM::GlobalOp CreateGlobalOp(mlir::Attribute value,
                                     mlir::ModuleOp module, bool is_constant,
                                     int addr_space,
                                     mlir::ImplicitLocOpBuilder& b) {
+  if (auto elements = mlir::dyn_cast_or_null<mlir::DenseElementsAttr>(value)) {
+    // The lowering to LLVM only works for 1d tensors or those with trailing
+    // unit dimensions.
+    value = elements.reshape(mlir::RankedTensorType::get(
+        {elements.getNumElements()}, elements.getElementType()));
+  }
+
   Type element_type = shaped_ty.getElementType();
   // Needed to support complex element type.
   mlir::LLVMTypeConverter converter(b.getContext());
diff --git a/third_party/xla/xla/service/gpu/fusions/mlir/tests/lower_tensors.mlir b/third_party/xla/xla/service/gpu/fusions/mlir/tests/lower_tensors.mlir
index 707cfc6f34196c..d93f3ecd514c24 100644
--- a/third_party/xla/xla/service/gpu/fusions/mlir/tests/lower_tensors.mlir
+++ b/third_party/xla/xla/service/gpu/fusions/mlir/tests/lower_tensors.mlir
@@ -163,8 +163,8 @@ module {
     return %0 : f32
   }
 }
-// CHECK: llvm.mlir.global private constant @global_cst_0(dense<[
-// CHECK-SAME: [1.000000e+00], [2.000000e+00]]> : tensor<2x1xf32>) {addr_space = 0 : i32} : !llvm.array<2 x f32>
+// CHECK: llvm.mlir.global private constant @global_cst_0(dense<
+// CHECK-SAME: [1.000000e+00, 2.000000e+00]> : tensor<2xf32>) {addr_space = 0 : i32} : !llvm.array<2 x f32>
 // CHECK: @extract_from_constant
 // CHECK: %[[ADDR_OF:.*]] = llvm.mlir.addressof @global_cst_0 : !llvm.ptr
 // CHECK: %[[GEP:.*]] = llvm.getelementptr inbounds %[[ADDR_OF]][%{{.*}}] : (!llvm.ptr, i64) -> !llvm.ptr, f32

From ad24c813bcf671fbb42da70efcf3e0ac4fe94a5f Mon Sep 17 00:00:00 2001
From: Ilia Sergachev <isergachev@nvidia.com>
Date: Mon, 13 May 2024 08:20:14 -0700
Subject: [PATCH 022/478] PR #12255: [GPU] Fix FMHA hangs by moving compilation
 to thunk initialization.

Imported from GitHub PR https://github.com/openxla/xla/pull/12255

Applies the same fix as in https://github.com/openxla/xla/pull/12228 to FMHA.
Copybara import of the project:

--
70a42828f86711a0e83a2eb37ee52833e1768187 by Ilia Sergachev <isergachev@nvidia.com>:

[GPU] Fix FMHA hangs by moving compilation to thunk initialization.

Merging this change closes #12255

PiperOrigin-RevId: 633217196
---
 .../xla/service/gpu/gpu_fused_mha_runner.cc   | 89 ++++++++++---------
 .../xla/service/gpu/gpu_fused_mha_runner.h    |  8 ++
 .../service/gpu/runtime/fused_mha_thunk.cc    | 14 +++
 .../xla/service/gpu/runtime/fused_mha_thunk.h |  2 +
 4 files changed, 70 insertions(+), 43 deletions(-)

diff --git a/third_party/xla/xla/service/gpu/gpu_fused_mha_runner.cc b/third_party/xla/xla/service/gpu/gpu_fused_mha_runner.cc
index 61be619b1d46d8..09af1e523ab7f1 100644
--- a/third_party/xla/xla/service/gpu/gpu_fused_mha_runner.cc
+++ b/third_party/xla/xla/service/gpu/gpu_fused_mha_runner.cc
@@ -68,29 +68,13 @@ absl::Status RunFusedMHA(GpufMHAParams params, se::Stream *stream,
     dropout_rate = *params.config->dropout_rate;
   }
 
-  double scale = 1.0;
-  if (params.config->fmha_scale) {
-    scale = *params.config->fmha_scale;
-  }
-
   std::optional<int64_t> seed;
   if (params.config->seed) {
     seed = *params.config->seed;
   }
-  TF_ASSIGN_OR_RETURN(
-      se::dnn::FMHAMaskKind mask_type,
-      GetDNNFmhaMaskKindFromCudnnFmhaMaskKind(params.config->mask_type));
-  se::dnn::FusedMHAOp::Config config{scale,
-                                     params.config->lhs_bmm1,
-                                     params.config->rhs_bmm1,
-                                     params.config->rhs_bmm2,
-                                     params.config->intermediate_lhs_bmm2,
-                                     params.config->output,
-                                     params.config->bias,
-                                     params.config->activation,
-                                     dropout_rate,
-                                     seed,
-                                     mask_type};
+
+  TF_ASSIGN_OR_RETURN(se::dnn::FusedMHAOp::Config config,
+                      params.config->AsDnnFusedMHAOpConfig());
   TF_ASSIGN_OR_RETURN(auto *runner,
                       lazy_runner->GetOrCreateRunner(config, stream));
   return (*runner)(stream, options.profile_result, scratch_memory,
@@ -183,35 +167,13 @@ absl::Status RunFusedMHABackward(
     dropout_rate = *params.config->dropout_rate;
   }
 
-  double scale = 1.0;
-  if (params.config->fmha_scale) {
-    scale = *params.config->fmha_scale;
-  }
-
   std::optional<int64_t> seed;
   if (params.config->seed) {
     seed = *params.config->seed;
   }
 
-  TF_ASSIGN_OR_RETURN(
-      se::dnn::FMHAMaskKind mask_type,
-      GetDNNFmhaMaskKindFromCudnnFmhaMaskKind(params.config->mask_type));
-  se::dnn::FusedMHABackwardOp::Config config{scale,
-                                             params.config->bmm1_grad_gemm1_rhs,
-                                             params.config->bmm1_grad_gemm2_rhs,
-                                             params.config->bmm2_grad_gemm1_lhs,
-                                             params.config->bmm2_grad_gemm2_rhs,
-                                             params.config->d_output,
-                                             params.config->d_bmm1_lhs,
-                                             params.config->d_bmm1_rhs,
-                                             params.config->d_bmm2_rhs,
-                                             params.config->d_s,
-                                             params.config->d_bias,
-                                             params.config->fwd_output,
-                                             params.config->bias,
-                                             dropout_rate,
-                                             seed,
-                                             mask_type};
+  TF_ASSIGN_OR_RETURN(se::dnn::FusedMHABackwardOp::Config config,
+                      params.config->AsDnnFusedMHABackwardOpConfig());
   TF_ASSIGN_OR_RETURN(auto *runner,
                       lazy_runner->GetOrCreateRunner(config, stream));
   // TODO: pass in real softmax_sum, dQ_accum, fwd_output
@@ -404,6 +366,21 @@ absl::Status RunGpuFMHABackwardImpl(const GpufMHABackwardParams &params,
   return config;
 }
 
+absl::StatusOr<se::dnn::FusedMHAOp::Config>
+GpufMHAConfig::AsDnnFusedMHAOpConfig() const {
+  double scale = 1.0;
+  if (fmha_scale.has_value()) {
+    scale = *fmha_scale;
+  }
+  TF_ASSIGN_OR_RETURN(se::dnn::FMHAMaskKind mask_type,
+                      GetDNNFmhaMaskKindFromCudnnFmhaMaskKind(mask_type));
+
+  return se::dnn::FusedMHAOp::Config{
+      scale,    lhs_bmm1, rhs_bmm1,   rhs_bmm2,     intermediate_lhs_bmm2,
+      output,   bias,     activation, dropout_rate, seed,
+      mask_type};
+}
+
 /*static*/ absl::StatusOr<GpufMHABackwardConfig> GpufMHABackwardConfig::For(
     const GpufMHABackwardDescriptor &desc) {
   // Get shapes from desc.
@@ -546,6 +523,32 @@ absl::Status RunGpuFMHABackwardImpl(const GpufMHABackwardParams &params,
   return config;
 }
 
+absl::StatusOr<se::dnn::FusedMHABackwardOp::Config>
+GpufMHABackwardConfig::AsDnnFusedMHABackwardOpConfig() const {
+  double scale = 1.0;
+  if (fmha_scale.has_value()) {
+    scale = *fmha_scale;
+  }
+  TF_ASSIGN_OR_RETURN(se::dnn::FMHAMaskKind mask_type,
+                      GetDNNFmhaMaskKindFromCudnnFmhaMaskKind(mask_type));
+  return se::dnn::FusedMHABackwardOp::Config{scale,
+                                             bmm1_grad_gemm1_rhs,
+                                             bmm1_grad_gemm2_rhs,
+                                             bmm2_grad_gemm1_lhs,
+                                             bmm2_grad_gemm2_rhs,
+                                             d_output,
+                                             d_bmm1_lhs,
+                                             d_bmm1_rhs,
+                                             d_bmm2_rhs,
+                                             d_s,
+                                             d_bias,
+                                             fwd_output,
+                                             bias,
+                                             dropout_rate,
+                                             seed,
+                                             mask_type};
+}
+
 /*static*/ absl::StatusOr<GpufMHAParams> GpufMHAParams::For(
     const GpufMHAConfig &config, se::DeviceMemoryBase lhs_bmm1_buffer,
     se::DeviceMemoryBase rhs_bmm1_buffer, se::DeviceMemoryBase rhs_bmm2_buffer,
diff --git a/third_party/xla/xla/service/gpu/gpu_fused_mha_runner.h b/third_party/xla/xla/service/gpu/gpu_fused_mha_runner.h
index 6538d542984325..7ca35805be251e 100644
--- a/third_party/xla/xla/service/gpu/gpu_fused_mha_runner.h
+++ b/third_party/xla/xla/service/gpu/gpu_fused_mha_runner.h
@@ -101,10 +101,14 @@ struct GpufMHABackwardDescriptor {
   std::optional<Shape> d_bias_shape;
   std::optional<Shape> bias_shape;
 };
+
 // Structure to describe static properties of a GPU fused Multi-Headed
 // Attention.
 struct GpufMHAConfig {
   static absl::StatusOr<GpufMHAConfig> For(const GpufMHADescriptor& fmha_desc);
+
+  absl::StatusOr<se::dnn::FusedMHAOp::Config> AsDnnFusedMHAOpConfig() const;
+
   PrimitiveType
       input_type;  // Capture the primitive type of one of the inputs of BMM1
   PrimitiveType output_type;
@@ -133,6 +137,10 @@ struct GpufMHAConfig {
 struct GpufMHABackwardConfig {
   static absl::StatusOr<GpufMHABackwardConfig> For(
       const GpufMHABackwardDescriptor& fmha_desc);
+
+  absl::StatusOr<se::dnn::FusedMHABackwardOp::Config>
+  AsDnnFusedMHABackwardOpConfig() const;
+
   PrimitiveType
       input_type;  // Capture the primitive type of one of the inputs of BMM1
   PrimitiveType output_type;
diff --git a/third_party/xla/xla/service/gpu/runtime/fused_mha_thunk.cc b/third_party/xla/xla/service/gpu/runtime/fused_mha_thunk.cc
index efa90bfc9feae3..41613f0121e65f 100644
--- a/third_party/xla/xla/service/gpu/runtime/fused_mha_thunk.cc
+++ b/third_party/xla/xla/service/gpu/runtime/fused_mha_thunk.cc
@@ -65,6 +65,13 @@ std::optional<se::DeviceMemoryBase> AssignBufferIfNotNull(
              : std::nullopt;
 }
 
+absl::Status FusedMHAThunk::Initialize(const InitializeParams& params) {
+  se::dnn::LazyOpRunner<se::dnn::FusedMHAOp>* lazy_runner =
+      GetOrCreateRunner(params.stream).AsFusedMHARunner();
+  TF_ASSIGN_OR_RETURN(auto config, config_.AsDnnFusedMHAOpConfig());
+  return lazy_runner->GetOrCreateRunner(config, params.stream).status();
+}
+
 absl::Status FusedMHAThunk::ExecuteOnStream(const ExecuteParams& params) {
   const auto& buffer_allocations = *params.buffer_allocations;
   se::DeviceMemoryBase lhs_bmm1_buffer =
@@ -143,6 +150,13 @@ FusedMHABackwardThunk::GetOrCreateRunner(
   return *it->second;
 }
 
+absl::Status FusedMHABackwardThunk::Initialize(const InitializeParams& params) {
+  se::dnn::LazyOpRunner<se::dnn::FusedMHABackwardOp>* lazy_runner =
+      GetOrCreateRunner(params.stream).AsFusedMHABackwardRunner();
+  TF_ASSIGN_OR_RETURN(auto config, config_.AsDnnFusedMHABackwardOpConfig());
+  return lazy_runner->GetOrCreateRunner(config, params.stream).status();
+}
+
 absl::Status FusedMHABackwardThunk::ExecuteOnStream(
     const ExecuteParams& params) {
   const auto& buffer_allocations = *params.buffer_allocations;
diff --git a/third_party/xla/xla/service/gpu/runtime/fused_mha_thunk.h b/third_party/xla/xla/service/gpu/runtime/fused_mha_thunk.h
index 32bfdd0ecd19f4..bf9cff354b5027 100644
--- a/third_party/xla/xla/service/gpu/runtime/fused_mha_thunk.h
+++ b/third_party/xla/xla/service/gpu/runtime/fused_mha_thunk.h
@@ -53,6 +53,7 @@ class FusedMHAThunk : public Thunk {
   FusedMHAThunk(const FusedMHAThunk&) = delete;
   FusedMHAThunk& operator=(const FusedMHAThunk&) = delete;
 
+  absl::Status Initialize(const InitializeParams& params) override;
   absl::Status ExecuteOnStream(const ExecuteParams& params) override;
 
  private:
@@ -101,6 +102,7 @@ class FusedMHABackwardThunk : public Thunk {
   FusedMHABackwardThunk(const FusedMHABackwardThunk&) = delete;
   FusedMHABackwardThunk& operator=(const FusedMHABackwardThunk&) = delete;
 
+  absl::Status Initialize(const InitializeParams& params) override;
   absl::Status ExecuteOnStream(const ExecuteParams& params) override;
 
  private:

From 1135035e5ee99295a1d675d179e4bd433e4192f9 Mon Sep 17 00:00:00 2001
From: Trevor Morris <tmorris@nvidia.com>
Date: Mon, 13 May 2024 08:31:43 -0700
Subject: [PATCH 023/478] PR #12228: [GPU] Fix hang with cudnn layer norm by
 moving build phase to Initialize()

Imported from GitHub PR https://github.com/openxla/xla/pull/12228

The first time that a NormThunk is executed, it will build a cudnn execution plan. This build step can hang if a NCCL collective is running at the same time. To fix this, I've moved the build step to take place during thunk initialization. We only observe this hang when using cudnn 9.

Here's a backtrace from the hang that will be fixed:
```
Thread 585 (Thread 0x7fb9391ff640 (LWP 41364) "main.py"):
#0  0x00007fd3d17cffd9 in ?? () from /lib/x86_64-linux-gnu/libc.so.6
#1  0x00007fd3d17da24f in pthread_rwlock_wrlock () from /lib/x86_64-linux-gnu/libc.so.6
#2  0x00007fd070967dfe in ?? () from /lib/x86_64-linux-gnu/libcuda.so.1
#3  0x00007fd0709c928a in ?? () from /lib/x86_64-linux-gnu/libcuda.so.1
#4  0x00007f1970d76102 in ?? () from /lib/x86_64-linux-gnu/libcudnn_engines_precompiled.so.9.1.0
#5  0x00007f1970f2c999 in ?? () from /lib/x86_64-linux-gnu/libcudnn_engines_precompiled.so.9.1.0
#6  0x00007f1970a7d4ab in ?? () from /lib/x86_64-linux-gnu/libcudnn_engines_precompiled.so.9.1.0
#7  0x00007f1970d0a9cb in ?? () from /lib/x86_64-linux-gnu/libcudnn_engines_precompiled.so.9.1.0
#8  0x00007fce60b2a98c in cudnn::backend::ExecutionPlan::finalize_internal() () from /lib/x86_64-linux-gnu/libcudnn_graph.so.9.1.0
#9  0x00007fce60aefbb1 in cudnn::backend::Descriptor::finalize() () from /lib/x86_64-linux-gnu/libcudnn_graph.so.9.1.0
#10 0x00007fce60b15bec in cudnnBackendFinalize () from /lib/x86_64-linux-gnu/libcudnn_graph.so.9.1.0
#11 0x00007fd2521b8f39 in cudnn_frontend::ExecutionPlanBuilder_v8::build() () from /usr/local/lib/python3.10/dist-packages/jaxlib/xla_extension.so
#12 0x00007fd2521734ba in stream_executor::gpu::(anonymous namespace)::GetExecPlanFromHeuristics(cudnn_frontend::OperationGraph_v8&&, stream_executor::gpu::(anonymous namespace)::CudnnHandle const&, bool) () from /usr/local/lib/python3.10/dist-packages/jaxlib/xla_extension.so
#13 0x00007fd25216ff9b in stream_executor::gpu::CudnnSupport::NormRunnerFromDesc(stream_executor::Stream*, stream_executor::dnn::AlgorithmDesc const&, stream_executor::dnn::NormKind, double, stream_executor::dnn::TensorDescriptor const&, stream_executor::dnn::TensorDescriptor const&, stream_executor::dnn::TensorDescriptor const&, std::optional<stream_executor::dnn::TensorDescriptor>, std::optional<stream_executor::dnn::TensorDescriptor>, std::optional<stream_executor::dnn::TensorDescriptor>, std::optional<stream_executor::dnn::TensorDescriptor>, std::optional<stream_executor::dnn::TensorDescriptor>, std::optional<stream_executor::dnn::TensorDescriptor>) () from /usr/local/lib/python3.10/dist-packages/jaxlib/xla_extension.so
#14 0x00007fd24e36b88b in stream_executor::dnn::NormOp::RunnerFromAlgorithmDesc(stream_executor::dnn::AlgorithmDesc const&, stream_executor::dnn::NormOp::Config, stream_executor::Stream*) () from /usr/local/lib/python3.10/dist-packages/jaxlib/xla_extension.so
#15 0x00007fd24e36ae37 in stream_executor::dnn::LazyOpRunner<stream_executor::dnn::NormOp>::GetOrCreateRunner(stream_executor::dnn::NormOp::Config, stream_executor::Stream*)::{lambda()#1}::operator()() const () from /usr/local/lib/python3.10/dist-packages/jaxlib/xla_extension.so
#16 0x00007fd24e36adbc in void absl::lts_20230802::base_internal::CallOnceImpl<stream_executor::dnn::LazyOpRunner<stream_executor::dnn::NormOp>::GetOrCreateRunner(stream_executor::dnn::NormOp::Config, stream_executor::Stream*)::{lambda()#1}>(std::atomic<unsigned int>*, absl::lts_20230802::base_internal::SchedulingMode, stream_executor::dnn::LazyOpRunner<stream_executor::dnn::NormOp>::GetOrCreateRunner(stream_executor::dnn::NormOp::Config, stream_executor::Stream*)::{lambda()#1}&&) () from /usr/local/lib/python3.10/dist-packages/jaxlib/xla_extension.so
#17 0x00007fd24e36a9bd in stream_executor::dnn::LazyOpRunner<stream_executor::dnn::NormOp>::GetOrCreateRunner(stream_executor::dnn::NormOp::Config, stream_executor::Stream*) () from /usr/local/lib/python3.10/dist-packages/jaxlib/xla_extension.so
#18 0x00007fd24e369d29 in xla::gpu::RunGpuNorm(xla::gpu::GpuNormConfig const&, stream_executor::DeviceMemoryBase const&, stream_executor::DeviceMemoryBase const&, stream_executor::DeviceMemoryBase const&, std::optional<stream_executor::DeviceMemoryBase>, std::optional<stream_executor::DeviceMemoryBase>, std::optional<stream_executor::DeviceMemoryBase>, std::optional<stream_executor::DeviceMemoryBase>, std::optional<stream_executor::DeviceMemoryBase>, std::optional<stream_executor::DeviceMemoryBase>, stream_executor::DeviceMemoryBase const&, stream_executor::Stream*, xla::gpu::RunNormOptions) () from /usr/local/lib/python3.10/dist-packages/jaxlib/xla_extension.so
#19 0x00007fd24e368be6 in xla::gpu::NormThunk::ExecuteOnStream(xla::gpu::Thunk::ExecuteParams const&) () from /usr/local/lib/python3.10/dist-packages/jaxlib/xla_extension.so
```
Copybara import of the project:

--
f53533087ba1ddcf65ad7cc6268ee89de4690d15 by Trevor Morris <tmorris@nvidia.com>:

Fix hang with cudnn layer norm by moving cudnn init to Initialize()

Merging this change closes #12228

PiperOrigin-RevId: 633220207
---
 .../xla/xla/service/gpu/gpu_norm_runner.cc     | 18 ++----------------
 .../xla/xla/service/gpu/gpu_norm_runner.h      | 16 ++++++++++++++++
 .../xla/xla/service/gpu/runtime/norm_thunk.cc  |  9 +++++++++
 .../xla/xla/service/gpu/runtime/norm_thunk.h   |  1 +
 4 files changed, 28 insertions(+), 16 deletions(-)

diff --git a/third_party/xla/xla/service/gpu/gpu_norm_runner.cc b/third_party/xla/xla/service/gpu/gpu_norm_runner.cc
index 5abb58af4e019d..9170de6016e60a 100644
--- a/third_party/xla/xla/service/gpu/gpu_norm_runner.cc
+++ b/third_party/xla/xla/service/gpu/gpu_norm_runner.cc
@@ -44,22 +44,8 @@ absl::Status RunGpuNorm(const gpu::GpuNormConfig& config,
                         se::Stream* stream, RunNormOptions options) {
   se::dnn::LazyOpRunner<se::dnn::NormOp>* lazy_runner =
       options.norm_runner->AsNormRunner();
-  std::optional<se::dnn::LazyOpRunner<se::dnn::NormOp>> local_runner;
-
-  TF_ASSIGN_OR_RETURN(se::dnn::NormKind kind,
-                      GetDNNNormKindFromCudnnNormKind(config.kind));
-
-  se::dnn::NormOp::Config ln_config{kind,
-                                    config.epsilon,
-                                    config.x_descriptor,
-                                    config.scale_descriptor,
-                                    config.y_or_dx_descriptor,
-                                    config.bias_descriptor,
-                                    config.dy_descriptor,
-                                    config.expectation_descriptor,
-                                    config.norm_factor_descriptor,
-                                    config.dscale_descriptor,
-                                    config.dbias_descriptor};
+  TF_ASSIGN_OR_RETURN(se::dnn::NormOp::Config ln_config,
+                      config.AsDnnNormOpConfig());
   TF_ASSIGN_OR_RETURN(auto* runner,
                       lazy_runner->GetOrCreateRunner(ln_config, stream));
 
diff --git a/third_party/xla/xla/service/gpu/gpu_norm_runner.h b/third_party/xla/xla/service/gpu/gpu_norm_runner.h
index 854e3c0892050c..8461671e86d037 100644
--- a/third_party/xla/xla/service/gpu/gpu_norm_runner.h
+++ b/third_party/xla/xla/service/gpu/gpu_norm_runner.h
@@ -118,6 +118,22 @@ struct GpuNormConfig {
     return config;
   }
 
+  absl::StatusOr<se::dnn::NormOp::Config> AsDnnNormOpConfig() const {
+    TF_ASSIGN_OR_RETURN(se::dnn::NormKind norm_kind,
+                        GetDNNNormKindFromCudnnNormKind(kind));
+    return se::dnn::NormOp::Config{norm_kind,
+                                   epsilon,
+                                   x_descriptor,
+                                   scale_descriptor,
+                                   y_or_dx_descriptor,
+                                   bias_descriptor,
+                                   dy_descriptor,
+                                   expectation_descriptor,
+                                   norm_factor_descriptor,
+                                   dscale_descriptor,
+                                   dbias_descriptor};
+  }
+
   double epsilon;
   CudnnNormKind kind;
   se::dnn::AlgorithmDesc algorithm;
diff --git a/third_party/xla/xla/service/gpu/runtime/norm_thunk.cc b/third_party/xla/xla/service/gpu/runtime/norm_thunk.cc
index d3862f7bfeac74..71c0744686e402 100644
--- a/third_party/xla/xla/service/gpu/runtime/norm_thunk.cc
+++ b/third_party/xla/xla/service/gpu/runtime/norm_thunk.cc
@@ -106,5 +106,14 @@ absl::Status NormThunk::ExecuteOnStream(const ExecuteParams& params) {
   return absl::OkStatus();
 }
 
+absl::Status NormThunk::Initialize(const InitializeParams& params) {
+  // Create the runner at initialization time to avoid hangs if we try to build
+  // the execution plan while a NCCL collective is running.
+  se::dnn::LazyOpRunner<se::dnn::NormOp>* lazy_runner =
+      GetOrCreateRunner(params.stream).AsNormRunner();
+  TF_ASSIGN_OR_RETURN(auto ln_config, config_.AsDnnNormOpConfig());
+  return lazy_runner->GetOrCreateRunner(ln_config, params.stream).status();
+}
+
 }  // namespace gpu
 }  // namespace xla
diff --git a/third_party/xla/xla/service/gpu/runtime/norm_thunk.h b/third_party/xla/xla/service/gpu/runtime/norm_thunk.h
index 602d504175fb3d..eca5deca3a68b7 100644
--- a/third_party/xla/xla/service/gpu/runtime/norm_thunk.h
+++ b/third_party/xla/xla/service/gpu/runtime/norm_thunk.h
@@ -49,6 +49,7 @@ class NormThunk : public Thunk {
   NormThunk& operator=(const NormThunk&) = delete;
 
   absl::Status ExecuteOnStream(const ExecuteParams& params) override;
+  absl::Status Initialize(const InitializeParams& params) override;
 
  private:
   BufferAllocation::Slice x_buffer_;

From f83f81c2cd4e57e1d6588ced9bbcb2cab011bd26 Mon Sep 17 00:00:00 2001
From: Vladyslav Tsilytskyi <tsilytskyi@google.com>
Date: Mon, 13 May 2024 08:44:40 -0700
Subject: [PATCH 024/478] Reverts 28bf7478c8481648dc3900e0ab52a5e0887f5145

PiperOrigin-RevId: 633223628
---
 .../xla/xla/stream_executor/host/BUILD        |  63 ++-
 .../host/host_execution_engine.cc             | 390 ++++++++++++++++++
 .../host/host_execution_engine.h              | 157 +++++++
 .../xla/stream_executor/host/host_executor.cc |  54 ++-
 .../xla/stream_executor/host/host_executor.h  |  11 +-
 .../xla/stream_executor/host/host_kernel.cc   |   5 +-
 .../xla/stream_executor/host/host_kernel.h    |  37 +-
 .../stream_executor/host/host_kernel_test.cc  |  67 +++
 8 files changed, 768 insertions(+), 16 deletions(-)
 create mode 100644 third_party/xla/xla/stream_executor/host/host_execution_engine.cc
 create mode 100644 third_party/xla/xla/stream_executor/host/host_execution_engine.h

diff --git a/third_party/xla/xla/stream_executor/host/BUILD b/third_party/xla/xla/stream_executor/host/BUILD
index 43af63f32502bf..da6b0d35ba7d5b 100644
--- a/third_party/xla/xla/stream_executor/host/BUILD
+++ b/third_party/xla/xla/stream_executor/host/BUILD
@@ -1,6 +1,14 @@
 # Description:
 #   Host-platform specific StreamExecutor support code.
 
+load(
+    "@local_tsl//tsl/platform:build_config_root.bzl",
+    "if_llvm_aarch32_available",
+    "if_llvm_aarch64_available",
+    "if_llvm_powerpc_available",
+    "if_llvm_system_z_available",
+    "if_llvm_x86_available",
+)
 load("@local_tsl//tsl/platform:rules_cc.bzl", "cc_library")
 load("//xla:xla.bzl", "xla_cc_test")
 load("//xla/stream_executor:build_defs.bzl", "stream_executor_friends")
@@ -81,15 +89,60 @@ cc_library(
     hdrs = ["host_kernel_c_api.h"],
 )
 
+cc_library(
+    name = "host_execution_engine",
+    srcs = ["host_execution_engine.cc"],
+    hdrs = ["host_execution_engine.h"],
+    deps = [
+        ":host_kernel_c_api",
+        "@com_google_absl//absl/log:check",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/status:statusor",
+        "@com_google_absl//absl/strings",
+        "@com_google_absl//absl/strings:str_format",
+        "@com_google_absl//absl/synchronization",
+        "@com_google_absl//absl/types:span",
+        "@llvm-project//llvm:Analysis",
+        "@llvm-project//llvm:AsmParser",
+        "@llvm-project//llvm:Core",
+        "@llvm-project//llvm:ExecutionEngine",
+        "@llvm-project//llvm:JITLink",
+        "@llvm-project//llvm:OrcJIT",
+        "@llvm-project//llvm:OrcShared",
+        "@llvm-project//llvm:Passes",
+        "@llvm-project//llvm:Support",
+        "@llvm-project//llvm:Target",
+        "@llvm-project//llvm:TargetParser",
+        "@llvm-project//llvm:TransformUtils",
+        "@llvm-project//llvm:ir_headers",
+        "@local_tsl//tsl/platform:statusor",
+    ] + if_llvm_aarch32_available([
+        "@llvm-project//llvm:ARMAsmParser",
+        "@llvm-project//llvm:ARMCodeGen",
+    ]) + if_llvm_aarch64_available([
+        "@llvm-project//llvm:AArch64AsmParser",
+        "@llvm-project//llvm:AArch64CodeGen",
+    ]) + if_llvm_powerpc_available([
+        "@llvm-project//llvm:PowerPCAsmParser",
+        "@llvm-project//llvm:PowerPCCodeGen",
+    ]) + if_llvm_system_z_available([
+        "@llvm-project//llvm:SystemZAsmParser",
+        "@llvm-project//llvm:SystemZCodeGen",
+    ]) + if_llvm_x86_available([
+        "@llvm-project//llvm:X86AsmParser",
+        "@llvm-project//llvm:X86CodeGen",
+    ]),
+)
+
 cc_library(
     name = "host_kernel",
     srcs = ["host_kernel.cc"],
     hdrs = ["host_kernel.h"],
     deps = [
+        ":host_execution_engine",
         ":host_kernel_c_api",
         "//xla/stream_executor",
         "//xla/stream_executor:device_memory",
-        "//xla/stream_executor:stream_executor_interface",
         "@com_google_absl//absl/status",
         "@com_google_absl//absl/status:statusor",
         "@com_google_absl//absl/types:span",
@@ -102,9 +155,12 @@ xla_cc_test(
     deps = [
         ":host_kernel",
         ":host_kernel_c_api",
+        ":host_platform",
         "//xla/stream_executor",
         "//xla/stream_executor:device_memory",
+        "@com_google_absl//absl/types:span",
         "@local_tsl//tsl/lib/core:status_test_util",
+        "@local_tsl//tsl/platform:statusor",
         "@local_tsl//tsl/platform:test",
         "@local_tsl//tsl/platform:test_main",
     ],
@@ -119,17 +175,22 @@ cc_library(
         "host_executor.h",
     ],
     deps = [
+        ":host_execution_engine",
+        ":host_kernel",
         ":host_stream",
         "//xla/stream_executor",
         "//xla/stream_executor:event_interface",
         "//xla/stream_executor:stream_executor_interface",
         "@com_google_absl//absl/functional:any_invocable",
+        "@com_google_absl//absl/log",
         "@com_google_absl//absl/log:check",
         "@com_google_absl//absl/status",
         "@com_google_absl//absl/status:statusor",
         "@com_google_absl//absl/strings",
         "@com_google_absl//absl/synchronization",
+        "@com_google_absl//absl/types:span",
         "@local_tsl//tsl/platform:platform_port",
+        "@local_tsl//tsl/platform:statusor",
         "@local_tsl//tsl/platform/profile_utils:profile_utils_cpu_utils",
     ],
     alwayslink = True,
diff --git a/third_party/xla/xla/stream_executor/host/host_execution_engine.cc b/third_party/xla/xla/stream_executor/host/host_execution_engine.cc
new file mode 100644
index 00000000000000..f48203bfb3166c
--- /dev/null
+++ b/third_party/xla/xla/stream_executor/host/host_execution_engine.cc
@@ -0,0 +1,390 @@
+/* Copyright 2024 The OpenXLA Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+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.
+==============================================================================*/
+
+#include "xla/stream_executor/host/host_execution_engine.h"
+
+#include <functional>
+#include <memory>
+#include <optional>
+#include <string>
+#include <utility>
+#include <vector>
+
+#include "absl/log/check.h"
+#include "absl/status/status.h"
+#include "absl/status/statusor.h"
+#include "absl/strings/str_format.h"
+#include "absl/strings/string_view.h"
+#include "absl/synchronization/mutex.h"
+#include "absl/types/span.h"
+#include "llvm/ADT/StringRef.h"
+#include "llvm/Analysis/CGSCCPassManager.h"
+#include "llvm/Analysis/LoopAnalysisManager.h"
+#include "llvm/AsmParser/Parser.h"
+#include "llvm/ExecutionEngine/JITEventListener.h"
+#include "llvm/ExecutionEngine/ObjectCache.h"
+#include "llvm/ExecutionEngine/Orc/CompileUtils.h"
+#include "llvm/ExecutionEngine/Orc/Core.h"
+#include "llvm/ExecutionEngine/Orc/IRCompileLayer.h"
+#include "llvm/ExecutionEngine/Orc/JITTargetMachineBuilder.h"
+#include "llvm/ExecutionEngine/Orc/LLJIT.h"
+#include "llvm/ExecutionEngine/Orc/Mangling.h"
+#include "llvm/ExecutionEngine/Orc/RTDyldObjectLinkingLayer.h"
+#include "llvm/ExecutionEngine/Orc/Shared/ExecutorAddress.h"
+#include "llvm/ExecutionEngine/Orc/TaskDispatch.h"
+#include "llvm/ExecutionEngine/Orc/ThreadSafeModule.h"
+#include "llvm/ExecutionEngine/SectionMemoryManager.h"
+#include "llvm/IR/BasicBlock.h"
+#include "llvm/IR/DerivedTypes.h"
+#include "llvm/IR/Instructions.h"
+#include "llvm/IR/LLVMContext.h"
+#include "llvm/IR/PassManager.h"
+#include "llvm/Passes/OptimizationLevel.h"
+#include "llvm/Passes/PassBuilder.h"
+#include "llvm/Support/Error.h"
+#include "llvm/Support/MemoryBuffer.h"
+#include "llvm/Support/SourceMgr.h"
+#include "llvm/Support/TargetSelect.h"
+#include "llvm/Support/raw_ostream.h"
+#include "llvm/Target/TargetMachine.h"
+#include "llvm/TargetParser/Triple.h"
+#include "llvm/Transforms/Utils/Cloning.h"
+#include "tsl/platform/statusor.h"
+
+namespace stream_executor {
+namespace host {
+
+using absl::InternalError;
+using absl::StatusOr;
+using absl::StrFormat;
+
+using llvm::Expected;
+using llvm::MemoryBuffer;
+using llvm::SectionMemoryManager;
+using llvm::Triple;
+
+using llvm::orc::ExecutionSession;
+using llvm::orc::ExecutorAddr;
+using llvm::orc::InPlaceTaskDispatcher;
+using llvm::orc::IRCompileLayer;
+using llvm::orc::JITTargetMachineBuilder;
+using llvm::orc::RTDyldObjectLinkingLayer;
+using llvm::orc::SelfExecutorProcessControl;
+using llvm::orc::SimpleCompiler;
+using llvm::orc::SymbolMap;
+using llvm::orc::ThreadSafeModule;
+
+namespace {
+
+// This compiler keeps weak pointers to the TargetMachine and the ObjectCache.
+//
+// This allows releasing the memory of those objects, even though the LLJIT
+// keeps the compiler alive.
+//
+// We wrote this class based on the code of llvm::orc::ConcurrentIRCompiler.
+class WeakCompiler : public IRCompileLayer::IRCompiler {
+ public:
+  static llvm::orc::IRSymbolMapper::ManglingOptions
+  IrManglingOptionsForWeakTargetMachine(
+      std::weak_ptr<llvm::TargetMachine> weak_target_machine) {
+    std::shared_ptr<llvm::TargetMachine> target_machine =
+        weak_target_machine.lock();
+    CHECK(target_machine != nullptr)
+        << "Compiler should not be used after the TargetMachine is destroyed.";
+
+    return llvm::orc::irManglingOptionsFromTargetOptions(
+        target_machine->Options);
+  }
+
+  // It's not recommended to allocate the parameters with std::make_shared,
+  // because that would allocate the object and the control block in one
+  // allocation, so the weak_ptr would keep alive the memory of the object as
+  // well.
+  explicit WeakCompiler(std::weak_ptr<llvm::TargetMachine> weak_target_machine)
+      : IRCompiler(IrManglingOptionsForWeakTargetMachine(weak_target_machine)),
+        weak_target_machine_(std::move(weak_target_machine)) {}
+
+  Expected<std::unique_ptr<MemoryBuffer>> operator()(
+      llvm::Module &module) override {
+    std::shared_ptr<llvm::TargetMachine> target_machine =
+        weak_target_machine_.lock();
+    CHECK(target_machine != nullptr)
+        << "Compiler should not be used after the TargetMachine is destroyed.";
+
+    SimpleCompiler compiler(*target_machine);
+    return compiler(module);
+  }
+
+ private:
+  std::weak_ptr<llvm::TargetMachine> weak_target_machine_;
+};
+
+}  // namespace
+
+RuntimeExecutionEngine::RuntimeExecutionEngine(bool enable_gdb_listener,
+                                               bool enable_perf_listener) {
+  if (enable_gdb_listener)
+    gdb_listener_ = llvm::JITEventListener::createGDBRegistrationListener();
+  if (enable_perf_listener)
+    perf_listener_ = llvm::JITEventListener::createPerfJITEventListener();
+}
+
+std::unique_ptr<MemoryBuffer> RuntimeExecutionEngine::obj_file() const {
+  return obj_file_ ? MemoryBuffer::getMemBuffer(obj_file_->getMemBufferRef())
+                   : nullptr;
+}
+
+// -------------------------------------------------------------------------- //
+
+namespace {
+using llvm::DenseMap;
+
+// Intercept object compilation to save the object file corresponding to the
+// XLA executable in the execution engine.
+class ExecutionEngineObjectCache : public llvm::ObjectCache {
+ public:
+  void notifyObjectCompiled(const llvm::Module *m,
+                            llvm::MemoryBufferRef objBuffer) override;
+
+  std::unique_ptr<llvm::MemoryBuffer> getObject(const llvm::Module *m) override;
+
+  // Transfer memory buffer from the cache to the caller.
+  std::unique_ptr<llvm::MemoryBuffer> stealObject(const llvm::Module *m);
+
+ private:
+  DenseMap<const llvm::Module *, std::unique_ptr<llvm::MemoryBuffer>> objs_;
+};
+}  // namespace
+
+void ExecutionEngineObjectCache::notifyObjectCompiled(
+    const llvm::Module *m, llvm::MemoryBufferRef objBuffer) {
+  objs_[m] = llvm::MemoryBuffer::getMemBufferCopy(
+      objBuffer.getBuffer(), objBuffer.getBufferIdentifier());
+}
+
+std::unique_ptr<llvm::MemoryBuffer> ExecutionEngineObjectCache::getObject(
+    const llvm::Module *m) {
+  auto it = objs_.find(m);
+  if (it == objs_.end()) return nullptr;
+  return llvm::MemoryBuffer::getMemBuffer(it->second->getMemBufferRef());
+}
+
+std::unique_ptr<llvm::MemoryBuffer> ExecutionEngineObjectCache::stealObject(
+    const llvm::Module *m) {
+  auto it = objs_.find(m);
+  if (it == objs_.end()) return nullptr;
+  return std::move(it->second);
+}
+
+// -------------------------------------------------------------------------- //
+
+// llvm_ir::DumpToString() is not used here, because we don't want to add too
+// many dependencies to the runtime.
+static std::string ToString(const llvm::Error &err) {
+  std::string str;
+  llvm::raw_string_ostream(str) << err;
+  return str;
+}
+
+/*static*/ StatusOr<std::unique_ptr<RuntimeExecutionEngine>>
+RuntimeExecutionEngine::CreateFromModule(
+    std::unique_ptr<llvm::LLVMContext> ctx,
+    std::unique_ptr<llvm::Module> module, JitOptions options,
+    absl::Span<const std::string_view> exported) {
+  auto engine =
+      std::unique_ptr<RuntimeExecutionEngine>(new RuntimeExecutionEngine(
+          options.enable_gdb_listener, options.enable_perf_listener));
+
+  // We'll need module pointer later to lookup object file in the cache.
+  llvm::Module *module_ptr = module.get();
+
+  // Set up the target machine details.
+  if (!options.target_machine)
+    return InternalError("target machine was not provided");
+  module->setDataLayout(options.target_machine->createDataLayout());
+  module->setTargetTriple(options.target_machine->getTargetTriple().str());
+
+  // Run an optimization pipeline over the LLVM module (alway run with default
+  // opt level independent of the options).
+  //
+  // TODO(ezhulenev): We should have out own optimizing transformer pipelines
+  // for different Xla backends, e.g. there is absolutely no need to run
+  // SLV vectorizer for Xla Gpi host side executable.
+  auto transformer =
+      options.make_optimizing_transformer(options.target_machine.get());
+  if (auto err = transformer(module_ptr))
+    return InternalError(
+        StrFormat("failed to run optimization pipeline: %s", ToString(err)));
+
+  // Callback to create the object layer with a user-provided section memory
+  // mapper and JIT event listeners.
+  auto obj_layer_creator = [&](ExecutionSession &session, const Triple &tt) {
+    auto obj_layer = std::make_unique<RTDyldObjectLinkingLayer>(
+        session, [section_memory_mapper = options.section_memory_mapper]() {
+          return std::make_unique<SectionMemoryManager>(section_memory_mapper);
+        });
+
+    // Register JIT event listeners if they are enabled.
+    if (engine->gdb_listener_)
+      obj_layer->registerJITEventListener(*engine->gdb_listener_);
+    if (engine->perf_listener_)
+      obj_layer->registerJITEventListener(*engine->perf_listener_);
+
+    return obj_layer;
+  };
+
+  // Callback to compile IR module on demand.
+  auto compile_function_creator =
+      [weak_target_machine = std::weak_ptr<llvm::TargetMachine>(
+           options.target_machine)](JITTargetMachineBuilder)
+      -> Expected<std::unique_ptr<IRCompileLayer::IRCompiler>> {
+    return std::make_unique<WeakCompiler>(weak_target_machine);
+  };
+
+  // Use in-process executor process control with in-place task dispatcher.
+  auto executorProcessControl = SelfExecutorProcessControl::Create(
+      nullptr, std::make_unique<InPlaceTaskDispatcher>());
+
+  if (auto err = executorProcessControl.takeError())
+    return InternalError(StrFormat(
+        "failed to create executor process control: %s", ToString(err)));
+
+  // TODO(b/286475799): Concurrent compilation leads to spurious memory
+  // corruptions and segfaults at run time, however nothing shows up in tsan
+  // or asan builds. This is a hack that for some unknown reason helps.
+  static auto *lljit_mu = new absl::Mutex();
+  std::optional<absl::MutexLock> lljit_lock(lljit_mu);
+
+  // Construct the LLJIT with the given compiler and object linking layers.
+  auto jit = llvm::orc::LLJITBuilder()
+                 .setCompileFunctionCreator(std::move(compile_function_creator))
+                 .setObjectLinkingLayerCreator(obj_layer_creator)
+                 .setExecutorProcessControl(std::move(*executorProcessControl))
+                 .setNumCompileThreads(0)  // disable multi-threading
+                 .create();
+
+  if (auto err = jit.takeError())
+    return InternalError(
+        StrFormat("failed to construct LLJIT: %s", ToString(err)));
+
+  lljit_lock.reset();
+
+  // Register input module with the LLJIT.
+  ThreadSafeModule tsm(std::move(module), std::move(ctx));
+  if (auto err = (*jit)->addIRModule(std::move(tsm)))
+    return InternalError(
+        StrFormat("failed to add source module: %s", ToString(err)));
+
+  llvm::DataLayout data_layout = (*jit)->getDataLayout();
+
+  // Resolve all exported functions to function pointers.
+  for (std::string_view name : exported) {
+    // Trigger compilation by looking up the exported function.
+    // TODO(tsilytskyi):
+    //   - Do we need to mangle function name?
+    //   - Do we need to verify/adapt function proto to expected API?
+    Expected<ExecutorAddr> addr = (*jit)->lookup(name);
+    if (auto err = addr.takeError())
+      return InternalError(StrFormat(
+          "failed to compile exported function %s: %s", name, ToString(err)));
+
+    // Check that we found an address of an exported function.
+    auto ptr = addr->toPtr<ExportedFunctionPtr>();
+    if (!ptr)
+      return InternalError(
+          StrFormat("exported function %s resolved to null", name));
+
+    engine->exported_.push_back(ptr);
+  }
+
+  // Fill remaining fields and return constructed ExecutionEngine to the caller.
+  engine->jit_ = std::move(*jit);
+  return std::move(engine);
+}
+
+static std::function<llvm::Error(llvm::Module *)>
+MakeOptimizingTransformerForJit(llvm::TargetMachine *targetMachine) {
+  return [targetMachine](llvm::Module *m) -> llvm::Error {
+    llvm::LoopAnalysisManager lam;
+    llvm::FunctionAnalysisManager fam;
+    llvm::CGSCCAnalysisManager cgam;
+    llvm::ModuleAnalysisManager mam;
+
+    llvm::PipelineTuningOptions tuningOptions;
+    // LLVM's loop unrolling isn't well tuned for the loops we emit. Turn it off
+    // as it consumes compile time with little benefit.
+    tuningOptions.LoopUnrolling = false;
+    // Vectorization happens at the MLIR level.
+    tuningOptions.LoopVectorization = false;
+    llvm::PassBuilder pb(targetMachine, tuningOptions);
+
+    pb.registerModuleAnalyses(mam);
+    pb.registerCGSCCAnalyses(cgam);
+    pb.registerFunctionAnalyses(fam);
+    pb.registerLoopAnalyses(lam);
+    pb.crossRegisterProxies(lam, fam, cgam, mam);
+
+    llvm::ModulePassManager mpm;
+    mpm.addPass(pb.buildPerModuleDefaultPipeline(llvm::OptimizationLevel::O2));
+    mpm.run(*m, mam);
+    return llvm::Error::success();
+  };
+}
+
+absl::StatusOr<std::unique_ptr<LlvmExecutionEngine>>
+LlvmExecutionEngine::CreateFromLlvmIr(absl::string_view name,
+                                      absl::string_view entry,
+                                      absl::string_view ir,
+                                      absl::Span<const std::string> options) {
+  llvm::InitializeNativeTarget();
+  llvm::InitializeNativeTargetAsmPrinter();
+  auto llvm_ctx = std::make_unique<llvm::LLVMContext>();
+  llvm::SMDiagnostic diagnostic;
+  llvm::MemoryBufferRef ir_buffer(ir, name);
+  std::unique_ptr<llvm::Module> llvm_module =
+      llvm::parseAssembly(ir_buffer, diagnostic, *llvm_ctx, nullptr);
+
+  // Prepare JIT target machine for code generation.
+  auto builder = llvm::orc::JITTargetMachineBuilder::detectHost();
+  if (!builder) return absl::InternalError(toString(builder.takeError()));
+
+  llvm::Expected<std::unique_ptr<llvm::TargetMachine>> target_machine =
+      builder->createTargetMachine();
+  if (!target_machine)
+    return absl::InternalError(toString(target_machine.takeError()));
+
+  // Set target triple
+  llvm_module->setTargetTriple(
+      llvm::StringRef(target_machine.get()->getTargetTriple().getTriple()));
+
+  // Construct options for the XLA runtime execution engine.
+  RuntimeExecutionEngine::JitOptions engine_options;
+  engine_options.target_machine = std::move(target_machine.get());
+  engine_options.make_optimizing_transformer = MakeOptimizingTransformerForJit;
+
+  std::vector<std::string_view> exported = {entry};
+
+  // Compile input module to the native function.
+  TF_ASSIGN_OR_RETURN(auto engine,
+                      RuntimeExecutionEngine::CreateFromModule(
+                          std::move(llvm_ctx), std::move(llvm_module),
+                          std::move(engine_options), exported));
+
+  return std::unique_ptr<LlvmExecutionEngine>(
+      new LlvmExecutionEngine(std::move(engine)));
+}
+
+}  // namespace host
+}  // namespace stream_executor
diff --git a/third_party/xla/xla/stream_executor/host/host_execution_engine.h b/third_party/xla/xla/stream_executor/host/host_execution_engine.h
new file mode 100644
index 00000000000000..9215ad5ed5e3af
--- /dev/null
+++ b/third_party/xla/xla/stream_executor/host/host_execution_engine.h
@@ -0,0 +1,157 @@
+/* Copyright 2024 The OpenXLA Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+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.
+==============================================================================*/
+
+#ifndef XLA_STREAM_EXECUTOR_HOST_HOST_EXECUTION_ENGINE_H_
+#define XLA_STREAM_EXECUTOR_HOST_HOST_EXECUTION_ENGINE_H_
+
+#include <functional>
+#include <memory>
+#include <string>
+#include <utility>
+#include <vector>
+
+#include "absl/status/statusor.h"
+#include "absl/strings/string_view.h"
+#include "absl/types/span.h"
+#include "llvm/ExecutionEngine/JITEventListener.h"
+#include "llvm/ExecutionEngine/Orc/LLJIT.h"
+#include "llvm/ExecutionEngine/SectionMemoryManager.h"
+#include "llvm/IR/LLVMContext.h"
+#include "llvm/Support/CodeGen.h"
+#include "llvm/Support/Error.h"
+#include "llvm/Support/MemoryBuffer.h"
+#include "llvm/Target/TargetMachine.h"
+#include "xla/stream_executor/host/host_kernel_c_api.h"
+
+namespace stream_executor::host {
+
+class RuntimeExecutionEngine {
+ public:
+  using ExportedFunctionPtr = const SE_HOST_Kernel *;
+
+  // Callback to run optimization passes on the compiled LLVM module.
+  using OptimizingTransformer = std::function<llvm::Error(llvm::Module *)>;
+
+  // Callback to construct an optimizing transformer for the given options.
+  using MakeOptimizingTransformer =
+      std::function<OptimizingTransformer(llvm::TargetMachine *targetMachine)>;
+
+  //------------------------------------------------------------------------- //
+  // Options for creating execution engine from an LLVM module.
+  //------------------------------------------------------------------------- //
+
+  struct JitOptions {
+    // User-provided codegen optimization level.
+    llvm::CodeGenOptLevel opt_level = llvm::CodeGenOptLevel::Default;
+
+    // User-provided target machine specification.
+    std::shared_ptr<llvm::TargetMachine> target_machine = nullptr;
+
+    // User-provided builder for the optimizing transformer.
+    MakeOptimizingTransformer make_optimizing_transformer;
+
+    // User-provided memory mapper for allocating memory for executables.
+    llvm::SectionMemoryManager::MemoryMapper *section_memory_mapper = nullptr;
+
+    // Notify the llvm's global GDB notifications listener.
+    bool enable_gdb_listener = false;
+
+    // Notify the llvm's global Perf notifications listener.
+    bool enable_perf_listener = false;
+  };
+
+  // Creates a new execution engine by compiling the provided LLVM module to
+  // a native executable using LLVM ORC stack.
+  static absl::StatusOr<std::unique_ptr<RuntimeExecutionEngine>>
+  CreateFromModule(std::unique_ptr<llvm::LLVMContext> ctx,
+                   std::unique_ptr<llvm::Module> module, JitOptions options,
+                   absl::Span<const std::string_view> exported);
+
+  //------------------------------------------------------------------------- //
+
+  // Returns a pointer to the exported function.
+  absl::Span<const ExportedFunctionPtr> exported() const { return exported_; }
+
+  ExportedFunctionPtr exported(unsigned ordinal) const {
+    return exported_[ordinal];
+  }
+
+  // Return a memory buffer with a object file behind this execution engine. Can
+  // be null if execution engine didn't save the compiled object file.
+  std::unique_ptr<llvm::MemoryBuffer> obj_file() const;
+
+ private:
+  RuntimeExecutionEngine(bool enable_gdb_listener, bool enable_perf_listener);
+
+  // We build execution engine on top of the ORC LLJIT API, which owns all
+  // compiled/loaded object files and does the linking at run time.
+  //
+  // TODO(ezhulenev): Instead of keeping LLJIT alive we should be able to keep
+  // only llvm::orc::JITDylibSP owning main dylib and the object layer owning
+  // memory-mapped regions holding object files. Once we are done with
+  // executable compilation this jit is defunct because it holds an expired
+  // weak_ptr to an llvm::orc::TargetMachine instance.
+  std::unique_ptr<llvm::orc::LLJIT> jit_;
+
+  // Pointers to resolved exported functions. Indexed by function ordinal.
+  std::vector<ExportedFunctionPtr> exported_;
+
+  // Object file behind the compiled executable. Can be null.
+  std::unique_ptr<llvm::MemoryBuffer> obj_file_;
+
+  llvm::JITEventListener *gdb_listener_ = nullptr;
+  llvm::JITEventListener *perf_listener_ = nullptr;
+};
+
+// Virtual base class that owns jit-compiled function.
+class HostExecutionEngine {
+ public:
+  virtual ~HostExecutionEngine() = default;
+  virtual SE_HOST_Kernel *kernel() const = 0;
+};
+
+class LlvmExecutionEngine : public HostExecutionEngine {
+ public:
+  SE_HOST_Kernel *kernel() const override { return kernel_; }
+  // TODO(tsilytskyi): clean up kernel_
+  ~LlvmExecutionEngine() override = default;
+
+  static absl::StatusOr<std::unique_ptr<LlvmExecutionEngine>> CreateFromLlvmIr(
+      absl::string_view name, absl::string_view entry, absl::string_view ir,
+      absl::Span<const std::string> options);
+
+ private:
+  explicit LlvmExecutionEngine(
+      std::unique_ptr<RuntimeExecutionEngine> exec_engine)
+      : engine_(std::move(exec_engine)) {
+    kernel_ = reinterpret_cast<SE_HOST_Kernel *>(engine_->exported(0));
+  };
+  std::unique_ptr<RuntimeExecutionEngine> engine_;
+  SE_HOST_Kernel *kernel_;
+};
+
+class CppExecutionEngine : public HostExecutionEngine {
+ public:
+  ~CppExecutionEngine() override = default;
+  SE_HOST_Kernel *kernel() const override { return kernel_; }
+
+ private:
+  CppExecutionEngine() = default;
+  SE_HOST_Kernel *kernel_ = nullptr;
+};
+
+}  // namespace stream_executor::host
+
+#endif  // XLA_STREAM_EXECUTOR_HOST_HOST_EXECUTION_ENGINE_H_
diff --git a/third_party/xla/xla/stream_executor/host/host_executor.cc b/third_party/xla/xla/stream_executor/host/host_executor.cc
index 8d6eb9dcdcf45e..3d59b1c69ea587 100644
--- a/third_party/xla/xla/stream_executor/host/host_executor.cc
+++ b/third_party/xla/xla/stream_executor/host/host_executor.cc
@@ -22,24 +22,31 @@ limitations under the License.
 
 #include <cstdint>
 #include <memory>
+#include <string>
 #include <utility>
 
 #include "absl/functional/any_invocable.h"
 #include "absl/log/check.h"
+#include "absl/log/log.h"
 #include "absl/status/status.h"
 #include "absl/status/statusor.h"
-#include "absl/strings/numbers.h"
-#include "absl/strings/str_cat.h"
+#include "absl/strings/string_view.h"
 #include "absl/synchronization/notification.h"
+#include "absl/types/span.h"
 #include "xla/stream_executor/device_description.h"
 #include "xla/stream_executor/device_memory.h"
 #include "xla/stream_executor/event.h"
 #include "xla/stream_executor/event_interface.h"
+#include "xla/stream_executor/host/host_execution_engine.h"
+#include "xla/stream_executor/host/host_kernel.h"
 #include "xla/stream_executor/host/host_stream.h"
+#include "xla/stream_executor/kernel_spec.h"
+#include "xla/stream_executor/launch_dim.h"
 #include "xla/stream_executor/stream_executor.h"
 #include "xla/stream_executor/stream_executor_interface.h"
 #include "tsl/platform/mem.h"
 #include "tsl/platform/profile_utils/cpu_utils.h"
+#include "tsl/platform/statusor.h"
 
 namespace stream_executor {
 namespace host {
@@ -51,6 +58,49 @@ HostStream* AsHostStream(Stream* stream) {
 
 absl::Status HostExecutor::Init() { return absl::OkStatus(); }
 
+absl::StatusOr<std::unique_ptr<Kernel>> HostExecutor::CreateKernel() {
+  return std::make_unique<HostKernel>();
+}
+
+absl::Status HostExecutor::GetKernel(const MultiKernelLoaderSpec& spec,
+                                     Kernel* kernel) {
+  HostKernel* host_kernel = AsHostKernel(kernel);
+  host_kernel->SetArity(spec.arity());
+
+  VLOG(3) << "GetKernel on kernel " << kernel << " : " << kernel->name();
+
+  if (spec.has_llvm_host_kernel()) {
+    const LlvmHostKernel& llvm_host_kernel = spec.llvm_host_kernel();
+    const absl::string_view name = llvm_host_kernel.kernel_name();
+    const absl::string_view entry = llvm_host_kernel.entrypoint();
+    const absl::string_view ir = llvm_host_kernel.ir();
+    const absl::Span<const std::string> options = llvm_host_kernel.options();
+
+    TF_ASSIGN_OR_RETURN(
+        auto execution_engine,
+        LlvmExecutionEngine::CreateFromLlvmIr(name, entry, ir, options));
+    host_kernel->SetExecutionEngine(std::move(execution_engine));
+    return absl::OkStatus();
+  } else if (false /* TODO(tsilytskyi): Implement CppHostKernel */) {
+    // host_kernel->SetExecutionEngine(std::make_unique<CppExecutionEngine>());
+  } else {
+    return absl::InternalError("No method of loading host kernel provided");
+  }
+
+  return absl::UnimplementedError("Not Implemented");
+}
+
+absl::Status HostExecutor::Launch(Stream* stream, const ThreadDim& thread_dims,
+                                  const BlockDim& block_dims,
+                                  const Kernel& kernel,
+                                  const KernelArgs& args) {
+  // const HostKernel* host_kernel = AsHostKernel(&kernel);
+
+  // TODO(tsilytskyi): convert args into proper format
+  // host_kernel->Launch(thread_dims, args);
+  return absl::UnimplementedError("Not Implemented");
+}
+
 bool HostExecutor::DeviceMemoryUsage(int64_t* free, int64_t* total) const {
   tsl::port::MemoryInfo mem_info = tsl::port::GetMemoryInfo();
   *free = (mem_info.free != INT64_MAX) ? mem_info.free : -1;
diff --git a/third_party/xla/xla/stream_executor/host/host_executor.h b/third_party/xla/xla/stream_executor/host/host_executor.h
index 8c781096d6c874..40cb0ccea89330 100644
--- a/third_party/xla/xla/stream_executor/host/host_executor.h
+++ b/third_party/xla/xla/stream_executor/host/host_executor.h
@@ -57,14 +57,13 @@ class HostExecutor : public StreamExecutor {
   absl::Status Init() override;
 
   absl::Status GetKernel(const MultiKernelLoaderSpec& spec,
-                         Kernel* kernel) override {
-    return absl::UnimplementedError("Not Implemented");
-  }
+                         Kernel* kernel) override;
+
+  absl::StatusOr<std::unique_ptr<Kernel>> CreateKernel() override;
+
   absl::Status Launch(Stream* stream, const ThreadDim& thread_dims,
                       const BlockDim& block_dims, const Kernel& kernel,
-                      const KernelArgs& args) override {
-    return absl::UnimplementedError("Not Implemented");
-  }
+                      const KernelArgs& args) override;
 
   DeviceMemoryBase Allocate(uint64_t size, int64_t memory_space) override;
   void Deallocate(DeviceMemoryBase* mem) override;
diff --git a/third_party/xla/xla/stream_executor/host/host_kernel.cc b/third_party/xla/xla/stream_executor/host/host_kernel.cc
index e2d37085a1c756..b208ec453b9e0e 100644
--- a/third_party/xla/xla/stream_executor/host/host_kernel.cc
+++ b/third_party/xla/xla/stream_executor/host/host_kernel.cc
@@ -29,8 +29,9 @@ namespace stream_executor::host {
 HostKernel::HostKernel(unsigned arity, SE_HOST_Kernel* kernel)
     : arity_(arity), kernel_(kernel) {}
 
-absl::Status HostKernel::Launch(const ThreadDim& thread_dims,
-                                absl::Span<const DeviceMemoryBase> buffers) {
+absl::Status HostKernel::Launch(
+    const ThreadDim& thread_dims,
+    absl::Span<const DeviceMemoryBase> buffers) const {
   SE_HOST_KernelThreadDim kernel_thread_dims = {thread_dims.x, thread_dims.y,
                                                 thread_dims.z};
 
diff --git a/third_party/xla/xla/stream_executor/host/host_kernel.h b/third_party/xla/xla/stream_executor/host/host_kernel.h
index ee8f67738bf08c..b06fbe76b244f2 100644
--- a/third_party/xla/xla/stream_executor/host/host_kernel.h
+++ b/third_party/xla/xla/stream_executor/host/host_kernel.h
@@ -18,26 +18,35 @@ limitations under the License.
 
 #include <cstddef>
 #include <cstdint>
+#include <memory>
+#include <type_traits>
+#include <utility>
 
 #include "absl/status/status.h"
 #include "absl/status/statusor.h"
 #include "absl/types/span.h"
 #include "xla/stream_executor/device_memory.h"
+#include "xla/stream_executor/host/host_execution_engine.h"
 #include "xla/stream_executor/host/host_kernel_c_api.h"
 #include "xla/stream_executor/kernel.h"
 #include "xla/stream_executor/launch_dim.h"
 
 namespace stream_executor::host {
 
+class HostExecutor;
+
 class HostKernel : public Kernel {
  public:
-  HostKernel(unsigned arity, SE_HOST_Kernel* kernel);
+  explicit HostKernel() = default;
+
+  // TODO(tsilytskyi): make this implementation detail private
+  explicit HostKernel(unsigned arity, SE_HOST_Kernel* kernel);
 
-  // TODO(b/331430625): Connect this API to Launch API defined at StreamExecutor
-  // level, which requires refactoring how arguments passed to kernels, as
-  // current KernelArgs structure tied to the GPU kernel ABI.
+  // TODO(b/331430625): Connect this API to Launch API defined at
+  // StreamExecutor level, which requires refactoring how arguments passed to
+  // kernels, as current KernelArgs structure tied to the GPU kernel ABI.
   absl::Status Launch(const ThreadDim& thread_dims,
-                      absl::Span<const DeviceMemoryBase> buffers);
+                      absl::Span<const DeviceMemoryBase> buffers) const;
 
   // For host platform, we assume that a core is a thread, and we can run at
   // most one instance of a kernel on a given thread.
@@ -46,13 +55,31 @@ class HostKernel : public Kernel {
     return 1;
   };
 
+  void SetArity(unsigned arity) { arity_ = arity; };
   unsigned Arity() const override { return arity_; };
 
+  template <typename T>
+  void SetExecutionEngine(std::unique_ptr<T> execution_engine) {
+    static_assert(std::is_base_of<HostExecutionEngine, T>::value,
+                  "T is not derived from HostExecutionEngine");
+    execution_engine_ = std::move(execution_engine);
+  }
+
  private:
+  std::unique_ptr<HostExecutionEngine> execution_engine_;
+
   unsigned arity_;
   SE_HOST_Kernel* kernel_ = nullptr;
 };
 
+inline const HostKernel* AsHostKernel(const Kernel* kernel) {
+  return static_cast<const HostKernel*>(kernel);
+}
+
+inline HostKernel* AsHostKernel(Kernel* kernel) {
+  return static_cast<HostKernel*>(kernel);
+}
+
 }  // namespace stream_executor::host
 
 #endif  // XLA_STREAM_EXECUTOR_HOST_HOST_KERNEL_H_
diff --git a/third_party/xla/xla/stream_executor/host/host_kernel_test.cc b/third_party/xla/xla/stream_executor/host/host_kernel_test.cc
index 6bf3439d2e95e9..2dcb432cf81bf4 100644
--- a/third_party/xla/xla/stream_executor/host/host_kernel_test.cc
+++ b/third_party/xla/xla/stream_executor/host/host_kernel_test.cc
@@ -16,12 +16,20 @@ limitations under the License.
 #include "xla/stream_executor/host/host_kernel.h"
 
 #include <cstdint>
+#include <memory>
+#include <string>
 #include <vector>
 
+#include "absl/types/span.h"
 #include "xla/stream_executor/device_memory.h"
 #include "xla/stream_executor/host/host_kernel_c_api.h"
+#include "xla/stream_executor/kernel_spec.h"
 #include "xla/stream_executor/launch_dim.h"
+#include "xla/stream_executor/platform.h"
+#include "xla/stream_executor/platform_manager.h"
+#include "xla/stream_executor/stream_executor.h"
 #include "tsl/lib/core/status_test_util.h"
+#include "tsl/platform/statusor.h"
 #include "tsl/platform/test.h"
 
 namespace stream_executor::host {
@@ -41,6 +49,38 @@ static SE_HOST_KernelError* AddI32(const SE_HOST_KernelCallFrame* call_frame) {
   return nullptr;
 }
 
+static const char* llvm_kernel_add = R"(
+%SE_HOST_KernelCallFrame = type { ptr, ptr, i64, ptr }
+%struct.SE_HOST_KernelArg = type { ptr, i64 }
+
+define ptr @LlvmAddI32(ptr noundef %0) {
+  %2 = getelementptr inbounds %SE_HOST_KernelCallFrame, ptr %0, i32 0, i32 3
+  %3 = load ptr, ptr %2, align 8
+  %4 = getelementptr inbounds %struct.SE_HOST_KernelArg, ptr %3, i64 1
+  %5 = getelementptr inbounds %struct.SE_HOST_KernelArg, ptr %3, i64 2
+  %6 = load ptr, ptr %3, align 8
+  %7 = load ptr, ptr %4, align 8
+  %8 = load ptr, ptr %5, align 8
+  %9 = getelementptr inbounds %SE_HOST_KernelCallFrame, ptr %0, i32 0, i32 1
+  %10 = load ptr, ptr %9, align 8
+  %11 = load i64, ptr %10, align 8
+  %12 = getelementptr inbounds i32, ptr %6, i64 %11
+  %13 = load i32, ptr %12, align 4
+  %14 = getelementptr inbounds i32, ptr %7, i64 %11
+  %15 = load i32, ptr %14, align 4
+  %16 = add nsw i32 %13, %15
+  %17 = getelementptr inbounds i32, ptr %8, i64 %11
+  store i32 %16, ptr %17, align 4
+  ret ptr null
+}
+)";
+
+static std::unique_ptr<StreamExecutor> NewStreamExecutor() {
+  Platform* platform = PlatformManager::PlatformWithName("Host").value();
+  StreamExecutorConfig config(/*ordinal=*/0);
+  return platform->GetUncachedExecutor(config).value();
+}
+
 TEST(HostKernelTest, Addition) {
   HostKernel kernel(/*arity=*/3, AddI32);
 
@@ -59,4 +99,31 @@ TEST(HostKernelTest, Addition) {
   EXPECT_EQ(out, expected);
 }
 
+TEST(HostKernelTest, LlvmAddition) {
+  std::vector<int32_t> lhs = {1, 2, 3, 4};
+  std::vector<int32_t> rhs = {5, 6, 7, 8};
+  std::vector<int32_t> out = {0, 0, 0, 0};
+
+  DeviceMemoryBase lhs_mem(lhs.data(), lhs.size() * sizeof(int32_t));
+  DeviceMemoryBase rhs_mem(rhs.data(), rhs.size() * sizeof(int32_t));
+  DeviceMemoryBase out_mem(out.data(), out.size() * sizeof(int32_t));
+  std::vector<DeviceMemoryBase> args = {lhs_mem, rhs_mem, out_mem};
+
+  MultiKernelLoaderSpec spec(/*arity=*/3);
+  spec.AddLlvmHostKernel(llvm_kernel_add, "LlvmAddI32", "LlvmAddI32",
+                         absl::Span<std::string>());
+
+  auto executor = NewStreamExecutor();
+  auto eg = executor.get();
+  EXPECT_NE(eg, nullptr);
+  TF_ASSERT_OK_AND_ASSIGN(auto add, HostKernel::Create(eg, spec));
+
+  // TODO(tsilytskyi): implement Launch part
+  // TF_ASSERT_OK(executor->Launch(ThreadDim(4), args));
+
+  // std::vector<int32_t> expected = {6, 8, 10, 12};
+  // EXPECT_EQ(out, expected);
+  // EXPECT_TRUE(true);
+}
+
 }  // namespace stream_executor::host

From c0cfee827b6d49e3aecd2de1837269bc2fbe9ad8 Mon Sep 17 00:00:00 2001
From: Greg Olechwierowicz <olechwierowicz@google.com>
Date: Mon, 13 May 2024 09:02:53 -0700
Subject: [PATCH 025/478] [XLA:GPU] Rename LoopDoubleBufferTransformer to
 DoubleBufferLoopUnrolling.

PiperOrigin-RevId: 633228767
---
 third_party/xla/xla/service/gpu/BUILD         | 14 ++++----
 ...mer.cc => double_buffer_loop_unrolling.cc} |  4 +--
 ...ormer.h => double_buffer_loop_unrolling.h} | 12 +++----
 ...c => double_buffer_loop_unrolling_test.cc} | 36 +++++++++----------
 .../xla/xla/service/gpu/gpu_compiler.cc       | 14 ++++----
 5 files changed, 39 insertions(+), 41 deletions(-)
 rename third_party/xla/xla/service/gpu/{loop_double_buffer_transformer.cc => double_buffer_loop_unrolling.cc} (99%)
 rename third_party/xla/xla/service/gpu/{loop_double_buffer_transformer.h => double_buffer_loop_unrolling.h} (88%)
 rename third_party/xla/xla/service/gpu/{loop_double_buffer_transformer_test.cc => double_buffer_loop_unrolling_test.cc} (97%)

diff --git a/third_party/xla/xla/service/gpu/BUILD b/third_party/xla/xla/service/gpu/BUILD
index fec7fc3ea976f5..a851ea99b9b182 100644
--- a/third_party/xla/xla/service/gpu/BUILD
+++ b/third_party/xla/xla/service/gpu/BUILD
@@ -3433,7 +3433,7 @@ cc_library(
         ":instruction_fusion",
         ":ir_emission_utils",
         ":ir_emitter",
-        ":loop_double_buffer_transformer",
+        ":double_buffer_loop_unrolling",
         ":matmul_utils",
         ":metrics",
         ":move_copy_to_users",
@@ -5803,9 +5803,9 @@ xla_cc_test(
 )
 
 cc_library(
-    name = "loop_double_buffer_transformer",
-    srcs = ["loop_double_buffer_transformer.cc"],
-    hdrs = ["loop_double_buffer_transformer.h"],
+    name = "double_buffer_loop_unrolling",
+    srcs = ["double_buffer_loop_unrolling.cc"],
+    hdrs = ["double_buffer_loop_unrolling.h"],
     deps = [
         "//xla:status",
         "//xla:util",
@@ -5829,10 +5829,10 @@ cc_library(
 )
 
 xla_cc_test(
-    name = "loop_double_buffer_transformer_test",
-    srcs = ["loop_double_buffer_transformer_test.cc"],
+    name = "double_buffer_loop_unrolling_test",
+    srcs = ["double_buffer_loop_unrolling_test.cc"],
     deps = [
-        ":loop_double_buffer_transformer",
+        ":double_buffer_loop_unrolling",
         "//xla:test",
         "//xla:xla_data_proto_cc",
         "//xla:xla_proto_cc",
diff --git a/third_party/xla/xla/service/gpu/loop_double_buffer_transformer.cc b/third_party/xla/xla/service/gpu/double_buffer_loop_unrolling.cc
similarity index 99%
rename from third_party/xla/xla/service/gpu/loop_double_buffer_transformer.cc
rename to third_party/xla/xla/service/gpu/double_buffer_loop_unrolling.cc
index 36f4d81a452a02..4c10da5b8aa047 100644
--- a/third_party/xla/xla/service/gpu/loop_double_buffer_transformer.cc
+++ b/third_party/xla/xla/service/gpu/double_buffer_loop_unrolling.cc
@@ -12,7 +12,7 @@ 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.
 ==============================================================================*/
-#include "xla/service/gpu/loop_double_buffer_transformer.h"
+#include "xla/service/gpu/double_buffer_loop_unrolling.h"
 
 #include <cstdint>
 #include <iterator>
@@ -362,7 +362,7 @@ absl::StatusOr<bool> DoubleBufferingUnroll(HloInstruction* while_instr,
 
 }  // namespace
 
-absl::StatusOr<bool> LoopDoubleBufferTransformer::Run(
+absl::StatusOr<bool> DoubleBufferLoopUnrolling::Run(
     HloModule* module,
     const absl::flat_hash_set<absl::string_view>& execution_threads) {
   bool changed = false;
diff --git a/third_party/xla/xla/service/gpu/loop_double_buffer_transformer.h b/third_party/xla/xla/service/gpu/double_buffer_loop_unrolling.h
similarity index 88%
rename from third_party/xla/xla/service/gpu/loop_double_buffer_transformer.h
rename to third_party/xla/xla/service/gpu/double_buffer_loop_unrolling.h
index f0d891a09ab685..120070dbccd452 100644
--- a/third_party/xla/xla/service/gpu/loop_double_buffer_transformer.h
+++ b/third_party/xla/xla/service/gpu/double_buffer_loop_unrolling.h
@@ -12,8 +12,8 @@ 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.
 ==============================================================================*/
-#ifndef XLA_SERVICE_GPU_LOOP_DOUBLE_BUFFER_TRANSFORMER_H_
-#define XLA_SERVICE_GPU_LOOP_DOUBLE_BUFFER_TRANSFORMER_H_
+#ifndef XLA_SERVICE_GPU_DOUBLE_BUFFER_LOOP_UNROLLING_H_
+#define XLA_SERVICE_GPU_DOUBLE_BUFFER_LOOP_UNROLLING_H_
 
 #include "absl/container/flat_hash_set.h"
 #include "absl/status/statusor.h"
@@ -47,14 +47,14 @@ namespace gpu {
 // unrolled.
 // TODO(olechwierowicz): Rename the loop unroller to something more generic like
 // 'DoubleBufferLoopUnrolling'.
-class LoopDoubleBufferTransformer : public HloModulePass {
+class DoubleBufferLoopUnrolling : public HloModulePass {
  public:
   enum class UnrollStrategy { kDoubleBuffer, kFullUnroll };
 
-  explicit LoopDoubleBufferTransformer(
+  explicit DoubleBufferLoopUnrolling(
       UnrollStrategy unroll_strategy = UnrollStrategy::kDoubleBuffer)
       : unroll_strategy_(unroll_strategy) {};
-  ~LoopDoubleBufferTransformer() override = default;
+  ~DoubleBufferLoopUnrolling() override = default;
 
   absl::string_view name() const override {
     return "loop-double-buffer-transformer";
@@ -72,4 +72,4 @@ class LoopDoubleBufferTransformer : public HloModulePass {
 }  // end namespace gpu
 }  // end namespace xla
 
-#endif  // XLA_SERVICE_GPU_LOOP_DOUBLE_BUFFER_TRANSFORMER_H_
+#endif  // XLA_SERVICE_GPU_DOUBLE_BUFFER_LOOP_UNROLLING_H_
diff --git a/third_party/xla/xla/service/gpu/loop_double_buffer_transformer_test.cc b/third_party/xla/xla/service/gpu/double_buffer_loop_unrolling_test.cc
similarity index 97%
rename from third_party/xla/xla/service/gpu/loop_double_buffer_transformer_test.cc
rename to third_party/xla/xla/service/gpu/double_buffer_loop_unrolling_test.cc
index 1cac4fa2c1991c..b7ff9a24de7e2e 100644
--- a/third_party/xla/xla/service/gpu/loop_double_buffer_transformer_test.cc
+++ b/third_party/xla/xla/service/gpu/double_buffer_loop_unrolling_test.cc
@@ -13,7 +13,7 @@ See the License for the specific language governing permissions and
 limitations under the License.
 ==============================================================================*/
 
-#include "xla/service/gpu/loop_double_buffer_transformer.h"
+#include "xla/service/gpu/double_buffer_loop_unrolling.h"
 
 #include <cstdint>
 #include <memory>
@@ -102,8 +102,8 @@ ENTRY main {
 
   TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr<xla::HloModule> module,
                           ParseAndReturnVerifiedModule(kModuleString));
-  LoopDoubleBufferTransformer double_buffer(
-      LoopDoubleBufferTransformer::UnrollStrategy::kFullUnroll);
+  DoubleBufferLoopUnrolling double_buffer(
+      DoubleBufferLoopUnrolling::UnrollStrategy::kFullUnroll);
   TupleSimplifier tuple_simp;
   bool changed;
   TF_ASSERT_OK_AND_ASSIGN(changed, double_buffer.Run(module.get()));
@@ -166,8 +166,8 @@ ENTRY main {
 
   TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr<xla::HloModule> module,
                           ParseAndReturnVerifiedModule(kModuleString));
-  LoopDoubleBufferTransformer double_buffer(
-      LoopDoubleBufferTransformer::UnrollStrategy::kFullUnroll);
+  DoubleBufferLoopUnrolling double_buffer(
+      DoubleBufferLoopUnrolling::UnrollStrategy::kFullUnroll);
   TupleSimplifier tuple_simp;
   bool changed;
   TF_ASSERT_OK_AND_ASSIGN(changed, double_buffer.Run(module.get()));
@@ -235,7 +235,7 @@ ENTRY main {
 
   TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr<xla::HloModule> module,
                           ParseAndReturnVerifiedModule(kModuleString));
-  LoopDoubleBufferTransformer double_buffer;
+  DoubleBufferLoopUnrolling double_buffer;
   TupleSimplifier tuple_simp;
   bool changed;
   TF_ASSERT_OK_AND_ASSIGN(changed, double_buffer.Run(module.get()));
@@ -307,7 +307,7 @@ ENTRY main {
 
   TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr<xla::HloModule> module,
                           ParseAndReturnVerifiedModule(kModuleString));
-  LoopDoubleBufferTransformer double_buffer;
+  DoubleBufferLoopUnrolling double_buffer;
   TupleSimplifier tuple_simp;
   EXPECT_THAT(double_buffer.Run(module.get()), IsOkAndHolds(true));
   EXPECT_THAT(tuple_simp.Run(module.get()), IsOkAndHolds(true));
@@ -372,7 +372,7 @@ ENTRY main {
 
   TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr<xla::HloModule> module,
                           ParseAndReturnVerifiedModule(kModuleString));
-  LoopDoubleBufferTransformer double_buffer;
+  DoubleBufferLoopUnrolling double_buffer;
   TupleSimplifier tuple_simp;
   EXPECT_THAT(double_buffer.Run(module.get()), IsOkAndHolds(true));
   EXPECT_THAT(tuple_simp.Run(module.get()), IsOkAndHolds(true));
@@ -444,7 +444,7 @@ ENTRY main {
 
   TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr<xla::HloModule> module,
                           ParseAndReturnVerifiedModule(kModuleString));
-  LoopDoubleBufferTransformer double_buffer;
+  DoubleBufferLoopUnrolling double_buffer;
   TupleSimplifier tuple_simp;
   EXPECT_THAT(double_buffer.Run(module.get()), IsOkAndHolds(true));
   EXPECT_THAT(tuple_simp.Run(module.get()), IsOkAndHolds(true));
@@ -517,8 +517,8 @@ ENTRY main {
 
   TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr<xla::HloModule> module,
                           ParseAndReturnVerifiedModule(kModuleString));
-  LoopDoubleBufferTransformer double_buffer(
-      LoopDoubleBufferTransformer::UnrollStrategy::kFullUnroll);
+  DoubleBufferLoopUnrolling double_buffer(
+      DoubleBufferLoopUnrolling::UnrollStrategy::kFullUnroll);
   TupleSimplifier tuple_simp;
   EXPECT_THAT(double_buffer.Run(module.get()), IsOkAndHolds(true));
   EXPECT_THAT(tuple_simp.Run(module.get()), IsOkAndHolds(true));
@@ -592,7 +592,7 @@ ENTRY main {
 
   TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr<xla::HloModule> module,
                           ParseAndReturnVerifiedModule(kModuleString));
-  LoopDoubleBufferTransformer double_buffer;
+  DoubleBufferLoopUnrolling double_buffer;
   EXPECT_THAT(double_buffer.Run(module.get()), IsOkAndHolds(true));
 
   absl::flat_hash_set<const HloComputation*> while_loops_callees;
@@ -651,7 +651,7 @@ ENTRY main {
 
   TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr<xla::HloModule> module,
                           ParseAndReturnVerifiedModule(kModuleString));
-  LoopDoubleBufferTransformer double_buffer;
+  DoubleBufferLoopUnrolling double_buffer;
   EXPECT_THAT(double_buffer.Run(module.get()), IsOkAndHolds(true));
 
   absl::flat_hash_set<const HloComputation*> while_loops_callees;
@@ -710,8 +710,8 @@ ENTRY main {
 
   TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr<xla::HloModule> module,
                           ParseAndReturnVerifiedModule(kModuleString));
-  LoopDoubleBufferTransformer double_buffer(
-      LoopDoubleBufferTransformer::UnrollStrategy::kFullUnroll);
+  DoubleBufferLoopUnrolling double_buffer(
+      DoubleBufferLoopUnrolling::UnrollStrategy::kFullUnroll);
   EXPECT_THAT(double_buffer.Run(module.get()), IsOkAndHolds(true));
 
   absl::flat_hash_set<const HloComputation*> while_loops_callees;
@@ -776,7 +776,7 @@ ENTRY main {
 
   TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr<xla::HloModule> module,
                           ParseAndReturnVerifiedModule(kModuleString));
-  LoopDoubleBufferTransformer double_buffer;
+  DoubleBufferLoopUnrolling double_buffer;
   EXPECT_THAT(double_buffer.Run(module.get()), IsOkAndHolds(true));
 
   int64_t num_whiles = 0;
@@ -830,8 +830,8 @@ ENTRY main {
 
   TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr<xla::HloModule> module,
                           ParseAndReturnVerifiedModule(kModuleString));
-  LoopDoubleBufferTransformer double_buffer(
-      LoopDoubleBufferTransformer::UnrollStrategy::kFullUnroll);
+  DoubleBufferLoopUnrolling double_buffer(
+      DoubleBufferLoopUnrolling::UnrollStrategy::kFullUnroll);
   EXPECT_THAT(double_buffer.Run(module.get()), IsOkAndHolds(true));
 
   int64_t num_whiles = 0;
diff --git a/third_party/xla/xla/service/gpu/gpu_compiler.cc b/third_party/xla/xla/service/gpu/gpu_compiler.cc
index 2d33d67d2d6064..02673628e26980 100644
--- a/third_party/xla/xla/service/gpu/gpu_compiler.cc
+++ b/third_party/xla/xla/service/gpu/gpu_compiler.cc
@@ -117,6 +117,7 @@ limitations under the License.
 #include "xla/service/gpu/custom_kernel_fusion_rewriter.h"
 #include "xla/service/gpu/dot_dimension_sorter.h"
 #include "xla/service/gpu/dot_operand_converter.h"
+#include "xla/service/gpu/double_buffer_loop_unrolling.h"
 #include "xla/service/gpu/fusion_pipeline.h"
 #include "xla/service/gpu/fusion_wrapper.h"
 #include "xla/service/gpu/gemm_broadcast_folding_rewriter.h"
@@ -141,7 +142,6 @@ limitations under the License.
 #include "xla/service/gpu/ir_emission_utils.h"
 #include "xla/service/gpu/ir_emitter_context.h"
 #include "xla/service/gpu/ir_emitter_unnested.h"
-#include "xla/service/gpu/loop_double_buffer_transformer.h"
 #include "xla/service/gpu/matmul_utils.h"
 #include "xla/service/gpu/metrics.h"
 #include "xla/service/gpu/model/gpu_cost_model_stats_collection.h"
@@ -1096,28 +1096,26 @@ absl::Status RunPostFusionPasses(
     pipeline.AddPass<AllReduceBlueConnect>(blueconnect_num_devices_per_host);
   }
 
-  std::optional<LoopDoubleBufferTransformer::UnrollStrategy> unroll_strategy =
+  std::optional<DoubleBufferLoopUnrolling::UnrollStrategy> unroll_strategy =
       std::nullopt;
   // Support old flag.
   if (opts.xla_gpu_enable_while_loop_double_buffering()) {
-    unroll_strategy =
-        LoopDoubleBufferTransformer::UnrollStrategy::kDoubleBuffer;
+    unroll_strategy = DoubleBufferLoopUnrolling::UnrollStrategy::kDoubleBuffer;
   }
   // Support new flag setting style, override the old one.
   if (opts.xla_gpu_enable_while_loop_unrolling() ==
       DebugOptions::WHILE_LOOP_UNROLLING_DOUBLE_BUFFER) {
-    unroll_strategy =
-        LoopDoubleBufferTransformer::UnrollStrategy::kDoubleBuffer;
+    unroll_strategy = DoubleBufferLoopUnrolling::UnrollStrategy::kDoubleBuffer;
   }
   if (opts.xla_gpu_enable_while_loop_unrolling() ==
       DebugOptions::WHILE_LOOP_UNROLLING_FULL_UNROLL) {
     LOG_IF(WARNING, unroll_strategy != std::nullopt)
         << "Overriding double buffering set via "
            "`xla_gpu_enable_while_loop_double_buffering` flag.";
-    unroll_strategy = LoopDoubleBufferTransformer::UnrollStrategy::kFullUnroll;
+    unroll_strategy = DoubleBufferLoopUnrolling::UnrollStrategy::kFullUnroll;
   }
   if (unroll_strategy != std::nullopt) {
-    pipeline.AddPass<LoopDoubleBufferTransformer>(*unroll_strategy);
+    pipeline.AddPass<DoubleBufferLoopUnrolling>(*unroll_strategy);
     pipeline.AddPass<TupleSimplifier>();
     pipeline.AddPass<HloDCE>();
   }

From 96ed1fe06391abb0ed9851c70066eb3b7edd5e34 Mon Sep 17 00:00:00 2001
From: Kyle Lucke <klucke@google.com>
Date: Mon, 13 May 2024 09:17:33 -0700
Subject: [PATCH 026/478] Remvoe unused, unimplemented methods on
 XlaInterpreterExecutor.

PiperOrigin-RevId: 633233078
---
 third_party/xla/xla/backends/interpreter/executor.h | 5 -----
 1 file changed, 5 deletions(-)

diff --git a/third_party/xla/xla/backends/interpreter/executor.h b/third_party/xla/xla/backends/interpreter/executor.h
index 909d7fbb1bc513..9f72e99dd0a69d 100644
--- a/third_party/xla/xla/backends/interpreter/executor.h
+++ b/third_party/xla/xla/backends/interpreter/executor.h
@@ -175,11 +175,6 @@ class XlaInterpreterExecutor : public StreamExecutor {
   // The device ordinal value that this executor was initialized with; recorded
   // for use in getting device metadata. Immutable post-initialization.
   int device_ordinal_;
-
-  DeviceMemoryBase AllocateSingleOutput(const xla::Shape &shape);
-
-  absl::StatusOr<DeviceMemoryBase> AllocateOutputBuffer(
-      const xla::Shape &shape);
 };
 
 }  // namespace interpreter

From d6222cd597f6c589aaace0719323433ff9b3f061 Mon Sep 17 00:00:00 2001
From: Bart Chrzaszcz <bartchr@google.com>
Date: Mon, 13 May 2024 09:17:46 -0700
Subject: [PATCH 027/478] Don't fail diagnostic if the severity isn't error.

If a pass emits a warning, then the status will be set to `UnknownError` causing pipelines to fail. But we shouldn't if a pass/pattern just emits a warning.

PiperOrigin-RevId: 633233147
---
 .../tsl/tsl/framework/mlir/status_scoped_diagnostic_handler.cc   | 1 +
 1 file changed, 1 insertion(+)

diff --git a/third_party/xla/third_party/tsl/tsl/framework/mlir/status_scoped_diagnostic_handler.cc b/third_party/xla/third_party/tsl/tsl/framework/mlir/status_scoped_diagnostic_handler.cc
index 8882ff8e66cf1d..b2e447a5e778a0 100644
--- a/third_party/xla/third_party/tsl/tsl/framework/mlir/status_scoped_diagnostic_handler.cc
+++ b/third_party/xla/third_party/tsl/tsl/framework/mlir/status_scoped_diagnostic_handler.cc
@@ -59,6 +59,7 @@ mlir::LogicalResult StatusScopedDiagnosticHandler::handleDiagnostic(
   // Emit non-errors to VLOG instead of the internal status.
   if (diag.getSeverity() != mlir::DiagnosticSeverity::Error) {
     VLOG(1) << diag_str_;
+    return mlir::success();
   }
 
   status_.Update(absl::UnknownError(diag_str_));

From daa1a97362889e183fb8df19f94666647073ce9e Mon Sep 17 00:00:00 2001
From: Kyle Lucke <klucke@google.com>
Date: Mon, 13 May 2024 09:19:44 -0700
Subject: [PATCH 028/478] Remove StreamExecutorInterface from all the
 non-static CommandBuffer methods.

Each CommandBuffer class is uniquely tied to a specific StreamExecutorInterface (parent_), which is tracked as member data.

PiperOrigin-RevId: 633233717
---
 .../service/gpu/runtime/command_buffer_cmd.cc |  19 +--
 .../xla/xla/stream_executor/command_buffer.h  |  43 ++----
 .../stream_executor/gpu/gpu_command_buffer.cc | 122 +++++++-----------
 .../stream_executor/gpu/gpu_command_buffer.h  |  38 ++----
 .../gpu/gpu_command_buffer_test.cc            |  50 ++++---
 5 files changed, 106 insertions(+), 166 deletions(-)

diff --git a/third_party/xla/xla/service/gpu/runtime/command_buffer_cmd.cc b/third_party/xla/xla/service/gpu/runtime/command_buffer_cmd.cc
index 8a4af21e1b3b10..b415fdd14e15a5 100644
--- a/third_party/xla/xla/service/gpu/runtime/command_buffer_cmd.cc
+++ b/third_party/xla/xla/service/gpu/runtime/command_buffer_cmd.cc
@@ -292,7 +292,6 @@ absl::Status CommandBufferCmdSequence::Record(
     }
   }
 
-  se::StreamExecutor* device = execute_params.stream->parent();
   const ModuleAnnotations* annotations = GetCurrentModuleAnnotations();
 
   // Track the number of commands recorded between barriers.
@@ -309,7 +308,7 @@ absl::Status CommandBufferCmdSequence::Record(
               << num_recorded_commands[execution_scope_id]
               << " recorded commands into the execution scope #"
               << execution_scope_id.value();
-      TF_RETURN_IF_ERROR(command_buffer->Barrier(device, execution_scope_id));
+      TF_RETURN_IF_ERROR(command_buffer->Barrier(execution_scope_id));
       num_recorded_commands.erase(execution_scope_id);
     }
     VLOG(5) << " Record command buffer with scope id "
@@ -849,8 +848,7 @@ absl::Status IfCmd::Record(const Thunk::ExecuteParams& execute_params,
   VLOG(5) << "  pred: " << pred_ << " (" << pred.opaque() << ")";
 
   return command_buffer->If(
-      execution_scope_id, execute_params.stream->parent(),
-      se::DeviceMemory<bool>(pred),
+      execution_scope_id, se::DeviceMemory<bool>(pred),
       CreateBuilder(&then_commands_, &execute_params, &record_params));
 }
 
@@ -893,8 +891,7 @@ absl::Status IfElseCmd::Record(const Thunk::ExecuteParams& execute_params,
   VLOG(5) << "  pred: " << pred_ << " (" << pred.opaque() << ")";
 
   return command_buffer->IfElse(
-      execution_scope_id, execute_params.stream->parent(),
-      se::DeviceMemory<bool>(pred),
+      execution_scope_id, se::DeviceMemory<bool>(pred),
       CreateBuilder(&then_commands_, &execute_params, &record_params),
       CreateBuilder(&else_commands_, &execute_params, &record_params));
 }
@@ -939,7 +936,6 @@ absl::Status CaseCmd::Record(const Thunk::ExecuteParams& execute_params,
   VLOG(5) << "  index: " << index_ << " (" << index.opaque() << ")";
 
   return command_buffer->Case(execution_scope_id,
-                              execute_params.stream->parent(),
                               se::DeviceMemory<int32_t>(index),
                               CreateBuilders(absl::MakeSpan(branches_commands_),
                                              &execute_params, &record_params));
@@ -985,7 +981,7 @@ absl::Status ForCmd::Record(const Thunk::ExecuteParams& execute_params,
           << loop_counter.opaque() << ")";
 
   return command_buffer->For(
-      execution_scope_id, execute_params.stream->parent(), num_iterations_,
+      execution_scope_id, num_iterations_,
       se::DeviceMemory<int32_t>(loop_counter),
       CreateBuilder(&body_commands_, &execute_params, &record_params));
 }
@@ -1030,8 +1026,7 @@ absl::Status WhileCmd::Record(const Thunk::ExecuteParams& execute_params,
   VLOG(5) << "  pred: " << pred_ << " (" << pred.opaque() << ")";
 
   return command_buffer->While(
-      execution_scope_id, execute_params.stream->parent(),
-      se::DeviceMemory<bool>(pred),
+      execution_scope_id, se::DeviceMemory<bool>(pred),
       CreateExecutionScopeBuilder(&cond_commands_, &execute_params,
                                   &record_params),
       CreateBuilder(&body_commands_, &execute_params, &record_params));
@@ -1340,7 +1335,6 @@ absl::Status BarrierCmd::Record(const Thunk::ExecuteParams& execute_params,
           << " to stream " << execution_stream_id().value();
   if (from_stream_id_ != execution_stream_id()) {
     TF_RETURN_IF_ERROR(command_buffer->Barrier(
-        execute_params.stream->parent(),
         CommandBufferCmd::GetExecutionScope(record_params, from_stream_id_),
         CommandBufferCmd::GetExecutionScope(record_params,
                                             execution_stream_id())));
@@ -1367,8 +1361,7 @@ absl::Status CollectiveCmd::BarrierIfAsync(
     const CommandBufferCmd::RecordParams& record_params) {
   if (IsAsync()) {
     TF_RETURN_IF_ERROR(
-        command_buffer->Barrier(executor,
-                                CommandBufferCmd::GetExecutionScope(
+        command_buffer->Barrier(CommandBufferCmd::GetExecutionScope(
                                     record_params, async_from_stream_id_),
                                 CommandBufferCmd::GetExecutionScope(
                                     record_params, execution_stream_id())));
diff --git a/third_party/xla/xla/stream_executor/command_buffer.h b/third_party/xla/xla/stream_executor/command_buffer.h
index 21ee0bc79f306e..1237493a2fb55e 100644
--- a/third_party/xla/xla/stream_executor/command_buffer.h
+++ b/third_party/xla/xla/stream_executor/command_buffer.h
@@ -202,25 +202,20 @@ class CommandBuffer {
   // Adds an execution barrier to a given execution scope: all commands added
   // before a barrier in a the execution scope will complete before any of the
   // commands added after a barrier in the same execution scope.
-  virtual absl::Status Barrier(StreamExecutorInterface* executor,
-                               ExecutionScopeId execution_scope_id) = 0;
+  virtual absl::Status Barrier(ExecutionScopeId execution_scope_id) = 0;
 
   // Adds an execution barrier that synchronizes commands across multiple
   // execution scopes. See example #2 in execution scope id documentation.
   virtual absl::Status Barrier(
-      StreamExecutorInterface* executor,
       absl::Span<const ExecutionScopeId> execution_scope_ids) = 0;
 
   // Adds an execution barrier from execution scope `from_execution_scope_id` to
   // execution scope `to_execution_scope_id`. See example #3 for details.
-  virtual absl::Status Barrier(StreamExecutorInterface* executor,
-                               ExecutionScopeId from_execution_scope_id,
+  virtual absl::Status Barrier(ExecutionScopeId from_execution_scope_id,
                                ExecutionScopeId to_execution_scope_id) = 0;
 
   // Adds an execution barrier to the default execution scope.
-  absl::Status Barrier(StreamExecutorInterface* executor) {
-    return Barrier(executor, kDefaulExecutionScope);
-  }
+  absl::Status Barrier() { return Barrier(kDefaulExecutionScope); }
 
   // Adds a kernel launch command.
   virtual absl::Status Launch(ExecutionScopeId execution_scope_id,
@@ -292,29 +287,24 @@ class CommandBuffer {
   // Adds a conditional operation that will execute a command buffer constructed
   // by `then_builder` if `pred` value is `true`.
   virtual absl::Status If(ExecutionScopeId execution_scope_id,
-                          StreamExecutorInterface* executor,
                           DeviceMemory<bool> pred, Builder then_builder) = 0;
 
   // Adds a conditional If operation to default execution scope.
-  absl::Status If(StreamExecutorInterface* executor, DeviceMemory<bool> pred,
-                  Builder then_builder) {
-    return If(kDefaulExecutionScope, executor, pred, then_builder);
+  absl::Status If(DeviceMemory<bool> pred, Builder then_builder) {
+    return If(kDefaulExecutionScope, pred, then_builder);
   }
 
   // Adds a conditional operation that will execute a command buffer constructed
   // by `then_builder` if `pred` value is `true`, or a command buffer
   // constructed by `else_builder` if `pred` is `false`.
   virtual absl::Status IfElse(ExecutionScopeId execution_scope_id,
-                              StreamExecutorInterface* executor,
                               DeviceMemory<bool> pred, Builder then_builder,
                               Builder else_builder) = 0;
 
   // Adds a conditional IfElse operation to default execution scope.
-  absl::Status IfElse(StreamExecutorInterface* executor,
-                      DeviceMemory<bool> pred, Builder then_builder,
+  absl::Status IfElse(DeviceMemory<bool> pred, Builder then_builder,
                       Builder else_builder) {
-    return IfElse(kDefaulExecutionScope, executor, pred, then_builder,
-                  else_builder);
+    return IfElse(kDefaulExecutionScope, pred, then_builder, else_builder);
   }
 
   // Adds a conditional operation that will execute a command buffer constructed
@@ -323,15 +313,13 @@ class CommandBuffer {
   //
   // See: https://github.com/openxla/stablehlo/blob/main/docs/spec.md#case
   virtual absl::Status Case(ExecutionScopeId execution_scope_id,
-                            StreamExecutorInterface* executor,
                             DeviceMemory<int32_t> index,
                             std::vector<Builder> branches) = 0;
 
   // Adds a conditional Case operation to default execution scope.
-  absl::Status Case(StreamExecutorInterface* executor,
-                    DeviceMemory<int32_t> index,
+  absl::Status Case(DeviceMemory<int32_t> index,
                     std::vector<Builder> branches) {
-    return Case(kDefaulExecutionScope, executor, index, branches);
+    return Case(kDefaulExecutionScope, index, branches);
   }
 
   // Adds a conditional operation that will execute a command buffer constructed
@@ -339,15 +327,14 @@ class CommandBuffer {
   // condition is known at compile time (`num_iteration` < `loop_counter`), and
   // does not require a `cond_builder`.
   virtual absl::Status For(ExecutionScopeId execution_scope_id,
-                           StreamExecutorInterface* executor,
                            int32_t num_iteration,
                            DeviceMemory<int32_t> loop_counter,
                            Builder body_builder) = 0;
 
   // Adds a conditional For operation to default execution scope.
-  absl::Status For(StreamExecutorInterface* executor, int32_t num_iteration,
-                   DeviceMemory<int32_t> loop_counter, Builder body_builder) {
-    return For(kDefaulExecutionScope, executor, num_iteration, loop_counter,
+  absl::Status For(int32_t num_iteration, DeviceMemory<int32_t> loop_counter,
+                   Builder body_builder) {
+    return For(kDefaulExecutionScope, num_iteration, loop_counter,
                body_builder);
   }
 
@@ -368,16 +355,14 @@ class CommandBuffer {
   // condition twice: (1) before the conditional node in the scope defined by
   // `execution_scope_id` (2) inside the loop body with default execution scope.
   virtual absl::Status While(ExecutionScopeId execution_scope_id,
-                             StreamExecutorInterface* executor,
                              DeviceMemory<bool> pred,
                              ExecutionScopeBuilder cond_builder,
                              Builder body_builder) = 0;
 
   // Adds a conditional While operation to default execution scope.
-  absl::Status While(StreamExecutorInterface* executor, DeviceMemory<bool> pred,
+  absl::Status While(DeviceMemory<bool> pred,
                      ExecutionScopeBuilder cond_builder, Builder body_builder) {
-    return While(kDefaulExecutionScope, executor, pred, cond_builder,
-                 body_builder);
+    return While(kDefaulExecutionScope, pred, cond_builder, body_builder);
   }
 
   //--------------------------------------------------------------------------//
diff --git a/third_party/xla/xla/stream_executor/gpu/gpu_command_buffer.cc b/third_party/xla/xla/stream_executor/gpu/gpu_command_buffer.cc
index 7d7ff5a4c6bdfb..c0788cf21a9e58 100644
--- a/third_party/xla/xla/stream_executor/gpu/gpu_command_buffer.cc
+++ b/third_party/xla/xla/stream_executor/gpu/gpu_command_buffer.cc
@@ -213,73 +213,71 @@ GpuCommandBuffer::Dependencies GpuCommandBuffer::GetBarrier(
 }
 
 absl::StatusOr<GpuCommandBuffer::SetIfConditionKernel*>
-GpuCommandBuffer::GetSetIfConditionKernel(StreamExecutorInterface* executor) {
+GpuCommandBuffer::GetSetIfConditionKernel() {
   if (!set_if_condition_kernel_) {
     MultiKernelLoaderSpec spec(/*arity=*/2);
     spec.AddCudaPtxInMemory(gpu::GetSetIfConditionKernel(), "set_if_condition");
     TF_ASSIGN_OR_RETURN(set_if_condition_kernel_,
-                        SetIfConditionKernel::Create(executor, spec));
+                        SetIfConditionKernel::Create(parent_, spec));
   }
   return &set_if_condition_kernel_;
 }
 
 absl::StatusOr<GpuCommandBuffer::SetIfElseConditionKernel*>
-GpuCommandBuffer::GetSetIfElseConditionKernel(
-    StreamExecutorInterface* executor) {
+GpuCommandBuffer::GetSetIfElseConditionKernel() {
   if (!set_if_else_condition_kernel_) {
     MultiKernelLoaderSpec spec(/*arity=*/3);
     spec.AddCudaPtxInMemory(gpu::GetSetIfElseConditionKernel(),
                             "set_if_else_condition");
     TF_ASSIGN_OR_RETURN(set_if_else_condition_kernel_,
-                        SetIfElseConditionKernel::Create(executor, spec));
+                        SetIfElseConditionKernel::Create(parent_, spec));
   }
   return &set_if_else_condition_kernel_;
 }
 
 absl::StatusOr<GpuCommandBuffer::SetCaseConditionKernel*>
-GpuCommandBuffer::GetSetCaseConditionKernel(StreamExecutorInterface* executor) {
+GpuCommandBuffer::GetSetCaseConditionKernel() {
   if (!set_case_condition_kernel_) {
     MultiKernelLoaderSpec spec(/*arity=*/10);
     spec.AddCudaPtxInMemory(gpu::GetSetCaseConditionKernel(),
                             "set_case_condition");
     TF_ASSIGN_OR_RETURN(set_case_condition_kernel_,
-                        SetCaseConditionKernel::Create(executor, spec));
+                        SetCaseConditionKernel::Create(parent_, spec));
   }
   return &set_case_condition_kernel_;
 }
 
 absl::StatusOr<GpuCommandBuffer::SetForConditionKernel*>
-GpuCommandBuffer::GetSetForConditionKernel(StreamExecutorInterface* executor) {
+GpuCommandBuffer::GetSetForConditionKernel() {
   if (!set_for_condition_kernel_) {
     MultiKernelLoaderSpec spec(/*arity=*/3);
     spec.AddCudaPtxInMemory(gpu::GetSetForConditionKernel(),
                             "set_for_condition");
     TF_ASSIGN_OR_RETURN(set_for_condition_kernel_,
-                        SetForConditionKernel::Create(executor, spec));
+                        SetForConditionKernel::Create(parent_, spec));
   }
   return &set_for_condition_kernel_;
 }
 
 absl::StatusOr<GpuCommandBuffer::SetWhileConditionKernel*>
-GpuCommandBuffer::GetSetWhileConditionKernel(
-    StreamExecutorInterface* executor) {
+GpuCommandBuffer::GetSetWhileConditionKernel() {
   if (!set_while_condition_kernel_) {
     MultiKernelLoaderSpec spec(/*arity=*/2);
     spec.AddCudaPtxInMemory(gpu::GetSetWhileConditionKernel(),
                             "set_while_condition");
     TF_ASSIGN_OR_RETURN(set_while_condition_kernel_,
-                        SetWhileConditionKernel::Create(executor, spec));
+                        SetWhileConditionKernel::Create(parent_, spec));
   }
   return &set_while_condition_kernel_;
 }
 
-absl::StatusOr<GpuCommandBuffer::NoOpKernel*> GpuCommandBuffer::GetNoOpKernel(
-    StreamExecutorInterface* executor) {
+absl::StatusOr<GpuCommandBuffer::NoOpKernel*>
+GpuCommandBuffer::GetNoOpKernel() {
 #if !defined(TENSORFLOW_USE_ROCM)
   if (!noop_kernel_) {
     MultiKernelLoaderSpec spec(/*arity=*/0);
     spec.AddCudaPtxInMemory(gpu::kNoOpKernel, "noop");
-    TF_ASSIGN_OR_RETURN(noop_kernel_, NoOpKernel::Create(executor, spec));
+    TF_ASSIGN_OR_RETURN(noop_kernel_, NoOpKernel::Create(parent_, spec));
   }
   return &noop_kernel_;
 #else
@@ -327,13 +325,13 @@ absl::Status GpuCommandBuffer::CheckNumCommandBuffers(
 }
 
 absl::StatusOr<GpuGraphNodeHandle> GpuCommandBuffer::CreateBarrierNode(
-    StreamExecutorInterface* executor, const Dependencies& dependencies) {
+    const Dependencies& dependencies) {
   GpuGraphNodeHandle barrier_handle = nullptr;
 #if !defined(TENSORFLOW_USE_ROCM)
   // TODO(b/316343054): Instead of empty nodes we create no-op kernel nodes as
   // barriers because CUDA 12.3 does not support empty nodes inside
   // conditional command buffers. This should be fixed in CUDA 12.4.
-  TF_ASSIGN_OR_RETURN(NoOpKernel * noop, GetNoOpKernel(executor));
+  TF_ASSIGN_OR_RETURN(NoOpKernel * noop, GetNoOpKernel());
 
   TF_RETURN_IF_ERROR(GpuDriver::GraphAddKernelNode(
       &barrier_handle, graph_, dependencies, "noop",
@@ -362,8 +360,7 @@ GpuCommandBuffer::Dependencies GpuCommandBuffer::GetBarrierDependencies(
   return dependencies;
 }
 
-absl::Status GpuCommandBuffer::Barrier(StreamExecutorInterface* executor,
-                                       ExecutionScopeId execution_scope_id) {
+absl::Status GpuCommandBuffer::Barrier(ExecutionScopeId execution_scope_id) {
   ExecutionScope& execution_scope = execution_scopes_[execution_scope_id];
 
   if (state_ == State::kCreate) {
@@ -391,8 +388,7 @@ absl::Status GpuCommandBuffer::Barrier(StreamExecutorInterface* executor,
 
     // If we have multiple dependencies or no existing barriers we have to
     // create a new empty node acting as an execution barrier.
-    TF_ASSIGN_OR_RETURN(auto barrier_handle,
-                        CreateBarrierNode(executor, dependencies));
+    TF_ASSIGN_OR_RETURN(auto barrier_handle, CreateBarrierNode(dependencies));
     execution_scope.barriers.push_back({barrier_handle, true, nodes_offset});
     return absl::OkStatus();
   }
@@ -414,19 +410,18 @@ absl::Status GpuCommandBuffer::Barrier(StreamExecutorInterface* executor,
 }
 
 absl::Status GpuCommandBuffer::Barrier(
-    StreamExecutorInterface* executor,
     absl::Span<const ExecutionScopeId> execution_scope_ids) {
   // Nothing to synchronize here.
   if (execution_scope_ids.empty()) return absl::OkStatus();
 
   // Do not create two-level barriers for single execution scope.
   if (execution_scope_ids.size() == 1) {
-    return Barrier(executor, execution_scope_ids[0]);
+    return Barrier(execution_scope_ids[0]);
   }
 
   // Add a new barrier to every synchronized execution scope.
   for (ExecutionScopeId execution_scope_id : execution_scope_ids) {
-    TF_RETURN_IF_ERROR(Barrier(executor, execution_scope_id));
+    TF_RETURN_IF_ERROR(Barrier(execution_scope_id));
   }
 
   if (state_ == State::kCreate) {
@@ -438,8 +433,7 @@ absl::Status GpuCommandBuffer::Barrier(
     }
 
     // Create a new barrier that joins all per-scope barriers together.
-    TF_ASSIGN_OR_RETURN(auto barrier_handle,
-                        CreateBarrierNode(executor, dependencies));
+    TF_ASSIGN_OR_RETURN(auto barrier_handle, CreateBarrierNode(dependencies));
 
     // Broadcast new barrier to all participating execution scopes.
     for (ExecutionScopeId execution_scope_id : execution_scope_ids) {
@@ -470,17 +464,16 @@ absl::Status GpuCommandBuffer::Barrier(
   return UnsupportedStateError(state_);
 }
 
-absl::Status GpuCommandBuffer::Barrier(StreamExecutorInterface* executor,
-                                       ExecutionScopeId from_execution_scope_id,
+absl::Status GpuCommandBuffer::Barrier(ExecutionScopeId from_execution_scope_id,
                                        ExecutionScopeId to_execution_scope_id) {
   // If scopes are the same simply add a barrier to it.
   if (from_execution_scope_id == to_execution_scope_id) {
-    return Barrier(executor, from_execution_scope_id);
+    return Barrier(from_execution_scope_id);
   }
 
   // Create new barriers in both execution scopes.
-  TF_RETURN_IF_ERROR(Barrier(executor, from_execution_scope_id));
-  TF_RETURN_IF_ERROR(Barrier(executor, to_execution_scope_id));
+  TF_RETURN_IF_ERROR(Barrier(from_execution_scope_id));
+  TF_RETURN_IF_ERROR(Barrier(to_execution_scope_id));
 
   if (state_ == State::kCreate) {
     // Collect barriers from each scope as dependencies.
@@ -489,8 +482,7 @@ absl::Status GpuCommandBuffer::Barrier(StreamExecutorInterface* executor,
         execution_scopes_[to_execution_scope_id].barriers.back().handle};
 
     // Create a new barrier that joins `from` and `to` scopes.
-    TF_ASSIGN_OR_RETURN(auto barrier_handle,
-                        CreateBarrierNode(executor, dependencies));
+    TF_ASSIGN_OR_RETURN(auto barrier_handle, CreateBarrierNode(dependencies));
 
     // Add a new barrier only to the `to_execution_scope_id`.
     ExecutionScope& execution_scope = execution_scopes_[to_execution_scope_id];
@@ -759,9 +751,8 @@ GpuCommandBuffer::CreateConditionalNodes(
 }
 
 absl::Status GpuCommandBuffer::CreateConditionalCommand(
-    ExecutionScopeId execution_scope_id, StreamExecutorInterface* executor,
-    ConditionType type, SetConditionFn set_condition,
-    absl::Span<const ConditionBuilder> builders) {
+    ExecutionScopeId execution_scope_id, ConditionType type,
+    SetConditionFn set_condition, absl::Span<const ConditionBuilder> builders) {
   ExecutionScope& execution_scope = execution_scopes_[execution_scope_id];
 
   TF_RETURN_IF_ERROR(CheckNotFinalized());
@@ -776,7 +767,7 @@ absl::Status GpuCommandBuffer::CreateConditionalCommand(
     TF_RETURN_IF_ERROR(set_condition(execution_scope_id, handles));
 
     // Add a barrier between conditional handles and conditional nodes.
-    TF_RETURN_IF_ERROR(Barrier(executor, execution_scope_id));
+    TF_RETURN_IF_ERROR(Barrier(execution_scope_id));
 
     // Create conditional command buffer for each builder.
     TF_ASSIGN_OR_RETURN(
@@ -804,7 +795,7 @@ absl::Status GpuCommandBuffer::CreateConditionalCommand(
         set_condition(execution_scope_id, cond_cmd_buffers.handles));
 
     // Update a barrier between conditional handles and conditional nodes.
-    TF_RETURN_IF_ERROR(Barrier(executor, execution_scope_id));
+    TF_RETURN_IF_ERROR(Barrier(execution_scope_id));
 
     // Skip updating conditional nodes.
     execution_scope.update_state.node_idx += num_handles;
@@ -818,13 +809,10 @@ absl::Status GpuCommandBuffer::CreateConditionalCommand(
 }
 
 absl::Status GpuCommandBuffer::If(ExecutionScopeId execution_scope_id,
-                                  StreamExecutorInterface* executor,
                                   DeviceMemory<bool> predicate,
                                   Builder then_builder) {
-  DCHECK(executor == parent_);
-
   TF_ASSIGN_OR_RETURN(SetIfConditionKernel * set_if_condition,
-                      GetSetIfConditionKernel(executor));
+                      GetSetIfConditionKernel());
 
   auto set_cond_fn = [&](ExecutionScopeId id, ConditionalHandles handles) {
     return CommandBuffer::Launch(*set_if_condition, id, ThreadDim(), BlockDim(),
@@ -834,19 +822,16 @@ absl::Status GpuCommandBuffer::If(ExecutionScopeId execution_scope_id,
   std::array<ConditionBuilder, 1> builders = {
       ToConditionBuilder(std::move(then_builder))};
 
-  return CreateConditionalCommand(execution_scope_id, executor,
-                                  ConditionType::kIf, set_cond_fn, builders);
+  return CreateConditionalCommand(execution_scope_id, ConditionType::kIf,
+                                  set_cond_fn, builders);
 }
 
 absl::Status GpuCommandBuffer::IfElse(ExecutionScopeId execution_scope_id,
-                                      StreamExecutorInterface* executor,
                                       DeviceMemory<bool> predicate,
                                       Builder then_builder,
                                       Builder else_builder) {
-  DCHECK(executor == parent_);
-
   TF_ASSIGN_OR_RETURN(SetIfElseConditionKernel * set_if_else_condition,
-                      GetSetIfElseConditionKernel(executor));
+                      GetSetIfElseConditionKernel());
 
   auto set_cond_fn = [&](ExecutionScopeId id, ConditionalHandles handles) {
     return CommandBuffer::Launch(*set_if_else_condition, id, ThreadDim(),
@@ -857,16 +842,13 @@ absl::Status GpuCommandBuffer::IfElse(ExecutionScopeId execution_scope_id,
       ToConditionBuilder(std::move(then_builder)),
       ToConditionBuilder(std::move(else_builder))};
 
-  return CreateConditionalCommand(execution_scope_id, executor,
-                                  ConditionType::kIf, set_cond_fn, builders);
+  return CreateConditionalCommand(execution_scope_id, ConditionType::kIf,
+                                  set_cond_fn, builders);
 }
 
 absl::Status GpuCommandBuffer::Case(ExecutionScopeId execution_scope_id,
-                                    StreamExecutorInterface* executor,
                                     DeviceMemory<int32_t> index,
                                     std::vector<Builder> branches) {
-  DCHECK(executor == parent_);
-
   // TODO(ezhulenev): Relax this constraint, we can launch multiple back to back
   // kernels to update conditional handles in batches of size 8.
   if (branches.size() > 8) {
@@ -875,7 +857,7 @@ absl::Status GpuCommandBuffer::Case(ExecutionScopeId execution_scope_id,
   }
 
   TF_ASSIGN_OR_RETURN(SetCaseConditionKernel * set_case_condition,
-                      GetSetCaseConditionKernel(executor));
+                      GetSetCaseConditionKernel());
 
   auto set_cond_fn = [&](ExecutionScopeId id, ConditionalHandles handles) {
     int32_t num_handles = handles.size();
@@ -899,23 +881,20 @@ absl::Status GpuCommandBuffer::Case(ExecutionScopeId execution_scope_id,
     builders.push_back(ToConditionBuilder(std::move(branch)));
   }
 
-  return CreateConditionalCommand(execution_scope_id, executor,
-                                  ConditionType::kIf, set_cond_fn, builders);
+  return CreateConditionalCommand(execution_scope_id, ConditionType::kIf,
+                                  set_cond_fn, builders);
 }
 
 absl::Status GpuCommandBuffer::For(ExecutionScopeId execution_scope_id,
-                                   StreamExecutorInterface* executor,
                                    int32_t num_iteration,
                                    DeviceMemory<int32_t> loop_counter,
                                    Builder body_builder) {
-  DCHECK(executor == parent_);
-
   TF_ASSIGN_OR_RETURN(SetForConditionKernel * set_for_condition,
-                      GetSetForConditionKernel(executor));
+                      GetSetForConditionKernel());
 
   // Reset loop counter to zero.
   TF_RETURN_IF_ERROR(Memset(execution_scope_id, &loop_counter, uint32_t{0}, 1));
-  TF_RETURN_IF_ERROR(Barrier(executor, execution_scope_id));
+  TF_RETURN_IF_ERROR(Barrier(execution_scope_id));
 
   auto set_cond_fn = [&](ExecutionScopeId id, ConditionalHandles handles) {
     return CommandBuffer::Launch(*set_for_condition, id, ThreadDim(),
@@ -925,7 +904,7 @@ absl::Status GpuCommandBuffer::For(ExecutionScopeId execution_scope_id,
 
   auto body = [&](CommandBuffer* body, GpuGraphConditionalHandle handle) {
     TF_RETURN_IF_ERROR(body_builder(body));
-    TF_RETURN_IF_ERROR(body->Barrier(executor));
+    TF_RETURN_IF_ERROR(body->Barrier());
 
     // Decide if we want to continue loop iteration.
     return body->Launch(*set_for_condition, ThreadDim(), BlockDim(), handle,
@@ -934,23 +913,20 @@ absl::Status GpuCommandBuffer::For(ExecutionScopeId execution_scope_id,
 
   std::array<ConditionBuilder, 1> builders = {std::move(body)};
 
-  return CreateConditionalCommand(execution_scope_id, executor,
-                                  ConditionType::kWhile, set_cond_fn, builders);
+  return CreateConditionalCommand(execution_scope_id, ConditionType::kWhile,
+                                  set_cond_fn, builders);
 }
 
 absl::Status GpuCommandBuffer::While(ExecutionScopeId execution_scope_id,
-                                     StreamExecutorInterface* executor,
                                      DeviceMemory<bool> pred,
                                      ExecutionScopeBuilder cond_builder,
                                      Builder body_builder) {
-  DCHECK(executor == parent_);
-
   TF_ASSIGN_OR_RETURN(SetWhileConditionKernel * set_while_condition,
-                      GetSetWhileConditionKernel(executor));
+                      GetSetWhileConditionKernel());
 
   // Record condition commands into the parent command buffer.
   TF_RETURN_IF_ERROR(cond_builder(execution_scope_id, this));
-  TF_RETURN_IF_ERROR(Barrier(executor, execution_scope_id));
+  TF_RETURN_IF_ERROR(Barrier(execution_scope_id));
 
   auto set_cond_fn = [&](ExecutionScopeId id, ConditionalHandles handles) {
     return CommandBuffer::Launch(*set_while_condition, id, ThreadDim(),
@@ -959,17 +935,17 @@ absl::Status GpuCommandBuffer::While(ExecutionScopeId execution_scope_id,
 
   auto body = [&](CommandBuffer* body, GpuGraphConditionalHandle handle) {
     TF_RETURN_IF_ERROR(body_builder(body));
-    TF_RETURN_IF_ERROR(body->Barrier(executor));
+    TF_RETURN_IF_ERROR(body->Barrier());
     TF_RETURN_IF_ERROR(cond_builder(kDefaulExecutionScope, body));
-    TF_RETURN_IF_ERROR(body->Barrier(executor));
+    TF_RETURN_IF_ERROR(body->Barrier());
     return body->Launch(*set_while_condition, ThreadDim(), BlockDim(), handle,
                         pred);
   };
 
   std::array<ConditionBuilder, 1> builders = {std::move(body)};
 
-  return CreateConditionalCommand(execution_scope_id, executor,
-                                  ConditionType::kWhile, set_cond_fn, builders);
+  return CreateConditionalCommand(execution_scope_id, ConditionType::kWhile,
+                                  set_cond_fn, builders);
 }
 
 absl::Status GpuCommandBuffer::Finalize() {
diff --git a/third_party/xla/xla/stream_executor/gpu/gpu_command_buffer.h b/third_party/xla/xla/stream_executor/gpu/gpu_command_buffer.h
index 2ee78b56a6e2f1..aaa0666a7f89bd 100644
--- a/third_party/xla/xla/stream_executor/gpu/gpu_command_buffer.h
+++ b/third_party/xla/xla/stream_executor/gpu/gpu_command_buffer.h
@@ -77,15 +77,12 @@ class GpuCommandBuffer : public CommandBuffer {
                    bool is_owned_graph = true);
   ~GpuCommandBuffer() override;
 
-  absl::Status Barrier(StreamExecutorInterface* executor,
-                       ExecutionScopeId execution_scope_id) override;
+  absl::Status Barrier(ExecutionScopeId execution_scope_id) override;
 
   absl::Status Barrier(
-      StreamExecutorInterface* executor,
       absl::Span<const ExecutionScopeId> execution_scope_ids) override;
 
-  absl::Status Barrier(StreamExecutorInterface* executor,
-                       ExecutionScopeId from_execution_scope_id,
+  absl::Status Barrier(ExecutionScopeId from_execution_scope_id,
                        ExecutionScopeId to_execution_scope_id) override;
 
   absl::Status Launch(ExecutionScopeId execution_scope_id,
@@ -105,26 +102,22 @@ class GpuCommandBuffer : public CommandBuffer {
                       size_t num_elements) override;
 
   absl::Status If(ExecutionScopeId execution_scope_id,
-                  StreamExecutorInterface* executor,
                   DeviceMemory<bool> predicate, Builder then_builder) override;
 
   absl::Status IfElse(ExecutionScopeId execution_scope_id,
-                      StreamExecutorInterface* executor,
                       DeviceMemory<bool> predicate, Builder then_builder,
                       Builder else_builder) override;
 
   absl::Status Case(ExecutionScopeId execution_scope_id,
-                    StreamExecutorInterface* executor,
                     DeviceMemory<int32_t> index,
                     std::vector<Builder> branches) override;
 
-  absl::Status For(ExecutionScopeId execution_scope_id,
-                   StreamExecutorInterface* executor, int32_t num_iteration,
+  absl::Status For(ExecutionScopeId execution_scope_id, int32_t num_iteration,
                    DeviceMemory<int32_t> loop_counter,
                    Builder body_builder) override;
 
   absl::Status While(ExecutionScopeId execution_scope_id,
-                     StreamExecutorInterface* executor, DeviceMemory<bool> pred,
+                     DeviceMemory<bool> pred,
                      ExecutionScopeBuilder cond_builder,
                      Builder body_builder) override;
 
@@ -250,25 +243,20 @@ class GpuCommandBuffer : public CommandBuffer {
       absl::Span<const GpuGraphConditionalHandle> handles);
 
   absl::Status CreateConditionalCommand(
-      ExecutionScopeId execution_scope_id, StreamExecutorInterface* executor,
-      ConditionType type, SetConditionFn set_condition,
+      ExecutionScopeId execution_scope_id, ConditionType type,
+      SetConditionFn set_condition,
       absl::Span<const ConditionBuilder> builders);
 
   Dependencies GetBarrier(ExecutionScopeId execution_scope_id);
 
   // Returns loaded auxiliary kernels, or loads them on a given stream executor.
   // Loaded kernels owned by a current command buffer.
-  absl::StatusOr<SetIfConditionKernel*> GetSetIfConditionKernel(
-      StreamExecutorInterface* executor);
-  absl::StatusOr<SetIfElseConditionKernel*> GetSetIfElseConditionKernel(
-      StreamExecutorInterface* executor);
-  absl::StatusOr<SetCaseConditionKernel*> GetSetCaseConditionKernel(
-      StreamExecutorInterface* executor);
-  absl::StatusOr<SetForConditionKernel*> GetSetForConditionKernel(
-      StreamExecutorInterface* executor);
-  absl::StatusOr<SetWhileConditionKernel*> GetSetWhileConditionKernel(
-      StreamExecutorInterface* executor);
-  absl::StatusOr<NoOpKernel*> GetNoOpKernel(StreamExecutorInterface* executor);
+  absl::StatusOr<SetIfConditionKernel*> GetSetIfConditionKernel();
+  absl::StatusOr<SetIfElseConditionKernel*> GetSetIfElseConditionKernel();
+  absl::StatusOr<SetCaseConditionKernel*> GetSetCaseConditionKernel();
+  absl::StatusOr<SetForConditionKernel*> GetSetForConditionKernel();
+  absl::StatusOr<SetWhileConditionKernel*> GetSetWhileConditionKernel();
+  absl::StatusOr<NoOpKernel*> GetNoOpKernel();
 
   // Recursively disable all nodes corresponding to barriers (including nested
   // conditional command buffers). This is work around the fact that we can't
@@ -293,7 +281,7 @@ class GpuCommandBuffer : public CommandBuffer {
 
   // Creates a new no-op node acting as a barrier.
   absl::StatusOr<GpuGraphNodeHandle> CreateBarrierNode(
-      StreamExecutorInterface* executor, const Dependencies& dependencies);
+      const Dependencies& dependencies);
 
   // Collects a set of dependencies for a new barrier.
   Dependencies GetBarrierDependencies(ExecutionScopeId execution_scope_id);
diff --git a/third_party/xla/xla/stream_executor/gpu/gpu_command_buffer_test.cc b/third_party/xla/xla/stream_executor/gpu/gpu_command_buffer_test.cc
index 7ff757aab9aa5d..4440393cc3bb8d 100644
--- a/third_party/xla/xla/stream_executor/gpu/gpu_command_buffer_test.cc
+++ b/third_party/xla/xla/stream_executor/gpu/gpu_command_buffer_test.cc
@@ -382,22 +382,22 @@ TEST(GpuCommandBufferTest, Barriers) {
 
   auto record = [&](CommandBuffer* cmd_buffer, uint32_t bit_pattern) {
     // Check that root barrier ignored.
-    TF_RETURN_IF_ERROR(cmd_buffer->Barrier(executor));
+    TF_RETURN_IF_ERROR(cmd_buffer->Barrier());
     TF_RETURN_IF_ERROR(cmd_buffer->Memset(&buffers[0], bit_pattern + 0, 1));
     // Check barrier after a single command.
-    TF_RETURN_IF_ERROR(cmd_buffer->Barrier(executor));
+    TF_RETURN_IF_ERROR(cmd_buffer->Barrier());
     TF_RETURN_IF_ERROR(cmd_buffer->Memset(&buffers[1], bit_pattern + 1, 1));
     // Check that repeated barriers are no-op.
-    TF_RETURN_IF_ERROR(cmd_buffer->Barrier(executor));
-    TF_RETURN_IF_ERROR(cmd_buffer->Barrier(executor));
+    TF_RETURN_IF_ERROR(cmd_buffer->Barrier());
+    TF_RETURN_IF_ERROR(cmd_buffer->Barrier());
     TF_RETURN_IF_ERROR(cmd_buffer->Memset(&buffers[2], bit_pattern + 2, 1));
     TF_RETURN_IF_ERROR(cmd_buffer->Memset(&buffers[3], bit_pattern + 3, 1));
     // Check that barrier can have multiple dependencies.
-    TF_RETURN_IF_ERROR(cmd_buffer->Barrier(executor));
+    TF_RETURN_IF_ERROR(cmd_buffer->Barrier());
     TF_RETURN_IF_ERROR(cmd_buffer->Memset(&buffers[4], bit_pattern + 4, 1));
     TF_RETURN_IF_ERROR(cmd_buffer->Memset(&buffers[5], bit_pattern + 5, 1));
     // Check that barrier can be that last command.
-    TF_RETURN_IF_ERROR(cmd_buffer->Barrier(executor));
+    TF_RETURN_IF_ERROR(cmd_buffer->Barrier());
     return cmd_buffer->Finalize();
   };
 
@@ -476,8 +476,8 @@ TEST(GpuCommandBufferTest, IndependentExecutionScopes) {
     TF_RETURN_IF_ERROR(cmd_buffer->Memset(s0, &buffers[1], bit_pattern + 1, 1));
     TF_RETURN_IF_ERROR(cmd_buffer->Memset(s1, &buffers[2], bit_pattern + 2, 1));
     TF_RETURN_IF_ERROR(cmd_buffer->Memset(s1, &buffers[3], bit_pattern + 3, 1));
-    TF_RETURN_IF_ERROR(cmd_buffer->Barrier(executor, s0));
-    TF_RETURN_IF_ERROR(cmd_buffer->Barrier(executor, s1));
+    TF_RETURN_IF_ERROR(cmd_buffer->Barrier(s0));
+    TF_RETURN_IF_ERROR(cmd_buffer->Barrier(s1));
     return cmd_buffer->Finalize();
   };
 
@@ -548,7 +548,7 @@ TEST(GpuCommandBufferTest, ExecutionScopeBarriers) {
     TF_RETURN_IF_ERROR(cmd_buffer->Memset(s1, &buffers[2], bit_pattern + 2, 1));
     TF_RETURN_IF_ERROR(cmd_buffer->Memset(s1, &buffers[3], bit_pattern + 3, 1));
     // This will synchronize scopes 0 and 1 and also create an empty scope 2.
-    TF_RETURN_IF_ERROR(cmd_buffer->Barrier(executor, {s0, s1, s2}));
+    TF_RETURN_IF_ERROR(cmd_buffer->Barrier({s0, s1, s2}));
     TF_RETURN_IF_ERROR(cmd_buffer->Memset(s0, &buffers[4], bit_pattern + 4, 1));
     TF_RETURN_IF_ERROR(cmd_buffer->Memset(s1, &buffers[5], bit_pattern + 5, 1));
     TF_RETURN_IF_ERROR(cmd_buffer->Memset(s2, &buffers[6], bit_pattern + 6, 1));
@@ -639,7 +639,7 @@ TEST(GpuCommandBufferTest, ExecutionScopeOneDirectionalBarriers) {
     TF_RETURN_IF_ERROR(cmd_buffer->Memset(s1, &buffers[2], bit_pattern + 2, 1));
     TF_RETURN_IF_ERROR(cmd_buffer->Memset(s1, &buffers[3], bit_pattern + 3, 1));
     // This will synchronize scopes 0 and 1.
-    TF_RETURN_IF_ERROR(cmd_buffer->Barrier(executor, s0, s1));
+    TF_RETURN_IF_ERROR(cmd_buffer->Barrier(s0, s1));
     TF_RETURN_IF_ERROR(cmd_buffer->Memset(s0, &buffers[4], bit_pattern + 4, 1));
     TF_RETURN_IF_ERROR(cmd_buffer->Memset(s1, &buffers[5], bit_pattern + 5, 1));
     return cmd_buffer->Finalize();
@@ -721,7 +721,7 @@ TEST(GpuCommandBufferTest, ConditionalIf) {
 
   // Create a command buffer with a single conditional operation.
   auto cmd_buffer = CommandBuffer::Create(executor).value();
-  TF_ASSERT_OK(cmd_buffer->If(executor, pred, then_builder));
+  TF_ASSERT_OK(cmd_buffer->If(pred, then_builder));
   TF_ASSERT_OK(cmd_buffer->Finalize());
 
   TF_ASSERT_OK(executor->Submit(stream.get(), *cmd_buffer));
@@ -760,7 +760,7 @@ TEST(GpuCommandBufferTest, ConditionalIf) {
 
   // Update command buffer with a conditional to use new builder.
   TF_ASSERT_OK(cmd_buffer->Update());
-  TF_ASSERT_OK(cmd_buffer->If(executor, pred, then_builder));
+  TF_ASSERT_OK(cmd_buffer->If(pred, then_builder));
   TF_ASSERT_OK(cmd_buffer->Finalize());
 
   TF_ASSERT_OK(executor->Submit(stream.get(), *cmd_buffer));
@@ -818,7 +818,7 @@ TEST(GpuCommandBufferTest, ConditionalIfElse) {
 
   // Create a command buffer with a single conditional operation.
   auto cmd_buffer = CommandBuffer::Create(executor).value();
-  TF_ASSERT_OK(cmd_buffer->IfElse(executor, pred, then_builder, else_builder));
+  TF_ASSERT_OK(cmd_buffer->IfElse(pred, then_builder, else_builder));
   TF_ASSERT_OK(cmd_buffer->Finalize());
 
   TF_ASSERT_OK(executor->Submit(stream.get(), *cmd_buffer));
@@ -855,7 +855,7 @@ TEST(GpuCommandBufferTest, ConditionalIfElse) {
 
   // Update command buffer with a conditional to use new `else` builder.
   TF_ASSERT_OK(cmd_buffer->Update());
-  TF_ASSERT_OK(cmd_buffer->IfElse(executor, pred, then_builder, else_builder));
+  TF_ASSERT_OK(cmd_buffer->IfElse(pred, then_builder, else_builder));
   TF_ASSERT_OK(cmd_buffer->Finalize());
 
   TF_ASSERT_OK(executor->Submit(stream.get(), *cmd_buffer));
@@ -913,7 +913,7 @@ TEST(GpuCommandBufferTest, ConditionalCase) {
 
   // Create a command buffer with a single conditional operation.
   auto cmd_buffer = CommandBuffer::Create(executor).value();
-  TF_ASSERT_OK(cmd_buffer->Case(executor, index, {branch0, branch1}));
+  TF_ASSERT_OK(cmd_buffer->Case(index, {branch0, branch1}));
   TF_ASSERT_OK(cmd_buffer->Finalize());
 
   TF_ASSERT_OK(executor->Submit(stream.get(), *cmd_buffer));
@@ -992,8 +992,7 @@ TEST(GpuCommandBufferTest, ConditionalFor) {
 
   // Create a command buffer with a single conditional operation.
   auto cmd_buffer = CommandBuffer::Create(executor).value();
-  TF_ASSERT_OK(
-      cmd_buffer->For(executor, num_iters, loop_counter, body_builder));
+  TF_ASSERT_OK(cmd_buffer->For(num_iters, loop_counter, body_builder));
   TF_ASSERT_OK(cmd_buffer->Finalize());
 
   TF_ASSERT_OK(executor->Submit(stream.get(), *cmd_buffer));
@@ -1060,7 +1059,7 @@ TEST(GpuCommandBufferTest, ConditionalWhile) {
 
   // Create a command buffer with a single conditional operation.
   auto cmd_buffer = CommandBuffer::Create(executor).value();
-  TF_ASSERT_OK(cmd_buffer->While(executor, pred, cond_builder, body_builder));
+  TF_ASSERT_OK(cmd_buffer->While(pred, cond_builder, body_builder));
   TF_ASSERT_OK(cmd_buffer->Finalize());
 
   TF_ASSERT_OK(executor->Submit(stream.get(), *cmd_buffer));
@@ -1112,16 +1111,15 @@ TEST(GpuCommandBufferTest, ConditionalIfInExecutionScope) {
     TF_RETURN_IF_ERROR(cmd_buffer->Memset(s0, &buffers[1], bit_pattern + 1, 1));
 
     // Record If in execution scope #1
-    TF_RETURN_IF_ERROR(
-        cmd_buffer->If(s1, executor, pred, [&](CommandBuffer* then_cmd) {
-          return then_cmd->Memset(&buffers[2], bit_pattern + 2, 1);
-        }));
+    TF_RETURN_IF_ERROR(cmd_buffer->If(s1, pred, [&](CommandBuffer* then_cmd) {
+      return then_cmd->Memset(&buffers[2], bit_pattern + 2, 1);
+    }));
 
     // Create a barrier in execution scope #0.
-    TF_RETURN_IF_ERROR(cmd_buffer->Barrier(executor, s0));
+    TF_RETURN_IF_ERROR(cmd_buffer->Barrier(s0));
 
     // Create a barrier between two execution scopes.
-    TF_RETURN_IF_ERROR(cmd_buffer->Barrier(executor, {s0, s1}));
+    TF_RETURN_IF_ERROR(cmd_buffer->Barrier({s0, s1}));
 
     return cmd_buffer->Finalize();
   };
@@ -1210,7 +1208,7 @@ TEST(GpuCommandBufferTest, ConditionalWhileInExecutionScope) {
 
     // Record While in execution scope #1
     TF_RETURN_IF_ERROR(cmd_buffer->While(
-        s1, executor, pred,
+        s1, pred,
         // Loop cond: loop_counter++ < num_iters;
         [&](ExecutionScopeId id, CommandBuffer* cond_cmd) {
           return cond_cmd->Launch(inc_and_cmp, id, ThreadDim(), BlockDim(),
@@ -1222,7 +1220,7 @@ TEST(GpuCommandBufferTest, ConditionalWhileInExecutionScope) {
         }));
 
     // Create a barrier between two execution scopes.
-    TF_RETURN_IF_ERROR(cmd_buffer->Barrier(executor, {s0, s1}));
+    TF_RETURN_IF_ERROR(cmd_buffer->Barrier({s0, s1}));
 
     return cmd_buffer->Finalize();
   };

From ad6690b18aa7f9ced5b4645e696d87df5e34b211 Mon Sep 17 00:00:00 2001
From: Eugene Zhulenev <ezhulenev@google.com>
Date: Mon, 13 May 2024 09:23:16 -0700
Subject: [PATCH 029/478] [xla:ffi] Unify external and internal user data
 registration and use i64 type ids

- treat external and internal user data uniformly as an opaque pointer + optional deleter
- use the same i64 TypeId for internal and external types to avoid expensive string comparison on a hot path

PiperOrigin-RevId: 633234788
---
 third_party/xla/xla/ffi/BUILD                 |   5 +
 third_party/xla/xla/ffi/api/api.h             |   6 +-
 third_party/xla/xla/ffi/api/c_api.h           |  45 ++++--
 third_party/xla/xla/ffi/api/ffi.h             |  41 ++++-
 third_party/xla/xla/ffi/api/ffi_test.cc       |  12 +-
 third_party/xla/xla/ffi/call_frame.cc         |   6 +-
 third_party/xla/xla/ffi/execution_context.cc  |  76 +++++----
 third_party/xla/xla/ffi/execution_context.h   | 146 ++++++++++--------
 .../xla/xla/ffi/execution_context_test.cc     |  52 ++++---
 third_party/xla/xla/ffi/ffi.h                 |   2 +-
 third_party/xla/xla/ffi/ffi_api.cc            |  41 +++--
 third_party/xla/xla/ffi/ffi_test.cc           |  10 +-
 .../xla/xla/service/gpu/custom_call_test.cc   |   4 +-
 13 files changed, 276 insertions(+), 170 deletions(-)

diff --git a/third_party/xla/xla/ffi/BUILD b/third_party/xla/xla/ffi/BUILD
index dbc3785e8d5e56..c53835bf83a88d 100644
--- a/third_party/xla/xla/ffi/BUILD
+++ b/third_party/xla/xla/ffi/BUILD
@@ -35,10 +35,15 @@ cc_library(
     srcs = ["execution_context.cc"],
     hdrs = ["execution_context.h"],
     deps = [
+        "@com_google_absl//absl/base:core_headers",
         "@com_google_absl//absl/container:flat_hash_map",
         "@com_google_absl//absl/status",
         "@com_google_absl//absl/status:statusor",
         "@com_google_absl//absl/strings",
+        "@com_google_absl//absl/synchronization",
+        "@local_tsl//tsl/lib/gtl:int_type",
+        "@local_tsl//tsl/platform:logging",
+        "@local_tsl//tsl/platform:statusor",
     ],
 )
 
diff --git a/third_party/xla/xla/ffi/api/api.h b/third_party/xla/xla/ffi/api/api.h
index 75b678d9490c10..655d0ba1b656d5 100644
--- a/third_party/xla/xla/ffi/api/api.h
+++ b/third_party/xla/xla/ffi/api/api.h
@@ -150,10 +150,8 @@ XLA_FFI_Error* Ffi::RegisterStaticHandler(const XLA_FFI_Api* api,
   XLA_FFI_Handler_Register_Args args;
   args.struct_size = XLA_FFI_Handler_Register_Args_STRUCT_SIZE;
   args.priv = nullptr;
-  args.name = XLA_FFI_ByteSpan{XLA_FFI_ByteSpan_STRUCT_SIZE, nullptr,
-                               name.data(), name.size()};
-  args.platform = XLA_FFI_ByteSpan{XLA_FFI_ByteSpan_STRUCT_SIZE, nullptr,
-                                   platform.data(), platform.size()};
+  args.name = XLA_FFI_ByteSpan{name.data(), name.size()};
+  args.platform = XLA_FFI_ByteSpan{platform.data(), platform.size()};
   args.handler = handler;
   args.traits = traits;
   return api->XLA_FFI_Handler_Register(&args);
diff --git a/third_party/xla/xla/ffi/api/c_api.h b/third_party/xla/xla/ffi/api/c_api.h
index 39d4d8feef825e..243e5d3e2bea51 100644
--- a/third_party/xla/xla/ffi/api/c_api.h
+++ b/third_party/xla/xla/ffi/api/c_api.h
@@ -217,44 +217,38 @@ typedef enum {
 typedef struct XLA_FFI_ExecutionContext XLA_FFI_ExecutionContext;
 
 //===----------------------------------------------------------------------===//
-// Call frame
+// Primitives.
 //===----------------------------------------------------------------------===//
 
+// TypeId uniquely identifies a user-defined type in a given XLA FFI instance.
+struct XLA_FFI_TypeId {
+  int64_t type_id;
+};
+
 // We use byte spans to pass strings to handlers because strings might not be
 // null terminated, and even if they are, looking for a null terminator can
 // become very expensive in tight loops.
 struct XLA_FFI_ByteSpan {
-  size_t struct_size;
-  void* priv;
-
   const char* ptr;
   size_t len;
 };
 
-XLA_FFI_DEFINE_STRUCT_TRAITS(XLA_FFI_ByteSpan, len);
-
 // A struct to pass a scalar value to FFI handler.
 struct XLA_FFI_Scalar {
-  size_t struct_size;
-  void* priv;
-
   XLA_FFI_DataType dtype;
   void* value;
 };
 
-XLA_FFI_DEFINE_STRUCT_TRAITS(XLA_FFI_Scalar, value);
-
 // A struct to pass a dense array to FFI handler.
 struct XLA_FFI_Array {
-  size_t struct_size;
-  void* priv;
-
   XLA_FFI_DataType dtype;
   size_t size;
   void* data;
 };
 
-XLA_FFI_DEFINE_STRUCT_TRAITS(XLA_FFI_Array, data);
+//===----------------------------------------------------------------------===//
+// Call frame
+//===----------------------------------------------------------------------===//
 
 struct XLA_FFI_Args {
   size_t struct_size;
@@ -336,6 +330,24 @@ XLA_FFI_DEFINE_STRUCT_TRAITS(XLA_FFI_Handler_Register_Args, traits);
 typedef XLA_FFI_Error* XLA_FFI_Handler_Register(
     XLA_FFI_Handler_Register_Args* args);
 
+//===----------------------------------------------------------------------===//
+// TypeId
+//===----------------------------------------------------------------------===//
+
+struct XLA_FFI_TypeId_Register_Args {
+  size_t struct_size;
+  void* priv;
+
+  XLA_FFI_ByteSpan name;
+  XLA_FFI_TypeId* type_id;  // out
+};
+
+XLA_FFI_DEFINE_STRUCT_TRAITS(XLA_FFI_TypeId_Register_Args, type_id);
+
+// Registers user type `name` and returns a unique `type_id`.
+typedef XLA_FFI_Error* XLA_FFI_TypeId_Register(
+    XLA_FFI_TypeId_Register_Args* args);
+
 //===----------------------------------------------------------------------===//
 // ExecutionContext
 //===----------------------------------------------------------------------===//
@@ -345,7 +357,7 @@ struct XLA_FFI_ExecutionContext_Get_Args {
   void* priv;
 
   XLA_FFI_ExecutionContext* ctx;
-  XLA_FFI_ByteSpan id;
+  XLA_FFI_TypeId* type_id;
   void* data;  // out
 };
 
@@ -390,6 +402,7 @@ struct XLA_FFI_Api {
   _XLA_FFI_API_STRUCT_FIELD(XLA_FFI_Error_Destroy);
   _XLA_FFI_API_STRUCT_FIELD(XLA_FFI_Handler_Register);
   _XLA_FFI_API_STRUCT_FIELD(XLA_FFI_Stream_Get);
+  _XLA_FFI_API_STRUCT_FIELD(XLA_FFI_TypeId_Register);
   _XLA_FFI_API_STRUCT_FIELD(XLA_FFI_ExecutionContext_Get);
 };
 
diff --git a/third_party/xla/xla/ffi/api/ffi.h b/third_party/xla/xla/ffi/api/ffi.h
index addb7af543c8cf..3d56c00b94bda2 100644
--- a/third_party/xla/xla/ffi/api/ffi.h
+++ b/third_party/xla/xla/ffi/api/ffi.h
@@ -16,12 +16,12 @@ limitations under the License.
 #ifndef XLA_FFI_API_FFI_H_
 #define XLA_FFI_API_FFI_H_
 
-#include <string_view>
 #ifdef XLA_FFI_FFI_H_
 #error Two different XLA FFI implementations cannot be included together
 #endif  // XLA_FFI_FFI_H_
 
 #include <algorithm>
+#include <cassert>
 #include <complex>
 #include <cstddef>
 #include <cstdint>
@@ -29,6 +29,7 @@ limitations under the License.
 #include <optional>
 #include <ostream>
 #include <string>
+#include <string_view>
 #include <type_traits>
 #include <utility>
 #include <vector>
@@ -567,27 +568,51 @@ struct CtxDecoding<PlatformStream<T>> {
 // UserData
 //===----------------------------------------------------------------------===//
 
+// All user data types that are passed via the execution context must be
+// registered with the XLA FFI ahead of time to get unique type id.
+using TypeId = XLA_FFI_TypeId;  // NOLINT
+
+inline XLA_FFI_Error* RegisterType(const XLA_FFI_Api* api,
+                                   std::string_view name,
+                                   XLA_FFI_TypeId* type_id) {
+  XLA_FFI_TypeId_Register_Args args;
+  args.struct_size = XLA_FFI_TypeId_Register_Args_STRUCT_SIZE;
+  args.priv = nullptr;
+  args.name = XLA_FFI_ByteSpan{name.data(), name.size()};
+  args.type_id = type_id;
+  return api->XLA_FFI_TypeId_Register(&args);
+}
+
+#define XLA_FFI_REGISTER_TYPE(API, NAME, TYPE_ID) \
+  XLA_FFI_REGISTER_TYPE_(API, NAME, TYPE_ID, __COUNTER__)
+#define XLA_FFI_REGISTER_TYPE_(API, NAME, TYPE_ID, N)  \
+  XLA_FFI_ATTRIBUTE_UNUSED static const XLA_FFI_Error* \
+      xla_ffi_type_##N##_registered_ =                 \
+          [] { return ::xla::ffi::RegisterType(API, NAME, TYPE_ID); }()
+
 // A type tag for automatic decoding user data passed via the execution context.
-template <const char* id, typename T>
+template <typename T>
 struct UserData {};
 
-template <const char* id, typename T>
-struct CtxDecoding<UserData<id, T>> {
+template <typename T>
+struct CtxDecoding<UserData<T>> {
   using Type = T*;
 
+  static_assert(std::is_same_v<decltype(T::id), TypeId>,
+                "UserData type must have a static `TypeId id` field");
+
   static std::optional<Type> Decode(const XLA_FFI_Api* api,
                                     XLA_FFI_ExecutionContext* ctx,
                                     DiagnosticEngine& diagnostic) {
-    static constexpr std::string_view id_view = {id};
-
     XLA_FFI_ExecutionContext_Get_Args args;
     args.struct_size = XLA_FFI_ExecutionContext_Get_Args_STRUCT_SIZE;
     args.priv = nullptr;
     args.ctx = ctx;
-    args.id = XLA_FFI_ByteSpan{XLA_FFI_ByteSpan_STRUCT_SIZE, nullptr,
-                               id_view.data(), id_view.size()};
+    args.type_id = &T::id;
     args.data = nullptr;
 
+    assert(args.type_id->type_id > 0 && "type must be registered with XLA FFI");
+
     if (XLA_FFI_Error* err = api->XLA_FFI_ExecutionContext_Get(&args); err) {
       diagnostic.Emit("Failed to get user data from execution context: ")
           << internal::ErrorUtil::GetErrorMessage(api, err);
diff --git a/third_party/xla/xla/ffi/api/ffi_test.cc b/third_party/xla/xla/ffi/api/ffi_test.cc
index c47e16b22e7de3..b963b99c5689fd 100644
--- a/third_party/xla/xla/ffi/api/ffi_test.cc
+++ b/third_party/xla/xla/ffi/api/ffi_test.cc
@@ -352,17 +352,19 @@ TEST(FfiTest, PointerAttr) {
 }
 
 struct MyData {
+  static TypeId id;
   std::string str;
 };
 
-TEST(FfiTest, UserData) {
-  static constexpr char kId[] = "my_data";
+TypeId MyData::id = {};  // zero-initialize type id
+XLA_FFI_REGISTER_TYPE(GetXlaFfiApi(), "my_data", &MyData::id);
 
+TEST(FfiTest, UserData) {
   MyData data{"foo"};
-  auto deleter = +[](void*) {};
 
   ExecutionContext execution_context;
-  TF_ASSERT_OK(execution_context.Emplace(kId, &data, deleter));
+  TF_ASSERT_OK(execution_context.Insert(
+      ExecutionContext::TypeId(MyData::id.type_id), &data));
 
   CallFrameBuilder builder;
   auto call_frame = builder.Build();
@@ -372,7 +374,7 @@ TEST(FfiTest, UserData) {
     return Error::Success();
   };
 
-  auto handler = Ffi::Bind().Ctx<UserData<kId, MyData>>().To(fn);
+  auto handler = Ffi::Bind().Ctx<UserData<MyData>>().To(fn);
 
   ServiceExecutableRunOptions service_run_options;
   service_run_options.mutable_run_options()->set_ffi_execution_context(
diff --git a/third_party/xla/xla/ffi/call_frame.cc b/third_party/xla/xla/ffi/call_frame.cc
index 867c7ccb35edf2..ee17f61c7e29fb 100644
--- a/third_party/xla/xla/ffi/call_frame.cc
+++ b/third_party/xla/xla/ffi/call_frame.cc
@@ -135,19 +135,19 @@ struct CallFrame::Dictionary {
 struct CallFrame::Array {
   CallFrameBuilder::Array value;  // XLA_FFI_Array::data
 
-  XLA_FFI_Array array = {XLA_FFI_Array_STRUCT_SIZE, nullptr};
+  XLA_FFI_Array array = {};
 };
 
 struct CallFrame::Scalar {
   CallFrameBuilder::Scalar value;  // XLA_FFI_Scalar::value
 
-  XLA_FFI_Scalar scalar = {XLA_FFI_Scalar_STRUCT_SIZE, nullptr};
+  XLA_FFI_Scalar scalar = {};
 };
 
 struct CallFrame::String {
   std::string value;  // XLA_FFI_ByteSpan::ptr
 
-  XLA_FFI_ByteSpan span = {XLA_FFI_ByteSpan_STRUCT_SIZE, nullptr};
+  XLA_FFI_ByteSpan span = {};
 };
 
 struct CallFrame::NamedAttribute {
diff --git a/third_party/xla/xla/ffi/execution_context.cc b/third_party/xla/xla/ffi/execution_context.cc
index 53aaeb7943a41c..2ea5f8bf37f736 100644
--- a/third_party/xla/xla/ffi/execution_context.cc
+++ b/third_party/xla/xla/ffi/execution_context.cc
@@ -15,73 +15,87 @@ limitations under the License.
 
 #include "xla/ffi/execution_context.h"
 
+#include <atomic>
 #include <cstdint>
 #include <memory>
 #include <string>
 #include <string_view>
 #include <utility>
 
+#include "absl/base/attributes.h"
+#include "absl/base/const_init.h"
+#include "absl/container/flat_hash_map.h"
 #include "absl/status/status.h"
 #include "absl/status/statusor.h"
 #include "absl/strings/str_cat.h"
+#include "absl/synchronization/mutex.h"
 
 namespace xla::ffi {
 
-ExecutionContext::OpaqueUserData::OpaqueUserData(
-    void* data, OpaqueUserData::Deleter deleter)
+ABSL_CONST_INIT absl::Mutex type_registry_mutex(absl::kConstInit);
+
+using TypeRegistry = absl::flat_hash_map<std::string, ExecutionContext::TypeId>;
+static TypeRegistry& StaticTypeRegistry() {
+  static auto* registry = new TypeRegistry();
+  return *registry;
+}
+
+ExecutionContext::TypeId ExecutionContext::GetNextTypeId() {
+  static auto* counter = new std::atomic<int64_t>(1);
+  return TypeId(counter->fetch_add(1));
+}
+
+ExecutionContext::UserData::UserData(void* data, Deleter<void> deleter)
     : data_(data), deleter_(std::move(deleter)) {}
 
-ExecutionContext::OpaqueUserData::~OpaqueUserData() {
+ExecutionContext::UserData::~UserData() {
   if (deleter_) deleter_(data_);
 }
 
-absl::Status ExecutionContext::Emplace(std::string id, void* data,
-                                       OpaqueUserData::Deleter deleter) {
-  if (!data) return absl::InvalidArgumentError("User data must be not null");
+absl::StatusOr<ExecutionContext::TypeId>
+ExecutionContext::RegisterExternalTypeId(std::string_view name) {
+  absl::MutexLock lock(&type_registry_mutex);
+  auto& registry = StaticTypeRegistry();
 
-  auto emplaced = opaque_.emplace(
-      id, std::make_shared<OpaqueUserData>(data, std::move(deleter)));
+  // Try to emplace with type id zero and fill it with real type id only if we
+  // successfully acquired an entry for a given name.
+  auto emplaced = registry.emplace(name, TypeId(0));
   if (!emplaced.second) {
     return absl::AlreadyExistsError(
-        absl::StrCat("Opaque user data with id ", id,
-                     " already exists in execution context"));
+        absl::StrCat("Type id ", emplaced.first->second.value(),
+                     " already registered for type name ", name));
   }
-
-  return absl::OkStatus();
+  return emplaced.first->second = GetNextTypeId();
 }
 
-absl::StatusOr<std::shared_ptr<ExecutionContext::OpaqueUserData>>
-ExecutionContext::Lookup(std::string_view id) const {
-  auto it = opaque_.find(id);
-  if (it == opaque_.end()) {
-    return absl::NotFoundError(absl::StrCat("Opaque user data with id ", id,
-                                            " not found in execution context"));
-  }
-  return it->second;
+absl::Status ExecutionContext::Insert(TypeId type_id, void* data,
+                                      Deleter<void> deleter) {
+  return InsertUserData(type_id,
+                        std::make_unique<UserData>(data, std::move(deleter)));
 }
 
-absl::Status ExecutionContext::Insert(int64_t type_id,
-                                      std::shared_ptr<UserData> data) {
+absl::Status ExecutionContext::InsertUserData(TypeId type_id,
+                                              std::unique_ptr<UserData> data) {
   if (!data) return absl::InvalidArgumentError("User data must be not null");
 
-  auto emplaced = typed_.emplace(type_id, std::move(data));
+  auto emplaced = user_data_.emplace(type_id, std::move(data));
   if (!emplaced.second) {
     return absl::AlreadyExistsError(
-        absl::StrCat("User data with type id ", type_id,
+        absl::StrCat("User data with type id ", type_id.value(),
                      " already exists in execution context"));
   }
-
   return absl::OkStatus();
 }
 
-absl::StatusOr<std::shared_ptr<ExecutionContext::UserData>>
-ExecutionContext::Lookup(int64_t type_id) const {
-  auto it = typed_.find(type_id);
-  if (it == typed_.end()) {
-    return absl::NotFoundError(absl::StrCat("User data with type id ", type_id,
+absl::StatusOr<ExecutionContext::UserData*> ExecutionContext::LookupUserData(
+    TypeId type_id) const {
+  auto it = user_data_.find(type_id);
+  if (it == user_data_.end()) {
+    return absl::NotFoundError(absl::StrCat("User data with type id ",
+                                            type_id.value(),
                                             " not found in execution context"));
   }
-  return it->second;
+  return it->second.get();
 }
 
 }  // namespace xla::ffi
diff --git a/third_party/xla/xla/ffi/execution_context.h b/third_party/xla/xla/ffi/execution_context.h
index 21d3c85e1facb4..1a5250913830c4 100644
--- a/third_party/xla/xla/ffi/execution_context.h
+++ b/third_party/xla/xla/ffi/execution_context.h
@@ -16,110 +16,128 @@ limitations under the License.
 #ifndef XLA_FFI_EXECUTION_CONTEXT_H_
 #define XLA_FFI_EXECUTION_CONTEXT_H_
 
+#include <algorithm>
 #include <cstdint>
 #include <functional>
 #include <memory>
-#include <string>
 #include <string_view>
-#include <type_traits>
 #include <utility>
 
 #include "absl/container/flat_hash_map.h"
 #include "absl/status/status.h"
 #include "absl/status/statusor.h"
+#include "tsl/lib/gtl/int_type.h"
+#include "tsl/platform/logging.h"
+#include "tsl/platform/statusor.h"
 
 namespace xla::ffi {
 
 // Execution context is a container for forwarding arbitrary user data to FFI
-// handlers in the scope of a single execution. Execution context allows to pass
-// arbitrary user data to FFI handlers via the side channel that does not
-// require modifying HLO modules. There are two kinds of user data that can be
-// passed to FFI handlers:
+// handlers in the scope of a single XLA execution. Execution context allows to
+// pass arbitrary user data to FFI handlers via the side channel that does not
+// require modifying HLO modules.
 //
-// 1. Opaque data. This is a wrapper for an opaque user data pointer that is
-//    useful when FFI handler is registered in the dynamically loaded library
-//    and we do not know the type of the data and can only work with the opaque
-//    pointer.
+// From XLA FFI perspective user data is an opaque pointer that can be
+// forwarded to the FFI handler. We rely on type id to guarantee that we forward
+// user data of correct type. There are kinds of type ids:
 //
-// 2. Typed data. This is useful when the FFI handler is registered in the same
-//    process and we can rely on global static variable to assign ids to types
-//    and we don't need to worry about breaking C++ ABI.
+// 1. External type id. When FFI handlers defined in a dynamically loaded
+//    library, they must register types used in the execution context ahead
+//    of time and explicitly get a unique type id for them.
 //
-// For internal FFI handlers we always use typed data, and use opaque data only
-// if FFI handler has to be defined in a separate dynamically loaded library.
+// 2. Internal type id. When FFI handler defined in the same binary we rely
+//    on a global static registry to automatically assing type ids.
 //
 // Examples: FFI handler can register a per-execution cache in the execution
 // context and get access to it in the FFI handler, with a guarantee that it is
 // unique between separate calls to XLA execute.
 class ExecutionContext {
  public:
-  // A base class for typed user data used for FFI handlers registered in the
-  // same process where we can safely pass around C++ objects.
-  class UserData {
-   public:
-    virtual ~UserData() = default;
-  };
+  template <typename T>
+  using Deleter = std::function<void(T*)>;
+
+  TSL_LIB_GTL_DEFINE_INT_TYPE(TypeId, int64_t);
+
+  // Registers external type with a given name in a static type registry.
+  static absl::StatusOr<TypeId> RegisterExternalTypeId(std::string_view name);
+
+  // Inserts opaque user data with a given type id and optional deleter.
+  absl::Status Insert(TypeId type_id, void* data,
+                      Deleter<void> deleter = nullptr);
 
+  // Inserts typed user data of type `T` and optional deleter.
   template <typename T>
-  using IsUserData = std::enable_if_t<std::is_base_of_v<UserData, T>>;
+  absl::Status Insert(T* data, Deleter<T> deleter = nullptr);
 
-  // An RAII wrapper for opaque user data that is useful when FFI handler is
-  // registered in the dynamically loaded library and we do not know the type of
-  // the data and can only work with the opaque pointer.
-  class OpaqueUserData {
-   public:
-    using Deleter = std::function<void(void*)>;
+  // Emplaces typed user data constructed from `args`. Execution context
+  // becomes the owner of the constructed object.
+  template <typename T, typename... Args>
+  absl::Status Emplace(Args&&... args);
+
+  // Looks up typed execution context data of type `T`.
+  template <typename T>
+  absl::StatusOr<T*> Lookup() const {
+    TF_ASSIGN_OR_RETURN(auto user_data, LookupUserData(GetTypeId<T>()));
+    return static_cast<T*>(user_data->data());
+  }
+
+  // Looks up opaque execution context data with given `type_id`.
+  absl::StatusOr<void*> Lookup(TypeId type_id) const {
+    TF_ASSIGN_OR_RETURN(auto user_data, LookupUserData(type_id));
+    return user_data->data();
+  }
 
-    OpaqueUserData(void* data, Deleter deleter);
-    ~OpaqueUserData();
+ private:
+  // An RAII wrapper for opaque user data. Optional deleter will be called when
+  // UserData is destroyed together with the execution context. If deleter is
+  // nullptr then the caller is responsible for making sure that the pointer
+  // stays valid during the XLA execution and correctly destroyed afterwards.
+  class UserData {
+   public:
+    UserData(void* data, Deleter<void> deleter);
+    ~UserData();
 
-    OpaqueUserData(OpaqueUserData&) = delete;
-    OpaqueUserData& operator=(const OpaqueUserData&) = delete;
+    UserData(UserData&) = delete;
+    UserData& operator=(const UserData&) = delete;
 
     void* data() const { return data_; }
 
    private:
     void* data_;
-    Deleter deleter_;
+    Deleter<void> deleter_;
   };
 
-  // Emplaces opaque user data keyed by `id`.
-  absl::Status Emplace(std::string id, void* data,
-                       OpaqueUserData::Deleter deleter);
-
-  // Looks up opaque user data keyed by `id`.
-  absl::StatusOr<std::shared_ptr<OpaqueUserData>> Lookup(
-      std::string_view id) const;
-
-  // Emplaces typed user data constructed from `args`.
-  template <typename T, typename... Args, IsUserData<T>* = nullptr>
-  absl::Status Emplace(Args&&... args) {
-    return Insert(GetTypeId<T>(),
-                  std::make_shared<T>(std::forward<Args>(args)...));
-  }
+  static TypeId GetNextTypeId();
 
-  // Looks up typed execution context data of type `T`.
-  template <typename T, IsUserData<T>* = nullptr>
-  absl::StatusOr<std::shared_ptr<T>> Lookup() const {
-    auto user_data = Lookup(GetTypeId<T>());
-    if (!user_data.ok()) return user_data.status();
-    return std::static_pointer_cast<T>(*std::move(user_data));
-  }
-
- private:
-  template <typename T, IsUserData<T>* = nullptr>
-  static int64_t GetTypeId() {
-    static const char id = 0;
-    return reinterpret_cast<int64_t>(&id);
+  template <typename T>
+  static TypeId GetTypeId() {
+    static const TypeId id = GetNextTypeId();
+    return id;
   }
 
-  absl::Status Insert(int64_t type_id, std::shared_ptr<UserData> data);
-  absl::StatusOr<std::shared_ptr<UserData>> Lookup(int64_t type_id) const;
+  absl::Status InsertUserData(TypeId type_id, std::unique_ptr<UserData> data);
+  absl::StatusOr<UserData*> LookupUserData(TypeId type_id) const;
 
-  absl::flat_hash_map<int64_t, std::shared_ptr<UserData>> typed_;
-  absl::flat_hash_map<std::string, std::shared_ptr<OpaqueUserData>> opaque_;
+  absl::flat_hash_map<TypeId, std::unique_ptr<UserData>> user_data_;
 };
 
+template <typename T>
+absl::Status ExecutionContext::Insert(T* data, Deleter<T> deleter) {
+  return InsertUserData(GetTypeId<T>(),
+                        std::make_unique<UserData>(
+                            data, [deleter = std::move(deleter)](void* data) {
+                              if (deleter) deleter(static_cast<T*>(data));
+                            }));
+}
+
+template <typename T, typename... Args>
+absl::Status ExecutionContext::Emplace(Args&&... args) {
+  return InsertUserData(GetTypeId<T>(),
+                        std::make_unique<UserData>(
+                            new T(std::forward<Args>(args)...),
+                            [](void* data) { delete static_cast<T*>(data); }));
+}
+
 }  // namespace xla::ffi
 
 #endif  // XLA_FFI_EXECUTION_CONTEXT_H_
diff --git a/third_party/xla/xla/ffi/execution_context_test.cc b/third_party/xla/xla/ffi/execution_context_test.cc
index 159e529396b7e7..5c4cb3fdaab666 100644
--- a/third_party/xla/xla/ffi/execution_context_test.cc
+++ b/third_party/xla/xla/ffi/execution_context_test.cc
@@ -25,36 +25,52 @@ limitations under the License.
 
 namespace xla::ffi {
 
-struct StringUserData {
-  std::string data;
-};
-
-struct I32UserData : public ExecutionContext::UserData {
+struct I32UserData {
   explicit I32UserData(int32_t value) : value(value) {}
   int32_t value;
 };
 
-TEST(ExecutionContextTest, OpaqueUserData) {
-  StringUserData string_data = {"foo"};
-  auto deleter = [](void*) {};
+struct StrUserData {
+  explicit StrUserData(std::string value) : value(value) {}
+  std::string value;
+};
 
+TEST(ExecutionContextTest, EmplaceUserData) {
   ExecutionContext context;
-  TF_ASSERT_OK(context.Emplace("foo", &string_data, deleter));
+  TF_ASSERT_OK(context.Emplace<I32UserData>(42));
+  TF_ASSERT_OK(context.Emplace<StrUserData>("hello"));
 
-  TF_ASSERT_OK_AND_ASSIGN(auto opaque_data, context.Lookup("foo"));
-  ASSERT_NE(opaque_data, nullptr);
+  TF_ASSERT_OK_AND_ASSIGN(auto* i32_data, context.Lookup<I32UserData>());
+  TF_ASSERT_OK_AND_ASSIGN(auto* str_data, context.Lookup<StrUserData>());
 
-  StringUserData* user_data = static_cast<StringUserData*>(opaque_data->data());
-  EXPECT_EQ(user_data, &string_data);
+  ASSERT_NE(i32_data, nullptr);
+  ASSERT_NE(str_data, nullptr);
+  ASSERT_EQ(i32_data->value, 42);
+  ASSERT_EQ(str_data->value, "hello");
 }
 
-TEST(ExecutionContextTest, UserData) {
+TEST(ExecutionContextTest, InsertUserOwned) {
+  I32UserData user_data(42);
+
   ExecutionContext context;
-  TF_ASSERT_OK(context.Emplace<I32UserData>(42));
+  TF_ASSERT_OK(context.Insert(&user_data));
 
-  TF_ASSERT_OK_AND_ASSIGN(auto i32_data, context.Lookup<I32UserData>());
-  ASSERT_NE(i32_data, nullptr);
-  ASSERT_EQ(i32_data->value, 42);
+  TF_ASSERT_OK_AND_ASSIGN(auto* i32_data, context.Lookup<I32UserData>());
+  ASSERT_EQ(i32_data, &user_data);
+}
+
+TEST(ExecutionContextTest, InsertUserOwnedWithTypeId) {
+  TF_ASSERT_OK_AND_ASSIGN(
+      ExecutionContext::TypeId type_id,
+      ExecutionContext::RegisterExternalTypeId("I32UserData"));
+
+  I32UserData user_data(42);
+
+  ExecutionContext context;
+  TF_ASSERT_OK(context.Insert(type_id, &user_data));
+
+  TF_ASSERT_OK_AND_ASSIGN(auto* i32_data, context.Lookup(type_id));
+  ASSERT_EQ(i32_data, &user_data);
 }
 
 TEST(ExecutionContextTest, UserDataNotFound) {
diff --git a/third_party/xla/xla/ffi/ffi.h b/third_party/xla/xla/ffi/ffi.h
index 4620f8ef63ccc2..629e0f02071d48 100644
--- a/third_party/xla/xla/ffi/ffi.h
+++ b/third_party/xla/xla/ffi/ffi.h
@@ -331,7 +331,7 @@ struct UserData {};
 
 template <typename T>
 struct CtxDecoding<UserData<T>> {
-  using Type = std::shared_ptr<T>;
+  using Type = T*;
 
   static std::optional<Type> Decode(const XLA_FFI_Api* api,
                                     XLA_FFI_ExecutionContext* ctx,
diff --git a/third_party/xla/xla/ffi/ffi_api.cc b/third_party/xla/xla/ffi/ffi_api.cc
index d9b3e3874c897d..73c6c7aba412af 100644
--- a/third_party/xla/xla/ffi/ffi_api.cc
+++ b/third_party/xla/xla/ffi/ffi_api.cc
@@ -289,21 +289,35 @@ static XLA_FFI_Error* XLA_FFI_Stream_Get(XLA_FFI_Stream_Get_Args* args) {
   return nullptr;
 }
 
+static XLA_FFI_Error* XLA_FFI_TypeId_Register(
+    XLA_FFI_TypeId_Register_Args* args) {
+  XLA_FFI_RETURN_IF_ERROR(ActualStructSizeIsGreaterOrEqual(
+      "XLA_FFI_ExecutionContext_Get_Args",
+      XLA_FFI_ExecutionContext_Get_Args_STRUCT_SIZE, args->struct_size));
+
+  auto type_id = ExecutionContext::RegisterExternalTypeId(
+      std::string_view(args->name.ptr, args->name.len));
+  if (!type_id.ok()) {
+    return new XLA_FFI_Error{std::move(type_id).status()};
+  }
+
+  args->type_id->type_id = type_id->value();
+  return nullptr;
+}
+
 static XLA_FFI_Error* XLA_FFI_ExecutionContext_Get(
     XLA_FFI_ExecutionContext_Get_Args* args) {
   XLA_FFI_RETURN_IF_ERROR(ActualStructSizeIsGreaterOrEqual(
       "XLA_FFI_ExecutionContext_Get_Args",
       XLA_FFI_ExecutionContext_Get_Args_STRUCT_SIZE, args->struct_size));
-  XLA_FFI_RETURN_IF_ERROR(ActualStructSizeIsGreaterOrEqual(
-      "XLA_FFI_ByteSpan", XLA_FFI_ByteSpan_STRUCT_SIZE, args->id.struct_size));
 
-  auto opaque = args->ctx->execution_context->Lookup(
-      std::string_view(args->id.ptr, args->id.len));
-  if (!opaque.ok()) {
-    return new XLA_FFI_Error{std::move(opaque).status()};
+  auto user_data = args->ctx->execution_context->Lookup(
+      ExecutionContext::TypeId(args->type_id->type_id));
+  if (!user_data.ok()) {
+    return new XLA_FFI_Error{std::move(user_data).status()};
   }
 
-  args->data = (*opaque)->data();
+  args->data = *user_data;
   return nullptr;
 }
 
@@ -360,12 +374,13 @@ static XLA_FFI_Api api = {
 
     &internal_api,
 
-    XLA_FFI_Error_Create,          // creates error
-    XLA_FFI_Error_GetMessage,      // get error message
-    XLA_FFI_Error_Destroy,         // frees error
-    XLA_FFI_Handler_Register,      // registers handler
-    XLA_FFI_Stream_Get,            // returns platform specific stream
-    XLA_FFI_ExecutionContext_Get,  // returns execution context data
+    XLA_FFI_Error_Create,
+    XLA_FFI_Error_GetMessage,
+    XLA_FFI_Error_Destroy,
+    XLA_FFI_Handler_Register,
+    XLA_FFI_Stream_Get,
+    XLA_FFI_TypeId_Register,
+    XLA_FFI_ExecutionContext_Get,
 };
 
 const XLA_FFI_Api* GetXlaFfiApi() { return &api; }
diff --git a/third_party/xla/xla/ffi/ffi_test.cc b/third_party/xla/xla/ffi/ffi_test.cc
index c95e7e3d28026b..f2c05eaf32313d 100644
--- a/third_party/xla/xla/ffi/ffi_test.cc
+++ b/third_party/xla/xla/ffi/ffi_test.cc
@@ -615,24 +615,24 @@ TEST(FfiTest, RunOptionsCtx) {
   TF_ASSERT_OK(status);
 }
 
-struct MyData : public ExecutionContext::UserData {
-  explicit MyData(std::string str) : str(std::move(str)) {}
+struct StrUserData {
+  explicit StrUserData(std::string str) : str(std::move(str)) {}
   std::string str;
 };
 
 TEST(FfiTest, UserData) {
   ExecutionContext execution_context;
-  TF_ASSERT_OK(execution_context.Emplace<MyData>("foo"));
+  TF_ASSERT_OK(execution_context.Emplace<StrUserData>("foo"));
 
   CallFrameBuilder builder;
   auto call_frame = builder.Build();
 
-  auto fn = [&](std::shared_ptr<MyData> data) {
+  auto fn = [&](StrUserData* data) {
     EXPECT_EQ(data->str, "foo");
     return absl::OkStatus();
   };
 
-  auto handler = Ffi::Bind().Ctx<UserData<MyData>>().To(fn);
+  auto handler = Ffi::Bind().Ctx<UserData<StrUserData>>().To(fn);
 
   ServiceExecutableRunOptions opts;
   opts.mutable_run_options()->set_ffi_execution_context(&execution_context);
diff --git a/third_party/xla/xla/service/gpu/custom_call_test.cc b/third_party/xla/xla/service/gpu/custom_call_test.cc
index d381f4de52f31c..7eb92d96cc1dd8 100644
--- a/third_party/xla/xla/service/gpu/custom_call_test.cc
+++ b/third_party/xla/xla/service/gpu/custom_call_test.cc
@@ -724,13 +724,13 @@ TEST_F(CustomCallTest, WithCalledComputation) {
 
 // Arbitrary user-defined context passed via the execution context side channel
 // to a custom call handlers.
-struct SomeExtraContext : public ffi::ExecutionContext::UserData {
+struct SomeExtraContext {
   explicit SomeExtraContext(int32_t value) : value(value) {}
   int32_t value;
 };
 
 static absl::Status ExecutionContext(ffi::Result<ffi::BufferBase>,
-                                     std::shared_ptr<SomeExtraContext> ctx) {
+                                     SomeExtraContext* ctx) {
   if (ctx->value != 42) return absl::InternalError("Unexpected value");
   return absl::OkStatus();
 }

From 68f1e149ebc32b11529a4306b6e109bd8ef85746 Mon Sep 17 00:00:00 2001
From: David Dunleavy <ddunleavy@google.com>
Date: Mon, 13 May 2024 10:05:25 -0700
Subject: [PATCH 030/478] Teach `xla_test` about specific GPU backends

Inspired by https://github.com/openxla/xla/pull/11753

PiperOrigin-RevId: 633247596
---
 third_party/xla/.kokoro/linux/build.sh        |   2 +-
 third_party/xla/xla/service/gpu/BUILD         |  70 ++++++-----
 third_party/xla/xla/service/gpu/tests/BUILD   |  31 ++---
 third_party/xla/xla/tests/build_defs.bzl      | 109 ++++++++++++++++--
 third_party/xla/xla/tools/hlo_opt/BUILD       |   2 +-
 .../xla/xla/tools/multihost_hlo_runner/BUILD  |   3 +-
 6 files changed, 156 insertions(+), 61 deletions(-)

diff --git a/third_party/xla/.kokoro/linux/build.sh b/third_party/xla/.kokoro/linux/build.sh
index 87bf19cd7ad71f..2a72099e5f5148 100644
--- a/third_party/xla/.kokoro/linux/build.sh
+++ b/third_party/xla/.kokoro/linux/build.sh
@@ -54,7 +54,7 @@ RBE_FLAGS=""
 TARGET_FILTERS="-@local_tsl//tsl/platform:subprocess_test -@local_tsl//tsl/platform/cloud:google_auth_provider_test -@local_tsl//tsl/platform/cloud:oauth_client_test"
 
 if is_linux_gpu_job ; then
-    TAGS_FILTER="$TAGS_FILTER,gpu,requires-gpu-nvidia,-no_gpu"
+    TAGS_FILTER="$TAGS_FILTER,gpu_any,requires-gpu-nvidia,-no_gpu"
 
     # We are currently running XLA presubmits on machines with NVIDIA T4 GPUs,
     # which have a compute compatibility of 7.5. Se we filter out all the tests
diff --git a/third_party/xla/xla/service/gpu/BUILD b/third_party/xla/xla/service/gpu/BUILD
index a851ea99b9b182..34f2210e1d1b04 100644
--- a/third_party/xla/xla/service/gpu/BUILD
+++ b/third_party/xla/xla/service/gpu/BUILD
@@ -1683,7 +1683,7 @@ xla_test(
     srcs = if_gpu_is_configured(["gemm_algorithm_picker_test.cc"]),
     backends = [
         "gpu_v100",
-    ] + if_oss(["gpu"]),
+    ] + if_oss(["gpu_any"]),
     deps = [
         ":autotuner_util",
         ":backend_configs_cc",
@@ -1976,7 +1976,7 @@ xla_test(
     srcs = if_gpu_is_configured(["conv_algorithm_picker_test.cc"]),
     backends = [
         "gpu_v100",
-    ] + if_oss(["gpu"]),
+    ] + if_oss(["gpu_any"]),
     tags = [
         "noasan",
         "nomsan",
@@ -4636,7 +4636,7 @@ xla_test(
     },
     backends = [
         "gpu_a100",
-    ] + if_oss(["gpu"]),
+    ] + if_oss(["gpu_any"]),
     local_defines = if_cuda_is_configured(["GOOGLE_CUDA=1"]),
     shard_count = 10,
     deps = [
@@ -5299,19 +5299,22 @@ xla_test(
         "gpu_p100",
         "gpu_v100",
     ],
-    deps = if_cuda_is_configured([
-        ":dot_operand_converter",
-        "@com_google_googletest//:gtest",
-        "@com_google_absl//absl/strings",
-        "@com_google_absl//absl/strings:string_view",
-        "//xla:shape_util",
-        "//xla/hlo/ir:hlo",
-        "//xla/hlo/utils:hlo_matchers",
-        "//xla/service:pattern_matcher",
-        "//xla/tests:hlo_test_base",
-        "//xla/tests:xla_internal_test_main",
-        "@local_tsl//tsl/platform:statusor",
-    ]),
+    deps = if_cuda_is_configured(
+        [
+            ":dot_operand_converter",
+            "@com_google_googletest//:gtest",
+            "@com_google_absl//absl/strings",
+            "@com_google_absl//absl/strings:string_view",
+            "//xla:shape_util",
+            "//xla/hlo/ir:hlo",
+            "//xla/hlo/utils:hlo_matchers",
+            "//xla/service:pattern_matcher",
+            "//xla/tests:hlo_test_base",
+            "//xla/tests:xla_internal_test_main",
+            "@local_tsl//tsl/platform:statusor",
+        ],
+        ["@local_tsl//tsl/platform:test_main"],  # b/317293391
+    ),
 )
 
 cc_library(
@@ -5848,28 +5851,31 @@ xla_cc_test(
 
 xla_test(
     name = "determinism_test",
-    srcs = ["determinism_test.cc"],
+    srcs = if_gpu_is_configured(["determinism_test.cc"]),
     backends = [
         "gpu_a100",
     ],
     local_defines = if_cuda_is_configured(["GOOGLE_CUDA=1"]) + if_rocm_is_configured([
         "TENSORFLOW_USE_ROCM=1",
     ]),
-    deps = [
-        ":autotuner_util",
-        "//xla:literal",
-        "//xla:xla_proto_cc",
-        "//xla/hlo/ir:hlo",
-        "//xla/service/gpu/tests:gpu_codegen_test",
-        "//xla/stream_executor:device_description",
-        "//xla/stream_executor/gpu:gpu_timer",
-        "//xla/tests:hlo_test_base",
-        "//xla/tests:literal_test_util",
-        "//xla/tests:test_utils",
-        "@com_google_absl//absl/strings",
-        "@com_google_googletest//:gtest_main",
-        "@local_tsl//tsl/platform:statusor",
-    ],
+    deps = if_gpu_is_configured(
+        [
+            ":autotuner_util",
+            "@com_google_googletest//:gtest_main",
+            "@com_google_absl//absl/strings",
+            "//xla:literal",
+            "//xla:xla_proto_cc",
+            "//xla/hlo/ir:hlo",
+            "//xla/service/gpu/tests:gpu_codegen_test",
+            "//xla/stream_executor:device_description",
+            "//xla/stream_executor/gpu:gpu_timer",
+            "//xla/tests:hlo_test_base",
+            "//xla/tests:literal_test_util",
+            "//xla/tests:test_utils",
+            "@local_tsl//tsl/platform:statusor",
+        ],
+        ["@local_tsl//tsl/platform:test_main"],  # b/317293391
+    ),
 )
 
 cc_library(
diff --git a/third_party/xla/xla/service/gpu/tests/BUILD b/third_party/xla/xla/service/gpu/tests/BUILD
index cc7597cbaf63f6..8b11e1de46be2f 100644
--- a/third_party/xla/xla/service/gpu/tests/BUILD
+++ b/third_party/xla/xla/service/gpu/tests/BUILD
@@ -530,7 +530,7 @@ xla_test(
     srcs = ["gpu_kernel_tiling_test.cc"],
     backends = [
         "gpu_p100",
-    ] + if_oss(["gpu"]),
+    ] + if_oss(["gpu_any"]),
     deps = [
         ":gpu_codegen_test",
         "//xla:error_spec",
@@ -820,7 +820,7 @@ lit_test_suite(
         "//xla/tools/hlo_opt:gpu_specs/p100.txtpb",
         "//xla/tools/hlo_opt:gpu_specs/v100.txtpb",
     ],
-    default_tags = tf_cuda_tests_tags(),
+    default_tags = tf_cuda_tests_tags() + ["gpu_any"],
     tags_override = {
         "element_wise_row_vectorization.hlo": ["no_rocm"],
         "scatter_bf16.hlo": ["no_rocm"],
@@ -910,7 +910,7 @@ xla_test(
     srcs = ["tensor_float_32_global_var_test.cc"],
     backends = [
         "gpu_a100",
-    ] + if_oss(["gpu"]),
+    ] + if_oss(["gpu_any"]),
     deps = [
         "//xla:error_spec",
         "//xla/tests:hlo_test_base",
@@ -921,21 +921,24 @@ xla_test(
 
 xla_test(
     name = "gpu_sparse_dot_test",
-    srcs = ["gpu_sparse_dot_test.cc"],
+    srcs = if_cuda_is_configured(["gpu_sparse_dot_test.cc"]),
     backends = [
         "gpu_a100",
         "gpu_h100",
     ],
-    deps = [
-        ":gpu_codegen_test",
-        "//third_party/half:includes",
-        "//xla:literal",
-        "//xla:literal_util",
-        "//xla/tests:xla_internal_test_main",
-        "@com_google_absl//absl/strings",
-        "@com_google_absl//absl/types:span",
-        "@com_google_googletest//:gtest",
-    ],
+    deps = if_cuda_is_configured(
+        [
+            ":gpu_codegen_test",
+            "@com_google_googletest//:gtest",
+            "@com_google_absl//absl/strings",
+            "@com_google_absl//absl/types:span",
+            "//third_party/half:includes",
+            "//xla:literal",
+            "//xla:literal_util",
+            "//xla/tests:xla_internal_test_main",
+        ],
+        ["@local_tsl//tsl/platform:test_main"],  # b/317293391
+    ),
 )
 
 xla_test(
diff --git a/third_party/xla/xla/tests/build_defs.bzl b/third_party/xla/xla/tests/build_defs.bzl
index 7674095cb20e68..ae74353e783306 100644
--- a/third_party/xla/xla/tests/build_defs.bzl
+++ b/third_party/xla/xla/tests/build_defs.bzl
@@ -11,7 +11,81 @@ load(
 )
 load("//xla/tests:plugin.bzl", "plugins")
 
-all_backends = ["cpu", "gpu"] + list(plugins.keys())
+# Possible backend values for the GPU family.
+GPU_BACKENDS = [
+    "gpu_any",
+    "gpu_p100",
+    "gpu_v100",
+    "gpu_a100",
+    "gpu_h100",
+]
+
+# The generic "gpu" backend includes the actual backends in this list.
+GPU_DEFAULT_BACKENDS = [
+    "gpu_any",
+    "gpu_a100",
+    "gpu_h100",
+]
+
+_DEFAULT_BACKENDS = ["cpu"] + GPU_DEFAULT_BACKENDS
+
+_ALL_BACKENDS = ["cpu"] + GPU_BACKENDS + list(plugins.keys())
+
+# buildifier: disable=function-docstring
+def prepare_gpu_backend_data(backends, disabled_backends, backend_tags, backend_args):
+    # Expand "gpu" backend name into device specific backend names.
+    new_backends = [name for name in backends if name != "gpu"]
+    if len(new_backends) < len(backends):
+        new_backends.extend(GPU_DEFAULT_BACKENDS)
+
+    new_disabled_backends = [name for name in disabled_backends if name != "gpu"]
+    if len(new_disabled_backends) < len(disabled_backends):
+        new_disabled_backends.extend(GPU_BACKENDS)
+
+    new_backend_tags = {key: value for key, value in backend_tags.items() if key != "gpu"}
+    gpu_backend_tags = backend_tags.get("gpu", [])
+    for key in GPU_BACKENDS:
+        new_backend_tags.setdefault(key, gpu_backend_tags[:])
+
+    new_backend_args = {key: value for key, value in backend_args.items() if key != "gpu"}
+    if "gpu" in backend_args:
+        for key in GPU_BACKENDS:
+            new_backend_args.setdefault(key, backend_args["gpu"])
+
+    # Disable backends that don't meet the device requirements.
+    sm_requirements = {
+        "gpu_any": (0, 0),
+        "gpu_p100": (6, 0),
+        "gpu_v100": (7, 0),
+        "gpu_a100": (8, 0),
+        "gpu_h100": (9, 0),
+    }
+    for gpu_backend in GPU_BACKENDS:
+        all_tags = new_backend_tags[gpu_backend]
+        requires_gpu = [t for t in all_tags if t.startswith("requires-gpu-")]
+        requires_sm, only = None, False
+        for tag in requires_gpu:
+            if tag.startswith("requires-gpu-sm"):
+                version = tag.split("-")[2][2:]
+                sm = (int(version[:-1]), int(version[-1]))
+                if not requires_sm or sm < requires_sm:
+                    requires_sm = sm
+                if tag.endswith("-only"):
+                    only = True
+        if only:
+            disable = requires_sm != sm_requirements[gpu_backend]
+        else:
+            disable = requires_sm and requires_sm > sm_requirements[gpu_backend]
+
+        if disable:
+            new_disabled_backends.append(gpu_backend)
+        else:
+            sm_major, sm_minor = sm_requirements[gpu_backend]
+            sm_tag = "requires-gpu-nvidia" if sm_major == 0 else "requires-gpu-sm%s%s-only" % (sm_major, sm_minor)
+            new_backend_tags[gpu_backend] = [t for t in all_tags if t not in requires_gpu]
+            new_backend_tags[gpu_backend].append(sm_tag)
+
+    return new_backends, new_disabled_backends, new_backend_tags, new_backend_args
 
 def xla_test(
         name,
@@ -94,7 +168,11 @@ def xla_test(
 
     test_names = []
     if not backends:
-        backends = all_backends
+        backends = _DEFAULT_BACKENDS
+
+    # Expand "gpu" backend name to specific GPU backends and update tags.
+    backends, disabled_backends, backend_tags, backend_args = \
+        prepare_gpu_backend_data(backends, disabled_backends, backend_tags, backend_args)
 
     backends = [
         backend
@@ -108,15 +186,20 @@ def xla_test(
         this_backend_copts = []
         this_backend_args = backend_args.get(backend, [])
         this_backend_data = []
+        backend_deps = []
         if backend == "cpu":
-            backend_deps = ["//xla/service:cpu_plugin"]
-            backend_deps += ["//xla/tests:test_macros_cpu"]  # buildifier: disable=list-append
-        elif backend == "gpu":
-            backend_deps = if_gpu_is_configured(["//xla/service:gpu_plugin"])
-            backend_deps += if_gpu_is_configured(["//xla/tests:test_macros_gpu"])  # buildifier: disable=list-append
+            backend_deps += [
+                "//xla/service:cpu_plugin",
+                "//xla/tests:test_macros_cpu",
+            ]
+        elif backend in GPU_BACKENDS:
+            backend_deps += if_gpu_is_configured([
+                "//xla/service:gpu_plugin",
+                "//xla/tests:test_macros_%s" % backend,
+            ])
             this_backend_tags += tf_gpu_tests_tags()
+            this_backend_copts.append("-DXLA_TEST_BACKEND_GPU=1")
         elif backend in plugins:
-            backend_deps = []
             backend_deps += plugins[backend]["deps"]
             this_backend_copts += plugins[backend]["copts"]
             this_backend_tags += plugins[backend]["tags"]
@@ -197,11 +280,13 @@ def xla_test_library(
     """
 
     if not backends:
-        backends = all_backends
+        backends = _ALL_BACKENDS
 
     for backend in backends:
         this_backend_copts = []
-        if backend in ["cpu", "gpu"]:
+        if backend == "cpu":
+            backend_deps = ["//xla/tests:test_macros_cpu"]
+        elif backend in GPU_BACKENDS:
             backend_deps = ["//xla/tests:test_macros_%s" % backend]
         elif backend in plugins:
             backend_deps = plugins[backend]["deps"]
@@ -230,7 +315,7 @@ def generate_backend_suites(backends = []):  # buildifier: disable=unnamed-macro
     """
 
     if not backends:
-        backends = all_backends
+        backends = _ALL_BACKENDS
     for backend in backends:
         native.test_suite(
             name = "%s_tests" % backend,
@@ -244,7 +329,7 @@ def generate_backend_test_macros(backends = []):  # buildifier: disable=unnamed-
       backends: The list of backends to generate libraries for.
     """
     if not backends:
-        backends = all_backends
+        backends = _ALL_BACKENDS
     for backend in backends:
         manifest = ""
         if backend in plugins:
diff --git a/third_party/xla/xla/tools/hlo_opt/BUILD b/third_party/xla/xla/tools/hlo_opt/BUILD
index dd7f62731a27f9..1e042b61849b5d 100644
--- a/third_party/xla/xla/tools/hlo_opt/BUILD
+++ b/third_party/xla/xla/tools/hlo_opt/BUILD
@@ -176,7 +176,7 @@ lit_test_suite(
     ]),
     cfg = "//xla:lit.cfg.py",
     data = [":test_utilities"],
-    default_tags = tf_cuda_tests_tags(),
+    default_tags = tf_cuda_tests_tags() + ["gpu_any"],
     tags_override = {
         "gpu_hlo_ptx.hlo": ["no_rocm"],
     },
diff --git a/third_party/xla/xla/tools/multihost_hlo_runner/BUILD b/third_party/xla/xla/tools/multihost_hlo_runner/BUILD
index 5cdc1ce69ca7a4..eb9b29fa9a3364 100644
--- a/third_party/xla/xla/tools/multihost_hlo_runner/BUILD
+++ b/third_party/xla/xla/tools/multihost_hlo_runner/BUILD
@@ -16,7 +16,7 @@ build_test(
     name = "hlo_runner_main_build_test",
     tags = [
         "cpu",
-        "gpu",
+        "gpu_any",
     ],
     targets = [
         ":hlo_runner_main",
@@ -29,6 +29,7 @@ xla_cc_binary(
     srcs = ["hlo_runner_main.cc"],
     tags = [
         "gpu",
+        "gpu_any",
         "noasan",  # Exceeds linker limit.
         "nomac",
     ],

From 9e28a4c10ee0673c43990ef3802113f27ef317b8 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Mon, 13 May 2024 10:14:07 -0700
Subject: [PATCH 031/478] Fix typo in TFLite conversion doc.

PiperOrigin-RevId: 633250884
---
 tensorflow/lite/g3doc/models/convert/index.md | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/tensorflow/lite/g3doc/models/convert/index.md b/tensorflow/lite/g3doc/models/convert/index.md
index 094758708a0faa..5aee45c5c1ae1f 100644
--- a/tensorflow/lite/g3doc/models/convert/index.md
+++ b/tensorflow/lite/g3doc/models/convert/index.md
@@ -97,7 +97,7 @@ for your model:
 1. [Optimization flags](../../performance/model_optimization) allow you to
    specify the type of optimization to apply
    during conversion. The most commonly used optimization technique is
-   [post-training quanitization]().
+   [post-training quantization]().
 1. [Metadata flags](metadata) allow you to add metadata to the converted model
    which makes it easier to create platform specific wrapper code when deploying
    models on devices.
@@ -142,7 +142,7 @@ format model and a custom runtime environment for that model.
   converting your model.
 * See the [optimization overview](../../performance/model_optimization) for
   guidance on how to optimize your converted model using techniques like
-  [post-training quanitization](../../performance/post_training_quantization).
+  [post-training quantization](../../performance/post_training_quantization).
 * See the [Adding metadata overview](metadata) to learn how to add metadata to
   your models. Metadata provides other uses a description of your model as well
   as information that can be leveraged by code generators.

From cafd38d0cda9d074cc7b9a5deb817ef73352a946 Mon Sep 17 00:00:00 2001
From: Eugene Zhulenev <ezhulenev@google.com>
Date: Mon, 13 May 2024 10:27:27 -0700
Subject: [PATCH 032/478] [pjrt][xla:ffi] Add XLA FFI execution context to PjRt
 ExecuteContext

- Extend PjRt ExecuteContext base class to always include XLA FFI execution context so that end-users can pass data to FFI handlers.

- Add an end-to-end test to SE GPU PjrtClient.

PJRT C API changes will be in the followup PRs.

PiperOrigin-RevId: 633255392
---
 third_party/xla/xla/pjrt/BUILD                |  1 +
 third_party/xla/xla/pjrt/gpu/BUILD            |  5 ++
 .../xla/pjrt/gpu/se_gpu_pjrt_client_test.cc   | 62 +++++++++++++++++++
 third_party/xla/xla/pjrt/pjrt_executable.h    | 14 ++++-
 .../xla/pjrt/pjrt_stream_executor_client.cc   |  4 +-
 .../xla/xla/service/gpu/gpu_executable.cc     | 11 +++-
 third_party/xla/xla/service/gpu/runtime/BUILD |  1 +
 .../gpu/runtime/command_buffer_thunk.cc       |  2 +-
 .../service/gpu/runtime/custom_call_thunk.cc  |  3 +-
 .../xla/xla/service/gpu/runtime/thunk.cc      |  6 +-
 .../xla/xla/service/gpu/runtime/thunk.h       |  8 +++
 11 files changed, 109 insertions(+), 8 deletions(-)

diff --git a/third_party/xla/xla/pjrt/BUILD b/third_party/xla/xla/pjrt/BUILD
index 7e6c924261a19b..662c77584df5e7 100644
--- a/third_party/xla/xla/pjrt/BUILD
+++ b/third_party/xla/xla/pjrt/BUILD
@@ -252,6 +252,7 @@ cc_library(
         "//xla:xla_data_proto_cc",
         "//xla:xla_proto_cc",
         "//xla/client:executable_build_options",
+        "//xla/ffi:execution_context",
         "//xla/hlo/ir:hlo",
         "//xla/service:compiler",
         "//xla/service:computation_layout",
diff --git a/third_party/xla/xla/pjrt/gpu/BUILD b/third_party/xla/xla/pjrt/gpu/BUILD
index 9aec3e3b52cbf8..f8991a93a189e5 100644
--- a/third_party/xla/xla/pjrt/gpu/BUILD
+++ b/third_party/xla/xla/pjrt/gpu/BUILD
@@ -145,7 +145,10 @@ xla_cc_test(
         "//xla:status_macros",
         "//xla:statusor",
         "//xla:test",
+        "//xla:xla_data_proto_cc",
         "//xla/client:xla_computation",
+        "//xla/ffi",
+        "//xla/ffi:ffi_api",
         "//xla/pjrt:pjrt_client",
         "//xla/pjrt:pjrt_executable",
         "//xla/pjrt:pjrt_future",
@@ -153,6 +156,8 @@ xla_cc_test(
         "//xla/pjrt/distributed:in_memory_key_value_store",
         "//xla/service:gpu_plugin",
         "//xla/service:hlo_parser",
+        "//xla/service:platform_util",
+        "//xla/stream_executor",
         "//xla/tests:literal_test_util",
         "@com_google_absl//absl/log:check",
         "@com_google_absl//absl/status",
diff --git a/third_party/xla/xla/pjrt/gpu/se_gpu_pjrt_client_test.cc b/third_party/xla/xla/pjrt/gpu/se_gpu_pjrt_client_test.cc
index 038c3cb0552bb4..5a14973d215cc2 100644
--- a/third_party/xla/xla/pjrt/gpu/se_gpu_pjrt_client_test.cc
+++ b/third_party/xla/xla/pjrt/gpu/se_gpu_pjrt_client_test.cc
@@ -19,6 +19,7 @@ limitations under the License.
 
 #include <array>
 #include <cstdint>
+#include <cstring>
 #include <memory>
 #include <numeric>
 #include <string>
@@ -34,6 +35,8 @@ limitations under the License.
 #include "absl/time/clock.h"
 #include "absl/time/time.h"
 #include "xla/client/xla_computation.h"
+#include "xla/ffi/ffi.h"
+#include "xla/ffi/ffi_api.h"
 #include "xla/literal.h"
 #include "xla/literal_util.h"
 #include "xla/pjrt/distributed/in_memory_key_value_store.h"
@@ -43,12 +46,15 @@ limitations under the License.
 #include "xla/pjrt/pjrt_future.h"
 #include "xla/pjrt/pjrt_stream_executor_client.h"
 #include "xla/service/hlo_parser.h"
+#include "xla/service/platform_util.h"
 #include "xla/shape.h"
 #include "xla/shape_util.h"
 #include "xla/status_macros.h"
 #include "xla/statusor.h"
+#include "xla/stream_executor/stream.h"
 #include "xla/test.h"
 #include "xla/tests/literal_test_util.h"
+#include "xla/xla_data.pb.h"
 #include "tsl/lib/core/status_test_util.h"
 #include "tsl/platform/env.h"
 #include "tsl/platform/errors.h"
@@ -257,6 +263,62 @@ TEST(StreamExecutorGpuClientTest, RecvErrorNoDeadLock) {
                                 "of size 8 (0 already transferred)"));
 }
 
+// User-defined data type to be passed to FFI handler via the execute context
+// side channel.
+struct MemsetValue {
+  explicit MemsetValue(float value) : value(value) {}
+  float value;
+};
+
+static absl::Status MemsetFromValue(
+    se::Stream* stream, ffi::Result<ffi::BufferR1<PrimitiveType::F32>> result,
+    MemsetValue* memset_value) {
+  uint32_t pattern;
+  std::memcpy(&pattern, &memset_value->value, sizeof(pattern));
+
+  se::DeviceMemoryBase base = result->data;
+  return stream->Memset32(&base, pattern, result->data.size());
+}
+
+XLA_FFI_DEFINE_HANDLER(kMemsetFromValue, MemsetFromValue,
+                       ffi::Ffi::Bind()
+                           .Ctx<ffi::Stream>()
+                           .Ret<ffi::BufferR1<PrimitiveType::F32>>()
+                           .Ctx<ffi::UserData<MemsetValue>>());
+
+XLA_FFI_REGISTER_HANDLER(ffi::GetXlaFfiApi(), "MemsetFromValue",
+                         PlatformUtil::CanonicalPlatformName("GPU").value(),
+                         kMemsetFromValue);
+
+TEST(StreamExecutorGpuClientTest, ForwardUserDataToFfiHandler) {
+  static constexpr char const* kProgram = R"(
+    HloModule ffi_handler
+    ENTRY main {
+      ROOT %custom-call = f32[4] custom-call(),
+                          custom_call_target="MemsetFromValue",
+                          api_version=API_VERSION_TYPED_FFI
+    })";
+
+  TF_ASSERT_OK_AND_ASSIGN(auto client,
+                          GetStreamExecutorGpuClient(GpuClientOptions()));
+  TF_ASSERT_OK_AND_ASSIGN(auto executable,
+                          CompileExecutable(kProgram, *client));
+
+  ExecuteContext context;
+  TF_ASSERT_OK(context.ffi_context().Emplace<MemsetValue>(42.0f));
+
+  ExecuteOptions opts;
+  opts.context = &context;
+
+  auto result = executable->Execute(/*argument_handles=*/{{}}, opts);
+
+  TF_ASSERT_OK_AND_ASSIGN(std::shared_ptr<xla::Literal> result_literal,
+                          ExtractSingleResult(result));
+  EXPECT_TRUE(LiteralTestUtil::Equal(
+      LiteralUtil::CreateR1<float>({42.0f, 42.0f, 42.0f, 42.0f}),
+      *result_literal));
+}
+
 TEST(StreamExecutorGpuClientTest, ToLiteralAsync) {
   TF_ASSERT_OK_AND_ASSIGN(auto client,
                           GetStreamExecutorGpuClient(GpuClientOptions()));
diff --git a/third_party/xla/xla/pjrt/pjrt_executable.h b/third_party/xla/xla/pjrt/pjrt_executable.h
index e4477cfde66f8f..176d1c11342007 100644
--- a/third_party/xla/xla/pjrt/pjrt_executable.h
+++ b/third_party/xla/xla/pjrt/pjrt_executable.h
@@ -32,6 +32,7 @@ limitations under the License.
 #include "absl/strings/string_view.h"
 #include "absl/types/span.h"
 #include "xla/client/executable_build_options.h"
+#include "xla/ffi/execution_context.h"
 #include "xla/hlo/ir/hlo_module.h"
 #include "xla/layout.h"
 #include "xla/pjrt/compile_options.pb.h"
@@ -152,6 +153,15 @@ struct LoadOptions {
 class ExecuteContext {
  public:
   virtual ~ExecuteContext() = default;
+
+  ffi::ExecutionContext& ffi_context() { return ffi_context_; }
+  const ffi::ExecutionContext& ffi_context() const { return ffi_context_; }
+
+ private:
+  // XLA FFI execution context is a mechanism to attach arbitrary user data to
+  // a particular call of PjRtLoadedExecutable::Execute and forward it to custom
+  // calls implemented as XLA FFI handlers.
+  ffi::ExecutionContext ffi_context_;
 };
 
 struct PjRtTransferMetadata {
@@ -213,7 +223,9 @@ struct ExecuteOptions {
   // the launch IDs may be used by the runtime to detect the mismatch.
   int32_t launch_id = 0;
   // If non-null, an opaque context passed to an execution that may be used to
-  // supply additional arguments to a derived class of PjRtExecutable.
+  // supply additional arguments to a derived class of PjRtExecutable. It is
+  // a caller responsibility to ensure that the context is valid for the
+  // duration of the execution.
   const ExecuteContext* context = nullptr;
   // If true, check that the PjRtBuffer argument shapes match the compiled
   // shapes. Otherwise, any shape with the right size on device may be passed.
diff --git a/third_party/xla/xla/pjrt/pjrt_stream_executor_client.cc b/third_party/xla/xla/pjrt/pjrt_stream_executor_client.cc
index c726474f054c00..b59fd9f145de05 100644
--- a/third_party/xla/xla/pjrt/pjrt_stream_executor_client.cc
+++ b/third_party/xla/xla/pjrt/pjrt_stream_executor_client.cc
@@ -2772,7 +2772,9 @@ PjRtStreamExecutorLoadedExecutable::EnqueueExecution(
   if (run_options.launch_id() != 0) {
     VLOG(3) << "launch id for " << name() << ": " << run_options.launch_id();
   }
-
+  if (options.context != nullptr) {
+    run_options.set_ffi_execution_context(&options.context->ffi_context());
+  }
   // The choice of where we wait is arbitrary; the reason for the wait is
   // pacing to avoid problems such as memory fragmentation and running ahead
   // too far, not for correctness. Placing it before the executable launch
diff --git a/third_party/xla/xla/service/gpu/gpu_executable.cc b/third_party/xla/xla/service/gpu/gpu_executable.cc
index 4c56c84cdadbe8..c4c55a0aed85af 100644
--- a/third_party/xla/xla/service/gpu/gpu_executable.cc
+++ b/third_party/xla/xla/service/gpu/gpu_executable.cc
@@ -471,9 +471,14 @@ absl::Status ExecuteThunks(
 
   {  // Initialize thunks using prepared resources before execution.
     Thunk::InitializeParams initialize_params{
-        executor,           executable_source,           &buffer_allocations,
-        main_stream,        command_buffer_trace_stream, &collective_params,
-        &collective_cliques};
+        executor,
+        executable_source,
+        &buffer_allocations,
+        main_stream,
+        command_buffer_trace_stream,
+        &collective_params,
+        &collective_cliques,
+        run_options->run_options().ffi_execution_context()};
 
     tsl::profiler::TraceMe trace([&] { return "Thunks::Initialize"; });
     for (const std::unique_ptr<Thunk>& thunk : thunk_sequence) {
diff --git a/third_party/xla/xla/service/gpu/runtime/BUILD b/third_party/xla/xla/service/gpu/runtime/BUILD
index b648e175da5b62..d5ad69a37b933c 100644
--- a/third_party/xla/xla/service/gpu/runtime/BUILD
+++ b/third_party/xla/xla/service/gpu/runtime/BUILD
@@ -1093,6 +1093,7 @@ cc_library(
         ":nccl_clique",
         ":nccl_clique_key",
         "//xla:executable_run_options",
+        "//xla/ffi:execution_context",
         "//xla/hlo/ir:hlo",
         "//xla/service:buffer_assignment",
         "//xla/service:executable",
diff --git a/third_party/xla/xla/service/gpu/runtime/command_buffer_thunk.cc b/third_party/xla/xla/service/gpu/runtime/command_buffer_thunk.cc
index 2447d279f4ad76..4e74fdcd821c8b 100644
--- a/third_party/xla/xla/service/gpu/runtime/command_buffer_thunk.cc
+++ b/third_party/xla/xla/service/gpu/runtime/command_buffer_thunk.cc
@@ -153,7 +153,7 @@ absl::Status CommandBufferThunk::Initialize(const InitializeParams& params) {
       params.collective_cliques, /*device_to_host_stream=*/nullptr,
       /*host_to_device_stream=*/nullptr,
       /*send_device_memory_function=*/nullptr,
-      /*recv_device_memory_function=*/nullptr);
+      /*recv_device_memory_function=*/nullptr, params.ffi_execution_context);
 
   // If command buffer is in `kCreate` state it means that command buffer
   // sequence was never recorded into it. We initialize all command buffers
diff --git a/third_party/xla/xla/service/gpu/runtime/custom_call_thunk.cc b/third_party/xla/xla/service/gpu/runtime/custom_call_thunk.cc
index 46656a65cccac5..9c25e984252519 100644
--- a/third_party/xla/xla/service/gpu/runtime/custom_call_thunk.cc
+++ b/third_party/xla/xla/service/gpu/runtime/custom_call_thunk.cc
@@ -150,12 +150,13 @@ absl::Status CustomCallThunk::ExecuteFfiHandler(const ExecuteParams& params) {
   builder.AddAttributes(attrs.Build());
   CallFrame call_frame = builder.Build();
 
-  // TODO(ezhulenev): Remove `ServiceExecutableRunOptions` from FFI handler
+  // TODO(b/340104720): Remove `ServiceExecutableRunOptions` from FFI handler
   // execution context, as apparently it's not easily accessible from Thunk.
   ExecutableRunOptions run_options;
   run_options.set_stream(params.stream);
   run_options.set_allocator(params.buffer_allocations->memory_allocator());
   run_options.set_device_ordinal(params.buffer_allocations->device_ordinal());
+  run_options.set_ffi_execution_context(params.ffi_execution_context);
   ServiceExecutableRunOptions service_run_options(run_options);
 
   CallOptions options = {&service_run_options, called_computation_};
diff --git a/third_party/xla/xla/service/gpu/runtime/thunk.cc b/third_party/xla/xla/service/gpu/runtime/thunk.cc
index 402890a1bc6a95..f066aed07aaa94 100644
--- a/third_party/xla/xla/service/gpu/runtime/thunk.cc
+++ b/third_party/xla/xla/service/gpu/runtime/thunk.cc
@@ -33,6 +33,7 @@ limitations under the License.
 #include "absl/types/span.h"
 #include "mlir/IR/Operation.h"  // from @llvm-project
 #include "xla/executable_run_options.h"
+#include "xla/ffi/execution_context.h"
 #include "xla/hlo/ir/hlo_instruction.h"
 #include "xla/service/global_device_id.h"
 #include "xla/service/gpu/backend_configs.pb.h"
@@ -185,6 +186,7 @@ Thunk::ExecuteParams Thunk::ExecuteParams::Create(
                        run_options.run_options().host_to_device_stream(),
                        run_options.run_options().send_device_memory_function(),
                        run_options.run_options().recv_device_memory_function(),
+                       run_options.run_options().ffi_execution_context(),
                        additional_compute_streams);
 }
 
@@ -196,7 +198,7 @@ Thunk::ExecuteParams Thunk::ExecuteParams::CloneWithNewAllocations(
       params.collective_params, params.collective_cliques,
       params.device_to_host_stream, params.host_to_device_stream,
       params.send_device_memory_function, params.recv_device_memory_function,
-      params.additional_compute_streams);
+      params.ffi_execution_context, params.additional_compute_streams);
 }
 
 Thunk::ExecuteParams::ExecuteParams(
@@ -207,6 +209,7 @@ Thunk::ExecuteParams::ExecuteParams(
     se::Stream* host_to_device_stream,
     SendDeviceMemoryFunction* send_device_memory_function,
     RecvDeviceMemoryFunction* recv_device_memory_function,
+    const ffi::ExecutionContext* ffi_execution_context,
     ExecutionStreamIdMap additional_compute_streams)
     : buffer_allocations(buffer_allocations),
       stream(stream),
@@ -217,6 +220,7 @@ Thunk::ExecuteParams::ExecuteParams(
       host_to_device_stream(host_to_device_stream),
       send_device_memory_function(send_device_memory_function),
       recv_device_memory_function(recv_device_memory_function),
+      ffi_execution_context(ffi_execution_context),
       additional_compute_streams(additional_compute_streams) {}
 
 //===----------------------------------------------------------------------===//
diff --git a/third_party/xla/xla/service/gpu/runtime/thunk.h b/third_party/xla/xla/service/gpu/runtime/thunk.h
index 1df2405de6da84..7acd446896fe91 100644
--- a/third_party/xla/xla/service/gpu/runtime/thunk.h
+++ b/third_party/xla/xla/service/gpu/runtime/thunk.h
@@ -34,6 +34,7 @@ limitations under the License.
 #include "absl/types/span.h"
 #include "mlir/IR/Operation.h"  // from @llvm-project
 #include "xla/executable_run_options.h"
+#include "xla/ffi/execution_context.h"
 #include "xla/hlo/ir/hlo_instruction.h"
 #include "xla/service/buffer_assignment.h"
 #include "xla/service/global_device_id.h"
@@ -327,6 +328,9 @@ class Thunk {
 
     // Collective cliques acquired based on resource requests.
     CollectiveCliques* collective_cliques = nullptr;
+
+    // XLA FFI execution context.
+    const ffi::ExecutionContext* ffi_execution_context = nullptr;
   };
 
   //===--------------------------------------------------------------------===//
@@ -376,6 +380,9 @@ class Thunk {
     SendDeviceMemoryFunction* send_device_memory_function;
     RecvDeviceMemoryFunction* recv_device_memory_function;
 
+    // XLA FFI execution context.
+    const ffi::ExecutionContext* ffi_execution_context;
+
     // Additional compute streams on which thunks launch operations.
     ExecutionStreamIdMap additional_compute_streams;
 
@@ -390,6 +397,7 @@ class Thunk {
                   se::Stream* host_to_device_stream,
                   SendDeviceMemoryFunction* send_device_memory_function,
                   RecvDeviceMemoryFunction* recv_device_memory_function,
+                  const ffi::ExecutionContext* ffi_execution_context,
                   ExecutionStreamIdMap additional_compute_streams = {});
   };
 

From 3adb649d6f8f21d657d7fa22036f89baf324c06a Mon Sep 17 00:00:00 2001
From: George Karpenkov <cheshire@google.com>
Date: Mon, 13 May 2024 10:41:29 -0700
Subject: [PATCH 033/478] [XLA] [NFC] Factor out MaybeOwning into a separate
 class.

PiperOrigin-RevId: 633260396
---
 third_party/xla/xla/BUILD                     |   1 +
 third_party/xla/xla/literal.cc                |   6 +-
 third_party/xla/xla/literal.h                 |  89 +--------------
 third_party/xla/xla/maybe_owning.h            | 107 ++++++++++++++++++
 .../xla/xla/service/gpu/gpu_compiler.cc       | 105 ++++-------------
 third_party/xla/xla/util_test.cc              |  25 ++++
 6 files changed, 162 insertions(+), 171 deletions(-)
 create mode 100644 third_party/xla/xla/maybe_owning.h

diff --git a/third_party/xla/xla/BUILD b/third_party/xla/xla/BUILD
index b22f7021257508..2219ec3a37e78c 100644
--- a/third_party/xla/xla/BUILD
+++ b/third_party/xla/xla/BUILD
@@ -341,6 +341,7 @@ cc_library(
     hdrs = [
         "iterator_util.h",
         "map_util.h",
+        "maybe_owning.h",
         "overflow_util.h",
         "util.h",
     ],
diff --git a/third_party/xla/xla/literal.cc b/third_party/xla/xla/literal.cc
index 9454cdc4ec2c8c..d7329d76b37120 100644
--- a/third_party/xla/xla/literal.cc
+++ b/third_party/xla/xla/literal.cc
@@ -223,7 +223,11 @@ std::ostream& operator<<(std::ostream& out, const Literal& literal) {
 
 Shape* MutableLiteralBase::mutable_shape_do_not_use() {
   const Shape* const_shape = shape_.get();
-  Shape* shape = shape_.get_mutable(/*ensure_owned=*/true);
+  if (!shape_.OwnsPtr()) {
+    shape_ = MaybeOwningShapePtr(std::make_unique<Shape>(*shape_));
+  }
+  Shape* shape = shape_.get_mutable();
+
   if (shape != const_shape) {
     std::function<void(const Shape&, Piece*)> set_piece_shapes =
         [&set_piece_shapes](const Shape& shape, Piece* piece) {
diff --git a/third_party/xla/xla/literal.h b/third_party/xla/xla/literal.h
index 8ba5163cfc7dc0..4f8ce8e57fa3b0 100644
--- a/third_party/xla/xla/literal.h
+++ b/third_party/xla/xla/literal.h
@@ -44,6 +44,7 @@ limitations under the License.
 #include "xla/index_util.h"
 #include "xla/layout.h"
 #include "xla/layout_util.h"
+#include "xla/maybe_owning.h"
 #include "xla/primitive_util.h"
 #include "xla/printer.h"
 #include "xla/shape.h"
@@ -1302,96 +1303,10 @@ class MutableLiteralBase : public LiteralBase {
                                absl::Span<const int64_t> dest_base,
                                absl::Span<const int64_t> copy_size);
 
-  // A unique_ptr like class which may or may not have ownership of its pointer.
   // The literal may or may not own the storage of the shape. Creating/copying a
   // shape can incur significant overhead which in many case we'd like to avoid,
   // esp. for small literals.
-  class MaybeOwningShapePtr {
-   public:
-    MaybeOwningShapePtr() = default;
-    explicit MaybeOwningShapePtr(std::unique_ptr<Shape> unique)
-        : ptr_and_owning_bit_(TakeUnique(std::move(unique))) {}
-
-    explicit MaybeOwningShapePtr(const Shape* borrowed)
-        : ptr_and_owning_bit_(Borrow(borrowed)) {}
-
-    ~MaybeOwningShapePtr() { MaybeDeleteOwned(); }
-
-    const Shape* get() const {
-      return reinterpret_cast<const Shape*>(ptr_and_owning_bit_ & kPointerMask);
-    }
-    Shape* get_mutable(bool ensure_owned = false) {
-      const Shape* const_ptr = get();
-      // TODO(b/67651157): Remove this copy on write logic and combine get() and
-      // get_mutable() once we remove mutable_shape_do_not_use().
-      if (const_ptr && !OwnsPtr()) {
-        ptr_and_owning_bit_ = TakeUnique(std::make_unique<Shape>(*const_ptr));
-        const_ptr = get();
-      }
-      DCHECK(OwnsPtr());
-      return const_cast<Shape*>(const_ptr);
-    }
-    const Shape* operator->() const { return get(); }
-    const Shape& operator*() const { return *get(); }
-
-    MaybeOwningShapePtr& operator=(std::unique_ptr<Shape> unique) {
-      MaybeDeleteOwned();
-      ptr_and_owning_bit_ = TakeUnique(std::move(std::move(unique)));
-      return *this;
-    }
-
-    MaybeOwningShapePtr& operator=(const Shape* borrowed) {
-      MaybeDeleteOwned();
-      ptr_and_owning_bit_ = Borrow(borrowed);
-      return *this;
-    }
-
-    MaybeOwningShapePtr& operator=(MaybeOwningShapePtr&& other) {
-      using std::swap;
-      swap(ptr_and_owning_bit_, other.ptr_and_owning_bit_);
-      return *this;
-    }
-
-    MaybeOwningShapePtr(const MaybeOwningShapePtr&) = delete;
-    MaybeOwningShapePtr(MaybeOwningShapePtr&& other)
-        : ptr_and_owning_bit_(other.ptr_and_owning_bit_) {
-      other.ptr_and_owning_bit_ = 0;
-    }
-
-    MaybeOwningShapePtr Clone() const {
-      const Shape* ptr = get();
-      if (ptr && OwnsPtr()) {
-        return MaybeOwningShapePtr(std::make_unique<Shape>(*ptr));
-      }
-      return MaybeOwningShapePtr(ptr);
-    }
-
-   private:
-    enum : uint64_t {
-      kOwningBitMask = 1UL,
-      kPointerMask = ~kOwningBitMask,
-    };
-    static intptr_t TakeUnique(std::unique_ptr<Shape> unique) {
-      Shape* released = unique.release();
-      DCHECK_EQ(reinterpret_cast<intptr_t>(released) & kOwningBitMask, 0);
-      return reinterpret_cast<intptr_t>(released) | kOwningBitMask;
-    }
-
-    static intptr_t Borrow(const Shape* borrowed) {
-      DCHECK_EQ(reinterpret_cast<intptr_t>(borrowed) & kOwningBitMask, 0);
-      return reinterpret_cast<intptr_t>(borrowed);
-    }
-
-    bool OwnsPtr() const { return kOwningBitMask & ptr_and_owning_bit_; }
-
-    void MaybeDeleteOwned() {
-      if (OwnsPtr()) {
-        delete get();
-      }
-    }
-
-    intptr_t ptr_and_owning_bit_ = 0;
-  };
+  using MaybeOwningShapePtr = MaybeOwning<Shape>;
 
   // The parent class borrows this shape.
   MaybeOwningShapePtr shape_;
diff --git a/third_party/xla/xla/maybe_owning.h b/third_party/xla/xla/maybe_owning.h
new file mode 100644
index 00000000000000..4f32472ecb2f95
--- /dev/null
+++ b/third_party/xla/xla/maybe_owning.h
@@ -0,0 +1,107 @@
+/* Copyright 2024 The OpenXLA Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+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.
+==============================================================================*/
+
+#ifndef XLA_MAYBE_OWNING_H_
+#define XLA_MAYBE_OWNING_H_
+
+#include <cstdint>
+#include <memory>
+
+// A unique_ptr like class which may or may not have ownership of its pointer.
+// Uses least significant bit of the pointer to indicate ownership.
+template <typename T>
+class MaybeOwning final {
+ public:
+  MaybeOwning() = default;
+  explicit MaybeOwning(std::unique_ptr<T> unique)
+      : ptr_and_owning_bit_(TakeUnique(std::move(unique))) {}
+
+  explicit MaybeOwning(const T* borrowed)
+      : ptr_and_owning_bit_(Borrow(borrowed)) {}
+
+  ~MaybeOwning() { MaybeDeleteOwned(); }
+
+  const T* get() const { return RemoveMask(); }
+
+  T* get_mutable() { return RemoveMask(); }
+
+  const T* operator->() const { return get(); }
+  const T& operator*() const { return *get(); }
+
+  MaybeOwning<T>& operator=(std::unique_ptr<T> unique) {
+    MaybeDeleteOwned();
+    ptr_and_owning_bit_ = TakeUnique(std::move(std::move(unique)));
+    return *this;
+  }
+
+  MaybeOwning& operator=(const T* borrowed) {
+    MaybeDeleteOwned();
+    ptr_and_owning_bit_ = Borrow(borrowed);
+    return *this;
+  }
+
+  MaybeOwning& operator=(MaybeOwning&& other) {
+    using std::swap;
+    swap(ptr_and_owning_bit_, other.ptr_and_owning_bit_);
+    return *this;
+  }
+
+  MaybeOwning(const MaybeOwning&) = delete;
+  MaybeOwning(MaybeOwning&& other)
+      : ptr_and_owning_bit_(other.ptr_and_owning_bit_) {
+    other.ptr_and_owning_bit_ = 0;
+  }
+
+  MaybeOwning Clone() const {
+    const T* ptr = get();
+    if (ptr && OwnsPtr()) {
+      return MaybeOwning(std::make_unique<T>(*ptr));
+    }
+    return MaybeOwning(ptr);
+  }
+
+  bool OwnsPtr() const { return kOwningBitMask & ptr_and_owning_bit_; }
+
+ private:
+  enum : uint64_t {
+    kOwningBitMask = 1UL,
+    kPointerMask = ~kOwningBitMask,
+  };
+
+  T* RemoveMask() const {
+    return reinterpret_cast<T*>(ptr_and_owning_bit_ & kPointerMask);
+  }
+
+  static intptr_t TakeUnique(std::unique_ptr<T> unique) {
+    T* released = unique.release();
+    DCHECK_EQ(reinterpret_cast<intptr_t>(released) & kOwningBitMask, 0);
+    return reinterpret_cast<intptr_t>(released) | kOwningBitMask;
+  }
+
+  static intptr_t Borrow(const T* borrowed) {
+    DCHECK_EQ(reinterpret_cast<intptr_t>(borrowed) & kOwningBitMask, 0);
+    return reinterpret_cast<intptr_t>(borrowed);
+  }
+
+  void MaybeDeleteOwned() {
+    if (OwnsPtr()) {
+      delete get();
+    }
+  }
+
+  intptr_t ptr_and_owning_bit_ = 0;
+};
+
+#endif  // XLA_MAYBE_OWNING_H_
diff --git a/third_party/xla/xla/service/gpu/gpu_compiler.cc b/third_party/xla/xla/service/gpu/gpu_compiler.cc
index 02673628e26980..8fc3b97cee7654 100644
--- a/third_party/xla/xla/service/gpu/gpu_compiler.cc
+++ b/third_party/xla/xla/service/gpu/gpu_compiler.cc
@@ -65,6 +65,7 @@ limitations under the License.
 #include "xla/hlo/ir/hlo_opcode.h"
 #include "xla/hlo/ir/hlo_schedule.h"
 #include "xla/hlo/transforms/hlo_constant_splitter.h"
+#include "xla/maybe_owning.h"
 #include "xla/service/all_gather_broadcast_reorder.h"
 #include "xla/service/all_gather_combiner.h"
 #include "xla/service/all_reduce_combiner.h"
@@ -259,38 +260,10 @@ limitations under the License.
 namespace xla {
 namespace gpu {
 namespace {
-// A class for storing either an owned thread pool or a non-owning pointer to an
-// external thread pool.
-class MaybeOwningThreadPool {
- public:
-  // Gets or creates a thread pool.
-  //
-  // See the code for the logic.
-  static MaybeOwningThreadPool GetOrCreate(
-      int parallelism, tsl::thread::ThreadPool* default_thread_pool,
-      int default_parallelism);
-
-  // Not owning (nullptr).
-  MaybeOwningThreadPool();
-  // Not owning.
-  explicit MaybeOwningThreadPool(tsl::thread::ThreadPool* thread_pool);
-  // Owning.
-  explicit MaybeOwningThreadPool(
-      std::unique_ptr<tsl::thread::ThreadPool> thread_pool);
-  tsl::thread::ThreadPool* get();
-  const tsl::thread::ThreadPool* get() const;
-  tsl::thread::ThreadPool* operator->();
-  const tsl::thread::ThreadPool* operator->() const;
-  explicit operator bool() const;
-  bool operator!() const;
 
- private:
-  std::variant<tsl::thread::ThreadPool*,
-               std::unique_ptr<tsl::thread::ThreadPool>>
-      thread_pool_;
-};
+using MaybeOwningThreadPool = MaybeOwning<tsl::thread::ThreadPool>;
 
-/*static*/ MaybeOwningThreadPool MaybeOwningThreadPool::GetOrCreate(
+MaybeOwningThreadPool CreateMaybeOwningThreadPool(
     int parallelism, tsl::thread::ThreadPool* default_thread_pool,
     int default_parallelism) {
   CHECK_GE(parallelism, 0);
@@ -320,41 +293,6 @@ class MaybeOwningThreadPool {
   }
 }
 
-MaybeOwningThreadPool::MaybeOwningThreadPool() : thread_pool_(nullptr) {}
-
-MaybeOwningThreadPool::MaybeOwningThreadPool(
-    tsl::thread::ThreadPool* thread_pool)
-    : thread_pool_(thread_pool) {}
-
-MaybeOwningThreadPool::MaybeOwningThreadPool(
-    std::unique_ptr<tsl::thread::ThreadPool> thread_pool)
-    : thread_pool_(std::move(thread_pool)) {}
-
-tsl::thread::ThreadPool* MaybeOwningThreadPool::get() {
-  if (std::holds_alternative<tsl::thread::ThreadPool*>(thread_pool_)) {
-    return std::get<tsl::thread::ThreadPool*>(thread_pool_);
-  }
-  return std::get<std::unique_ptr<tsl::thread::ThreadPool>>(thread_pool_).get();
-}
-
-const tsl::thread::ThreadPool* MaybeOwningThreadPool::get() const {
-  return const_cast<MaybeOwningThreadPool*>(this)->get();
-}
-
-tsl::thread::ThreadPool* MaybeOwningThreadPool::operator->() {
-  tsl::thread::ThreadPool* thread_pool = get();
-  CHECK_NE(thread_pool, nullptr);
-  return thread_pool;
-}
-
-const tsl::thread::ThreadPool* MaybeOwningThreadPool::operator->() const {
-  return const_cast<MaybeOwningThreadPool*>(this)->operator->();
-}
-
-MaybeOwningThreadPool::operator bool() const { return get() != nullptr; }
-
-bool MaybeOwningThreadPool::operator!() const { return get() == nullptr; }
-
 absl::StatusOr<AutotuneConfig> GetAutotuneConfig(
     se::StreamExecutor* stream_exec, const DebugOptions& debug_options,
     const GpuCompiler::CompileOptions& options,
@@ -1230,7 +1168,7 @@ absl::Status GpuCompiler::OptimizeHloModule(
   CheckNotScheduled(hlo_module);
   LogDebugOptions(hlo_module);
 
-  MaybeOwningThreadPool thread_pool = MaybeOwningThreadPool::GetOrCreate(
+  MaybeOwningThreadPool thread_pool = CreateMaybeOwningThreadPool(
       /*parallelism=*/hlo_module->config()
           .debug_options()
           .xla_gpu_force_compilation_parallelism(),
@@ -1300,7 +1238,8 @@ absl::Status GpuCompiler::OptimizeHloModule(
   TF_RETURN_IF_ERROR(layout_normalization_pipeline.Run(hlo_module).status());
   // Run target-specific HLO optimization passes after layout assignment.
   TF_RETURN_IF_ERROR(OptimizeHloPostLayoutAssignment(
-      hlo_module, stream_exec, options, gpu_target_config, thread_pool.get()));
+      hlo_module, stream_exec, options, gpu_target_config,
+      thread_pool.get_mutable()));
 
   // This is a "low effort, high impact" fusion that should be run first.
   if (hlo_module->config()
@@ -1314,7 +1253,7 @@ absl::Status GpuCompiler::OptimizeHloModule(
   }
 
   TF_RETURN_IF_ERROR(RunFusionPasses(hlo_module, gpu_target_config,
-                                     thread_pool.get(),
+                                     thread_pool.get_mutable(),
                                      ShapeSizeBytesFunction()));
   TF_RETURN_IF_ERROR(RunPostFusionPasses(
       hlo_module,
@@ -1831,7 +1770,7 @@ GpuCompiler::CompileToTargetBinary(const HloModuleConfig& module_config,
   MaybeOwningThreadPool thread_pool =
       module_config.debug_options()
               .xla_gpu_enable_llvm_module_compilation_parallelism()
-          ? MaybeOwningThreadPool::GetOrCreate(
+          ? CreateMaybeOwningThreadPool(
                 /*parallelism=*/module_config.debug_options()
                     .xla_gpu_force_compilation_parallelism(),
                 /*default_thread_pool=*/options.thread_pool,
@@ -1845,7 +1784,7 @@ GpuCompiler::CompileToTargetBinary(const HloModuleConfig& module_config,
   // Disable multi-threading during deviceless AOT compilation.
   // TODO(anlunx): Enable multi-threading once deviceless AOT compilation is
   // enabled.
-  if (!can_use_link_modules || !thread_pool || !stream_exec) {
+  if (!can_use_link_modules || !thread_pool.get() || !stream_exec) {
     return CompileSingleModule(module_config, gpu_version, debug_module,
 
                                llvm_module, /*relocatable=*/false, options,
@@ -1904,19 +1843,19 @@ GpuCompiler::CompileToTargetBinary(const HloModuleConfig& module_config,
       llvm_modules.size());
   tsl::BlockingCounter counter(llvm_modules.size());
   for (int i = 0; i < llvm_modules.size(); i++) {
-    thread_pool->Schedule([&compile_results, i, &llvm_modules, &counter, this,
-                           &module_config, &gpu_version, &debug_module,
-                           &options] {
-      // Each thread has its own context to avoid race conditions.
-      llvm::LLVMContext new_context;
-      std::unique_ptr<llvm::Module> new_module =
-          CopyToContext(*llvm_modules.at(i), new_context);
-      compile_results.at(i) = CompileSingleModule(
-          module_config, gpu_version, debug_module, new_module.get(),
-          /*relocatable=*/true, options,
-          /*shard_number=*/i);
-      counter.DecrementCount();
-    });
+    thread_pool.get_mutable()->Schedule(
+        [&compile_results, i, &llvm_modules, &counter, this, &module_config,
+         &gpu_version, &debug_module, &options] {
+          // Each thread has its own context to avoid race conditions.
+          llvm::LLVMContext new_context;
+          std::unique_ptr<llvm::Module> new_module =
+              CopyToContext(*llvm_modules.at(i), new_context);
+          compile_results.at(i) = CompileSingleModule(
+              module_config, gpu_version, debug_module, new_module.get(),
+              /*relocatable=*/true, options,
+              /*shard_number=*/i);
+          counter.DecrementCount();
+        });
   }
   counter.Wait();
 
diff --git a/third_party/xla/xla/util_test.cc b/third_party/xla/xla/util_test.cc
index 052d8cd2fe7eae..8e442edd68254c 100644
--- a/third_party/xla/xla/util_test.cc
+++ b/third_party/xla/xla/util_test.cc
@@ -19,6 +19,7 @@ limitations under the License.
 #include <cstdint>
 #include <limits>
 #include <list>
+#include <memory>
 #include <numeric>
 #include <set>
 #include <string>
@@ -26,6 +27,7 @@ limitations under the License.
 #include <utility>
 #include <vector>
 
+#include "xla/maybe_owning.h"
 #include "xla/test.h"
 #include "xla/types.h"
 #include "tsl/platform/logging.h"
@@ -318,5 +320,28 @@ TEST(UtilTest, PackInt4) {
   }
 }
 
+TEST(UtilTest, MaybeOwningTestNull) {
+  MaybeOwning<char> m(nullptr);
+  EXPECT_EQ(m.get(), nullptr);
+  EXPECT_EQ(m.get_mutable(), nullptr);
+}
+
+TEST(UtilTest, MaybeOwningTestOwning) {
+  MaybeOwning<char> m(std::make_unique<char>());
+  *m.get_mutable() = 'a';
+  EXPECT_EQ(*m, 'a');
+}
+
+TEST(UtilTest, MaybeOwningTestShared) {
+  auto owner = std::make_unique<char>();
+  *owner = 'x';
+  MaybeOwning<char> c1(owner.get());
+  MaybeOwning<char> c2(owner.get());
+
+  EXPECT_EQ(*c1, 'x');
+  EXPECT_EQ(*c2, 'x');
+  EXPECT_EQ(c1.get(), c2.get());
+}
+
 }  // namespace
 }  // namespace xla

From 06490c12157d7fa5ada4e529c83a3b3333be7602 Mon Sep 17 00:00:00 2001
From: Kyle Lucke <klucke@google.com>
Date: Mon, 13 May 2024 10:49:31 -0700
Subject: [PATCH 034/478] Remove unused constant from stream_executor_pimpl.h.

PiperOrigin-RevId: 633263139
---
 third_party/xla/xla/stream_executor/stream_executor_pimpl.h | 4 ----
 1 file changed, 4 deletions(-)

diff --git a/third_party/xla/xla/stream_executor/stream_executor_pimpl.h b/third_party/xla/xla/stream_executor/stream_executor_pimpl.h
index a75c23f6348b0a..f57ace700cc9a0 100644
--- a/third_party/xla/xla/stream_executor/stream_executor_pimpl.h
+++ b/third_party/xla/xla/stream_executor/stream_executor_pimpl.h
@@ -87,10 +87,6 @@ class StreamExecutor : public StreamExecutorInterface {
   mutable std::unique_ptr<DeviceDescription> device_description_
       ABSL_GUARDED_BY(mu_);
 
-  // Only one worker thread is needed; little work will be done by the
-  // executor.
-  static constexpr int kNumBackgroundThreads = 1;
-
   // Memory limit in bytes. Value less or equal to 0 indicates there is no
   // limit.
   int64_t memory_limit_bytes_;

From baa2cd805d9bb3a748a35f4a02eedd82d1d55056 Mon Sep 17 00:00:00 2001
From: David Dunleavy <ddunleavy@google.com>
Date: Mon, 13 May 2024 10:51:53 -0700
Subject: [PATCH 035/478] Teach `xla_test` about the interpreter backend, fail
 on unknown backends

PiperOrigin-RevId: 633263903
---
 third_party/xla/xla/tests/build_defs.bzl | 12 +++++++++---
 1 file changed, 9 insertions(+), 3 deletions(-)

diff --git a/third_party/xla/xla/tests/build_defs.bzl b/third_party/xla/xla/tests/build_defs.bzl
index ae74353e783306..ef4570ed47c18e 100644
--- a/third_party/xla/xla/tests/build_defs.bzl
+++ b/third_party/xla/xla/tests/build_defs.bzl
@@ -29,7 +29,7 @@ GPU_DEFAULT_BACKENDS = [
 
 _DEFAULT_BACKENDS = ["cpu"] + GPU_DEFAULT_BACKENDS
 
-_ALL_BACKENDS = ["cpu"] + GPU_BACKENDS + list(plugins.keys())
+_ALL_BACKENDS = ["cpu", "interpreter"] + GPU_BACKENDS + list(plugins.keys())
 
 # buildifier: disable=function-docstring
 def prepare_gpu_backend_data(backends, disabled_backends, backend_tags, backend_args):
@@ -199,6 +199,11 @@ def xla_test(
             ])
             this_backend_tags += tf_gpu_tests_tags()
             this_backend_copts.append("-DXLA_TEST_BACKEND_GPU=1")
+        elif backend == "interpreter":
+            backend_deps += [
+                "//xla/service:interpreter_plugin",
+                "//xla/tests:test_macros_interpreter",
+            ]
         elif backend in plugins:
             backend_deps += plugins[backend]["deps"]
             this_backend_copts += plugins[backend]["copts"]
@@ -206,8 +211,7 @@ def xla_test(
             this_backend_args += plugins[backend]["args"]
             this_backend_data += plugins[backend]["data"]
         else:
-            # Ignore unknown backends. TODO(b/289028518): Change back to fail.
-            continue
+            fail("Unknown backend %s" % backend)
 
         if xla_test_library_deps:
             for lib_dep in xla_test_library_deps:
@@ -288,6 +292,8 @@ def xla_test_library(
             backend_deps = ["//xla/tests:test_macros_cpu"]
         elif backend in GPU_BACKENDS:
             backend_deps = ["//xla/tests:test_macros_%s" % backend]
+        elif backend == "interpreter":
+            backend_deps = ["//xla/tests:test_macros_interpreter"]
         elif backend in plugins:
             backend_deps = plugins[backend]["deps"]
             this_backend_copts += plugins[backend]["copts"]

From d90ec6a2b507ea3fd12e5b35c3484ea9ae0d826a Mon Sep 17 00:00:00 2001
From: George Karpenkov <cheshire@google.com>
Date: Mon, 13 May 2024 10:55:23 -0700
Subject: [PATCH 036/478] [XLA] [NFC] Remove unused var

PiperOrigin-RevId: 633265210
---
 third_party/xla/xla/service/gpu/gpu_compiler.cc | 4 ----
 1 file changed, 4 deletions(-)

diff --git a/third_party/xla/xla/service/gpu/gpu_compiler.cc b/third_party/xla/xla/service/gpu/gpu_compiler.cc
index 8fc3b97cee7654..297eb7b98fdb71 100644
--- a/third_party/xla/xla/service/gpu/gpu_compiler.cc
+++ b/third_party/xla/xla/service/gpu/gpu_compiler.cc
@@ -421,10 +421,6 @@ GpuThunkAotCompilationResult::LoadExecutable(
   TF_ASSIGN_OR_RETURN(auto output_info,
                       GetOutputInfo(*hlo_module, *buffer_assignment));
   const Shape& output_shape = hlo_module->result_shape();
-  std::function<std::string()> buffer_assignment_dumper = [] {
-    return std::string();
-  };
-
   int64_t debug_buffer_assignment_show_max =
       hlo_module->config()
           .debug_options()

From 3a61f9eee0024ee36d2a4be6026e339bde837011 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Mon, 13 May 2024 11:13:32 -0700
Subject: [PATCH 037/478] Add an option (set to false by default, for now) and
 support for adding strategies for dot operators that trigger windowed einsum.

PiperOrigin-RevId: 633271761
---
 .../auto_sharding_dot_handler.cc              | 180 +++++++++++++++++-
 .../auto_sharding/auto_sharding_option.cc     |   3 +
 .../auto_sharding/auto_sharding_option.h      |   6 +
 3 files changed, 188 insertions(+), 1 deletion(-)

diff --git a/third_party/xla/xla/hlo/experimental/auto_sharding/auto_sharding_dot_handler.cc b/third_party/xla/xla/hlo/experimental/auto_sharding/auto_sharding_dot_handler.cc
index 84c22df6fa0746..8a8bf39312fd5a 100644
--- a/third_party/xla/xla/hlo/experimental/auto_sharding/auto_sharding_dot_handler.cc
+++ b/third_party/xla/xla/hlo/experimental/auto_sharding/auto_sharding_dot_handler.cc
@@ -23,8 +23,11 @@ limitations under the License.
 #include <vector>
 
 #include "absl/algorithm/container.h"
+#include "absl/container/flat_hash_set.h"
 #include "absl/log/check.h"
 #include "absl/log/log.h"
+#include "absl/strings/match.h"
+#include "absl/strings/str_cat.h"
 #include "absl/strings/str_format.h"
 #include "absl/strings/str_join.h"
 #include "absl/types/span.h"
@@ -87,6 +90,8 @@ class HandlerBase {
         lhs_(ins->operand(0)),
         rhs_(ins->operand(1)) {}
 
+  virtual ~HandlerBase() = default;
+
   void AppendNewStrategy(const std::string& name,
                          const HloSharding& output_spec,
                          absl::Span<const HloSharding> input_specs,
@@ -113,6 +118,34 @@ class HandlerBase {
       const std::optional<std::function<double(const HloSharding&)>>&
           communication_cost_fn = std::nullopt);
 
+  // Given lhs and rhs dim maps, infers a sharding for the output by relying on
+  // the sharding_propagation pass.
+  void MaybeAppendInternal(
+      const std::string& name, const DimMap& lhs_dim_map,
+      const DimMap& rhs_dim_map,
+      const std::optional<DimMap>& expected_output_dim_map,
+      const Array<int64_t>& device_mesh, double compute_cost = 0,
+      const std::optional<std::function<double(const HloSharding&)>>&
+          communication_cost_fn = std::nullopt);
+
+  // Given an existing (non-allreduce) sharding candidate, generate a
+  // corresponding candidate by additionally sharding (if possible) the passed
+  // in operand, such that, the generated candidate can trigger all-gather
+  // windowed einsum during partitioning.
+  virtual void AppendAllGatherWindowedEinsumStrategyForOperand(
+      int operand_num, const std::string& name, const DimMap& lhs_dim_map,
+      const DimMap& rhs_dim_map, const DimMap& output_dim_map,
+      const Array<int64_t>& device_mesh, double compute_cost) {}
+
+  // Given an existing (allreduce) sharding candidate, generate a corresponding
+  // candidate by additionally sharding (if possible) the dot/conv output, such
+  // that, the generated candidate can trigger reduce-scatter windowed einsum
+  // during partitioning.
+  virtual void AppendReduceScatterWindowedEinsumStrategy(
+      const std::string& name, const DimMap& lhs_dim_map,
+      const DimMap& rhs_dim_map, const DimMap& output_dim_map,
+      const Array<int64_t>& device_mesh, double compute_cost) {}
+
   std::optional<HloSharding> GetShardingFromUser(const HloSharding& lhs_spec,
                                                  const HloSharding& rhs_spec);
 
@@ -177,6 +210,8 @@ class DotHandler : public HandlerBase {
       const InstructionBatchDimMap& batch_map, const AutoShardingOption& option,
       const CallGraph& call_graph);
 
+  ~DotHandler() override = default;
+
   void SplitLhsSpaceRhsSpace();
 
   void SplitLhsSpaceOnly();
@@ -205,6 +240,16 @@ class DotHandler : public HandlerBase {
 
   void Add1DBatchSplit();
 
+  void AppendAllGatherWindowedEinsumStrategyForOperand(
+      int operand_num, const std::string& name, const DimMap& lhs_dim_map,
+      const DimMap& rhs_dim_map, const DimMap& output_dim_map,
+      const Array<int64_t>& device_mesh, double compute_cost) override;
+
+  void AppendReduceScatterWindowedEinsumStrategy(
+      const std::string& name, const DimMap& lhs_dim_map,
+      const DimMap& rhs_dim_map, const DimMap& output_dim_map,
+      const Array<int64_t>& device_mesh, double compute_cost) override;
+
   Status RegisterStrategies();
 
   // Dimension information
@@ -228,6 +273,8 @@ class ConvHandler : public HandlerBase {
               const InstructionBatchDimMap& batch_map,
               const AutoShardingOption& option, const CallGraph& call_graph);
 
+  ~ConvHandler() override = default;
+
   void SplitLhsBatchRhsOutchannel();
 
   void SplitLhsBatchBothInchannel();
@@ -287,7 +334,7 @@ void HandlerBase::AppendNewStrategy(const std::string& name,
 // TODO(b/309638633) As we build more confidence in this, we should remove
 // this expected_output_dim_map argument and fully rely on sharding
 // propagation.
-void HandlerBase::MaybeAppend(
+void HandlerBase::MaybeAppendInternal(
     const std::string& name, const DimMap& lhs_dim_map,
     const DimMap& rhs_dim_map,
     const std::optional<DimMap>& expected_output_dim_map,
@@ -336,6 +383,35 @@ void HandlerBase::MaybeAppend(
                     communication_cost);
 }
 
+void HandlerBase::MaybeAppend(
+    const std::string& name, const DimMap& lhs_dim_map,
+    const DimMap& rhs_dim_map,
+    const std::optional<DimMap>& expected_output_dim_map,
+    const Array<int64_t>& device_mesh, double compute_cost,
+    const std::optional<std::function<double(const HloSharding&)>>&
+        communication_cost_fn) {
+  MaybeAppendInternal(name, lhs_dim_map, rhs_dim_map, expected_output_dim_map,
+                      device_mesh, compute_cost, communication_cost_fn);
+  if (!option_.generate_windowed_einsum_strategies ||
+      !expected_output_dim_map.has_value()) {
+    return;
+  }
+  if (absl::StrContains(name, "allreduce")) {
+    CHECK(communication_cost_fn.has_value());
+    AppendReduceScatterWindowedEinsumStrategy(name, lhs_dim_map, rhs_dim_map,
+                                              *expected_output_dim_map,
+                                              device_mesh, compute_cost);
+  } else {
+    CHECK(!communication_cost_fn.has_value());
+    AppendAllGatherWindowedEinsumStrategyForOperand(
+        0, name, lhs_dim_map, rhs_dim_map, *expected_output_dim_map,
+        device_mesh, compute_cost);
+    AppendAllGatherWindowedEinsumStrategyForOperand(
+        1, name, lhs_dim_map, rhs_dim_map, *expected_output_dim_map,
+        device_mesh, compute_cost);
+  }
+}
+
 std::optional<HloSharding> HandlerBase::GetShardingFromUser(
     const HloSharding& lhs_spec, const HloSharding& rhs_spec) {
   std::unique_ptr<HloInstruction> ins_clone = ins_->Clone();
@@ -771,6 +847,108 @@ void DotHandler::Add1DBatchSplit() {
   }
 }
 
+void DotHandler::AppendAllGatherWindowedEinsumStrategyForOperand(
+    int operand_num, const std::string& name, const DimMap& lhs_dim_map,
+    const DimMap& rhs_dim_map, const DimMap& output_dim_map,
+    const Array<int64_t>& device_mesh, double compute_cost) {
+  const HloInstruction* operand = ins_->operand(operand_num);
+  const DimMap& operand_dim_map = operand_num == 0 ? lhs_dim_map : rhs_dim_map;
+  absl::flat_hash_set<int64_t> sharded_tensor_dims;
+  absl::flat_hash_set<int64_t> used_mesh_dims;
+  for (const auto [tensor_dim, mesh_dim] : operand_dim_map) {
+    if (device_mesh.dim(mesh_dim) == 1) {
+      continue;
+    }
+    sharded_tensor_dims.insert(tensor_dim);
+    used_mesh_dims.insert(mesh_dim);
+  }
+  if (used_mesh_dims.size() == device_mesh_.num_dimensions() ||
+      sharded_tensor_dims.size() == operand->shape().rank()) {
+    return;
+  }
+
+  for (int64_t tensor_dim = 0; tensor_dim < operand->shape().rank();
+       ++tensor_dim) {
+    if (sharded_tensor_dims.contains(tensor_dim)) {
+      continue;
+    }
+    for (int64_t mesh_dim = 0; mesh_dim < device_mesh_.num_dimensions();
+         ++mesh_dim) {
+      if (used_mesh_dims.contains(mesh_dim) ||
+          (device_mesh.dim(mesh_dim) == 1)) {
+        continue;
+      }
+      DimMap further_sharded_dim_map = operand_dim_map;
+      further_sharded_dim_map[tensor_dim] = mesh_dim;
+
+      auto updated_communication_cost_fn =
+          [](const HloSharding& output_sharding) -> double {
+        // TODO(331684721): Model costs for windowed einsum
+        return 100.0;
+      };
+
+      std::string updated_name =
+          absl::StrCat(absl::StrFormat("WindowedEinsum @ {%d,%d,%d}",
+                                       operand_num, tensor_dim, mesh_dim),
+                       name);
+      MaybeAppendInternal(
+          updated_name,
+          operand_num == 0 ? further_sharded_dim_map : lhs_dim_map,
+          operand_num == 1 ? further_sharded_dim_map : rhs_dim_map,
+          output_dim_map, device_mesh, compute_cost,
+          updated_communication_cost_fn);
+    }
+  }
+}
+
+void DotHandler::AppendReduceScatterWindowedEinsumStrategy(
+    const std::string& name, const DimMap& lhs_dim_map,
+    const DimMap& rhs_dim_map, const DimMap& output_dim_map,
+    const Array<int64_t>& device_mesh, double compute_cost) {
+  absl::flat_hash_set<int64_t> sharded_tensor_dims;
+  absl::flat_hash_set<int64_t> used_mesh_dims;
+  for (const auto [tensor_dim, mesh_dim] : output_dim_map) {
+    if (device_mesh.dim(mesh_dim) == 1) {
+      continue;
+    }
+    sharded_tensor_dims.insert(tensor_dim);
+    used_mesh_dims.insert(mesh_dim);
+  }
+  if (used_mesh_dims.size() == device_mesh_.num_dimensions() ||
+      sharded_tensor_dims.size() == ins_->shape().rank()) {
+    return;
+  }
+
+  for (int64_t tensor_dim = 0; tensor_dim < ins_->shape().rank();
+       ++tensor_dim) {
+    if (sharded_tensor_dims.contains(tensor_dim)) {
+      continue;
+    }
+    for (int64_t mesh_dim = 0; mesh_dim < device_mesh_.num_dimensions();
+         ++mesh_dim) {
+      if (used_mesh_dims.contains(mesh_dim) ||
+          (device_mesh.dim(mesh_dim) == 1)) {
+        continue;
+      }
+      DimMap further_sharded_dim_map = output_dim_map;
+      further_sharded_dim_map[tensor_dim] = mesh_dim;
+
+      auto updated_communication_cost_fn =
+          [](const HloSharding& output_sharding) -> double {
+        // TODO(331684721): Model costs for windowed einsum
+        return 100.0;
+      };
+
+      std::string updated_name = absl::StrCat(
+          absl::StrFormat("WindowedEinsum @ {%d,%d}", tensor_dim, mesh_dim),
+          name);
+      MaybeAppendInternal(updated_name, lhs_dim_map, rhs_dim_map,
+                          further_sharded_dim_map, device_mesh, compute_cost,
+                          updated_communication_cost_fn);
+    }
+  }
+}
+
 Status DotHandler::RegisterStrategies() {
   // SS = SR x RS
   // Split lhs space dim and rhs space dim.
diff --git a/third_party/xla/xla/hlo/experimental/auto_sharding/auto_sharding_option.cc b/third_party/xla/xla/hlo/experimental/auto_sharding/auto_sharding_option.cc
index e9648436047dd7..857110d78483d9 100644
--- a/third_party/xla/xla/hlo/experimental/auto_sharding/auto_sharding_option.cc
+++ b/third_party/xla/xla/hlo/experimental/auto_sharding/auto_sharding_option.cc
@@ -136,6 +136,9 @@ std::string AutoShardingOption::ToString() const {
   lines.push_back(absl::StrCat("model_resharding_memory_costs: ",
                                model_resharding_memory_costs));
 
+  lines.push_back(absl::StrCat("generate_windowed_einsum_strategies: ",
+                               generate_windowed_einsum_strategies));
+
   return absl::StrJoin(lines, "\n");
 }
 
diff --git a/third_party/xla/xla/hlo/experimental/auto_sharding/auto_sharding_option.h b/third_party/xla/xla/hlo/experimental/auto_sharding/auto_sharding_option.h
index a858ecfce1d4da..2983d5ac9d7999 100644
--- a/third_party/xla/xla/hlo/experimental/auto_sharding/auto_sharding_option.h
+++ b/third_party/xla/xla/hlo/experimental/auto_sharding/auto_sharding_option.h
@@ -191,6 +191,12 @@ struct AutoShardingOption {
   // for resharding edges.
   bool model_resharding_memory_costs = true;
 
+  // Whether or not to generate strategies that model the windowed einsum (or
+  // collective matmul) optimization
+  // TODO(331684721,329508561): Generate windowed-einsum strategies by default
+  // once it is fully implemented.
+  bool generate_windowed_einsum_strategies = false;
+
   // Prints a debug string.
   std::string ToString() const;
 

From 3cefc36d7cb4b065a8ebe87c11c9577b9f49bdd5 Mon Sep 17 00:00:00 2001
From: Junwhan Ahn <junwhan@google.com>
Date: Mon, 13 May 2024 11:17:14 -0700
Subject: [PATCH 038/478] Implement basic memory space support for PjRt CPU

This CL implements minimal memory support for PjRt CPU so that memory space APIs can be used consistently across device types. Like GPU, each CPU device for now has just one memory space kind (`device`) and all memory space flavors of transfer APIs redirect the calls to the corresponding versions that take a device instead, assuming that there's only one device to which each memory space is attached.

PiperOrigin-RevId: 633272982
---
 third_party/xla/xla/pjrt/cpu/BUILD            |   3 +
 third_party/xla/xla/pjrt/cpu/cpu_client.cc    | 131 ++++++++++++++++--
 third_party/xla/xla/pjrt/cpu/cpu_client.h     |  45 +++++-
 .../xla/xla/pjrt/cpu/cpu_client_test.cc       |  20 +++
 4 files changed, 180 insertions(+), 19 deletions(-)

diff --git a/third_party/xla/xla/pjrt/cpu/BUILD b/third_party/xla/xla/pjrt/cpu/BUILD
index fdc4acf62abae4..a06396197fd617 100644
--- a/third_party/xla/xla/pjrt/cpu/BUILD
+++ b/third_party/xla/xla/pjrt/cpu/BUILD
@@ -154,6 +154,7 @@ cc_library(
         "//xla/client:xla_computation",
         "//xla/hlo/ir:hlo",
         "//xla/pjrt:compile_options_proto_cc",
+        "//xla/pjrt:host_memory_spaces",
         "//xla/pjrt:mlir_to_hlo",
         "//xla/pjrt:pjrt_client",
         "//xla/pjrt:pjrt_common",
@@ -186,6 +187,7 @@ cc_library(
         "//xla/service/cpu:simple_orc_jit",
         "//xla/tsl/concurrency:async_value",
         "//xla/tsl/concurrency:ref_count",
+        "@com_google_absl//absl/algorithm:container",
         "@com_google_absl//absl/base:core_headers",
         "@com_google_absl//absl/base:dynamic_annotations",
         "@com_google_absl//absl/container:flat_hash_map",
@@ -229,6 +231,7 @@ xla_cc_test(
         "//xla:util",
         "//xla/ffi",
         "//xla/ffi:ffi_api",
+        "//xla/pjrt:host_memory_spaces",
         "//xla/service:hlo_parser",
         "//xla/tests:test_utils",
         "@com_google_absl//absl/status",
diff --git a/third_party/xla/xla/pjrt/cpu/cpu_client.cc b/third_party/xla/xla/pjrt/cpu/cpu_client.cc
index b0b63a257c7fdd..3746dd4fd1c3d5 100644
--- a/third_party/xla/xla/pjrt/cpu/cpu_client.cc
+++ b/third_party/xla/xla/pjrt/cpu/cpu_client.cc
@@ -27,7 +27,9 @@ limitations under the License.
 #include <utility>
 #include <vector>
 
+#include "absl/algorithm/container.h"
 #include "xla/pjrt/cpu/cpu_topology.h"
+#include "xla/pjrt/host_memory_spaces.h"
 #include "xla/pjrt/pjrt_compiler.h"
 
 #define EIGEN_USE_THREADS
@@ -123,7 +125,8 @@ absl::StatusOr<std::unique_ptr<TfrtCpuBuffer>> AllocateDestinationBuffer(
       AbstractTfrtCpuBuffer::AllocateTrackedDeviceBuffer(
           on_device_shape, std::move(definition_events)));
   return std::make_unique<TfrtCpuBuffer>(
-      on_device_shape, std::move(tracked_device_buffer), client, device);
+      on_device_shape, std::move(tracked_device_buffer), client, device,
+      *device->default_memory_space());
 }
 
 absl::StatusOr<std::unique_ptr<TfrtCpuBuffer>> AllocateDestinationBufferAndAvs(
@@ -309,12 +312,47 @@ Status TfrtCpuDevice::TransferFromOutfeed(MutableBorrowingLiteral literal) {
   return TransferLiteralFromOutfeedOnCpu(local_hardware_id(), literal);
 }
 
+void TfrtCpuDevice::AttachMemorySpace(PjRtMemorySpace* memory_space) {
+  CHECK(memory_space != nullptr);
+  CHECK(client_ == memory_space->client()) << absl::StrFormat(
+      "Could not attach a TfrtCpuDevice to a PjRtMemorySpace owned by a "
+      "different client, the device's client: %s, the memory space's client: "
+      "%s.",
+      client_->platform_name(), memory_space->client()->platform_name());
+
+  memory_spaces_.push_back(memory_space);
+  memory_spaces_by_id_.emplace(memory_space->kind_id(), memory_space);
+}
+
 absl::Span<PjRtMemorySpace* const> TfrtCpuDevice::memory_spaces() const {
-  return {};
+  return memory_spaces_;
 }
 
 absl::StatusOr<PjRtMemorySpace*> TfrtCpuDevice::default_memory_space() const {
-  return Unimplemented("default_memory_space is not supported");
+  return memory_space_by_kind_id(UnpinnedHostMemorySpace::kKindId);
+}
+
+absl::StatusOr<PjRtMemorySpace*> TfrtCpuDevice::memory_space_by_kind(
+    absl::string_view memory_space_kind) const {
+  auto it =
+      absl::c_find_if(memory_spaces_, [memory_space_kind](PjRtMemorySpace* ms) {
+        return ms->kind() == memory_space_kind;
+      });
+  if (it != memory_spaces_.end()) {
+    return *it;
+  }
+  return absl::InternalError(
+      absl::StrCat("No memory space found (kind: ", memory_space_kind, ")"));
+}
+
+absl::StatusOr<PjRtMemorySpace*> TfrtCpuDevice::memory_space_by_kind_id(
+    int id) const {
+  auto it = memory_spaces_by_id_.find(id);
+  if (it == memory_spaces_by_id_.end()) {
+    return absl::InternalError(
+        absl::StrCat("No memory space found (kind_id: ", id, ")"));
+  }
+  return it->second;
 }
 
 static int CpuDeviceCount() {
@@ -417,8 +455,19 @@ TfrtCpuClient::TfrtCpuClient(
     }
   }
   for (int idx = 0; idx < addressable_devices_.size(); ++idx) {
-    CHECK(addressable_devices_[idx] != nullptr) << idx;
+    auto* const device = addressable_devices_[idx];
+    CHECK(device != nullptr) << idx;
+
+    // Use the device id to construct a globally unique memory space id. We
+    // do not promise that memory space ids and device ids are the same.
+    const int id = device->id();
+    auto memory_space = std::make_unique<UnpinnedHostMemorySpace>(id, device);
+    tensorflow::down_cast<TfrtCpuDevice*>(device)->AttachMemorySpace(
+        memory_space.get());
+    memory_spaces_.push_back(memory_space.get());
+    owned_memory_spaces_.push_back(std::move(memory_space));
   }
+
   LOG(INFO) << "TfrtCpuClient created.";
 }
 
@@ -451,7 +500,7 @@ absl::StatusOr<PjRtDevice*> TfrtCpuClient::LookupAddressableDevice(
 }
 
 absl::Span<PjRtMemorySpace* const> TfrtCpuClient::memory_spaces() const {
-  return {};
+  return memory_spaces_;
 }
 
 absl::StatusOr<DeviceAssignment> TfrtCpuClient::GetDefaultDeviceAssignment(
@@ -844,7 +893,8 @@ TfrtCpuClient::CreateViewOfDeviceBuffer(
       std::move(on_delete_callback));
   return std::unique_ptr<PjRtBuffer>(std::make_unique<TfrtCpuBuffer>(
       shape, std::move(tracked_device_buffer), this,
-      tensorflow::down_cast<TfrtCpuDevice*>(device)));
+      tensorflow::down_cast<TfrtCpuDevice*>(device),
+      *device->default_memory_space()));
 }
 
 absl::StatusOr<std::unique_ptr<PjRtBuffer>> TfrtCpuClient::CreateErrorBuffer(
@@ -860,7 +910,8 @@ absl::StatusOr<std::unique_ptr<PjRtBuffer>> TfrtCpuClient::CreateErrorBuffer(
           absl::InlinedVector<tsl::AsyncValueRef<CpuEvent>, 4>{
               tsl::AsyncValueRef<CpuEvent>(
                   tsl::MakeErrorAsyncValueRef(std::move(error)))}),
-      this, tensorflow::down_cast<TfrtCpuDevice*>(device));
+      this, tensorflow::down_cast<TfrtCpuDevice*>(device),
+      *device->default_memory_space());
 }
 
 absl::StatusOr<std::unique_ptr<PjRtBuffer>> TfrtCpuClient::CreateErrorBuffer(
@@ -887,6 +938,13 @@ TfrtCpuClient::CreateBuffersForAsyncHostToDevice(absl::Span<const Shape> shapes,
                                                          this);
 }
 
+absl::StatusOr<std::unique_ptr<PjRtClient::AsyncHostToDeviceTransferManager>>
+TfrtCpuClient::CreateBuffersForAsyncHostToDevice(
+    absl::Span<const Shape> shapes, PjRtMemorySpace* memory_space) {
+  CHECK_EQ(memory_space->devices().size(), 1);
+  return CreateBuffersForAsyncHostToDevice(shapes, memory_space->devices()[0]);
+}
+
 absl::StatusOr<std::unique_ptr<PjRtBuffer>> TfrtCpuClient::BufferFromHostBuffer(
     const void* data, PrimitiveType type, absl::Span<int64_t const> dims,
     std::optional<absl::Span<int64_t const>> byte_strides,
@@ -911,7 +969,38 @@ absl::StatusOr<std::unique_ptr<PjRtBuffer>> TfrtCpuClient::BufferFromHostBuffer(
 
   return std::unique_ptr<PjRtBuffer>(std::make_unique<TfrtCpuBuffer>(
       shape, std::move(tracked_device_buffer), this,
-      tensorflow::down_cast<TfrtCpuDevice*>(device)));
+      tensorflow::down_cast<TfrtCpuDevice*>(device),
+      *device->default_memory_space()));
+}
+
+absl::StatusOr<std::unique_ptr<PjRtBuffer>> TfrtCpuClient::BufferFromHostBuffer(
+    const void* data, PrimitiveType type, absl::Span<int64_t const> dims,
+    std::optional<absl::Span<int64_t const>> byte_strides,
+    HostBufferSemantics host_buffer_semantics,
+    absl::AnyInvocable<void() &&> on_done_with_host_buffer, PjRtDevice* device,
+    const Layout* device_layout) {
+  if (device_layout != nullptr) {
+    return absl::UnimplementedError(absl::StrCat(
+        "BufferFromHostBuffer with an optional device layout is not "
+        "implemented on platform: ",
+        platform_name()));
+  }
+  return BufferFromHostBuffer(data, type, dims, byte_strides,
+                              host_buffer_semantics,
+                              std::move(on_done_with_host_buffer), device);
+}
+
+absl::StatusOr<std::unique_ptr<PjRtBuffer>> TfrtCpuClient::BufferFromHostBuffer(
+    const void* data, PrimitiveType type, absl::Span<int64_t const> dims,
+    std::optional<absl::Span<int64_t const>> byte_strides,
+    HostBufferSemantics host_buffer_semantics,
+    absl::AnyInvocable<void() &&> on_done_with_host_buffer,
+    PjRtMemorySpace* memory_space, const Layout* device_layout) {
+  CHECK_EQ(memory_space->devices().size(), 1);
+  return BufferFromHostBuffer(data, type, dims, byte_strides,
+                              host_buffer_semantics,
+                              std::move(on_done_with_host_buffer),
+                              memory_space->devices()[0], device_layout);
 }
 
 absl::StatusOr<std::unique_ptr<PjRtBuffer>>
@@ -934,14 +1023,22 @@ TfrtCpuClient::BufferFromHostLiteral(const LiteralSlice& literal,
   return std::unique_ptr<PjRtBuffer>(std::move(output_buffer));
 }
 
+absl::StatusOr<std::unique_ptr<PjRtBuffer>>
+TfrtCpuClient::BufferFromHostLiteral(const LiteralSlice& literal,
+                                     PjRtMemorySpace* memory_space) {
+  CHECK_EQ(memory_space->devices().size(), 1);
+  return BufferFromHostLiteral(literal, memory_space->devices()[0]);
+}
+
 TfrtCpuBuffer::TfrtCpuBuffer(
     Shape on_device_shape,
     std::unique_ptr<TrackedTfrtCpuDeviceBuffer> tracked_device_buffer,
-    TfrtCpuClient* client, TfrtCpuDevice* device)
+    TfrtCpuClient* client, TfrtCpuDevice* device, PjRtMemorySpace* memory_space)
     : AbstractTfrtCpuBuffer(std::move(on_device_shape),
                             std::move(tracked_device_buffer)),
       client_(client),
-      device_(device) {}
+      device_(device),
+      memory_space_(memory_space) {}
 
 static std::vector<tsl::RCReference<tsl::AsyncValue>> CopyAsyncValues(
     absl::Span<const tsl::RCReference<tsl::AsyncValue>> events) {
@@ -994,7 +1091,14 @@ absl::StatusOr<std::unique_ptr<PjRtBuffer>> TfrtCpuBuffer::CopyToDevice(
 
   return std::unique_ptr<PjRtBuffer>(std::make_unique<TfrtCpuBuffer>(
       on_device_shape_, std::move(tracked_device_buffer), client(),
-      tensorflow::down_cast<TfrtCpuDevice*>(dst_device)));
+      tensorflow::down_cast<TfrtCpuDevice*>(dst_device),
+      *dst_device->default_memory_space()));
+}
+
+absl::StatusOr<std::unique_ptr<PjRtBuffer>> TfrtCpuBuffer::CopyToMemorySpace(
+    PjRtMemorySpace* dst_memory_space) {
+  CHECK_EQ(dst_memory_space->devices().size(), 1);
+  return CopyToDevice(dst_memory_space->devices()[0]);
 }
 
 TfrtCpuExecutable::TfrtCpuExecutable(
@@ -1580,7 +1684,7 @@ absl::StatusOr<PjRtLoadedExecutable::Result> TfrtCpuExecutable::ExecuteHelper(
               std::move(definition_events));
       auto leaf_buffer = std::make_unique<TfrtCpuBuffer>(
           result_shape.tuple_shapes(i), std::move(leaf_tracked_device_buffer),
-          client_, device);
+          client_, device, *device->default_memory_space());
       res.push_back(std::move(leaf_buffer));
     }
   } else {
@@ -1601,7 +1705,8 @@ absl::StatusOr<PjRtLoadedExecutable::Result> TfrtCpuExecutable::ExecuteHelper(
         std::move(sub_buffers), std::move(sub_buffer_sizes),
         /*definition_event=*/execute_event);
     auto tfrt_output_buffer = std::make_unique<TfrtCpuBuffer>(
-        result_shape, std::move(tracked_device_buffer), client_, device);
+        result_shape, std::move(tracked_device_buffer), client_, device,
+        *device->default_memory_space());
     res.push_back(std::move(tfrt_output_buffer));
   }
 
diff --git a/third_party/xla/xla/pjrt/cpu/cpu_client.h b/third_party/xla/xla/pjrt/cpu/cpu_client.h
index 4dad5dd3efb83f..554d7bb9c3b0d6 100644
--- a/third_party/xla/xla/pjrt/cpu/cpu_client.h
+++ b/third_party/xla/xla/pjrt/cpu/cpu_client.h
@@ -233,10 +233,17 @@ class TfrtCpuDevice final : public PjRtDevice {
 
   Status TransferFromOutfeed(MutableBorrowingLiteral literal) override;
 
+  void AttachMemorySpace(PjRtMemorySpace* memory_space);
+
   absl::Span<PjRtMemorySpace* const> memory_spaces() const override;
 
   absl::StatusOr<PjRtMemorySpace*> default_memory_space() const override;
 
+  absl::StatusOr<PjRtMemorySpace*> memory_space_by_kind(
+      absl::string_view memory_space_kind) const override;
+
+  absl::StatusOr<PjRtMemorySpace*> memory_space_by_kind_id(int id) const;
+
   // Returns a semaphore for admission control on inflight computations.
   Semaphore& max_inflight_computations_semaphore() {
     return max_inflight_computations_semaphore_;
@@ -250,6 +257,8 @@ class TfrtCpuDevice final : public PjRtDevice {
  private:
   PjRtClient* client_ = nullptr;
   TfrtCpuDeviceDescription description_;
+  absl::InlinedVector<PjRtMemorySpace*, 1> memory_spaces_;
+  absl::flat_hash_map<int, PjRtMemorySpace*> memory_spaces_by_id_;
 
   // TODO(zhangqiaorjc): Optimize semaphore related overhead.
   // Semaphore used to limit how many programs can be enqueued by the host
@@ -335,10 +344,7 @@ class TfrtCpuClient final : public PjRtClient {
 
   absl::StatusOr<std::unique_ptr<PjRtClient::AsyncHostToDeviceTransferManager>>
   CreateBuffersForAsyncHostToDevice(absl::Span<const Shape> shapes,
-                                    PjRtMemorySpace* memory_space) override {
-    return Unimplemented(
-        "CreateBuffersForAsyncHostToDevice with memory_space not implemented.");
-  }
+                                    PjRtMemorySpace* memory_space) override;
 
   absl::StatusOr<std::unique_ptr<PjRtBuffer>> BufferFromHostBuffer(
       const void* data, PrimitiveType type, absl::Span<int64_t const> dims,
@@ -347,9 +353,26 @@ class TfrtCpuClient final : public PjRtClient {
       absl::AnyInvocable<void() &&> on_done_with_host_buffer,
       PjRtDevice* device) override;
 
+  absl::StatusOr<std::unique_ptr<PjRtBuffer>> BufferFromHostBuffer(
+      const void* data, PrimitiveType type, absl::Span<int64_t const> dims,
+      std::optional<absl::Span<int64_t const>> byte_strides,
+      HostBufferSemantics host_buffer_semantics,
+      absl::AnyInvocable<void() &&> on_done_with_host_buffer,
+      PjRtDevice* device, const Layout* device_layout) override;
+
+  absl::StatusOr<std::unique_ptr<PjRtBuffer>> BufferFromHostBuffer(
+      const void* data, PrimitiveType type, absl::Span<int64_t const> dims,
+      std::optional<absl::Span<int64_t const>> byte_strides,
+      HostBufferSemantics host_buffer_semantics,
+      absl::AnyInvocable<void() &&> on_done_with_host_buffer,
+      PjRtMemorySpace* memory_space, const Layout* device_layout) override;
+
   absl::StatusOr<std::unique_ptr<PjRtBuffer>> BufferFromHostLiteral(
       const LiteralSlice& literal, PjRtDevice* device) override;
 
+  absl::StatusOr<std::unique_ptr<PjRtBuffer>> BufferFromHostLiteral(
+      const LiteralSlice& literal, PjRtMemorySpace* memory_space) override;
+
   absl::StatusOr<std::vector<std::unique_ptr<PjRtBuffer>>>
   MakeCrossHostReceiveBuffers(absl::Span<const Shape> shapes,
                               PjRtDevice* device,
@@ -433,6 +456,11 @@ class TfrtCpuClient final : public PjRtClient {
   std::vector<PjRtDevice*> addressable_devices_;
   std::unique_ptr<ComputationPlacer> computation_placer_;
 
+  // Addressable memory spaces.
+  std::vector<std::unique_ptr<PjRtMemorySpace>> owned_memory_spaces_;
+  // Pointers to `owned_memory_spaces_`.
+  std::vector<PjRtMemorySpace*> memory_spaces_;
+
   // Thread pool for running PjRtClient tasks.
   std::unique_ptr<tsl::thread::ThreadPool> pjrt_client_thread_pool_;
   std::unique_ptr<AsyncWorkRunner> async_work_runner_;
@@ -482,14 +510,15 @@ class TfrtCpuBuffer final : public AbstractTfrtCpuBuffer {
   TfrtCpuBuffer(
       Shape on_device_shape,
       std::unique_ptr<TrackedTfrtCpuDeviceBuffer> tracked_device_buffer,
-      TfrtCpuClient* client, TfrtCpuDevice* device);
+      TfrtCpuClient* client, TfrtCpuDevice* device,
+      PjRtMemorySpace* memory_space);
 
   TfrtCpuBuffer(const TfrtCpuBuffer&) = delete;
   TfrtCpuBuffer(TfrtCpuBuffer&&) = delete;
   TfrtCpuBuffer& operator=(const TfrtCpuBuffer&) = delete;
   TfrtCpuBuffer& operator=(TfrtCpuBuffer&&) = delete;
 
-  PjRtMemorySpace* memory_space() const override { return nullptr; }
+  PjRtMemorySpace* memory_space() const override { return memory_space_; }
   TfrtCpuDevice* device() const override { return device_; }
   TfrtCpuClient* client() const override { return client_; }
 
@@ -502,11 +531,15 @@ class TfrtCpuBuffer final : public AbstractTfrtCpuBuffer {
   absl::StatusOr<std::unique_ptr<PjRtBuffer>> CopyToDevice(
       PjRtDevice* dst_device) override;
 
+  absl::StatusOr<std::unique_ptr<PjRtBuffer>> CopyToMemorySpace(
+      PjRtMemorySpace* dst_memory_space) override;
+
  private:
   absl::string_view buffer_name() const override { return "TfrtCpuBuffer"; }
 
   TfrtCpuClient* client_;
   TfrtCpuDevice* const device_;
+  PjRtMemorySpace* const memory_space_;
 };
 
 class TfrtCpuExecutable final : public PjRtLoadedExecutable {
diff --git a/third_party/xla/xla/pjrt/cpu/cpu_client_test.cc b/third_party/xla/xla/pjrt/cpu/cpu_client_test.cc
index f8b343d4313f10..1b8075b11013d5 100644
--- a/third_party/xla/xla/pjrt/cpu/cpu_client_test.cc
+++ b/third_party/xla/xla/pjrt/cpu/cpu_client_test.cc
@@ -15,6 +15,8 @@ limitations under the License.
 
 #include "xla/pjrt/cpu/cpu_client.h"
 
+#include "xla/pjrt/host_memory_spaces.h"
+
 #ifndef _WIN32
 #include <unistd.h>
 #endif
@@ -50,9 +52,11 @@ namespace xla {
 namespace {
 
 using ::testing::Each;
+using ::testing::ElementsAre;
 using ::testing::ElementsAreArray;
 using ::testing::HasSubstr;
 using ::testing::IsFalse;
+using ::tsl::testing::IsOkAndHolds;
 
 static absl::Status TestError(ffi::BufferBase, ffi::Result<ffi::BufferBase>,
                               ffi::Result<ffi::BufferBase>) {
@@ -69,6 +73,22 @@ XLA_FFI_DEFINE_HANDLER(kTestError, TestError,
 XLA_FFI_REGISTER_HANDLER(ffi::GetXlaFfiApi(), "__xla_test$$TestError", "Host",
                          kTestError);
 
+TEST(TfrtCpuClientTest, MemorySpace) {
+  TF_ASSERT_OK_AND_ASSIGN(auto client, GetTfrtCpuClient(CpuClientOptions()));
+  ASSERT_GE(client->devices().size(), 1);
+
+  ASSERT_EQ(client->memory_spaces().size(),
+            client->addressable_devices().size());
+  for (auto* device : client->devices()) {
+    TF_ASSERT_OK_AND_ASSIGN(auto* memory_space, device->default_memory_space());
+    EXPECT_THAT(device->memory_spaces(), ElementsAre(memory_space));
+    EXPECT_EQ(memory_space->kind(), UnpinnedHostMemorySpace::kKind);
+    EXPECT_EQ(memory_space->kind_id(), UnpinnedHostMemorySpace::kKindId);
+    EXPECT_THAT(device->memory_space_by_kind(UnpinnedHostMemorySpace::kKind),
+                IsOkAndHolds(memory_space));
+  }
+}
+
 TEST(TfrtCpuClientTest, DonationWithExecutionError) {
   constexpr char kProgram[] =
       R"(HloModule DonationWithExecutionError, input_output_alias={ {}: (0, {}, must-alias) }

From 7e1a1b30710826abb7ab43a9bff765fee1c2cbb1 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Mon, 13 May 2024 11:18:18 -0700
Subject: [PATCH 039/478] Prevent heartbeat cancelled errors at shutdown.

Clients disconnect before the service shuts down. However, there is a short time window where the clients send out heartbeat RPCs right before shutting down, which might reach the server after it shuts down.

PiperOrigin-RevId: 633273384
---
 .../coordination_service_agent.cc             | 35 +++++++++++--------
 1 file changed, 21 insertions(+), 14 deletions(-)

diff --git a/third_party/xla/xla/tsl/distributed_runtime/coordination/coordination_service_agent.cc b/third_party/xla/xla/tsl/distributed_runtime/coordination/coordination_service_agent.cc
index 59d1abf1b80956..36969eb0025205 100644
--- a/third_party/xla/xla/tsl/distributed_runtime/coordination/coordination_service_agent.cc
+++ b/third_party/xla/xla/tsl/distributed_runtime/coordination/coordination_service_agent.cc
@@ -45,7 +45,6 @@ limitations under the License.
 #include "tsl/platform/mutex.h"
 #include "tsl/platform/random.h"
 #include "tsl/platform/status.h"
-#include "tsl/platform/statusor.h"
 #include "tsl/platform/thread_annotations.h"
 #include "tsl/protobuf/coordination_config.pb.h"
 #include "tsl/protobuf/coordination_service.pb.h"
@@ -72,7 +71,9 @@ class CoordinationServiceAgentImpl : public CoordinationServiceAgent {
  public:
   CoordinationServiceAgentImpl() = default;
   ~CoordinationServiceAgentImpl() override {
-    absl::Status s = Shutdown();
+    // TODO(b/339231167): Fix the lint.
+    absl::Status s =
+        Shutdown();  // NOLINT(clang-analyzer-optin.cplusplus.VirtualCall)
     VLOG(3) << "Coordination agent dtor failed with status: " << s;
   }
   absl::Status Initialize(Env* env, std::string_view job_name, int task_id,
@@ -230,7 +231,7 @@ void CoordinationServiceAgentImpl::StopHeartbeat() {
     shutting_down_ = true;
     heartbeat_thread_cv_.notify_all();
   }
-  heartbeat_thread_.reset();
+  heartbeat_thread_ = nullptr;
 }
 
 absl::Status CoordinationServiceAgentImpl::Connect() {
@@ -331,17 +332,21 @@ absl::Status CoordinationServiceAgentImpl::Connect() {
                                          });
           n.WaitForNotification();
           VLOG(10) << "HeartbeatResponse: " << status;
-          {
-            mutex_lock l(heartbeat_thread_shutdown_mu_);
+          if (!status.ok()) {
             // Ignore heartbeat errors and exit thread if shutting down. For
-            // example, the agent may send a heartbeat right after Shutdown(),
-            // but before StopHeartbeat(). This results in an unexpected
-            // heartbeat error.
-            if (shutting_down_) {
-              return;
+            // example, the agent may send a heartbeat right after Shutdown()
+            // started, but before StopHeartbeat() and end of Shutdown(). This
+            // results in an unexpected heartbeat error.
+            // Waiting for a second allows us to identify if errors are due to
+            // inflight heartbeats sent during shutdown and can be ignored.
+            absl::SleepFor(absl::Seconds(1));
+            {
+              mutex_lock l(heartbeat_thread_shutdown_mu_);
+
+              if (shutting_down_) {
+                return;
+              }
             }
-          }
-          if (!status.ok()) {
             SetError(status);
           } else if (response.leader_incarnation() != leader_incarnation_) {
             SetError(MakeCoordinationError(
@@ -351,8 +356,10 @@ absl::Status CoordinationServiceAgentImpl::Connect() {
           // Send next heartbeat after an interval.
           {
             mutex_lock l(heartbeat_thread_shutdown_mu_);
+            // TODO(b/339231167): Fix the lint.
             heartbeat_thread_cv_.wait_for(
-                l, std::chrono::milliseconds(heartbeat_interval_ms));
+                l, std::chrono::milliseconds(  // NOLINT(misc-include-cleaner)
+                       heartbeat_interval_ms));
             if (shutting_down_) {
               return;
             }
@@ -583,7 +590,7 @@ absl::StatusOr<std::string> CoordinationServiceAgentImpl::GetKeyValue(
 absl::StatusOr<std::string> CoordinationServiceAgentImpl::GetKeyValue(
     std::string_view key, absl::Duration timeout) {
   auto n = std::make_shared<absl::Notification>();
-  auto result = std::make_shared<StatusOr<std::string>>();
+  auto result = std::make_shared<absl::StatusOr<std::string>>();
   GetKeyValueAsync(
       key, [n, result](const absl::StatusOr<std::string>& status_or_value) {
         *result = status_or_value;

From a53fd126c00f96c8a8b53e6853b17f9825aa5a26 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Mon, 13 May 2024 11:39:45 -0700
Subject: [PATCH 040/478] Add fields to GpuTopology for multi-host support.

PiperOrigin-RevId: 633280790
---
 third_party/xla/xla/pjrt/gpu/gpu_topology.cc  |  7 ++++-
 third_party/xla/xla/pjrt/gpu/gpu_topology.h   | 29 +++++++++++++++++--
 .../xla/xla/pjrt/gpu/se_gpu_pjrt_client.h     |  8 +++--
 .../xla/pjrt/gpu/se_gpu_pjrt_client_test.cc   | 14 ++++++++-
 4 files changed, 50 insertions(+), 8 deletions(-)

diff --git a/third_party/xla/xla/pjrt/gpu/gpu_topology.cc b/third_party/xla/xla/pjrt/gpu/gpu_topology.cc
index b7593448646d99..e9baf5f359ba66 100644
--- a/third_party/xla/xla/pjrt/gpu/gpu_topology.cc
+++ b/third_party/xla/xla/pjrt/gpu/gpu_topology.cc
@@ -25,13 +25,18 @@ std::unique_ptr<const GpuTopology> GpuTopology::FromProto(
   return std::make_unique<GpuTopology>(
       std::vector<int>{gpu_topology_proto.device_ids().begin(),
                        gpu_topology_proto.device_ids().end()},
-      gpu_topology_proto.platform_version());
+      gpu_topology_proto.platform_version(), gpu_topology_proto.num_slices(),
+      gpu_topology_proto.num_hosts_per_slice(),
+      gpu_topology_proto.num_devices_per_host());
 }
 
 GpuTopologyProto GpuTopology::ToProto() const {
   GpuTopologyProto proto;
   proto.mutable_device_ids()->Add(device_ids().begin(), device_ids().end());
   proto.set_platform_version(platform_version());
+  proto.set_num_slices(num_slices());
+  proto.set_num_hosts_per_slice(num_hosts_per_slice());
+  proto.set_num_devices_per_host(num_devices_per_host());
   return proto;
 }
 
diff --git a/third_party/xla/xla/pjrt/gpu/gpu_topology.h b/third_party/xla/xla/pjrt/gpu/gpu_topology.h
index 1412376b8a7be7..9c1e191bc4b72b 100644
--- a/third_party/xla/xla/pjrt/gpu/gpu_topology.h
+++ b/third_party/xla/xla/pjrt/gpu/gpu_topology.h
@@ -26,21 +26,44 @@ namespace xla {
 class GpuTopology {
  public:
   explicit GpuTopology(const std::vector<int>& gpu_device_ids,
-                       absl::string_view platform_version)
-      : devices_ids_(gpu_device_ids), platform_version_(platform_version) {}
+                       absl::string_view platform_version, int32_t num_slices,
+                       int32_t num_hosts_per_slice,
+                       int32_t num_devices_per_host)
+      : devices_ids_(gpu_device_ids),
+        platform_version_(platform_version),
+        num_slices_(num_slices),
+        num_hosts_per_slice_(num_hosts_per_slice),
+        num_devices_per_host_(num_devices_per_host) {}
 
-  int number_of_devices() const { return devices_ids_.size(); }
+  bool operator==(const GpuTopology& other) const {
+    return devices_ids_ == other.devices_ids_ &&
+           platform_version_ == other.platform_version_ &&
+           num_slices_ == other.num_slices_ &&
+           num_hosts_per_slice_ == other.num_hosts_per_slice_ &&
+           num_devices_per_host_ == other.num_devices_per_host_;
+  }
+
+  int number_of_devices() const {
+    return number_of_hosts() * num_devices_per_host_;
+  }
   const std::vector<int>& device_ids() const { return devices_ids_; }
+  int number_of_hosts() const { return num_slices_ * num_hosts_per_slice_; }
 
   static std::unique_ptr<const GpuTopology> FromProto(
       const GpuTopologyProto& proto);
   GpuTopologyProto ToProto() const;
 
   std::string platform_version() const { return platform_version_; }
+  int32_t num_slices() const { return num_slices_; }
+  int32_t num_hosts_per_slice() const { return num_hosts_per_slice_; }
+  int32_t num_devices_per_host() const { return num_devices_per_host_; }
 
  private:
   const std::vector<int> devices_ids_;
   const std::string platform_version_;
+  const int32_t num_slices_;
+  const int32_t num_hosts_per_slice_;
+  const int32_t num_devices_per_host_;
 };
 
 }  // namespace xla
diff --git a/third_party/xla/xla/pjrt/gpu/se_gpu_pjrt_client.h b/third_party/xla/xla/pjrt/gpu/se_gpu_pjrt_client.h
index 90717f87fbe6c5..2101501050a670 100644
--- a/third_party/xla/xla/pjrt/gpu/se_gpu_pjrt_client.h
+++ b/third_party/xla/xla/pjrt/gpu/se_gpu_pjrt_client.h
@@ -87,15 +87,17 @@ class StreamExecutorGpuTopologyDescription : public PjRtTopologyDescription {
       : platform_id_(platform_id),
         platform_name_(platform_name),
         platform_version_(platform_version),
-        gpu_topology_(gpu_device_ids, platform_version),
+        // TODO(b/331224674): Add support for multi-host.
+        gpu_topology_(gpu_device_ids, platform_version, /*num_slices=*/1,
+                      /*num_hosts_per_slice=*/1,
+                      /*num_devices_per_host=*/gpu_device_ids.size()),
         attributes_(attributes) {}
 
   bool operator==(const StreamExecutorGpuTopologyDescription& other) const {
     return this->platform_id() == other.platform_id() &&
            this->platform_name() == other.platform_name() &&
            this->platform_version() == other.platform_version() &&
-           this->gpu_topology().device_ids() ==
-               other.gpu_topology().device_ids();
+           this->gpu_topology() == other.gpu_topology();
   }
 
   PjRtPlatformId platform_id() const override { return platform_id_; }
diff --git a/third_party/xla/xla/pjrt/gpu/se_gpu_pjrt_client_test.cc b/third_party/xla/xla/pjrt/gpu/se_gpu_pjrt_client_test.cc
index 5a14973d215cc2..d3bb175754e1fb 100644
--- a/third_party/xla/xla/pjrt/gpu/se_gpu_pjrt_client_test.cc
+++ b/third_party/xla/xla/pjrt/gpu/se_gpu_pjrt_client_test.cc
@@ -640,20 +640,32 @@ TEST(GpuTopology, FromProto) {
       R"pb(
         device_ids: [ 3, 2, 1 ]
         platform_version: "platform_version"
+        num_slices: 2
+        num_hosts_per_slice: 1
+        num_devices_per_host: 3
       )pb",
       &msg));
 
   std::unique_ptr<const GpuTopology> gpu_topology = GpuTopology::FromProto(msg);
   EXPECT_THAT(gpu_topology->device_ids(), ElementsAre(3, 2, 1));
   EXPECT_THAT(gpu_topology->platform_version(), "platform_version");
+  EXPECT_THAT(gpu_topology->num_slices(), 2);
+  EXPECT_THAT(gpu_topology->num_hosts_per_slice(), 1);
+  EXPECT_THAT(gpu_topology->num_devices_per_host(), 3);
 }
 
 TEST(GpuTopology, ToProto) {
   GpuTopology gpu_topology(/*gpu_device_ids=*/{3, 2, 1},
-                           /*platform_version=*/"platform_version");
+                           /*platform_version=*/"platform_version",
+                           /*num_slices=*/2,
+                           /*num_hosts_per_slice=*/1,
+                           /*num_devices_per_host=*/3);
   GpuTopologyProto msg = gpu_topology.ToProto();
   EXPECT_THAT(msg.device_ids(), ElementsAre(3, 2, 1));
   EXPECT_THAT(msg.platform_version(), "platform_version");
+  EXPECT_THAT(msg.num_slices(), 2);
+  EXPECT_THAT(msg.num_hosts_per_slice(), 1);
+  EXPECT_THAT(msg.num_devices_per_host(), 3);
 }
 
 TEST(StreamExecutorGpuClientTest, DistributedInit) {

From f24bdda19768b8fcd230326d5b6e1b0b27be4fce Mon Sep 17 00:00:00 2001
From: Parker Schuh <parkers@google.com>
Date: Mon, 13 May 2024 11:44:31 -0700
Subject: [PATCH 041/478] manual dims must be excluded when computing the new
 tile assignment, or num_new_tiles will always be > NumTiles() + 1 and
 sharding propagation will infinitely reassign the same partial manual
 shardings.

PiperOrigin-RevId: 633282263
---
 .../xla/xla/hlo/utils/hlo_sharding_util.cc    | 19 +++++++++++++------
 .../xla/hlo/utils/hlo_sharding_util_test.cc   | 11 +++++++++++
 2 files changed, 24 insertions(+), 6 deletions(-)

diff --git a/third_party/xla/xla/hlo/utils/hlo_sharding_util.cc b/third_party/xla/xla/hlo/utils/hlo_sharding_util.cc
index 96a2506d59ce5b..0133703e0abd68 100644
--- a/third_party/xla/xla/hlo/utils/hlo_sharding_util.cc
+++ b/third_party/xla/xla/hlo/utils/hlo_sharding_util.cc
@@ -352,6 +352,13 @@ bool MergeShardingIfCompatible(const HloSharding& to_merge,
     }
   }
 
+  const int64_t num_devices = to_merge.tile_assignment().num_elements();
+  const int64_t new_num_tiles = Product(merged_tile_dims);
+  if (num_devices % new_num_tiles != 0 || new_num_tiles < minimum_tiles) {
+    return false;
+  }
+  int64_t replication;
+
   if (to_merge_man_dim >= 0) {
     int64_t man_group_size = to_merge.tile_assignment().dim(to_merge_man_dim);
     if (man_group_size != dst->tile_assignment().dim(dst_man_dim)) {
@@ -365,14 +372,14 @@ bool MergeShardingIfCompatible(const HloSharding& to_merge,
     merged_tile_dims.push_back(man_group_size);
     num_merge_groups *= man_group_size;
     num_dst_groups *= man_group_size;
+    if (num_devices % (new_num_tiles * man_group_size) != 0) {
+      return false;
+    }
+    replication = num_devices / (new_num_tiles * man_group_size);
+  } else {
+    replication = num_devices / new_num_tiles;
   }
 
-  const int64_t num_devices = to_merge.tile_assignment().num_elements();
-  const int64_t new_num_tiles = Product(merged_tile_dims);
-  if (num_devices % new_num_tiles != 0 || new_num_tiles < minimum_tiles) {
-    return false;
-  }
-  const int64_t replication = num_devices / new_num_tiles;
   if (replication > 1) {
     merged_tile_dims.push_back(replication);
   }
diff --git a/third_party/xla/xla/hlo/utils/hlo_sharding_util_test.cc b/third_party/xla/xla/hlo/utils/hlo_sharding_util_test.cc
index 126f414c789609..c9902e6a6ad14c 100644
--- a/third_party/xla/xla/hlo/utils/hlo_sharding_util_test.cc
+++ b/third_party/xla/xla/hlo/utils/hlo_sharding_util_test.cc
@@ -529,6 +529,17 @@ TEST(HloShardingUtilTest,
   EXPECT_EQ(result, output_sharding);
 }
 
+TEST(HloShardingUtilTest, MergeManualSubgroupSharding) {
+  TileAssignment tile_assignment({16, 4});
+  std::vector<OpSharding::Type> subgroup_types = {OpSharding::MANUAL,
+                                                  OpSharding::REPLICATED};
+  // Subgroup sharding
+  //  {devices=[16,4]<=[64] last_tile_dims={manual, replicated}}
+  HloSharding dst = HloSharding::Subgroup(tile_assignment, subgroup_types);
+  HloSharding to_merge = dst;
+  EXPECT_FALSE(MergeShardingIfCompatible(to_merge, dst.NumTiles() + 1, &dst));
+}
+
 TEST(HloShardingUtilTest, GetManualSubgroupSharding_ManualOnly) {
   TileAssignment tile_assignment({1, 2, 2});
   std::vector<OpSharding::Type> subgroup_types = {OpSharding::MANUAL};

From dadff518bb1eace98ad3f3d7ceed81319736cfbe Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Mon, 13 May 2024 11:53:10 -0700
Subject: [PATCH 042/478] Update ops-related pbtxt files.

PiperOrigin-RevId: 633284964
---
 tensorflow/core/ops/ops.pbtxt | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorflow/core/ops/ops.pbtxt b/tensorflow/core/ops/ops.pbtxt
index 68d78d0329a9ac..53c362ed0ab819 100644
--- a/tensorflow/core/ops/ops.pbtxt
+++ b/tensorflow/core/ops/ops.pbtxt
@@ -1,4 +1,4 @@
-go/nodeserialize   
+go/debugstr    
 op {
   name: "Abort"
   attr {

From 335abf5f52a0f2bb678cf93ca701de1245c5c86f Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Mon, 13 May 2024 12:09:05 -0700
Subject: [PATCH 043/478] Fix populate of supported_subgroup_sizes

PiperOrigin-RevId: 633290357
---
 tensorflow/lite/delegates/gpu/cl/cl_device.cc | 27 +++++++++----------
 1 file changed, 13 insertions(+), 14 deletions(-)

diff --git a/tensorflow/lite/delegates/gpu/cl/cl_device.cc b/tensorflow/lite/delegates/gpu/cl/cl_device.cc
index 852cc95eedf261..4f3ea93b1cf792 100644
--- a/tensorflow/lite/delegates/gpu/cl/cl_device.cc
+++ b/tensorflow/lite/delegates/gpu/cl/cl_device.cc
@@ -344,21 +344,20 @@ GpuInfo GpuInfoFromDeviceID(cl_device_id id, cl_platform_id platform_id) {
     }
   }
 
-  if (info.IsIntel()) {
-    if (info.SupportsExtension("cl_intel_required_subgroup_size")) {
-      size_t sub_groups_count;
-      cl_int status =
-          clGetDeviceInfo(id, 0x4108 /*CL_DEVICE_SUB_GROUP_SIZES_INTEL*/, 0,
-                          nullptr, &sub_groups_count);
+  if (info.SupportsExtension("cl_intel_required_subgroup_size")) {
+    size_t sub_groups_ret_size;
+    cl_int status =
+        clGetDeviceInfo(id, 0x4108 /*CL_DEVICE_SUB_GROUP_SIZES_INTEL*/, 0,
+                        nullptr, &sub_groups_ret_size);
+    if (status == CL_SUCCESS) {
+      size_t sub_groups_count = sub_groups_ret_size / sizeof(size_t);
+      std::vector<size_t> sub_group_sizes(sub_groups_count);
+      status =
+          clGetDeviceInfo(id, 0x4108 /*CL_DEVICE_SUB_GROUP_SIZES_INTEL*/,
+                          sub_groups_ret_size, sub_group_sizes.data(), nullptr);
       if (status == CL_SUCCESS) {
-        std::vector<size_t> sub_group_sizes(sub_groups_count);
-        status = clGetDeviceInfo(id, 0x4108 /*CL_DEVICE_SUB_GROUP_SIZES_INTEL*/,
-                                 sizeof(size_t) * sub_groups_count,
-                                 sub_group_sizes.data(), nullptr);
-        if (status == CL_SUCCESS) {
-          for (int i = 0; i < sub_groups_count; ++i) {
-            info.supported_subgroup_sizes.push_back(sub_group_sizes[i]);
-          }
+        for (int i = 0; i < sub_groups_count; ++i) {
+          info.supported_subgroup_sizes.push_back(sub_group_sizes[i]);
         }
       }
     }

From 88c4ac653ec8e04644a28dbbb25dbe5987a20a83 Mon Sep 17 00:00:00 2001
From: Mason Chang <masonchang@google.com>
Date: Mon, 13 May 2024 12:11:09 -0700
Subject: [PATCH 044/478] Keep identity ops that have the same input/output
 device assignment on the same devices

PiperOrigin-RevId: 633290991
---
 tensorflow/core/common_runtime/placer.cc      |  39 +++++-
 tensorflow/core/common_runtime/placer_test.cc | 126 ++++++++++++++++++
 2 files changed, 164 insertions(+), 1 deletion(-)

diff --git a/tensorflow/core/common_runtime/placer.cc b/tensorflow/core/common_runtime/placer.cc
index 05dc029cc74756..5908b26a898690 100644
--- a/tensorflow/core/common_runtime/placer.cc
+++ b/tensorflow/core/common_runtime/placer.cc
@@ -27,6 +27,7 @@ limitations under the License.
 #include "tensorflow/core/framework/graph.pb.h"
 #include "tensorflow/core/framework/types.h"
 #include "tensorflow/core/framework/types.pb.h"
+#include "tensorflow/core/graph/graph.h"
 #include "tensorflow/core/graph/graph_node_util.h"
 #include "tensorflow/core/lib/core/errors.h"
 #include "tensorflow/core/platform/path.h"
@@ -117,6 +118,42 @@ bool IsGeneratorNode(const Node* node) {
          !IsRefType(node->output_type(0));
 }
 
+// If a node is an Identity op with input and output on the same device,
+// assign this Identity the same device. If the node already has a requested
+// or assigned device, don't touch it.
+bool MatchIdentityOperation(const Node* node) {
+  if (!node) {
+    return false;
+  }
+
+  if (!node->IsIdentity()) {
+    return false;
+  }
+
+  if (node->has_assigned_device_name()) {
+    return false;
+  }
+
+  if (!node->requested_device().empty()) {
+    return false;
+  }
+
+  // Strictly only check for IDENTITY nodes with only 1 input and
+  // 1 output edge.
+  if (node->in_edges().size() != 1) {
+    return false;
+  }
+
+  if (node->out_edges().size() != 1) {
+    return false;
+  }
+
+  const Node* input = *node->in_nodes().begin();
+  const Node* output = *node->out_nodes().begin();
+
+  return input->requested_device() == output->requested_device();
+}
+
 void LogDeviceAssignment(const Node* node, bool log_device_placement) {
   // Log placement if log_device_placement is set.
   if (log_device_placement) {
@@ -257,7 +294,7 @@ Status Placer::Run(const GraphOptimizationPassOptions& options) {
     // Heuristic B: If the node only operates on metadata, not data,
     // then it is desirable to place that metadata node with its
     // input.
-    if (IsMetadata(node)) {
+    if (IsMetadata(node) || MatchIdentityOperation(node)) {
       // Make sure that the input device type is in the list of supported
       // device types for this node.
       const Node* input = (*node->in_edges().begin())->src();
diff --git a/tensorflow/core/common_runtime/placer_test.cc b/tensorflow/core/common_runtime/placer_test.cc
index aeca73546468c0..6d76bc7d12c86e 100644
--- a/tensorflow/core/common_runtime/placer_test.cc
+++ b/tensorflow/core/common_runtime/placer_test.cc
@@ -213,6 +213,8 @@ REGISTER_OP("TestTypedConsumer").Input("i: variant");
 REGISTER_KERNEL_BUILDER(Name("TestTypedConsumer").Device("FakeCPU"), DummyOp);
 REGISTER_KERNEL_BUILDER(Name("TestTypedConsumer").Device("FakeGPU"), DummyOp);
 
+REGISTER_OP("ConvertToListOfCooTensorsV2").Input("i: int32");
+
 ////////////////////////////////////////////////////////////////////////////////
 //
 // A PlacerTest method has three phases:
@@ -1948,6 +1950,9 @@ REGISTER_KERNEL_BUILDER(Name("Add").Device("FakeCPU"), DummyOp);
 REGISTER_KERNEL_BUILDER(Name("Add").Device("FakeGPU"), DummyOp);
 REGISTER_KERNEL_BUILDER(Name("PartitionedCall").Device("FakeCPU"), DummyOp);
 REGISTER_KERNEL_BUILDER(Name("PartitionedCall").Device("FakeGPU"), DummyOp);
+REGISTER_KERNEL_BUILDER(Name("ConvertToListOfCooTensorsV2").Device("FakeCPU"),
+                        DummyOp);
+REGISTER_KERNEL_BUILDER(Name("Cast").Device("FakeCPU"), DummyOp);
 
 TEST_P(SoftPlacementPlacerTest,
        RequestedDeviceOnResourceGeneratorIsTreatedAsAssigned) {
@@ -3108,5 +3113,126 @@ TEST_F(NestedPlacerTest, IndirectRecursion) {
       << s.ToString();
 }
 
+TEST_F(PlacerTest, IdentityMatchesInputAndOutputPlacement) {
+  /*
+   *     Op Input (assigned to task:1)
+   *       |
+   *       v
+   *     // Tests that this gets reassigned to task:1
+   *     Identity (No Assignment)
+   *       |
+   *       v
+   *     Op Output (assigned to task:1)
+   */
+  const std::string task0_device = "/job:b/replica:0/task:0/device:FakeCPU:0";
+  const std::string task1_device = "/job:b/replica:0/task:1/device:FakeCPU:0";
+
+  GraphDef graph = GDef({
+      NDef("a", "_Arg", {}, {{"T", DT_FLOAT}}, task1_device),
+      NDef("identity1", "Identity", {"a"}, {{"T", DT_FLOAT}}, task1_device),
+      NDef("identity2", "Identity", {"identity1:0"}, {{"T", DT_FLOAT}}),
+      NDef("cast", "Cast", {"identity2:0"},
+           {{"SrcT", DT_FLOAT}, {"DstT", DT_INT32}}, task1_device),
+      NDef("COO", "ConvertToListOfCooTensorsV2", {"cast:0"}, {{"T", DT_INT32}},
+           task1_device),
+  });
+
+  Graph g(OpRegistry::Global());
+
+  DeviceSet multiple_tasks;
+  std::unique_ptr<Device> task0_cpu(FakeDevice::MakeCPU(task0_device));
+  multiple_tasks.AddDevice(task0_cpu.get());
+
+  std::unique_ptr<Device> task1_cpu(FakeDevice::MakeCPU(task1_device));
+  multiple_tasks.AddDevice(task1_cpu.get());
+
+  TF_ASSERT_OK(BuildGraph(graph, &g));
+
+  absl::Status s = Place(&g, &multiple_tasks);
+  TF_ASSERT_OK(s);
+
+  Node* identity2 = GetNodeByName(g, "identity2");
+  EXPECT_EQ(identity2->assigned_device_name().c_str(), task1_device);
+}
+
+TEST_F(PlacerTest, IdentityWithoutOutputDoesntCrash) {
+  /*
+   *     Op Input (assigned to task:1)
+   *       |
+   *       v
+   *     // Tests that this doesn't crash.
+   *     Identity (No output)
+   */
+  const std::string task0_device = "/job:b/replica:0/task:0/device:FakeCPU:0";
+  const std::string task1_device = "/job:b/replica:0/task:1/device:FakeCPU:0";
+
+  GraphDef graph = GDef({
+      NDef("a", "_Arg", {}, {{"T", DT_FLOAT}}, task1_device),
+      NDef("identity1", "Identity", {"a"}, {{"T", DT_FLOAT}}, task1_device),
+      NDef("identity2", "Identity", {"identity1:0"}, {{"T", DT_FLOAT}}),
+  });
+
+  Graph g(OpRegistry::Global());
+
+  DeviceSet multiple_tasks;
+  std::unique_ptr<Device> task0_cpu(FakeDevice::MakeCPU(task0_device));
+  multiple_tasks.AddDevice(task0_cpu.get());
+
+  std::unique_ptr<Device> task1_cpu(FakeDevice::MakeCPU(task1_device));
+  multiple_tasks.AddDevice(task1_cpu.get());
+
+  TF_ASSERT_OK(BuildGraph(graph, &g));
+  Node* identity2 = GetNodeByName(g, "identity2");
+  const Edge* out_edge = *identity2->out_edges().begin();
+
+  g.RemoveEdge(out_edge);
+
+  absl::Status s = Place(&g, &multiple_tasks);
+  TF_ASSERT_OK(s);
+}
+
+TEST_F(PlacerTest, IdentityDoesntMatchWithMultipleOutput) {
+  /*
+   *     Op Input (assigned to task:1)
+   *       |
+   *       v
+   *     // Tests that identity gets assigned to default task:0
+   *     Identity (No Assignment)
+   *       |
+   *       v
+   *     Multiple Op Output (assigned to task:1)
+   */
+  const std::string task0_device = "/job:b/replica:0/task:0/device:FakeCPU:0";
+  const std::string task1_device = "/job:b/replica:0/task:1/device:FakeCPU:0";
+
+  GraphDef graph = GDef({
+      NDef("a", "_Arg", {}, {{"T", DT_FLOAT}}, task1_device),
+      NDef("identity1", "Identity", {"a"}, {{"T", DT_FLOAT}}, task1_device),
+      NDef("identity2", "Identity", {"identity1:0"}, {{"T", DT_FLOAT}}),
+      NDef("cast", "Cast", {"identity2:0"},
+           {{"SrcT", DT_FLOAT}, {"DstT", DT_INT32}}, task1_device),
+      NDef("COO", "ConvertToListOfCooTensorsV2", {"cast:0"}, {{"T", DT_INT32}},
+           task1_device),
+      NDef("identity3", "Identity", {"identity2:0"}, {{"T", DT_FLOAT}}),
+  });
+
+  Graph g(OpRegistry::Global());
+
+  DeviceSet multiple_tasks;
+  std::unique_ptr<Device> task0_cpu(FakeDevice::MakeCPU(task0_device));
+  multiple_tasks.AddDevice(task0_cpu.get());
+
+  std::unique_ptr<Device> task1_cpu(FakeDevice::MakeCPU(task1_device));
+  multiple_tasks.AddDevice(task1_cpu.get());
+
+  TF_ASSERT_OK(BuildGraph(graph, &g));
+
+  absl::Status s = Place(&g, &multiple_tasks);
+  TF_ASSERT_OK(s);
+
+  Node* identity2 = GetNodeByName(g, "identity2");
+  EXPECT_EQ(identity2->assigned_device_name().c_str(), task0_device);
+}
+
 }  // namespace
 }  // namespace tensorflow

From 6607b1f2e6d2aa9b06669bf8a69b9f628212e6c6 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Mon, 13 May 2024 12:26:49 -0700
Subject: [PATCH 045/478] Fix C++ dtor bypass virtual call dispatch lint.

PiperOrigin-RevId: 633295761
---
 .../coordination/coordination_service_agent.cc         | 10 +++++++---
 1 file changed, 7 insertions(+), 3 deletions(-)

diff --git a/third_party/xla/xla/tsl/distributed_runtime/coordination/coordination_service_agent.cc b/third_party/xla/xla/tsl/distributed_runtime/coordination/coordination_service_agent.cc
index 36969eb0025205..1d090c7a0ffbc7 100644
--- a/third_party/xla/xla/tsl/distributed_runtime/coordination/coordination_service_agent.cc
+++ b/third_party/xla/xla/tsl/distributed_runtime/coordination/coordination_service_agent.cc
@@ -71,9 +71,7 @@ class CoordinationServiceAgentImpl : public CoordinationServiceAgent {
  public:
   CoordinationServiceAgentImpl() = default;
   ~CoordinationServiceAgentImpl() override {
-    // TODO(b/339231167): Fix the lint.
-    absl::Status s =
-        Shutdown();  // NOLINT(clang-analyzer-optin.cplusplus.VirtualCall)
+    absl::Status s = ShutdownInternal();
     VLOG(3) << "Coordination agent dtor failed with status: " << s;
   }
   absl::Status Initialize(Env* env, std::string_view job_name, int task_id,
@@ -139,6 +137,8 @@ class CoordinationServiceAgentImpl : public CoordinationServiceAgent {
   void StopHeartbeat();
 
  private:
+  absl::Status ShutdownInternal();
+
   Env* env_ = nullptr;  // Not owned.
   const uint64_t incarnation_id_ = random::New64();
   CoordinatedTask task_;
@@ -477,6 +477,10 @@ absl::Status CoordinationServiceAgentImpl::ReportError(
 }
 
 absl::Status CoordinationServiceAgentImpl::Shutdown() {
+  return ShutdownInternal();
+}
+
+absl::Status CoordinationServiceAgentImpl::ShutdownInternal() {
   absl::Status status = absl::OkStatus();
   bool is_connected = false;
   {

From 86ea1d18bf750ba46c3998422b84eb2c6c3cd306 Mon Sep 17 00:00:00 2001
From: Antonio Sanchez <cantonios@google.com>
Date: Mon, 13 May 2024 12:29:26 -0700
Subject: [PATCH 046/478] Propagate status from sparse reduction op to avoid
 crash.

Fixes #65865

PiperOrigin-RevId: 633296435
---
 tensorflow/core/kernels/BUILD               |  2 +-
 tensorflow/core/kernels/sparse_reduce_op.cc | 26 ++++++++++++++-------
 2 files changed, 19 insertions(+), 9 deletions(-)

diff --git a/tensorflow/core/kernels/BUILD b/tensorflow/core/kernels/BUILD
index c541597c8d9d80..158ace64ba184b 100644
--- a/tensorflow/core/kernels/BUILD
+++ b/tensorflow/core/kernels/BUILD
@@ -5138,7 +5138,7 @@ tf_kernel_library(
 tf_kernel_library(
     name = "sparse_reduce_op",
     prefix = "sparse_reduce_op",
-    deps = SPARSE_DEPS,
+    deps = SPARSE_DEPS + ["@com_google_absl//absl/status"],
 )
 
 tf_kernel_library(
diff --git a/tensorflow/core/kernels/sparse_reduce_op.cc b/tensorflow/core/kernels/sparse_reduce_op.cc
index 348a73e0816280..8f875053c66579 100644
--- a/tensorflow/core/kernels/sparse_reduce_op.cc
+++ b/tensorflow/core/kernels/sparse_reduce_op.cc
@@ -17,6 +17,7 @@ limitations under the License.
 
 #define EIGEN_USE_THREADS
 
+#include "absl/status/status.h"
 #include "tensorflow/core/framework/op_kernel.h"
 #include "tensorflow/core/framework/op_requires.h"
 #include "tensorflow/core/framework/register_types.h"
@@ -47,13 +48,14 @@ struct ReduceDetails {
 
 // Compute common reduce parameters that'll be used for SparseTensor
 // reductions. Usage:
-// ReduceDetails reduction = SparseTensorReduceHelper(sp, axes, keep_dims);
-// sp.Reorder(reduction.reorder_dims);
-// for (const auto& g : sp.group(reduction.group_by_dims)) {
+// StatusOr<ReduceDetails> reduction =
+//     SparseTensorReduceHelper(sp, axes, keep_dims);
+// sp.Reorder(reduction->reorder_dims);
+// for (const auto& g : sp.group(reduction->group_by_dims)) {
 //   ...
 // }
-// // Set output shape to reduction.reduced_shape.
-ReduceDetails SparseTensorReduceHelper(const SparseTensor &sp,
+// // Set output shape to reduction->reduced_shape.
+absl::StatusOr<ReduceDetails> SparseTensorReduceHelper(const SparseTensor &sp,
                                        absl::Span<const int32> axes_slice,
                                        bool keep_dims) {
   ReduceDetails reduction;
@@ -101,7 +103,11 @@ ReduceDetails SparseTensorReduceHelper(const SparseTensor &sp,
     out_dim_sizes = sp.PickDims(reduction.group_by_dims);
   }
 
-  reduction.reduced_shape = TensorShape(out_dim_sizes);
+  absl::Status success =
+      TensorShape::BuildTensorShape(out_dim_sizes, &reduction.reduced_shape);
+  if (!success.ok()) {
+    return success;
+  }
   return reduction;
 }
 
@@ -181,8 +187,10 @@ class SparseReduceOp : public OpKernel {
     OP_REQUIRES_OK(ctx, SparseTensor::Create(
         tensor::DeepCopy(*indices_t), tensor::DeepCopy(*values_t),
                     shape, &sp));
-    ReduceDetails reduction = SparseTensorReduceHelper(
+    absl::StatusOr<ReduceDetails> reduction_or = SparseTensorReduceHelper(
         sp, reduction_axes_t->flat<int32>(), keep_dims_);
+    OP_REQUIRES_OK(ctx, reduction_or.status());
+    ReduceDetails reduction = *reduction_or;
 
     Tensor *out_values;
     OP_REQUIRES_OK(
@@ -287,8 +295,10 @@ class SparseReduceSparseOp : public OpKernel {
     OP_REQUIRES_OK(ctx, SparseTensor::Create(tensor::DeepCopy(*indices_t),
                                          tensor::DeepCopy(*values_t),
                     shape, &sp));
-    ReduceDetails reduction = SparseTensorReduceHelper(
+    absl::StatusOr<ReduceDetails> reduction_or = SparseTensorReduceHelper(
         sp, reduction_axes_t->flat<int32>(), keep_dims_);
+    OP_REQUIRES_OK(ctx, reduction_or.status());
+    ReduceDetails reduction = *reduction_or;
 
     sp.Reorder<T>(reduction.reorder_dims);
     // Count nnzs in the output SparseTensor.

From 71c1529afeb44d7a6428a7681b8c0294c67dd67b Mon Sep 17 00:00:00 2001
From: Kuy Mainwaring <kuym@google.com>
Date: Mon, 13 May 2024 12:53:59 -0700
Subject: [PATCH 047/478] [XLA:GPU] Clang-tidy fixes for
 xla/service/gpu/fusions/transpose_test.cc

PiperOrigin-RevId: 633303182
---
 third_party/xla/xla/service/gpu/fusions/transpose_test.cc | 2 --
 1 file changed, 2 deletions(-)

diff --git a/third_party/xla/xla/service/gpu/fusions/transpose_test.cc b/third_party/xla/xla/service/gpu/fusions/transpose_test.cc
index d7363bbd39f382..dbafd613d15b2b 100644
--- a/third_party/xla/xla/service/gpu/fusions/transpose_test.cc
+++ b/third_party/xla/xla/service/gpu/fusions/transpose_test.cc
@@ -34,8 +34,6 @@ namespace xla {
 namespace gpu {
 namespace {
 
-using ::testing::HasSubstr;
-
 class TransposeTest : public HloTestBase {
  protected:
   stream_executor::DeviceDescription device_info_ =

From c945821b669c772f7ad2858eb11b27757685c535 Mon Sep 17 00:00:00 2001
From: Fangrui Song <maskray@google.com>
Date: Mon, 13 May 2024 13:03:00 -0700
Subject: [PATCH 048/478] Integrate LLVM at llvm/llvm-project@1066eb554770

Updates LLVM usage to match
[1066eb554770](https://github.com/llvm/llvm-project/commit/1066eb554770)

PiperOrigin-RevId: 633306010
---
 third_party/llvm/workspace.bzl | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/third_party/llvm/workspace.bzl b/third_party/llvm/workspace.bzl
index cd453a40ec3bc0..d795d29b2a9434 100644
--- a/third_party/llvm/workspace.bzl
+++ b/third_party/llvm/workspace.bzl
@@ -4,8 +4,8 @@ load("//third_party:repo.bzl", "tf_http_archive")
 
 def repo(name):
     """Imports LLVM."""
-    LLVM_COMMIT = "fc57f88f007497a4ead0ec8607ac66e1847b02d6"
-    LLVM_SHA256 = "0b66773795454d466ef4dcfae7cf38c8200ac4ee431e069ddf68313b3486b004"
+    LLVM_COMMIT = "1066eb55477044a3a92f3a40471375194dfcdbc8"
+    LLVM_SHA256 = "dc0fc82c184cb9ccd094f5ff821914d610d79529e82f0cbcb91ddedebff29a0a"
 
     tf_http_archive(
         name = name,

From b0fd358c314a7b7a97546ceecd59041fbf1ee3b0 Mon Sep 17 00:00:00 2001
From: Kuy Mainwaring <kuym@google.com>
Date: Mon, 13 May 2024 13:12:44 -0700
Subject: [PATCH 049/478] [XLA:GPU] Clang-tidy fixes for
 xla/service/gpu/fusions/tiling_util.h

PiperOrigin-RevId: 633308984
---
 third_party/xla/xla/service/gpu/fusions/tiling_util.cc | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/third_party/xla/xla/service/gpu/fusions/tiling_util.cc b/third_party/xla/xla/service/gpu/fusions/tiling_util.cc
index 24456209e521fb..83f9e5fcfb1da7 100644
--- a/third_party/xla/xla/service/gpu/fusions/tiling_util.cc
+++ b/third_party/xla/xla/service/gpu/fusions/tiling_util.cc
@@ -201,7 +201,7 @@ absl::StatusOr<TilingThreadIdInfo> EmitThreadIdInfo(llvm::IRBuilder<>* builder,
 
 absl::StatusOr<TilingKernelInfo> EmitTilingKernel(
     llvm::IRBuilder<>* builder, const Tiling& tiling, llvm::Type* index_ty,
-    const TileGenerator& tile_generator) {
+    const TileGenerator& tile_element_generator) {
   absl::Span<const int64_t> dims_in_elems = tiling.GetShape();
   const auto& block_counts = tiling.GetBlockCounts();
   auto constant = [&](uint64_t c) -> llvm::Constant* {
@@ -249,7 +249,7 @@ absl::StatusOr<TilingKernelInfo> EmitTilingKernel(
                                    index_ty);
   }();
 
-  tile_generator(thread_id_info, tile_offset, tile_dimensions);
+  tile_element_generator(thread_id_info, tile_offset, tile_dimensions);
   return {{tile_dimensions, tile_offset, thread_id_info}};
 }
 

From eac02b8bc7efcdd922d6095393b090f5c9f344d1 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Mon, 13 May 2024 13:15:51 -0700
Subject: [PATCH 050/478] Migrate coord utils to use absl libraries directly +
 some clang fixes.

PiperOrigin-RevId: 633310000
---
 .../distributed_runtime/coordination/BUILD    |  12 +-
 .../coordination_service_error_util.h         |   4 +-
 .../coordination_service_rpc_handler.cc       | 143 +++++++++---------
 .../coordination_service_rpc_handler.h        |   7 +-
 4 files changed, 86 insertions(+), 80 deletions(-)

diff --git a/third_party/xla/xla/tsl/distributed_runtime/coordination/BUILD b/third_party/xla/xla/tsl/distributed_runtime/coordination/BUILD
index d7d5e7570970e5..33cf2f09025912 100644
--- a/third_party/xla/xla/tsl/distributed_runtime/coordination/BUILD
+++ b/third_party/xla/xla/tsl/distributed_runtime/coordination/BUILD
@@ -14,9 +14,9 @@ cc_library(
     name = "coordination_service_error_util",
     hdrs = ["coordination_service_error_util.h"],
     deps = [
+        "@com_google_absl//absl/status",
         "@com_google_absl//absl/strings",
-        "@local_tsl//tsl/platform:errors",
-        "@local_tsl//tsl/platform:status",
+        "@com_google_absl//absl/strings:cord",
         "@local_tsl//tsl/protobuf:coordination_service_proto_cc",
     ],
 )
@@ -195,12 +195,14 @@ cc_library(
         ":coordination_service_agent",
         ":coordination_service_error_util",
         "@com_google_absl//absl/algorithm:container",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/status:statusor",
+        "@com_google_absl//absl/strings",
+        "@com_google_absl//absl/synchronization",
         "@com_google_absl//absl/time",
         "@local_tsl//tsl/platform:casts",
-        "@local_tsl//tsl/platform:errors",
-        "@local_tsl//tsl/platform:mutex",
+        "@local_tsl//tsl/platform:protobuf",
         "@local_tsl//tsl/platform:status",
-        "@local_tsl//tsl/platform:statusor",
         "@local_tsl//tsl/platform:thread_annotations",
         "@local_tsl//tsl/protobuf:coordination_service_proto_cc",
     ],
diff --git a/third_party/xla/xla/tsl/distributed_runtime/coordination/coordination_service_error_util.h b/third_party/xla/xla/tsl/distributed_runtime/coordination/coordination_service_error_util.h
index a777d121c4ce93..4555a4e90e3a97 100644
--- a/third_party/xla/xla/tsl/distributed_runtime/coordination/coordination_service_error_util.h
+++ b/third_party/xla/xla/tsl/distributed_runtime/coordination/coordination_service_error_util.h
@@ -15,9 +15,9 @@ limitations under the License.
 #ifndef XLA_TSL_DISTRIBUTED_RUNTIME_COORDINATION_COORDINATION_SERVICE_ERROR_UTIL_H_
 #define XLA_TSL_DISTRIBUTED_RUNTIME_COORDINATION_COORDINATION_SERVICE_ERROR_UTIL_H_
 
+#include "absl/status/status.h"
+#include "absl/strings/cord.h"
 #include "absl/strings/string_view.h"
-#include "tsl/platform/errors.h"
-#include "tsl/platform/status.h"
 #include "tsl/protobuf/coordination_service.pb.h"
 
 namespace tsl {
diff --git a/third_party/xla/xla/tsl/distributed_runtime/coordination/coordination_service_rpc_handler.cc b/third_party/xla/xla/tsl/distributed_runtime/coordination/coordination_service_rpc_handler.cc
index ea15408886197c..920d7840486a9e 100644
--- a/third_party/xla/xla/tsl/distributed_runtime/coordination/coordination_service_rpc_handler.cc
+++ b/third_party/xla/xla/tsl/distributed_runtime/coordination/coordination_service_rpc_handler.cc
@@ -15,20 +15,23 @@ limitations under the License.
 
 #include "xla/tsl/distributed_runtime/coordination/coordination_service_rpc_handler.h"
 
+#include <cstdint>
 #include <iterator>
 #include <string>
 #include <utility>
 #include <vector>
 
 #include "absl/algorithm/container.h"
+#include "absl/status/status.h"
+#include "absl/status/statusor.h"
+#include "absl/strings/str_cat.h"
+#include "absl/synchronization/mutex.h"
 #include "absl/time/time.h"
 #include "xla/tsl/distributed_runtime/coordination/coordination_service.h"
 #include "xla/tsl/distributed_runtime/coordination/coordination_service_agent.h"
 #include "xla/tsl/distributed_runtime/coordination/coordination_service_error_util.h"
-#include "tsl/platform/casts.h"
-#include "tsl/platform/errors.h"
-#include "tsl/platform/mutex.h"
-#include "tsl/platform/statusor.h"
+#include "tsl/platform/protobuf.h"
+#include "tsl/platform/status.h"
 #include "tsl/protobuf/coordination_service.pb.h"
 
 namespace tsl {
@@ -40,23 +43,23 @@ using tensorflow::KeyValueEntry;
 
 void CoordinationServiceRpcHandler::SetAgentInstance(
     CoordinationServiceAgent* agent) {
-  mutex_lock l(mu_);
+  absl::MutexLock l(&mu_);
   agent_ = agent;
 }
 
 void CoordinationServiceRpcHandler::SetServiceInstance(
     CoordinationServiceInterface* service) {
-  mutex_lock l(mu_);
+  absl::MutexLock l(&mu_);
   service_ = service;
 }
 
 void CoordinationServiceRpcHandler::RegisterTaskAsync(
-    const RegisterTaskRequest* request, RegisterTaskResponse* response,
-    StatusCallback done) {
-  tf_shared_lock l(mu_);
+    const tensorflow::RegisterTaskRequest* request,
+    tensorflow::RegisterTaskResponse* response, StatusCallback done) {
+  absl::ReaderMutexLock l(&mu_);
   if (service_ == nullptr) {
     done(MakeCoordinationError(
-        errors::Internal("Coordination service is not enabled.")));
+        absl::InternalError("Coordination service is not enabled.")));
     return;
   }
   const CoordinatedTask& task = request->source_task();
@@ -67,12 +70,12 @@ void CoordinationServiceRpcHandler::RegisterTaskAsync(
 }
 
 void CoordinationServiceRpcHandler::HeartbeatAsync(
-    const HeartbeatRequest* request, HeartbeatResponse* response,
-    StatusCallback done) {
-  tf_shared_lock l(mu_);
+    const tensorflow::HeartbeatRequest* request,
+    tensorflow::HeartbeatResponse* response, StatusCallback done) {
+  absl::ReaderMutexLock l(&mu_);
   if (service_ == nullptr) {
     done(MakeCoordinationError(
-        errors::Internal("Coordination service is not enabled.")));
+        absl::InternalError("Coordination service is not enabled.")));
     return;
   }
   const CoordinatedTask& task = request->source_task();
@@ -88,12 +91,12 @@ void CoordinationServiceRpcHandler::HeartbeatAsync(
 }
 
 void CoordinationServiceRpcHandler::WaitForAllTasksAsync(
-    const WaitForAllTasksRequest* request, WaitForAllTasksResponse* response,
-    StatusCallback done) {
-  tf_shared_lock l(mu_);
+    const tensorflow::WaitForAllTasksRequest* request,
+    tensorflow::WaitForAllTasksResponse* response, StatusCallback done) {
+  absl::ReaderMutexLock l(&mu_);
   if (service_ == nullptr) {
     done(MakeCoordinationError(
-        errors::Internal("Coordination service is not enabled.")));
+        absl::InternalError("Coordination service is not enabled.")));
     return;
   }
   service_->WaitForAllTasks(
@@ -107,12 +110,12 @@ void CoordinationServiceRpcHandler::WaitForAllTasksAsync(
 }
 
 void CoordinationServiceRpcHandler::ShutdownTaskAsync(
-    const ShutdownTaskRequest* request, ShutdownTaskResponse* response,
-    StatusCallback done) {
-  tf_shared_lock l(mu_);
+    const tensorflow::ShutdownTaskRequest* request,
+    tensorflow::ShutdownTaskResponse* response, StatusCallback done) {
+  absl::ReaderMutexLock l(&mu_);
   if (service_ == nullptr) {
     done(MakeCoordinationError(
-        errors::Internal("Coordination service is not enabled.")));
+        absl::InternalError("Coordination service is not enabled.")));
     return;
   }
   service_->ShutdownTaskAsync(request->source_task(),
@@ -120,30 +123,30 @@ void CoordinationServiceRpcHandler::ShutdownTaskAsync(
 }
 
 void CoordinationServiceRpcHandler::ResetTaskAsync(
-    const ResetTaskRequest* request, ResetTaskResponse* response,
-    StatusCallback done) {
-  tf_shared_lock l(mu_);
+    const tensorflow::ResetTaskRequest* request,
+    tensorflow::ResetTaskResponse* response, StatusCallback done) {
+  absl::ReaderMutexLock l(&mu_);
   if (service_ == nullptr) {
     done(MakeCoordinationError(
-        errors::Internal("Coordination service is not enabled.")));
+        absl::InternalError("Coordination service is not enabled.")));
     return;
   }
   done(service_->ResetTask(request->source_task()));
 }
 
 void CoordinationServiceRpcHandler::ReportErrorToTaskAsync(
-    const ReportErrorToTaskRequest* request,
-    ReportErrorToTaskResponse* response, StatusCallback done) {
-  tf_shared_lock l(mu_);
+    const tensorflow::ReportErrorToTaskRequest* request,
+    tensorflow::ReportErrorToTaskResponse* response, StatusCallback done) {
+  absl::ReaderMutexLock l(&mu_);
   if (agent_ == nullptr) {
-    done(MakeCoordinationError(errors::Internal(
+    done(MakeCoordinationError(absl::InternalError(
         "CoordinationServiceAgent is uninitialized or has already shutdown.")));
     return;
   }
   const CoordinationServiceError& error_payload = request->error_payload();
   absl::Status error(
       static_cast<absl::StatusCode>(request->error_code()),
-      strings::StrCat(
+      absl::StrCat(
           "Error reported from /job:", error_payload.source_task().job_name(),
           "/task:", error_payload.source_task().task_id(), ": ",
           request->error_message()));
@@ -153,12 +156,12 @@ void CoordinationServiceRpcHandler::ReportErrorToTaskAsync(
 }
 
 void CoordinationServiceRpcHandler::ReportErrorToServiceAsync(
-    const ReportErrorToServiceRequest* request,
-    ReportErrorToServiceResponse* response, StatusCallback done) {
-  tf_shared_lock l(mu_);
+    const tensorflow::ReportErrorToServiceRequest* request,
+    tensorflow::ReportErrorToServiceResponse* response, StatusCallback done) {
+  absl::ReaderMutexLock l(&mu_);
   if (service_ == nullptr) {
     done(MakeCoordinationError(
-        errors::Internal("Coordination service is not enabled.")));
+        absl::InternalError("Coordination service is not enabled.")));
     return;
   }
   done(service_->ReportTaskError(
@@ -171,40 +174,40 @@ void CoordinationServiceRpcHandler::ReportErrorToServiceAsync(
 }
 
 void CoordinationServiceRpcHandler::GetTaskStateAsync(
-    const GetTaskStateRequest* request, GetTaskStateResponse* response,
-    StatusCallback done) {
-  tf_shared_lock l(mu_);
+    const tensorflow::GetTaskStateRequest* request,
+    tensorflow::GetTaskStateResponse* response, StatusCallback done) {
+  absl::ReaderMutexLock l(&mu_);
   if (service_ == nullptr) {
     done(MakeCoordinationError(
-        errors::Internal("Coordination service is not enabled.")));
+        absl::InternalError("Coordination service is not enabled.")));
     return;
   }
   auto result = service_->GetTaskState(
       {request->source_task().begin(), request->source_task().end()});
-  absl::c_move(result,
-               RepeatedFieldBackInserter(response->mutable_task_state()));
+  absl::c_move(result, tsl::protobuf::RepeatedFieldBackInserter(
+                           response->mutable_task_state()));
   done(absl::OkStatus());
 }
 
 void CoordinationServiceRpcHandler::InsertKeyValueAsync(
-    const InsertKeyValueRequest* request, InsertKeyValueResponse* response,
-    StatusCallback done) {
-  tf_shared_lock l(mu_);
+    const tensorflow::InsertKeyValueRequest* request,
+    tensorflow::InsertKeyValueResponse* response, StatusCallback done) {
+  absl::ReaderMutexLock l(&mu_);
   if (service_ == nullptr) {
     done(MakeCoordinationError(
-        errors::Internal("Coordination service is not enabled.")));
+        absl::InternalError("Coordination service is not enabled.")));
     return;
   }
   done(service_->InsertKeyValue(request->kv().key(), request->kv().value()));
 }
 
 void CoordinationServiceRpcHandler::GetKeyValueAsync(
-    const GetKeyValueRequest* request, GetKeyValueResponse* response,
-    StatusCallback done) {
-  tf_shared_lock l(mu_);
+    const tensorflow::GetKeyValueRequest* request,
+    tensorflow::GetKeyValueResponse* response, StatusCallback done) {
+  absl::ReaderMutexLock l(&mu_);
   if (service_ == nullptr) {
     done(MakeCoordinationError(
-        errors::Internal("Coordination service is not enabled.")));
+        absl::InternalError("Coordination service is not enabled.")));
     return;
   }
   response->mutable_kv()->set_key(request->key());
@@ -219,12 +222,12 @@ void CoordinationServiceRpcHandler::GetKeyValueAsync(
 }
 
 void CoordinationServiceRpcHandler::TryGetKeyValueAsync(
-    const TryGetKeyValueRequest* request, TryGetKeyValueResponse* response,
-    StatusCallback done) {
-  tf_shared_lock l(mu_);
+    const tensorflow::TryGetKeyValueRequest* request,
+    tensorflow::TryGetKeyValueResponse* response, StatusCallback done) {
+  absl::ReaderMutexLock l(&mu_);
   if (service_ == nullptr) {
     done(MakeCoordinationError(
-        errors::Internal("Coordination service is not enabled.")));
+        absl::InternalError("Coordination service is not enabled.")));
     return;
   }
   auto result = service_->TryGetKeyValue(request->key());
@@ -238,12 +241,12 @@ void CoordinationServiceRpcHandler::TryGetKeyValueAsync(
 }
 
 void CoordinationServiceRpcHandler::GetKeyValueDirAsync(
-    const GetKeyValueDirRequest* request, GetKeyValueDirResponse* response,
-    StatusCallback done) {
-  tf_shared_lock l(mu_);
+    const tensorflow::GetKeyValueDirRequest* request,
+    tensorflow::GetKeyValueDirResponse* response, StatusCallback done) {
+  absl::ReaderMutexLock l(&mu_);
   if (service_ == nullptr) {
     done(MakeCoordinationError(
-        errors::Internal("Coordination service is not enabled.")));
+        absl::InternalError("Coordination service is not enabled.")));
     return;
   }
   std::vector<KeyValueEntry> results =
@@ -254,24 +257,24 @@ void CoordinationServiceRpcHandler::GetKeyValueDirAsync(
 }
 
 void CoordinationServiceRpcHandler::DeleteKeyValueAsync(
-    const DeleteKeyValueRequest* request, DeleteKeyValueResponse* response,
-    StatusCallback done) {
-  tf_shared_lock l(mu_);
+    const tensorflow::DeleteKeyValueRequest* request,
+    tensorflow::DeleteKeyValueResponse* response, StatusCallback done) {
+  absl::ReaderMutexLock l(&mu_);
   if (service_ == nullptr) {
     done(MakeCoordinationError(
-        errors::Internal("Coordination service is not enabled.")));
+        absl::InternalError("Coordination service is not enabled.")));
     return;
   }
   done(service_->DeleteKeyValue(request->key()));
 }
 
-void CoordinationServiceRpcHandler::BarrierAsync(const BarrierRequest* request,
-                                                 BarrierResponse* response,
-                                                 StatusCallback done) {
-  tf_shared_lock l(mu_);
+void CoordinationServiceRpcHandler::BarrierAsync(
+    const tensorflow::BarrierRequest* request,
+    tensorflow::BarrierResponse* response, StatusCallback done) {
+  absl::ReaderMutexLock l(&mu_);
   if (service_ == nullptr) {
     done(MakeCoordinationError(
-        errors::Internal("Coordination service is not enabled.")));
+        absl::InternalError("Coordination service is not enabled.")));
     return;
   }
   std::vector<CoordinatedTask> tasks = {request->tasks().begin(),
@@ -284,12 +287,12 @@ void CoordinationServiceRpcHandler::BarrierAsync(const BarrierRequest* request,
 }
 
 void CoordinationServiceRpcHandler::CancelBarrierAsync(
-    const CancelBarrierRequest* request, CancelBarrierResponse* response,
-    StatusCallback done) {
-  tf_shared_lock l(mu_);
+    const tensorflow::CancelBarrierRequest* request,
+    tensorflow::CancelBarrierResponse* response, StatusCallback done) {
+  absl::ReaderMutexLock l(&mu_);
   if (service_ == nullptr) {
     done(MakeCoordinationError(
-        errors::Internal("Coordination service is not enabled.")));
+        absl::InternalError("Coordination service is not enabled.")));
     return;
   }
   done(service_->CancelBarrier(request->barrier_id(), request->source_task()));
diff --git a/third_party/xla/xla/tsl/distributed_runtime/coordination/coordination_service_rpc_handler.h b/third_party/xla/xla/tsl/distributed_runtime/coordination/coordination_service_rpc_handler.h
index 2895467f6e89d6..51d4f9f6901dc6 100644
--- a/third_party/xla/xla/tsl/distributed_runtime/coordination/coordination_service_rpc_handler.h
+++ b/third_party/xla/xla/tsl/distributed_runtime/coordination/coordination_service_rpc_handler.h
@@ -16,9 +16,10 @@ limitations under the License.
 #ifndef XLA_TSL_DISTRIBUTED_RUNTIME_COORDINATION_COORDINATION_SERVICE_RPC_HANDLER_H_
 #define XLA_TSL_DISTRIBUTED_RUNTIME_COORDINATION_COORDINATION_SERVICE_RPC_HANDLER_H_
 
+#include "absl/status/status.h"
+#include "absl/synchronization/mutex.h"
 #include "xla/tsl/distributed_runtime/coordination/coordination_service.h"
 #include "xla/tsl/distributed_runtime/coordination/coordination_service_agent.h"
-#include "tsl/platform/mutex.h"
 #include "tsl/platform/status.h"
 #include "tsl/platform/thread_annotations.h"
 #include "tsl/protobuf/coordination_service.pb.h"
@@ -26,7 +27,7 @@ limitations under the License.
 namespace tsl {
 class CoordinationServiceRpcHandler {
  public:
-  explicit CoordinationServiceRpcHandler() {}
+  explicit CoordinationServiceRpcHandler() = default;
 
   void SetAgentInstance(CoordinationServiceAgent* agent);
 
@@ -92,7 +93,7 @@ class CoordinationServiceRpcHandler {
                           StatusCallback done);
 
  private:
-  mutex mu_;
+  absl::Mutex mu_;
   CoordinationServiceAgent* agent_ TF_GUARDED_BY(mu_) = nullptr;
   CoordinationServiceInterface* service_ TF_GUARDED_BY(mu_) = nullptr;
 };

From 362bb373366c2a377e719cedf9867845aee807f5 Mon Sep 17 00:00:00 2001
From: Dimitris Vardoulakis <dvardoulakis@nvidia.com>
Date: Mon, 13 May 2024 13:17:38 -0700
Subject: [PATCH 051/478] PR #12049: Detect when we are running on an NVIDIA
 simulator and check that we pull the device description from gpu_specs.

Imported from GitHub PR https://github.com/openxla/xla/pull/12049

This helps when testing the compiler on new hardware platforms. The simulator may have nonsensical values for a few fields in the device description (such as device_memory_size == -1), so we want to use a hard-coded device description that contains the correct values.
Copybara import of the project:

--
1a488e45b3a7d0f7729fde9c28915c40188d17b3 by Dimitris Vardoulakis <dvardoulakis@nvidia.com>:

On Nvidia simulation, fail if target config hasn't been supplied using --xla_gpu_target_config_filename.

Merging this change closes #12049

PiperOrigin-RevId: 633310642
---
 third_party/xla/xla/service/gpu/gpu_compiler.cc | 15 ++++++++++++++-
 1 file changed, 14 insertions(+), 1 deletion(-)

diff --git a/third_party/xla/xla/service/gpu/gpu_compiler.cc b/third_party/xla/xla/service/gpu/gpu_compiler.cc
index 297eb7b98fdb71..35f6bf612f7ee6 100644
--- a/third_party/xla/xla/service/gpu/gpu_compiler.cc
+++ b/third_party/xla/xla/service/gpu/gpu_compiler.cc
@@ -1533,7 +1533,20 @@ absl::Status GpuCompiler::OptimizeHloPostLayoutAssignment(
     return Compiler::TargetConfig{gpu_target_config_proto};
   }
   if (executor) {
-    return Compiler::TargetConfig{executor};
+    Compiler::TargetConfig target_config = Compiler::TargetConfig{executor};
+    int64_t device_memory_size =
+        target_config.device_description.device_memory_size();
+    // Checking for device_memory_size == -1 is how we detect that we are
+    // running on Nvidia's software simulator. When running on simulation,
+    // the config from StreamExecutor is inaccurate, so we must load the
+    // hard-coded config from a file.
+    if (device_memory_size == -1) {
+      return absl::FailedPreconditionError(
+          "When running on an NVIDIA simulation device, you must use "
+          "--xla_gpu_target_config_filename to pass in target information. "
+          "The target config from StreamExecutor is inaccurate.");
+    }
+    return target_config;
   }
   return absl::InternalError(
       "Either GPU has to be attached, or --xla_gpu_target_config_filename "

From 30784c22ee80465fb7eff793c79d81ce800bc2b3 Mon Sep 17 00:00:00 2001
From: Kuy Mainwaring <kuym@google.com>
Date: Mon, 13 May 2024 13:18:29 -0700
Subject: [PATCH 052/478] [XLA:GPU] Clang-tidy fixes for
 xla/service/gpu/fusions/mlir/lower_tensors.cc

PiperOrigin-RevId: 633310993
---
 third_party/xla/xla/service/gpu/fusions/mlir/lower_tensors.cc | 1 -
 1 file changed, 1 deletion(-)

diff --git a/third_party/xla/xla/service/gpu/fusions/mlir/lower_tensors.cc b/third_party/xla/xla/service/gpu/fusions/mlir/lower_tensors.cc
index 5677b45e5342bf..43bc759b067221 100644
--- a/third_party/xla/xla/service/gpu/fusions/mlir/lower_tensors.cc
+++ b/third_party/xla/xla/service/gpu/fusions/mlir/lower_tensors.cc
@@ -78,7 +78,6 @@ using mlir::TypedValue;
 using mlir::TypeRange;
 using mlir::Value;
 using mlir::ValueRange;
-using mlir::arith::AtomicRMWKind;
 
 namespace arith = ::mlir::arith;
 namespace scf = ::mlir::scf;

From f9b3c69fb1dc02ec06962edf1ce660fe6efbb6e7 Mon Sep 17 00:00:00 2001
From: Kuy Mainwaring <kuym@google.com>
Date: Mon, 13 May 2024 13:18:39 -0700
Subject: [PATCH 053/478] [XLA:GPU] Clang-tidy fixes for
 xla/service/gpu/fusions/mlir/convert_xla_gpu_pure_call_ops.cc

PiperOrigin-RevId: 633311058
---
 .../service/gpu/fusions/mlir/convert_xla_gpu_pure_call_ops.cc    | 1 +
 1 file changed, 1 insertion(+)

diff --git a/third_party/xla/xla/service/gpu/fusions/mlir/convert_xla_gpu_pure_call_ops.cc b/third_party/xla/xla/service/gpu/fusions/mlir/convert_xla_gpu_pure_call_ops.cc
index 181005130b6a56..7f06615b91a9d2 100644
--- a/third_party/xla/xla/service/gpu/fusions/mlir/convert_xla_gpu_pure_call_ops.cc
+++ b/third_party/xla/xla/service/gpu/fusions/mlir/convert_xla_gpu_pure_call_ops.cc
@@ -10,6 +10,7 @@ See the License for the specific language governing permissions and
 limitations under the License.
 ==============================================================================*/
 #include <memory>
+#include <utility>
 
 #include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
 #include "mlir/IR/PatternMatch.h"  // from @llvm-project

From 329241d3241de9291d07e574668080f4ff61d890 Mon Sep 17 00:00:00 2001
From: Kuy Mainwaring <kuym@google.com>
Date: Mon, 13 May 2024 13:19:45 -0700
Subject: [PATCH 054/478] [XLA:GPU] Clang-tidy fixes for
 xla/service/gpu/fusions/triton.cc

PiperOrigin-RevId: 633311473
---
 third_party/xla/xla/service/gpu/fusions/triton.cc | 1 +
 1 file changed, 1 insertion(+)

diff --git a/third_party/xla/xla/service/gpu/fusions/triton.cc b/third_party/xla/xla/service/gpu/fusions/triton.cc
index ebbaccdb0bd742..30a98194e04200 100644
--- a/third_party/xla/xla/service/gpu/fusions/triton.cc
+++ b/third_party/xla/xla/service/gpu/fusions/triton.cc
@@ -18,6 +18,7 @@ limitations under the License.
 #include <optional>
 #include <string>
 #include <variant>
+#include <vector>
 
 #include "absl/log/check.h"
 #include "absl/log/log.h"

From a8fe248d93b5608bc1ed1a8fda14c60b131b7371 Mon Sep 17 00:00:00 2001
From: Kyle Lucke <klucke@google.com>
Date: Mon, 13 May 2024 13:29:47 -0700
Subject: [PATCH 055/478] Eliminate CommandBuffer::Create method in favor of
 the CreateCommandBuffer on StreamExecutorInterface.  This is a step of
 reducing circular dependencies between command_buffer.h and
 stream_executor_interface.h.

PiperOrigin-RevId: 633314618
---
 .../gpu/runtime/command_buffer_cmd_test.cc    |  9 +++--
 .../gpu/runtime/command_buffer_thunk.cc       |  4 +-
 .../xla/xla/stream_executor/command_buffer.cc |  7 +---
 .../xla/xla/stream_executor/command_buffer.h  |  8 ----
 .../gpu/gpu_command_buffer_test.cc            | 38 +++++++++----------
 5 files changed, 29 insertions(+), 37 deletions(-)

diff --git a/third_party/xla/xla/service/gpu/runtime/command_buffer_cmd_test.cc b/third_party/xla/xla/service/gpu/runtime/command_buffer_cmd_test.cc
index 7e1fa0d0cfd0c3..f96f4fa5d9e246 100644
--- a/third_party/xla/xla/service/gpu/runtime/command_buffer_cmd_test.cc
+++ b/third_party/xla/xla/service/gpu/runtime/command_buffer_cmd_test.cc
@@ -214,7 +214,8 @@ TEST(CommandBufferCmdTest, MemcpyCmd) {
 
   CommandBufferCmd::RecordParams record_params = {state};
 
-  auto command_buffer = se::CommandBuffer::Create(executor).value();
+  auto command_buffer =
+      executor->CreateCommandBuffer(se::CommandBuffer::Mode::kPrimary).value();
   TF_ASSERT_OK(commands.Record(params, record_params, command_buffer.get()));
 
   // Execute command buffer and verify that it copied the memory.
@@ -282,7 +283,8 @@ TEST(CommandBufferCmdTest, BarrierCmd) {
 
   CommandBufferCmd::RecordParams record_params = {state};
 
-  auto command_buffer = se::CommandBuffer::Create(executor).value();
+  auto command_buffer =
+      executor->CreateCommandBuffer(se::CommandBuffer::Mode::kPrimary).value();
   TF_ASSERT_OK(commands.Record(params, record_params, command_buffer.get()));
 
   // Execute command buffer and verify that it copied the memory.
@@ -359,7 +361,8 @@ TEST(CommandBufferCmdTest, LaunchCmd) {
 
   CommandBufferCmd::RecordParams record_params = {state};
 
-  auto command_buffer = se::CommandBuffer::Create(executor).value();
+  auto command_buffer =
+      executor->CreateCommandBuffer(se::CommandBuffer::Mode::kPrimary).value();
   TF_ASSERT_OK(commands.Record(params, record_params, command_buffer.get()));
 
   // Execute command buffer and verify that it copied the memory.
diff --git a/third_party/xla/xla/service/gpu/runtime/command_buffer_thunk.cc b/third_party/xla/xla/service/gpu/runtime/command_buffer_thunk.cc
index 4e74fdcd821c8b..8236ad2a65cb08 100644
--- a/third_party/xla/xla/service/gpu/runtime/command_buffer_thunk.cc
+++ b/third_party/xla/xla/service/gpu/runtime/command_buffer_thunk.cc
@@ -274,7 +274,9 @@ CommandBufferThunk::GetOrCreateCommandBuffer(se::StreamExecutor* executor) {
   }
 
   // Create a new empty command buffer.
-  TF_ASSIGN_OR_RETURN(auto command_buffer, se::CommandBuffer::Create(executor));
+  TF_ASSIGN_OR_RETURN(
+      auto command_buffer,
+      executor->CreateCommandBuffer(se::CommandBuffer::Mode::kPrimary));
   auto emplaced = state_->command_buffers.emplace(
       executor,
       std::make_shared<ExecutorCommandBuffer>(std::move(command_buffer)));
diff --git a/third_party/xla/xla/stream_executor/command_buffer.cc b/third_party/xla/xla/stream_executor/command_buffer.cc
index 1c5be3d8e0d80a..cd28a135bf519c 100644
--- a/third_party/xla/xla/stream_executor/command_buffer.cc
+++ b/third_party/xla/xla/stream_executor/command_buffer.cc
@@ -29,11 +29,6 @@ limitations under the License.
 
 namespace stream_executor {
 
-absl::StatusOr<std::unique_ptr<CommandBuffer>> CommandBuffer::Create(
-    StreamExecutorInterface* executor, Mode mode) {
-  return executor->CreateCommandBuffer(mode);
-}
-
 absl::StatusOr<std::unique_ptr<CommandBuffer>> CommandBuffer::Trace(
     StreamExecutorInterface* executor,
     absl::AnyInvocable<absl::Status(Stream*)> function, Mode mode) {
@@ -50,7 +45,7 @@ absl::StatusOr<std::unique_ptr<CommandBuffer>> CommandBuffer::Trace(
 
   // Prepare an empty command buffer instance.
   TF_ASSIGN_OR_RETURN(std::unique_ptr<CommandBuffer> command_buffer,
-                      CommandBuffer::Create(executor, mode));
+                      executor->CreateCommandBuffer(mode));
 
   // Trace and finalize the command buffer.
   TF_RETURN_IF_ERROR(
diff --git a/third_party/xla/xla/stream_executor/command_buffer.h b/third_party/xla/xla/stream_executor/command_buffer.h
index 1237493a2fb55e..a91963d46a0880 100644
--- a/third_party/xla/xla/stream_executor/command_buffer.h
+++ b/third_party/xla/xla/stream_executor/command_buffer.h
@@ -165,14 +165,6 @@ class CommandBuffer {
   // Command buffer constructors
   //===--------------------------------------------------------------------===//
 
-  // TODO(b/323534971): Command buffer constructors should be moved to
-  // StreamExecutor or a dedicated CommandBufferFactory accessible via
-  // StreamExecutor.
-
-  // Creates a new empty command buffer on the given executor.
-  static absl::StatusOr<std::unique_ptr<CommandBuffer>> Create(
-      StreamExecutorInterface* executor, Mode mode = Mode::kPrimary);
-
   // Creates a new command buffer on the given executor by tracing `function`
   // invocation. All StreamExecutor operations on a Stream argument will be
   // recorded into the command buffer. Returned command buffer is finalized, and
diff --git a/third_party/xla/xla/stream_executor/gpu/gpu_command_buffer_test.cc b/third_party/xla/xla/stream_executor/gpu/gpu_command_buffer_test.cc
index 4440393cc3bb8d..bd50ac29429b0e 100644
--- a/third_party/xla/xla/stream_executor/gpu/gpu_command_buffer_test.cc
+++ b/third_party/xla/xla/stream_executor/gpu/gpu_command_buffer_test.cc
@@ -125,7 +125,7 @@ TEST(GpuCommandBufferTest, LaunchSingleKernel) {
   TF_ASSERT_OK(stream->MemZero(&c, byte_length));
 
   // Create a command buffer with a single kernel launch.
-  auto cmd_buffer = CommandBuffer::Create(executor).value();
+  auto cmd_buffer = executor->CreateCommandBuffer(primary).value();
   TF_ASSERT_OK(cmd_buffer->Launch(add, ThreadDim(), BlockDim(4), a, b, c));
   TF_ASSERT_OK(cmd_buffer->Finalize());
 
@@ -239,8 +239,8 @@ TEST(GpuCommandBufferTest, LaunchNestedCommandBuffer) {
   TF_ASSERT_OK(stream->MemZero(&c, byte_length));
 
   // Create a command buffer with a single kernel launch.
-  auto primary_cmd = CommandBuffer::Create(executor).value();
-  auto nested_cmd = CommandBuffer::Create(executor, nested).value();
+  auto primary_cmd = executor->CreateCommandBuffer(primary).value();
+  auto nested_cmd = executor->CreateCommandBuffer(nested).value();
   TF_ASSERT_OK(nested_cmd->Launch(add, ThreadDim(), BlockDim(4), a, b, c));
   TF_ASSERT_OK(primary_cmd->AddNestedCommandBuffer(*nested_cmd));
   TF_ASSERT_OK(primary_cmd->Finalize());
@@ -260,7 +260,7 @@ TEST(GpuCommandBufferTest, LaunchNestedCommandBuffer) {
 
   // Update command buffer to write into `d` buffer by creating a new nested
   // command buffer.
-  nested_cmd = CommandBuffer::Create(executor, nested).value();
+  nested_cmd = executor->CreateCommandBuffer(nested).value();
   TF_ASSERT_OK(nested_cmd->Launch(add, ThreadDim(), BlockDim(4), a, b, d));
   TF_ASSERT_OK(primary_cmd->Update());
   TF_ASSERT_OK(primary_cmd->AddNestedCommandBuffer(*nested_cmd));
@@ -290,7 +290,7 @@ TEST(GpuCommandBufferTest, MemcpyDeviceToDevice) {
   TF_ASSERT_OK(stream->Memset32(&a, 42, byte_length));
 
   // Create a command buffer with a single a to b memcpy command.
-  auto cmd_buffer = CommandBuffer::Create(executor).value();
+  auto cmd_buffer = executor->CreateCommandBuffer(primary).value();
   TF_ASSERT_OK(cmd_buffer->MemcpyDeviceToDevice(&b, a, byte_length));
   TF_ASSERT_OK(cmd_buffer->Finalize());
 
@@ -331,7 +331,7 @@ TEST(GpuCommandBufferTest, Memset) {
   DeviceMemory<int32_t> a = executor->AllocateArray<int32_t>(length, 0);
 
   // Create a command buffer with a single memset command.
-  auto cmd_buffer = CommandBuffer::Create(executor).value();
+  auto cmd_buffer = executor->CreateCommandBuffer(primary).value();
   TF_ASSERT_OK(cmd_buffer->Memset(&a, uint32_t{42}, length));
   TF_ASSERT_OK(cmd_buffer->Finalize());
 
@@ -402,7 +402,7 @@ TEST(GpuCommandBufferTest, Barriers) {
   };
 
   // Create a command buffer with a DAG of memset commands.
-  auto cmd_buffer = CommandBuffer::Create(executor).value();
+  auto cmd_buffer = executor->CreateCommandBuffer(primary).value();
   TF_ASSERT_OK(record(cmd_buffer.get(), 42));
   TF_ASSERT_OK(executor->Submit(stream.get(), *cmd_buffer));
 
@@ -482,7 +482,7 @@ TEST(GpuCommandBufferTest, IndependentExecutionScopes) {
   };
 
   // Create a command buffer with a DAG of memset commands.
-  auto cmd_buffer = CommandBuffer::Create(executor).value();
+  auto cmd_buffer = executor->CreateCommandBuffer(primary).value();
   TF_ASSERT_OK(record(cmd_buffer.get(), 42));
   TF_ASSERT_OK(executor->Submit(stream.get(), *cmd_buffer));
 
@@ -556,7 +556,7 @@ TEST(GpuCommandBufferTest, ExecutionScopeBarriers) {
   };
 
   // Create a command buffer with a DAG of memset commands.
-  auto cmd_buffer = CommandBuffer::Create(executor).value();
+  auto cmd_buffer = executor->CreateCommandBuffer(primary).value();
   TF_ASSERT_OK(record(cmd_buffer.get(), 42));
   TF_ASSERT_OK(executor->Submit(stream.get(), *cmd_buffer));
 
@@ -646,7 +646,7 @@ TEST(GpuCommandBufferTest, ExecutionScopeOneDirectionalBarriers) {
   };
 
   // Create a command buffer with a DAG of memset commands.
-  auto cmd_buffer = CommandBuffer::Create(executor).value();
+  auto cmd_buffer = executor->CreateCommandBuffer(primary).value();
   TF_ASSERT_OK(record(cmd_buffer.get(), 42));
   TF_ASSERT_OK(executor->Submit(stream.get(), *cmd_buffer));
 
@@ -720,7 +720,7 @@ TEST(GpuCommandBufferTest, ConditionalIf) {
   };
 
   // Create a command buffer with a single conditional operation.
-  auto cmd_buffer = CommandBuffer::Create(executor).value();
+  auto cmd_buffer = executor->CreateCommandBuffer(primary).value();
   TF_ASSERT_OK(cmd_buffer->If(pred, then_builder));
   TF_ASSERT_OK(cmd_buffer->Finalize());
 
@@ -817,7 +817,7 @@ TEST(GpuCommandBufferTest, ConditionalIfElse) {
   };
 
   // Create a command buffer with a single conditional operation.
-  auto cmd_buffer = CommandBuffer::Create(executor).value();
+  auto cmd_buffer = executor->CreateCommandBuffer(primary).value();
   TF_ASSERT_OK(cmd_buffer->IfElse(pred, then_builder, else_builder));
   TF_ASSERT_OK(cmd_buffer->Finalize());
 
@@ -912,7 +912,7 @@ TEST(GpuCommandBufferTest, ConditionalCase) {
   };
 
   // Create a command buffer with a single conditional operation.
-  auto cmd_buffer = CommandBuffer::Create(executor).value();
+  auto cmd_buffer = executor->CreateCommandBuffer(primary).value();
   TF_ASSERT_OK(cmd_buffer->Case(index, {branch0, branch1}));
   TF_ASSERT_OK(cmd_buffer->Finalize());
 
@@ -991,7 +991,7 @@ TEST(GpuCommandBufferTest, ConditionalFor) {
   int32_t num_iters = 10;
 
   // Create a command buffer with a single conditional operation.
-  auto cmd_buffer = CommandBuffer::Create(executor).value();
+  auto cmd_buffer = executor->CreateCommandBuffer(primary).value();
   TF_ASSERT_OK(cmd_buffer->For(num_iters, loop_counter, body_builder));
   TF_ASSERT_OK(cmd_buffer->Finalize());
 
@@ -1058,7 +1058,7 @@ TEST(GpuCommandBufferTest, ConditionalWhile) {
   };
 
   // Create a command buffer with a single conditional operation.
-  auto cmd_buffer = CommandBuffer::Create(executor).value();
+  auto cmd_buffer = executor->CreateCommandBuffer(primary).value();
   TF_ASSERT_OK(cmd_buffer->While(pred, cond_builder, body_builder));
   TF_ASSERT_OK(cmd_buffer->Finalize());
 
@@ -1125,7 +1125,7 @@ TEST(GpuCommandBufferTest, ConditionalIfInExecutionScope) {
   };
 
   // Create a command buffer with a DAG of memset commands.
-  auto cmd_buffer = CommandBuffer::Create(executor).value();
+  auto cmd_buffer = executor->CreateCommandBuffer(primary).value();
   TF_ASSERT_OK(record(cmd_buffer.get(), 42));
   TF_ASSERT_OK(executor->Submit(stream.get(), *cmd_buffer));
 
@@ -1226,7 +1226,7 @@ TEST(GpuCommandBufferTest, ConditionalWhileInExecutionScope) {
   };
 
   // Create a command buffer with a single conditional operation.
-  auto cmd_buffer = CommandBuffer::Create(executor).value();
+  auto cmd_buffer = executor->CreateCommandBuffer(primary).value();
   TF_ASSERT_OK(record(cmd_buffer.get(), 42, 10));
   TF_ASSERT_OK(executor->Submit(stream.get(), *cmd_buffer));
 
@@ -1290,7 +1290,7 @@ static void BM_CreateCommandBuffer(benchmark::State& state) {
   DeviceMemory<int32_t> b = executor->AllocateArray<int32_t>(1, 0);
 
   for (auto s : state) {
-    auto cmd_buffer = CommandBuffer::Create(executor, nested).value();
+    auto cmd_buffer = executor->CreateCommandBuffer(nested).value();
     for (int i = 1; i < state.range(0); ++i) {
       CHECK_OK(cmd_buffer->Launch(add, ThreadDim(), BlockDim(4), b, b, b));
     }
@@ -1336,7 +1336,7 @@ static void BM_UpdateCommandBuffer(benchmark::State& state) {
 
   DeviceMemory<int32_t> b = executor->AllocateArray<int32_t>(1, 0);
 
-  auto cmd_buffer = CommandBuffer::Create(executor, primary).value();
+  auto cmd_buffer = executor->CreateCommandBuffer(primary).value();
   for (int i = 1; i < state.range(0); ++i) {
     CHECK_OK(cmd_buffer->Launch(add, ThreadDim(), BlockDim(4), b, b, b));
   }

From 3cf76f38d220c81a03b71bf464232e908093004f Mon Sep 17 00:00:00 2001
From: Eugene Zhulenev <ezhulenev@google.com>
Date: Mon, 13 May 2024 13:39:34 -0700
Subject: [PATCH 056/478] [xla:cpu] Fix msan warnings from arguments coming
 from jit compiled function

PiperOrigin-RevId: 633317729
---
 third_party/xla/xla/service/cpu/BUILD              |  1 +
 .../xla/xla/service/cpu/runtime_handle_ffi_call.cc | 14 ++++++++++++++
 2 files changed, 15 insertions(+)

diff --git a/third_party/xla/xla/service/cpu/BUILD b/third_party/xla/xla/service/cpu/BUILD
index 9d115f91eb2092..125af6dbce9170 100644
--- a/third_party/xla/xla/service/cpu/BUILD
+++ b/third_party/xla/xla/service/cpu/BUILD
@@ -1189,6 +1189,7 @@ cc_library(
         "//xla/service:custom_call_status_public_headers",
         "@com_google_absl//absl/algorithm:container",
         "@com_google_absl//absl/base:core_headers",
+        "@com_google_absl//absl/base:dynamic_annotations",
         "@com_google_absl//absl/log",
         "@com_google_absl//absl/status",
         "@com_google_absl//absl/strings",
diff --git a/third_party/xla/xla/service/cpu/runtime_handle_ffi_call.cc b/third_party/xla/xla/service/cpu/runtime_handle_ffi_call.cc
index d0e26cc5cfc232..7a7963593f592a 100644
--- a/third_party/xla/xla/service/cpu/runtime_handle_ffi_call.cc
+++ b/third_party/xla/xla/service/cpu/runtime_handle_ffi_call.cc
@@ -22,8 +22,10 @@ limitations under the License.
 
 #include "absl/algorithm/container.h"
 #include "absl/base/attributes.h"
+#include "absl/base/dynamic_annotations.h"
 #include "absl/status/status.h"
 #include "absl/strings/str_cat.h"
+#include "absl/strings/string_view.h"
 #include "absl/types/span.h"
 #include "llvm/ADT/TypeSwitch.h"
 #include "mlir/AsmParser/AsmParser.h"  // from @llvm-project
@@ -134,7 +136,12 @@ absl::StatusOr<AttributesMap> BuildAttributesMap(mlir::DictionaryAttr dict) {
 }
 
 absl::Span<const int64_t> DecodeDims(int64_t* encoded_dims_data) {
+  // Annotate memory coming from jit compiled function as initialized to
+  // suppress false positives from msan sanitizer.
+  ABSL_ANNOTATE_MEMORY_IS_INITIALIZED(encoded_dims_data, sizeof(int64_t));
   auto dims_count = encoded_dims_data[0];
+  ABSL_ANNOTATE_MEMORY_IS_INITIALIZED(encoded_dims_data,
+                                      dims_count * sizeof(int64_t));
   auto dims_begin = encoded_dims_data + 1;
   return absl::MakeSpan(dims_begin, dims_begin + dims_count);
 }
@@ -263,6 +270,13 @@ ABSL_ATTRIBUTE_NO_SANITIZE_MEMORY void __xla_cpu_runtime_HandleFfiCall(
     outputs = reinterpret_cast<void**>(output);
   }
 
+  // Annotate memory coming from jit compiled function as initialized to
+  // suppress false positives from msan sanitizer.
+  ABSL_ANNOTATE_MEMORY_IS_INITIALIZED(result_types,
+                                      result_count * sizeof(int32_t));
+  ABSL_ANNOTATE_MEMORY_IS_INITIALIZED(operand_types,
+                                      operand_count * sizeof(int32_t));
+
   absl::Status status = BuildAndCallFfi(
       target_name, backend_config, absl::MakeSpan(outputs, result_count),
       absl::MakeSpan(inputs, operand_count),

From a24156129530420b7e97d873f8d80be6158ab6d8 Mon Sep 17 00:00:00 2001
From: Kuy Mainwaring <kuym@google.com>
Date: Mon, 13 May 2024 13:42:45 -0700
Subject: [PATCH 057/478] [XLA:GPU] Clang-tidy fixes for
 xla/service/gpu/runtime/address_computation_thunk.h

PiperOrigin-RevId: 633318572
---
 .../xla/xla/service/gpu/runtime/address_computation_thunk.h      | 1 +
 1 file changed, 1 insertion(+)

diff --git a/third_party/xla/xla/service/gpu/runtime/address_computation_thunk.h b/third_party/xla/xla/service/gpu/runtime/address_computation_thunk.h
index 8d36751b9d830d..dffe9c5da1d8f6 100644
--- a/third_party/xla/xla/service/gpu/runtime/address_computation_thunk.h
+++ b/third_party/xla/xla/service/gpu/runtime/address_computation_thunk.h
@@ -18,6 +18,7 @@ limitations under the License.
 
 #include <cstdint>
 #include <memory>
+#include <optional>
 #include <vector>
 
 #include "absl/base/thread_annotations.h"

From dbcbb3eef6278cee0e614926f64fbb795308198f Mon Sep 17 00:00:00 2001
From: Ionel Gog <icgog@google.com>
Date: Mon, 13 May 2024 13:45:54 -0700
Subject: [PATCH 058/478] Allow ifrt.Call of functions that consists only of
 ifrt.Reshard ops.

PiperOrigin-RevId: 633319317
---
 third_party/xla/xla/python/ifrt/ir/constants.h        | 6 ++++++
 third_party/xla/xla/python/ifrt/ir/ifrt_interfaces.cc | 9 ++++++---
 2 files changed, 12 insertions(+), 3 deletions(-)

diff --git a/third_party/xla/xla/python/ifrt/ir/constants.h b/third_party/xla/xla/python/ifrt/ir/constants.h
index cd1bc06bf85ee0..27e9d11fb6a1cf 100644
--- a/third_party/xla/xla/python/ifrt/ir/constants.h
+++ b/third_party/xla/xla/python/ifrt/ir/constants.h
@@ -25,6 +25,12 @@ namespace ifrt {
 // apart from atom program FuncOps (callee of `ifrt.Call`).
 inline constexpr llvm::StringLiteral kIfrtFunctionAttrName = "ifrt.function";
 
+// Name of UnitAttr on FuncOp to indicate it's an IFRT IR function that
+// only reshards arrays. While functions with kIfrtFunctionAttrName attribute
+// cannot be `ifrt.Call`ed, kIfrtReshardFunctionAttrName can be called.
+inline constexpr llvm::StringLiteral kIfrtReshardFunctionAttrName =
+    "ifrt.reshard_function";
+
 // Name of UnitAttr on arguments of FuncOp to indicate a donated input.
 // Must be used in a FuncOp with `ifrt.function` attr.
 inline constexpr llvm::StringLiteral kIfrtDonatedArgAttrName = "ifrt.donated";
diff --git a/third_party/xla/xla/python/ifrt/ir/ifrt_interfaces.cc b/third_party/xla/xla/python/ifrt/ir/ifrt_interfaces.cc
index a079f75fbdd597..db8844da11e36b 100644
--- a/third_party/xla/xla/python/ifrt/ir/ifrt_interfaces.cc
+++ b/third_party/xla/xla/python/ifrt/ir/ifrt_interfaces.cc
@@ -37,9 +37,12 @@ namespace impl {
 LogicalResult verifyNestedInIfrtFunc(Operation* op) {
   auto func_op = op->getParentOfType<func::FuncOp>();
   if (func_op != nullptr &&
-      !func_op->hasAttr(::xla::ifrt::kIfrtFunctionAttrName)) {
-    return op->emitOpError() << "must be in a FuncOp with attr `"
-                             << ::xla::ifrt::kIfrtFunctionAttrName << "`";
+      !func_op->hasAttr(::xla::ifrt::kIfrtFunctionAttrName) &&
+      !func_op->hasAttr(::xla::ifrt::kIfrtReshardFunctionAttrName)) {
+    return op->emitOpError()
+           << "must be in a FuncOp with attr `"
+           << ::xla::ifrt::kIfrtFunctionAttrName << "` or atttr `"
+           << ::xla::ifrt::kIfrtReshardFunctionAttrName << "`";
   }
   return success();
 }

From 11beb9d544e7b18449a958e4ba877346d697db21 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Mon, 13 May 2024 13:49:17 -0700
Subject: [PATCH 059/478] Update ops-related pbtxt files.

PiperOrigin-RevId: 633320442
---
 tensorflow/core/ops/ops.pbtxt | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorflow/core/ops/ops.pbtxt b/tensorflow/core/ops/ops.pbtxt
index 53c362ed0ab819..c8c55e08ca0d32 100644
--- a/tensorflow/core/ops/ops.pbtxt
+++ b/tensorflow/core/ops/ops.pbtxt
@@ -1,4 +1,4 @@
-go/debugstr    
+go/debugonly    
 op {
   name: "Abort"
   attr {

From 696e6819516e973ee2e0ab7be439d1a747b55d32 Mon Sep 17 00:00:00 2001
From: Blake Hechtman <blakehechtman@google.com>
Date: Mon, 13 May 2024 13:55:14 -0700
Subject: [PATCH 060/478] [XLA:SPACE_TO_BATCH] correctly propagate on dot

PiperOrigin-RevId: 633322160
---
 .../xla/service/space_to_batch_converter.cc   | 101 +++++++++++++++++-
 .../service/space_to_batch_converter_test.cc  |  13 +--
 2 files changed, 104 insertions(+), 10 deletions(-)

diff --git a/third_party/xla/xla/service/space_to_batch_converter.cc b/third_party/xla/xla/service/space_to_batch_converter.cc
index 88cbb5ecad505e..a6149b197805ec 100644
--- a/third_party/xla/xla/service/space_to_batch_converter.cc
+++ b/third_party/xla/xla/service/space_to_batch_converter.cc
@@ -16,6 +16,7 @@ limitations under the License.
 
 #include <algorithm>
 #include <cstddef>
+#include <cstdint>
 #include <iterator>
 #include <map>
 #include <memory>
@@ -112,6 +113,8 @@ class ConvolutionVisitor {
   // This function checks if the HLO instruction supports propagation.
   bool SupportedOpForPropagation(HloInstruction* consumer,
                                  HloInstruction* producer);
+  bool SupportedDotForPropagation(HloInstruction* consumer,
+                                  HloInstruction* producer);
 
   // Method that checks validity of Broadcast propagation.
   bool IsBroadcastPropagatable(HloInstruction* broadcast,
@@ -1561,13 +1564,55 @@ bool ConvolutionVisitor::IsOpcodeNonPropagatable(HloInstruction* consumer) {
   switch (consumer->opcode()) {
     case HloOpcode::kCustomCall:
       return true;
-    case HloOpcode::kDot:
-      return !ctrl_.enable_propagations_on_dots;
     default:
       return false;
   }
 }
 
+bool ConvolutionVisitor::SupportedDotForPropagation(HloInstruction* consumer,
+                                                    HloInstruction* producer) {
+  if (consumer->opcode() != HloOpcode::kDot) {
+    return false;
+  }
+  auto operand = consumer->mutable_operand(0);
+  if (operand != producer || !instr_to_dim_map_.contains(operand)) {
+    return false;
+  }
+  const auto& dnums = consumer->dot_dimension_numbers();
+  const auto& contracting_dims = dnums.lhs_contracting_dimensions();
+  const auto& batch_dims = dnums.lhs_batch_dimensions();
+  auto result = instr_to_dim_map_[operand];
+  const int64_t old_batch_dim = result[DimMapper(SpaceToBatchDimMap::kBatch)];
+  const int64_t old_space_dim = result[DimMapper(SpaceToBatchDimMap::kSpace0)];
+  const int64_t old_feature_dim =
+      result[DimMapper(SpaceToBatchDimMap::kFeature)];
+  // No feature dimension in output
+  if (consumer->operand(1)->shape().rank() ==
+      batch_dims.size() + contracting_dims.size()) {
+    return false;
+  }
+  // If the convolution space or batch dimension are contracting or batch on
+  // the dot, do not propagate.
+  bool found = false;
+  for (auto dim : batch_dims) {
+    if (dim == old_batch_dim || dim == old_space_dim) {
+      return false;
+    }
+    if (dim == old_feature_dim) {
+      found = true;
+    }
+  }
+  if (!found) {
+    return false;
+  }
+  for (auto dim : contracting_dims) {
+    if (dim == old_batch_dim || dim == old_space_dim) {
+      return false;
+    }
+  }
+  return true;
+}
+
 bool ConvolutionVisitor::SupportedOpForPropagation(HloInstruction* consumer,
                                                    HloInstruction* producer) {
   if (IsOpcodeNonPropagatable(consumer)) {
@@ -1682,6 +1727,10 @@ bool ConvolutionVisitor::SupportedOpForPropagation(HloInstruction* consumer,
     return true;
   }
 
+  if (SupportedDotForPropagation(consumer, producer)) {
+    return true;
+  }
+
   if (consumer->opcode() == HloOpcode::kReduce) {
     // Support only the trivial case where both batch and split spatial dim are
     // being reduced
@@ -1964,6 +2013,50 @@ absl::StatusOr<bool> ConvolutionVisitor::Propagate(HloInstruction* consumer,
     return true;
   }
 
+  if (consumer->opcode() == HloOpcode::kDot) {
+    auto dim_map_val = instr_to_dim_map_[producer];
+    const int64_t old_batch_dim =
+        dim_map_val[DimMapper(SpaceToBatchDimMap::kBatch)];
+    const int64_t old_space_dim =
+        dim_map_val[DimMapper(SpaceToBatchDimMap::kSpace0)];
+    int64_t new_batch_dim = -1;
+    int64_t new_space_dim = -1;
+    int64_t outer = 0;
+    for (int64_t i = 0; i < producer->shape().rank(); ++i) {
+      if (absl::c_linear_search(
+              consumer->dot_dimension_numbers().lhs_batch_dimensions(), i) ||
+          absl::c_linear_search(
+              consumer->dot_dimension_numbers().lhs_contracting_dimensions(),
+              i)) {
+        continue;
+      }
+      if (i == old_batch_dim) {
+        new_batch_dim =
+            outer +
+            consumer->dot_dimension_numbers().lhs_batch_dimensions_size();
+      }
+      if (i == old_space_dim) {
+        new_batch_dim =
+            outer +
+            consumer->dot_dimension_numbers().lhs_batch_dimensions_size();
+      }
+      ++outer;
+    }
+    std::vector<int64_t> dim_map(NumMappedDims());
+    dim_map[DimMapper(SpaceToBatchDimMap::kBatch)] = new_batch_dim;
+    dim_map[DimMapper(SpaceToBatchDimMap::kSpace0)] = new_space_dim;
+    dim_map[DimMapper(SpaceToBatchDimMap::kFeature)] =
+        consumer->shape().rank() - 1;
+    instr_to_dim_map_[consumer] = dim_map;
+    auto new_consumer = computation->AddInstruction(consumer->Clone());
+    new_consumer->mutable_shape()->mutable_dimensions()[new_batch_dim] =
+        producer->shape().dimensions(old_batch_dim);
+    new_consumer->mutable_shape()->mutable_dimensions()[new_space_dim] =
+        producer->shape().dimensions(old_space_dim);
+    old_to_new_instrs_[consumer] = new_consumer;
+    return true;
+  }
+
   // TODO(b/189500737) : Consider a common way of propagation for
   // slice/pad/reduce-window.
   if (consumer->opcode() == HloOpcode::kPad) {
@@ -3619,7 +3712,8 @@ ConvolutionVisitor::DoesConvolutionFeedReduceWindowOrSelectAndScatter(
     // Stop the search if these ops are encountered.
     if (user->opcode() == HloOpcode::kConvolution ||
         user->opcode() == HloOpcode::kPad ||
-        user->opcode() == HloOpcode::kTranspose) {
+        user->opcode() == HloOpcode::kTranspose ||
+        user->opcode() == HloOpcode::kDot) {
       continue;
     }
     auto ret =
@@ -3986,7 +4080,6 @@ Status ConvolutionVisitor::PerformSpaceToBatchOnConvolution(
   }
   TF_CHECK_OK(PropagateOnUsers(original_conv));
 
-
   return OkStatus();
 }
 
diff --git a/third_party/xla/xla/service/space_to_batch_converter_test.cc b/third_party/xla/xla/service/space_to_batch_converter_test.cc
index cf8dce26211b83..e2ed3314bc4f6f 100644
--- a/third_party/xla/xla/service/space_to_batch_converter_test.cc
+++ b/third_party/xla/xla/service/space_to_batch_converter_test.cc
@@ -247,7 +247,7 @@ ENTRY computation {
   EXPECT_GT(previous_reshape->operand(0)->shape().dimensions(batch_dim), 4);
 }
 
-TEST_F(SpaceToBatchConverterTest, NoPropagateThroughDot) {
+TEST_F(SpaceToBatchConverterTest, PropagateThroughDot) {
   std::string hlo_string = R"(
   HloModule module
 
@@ -256,9 +256,10 @@ TEST_F(SpaceToBatchConverterTest, NoPropagateThroughDot) {
     %p1 = bf16[3,3,32,32] parameter(1)
     %convolution = bf16[1,256,256,32] convolution(%p0, %p1), window={size=3x3},
     dim_labels=b01f_01io->b01f
-    %p2 = bf16[1,256,256,32] parameter(2)
-    ROOT %dot.5010 = bf16[1,256,32,32] dot(%convolution, %p2), lhs_batch_dims={0,1},
-    lhs_contracting_dims={2}, rhs_batch_dims={0,2}, rhs_contracting_dims={1}
+    %p2 = bf16[32,32] parameter(2)
+    ROOT %dot.5010 = bf16[1,256,256,32] dot(%convolution, %p2),
+      lhs_contracting_dims={3},
+      rhs_contracting_dims={0}
   }
 
   )";
@@ -267,8 +268,8 @@ TEST_F(SpaceToBatchConverterTest, NoPropagateThroughDot) {
 
   SpaceToBatchConverter converter(
       SpaceToBatchController{true, true, true, true, 8});
-  // Test that we do not start space-to-batch on conv->dot chains
-  ASSERT_FALSE(converter.Run(module.get()).value());
+  // Test that we do not start space-to-batch on conv->dot chains.
+  ASSERT_TRUE(converter.Run(module.get()).value());
 }
 
 }  // namespace

From 4beabf4430b0950890f31298012edd8aaaba0510 Mon Sep 17 00:00:00 2001
From: Mason Chang <masonchang@google.com>
Date: Mon, 13 May 2024 13:55:49 -0700
Subject: [PATCH 061/478] Update comment to reflect placement

PiperOrigin-RevId: 633322446
---
 tensorflow/core/common_runtime/placer.cc | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/tensorflow/core/common_runtime/placer.cc b/tensorflow/core/common_runtime/placer.cc
index 5908b26a898690..2d03de9f2434bc 100644
--- a/tensorflow/core/common_runtime/placer.cc
+++ b/tensorflow/core/common_runtime/placer.cc
@@ -291,8 +291,8 @@ Status Placer::Run(const GraphOptimizationPassOptions& options) {
     // to perform good placement we can add an interface for this.
     int assigned_device = -1;
 
-    // Heuristic B: If the node only operates on metadata, not data,
-    // then it is desirable to place that metadata node with its
+    // Heuristic B: If the node only operates on metadata (not data) or is
+    // an identity node, then it is desirable to place that node with its
     // input.
     if (IsMetadata(node) || MatchIdentityOperation(node)) {
       // Make sure that the input device type is in the list of supported

From 3bd8e82f856430f30276020f2342c1e797534db4 Mon Sep 17 00:00:00 2001
From: Sania Nagpal <sanianagpal@google.com>
Date: Mon, 13 May 2024 14:13:17 -0700
Subject: [PATCH 062/478] No public description

PiperOrigin-RevId: 633328239
---
 tensorflow/core/tfrt/runtime/BUILD          |  3 +
 tensorflow/core/tfrt/runtime/stream_test.cc | 87 ++++++++++++++++++---
 2 files changed, 77 insertions(+), 13 deletions(-)

diff --git a/tensorflow/core/tfrt/runtime/BUILD b/tensorflow/core/tfrt/runtime/BUILD
index 95712a19629fd8..c4cf2dda41b87e 100644
--- a/tensorflow/core/tfrt/runtime/BUILD
+++ b/tensorflow/core/tfrt/runtime/BUILD
@@ -153,6 +153,7 @@ tf_cc_shared_test(
     srcs = ["stream_test.cc"],
     tags = ["no_oss"],
     deps = [
+        ":step_id",
         ":stream",
         "//tensorflow/core/framework:tensor",
         "//tensorflow/core/framework:tensor_testutil",
@@ -162,6 +163,8 @@ tf_cc_shared_test(
         "@com_google_absl//absl/log",
         "@com_google_absl//absl/log:check",
         "@com_google_absl//absl/memory",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/strings:string_view",
         "@com_google_absl//absl/time",
         "@com_google_googletest//:gtest_main",
         "@local_tsl//tsl/platform:env",
diff --git a/tensorflow/core/tfrt/runtime/stream_test.cc b/tensorflow/core/tfrt/runtime/stream_test.cc
index cac9113053bfab..bcb8a14a553675 100644
--- a/tensorflow/core/tfrt/runtime/stream_test.cc
+++ b/tensorflow/core/tfrt/runtime/stream_test.cc
@@ -15,6 +15,7 @@ limitations under the License.
 #include "tensorflow/core/tfrt/runtime/stream.h"
 
 #include <cstdint>
+#include <functional>
 #include <memory>
 #include <string>
 #include <utility>
@@ -26,10 +27,13 @@ limitations under the License.
 #include "absl/log/check.h"
 #include "absl/log/log.h"
 #include "absl/memory/memory.h"
+#include "absl/status/status.h"
+#include "absl/strings/string_view.h"
 #include "absl/time/clock.h"
 #include "absl/time/time.h"
 #include "tensorflow/core/framework/tensor.h"
 #include "tensorflow/core/framework/tensor_testutil.h"
+#include "tensorflow/core/tfrt/runtime/step_id.h"
 #include "tensorflow/core/tfrt/saved_model/saved_model_testutil.h"
 #include "tensorflow/core/tfrt/utils/thread_pool.h"
 #include "tsl/platform/env.h"
@@ -44,27 +48,17 @@ using ::testing::AnyOf;
 using ::testing::ElementsAreArray;
 using ::testing::Pair;
 using ::testing::UnorderedElementsAre;
-
-class TestStreamInterface : public StreamControllerInterface {
- public:
-  TestStreamInterface() : StreamControllerInterface("test_address") {}
-};
-
-const bool kUnused = []() {
-  GetGlobalStreamInterfaceFactory().RegisterController(
-      []() { return std::make_unique<TestStreamInterface>(); });
-  return true;
-}();
+using ::testing::status::StatusIs;
 
 TEST(StreamTest, Simple) {
   StreamCallbackId callback_id(1234);
   StepId step_id(5678);
 
   std::vector<absl::flat_hash_map<std::string, tensorflow::Tensor>> outputs;
-
+  ScopedStreamCallback scoped_stream_callback;
   {
     TF_ASSERT_OK_AND_ASSIGN(
-        auto scoped_stream_callback,
+        scoped_stream_callback,
         GetGlobalStreamCallbackRegistry().Register(
             "test_model", callback_id, step_id,
             [&](absl::flat_hash_map<std::string, tensorflow::Tensor> arg) {
@@ -91,6 +85,16 @@ TEST(StreamTest, Simple) {
               ElementsAreArray({200}));
   EXPECT_THAT(GetTfTensorData<int32_t>(outputs[1]["c"]),
               ElementsAreArray({300}));
+
+  ScopedStreamCallback scoped_stream_callback_copy;
+  scoped_stream_callback_copy = std::move(scoped_stream_callback);
+
+  auto status = GetGlobalStreamCallbackRegistry().Register(
+      "test_model", callback_id, step_id,
+      [&](absl::flat_hash_map<std::string, tensorflow::Tensor> arg) {
+        outputs.push_back(std::move(arg));
+      });
+  EXPECT_THAT(status, StatusIs(absl::StatusCode::kAlreadyExists));
 }
 
 TEST(StreamTest, MultipleWriters) {
@@ -142,6 +146,63 @@ TEST(StreamTest, MultipleWriters) {
   }
 }
 
+class TestStreamControllerInterface : public StreamControllerInterface {
+ public:
+  TestStreamControllerInterface()
+      : StreamControllerInterface("test_controller_address") {}
+};
+
+TEST(StreamControllerInterface, Initialize) {
+  GetGlobalStreamInterfaceFactory().RegisterController(
+      []() { return std::make_unique<TestStreamControllerInterface>(); });
+  TF_ASSERT_OK_AND_ASSIGN(
+      auto controller_interface,
+      GetGlobalStreamInterfaceFactory().CreateControllerStreamInterface());
+  EXPECT_EQ(controller_interface->controller_address(),
+            "test_controller_address");
+}
+
+class TestStreamWorkerInterface : public StreamWorkerInterface {
+ public:
+  explicit TestStreamWorkerInterface(std::string worker_address)
+      : StreamWorkerInterface(worker_address) {}
+  absl::Status InvokeStreamCallback(
+      const StreamCallbackId& callback_id,
+      const std::vector<std::string>& names,
+      const std::vector<std::pair<int64_t, std::vector<tensorflow::Tensor>>>&
+          responses) override {
+    return absl::OkStatus();
+  }
+};
+
+TEST(StreamWorkerInterface, Initialize) {
+  GetGlobalStreamInterfaceFactory().RegisterWorker(
+      [](absl::string_view address)
+          -> absl::StatusOr<std::unique_ptr<TestStreamWorkerInterface>> {
+        return std::make_unique<TestStreamWorkerInterface>(
+            "test_worker_address");
+      });
+  TF_ASSERT_OK_AND_ASSIGN(
+      auto worker_interface,
+      GetGlobalStreamInterfaceFactory().CreateWorkerStreamInterface()(
+          "test_worker_address"));
+  EXPECT_EQ(worker_interface->controller_address(), "test_worker_address");
+}
+
+TEST(StepId, Generate) {
+  StepId step_id(1234);
+  EXPECT_EQ(step_id.id, 1234);
+  StepIdGenerator step_id_generator;
+  EXPECT_EQ(step_id_generator.GetNextStepId(), StepId(1));
+  EXPECT_EQ(step_id_generator.GetNextStepId(), StepId(2));
+  EXPECT_EQ(step_id_generator.GetNextStepId(), StepId(3));
+}
+
+TEST(StepId, GlobalInitial) {
+  EXPECT_EQ(GetGlobalInitialStepId(), 0);
+  TEST_ScopedInitialStepId test_id(127);
+  EXPECT_EQ(GetGlobalInitialStepId(), 127);
+}
 }  // namespace
 }  // namespace tfrt_stub
 }  // namespace tensorflow

From cad3f2972e55027886d2e121e4d27a1b6ba60f97 Mon Sep 17 00:00:00 2001
From: Yifan Jiang <yifanjiang@google.com>
Date: Mon, 13 May 2024 14:18:03 -0700
Subject: [PATCH 063/478] Fix a potential memory corruption in PJRT GPU client.
 Make sure the staging buffer outlives the memcpys to and from the staging
 buffer.

PiperOrigin-RevId: 633329702
---
 third_party/xla/xla/pjrt/gpu/se_gpu_pjrt_client.cc | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/third_party/xla/xla/pjrt/gpu/se_gpu_pjrt_client.cc b/third_party/xla/xla/pjrt/gpu/se_gpu_pjrt_client.cc
index 4e818b7f6fddd5..bfa11cb315c8c5 100644
--- a/third_party/xla/xla/pjrt/gpu/se_gpu_pjrt_client.cc
+++ b/third_party/xla/xla/pjrt/gpu/se_gpu_pjrt_client.cc
@@ -393,8 +393,8 @@ class AsyncHostToDeviceTransferManager
                                                                 event.value());
 
     auto cleanup = [this, buffer_index, event = std::move(event).value(),
-                    stream, is_last_transfer,
-                    on_done = std::move(on_done)]() mutable {
+                    stream, is_last_transfer, on_done = std::move(on_done),
+                    staging_buffer = std::move(staging_buffer)]() mutable {
       CleanUp(buffer_index, std::move(event), stream, is_last_transfer,
               std::move(on_done));
     };

From b3a856c2c25fe0e7d6df49b058b1a1bf31de820d Mon Sep 17 00:00:00 2001
From: Junwhan Ahn <junwhan@google.com>
Date: Mon, 13 May 2024 14:36:43 -0700
Subject: [PATCH 064/478] Enable JAX memory tests for GPUs and CPUs

PjRt GPU and CPU has recently gotten memory space support with just one memory space per device, so enabling relevant JAX memory tests. Most tests cannot be enabled yet because they rely on `unpinned_host`, so only enabling `ShardingMemoriesTest` for now.

PiperOrigin-RevId: 633335638
---
 third_party/xla/xla/python/xla_client.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/third_party/xla/xla/python/xla_client.py b/third_party/xla/xla/python/xla_client.py
index 37c462efc3e0f6..af0590c12c0f8a 100644
--- a/third_party/xla/xla/python/xla_client.py
+++ b/third_party/xla/xla/python/xla_client.py
@@ -49,7 +49,7 @@
 
 # Just an internal arbitrary increasing number to help with backward-compatible
 # changes. In JAX, reference this via jax._src.lib.xla_extension_version.
-_version = 264
+_version = 265
 
 # Version number for MLIR:Python components.
 mlir_api_version = 56

From 9a97b727bcb6536663878404f06bb71e7cc4113b Mon Sep 17 00:00:00 2001
From: Swachhand Lokhande <swachhand@google.com>
Date: Mon, 13 May 2024 14:41:53 -0700
Subject: [PATCH 065/478] Add TFExecutorGraphPruningPass to TF dialect to
 executor export pipeline

This is to prune out unused ops/nodes in the graph. This doesn't affect correctness but in some cases will drastically reduce memory/network usage and improve performance depending on the number and type of dead nodes that are pruned.

PiperOrigin-RevId: 633337115
---
 tensorflow/compiler/mlir/tf2xla/api/v2/BUILD  |  4 +-
 .../api/v2/testdata/func_with_dead_ops.mlir   | 62 +++++++++++++++++++
 .../tf2xla/api/v2/tf_dialect_to_executor.cc   |  2 +
 .../api/v2/tf_dialect_to_executor_test.cc     | 32 +++++++++-
 4 files changed, 97 insertions(+), 3 deletions(-)
 create mode 100644 tensorflow/compiler/mlir/tf2xla/api/v2/testdata/func_with_dead_ops.mlir

diff --git a/tensorflow/compiler/mlir/tf2xla/api/v2/BUILD b/tensorflow/compiler/mlir/tf2xla/api/v2/BUILD
index 545203ad20ea23..709a63bea84ebe 100644
--- a/tensorflow/compiler/mlir/tf2xla/api/v2/BUILD
+++ b/tensorflow/compiler/mlir/tf2xla/api/v2/BUILD
@@ -210,6 +210,7 @@ tf_cc_test(
     srcs = ["tf_dialect_to_executor_test.cc"],
     data = [
         "testdata/empty_func.mlir",
+        "testdata/func_with_dead_ops.mlir",
         "testdata/invalid_executor.mlir",
     ],
     deps = [
@@ -220,10 +221,9 @@ tf_cc_test(
         "@com_google_absl//absl/status",
         "@com_google_absl//absl/strings",
         "@com_google_googletest//:gtest_main",
+        "@llvm-project//llvm:Support",
         "@llvm-project//mlir:IR",
         "@llvm-project//mlir:Parser",
         "@local_tsl//tsl/lib/core:status_test_util",
-        "@local_tsl//tsl/lib/monitoring:test_utils",
-        "@local_tsl//tsl/platform:status",
     ],
 )
diff --git a/tensorflow/compiler/mlir/tf2xla/api/v2/testdata/func_with_dead_ops.mlir b/tensorflow/compiler/mlir/tf2xla/api/v2/testdata/func_with_dead_ops.mlir
new file mode 100644
index 00000000000000..f8dd51f4e12d3c
--- /dev/null
+++ b/tensorflow/compiler/mlir/tf2xla/api/v2/testdata/func_with_dead_ops.mlir
@@ -0,0 +1,62 @@
+module attributes {tf.devices = {"/job:tpu_host_worker/replica:0/task:0/device:CPU:0", "/job:tpu_host_worker/replica:0/task:0/device:TPU:0", "/job:tpu_host_worker/replica:0/task:0/device:TPU:1", "/job:tpu_host_worker/replica:0/task:0/device:TPU_SYSTEM:0", "/job:tpu_host_worker/replica:0/task:1/device:CPU:0", "/job:tpu_host_worker/replica:0/task:1/device:TPU:0", "/job:tpu_host_worker/replica:0/task:1/device:TPU:1", "/job:tpu_host_worker/replica:0/task:1/device:TPU_SYSTEM:0", "/job:tpu_host_worker/replica:0/task:2/device:CPU:0", "/job:tpu_host_worker/replica:0/task:2/device:TPU:0", "/job:tpu_host_worker/replica:0/task:2/device:TPU:1", "/job:tpu_host_worker/replica:0/task:2/device:TPU_SYSTEM:0", "/job:tpu_host_worker/replica:0/task:3/device:CPU:0", "/job:tpu_host_worker/replica:0/task:3/device:TPU:0", "/job:tpu_host_worker/replica:0/task:3/device:TPU:1", "/job:tpu_host_worker/replica:0/task:3/device:TPU_SYSTEM:0"}, tf.versions = {bad_consumers = [], min_consumer = 0 : i32, producer = 1847 : i32}} {
+  func.func @main(%arg0: tensor<i32> {tf._user_specified_name = "steps", tf.device = "/job:tpu_host_worker/replica:0/task:0/device:CPU:0"}, %arg1: tensor<*x!tf_type.resource<tensor<i64>>> {tf._user_specified_name = "899", tf.device = "/job:tpu_host_worker/replica:0/task:0/device:CPU:0"}, %arg2: tensor<*x!tf_type.resource<tensor<i64>>> {tf._user_specified_name = "901", tf.device = "/job:tpu_host_worker/replica:0/task:0/device:CPU:0"}, %arg3: tensor<*x!tf_type.resource<tensor<128x1024xf32>>> {tf._user_specified_name = "903", tf.device = "/job:tpu_host_worker/replica:0/task:0/device:CPU:0"}, %arg4: tensor<*x!tf_type.resource<tensor<1024xf32>>> {tf._user_specified_name = "905", tf.device = "/job:tpu_host_worker/replica:0/task:0/device:CPU:0"}, %arg5: tensor<*x!tf_type.resource<tensor<1024x1xf32>>> {tf._user_specified_name = "907", tf.device = "/job:tpu_host_worker/replica:0/task:0/device:CPU:0"}, %arg6: tensor<*x!tf_type.resource<tensor<i64>>> {tf._user_specified_name = "909", tf.device = "/job:tpu_host_worker/replica:0/task:0/device:CPU:0"}, %arg7: tensor<*x!tf_type.resource<tensor<25001x64xf32>>> {tf._user_specified_name = "911", tf.device = "/job:tpu_host_worker/replica:0/task:0/device:CPU:0"}, %arg8: tensor<*x!tf_type.resource<tensor<25001x64xf32>>> {tf._user_specified_name = "913", tf.device = "/job:tpu_host_worker/replica:0/task:1/device:CPU:0"}, %arg9: tensor<*x!tf_type.resource<tensor<25001x64xf32>>> {tf._user_specified_name = "915", tf.device = "/job:tpu_host_worker/replica:0/task:2/device:CPU:0"}, %arg10: tensor<*x!tf_type.resource<tensor<25001x64xf32>>> {tf._user_specified_name = "917", tf.device = "/job:tpu_host_worker/replica:0/task:3/device:CPU:0"}, %arg11: tensor<*x!tf_type.resource<tensor<25001x32xf32>>> {tf._user_specified_name = "919", tf.device = "/job:tpu_host_worker/replica:0/task:0/device:CPU:0"}, %arg12: tensor<*x!tf_type.resource<tensor<25001x32xf32>>> {tf._user_specified_name = "921", tf.device = "/job:tpu_host_worker/replica:0/task:1/device:CPU:0"}, %arg13: tensor<*x!tf_type.resource<tensor<25001x32xf32>>> {tf._user_specified_name = "923", tf.device = "/job:tpu_host_worker/replica:0/task:2/device:CPU:0"}, %arg14: tensor<*x!tf_type.resource<tensor<25001x32xf32>>> {tf._user_specified_name = "925", tf.device = "/job:tpu_host_worker/replica:0/task:3/device:CPU:0"}, %arg15: tensor<*x!tf_type.resource<tensor<6x32xf32>>> {tf._user_specified_name = "927", tf.device = "/job:tpu_host_worker/replica:0/task:0/device:CPU:0"}, %arg16: tensor<*x!tf_type.resource<tensor<6x32xf32>>> {tf._user_specified_name = "929", tf.device = "/job:tpu_host_worker/replica:0/task:1/device:CPU:0"}, %arg17: tensor<*x!tf_type.resource<tensor<6x32xf32>>> {tf._user_specified_name = "931", tf.device = "/job:tpu_host_worker/replica:0/task:2/device:CPU:0"}, %arg18: tensor<*x!tf_type.resource<tensor<6x32xf32>>> {tf._user_specified_name = "933", tf.device = "/job:tpu_host_worker/replica:0/task:3/device:CPU:0"}, %arg19: tensor<*x!tf_type.resource<tensor<128x1024xf32>>> {tf._user_specified_name = "935", tf.device = "/job:tpu_host_worker/replica:0/task:0/device:CPU:0"}, %arg20: tensor<*x!tf_type.resource<tensor<1024xf32>>> {tf._user_specified_name = "937", tf.device = "/job:tpu_host_worker/replica:0/task:0/device:CPU:0"}, %arg21: tensor<*x!tf_type.resource<tensor<1024x1xf32>>> {tf._user_specified_name = "939", tf.device = "/job:tpu_host_worker/replica:0/task:0/device:CPU:0"}) -> tensor<i64> attributes {allow_soft_placement = false, tf.entry_function = {control_outputs = "", inputs = "steps,unknown,unknown_0,unknown_1,unknown_2,unknown_3,unknown_4,unknown_5,unknown_6,unknown_7,unknown_8,unknown_9,unknown_10,unknown_11,unknown_12,unknown_13,unknown_14,unknown_15,unknown_16,unknown_17,unknown_18,unknown_19", outputs = "statefulpartitionedcall_RetVal"}} {
+    %0 = "tf.ReadVariableOp"(%arg19) : (tensor<*x!tf_type.resource<tensor<128x1024xf32>>>) -> tensor<128x1024xf32>
+    %1 = "tf.ReadVariableOp"(%arg1) : (tensor<*x!tf_type.resource<tensor<i64>>>) -> tensor<i64>
+    %2 = "tf.ReadVariableOp"(%arg2) : (tensor<*x!tf_type.resource<tensor<i64>>>) -> tensor<i64>
+    %3 = "tf.ReadVariableOp"(%arg4) : (tensor<*x!tf_type.resource<tensor<1024xf32>>>) -> tensor<1024xf32>
+    %4 = "tf.ReadVariableOp"(%arg3) : (tensor<*x!tf_type.resource<tensor<128x1024xf32>>>) -> tensor<128x1024xf32>
+    %5 = "tf.ReadVariableOp"(%arg5) : (tensor<*x!tf_type.resource<tensor<1024x1xf32>>>) -> tensor<1024x1xf32>
+    %6 = "tf.ReadVariableOp"(%arg20) : (tensor<*x!tf_type.resource<tensor<1024xf32>>>) -> tensor<1024xf32>
+    %7 = "tf.ReadVariableOp"(%arg21) : (tensor<*x!tf_type.resource<tensor<1024x1xf32>>>) -> tensor<1024x1xf32>
+    %8 = "tf.ReadVariableOp"(%arg6) : (tensor<*x!tf_type.resource<tensor<i64>>>) -> tensor<i64>
+    %9 = "tf.Const"() <{value = dense<"test"> : tensor<3x!tf_type.string>}> : () -> tensor<3x!tf_type.string>
+    %cst = "tf.Const"() <{value = dense<0> : tensor<i32>}> : () -> tensor<i32>
+    %11:4 = "tf.Split"(%cst, %0) {num_split = 4 : i32} : (tensor<i32>, tensor<128x1024xf32>) -> (tensor<32x1024xf32>, tensor<32x1024xf32>, tensor<32x1024xf32>, tensor<32x1024xf32>)
+    %cst_0 = "tf.Const"() <{value = dense<0> : tensor<i32>}> : () -> tensor<i32>
+    %12:4 = "tf.Split"(%cst_0, %4) {num_split = 4 : i32} : (tensor<i32>, tensor<128x1024xf32>) -> (tensor<32x1024xf32>, tensor<32x1024xf32>, tensor<32x1024xf32>, tensor<32x1024xf32>)
+    %cst_1 = "tf.Const"() <{value = dense<0> : tensor<i32>}> : () -> tensor<i32>
+    %cst_2 = "tf.Const"() <{value = dense<0> : tensor<i32>}> : () -> tensor<i32>
+    %13:20 = tf_device.replicate {devices = {TPU_REPLICATED_CORE_0 = ["/job:tpu_host_worker/replica:0/task:0/device:TPU:0", "/job:tpu_host_worker/replica:0/task:2/device:TPU:0"], TPU_REPLICATED_CORE_1 = ["/job:tpu_host_worker/replica:0/task:0/device:TPU:1", "/job:tpu_host_worker/replica:0/task:2/device:TPU:1"], TPU_REPLICATED_CORE_2 = ["/job:tpu_host_worker/replica:0/task:1/device:TPU:0", "/job:tpu_host_worker/replica:0/task:3/device:TPU:0"], TPU_REPLICATED_CORE_3 = ["/job:tpu_host_worker/replica:0/task:1/device:TPU:1", "/job:tpu_host_worker/replica:0/task:3/device:TPU:1"], TPU_REPLICATED_HOST_0 = ["/job:tpu_host_worker/replica:0/task:0/device:CPU:0", "/job:tpu_host_worker/replica:0/task:2/device:CPU:0"], TPU_REPLICATED_HOST_1 = ["/job:tpu_host_worker/replica:0/task:0/device:CPU:0", "/job:tpu_host_worker/replica:0/task:2/device:CPU:0"], TPU_REPLICATED_HOST_2 = ["/job:tpu_host_worker/replica:0/task:1/device:CPU:0", "/job:tpu_host_worker/replica:0/task:3/device:CPU:0"], TPU_REPLICATED_HOST_3 = ["/job:tpu_host_worker/replica:0/task:1/device:CPU:0", "/job:tpu_host_worker/replica:0/task:3/device:CPU:0"]}, n = 2 : i32} {
+      %16:40 = "tf_device.parallel_execute"() ({
+        %19:10 = "tf_device.launch"() <{device = "TPU_REPLICATED_CORE_0"}> ({
+          %20:10 = "tf.TPUExecute"(%arg0, %11#0, %1, %2, %3, %12#0, %5, %6, %7, %8, %9) : (tensor<i32>, tensor<32x1024xf32>, tensor<i64>, tensor<i64>, tensor<1024xf32>, tensor<32x1024xf32>, tensor<1024x1xf32>, tensor<1024xf32>, tensor<1024x1xf32>, tensor<i64>, tensor<3x!tf_type.string>) -> (tensor<i32>, tensor<32x1024xf32>, tensor<i64>, tensor<i64>, tensor<1024xf32>, tensor<32x1024xf32>, tensor<1024x1xf32>, tensor<1024xf32>, tensor<1024x1xf32>, tensor<i64>)
+          tf_device.return %20#0, %20#1, %20#2, %20#3, %20#4, %20#5, %20#6, %20#7, %20#8, %20#9 : tensor<i32>, tensor<32x1024xf32>, tensor<i64>, tensor<i64>, tensor<1024xf32>, tensor<32x1024xf32>, tensor<1024x1xf32>, tensor<1024xf32>, tensor<1024x1xf32>, tensor<i64>
+        }) : () -> (tensor<i32>, tensor<32x1024xf32>, tensor<i64>, tensor<i64>, tensor<1024xf32>, tensor<32x1024xf32>, tensor<1024x1xf32>, tensor<1024xf32>, tensor<1024x1xf32>, tensor<i64>)
+        tf_device.return %19#0, %19#1, %19#2, %19#3, %19#4, %19#5, %19#6, %19#7, %19#8, %19#9 : tensor<i32>, tensor<32x1024xf32>, tensor<i64>, tensor<i64>, tensor<1024xf32>, tensor<32x1024xf32>, tensor<1024x1xf32>, tensor<1024xf32>, tensor<1024x1xf32>, tensor<i64>
+      }, {
+        %19:10 = "tf_device.launch"() <{device = "TPU_REPLICATED_CORE_1"}> ({
+          %20:10 = "tf.TPUExecute"(%arg0, %11#1, %1, %2, %3, %12#1, %5, %6, %7, %8, %9) : (tensor<i32>, tensor<32x1024xf32>, tensor<i64>, tensor<i64>, tensor<1024xf32>, tensor<32x1024xf32>, tensor<1024x1xf32>, tensor<1024xf32>, tensor<1024x1xf32>, tensor<i64>, tensor<3x!tf_type.string>) -> (tensor<i32>, tensor<32x1024xf32>, tensor<i64>, tensor<i64>, tensor<1024xf32>, tensor<32x1024xf32>, tensor<1024x1xf32>, tensor<1024xf32>, tensor<1024x1xf32>, tensor<i64>)
+          tf_device.return %20#0, %20#1, %20#2, %20#3, %20#4, %20#5, %20#6, %20#7, %20#8, %20#9 : tensor<i32>, tensor<32x1024xf32>, tensor<i64>, tensor<i64>, tensor<1024xf32>, tensor<32x1024xf32>, tensor<1024x1xf32>, tensor<1024xf32>, tensor<1024x1xf32>, tensor<i64>
+        }) : () -> (tensor<i32>, tensor<32x1024xf32>, tensor<i64>, tensor<i64>, tensor<1024xf32>, tensor<32x1024xf32>, tensor<1024x1xf32>, tensor<1024xf32>, tensor<1024x1xf32>, tensor<i64>)
+        tf_device.return %19#0, %19#1, %19#2, %19#3, %19#4, %19#5, %19#6, %19#7, %19#8, %19#9 : tensor<i32>, tensor<32x1024xf32>, tensor<i64>, tensor<i64>, tensor<1024xf32>, tensor<32x1024xf32>, tensor<1024x1xf32>, tensor<1024xf32>, tensor<1024x1xf32>, tensor<i64>
+      }, {
+        %19:10 = "tf_device.launch"() <{device = "TPU_REPLICATED_CORE_2"}> ({
+          %20:10 = "tf.TPUExecute"(%arg0, %11#2, %1, %2, %3, %12#2, %5, %6, %7, %8, %9) : (tensor<i32>, tensor<32x1024xf32>, tensor<i64>, tensor<i64>, tensor<1024xf32>, tensor<32x1024xf32>, tensor<1024x1xf32>, tensor<1024xf32>, tensor<1024x1xf32>, tensor<i64>, tensor<3x!tf_type.string>) -> (tensor<i32>, tensor<32x1024xf32>, tensor<i64>, tensor<i64>, tensor<1024xf32>, tensor<32x1024xf32>, tensor<1024x1xf32>, tensor<1024xf32>, tensor<1024x1xf32>, tensor<i64>)
+          tf_device.return %20#0, %20#1, %20#2, %20#3, %20#4, %20#5, %20#6, %20#7, %20#8, %20#9 : tensor<i32>, tensor<32x1024xf32>, tensor<i64>, tensor<i64>, tensor<1024xf32>, tensor<32x1024xf32>, tensor<1024x1xf32>, tensor<1024xf32>, tensor<1024x1xf32>, tensor<i64>
+        }) : () -> (tensor<i32>, tensor<32x1024xf32>, tensor<i64>, tensor<i64>, tensor<1024xf32>, tensor<32x1024xf32>, tensor<1024x1xf32>, tensor<1024xf32>, tensor<1024x1xf32>, tensor<i64>)
+        tf_device.return %19#0, %19#1, %19#2, %19#3, %19#4, %19#5, %19#6, %19#7, %19#8, %19#9 : tensor<i32>, tensor<32x1024xf32>, tensor<i64>, tensor<i64>, tensor<1024xf32>, tensor<32x1024xf32>, tensor<1024x1xf32>, tensor<1024xf32>, tensor<1024x1xf32>, tensor<i64>
+      }, {
+        %19:10 = "tf_device.launch"() <{device = "TPU_REPLICATED_CORE_3"}> ({
+          %20:10 = "tf.TPUExecute"(%arg0, %11#3, %1, %2, %3, %12#3, %5, %6, %7, %8, %9) : (tensor<i32>, tensor<32x1024xf32>, tensor<i64>, tensor<i64>, tensor<1024xf32>, tensor<32x1024xf32>, tensor<1024x1xf32>, tensor<1024xf32>, tensor<1024x1xf32>, tensor<i64>, tensor<3x!tf_type.string>) -> (tensor<i32>, tensor<32x1024xf32>, tensor<i64>, tensor<i64>, tensor<1024xf32>, tensor<32x1024xf32>, tensor<1024x1xf32>, tensor<1024xf32>, tensor<1024x1xf32>, tensor<i64>)
+          tf_device.return %20#0, %20#1, %20#2, %20#3, %20#4, %20#5, %20#6, %20#7, %20#8, %20#9 : tensor<i32>, tensor<32x1024xf32>, tensor<i64>, tensor<i64>, tensor<1024xf32>, tensor<32x1024xf32>, tensor<1024x1xf32>, tensor<1024xf32>, tensor<1024x1xf32>, tensor<i64>
+        }) : () -> (tensor<i32>, tensor<32x1024xf32>, tensor<i64>, tensor<i64>, tensor<1024xf32>, tensor<32x1024xf32>, tensor<1024x1xf32>, tensor<1024xf32>, tensor<1024x1xf32>, tensor<i64>)
+        tf_device.return %19#0, %19#1, %19#2, %19#3, %19#4, %19#5, %19#6, %19#7, %19#8, %19#9 : tensor<i32>, tensor<32x1024xf32>, tensor<i64>, tensor<i64>, tensor<1024xf32>, tensor<32x1024xf32>, tensor<1024x1xf32>, tensor<1024xf32>, tensor<1024x1xf32>, tensor<i64>
+      }) : () -> (tensor<i32>, tensor<32x1024xf32>, tensor<i64>, tensor<i64>, tensor<1024xf32>, tensor<32x1024xf32>, tensor<1024x1xf32>, tensor<1024xf32>, tensor<1024x1xf32>, tensor<i64>, tensor<i32>, tensor<32x1024xf32>, tensor<i64>, tensor<i64>, tensor<1024xf32>, tensor<32x1024xf32>, tensor<1024x1xf32>, tensor<1024xf32>, tensor<1024x1xf32>, tensor<i64>, tensor<i32>, tensor<32x1024xf32>, tensor<i64>, tensor<i64>, tensor<1024xf32>, tensor<32x1024xf32>, tensor<1024x1xf32>, tensor<1024xf32>, tensor<1024x1xf32>, tensor<i64>, tensor<i32>, tensor<32x1024xf32>, tensor<i64>, tensor<i64>, tensor<1024xf32>, tensor<32x1024xf32>, tensor<1024x1xf32>, tensor<1024xf32>, tensor<1024x1xf32>, tensor<i64>)
+      %17 = "tf.Concat"(%cst_1, %16#5, %16#15, %16#25, %16#35) : (tensor<i32>, tensor<32x1024xf32>, tensor<32x1024xf32>, tensor<32x1024xf32>, tensor<32x1024xf32>) -> tensor<128x1024xf32>
+      %18 = "tf.Concat"(%cst_2, %16#1, %16#11, %16#21, %16#31) : (tensor<i32>, tensor<32x1024xf32>, tensor<32x1024xf32>, tensor<32x1024xf32>, tensor<32x1024xf32>) -> tensor<128x1024xf32>
+      tf_device.return %16#0, %16#9, %16#8, %16#7, %16#6, %17, %16#4, %16#3, %16#2, %18 : tensor<i32>, tensor<i64>, tensor<1024x1xf32>, tensor<1024xf32>, tensor<1024x1xf32>, tensor<128x1024xf32>, tensor<1024xf32>, tensor<i64>, tensor<i64>, tensor<128x1024xf32>
+    }
+    "tf.AssignVariableOp"(%arg19, %13#18) <{validate_shape = false}> : (tensor<*x!tf_type.resource<tensor<128x1024xf32>>>, tensor<128x1024xf32>) -> ()
+    "tf.AssignVariableOp"(%arg1, %13#16) <{validate_shape = false}> : (tensor<*x!tf_type.resource<tensor<i64>>>, tensor<i64>) -> ()
+    "tf.AssignVariableOp"(%arg2, %13#14) <{validate_shape = false}> : (tensor<*x!tf_type.resource<tensor<i64>>>, tensor<i64>) -> ()
+    "tf.AssignVariableOp"(%arg4, %13#12) <{validate_shape = false}> : (tensor<*x!tf_type.resource<tensor<1024xf32>>>, tensor<1024xf32>) -> ()
+    "tf.AssignVariableOp"(%arg3, %13#10) <{validate_shape = false}> : (tensor<*x!tf_type.resource<tensor<128x1024xf32>>>, tensor<128x1024xf32>) -> ()
+    "tf.AssignVariableOp"(%arg5, %13#8) <{validate_shape = false}> : (tensor<*x!tf_type.resource<tensor<1024x1xf32>>>, tensor<1024x1xf32>) -> ()
+    "tf.AssignVariableOp"(%arg20, %13#6) <{validate_shape = false}> : (tensor<*x!tf_type.resource<tensor<1024xf32>>>, tensor<1024xf32>) -> ()
+    "tf.AssignVariableOp"(%arg21, %13#4) <{validate_shape = false}> : (tensor<*x!tf_type.resource<tensor<1024x1xf32>>>, tensor<1024x1xf32>) -> ()
+    "tf.AssignVariableOp"(%arg6, %13#2) <{validate_shape = true}> {_has_manual_control_dependencies = true} : (tensor<*x!tf_type.resource<tensor<i64>>>, tensor<i64>) -> ()
+    %14 = "tf.ReadVariableOp"(%arg2) {device = ""} : (tensor<*x!tf_type.resource<tensor<i64>>>) -> tensor<i64>
+    %15 = "tf.Identity"(%14) {device = ""} : (tensor<i64>) -> tensor<i64>
+    return %15 : tensor<i64>
+  }
+}
diff --git a/tensorflow/compiler/mlir/tf2xla/api/v2/tf_dialect_to_executor.cc b/tensorflow/compiler/mlir/tf2xla/api/v2/tf_dialect_to_executor.cc
index c92fd85d3567b4..cd13e869e811dd 100644
--- a/tensorflow/compiler/mlir/tf2xla/api/v2/tf_dialect_to_executor.cc
+++ b/tensorflow/compiler/mlir/tf2xla/api/v2/tf_dialect_to_executor.cc
@@ -88,6 +88,8 @@ void AddTfDialectToExecutorPasses(OpPassManager &pm) {
   pm.addNestedPass<FuncOp>(mlir::TFTPU::CreateTPUDevicePropagationPass());
   pm.addNestedPass<FuncOp>(mlir::TFTPU::CreateTPUColocateSplitsPass());
   pm.addPass(mlir::createSymbolDCEPass());
+  pm.addNestedPass<FuncOp>(
+      mlir::tf_executor::CreateTFExecutorGraphPruningPass());
   if (tensorflow::GetMlirCommonFlags()
           ->tf_mlir_enable_convert_control_to_data_outputs_pass) {
     bool composite_tpuexecute_side_effects =
diff --git a/tensorflow/compiler/mlir/tf2xla/api/v2/tf_dialect_to_executor_test.cc b/tensorflow/compiler/mlir/tf2xla/api/v2/tf_dialect_to_executor_test.cc
index 0c64dd3dcbe1a3..897c800d9e4cd7 100644
--- a/tensorflow/compiler/mlir/tf2xla/api/v2/tf_dialect_to_executor_test.cc
+++ b/tensorflow/compiler/mlir/tf2xla/api/v2/tf_dialect_to_executor_test.cc
@@ -15,12 +15,16 @@ limitations under the License.
 
 #include "tensorflow/compiler/mlir/tf2xla/api/v2/tf_dialect_to_executor.h"
 
+#include <stdlib.h>
+
 #include <cstdint>
 #include <string>
 
 #include <gtest/gtest.h>
 #include "absl/status/status.h"
 #include "absl/strings/str_cat.h"
+#include "absl/strings/string_view.h"
+#include "llvm/Support/raw_ostream.h"
 #include "mlir/IR/BuiltinOps.h"  // from @llvm-project
 #include "mlir/IR/DialectRegistry.h"  // from @llvm-project
 #include "mlir/IR/MLIRContext.h"  // from @llvm-project
@@ -30,7 +34,6 @@ limitations under the License.
 #include "tensorflow/core/lib/monitoring/cell_reader.h"
 #include "tensorflow/core/platform/resource_loader.h"
 #include "tsl/lib/core/status_test_util.h"
-#include "tsl/platform/status.h"
 
 namespace tensorflow {
 namespace tf2xla {
@@ -53,6 +56,16 @@ std::string TestDataPath() {
       "tensorflow/compiler/mlir/tf2xla/api/v2/testdata/");
 }
 
+size_t CountSubstring(absl::string_view str, absl::string_view substr) {
+  size_t count = 0;
+  size_t idx = str.find(substr);
+  while (idx != std::string::npos) {
+    count++;
+    idx = str.find(substr, idx + 1);
+  }
+  return count;
+}
+
 class TensorflowDialectToExecutorTest : public ::testing::Test {
  public:
   TensorflowDialectToExecutorTest() {
@@ -100,6 +113,23 @@ TEST_F(TensorflowDialectToExecutorTest, ErrorsWhenCannotConvert) {
   EXPECT_EQ(compilation_status.Delta(kExportFailed), 1);
 }
 
+TEST_F(TensorflowDialectToExecutorTest, PrunesDeadOps) {
+  CellReader<int64_t> compilation_status(kExportStreamzName);
+
+  TF_ASSERT_OK(CreateMlirModule("func_with_dead_ops.mlir"));
+
+  TF_EXPECT_OK(ExportFromTensorflowDialectToExecutor(*mlir_module_));
+
+  std::string module_dump;
+  llvm::raw_string_ostream raw_stream(module_dump);
+  mlir_module_->print(raw_stream);
+
+  EXPECT_EQ(compilation_status.Delta(kExportSuccess), 1);
+  EXPECT_EQ(compilation_status.Delta(kExportFailed), 0);
+  EXPECT_EQ(
+      CountSubstring(module_dump, "tf_executor.island wraps \"tf.Concat\""), 2);
+}
+
 }  // namespace
 }  // namespace v2
 }  // namespace tf2xla

From 1a59524ded5f73beb09fa78ac2a4eed8673ba6bb Mon Sep 17 00:00:00 2001
From: Eugene Zhulenev <ezhulenev@google.com>
Date: Mon, 13 May 2024 14:45:39 -0700
Subject: [PATCH 066/478] [xla] NFC: Delete deprecated flags and remains of xla
 runtime

PiperOrigin-RevId: 633338202
---
 third_party/xla/xla/debug_options_flags.cc    | 36 -------------------
 .../xla/xla/service/gpu/gpu_compiler.cc       |  2 +-
 .../xla/xla/service/gpu/gpu_executable.cc     | 13 -------
 .../xla/xla/service/gpu/gpu_executable.h      |  3 --
 .../xla/service/gpu/ir_emitter_unnested.cc    |  6 ----
 third_party/xla/xla/xla.proto                 | 26 ++++----------
 6 files changed, 8 insertions(+), 78 deletions(-)

diff --git a/third_party/xla/xla/debug_options_flags.cc b/third_party/xla/xla/debug_options_flags.cc
index 3eb1a7a2fd4292..831a26b2c96793 100644
--- a/third_party/xla/xla/debug_options_flags.cc
+++ b/third_party/xla/xla/debug_options_flags.cc
@@ -78,8 +78,6 @@ DebugOptions DefaultDebugOptionsIgnoringFlags() {
 #ifdef XLA_CPU_USE_ACL
   opts.set_xla_cpu_use_acl(true);
 #endif
-  opts.set_xla_cpu_use_xla_runtime(false);
-  opts.set_xla_cpu_sparse_cuda_threads(0);
 
   opts.set_xla_cpu_enable_fast_math(false);
   // Disable forms of fast math that have caused users problems in the past.
@@ -111,10 +109,8 @@ DebugOptions DefaultDebugOptionsIgnoringFlags() {
   opts.add_xla_gpu_enable_command_buffer(DebugOptions::CUBLAS);
   opts.add_xla_gpu_enable_command_buffer(DebugOptions::CUSTOM_CALL);
   opts.add_xla_gpu_enable_command_buffer(DebugOptions::CUDNN);
-  opts.set_xla_gpu_graph_num_runs_to_instantiate(-1);
   opts.set_xla_gpu_graph_min_graph_size(5);
   opts.set_xla_gpu_graph_enable_concurrent_region(false);
-  opts.set_xla_gpu_graph_eviction_timeout_seconds(60);
 
   // Despite the name, fast min/max on GPUs does not seem to be any faster, and
   // adds very counter-intuitive "NaN-swallowing" behavior.
@@ -138,7 +134,6 @@ DebugOptions DefaultDebugOptionsIgnoringFlags() {
   opts.set_xla_detailed_logging(true);
   opts.set_xla_enable_dumping(true);
 
-  opts.set_xla_gpu_enable_xla_runtime_executable(false);
   opts.set_xla_gpu_enable_custom_fusions(false);
   opts.set_xla_gpu_enable_address_computation_fusion(true);
   opts.set_xla_gpu_nccl_termination_timeout_seconds(-1);
@@ -766,17 +761,6 @@ void MakeDebugOptionsFlags(std::vector<tsl::Flag>* flag_list,
       "xla_cpu_use_acl", bool_setter_for(&DebugOptions::set_xla_cpu_use_acl),
       debug_options->xla_cpu_use_acl(),
       "Generate calls to ACL (Arm Compute Library) in the CPU backend."));
-  flag_list->push_back(
-      tsl::Flag("xla_cpu_use_xla_runtime",
-                bool_setter_for(&DebugOptions::set_xla_cpu_use_xla_runtime),
-                debug_options->xla_cpu_use_xla_runtime(),
-                "Enable XLA Runtime in the CPU backend."));
-  flag_list->push_back(tsl::Flag(
-      "xla_cpu_sparse_cuda_threads",
-      int32_setter_for(&DebugOptions::set_xla_cpu_sparse_cuda_threads),
-      debug_options->xla_cpu_sparse_cuda_threads(),
-      "Sets number fo CUDA threads for sparse GPU acceleration in the CPU "
-      "backend (0 = off)."));
   flag_list->push_back(tsl::Flag(
       "xla_gpu_crash_on_verification_failures",
       bool_setter_for(
@@ -1148,13 +1132,6 @@ void MakeDebugOptionsFlags(std::vector<tsl::Flag>* flag_list,
       " can either be a list of command types or a list of command types with"
       " + and - as prefix, which indicate adding or removing a command type"
       " to/from the default list."));
-  flag_list->push_back(tsl::Flag(
-      "xla_gpu_graph_num_runs_to_instantiate",
-      int32_setter_for(
-          &DebugOptions::set_xla_gpu_graph_num_runs_to_instantiate),
-      debug_options->xla_gpu_graph_num_runs_to_instantiate(),
-      "Instantiate a gpu graph after the time a captured function is executed "
-      "reaches the threshold."));
   flag_list->push_back(tsl::Flag(
       "xla_gpu_graph_min_graph_size",
       int32_setter_for(&DebugOptions::set_xla_gpu_graph_min_graph_size),
@@ -1168,14 +1145,6 @@ void MakeDebugOptionsFlags(std::vector<tsl::Flag>* flag_list,
                 debug_options->xla_gpu_graph_enable_concurrent_region(),
                 "Identify concurrent regions in gpu graphs and execute them "
                 "concurrently."));
-  flag_list->push_back(tsl::Flag(
-      "xla_gpu_graph_eviction_timeout_seconds",
-      int32_setter_for(
-          &DebugOptions::set_xla_gpu_graph_eviction_timeout_seconds),
-      debug_options->xla_gpu_graph_eviction_timeout_seconds(),
-      "Timeout in seconds to evict instantiated Gpu graphs from device. When "
-      "XLA instantiates new Gpu graphs, it evicts graphs that were not "
-      "recently executed to free space on device."));
 
   flag_list->push_back(
       tsl::Flag("xla_dump_disable_metadata",
@@ -1197,11 +1166,6 @@ void MakeDebugOptionsFlags(std::vector<tsl::Flag>* flag_list,
       "MLIR will be in the llvm-parsable format and can be processed by "
       "mlir-opt tools. "
       "Pretty print form is not legal MLIR."));
-  flag_list->push_back(tsl::Flag(
-      "xla_gpu_enable_xla_runtime_executable",
-      bool_setter_for(&DebugOptions::set_xla_gpu_enable_xla_runtime_executable),
-      debug_options->xla_gpu_enable_xla_runtime_executable(),
-      "Whether to enable XLA runtime for XLA:GPU backend"));
   flag_list->push_back(tsl::Flag(
       "xla_gpu_enable_custom_fusions",
       bool_setter_for(&DebugOptions::set_xla_gpu_enable_custom_fusions),
diff --git a/third_party/xla/xla/service/gpu/gpu_compiler.cc b/third_party/xla/xla/service/gpu/gpu_compiler.cc
index 35f6bf612f7ee6..276a4f2c9929a7 100644
--- a/third_party/xla/xla/service/gpu/gpu_compiler.cc
+++ b/third_party/xla/xla/service/gpu/gpu_compiler.cc
@@ -2210,7 +2210,7 @@ absl::Status GpuCompiler::RunPostSchedulingPipelines(
 
   // After we have a scheduled module and all operations wrapped into fusions we
   // can decide how to wrap them into command buffers.
-  if (!IsXlaRuntimeExecutableEnabled(module->config())) {
+  {
     HloPassPipeline pipeline("command-buffer-scheduling");
     auto driver_version = se::gpu::GpuDriver::GetDriverVersion();
 #if GOOGLE_CUDA
diff --git a/third_party/xla/xla/service/gpu/gpu_executable.cc b/third_party/xla/xla/service/gpu/gpu_executable.cc
index c4c55a0aed85af..25343eb9c93512 100644
--- a/third_party/xla/xla/service/gpu/gpu_executable.cc
+++ b/third_party/xla/xla/service/gpu/gpu_executable.cc
@@ -109,19 +109,6 @@ namespace gpu {
 
 using ::tsl::profiler::ScopedAnnotation;
 
-bool IsXlaRuntimeExecutableEnabled(const HloModuleConfig& config) {
-  bool enabled = config.debug_options().xla_gpu_enable_xla_runtime_executable();
-  if (enabled) {
-    LOG(ERROR)
-        << "XLA:GPU tried to use deprecated xla runtime by setting "
-           "--xla_gpu_enable_xla_runtime_executable flag to `true` but the "
-           "flag value was ignored as XLA:GPU uses default runtime. This flag "
-           "together with the deprecated code will be removed soon. Please "
-           "report bugs to XLA team if this breaks your workloads.";
-  }
-  return false;
-}
-
 static bool NeedsAsyncCommsStream(Thunk& thunk) {
   switch (thunk.kind()) {
     case Thunk::Kind::kNcclAllReduceStart:
diff --git a/third_party/xla/xla/service/gpu/gpu_executable.h b/third_party/xla/xla/service/gpu/gpu_executable.h
index ae481f1b45519d..7b5bfd7dcc2a8c 100644
--- a/third_party/xla/xla/service/gpu/gpu_executable.h
+++ b/third_party/xla/xla/service/gpu/gpu_executable.h
@@ -56,9 +56,6 @@ limitations under the License.
 namespace xla {
 namespace gpu {
 
-// Returns whether GpuExecutable runs with Xla Runtime.
-bool IsXlaRuntimeExecutableEnabled(const HloModuleConfig& config);
-
 // GPU-targeting implementation of the XLA Executable interface.
 //
 // Launches the given GPU kernel via the StreamExecutor.
diff --git a/third_party/xla/xla/service/gpu/ir_emitter_unnested.cc b/third_party/xla/xla/service/gpu/ir_emitter_unnested.cc
index 9ba1be41febde8..308a7c479822f5 100644
--- a/third_party/xla/xla/service/gpu/ir_emitter_unnested.cc
+++ b/third_party/xla/xla/service/gpu/ir_emitter_unnested.cc
@@ -1307,12 +1307,6 @@ absl::Status IrEmitterUnnested::EmitCustomCallThunk(
       return absl::OkStatus();
     }
 
-    // TODO(ezhulenev): Custom calls registered with an XLA runtime are not part
-    // of a legacy registry, or an FFI registry. For now we simply ignore them.
-    if (debug_options.xla_gpu_enable_xla_runtime_executable()) {
-      return absl::OkStatus();
-    }
-
     return absl::UnimplementedError(
         absl::StrCat("No registered implementation for custom call to ",
                      call_target_name, " for platform ", platform_name()));
diff --git a/third_party/xla/xla/xla.proto b/third_party/xla/xla/xla.proto
index c49d4a509a2618..b5044afe4d6292 100644
--- a/third_party/xla/xla/xla.proto
+++ b/third_party/xla/xla/xla.proto
@@ -114,8 +114,7 @@ message DebugOptions {
   // Generate calls to MKL-DNN in the CPU backend.
   bool xla_cpu_use_mkl_dnn = 97;
 
-  // Enable XLA Runtime in the CPU backend.
-  bool xla_cpu_use_xla_runtime = 177;
+  reserved 177;  // Was xla_cpu_use_xla_runtime
 
   reserved 98;  // Was xla_gpu_max_kernel_unroll_factor
 
@@ -165,12 +164,7 @@ message DebugOptions {
   // above!
   bool xla_gpu_enable_fast_min_max = 100;
 
-  // Defines the number of CUDA threads that can be used to accelerate
-  // a sparse computation compiled for the XLA Runtime and CPU backend.
-  // By default (value 0), no acceleration is used. Otherwise, this
-  // many threads may be used to accelerate sparse operations, typically
-  // useful when accelerating structured sparsity.
-  int32 xla_cpu_sparse_cuda_threads = 207;
+  reserved 207;  // Was xla_cpu_sparse_cuda_threads
 
   // Allows xla to increase the output precision of floating point operations
   // and all floating-point conversions to be simplified, including those
@@ -454,8 +448,7 @@ message DebugOptions {
   reserved 161;  // Was xla_gpu_bef_executable
   reserved 162;  // Was xla_gpu_bef_thunk
 
-  // If true, use XLA runtime for XLA:GPU backend.
-  bool xla_gpu_enable_xla_runtime_executable = 169;
+  reserved 169;  // Was xla_gpu_enable_xla_runtime_executable
 
   // If true, XLA will try to pattern match subgraphs of HLO operations into
   // custom fusions registered in the current process (pre-compiled hand written
@@ -466,7 +459,8 @@ message DebugOptions {
   // if `xla_gpu_enable_custom_fusion` set to true.
   string xla_gpu_enable_custom_fusions_re = 264;
 
-  // If true, use XLA runtime for XLA:GPU backend.
+  // Enables address computation fusion to optimize dynamic-slice and
+  // dynamic-update-slice operations around library calls.
   bool xla_gpu_enable_address_computation_fusion = 105;
 
   reserved 233;  // was xla_gpu_enable_gpu2_runtime
@@ -500,10 +494,7 @@ message DebugOptions {
   // Determine the types of commands that are recorded into command buffers.
   repeated CommandBufferCmdType xla_gpu_enable_command_buffer = 258;
 
-  // Only instantiates a GPU graph after the captured function execution count
-  // reaches the threshold. This constant is a heuristic to avoid creating a
-  // large number of CUDA graph instances in memory.
-  int32 xla_gpu_graph_num_runs_to_instantiate = 202;
+  reserved 202;  // Was xla_gpu_graph_num_runs_to_instantiate
 
   // This number determines how many moved instructions like fusion kernels are
   // required for a region to be captured as a function to be launched as a GPU
@@ -513,10 +504,7 @@ message DebugOptions {
   // Identify concurrent regions in GPU graphs and execute them concurrently.
   bool xla_gpu_graph_enable_concurrent_region = 215;
 
-  // Timeout in seconds to evict instantiated Gpu graphs from device. When XLA
-  // instantiates new Gpu graphs, it evicts graphs that were not recently
-  // executed to free space on device.
-  int32 xla_gpu_graph_eviction_timeout_seconds = 230;
+  reserved 230;  // Was xla_gpu_graph_eviction_timeout_seconds
 
   // Size threshold (in megabytes) for the GPU redzone scratch allocator.
   int64 xla_gpu_redzone_scratch_max_megabytes = 167;

From 09ce711fffb3c4268af4a1b55d39e1d37eb6e822 Mon Sep 17 00:00:00 2001
From: Fergus Henderson <fergus@google.com>
Date: Mon, 13 May 2024 14:46:48 -0700
Subject: [PATCH 067/478] Remove EmitErrorReporter class, since it isn't used.

PiperOrigin-RevId: 633338523
---
 tensorflow/compiler/mlir/lite/BUILD           | 15 -------
 .../compiler/mlir/lite/emit_error_reporter.cc | 30 --------------
 .../compiler/mlir/lite/emit_error_reporter.h  | 39 -------------------
 3 files changed, 84 deletions(-)
 delete mode 100644 tensorflow/compiler/mlir/lite/emit_error_reporter.cc
 delete mode 100644 tensorflow/compiler/mlir/lite/emit_error_reporter.h

diff --git a/tensorflow/compiler/mlir/lite/BUILD b/tensorflow/compiler/mlir/lite/BUILD
index 72970c2d07d4b5..2cfbecb6cbbc89 100644
--- a/tensorflow/compiler/mlir/lite/BUILD
+++ b/tensorflow/compiler/mlir/lite/BUILD
@@ -1070,21 +1070,6 @@ tf_native_cc_binary(
     ],
 )
 
-cc_library(
-    name = "emit_error_reporter",
-    srcs = [
-        "emit_error_reporter.cc",
-    ],
-    hdrs = [
-        "emit_error_reporter.h",
-    ],
-    deps = [
-        "//tensorflow/lite/core/api",
-        "@llvm-project//mlir:FuncDialect",
-        "@llvm-project//mlir:IR",
-    ],
-)
-
 cc_library(
     name = "flatbuffer_export",
     srcs = [
diff --git a/tensorflow/compiler/mlir/lite/emit_error_reporter.cc b/tensorflow/compiler/mlir/lite/emit_error_reporter.cc
deleted file mode 100644
index f9c4760326b5d6..00000000000000
--- a/tensorflow/compiler/mlir/lite/emit_error_reporter.cc
+++ /dev/null
@@ -1,30 +0,0 @@
-/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
-
-Licensed under the Apache License, Version 2.0 (the "License");
-you may not use this file except in compliance with the License.
-You may obtain a copy of the License at
-
-    http://www.apache.org/licenses/LICENSE-2.0
-
-Unless required by applicable law or agreed to in writing, software
-distributed under the License is distributed on an "AS IS" BASIS,
-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.
-==============================================================================*/
-
-#include "tensorflow/compiler/mlir/lite/emit_error_reporter.h"
-
-#include <cstdio>
-#include <vector>
-
-namespace tflite {
-
-int EmitErrorReporter::Report(const char* format, va_list args) {
-  std::vector<char> buf(1 + snprintf(nullptr, 0, format, args));
-  std::vsnprintf(buf.data(), buf.size(), format, args);
-  module_.emitError() << std::string(buf.begin(), buf.end());
-  return 0;
-}
-
-}  // namespace tflite
diff --git a/tensorflow/compiler/mlir/lite/emit_error_reporter.h b/tensorflow/compiler/mlir/lite/emit_error_reporter.h
deleted file mode 100644
index 9e9a5925600fc2..00000000000000
--- a/tensorflow/compiler/mlir/lite/emit_error_reporter.h
+++ /dev/null
@@ -1,39 +0,0 @@
-/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
-
-Licensed under the Apache License, Version 2.0 (the "License");
-you may not use this file except in compliance with the License.
-You may obtain a copy of the License at
-
-    http://www.apache.org/licenses/LICENSE-2.0
-
-Unless required by applicable law or agreed to in writing, software
-distributed under the License is distributed on an "AS IS" BASIS,
-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.
-==============================================================================*/
-
-#ifndef TENSORFLOW_COMPILER_MLIR_LITE_EMIT_ERROR_REPORTER_H_
-#define TENSORFLOW_COMPILER_MLIR_LITE_EMIT_ERROR_REPORTER_H_
-
-#include <cstdarg>
-
-#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
-#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
-#include "tensorflow/lite/core/api/error_reporter.h"
-
-namespace tflite {
-
-// Error reporter that reports errors via the module's emitError.
-class EmitErrorReporter : public ErrorReporter {
- public:
-  explicit EmitErrorReporter(mlir::ModuleOp module) : module_(module) {}
-  int Report(const char* format, va_list args) override;
-
- private:
-  mlir::ModuleOp module_;
-};
-
-}  // namespace tflite
-
-#endif  // TENSORFLOW_COMPILER_MLIR_LITE_EMIT_ERROR_REPORTER_H_

From 7c51020c4efc2ec0ab18d9d9c3c8a98408c46e08 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Mon, 13 May 2024 14:59:52 -0700
Subject: [PATCH 068/478] Move #include <stddef.h> outside of ifdef.

PiperOrigin-RevId: 633342149
---
 .../internal/optimized/integer_ops/depthwise_conv_3x3_filter.h | 3 ++-
 .../optimized/integer_ops/depthwise_conv_hybrid_3x3_filter.h   | 3 ++-
 2 files changed, 4 insertions(+), 2 deletions(-)

diff --git a/tensorflow/lite/kernels/internal/optimized/integer_ops/depthwise_conv_3x3_filter.h b/tensorflow/lite/kernels/internal/optimized/integer_ops/depthwise_conv_3x3_filter.h
index 0cb1a23e5567b9..8d761dd3dc3e10 100644
--- a/tensorflow/lite/kernels/internal/optimized/integer_ops/depthwise_conv_3x3_filter.h
+++ b/tensorflow/lite/kernels/internal/optimized/integer_ops/depthwise_conv_3x3_filter.h
@@ -15,6 +15,8 @@ limitations under the License.
 #ifndef TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_INTEGER_OPS_DEPTHWISE_CONV_3X3_FILTER_H_
 #define TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_INTEGER_OPS_DEPTHWISE_CONV_3X3_FILTER_H_
 
+#include <stddef.h>
+
 #include <memory>
 
 #include "ruy/profiler/instrumentation.h"  // from @ruy
@@ -32,7 +34,6 @@ namespace depthwise_conv {
 // Enable for arm64 except for the Nvidia Linux 4 Tegra (L4T) running on
 // Jetson TX-2. This compiler does not support the offsetof() macro.
 #if defined(__aarch64__) && !defined(GOOGLE_L4T)
-#include <stddef.h>
 
 // Represents the number of bytes offset from the start of the
 // DepthwiseConvParams struct. This is used in the asm to load parameters.
diff --git a/tensorflow/lite/kernels/internal/optimized/integer_ops/depthwise_conv_hybrid_3x3_filter.h b/tensorflow/lite/kernels/internal/optimized/integer_ops/depthwise_conv_hybrid_3x3_filter.h
index f9472515417e85..5c7abda84fcfa0 100644
--- a/tensorflow/lite/kernels/internal/optimized/integer_ops/depthwise_conv_hybrid_3x3_filter.h
+++ b/tensorflow/lite/kernels/internal/optimized/integer_ops/depthwise_conv_hybrid_3x3_filter.h
@@ -15,6 +15,8 @@ limitations under the License.
 #ifndef TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_INTEGER_OPS_DEPTHWISE_CONV_HYBRID_3X3_FILTER_H_
 #define TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_INTEGER_OPS_DEPTHWISE_CONV_HYBRID_3X3_FILTER_H_
 
+#include <stddef.h>
+
 #include <memory>
 
 #include "ruy/profiler/instrumentation.h"  // from @ruy
@@ -32,7 +34,6 @@ namespace depthwise_conv {
 // Enable for arm64 except for the Nvidia Linux 4 Tegra (L4T) running on
 // Jetson TX-2. This compiler does not support the offsetof() macro.
 #if defined(__aarch64__) && !defined(GOOGLE_L4T)
-#include <stddef.h>
 
 // Represents the number of bytes offset from the start of the
 // DepthwiseConvParams struct. This is used in the asm to load parameters.

From 0025bdd212d4c153ea87ce40cd098e094f3f944f Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Mon, 13 May 2024 15:25:03 -0700
Subject: [PATCH 069/478] Update ops-related pbtxt files.

PiperOrigin-RevId: 633349564
---
 tensorflow/core/ops/ops.pbtxt | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorflow/core/ops/ops.pbtxt b/tensorflow/core/ops/ops.pbtxt
index c8c55e08ca0d32..fb2f485604495b 100644
--- a/tensorflow/core/ops/ops.pbtxt
+++ b/tensorflow/core/ops/ops.pbtxt
@@ -1,4 +1,4 @@
-go/debugonly    
+go/nodeserialize  
 op {
   name: "Abort"
   attr {

From 1144753a8fb45c6dda8fec7c085bf307b6903aef Mon Sep 17 00:00:00 2001
From: Fangrui Song <maskray@google.com>
Date: Mon, 13 May 2024 15:41:41 -0700
Subject: [PATCH 070/478] Integrate LLVM at llvm/llvm-project@a6d7828f4c50

Updates LLVM usage to match
[a6d7828f4c50](https://github.com/llvm/llvm-project/commit/a6d7828f4c50)

PiperOrigin-RevId: 633354390
---
 third_party/llvm/workspace.bzl | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/third_party/llvm/workspace.bzl b/third_party/llvm/workspace.bzl
index d795d29b2a9434..243cc54703c4ec 100644
--- a/third_party/llvm/workspace.bzl
+++ b/third_party/llvm/workspace.bzl
@@ -4,8 +4,8 @@ load("//third_party:repo.bzl", "tf_http_archive")
 
 def repo(name):
     """Imports LLVM."""
-    LLVM_COMMIT = "1066eb55477044a3a92f3a40471375194dfcdbc8"
-    LLVM_SHA256 = "dc0fc82c184cb9ccd094f5ff821914d610d79529e82f0cbcb91ddedebff29a0a"
+    LLVM_COMMIT = "a6d7828f4c50c1ec7b0b5f61fe59d7a768175dcc"
+    LLVM_SHA256 = "c09ec3020fb6b136064ff32f53ac07067b0f12ccbf016ac69965e4e38d61a9c0"
 
     tf_http_archive(
         name = name,

From 6b9a25f438944132acc54ecf2c59af53c10fd6a1 Mon Sep 17 00:00:00 2001
From: Peter Hawkins <phawkins@google.com>
Date: Mon, 13 May 2024 15:48:45 -0700
Subject: [PATCH 071/478] [IFRT] Rename XlaProgram to HloProgram, move it from
 PJRT-IFRT to IFRT.

HLO programs represented in MLIR are independent of PJRT or XLA as a compiler, so they don't belong in PJRT-IFRT.

Also rename the Python binding API from get_xla_program to get_hlo_program, but this API is not used yet.

Cleanup; no functional changes intended.

PiperOrigin-RevId: 633356170
---
 tensorflow/core/tfrt/ifrt/BUILD               |  1 +
 .../core/tfrt/ifrt/ifrt_serving_executable.cc |  3 +-
 third_party/xla/xla/python/BUILD              |  1 +
 third_party/xla/xla/python/ifrt/hlo/BUILD     | 62 +++++++++++++++++++
 .../xla/xla/python/ifrt/hlo/hlo_program.cc    | 22 +++++++
 .../xla/xla/python/ifrt/hlo/hlo_program.h     | 54 ++++++++++++++++
 .../hlo/hlo_program_serdes.cc}                | 20 +++---
 .../hlo/hlo_program_serdes_test.cc}           | 10 +--
 .../xla/xla/python/ifrt/ir/tests/BUILD        |  1 +
 .../ifrt/ir/tests/executable_impl_test_lib.cc |  3 +-
 .../xla/xla/python/ifrt_proxy/common/BUILD    |  2 +-
 third_party/xla/xla/python/pjrt_ifrt/BUILD    | 41 +-----------
 .../xla/xla/python/pjrt_ifrt/pjrt_compiler.cc |  5 +-
 .../xla/xla/python/pjrt_ifrt/xla_compiler.cc  |  1 -
 .../xla/xla/python/pjrt_ifrt/xla_compiler.h   | 18 ------
 .../pjrt_ifrt/xla_executable_impl_test_lib.cc |  3 +-
 third_party/xla/xla/python/py_client.cc       |  3 +-
 third_party/xla/xla/python/py_program.cc      | 21 ++++---
 third_party/xla/xla/python/xla_client_test.py |  6 +-
 .../python/xla_extension/ifrt_programs.pyi    |  2 +-
 20 files changed, 186 insertions(+), 93 deletions(-)
 create mode 100644 third_party/xla/xla/python/ifrt/hlo/BUILD
 create mode 100644 third_party/xla/xla/python/ifrt/hlo/hlo_program.cc
 create mode 100644 third_party/xla/xla/python/ifrt/hlo/hlo_program.h
 rename third_party/xla/xla/python/{pjrt_ifrt/xla_program_serdes.cc => ifrt/hlo/hlo_program_serdes.cc} (85%)
 rename third_party/xla/xla/python/{pjrt_ifrt/xla_program_serdes_test.cc => ifrt/hlo/hlo_program_serdes_test.cc} (91%)

diff --git a/tensorflow/core/tfrt/ifrt/BUILD b/tensorflow/core/tfrt/ifrt/BUILD
index 369fecac053dbe..e0a3b4ed4553bc 100644
--- a/tensorflow/core/tfrt/ifrt/BUILD
+++ b/tensorflow/core/tfrt/ifrt/BUILD
@@ -95,6 +95,7 @@ cc_library(
         "@local_xla//xla/pjrt:host_callback",
         "@local_xla//xla/pjrt:pjrt_executable",
         "@local_xla//xla/python/ifrt",
+        "@local_xla//xla/python/ifrt/hlo:hlo_program",
         "@local_xla//xla/python/pjrt_ifrt",
         "@local_xla//xla/python/pjrt_ifrt:xla_ifrt",
         "@local_xla//xla/service:computation_placer_hdr",
diff --git a/tensorflow/core/tfrt/ifrt/ifrt_serving_executable.cc b/tensorflow/core/tfrt/ifrt/ifrt_serving_executable.cc
index c85c9a9d223b08..b65b9c68eb1f78 100644
--- a/tensorflow/core/tfrt/ifrt/ifrt_serving_executable.cc
+++ b/tensorflow/core/tfrt/ifrt/ifrt_serving_executable.cc
@@ -51,6 +51,7 @@ limitations under the License.
 #include "xla/python/ifrt/device.h"
 #include "xla/python/ifrt/executable.h"
 #include "xla/python/ifrt/future.h"
+#include "xla/python/ifrt/hlo/hlo_program.h"
 #include "xla/python/ifrt/host_callback.h"
 #include "xla/python/ifrt/shape.h"
 #include "xla/python/ifrt/sharding.h"
@@ -386,7 +387,7 @@ IfrtServingExecutable::CreateExecutableSynchronously(
   TF_ASSIGN_OR_RETURN(
       std::unique_ptr<xla::ifrt::LoadedExecutable> ifrt_executable,
       ifrt_client_->GetDefaultCompiler()->Compile(
-          std::make_unique<xla::ifrt::XlaProgram>(
+          std::make_unique<xla::ifrt::HloProgram>(
               tf2hlo_result.mlir_hlo_module.get()),
           std::make_unique<xla::ifrt::XlaCompileOptions>(
               xla_compile_options, loaded_host_callbacks)));
diff --git a/third_party/xla/xla/python/BUILD b/third_party/xla/xla/python/BUILD
index c1a36af8359a20..ec8631655c1243 100644
--- a/third_party/xla/xla/python/BUILD
+++ b/third_party/xla/xla/python/BUILD
@@ -389,6 +389,7 @@ cc_library(
         "//xla/python/ifrt",
         "//xla/python/ifrt:plugin_program",
         "//xla/python/ifrt:plugin_program_serdes",
+        "//xla/python/ifrt/hlo:hlo_program",
         "//xla/python/pjrt_ifrt",
         "//xla/python/pjrt_ifrt:xla_host_callback_proto_cc",
         "//xla/python/pjrt_ifrt:xla_ifrt",
diff --git a/third_party/xla/xla/python/ifrt/hlo/BUILD b/third_party/xla/xla/python/ifrt/hlo/BUILD
new file mode 100644
index 00000000000000..9925bf37459263
--- /dev/null
+++ b/third_party/xla/xla/python/ifrt/hlo/BUILD
@@ -0,0 +1,62 @@
+load("//xla:xla.bzl", "xla_cc_test")
+load("//xla/tsl:tsl.bzl", "internal_visibility")
+load("//xla/tsl:tsl.default.bzl", "get_compatible_with_portable")
+
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
+    default_visibility = internal_visibility([
+        "//xla/python/ifrt:friends",
+        "//xla/python/ifrt:internal",
+    ]),
+)
+
+cc_library(
+    name = "hlo_program",
+    srcs = ["hlo_program.cc"],
+    hdrs = ["hlo_program.h"],
+    compatible_with = get_compatible_with_portable(),
+    deps = [
+        "//xla/python/ifrt",
+        "@llvm-project//llvm:Support",
+        "@llvm-project//mlir:IR",
+    ],
+)
+
+cc_library(
+    name = "hlo_program_serdes",
+    srcs = ["hlo_program_serdes.cc"],
+    compatible_with = get_compatible_with_portable(),
+    deps = [
+        ":hlo_program",
+        "//xla/mlir_hlo:mhlo_passes",
+        "//xla/pjrt:mlir_to_hlo",
+        "//xla/python/ifrt:serdes",
+        "@com_google_absl//absl/status:statusor",
+        "@com_google_absl//absl/strings",
+        "@llvm-project//llvm:Support",
+        "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:Pass",
+        "@llvm-project//mlir:Support",
+        "@local_tsl//tsl/platform:status",
+        "@stablehlo//:stablehlo_portable_api",
+        "@stablehlo//:stablehlo_serialization",
+    ],
+    alwayslink = True,
+)
+
+xla_cc_test(
+    name = "hlo_program_serdes_test",
+    srcs = ["hlo_program_serdes_test.cc"],
+    deps = [
+        ":hlo_program",
+        ":hlo_program_serdes",
+        "//xla/mlir_hlo",
+        "//xla/pjrt:mlir_to_hlo",
+        "//xla/python/ifrt:serdes",
+        "@com_google_googletest//:gtest_main",
+        "@llvm-project//llvm:Support",
+        "@llvm-project//mlir:FuncDialect",
+        "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:Support",
+    ],
+)
diff --git a/third_party/xla/xla/python/ifrt/hlo/hlo_program.cc b/third_party/xla/xla/python/ifrt/hlo/hlo_program.cc
new file mode 100644
index 00000000000000..3e79c3d4ec6079
--- /dev/null
+++ b/third_party/xla/xla/python/ifrt/hlo/hlo_program.cc
@@ -0,0 +1,22 @@
+/* Copyright 2023 The OpenXLA Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+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.
+==============================================================================*/
+
+#include "xla/python/ifrt/hlo/hlo_program.h"
+
+namespace xla::ifrt {
+
+char HloProgram::ID = 0;
+
+}  // namespace xla::ifrt
diff --git a/third_party/xla/xla/python/ifrt/hlo/hlo_program.h b/third_party/xla/xla/python/ifrt/hlo/hlo_program.h
new file mode 100644
index 00000000000000..79f291cb062c3b
--- /dev/null
+++ b/third_party/xla/xla/python/ifrt/hlo/hlo_program.h
@@ -0,0 +1,54 @@
+/* Copyright 2023 The OpenXLA Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+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.
+==============================================================================*/
+
+#ifndef XLA_PYTHON_IFRT_HLO_HLO_PROGRAM_H_
+#define XLA_PYTHON_IFRT_HLO_HLO_PROGRAM_H_
+
+#include <memory>
+#include <optional>
+#include <utility>
+#include <vector>
+
+#include "llvm/Support/ExtensibleRTTI.h"
+#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
+#include "mlir/IR/MLIRContext.h"  // from @llvm-project
+#include "mlir/IR/OwningOpRef.h"  // from @llvm-project
+#include "xla/python/ifrt/program.h"
+
+namespace xla {
+namespace ifrt {
+
+struct HloProgram : llvm::RTTIExtends<HloProgram, Program> {
+  HloProgram() = default;
+  explicit HloProgram(mlir::ModuleOp module) : mlir_module(module) {}
+  HloProgram(std::unique_ptr<mlir::MLIRContext> context,
+             mlir::OwningOpRef<mlir::ModuleOp> module)
+      : mlir_module(*module),
+        mlir_context(std::move(context)),
+        owning_mlir_module(std::move(module)) {}
+
+  mlir::ModuleOp mlir_module;
+
+  static char ID;  // NOLINT
+
+ private:
+  std::unique_ptr<mlir::MLIRContext> mlir_context;
+  mlir::OwningOpRef<mlir::ModuleOp> owning_mlir_module;
+};
+
+}  // namespace ifrt
+}  // namespace xla
+
+#endif  // XLA_PYTHON_IFRT_HLO_HLO_PROGRAM_H_
diff --git a/third_party/xla/xla/python/pjrt_ifrt/xla_program_serdes.cc b/third_party/xla/xla/python/ifrt/hlo/hlo_program_serdes.cc
similarity index 85%
rename from third_party/xla/xla/python/pjrt_ifrt/xla_program_serdes.cc
rename to third_party/xla/xla/python/ifrt/hlo/hlo_program_serdes.cc
index 83348b7b8b6e27..1360bfa0df5875 100644
--- a/third_party/xla/xla/python/pjrt_ifrt/xla_program_serdes.cc
+++ b/third_party/xla/xla/python/ifrt/hlo/hlo_program_serdes.cc
@@ -29,8 +29,8 @@ limitations under the License.
 #include "stablehlo/dialect/Serialization.h"  // from @stablehlo
 #include "xla/mlir_hlo/mhlo/transforms/passes.h"
 #include "xla/pjrt/mlir_to_hlo.h"
+#include "xla/python/ifrt/hlo/hlo_program.h"
 #include "xla/python/ifrt/serdes.h"
-#include "xla/python/pjrt_ifrt/xla_compiler.h"
 #include "tsl/platform/status.h"
 
 namespace xla {
@@ -39,7 +39,7 @@ namespace ifrt {
 namespace {
 
 // Library that provides stable serialization and deserialization of
-// `xla::ifrt::XlaProgram`. Both serialization and deserialization require
+// `xla::ifrt::HloProgram`. Both serialization and deserialization require
 // linking in this library.
 //
 // Serialization:
@@ -50,21 +50,23 @@ namespace {
 // Deserialization:
 // ```
 // TF_ASSIGN_OR_RETURN(auto deserialized, Deserialize(serialized));
-// auto xla_program = llvm::dyn_cast<XlaProgram>(deserialized);
+// auto xla_program = llvm::dyn_cast<HloProgram>(deserialized);
 // ```
 
-class XlaProgramSerDes : public llvm::RTTIExtends<XlaProgramSerDes, SerDes> {
+class HloProgramSerDes : public llvm::RTTIExtends<HloProgramSerDes, SerDes> {
  public:
   absl::string_view type_name() const override {
+    // TODO(phawkins): whenever we next break compatibility, change this to
+    // "xla::ifrt::HloProgram".
     return "xla::ifrt::XlaProgram";
   }
 
   absl::StatusOr<std::string> Serialize(Serializable& serializable) override {
-    // Currently, PjRT-IFRT accepts an `XlaProgram` that contains C/MHLO. Since
+    // Currently, PjRT-IFRT accepts an `HloProgram` that contains C/MHLO. Since
     // these dialects don't provide version compatibility, the following
     // converts the module into StableHLO and use its portable serialization.
 
-    const auto& program = llvm::cast<XlaProgram>(serializable);
+    const auto& program = llvm::cast<HloProgram>(serializable);
     if (program.mlir_module == nullptr) {
       return absl::InvalidArgumentError("Unable to serialize null MLIR module");
     }
@@ -97,17 +99,17 @@ class XlaProgramSerDes : public llvm::RTTIExtends<XlaProgramSerDes, SerDes> {
       return absl::InvalidArgumentError("StableHLO => MHLO failed");
     }
 
-    return std::make_unique<XlaProgram>(std::move(context), std::move(module));
+    return std::make_unique<HloProgram>(std::move(context), std::move(module));
   }
 
   static char ID;  // NOLINT
 };
 
-char XlaProgramSerDes::ID = 0;  // NOLINT
+char HloProgramSerDes::ID = 0;  // NOLINT
 
 // clang-format off
 bool register_xla_program_serdes = ([]() {
-  RegisterSerDes<XlaProgram>(std::make_unique<XlaProgramSerDes>());
+  RegisterSerDes<HloProgram>(std::make_unique<HloProgramSerDes>());
 }(), true);
 // clang-format on
 
diff --git a/third_party/xla/xla/python/pjrt_ifrt/xla_program_serdes_test.cc b/third_party/xla/xla/python/ifrt/hlo/hlo_program_serdes_test.cc
similarity index 91%
rename from third_party/xla/xla/python/pjrt_ifrt/xla_program_serdes_test.cc
rename to third_party/xla/xla/python/ifrt/hlo/hlo_program_serdes_test.cc
index 41ce1c6ab67a3d..2febe7ddf31896 100644
--- a/third_party/xla/xla/python/pjrt_ifrt/xla_program_serdes_test.cc
+++ b/third_party/xla/xla/python/ifrt/hlo/hlo_program_serdes_test.cc
@@ -27,8 +27,8 @@ limitations under the License.
 #include "mlir/Support/DebugStringHelper.h"  // from @llvm-project
 #include "xla/mlir_hlo/mhlo/IR/hlo_ops.h"
 #include "xla/pjrt/mlir_to_hlo.h"
+#include "xla/python/ifrt/hlo/hlo_program.h"
 #include "xla/python/ifrt/serdes.h"
-#include "xla/python/pjrt_ifrt/xla_compiler.h"
 
 namespace xla {
 namespace ifrt {
@@ -37,7 +37,7 @@ namespace {
 using ::testing::IsNull;
 using ::testing::Not;
 
-TEST(XlaProgramSerDesTest, RoundTrip) {
+TEST(HloProgramSerDesTest, RoundTrip) {
   static constexpr absl::string_view kMlirModuleStr = R"(
 module {
   func.func @main(%arg0: tensor<2x3xf32>) -> tensor<2x3xf32> {
@@ -56,13 +56,13 @@ module {
         mlir::OwningOpRef<mlir::ModuleOp> module,
         xla::ParseMlirModuleString(kMlirModuleStr, *context));
     auto program =
-        std::make_unique<XlaProgram>(std::move(context), std::move(module));
+        std::make_unique<HloProgram>(std::move(context), std::move(module));
     TF_ASSERT_OK_AND_ASSIGN(serialized, Serialize(*program));
   }
 
   TF_ASSERT_OK_AND_ASSIGN(
-      std::unique_ptr<XlaProgram> xla_program,
-      Deserialize<XlaProgram>(serialized, /*options=*/nullptr));
+      std::unique_ptr<HloProgram> xla_program,
+      Deserialize<HloProgram>(serialized, /*options=*/nullptr));
 
   // Verify that the deserialized program has no StableHLO ops.
   bool has_unsupported_dialect = false;
diff --git a/third_party/xla/xla/python/ifrt/ir/tests/BUILD b/third_party/xla/xla/python/ifrt/ir/tests/BUILD
index 6b173ba429608c..65fb16d0a7813a 100644
--- a/third_party/xla/xla/python/ifrt/ir/tests/BUILD
+++ b/third_party/xla/xla/python/ifrt/ir/tests/BUILD
@@ -86,6 +86,7 @@ cc_library(
         "//xla/pjrt:pjrt_executable",
         "//xla/python/ifrt",
         "//xla/python/ifrt:test_util",
+        "//xla/python/ifrt/hlo:hlo_program",
         "//xla/python/ifrt/ir:compiler",
         "//xla/python/ifrt/ir:sharding_param",
         "//xla/python/pjrt_ifrt:xla_ifrt",
diff --git a/third_party/xla/xla/python/ifrt/ir/tests/executable_impl_test_lib.cc b/third_party/xla/xla/python/ifrt/ir/tests/executable_impl_test_lib.cc
index 25d0c05f3fcb27..8cbb44f908f0e2 100644
--- a/third_party/xla/xla/python/ifrt/ir/tests/executable_impl_test_lib.cc
+++ b/third_party/xla/xla/python/ifrt/ir/tests/executable_impl_test_lib.cc
@@ -28,6 +28,7 @@ limitations under the License.
 #include "xla/python/ifrt/device.h"
 #include "xla/python/ifrt/dtype.h"
 #include "xla/python/ifrt/executable.h"
+#include "xla/python/ifrt/hlo/hlo_program.h"
 #include "xla/python/ifrt/ir/compiler.h"
 #include "xla/python/ifrt/ir/sharding_param.h"
 #include "xla/python/ifrt/ir/tests/executable_impl_test_base.h"
@@ -304,7 +305,7 @@ module {
   TF_ASSERT_OK_AND_ASSIGN(
       std::unique_ptr<LoadedExecutable> child_exec,
       client_->GetDefaultCompiler()->Compile(
-          std::make_unique<xla::ifrt::XlaProgram>(*mhlo_module),
+          std::make_unique<xla::ifrt::HloProgram>(*mhlo_module),
           std::make_unique<XlaCompileOptions>(std::move(xla_options))));
 
   std::string source = R"(
diff --git a/third_party/xla/xla/python/ifrt_proxy/common/BUILD b/third_party/xla/xla/python/ifrt_proxy/common/BUILD
index 8bca5749c2aa98..30e93ada9ccfda 100644
--- a/third_party/xla/xla/python/ifrt_proxy/common/BUILD
+++ b/third_party/xla/xla/python/ifrt_proxy/common/BUILD
@@ -160,7 +160,7 @@ cc_library(
     name = "common_serdes",
     deps = [
         "//xla/python/ifrt:plugin_program_serdes",
-        "//xla/python/pjrt_ifrt:xla_program_serdes",
+        "//xla/python/ifrt/hlo:hlo_program_serdes",
     ],
     alwayslink = True,
 )
diff --git a/third_party/xla/xla/python/pjrt_ifrt/BUILD b/third_party/xla/xla/python/pjrt_ifrt/BUILD
index fc841a77512c7b..2ce7d49e13e4df 100644
--- a/third_party/xla/xla/python/pjrt_ifrt/BUILD
+++ b/third_party/xla/xla/python/pjrt_ifrt/BUILD
@@ -77,45 +77,6 @@ tf_proto_library(
     protodeps = ["//xla/pjrt:compile_options_proto"],
 )
 
-cc_library(
-    name = "xla_program_serdes",
-    srcs = ["xla_program_serdes.cc"],
-    compatible_with = get_compatible_with_portable(),
-    deps = [
-        ":xla_ifrt",
-        "//xla/mlir_hlo:mhlo_passes",
-        "//xla/pjrt:mlir_to_hlo",
-        "//xla/python/ifrt:serdes",
-        "@com_google_absl//absl/status:statusor",
-        "@com_google_absl//absl/strings",
-        "@llvm-project//llvm:Support",
-        "@llvm-project//mlir:IR",
-        "@llvm-project//mlir:Pass",
-        "@llvm-project//mlir:Support",
-        "@local_tsl//tsl/platform:status",
-        "@stablehlo//:stablehlo_portable_api",
-        "@stablehlo//:stablehlo_serialization",
-    ],
-    alwayslink = True,
-)
-
-xla_cc_test(
-    name = "xla_program_serdes_test",
-    srcs = ["xla_program_serdes_test.cc"],
-    deps = [
-        ":xla_ifrt",
-        ":xla_program_serdes",
-        "//xla/mlir_hlo",
-        "//xla/pjrt:mlir_to_hlo",
-        "//xla/python/ifrt:serdes",
-        "@com_google_googletest//:gtest_main",
-        "@llvm-project//llvm:Support",
-        "@llvm-project//mlir:FuncDialect",
-        "@llvm-project//mlir:IR",
-        "@llvm-project//mlir:Support",
-    ],
-)
-
 tf_proto_library(
     name = "xla_sharding_proto",
     srcs = ["xla_sharding.proto"],
@@ -174,6 +135,7 @@ cc_library(
         "//xla/pjrt:pjrt_executable",
         "//xla/python/ifrt",
         "//xla/python/ifrt:test_util",
+        "//xla/python/ifrt/hlo:hlo_program",
         "@com_google_absl//absl/status:statusor",
         "@com_google_absl//absl/strings",
         "@com_google_absl//absl/types:span",
@@ -259,6 +221,7 @@ cc_library(
         "//xla/pjrt:pjrt_layout",
         "//xla/pjrt:utils",
         "//xla/python/ifrt",
+        "//xla/python/ifrt/hlo:hlo_program",
         "//xla/service:hlo_proto_cc",
         "//xla/translate/mhlo_to_hlo:type_to_shape",
         "//xla/tsl/concurrency:ref_count",
diff --git a/third_party/xla/xla/python/pjrt_ifrt/pjrt_compiler.cc b/third_party/xla/xla/python/pjrt_ifrt/pjrt_compiler.cc
index bb7aaad5e02df0..7245857e70cdb5 100644
--- a/third_party/xla/xla/python/pjrt_ifrt/pjrt_compiler.cc
+++ b/third_party/xla/xla/python/pjrt_ifrt/pjrt_compiler.cc
@@ -21,6 +21,7 @@ limitations under the License.
 
 #include "absl/status/status.h"
 #include "llvm/Support/Casting.h"
+#include "xla/python/ifrt/hlo/hlo_program.h"
 #include "xla/python/pjrt_ifrt/pjrt_client.h"
 #include "xla/python/pjrt_ifrt/pjrt_executable.h"
 #include "xla/python/pjrt_ifrt/xla_compiler.h"
@@ -34,9 +35,9 @@ char PjRtCompiler::ID = 0;
 absl::StatusOr<std::unique_ptr<LoadedExecutable>> PjRtCompiler::Compile(
     std::unique_ptr<Program> program, std::unique_ptr<CompileOptions> options) {
   DCHECK(this);
-  const auto* xla_program = llvm::dyn_cast<XlaProgram>(program.get());
+  const auto* xla_program = llvm::dyn_cast<HloProgram>(program.get());
   if (xla_program == nullptr) {
-    return absl::InvalidArgumentError("PjRtCompiler requires an XlaProgram");
+    return absl::InvalidArgumentError("PjRtCompiler requires an HloProgram");
   }
   TF_ASSIGN_OR_RETURN(auto xla_compile_options,
                       GetXlaCompileOptions(std::move(options)));
diff --git a/third_party/xla/xla/python/pjrt_ifrt/xla_compiler.cc b/third_party/xla/xla/python/pjrt_ifrt/xla_compiler.cc
index 165742ad28904c..279fb27ad9d864 100644
--- a/third_party/xla/xla/python/pjrt_ifrt/xla_compiler.cc
+++ b/third_party/xla/xla/python/pjrt_ifrt/xla_compiler.cc
@@ -82,7 +82,6 @@ bool register_xla_compile_options_serdes = ([]{
 
 }  // namespace
 
-char XlaProgram::ID = 0;
 char XlaCompileOptions::ID = 0;
 char XlaDeserializeExecutableOptions::ID = 0;
 
diff --git a/third_party/xla/xla/python/pjrt_ifrt/xla_compiler.h b/third_party/xla/xla/python/pjrt_ifrt/xla_compiler.h
index dc670d52418a6d..ad6a8442249866 100644
--- a/third_party/xla/xla/python/pjrt_ifrt/xla_compiler.h
+++ b/third_party/xla/xla/python/pjrt_ifrt/xla_compiler.h
@@ -31,24 +31,6 @@ limitations under the License.
 namespace xla {
 namespace ifrt {
 
-struct XlaProgram : llvm::RTTIExtends<XlaProgram, Program> {
-  XlaProgram() = default;
-  explicit XlaProgram(mlir::ModuleOp module) : mlir_module(module) {}
-  XlaProgram(std::unique_ptr<mlir::MLIRContext> context,
-             mlir::OwningOpRef<mlir::ModuleOp> module)
-      : mlir_module(*module),
-        mlir_context(std::move(context)),
-        owning_mlir_module(std::move(module)) {}
-
-  mlir::ModuleOp mlir_module;
-
-  static char ID;  // NOLINT
-
- private:
-  std::unique_ptr<mlir::MLIRContext> mlir_context;
-  mlir::OwningOpRef<mlir::ModuleOp> owning_mlir_module;
-};
-
 // Wraps compilation options for an XLA computation.
 //
 // TODO(hyeontaek): Move this class out of pjrt_ifrt.
diff --git a/third_party/xla/xla/python/pjrt_ifrt/xla_executable_impl_test_lib.cc b/third_party/xla/xla/python/pjrt_ifrt/xla_executable_impl_test_lib.cc
index 077b635285bd12..e52c9c2a169290 100644
--- a/third_party/xla/xla/python/pjrt_ifrt/xla_executable_impl_test_lib.cc
+++ b/third_party/xla/xla/python/pjrt_ifrt/xla_executable_impl_test_lib.cc
@@ -34,6 +34,7 @@ limitations under the License.
 #include "xla/python/ifrt/device.h"
 #include "xla/python/ifrt/dtype.h"
 #include "xla/python/ifrt/executable.h"
+#include "xla/python/ifrt/hlo/hlo_program.h"
 #include "xla/python/ifrt/memory.h"
 #include "xla/python/ifrt/shape.h"
 #include "xla/python/ifrt/sharding.h"
@@ -97,7 +98,7 @@ absl::StatusOr<std::unique_ptr<LoadedExecutable>> CompileOnDevices(
       }
     }
   }
-  return compiler->Compile(std::make_unique<XlaProgram>(*module),
+  return compiler->Compile(std::make_unique<HloProgram>(*module),
                            std::move(compile_options));
 }
 
diff --git a/third_party/xla/xla/python/py_client.cc b/third_party/xla/xla/python/py_client.cc
index 5a67786bbe39d3..2321ec18abd3c9 100644
--- a/third_party/xla/xla/python/py_client.cc
+++ b/third_party/xla/xla/python/py_client.cc
@@ -64,6 +64,7 @@ limitations under the License.
 #include "xla/python/ifrt/device.h"
 #include "xla/python/ifrt/dtype.h"
 #include "xla/python/ifrt/executable.h"
+#include "xla/python/ifrt/hlo/hlo_program.h"
 #include "xla/python/ifrt/host_callback.h"
 #include "xla/python/ifrt/memory.h"
 #include "xla/python/nb_absl_span.h"  // IWYU pragma: keep
@@ -437,7 +438,7 @@ PyClient::CompileIfrtProgram(
   TF_ASSIGN_OR_RETURN(mlir::OwningOpRef<mlir::ModuleOp> module,
                       ParseMlirModuleString(mlir_module, context));
   return CompileIfrtProgram(
-      client, std::make_unique<xla::ifrt::XlaProgram>(module.get()),
+      client, std::make_unique<xla::ifrt::HloProgram>(module.get()),
       MakeIfrtCompileOptions(std::move(options), std::move(host_callbacks)));
 }
 
diff --git a/third_party/xla/xla/python/py_program.cc b/third_party/xla/xla/python/py_program.cc
index 230606f879038f..e85b0a5fda8770 100644
--- a/third_party/xla/xla/python/py_program.cc
+++ b/third_party/xla/xla/python/py_program.cc
@@ -34,6 +34,7 @@ limitations under the License.
 #include "xla/pjrt/pjrt_executable.h"
 #include "xla/pjrt/status_casters.h"
 #include "xla/python/ifrt/compiler.h"
+#include "xla/python/ifrt/hlo/hlo_program.h"
 #include "xla/python/ifrt/host_callback.h"
 #include "xla/python/ifrt/plugin_program.h"
 #include "xla/python/pjrt_ifrt/xla_compiler.h"
@@ -65,23 +66,23 @@ MakePluginCompileOptions() {
   return std::make_unique<ifrt::PluginCompileOptions>();
 }
 
-absl::StatusOr<std::unique_ptr<ifrt::Program>> MakeXlaProgram(
+absl::StatusOr<std::unique_ptr<ifrt::Program>> MakeHloProgram(
     absl::string_view mlir_module) {
   auto context = std::make_unique<mlir::MLIRContext>();
   TF_ASSIGN_OR_RETURN(mlir::OwningOpRef<mlir::ModuleOp> module,
                       ParseMlirModuleString(mlir_module, *context));
-  return std::make_unique<xla::ifrt::XlaProgram>(std::move(context),
+  return std::make_unique<xla::ifrt::HloProgram>(std::move(context),
                                                  std::move(module));
 }
 
-absl::StatusOr<std::unique_ptr<ifrt::Program>> MakeXlaProgramFromString(
+absl::StatusOr<std::unique_ptr<ifrt::Program>> MakeHloProgramFromString(
     std::string mlir_module) {
-  return MakeXlaProgram(mlir_module);
+  return MakeHloProgram(mlir_module);
 }
 
-absl::StatusOr<std::unique_ptr<ifrt::Program>> MakeXlaProgramFromBytes(
+absl::StatusOr<std::unique_ptr<ifrt::Program>> MakeHloProgramFromBytes(
     nb::bytes mlir_module) {
-  return MakeXlaProgram(
+  return MakeHloProgram(
       absl::string_view(mlir_module.c_str(), mlir_module.size()));
 }
 
@@ -109,10 +110,10 @@ void BuildIfrtProgramsSubmodule(nanobind::module_& m) {
   nb::class_<xla::ifrt::CompileOptions> ifrt_compile_options_base_class(
       sub_module, "CompileOptions");
   sub_module
-      .def("make_xla_program",
-           xla::ValueOrThrowWrapper(MakeXlaProgramFromString))
-      .def("make_xla_program",
-           xla::ValueOrThrowWrapper(MakeXlaProgramFromBytes))
+      .def("make_hlo_program",
+           xla::ValueOrThrowWrapper(MakeHloProgramFromString))
+      .def("make_hlo_program",
+           xla::ValueOrThrowWrapper(MakeHloProgramFromBytes))
       .def("make_plugin_program",
            xla::ValueOrThrowWrapper(MakePluginProgramFromString))
       .def("make_plugin_program",
diff --git a/third_party/xla/xla/python/xla_client_test.py b/third_party/xla/xla/python/xla_client_test.py
index afcacfdeafe80e..71549ecb3e2979 100644
--- a/third_party/xla/xla/python/xla_client_test.py
+++ b/third_party/xla/xla/python/xla_client_test.py
@@ -3076,15 +3076,15 @@ def testPluginProgramDoesNotCompile(self):
       program = xla_client.ifrt_programs.make_plugin_program("foobar")
       options = xla_client.ifrt_programs.make_plugin_compile_options()
       with self.assertRaisesRegex(
-          xla_client.XlaRuntimeError, "PjRtCompiler requires an XlaProgram"
+          xla_client.XlaRuntimeError, "PjRtCompiler requires an HloProgram"
       ):
         self.backend.compile_ifrt_program(program, options)
 
     @unittest.skipIf(pathways, "does not work with non-ifrt legacy pathways")
-    def testXlaProgramViaIfrtProgram(self):
+    def testHloProgramViaIfrtProgram(self):
       c = self._NewComputation()
       ops.Iota(c, xla_client.PrimitiveType.F32, 10)
-      program = xla_client.ifrt_programs.make_xla_program(
+      program = xla_client.ifrt_programs.make_hlo_program(
           xla_computation_to_mlir_module(c.build())
       )
       options = xla_client.ifrt_programs.make_xla_compile_options(
diff --git a/third_party/xla/xla/python/xla_extension/ifrt_programs.pyi b/third_party/xla/xla/python/xla_extension/ifrt_programs.pyi
index 58b0996b75797b..830df593b6e4fc 100644
--- a/third_party/xla/xla/python/xla_extension/ifrt_programs.pyi
+++ b/third_party/xla/xla/python/xla_extension/ifrt_programs.pyi
@@ -21,7 +21,7 @@ class Program:  ...
 
 class CompileOptions:  ...
 
-def make_xla_program(mlir_module: Union[str, bytes]) -> Program: ...
+def make_hlo_program(mlir_module: Union[str, bytes]) -> Program: ...
 
 def make_plugin_program(data: Union[str, bytes]) -> Program: ...
 

From 8c004713d87c04183b4e3a6dc44838aca5bd488c Mon Sep 17 00:00:00 2001
From: Kuy Mainwaring <kuym@google.com>
Date: Mon, 13 May 2024 16:08:23 -0700
Subject: [PATCH 072/478] [XLA:GPU] Clang-tidy fixes for
 xla/service/gpu/fusions/copy.h

PiperOrigin-RevId: 633361839
---
 third_party/xla/xla/service/gpu/fusions/copy.h | 1 +
 1 file changed, 1 insertion(+)

diff --git a/third_party/xla/xla/service/gpu/fusions/copy.h b/third_party/xla/xla/service/gpu/fusions/copy.h
index 00b9ce4506a5a3..53fd1fb3ddfd31 100644
--- a/third_party/xla/xla/service/gpu/fusions/copy.h
+++ b/third_party/xla/xla/service/gpu/fusions/copy.h
@@ -15,6 +15,7 @@ limitations under the License.
 #ifndef XLA_SERVICE_GPU_FUSIONS_COPY_H_
 #define XLA_SERVICE_GPU_FUSIONS_COPY_H_
 
+#include <utility>
 #include <vector>
 
 #include "absl/status/statusor.h"

From 96737ab5a239b663871eea6d55581ac8b2182ef2 Mon Sep 17 00:00:00 2001
From: Kuangyuan Chen <chky@google.com>
Date: Mon, 13 May 2024 16:08:46 -0700
Subject: [PATCH 073/478] Stop generating corert const ops in tf_to_tfrt

PiperOrigin-RevId: 633361967
---
 .../tfrt/tests/tf_to_corert/const_tensor.mlir |  6 ++---
 .../tfrt/tests/tf_to_corert/control_flow.mlir |  4 +--
 .../tf_to_corert/decompose_resource_op.mlir   |  3 +--
 .../mlir/tfrt/transforms/tf_to_tfrt.cc        | 25 ++++++++++---------
 4 files changed, 18 insertions(+), 20 deletions(-)

diff --git a/tensorflow/compiler/mlir/tfrt/tests/tf_to_corert/const_tensor.mlir b/tensorflow/compiler/mlir/tfrt/tests/tf_to_corert/const_tensor.mlir
index b208fe390acc3f..6596d650889384 100644
--- a/tensorflow/compiler/mlir/tfrt/tests/tf_to_corert/const_tensor.mlir
+++ b/tensorflow/compiler/mlir/tfrt/tests/tf_to_corert/const_tensor.mlir
@@ -9,12 +9,12 @@ func.func @string_tensor() -> (tensor<0x!tf_type.string>, tensor<7x!tf_type.stri
   func.return %0, %1 : tensor<0x!tf_type.string>, tensor<7x!tf_type.string>
 }
 
-// Convert tf.Const to corert.const_dense_tensor only on cpu device
+// Convert tf.Const to tfrt_fallback_async.const_dense_tensor only on cpu device
 // CHECK-LABEL: func @dense_tensor
 func.func @dense_tensor() -> tensor<4xui64> {
-  // CHECK: corert.const_dense_tensor dense<[1, 2, 3, 4]> : tensor<4xui64>
+  // CHECK: tfrt_fallback_async.const_dense_tensor dense<[1, 2, 3, 4]> : tensor<4xui64>
   %0 = "tf.Const"() {value = dense<[1, 2, 3, 4]> : tensor<4xui64>} : () -> tensor<4xui64>
-  // CHECK: corert.const_dense_tensor  dense<1.000000e+00> : tensor<1xbf16>
+  // CHECK: tfrt_fallback_async.const_dense_tensor  dense<1.000000e+00> : tensor<1xbf16>
   %1 = "tf.Const"() {device = "/device:CPU:0", value = dense<[1.0]> : tensor<1xbf16>} : () -> tensor<4xbf16>
   // CHECK: corert.executeop({{.*}}) "tf.Const"() {dtype = ui64, value = dense<[1, 2, 3, 4]> : tensor<4xui64>} : 1
   %2 = "tf.Const"() {device = "/device:GPU:0", value = dense<[1, 2, 3, 4]> : tensor<4xui64>} : () -> tensor<4xui64>
diff --git a/tensorflow/compiler/mlir/tfrt/tests/tf_to_corert/control_flow.mlir b/tensorflow/compiler/mlir/tfrt/tests/tf_to_corert/control_flow.mlir
index e3562b286c47f3..ad3232042ca5e7 100644
--- a/tensorflow/compiler/mlir/tfrt/tests/tf_to_corert/control_flow.mlir
+++ b/tensorflow/compiler/mlir/tfrt/tests/tf_to_corert/control_flow.mlir
@@ -56,10 +56,8 @@ func.func @while_body_add2(%arg0: tensor<i32>) -> tensor<i32> {
 // CHECK-LABEL: func @while_test
 // CHECK-SAME: ([[ARG0:%.+]]: !tfrt.chain) -> (!tfrt.chain, !tfrt_fallback.tf_tensor)
 func.func @while_test() -> (tensor<i32>) {
-  // CHECK: [[CONST_TH:%.*]] = corert.const_dense_tensor dense<0> : tensor<i32>
+  // CHECK: [[CONST:%.*]] = tfrt_fallback_async.const_dense_tensor dense<0> : tensor<i32>
   %0 = "tf.Const"() {device = "/device:CPU:0", value = dense<0> : tensor<i32>} : () -> tensor<i32>
-  // CHECK: [[CONST:%.*]] = tfrt_fallback_async.corert_tensorhandle_to_fallback_tensor [[CONST_TH]]
-  // CHECK: (!corert.tensorhandle) -> (!tfrt_fallback.tf_tensor)
   // CHECK: [[pred_res:%.*]]:2 = tfrt.call @"while_cond_lt9/tfrt_predicate"([[ARG0]], [[CONST]]) : (!tfrt.chain, !tfrt_fallback.tf_tensor) -> (!tfrt.chain, i1)
   // CHECK: [[while_res:%.]]:2 = tfrt.while [[pred_res]]#1 @"while_body_add2/tfrt_body_1"([[pred_res]]#0, [[CONST]])
   // CHECK-SAME: (!tfrt.chain, !tfrt_fallback.tf_tensor) -> (!tfrt.chain, !tfrt_fallback.tf_tensor)
diff --git a/tensorflow/compiler/mlir/tfrt/tests/tf_to_corert/decompose_resource_op.mlir b/tensorflow/compiler/mlir/tfrt/tests/tf_to_corert/decompose_resource_op.mlir
index ed308d02ad6e0b..ff0f0e7dbfd2cd 100644
--- a/tensorflow/compiler/mlir/tfrt/tests/tf_to_corert/decompose_resource_op.mlir
+++ b/tensorflow/compiler/mlir/tfrt/tests/tf_to_corert/decompose_resource_op.mlir
@@ -5,8 +5,7 @@ module attributes {tf.versions = {bad_consumers = [], min_consumer = 12 : i32, p
 // CHECK-LABEL: func @gather
 // CHECK-SAME: ([[in_chain:%.*]]: !tfrt.chain
 // CHECK-SAME: [[arg0:%.*]]: !tfrt_fallback.tf_tensor, [[arg1:%.*]]: !tfrt_fallback.tf_tensor)
-// CHECK: [[const_th:%.*]] = corert.const_dense_tensor
-// CHECK-NEXT: [[const:%.*]] = tfrt_fallback_async.corert_tensorhandle_to_fallback_tensor [[const_th]] {device = "/job:localhost/replica:0/task:0/device:CPU:0"}
+// CHECK: [[const:%.*]] = tfrt_fallback_async.const_dense_tensor
 // CHECK-NEXT: [[out_chain:%.*]], [[value:%.*]] = tfrt_fallback_async.executeop.seq([[in_chain]]) key(0) cost({{.*}}) device("/job:localhost/replica:0/task:0/device:CPU:0") "tf.ReadVariableOp"({{.*}})
 // CHECK-NEXT: [[res:%.*]] = tfrt_fallback_async.executeop key(1) cost({{.*}}) device("/job:localhost/replica:0/task:0/device:CPU:0") "tf.GatherV2"([[value]], {{.*}}, [[const]])
 // CHECK-NEXT: tfrt.return [[out_chain]], [[res]] : !tfrt.chain, !tfrt_fallback.tf_tensor
diff --git a/tensorflow/compiler/mlir/tfrt/transforms/tf_to_tfrt.cc b/tensorflow/compiler/mlir/tfrt/transforms/tf_to_tfrt.cc
index 024baa43c73132..f090745e0ae1c4 100644
--- a/tensorflow/compiler/mlir/tfrt/transforms/tf_to_tfrt.cc
+++ b/tensorflow/compiler/mlir/tfrt/transforms/tf_to_tfrt.cc
@@ -738,11 +738,11 @@ class FallbackBatchFunctionOpConversion
 
 // Lower a tf.Const op that creates a string tensor to a native
 // corert.create_string_tensor op.
-class CoreRTConstDenseTensorOpConversion
+class FallbackConstDenseTensorOpConversion
     : public mlir::OpConversionPattern<mlir::TF::ConstOp> {
  public:
-  CoreRTConstDenseTensorOpConversion(mlir::MLIRContext *context,
-                                     CoreRTConverter *corert_converter)
+  FallbackConstDenseTensorOpConversion(mlir::MLIRContext *context,
+                                       CoreRTConverter *corert_converter)
       : mlir::OpConversionPattern<mlir::TF::ConstOp>(context, kCoreRTBenefit),
         corert_converter_(*corert_converter) {}
 
@@ -756,8 +756,8 @@ class CoreRTConstDenseTensorOpConversion
     if (auto parsed_device_name = corert_converter_.ParseDeviceName(op))
       if (parsed_device_name->device_type != DEVICE_CPU) return failure();
 
-    auto new_op = rewriter.create<tfrt::corert::ConstDenseTensorOp>(
-        op.getLoc(), corert_converter_.tensor_handle_type(),
+    auto new_op = rewriter.create<tfrt::fallback_async::ConstDenseTensorOp>(
+        op.getLoc(), rewriter.getType<tfrt::fallback::TFTensorType>(),
         mlir::cast<DenseElementsAttr>(op.getValue()));
     rewriter.replaceOp(op, new_op->getResult(0));
     return success();
@@ -860,11 +860,11 @@ class TFRTFuncOpSignatureConversion
 
 // Lower a tf.Const op that creates a string tensor to a native
 // corert.create_string_tensor op.
-class CoreRTConstStringTensorOpConversion
+class FallbackConstStringTensorOpConversion
     : public mlir::OpConversionPattern<mlir::TF::ConstOp> {
  public:
-  CoreRTConstStringTensorOpConversion(mlir::MLIRContext *context,
-                                      CoreRTConverter *corert_converter)
+  FallbackConstStringTensorOpConversion(mlir::MLIRContext *context,
+                                        CoreRTConverter *corert_converter)
       : mlir::OpConversionPattern<mlir::TF::ConstOp>(context, kCoreRTBenefit),
         corert_converter_(*corert_converter) {}
 
@@ -890,8 +890,8 @@ class CoreRTConstStringTensorOpConversion
     for (auto dim : shape)
       dims.push_back(rewriter.getIntegerAttr(i64_type, dim));
 
-    auto new_op = rewriter.create<tfrt::corert::ConstStringTensorOp>(
-        op.getLoc(), corert_converter_.tensor_handle_type(),
+    auto new_op = rewriter.create<tfrt::fallback_async::ConstStringTensorOp>(
+        op.getLoc(), rewriter.getType<tfrt::fallback::TFTensorType>(),
         rewriter.getArrayAttr(dims), rewriter.getArrayAttr(values));
 
     rewriter.replaceOp(op, new_op.getResult());
@@ -1532,8 +1532,9 @@ void PopulateTFToTFRTConversionPatterns(
 
   // Here we use specialized patterns for tf.Const on CPU as it is incorrect to
   // use ExecuteOp pattern to convert string tensor attribute.
-  patterns->add<CoreRTConstStringTensorOpConversion,
-                CoreRTConstDenseTensorOpConversion>(context, corert_converter);
+  patterns->add<FallbackConstStringTensorOpConversion,
+                FallbackConstDenseTensorOpConversion>(context,
+                                                      corert_converter);
 }
 
 // Lower TF dialect MLIR to TFRT dialect.

From 53cc802c4002a59aa11032423c927494282039f7 Mon Sep 17 00:00:00 2001
From: Yang Chen <yangchen@google.com>
Date: Mon, 13 May 2024 16:12:02 -0700
Subject: [PATCH 074/478] #tf-buildcop Fix test timeout.

Error:
```
//tensorflow/core/data/service/snapshot:snapshot_stream_writer_checkpoint_test TIMEOUT in 3 out of 3 in 96.0s
  Stats over 3 runs: max = 96.0s, min = 66.2s, avg = 79.4s, dev = 12.4s
  tensorflow/core/data/service/snapshot/snapshot_stream_writer_checkpoint_test/test.log
  tensorflow/core/data/service/snapshot/snapshot_stream_writer_checkpoint_test/test_attempts/attempt_1.log
  tensorflow/core/data/service/snapshot/snapshot_stream_writer_checkpoint_test/test_attempts/attempt_2.log
```
PiperOrigin-RevId: 633362798
---
 tensorflow/core/data/service/snapshot/BUILD | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorflow/core/data/service/snapshot/BUILD b/tensorflow/core/data/service/snapshot/BUILD
index 523bfebc44dcb5..40b5cbaa6873aa 100644
--- a/tensorflow/core/data/service/snapshot/BUILD
+++ b/tensorflow/core/data/service/snapshot/BUILD
@@ -462,7 +462,7 @@ cc_library(
 
 tf_cc_test(
     name = "snapshot_stream_writer_checkpoint_test",
-    size = "small",
+    size = "medium",
     srcs = ["snapshot_stream_writer_checkpoint_test.cc"],
     deps = [
         ":path_utils",

From 7fd8e285c203dde472a2671a7cda8ba55bcb49d6 Mon Sep 17 00:00:00 2001
From: Kuy Mainwaring <kuym@google.com>
Date: Mon, 13 May 2024 16:12:02 -0700
Subject: [PATCH 075/478] [XLA:GPU] Clang-tidy fixes for
 xla/service/gpu/fusions/mlir/mlir_fusion_emitter.cc

PiperOrigin-RevId: 633362802
---
 .../xla/xla/service/gpu/fusions/mlir/mlir_fusion_emitter.cc      | 1 +
 .../xla/xla/service/gpu/fusions/mlir/mlir_fusion_emitter.h       | 1 +
 2 files changed, 2 insertions(+)

diff --git a/third_party/xla/xla/service/gpu/fusions/mlir/mlir_fusion_emitter.cc b/third_party/xla/xla/service/gpu/fusions/mlir/mlir_fusion_emitter.cc
index 41108b774751b9..7daca77a7bd37a 100644
--- a/third_party/xla/xla/service/gpu/fusions/mlir/mlir_fusion_emitter.cc
+++ b/third_party/xla/xla/service/gpu/fusions/mlir/mlir_fusion_emitter.cc
@@ -22,6 +22,7 @@ limitations under the License.
 #include <string>
 #include <utility>
 #include <variant>
+#include <vector>
 
 #include "absl/algorithm/container.h"
 #include "absl/container/flat_hash_map.h"
diff --git a/third_party/xla/xla/service/gpu/fusions/mlir/mlir_fusion_emitter.h b/third_party/xla/xla/service/gpu/fusions/mlir/mlir_fusion_emitter.h
index 1c15af1f901d58..1ea02ac98cd120 100644
--- a/third_party/xla/xla/service/gpu/fusions/mlir/mlir_fusion_emitter.h
+++ b/third_party/xla/xla/service/gpu/fusions/mlir/mlir_fusion_emitter.h
@@ -18,6 +18,7 @@ limitations under the License.
 #include <functional>
 #include <memory>
 #include <string>
+#include <vector>
 
 #include "absl/container/flat_hash_set.h"
 #include "absl/status/status.h"

From a7f0e0261b48393c207782f89c0cf991e45152b2 Mon Sep 17 00:00:00 2001
From: David Dunleavy <ddunleavy@google.com>
Date: Mon, 13 May 2024 16:12:57 -0700
Subject: [PATCH 076/478] Use `--nobuild_tests_only` even on ARM builds of XLA

PiperOrigin-RevId: 633363011
---
 third_party/xla/.kokoro/linux/build.sh | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/third_party/xla/.kokoro/linux/build.sh b/third_party/xla/.kokoro/linux/build.sh
index 2a72099e5f5148..81251bd3f44a6c 100644
--- a/third_party/xla/.kokoro/linux/build.sh
+++ b/third_party/xla/.kokoro/linux/build.sh
@@ -62,7 +62,7 @@ if is_linux_gpu_job ; then
     UNSUPPORTED_GPU_TAGS="$(echo -requires-gpu-sm{80,86,89,90}{,-only})"
     TAGS_FILTER="${TAGS_FILTER},${UNSUPPORTED_GPU_TAGS// /,}"
 
-    ADDITIONAL_FLAGS="$ADDITIONAL_FLAGS --nobuild_tests_only --run_under=//tools/ci_build/gpu_build:parallel_gpu_execute"
+    ADDITIONAL_FLAGS="$ADDITIONAL_FLAGS --run_under=//tools/ci_build/gpu_build:parallel_gpu_execute"
     RBE_FLAGS="--config=rbe_linux_cuda_nvcc --jobs=150"
     echo "***NOTE: nvidia-smi lists the highest CUDA version the driver supports, which may be different than the version of CUDA actually used!!***"
     nvidia-smi
@@ -79,17 +79,17 @@ else
         RBE_FLAGS="--config=rbe_cross_compile_linux_arm64_xla --jobs=150"
     else
         RBE_FLAGS="--config=rbe_linux_cpu --jobs=150"
-        ADDITIONAL_FLAGS="$ADDITIONAL_FLAGS --nobuild_tests_only"
     fi
 fi
 
 # Build & test XLA
 docker exec xla bazel \
         test \
-        --build_tag_filters=$TAGS_FILTER  \
+        --build_tag_filters=$TAGS_FILTER \
         --test_tag_filters=$TAGS_FILTER \
         --test_output=errors \
         --keep_going \
+        --nobuild_tests_only \
         --features=layering_check \
         --profile=/tf/pkg/profile.json.gz \
         --flaky_test_attempts=3 \

From e4a0ecea4889d268e5f44f4532348df79c42a8f8 Mon Sep 17 00:00:00 2001
From: Kuy Mainwaring <kuym@google.com>
Date: Mon, 13 May 2024 16:21:02 -0700
Subject: [PATCH 077/478] [XLA:GPU] Clang-tidy fixes for
 xla/service/gpu/runtime/convolution_thunk.cc

PiperOrigin-RevId: 633365188
---
 third_party/xla/xla/service/gpu/runtime/BUILD         |  5 +++--
 .../xla/xla/service/gpu/runtime/convolution_thunk.cc  | 11 ++++++++++-
 2 files changed, 13 insertions(+), 3 deletions(-)

diff --git a/third_party/xla/xla/service/gpu/runtime/BUILD b/third_party/xla/xla/service/gpu/runtime/BUILD
index d5ad69a37b933c..d0b940e88b523e 100644
--- a/third_party/xla/xla/service/gpu/runtime/BUILD
+++ b/third_party/xla/xla/service/gpu/runtime/BUILD
@@ -512,18 +512,19 @@ cc_library(
         "TENSORFLOW_USE_ROCM=1",
     ]),
     deps = [
+        ":thunk",
         "//xla:util",
         "//xla/service:buffer_assignment",
         "//xla/service/gpu:gpu_conv_runner",
-        "//xla/service/gpu:stream_executor_util",
-        "//xla/service/gpu/runtime:thunk",
         "//xla/stream_executor",
         "@com_google_absl//absl/base:core_headers",
         "@com_google_absl//absl/container:flat_hash_map",
         "@com_google_absl//absl/container:inlined_vector",
+        "@com_google_absl//absl/log:check",
         "@com_google_absl//absl/status",
         "@com_google_absl//absl/synchronization",
         "@com_google_absl//absl/types:span",
+        "@local_tsl//tsl/platform:errors",
     ],
 )
 
diff --git a/third_party/xla/xla/service/gpu/runtime/convolution_thunk.cc b/third_party/xla/xla/service/gpu/runtime/convolution_thunk.cc
index 6e8158d866aaf5..b7bba0fab3b418 100644
--- a/third_party/xla/xla/service/gpu/runtime/convolution_thunk.cc
+++ b/third_party/xla/xla/service/gpu/runtime/convolution_thunk.cc
@@ -15,17 +15,26 @@ limitations under the License.
 
 #include "xla/service/gpu/runtime/convolution_thunk.h"
 
+#include <cstdint>
 #include <memory>
 #include <optional>
+#include <utility>
+#include <vector>
 
 #include "absl/container/inlined_vector.h"
+#include "absl/log/check.h"
 #include "absl/status/status.h"
+#include "absl/synchronization/mutex.h"
+#include "absl/types/span.h"
 #include "xla/service/buffer_assignment.h"
 #include "xla/service/gpu/gpu_conv_runner.h"
-#include "xla/service/gpu/stream_executor_util.h"
+#include "xla/service/gpu/runtime/thunk.h"
+#include "xla/stream_executor/device_memory.h"
+#include "xla/stream_executor/dnn.h"
 #include "xla/stream_executor/scratch_allocator.h"
 #include "xla/stream_executor/stream_executor.h"
 #include "xla/util.h"
+#include "tsl/platform/errors.h"
 
 namespace xla {
 namespace gpu {

From 60645902f1745b90232cce4e0116faf87e05948e Mon Sep 17 00:00:00 2001
From: Siqiao Wu <siqiaowu@google.com>
Date: Mon, 13 May 2024 16:38:50 -0700
Subject: [PATCH 078/478] Add some internal change.

PiperOrigin-RevId: 633369641
---
 tensorflow/core/tfrt/ifrt/BUILD               |  20 ++
 .../ifrt/ifrt_restore_tensor_registry_test.cc | 190 ++++++++++++++++++
 2 files changed, 210 insertions(+)
 create mode 100644 tensorflow/core/tfrt/ifrt/ifrt_restore_tensor_registry_test.cc

diff --git a/tensorflow/core/tfrt/ifrt/BUILD b/tensorflow/core/tfrt/ifrt/BUILD
index e0a3b4ed4553bc..9db13a6387a286 100644
--- a/tensorflow/core/tfrt/ifrt/BUILD
+++ b/tensorflow/core/tfrt/ifrt/BUILD
@@ -299,6 +299,26 @@ cc_library(
     ],
 )
 
+tf_cc_test(
+    name = "ifrt_restore_tensor_registry_test",
+    srcs = ["ifrt_restore_tensor_registry_test.cc"],
+    tags = ["no_oss"],
+    deps = [
+        ":ifrt_restore_tensor_registry",
+        "//tensorflow/compiler/mlir/tfrt/transforms/ifrt:ifrt_types",
+        "//tensorflow/core:framework",
+        "//tensorflow/core/framework:tensor",
+        "//tensorflow/core/framework:tensor_testutil",
+        "//tensorflow/core/framework:types_proto_cc",
+        "@com_google_absl//absl/status",
+        "@com_google_googletest//:gtest_main",
+        "@local_tsl//tsl/lib/core:status_test_util",
+        "@local_tsl//tsl/platform:status_matchers",
+        "@local_tsl//tsl/platform:statusor",
+        "@local_xla//xla/python/ifrt",
+    ],
+)
+
 tf_cc_test(
     name = "ifrt_loaded_variable_utils_test",
     srcs = ["ifrt_loaded_variable_utils_test.cc"],
diff --git a/tensorflow/core/tfrt/ifrt/ifrt_restore_tensor_registry_test.cc b/tensorflow/core/tfrt/ifrt/ifrt_restore_tensor_registry_test.cc
new file mode 100644
index 00000000000000..de0a27aecc4104
--- /dev/null
+++ b/tensorflow/core/tfrt/ifrt/ifrt_restore_tensor_registry_test.cc
@@ -0,0 +1,190 @@
+/* Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+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.
+==============================================================================*/
+#include "tensorflow/core/tfrt/ifrt/ifrt_restore_tensor_registry.h"
+
+#include <cstdint>
+
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+#include "absl/status/status.h"
+#include "tensorflow/compiler/mlir/tfrt/transforms/ifrt/ifrt_types.h"
+#include "xla/python/ifrt/future.h"
+#include "tensorflow/core/framework/tensor.h"
+#include "tensorflow/core/framework/tensor_shape.h"
+#include "tensorflow/core/framework/tensor_testutil.h"
+#include "tensorflow/core/framework/types.pb.h"
+#include "tsl/lib/core/status_test_util.h"
+#include "tsl/platform/status_matchers.h"
+#include "tsl/platform/statusor.h"
+
+using tsl::testing::IsOk;
+using tsl::testing::StatusIs;
+
+namespace tensorflow {
+namespace ifrt_serving {
+namespace {
+
+TEST(IfrtRestoreTensorRegistryTest, RetrieveNonRegisteredTensorFails) {
+  IfrtRestoreTensorRegistry registry;
+  EXPECT_THAT(registry.GetRestoredTensor("input_tensor_1").Await(),
+              StatusIs(absl::StatusCode::kNotFound));
+}
+
+TEST(IfrtRestoreTensorRegistryTest,
+     RetrieveNonRegisteredTensorDTypeAndShapeFails) {
+  IfrtRestoreTensorRegistry registry;
+  EXPECT_THAT(registry.GetDtypeAndShape("input_tensor_1"),
+              StatusIs(absl::StatusCode::kNotFound));
+}
+
+TEST(IfrtRestoreTensorRegistryTest, SetNonExistedTensorAsUsedByHostFails) {
+  IfrtRestoreTensorRegistry registry;
+  EXPECT_THAT(registry.SetUsedByHost("input_tensor_1"),
+              StatusIs(absl::StatusCode::kNotFound));
+}
+
+TEST(IfrtRestoreTensorRegistryTest, RegisteredExistedTensorFails) {
+  auto input_tensor =
+      test::AsTensor<int32_t>({1, 2, 3, 4}, tensorflow::TensorShape({2, 2}));
+  auto promise = xla::ifrt::Future<tensorflow::Tensor>::CreatePromise();
+  auto future = xla::ifrt::Future<tensorflow::Tensor>(promise);
+
+  IfrtRestoreTensorRegistry::RestoredTensorInfo restored_tensor_info = {
+      .used_by_host = false,
+      .dtype_and_shape =
+          {
+              .dtype = DT_INT32,
+              .shape = tensorflow::TensorShape({2, 2}),
+          },
+      .tensor_future = future};
+  IfrtRestoreTensorRegistry registry;
+  EXPECT_THAT(registry.TryRegister("input_tensor_2", restored_tensor_info),
+              IsOk());
+  promise.Set(input_tensor);
+  EXPECT_THAT(registry.TryRegister("input_tensor_2", restored_tensor_info),
+              StatusIs(absl::StatusCode::kAlreadyExists));
+}
+
+TEST(IfrtRestoreTensorRegistryTest, SetTensorAsUsedByHost) {
+  auto promise = xla::ifrt::Future<tensorflow::Tensor>::CreatePromise();
+  auto future = xla::ifrt::Future<tensorflow::Tensor>(promise);
+  IfrtRestoreTensorRegistry::RestoredTensorInfo restored_tensor_info = {
+      .used_by_host = false,
+      .dtype_and_shape =
+          {
+              .dtype = DT_INT32,
+              .shape = tensorflow::TensorShape({2, 2}),
+          },
+      .tensor_future = future};
+  IfrtRestoreTensorRegistry registry;
+  EXPECT_THAT(registry.TryRegister("input_tensor_1", restored_tensor_info),
+              IsOk());
+  EXPECT_THAT(registry.SetUsedByHost("input_tensor_1"), IsOk());
+}
+
+TEST(IfrtRestoreTensorRegistryTest, RegisteredTensorCanBeRetrieved) {
+  auto input_tensor =
+      test::AsTensor<int32_t>({1, 2, 3, 4}, tensorflow::TensorShape({2, 2}));
+  auto promise = xla::ifrt::Future<tensorflow::Tensor>::CreatePromise();
+  auto future = xla::ifrt::Future<tensorflow::Tensor>(promise);
+
+  IfrtRestoreTensorRegistry::RestoredTensorInfo restored_tensor_info = {
+      .used_by_host = false,
+      .dtype_and_shape =
+          {
+              .dtype = DT_INT32,
+              .shape = tensorflow::TensorShape({2, 2}),
+          },
+      .tensor_future = future};
+  IfrtRestoreTensorRegistry registry;
+  EXPECT_THAT(registry.TryRegister("input_tensor_1", restored_tensor_info),
+              IsOk());
+  promise.Set(input_tensor);
+  TF_ASSERT_OK_AND_ASSIGN(tensorflow::Tensor retrieved,
+                          registry.GetRestoredTensor("input_tensor_1").Await());
+  test::ExpectEqual(retrieved, input_tensor);
+  TF_ASSERT_OK_AND_ASSIGN(DtypeAndShape dtype_and_shape,
+                          registry.GetDtypeAndShape("input_tensor_1"));
+  EXPECT_TRUE(
+      dtype_and_shape.shape.IsSameSize(tensorflow::TensorShape({2, 2})));
+  EXPECT_EQ(dtype_and_shape.dtype, DT_INT32);
+}
+
+TEST(IfrtRestoreTensorRegistryTest,
+     RegisteredTensorDTypeAndShapeCanBeRetrieved) {
+  auto input_tensor =
+      test::AsTensor<int32_t>({1, 2, 3, 4}, tensorflow::TensorShape({2, 2}));
+  auto promise = xla::ifrt::Future<tensorflow::Tensor>::CreatePromise();
+  auto future = xla::ifrt::Future<tensorflow::Tensor>(promise);
+
+  IfrtRestoreTensorRegistry::RestoredTensorInfo restored_tensor_info = {
+      .used_by_host = false,
+      .dtype_and_shape =
+          {
+              .dtype = DT_INT32,
+              .shape = tensorflow::TensorShape({2, 2}),
+          },
+      .tensor_future = future};
+  IfrtRestoreTensorRegistry registry;
+  EXPECT_THAT(registry.TryRegister("input_tensor_1", restored_tensor_info),
+              IsOk());
+  TF_ASSERT_OK_AND_ASSIGN(DtypeAndShape dtype_and_shape,
+                          registry.GetDtypeAndShape("input_tensor_1"));
+  EXPECT_TRUE(
+      dtype_and_shape.shape.IsSameSize(tensorflow::TensorShape({2, 2})));
+  EXPECT_EQ(dtype_and_shape.dtype, DT_INT32);
+}
+
+TEST(IfrtRestoreTensorRegistryTest, FeezeTensorRegistry) {
+  auto input_tensor =
+      test::AsTensor<int32_t>({1, 2, 3, 4}, tensorflow::TensorShape({2, 2}));
+  auto promise1 = xla::ifrt::Future<tensorflow::Tensor>::CreatePromise();
+  auto future1 = xla::ifrt::Future<tensorflow::Tensor>(promise1);
+  auto promise2 = xla::ifrt::Future<tensorflow::Tensor>::CreatePromise();
+  auto future2 = xla::ifrt::Future<tensorflow::Tensor>(promise2);
+
+  IfrtRestoreTensorRegistry::RestoredTensorInfo restored_tensor_info1 = {
+      .used_by_host = false,
+      .dtype_and_shape =
+          {
+              .dtype = DT_INT32,
+              .shape = tensorflow::TensorShape({2, 2}),
+          },
+      .tensor_future = future1};
+  IfrtRestoreTensorRegistry::RestoredTensorInfo restored_tensor_info2 = {
+      .used_by_host = true,
+      .dtype_and_shape =
+          {
+              .dtype = DT_INT32,
+              .shape = tensorflow::TensorShape({2, 2}),
+          },
+      .tensor_future = future2};
+  IfrtRestoreTensorRegistry registry;
+  TF_ASSERT_OK(registry.TryRegister("input_tensor_1", restored_tensor_info1));
+  TF_ASSERT_OK(registry.TryRegister("input_tensor_2", restored_tensor_info2));
+  promise1.Set(input_tensor);
+  promise2.Set(input_tensor);
+  registry.Freeze();
+  // Tensor with `used_by_host` set to false will be freed after freeze.
+  EXPECT_THAT(registry.GetRestoredTensor("input_tensor_1").Await(),
+              StatusIs(absl::StatusCode::kUnavailable));
+  // Tensor with `used_by_host` set to true will be kept after freeze.
+  TF_ASSERT_OK_AND_ASSIGN(tensorflow::Tensor retrieved,
+                          registry.GetRestoredTensor("input_tensor_2").Await());
+  test::ExpectEqual(retrieved, input_tensor);
+}
+}  // namespace
+}  // namespace ifrt_serving
+}  // namespace tensorflow

From 0cb4750928311f7d901f99739377f9a62c5f1d92 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Mon, 13 May 2024 16:54:33 -0700
Subject: [PATCH 079/478] Migrate coord agent and service to use absl 
 mutex/condvar libraries directly + some clang fixes.

PiperOrigin-RevId: 633373388
---
 .../distributed_runtime/coordination/BUILD    | 10 ++-
 .../coordination/coordination_service.cc      | 80 ++++++++++---------
 .../coordination/coordination_service.h       |  7 +-
 .../coordination_service_agent.cc             | 53 ++++++------
 .../coordination/coordination_service_agent.h |  5 +-
 5 files changed, 80 insertions(+), 75 deletions(-)

diff --git a/third_party/xla/xla/tsl/distributed_runtime/coordination/BUILD b/third_party/xla/xla/tsl/distributed_runtime/coordination/BUILD
index 33cf2f09025912..bf67649bb5474f 100644
--- a/third_party/xla/xla/tsl/distributed_runtime/coordination/BUILD
+++ b/third_party/xla/xla/tsl/distributed_runtime/coordination/BUILD
@@ -49,11 +49,13 @@ cc_library(
     hdrs = ["coordination_service.h"],
     deps = [
         ":coordination_client",
+        "@com_google_absl//absl/log",
         "@com_google_absl//absl/status",
+        "@com_google_absl//absl/status:statusor",
         "@com_google_absl//absl/strings",
         "@com_google_absl//absl/time",
+        "@local_tsl//tsl/platform:macros",
         "@local_tsl//tsl/platform:status",
-        "@local_tsl//tsl/platform:statusor",
         "@local_tsl//tsl/protobuf:coordination_config_proto_cc",
     ],
 )
@@ -72,15 +74,16 @@ tsl_gpu_library(
         "//xla/tsl/util:device_name_utils",
         "@com_google_absl//absl/container:flat_hash_map",
         "@com_google_absl//absl/container:flat_hash_set",
+        "@com_google_absl//absl/hash",
         "@com_google_absl//absl/log",
         "@com_google_absl//absl/status",
+        "@com_google_absl//absl/status:statusor",
         "@com_google_absl//absl/strings",
         "@com_google_absl//absl/synchronization",
         "@com_google_absl//absl/time",
         "@local_tsl//tsl/platform:env",
         "@local_tsl//tsl/platform:errors",
         "@local_tsl//tsl/platform:macros",
-        "@local_tsl//tsl/platform:mutex",
         "@local_tsl//tsl/platform:random",
         "@local_tsl//tsl/platform:status",
         "@local_tsl//tsl/platform:strcat",
@@ -142,16 +145,15 @@ tsl_gpu_library(
         "@com_google_absl//absl/container:flat_hash_set",
         "@com_google_absl//absl/log",
         "@com_google_absl//absl/status",
+        "@com_google_absl//absl/status:statusor",
         "@com_google_absl//absl/strings",
         "@com_google_absl//absl/synchronization",
         "@com_google_absl//absl/time",
         "@local_tsl//tsl/framework:cancellation",
         "@local_tsl//tsl/lib/monitoring:gauge",
         "@local_tsl//tsl/platform:env",
-        "@local_tsl//tsl/platform:mutex",
         "@local_tsl//tsl/platform:random",
         "@local_tsl//tsl/platform:status",
-        "@local_tsl//tsl/platform:statusor",
         "@local_tsl//tsl/platform:thread_annotations",
         "@local_tsl//tsl/protobuf:coordination_config_proto_cc",
         "@local_tsl//tsl/protobuf:coordination_service_proto_cc",
diff --git a/third_party/xla/xla/tsl/distributed_runtime/coordination/coordination_service.cc b/third_party/xla/xla/tsl/distributed_runtime/coordination/coordination_service.cc
index 09622fe037e934..8c3a0fb1a57edf 100644
--- a/third_party/xla/xla/tsl/distributed_runtime/coordination/coordination_service.cc
+++ b/third_party/xla/xla/tsl/distributed_runtime/coordination/coordination_service.cc
@@ -16,6 +16,8 @@ limitations under the License.
 #include "xla/tsl/distributed_runtime/coordination/coordination_service.h"
 
 #include <algorithm>
+#include <cassert>
+#include <cstddef>
 #include <cstdint>
 #include <functional>
 #include <map>
@@ -26,11 +28,13 @@ limitations under the License.
 
 #include "absl/container/flat_hash_map.h"
 #include "absl/container/flat_hash_set.h"
+#include "absl/hash/hash.h"
 #include "absl/log/log.h"
 #include "absl/status/status.h"
 #include "absl/strings/str_cat.h"
 #include "absl/strings/str_join.h"
 #include "absl/strings/string_view.h"
+#include "absl/synchronization/mutex.h"
 #include "absl/synchronization/notification.h"
 #include "absl/time/time.h"
 #include "xla/tsl/distributed_runtime/call_options.h"
@@ -39,11 +43,8 @@ limitations under the License.
 #include "xla/tsl/util/device_name_utils.h"
 #include "tsl/platform/env.h"
 #include "tsl/platform/errors.h"
-#include "tsl/platform/macros.h"
-#include "tsl/platform/mutex.h"
 #include "tsl/platform/random.h"
 #include "tsl/platform/status.h"
-#include "tsl/platform/strcat.h"
 #include "tsl/platform/thread_annotations.h"
 #include "tsl/protobuf/coordination_config.pb.h"
 #include "tsl/protobuf/coordination_service.pb.h"
@@ -223,7 +224,7 @@ class CoordinationServiceStandaloneImpl : public CoordinationServiceInterface {
 
     CoordinatedTaskState state_ = CoordinatedTaskState::TASKSTATE_DISCONNECTED;
     absl::Status status_;
-    mutex last_heartbeat_mu_;
+    absl::Mutex last_heartbeat_mu_;
     uint64_t last_heartbeat_us_ TF_GUARDED_BY(last_heartbeat_mu_);
     // This denotes the deadline after which we stop accepting heartbeats from a
     // disconnected task. This grace period accounts for the lag time between
@@ -252,19 +253,19 @@ class CoordinationServiceStandaloneImpl : public CoordinationServiceInterface {
   const std::string shutdown_barrier_id_ =
       absl::StrCat("Shutdown::", std::to_string(service_incarnation_));
 
-  mutex state_mu_;
+  absl::Mutex state_mu_;
   absl::flat_hash_map<std::string, std::unique_ptr<TaskState>> cluster_state_
       TF_GUARDED_BY(state_mu_);
   DeviceInfo cluster_devices_ TF_GUARDED_BY(state_mu_);
 
-  mutex kv_mu_;
+  absl::Mutex kv_mu_;
   // Ordered map to store config key-values
   std::map<std::string, std::string> kv_store_ TF_GUARDED_BY(kv_mu_);
   absl::flat_hash_map<std::string, std::vector<StatusOrValueCallback>> get_cb_
       TF_GUARDED_BY(kv_mu_);
 
-  mutex check_staleness_thread_shutdown_mu_;
-  condition_variable check_staleness_thread_cv_;
+  absl::Mutex check_staleness_thread_shutdown_mu_;
+  absl::CondVar check_staleness_thread_cv_;
   bool shutting_down_ TF_GUARDED_BY(check_staleness_thread_shutdown_mu_) =
       false;
   std::unique_ptr<Thread> check_staleness_thread_;
@@ -287,7 +288,7 @@ void CoordinationServiceStandaloneImpl::TaskState::SetConnected(
   state_ = CoordinatedTaskState::TASKSTATE_CONNECTED;
   status_ = absl::OkStatus();
   task_incarnation_ = task_incarnation;
-  mutex_lock l(last_heartbeat_mu_);
+  absl::MutexLock l(&last_heartbeat_mu_);
   last_heartbeat_us_ = Env::Default()->NowMicros();
 }
 
@@ -314,14 +315,14 @@ absl::Status CoordinationServiceStandaloneImpl::TaskState::RecordHeartbeat(
         "Incarnation ID mismatch: expecting ", task_incarnation_, " but got ",
         task_incarnation, ". This means the remote task has restarted."));
   }
-  mutex_lock l(last_heartbeat_mu_);
+  absl::MutexLock l(&last_heartbeat_mu_);
   last_heartbeat_us_ = Env::Default()->NowMicros();
   return absl::OkStatus();
 }
 
 int64_t
 CoordinationServiceStandaloneImpl::TaskState::TimeSinceLastHeartbeatMs() {
-  mutex_lock l(last_heartbeat_mu_);
+  absl::MutexLock l(&last_heartbeat_mu_);
   return (Env::Default()->NowMicros() - last_heartbeat_us_) / 1000;
 }
 
@@ -388,8 +389,9 @@ void CoordinationServiceStandaloneImpl::StartCheckStaleness() {
         absl::flat_hash_map<std::string, BarrierState*> expired_barriers;
         while (true) {
           {
-            mutex_lock l(check_staleness_thread_shutdown_mu_);
-            check_staleness_thread_cv_.wait_for(l, std::chrono::seconds(1));
+            absl::MutexLock l(&check_staleness_thread_shutdown_mu_);
+            check_staleness_thread_cv_.WaitWithTimeout(
+                &check_staleness_thread_shutdown_mu_, absl::Seconds(1));
             if (shutting_down_) {
               return;
             }
@@ -397,7 +399,7 @@ void CoordinationServiceStandaloneImpl::StartCheckStaleness() {
           // Heartbeat check.
           absl::Status status = absl::OkStatus();
           {
-            mutex_lock l(state_mu_);
+            absl::MutexLock l(&state_mu_);
             for (const auto& [task_name, task_state] : cluster_state_) {
               // Skip tasks that are not registered or in error state
               if (task_state->GetState() !=
@@ -444,7 +446,7 @@ void CoordinationServiceStandaloneImpl::StartCheckStaleness() {
           // Barrier timeout check.
           uint64_t current_time_micros = Env::Default()->NowMicros();
           {
-            mutex_lock l(state_mu_);
+            absl::MutexLock l(&state_mu_);
             // Gather barriers which have timed out.
             for (const std::string& barrier_id : ongoing_barriers_) {
               auto* barrier = &barriers_[barrier_id];
@@ -492,7 +494,7 @@ void CoordinationServiceStandaloneImpl::StartCheckStaleness() {
 
 void CoordinationServiceStandaloneImpl::Stop(bool shut_staleness_thread) {
   {
-    mutex_lock l(kv_mu_);
+    absl::MutexLock l(&kv_mu_);
     for (const auto& [key, get_kv_callbacks] : get_cb_) {
       for (const auto& get_kv_callback : get_kv_callbacks) {
         get_kv_callback(errors::Cancelled(
@@ -504,7 +506,7 @@ void CoordinationServiceStandaloneImpl::Stop(bool shut_staleness_thread) {
     get_cb_.clear();
   }
   {
-    mutex_lock l(state_mu_);
+    absl::MutexLock l(&state_mu_);
     for (auto& [barrier_id, barrier] : barriers_) {
       if (!barrier.passed) {
         absl::Status error = MakeCoordinationError(errors::Aborted(absl::StrCat(
@@ -519,9 +521,9 @@ void CoordinationServiceStandaloneImpl::Stop(bool shut_staleness_thread) {
     cluster_state_.clear();
   }
   {
-    mutex_lock l(check_staleness_thread_shutdown_mu_);
+    absl::MutexLock l(&check_staleness_thread_shutdown_mu_);
     shutting_down_ = true;
-    check_staleness_thread_cv_.notify_all();
+    check_staleness_thread_cv_.SignalAll();
   }
   if (shut_staleness_thread) {
     check_staleness_thread_.reset();
@@ -559,7 +561,7 @@ absl::Status CoordinationServiceStandaloneImpl::RegisterTask(
   absl::Status error;
   std::string error_message;
   {
-    mutex_lock l(state_mu_);
+    absl::MutexLock l(&state_mu_);
     if (!cluster_state_.contains(task_name)) {
       // Note: return early here as unexpected task register errors should not
       // be propagated to other tasks.
@@ -629,7 +631,7 @@ void CoordinationServiceStandaloneImpl::WaitForAllTasks(
     const CoordinatedTask& task, const DeviceInfo& devices,
     StatusCallback done) {
   {
-    mutex_lock l(state_mu_);
+    absl::MutexLock l(&state_mu_);
     const auto& task_state = cluster_state_.find(GetTaskName(task));
     // Collect task device info for the first time that task
     // has called WaitForAllTasks(). This will be aggregated when the barrier
@@ -652,7 +654,7 @@ void CoordinationServiceStandaloneImpl::ShutdownTaskAsync(
   } else {
     absl::Status status;
     {
-      mutex_lock l(state_mu_);
+      absl::MutexLock l(&state_mu_);
       // Disconnect task from service individually.
       status = DisconnectTask(task);
     }
@@ -662,7 +664,7 @@ void CoordinationServiceStandaloneImpl::ShutdownTaskAsync(
 
 absl::Status CoordinationServiceStandaloneImpl::ResetTask(
     const CoordinatedTask& task) {
-  mutex_lock l(state_mu_);
+  absl::MutexLock l(&state_mu_);
   return DisconnectTask(task);
 }
 
@@ -706,7 +708,7 @@ absl::Status CoordinationServiceStandaloneImpl::ReportTaskError(
     const CoordinatedTask& task, absl::Status error) {
   const std::string task_name = GetTaskName(task);
   {
-    mutex_lock l(state_mu_);
+    absl::MutexLock l(&state_mu_);
     if (!cluster_state_.contains(task_name)) {
       return MakeCoordinationError(
           errors::InvalidArgument("Unexpected request from task ", task_name));
@@ -731,7 +733,7 @@ CoordinationServiceStandaloneImpl::GetTaskState(
     auto& state_info = states_info.emplace_back();
     absl::Status error;
     {
-      mutex_lock l(state_mu_);
+      absl::MutexLock l(&state_mu_);
       state_info.set_state(cluster_state_[task_name]->GetState());
       error = cluster_state_[task_name]->GetStatus();
     }
@@ -751,7 +753,7 @@ absl::Status CoordinationServiceStandaloneImpl::RecordHeartbeat(
   const std::string task_name = GetTaskName(task);
   absl::Status s = absl::OkStatus();
   {
-    mutex_lock l(state_mu_);
+    absl::MutexLock l(&state_mu_);
     if (!cluster_state_.contains(task_name)) {
       return MakeCoordinationError(errors::InvalidArgument(
           "Unexpected heartbeat request from task: ", task_name,
@@ -779,7 +781,7 @@ absl::Status CoordinationServiceStandaloneImpl::RecordHeartbeat(
   // Set and propagate any heartbeat errors.
   if (!s.ok()) {
     {
-      mutex_lock l(state_mu_);
+      absl::MutexLock l(&state_mu_);
       SetTaskError(task_name, s);
     }
     PropagateError(task);
@@ -827,7 +829,7 @@ void CoordinationServiceStandaloneImpl::PropagateError(
   if (isRecoverableJob(source_task.job_name())) return;
   absl::Status error;
   {
-    mutex_lock l(state_mu_);
+    absl::MutexLock l(&state_mu_);
     error = cluster_state_[GetTaskName(source_task)]->GetStatus();
   }
   assert(!error.ok());
@@ -843,7 +845,7 @@ void CoordinationServiceStandaloneImpl::PropagateError(
 
   std::vector<absl::string_view> task_names;
   {
-    tf_shared_lock l(state_mu_);
+    absl::ReaderMutexLock l(&state_mu_);
     task_names.reserve(cluster_state_.size());
     for (const auto& pair : cluster_state_) {
       task_names.emplace_back(pair.first);
@@ -851,7 +853,7 @@ void CoordinationServiceStandaloneImpl::PropagateError(
   }
   for (absl::string_view task : task_names) {
     {
-      mutex_lock l(state_mu_);
+      absl::MutexLock l(&state_mu_);
       // Propagate error only to tasks that are connected
       if (cluster_state_[task]->GetState() !=
           CoordinatedTaskState::TASKSTATE_CONNECTED)
@@ -890,7 +892,7 @@ void CoordinationServiceStandaloneImpl::PropagateError(
 // The normalized key will not have leading or trailing slashes, and all parts
 // in the key path are separated by exactly one slack ('/').
 // E.g., ///a//b/c// --> a/b/c
-std::string NormalizeKey(const StringPiece orig_key) {
+std::string NormalizeKey(absl::string_view orig_key) {
   std::string norm_key = std::string(orig_key);
   const char* src = norm_key.c_str();
   std::string::iterator dst = norm_key.begin();
@@ -918,7 +920,7 @@ absl::Status CoordinationServiceStandaloneImpl::InsertKeyValue(
     const std::string& key, const std::string& value) {
   VLOG(3) << "InsertKeyValue(): " << key << ": " << value;
   const std::string norm_key = NormalizeKey(key);
-  mutex_lock l(kv_mu_);
+  absl::MutexLock l(&kv_mu_);
   if (kv_store_.find(norm_key) != kv_store_.end()) {
     return MakeCoordinationError(
         errors::AlreadyExists("Config key ", key, " already exists."));
@@ -938,7 +940,7 @@ void CoordinationServiceStandaloneImpl::GetKeyValueAsync(
     const std::string& key, StatusOrValueCallback done) {
   VLOG(3) << "GetKeyValue(): " << key;
   const std::string norm_key = NormalizeKey(key);
-  mutex_lock l(kv_mu_);
+  absl::MutexLock l(&kv_mu_);
   const auto& iter = kv_store_.find(norm_key);
   if (iter != kv_store_.end()) {
     done(iter->second);
@@ -956,7 +958,7 @@ absl::StatusOr<std::string> CoordinationServiceStandaloneImpl::TryGetKeyValue(
     const std::string& key) {
   VLOG(3) << "TryGetKeyValue(): " << key;
   const std::string norm_key = NormalizeKey(key);
-  mutex_lock l(kv_mu_);
+  absl::MutexLock l(&kv_mu_);
   const auto& iter = kv_store_.find(norm_key);
   if (iter == kv_store_.end()) {
     return errors::NotFound("Config key ", key, " not found.");
@@ -971,7 +973,7 @@ std::vector<KeyValueEntry> CoordinationServiceStandaloneImpl::GetKeyValueDir(
   const std::string norm_key = NormalizeKey(directory_key);
   const std::string dir = absl::StrCat(norm_key, "/");
 
-  mutex_lock l(kv_mu_);
+  absl::MutexLock l(&kv_mu_);
   // Find first key in ordered map that has the directory prefix.
   auto begin = kv_store_.lower_bound(dir);
   std::map<std::string, std::string>::iterator it;
@@ -996,7 +998,7 @@ absl::Status CoordinationServiceStandaloneImpl::DeleteKeyValue(
     const std::string& key) {
   VLOG(3) << "DeleteKeyValue(): " << key;
   const std::string norm_key = NormalizeKey(key);
-  mutex_lock l(kv_mu_);
+  absl::MutexLock l(&kv_mu_);
   // Delete directory: find key range that match directory prefix
   const std::string dir = absl::StrCat(norm_key, "/");
   auto begin = kv_store_.lower_bound(dir);
@@ -1054,7 +1056,7 @@ void CoordinationServiceStandaloneImpl::BarrierAsync(
         absl::StrCat("A non-participating task (", GetTaskName(task),
                      ") called the barrier: ", barrier_id)));
     {
-      mutex_lock l(state_mu_);
+      absl::MutexLock l(&state_mu_);
       // Check if coordination service has stopped. If so, return an error
       // immediately.
       if (ServiceHasStopped()) {
@@ -1072,7 +1074,7 @@ void CoordinationServiceStandaloneImpl::BarrierAsync(
     done(error);
     return;
   }
-  mutex_lock l(state_mu_);
+  absl::MutexLock l(&state_mu_);
   // Check if coordination service has stopped. If so, return an error
   // immediately.
   if (ServiceHasStopped()) {
@@ -1191,7 +1193,7 @@ void CoordinationServiceStandaloneImpl::BarrierAsync(
 
 absl::Status CoordinationServiceStandaloneImpl::CancelBarrier(
     const std::string& barrier_id, const CoordinatedTask& task) {
-  mutex_lock l(state_mu_);
+  absl::MutexLock l(&state_mu_);
   auto [it, inserted] = barriers_.try_emplace(barrier_id);
   auto* barrier = &it->second;
   if (inserted) {
diff --git a/third_party/xla/xla/tsl/distributed_runtime/coordination/coordination_service.h b/third_party/xla/xla/tsl/distributed_runtime/coordination/coordination_service.h
index 29abe4cbb6aaaa..fbf3206c5e95d3 100644
--- a/third_party/xla/xla/tsl/distributed_runtime/coordination/coordination_service.h
+++ b/third_party/xla/xla/tsl/distributed_runtime/coordination/coordination_service.h
@@ -16,6 +16,7 @@ limitations under the License.
 #ifndef XLA_TSL_DISTRIBUTED_RUNTIME_COORDINATION_COORDINATION_SERVICE_H_
 #define XLA_TSL_DISTRIBUTED_RUNTIME_COORDINATION_COORDINATION_SERVICE_H_
 
+#include <cstdint>
 #include <functional>
 #include <memory>
 #include <string>
@@ -23,12 +24,14 @@ limitations under the License.
 #include <utility>
 #include <vector>
 
+#include "absl/log/log.h"
 #include "absl/status/status.h"
+#include "absl/status/statusor.h"
 #include "absl/strings/string_view.h"
 #include "absl/time/time.h"
 #include "xla/tsl/distributed_runtime/coordination/coordination_client.h"
+#include "tsl/platform/macros.h"
 #include "tsl/platform/status.h"
-#include "tsl/platform/statusor.h"
 #include "tsl/protobuf/coordination_config.pb.h"
 
 namespace tsl {
@@ -51,7 +54,7 @@ class Env;
 // execution in a cluster of multiple tasks.
 //
 // When enabled, the service keeps track of cluster configurations and the state
-// of cluster members. TF runtime and libraries can use it to orchastrate
+// of cluster members. TF runtime and libraries can use it to orchestrate
 // cluster initialization, check the healthiness of tasks, and propagate error
 // messages to the cluster.
 //
diff --git a/third_party/xla/xla/tsl/distributed_runtime/coordination/coordination_service_agent.cc b/third_party/xla/xla/tsl/distributed_runtime/coordination/coordination_service_agent.cc
index 1d090c7a0ffbc7..860d4fd65ab90b 100644
--- a/third_party/xla/xla/tsl/distributed_runtime/coordination/coordination_service_agent.cc
+++ b/third_party/xla/xla/tsl/distributed_runtime/coordination/coordination_service_agent.cc
@@ -33,6 +33,7 @@ limitations under the License.
 #include "absl/status/status.h"
 #include "absl/strings/str_cat.h"
 #include "absl/strings/substitute.h"
+#include "absl/synchronization/mutex.h"
 #include "absl/synchronization/notification.h"
 #include "absl/time/clock.h"
 #include "absl/time/time.h"
@@ -42,7 +43,6 @@ limitations under the License.
 #include "tsl/framework/cancellation.h"
 #include "tsl/lib/monitoring/gauge.h"
 #include "tsl/platform/env.h"
-#include "tsl/platform/mutex.h"
 #include "tsl/platform/random.h"
 #include "tsl/platform/status.h"
 #include "tsl/platform/thread_annotations.h"
@@ -145,7 +145,7 @@ class CoordinationServiceAgentImpl : public CoordinationServiceAgent {
   CoordinationServiceConfig configs_;
   StatusCallback error_fn_;
 
-  mutable mutex state_mu_;
+  mutable absl::Mutex state_mu_;
   CoordinatedTaskState state_ TF_GUARDED_BY(state_mu_) =
       CoordinatedTaskState::TASKSTATE_UNINITIALIZED;
   absl::Status status_ TF_GUARDED_BY(state_mu_) = absl::OkStatus();
@@ -157,8 +157,8 @@ class CoordinationServiceAgentImpl : public CoordinationServiceAgent {
   uint64_t leader_incarnation_ = 0;
   DeviceInfo cluster_devices_;
 
-  mutex heartbeat_thread_shutdown_mu_;
-  condition_variable heartbeat_thread_cv_;
+  absl::Mutex heartbeat_thread_shutdown_mu_;
+  absl::CondVar heartbeat_thread_cv_;
   bool shutting_down_ TF_GUARDED_BY(heartbeat_thread_shutdown_mu_) = false;
   std::unique_ptr<Thread> heartbeat_thread_;
   // Must outlive coordination client which may need to access it within
@@ -187,7 +187,7 @@ absl::Status CoordinationServiceAgentImpl::Initialize(
     std::unique_ptr<CoordinationClient> leader_client,
     StatusCallback error_fn) {
   enabled_usage_metric->GetCell()->Set(true);
-  mutex_lock l(state_mu_);
+  absl::MutexLock l(&state_mu_);
   if (state_ != CoordinatedTaskState::TASKSTATE_UNINITIALIZED) {
     return MakeCoordinationError(absl::FailedPreconditionError(
         "Coordination service agent has already been initialized."));
@@ -211,25 +211,25 @@ absl::Status CoordinationServiceAgentImpl::Initialize(
 }
 
 bool CoordinationServiceAgentImpl::IsInitialized() {
-  mutex_lock l(state_mu_);
+  absl::MutexLock l(&state_mu_);
   return state_ != CoordinatedTaskState::TASKSTATE_UNINITIALIZED;
 }
 
 bool CoordinationServiceAgentImpl::IsConnected() {
-  mutex_lock l(state_mu_);
+  absl::MutexLock l(&state_mu_);
   return state_ == CoordinatedTaskState::TASKSTATE_CONNECTED;
 }
 
 bool CoordinationServiceAgentImpl::IsError() {
-  mutex_lock l(state_mu_);
+  absl::MutexLock l(&state_mu_);
   return state_ == CoordinatedTaskState::TASKSTATE_ERROR;
 }
 
 void CoordinationServiceAgentImpl::StopHeartbeat() {
   {
-    mutex_lock l(heartbeat_thread_shutdown_mu_);
+    absl::MutexLock l(&heartbeat_thread_shutdown_mu_);
     shutting_down_ = true;
-    heartbeat_thread_cv_.notify_all();
+    heartbeat_thread_cv_.SignalAll();
   }
   heartbeat_thread_ = nullptr;
 }
@@ -237,7 +237,7 @@ void CoordinationServiceAgentImpl::StopHeartbeat() {
 absl::Status CoordinationServiceAgentImpl::Connect() {
   VLOG(3) << "Agent has started trying to Connect().";
   {
-    mutex_lock l(state_mu_);
+    absl::MutexLock l(&state_mu_);
     if (state_ != CoordinatedTaskState::TASKSTATE_DISCONNECTED) {
       return MakeCoordinationError(absl::FailedPreconditionError(
           "Coordination service agent is not in DISCONNECTED state."));
@@ -270,7 +270,7 @@ absl::Status CoordinationServiceAgentImpl::Connect() {
           if (s.ok()) {
             leader_incarnation_ = response.leader_incarnation();
             {
-              mutex_lock l(state_mu_);
+              absl::MutexLock l(&state_mu_);
               state_ = CoordinatedTaskState::TASKSTATE_CONNECTED;
             }
           }
@@ -341,7 +341,7 @@ absl::Status CoordinationServiceAgentImpl::Connect() {
             // inflight heartbeats sent during shutdown and can be ignored.
             absl::SleepFor(absl::Seconds(1));
             {
-              mutex_lock l(heartbeat_thread_shutdown_mu_);
+              absl::MutexLock l(&heartbeat_thread_shutdown_mu_);
 
               if (shutting_down_) {
                 return;
@@ -355,11 +355,10 @@ absl::Status CoordinationServiceAgentImpl::Connect() {
           }
           // Send next heartbeat after an interval.
           {
-            mutex_lock l(heartbeat_thread_shutdown_mu_);
-            // TODO(b/339231167): Fix the lint.
-            heartbeat_thread_cv_.wait_for(
-                l, std::chrono::milliseconds(  // NOLINT(misc-include-cleaner)
-                       heartbeat_interval_ms));
+            absl::MutexLock l(&heartbeat_thread_shutdown_mu_);
+            heartbeat_thread_cv_.WaitWithTimeout(
+                &heartbeat_thread_shutdown_mu_,
+                absl::Milliseconds(heartbeat_interval_ms));
             if (shutting_down_) {
               return;
             }
@@ -437,7 +436,7 @@ CoordinationServiceAgentImpl::GetTaskState(
 absl::Status CoordinationServiceAgentImpl::ReportError(
     const absl::Status& error) {
   {
-    mutex_lock l(state_mu_);
+    absl::MutexLock l(&state_mu_);
     if (state_ == CoordinatedTaskState::TASKSTATE_UNINITIALIZED) {
       return MakeCoordinationError(absl::FailedPreconditionError(
           "Coordination service agent must be initialized first before "
@@ -484,7 +483,7 @@ absl::Status CoordinationServiceAgentImpl::ShutdownInternal() {
   absl::Status status = absl::OkStatus();
   bool is_connected = false;
   {
-    mutex_lock l(state_mu_);
+    absl::MutexLock l(&state_mu_);
     is_connected = state_ == CoordinatedTaskState::TASKSTATE_CONNECTED;
   }
   // Disconnect agent from service.
@@ -522,7 +521,7 @@ absl::Status CoordinationServiceAgentImpl::ShutdownInternal() {
   // Tear down agent.
   StopHeartbeat();
   {
-    mutex_lock l(state_mu_);
+    absl::MutexLock l(&state_mu_);
     if (state_ == CoordinatedTaskState::TASKSTATE_ERROR) {
       const std::string status_message = absl::StrCat(
           "Shutdown() was called while coordination agent is in error state, "
@@ -546,7 +545,7 @@ absl::Status CoordinationServiceAgentImpl::ShutdownInternal() {
 
 absl::Status CoordinationServiceAgentImpl::Reset() {
   {
-    mutex_lock l(state_mu_);
+    absl::MutexLock l(&state_mu_);
     if (state_ != CoordinatedTaskState::TASKSTATE_ERROR) {
       return MakeCoordinationError(absl::FailedPreconditionError(
           "Reset() failed: coordination service agent is not in ERROR state."));
@@ -574,11 +573,11 @@ absl::Status CoordinationServiceAgentImpl::Reset() {
   // Reset agent state.
   StopHeartbeat();
   {
-    mutex_lock l(state_mu_);
+    absl::MutexLock l(&state_mu_);
     state_ = CoordinatedTaskState::TASKSTATE_DISCONNECTED;
   }
   {
-    mutex_lock l(heartbeat_thread_shutdown_mu_);
+    absl::MutexLock l(&heartbeat_thread_shutdown_mu_);
     shutting_down_ = false;
   }
 
@@ -766,7 +765,7 @@ absl::Status CoordinationServiceAgentImpl::StopWatchKey(std::string_view key) {
 
 void CoordinationServiceAgentImpl::SetError(const absl::Status& error) {
   assert(!error.ok());
-  mutex_lock l(state_mu_);
+  absl::MutexLock l(&state_mu_);
   if (state_ == CoordinatedTaskState::TASKSTATE_ERROR) return;
 
   LOG(ERROR) << "Coordination agent is set to ERROR: " << error;
@@ -804,7 +803,7 @@ void CoordinationServiceAgentImpl::WaitAtBarrierAsync(
     return;
   }
   {
-    mutex_lock l(state_mu_);
+    absl::MutexLock l(&state_mu_);
     auto [it, inserted] = used_barrier_ids_.insert(std::string(barrier_id));
     if (!inserted) {
       done(absl::FailedPreconditionError(absl::StrCat(
@@ -865,7 +864,7 @@ void CoordinationServiceAgentImpl::CancelBarrierAsync(
 // Returns an error if agent is not running.
 absl::Status CoordinationServiceAgentImpl::ValidateRunningAgent(
     bool allow_disconnected) {
-  mutex_lock l(state_mu_);
+  absl::MutexLock l(&state_mu_);
   switch (state_) {
     case CoordinatedTaskState::TASKSTATE_CONNECTED:
       return absl::OkStatus();
diff --git a/third_party/xla/xla/tsl/distributed_runtime/coordination/coordination_service_agent.h b/third_party/xla/xla/tsl/distributed_runtime/coordination/coordination_service_agent.h
index 0f854ae42a101b..418ca50b0db122 100644
--- a/third_party/xla/xla/tsl/distributed_runtime/coordination/coordination_service_agent.h
+++ b/third_party/xla/xla/tsl/distributed_runtime/coordination/coordination_service_agent.h
@@ -16,20 +16,19 @@ limitations under the License.
 #ifndef XLA_TSL_DISTRIBUTED_RUNTIME_COORDINATION_COORDINATION_SERVICE_AGENT_H_
 #define XLA_TSL_DISTRIBUTED_RUNTIME_COORDINATION_COORDINATION_SERVICE_AGENT_H_
 
-#include <cstdint>
 #include <functional>
 #include <map>
 #include <memory>
 #include <string>
 #include <string_view>
-#include <utility>
 #include <vector>
 
+#include "absl/status/status.h"
+#include "absl/status/statusor.h"
 #include "absl/time/time.h"
 #include "xla/tsl/distributed_runtime/call_options.h"
 #include "xla/tsl/distributed_runtime/coordination/coordination_client.h"
 #include "tsl/platform/status.h"
-#include "tsl/platform/statusor.h"
 #include "tsl/protobuf/coordination_service.pb.h"
 
 namespace tensorflow {

From 4650838ce615ff4b2d2a74882cbeb0add745a82e Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Mon, 13 May 2024 17:18:54 -0700
Subject: [PATCH 080/478] Update ops-related pbtxt files.

PiperOrigin-RevId: 633379311
---
 tensorflow/core/ops/ops.pbtxt | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorflow/core/ops/ops.pbtxt b/tensorflow/core/ops/ops.pbtxt
index fb2f485604495b..7322cef85415f5 100644
--- a/tensorflow/core/ops/ops.pbtxt
+++ b/tensorflow/core/ops/ops.pbtxt
@@ -1,4 +1,4 @@
-go/nodeserialize  
+go/debugstr   
 op {
   name: "Abort"
   attr {

From 4c5ee99fb2f6c23aae35883dbc9dcf8e648d85d6 Mon Sep 17 00:00:00 2001
From: Kyle Lucke <klucke@google.com>
Date: Mon, 13 May 2024 17:23:26 -0700
Subject: [PATCH 081/478] Move Trace CommandBuffer creation into
 TraceCommandBufferFactory to remove circular dependencies.

PiperOrigin-RevId: 633380264
---
 third_party/xla/xla/service/gpu/runtime/BUILD |  1 +
 .../service/gpu/runtime/command_buffer_cmd.cc | 17 +++---
 third_party/xla/xla/stream_executor/BUILD     | 19 ++++++-
 .../xla/xla/stream_executor/command_buffer.h  | 31 +---------
 third_party/xla/xla/stream_executor/gpu/BUILD |  1 +
 .../gpu/gpu_command_buffer_test.cc            |  6 +-
 ...fer.cc => trace_command_buffer_factory.cc} | 22 +++++---
 .../trace_command_buffer_factory.h            | 56 +++++++++++++++++++
 8 files changed, 102 insertions(+), 51 deletions(-)
 rename third_party/xla/xla/stream_executor/{command_buffer.cc => trace_command_buffer_factory.cc} (74%)
 create mode 100644 third_party/xla/xla/stream_executor/trace_command_buffer_factory.h

diff --git a/third_party/xla/xla/service/gpu/runtime/BUILD b/third_party/xla/xla/service/gpu/runtime/BUILD
index d0b940e88b523e..c259a044cb8308 100644
--- a/third_party/xla/xla/service/gpu/runtime/BUILD
+++ b/third_party/xla/xla/service/gpu/runtime/BUILD
@@ -85,6 +85,7 @@ cc_library(
         "//xla/service/gpu/kernels:custom_kernel",
         "//xla/service/gpu/runtime:thunk",
         "//xla/stream_executor",
+        "//xla/stream_executor:trace_command_buffer_factory",
         "//xla/stream_executor/gpu:gpu_stream_header",
         "//xla/stream_executor/gpu:gpu_types_header",
         "//xla/tsl/concurrency:ref_count",
diff --git a/third_party/xla/xla/service/gpu/runtime/command_buffer_cmd.cc b/third_party/xla/xla/service/gpu/runtime/command_buffer_cmd.cc
index b415fdd14e15a5..5a6263fe3ef557 100644
--- a/third_party/xla/xla/service/gpu/runtime/command_buffer_cmd.cc
+++ b/third_party/xla/xla/service/gpu/runtime/command_buffer_cmd.cc
@@ -66,6 +66,7 @@ limitations under the License.
 #include "xla/stream_executor/launch_dim.h"
 #include "xla/stream_executor/stream.h"
 #include "xla/stream_executor/stream_executor.h"
+#include "xla/stream_executor/trace_command_buffer_factory.h"
 #include "xla/tsl/concurrency/ref_count.h"
 #include "xla/types.h"  // IWYU pragma: keep
 #include "xla/util.h"
@@ -397,8 +398,9 @@ absl::StatusOr<se::CommandBuffer*> TracedCommandBuffer::GetOrTraceCommandBuffer(
     // Create a new entry by calling a user-provided tracing function, move it
     // to front and return a pointer to cached command buffer.
     if (entries_[i].command_buffer == nullptr) {
-      TF_ASSIGN_OR_RETURN(entries_[i].command_buffer,
-                          se::CommandBuffer::Trace(executor, stream, trace));
+      TF_ASSIGN_OR_RETURN(
+          entries_[i].command_buffer,
+          se::TraceCommandBufferFactory::Create(executor, stream, trace));
       entries_[i].recorded_allocs.assign(allocs.begin(), allocs.end());
       return shift_right(i).command_buffer.get();
     }
@@ -407,8 +409,9 @@ absl::StatusOr<se::CommandBuffer*> TracedCommandBuffer::GetOrTraceCommandBuffer(
   // Create a new entry by calling a user-provided tracing function, replace the
   // last entry with it, move it to front and return a pointer to cached command
   // buffer.
-  TF_ASSIGN_OR_RETURN(entries_[capacity_ - 1].command_buffer,
-                      se::CommandBuffer::Trace(executor, stream, trace));
+  TF_ASSIGN_OR_RETURN(
+      entries_[capacity_ - 1].command_buffer,
+      se::TraceCommandBufferFactory::Create(executor, stream, trace));
   entries_[capacity_ - 1].recorded_allocs.assign(allocs.begin(), allocs.end());
   return shift_right(capacity_ - 1).command_buffer.get();
 }
@@ -1209,7 +1212,7 @@ absl::Status CustomCallCmd::RecordLegacyCustomCall(
 #if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
   TF_ASSIGN_OR_RETURN(
       auto nested_cmd,
-      se::CommandBuffer::Trace(
+      se::TraceCommandBufferFactory::Create(
           execute_params.stream->parent(),
           execute_params.command_buffer_trace_stream, [&](se::Stream* stream) {
             se::gpu::GpuStreamHandle gpu_stream =
@@ -1289,7 +1292,7 @@ absl::Status CustomCallCmd::RecordXlaFfiCall(
 #if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
   TF_ASSIGN_OR_RETURN(
       auto nested_cmd,
-      se::CommandBuffer::Trace(
+      se::TraceCommandBufferFactory::Create(
           execute_params.stream->parent(),
           execute_params.command_buffer_trace_stream, [&](se::Stream* stream) {
             ExecutableRunOptions run_options;
@@ -1406,7 +1409,7 @@ absl::Status CollectiveCmd::AddTracedCommandBuffer(
     const RecordParams& record_params, se::CommandBuffer* command_buffer,
     absl::FunctionRef<absl::Status(se::Stream*)> trace) {
   TF_ASSIGN_OR_RETURN(std::unique_ptr<se::CommandBuffer> nested_cmd,
-                      se::CommandBuffer::Trace(
+                      se::TraceCommandBufferFactory::Create(
                           execute_params.stream->parent(),
                           execute_params.command_buffer_trace_stream, trace));
 
diff --git a/third_party/xla/xla/stream_executor/BUILD b/third_party/xla/xla/stream_executor/BUILD
index eb14242b323b00..b73e82b9e42607 100644
--- a/third_party/xla/xla/stream_executor/BUILD
+++ b/third_party/xla/xla/stream_executor/BUILD
@@ -521,20 +521,33 @@ cc_library(
 
 cc_library(
     name = "command_buffer",
-    srcs = ["command_buffer.cc"],
     hdrs = ["command_buffer.h"],
     visibility = ["//visibility:private"],
     deps = [
         ":device_memory",
+        ":kernel",
         ":launch_dim",
         ":platform",
+        "@com_google_absl//absl/functional:any_invocable",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/types:span",
+        "@local_tsl//tsl/lib/gtl:int_type",
+        "@local_tsl//tsl/platform:errors",
+    ],
+)
+
+cc_library(
+    name = "trace_command_buffer_factory",
+    srcs = ["trace_command_buffer_factory.cc"],
+    hdrs = ["trace_command_buffer_factory.h"],
+    visibility = [":internal"],
+    deps = [
+        ":command_buffer",
         ":stream_executor_headers",
         ":stream_executor_interface",
         "@com_google_absl//absl/functional:any_invocable",
         "@com_google_absl//absl/status",
         "@com_google_absl//absl/status:statusor",
-        "@com_google_absl//absl/types:span",
-        "@local_tsl//tsl/lib/gtl:int_type",
         "@local_tsl//tsl/platform:errors",
         "@local_tsl//tsl/platform:statusor",
     ],
diff --git a/third_party/xla/xla/stream_executor/command_buffer.h b/third_party/xla/xla/stream_executor/command_buffer.h
index a91963d46a0880..5cb39e857f7fbb 100644
--- a/third_party/xla/xla/stream_executor/command_buffer.h
+++ b/third_party/xla/xla/stream_executor/command_buffer.h
@@ -19,25 +19,21 @@ limitations under the License.
 #include <cstddef>
 #include <cstdint>
 #include <functional>
-#include <memory>
 #include <variant>
 #include <vector>
 
 #include "absl/functional/any_invocable.h"
 #include "absl/status/status.h"
-#include "absl/status/statusor.h"
 #include "absl/types/span.h"
 #include "xla/stream_executor/device_memory.h"
 #include "xla/stream_executor/kernel.h"
 #include "xla/stream_executor/launch_dim.h"
-#include "xla/stream_executor/platform.h"
 #include "tsl/lib/gtl/int_type.h"
 #include "tsl/platform/errors.h"
 
 namespace stream_executor {
 
 class Stream;
-class StreamExecutorInterface;
 
 //===----------------------------------------------------------------------===//
 // CommandBuffer
@@ -161,32 +157,6 @@ class CommandBuffer {
   //
   enum class Mode { kPrimary, kNested };
 
-  //===--------------------------------------------------------------------===//
-  // Command buffer constructors
-  //===--------------------------------------------------------------------===//
-
-  // Creates a new command buffer on the given executor by tracing `function`
-  // invocation. All StreamExecutor operations on a Stream argument will be
-  // recorded into the command buffer. Returned command buffer is finalized, and
-  // can't be updated.
-  //
-  // Command buffer tracing should be used only when it is impossible to use
-  // explicit construction APIs, e.g. when calling external libraries. By
-  // default we construct traced command buffers in nested mode because the
-  // primary use case for traced command buffers is to be inserted into primary
-  // command buffers constructed with explicit APIs.
-  static absl::StatusOr<std::unique_ptr<CommandBuffer>> Trace(
-      StreamExecutorInterface* executor,
-      absl::AnyInvocable<absl::Status(Stream*)> function,
-      Mode mode = Mode::kNested);
-
-  // Creates a new command buffer on the given executor by tracing `function`
-  // invocation using a user provided stream that will be passed to `function`.
-  static absl::StatusOr<std::unique_ptr<CommandBuffer>> Trace(
-      StreamExecutorInterface* executor, Stream* stream,
-      absl::AnyInvocable<absl::Status(Stream*)> function,
-      Mode mode = Mode::kNested);
-
   //===--------------------------------------------------------------------===//
   // Command buffer API
   //===--------------------------------------------------------------------===//
@@ -379,6 +349,7 @@ class CommandBuffer {
   // Command buffer tracing API
   //--------------------------------------------------------------------------//
  private:
+  friend class TraceCommandBufferFactory;
   // Tracing APIs are private because they do not compose with command buffer
   // updates. Instead of tracing directly into the command buffer users should
   // create traced command buffers using factory methods and add them to primary
diff --git a/third_party/xla/xla/stream_executor/gpu/BUILD b/third_party/xla/xla/stream_executor/gpu/BUILD
index 48e951e0ed9c99..19b33618b223d3 100644
--- a/third_party/xla/xla/stream_executor/gpu/BUILD
+++ b/third_party/xla/xla/stream_executor/gpu/BUILD
@@ -694,6 +694,7 @@ xla_test(
         "//xla/stream_executor",
         "//xla/stream_executor:platform",
         "//xla/stream_executor:platform_manager",
+        "//xla/stream_executor:trace_command_buffer_factory",
         "//xla/stream_executor/gpu:gpu_driver_header",
         "@com_google_absl//absl/log:check",
         "@com_google_absl//absl/status",
diff --git a/third_party/xla/xla/stream_executor/gpu/gpu_command_buffer_test.cc b/third_party/xla/xla/stream_executor/gpu/gpu_command_buffer_test.cc
index bd50ac29429b0e..d1999c819f9927 100644
--- a/third_party/xla/xla/stream_executor/gpu/gpu_command_buffer_test.cc
+++ b/third_party/xla/xla/stream_executor/gpu/gpu_command_buffer_test.cc
@@ -35,6 +35,7 @@ limitations under the License.
 #include "xla/stream_executor/platform_manager.h"
 #include "xla/stream_executor/stream.h"
 #include "xla/stream_executor/stream_executor.h"
+#include "xla/stream_executor/trace_command_buffer_factory.h"
 #include "tsl/lib/core/status_test_util.h"
 #include "tsl/platform/errors.h"
 #include "tsl/platform/status.h"
@@ -199,7 +200,7 @@ TEST(CudaCommandBufferTest, TraceSingleKernel) {
   KernelArgsDeviceMemoryArray args({a, b, c}, 0);
 
   // Create a command buffer by tracing kernel launch operations.
-  auto cmd_buffer = CommandBuffer::Trace(
+  auto cmd_buffer = TraceCommandBufferFactory::Create(
       executor,
       [&](Stream* stream) {
         return executor->Launch(stream, ThreadDim(), BlockDim(4), *add, args);
@@ -1320,7 +1321,8 @@ static void BM_TraceCommandBuffer(benchmark::State& state) {
       return absl::OkStatus();
     };
 
-    CHECK_OK(CommandBuffer::Trace(executor, launch_kernels, nested));
+    CHECK_OK(
+        TraceCommandBufferFactory::Create(executor, launch_kernels, nested));
   }
 }
 
diff --git a/third_party/xla/xla/stream_executor/command_buffer.cc b/third_party/xla/xla/stream_executor/trace_command_buffer_factory.cc
similarity index 74%
rename from third_party/xla/xla/stream_executor/command_buffer.cc
rename to third_party/xla/xla/stream_executor/trace_command_buffer_factory.cc
index cd28a135bf519c..6e66b897e32c28 100644
--- a/third_party/xla/xla/stream_executor/command_buffer.cc
+++ b/third_party/xla/xla/stream_executor/trace_command_buffer_factory.cc
@@ -1,4 +1,4 @@
-/* Copyright 2023 The OpenXLA Authors.
+/* Copyright 2024 The OpenXLA Authors.
 
 Licensed under the Apache License, Version 2.0 (the "License");
 you may not use this file except in compliance with the License.
@@ -13,7 +13,7 @@ See the License for the specific language governing permissions and
 limitations under the License.
 ==============================================================================*/
 
-#include "xla/stream_executor/command_buffer.h"
+#include "xla/stream_executor/trace_command_buffer_factory.h"
 
 #include <memory>
 #include <utility>
@@ -21,24 +21,28 @@ limitations under the License.
 #include "absl/functional/any_invocable.h"
 #include "absl/status/status.h"
 #include "absl/status/statusor.h"
-#include "xla/stream_executor/kernel.h"
-#include "xla/stream_executor/stream.h"
+#include "xla/stream_executor/command_buffer.h"
 #include "xla/stream_executor/stream_executor_interface.h"
 #include "tsl/platform/errors.h"
 #include "tsl/platform/statusor.h"
 
 namespace stream_executor {
 
-absl::StatusOr<std::unique_ptr<CommandBuffer>> CommandBuffer::Trace(
+absl::StatusOr<std::unique_ptr<CommandBuffer>>
+TraceCommandBufferFactory::Create(
     StreamExecutorInterface* executor,
-    absl::AnyInvocable<absl::Status(Stream*)> function, Mode mode) {
+    absl::AnyInvocable<absl::Status(Stream*)> function,
+    CommandBuffer::Mode mode) {
   TF_ASSIGN_OR_RETURN(auto stream, executor->CreateStream());
-  return Trace(executor, stream.get(), std::move(function), mode);
+  return TraceCommandBufferFactory::Create(executor, stream.get(),
+                                           std::move(function), mode);
 }
 
-absl::StatusOr<std::unique_ptr<CommandBuffer>> CommandBuffer::Trace(
+absl::StatusOr<std::unique_ptr<CommandBuffer>>
+TraceCommandBufferFactory::Create(
     StreamExecutorInterface* executor, Stream* stream,
-    absl::AnyInvocable<absl::Status(Stream*)> function, Mode mode) {
+    absl::AnyInvocable<absl::Status(Stream*)> function,
+    CommandBuffer::Mode mode) {
   if (stream == nullptr)
     return absl::InvalidArgumentError(
         "Can't trace command buffer on a null stream");
diff --git a/third_party/xla/xla/stream_executor/trace_command_buffer_factory.h b/third_party/xla/xla/stream_executor/trace_command_buffer_factory.h
new file mode 100644
index 00000000000000..9a88770c8d8b58
--- /dev/null
+++ b/third_party/xla/xla/stream_executor/trace_command_buffer_factory.h
@@ -0,0 +1,56 @@
+/* Copyright 2024 The OpenXLA Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+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.
+==============================================================================*/
+#ifndef XLA_STREAM_EXECUTOR_TRACE_COMMAND_BUFFER_FACTORY_H_
+#define XLA_STREAM_EXECUTOR_TRACE_COMMAND_BUFFER_FACTORY_H_
+
+#include <memory>
+
+#include "absl/functional/any_invocable.h"
+#include "absl/status/status.h"
+#include "absl/status/statusor.h"
+#include "xla/stream_executor/command_buffer.h"
+#include "xla/stream_executor/stream.h"
+#include "xla/stream_executor/stream_executor_interface.h"
+
+namespace stream_executor {
+
+class TraceCommandBufferFactory {
+ public:
+  // Creates a new command buffer on the given executor by tracing `function`
+  // invocation. All StreamExecutor operations on a Stream argument will be
+  // recorded into the command buffer. Returned command buffer is finalized, and
+  // can't be updated.
+  //
+  // Command buffer tracing should be used only when it is impossible to use
+  // explicit construction APIs, e.g. when calling external libraries. By
+  // default we construct traced command buffers in nested mode because the
+  // primary use case for traced command buffers is to be inserted into primary
+  // command buffers constructed with explicit APIs.
+  static absl::StatusOr<std::unique_ptr<CommandBuffer>> Create(
+      StreamExecutorInterface* executor,
+      absl::AnyInvocable<absl::Status(Stream*)> function,
+      CommandBuffer::Mode mode = CommandBuffer::Mode::kNested);
+
+  // Creates a new command buffer on the given executor by tracing `function`
+  // invocation using a user provided stream that will be passed to `function`.
+  static absl::StatusOr<std::unique_ptr<CommandBuffer>> Create(
+      StreamExecutorInterface* executor, Stream* stream,
+      absl::AnyInvocable<absl::Status(Stream*)> function,
+      CommandBuffer::Mode mode = CommandBuffer::Mode::kNested);
+};
+
+}  // namespace stream_executor
+
+#endif  // XLA_STREAM_EXECUTOR_TRACE_COMMAND_BUFFER_FACTORY_H_

From fa1ec11fc5277c5890d177cb91b8ed5878f5bbb9 Mon Sep 17 00:00:00 2001
From: Peter Hawkins <phawkins@google.com>
Date: Mon, 13 May 2024 18:19:47 -0700
Subject: [PATCH 082/478] [PJRT:IFRT] Remove overload of
 PjRtLoadedExecutable::Create() that takes an unique_ptr argument.

Cleanup only, NFC intended.

PiperOrigin-RevId: 633393546
---
 .../xla/xla/python/pjrt_ifrt/pjrt_compiler.cc    | 13 +++++++++++--
 .../xla/xla/python/pjrt_ifrt/pjrt_executable.cc  | 16 ----------------
 .../xla/xla/python/pjrt_ifrt/pjrt_executable.h   |  6 ------
 3 files changed, 11 insertions(+), 24 deletions(-)

diff --git a/third_party/xla/xla/python/pjrt_ifrt/pjrt_compiler.cc b/third_party/xla/xla/python/pjrt_ifrt/pjrt_compiler.cc
index 7245857e70cdb5..028e480b79af2c 100644
--- a/third_party/xla/xla/python/pjrt_ifrt/pjrt_compiler.cc
+++ b/third_party/xla/xla/python/pjrt_ifrt/pjrt_compiler.cc
@@ -20,11 +20,18 @@ limitations under the License.
 #include <utility>
 
 #include "absl/status/status.h"
+#include "absl/status/statusor.h"
+#include "absl/strings/string_view.h"
 #include "llvm/Support/Casting.h"
+#include "xla/pjrt/pjrt_client.h"
+#include "xla/python/ifrt/compiler.h"
+#include "xla/python/ifrt/executable.h"
 #include "xla/python/ifrt/hlo/hlo_program.h"
+#include "xla/python/ifrt/program.h"
 #include "xla/python/pjrt_ifrt/pjrt_client.h"
 #include "xla/python/pjrt_ifrt/pjrt_executable.h"
 #include "xla/python/pjrt_ifrt/xla_compiler.h"
+#include "tsl/platform/logging.h"
 #include "tsl/platform/statusor.h"
 
 namespace xla {
@@ -55,11 +62,13 @@ PjRtCompiler::DeserializeLoadedExecutable(
   TF_ASSIGN_OR_RETURN(auto xla_deserialize_options,
                       GetXlaDeserializeExecutableOptions(std::move(options)));
   TF_ASSIGN_OR_RETURN(
-      auto pjrt_loaded_executble,
+      auto pjrt_loaded_executable,
       client_->pjrt_client()->DeserializeExecutable(
           serialized, std::move(xla_deserialize_options->compile_options)));
   return PjRtLoadedExecutable::Create(
-      client_, std::move(pjrt_loaded_executble),
+      client_,
+      std::shared_ptr<xla::PjRtLoadedExecutable>(
+          std::move(pjrt_loaded_executable)),
       std::move(xla_deserialize_options->loaded_host_callbacks));
 }
 
diff --git a/third_party/xla/xla/python/pjrt_ifrt/pjrt_executable.cc b/third_party/xla/xla/python/pjrt_ifrt/pjrt_executable.cc
index 7ae786400cd474..bf2930516a61bb 100644
--- a/third_party/xla/xla/python/pjrt_ifrt/pjrt_executable.cc
+++ b/third_party/xla/xla/python/pjrt_ifrt/pjrt_executable.cc
@@ -181,12 +181,6 @@ char PjRtCompatibleLoadedExecutable::ID = 0;
 char PjRtExecutable::ID = 0;
 char PjRtLoadedExecutable::ID = 0;
 
-absl::StatusOr<std::unique_ptr<Executable>> PjRtExecutable::Create(
-    std::unique_ptr<xla::PjRtExecutable> pjrt_executable) {
-  return std::unique_ptr<Executable>(new PjRtExecutable(
-      std::shared_ptr<xla::PjRtExecutable>(pjrt_executable.release())));
-}
-
 absl::StatusOr<std::unique_ptr<Executable>> PjRtExecutable::Create(
     std::shared_ptr<xla::PjRtExecutable> pjrt_executable) {
   return std::unique_ptr<Executable>(
@@ -203,16 +197,6 @@ absl::StatusOr<std::string> PjRtExecutable::Serialize() const {
   return pjrt_executable_->SerializeExecutable();
 }
 
-absl::StatusOr<std::unique_ptr<LoadedExecutable>> PjRtLoadedExecutable::Create(
-    PjRtCompatibleClient* client,
-    std::unique_ptr<xla::PjRtLoadedExecutable> pjrt_loaded_executable,
-    std::vector<tsl::RCReference<LoadedHostCallback>> loaded_host_callbacks) {
-  return Create(client,
-                std::shared_ptr<xla::PjRtLoadedExecutable>(
-                    pjrt_loaded_executable.release()),
-                std::move(loaded_host_callbacks));
-}
-
 absl::StatusOr<std::unique_ptr<LoadedExecutable>> PjRtLoadedExecutable::Create(
     PjRtCompatibleClient* client,
     std::shared_ptr<xla::PjRtLoadedExecutable> pjrt_loaded_executable,
diff --git a/third_party/xla/xla/python/pjrt_ifrt/pjrt_executable.h b/third_party/xla/xla/python/pjrt_ifrt/pjrt_executable.h
index 1aa8e2479126e4..fd59a236b15e20 100644
--- a/third_party/xla/xla/python/pjrt_ifrt/pjrt_executable.h
+++ b/third_party/xla/xla/python/pjrt_ifrt/pjrt_executable.h
@@ -81,8 +81,6 @@ class PjRtExecutable final
     : public llvm::RTTIExtends<PjRtExecutable, PjRtCompatibleExecutable> {
  public:
   // Creates PjRtExecutable from xla::PjRtExecutable.
-  static absl::StatusOr<std::unique_ptr<Executable>> Create(
-      std::unique_ptr<xla::PjRtExecutable> pjrt_executable);
   static absl::StatusOr<std::unique_ptr<Executable>> Create(
       std::shared_ptr<xla::PjRtExecutable> pjrt_executable);
 
@@ -175,10 +173,6 @@ class PjRtLoadedExecutable final
   // Creates PjRtExecutable from xla::PjRtLoadedExecutable. We expect that
   // xla::PjRtLoadedExecutable has fixed output dtypes/shapes/shardings.
   // PjRtLoadedExecutable::GetHloModules() must be implemented.
-  static absl::StatusOr<std::unique_ptr<LoadedExecutable>> Create(
-      PjRtCompatibleClient* client,
-      std::unique_ptr<xla::PjRtLoadedExecutable> pjrt_loaded_executable,
-      std::vector<tsl::RCReference<LoadedHostCallback>> loaded_host_callbacks);
   static absl::StatusOr<std::unique_ptr<LoadedExecutable>> Create(
       PjRtCompatibleClient* client,
       std::shared_ptr<xla::PjRtLoadedExecutable> pjrt_loaded_executable,

From 7a49d6de1169a87b199b0face45b46af9fcbbe75 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Mon, 13 May 2024 19:17:41 -0700
Subject: [PATCH 083/478] Update ops-related pbtxt files.

PiperOrigin-RevId: 633404455
---
 tensorflow/core/ops/ops.pbtxt | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorflow/core/ops/ops.pbtxt b/tensorflow/core/ops/ops.pbtxt
index 7322cef85415f5..2c64cab5e211e2 100644
--- a/tensorflow/core/ops/ops.pbtxt
+++ b/tensorflow/core/ops/ops.pbtxt
@@ -1,4 +1,4 @@
-go/debugstr   
+go/nodeserialize    
 op {
   name: "Abort"
   attr {

From 3e4ed32a8ad25709dfd561d5c34cc13de537578d Mon Sep 17 00:00:00 2001
From: David Dunleavy <ddunleavy@google.com>
Date: Mon, 13 May 2024 19:39:55 -0700
Subject: [PATCH 084/478] Mark `gpu_sparse_dot_test` as `no_oss` due to missing
 dep

PiperOrigin-RevId: 633408111
---
 third_party/xla/xla/service/gpu/tests/BUILD | 1 +
 1 file changed, 1 insertion(+)

diff --git a/third_party/xla/xla/service/gpu/tests/BUILD b/third_party/xla/xla/service/gpu/tests/BUILD
index 8b11e1de46be2f..1c3dfbf7311549 100644
--- a/third_party/xla/xla/service/gpu/tests/BUILD
+++ b/third_party/xla/xla/service/gpu/tests/BUILD
@@ -926,6 +926,7 @@ xla_test(
         "gpu_a100",
         "gpu_h100",
     ],
+    tags = ["no_oss"],  # b/340304923
     deps = if_cuda_is_configured(
         [
             ":gpu_codegen_test",

From 820f077cc176bf9942d7306bf11ef54e585b3474 Mon Sep 17 00:00:00 2001
From: Siqiao Wu <siqiaowu@google.com>
Date: Mon, 13 May 2024 19:53:26 -0700
Subject: [PATCH 085/478] Add some internal change.

PiperOrigin-RevId: 633410091
---
 tensorflow/core/tfrt/ifrt/BUILD               |  1 +
 .../ifrt/ifrt_executable_registry_test.cc     | 60 +++++++++++++++++++
 2 files changed, 61 insertions(+)

diff --git a/tensorflow/core/tfrt/ifrt/BUILD b/tensorflow/core/tfrt/ifrt/BUILD
index 9db13a6387a286..0bfb4aa0cb5752 100644
--- a/tensorflow/core/tfrt/ifrt/BUILD
+++ b/tensorflow/core/tfrt/ifrt/BUILD
@@ -490,6 +490,7 @@ tf_cc_test(
         "//tensorflow/core/framework:tensor",
         "//tensorflow/core/framework:types_proto_cc",
         "//tensorflow/core/platform:resource_loader",
+        "//tensorflow/core/platform:status_matchers",
         "@com_google_absl//absl/status",
         "@com_google_absl//absl/status:statusor",
         "@com_google_absl//absl/strings",
diff --git a/tensorflow/core/tfrt/ifrt/ifrt_executable_registry_test.cc b/tensorflow/core/tfrt/ifrt/ifrt_executable_registry_test.cc
index 5dcdc6e80cca2c..e3ee0e43547771 100644
--- a/tensorflow/core/tfrt/ifrt/ifrt_executable_registry_test.cc
+++ b/tensorflow/core/tfrt/ifrt/ifrt_executable_registry_test.cc
@@ -40,6 +40,7 @@ limitations under the License.
 #include "tensorflow/core/common_runtime/device_mgr.h"
 #include "tensorflow/core/framework/types.pb.h"
 #include "tensorflow/core/platform/resource_loader.h"
+#include "tensorflow/core/platform/status_matchers.h"
 #include "tensorflow/core/platform/test.h"
 #include "tensorflow/core/tfrt/ifrt/ifrt_loaded_variable_registry.h"
 #include "tensorflow/core/tfrt/ifrt/ifrt_restore_tensor_registry.h"
@@ -118,6 +119,44 @@ TEST(IfrtExecutableRegistry, Basic) {
   ASSERT_EQ(executable_ptr, raw_ptr);
 }
 
+TEST(IfrtExecutableRegistry, DuplicateRegistrationFails) {
+  mlir::DialectRegistry registry;
+  mlir::registerAllDialects(registry);
+  mlir::RegisterAllTensorFlowDialects(registry);
+
+  mlir::MLIRContext context(registry);
+
+  int64_t program_id = 1234;
+
+  TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr<IfrtServingExecutable> executable,
+                          CreateIfrtServingExecutable(context, program_id));
+  TF_ASSERT_OK_AND_ASSIGN(auto handle, ServingExecutableRegistry::Register(
+                                           program_id, std::move(executable)));
+
+  EXPECT_THAT(
+      ServingExecutableRegistry::Register(program_id, std::move(executable)),
+      testing::StatusIs(absl::StatusCode::kAlreadyExists));
+}
+
+TEST(IfrtExecutableRegistry, ReleaseOk) {
+  mlir::DialectRegistry registry;
+  mlir::registerAllDialects(registry);
+  mlir::RegisterAllTensorFlowDialects(registry);
+
+  mlir::MLIRContext context(registry);
+
+  int64_t program_id = 1234;
+
+  TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr<IfrtServingExecutable> executable,
+                          CreateIfrtServingExecutable(context, program_id));
+  TF_ASSERT_OK_AND_ASSIGN(auto handle, ServingExecutableRegistry::Register(
+                                           program_id, std::move(executable)));
+
+  handle.Release();
+
+  EXPECT_EQ(ServingExecutableRegistry::Lookup(program_id), nullptr);
+}
+
 TEST(IfrtExecutableRegistry, FreezeOk) {
   mlir::DialectRegistry registry;
   mlir::registerAllDialects(registry);
@@ -142,6 +181,27 @@ TEST(IfrtExecutableRegistry, FreezeOk) {
   ASSERT_EQ(executable_ptr, raw_ptr);
 }
 
+TEST(IfrtExecutableRegistry, FreezeFailedProgramNotRegistered) {
+  mlir::DialectRegistry registry;
+  mlir::registerAllDialects(registry);
+  mlir::RegisterAllTensorFlowDialects(registry);
+
+  mlir::MLIRContext context(registry);
+
+  int64_t program_id = 1234;
+
+  TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr<IfrtServingExecutable> executable,
+                          CreateIfrtServingExecutable(context, program_id));
+
+  TF_ASSERT_OK_AND_ASSIGN(auto handle, ServingExecutableRegistry::Register(
+                                           program_id, std::move(executable)));
+
+  handle.Release();
+
+  EXPECT_THAT(handle.Freeze(),
+              testing::StatusIs(absl::StatusCode::kFailedPrecondition));
+}
+
 TEST(IfrtExecutableRegistry, InvalidProgramIdShallReturnNull) {
   int64_t program_id = 1234;
 

From c76d67465b432d1d27fbfb737c63027302f1f7a6 Mon Sep 17 00:00:00 2001
From: Kuangyuan Chen <chky@google.com>
Date: Mon, 13 May 2024 20:12:26 -0700
Subject: [PATCH 086/478] Improve test coverage of fallback_state.cc

PiperOrigin-RevId: 633415502
---
 tensorflow/core/tfrt/fallback/BUILD           |  7 +++-
 .../core/tfrt/fallback/fallback_state_test.cc | 40 ++++++++++++++++++-
 2 files changed, 43 insertions(+), 4 deletions(-)

diff --git a/tensorflow/core/tfrt/fallback/BUILD b/tensorflow/core/tfrt/fallback/BUILD
index 454c102deec081..c77b4133def642 100644
--- a/tensorflow/core/tfrt/fallback/BUILD
+++ b/tensorflow/core/tfrt/fallback/BUILD
@@ -57,10 +57,13 @@ tf_cc_test(
     srcs = ["fallback_state_test.cc"],
     deps = [
         ":fallback_state",
-        "//tensorflow/core:framework",
+        "//tensorflow/cc:cc_ops",
+        "//tensorflow/cc:const_op",
+        "//tensorflow/cc:ops",
+        "//tensorflow/cc:scope",
+        "//tensorflow/core:all_kernels",
         "//tensorflow/core:test",
         "//tensorflow/core:test_main",
-        "//tensorflow/core/framework:function_proto_cc",
         "//tensorflow/core/platform:status_matchers",
         "//tensorflow/core/protobuf:error_codes_proto_impl_cc",
     ],
diff --git a/tensorflow/core/tfrt/fallback/fallback_state_test.cc b/tensorflow/core/tfrt/fallback/fallback_state_test.cc
index 352a5119e69460..d7d55311e7ffd4 100644
--- a/tensorflow/core/tfrt/fallback/fallback_state_test.cc
+++ b/tensorflow/core/tfrt/fallback/fallback_state_test.cc
@@ -14,11 +14,15 @@ limitations under the License.
 ==============================================================================*/
 #include "tensorflow/core/tfrt/fallback/fallback_state.h"
 
-#include "tensorflow/core/framework/function.h"
-#include "tensorflow/core/framework/function.pb.h"
+#include <utility>
+
+#include "tensorflow/cc/framework/ops.h"
+#include "tensorflow/cc/framework/scope.h"
+#include "tensorflow/cc/ops/const_op.h"
 #include "tensorflow/core/platform/status_matchers.h"
 #include "tensorflow/core/platform/test.h"
 #include "tensorflow/core/protobuf/error_codes.pb.h"
+#include "tsl/lib/core/status_test_util.h"
 
 namespace tensorflow {
 namespace {
@@ -54,5 +58,37 @@ TEST(FallbackStateTest, CreateRendezvous) {
                                    HasSubstr("rendezvous"))));
 }
 
+TEST(FallbackStateTest, CreateGraphExecutionState) {
+  tensorflow::SessionOptions session_options;
+  tensorflow::FunctionDefLibrary fdef_lib;
+  TF_ASSERT_OK_AND_ASSIGN(
+      auto fallback_state,
+      tfrt_stub::FallbackState::CreateWithCpuDevice(session_options, fdef_lib));
+
+  GraphDef graphdef;
+  {
+    auto scope = tensorflow::Scope::NewRootScope().WithDevice(
+        "/job:localhost/replica:0/task:0/device:CPU:0");
+
+    Output a = ops::Const(scope.WithOpName("a"), 2.0, {1, 1});
+
+    TF_ASSERT_OK(scope.ToGraphDef(&graphdef));
+  }
+
+  TF_ASSERT_OK_AND_ASSIGN(
+      auto graph_execution_state,
+      fallback_state->CreateGraphExecutionState(std::move(graphdef)));
+}
+
+TEST(FallbackStateTest, CreateWithMockGpuDevice) {
+  tensorflow::SessionOptions session_options;
+  tensorflow::FunctionDefLibrary fdef_lib;
+  TF_ASSERT_OK_AND_ASSIGN(auto fallback_state,
+                          tfrt_stub::FallbackState::CreateWithMockGpuDevice(
+                              session_options, fdef_lib));
+  const auto& device_manager = fallback_state->device_manager();
+  EXPECT_GT(device_manager.NumDeviceType("GPU"), 0);
+}
+
 }  // namespace
 }  // namespace tensorflow

From ff8a035af1e1364c29601b2bf7d5046eeeaeea23 Mon Sep 17 00:00:00 2001
From: Fangrui Song <maskray@google.com>
Date: Mon, 13 May 2024 20:52:17 -0700
Subject: [PATCH 087/478] Integrate LLVM at llvm/llvm-project@e6b2197a89f5

Updates LLVM usage to match
[e6b2197a89f5](https://github.com/llvm/llvm-project/commit/e6b2197a89f5)

PiperOrigin-RevId: 633423574
---
 third_party/llvm/workspace.bzl | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/third_party/llvm/workspace.bzl b/third_party/llvm/workspace.bzl
index 243cc54703c4ec..669cff9644af1a 100644
--- a/third_party/llvm/workspace.bzl
+++ b/third_party/llvm/workspace.bzl
@@ -4,8 +4,8 @@ load("//third_party:repo.bzl", "tf_http_archive")
 
 def repo(name):
     """Imports LLVM."""
-    LLVM_COMMIT = "a6d7828f4c50c1ec7b0b5f61fe59d7a768175dcc"
-    LLVM_SHA256 = "c09ec3020fb6b136064ff32f53ac07067b0f12ccbf016ac69965e4e38d61a9c0"
+    LLVM_COMMIT = "e6b2197a89f5d6d0f56a03c03b8afda561eee899"
+    LLVM_SHA256 = "a2398fc87e7c4ef96f17999a83240d421686718a8f095fed718b7539576c1fb7"
 
     tf_http_archive(
         name = name,

From f29f62111ecb98a54951377bb0d7069a7d8c2314 Mon Sep 17 00:00:00 2001
From: Terry Heo <terryheo@google.com>
Date: Mon, 13 May 2024 21:09:52 -0700
Subject: [PATCH 088/478] gpu_delegate: Update ADD / MUL broadcast support

4D + 3D works with broadcast when 4D batch is 1.

PiperOrigin-RevId: 633426993
---
 .../tools/versioning/gpu_compatibility.cc     | 51 +++++++++++++++----
 .../versioning/gpu_compatibility_test.cc      | 42 +++++++++++++--
 2 files changed, 77 insertions(+), 16 deletions(-)

diff --git a/tensorflow/lite/tools/versioning/gpu_compatibility.cc b/tensorflow/lite/tools/versioning/gpu_compatibility.cc
index 29d5e27c4266c3..d8e1b12879a85b 100644
--- a/tensorflow/lite/tools/versioning/gpu_compatibility.cc
+++ b/tensorflow/lite/tools/versioning/gpu_compatibility.cc
@@ -435,6 +435,36 @@ absl::Status CheckCustomOpsGpuDelegateCompatibility(const OpSignature& op_sig) {
       absl::StrCat("Not supported custom op ", op_sig.custom_name));
 }
 
+absl::Status CheckAddMulBroadcastCompatibility(
+    const OpSignatureTensorSpec& input0, const OpSignatureTensorSpec& input1) {
+  if (input0.dims.size() > 1 && input1.dims.size() > 1 &&
+      input0.dims.size() != input1.dims.size()) {
+    const std::vector<int32_t>*longer_dims, *shorter_dims;
+    if (input0.dims.size() >= input1.dims.size()) {
+      longer_dims = &input0.dims;
+      shorter_dims = &input1.dims;
+    } else {
+      longer_dims = &input1.dims;
+      shorter_dims = &input0.dims;
+    }
+    bool is_broadcastable = false;
+
+    // Broadcasting 3D to 4D with batch 1 works.
+    if (longer_dims->size() == 4 && shorter_dims->size() == 3 &&
+        longer_dims->at(0) == 1) {
+      is_broadcastable = true;
+    }
+
+    if (!is_broadcastable) {
+      return absl::UnimplementedError(
+          absl::StrCat("Doesn't support broadcasting - input0: [",
+                       absl::StrJoin(input0.dims, ","), "], input1: [",
+                       absl::StrJoin(input1.dims, ","), "]"));
+    }
+  }
+  return absl::OkStatus();
+}
+
 }  // namespace
 
 // Logics here used to be in TFLiteOperationParser:IsSupported()
@@ -449,12 +479,9 @@ absl::Status CheckGpuDelegateCompatibility(const OpSignature& op_sig) {
       }
       const auto& input0 = op_sig.inputs.at(0);
       const auto& input1 = op_sig.inputs.at(1);
-      if (input0.dims.size() > 1 && input1.dims.size() > 1 &&
-          input0.dims.size() != input1.dims.size()) {
-        return absl::UnimplementedError(
-            absl::StrCat("ADD doesn't support broadcasting - input0: [",
-                         absl::StrJoin(input0.dims, ","), "], input1: [",
-                         absl::StrJoin(input1.dims, ","), "]"));
+      auto broadcastable = CheckAddMulBroadcastCompatibility(input0, input1);
+      if (!broadcastable.ok()) {
+        return broadcastable;
       }
       const TfLiteAddParams* tf_options;
       return RetrieveBuiltinData(op_sig, &tf_options);
@@ -700,11 +727,13 @@ absl::Status CheckGpuDelegateCompatibility(const OpSignature& op_sig) {
               "MUL requires one tensor that not less than second in all "
               "dimensions.");
         }
-      } else if (input0.dims.size() > 1 && input1.dims.size() > 1) {
-        return absl::UnimplementedError(
-            absl::StrCat("MUL doesn't support broadcasting - input0: [",
-                         absl::StrJoin(input0.dims, ","), "], input1: [",
-                         absl::StrJoin(input1.dims, ","), "]"));
+      } else {
+        const auto& input0 = op_sig.inputs.at(0);
+        const auto& input1 = op_sig.inputs.at(1);
+        auto broadcastable = CheckAddMulBroadcastCompatibility(input0, input1);
+        if (!broadcastable.ok()) {
+          return broadcastable;
+        }
       }
       const TfLiteMulParams* tf_options;
       RETURN_IF_ERROR(RetrieveBuiltinData(op_sig, &tf_options));
diff --git a/tensorflow/lite/tools/versioning/gpu_compatibility_test.cc b/tensorflow/lite/tools/versioning/gpu_compatibility_test.cc
index f3ca69dfc7e98d..8493137c1e0f90 100644
--- a/tensorflow/lite/tools/versioning/gpu_compatibility_test.cc
+++ b/tensorflow/lite/tools/versioning/gpu_compatibility_test.cc
@@ -94,7 +94,7 @@ TEST(CheckGpuDelegateCompatibility, FCConstInput) {
             "FullyConnected doesn't support constant input.");
 }
 
-TEST(CheckGpuDelegateCompatibility, Add1DBroadCastSuccess) {
+TEST(CheckGpuDelegateCompatibility, Add1Dto3DBroadcastSuccess) {
   OpSignature op_sig = OpSignature();
   op_sig.op = BuiltinOperator_ADD;
   auto params = std::make_unique<TfLiteAddParams>();
@@ -108,20 +108,52 @@ TEST(CheckGpuDelegateCompatibility, Add1DBroadCastSuccess) {
   EXPECT_TRUE(CheckGpuDelegateCompatibility(op_sig).message().empty());
 }
 
-TEST(CheckGpuDelegateCompatibility, Add2DBroadCastFail) {
+TEST(CheckGpuDelegateCompatibility, Add2Dto3DBroadcastFail) {
   OpSignature op_sig = OpSignature();
   op_sig.op = BuiltinOperator_ADD;
   auto params = std::make_unique<TfLiteAddParams>();
   op_sig.builtin_data = static_cast<void*>(params.get());
   op_sig.inputs = std::vector<OpSignatureTensorSpec>(2);
   op_sig.inputs[0] = OpSignatureTensorSpec();
-  op_sig.inputs[0].dims = {4, 1, 2};
+  op_sig.inputs[0].dims = {1, 100, 256};
+  op_sig.inputs[1] = OpSignatureTensorSpec();
+  op_sig.inputs[1].dims = {100, 256};
+
+  EXPECT_EQ(CheckGpuDelegateCompatibility(op_sig).message(),
+            "Doesn't support broadcasting - input0: [1,100,256], input1: "
+            "[100,256]");
+}
+
+TEST(CheckGpuDelegateCompatibility, Add3Dto4DBroadcastFail) {
+  OpSignature op_sig = OpSignature();
+  op_sig.op = BuiltinOperator_ADD;
+  auto params = std::make_unique<TfLiteAddParams>();
+  op_sig.builtin_data = static_cast<void*>(params.get());
+  op_sig.inputs = std::vector<OpSignatureTensorSpec>(2);
+  op_sig.inputs[0] = OpSignatureTensorSpec();
+  op_sig.inputs[0].dims = {4, 1, 1, 2};
   op_sig.inputs[1] = OpSignatureTensorSpec();
-  op_sig.inputs[1].dims = {2, 2};
+  // Can't broadcast using batch of 4
+  op_sig.inputs[1].dims = {1, 1, 2};
 
   EXPECT_EQ(
       CheckGpuDelegateCompatibility(op_sig).message(),
-      "ADD doesn't support broadcasting - input0: [4,1,2], input1: [2,2]");
+      "Doesn't support broadcasting - input0: [4,1,1,2], input1: [1,1,2]");
+}
+
+TEST(CheckGpuDelegateCompatibility, Add3Dto4DBroadcastSuccess) {
+  OpSignature op_sig = OpSignature();
+  op_sig.op = BuiltinOperator_ADD;
+  auto params = std::make_unique<TfLiteAddParams>();
+  op_sig.builtin_data = static_cast<void*>(params.get());
+  op_sig.inputs = std::vector<OpSignatureTensorSpec>(2);
+  op_sig.inputs[0] = OpSignatureTensorSpec();
+  op_sig.inputs[0].dims = {1, 128, 513, 3};
+  op_sig.inputs[1] = OpSignatureTensorSpec();
+  // Can be broadcasted to {1, 128, 513, 3}
+  op_sig.inputs[1].dims = {128, 513, 3};
+
+  EXPECT_TRUE(CheckGpuDelegateCompatibility(op_sig).message().empty());
 }
 
 }  // namespace tflite

From 53f8d3cd09d6716f51b8bef82b7b4d281c3b4850 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Mon, 13 May 2024 21:19:27 -0700
Subject: [PATCH 089/478] Update ops-related pbtxt files.

PiperOrigin-RevId: 633428831
---
 tensorflow/core/ops/ops.pbtxt | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorflow/core/ops/ops.pbtxt b/tensorflow/core/ops/ops.pbtxt
index 2c64cab5e211e2..53c362ed0ab819 100644
--- a/tensorflow/core/ops/ops.pbtxt
+++ b/tensorflow/core/ops/ops.pbtxt
@@ -1,4 +1,4 @@
-go/nodeserialize    
+go/debugstr    
 op {
   name: "Abort"
   attr {

From b52925b10aa2bcd7503e6ed3af7a40335a751091 Mon Sep 17 00:00:00 2001
From: Thai Nguyen <thaink@google.com>
Date: Mon, 13 May 2024 21:54:14 -0700
Subject: [PATCH 090/478] Support reusing calibration data if exists

Calibration is the most time-consuming step in quantization. This cl will help avoiding it in case of users already quantized it before and the changes in configuration if have doesn't affect the calibration results.

PiperOrigin-RevId: 633435225
---
 .../stablehlo/cc/calibration/BUILD            |   2 +
 .../stablehlo/cc/calibration/component.cc     |  99 ++++++++++------
 .../stablehlo/cc/calibration/component.h      |  12 +-
 .../stablehlo/cc/calibration/statistics.cc    |  14 ++-
 .../stablehlo/cc/calibration/statistics.h     |  13 +++
 .../insert_calibration_statistics_saver.cc    |  44 ++++---
 .../quantization/stablehlo/passes/passes.h    |   5 +-
 .../quantization/stablehlo/passes/passes.td   |   7 ++
 .../mlir/quantization/stablehlo/python/BUILD  |   1 +
 .../integration_test/quantize_model_test.py   | 108 +++++++++++++++++
 .../quantize_model_test_base.py               |  23 ++++
 .../stablehlo/quantization_config.proto       |   9 +-
 ...ration_statistics_saver_with_skipping.mlir |  47 ++++++++
 .../integration_test/quantize_model_test.py   | 110 ++++++++++++++++++
 .../tensorflow/python/quantize_model.cc       |  28 +++--
 .../tensorflow/python/quantize_model.py       |   2 +-
 16 files changed, 453 insertions(+), 71 deletions(-)
 create mode 100644 tensorflow/compiler/mlir/quantization/stablehlo/tests/passes/insert_calibration_statistics_saver_with_skipping.mlir

diff --git a/tensorflow/compiler/mlir/quantization/stablehlo/cc/calibration/BUILD b/tensorflow/compiler/mlir/quantization/stablehlo/cc/calibration/BUILD
index 7b7653c9db12c8..9926546f8c47a8 100644
--- a/tensorflow/compiler/mlir/quantization/stablehlo/cc/calibration/BUILD
+++ b/tensorflow/compiler/mlir/quantization/stablehlo/cc/calibration/BUILD
@@ -61,9 +61,11 @@ cc_library(
         "//tensorflow/compiler/mlir/quantization/stablehlo/cc:types",
         "//tensorflow/compiler/mlir/quantization/tensorflow:exported_model_proto_cc",
         "//tensorflow/compiler/mlir/quantization/tensorflow:quantization_options_proto_cc",
+        "//tensorflow/compiler/mlir/quantization/tensorflow/calibrator:calibration_statistics_proto_cc",
         "//tensorflow/compiler/mlir/quantization/tensorflow/cc:run_passes",
         "//tensorflow/compiler/mlir/quantization/tensorflow/python:py_function_lib",
         "//tensorflow/core/protobuf:for_core_protos_cc",
+        "@com_google_absl//absl/algorithm:container",
         "@com_google_absl//absl/base:nullability",
         "@com_google_absl//absl/container:flat_hash_map",
         "@com_google_absl//absl/log",
diff --git a/tensorflow/compiler/mlir/quantization/stablehlo/cc/calibration/component.cc b/tensorflow/compiler/mlir/quantization/stablehlo/cc/calibration/component.cc
index 2ec9b38865edfb..52db906e512391 100644
--- a/tensorflow/compiler/mlir/quantization/stablehlo/cc/calibration/component.cc
+++ b/tensorflow/compiler/mlir/quantization/stablehlo/cc/calibration/component.cc
@@ -20,6 +20,7 @@ limitations under the License.
 #include <utility>
 #include <vector>
 
+#include "absl/algorithm/container.h"
 #include "absl/base/nullability.h"
 #include "absl/container/flat_hash_map.h"
 #include "absl/log/die_if_null.h"
@@ -42,6 +43,7 @@ limitations under the License.
 #include "tensorflow/compiler/mlir/quantization/stablehlo/cc/types.h"
 #include "tensorflow/compiler/mlir/quantization/stablehlo/passes/passes.h"
 #include "tensorflow/compiler/mlir/quantization/stablehlo/quantization_config.pb.h"
+#include "tensorflow/compiler/mlir/quantization/tensorflow/calibrator/calibration_statistics.pb.h"
 #include "tensorflow/compiler/mlir/quantization/tensorflow/cc/run_passes.h"
 #include "tensorflow/compiler/mlir/quantization/tensorflow/exported_model.pb.h"
 #include "tensorflow/compiler/mlir/quantization/tensorflow/python/py_function_lib.h"
@@ -51,32 +53,52 @@ limitations under the License.
 #include "tsl/platform/statusor.h"
 
 namespace mlir::quant::stablehlo {
+namespace {
 
 using ::stablehlo::quantization::AddCalibrationStatistics;
 using ::stablehlo::quantization::CreateRepresentativeDatasetFileMap;
 using ::stablehlo::quantization::DisableDebugging;
+using ::stablehlo::quantization::IsCalibrationRequired;
 using ::stablehlo::quantization::QuantizationConfig;
+using ::stablehlo::quantization::ReadStatistics;
 using ::stablehlo::quantization::RepresentativeDatasetConfig;
 using ::stablehlo::quantization::io::CreateTmpDir;
 using ::stablehlo::quantization::io::GetLocalTmpFileName;
+using ::stablehlo::quantization::io::ListDirectory;
 using ::tensorflow::AssetFileDef;
 using ::tensorflow::SignatureDef;
+using ::tensorflow::calibrator::CalibrationStatistics;
 using ::tensorflow::quantization::ExportedModel;
 using ::tensorflow::quantization::PyFunctionLibrary;
 using ::tensorflow::quantization::RunPasses;
+using CalibrationStatisticsFlatMap =
+    absl::flat_hash_map<std::string, CalibrationStatistics>;
 
-absl::Status RunCalibrationPasses(mlir::ModuleOp module_op, MLIRContext& ctx,
-                                  absl::string_view calibration_data_dir) {
+}  // namespace
+
+absl::Status RunCalibrationPasses(
+    mlir::ModuleOp module_op, MLIRContext& ctx,
+    absl::string_view calibration_data_dir,
+    const bool force_regenerate_calibration_data) {
   // Disable DumpTensor ops when running calibration.
   DisableDebugging(module_op);
 
+  std::vector<std::string> skipping_aggregator_ops;
+  if (!force_regenerate_calibration_data) {
+    TF_ASSIGN_OR_RETURN(const CalibrationStatisticsFlatMap statistics_map,
+                        ReadStatistics(calibration_data_dir));
+    absl::c_for_each(statistics_map, [&](const auto& iter) {
+      return skipping_aggregator_ops.push_back(iter.first);
+    });
+  }
+
   return RunPasses(
       /*name=*/
       CalibrationComponent::kName,
       /*add_passes_func=*/
-      [calibration_data_dir](PassManager& pm) {
-        pm.addPass(
-            CreateInsertCalibrationStatisticsSaverPass(calibration_data_dir));
+      [calibration_data_dir, &skipping_aggregator_ops](PassManager& pm) {
+        pm.addPass(CreateInsertCalibrationStatisticsSaverPass(
+            calibration_data_dir, skipping_aggregator_ops));
       },
       ctx, module_op);
 }
@@ -97,8 +119,9 @@ CalibrationComponent::CalibrationComponent(
       signature_def_map_(std::move(signature_def_map)),
       signature_keys_(std::move(signature_keys)) {}
 
-absl::StatusOr<ExportedModel> CalibrationComponent::ExportToSavedModel(
+absl::Status CalibrationComponent::ExportToSavedModel(
     ModuleOp module_op, absl::string_view calibration_data_dir,
+    const bool force_regenerate_calibration_data,
     const absl::string_view dst_saved_model_path) {
   TF_ASSIGN_OR_RETURN(const std::string checkpoint_dir, GetLocalTmpFileName());
 
@@ -106,8 +129,13 @@ absl::StatusOr<ExportedModel> CalibrationComponent::ExportToSavedModel(
   // be reflected in the original values.
   mlir::OwningOpRef<mlir::ModuleOp> cloned_module_ref(module_op.clone());
 
-  TF_RETURN_IF_ERROR(
-      RunCalibrationPasses(*cloned_module_ref, *ctx_, calibration_data_dir));
+  TF_RETURN_IF_ERROR(RunCalibrationPasses(*cloned_module_ref, *ctx_,
+                                          calibration_data_dir,
+                                          force_regenerate_calibration_data));
+
+  const bool is_calibration_required =
+      IsCalibrationRequired(*cloned_module_ref);
+  if (!is_calibration_required) return absl::OkStatus();
 
   // `duplicate_shape_determining_constants = false` because the
   // resulting graph of this step is not expected to be loaded on TPU.
@@ -128,13 +156,13 @@ absl::StatusOr<ExportedModel> CalibrationComponent::ExportToSavedModel(
                                       src_saved_model_path_, tags_,
                                       signature_def_map_);
 
-  return exported_model;
+  return absl::OkStatus();
 }
 
 absl::StatusOr<ModuleOp> CalibrationComponent::Run(
     ModuleOp module_op, const QuantizationConfig& config) {
-  // Exports the pre-calibrated model to SavedModel.
-  TF_ASSIGN_OR_RETURN(const std::string precalibrated_saved_model_dir,
+  // Export the calibration model to SavedModel.
+  TF_ASSIGN_OR_RETURN(const std::string calibration_saved_model_dir,
                       CreateTmpDir());
 
   std::string calibration_data_dir =
@@ -143,29 +171,32 @@ absl::StatusOr<ModuleOp> CalibrationComponent::Run(
     TF_ASSIGN_OR_RETURN(calibration_data_dir, CreateTmpDir());
   }
 
-  TF_ASSIGN_OR_RETURN(ExportedModel exported_model,
-                      ExportToSavedModel(module_op, calibration_data_dir,
-                                         precalibrated_saved_model_dir));
-
-  // Translates `RepresentativeDatasetConfig`s to signature key ->
-  // `RepresentativeDatasetFile` mapping.
-  const auto dataset_configs =
-      config.calibration_options().representative_datasets();
-  const std::vector<RepresentativeDatasetConfig> dataset_config_vector(
-      dataset_configs.begin(), dataset_configs.end());
-  TF_ASSIGN_OR_RETURN(
-      const auto representative_dataset_file_map,
-      CreateRepresentativeDatasetFileMap(dataset_config_vector));
-
-  // Runs calibration on the exported model. The statistics will be stored in a
-  // separate singleton object `CalibratorSingleton` and are directly added to
-  // `exported_model` without re-importing it.
-  if (py_function_lib_->RunCalibration(
-          precalibrated_saved_model_dir, signature_keys_, tags_,
-          /*force_graph_mode_calibration=*/true,
-          representative_dataset_file_map) == std::nullopt) {
-    return absl::InternalError(
-        "CalibrationComponent error: Failed to run calibration.");
+  TF_RETURN_IF_ERROR(ExportToSavedModel(
+      module_op, calibration_data_dir,
+      config.calibration_options().force_regenerate_calibration_data(),
+      calibration_saved_model_dir));
+
+  TF_ASSIGN_OR_RETURN(std::vector<std::string> calibration_saved_model_files,
+                      ListDirectory(calibration_saved_model_dir));
+  if (!calibration_saved_model_files.empty()) {
+    // Translate `RepresentativeDatasetConfig`s to signature key ->
+    // `RepresentativeDatasetFile` mapping.
+    const auto dataset_configs =
+        config.calibration_options().representative_datasets();
+    const std::vector<RepresentativeDatasetConfig> dataset_config_vector(
+        dataset_configs.begin(), dataset_configs.end());
+    TF_ASSIGN_OR_RETURN(
+        const auto representative_dataset_file_map,
+        CreateRepresentativeDatasetFileMap(dataset_config_vector));
+
+    // Run calibration on the exported model.
+    if (py_function_lib_->RunCalibration(
+            calibration_saved_model_dir, signature_keys_, tags_,
+            /*force_graph_mode_calibration=*/true,
+            representative_dataset_file_map) == std::nullopt) {
+      return absl::InternalError(
+          "CalibrationComponent error: Failed to run calibration.");
+    }
   }
 
   if (absl::Status status = AddCalibrationStatistics(
diff --git a/tensorflow/compiler/mlir/quantization/stablehlo/cc/calibration/component.h b/tensorflow/compiler/mlir/quantization/stablehlo/cc/calibration/component.h
index a7f94e9f0a37bf..03d2dd933732d4 100644
--- a/tensorflow/compiler/mlir/quantization/stablehlo/cc/calibration/component.h
+++ b/tensorflow/compiler/mlir/quantization/stablehlo/cc/calibration/component.h
@@ -77,10 +77,11 @@ class CalibrationComponent : public Component {
   // Exports `module_op` to SavedModel at `dst_saved_model_path`. This is used
   // to export the pre-calibrated `module_op` to SavedModel so that the
   // calibration process can use it to load and run the graph with the
-  // representative dataset.
-  absl::StatusOr<tensorflow::quantization::ExportedModel> ExportToSavedModel(
-      ModuleOp module_op, absl::string_view calibration_data_dir,
-      absl::string_view dst_saved_model_path);
+  // representative dataset. Returns a failure status if the export fails.
+  absl::Status ExportToSavedModel(ModuleOp module_op,
+                                  absl::string_view calibration_data_dir,
+                                  bool force_regenerate_calibration_data,
+                                  absl::string_view dst_saved_model_path);
 
   // Imports the SavedModel at `calibrated_saved_model_path` to `ModuleOp` after
   // running calibration.
@@ -113,7 +114,8 @@ class CalibrationComponent : public Component {
 
 // Runs passes to prepare the calibration model.
 absl::Status RunCalibrationPasses(mlir::ModuleOp module_op, MLIRContext& ctx,
-                                  absl::string_view calibration_data_dir);
+                                  absl::string_view calibration_data_dir,
+                                  bool force_regenerate_calibration_data);
 
 }  // namespace mlir::quant::stablehlo
 
diff --git a/tensorflow/compiler/mlir/quantization/stablehlo/cc/calibration/statistics.cc b/tensorflow/compiler/mlir/quantization/stablehlo/cc/calibration/statistics.cc
index 141af8e06fc6db..ea96bd029b079e 100644
--- a/tensorflow/compiler/mlir/quantization/stablehlo/cc/calibration/statistics.cc
+++ b/tensorflow/compiler/mlir/quantization/stablehlo/cc/calibration/statistics.cc
@@ -45,6 +45,8 @@ using ::tensorflow::quantization::PyFunctionLibrary;
 using CalibrationStatisticsFlatMap =
     absl::flat_hash_map<std::string, CalibrationStatistics>;
 
+}  // namespace
+
 // Reads the calibration statistics from the given directory.
 absl::StatusOr<CalibrationStatisticsFlatMap> ReadStatistics(
     absl::string_view calibration_data_dir) {
@@ -63,8 +65,6 @@ absl::StatusOr<CalibrationStatisticsFlatMap> ReadStatistics(
   return statistics_map;
 }
 
-}  // namespace
-
 absl::Status AddCalibrationStatistics(
     mlir::ModuleOp module_op, absl::string_view calibration_data_dir,
     const CalibrationOptions& calibration_options,
@@ -102,4 +102,14 @@ absl::Status AddCalibrationStatistics(
   return status;
 }
 
+bool IsCalibrationRequired(mlir::ModuleOp module_op) {
+  bool calibration_required = false;
+  module_op.walk(
+      [&calibration_required](
+          mlir::TF::CalibrationStatisticsSaverOp statistics_saver_op) {
+        calibration_required = true;
+      });
+  return calibration_required;
+}
+
 }  // namespace stablehlo::quantization
diff --git a/tensorflow/compiler/mlir/quantization/stablehlo/cc/calibration/statistics.h b/tensorflow/compiler/mlir/quantization/stablehlo/cc/calibration/statistics.h
index 48112b1eafc816..41f78be3578bca 100644
--- a/tensorflow/compiler/mlir/quantization/stablehlo/cc/calibration/statistics.h
+++ b/tensorflow/compiler/mlir/quantization/stablehlo/cc/calibration/statistics.h
@@ -15,14 +15,24 @@ limitations under the License.
 #ifndef TENSORFLOW_COMPILER_MLIR_QUANTIZATION_STABLEHLO_CC_CALIBRATION_STATISTICS_H_
 #define TENSORFLOW_COMPILER_MLIR_QUANTIZATION_STABLEHLO_CC_CALIBRATION_STATISTICS_H_
 
+#include <string>
+
+#include "absl/container/flat_hash_map.h"
 #include "absl/status/status.h"
+#include "absl/status/statusor.h"
 #include "absl/strings/string_view.h"
 #include "mlir/IR/BuiltinOps.h"  // from @llvm-project
 #include "tensorflow/compiler/mlir/quantization/stablehlo/quantization_config.pb.h"
+#include "tensorflow/compiler/mlir/quantization/tensorflow/calibrator/calibration_statistics.pb.h"
 #include "tensorflow/compiler/mlir/quantization/tensorflow/python/py_function_lib.h"
 
 namespace stablehlo::quantization {
 
+// Reads the calibration statistics from the given directory.
+absl::StatusOr<absl::flat_hash_map<
+    std::string, tensorflow::calibrator::CalibrationStatistics>>
+ReadStatistics(absl::string_view calibration_data_dir);
+
 // Adds calibrated min / max values to CustomAggregator nodes in `graph_def`.
 // The min and max values will be added to the "min" and "max" attributes,
 // respectively. `calibration_options` provides the strategy to retrieve min and
@@ -32,6 +42,9 @@ absl::Status AddCalibrationStatistics(
     const stablehlo::quantization::CalibrationOptions& calibration_options,
     const tensorflow::quantization::PyFunctionLibrary& py_function_library);
 
+// Checks if the model required calibration.
+bool IsCalibrationRequired(mlir::ModuleOp module_op);
+
 }  // namespace stablehlo::quantization
 
 #endif  // TENSORFLOW_COMPILER_MLIR_QUANTIZATION_STABLEHLO_CC_CALIBRATION_STATISTICS_H_
diff --git a/tensorflow/compiler/mlir/quantization/stablehlo/passes/insert_calibration_statistics_saver.cc b/tensorflow/compiler/mlir/quantization/stablehlo/passes/insert_calibration_statistics_saver.cc
index 9e6bf25fb44811..8cb0b645c312cf 100644
--- a/tensorflow/compiler/mlir/quantization/stablehlo/passes/insert_calibration_statistics_saver.cc
+++ b/tensorflow/compiler/mlir/quantization/stablehlo/passes/insert_calibration_statistics_saver.cc
@@ -15,6 +15,8 @@ limitations under the License.
 #include <cstdint>
 #include <memory>
 #include <string>
+#include <unordered_set>
+#include <vector>
 
 #include "absl/strings/string_view.h"
 #include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
@@ -48,11 +50,14 @@ std::string GetOutputFilePath(absl::string_view calibration_data_dir,
 }
 
 // Finds `CustomAggregator` ops and collects their outputs and attributes.
-void FindCustomAggregatorOps(Region& region,
-                             SmallVector<Value>& statistics_outputs,
-                             SmallVector<StringRef>& ids,
-                             SmallVector<int32_t>& calibration_methods) {
+void FindCustomAggregatorOps(
+    Region& region,
+    const std::unordered_set<std::string>& aggregator_ops_to_ignore,
+    SmallVector<Value>& statistics_outputs, SmallVector<StringRef>& ids,
+    SmallVector<int32_t>& calibration_methods) {
   for (auto op : region.getOps<TF::CustomAggregatorOp>()) {
+    if (aggregator_ops_to_ignore.count(op.getId().str())) continue;
+
     ids.push_back(op.getId());
     calibration_methods.push_back(op.getCalibrationMethod());
     statistics_outputs.push_back(op.getMin());
@@ -63,11 +68,13 @@ void FindCustomAggregatorOps(Region& region,
 
 // Inserts a `CalibrationStatisticsSaverOp` to the end of the region.
 LogicalResult InsertCalibrationStatisticsSaverOp(
-    Region& region, MLIRContext& ctx, absl::string_view output_file_path) {
+    Region& region, MLIRContext& ctx, absl::string_view output_file_path,
+    const std::unordered_set<std::string>& aggregator_ops_to_ignore) {
   SmallVector<Value> statistics_outputs;
   SmallVector<StringRef> ids;
   SmallVector<int32_t> calibration_methods;
-  FindCustomAggregatorOps(region, statistics_outputs, ids, calibration_methods);
+  FindCustomAggregatorOps(region, aggregator_ops_to_ignore, statistics_outputs,
+                          ids, calibration_methods);
   if (statistics_outputs.empty()) return failure();
 
   OpBuilder builder(&ctx);
@@ -115,6 +122,7 @@ bool ContainCalibrationStatisticsSaverOp(Operation* op) {
 
 }  // namespace
 
+#define GEN_PASS_DECL_INSERTCALIBRATIONSTATISTICSSAVERPASS
 #define GEN_PASS_DEF_INSERTCALIBRATIONSTATISTICSSAVERPASS
 #include "tensorflow/compiler/mlir/quantization/stablehlo/passes/passes.h.inc"
 
@@ -126,11 +134,7 @@ class InsertCalibrationStatisticsSaverPass
       InsertCalibrationStatisticsSaverPass>::
       InsertCalibrationStatisticsSaverPassBase;
 
-  explicit InsertCalibrationStatisticsSaverPass(StringRef calibration_data_dir)
-      : calibration_data_dir_(calibration_data_dir) {}
-
  private:
-  std::string calibration_data_dir_;
   void runOnOperation() override;
 };
 
@@ -138,17 +142,22 @@ void InsertCalibrationStatisticsSaverPass::runOnOperation() {
   ModuleOp module_op = getOperation();
   MLIRContext& ctx = getContext();
 
+  std::unordered_set<std::string> aggregator_ops_to_ignore(
+      aggregator_ops_to_ignore_.begin(), aggregator_ops_to_ignore_.end());
+
   // Insert CalibrationStatisticsSaverOp to the end of each region.
   for (auto func_op : module_op.getOps<func::FuncOp>()) {
     int32_t output_file_idx = 0;
     StringRef func_name = func_op.getSymName();
 
-    func_op.walk([&output_file_idx, &ctx, &func_name, this](Operation* op) {
+    func_op.walk([&output_file_idx, &ctx, &func_name, &aggregator_ops_to_ignore,
+                  this](Operation* op) {
       for (Region& region : op->getRegions()) {
         if (succeeded(InsertCalibrationStatisticsSaverOp(
                 region, ctx,
                 GetOutputFilePath(calibration_data_dir_, func_name,
-                                  output_file_idx)))) {
+                                  output_file_idx),
+                aggregator_ops_to_ignore))) {
           ++output_file_idx;
         };
       }
@@ -167,9 +176,14 @@ void InsertCalibrationStatisticsSaverPass::runOnOperation() {
 }
 
 std::unique_ptr<OperationPass<ModuleOp>>
-CreateInsertCalibrationStatisticsSaverPass(StringRef calibration_data_dir) {
-  return std::make_unique<InsertCalibrationStatisticsSaverPass>(
-      calibration_data_dir);
+CreateInsertCalibrationStatisticsSaverPass(
+    StringRef calibration_data_dir,
+    const std::vector<std::string>& aggregator_ops_to_ignore) {
+  InsertCalibrationStatisticsSaverPassOptions options = {
+      .aggregator_ops_to_ignore_ = aggregator_ops_to_ignore,
+      .calibration_data_dir_ = calibration_data_dir.str(),
+  };
+  return std::make_unique<InsertCalibrationStatisticsSaverPass>(options);
 }
 
 }  // namespace mlir::quant::stablehlo
diff --git a/tensorflow/compiler/mlir/quantization/stablehlo/passes/passes.h b/tensorflow/compiler/mlir/quantization/stablehlo/passes/passes.h
index a6836eaf1009cd..d13c589c2ba890 100644
--- a/tensorflow/compiler/mlir/quantization/stablehlo/passes/passes.h
+++ b/tensorflow/compiler/mlir/quantization/stablehlo/passes/passes.h
@@ -18,6 +18,7 @@ limitations under the License.
 
 #include <memory>
 #include <string>
+#include <vector>
 
 #include "absl/status/statusor.h"
 #include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
@@ -46,7 +47,9 @@ CreateLiftQuantizableSpotsAsFunctionsPass(
 
 // Creates a pass that inserts CalibrationStatisticsSaverOp.
 std::unique_ptr<OperationPass<ModuleOp>>
-CreateInsertCalibrationStatisticsSaverPass(StringRef calibration_data_dir);
+CreateInsertCalibrationStatisticsSaverPass(
+    StringRef calibration_data_dir,
+    const std::vector<std::string>& aggregator_ops_to_ignore);
 
 // Adds generated pass default constructors or options definitions.
 #define GEN_PASS_DECL
diff --git a/tensorflow/compiler/mlir/quantization/stablehlo/passes/passes.td b/tensorflow/compiler/mlir/quantization/stablehlo/passes/passes.td
index b61550e5347e36..7661e8d562fbe9 100644
--- a/tensorflow/compiler/mlir/quantization/stablehlo/passes/passes.td
+++ b/tensorflow/compiler/mlir/quantization/stablehlo/passes/passes.td
@@ -228,5 +228,12 @@ def InsertCalibrationStatisticsSaverPass : Pass<"stablehlo-insert-calibration-st
     `CalibrationStatisticsSaver` op at the end of the function to collect their
     statistics.
   }];
+  let options = [
+    ListOption<"aggregator_ops_to_ignore_", "aggregator-ops-to-ignore", "std::string",
+               "Ops to ignore when inserting CalibrationStatisticsSaver.">,
+    Option<"calibration_data_dir_", "calibration-data-dir",
+        "std::string", /*default=*/"",
+        "The directory to save calibration data.">,
+  ];
   let dependentDialects = ["TF::TensorFlowDialect"];
 }
diff --git a/tensorflow/compiler/mlir/quantization/stablehlo/python/BUILD b/tensorflow/compiler/mlir/quantization/stablehlo/python/BUILD
index df5252b986adf5..a4bdf7da09f8ee 100644
--- a/tensorflow/compiler/mlir/quantization/stablehlo/python/BUILD
+++ b/tensorflow/compiler/mlir/quantization/stablehlo/python/BUILD
@@ -62,6 +62,7 @@ pytype_strict_library(
 #         "//tensorflow/python/ops:nn_ops",
 #         "//tensorflow/python/ops:variables",
 #         "//tensorflow/python/platform:client_testlib",
+#         "//tensorflow/python/platform:tf_logging",
 #         "//tensorflow/python/saved_model:load",
 #         "//tensorflow/python/saved_model:loader",
 #         "//tensorflow/python/saved_model:save",
diff --git a/tensorflow/compiler/mlir/quantization/stablehlo/python/integration_test/quantize_model_test.py b/tensorflow/compiler/mlir/quantization/stablehlo/python/integration_test/quantize_model_test.py
index 5e9f9955f329f3..ab0fb1d5662bba 100644
--- a/tensorflow/compiler/mlir/quantization/stablehlo/python/integration_test/quantize_model_test.py
+++ b/tensorflow/compiler/mlir/quantization/stablehlo/python/integration_test/quantize_model_test.py
@@ -577,6 +577,114 @@ def data_gen() -> repr_dataset.RepresentativeDataset:
         0.65,
     )
 
+  def test_reuse_calibration_data(self):
+    _, y_shape, bias_shape, x_signature, y_signature = (
+        self._prepare_sample_einsum_datashapes('abc,cde->abde', use_bias=True)
+    )
+
+    self._create_einsum_model(
+        self._input_saved_model_path,
+        'abc,cde->abde',
+        y_shape,
+        x_signature,
+        y_signature,
+        bias_shape,
+    )
+
+    # Generate model input data.
+    rng = np.random.default_rng(seed=42)
+    input_data = ops.convert_to_tensor(
+        rng.uniform(low=0.0, high=1.0, size=x_signature).astype('f4')
+    )
+
+    def data_gen() -> repr_dataset.RepresentativeDataset:
+      for _ in range(100):
+        yield {
+            'x': ops.convert_to_tensor(
+                np.random.uniform(low=0.0, high=1.0, size=x_signature).astype(
+                    'f4'
+                )
+            ),
+        }
+
+    dataset_path = self.create_tempfile('tfrecord').full_path
+    path_map = {'serving_default': dataset_path}
+    repr_dataset.TfRecordRepresentativeDatasetSaver(path_map).save(
+        {'serving_default': data_gen()}
+    )
+
+    calibration_data_dir = self.create_tempdir('calibration_data').full_path
+    config = qc.QuantizationConfig(
+        static_range_ptq_preset=qc.StaticRangePtqPreset(
+            representative_datasets=[
+                qc.RepresentativeDatasetConfig(
+                    tf_record=qc.TfRecordFile(path=dataset_path)
+                )
+            ]
+        ),
+        tf_saved_model=qc.TfSavedModelConfig(tags=[tag_constants.SERVING]),
+        calibration_options=qc.CalibrationOptions(
+            calibration_method=_CalibrationMethod.CALIBRATION_METHOD_MIN_MAX,
+            calibration_data_dir=calibration_data_dir,
+        ),
+    )
+
+    # Run quantization the first time, calibration is expected to be run.
+    with self.assertLogs(level='INFO') as info_logs:
+      quantization.quantize_saved_model(
+          self._input_saved_model_path,
+          self._output_saved_model_path,
+          config,
+      )
+      self.assertTrue(
+          self._any_log_contains(
+              'Calibration step is executed in graph mode.',
+              info_logs.records,
+          )
+      )
+      module_str = self._extract_first_xla_call_module_op(
+          self._output_saved_model_path
+      )
+      self.assertTrue(
+          re.search('stablehlo.dot_general.*xi8>.*xi8>.*xi32>', module_str)
+      )
+
+    # Run quantization the first time, calibration is expected to be skipped.
+    output_saved_model_path_2 = self.create_tempdir('output2').full_path
+    with self.assertLogs(level='INFO') as info_logs:
+      quantization.quantize_saved_model(
+          self._input_saved_model_path,
+          output_saved_model_path_2,
+          config,
+      )
+      self.assertFalse(
+          self._any_log_contains(
+              'Calibration step is executed in graph mode.',
+              info_logs.records,
+          )
+      )
+      module_str = self._extract_first_xla_call_module_op(
+          output_saved_model_path_2
+      )
+      self.assertTrue(
+          re.search('stablehlo.dot_general.*xi8>.*xi8>.*xi32>', module_str)
+      )
+
+    # Expect both quantized model to produce the same results.
+    root = load.load(self._output_saved_model_path)
+    self.assertCountEqual(root.signatures.keys(), {'serving_default'})
+    new_outputs_1 = root.signatures['serving_default'](
+        x=ops.convert_to_tensor(input_data)
+    )
+
+    root = load.load(output_saved_model_path_2)
+    self.assertCountEqual(root.signatures.keys(), {'serving_default'})
+    new_outputs_2 = root.signatures['serving_default'](
+        x=ops.convert_to_tensor(input_data)
+    )
+
+    self.assertAllClose(new_outputs_1, new_outputs_2)
+
   @parameterized.named_parameters(
       ('use_constant_with_int32_input', np.int32, False),
       ('use_variable_with_int32_input', np.int32, True),
diff --git a/tensorflow/compiler/mlir/quantization/stablehlo/python/integration_test/quantize_model_test_base.py b/tensorflow/compiler/mlir/quantization/stablehlo/python/integration_test/quantize_model_test_base.py
index 31c53a4cf20fe9..fef1784fec9370 100644
--- a/tensorflow/compiler/mlir/quantization/stablehlo/python/integration_test/quantize_model_test_base.py
+++ b/tensorflow/compiler/mlir/quantization/stablehlo/python/integration_test/quantize_model_test_base.py
@@ -33,11 +33,13 @@
 from tensorflow.python.ops import nn_ops
 from tensorflow.python.ops import variables
 from tensorflow.python.platform import test
+from tensorflow.python.platform import tf_logging as logging
 from tensorflow.python.saved_model import load
 from tensorflow.python.saved_model import loader_impl
 from tensorflow.python.saved_model import save as saved_model_save
 from tensorflow.python.types import core
 
+
 FUNC_ALIAS = 'some_alias'
 
 
@@ -164,6 +166,27 @@ def matmul(self, input_tensor: core.Tensor) -> Mapping[str, core.Tensor]:
     )
     return model
 
+  def _any_log_contains(
+      self, substring: str, log_record_list: List['logging.LogRecord']
+  ) -> bool:
+    """Returns True if any of the log contains a given substring.
+
+    Args:
+      substring: A piece of string to check whether it exists in the log
+        message.
+      log_record_list: A list of `absl.logging.LogRecord`s.
+
+    Returns:
+      True if and only if the substring exists in any of the log in
+      `log_record_list`.
+    """
+    return any(
+        map(
+            lambda log_record: substring in str(log_record.message),
+            log_record_list,
+        )
+    )
+
   def _create_matmul_and_same_scale_model(
       self,
       input_shape: Sequence[int],
diff --git a/tensorflow/compiler/mlir/quantization/stablehlo/quantization_config.proto b/tensorflow/compiler/mlir/quantization/stablehlo/quantization_config.proto
index f156b66997211d..a307e1d927eea3 100644
--- a/tensorflow/compiler/mlir/quantization/stablehlo/quantization_config.proto
+++ b/tensorflow/compiler/mlir/quantization/stablehlo/quantization_config.proto
@@ -278,7 +278,7 @@ message DebuggerConfig {
 }
 
 // Defines various calibration options.
-// Next ID: 5
+// Next ID: 6
 message CalibrationOptions {
   // Configurations for calibration methods.
   // Next ID: 7
@@ -332,8 +332,13 @@ message CalibrationOptions {
   // representative dataset used to calibrate a function.
   repeated RepresentativeDatasetConfig representative_datasets = 3;
 
-  // The path to save calibration statistics data.
+  // The path to save calibration statistics data. If not set, use a temporary
+  // directory.
   string calibration_data_dir = 4;
+
+  // Whether to reuse the existing calibration data in `calibration_data_dir`.
+  // Default to False.
+  bool force_regenerate_calibration_data = 5;
 }
 
 // Quantization configuration for StableHLO Quantizer. This is the primary
diff --git a/tensorflow/compiler/mlir/quantization/stablehlo/tests/passes/insert_calibration_statistics_saver_with_skipping.mlir b/tensorflow/compiler/mlir/quantization/stablehlo/tests/passes/insert_calibration_statistics_saver_with_skipping.mlir
new file mode 100644
index 00000000000000..97d546afe2b723
--- /dev/null
+++ b/tensorflow/compiler/mlir/quantization/stablehlo/tests/passes/insert_calibration_statistics_saver_with_skipping.mlir
@@ -0,0 +1,47 @@
+// RUN: stablehlo-quant-opt %s -split-input-file -stablehlo-insert-calibration-statistics-saver='aggregator-ops-to-ignore=skipping_id' | FileCheck %s
+
+func.func @serving_default(%arg0: tensor<1x3x4x3xf32>) -> (tensor<1x2x2x2xf32>) attributes {tf.entry_function = {control_outputs = "", inputs = "serving_default_input_tensor:0", outputs = "PartitionedCall:0"}} {
+  %cst = "tf.Const"() <{value = dense<[[[[-0.891899645, 0.392044574], [0.77720493, 1.31188095], [0.255048186, 2.700150e+00]], [[-1.08111858, -0.406604826], [-0.298575521, -2.25356531], [-1.00201964, 2.54532099]], [[-1.34911358, 0.279911458], [-0.868258893, -1.36708188], [0.866317451, -2.05804896]]], [[[-0.591397941, 0.331505477], [0.715151429, 2.64073896], [1.27163255, 0.206143498]], [[0.474211812, 1.45044816], [0.119936548, 2.54149938], [-0.939900994, 0.438387245]], [[-1.12486279, -1.09022558], [0.82202208, 1.04652023], [1.30316162, 2.62054276]]]]> : tensor<2x3x3x2xf32>}> : () -> tensor<2x3x3x2xf32>
+  %output, %min, %max, %histogram = "tf.CustomAggregator"(%arg0) <{calibration_method = 5 : i32, id = "skipping_id", num_bins = 32 : i32, max_percentile = 0.000000e+00 : f32, min_percentile = 0.000000e+00 : f32}> : (tensor<1x3x4x3xf32>) -> (tensor<1x3x4x3xf32>, tensor<f32>, tensor<f32>, tensor<512xi64>)
+  %0 = "tf.Conv2D"(%output, %cst) <{data_format = "NHWC", dilations = [1, 1, 1, 1], explicit_paddings = [], padding = "SAME", strides = [1, 2, 2, 1], use_cudnn_on_gpu = true}> {attr_map = "0:strides,1:use_cudnn_on_gpu,2:padding,3:explicit_paddings,4:dilations", device = ""} : (tensor<1x3x4x3xf32>, tensor<2x3x3x2xf32>) -> tensor<1x2x2x2xf32>
+  %output_1, %min_2, %max_3, %histogram_4 = "tf.CustomAggregator"(%0) <{calibration_method = 5 : i32, id = "keeping_id", num_bins = 32 : i32, max_percentile = 0.000000e+00 : f32, min_percentile = 0.000000e+00 : f32}> : (tensor<1x2x2x2xf32>) -> (tensor<1x2x2x2xf32>, tensor<f32>, tensor<f32>, tensor<512xi64>)
+  %1 = "tf.Identity"(%output_1) {device = ""} : (tensor<1x2x2x2xf32>) -> tensor<1x2x2x2xf32>
+  return %1 : tensor<1x2x2x2xf32>
+}
+// CHECK-LABEL: @serving_default
+// CHECK: %[[CUSTOM_AGGREGATOR_0:.*]], %[[MIN_O:.*]], %[[MAX_O:.*]], %[[HISTOGRAM_0:.*]] = "tf.CustomAggregator"
+// CKECK-SAME: <{calibration_method = 5 : i32, id = "skipping_id", num_bins = 32 : i32, max_percentile = 0.000000e+00 : f32, min_percentile = 0.000000e+00 : f32}> : (tensor<1x3x4x3xf32>) -> (tensor<1x3x4x3xf32>, tensor<f32>, tensor<f32>, tensor<512xi64>)
+// CHECK: %[[CUSTOM_AGGREGATOR_1:.*]], %[[MIN_1:.*]], %[[MAX_1:.*]], %[[HISTOGRAM_1:.*]] = "tf.CustomAggregator"
+// CKECK-SAME: <{calibration_method = 5 : i32, id = "keeping_id", num_bins = 32 : i32, max_percentile = 0.000000e+00 : f32, min_percentile = 0.000000e+00 : f32}> : (tensor<1x3x4x3xf32>) -> (tensor<1x3x4x3xf32>, tensor<f32>, tensor<f32>, tensor<512xi64>)
+// CHECK: "tf.CalibrationStatisticsSaver"(%[[MIN_1]], %[[MAX_1]], %[[HISTOGRAM_1]])
+// CHECK-SAME: <{calibration_methods = [5 : i32], ids = ["keeping_id"], output_file_path = "serving_default_0.pb"}>  : (tensor<f32>, tensor<f32>, tensor<512xi64>) -> ()
+// CHECK: return
+
+// -----
+
+module attributes {tf.versions = {bad_consumers = [], min_consumer = 12 : i32, producer = 1836 : i32}, tf_saved_model.semantics} {
+  func.func @main(%arg0: tensor<10x1x1024xf32> {tf_saved_model.index_path = ["input_tensor"]}) -> (tensor<10x1x3xf32> {tf_saved_model.index_path = ["output"]}) attributes {tf.entry_function = {control_outputs = "", inputs = "serving_default_input_tensor:0", outputs = "PartitionedCall:0"}, tf_saved_model.exported_names = ["serving_default"]} {
+    %cst = stablehlo.constant dense<0.000000e+00>: tensor<10x1024x3xf32>
+    %output, %min, %max, %histogram = "tf.CustomAggregator"(%arg0) <{calibration_method = 1 : i32, id = "skipping_id", max_percentile = 0.000000e+00 : f32, min_percentile = 0.000000e+00 : f32, num_bins = 0 : i32}> : (tensor<10x1x1024xf32>) -> (tensor<10x1x1024xf32>, tensor<f32>, tensor<f32>, tensor<0xi64>)
+    %0 = "tf.XlaCallModule"(%output, %cst) <{Sout = [#tf_type.shape<10x1x3>], dim_args_spec = [], disabled_checks = [], function_list = [], has_token_input_output = false, module = "", platforms = ["CPU"], version = 9 : i64}> {_entry_function = @composite_dot_general_with_relu_fn_1, _original_entry_function = "composite_dot_general_with_relu_fn_1", _quantization_method = "static_range_ptq { }", _stablehlo_module_attrs = {jax.uses_shape_polymorphism = true}, _tfl_quant_trait = "fully_quantizable"} : (tensor<10x1x1024xf32>, tensor<10x1024x3xf32>) -> tensor<10x1x3xf32>
+    %output_0, %min_1, %max_2, %histogram_3 = "tf.CustomAggregator"(%0) <{calibration_method = 1 : i32, id = "keeping_id", max_percentile = 0.000000e+00 : f32, min_percentile = 0.000000e+00 : f32, num_bins = 0 : i32}> : (tensor<10x1x3xf32>) -> (tensor<10x1x3xf32>, tensor<f32>, tensor<f32>, tensor<0xi64>)
+    return %output_0 : tensor<10x1x3xf32>
+  }
+  // CHECK-LABEL: @main
+  // CHECK: %[[CUSTOM_AGGREGATOR_0:.*]], %[[MIN_O:.*]], %[[MAX_O:.*]], %[[HISTOGRAM_0:.*]] = "tf.CustomAggregator"
+  // CKECK-SAME: <{calibration_method = 1 : i32, id = "skipping_id", max_percentile = 0.000000e+00 : f32, min_percentile = 0.000000e+00 : f32, num_bins = 0 : i32}>
+  // CHECK: %[[CUSTOM_AGGREGATOR_1:.*]], %[[MIN_1:.*]], %[[MAX_1:.*]], %[[HISTOGRAM_1:.*]] = "tf.CustomAggregator"
+  // CKECK-SAME: <{calibration_method = 1 : i32, id = "keeping_id", max_percentile = 0.000000e+00 : f32, min_percentile = 0.000000e+00 : f32, num_bins = 0 : i32}>
+  // CHECK: "tf.CalibrationStatisticsSaver"(%[[MIN_1]], %[[MAX_1]], %[[HISTOGRAM_1]])
+  // CHECK-SAME: <{calibration_methods = [1 : i32], ids = ["keeping_id"], output_file_path = "main_0.pb"}> : (tensor<f32>, tensor<f32>, tensor<0xi64>) -> ()
+  // CHECK: return
+
+  func.func private @composite_dot_general_with_relu_fn_1(%arg0: tensor<10x1x1024xf32>, %arg1: tensor<10x1024x3xf32>) -> tensor<10x1x3xf32> attributes {_from_xla_call_module, tf_quant.composite_function} {
+    %cst = stablehlo.constant dense<0.000000e+00> : tensor<10x1x3xf32>
+    %0 = stablehlo.dot_general %arg0, %arg1, batching_dims = [0] x [0], contracting_dims = [2] x [1], precision = [DEFAULT, DEFAULT] {mhlo.frontend_attributes = {grad_x = "false", grad_y = "false"}} : (tensor<10x1x1024xf32>, tensor<10x1024x3xf32>) -> tensor<10x1x3xf32>
+    %1 = stablehlo.maximum %0, %cst : tensor<10x1x3xf32>
+    return %1 : tensor<10x1x3xf32>
+  }
+  // CHECK-LABEL: func.func private @composite_dot_general_with_relu_fn_1
+  // CHECK-NOT: "tf.CalibrationStatisticsSaver"
+}
\ No newline at end of file
diff --git a/tensorflow/compiler/mlir/quantization/tensorflow/python/integration_test/quantize_model_test.py b/tensorflow/compiler/mlir/quantization/tensorflow/python/integration_test/quantize_model_test.py
index e9c18a66943ff9..53228d30bd53fc 100644
--- a/tensorflow/compiler/mlir/quantization/tensorflow/python/integration_test/quantize_model_test.py
+++ b/tensorflow/compiler/mlir/quantization/tensorflow/python/integration_test/quantize_model_test.py
@@ -2710,6 +2710,116 @@ def data_gen() -> repr_dataset.RepresentativeDataset:
     self.assertAllClose(new_outputs, got_outputs, atol=0.097)
     self.assertAllClose(new_outputs, expected_outputs, atol=0.057)
 
+  def test_reuse_calibration_data(self):
+    model = self._create_simple_gather_and_conv_model(
+        dtypes.int32, filter_shape=(2, 3, 3, 1024)
+    )
+    saved_model_save.save(model, self._input_saved_model_path)
+
+    data_gen = self._create_data_generator(
+        input_key='input_tensor',
+        shape=[50],
+        minval=0,
+        maxval=64,
+        dtype=dtypes.int32,
+    )
+
+    tags = {tag_constants.SERVING}
+
+    calibration_data_dir = self.create_tempdir('calibration_data').full_path
+    quantization_options = quant_opts_pb2.QuantizationOptions(
+        quantization_method=quant_opts_pb2.QuantizationMethod(
+            preset_method=_PresetMethod.METHOD_STATIC_RANGE_INT8
+        ),
+        tags=tags,
+        signature_keys=['serving_default'],
+        op_set=quant_opts_pb2.XLA,
+        force_graph_mode_calibration=True,
+        calibration_options=stablehlo_quant_config_pb2.CalibrationOptions(
+            calibration_method=_CalibrationMethod.CALIBRATION_METHOD_MIN_MAX,
+            calibration_data_dir=calibration_data_dir,
+        ),
+    )
+
+    # Run quantization the first time, calibration is expected to be run.
+    with self.assertLogs(level='INFO') as info_logs:
+      # Save the logger verbosity.
+      prev_log_level = logging.get_verbosity()
+      logging.set_verbosity(logging.INFO)
+      try:
+        converted_model1 = quantize_model.quantize(
+            self._input_saved_model_path,
+            self._output_saved_model_path,
+            quantization_options,
+            representative_dataset=data_gen,
+        )
+      finally:
+        # Restore the logger verbosity.
+        logging.set_verbosity(prev_log_level)
+
+      self.assertNotEmpty(info_logs.records)
+      self.assertTrue(
+          self._any_log_contains(
+              'Calibration step is executed in graph mode.',
+              info_logs.records,
+          )
+      )
+      self.assertIsNotNone(converted_model1)
+      self.assertCountEqual(
+          converted_model1.signatures._signatures.keys(), {'serving_default'}
+      )
+
+      output_loader = saved_model_loader.SavedModelLoader(
+          self._output_saved_model_path
+      )
+      output_graphdef = output_loader.get_meta_graph_def_from_tags(
+          tags
+      ).graph_def
+      self.assertTrue(self._contains_op(output_graphdef, 'XlaConvV2'))
+
+    # Run quantization the first time, calibration is expected to be skipped.
+    with self.assertLogs(level='INFO') as info_logs:
+      # Save the logger verbosity.
+      prev_log_level = logging.get_verbosity()
+      logging.set_verbosity(logging.INFO)
+      try:
+        converted_model2 = quantize_model.quantize(
+            self._input_saved_model_path,
+            self._output_saved_model_path,
+            quantization_options,
+            representative_dataset=data_gen,
+            overwrite_output_directory=True,
+        )
+      finally:
+        # Restore the logger verbosity.
+        logging.set_verbosity(prev_log_level)
+
+      self.assertNotEmpty(info_logs.records)
+      self.assertFalse(
+          self._any_log_contains(
+              'Calibration step is executed in graph mode.',
+              info_logs.records,
+          )
+      )
+      self.assertIsNotNone(converted_model2)
+      self.assertCountEqual(
+          converted_model2.signatures._signatures.keys(), {'serving_default'}
+      )
+
+      # Expect two models to produce the same results.
+      test_data = ops.convert_to_tensor(
+          np.random.uniform(low=0, high=64, size=(32)).astype(
+              dtypes.int32.as_numpy_dtype
+          )
+      )
+      new_outputs_1 = converted_model1.signatures['serving_default'](
+          input_tensor=test_data
+      )['output']
+      new_outputs_2 = converted_model2.signatures['serving_default'](
+          input_tensor=test_data
+      )['output']
+      self.assertAllClose(new_outputs_1, new_outputs_2)
+
   @test_util.run_in_graph_and_eager_modes
   def test_function_alias_preserved(self):
     model = self._create_conv2d_model(
diff --git a/tensorflow/compiler/mlir/quantization/tensorflow/python/quantize_model.cc b/tensorflow/compiler/mlir/quantization/tensorflow/python/quantize_model.cc
index e055d9dda95bbb..aef2732549ede2 100644
--- a/tensorflow/compiler/mlir/quantization/tensorflow/python/quantize_model.cc
+++ b/tensorflow/compiler/mlir/quantization/tensorflow/python/quantize_model.cc
@@ -84,6 +84,7 @@ using ::stablehlo::quantization::AddCalibrationStatistics;
 using ::stablehlo::quantization::ChangeToQuantizedFilename;
 using ::stablehlo::quantization::DebuggerConfig;
 using ::stablehlo::quantization::ExpandPresets;
+using ::stablehlo::quantization::IsCalibrationRequired;
 using ::stablehlo::quantization::PopulateDefaults;
 using ::stablehlo::quantization::QuantizationConfig;
 using ::stablehlo::quantization::io::CreateTmpDir;
@@ -163,7 +164,10 @@ absl::StatusOr<ExportedModel> ExportCalibrationModel(
   mlir::OwningOpRef<mlir::ModuleOp> cloned_module_ref(module_op.clone());
 
   TF_RETURN_IF_ERROR(
-      RunCalibrationPasses(*cloned_module_ref, *context, calibration_data_dir));
+      RunCalibrationPasses(*cloned_module_ref, *context, calibration_data_dir,
+                           quantization_options.calibration_options()
+                               .force_regenerate_calibration_data()));
+  if (!IsCalibrationRequired(*cloned_module_ref)) return ExportedModel();
 
   absl::StatusOr<ExportedModel> exported_model = ModuleOpToExportedModel(
       *cloned_module_ref, context, kTfQuantPtqPreCalibrationStepName,
@@ -457,16 +461,18 @@ absl::StatusOr<ExportedModel> QuantizeStaticRangePtq(
                           *function_aliases, calibration_data_dir));
 
   // Save and run the calibration model.
-  TF_ASSIGN_OR_RETURN(std::string precalibrated_saved_model_dir,
-                      CreateTmpDir());
-  py_function_library.SaveExportedModel(
-      precalibrated_saved_model_dir, calibration_exported_model,
-      saved_model_path, tags, signature_def_map);
-
-  py_function_library.RunCalibration(
-      precalibrated_saved_model_dir, signature_keys, tags,
-      quantization_options.force_graph_mode_calibration(),
-      representative_dataset_file_map_serialized);
+  if (calibration_exported_model.has_graph_def()) {
+    TF_ASSIGN_OR_RETURN(std::string calibration_saved_model_dir,
+                        CreateTmpDir());
+    py_function_library.SaveExportedModel(
+        calibration_saved_model_dir, calibration_exported_model,
+        saved_model_path, tags, signature_def_map);
+
+    py_function_library.RunCalibration(
+        calibration_saved_model_dir, signature_keys, tags,
+        quantization_options.force_graph_mode_calibration(),
+        representative_dataset_file_map_serialized);
+  }
 
   if (absl::Status status = AddCalibrationStatistics(
           *module_ref, calibration_data_dir,
diff --git a/tensorflow/compiler/mlir/quantization/tensorflow/python/quantize_model.py b/tensorflow/compiler/mlir/quantization/tensorflow/python/quantize_model.py
index 92ee947c1c3c2c..91c82710c489a9 100644
--- a/tensorflow/compiler/mlir/quantization/tensorflow/python/quantize_model.py
+++ b/tensorflow/compiler/mlir/quantization/tensorflow/python/quantize_model.py
@@ -638,7 +638,7 @@ def _populate_calibration_options(
   if calib_opts.calibration_data_dir:
     save_model.create_empty_output_dir(
         calib_opts.calibration_data_dir,
-        overwrite=True,
+        overwrite=calib_opts.force_regenerate_calibration_data,
     )
 
 

From b589eadd14cf20e46726d6d10fe98f50d11d8ccf Mon Sep 17 00:00:00 2001
From: Adrian Kuegel <akuegel@google.com>
Date: Mon, 13 May 2024 22:48:43 -0700
Subject: [PATCH 091/478] Add AlgebraicSimplifier pass directly in front of
 HloConstantFolding.

We have cases with DynamicSlice(Broadcast) which can be simplified away which
would otherwise be constant folded. The HloEvaluator is really slow for
DynamicSlice and Slice, as it will evalute the whole first operand, so if we
can simplify it away, it will help.

PiperOrigin-RevId: 633445573
---
 third_party/xla/xla/service/gpu/gpu_compiler.cc | 10 +++++++++-
 1 file changed, 9 insertions(+), 1 deletion(-)

diff --git a/third_party/xla/xla/service/gpu/gpu_compiler.cc b/third_party/xla/xla/service/gpu/gpu_compiler.cc
index 276a4f2c9929a7..a4086262b96a7a 100644
--- a/third_party/xla/xla/service/gpu/gpu_compiler.cc
+++ b/third_party/xla/xla/service/gpu/gpu_compiler.cc
@@ -574,7 +574,6 @@ absl::Status RunSPMDPasses(
         gpu_target_config.device_description.gpu_compute_capability();
     spmd_simplify.AddPass<GpuAlgebraicSimplifier>(
         layout_insensitive_algsimp_opts, gpu_version);
-
     spmd_simplify.AddPass<SortSimplifier>();
     spmd_simplify.AddPass<TupleSimplifier>();
     spmd_simplify.AddPass<ScatterExpander>(
@@ -587,6 +586,15 @@ absl::Status RunSPMDPasses(
     ReshapeMoverOptions reshape_mover_options;
     reshape_mover_options.reshape_of_1d_broadcast_is_cheap = true;
     spmd_simplify.AddPass<ReshapeMover>(reshape_mover_options);
+    // Run AlgebraicSimplifier directly before HloConstantFolding, because we
+    // need to simplify DynamicSlice(Broadcast) away. Constant folding of
+    // DynamicSlice can be quite costly, as the whole operand will be evaluated.
+    // We run AlgebraicSimplifier as HloPassFix to make sure all simplifications
+    // have been done before running HloConstantFolding. This is necessary
+    // because simplifications create new instructions which may not be visited
+    // in the same iteration of AlgebraicSimplifier.
+    spmd_simplify.AddPass<HloPassFix<AlgebraicSimplifier>>(
+        layout_insensitive_algsimp_opts);
     spmd_simplify.AddPass<HloConstantFolding>();
     spmd_simplify.AddPass<ConditionalSimplifier>();
 

From ac2c81c3d492420550bce77439ae5ebc77c39150 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Mon, 13 May 2024 23:19:02 -0700
Subject: [PATCH 092/478] Update ops-related pbtxt files.

PiperOrigin-RevId: 633451656
---
 tensorflow/core/ops/ops.pbtxt | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorflow/core/ops/ops.pbtxt b/tensorflow/core/ops/ops.pbtxt
index 53c362ed0ab819..68d78d0329a9ac 100644
--- a/tensorflow/core/ops/ops.pbtxt
+++ b/tensorflow/core/ops/ops.pbtxt
@@ -1,4 +1,4 @@
-go/debugstr    
+go/nodeserialize   
 op {
   name: "Abort"
   attr {

From f32e67e533bfe2d73c3658e1c2716bbc1b2c7baa Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Tue, 14 May 2024 00:11:15 -0700
Subject: [PATCH 093/478] Improve documentation of mlir ops.

PiperOrigin-RevId: 633462768
---
 .../tensorflow/ir/host_runtime/tfrt_ops.td    | 20 +++++++------------
 .../compiler/mlir/tfrt/ir/mlrt/tf_mlrt_ops.td | 19 ++++++------------
 2 files changed, 13 insertions(+), 26 deletions(-)

diff --git a/tensorflow/compiler/mlir/tensorflow/ir/host_runtime/tfrt_ops.td b/tensorflow/compiler/mlir/tensorflow/ir/host_runtime/tfrt_ops.td
index e3736e96851262..e46a6500dfd516 100644
--- a/tensorflow/compiler/mlir/tensorflow/ir/host_runtime/tfrt_ops.td
+++ b/tensorflow/compiler/mlir/tensorflow/ir/host_runtime/tfrt_ops.td
@@ -94,23 +94,17 @@ Empty strings indicate that they are non-partitioned tensors.}]>:$shape_and_slic
 
 
 def TF_IfrtLoadVariableOp : TF_Op<"IfrtLoadVariable", [Pure]> {
-  let summary = "Loads a restored variable tensor as an IFRT array and tensor future";
+  let summary = "Loads a restored variable tensor as a tensor future";
 
   let description = [{
-    This op loads a variable tensor as an IFRT array and binds it with the specified name.
+    This op loads a restored variable tensor as a tensor future. It is a
+    replacement of `tf.ReadVariableOp`.
 
-    This op is an replacement of `tf.ReadVariableOp` in the case that a constant
-    variable tensor is an input to the tpu program invoked by `tf.IfrtCall`.
+    This op returns a scalar string tensor containing the restored variable name, which can be
+    used as a key within the runtime, as well as a future for the tensor.
 
-    After a `tf.ReadVariableOp` is lowered into `tf.IfrtLoadVariableOp`, the `tf.IfrtCall` kernel
-    will bind the loaded IFRT array by name with the tpu program's input.
-
-    `tf.IfrtLoadVariableOp` converts the tensor into an IFRT array based on device and sharding
-    configuration specified in `VariableDeviceShardingConfigProto`.
-
-    This op returns a scalar string tensor containing the loaded variable name, which can be
-    used as a key to look for the loaded IFRT array in runtime and a restored tensor, which
-    maybe lowered to a future by runtime.
+    The `tf.IfrtCall` kernel uses the output $array_key.
+    Other ops executed by TFRT may make use of $tensor_future.
   }];
 
   // TODO(b/339423851) Redefine the IfrtLoadVariableOp as it doesn't require the
diff --git a/tensorflow/compiler/mlir/tfrt/ir/mlrt/tf_mlrt_ops.td b/tensorflow/compiler/mlir/tfrt/ir/mlrt/tf_mlrt_ops.td
index ae4fa4be4f78c9..0659143f49b39b 100644
--- a/tensorflow/compiler/mlir/tfrt/ir/mlrt/tf_mlrt_ops.td
+++ b/tensorflow/compiler/mlir/tfrt/ir/mlrt/tf_mlrt_ops.td
@@ -449,21 +449,14 @@ def IfrtLoadVariableOp: TensorflowMlrt_Op<"ifrt_load_variable", [Pure]> {
   let summary = "Loads a variable tensor as an IFRT array for mlrt";
 
   let description = [{
-    This is the MLRT version of tf.IfrtLoadVariableOp.
+    This op loads a restored variable tensor as a tensor future. It is a
+    replacement of `tf.ReadVariableOp`.
 
-    This op loads a variable tensor as an IFRT array and binds it with the specified name.
+    This op returns a scalar string tensor containing the restored variable name, which can be
+    used as a key within the runtime, as well as a future for the tensor.
 
-    This op is an replacement of `tf.ReadVariableOp` in the case that a constant
-    variable tensor is an input to the tpu program invoked by `tf.IfrtCall`.
-
-    After a `tf.ReadVariableOp` is lowered into `tf.IfrtLoadVariableOp`, the `tf.IfrtCall` kernel
-    will bind the loaded IFRT array by name with the tpu program's input.
-
-    `tf.IfrtLoadVariableOp` converts the tensor into an IFRT array based on device and sharding
-    configuration specified in `VariableDeviceShardingConfigProto`.
-
-    This op returns a scalar string tensor as a key for user to look for the loaded array
-    and a future containing the restored tensor.
+    The `tf.IfrtCall` kernel uses the output $array_key.
+    Other ops executed by TFRT may make use of $tensor_future.
   }];
 
   let arguments = (ins

From 24fb166319aec315431bb4fc826c1bfa3a9da43a Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Tue, 14 May 2024 01:17:37 -0700
Subject: [PATCH 094/478] Update ops-related pbtxt files.

PiperOrigin-RevId: 633477018
---
 tensorflow/core/ops/ops.pbtxt | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorflow/core/ops/ops.pbtxt b/tensorflow/core/ops/ops.pbtxt
index 68d78d0329a9ac..5019fe2a5b2fad 100644
--- a/tensorflow/core/ops/ops.pbtxt
+++ b/tensorflow/core/ops/ops.pbtxt
@@ -1,4 +1,4 @@
-go/nodeserialize   
+go/debugproto    
 op {
   name: "Abort"
   attr {

From d83f8918629653a50d239388a84adc676620bc1c Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Tue, 14 May 2024 01:19:17 -0700
Subject: [PATCH 095/478] Automated Code Change

PiperOrigin-RevId: 633477369
---
 tensorflow/lite/delegates/gpu/common/data_type_test.cc | 1 -
 tensorflow/lite/delegates/gpu/common/flops_util.cc     | 2 ++
 2 files changed, 2 insertions(+), 1 deletion(-)

diff --git a/tensorflow/lite/delegates/gpu/common/data_type_test.cc b/tensorflow/lite/delegates/gpu/common/data_type_test.cc
index a88eb38047eda8..f7ee35c73145dd 100644
--- a/tensorflow/lite/delegates/gpu/common/data_type_test.cc
+++ b/tensorflow/lite/delegates/gpu/common/data_type_test.cc
@@ -15,7 +15,6 @@ limitations under the License.
 
 #include "tensorflow/lite/delegates/gpu/common/data_type.h"
 
-#include <gmock/gmock.h>
 #include <gtest/gtest.h>
 
 namespace tflite {
diff --git a/tensorflow/lite/delegates/gpu/common/flops_util.cc b/tensorflow/lite/delegates/gpu/common/flops_util.cc
index 88d7edad88fc57..3807f6b6a46274 100644
--- a/tensorflow/lite/delegates/gpu/common/flops_util.cc
+++ b/tensorflow/lite/delegates/gpu/common/flops_util.cc
@@ -15,6 +15,8 @@ limitations under the License.
 
 #include "tensorflow/lite/delegates/gpu/common/flops_util.h"
 
+#include "tensorflow/lite/delegates/gpu/common/shape.h"
+
 namespace tflite {
 namespace gpu {
 

From 16244310e5a73e8f187d6dab487fee0e80d4afd9 Mon Sep 17 00:00:00 2001
From: TJ Xu <tjx@nvidia.com>
Date: Tue, 14 May 2024 01:48:26 -0700
Subject: [PATCH 096/478] PR #11563: [NVIDIA GPU] Improve GPU collective matmul
 to support all-gather having multiple users

Imported from GitHub PR https://github.com/openxla/xla/pull/11563

We have identified another optimization opportunity for gpt-3 using collective matmul, in the backward pass, the all-gather has multiple dot users but current spmd will duplicate multiple collective matmul loops. We'd like this transformation:
before:
```
  //                       input
  //                       /    |
  //                      /     |
  //                     AG    windowed loop
  //                     /
  //                    /
  //                   dot

```
after:
```
  //                       input
  //                       |
  //                       |
  //                     windowed loop
  //                       |
  //                       |
  //                      dot
```
This is advantageous since the chained dot can fully utilize all the resource on the GPU while comm is hidden by the first collective matmul loop.

We introduced an option to turn off CM loop duplication in SPMD and rewrite the graph to desired pattern in the gpu_windowed_einsum_handler pass.
Copybara import of the project:

--
986ac94ab44d31f6d11ec6f135f6cfb2e5636d80 by TJ <tjx@nvidia.com>:

Moved most of changes to gpu pass

--
44e81df91c235cac635f334c89d1d8a117ac6511 by TJ <tjx@nvidia.com>:

Added e2e test for windowed einsum
Minimized unit test hlo

--
8fc24a479de7515f532f36de8ffbcce49516c154 by TJ <tjx@nvidia.com>:

Added explanations for spmd tests and dot_handler to skip multiple
consumers

--
142d84d54db2b6291484443e43913d86c44a485c by TJ <tjx@nvidia.com>:

move windowed einsum test to stateful_rng_spmd_partitioner_test

--
8b9fc43746136b40a814d93bf8086a687490fd7f by TJ <tjx@nvidia.com>:

Changed e2e test back to include reducescatter

Merging this change closes #11563

PiperOrigin-RevId: 633483864
---
 third_party/xla/xla/service/gpu/BUILD         |   2 +
 .../xla/xla/service/gpu/gpu_compiler.cc       |   4 +-
 .../gpu/gpu_windowed_einsum_handler.cc        | 162 ++++++++++++++++++
 .../service/gpu/gpu_windowed_einsum_handler.h |  10 +-
 .../gpu/gpu_windowed_einsum_handler_test.cc   |  95 ++++++++++
 .../xla/xla/service/spmd/dot_handler.cc       |  61 +++++--
 .../xla/xla/service/spmd/spmd_partitioner.h   |   4 +
 .../spmd/stateful_rng_spmd_partitioner.h      |  27 ++-
 .../stateful_rng_spmd_partitioner_test.cc     |  60 ++++++-
 .../xla/xla/tests/collective_ops_test_e2e.cc  |  80 +++++++++
 10 files changed, 479 insertions(+), 26 deletions(-)

diff --git a/third_party/xla/xla/service/gpu/BUILD b/third_party/xla/xla/service/gpu/BUILD
index 34f2210e1d1b04..5faf098331bb14 100644
--- a/third_party/xla/xla/service/gpu/BUILD
+++ b/third_party/xla/xla/service/gpu/BUILD
@@ -6070,6 +6070,8 @@ xla_cc_test(
         "//xla:util",
         "//xla:xla_data_proto_cc",
         "//xla/hlo/ir:hlo",
+        "//xla/service:pattern_matcher",
+        "//xla/service:pattern_matcher_gmock",
         "//xla/tests:hlo_test_base",
         "@com_google_absl//absl/algorithm:container",
         "@com_google_absl//absl/strings:string_view",
diff --git a/third_party/xla/xla/service/gpu/gpu_compiler.cc b/third_party/xla/xla/service/gpu/gpu_compiler.cc
index a4086262b96a7a..6c3e2e70e05422 100644
--- a/third_party/xla/xla/service/gpu/gpu_compiler.cc
+++ b/third_party/xla/xla/service/gpu/gpu_compiler.cc
@@ -640,7 +640,9 @@ absl::Status RunSPMDPasses(
             .xla_gpu_threshold_for_windowed_einsum_mib(),
         hlo_module->config()
             .debug_options()
-            .xla_gpu_multi_streamed_windowed_einsum());
+            .xla_gpu_multi_streamed_windowed_einsum(),
+        /*skip_checking_windowed_einsum_users=*/true,
+        /*disable_ag_rewrite_for_multiple_consumers=*/true);
     spmd_pipeline.AddPass<CollectivePermuteMotion>();
     return spmd_pipeline.Run(hlo_module).status();
   } else {
diff --git a/third_party/xla/xla/service/gpu/gpu_windowed_einsum_handler.cc b/third_party/xla/xla/service/gpu/gpu_windowed_einsum_handler.cc
index fa5339fc007f44..875d4d5eed8bfb 100644
--- a/third_party/xla/xla/service/gpu/gpu_windowed_einsum_handler.cc
+++ b/third_party/xla/xla/service/gpu/gpu_windowed_einsum_handler.cc
@@ -20,6 +20,7 @@ limitations under the License.
 #include "absl/container/flat_hash_set.h"
 #include "absl/status/status.h"
 #include "absl/strings/string_view.h"
+#include "xla/hlo/ir/dfs_hlo_visitor_with_default.h"
 #include "xla/hlo/ir/hlo_computation.h"
 #include "xla/hlo/ir/hlo_instruction.h"
 #include "xla/hlo/ir/hlo_module.h"
@@ -142,6 +143,136 @@ absl::StatusOr<bool> HandleAgWindowedEinsumLoop(HloComputation* comp,
   return changed;
 }
 
+absl::Status ProcessWindowedEinsumLoopForActivationCaching(
+    GpuWindowedEinsumHandler::WindowedEinsumAgLoops& ag_loop) {
+  HloInstruction* loop = ag_loop.loop;
+  // Transform the while body to cache the allgathered result in the
+  // output buffer to be consumed by the dot
+  HloComputation* while_body = loop->while_body();
+  HloInstruction* input_gte;
+  for (HloInstruction* gte : while_body->parameter_instruction(0)->users()) {
+    if (gte->tuple_index() == 0) {
+      input_gte = gte;
+    }
+  }
+  // Get the output operand of the full buffer.
+  HloInstruction* root = while_body->root_instruction();
+  // The full buffer that we will use to cache the accumulated activation
+  // is the 4th operand in the output tuple.
+  int64_t full_cache_buffer_index = 3;
+  HloInstruction* full_buffer_output_gte =
+      root->mutable_operand(full_cache_buffer_index);
+  HloInstruction* new_full_buffer_output;
+  // Find the DUS in the loop body and re-use the slice indices
+  // This should just be a constant(0)
+  HloInstruction* dus_boundary_constant;
+  for (HloInstruction* inst : while_body->MakeInstructionPostOrder()) {
+    HloInstruction* slice_indices;
+    // If we have a DUS(PARAM,DS) pattern, we need to update the output
+    // buffer with the first slice.
+    if (Match(inst,
+              m::DynamicUpdateSlice(
+                  m::GetTupleElement(m::Parameter()), m::Op(),
+                  m::Constant(&dus_boundary_constant),
+                  m::Reshape(m::DynamicSlice(&slice_indices, m::Op(), m::Op())),
+                  m::Op()))) {
+      slice_indices = while_body->AddInstruction(HloInstruction::CreateReshape(
+          dus_boundary_constant->shape(), slice_indices));
+      VLOG(5) << "Created slice op for first slice: "
+              << slice_indices->ToString();
+      full_buffer_output_gte =
+          while_body->AddInstruction(HloInstruction::CreateDynamicUpdateSlice(
+              full_buffer_output_gte->shape(), full_buffer_output_gte,
+              input_gte,
+              {dus_boundary_constant, slice_indices, dus_boundary_constant}));
+    }
+    // If we have a DUS(DUS,DS) pattern, then the einsum loop is
+    // unrolled, we need to update the output buffer again with the
+    // second slice. Since the second slice will have different indices,
+    // we need to re-capture slice_indices.
+    if (Match(inst,
+              m::DynamicUpdateSlice(
+                  m::DynamicUpdateSlice(), m::Op(), m::Constant(),
+                  m::Reshape(m::DynamicSlice(&slice_indices, m::Op(), m::Op())),
+                  m::Op()))) {
+      slice_indices = while_body->AddInstruction(HloInstruction::CreateReshape(
+          dus_boundary_constant->shape(), slice_indices));
+      VLOG(5) << "Created slice op for second slice: "
+              << slice_indices->ToString();
+      // The slice we need this time is the output of the first
+      // collective-permute
+      HloInstruction* cp_output;
+      for (HloInstruction* gte_user : input_gte->users()) {
+        if (gte_user->opcode() == HloOpcode::kCollectivePermute) {
+          cp_output = gte_user;
+          break;
+        }
+      }
+      new_full_buffer_output =
+          while_body->AddInstruction(HloInstruction::CreateDynamicUpdateSlice(
+              full_buffer_output_gte->shape(), full_buffer_output_gte,
+              cp_output,
+              {dus_boundary_constant, slice_indices, dus_boundary_constant}));
+    }
+  }
+  TF_RETURN_IF_ERROR(root->ReplaceOperandWith(full_cache_buffer_index,
+                                              new_full_buffer_output));
+  return OkStatus();
+}
+
+class WindowedEinsumVisitor : public DfsHloRewriteVisitor {
+ public:
+  explicit WindowedEinsumVisitor(
+      std::vector<GpuWindowedEinsumHandler::WindowedEinsumAgLoops>&
+          all_ag_loops)
+      : all_ag_loops_(all_ag_loops) {}
+  // Rewrites a allgather-dot pattern that shares the same operand
+  // with a windowed einsum loop to consume the output of the loop
+  // and remove the all-gather.
+  absl::Status HandleDot(HloInstruction* dot) override {
+    CHECK_EQ(dot->opcode(), HloOpcode::kDot);
+    for (GpuWindowedEinsumHandler::WindowedEinsumAgLoops ag_loop :
+         all_ag_loops_) {
+      HloInstruction* loop = ag_loop.loop;
+      HloInstruction* ag_operand = nullptr;
+
+      if (Match(dot, m::Dot(m::AllGather(&ag_operand), m::Op())) ||
+          Match(dot, m::Dot(m::Op(), m::AllGather(&ag_operand)))) {
+        HloInstruction* windowed_lhs =
+            loop->mutable_operand(0)->mutable_operand(0);
+        HloInstruction* ag_with_shared_operand = nullptr;
+        if (ag_operand && ag_operand->mutable_operand(0) == windowed_lhs) {
+          ag_with_shared_operand = ag_operand;
+        }
+
+        if (!ag_with_shared_operand) {
+          continue;
+        }
+
+        VLOG(5) << "Found all-gather that shares the same operand with a "
+                   "windowed einsum loop : "
+                << loop->ToString();
+        int64_t cache_output_index = dot->operand_index(ag_with_shared_operand);
+        HloComputation* comp = dot->parent();
+        HloInstruction* new_gte = comp->AddInstruction(
+            HloInstruction::CreateGetTupleElement(loop, 3));
+        TF_RETURN_IF_ERROR(
+            dot->ReplaceOperandWith(cache_output_index, new_gte));
+        TF_RETURN_IF_ERROR(comp->RemoveInstruction(ag_with_shared_operand));
+        if (!ag_loop.consumed) {
+          TF_RETURN_IF_ERROR(
+              ProcessWindowedEinsumLoopForActivationCaching(ag_loop));
+          ag_loop.consumed = true;
+        }
+      }
+    }
+    return absl::OkStatus();
+  }
+
+ private:
+  std::vector<GpuWindowedEinsumHandler::WindowedEinsumAgLoops>& all_ag_loops_;
+};
+
 }  // namespace
 
 absl::StatusOr<bool> GpuWindowedEinsumHandler::Run(
@@ -163,9 +294,40 @@ absl::StatusOr<bool> GpuWindowedEinsumHandler::Run(
       VLOG(5) << "Processing computation: " << comp->name();
       TF_ASSIGN_OR_RETURN(bool comp_result,
                           HandleAgWindowedEinsumLoop(comp, stream_id));
+      all_ag_loops_.push_back(
+          WindowedEinsumAgLoops(comp->WhileCallInstruction()));
       changed = comp_result;
     }
   }
+  // Now that we have processed all loops, we can check if there are any
+  // allgather-dot pattern that we can optimize. We'd want to transform:
+  //                       input
+  //                       /    |
+  //                      /     |
+  //                     AG    windowed loop
+  //                     /
+  //                    /
+  //                   dot
+  // to:
+  //                       input
+  //                       |
+  //                       |
+  //                     windowed loop
+  //                       |
+  //                       |
+  //                      dot
+  // The windowed einsum loop will also be rewritten to output the full input to
+  // be consumed by the dot.
+  // This is advantageous since the chained dot can fully utilize all the
+  // resources on the GPU while comm is hidden by the first collective matmul
+  // loop.
+  for (HloComputation* comp :
+       module->MakeNonfusionComputations(execution_threads)) {
+    WindowedEinsumVisitor visitor(all_ag_loops_);
+    TF_RETURN_IF_ERROR(comp->Accept(&visitor));
+    changed |= visitor.changed();
+  }
+
   XLA_VLOG_LINES(
       5, "GpuWindowedEinsumHandler::Run(), after:\n" + module->ToString());
   return changed;
diff --git a/third_party/xla/xla/service/gpu/gpu_windowed_einsum_handler.h b/third_party/xla/xla/service/gpu/gpu_windowed_einsum_handler.h
index 87ec1474d576f7..8f42cbc00fb3ab 100644
--- a/third_party/xla/xla/service/gpu/gpu_windowed_einsum_handler.h
+++ b/third_party/xla/xla/service/gpu/gpu_windowed_einsum_handler.h
@@ -38,16 +38,24 @@ class GpuWindowedEinsumHandler : public HloModulePass {
     return "gpu-windowed-einsum-handler";
   }
 
+  struct WindowedEinsumAgLoops {
+    WindowedEinsumAgLoops(HloInstruction* loop) : loop(loop) {}
+    HloInstruction* loop;
+    bool consumed = false;
+  };
+
   using HloPassInterface::Run;
   absl::StatusOr<bool> Run(
       HloModule* module,
       const absl::flat_hash_set<absl::string_view>& execution_threads) override;
 
- private:
   constexpr static const char* kWindowedEinsumRsLoopName =
       "windowed_dot_general_body_rs";
   constexpr static const char* kWindowedEinsumAgLoopName =
       "windowed_dot_general_body_ag";
+
+ private:
+  std::vector<WindowedEinsumAgLoops> all_ag_loops_;
 };
 
 }  // namespace xla::gpu
diff --git a/third_party/xla/xla/service/gpu/gpu_windowed_einsum_handler_test.cc b/third_party/xla/xla/service/gpu/gpu_windowed_einsum_handler_test.cc
index c70fbf2b08d126..b0bce884392253 100644
--- a/third_party/xla/xla/service/gpu/gpu_windowed_einsum_handler_test.cc
+++ b/third_party/xla/xla/service/gpu/gpu_windowed_einsum_handler_test.cc
@@ -23,12 +23,16 @@ limitations under the License.
 #include "xla/hlo/ir/hlo_instruction.h"
 #include "xla/hlo/ir/hlo_module.h"
 #include "xla/service/gpu/backend_configs.pb.h"
+#include "xla/service/pattern_matcher.h"
+#include "xla/service/pattern_matcher_gmock.h"
 #include "xla/tests/hlo_test_base.h"
 #include "tsl/platform/statusor.h"
 
 namespace xla::gpu {
 namespace {
 
+namespace m = ::xla::match;
+
 using GpuWindowedEinsumHanlderTest = HloTestBase;
 
 HloInstruction* FindInstructionByName(HloComputation* comp, std::string name) {
@@ -193,5 +197,96 @@ ENTRY main.9_spmd {
       cp1->backend_config<GpuBackendConfig>()->force_earliest_schedule());
 }
 
+TEST_F(GpuWindowedEinsumHanlderTest, AgLoopsMultipleConsumersAreChained) {
+  constexpr absl::string_view kHloString = R"(
+HloModule pjit__unnamed_wrapped_function_, entry_computation_layout={(bf16[2,512,24576]{2,1,0}, bf16[24576,24576]{1,0}, bf16[24576,24576]{1,0})->bf16[2,2048,24576]{2,1,0}}, num_partitions=4
+
+windowed_dot_general_body_ag {
+  param.1 = (bf16[2,512,24576]{2,1,0}, bf16[24576,24576]{1,0}, bf16[2,2048,24576]{2,1,0}, bf16[2,2048,24576]{2,1,0}, u32[]) parameter(0)
+  get-tuple-element.1 = bf16[2,512,24576]{2,1,0} get-tuple-element(param.1), index=0
+  collective-permute = bf16[2,512,24576]{2,1,0} collective-permute(get-tuple-element.1), channel_id=2, source_target_pairs={{0,3},{1,0},{2,1},{3,2}}
+  collective-permute.1 = bf16[2,512,24576]{2,1,0} collective-permute(collective-permute), channel_id=3, source_target_pairs={{0,3},{1,0},{2,1},{3,2}}
+  get-tuple-element.2 = bf16[24576,24576]{1,0} get-tuple-element(param.1), index=1
+  get-tuple-element.3 = bf16[2,2048,24576]{2,1,0} get-tuple-element(param.1), index=2
+  dot = bf16[2,512,24576]{2,1,0} dot(get-tuple-element.1, get-tuple-element.2), lhs_contracting_dims={2}, rhs_contracting_dims={0}
+  constant.2 = s32[] constant(0)
+  constant.3 = s32[4]{0} constant({0, 512, 1024, 1536})
+  get-tuple-element.5 = u32[] get-tuple-element(param.1), index=4
+  partition-id = u32[] partition-id()
+  add = u32[] add(get-tuple-element.5, partition-id)
+  constant.1 = u32[] constant(4)
+  remainder = u32[] remainder(add, constant.1)
+  dynamic-slice = s32[1]{0} dynamic-slice(constant.3, remainder), dynamic_slice_sizes={1}
+  reshape = s32[] reshape(dynamic-slice)
+  dynamic-update-slice = bf16[2,2048,24576]{2,1,0} dynamic-update-slice(get-tuple-element.3, dot, constant.2, reshape, constant.2)
+  dot.1 = bf16[2,512,24576]{2,1,0} dot(collective-permute, get-tuple-element.2), lhs_contracting_dims={2}, rhs_contracting_dims={0}
+  constant.5 = u32[] constant(1)
+  add.1 = u32[] add(get-tuple-element.5, constant.5)
+  add.2 = u32[] add(add.1, partition-id)
+  remainder.1 = u32[] remainder(add.2, constant.1)
+  dynamic-slice.1 = s32[1]{0} dynamic-slice(constant.3, remainder.1), dynamic_slice_sizes={1}
+  reshape.1 = s32[] reshape(dynamic-slice.1)
+  dynamic-update-slice.1 = bf16[2,2048,24576]{2,1,0} dynamic-update-slice(dynamic-update-slice, dot.1, constant.2, reshape.1, constant.2)
+  get-tuple-element.4 = bf16[2,2048,24576]{2,1,0} get-tuple-element(param.1), index=3
+  add.3 = u32[] add(add.1, constant.5)
+  ROOT tuple = (bf16[2,512,24576]{2,1,0}, bf16[24576,24576]{1,0}, bf16[2,2048,24576]{2,1,0}, bf16[2,2048,24576]{2,1,0}, u32[]) tuple(collective-permute.1, get-tuple-element.2, dynamic-update-slice.1, get-tuple-element.4, add.3)
+} // windowed_dot_general_body_ag
+
+windowed_dot_general_cond_ag {
+  param = (bf16[2,512,24576]{2,1,0}, bf16[24576,24576]{1,0}, bf16[2,2048,24576]{2,1,0}, bf16[2,2048,24576]{2,1,0}, u32[]) parameter(0)
+  get-tuple-element = u32[] get-tuple-element(param), index=4
+  constant = u32[] constant(4)
+  ROOT compare = pred[] compare(get-tuple-element, constant), direction=LT
+}
+
+ENTRY main.12_spmd {
+  param.4 = bf16[2,512,24576]{2,1,0} parameter(0), sharding={devices=[1,4,1]<=[4]}
+  param.5 = bf16[24576,24576]{1,0} parameter(1), sharding={devices=[1,4]<=[4]}
+  constant.22 = bf16[] constant(0)
+  broadcast = bf16[2,2048,24576]{2,1,0} broadcast(constant.22), dimensions={}
+  constant.24 = u32[] constant(0)
+  tuple.2 = (bf16[2,512,24576]{2,1,0}, bf16[24576,24576]{1,0}, bf16[2,2048,24576]{2,1,0}, bf16[2,2048,24576]{2,1,0}, u32[]) tuple(param.4, param.5, broadcast, broadcast, constant.24)
+  while = (bf16[2,512,24576]{2,1,0}, bf16[24576,24576]{1,0}, bf16[2,2048,24576]{2,1,0}, bf16[2,2048,24576]{2,1,0}, u32[]) while(tuple.2), condition=windowed_dot_general_cond_ag, body=windowed_dot_general_body_ag
+  get-tuple-element.13 = bf16[2,2048,24576]{2,1,0} get-tuple-element(while), index=2
+  copy.1 = bf16[2,2048,24576]{2,1,0} copy(get-tuple-element.13)
+  all-gather = bf16[2,2048,24576]{2,1,0} all-gather(param.4), channel_id=1, replica_groups={{0,1,2,3}}, dimensions={1}, use_global_device_ids=true
+  param.6 = bf16[24576,24576]{1,0} parameter(2), sharding={devices=[1,4]<=[4]}
+  ROOT dot.7 = bf16[2,2048,24576]{2,1,0} dot(all-gather, param.6), lhs_contracting_dims={2}, rhs_contracting_dims={0}
+}
+)";
+
+  TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
+                          ParseAndReturnVerifiedModule(kHloString));
+
+  GpuWindowedEinsumHandler gpu_handler;
+  bool changed;
+  TF_ASSERT_OK_AND_ASSIGN(changed, gpu_handler.Run(module.get()));
+  EXPECT_TRUE(changed);
+
+  HloInstruction* ag_loop =
+      FindInstructionByName(module->entry_computation(), "while");
+  HloInstruction* inst =
+      FindInstructionByName(module->entry_computation(), "dot.7");
+  // dot.7 should now consume output of the windowed einsum while loop.
+  EXPECT_EQ(inst->operand(0)->opcode(), HloOpcode::kGetTupleElement);
+  EXPECT_EQ(inst->operand(0)->tuple_index(), 3);
+  EXPECT_EQ(inst->operand(0)->operand(0), ag_loop);
+
+  // while loop's root should now have a chain of DUS.
+  HloInstruction* ag_while_root = ag_loop->while_body()->root_instruction();
+  EXPECT_THAT(ag_while_root,
+              GmockMatch(m::Tuple(
+                  m::Op(), m::Op(), m::Op(),
+                  m::DynamicUpdateSlice(
+                      m::DynamicUpdateSlice(
+                          m::GetTupleElement(m::Parameter())
+                              .WithPredicate([](const HloInstruction* instr) {
+                                return instr->tuple_index() == 3;
+                              }),
+                          m::Op(), m::Op(), m::Op(), m::Op()),
+                      m::Op(), m::Op(), m::Op(), m::Op()),
+                  m::Op())));
+}
+
 }  // namespace
 }  // namespace xla::gpu
diff --git a/third_party/xla/xla/service/spmd/dot_handler.cc b/third_party/xla/xla/service/spmd/dot_handler.cc
index 8115a78794672d..47909ff72a5801 100644
--- a/third_party/xla/xla/service/spmd/dot_handler.cc
+++ b/third_party/xla/xla/service/spmd/dot_handler.cc
@@ -1909,22 +1909,57 @@ absl::StatusOr<HloInstruction*> PartitionBaseCase(
     }
   }
 
+  // If we see a dot that shares the same operand with a windowed einsum ag loop
+  // and disable_ag_rewrite_for_multiple_consumers is true. We skip rewriting
+  // the current dot. We also skip any reshape operand as long as it only has
+  // the lhs or rhs of the dot as the only user since reshape ops won't change
+  // the functional meaning of the pattern.
+  auto has_reshape_operand = [](PartitionedHlo& hlo) -> bool {
+    return hlo.hlo()->opcode() == HloOpcode::kReshape ||
+           hlo.hlo()->opcode() == HloOpcode::kBitcast ||
+           hlo.hlo()->opcode() == HloOpcode::kTranspose;
+  };
+  bool should_skip_windowed_einsum = false;
+  if (options.disable_ag_rewrite_for_multiple_consumers) {
+    auto lhs_operand =
+        has_reshape_operand(lhs) ? lhs.hlo()->operand(0) : lhs.hlo();
+    auto rhs_operand =
+        has_reshape_operand(rhs) ? rhs.hlo()->operand(0) : rhs.hlo();
+    for (auto loop : *windowed_dot_general_loops) {
+      if (loop.while_loop->while_body()->name().find(
+              "windowed_dot_general_body_ag") == 0) {
+        auto cm_lhs = loop.while_loop->operand(0)->operand(0);
+        if (cm_lhs == lhs_operand || cm_lhs == rhs_operand) {
+          VLOG(2) << "Skip processing: " << original_hlo->ToString();
+          VLOG(2) << "It shares the same operand with "
+                  << loop.while_loop->ToString()
+                  << " and disable_ag_rewrite_for_multiple_consumers is set to "
+                     "true.";
+          should_skip_windowed_einsum = true;
+        }
+      }
+    }
+  }
+
   // Hard limit on iteration count based on empirical data (above this amount
   // there's pretty significant overhead).
   constexpr int64_t kMaxIterations = 32;
-  std::optional<WindowedEinsumConfig> e_config = GetWindowedEinsumConfiguration(
-      num_partitions, output_lhs_non_contracting_partitions,
-      output_rhs_non_contracting_partitions, rhs_contracting_partitions,
-      rhs_non_contracting_partitions, rhs_batch_partitions,
-      lhs_contracting_partitions, lhs_non_contracting_partitions,
-      lhs_batch_partitions, ShapeSizeInBytes(rhs.base_shape()),
-      ShapeSizeInBytes(lhs.base_shape()), ShapeSizeInBytes(output_base_shape),
-      options, output_sharding_transposed_to_match_lhs,
-      output_sharding_transposed_to_match_rhs,
-      lhs_sharding_transposed_to_match_rhs,
-      rhs_sharding_transposed_to_match_lhs, lhs_sharding, rhs_sharding,
-      conv_window, dims_mapping, visitor->call_graph(), kMaxIterations,
-      original_hlo, &lhs, &rhs, create_sharded_dot, b, module, visitor);
+  std::optional<WindowedEinsumConfig> e_config = std::nullopt;
+  if (!should_skip_windowed_einsum) {
+    e_config = GetWindowedEinsumConfiguration(
+        num_partitions, output_lhs_non_contracting_partitions,
+        output_rhs_non_contracting_partitions, rhs_contracting_partitions,
+        rhs_non_contracting_partitions, rhs_batch_partitions,
+        lhs_contracting_partitions, lhs_non_contracting_partitions,
+        lhs_batch_partitions, ShapeSizeInBytes(rhs.base_shape()),
+        ShapeSizeInBytes(lhs.base_shape()), ShapeSizeInBytes(output_base_shape),
+        options, output_sharding_transposed_to_match_lhs,
+        output_sharding_transposed_to_match_rhs,
+        lhs_sharding_transposed_to_match_rhs,
+        rhs_sharding_transposed_to_match_lhs, lhs_sharding, rhs_sharding,
+        conv_window, dims_mapping, visitor->call_graph(), kMaxIterations,
+        original_hlo, &lhs, &rhs, create_sharded_dot, b, module, visitor);
+  }
   if (e_config) {
     VLOG(2) << "Emit windowed dot.";
     return EmitWindowedDotGeneral(
diff --git a/third_party/xla/xla/service/spmd/spmd_partitioner.h b/third_party/xla/xla/service/spmd/spmd_partitioner.h
index 4a8dc50b8b7fec..c137a821c333a8 100644
--- a/third_party/xla/xla/service/spmd/spmd_partitioner.h
+++ b/third_party/xla/xla/service/spmd/spmd_partitioner.h
@@ -89,6 +89,10 @@ struct SpmdPartitionerOptions {
   bool enable_windowed_einsum_for_all_gather = true;
   // Enables windowed einsum for result reduce-scatter.
   bool enable_windowed_einsum_for_reduce_scatter = true;
+
+  // Whether disable rewrite for dots that share the same
+  // operand as an already rewritten windowed einsum loop.
+  bool disable_ag_rewrite_for_multiple_consumers = false;
 };
 
 // Class to wrap the computation builder to capture information during SPMD
diff --git a/third_party/xla/xla/service/spmd/stateful_rng_spmd_partitioner.h b/third_party/xla/xla/service/spmd/stateful_rng_spmd_partitioner.h
index 5d9170e67e6277..dd3e2635076a8a 100644
--- a/third_party/xla/xla/service/spmd/stateful_rng_spmd_partitioner.h
+++ b/third_party/xla/xla/service/spmd/stateful_rng_spmd_partitioner.h
@@ -45,13 +45,18 @@ class StatefulRngSpmdPartitioningVisitor
 
 class StatefulRngSpmdPartitioner : public spmd::SpmdPartitioner {
  public:
-  StatefulRngSpmdPartitioner(int64_t num_partitions, int64_t num_replicas,
-                             int64_t threshold_for_windowed_einsum_mib = 100000,
-                             bool windowed_einsum_use_multiple_streams = false)
-      : spmd::SpmdPartitioner(
-            num_partitions, num_replicas,
-            GetSpmdPartitionerOptions(threshold_for_windowed_einsum_mib,
-                                      windowed_einsum_use_multiple_streams)) {}
+  StatefulRngSpmdPartitioner(
+      int64_t num_partitions, int64_t num_replicas,
+      int64_t threshold_for_windowed_einsum_mib = 100000,
+      bool windowed_einsum_use_multiple_streams = false,
+      bool skip_checking_windowed_einsum_users = false,
+      bool disable_ag_rewrite_for_multiple_consumers = false)
+      : spmd::SpmdPartitioner(num_partitions, num_replicas,
+                              GetSpmdPartitionerOptions(
+                                  threshold_for_windowed_einsum_mib,
+                                  windowed_einsum_use_multiple_streams,
+                                  skip_checking_windowed_einsum_users,
+                                  disable_ag_rewrite_for_multiple_consumers)) {}
 
  protected:
   std::unique_ptr<spmd::SpmdPartitioningVisitor> CreateVisitor(
@@ -70,12 +75,18 @@ class StatefulRngSpmdPartitioner : public spmd::SpmdPartitioner {
  private:
   static spmd::SpmdPartitionerOptions GetSpmdPartitionerOptions(
       int64_t threshold_for_windowed_einsum_mib,
-      bool windowed_einsum_use_multiple_streams = false) {
+      bool windowed_einsum_use_multiple_streams = false,
+      bool skip_checking_windowed_einsum_users = false,
+      bool disable_ag_rewrite_for_multiple_consumers = false) {
     spmd::SpmdPartitionerOptions options;
     options.allow_module_signature_change = true;
     options.threshold_for_windowed_einsum_mib =
         threshold_for_windowed_einsum_mib;
     options.unroll_windowed_einsum = windowed_einsum_use_multiple_streams;
+    options.skip_checking_windowed_einsum_users =
+        skip_checking_windowed_einsum_users;
+    options.disable_ag_rewrite_for_multiple_consumers =
+        disable_ag_rewrite_for_multiple_consumers;
     return options;
   }
 };
diff --git a/third_party/xla/xla/service/spmd/stateful_rng_spmd_partitioner_test.cc b/third_party/xla/xla/service/spmd/stateful_rng_spmd_partitioner_test.cc
index d301f0ccdcdd74..0ae7efe0daf4f4 100644
--- a/third_party/xla/xla/service/spmd/stateful_rng_spmd_partitioner_test.cc
+++ b/third_party/xla/xla/service/spmd/stateful_rng_spmd_partitioner_test.cc
@@ -34,11 +34,25 @@ namespace xla {
 namespace spmd {
 namespace {
 
+int64_t CountInstructions(const HloComputation &computation, HloOpcode opcode) {
+  int64_t count = 0;
+  for (const auto &instruction : computation.instructions()) {
+    if (instruction->opcode() == opcode) {
+      count++;
+    }
+  }
+  return count;
+}
+
 class StatefulRngSpmdPartitionerTest : public HloTestBase {
  public:
   absl::StatusOr<std::unique_ptr<HloModule>> PartitionComputation(
       absl::string_view hlo_module, int64_t num_partitions,
-      std::function<void(HloPassPipeline &pipeline)> add_passes = nullptr) {
+      std::function<void(HloPassPipeline &pipeline)> add_passes = nullptr,
+      int64_t threshold_for_windowed_einsum_mib = 1000000,
+      bool windowed_einsum_use_multiple_streams = false,
+      bool skip_checking_windowed_einsum_users = false,
+      bool disable_ag_rewrite_for_multiple_consumers = false) {
     TF_ASSIGN_OR_RETURN(
         auto module, ParseAndReturnVerifiedModule(
                          hlo_module, GetModuleConfigForTest(
@@ -51,8 +65,12 @@ class StatefulRngSpmdPartitionerTest : public HloTestBase {
       add_passes(pass);
     }
     pass.AddPass<ShardingPropagation>(/*is_spmd=*/true);
-    pass.AddPass<StatefulRngSpmdPartitioner>(num_partitions,
-                                             /*num_replicas=*/1);
+    pass.AddPass<StatefulRngSpmdPartitioner>(
+        num_partitions,
+        /*num_replicas=*/1, threshold_for_windowed_einsum_mib,
+        windowed_einsum_use_multiple_streams,
+        skip_checking_windowed_einsum_users,
+        disable_ag_rewrite_for_multiple_consumers);
     pass.AddPass<HloVerifier>(/*layout_sensitive=*/false,
                               /*allow_mixed_precision=*/false);
     TF_RETURN_IF_ERROR(pass.Run(module.get()).status());
@@ -116,6 +134,42 @@ ENTRY entry {
   VerifyNoAllReduce(module.get());
 }
 
+TEST_F(StatefulRngSpmdPartitionerTest,
+       EinsumDisableRewriteForAgWithMultipleConsumers) {
+  absl::string_view hlo_string = R"(
+HloModule test, entry_computation_layout={(bf16[2,2048,24576]{2,1,0}, bf16[24576,98304]{1,0}, bf16[24576,98304]{1,0})->bf16[2,2048,98304]{2,1,0}}, num_partitions=4
+
+ENTRY main {
+  Arg_0.1 = bf16[2,2048,24576]{2,1,0} parameter(0), sharding={devices=[1,4,1]<=[4]}
+  Arg_1.2 = bf16[24576,98304]{1,0} parameter(1), sharding={devices=[1,4]<=[4]}
+  dot.5 = bf16[2,2048,98304]{2,1,0} dot(Arg_0.1, Arg_1.2), lhs_contracting_dims={2}, rhs_contracting_dims={0}, sharding={devices=[1,1,4]<=[4]}
+  Arg_2.3 = bf16[24576,98304]{1,0} parameter(2), sharding={devices=[1,4]<=[4]}
+  dot.6 = bf16[2,2048,98304]{2,1,0} dot(Arg_0.1, Arg_2.3), lhs_contracting_dims={2}, rhs_contracting_dims={0}, sharding={devices=[1,1,4]<=[4]}
+  ROOT add.8 = bf16[2,2048,98304]{2,1,0} add(dot.5, dot.6), sharding={devices=[1,1,4]<=[4]}
+}
+
+)";
+  // With disable_ag_rewrite_for_multiple_consumers set to true, we expect only
+  // 1 while loop to exist which is the rewritten windowed einsum loop for the
+  // first ag->dot pattern. The second dot which shares the same operand with
+  // the loop will remain as is.
+  TF_ASSERT_OK_AND_ASSIGN(
+      auto module, PartitionComputation(
+                       hlo_string, /*num_partitions=*/4, /*add_passes=*/nullptr,
+                       /*threshold_for_windowed_einsum_mib=*/0,
+                       /*windowed_einsum_use_multiple_streams=*/true,
+                       /*skip_checking_windowed_einsum_users=*/true,
+                       /*disable_ag_rewrite_for_multiple_consumers=*/true));
+  XLA_VLOG_LINES(1, module->ToString());
+  EXPECT_EQ(CountInstructions(*module->entry_computation(), HloOpcode::kWhile),
+            1);
+  EXPECT_EQ(CountInstructions(*module->entry_computation(), HloOpcode::kDot),
+            1);
+  EXPECT_EQ(
+      CountInstructions(*module->entry_computation(), HloOpcode::kAllGather),
+      1);
+}
+
 TEST_F(StatefulRngSpmdPartitionerTest, VerifyThresholdSetCorrectly) {
   auto debug_options = HloTestBase::GetDebugOptionsForTest();
   int64_t threshold = 400;
diff --git a/third_party/xla/xla/tests/collective_ops_test_e2e.cc b/third_party/xla/xla/tests/collective_ops_test_e2e.cc
index c6885ab13ee5be..d40ce579a41e06 100644
--- a/third_party/xla/xla/tests/collective_ops_test_e2e.cc
+++ b/third_party/xla/xla/tests/collective_ops_test_e2e.cc
@@ -647,5 +647,85 @@ TEST_F(CollectiveOpsTestE2E, NoAllToAllDecomposition) {
   LiteralTestUtil::ExpectR1Equal<uint32_t>({20, 25, 21, 26}, results[1]);
 }
 
+TEST_F(CollectiveOpsTestE2E, WindowedEinsumE2EAllgatherMultiConsumer) {
+  absl::string_view kModuleReplicatedStr = R"(
+HloModule pjit__unnamed_wrapped_function_, entry_computation_layout={(bf16[2,16,48]{2,1,0}, bf16[48,192]{1,0}, bf16[48,192]{1,0}, bf16[192,48]{1,0})->bf16[2,16,48]{2,1,0}}, allow_spmd_sharding_propagation_to_parameters={false,false,false,false}, num_partitions=4
+
+ENTRY main.12 {
+  Arg_0.1 = bf16[2,16,48]{2,1,0} parameter(0), sharding={devices=[1,4,1]<=[4]}
+  Arg_1.2 = bf16[48,192]{1,0} parameter(1), sharding={devices=[1,4]<=[4]}
+  dot.5 = bf16[2,16,192]{2,1,0} dot(Arg_0.1, Arg_1.2), lhs_contracting_dims={2}, rhs_contracting_dims={0}
+  custom-call.7 = bf16[2,16,192]{2,1,0} custom-call(dot.5), custom_call_target="Sharding", sharding={devices=[1,1,4]<=[4]}
+  Arg_2.3 = bf16[48,192]{1,0} parameter(2), sharding={devices=[1,4]<=[4]}
+  dot.6 = bf16[2,16,192]{2,1,0} dot(Arg_0.1, Arg_2.3), lhs_contracting_dims={2}, rhs_contracting_dims={0}
+  add.8 = bf16[2,16,192]{2,1,0} add(custom-call.7, dot.6)
+  Arg_3.4 = bf16[192,48]{1,0} parameter(3), sharding={devices=[4,1]<=[4]}
+  dot.9 = bf16[2,16,48]{2,1,0} dot(add.8, Arg_3.4), lhs_contracting_dims={2}, rhs_contracting_dims={0}
+  tuple.10 = (bf16[2,16,48]{2,1,0}) tuple(dot.9)
+  ROOT get-tuple-element.11 = bf16[2,16,48]{2,1,0} get-tuple-element(tuple.10), index=0, sharding={devices=[1,4,1]<=[4]}
+} // main.12
+)";
+
+  const int64_t kNumReplicas = 1;
+  const int64_t kNumPartitions = 4;
+
+  HloModuleConfig config =
+      GetModuleConfigForTest(/*replica_count=*/kNumReplicas);
+  auto opts = GetDebugOptionsForTest();
+  opts.set_xla_gpu_threshold_for_windowed_einsum_mib(0);
+  opts.set_xla_gpu_multi_streamed_windowed_einsum(true);
+  opts.set_xla_gpu_graph_min_graph_size(200);
+  opts.set_xla_gpu_enable_triton_gemm(false);
+  config.set_debug_options(opts);
+  config.set_num_partitions(kNumPartitions);
+  TF_ASSERT_OK_AND_ASSIGN(
+      auto module, ParseAndReturnVerifiedModule(kModuleReplicatedStr, config));
+  DeviceAssignment assn(/*replica_count=*/kNumReplicas,
+                        /*computation_count=*/kNumPartitions);
+  config.set_replica_count(kNumReplicas);
+  for (int64_t i = 0; i < kNumPartitions; ++i) {
+    assn(0, i) = i;
+  }
+
+  auto fake_arguments = xla::MakeFakeArguments(module.get()).value();
+  std::vector<Literal*> fake_ptrs(fake_arguments.size());
+  for (int i = 0; i < fake_arguments.size(); i++) {
+    fake_ptrs[i] = &fake_arguments[i];
+  }
+
+  TF_ASSERT_OK_AND_ASSIGN(
+      std::vector<Literal> results,
+      HloTestBase::ExecuteReplicated(
+          std::move(module), fake_ptrs, kNumPartitions, &assn,
+          true /*run_hlo_passes*/, true /*use-threads*/));
+  ASSERT_EQ(results.size(), kNumPartitions);
+  HloModuleConfig ref_config =
+      GetModuleConfigForTest(/*replica_count=*/kNumReplicas);
+  auto ref_opts = GetDebugOptionsForTest();
+  ref_opts.set_xla_gpu_graph_min_graph_size(200);
+  ref_opts.set_xla_gpu_enable_triton_gemm(false);
+  ref_config.set_debug_options(ref_opts);
+  ref_config.set_num_partitions(kNumPartitions);
+  TF_ASSERT_OK_AND_ASSIGN(
+      auto ref_module,
+      ParseAndReturnVerifiedModule(kModuleReplicatedStr, ref_config));
+  auto fake_ref_arguments = xla::MakeFakeArguments(ref_module.get()).value();
+  std::vector<Literal*> ref_fake_ptrs(fake_ref_arguments.size());
+  for (int i = 0; i < fake_ref_arguments.size(); i++) {
+    ref_fake_ptrs[i] = &fake_ref_arguments[i];
+  }
+
+  TF_ASSERT_OK_AND_ASSIGN(
+      std::vector<Literal> ref_results,
+      HloTestBase::ExecuteReplicated(
+          std::move(ref_module), ref_fake_ptrs, kNumPartitions, &assn,
+          true /*run_hlo_passes*/, true /*use-threads*/));
+  ASSERT_EQ(ref_results.size(), kNumPartitions);
+  ErrorSpec error_spec{1e-2, 1e-2};
+  // Results should be the same between windowed einsum and non-windowed cases
+  for (int i = 0; i < kNumPartitions; i++) {
+    EXPECT_TRUE(LiteralTestUtil::Near(ref_results[i], results[i], error_spec));
+  }
+}
 }  // namespace
 }  // namespace xla

From 934c0de9e0c19f8b93ce2f290764991150f5bb8a Mon Sep 17 00:00:00 2001
From: Mohammed Anany <manany@google.com>
Date: Tue, 14 May 2024 01:53:21 -0700
Subject: [PATCH 097/478] [XLA:GPU] Refactor inputPrecision condition for
 Triton to use HloAnyOf without requiring adaptors

PiperOrigin-RevId: 633484897
---
 .../xla/xla/service/gpu/ir_emitter_triton.cc    | 17 +++++++----------
 1 file changed, 7 insertions(+), 10 deletions(-)

diff --git a/third_party/xla/xla/service/gpu/ir_emitter_triton.cc b/third_party/xla/xla/service/gpu/ir_emitter_triton.cc
index c190665f71e4e3..20f290e27ff18a 100644
--- a/third_party/xla/xla/service/gpu/ir_emitter_triton.cc
+++ b/third_party/xla/xla/service/gpu/ir_emitter_triton.cc
@@ -1942,18 +1942,15 @@ absl::Status EmitMatMul(mlir::OpBuilder builder,
   const HloInstruction* root = dot_instr->parent()->root_instruction();
   TF_RET_CHECK(!root->shape().IsTuple());
 
-  auto fusion_adaptor = HloFusionAdaptor::ForComputation(computation);
-  HloInstructionAdaptor instr_adaptor{*instr, fusion_adaptor.get()};
   // TODO(b/320659359) Allow TF32 for 8-bit or less types with F32.
-  bool is_8_bit_or_less_dot_with_F32 = HloAnyOf(
-      instr_adaptor.GetOperands(), *fusion_adaptor,
-      [&](HloInstructionAdaptor node) {
-        if (node.opcode() != HloOpcode::kConvert) {
+  bool is_unsupported_bitwidth =
+      HloAnyOf({dot_instr}, [&](const HloInstruction* node) {
+        if (node->opcode() != HloOpcode::kConvert) {
           return false;
         }
-        Type in_type =
-            TritonType(builder, node.GetOperand(0).shape().element_type());
-        Type out_type = TritonType(builder, node.shape().element_type());
+        auto in_type =
+            TritonType(builder, node->operand(0)->shape().element_type());
+        Type out_type = TritonType(builder, node->shape().element_type());
         return in_type.getIntOrFloatBitWidth() <= 8 && out_type.isF32();
       });
 
@@ -2168,7 +2165,7 @@ absl::Status EmitMatMul(mlir::OpBuilder builder,
       // lower precision than the output type. The change was introduced here:
       // https://github.com/openai/triton/commit/31b0c521427109a8eda609b58d756c380b21599a
       auto input_precision =
-          IsTf32Allowed(dot_instr) && !is_8_bit_or_less_dot_with_F32
+          IsTf32Allowed(dot_instr) && !is_unsupported_bitwidth
               ? mt::InputPrecision::TF32
               : mt::InputPrecision::IEEE;
       accumulator_next =

From ca03c7863ceadfe2247e0d0a1f9b1e9e3854f4a7 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Tue, 14 May 2024 02:02:33 -0700
Subject: [PATCH 098/478] compat: Update forward compatibility horizon to
 2024-05-14

PiperOrigin-RevId: 633487032
---
 tensorflow/python/compat/compat.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorflow/python/compat/compat.py b/tensorflow/python/compat/compat.py
index 2a9869afa5dfb4..f372d9e16c9363 100644
--- a/tensorflow/python/compat/compat.py
+++ b/tensorflow/python/compat/compat.py
@@ -29,7 +29,7 @@
 # This value changes every day with an automatic CL. It can be modified in code
 # via `forward_compatibility_horizon()` or with the environment variable
 # TF_FORWARD_COMPATIBILITY_DELTA_DAYS, which is added to the compatibility date.
-_FORWARD_COMPATIBILITY_HORIZON = datetime.date(2024, 5, 13)
+_FORWARD_COMPATIBILITY_HORIZON = datetime.date(2024, 5, 14)
 _FORWARD_COMPATIBILITY_DELTA_DAYS_VAR_NAME = "TF_FORWARD_COMPATIBILITY_DELTA_DAYS"
 _FORWARD_COMPATIBILITY_DATE_NUMBER = None
 

From 1581da5b8d475216d3f1938e8f001db2f03c508b Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Tue, 14 May 2024 02:02:34 -0700
Subject: [PATCH 099/478] Update GraphDef version to 1862.

PiperOrigin-RevId: 633487039
---
 tensorflow/core/public/version.h | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorflow/core/public/version.h b/tensorflow/core/public/version.h
index aafe138460a6db..e6d5b27b16cb2c 100644
--- a/tensorflow/core/public/version.h
+++ b/tensorflow/core/public/version.h
@@ -108,7 +108,7 @@ limitations under the License.
 
 #define TF_GRAPH_DEF_VERSION_MIN_PRODUCER 0
 #define TF_GRAPH_DEF_VERSION_MIN_CONSUMER 0
-#define TF_GRAPH_DEF_VERSION 1861  // Updated: 2024/5/13
+#define TF_GRAPH_DEF_VERSION 1862  // Updated: 2024/5/14
 
 // Checkpoint compatibility versions (the versions field in SavedSliceMeta).
 //

From 360088f54467b4fd7491f7f9d9f32da59ca6c958 Mon Sep 17 00:00:00 2001
From: Sergey Kozub <sergeykozub@google.com>
Date: Tue, 14 May 2024 02:52:44 -0700
Subject: [PATCH 100/478] Fix tile size for matmuls with 8-bit operands in
 autotuner

The result of `GetExhaustiveTritonConfigs` is cached in the autotuner class instance, so should not depend on the input dot instruction. Therefore, removing the `has_8_bit_operand` parameter from the method signature.
Instead, adjust the `block_k` attribute in the config normalization phase.

Additionally, treat 8-bit floating point types in the same manner as integer types.

PiperOrigin-RevId: 633506809
---
 .../xla/service/gpu/gemm_fusion_autotuner.cc  | 38 +++++++++----------
 .../xla/service/gpu/gemm_fusion_autotuner.h   |  3 +-
 2 files changed, 20 insertions(+), 21 deletions(-)

diff --git a/third_party/xla/xla/service/gpu/gemm_fusion_autotuner.cc b/third_party/xla/xla/service/gpu/gemm_fusion_autotuner.cc
index d3d09a21c048f1..d5b34fb877add3 100644
--- a/third_party/xla/xla/service/gpu/gemm_fusion_autotuner.cc
+++ b/third_party/xla/xla/service/gpu/gemm_fusion_autotuner.cc
@@ -51,6 +51,7 @@ limitations under the License.
 #include "xla/hlo/ir/hlo_module.h"
 #include "xla/hlo/ir/hlo_opcode.h"
 #include "xla/hlo/utils/hlo_query.h"
+#include "xla/primitive_util.h"
 #include "xla/service/algorithm_util.h"
 #include "xla/service/dump.h"
 #include "xla/service/executable.h"
@@ -115,8 +116,6 @@ namespace {
 
 // Minimum tile size.
 constexpr int kMinTileSize = 16;
-constexpr int kDimKMinTileSize = 32;
-constexpr std::array<int, 6> kDimKBlockSizes = {32, 64, 128, 256, 512};
 
 // Default tiling when autotuning is disabled.
 constexpr TritonGemmConfig kDefaultGemmTiling = {32, 32, 32, 1, 1, 4};
@@ -277,8 +276,7 @@ struct TileSizeLimit {
   int block_k = 0;
 };
 
-absl::StatusOr<TileSizeLimit> GetLimits(const HloDotInstruction& dot,
-                                        bool has_8_bit_operand) {
+absl::StatusOr<TileSizeLimit> GetLimits(const HloDotInstruction& dot) {
   TF_ASSIGN_OR_RETURN(int64_t non_contracting_index_lhs,
                       NonContractingDimensionIndex(dot, /*operand_number=*/0));
   TF_ASSIGN_OR_RETURN(int64_t non_contracting_index_rhs,
@@ -306,8 +304,7 @@ absl::StatusOr<TileSizeLimit> GetLimits(const HloDotInstruction& dot,
   return TileSizeLimit{
       /*block_m=*/std::max(max_m, kMinTileSize),
       /*block_n=*/std::max(max_n, kMinTileSize),
-      /*block_k=*/
-      std::max(max_k, has_8_bit_operand ? kDimKMinTileSize : kMinTileSize),
+      /*block_k=*/std::max(max_k, kMinTileSize),
   };
 }
 
@@ -587,18 +584,18 @@ GemmFusionAutotunerImpl::GenerateTritonConfigs(const HloDotInstruction& dot) {
       return false;
     }
     auto in_type = node->operand(0)->shape().element_type();
-    return (in_type == PrimitiveType::S8) || (in_type == PrimitiveType::U8);
+    return primitive_util::BitWidth(in_type) == 8;
   });
 
   std::vector<TritonGemmConfig> result_configs;
-  TF_ASSIGN_OR_RETURN(TileSizeLimit limits, GetLimits(dot, has_8_bit_operand));
+  TF_ASSIGN_OR_RETURN(TileSizeLimit limits, GetLimits(dot));
 
   // Generate the list of configurations (once).
   if (triton_configs_.empty()) {
     triton_configs_ = !IsAutotuningEnabled()
                           ? std::vector(1, kDefaultGemmTiling)
                       : debug_options_.xla_gpu_exhaustive_tiling_search()
-                          ? GetExhaustiveTritonConfigs(has_8_bit_operand)
+                          ? GetExhaustiveTritonConfigs()
                           : GetDefaultTritonConfigs();
   }
 
@@ -639,6 +636,15 @@ GemmFusionAutotunerImpl::GenerateTritonConfigs(const HloDotInstruction& dot) {
     }
     config.split_k = std::min(config.split_k, max_split_k);
 
+    // TODO(b/337839570): block_k = 16 is bugged in Triton for dots with 8-bit
+    // input. Setting minimum to 32 instead of 16 for these cases.
+    // TODO(b/337838200): Write the restriction on the minimum tile size to be
+    // generic. Currently we only handle the 8-bit case as this was the bug we
+    // ran into.
+    if (has_8_bit_operand && config.block_k == kMinTileSize) {
+      config.block_k *= 2;
+    }
+
     // Hopper `wgmma` instruction requires at least 4 warps and 64 elements
     // for LHS tile height.
     if (is_hopper) {
@@ -649,7 +655,8 @@ GemmFusionAutotunerImpl::GenerateTritonConfigs(const HloDotInstruction& dot) {
     // Sparse meta should have at least one element per thread.
     // Note: only 2:4 structured sparsity is currently supported.
     if (dot.sparse_operands()) {
-      config.block_k = std::max(config.block_k, 4 * kMinTileSize);
+      config.block_k =
+          std::max(config.block_k, kMinTileSize * (has_8_bit_operand ? 4 : 2));
       int meta_elements = config.block_m * config.block_k / 16;
       config.num_warps =
           std::min<int>(config.num_warps, meta_elements / WarpSize());
@@ -909,19 +916,12 @@ absl::StatusOr<std::vector<AutotuneResult>> GemmFusionAutotunerImpl::Profile(
 }
 
 std::vector<TritonGemmConfig>
-GemmFusionAutotunerImpl::GetExhaustiveTritonConfigs(
-    bool has_8_bit_operand) const {
+GemmFusionAutotunerImpl::GetExhaustiveTritonConfigs() const {
   std::vector<TritonGemmConfig> configs;
   se::CudaComputeCapability cc = GetComputeCapability();
   bool tune_ctas =
       debug_options_.xla_gpu_enable_triton_hopper() && cc.IsAtLeastHopper();
 
-  // TODO(b/337839570): block_k = 16 is bugged in Triton for dots with 8-bit
-  // input. Setting minimum to 32 instead of 16 for these cases.
-  // TODO(b/337838200): Write the restriction on the minimum tile size to be
-  // generic. Currently we only handle the 8-bit case as this was the bug we
-  // ran into.
-  auto kDimBlockSizes = has_8_bit_operand ? kDimKBlockSizes : kBlockSizes;
   for (int num_stages : kNumStages) {
     // Volta doesn't support num_stages > 2.
     if (!cc.IsAtLeastAmpere() && num_stages > 2) {
@@ -929,7 +929,7 @@ GemmFusionAutotunerImpl::GetExhaustiveTritonConfigs(
     }
     for (int tile_m : kBlockSizes) {
       for (int tile_n : kBlockSizes) {
-        for (int tile_k : kDimBlockSizes) {
+        for (int tile_k : kBlockSizes) {
           const int tile_lhs = tile_m * tile_k;
           const int tile_rhs = tile_k * tile_n;
           for (int num_warps : kNumWarps) {
diff --git a/third_party/xla/xla/service/gpu/gemm_fusion_autotuner.h b/third_party/xla/xla/service/gpu/gemm_fusion_autotuner.h
index 251826d79e7756..3f2a7a7d19a8e4 100644
--- a/third_party/xla/xla/service/gpu/gemm_fusion_autotuner.h
+++ b/third_party/xla/xla/service/gpu/gemm_fusion_autotuner.h
@@ -123,8 +123,7 @@ class GemmFusionAutotunerImpl {
   }
 
   std::vector<TritonGemmConfig> GetDefaultTritonConfigs() const;
-  std::vector<TritonGemmConfig> GetExhaustiveTritonConfigs(
-      bool has_8_bit_operand) const;
+  std::vector<TritonGemmConfig> GetExhaustiveTritonConfigs() const;
 
   const AutotuneConfig config_;
   const DebugOptions debug_options_;

From f55e2a43033fd921173bf67486ef8645574e6b35 Mon Sep 17 00:00:00 2001
From: Sergey Kozub <sergeykozub@google.com>
Date: Tue, 14 May 2024 03:07:12 -0700
Subject: [PATCH 101/478] Fix sparse dot metadata loader

Metadata loader was using incorrect warp assignment, which resulted in incorrect addresses with num_warps>4. This was previously missed, as the autotuner rarely selected such configs.

PiperOrigin-RevId: 633513110
---
 .../sparse_dot_fixes_y24w19.patch              | 18 ++++++++++++++++++
 .../sparse_dot_fixes_y24w19.patch              | 18 ++++++++++++++++++
 2 files changed, 36 insertions(+)

diff --git a/third_party/triton/xla_extensions/sparse_dot_fixes_y24w19.patch b/third_party/triton/xla_extensions/sparse_dot_fixes_y24w19.patch
index 8ac91d153690fd..775ed317d1f9b9 100644
--- a/third_party/triton/xla_extensions/sparse_dot_fixes_y24w19.patch
+++ b/third_party/triton/xla_extensions/sparse_dot_fixes_y24w19.patch
@@ -11,3 +11,21 @@ diff --git a/lib/Dialect/TritonGPU/Transforms/AccelerateMatmul.cpp b/lib/Dialect
  
      // get MMA encoding for the given number of warps
      auto retShapePerCTA = ttg::getShapePerCTA(oldRetType);
+diff --git a/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/ConvertLayoutOpToLLVM/SharedToSparseDotOperand.cpp b/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/ConvertLayoutOpToLLVM/SharedToSparseDotOperand.cpp
+--- a/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/ConvertLayoutOpToLLVM/SharedToSparseDotOperand.cpp
++++ b/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/ConvertLayoutOpToLLVM/SharedToSparseDotOperand.cpp
+@@ -31,7 +31,13 @@ Value convertLayout(
+   // Calculate offset in the tile for the current thread.
+   Value threadsPerWarp = i32_val(kThreadsPerWarp);
+   Value warpId = udiv(thread, threadsPerWarp);
+-  Value warpGroupId = udiv(warpId, i32_val(warpsPerCTA[1]));
++  Value warpGroupId;
++  if (mmaLayout.isHopper()) {
++    warpGroupId = urem(warpId, i32_val(warpsPerCTA[0]));
++  } else {
++    assert(mmaLayout.isAmpere());
++    warpGroupId = udiv(warpId, i32_val(warpsPerCTA[1]));
++  }
+   Value laneId = urem(thread, threadsPerWarp);
+   Value laneGroupId = udiv(laneId, i32_val(kThreadsInGroup));
+   Value columnId = urem(laneId, i32_val(shapePerCTATile[1]));
diff --git a/third_party/xla/third_party/triton/xla_extensions/sparse_dot_fixes_y24w19.patch b/third_party/xla/third_party/triton/xla_extensions/sparse_dot_fixes_y24w19.patch
index 8ac91d153690fd..775ed317d1f9b9 100644
--- a/third_party/xla/third_party/triton/xla_extensions/sparse_dot_fixes_y24w19.patch
+++ b/third_party/xla/third_party/triton/xla_extensions/sparse_dot_fixes_y24w19.patch
@@ -11,3 +11,21 @@ diff --git a/lib/Dialect/TritonGPU/Transforms/AccelerateMatmul.cpp b/lib/Dialect
  
      // get MMA encoding for the given number of warps
      auto retShapePerCTA = ttg::getShapePerCTA(oldRetType);
+diff --git a/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/ConvertLayoutOpToLLVM/SharedToSparseDotOperand.cpp b/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/ConvertLayoutOpToLLVM/SharedToSparseDotOperand.cpp
+--- a/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/ConvertLayoutOpToLLVM/SharedToSparseDotOperand.cpp
++++ b/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/ConvertLayoutOpToLLVM/SharedToSparseDotOperand.cpp
+@@ -31,7 +31,13 @@ Value convertLayout(
+   // Calculate offset in the tile for the current thread.
+   Value threadsPerWarp = i32_val(kThreadsPerWarp);
+   Value warpId = udiv(thread, threadsPerWarp);
+-  Value warpGroupId = udiv(warpId, i32_val(warpsPerCTA[1]));
++  Value warpGroupId;
++  if (mmaLayout.isHopper()) {
++    warpGroupId = urem(warpId, i32_val(warpsPerCTA[0]));
++  } else {
++    assert(mmaLayout.isAmpere());
++    warpGroupId = udiv(warpId, i32_val(warpsPerCTA[1]));
++  }
+   Value laneId = urem(thread, threadsPerWarp);
+   Value laneGroupId = udiv(laneId, i32_val(kThreadsInGroup));
+   Value columnId = urem(laneId, i32_val(shapePerCTATile[1]));

From 45c2071da17491a46469e894561ea823979b2029 Mon Sep 17 00:00:00 2001
From: Mohammed Anany <manany@google.com>
Date: Tue, 14 May 2024 03:08:23 -0700
Subject: [PATCH 102/478] [Triton] Limit visibility of original bit-width
 computation during computing mma and dot_operand layouts to not propagate
 through predicates. Predicates are not supported for these layouts during
 lowering from TritonGPU to LLVM and cause crashes.

PiperOrigin-RevId: 633513652
---
 .../triton/temporary/mma_limit_pred.patch     | 23 +++++++++++++++++++
 third_party/triton/temporary/series.bzl       |  1 +
 .../triton/temporary/mma_limit_pred.patch     | 23 +++++++++++++++++++
 .../third_party/triton/temporary/series.bzl   |  1 +
 .../xla/service/gpu/ir_emitter_triton_test.cc | 23 +++++++++++++++++++
 5 files changed, 71 insertions(+)
 create mode 100644 third_party/triton/temporary/mma_limit_pred.patch
 create mode 100644 third_party/xla/third_party/triton/temporary/mma_limit_pred.patch

diff --git a/third_party/triton/temporary/mma_limit_pred.patch b/third_party/triton/temporary/mma_limit_pred.patch
new file mode 100644
index 00000000000000..cde501d0a9b6fd
--- /dev/null
+++ b/third_party/triton/temporary/mma_limit_pred.patch
@@ -0,0 +1,23 @@
+// This patch fixes issues that broke internal benchmarks as well as individual
+// fusions reported. It should be moved to our internal patches in the next
+// integration.
+
+diff --git a/lib/Dialect/TritonGPU/Transforms/AccelerateMatmul.cpp b/lib/Dialect/TritonGPU/Transforms/AccelerateMatmul.cpp
+--- a/lib/Dialect/TritonGPU/Transforms/AccelerateMatmul.cpp
++++ b/lib/Dialect/TritonGPU/Transforms/AccelerateMatmul.cpp
+@@ -133,6 +133,15 @@ class BlockedToMMA : public mlir::Rewrit
+   mutable llvm::DenseMap<Operation *, unsigned> dotOpInstNs;
+ 
+   static bool bwdFilter(Operation *op) {
++    // Dot operand layout assignment to Predicates are not currently supported
++    // during lowering from TritonGPU to LLVM in Triton for MMA cases. This
++    // condition limits visibility of the original bit-width so that predicate
++    // are not considered, hence, kwidth can never be = 32.
++    if (isa<arith::UIToFPOp>(op)) {
++      Type srcType = getElementTypeOrSelf(op->getOperand(0));
++      if (srcType.isInteger(1))
++        return false;
++    }
+     return op->getNumOperands() == 1 &&
+            (isa<tt::FpToFpOp, tt::BitcastOp, ttg::ConvertLayoutOp>(op) ||
+             isPureUnaryInlineAsm(op) ||
diff --git a/third_party/triton/temporary/series.bzl b/third_party/triton/temporary/series.bzl
index b6002f83e1d7bb..70313c9b436d2a 100644
--- a/third_party/triton/temporary/series.bzl
+++ b/third_party/triton/temporary/series.bzl
@@ -8,4 +8,5 @@ internal patch during the next triton integration process.
 temporary_patch_list = [
     "//third_party/triton/temporary:pipelining.patch",
     "//third_party/triton/temporary:support_ceil_op.patch",
+    "//third_party/triton/temporary:mma_limit_pred.patch",
 ]
diff --git a/third_party/xla/third_party/triton/temporary/mma_limit_pred.patch b/third_party/xla/third_party/triton/temporary/mma_limit_pred.patch
new file mode 100644
index 00000000000000..cde501d0a9b6fd
--- /dev/null
+++ b/third_party/xla/third_party/triton/temporary/mma_limit_pred.patch
@@ -0,0 +1,23 @@
+// This patch fixes issues that broke internal benchmarks as well as individual
+// fusions reported. It should be moved to our internal patches in the next
+// integration.
+
+diff --git a/lib/Dialect/TritonGPU/Transforms/AccelerateMatmul.cpp b/lib/Dialect/TritonGPU/Transforms/AccelerateMatmul.cpp
+--- a/lib/Dialect/TritonGPU/Transforms/AccelerateMatmul.cpp
++++ b/lib/Dialect/TritonGPU/Transforms/AccelerateMatmul.cpp
+@@ -133,6 +133,15 @@ class BlockedToMMA : public mlir::Rewrit
+   mutable llvm::DenseMap<Operation *, unsigned> dotOpInstNs;
+ 
+   static bool bwdFilter(Operation *op) {
++    // Dot operand layout assignment to Predicates are not currently supported
++    // during lowering from TritonGPU to LLVM in Triton for MMA cases. This
++    // condition limits visibility of the original bit-width so that predicate
++    // are not considered, hence, kwidth can never be = 32.
++    if (isa<arith::UIToFPOp>(op)) {
++      Type srcType = getElementTypeOrSelf(op->getOperand(0));
++      if (srcType.isInteger(1))
++        return false;
++    }
+     return op->getNumOperands() == 1 &&
+            (isa<tt::FpToFpOp, tt::BitcastOp, ttg::ConvertLayoutOp>(op) ||
+             isPureUnaryInlineAsm(op) ||
diff --git a/third_party/xla/third_party/triton/temporary/series.bzl b/third_party/xla/third_party/triton/temporary/series.bzl
index b6002f83e1d7bb..70313c9b436d2a 100644
--- a/third_party/xla/third_party/triton/temporary/series.bzl
+++ b/third_party/xla/third_party/triton/temporary/series.bzl
@@ -8,4 +8,5 @@ internal patch during the next triton integration process.
 temporary_patch_list = [
     "//third_party/triton/temporary:pipelining.patch",
     "//third_party/triton/temporary:support_ceil_op.patch",
+    "//third_party/triton/temporary:mma_limit_pred.patch",
 ]
diff --git a/third_party/xla/xla/service/gpu/ir_emitter_triton_test.cc b/third_party/xla/xla/service/gpu/ir_emitter_triton_test.cc
index 8ebf9e94386885..111f8e276e1de8 100644
--- a/third_party/xla/xla/service/gpu/ir_emitter_triton_test.cc
+++ b/third_party/xla/xla/service/gpu/ir_emitter_triton_test.cc
@@ -1375,6 +1375,29 @@ ENTRY e {
   EXPECT_EQ(paths.size(), 1);
 }
 
+TEST_F(TritonGemmTest, DotWithPredFromCompareProducesCorrectResult) {
+  const std::string hlo_text = R"(
+triton_dot {
+  parameter_0 = s32[4,128]{1,0} parameter(0)
+  broadcast.255 = s32[4,128,64]{2,1,0} broadcast(parameter_0), dimensions={0,1}
+  parameter_1 = s32[4,128,64]{2,1,0} parameter(1)
+  compare.39 = pred[4,128,64]{2,1,0} compare(broadcast.255, parameter_1), direction=EQ
+  bitcast.1097 = pred[512,64]{1,0} reshape(compare.39)
+  convert.229 = bf16[512,64]{1,0} convert(bitcast.1097)
+  parameter_2 = bf16[64,256]{0,1} parameter(2)
+  ROOT dot.21 = bf16[512,256]{1,0} dot(convert.229, parameter_2),
+    lhs_contracting_dims={1}, rhs_contracting_dims={0}
+}
+ENTRY main {
+  p0 = s32[4,128]{1,0} parameter(0)
+  p1 = s32[4,128,64]{2,1,0} parameter(1)
+  p2 = bf16[64,256]{0,1} parameter(2)
+  ROOT gemm_fusion_dot.0 = bf16[512,256]{1,0} fusion(p0, p1, p2), kind=kCustom, calls=triton_dot, backend_config={"fusion_backend_config":{"kind":"__triton_gemm","triton_gemm_config":{"block_m":"64","block_n":"128","block_k":"32","split_k":"1","num_stages":"4","num_warps":"4","num_ctas":"1"}}}
+})";
+
+  EXPECT_TRUE(RunAndCompare(hlo_text, ErrorSpec{/*aabs=*/1e-3, /*arel=*/1e-3}));
+}
+
 TEST_F(TritonGemmTest, UseTensorCoresForF32OnAmpere) {
   const std::string kHloText = R"(
 triton_gemm_r {

From 3734929103b83e6b5a93a77a58d361b387d1feaf Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Tue, 14 May 2024 03:17:43 -0700
Subject: [PATCH 103/478] Update ops-related pbtxt files.

PiperOrigin-RevId: 633517388
---
 tensorflow/core/ops/ops.pbtxt | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorflow/core/ops/ops.pbtxt b/tensorflow/core/ops/ops.pbtxt
index 5019fe2a5b2fad..0556051577f5e4 100644
--- a/tensorflow/core/ops/ops.pbtxt
+++ b/tensorflow/core/ops/ops.pbtxt
@@ -1,4 +1,4 @@
-go/debugproto    
+go/debugstr  
 op {
   name: "Abort"
   attr {

From d6bfe26897528425030ff0cd9b394159d0100504 Mon Sep 17 00:00:00 2001
From: Leo Heinsaar <heinsaar@google.com>
Date: Tue, 14 May 2024 03:30:39 -0700
Subject: [PATCH 104/478] Move `BuildAttributesMap()` to a common place

So that both CPU and GPU code can reuse the same implementation. At the moment each uses its own version that are identical.

PiperOrigin-RevId: 633522046
---
 third_party/xla/xla/ffi/BUILD                 |  17 +++
 third_party/xla/xla/ffi/attribute_map.cc      | 130 ++++++++++++++++++
 third_party/xla/xla/ffi/attribute_map.h       |  32 +++++
 third_party/xla/xla/service/cpu/BUILD         |   1 +
 .../service/cpu/runtime_handle_ffi_call.cc    |  92 +------------
 third_party/xla/xla/service/gpu/BUILD         |   1 +
 third_party/xla/xla/service/gpu/fusions/BUILD |   1 +
 .../xla/xla/service/gpu/fusions/custom.cc     |   3 +-
 .../xla/service/gpu/ir_emitter_unnested.cc    |   3 +-
 .../service/gpu/runtime/custom_call_thunk.cc  |  94 -------------
 .../service/gpu/runtime/custom_call_thunk.h   |   5 -
 third_party/xla/xla/tests/BUILD               |   1 +
 12 files changed, 189 insertions(+), 191 deletions(-)
 create mode 100644 third_party/xla/xla/ffi/attribute_map.cc
 create mode 100644 third_party/xla/xla/ffi/attribute_map.h

diff --git a/third_party/xla/xla/ffi/BUILD b/third_party/xla/xla/ffi/BUILD
index c53835bf83a88d..4da1564a334ad1 100644
--- a/third_party/xla/xla/ffi/BUILD
+++ b/third_party/xla/xla/ffi/BUILD
@@ -102,6 +102,23 @@ cc_library(
     ],
 )
 
+cc_library(
+    name = "attribute_map",
+    srcs = ["attribute_map.cc"],
+    hdrs = ["attribute_map.h"],
+    deps = [
+        ":call_frame",
+        "//xla:xla_data_proto_cc_impl",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/status:statusor",
+        "@com_google_absl//absl/strings",
+        "@llvm-project//llvm:Support",
+        "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:Support",
+        "@local_tsl//tsl/platform:errors",
+    ],
+)
+
 xla_cc_test(
     name = "ffi_test",
     srcs = ["ffi_test.cc"],
diff --git a/third_party/xla/xla/ffi/attribute_map.cc b/third_party/xla/xla/ffi/attribute_map.cc
new file mode 100644
index 00000000000000..9443a8c010d542
--- /dev/null
+++ b/third_party/xla/xla/ffi/attribute_map.cc
@@ -0,0 +1,130 @@
+/* Copyright 2024 The OpenXLA Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+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.
+==============================================================================*/
+
+#include "xla/ffi/attribute_map.h"
+
+#include <cstdint>
+#include <string_view>
+
+#include "absl/status/status.h"
+#include "absl/status/statusor.h"
+#include "absl/strings/str_cat.h"
+#include "llvm/ADT/TypeSwitch.h"
+#include "mlir/IR/Attributes.h"  // from @llvm-project
+#include "mlir/IR/BuiltinAttributes.h"  // from @llvm-project
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+#include "xla/ffi/call_frame.h"
+#include "xla/xla_data.pb.h"
+#include "tsl/platform/errors.h"
+
+using FlatAttribute = xla::ffi::CallFrameBuilder::FlatAttribute;
+using FlatAttributesMap = xla::ffi::CallFrameBuilder::FlatAttributesMap;
+
+namespace xla::ffi {
+
+absl::StatusOr<FlatAttributesMap> BuildAttributesMap(
+    mlir::DictionaryAttr dict) {
+  FlatAttributesMap attributes;
+  for (auto& kv : dict) {
+    std::string_view name = kv.getName().strref();
+
+    auto boolean = [&](mlir::BoolAttr boolean) {
+      attributes[name] = static_cast<bool>(boolean.getValue());
+      return absl::OkStatus();
+    };
+
+    auto integer = [&](mlir::IntegerAttr integer) {
+      switch (integer.getType().getIntOrFloatBitWidth()) {
+        case 1:
+          attributes[name] = static_cast<bool>(integer.getInt());
+          return absl::OkStatus();
+        case 8:
+          attributes[name] = static_cast<int8_t>(integer.getInt());
+          return absl::OkStatus();
+        case 16:
+          attributes[name] = static_cast<int16_t>(integer.getInt());
+          return absl::OkStatus();
+        case 32:
+          attributes[name] = static_cast<int32_t>(integer.getInt());
+          return absl::OkStatus();
+        case 64:
+          attributes[name] = static_cast<int64_t>(integer.getInt());
+          return absl::OkStatus();
+        default:
+          return absl::InvalidArgumentError(absl::StrCat(
+              "Unsupported integer attribute bit width for attribute: ", name));
+      }
+    };
+
+    auto fp = [&](mlir::FloatAttr fp) {
+      switch (fp.getType().getIntOrFloatBitWidth()) {
+        case 32:
+          attributes[name] = static_cast<float>(fp.getValue().convertToFloat());
+          return absl::OkStatus();
+        case 64:
+          attributes[name] =
+              static_cast<double>(fp.getValue().convertToDouble());
+          return absl::OkStatus();
+        default:
+          return absl::InvalidArgumentError(absl::StrCat(
+              "Unsupported float attribute bit width for attribute: ", name));
+      }
+    };
+
+    auto arr = [&](mlir::DenseArrayAttr arr) {
+      if (auto dense = mlir::dyn_cast<mlir::DenseI8ArrayAttr>(arr)) {
+        attributes[name] = dense.asArrayRef().vec();
+        return absl::OkStatus();
+      } else if (auto dense = mlir::dyn_cast<mlir::DenseI16ArrayAttr>(arr)) {
+        attributes[name] = dense.asArrayRef().vec();
+        return absl::OkStatus();
+      } else if (auto dense = mlir::dyn_cast<mlir::DenseI32ArrayAttr>(arr)) {
+        attributes[name] = dense.asArrayRef().vec();
+        return absl::OkStatus();
+      } else if (auto dense = mlir::dyn_cast<mlir::DenseI64ArrayAttr>(arr)) {
+        attributes[name] = dense.asArrayRef().vec();
+        return absl::OkStatus();
+      } else if (auto dense = mlir::dyn_cast<mlir::DenseF32ArrayAttr>(arr)) {
+        attributes[name] = dense.asArrayRef().vec();
+        return absl::OkStatus();
+      } else if (auto dense = mlir::dyn_cast<mlir::DenseF64ArrayAttr>(arr)) {
+        attributes[name] = dense.asArrayRef().vec();
+        return absl::OkStatus();
+      } else {
+        return absl::InvalidArgumentError(absl::StrCat(
+            "Unsupported array element type for attribute: ", name));
+      }
+    };
+
+    auto str = [&](mlir::StringAttr str) {
+      attributes[name] = str.getValue().str();
+      return absl::OkStatus();
+    };
+
+    TF_RETURN_IF_ERROR(
+        llvm::TypeSwitch<mlir::Attribute, absl::Status>(kv.getValue())
+            .Case<mlir::BoolAttr>(boolean)
+            .Case<mlir::IntegerAttr>(integer)
+            .Case<mlir::FloatAttr>(fp)
+            .Case<mlir::DenseArrayAttr>(arr)
+            .Case<mlir::StringAttr>(str)
+            .Default([&](mlir::Attribute) {
+              return absl::InvalidArgumentError(absl::StrCat(
+                  "Unsupported attribute type for attribute: ", name));
+            }));
+  }
+  return attributes;
+}
+}  // namespace xla::ffi
diff --git a/third_party/xla/xla/ffi/attribute_map.h b/third_party/xla/xla/ffi/attribute_map.h
new file mode 100644
index 00000000000000..d6c37b31c5522b
--- /dev/null
+++ b/third_party/xla/xla/ffi/attribute_map.h
@@ -0,0 +1,32 @@
+/* Copyright 2024 The OpenXLA Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+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.
+==============================================================================*/
+
+#ifndef XLA_FFI_ATTRIBUTE_MAP_H_
+#define XLA_FFI_ATTRIBUTE_MAP_H_
+
+#include "absl/status/statusor.h"
+#include "mlir/IR/BuiltinAttributes.h"  // from @llvm-project
+#include "xla/ffi/call_frame.h"
+
+namespace xla::ffi {
+
+// Converts MLIR dictionary attribute attached to a custom call operation to a
+// custom call handler attributes that are forwarded to the FFI handler.
+absl::StatusOr<CallFrameBuilder::FlatAttributesMap> BuildAttributesMap(
+    mlir::DictionaryAttr dict);
+
+}  // namespace xla::ffi
+
+#endif  // XLA_FFI_ATTRIBUTE_MAP_H_
diff --git a/third_party/xla/xla/service/cpu/BUILD b/third_party/xla/xla/service/cpu/BUILD
index 125af6dbce9170..1d0defbdcced10 100644
--- a/third_party/xla/xla/service/cpu/BUILD
+++ b/third_party/xla/xla/service/cpu/BUILD
@@ -1184,6 +1184,7 @@ cc_library(
     deps = [
         "//xla:shape_util",
         "//xla:xla_data_proto_cc",
+        "//xla/ffi:attribute_map",
         "//xla/ffi:call_frame",
         "//xla/ffi:ffi_api",
         "//xla/service:custom_call_status_public_headers",
diff --git a/third_party/xla/xla/service/cpu/runtime_handle_ffi_call.cc b/third_party/xla/xla/service/cpu/runtime_handle_ffi_call.cc
index 7a7963593f592a..7d3271db600634 100644
--- a/third_party/xla/xla/service/cpu/runtime_handle_ffi_call.cc
+++ b/third_party/xla/xla/service/cpu/runtime_handle_ffi_call.cc
@@ -33,6 +33,7 @@ limitations under the License.
 #include "mlir/IR/Builders.h"  // from @llvm-project
 #include "mlir/IR/BuiltinAttributes.h"  // from @llvm-project
 #include "mlir/IR/MLIRContext.h"  // from @llvm-project
+#include "xla/ffi/attribute_map.h"
 #include "xla/ffi/call_frame.h"
 #include "xla/ffi/ffi_api.h"
 #include "xla/primitive_util.h"
@@ -46,95 +47,6 @@ namespace ffi = xla::ffi;
 
 namespace {
 
-using Attribute = ffi::CallFrameBuilder::FlatAttribute;
-using AttributesMap = ffi::CallFrameBuilder::FlatAttributesMap;
-
-// TODO(heinsaar): This BuildAttributesMap() is originally an identical
-//                 copy-paste of the same function in custom_call_thunk.cc
-//                 May make sense to have one in a common place & reuse.
-absl::StatusOr<AttributesMap> BuildAttributesMap(mlir::DictionaryAttr dict) {
-  AttributesMap attributes;
-  for (auto& kv : dict) {
-    std::string_view name = kv.getName().strref();
-
-    auto boolean = [&](mlir::BoolAttr boolean) {
-      attributes[name] = static_cast<bool>(boolean.getValue());
-      return absl::OkStatus();
-    };
-
-    auto integer = [&](mlir::IntegerAttr integer) {
-      const bool is_unsigned = integer.getType().isUnsignedInteger();
-      if (is_unsigned) {
-        switch (integer.getType().getIntOrFloatBitWidth()) {
-          case 8:
-            attributes[name] = static_cast<uint8_t>(integer.getUInt());
-            return absl::OkStatus();
-          case 16:
-            attributes[name] = static_cast<uint16_t>(integer.getUInt());
-            return absl::OkStatus();
-          case 32:
-            attributes[name] = static_cast<uint32_t>(integer.getUInt());
-            return absl::OkStatus();
-          case 64:
-            attributes[name] = static_cast<uint64_t>(integer.getUInt());
-            return absl::OkStatus();
-          default:
-            return absl::InvalidArgumentError(absl::StrCat(
-                "Unsupported integer attribute bit width for attribute: ",
-                name));
-        }
-      } else {
-        switch (integer.getType().getIntOrFloatBitWidth()) {
-          case 8:
-            attributes[name] = static_cast<int8_t>(integer.getInt());
-            return absl::OkStatus();
-          case 16:
-            attributes[name] = static_cast<int16_t>(integer.getInt());
-            return absl::OkStatus();
-          case 32:
-            attributes[name] = static_cast<int32_t>(integer.getInt());
-            return absl::OkStatus();
-          case 64:
-            attributes[name] = static_cast<int64_t>(integer.getInt());
-            return absl::OkStatus();
-          default:
-            return absl::InvalidArgumentError(absl::StrCat(
-                "Unsupported integer attribute bit width for attribute: ",
-                name));
-        }
-      }
-    };
-
-    auto fp = [&](mlir::FloatAttr fp) {
-      switch (fp.getType().getIntOrFloatBitWidth()) {
-        case 32:
-          attributes[name] = static_cast<float>(fp.getValue().convertToFloat());
-          return absl::OkStatus();
-        default:
-          return absl::InvalidArgumentError(absl::StrCat(
-              "Unsupported float attribute bit width for attribute: ", name));
-      }
-    };
-
-    auto str = [&](mlir::StringAttr str) {
-      attributes[name] = str.getValue().str();
-      return absl::OkStatus();
-    };
-
-    TF_RETURN_IF_ERROR(
-        llvm::TypeSwitch<mlir::Attribute, absl::Status>(kv.getValue())
-            .Case<mlir::BoolAttr>(boolean)
-            .Case<mlir::IntegerAttr>(integer)
-            .Case<mlir::FloatAttr>(fp)
-            .Case<mlir::StringAttr>(str)
-            .Default([&](mlir::Attribute) {
-              return absl::InvalidArgumentError(absl::StrCat(
-                  "Unsupported attribute type for attribute: ", name));
-            }));
-  }
-  return attributes;
-}
-
 absl::Span<const int64_t> DecodeDims(int64_t* encoded_dims_data) {
   // Annotate memory coming from jit compiled function as initialized to
   // suppress false positives from msan sanitizer.
@@ -230,7 +142,7 @@ inline absl::Status BuildAndCallFfi(
     // and build an MLIR compatible map of attributes out of it.
     mlir::Attribute attr = mlir::parseAttribute(backend_config, &mlir_context);
     if (auto dict = attr.dyn_cast_or_null<mlir::DictionaryAttr>()) {
-      TF_ASSIGN_OR_RETURN(attributes, BuildAttributesMap(dict));
+      TF_ASSIGN_OR_RETURN(attributes, xla::ffi::BuildAttributesMap(dict));
     } else {
       return absl::InternalError(
           "Unsupported backend config. Expected a string parsable into "
diff --git a/third_party/xla/xla/service/gpu/BUILD b/third_party/xla/xla/service/gpu/BUILD
index 5faf098331bb14..e665640e8ce5eb 100644
--- a/third_party/xla/xla/service/gpu/BUILD
+++ b/third_party/xla/xla/service/gpu/BUILD
@@ -347,6 +347,7 @@ cc_library(
         "//xla:statusor",
         "//xla:util",
         "//xla:xla_data_proto_cc",
+        "//xla/ffi:attribute_map",
         "//xla/ffi:ffi_api",
         "//xla/ffi/api:c_api",
         "//xla/hlo/ir:hlo",
diff --git a/third_party/xla/xla/service/gpu/fusions/BUILD b/third_party/xla/xla/service/gpu/fusions/BUILD
index 02752d4bfffe04..2e0a1202157486 100644
--- a/third_party/xla/xla/service/gpu/fusions/BUILD
+++ b/third_party/xla/xla/service/gpu/fusions/BUILD
@@ -126,6 +126,7 @@ cc_library(
         "//xla:status",
         "//xla:statusor",
         "//xla:util",
+        "//xla/ffi:attribute_map",
         "//xla/ffi:ffi_api",
         "//xla/hlo/ir:hlo",
         "//xla/service:buffer_assignment",
diff --git a/third_party/xla/xla/service/gpu/fusions/custom.cc b/third_party/xla/xla/service/gpu/fusions/custom.cc
index 5d5cc5d3898e08..e6dc11a46cb0be 100644
--- a/third_party/xla/xla/service/gpu/fusions/custom.cc
+++ b/third_party/xla/xla/service/gpu/fusions/custom.cc
@@ -34,6 +34,7 @@ limitations under the License.
 #include "mlir/IR/Attributes.h"  // from @llvm-project
 #include "mlir/IR/BuiltinAttributes.h"  // from @llvm-project
 #include "mlir/Support/LLVM.h"  // from @llvm-project
+#include "xla/ffi/attribute_map.h"
 #include "xla/ffi/ffi_api.h"
 #include "xla/hlo/ir/hlo_casting_utils.h"
 #include "xla/hlo/ir/hlo_instruction.h"
@@ -563,7 +564,7 @@ absl::StatusOr<FusionEmissionResult> EmitCustomCall(
         mlir::Attribute attr = mlir::parseAttribute(
             backend_config_str, ir_emitter_context.mlir_context());
         if (auto dict = mlir::dyn_cast_or_null<mlir::DictionaryAttr>(attr)) {
-          TF_ASSIGN_OR_RETURN(attributes, BuildAttributesMap(dict));
+          TF_ASSIGN_OR_RETURN(attributes, xla::ffi::BuildAttributesMap(dict));
           break;
         }
         return absl::InternalError(
diff --git a/third_party/xla/xla/service/gpu/ir_emitter_unnested.cc b/third_party/xla/xla/service/gpu/ir_emitter_unnested.cc
index 308a7c479822f5..85e39381bac1d9 100644
--- a/third_party/xla/xla/service/gpu/ir_emitter_unnested.cc
+++ b/third_party/xla/xla/service/gpu/ir_emitter_unnested.cc
@@ -72,6 +72,7 @@ limitations under the License.
 #include "mlir/Target/LLVMIR/Dialect/ROCDL/ROCDLToLLVMIRTranslation.h"  // from @llvm-project
 #include "mlir/Target/LLVMIR/Export.h"  // from @llvm-project
 #include "xla/ffi/api/c_api.h"
+#include "xla/ffi/attribute_map.h"
 #include "xla/ffi/ffi_api.h"
 #include "xla/hlo/ir/hlo_casting_utils.h"
 #include "xla/hlo/ir/hlo_computation.h"
@@ -1405,7 +1406,7 @@ absl::Status IrEmitterUnnested::EmitCustomCallThunk(
         mlir::Attribute attr = mlir::parseAttribute(
             backend_config_str, ir_emitter_context_->mlir_context());
         if (auto dict = mlir::dyn_cast_or_null<mlir::DictionaryAttr>(attr)) {
-          TF_ASSIGN_OR_RETURN(attributes, BuildAttributesMap(dict));
+          TF_ASSIGN_OR_RETURN(attributes, xla::ffi::BuildAttributesMap(dict));
           break;
         }
         return absl::InternalError(
diff --git a/third_party/xla/xla/service/gpu/runtime/custom_call_thunk.cc b/third_party/xla/xla/service/gpu/runtime/custom_call_thunk.cc
index 9c25e984252519..b110384ce35563 100644
--- a/third_party/xla/xla/service/gpu/runtime/custom_call_thunk.cc
+++ b/third_party/xla/xla/service/gpu/runtime/custom_call_thunk.cc
@@ -167,99 +167,5 @@ absl::Status CustomCallThunk::ExecuteOnStream(const ExecuteParams& params) {
   return handler_ ? ExecuteFfiHandler(params) : ExecuteCustomCall(params);
 }
 
-absl::StatusOr<CustomCallThunk::AttributesMap> BuildAttributesMap(
-    mlir::DictionaryAttr dict) {
-  CustomCallThunk::AttributesMap attributes;
-  for (auto& kv : dict) {
-    std::string_view name = kv.getName().strref();
-
-    auto boolean = [&](mlir::BoolAttr boolean) {
-      attributes[name] = static_cast<bool>(boolean.getValue());
-      return absl::OkStatus();
-    };
-
-    auto integer = [&](mlir::IntegerAttr integer) {
-      switch (integer.getType().getIntOrFloatBitWidth()) {
-        case 1:
-          attributes[name] = static_cast<bool>(integer.getInt());
-          return absl::OkStatus();
-        case 8:
-          attributes[name] = static_cast<int8_t>(integer.getInt());
-          return absl::OkStatus();
-        case 16:
-          attributes[name] = static_cast<int16_t>(integer.getInt());
-          return absl::OkStatus();
-        case 32:
-          attributes[name] = static_cast<int32_t>(integer.getInt());
-          return absl::OkStatus();
-        case 64:
-          attributes[name] = static_cast<int64_t>(integer.getInt());
-          return absl::OkStatus();
-        default:
-          return absl::InvalidArgumentError(absl::StrCat(
-              "Unsupported integer attribute bit width for attribute: ", name));
-      }
-    };
-
-    auto fp = [&](mlir::FloatAttr fp) {
-      switch (fp.getType().getIntOrFloatBitWidth()) {
-        case 32:
-          attributes[name] = static_cast<float>(fp.getValue().convertToFloat());
-          return absl::OkStatus();
-        case 64:
-          attributes[name] =
-              static_cast<double>(fp.getValue().convertToDouble());
-          return absl::OkStatus();
-        default:
-          return absl::InvalidArgumentError(absl::StrCat(
-              "Unsupported float attribute bit width for attribute: ", name));
-      }
-    };
-
-    auto arr = [&](mlir::DenseArrayAttr arr) {
-      if (auto dense = mlir::dyn_cast<mlir::DenseI8ArrayAttr>(arr)) {
-        attributes[name] = dense.asArrayRef().vec();
-        return absl::OkStatus();
-      } else if (auto dense = mlir::dyn_cast<mlir::DenseI16ArrayAttr>(arr)) {
-        attributes[name] = dense.asArrayRef().vec();
-        return absl::OkStatus();
-      } else if (auto dense = mlir::dyn_cast<mlir::DenseI32ArrayAttr>(arr)) {
-        attributes[name] = dense.asArrayRef().vec();
-        return absl::OkStatus();
-      } else if (auto dense = mlir::dyn_cast<mlir::DenseI64ArrayAttr>(arr)) {
-        attributes[name] = dense.asArrayRef().vec();
-        return absl::OkStatus();
-      } else if (auto dense = mlir::dyn_cast<mlir::DenseF32ArrayAttr>(arr)) {
-        attributes[name] = dense.asArrayRef().vec();
-        return absl::OkStatus();
-      } else if (auto dense = mlir::dyn_cast<mlir::DenseF64ArrayAttr>(arr)) {
-        attributes[name] = dense.asArrayRef().vec();
-        return absl::OkStatus();
-      } else {
-        return absl::InvalidArgumentError(absl::StrCat(
-            "Unsupported array element type for attribute: ", name));
-      }
-    };
-
-    auto str = [&](mlir::StringAttr str) {
-      attributes[name] = str.getValue().str();
-      return absl::OkStatus();
-    };
-
-    TF_RETURN_IF_ERROR(
-        llvm::TypeSwitch<mlir::Attribute, Status>(kv.getValue())
-            .Case<mlir::BoolAttr>(boolean)
-            .Case<mlir::IntegerAttr>(integer)
-            .Case<mlir::FloatAttr>(fp)
-            .Case<mlir::DenseArrayAttr>(arr)
-            .Case<mlir::StringAttr>(str)
-            .Default([&](mlir::Attribute) {
-              return absl::InvalidArgumentError(absl::StrCat(
-                  "Unsupported attribute type for attribute: ", name));
-            }));
-  }
-  return attributes;
-}
-
 }  // namespace gpu
 }  // namespace xla
diff --git a/third_party/xla/xla/service/gpu/runtime/custom_call_thunk.h b/third_party/xla/xla/service/gpu/runtime/custom_call_thunk.h
index 02679d2e0d21ff..2d797ecea01a6c 100644
--- a/third_party/xla/xla/service/gpu/runtime/custom_call_thunk.h
+++ b/third_party/xla/xla/service/gpu/runtime/custom_call_thunk.h
@@ -120,11 +120,6 @@ class CustomCallThunk : public Thunk {
   const HloComputation* called_computation_ = nullptr;
 };
 
-// Converts MLIR dictionary attribute attached to a custom call operation to a
-// custom call thunk attributes that are forwarded to the FFI handler.
-absl::StatusOr<CustomCallThunk::AttributesMap> BuildAttributesMap(
-    mlir::DictionaryAttr dict);
-
 }  // namespace gpu
 }  // namespace xla
 
diff --git a/third_party/xla/xla/tests/BUILD b/third_party/xla/xla/tests/BUILD
index 2bffc0a1316e5f..58600c91d25852 100644
--- a/third_party/xla/xla/tests/BUILD
+++ b/third_party/xla/xla/tests/BUILD
@@ -1783,6 +1783,7 @@ xla_test(
         "//xla/client:xla_builder",
         "//xla/client/lib:constants",
         "//xla/ffi",
+        "//xla/ffi:attribute_map",
         "//xla/ffi:ffi_api",
         "//xla/hlo/ir:hlo",
         "//xla/service:custom_call_status",

From b8629e1a9a0df24df2b9ce7f51a594ffd902519a Mon Sep 17 00:00:00 2001
From: Johannes Reifferscheid <jreiffers@google.com>
Date: Tue, 14 May 2024 04:04:35 -0700
Subject: [PATCH 105/478] Canonicalize affine expression trees in simplifier.

I don't know where the current non-determinism comes from, but this should
fix it.

The particular order is not important, but it being canonical is. These
expressions are used for codegen in the new emitters, so the simplification
must be deterministic.

PiperOrigin-RevId: 633531399
---
 .../gpu/fusions/concatenate_mlir_test.cc      |  2 +-
 .../service/gpu/fusions/concatenate_test.cc   |  2 +-
 .../service/gpu/fusions/input_slices_test.cc  |  4 +-
 .../xla/service/gpu/fusions/loop_mlir_test.cc |  4 +-
 .../xla/xla/service/gpu/fusions/loop_test.cc  |  4 +-
 .../gpu/fusions/reduction_base_test.cc        |  2 +-
 .../service/gpu/fusions/scatter_mlir_test.cc  |  2 +-
 .../xla/service/gpu/fusions/scatter_test.cc   |  2 +-
 .../gpu/fusions/transpose_mlir_test.cc        | 12 +--
 .../xla/service/gpu/fusions/transpose_test.cc |  4 +-
 .../gpu/model/indexing_analysis_test.cc       |  8 +-
 .../xla/xla/service/gpu/model/indexing_map.cc | 70 ++++++++++++++++-
 .../service/gpu/model/indexing_map_test.cc    |  6 +-
 .../service/gpu/model/indexing_test_utils.cc  | 77 +++----------------
 14 files changed, 107 insertions(+), 92 deletions(-)

diff --git a/third_party/xla/xla/service/gpu/fusions/concatenate_mlir_test.cc b/third_party/xla/xla/service/gpu/fusions/concatenate_mlir_test.cc
index eea090ed80474a..781ebf563f66ea 100644
--- a/third_party/xla/xla/service/gpu/fusions/concatenate_mlir_test.cc
+++ b/third_party/xla/xla/service/gpu/fusions/concatenate_mlir_test.cc
@@ -57,7 +57,7 @@ TEST_F(MlirConcatenateFusionTest, ThreadIdIndexing) {
 
   constexpr auto kIndexing = R"(
     (th_x, th_y, th_z, bl_x, bl_y, bl_z)[chunk_id, unroll_id] -> (
-    (th_x + bl_x * 128) mod 400)
+    (bl_x * 128 + th_x) mod 400)
     domain:
     th_x in [0, 127]
     th_y in [0, 0]
diff --git a/third_party/xla/xla/service/gpu/fusions/concatenate_test.cc b/third_party/xla/xla/service/gpu/fusions/concatenate_test.cc
index b617fd6513d107..4ac8ca3c665497 100644
--- a/third_party/xla/xla/service/gpu/fusions/concatenate_test.cc
+++ b/third_party/xla/xla/service/gpu/fusions/concatenate_test.cc
@@ -80,7 +80,7 @@ TEST_F(ConcatenateTest, ThreadIndexing) {
 
   constexpr auto kIndexing = R"(
     (th_x, th_y, th_z, bl_x, bl_y, bl_z)[chunk_id, unroll_id] -> (
-    (th_x + bl_x * 128) mod 400)
+    (bl_x * 128 + th_x) mod 400)
     domain:
     th_x in [0, 127]
     th_y in [0, 0]
diff --git a/third_party/xla/xla/service/gpu/fusions/input_slices_test.cc b/third_party/xla/xla/service/gpu/fusions/input_slices_test.cc
index 094bbfac7a27a9..9c10beb8be1e12 100644
--- a/third_party/xla/xla/service/gpu/fusions/input_slices_test.cc
+++ b/third_party/xla/xla/service/gpu/fusions/input_slices_test.cc
@@ -80,8 +80,8 @@ TEST_F(InputSlicesTest, ThreadIndexing) {
   EXPECT_THAT(thread_id_to_output_indexing->ToString(printer_),
               MatchIndexingString(R"(
     (th_x, th_y, th_z, bl_x, bl_y, bl_z)[chunk_id, unroll_id] -> (0,
-      ((th_x + bl_x * 128) floordiv 3) mod 2,
-       (th_x + bl_x * 128) mod 3,
+      ((bl_x * 128 + th_x) floordiv 3) mod 2,
+       (bl_x * 128 + th_x) mod 3,
        ((bl_x * 64 + th_x floordiv 2) floordiv 3) mod 5)
     domain:
     th_x in [0, 127]
diff --git a/third_party/xla/xla/service/gpu/fusions/loop_mlir_test.cc b/third_party/xla/xla/service/gpu/fusions/loop_mlir_test.cc
index 4a283349779fdc..bd319d9d86d60d 100644
--- a/third_party/xla/xla/service/gpu/fusions/loop_mlir_test.cc
+++ b/third_party/xla/xla/service/gpu/fusions/loop_mlir_test.cc
@@ -55,7 +55,7 @@ TEST_F(MlirLoopFusionTest, ThreadId_IndexingUnrolled) {
               MatchIndexingString(R"(
   (th_x, th_y, th_z, bl_x, bl_y, bl_z)[chunk_id, unroll_id] -> (
    (((bl_x * 16 + th_x floordiv 8) floordiv 3 + chunk_id * 5376) floordiv 625) mod 100,
-   (((th_x + bl_x * 128) floordiv 3 + chunk_id * 43008) floordiv 25) mod 200,
+   (((bl_x * 128 + th_x) floordiv 3 + chunk_id * 43008) floordiv 25) mod 200,
    (th_x * 4 + bl_x * 512 + chunk_id * 516096) mod 300 + unroll_id
   )
   domain:
@@ -150,7 +150,7 @@ TEST_F(MlirLoopFusionTest, ThreadId_Broadcast) {
               (th_x, th_y, th_z, bl_x, bl_y, bl_z)[chunk_id, unroll_id] -> (
                 ((bl_x * 16 + th_x floordiv 8) floordiv 75) mod 10,
                 ((bl_x * 64 + th_x floordiv 2) floordiv 15) mod 20,
-                (th_x + bl_x * 128) mod 30)
+                (bl_x * 128 + th_x) mod 30)
                 domain:
                 th_x in [0, 127]
                 th_y in [0, 0]
diff --git a/third_party/xla/xla/service/gpu/fusions/loop_test.cc b/third_party/xla/xla/service/gpu/fusions/loop_test.cc
index d87bc8fa0db98e..4507879c2310ea 100644
--- a/third_party/xla/xla/service/gpu/fusions/loop_test.cc
+++ b/third_party/xla/xla/service/gpu/fusions/loop_test.cc
@@ -90,7 +90,7 @@ TEST_F(LoopTest, ThreadIndexingUnrolled) {
               MatchIndexingString(R"(
   (th_x, th_y, th_z, bl_x, bl_y, bl_z)[chunk_id, unroll_id] -> (
    (((bl_x * 16 + th_x floordiv 8) floordiv 3 + chunk_id * 5376) floordiv 625) mod 100,
-   (((th_x + bl_x * 128) floordiv 3 + chunk_id * 43008) floordiv 25) mod 200,
+   (((bl_x * 128 + th_x) floordiv 3 + chunk_id * 43008) floordiv 25) mod 200,
    (th_x * 4 + bl_x * 512 + chunk_id * 516096) mod 300 + unroll_id
   )
   domain:
@@ -186,7 +186,7 @@ TEST_F(LoopTest, Broadcast) {
               (th_x, th_y, th_z, bl_x, bl_y, bl_z)[chunk_id, unroll_id] -> (
                 ((bl_x * 16 + th_x floordiv 8) floordiv 75) mod 10,
                 ((bl_x * 64 + th_x floordiv 2) floordiv 15) mod 20,
-                (th_x + bl_x * 128) mod 30)
+                (bl_x * 128 + th_x) mod 30)
                 domain:
                 th_x in [0, 127]
                 th_y in [0, 0]
diff --git a/third_party/xla/xla/service/gpu/fusions/reduction_base_test.cc b/third_party/xla/xla/service/gpu/fusions/reduction_base_test.cc
index aefb4ad8219313..6a43851f555bc7 100644
--- a/third_party/xla/xla/service/gpu/fusions/reduction_base_test.cc
+++ b/third_party/xla/xla/service/gpu/fusions/reduction_base_test.cc
@@ -330,7 +330,7 @@ TEST_F(ReductionTest, ThreadIndexingSideOutput) {
   constexpr char kExpectedIndexing[] = R"(
       (d0, d1, d2, d3, d4, d5)[s0, s1, s2, s3] -> (
         d3 floordiv 8,
-        d0 floordiv 32 + (d3 mod 8) * 8,
+        (d3 mod 8) * 8 + d0 floordiv 32,
         (d0 mod 32) * 2 + s2 * 64 + s3
       )
       domain:
diff --git a/third_party/xla/xla/service/gpu/fusions/scatter_mlir_test.cc b/third_party/xla/xla/service/gpu/fusions/scatter_mlir_test.cc
index 713314969c8d82..86ce79992c819c 100644
--- a/third_party/xla/xla/service/gpu/fusions/scatter_mlir_test.cc
+++ b/third_party/xla/xla/service/gpu/fusions/scatter_mlir_test.cc
@@ -82,7 +82,7 @@ TEST_F(MlirScatterFusionTest, ThreadIdIndexing) {
     (th_x, th_y, th_z, bl_x, bl_y, bl_z)[chunk_id, unroll_id] -> (
     ((bl_x * 16 + th_x floordiv 8) floordiv 25) mod 42,
     ((bl_x * 32 + th_x floordiv 4) floordiv 5) mod 10,
-    (th_x + bl_x * 128) mod 20)
+    (bl_x * 128 + th_x) mod 20)
     domain:
     th_x in [0, 127]
     th_y in [0, 0]
diff --git a/third_party/xla/xla/service/gpu/fusions/scatter_test.cc b/third_party/xla/xla/service/gpu/fusions/scatter_test.cc
index 2be8dc86d75540..07e2da47a6a6b0 100644
--- a/third_party/xla/xla/service/gpu/fusions/scatter_test.cc
+++ b/third_party/xla/xla/service/gpu/fusions/scatter_test.cc
@@ -148,7 +148,7 @@ TEST_F(ScatterFusionTest, ThreadIdIndexing) {
     (th_x, th_y, th_z, bl_x, bl_y, bl_z)[chunk_id, unroll_id] -> (
     ((bl_x * 16 + th_x floordiv 8) floordiv 25) mod 42,
     ((bl_x * 32 + th_x floordiv 4) floordiv 5) mod 10,
-    (th_x + bl_x * 128) mod 20)
+    (bl_x * 128 + th_x) mod 20)
     domain:
     th_x in [0, 127]
     th_y in [0, 0]
diff --git a/third_party/xla/xla/service/gpu/fusions/transpose_mlir_test.cc b/third_party/xla/xla/service/gpu/fusions/transpose_mlir_test.cc
index 64c66419c48a1b..086ebf8f2ad911 100644
--- a/third_party/xla/xla/service/gpu/fusions/transpose_mlir_test.cc
+++ b/third_party/xla/xla/service/gpu/fusions/transpose_mlir_test.cc
@@ -51,8 +51,8 @@ TEST_F(MlirTransposeFusionTest, ThreadIndexing021) {
       MatchIndexingString(R"(
         (d0, d1, d2, d3, d4, d5)[s0, s1, s2] -> (
           d3 floordiv 2,
-          ### d0 floordiv 32 + s1 * 4 ###,
-          ### (d3 mod 2) * 32 + d0 mod 32 ###
+          d0 floordiv 32 + s1 * 4,
+          (d3 mod 2) * 32 + d0 mod 32
         )
         domain:
         d0 in [0, 127]
@@ -71,7 +71,7 @@ TEST_F(MlirTransposeFusionTest, ThreadIndexing021) {
       MatchIndexingString(R"(
         (d0, d1, d2, d3, d4, d5)[s0, s1, s2] -> (
           d3 floordiv 2,
-          ### d0 floordiv 32 + (d3 mod 2) * 32 + s1 * 4 ###,
+          (d3 mod 2) * 32 + s1 * 4 + d0 floordiv 32,
           d0 mod 32
         )
         domain:
@@ -110,7 +110,7 @@ TEST_F(MlirTransposeFusionTest, ThreadIndexing201) {
       MatchIndexingString(R"(
         (d0, d1, d2, d3, d4, d5)[s0, s1, s2] -> (
           d3 floordiv 2,
-          d0 floordiv 32 + (d3 * 32 + s1 * 4) mod 64,
+          (d3 * 32 + s1 * 4) mod 64 + d0 floordiv 32,
           d0 mod 32
         )
         domain:
@@ -129,9 +129,9 @@ TEST_F(MlirTransposeFusionTest, ThreadIndexing201) {
       fusion.ComputeThreadIdToOutputIndexing(0, &mlir_context_)->ToString(),
       MatchIndexingString(R"(
         (d0, d1, d2, d3, d4, d5)[s0, s1, s2] -> (
-          ### d0 floordiv 32 + s1 * 4 ###,
+          d0 floordiv 32 + s1 * 4,
           d3 floordiv 2,
-          ### (d3 mod 2) * 32 + d0 mod 32 ###
+          (d3 mod 2) * 32 + d0 mod 32
         )
         domain:
         d0 in [0, 127]
diff --git a/third_party/xla/xla/service/gpu/fusions/transpose_test.cc b/third_party/xla/xla/service/gpu/fusions/transpose_test.cc
index dbafd613d15b2b..a9b40704ff7842 100644
--- a/third_party/xla/xla/service/gpu/fusions/transpose_test.cc
+++ b/third_party/xla/xla/service/gpu/fusions/transpose_test.cc
@@ -97,7 +97,7 @@ TEST_F(TransposeTest, ThreadIndexing021) {
       MatchIndexingString(R"(
         (d0, d1, d2, d3, d4, d5)[s0, s1, s2] -> (
           d3 floordiv 2,
-          d0 floordiv 32 + (d3 mod 2) * 32 + s1 * 4,
+          (d3 mod 2) * 32 + s1 * 4 + d0 floordiv 32,
           d0 mod 32
         )
         domain:
@@ -139,7 +139,7 @@ TEST_F(TransposeTest, ThreadIndexing201) {
       MatchIndexingString(R"(
         (d0, d1, d2, d3, d4, d5)[s0, s1, s2] -> (
           d3 floordiv 2,
-          d0 floordiv 32 + (d3 * 32 + s1 * 4) mod 64,
+          (d3 * 32 + s1 * 4) mod 64 + d0 floordiv 32,
           d0 mod 32
         )
         domain:
diff --git a/third_party/xla/xla/service/gpu/model/indexing_analysis_test.cc b/third_party/xla/xla/service/gpu/model/indexing_analysis_test.cc
index 0413c8280786cc..daddff1f4c1436 100644
--- a/third_party/xla/xla/service/gpu/model/indexing_analysis_test.cc
+++ b/third_party/xla/xla/service/gpu/model/indexing_analysis_test.cc
@@ -1643,7 +1643,7 @@ TEST_F(IndexingAnalysisTest, ReshapeOpGenericReshape2DTo3D) {
   EXPECT_THAT(input_indexing.indexing_maps,
               ElementsAre(ElementsAre(MatchIndexingMap(R"(
                 (d0, d1, d2) -> (d0 * 2 + d1 floordiv 2,
-                                d2 + (d1 mod 2) * 4)
+                                (d1 mod 2) * 4 + d2)
                 domain:
                 d0 in [0, 1]
                 d1 in [0, 3]
@@ -1662,7 +1662,7 @@ TEST_F(IndexingAnalysisTest, ReshapeOpGenericReshape3DTo2D) {
   EXPECT_THAT(input_indexing.indexing_maps,
               ElementsAre(ElementsAre(MatchIndexingMap(R"(
                             (d0, d1) -> (d0 floordiv 2,
-                                        d1 floordiv 4 + (d0 mod 2) * 2,
+                                        (d0 mod 2) * 2 + d1 floordiv 4,
                                         d1 mod 4)
                             domain:
                             d0 in [0, 3]
@@ -2615,7 +2615,7 @@ TEST_F(IndexingAnalysisTest, TilingIndexing) {
   EXPECT_THAT(indexing_map.ToString(), MatchIndexingString(R"(
         (d0, d1, d2, d3, d4, d5)[s0, s1, s2] -> (
           (d3 floordiv 64) * 8 + s0,
-          d0 floordiv 4 + (d3 mod 64) * 4,
+          (d3 mod 64) * 4 + d0 floordiv 4,
           d0 mod 4 + s2 * 4
         )
         domain:
@@ -2659,7 +2659,7 @@ TEST_F(IndexingAnalysisTest, EpilogueIndexing) {
       ComputeEpilogueInputToOutputIndexing(transpose, log, &mlir_context_)
           .ToString(),
       MatchIndexingString(R"(
-                  (d0, d1) -> (d0 + d1 * 1000)
+                  (d0, d1) -> (d1 * 1000 + d0)
                   domain:
                   d0 in [0, 999]
                   d1 in [0, 999]
diff --git a/third_party/xla/xla/service/gpu/model/indexing_map.cc b/third_party/xla/xla/service/gpu/model/indexing_map.cc
index f12f923342276c..80390bb6c119e8 100644
--- a/third_party/xla/xla/service/gpu/model/indexing_map.cc
+++ b/third_party/xla/xla/service/gpu/model/indexing_map.cc
@@ -32,6 +32,7 @@ limitations under the License.
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/SmallBitVector.h"
 #include "llvm/ADT/SmallVector.h"
+#include "llvm/Support/raw_ostream.h"
 #include "mlir/IR/AffineExpr.h"  // from @llvm-project
 #include "mlir/IR/AffineMap.h"  // from @llvm-project
 #include "mlir/IR/MLIRContext.h"  // from @llvm-project
@@ -359,6 +360,69 @@ AffineExpr AffineExprSimplifier::RewriteSumIf(
   return pred(expr) ? expr : mlir::getAffineConstantExpr(0, expr.getContext());
 }
 
+// Compares the two expression by their AST. The ordering is arbitrary but
+// similar to what MLIR's simplifier does.
+int CompareExprs(AffineExpr a, AffineExpr b) {
+  if ((b.getKind() == AffineExprKind::Constant) !=
+      (a.getKind() == AffineExprKind::Constant)) {
+    return a.getKind() == AffineExprKind::Constant ? 1 : -1;
+  }
+  if (a.getKind() < b.getKind()) {
+    return -1;
+  }
+  if (a.getKind() > b.getKind()) {
+    return 1;
+  }
+  assert(a.getKind() == b.getKind());
+  int64_t a_value = 0;
+  int64_t b_value = 0;
+  switch (a.getKind()) {
+    case AffineExprKind::Add:
+    case AffineExprKind::FloorDiv:
+    case AffineExprKind::CeilDiv:
+    case AffineExprKind::Mul:
+    case AffineExprKind::Mod: {
+      auto a_bin = mlir::cast<AffineBinaryOpExpr>(a);
+      auto b_bin = mlir::cast<AffineBinaryOpExpr>(b);
+      auto lhs = CompareExprs(a_bin.getLHS(), b_bin.getLHS());
+      if (lhs != 0) {
+        return lhs;
+      }
+      return CompareExprs(a_bin.getRHS(), b_bin.getRHS());
+    }
+    case AffineExprKind::Constant: {
+      a_value = mlir::cast<AffineConstantExpr>(a).getValue();
+      b_value = mlir::cast<AffineConstantExpr>(b).getValue();
+      break;
+    }
+    case AffineExprKind::SymbolId: {
+      a_value = mlir::cast<AffineSymbolExpr>(a).getPosition();
+      b_value = mlir::cast<AffineSymbolExpr>(b).getPosition();
+      break;
+    }
+    case AffineExprKind::DimId: {
+      a_value = mlir::cast<AffineDimExpr>(a).getPosition();
+      b_value = mlir::cast<AffineDimExpr>(b).getPosition();
+      break;
+    }
+  }
+  return a_value < b_value ? -1 : (a_value > b_value ? 1 : 0);
+}
+
+AffineExpr CanonicalizeOrder(AffineExpr in) {
+  if (auto binop = mlir::dyn_cast<AffineBinaryOpExpr>(in)) {
+    auto lhs = CanonicalizeOrder(binop.getLHS());
+    auto rhs = CanonicalizeOrder(binop.getRHS());
+    if ((binop.getKind() == AffineExprKind::Add ||
+         binop.getKind() == AffineExprKind::Mul) &&
+        CompareExprs(lhs, rhs) > 0) {
+      std::swap(lhs, rhs);
+    }
+    return getAffineBinaryOpExpr(binop.getKind(), lhs, rhs);
+  }
+  return in;
+}
+
 AffineExpr AffineExprSimplifier::SimplifyOnce(AffineExpr expr) {
   switch (expr.getKind()) {
     case AffineExprKind::Mul: {
@@ -478,7 +542,11 @@ AffineMap AffineExprSimplifier::Simplify(AffineMap affine_map) {
     results.push_back(simplified);
   }
   if (nothing_changed) {
-    return affine_map;
+    for (auto& result : results) {
+      result = CanonicalizeOrder(result);
+    }
+    return AffineMap::get(affine_map.getNumDims(), affine_map.getNumSymbols(),
+                          results, affine_map.getContext());
   }
   return Simplify(AffineMap::get(affine_map.getNumDims(),
                                  affine_map.getNumSymbols(), results,
diff --git a/third_party/xla/xla/service/gpu/model/indexing_map_test.cc b/third_party/xla/xla/service/gpu/model/indexing_map_test.cc
index 20807546347c0c..f753d945cf15d2 100644
--- a/third_party/xla/xla/service/gpu/model/indexing_map_test.cc
+++ b/third_party/xla/xla/service/gpu/model/indexing_map_test.cc
@@ -575,7 +575,7 @@ TEST_F(IndexingMapTest,
       ParseAffineMap(serialized_map, &mlir_context_), {10, 10, 10}, {});
   indexing_map.Simplify(GetIndexingMapForInstruction);
   EXPECT_THAT(indexing_map.ToString(printer_), MatchIndexingString(R"(
-    (d0, d1, d2) -> (d0 * 2 + (d1 + d2 floordiv 4) floordiv 2,
+    (d0, d1, d2) -> (d0 * 2 + (d2 floordiv 4 + d1) floordiv 2,
                      (d1 * 4 + d2) mod 8)
     domain:
     d0 in [0, 9]
@@ -619,7 +619,7 @@ TEST_F(IndexingMapTest, AffineMapSimplification_SimplifyReshape_Regression) {
       ParseAffineMap(serialized_map, &mlir_context_), {}, {128});
   indexing_map.Simplify(GetIndexingMapForInstruction);
   EXPECT_THAT(indexing_map.ToString(printer_), MatchIndexingString(R"(
-      ()[s0] -> ((s0 * 128) mod 715 + ((s0 * 64) floordiv 715) * 715)
+      ()[s0] -> (((s0 * 64) floordiv 715) * 715 + (s0 * 128) mod 715)
       domain: s0 in [0, 127]
   )"));
 }
@@ -663,7 +663,7 @@ TEST_F(IndexingMapTest, AffineMapSimplification_ExtractFromMod) {
   indexing_map.Simplify(GetIndexingMapForInstruction);
   EXPECT_THAT(indexing_map.ToString(printer_), MatchIndexingString(R"(
       ()[s0, s1, s2, s3] -> (
-        s1 + (s0 * 458752 + s2 * 4 + s3 * 512) mod 20000
+        (s0 * 458752 + s2 * 4 + s3 * 512) mod 20000 + s1
       )
       domain:
       s0 in [0, 871]
diff --git a/third_party/xla/xla/service/gpu/model/indexing_test_utils.cc b/third_party/xla/xla/service/gpu/model/indexing_test_utils.cc
index 67163320e04275..72becac80c075e 100644
--- a/third_party/xla/xla/service/gpu/model/indexing_test_utils.cc
+++ b/third_party/xla/xla/service/gpu/model/indexing_test_utils.cc
@@ -140,78 +140,25 @@ AffineExpr ParseAffineExpr(absl::string_view serialized_affine_expr,
       .getResult(0);
 }
 
-inline std::vector<std::string> split_string(std::string s,
-                                             std::string pattern) {
-  std::vector<std::string> result;
-  size_t pos = 0;
-  while ((pos = s.find(pattern)) != std::string::npos) {
-    result.push_back(s.substr(0, pos));
-    s.erase(0, pos + pattern.length());
-  }
-  if (!s.empty()) result.push_back(s);
-  return result;
-}
-
-inline bool startswith(const std::string& s, const std::string& pattern) {
-  return s.substr(0, pattern.size()) == pattern;
-}
-
 bool ApproximateMatch(std::string_view lhs, std::string_view rhs) {
-  std::string lhs_unspaced, rhs_unspaced;
-  for (auto c : lhs) {
-    if (!std::isspace(c)) {
-      lhs_unspaced += c;
+  size_t lhs_length = lhs.size();
+  size_t rhs_length = rhs.size();
+  size_t l = 0, r = 0;
+  while (l < lhs_length && r < rhs_length) {
+    while (l < lhs_length && std::isspace(lhs[l])) {
+      ++l;
     }
-  }
-  for (auto c : rhs) {
-    if (!std::isspace(c)) {
-      rhs_unspaced += c;
-    }
-  }
-
-  if (lhs_unspaced.find("###") == std::string::npos)
-    return lhs_unspaced == rhs_unspaced;
-
-  std::vector<std::string> frags = split_string(lhs_unspaced, "###");
-
-  while (frags.size() >= 2) {
-    if (!startswith(rhs_unspaced, frags[0])) {
-      return false;
+    while (r < rhs_length && std::isspace(rhs[r])) {
+      ++r;
     }
-
-    rhs_unspaced = rhs_unspaced.substr(frags[0].size());
-
-    auto terms = split_string(frags[1], "+");
-    // iterate through permutations of terms
-    std::vector<int> indexes(terms.size());
-    for (auto i = 0; i < terms.size(); i++) {
-      indexes[i] = i;
+    if (l == lhs_length || r == rhs_length) {
+      continue;
     }
-    bool match = false;
-    do {
-      std::string permuted = "";
-      for (auto i : indexes) {
-        permuted += terms[i] + "+";
-      }
-      permuted.pop_back();
-      if (startswith(rhs_unspaced, permuted)) {
-        match = true;
-        break;
-      }
-    } while (std::next_permutation(indexes.begin(), indexes.end()));
-
-    if (!match) {
+    if (lhs[l++] != rhs[r++]) {
       return false;
     }
-
-    rhs_unspaced = rhs_unspaced.substr(frags[1].size());
-    frags.erase(frags.begin());
-    frags.erase(frags.begin());
   }
-  if (frags.empty())
-    return rhs_unspaced.empty();
-  else
-    return rhs_unspaced == frags[0];
+  return l == lhs_length && r == rhs_length;
 }
 
 }  // namespace gpu

From 055675be44b77097a01b8fff9275ce11df238918 Mon Sep 17 00:00:00 2001
From: Ilia Sergachev <isergachev@nvidia.com>
Date: Tue, 14 May 2024 04:13:51 -0700
Subject: [PATCH 106/478] PR #12434: [GPU] Fix OSS compilation problems in a
 previously disabled test.

Imported from GitHub PR https://github.com/openxla/xla/pull/12434

Copybara import of the project:

--
723c9bb29adfcc33c015b74f90ce8024c2f79255 by Ilia Sergachev <isergachev@nvidia.com>:

[GPU] Fix OSS compilation problems in a previously disabled test.

Merging this change closes #12434

PiperOrigin-RevId: 633533757
---
 .../xla/xla/service/gpu/ir_emitter_triton_test.cc      | 10 +++++-----
 1 file changed, 5 insertions(+), 5 deletions(-)

diff --git a/third_party/xla/xla/service/gpu/ir_emitter_triton_test.cc b/third_party/xla/xla/service/gpu/ir_emitter_triton_test.cc
index 111f8e276e1de8..972b5cb57c8b6c 100644
--- a/third_party/xla/xla/service/gpu/ir_emitter_triton_test.cc
+++ b/third_party/xla/xla/service/gpu/ir_emitter_triton_test.cc
@@ -284,7 +284,7 @@ ENTRY e {
 })";
 
   TritonGemmConfig config(16, 16, 32, 1, 1, 1);
-  EXPECT_OK(
+  TF_EXPECT_OK(
       CreateTritonIrAndFileCheck(kHloText, config, EmitMatMul, "triton_dot", R"(
 CHECK:    tt.func @triton_fn(%[[LHS:.*]]: !tt.ptr<f32> {tt.divisibility = 16 : i32}, %[[RHS:.*]]: !tt.ptr<f32> {tt.divisibility = 16 : i32}, %[[OUT:.*]]: !tt.ptr<f32> {tt.divisibility = 16 : i32}) {
 CHECK-DAG:  %[[ZERO_KN:.*]] = arith.constant dense<0.000000e+00> : tensor<32x16xf32>
@@ -1083,8 +1083,8 @@ ENTRY main {
                           ParseAndReturnVerifiedModule(kHloText));
 
   TritonGemmConfig config(16, 64, 32, 1, 1, 1);
-  ASSERT_OK(CreateTritonIrAndFileCheck(kHloText, config, EmitSoftMax,
-                                       "triton_softmax_computation", R"(
+  TF_ASSERT_OK(CreateTritonIrAndFileCheck(kHloText, config, EmitSoftMax,
+                                          "triton_softmax_computation", R"(
 // CHECK: #[[MAP:.*]] = affine_map<()[s0] -> (s0 * 16)>
 // CHECK-LABEL:   tt.func @triton_fn(
 // CHECK-SAME:                       %[[P0:[^:]*]]: !tt.ptr<f32> {tt.divisibility = 16 : i32},
@@ -4550,7 +4550,7 @@ ENTRY e {
 }
 )";
   TritonGemmConfig config(32, 32, 32, 1, 1, 1);
-  ASSERT_OK(
+  TF_ASSERT_OK(
       CreateTritonIrAndFileCheck(kHloText, config, EmitMatMul, "triton_dot", R"(
 CHECK:          %[[INFINITY:.*]] = arith.constant dense<0x7F800000> : tensor<32x32xf32>
 CHECK:          %[[C_MASK:.*]] = arith.constant dense<-65536> : tensor<32x32xi32>
@@ -4881,7 +4881,7 @@ ENTRY e {
 }
 )";
   TritonGemmConfig config(32, 32, 32, 1, 1, 1);
-  ASSERT_OK(
+  TF_ASSERT_OK(
       CreateTritonIrAndFileCheck(kHloText, config, EmitMatMul, "triton_dot", R"(
 CHECK:          %[[INFINITY:.*]] = arith.constant dense<0x7F800000> : tensor<32x32xf32>
 CHECK:          %[[C_MASK:.*]] = arith.constant dense<-65536> : tensor<32x32xi32>

From 614ea89256e015f28294ea081a76423ae480f6b9 Mon Sep 17 00:00:00 2001
From: Dmitri Gribenko <dmitrig@google.com>
Date: Tue, 14 May 2024 04:36:11 -0700
Subject: [PATCH 107/478] Integrate LLVM at llvm/llvm-project@79a6a7e28fff

Updates LLVM usage to match
[79a6a7e28fff](https://github.com/llvm/llvm-project/commit/79a6a7e28fff)

PiperOrigin-RevId: 633538634
---
 third_party/llvm/workspace.bzl                |  4 ++--
 third_party/stablehlo/temporary.patch         | 22 +++++++++++++++++++
 .../xla/third_party/stablehlo/temporary.patch | 22 +++++++++++++++++++
 3 files changed, 46 insertions(+), 2 deletions(-)

diff --git a/third_party/llvm/workspace.bzl b/third_party/llvm/workspace.bzl
index 669cff9644af1a..eb3a0f80219c81 100644
--- a/third_party/llvm/workspace.bzl
+++ b/third_party/llvm/workspace.bzl
@@ -4,8 +4,8 @@ load("//third_party:repo.bzl", "tf_http_archive")
 
 def repo(name):
     """Imports LLVM."""
-    LLVM_COMMIT = "e6b2197a89f5d6d0f56a03c03b8afda561eee899"
-    LLVM_SHA256 = "a2398fc87e7c4ef96f17999a83240d421686718a8f095fed718b7539576c1fb7"
+    LLVM_COMMIT = "79a6a7e28fffd14e54a9a208af12d724b6eeb2d4"
+    LLVM_SHA256 = "ce0caa77cca929e29b0a6304820b833f4514a56bbfaaa533df846077916aa72d"
 
     tf_http_archive(
         name = name,
diff --git a/third_party/stablehlo/temporary.patch b/third_party/stablehlo/temporary.patch
index 2e395d99d26cf4..f906a856ae69e1 100755
--- a/third_party/stablehlo/temporary.patch
+++ b/third_party/stablehlo/temporary.patch
@@ -361,6 +361,28 @@ diff --ruN a/stablehlo/stablehlo/conversions/tosa/tests/binary.mlir b/stablehlo/
    // CHECK: stablehlo.divide
    %0 = "stablehlo.divide"(%arg0, %arg1) : (tensor<10xf32>, tensor<10xf32>) -> tensor<10xf32>
    return %0 : tensor<10xf32>
+@@ -123,7 +123,7 @@
+ 
+ // CHECK-LABEL: @maximum_f64
+ func.func @maximum_f64(%arg0 : tensor<10xf64>, %arg1 : tensor<10xf64>) -> tensor<10xf64> {
+-  // CHECK: stablehlo.maximum
++  // CHECK: tosa.maximum
+   %0 = "stablehlo.maximum"(%arg0, %arg1) : (tensor<10xf64>, tensor<10xf64>) -> tensor<10xf64>
+   return %0 : tensor<10xf64>
+ }
+diff --ruN a/stablehlo/stablehlo/conversions/tosa/tests/nullary.mlir b/stablehlo/stablehlo/conversions/tosa/tests/nullary.mlir
+--- stablehlo/stablehlo/conversions/tosa/tests/nullary.mlir
++++ stablehlo/stablehlo/conversions/tosa/tests/nullary.mlir
+@@ -9,8 +9,7 @@
+ 
+ // CHECK-LABEL: @constant_f64
+ func.func @constant_f64() -> tensor<10xf64> {
+-  // TOSA does not support 64-bit types, so this should not legalize.
+-  // CHECK: stablehlo.constant
++  // CHECK: tosa.const
+   %0 = stablehlo.constant dense<0.000000e+00> : tensor<10xf64>
+   return %0 : tensor<10xf64>
+ }
 diff --ruN a/stablehlo/stablehlo/conversions/tosa/transforms/StablehloLegalizeToTosa.pdll b/stablehlo/stablehlo/conversions/tosa/transforms/StablehloLegalizeToTosa.pdll
 --- stablehlo/stablehlo/conversions/tosa/transforms/StablehloLegalizeToTosa.pdll
 +++ stablehlo/stablehlo/conversions/tosa/transforms/StablehloLegalizeToTosa.pdll
diff --git a/third_party/xla/third_party/stablehlo/temporary.patch b/third_party/xla/third_party/stablehlo/temporary.patch
index 2e395d99d26cf4..f906a856ae69e1 100755
--- a/third_party/xla/third_party/stablehlo/temporary.patch
+++ b/third_party/xla/third_party/stablehlo/temporary.patch
@@ -361,6 +361,28 @@ diff --ruN a/stablehlo/stablehlo/conversions/tosa/tests/binary.mlir b/stablehlo/
    // CHECK: stablehlo.divide
    %0 = "stablehlo.divide"(%arg0, %arg1) : (tensor<10xf32>, tensor<10xf32>) -> tensor<10xf32>
    return %0 : tensor<10xf32>
+@@ -123,7 +123,7 @@
+ 
+ // CHECK-LABEL: @maximum_f64
+ func.func @maximum_f64(%arg0 : tensor<10xf64>, %arg1 : tensor<10xf64>) -> tensor<10xf64> {
+-  // CHECK: stablehlo.maximum
++  // CHECK: tosa.maximum
+   %0 = "stablehlo.maximum"(%arg0, %arg1) : (tensor<10xf64>, tensor<10xf64>) -> tensor<10xf64>
+   return %0 : tensor<10xf64>
+ }
+diff --ruN a/stablehlo/stablehlo/conversions/tosa/tests/nullary.mlir b/stablehlo/stablehlo/conversions/tosa/tests/nullary.mlir
+--- stablehlo/stablehlo/conversions/tosa/tests/nullary.mlir
++++ stablehlo/stablehlo/conversions/tosa/tests/nullary.mlir
+@@ -9,8 +9,7 @@
+ 
+ // CHECK-LABEL: @constant_f64
+ func.func @constant_f64() -> tensor<10xf64> {
+-  // TOSA does not support 64-bit types, so this should not legalize.
+-  // CHECK: stablehlo.constant
++  // CHECK: tosa.const
+   %0 = stablehlo.constant dense<0.000000e+00> : tensor<10xf64>
+   return %0 : tensor<10xf64>
+ }
 diff --ruN a/stablehlo/stablehlo/conversions/tosa/transforms/StablehloLegalizeToTosa.pdll b/stablehlo/stablehlo/conversions/tosa/transforms/StablehloLegalizeToTosa.pdll
 --- stablehlo/stablehlo/conversions/tosa/transforms/StablehloLegalizeToTosa.pdll
 +++ stablehlo/stablehlo/conversions/tosa/transforms/StablehloLegalizeToTosa.pdll

From 63bb7339a2bc98af97c61e86eedbbd60c123529c Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Tue, 14 May 2024 05:17:35 -0700
Subject: [PATCH 108/478] Update ops-related pbtxt files.

PiperOrigin-RevId: 633547952
---
 tensorflow/core/ops/ops.pbtxt | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorflow/core/ops/ops.pbtxt b/tensorflow/core/ops/ops.pbtxt
index 0556051577f5e4..7322cef85415f5 100644
--- a/tensorflow/core/ops/ops.pbtxt
+++ b/tensorflow/core/ops/ops.pbtxt
@@ -1,4 +1,4 @@
-go/debugstr  
+go/debugstr   
 op {
   name: "Abort"
   attr {

From e0e64ad7bbae544a07df6e52ca9794c4138b0554 Mon Sep 17 00:00:00 2001
From: Thomas Joerg <tjoerg@google.com>
Date: Tue, 14 May 2024 05:20:31 -0700
Subject: [PATCH 109/478] [XLA:GPU] Introduce `GpuCompiler::GetToolkit` method
 for subclasses to provide the CUDA or ROCm versions.

Remove all `#if GOOGLE_CUDA` and friends from gpu_compiler.cc

PiperOrigin-RevId: 633548576
---
 third_party/xla/xla/service/gpu/BUILD         | 10 ++++-----
 .../xla/xla/service/gpu/amdgpu_compiler.cc    |  8 +++++++
 .../xla/xla/service/gpu/amdgpu_compiler.h     |  2 ++
 .../xla/xla/service/gpu/gpu_compiler.cc       | 21 ++++---------------
 .../xla/xla/service/gpu/gpu_compiler.h        |  2 ++
 .../xla/xla/service/gpu/nvptx_compiler.cc     |  3 +++
 .../xla/xla/service/gpu/nvptx_compiler.h      |  2 ++
 .../xla/xla/tests/llvm_compiler_test.cc       |  4 ++++
 8 files changed, 29 insertions(+), 23 deletions(-)

diff --git a/third_party/xla/xla/service/gpu/BUILD b/third_party/xla/xla/service/gpu/BUILD
index e665640e8ce5eb..223aa418be67e7 100644
--- a/third_party/xla/xla/service/gpu/BUILD
+++ b/third_party/xla/xla/service/gpu/BUILD
@@ -3387,9 +3387,6 @@ cc_library(
     hdrs = if_gpu_is_configured([
         "gpu_compiler.h",
     ]),
-    local_defines = if_cuda_is_configured(["GOOGLE_CUDA=1"]) + if_rocm_is_configured([
-        "TENSORFLOW_USE_ROCM=1",
-    ]),
     deps = if_gpu_is_configured([
         ":gpu_p2p_pipeliner",
         ":pipelined_p2p_rewriter",
@@ -3460,7 +3457,6 @@ cc_library(
         "@com_google_absl//absl/strings",
         "@com_google_absl//absl/strings:str_format",
         "@com_google_absl//absl/types:variant",
-        "@local_config_cuda//cuda:cuda_headers",
         "@llvm-project//llvm:AsmParser",
         "@llvm-project//llvm:BitReader",
         "@llvm-project//llvm:BitWriter",
@@ -3579,7 +3575,6 @@ cc_library(
         "//xla/stream_executor:device_description",
         "//xla/stream_executor",
         "//xla/stream_executor/gpu:gpu_driver_header",
-        "//xla/stream_executor/cuda:cuda_platform_id",
         "//xla/stream_executor/integrations:device_mem_allocator",
         "//xla/translate/hlo_to_mhlo:hlo_utils",
         "//xla/translate/mhlo_to_hlo:location_exporter",
@@ -3610,7 +3605,6 @@ cc_library(
         "//xla/service:hlo_ordering",
         "//xla/service:layout_assignment",
         "//xla/service:logical_buffer",
-        "//xla/stream_executor/rocm:rocm_platform_id",
         "@local_tsl//tsl/platform:numbers",
     ]) + xla_internal(["service:export_hlo"]) + [
         ":command_buffer_scheduling",
@@ -3787,6 +3781,7 @@ cc_library(
         "@com_google_absl//absl/cleanup",
         "@com_google_absl//absl/container:flat_hash_map",
         "@com_google_absl//absl/container:node_hash_map",
+        "@local_config_cuda//cuda:cuda_headers",
         "@com_google_absl//absl/log",
         "@com_google_absl//absl/log:check",
         "@com_google_absl//absl/status",
@@ -3929,6 +3924,7 @@ cc_library(
     srcs = [
         "amdgpu_compiler_registration.cc",
     ],
+    local_defines = if_rocm_is_configured(["TENSORFLOW_USE_ROCM=1"]),
     tags = ["manual"],
     deps = [
         ":amdgpu_compiler_impl",
@@ -4017,6 +4013,8 @@ cc_library(
         "//xla/stream_executor:device_description",
         "//xla/stream_executor:stream_executor_headers",
         "//xla/stream_executor/rocm:rocm_platform_id",
+        "@com_google_absl//absl/log",
+        "@com_google_absl//absl/log:check",
         "@com_google_absl//absl/status",
         "@com_google_absl//absl/status:statusor",
         "@llvm-project//llvm:ir_headers",
diff --git a/third_party/xla/xla/service/gpu/amdgpu_compiler.cc b/third_party/xla/xla/service/gpu/amdgpu_compiler.cc
index 2a59d086d0223a..f8621e76610000 100644
--- a/third_party/xla/xla/service/gpu/amdgpu_compiler.cc
+++ b/third_party/xla/xla/service/gpu/amdgpu_compiler.cc
@@ -18,6 +18,7 @@ limitations under the License.
 #include <utility>
 #include <vector>
 
+#include "absl/log/log.h"
 #include "absl/status/status.h"
 #include "absl/status/statusor.h"
 #include "llvm/IR/Module.h"
@@ -90,6 +91,13 @@ struct ConvBfloat16Support : public FloatSupport {
 
 }  // namespace
 
+int32_t AMDGPUCompiler::GetToolkitVersion() const {
+#if TENSORFLOW_USE_ROCM
+  return TF_ROCM_VERSION;
+#endif
+  LOG(FATAL) << "Failed to get ROCm version.";
+}
+
 absl::Status AMDGPUCompiler::OptimizeHloConvolutionCanonicalization(
     HloModule* hlo_module, se::GpuComputeCapability gpu_version,
     se::dnn::VersionInfo dnn_version,
diff --git a/third_party/xla/xla/service/gpu/amdgpu_compiler.h b/third_party/xla/xla/service/gpu/amdgpu_compiler.h
index 9f4ad9f656256f..477ac1aeb1028d 100644
--- a/third_party/xla/xla/service/gpu/amdgpu_compiler.h
+++ b/third_party/xla/xla/service/gpu/amdgpu_compiler.h
@@ -39,6 +39,8 @@ class AMDGPUCompiler : public GpuCompiler {
  public:
   AMDGPUCompiler();
 
+  int32_t GetToolkitVersion() const override;
+
   absl::Status OptimizeHloConvolutionCanonicalization(
       HloModule* hlo_module, se::GpuComputeCapability gpu_version,
       se::dnn::VersionInfo dnn_version,
diff --git a/third_party/xla/xla/service/gpu/gpu_compiler.cc b/third_party/xla/xla/service/gpu/gpu_compiler.cc
index 6c3e2e70e05422..55edb1e4b00336 100644
--- a/third_party/xla/xla/service/gpu/gpu_compiler.cc
+++ b/third_party/xla/xla/service/gpu/gpu_compiler.cc
@@ -246,13 +246,6 @@ limitations under the License.
 #include "tsl/platform/threadpool.h"
 #include "tsl/profiler/lib/traceme.h"
 
-#if GOOGLE_CUDA
-#include "third_party/gpus/cuda/include/cuda.h"
-#include "xla/stream_executor/cuda/cuda_platform_id.h"
-#elif TENSORFLOW_USE_ROCM
-#include "xla/stream_executor/rocm/rocm_platform_id.h"
-#endif
-
 #ifdef PLATFORM_GOOGLE
 #include "xla/hlo/experimental/auto_sharding/auto_sharding.h"
 #endif  // PLATFORM_GOOGLE
@@ -2084,11 +2077,9 @@ absl::StatusOr<std::unique_ptr<Executable>> GpuCompiler::RunBackend(
 absl::StatusOr<std::vector<std::unique_ptr<AotCompilationResult>>>
 GpuCompiler::CompileAheadOfTime(std::unique_ptr<HloModuleGroup> module_group,
                                 const AotCompilationOptions& options) {
-#if GOOGLE_CUDA
-  CHECK(options.PlatformId() == se::cuda::kCudaPlatformId);
-#elif TENSORFLOW_USE_ROCM
-  CHECK(options.PlatformId() == se::rocm::kROCmPlatformId);
-#endif
+  // Check that we are on the platform (CUDA or ROCm) that was chosen for AOT
+  // compilation.
+  CHECK_EQ(options.PlatformId(), PlatformId());
 
   std::vector<std::unique_ptr<HloModule>> modules =
       module_group->ConsumeModules();
@@ -2223,11 +2214,7 @@ absl::Status GpuCompiler::RunPostSchedulingPipelines(
   {
     HloPassPipeline pipeline("command-buffer-scheduling");
     auto driver_version = se::gpu::GpuDriver::GetDriverVersion();
-#if GOOGLE_CUDA
-    constexpr int toolkit_version = CUDA_VERSION;
-#else
-    constexpr int toolkit_version = TF_ROCM_VERSION;
-#endif
+    const int32_t toolkit_version = GetToolkitVersion();
     pipeline.AddPass<CommandBufferScheduling>(
         gpu_device_info, toolkit_version,
         driver_version.value_or(toolkit_version));
diff --git a/third_party/xla/xla/service/gpu/gpu_compiler.h b/third_party/xla/xla/service/gpu/gpu_compiler.h
index 9d30a471deda8a..6ee9de66fc851f 100644
--- a/third_party/xla/xla/service/gpu/gpu_compiler.h
+++ b/third_party/xla/xla/service/gpu/gpu_compiler.h
@@ -116,6 +116,8 @@ class GpuCompiler : public LLVMCompiler {
     return &FusionCanShareBufferHint;
   }
 
+  virtual int32_t GetToolkitVersion() const = 0;
+
  protected:
   struct BackendCompileResult {
     std::string asm_text;
diff --git a/third_party/xla/xla/service/gpu/nvptx_compiler.cc b/third_party/xla/xla/service/gpu/nvptx_compiler.cc
index 6a1fa45926b64c..94b1314fd6cac5 100644
--- a/third_party/xla/xla/service/gpu/nvptx_compiler.cc
+++ b/third_party/xla/xla/service/gpu/nvptx_compiler.cc
@@ -36,6 +36,7 @@ limitations under the License.
 #include "absl/strings/str_format.h"
 #include "absl/strings/string_view.h"
 #include "absl/synchronization/mutex.h"
+#include "third_party/gpus/cuda/include/cuda.h"
 #include "llvm/IRReader/IRReader.h"
 #include "llvm/Support/SourceMgr.h"
 #include "llvm/Support/raw_ostream.h"
@@ -175,6 +176,8 @@ class MatmulBfloat16Support : public FloatSupport {
 
 }  // namespace
 
+int32_t NVPTXCompiler::GetToolkitVersion() const { return CUDA_VERSION; }
+
 absl::Status NVPTXCompiler::OptimizeHloConvolutionCanonicalization(
     HloModule* hlo_module, se::GpuComputeCapability gpu_version,
     se::dnn::VersionInfo dnn_version,
diff --git a/third_party/xla/xla/service/gpu/nvptx_compiler.h b/third_party/xla/xla/service/gpu/nvptx_compiler.h
index 3d7a770282b134..c25c850a915003 100644
--- a/third_party/xla/xla/service/gpu/nvptx_compiler.h
+++ b/third_party/xla/xla/service/gpu/nvptx_compiler.h
@@ -53,6 +53,8 @@ class NVPTXCompiler : public GpuCompiler {
  public:
   NVPTXCompiler();
 
+  int32_t GetToolkitVersion() const override;
+
   absl::Status OptimizeHloConvolutionCanonicalization(
       HloModule* hlo_module, se::GpuComputeCapability gpu_version,
       se::dnn::VersionInfo dnn_version,
diff --git a/third_party/xla/xla/tests/llvm_compiler_test.cc b/third_party/xla/xla/tests/llvm_compiler_test.cc
index 5f826254f2524b..32b2a5be85e3d1 100644
--- a/third_party/xla/xla/tests/llvm_compiler_test.cc
+++ b/third_party/xla/xla/tests/llvm_compiler_test.cc
@@ -15,6 +15,7 @@ limitations under the License.
 
 #include "xla/service/llvm_compiler.h"
 
+#include <cstdint>
 #include <memory>
 #include <string>
 #include <utility>
@@ -55,6 +56,9 @@ class GpuDummyCompiler : public GpuCompiler {
  public:
   GpuDummyCompiler()
       : GpuCompiler(kGpuPlatformId, kDummyTriple, kDummyLayout) {}
+
+  int32_t GetToolkitVersion() const override { return 0; }
+
   Status OptimizeHloConvolutionCanonicalization(
       HloModule* hlo_module, se::GpuComputeCapability gpu_version,
       se::dnn::VersionInfo dnn_version,

From 08a13d2991724a4b0039237ce07340e16b760656 Mon Sep 17 00:00:00 2001
From: Dmitri Gribenko <dmitrig@google.com>
Date: Tue, 14 May 2024 06:16:08 -0700
Subject: [PATCH 110/478] Integrate LLVM at llvm/llvm-project@e6d3a4212d20

Updates LLVM usage to match
[e6d3a4212d20](https://github.com/llvm/llvm-project/commit/e6d3a4212d20)

PiperOrigin-RevId: 633561628
---
 third_party/llvm/workspace.bzl | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/third_party/llvm/workspace.bzl b/third_party/llvm/workspace.bzl
index eb3a0f80219c81..8d2c0c7e5a2254 100644
--- a/third_party/llvm/workspace.bzl
+++ b/third_party/llvm/workspace.bzl
@@ -4,8 +4,8 @@ load("//third_party:repo.bzl", "tf_http_archive")
 
 def repo(name):
     """Imports LLVM."""
-    LLVM_COMMIT = "79a6a7e28fffd14e54a9a208af12d724b6eeb2d4"
-    LLVM_SHA256 = "ce0caa77cca929e29b0a6304820b833f4514a56bbfaaa533df846077916aa72d"
+    LLVM_COMMIT = "e6d3a4212d20b49a8e63f11fedea79cccf261479"
+    LLVM_SHA256 = "a9581601d91836d182180e35a698550b9c8257eacead9e5cc7ec956128200fce"
 
     tf_http_archive(
         name = name,

From 84fc9613ab556faeec2f3c7453e52f68cd083c51 Mon Sep 17 00:00:00 2001
From: Christian Sigg <csigg@google.com>
Date: Tue, 14 May 2024 06:33:36 -0700
Subject: [PATCH 111/478] Integrate Triton up to
 [@4caedc2](https://github.com/openai/triton/commits/4caedc2a51f5c81f7b63212b9bcb76b430102913)

PiperOrigin-RevId: 633565503
---
 third_party/triton/temporary/pipelining.patch | 472 ------------------
 third_party/triton/temporary/series.bzl       |   2 -
 .../triton/temporary/support_ceil_op.patch    | 138 -----
 third_party/triton/workspace.bzl              |   4 +-
 .../triton/xla_extensions/env_vars.patch      |  14 -
 third_party/triton/xla_extensions/series.bzl  |   1 -
 .../xla_extensions/sparse_dot_base.patch      |  99 ++--
 .../sparse_dot_fixes_y24w17.patch             |  58 +--
 .../xla_extensions/sparse_dot_nvgpu.patch     |  32 +-
 .../xla_extensions/sparse_dot_passes.patch    | 142 +++---
 .../triton/temporary/pipelining.patch         | 472 ------------------
 .../third_party/triton/temporary/series.bzl   |   2 -
 .../triton/temporary/support_ceil_op.patch    | 138 -----
 .../xla/third_party/triton/workspace.bzl      |   4 +-
 .../triton/xla_extensions/env_vars.patch      |  14 -
 .../triton/xla_extensions/series.bzl          |   1 -
 .../xla_extensions/sparse_dot_base.patch      |  99 ++--
 .../sparse_dot_fixes_y24w17.patch             |  58 +--
 .../xla_extensions/sparse_dot_nvgpu.patch     |  32 +-
 .../xla_extensions/sparse_dot_passes.patch    | 142 +++---
 .../xla/service/gpu/ir_emitter_triton_cuda.cc |   9 +-
 .../service/gpu/tests/sparse_add_layout.mlir  |   4 +-
 22 files changed, 346 insertions(+), 1591 deletions(-)
 delete mode 100644 third_party/triton/temporary/pipelining.patch
 delete mode 100644 third_party/triton/temporary/support_ceil_op.patch
 delete mode 100644 third_party/triton/xla_extensions/env_vars.patch
 delete mode 100644 third_party/xla/third_party/triton/temporary/pipelining.patch
 delete mode 100644 third_party/xla/third_party/triton/temporary/support_ceil_op.patch
 delete mode 100644 third_party/xla/third_party/triton/xla_extensions/env_vars.patch

diff --git a/third_party/triton/temporary/pipelining.patch b/third_party/triton/temporary/pipelining.patch
deleted file mode 100644
index 9f5f36aeb5099d..00000000000000
--- a/third_party/triton/temporary/pipelining.patch
+++ /dev/null
@@ -1,472 +0,0 @@
-This is patching changes upstream from different PRs that fix issues with
-pipelining internally. Required changes are upto and including this commit
-https://github.com/openai/triton/commit/70f0b7b6e333fe2155c79dfa8bec6ad388073670
-The patch can be removed with the integration that includes these changes.
-
-diff --git a/include/triton/Analysis/Utility.h b/include/triton/Analysis/Utility.h
---- a/include/triton/Analysis/Utility.h
-+++ b/include/triton/Analysis/Utility.h
-@@ -8,6 +8,18 @@
- 
- namespace mlir {
- 
-+inline bool isZeroConst(Value v) {
-+  auto constantOp = v.getDefiningOp<arith::ConstantOp>();
-+  if (!constantOp)
-+    return false;
-+  if (auto denseAttr = dyn_cast<DenseFPElementsAttr>(constantOp.getValueAttr()))
-+    return denseAttr.isSplat() && denseAttr.getSplatValue<APFloat>().isZero();
-+  if (auto denseAttr =
-+          dyn_cast<DenseIntElementsAttr>(constantOp.getValueAttr()))
-+    return denseAttr.isSplat() && denseAttr.getSplatValue<APInt>().isZero();
-+  return false;
-+}
-+
- class ReduceOpHelper {
- public:
-   explicit ReduceOpHelper(triton::ReduceOp op)
-diff --git a/include/triton/Dialect/TritonGPU/IR/TritonGPUOps.td b/include/triton/Dialect/TritonGPU/IR/TritonGPUOps.td
---- a/include/triton/Dialect/TritonGPU/IR/TritonGPUOps.td
-+++ b/include/triton/Dialect/TritonGPU/IR/TritonGPUOps.td
-@@ -45,6 +45,8 @@ def TTG_AsyncWaitOp : TTG_Op<"async_wait
- 
-   let arguments = (ins Variadic<TTG_AsyncToken>:$asyncToken, I32Attr:$num);
- 
-+  let results = (outs TTG_AsyncToken:$retToken);
-+
-   let assemblyFormat = "$asyncToken attr-dict";
- 
-   let extraClassDeclaration = [{
-@@ -229,10 +231,16 @@ def TTG_LocalLoadOp : TTG_Op<"local_load
-   let description = [{
-     Load a tensor from the local memory descriptor into a distributed tensor.
-   }];
--  let arguments = (ins TT_MemDescType:$src);
-+  let arguments = (ins TT_MemDescType:$src, Optional<TTG_AsyncToken> :$token);
-+
-+  let builders = [
-+      OpBuilder<(ins "Type":$retType, "Value":$src),
-+      [{
-+      build($_builder, $_state, retType, src, /*token=*/static_cast<mlir::Value>(nullptr));
-+      }]>];
- 
-   // Use qualified() otherwise "!tt.memdesc<X>" is printed as "<X>".
--  let assemblyFormat = [{$src attr-dict `:` qualified(type($src)) `->` type($result)}];
-+  let assemblyFormat = [{$src (`token` $token^)? attr-dict `:` qualified(type($src)) `->` type($result)}];
- 
-   let results = (outs TT_Tensor:$result);
- }
-diff --git a/lib/Dialect/TritonGPU/Transforms/Pipeliner/MatmulLoopPipeline.cpp b/lib/Dialect/TritonGPU/Transforms/Pipeliner/MatmulLoopPipeline.cpp
---- a/lib/Dialect/TritonGPU/Transforms/Pipeliner/MatmulLoopPipeline.cpp
-+++ b/lib/Dialect/TritonGPU/Transforms/Pipeliner/MatmulLoopPipeline.cpp
-@@ -8,6 +8,7 @@
- #include "mlir/Interfaces/SideEffectInterfaces.h"
- #include "mlir/Support/LLVM.h"
- #include "triton/Analysis/AxisInfo.h"
-+#include "triton/Analysis/Utility.h"
- #include "triton/Dialect/Triton/IR/Types.h"
- #include "triton/Dialect/Triton/IR/Utility.h"
- #include "triton/Dialect/TritonGPU/IR/Attributes.h"
-@@ -84,12 +85,13 @@ createAsyncCopy(scf::ForOp &forOp, tt::L
-   Location loc = loadOp.getLoc();
-   Value src = loadOp.getPtr();
-   Value mask = loadOp.getMask();
-+  Value other = loadOp.getOther();
-   if (!isExpensiveLoadOrStore(loadOp) && opToInfo[loadOp].blockedEncoding) {
-     // For inexpensive loads that do not directly feed into dot ops
-     // we want to use optimal layout for the data.
-     ttg::BlockedEncodingAttr encoding = opToInfo[loadOp].blockedEncoding;
-     auto convertBlockLayout = [&](Value src) {
--      auto ty = src.getType().cast<RankedTensorType>();
-+      auto ty = cast<RankedTensorType>(src.getType());
-       auto newTy =
-           RankedTensorType::get(ty.getShape(), ty.getElementType(), encoding);
-       auto cvt =
-@@ -99,9 +101,11 @@ createAsyncCopy(scf::ForOp &forOp, tt::L
-     src = convertBlockLayout(src);
-     if (mask)
-       mask = convertBlockLayout(mask);
-+    if (other)
-+      other = convertBlockLayout(other);
-   }
- 
--  tt::MemDescType allocTy = alloc.getType().cast<tt::MemDescType>();
-+  tt::MemDescType allocTy = cast<tt::MemDescType>(alloc.getType());
-   SmallVector<Value> copyOffsets(allocTy.getRank(), zero);
-   copyOffsets[0] = insertIdx;
-   tt::MemDescType subviewTy = tt::MemDescType::get(
-@@ -110,11 +114,12 @@ createAsyncCopy(scf::ForOp &forOp, tt::L
-   auto view =
-       builder.create<ttg::MemDescSubviewOp>(loc, subviewTy, alloc, copyOffsets);
-   Operation *copy = builder.create<ttg::AsyncCopyGlobalToLocalOp>(
--      loc, src, view, mask, loadOp.getOther(), loadOp.getCache(),
--      loadOp.getEvict(), loadOp.getIsVolatile());
-+      loc, src, view, mask, other, loadOp.getCache(), loadOp.getEvict(),
-+      loadOp.getIsVolatile());
-   Operation *commmit =
-       builder.create<ttg::AsyncCommitGroupOp>(loc, copy->getResult(0));
--  builder.create<ttg::AsyncWaitOp>(loc, commmit->getResult(0), 0);
-+  Operation *wait =
-+      builder.create<ttg::AsyncWaitOp>(loc, commmit->getResult(0), 0);
- 
-   int stage = opToInfo[loadOp].stage;
-   bool isMMV3Load = opToInfo[loadOp].loadIsMMAV3;
-@@ -142,9 +147,21 @@ createAsyncCopy(scf::ForOp &forOp, tt::L
-     for (auto alloc : allocsToErase) {
-       alloc.erase();
-     }
--    auto sharedLoad =
--        builder.create<ttg::LocalLoadOp>(loc, loadOp.getType(), viewLoad);
--    loadOp->replaceAllUsesWith(sharedLoad->getResults());
-+
-+    auto sharedLoad = builder.create<ttg::LocalLoadOp>(
-+        loc, loadOp.getType(), viewLoad, wait->getResult(0));
-+    auto result = sharedLoad->getResults();
-+
-+    // Create a select for non-zero other values as they are not handled by
-+    // AsyncCopyGlobalToLocalOp for now.
-+    Value other = loadOp.getOther();
-+    if (other && !isZeroConst(other)) {
-+      auto select = builder.create<arith::SelectOp>(
-+          loc, loadOp.getType(), mask, sharedLoad.getResult(), other);
-+      result = select->getResults();
-+    }
-+
-+    loadOp->replaceAllUsesWith(result);
-   }
-   loadOp.erase();
- }
-@@ -160,7 +177,7 @@ getSharedEncIfAllUsersAreDotEnc(Value va
-     if (user->getNumResults() != 1)
-       return std::nullopt;
-     if (auto memDesc =
--            user->getResult(0).getType().dyn_cast<triton::MemDescType>()) {
-+            dyn_cast<triton::MemDescType>(user->getResult(0).getType())) {
-       // First time we find a shared encoding in the chain, save it and try to
-       // use it if it is compatible with the other users.
-       tempAttr = memDesc.getEncoding().cast<ttg::SharedEncodingAttr>();
-@@ -203,7 +220,7 @@ getSharedEncIfAllUsersAreDotEnc(Value va
- static ttg::BlockedEncodingAttr
- getBlockedEncoding(tt::LoadOp loadOp, tt::ModuleAxisInfoAnalysis &axisInfo) {
-   Value src = loadOp.getPtr();
--  auto ty = src.getType().cast<RankedTensorType>();
-+  auto ty = cast<RankedTensorType>(src.getType());
-   auto mod = loadOp->getParentOfType<ModuleOp>();
-   int numWarps = ttg::TritonGPUDialect::getNumWarps(mod);
-   int threadsPerWarp = ttg::TritonGPUDialect::getThreadsPerWarp(mod);
-@@ -221,7 +238,7 @@ getBlockedEncoding(tt::LoadOp loadOp, tt
- 
- static std::optional<ttg::SharedEncodingAttr>
- getSharedEncoding(tt::LoadOp loadOp, bool isMMAV3) {
--  auto ty = loadOp.getType().cast<RankedTensorType>();
-+  auto ty = cast<RankedTensorType>(loadOp.getType());
-   auto ctaLayout = ttg::getCTALayout(ty.getEncoding());
-   auto blockedOrder = ttg::getOrder(ty.getEncoding());
-   SmallVector<unsigned> order;
-@@ -285,11 +302,10 @@ loadOpsToDistanceAndUse(scf::ForOp forOp
-     if (auto mask = loadOp.getMask())
-       vec = std::min<unsigned>(vec, axisInfoAnalysis.getMaskAlignment(mask));
- 
--    auto tensorTy = ptr.getType().dyn_cast<RankedTensorType>();
-+    auto tensorTy = dyn_cast<RankedTensorType>(ptr.getType());
-     if (!tensorTy)
-       return false;
--    auto ty =
--        tensorTy.getElementType().cast<tt::PointerType>().getPointeeType();
-+    auto ty = cast<tt::PointerType>(tensorTy.getElementType()).getPointeeType();
-     unsigned width = vec * ty.getIntOrFloatBitWidth();
- 
-     // We do not pipeline all loads for the following reasons:
-@@ -353,7 +369,7 @@ static bool loadIsMMAv3(tt::LoadOp loadO
- 
-   // MMA V3 case.
-   auto newOrder = sharedEnc.getOrder();
--  auto ty = loadOp.getType().cast<RankedTensorType>();
-+  auto ty = cast<RankedTensorType>(loadOp.getType());
-   auto oldOrder = ttg::getOrder(ty.getEncoding());
- 
-   // The operand of MMAv3 is in SharedEncoding and its order should not
-@@ -497,7 +513,7 @@ collectOpsToPipeline(scf::ForOp forOp,
- static Value createAlloc(scf::ForOp &forOp, tt::LoadOp loadOp,
-                          ttg::SharedEncodingAttr sharedEnc, unsigned distance) {
-   OpBuilder builder(forOp);
--  auto ty = loadOp.getType().cast<RankedTensorType>();
-+  auto ty = cast<RankedTensorType>(loadOp.getType());
-   SmallVector<int64_t> bufferShape(ty.getShape().begin(), ty.getShape().end());
-   bufferShape.insert(bufferShape.begin(), distance);
-   Type memdescType = mlir::triton::MemDescType::get(
-@@ -669,12 +685,23 @@ createSchedule(scf::ForOp forOp, int num
-     }
-   });
- 
-+  auto getNestedOperands = [](Operation *op) -> SmallVector<Value> {
-+    SmallVector<Value> operands;
-+    op->walk([&](Operation *nestedOp) {
-+      for (Value operand : nestedOp->getOperands()) {
-+        if (operand.getParentBlock()->getParentOp()->isAncestor(nestedOp))
-+          operands.push_back(operand);
-+      }
-+    });
-+    return operands;
-+  };
-+
-   // Find dependencies with distance of 1.
-   SmallVector<DenseSet<Operation *>> distanceOneUsers(numStages);
-   for (int stage = 0; stage < numStages - 1; stage++) {
-     auto &group = insertAndDeps[stage];
-     for (Operation *op : group) {
--      for (Value operand : op->getOperands()) {
-+      for (Value operand : getNestedOperands(op)) {
-         if (auto arg = operand.dyn_cast<BlockArgument>()) {
-           if (arg.getArgNumber() > 0 && arg.getOwner() == op->getBlock()) {
-             auto yieldOp = op->getBlock()->getTerminator();
-@@ -905,7 +932,7 @@ static int minNumInterleavedCommitOps(Op
- // Look for consecutive wait ops and combine them into a single wait op.
- static void
- combineRedundantWaitOps(llvm::SmallSetVector<ttg::AsyncWaitOp, 8> &waitOps) {
--  llvm::SmallSetVector<ttg::AsyncWaitOp, 8> toDelete;
-+  llvm::MapVector<ttg::AsyncWaitOp, ttg::AsyncWaitOp> toDelete;
-   for (auto waitOp : waitOps) {
-     if (toDelete.count(waitOp))
-       continue;
-@@ -927,10 +954,13 @@ combineRedundantWaitOps(llvm::SmallSetVe
-     OpBuilder builder(waitGroup.back());
-     auto newWaitOp = builder.create<ttg::AsyncWaitOp>(waitOp.getLoc(),
-                                                       depTokens, minWaitNumber);
--    toDelete.insert(waitGroup.begin(), waitGroup.end());
-+    for (auto waitOp : waitGroup) {
-+      toDelete[waitOp] = newWaitOp;
-+    }
-   }
-   for (auto waitOp : toDelete) {
--    waitOp->erase();
-+    waitOp.first->replaceAllUsesWith(waitOp.second);
-+    waitOp.first->erase();
-   }
- }
- 
-@@ -1010,7 +1040,7 @@ static void threadValuesThroughWait(ttng
- 
-   for (ttng::DotAsyncOp dot : asyncDots) {
-     for (Value operand : dot.getOperands()) {
--      if (operand.getType().isa<tt::MemDescType>()) {
-+      if (isa<tt::MemDescType>(operand.getType())) {
-         newOperands.insert(operand);
-       }
-     }
-@@ -1020,15 +1050,21 @@ static void threadValuesThroughWait(ttng
-   // values in the operation.
-   auto newWait = builder.create<ttng::DotWaitOp>(
-       wait.getLoc(), llvm::to_vector(newOperands), wait.getPendings());
-+
-+  auto dominatedByNewWait = [&](OpOperand &operand) {
-+    auto opInThisBlock =
-+        newWait->getBlock()->findAncestorOpInBlock(*operand.getOwner());
-+    return opInThisBlock && newWait->isBeforeInBlock(opInThisBlock);
-+  };
-   for (int i = 0; i < origNumOperands; i++) {
-     Value operand = wait.getResult(i);
--    if (!operand.getType().isa<tt::MemDescType>())
-+    if (!isa<tt::MemDescType>(operand.getType()))
-       operand.replaceAllUsesWith(newWait.getResult(i));
-   }
-   for (int i = origNumOperands; i < newOperands.size(); i++) {
-     Value operand = newWait.getOperand(i);
--    if (!operand.getType().isa<tt::MemDescType>())
--      operand.replaceAllUsesExcept(newWait.getResult(i), newWait);
-+    if (!isa<tt::MemDescType>(operand.getType()))
-+      operand.replaceUsesWithIf(newWait.getResult(i), dominatedByNewWait);
-   }
-   wait->erase();
- }
-@@ -1047,8 +1083,8 @@ static void threadValuesThroughWait(ttng
- //  1. All operands that touch shared memory are multi-buffered, i.e. can't read
- //     an incomplete value while it's being written asynchronously by a load.
- //
--//  2. During iteration i, nothing other than the loop's `yield` reads the
--//     result of the dot.
-+//  2. If the dot is used by any op in the loop, it must be used under an `if`,
-+//     and will be synced with a `wait 0` at the beginning of the `if` block.
- //
- //  3. During iteration i, between the start of the loop up until the first
- //     `ttng.dot_wait {pendings=0}` op, the result of the dot from iteration i-1
-@@ -1079,7 +1115,7 @@ static std::optional<int> dotCanBeProper
-   // Rule 1: All shmem operands are multi-buffered.
-   auto checkOperand = [&](Value operand) {
-     if (!isa<ttg::SharedEncodingAttr>(
--            operand.getType().cast<TensorOrMemDesc>().getEncoding())) {
-+            cast<TensorOrMemDesc>(operand.getType()).getEncoding())) {
-       return true;
-     }
- 
-@@ -1103,17 +1139,41 @@ static std::optional<int> dotCanBeProper
-     return std::nullopt;
-   }
- 
--  // Rule 2: The dot should only be used by the for loop's `yield`.
--  if (!dotOp->hasOneUse() ||
--      *dotOp->getUsers().begin() != forOp.getBody()->getTerminator()) {
--    LDBG("Can't make dot async because it is not used only by the loop's "
--         "`yield`.");
--    return std::nullopt;
-+  // Rule 2: The dot cannot be unconditionally used by any op in the loop.
-+  // Uses under `if` are allowed, as can be explicitly synced with a `wait 0`.
-+  int iterArgIdx = -1;
-+  Value iterArg = nullptr;
-+  SmallVector<std::pair<Operation *, int>> queue;
-+  for (auto &use : dotOp->getUses()) {
-+    queue.push_back({use.getOwner(), use.getOperandNumber()});
-   }
--
--  // The result of the dot becomes this loop carry value.
--  auto iterArgIdx = dotOp->getUses().begin()->getOperandNumber();
--  auto iterArg = forOp.getRegionIterArg(iterArgIdx);
-+  while (!queue.empty()) {
-+    auto [user, argIdx] = queue.pop_back_val();
-+    if (user->getParentOp() == forOp) {
-+      if (isa<scf::YieldOp>(user)) {
-+        if (iterArg) {
-+          // The dot is used by the loop's yield, but we can't have any other
-+          // uses.
-+          return std::nullopt;
-+        }
-+        iterArgIdx = argIdx;
-+        iterArg = forOp.getRegionIterArg(argIdx);
-+        continue;
-+      }
-+      return std::nullopt;
-+    }
-+    if (auto ifOp = dyn_cast<scf::IfOp>(user->getParentOp())) {
-+      if (isa<scf::YieldOp>(user)) {
-+        // The result is returned by the if, follow it further.
-+        auto uses = ifOp.getResult(argIdx).getUses();
-+        for (auto &use : uses) {
-+          queue.push_back({use.getOwner(), use.getOperandNumber()});
-+        }
-+      }
-+    } else {
-+      return std::nullopt;
-+    }
-+  }
- 
-   // Rule 3a: Are the only users of the dot's result from iteration i-1 other
-   // MMAv3 dots?  If so, we're done, this dot can be properly async.
-@@ -1181,6 +1241,32 @@ static void insertAsyncDotWaitInLoop(
-     return;
-   }
- 
-+  // Insert waits before the users of the properly async dots other than loop
-+  // yield.
-+  for (auto [asyncDot, iterArgIdx] : properlyAsyncDots) {
-+    SmallVector<OpOperand *> uses;
-+    for (auto &use : asyncDot->getUses()) {
-+      if (auto yieldOp = dyn_cast<scf::YieldOp>(use.getOwner())) {
-+        continue;
-+      }
-+      uses.push_back(&use);
-+    }
-+
-+    DenseMap<Block *, SmallVector<Value>> blockToUsers;
-+    for (auto use : uses) {
-+      auto block = use->getOwner()->getBlock();
-+      blockToUsers[block].push_back(use->get());
-+    }
-+
-+    for (auto [block, users] : blockToUsers) {
-+      OpBuilder builder(block, block->begin());
-+      auto newWait = builder.create<ttng::DotWaitOp>(asyncDot->getLoc(),
-+                                                     ArrayRef<Value>{}, 0);
-+
-+      threadValuesThroughWait(newWait, users);
-+    }
-+  }
-+
-   // Add the wait right after the last properly-async dot.  This only needs to
-   // wait for all properly-async dots from the i-1'th iteration to complete, IOW
-   // we wait until there are most `asyncDots.size()` dots in flight.
-diff --git a/test/TritonGPU/loop-pipeline.mlir b/test/TritonGPU/loop-pipeline.mlir
---- a/test/TritonGPU/loop-pipeline.mlir
-+++ b/test/TritonGPU/loop-pipeline.mlir
-@@ -349,16 +349,21 @@ tt.func @indirect_bmm_scalar_dist_one(%7
- // CHECK: triton_gpu.async_copy_global_to_local
- // CHECK: triton_gpu.async_copy_global_to_local
- // CHECK: triton_gpu.async_commit_group
-+// CHECK: triton_gpu.async_wait {{.*}} {num = 1 : i32}
-+// CHECK: scf.for
-+// CHECK: tt.dot
- // CHECK: %[[NEXT_BUFFER_1:.*]] = tt.addptr %{{.*}}, {{.*}}
- // CHECK: triton_gpu.async_copy_global_to_local %[[NEXT_BUFFER_1]]
--// CHECK: %[[IND_BUFFER_0:.*]] = triton_gpu.memdesc_subview
--// CHECK: %[[IND_BUFFER_1:.*]] = triton_gpu.local_load %[[IND_BUFFER_0]]
-+// CHECK-DAG: %[[IND_BUFFER_WAIT_TOKEN:.*]] = triton_gpu.async_wait {{.*}} {num = 1 : i32}
-+// CHECK-DAG: %[[IND_BUFFER_0:.*]] = triton_gpu.memdesc_subview
-+// CHECK: %[[IND_BUFFER_1:.*]] = triton_gpu.local_load %[[IND_BUFFER_0]] token %[[IND_BUFFER_WAIT_TOKEN]]
- // CHECK: %[[IND_BUFFER_2:.*]] = tt.expand_dims %[[IND_BUFFER_1]] {axis = 1 : i32}
- // CHECK: %[[IND_BUFFER_3:.*]] = tt.broadcast %[[IND_BUFFER_2]]
- // CHECK: %[[IND_BUFFER_4:.*]] = arith.muli {{.*}}, %[[IND_BUFFER_3]]
- // CHECK: %[[NEXT_BUFFER_0:.*]] = tt.addptr {{.*}}, %[[IND_BUFFER_4]]
- // CHECK: triton_gpu.async_copy_global_to_local %[[NEXT_BUFFER_0]]
- // CHECK: triton_gpu.async_wait {{.*}} {num = 1 : i32}
-+// CHECK: scf.yield
- tt.func @indirect_bmm_vector(%77: tensor<16x16xi64, #BL> {tt.divisibility=16: i32, tt.constancy=16: i32},
-                    %76: index,
-                    %49: tensor<16x16x!tt.ptr<f16>, #AL> {tt.divisibility=16: i32, tt.contiguity=2 : i32},
-diff --git a/test/TritonGPU/reorder-instructions.mlir b/test/TritonGPU/reorder-instructions.mlir
---- a/test/TritonGPU/reorder-instructions.mlir
-+++ b/test/TritonGPU/reorder-instructions.mlir
-@@ -28,7 +28,7 @@ module attributes {"triton_gpu.num-warps
- //       CHECK: triton_gpu.async_wait {num = 0 : i32}
- //       CHECK: triton_gpu.local_dealloc %0 : !tt.memdesc<4x128x64xf16, #shared>
- //       CHECK: triton_gpu.local_dealloc %1 : !tt.memdesc<4x128x64xf16, #shared>
--//       CHECK: %2 = triton_gpu.convert_layout %arg0 : tensor<32x32xf32, #blocked> -> tensor<32x32xf32, #blocked1>
-+//       CHECK: %3 = triton_gpu.convert_layout %arg0 : tensor<32x32xf32, #blocked> -> tensor<32x32xf32, #blocked1>
- #blocked = #triton_gpu.blocked<{sizePerThread = [1, 1], threadsPerWarp = [32, 1], warpsPerCTA = [1, 4], order = [0, 1]}>
- #blocked1 = #triton_gpu.blocked<{sizePerThread = [1, 1], threadsPerWarp = [32, 1], warpsPerCTA = [1, 4], order = [1, 0]}>
- #shared = #triton_gpu.shared<{vec = 8, perPhase = 1, maxPhase = 4, order = [0, 1]}>
-diff --git a/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/DotOpToLLVM/WGMMA.cpp b/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/DotOpToLLVM/WGMMA.cpp
---- a/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/DotOpToLLVM/WGMMA.cpp
-+++ b/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/DotOpToLLVM/WGMMA.cpp
-@@ -333,17 +333,6 @@ static Value faddAccumulate(ConversionPa
-   return newStruct;
- }
- 
--static bool isZero(Value v) {
--  auto constantOp = v.getDefiningOp<arith::ConstantOp>();
--  if (!constantOp)
--    return false;
--  if (auto denseAttr = dyn_cast<DenseFPElementsAttr>(constantOp.getValueAttr()))
--    return denseAttr.isSplat() && denseAttr.getSplatValue<APFloat>().isZero();
--  if (auto denseAttr =
--          dyn_cast<DenseIntElementsAttr>(constantOp.getValueAttr()))
--    return denseAttr.isSplat() && denseAttr.getSplatValue<APInt>().isZero();
--  return false;
--}
- 
- static SmallVector<Value> emitWait(ConversionPatternRewriter &rewriter,
-                                    Location loc, SmallVector<Value> acc,
-@@ -402,7 +391,7 @@ LogicalResult convertDot(const LLVMTypeC
-   int M = 4 * instrShape[0];
-   int N = instrShape[1];
-   int K = instrShape[2];
--  bool zeroAcc = isZero(c);
-+  bool zeroAcc = isZeroConst(c);
-   auto shapePerCTATile = getShapePerCTATile(mmaEncoding);
-   int numRepM = ceil<unsigned>(dShapePerCTA[0], shapePerCTATile[0]);
-   int numRepN = ceil<unsigned>(dShapePerCTA[1], shapePerCTATile[1]);
-diff --git a/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/LoadStoreOpToLLVM.cpp b/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/LoadStoreOpToLLVM.cpp
---- a/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/LoadStoreOpToLLVM.cpp
-+++ b/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/LoadStoreOpToLLVM.cpp
-@@ -924,8 +924,11 @@ struct AsyncWaitOpConversion
-     auto voidTy = void_ty(ctx);
-     ptxBuilder.launch(rewriter, loc, voidTy);
- 
--    // Safe to remove the op since it doesn't have any return value.
--    rewriter.eraseOp(op);
-+    // Drop the result token.
-+    Value zero = rewriter.create<LLVM::ConstantOp>(
-+        op.getLoc(), IntegerType::get(op.getContext(), 32),
-+        rewriter.getI32IntegerAttr(0));
-+    rewriter.replaceOp(op, zero);
-     return success();
-   }
- };
diff --git a/third_party/triton/temporary/series.bzl b/third_party/triton/temporary/series.bzl
index 70313c9b436d2a..214666767949de 100644
--- a/third_party/triton/temporary/series.bzl
+++ b/third_party/triton/temporary/series.bzl
@@ -6,7 +6,5 @@ internal patch during the next triton integration process.
 """
 
 temporary_patch_list = [
-    "//third_party/triton/temporary:pipelining.patch",
-    "//third_party/triton/temporary:support_ceil_op.patch",
     "//third_party/triton/temporary:mma_limit_pred.patch",
 ]
diff --git a/third_party/triton/temporary/support_ceil_op.patch b/third_party/triton/temporary/support_ceil_op.patch
deleted file mode 100644
index 71b323d9fccdca..00000000000000
--- a/third_party/triton/temporary/support_ceil_op.patch
+++ /dev/null
@@ -1,138 +0,0 @@
-Cherry-picking https://github.com/openai/triton/commit/62706e8c518c8c56e56460a43732d8e375217860
-until the next integration lands it. Can be removed as it is already merged.
-
-diff --git a/lib/Conversion/TritonGPUToLLVM/ElementwiseOpToLLVM.cpp b/lib/Conversion/TritonGPUToLLVM/ElementwiseOpToLLVM.cpp
---- a/lib/Conversion/TritonGPUToLLVM/ElementwiseOpToLLVM.cpp
-+++ b/lib/Conversion/TritonGPUToLLVM/ElementwiseOpToLLVM.cpp
-@@ -805,6 +805,7 @@ void mlir::triton::populateElementwiseOp
-   POPULATE_UNARY_OP(arith::FPToUIOp, LLVM::FPToUIOp)
-   POPULATE_UNARY_OP(arith::UIToFPOp, LLVM::UIToFPOp)
-   POPULATE_UNARY_OP(math::FloorOp, math::FloorOp)
-+  POPULATE_UNARY_OP(math::CeilOp, math::CeilOp)
-   POPULATE_UNARY_OP(math::LogOp, math::LogOp)
-   POPULATE_UNARY_OP(math::Log2Op, math::Log2Op)
-   POPULATE_UNARY_OP(math::CosOp, math::CosOp)
-diff --git a/lib/Conversion/TritonToTritonGPU/TritonToTritonGPUPass.cpp b/lib/Conversion/TritonToTritonGPU/TritonToTritonGPUPass.cpp
---- a/lib/Conversion/TritonToTritonGPU/TritonToTritonGPUPass.cpp
-+++ b/lib/Conversion/TritonToTritonGPU/TritonToTritonGPUPass.cpp
-@@ -125,12 +125,13 @@ void populateMathPatternsAndLegality(Tri
-   MLIRContext *context = patterns.getContext();
-   // Rewrite rule
-   patterns.add<GenericOpPattern<math::ExpOp>, GenericOpPattern<math::Exp2Op>,
--               GenericOpPattern<math::FloorOp>, GenericOpPattern<math::CosOp>,
--               GenericOpPattern<math::SinOp>, GenericOpPattern<math::LogOp>,
--               GenericOpPattern<math::Log2Op>, GenericOpPattern<math::ErfOp>,
--               GenericOpPattern<math::AbsFOp>, GenericOpPattern<math::AbsIOp>,
--               GenericOpPattern<math::SqrtOp>, GenericOpPattern<math::RsqrtOp>,
--               GenericOpPattern<math::FmaOp>>(typeConverter, context);
-+               GenericOpPattern<math::FloorOp>, GenericOpPattern<math::CeilOp>,
-+               GenericOpPattern<math::CosOp>, GenericOpPattern<math::SinOp>,
-+               GenericOpPattern<math::LogOp>, GenericOpPattern<math::Log2Op>,
-+               GenericOpPattern<math::ErfOp>, GenericOpPattern<math::AbsFOp>,
-+               GenericOpPattern<math::AbsIOp>, GenericOpPattern<math::SqrtOp>,
-+               GenericOpPattern<math::RsqrtOp>, GenericOpPattern<math::FmaOp>>(
-+      typeConverter, context);
- }
- 
- //
-diff --git a/lib/Dialect/TritonNvidiaGPU/Transforms/PlanCTA.cpp b/lib/Dialect/TritonNvidiaGPU/Transforms/PlanCTA.cpp
---- a/lib/Dialect/TritonNvidiaGPU/Transforms/PlanCTA.cpp
-+++ b/lib/Dialect/TritonNvidiaGPU/Transforms/PlanCTA.cpp
-@@ -651,10 +651,9 @@ bool CTAPlanner::isElementwiseOp(Operati
-                 math::CeilOp, math::CopySignOp, math::CosOp, math::SinOp,
-                 math::CountLeadingZerosOp, math::CountTrailingZerosOp,
-                 math::CtPopOp, math::ErfOp, math::ExpOp, math::Exp2Op,
--                math::FloorOp, math::ExpM1Op, math::FloorOp, math::FmaOp,
--                math::LogOp, math::Log10Op, math::Log1pOp, math::Log2Op,
--                math::PowFOp, math::RsqrtOp, math::SqrtOp, math::RsqrtOp,
--                math::TanhOp>(op))
-+                math::FloorOp, math::ExpM1Op, math::FmaOp, math::LogOp,
-+                math::Log10Op, math::Log1pOp, math::Log2Op, math::PowFOp,
-+                math::RsqrtOp, math::SqrtOp, math::RsqrtOp, math::TanhOp>(op))
-     return true;
-   if (llvm::isa<triton::IntToPtrOp, triton::PtrToIntOp, triton::BitcastOp,
-                 triton::FpToFpOp, triton::AddPtrOp, triton::PreciseSqrtOp,
-diff --git a/python/src/ir.cc b/python/src/ir.cc
---- a/python/src/ir.cc
-+++ b/python/src/ir.cc
-@@ -1379,6 +1379,10 @@ void init_triton_ir(py::module &&m) {
-            [](TritonOpBuilder &self, Value &val) -> Value {
-              return self.create<math::FloorOp>(val);
-            })
-+      .def("create_ceil",
-+           [](TritonOpBuilder &self, Value &val) -> Value {
-+             return self.create<math::CeilOp>(val);
-+           })
-       .def("create_exp",
-            [](TritonOpBuilder &self, Value &val) -> Value {
-              return self.create<math::ExpOp>(val);
-diff --git a/python/test/unit/language/test_core.py b/python/test/unit/language/test_core.py
---- a/python/test/unit/language/test_core.py
-+++ b/python/test/unit/language/test_core.py
-@@ -915,10 +915,11 @@ def test_unary_op(dtype_x, expr, num_cta
- 
- 
- @pytest.mark.interpreter
--@pytest.mark.parametrize("dtype_x, expr, x", [(dtype_x, expr, x)
--                                              for dtype_x in ["float32", "float64"]
--                                              for expr in ['exp', 'log', 'cos', 'sin', 'exp2', 'log2', 'sqrt', 'floor']
--                                              for x in ['x', '3.0']])
-+@pytest.mark.parametrize("dtype_x, expr, x",
-+                         [(dtype_x, expr, x)
-+                          for dtype_x in ["float32", "float64"]
-+                          for expr in ['exp', 'log', 'cos', 'sin', 'exp2', 'log2', 'sqrt', 'floor', 'ceil']
-+                          for x in ['x', '3.0']])
- def test_math_op(dtype_x, expr, x, device):
-     _test_unary(dtype_x, f'tl.{expr}({x})', f'np.{expr}({x}) ', device=device)
- 
-diff --git a/python/triton/language/__init__.py b/python/triton/language/__init__.py
---- a/python/triton/language/__init__.py
-+++ b/python/triton/language/__init__.py
-@@ -102,7 +102,8 @@ from .core import (
-     void,
-     where,
- )
--from .math import (umulhi, exp, exp2, fma, log, log2, cos, rsqrt, sin, sqrt, sqrt_rn, abs, fdiv, div_rn, erf, floor)
-+from .math import (umulhi, exp, exp2, fma, log, log2, cos, rsqrt, sin, sqrt, sqrt_rn, abs, fdiv, div_rn, erf, floor,
-+                   ceil)
- from .random import (
-     pair_uniform_to_normal,
-     philox,
-@@ -142,6 +143,7 @@ from .random import (
-     "builtin",
-     "cat",
-     "cdiv",
-+    "ceil",
-     "clamp",
-     "const",
-     "const_pointer_type",
-diff --git a/python/triton/language/math.py b/python/triton/language/math.py
---- a/python/triton/language/math.py
-+++ b/python/triton/language/math.py
-@@ -230,6 +230,15 @@ def floor(x, _builder=None):
- 
- 
- @core.builtin
-+@_check_dtype(dtypes=["fp32", "fp64"])
-+@_add_math_1arg_docstr("ceil")
-+@core._tensor_member_fn
-+def ceil(x, _builder=None):
-+    x = core._to_tensor(x, _builder)
-+    return core.tensor(_builder.create_ceil(x.handle), x.type)
-+
-+
-+@core.builtin
- @_add_math_3arg_docstr("fused multiply-add")
- def fma(x, y, z, _builder=None):
-     x = core._to_tensor(x, _builder)
-diff --git a/python/triton/runtime/interpreter.py b/python/triton/runtime/interpreter.py
---- a/python/triton/runtime/interpreter.py
-+++ b/python/triton/runtime/interpreter.py
-@@ -391,6 +391,7 @@ class InterpreterBuilder:
-     create_fabs = lambda self, arg: self.unary_op(arg, np.abs)
-     create_iabs = lambda self, arg: self.unary_op(arg, np.abs)
-     create_floor = lambda self, arg: self.unary_op(arg, np.floor)
-+    create_ceil = lambda self, arg: self.unary_op(arg, np.ceil)
-     create_log = lambda self, arg: self.unary_op(arg, np.log)
-     create_log2 = lambda self, arg: self.unary_op(arg, np.log2)
-     create_precise_sqrt = lambda self, arg: self.unary_op(arg, np.sqrt)
diff --git a/third_party/triton/workspace.bzl b/third_party/triton/workspace.bzl
index 45daf7974a022e..933b815253ffcd 100644
--- a/third_party/triton/workspace.bzl
+++ b/third_party/triton/workspace.bzl
@@ -8,8 +8,8 @@ load("//third_party/triton/xla_extensions:series.bzl", "extensions_files_patch_l
 def repo():
     """Imports Triton."""
 
-    TRITON_COMMIT = "cl623533461"
-    TRITON_SHA256 = "7aa74e82e4417a91fc7a7a84b4f6ad2b7e4e58512758d6c78ca3cd1c8771326b"
+    TRITON_COMMIT = "cl632952902"
+    TRITON_SHA256 = "f83c4f197cc2ae3b034070ec8189007451093edc445d1e383eb17a1e3808de9a"
     tf_http_archive(
         name = "triton",
         sha256 = TRITON_SHA256,
diff --git a/third_party/triton/xla_extensions/env_vars.patch b/third_party/triton/xla_extensions/env_vars.patch
deleted file mode 100644
index 955eb6db8da68e..00000000000000
--- a/third_party/triton/xla_extensions/env_vars.patch
+++ /dev/null
@@ -1,14 +0,0 @@
-Long standing patch due to licensing issues.
-diff --git a/include/triton/Tools/Sys/GetEnv.hpp b/include/triton/Tools/Sys/GetEnv.hpp
-index 31bc03fe1..a19a432df 100644
---- a/include/triton/Tools/Sys/GetEnv.hpp
-+++ b/include/triton/Tools/Sys/GetEnv.hpp
-@@ -34,7 +34,7 @@ inline const std::set<std::string> ENV_VARS = {
-     "AMDGCN_ENABLE_DUMP",
-     "DISABLE_FAST_REDUCTION",
-     "DISABLE_LLVM_OPT",
--    "DISABLE_MMA_V3",
-+    "ENABLE_MMA_V3",
-     "DISABLE_PTXAS_OPT",
-     "LLVM_IR_ENABLE_DUMP",
-     "MLIR_ENABLE_DUMP",
diff --git a/third_party/triton/xla_extensions/series.bzl b/third_party/triton/xla_extensions/series.bzl
index b858da203fb094..af524fb253cbef 100644
--- a/third_party/triton/xla_extensions/series.bzl
+++ b/third_party/triton/xla_extensions/series.bzl
@@ -4,7 +4,6 @@ applied in the previous copybara workflow.
 """
 
 extensions_files_patch_list = [
-    "//third_party/triton/xla_extensions:env_vars.patch",  # File not exported to google
     "//third_party/triton/xla_extensions:sparse_dot_nvgpu.patch",  # Sparsity internal patch
     "//third_party/triton/xla_extensions:sparse_dot_base.patch",  # Sparsity internal patch
     "//third_party/triton/xla_extensions:sparse_dot_passes.patch",  # Sparsity internal patch
diff --git a/third_party/triton/xla_extensions/sparse_dot_base.patch b/third_party/triton/xla_extensions/sparse_dot_base.patch
index dcacd99740b18f..08b7dd6f7ada87 100644
--- a/third_party/triton/xla_extensions/sparse_dot_base.patch
+++ b/third_party/triton/xla_extensions/sparse_dot_base.patch
@@ -1,8 +1,9 @@
 diff --git a/include/triton/Dialect/TritonGPU/IR/TritonGPUAttrDefs.td b/include/triton/Dialect/TritonGPU/IR/TritonGPUAttrDefs.td
+index 56f0b6b49..aa91ea9b8 100644
 --- a/include/triton/Dialect/TritonGPU/IR/TritonGPUAttrDefs.td
 +++ b/include/triton/Dialect/TritonGPU/IR/TritonGPUAttrDefs.td
-@@ -1158,4 +1158,12 @@ section 9.7.13.4.1 for more details.
-   let extraClassDeclaration = extraDistributedDeclaration;
+@@ -1262,4 +1262,16 @@ section 9.7.13.4.1 for more details.
+   }];
  }
  
 +def SparseDotMetaEncodingAttr : DistributedEncoding<"SparseDotMetaEncoding", "sparse_dot_meta_encoding"> {
@@ -10,14 +11,19 @@ diff --git a/include/triton/Dialect/TritonGPU/IR/TritonGPUAttrDefs.td b/include/
 +
 +  let parameters = (ins "Attribute":$parent);
 +  let assemblyFormat = "`<``{` struct(params) `}``>`";
-+  let extraClassDeclaration = extraDistributedDeclaration;
++  let extraClassDeclaration = extraDistributedDeclaration # [{
++    SmallVector<unsigned> getContigPerThread() {
++      return getSizePerThread();
++    };
++  }];
 +}
 +
  #endif
 diff --git a/include/triton/Dialect/TritonGPU/IR/TritonGPUOps.td b/include/triton/Dialect/TritonGPU/IR/TritonGPUOps.td
+index 4966a5f73..d2bb33cfa 100644
 --- a/include/triton/Dialect/TritonGPU/IR/TritonGPUOps.td
 +++ b/include/triton/Dialect/TritonGPU/IR/TritonGPUOps.td
-@@ -7,6 +7,7 @@ include "triton/Dialect/TritonGPU/IR/Tri
+@@ -7,6 +7,7 @@ include "triton/Dialect/TritonGPU/IR/TritonGPUAttrDefs.td"
  include "mlir/Dialect/Arith/IR/ArithBase.td"
  include "triton/Dialect/Triton/IR/TritonTypes.td"
  include "triton/Dialect/Triton/IR/TritonAttrDefs.td"
@@ -25,8 +31,8 @@ diff --git a/include/triton/Dialect/TritonGPU/IR/TritonGPUOps.td b/include/trito
  include "mlir/IR/OpBase.td"
  include "mlir/Interfaces/SideEffectInterfaces.td" // Pure
  include "mlir/Interfaces/InferTypeOpInterface.td" // SameOperandsAndResultType
-@@ -214,4 +215,19 @@ def TTG_LocalLoadOp : TTG_Op<"local_load
-   let results = (outs TT_Tensor:$result);
+@@ -232,4 +233,19 @@ def TTG_LocalStoreOp : TTG_Op<"local_store", [MemoryEffects<[MemWrite<SharedMemo
+   }];
  }
  
 +def TTNG_SparseDotOp : TTG_Op<"sparse_dot", [
@@ -46,9 +52,10 @@ diff --git a/include/triton/Dialect/TritonGPU/IR/TritonGPUOps.td b/include/trito
 +
  #endif
 diff --git a/lib/Dialect/TritonGPU/IR/Dialect.cpp b/lib/Dialect/TritonGPU/IR/Dialect.cpp
+index 0ce7ecf18..3736a1551 100644
 --- a/lib/Dialect/TritonGPU/IR/Dialect.cpp
 +++ b/lib/Dialect/TritonGPU/IR/Dialect.cpp
-@@ -479,6 +479,119 @@ getDefaultBlockedEncoding(MLIRContext *c
+@@ -483,6 +483,119 @@ getDefaultBlockedEncoding(MLIRContext *context, ArrayRef<int64_t> shape,
    return encoding;
  }
  
@@ -70,7 +77,7 @@ diff --git a/lib/Dialect/TritonGPU/IR/Dialect.cpp b/lib/Dialect/TritonGPU/IR/Dia
 +
 +LogicalResult SparseDotOp::verify() {
 +  // Verify operand A.
-+  auto aTensorTy = getOperand(0).getType().cast<TensorOrMemDesc>();
++  auto aTensorTy = cast<TensorOrMemDesc>(getOperand(0).getType());
 +  auto aElemTy = aTensorTy.getElementType();
 +  if (!aElemTy.isF16() && !aElemTy.isBF16())
 +    return emitError("element type of operand A is not supported");
@@ -78,7 +85,7 @@ diff --git a/lib/Dialect/TritonGPU/IR/Dialect.cpp b/lib/Dialect/TritonGPU/IR/Dia
 +  if (aShape.size() != 2) return emitError("shape of operand A is incorrect");
 +
 +  // Verify operand B.
-+  auto bTensorTy = getOperand(1).getType().cast<TensorOrMemDesc>();
++  auto bTensorTy = cast<TensorOrMemDesc>(getOperand(1).getType());
 +  auto bElemTy = bTensorTy.getElementType();
 +  if (!bElemTy.isF16() && !bElemTy.isBF16())
 +    return emitError("element type of operand B is not supported");
@@ -86,7 +93,7 @@ diff --git a/lib/Dialect/TritonGPU/IR/Dialect.cpp b/lib/Dialect/TritonGPU/IR/Dia
 +  if (bShape.size() != 2) return emitError("shape of operand B is incorrect");
 +
 +  // Verify operand C.
-+  auto cTensorTy = getOperand(2).getType().cast<RankedTensorType>();
++  auto cTensorTy = cast<RankedTensorType>(getOperand(2).getType());
 +  auto cElemTy = cTensorTy.getElementType();
 +  if (!cElemTy.isF32())
 +    return emitError("element type of operand C is not supported");
@@ -101,7 +108,7 @@ diff --git a/lib/Dialect/TritonGPU/IR/Dialect.cpp b/lib/Dialect/TritonGPU/IR/Dia
 +    return emitError("operand element types do not match");
 +
 +  // Verify sparse metadata.
-+  auto metaTy = getOperand(3).getType().cast<RankedTensorType>();
++  auto metaTy = cast<RankedTensorType>(getOperand(3).getType());
 +  auto metaShape = metaTy.getShape();
 +  if (!metaTy.getElementType().isInteger(16) || metaShape.size() != 2)
 +    return emitError("sparse metadata tensor is invalid");
@@ -125,7 +132,7 @@ diff --git a/lib/Dialect/TritonGPU/IR/Dialect.cpp b/lib/Dialect/TritonGPU/IR/Dia
 +//--- SparseDotMetaEncodingAttr ---
 +unsigned SparseDotMetaEncodingAttr::getTotalElemsPerThread(
 +    ArrayRef<int64_t> shape, Type eltTy) const {
-+  auto mmaLayout = getParent().cast<NvidiaMmaEncodingAttr>();
++  auto mmaLayout = mlir::cast<NvidiaMmaEncodingAttr>(getParent());
 +  return product<int64_t>(shape) /
 +         (mmaLayout.getWarpsPerCTA()[0] * kMetadataElementsPerWarp);
 +}
@@ -169,9 +176,10 @@ diff --git a/lib/Dialect/TritonGPU/IR/Dialect.cpp b/lib/Dialect/TritonGPU/IR/Dia
  } // namespace triton
  } // namespace mlir
 diff --git a/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/ConvertLayoutOpToLLVM.cpp b/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/ConvertLayoutOpToLLVM.cpp
+index f8ece0f1c..435610817 100644
 --- a/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/ConvertLayoutOpToLLVM.cpp
 +++ b/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/ConvertLayoutOpToLLVM.cpp
-@@ -38,6 +38,14 @@ Value convertLayout(int opIdx, Conversio
+@@ -43,6 +43,14 @@ Value convertLayout(int opIdx, ConversionPatternRewriter &rewriter,
                      const LLVMTypeConverter *typeConverter, Value thread);
  }
  
@@ -185,19 +193,19 @@ diff --git a/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/ConvertLayoutOpToLLVM.
 +
  namespace {
  
- struct LocalLoadOpConversion
-@@ -59,6 +67,10 @@ public:
-             .isa<NvidiaMmaEncodingAttr>()) {
+ using namespace mlir;
+@@ -67,6 +75,10 @@ public:
+             cast<DotOperandEncodingAttr>(dstLayout).getParent())) {
        return lowerSharedToDotOperand(op, adaptor, getTypeConverter(), rewriter);
      }
-+    if (srcLayout.isa<SharedEncodingAttr>() &&
-+        dstLayout.isa<triton::gpu::SparseDotMetaEncodingAttr>()) {
++    if (isa<SharedEncodingAttr>(srcLayout) &&
++        isa<triton::gpu::SparseDotMetaEncodingAttr>(dstLayout)) {
 +      return lowerSharedToSparseMeta(op, adaptor, getTypeConverter(), rewriter);
 +    }
      return failure();
    }
  
-@@ -130,6 +142,29 @@ private:
+@@ -138,6 +150,26 @@ private:
      rewriter.replaceOp(op, res);
      return success();
    }
@@ -208,13 +216,10 @@ diff --git a/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/ConvertLayoutOpToLLVM.
 +      const LLVMTypeConverter *typeConverter,
 +      ConversionPatternRewriter &rewriter) const {
 +    auto loc = op.getLoc();
-+    auto sparseEncoding = op.getResult()
-+                              .getType()
-+                              .cast<RankedTensorType>()
-+                              .getEncoding()
-+                              .cast<triton::gpu::SparseDotMetaEncodingAttr>();
++    auto sparseEncoding = cast<triton::gpu::SparseDotMetaEncodingAttr>(
++        cast<RankedTensorType>(op.getResult().getType()).getEncoding());
 +    auto llvmElemTy = typeConverter->convertType(
-+        op.getSrc().getType().cast<MemDescType>().getElementType());
++        cast<MemDescType>(op.getSrc().getType()).getElementType());
 +    auto smemObj = getSharedMemoryObjectFromStruct(loc, adaptor.getSrc(),
 +                                                   llvmElemTy, rewriter);
 +    Value res = SharedToSparseDotOperand::convertLayout(
@@ -229,6 +234,7 @@ diff --git a/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/ConvertLayoutOpToLLVM.
  struct ConvertLayoutOpOptimizedConversion
 diff --git a/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/ConvertLayoutOpToLLVM/SharedToSparseDotOperand.cpp b/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/ConvertLayoutOpToLLVM/SharedToSparseDotOperand.cpp
 new file mode 100644
+index 000000000..3011cf73d
 --- /dev/null
 +++ b/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/ConvertLayoutOpToLLVM/SharedToSparseDotOperand.cpp
 @@ -0,0 +1,69 @@
@@ -255,7 +261,7 @@ new file mode 100644
 +    Value thread) {
 +  // Calculate tile size as number of mask elements (4xi4).
 +  NvidiaMmaEncodingAttr mmaLayout =
-+      sparseEncoding.getParent().cast<NvidiaMmaEncodingAttr>();
++      cast<NvidiaMmaEncodingAttr>(sparseEncoding.getParent());
 +  SmallVector<unsigned> shapePerCTATile = {
 +      kTileSize * mmaLayout.getWarpsPerCTA()[0],
 +      kTileSize / kMetadataElementsPerPackedValue};
@@ -272,7 +278,7 @@ new file mode 100644
 +  Value rowId = add(mul(warpGroupId, i32_val(kTileSize)), laneGroupId);
 +
 +  // Calculate number of tile repetitions.
-+  auto shape = tensor.getType().cast<MemDescType>().getShape();
++  auto shape = cast<MemDescType>(tensor.getType()).getShape();
 +  int repM = shape[0] / shapePerCTATile[0];
 +  int repK = shape[1] / shapePerCTATile[1];
 +  assert(repM > 0 && repK > 0);
@@ -302,9 +308,10 @@ new file mode 100644
 +}
 +}  // namespace SharedToSparseDotOperand
 diff --git a/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/DotOpToLLVM.cpp b/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/DotOpToLLVM.cpp
+index 374b9ec9e..1601806b4 100644
 --- a/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/DotOpToLLVM.cpp
 +++ b/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/DotOpToLLVM.cpp
-@@ -32,6 +32,12 @@ LogicalResult convertAsyncWGMMA(triton::
+@@ -32,6 +32,12 @@ LogicalResult convertAsyncWGMMA(triton::nvidia_gpu::DotAsyncOp op,
                                  const LLVMTypeConverter *typeConverter,
                                  ConversionPatternRewriter &rewriter,
                                  Value thread);
@@ -317,7 +324,7 @@ diff --git a/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/DotOpToLLVM.cpp b/thir
  namespace {
  struct DotOpConversion : public ConvertOpToLLVMPattern<triton::DotOp> {
    using ConvertOpToLLVMPattern<triton::DotOp>::ConvertOpToLLVMPattern;
-@@ -180,6 +186,18 @@ struct DotWaitOpConversion
+@@ -174,6 +180,18 @@ struct DotWaitOpConversion
      return success();
    }
  };
@@ -336,7 +343,7 @@ diff --git a/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/DotOpToLLVM.cpp b/thir
  } // namespace
  
  void mlir::triton::NVIDIA::populateDotOpToLLVMPatterns(
-@@ -188,4 +206,5 @@ void mlir::triton::NVIDIA::populateDotOp
+@@ -182,4 +200,5 @@ void mlir::triton::NVIDIA::populateDotOpToLLVMPatterns(
    patterns.add<DotOpConversion>(typeConverter, benefit);
    patterns.add<DotAsyncOpConversion>(typeConverter, benefit);
    patterns.add<DotWaitOpConversion>(typeConverter, benefit);
@@ -344,6 +351,7 @@ diff --git a/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/DotOpToLLVM.cpp b/thir
  }
 diff --git a/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/DotOpToLLVM/Sparse.cpp b/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/DotOpToLLVM/Sparse.cpp
 new file mode 100644
+index 000000000..34d9212d2
 --- /dev/null
 +++ b/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/DotOpToLLVM/Sparse.cpp
 @@ -0,0 +1,339 @@
@@ -397,15 +405,15 @@ new file mode 100644
 +                               const LLVMTypeConverter *typeConverter,
 +                               ConversionPatternRewriter &rewriter) {
 +  // Get number of repetitions across the dimensions.
-+  auto aTensorTy = op.getA().getType().cast<RankedTensorType>();
-+  auto bTensorTy = op.getB().getType().cast<RankedTensorType>();
++  auto aTensorTy = cast<RankedTensorType>(op.getA().getType());
++  auto bTensorTy = cast<RankedTensorType>(op.getB().getType());
 +
-+  auto layoutA = aTensorTy.getEncoding().dyn_cast<DotOperandEncodingAttr>();
-+  auto layoutB = bTensorTy.getEncoding().dyn_cast<DotOperandEncodingAttr>();
++  auto layoutA = dyn_cast<DotOperandEncodingAttr>(aTensorTy.getEncoding());
++  auto layoutB = dyn_cast<DotOperandEncodingAttr>(bTensorTy.getEncoding());
 +  assert(layoutA != nullptr && layoutB != nullptr);
 +
 +  int bitwidth = aTensorTy.getElementType().getIntOrFloatBitWidth();
-+  auto mmaEnc = layoutA.getParent().cast<NvidiaMmaEncodingAttr>();
++  auto mmaEnc = cast<NvidiaMmaEncodingAttr>(layoutA.getParent());
 +  auto repA = mmaEnc.getMMAv2Rep(triton::gpu::getShapePerCTA(aTensorTy),
 +                                 bitwidth, layoutA.getOpIdx());
 +  auto repB = mmaEnc.getMMAv2Rep(triton::gpu::getShapePerCTA(bTensorTy),
@@ -435,7 +443,7 @@ new file mode 100644
 +  }
 +
 +  // Flatten accumulator values.
-+  auto dTensorTy = op.getD().getType().cast<RankedTensorType>();
++  auto dTensorTy = cast<RankedTensorType>(op.getD().getType());
 +  auto fc = unpackLLElements(loc, adaptor.getC(), rewriter);
 +
 +  // Create `mma.sp` instruction for 4/8 core matrices.
@@ -512,7 +520,7 @@ new file mode 100644
 +                     Location loc, std::vector<unsigned int> instrShape,
 +                     bool trans, int dimWpt, Value warpId, MemDescType tensorTy,
 +                     Value baseDesc, int minor) {
-+  auto sharedLayout = tensorTy.getEncoding().cast<SharedEncodingAttr>();
++  auto sharedLayout = cast<SharedEncodingAttr>(tensorTy.getEncoding());
 +  int elemBytes = tensorTy.getElementTypeBitWidth() / 8;
 +  int elemsPerSwizzlingRow =
 +      kMmaLineSize / sharedLayout.getPerPhase() / elemBytes;
@@ -541,10 +549,10 @@ new file mode 100644
 +                                 ConversionPatternRewriter &rewriter,
 +                                 Value thread) {
 +  // Get number of repetitions across the dimensions.
-+  auto aTensorTy = op.getA().getType().cast<MemDescType>();
-+  auto bTensorTy = op.getB().getType().cast<MemDescType>();
-+  auto dTensorTy = op.getD().getType().cast<RankedTensorType>();
-+  auto mmaEnc = dTensorTy.getEncoding().cast<NvidiaMmaEncodingAttr>();
++  auto aTensorTy = cast<MemDescType>(op.getA().getType());
++  auto bTensorTy = cast<MemDescType>(op.getB().getType());
++  auto dTensorTy = cast<RankedTensorType>(op.getD().getType());
++  auto mmaEnc = cast<NvidiaMmaEncodingAttr>(dTensorTy.getEncoding());
 +
 +  auto shapePerCTA = getShapePerCTA(dTensorTy);
 +  auto shapePerCTATile = getShapePerCTATile(mmaEnc);
@@ -573,7 +581,7 @@ new file mode 100644
 +    auto sharedObj = getSharedMemoryObjectFromStruct(
 +        loc, arg, typeConverter->convertType(tensorTy.getElementType()),
 +        rewriter);
-+    auto sharedLayout = tensorTy.getEncoding().cast<SharedEncodingAttr>();
++    auto sharedLayout = cast<SharedEncodingAttr>(tensorTy.getEncoding());
 +    auto shape = getShapePerCTA(tensorTy);
 +    auto ord = sharedLayout.getOrder();
 +    int byteSize = aTensorTy.getElementTypeBitWidth() / 8;
@@ -671,9 +679,9 @@ new file mode 100644
 +                                 SparseDotOp::Adaptor adaptor,
 +                                 const LLVMTypeConverter *typeConverter,
 +                                 ConversionPatternRewriter &rewriter) {
-+  auto resultTy = op.getResult().getType().cast<RankedTensorType>();
++  auto resultTy = cast<RankedTensorType>(op.getResult().getType());
 +  NvidiaMmaEncodingAttr mmaLayout =
-+      resultTy.getEncoding().cast<NvidiaMmaEncodingAttr>();
++      cast<NvidiaMmaEncodingAttr>(resultTy.getEncoding());
 +
 +  if (mmaLayout.isAmpere()) {
 +    return convertSparseMMA(op, adaptor, typeConverter, rewriter);
@@ -687,9 +695,10 @@ new file mode 100644
 +      "Unsupported SparseDotOp found when converting TritonGPU to LLVM.");
 +}
 diff --git a/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/DotOpToLLVM/WGMMA.cpp b/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/DotOpToLLVM/WGMMA.cpp
+index 738f0fe04..867939f65 100644
 --- a/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/DotOpToLLVM/WGMMA.cpp
 +++ b/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/DotOpToLLVM/WGMMA.cpp
-@@ -87,8 +87,8 @@ int64_t getSwizzlingFromLayout(const Sha
+@@ -88,8 +88,8 @@ int64_t getSwizzlingFromLayout(const SharedEncodingAttr &layout,
    return swizzlingByteWidth;
  }
  
diff --git a/third_party/triton/xla_extensions/sparse_dot_fixes_y24w17.patch b/third_party/triton/xla_extensions/sparse_dot_fixes_y24w17.patch
index 9d1ae2e91cae3f..ce009aa688e9bf 100644
--- a/third_party/triton/xla_extensions/sparse_dot_fixes_y24w17.patch
+++ b/third_party/triton/xla_extensions/sparse_dot_fixes_y24w17.patch
@@ -1,30 +1,8 @@
-diff --git a/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/ConvertLayoutOpToLLVM/SharedToSparseDotOperand.cpp b/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/ConvertLayoutOpToLLVM/SharedToSparseDotOperand.cpp
---- a/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/ConvertLayoutOpToLLVM/SharedToSparseDotOperand.cpp
-+++ b/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/ConvertLayoutOpToLLVM/SharedToSparseDotOperand.cpp
-@@ -22,16 +22,16 @@ Value convertLayout(
-   // Calculate tile size as number of mask elements (4xi4).
-   NvidiaMmaEncodingAttr mmaLayout =
-       sparseEncoding.getParent().cast<NvidiaMmaEncodingAttr>();
-+  SmallVector<unsigned> warpsPerCTA = mmaLayout.getWarpsPerCTA();
-   SmallVector<unsigned> shapePerCTATile = {
--      kTileSize * mmaLayout.getWarpsPerCTA()[0],
--      kTileSize / kMetadataElementsPerPackedValue};
-+      kTileSize * warpsPerCTA[0], kTileSize / kMetadataElementsPerPackedValue};
-   Value strideM = smemObj.strides[0];
-   Value strideK = smemObj.strides[1];
- 
-   // Calculate offset in the tile for the current thread.
-   Value threadsPerWarp = i32_val(kThreadsPerWarp);
-   Value warpId = udiv(thread, threadsPerWarp);
--  Value warpGroupId = urem(warpId, i32_val(shapePerCTATile[0] / kTileSize));
-+  Value warpGroupId = udiv(warpId, i32_val(warpsPerCTA[1]));
-   Value laneId = urem(thread, threadsPerWarp);
-   Value laneGroupId = udiv(laneId, i32_val(kThreadsInGroup));
-   Value columnId = urem(laneId, i32_val(shapePerCTATile[1]));
 diff --git a/lib/Dialect/TritonGPU/Transforms/AccelerateMatmul.cpp b/lib/Dialect/TritonGPU/Transforms/AccelerateMatmul.cpp
+index 0516fc56f..1f27f8a43 100644
 --- a/lib/Dialect/TritonGPU/Transforms/AccelerateMatmul.cpp
 +++ b/lib/Dialect/TritonGPU/Transforms/AccelerateMatmul.cpp
-@@ -139,6 +139,7 @@ class BlockedToMMA : public mlir::Rewrit
+@@ -142,6 +142,7 @@ class BlockedToMMA : public mlir::RewritePattern {
                  mlir::TypeID::get<arith::ArithDialect>());
    }
  
@@ -32,7 +10,7 @@ diff --git a/lib/Dialect/TritonGPU/Transforms/AccelerateMatmul.cpp b/lib/Dialect
    // Finds the first different bitwidth in the chain of shape-preserving
    // unary ops that x depends on.
    // There are two primary scenarios:
-@@ -172,7 +173,6 @@ class BlockedToMMA : public mlir::Rewrit
+@@ -175,7 +176,6 @@ class BlockedToMMA : public mlir::RewritePattern {
      return origBitWidth;
    }
  
@@ -40,7 +18,7 @@ diff --git a/lib/Dialect/TritonGPU/Transforms/AccelerateMatmul.cpp b/lib/Dialect
    BlockedToMMA(mlir::MLIRContext *context, int computeCapability)
        : mlir::RewritePattern(tt::DotOp::getOperationName(), 2, context),
          computeCapability(computeCapability) {}
-@@ -388,18 +388,22 @@ class SparseBlockedToMMA : public mlir::
+@@ -389,18 +389,22 @@ class SparseBlockedToMMA : public mlir::RewritePattern {
                                                          newRetType, oldAcc);
  
      if (versionMajor == 2) {
@@ -49,7 +27,7 @@ diff --git a/lib/Dialect/TritonGPU/Transforms/AccelerateMatmul.cpp b/lib/Dialect
 +      int kWidth = 32 / minBitwidth;
 +
        // convert A operand
-       auto oldAType = a.getType().cast<RankedTensorType>();
+       auto oldAType = cast<RankedTensorType>(a.getType());
 -      auto newAEncoding = ttg::DotOperandEncodingAttr::get(
 -          ctx, 0, mmaEnc, oldAType.getElementType());
 +      auto newAEncoding =
@@ -59,7 +37,7 @@ diff --git a/lib/Dialect/TritonGPU/Transforms/AccelerateMatmul.cpp b/lib/Dialect
        a = rewriter.create<ttg::ConvertLayoutOp>(a.getLoc(), newAType, a);
  
        // convert B operand
-       auto oldBType = b.getType().cast<RankedTensorType>();
+       auto oldBType = cast<RankedTensorType>(b.getType());
 -      auto newBEncoding = ttg::DotOperandEncodingAttr::get(
 -          ctx, 1, mmaEnc, oldBType.getElementType());
 +      auto newBEncoding =
@@ -67,3 +45,27 @@ diff --git a/lib/Dialect/TritonGPU/Transforms/AccelerateMatmul.cpp b/lib/Dialect
        auto newBType = RankedTensorType::get(
            oldBType.getShape(), oldBType.getElementType(), newBEncoding);
        b = rewriter.create<ttg::ConvertLayoutOp>(b.getLoc(), newBType, b);
+diff --git a/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/ConvertLayoutOpToLLVM/SharedToSparseDotOperand.cpp b/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/ConvertLayoutOpToLLVM/SharedToSparseDotOperand.cpp
+index 3011cf73d..ea587dced 100644
+--- a/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/ConvertLayoutOpToLLVM/SharedToSparseDotOperand.cpp
++++ b/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/ConvertLayoutOpToLLVM/SharedToSparseDotOperand.cpp
+@@ -22,16 +22,16 @@ Value convertLayout(
+   // Calculate tile size as number of mask elements (4xi4).
+   NvidiaMmaEncodingAttr mmaLayout =
+       cast<NvidiaMmaEncodingAttr>(sparseEncoding.getParent());
++  SmallVector<unsigned> warpsPerCTA = mmaLayout.getWarpsPerCTA();
+   SmallVector<unsigned> shapePerCTATile = {
+-      kTileSize * mmaLayout.getWarpsPerCTA()[0],
+-      kTileSize / kMetadataElementsPerPackedValue};
++      kTileSize * warpsPerCTA[0], kTileSize / kMetadataElementsPerPackedValue};
+   Value strideM = smemObj.strides[0];
+   Value strideK = smemObj.strides[1];
+ 
+   // Calculate offset in the tile for the current thread.
+   Value threadsPerWarp = i32_val(kThreadsPerWarp);
+   Value warpId = udiv(thread, threadsPerWarp);
+-  Value warpGroupId = urem(warpId, i32_val(shapePerCTATile[0] / kTileSize));
++  Value warpGroupId = udiv(warpId, i32_val(warpsPerCTA[1]));
+   Value laneId = urem(thread, threadsPerWarp);
+   Value laneGroupId = udiv(laneId, i32_val(kThreadsInGroup));
+   Value columnId = urem(laneId, i32_val(shapePerCTATile[1]));
diff --git a/third_party/triton/xla_extensions/sparse_dot_nvgpu.patch b/third_party/triton/xla_extensions/sparse_dot_nvgpu.patch
index ced13cff33fd16..791618363b2f34 100644
--- a/third_party/triton/xla_extensions/sparse_dot_nvgpu.patch
+++ b/third_party/triton/xla_extensions/sparse_dot_nvgpu.patch
@@ -1,7 +1,8 @@
-diff --git a/include/triton/Dialect/NVGPU/IR/NVGPUOps.td b/include/triton/Dialect/NVGPU/IR/NVGPUOps.td
---- a/include/triton/Dialect/NVGPU/IR/NVGPUOps.td
-+++ b/include/triton/Dialect/NVGPU/IR/NVGPUOps.td
-@@ -87,6 +87,15 @@ def NVGPU_WGMMAOp : NVGPU_Op<"wgmma", []
+diff --git a/third_party/nvidia/include/Dialect/NVGPU/IR/NVGPUOps.td b/third_party/nvidia/include/Dialect/NVGPU/IR/NVGPUOps.td
+index ca9d18873..d39bc6ec4 100644
+--- a/third_party/nvidia/include/Dialect/NVGPU/IR/NVGPUOps.td
++++ b/third_party/nvidia/include/Dialect/NVGPU/IR/NVGPUOps.td
+@@ -87,6 +87,15 @@ def NVGPU_WGMMAOp : NVGPU_Op<"wgmma", []> {
    let assemblyFormat = "$opA `,` $opB (`,` $opC^)? attr-dict `:` functional-type(operands, $res)";
  }
  
@@ -18,9 +19,10 @@ diff --git a/include/triton/Dialect/NVGPU/IR/NVGPUOps.td b/include/triton/Dialec
    let arguments = (ins LLVM_AnyPointer:$addr, I32:$ctaId, I32Attr:$bitwidth, I32Attr:$vec);
    let builders = [
 diff --git a/third_party/nvidia/lib/NVGPUToLLVM/NVGPUToLLVMPass.cpp b/third_party/nvidia/lib/NVGPUToLLVM/NVGPUToLLVMPass.cpp
+index e19216520..aacbfb569 100644
 --- a/third_party/nvidia/lib/NVGPUToLLVM/NVGPUToLLVMPass.cpp
 +++ b/third_party/nvidia/lib/NVGPUToLLVM/NVGPUToLLVMPass.cpp
-@@ -688,6 +688,84 @@ public:
+@@ -668,6 +668,84 @@ public:
    }
  };
  
@@ -31,7 +33,7 @@ diff --git a/third_party/nvidia/lib/NVGPUToLLVM/NVGPUToLLVMPass.cpp b/third_part
 +  using Base::Base;
 +
 +  std::vector<std::string> getOutputConstraints(ttn::SparseWGMMAOp op) const {
-+    auto outputStructType = op.getType().cast<LLVM::LLVMStructType>();
++    auto outputStructType = cast<LLVM::LLVMStructType>(op.getType());
 +    uint32_t numOutputRegs = outputStructType.getBody().size();
 +    std::string output =
 +        outputStructType.getBody().front().isF32() ? "=f" : "=r";
@@ -71,7 +73,7 @@ diff --git a/third_party/nvidia/lib/NVGPUToLLVM/NVGPUToLLVMPass.cpp b/third_part
 +
 +    // Output and operand C
 +    uint32_t numCRegs =
-+        op.getType().cast<LLVM::LLVMStructType>().getBody().size();
++        cast<LLVM::LLVMStructType>(op.getType()).getBody().size();
 +    args += "{";
 +    for (uint32_t i = 0; i < numCRegs; ++i) {
 +      args += "$" + std::to_string(asmOpIdx++) + (i == numCRegs - 1 ? "" : ",");
@@ -105,13 +107,17 @@ diff --git a/third_party/nvidia/lib/NVGPUToLLVM/NVGPUToLLVMPass.cpp b/third_part
  class ConvertNVGPUToLLVM : public ConvertNVGPUToLLVMBase<ConvertNVGPUToLLVM> {
  
  public:
-@@ -711,7 +789,8 @@ public:
+@@ -688,10 +766,9 @@ public:
+     patterns.add<NVGPUOpGenericPattern<ttn::ClusterCTAIdOp>>(
+         context, Cluster_Cta_Id_Op, Constraints({"=r"}), Constraints());
  
-     patterns.add<FenceAsyncSharedOpPattern, StoreMatrixOpPattern,
-                  ClusterArriveOpPattern, LoadDSmemOpPattern, WGMMAOpPattern,
--                 WGMMAWaitGroupOpPattern, StoreDSmemOpPattern>(context);
-+                 WGMMAWaitGroupOpPattern, StoreDSmemOpPattern,
-+                 SparseWGMMAOpPattern>(context);
+-    patterns
+-        .add<FenceAsyncSharedOpPattern, StoreMatrixOpPattern,
+-             ClusterArriveOpPattern, WGMMAOpPattern, WGMMAWaitGroupOpPattern>(
+-            context);
++    patterns.add<FenceAsyncSharedOpPattern, StoreMatrixOpPattern,
++                 ClusterArriveOpPattern, WGMMAOpPattern,
++                 WGMMAWaitGroupOpPattern, SparseWGMMAOpPattern>(context);
  
      if (applyPatternsAndFoldGreedily(mod, std::move(patterns)).failed())
        signalPassFailure();
diff --git a/third_party/triton/xla_extensions/sparse_dot_passes.patch b/third_party/triton/xla_extensions/sparse_dot_passes.patch
index 74662cf3c90dc9..9136cb84b24254 100644
--- a/third_party/triton/xla_extensions/sparse_dot_passes.patch
+++ b/third_party/triton/xla_extensions/sparse_dot_passes.patch
@@ -1,7 +1,8 @@
 diff --git a/lib/Conversion/TritonToTritonGPU/TritonToTritonGPUPass.cpp b/lib/Conversion/TritonToTritonGPU/TritonToTritonGPUPass.cpp
+index 4aa2712ec..16a6253d7 100644
 --- a/lib/Conversion/TritonToTritonGPU/TritonToTritonGPUPass.cpp
 +++ b/lib/Conversion/TritonToTritonGPU/TritonToTritonGPUPass.cpp
-@@ -277,6 +277,89 @@ struct TritonDotPattern : public OpConve
+@@ -279,6 +279,89 @@ struct TritonDotPattern : public OpConversionPattern<triton::DotOp> {
    }
  };
  
@@ -12,7 +13,7 @@ diff --git a/lib/Conversion/TritonToTritonGPU/TritonToTritonGPUPass.cpp b/lib/Co
 +  LogicalResult matchAndRewrite(
 +      triton::gpu::SparseDotOp op, OpAdaptor adaptor,
 +      ConversionPatternRewriter &rewriter) const override {
-+    RankedTensorType origType = op.getType().cast<RankedTensorType>();
++    RankedTensorType origType = cast<RankedTensorType>(op.getType());
 +    auto origShape = origType.getShape();
 +    auto typeConverter = getTypeConverter<TritonGPUTypeConverter>();
 +    int numWarps = typeConverter->getNumWarps();
@@ -40,8 +41,8 @@ diff --git a/lib/Conversion/TritonToTritonGPU/TritonToTritonGPUPass.cpp b/lib/Co
 +        RankedTensorType::get(origShape, origType.getElementType(), dEncoding);
 +
 +    // a & b must be of smem layout
-+    auto aType = adaptor.getA().getType().cast<RankedTensorType>();
-+    auto bType = adaptor.getB().getType().cast<RankedTensorType>();
++    auto aType = cast<RankedTensorType>(adaptor.getA().getType());
++    auto bType = cast<RankedTensorType>(adaptor.getB().getType());
 +    Type aEltType = aType.getElementType();
 +    Type bEltType = bType.getElementType();
 +    Attribute aEncoding = aType.getEncoding();
@@ -51,14 +52,14 @@ diff --git a/lib/Conversion/TritonToTritonGPU/TritonToTritonGPUPass.cpp b/lib/Co
 +    Value a = adaptor.getA();
 +    Value b = adaptor.getB();
 +    Value c = adaptor.getC();
-+    if (!aEncoding.isa<triton::gpu::DotOperandEncodingAttr>()) {
++    if (!isa<triton::gpu::DotOperandEncodingAttr>(aEncoding)) {
 +      Attribute encoding = triton::gpu::DotOperandEncodingAttr::get(
 +          getContext(), 0, dEncoding, aEltType);
 +      auto dstType =
 +          RankedTensorType::get(aType.getShape(), aEltType, encoding);
 +      a = rewriter.create<triton::gpu::ConvertLayoutOp>(a.getLoc(), dstType, a);
 +    }
-+    if (!bEncoding.isa<triton::gpu::DotOperandEncodingAttr>()) {
++    if (!isa<triton::gpu::DotOperandEncodingAttr>(bEncoding)) {
 +      Attribute encoding = triton::gpu::DotOperandEncodingAttr::get(
 +          getContext(), 1, dEncoding, bEltType);
 +      auto dstType =
@@ -68,11 +69,11 @@ diff --git a/lib/Conversion/TritonToTritonGPU/TritonToTritonGPUPass.cpp b/lib/Co
 +    c = rewriter.create<triton::gpu::ConvertLayoutOp>(c.getLoc(), retType, c);
 +
 +    // aMeta must be of smem layout
-+    auto aMetaType = adaptor.getAMeta().getType().cast<RankedTensorType>();
++    auto aMetaType = cast<RankedTensorType>(adaptor.getAMeta().getType());
 +    Attribute aMetaEncoding = aMetaType.getEncoding();
 +    if (!aMetaEncoding) return failure();
 +    Value aMeta = adaptor.getAMeta();
-+    if (!aMetaEncoding.isa<triton::gpu::SparseDotMetaEncodingAttr>()) {
++    if (!isa<triton::gpu::SparseDotMetaEncodingAttr>(aMetaEncoding)) {
 +      Attribute encoding =
 +          triton::gpu::SparseDotMetaEncodingAttr::get(getContext(), dEncoding);
 +      auto dstType = RankedTensorType::get(
@@ -91,17 +92,17 @@ diff --git a/lib/Conversion/TritonToTritonGPU/TritonToTritonGPUPass.cpp b/lib/Co
  struct TritonCatPattern : public OpConversionPattern<triton::CatOp> {
    using OpConversionPattern::OpConversionPattern;
  
-@@ -550,6 +633,7 @@ void populateTritonPatterns(TritonGPUTyp
-       GenericOpPattern<triton::AtomicRMWOp>, GenericOpPattern<ReturnOp>,
+@@ -553,6 +636,7 @@ void populateTritonPatterns(TritonGPUTypeConverter &typeConverter,
+       GenericOpPattern<triton::ExperimentalDescriptorStoreOp>,
        GenericOpPattern<triton::CallOp>, TritonFuncOpPattern>(typeConverter,
                                                               context);
 +  patterns.insert<TritonSparseDotPattern>(typeConverter, context);
  }
  
  //
-@@ -788,6 +872,12 @@ public:
-                  IntegerAttr::get(
-                      i32_ty, llvm::APInt(32, computeCapability.getValue())));
+@@ -794,6 +878,12 @@ public:
+     mod->setAttr(AttrTargetName,
+                  StringAttr::get(context, this->target.getValue()));
  
 +    // Only transform sparse dot op with undefined layout.
 +    target.addDynamicallyLegalOp<triton::gpu::SparseDotOp>(
@@ -113,9 +114,10 @@ diff --git a/lib/Conversion/TritonToTritonGPU/TritonToTritonGPUPass.cpp b/lib/Co
        return signalPassFailure();
  
 diff --git a/lib/Dialect/TritonGPU/Transforms/AccelerateMatmul.cpp b/lib/Dialect/TritonGPU/Transforms/AccelerateMatmul.cpp
+index 098ee85e4..0516fc56f 100644
 --- a/lib/Dialect/TritonGPU/Transforms/AccelerateMatmul.cpp
 +++ b/lib/Dialect/TritonGPU/Transforms/AccelerateMatmul.cpp
-@@ -42,8 +42,9 @@ static int getMMAVersionSafe(int compute
+@@ -44,8 +44,9 @@ static int getMMAVersionSafe(int computeCapability, tt::DotOp op) {
    return 0;
  }
  
@@ -126,7 +128,7 @@ diff --git a/lib/Dialect/TritonGPU/Transforms/AccelerateMatmul.cpp b/lib/Dialect
    auto rank = shape.size();
    // Early exit for batched matmul
    if (rank == 3)
-@@ -56,14 +57,14 @@ warpsPerTileV2(tt::DotOp dotOp, const Ar
+@@ -58,8 +59,8 @@ warpsPerTileV2(tt::DotOp dotOp, const ArrayRef<int64_t> shape, int numWarps) {
    auto slices = multiRootGetSlice(dotOp, {filter}, {filter});
    bool hasChainedDot = false;
    for (Operation *op : slices) {
@@ -137,14 +139,7 @@ diff --git a/lib/Dialect/TritonGPU/Transforms/AccelerateMatmul.cpp b/lib/Dialect
        auto resTy = chainedDot.getResult().getType();
        if (resTy.getRank() != rank) {
          continue;
-       }
-       if (auto mmaEncoding =
--              resTy.getEncoding().dyn_cast<NvidiaMmaEncodingAttr>()) {
-+              resTy.getEncoding().template dyn_cast<NvidiaMmaEncodingAttr>()) {
-         return ttg::getWarpsPerCTA(mmaEncoding);
-       }
-       hasChainedDot = true;
-@@ -101,12 +102,13 @@ warpsPerTileV2(tt::DotOp dotOp, const Ar
+@@ -103,12 +104,13 @@ warpsPerTileV2(tt::DotOp dotOp, const ArrayRef<int64_t> shape, int numWarps) {
    return ret;
  }
  
@@ -162,7 +157,7 @@ diff --git a/lib/Dialect/TritonGPU/Transforms/AccelerateMatmul.cpp b/lib/Dialect
        slices.end())
      return {(unsigned)numWarps, 1};
  
-@@ -175,9 +177,10 @@ public:
+@@ -178,9 +180,10 @@ public:
        : mlir::RewritePattern(tt::DotOp::getOperationName(), 2, context),
          computeCapability(computeCapability) {}
  
@@ -176,7 +171,7 @@ diff --git a/lib/Dialect/TritonGPU/Transforms/AccelerateMatmul.cpp b/lib/Dialect
      switch (version) {
      case 2:
        return warpsPerTileV2(dotOp, shape, numWarps);
-@@ -337,6 +340,98 @@ public:
+@@ -335,6 +338,98 @@ public:
      return success();
    }
  };
@@ -201,7 +196,7 @@ diff --git a/lib/Dialect/TritonGPU/Transforms/AccelerateMatmul.cpp b/lib/Dialect
 +    // Check data-types and SM compatibility
 +    RankedTensorType oldRetType = dotOp.getType();
 +    if (!oldRetType.getEncoding() ||
-+        oldRetType.getEncoding().isa<ttg::NvidiaMmaEncodingAttr>())
++        isa<ttg::NvidiaMmaEncodingAttr>(oldRetType.getEncoding()))
 +      return failure();
 +
 +    assert(computeCapability >= 80 &&
@@ -216,7 +211,7 @@ diff --git a/lib/Dialect/TritonGPU/Transforms/AccelerateMatmul.cpp b/lib/Dialect
 +
 +    auto instrShape =
 +        mmaVersionToInstrShape(versionMajor, retShapePerCTA,
-+                               a.getType().cast<TensorOrMemDesc>(), numWarps);
++                               cast<TensorOrMemDesc>(a.getType()), numWarps);
 +    auto warpsPerTile = BlockedToMMA::getWarpsPerTile(
 +        dotOp, retShapePerCTA, versionMajor, numWarps, instrShape);
 +    ttg::NvidiaMmaEncodingAttr mmaEnc =
@@ -232,7 +227,7 @@ diff --git a/lib/Dialect/TritonGPU/Transforms/AccelerateMatmul.cpp b/lib/Dialect
 +
 +    if (versionMajor == 2) {
 +      // convert A operand
-+      auto oldAType = a.getType().cast<RankedTensorType>();
++      auto oldAType = cast<RankedTensorType>(a.getType());
 +      auto newAEncoding = ttg::DotOperandEncodingAttr::get(
 +          ctx, 0, mmaEnc, oldAType.getElementType());
 +      auto newAType = RankedTensorType::get(
@@ -240,7 +235,7 @@ diff --git a/lib/Dialect/TritonGPU/Transforms/AccelerateMatmul.cpp b/lib/Dialect
 +      a = rewriter.create<ttg::ConvertLayoutOp>(a.getLoc(), newAType, a);
 +
 +      // convert B operand
-+      auto oldBType = b.getType().cast<RankedTensorType>();
++      auto oldBType = cast<RankedTensorType>(b.getType());
 +      auto newBEncoding = ttg::DotOperandEncodingAttr::get(
 +          ctx, 1, mmaEnc, oldBType.getElementType());
 +      auto newBType = RankedTensorType::get(
@@ -253,7 +248,7 @@ diff --git a/lib/Dialect/TritonGPU/Transforms/AccelerateMatmul.cpp b/lib/Dialect
 +
 +    // convert metadata
 +    Value meta = dotOp.getAMeta();
-+    auto oldMetaType = meta.getType().cast<RankedTensorType>();
++    auto oldMetaType = cast<RankedTensorType>(meta.getType());
 +    auto newMetaType = RankedTensorType::get(
 +        oldMetaType.getShape(), oldMetaType.getElementType(),
 +        SparseDotMetaEncodingAttr::get(ctx, mmaEnc));
@@ -275,7 +270,7 @@ diff --git a/lib/Dialect/TritonGPU/Transforms/AccelerateMatmul.cpp b/lib/Dialect
  } // namespace
  
  static Value promoteOperand(OpBuilder &builder, Location loc, Value operand,
-@@ -397,6 +491,7 @@ public:
+@@ -394,6 +489,7 @@ public:
  
      mlir::RewritePatternSet patterns(context);
      patterns.add<::BlockedToMMA>(context, computeCapability);
@@ -284,33 +279,31 @@ diff --git a/lib/Dialect/TritonGPU/Transforms/AccelerateMatmul.cpp b/lib/Dialect
        signalPassFailure();
      }
 diff --git a/lib/Dialect/TritonGPU/Transforms/Pipeliner/MatmulLoopPipeline.cpp b/lib/Dialect/TritonGPU/Transforms/Pipeliner/MatmulLoopPipeline.cpp
+index 97ca6a840..f0ef124ff 100644
 --- a/lib/Dialect/TritonGPU/Transforms/Pipeliner/MatmulLoopPipeline.cpp
 +++ b/lib/Dialect/TritonGPU/Transforms/Pipeliner/MatmulLoopPipeline.cpp
-@@ -47,6 +47,10 @@ struct PipelinedOpInfo {
-   bool loadIsMMAV3 = false;
+@@ -188,6 +188,10 @@ public:
+   }
  };
  
-+bool isDotOp(Operation* op) {
++static bool isDotOp(Operation* op) {
 +  return isa<tt::DotOp, ttg::SparseDotOp>(op);
 +}
 +
- } // namespace
- 
  static bool isMMAv3Dot(Operation *op) {
-@@ -165,22 +169,28 @@ getSharedEncIfAllUsersAreDotEnc(Value val) {
+   auto dot = dyn_cast<tt::DotOp>(op);
+   if (!dot)
+@@ -399,19 +403,28 @@ getSharedEncIfAllUsersAreDotEnc(Value val) {
      } else {
        if (!isa<ttg::LocalLoadOp, ttg::ConvertLayoutOp>(user))
          return std::nullopt;
--      auto dotOpEnc = user->getResult(0)
--                          .getType()
--                          .cast<TensorOrMemDesc>()
--                          .getEncoding()
--                          .dyn_cast<ttg::DotOperandEncodingAttr>();
+-      auto dotOpEnc = dyn_cast<ttg::DotOperandEncodingAttr>(
+-          cast<TensorOrMemDesc>(user->getResult(0).getType()).getEncoding());
 -      if (!dotOpEnc)
 +      auto enc =
-+          user->getResult(0).getType().cast<TensorOrMemDesc>().getEncoding();
++          cast<TensorOrMemDesc>(user->getResult(0).getType()).getEncoding();
 +      if (isa<ttg::DotOperandEncodingAttr>(enc)) {
-+        auto srcTy = val.getType().cast<TensorOrMemDesc>();
++        auto srcTy = cast<TensorOrMemDesc>(val.getType());
 +        auto CTALayout = ttg::getCTALayout(srcTy.getEncoding());
 +        auto order = ttg::getOrder(srcTy.getEncoding());
 +        unsigned bitWidth = srcTy.getElementType().getIntOrFloatBitWidth();
@@ -321,14 +314,14 @@ diff --git a/lib/Dialect/TritonGPU/Transforms/Pipeliner/MatmulLoopPipeline.cpp b
 +            srcTy.getElementType().getIntOrFloatBitWidth(),
 +            /*needTrans=*/false);
 +      } else if (isa<ttg::SparseDotMetaEncodingAttr>(enc)) {
-+        auto srcTy = val.getType().cast<TensorOrMemDesc>();
++        auto srcTy = cast<TensorOrMemDesc>(val.getType());
 +        tempAttr = ttg::SharedEncodingAttr::get(
 +            val.getContext(), /*vec=*/1, /*perPhase=*/1, /*maxPhase=*/1,
 +            ttg::getOrder(srcTy.getEncoding()),
 +            ttg::getCTALayout(srcTy.getEncoding()));
 +      } else {
          return std::nullopt;
--      auto srcTy = val.getType().cast<TensorOrMemDesc>();
+-      auto srcTy = cast<TensorOrMemDesc>(val.getType());
 -      auto CTALayout = ttg::getCTALayout(srcTy.getEncoding());
 -      auto order = ttg::getOrder(srcTy.getEncoding());
 -      unsigned bitWidth = srcTy.getElementType().getIntOrFloatBitWidth();
@@ -341,71 +334,63 @@ diff --git a/lib/Dialect/TritonGPU/Transforms/Pipeliner/MatmulLoopPipeline.cpp b
      }
      // Check that the shared encodings needed by the users are compatible.
      if (!tempAttr || (attr != nullptr && attr != tempAttr))
-@@ -313,7 +323,7 @@ loadOpsToDistanceAndUse(scf::ForOp forOp) {
+@@ -518,7 +531,7 @@ loadOpsToIndirectionLevelAndUse(scf::ForOp forOp) {
        };
  
    for (Operation &op : forOp.getBody()->without_terminator()) {
 -    if (!isa<tt::DotOp>(op))
 +    if (!isDotOp(&op))
        continue;
+     seen.clear();
      dfs(&op, 0, &op);
-   }
-@@ -391,7 +401,8 @@ collectOpsToPipeline(scf::ForOp forOp,
-   // loads.
-   for (auto &[loadOp, distAndUse] : loadOpToDistAndUse) {
-     PipelinedOpInfo loadInfo;
--    if (auto dot = dyn_cast<tt::DotOp>(distAndUse.second)) {
-+    if (isDotOp(distAndUse.second)) {
-+      auto dot = dyn_cast<tt::DotOp>(distAndUse.second);
-       if (loadIsMMAv3(loadOp)) {
+@@ -595,7 +608,8 @@ assignMemoryLayouts(llvm::SmallVector<std::tuple<Operation *, int, Operation *>>
+         continue;
+     }
+ 
+-    if (auto dot = dyn_cast<tt::DotOp>(use)) {
++    if (isDotOp(use)) {
++      auto dot = dyn_cast<tt::DotOp>(use);
+       loadInfo.usedByDot = true;
+       if (loadIsMMAv3(op)) {
          loadInfo.loadIsMMAV3 = true;
-         loadInfo.sharedEncoding =
-@@ -410,7 +421,7 @@ collectOpsToPipeline(scf::ForOp forOp,
+@@ -614,7 +628,7 @@ assignMemoryLayouts(llvm::SmallVector<std::tuple<Operation *, int, Operation *>>
          // The codegen bug is caught by an assertion, so if you think you've
          // fixed it, feel free to delete this code and see if the assert still
          // fails.  :)
 -        if (!loadInfo.sharedEncoding) {
 +        if (dot && !loadInfo.sharedEncoding) {
-           if (auto dotEnc = dot.getResult()
-                                 .getType()
-                                 .getEncoding()
-@@ -788,7 +799,7 @@ bool mlir::triton::preProcessLoopAndGetSchedule(
-     int useStage = opToInfo[info.use].stage;
-     int numBuffers = useStage - defStage;
- 
--    if (hasMMAV3 && isa<tt::DotOp>(info.use)) {
-+    if (hasMMAV3 && isDotOp(info.use)) {
-       // For MMAv3, we need an extra buffer as this is assumed in the wgmma
-       // pipelining post-processing.
-       numBuffers++;
+           if (auto dotEnc = dyn_cast<ttg::NvidiaMmaEncodingAttr>(
+                   dot.getResult().getType().getEncoding())) {
+             auto loadTy = cast<RankedTensorType>(op->getResultTypes()[0]);
 diff --git a/lib/Dialect/TritonGPU/Transforms/ReduceDataDuplication.cpp b/lib/Dialect/TritonGPU/Transforms/ReduceDataDuplication.cpp
+index 2211df31b..ee5ff44d8 100644
 --- a/lib/Dialect/TritonGPU/Transforms/ReduceDataDuplication.cpp
 +++ b/lib/Dialect/TritonGPU/Transforms/ReduceDataDuplication.cpp
-@@ -36,6 +36,10 @@ public:
+@@ -37,6 +37,10 @@ public:
        auto srcEncoding = srcType.getEncoding();
-       if (srcEncoding.isa<triton::gpu::SharedEncodingAttr>())
+       if (isa<triton::gpu::SharedEncodingAttr>(srcEncoding))
          return;
-+      if (dstType.getEncoding().isa<triton::gpu::SparseDotMetaEncodingAttr>()) {
++      if (isa<triton::gpu::SparseDotMetaEncodingAttr>(dstType.getEncoding())) {
 +        replaceSparseMetaEncoding(cvtOp);
 +        return;
 +      }
        auto dstDotOp =
-           dstType.getEncoding().dyn_cast<triton::gpu::DotOperandEncodingAttr>();
+           dyn_cast<triton::gpu::DotOperandEncodingAttr>(dstType.getEncoding());
        if (!dstDotOp)
-@@ -74,6 +78,27 @@ public:
+@@ -83,6 +87,27 @@ public:
        cvtOp.erase();
      });
    }
 +
 + private:
 +  void replaceSparseMetaEncoding(triton::gpu::ConvertLayoutOp cvtOp) {
-+    auto srcType = cvtOp.getOperand().getType().cast<RankedTensorType>();
++    auto srcType = cast<RankedTensorType>(cvtOp.getOperand().getType());
 +    auto srcEncoding = srcType.getEncoding();
 +    auto sharedLayout = triton::gpu::SharedEncodingAttr::get(
 +        cvtOp.getContext(), 8, 1, 1, triton::gpu::getOrder(srcEncoding),
 +        triton::gpu::getCTALayout(srcEncoding));
 +
-+    auto dstType = cvtOp.getType().cast<RankedTensorType>();
++    auto dstType = cast<RankedTensorType>(cvtOp.getType());
 +    auto tmpType = triton::MemDescType::get(
 +        dstType.getShape(), dstType.getElementType(), sharedLayout);
 +
@@ -421,6 +406,7 @@ diff --git a/lib/Dialect/TritonGPU/Transforms/ReduceDataDuplication.cpp b/lib/Di
  
  std::unique_ptr<Pass> mlir::triton::gpu::createReduceDataDuplicationPass() {
 diff --git a/lib/Dialect/TritonNvidiaGPU/Transforms/FenceInsertion.cpp b/lib/Dialect/TritonNvidiaGPU/Transforms/FenceInsertion.cpp
+index f456d36a6..a1dac2b72 100644
 --- a/lib/Dialect/TritonNvidiaGPU/Transforms/FenceInsertion.cpp
 +++ b/lib/Dialect/TritonNvidiaGPU/Transforms/FenceInsertion.cpp
 @@ -45,7 +45,7 @@ public:
@@ -432,7 +418,7 @@ diff --git a/lib/Dialect/TritonNvidiaGPU/Transforms/FenceInsertion.cpp b/lib/Dia
          return WalkResult::advance();
        OpBuilder builder(op);
        auto a = op->getOperand(0);
-@@ -83,7 +83,7 @@ private:
+@@ -80,7 +80,7 @@ private:
      static DenseSet<std::pair<Operation *, unsigned>> trace;
      auto op = operand.getDefiningOp();
      // avoid redundant insertion
diff --git a/third_party/xla/third_party/triton/temporary/pipelining.patch b/third_party/xla/third_party/triton/temporary/pipelining.patch
deleted file mode 100644
index 9f5f36aeb5099d..00000000000000
--- a/third_party/xla/third_party/triton/temporary/pipelining.patch
+++ /dev/null
@@ -1,472 +0,0 @@
-This is patching changes upstream from different PRs that fix issues with
-pipelining internally. Required changes are upto and including this commit
-https://github.com/openai/triton/commit/70f0b7b6e333fe2155c79dfa8bec6ad388073670
-The patch can be removed with the integration that includes these changes.
-
-diff --git a/include/triton/Analysis/Utility.h b/include/triton/Analysis/Utility.h
---- a/include/triton/Analysis/Utility.h
-+++ b/include/triton/Analysis/Utility.h
-@@ -8,6 +8,18 @@
- 
- namespace mlir {
- 
-+inline bool isZeroConst(Value v) {
-+  auto constantOp = v.getDefiningOp<arith::ConstantOp>();
-+  if (!constantOp)
-+    return false;
-+  if (auto denseAttr = dyn_cast<DenseFPElementsAttr>(constantOp.getValueAttr()))
-+    return denseAttr.isSplat() && denseAttr.getSplatValue<APFloat>().isZero();
-+  if (auto denseAttr =
-+          dyn_cast<DenseIntElementsAttr>(constantOp.getValueAttr()))
-+    return denseAttr.isSplat() && denseAttr.getSplatValue<APInt>().isZero();
-+  return false;
-+}
-+
- class ReduceOpHelper {
- public:
-   explicit ReduceOpHelper(triton::ReduceOp op)
-diff --git a/include/triton/Dialect/TritonGPU/IR/TritonGPUOps.td b/include/triton/Dialect/TritonGPU/IR/TritonGPUOps.td
---- a/include/triton/Dialect/TritonGPU/IR/TritonGPUOps.td
-+++ b/include/triton/Dialect/TritonGPU/IR/TritonGPUOps.td
-@@ -45,6 +45,8 @@ def TTG_AsyncWaitOp : TTG_Op<"async_wait
- 
-   let arguments = (ins Variadic<TTG_AsyncToken>:$asyncToken, I32Attr:$num);
- 
-+  let results = (outs TTG_AsyncToken:$retToken);
-+
-   let assemblyFormat = "$asyncToken attr-dict";
- 
-   let extraClassDeclaration = [{
-@@ -229,10 +231,16 @@ def TTG_LocalLoadOp : TTG_Op<"local_load
-   let description = [{
-     Load a tensor from the local memory descriptor into a distributed tensor.
-   }];
--  let arguments = (ins TT_MemDescType:$src);
-+  let arguments = (ins TT_MemDescType:$src, Optional<TTG_AsyncToken> :$token);
-+
-+  let builders = [
-+      OpBuilder<(ins "Type":$retType, "Value":$src),
-+      [{
-+      build($_builder, $_state, retType, src, /*token=*/static_cast<mlir::Value>(nullptr));
-+      }]>];
- 
-   // Use qualified() otherwise "!tt.memdesc<X>" is printed as "<X>".
--  let assemblyFormat = [{$src attr-dict `:` qualified(type($src)) `->` type($result)}];
-+  let assemblyFormat = [{$src (`token` $token^)? attr-dict `:` qualified(type($src)) `->` type($result)}];
- 
-   let results = (outs TT_Tensor:$result);
- }
-diff --git a/lib/Dialect/TritonGPU/Transforms/Pipeliner/MatmulLoopPipeline.cpp b/lib/Dialect/TritonGPU/Transforms/Pipeliner/MatmulLoopPipeline.cpp
---- a/lib/Dialect/TritonGPU/Transforms/Pipeliner/MatmulLoopPipeline.cpp
-+++ b/lib/Dialect/TritonGPU/Transforms/Pipeliner/MatmulLoopPipeline.cpp
-@@ -8,6 +8,7 @@
- #include "mlir/Interfaces/SideEffectInterfaces.h"
- #include "mlir/Support/LLVM.h"
- #include "triton/Analysis/AxisInfo.h"
-+#include "triton/Analysis/Utility.h"
- #include "triton/Dialect/Triton/IR/Types.h"
- #include "triton/Dialect/Triton/IR/Utility.h"
- #include "triton/Dialect/TritonGPU/IR/Attributes.h"
-@@ -84,12 +85,13 @@ createAsyncCopy(scf::ForOp &forOp, tt::L
-   Location loc = loadOp.getLoc();
-   Value src = loadOp.getPtr();
-   Value mask = loadOp.getMask();
-+  Value other = loadOp.getOther();
-   if (!isExpensiveLoadOrStore(loadOp) && opToInfo[loadOp].blockedEncoding) {
-     // For inexpensive loads that do not directly feed into dot ops
-     // we want to use optimal layout for the data.
-     ttg::BlockedEncodingAttr encoding = opToInfo[loadOp].blockedEncoding;
-     auto convertBlockLayout = [&](Value src) {
--      auto ty = src.getType().cast<RankedTensorType>();
-+      auto ty = cast<RankedTensorType>(src.getType());
-       auto newTy =
-           RankedTensorType::get(ty.getShape(), ty.getElementType(), encoding);
-       auto cvt =
-@@ -99,9 +101,11 @@ createAsyncCopy(scf::ForOp &forOp, tt::L
-     src = convertBlockLayout(src);
-     if (mask)
-       mask = convertBlockLayout(mask);
-+    if (other)
-+      other = convertBlockLayout(other);
-   }
- 
--  tt::MemDescType allocTy = alloc.getType().cast<tt::MemDescType>();
-+  tt::MemDescType allocTy = cast<tt::MemDescType>(alloc.getType());
-   SmallVector<Value> copyOffsets(allocTy.getRank(), zero);
-   copyOffsets[0] = insertIdx;
-   tt::MemDescType subviewTy = tt::MemDescType::get(
-@@ -110,11 +114,12 @@ createAsyncCopy(scf::ForOp &forOp, tt::L
-   auto view =
-       builder.create<ttg::MemDescSubviewOp>(loc, subviewTy, alloc, copyOffsets);
-   Operation *copy = builder.create<ttg::AsyncCopyGlobalToLocalOp>(
--      loc, src, view, mask, loadOp.getOther(), loadOp.getCache(),
--      loadOp.getEvict(), loadOp.getIsVolatile());
-+      loc, src, view, mask, other, loadOp.getCache(), loadOp.getEvict(),
-+      loadOp.getIsVolatile());
-   Operation *commmit =
-       builder.create<ttg::AsyncCommitGroupOp>(loc, copy->getResult(0));
--  builder.create<ttg::AsyncWaitOp>(loc, commmit->getResult(0), 0);
-+  Operation *wait =
-+      builder.create<ttg::AsyncWaitOp>(loc, commmit->getResult(0), 0);
- 
-   int stage = opToInfo[loadOp].stage;
-   bool isMMV3Load = opToInfo[loadOp].loadIsMMAV3;
-@@ -142,9 +147,21 @@ createAsyncCopy(scf::ForOp &forOp, tt::L
-     for (auto alloc : allocsToErase) {
-       alloc.erase();
-     }
--    auto sharedLoad =
--        builder.create<ttg::LocalLoadOp>(loc, loadOp.getType(), viewLoad);
--    loadOp->replaceAllUsesWith(sharedLoad->getResults());
-+
-+    auto sharedLoad = builder.create<ttg::LocalLoadOp>(
-+        loc, loadOp.getType(), viewLoad, wait->getResult(0));
-+    auto result = sharedLoad->getResults();
-+
-+    // Create a select for non-zero other values as they are not handled by
-+    // AsyncCopyGlobalToLocalOp for now.
-+    Value other = loadOp.getOther();
-+    if (other && !isZeroConst(other)) {
-+      auto select = builder.create<arith::SelectOp>(
-+          loc, loadOp.getType(), mask, sharedLoad.getResult(), other);
-+      result = select->getResults();
-+    }
-+
-+    loadOp->replaceAllUsesWith(result);
-   }
-   loadOp.erase();
- }
-@@ -160,7 +177,7 @@ getSharedEncIfAllUsersAreDotEnc(Value va
-     if (user->getNumResults() != 1)
-       return std::nullopt;
-     if (auto memDesc =
--            user->getResult(0).getType().dyn_cast<triton::MemDescType>()) {
-+            dyn_cast<triton::MemDescType>(user->getResult(0).getType())) {
-       // First time we find a shared encoding in the chain, save it and try to
-       // use it if it is compatible with the other users.
-       tempAttr = memDesc.getEncoding().cast<ttg::SharedEncodingAttr>();
-@@ -203,7 +220,7 @@ getSharedEncIfAllUsersAreDotEnc(Value va
- static ttg::BlockedEncodingAttr
- getBlockedEncoding(tt::LoadOp loadOp, tt::ModuleAxisInfoAnalysis &axisInfo) {
-   Value src = loadOp.getPtr();
--  auto ty = src.getType().cast<RankedTensorType>();
-+  auto ty = cast<RankedTensorType>(src.getType());
-   auto mod = loadOp->getParentOfType<ModuleOp>();
-   int numWarps = ttg::TritonGPUDialect::getNumWarps(mod);
-   int threadsPerWarp = ttg::TritonGPUDialect::getThreadsPerWarp(mod);
-@@ -221,7 +238,7 @@ getBlockedEncoding(tt::LoadOp loadOp, tt
- 
- static std::optional<ttg::SharedEncodingAttr>
- getSharedEncoding(tt::LoadOp loadOp, bool isMMAV3) {
--  auto ty = loadOp.getType().cast<RankedTensorType>();
-+  auto ty = cast<RankedTensorType>(loadOp.getType());
-   auto ctaLayout = ttg::getCTALayout(ty.getEncoding());
-   auto blockedOrder = ttg::getOrder(ty.getEncoding());
-   SmallVector<unsigned> order;
-@@ -285,11 +302,10 @@ loadOpsToDistanceAndUse(scf::ForOp forOp
-     if (auto mask = loadOp.getMask())
-       vec = std::min<unsigned>(vec, axisInfoAnalysis.getMaskAlignment(mask));
- 
--    auto tensorTy = ptr.getType().dyn_cast<RankedTensorType>();
-+    auto tensorTy = dyn_cast<RankedTensorType>(ptr.getType());
-     if (!tensorTy)
-       return false;
--    auto ty =
--        tensorTy.getElementType().cast<tt::PointerType>().getPointeeType();
-+    auto ty = cast<tt::PointerType>(tensorTy.getElementType()).getPointeeType();
-     unsigned width = vec * ty.getIntOrFloatBitWidth();
- 
-     // We do not pipeline all loads for the following reasons:
-@@ -353,7 +369,7 @@ static bool loadIsMMAv3(tt::LoadOp loadO
- 
-   // MMA V3 case.
-   auto newOrder = sharedEnc.getOrder();
--  auto ty = loadOp.getType().cast<RankedTensorType>();
-+  auto ty = cast<RankedTensorType>(loadOp.getType());
-   auto oldOrder = ttg::getOrder(ty.getEncoding());
- 
-   // The operand of MMAv3 is in SharedEncoding and its order should not
-@@ -497,7 +513,7 @@ collectOpsToPipeline(scf::ForOp forOp,
- static Value createAlloc(scf::ForOp &forOp, tt::LoadOp loadOp,
-                          ttg::SharedEncodingAttr sharedEnc, unsigned distance) {
-   OpBuilder builder(forOp);
--  auto ty = loadOp.getType().cast<RankedTensorType>();
-+  auto ty = cast<RankedTensorType>(loadOp.getType());
-   SmallVector<int64_t> bufferShape(ty.getShape().begin(), ty.getShape().end());
-   bufferShape.insert(bufferShape.begin(), distance);
-   Type memdescType = mlir::triton::MemDescType::get(
-@@ -669,12 +685,23 @@ createSchedule(scf::ForOp forOp, int num
-     }
-   });
- 
-+  auto getNestedOperands = [](Operation *op) -> SmallVector<Value> {
-+    SmallVector<Value> operands;
-+    op->walk([&](Operation *nestedOp) {
-+      for (Value operand : nestedOp->getOperands()) {
-+        if (operand.getParentBlock()->getParentOp()->isAncestor(nestedOp))
-+          operands.push_back(operand);
-+      }
-+    });
-+    return operands;
-+  };
-+
-   // Find dependencies with distance of 1.
-   SmallVector<DenseSet<Operation *>> distanceOneUsers(numStages);
-   for (int stage = 0; stage < numStages - 1; stage++) {
-     auto &group = insertAndDeps[stage];
-     for (Operation *op : group) {
--      for (Value operand : op->getOperands()) {
-+      for (Value operand : getNestedOperands(op)) {
-         if (auto arg = operand.dyn_cast<BlockArgument>()) {
-           if (arg.getArgNumber() > 0 && arg.getOwner() == op->getBlock()) {
-             auto yieldOp = op->getBlock()->getTerminator();
-@@ -905,7 +932,7 @@ static int minNumInterleavedCommitOps(Op
- // Look for consecutive wait ops and combine them into a single wait op.
- static void
- combineRedundantWaitOps(llvm::SmallSetVector<ttg::AsyncWaitOp, 8> &waitOps) {
--  llvm::SmallSetVector<ttg::AsyncWaitOp, 8> toDelete;
-+  llvm::MapVector<ttg::AsyncWaitOp, ttg::AsyncWaitOp> toDelete;
-   for (auto waitOp : waitOps) {
-     if (toDelete.count(waitOp))
-       continue;
-@@ -927,10 +954,13 @@ combineRedundantWaitOps(llvm::SmallSetVe
-     OpBuilder builder(waitGroup.back());
-     auto newWaitOp = builder.create<ttg::AsyncWaitOp>(waitOp.getLoc(),
-                                                       depTokens, minWaitNumber);
--    toDelete.insert(waitGroup.begin(), waitGroup.end());
-+    for (auto waitOp : waitGroup) {
-+      toDelete[waitOp] = newWaitOp;
-+    }
-   }
-   for (auto waitOp : toDelete) {
--    waitOp->erase();
-+    waitOp.first->replaceAllUsesWith(waitOp.second);
-+    waitOp.first->erase();
-   }
- }
- 
-@@ -1010,7 +1040,7 @@ static void threadValuesThroughWait(ttng
- 
-   for (ttng::DotAsyncOp dot : asyncDots) {
-     for (Value operand : dot.getOperands()) {
--      if (operand.getType().isa<tt::MemDescType>()) {
-+      if (isa<tt::MemDescType>(operand.getType())) {
-         newOperands.insert(operand);
-       }
-     }
-@@ -1020,15 +1050,21 @@ static void threadValuesThroughWait(ttng
-   // values in the operation.
-   auto newWait = builder.create<ttng::DotWaitOp>(
-       wait.getLoc(), llvm::to_vector(newOperands), wait.getPendings());
-+
-+  auto dominatedByNewWait = [&](OpOperand &operand) {
-+    auto opInThisBlock =
-+        newWait->getBlock()->findAncestorOpInBlock(*operand.getOwner());
-+    return opInThisBlock && newWait->isBeforeInBlock(opInThisBlock);
-+  };
-   for (int i = 0; i < origNumOperands; i++) {
-     Value operand = wait.getResult(i);
--    if (!operand.getType().isa<tt::MemDescType>())
-+    if (!isa<tt::MemDescType>(operand.getType()))
-       operand.replaceAllUsesWith(newWait.getResult(i));
-   }
-   for (int i = origNumOperands; i < newOperands.size(); i++) {
-     Value operand = newWait.getOperand(i);
--    if (!operand.getType().isa<tt::MemDescType>())
--      operand.replaceAllUsesExcept(newWait.getResult(i), newWait);
-+    if (!isa<tt::MemDescType>(operand.getType()))
-+      operand.replaceUsesWithIf(newWait.getResult(i), dominatedByNewWait);
-   }
-   wait->erase();
- }
-@@ -1047,8 +1083,8 @@ static void threadValuesThroughWait(ttng
- //  1. All operands that touch shared memory are multi-buffered, i.e. can't read
- //     an incomplete value while it's being written asynchronously by a load.
- //
--//  2. During iteration i, nothing other than the loop's `yield` reads the
--//     result of the dot.
-+//  2. If the dot is used by any op in the loop, it must be used under an `if`,
-+//     and will be synced with a `wait 0` at the beginning of the `if` block.
- //
- //  3. During iteration i, between the start of the loop up until the first
- //     `ttng.dot_wait {pendings=0}` op, the result of the dot from iteration i-1
-@@ -1079,7 +1115,7 @@ static std::optional<int> dotCanBeProper
-   // Rule 1: All shmem operands are multi-buffered.
-   auto checkOperand = [&](Value operand) {
-     if (!isa<ttg::SharedEncodingAttr>(
--            operand.getType().cast<TensorOrMemDesc>().getEncoding())) {
-+            cast<TensorOrMemDesc>(operand.getType()).getEncoding())) {
-       return true;
-     }
- 
-@@ -1103,17 +1139,41 @@ static std::optional<int> dotCanBeProper
-     return std::nullopt;
-   }
- 
--  // Rule 2: The dot should only be used by the for loop's `yield`.
--  if (!dotOp->hasOneUse() ||
--      *dotOp->getUsers().begin() != forOp.getBody()->getTerminator()) {
--    LDBG("Can't make dot async because it is not used only by the loop's "
--         "`yield`.");
--    return std::nullopt;
-+  // Rule 2: The dot cannot be unconditionally used by any op in the loop.
-+  // Uses under `if` are allowed, as can be explicitly synced with a `wait 0`.
-+  int iterArgIdx = -1;
-+  Value iterArg = nullptr;
-+  SmallVector<std::pair<Operation *, int>> queue;
-+  for (auto &use : dotOp->getUses()) {
-+    queue.push_back({use.getOwner(), use.getOperandNumber()});
-   }
--
--  // The result of the dot becomes this loop carry value.
--  auto iterArgIdx = dotOp->getUses().begin()->getOperandNumber();
--  auto iterArg = forOp.getRegionIterArg(iterArgIdx);
-+  while (!queue.empty()) {
-+    auto [user, argIdx] = queue.pop_back_val();
-+    if (user->getParentOp() == forOp) {
-+      if (isa<scf::YieldOp>(user)) {
-+        if (iterArg) {
-+          // The dot is used by the loop's yield, but we can't have any other
-+          // uses.
-+          return std::nullopt;
-+        }
-+        iterArgIdx = argIdx;
-+        iterArg = forOp.getRegionIterArg(argIdx);
-+        continue;
-+      }
-+      return std::nullopt;
-+    }
-+    if (auto ifOp = dyn_cast<scf::IfOp>(user->getParentOp())) {
-+      if (isa<scf::YieldOp>(user)) {
-+        // The result is returned by the if, follow it further.
-+        auto uses = ifOp.getResult(argIdx).getUses();
-+        for (auto &use : uses) {
-+          queue.push_back({use.getOwner(), use.getOperandNumber()});
-+        }
-+      }
-+    } else {
-+      return std::nullopt;
-+    }
-+  }
- 
-   // Rule 3a: Are the only users of the dot's result from iteration i-1 other
-   // MMAv3 dots?  If so, we're done, this dot can be properly async.
-@@ -1181,6 +1241,32 @@ static void insertAsyncDotWaitInLoop(
-     return;
-   }
- 
-+  // Insert waits before the users of the properly async dots other than loop
-+  // yield.
-+  for (auto [asyncDot, iterArgIdx] : properlyAsyncDots) {
-+    SmallVector<OpOperand *> uses;
-+    for (auto &use : asyncDot->getUses()) {
-+      if (auto yieldOp = dyn_cast<scf::YieldOp>(use.getOwner())) {
-+        continue;
-+      }
-+      uses.push_back(&use);
-+    }
-+
-+    DenseMap<Block *, SmallVector<Value>> blockToUsers;
-+    for (auto use : uses) {
-+      auto block = use->getOwner()->getBlock();
-+      blockToUsers[block].push_back(use->get());
-+    }
-+
-+    for (auto [block, users] : blockToUsers) {
-+      OpBuilder builder(block, block->begin());
-+      auto newWait = builder.create<ttng::DotWaitOp>(asyncDot->getLoc(),
-+                                                     ArrayRef<Value>{}, 0);
-+
-+      threadValuesThroughWait(newWait, users);
-+    }
-+  }
-+
-   // Add the wait right after the last properly-async dot.  This only needs to
-   // wait for all properly-async dots from the i-1'th iteration to complete, IOW
-   // we wait until there are most `asyncDots.size()` dots in flight.
-diff --git a/test/TritonGPU/loop-pipeline.mlir b/test/TritonGPU/loop-pipeline.mlir
---- a/test/TritonGPU/loop-pipeline.mlir
-+++ b/test/TritonGPU/loop-pipeline.mlir
-@@ -349,16 +349,21 @@ tt.func @indirect_bmm_scalar_dist_one(%7
- // CHECK: triton_gpu.async_copy_global_to_local
- // CHECK: triton_gpu.async_copy_global_to_local
- // CHECK: triton_gpu.async_commit_group
-+// CHECK: triton_gpu.async_wait {{.*}} {num = 1 : i32}
-+// CHECK: scf.for
-+// CHECK: tt.dot
- // CHECK: %[[NEXT_BUFFER_1:.*]] = tt.addptr %{{.*}}, {{.*}}
- // CHECK: triton_gpu.async_copy_global_to_local %[[NEXT_BUFFER_1]]
--// CHECK: %[[IND_BUFFER_0:.*]] = triton_gpu.memdesc_subview
--// CHECK: %[[IND_BUFFER_1:.*]] = triton_gpu.local_load %[[IND_BUFFER_0]]
-+// CHECK-DAG: %[[IND_BUFFER_WAIT_TOKEN:.*]] = triton_gpu.async_wait {{.*}} {num = 1 : i32}
-+// CHECK-DAG: %[[IND_BUFFER_0:.*]] = triton_gpu.memdesc_subview
-+// CHECK: %[[IND_BUFFER_1:.*]] = triton_gpu.local_load %[[IND_BUFFER_0]] token %[[IND_BUFFER_WAIT_TOKEN]]
- // CHECK: %[[IND_BUFFER_2:.*]] = tt.expand_dims %[[IND_BUFFER_1]] {axis = 1 : i32}
- // CHECK: %[[IND_BUFFER_3:.*]] = tt.broadcast %[[IND_BUFFER_2]]
- // CHECK: %[[IND_BUFFER_4:.*]] = arith.muli {{.*}}, %[[IND_BUFFER_3]]
- // CHECK: %[[NEXT_BUFFER_0:.*]] = tt.addptr {{.*}}, %[[IND_BUFFER_4]]
- // CHECK: triton_gpu.async_copy_global_to_local %[[NEXT_BUFFER_0]]
- // CHECK: triton_gpu.async_wait {{.*}} {num = 1 : i32}
-+// CHECK: scf.yield
- tt.func @indirect_bmm_vector(%77: tensor<16x16xi64, #BL> {tt.divisibility=16: i32, tt.constancy=16: i32},
-                    %76: index,
-                    %49: tensor<16x16x!tt.ptr<f16>, #AL> {tt.divisibility=16: i32, tt.contiguity=2 : i32},
-diff --git a/test/TritonGPU/reorder-instructions.mlir b/test/TritonGPU/reorder-instructions.mlir
---- a/test/TritonGPU/reorder-instructions.mlir
-+++ b/test/TritonGPU/reorder-instructions.mlir
-@@ -28,7 +28,7 @@ module attributes {"triton_gpu.num-warps
- //       CHECK: triton_gpu.async_wait {num = 0 : i32}
- //       CHECK: triton_gpu.local_dealloc %0 : !tt.memdesc<4x128x64xf16, #shared>
- //       CHECK: triton_gpu.local_dealloc %1 : !tt.memdesc<4x128x64xf16, #shared>
--//       CHECK: %2 = triton_gpu.convert_layout %arg0 : tensor<32x32xf32, #blocked> -> tensor<32x32xf32, #blocked1>
-+//       CHECK: %3 = triton_gpu.convert_layout %arg0 : tensor<32x32xf32, #blocked> -> tensor<32x32xf32, #blocked1>
- #blocked = #triton_gpu.blocked<{sizePerThread = [1, 1], threadsPerWarp = [32, 1], warpsPerCTA = [1, 4], order = [0, 1]}>
- #blocked1 = #triton_gpu.blocked<{sizePerThread = [1, 1], threadsPerWarp = [32, 1], warpsPerCTA = [1, 4], order = [1, 0]}>
- #shared = #triton_gpu.shared<{vec = 8, perPhase = 1, maxPhase = 4, order = [0, 1]}>
-diff --git a/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/DotOpToLLVM/WGMMA.cpp b/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/DotOpToLLVM/WGMMA.cpp
---- a/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/DotOpToLLVM/WGMMA.cpp
-+++ b/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/DotOpToLLVM/WGMMA.cpp
-@@ -333,17 +333,6 @@ static Value faddAccumulate(ConversionPa
-   return newStruct;
- }
- 
--static bool isZero(Value v) {
--  auto constantOp = v.getDefiningOp<arith::ConstantOp>();
--  if (!constantOp)
--    return false;
--  if (auto denseAttr = dyn_cast<DenseFPElementsAttr>(constantOp.getValueAttr()))
--    return denseAttr.isSplat() && denseAttr.getSplatValue<APFloat>().isZero();
--  if (auto denseAttr =
--          dyn_cast<DenseIntElementsAttr>(constantOp.getValueAttr()))
--    return denseAttr.isSplat() && denseAttr.getSplatValue<APInt>().isZero();
--  return false;
--}
- 
- static SmallVector<Value> emitWait(ConversionPatternRewriter &rewriter,
-                                    Location loc, SmallVector<Value> acc,
-@@ -402,7 +391,7 @@ LogicalResult convertDot(const LLVMTypeC
-   int M = 4 * instrShape[0];
-   int N = instrShape[1];
-   int K = instrShape[2];
--  bool zeroAcc = isZero(c);
-+  bool zeroAcc = isZeroConst(c);
-   auto shapePerCTATile = getShapePerCTATile(mmaEncoding);
-   int numRepM = ceil<unsigned>(dShapePerCTA[0], shapePerCTATile[0]);
-   int numRepN = ceil<unsigned>(dShapePerCTA[1], shapePerCTATile[1]);
-diff --git a/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/LoadStoreOpToLLVM.cpp b/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/LoadStoreOpToLLVM.cpp
---- a/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/LoadStoreOpToLLVM.cpp
-+++ b/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/LoadStoreOpToLLVM.cpp
-@@ -924,8 +924,11 @@ struct AsyncWaitOpConversion
-     auto voidTy = void_ty(ctx);
-     ptxBuilder.launch(rewriter, loc, voidTy);
- 
--    // Safe to remove the op since it doesn't have any return value.
--    rewriter.eraseOp(op);
-+    // Drop the result token.
-+    Value zero = rewriter.create<LLVM::ConstantOp>(
-+        op.getLoc(), IntegerType::get(op.getContext(), 32),
-+        rewriter.getI32IntegerAttr(0));
-+    rewriter.replaceOp(op, zero);
-     return success();
-   }
- };
diff --git a/third_party/xla/third_party/triton/temporary/series.bzl b/third_party/xla/third_party/triton/temporary/series.bzl
index 70313c9b436d2a..214666767949de 100644
--- a/third_party/xla/third_party/triton/temporary/series.bzl
+++ b/third_party/xla/third_party/triton/temporary/series.bzl
@@ -6,7 +6,5 @@ internal patch during the next triton integration process.
 """
 
 temporary_patch_list = [
-    "//third_party/triton/temporary:pipelining.patch",
-    "//third_party/triton/temporary:support_ceil_op.patch",
     "//third_party/triton/temporary:mma_limit_pred.patch",
 ]
diff --git a/third_party/xla/third_party/triton/temporary/support_ceil_op.patch b/third_party/xla/third_party/triton/temporary/support_ceil_op.patch
deleted file mode 100644
index 71b323d9fccdca..00000000000000
--- a/third_party/xla/third_party/triton/temporary/support_ceil_op.patch
+++ /dev/null
@@ -1,138 +0,0 @@
-Cherry-picking https://github.com/openai/triton/commit/62706e8c518c8c56e56460a43732d8e375217860
-until the next integration lands it. Can be removed as it is already merged.
-
-diff --git a/lib/Conversion/TritonGPUToLLVM/ElementwiseOpToLLVM.cpp b/lib/Conversion/TritonGPUToLLVM/ElementwiseOpToLLVM.cpp
---- a/lib/Conversion/TritonGPUToLLVM/ElementwiseOpToLLVM.cpp
-+++ b/lib/Conversion/TritonGPUToLLVM/ElementwiseOpToLLVM.cpp
-@@ -805,6 +805,7 @@ void mlir::triton::populateElementwiseOp
-   POPULATE_UNARY_OP(arith::FPToUIOp, LLVM::FPToUIOp)
-   POPULATE_UNARY_OP(arith::UIToFPOp, LLVM::UIToFPOp)
-   POPULATE_UNARY_OP(math::FloorOp, math::FloorOp)
-+  POPULATE_UNARY_OP(math::CeilOp, math::CeilOp)
-   POPULATE_UNARY_OP(math::LogOp, math::LogOp)
-   POPULATE_UNARY_OP(math::Log2Op, math::Log2Op)
-   POPULATE_UNARY_OP(math::CosOp, math::CosOp)
-diff --git a/lib/Conversion/TritonToTritonGPU/TritonToTritonGPUPass.cpp b/lib/Conversion/TritonToTritonGPU/TritonToTritonGPUPass.cpp
---- a/lib/Conversion/TritonToTritonGPU/TritonToTritonGPUPass.cpp
-+++ b/lib/Conversion/TritonToTritonGPU/TritonToTritonGPUPass.cpp
-@@ -125,12 +125,13 @@ void populateMathPatternsAndLegality(Tri
-   MLIRContext *context = patterns.getContext();
-   // Rewrite rule
-   patterns.add<GenericOpPattern<math::ExpOp>, GenericOpPattern<math::Exp2Op>,
--               GenericOpPattern<math::FloorOp>, GenericOpPattern<math::CosOp>,
--               GenericOpPattern<math::SinOp>, GenericOpPattern<math::LogOp>,
--               GenericOpPattern<math::Log2Op>, GenericOpPattern<math::ErfOp>,
--               GenericOpPattern<math::AbsFOp>, GenericOpPattern<math::AbsIOp>,
--               GenericOpPattern<math::SqrtOp>, GenericOpPattern<math::RsqrtOp>,
--               GenericOpPattern<math::FmaOp>>(typeConverter, context);
-+               GenericOpPattern<math::FloorOp>, GenericOpPattern<math::CeilOp>,
-+               GenericOpPattern<math::CosOp>, GenericOpPattern<math::SinOp>,
-+               GenericOpPattern<math::LogOp>, GenericOpPattern<math::Log2Op>,
-+               GenericOpPattern<math::ErfOp>, GenericOpPattern<math::AbsFOp>,
-+               GenericOpPattern<math::AbsIOp>, GenericOpPattern<math::SqrtOp>,
-+               GenericOpPattern<math::RsqrtOp>, GenericOpPattern<math::FmaOp>>(
-+      typeConverter, context);
- }
- 
- //
-diff --git a/lib/Dialect/TritonNvidiaGPU/Transforms/PlanCTA.cpp b/lib/Dialect/TritonNvidiaGPU/Transforms/PlanCTA.cpp
---- a/lib/Dialect/TritonNvidiaGPU/Transforms/PlanCTA.cpp
-+++ b/lib/Dialect/TritonNvidiaGPU/Transforms/PlanCTA.cpp
-@@ -651,10 +651,9 @@ bool CTAPlanner::isElementwiseOp(Operati
-                 math::CeilOp, math::CopySignOp, math::CosOp, math::SinOp,
-                 math::CountLeadingZerosOp, math::CountTrailingZerosOp,
-                 math::CtPopOp, math::ErfOp, math::ExpOp, math::Exp2Op,
--                math::FloorOp, math::ExpM1Op, math::FloorOp, math::FmaOp,
--                math::LogOp, math::Log10Op, math::Log1pOp, math::Log2Op,
--                math::PowFOp, math::RsqrtOp, math::SqrtOp, math::RsqrtOp,
--                math::TanhOp>(op))
-+                math::FloorOp, math::ExpM1Op, math::FmaOp, math::LogOp,
-+                math::Log10Op, math::Log1pOp, math::Log2Op, math::PowFOp,
-+                math::RsqrtOp, math::SqrtOp, math::RsqrtOp, math::TanhOp>(op))
-     return true;
-   if (llvm::isa<triton::IntToPtrOp, triton::PtrToIntOp, triton::BitcastOp,
-                 triton::FpToFpOp, triton::AddPtrOp, triton::PreciseSqrtOp,
-diff --git a/python/src/ir.cc b/python/src/ir.cc
---- a/python/src/ir.cc
-+++ b/python/src/ir.cc
-@@ -1379,6 +1379,10 @@ void init_triton_ir(py::module &&m) {
-            [](TritonOpBuilder &self, Value &val) -> Value {
-              return self.create<math::FloorOp>(val);
-            })
-+      .def("create_ceil",
-+           [](TritonOpBuilder &self, Value &val) -> Value {
-+             return self.create<math::CeilOp>(val);
-+           })
-       .def("create_exp",
-            [](TritonOpBuilder &self, Value &val) -> Value {
-              return self.create<math::ExpOp>(val);
-diff --git a/python/test/unit/language/test_core.py b/python/test/unit/language/test_core.py
---- a/python/test/unit/language/test_core.py
-+++ b/python/test/unit/language/test_core.py
-@@ -915,10 +915,11 @@ def test_unary_op(dtype_x, expr, num_cta
- 
- 
- @pytest.mark.interpreter
--@pytest.mark.parametrize("dtype_x, expr, x", [(dtype_x, expr, x)
--                                              for dtype_x in ["float32", "float64"]
--                                              for expr in ['exp', 'log', 'cos', 'sin', 'exp2', 'log2', 'sqrt', 'floor']
--                                              for x in ['x', '3.0']])
-+@pytest.mark.parametrize("dtype_x, expr, x",
-+                         [(dtype_x, expr, x)
-+                          for dtype_x in ["float32", "float64"]
-+                          for expr in ['exp', 'log', 'cos', 'sin', 'exp2', 'log2', 'sqrt', 'floor', 'ceil']
-+                          for x in ['x', '3.0']])
- def test_math_op(dtype_x, expr, x, device):
-     _test_unary(dtype_x, f'tl.{expr}({x})', f'np.{expr}({x}) ', device=device)
- 
-diff --git a/python/triton/language/__init__.py b/python/triton/language/__init__.py
---- a/python/triton/language/__init__.py
-+++ b/python/triton/language/__init__.py
-@@ -102,7 +102,8 @@ from .core import (
-     void,
-     where,
- )
--from .math import (umulhi, exp, exp2, fma, log, log2, cos, rsqrt, sin, sqrt, sqrt_rn, abs, fdiv, div_rn, erf, floor)
-+from .math import (umulhi, exp, exp2, fma, log, log2, cos, rsqrt, sin, sqrt, sqrt_rn, abs, fdiv, div_rn, erf, floor,
-+                   ceil)
- from .random import (
-     pair_uniform_to_normal,
-     philox,
-@@ -142,6 +143,7 @@ from .random import (
-     "builtin",
-     "cat",
-     "cdiv",
-+    "ceil",
-     "clamp",
-     "const",
-     "const_pointer_type",
-diff --git a/python/triton/language/math.py b/python/triton/language/math.py
---- a/python/triton/language/math.py
-+++ b/python/triton/language/math.py
-@@ -230,6 +230,15 @@ def floor(x, _builder=None):
- 
- 
- @core.builtin
-+@_check_dtype(dtypes=["fp32", "fp64"])
-+@_add_math_1arg_docstr("ceil")
-+@core._tensor_member_fn
-+def ceil(x, _builder=None):
-+    x = core._to_tensor(x, _builder)
-+    return core.tensor(_builder.create_ceil(x.handle), x.type)
-+
-+
-+@core.builtin
- @_add_math_3arg_docstr("fused multiply-add")
- def fma(x, y, z, _builder=None):
-     x = core._to_tensor(x, _builder)
-diff --git a/python/triton/runtime/interpreter.py b/python/triton/runtime/interpreter.py
---- a/python/triton/runtime/interpreter.py
-+++ b/python/triton/runtime/interpreter.py
-@@ -391,6 +391,7 @@ class InterpreterBuilder:
-     create_fabs = lambda self, arg: self.unary_op(arg, np.abs)
-     create_iabs = lambda self, arg: self.unary_op(arg, np.abs)
-     create_floor = lambda self, arg: self.unary_op(arg, np.floor)
-+    create_ceil = lambda self, arg: self.unary_op(arg, np.ceil)
-     create_log = lambda self, arg: self.unary_op(arg, np.log)
-     create_log2 = lambda self, arg: self.unary_op(arg, np.log2)
-     create_precise_sqrt = lambda self, arg: self.unary_op(arg, np.sqrt)
diff --git a/third_party/xla/third_party/triton/workspace.bzl b/third_party/xla/third_party/triton/workspace.bzl
index 45daf7974a022e..933b815253ffcd 100644
--- a/third_party/xla/third_party/triton/workspace.bzl
+++ b/third_party/xla/third_party/triton/workspace.bzl
@@ -8,8 +8,8 @@ load("//third_party/triton/xla_extensions:series.bzl", "extensions_files_patch_l
 def repo():
     """Imports Triton."""
 
-    TRITON_COMMIT = "cl623533461"
-    TRITON_SHA256 = "7aa74e82e4417a91fc7a7a84b4f6ad2b7e4e58512758d6c78ca3cd1c8771326b"
+    TRITON_COMMIT = "cl632952902"
+    TRITON_SHA256 = "f83c4f197cc2ae3b034070ec8189007451093edc445d1e383eb17a1e3808de9a"
     tf_http_archive(
         name = "triton",
         sha256 = TRITON_SHA256,
diff --git a/third_party/xla/third_party/triton/xla_extensions/env_vars.patch b/third_party/xla/third_party/triton/xla_extensions/env_vars.patch
deleted file mode 100644
index 955eb6db8da68e..00000000000000
--- a/third_party/xla/third_party/triton/xla_extensions/env_vars.patch
+++ /dev/null
@@ -1,14 +0,0 @@
-Long standing patch due to licensing issues.
-diff --git a/include/triton/Tools/Sys/GetEnv.hpp b/include/triton/Tools/Sys/GetEnv.hpp
-index 31bc03fe1..a19a432df 100644
---- a/include/triton/Tools/Sys/GetEnv.hpp
-+++ b/include/triton/Tools/Sys/GetEnv.hpp
-@@ -34,7 +34,7 @@ inline const std::set<std::string> ENV_VARS = {
-     "AMDGCN_ENABLE_DUMP",
-     "DISABLE_FAST_REDUCTION",
-     "DISABLE_LLVM_OPT",
--    "DISABLE_MMA_V3",
-+    "ENABLE_MMA_V3",
-     "DISABLE_PTXAS_OPT",
-     "LLVM_IR_ENABLE_DUMP",
-     "MLIR_ENABLE_DUMP",
diff --git a/third_party/xla/third_party/triton/xla_extensions/series.bzl b/third_party/xla/third_party/triton/xla_extensions/series.bzl
index b858da203fb094..af524fb253cbef 100644
--- a/third_party/xla/third_party/triton/xla_extensions/series.bzl
+++ b/third_party/xla/third_party/triton/xla_extensions/series.bzl
@@ -4,7 +4,6 @@ applied in the previous copybara workflow.
 """
 
 extensions_files_patch_list = [
-    "//third_party/triton/xla_extensions:env_vars.patch",  # File not exported to google
     "//third_party/triton/xla_extensions:sparse_dot_nvgpu.patch",  # Sparsity internal patch
     "//third_party/triton/xla_extensions:sparse_dot_base.patch",  # Sparsity internal patch
     "//third_party/triton/xla_extensions:sparse_dot_passes.patch",  # Sparsity internal patch
diff --git a/third_party/xla/third_party/triton/xla_extensions/sparse_dot_base.patch b/third_party/xla/third_party/triton/xla_extensions/sparse_dot_base.patch
index dcacd99740b18f..08b7dd6f7ada87 100644
--- a/third_party/xla/third_party/triton/xla_extensions/sparse_dot_base.patch
+++ b/third_party/xla/third_party/triton/xla_extensions/sparse_dot_base.patch
@@ -1,8 +1,9 @@
 diff --git a/include/triton/Dialect/TritonGPU/IR/TritonGPUAttrDefs.td b/include/triton/Dialect/TritonGPU/IR/TritonGPUAttrDefs.td
+index 56f0b6b49..aa91ea9b8 100644
 --- a/include/triton/Dialect/TritonGPU/IR/TritonGPUAttrDefs.td
 +++ b/include/triton/Dialect/TritonGPU/IR/TritonGPUAttrDefs.td
-@@ -1158,4 +1158,12 @@ section 9.7.13.4.1 for more details.
-   let extraClassDeclaration = extraDistributedDeclaration;
+@@ -1262,4 +1262,16 @@ section 9.7.13.4.1 for more details.
+   }];
  }
  
 +def SparseDotMetaEncodingAttr : DistributedEncoding<"SparseDotMetaEncoding", "sparse_dot_meta_encoding"> {
@@ -10,14 +11,19 @@ diff --git a/include/triton/Dialect/TritonGPU/IR/TritonGPUAttrDefs.td b/include/
 +
 +  let parameters = (ins "Attribute":$parent);
 +  let assemblyFormat = "`<``{` struct(params) `}``>`";
-+  let extraClassDeclaration = extraDistributedDeclaration;
++  let extraClassDeclaration = extraDistributedDeclaration # [{
++    SmallVector<unsigned> getContigPerThread() {
++      return getSizePerThread();
++    };
++  }];
 +}
 +
  #endif
 diff --git a/include/triton/Dialect/TritonGPU/IR/TritonGPUOps.td b/include/triton/Dialect/TritonGPU/IR/TritonGPUOps.td
+index 4966a5f73..d2bb33cfa 100644
 --- a/include/triton/Dialect/TritonGPU/IR/TritonGPUOps.td
 +++ b/include/triton/Dialect/TritonGPU/IR/TritonGPUOps.td
-@@ -7,6 +7,7 @@ include "triton/Dialect/TritonGPU/IR/Tri
+@@ -7,6 +7,7 @@ include "triton/Dialect/TritonGPU/IR/TritonGPUAttrDefs.td"
  include "mlir/Dialect/Arith/IR/ArithBase.td"
  include "triton/Dialect/Triton/IR/TritonTypes.td"
  include "triton/Dialect/Triton/IR/TritonAttrDefs.td"
@@ -25,8 +31,8 @@ diff --git a/include/triton/Dialect/TritonGPU/IR/TritonGPUOps.td b/include/trito
  include "mlir/IR/OpBase.td"
  include "mlir/Interfaces/SideEffectInterfaces.td" // Pure
  include "mlir/Interfaces/InferTypeOpInterface.td" // SameOperandsAndResultType
-@@ -214,4 +215,19 @@ def TTG_LocalLoadOp : TTG_Op<"local_load
-   let results = (outs TT_Tensor:$result);
+@@ -232,4 +233,19 @@ def TTG_LocalStoreOp : TTG_Op<"local_store", [MemoryEffects<[MemWrite<SharedMemo
+   }];
  }
  
 +def TTNG_SparseDotOp : TTG_Op<"sparse_dot", [
@@ -46,9 +52,10 @@ diff --git a/include/triton/Dialect/TritonGPU/IR/TritonGPUOps.td b/include/trito
 +
  #endif
 diff --git a/lib/Dialect/TritonGPU/IR/Dialect.cpp b/lib/Dialect/TritonGPU/IR/Dialect.cpp
+index 0ce7ecf18..3736a1551 100644
 --- a/lib/Dialect/TritonGPU/IR/Dialect.cpp
 +++ b/lib/Dialect/TritonGPU/IR/Dialect.cpp
-@@ -479,6 +479,119 @@ getDefaultBlockedEncoding(MLIRContext *c
+@@ -483,6 +483,119 @@ getDefaultBlockedEncoding(MLIRContext *context, ArrayRef<int64_t> shape,
    return encoding;
  }
  
@@ -70,7 +77,7 @@ diff --git a/lib/Dialect/TritonGPU/IR/Dialect.cpp b/lib/Dialect/TritonGPU/IR/Dia
 +
 +LogicalResult SparseDotOp::verify() {
 +  // Verify operand A.
-+  auto aTensorTy = getOperand(0).getType().cast<TensorOrMemDesc>();
++  auto aTensorTy = cast<TensorOrMemDesc>(getOperand(0).getType());
 +  auto aElemTy = aTensorTy.getElementType();
 +  if (!aElemTy.isF16() && !aElemTy.isBF16())
 +    return emitError("element type of operand A is not supported");
@@ -78,7 +85,7 @@ diff --git a/lib/Dialect/TritonGPU/IR/Dialect.cpp b/lib/Dialect/TritonGPU/IR/Dia
 +  if (aShape.size() != 2) return emitError("shape of operand A is incorrect");
 +
 +  // Verify operand B.
-+  auto bTensorTy = getOperand(1).getType().cast<TensorOrMemDesc>();
++  auto bTensorTy = cast<TensorOrMemDesc>(getOperand(1).getType());
 +  auto bElemTy = bTensorTy.getElementType();
 +  if (!bElemTy.isF16() && !bElemTy.isBF16())
 +    return emitError("element type of operand B is not supported");
@@ -86,7 +93,7 @@ diff --git a/lib/Dialect/TritonGPU/IR/Dialect.cpp b/lib/Dialect/TritonGPU/IR/Dia
 +  if (bShape.size() != 2) return emitError("shape of operand B is incorrect");
 +
 +  // Verify operand C.
-+  auto cTensorTy = getOperand(2).getType().cast<RankedTensorType>();
++  auto cTensorTy = cast<RankedTensorType>(getOperand(2).getType());
 +  auto cElemTy = cTensorTy.getElementType();
 +  if (!cElemTy.isF32())
 +    return emitError("element type of operand C is not supported");
@@ -101,7 +108,7 @@ diff --git a/lib/Dialect/TritonGPU/IR/Dialect.cpp b/lib/Dialect/TritonGPU/IR/Dia
 +    return emitError("operand element types do not match");
 +
 +  // Verify sparse metadata.
-+  auto metaTy = getOperand(3).getType().cast<RankedTensorType>();
++  auto metaTy = cast<RankedTensorType>(getOperand(3).getType());
 +  auto metaShape = metaTy.getShape();
 +  if (!metaTy.getElementType().isInteger(16) || metaShape.size() != 2)
 +    return emitError("sparse metadata tensor is invalid");
@@ -125,7 +132,7 @@ diff --git a/lib/Dialect/TritonGPU/IR/Dialect.cpp b/lib/Dialect/TritonGPU/IR/Dia
 +//--- SparseDotMetaEncodingAttr ---
 +unsigned SparseDotMetaEncodingAttr::getTotalElemsPerThread(
 +    ArrayRef<int64_t> shape, Type eltTy) const {
-+  auto mmaLayout = getParent().cast<NvidiaMmaEncodingAttr>();
++  auto mmaLayout = mlir::cast<NvidiaMmaEncodingAttr>(getParent());
 +  return product<int64_t>(shape) /
 +         (mmaLayout.getWarpsPerCTA()[0] * kMetadataElementsPerWarp);
 +}
@@ -169,9 +176,10 @@ diff --git a/lib/Dialect/TritonGPU/IR/Dialect.cpp b/lib/Dialect/TritonGPU/IR/Dia
  } // namespace triton
  } // namespace mlir
 diff --git a/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/ConvertLayoutOpToLLVM.cpp b/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/ConvertLayoutOpToLLVM.cpp
+index f8ece0f1c..435610817 100644
 --- a/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/ConvertLayoutOpToLLVM.cpp
 +++ b/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/ConvertLayoutOpToLLVM.cpp
-@@ -38,6 +38,14 @@ Value convertLayout(int opIdx, Conversio
+@@ -43,6 +43,14 @@ Value convertLayout(int opIdx, ConversionPatternRewriter &rewriter,
                      const LLVMTypeConverter *typeConverter, Value thread);
  }
  
@@ -185,19 +193,19 @@ diff --git a/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/ConvertLayoutOpToLLVM.
 +
  namespace {
  
- struct LocalLoadOpConversion
-@@ -59,6 +67,10 @@ public:
-             .isa<NvidiaMmaEncodingAttr>()) {
+ using namespace mlir;
+@@ -67,6 +75,10 @@ public:
+             cast<DotOperandEncodingAttr>(dstLayout).getParent())) {
        return lowerSharedToDotOperand(op, adaptor, getTypeConverter(), rewriter);
      }
-+    if (srcLayout.isa<SharedEncodingAttr>() &&
-+        dstLayout.isa<triton::gpu::SparseDotMetaEncodingAttr>()) {
++    if (isa<SharedEncodingAttr>(srcLayout) &&
++        isa<triton::gpu::SparseDotMetaEncodingAttr>(dstLayout)) {
 +      return lowerSharedToSparseMeta(op, adaptor, getTypeConverter(), rewriter);
 +    }
      return failure();
    }
  
-@@ -130,6 +142,29 @@ private:
+@@ -138,6 +150,26 @@ private:
      rewriter.replaceOp(op, res);
      return success();
    }
@@ -208,13 +216,10 @@ diff --git a/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/ConvertLayoutOpToLLVM.
 +      const LLVMTypeConverter *typeConverter,
 +      ConversionPatternRewriter &rewriter) const {
 +    auto loc = op.getLoc();
-+    auto sparseEncoding = op.getResult()
-+                              .getType()
-+                              .cast<RankedTensorType>()
-+                              .getEncoding()
-+                              .cast<triton::gpu::SparseDotMetaEncodingAttr>();
++    auto sparseEncoding = cast<triton::gpu::SparseDotMetaEncodingAttr>(
++        cast<RankedTensorType>(op.getResult().getType()).getEncoding());
 +    auto llvmElemTy = typeConverter->convertType(
-+        op.getSrc().getType().cast<MemDescType>().getElementType());
++        cast<MemDescType>(op.getSrc().getType()).getElementType());
 +    auto smemObj = getSharedMemoryObjectFromStruct(loc, adaptor.getSrc(),
 +                                                   llvmElemTy, rewriter);
 +    Value res = SharedToSparseDotOperand::convertLayout(
@@ -229,6 +234,7 @@ diff --git a/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/ConvertLayoutOpToLLVM.
  struct ConvertLayoutOpOptimizedConversion
 diff --git a/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/ConvertLayoutOpToLLVM/SharedToSparseDotOperand.cpp b/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/ConvertLayoutOpToLLVM/SharedToSparseDotOperand.cpp
 new file mode 100644
+index 000000000..3011cf73d
 --- /dev/null
 +++ b/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/ConvertLayoutOpToLLVM/SharedToSparseDotOperand.cpp
 @@ -0,0 +1,69 @@
@@ -255,7 +261,7 @@ new file mode 100644
 +    Value thread) {
 +  // Calculate tile size as number of mask elements (4xi4).
 +  NvidiaMmaEncodingAttr mmaLayout =
-+      sparseEncoding.getParent().cast<NvidiaMmaEncodingAttr>();
++      cast<NvidiaMmaEncodingAttr>(sparseEncoding.getParent());
 +  SmallVector<unsigned> shapePerCTATile = {
 +      kTileSize * mmaLayout.getWarpsPerCTA()[0],
 +      kTileSize / kMetadataElementsPerPackedValue};
@@ -272,7 +278,7 @@ new file mode 100644
 +  Value rowId = add(mul(warpGroupId, i32_val(kTileSize)), laneGroupId);
 +
 +  // Calculate number of tile repetitions.
-+  auto shape = tensor.getType().cast<MemDescType>().getShape();
++  auto shape = cast<MemDescType>(tensor.getType()).getShape();
 +  int repM = shape[0] / shapePerCTATile[0];
 +  int repK = shape[1] / shapePerCTATile[1];
 +  assert(repM > 0 && repK > 0);
@@ -302,9 +308,10 @@ new file mode 100644
 +}
 +}  // namespace SharedToSparseDotOperand
 diff --git a/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/DotOpToLLVM.cpp b/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/DotOpToLLVM.cpp
+index 374b9ec9e..1601806b4 100644
 --- a/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/DotOpToLLVM.cpp
 +++ b/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/DotOpToLLVM.cpp
-@@ -32,6 +32,12 @@ LogicalResult convertAsyncWGMMA(triton::
+@@ -32,6 +32,12 @@ LogicalResult convertAsyncWGMMA(triton::nvidia_gpu::DotAsyncOp op,
                                  const LLVMTypeConverter *typeConverter,
                                  ConversionPatternRewriter &rewriter,
                                  Value thread);
@@ -317,7 +324,7 @@ diff --git a/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/DotOpToLLVM.cpp b/thir
  namespace {
  struct DotOpConversion : public ConvertOpToLLVMPattern<triton::DotOp> {
    using ConvertOpToLLVMPattern<triton::DotOp>::ConvertOpToLLVMPattern;
-@@ -180,6 +186,18 @@ struct DotWaitOpConversion
+@@ -174,6 +180,18 @@ struct DotWaitOpConversion
      return success();
    }
  };
@@ -336,7 +343,7 @@ diff --git a/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/DotOpToLLVM.cpp b/thir
  } // namespace
  
  void mlir::triton::NVIDIA::populateDotOpToLLVMPatterns(
-@@ -188,4 +206,5 @@ void mlir::triton::NVIDIA::populateDotOp
+@@ -182,4 +200,5 @@ void mlir::triton::NVIDIA::populateDotOpToLLVMPatterns(
    patterns.add<DotOpConversion>(typeConverter, benefit);
    patterns.add<DotAsyncOpConversion>(typeConverter, benefit);
    patterns.add<DotWaitOpConversion>(typeConverter, benefit);
@@ -344,6 +351,7 @@ diff --git a/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/DotOpToLLVM.cpp b/thir
  }
 diff --git a/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/DotOpToLLVM/Sparse.cpp b/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/DotOpToLLVM/Sparse.cpp
 new file mode 100644
+index 000000000..34d9212d2
 --- /dev/null
 +++ b/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/DotOpToLLVM/Sparse.cpp
 @@ -0,0 +1,339 @@
@@ -397,15 +405,15 @@ new file mode 100644
 +                               const LLVMTypeConverter *typeConverter,
 +                               ConversionPatternRewriter &rewriter) {
 +  // Get number of repetitions across the dimensions.
-+  auto aTensorTy = op.getA().getType().cast<RankedTensorType>();
-+  auto bTensorTy = op.getB().getType().cast<RankedTensorType>();
++  auto aTensorTy = cast<RankedTensorType>(op.getA().getType());
++  auto bTensorTy = cast<RankedTensorType>(op.getB().getType());
 +
-+  auto layoutA = aTensorTy.getEncoding().dyn_cast<DotOperandEncodingAttr>();
-+  auto layoutB = bTensorTy.getEncoding().dyn_cast<DotOperandEncodingAttr>();
++  auto layoutA = dyn_cast<DotOperandEncodingAttr>(aTensorTy.getEncoding());
++  auto layoutB = dyn_cast<DotOperandEncodingAttr>(bTensorTy.getEncoding());
 +  assert(layoutA != nullptr && layoutB != nullptr);
 +
 +  int bitwidth = aTensorTy.getElementType().getIntOrFloatBitWidth();
-+  auto mmaEnc = layoutA.getParent().cast<NvidiaMmaEncodingAttr>();
++  auto mmaEnc = cast<NvidiaMmaEncodingAttr>(layoutA.getParent());
 +  auto repA = mmaEnc.getMMAv2Rep(triton::gpu::getShapePerCTA(aTensorTy),
 +                                 bitwidth, layoutA.getOpIdx());
 +  auto repB = mmaEnc.getMMAv2Rep(triton::gpu::getShapePerCTA(bTensorTy),
@@ -435,7 +443,7 @@ new file mode 100644
 +  }
 +
 +  // Flatten accumulator values.
-+  auto dTensorTy = op.getD().getType().cast<RankedTensorType>();
++  auto dTensorTy = cast<RankedTensorType>(op.getD().getType());
 +  auto fc = unpackLLElements(loc, adaptor.getC(), rewriter);
 +
 +  // Create `mma.sp` instruction for 4/8 core matrices.
@@ -512,7 +520,7 @@ new file mode 100644
 +                     Location loc, std::vector<unsigned int> instrShape,
 +                     bool trans, int dimWpt, Value warpId, MemDescType tensorTy,
 +                     Value baseDesc, int minor) {
-+  auto sharedLayout = tensorTy.getEncoding().cast<SharedEncodingAttr>();
++  auto sharedLayout = cast<SharedEncodingAttr>(tensorTy.getEncoding());
 +  int elemBytes = tensorTy.getElementTypeBitWidth() / 8;
 +  int elemsPerSwizzlingRow =
 +      kMmaLineSize / sharedLayout.getPerPhase() / elemBytes;
@@ -541,10 +549,10 @@ new file mode 100644
 +                                 ConversionPatternRewriter &rewriter,
 +                                 Value thread) {
 +  // Get number of repetitions across the dimensions.
-+  auto aTensorTy = op.getA().getType().cast<MemDescType>();
-+  auto bTensorTy = op.getB().getType().cast<MemDescType>();
-+  auto dTensorTy = op.getD().getType().cast<RankedTensorType>();
-+  auto mmaEnc = dTensorTy.getEncoding().cast<NvidiaMmaEncodingAttr>();
++  auto aTensorTy = cast<MemDescType>(op.getA().getType());
++  auto bTensorTy = cast<MemDescType>(op.getB().getType());
++  auto dTensorTy = cast<RankedTensorType>(op.getD().getType());
++  auto mmaEnc = cast<NvidiaMmaEncodingAttr>(dTensorTy.getEncoding());
 +
 +  auto shapePerCTA = getShapePerCTA(dTensorTy);
 +  auto shapePerCTATile = getShapePerCTATile(mmaEnc);
@@ -573,7 +581,7 @@ new file mode 100644
 +    auto sharedObj = getSharedMemoryObjectFromStruct(
 +        loc, arg, typeConverter->convertType(tensorTy.getElementType()),
 +        rewriter);
-+    auto sharedLayout = tensorTy.getEncoding().cast<SharedEncodingAttr>();
++    auto sharedLayout = cast<SharedEncodingAttr>(tensorTy.getEncoding());
 +    auto shape = getShapePerCTA(tensorTy);
 +    auto ord = sharedLayout.getOrder();
 +    int byteSize = aTensorTy.getElementTypeBitWidth() / 8;
@@ -671,9 +679,9 @@ new file mode 100644
 +                                 SparseDotOp::Adaptor adaptor,
 +                                 const LLVMTypeConverter *typeConverter,
 +                                 ConversionPatternRewriter &rewriter) {
-+  auto resultTy = op.getResult().getType().cast<RankedTensorType>();
++  auto resultTy = cast<RankedTensorType>(op.getResult().getType());
 +  NvidiaMmaEncodingAttr mmaLayout =
-+      resultTy.getEncoding().cast<NvidiaMmaEncodingAttr>();
++      cast<NvidiaMmaEncodingAttr>(resultTy.getEncoding());
 +
 +  if (mmaLayout.isAmpere()) {
 +    return convertSparseMMA(op, adaptor, typeConverter, rewriter);
@@ -687,9 +695,10 @@ new file mode 100644
 +      "Unsupported SparseDotOp found when converting TritonGPU to LLVM.");
 +}
 diff --git a/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/DotOpToLLVM/WGMMA.cpp b/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/DotOpToLLVM/WGMMA.cpp
+index 738f0fe04..867939f65 100644
 --- a/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/DotOpToLLVM/WGMMA.cpp
 +++ b/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/DotOpToLLVM/WGMMA.cpp
-@@ -87,8 +87,8 @@ int64_t getSwizzlingFromLayout(const Sha
+@@ -88,8 +88,8 @@ int64_t getSwizzlingFromLayout(const SharedEncodingAttr &layout,
    return swizzlingByteWidth;
  }
  
diff --git a/third_party/xla/third_party/triton/xla_extensions/sparse_dot_fixes_y24w17.patch b/third_party/xla/third_party/triton/xla_extensions/sparse_dot_fixes_y24w17.patch
index 9d1ae2e91cae3f..ce009aa688e9bf 100644
--- a/third_party/xla/third_party/triton/xla_extensions/sparse_dot_fixes_y24w17.patch
+++ b/third_party/xla/third_party/triton/xla_extensions/sparse_dot_fixes_y24w17.patch
@@ -1,30 +1,8 @@
-diff --git a/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/ConvertLayoutOpToLLVM/SharedToSparseDotOperand.cpp b/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/ConvertLayoutOpToLLVM/SharedToSparseDotOperand.cpp
---- a/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/ConvertLayoutOpToLLVM/SharedToSparseDotOperand.cpp
-+++ b/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/ConvertLayoutOpToLLVM/SharedToSparseDotOperand.cpp
-@@ -22,16 +22,16 @@ Value convertLayout(
-   // Calculate tile size as number of mask elements (4xi4).
-   NvidiaMmaEncodingAttr mmaLayout =
-       sparseEncoding.getParent().cast<NvidiaMmaEncodingAttr>();
-+  SmallVector<unsigned> warpsPerCTA = mmaLayout.getWarpsPerCTA();
-   SmallVector<unsigned> shapePerCTATile = {
--      kTileSize * mmaLayout.getWarpsPerCTA()[0],
--      kTileSize / kMetadataElementsPerPackedValue};
-+      kTileSize * warpsPerCTA[0], kTileSize / kMetadataElementsPerPackedValue};
-   Value strideM = smemObj.strides[0];
-   Value strideK = smemObj.strides[1];
- 
-   // Calculate offset in the tile for the current thread.
-   Value threadsPerWarp = i32_val(kThreadsPerWarp);
-   Value warpId = udiv(thread, threadsPerWarp);
--  Value warpGroupId = urem(warpId, i32_val(shapePerCTATile[0] / kTileSize));
-+  Value warpGroupId = udiv(warpId, i32_val(warpsPerCTA[1]));
-   Value laneId = urem(thread, threadsPerWarp);
-   Value laneGroupId = udiv(laneId, i32_val(kThreadsInGroup));
-   Value columnId = urem(laneId, i32_val(shapePerCTATile[1]));
 diff --git a/lib/Dialect/TritonGPU/Transforms/AccelerateMatmul.cpp b/lib/Dialect/TritonGPU/Transforms/AccelerateMatmul.cpp
+index 0516fc56f..1f27f8a43 100644
 --- a/lib/Dialect/TritonGPU/Transforms/AccelerateMatmul.cpp
 +++ b/lib/Dialect/TritonGPU/Transforms/AccelerateMatmul.cpp
-@@ -139,6 +139,7 @@ class BlockedToMMA : public mlir::Rewrit
+@@ -142,6 +142,7 @@ class BlockedToMMA : public mlir::RewritePattern {
                  mlir::TypeID::get<arith::ArithDialect>());
    }
  
@@ -32,7 +10,7 @@ diff --git a/lib/Dialect/TritonGPU/Transforms/AccelerateMatmul.cpp b/lib/Dialect
    // Finds the first different bitwidth in the chain of shape-preserving
    // unary ops that x depends on.
    // There are two primary scenarios:
-@@ -172,7 +173,6 @@ class BlockedToMMA : public mlir::Rewrit
+@@ -175,7 +176,6 @@ class BlockedToMMA : public mlir::RewritePattern {
      return origBitWidth;
    }
  
@@ -40,7 +18,7 @@ diff --git a/lib/Dialect/TritonGPU/Transforms/AccelerateMatmul.cpp b/lib/Dialect
    BlockedToMMA(mlir::MLIRContext *context, int computeCapability)
        : mlir::RewritePattern(tt::DotOp::getOperationName(), 2, context),
          computeCapability(computeCapability) {}
-@@ -388,18 +388,22 @@ class SparseBlockedToMMA : public mlir::
+@@ -389,18 +389,22 @@ class SparseBlockedToMMA : public mlir::RewritePattern {
                                                          newRetType, oldAcc);
  
      if (versionMajor == 2) {
@@ -49,7 +27,7 @@ diff --git a/lib/Dialect/TritonGPU/Transforms/AccelerateMatmul.cpp b/lib/Dialect
 +      int kWidth = 32 / minBitwidth;
 +
        // convert A operand
-       auto oldAType = a.getType().cast<RankedTensorType>();
+       auto oldAType = cast<RankedTensorType>(a.getType());
 -      auto newAEncoding = ttg::DotOperandEncodingAttr::get(
 -          ctx, 0, mmaEnc, oldAType.getElementType());
 +      auto newAEncoding =
@@ -59,7 +37,7 @@ diff --git a/lib/Dialect/TritonGPU/Transforms/AccelerateMatmul.cpp b/lib/Dialect
        a = rewriter.create<ttg::ConvertLayoutOp>(a.getLoc(), newAType, a);
  
        // convert B operand
-       auto oldBType = b.getType().cast<RankedTensorType>();
+       auto oldBType = cast<RankedTensorType>(b.getType());
 -      auto newBEncoding = ttg::DotOperandEncodingAttr::get(
 -          ctx, 1, mmaEnc, oldBType.getElementType());
 +      auto newBEncoding =
@@ -67,3 +45,27 @@ diff --git a/lib/Dialect/TritonGPU/Transforms/AccelerateMatmul.cpp b/lib/Dialect
        auto newBType = RankedTensorType::get(
            oldBType.getShape(), oldBType.getElementType(), newBEncoding);
        b = rewriter.create<ttg::ConvertLayoutOp>(b.getLoc(), newBType, b);
+diff --git a/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/ConvertLayoutOpToLLVM/SharedToSparseDotOperand.cpp b/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/ConvertLayoutOpToLLVM/SharedToSparseDotOperand.cpp
+index 3011cf73d..ea587dced 100644
+--- a/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/ConvertLayoutOpToLLVM/SharedToSparseDotOperand.cpp
++++ b/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/ConvertLayoutOpToLLVM/SharedToSparseDotOperand.cpp
+@@ -22,16 +22,16 @@ Value convertLayout(
+   // Calculate tile size as number of mask elements (4xi4).
+   NvidiaMmaEncodingAttr mmaLayout =
+       cast<NvidiaMmaEncodingAttr>(sparseEncoding.getParent());
++  SmallVector<unsigned> warpsPerCTA = mmaLayout.getWarpsPerCTA();
+   SmallVector<unsigned> shapePerCTATile = {
+-      kTileSize * mmaLayout.getWarpsPerCTA()[0],
+-      kTileSize / kMetadataElementsPerPackedValue};
++      kTileSize * warpsPerCTA[0], kTileSize / kMetadataElementsPerPackedValue};
+   Value strideM = smemObj.strides[0];
+   Value strideK = smemObj.strides[1];
+ 
+   // Calculate offset in the tile for the current thread.
+   Value threadsPerWarp = i32_val(kThreadsPerWarp);
+   Value warpId = udiv(thread, threadsPerWarp);
+-  Value warpGroupId = urem(warpId, i32_val(shapePerCTATile[0] / kTileSize));
++  Value warpGroupId = udiv(warpId, i32_val(warpsPerCTA[1]));
+   Value laneId = urem(thread, threadsPerWarp);
+   Value laneGroupId = udiv(laneId, i32_val(kThreadsInGroup));
+   Value columnId = urem(laneId, i32_val(shapePerCTATile[1]));
diff --git a/third_party/xla/third_party/triton/xla_extensions/sparse_dot_nvgpu.patch b/third_party/xla/third_party/triton/xla_extensions/sparse_dot_nvgpu.patch
index ced13cff33fd16..791618363b2f34 100644
--- a/third_party/xla/third_party/triton/xla_extensions/sparse_dot_nvgpu.patch
+++ b/third_party/xla/third_party/triton/xla_extensions/sparse_dot_nvgpu.patch
@@ -1,7 +1,8 @@
-diff --git a/include/triton/Dialect/NVGPU/IR/NVGPUOps.td b/include/triton/Dialect/NVGPU/IR/NVGPUOps.td
---- a/include/triton/Dialect/NVGPU/IR/NVGPUOps.td
-+++ b/include/triton/Dialect/NVGPU/IR/NVGPUOps.td
-@@ -87,6 +87,15 @@ def NVGPU_WGMMAOp : NVGPU_Op<"wgmma", []
+diff --git a/third_party/nvidia/include/Dialect/NVGPU/IR/NVGPUOps.td b/third_party/nvidia/include/Dialect/NVGPU/IR/NVGPUOps.td
+index ca9d18873..d39bc6ec4 100644
+--- a/third_party/nvidia/include/Dialect/NVGPU/IR/NVGPUOps.td
++++ b/third_party/nvidia/include/Dialect/NVGPU/IR/NVGPUOps.td
+@@ -87,6 +87,15 @@ def NVGPU_WGMMAOp : NVGPU_Op<"wgmma", []> {
    let assemblyFormat = "$opA `,` $opB (`,` $opC^)? attr-dict `:` functional-type(operands, $res)";
  }
  
@@ -18,9 +19,10 @@ diff --git a/include/triton/Dialect/NVGPU/IR/NVGPUOps.td b/include/triton/Dialec
    let arguments = (ins LLVM_AnyPointer:$addr, I32:$ctaId, I32Attr:$bitwidth, I32Attr:$vec);
    let builders = [
 diff --git a/third_party/nvidia/lib/NVGPUToLLVM/NVGPUToLLVMPass.cpp b/third_party/nvidia/lib/NVGPUToLLVM/NVGPUToLLVMPass.cpp
+index e19216520..aacbfb569 100644
 --- a/third_party/nvidia/lib/NVGPUToLLVM/NVGPUToLLVMPass.cpp
 +++ b/third_party/nvidia/lib/NVGPUToLLVM/NVGPUToLLVMPass.cpp
-@@ -688,6 +688,84 @@ public:
+@@ -668,6 +668,84 @@ public:
    }
  };
  
@@ -31,7 +33,7 @@ diff --git a/third_party/nvidia/lib/NVGPUToLLVM/NVGPUToLLVMPass.cpp b/third_part
 +  using Base::Base;
 +
 +  std::vector<std::string> getOutputConstraints(ttn::SparseWGMMAOp op) const {
-+    auto outputStructType = op.getType().cast<LLVM::LLVMStructType>();
++    auto outputStructType = cast<LLVM::LLVMStructType>(op.getType());
 +    uint32_t numOutputRegs = outputStructType.getBody().size();
 +    std::string output =
 +        outputStructType.getBody().front().isF32() ? "=f" : "=r";
@@ -71,7 +73,7 @@ diff --git a/third_party/nvidia/lib/NVGPUToLLVM/NVGPUToLLVMPass.cpp b/third_part
 +
 +    // Output and operand C
 +    uint32_t numCRegs =
-+        op.getType().cast<LLVM::LLVMStructType>().getBody().size();
++        cast<LLVM::LLVMStructType>(op.getType()).getBody().size();
 +    args += "{";
 +    for (uint32_t i = 0; i < numCRegs; ++i) {
 +      args += "$" + std::to_string(asmOpIdx++) + (i == numCRegs - 1 ? "" : ",");
@@ -105,13 +107,17 @@ diff --git a/third_party/nvidia/lib/NVGPUToLLVM/NVGPUToLLVMPass.cpp b/third_part
  class ConvertNVGPUToLLVM : public ConvertNVGPUToLLVMBase<ConvertNVGPUToLLVM> {
  
  public:
-@@ -711,7 +789,8 @@ public:
+@@ -688,10 +766,9 @@ public:
+     patterns.add<NVGPUOpGenericPattern<ttn::ClusterCTAIdOp>>(
+         context, Cluster_Cta_Id_Op, Constraints({"=r"}), Constraints());
  
-     patterns.add<FenceAsyncSharedOpPattern, StoreMatrixOpPattern,
-                  ClusterArriveOpPattern, LoadDSmemOpPattern, WGMMAOpPattern,
--                 WGMMAWaitGroupOpPattern, StoreDSmemOpPattern>(context);
-+                 WGMMAWaitGroupOpPattern, StoreDSmemOpPattern,
-+                 SparseWGMMAOpPattern>(context);
+-    patterns
+-        .add<FenceAsyncSharedOpPattern, StoreMatrixOpPattern,
+-             ClusterArriveOpPattern, WGMMAOpPattern, WGMMAWaitGroupOpPattern>(
+-            context);
++    patterns.add<FenceAsyncSharedOpPattern, StoreMatrixOpPattern,
++                 ClusterArriveOpPattern, WGMMAOpPattern,
++                 WGMMAWaitGroupOpPattern, SparseWGMMAOpPattern>(context);
  
      if (applyPatternsAndFoldGreedily(mod, std::move(patterns)).failed())
        signalPassFailure();
diff --git a/third_party/xla/third_party/triton/xla_extensions/sparse_dot_passes.patch b/third_party/xla/third_party/triton/xla_extensions/sparse_dot_passes.patch
index 74662cf3c90dc9..9136cb84b24254 100644
--- a/third_party/xla/third_party/triton/xla_extensions/sparse_dot_passes.patch
+++ b/third_party/xla/third_party/triton/xla_extensions/sparse_dot_passes.patch
@@ -1,7 +1,8 @@
 diff --git a/lib/Conversion/TritonToTritonGPU/TritonToTritonGPUPass.cpp b/lib/Conversion/TritonToTritonGPU/TritonToTritonGPUPass.cpp
+index 4aa2712ec..16a6253d7 100644
 --- a/lib/Conversion/TritonToTritonGPU/TritonToTritonGPUPass.cpp
 +++ b/lib/Conversion/TritonToTritonGPU/TritonToTritonGPUPass.cpp
-@@ -277,6 +277,89 @@ struct TritonDotPattern : public OpConve
+@@ -279,6 +279,89 @@ struct TritonDotPattern : public OpConversionPattern<triton::DotOp> {
    }
  };
  
@@ -12,7 +13,7 @@ diff --git a/lib/Conversion/TritonToTritonGPU/TritonToTritonGPUPass.cpp b/lib/Co
 +  LogicalResult matchAndRewrite(
 +      triton::gpu::SparseDotOp op, OpAdaptor adaptor,
 +      ConversionPatternRewriter &rewriter) const override {
-+    RankedTensorType origType = op.getType().cast<RankedTensorType>();
++    RankedTensorType origType = cast<RankedTensorType>(op.getType());
 +    auto origShape = origType.getShape();
 +    auto typeConverter = getTypeConverter<TritonGPUTypeConverter>();
 +    int numWarps = typeConverter->getNumWarps();
@@ -40,8 +41,8 @@ diff --git a/lib/Conversion/TritonToTritonGPU/TritonToTritonGPUPass.cpp b/lib/Co
 +        RankedTensorType::get(origShape, origType.getElementType(), dEncoding);
 +
 +    // a & b must be of smem layout
-+    auto aType = adaptor.getA().getType().cast<RankedTensorType>();
-+    auto bType = adaptor.getB().getType().cast<RankedTensorType>();
++    auto aType = cast<RankedTensorType>(adaptor.getA().getType());
++    auto bType = cast<RankedTensorType>(adaptor.getB().getType());
 +    Type aEltType = aType.getElementType();
 +    Type bEltType = bType.getElementType();
 +    Attribute aEncoding = aType.getEncoding();
@@ -51,14 +52,14 @@ diff --git a/lib/Conversion/TritonToTritonGPU/TritonToTritonGPUPass.cpp b/lib/Co
 +    Value a = adaptor.getA();
 +    Value b = adaptor.getB();
 +    Value c = adaptor.getC();
-+    if (!aEncoding.isa<triton::gpu::DotOperandEncodingAttr>()) {
++    if (!isa<triton::gpu::DotOperandEncodingAttr>(aEncoding)) {
 +      Attribute encoding = triton::gpu::DotOperandEncodingAttr::get(
 +          getContext(), 0, dEncoding, aEltType);
 +      auto dstType =
 +          RankedTensorType::get(aType.getShape(), aEltType, encoding);
 +      a = rewriter.create<triton::gpu::ConvertLayoutOp>(a.getLoc(), dstType, a);
 +    }
-+    if (!bEncoding.isa<triton::gpu::DotOperandEncodingAttr>()) {
++    if (!isa<triton::gpu::DotOperandEncodingAttr>(bEncoding)) {
 +      Attribute encoding = triton::gpu::DotOperandEncodingAttr::get(
 +          getContext(), 1, dEncoding, bEltType);
 +      auto dstType =
@@ -68,11 +69,11 @@ diff --git a/lib/Conversion/TritonToTritonGPU/TritonToTritonGPUPass.cpp b/lib/Co
 +    c = rewriter.create<triton::gpu::ConvertLayoutOp>(c.getLoc(), retType, c);
 +
 +    // aMeta must be of smem layout
-+    auto aMetaType = adaptor.getAMeta().getType().cast<RankedTensorType>();
++    auto aMetaType = cast<RankedTensorType>(adaptor.getAMeta().getType());
 +    Attribute aMetaEncoding = aMetaType.getEncoding();
 +    if (!aMetaEncoding) return failure();
 +    Value aMeta = adaptor.getAMeta();
-+    if (!aMetaEncoding.isa<triton::gpu::SparseDotMetaEncodingAttr>()) {
++    if (!isa<triton::gpu::SparseDotMetaEncodingAttr>(aMetaEncoding)) {
 +      Attribute encoding =
 +          triton::gpu::SparseDotMetaEncodingAttr::get(getContext(), dEncoding);
 +      auto dstType = RankedTensorType::get(
@@ -91,17 +92,17 @@ diff --git a/lib/Conversion/TritonToTritonGPU/TritonToTritonGPUPass.cpp b/lib/Co
  struct TritonCatPattern : public OpConversionPattern<triton::CatOp> {
    using OpConversionPattern::OpConversionPattern;
  
-@@ -550,6 +633,7 @@ void populateTritonPatterns(TritonGPUTyp
-       GenericOpPattern<triton::AtomicRMWOp>, GenericOpPattern<ReturnOp>,
+@@ -553,6 +636,7 @@ void populateTritonPatterns(TritonGPUTypeConverter &typeConverter,
+       GenericOpPattern<triton::ExperimentalDescriptorStoreOp>,
        GenericOpPattern<triton::CallOp>, TritonFuncOpPattern>(typeConverter,
                                                               context);
 +  patterns.insert<TritonSparseDotPattern>(typeConverter, context);
  }
  
  //
-@@ -788,6 +872,12 @@ public:
-                  IntegerAttr::get(
-                      i32_ty, llvm::APInt(32, computeCapability.getValue())));
+@@ -794,6 +878,12 @@ public:
+     mod->setAttr(AttrTargetName,
+                  StringAttr::get(context, this->target.getValue()));
  
 +    // Only transform sparse dot op with undefined layout.
 +    target.addDynamicallyLegalOp<triton::gpu::SparseDotOp>(
@@ -113,9 +114,10 @@ diff --git a/lib/Conversion/TritonToTritonGPU/TritonToTritonGPUPass.cpp b/lib/Co
        return signalPassFailure();
  
 diff --git a/lib/Dialect/TritonGPU/Transforms/AccelerateMatmul.cpp b/lib/Dialect/TritonGPU/Transforms/AccelerateMatmul.cpp
+index 098ee85e4..0516fc56f 100644
 --- a/lib/Dialect/TritonGPU/Transforms/AccelerateMatmul.cpp
 +++ b/lib/Dialect/TritonGPU/Transforms/AccelerateMatmul.cpp
-@@ -42,8 +42,9 @@ static int getMMAVersionSafe(int compute
+@@ -44,8 +44,9 @@ static int getMMAVersionSafe(int computeCapability, tt::DotOp op) {
    return 0;
  }
  
@@ -126,7 +128,7 @@ diff --git a/lib/Dialect/TritonGPU/Transforms/AccelerateMatmul.cpp b/lib/Dialect
    auto rank = shape.size();
    // Early exit for batched matmul
    if (rank == 3)
-@@ -56,14 +57,14 @@ warpsPerTileV2(tt::DotOp dotOp, const Ar
+@@ -58,8 +59,8 @@ warpsPerTileV2(tt::DotOp dotOp, const ArrayRef<int64_t> shape, int numWarps) {
    auto slices = multiRootGetSlice(dotOp, {filter}, {filter});
    bool hasChainedDot = false;
    for (Operation *op : slices) {
@@ -137,14 +139,7 @@ diff --git a/lib/Dialect/TritonGPU/Transforms/AccelerateMatmul.cpp b/lib/Dialect
        auto resTy = chainedDot.getResult().getType();
        if (resTy.getRank() != rank) {
          continue;
-       }
-       if (auto mmaEncoding =
--              resTy.getEncoding().dyn_cast<NvidiaMmaEncodingAttr>()) {
-+              resTy.getEncoding().template dyn_cast<NvidiaMmaEncodingAttr>()) {
-         return ttg::getWarpsPerCTA(mmaEncoding);
-       }
-       hasChainedDot = true;
-@@ -101,12 +102,13 @@ warpsPerTileV2(tt::DotOp dotOp, const Ar
+@@ -103,12 +104,13 @@ warpsPerTileV2(tt::DotOp dotOp, const ArrayRef<int64_t> shape, int numWarps) {
    return ret;
  }
  
@@ -162,7 +157,7 @@ diff --git a/lib/Dialect/TritonGPU/Transforms/AccelerateMatmul.cpp b/lib/Dialect
        slices.end())
      return {(unsigned)numWarps, 1};
  
-@@ -175,9 +177,10 @@ public:
+@@ -178,9 +180,10 @@ public:
        : mlir::RewritePattern(tt::DotOp::getOperationName(), 2, context),
          computeCapability(computeCapability) {}
  
@@ -176,7 +171,7 @@ diff --git a/lib/Dialect/TritonGPU/Transforms/AccelerateMatmul.cpp b/lib/Dialect
      switch (version) {
      case 2:
        return warpsPerTileV2(dotOp, shape, numWarps);
-@@ -337,6 +340,98 @@ public:
+@@ -335,6 +338,98 @@ public:
      return success();
    }
  };
@@ -201,7 +196,7 @@ diff --git a/lib/Dialect/TritonGPU/Transforms/AccelerateMatmul.cpp b/lib/Dialect
 +    // Check data-types and SM compatibility
 +    RankedTensorType oldRetType = dotOp.getType();
 +    if (!oldRetType.getEncoding() ||
-+        oldRetType.getEncoding().isa<ttg::NvidiaMmaEncodingAttr>())
++        isa<ttg::NvidiaMmaEncodingAttr>(oldRetType.getEncoding()))
 +      return failure();
 +
 +    assert(computeCapability >= 80 &&
@@ -216,7 +211,7 @@ diff --git a/lib/Dialect/TritonGPU/Transforms/AccelerateMatmul.cpp b/lib/Dialect
 +
 +    auto instrShape =
 +        mmaVersionToInstrShape(versionMajor, retShapePerCTA,
-+                               a.getType().cast<TensorOrMemDesc>(), numWarps);
++                               cast<TensorOrMemDesc>(a.getType()), numWarps);
 +    auto warpsPerTile = BlockedToMMA::getWarpsPerTile(
 +        dotOp, retShapePerCTA, versionMajor, numWarps, instrShape);
 +    ttg::NvidiaMmaEncodingAttr mmaEnc =
@@ -232,7 +227,7 @@ diff --git a/lib/Dialect/TritonGPU/Transforms/AccelerateMatmul.cpp b/lib/Dialect
 +
 +    if (versionMajor == 2) {
 +      // convert A operand
-+      auto oldAType = a.getType().cast<RankedTensorType>();
++      auto oldAType = cast<RankedTensorType>(a.getType());
 +      auto newAEncoding = ttg::DotOperandEncodingAttr::get(
 +          ctx, 0, mmaEnc, oldAType.getElementType());
 +      auto newAType = RankedTensorType::get(
@@ -240,7 +235,7 @@ diff --git a/lib/Dialect/TritonGPU/Transforms/AccelerateMatmul.cpp b/lib/Dialect
 +      a = rewriter.create<ttg::ConvertLayoutOp>(a.getLoc(), newAType, a);
 +
 +      // convert B operand
-+      auto oldBType = b.getType().cast<RankedTensorType>();
++      auto oldBType = cast<RankedTensorType>(b.getType());
 +      auto newBEncoding = ttg::DotOperandEncodingAttr::get(
 +          ctx, 1, mmaEnc, oldBType.getElementType());
 +      auto newBType = RankedTensorType::get(
@@ -253,7 +248,7 @@ diff --git a/lib/Dialect/TritonGPU/Transforms/AccelerateMatmul.cpp b/lib/Dialect
 +
 +    // convert metadata
 +    Value meta = dotOp.getAMeta();
-+    auto oldMetaType = meta.getType().cast<RankedTensorType>();
++    auto oldMetaType = cast<RankedTensorType>(meta.getType());
 +    auto newMetaType = RankedTensorType::get(
 +        oldMetaType.getShape(), oldMetaType.getElementType(),
 +        SparseDotMetaEncodingAttr::get(ctx, mmaEnc));
@@ -275,7 +270,7 @@ diff --git a/lib/Dialect/TritonGPU/Transforms/AccelerateMatmul.cpp b/lib/Dialect
  } // namespace
  
  static Value promoteOperand(OpBuilder &builder, Location loc, Value operand,
-@@ -397,6 +491,7 @@ public:
+@@ -394,6 +489,7 @@ public:
  
      mlir::RewritePatternSet patterns(context);
      patterns.add<::BlockedToMMA>(context, computeCapability);
@@ -284,33 +279,31 @@ diff --git a/lib/Dialect/TritonGPU/Transforms/AccelerateMatmul.cpp b/lib/Dialect
        signalPassFailure();
      }
 diff --git a/lib/Dialect/TritonGPU/Transforms/Pipeliner/MatmulLoopPipeline.cpp b/lib/Dialect/TritonGPU/Transforms/Pipeliner/MatmulLoopPipeline.cpp
+index 97ca6a840..f0ef124ff 100644
 --- a/lib/Dialect/TritonGPU/Transforms/Pipeliner/MatmulLoopPipeline.cpp
 +++ b/lib/Dialect/TritonGPU/Transforms/Pipeliner/MatmulLoopPipeline.cpp
-@@ -47,6 +47,10 @@ struct PipelinedOpInfo {
-   bool loadIsMMAV3 = false;
+@@ -188,6 +188,10 @@ public:
+   }
  };
  
-+bool isDotOp(Operation* op) {
++static bool isDotOp(Operation* op) {
 +  return isa<tt::DotOp, ttg::SparseDotOp>(op);
 +}
 +
- } // namespace
- 
  static bool isMMAv3Dot(Operation *op) {
-@@ -165,22 +169,28 @@ getSharedEncIfAllUsersAreDotEnc(Value val) {
+   auto dot = dyn_cast<tt::DotOp>(op);
+   if (!dot)
+@@ -399,19 +403,28 @@ getSharedEncIfAllUsersAreDotEnc(Value val) {
      } else {
        if (!isa<ttg::LocalLoadOp, ttg::ConvertLayoutOp>(user))
          return std::nullopt;
--      auto dotOpEnc = user->getResult(0)
--                          .getType()
--                          .cast<TensorOrMemDesc>()
--                          .getEncoding()
--                          .dyn_cast<ttg::DotOperandEncodingAttr>();
+-      auto dotOpEnc = dyn_cast<ttg::DotOperandEncodingAttr>(
+-          cast<TensorOrMemDesc>(user->getResult(0).getType()).getEncoding());
 -      if (!dotOpEnc)
 +      auto enc =
-+          user->getResult(0).getType().cast<TensorOrMemDesc>().getEncoding();
++          cast<TensorOrMemDesc>(user->getResult(0).getType()).getEncoding();
 +      if (isa<ttg::DotOperandEncodingAttr>(enc)) {
-+        auto srcTy = val.getType().cast<TensorOrMemDesc>();
++        auto srcTy = cast<TensorOrMemDesc>(val.getType());
 +        auto CTALayout = ttg::getCTALayout(srcTy.getEncoding());
 +        auto order = ttg::getOrder(srcTy.getEncoding());
 +        unsigned bitWidth = srcTy.getElementType().getIntOrFloatBitWidth();
@@ -321,14 +314,14 @@ diff --git a/lib/Dialect/TritonGPU/Transforms/Pipeliner/MatmulLoopPipeline.cpp b
 +            srcTy.getElementType().getIntOrFloatBitWidth(),
 +            /*needTrans=*/false);
 +      } else if (isa<ttg::SparseDotMetaEncodingAttr>(enc)) {
-+        auto srcTy = val.getType().cast<TensorOrMemDesc>();
++        auto srcTy = cast<TensorOrMemDesc>(val.getType());
 +        tempAttr = ttg::SharedEncodingAttr::get(
 +            val.getContext(), /*vec=*/1, /*perPhase=*/1, /*maxPhase=*/1,
 +            ttg::getOrder(srcTy.getEncoding()),
 +            ttg::getCTALayout(srcTy.getEncoding()));
 +      } else {
          return std::nullopt;
--      auto srcTy = val.getType().cast<TensorOrMemDesc>();
+-      auto srcTy = cast<TensorOrMemDesc>(val.getType());
 -      auto CTALayout = ttg::getCTALayout(srcTy.getEncoding());
 -      auto order = ttg::getOrder(srcTy.getEncoding());
 -      unsigned bitWidth = srcTy.getElementType().getIntOrFloatBitWidth();
@@ -341,71 +334,63 @@ diff --git a/lib/Dialect/TritonGPU/Transforms/Pipeliner/MatmulLoopPipeline.cpp b
      }
      // Check that the shared encodings needed by the users are compatible.
      if (!tempAttr || (attr != nullptr && attr != tempAttr))
-@@ -313,7 +323,7 @@ loadOpsToDistanceAndUse(scf::ForOp forOp) {
+@@ -518,7 +531,7 @@ loadOpsToIndirectionLevelAndUse(scf::ForOp forOp) {
        };
  
    for (Operation &op : forOp.getBody()->without_terminator()) {
 -    if (!isa<tt::DotOp>(op))
 +    if (!isDotOp(&op))
        continue;
+     seen.clear();
      dfs(&op, 0, &op);
-   }
-@@ -391,7 +401,8 @@ collectOpsToPipeline(scf::ForOp forOp,
-   // loads.
-   for (auto &[loadOp, distAndUse] : loadOpToDistAndUse) {
-     PipelinedOpInfo loadInfo;
--    if (auto dot = dyn_cast<tt::DotOp>(distAndUse.second)) {
-+    if (isDotOp(distAndUse.second)) {
-+      auto dot = dyn_cast<tt::DotOp>(distAndUse.second);
-       if (loadIsMMAv3(loadOp)) {
+@@ -595,7 +608,8 @@ assignMemoryLayouts(llvm::SmallVector<std::tuple<Operation *, int, Operation *>>
+         continue;
+     }
+ 
+-    if (auto dot = dyn_cast<tt::DotOp>(use)) {
++    if (isDotOp(use)) {
++      auto dot = dyn_cast<tt::DotOp>(use);
+       loadInfo.usedByDot = true;
+       if (loadIsMMAv3(op)) {
          loadInfo.loadIsMMAV3 = true;
-         loadInfo.sharedEncoding =
-@@ -410,7 +421,7 @@ collectOpsToPipeline(scf::ForOp forOp,
+@@ -614,7 +628,7 @@ assignMemoryLayouts(llvm::SmallVector<std::tuple<Operation *, int, Operation *>>
          // The codegen bug is caught by an assertion, so if you think you've
          // fixed it, feel free to delete this code and see if the assert still
          // fails.  :)
 -        if (!loadInfo.sharedEncoding) {
 +        if (dot && !loadInfo.sharedEncoding) {
-           if (auto dotEnc = dot.getResult()
-                                 .getType()
-                                 .getEncoding()
-@@ -788,7 +799,7 @@ bool mlir::triton::preProcessLoopAndGetSchedule(
-     int useStage = opToInfo[info.use].stage;
-     int numBuffers = useStage - defStage;
- 
--    if (hasMMAV3 && isa<tt::DotOp>(info.use)) {
-+    if (hasMMAV3 && isDotOp(info.use)) {
-       // For MMAv3, we need an extra buffer as this is assumed in the wgmma
-       // pipelining post-processing.
-       numBuffers++;
+           if (auto dotEnc = dyn_cast<ttg::NvidiaMmaEncodingAttr>(
+                   dot.getResult().getType().getEncoding())) {
+             auto loadTy = cast<RankedTensorType>(op->getResultTypes()[0]);
 diff --git a/lib/Dialect/TritonGPU/Transforms/ReduceDataDuplication.cpp b/lib/Dialect/TritonGPU/Transforms/ReduceDataDuplication.cpp
+index 2211df31b..ee5ff44d8 100644
 --- a/lib/Dialect/TritonGPU/Transforms/ReduceDataDuplication.cpp
 +++ b/lib/Dialect/TritonGPU/Transforms/ReduceDataDuplication.cpp
-@@ -36,6 +36,10 @@ public:
+@@ -37,6 +37,10 @@ public:
        auto srcEncoding = srcType.getEncoding();
-       if (srcEncoding.isa<triton::gpu::SharedEncodingAttr>())
+       if (isa<triton::gpu::SharedEncodingAttr>(srcEncoding))
          return;
-+      if (dstType.getEncoding().isa<triton::gpu::SparseDotMetaEncodingAttr>()) {
++      if (isa<triton::gpu::SparseDotMetaEncodingAttr>(dstType.getEncoding())) {
 +        replaceSparseMetaEncoding(cvtOp);
 +        return;
 +      }
        auto dstDotOp =
-           dstType.getEncoding().dyn_cast<triton::gpu::DotOperandEncodingAttr>();
+           dyn_cast<triton::gpu::DotOperandEncodingAttr>(dstType.getEncoding());
        if (!dstDotOp)
-@@ -74,6 +78,27 @@ public:
+@@ -83,6 +87,27 @@ public:
        cvtOp.erase();
      });
    }
 +
 + private:
 +  void replaceSparseMetaEncoding(triton::gpu::ConvertLayoutOp cvtOp) {
-+    auto srcType = cvtOp.getOperand().getType().cast<RankedTensorType>();
++    auto srcType = cast<RankedTensorType>(cvtOp.getOperand().getType());
 +    auto srcEncoding = srcType.getEncoding();
 +    auto sharedLayout = triton::gpu::SharedEncodingAttr::get(
 +        cvtOp.getContext(), 8, 1, 1, triton::gpu::getOrder(srcEncoding),
 +        triton::gpu::getCTALayout(srcEncoding));
 +
-+    auto dstType = cvtOp.getType().cast<RankedTensorType>();
++    auto dstType = cast<RankedTensorType>(cvtOp.getType());
 +    auto tmpType = triton::MemDescType::get(
 +        dstType.getShape(), dstType.getElementType(), sharedLayout);
 +
@@ -421,6 +406,7 @@ diff --git a/lib/Dialect/TritonGPU/Transforms/ReduceDataDuplication.cpp b/lib/Di
  
  std::unique_ptr<Pass> mlir::triton::gpu::createReduceDataDuplicationPass() {
 diff --git a/lib/Dialect/TritonNvidiaGPU/Transforms/FenceInsertion.cpp b/lib/Dialect/TritonNvidiaGPU/Transforms/FenceInsertion.cpp
+index f456d36a6..a1dac2b72 100644
 --- a/lib/Dialect/TritonNvidiaGPU/Transforms/FenceInsertion.cpp
 +++ b/lib/Dialect/TritonNvidiaGPU/Transforms/FenceInsertion.cpp
 @@ -45,7 +45,7 @@ public:
@@ -432,7 +418,7 @@ diff --git a/lib/Dialect/TritonNvidiaGPU/Transforms/FenceInsertion.cpp b/lib/Dia
          return WalkResult::advance();
        OpBuilder builder(op);
        auto a = op->getOperand(0);
-@@ -83,7 +83,7 @@ private:
+@@ -80,7 +80,7 @@ private:
      static DenseSet<std::pair<Operation *, unsigned>> trace;
      auto op = operand.getDefiningOp();
      // avoid redundant insertion
diff --git a/third_party/xla/xla/service/gpu/ir_emitter_triton_cuda.cc b/third_party/xla/xla/service/gpu/ir_emitter_triton_cuda.cc
index 2ad76a6391226c..0606d115ced97f 100644
--- a/third_party/xla/xla/service/gpu/ir_emitter_triton_cuda.cc
+++ b/third_party/xla/xla/service/gpu/ir_emitter_triton_cuda.cc
@@ -18,6 +18,7 @@ limitations under the License.
 #include "nvidia/include/NVGPUToLLVM/NVGPUToLLVMPass.h"
 #include "nvidia/include/TritonNVIDIAGPUToLLVM/Passes.h"
 #include "absl/status/status.h"
+#include "absl/strings/str_format.h"
 #include "mlir/Conversion/ArithToLLVM/ArithToLLVM.h"  // from @llvm-project
 #include "mlir/Conversion/IndexToLLVM/IndexToLLVM.h"  // from @llvm-project
 #include "mlir/Conversion/SCFToControlFlow/SCFToControlFlow.h"  // from @llvm-project
@@ -60,7 +61,8 @@ absl::Status CreateTritonPipeline(
   // Based on make_ttgir() in
   // @triton//:third_party/nvidia/backend/compiler.py
   pm.addPass(mt::createConvertTritonToTritonGPUPass(
-      config.num_warps, threadsPerWarp, config.num_ctas, ccAsInt));
+      absl::StrFormat("cuda:%u", ccAsInt), config.num_warps, threadsPerWarp,
+      config.num_ctas));
   pm.addPass(mt::gpu::createCoalescePass());
   if (ccCuda.IsAtLeastAmpere()) {
     pm.addPass(mt::gpu::createF32DotTCPass());
@@ -70,7 +72,7 @@ absl::Status CreateTritonPipeline(
   pm.addPass(mt::gpu::createOptimizeThreadLocalityPass());
   pm.addPass(mt::gpu::createAccelerateMatmulPass(ccAsInt));
   pm.addPass(mt::gpu::createRemoveLayoutConversionsPass());
-  pm.addPass(mt::gpu::createOptimizeDotOperandsPass());
+  pm.addPass(mt::gpu::createOptimizeDotOperandsPass(ccCuda.IsAtLeastAmpere()));
   pm.addPass(mlir::createCSEPass());
 
   pm.addPass(mt::gpu::createPipelinePass(config.num_stages, config.num_warps,
@@ -80,7 +82,7 @@ absl::Status CreateTritonPipeline(
     pm.addPass(mt::gpu::createPrefetchPass());
   }
 
-  pm.addPass(mt::gpu::createOptimizeDotOperandsPass());
+  pm.addPass(mt::gpu::createOptimizeDotOperandsPass(ccCuda.IsAtLeastAmpere()));
   // We need to disable this pass because it undoes the hoisting of dot_operand
   // layout conversion done in
   // triton/lib/Dialect/TritonGPU/Transforms/OptimizeDotOperands.cpp in
@@ -93,6 +95,7 @@ absl::Status CreateTritonPipeline(
   pm.addPass(mlir::createSymbolDCEPass());
   if (ccCuda.IsAtLeastHopper()) {
     pm.addPass(mlir::createTritonNvidiaGPUFenceInsertionPass(ccAsInt));
+    pm.addPass(mlir::createTritonNvidiaGPUTMALoweringPass());
   }
   pm.addPass(mlir::createCanonicalizerPass());
 
diff --git a/third_party/xla/xla/service/gpu/tests/sparse_add_layout.mlir b/third_party/xla/xla/service/gpu/tests/sparse_add_layout.mlir
index 3d1e42b46c80d6..55faf1b76bda3d 100644
--- a/third_party/xla/xla/service/gpu/tests/sparse_add_layout.mlir
+++ b/third_party/xla/xla/service/gpu/tests/sparse_add_layout.mlir
@@ -1,7 +1,7 @@
-// RUN: triton-opt %s -split-input-file -convert-triton-to-tritongpu | FileCheck %s
+// RUN: triton-opt %s -split-input-file -convert-triton-to-tritongpu='target=cuda:80 num-warps=4' | FileCheck %s
 
 // CHECK-COUNT-4: #triton_gpu.blocked
-module attributes {"triton_gpu.num-warps" = 4 : i32} {
+module {
   tt.func @sparse_dot() {
     %A = arith.constant dense<1.00e+00> : tensor<64x32xf16>
     %meta = arith.constant dense<0x3333> : tensor<64x4xi16>

From 322831ccc5a8fef9a98ab068e6fc77ed46a25f37 Mon Sep 17 00:00:00 2001
From: Adrian Kuegel <akuegel@google.com>
Date: Tue, 14 May 2024 07:04:41 -0700
Subject: [PATCH 112/478] Avoid FindIndex function by using the inverse
 permutation.

Let p be a permutation. FindIndex(p, x) should return y where x = p[y].
This is equivalent to p^-1[x] = p^-1[p[y]] = y, so y = p^-1[x].

PiperOrigin-RevId: 633572618
---
 .../xla/xla/service/layout_normalization.cc   | 23 +++++++++++--------
 .../xla/service/layout_normalization_test.cc  | 18 +++++++++++++++
 2 files changed, 31 insertions(+), 10 deletions(-)

diff --git a/third_party/xla/xla/service/layout_normalization.cc b/third_party/xla/xla/service/layout_normalization.cc
index 6837437f8f93e2..9974feb9541724 100644
--- a/third_party/xla/xla/service/layout_normalization.cc
+++ b/third_party/xla/xla/service/layout_normalization.cc
@@ -156,7 +156,7 @@ class LayoutNormalizationVisitor : public DfsHloRewriteVisitor {
     auto normalized_shape = Normalize(s);
     auto layout_as_permutation = ToTransposeDimensions(s.layout());
     int64_t normalized_concat_dim =
-        FindIndex(layout_as_permutation, orig_concat_dim);
+        InversePermutation(layout_as_permutation)[orig_concat_dim];
     auto normalized_concat =
         hlo->AddInstruction(HloInstruction::CreateConcatenate(
             normalized_shape, normalized_inputs, normalized_concat_dim));
@@ -224,10 +224,12 @@ class LayoutNormalizationVisitor : public DfsHloRewriteVisitor {
         ToTransposeDimensions(s.layout());
     std::vector<int64_t> br_dimensions;
     if (!hlo->dimensions().empty()) {
-      br_dimensions = Permute(hlo->dimensions(), layout_as_permutation);
-    }
-    for (int64_t& d : br_dimensions) {
-      d = FindIndex(orig_output_layout_as_permutation, d);
+      br_dimensions.reserve(hlo->dimensions().size());
+      auto inverse_perm = InversePermutation(orig_output_layout_as_permutation);
+      for (int64_t dim :
+           ComposePermutations(hlo->dimensions(), layout_as_permutation)) {
+        br_dimensions.push_back(inverse_perm[dim]);
+      }
     }
     auto normalized_broadcast = MakeBroadcastHlo(
         normalized_input, br_dimensions, normalized_shape, &hlo->metadata());
@@ -244,9 +246,8 @@ class LayoutNormalizationVisitor : public DfsHloRewriteVisitor {
     auto normalized_shape = Normalize(s);
     std::vector<int64_t> orig_output_layout_as_permutation =
         ToTransposeDimensions(s.layout());
-    int64_t iota_dimension = hlo->dimensions()[0];
-    int64_t new_iota_dimension =
-        FindIndex(orig_output_layout_as_permutation, iota_dimension);
+    int64_t new_iota_dimension = InversePermutation(
+        orig_output_layout_as_permutation)[hlo->dimensions()[0]];
     auto normalized_iota = hlo->AddInstruction(
         HloInstruction::CreateIota(normalized_shape, new_iota_dimension));
     SetVisited(*normalized_iota);
@@ -468,8 +469,9 @@ class LayoutNormalizationVisitor : public DfsHloRewriteVisitor {
         ToTransposeDimensions(hlo->shape().layout());
     std::vector<int64_t> new_dimensions;
     new_dimensions.reserve(hlo->dimensions().size());
+    auto inverse_perm = InversePermutation(layout_as_permutation);
     for (int64_t dim : hlo->dimensions()) {
-      new_dimensions.push_back(FindIndex(layout_as_permutation, dim));
+      new_dimensions.push_back(inverse_perm[dim]);
     }
     absl::c_sort(new_dimensions);
     auto normalized_reverse = hlo->AddInstruction(
@@ -499,8 +501,9 @@ class LayoutNormalizationVisitor : public DfsHloRewriteVisitor {
       new_padding.add_dimensions();
     }
 
+    auto inverse_perm = InversePermutation(layout_as_permutation);
     for (int dim = 0; dim < s.dimensions_size(); dim++) {
-      int tr_dim = static_cast<int>(FindIndex(layout_as_permutation, dim));
+      int tr_dim = static_cast<int>(inverse_perm[dim]);
       *new_padding.mutable_dimensions(tr_dim) = padded_config.dimensions(dim);
     }
 
diff --git a/third_party/xla/xla/service/layout_normalization_test.cc b/third_party/xla/xla/service/layout_normalization_test.cc
index dd528096ffa312..6c5da347cb1d58 100644
--- a/third_party/xla/xla/service/layout_normalization_test.cc
+++ b/third_party/xla/xla/service/layout_normalization_test.cc
@@ -235,6 +235,24 @@ ENTRY main {
 )");
 }
 
+TEST_F(LayoutNormalizationTest, BroadcastOperandLayoutNotInverseOfItself) {
+  const char* hlo = R"(
+HloModule module
+
+ENTRY main {
+  a = f32[4,3,5]{0,2,1} parameter(0)
+  b = f32[4,3,2,5]{0,1,2,3} broadcast(a), dimensions={0,1,3}
+  ROOT out = abs(b)
+}
+)";
+
+  CheckLayoutNormalization(hlo, R"(
+// CHECK: [[bitcast_1:%[^ ]+]] = f32[3,5,4]{2,1,0} bitcast
+// CHECK: [[broadcast_0:%[^ ]+]] = f32[5,2,3,4]{3,2,1,0} broadcast([[bitcast_1]]), dimensions={2,0,3}
+// CHECK: [[abs_2:%[^ ]+]] = f32[5,2,3,4]{3,2,1,0} abs([[broadcast_0]])
+)");
+}
+
 TEST_F(LayoutNormalizationTest, BroadcastCustomOutputLayout) {
   const char* hlo = R"(
 HloModule module

From ece868ec01922cc9faf3c5746ce4dbc67100d3de Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Tue, 14 May 2024 07:18:11 -0700
Subject: [PATCH 113/478] Update ops-related pbtxt files.

PiperOrigin-RevId: 633575662
---
 tensorflow/core/ops/ops.pbtxt | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorflow/core/ops/ops.pbtxt b/tensorflow/core/ops/ops.pbtxt
index 7322cef85415f5..5019fe2a5b2fad 100644
--- a/tensorflow/core/ops/ops.pbtxt
+++ b/tensorflow/core/ops/ops.pbtxt
@@ -1,4 +1,4 @@
-go/debugstr   
+go/debugproto    
 op {
   name: "Abort"
   attr {

From 9e2a14d27e834aef4f46bd38b849801df64f6cb1 Mon Sep 17 00:00:00 2001
From: Greg Olechwierowicz <olechwierowicz@google.com>
Date: Tue, 14 May 2024 07:25:41 -0700
Subject: [PATCH 114/478] [XLA:GPU] Add utility functions to iterate over
 instructions matching opcode.

PiperOrigin-RevId: 633577273
---
 third_party/xla/xla/hlo/utils/BUILD           |  16 +++
 third_party/xla/xla/hlo/utils/hlo_query.h     |  20 ++++
 .../xla/xla/hlo/utils/hlo_query_test.cc       | 113 ++++++++++++++++++
 3 files changed, 149 insertions(+)
 create mode 100644 third_party/xla/xla/hlo/utils/hlo_query_test.cc

diff --git a/third_party/xla/xla/hlo/utils/BUILD b/third_party/xla/xla/hlo/utils/BUILD
index 5b08563b2d483b..39153f7e1eda67 100644
--- a/third_party/xla/xla/hlo/utils/BUILD
+++ b/third_party/xla/xla/hlo/utils/BUILD
@@ -154,3 +154,19 @@ cc_library(
         "@com_google_absl//absl/container:flat_hash_set",
     ],
 )
+
+xla_cc_test(
+    name = "hlo_query_test",
+    srcs = [
+        "hlo_query_test.cc",
+    ],
+    deps = [
+        ":hlo_query",
+        "//xla/hlo/ir:hlo",
+        "//xla/service:hlo_parser",
+        "//xla/tests:hlo_test_base",
+        "@com_google_absl//absl/strings:string_view",
+        "@com_google_googletest//:gtest_main",
+        "@local_tsl//tsl/platform:statusor",
+    ],
+)
diff --git a/third_party/xla/xla/hlo/utils/hlo_query.h b/third_party/xla/xla/hlo/utils/hlo_query.h
index e12e3306651375..8343df4dc24472 100644
--- a/third_party/xla/xla/hlo/utils/hlo_query.h
+++ b/third_party/xla/xla/hlo/utils/hlo_query.h
@@ -78,6 +78,26 @@ bool IsBroadcastOfParameter(const HloInstruction& instr);
 HloInstruction* GetFirstInstructionWithOpcode(const HloComputation& computation,
                                               HloOpcode opcode);
 
+// Applies `fn` to a collection of instruction for a given `computation`.
+template <typename Fn>
+void ForEachInstructionWithOpcode(HloComputation& computation, HloOpcode opcode,
+                                  Fn&& fn) {
+  for (HloInstruction* instr : computation.instructions()) {
+    if (instr->opcode() == opcode) {
+      fn(instr);
+    }
+  }
+}
+
+// Applies `fn` to a collection of instruction for a given `module`.
+template <typename Fn>
+void ForEachInstructionWithOpcode(HloModule& module, HloOpcode opcode,
+                                  Fn&& fn) {
+  for (HloComputation* computation : module.computations()) {
+    ForEachInstructionWithOpcode(*computation, opcode, fn);
+  }
+}
+
 // Determines whether the given computation contains an instruction with one of
 // the given opcodes.  Checks both comp's instructions and the instructions of
 // any computations nested within it.
diff --git a/third_party/xla/xla/hlo/utils/hlo_query_test.cc b/third_party/xla/xla/hlo/utils/hlo_query_test.cc
new file mode 100644
index 00000000000000..7697bffc855806
--- /dev/null
+++ b/third_party/xla/xla/hlo/utils/hlo_query_test.cc
@@ -0,0 +1,113 @@
+/* Copyright 2024 The OpenXLA Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+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.
+==============================================================================*/
+
+#include "xla/hlo/utils/hlo_query.h"
+
+#include <memory>
+
+#include <gtest/gtest.h>
+#include "absl/strings/string_view.h"
+#include "xla/hlo/ir/hlo_computation.h"
+#include "xla/hlo/ir/hlo_module.h"
+#include "xla/hlo/ir/hlo_opcode.h"
+#include "xla/service/hlo_parser.h"
+#include "xla/tests/hlo_test_base.h"
+#include "tsl/platform/statusor.h"
+
+namespace xla {
+namespace {
+
+using HloQueryTest = HloTestBase;
+
+template <typename Hlo>
+int CountInstructions(Hlo& module, HloOpcode opcode) {
+  int counter = 0;
+  hlo_query::ForEachInstructionWithOpcode(
+      module, opcode, [&counter](auto& instr) { counter++; });
+  return counter;
+}
+
+TEST_F(HloQueryTest,
+       GetInstructionWithOpCodeReturnsMatchingInstructionForModule) {
+  constexpr absl::string_view kHloString = R"(
+HloModule m
+
+computation.0 {
+  param.0 = f32[32]{0} parameter(0)
+  ROOT _ = f32[32]{0} rsqrt(param.0)
+}
+
+ENTRY main {
+  param.0 = f32[32]{0} parameter(0)
+  param.1 = f32[32]{0} parameter(1)
+  param.2 = f32[32]{0} parameter(2)
+  param.3 = f32[32]{0} parameter(3)
+  add.0 = f32[32]{0} add(param.0,param.1)
+  add.1 = f32[32]{0} add(param.1,param.2)
+  sub.0 = f32[32]{0} subtract(param.0,param.1)
+  mul.0 = f32[32]{0} multiply(param.0,param.1)
+  mul.1 = f32[32]{0} multiply(param.1,param.2)
+  mul.2 = f32[32]{0} multiply(param.2,param.3)
+  comp.0 = call(param.0), to_apply=computation.0
+  ROOT _ = (f32[32],f32[32],f32[32],f32[32],f32[32],f32[32],f32[32]) tuple(comp.0,add.0,add.1,sub.0,mul.0,mul.1,mul.2)
+})";
+
+  TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
+                          ParseAndReturnUnverifiedModule(kHloString));
+  EXPECT_EQ(CountInstructions(*module, HloOpcode::kAdd), 2);
+  EXPECT_EQ(CountInstructions(*module, HloOpcode::kSubtract), 1);
+  EXPECT_EQ(CountInstructions(*module, HloOpcode::kMultiply), 3);
+}
+
+TEST_F(HloQueryTest,
+       GetInstructionWithOpCodeReturnsMatchingInstructionForComputation) {
+  constexpr absl::string_view kHloString = R"(
+HloModule m
+
+computation.0 {
+  param.0 = f32[32]{0} parameter(0)
+  param.1 = f32[32]{0} parameter(1)
+  param.2 = f32[32]{0} parameter(2)
+  param.3 = f32[32]{0} parameter(3)
+  add.0 = f32[32]{0} add(param.0,param.1)
+  add.1 = f32[32]{0} add(param.1,param.2)
+  sub.0 = f32[32]{0} subtract(param.0,param.1)
+  mul.0 = f32[32]{0} multiply(param.0,param.1)
+  mul.1 = f32[32]{0} multiply(param.1,param.2)
+  ROOT mul.2 = f32[32]{0} multiply(param.2,param.3)
+}
+
+ENTRY main {
+  param.0 = f32[32]{0} parameter(0)
+  param.1 = f32[32]{0} parameter(1)
+  param.2 = f32[32]{0} parameter(2)
+  param.3 = f32[32]{0} parameter(3)
+  add.0 = f32[32]{0} add(param.0,param.1)
+  sub.0 = f32[32]{0} subtract(param.0,param.1)
+  mul.0 = f32[32]{0} multiply(param.0,param.1)
+  comp.0 = f32[32]{0} call(param.0,param.1,param.2), to_apply=computation.0
+  ROOT _ = (f32[32],f32[32],f32[32],f32[32]) tuple(add.0,sub.0,mul.0,comp.0)
+})";
+
+  TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
+                          ParseAndReturnUnverifiedModule(kHloString));
+  HloComputation* computation = module->GetComputationWithName("computation.0");
+  EXPECT_EQ(CountInstructions(*computation, HloOpcode::kAdd), 2);
+  EXPECT_EQ(CountInstructions(*computation, HloOpcode::kSubtract), 1);
+  EXPECT_EQ(CountInstructions(*computation, HloOpcode::kMultiply), 3);
+}
+
+}  // namespace
+}  // namespace xla

From a26c03d30732a3c17b2a18dd7156e26920d1cf61 Mon Sep 17 00:00:00 2001
From: Ilia Sergachev <isergachev@nvidia.com>
Date: Tue, 14 May 2024 08:22:56 -0700
Subject: [PATCH 115/478] PR #12433: [GPU] Make cuDNN fusion test run on H100.

Imported from GitHub PR https://github.com/openxla/xla/pull/12433

Copybara import of the project:

--
856f1ddb05b0cecf1092aa3547611274b27f0796 by Ilia Sergachev <isergachev@nvidia.com>:

[GPU] Make cuDNN fusion test run on H100.

Merging this change closes #12433

PiperOrigin-RevId: 633592263
---
 third_party/xla/xla/service/gpu/fusions/BUILD | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/third_party/xla/xla/service/gpu/fusions/BUILD b/third_party/xla/xla/service/gpu/fusions/BUILD
index 2e0a1202157486..1b1d9c8a5ef0e5 100644
--- a/third_party/xla/xla/service/gpu/fusions/BUILD
+++ b/third_party/xla/xla/service/gpu/fusions/BUILD
@@ -665,7 +665,7 @@ xla_test(
     name = "cudnn_test",
     srcs = if_cuda_is_configured(["cudnn_test.cc"]),
     backend_tags = {"gpu": [
-        "requires-gpu-sm80",
+        "requires-gpu-sm90",
     ]},
     backends = [
         "gpu",

From d030362628a694b8b499e4e7d9ed835267be43d4 Mon Sep 17 00:00:00 2001
From: Kyle Lucke <klucke@google.com>
Date: Tue, 14 May 2024 08:59:27 -0700
Subject: [PATCH 116/478] Lift Kernel::Create and TypedKernel::Create methods
 out into separate factory objects to remove circular dependencies.

Existing Create methods left in place until all uses are eliminated.

PiperOrigin-RevId: 633601919
---
 third_party/xla/xla/service/gpu/BUILD         |  1 +
 third_party/xla/xla/service/gpu/kernels/BUILD |  3 +
 .../cutlass_gemm_custom_kernel_benchmarks.cc  |  4 +-
 .../cutlass_gemm_custom_kernel_test.cc        |  7 +-
 .../gpu/kernels/topk_custom_kernel_test.cc    |  7 +-
 third_party/xla/xla/service/gpu/runtime/BUILD |  5 +-
 .../service/gpu/runtime/command_buffer_cmd.cc |  3 +-
 .../xla/service/gpu/runtime/kernel_thunk.cc   |  7 +-
 .../xla/service/gpu/stream_executor_util.cc   |  3 +-
 third_party/xla/xla/stream_executor/BUILD     | 30 ++++++
 third_party/xla/xla/stream_executor/gpu/BUILD |  2 +
 .../gpu/gpu_command_buffer_test.cc            | 15 +--
 .../stream_executor/gpu/gpu_kernel_test.cc    |  6 +-
 third_party/xla/xla/stream_executor/kernel.h  |  9 ++
 .../xla/xla/stream_executor/kernel_factory.h  | 44 +++++++++
 .../stream_executor/typed_kernel_factory.h    | 95 +++++++++++++++++++
 16 files changed, 221 insertions(+), 20 deletions(-)
 create mode 100644 third_party/xla/xla/stream_executor/kernel_factory.h
 create mode 100644 third_party/xla/xla/stream_executor/typed_kernel_factory.h

diff --git a/third_party/xla/xla/service/gpu/BUILD b/third_party/xla/xla/service/gpu/BUILD
index 223aa418be67e7..91776ad61dee01 100644
--- a/third_party/xla/xla/service/gpu/BUILD
+++ b/third_party/xla/xla/service/gpu/BUILD
@@ -4364,6 +4364,7 @@ cc_library(
         "//xla/hlo/ir:hlo",
         "//xla/service:hlo_module_config",
         "//xla/stream_executor",
+        "//xla/stream_executor:kernel_factory",
         "//xla/stream_executor:launch_dim",
         "//xla/tsl/util:env_var",
         "//xla/tsl/util/proto:proto_utils",
diff --git a/third_party/xla/xla/service/gpu/kernels/BUILD b/third_party/xla/xla/service/gpu/kernels/BUILD
index 8d950a052f2aef..fc9c7ec54b2aff 100644
--- a/third_party/xla/xla/service/gpu/kernels/BUILD
+++ b/third_party/xla/xla/service/gpu/kernels/BUILD
@@ -232,6 +232,7 @@ xla_test(
         "//xla:xla_data_proto_cc",
         "//xla/service:platform_util",
         "//xla/stream_executor",
+        "//xla/stream_executor:kernel_factory",
         "//xla/stream_executor:platform",
         "//xla/stream_executor:platform_manager",
         "//xla/stream_executor/cuda:cuda_platform",
@@ -281,6 +282,7 @@ xla_test(
         ":cutlass_gemm_custom_kernel",
         "//xla:xla_data_proto_cc",
         "//xla/stream_executor",
+        "//xla/stream_executor:kernel_factory",
         "//xla/stream_executor:platform",
         "//xla/stream_executor:platform_manager",
         "//xla/stream_executor/cuda:cuda_platform",
@@ -301,6 +303,7 @@ cc_binary(
         "//xla:xla_data_proto_cc",
         "//xla/service:gpu_plugin",
         "//xla/stream_executor",
+        "//xla/stream_executor:kernel_factory",
         "//xla/stream_executor:platform",
         "//xla/stream_executor:platform_manager",
         "//xla/stream_executor/cuda:cuda_platform",
diff --git a/third_party/xla/xla/service/gpu/kernels/cutlass_gemm_custom_kernel_benchmarks.cc b/third_party/xla/xla/service/gpu/kernels/cutlass_gemm_custom_kernel_benchmarks.cc
index e2f7bccea9fe57..b843fac733154c 100644
--- a/third_party/xla/xla/service/gpu/kernels/cutlass_gemm_custom_kernel_benchmarks.cc
+++ b/third_party/xla/xla/service/gpu/kernels/cutlass_gemm_custom_kernel_benchmarks.cc
@@ -20,6 +20,7 @@ limitations under the License.
 #include "xla/service/gpu/kernels/cutlass_gemm_custom_kernel.h"
 #include "xla/stream_executor/device_description.h"
 #include "xla/stream_executor/kernel.h"
+#include "xla/stream_executor/kernel_factory.h"
 #include "xla/stream_executor/platform.h"
 #include "xla/stream_executor/platform_manager.h"
 #include "xla/stream_executor/stream.h"
@@ -56,7 +57,8 @@ static void BM_RowMajorGemm(benchmark::State& state) {
                            /*indices=*/{0, 1, 2}, /*slices=*/{}, device);
 
   TF_ASSERT_OK_AND_ASSIGN(
-      auto gemm, se::Kernel::Create(executor, custom_kernel->kernel_spec()));
+      auto gemm,
+      se::KernelFactory::Create(executor, custom_kernel->kernel_spec()));
 
   // Prepare arguments: a=1.1, b=1.2, c=0.0
   se::DeviceMemory<float> a = executor->AllocateArray<float>(m * k, 0);
diff --git a/third_party/xla/xla/service/gpu/kernels/cutlass_gemm_custom_kernel_test.cc b/third_party/xla/xla/service/gpu/kernels/cutlass_gemm_custom_kernel_test.cc
index 3748ed5251564e..b566fceab8511c 100644
--- a/third_party/xla/xla/service/gpu/kernels/cutlass_gemm_custom_kernel_test.cc
+++ b/third_party/xla/xla/service/gpu/kernels/cutlass_gemm_custom_kernel_test.cc
@@ -21,6 +21,7 @@ limitations under the License.
 #include <vector>
 
 #include "xla/stream_executor/kernel.h"
+#include "xla/stream_executor/kernel_factory.h"
 #include "xla/stream_executor/platform.h"
 #include "xla/stream_executor/platform_manager.h"
 #include "xla/stream_executor/stream.h"
@@ -46,7 +47,8 @@ TEST(CutlassGemmKernelTest, SimpleGemm) {
       /*indices=*/{0, 1, 2}, /*slices=*/{}, executor->GetDeviceDescription());
 
   TF_ASSERT_OK_AND_ASSIGN(
-      auto gemm, se::Kernel::Create(executor, custom_kernel->kernel_spec()));
+      auto gemm,
+      se::KernelFactory::Create(executor, custom_kernel->kernel_spec()));
 
   int64_t length = 4 * 4;
   int64_t byte_length = sizeof(float) * length;
@@ -96,7 +98,8 @@ TEST(CutlassGemmKernelTest, LoadFromSharedLibrary) {
       /*indices=*/{0, 1, 2}, /*slices=*/{}, executor->GetDeviceDescription());
 
   TF_ASSERT_OK_AND_ASSIGN(
-      auto gemm, se::Kernel::Create(executor, custom_kernel->kernel_spec()));
+      auto gemm,
+      se::KernelFactory::Create(executor, custom_kernel->kernel_spec()));
 
   int64_t length = 4 * 4;
   int64_t byte_length = sizeof(float) * length;
diff --git a/third_party/xla/xla/service/gpu/kernels/topk_custom_kernel_test.cc b/third_party/xla/xla/service/gpu/kernels/topk_custom_kernel_test.cc
index d50fd054df2ac3..b083118d0a8d7c 100644
--- a/third_party/xla/xla/service/gpu/kernels/topk_custom_kernel_test.cc
+++ b/third_party/xla/xla/service/gpu/kernels/topk_custom_kernel_test.cc
@@ -28,6 +28,7 @@ limitations under the License.
 #include "absl/strings/substitute.h"
 #include "xla/service/platform_util.h"
 #include "xla/stream_executor/kernel.h"
+#include "xla/stream_executor/kernel_factory.h"
 #include "xla/stream_executor/platform.h"
 #include "xla/stream_executor/platform_manager.h"
 #include "xla/stream_executor/stream.h"
@@ -111,7 +112,8 @@ TEST_P(TopKKernelTest, TopKFloat) {
       GetTopKKernel("topk", PrimitiveType::F32, n, k, batch_size);
 
   TF_ASSERT_OK_AND_ASSIGN(
-      auto kernel, se::Kernel::Create(executor, custom_kernel->kernel_spec()));
+      auto kernel,
+      se::KernelFactory::Create(executor, custom_kernel->kernel_spec()));
 
   // Launch topk kernel with device memory arguments.
   se::KernelArgsDeviceMemoryArray arr(
@@ -165,7 +167,8 @@ TEST_P(TopKKernelTest, TopKPackedNegative) {
       GetTopKKernel("topk", PrimitiveType::F32, n, k, batch_size);
 
   TF_ASSERT_OK_AND_ASSIGN(
-      auto kernel, se::Kernel::Create(executor, custom_kernel->kernel_spec()));
+      auto kernel,
+      se::KernelFactory::Create(executor, custom_kernel->kernel_spec()));
 
   // Launch topk kernel with device memory arguments.
   se::KernelArgsDeviceMemoryArray arr(
diff --git a/third_party/xla/xla/service/gpu/runtime/BUILD b/third_party/xla/xla/service/gpu/runtime/BUILD
index c259a044cb8308..4175990422e7bb 100644
--- a/third_party/xla/xla/service/gpu/runtime/BUILD
+++ b/third_party/xla/xla/service/gpu/runtime/BUILD
@@ -85,6 +85,7 @@ cc_library(
         "//xla/service/gpu/kernels:custom_kernel",
         "//xla/service/gpu/runtime:thunk",
         "//xla/stream_executor",
+        "//xla/stream_executor:kernel_factory",
         "//xla/stream_executor:trace_command_buffer_factory",
         "//xla/stream_executor/gpu:gpu_stream_header",
         "//xla/stream_executor/gpu:gpu_types_header",
@@ -712,6 +713,7 @@ cc_library(
     srcs = ["kernel_thunk.cc"],
     hdrs = ["kernel_thunk.h"],
     deps = [
+        ":thunk",
         "//xla:status",
         "//xla:types",
         "//xla/hlo/ir:hlo",
@@ -720,8 +722,8 @@ cc_library(
         "//xla/service/gpu:launch_dimensions",
         "//xla/service/gpu:stream_executor_util",
         "//xla/service/gpu/kernels:custom_kernel",
-        "//xla/service/gpu/runtime:thunk",
         "//xla/stream_executor",
+        "//xla/stream_executor:kernel_factory",
         "@com_google_absl//absl/base:core_headers",
         "@com_google_absl//absl/container:flat_hash_map",
         "@com_google_absl//absl/container:inlined_vector",
@@ -729,7 +731,6 @@ cc_library(
         "@com_google_absl//absl/strings:str_format",
         "@com_google_absl//absl/synchronization",
         "@com_google_absl//absl/types:span",
-        "@llvm-project//mlir:IR",
         "@local_tsl//tsl/platform:errors",
         "@local_tsl//tsl/platform:logging",
         "@local_tsl//tsl/platform:statusor",
diff --git a/third_party/xla/xla/service/gpu/runtime/command_buffer_cmd.cc b/third_party/xla/xla/service/gpu/runtime/command_buffer_cmd.cc
index 5a6263fe3ef557..75a5d704c33bff 100644
--- a/third_party/xla/xla/service/gpu/runtime/command_buffer_cmd.cc
+++ b/third_party/xla/xla/service/gpu/runtime/command_buffer_cmd.cc
@@ -63,6 +63,7 @@ limitations under the License.
 #include "xla/stream_executor/device_memory.h"
 #include "xla/stream_executor/dnn.h"
 #include "xla/stream_executor/kernel.h"
+#include "xla/stream_executor/kernel_factory.h"
 #include "xla/stream_executor/launch_dim.h"
 #include "xla/stream_executor/stream.h"
 #include "xla/stream_executor/stream_executor.h"
@@ -670,7 +671,7 @@ absl::Status CustomKernelLaunchCmd::Initialize(
 
   TF_ASSIGN_OR_RETURN(
       std::unique_ptr<se::Kernel> kernel,
-      se::Kernel::Create(params.executor, custom_kernel_.kernel_spec()));
+      se::KernelFactory::Create(params.executor, custom_kernel_.kernel_spec()));
 
   absl::MutexLock lock(&mutex_);
   kernels_.emplace(params.executor, std::move(kernel));
diff --git a/third_party/xla/xla/service/gpu/runtime/kernel_thunk.cc b/third_party/xla/xla/service/gpu/runtime/kernel_thunk.cc
index 063940d8ab7cd4..e5cca1703a8e1d 100644
--- a/third_party/xla/xla/service/gpu/runtime/kernel_thunk.cc
+++ b/third_party/xla/xla/service/gpu/runtime/kernel_thunk.cc
@@ -38,6 +38,7 @@ limitations under the License.
 #include "xla/status.h"
 #include "xla/stream_executor/device_memory.h"
 #include "xla/stream_executor/kernel.h"
+#include "xla/stream_executor/kernel_factory.h"
 #include "xla/stream_executor/launch_dim.h"
 #include "xla/stream_executor/stream_executor.h"
 #include "tsl/platform/errors.h"
@@ -188,9 +189,9 @@ absl::Status CustomKernelThunk::Initialize(const InitializeParams& params) {
 
   auto it = kernel_cache_.find(params.executor);
   if (kernel_cache_.end() == it) {
-    TF_ASSIGN_OR_RETURN(
-        std::unique_ptr<se::Kernel> kernel,
-        se::Kernel::Create(params.executor, custom_kernel_.kernel_spec()));
+    TF_ASSIGN_OR_RETURN(std::unique_ptr<se::Kernel> kernel,
+                        se::KernelFactory::Create(
+                            params.executor, custom_kernel_.kernel_spec()));
     kernel_cache_.emplace(params.executor, std::move(kernel));
   }
 
diff --git a/third_party/xla/xla/service/gpu/stream_executor_util.cc b/third_party/xla/xla/service/gpu/stream_executor_util.cc
index 221b36c95683d2..8760faa725bcef 100644
--- a/third_party/xla/xla/service/gpu/stream_executor_util.cc
+++ b/third_party/xla/xla/service/gpu/stream_executor_util.cc
@@ -52,6 +52,7 @@ limitations under the License.
 #include "xla/stream_executor/device_memory.h"
 #include "xla/stream_executor/dnn.h"
 #include "xla/stream_executor/kernel.h"
+#include "xla/stream_executor/kernel_factory.h"
 #include "xla/stream_executor/kernel_spec.h"
 #include "xla/stream_executor/launch_dim.h"
 #include "xla/stream_executor/platform.h"
@@ -368,7 +369,7 @@ absl::StatusOr<std::unique_ptr<se::Kernel>> CreateKernel(
   }
 
   TF_ASSIGN_OR_RETURN(std::unique_ptr<se::Kernel> kernel,
-                      se::Kernel::Create(stream_exec, loader_spec));
+                      se::KernelFactory::Create(stream_exec, loader_spec));
 
   se::KernelMetadata m;
   m.set_shared_memory_bytes(shared_mem_bytes);
diff --git a/third_party/xla/xla/stream_executor/BUILD b/third_party/xla/xla/stream_executor/BUILD
index b73e82b9e42607..3899b903553264 100644
--- a/third_party/xla/xla/stream_executor/BUILD
+++ b/third_party/xla/xla/stream_executor/BUILD
@@ -607,6 +607,7 @@ cc_library(
         ":kernel_spec",
         ":platform",
         ":stream_executor_headers",
+        "@com_google_absl//absl/base:core_headers",
         "@com_google_absl//absl/container:inlined_vector",
         "@com_google_absl//absl/meta:type_traits",
         "@com_google_absl//absl/status",
@@ -620,6 +621,35 @@ cc_library(
     ],
 )
 
+cc_library(
+    name = "kernel_factory",
+    hdrs = ["kernel_factory.h"],
+    deps = [
+        ":kernel",
+        ":kernel_spec",
+        ":stream_executor_interface",
+        "@com_google_absl//absl/status:statusor",
+        "@local_tsl//tsl/platform:errors",
+        "@local_tsl//tsl/platform:statusor",
+    ],
+)
+
+cc_library(
+    name = "typed_kernel_factory",
+    hdrs = ["typed_kernel_factory.h"],
+    deps = [
+        ":kernel",
+        ":kernel_factory",
+        ":kernel_spec",
+        ":stream_executor_interface",
+        "@com_google_absl//absl/status:statusor",
+        "@com_google_absl//absl/strings:string_view",
+        "@com_google_absl//absl/types:span",
+        "@local_tsl//tsl/platform:errors",
+        "@local_tsl//tsl/platform:statusor",
+    ],
+)
+
 cc_library(
     name = "scratch_allocator",
     hdrs = ["scratch_allocator.h"],
diff --git a/third_party/xla/xla/stream_executor/gpu/BUILD b/third_party/xla/xla/stream_executor/gpu/BUILD
index 19b33618b223d3..8e70a01a43a6ed 100644
--- a/third_party/xla/xla/stream_executor/gpu/BUILD
+++ b/third_party/xla/xla/stream_executor/gpu/BUILD
@@ -668,6 +668,7 @@ xla_test(
         "//xla/stream_executor",
         "//xla/stream_executor:platform",
         "//xla/stream_executor:platform_manager",
+        "//xla/stream_executor:typed_kernel_factory",
         "@com_google_absl//absl/strings",
         "@com_google_googletest//:gtest",
         "@local_tsl//tsl/lib/core:status_test_util",
@@ -695,6 +696,7 @@ xla_test(
         "//xla/stream_executor:platform",
         "//xla/stream_executor:platform_manager",
         "//xla/stream_executor:trace_command_buffer_factory",
+        "//xla/stream_executor:typed_kernel_factory",
         "//xla/stream_executor/gpu:gpu_driver_header",
         "@com_google_absl//absl/log:check",
         "@com_google_absl//absl/status",
diff --git a/third_party/xla/xla/stream_executor/gpu/gpu_command_buffer_test.cc b/third_party/xla/xla/stream_executor/gpu/gpu_command_buffer_test.cc
index d1999c819f9927..14c5897118ead7 100644
--- a/third_party/xla/xla/stream_executor/gpu/gpu_command_buffer_test.cc
+++ b/third_party/xla/xla/stream_executor/gpu/gpu_command_buffer_test.cc
@@ -36,6 +36,7 @@ limitations under the License.
 #include "xla/stream_executor/stream.h"
 #include "xla/stream_executor/stream_executor.h"
 #include "xla/stream_executor/trace_command_buffer_factory.h"
+#include "xla/stream_executor/typed_kernel_factory.h"
 #include "tsl/lib/core/status_test_util.h"
 #include "tsl/platform/errors.h"
 #include "tsl/platform/status.h"
@@ -67,14 +68,16 @@ static MultiKernelLoaderSpec GetAddI32KernelSpec() {
   return spec;
 }
 
-using AddI32Kernel = TypedKernel<DeviceMemory<int32_t>, DeviceMemory<int32_t>,
-                                 DeviceMemory<int32_t>>;
-using MulI32Kernel = TypedKernel<DeviceMemory<int32_t>, DeviceMemory<int32_t>,
-                                 DeviceMemory<int32_t>>;
+using AddI32Kernel =
+    TypedKernelFactory<DeviceMemory<int32_t>, DeviceMemory<int32_t>,
+                       DeviceMemory<int32_t>>;
+using MulI32Kernel =
+    TypedKernelFactory<DeviceMemory<int32_t>, DeviceMemory<int32_t>,
+                       DeviceMemory<int32_t>>;
 using IncAndCmpKernel =
-    TypedKernel<DeviceMemory<int32_t>, DeviceMemory<bool>, int32_t>;
+    TypedKernelFactory<DeviceMemory<int32_t>, DeviceMemory<bool>, int32_t>;
 
-using AddI32Ptrs3 = TypedKernel<internal::Ptrs3<int32_t>>;
+using AddI32Ptrs3 = TypedKernelFactory<internal::Ptrs3<int32_t>>;
 
 static constexpr auto nested = CommandBuffer::Mode::kNested;    // NOLINT
 static constexpr auto primary = CommandBuffer::Mode::kPrimary;  // NOLINT
diff --git a/third_party/xla/xla/stream_executor/gpu/gpu_kernel_test.cc b/third_party/xla/xla/stream_executor/gpu/gpu_kernel_test.cc
index 591d417348776f..507fbfa477520f 100644
--- a/third_party/xla/xla/stream_executor/gpu/gpu_kernel_test.cc
+++ b/third_party/xla/xla/stream_executor/gpu/gpu_kernel_test.cc
@@ -26,6 +26,7 @@ limitations under the License.
 #include "xla/stream_executor/platform_manager.h"
 #include "xla/stream_executor/stream.h"
 #include "xla/stream_executor/stream_executor.h"
+#include "xla/stream_executor/typed_kernel_factory.h"
 #include "tsl/lib/core/status_test_util.h"
 #include "tsl/platform/statusor.h"
 #include "tsl/platform/test.h"
@@ -33,8 +34,9 @@ limitations under the License.
 namespace stream_executor::gpu {
 
 TEST(GpuKernelTest, Add) {
-  using AddI32Kernel = TypedKernel<DeviceMemory<int32_t>, DeviceMemory<int32_t>,
-                                   DeviceMemory<int32_t>>;
+  using AddI32Kernel =
+      TypedKernelFactory<DeviceMemory<int32_t>, DeviceMemory<int32_t>,
+                         DeviceMemory<int32_t>>;
   auto name = absl::AsciiStrToUpper(
       xla::PlatformUtil::CanonicalPlatformName("gpu").value());
   Platform* platform = PlatformManager::PlatformWithName(name).value();
diff --git a/third_party/xla/xla/stream_executor/kernel.h b/third_party/xla/xla/stream_executor/kernel.h
index 67871cc2ba1876..f90ad3fc50fbf5 100644
--- a/third_party/xla/xla/stream_executor/kernel.h
+++ b/third_party/xla/xla/stream_executor/kernel.h
@@ -1,3 +1,4 @@
+#include "absl/base/attributes.h"
 /* Copyright 2015 The OpenXLA Authors.
 
 Licensed under the Apache License, Version 2.0 (the "License");
@@ -232,6 +233,7 @@ class Kernel {
   // a dedicated KernelFactory accessible via StreamExecutor.
 
   // Creates kernel on a given executor from a given kernel specification.
+  ABSL_DEPRECATED("Use KernelFactory::Create instead.")
   static absl::StatusOr<std::unique_ptr<Kernel>> Create(
       StreamExecutorInterface *executor, const MultiKernelLoaderSpec &spec);
 
@@ -282,6 +284,8 @@ class Kernel {
 //===----------------------------------------------------------------------===//
 // Typed kernel
 //===----------------------------------------------------------------------===//
+template <typename... Params>
+class TypedKernelFactory;
 
 // Typed kernel is a typed smart-pointer-like wrapper around untyped Kernel.
 template <typename... Params>
@@ -290,6 +294,7 @@ class TypedKernel {
   static constexpr size_t kNumberOfParameters = sizeof...(Params);
 
   // Creates a typed kernel on a given executor from a kernel specification.
+  ABSL_DEPRECATED("Use TypedKernelFactory::Create instead.")
   static absl::StatusOr<TypedKernel> Create(StreamExecutorInterface *executor,
                                             const MultiKernelLoaderSpec &spec) {
     TF_ASSIGN_OR_RETURN(std::unique_ptr<Kernel> kernel,
@@ -302,18 +307,21 @@ class TypedKernel {
   // launch would have to match types of the arguments provided at creation
   // time. The canonical storage for both ptx and cubin_data should outlive the
   // lifetime of the kernel.
+  ABSL_DEPRECATED("Use TypedKernelFactory::Create instead.")
   static absl::StatusOr<TypedKernel> Create(
       StreamExecutorInterface *executor, absl::string_view kernel_name,
       absl::string_view ptx, absl::Span<const uint8_t> cubin_data);
 
   // Creates a kernel which can be launched with `stream.ThenLaunch(...)` from
   // an in-process symbol pointer.
+  ABSL_DEPRECATED("Use TypedKernelFactory::Create instead.")
   static absl::StatusOr<TypedKernel> Create(StreamExecutorInterface *executor,
                                             absl::string_view kernel_name,
                                             void *symbol);
 
   // Creates a kernel which can be launched with `stream.ThenLaunch(...)` from
   // an LLVM IR.
+  ABSL_DEPRECATED("Use TypedKernelFactory::Create instead.")
   static absl::StatusOr<TypedKernel> Create(StreamExecutorInterface *executor,
                                             absl::string_view ir,
                                             absl::string_view entrypoint,
@@ -331,6 +339,7 @@ class TypedKernel {
   operator bool() const { return static_cast<bool>(kernel_); }  // NOLINT
 
  private:
+  friend class TypedKernelFactory<Params...>;
   explicit TypedKernel(std::unique_ptr<Kernel> kernel)
       : kernel_(std::move(kernel)) {}
 
diff --git a/third_party/xla/xla/stream_executor/kernel_factory.h b/third_party/xla/xla/stream_executor/kernel_factory.h
new file mode 100644
index 00000000000000..2dd955efe5ed2c
--- /dev/null
+++ b/third_party/xla/xla/stream_executor/kernel_factory.h
@@ -0,0 +1,44 @@
+/* Copyright 2024 The OpenXLA Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+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.
+==============================================================================*/
+
+#ifndef XLA_STREAM_EXECUTOR_KERNEL_FACTORY_H_
+#define XLA_STREAM_EXECUTOR_KERNEL_FACTORY_H_
+
+#include <memory>
+
+#include "absl/status/statusor.h"
+#include "xla/stream_executor/kernel.h"
+#include "xla/stream_executor/kernel_spec.h"
+#include "xla/stream_executor/stream_executor_interface.h"
+#include "tsl/platform/errors.h"
+#include "tsl/platform/statusor.h"
+
+namespace stream_executor {
+
+// Creates Kernels from kernel specifications.
+class KernelFactory {
+ public:
+  // Creates kernel on a given executor from a given kernel specification.
+  static inline absl::StatusOr<std::unique_ptr<Kernel>> Create(
+      StreamExecutorInterface *executor, const MultiKernelLoaderSpec &spec) {
+    TF_ASSIGN_OR_RETURN(auto kernel, executor->CreateKernel());
+    TF_RETURN_IF_ERROR(executor->GetKernel(spec, kernel.get()));
+    return kernel;
+  }
+};
+
+}  // namespace stream_executor
+
+#endif  // XLA_STREAM_EXECUTOR_KERNEL_FACTORY_H_
diff --git a/third_party/xla/xla/stream_executor/typed_kernel_factory.h b/third_party/xla/xla/stream_executor/typed_kernel_factory.h
new file mode 100644
index 00000000000000..ac908b13b08e0c
--- /dev/null
+++ b/third_party/xla/xla/stream_executor/typed_kernel_factory.h
@@ -0,0 +1,95 @@
+/* Copyright 2024 The OpenXLA Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+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.
+==============================================================================*/
+
+#ifndef XLA_STREAM_EXECUTOR_TYPED_KERNEL_FACTORY_H_
+#define XLA_STREAM_EXECUTOR_TYPED_KERNEL_FACTORY_H_
+
+#include <cstdint>
+#include <memory>
+#include <string>
+#include <utility>
+
+#include "absl/status/statusor.h"
+#include "absl/strings/string_view.h"
+#include "absl/types/span.h"
+#include "xla/stream_executor/kernel.h"
+#include "xla/stream_executor/kernel_factory.h"
+#include "xla/stream_executor/kernel_spec.h"
+#include "xla/stream_executor/stream_executor_interface.h"
+#include "tsl/platform/statusor.h"
+
+namespace stream_executor {
+
+// This class creates TypedKernel objects for stream executors based on the
+// specification.
+template <typename... Params>
+class TypedKernelFactory {
+ public:
+  // Creates a typed kernel on a given executor from a kernel specification.
+  static absl::StatusOr<TypedKernel<Params...>> Create(
+      StreamExecutorInterface *executor, const MultiKernelLoaderSpec &spec) {
+    TF_ASSIGN_OR_RETURN(std::unique_ptr<Kernel> kernel,
+                        KernelFactory::Create(executor, spec));
+    return TypedKernel<Params...>(std::move(kernel));
+  }
+
+  // Creates a kernel which can be launched with `stream.ThenLaunch(...)` from a
+  // PTX (and optional CUBIN), such that the types of the arguments provided for
+  // launch would have to match types of the arguments provided at creation
+  // time. The canonical storage for both ptx and cubin_data should outlive the
+  // lifetime of the kernel.
+  static absl::StatusOr<TypedKernel<Params...>> Create(
+      StreamExecutorInterface *executor, absl::string_view kernel_name,
+      absl::string_view ptx, absl::Span<const uint8_t> cubin_data) {
+    MultiKernelLoaderSpec loader_spec(
+        TypedKernel<Params...>::kNumberOfParameters);
+    loader_spec.AddCudaPtxInMemory(ptx, kernel_name);
+
+    if (!cubin_data.empty()) {
+      loader_spec.AddCudaCubinInMemory(cubin_data, kernel_name);
+    }
+
+    return TypedKernel<Params...>::Create(executor, loader_spec);
+  }
+
+  // Creates a kernel which can be launched with `stream.ThenLaunch(...)` from
+  // an in-process symbol pointer.
+  static absl::StatusOr<TypedKernel<Params...>> Create(
+      StreamExecutorInterface *executor, absl::string_view kernel_name,
+      void *symbol) {
+    MultiKernelLoaderSpec loader_spec(
+        TypedKernel<Params...>::kNumberOfParameters);
+    loader_spec.AddInProcessSymbol(symbol, kernel_name);
+
+    return TypedKernel<Params...>::Create(executor, loader_spec);
+  }
+
+  // Creates a kernel which can be launched with `stream.ThenLaunch(...)` from
+  // an LLVM IR.
+  static absl::StatusOr<TypedKernel<Params...>> Create(
+      StreamExecutorInterface *executor, absl::string_view ir,
+      absl::string_view entrypoint, absl::string_view kernel_name,
+      absl::Span<std::string> options) {
+    MultiKernelLoaderSpec loader_spec(
+        TypedKernel<Params...>::kNumberOfParameters);
+    loader_spec.AddLlvmHostKernel(ir, entrypoint, kernel_name, options);
+
+    return TypedKernel<Params...>::Create(executor, loader_spec);
+  }
+};
+
+}  // namespace stream_executor
+
+#endif  // XLA_STREAM_EXECUTOR_TYPED_KERNEL_FACTORY_H_

From 872ada51766d5e6185dbcea16ed568eb42ee7c8b Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Tue, 14 May 2024 09:17:42 -0700
Subject: [PATCH 117/478] Update ops-related pbtxt files.

PiperOrigin-RevId: 633607491
---
 tensorflow/core/ops/ops.pbtxt | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorflow/core/ops/ops.pbtxt b/tensorflow/core/ops/ops.pbtxt
index 5019fe2a5b2fad..c8c55e08ca0d32 100644
--- a/tensorflow/core/ops/ops.pbtxt
+++ b/tensorflow/core/ops/ops.pbtxt
@@ -1,4 +1,4 @@
-go/debugproto    
+go/debugonly    
 op {
   name: "Abort"
   attr {

From 800071f14eab24e141392abb72671c861d4ab6b9 Mon Sep 17 00:00:00 2001
From: Kevin Gleason <gleasonk@google.com>
Date: Tue, 14 May 2024 10:03:21 -0700
Subject: [PATCH 118/478] Integrate StableHLO at openxla/stablehlo@797bee21

PiperOrigin-RevId: 633621411
---
 .../tests/compose-uniform-quantized-type.mlir |    14 +-
 .../uniform-quantized-stablehlo-to-tfl.mlir   |    34 +-
 .../common/attrs_and_constraints_test.cc      |     4 +-
 .../passes/quantize/quantize_same_scale.mlir  |    12 +-
 .../passes/quantize/quantize_weight_only.mlir |    12 +-
 third_party/stablehlo/temporary.patch         | 23389 +---------------
 third_party/stablehlo/workspace.bzl           |     4 +-
 .../xla/third_party/stablehlo/temporary.patch | 23389 +---------------
 .../xla/third_party/stablehlo/workspace.bzl   |     4 +-
 .../xla/xla/mlir_hlo/mhlo/IR/hlo_base.td      |     2 +
 .../xla/xla/mlir_hlo/mhlo/IR/hlo_ops.td       |    48 +-
 .../mhlo/transforms/map_stablehlo_to_hlo_op.h |     1 -
 .../mhlo/hlo-legalize-to-stablehlo.mlir       |    19 +-
 .../mhlo/mhlo-quant-legalize-to-int.mlir      |   200 +-
 .../xla/mlir_hlo/tests/Dialect/mhlo/ops.mlir  |   297 +-
 .../mhlo/stablehlo-legalize-to-hlo.mlir       |    19 +-
 16 files changed, 380 insertions(+), 47068 deletions(-)

diff --git a/tensorflow/compiler/mlir/lite/stablehlo/tests/compose-uniform-quantized-type.mlir b/tensorflow/compiler/mlir/lite/stablehlo/tests/compose-uniform-quantized-type.mlir
index 20d56eea578d79..7e60dc85a487a6 100644
--- a/tensorflow/compiler/mlir/lite/stablehlo/tests/compose-uniform-quantized-type.mlir
+++ b/tensorflow/compiler/mlir/lite/stablehlo/tests/compose-uniform-quantized-type.mlir
@@ -431,8 +431,8 @@ module {
     %2 = stablehlo.constant dense<-128> : tensor<1x1x1xi8>  // Input 1 zero point (z1).
     %3 = stablehlo.constant dense<-128> : tensor<1x1x1xi32>  // Input 1 zero point (z1) (upcast & folded into i32).
     %4 = stablehlo.constant dense<4.000000e-01> : tensor<1x1x1xf32>  // Input 2 inverse scale (1 / s2).
-    %5 = stablehlo.constant dense<-3> : tensor<1x1x1xi8>  // Input 2 zero point (z2).
-    %6 = stablehlo.constant dense<-3> : tensor<1x1x1xi32>  // Input 2 zero point (z2) (upcast & folded into i32).
+    %5 = stablehlo.constant dense<0> : tensor<1x1x1xi8>  // Input 2 zero point (z2).
+    %6 = stablehlo.constant dense<0> : tensor<1x1x1xi32>  // Input 2 zero point (z2) (upcast & folded into i32).
     %7 = stablehlo.constant dense<5.000000e-01> : tensor<1x1x1xf32>  // Output inverse scale (1 / s3).
     %8 = stablehlo.constant dense<-5> : tensor<1x1x1xi8>  // Output zero point (z3).
     %9 = stablehlo.constant dense<1.250000e+01> : tensor<1x1x1xf32>  // Merged scale (s1 * s2).
@@ -454,8 +454,8 @@ module {
     return %23 : tensor<8x16x4xf32>
   }
 // CHECK: %[[UQ_0:.*]] = stablehlo.uniform_quantize %[[ARG_0]] : (tensor<8x16x16xf32>) -> tensor<8x16x16x!quant.uniform<i8:f32, 5.000000e+00:-128>>
-// CHECK: %[[UQ_1:.*]] = stablehlo.uniform_quantize %[[ARG_1]] : (tensor<8x16x4xf32>) -> tensor<8x16x4x!quant.uniform<i8:f32, 2.500000e+00:-3>>
-// CHECK: %[[DOT_GENERAL:.*]] = stablehlo.dot_general %[[UQ_0]], %[[UQ_1]], batching_dims = [0] x [0], contracting_dims = [2] x [1] : (tensor<8x16x16x!quant.uniform<i8:f32, 5.000000e+00:-128>>, tensor<8x16x4x!quant.uniform<i8:f32, 2.500000e+00:-3>>) -> tensor<8x16x4x!quant.uniform<i8:f32, 2.000000e+00:-5>>
+// CHECK: %[[UQ_1:.*]] = stablehlo.uniform_quantize %[[ARG_1]] : (tensor<8x16x4xf32>) -> tensor<8x16x4x!quant.uniform<i8:f32, 2.500000e+00>>
+// CHECK: %[[DOT_GENERAL:.*]] = stablehlo.dot_general %[[UQ_0]], %[[UQ_1]], batching_dims = [0] x [0], contracting_dims = [2] x [1] : (tensor<8x16x16x!quant.uniform<i8:f32, 5.000000e+00:-128>>, tensor<8x16x4x!quant.uniform<i8:f32, 2.500000e+00>>) -> tensor<8x16x4x!quant.uniform<i8:f32, 2.000000e+00:-5>>
 // CHECK: %[[DQ_0:.*]] = stablehlo.uniform_dequantize %[[DOT_GENERAL]] : (tensor<8x16x4x!quant.uniform<i8:f32, 2.000000e+00:-5>>) -> tensor<8x16x4xf32>
 // CHECK: return %[[DQ_0]]
 
@@ -492,7 +492,7 @@ module {
     %1 = stablehlo.constant dense<2.000000e-01> : tensor<1x1x1xf32>  // Input 1 inverse scale (1 / s1).
     %2 = stablehlo.constant dense<-128> : tensor<1x1x1xi8>  // Input 1 zero point (z1).
     %3 = stablehlo.constant dense<4.000000e-01> : tensor<1x1x1xf32>  // Input 2 inverse scale (1 / s2).
-    %4 = stablehlo.constant dense<-3> : tensor<1x1x1xi8>  // Input 2 zero point (z2).
+    %4 = stablehlo.constant dense<0> : tensor<1x1x1xi8>  // Input 2 zero point (z2).
     %5 = stablehlo.constant dense<5.000000e-01> : tensor<1x1x1xf32>  // Output inverse scale (1 / s3).
     %6 = stablehlo.constant dense<-5> : tensor<1x1x1xi8>  // Output zero point (z3).
     %7 = stablehlo.constant dense<1.250000e+01> : tensor<1x1x1xf32>  // Merged scale (s1 * s2).
@@ -516,8 +516,8 @@ module {
     return %23 : tensor<8x16x4xf32>
   }
 // CHECK: %[[UQ_0:.*]] = stablehlo.uniform_quantize %[[ARG_0]] : (tensor<8x16x16xf32>) -> tensor<8x16x16x!quant.uniform<i8:f32, 5.000000e+00:-128>>
-// CHECK: %[[UQ_1:.*]] = stablehlo.uniform_quantize %[[ARG_1]] : (tensor<8x16x4xf32>) -> tensor<8x16x4x!quant.uniform<i8:f32, 2.500000e+00:-3>>
-// CHECK: %[[DOT_GENERAL:.*]] = stablehlo.dot_general %[[UQ_0]], %[[UQ_1]], batching_dims = [0] x [0], contracting_dims = [2] x [1] : (tensor<8x16x16x!quant.uniform<i8:f32, 5.000000e+00:-128>>, tensor<8x16x4x!quant.uniform<i8:f32, 2.500000e+00:-3>>) -> tensor<8x16x4x!quant.uniform<i8:f32, 2.000000e+00:-5>>
+// CHECK: %[[UQ_1:.*]] = stablehlo.uniform_quantize %[[ARG_1]] : (tensor<8x16x4xf32>) -> tensor<8x16x4x!quant.uniform<i8:f32, 2.500000e+00>>
+// CHECK: %[[DOT_GENERAL:.*]] = stablehlo.dot_general %[[UQ_0]], %[[UQ_1]], batching_dims = [0] x [0], contracting_dims = [2] x [1] : (tensor<8x16x16x!quant.uniform<i8:f32, 5.000000e+00:-128>>, tensor<8x16x4x!quant.uniform<i8:f32, 2.500000e+00>>) -> tensor<8x16x4x!quant.uniform<i8:f32, 2.000000e+00:-5>>
 // CHECK: %[[DQ_0:.*]] = stablehlo.uniform_dequantize %[[DOT_GENERAL]] : (tensor<8x16x4x!quant.uniform<i8:f32, 2.000000e+00:-5>>) -> tensor<8x16x4xf32>
 // CHECK: return %[[DQ_0]]
 
diff --git a/tensorflow/compiler/mlir/lite/stablehlo/tests/uniform-quantized-stablehlo-to-tfl.mlir b/tensorflow/compiler/mlir/lite/stablehlo/tests/uniform-quantized-stablehlo-to-tfl.mlir
index dde72965ae65fc..64b14b85fc7c71 100644
--- a/tensorflow/compiler/mlir/lite/stablehlo/tests/uniform-quantized-stablehlo-to-tfl.mlir
+++ b/tensorflow/compiler/mlir/lite/stablehlo/tests/uniform-quantized-stablehlo-to-tfl.mlir
@@ -393,13 +393,13 @@ func.func @dot_general_upstream_srq_float_operands(%arg0: tensor<1x2x3x4xf32>, %
 
 // CHECK-LABEL: dot_general_upstream_srq_asym_weight
 func.func @dot_general_upstream_srq_asym_weight(%arg0: tensor<1x2x3x4x!quant.uniform<i8:f32, 1.000000e+0:-100>>) -> tensor<1x2x3x5x!quant.uniform<i8:f32, 4.000000e+0>> {
-  %0 = stablehlo.constant() {value = dense<1> : tensor<1x2x4x5xi8>} : () -> tensor<1x2x4x5x!quant.uniform<i8:f32, 1.000000e+0:5>>
-  %1 = "stablehlo.dot_general"(%arg0, %0) {dot_dimension_numbers = #stablehlo.dot<lhs_batching_dimensions = [0, 1], rhs_batching_dimensions = [0, 1], lhs_contracting_dimensions = [3], rhs_contracting_dimensions = [2]>, precision_config = [#stablehlo<precision DEFAULT>, #stablehlo<precision DEFAULT>]} : (tensor<1x2x3x4x!quant.uniform<i8:f32, 1.000000e+0:-100>>, tensor<1x2x4x5x!quant.uniform<i8:f32, 1.000000e+0:5>>) -> tensor<1x2x3x5x!quant.uniform<i8:f32, 4.000000e+0>>
+  %0 = stablehlo.constant() {value = dense<1> : tensor<1x2x4x5xi8>} : () -> tensor<1x2x4x5x!quant.uniform<i8:f32, 1.000000e+0:0>>
+  %1 = "stablehlo.dot_general"(%arg0, %0) {dot_dimension_numbers = #stablehlo.dot<lhs_batching_dimensions = [0, 1], rhs_batching_dimensions = [0, 1], lhs_contracting_dimensions = [3], rhs_contracting_dimensions = [2]>, precision_config = [#stablehlo<precision DEFAULT>, #stablehlo<precision DEFAULT>]} : (tensor<1x2x3x4x!quant.uniform<i8:f32, 1.000000e+0:-100>>, tensor<1x2x4x5x!quant.uniform<i8:f32, 1.000000e+0:0>>) -> tensor<1x2x3x5x!quant.uniform<i8:f32, 4.000000e+0>>
   return %1 : tensor<1x2x3x5x!quant.uniform<i8:f32, 4.000000e+0>>
 }
 // CHECK-SAME: %[[ARG:.+]]: tensor<1x2x3x4x!quant.uniform<i8:f32, 1.000000e+00:-100>>
-// CHECK: %[[QCONST_0:.+]] =  "tfl.pseudo_qconst"() <{qtype = tensor<1x2x4x5x!quant.uniform<i8:f32, 1.000000e+00:5>>, value = dense<1> : tensor<1x2x4x5xi8>}> : () -> tensor<1x2x4x5x!quant.uniform<i8:f32, 1.000000e+00:5>>
-// CHECK: %[[BMM:.+]] = "tfl.batch_matmul"(%[[ARG]], %[[QCONST_0]]) <{adj_x = false, adj_y = false}> : (tensor<1x2x3x4x!quant.uniform<i8:f32, 1.000000e+00:-100>>, tensor<1x2x4x5x!quant.uniform<i8:f32, 1.000000e+00:5>>) -> tensor<1x2x3x5x!quant.uniform<i8:f32, 4.000000e+00>>
+// CHECK: %[[QCONST_0:.+]] =  "tfl.pseudo_qconst"() <{qtype = tensor<1x2x4x5x!quant.uniform<i8:f32, 1.000000e+00>>, value = dense<1> : tensor<1x2x4x5xi8>}> : () -> tensor<1x2x4x5x!quant.uniform<i8:f32, 1.000000e+00>>
+// CHECK: %[[BMM:.+]] = "tfl.batch_matmul"(%[[ARG]], %[[QCONST_0]]) <{adj_x = false, adj_y = false}> : (tensor<1x2x3x4x!quant.uniform<i8:f32, 1.000000e+00:-100>>, tensor<1x2x4x5x!quant.uniform<i8:f32, 1.000000e+00>>) -> tensor<1x2x3x5x!quant.uniform<i8:f32, 4.000000e+00>>
 
 // -----
 
@@ -427,8 +427,8 @@ func.func @dot_general_upstream_srq_per_axis_quantized_filter(%arg0: tensor<1x3x
 
 // CHECK-LABEL: dot_general_upstream_srq_per_axis_quantized_filter_with_batch_dim
 func.func @dot_general_upstream_srq_per_axis_quantized_filter_with_batch_dim(%arg0: tensor<1x1x3x!quant.uniform<i8:f32, 5.000000e+05:-100>>) -> tensor<1x1x2x!quant.uniform<i8:f32, 4.000000e+04:127>> {
-  %0 = stablehlo.constant() {value = dense<1> : tensor<1x3x2xi8>} : () -> tensor<1x3x2x!quant.uniform<i8:f32:2,{2.000000e+02, 3.000000e+03}>>
-  %1 = stablehlo.dot_general %arg0, %0, batching_dims = [0] x [0], contracting_dims = [2] x [1] : (tensor<1x1x3x!quant.uniform<i8:f32, 5.000000e+05:-100>>, tensor<1x3x2x!quant.uniform<i8:f32:2,{2.000000e+02, 3.000000e+03}>>) -> tensor<1x1x2x!quant.uniform<i8:f32, 4.000000e+04:127>>
+  %0 = stablehlo.constant() {value = dense<1> : tensor<1x3x2xi8>} : () -> tensor<1x3x2x!quant.uniform<i8:f32:0,{2.000000e+02}>>
+  %1 = stablehlo.dot_general %arg0, %0, batching_dims = [0] x [0], contracting_dims = [2] x [1] : (tensor<1x1x3x!quant.uniform<i8:f32, 5.000000e+05:-100>>, tensor<1x3x2x!quant.uniform<i8:f32:0,{2.000000e+02}>>) -> tensor<1x1x2x!quant.uniform<i8:f32, 4.000000e+04:127>>
   return %1 : tensor<1x1x2x!quant.uniform<i8:f32, 4.000000e+04:127>>
 }
 // Nothing changes.
@@ -459,8 +459,8 @@ func.func @dot_general_upstream_srq_per_axis_quantized_filter_multibatch(%arg0:
 
 // CHECK-LABEL: dot_general_upstream_srq_per_axis_quantized_filter_with_multiple_contracting_dims
 func.func @dot_general_upstream_srq_per_axis_quantized_filter_with_multiple_contracting_dims(%arg0: tensor<1x2x3x!quant.uniform<i8:f32, 5.000000e+05:-100>>) -> tensor<1x1x!quant.uniform<i8:f32, 4.000000e+04:127>> {
-  %0 = stablehlo.constant() {value = dense<1> : tensor<1x3x2xi8>} : () -> tensor<1x3x2x!quant.uniform<i8:f32:2,{2.000000e+02, 3.000000e+03}>>
-  %1 = stablehlo.dot_general %arg0, %0, contracting_dims = [1, 2] x [2, 1] : (tensor<1x2x3x!quant.uniform<i8:f32, 5.000000e+05:-100>>, tensor<1x3x2x!quant.uniform<i8:f32:2,{2.000000e+02, 3.000000e+03}>>) -> tensor<1x1x!quant.uniform<i8:f32, 4.000000e+04:127>>
+  %0 = stablehlo.constant() {value = dense<1> : tensor<1x3x2xi8>} : () -> tensor<1x3x2x!quant.uniform<i8:f32:0,{2.000000e+02}>>
+  %1 = stablehlo.dot_general %arg0, %0, contracting_dims = [1, 2] x [2, 1] : (tensor<1x2x3x!quant.uniform<i8:f32, 5.000000e+05:-100>>, tensor<1x3x2x!quant.uniform<i8:f32:0,{2.000000e+02}>>) -> tensor<1x1x!quant.uniform<i8:f32, 4.000000e+04:127>>
   return %1 : tensor<1x1x!quant.uniform<i8:f32, 4.000000e+04:127>>
 }
 // Nothing changes.
@@ -557,9 +557,9 @@ func.func @dot_general_srq_constant_transpose_rhs(%arg0: tensor<1x3x!quant.unifo
 // (e.g. argument), the conversion to `tfl.fully_connected` doesn't happen.
 
 // CHECK-LABEL: dot_general_srq_arg_transpose_rhs
-func.func @dot_general_srq_arg_transpose_rhs(%arg0: tensor<1x3x!quant.uniform<i8:f32, 5.000000e+00:-128>>, %arg1: tensor<2x3x!quant.uniform<i8:f32, 3.000000e+00:-23>>) -> tensor<1x2x!quant.uniform<i8:f32, 4.000000e+00:7>> {
-  %1 = stablehlo.transpose %arg1, dims = [1, 0] : (tensor<2x3x!quant.uniform<i8:f32, 3.000000e+00:-23>>) -> tensor<3x2x!quant.uniform<i8:f32, 3.000000e+00:-23>>
-  %2 = stablehlo.dot_general %arg0, %1, contracting_dims = [1] x [0] : (tensor<1x3x!quant.uniform<i8:f32, 5.000000e+00:-128>>, tensor<3x2x!quant.uniform<i8:f32, 3.000000e+00:-23>>) -> tensor<1x2x!quant.uniform<i32:f32, 2.000000e+00>>
+func.func @dot_general_srq_arg_transpose_rhs(%arg0: tensor<1x3x!quant.uniform<i8:f32, 5.000000e+00:-128>>, %arg1: tensor<2x3x!quant.uniform<i8:f32, 3.000000e+00>>) -> tensor<1x2x!quant.uniform<i8:f32, 4.000000e+00:7>> {
+  %1 = stablehlo.transpose %arg1, dims = [1, 0] : (tensor<2x3x!quant.uniform<i8:f32, 3.000000e+00>>) -> tensor<3x2x!quant.uniform<i8:f32, 3.000000e+00>>
+  %2 = stablehlo.dot_general %arg0, %1, contracting_dims = [1] x [0] : (tensor<1x3x!quant.uniform<i8:f32, 5.000000e+00:-128>>, tensor<3x2x!quant.uniform<i8:f32, 3.000000e+00>>) -> tensor<1x2x!quant.uniform<i32:f32, 2.000000e+00>>
   %3 = stablehlo.uniform_quantize %2 : (tensor<1x2x!quant.uniform<i32:f32, 2.000000e+00>>) -> tensor<1x2x!quant.uniform<i8:f32, 4.000000e+00:7>>
   return %3 : tensor<1x2x!quant.uniform<i8:f32, 4.000000e+00:7>>
 }
@@ -577,7 +577,7 @@ func.func @dot_general_srq_arg_transpose_rhs(%arg0: tensor<1x3x!quant.uniform<i8
 // Tests static range quantized dot_general with qi32 -> qi8 requantization is
 // properly lowered to `tfl.batch_matmul`.
 
-func.func @dot_general_srq_to_batch_matmul(%arg0: tensor<1x2x3x4x!quant.uniform<i8:f32, 1.000000e+00:3>>, %arg1: tensor<1x2x4x5x!quant.uniform<i8:f32, 2.000000e+00:4>>) -> tensor<1x2x3x5x!quant.uniform<i8:f32, 6.000000e+00:5>> {
+func.func @dot_general_srq_to_batch_matmul(%arg0: tensor<1x2x3x4x!quant.uniform<i8:f32, 1.000000e+00:3>>, %arg1: tensor<1x2x4x5x!quant.uniform<i8:f32, 2.000000e+00>>) -> tensor<1x2x3x5x!quant.uniform<i8:f32, 6.000000e+00:5>> {
   %0 = "stablehlo.dot_general"(%arg0, %arg1) {
     dot_dimension_numbers = #stablehlo.dot<
       lhs_batching_dimensions = [0, 1],
@@ -586,14 +586,14 @@ func.func @dot_general_srq_to_batch_matmul(%arg0: tensor<1x2x3x4x!quant.uniform<
       rhs_contracting_dimensions = [2]
     >,
     precision_config = [#stablehlo<precision DEFAULT>, #stablehlo<precision DEFAULT>]
-  } : (tensor<1x2x3x4x!quant.uniform<i8:f32, 1.000000e+00:3>>, tensor<1x2x4x5x!quant.uniform<i8:f32, 2.000000e+00:4>>) -> tensor<1x2x3x5x!quant.uniform<i32:f32, 4.000000e+00:1>>
+  } : (tensor<1x2x3x4x!quant.uniform<i8:f32, 1.000000e+00:3>>, tensor<1x2x4x5x!quant.uniform<i8:f32, 2.000000e+00>>) -> tensor<1x2x3x5x!quant.uniform<i32:f32, 4.000000e+00:1>>
   %1 = stablehlo.uniform_quantize %0 : (tensor<1x2x3x5x!quant.uniform<i32:f32, 4.000000e+00:1>>) -> tensor<1x2x3x5x!quant.uniform<i8:f32, 6.000000e+00:5>>
   return %1 : tensor<1x2x3x5x!quant.uniform<i8:f32, 6.000000e+00:5>>
 }
 
 // CHECK-LABEL: dot_general_srq_to_batch_matmul
-// CHECK-SAME: (%[[ARG_0:.+]]: tensor<1x2x3x4x!quant.uniform<i8:f32, 1.000000e+00:3>>, %[[ARG_1:.+]]: tensor<1x2x4x5x!quant.uniform<i8:f32, 2.000000e+00:4>>) -> tensor<1x2x3x5x!quant.uniform<i8:f32, 6.000000e+00:5>>
-// CHECK: %[[BMM:.+]] = "tfl.batch_matmul"(%[[ARG_0]], %[[ARG_1]]) <{adj_x = false, adj_y = false}> : (tensor<1x2x3x4x!quant.uniform<i8:f32, 1.000000e+00:3>>, tensor<1x2x4x5x!quant.uniform<i8:f32, 2.000000e+00:4>>) -> tensor<1x2x3x5x!quant.uniform<i8:f32, 6.000000e+00:5>>
+// CHECK-SAME: (%[[ARG_0:.+]]: tensor<1x2x3x4x!quant.uniform<i8:f32, 1.000000e+00:3>>, %[[ARG_1:.+]]: tensor<1x2x4x5x!quant.uniform<i8:f32, 2.000000e+00>>) -> tensor<1x2x3x5x!quant.uniform<i8:f32, 6.000000e+00:5>>
+// CHECK: %[[BMM:.+]] = "tfl.batch_matmul"(%[[ARG_0]], %[[ARG_1]]) <{adj_x = false, adj_y = false}> : (tensor<1x2x3x4x!quant.uniform<i8:f32, 1.000000e+00:3>>, tensor<1x2x4x5x!quant.uniform<i8:f32, 2.000000e+00>>) -> tensor<1x2x3x5x!quant.uniform<i8:f32, 6.000000e+00:5>>
 // CHECK-NOT: stablehlo.dot_general
 // CHECK-NOT: stablehlo.uniform_quantize
 // CHECK-NOT: tfl.fully_connected
@@ -606,7 +606,7 @@ func.func @dot_general_srq_to_batch_matmul(%arg0: tensor<1x2x3x4x!quant.uniform<
 // not converted to `tfl.batch_matmul` when there are multiple use of the
 // intermediate result.
 
-func.func @dot_general_srq_multiple_use_of_intermediate_result(%arg0: tensor<1x2x3x4x!quant.uniform<i8:f32, 1.000000e+00:3>>, %arg1: tensor<1x2x4x5x!quant.uniform<i8:f32, 2.000000e+00:4>>) -> tensor<1x2x3x5x!quant.uniform<i8:f32, 6.000000e+00:5>> {
+func.func @dot_general_srq_multiple_use_of_intermediate_result(%arg0: tensor<1x2x3x4x!quant.uniform<i8:f32, 1.000000e+00:3>>, %arg1: tensor<1x2x4x5x!quant.uniform<i8:f32, 2.000000e+00>>) -> tensor<1x2x3x5x!quant.uniform<i8:f32, 6.000000e+00:5>> {
   %0 = "stablehlo.dot_general"(%arg0, %arg1) {
     dot_dimension_numbers = #stablehlo.dot<
       lhs_batching_dimensions = [0, 1],
@@ -615,7 +615,7 @@ func.func @dot_general_srq_multiple_use_of_intermediate_result(%arg0: tensor<1x2
       rhs_contracting_dimensions = [2]
     >,
     precision_config = [#stablehlo<precision DEFAULT>, #stablehlo<precision DEFAULT>]
-  } : (tensor<1x2x3x4x!quant.uniform<i8:f32, 1.000000e+00:3>>, tensor<1x2x4x5x!quant.uniform<i8:f32, 2.000000e+00:4>>) -> tensor<1x2x3x5x!quant.uniform<i32:f32, 4.000000e+00:1>>
+  } : (tensor<1x2x3x4x!quant.uniform<i8:f32, 1.000000e+00:3>>, tensor<1x2x4x5x!quant.uniform<i8:f32, 2.000000e+00>>) -> tensor<1x2x3x5x!quant.uniform<i32:f32, 4.000000e+00:1>>
   %1 = stablehlo.uniform_quantize %0 : (tensor<1x2x3x5x!quant.uniform<i32:f32, 4.000000e+00:1>>) -> tensor<1x2x3x5x!quant.uniform<i8:f32, 6.000000e+00:5>>
   %2 = stablehlo.uniform_quantize %0 : (tensor<1x2x3x5x!quant.uniform<i32:f32, 4.000000e+00:1>>) -> tensor<1x2x3x5x!quant.uniform<i8:f32, 6.000000e+00:5>>
   %3 = stablehlo.add %1, %2 : tensor<1x2x3x5x!quant.uniform<i8:f32, 6.000000e+00:5>>
diff --git a/tensorflow/compiler/mlir/quantization/common/attrs_and_constraints_test.cc b/tensorflow/compiler/mlir/quantization/common/attrs_and_constraints_test.cc
index ad0fc81157401f..720616309afe38 100644
--- a/tensorflow/compiler/mlir/quantization/common/attrs_and_constraints_test.cc
+++ b/tensorflow/compiler/mlir/quantization/common/attrs_and_constraints_test.cc
@@ -143,8 +143,8 @@ constexpr absl::string_view kModulePartitionedCall = R"mlir(
 
 constexpr absl::string_view kModuleHybridQuantized = R"mlir(
   module {
-    func.func @main(%arg0: tensor<1x2xf32>, %arg1: tensor<2x3x!quant.uniform<i8:f32, 6.000000e-03:-128>> {tf_saved_model.index_path = ["input_tensor"]}) -> (tensor<1x3xf32>) {
-      %0 = stablehlo.dot_general %arg0, %arg1, contracting_dims = [1] x [0] : (tensor<1x2xf32>, tensor<2x3x!quant.uniform<i8:f32, 6.000000e-03:-128>>) -> tensor<1x3xf32>
+    func.func @main(%arg0: tensor<1x2xf32>, %arg1: tensor<2x3x!quant.uniform<i8:f32, 6.000000e-03:0>> {tf_saved_model.index_path = ["input_tensor"]}) -> (tensor<1x3xf32>) {
+      %0 = stablehlo.dot_general %arg0, %arg1, contracting_dims = [1] x [0] : (tensor<1x2xf32>, tensor<2x3x!quant.uniform<i8:f32, 6.000000e-03:0>>) -> tensor<1x3xf32>
       return %0 : tensor<1x3xf32>
     }
   }
diff --git a/tensorflow/compiler/mlir/quantization/stablehlo/tests/passes/quantize/quantize_same_scale.mlir b/tensorflow/compiler/mlir/quantization/stablehlo/tests/passes/quantize/quantize_same_scale.mlir
index 25aab3044a3496..7a905dfbe58a9e 100644
--- a/tensorflow/compiler/mlir/quantization/stablehlo/tests/passes/quantize/quantize_same_scale.mlir
+++ b/tensorflow/compiler/mlir/quantization/stablehlo/tests/passes/quantize/quantize_same_scale.mlir
@@ -291,17 +291,17 @@ module attributes {tf_saved_model.semantics} {
   // CHECK-SAME: %[[ARG2:.*]]: tensor<2x3x2xi64>
   func.func private @composite_and_gather(%arg0: tensor<3x4x5xf32>, %arg1: tensor<3x5x2xf32>, %arg2: tensor<2x3x2xi64>) -> tensor<2x3x2x2xf32> {
     // CHECK: %[[Q1:.*]] = "quantfork.qcast"(%[[ARG0]]) {volatile} : (tensor<3x4x5xf32>) -> tensor<3x4x5x!quant.uniform<i8:f32, 5.000000e-03>>
-    // CHECK: %[[Q2:.*]] = "quantfork.qcast"(%[[ARG1]]) {volatile} : (tensor<3x5x2xf32>) -> tensor<3x5x2x!quant.uniform<i8<-127:127>:f32, 6.000000e-03:13>>
+    // CHECK: %[[Q2:.*]] = "quantfork.qcast"(%[[ARG1]]) {volatile} : (tensor<3x5x2xf32>) -> tensor<3x5x2x!quant.uniform<i8<-127:127>:f32, 6.000000e-03>>
     // CHECK: %[[CALL:.*]] = call @quantized_dot_general_fn_1(%[[Q1]], %[[Q2]])
-    // CHECK-SAME: (tensor<3x4x5x!quant.uniform<i8:f32, 5.000000e-03>>, tensor<3x5x2x!quant.uniform<i8<-127:127>:f32, 6.000000e-03:13>>) -> tensor<3x4x2x!quant.uniform<i8:f32, 0.13170163023705575:-1>>
+    // CHECK-SAME: (tensor<3x4x5x!quant.uniform<i8:f32, 5.000000e-03>>, tensor<3x5x2x!quant.uniform<i8<-127:127>:f32, 6.000000e-03>>) -> tensor<3x4x2x!quant.uniform<i8:f32, 0.13170163023705575:-1>>
     // CHECK: %[[GATHER:.*]] = "stablehlo.gather"(%[[CALL]], %[[ARG2]])
     // CHECK-SAME: (tensor<3x4x2x!quant.uniform<i8:f32, 0.13170163023705575:-1>>, tensor<2x3x2xi64>) -> tensor<2x3x2x2x!quant.uniform<i8:f32, 0.13170163023705575:-1>>
     // CHECK: %[[DQ:.*]] = "quantfork.dcast"(%[[GATHER]]) : (tensor<2x3x2x2x!quant.uniform<i8:f32, 0.13170163023705575:-1>>) -> tensor<2x3x2x2xf32>
     // CHECK: return %[[DQ]]
     %0 = "quantfork.qcast"(%arg0) {volatile} : (tensor<3x4x5xf32>) -> tensor<3x4x5x!quant.uniform<i8:f32, 5.000000e-03>>
     %1 = "quantfork.dcast"(%0) : (tensor<3x4x5x!quant.uniform<i8:f32, 5.000000e-03>>) -> tensor<3x4x5xf32>
-    %2 = "quantfork.qcast"(%arg1) {volatile} : (tensor<3x5x2xf32>) -> tensor<3x5x2x!quant.uniform<i8<-127:127>:f32, 6.000000e-03:13>>
-    %3 = "quantfork.dcast"(%2) : (tensor<3x5x2x!quant.uniform<i8<-127:127>:f32, 6.000000e-03:13>>) -> tensor<3x5x2xf32>
+    %2 = "quantfork.qcast"(%arg1) {volatile} : (tensor<3x5x2xf32>) -> tensor<3x5x2x!quant.uniform<i8<-127:127>:f32, 6.000000e-03>>
+    %3 = "quantfork.dcast"(%2) : (tensor<3x5x2x!quant.uniform<i8<-127:127>:f32, 6.000000e-03>>) -> tensor<3x5x2xf32>
     %4 = "tf.XlaCallModule"(%1, %3) {Sout = [#tf_type.shape<1x3>], _entry_function = @composite_dot_general_fn_1, _original_entry_function = "composite_dot_general_fn_1", _quantization_method = "static_range_ptq {}", _stablehlo_module_attrs = {}, _tfl_quant_trait = "fully_quantizable",   device = "", dim_args_spec = [], disabled_checks = [], has_token_input_output = false, module = "", platforms = [], version = 5 : i64} : (tensor<3x4x5xf32>, tensor<3x5x2xf32>) -> tensor<3x4x2xf32>
     %5 = "quantfork.qcast"(%4) {volatile} : (tensor<3x4x2xf32>) -> tensor<3x4x2x!quant.uniform<i8:f32, 0.13170163023705575:-1>>
     %6 = "quantfork.dcast"(%5) : (tensor<3x4x2x!quant.uniform<i8:f32, 0.13170163023705575:-1>>) -> tensor<3x4x2xf32>
@@ -321,10 +321,10 @@ module attributes {tf_saved_model.semantics} {
 
   // CHECK: quantized_dot_general_fn_1
   // CHECK-SAME: %[[ARG2:.*]]: tensor<3x4x5x!quant.uniform<i8:f32, 5.000000e-03>>
-  // CHECK-SAME: %[[ARG3:.*]]: tensor<3x5x2x!quant.uniform<i8<-127:127>:f32, 6.000000e-03:13>>
+  // CHECK-SAME: %[[ARG3:.*]]: tensor<3x5x2x!quant.uniform<i8<-127:127>:f32, 6.000000e-03>>
   func.func private @composite_dot_general_fn_1(%arg0: tensor<3x4x5xf32>, %arg1: tensor<3x5x2xf32>) -> tensor<3x4x2xf32> attributes {_from_xla_call_module} {
     // CHECK: %[[DOT:.*]] = stablehlo.dot_general %[[ARG2]], %[[ARG3]]
-    // CHECK-SAME: (tensor<3x4x5x!quant.uniform<i8:f32, 5.000000e-03>>, tensor<3x5x2x!quant.uniform<i8<-127:127>:f32, 6.000000e-03:13>>) -> tensor<3x4x2x!quant.uniform<i32:f32, 3.000000e-05>>
+    // CHECK-SAME: (tensor<3x4x5x!quant.uniform<i8:f32, 5.000000e-03>>, tensor<3x5x2x!quant.uniform<i8<-127:127>:f32, 6.000000e-03>>) -> tensor<3x4x2x!quant.uniform<i32:f32, 3.000000e-05>>
     // CHECK: %[[Q3:.*]] = stablehlo.uniform_quantize %0 : (tensor<3x4x2x!quant.uniform<i32:f32, 3.000000e-05>>) -> tensor<3x4x2x!quant.uniform<i8:f32, 0.13170163023705575:-1>>
     // CHECK: return %[[Q3]]
     %0 = stablehlo.dot_general %arg0, %arg1, batching_dims = [0] x [0], contracting_dims = [2] x [1] : (tensor<3x4x5xf32>, tensor<3x5x2xf32>) -> tensor<3x4x2xf32>
diff --git a/tensorflow/compiler/mlir/quantization/stablehlo/tests/passes/quantize/quantize_weight_only.mlir b/tensorflow/compiler/mlir/quantization/stablehlo/tests/passes/quantize/quantize_weight_only.mlir
index 78a48a2d9373e1..e152a90ce72c3a 100644
--- a/tensorflow/compiler/mlir/quantization/stablehlo/tests/passes/quantize/quantize_weight_only.mlir
+++ b/tensorflow/compiler/mlir/quantization/stablehlo/tests/passes/quantize/quantize_weight_only.mlir
@@ -6,8 +6,8 @@
 module attributes {tf_saved_model.semantics} {
   func.func private @quantize_dot_general_fn(%arg0: tensor<1x2xf32>) -> tensor<1x3xf32> attributes {tf._original_func_name = "main_0"} {
     %cst = stablehlo.constant dense<3.000000e-01> : tensor<2x3xf32>
-    %0 = "quantfork.qcast"(%cst) : (tensor<2x3xf32>) -> tensor<2x3x!quant.uniform<i8:f32, 6.000000e-03:-128>>
-    %1 = "quantfork.dcast"(%0) : (tensor<2x3x!quant.uniform<i8:f32, 6.000000e-03:-128>>) -> tensor<2x3xf32>
+    %0 = "quantfork.qcast"(%cst) : (tensor<2x3xf32>) -> tensor<2x3x!quant.uniform<i8:f32, 6.000000e-03>>
+    %1 = "quantfork.dcast"(%0) : (tensor<2x3x!quant.uniform<i8:f32, 6.000000e-03>>) -> tensor<2x3xf32>
     %2 = "tf.XlaCallModule"(%arg0, %1) <{Sout = [#tf_type.shape<1x3>], dim_args_spec = [], disabled_checks = [], has_token_input_output = false, module = "", platforms = [], version = 5 : i64}> {_entry_function = @composite_dot_general_fn, _original_entry_function = "composite_dot_general_fn", _quantization_method = "weight_only_ptq { }", _stablehlo_module_attrs = {}, _tfl_quant_trait = "fully_quantizable", device = ""} : (tensor<1x2xf32>, tensor<2x3xf32>) -> tensor<1x3xf32>
     return %2 : tensor<1x3xf32>
   }
@@ -21,15 +21,15 @@ module attributes {tf_saved_model.semantics} {
 // CHECK-LABEL: quantize_dot_general_fn
 // CHECK-SAME: %[[ARG0:.+]]: tensor<1x2xf32>
 // CHECK: %[[CST:.+]] = stablehlo.constant dense<3.000000e-01> : tensor<2x3xf32>
-// CHECK: %[[Q:.+]] = "quantfork.qcast"(%[[CST]]) : (tensor<2x3xf32>) -> tensor<2x3x!quant.uniform<i8:f32, 6.000000e-03:-128>>
+// CHECK: %[[Q:.+]] = "quantfork.qcast"(%[[CST]]) : (tensor<2x3xf32>) -> tensor<2x3x!quant.uniform<i8:f32, 6.000000e-03>>
 // CHECK: %[[CALL:.+]] = call @quantized_dot_general_fn(%[[ARG0]], %[[Q]])
-// CHECK-SAME: {_quantization_method = "weight_only_ptq { }"} : (tensor<1x2xf32>, tensor<2x3x!quant.uniform<i8:f32, 6.000000e-03:-128>>) -> tensor<1x3xf32>
+// CHECK-SAME: {_quantization_method = "weight_only_ptq { }"} : (tensor<1x2xf32>, tensor<2x3x!quant.uniform<i8:f32, 6.000000e-03>>) -> tensor<1x3xf32>
 // CHECK: return %[[CALL]]
 
 // CHECK: quantized_dot_general_fn
-// CHECK-SAME: (%[[ARG1:.+]]: tensor<1x2xf32>,  %[[ARG2:.+]]: tensor<2x3x!quant.uniform<i8:f32, 6.000000e-03:-128>>) -> tensor<1x3xf32>
+// CHECK-SAME: (%[[ARG1:.+]]: tensor<1x2xf32>,  %[[ARG2:.+]]: tensor<2x3x!quant.uniform<i8:f32, 6.000000e-03>>) -> tensor<1x3xf32>
 // CHECK: %[[DOT:.+]] = stablehlo.dot_general %[[ARG1]], %[[ARG2]]
-// CHECK-SAME: (tensor<1x2xf32>, tensor<2x3x!quant.uniform<i8:f32, 6.000000e-03:-128>>) -> tensor<1x3xf32>
+// CHECK-SAME: (tensor<1x2xf32>, tensor<2x3x!quant.uniform<i8:f32, 6.000000e-03>>) -> tensor<1x3xf32>
 // CHECK: return %[[DOT]]
 
 // -----
diff --git a/third_party/stablehlo/temporary.patch b/third_party/stablehlo/temporary.patch
index f906a856ae69e1..abb329aec579e4 100755
--- a/third_party/stablehlo/temporary.patch
+++ b/third_party/stablehlo/temporary.patch
@@ -164,81 +164,6 @@ diff --ruN a/stablehlo/CMakeLists.txt b/stablehlo/CMakeLists.txt
  
  #-------------------------------------------------------------------------------
  # Directory setup
-diff --ruN a/stablehlo/docs/spec.md b/stablehlo/docs/spec.md
---- stablehlo/docs/spec.md
-+++ stablehlo/docs/spec.md
-@@ -2532,10 +2532,10 @@
-         rhs_batching_dimensions, lhs_contracting_dimensions,
-         rhs_contracting_dimensions, precision_config), lhs, rhs, type(result))`.
- 
--For hybrid quantized types, performs `hybrid_dequantize_then_op(
--    lambda lhs, rhs: dot_general(lhs, rhs, lhs_batching_dimensions,
--        rhs_batching_dimensions, lhs_contracting_dimensions,
--        rhs_contracting_dimensions, precision_config), lhs, rhs)`.
-+This only specifies semantics for per-tensor quantization. Per-axis quantization
-+is work in progress ([#1574](https://github.com/openxla/stablehlo/issues/1574)).
-+Also, in the future we may consider adding support for hybrid quantization
-+ ([#1575](https://github.com/openxla/stablehlo/issues/1575)).
- 
- `precision_config` controls the tradeoff between speed and accuracy for
- computations on accelerator backends. This can be one of the following (at the
-@@ -2552,21 +2552,21 @@
- 
- #### Inputs
- 
--| Label | Name                         | Type                                                         | Constraints                                    |
--|-------|------------------------------|--------------------------------------------------------------|------------------------------------------------|
--| (I1)  | `lhs`                        | tensor or per-tensor quantized tensor                        | (C5-C6), (C9-C10), (C12-C14), (C17-C18), (C20) |
--| (I2)  | `rhs`                        | tensor or quantized tensor                                   | (C7-C10), (C12-C20)                            |
--| (I3)  | `lhs_batching_dimensions`    | 1-dimensional tensor constant of type `si64`                 | (C1), (C3), (C5), (C9), (C12)                  |
--| (I4)  | `rhs_batching_dimensions`    | 1-dimensional tensor constant of type `si64`                 | (C1), (C4), (C7), (C9)                         |
--| (I5)  | `lhs_contracting_dimensions` | 1-dimensional tensor constant of type `si64`                 | (C2), (C3), (C6), (C10)                        |
--| (I6)  | `rhs_contracting_dimensions` | 1-dimensional tensor constant of type `si64`                 | (C2), (C4), (C8), (C10), (C16)                 |
--| (I7)  | `precision_config`           | variadic number of enums of `DEFAULT`, `HIGH`, and `HIGHEST` | (C11)                                          |
-+| Label | Name                         | Type                                                         | Constraints                    |
-+|-------|------------------------------|--------------------------------------------------------------|--------------------------------|
-+| (I1)  | `lhs`                        | tensor or per-tensor quantized tensor                        | (C5-C6), (C9-C10), (C12-C16)   |
-+| (I2)  | `rhs`                        | tensor or quantized tensor                                   | (C7-C10), (C12), (C18-C19)     |
-+| (I3)  | `lhs_batching_dimensions`    | 1-dimensional tensor constant of type `si64`                 | (C1), (C3), (C5), (C9), (C12)  |
-+| (I4)  | `rhs_batching_dimensions`    | 1-dimensional tensor constant of type `si64`                 | (C1), (C4), (C7), (C9)         |
-+| (I5)  | `lhs_contracting_dimensions` | 1-dimensional tensor constant of type `si64`                 | (C2), (C3), (C6), (C10)        |
-+| (I6)  | `rhs_contracting_dimensions` | 1-dimensional tensor constant of type `si64`                 | (C2), (C4), (C8), (C10), (C19) |
-+| (I7)  | `precision_config`           | variadic number of enums of `DEFAULT`, `HIGH`, and `HIGHEST` | (C11)                          |
- 
- #### Outputs
- 
- | Name     | Type                       | Constraints                |
- |----------|----------------------------|----------------------------|
--| `result` | tensor or quantized tensor | (C12), (C14), (C18-C20)    |
-+| `result` | tensor or quantized tensor | (C12), (C14), (C16), (C18) |
- 
- #### Constraints
- 
-@@ -2589,17 +2589,14 @@
- * If the operation uses non-quantized tensors:
-   * (C13) `element_type(lhs) = element_type(rhs)`.
- * If the operation uses quantized tensors:
--  * (C14) `is_quantized(lhs) = is_quantized(result) and is_quantized(rhs)`.
--  * (C15) `zero_points(rhs) = 0`.
--  * (C16) If `is_per_axis_quantized(rhs)`, then
-+  * (C14) `is_quantized(lhs) and is_quantized(rhs) and is_quantized(result)`.
-+  * (C15) `storage_type(lhs) = storage_type(rhs)`.
-+  * (C16) `expressed_type(lhs) = expressed_type(rhs) = expressed_type(result)`.
-+  * (C17) `zero_points(rhs) = 0`.
-+  * (C18) If `is_per_tensor_quantized(rhs)`, then
-+    `is_per_tensor_quantized(result)`.
-+  * (C19) If `is_per_axis_quantized(rhs)`, then
-     `quantization_dimension(rhs)` not in `rhs_contracting_dimensions`.
--  * If `is_quantized(lhs)`:
--    * (C17) `storage_type(lhs) = storage_type(rhs)`.
--    * (C18) `expressed_type(lhs) = expressed_type(rhs) = expressed_type(result)`.
--    * (C19) If `is_per_tensor_quantized(rhs)`, then
--      `is_per_tensor_quantized(result)`.
--  * If `!is_quantized(lhs)`:
--    * (C20) `element_type(lhs) = expressed_type(rhs) = element_type(result)`.
- 
- #### Examples
- 
 diff --ruN a/stablehlo/stablehlo/CMakeLists.txt b/stablehlo/stablehlo/CMakeLists.txt
 --- stablehlo/stablehlo/CMakeLists.txt
 +++ stablehlo/stablehlo/CMakeLists.txt
@@ -250,118 +175,10 @@ diff --ruN a/stablehlo/stablehlo/CMakeLists.txt b/stablehlo/stablehlo/CMakeLists
  add_subdirectory(integrations)
  add_subdirectory(reference)
  add_subdirectory(tests)
-diff --ruN a/stablehlo/stablehlo/conversions/linalg/tests/convolution.mlir b/stablehlo/stablehlo/conversions/linalg/tests/convolution.mlir
---- stablehlo/stablehlo/conversions/linalg/tests/convolution.mlir
-+++ stablehlo/stablehlo/conversions/linalg/tests/convolution.mlir
-@@ -356,7 +356,7 @@
- }
- // CHECK-DAG:       %[[CST:.+]] = arith.constant 0.000000e+00 : f32
- // CHECK:       %[[COLLAPSE:.+]] = tensor.collapse_shape %[[FILTER]] {{\[}}[0, 1, 2, 3]] : tensor<2x2x1x6xf32> into tensor<24xf32>
--// CHECK:       %[[EXPAND:.+]] = tensor.expand_shape %[[COLLAPSE]] {{\[}}[0, 1, 2, 3]] : tensor<24xf32> into tensor<2x2x2x3xf32>
-+// CHECK:       %[[EXPAND:.+]] = tensor.expand_shape %[[COLLAPSE]] {{\[}}[0, 1, 2, 3]] output_shape [2, 2, 2, 3] : tensor<24xf32> into tensor<2x2x2x3xf32>
- // CHECK:       %[[INIT:.+]] = tensor.empty() : tensor<2x3x4x2x3xf32>
- // CHECK:       %[[FILL:.+]] = linalg.fill ins(%[[CST]] : f32) outs(%[[INIT]] : tensor<2x3x4x2x3xf32>) -> tensor<2x3x4x2x3xf32>
- // CHECK:       %[[OUT:.+]] = linalg.depthwise_conv_2d_nhwc_hwcm
-diff --ruN a/stablehlo/stablehlo/conversions/linalg/tests/miscellaneous.mlir b/stablehlo/stablehlo/conversions/linalg/tests/miscellaneous.mlir
---- stablehlo/stablehlo/conversions/linalg/tests/miscellaneous.mlir
-+++ stablehlo/stablehlo/conversions/linalg/tests/miscellaneous.mlir
-@@ -865,7 +865,7 @@
-   %0 = "stablehlo.reshape"(%arg0) : (tensor<i32>) -> tensor<1xi32>
-   func.return %0 : tensor<1xi32>
- }
--// CHECK: tensor.expand_shape %{{.*}} [] : tensor<i32> into tensor<1xi32>
-+// CHECK: tensor.expand_shape %{{.*}} [] output_shape [1] : tensor<i32> into tensor<1xi32>
- 
- // -----
- 
-@@ -876,7 +876,7 @@
-   func.return %0 : tensor<1xui32>
- }
- // CHECK:         %[[ARG_SIGNLESS:.*]] = builtin.unrealized_conversion_cast %[[ARG_UNSIGNED]] : tensor<ui32> to tensor<i32>
--// CHECK:         %[[RET_SIGNLESS:.*]] = tensor.expand_shape %[[ARG_SIGNLESS]] [] : tensor<i32> into tensor<1xi32>
-+// CHECK:         %[[RET_SIGNLESS:.*]] = tensor.expand_shape %[[ARG_SIGNLESS]] [] output_shape [1] : tensor<i32> into tensor<1xi32>
- // CHECK:         %[[RET_UNSIGNED:.*]] = builtin.unrealized_conversion_cast %[[RET_SIGNLESS]] : tensor<1xi32> to tensor<1xui32>
- // CHECK:         return %[[RET_UNSIGNED]] : tensor<1xui32>
- 
-@@ -978,7 +978,7 @@
- }
- // CHECK: %[[FLATTEN:.*]] = tensor.collapse_shape %{{.*}} {{\[}}[0, 1]] : tensor<?x?xf32> into tensor<?xf32>
- // CHECK: %[[CAST:.*]] = tensor.cast %[[FLATTEN]] : tensor<?xf32> to tensor<40xf32>
--// CHECK: tensor.expand_shape %[[CAST]] {{\[}}[0, 1, 2]] : tensor<40xf32> into tensor<2x4x5xf32>
-+// CHECK: tensor.expand_shape %[[CAST]] {{\[}}[0, 1, 2]] output_shape [2, 4, 5] : tensor<40xf32> into tensor<2x4x5xf32>
- 
- // -----
- 
-@@ -988,7 +988,7 @@
-   func.return %0 : tensor<1x3xi32>
- }
- // CHECK: %[[CAST:.*]] = tensor.cast %{{.*}} : tensor<?xi32> to tensor<3xi32>
--// CHECK: tensor.expand_shape %[[CAST]] {{\[}}[0, 1]] : tensor<3xi32> into tensor<1x3xi32>
-+// CHECK: tensor.expand_shape %[[CAST]] {{\[}}[0, 1]] output_shape [1, 3] : tensor<3xi32> into tensor<1x3xi32>
- 
- // -----
- 
-diff --ruN a/stablehlo/stablehlo/conversions/linalg/tests/random.mlir b/stablehlo/stablehlo/conversions/linalg/tests/random.mlir
---- stablehlo/stablehlo/conversions/linalg/tests/random.mlir
-+++ stablehlo/stablehlo/conversions/linalg/tests/random.mlir
-@@ -480,8 +480,8 @@
- // CHECK-DAG:   %[[VAL_101:.*]] = arith.xori %[[VAL_100]], %[[VAL_87]] : i32
- 
- // CHECK: linalg.yield %[[YIELDED_1:.*]], %[[YIELDED_2:.*]] : i64, i64
--// CHECK-DAG: %[[VAL_206:.*]] = tensor.expand_shape %[[VAL_207:.*]]#0 {{\[\[}}0, 1]] : tensor<4xi64> into tensor<4x1xi64>
--// CHECK-DAG: %[[VAL_208:.*]] = tensor.expand_shape %[[VAL_207]]#1 {{\[\[}}0, 1]] : tensor<4xi64> into tensor<4x1xi64>
-+// CHECK-DAG: %[[VAL_206:.*]] = tensor.expand_shape %[[VAL_207:.*]]#0 {{\[\[}}0, 1]]
-+// CHECK-DAG: %[[VAL_208:.*]] = tensor.expand_shape %[[VAL_207]]#1 {{\[\[}}0, 1]]
- // CHECK-DAG: %[[VAL_209:.*]] = tensor.empty() : tensor<4x2xi64>
- // CHECK-DAG: %[[VAL_213:.*]] = tensor.insert %[[VAL_30]] into %[[VAL_0]]{{\[}}%[[VAL_19]]] : tensor<2xi64>
- 
-@@ -575,10 +575,10 @@
- // CHECK: %[[COLLAPSE:.+]] = tensor.collapse_shape %[[CONCAT]]
- 
- 
--// CHECK: %[[VAL_213:.*]] = tensor.expand_shape %[[COLLAPSE]] {{\[\[}}0, 1]] : tensor<80xi32> into tensor<80x1xi32>
-+// CHECK: %[[VAL_213:.*]] = tensor.expand_shape %[[COLLAPSE]] {{\[\[}}0, 1]]
- // CHECK: %[[VAL_214:.*]] = tensor.extract_slice %[[VAL_213]][0, 0] [77, 1] [1, 1] : tensor<80x1xi32> to tensor<77x1xi32>
- // CHECK: %[[VAL_215:.*]] = tensor.collapse_shape %[[VAL_214]] {{\[\[}}0, 1]] : tensor<77x1xi32> into tensor<77xi32>
--// CHECK: %[[VAL_216:.*]] = tensor.expand_shape %[[VAL_215]] {{\[\[}}0, 1]] : tensor<77xi32> into tensor<7x11xi32>
-+// CHECK: %[[VAL_216:.*]] = tensor.expand_shape %[[VAL_215]] {{\[\[}}0, 1]]
- // CHECK: %[[VAL_217:.*]] = tensor.insert %[[NEWSTATE]] into %[[ARG0]]{{\[}}%[[C1]]] : tensor<2xi64>
- // CHECK: return %[[VAL_217]], %[[VAL_216]] : tensor<2xi64>, tensor<7x11xi32>
- 
-@@ -616,10 +616,10 @@
- // CHECK-DAG: %[[COLLAPSE:.+]] = tensor.collapse_shape %[[CONCAT]] {{\[\[}}0, 1]] : tensor<8x2xi64> into tensor<16xi64>
- 
- 
--// CHECK-DAG: %[[EXPANDED:.*]] = tensor.expand_shape %[[COLLAPSE]] {{\[\[}}0, 1]] : tensor<16xi64> into tensor<16x1xi64>
-+// CHECK-DAG: %[[EXPANDED:.*]] = tensor.expand_shape %[[COLLAPSE]] {{\[\[}}0, 1]]
- // CHECK-DAG: %[[SLICE:.*]] = tensor.extract_slice %[[EXPANDED]][0, 0] [15, 1] [1, 1] : tensor<16x1xi64> to tensor<15x1xi64>
- // CHECK-DAG: %[[EXPAND_2:.*]] = tensor.collapse_shape %[[SLICE]] {{\[\[}}0, 1]] : tensor<15x1xi64> into tensor<15xi64>
--// CHECK-DAG: %[[RESHAPE:.*]] = tensor.expand_shape %[[EXPAND_2]] {{\[\[}}0, 1]] : tensor<15xi64> into tensor<3x5xi64>
-+// CHECK-DAG: %[[RESHAPE:.*]] = tensor.expand_shape %[[EXPAND_2]] {{\[\[}}0, 1]]
- // CHECK-DAG: %[[INSERTED:.+]] = tensor.insert %[[NEWSTATE]] into %[[ARG0]][%[[C1]]] : tensor<2xi64>
- // CHECK: return %[[INSERTED]], %[[RESHAPE]]
- 
 diff --ruN a/stablehlo/stablehlo/conversions/tosa/tests/binary.mlir b/stablehlo/stablehlo/conversions/tosa/tests/binary.mlir
 --- stablehlo/stablehlo/conversions/tosa/tests/binary.mlir
 +++ stablehlo/stablehlo/conversions/tosa/tests/binary.mlir
-@@ -45,14 +45,14 @@
- 
- // CHECK-LABEL: @divide
- func.func @divide(%arg0 : tensor<10xi32>, %arg1 : tensor<10xi32>) -> tensor<10xi32> {
--  // CHECK: tosa.div
-+  // CHECK: tosa.int_div
-   %0 = "stablehlo.divide"(%arg0, %arg1) : (tensor<10xi32>, tensor<10xi32>) -> tensor<10xi32>
-   return %0 : tensor<10xi32>
- }
- 
- // CHECK-LABEL: @divide_f32
- func.func @divide_f32(%arg0 : tensor<10xf32>, %arg1 : tensor<10xf32>) -> tensor<10xf32> {
--  // tosa.div only supports i32, so this should not legalize.
-+  // tosa.int_div only supports i32, so this should not legalize.
-   // CHECK: stablehlo.divide
-   %0 = "stablehlo.divide"(%arg0, %arg1) : (tensor<10xf32>, tensor<10xf32>) -> tensor<10xf32>
-   return %0 : tensor<10xf32>
-@@ -123,7 +123,7 @@
+@@ -155,7 +155,7 @@
  
  // CHECK-LABEL: @maximum_f64
  func.func @maximum_f64(%arg0 : tensor<10xf64>, %arg1 : tensor<10xf64>) -> tensor<10xf64> {
@@ -383,162 +200,6 @@ diff --ruN a/stablehlo/stablehlo/conversions/tosa/tests/nullary.mlir b/stablehlo
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<10xf64>
    return %0 : tensor<10xf64>
  }
-diff --ruN a/stablehlo/stablehlo/conversions/tosa/transforms/StablehloLegalizeToTosa.pdll b/stablehlo/stablehlo/conversions/tosa/transforms/StablehloLegalizeToTosa.pdll
---- stablehlo/stablehlo/conversions/tosa/transforms/StablehloLegalizeToTosa.pdll
-+++ stablehlo/stablehlo/conversions/tosa/transforms/StablehloLegalizeToTosa.pdll
-@@ -125,7 +125,7 @@
- Pattern =>
-   replace op<stablehlo.divide>(input0 : Value<_: Tosa_Int32Tensor>,
-                           input1 : Value<_: Tosa_Int32Tensor>)
--     with op<tosa.div>(input0, input1);
-+     with op<tosa.int_div>(input0, input1);
- Pattern =>
-   replace op<stablehlo.maximum>(input0 : Value<_: Tosa_Tensor>,
-                            input1 : Value<_: Tosa_Tensor>)
-diff --ruN a/stablehlo/stablehlo/dialect/Base.cpp b/stablehlo/stablehlo/dialect/Base.cpp
---- stablehlo/stablehlo/dialect/Base.cpp
-+++ stablehlo/stablehlo/dialect/Base.cpp
-@@ -651,14 +651,14 @@
- 
-   // quantized_type_c5
-   auto maxPosFiniteNum =
--      APFloat::getLargest(quantizedElementType.getExpressedType()
--                              .cast<FloatType>()
--                              .getFloatSemantics())
-+      APFloat::getLargest(
-+          cast<FloatType>(quantizedElementType.getExpressedType())
-+              .getFloatSemantics())
-           .convertToDouble();
-   auto minPosFiniteNum =
--      APFloat::getSmallest(quantizedElementType.getExpressedType()
--                               .cast<FloatType>()
--                               .getFloatSemantics())
-+      APFloat::getSmallest(
-+          cast<FloatType>(quantizedElementType.getExpressedType())
-+              .getFloatSemantics())
-           .convertToDouble();
-   if (llvm::any_of(scales, [&](double scale) {
-         return scale < minPosFiniteNum || scale > maxPosFiniteNum;
-diff --ruN a/stablehlo/stablehlo/dialect/TypeInference.cpp b/stablehlo/stablehlo/dialect/TypeInference.cpp
---- stablehlo/stablehlo/dialect/TypeInference.cpp
-+++ stablehlo/stablehlo/dialect/TypeInference.cpp
-@@ -171,7 +171,7 @@
- LogicalResult verifyConvolutionDotGeneralCommonQuantizationConstraints(
-     std::optional<Location> location, Type lhsElementType, Type rhsElementType,
-     Type resultElementType) {
--  // convolution_c28 and dot_general_c14
-+  // convolution_c28
-   if (!isa<quant::QuantizedType>(rhsElementType) ||
-       (isa<quant::QuantizedType>(lhsElementType) !=
-        isa<quant::QuantizedType>(resultElementType))) {
-@@ -184,19 +184,19 @@
-   auto rhsQuantType = cast<quant::QuantizedType>(rhsElementType);
-   if (auto lhsQuantType = dyn_cast<quant::QuantizedType>(lhsElementType)) {
-     auto resultQuantType = cast<quant::QuantizedType>(resultElementType);
--    // convolution_c31 and dot_general_c17
-+    // convolution_c31
-     if (lhsQuantType.getStorageType() != rhsQuantType.getStorageType()) {
-       return emitOptionalError(
-           location, "mismatched lhs and rhs quantization storage types");
-     }
--    // convolution_c32 and dot_general_c18
-+    // convolution_c32
-     if (lhsQuantType.getExpressedType() != rhsQuantType.getExpressedType() ||
-         lhsQuantType.getExpressedType() != resultQuantType.getExpressedType()) {
-       return emitOptionalError(
-           location,
-           "mismatched lhs, rhs and result quantization expressed types");
-     }
--    // convolution_c33 and dot_general_c19
-+    // convolution_c33
-     if (isa<quant::UniformQuantizedType>(rhsQuantType) &&
-         !isa<quant::UniformQuantizedType>(resultQuantType)) {
-       return emitOptionalError(
-@@ -204,7 +204,7 @@
-     }
-   } else {
-     Type rhsExpressedType = rhsQuantType.getExpressedType();
--    // convolution_c34 and dot_general_c20
-+    // convolution_c34
-     if (lhsElementType != rhsExpressedType ||
-         lhsElementType != resultElementType) {
-       return emitOptionalError(location,
-@@ -3559,7 +3559,7 @@
-     }
-   }
- 
--  // convolution_c28, convolution_c31 - convolution_c34
-+  // convolution_c31 - convolution_c34
-   return verifyConvolutionDotGeneralCommonQuantizationConstraints(
-       location, lhsElementType, rhsElementType, resultElementType);
- }
-@@ -3626,41 +3626,6 @@
-   return success();
- }
- 
--LogicalResult verifyDotGeneralOpQuantizationConstraints(
--    std::optional<Location> location, Type lhsType, Type rhsType,
--    Type resultType, ArrayRef<int64_t> rhsContractingDimensions) {
--  Type lhsElementType = getElementTypeOrSelf(lhsType);
--  Type rhsElementType = getElementTypeOrSelf(rhsType);
--  Type resultElementType = getElementTypeOrSelf(resultType);
--
--  // dot_general_c15
--  if (auto rhsPerTensorQuantType =
--          dyn_cast<quant::UniformQuantizedType>(rhsElementType)) {
--    if (rhsPerTensorQuantType.getZeroPoint() != 0) {
--      return emitOptionalError(location, "Zero point of rhs should be 0");
--    }
--  } else if (auto rhsPerAxisQuantType =
--                 dyn_cast<quant::UniformQuantizedPerAxisType>(rhsElementType)) {
--    if (llvm::any_of(rhsPerAxisQuantType.getZeroPoints(),
--                     [](int64_t zero_point) { return zero_point != 0; })) {
--      return emitOptionalError(location, "Zero points of rhs should be 0");
--    }
--
--    // dot_general_c16
--    if (llvm::is_contained(rhsContractingDimensions,
--                           rhsPerAxisQuantType.getQuantizedDimension())) {
--      return emitOptionalError(
--          location,
--          "Quantization dimension of rhs should not be in the "
--          "contracting dimension of rhs");
--    }
--  }
--
--  // dot_general_c14, dot_general_c17 - dot_general_c20
--  return verifyConvolutionDotGeneralCommonQuantizationConstraints(
--      location, lhsElementType, rhsElementType, resultElementType);
--}
--
- LogicalResult verifyDotGeneralOp(std::optional<Location> location, Value lhs,
-                                  Value rhs,
-                                  ArrayRef<int64_t> lhsBatchingDimensions,
-@@ -3683,13 +3648,6 @@
-     return emitOptionalError(
-         location, "inferred shape '", dimSizesToString(inferredShape.getDims()),
-         "' ", "is incompatible with return type of operation ", resultType, "");
--
--  Type lhsType = lhs.getType();
--  Type rhsType = rhs.getType();
--  if (anyQuantized<quant::QuantizedType>({lhsType, rhsType, resultType})) {
--    return verifyDotGeneralOpQuantizationConstraints(
--        location, lhsType, rhsType, resultType, rhsContractingDimensions);
--  }
-   return success();
- }
- 
-@@ -3861,8 +3819,8 @@
-   if (SmallVector<int64_t> shape; operandType.hasStaticShape() &&
-                                   matchInts(outputShape, shape).succeeded()) {
-     int64_t operandCount = operandType.getNumElements();
--    int64_t shapeCount = std::accumulate(shape.begin(), shape.end(), 1,
--                                         std::multiplies<int64_t>());
-+    int64_t shapeCount = std::accumulate(
-+        shape.begin(), shape.end(), int64_t{1}, std::multiplies<int64_t>());
-     if (operandCount != shapeCount) {
-       return emitOptionalError(location,
-                                "output_shape is incompatible with input type "
 diff --ruN a/stablehlo/stablehlo/experimental/BUILD.bazel b/stablehlo/stablehlo/experimental/BUILD.bazel
 --- stablehlo/stablehlo/experimental/BUILD.bazel
 +++ stablehlo/stablehlo/experimental/BUILD.bazel
@@ -2928,23052 +2589,4 @@ diff --ruN a/stablehlo/stablehlo/experimental/transforms/StablehloRefineShapes.c
 +}  // namespace experimental
 +}  // namespace stablehlo
 +}  // namespace mlir
-diff --ruN a/stablehlo/stablehlo/integrations/python/tests/stablehlo.py b/stablehlo/stablehlo/integrations/python/tests/stablehlo.py
---- stablehlo/stablehlo/integrations/python/tests/stablehlo.py
-+++ stablehlo/stablehlo/integrations/python/tests/stablehlo.py
-@@ -241,18 +241,18 @@
-   # Formatted as (tensor_type, np_value)
-   # Program runs arg + arg, which is used for expected value
-   tests = [
--    # No numpy types for f8 - skipping fp8 tests
--    ("f16", np.asarray(1, np.float16)),
--    ("f32", np.asarray(2, np.float32)),
--    ("f64", np.asarray(3, np.double)),
--    ("1xi8", np.asarray([4], np.int8)),
--    ("1xi16", np.asarray([5], np.int16)),
--    ("1xi32", np.asarray([-6], np.int32)),
--    # Numpy's uint treated as int by DenseElementsAttr, skipping np.uint tests
--    ("2x2xf16", np.asarray([1, 2, 3, 4], np.float16).reshape(2,2)),
--    ("2x1x2xf16", np.asarray([1, 2, 3, 4], np.float16).reshape(2,1,2)),
--    ("?x?xf16", np.asarray([1, 2, 3, 4], np.float16).reshape(2,2)),
--    ("?x2xf16", np.asarray([1, 2, 3, 4], np.float16).reshape(2,2)),
-+      # No numpy types for f8 - skipping fp8 tests
-+      ("f16", np.asarray(1, np.float16)),
-+      ("f32", np.asarray(2, np.float32)),
-+      ("f64", np.asarray(3, np.double)),
-+      ("1xi8", np.asarray([4], np.int8)),
-+      ("1xi16", np.asarray([5], np.int16)),
-+      ("1xi32", np.asarray([-6], np.int32)),
-+      # Numpy's uint treated as int by DenseElementsAttr, skipping np.uint tests
-+      ("2x2xf16", np.asarray([1, 2, 3, 4], np.float16).reshape(2, 2)),
-+      ("2x1x2xf16", np.asarray([1, 2, 3, 4], np.float16).reshape(2, 1, 2)),
-+      ("?x?xf16", np.asarray([1, 2, 3, 4], np.float16).reshape(2, 2)),
-+      ("?x2xf16", np.asarray([1, 2, 3, 4], np.float16).reshape(2, 2)),
-   ]
-   for test in tests:
-     tensor_type, arg = test
-diff --ruN a/stablehlo/stablehlo/tests/ops_stablehlo_quantized.mlir b/stablehlo/stablehlo/tests/ops_stablehlo_quantized.mlir
---- stablehlo/stablehlo/tests/ops_stablehlo_quantized.mlir
-+++ stablehlo/stablehlo/tests/ops_stablehlo_quantized.mlir
-@@ -1066,146 +1066,6 @@
- 
- // -----
- 
--func.func @dot_general_hybrid_quantized(%arg0: tensor<2x3x4xf32>, %arg1: tensor<2x3x5x!quant.uniform<i8:f32, 1.0:0>>) -> tensor<2x4x5xf32> {
--  %0 = "stablehlo.dot_general"(%arg0, %arg1) {
--    dot_dimension_numbers = #stablehlo.dot<
--      lhs_batching_dimensions = [0],
--      rhs_batching_dimensions = [0],
--      lhs_contracting_dimensions = [1],
--      rhs_contracting_dimensions = [1]
--    >
--  } : (tensor<2x3x4xf32>, tensor<2x3x5x!quant.uniform<i8:f32, 1.0:0>>) -> tensor<2x4x5xf32>
--  func.return %0 : tensor<2x4x5xf32>
--}
--
--// -----
--
--func.func @dot_general_c14(%arg0: tensor<2x3x4xf32>, %arg1: tensor<2x3x5x!quant.uniform<i8:f32, 1.0:0>>) -> tensor<2x4x5x!quant.uniform<i8:f32, 1.0:17>> {
--  // expected-error@+1 {{rhs should be quantized for quantized operations and is_quantized(lhs)=is_quantized(result) should hold}}
--  %0 = "stablehlo.dot_general"(%arg0, %arg1) {
--    dot_dimension_numbers = #stablehlo.dot<
--      lhs_batching_dimensions = [0],
--      rhs_batching_dimensions = [0],
--      lhs_contracting_dimensions = [1],
--      rhs_contracting_dimensions = [1]
--    >
--  } : (tensor<2x3x4xf32>, tensor<2x3x5x!quant.uniform<i8:f32, 1.0:0>>) -> tensor<2x4x5x!quant.uniform<i8:f32, 1.0:17>>
--  func.return %0 : tensor<2x4x5x!quant.uniform<i8:f32, 1.0:17>>
--}
--
--// -----
--
--func.func @dot_general_c15_per_tensor(%arg0: tensor<2x3x4xf32>, %arg1: tensor<2x3x5x!quant.uniform<i8:f32, 1.0:30>>) -> tensor<2x4x5xf32> {
--  // expected-error@+1 {{Zero point of rhs should be 0}}
--  %0 = "stablehlo.dot_general"(%arg0, %arg1) {
--    dot_dimension_numbers = #stablehlo.dot<
--      lhs_batching_dimensions = [0],
--      rhs_batching_dimensions = [0],
--      lhs_contracting_dimensions = [1],
--      rhs_contracting_dimensions = [1]
--    >
--  } : (tensor<2x3x4xf32>, tensor<2x3x5x!quant.uniform<i8:f32, 1.0:30>>) -> tensor<2x4x5xf32>
--  func.return %0 : tensor<2x4x5xf32>
--}
--
--// -----
--
--func.func @dot_general_c15_per_axis(
--  %arg0: tensor<2x3x4x!quant.uniform<i8:f32, 1.0:17>>,
--  %arg1: tensor<2x3x5x!quant.uniform<i8:f32:0, {0.1:10, 0.1:10}>>) -> tensor<2x4x5x!quant.uniform<i8:f32:0, {0.1:-30, 0.1:-30}>> {
--  // expected-error@+1 {{Zero points of rhs should be 0}}
--  %0 = "stablehlo.dot_general"(%arg0, %arg1) {
--    dot_dimension_numbers = #stablehlo.dot<
--      lhs_batching_dimensions = [0],
--      rhs_batching_dimensions = [0],
--      lhs_contracting_dimensions = [1],
--      rhs_contracting_dimensions = [1]
--    >
--  } : (tensor<2x3x4x!quant.uniform<i8:f32, 1.0:17>>,
--  tensor<2x3x5x!quant.uniform<i8:f32:0, {0.1:10, 0.1:10}>>) -> tensor<2x4x5x!quant.uniform<i8:f32:0, {0.1:-30, 0.1:-30}>>
--  func.return %0 : tensor<2x4x5x!quant.uniform<i8:f32:0, {0.1:-30, 0.1:-30}>>
--}
--
--// -----
--
--func.func @dot_general_c16(
--  %arg0: tensor<2x3x4x!quant.uniform<i8:f32, 1.0:17>>,
--  %arg1: tensor<2x3x5x!quant.uniform<i8:f32:0, {0.1:0, 0.1:0}>>) -> tensor<2x4x5x!quant.uniform<i8:f32:0, {0.1:-30, 0.1:-30, 0.1:-30}>> {
--  // expected-error@+1 {{Quantization dimension of rhs should not be in the contracting dimension of rhs}}
--  %0 = "stablehlo.dot_general"(%arg0, %arg1) {
--    dot_dimension_numbers = #stablehlo.dot<
--      lhs_batching_dimensions = [1],
--      rhs_batching_dimensions = [1],
--      lhs_contracting_dimensions = [0],
--      rhs_contracting_dimensions = [0]
--    >
--  } : (tensor<2x3x4x!quant.uniform<i8:f32, 1.0:17>>,
--  tensor<2x3x5x!quant.uniform<i8:f32:0, {0.1:0, 0.1:0}>>) -> tensor<3x4x5x!quant.uniform<i8:f32:0, {0.1:-30, 0.1:-30, 0.1:-30}>>
--  func.return %0 : tensor<3x4x5x!quant.uniform<i8:f32:0, {0.1:-30, 0.1:-30, 0.1:-30}>>
--}
--
--// -----
--
--func.func @dot_general_c17(%arg0: tensor<2x3x4x!quant.uniform<i16:f32, 1.0:17>>, %arg1: tensor<2x3x5x!quant.uniform<i8:f32, 1.0:0>>) -> tensor<2x4x5x!quant.uniform<i8:f32, 1.0:17>> {
--  // expected-error@+1 {{mismatched lhs and rhs quantization storage types}}
--  %0 = "stablehlo.dot_general"(%arg0, %arg1) {
--    dot_dimension_numbers = #stablehlo.dot<
--      lhs_batching_dimensions = [0],
--      rhs_batching_dimensions = [0],
--      lhs_contracting_dimensions = [1],
--      rhs_contracting_dimensions = [1]
--    >
--  } : (tensor<2x3x4x!quant.uniform<i16:f32, 1.0:17>>, tensor<2x3x5x!quant.uniform<i8:f32, 1.0:0>>) -> tensor<2x4x5x!quant.uniform<i8:f32, 1.0:17>>
--  func.return %0 : tensor<2x4x5x!quant.uniform<i8:f32, 1.0:17>>
--}
--
--// -----
--
--func.func @dot_general_c18(%arg0: tensor<2x3x4x!quant.uniform<i8:f32, 1.0:17>>, %arg1: tensor<2x3x5x!quant.uniform<i8:f16, 1.0:0>>) -> tensor<2x4x5x!quant.uniform<i8:f32, 1.0:17>> {
--  // expected-error@+1 {{mismatched lhs, rhs and result quantization expressed types}}
--  %0 = "stablehlo.dot_general"(%arg0, %arg1) {
--    dot_dimension_numbers = #stablehlo.dot<
--      lhs_batching_dimensions = [0],
--      rhs_batching_dimensions = [0],
--      lhs_contracting_dimensions = [1],
--      rhs_contracting_dimensions = [1]
--    >
--  } : (tensor<2x3x4x!quant.uniform<i8:f32, 1.0:17>>, tensor<2x3x5x!quant.uniform<i8:f16, 1.0:0>>) -> tensor<2x4x5x!quant.uniform<i8:f32, 1.0:17>>
--  func.return %0 : tensor<2x4x5x!quant.uniform<i8:f32, 1.0:17>>
--}
--
--// -----
--
--func.func @dot_general_c19(%arg0: tensor<2x3x4x!quant.uniform<i8:f32, 1.0:17>>, %arg1: tensor<2x3x5x!quant.uniform<i8:f32, 1.0:0>>) -> tensor<2x4x5x!quant.uniform<i8:f32:0, {0.1:-30, 0.1:-30}>> {
--  // expected-error@+1 {{mismatched rhs and result quantization granularity}}
--  %0 = "stablehlo.dot_general"(%arg0, %arg1) {
--    dot_dimension_numbers = #stablehlo.dot<
--      lhs_batching_dimensions = [0],
--      rhs_batching_dimensions = [0],
--      lhs_contracting_dimensions = [1],
--      rhs_contracting_dimensions = [1]
--    >
--  } : (tensor<2x3x4x!quant.uniform<i8:f32, 1.0:17>>, tensor<2x3x5x!quant.uniform<i8:f32, 1.0:0>>) -> tensor<2x4x5x!quant.uniform<i8:f32:0, {0.1:-30, 0.1:-30}>>
--  func.return %0 : tensor<2x4x5x!quant.uniform<i8:f32:0, {0.1:-30, 0.1:-30}>>
--}
--
--// -----
--
--func.func @dot_general_c20(%arg0: tensor<2x3x4xf32>, %arg1: tensor<2x3x5x!quant.uniform<i8:f16, 1.0:0>>) -> tensor<2x4x5xf32> {
--  // expected-error@+1 {{mismatched rhs quantization expressed type and lhs and result element type}}
--  %0 = "stablehlo.dot_general"(%arg0, %arg1) {
--    dot_dimension_numbers = #stablehlo.dot<
--      lhs_batching_dimensions = [0],
--      rhs_batching_dimensions = [0],
--      lhs_contracting_dimensions = [1],
--      rhs_contracting_dimensions = [1]
--    >
--  } : (tensor<2x3x4xf32>, tensor<2x3x5x!quant.uniform<i8:f16, 1.0:0>>) -> tensor<2x4x5xf32>
--  func.return %0 : tensor<2x4x5xf32>
--}
--
--// -----
--
- func.func @quantized_element_type_c8(%arg0: tensor<1x2x!quant.uniform<i8<-128:127>:f32, 1.0:300>>) {
-   // expected-error-re@+1 {{operand #0 must be ranked tensor of {{.*}} 4/8/16/32-bit uniform quantized signed integer or 4/8/16/32-bit uniform quantized unsigned integer or 4/8/16/32-bit uniform quantized per axis signed integer or 4/8/16/32-bit uniform quantized per axis unsigned integer values, but got 'tensor<1x2x!quant.uniform<i8:f32, 1.000000e+00:300>>'}}
-   %0 = stablehlo.add %arg0,  %arg0 : tensor<1x2x!quant.uniform<i8<-128:127>:f32, 1.0:300>>
-diff --ruN a/stablehlo/stablehlo/tests/vhlo/stablehlo_legalize_to_vhlo.0_10_0.mlir b/stablehlo/stablehlo/tests/vhlo/stablehlo_legalize_to_vhlo.0_10_0.mlir
---- stablehlo/stablehlo/tests/vhlo/stablehlo_legalize_to_vhlo.0_10_0.mlir
-+++ stablehlo/stablehlo/tests/vhlo/stablehlo_legalize_to_vhlo.0_10_0.mlir
-@@ -13,6 +13,7 @@
- // ============ ATTRIBUTES ============
- 
- // CHECK-LABEL: "attr_comparison_direction_eq"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_direction_eq(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     // CHECK: comparison_direction = #vhlo<comparison_direction_v1 EQ>
-@@ -22,6 +23,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_direction_ne"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_direction_ne(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     // CHECK: comparison_direction = #vhlo<comparison_direction_v1 NE>
-@@ -31,6 +33,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_direction_ge"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_direction_ge(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     // CHECK: comparison_direction = #vhlo<comparison_direction_v1 GE>
-@@ -40,6 +43,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_direction_gt"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_direction_gt(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     // CHECK: comparison_direction = #vhlo<comparison_direction_v1 GT>
-@@ -49,6 +53,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_direction_le"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_direction_le(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     // CHECK: comparison_direction = #vhlo<comparison_direction_v1 LE>
-@@ -58,6 +63,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_direction_lt"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_direction_lt(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     // CHECK: comparison_direction = #vhlo<comparison_direction_v1 LT>
-@@ -67,6 +73,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_type_notype"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_type_notype(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     comparison_direction = #stablehlo<comparison_direction EQ>
-@@ -76,6 +83,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_type_float"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_type_float(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     comparison_direction = #stablehlo<comparison_direction EQ>,
-@@ -86,6 +94,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_type_totalorder"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_type_totalorder(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     comparison_direction = #stablehlo<comparison_direction EQ>,
-@@ -96,6 +105,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_type_signed"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_type_signed(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     comparison_direction = #stablehlo<comparison_direction EQ>,
-@@ -106,6 +116,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_type_unsigned"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_type_unsigned(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     comparison_direction = #stablehlo<comparison_direction EQ>,
-@@ -118,6 +129,7 @@
- // ConvDimensionNumbers aka #stablehlo.conv is covered below.
- 
- // CHECK-LABEL: "attr_custom_call_api_version_unspecified"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_custom_call_api_version_unspecified(%arg0: tensor<f32>) -> tensor<f32> {
-   %0 = "stablehlo.custom_call"(%arg0) {
-     call_target_name = "foo",
-@@ -128,6 +140,7 @@
- }
- 
- // CHECK-LABEL: "attr_custom_call_api_version_original"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_custom_call_api_version_original(%arg0: tensor<f32>) -> tensor<f32> {
-   %0 = "stablehlo.custom_call"(%arg0) {
-     call_target_name = "foo",
-@@ -138,6 +151,7 @@
- }
- 
- // CHECK-LABEL: "attr_custom_call_api_version_status_returning"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_custom_call_api_version_status_returning(%arg0: tensor<f32>) -> tensor<f32> {
-   %0 = "stablehlo.custom_call"(%arg0) {
-     call_target_name = "foo",
-@@ -148,6 +162,7 @@
- }
- 
- // CHECK-LABEL: "attr_custom_call_api_version_status_returning_unified"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_custom_call_api_version_status_returning_unified(%arg0: tensor<f32>) -> tensor<f32> {
-   %0 = "stablehlo.custom_call"(%arg0) {
-     call_target_name = "foo",
-@@ -160,6 +175,7 @@
- // DotDimensionNumbers aka #stablehlo.dot is covered below.
- 
- // CHECK-LABEL: "attr_fft_type_fft"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_fft_type_fft(%arg0: tensor<16xcomplex<f32>>) -> tensor<16xcomplex<f32>> {
-   %0 = "stablehlo.fft"(%arg0) {
-     // CHECK: fft_type = #vhlo<fft_type_v1 FFT>
-@@ -170,6 +186,7 @@
- }
- 
- // CHECK-LABEL: "attr_fft_type_ifft"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_fft_type_ifft(%arg0: tensor<16xcomplex<f32>>) -> tensor<16xcomplex<f32>> {
-   %0 = "stablehlo.fft"(%arg0) {
-     // CHECK: fft_type = #vhlo<fft_type_v1 IFFT>
-@@ -180,6 +197,7 @@
- }
- 
- // CHECK-LABEL: "attr_fft_type_rfft"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_fft_type_rfft(%arg0: tensor<16xf32>) -> tensor<9xcomplex<f32>> {
-   %0 = "stablehlo.fft"(%arg0) {
-     // CHECK: fft_type = #vhlo<fft_type_v1 RFFT>
-@@ -190,6 +208,7 @@
- }
- 
- // CHECK-LABEL: "attr_fft_type_irfft"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_fft_type_irfft(%arg0: tensor<9xcomplex<f32>>) -> tensor<16xf32> {
-   %0 = "stablehlo.fft"(%arg0) {
-     // CHECK: fft_type = #vhlo<fft_type_v1 IRFFT>
-@@ -202,6 +221,7 @@
- // GatherDimensionNumbers aka #stablehlo.gather is covered below.
- 
- // CHECK-LABEL: "attr_precision_config_default"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_precision_config_default(%arg0: tensor<8x16xf32>, %arg1: tensor<16x8xf32>) -> tensor<8x8xf32> {
-   %0 = "stablehlo.dot"(%arg0, %arg1) {
-     // CHECK: precision_config = #vhlo.array_v1<[#vhlo<precision_v1 DEFAULT>, #vhlo<precision_v1 DEFAULT>]>
-@@ -210,6 +230,7 @@
- }
- 
- // CHECK-LABEL: "attr_precision_config_high"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_precision_config_high(%arg0: tensor<8x16xf32>, %arg1: tensor<16x8xf32>) -> tensor<8x8xf32> {
-   %0 = "stablehlo.dot"(%arg0, %arg1) {
-     // CHECK: precision_config = #vhlo.array_v1<[#vhlo<precision_v1 HIGH>, #vhlo<precision_v1 HIGH>]>
-@@ -219,6 +240,7 @@
- }
- 
- // CHECK-LABEL: "attr_precision_config_highest"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_precision_config_highest(%arg0: tensor<8x16xf32>, %arg1: tensor<16x8xf32>) -> tensor<8x8xf32> {
-   %0 = "stablehlo.dot"(%arg0, %arg1) {
-     // CHECK: precision_config = #vhlo.array_v1<[#vhlo<precision_v1 HIGHEST>, #vhlo<precision_v1 HIGHEST>]>
-@@ -228,6 +250,7 @@
- }
- 
- // CHECK-LABEL: "attr_rng_algorithm_default"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_rng_algorithm_default(%arg0: tensor<f32>) -> (tensor<f32>, tensor<f32>) {
-   %0:2 = "stablehlo.rng_bit_generator"(%arg0) {
-     // CHECK: rng_algorithm = #vhlo<rng_algorithm_v1 DEFAULT>
-@@ -237,6 +260,7 @@
- }
- 
- // CHECK-LABEL: "attr_rng_algorithm_three_fry"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_rng_algorithm_three_fry(%arg0: tensor<f32>) -> (tensor<f32>, tensor<f32>) {
-   %0:2 = "stablehlo.rng_bit_generator"(%arg0) {
-     // CHECK: rng_algorithm = #vhlo<rng_algorithm_v1 THREE_FRY>
-@@ -246,6 +270,7 @@
- }
- 
- // CHECK-LABEL: "attr_rng_algorithm_philox"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_rng_algorithm_philox(%arg0: tensor<f32>) -> (tensor<f32>, tensor<f32>) {
-   %0:2 = "stablehlo.rng_bit_generator"(%arg0) {
-     // CHECK: rng_algorithm = #vhlo<rng_algorithm_v1 PHILOX>
-@@ -255,6 +280,7 @@
- }
- 
- // CHECK-LABEL: "attr_rng_distribution_uniform"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @attr_rng_distribution_uniform(%arg0: tensor<f32>, %arg1: tensor<f32>, %arg2: tensor<0xindex>) -> tensor<f32> {
-   %0 = "stablehlo.rng"(%arg0, %arg1, %arg2) {
-     // CHECK: rng_distribution = #vhlo<rng_distribution_v1 UNIFORM>
-@@ -264,6 +290,7 @@
- }
- 
- // CHECK-LABEL: "attr_rng_distribution_normal"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @attr_rng_distribution_normal(%arg0: tensor<f32>, %arg1: tensor<f32>, %arg2: tensor<0xindex>) -> tensor<f32> {
-   %0 = "stablehlo.rng"(%arg0, %arg1, %arg2) {
-     // CHECK: rng_distribution = #vhlo<rng_distribution_v1 NORMAL>
-@@ -275,6 +302,7 @@
- // ScatterDimensionNumbers aka #stablehlo.scatter is covered below.
- 
- // CHECK-LABEL: "attr_transpose_no_transpose"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_transpose_no_transpose(%arg0: tensor<16x16xf32>, %arg1: tensor<16x16xf32>) ->  tensor<16x16xf32> {
-   %0 = "stablehlo.triangular_solve"(%arg0, %arg1) {
-     left_side = true,
-@@ -287,6 +315,7 @@
- }
- 
- // CHECK-LABEL: "attr_transpose_transpose"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_transpose_transpose(%arg0: tensor<16x16xf32>, %arg1: tensor<16x16xf32>) ->  tensor<16x16xf32> {
-   %0 = "stablehlo.triangular_solve"(%arg0, %arg1) {
-     left_side = true,
-@@ -299,6 +328,7 @@
- }
- 
- // CHECK-LABEL: "attr_transpose_adjoint"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_transpose_adjoint(%arg0: tensor<16x16xf32>, %arg1: tensor<16x16xf32>) ->  tensor<16x16xf32> {
-   %0 = "stablehlo.triangular_solve"(%arg0, %arg1) {
-     left_side = true,
-@@ -313,10 +343,9 @@
- // TypeExtensionsAttr aka #stablehlo.type_extensions is covered below.
- 
- // CHECK-LABEL: "attr_type_extensions_bounds"
--func.func @attr_type_extensions_bounds(
--    %arg0: tensor<?x?xf32, #stablehlo.type_extensions<bounds = [16, ?]>>)
--    -> tensor<?x?xf32, #stablehlo.type_extensions<bounds = [16, ?]>> {
--  // CHECK: "vhlo.return_v1"(%arg0) : (!vhlo.tensor_v1<?x?x!vhlo.f32_v1, #vhlo.type_extensions_v1<bounds = [16, ?]>>) -> ()
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @attr_type_extensions_bounds(%arg0: tensor<?x?xf32, #stablehlo.type_extensions<bounds = [16, ?]>>) -> tensor<?x?xf32, #stablehlo.type_extensions<bounds = [16, ?]>> {
-+  // CHECK: "vhlo.return_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<?x?x!vhlo.f32_v1, #vhlo.type_extensions_v1<bounds = [16, ?]>>) -> ()
-   func.return %arg0 : tensor<?x?xf32, #stablehlo.type_extensions<bounds = [16, ?]>>
- }
- 
-@@ -324,8 +353,9 @@
- // ============ DEFAULTS ============
- 
- // CHECK-LABEL: "default_all_gather"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_all_gather(%arg0: tensor<16x8xf32>) -> tensor<16x16xf32> {
--  //               CHECK: "vhlo.all_gather_v1"(%arg0) <{
-+  //               CHECK: "vhlo.all_gather_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   all_gather_dim = #vhlo.integer_v1<1 : i64>
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0], [1]]> : tensor<2x1xi64>>,
-@@ -339,8 +369,9 @@
- }
- 
- // CHECK-LABEL: "default_all_reduce"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_all_reduce(%arg0: tensor<f32>) -> tensor<f32> {
--  //               CHECK: "vhlo.all_reduce_v1"(%arg0)
-+  //               CHECK: "vhlo.all_reduce_v1"(%[[ARG0]])
-   //          CHECK-SAME: <{
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0], [1]]> : tensor<2x1xi64>>,
-@@ -362,8 +393,9 @@
- }
- 
- // CHECK-LABEL: "default_all_to_all"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_all_to_all(%arg0: tensor<4x16xf32>) -> tensor<16x4xf32> {
--  //               CHECK: "vhlo.all_to_all_v1"(%arg0) <{
-+  //               CHECK: "vhlo.all_to_all_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   //          CHECK-SAME:   concat_dimension = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0, 1, 2, 3]]> : tensor<1x4xi64>>,
-@@ -380,8 +412,9 @@
- }
- 
- // CHECK-LABEL: "default_cholesky"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_cholesky(%arg0: tensor<1x16x16xf32>) -> tensor<1x16x16xf32> {
--  //      CHECK: "vhlo.cholesky_v1"(%arg0) <{
-+  //      CHECK: "vhlo.cholesky_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   lower = #vhlo.bool_v1<false>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<1x16x16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<1x16x16x!vhlo.f32_v1>
-   %0 = "stablehlo.cholesky"(%arg0) : (tensor<1x16x16xf32>) -> tensor<1x16x16xf32>
-@@ -389,8 +422,9 @@
- }
- 
- // CHECK-LABEL: "default_collective_permute"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_collective_permute(%arg0: tensor<16x8xf32>) -> tensor<16x8xf32> {
--  //               CHECK: "vhlo.collective_permute_v1"(%arg0) <{
-+  //               CHECK: "vhlo.collective_permute_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME{LITERAL}:   source_target_pairs = #vhlo.tensor_v1<dense<[[0, 1], [1, 2], [2, 3]]> : tensor<3x2xi64>>
-   //          CHECK-SAME: }> : (!vhlo.tensor_v1<16x8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x8x!vhlo.f32_v1>
-@@ -401,8 +435,9 @@
- }
- 
- // CHECK-LABEL: "default_compare"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @default_compare(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
--  //      CHECK: "vhlo.compare_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.compare_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   compare_type = #vhlo<comparison_type_v1 NOTYPE>,
-   // CHECK-SAME:   comparison_direction = #vhlo<comparison_direction_v1 EQ>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-@@ -413,8 +448,9 @@
- }
- 
- // CHECK-LABEL: "default_convolution"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @default_convolution(%arg0: tensor<1x8x8x207xf32>, %arg1: tensor<3x3x207x16xf32>) -> tensor<1x6x6x16xf32> {
--  //      CHECK: "vhlo.convolution_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.convolution_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   batch_group_count = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   feature_group_count = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   input_batch_dimension = #vhlo.integer_v1<0 : i64>,
-@@ -442,8 +478,9 @@
- }
- 
- // CHECK-LABEL: "default_custom_call"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_custom_call(%arg0: tensor<f32>) -> tensor<f32> {
--  //      CHECK: "vhlo.custom_call_v1"(%arg0) <{
-+  //      CHECK: "vhlo.custom_call_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   api_version = #vhlo<api_version_v1 API_VERSION_ORIGINAL>,
-   // CHECK-SAME:   backend_config = #vhlo.string_v1<"">,
-   // CHECK-SAME:   call_target_name = #vhlo.string_v1<"foo">,
-@@ -460,8 +497,9 @@
- }
- 
- // CHECK-LABEL: "default_dot_general"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @default_dot_general(%arg0: tensor<8x8x16xf32>, %arg1: tensor<8x16x8xf32>) -> tensor<8x8x8xf32> {
--  //      CHECK: "vhlo.dot_general_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.dot_general_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   lhs_batching_dimensions = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>,
-   // CHECK-SAME:   lhs_contracting_dimensions = #vhlo.tensor_v1<dense<2> : tensor<1xi64>>,
-   // CHECK-SAME:   precision_config = #vhlo.array_v1<[#vhlo<precision_v1 DEFAULT>, #vhlo<precision_v1 DEFAULT>]>,
-@@ -480,8 +518,9 @@
- }
- 
- // CHECK-LABEL: "default_dot"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @default_dot(%arg0: tensor<8x16xf32>, %arg1: tensor<16x8xf32>) -> tensor<8x8xf32> {
--  //      CHECK: "vhlo.dot_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.dot_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   precision_config = #vhlo.array_v1<[#vhlo<precision_v1 DEFAULT>, #vhlo<precision_v1 DEFAULT>]>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<8x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<8x8x!vhlo.f32_v1>
-   %0 = "stablehlo.dot"(%arg0, %arg1) : (tensor<8x16xf32>, tensor<16x8xf32>) -> tensor<8x8xf32>
-@@ -489,8 +528,9 @@
- }
- 
- // CHECK-LABEL: "default_dynamic_broadcast_in_dim"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @default_dynamic_broadcast_in_dim(%arg0: tensor<?x?xf32>, %arg1: tensor<2xindex>) -> tensor<?x?xf32> {
--  //      CHECK: "vhlo.dynamic_broadcast_in_dim_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.dynamic_broadcast_in_dim_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   broadcast_dimensions = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-   // CHECK-SAME:   known_expanding_dimensions = #vhlo.tensor_v1<dense<> : tensor<0xi64>>,
-   // CHECK-SAME:   known_nonexpanding_dimensions = #vhlo.tensor_v1<dense<> : tensor<0xi64>>
-@@ -502,8 +542,9 @@
- }
- 
- // CHECK-LABEL: "default_dynamic_conv"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @default_dynamic_conv(%arg0: tensor<1x8x8x207xf32>, %arg1: tensor<3x3x207x16xf32>, %arg2: tensor<4xi32>) -> tensor<1x?x?x16xf32> {
--  //      CHECK: "vhlo.dynamic_conv_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.dynamic_conv_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   batch_group_count = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   feature_group_count = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   input_batch_dimension = #vhlo.integer_v1<0 : i64>,
-@@ -531,8 +572,9 @@
- }
- 
- // CHECK-LABEL: "default_dynamic_gather"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @default_dynamic_gather(%arg0 : tensor<2x4x9xf32>, %arg1 : tensor<1x5x2xi32>, %arg2 : tensor<3xi32>) -> tensor<1x5x8xf32> {
--  //      CHECK: "vhlo.dynamic_gather_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.dynamic_gather_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   collapsed_slice_dims = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-   // CHECK-SAME:   index_vector_dim = #vhlo.integer_v1<2 : i64>,
-   // CHECK-SAME:   indices_are_sorted = #vhlo.bool_v1<false>,
-@@ -558,15 +600,16 @@
-   // CHECK-SAME:   sym_name = #vhlo.string_v1<"default_func">,
-   // CHECK-SAME:   sym_visibility = #vhlo.string_v1<"">
-   // CHECK-SAME: }> ({
--  // CHECK-NEXT:   ^[[BB:bb.*]](%arg0: !vhlo.tensor_v1<!vhlo.f32_v1>):
--  // CHECK-NEXT:     "vhlo.return_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-+  // CHECK-NEXT:   ^[[BB:bb.*]](%[[ARG0:.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
-+  // CHECK-NEXT:     "vhlo.return_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-   // CHECK-NEXT: }) : () -> ()
-   func.return %arg0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "dynamic_gather"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @dynamic_gather(%arg0 : tensor<2x4x9xf32>, %arg1 : tensor<1x5x2xi32>) -> tensor<1x5x1xf32> {
--  //      CHECK: "vhlo.gather_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.gather_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   collapsed_slice_dims = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-   // CHECK-SAME:   index_vector_dim = #vhlo.integer_v1<2 : i64>,
-   // CHECK-SAME:   indices_are_sorted = #vhlo.bool_v1<false>,
-@@ -587,8 +630,9 @@
- }
- 
- // CHECK-LABEL: "default_infeed"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_infeed(%arg0: !stablehlo.token) -> (tensor<f32>, !stablehlo.token) {
--  //               CHECK: "vhlo.infeed_v1"(%arg0) <{
-+  //               CHECK: "vhlo.infeed_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   infeed_config = #vhlo.string_v1<"">,
-   // CHECK-SAME{LITERAL}:   layout = #vhlo.array_v1<[]>
-   //          CHECK-SAME: }> : (!vhlo.token_v1) -> (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.token_v1)
-@@ -597,8 +641,9 @@
- }
- 
- // CHECK-LABEL: "default_outfeed"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @default_outfeed(%arg0: tensor<f32>, %arg1: !stablehlo.token) -> !stablehlo.token {
--  //      CHECK: "vhlo.outfeed_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.outfeed_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   outfeed_config = #vhlo.string_v1<"">
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.token_v1) -> !vhlo.token_v1
-   %0 = "stablehlo.outfeed"(%arg0, %arg1) : (tensor<f32>, !stablehlo.token) -> !stablehlo.token
-@@ -606,8 +651,9 @@
- }
- 
- // CHECK-LABEL: "default_recv"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_recv(%arg0: !stablehlo.token) -> (tensor<f32>, !stablehlo.token) {
--  //      CHECK: "vhlo.recv_v1"(%arg0) <{
-+  //      CHECK: "vhlo.recv_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME:   channel_type = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   is_host_transfer = #vhlo.bool_v1<false>
-@@ -619,8 +665,9 @@
- }
- 
- // CHECK-LABEL: "default_send"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @default_send(%arg0: tensor<f32>, %arg1: !stablehlo.token) -> !stablehlo.token {
--  //      CHECK: "vhlo.send_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.send_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME:   channel_type = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   is_host_transfer = #vhlo.bool_v1<false>
-@@ -632,8 +679,9 @@
- }
- 
- // CHECK-LABEL: "default_reduce_scatter"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_reduce_scatter(%arg0: tensor<16xf32>) -> tensor<16xf32> {
--  //               CHECK: "vhlo.reduce_scatter_v1"(%arg0) <{
-+  //               CHECK: "vhlo.reduce_scatter_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0], [1]]> : tensor<2x1xi64>>,
-   //          CHECK-SAME:   scatter_dimension = #vhlo.integer_v1<0 : i64>
-@@ -655,8 +703,9 @@
- }
- 
- // CHECK-LABEL: "default_reduce_window"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @default_reduce_window(%arg0: tensor<2x17x31x7xf32>, %arg1: tensor<f32>) -> tensor<2x16x30x7xf32> {
--  //               CHECK: "vhlo.reduce_window_v1"(%arg0, %arg1)  <{
-+  //               CHECK: "vhlo.reduce_window_v1"(%[[ARG0]], %[[ARG1]])  <{
-   //          CHECK-SAME:   base_dilations = #vhlo.tensor_v1<dense<1> : tensor<4xi64>>,
-   // CHECK-SAME{LITERAL}:   padding = #vhlo.tensor_v1<dense<0> : tensor<4x2xi64>>,
-   //          CHECK-SAME:   window_dilations = #vhlo.tensor_v1<dense<1> : tensor<4xi64>>,
-@@ -678,8 +727,9 @@
- }
- 
- // CHECK-LABEL: "default_scatter"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @default_scatter(%arg0: tensor<200x100x300xf32>, %arg1: tensor<10x2xi32>, %arg2: tensor<10x300xf32>) -> tensor<200x100x300xf32> {
--  //      CHECK: "vhlo.scatter_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.scatter_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   index_vector_dim = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   indices_are_sorted = #vhlo.bool_v1<false>,
-   // CHECK-SAME:   inserted_window_dims = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-@@ -707,8 +757,9 @@
- }
- 
- // CHECK-LABEL: "default_select_and_scatter"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @default_select_and_scatter(%arg0: tensor<10x24x24x64xf32>, %arg1: tensor<10x23x23x64xf32>, %arg2: tensor<f32>) -> tensor<10x24x24x64xf32> {
--  //      CHECK: "vhlo.select_and_scatter_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.select_and_scatter_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   padding = #vhlo.tensor_v1<dense<0> : tensor<4x2xi64>>,
-   // CHECK-SAME:   window_dimensions = #vhlo.tensor_v1<dense<[1, 2, 2, 1]> : tensor<4xi64>>,
-   // CHECK-SAME:   window_strides = #vhlo.tensor_v1<dense<1> : tensor<4xi64>>
-@@ -736,8 +787,9 @@
- }
- 
- // CHECK-LABEL: "default_sort"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_sort(%arg0: tensor<16xf32>) -> tensor<16xf32> {
--  //      CHECK: "vhlo.sort_v1"(%arg0) <{
-+  //      CHECK: "vhlo.sort_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   dimension = #vhlo.integer_v1<-1 : i64>
-   // CHECK-SAME:   is_stable = #vhlo.bool_v1<false>
-   // CHECK-SAME: }> ({
-@@ -756,29 +808,33 @@
- // ============ OPS ============
- 
- // CHECK-LABEL: "op_abs"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_abs(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.abs_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.abs_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.abs"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_add"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_add(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_after_all"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_after_all(%arg0: !stablehlo.token) -> !stablehlo.token {
--  // CHECK: "vhlo.after_all_v1"(%arg0) : (!vhlo.token_v1) -> !vhlo.token_v1
-+  // CHECK: "vhlo.after_all_v1"(%[[ARG0]]) : (!vhlo.token_v1) -> !vhlo.token_v1
-   %0 = "stablehlo.after_all"(%arg0) : (!stablehlo.token) -> !stablehlo.token
-   func.return %0 : !stablehlo.token
- }
- 
- // CHECK-LABEL: "op_all_gather"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_all_gather(%arg0: tensor<16x8xf32>) -> tensor<16x16xf32> {
--  //               CHECK: "vhlo.all_gather_v1"(%arg0) <{
-+  //               CHECK: "vhlo.all_gather_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   all_gather_dim = #vhlo.integer_v1<1 : i64>
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0], [1]]> : tensor<2x1xi64>>,
-@@ -794,8 +850,9 @@
- }
- 
- // CHECK-LABEL: "op_all_reduce"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_all_reduce(%arg0: tensor<f32>) -> tensor<f32> {
--  //               CHECK: "vhlo.all_reduce_v1"(%arg0) <{
-+  //               CHECK: "vhlo.all_reduce_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0], [1]]> : tensor<2x1xi64>>,
-   //          CHECK-SAME:   use_global_device_ids = #vhlo.bool_v1<true>
-@@ -817,8 +874,9 @@
- }
- 
- // CHECK-LABEL: "op_all_to_all"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_all_to_all(%arg0: tensor<4x16xf32>) -> tensor<16x4xf32> {
--  //               CHECK: "vhlo.all_to_all_v1"(%arg0) <{
-+  //               CHECK: "vhlo.all_to_all_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<1 : i64>,
-   //          CHECK-SAME:   concat_dimension = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0, 1, 2, 3]]> : tensor<1x4xi64>>,
-@@ -836,22 +894,25 @@
- }
- 
- // CHECK-LABEL: "op_and"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_and(%arg0: tensor<i1>, %arg1: tensor<i1>) -> tensor<i1> {
--  // CHECK: "vhlo.and_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-+  // CHECK: "vhlo.and_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-   %0 = "stablehlo.and"(%arg0, %arg1) : (tensor<i1>, tensor<i1>) -> tensor<i1>
-   func.return %0 : tensor<i1>
- }
- 
- // CHECK-LABEL: "op_atan2"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_atan2(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.atan2_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.atan2_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.atan2"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_batch_norm_grad"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}}, %[[ARG3:.*]]: {{.*}}, %[[ARG4:.*]]: {{.*}})
- func.func @op_batch_norm_grad(%arg0: tensor<16x16x16x16xf32>, %arg1: tensor<16xf32>, %arg2: tensor<16xf32>, %arg3: tensor<16xf32>, %arg4: tensor<16x16x16x16xf32>) -> (tensor<16x16x16x16xf32>, tensor<16xf32>, tensor<16xf32>) {
--  //      CHECK: "vhlo.batch_norm_grad_v1"(%arg0, %arg1, %arg2, %arg3, %arg4) <{
-+  //      CHECK: "vhlo.batch_norm_grad_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]], %[[ARG3]], %[[ARG4]]) <{
-   // CHECK-SAME:   epsilon = #vhlo.float_v1<1.000000e-03 : !vhlo.f32_v1>,
-   // CHECK-SAME:   feature_index = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<16x16x16x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x16x16x16x!vhlo.f32_v1>) -> (!vhlo.tensor_v1<16x16x16x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>)
-@@ -863,8 +924,9 @@
- }
- 
- // CHECK-LABEL: "op_batch_norm_inference"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}}, %[[ARG3:.*]]: {{.*}}, %[[ARG4:.*]]: {{.*}})
- func.func @op_batch_norm_inference(%arg0: tensor<16x16x16x16xf32>, %arg1: tensor<16xf32>, %arg2: tensor<16xf32>, %arg3: tensor<16xf32>, %arg4: tensor<16xf32>) -> tensor<16x16x16x16xf32> {
--  //      CHECK: "vhlo.batch_norm_inference_v1"(%arg0, %arg1, %arg2, %arg3, %arg4) <{
-+  //      CHECK: "vhlo.batch_norm_inference_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]], %[[ARG3]], %[[ARG4]]) <{
-   // CHECK-SAME:   epsilon = #vhlo.float_v1<1.000000e-03 : !vhlo.f32_v1>,
-   // CHECK-SAME:   feature_index = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<16x16x16x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x16x16x16x!vhlo.f32_v1>
-@@ -876,8 +938,9 @@
- }
- 
- // CHECK-LABEL: "op_batch_norm_training"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_batch_norm_training(%arg0: tensor<16x16x16x16xf32>, %arg1: tensor<16xf32>, %arg2: tensor<16xf32>) -> (tensor<16x16x16x16xf32>, tensor<16xf32>, tensor<16xf32>) {
--  //      CHECK: "vhlo.batch_norm_training_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.batch_norm_training_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   epsilon = #vhlo.float_v1<1.000000e-03 : !vhlo.f32_v1>,
-   // CHECK-SAME:   feature_index = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<16x16x16x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>) -> (!vhlo.tensor_v1<16x16x16x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>)
-@@ -889,15 +952,17 @@
- }
- 
- // CHECK-LABEL: "op_bitcast_convert"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_bitcast_convert(%arg0: tensor<i32>) -> tensor<f32> {
--  // CHECK: "vhlo.bitcast_convert_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.bitcast_convert_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.bitcast_convert"(%arg0) : (tensor<i32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_broadcast_in_dim"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_broadcast_in_dim(%arg0: tensor<16xf32>) -> tensor<16x16xf32> {
--  //      CHECK: "vhlo.broadcast_in_dim_v1"(%arg0) <{
-+  //      CHECK: "vhlo.broadcast_in_dim_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   broadcast_dimensions = #vhlo.tensor_v1<dense<1> : tensor<1xi64>>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x16x!vhlo.f32_v1>
-   %0 = "stablehlo.broadcast_in_dim"(%arg0) {
-@@ -907,8 +972,9 @@
- }
- 
- // CHECK-LABEL: "op_broadcast"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_broadcast(%arg0: tensor<16xf32>) -> tensor<16x16xf32> {
--  //      CHECK: "vhlo.broadcast_v1"(%arg0) <{
-+  //      CHECK: "vhlo.broadcast_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   broadcast_sizes = #vhlo.tensor_v1<dense<16> : tensor<1xi64>>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x16x!vhlo.f32_v1>
-   %0 = "stablehlo.broadcast"(%arg0) {
-@@ -918,9 +984,10 @@
- }
- 
- // CHECK-LABEL: "op_case"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_case(%arg0: tensor<i32>, %arg1: tensor<f32>) -> tensor<f32> {
--  //      CHECK: "vhlo.case_v1"(%arg0) ({
--  // CHECK-NEXT:   "vhlo.return_v1"(%arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-+  //      CHECK: "vhlo.case_v1"(%[[ARG0]]) ({
-+  // CHECK-NEXT:   "vhlo.return_v1"(%[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-   // CHECK-NEXT: }) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.case"(%arg0) ({
-     "stablehlo.return"(%arg1) : (tensor<f32>) -> ()
-@@ -929,22 +996,25 @@
- }
- 
- // CHECK-LABEL: "op_cbrt"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_cbrt(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.cbrt_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.cbrt_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.cbrt"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_ceil"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_ceil(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.ceil_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.ceil_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.ceil"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_cholesky"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_cholesky(%arg0: tensor<1x16x16xf32>) -> tensor<1x16x16xf32> {
--  //      CHECK: "vhlo.cholesky_v1"(%arg0) <{
-+  //      CHECK: "vhlo.cholesky_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   lower = #vhlo.bool_v1<true>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<1x16x16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<1x16x16x!vhlo.f32_v1>
-   %0 = "stablehlo.cholesky"(%arg0) {
-@@ -954,22 +1024,25 @@
- }
- 
- // CHECK-LABEL: "op_clamp"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_clamp(%arg0: tensor<f32>, %arg1: tensor<f32>, %arg2: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.clamp_v1"(%arg0, %arg1, %arg2) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.clamp_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.clamp"(%arg0, %arg1, %arg2) : (tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_count_leading_zeros"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_count_leading_zeros(%arg0: tensor<i32>) -> tensor<i32> {
--  // CHECK: "vhlo.count_leading_zeros_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-+  // CHECK: "vhlo.count_leading_zeros_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-   %0 = "stablehlo.count_leading_zeros"(%arg0) : (tensor<i32>) -> tensor<i32>
-   func.return %0 : tensor<i32>
- }
- 
- // CHECK-LABEL: "op_collective_permute"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_collective_permute(%arg0: tensor<16x8xf32>) -> tensor<16x8xf32> {
--  //               CHECK: "vhlo.collective_permute_v1"(%arg0) <{
-+  //               CHECK: "vhlo.collective_permute_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME{LITERAL}:   source_target_pairs = #vhlo.tensor_v1<dense<[[0, 1], [1, 2], [2, 3]]> : tensor<3x2xi64>>
-   //          CHECK-SAME: }> : (!vhlo.tensor_v1<16x8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x8x!vhlo.f32_v1>
-@@ -981,8 +1054,9 @@
- }
- 
- // CHECK-LABEL: "op_compare"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_compare(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
--  //      CHECK: "vhlo.compare_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.compare_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   compare_type = #vhlo<comparison_type_v1 TOTALORDER>,
-   // CHECK-SAME:   comparison_direction = #vhlo<comparison_direction_v1 EQ>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-@@ -994,15 +1068,17 @@
- }
- 
- // CHECK-LABEL: "op_complex"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_complex(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<complex<f32>> {
--  // CHECK: "vhlo.complex_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>
-+  // CHECK: "vhlo.complex_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>
-   %0 = "stablehlo.complex"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<complex<f32>>
-   func.return %0 : tensor<complex<f32>>
- }
- 
- // CHECK-LABEL: "op_concatenate"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_concatenate(%arg0: tensor<8xf32>, %arg1: tensor<8xf32>) -> tensor<16xf32> {
--  //      CHECK: "vhlo.concatenate_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.concatenate_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   dimension = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<8x!vhlo.f32_v1>, !vhlo.tensor_v1<8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x!vhlo.f32_v1>
-   %0 = "stablehlo.concatenate"(%arg0, %arg1) {
-@@ -1012,6 +1088,7 @@
- }
- 
- // CHECK-LABEL: "op_constant"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_constant(%arg0: tensor<f32>) -> tensor<f32> {
-   //      CHECK: "vhlo.constant_v1"() <{
-   // CHECK-SAME:   value = #vhlo.tensor_v1<dense<0.000000e+00> : tensor<f32>>
-@@ -1023,15 +1100,17 @@
- }
- 
- // CHECK-LABEL: "op_convert"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_convert(%arg0: tensor<i32>) -> tensor<f32> {
--  // CHECK: "vhlo.convert_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.convert_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.convert"(%arg0) : (tensor<i32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_convolution"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_convolution(%arg0: tensor<1x8x8x207xf32>, %arg1: tensor<3x3x207x16xf32>) -> tensor<1x7x7x16xf32> {
--  //      CHECK: "vhlo.convolution_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.convolution_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   batch_group_count = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   feature_group_count = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   input_batch_dimension = #vhlo.integer_v1<0 : i64>,
-@@ -1065,8 +1144,9 @@
- }
- 
- // CHECK-LABEL: "op_cosine"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_cosine(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.cosine_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.cosine_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.cosine"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
-@@ -1079,8 +1159,9 @@
- }
- 
- // CHECK-LABEL: "op_cross_replica_sum"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_cross_replica_sum(%arg0: tensor<f32>) -> tensor<f32> {
--  //               CHECK: "vhlo.cross-replica-sum_v1"(%arg0) <{
-+  //               CHECK: "vhlo.cross-replica-sum_v1"(%[[ARG0]]) <{
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0], [1]]> : tensor<2x1xi64>>
-   //          CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.cross-replica-sum"(%arg0) {
-@@ -1090,8 +1171,9 @@
- }
- 
- // CHECK-LABEL: "op_custom_call"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_custom_call(%arg0: tensor<f32>) -> tensor<f32> {
--  //      CHECK: "vhlo.custom_call_v1"(%arg0) <{
-+  //      CHECK: "vhlo.custom_call_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   api_version = #vhlo<api_version_v1 API_VERSION_STATUS_RETURNING>,
-   // CHECK-SAME:   backend_config = #vhlo.string_v1<"\08\03\1A\02">,
-   // CHECK-SAME:   call_target_name = #vhlo.string_v1<"foo">,
-@@ -1122,15 +1204,17 @@
- }
- 
- // CHECK-LABEL: "op_divide"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_divide(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.divide_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.divide_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.divide"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_dot_general"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_dot_general(%arg0: tensor<8x8x16xf32>, %arg1: tensor<8x16x8xf32>) -> tensor<8x8x8xf32> {
--  //      CHECK: "vhlo.dot_general_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.dot_general_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   lhs_batching_dimensions = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>,
-   // CHECK-SAME:   lhs_contracting_dimensions = #vhlo.tensor_v1<dense<2> : tensor<1xi64>>,
-   // CHECK-SAME:   precision_config = #vhlo.array_v1<[#vhlo<precision_v1 HIGHEST>, #vhlo<precision_v1 HIGHEST>]>,
-@@ -1150,8 +1234,9 @@
- }
- 
- // CHECK-LABEL: "op_dot"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_dot(%arg0: tensor<8x16xf32>, %arg1: tensor<16x8xf32>) -> tensor<8x8xf32> {
--  //      CHECK: "vhlo.dot_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.dot_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   precision_config = #vhlo.array_v1<[#vhlo<precision_v1 HIGHEST>, #vhlo<precision_v1 HIGHEST>]>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<8x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<8x8x!vhlo.f32_v1>
-   %0 = "stablehlo.dot"(%arg0, %arg1) {
-@@ -1161,8 +1246,9 @@
- }
- 
- // CHECK-LABEL: "op_dynamic_broadcast_in_dim"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_dynamic_broadcast_in_dim(%arg0: tensor<?x?xf32>, %arg1: tensor<2xindex>) -> tensor<?x?xf32> {
--  //      CHECK: "vhlo.dynamic_broadcast_in_dim_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.dynamic_broadcast_in_dim_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   broadcast_dimensions = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-   // CHECK-SAME:   known_expanding_dimensions = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>,
-   // CHECK-SAME:   known_nonexpanding_dimensions = #vhlo.tensor_v1<dense<1> : tensor<1xi64>>
-@@ -1176,8 +1262,9 @@
- }
- 
- // CHECK-LABEL: "op_dynamic_conv"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_dynamic_conv(%arg0: tensor<1x8x8x207xf32>, %arg1: tensor<3x3x207x16xf32>, %arg2: tensor<4xi32>) -> tensor<1x?x?x16xf32> {
--  //      CHECK: "vhlo.dynamic_conv_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.dynamic_conv_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   batch_group_count = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   feature_group_count = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   input_batch_dimension = #vhlo.integer_v1<0 : i64>,
-@@ -1211,8 +1298,9 @@
- }
- 
- // CHECK-LABEL: "op_dynamic_gather"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_dynamic_gather(%arg0 : tensor<2x4x9xf32>, %arg1 : tensor<1x5x2xi32>, %arg2 : tensor<3xi32>) -> tensor<1x5x8xf32> {
--  //      CHECK: "vhlo.dynamic_gather_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.dynamic_gather_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   collapsed_slice_dims = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-   // CHECK-SAME:   index_vector_dim = #vhlo.integer_v1<2 : i64>,
-   // CHECK-SAME:   indices_are_sorted = #vhlo.bool_v1<true>,
-@@ -1232,8 +1320,9 @@
- }
- 
- // CHECK-LABEL: "op_dynamic_iota"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_dynamic_iota(%arg0: tensor<1xindex>) -> tensor<?xf32> {
--  //      CHECK: "vhlo.dynamic_iota_v1"(%arg0) <{
-+  //      CHECK: "vhlo.dynamic_iota_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   iota_dimension = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<1x!vhlo.index_v1>) -> !vhlo.tensor_v1<?x!vhlo.f32_v1>
-   %0 = "stablehlo.dynamic_iota"(%arg0) {
-@@ -1243,22 +1332,25 @@
- }
- 
- // CHECK-LABEL: "op_dynamic_pad"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}}, %[[ARG3:.*]]: {{.*}}, %[[ARG4:.*]]: {{.*}})
- func.func @op_dynamic_pad(%arg0: tensor<?xf32>, %arg1: tensor<f32>, %arg2: tensor<1xindex>, %arg3: tensor<1xindex>, %arg4: tensor<1xindex>) -> tensor<?xf32> {
--  // CHECK: "vhlo.dynamic_pad_v1"(%arg0, %arg1, %arg2, %arg3, %arg4) : (!vhlo.tensor_v1<?x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>) -> !vhlo.tensor_v1<?x!vhlo.f32_v1>
-+  // CHECK: "vhlo.dynamic_pad_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]], %[[ARG3]], %[[ARG4]]) : (!vhlo.tensor_v1<?x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>) -> !vhlo.tensor_v1<?x!vhlo.f32_v1>
-   %0 = "stablehlo.dynamic_pad"(%arg0, %arg1, %arg2, %arg3, %arg4) : (tensor<?xf32>, tensor<f32>, tensor<1xindex>, tensor<1xindex>, tensor<1xindex>) -> tensor<?xf32>
-   func.return %0 : tensor<?xf32>
- }
- 
- // CHECK-LABEL: "op_dynamic_reshape"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_dynamic_reshape(%arg0: tensor<16xf32>, %arg1: tensor<2xindex>) -> tensor<?x?xf32> {
--  // CHECK: "vhlo.dynamic_reshape_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<2x!vhlo.index_v1>) -> !vhlo.tensor_v1<?x?x!vhlo.f32_v1>
-+  // CHECK: "vhlo.dynamic_reshape_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<2x!vhlo.index_v1>) -> !vhlo.tensor_v1<?x?x!vhlo.f32_v1>
-   %0 = "stablehlo.dynamic_reshape"(%arg0, %arg1) : (tensor<16xf32>, tensor<2xindex>) -> tensor<?x?xf32>
-   func.return %0 : tensor<?x?xf32>
- }
- 
- // CHECK-LABEL: "op_dynamic_slice"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_dynamic_slice(%arg0: tensor<16xf32>, %arg1: tensor<i64>) -> tensor<4xf32> {
--  //      CHECK: "vhlo.dynamic_slice_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.dynamic_slice_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   slice_sizes = #vhlo.tensor_v1<dense<4> : tensor<1xi64>>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.i64_v1>) -> !vhlo.tensor_v1<4x!vhlo.f32_v1>
-   %0 = "stablehlo.dynamic_slice"(%arg0, %arg1) {
-@@ -1268,15 +1360,17 @@
- }
- 
- // CHECK-LABEL: "op_dynamic_update_slice"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_dynamic_update_slice(%arg0: tensor<16xf32>, %arg1: tensor<4xf32>, %arg2: tensor<i64>) -> tensor<16xf32> {
--  // CHECK: "vhlo.dynamic_update_slice_v1"(%arg0, %arg1, %arg2) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<4x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.i64_v1>) -> !vhlo.tensor_v1<16x!vhlo.f32_v1>
-+  // CHECK: "vhlo.dynamic_update_slice_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<4x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.i64_v1>) -> !vhlo.tensor_v1<16x!vhlo.f32_v1>
-   %0 = "stablehlo.dynamic_update_slice"(%arg0, %arg1, %arg2) : (tensor<16xf32>, tensor<4xf32>, tensor<i64>) -> tensor<16xf32>
-   func.return %0 : tensor<16xf32>
- }
- 
- // CHECK-LABEL: "op_einsum"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_einsum(%arg0: tensor<8x16xf32>, %arg1: tensor<16x8xf32>) -> tensor<8x8xf32> {
--  //      CHECK: "vhlo.einsum_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.einsum_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   einsum_config = #vhlo.string_v1<"ab,bc->ac">
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<8x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<8x8x!vhlo.f32_v1>
-   %0 = "stablehlo.einsum"(%arg0, %arg1) {
-@@ -1286,22 +1380,25 @@
- }
- 
- // CHECK-LABEL: "op_exponential_minus_one"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_exponential_minus_one(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.exponential_minus_one_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.exponential_minus_one_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.exponential_minus_one"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_exponential"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_exponential(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.exponential_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.exponential_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.exponential"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_fft"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_fft(%arg0: tensor<16xcomplex<f32>>) -> tensor<16xcomplex<f32>> {
--  //      CHECK: "vhlo.fft_v1"(%arg0) <{
-+  //      CHECK: "vhlo.fft_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   fft_length = #vhlo.tensor_v1<dense<16> : tensor<1xi64>>,
-   // CHECK-SAME:   fft_type = #vhlo<fft_type_v1 FFT>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<16x!vhlo.complex_v1<!vhlo.f32_v1>>) -> !vhlo.tensor_v1<16x!vhlo.complex_v1<!vhlo.f32_v1>>
-@@ -1313,8 +1410,9 @@
- }
- 
- // CHECK-LABEL: "op_floor"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_floor(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.floor_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.floor_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.floor"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
-@@ -1327,16 +1425,17 @@
-   // CHECK-SAME:   sym_name = #vhlo.string_v1<"op_func">,
-   // CHECK-SAME:   sym_visibility = #vhlo.string_v1<"private">
-   // CHECK-SAME: }> ({
--  // CHECK-NEXT:   ^[[BB:bb.*]](%arg0: !vhlo.tensor_v1<!vhlo.f32_v1>):
--  // CHECK-NEXT:     "vhlo.return_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-+  // CHECK-NEXT:   ^[[BB:bb.*]](%[[ARG0:.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
-+  // CHECK-NEXT:     "vhlo.return_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-   // CHECK-NEXT: }) : () -> ()
- 
-   func.return %arg0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_gather"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_gather(%arg0 : tensor<2x4x9xf32>, %arg1 : tensor<1x5x2xi32>) -> tensor<1x5x1xf32> {
--  //      CHECK: "vhlo.gather_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.gather_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   collapsed_slice_dims = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-   // CHECK-SAME:   index_vector_dim = #vhlo.integer_v1<2 : i64>,
-   // CHECK-SAME:   indices_are_sorted = #vhlo.bool_v1<true>,
-@@ -1358,8 +1457,9 @@
- }
- 
- // CHECK-LABEL: "op_get_dimension_size"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_get_dimension_size(%arg0: tensor<?xf32>) -> tensor<i32> {
--  //      CHECK: "vhlo.get_dimension_size_v1"(%arg0) <{
-+  //      CHECK: "vhlo.get_dimension_size_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   dimension = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<?x!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-   %0 = "stablehlo.get_dimension_size"(%arg0) {
-@@ -1369,8 +1469,9 @@
- }
- 
- // CHECK-LABEL: "op_get_tuple_element"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_get_tuple_element(%arg0: tuple<tensor<f32>, tensor<i32>>) -> tensor<f32> {
--  //      CHECK: "vhlo.get_tuple_element_v1"(%arg0) <{
-+  //      CHECK: "vhlo.get_tuple_element_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   index = #vhlo.integer_v1<0 : i32>
-   // CHECK-SAME: }> : (!vhlo.tuple_v1<!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.get_tuple_element"(%arg0) {
-@@ -1380,11 +1481,12 @@
- }
- 
- // CHECK-LABEL: "op_if"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_if(%arg0: tensor<i1>, %arg1: tensor<f32>, %arg2: tensor<f32>) -> tensor<f32> {
--  //      CHECK: "vhlo.if_v1"(%arg0) ({
--  // CHECK-NEXT:   "vhlo.return_v1"(%arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-+  //      CHECK: "vhlo.if_v1"(%[[ARG0]]) ({
-+  // CHECK-NEXT:   "vhlo.return_v1"(%[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-   // CHECK-NEXT: }, {
--  // CHECK-NEXT:   "vhlo.return_v1"(%arg2) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-+  // CHECK-NEXT:   "vhlo.return_v1"(%[[ARG2]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-   // CHECK-NEXT: }) : (!vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.if"(%arg0) ({
-     "stablehlo.return"(%arg1) : (tensor<f32>) -> ()
-@@ -1395,15 +1497,17 @@
- }
- 
- // CHECK-LABEL: "op_imag"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_imag(%arg0: tensor<complex<f32>>) -> tensor<f32> {
--  // CHECK: "vhlo.imag_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.imag_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.imag"(%arg0) : (tensor<complex<f32>>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_infeed"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_infeed(%arg0: !stablehlo.token) -> (tensor<f32>, !stablehlo.token) {
--  //               CHECK: "vhlo.infeed_v1"(%arg0) <{
-+  //               CHECK: "vhlo.infeed_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   infeed_config = #vhlo.string_v1<"foo">,
-   // CHECK-SAME{LITERAL}:   layout = #vhlo.array_v1<[#vhlo.array_v1<[]>]>
-   //          CHECK-SAME: }> : (!vhlo.token_v1) -> (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.token_v1)
-@@ -1426,36 +1530,41 @@
- }
- 
- // CHECK-LABEL: "op_is_finite"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_is_finite(%arg0: tensor<f32>) -> tensor<i1> {
--  // CHECK: "vhlo.is_finite_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-+  // CHECK: "vhlo.is_finite_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-   %0 = "stablehlo.is_finite"(%arg0) : (tensor<f32>) -> tensor<i1>
-   func.return %0 : tensor<i1>
- }
- 
- // CHECK-LABEL: "op_log"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_log(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.log_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.log_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.log"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_log_plus_one"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_log_plus_one(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.log_plus_one_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.log_plus_one_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.log_plus_one"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_logistic"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_logistic(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.logistic_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.logistic_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.logistic"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_map"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_map(%arg0: tensor<16xf32>) -> tensor<16xf32> {
--  //      CHECK: "vhlo.map_v1"(%arg0) <{
-+  //      CHECK: "vhlo.map_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   dimensions = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>
-   // CHECK-SAME: }> ({
-   // CHECK-NEXT:   ^[[BB:bb.*]](%[[ARG1:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
-@@ -1473,57 +1582,65 @@
- }
- 
- // CHECK-LABEL: "op_maximum"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_maximum(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.maximum_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.maximum_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.maximum"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_minimum"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_minimum(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.minimum_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.minimum_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.minimum"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_multiply"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_multiply(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.multiply_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.multiply_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.multiply"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_negate"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_negate(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.negate_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.negate_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.negate"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_not"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_not(%arg0: tensor<i1>) -> tensor<i1> {
--  // CHECK: "vhlo.not_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-+  // CHECK: "vhlo.not_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-   %0 = "stablehlo.not"(%arg0) : (tensor<i1>) -> tensor<i1>
-   func.return %0 : tensor<i1>
- }
- 
- // CHECK-LABEL: "op_optimization_barrier"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_optimization_barrier(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.optimization_barrier_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.optimization_barrier_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.optimization_barrier"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_or"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_or(%arg0: tensor<i1>, %arg1: tensor<i1>) -> tensor<i1> {
--  // CHECK: "vhlo.or_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-+  // CHECK: "vhlo.or_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-   %0 = "stablehlo.or"(%arg0, %arg1) : (tensor<i1>, tensor<i1>) -> tensor<i1>
-   func.return %0 : tensor<i1>
- }
- 
- // CHECK-LABEL: "op_outfeed"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_outfeed(%arg0: tensor<f32>, %arg1: !stablehlo.token) -> !stablehlo.token {
--  //      CHECK: "vhlo.outfeed_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.outfeed_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   outfeed_config = #vhlo.string_v1<"foo">
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.token_v1) -> !vhlo.token_v1
-   %0 = "stablehlo.outfeed"(%arg0, %arg1) {
-@@ -1533,8 +1650,9 @@
- }
- 
- // CHECK-LABEL: "op_pad"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_pad(%arg0: tensor<8xf32>, %arg1: tensor<f32>) -> tensor<16xf32> {
--  //      CHECK: "vhlo.pad_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.pad_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   edge_padding_high = #vhlo.tensor_v1<dense<4> : tensor<1xi64>>,
-   // CHECK-SAME:   edge_padding_low = #vhlo.tensor_v1<dense<4> : tensor<1xi64>>,
-   // CHECK-SAME:   interior_padding = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>
-@@ -1548,36 +1666,41 @@
- }
- 
- // CHECK-LABEL: "op_popcnt"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_popcnt(%arg0: tensor<i32>) -> tensor<i32> {
--  // CHECK: "vhlo.popcnt_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-+  // CHECK: "vhlo.popcnt_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-   %0 = "stablehlo.popcnt"(%arg0) : (tensor<i32>) -> tensor<i32>
-   func.return %0 : tensor<i32>
- }
- 
- // CHECK-LABEL: "op_power"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_power(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.power_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.power_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.power"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_real_dynamic_slice"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}}, %[[ARG3:.*]]: {{.*}})
- func.func @op_real_dynamic_slice(%arg0: tensor<?xf32>, %arg1: tensor<1xindex>, %arg2: tensor<1xindex>, %arg3: tensor<1xindex>) -> tensor<?xf32> {
--  // CHECK: "vhlo.real_dynamic_slice_v1"(%arg0, %arg1, %arg2, %arg3) : (!vhlo.tensor_v1<?x!vhlo.f32_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>) -> !vhlo.tensor_v1<?x!vhlo.f32_v1>
-+  // CHECK: "vhlo.real_dynamic_slice_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]], %[[ARG3]]) : (!vhlo.tensor_v1<?x!vhlo.f32_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>) -> !vhlo.tensor_v1<?x!vhlo.f32_v1>
-   %0 = "stablehlo.real_dynamic_slice"(%arg0, %arg1, %arg2, %arg3) : (tensor<?xf32>, tensor<1xindex>, tensor<1xindex>, tensor<1xindex>) -> tensor<?xf32>
-   func.return %0 : tensor<?xf32>
- }
- 
- // CHECK-LABEL: "op_real"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_real(%arg0: tensor<complex<f32>>) -> tensor<f32> {
--  // CHECK: "vhlo.real_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.real_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.real"(%arg0) : (tensor<complex<f32>>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_recv"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_recv(%arg0: !stablehlo.token) -> (tensor<f32>, !stablehlo.token) {
--  //      CHECK: "vhlo.recv_v1"(%arg0) <{
-+  //      CHECK: "vhlo.recv_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME:   channel_type = #vhlo.integer_v1<3 : i64>,
-   // CHECK-SAME:   is_host_transfer = #vhlo.bool_v1<true>
-@@ -1590,8 +1713,9 @@
- }
- 
- // CHECK-LABEL: "op_reduce"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_reduce(%arg0: tensor<16xf32>, %arg1: tensor<f32>) -> tensor<f32> {
--  //  CHECK: "vhlo.reduce_v1"(%arg0, %arg1)
-+  //  CHECK: "vhlo.reduce_v1"(%[[ARG0]], %[[ARG1]])
-   //  CHECK:   ^[[BB:bb.*]](%[[ARG1:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>, %[[ARG2:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
-   //  CHECK:     "vhlo.return_v1"(%[[VAL1:.*]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-   //  CHECK: }) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-@@ -1606,8 +1730,9 @@
- }
- 
- // CHECK-LABEL: "op_reduce_precision"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_reduce_precision(%arg0: tensor<f32>) -> tensor<f32> {
--  //      CHECK: "vhlo.reduce_precision_v1"(%arg0) <{
-+  //      CHECK: "vhlo.reduce_precision_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   exponent_bits = #vhlo.integer_v1<8 : i32>
-   // CHECK-SAME:   mantissa_bits = #vhlo.integer_v1<10 : i32>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-@@ -1619,8 +1744,9 @@
- }
- 
- // CHECK-LABEL: "op_reduce_scatter"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_reduce_scatter(%arg0: tensor<16xf32>) -> tensor<16xf32> {
--  //               CHECK: "vhlo.reduce_scatter_v1"(%arg0) <{
-+  //               CHECK: "vhlo.reduce_scatter_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0], [1]]> : tensor<2x1xi64>>,
-   //          CHECK-SAME:   scatter_dimension = #vhlo.integer_v1<0 : i64>
-@@ -1644,8 +1770,9 @@
- }
- 
- // CHECK-LABEL: "op_reduce_window"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_reduce_window(%arg0: tensor<2x17x31x7xf32>, %arg1: tensor<f32>) -> tensor<2x9x16x7xf32> {
--  //               CHECK: "vhlo.reduce_window_v1"(%arg0, %arg1) <{
-+  //               CHECK: "vhlo.reduce_window_v1"(%[[ARG0]], %[[ARG1]]) <{
-   //          CHECK-SAME:   base_dilations = #vhlo.tensor_v1<dense<[1, 2, 2, 1]> : tensor<4xi64>>,
-   // CHECK-SAME{LITERAL}:   padding = #vhlo.tensor_v1<dense<[[0, 0], [2, 0], [0, 2], [0, 0]]> : tensor<4x2xi64>>,
-   //          CHECK-SAME:   window_dilations = #vhlo.tensor_v1<dense<[1, 2, 2, 1]> : tensor<4xi64>>,
-@@ -1671,8 +1798,9 @@
- }
- 
- // CHECK-LABEL: "op_remainder"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_remainder(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.remainder_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.remainder_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.remainder"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
-@@ -1692,16 +1820,18 @@
- }
- 
- // CHECK-LABEL: "op_reshape"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_reshape(%arg0: tensor<16xf32>) -> tensor<4x4xf32> {
--  // CHECK: "vhlo.reshape_v1"(%arg0) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<4x4x!vhlo.f32_v1>
-+  // CHECK: "vhlo.reshape_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<4x4x!vhlo.f32_v1>
-   %0 = "stablehlo.reshape"(%arg0) : (tensor<16xf32>) -> tensor<4x4xf32>
-   func.return %0 : tensor<4x4xf32>
- }
- 
- // CHECK-LABEL: "op_return"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_return(%arg0: tensor<i32>, %arg1: tensor<f32>) -> tensor<f32> {
--  //      CHECK: "vhlo.case_v1"(%arg0) ({
--  // CHECK-NEXT:   "vhlo.return_v1"(%arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-+  //      CHECK: "vhlo.case_v1"(%[[ARG0]]) ({
-+  // CHECK-NEXT:   "vhlo.return_v1"(%[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-   // CHECK-NEXT: }) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.case"(%arg0) ({
-     "stablehlo.return"(%arg1) : (tensor<f32>) -> ()
-@@ -1710,8 +1840,9 @@
- }
- 
- // CHECK-LABEL: "op_reverse"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_reverse(%arg0: tensor<16xf32>) -> tensor<16xf32> {
--  //      CHECK: "vhlo.reverse_v1"(%arg0) <{
-+  //      CHECK: "vhlo.reverse_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   dimensions = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x!vhlo.f32_v1>
-   %0 = "stablehlo.reverse"(%arg0) {
-@@ -1721,8 +1852,9 @@
- }
- 
- // CHECK-LABEL: "op_rng_bit_generator"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_rng_bit_generator(%arg0: tensor<f32>) -> (tensor<f32>, tensor<f32>) {
--  //      CHECK: "vhlo.rng_bit_generator_v1"(%arg0) <{
-+  //      CHECK: "vhlo.rng_bit_generator_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   rng_algorithm = #vhlo<rng_algorithm_v1 PHILOX>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>)
-   %0:2 = "stablehlo.rng_bit_generator"(%arg0) {
-@@ -1732,8 +1864,9 @@
- }
- 
- // CHECK-LABEL: "op_rng"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_rng(%arg0: tensor<f32>, %arg1: tensor<f32>, %arg2: tensor<0xindex>) -> tensor<f32> {
--  //      CHECK: "vhlo.rng_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.rng_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   rng_distribution = #vhlo<rng_distribution_v1 NORMAL>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<0x!vhlo.index_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.rng"(%arg0, %arg1, %arg2) {
-@@ -1743,29 +1876,33 @@
- }
- 
- // CHECK-LABEL: "op_round_nearest_afz"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_round_nearest_afz(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.round_nearest_afz_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.round_nearest_afz_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.round_nearest_afz"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_round_nearest_even"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_round_nearest_even(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.round_nearest_even_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.round_nearest_even_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.round_nearest_even"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_rsqrt"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_rsqrt(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.rsqrt_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.rsqrt_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.rsqrt"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_scatter"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_scatter(%arg0: tensor<200x100x300xf32>, %arg1: tensor<10x2xi32>, %arg2: tensor<10x300xf32>) -> tensor<200x100x300xf32> {
--  //      CHECK: "vhlo.scatter_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.scatter_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   index_vector_dim = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   indices_are_sorted = #vhlo.bool_v1<true>,
-   // CHECK-SAME:   inserted_window_dims = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-@@ -1795,8 +1932,9 @@
- }
- 
- // CHECK-LABEL: "op_select_and_scatter"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_select_and_scatter(%arg0: tensor<10x24x24x64xf32>, %arg1: tensor<12x13x13x66xf32>, %arg2: tensor<f32>) -> tensor<10x24x24x64xf32> {
--  //      CHECK: "vhlo.select_and_scatter_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.select_and_scatter_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   padding = #vhlo.tensor_v1<dense<1> : tensor<4x2xi64>>,
-   // CHECK-SAME:   window_dimensions = #vhlo.tensor_v1<dense<[1, 2, 2, 1]> : tensor<4xi64>>,
-   // CHECK-SAME:   window_strides = #vhlo.tensor_v1<dense<[1, 2, 2, 1]> : tensor<4xi64>>
-@@ -1826,15 +1964,17 @@
- }
- 
- // CHECK-LABEL: "op_select"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_select(%arg0: tensor<i1>, %arg1: tensor<f32>, %arg2: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.select_v1"(%arg0, %arg1, %arg2) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.select_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.select"(%arg0, %arg1, %arg2) : (tensor<i1>, tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_send"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_send(%arg0: tensor<f32>, %arg1: !stablehlo.token) -> !stablehlo.token {
--  //      CHECK: "vhlo.send_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.send_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME:   channel_type = #vhlo.integer_v1<2 : i64>,
-   // CHECK-SAME:   is_host_transfer = #vhlo.bool_v1<true>
-@@ -1847,8 +1987,9 @@
- }
- 
- // CHECK-LABEL: "op_set_dimension_size"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_set_dimension_size(%arg0: tensor<?xf32>, %arg1: tensor<i32>) -> tensor<16xf32> {
--  //      CHECK: "vhlo.set_dimension_size_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.set_dimension_size_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   dimension = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<?x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<16x!vhlo.f32_v1>
-   %0 = "stablehlo.set_dimension_size"(%arg0, %arg1) {
-@@ -1858,43 +1999,49 @@
- }
- 
- // CHECK-LABEL: "op_shift_left"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_shift_left(%arg0: tensor<i32>, %arg1: tensor<i32>) -> tensor<i32> {
--  // CHECK: "vhlo.shift_left_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.i32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-+  // CHECK: "vhlo.shift_left_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-   %0 = "stablehlo.shift_left"(%arg0, %arg1) : (tensor<i32>, tensor<i32>) -> tensor<i32>
-   func.return %0 : tensor<i32>
- }
- 
- // CHECK-LABEL: "op_shift_right_arithmetic"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_shift_right_arithmetic(%arg0: tensor<i32>, %arg1: tensor<i32>) -> tensor<i32> {
--  // CHECK: "vhlo.shift_right_arithmetic_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.i32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-+  // CHECK: "vhlo.shift_right_arithmetic_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-   %0 = "stablehlo.shift_right_arithmetic"(%arg0, %arg1) : (tensor<i32>, tensor<i32>) -> tensor<i32>
-   func.return %0 : tensor<i32>
- }
- 
- // CHECK-LABEL: "op_shift_right_logical"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_shift_right_logical(%arg0: tensor<i32>, %arg1: tensor<i32>) -> tensor<i32> {
--  // CHECK: "vhlo.shift_right_logical_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.i32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-+  // CHECK: "vhlo.shift_right_logical_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-   %0 = "stablehlo.shift_right_logical"(%arg0, %arg1) : (tensor<i32>, tensor<i32>) -> tensor<i32>
-   func.return %0 : tensor<i32>
- }
- 
- // CHECK-LABEL: "op_sign"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_sign(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.sign_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.sign_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.sign"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_sine"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_sine(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.sine_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.sine_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.sine"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_slice"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_slice(%arg0: tensor<16xf32>) -> tensor<4xf32> {
--  //      CHECK: "vhlo.slice_v1"(%arg0) <{
-+  //      CHECK: "vhlo.slice_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   limit_indices = #vhlo.tensor_v1<dense<4> : tensor<1xi64>>,
-   // CHECK-SAME:   start_indices = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>,
-   // CHECK-SAME:   strides = #vhlo.tensor_v1<dense<1> : tensor<1xi64>>
-@@ -1908,8 +2055,9 @@
- }
- 
- // CHECK-LABEL: "op_sort"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_sort(%arg0: tensor<16xf32>) -> tensor<16xf32> {
--  //      CHECK: "vhlo.sort_v1"(%arg0) <{
-+  //      CHECK: "vhlo.sort_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   dimension = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME:   is_stable = #vhlo.bool_v1<true>
-   // CHECK-SAME: }> ({
-@@ -1929,29 +2077,33 @@
- }
- 
- // CHECK-LABEL: "op_sqrt"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_sqrt(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.sqrt_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.sqrt_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.sqrt"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_subtract"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_subtract(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.subtract_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.subtract_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.subtract"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_tanh"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_tanh(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.tanh_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.tanh_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.tanh"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_torch_index_select"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_torch_index_select(%arg0: tensor<5x1x5xf32>, %arg1: tensor<2xi32>) ->  tensor<2x1x5xf32> {
--  //      CHECK: "vhlo.torch_index_select_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.torch_index_select_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   batch_dims = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME:   dim = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<5x1x5x!vhlo.f32_v1>, !vhlo.tensor_v1<2x!vhlo.i32_v1>) -> !vhlo.tensor_v1<2x1x5x!vhlo.f32_v1>
-@@ -1963,8 +2115,9 @@
- }
- 
- // CHECK-LABEL: "op_trace"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_trace(%arg0: tensor<f32>) {
--  //      CHECK: "vhlo.trace_v1"(%arg0) <{
-+  //      CHECK: "vhlo.trace_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   tag = #vhlo.string_v1<"foo">
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-   "stablehlo.trace"(%arg0) {
-@@ -1974,8 +2127,9 @@
- }
- 
- // CHECK-LABEL: "op_transpose"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_transpose(%arg0: tensor<16x8xf32>) ->  tensor<8x16xf32> {
--  //      CHECK: "vhlo.transpose_v1"(%arg0) <{
-+  //      CHECK: "vhlo.transpose_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   permutation = #vhlo.tensor_v1<dense<[1, 0]> : tensor<2xi64>>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<16x8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<8x16x!vhlo.f32_v1>
-   %0 = "stablehlo.transpose"(%arg0) {
-@@ -1985,8 +2139,9 @@
- }
- 
- // CHECK-LABEL: "op_triangular_solve"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_triangular_solve(%arg0: tensor<16x16xf32>, %arg1: tensor<16x16xf32>) ->  tensor<16x16xf32> {
--  //      CHECK: "vhlo.triangular_solve_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.triangular_solve_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   left_side = #vhlo.bool_v1<true>,
-   // CHECK-SAME:   lower = #vhlo.bool_v1<true>,
-   // CHECK-SAME:   transpose_a = #vhlo<transpose_v1 NO_TRANSPOSE>,
-@@ -2002,15 +2157,17 @@
- }
- 
- // CHECK-LABEL: "op_tuple"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_tuple(%arg0: tensor<f32>) -> tuple<tensor<f32>> {
--  // CHECK: "vhlo.tuple_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tuple_v1<!vhlo.tensor_v1<!vhlo.f32_v1>>
-+  // CHECK: "vhlo.tuple_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tuple_v1<!vhlo.tensor_v1<!vhlo.f32_v1>>
-   %0 = "stablehlo.tuple"(%arg0) : (tensor<f32>) -> tuple<tensor<f32>>
-   func.return %0 : tuple<tensor<f32>>
- }
- 
- // CHECK-LABEL: "op_unary_einsum"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_unary_einsum(%arg0: tensor<8x16xf32>) -> tensor<8xf32> {
--  //      CHECK: "vhlo.unary_einsum_v1"(%arg0) <{
-+  //      CHECK: "vhlo.unary_einsum_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   einsum_config = #vhlo.string_v1<"ab->a">
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<8x16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<8x!vhlo.f32_v1>
-   %0 = "stablehlo.unary_einsum"(%arg0) {
-@@ -2020,22 +2177,25 @@
- }
- 
- // CHECK-LABEL: "op_uniform_dequantize"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_uniform_dequantize(%arg0: tensor<!quant.uniform<i8:f32, 34.0:16>>) -> tensor<f32> {
--  // CHECK: "vhlo.uniform_dequantize_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.uniform_dequantize_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.uniform_dequantize"(%arg0) : (tensor<!quant.uniform<i8:f32, 34.0:16>>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_uniform_quantize"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_uniform_quantize(%arg0: tensor<f32>) -> tensor<!quant.uniform<i8:f32, 34.0:16>> {
--  // CHECK: "vhlo.uniform_quantize_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>
-+  // CHECK: "vhlo.uniform_quantize_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>
-   %0 = "stablehlo.uniform_quantize"(%arg0) : (tensor<f32>) -> tensor<!quant.uniform<i8:f32, 34.0:16>>
-   func.return %0 : tensor<!quant.uniform<i8:f32, 34.0:16>>
- }
- 
- // CHECK-LABEL: "op_while"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_while(%arg0: tensor<i1>) -> tensor<i1> {
--  //      CHECK: "vhlo.while_v1"(%arg0) ({
-+  //      CHECK: "vhlo.while_v1"(%[[ARG0]]) ({
-   // CHECK-NEXT:   ^[[BB:bb.*]](%[[ARG1:arg.*]]: !vhlo.tensor_v1<!vhlo.bool_v1>):
-   // CHECK-NEXT:     "vhlo.return_v1"(%[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>) -> ()
-   // CHECK-NEXT:   }, {
-@@ -2053,8 +2213,9 @@
- }
- 
- // CHECK-LABEL: "op_xor"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_xor(%arg0: tensor<i1>, %arg1: tensor<i1>) -> tensor<i1> {
--  // CHECK: "vhlo.xor_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-+  // CHECK: "vhlo.xor_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-   %0 = "stablehlo.xor"(%arg0, %arg1) : (tensor<i1>, tensor<i1>) -> tensor<i1>
-   func.return %0 : tensor<i1>
- }
-@@ -2062,183 +2223,209 @@
- // ============ TYPES ============
- 
- // CHECK-LABEL: "type_i1"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_i1(%arg0: tensor<i1>, %arg1: tensor<i1>) -> tensor<i1> {
--  // CHECK: "vhlo.and_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-+  // CHECK: "vhlo.and_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-   %0 = "stablehlo.and"(%arg0, %arg1) : (tensor<i1>, tensor<i1>) -> tensor<i1>
-   func.return %0 : tensor<i1>
- }
- 
- // CHECK-LABEL: "type_i4"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_i4(%arg0: tensor<i4>, %arg1: tensor<i4>) -> tensor<i4> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.i4_v1>, !vhlo.tensor_v1<!vhlo.i4_v1>) -> !vhlo.tensor_v1<!vhlo.i4_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i4_v1>, !vhlo.tensor_v1<!vhlo.i4_v1>) -> !vhlo.tensor_v1<!vhlo.i4_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<i4>, tensor<i4>) -> tensor<i4>
-   func.return %0 : tensor<i4>
- }
- 
- // CHECK-LABEL: "type_i8"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_i8(%arg0: tensor<i8>, %arg1: tensor<i8>) -> tensor<i8> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.i8_v1>, !vhlo.tensor_v1<!vhlo.i8_v1>) -> !vhlo.tensor_v1<!vhlo.i8_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i8_v1>, !vhlo.tensor_v1<!vhlo.i8_v1>) -> !vhlo.tensor_v1<!vhlo.i8_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<i8>, tensor<i8>) -> tensor<i8>
-   func.return %0 : tensor<i8>
- }
- 
- // CHECK-LABEL: "type_i16"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_i16(%arg0: tensor<i16>, %arg1: tensor<i16>) -> tensor<i16> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.i16_v1>, !vhlo.tensor_v1<!vhlo.i16_v1>) -> !vhlo.tensor_v1<!vhlo.i16_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i16_v1>, !vhlo.tensor_v1<!vhlo.i16_v1>) -> !vhlo.tensor_v1<!vhlo.i16_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<i16>, tensor<i16>) -> tensor<i16>
-   func.return %0 : tensor<i16>
- }
- 
- // CHECK-LABEL: "type_i32"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_i32(%arg0: tensor<i32>, %arg1: tensor<i32>) -> tensor<i32> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.i32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<i32>, tensor<i32>) -> tensor<i32>
-   func.return %0 : tensor<i32>
- }
- 
- // CHECK-LABEL: "type_i64"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_i64(%arg0: tensor<i64>, %arg1: tensor<i64>) -> tensor<i64> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.i64_v1>, !vhlo.tensor_v1<!vhlo.i64_v1>) -> !vhlo.tensor_v1<!vhlo.i64_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i64_v1>, !vhlo.tensor_v1<!vhlo.i64_v1>) -> !vhlo.tensor_v1<!vhlo.i64_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<i64>, tensor<i64>) -> tensor<i64>
-   func.return %0 : tensor<i64>
- }
- 
- // CHECK-LABEL: "type_ui4"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_ui4(%arg0: tensor<ui4>, %arg1: tensor<ui4>) -> tensor<ui4> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.ui4_v1>, !vhlo.tensor_v1<!vhlo.ui4_v1>) -> !vhlo.tensor_v1<!vhlo.ui4_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.ui4_v1>, !vhlo.tensor_v1<!vhlo.ui4_v1>) -> !vhlo.tensor_v1<!vhlo.ui4_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<ui4>, tensor<ui4>) -> tensor<ui4>
-   func.return %0 : tensor<ui4>
- }
- 
- // CHECK-LABEL: "type_ui8"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_ui8(%arg0: tensor<ui8>, %arg1: tensor<ui8>) -> tensor<ui8> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.ui8_v1>, !vhlo.tensor_v1<!vhlo.ui8_v1>) -> !vhlo.tensor_v1<!vhlo.ui8_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.ui8_v1>, !vhlo.tensor_v1<!vhlo.ui8_v1>) -> !vhlo.tensor_v1<!vhlo.ui8_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<ui8>, tensor<ui8>) -> tensor<ui8>
-   func.return %0 : tensor<ui8>
- }
- 
- // CHECK-LABEL: "type_ui16"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_ui16(%arg0: tensor<ui16>, %arg1: tensor<ui16>) -> tensor<ui16> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.ui16_v1>, !vhlo.tensor_v1<!vhlo.ui16_v1>) -> !vhlo.tensor_v1<!vhlo.ui16_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.ui16_v1>, !vhlo.tensor_v1<!vhlo.ui16_v1>) -> !vhlo.tensor_v1<!vhlo.ui16_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<ui16>, tensor<ui16>) -> tensor<ui16>
-   func.return %0 : tensor<ui16>
- }
- 
- // CHECK-LABEL: "type_ui32"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_ui32(%arg0: tensor<ui32>, %arg1: tensor<ui32>) -> tensor<ui32> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.ui32_v1>, !vhlo.tensor_v1<!vhlo.ui32_v1>) -> !vhlo.tensor_v1<!vhlo.ui32_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.ui32_v1>, !vhlo.tensor_v1<!vhlo.ui32_v1>) -> !vhlo.tensor_v1<!vhlo.ui32_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<ui32>, tensor<ui32>) -> tensor<ui32>
-   func.return %0 : tensor<ui32>
- }
- 
- // CHECK-LABEL: "type_ui64"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_ui64(%arg0: tensor<ui64>, %arg1: tensor<ui64>) -> tensor<ui64> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.ui64_v1>, !vhlo.tensor_v1<!vhlo.ui64_v1>) -> !vhlo.tensor_v1<!vhlo.ui64_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.ui64_v1>, !vhlo.tensor_v1<!vhlo.ui64_v1>) -> !vhlo.tensor_v1<!vhlo.ui64_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<ui64>, tensor<ui64>) -> tensor<ui64>
-   func.return %0 : tensor<ui64>
- }
- 
- // CHECK-LABEL: "type_f8E4M3FN"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_f8E4M3FN(%arg0: tensor<f8E4M3FN>, %arg1: tensor<f8E4M3FN>) -> tensor<f8E4M3FN> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f8E4M3FN_v1>, !vhlo.tensor_v1<!vhlo.f8E4M3FN_v1>) -> !vhlo.tensor_v1<!vhlo.f8E4M3FN_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f8E4M3FN_v1>, !vhlo.tensor_v1<!vhlo.f8E4M3FN_v1>) -> !vhlo.tensor_v1<!vhlo.f8E4M3FN_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f8E4M3FN>, tensor<f8E4M3FN>) -> tensor<f8E4M3FN>
-   func.return %0 : tensor<f8E4M3FN>
- }
- 
- // CHECK-LABEL: "type_f8E5M2"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_f8E5M2(%arg0: tensor<f8E5M2>, %arg1: tensor<f8E5M2>) -> tensor<f8E5M2> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f8E5M2_v1>, !vhlo.tensor_v1<!vhlo.f8E5M2_v1>) -> !vhlo.tensor_v1<!vhlo.f8E5M2_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f8E5M2_v1>, !vhlo.tensor_v1<!vhlo.f8E5M2_v1>) -> !vhlo.tensor_v1<!vhlo.f8E5M2_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f8E5M2>, tensor<f8E5M2>) -> tensor<f8E5M2>
-   func.return %0 : tensor<f8E5M2>
- }
- 
- // CHECK-LABEL: "type_f8E4M3FNUZ"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_f8E4M3FNUZ(%arg0: tensor<f8E4M3FNUZ>, %arg1: tensor<f8E4M3FNUZ>) -> tensor<f8E4M3FNUZ> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f8E4M3FNUZ_v1>, !vhlo.tensor_v1<!vhlo.f8E4M3FNUZ_v1>) -> !vhlo.tensor_v1<!vhlo.f8E4M3FNUZ_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f8E4M3FNUZ_v1>, !vhlo.tensor_v1<!vhlo.f8E4M3FNUZ_v1>) -> !vhlo.tensor_v1<!vhlo.f8E4M3FNUZ_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f8E4M3FNUZ>, tensor<f8E4M3FNUZ>) -> tensor<f8E4M3FNUZ>
-   func.return %0 : tensor<f8E4M3FNUZ>
- }
- 
- // CHECK-LABEL: "type_f8E5M2FNUZ"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_f8E5M2FNUZ(%arg0: tensor<f8E5M2FNUZ>, %arg1: tensor<f8E5M2FNUZ>) -> tensor<f8E5M2FNUZ> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f8E5M2FNUZ_v1>, !vhlo.tensor_v1<!vhlo.f8E5M2FNUZ_v1>) -> !vhlo.tensor_v1<!vhlo.f8E5M2FNUZ_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f8E5M2FNUZ_v1>, !vhlo.tensor_v1<!vhlo.f8E5M2FNUZ_v1>) -> !vhlo.tensor_v1<!vhlo.f8E5M2FNUZ_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f8E5M2FNUZ>, tensor<f8E5M2FNUZ>) -> tensor<f8E5M2FNUZ>
-   func.return %0 : tensor<f8E5M2FNUZ>
- }
- 
- // CHECK-LABEL: "type_bf16"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_bf16(%arg0: tensor<bf16>, %arg1: tensor<bf16>) -> tensor<bf16> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.bf16_v1>, !vhlo.tensor_v1<!vhlo.bf16_v1>) -> !vhlo.tensor_v1<!vhlo.bf16_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.bf16_v1>, !vhlo.tensor_v1<!vhlo.bf16_v1>) -> !vhlo.tensor_v1<!vhlo.bf16_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<bf16>, tensor<bf16>) -> tensor<bf16>
-   func.return %0 : tensor<bf16>
- }
- 
- // CHECK-LABEL: "type_f16"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_f16(%arg0: tensor<f16>, %arg1: tensor<f16>) -> tensor<f16> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f16_v1>, !vhlo.tensor_v1<!vhlo.f16_v1>) -> !vhlo.tensor_v1<!vhlo.f16_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f16_v1>, !vhlo.tensor_v1<!vhlo.f16_v1>) -> !vhlo.tensor_v1<!vhlo.f16_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f16>, tensor<f16>) -> tensor<f16>
-   func.return %0 : tensor<f16>
- }
- 
- // CHECK-LABEL: "type_f32"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_f32(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "type_f64"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_f64(%arg0: tensor<f64>, %arg1: tensor<f64>) -> tensor<f64> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f64_v1>, !vhlo.tensor_v1<!vhlo.f64_v1>) -> !vhlo.tensor_v1<!vhlo.f64_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f64_v1>, !vhlo.tensor_v1<!vhlo.f64_v1>) -> !vhlo.tensor_v1<!vhlo.f64_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f64>, tensor<f64>) -> tensor<f64>
-   func.return %0 : tensor<f64>
- }
- 
- // CHECK-LABEL: "type_complex_f32"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_complex_f32(%arg0: tensor<complex<f32>>, %arg1: tensor<complex<f32>>) -> tensor<complex<f32>> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>, !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>) -> !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>, !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>) -> !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<complex<f32>>, tensor<complex<f32>>) -> tensor<complex<f32>>
-   func.return %0 : tensor<complex<f32>>
- }
- 
- // CHECK-LABEL: "type_complex_f64"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_complex_f64(%arg0: tensor<complex<f64>>, %arg1: tensor<complex<f64>>) -> tensor<complex<f64>> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f64_v1>>, !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f64_v1>>) -> !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f64_v1>>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f64_v1>>, !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f64_v1>>) -> !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f64_v1>>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<complex<f64>>, tensor<complex<f64>>) -> tensor<complex<f64>>
-   func.return %0 : tensor<complex<f64>>
- }
- 
- // CHECK-LABEL: "type_dynamism_ranked"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @type_dynamism_ranked(%arg0: tensor<?xf32>) -> tensor<?xf32> {
--  // CHECK: "vhlo.abs_v1"(%arg0) : (!vhlo.tensor_v1<?x!vhlo.f32_v1>) -> !vhlo.tensor_v1<?x!vhlo.f32_v1>
-+  // CHECK: "vhlo.abs_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<?x!vhlo.f32_v1>) -> !vhlo.tensor_v1<?x!vhlo.f32_v1>
-   %0 = "stablehlo.abs"(%arg0) : (tensor<?xf32>) -> tensor<?xf32>
-   func.return %0 : tensor<?xf32>
- }
- 
- // CHECK-LABEL: "type_quantization"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_quantization(%arg0: tensor<!quant.uniform<i8:f32, 34.0:16>>, %arg1: tensor<!quant.uniform<i8:f32, 34.0:16>>) -> tensor<!quant.uniform<i8:f32, 34.0:16>> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>, !vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>) -> !vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>, !vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>) -> !vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<!quant.uniform<i8:f32, 34.0:16>>, tensor<!quant.uniform<i8:f32, 34.0:16>>) -> tensor<!quant.uniform<i8:f32, 34.0:16>>
-   func.return %0 : tensor<!quant.uniform<i8:f32, 34.0:16>>
- }
- 
- //       CHECK: function_type = #vhlo.type_v1<!vhlo.func_v1<(!vhlo.token_v1) -> !vhlo.token_v1>>
- // CHECK-LABEL: "type_token_callee"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @type_token_callee(%arg0: !stablehlo.token) -> !stablehlo.token {
--  // CHECK: "vhlo.return_v1"(%arg0) : (!vhlo.token_v1) -> ()
-+  // CHECK: "vhlo.return_v1"(%[[ARG0]]) : (!vhlo.token_v1) -> ()
-   return %arg0 : !stablehlo.token
- }
- 
- //       CHECK: function_type = #vhlo.type_v1<!vhlo.func_v1<(!vhlo.token_v1) -> !vhlo.token_v1>>
- // CHECK-LABEL: "type_token_caller"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @type_token_caller(%arg0: !stablehlo.token) -> !stablehlo.token {
--  // CHECK:      "vhlo.call_v1"(%arg0) <{callee = #vhlo.string_v1<"type_token_callee">}
-+  // CHECK:      "vhlo.call_v1"(%[[ARG0]]) <{callee = #vhlo.string_v1<"type_token_callee">}
-   // CHECK-SAME: (!vhlo.token_v1) -> !vhlo.token_v1
-   %0 = func.call @type_token_callee(%arg0) : (!stablehlo.token) -> !stablehlo.token
-   return %0 : !stablehlo.token
- }
- 
- // CHECK-LABEL: "type_tuple"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @type_tuple(%arg0: tuple<tensor<f32>>) -> tuple<!stablehlo.token> {
-   %0 = "stablehlo.custom_call"(%arg0) {
-     call_target_name = "foo"
-diff --ruN a/stablehlo/stablehlo/tests/vhlo/stablehlo_legalize_to_vhlo.0_11_0.mlir b/stablehlo/stablehlo/tests/vhlo/stablehlo_legalize_to_vhlo.0_11_0.mlir
---- stablehlo/stablehlo/tests/vhlo/stablehlo_legalize_to_vhlo.0_11_0.mlir
-+++ stablehlo/stablehlo/tests/vhlo/stablehlo_legalize_to_vhlo.0_11_0.mlir
-@@ -13,6 +13,7 @@
- // ============ ATTRIBUTES ============
- 
- // CHECK-LABEL: "attr_comparison_direction_eq"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_direction_eq(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     // CHECK: comparison_direction = #vhlo<comparison_direction_v1 EQ>
-@@ -22,6 +23,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_direction_ne"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_direction_ne(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     // CHECK: comparison_direction = #vhlo<comparison_direction_v1 NE>
-@@ -31,6 +33,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_direction_ge"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_direction_ge(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     // CHECK: comparison_direction = #vhlo<comparison_direction_v1 GE>
-@@ -40,6 +43,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_direction_gt"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_direction_gt(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     // CHECK: comparison_direction = #vhlo<comparison_direction_v1 GT>
-@@ -49,6 +53,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_direction_le"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_direction_le(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     // CHECK: comparison_direction = #vhlo<comparison_direction_v1 LE>
-@@ -58,6 +63,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_direction_lt"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_direction_lt(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     // CHECK: comparison_direction = #vhlo<comparison_direction_v1 LT>
-@@ -67,6 +73,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_type_notype"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_type_notype(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     comparison_direction = #stablehlo<comparison_direction EQ>
-@@ -76,6 +83,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_type_float"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_type_float(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     comparison_direction = #stablehlo<comparison_direction EQ>,
-@@ -86,6 +94,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_type_totalorder"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_type_totalorder(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     comparison_direction = #stablehlo<comparison_direction EQ>,
-@@ -96,6 +105,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_type_signed"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_type_signed(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     comparison_direction = #stablehlo<comparison_direction EQ>,
-@@ -106,6 +116,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_type_unsigned"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_type_unsigned(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     comparison_direction = #stablehlo<comparison_direction EQ>,
-@@ -118,6 +129,7 @@
- // ConvDimensionNumbers aka #stablehlo.conv is covered below.
- 
- // CHECK-LABEL: "attr_custom_call_api_version_unspecified"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_custom_call_api_version_unspecified(%arg0: tensor<f32>) -> tensor<f32> {
-   %0 = "stablehlo.custom_call"(%arg0) {
-     call_target_name = "foo",
-@@ -128,6 +140,7 @@
- }
- 
- // CHECK-LABEL: "attr_custom_call_api_version_original"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_custom_call_api_version_original(%arg0: tensor<f32>) -> tensor<f32> {
-   %0 = "stablehlo.custom_call"(%arg0) {
-     call_target_name = "foo",
-@@ -138,6 +151,7 @@
- }
- 
- // CHECK-LABEL: "attr_custom_call_api_version_status_returning"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_custom_call_api_version_status_returning(%arg0: tensor<f32>) -> tensor<f32> {
-   %0 = "stablehlo.custom_call"(%arg0) {
-     call_target_name = "foo",
-@@ -148,6 +162,7 @@
- }
- 
- // CHECK-LABEL: "attr_custom_call_api_version_status_returning_unified"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_custom_call_api_version_status_returning_unified(%arg0: tensor<f32>) -> tensor<f32> {
-   %0 = "stablehlo.custom_call"(%arg0) {
-     call_target_name = "foo",
-@@ -160,6 +175,7 @@
- // DotDimensionNumbers aka #stablehlo.dot is covered below.
- 
- // CHECK-LABEL: "attr_fft_type_fft"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_fft_type_fft(%arg0: tensor<16xcomplex<f32>>) -> tensor<16xcomplex<f32>> {
-   %0 = "stablehlo.fft"(%arg0) {
-     // CHECK: fft_type = #vhlo<fft_type_v1 FFT>
-@@ -170,6 +186,7 @@
- }
- 
- // CHECK-LABEL: "attr_fft_type_ifft"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_fft_type_ifft(%arg0: tensor<16xcomplex<f32>>) -> tensor<16xcomplex<f32>> {
-   %0 = "stablehlo.fft"(%arg0) {
-     // CHECK: fft_type = #vhlo<fft_type_v1 IFFT>
-@@ -180,6 +197,7 @@
- }
- 
- // CHECK-LABEL: "attr_fft_type_rfft"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_fft_type_rfft(%arg0: tensor<16xf32>) -> tensor<9xcomplex<f32>> {
-   %0 = "stablehlo.fft"(%arg0) {
-     // CHECK: fft_type = #vhlo<fft_type_v1 RFFT>
-@@ -190,6 +208,7 @@
- }
- 
- // CHECK-LABEL: "attr_fft_type_irfft"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_fft_type_irfft(%arg0: tensor<9xcomplex<f32>>) -> tensor<16xf32> {
-   %0 = "stablehlo.fft"(%arg0) {
-     // CHECK: fft_type = #vhlo<fft_type_v1 IRFFT>
-@@ -202,6 +221,7 @@
- // GatherDimensionNumbers aka #stablehlo.gather is covered below.
- 
- // CHECK-LABEL: "attr_precision_config_default"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_precision_config_default(%arg0: tensor<8x16xf32>, %arg1: tensor<16x8xf32>) -> tensor<8x8xf32> {
-   %0 = "stablehlo.dot"(%arg0, %arg1) {
-     // CHECK: precision_config = #vhlo.array_v1<[#vhlo<precision_v1 DEFAULT>, #vhlo<precision_v1 DEFAULT>]>
-@@ -210,6 +230,7 @@
- }
- 
- // CHECK-LABEL: "attr_precision_config_high"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_precision_config_high(%arg0: tensor<8x16xf32>, %arg1: tensor<16x8xf32>) -> tensor<8x8xf32> {
-   %0 = "stablehlo.dot"(%arg0, %arg1) {
-     // CHECK: precision_config = #vhlo.array_v1<[#vhlo<precision_v1 HIGH>, #vhlo<precision_v1 HIGH>]>
-@@ -219,6 +240,7 @@
- }
- 
- // CHECK-LABEL: "attr_precision_config_highest"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_precision_config_highest(%arg0: tensor<8x16xf32>, %arg1: tensor<16x8xf32>) -> tensor<8x8xf32> {
-   %0 = "stablehlo.dot"(%arg0, %arg1) {
-     // CHECK: precision_config = #vhlo.array_v1<[#vhlo<precision_v1 HIGHEST>, #vhlo<precision_v1 HIGHEST>]>
-@@ -228,6 +250,7 @@
- }
- 
- // CHECK-LABEL: "attr_rng_algorithm_default"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_rng_algorithm_default(%arg0: tensor<f32>) -> (tensor<f32>, tensor<f32>) {
-   %0:2 = "stablehlo.rng_bit_generator"(%arg0) {
-     // CHECK: rng_algorithm = #vhlo<rng_algorithm_v1 DEFAULT>
-@@ -237,6 +260,7 @@
- }
- 
- // CHECK-LABEL: "attr_rng_algorithm_three_fry"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_rng_algorithm_three_fry(%arg0: tensor<f32>) -> (tensor<f32>, tensor<f32>) {
-   %0:2 = "stablehlo.rng_bit_generator"(%arg0) {
-     // CHECK: rng_algorithm = #vhlo<rng_algorithm_v1 THREE_FRY>
-@@ -246,6 +270,7 @@
- }
- 
- // CHECK-LABEL: "attr_rng_algorithm_philox"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_rng_algorithm_philox(%arg0: tensor<f32>) -> (tensor<f32>, tensor<f32>) {
-   %0:2 = "stablehlo.rng_bit_generator"(%arg0) {
-     // CHECK: rng_algorithm = #vhlo<rng_algorithm_v1 PHILOX>
-@@ -255,6 +280,7 @@
- }
- 
- // CHECK-LABEL: "attr_rng_distribution_uniform"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @attr_rng_distribution_uniform(%arg0: tensor<f32>, %arg1: tensor<f32>, %arg2: tensor<0xindex>) -> tensor<f32> {
-   %0 = "stablehlo.rng"(%arg0, %arg1, %arg2) {
-     // CHECK: rng_distribution = #vhlo<rng_distribution_v1 UNIFORM>
-@@ -264,6 +290,7 @@
- }
- 
- // CHECK-LABEL: "attr_rng_distribution_normal"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @attr_rng_distribution_normal(%arg0: tensor<f32>, %arg1: tensor<f32>, %arg2: tensor<0xindex>) -> tensor<f32> {
-   %0 = "stablehlo.rng"(%arg0, %arg1, %arg2) {
-     // CHECK: rng_distribution = #vhlo<rng_distribution_v1 NORMAL>
-@@ -275,6 +302,7 @@
- // ScatterDimensionNumbers aka #stablehlo.scatter is covered below.
- 
- // CHECK-LABEL: "attr_transpose_no_transpose"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_transpose_no_transpose(%arg0: tensor<16x16xf32>, %arg1: tensor<16x16xf32>) ->  tensor<16x16xf32> {
-   %0 = "stablehlo.triangular_solve"(%arg0, %arg1) {
-     left_side = true,
-@@ -287,6 +315,7 @@
- }
- 
- // CHECK-LABEL: "attr_transpose_transpose"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_transpose_transpose(%arg0: tensor<16x16xf32>, %arg1: tensor<16x16xf32>) ->  tensor<16x16xf32> {
-   %0 = "stablehlo.triangular_solve"(%arg0, %arg1) {
-     left_side = true,
-@@ -299,6 +328,7 @@
- }
- 
- // CHECK-LABEL: "attr_transpose_adjoint"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_transpose_adjoint(%arg0: tensor<16x16xf32>, %arg1: tensor<16x16xf32>) ->  tensor<16x16xf32> {
-   %0 = "stablehlo.triangular_solve"(%arg0, %arg1) {
-     left_side = true,
-@@ -313,10 +343,9 @@
- // TypeExtensionsAttr aka #stablehlo.type_extensions is covered below.
- 
- // CHECK-LABEL: "attr_type_extensions_bounds"
--func.func @attr_type_extensions_bounds(
--    %arg0: tensor<?x?xf32, #stablehlo.type_extensions<bounds = [16, ?]>>)
--    -> tensor<?x?xf32, #stablehlo.type_extensions<bounds = [16, ?]>> {
--  // CHECK: "vhlo.return_v1"(%arg0) : (!vhlo.tensor_v1<?x?x!vhlo.f32_v1, #vhlo.type_extensions_v1<bounds = [16, ?]>>) -> ()
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @attr_type_extensions_bounds(%arg0: tensor<?x?xf32, #stablehlo.type_extensions<bounds = [16, ?]>>) -> tensor<?x?xf32, #stablehlo.type_extensions<bounds = [16, ?]>> {
-+  // CHECK: "vhlo.return_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<?x?x!vhlo.f32_v1, #vhlo.type_extensions_v1<bounds = [16, ?]>>) -> ()
-   func.return %arg0 : tensor<?x?xf32, #stablehlo.type_extensions<bounds = [16, ?]>>
- }
- 
-@@ -324,8 +353,9 @@
- // ============ DEFAULTS ============
- 
- // CHECK-LABEL: "default_all_gather"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_all_gather(%arg0: tensor<16x8xf32>) -> tensor<16x16xf32> {
--  //               CHECK: "vhlo.all_gather_v1"(%arg0) <{
-+  //               CHECK: "vhlo.all_gather_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   all_gather_dim = #vhlo.integer_v1<1 : i64>
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0], [1]]> : tensor<2x1xi64>>,
-@@ -339,8 +369,9 @@
- }
- 
- // CHECK-LABEL: "default_all_reduce"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_all_reduce(%arg0: tensor<f32>) -> tensor<f32> {
--  //               CHECK: "vhlo.all_reduce_v1"(%arg0)
-+  //               CHECK: "vhlo.all_reduce_v1"(%[[ARG0]])
-   //          CHECK-SAME: <{
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0], [1]]> : tensor<2x1xi64>>,
-@@ -362,8 +393,9 @@
- }
- 
- // CHECK-LABEL: "default_all_to_all"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_all_to_all(%arg0: tensor<4x16xf32>) -> tensor<16x4xf32> {
--  //               CHECK: "vhlo.all_to_all_v1"(%arg0) <{
-+  //               CHECK: "vhlo.all_to_all_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   //          CHECK-SAME:   concat_dimension = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0, 1, 2, 3]]> : tensor<1x4xi64>>,
-@@ -380,8 +412,9 @@
- }
- 
- // CHECK-LABEL: "default_cholesky"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_cholesky(%arg0: tensor<1x16x16xf32>) -> tensor<1x16x16xf32> {
--  //      CHECK: "vhlo.cholesky_v1"(%arg0) <{
-+  //      CHECK: "vhlo.cholesky_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   lower = #vhlo.bool_v1<false>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<1x16x16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<1x16x16x!vhlo.f32_v1>
-   %0 = "stablehlo.cholesky"(%arg0) : (tensor<1x16x16xf32>) -> tensor<1x16x16xf32>
-@@ -389,8 +422,9 @@
- }
- 
- // CHECK-LABEL: "default_collective_permute"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_collective_permute(%arg0: tensor<16x8xf32>) -> tensor<16x8xf32> {
--  //               CHECK: "vhlo.collective_permute_v1"(%arg0) <{
-+  //               CHECK: "vhlo.collective_permute_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME{LITERAL}:   source_target_pairs = #vhlo.tensor_v1<dense<[[0, 1], [1, 2], [2, 3]]> : tensor<3x2xi64>>
-   //          CHECK-SAME: }> : (!vhlo.tensor_v1<16x8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x8x!vhlo.f32_v1>
-@@ -401,8 +435,9 @@
- }
- 
- // CHECK-LABEL: "default_compare"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @default_compare(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
--  //      CHECK: "vhlo.compare_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.compare_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   compare_type = #vhlo<comparison_type_v1 NOTYPE>,
-   // CHECK-SAME:   comparison_direction = #vhlo<comparison_direction_v1 EQ>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-@@ -413,8 +448,9 @@
- }
- 
- // CHECK-LABEL: "default_convolution"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @default_convolution(%arg0: tensor<1x8x8x207xf32>, %arg1: tensor<3x3x207x16xf32>) -> tensor<1x6x6x16xf32> {
--  //      CHECK: "vhlo.convolution_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.convolution_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   batch_group_count = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   feature_group_count = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   input_batch_dimension = #vhlo.integer_v1<0 : i64>,
-@@ -442,8 +478,9 @@
- }
- 
- // CHECK-LABEL: "default_custom_call"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_custom_call(%arg0: tensor<f32>) -> tensor<f32> {
--  //      CHECK: "vhlo.custom_call_v1"(%arg0) <{
-+  //      CHECK: "vhlo.custom_call_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   api_version = #vhlo<api_version_v1 API_VERSION_ORIGINAL>,
-   // CHECK-SAME:   backend_config = #vhlo.string_v1<"">,
-   // CHECK-SAME:   call_target_name = #vhlo.string_v1<"foo">,
-@@ -460,8 +497,9 @@
- }
- 
- // CHECK-LABEL: "default_dot_general"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @default_dot_general(%arg0: tensor<8x8x16xf32>, %arg1: tensor<8x16x8xf32>) -> tensor<8x8x8xf32> {
--  //      CHECK: "vhlo.dot_general_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.dot_general_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   lhs_batching_dimensions = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>,
-   // CHECK-SAME:   lhs_contracting_dimensions = #vhlo.tensor_v1<dense<2> : tensor<1xi64>>,
-   // CHECK-SAME:   precision_config = #vhlo.array_v1<[#vhlo<precision_v1 DEFAULT>, #vhlo<precision_v1 DEFAULT>]>,
-@@ -480,8 +518,9 @@
- }
- 
- // CHECK-LABEL: "default_dot"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @default_dot(%arg0: tensor<8x16xf32>, %arg1: tensor<16x8xf32>) -> tensor<8x8xf32> {
--  //      CHECK: "vhlo.dot_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.dot_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   precision_config = #vhlo.array_v1<[#vhlo<precision_v1 DEFAULT>, #vhlo<precision_v1 DEFAULT>]>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<8x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<8x8x!vhlo.f32_v1>
-   %0 = "stablehlo.dot"(%arg0, %arg1) : (tensor<8x16xf32>, tensor<16x8xf32>) -> tensor<8x8xf32>
-@@ -489,8 +528,9 @@
- }
- 
- // CHECK-LABEL: "default_dynamic_broadcast_in_dim"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @default_dynamic_broadcast_in_dim(%arg0: tensor<?x?xf32>, %arg1: tensor<2xindex>) -> tensor<?x?xf32> {
--  //      CHECK: "vhlo.dynamic_broadcast_in_dim_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.dynamic_broadcast_in_dim_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   broadcast_dimensions = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-   // CHECK-SAME:   known_expanding_dimensions = #vhlo.tensor_v1<dense<> : tensor<0xi64>>,
-   // CHECK-SAME:   known_nonexpanding_dimensions = #vhlo.tensor_v1<dense<> : tensor<0xi64>>
-@@ -502,8 +542,9 @@
- }
- 
- // CHECK-LABEL: "default_dynamic_conv"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @default_dynamic_conv(%arg0: tensor<1x8x8x207xf32>, %arg1: tensor<3x3x207x16xf32>, %arg2: tensor<4xi32>) -> tensor<1x?x?x16xf32> {
--  //      CHECK: "vhlo.dynamic_conv_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.dynamic_conv_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   batch_group_count = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   feature_group_count = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   input_batch_dimension = #vhlo.integer_v1<0 : i64>,
-@@ -531,8 +572,9 @@
- }
- 
- // CHECK-LABEL: "default_dynamic_gather"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @default_dynamic_gather(%arg0 : tensor<2x4x9xf32>, %arg1 : tensor<1x5x2xi32>, %arg2 : tensor<3xi32>) -> tensor<1x5x8xf32> {
--  //      CHECK: "vhlo.dynamic_gather_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.dynamic_gather_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   collapsed_slice_dims = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-   // CHECK-SAME:   index_vector_dim = #vhlo.integer_v1<2 : i64>,
-   // CHECK-SAME:   indices_are_sorted = #vhlo.bool_v1<false>,
-@@ -558,15 +600,16 @@
-   // CHECK-SAME:   sym_name = #vhlo.string_v1<"default_func">,
-   // CHECK-SAME:   sym_visibility = #vhlo.string_v1<"">
-   // CHECK-SAME: }> ({
--  // CHECK-NEXT:   ^[[BB:bb.*]](%arg0: !vhlo.tensor_v1<!vhlo.f32_v1>):
--  // CHECK-NEXT:     "vhlo.return_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-+  // CHECK-NEXT:   ^[[BB:bb.*]](%[[ARG0:.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
-+  // CHECK-NEXT:     "vhlo.return_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-   // CHECK-NEXT: }) : () -> ()
-   func.return %arg0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "dynamic_gather"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @dynamic_gather(%arg0 : tensor<2x4x9xf32>, %arg1 : tensor<1x5x2xi32>) -> tensor<1x5x1xf32> {
--  //      CHECK: "vhlo.gather_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.gather_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   collapsed_slice_dims = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-   // CHECK-SAME:   index_vector_dim = #vhlo.integer_v1<2 : i64>,
-   // CHECK-SAME:   indices_are_sorted = #vhlo.bool_v1<false>,
-@@ -587,8 +630,9 @@
- }
- 
- // CHECK-LABEL: "default_infeed"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_infeed(%arg0: !stablehlo.token) -> (tensor<f32>, !stablehlo.token) {
--  //               CHECK: "vhlo.infeed_v1"(%arg0) <{
-+  //               CHECK: "vhlo.infeed_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   infeed_config = #vhlo.string_v1<"">,
-   // CHECK-SAME{LITERAL}:   layout = #vhlo.array_v1<[]>
-   //          CHECK-SAME: }> : (!vhlo.token_v1) -> (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.token_v1)
-@@ -597,8 +641,9 @@
- }
- 
- // CHECK-LABEL: "default_outfeed"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @default_outfeed(%arg0: tensor<f32>, %arg1: !stablehlo.token) -> !stablehlo.token {
--  //      CHECK: "vhlo.outfeed_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.outfeed_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   outfeed_config = #vhlo.string_v1<"">
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.token_v1) -> !vhlo.token_v1
-   %0 = "stablehlo.outfeed"(%arg0, %arg1) : (tensor<f32>, !stablehlo.token) -> !stablehlo.token
-@@ -606,8 +651,9 @@
- }
- 
- // CHECK-LABEL: "default_recv"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_recv(%arg0: !stablehlo.token) -> (tensor<f32>, !stablehlo.token) {
--  //      CHECK: "vhlo.recv_v1"(%arg0) <{
-+  //      CHECK: "vhlo.recv_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME:   channel_type = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   is_host_transfer = #vhlo.bool_v1<false>
-@@ -619,8 +665,9 @@
- }
- 
- // CHECK-LABEL: "default_send"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @default_send(%arg0: tensor<f32>, %arg1: !stablehlo.token) -> !stablehlo.token {
--  //      CHECK: "vhlo.send_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.send_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME:   channel_type = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   is_host_transfer = #vhlo.bool_v1<false>
-@@ -632,8 +679,9 @@
- }
- 
- // CHECK-LABEL: "default_reduce_scatter"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_reduce_scatter(%arg0: tensor<16xf32>) -> tensor<16xf32> {
--  //               CHECK: "vhlo.reduce_scatter_v1"(%arg0) <{
-+  //               CHECK: "vhlo.reduce_scatter_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0], [1]]> : tensor<2x1xi64>>,
-   //          CHECK-SAME:   scatter_dimension = #vhlo.integer_v1<0 : i64>
-@@ -655,8 +703,9 @@
- }
- 
- // CHECK-LABEL: "default_reduce_window"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @default_reduce_window(%arg0: tensor<2x17x31x7xf32>, %arg1: tensor<f32>) -> tensor<2x16x30x7xf32> {
--  //               CHECK: "vhlo.reduce_window_v1"(%arg0, %arg1)  <{
-+  //               CHECK: "vhlo.reduce_window_v1"(%[[ARG0]], %[[ARG1]])  <{
-   //          CHECK-SAME:   base_dilations = #vhlo.tensor_v1<dense<1> : tensor<4xi64>>,
-   // CHECK-SAME{LITERAL}:   padding = #vhlo.tensor_v1<dense<0> : tensor<4x2xi64>>,
-   //          CHECK-SAME:   window_dilations = #vhlo.tensor_v1<dense<1> : tensor<4xi64>>,
-@@ -678,8 +727,9 @@
- }
- 
- // CHECK-LABEL: "default_scatter"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @default_scatter(%arg0: tensor<200x100x300xf32>, %arg1: tensor<10x2xi32>, %arg2: tensor<10x300xf32>) -> tensor<200x100x300xf32> {
--  //      CHECK: "vhlo.scatter_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.scatter_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   index_vector_dim = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   indices_are_sorted = #vhlo.bool_v1<false>,
-   // CHECK-SAME:   inserted_window_dims = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-@@ -707,8 +757,9 @@
- }
- 
- // CHECK-LABEL: "default_select_and_scatter"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @default_select_and_scatter(%arg0: tensor<10x24x24x64xf32>, %arg1: tensor<10x23x23x64xf32>, %arg2: tensor<f32>) -> tensor<10x24x24x64xf32> {
--  //      CHECK: "vhlo.select_and_scatter_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.select_and_scatter_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   padding = #vhlo.tensor_v1<dense<0> : tensor<4x2xi64>>,
-   // CHECK-SAME:   window_dimensions = #vhlo.tensor_v1<dense<[1, 2, 2, 1]> : tensor<4xi64>>,
-   // CHECK-SAME:   window_strides = #vhlo.tensor_v1<dense<1> : tensor<4xi64>>
-@@ -736,8 +787,9 @@
- }
- 
- // CHECK-LABEL: "default_sort"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_sort(%arg0: tensor<16xf32>) -> tensor<16xf32> {
--  //      CHECK: "vhlo.sort_v1"(%arg0) <{
-+  //      CHECK: "vhlo.sort_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   dimension = #vhlo.integer_v1<-1 : i64>
-   // CHECK-SAME:   is_stable = #vhlo.bool_v1<false>
-   // CHECK-SAME: }> ({
-@@ -756,29 +808,33 @@
- // ============ OPS ============
- 
- // CHECK-LABEL: "op_abs"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_abs(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.abs_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.abs_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.abs"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_add"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_add(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_after_all"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_after_all(%arg0: !stablehlo.token) -> !stablehlo.token {
--  // CHECK: "vhlo.after_all_v1"(%arg0) : (!vhlo.token_v1) -> !vhlo.token_v1
-+  // CHECK: "vhlo.after_all_v1"(%[[ARG0]]) : (!vhlo.token_v1) -> !vhlo.token_v1
-   %0 = "stablehlo.after_all"(%arg0) : (!stablehlo.token) -> !stablehlo.token
-   func.return %0 : !stablehlo.token
- }
- 
- // CHECK-LABEL: "op_all_gather"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_all_gather(%arg0: tensor<16x8xf32>) -> tensor<16x16xf32> {
--  //               CHECK: "vhlo.all_gather_v1"(%arg0) <{
-+  //               CHECK: "vhlo.all_gather_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   all_gather_dim = #vhlo.integer_v1<1 : i64>
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0], [1]]> : tensor<2x1xi64>>,
-@@ -794,8 +850,9 @@
- }
- 
- // CHECK-LABEL: "op_all_reduce"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_all_reduce(%arg0: tensor<f32>) -> tensor<f32> {
--  //               CHECK: "vhlo.all_reduce_v1"(%arg0) <{
-+  //               CHECK: "vhlo.all_reduce_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0], [1]]> : tensor<2x1xi64>>,
-   //          CHECK-SAME:   use_global_device_ids = #vhlo.bool_v1<true>
-@@ -817,8 +874,9 @@
- }
- 
- // CHECK-LABEL: "op_all_to_all"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_all_to_all(%arg0: tensor<4x16xf32>) -> tensor<16x4xf32> {
--  //               CHECK: "vhlo.all_to_all_v1"(%arg0) <{
-+  //               CHECK: "vhlo.all_to_all_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<1 : i64>,
-   //          CHECK-SAME:   concat_dimension = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0, 1, 2, 3]]> : tensor<1x4xi64>>,
-@@ -836,22 +894,25 @@
- }
- 
- // CHECK-LABEL: "op_and"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_and(%arg0: tensor<i1>, %arg1: tensor<i1>) -> tensor<i1> {
--  // CHECK: "vhlo.and_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-+  // CHECK: "vhlo.and_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-   %0 = "stablehlo.and"(%arg0, %arg1) : (tensor<i1>, tensor<i1>) -> tensor<i1>
-   func.return %0 : tensor<i1>
- }
- 
- // CHECK-LABEL: "op_atan2"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_atan2(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.atan2_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.atan2_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.atan2"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_batch_norm_grad"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}}, %[[ARG3:.*]]: {{.*}}, %[[ARG4:.*]]: {{.*}})
- func.func @op_batch_norm_grad(%arg0: tensor<16x16x16x16xf32>, %arg1: tensor<16xf32>, %arg2: tensor<16xf32>, %arg3: tensor<16xf32>, %arg4: tensor<16x16x16x16xf32>) -> (tensor<16x16x16x16xf32>, tensor<16xf32>, tensor<16xf32>) {
--  //      CHECK: "vhlo.batch_norm_grad_v1"(%arg0, %arg1, %arg2, %arg3, %arg4) <{
-+  //      CHECK: "vhlo.batch_norm_grad_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]], %[[ARG3]], %[[ARG4]]) <{
-   // CHECK-SAME:   epsilon = #vhlo.float_v1<1.000000e-03 : !vhlo.f32_v1>,
-   // CHECK-SAME:   feature_index = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<16x16x16x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x16x16x16x!vhlo.f32_v1>) -> (!vhlo.tensor_v1<16x16x16x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>)
-@@ -863,8 +924,9 @@
- }
- 
- // CHECK-LABEL: "op_batch_norm_inference"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}}, %[[ARG3:.*]]: {{.*}}, %[[ARG4:.*]]: {{.*}})
- func.func @op_batch_norm_inference(%arg0: tensor<16x16x16x16xf32>, %arg1: tensor<16xf32>, %arg2: tensor<16xf32>, %arg3: tensor<16xf32>, %arg4: tensor<16xf32>) -> tensor<16x16x16x16xf32> {
--  //      CHECK: "vhlo.batch_norm_inference_v1"(%arg0, %arg1, %arg2, %arg3, %arg4) <{
-+  //      CHECK: "vhlo.batch_norm_inference_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]], %[[ARG3]], %[[ARG4]]) <{
-   // CHECK-SAME:   epsilon = #vhlo.float_v1<1.000000e-03 : !vhlo.f32_v1>,
-   // CHECK-SAME:   feature_index = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<16x16x16x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x16x16x16x!vhlo.f32_v1>
-@@ -876,8 +938,9 @@
- }
- 
- // CHECK-LABEL: "op_batch_norm_training"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_batch_norm_training(%arg0: tensor<16x16x16x16xf32>, %arg1: tensor<16xf32>, %arg2: tensor<16xf32>) -> (tensor<16x16x16x16xf32>, tensor<16xf32>, tensor<16xf32>) {
--  //      CHECK: "vhlo.batch_norm_training_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.batch_norm_training_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   epsilon = #vhlo.float_v1<1.000000e-03 : !vhlo.f32_v1>,
-   // CHECK-SAME:   feature_index = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<16x16x16x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>) -> (!vhlo.tensor_v1<16x16x16x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>)
-@@ -889,15 +952,17 @@
- }
- 
- // CHECK-LABEL: "op_bitcast_convert"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_bitcast_convert(%arg0: tensor<i32>) -> tensor<f32> {
--  // CHECK: "vhlo.bitcast_convert_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.bitcast_convert_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.bitcast_convert"(%arg0) : (tensor<i32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_broadcast_in_dim"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_broadcast_in_dim(%arg0: tensor<16xf32>) -> tensor<16x16xf32> {
--  //      CHECK: "vhlo.broadcast_in_dim_v1"(%arg0) <{
-+  //      CHECK: "vhlo.broadcast_in_dim_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   broadcast_dimensions = #vhlo.tensor_v1<dense<1> : tensor<1xi64>>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x16x!vhlo.f32_v1>
-   %0 = "stablehlo.broadcast_in_dim"(%arg0) {
-@@ -907,8 +972,9 @@
- }
- 
- // CHECK-LABEL: "op_broadcast"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_broadcast(%arg0: tensor<16xf32>) -> tensor<16x16xf32> {
--  //      CHECK: "vhlo.broadcast_v1"(%arg0) <{
-+  //      CHECK: "vhlo.broadcast_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   broadcast_sizes = #vhlo.tensor_v1<dense<16> : tensor<1xi64>>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x16x!vhlo.f32_v1>
-   %0 = "stablehlo.broadcast"(%arg0) {
-@@ -918,9 +984,10 @@
- }
- 
- // CHECK-LABEL: "op_case"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_case(%arg0: tensor<i32>, %arg1: tensor<f32>) -> tensor<f32> {
--  //      CHECK: "vhlo.case_v1"(%arg0) ({
--  // CHECK-NEXT:   "vhlo.return_v1"(%arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-+  //      CHECK: "vhlo.case_v1"(%[[ARG0]]) ({
-+  // CHECK-NEXT:   "vhlo.return_v1"(%[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-   // CHECK-NEXT: }) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.case"(%arg0) ({
-     "stablehlo.return"(%arg1) : (tensor<f32>) -> ()
-@@ -929,22 +996,25 @@
- }
- 
- // CHECK-LABEL: "op_cbrt"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_cbrt(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.cbrt_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.cbrt_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.cbrt"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_ceil"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_ceil(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.ceil_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.ceil_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.ceil"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_cholesky"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_cholesky(%arg0: tensor<1x16x16xf32>) -> tensor<1x16x16xf32> {
--  //      CHECK: "vhlo.cholesky_v1"(%arg0) <{
-+  //      CHECK: "vhlo.cholesky_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   lower = #vhlo.bool_v1<true>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<1x16x16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<1x16x16x!vhlo.f32_v1>
-   %0 = "stablehlo.cholesky"(%arg0) {
-@@ -954,22 +1024,25 @@
- }
- 
- // CHECK-LABEL: "op_clamp"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_clamp(%arg0: tensor<f32>, %arg1: tensor<f32>, %arg2: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.clamp_v1"(%arg0, %arg1, %arg2) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.clamp_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.clamp"(%arg0, %arg1, %arg2) : (tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_count_leading_zeros"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_count_leading_zeros(%arg0: tensor<i32>) -> tensor<i32> {
--  // CHECK: "vhlo.count_leading_zeros_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-+  // CHECK: "vhlo.count_leading_zeros_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-   %0 = "stablehlo.count_leading_zeros"(%arg0) : (tensor<i32>) -> tensor<i32>
-   func.return %0 : tensor<i32>
- }
- 
- // CHECK-LABEL: "op_collective_permute"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_collective_permute(%arg0: tensor<16x8xf32>) -> tensor<16x8xf32> {
--  //               CHECK: "vhlo.collective_permute_v1"(%arg0) <{
-+  //               CHECK: "vhlo.collective_permute_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME{LITERAL}:   source_target_pairs = #vhlo.tensor_v1<dense<[[0, 1], [1, 2], [2, 3]]> : tensor<3x2xi64>>
-   //          CHECK-SAME: }> : (!vhlo.tensor_v1<16x8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x8x!vhlo.f32_v1>
-@@ -981,8 +1054,9 @@
- }
- 
- // CHECK-LABEL: "op_compare"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_compare(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
--  //      CHECK: "vhlo.compare_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.compare_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   compare_type = #vhlo<comparison_type_v1 TOTALORDER>,
-   // CHECK-SAME:   comparison_direction = #vhlo<comparison_direction_v1 EQ>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-@@ -994,15 +1068,17 @@
- }
- 
- // CHECK-LABEL: "op_complex"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_complex(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<complex<f32>> {
--  // CHECK: "vhlo.complex_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>
-+  // CHECK: "vhlo.complex_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>
-   %0 = "stablehlo.complex"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<complex<f32>>
-   func.return %0 : tensor<complex<f32>>
- }
- 
- // CHECK-LABEL: "op_concatenate"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_concatenate(%arg0: tensor<8xf32>, %arg1: tensor<8xf32>) -> tensor<16xf32> {
--  //      CHECK: "vhlo.concatenate_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.concatenate_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   dimension = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<8x!vhlo.f32_v1>, !vhlo.tensor_v1<8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x!vhlo.f32_v1>
-   %0 = "stablehlo.concatenate"(%arg0, %arg1) {
-@@ -1012,6 +1088,7 @@
- }
- 
- // CHECK-LABEL: "op_constant"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_constant(%arg0: tensor<f32>) -> tensor<f32> {
-   //      CHECK: "vhlo.constant_v1"() <{
-   // CHECK-SAME:   value = #vhlo.tensor_v1<dense<0.000000e+00> : tensor<f32>>
-@@ -1023,15 +1100,17 @@
- }
- 
- // CHECK-LABEL: "op_convert"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_convert(%arg0: tensor<i32>) -> tensor<f32> {
--  // CHECK: "vhlo.convert_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.convert_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.convert"(%arg0) : (tensor<i32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_convolution"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_convolution(%arg0: tensor<1x8x8x207xf32>, %arg1: tensor<3x3x207x16xf32>) -> tensor<1x7x7x16xf32> {
--  //      CHECK: "vhlo.convolution_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.convolution_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   batch_group_count = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   feature_group_count = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   input_batch_dimension = #vhlo.integer_v1<0 : i64>,
-@@ -1065,8 +1144,9 @@
- }
- 
- // CHECK-LABEL: "op_cosine"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_cosine(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.cosine_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.cosine_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.cosine"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
-@@ -1079,8 +1159,9 @@
- }
- 
- // CHECK-LABEL: "op_cross_replica_sum"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_cross_replica_sum(%arg0: tensor<f32>) -> tensor<f32> {
--  //               CHECK: "vhlo.cross-replica-sum_v1"(%arg0) <{
-+  //               CHECK: "vhlo.cross-replica-sum_v1"(%[[ARG0]]) <{
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0], [1]]> : tensor<2x1xi64>>
-   //          CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.cross-replica-sum"(%arg0) {
-@@ -1090,8 +1171,9 @@
- }
- 
- // CHECK-LABEL: "op_custom_call"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_custom_call(%arg0: tensor<f32>) -> tensor<f32> {
--  //      CHECK: "vhlo.custom_call_v1"(%arg0) <{
-+  //      CHECK: "vhlo.custom_call_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   api_version = #vhlo<api_version_v1 API_VERSION_STATUS_RETURNING>,
-   // CHECK-SAME:   backend_config = #vhlo.string_v1<"\08\03\1A\02">,
-   // CHECK-SAME:   call_target_name = #vhlo.string_v1<"foo">,
-@@ -1122,15 +1204,17 @@
- }
- 
- // CHECK-LABEL: "op_divide"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_divide(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.divide_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.divide_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.divide"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_dot_general"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_dot_general(%arg0: tensor<8x8x16xf32>, %arg1: tensor<8x16x8xf32>) -> tensor<8x8x8xf32> {
--  //      CHECK: "vhlo.dot_general_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.dot_general_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   lhs_batching_dimensions = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>,
-   // CHECK-SAME:   lhs_contracting_dimensions = #vhlo.tensor_v1<dense<2> : tensor<1xi64>>,
-   // CHECK-SAME:   precision_config = #vhlo.array_v1<[#vhlo<precision_v1 HIGHEST>, #vhlo<precision_v1 HIGHEST>]>,
-@@ -1150,8 +1234,9 @@
- }
- 
- // CHECK-LABEL: "op_dot"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_dot(%arg0: tensor<8x16xf32>, %arg1: tensor<16x8xf32>) -> tensor<8x8xf32> {
--  //      CHECK: "vhlo.dot_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.dot_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   precision_config = #vhlo.array_v1<[#vhlo<precision_v1 HIGHEST>, #vhlo<precision_v1 HIGHEST>]>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<8x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<8x8x!vhlo.f32_v1>
-   %0 = "stablehlo.dot"(%arg0, %arg1) {
-@@ -1161,8 +1246,9 @@
- }
- 
- // CHECK-LABEL: "op_dynamic_broadcast_in_dim"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_dynamic_broadcast_in_dim(%arg0: tensor<?x?xf32>, %arg1: tensor<2xindex>) -> tensor<?x?xf32> {
--  //      CHECK: "vhlo.dynamic_broadcast_in_dim_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.dynamic_broadcast_in_dim_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   broadcast_dimensions = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-   // CHECK-SAME:   known_expanding_dimensions = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>,
-   // CHECK-SAME:   known_nonexpanding_dimensions = #vhlo.tensor_v1<dense<1> : tensor<1xi64>>
-@@ -1176,8 +1262,9 @@
- }
- 
- // CHECK-LABEL: "op_dynamic_conv"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_dynamic_conv(%arg0: tensor<1x8x8x207xf32>, %arg1: tensor<3x3x207x16xf32>, %arg2: tensor<4xi32>) -> tensor<1x?x?x16xf32> {
--  //      CHECK: "vhlo.dynamic_conv_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.dynamic_conv_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   batch_group_count = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   feature_group_count = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   input_batch_dimension = #vhlo.integer_v1<0 : i64>,
-@@ -1211,8 +1298,9 @@
- }
- 
- // CHECK-LABEL: "op_dynamic_gather"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_dynamic_gather(%arg0 : tensor<2x4x9xf32>, %arg1 : tensor<1x5x2xi32>, %arg2 : tensor<3xi32>) -> tensor<1x5x8xf32> {
--  //      CHECK: "vhlo.dynamic_gather_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.dynamic_gather_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   collapsed_slice_dims = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-   // CHECK-SAME:   index_vector_dim = #vhlo.integer_v1<2 : i64>,
-   // CHECK-SAME:   indices_are_sorted = #vhlo.bool_v1<true>,
-@@ -1232,8 +1320,9 @@
- }
- 
- // CHECK-LABEL: "op_dynamic_iota"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_dynamic_iota(%arg0: tensor<1xindex>) -> tensor<?xf32> {
--  //      CHECK: "vhlo.dynamic_iota_v1"(%arg0) <{
-+  //      CHECK: "vhlo.dynamic_iota_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   iota_dimension = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<1x!vhlo.index_v1>) -> !vhlo.tensor_v1<?x!vhlo.f32_v1>
-   %0 = "stablehlo.dynamic_iota"(%arg0) {
-@@ -1243,22 +1332,25 @@
- }
- 
- // CHECK-LABEL: "op_dynamic_pad"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}}, %[[ARG3:.*]]: {{.*}}, %[[ARG4:.*]]: {{.*}})
- func.func @op_dynamic_pad(%arg0: tensor<?xf32>, %arg1: tensor<f32>, %arg2: tensor<1xindex>, %arg3: tensor<1xindex>, %arg4: tensor<1xindex>) -> tensor<?xf32> {
--  // CHECK: "vhlo.dynamic_pad_v1"(%arg0, %arg1, %arg2, %arg3, %arg4) : (!vhlo.tensor_v1<?x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>) -> !vhlo.tensor_v1<?x!vhlo.f32_v1>
-+  // CHECK: "vhlo.dynamic_pad_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]], %[[ARG3]], %[[ARG4]]) : (!vhlo.tensor_v1<?x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>) -> !vhlo.tensor_v1<?x!vhlo.f32_v1>
-   %0 = "stablehlo.dynamic_pad"(%arg0, %arg1, %arg2, %arg3, %arg4) : (tensor<?xf32>, tensor<f32>, tensor<1xindex>, tensor<1xindex>, tensor<1xindex>) -> tensor<?xf32>
-   func.return %0 : tensor<?xf32>
- }
- 
- // CHECK-LABEL: "op_dynamic_reshape"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_dynamic_reshape(%arg0: tensor<16xf32>, %arg1: tensor<2xindex>) -> tensor<?x?xf32> {
--  // CHECK: "vhlo.dynamic_reshape_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<2x!vhlo.index_v1>) -> !vhlo.tensor_v1<?x?x!vhlo.f32_v1>
-+  // CHECK: "vhlo.dynamic_reshape_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<2x!vhlo.index_v1>) -> !vhlo.tensor_v1<?x?x!vhlo.f32_v1>
-   %0 = "stablehlo.dynamic_reshape"(%arg0, %arg1) : (tensor<16xf32>, tensor<2xindex>) -> tensor<?x?xf32>
-   func.return %0 : tensor<?x?xf32>
- }
- 
- // CHECK-LABEL: "op_dynamic_slice"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_dynamic_slice(%arg0: tensor<16xf32>, %arg1: tensor<i64>) -> tensor<4xf32> {
--  //      CHECK: "vhlo.dynamic_slice_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.dynamic_slice_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   slice_sizes = #vhlo.tensor_v1<dense<4> : tensor<1xi64>>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.i64_v1>) -> !vhlo.tensor_v1<4x!vhlo.f32_v1>
-   %0 = "stablehlo.dynamic_slice"(%arg0, %arg1) {
-@@ -1268,15 +1360,17 @@
- }
- 
- // CHECK-LABEL: "op_dynamic_update_slice"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_dynamic_update_slice(%arg0: tensor<16xf32>, %arg1: tensor<4xf32>, %arg2: tensor<i64>) -> tensor<16xf32> {
--  // CHECK: "vhlo.dynamic_update_slice_v1"(%arg0, %arg1, %arg2) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<4x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.i64_v1>) -> !vhlo.tensor_v1<16x!vhlo.f32_v1>
-+  // CHECK: "vhlo.dynamic_update_slice_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<4x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.i64_v1>) -> !vhlo.tensor_v1<16x!vhlo.f32_v1>
-   %0 = "stablehlo.dynamic_update_slice"(%arg0, %arg1, %arg2) : (tensor<16xf32>, tensor<4xf32>, tensor<i64>) -> tensor<16xf32>
-   func.return %0 : tensor<16xf32>
- }
- 
- // CHECK-LABEL: "op_einsum"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_einsum(%arg0: tensor<8x16xf32>, %arg1: tensor<16x8xf32>) -> tensor<8x8xf32> {
--  //      CHECK: "vhlo.einsum_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.einsum_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   einsum_config = #vhlo.string_v1<"ab,bc->ac">
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<8x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<8x8x!vhlo.f32_v1>
-   %0 = "stablehlo.einsum"(%arg0, %arg1) {
-@@ -1286,22 +1380,25 @@
- }
- 
- // CHECK-LABEL: "op_exponential_minus_one"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_exponential_minus_one(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.exponential_minus_one_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.exponential_minus_one_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.exponential_minus_one"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_exponential"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_exponential(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.exponential_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.exponential_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.exponential"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_fft"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_fft(%arg0: tensor<16xcomplex<f32>>) -> tensor<16xcomplex<f32>> {
--  //      CHECK: "vhlo.fft_v1"(%arg0) <{
-+  //      CHECK: "vhlo.fft_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   fft_length = #vhlo.tensor_v1<dense<16> : tensor<1xi64>>,
-   // CHECK-SAME:   fft_type = #vhlo<fft_type_v1 FFT>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<16x!vhlo.complex_v1<!vhlo.f32_v1>>) -> !vhlo.tensor_v1<16x!vhlo.complex_v1<!vhlo.f32_v1>>
-@@ -1313,8 +1410,9 @@
- }
- 
- // CHECK-LABEL: "op_floor"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_floor(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.floor_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.floor_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.floor"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
-@@ -1327,16 +1425,17 @@
-   // CHECK-SAME:   sym_name = #vhlo.string_v1<"op_func">,
-   // CHECK-SAME:   sym_visibility = #vhlo.string_v1<"private">
-   // CHECK-SAME: }> ({
--  // CHECK-NEXT:   ^[[BB:bb.*]](%arg0: !vhlo.tensor_v1<!vhlo.f32_v1>):
--  // CHECK-NEXT:     "vhlo.return_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-+  // CHECK-NEXT:   ^[[BB:bb.*]](%[[ARG0:.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
-+  // CHECK-NEXT:     "vhlo.return_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-   // CHECK-NEXT: }) : () -> ()
- 
-   func.return %arg0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_gather"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_gather(%arg0 : tensor<2x4x9xf32>, %arg1 : tensor<1x5x2xi32>) -> tensor<1x5x1xf32> {
--  //      CHECK: "vhlo.gather_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.gather_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   collapsed_slice_dims = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-   // CHECK-SAME:   index_vector_dim = #vhlo.integer_v1<2 : i64>,
-   // CHECK-SAME:   indices_are_sorted = #vhlo.bool_v1<true>,
-@@ -1358,8 +1457,9 @@
- }
- 
- // CHECK-LABEL: "op_get_dimension_size"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_get_dimension_size(%arg0: tensor<?xf32>) -> tensor<i32> {
--  //      CHECK: "vhlo.get_dimension_size_v1"(%arg0) <{
-+  //      CHECK: "vhlo.get_dimension_size_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   dimension = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<?x!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-   %0 = "stablehlo.get_dimension_size"(%arg0) {
-@@ -1369,8 +1469,9 @@
- }
- 
- // CHECK-LABEL: "op_get_tuple_element"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_get_tuple_element(%arg0: tuple<tensor<f32>, tensor<i32>>) -> tensor<f32> {
--  //      CHECK: "vhlo.get_tuple_element_v1"(%arg0) <{
-+  //      CHECK: "vhlo.get_tuple_element_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   index = #vhlo.integer_v1<0 : i32>
-   // CHECK-SAME: }> : (!vhlo.tuple_v1<!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.get_tuple_element"(%arg0) {
-@@ -1380,11 +1481,12 @@
- }
- 
- // CHECK-LABEL: "op_if"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_if(%arg0: tensor<i1>, %arg1: tensor<f32>, %arg2: tensor<f32>) -> tensor<f32> {
--  //      CHECK: "vhlo.if_v1"(%arg0) ({
--  // CHECK-NEXT:   "vhlo.return_v1"(%arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-+  //      CHECK: "vhlo.if_v1"(%[[ARG0]]) ({
-+  // CHECK-NEXT:   "vhlo.return_v1"(%[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-   // CHECK-NEXT: }, {
--  // CHECK-NEXT:   "vhlo.return_v1"(%arg2) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-+  // CHECK-NEXT:   "vhlo.return_v1"(%[[ARG2]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-   // CHECK-NEXT: }) : (!vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.if"(%arg0) ({
-     "stablehlo.return"(%arg1) : (tensor<f32>) -> ()
-@@ -1395,15 +1497,17 @@
- }
- 
- // CHECK-LABEL: "op_imag"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_imag(%arg0: tensor<complex<f32>>) -> tensor<f32> {
--  // CHECK: "vhlo.imag_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.imag_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.imag"(%arg0) : (tensor<complex<f32>>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_infeed"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_infeed(%arg0: !stablehlo.token) -> (tensor<f32>, !stablehlo.token) {
--  //               CHECK: "vhlo.infeed_v1"(%arg0) <{
-+  //               CHECK: "vhlo.infeed_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   infeed_config = #vhlo.string_v1<"foo">,
-   // CHECK-SAME{LITERAL}:   layout = #vhlo.array_v1<[#vhlo.array_v1<[]>]>
-   //          CHECK-SAME: }> : (!vhlo.token_v1) -> (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.token_v1)
-@@ -1426,36 +1530,41 @@
- }
- 
- // CHECK-LABEL: "op_is_finite"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_is_finite(%arg0: tensor<f32>) -> tensor<i1> {
--  // CHECK: "vhlo.is_finite_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-+  // CHECK: "vhlo.is_finite_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-   %0 = "stablehlo.is_finite"(%arg0) : (tensor<f32>) -> tensor<i1>
-   func.return %0 : tensor<i1>
- }
- 
- // CHECK-LABEL: "op_log"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_log(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.log_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.log_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.log"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_log_plus_one"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_log_plus_one(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.log_plus_one_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.log_plus_one_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.log_plus_one"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_logistic"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_logistic(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.logistic_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.logistic_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.logistic"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_map"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_map(%arg0: tensor<16xf32>) -> tensor<16xf32> {
--  //      CHECK: "vhlo.map_v1"(%arg0) <{
-+  //      CHECK: "vhlo.map_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   dimensions = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>
-   // CHECK-SAME: }> ({
-   // CHECK-NEXT:   ^[[BB:bb.*]](%[[ARG1:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
-@@ -1473,57 +1582,65 @@
- }
- 
- // CHECK-LABEL: "op_maximum"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_maximum(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.maximum_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.maximum_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.maximum"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_minimum"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_minimum(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.minimum_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.minimum_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.minimum"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_multiply"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_multiply(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.multiply_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.multiply_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.multiply"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_negate"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_negate(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.negate_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.negate_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.negate"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_not"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_not(%arg0: tensor<i1>) -> tensor<i1> {
--  // CHECK: "vhlo.not_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-+  // CHECK: "vhlo.not_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-   %0 = "stablehlo.not"(%arg0) : (tensor<i1>) -> tensor<i1>
-   func.return %0 : tensor<i1>
- }
- 
- // CHECK-LABEL: "op_optimization_barrier"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_optimization_barrier(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.optimization_barrier_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.optimization_barrier_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.optimization_barrier"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_or"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_or(%arg0: tensor<i1>, %arg1: tensor<i1>) -> tensor<i1> {
--  // CHECK: "vhlo.or_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-+  // CHECK: "vhlo.or_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-   %0 = "stablehlo.or"(%arg0, %arg1) : (tensor<i1>, tensor<i1>) -> tensor<i1>
-   func.return %0 : tensor<i1>
- }
- 
- // CHECK-LABEL: "op_outfeed"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_outfeed(%arg0: tensor<f32>, %arg1: !stablehlo.token) -> !stablehlo.token {
--  //      CHECK: "vhlo.outfeed_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.outfeed_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   outfeed_config = #vhlo.string_v1<"foo">
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.token_v1) -> !vhlo.token_v1
-   %0 = "stablehlo.outfeed"(%arg0, %arg1) {
-@@ -1533,8 +1650,9 @@
- }
- 
- // CHECK-LABEL: "op_pad"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_pad(%arg0: tensor<8xf32>, %arg1: tensor<f32>) -> tensor<16xf32> {
--  //      CHECK: "vhlo.pad_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.pad_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   edge_padding_high = #vhlo.tensor_v1<dense<4> : tensor<1xi64>>,
-   // CHECK-SAME:   edge_padding_low = #vhlo.tensor_v1<dense<4> : tensor<1xi64>>,
-   // CHECK-SAME:   interior_padding = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>
-@@ -1548,36 +1666,41 @@
- }
- 
- // CHECK-LABEL: "op_popcnt"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_popcnt(%arg0: tensor<i32>) -> tensor<i32> {
--  // CHECK: "vhlo.popcnt_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-+  // CHECK: "vhlo.popcnt_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-   %0 = "stablehlo.popcnt"(%arg0) : (tensor<i32>) -> tensor<i32>
-   func.return %0 : tensor<i32>
- }
- 
- // CHECK-LABEL: "op_power"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_power(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.power_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.power_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.power"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_real_dynamic_slice"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}}, %[[ARG3:.*]]: {{.*}})
- func.func @op_real_dynamic_slice(%arg0: tensor<?xf32>, %arg1: tensor<1xindex>, %arg2: tensor<1xindex>, %arg3: tensor<1xindex>) -> tensor<?xf32> {
--  // CHECK: "vhlo.real_dynamic_slice_v1"(%arg0, %arg1, %arg2, %arg3) : (!vhlo.tensor_v1<?x!vhlo.f32_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>) -> !vhlo.tensor_v1<?x!vhlo.f32_v1>
-+  // CHECK: "vhlo.real_dynamic_slice_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]], %[[ARG3]]) : (!vhlo.tensor_v1<?x!vhlo.f32_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>) -> !vhlo.tensor_v1<?x!vhlo.f32_v1>
-   %0 = "stablehlo.real_dynamic_slice"(%arg0, %arg1, %arg2, %arg3) : (tensor<?xf32>, tensor<1xindex>, tensor<1xindex>, tensor<1xindex>) -> tensor<?xf32>
-   func.return %0 : tensor<?xf32>
- }
- 
- // CHECK-LABEL: "op_real"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_real(%arg0: tensor<complex<f32>>) -> tensor<f32> {
--  // CHECK: "vhlo.real_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.real_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.real"(%arg0) : (tensor<complex<f32>>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_recv"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_recv(%arg0: !stablehlo.token) -> (tensor<f32>, !stablehlo.token) {
--  //      CHECK: "vhlo.recv_v1"(%arg0) <{
-+  //      CHECK: "vhlo.recv_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME:   channel_type = #vhlo.integer_v1<3 : i64>,
-   // CHECK-SAME:   is_host_transfer = #vhlo.bool_v1<true>
-@@ -1590,8 +1713,9 @@
- }
- 
- // CHECK-LABEL: "op_reduce"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_reduce(%arg0: tensor<16xf32>, %arg1: tensor<f32>) -> tensor<f32> {
--  //  CHECK: "vhlo.reduce_v1"(%arg0, %arg1)
-+  //  CHECK: "vhlo.reduce_v1"(%[[ARG0]], %[[ARG1]])
-   //  CHECK:   ^[[BB:bb.*]](%[[ARG1:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>, %[[ARG2:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
-   //  CHECK:     "vhlo.return_v1"(%[[VAL1:.*]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-   //  CHECK: }) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-@@ -1606,8 +1730,9 @@
- }
- 
- // CHECK-LABEL: "op_reduce_precision"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_reduce_precision(%arg0: tensor<f32>) -> tensor<f32> {
--  //      CHECK: "vhlo.reduce_precision_v1"(%arg0) <{
-+  //      CHECK: "vhlo.reduce_precision_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   exponent_bits = #vhlo.integer_v1<8 : i32>
-   // CHECK-SAME:   mantissa_bits = #vhlo.integer_v1<10 : i32>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-@@ -1619,8 +1744,9 @@
- }
- 
- // CHECK-LABEL: "op_reduce_scatter"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_reduce_scatter(%arg0: tensor<16xf32>) -> tensor<16xf32> {
--  //               CHECK: "vhlo.reduce_scatter_v1"(%arg0) <{
-+  //               CHECK: "vhlo.reduce_scatter_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0], [1]]> : tensor<2x1xi64>>,
-   //          CHECK-SAME:   scatter_dimension = #vhlo.integer_v1<0 : i64>
-@@ -1644,8 +1770,9 @@
- }
- 
- // CHECK-LABEL: "op_reduce_window"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_reduce_window(%arg0: tensor<2x17x31x7xf32>, %arg1: tensor<f32>) -> tensor<2x9x16x7xf32> {
--  //               CHECK: "vhlo.reduce_window_v1"(%arg0, %arg1) <{
-+  //               CHECK: "vhlo.reduce_window_v1"(%[[ARG0]], %[[ARG1]]) <{
-   //          CHECK-SAME:   base_dilations = #vhlo.tensor_v1<dense<[1, 2, 2, 1]> : tensor<4xi64>>,
-   // CHECK-SAME{LITERAL}:   padding = #vhlo.tensor_v1<dense<[[0, 0], [2, 0], [0, 2], [0, 0]]> : tensor<4x2xi64>>,
-   //          CHECK-SAME:   window_dilations = #vhlo.tensor_v1<dense<[1, 2, 2, 1]> : tensor<4xi64>>,
-@@ -1671,8 +1798,9 @@
- }
- 
- // CHECK-LABEL: "op_remainder"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_remainder(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.remainder_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.remainder_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.remainder"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
-@@ -1692,16 +1820,18 @@
- }
- 
- // CHECK-LABEL: "op_reshape"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_reshape(%arg0: tensor<16xf32>) -> tensor<4x4xf32> {
--  // CHECK: "vhlo.reshape_v1"(%arg0) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<4x4x!vhlo.f32_v1>
-+  // CHECK: "vhlo.reshape_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<4x4x!vhlo.f32_v1>
-   %0 = "stablehlo.reshape"(%arg0) : (tensor<16xf32>) -> tensor<4x4xf32>
-   func.return %0 : tensor<4x4xf32>
- }
- 
- // CHECK-LABEL: "op_return"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_return(%arg0: tensor<i32>, %arg1: tensor<f32>) -> tensor<f32> {
--  //      CHECK: "vhlo.case_v1"(%arg0) ({
--  // CHECK-NEXT:   "vhlo.return_v1"(%arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-+  //      CHECK: "vhlo.case_v1"(%[[ARG0]]) ({
-+  // CHECK-NEXT:   "vhlo.return_v1"(%[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-   // CHECK-NEXT: }) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.case"(%arg0) ({
-     "stablehlo.return"(%arg1) : (tensor<f32>) -> ()
-@@ -1710,8 +1840,9 @@
- }
- 
- // CHECK-LABEL: "op_reverse"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_reverse(%arg0: tensor<16xf32>) -> tensor<16xf32> {
--  //      CHECK: "vhlo.reverse_v1"(%arg0) <{
-+  //      CHECK: "vhlo.reverse_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   dimensions = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x!vhlo.f32_v1>
-   %0 = "stablehlo.reverse"(%arg0) {
-@@ -1721,8 +1852,9 @@
- }
- 
- // CHECK-LABEL: "op_rng_bit_generator"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_rng_bit_generator(%arg0: tensor<f32>) -> (tensor<f32>, tensor<f32>) {
--  //      CHECK: "vhlo.rng_bit_generator_v1"(%arg0) <{
-+  //      CHECK: "vhlo.rng_bit_generator_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   rng_algorithm = #vhlo<rng_algorithm_v1 PHILOX>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>)
-   %0:2 = "stablehlo.rng_bit_generator"(%arg0) {
-@@ -1732,8 +1864,9 @@
- }
- 
- // CHECK-LABEL: "op_rng"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_rng(%arg0: tensor<f32>, %arg1: tensor<f32>, %arg2: tensor<0xindex>) -> tensor<f32> {
--  //      CHECK: "vhlo.rng_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.rng_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   rng_distribution = #vhlo<rng_distribution_v1 NORMAL>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<0x!vhlo.index_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.rng"(%arg0, %arg1, %arg2) {
-@@ -1743,29 +1876,33 @@
- }
- 
- // CHECK-LABEL: "op_round_nearest_afz"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_round_nearest_afz(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.round_nearest_afz_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.round_nearest_afz_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.round_nearest_afz"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_round_nearest_even"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_round_nearest_even(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.round_nearest_even_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.round_nearest_even_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.round_nearest_even"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_rsqrt"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_rsqrt(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.rsqrt_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.rsqrt_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.rsqrt"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_scatter"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_scatter(%arg0: tensor<200x100x300xf32>, %arg1: tensor<10x2xi32>, %arg2: tensor<10x300xf32>) -> tensor<200x100x300xf32> {
--  //      CHECK: "vhlo.scatter_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.scatter_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   index_vector_dim = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   indices_are_sorted = #vhlo.bool_v1<true>,
-   // CHECK-SAME:   inserted_window_dims = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-@@ -1795,8 +1932,9 @@
- }
- 
- // CHECK-LABEL: "op_select_and_scatter"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_select_and_scatter(%arg0: tensor<10x24x24x64xf32>, %arg1: tensor<12x13x13x66xf32>, %arg2: tensor<f32>) -> tensor<10x24x24x64xf32> {
--  //      CHECK: "vhlo.select_and_scatter_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.select_and_scatter_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   padding = #vhlo.tensor_v1<dense<1> : tensor<4x2xi64>>,
-   // CHECK-SAME:   window_dimensions = #vhlo.tensor_v1<dense<[1, 2, 2, 1]> : tensor<4xi64>>,
-   // CHECK-SAME:   window_strides = #vhlo.tensor_v1<dense<[1, 2, 2, 1]> : tensor<4xi64>>
-@@ -1826,15 +1964,17 @@
- }
- 
- // CHECK-LABEL: "op_select"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_select(%arg0: tensor<i1>, %arg1: tensor<f32>, %arg2: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.select_v1"(%arg0, %arg1, %arg2) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.select_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.select"(%arg0, %arg1, %arg2) : (tensor<i1>, tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_send"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_send(%arg0: tensor<f32>, %arg1: !stablehlo.token) -> !stablehlo.token {
--  //      CHECK: "vhlo.send_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.send_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME:   channel_type = #vhlo.integer_v1<2 : i64>,
-   // CHECK-SAME:   is_host_transfer = #vhlo.bool_v1<true>
-@@ -1847,8 +1987,9 @@
- }
- 
- // CHECK-LABEL: "op_set_dimension_size"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_set_dimension_size(%arg0: tensor<?xf32>, %arg1: tensor<i32>) -> tensor<16xf32> {
--  //      CHECK: "vhlo.set_dimension_size_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.set_dimension_size_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   dimension = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<?x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<16x!vhlo.f32_v1>
-   %0 = "stablehlo.set_dimension_size"(%arg0, %arg1) {
-@@ -1858,43 +1999,49 @@
- }
- 
- // CHECK-LABEL: "op_shift_left"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_shift_left(%arg0: tensor<i32>, %arg1: tensor<i32>) -> tensor<i32> {
--  // CHECK: "vhlo.shift_left_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.i32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-+  // CHECK: "vhlo.shift_left_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-   %0 = "stablehlo.shift_left"(%arg0, %arg1) : (tensor<i32>, tensor<i32>) -> tensor<i32>
-   func.return %0 : tensor<i32>
- }
- 
- // CHECK-LABEL: "op_shift_right_arithmetic"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_shift_right_arithmetic(%arg0: tensor<i32>, %arg1: tensor<i32>) -> tensor<i32> {
--  // CHECK: "vhlo.shift_right_arithmetic_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.i32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-+  // CHECK: "vhlo.shift_right_arithmetic_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-   %0 = "stablehlo.shift_right_arithmetic"(%arg0, %arg1) : (tensor<i32>, tensor<i32>) -> tensor<i32>
-   func.return %0 : tensor<i32>
- }
- 
- // CHECK-LABEL: "op_shift_right_logical"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_shift_right_logical(%arg0: tensor<i32>, %arg1: tensor<i32>) -> tensor<i32> {
--  // CHECK: "vhlo.shift_right_logical_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.i32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-+  // CHECK: "vhlo.shift_right_logical_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-   %0 = "stablehlo.shift_right_logical"(%arg0, %arg1) : (tensor<i32>, tensor<i32>) -> tensor<i32>
-   func.return %0 : tensor<i32>
- }
- 
- // CHECK-LABEL: "op_sign"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_sign(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.sign_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.sign_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.sign"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_sine"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_sine(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.sine_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.sine_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.sine"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_slice"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_slice(%arg0: tensor<16xf32>) -> tensor<4xf32> {
--  //      CHECK: "vhlo.slice_v1"(%arg0) <{
-+  //      CHECK: "vhlo.slice_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   limit_indices = #vhlo.tensor_v1<dense<4> : tensor<1xi64>>,
-   // CHECK-SAME:   start_indices = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>,
-   // CHECK-SAME:   strides = #vhlo.tensor_v1<dense<1> : tensor<1xi64>>
-@@ -1908,8 +2055,9 @@
- }
- 
- // CHECK-LABEL: "op_sort"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_sort(%arg0: tensor<16xf32>) -> tensor<16xf32> {
--  //      CHECK: "vhlo.sort_v1"(%arg0) <{
-+  //      CHECK: "vhlo.sort_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   dimension = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME:   is_stable = #vhlo.bool_v1<true>
-   // CHECK-SAME: }> ({
-@@ -1929,29 +2077,33 @@
- }
- 
- // CHECK-LABEL: "op_sqrt"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_sqrt(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.sqrt_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.sqrt_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.sqrt"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_subtract"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_subtract(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.subtract_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.subtract_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.subtract"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_tanh"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_tanh(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.tanh_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.tanh_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.tanh"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_torch_index_select"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_torch_index_select(%arg0: tensor<5x1x5xf32>, %arg1: tensor<2xi32>) ->  tensor<2x1x5xf32> {
--  //      CHECK: "vhlo.torch_index_select_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.torch_index_select_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   batch_dims = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME:   dim = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<5x1x5x!vhlo.f32_v1>, !vhlo.tensor_v1<2x!vhlo.i32_v1>) -> !vhlo.tensor_v1<2x1x5x!vhlo.f32_v1>
-@@ -1963,8 +2115,9 @@
- }
- 
- // CHECK-LABEL: "op_trace"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_trace(%arg0: tensor<f32>) {
--  //      CHECK: "vhlo.trace_v1"(%arg0) <{
-+  //      CHECK: "vhlo.trace_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   tag = #vhlo.string_v1<"foo">
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-   "stablehlo.trace"(%arg0) {
-@@ -1974,8 +2127,9 @@
- }
- 
- // CHECK-LABEL: "op_transpose"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_transpose(%arg0: tensor<16x8xf32>) ->  tensor<8x16xf32> {
--  //      CHECK: "vhlo.transpose_v1"(%arg0) <{
-+  //      CHECK: "vhlo.transpose_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   permutation = #vhlo.tensor_v1<dense<[1, 0]> : tensor<2xi64>>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<16x8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<8x16x!vhlo.f32_v1>
-   %0 = "stablehlo.transpose"(%arg0) {
-@@ -1985,8 +2139,9 @@
- }
- 
- // CHECK-LABEL: "op_triangular_solve"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_triangular_solve(%arg0: tensor<16x16xf32>, %arg1: tensor<16x16xf32>) ->  tensor<16x16xf32> {
--  //      CHECK: "vhlo.triangular_solve_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.triangular_solve_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   left_side = #vhlo.bool_v1<true>,
-   // CHECK-SAME:   lower = #vhlo.bool_v1<true>,
-   // CHECK-SAME:   transpose_a = #vhlo<transpose_v1 NO_TRANSPOSE>,
-@@ -2002,15 +2157,17 @@
- }
- 
- // CHECK-LABEL: "op_tuple"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_tuple(%arg0: tensor<f32>) -> tuple<tensor<f32>> {
--  // CHECK: "vhlo.tuple_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tuple_v1<!vhlo.tensor_v1<!vhlo.f32_v1>>
-+  // CHECK: "vhlo.tuple_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tuple_v1<!vhlo.tensor_v1<!vhlo.f32_v1>>
-   %0 = "stablehlo.tuple"(%arg0) : (tensor<f32>) -> tuple<tensor<f32>>
-   func.return %0 : tuple<tensor<f32>>
- }
- 
- // CHECK-LABEL: "op_unary_einsum"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_unary_einsum(%arg0: tensor<8x16xf32>) -> tensor<8xf32> {
--  //      CHECK: "vhlo.unary_einsum_v1"(%arg0) <{
-+  //      CHECK: "vhlo.unary_einsum_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   einsum_config = #vhlo.string_v1<"ab->a">
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<8x16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<8x!vhlo.f32_v1>
-   %0 = "stablehlo.unary_einsum"(%arg0) {
-@@ -2020,22 +2177,25 @@
- }
- 
- // CHECK-LABEL: "op_uniform_dequantize"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_uniform_dequantize(%arg0: tensor<!quant.uniform<i8:f32, 34.0:16>>) -> tensor<f32> {
--  // CHECK: "vhlo.uniform_dequantize_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.uniform_dequantize_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.uniform_dequantize"(%arg0) : (tensor<!quant.uniform<i8:f32, 34.0:16>>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_uniform_quantize"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_uniform_quantize(%arg0: tensor<f32>) -> tensor<!quant.uniform<i8:f32, 34.0:16>> {
--  // CHECK: "vhlo.uniform_quantize_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>
-+  // CHECK: "vhlo.uniform_quantize_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>
-   %0 = "stablehlo.uniform_quantize"(%arg0) : (tensor<f32>) -> tensor<!quant.uniform<i8:f32, 34.0:16>>
-   func.return %0 : tensor<!quant.uniform<i8:f32, 34.0:16>>
- }
- 
- // CHECK-LABEL: "op_while"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_while(%arg0: tensor<i1>) -> tensor<i1> {
--  //      CHECK: "vhlo.while_v1"(%arg0) ({
-+  //      CHECK: "vhlo.while_v1"(%[[ARG0]]) ({
-   // CHECK-NEXT:   ^[[BB:bb.*]](%[[ARG1:arg.*]]: !vhlo.tensor_v1<!vhlo.bool_v1>):
-   // CHECK-NEXT:     "vhlo.return_v1"(%[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>) -> ()
-   // CHECK-NEXT:   }, {
-@@ -2053,8 +2213,9 @@
- }
- 
- // CHECK-LABEL: "op_xor"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_xor(%arg0: tensor<i1>, %arg1: tensor<i1>) -> tensor<i1> {
--  // CHECK: "vhlo.xor_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-+  // CHECK: "vhlo.xor_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-   %0 = "stablehlo.xor"(%arg0, %arg1) : (tensor<i1>, tensor<i1>) -> tensor<i1>
-   func.return %0 : tensor<i1>
- }
-@@ -2062,190 +2223,217 @@
- // ============ TYPES ============
- 
- // CHECK-LABEL: "type_i1"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_i1(%arg0: tensor<i1>, %arg1: tensor<i1>) -> tensor<i1> {
--  // CHECK: "vhlo.and_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-+  // CHECK: "vhlo.and_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-   %0 = "stablehlo.and"(%arg0, %arg1) : (tensor<i1>, tensor<i1>) -> tensor<i1>
-   func.return %0 : tensor<i1>
- }
- 
- // CHECK-LABEL: "type_i4"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_i4(%arg0: tensor<i4>, %arg1: tensor<i4>) -> tensor<i4> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.i4_v1>, !vhlo.tensor_v1<!vhlo.i4_v1>) -> !vhlo.tensor_v1<!vhlo.i4_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i4_v1>, !vhlo.tensor_v1<!vhlo.i4_v1>) -> !vhlo.tensor_v1<!vhlo.i4_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<i4>, tensor<i4>) -> tensor<i4>
-   func.return %0 : tensor<i4>
- }
- 
- // CHECK-LABEL: "type_i8"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_i8(%arg0: tensor<i8>, %arg1: tensor<i8>) -> tensor<i8> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.i8_v1>, !vhlo.tensor_v1<!vhlo.i8_v1>) -> !vhlo.tensor_v1<!vhlo.i8_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i8_v1>, !vhlo.tensor_v1<!vhlo.i8_v1>) -> !vhlo.tensor_v1<!vhlo.i8_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<i8>, tensor<i8>) -> tensor<i8>
-   func.return %0 : tensor<i8>
- }
- 
- // CHECK-LABEL: "type_i16"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_i16(%arg0: tensor<i16>, %arg1: tensor<i16>) -> tensor<i16> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.i16_v1>, !vhlo.tensor_v1<!vhlo.i16_v1>) -> !vhlo.tensor_v1<!vhlo.i16_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i16_v1>, !vhlo.tensor_v1<!vhlo.i16_v1>) -> !vhlo.tensor_v1<!vhlo.i16_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<i16>, tensor<i16>) -> tensor<i16>
-   func.return %0 : tensor<i16>
- }
- 
- // CHECK-LABEL: "type_i32"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_i32(%arg0: tensor<i32>, %arg1: tensor<i32>) -> tensor<i32> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.i32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<i32>, tensor<i32>) -> tensor<i32>
-   func.return %0 : tensor<i32>
- }
- 
- // CHECK-LABEL: "type_i64"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_i64(%arg0: tensor<i64>, %arg1: tensor<i64>) -> tensor<i64> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.i64_v1>, !vhlo.tensor_v1<!vhlo.i64_v1>) -> !vhlo.tensor_v1<!vhlo.i64_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i64_v1>, !vhlo.tensor_v1<!vhlo.i64_v1>) -> !vhlo.tensor_v1<!vhlo.i64_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<i64>, tensor<i64>) -> tensor<i64>
-   func.return %0 : tensor<i64>
- }
- 
- // CHECK-LABEL: "type_ui4"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_ui4(%arg0: tensor<ui4>, %arg1: tensor<ui4>) -> tensor<ui4> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.ui4_v1>, !vhlo.tensor_v1<!vhlo.ui4_v1>) -> !vhlo.tensor_v1<!vhlo.ui4_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.ui4_v1>, !vhlo.tensor_v1<!vhlo.ui4_v1>) -> !vhlo.tensor_v1<!vhlo.ui4_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<ui4>, tensor<ui4>) -> tensor<ui4>
-   func.return %0 : tensor<ui4>
- }
- 
- // CHECK-LABEL: "type_ui8"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_ui8(%arg0: tensor<ui8>, %arg1: tensor<ui8>) -> tensor<ui8> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.ui8_v1>, !vhlo.tensor_v1<!vhlo.ui8_v1>) -> !vhlo.tensor_v1<!vhlo.ui8_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.ui8_v1>, !vhlo.tensor_v1<!vhlo.ui8_v1>) -> !vhlo.tensor_v1<!vhlo.ui8_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<ui8>, tensor<ui8>) -> tensor<ui8>
-   func.return %0 : tensor<ui8>
- }
- 
- // CHECK-LABEL: "type_ui16"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_ui16(%arg0: tensor<ui16>, %arg1: tensor<ui16>) -> tensor<ui16> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.ui16_v1>, !vhlo.tensor_v1<!vhlo.ui16_v1>) -> !vhlo.tensor_v1<!vhlo.ui16_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.ui16_v1>, !vhlo.tensor_v1<!vhlo.ui16_v1>) -> !vhlo.tensor_v1<!vhlo.ui16_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<ui16>, tensor<ui16>) -> tensor<ui16>
-   func.return %0 : tensor<ui16>
- }
- 
- // CHECK-LABEL: "type_ui32"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_ui32(%arg0: tensor<ui32>, %arg1: tensor<ui32>) -> tensor<ui32> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.ui32_v1>, !vhlo.tensor_v1<!vhlo.ui32_v1>) -> !vhlo.tensor_v1<!vhlo.ui32_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.ui32_v1>, !vhlo.tensor_v1<!vhlo.ui32_v1>) -> !vhlo.tensor_v1<!vhlo.ui32_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<ui32>, tensor<ui32>) -> tensor<ui32>
-   func.return %0 : tensor<ui32>
- }
- 
- // CHECK-LABEL: "type_ui64"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_ui64(%arg0: tensor<ui64>, %arg1: tensor<ui64>) -> tensor<ui64> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.ui64_v1>, !vhlo.tensor_v1<!vhlo.ui64_v1>) -> !vhlo.tensor_v1<!vhlo.ui64_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.ui64_v1>, !vhlo.tensor_v1<!vhlo.ui64_v1>) -> !vhlo.tensor_v1<!vhlo.ui64_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<ui64>, tensor<ui64>) -> tensor<ui64>
-   func.return %0 : tensor<ui64>
- }
- 
- // CHECK-LABEL: "type_f8E4M3FN"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_f8E4M3FN(%arg0: tensor<f8E4M3FN>, %arg1: tensor<f8E4M3FN>) -> tensor<f8E4M3FN> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f8E4M3FN_v1>, !vhlo.tensor_v1<!vhlo.f8E4M3FN_v1>) -> !vhlo.tensor_v1<!vhlo.f8E4M3FN_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f8E4M3FN_v1>, !vhlo.tensor_v1<!vhlo.f8E4M3FN_v1>) -> !vhlo.tensor_v1<!vhlo.f8E4M3FN_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f8E4M3FN>, tensor<f8E4M3FN>) -> tensor<f8E4M3FN>
-   func.return %0 : tensor<f8E4M3FN>
- }
- 
- // CHECK-LABEL: "type_f8E5M2"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_f8E5M2(%arg0: tensor<f8E5M2>, %arg1: tensor<f8E5M2>) -> tensor<f8E5M2> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f8E5M2_v1>, !vhlo.tensor_v1<!vhlo.f8E5M2_v1>) -> !vhlo.tensor_v1<!vhlo.f8E5M2_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f8E5M2_v1>, !vhlo.tensor_v1<!vhlo.f8E5M2_v1>) -> !vhlo.tensor_v1<!vhlo.f8E5M2_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f8E5M2>, tensor<f8E5M2>) -> tensor<f8E5M2>
-   func.return %0 : tensor<f8E5M2>
- }
- 
- // CHECK-LABEL: "type_f8E4M3FNUZ"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_f8E4M3FNUZ(%arg0: tensor<f8E4M3FNUZ>, %arg1: tensor<f8E4M3FNUZ>) -> tensor<f8E4M3FNUZ> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f8E4M3FNUZ_v1>, !vhlo.tensor_v1<!vhlo.f8E4M3FNUZ_v1>) -> !vhlo.tensor_v1<!vhlo.f8E4M3FNUZ_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f8E4M3FNUZ_v1>, !vhlo.tensor_v1<!vhlo.f8E4M3FNUZ_v1>) -> !vhlo.tensor_v1<!vhlo.f8E4M3FNUZ_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f8E4M3FNUZ>, tensor<f8E4M3FNUZ>) -> tensor<f8E4M3FNUZ>
-   func.return %0 : tensor<f8E4M3FNUZ>
- }
- 
- // CHECK-LABEL: "type_f8E4M3B11FNUZ"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_f8E4M3B11FNUZ(%arg0: tensor<f8E4M3B11FNUZ>, %arg1: tensor<f8E4M3B11FNUZ>) -> tensor<f8E4M3B11FNUZ> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f8E4M3B11FNUZ_v1>, !vhlo.tensor_v1<!vhlo.f8E4M3B11FNUZ_v1>) -> !vhlo.tensor_v1<!vhlo.f8E4M3B11FNUZ_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f8E4M3B11FNUZ_v1>, !vhlo.tensor_v1<!vhlo.f8E4M3B11FNUZ_v1>) -> !vhlo.tensor_v1<!vhlo.f8E4M3B11FNUZ_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f8E4M3B11FNUZ>, tensor<f8E4M3B11FNUZ>) -> tensor<f8E4M3B11FNUZ>
-   func.return %0 : tensor<f8E4M3B11FNUZ>
- }
- 
- // CHECK-LABEL: "type_f8E5M2FNUZ"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_f8E5M2FNUZ(%arg0: tensor<f8E5M2FNUZ>, %arg1: tensor<f8E5M2FNUZ>) -> tensor<f8E5M2FNUZ> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f8E5M2FNUZ_v1>, !vhlo.tensor_v1<!vhlo.f8E5M2FNUZ_v1>) -> !vhlo.tensor_v1<!vhlo.f8E5M2FNUZ_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f8E5M2FNUZ_v1>, !vhlo.tensor_v1<!vhlo.f8E5M2FNUZ_v1>) -> !vhlo.tensor_v1<!vhlo.f8E5M2FNUZ_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f8E5M2FNUZ>, tensor<f8E5M2FNUZ>) -> tensor<f8E5M2FNUZ>
-   func.return %0 : tensor<f8E5M2FNUZ>
- }
- 
- // CHECK-LABEL: "type_bf16"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_bf16(%arg0: tensor<bf16>, %arg1: tensor<bf16>) -> tensor<bf16> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.bf16_v1>, !vhlo.tensor_v1<!vhlo.bf16_v1>) -> !vhlo.tensor_v1<!vhlo.bf16_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.bf16_v1>, !vhlo.tensor_v1<!vhlo.bf16_v1>) -> !vhlo.tensor_v1<!vhlo.bf16_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<bf16>, tensor<bf16>) -> tensor<bf16>
-   func.return %0 : tensor<bf16>
- }
- 
- // CHECK-LABEL: "type_f16"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_f16(%arg0: tensor<f16>, %arg1: tensor<f16>) -> tensor<f16> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f16_v1>, !vhlo.tensor_v1<!vhlo.f16_v1>) -> !vhlo.tensor_v1<!vhlo.f16_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f16_v1>, !vhlo.tensor_v1<!vhlo.f16_v1>) -> !vhlo.tensor_v1<!vhlo.f16_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f16>, tensor<f16>) -> tensor<f16>
-   func.return %0 : tensor<f16>
- }
- 
- // CHECK-LABEL: "type_f32"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_f32(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "type_f64"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_f64(%arg0: tensor<f64>, %arg1: tensor<f64>) -> tensor<f64> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f64_v1>, !vhlo.tensor_v1<!vhlo.f64_v1>) -> !vhlo.tensor_v1<!vhlo.f64_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f64_v1>, !vhlo.tensor_v1<!vhlo.f64_v1>) -> !vhlo.tensor_v1<!vhlo.f64_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f64>, tensor<f64>) -> tensor<f64>
-   func.return %0 : tensor<f64>
- }
- 
- // CHECK-LABEL: "type_complex_f32"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_complex_f32(%arg0: tensor<complex<f32>>, %arg1: tensor<complex<f32>>) -> tensor<complex<f32>> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>, !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>) -> !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>, !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>) -> !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<complex<f32>>, tensor<complex<f32>>) -> tensor<complex<f32>>
-   func.return %0 : tensor<complex<f32>>
- }
- 
- // CHECK-LABEL: "type_complex_f64"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_complex_f64(%arg0: tensor<complex<f64>>, %arg1: tensor<complex<f64>>) -> tensor<complex<f64>> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f64_v1>>, !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f64_v1>>) -> !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f64_v1>>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f64_v1>>, !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f64_v1>>) -> !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f64_v1>>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<complex<f64>>, tensor<complex<f64>>) -> tensor<complex<f64>>
-   func.return %0 : tensor<complex<f64>>
- }
- 
- // CHECK-LABEL: "type_dynamism_ranked"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @type_dynamism_ranked(%arg0: tensor<?xf32>) -> tensor<?xf32> {
--  // CHECK: "vhlo.abs_v1"(%arg0) : (!vhlo.tensor_v1<?x!vhlo.f32_v1>) -> !vhlo.tensor_v1<?x!vhlo.f32_v1>
-+  // CHECK: "vhlo.abs_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<?x!vhlo.f32_v1>) -> !vhlo.tensor_v1<?x!vhlo.f32_v1>
-   %0 = "stablehlo.abs"(%arg0) : (tensor<?xf32>) -> tensor<?xf32>
-   func.return %0 : tensor<?xf32>
- }
- 
- // CHECK-LABEL: "type_quantization"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_quantization(%arg0: tensor<!quant.uniform<i8:f32, 34.0:16>>, %arg1: tensor<!quant.uniform<i8:f32, 34.0:16>>) -> tensor<!quant.uniform<i8:f32, 34.0:16>> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>, !vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>) -> !vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>, !vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>) -> !vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<!quant.uniform<i8:f32, 34.0:16>>, tensor<!quant.uniform<i8:f32, 34.0:16>>) -> tensor<!quant.uniform<i8:f32, 34.0:16>>
-   func.return %0 : tensor<!quant.uniform<i8:f32, 34.0:16>>
- }
- 
- //       CHECK: function_type = #vhlo.type_v1<!vhlo.func_v1<(!vhlo.token_v1) -> !vhlo.token_v1>>
- // CHECK-LABEL: "type_token_callee"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @type_token_callee(%arg0: !stablehlo.token) -> !stablehlo.token {
--  // CHECK: "vhlo.return_v1"(%arg0) : (!vhlo.token_v1) -> ()
-+  // CHECK: "vhlo.return_v1"(%[[ARG0]]) : (!vhlo.token_v1) -> ()
-   return %arg0 : !stablehlo.token
- }
- 
- //       CHECK: function_type = #vhlo.type_v1<!vhlo.func_v1<(!vhlo.token_v1) -> !vhlo.token_v1>>
- // CHECK-LABEL: "type_token_caller"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @type_token_caller(%arg0: !stablehlo.token) -> !stablehlo.token {
--  // CHECK:      "vhlo.call_v1"(%arg0) <{callee = #vhlo.string_v1<"type_token_callee">}
-+  // CHECK:      "vhlo.call_v1"(%[[ARG0]]) <{callee = #vhlo.string_v1<"type_token_callee">}
-   // CHECK-SAME: (!vhlo.token_v1) -> !vhlo.token_v1
-   %0 = func.call @type_token_callee(%arg0) : (!stablehlo.token) -> !stablehlo.token
-   return %0 : !stablehlo.token
- }
- 
- // CHECK-LABEL: "type_tuple"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @type_tuple(%arg0: tuple<tensor<f32>>) -> tuple<!stablehlo.token> {
-   %0 = "stablehlo.custom_call"(%arg0) {
-     call_target_name = "foo"
-diff --ruN a/stablehlo/stablehlo/tests/vhlo/stablehlo_legalize_to_vhlo.0_12_0.mlir b/stablehlo/stablehlo/tests/vhlo/stablehlo_legalize_to_vhlo.0_12_0.mlir
---- stablehlo/stablehlo/tests/vhlo/stablehlo_legalize_to_vhlo.0_12_0.mlir
-+++ stablehlo/stablehlo/tests/vhlo/stablehlo_legalize_to_vhlo.0_12_0.mlir
-@@ -13,6 +13,7 @@
- // ============ ATTRIBUTES ============
- 
- // CHECK-LABEL: "attr_comparison_direction_eq"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_direction_eq(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     // CHECK: comparison_direction = #vhlo<comparison_direction_v1 EQ>
-@@ -22,6 +23,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_direction_ne"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_direction_ne(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     // CHECK: comparison_direction = #vhlo<comparison_direction_v1 NE>
-@@ -31,6 +33,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_direction_ge"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_direction_ge(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     // CHECK: comparison_direction = #vhlo<comparison_direction_v1 GE>
-@@ -40,6 +43,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_direction_gt"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_direction_gt(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     // CHECK: comparison_direction = #vhlo<comparison_direction_v1 GT>
-@@ -49,6 +53,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_direction_le"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_direction_le(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     // CHECK: comparison_direction = #vhlo<comparison_direction_v1 LE>
-@@ -58,6 +63,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_direction_lt"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_direction_lt(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     // CHECK: comparison_direction = #vhlo<comparison_direction_v1 LT>
-@@ -67,6 +73,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_type_notype"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_type_notype(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     comparison_direction = #stablehlo<comparison_direction EQ>
-@@ -76,6 +83,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_type_float"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_type_float(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     comparison_direction = #stablehlo<comparison_direction EQ>,
-@@ -86,6 +94,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_type_totalorder"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_type_totalorder(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     comparison_direction = #stablehlo<comparison_direction EQ>,
-@@ -96,6 +105,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_type_signed"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_type_signed(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     comparison_direction = #stablehlo<comparison_direction EQ>,
-@@ -106,6 +116,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_type_unsigned"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_type_unsigned(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     comparison_direction = #stablehlo<comparison_direction EQ>,
-@@ -118,6 +129,7 @@
- // ConvDimensionNumbers aka #stablehlo.conv is covered below.
- 
- // CHECK-LABEL: "attr_custom_call_api_version_unspecified"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_custom_call_api_version_unspecified(%arg0: tensor<f32>) -> tensor<f32> {
-   %0 = "stablehlo.custom_call"(%arg0) {
-     call_target_name = "foo",
-@@ -128,6 +140,7 @@
- }
- 
- // CHECK-LABEL: "attr_custom_call_api_version_original"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_custom_call_api_version_original(%arg0: tensor<f32>) -> tensor<f32> {
-   %0 = "stablehlo.custom_call"(%arg0) {
-     call_target_name = "foo",
-@@ -138,6 +151,7 @@
- }
- 
- // CHECK-LABEL: "attr_custom_call_api_version_status_returning"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_custom_call_api_version_status_returning(%arg0: tensor<f32>) -> tensor<f32> {
-   %0 = "stablehlo.custom_call"(%arg0) {
-     call_target_name = "foo",
-@@ -148,6 +162,7 @@
- }
- 
- // CHECK-LABEL: "attr_custom_call_api_version_status_returning_unified"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_custom_call_api_version_status_returning_unified(%arg0: tensor<f32>) -> tensor<f32> {
-   %0 = "stablehlo.custom_call"(%arg0) {
-     call_target_name = "foo",
-@@ -160,6 +175,7 @@
- // DotDimensionNumbers aka #stablehlo.dot is covered below.
- 
- // CHECK-LABEL: "attr_fft_type_fft"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_fft_type_fft(%arg0: tensor<16xcomplex<f32>>) -> tensor<16xcomplex<f32>> {
-   %0 = "stablehlo.fft"(%arg0) {
-     // CHECK: fft_type = #vhlo<fft_type_v1 FFT>
-@@ -170,6 +186,7 @@
- }
- 
- // CHECK-LABEL: "attr_fft_type_ifft"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_fft_type_ifft(%arg0: tensor<16xcomplex<f32>>) -> tensor<16xcomplex<f32>> {
-   %0 = "stablehlo.fft"(%arg0) {
-     // CHECK: fft_type = #vhlo<fft_type_v1 IFFT>
-@@ -180,6 +197,7 @@
- }
- 
- // CHECK-LABEL: "attr_fft_type_rfft"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_fft_type_rfft(%arg0: tensor<16xf32>) -> tensor<9xcomplex<f32>> {
-   %0 = "stablehlo.fft"(%arg0) {
-     // CHECK: fft_type = #vhlo<fft_type_v1 RFFT>
-@@ -190,6 +208,7 @@
- }
- 
- // CHECK-LABEL: "attr_fft_type_irfft"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_fft_type_irfft(%arg0: tensor<9xcomplex<f32>>) -> tensor<16xf32> {
-   %0 = "stablehlo.fft"(%arg0) {
-     // CHECK: fft_type = #vhlo<fft_type_v1 IRFFT>
-@@ -202,6 +221,7 @@
- // GatherDimensionNumbers aka #stablehlo.gather is covered below.
- 
- // CHECK-LABEL: "attr_precision_config_default"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_precision_config_default(%arg0: tensor<8x16xf32>, %arg1: tensor<16x8xf32>) -> tensor<8x8xf32> {
-   %0 = "stablehlo.dot"(%arg0, %arg1) {
-     // CHECK: precision_config = #vhlo.array_v1<[#vhlo<precision_v1 DEFAULT>, #vhlo<precision_v1 DEFAULT>]>
-@@ -210,6 +230,7 @@
- }
- 
- // CHECK-LABEL: "attr_precision_config_high"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_precision_config_high(%arg0: tensor<8x16xf32>, %arg1: tensor<16x8xf32>) -> tensor<8x8xf32> {
-   %0 = "stablehlo.dot"(%arg0, %arg1) {
-     // CHECK: precision_config = #vhlo.array_v1<[#vhlo<precision_v1 HIGH>, #vhlo<precision_v1 HIGH>]>
-@@ -219,6 +240,7 @@
- }
- 
- // CHECK-LABEL: "attr_precision_config_highest"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_precision_config_highest(%arg0: tensor<8x16xf32>, %arg1: tensor<16x8xf32>) -> tensor<8x8xf32> {
-   %0 = "stablehlo.dot"(%arg0, %arg1) {
-     // CHECK: precision_config = #vhlo.array_v1<[#vhlo<precision_v1 HIGHEST>, #vhlo<precision_v1 HIGHEST>]>
-@@ -228,6 +250,7 @@
- }
- 
- // CHECK-LABEL: "attr_rng_algorithm_default"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_rng_algorithm_default(%arg0: tensor<f32>) -> (tensor<f32>, tensor<f32>) {
-   %0:2 = "stablehlo.rng_bit_generator"(%arg0) {
-     // CHECK: rng_algorithm = #vhlo<rng_algorithm_v1 DEFAULT>
-@@ -237,6 +260,7 @@
- }
- 
- // CHECK-LABEL: "attr_rng_algorithm_three_fry"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_rng_algorithm_three_fry(%arg0: tensor<f32>) -> (tensor<f32>, tensor<f32>) {
-   %0:2 = "stablehlo.rng_bit_generator"(%arg0) {
-     // CHECK: rng_algorithm = #vhlo<rng_algorithm_v1 THREE_FRY>
-@@ -246,6 +270,7 @@
- }
- 
- // CHECK-LABEL: "attr_rng_algorithm_philox"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_rng_algorithm_philox(%arg0: tensor<f32>) -> (tensor<f32>, tensor<f32>) {
-   %0:2 = "stablehlo.rng_bit_generator"(%arg0) {
-     // CHECK: rng_algorithm = #vhlo<rng_algorithm_v1 PHILOX>
-@@ -255,6 +280,7 @@
- }
- 
- // CHECK-LABEL: "attr_rng_distribution_uniform"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @attr_rng_distribution_uniform(%arg0: tensor<f32>, %arg1: tensor<f32>, %arg2: tensor<0xindex>) -> tensor<f32> {
-   %0 = "stablehlo.rng"(%arg0, %arg1, %arg2) {
-     // CHECK: rng_distribution = #vhlo<rng_distribution_v1 UNIFORM>
-@@ -264,6 +290,7 @@
- }
- 
- // CHECK-LABEL: "attr_rng_distribution_normal"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @attr_rng_distribution_normal(%arg0: tensor<f32>, %arg1: tensor<f32>, %arg2: tensor<0xindex>) -> tensor<f32> {
-   %0 = "stablehlo.rng"(%arg0, %arg1, %arg2) {
-     // CHECK: rng_distribution = #vhlo<rng_distribution_v1 NORMAL>
-@@ -275,6 +302,7 @@
- // ScatterDimensionNumbers aka #stablehlo.scatter is covered below.
- 
- // CHECK-LABEL: "attr_transpose_no_transpose"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_transpose_no_transpose(%arg0: tensor<16x16xf32>, %arg1: tensor<16x16xf32>) ->  tensor<16x16xf32> {
-   %0 = "stablehlo.triangular_solve"(%arg0, %arg1) {
-     left_side = true,
-@@ -287,6 +315,7 @@
- }
- 
- // CHECK-LABEL: "attr_transpose_transpose"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_transpose_transpose(%arg0: tensor<16x16xf32>, %arg1: tensor<16x16xf32>) ->  tensor<16x16xf32> {
-   %0 = "stablehlo.triangular_solve"(%arg0, %arg1) {
-     left_side = true,
-@@ -299,6 +328,7 @@
- }
- 
- // CHECK-LABEL: "attr_transpose_adjoint"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_transpose_adjoint(%arg0: tensor<16x16xf32>, %arg1: tensor<16x16xf32>) ->  tensor<16x16xf32> {
-   %0 = "stablehlo.triangular_solve"(%arg0, %arg1) {
-     left_side = true,
-@@ -313,10 +343,9 @@
- // TypeExtensionsAttr aka #stablehlo.type_extensions is covered below.
- 
- // CHECK-LABEL: "attr_type_extensions_bounds"
--func.func @attr_type_extensions_bounds(
--    %arg0: tensor<?x?xf32, #stablehlo.type_extensions<bounds = [16, ?]>>)
--    -> tensor<?x?xf32, #stablehlo.type_extensions<bounds = [16, ?]>> {
--  // CHECK: "vhlo.return_v1"(%arg0) : (!vhlo.tensor_v1<?x?x!vhlo.f32_v1, #vhlo.type_extensions_v1<bounds = [16, ?]>>) -> ()
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @attr_type_extensions_bounds(%arg0: tensor<?x?xf32, #stablehlo.type_extensions<bounds = [16, ?]>>) -> tensor<?x?xf32, #stablehlo.type_extensions<bounds = [16, ?]>> {
-+  // CHECK: "vhlo.return_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<?x?x!vhlo.f32_v1, #vhlo.type_extensions_v1<bounds = [16, ?]>>) -> ()
-   func.return %arg0 : tensor<?x?xf32, #stablehlo.type_extensions<bounds = [16, ?]>>
- }
- 
-@@ -324,8 +353,9 @@
- // ============ DEFAULTS ============
- 
- // CHECK-LABEL: "default_all_gather"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_all_gather(%arg0: tensor<16x8xf32>) -> tensor<16x16xf32> {
--  //               CHECK: "vhlo.all_gather_v1"(%arg0) <{
-+  //               CHECK: "vhlo.all_gather_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   all_gather_dim = #vhlo.integer_v1<1 : i64>
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0], [1]]> : tensor<2x1xi64>>,
-@@ -339,8 +369,9 @@
- }
- 
- // CHECK-LABEL: "default_all_reduce"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_all_reduce(%arg0: tensor<f32>) -> tensor<f32> {
--  //               CHECK: "vhlo.all_reduce_v1"(%arg0)
-+  //               CHECK: "vhlo.all_reduce_v1"(%[[ARG0]])
-   //          CHECK-SAME: <{
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0], [1]]> : tensor<2x1xi64>>,
-@@ -362,8 +393,9 @@
- }
- 
- // CHECK-LABEL: "default_all_to_all"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_all_to_all(%arg0: tensor<4x16xf32>) -> tensor<16x4xf32> {
--  //               CHECK: "vhlo.all_to_all_v1"(%arg0) <{
-+  //               CHECK: "vhlo.all_to_all_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   //          CHECK-SAME:   concat_dimension = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0, 1, 2, 3]]> : tensor<1x4xi64>>,
-@@ -380,8 +412,9 @@
- }
- 
- // CHECK-LABEL: "default_cholesky"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_cholesky(%arg0: tensor<1x16x16xf32>) -> tensor<1x16x16xf32> {
--  //      CHECK: "vhlo.cholesky_v1"(%arg0) <{
-+  //      CHECK: "vhlo.cholesky_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   lower = #vhlo.bool_v1<false>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<1x16x16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<1x16x16x!vhlo.f32_v1>
-   %0 = "stablehlo.cholesky"(%arg0) : (tensor<1x16x16xf32>) -> tensor<1x16x16xf32>
-@@ -389,8 +422,9 @@
- }
- 
- // CHECK-LABEL: "default_collective_permute"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_collective_permute(%arg0: tensor<16x8xf32>) -> tensor<16x8xf32> {
--  //               CHECK: "vhlo.collective_permute_v1"(%arg0) <{
-+  //               CHECK: "vhlo.collective_permute_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME{LITERAL}:   source_target_pairs = #vhlo.tensor_v1<dense<[[0, 1], [1, 2], [2, 3]]> : tensor<3x2xi64>>
-   //          CHECK-SAME: }> : (!vhlo.tensor_v1<16x8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x8x!vhlo.f32_v1>
-@@ -401,8 +435,9 @@
- }
- 
- // CHECK-LABEL: "default_compare"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @default_compare(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
--  //      CHECK: "vhlo.compare_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.compare_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   compare_type = #vhlo<comparison_type_v1 NOTYPE>,
-   // CHECK-SAME:   comparison_direction = #vhlo<comparison_direction_v1 EQ>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-@@ -413,8 +448,9 @@
- }
- 
- // CHECK-LABEL: "default_convolution"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @default_convolution(%arg0: tensor<1x8x8x207xf32>, %arg1: tensor<3x3x207x16xf32>) -> tensor<1x6x6x16xf32> {
--  //      CHECK: "vhlo.convolution_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.convolution_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   batch_group_count = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   feature_group_count = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   input_batch_dimension = #vhlo.integer_v1<0 : i64>,
-@@ -442,8 +478,9 @@
- }
- 
- // CHECK-LABEL: "default_custom_call"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_custom_call(%arg0: tensor<f32>) -> tensor<f32> {
--  //      CHECK: "vhlo.custom_call_v1"(%arg0) <{
-+  //      CHECK: "vhlo.custom_call_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   api_version = #vhlo<api_version_v1 API_VERSION_ORIGINAL>,
-   // CHECK-SAME:   backend_config = #vhlo.string_v1<"">,
-   // CHECK-SAME:   call_target_name = #vhlo.string_v1<"foo">,
-@@ -460,8 +497,9 @@
- }
- 
- // CHECK-LABEL: "default_dot_general"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @default_dot_general(%arg0: tensor<8x8x16xf32>, %arg1: tensor<8x16x8xf32>) -> tensor<8x8x8xf32> {
--  //      CHECK: "vhlo.dot_general_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.dot_general_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   lhs_batching_dimensions = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>,
-   // CHECK-SAME:   lhs_contracting_dimensions = #vhlo.tensor_v1<dense<2> : tensor<1xi64>>,
-   // CHECK-SAME:   precision_config = #vhlo.array_v1<[#vhlo<precision_v1 DEFAULT>, #vhlo<precision_v1 DEFAULT>]>,
-@@ -480,8 +518,9 @@
- }
- 
- // CHECK-LABEL: "default_dot"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @default_dot(%arg0: tensor<8x16xf32>, %arg1: tensor<16x8xf32>) -> tensor<8x8xf32> {
--  //      CHECK: "vhlo.dot_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.dot_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   precision_config = #vhlo.array_v1<[#vhlo<precision_v1 DEFAULT>, #vhlo<precision_v1 DEFAULT>]>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<8x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<8x8x!vhlo.f32_v1>
-   %0 = "stablehlo.dot"(%arg0, %arg1) : (tensor<8x16xf32>, tensor<16x8xf32>) -> tensor<8x8xf32>
-@@ -489,8 +528,9 @@
- }
- 
- // CHECK-LABEL: "default_dynamic_broadcast_in_dim"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @default_dynamic_broadcast_in_dim(%arg0: tensor<?x?xf32>, %arg1: tensor<2xindex>) -> tensor<?x?xf32> {
--  //      CHECK: "vhlo.dynamic_broadcast_in_dim_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.dynamic_broadcast_in_dim_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   broadcast_dimensions = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-   // CHECK-SAME:   known_expanding_dimensions = #vhlo.tensor_v1<dense<> : tensor<0xi64>>,
-   // CHECK-SAME:   known_nonexpanding_dimensions = #vhlo.tensor_v1<dense<> : tensor<0xi64>>
-@@ -502,8 +542,9 @@
- }
- 
- // CHECK-LABEL: "default_dynamic_conv"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @default_dynamic_conv(%arg0: tensor<1x8x8x207xf32>, %arg1: tensor<3x3x207x16xf32>, %arg2: tensor<4xi32>) -> tensor<1x?x?x16xf32> {
--  //      CHECK: "vhlo.dynamic_conv_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.dynamic_conv_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   batch_group_count = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   feature_group_count = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   input_batch_dimension = #vhlo.integer_v1<0 : i64>,
-@@ -531,8 +572,9 @@
- }
- 
- // CHECK-LABEL: "default_dynamic_gather"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @default_dynamic_gather(%arg0 : tensor<2x4x9xf32>, %arg1 : tensor<1x5x2xi32>, %arg2 : tensor<3xi32>) -> tensor<1x5x8xf32> {
--  //      CHECK: "vhlo.dynamic_gather_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.dynamic_gather_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   collapsed_slice_dims = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-   // CHECK-SAME:   index_vector_dim = #vhlo.integer_v1<2 : i64>,
-   // CHECK-SAME:   indices_are_sorted = #vhlo.bool_v1<false>,
-@@ -558,15 +600,16 @@
-   // CHECK-SAME:   sym_name = #vhlo.string_v1<"default_func">,
-   // CHECK-SAME:   sym_visibility = #vhlo.string_v1<"">
-   // CHECK-SAME: }> ({
--  // CHECK-NEXT:   ^[[BB:bb.*]](%arg0: !vhlo.tensor_v1<!vhlo.f32_v1>):
--  // CHECK-NEXT:     "vhlo.return_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-+  // CHECK-NEXT:   ^[[BB:bb.*]](%[[ARG0:.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
-+  // CHECK-NEXT:     "vhlo.return_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-   // CHECK-NEXT: }) : () -> ()
-   func.return %arg0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "dynamic_gather"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @dynamic_gather(%arg0 : tensor<2x4x9xf32>, %arg1 : tensor<1x5x2xi32>) -> tensor<1x5x1xf32> {
--  //      CHECK: "vhlo.gather_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.gather_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   collapsed_slice_dims = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-   // CHECK-SAME:   index_vector_dim = #vhlo.integer_v1<2 : i64>,
-   // CHECK-SAME:   indices_are_sorted = #vhlo.bool_v1<false>,
-@@ -587,8 +630,9 @@
- }
- 
- // CHECK-LABEL: "default_infeed"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_infeed(%arg0: !stablehlo.token) -> (tensor<f32>, !stablehlo.token) {
--  //               CHECK: "vhlo.infeed_v1"(%arg0) <{
-+  //               CHECK: "vhlo.infeed_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   infeed_config = #vhlo.string_v1<"">,
-   // CHECK-SAME{LITERAL}:   layout = #vhlo.array_v1<[]>
-   //          CHECK-SAME: }> : (!vhlo.token_v1) -> (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.token_v1)
-@@ -597,8 +641,9 @@
- }
- 
- // CHECK-LABEL: "default_outfeed"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @default_outfeed(%arg0: tensor<f32>, %arg1: !stablehlo.token) -> !stablehlo.token {
--  //      CHECK: "vhlo.outfeed_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.outfeed_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   outfeed_config = #vhlo.string_v1<"">
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.token_v1) -> !vhlo.token_v1
-   %0 = "stablehlo.outfeed"(%arg0, %arg1) : (tensor<f32>, !stablehlo.token) -> !stablehlo.token
-@@ -606,8 +651,9 @@
- }
- 
- // CHECK-LABEL: "default_recv"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_recv(%arg0: !stablehlo.token) -> (tensor<f32>, !stablehlo.token) {
--  //      CHECK: "vhlo.recv_v1"(%arg0) <{
-+  //      CHECK: "vhlo.recv_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME:   channel_type = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   is_host_transfer = #vhlo.bool_v1<false>
-@@ -619,8 +665,9 @@
- }
- 
- // CHECK-LABEL: "default_send"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @default_send(%arg0: tensor<f32>, %arg1: !stablehlo.token) -> !stablehlo.token {
--  //      CHECK: "vhlo.send_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.send_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME:   channel_type = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   is_host_transfer = #vhlo.bool_v1<false>
-@@ -632,8 +679,9 @@
- }
- 
- // CHECK-LABEL: "default_reduce_scatter"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_reduce_scatter(%arg0: tensor<16xf32>) -> tensor<16xf32> {
--  //               CHECK: "vhlo.reduce_scatter_v1"(%arg0) <{
-+  //               CHECK: "vhlo.reduce_scatter_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0], [1]]> : tensor<2x1xi64>>,
-   //          CHECK-SAME:   scatter_dimension = #vhlo.integer_v1<0 : i64>
-@@ -655,8 +703,9 @@
- }
- 
- // CHECK-LABEL: "default_reduce_window"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @default_reduce_window(%arg0: tensor<2x17x31x7xf32>, %arg1: tensor<f32>) -> tensor<2x16x30x7xf32> {
--  //               CHECK: "vhlo.reduce_window_v1"(%arg0, %arg1)  <{
-+  //               CHECK: "vhlo.reduce_window_v1"(%[[ARG0]], %[[ARG1]])  <{
-   //          CHECK-SAME:   base_dilations = #vhlo.tensor_v1<dense<1> : tensor<4xi64>>,
-   // CHECK-SAME{LITERAL}:   padding = #vhlo.tensor_v1<dense<0> : tensor<4x2xi64>>,
-   //          CHECK-SAME:   window_dilations = #vhlo.tensor_v1<dense<1> : tensor<4xi64>>,
-@@ -678,8 +727,9 @@
- }
- 
- // CHECK-LABEL: "default_scatter"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @default_scatter(%arg0: tensor<200x100x300xf32>, %arg1: tensor<10x2xi32>, %arg2: tensor<10x300xf32>) -> tensor<200x100x300xf32> {
--  //      CHECK: "vhlo.scatter_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.scatter_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   index_vector_dim = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   indices_are_sorted = #vhlo.bool_v1<false>,
-   // CHECK-SAME:   inserted_window_dims = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-@@ -707,8 +757,9 @@
- }
- 
- // CHECK-LABEL: "default_select_and_scatter"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @default_select_and_scatter(%arg0: tensor<10x24x24x64xf32>, %arg1: tensor<10x23x23x64xf32>, %arg2: tensor<f32>) -> tensor<10x24x24x64xf32> {
--  //      CHECK: "vhlo.select_and_scatter_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.select_and_scatter_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   padding = #vhlo.tensor_v1<dense<0> : tensor<4x2xi64>>,
-   // CHECK-SAME:   window_dimensions = #vhlo.tensor_v1<dense<[1, 2, 2, 1]> : tensor<4xi64>>,
-   // CHECK-SAME:   window_strides = #vhlo.tensor_v1<dense<1> : tensor<4xi64>>
-@@ -736,8 +787,9 @@
- }
- 
- // CHECK-LABEL: "default_sort"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_sort(%arg0: tensor<16xf32>) -> tensor<16xf32> {
--  //      CHECK: "vhlo.sort_v1"(%arg0) <{
-+  //      CHECK: "vhlo.sort_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   dimension = #vhlo.integer_v1<-1 : i64>
-   // CHECK-SAME:   is_stable = #vhlo.bool_v1<false>
-   // CHECK-SAME: }> ({
-@@ -756,29 +808,33 @@
- // ============ OPS ============
- 
- // CHECK-LABEL: "op_abs"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_abs(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.abs_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.abs_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.abs"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_add"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_add(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_after_all"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_after_all(%arg0: !stablehlo.token) -> !stablehlo.token {
--  // CHECK: "vhlo.after_all_v1"(%arg0) : (!vhlo.token_v1) -> !vhlo.token_v1
-+  // CHECK: "vhlo.after_all_v1"(%[[ARG0]]) : (!vhlo.token_v1) -> !vhlo.token_v1
-   %0 = "stablehlo.after_all"(%arg0) : (!stablehlo.token) -> !stablehlo.token
-   func.return %0 : !stablehlo.token
- }
- 
- // CHECK-LABEL: "op_all_gather"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_all_gather(%arg0: tensor<16x8xf32>) -> tensor<16x16xf32> {
--  //               CHECK: "vhlo.all_gather_v1"(%arg0) <{
-+  //               CHECK: "vhlo.all_gather_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   all_gather_dim = #vhlo.integer_v1<1 : i64>
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0], [1]]> : tensor<2x1xi64>>,
-@@ -794,8 +850,9 @@
- }
- 
- // CHECK-LABEL: "op_all_reduce"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_all_reduce(%arg0: tensor<f32>) -> tensor<f32> {
--  //               CHECK: "vhlo.all_reduce_v1"(%arg0) <{
-+  //               CHECK: "vhlo.all_reduce_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0], [1]]> : tensor<2x1xi64>>,
-   //          CHECK-SAME:   use_global_device_ids = #vhlo.bool_v1<true>
-@@ -817,8 +874,9 @@
- }
- 
- // CHECK-LABEL: "op_all_to_all"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_all_to_all(%arg0: tensor<4x16xf32>) -> tensor<16x4xf32> {
--  //               CHECK: "vhlo.all_to_all_v1"(%arg0) <{
-+  //               CHECK: "vhlo.all_to_all_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<1 : i64>,
-   //          CHECK-SAME:   concat_dimension = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0, 1, 2, 3]]> : tensor<1x4xi64>>,
-@@ -836,22 +894,25 @@
- }
- 
- // CHECK-LABEL: "op_and"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_and(%arg0: tensor<i1>, %arg1: tensor<i1>) -> tensor<i1> {
--  // CHECK: "vhlo.and_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-+  // CHECK: "vhlo.and_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-   %0 = "stablehlo.and"(%arg0, %arg1) : (tensor<i1>, tensor<i1>) -> tensor<i1>
-   func.return %0 : tensor<i1>
- }
- 
- // CHECK-LABEL: "op_atan2"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_atan2(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.atan2_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.atan2_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.atan2"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_batch_norm_grad"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}}, %[[ARG3:.*]]: {{.*}}, %[[ARG4:.*]]: {{.*}})
- func.func @op_batch_norm_grad(%arg0: tensor<16x16x16x16xf32>, %arg1: tensor<16xf32>, %arg2: tensor<16xf32>, %arg3: tensor<16xf32>, %arg4: tensor<16x16x16x16xf32>) -> (tensor<16x16x16x16xf32>, tensor<16xf32>, tensor<16xf32>) {
--  //      CHECK: "vhlo.batch_norm_grad_v1"(%arg0, %arg1, %arg2, %arg3, %arg4) <{
-+  //      CHECK: "vhlo.batch_norm_grad_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]], %[[ARG3]], %[[ARG4]]) <{
-   // CHECK-SAME:   epsilon = #vhlo.float_v1<1.000000e-03 : !vhlo.f32_v1>,
-   // CHECK-SAME:   feature_index = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<16x16x16x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x16x16x16x!vhlo.f32_v1>) -> (!vhlo.tensor_v1<16x16x16x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>)
-@@ -863,8 +924,9 @@
- }
- 
- // CHECK-LABEL: "op_batch_norm_inference"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}}, %[[ARG3:.*]]: {{.*}}, %[[ARG4:.*]]: {{.*}})
- func.func @op_batch_norm_inference(%arg0: tensor<16x16x16x16xf32>, %arg1: tensor<16xf32>, %arg2: tensor<16xf32>, %arg3: tensor<16xf32>, %arg4: tensor<16xf32>) -> tensor<16x16x16x16xf32> {
--  //      CHECK: "vhlo.batch_norm_inference_v1"(%arg0, %arg1, %arg2, %arg3, %arg4) <{
-+  //      CHECK: "vhlo.batch_norm_inference_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]], %[[ARG3]], %[[ARG4]]) <{
-   // CHECK-SAME:   epsilon = #vhlo.float_v1<1.000000e-03 : !vhlo.f32_v1>,
-   // CHECK-SAME:   feature_index = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<16x16x16x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x16x16x16x!vhlo.f32_v1>
-@@ -876,8 +938,9 @@
- }
- 
- // CHECK-LABEL: "op_batch_norm_training"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_batch_norm_training(%arg0: tensor<16x16x16x16xf32>, %arg1: tensor<16xf32>, %arg2: tensor<16xf32>) -> (tensor<16x16x16x16xf32>, tensor<16xf32>, tensor<16xf32>) {
--  //      CHECK: "vhlo.batch_norm_training_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.batch_norm_training_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   epsilon = #vhlo.float_v1<1.000000e-03 : !vhlo.f32_v1>,
-   // CHECK-SAME:   feature_index = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<16x16x16x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>) -> (!vhlo.tensor_v1<16x16x16x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>)
-@@ -889,15 +952,17 @@
- }
- 
- // CHECK-LABEL: "op_bitcast_convert"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_bitcast_convert(%arg0: tensor<i32>) -> tensor<f32> {
--  // CHECK: "vhlo.bitcast_convert_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.bitcast_convert_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.bitcast_convert"(%arg0) : (tensor<i32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_broadcast_in_dim"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_broadcast_in_dim(%arg0: tensor<16xf32>) -> tensor<16x16xf32> {
--  //      CHECK: "vhlo.broadcast_in_dim_v1"(%arg0) <{
-+  //      CHECK: "vhlo.broadcast_in_dim_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   broadcast_dimensions = #vhlo.tensor_v1<dense<1> : tensor<1xi64>>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x16x!vhlo.f32_v1>
-   %0 = "stablehlo.broadcast_in_dim"(%arg0) {
-@@ -907,8 +972,9 @@
- }
- 
- // CHECK-LABEL: "op_broadcast"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_broadcast(%arg0: tensor<16xf32>) -> tensor<16x16xf32> {
--  //      CHECK: "vhlo.broadcast_v1"(%arg0) <{
-+  //      CHECK: "vhlo.broadcast_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   broadcast_sizes = #vhlo.tensor_v1<dense<16> : tensor<1xi64>>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x16x!vhlo.f32_v1>
-   %0 = "stablehlo.broadcast"(%arg0) {
-@@ -918,9 +984,10 @@
- }
- 
- // CHECK-LABEL: "op_case"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_case(%arg0: tensor<i32>, %arg1: tensor<f32>) -> tensor<f32> {
--  //      CHECK: "vhlo.case_v1"(%arg0) ({
--  // CHECK-NEXT:   "vhlo.return_v1"(%arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-+  //      CHECK: "vhlo.case_v1"(%[[ARG0]]) ({
-+  // CHECK-NEXT:   "vhlo.return_v1"(%[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-   // CHECK-NEXT: }) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.case"(%arg0) ({
-     "stablehlo.return"(%arg1) : (tensor<f32>) -> ()
-@@ -929,22 +996,25 @@
- }
- 
- // CHECK-LABEL: "op_cbrt"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_cbrt(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.cbrt_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.cbrt_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.cbrt"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_ceil"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_ceil(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.ceil_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.ceil_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.ceil"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_cholesky"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_cholesky(%arg0: tensor<1x16x16xf32>) -> tensor<1x16x16xf32> {
--  //      CHECK: "vhlo.cholesky_v1"(%arg0) <{
-+  //      CHECK: "vhlo.cholesky_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   lower = #vhlo.bool_v1<true>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<1x16x16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<1x16x16x!vhlo.f32_v1>
-   %0 = "stablehlo.cholesky"(%arg0) {
-@@ -954,22 +1024,25 @@
- }
- 
- // CHECK-LABEL: "op_clamp"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_clamp(%arg0: tensor<f32>, %arg1: tensor<f32>, %arg2: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.clamp_v1"(%arg0, %arg1, %arg2) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.clamp_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.clamp"(%arg0, %arg1, %arg2) : (tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_count_leading_zeros"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_count_leading_zeros(%arg0: tensor<i32>) -> tensor<i32> {
--  // CHECK: "vhlo.count_leading_zeros_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-+  // CHECK: "vhlo.count_leading_zeros_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-   %0 = "stablehlo.count_leading_zeros"(%arg0) : (tensor<i32>) -> tensor<i32>
-   func.return %0 : tensor<i32>
- }
- 
- // CHECK-LABEL: "op_collective_permute"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_collective_permute(%arg0: tensor<16x8xf32>) -> tensor<16x8xf32> {
--  //               CHECK: "vhlo.collective_permute_v1"(%arg0) <{
-+  //               CHECK: "vhlo.collective_permute_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME{LITERAL}:   source_target_pairs = #vhlo.tensor_v1<dense<[[0, 1], [1, 2], [2, 3]]> : tensor<3x2xi64>>
-   //          CHECK-SAME: }> : (!vhlo.tensor_v1<16x8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x8x!vhlo.f32_v1>
-@@ -981,8 +1054,9 @@
- }
- 
- // CHECK-LABEL: "op_compare"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_compare(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
--  //      CHECK: "vhlo.compare_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.compare_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   compare_type = #vhlo<comparison_type_v1 TOTALORDER>,
-   // CHECK-SAME:   comparison_direction = #vhlo<comparison_direction_v1 EQ>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-@@ -994,15 +1068,17 @@
- }
- 
- // CHECK-LABEL: "op_complex"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_complex(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<complex<f32>> {
--  // CHECK: "vhlo.complex_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>
-+  // CHECK: "vhlo.complex_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>
-   %0 = "stablehlo.complex"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<complex<f32>>
-   func.return %0 : tensor<complex<f32>>
- }
- 
- // CHECK-LABEL: "op_concatenate"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_concatenate(%arg0: tensor<8xf32>, %arg1: tensor<8xf32>) -> tensor<16xf32> {
--  //      CHECK: "vhlo.concatenate_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.concatenate_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   dimension = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<8x!vhlo.f32_v1>, !vhlo.tensor_v1<8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x!vhlo.f32_v1>
-   %0 = "stablehlo.concatenate"(%arg0, %arg1) {
-@@ -1012,6 +1088,7 @@
- }
- 
- // CHECK-LABEL: "op_constant"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_constant(%arg0: tensor<f32>) -> tensor<f32> {
-   //      CHECK: "vhlo.constant_v1"() <{
-   // CHECK-SAME:   value = #vhlo.tensor_v1<dense<0.000000e+00> : tensor<f32>>
-@@ -1023,15 +1100,17 @@
- }
- 
- // CHECK-LABEL: "op_convert"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_convert(%arg0: tensor<i32>) -> tensor<f32> {
--  // CHECK: "vhlo.convert_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.convert_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.convert"(%arg0) : (tensor<i32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_convolution"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_convolution(%arg0: tensor<1x8x8x207xf32>, %arg1: tensor<3x3x207x16xf32>) -> tensor<1x7x7x16xf32> {
--  //      CHECK: "vhlo.convolution_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.convolution_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   batch_group_count = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   feature_group_count = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   input_batch_dimension = #vhlo.integer_v1<0 : i64>,
-@@ -1065,8 +1144,9 @@
- }
- 
- // CHECK-LABEL: "op_cosine"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_cosine(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.cosine_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.cosine_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.cosine"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
-@@ -1079,8 +1159,9 @@
- }
- 
- // CHECK-LABEL: "op_cross_replica_sum"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_cross_replica_sum(%arg0: tensor<f32>) -> tensor<f32> {
--  //               CHECK: "vhlo.cross-replica-sum_v1"(%arg0) <{
-+  //               CHECK: "vhlo.cross-replica-sum_v1"(%[[ARG0]]) <{
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0], [1]]> : tensor<2x1xi64>>
-   //          CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.cross-replica-sum"(%arg0) {
-@@ -1090,8 +1171,9 @@
- }
- 
- // CHECK-LABEL: "op_custom_call"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_custom_call(%arg0: tensor<f32>) -> tensor<f32> {
--  //      CHECK: "vhlo.custom_call_v1"(%arg0) <{
-+  //      CHECK: "vhlo.custom_call_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   api_version = #vhlo<api_version_v1 API_VERSION_STATUS_RETURNING>,
-   // CHECK-SAME:   backend_config = #vhlo.string_v1<"\08\03\1A\02">,
-   // CHECK-SAME:   call_target_name = #vhlo.string_v1<"foo">,
-@@ -1122,15 +1204,17 @@
- }
- 
- // CHECK-LABEL: "op_divide"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_divide(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.divide_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.divide_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.divide"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_dot_general"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_dot_general(%arg0: tensor<8x8x16xf32>, %arg1: tensor<8x16x8xf32>) -> tensor<8x8x8xf32> {
--  //      CHECK: "vhlo.dot_general_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.dot_general_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   lhs_batching_dimensions = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>,
-   // CHECK-SAME:   lhs_contracting_dimensions = #vhlo.tensor_v1<dense<2> : tensor<1xi64>>,
-   // CHECK-SAME:   precision_config = #vhlo.array_v1<[#vhlo<precision_v1 HIGHEST>, #vhlo<precision_v1 HIGHEST>]>,
-@@ -1150,8 +1234,9 @@
- }
- 
- // CHECK-LABEL: "op_dot"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_dot(%arg0: tensor<8x16xf32>, %arg1: tensor<16x8xf32>) -> tensor<8x8xf32> {
--  //      CHECK: "vhlo.dot_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.dot_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   precision_config = #vhlo.array_v1<[#vhlo<precision_v1 HIGHEST>, #vhlo<precision_v1 HIGHEST>]>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<8x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<8x8x!vhlo.f32_v1>
-   %0 = "stablehlo.dot"(%arg0, %arg1) {
-@@ -1161,8 +1246,9 @@
- }
- 
- // CHECK-LABEL: "op_dynamic_broadcast_in_dim"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_dynamic_broadcast_in_dim(%arg0: tensor<?x?xf32>, %arg1: tensor<2xindex>) -> tensor<?x?xf32> {
--  //      CHECK: "vhlo.dynamic_broadcast_in_dim_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.dynamic_broadcast_in_dim_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   broadcast_dimensions = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-   // CHECK-SAME:   known_expanding_dimensions = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>,
-   // CHECK-SAME:   known_nonexpanding_dimensions = #vhlo.tensor_v1<dense<1> : tensor<1xi64>>
-@@ -1176,8 +1262,9 @@
- }
- 
- // CHECK-LABEL: "op_dynamic_conv"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_dynamic_conv(%arg0: tensor<1x8x8x207xf32>, %arg1: tensor<3x3x207x16xf32>, %arg2: tensor<4xi32>) -> tensor<1x?x?x16xf32> {
--  //      CHECK: "vhlo.dynamic_conv_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.dynamic_conv_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   batch_group_count = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   feature_group_count = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   input_batch_dimension = #vhlo.integer_v1<0 : i64>,
-@@ -1211,8 +1298,9 @@
- }
- 
- // CHECK-LABEL: "op_dynamic_gather"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_dynamic_gather(%arg0 : tensor<2x4x9xf32>, %arg1 : tensor<1x5x2xi32>, %arg2 : tensor<3xi32>) -> tensor<1x5x8xf32> {
--  //      CHECK: "vhlo.dynamic_gather_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.dynamic_gather_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   collapsed_slice_dims = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-   // CHECK-SAME:   index_vector_dim = #vhlo.integer_v1<2 : i64>,
-   // CHECK-SAME:   indices_are_sorted = #vhlo.bool_v1<true>,
-@@ -1232,8 +1320,9 @@
- }
- 
- // CHECK-LABEL: "op_dynamic_iota"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_dynamic_iota(%arg0: tensor<1xindex>) -> tensor<?xf32> {
--  //      CHECK: "vhlo.dynamic_iota_v1"(%arg0) <{
-+  //      CHECK: "vhlo.dynamic_iota_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   iota_dimension = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<1x!vhlo.index_v1>) -> !vhlo.tensor_v1<?x!vhlo.f32_v1>
-   %0 = "stablehlo.dynamic_iota"(%arg0) {
-@@ -1243,22 +1332,25 @@
- }
- 
- // CHECK-LABEL: "op_dynamic_pad"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}}, %[[ARG3:.*]]: {{.*}}, %[[ARG4:.*]]: {{.*}})
- func.func @op_dynamic_pad(%arg0: tensor<?xf32>, %arg1: tensor<f32>, %arg2: tensor<1xindex>, %arg3: tensor<1xindex>, %arg4: tensor<1xindex>) -> tensor<?xf32> {
--  // CHECK: "vhlo.dynamic_pad_v1"(%arg0, %arg1, %arg2, %arg3, %arg4) : (!vhlo.tensor_v1<?x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>) -> !vhlo.tensor_v1<?x!vhlo.f32_v1>
-+  // CHECK: "vhlo.dynamic_pad_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]], %[[ARG3]], %[[ARG4]]) : (!vhlo.tensor_v1<?x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>) -> !vhlo.tensor_v1<?x!vhlo.f32_v1>
-   %0 = "stablehlo.dynamic_pad"(%arg0, %arg1, %arg2, %arg3, %arg4) : (tensor<?xf32>, tensor<f32>, tensor<1xindex>, tensor<1xindex>, tensor<1xindex>) -> tensor<?xf32>
-   func.return %0 : tensor<?xf32>
- }
- 
- // CHECK-LABEL: "op_dynamic_reshape"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_dynamic_reshape(%arg0: tensor<16xf32>, %arg1: tensor<2xindex>) -> tensor<?x?xf32> {
--  // CHECK: "vhlo.dynamic_reshape_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<2x!vhlo.index_v1>) -> !vhlo.tensor_v1<?x?x!vhlo.f32_v1>
-+  // CHECK: "vhlo.dynamic_reshape_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<2x!vhlo.index_v1>) -> !vhlo.tensor_v1<?x?x!vhlo.f32_v1>
-   %0 = "stablehlo.dynamic_reshape"(%arg0, %arg1) : (tensor<16xf32>, tensor<2xindex>) -> tensor<?x?xf32>
-   func.return %0 : tensor<?x?xf32>
- }
- 
- // CHECK-LABEL: "op_dynamic_slice"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_dynamic_slice(%arg0: tensor<16xf32>, %arg1: tensor<i64>) -> tensor<4xf32> {
--  //      CHECK: "vhlo.dynamic_slice_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.dynamic_slice_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   slice_sizes = #vhlo.tensor_v1<dense<4> : tensor<1xi64>>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.i64_v1>) -> !vhlo.tensor_v1<4x!vhlo.f32_v1>
-   %0 = "stablehlo.dynamic_slice"(%arg0, %arg1) {
-@@ -1268,15 +1360,17 @@
- }
- 
- // CHECK-LABEL: "op_dynamic_update_slice"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_dynamic_update_slice(%arg0: tensor<16xf32>, %arg1: tensor<4xf32>, %arg2: tensor<i64>) -> tensor<16xf32> {
--  // CHECK: "vhlo.dynamic_update_slice_v1"(%arg0, %arg1, %arg2) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<4x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.i64_v1>) -> !vhlo.tensor_v1<16x!vhlo.f32_v1>
-+  // CHECK: "vhlo.dynamic_update_slice_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<4x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.i64_v1>) -> !vhlo.tensor_v1<16x!vhlo.f32_v1>
-   %0 = "stablehlo.dynamic_update_slice"(%arg0, %arg1, %arg2) : (tensor<16xf32>, tensor<4xf32>, tensor<i64>) -> tensor<16xf32>
-   func.return %0 : tensor<16xf32>
- }
- 
- // CHECK-LABEL: "op_einsum"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_einsum(%arg0: tensor<8x16xf32>, %arg1: tensor<16x8xf32>) -> tensor<8x8xf32> {
--  //      CHECK: "vhlo.einsum_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.einsum_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   einsum_config = #vhlo.string_v1<"ab,bc->ac">
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<8x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<8x8x!vhlo.f32_v1>
-   %0 = "stablehlo.einsum"(%arg0, %arg1) {
-@@ -1286,22 +1380,25 @@
- }
- 
- // CHECK-LABEL: "op_exponential_minus_one"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_exponential_minus_one(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.exponential_minus_one_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.exponential_minus_one_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.exponential_minus_one"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_exponential"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_exponential(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.exponential_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.exponential_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.exponential"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_fft"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_fft(%arg0: tensor<16xcomplex<f32>>) -> tensor<16xcomplex<f32>> {
--  //      CHECK: "vhlo.fft_v1"(%arg0) <{
-+  //      CHECK: "vhlo.fft_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   fft_length = #vhlo.tensor_v1<dense<16> : tensor<1xi64>>,
-   // CHECK-SAME:   fft_type = #vhlo<fft_type_v1 FFT>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<16x!vhlo.complex_v1<!vhlo.f32_v1>>) -> !vhlo.tensor_v1<16x!vhlo.complex_v1<!vhlo.f32_v1>>
-@@ -1313,8 +1410,9 @@
- }
- 
- // CHECK-LABEL: "op_floor"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_floor(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.floor_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.floor_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.floor"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
-@@ -1327,16 +1425,17 @@
-   // CHECK-SAME:   sym_name = #vhlo.string_v1<"op_func">,
-   // CHECK-SAME:   sym_visibility = #vhlo.string_v1<"private">
-   // CHECK-SAME: }> ({
--  // CHECK-NEXT:   ^[[BB:bb.*]](%arg0: !vhlo.tensor_v1<!vhlo.f32_v1>):
--  // CHECK-NEXT:     "vhlo.return_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-+  // CHECK-NEXT:   ^[[BB:bb.*]](%[[ARG0:.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
-+  // CHECK-NEXT:     "vhlo.return_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-   // CHECK-NEXT: }) : () -> ()
- 
-   func.return %arg0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_gather"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_gather(%arg0 : tensor<2x4x9xf32>, %arg1 : tensor<1x5x2xi32>) -> tensor<1x5x1xf32> {
--  //      CHECK: "vhlo.gather_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.gather_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   collapsed_slice_dims = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-   // CHECK-SAME:   index_vector_dim = #vhlo.integer_v1<2 : i64>,
-   // CHECK-SAME:   indices_are_sorted = #vhlo.bool_v1<true>,
-@@ -1358,8 +1457,9 @@
- }
- 
- // CHECK-LABEL: "op_get_dimension_size"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_get_dimension_size(%arg0: tensor<?xf32>) -> tensor<i32> {
--  //      CHECK: "vhlo.get_dimension_size_v1"(%arg0) <{
-+  //      CHECK: "vhlo.get_dimension_size_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   dimension = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<?x!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-   %0 = "stablehlo.get_dimension_size"(%arg0) {
-@@ -1369,8 +1469,9 @@
- }
- 
- // CHECK-LABEL: "op_get_tuple_element"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_get_tuple_element(%arg0: tuple<tensor<f32>, tensor<i32>>) -> tensor<f32> {
--  //      CHECK: "vhlo.get_tuple_element_v1"(%arg0) <{
-+  //      CHECK: "vhlo.get_tuple_element_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   index = #vhlo.integer_v1<0 : i32>
-   // CHECK-SAME: }> : (!vhlo.tuple_v1<!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.get_tuple_element"(%arg0) {
-@@ -1380,11 +1481,12 @@
- }
- 
- // CHECK-LABEL: "op_if"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_if(%arg0: tensor<i1>, %arg1: tensor<f32>, %arg2: tensor<f32>) -> tensor<f32> {
--  //      CHECK: "vhlo.if_v1"(%arg0) ({
--  // CHECK-NEXT:   "vhlo.return_v1"(%arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-+  //      CHECK: "vhlo.if_v1"(%[[ARG0]]) ({
-+  // CHECK-NEXT:   "vhlo.return_v1"(%[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-   // CHECK-NEXT: }, {
--  // CHECK-NEXT:   "vhlo.return_v1"(%arg2) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-+  // CHECK-NEXT:   "vhlo.return_v1"(%[[ARG2]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-   // CHECK-NEXT: }) : (!vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.if"(%arg0) ({
-     "stablehlo.return"(%arg1) : (tensor<f32>) -> ()
-@@ -1395,15 +1497,17 @@
- }
- 
- // CHECK-LABEL: "op_imag"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_imag(%arg0: tensor<complex<f32>>) -> tensor<f32> {
--  // CHECK: "vhlo.imag_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.imag_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.imag"(%arg0) : (tensor<complex<f32>>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_infeed"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_infeed(%arg0: !stablehlo.token) -> (tensor<f32>, !stablehlo.token) {
--  //               CHECK: "vhlo.infeed_v1"(%arg0) <{
-+  //               CHECK: "vhlo.infeed_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   infeed_config = #vhlo.string_v1<"foo">,
-   // CHECK-SAME{LITERAL}:   layout = #vhlo.array_v1<[#vhlo.array_v1<[]>]>
-   //          CHECK-SAME: }> : (!vhlo.token_v1) -> (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.token_v1)
-@@ -1426,36 +1530,41 @@
- }
- 
- // CHECK-LABEL: "op_is_finite"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_is_finite(%arg0: tensor<f32>) -> tensor<i1> {
--  // CHECK: "vhlo.is_finite_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-+  // CHECK: "vhlo.is_finite_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-   %0 = "stablehlo.is_finite"(%arg0) : (tensor<f32>) -> tensor<i1>
-   func.return %0 : tensor<i1>
- }
- 
- // CHECK-LABEL: "op_log"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_log(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.log_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.log_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.log"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_log_plus_one"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_log_plus_one(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.log_plus_one_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.log_plus_one_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.log_plus_one"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_logistic"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_logistic(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.logistic_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.logistic_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.logistic"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_map"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_map(%arg0: tensor<16xf32>) -> tensor<16xf32> {
--  //      CHECK: "vhlo.map_v1"(%arg0) <{
-+  //      CHECK: "vhlo.map_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   dimensions = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>
-   // CHECK-SAME: }> ({
-   // CHECK-NEXT:   ^[[BB:bb.*]](%[[ARG1:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
-@@ -1473,57 +1582,65 @@
- }
- 
- // CHECK-LABEL: "op_maximum"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_maximum(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.maximum_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.maximum_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.maximum"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_minimum"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_minimum(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.minimum_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.minimum_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.minimum"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_multiply"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_multiply(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.multiply_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.multiply_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.multiply"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_negate"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_negate(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.negate_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.negate_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.negate"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_not"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_not(%arg0: tensor<i1>) -> tensor<i1> {
--  // CHECK: "vhlo.not_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-+  // CHECK: "vhlo.not_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-   %0 = "stablehlo.not"(%arg0) : (tensor<i1>) -> tensor<i1>
-   func.return %0 : tensor<i1>
- }
- 
- // CHECK-LABEL: "op_optimization_barrier"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_optimization_barrier(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.optimization_barrier_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.optimization_barrier_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.optimization_barrier"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_or"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_or(%arg0: tensor<i1>, %arg1: tensor<i1>) -> tensor<i1> {
--  // CHECK: "vhlo.or_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-+  // CHECK: "vhlo.or_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-   %0 = "stablehlo.or"(%arg0, %arg1) : (tensor<i1>, tensor<i1>) -> tensor<i1>
-   func.return %0 : tensor<i1>
- }
- 
- // CHECK-LABEL: "op_outfeed"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_outfeed(%arg0: tensor<f32>, %arg1: !stablehlo.token) -> !stablehlo.token {
--  //      CHECK: "vhlo.outfeed_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.outfeed_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   outfeed_config = #vhlo.string_v1<"foo">
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.token_v1) -> !vhlo.token_v1
-   %0 = "stablehlo.outfeed"(%arg0, %arg1) {
-@@ -1533,8 +1650,9 @@
- }
- 
- // CHECK-LABEL: "op_pad"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_pad(%arg0: tensor<8xf32>, %arg1: tensor<f32>) -> tensor<16xf32> {
--  //      CHECK: "vhlo.pad_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.pad_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   edge_padding_high = #vhlo.tensor_v1<dense<4> : tensor<1xi64>>,
-   // CHECK-SAME:   edge_padding_low = #vhlo.tensor_v1<dense<4> : tensor<1xi64>>,
-   // CHECK-SAME:   interior_padding = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>
-@@ -1548,36 +1666,41 @@
- }
- 
- // CHECK-LABEL: "op_popcnt"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_popcnt(%arg0: tensor<i32>) -> tensor<i32> {
--  // CHECK: "vhlo.popcnt_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-+  // CHECK: "vhlo.popcnt_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-   %0 = "stablehlo.popcnt"(%arg0) : (tensor<i32>) -> tensor<i32>
-   func.return %0 : tensor<i32>
- }
- 
- // CHECK-LABEL: "op_power"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_power(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.power_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.power_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.power"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_real_dynamic_slice"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}}, %[[ARG3:.*]]: {{.*}})
- func.func @op_real_dynamic_slice(%arg0: tensor<?xf32>, %arg1: tensor<1xindex>, %arg2: tensor<1xindex>, %arg3: tensor<1xindex>) -> tensor<?xf32> {
--  // CHECK: "vhlo.real_dynamic_slice_v1"(%arg0, %arg1, %arg2, %arg3) : (!vhlo.tensor_v1<?x!vhlo.f32_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>) -> !vhlo.tensor_v1<?x!vhlo.f32_v1>
-+  // CHECK: "vhlo.real_dynamic_slice_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]], %[[ARG3]]) : (!vhlo.tensor_v1<?x!vhlo.f32_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>) -> !vhlo.tensor_v1<?x!vhlo.f32_v1>
-   %0 = "stablehlo.real_dynamic_slice"(%arg0, %arg1, %arg2, %arg3) : (tensor<?xf32>, tensor<1xindex>, tensor<1xindex>, tensor<1xindex>) -> tensor<?xf32>
-   func.return %0 : tensor<?xf32>
- }
- 
- // CHECK-LABEL: "op_real"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_real(%arg0: tensor<complex<f32>>) -> tensor<f32> {
--  // CHECK: "vhlo.real_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.real_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.real"(%arg0) : (tensor<complex<f32>>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_recv"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_recv(%arg0: !stablehlo.token) -> (tensor<f32>, !stablehlo.token) {
--  //      CHECK: "vhlo.recv_v1"(%arg0) <{
-+  //      CHECK: "vhlo.recv_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME:   channel_type = #vhlo.integer_v1<3 : i64>,
-   // CHECK-SAME:   is_host_transfer = #vhlo.bool_v1<true>
-@@ -1590,8 +1713,9 @@
- }
- 
- // CHECK-LABEL: "op_reduce"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_reduce(%arg0: tensor<16xf32>, %arg1: tensor<f32>) -> tensor<f32> {
--  //  CHECK: "vhlo.reduce_v1"(%arg0, %arg1)
-+  //  CHECK: "vhlo.reduce_v1"(%[[ARG0]], %[[ARG1]])
-   //  CHECK:   ^[[BB:bb.*]](%[[ARG1:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>, %[[ARG2:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
-   //  CHECK:     "vhlo.return_v1"(%[[VAL1:.*]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-   //  CHECK: }) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-@@ -1606,8 +1730,9 @@
- }
- 
- // CHECK-LABEL: "op_reduce_precision"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_reduce_precision(%arg0: tensor<f32>) -> tensor<f32> {
--  //      CHECK: "vhlo.reduce_precision_v1"(%arg0) <{
-+  //      CHECK: "vhlo.reduce_precision_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   exponent_bits = #vhlo.integer_v1<8 : i32>
-   // CHECK-SAME:   mantissa_bits = #vhlo.integer_v1<10 : i32>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-@@ -1619,8 +1744,9 @@
- }
- 
- // CHECK-LABEL: "op_reduce_scatter"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_reduce_scatter(%arg0: tensor<16xf32>) -> tensor<16xf32> {
--  //               CHECK: "vhlo.reduce_scatter_v1"(%arg0) <{
-+  //               CHECK: "vhlo.reduce_scatter_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0], [1]]> : tensor<2x1xi64>>,
-   //          CHECK-SAME:   scatter_dimension = #vhlo.integer_v1<0 : i64>
-@@ -1644,8 +1770,9 @@
- }
- 
- // CHECK-LABEL: "op_reduce_window"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_reduce_window(%arg0: tensor<2x17x31x7xf32>, %arg1: tensor<f32>) -> tensor<2x9x16x7xf32> {
--  //               CHECK: "vhlo.reduce_window_v1"(%arg0, %arg1) <{
-+  //               CHECK: "vhlo.reduce_window_v1"(%[[ARG0]], %[[ARG1]]) <{
-   //          CHECK-SAME:   base_dilations = #vhlo.tensor_v1<dense<[1, 2, 2, 1]> : tensor<4xi64>>,
-   // CHECK-SAME{LITERAL}:   padding = #vhlo.tensor_v1<dense<[[0, 0], [2, 0], [0, 2], [0, 0]]> : tensor<4x2xi64>>,
-   //          CHECK-SAME:   window_dilations = #vhlo.tensor_v1<dense<[1, 2, 2, 1]> : tensor<4xi64>>,
-@@ -1671,8 +1798,9 @@
- }
- 
- // CHECK-LABEL: "op_remainder"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_remainder(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.remainder_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.remainder_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.remainder"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
-@@ -1692,16 +1820,18 @@
- }
- 
- // CHECK-LABEL: "op_reshape"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_reshape(%arg0: tensor<16xf32>) -> tensor<4x4xf32> {
--  // CHECK: "vhlo.reshape_v1"(%arg0) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<4x4x!vhlo.f32_v1>
-+  // CHECK: "vhlo.reshape_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<4x4x!vhlo.f32_v1>
-   %0 = "stablehlo.reshape"(%arg0) : (tensor<16xf32>) -> tensor<4x4xf32>
-   func.return %0 : tensor<4x4xf32>
- }
- 
- // CHECK-LABEL: "op_return"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_return(%arg0: tensor<i32>, %arg1: tensor<f32>) -> tensor<f32> {
--  //      CHECK: "vhlo.case_v1"(%arg0) ({
--  // CHECK-NEXT:   "vhlo.return_v1"(%arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-+  //      CHECK: "vhlo.case_v1"(%[[ARG0]]) ({
-+  // CHECK-NEXT:   "vhlo.return_v1"(%[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-   // CHECK-NEXT: }) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.case"(%arg0) ({
-     "stablehlo.return"(%arg1) : (tensor<f32>) -> ()
-@@ -1710,8 +1840,9 @@
- }
- 
- // CHECK-LABEL: "op_reverse"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_reverse(%arg0: tensor<16xf32>) -> tensor<16xf32> {
--  //      CHECK: "vhlo.reverse_v1"(%arg0) <{
-+  //      CHECK: "vhlo.reverse_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   dimensions = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x!vhlo.f32_v1>
-   %0 = "stablehlo.reverse"(%arg0) {
-@@ -1721,8 +1852,9 @@
- }
- 
- // CHECK-LABEL: "op_rng_bit_generator"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_rng_bit_generator(%arg0: tensor<f32>) -> (tensor<f32>, tensor<f32>) {
--  //      CHECK: "vhlo.rng_bit_generator_v1"(%arg0) <{
-+  //      CHECK: "vhlo.rng_bit_generator_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   rng_algorithm = #vhlo<rng_algorithm_v1 PHILOX>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>)
-   %0:2 = "stablehlo.rng_bit_generator"(%arg0) {
-@@ -1732,8 +1864,9 @@
- }
- 
- // CHECK-LABEL: "op_rng"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_rng(%arg0: tensor<f32>, %arg1: tensor<f32>, %arg2: tensor<0xindex>) -> tensor<f32> {
--  //      CHECK: "vhlo.rng_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.rng_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   rng_distribution = #vhlo<rng_distribution_v1 NORMAL>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<0x!vhlo.index_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.rng"(%arg0, %arg1, %arg2) {
-@@ -1743,29 +1876,33 @@
- }
- 
- // CHECK-LABEL: "op_round_nearest_afz"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_round_nearest_afz(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.round_nearest_afz_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.round_nearest_afz_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.round_nearest_afz"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_round_nearest_even"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_round_nearest_even(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.round_nearest_even_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.round_nearest_even_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.round_nearest_even"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_rsqrt"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_rsqrt(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.rsqrt_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.rsqrt_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.rsqrt"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_scatter"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_scatter(%arg0: tensor<200x100x300xf32>, %arg1: tensor<10x2xi32>, %arg2: tensor<10x300xf32>) -> tensor<200x100x300xf32> {
--  //      CHECK: "vhlo.scatter_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.scatter_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   index_vector_dim = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   indices_are_sorted = #vhlo.bool_v1<true>,
-   // CHECK-SAME:   inserted_window_dims = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-@@ -1795,8 +1932,9 @@
- }
- 
- // CHECK-LABEL: "op_select_and_scatter"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_select_and_scatter(%arg0: tensor<10x24x24x64xf32>, %arg1: tensor<12x13x13x66xf32>, %arg2: tensor<f32>) -> tensor<10x24x24x64xf32> {
--  //      CHECK: "vhlo.select_and_scatter_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.select_and_scatter_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   padding = #vhlo.tensor_v1<dense<1> : tensor<4x2xi64>>,
-   // CHECK-SAME:   window_dimensions = #vhlo.tensor_v1<dense<[1, 2, 2, 1]> : tensor<4xi64>>,
-   // CHECK-SAME:   window_strides = #vhlo.tensor_v1<dense<[1, 2, 2, 1]> : tensor<4xi64>>
-@@ -1826,15 +1964,17 @@
- }
- 
- // CHECK-LABEL: "op_select"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_select(%arg0: tensor<i1>, %arg1: tensor<f32>, %arg2: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.select_v1"(%arg0, %arg1, %arg2) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.select_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.select"(%arg0, %arg1, %arg2) : (tensor<i1>, tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_send"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_send(%arg0: tensor<f32>, %arg1: !stablehlo.token) -> !stablehlo.token {
--  //      CHECK: "vhlo.send_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.send_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME:   channel_type = #vhlo.integer_v1<2 : i64>,
-   // CHECK-SAME:   is_host_transfer = #vhlo.bool_v1<true>
-@@ -1847,8 +1987,9 @@
- }
- 
- // CHECK-LABEL: "op_set_dimension_size"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_set_dimension_size(%arg0: tensor<?xf32>, %arg1: tensor<i32>) -> tensor<16xf32> {
--  //      CHECK: "vhlo.set_dimension_size_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.set_dimension_size_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   dimension = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<?x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<16x!vhlo.f32_v1>
-   %0 = "stablehlo.set_dimension_size"(%arg0, %arg1) {
-@@ -1858,43 +1999,49 @@
- }
- 
- // CHECK-LABEL: "op_shift_left"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_shift_left(%arg0: tensor<i32>, %arg1: tensor<i32>) -> tensor<i32> {
--  // CHECK: "vhlo.shift_left_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.i32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-+  // CHECK: "vhlo.shift_left_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-   %0 = "stablehlo.shift_left"(%arg0, %arg1) : (tensor<i32>, tensor<i32>) -> tensor<i32>
-   func.return %0 : tensor<i32>
- }
- 
- // CHECK-LABEL: "op_shift_right_arithmetic"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_shift_right_arithmetic(%arg0: tensor<i32>, %arg1: tensor<i32>) -> tensor<i32> {
--  // CHECK: "vhlo.shift_right_arithmetic_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.i32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-+  // CHECK: "vhlo.shift_right_arithmetic_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-   %0 = "stablehlo.shift_right_arithmetic"(%arg0, %arg1) : (tensor<i32>, tensor<i32>) -> tensor<i32>
-   func.return %0 : tensor<i32>
- }
- 
- // CHECK-LABEL: "op_shift_right_logical"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_shift_right_logical(%arg0: tensor<i32>, %arg1: tensor<i32>) -> tensor<i32> {
--  // CHECK: "vhlo.shift_right_logical_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.i32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-+  // CHECK: "vhlo.shift_right_logical_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-   %0 = "stablehlo.shift_right_logical"(%arg0, %arg1) : (tensor<i32>, tensor<i32>) -> tensor<i32>
-   func.return %0 : tensor<i32>
- }
- 
- // CHECK-LABEL: "op_sign"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_sign(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.sign_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.sign_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.sign"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_sine"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_sine(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.sine_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.sine_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.sine"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_slice"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_slice(%arg0: tensor<16xf32>) -> tensor<4xf32> {
--  //      CHECK: "vhlo.slice_v1"(%arg0) <{
-+  //      CHECK: "vhlo.slice_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   limit_indices = #vhlo.tensor_v1<dense<4> : tensor<1xi64>>,
-   // CHECK-SAME:   start_indices = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>,
-   // CHECK-SAME:   strides = #vhlo.tensor_v1<dense<1> : tensor<1xi64>>
-@@ -1908,8 +2055,9 @@
- }
- 
- // CHECK-LABEL: "op_sort"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_sort(%arg0: tensor<16xf32>) -> tensor<16xf32> {
--  //      CHECK: "vhlo.sort_v1"(%arg0) <{
-+  //      CHECK: "vhlo.sort_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   dimension = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME:   is_stable = #vhlo.bool_v1<true>
-   // CHECK-SAME: }> ({
-@@ -1929,29 +2077,33 @@
- }
- 
- // CHECK-LABEL: "op_sqrt"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_sqrt(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.sqrt_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.sqrt_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.sqrt"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_subtract"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_subtract(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.subtract_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.subtract_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.subtract"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_tanh"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_tanh(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.tanh_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.tanh_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.tanh"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_torch_index_select"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_torch_index_select(%arg0: tensor<5x1x5xf32>, %arg1: tensor<2xi32>) ->  tensor<2x1x5xf32> {
--  //      CHECK: "vhlo.torch_index_select_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.torch_index_select_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   batch_dims = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME:   dim = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<5x1x5x!vhlo.f32_v1>, !vhlo.tensor_v1<2x!vhlo.i32_v1>) -> !vhlo.tensor_v1<2x1x5x!vhlo.f32_v1>
-@@ -1963,8 +2115,9 @@
- }
- 
- // CHECK-LABEL: "op_trace"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_trace(%arg0: tensor<f32>) {
--  //      CHECK: "vhlo.trace_v1"(%arg0) <{
-+  //      CHECK: "vhlo.trace_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   tag = #vhlo.string_v1<"foo">
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-   "stablehlo.trace"(%arg0) {
-@@ -1974,8 +2127,9 @@
- }
- 
- // CHECK-LABEL: "op_transpose"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_transpose(%arg0: tensor<16x8xf32>) ->  tensor<8x16xf32> {
--  //      CHECK: "vhlo.transpose_v1"(%arg0) <{
-+  //      CHECK: "vhlo.transpose_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   permutation = #vhlo.tensor_v1<dense<[1, 0]> : tensor<2xi64>>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<16x8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<8x16x!vhlo.f32_v1>
-   %0 = "stablehlo.transpose"(%arg0) {
-@@ -1985,8 +2139,9 @@
- }
- 
- // CHECK-LABEL: "op_triangular_solve"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_triangular_solve(%arg0: tensor<16x16xf32>, %arg1: tensor<16x16xf32>) ->  tensor<16x16xf32> {
--  //      CHECK: "vhlo.triangular_solve_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.triangular_solve_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   left_side = #vhlo.bool_v1<true>,
-   // CHECK-SAME:   lower = #vhlo.bool_v1<true>,
-   // CHECK-SAME:   transpose_a = #vhlo<transpose_v1 NO_TRANSPOSE>,
-@@ -2002,15 +2157,17 @@
- }
- 
- // CHECK-LABEL: "op_tuple"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_tuple(%arg0: tensor<f32>) -> tuple<tensor<f32>> {
--  // CHECK: "vhlo.tuple_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tuple_v1<!vhlo.tensor_v1<!vhlo.f32_v1>>
-+  // CHECK: "vhlo.tuple_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tuple_v1<!vhlo.tensor_v1<!vhlo.f32_v1>>
-   %0 = "stablehlo.tuple"(%arg0) : (tensor<f32>) -> tuple<tensor<f32>>
-   func.return %0 : tuple<tensor<f32>>
- }
- 
- // CHECK-LABEL: "op_unary_einsum"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_unary_einsum(%arg0: tensor<8x16xf32>) -> tensor<8xf32> {
--  //      CHECK: "vhlo.unary_einsum_v1"(%arg0) <{
-+  //      CHECK: "vhlo.unary_einsum_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   einsum_config = #vhlo.string_v1<"ab->a">
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<8x16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<8x!vhlo.f32_v1>
-   %0 = "stablehlo.unary_einsum"(%arg0) {
-@@ -2020,22 +2177,25 @@
- }
- 
- // CHECK-LABEL: "op_uniform_dequantize"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_uniform_dequantize(%arg0: tensor<!quant.uniform<i8:f32, 34.0:16>>) -> tensor<f32> {
--  // CHECK: "vhlo.uniform_dequantize_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.uniform_dequantize_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.uniform_dequantize"(%arg0) : (tensor<!quant.uniform<i8:f32, 34.0:16>>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_uniform_quantize"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_uniform_quantize(%arg0: tensor<f32>) -> tensor<!quant.uniform<i8:f32, 34.0:16>> {
--  // CHECK: "vhlo.uniform_quantize_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>
-+  // CHECK: "vhlo.uniform_quantize_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>
-   %0 = "stablehlo.uniform_quantize"(%arg0) : (tensor<f32>) -> tensor<!quant.uniform<i8:f32, 34.0:16>>
-   func.return %0 : tensor<!quant.uniform<i8:f32, 34.0:16>>
- }
- 
- // CHECK-LABEL: "op_while"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_while(%arg0: tensor<i1>) -> tensor<i1> {
--  //      CHECK: "vhlo.while_v1"(%arg0) ({
-+  //      CHECK: "vhlo.while_v1"(%[[ARG0]]) ({
-   // CHECK-NEXT:   ^[[BB:bb.*]](%[[ARG1:arg.*]]: !vhlo.tensor_v1<!vhlo.bool_v1>):
-   // CHECK-NEXT:     "vhlo.return_v1"(%[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>) -> ()
-   // CHECK-NEXT:   }, {
-@@ -2053,8 +2213,9 @@
- }
- 
- // CHECK-LABEL: "op_xor"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_xor(%arg0: tensor<i1>, %arg1: tensor<i1>) -> tensor<i1> {
--  // CHECK: "vhlo.xor_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-+  // CHECK: "vhlo.xor_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-   %0 = "stablehlo.xor"(%arg0, %arg1) : (tensor<i1>, tensor<i1>) -> tensor<i1>
-   func.return %0 : tensor<i1>
- }
-@@ -2062,190 +2223,217 @@
- // ============ TYPES ============
- 
- // CHECK-LABEL: "type_i1"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_i1(%arg0: tensor<i1>, %arg1: tensor<i1>) -> tensor<i1> {
--  // CHECK: "vhlo.and_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-+  // CHECK: "vhlo.and_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-   %0 = "stablehlo.and"(%arg0, %arg1) : (tensor<i1>, tensor<i1>) -> tensor<i1>
-   func.return %0 : tensor<i1>
- }
- 
- // CHECK-LABEL: "type_i4"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_i4(%arg0: tensor<i4>, %arg1: tensor<i4>) -> tensor<i4> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.i4_v1>, !vhlo.tensor_v1<!vhlo.i4_v1>) -> !vhlo.tensor_v1<!vhlo.i4_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i4_v1>, !vhlo.tensor_v1<!vhlo.i4_v1>) -> !vhlo.tensor_v1<!vhlo.i4_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<i4>, tensor<i4>) -> tensor<i4>
-   func.return %0 : tensor<i4>
- }
- 
- // CHECK-LABEL: "type_i8"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_i8(%arg0: tensor<i8>, %arg1: tensor<i8>) -> tensor<i8> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.i8_v1>, !vhlo.tensor_v1<!vhlo.i8_v1>) -> !vhlo.tensor_v1<!vhlo.i8_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i8_v1>, !vhlo.tensor_v1<!vhlo.i8_v1>) -> !vhlo.tensor_v1<!vhlo.i8_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<i8>, tensor<i8>) -> tensor<i8>
-   func.return %0 : tensor<i8>
- }
- 
- // CHECK-LABEL: "type_i16"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_i16(%arg0: tensor<i16>, %arg1: tensor<i16>) -> tensor<i16> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.i16_v1>, !vhlo.tensor_v1<!vhlo.i16_v1>) -> !vhlo.tensor_v1<!vhlo.i16_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i16_v1>, !vhlo.tensor_v1<!vhlo.i16_v1>) -> !vhlo.tensor_v1<!vhlo.i16_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<i16>, tensor<i16>) -> tensor<i16>
-   func.return %0 : tensor<i16>
- }
- 
- // CHECK-LABEL: "type_i32"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_i32(%arg0: tensor<i32>, %arg1: tensor<i32>) -> tensor<i32> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.i32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<i32>, tensor<i32>) -> tensor<i32>
-   func.return %0 : tensor<i32>
- }
- 
- // CHECK-LABEL: "type_i64"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_i64(%arg0: tensor<i64>, %arg1: tensor<i64>) -> tensor<i64> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.i64_v1>, !vhlo.tensor_v1<!vhlo.i64_v1>) -> !vhlo.tensor_v1<!vhlo.i64_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i64_v1>, !vhlo.tensor_v1<!vhlo.i64_v1>) -> !vhlo.tensor_v1<!vhlo.i64_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<i64>, tensor<i64>) -> tensor<i64>
-   func.return %0 : tensor<i64>
- }
- 
- // CHECK-LABEL: "type_ui4"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_ui4(%arg0: tensor<ui4>, %arg1: tensor<ui4>) -> tensor<ui4> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.ui4_v1>, !vhlo.tensor_v1<!vhlo.ui4_v1>) -> !vhlo.tensor_v1<!vhlo.ui4_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.ui4_v1>, !vhlo.tensor_v1<!vhlo.ui4_v1>) -> !vhlo.tensor_v1<!vhlo.ui4_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<ui4>, tensor<ui4>) -> tensor<ui4>
-   func.return %0 : tensor<ui4>
- }
- 
- // CHECK-LABEL: "type_ui8"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_ui8(%arg0: tensor<ui8>, %arg1: tensor<ui8>) -> tensor<ui8> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.ui8_v1>, !vhlo.tensor_v1<!vhlo.ui8_v1>) -> !vhlo.tensor_v1<!vhlo.ui8_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.ui8_v1>, !vhlo.tensor_v1<!vhlo.ui8_v1>) -> !vhlo.tensor_v1<!vhlo.ui8_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<ui8>, tensor<ui8>) -> tensor<ui8>
-   func.return %0 : tensor<ui8>
- }
- 
- // CHECK-LABEL: "type_ui16"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_ui16(%arg0: tensor<ui16>, %arg1: tensor<ui16>) -> tensor<ui16> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.ui16_v1>, !vhlo.tensor_v1<!vhlo.ui16_v1>) -> !vhlo.tensor_v1<!vhlo.ui16_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.ui16_v1>, !vhlo.tensor_v1<!vhlo.ui16_v1>) -> !vhlo.tensor_v1<!vhlo.ui16_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<ui16>, tensor<ui16>) -> tensor<ui16>
-   func.return %0 : tensor<ui16>
- }
- 
- // CHECK-LABEL: "type_ui32"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_ui32(%arg0: tensor<ui32>, %arg1: tensor<ui32>) -> tensor<ui32> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.ui32_v1>, !vhlo.tensor_v1<!vhlo.ui32_v1>) -> !vhlo.tensor_v1<!vhlo.ui32_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.ui32_v1>, !vhlo.tensor_v1<!vhlo.ui32_v1>) -> !vhlo.tensor_v1<!vhlo.ui32_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<ui32>, tensor<ui32>) -> tensor<ui32>
-   func.return %0 : tensor<ui32>
- }
- 
- // CHECK-LABEL: "type_ui64"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_ui64(%arg0: tensor<ui64>, %arg1: tensor<ui64>) -> tensor<ui64> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.ui64_v1>, !vhlo.tensor_v1<!vhlo.ui64_v1>) -> !vhlo.tensor_v1<!vhlo.ui64_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.ui64_v1>, !vhlo.tensor_v1<!vhlo.ui64_v1>) -> !vhlo.tensor_v1<!vhlo.ui64_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<ui64>, tensor<ui64>) -> tensor<ui64>
-   func.return %0 : tensor<ui64>
- }
- 
- // CHECK-LABEL: "type_f8E4M3FN"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_f8E4M3FN(%arg0: tensor<f8E4M3FN>, %arg1: tensor<f8E4M3FN>) -> tensor<f8E4M3FN> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f8E4M3FN_v1>, !vhlo.tensor_v1<!vhlo.f8E4M3FN_v1>) -> !vhlo.tensor_v1<!vhlo.f8E4M3FN_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f8E4M3FN_v1>, !vhlo.tensor_v1<!vhlo.f8E4M3FN_v1>) -> !vhlo.tensor_v1<!vhlo.f8E4M3FN_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f8E4M3FN>, tensor<f8E4M3FN>) -> tensor<f8E4M3FN>
-   func.return %0 : tensor<f8E4M3FN>
- }
- 
- // CHECK-LABEL: "type_f8E5M2"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_f8E5M2(%arg0: tensor<f8E5M2>, %arg1: tensor<f8E5M2>) -> tensor<f8E5M2> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f8E5M2_v1>, !vhlo.tensor_v1<!vhlo.f8E5M2_v1>) -> !vhlo.tensor_v1<!vhlo.f8E5M2_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f8E5M2_v1>, !vhlo.tensor_v1<!vhlo.f8E5M2_v1>) -> !vhlo.tensor_v1<!vhlo.f8E5M2_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f8E5M2>, tensor<f8E5M2>) -> tensor<f8E5M2>
-   func.return %0 : tensor<f8E5M2>
- }
- 
- // CHECK-LABEL: "type_f8E4M3FNUZ"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_f8E4M3FNUZ(%arg0: tensor<f8E4M3FNUZ>, %arg1: tensor<f8E4M3FNUZ>) -> tensor<f8E4M3FNUZ> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f8E4M3FNUZ_v1>, !vhlo.tensor_v1<!vhlo.f8E4M3FNUZ_v1>) -> !vhlo.tensor_v1<!vhlo.f8E4M3FNUZ_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f8E4M3FNUZ_v1>, !vhlo.tensor_v1<!vhlo.f8E4M3FNUZ_v1>) -> !vhlo.tensor_v1<!vhlo.f8E4M3FNUZ_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f8E4M3FNUZ>, tensor<f8E4M3FNUZ>) -> tensor<f8E4M3FNUZ>
-   func.return %0 : tensor<f8E4M3FNUZ>
- }
- 
- // CHECK-LABEL: "type_f8E4M3B11FNUZ"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_f8E4M3B11FNUZ(%arg0: tensor<f8E4M3B11FNUZ>, %arg1: tensor<f8E4M3B11FNUZ>) -> tensor<f8E4M3B11FNUZ> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f8E4M3B11FNUZ_v1>, !vhlo.tensor_v1<!vhlo.f8E4M3B11FNUZ_v1>) -> !vhlo.tensor_v1<!vhlo.f8E4M3B11FNUZ_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f8E4M3B11FNUZ_v1>, !vhlo.tensor_v1<!vhlo.f8E4M3B11FNUZ_v1>) -> !vhlo.tensor_v1<!vhlo.f8E4M3B11FNUZ_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f8E4M3B11FNUZ>, tensor<f8E4M3B11FNUZ>) -> tensor<f8E4M3B11FNUZ>
-   func.return %0 : tensor<f8E4M3B11FNUZ>
- }
- 
- // CHECK-LABEL: "type_f8E5M2FNUZ"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_f8E5M2FNUZ(%arg0: tensor<f8E5M2FNUZ>, %arg1: tensor<f8E5M2FNUZ>) -> tensor<f8E5M2FNUZ> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f8E5M2FNUZ_v1>, !vhlo.tensor_v1<!vhlo.f8E5M2FNUZ_v1>) -> !vhlo.tensor_v1<!vhlo.f8E5M2FNUZ_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f8E5M2FNUZ_v1>, !vhlo.tensor_v1<!vhlo.f8E5M2FNUZ_v1>) -> !vhlo.tensor_v1<!vhlo.f8E5M2FNUZ_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f8E5M2FNUZ>, tensor<f8E5M2FNUZ>) -> tensor<f8E5M2FNUZ>
-   func.return %0 : tensor<f8E5M2FNUZ>
- }
- 
- // CHECK-LABEL: "type_bf16"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_bf16(%arg0: tensor<bf16>, %arg1: tensor<bf16>) -> tensor<bf16> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.bf16_v1>, !vhlo.tensor_v1<!vhlo.bf16_v1>) -> !vhlo.tensor_v1<!vhlo.bf16_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.bf16_v1>, !vhlo.tensor_v1<!vhlo.bf16_v1>) -> !vhlo.tensor_v1<!vhlo.bf16_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<bf16>, tensor<bf16>) -> tensor<bf16>
-   func.return %0 : tensor<bf16>
- }
- 
- // CHECK-LABEL: "type_f16"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_f16(%arg0: tensor<f16>, %arg1: tensor<f16>) -> tensor<f16> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f16_v1>, !vhlo.tensor_v1<!vhlo.f16_v1>) -> !vhlo.tensor_v1<!vhlo.f16_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f16_v1>, !vhlo.tensor_v1<!vhlo.f16_v1>) -> !vhlo.tensor_v1<!vhlo.f16_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f16>, tensor<f16>) -> tensor<f16>
-   func.return %0 : tensor<f16>
- }
- 
- // CHECK-LABEL: "type_f32"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_f32(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "type_f64"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_f64(%arg0: tensor<f64>, %arg1: tensor<f64>) -> tensor<f64> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f64_v1>, !vhlo.tensor_v1<!vhlo.f64_v1>) -> !vhlo.tensor_v1<!vhlo.f64_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f64_v1>, !vhlo.tensor_v1<!vhlo.f64_v1>) -> !vhlo.tensor_v1<!vhlo.f64_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f64>, tensor<f64>) -> tensor<f64>
-   func.return %0 : tensor<f64>
- }
- 
- // CHECK-LABEL: "type_complex_f32"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_complex_f32(%arg0: tensor<complex<f32>>, %arg1: tensor<complex<f32>>) -> tensor<complex<f32>> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>, !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>) -> !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>, !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>) -> !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<complex<f32>>, tensor<complex<f32>>) -> tensor<complex<f32>>
-   func.return %0 : tensor<complex<f32>>
- }
- 
- // CHECK-LABEL: "type_complex_f64"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_complex_f64(%arg0: tensor<complex<f64>>, %arg1: tensor<complex<f64>>) -> tensor<complex<f64>> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f64_v1>>, !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f64_v1>>) -> !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f64_v1>>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f64_v1>>, !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f64_v1>>) -> !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f64_v1>>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<complex<f64>>, tensor<complex<f64>>) -> tensor<complex<f64>>
-   func.return %0 : tensor<complex<f64>>
- }
- 
- // CHECK-LABEL: "type_dynamism_ranked"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @type_dynamism_ranked(%arg0: tensor<?xf32>) -> tensor<?xf32> {
--  // CHECK: "vhlo.abs_v1"(%arg0) : (!vhlo.tensor_v1<?x!vhlo.f32_v1>) -> !vhlo.tensor_v1<?x!vhlo.f32_v1>
-+  // CHECK: "vhlo.abs_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<?x!vhlo.f32_v1>) -> !vhlo.tensor_v1<?x!vhlo.f32_v1>
-   %0 = "stablehlo.abs"(%arg0) : (tensor<?xf32>) -> tensor<?xf32>
-   func.return %0 : tensor<?xf32>
- }
- 
- // CHECK-LABEL: "type_quantization"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_quantization(%arg0: tensor<!quant.uniform<i8:f32, 34.0:16>>, %arg1: tensor<!quant.uniform<i8:f32, 34.0:16>>) -> tensor<!quant.uniform<i8:f32, 34.0:16>> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>, !vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>) -> !vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>, !vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>) -> !vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<!quant.uniform<i8:f32, 34.0:16>>, tensor<!quant.uniform<i8:f32, 34.0:16>>) -> tensor<!quant.uniform<i8:f32, 34.0:16>>
-   func.return %0 : tensor<!quant.uniform<i8:f32, 34.0:16>>
- }
- 
- //       CHECK: function_type = #vhlo.type_v1<!vhlo.func_v1<(!vhlo.token_v1) -> !vhlo.token_v1>>
- // CHECK-LABEL: "type_token_callee"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @type_token_callee(%arg0: !stablehlo.token) -> !stablehlo.token {
--  // CHECK: "vhlo.return_v1"(%arg0) : (!vhlo.token_v1) -> ()
-+  // CHECK: "vhlo.return_v1"(%[[ARG0]]) : (!vhlo.token_v1) -> ()
-   return %arg0 : !stablehlo.token
- }
- 
- //       CHECK: function_type = #vhlo.type_v1<!vhlo.func_v1<(!vhlo.token_v1) -> !vhlo.token_v1>>
- // CHECK-LABEL: "type_token_caller"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @type_token_caller(%arg0: !stablehlo.token) -> !stablehlo.token {
--  // CHECK:      "vhlo.call_v1"(%arg0) <{callee = #vhlo.string_v1<"type_token_callee">}
-+  // CHECK:      "vhlo.call_v1"(%[[ARG0]]) <{callee = #vhlo.string_v1<"type_token_callee">}
-   // CHECK-SAME: (!vhlo.token_v1) -> !vhlo.token_v1
-   %0 = func.call @type_token_callee(%arg0) : (!stablehlo.token) -> !stablehlo.token
-   return %0 : !stablehlo.token
- }
- 
- // CHECK-LABEL: "type_tuple"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @type_tuple(%arg0: tuple<tensor<f32>>) -> tuple<!stablehlo.token> {
-   %0 = "stablehlo.custom_call"(%arg0) {
-     call_target_name = "foo"
-diff --ruN a/stablehlo/stablehlo/tests/vhlo/stablehlo_legalize_to_vhlo.0_13_0.mlir b/stablehlo/stablehlo/tests/vhlo/stablehlo_legalize_to_vhlo.0_13_0.mlir
---- stablehlo/stablehlo/tests/vhlo/stablehlo_legalize_to_vhlo.0_13_0.mlir
-+++ stablehlo/stablehlo/tests/vhlo/stablehlo_legalize_to_vhlo.0_13_0.mlir
-@@ -13,6 +13,7 @@
- // ============ ATTRIBUTES ============
- 
- // CHECK-LABEL: "attr_comparison_direction_eq"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_direction_eq(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     // CHECK: comparison_direction = #vhlo<comparison_direction_v1 EQ>
-@@ -22,6 +23,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_direction_ne"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_direction_ne(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     // CHECK: comparison_direction = #vhlo<comparison_direction_v1 NE>
-@@ -31,6 +33,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_direction_ge"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_direction_ge(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     // CHECK: comparison_direction = #vhlo<comparison_direction_v1 GE>
-@@ -40,6 +43,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_direction_gt"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_direction_gt(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     // CHECK: comparison_direction = #vhlo<comparison_direction_v1 GT>
-@@ -49,6 +53,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_direction_le"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_direction_le(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     // CHECK: comparison_direction = #vhlo<comparison_direction_v1 LE>
-@@ -58,6 +63,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_direction_lt"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_direction_lt(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     // CHECK: comparison_direction = #vhlo<comparison_direction_v1 LT>
-@@ -67,6 +73,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_type_notype"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_type_notype(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     comparison_direction = #stablehlo<comparison_direction EQ>
-@@ -76,6 +83,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_type_float"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_type_float(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     comparison_direction = #stablehlo<comparison_direction EQ>,
-@@ -86,6 +94,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_type_totalorder"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_type_totalorder(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     comparison_direction = #stablehlo<comparison_direction EQ>,
-@@ -96,6 +105,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_type_signed"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_type_signed(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     comparison_direction = #stablehlo<comparison_direction EQ>,
-@@ -106,6 +116,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_type_unsigned"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_type_unsigned(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     comparison_direction = #stablehlo<comparison_direction EQ>,
-@@ -118,6 +129,7 @@
- // ConvDimensionNumbers aka #stablehlo.conv is covered below.
- 
- // CHECK-LABEL: "attr_custom_call_api_version_unspecified"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_custom_call_api_version_unspecified(%arg0: tensor<f32>) -> tensor<f32> {
-   %0 = "stablehlo.custom_call"(%arg0) {
-     call_target_name = "foo",
-@@ -128,6 +140,7 @@
- }
- 
- // CHECK-LABEL: "attr_custom_call_api_version_original"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_custom_call_api_version_original(%arg0: tensor<f32>) -> tensor<f32> {
-   %0 = "stablehlo.custom_call"(%arg0) {
-     call_target_name = "foo",
-@@ -138,6 +151,7 @@
- }
- 
- // CHECK-LABEL: "attr_custom_call_api_version_status_returning"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_custom_call_api_version_status_returning(%arg0: tensor<f32>) -> tensor<f32> {
-   %0 = "stablehlo.custom_call"(%arg0) {
-     call_target_name = "foo",
-@@ -148,6 +162,7 @@
- }
- 
- // CHECK-LABEL: "attr_custom_call_api_version_status_returning_unified"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_custom_call_api_version_status_returning_unified(%arg0: tensor<f32>) -> tensor<f32> {
-   %0 = "stablehlo.custom_call"(%arg0) {
-     call_target_name = "foo",
-@@ -160,6 +175,7 @@
- // DotDimensionNumbers aka #stablehlo.dot is covered below.
- 
- // CHECK-LABEL: "attr_fft_type_fft"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_fft_type_fft(%arg0: tensor<16xcomplex<f32>>) -> tensor<16xcomplex<f32>> {
-   %0 = "stablehlo.fft"(%arg0) {
-     // CHECK: fft_type = #vhlo<fft_type_v1 FFT>
-@@ -170,6 +186,7 @@
- }
- 
- // CHECK-LABEL: "attr_fft_type_ifft"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_fft_type_ifft(%arg0: tensor<16xcomplex<f32>>) -> tensor<16xcomplex<f32>> {
-   %0 = "stablehlo.fft"(%arg0) {
-     // CHECK: fft_type = #vhlo<fft_type_v1 IFFT>
-@@ -180,6 +197,7 @@
- }
- 
- // CHECK-LABEL: "attr_fft_type_rfft"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_fft_type_rfft(%arg0: tensor<16xf32>) -> tensor<9xcomplex<f32>> {
-   %0 = "stablehlo.fft"(%arg0) {
-     // CHECK: fft_type = #vhlo<fft_type_v1 RFFT>
-@@ -190,6 +208,7 @@
- }
- 
- // CHECK-LABEL: "attr_fft_type_irfft"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_fft_type_irfft(%arg0: tensor<9xcomplex<f32>>) -> tensor<16xf32> {
-   %0 = "stablehlo.fft"(%arg0) {
-     // CHECK: fft_type = #vhlo<fft_type_v1 IRFFT>
-@@ -202,6 +221,7 @@
- // GatherDimensionNumbers aka #stablehlo.gather is covered below.
- 
- // CHECK-LABEL: "attr_precision_config_default"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_precision_config_default(%arg0: tensor<8x16xf32>, %arg1: tensor<16x8xf32>) -> tensor<8x8xf32> {
-   %0 = "stablehlo.dot"(%arg0, %arg1) {
-     // CHECK: precision_config = #vhlo.array_v1<[#vhlo<precision_v1 DEFAULT>, #vhlo<precision_v1 DEFAULT>]>
-@@ -210,6 +230,7 @@
- }
- 
- // CHECK-LABEL: "attr_precision_config_high"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_precision_config_high(%arg0: tensor<8x16xf32>, %arg1: tensor<16x8xf32>) -> tensor<8x8xf32> {
-   %0 = "stablehlo.dot"(%arg0, %arg1) {
-     // CHECK: precision_config = #vhlo.array_v1<[#vhlo<precision_v1 HIGH>, #vhlo<precision_v1 HIGH>]>
-@@ -219,6 +240,7 @@
- }
- 
- // CHECK-LABEL: "attr_precision_config_highest"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_precision_config_highest(%arg0: tensor<8x16xf32>, %arg1: tensor<16x8xf32>) -> tensor<8x8xf32> {
-   %0 = "stablehlo.dot"(%arg0, %arg1) {
-     // CHECK: precision_config = #vhlo.array_v1<[#vhlo<precision_v1 HIGHEST>, #vhlo<precision_v1 HIGHEST>]>
-@@ -228,6 +250,7 @@
- }
- 
- // CHECK-LABEL: "attr_rng_algorithm_default"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_rng_algorithm_default(%arg0: tensor<f32>) -> (tensor<f32>, tensor<f32>) {
-   %0:2 = "stablehlo.rng_bit_generator"(%arg0) {
-     // CHECK: rng_algorithm = #vhlo<rng_algorithm_v1 DEFAULT>
-@@ -237,6 +260,7 @@
- }
- 
- // CHECK-LABEL: "attr_rng_algorithm_three_fry"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_rng_algorithm_three_fry(%arg0: tensor<f32>) -> (tensor<f32>, tensor<f32>) {
-   %0:2 = "stablehlo.rng_bit_generator"(%arg0) {
-     // CHECK: rng_algorithm = #vhlo<rng_algorithm_v1 THREE_FRY>
-@@ -246,6 +270,7 @@
- }
- 
- // CHECK-LABEL: "attr_rng_algorithm_philox"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_rng_algorithm_philox(%arg0: tensor<f32>) -> (tensor<f32>, tensor<f32>) {
-   %0:2 = "stablehlo.rng_bit_generator"(%arg0) {
-     // CHECK: rng_algorithm = #vhlo<rng_algorithm_v1 PHILOX>
-@@ -255,6 +280,7 @@
- }
- 
- // CHECK-LABEL: "attr_rng_distribution_uniform"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @attr_rng_distribution_uniform(%arg0: tensor<f32>, %arg1: tensor<f32>, %arg2: tensor<0xindex>) -> tensor<f32> {
-   %0 = "stablehlo.rng"(%arg0, %arg1, %arg2) {
-     // CHECK: rng_distribution = #vhlo<rng_distribution_v1 UNIFORM>
-@@ -264,6 +290,7 @@
- }
- 
- // CHECK-LABEL: "attr_rng_distribution_normal"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @attr_rng_distribution_normal(%arg0: tensor<f32>, %arg1: tensor<f32>, %arg2: tensor<0xindex>) -> tensor<f32> {
-   %0 = "stablehlo.rng"(%arg0, %arg1, %arg2) {
-     // CHECK: rng_distribution = #vhlo<rng_distribution_v1 NORMAL>
-@@ -275,6 +302,7 @@
- // ScatterDimensionNumbers aka #stablehlo.scatter is covered below.
- 
- // CHECK-LABEL: "attr_transpose_no_transpose"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_transpose_no_transpose(%arg0: tensor<16x16xf32>, %arg1: tensor<16x16xf32>) ->  tensor<16x16xf32> {
-   %0 = "stablehlo.triangular_solve"(%arg0, %arg1) {
-     left_side = true,
-@@ -287,6 +315,7 @@
- }
- 
- // CHECK-LABEL: "attr_transpose_transpose"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_transpose_transpose(%arg0: tensor<16x16xf32>, %arg1: tensor<16x16xf32>) ->  tensor<16x16xf32> {
-   %0 = "stablehlo.triangular_solve"(%arg0, %arg1) {
-     left_side = true,
-@@ -299,6 +328,7 @@
- }
- 
- // CHECK-LABEL: "attr_transpose_adjoint"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_transpose_adjoint(%arg0: tensor<16x16xf32>, %arg1: tensor<16x16xf32>) ->  tensor<16x16xf32> {
-   %0 = "stablehlo.triangular_solve"(%arg0, %arg1) {
-     left_side = true,
-@@ -313,10 +343,9 @@
- // TypeExtensionsAttr aka #stablehlo.type_extensions is covered below.
- 
- // CHECK-LABEL: "attr_type_extensions_bounds"
--func.func @attr_type_extensions_bounds(
--    %arg0: tensor<?x?xf32, #stablehlo.type_extensions<bounds = [16, ?]>>)
--    -> tensor<?x?xf32, #stablehlo.type_extensions<bounds = [16, ?]>> {
--  // CHECK: "vhlo.return_v1"(%arg0) : (!vhlo.tensor_v1<?x?x!vhlo.f32_v1, #vhlo.type_extensions_v1<bounds = [16, ?]>>) -> ()
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @attr_type_extensions_bounds(%arg0: tensor<?x?xf32, #stablehlo.type_extensions<bounds = [16, ?]>>) -> tensor<?x?xf32, #stablehlo.type_extensions<bounds = [16, ?]>> {
-+  // CHECK: "vhlo.return_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<?x?x!vhlo.f32_v1, #vhlo.type_extensions_v1<bounds = [16, ?]>>) -> ()
-   func.return %arg0 : tensor<?x?xf32, #stablehlo.type_extensions<bounds = [16, ?]>>
- }
- 
-@@ -324,8 +353,9 @@
- // ============ DEFAULTS ============
- 
- // CHECK-LABEL: "default_all_gather"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_all_gather(%arg0: tensor<16x8xf32>) -> tensor<16x16xf32> {
--  //               CHECK: "vhlo.all_gather_v1"(%arg0) <{
-+  //               CHECK: "vhlo.all_gather_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   all_gather_dim = #vhlo.integer_v1<1 : i64>
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0], [1]]> : tensor<2x1xi64>>,
-@@ -339,8 +369,9 @@
- }
- 
- // CHECK-LABEL: "default_all_reduce"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_all_reduce(%arg0: tensor<f32>) -> tensor<f32> {
--  //               CHECK: "vhlo.all_reduce_v1"(%arg0)
-+  //               CHECK: "vhlo.all_reduce_v1"(%[[ARG0]])
-   //          CHECK-SAME: <{
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0], [1]]> : tensor<2x1xi64>>,
-@@ -362,8 +393,9 @@
- }
- 
- // CHECK-LABEL: "default_all_to_all"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_all_to_all(%arg0: tensor<4x16xf32>) -> tensor<16x4xf32> {
--  //               CHECK: "vhlo.all_to_all_v1"(%arg0) <{
-+  //               CHECK: "vhlo.all_to_all_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   //          CHECK-SAME:   concat_dimension = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0, 1, 2, 3]]> : tensor<1x4xi64>>,
-@@ -380,8 +412,9 @@
- }
- 
- // CHECK-LABEL: "default_cholesky"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_cholesky(%arg0: tensor<1x16x16xf32>) -> tensor<1x16x16xf32> {
--  //      CHECK: "vhlo.cholesky_v1"(%arg0) <{
-+  //      CHECK: "vhlo.cholesky_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   lower = #vhlo.bool_v1<false>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<1x16x16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<1x16x16x!vhlo.f32_v1>
-   %0 = "stablehlo.cholesky"(%arg0) : (tensor<1x16x16xf32>) -> tensor<1x16x16xf32>
-@@ -389,8 +422,9 @@
- }
- 
- // CHECK-LABEL: "default_collective_permute"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_collective_permute(%arg0: tensor<16x8xf32>) -> tensor<16x8xf32> {
--  //               CHECK: "vhlo.collective_permute_v1"(%arg0) <{
-+  //               CHECK: "vhlo.collective_permute_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME{LITERAL}:   source_target_pairs = #vhlo.tensor_v1<dense<[[0, 1], [1, 2], [2, 3]]> : tensor<3x2xi64>>
-   //          CHECK-SAME: }> : (!vhlo.tensor_v1<16x8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x8x!vhlo.f32_v1>
-@@ -401,8 +435,9 @@
- }
- 
- // CHECK-LABEL: "default_compare"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @default_compare(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
--  //      CHECK: "vhlo.compare_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.compare_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   compare_type = #vhlo<comparison_type_v1 NOTYPE>,
-   // CHECK-SAME:   comparison_direction = #vhlo<comparison_direction_v1 EQ>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-@@ -413,8 +448,9 @@
- }
- 
- // CHECK-LABEL: "default_convolution"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @default_convolution(%arg0: tensor<1x8x8x207xf32>, %arg1: tensor<3x3x207x16xf32>) -> tensor<1x6x6x16xf32> {
--  //      CHECK: "vhlo.convolution_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.convolution_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   batch_group_count = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   feature_group_count = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   input_batch_dimension = #vhlo.integer_v1<0 : i64>,
-@@ -442,8 +478,9 @@
- }
- 
- // CHECK-LABEL: "default_custom_call"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_custom_call(%arg0: tensor<f32>) -> tensor<f32> {
--  //      CHECK: "vhlo.custom_call_v1"(%arg0) <{
-+  //      CHECK: "vhlo.custom_call_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   api_version = #vhlo<api_version_v1 API_VERSION_ORIGINAL>,
-   // CHECK-SAME:   backend_config = #vhlo.string_v1<"">,
-   // CHECK-SAME:   call_target_name = #vhlo.string_v1<"foo">,
-@@ -460,8 +497,9 @@
- }
- 
- // CHECK-LABEL: "default_dot_general"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @default_dot_general(%arg0: tensor<8x8x16xf32>, %arg1: tensor<8x16x8xf32>) -> tensor<8x8x8xf32> {
--  //      CHECK: "vhlo.dot_general_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.dot_general_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   lhs_batching_dimensions = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>,
-   // CHECK-SAME:   lhs_contracting_dimensions = #vhlo.tensor_v1<dense<2> : tensor<1xi64>>,
-   // CHECK-SAME:   precision_config = #vhlo.array_v1<[#vhlo<precision_v1 DEFAULT>, #vhlo<precision_v1 DEFAULT>]>,
-@@ -480,8 +518,9 @@
- }
- 
- // CHECK-LABEL: "default_dot"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @default_dot(%arg0: tensor<8x16xf32>, %arg1: tensor<16x8xf32>) -> tensor<8x8xf32> {
--  //      CHECK: "vhlo.dot_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.dot_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   precision_config = #vhlo.array_v1<[#vhlo<precision_v1 DEFAULT>, #vhlo<precision_v1 DEFAULT>]>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<8x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<8x8x!vhlo.f32_v1>
-   %0 = "stablehlo.dot"(%arg0, %arg1) : (tensor<8x16xf32>, tensor<16x8xf32>) -> tensor<8x8xf32>
-@@ -489,8 +528,9 @@
- }
- 
- // CHECK-LABEL: "default_dynamic_broadcast_in_dim"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @default_dynamic_broadcast_in_dim(%arg0: tensor<?x?xf32>, %arg1: tensor<2xindex>) -> tensor<?x?xf32> {
--  //      CHECK: "vhlo.dynamic_broadcast_in_dim_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.dynamic_broadcast_in_dim_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   broadcast_dimensions = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-   // CHECK-SAME:   known_expanding_dimensions = #vhlo.tensor_v1<dense<> : tensor<0xi64>>,
-   // CHECK-SAME:   known_nonexpanding_dimensions = #vhlo.tensor_v1<dense<> : tensor<0xi64>>
-@@ -502,8 +542,9 @@
- }
- 
- // CHECK-LABEL: "default_dynamic_conv"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @default_dynamic_conv(%arg0: tensor<1x8x8x207xf32>, %arg1: tensor<3x3x207x16xf32>, %arg2: tensor<4xi32>) -> tensor<1x?x?x16xf32> {
--  //      CHECK: "vhlo.dynamic_conv_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.dynamic_conv_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   batch_group_count = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   feature_group_count = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   input_batch_dimension = #vhlo.integer_v1<0 : i64>,
-@@ -531,8 +572,9 @@
- }
- 
- // CHECK-LABEL: "default_dynamic_gather"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @default_dynamic_gather(%arg0 : tensor<2x4x9xf32>, %arg1 : tensor<1x5x2xi32>, %arg2 : tensor<3xi32>) -> tensor<1x5x8xf32> {
--  //      CHECK: "vhlo.dynamic_gather_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.dynamic_gather_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   collapsed_slice_dims = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-   // CHECK-SAME:   index_vector_dim = #vhlo.integer_v1<2 : i64>,
-   // CHECK-SAME:   indices_are_sorted = #vhlo.bool_v1<false>,
-@@ -558,15 +600,16 @@
-   // CHECK-SAME:   sym_name = #vhlo.string_v1<"default_func">,
-   // CHECK-SAME:   sym_visibility = #vhlo.string_v1<"">
-   // CHECK-SAME: }> ({
--  // CHECK-NEXT:   ^[[BB:bb.*]](%arg0: !vhlo.tensor_v1<!vhlo.f32_v1>):
--  // CHECK-NEXT:     "vhlo.return_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-+  // CHECK-NEXT:   ^[[BB:bb.*]](%[[ARG0:.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
-+  // CHECK-NEXT:     "vhlo.return_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-   // CHECK-NEXT: }) : () -> ()
-   func.return %arg0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "dynamic_gather"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @dynamic_gather(%arg0 : tensor<2x4x9xf32>, %arg1 : tensor<1x5x2xi32>) -> tensor<1x5x1xf32> {
--  //      CHECK: "vhlo.gather_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.gather_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   collapsed_slice_dims = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-   // CHECK-SAME:   index_vector_dim = #vhlo.integer_v1<2 : i64>,
-   // CHECK-SAME:   indices_are_sorted = #vhlo.bool_v1<false>,
-@@ -587,8 +630,9 @@
- }
- 
- // CHECK-LABEL: "default_infeed"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_infeed(%arg0: !stablehlo.token) -> (tensor<f32>, !stablehlo.token) {
--  //               CHECK: "vhlo.infeed_v1"(%arg0) <{
-+  //               CHECK: "vhlo.infeed_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   infeed_config = #vhlo.string_v1<"">,
-   // CHECK-SAME{LITERAL}:   layout = #vhlo.array_v1<[]>
-   //          CHECK-SAME: }> : (!vhlo.token_v1) -> (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.token_v1)
-@@ -597,8 +641,9 @@
- }
- 
- // CHECK-LABEL: "default_outfeed"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @default_outfeed(%arg0: tensor<f32>, %arg1: !stablehlo.token) -> !stablehlo.token {
--  //      CHECK: "vhlo.outfeed_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.outfeed_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   outfeed_config = #vhlo.string_v1<"">
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.token_v1) -> !vhlo.token_v1
-   %0 = "stablehlo.outfeed"(%arg0, %arg1) : (tensor<f32>, !stablehlo.token) -> !stablehlo.token
-@@ -606,8 +651,9 @@
- }
- 
- // CHECK-LABEL: "default_recv"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_recv(%arg0: !stablehlo.token) -> (tensor<f32>, !stablehlo.token) {
--  //      CHECK: "vhlo.recv_v1"(%arg0) <{
-+  //      CHECK: "vhlo.recv_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME:   channel_type = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   is_host_transfer = #vhlo.bool_v1<false>
-@@ -619,8 +665,9 @@
- }
- 
- // CHECK-LABEL: "default_send"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @default_send(%arg0: tensor<f32>, %arg1: !stablehlo.token) -> !stablehlo.token {
--  //      CHECK: "vhlo.send_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.send_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME:   channel_type = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   is_host_transfer = #vhlo.bool_v1<false>
-@@ -632,8 +679,9 @@
- }
- 
- // CHECK-LABEL: "default_reduce_scatter"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_reduce_scatter(%arg0: tensor<16xf32>) -> tensor<16xf32> {
--  //               CHECK: "vhlo.reduce_scatter_v1"(%arg0) <{
-+  //               CHECK: "vhlo.reduce_scatter_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0], [1]]> : tensor<2x1xi64>>,
-   //          CHECK-SAME:   scatter_dimension = #vhlo.integer_v1<0 : i64>
-@@ -655,8 +703,9 @@
- }
- 
- // CHECK-LABEL: "default_reduce_window"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @default_reduce_window(%arg0: tensor<2x17x31x7xf32>, %arg1: tensor<f32>) -> tensor<2x16x30x7xf32> {
--  //               CHECK: "vhlo.reduce_window_v1"(%arg0, %arg1)  <{
-+  //               CHECK: "vhlo.reduce_window_v1"(%[[ARG0]], %[[ARG1]])  <{
-   //          CHECK-SAME:   base_dilations = #vhlo.tensor_v1<dense<1> : tensor<4xi64>>,
-   // CHECK-SAME{LITERAL}:   padding = #vhlo.tensor_v1<dense<0> : tensor<4x2xi64>>,
-   //          CHECK-SAME:   window_dilations = #vhlo.tensor_v1<dense<1> : tensor<4xi64>>,
-@@ -678,8 +727,9 @@
- }
- 
- // CHECK-LABEL: "default_scatter"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @default_scatter(%arg0: tensor<200x100x300xf32>, %arg1: tensor<10x2xi32>, %arg2: tensor<10x300xf32>) -> tensor<200x100x300xf32> {
--  //      CHECK: "vhlo.scatter_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.scatter_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   index_vector_dim = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   indices_are_sorted = #vhlo.bool_v1<false>,
-   // CHECK-SAME:   inserted_window_dims = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-@@ -707,8 +757,9 @@
- }
- 
- // CHECK-LABEL: "default_select_and_scatter"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @default_select_and_scatter(%arg0: tensor<10x24x24x64xf32>, %arg1: tensor<10x23x23x64xf32>, %arg2: tensor<f32>) -> tensor<10x24x24x64xf32> {
--  //      CHECK: "vhlo.select_and_scatter_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.select_and_scatter_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   padding = #vhlo.tensor_v1<dense<0> : tensor<4x2xi64>>,
-   // CHECK-SAME:   window_dimensions = #vhlo.tensor_v1<dense<[1, 2, 2, 1]> : tensor<4xi64>>,
-   // CHECK-SAME:   window_strides = #vhlo.tensor_v1<dense<1> : tensor<4xi64>>
-@@ -736,8 +787,9 @@
- }
- 
- // CHECK-LABEL: "default_sort"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_sort(%arg0: tensor<16xf32>) -> tensor<16xf32> {
--  //      CHECK: "vhlo.sort_v1"(%arg0) <{
-+  //      CHECK: "vhlo.sort_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   dimension = #vhlo.integer_v1<-1 : i64>
-   // CHECK-SAME:   is_stable = #vhlo.bool_v1<false>
-   // CHECK-SAME: }> ({
-@@ -756,29 +808,33 @@
- // ============ OPS ============
- 
- // CHECK-LABEL: "op_abs"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_abs(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.abs_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.abs_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.abs"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_add"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_add(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_after_all"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_after_all(%arg0: !stablehlo.token) -> !stablehlo.token {
--  // CHECK: "vhlo.after_all_v1"(%arg0) : (!vhlo.token_v1) -> !vhlo.token_v1
-+  // CHECK: "vhlo.after_all_v1"(%[[ARG0]]) : (!vhlo.token_v1) -> !vhlo.token_v1
-   %0 = "stablehlo.after_all"(%arg0) : (!stablehlo.token) -> !stablehlo.token
-   func.return %0 : !stablehlo.token
- }
- 
- // CHECK-LABEL: "op_all_gather"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_all_gather(%arg0: tensor<16x8xf32>) -> tensor<16x16xf32> {
--  //               CHECK: "vhlo.all_gather_v1"(%arg0) <{
-+  //               CHECK: "vhlo.all_gather_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   all_gather_dim = #vhlo.integer_v1<1 : i64>
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0], [1]]> : tensor<2x1xi64>>,
-@@ -794,8 +850,9 @@
- }
- 
- // CHECK-LABEL: "op_all_reduce"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_all_reduce(%arg0: tensor<f32>) -> tensor<f32> {
--  //               CHECK: "vhlo.all_reduce_v1"(%arg0) <{
-+  //               CHECK: "vhlo.all_reduce_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0], [1]]> : tensor<2x1xi64>>,
-   //          CHECK-SAME:   use_global_device_ids = #vhlo.bool_v1<true>
-@@ -817,8 +874,9 @@
- }
- 
- // CHECK-LABEL: "op_all_to_all"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_all_to_all(%arg0: tensor<4x16xf32>) -> tensor<16x4xf32> {
--  //               CHECK: "vhlo.all_to_all_v1"(%arg0) <{
-+  //               CHECK: "vhlo.all_to_all_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<1 : i64>,
-   //          CHECK-SAME:   concat_dimension = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0, 1, 2, 3]]> : tensor<1x4xi64>>,
-@@ -836,22 +894,25 @@
- }
- 
- // CHECK-LABEL: "op_and"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_and(%arg0: tensor<i1>, %arg1: tensor<i1>) -> tensor<i1> {
--  // CHECK: "vhlo.and_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-+  // CHECK: "vhlo.and_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-   %0 = "stablehlo.and"(%arg0, %arg1) : (tensor<i1>, tensor<i1>) -> tensor<i1>
-   func.return %0 : tensor<i1>
- }
- 
- // CHECK-LABEL: "op_atan2"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_atan2(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.atan2_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.atan2_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.atan2"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_batch_norm_grad"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}}, %[[ARG3:.*]]: {{.*}}, %[[ARG4:.*]]: {{.*}})
- func.func @op_batch_norm_grad(%arg0: tensor<16x16x16x16xf32>, %arg1: tensor<16xf32>, %arg2: tensor<16xf32>, %arg3: tensor<16xf32>, %arg4: tensor<16x16x16x16xf32>) -> (tensor<16x16x16x16xf32>, tensor<16xf32>, tensor<16xf32>) {
--  //      CHECK: "vhlo.batch_norm_grad_v1"(%arg0, %arg1, %arg2, %arg3, %arg4) <{
-+  //      CHECK: "vhlo.batch_norm_grad_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]], %[[ARG3]], %[[ARG4]]) <{
-   // CHECK-SAME:   epsilon = #vhlo.float_v1<1.000000e-03 : !vhlo.f32_v1>,
-   // CHECK-SAME:   feature_index = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<16x16x16x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x16x16x16x!vhlo.f32_v1>) -> (!vhlo.tensor_v1<16x16x16x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>)
-@@ -863,8 +924,9 @@
- }
- 
- // CHECK-LABEL: "op_batch_norm_inference"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}}, %[[ARG3:.*]]: {{.*}}, %[[ARG4:.*]]: {{.*}})
- func.func @op_batch_norm_inference(%arg0: tensor<16x16x16x16xf32>, %arg1: tensor<16xf32>, %arg2: tensor<16xf32>, %arg3: tensor<16xf32>, %arg4: tensor<16xf32>) -> tensor<16x16x16x16xf32> {
--  //      CHECK: "vhlo.batch_norm_inference_v1"(%arg0, %arg1, %arg2, %arg3, %arg4) <{
-+  //      CHECK: "vhlo.batch_norm_inference_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]], %[[ARG3]], %[[ARG4]]) <{
-   // CHECK-SAME:   epsilon = #vhlo.float_v1<1.000000e-03 : !vhlo.f32_v1>,
-   // CHECK-SAME:   feature_index = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<16x16x16x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x16x16x16x!vhlo.f32_v1>
-@@ -876,8 +938,9 @@
- }
- 
- // CHECK-LABEL: "op_batch_norm_training"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_batch_norm_training(%arg0: tensor<16x16x16x16xf32>, %arg1: tensor<16xf32>, %arg2: tensor<16xf32>) -> (tensor<16x16x16x16xf32>, tensor<16xf32>, tensor<16xf32>) {
--  //      CHECK: "vhlo.batch_norm_training_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.batch_norm_training_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   epsilon = #vhlo.float_v1<1.000000e-03 : !vhlo.f32_v1>,
-   // CHECK-SAME:   feature_index = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<16x16x16x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>) -> (!vhlo.tensor_v1<16x16x16x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>)
-@@ -889,15 +952,17 @@
- }
- 
- // CHECK-LABEL: "op_bitcast_convert"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_bitcast_convert(%arg0: tensor<i32>) -> tensor<f32> {
--  // CHECK: "vhlo.bitcast_convert_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.bitcast_convert_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.bitcast_convert"(%arg0) : (tensor<i32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_broadcast_in_dim"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_broadcast_in_dim(%arg0: tensor<16xf32>) -> tensor<16x16xf32> {
--  //      CHECK: "vhlo.broadcast_in_dim_v1"(%arg0) <{
-+  //      CHECK: "vhlo.broadcast_in_dim_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   broadcast_dimensions = #vhlo.tensor_v1<dense<1> : tensor<1xi64>>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x16x!vhlo.f32_v1>
-   %0 = "stablehlo.broadcast_in_dim"(%arg0) {
-@@ -907,8 +972,9 @@
- }
- 
- // CHECK-LABEL: "op_broadcast"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_broadcast(%arg0: tensor<16xf32>) -> tensor<16x16xf32> {
--  //      CHECK: "vhlo.broadcast_v1"(%arg0) <{
-+  //      CHECK: "vhlo.broadcast_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   broadcast_sizes = #vhlo.tensor_v1<dense<16> : tensor<1xi64>>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x16x!vhlo.f32_v1>
-   %0 = "stablehlo.broadcast"(%arg0) {
-@@ -918,9 +984,10 @@
- }
- 
- // CHECK-LABEL: "op_case"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_case(%arg0: tensor<i32>, %arg1: tensor<f32>) -> tensor<f32> {
--  //      CHECK: "vhlo.case_v1"(%arg0) ({
--  // CHECK-NEXT:   "vhlo.return_v1"(%arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-+  //      CHECK: "vhlo.case_v1"(%[[ARG0]]) ({
-+  // CHECK-NEXT:   "vhlo.return_v1"(%[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-   // CHECK-NEXT: }) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.case"(%arg0) ({
-     "stablehlo.return"(%arg1) : (tensor<f32>) -> ()
-@@ -929,22 +996,25 @@
- }
- 
- // CHECK-LABEL: "op_cbrt"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_cbrt(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.cbrt_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.cbrt_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.cbrt"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_ceil"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_ceil(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.ceil_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.ceil_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.ceil"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_cholesky"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_cholesky(%arg0: tensor<1x16x16xf32>) -> tensor<1x16x16xf32> {
--  //      CHECK: "vhlo.cholesky_v1"(%arg0) <{
-+  //      CHECK: "vhlo.cholesky_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   lower = #vhlo.bool_v1<true>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<1x16x16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<1x16x16x!vhlo.f32_v1>
-   %0 = "stablehlo.cholesky"(%arg0) {
-@@ -954,22 +1024,25 @@
- }
- 
- // CHECK-LABEL: "op_clamp"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_clamp(%arg0: tensor<f32>, %arg1: tensor<f32>, %arg2: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.clamp_v1"(%arg0, %arg1, %arg2) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.clamp_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.clamp"(%arg0, %arg1, %arg2) : (tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_count_leading_zeros"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_count_leading_zeros(%arg0: tensor<i32>) -> tensor<i32> {
--  // CHECK: "vhlo.count_leading_zeros_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-+  // CHECK: "vhlo.count_leading_zeros_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-   %0 = "stablehlo.count_leading_zeros"(%arg0) : (tensor<i32>) -> tensor<i32>
-   func.return %0 : tensor<i32>
- }
- 
- // CHECK-LABEL: "op_collective_permute"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_collective_permute(%arg0: tensor<16x8xf32>) -> tensor<16x8xf32> {
--  //               CHECK: "vhlo.collective_permute_v1"(%arg0) <{
-+  //               CHECK: "vhlo.collective_permute_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME{LITERAL}:   source_target_pairs = #vhlo.tensor_v1<dense<[[0, 1], [1, 2], [2, 3]]> : tensor<3x2xi64>>
-   //          CHECK-SAME: }> : (!vhlo.tensor_v1<16x8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x8x!vhlo.f32_v1>
-@@ -981,8 +1054,9 @@
- }
- 
- // CHECK-LABEL: "op_compare"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_compare(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
--  //      CHECK: "vhlo.compare_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.compare_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   compare_type = #vhlo<comparison_type_v1 TOTALORDER>,
-   // CHECK-SAME:   comparison_direction = #vhlo<comparison_direction_v1 EQ>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-@@ -994,15 +1068,17 @@
- }
- 
- // CHECK-LABEL: "op_complex"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_complex(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<complex<f32>> {
--  // CHECK: "vhlo.complex_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>
-+  // CHECK: "vhlo.complex_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>
-   %0 = "stablehlo.complex"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<complex<f32>>
-   func.return %0 : tensor<complex<f32>>
- }
- 
- // CHECK-LABEL: "op_concatenate"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_concatenate(%arg0: tensor<8xf32>, %arg1: tensor<8xf32>) -> tensor<16xf32> {
--  //      CHECK: "vhlo.concatenate_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.concatenate_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   dimension = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<8x!vhlo.f32_v1>, !vhlo.tensor_v1<8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x!vhlo.f32_v1>
-   %0 = "stablehlo.concatenate"(%arg0, %arg1) {
-@@ -1012,6 +1088,7 @@
- }
- 
- // CHECK-LABEL: "op_constant"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_constant(%arg0: tensor<f32>) -> tensor<f32> {
-   //      CHECK: "vhlo.constant_v1"() <{
-   // CHECK-SAME:   value = #vhlo.tensor_v1<dense<0.000000e+00> : tensor<f32>>
-@@ -1023,15 +1100,17 @@
- }
- 
- // CHECK-LABEL: "op_convert"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_convert(%arg0: tensor<i32>) -> tensor<f32> {
--  // CHECK: "vhlo.convert_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.convert_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.convert"(%arg0) : (tensor<i32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_convolution"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_convolution(%arg0: tensor<1x8x8x207xf32>, %arg1: tensor<3x3x207x16xf32>) -> tensor<1x7x7x16xf32> {
--  //      CHECK: "vhlo.convolution_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.convolution_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   batch_group_count = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   feature_group_count = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   input_batch_dimension = #vhlo.integer_v1<0 : i64>,
-@@ -1065,8 +1144,9 @@
- }
- 
- // CHECK-LABEL: "op_cosine"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_cosine(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.cosine_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.cosine_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.cosine"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
-@@ -1079,8 +1159,9 @@
- }
- 
- // CHECK-LABEL: "op_cross_replica_sum"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_cross_replica_sum(%arg0: tensor<f32>) -> tensor<f32> {
--  //               CHECK: "vhlo.cross-replica-sum_v1"(%arg0) <{
-+  //               CHECK: "vhlo.cross-replica-sum_v1"(%[[ARG0]]) <{
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0], [1]]> : tensor<2x1xi64>>
-   //          CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.cross-replica-sum"(%arg0) {
-@@ -1090,8 +1171,9 @@
- }
- 
- // CHECK-LABEL: "op_custom_call"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_custom_call(%arg0: tensor<f32>) -> tensor<f32> {
--  //      CHECK: "vhlo.custom_call_v1"(%arg0) <{
-+  //      CHECK: "vhlo.custom_call_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   api_version = #vhlo<api_version_v1 API_VERSION_STATUS_RETURNING>,
-   // CHECK-SAME:   backend_config = #vhlo.string_v1<"\08\03\1A\02">,
-   // CHECK-SAME:   call_target_name = #vhlo.string_v1<"foo">,
-@@ -1122,15 +1204,17 @@
- }
- 
- // CHECK-LABEL: "op_divide"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_divide(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.divide_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.divide_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.divide"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_dot_general"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_dot_general(%arg0: tensor<8x8x16xf32>, %arg1: tensor<8x16x8xf32>) -> tensor<8x8x8xf32> {
--  //      CHECK: "vhlo.dot_general_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.dot_general_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   lhs_batching_dimensions = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>,
-   // CHECK-SAME:   lhs_contracting_dimensions = #vhlo.tensor_v1<dense<2> : tensor<1xi64>>,
-   // CHECK-SAME:   precision_config = #vhlo.array_v1<[#vhlo<precision_v1 HIGHEST>, #vhlo<precision_v1 HIGHEST>]>,
-@@ -1150,8 +1234,9 @@
- }
- 
- // CHECK-LABEL: "op_dot"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_dot(%arg0: tensor<8x16xf32>, %arg1: tensor<16x8xf32>) -> tensor<8x8xf32> {
--  //      CHECK: "vhlo.dot_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.dot_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   precision_config = #vhlo.array_v1<[#vhlo<precision_v1 HIGHEST>, #vhlo<precision_v1 HIGHEST>]>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<8x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<8x8x!vhlo.f32_v1>
-   %0 = "stablehlo.dot"(%arg0, %arg1) {
-@@ -1161,8 +1246,9 @@
- }
- 
- // CHECK-LABEL: "op_dynamic_broadcast_in_dim"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_dynamic_broadcast_in_dim(%arg0: tensor<?x?xf32>, %arg1: tensor<2xindex>) -> tensor<?x?xf32> {
--  //      CHECK: "vhlo.dynamic_broadcast_in_dim_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.dynamic_broadcast_in_dim_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   broadcast_dimensions = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-   // CHECK-SAME:   known_expanding_dimensions = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>,
-   // CHECK-SAME:   known_nonexpanding_dimensions = #vhlo.tensor_v1<dense<1> : tensor<1xi64>>
-@@ -1176,8 +1262,9 @@
- }
- 
- // CHECK-LABEL: "op_dynamic_conv"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_dynamic_conv(%arg0: tensor<1x8x8x207xf32>, %arg1: tensor<3x3x207x16xf32>, %arg2: tensor<4xi32>) -> tensor<1x?x?x16xf32> {
--  //      CHECK: "vhlo.dynamic_conv_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.dynamic_conv_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   batch_group_count = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   feature_group_count = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   input_batch_dimension = #vhlo.integer_v1<0 : i64>,
-@@ -1211,8 +1298,9 @@
- }
- 
- // CHECK-LABEL: "op_dynamic_gather"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_dynamic_gather(%arg0 : tensor<2x4x9xf32>, %arg1 : tensor<1x5x2xi32>, %arg2 : tensor<3xi32>) -> tensor<1x5x8xf32> {
--  //      CHECK: "vhlo.dynamic_gather_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.dynamic_gather_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   collapsed_slice_dims = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-   // CHECK-SAME:   index_vector_dim = #vhlo.integer_v1<2 : i64>,
-   // CHECK-SAME:   indices_are_sorted = #vhlo.bool_v1<true>,
-@@ -1232,8 +1320,9 @@
- }
- 
- // CHECK-LABEL: "op_dynamic_iota"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_dynamic_iota(%arg0: tensor<1xindex>) -> tensor<?xf32> {
--  //      CHECK: "vhlo.dynamic_iota_v1"(%arg0) <{
-+  //      CHECK: "vhlo.dynamic_iota_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   iota_dimension = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<1x!vhlo.index_v1>) -> !vhlo.tensor_v1<?x!vhlo.f32_v1>
-   %0 = "stablehlo.dynamic_iota"(%arg0) {
-@@ -1243,22 +1332,25 @@
- }
- 
- // CHECK-LABEL: "op_dynamic_pad"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}}, %[[ARG3:.*]]: {{.*}}, %[[ARG4:.*]]: {{.*}})
- func.func @op_dynamic_pad(%arg0: tensor<?xf32>, %arg1: tensor<f32>, %arg2: tensor<1xindex>, %arg3: tensor<1xindex>, %arg4: tensor<1xindex>) -> tensor<?xf32> {
--  // CHECK: "vhlo.dynamic_pad_v1"(%arg0, %arg1, %arg2, %arg3, %arg4) : (!vhlo.tensor_v1<?x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>) -> !vhlo.tensor_v1<?x!vhlo.f32_v1>
-+  // CHECK: "vhlo.dynamic_pad_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]], %[[ARG3]], %[[ARG4]]) : (!vhlo.tensor_v1<?x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>) -> !vhlo.tensor_v1<?x!vhlo.f32_v1>
-   %0 = "stablehlo.dynamic_pad"(%arg0, %arg1, %arg2, %arg3, %arg4) : (tensor<?xf32>, tensor<f32>, tensor<1xindex>, tensor<1xindex>, tensor<1xindex>) -> tensor<?xf32>
-   func.return %0 : tensor<?xf32>
- }
- 
- // CHECK-LABEL: "op_dynamic_reshape"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_dynamic_reshape(%arg0: tensor<16xf32>, %arg1: tensor<2xindex>) -> tensor<?x?xf32> {
--  // CHECK: "vhlo.dynamic_reshape_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<2x!vhlo.index_v1>) -> !vhlo.tensor_v1<?x?x!vhlo.f32_v1>
-+  // CHECK: "vhlo.dynamic_reshape_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<2x!vhlo.index_v1>) -> !vhlo.tensor_v1<?x?x!vhlo.f32_v1>
-   %0 = "stablehlo.dynamic_reshape"(%arg0, %arg1) : (tensor<16xf32>, tensor<2xindex>) -> tensor<?x?xf32>
-   func.return %0 : tensor<?x?xf32>
- }
- 
- // CHECK-LABEL: "op_dynamic_slice"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_dynamic_slice(%arg0: tensor<16xf32>, %arg1: tensor<i64>) -> tensor<4xf32> {
--  //      CHECK: "vhlo.dynamic_slice_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.dynamic_slice_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   slice_sizes = #vhlo.tensor_v1<dense<4> : tensor<1xi64>>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.i64_v1>) -> !vhlo.tensor_v1<4x!vhlo.f32_v1>
-   %0 = "stablehlo.dynamic_slice"(%arg0, %arg1) {
-@@ -1268,15 +1360,17 @@
- }
- 
- // CHECK-LABEL: "op_dynamic_update_slice"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_dynamic_update_slice(%arg0: tensor<16xf32>, %arg1: tensor<4xf32>, %arg2: tensor<i64>) -> tensor<16xf32> {
--  // CHECK: "vhlo.dynamic_update_slice_v1"(%arg0, %arg1, %arg2) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<4x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.i64_v1>) -> !vhlo.tensor_v1<16x!vhlo.f32_v1>
-+  // CHECK: "vhlo.dynamic_update_slice_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<4x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.i64_v1>) -> !vhlo.tensor_v1<16x!vhlo.f32_v1>
-   %0 = "stablehlo.dynamic_update_slice"(%arg0, %arg1, %arg2) : (tensor<16xf32>, tensor<4xf32>, tensor<i64>) -> tensor<16xf32>
-   func.return %0 : tensor<16xf32>
- }
- 
- // CHECK-LABEL: "op_einsum"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_einsum(%arg0: tensor<8x16xf32>, %arg1: tensor<16x8xf32>) -> tensor<8x8xf32> {
--  //      CHECK: "vhlo.einsum_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.einsum_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   einsum_config = #vhlo.string_v1<"ab,bc->ac">
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<8x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<8x8x!vhlo.f32_v1>
-   %0 = "stablehlo.einsum"(%arg0, %arg1) {
-@@ -1286,22 +1380,25 @@
- }
- 
- // CHECK-LABEL: "op_exponential_minus_one"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_exponential_minus_one(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.exponential_minus_one_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.exponential_minus_one_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.exponential_minus_one"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_exponential"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_exponential(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.exponential_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.exponential_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.exponential"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_fft"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_fft(%arg0: tensor<16xcomplex<f32>>) -> tensor<16xcomplex<f32>> {
--  //      CHECK: "vhlo.fft_v1"(%arg0) <{
-+  //      CHECK: "vhlo.fft_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   fft_length = #vhlo.tensor_v1<dense<16> : tensor<1xi64>>,
-   // CHECK-SAME:   fft_type = #vhlo<fft_type_v1 FFT>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<16x!vhlo.complex_v1<!vhlo.f32_v1>>) -> !vhlo.tensor_v1<16x!vhlo.complex_v1<!vhlo.f32_v1>>
-@@ -1313,8 +1410,9 @@
- }
- 
- // CHECK-LABEL: "op_floor"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_floor(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.floor_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.floor_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.floor"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
-@@ -1327,16 +1425,17 @@
-   // CHECK-SAME:   sym_name = #vhlo.string_v1<"op_func">,
-   // CHECK-SAME:   sym_visibility = #vhlo.string_v1<"private">
-   // CHECK-SAME: }> ({
--  // CHECK-NEXT:   ^[[BB:bb.*]](%arg0: !vhlo.tensor_v1<!vhlo.f32_v1>):
--  // CHECK-NEXT:     "vhlo.return_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-+  // CHECK-NEXT:   ^[[BB:bb.*]](%[[ARG0:.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
-+  // CHECK-NEXT:     "vhlo.return_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-   // CHECK-NEXT: }) : () -> ()
- 
-   func.return %arg0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_gather"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_gather(%arg0 : tensor<2x4x9xf32>, %arg1 : tensor<1x5x2xi32>) -> tensor<1x5x1xf32> {
--  //      CHECK: "vhlo.gather_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.gather_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   collapsed_slice_dims = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-   // CHECK-SAME:   index_vector_dim = #vhlo.integer_v1<2 : i64>,
-   // CHECK-SAME:   indices_are_sorted = #vhlo.bool_v1<true>,
-@@ -1358,8 +1457,9 @@
- }
- 
- // CHECK-LABEL: "op_get_dimension_size"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_get_dimension_size(%arg0: tensor<?xf32>) -> tensor<i32> {
--  //      CHECK: "vhlo.get_dimension_size_v1"(%arg0) <{
-+  //      CHECK: "vhlo.get_dimension_size_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   dimension = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<?x!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-   %0 = "stablehlo.get_dimension_size"(%arg0) {
-@@ -1369,8 +1469,9 @@
- }
- 
- // CHECK-LABEL: "op_get_tuple_element"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_get_tuple_element(%arg0: tuple<tensor<f32>, tensor<i32>>) -> tensor<f32> {
--  //      CHECK: "vhlo.get_tuple_element_v1"(%arg0) <{
-+  //      CHECK: "vhlo.get_tuple_element_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   index = #vhlo.integer_v1<0 : i32>
-   // CHECK-SAME: }> : (!vhlo.tuple_v1<!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.get_tuple_element"(%arg0) {
-@@ -1380,11 +1481,12 @@
- }
- 
- // CHECK-LABEL: "op_if"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_if(%arg0: tensor<i1>, %arg1: tensor<f32>, %arg2: tensor<f32>) -> tensor<f32> {
--  //      CHECK: "vhlo.if_v1"(%arg0) ({
--  // CHECK-NEXT:   "vhlo.return_v1"(%arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-+  //      CHECK: "vhlo.if_v1"(%[[ARG0]]) ({
-+  // CHECK-NEXT:   "vhlo.return_v1"(%[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-   // CHECK-NEXT: }, {
--  // CHECK-NEXT:   "vhlo.return_v1"(%arg2) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-+  // CHECK-NEXT:   "vhlo.return_v1"(%[[ARG2]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-   // CHECK-NEXT: }) : (!vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.if"(%arg0) ({
-     "stablehlo.return"(%arg1) : (tensor<f32>) -> ()
-@@ -1395,15 +1497,17 @@
- }
- 
- // CHECK-LABEL: "op_imag"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_imag(%arg0: tensor<complex<f32>>) -> tensor<f32> {
--  // CHECK: "vhlo.imag_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.imag_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.imag"(%arg0) : (tensor<complex<f32>>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_infeed"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_infeed(%arg0: !stablehlo.token) -> (tensor<f32>, !stablehlo.token) {
--  //               CHECK: "vhlo.infeed_v1"(%arg0) <{
-+  //               CHECK: "vhlo.infeed_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   infeed_config = #vhlo.string_v1<"foo">,
-   // CHECK-SAME{LITERAL}:   layout = #vhlo.array_v1<[#vhlo.array_v1<[]>]>
-   //          CHECK-SAME: }> : (!vhlo.token_v1) -> (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.token_v1)
-@@ -1426,36 +1530,41 @@
- }
- 
- // CHECK-LABEL: "op_is_finite"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_is_finite(%arg0: tensor<f32>) -> tensor<i1> {
--  // CHECK: "vhlo.is_finite_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-+  // CHECK: "vhlo.is_finite_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-   %0 = "stablehlo.is_finite"(%arg0) : (tensor<f32>) -> tensor<i1>
-   func.return %0 : tensor<i1>
- }
- 
- // CHECK-LABEL: "op_log"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_log(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.log_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.log_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.log"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_log_plus_one"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_log_plus_one(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.log_plus_one_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.log_plus_one_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.log_plus_one"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_logistic"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_logistic(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.logistic_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.logistic_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.logistic"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_map"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_map(%arg0: tensor<16xf32>) -> tensor<16xf32> {
--  //      CHECK: "vhlo.map_v1"(%arg0) <{
-+  //      CHECK: "vhlo.map_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   dimensions = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>
-   // CHECK-SAME: }> ({
-   // CHECK-NEXT:   ^[[BB:bb.*]](%[[ARG1:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
-@@ -1473,57 +1582,65 @@
- }
- 
- // CHECK-LABEL: "op_maximum"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_maximum(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.maximum_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.maximum_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.maximum"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_minimum"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_minimum(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.minimum_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.minimum_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.minimum"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_multiply"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_multiply(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.multiply_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.multiply_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.multiply"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_negate"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_negate(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.negate_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.negate_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.negate"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_not"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_not(%arg0: tensor<i1>) -> tensor<i1> {
--  // CHECK: "vhlo.not_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-+  // CHECK: "vhlo.not_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-   %0 = "stablehlo.not"(%arg0) : (tensor<i1>) -> tensor<i1>
-   func.return %0 : tensor<i1>
- }
- 
- // CHECK-LABEL: "op_optimization_barrier"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_optimization_barrier(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.optimization_barrier_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.optimization_barrier_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.optimization_barrier"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_or"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_or(%arg0: tensor<i1>, %arg1: tensor<i1>) -> tensor<i1> {
--  // CHECK: "vhlo.or_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-+  // CHECK: "vhlo.or_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-   %0 = "stablehlo.or"(%arg0, %arg1) : (tensor<i1>, tensor<i1>) -> tensor<i1>
-   func.return %0 : tensor<i1>
- }
- 
- // CHECK-LABEL: "op_outfeed"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_outfeed(%arg0: tensor<f32>, %arg1: !stablehlo.token) -> !stablehlo.token {
--  //      CHECK: "vhlo.outfeed_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.outfeed_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   outfeed_config = #vhlo.string_v1<"foo">
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.token_v1) -> !vhlo.token_v1
-   %0 = "stablehlo.outfeed"(%arg0, %arg1) {
-@@ -1533,8 +1650,9 @@
- }
- 
- // CHECK-LABEL: "op_pad"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_pad(%arg0: tensor<8xf32>, %arg1: tensor<f32>) -> tensor<16xf32> {
--  //      CHECK: "vhlo.pad_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.pad_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   edge_padding_high = #vhlo.tensor_v1<dense<4> : tensor<1xi64>>,
-   // CHECK-SAME:   edge_padding_low = #vhlo.tensor_v1<dense<4> : tensor<1xi64>>,
-   // CHECK-SAME:   interior_padding = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>
-@@ -1548,36 +1666,41 @@
- }
- 
- // CHECK-LABEL: "op_popcnt"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_popcnt(%arg0: tensor<i32>) -> tensor<i32> {
--  // CHECK: "vhlo.popcnt_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-+  // CHECK: "vhlo.popcnt_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-   %0 = "stablehlo.popcnt"(%arg0) : (tensor<i32>) -> tensor<i32>
-   func.return %0 : tensor<i32>
- }
- 
- // CHECK-LABEL: "op_power"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_power(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.power_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.power_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.power"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_real_dynamic_slice"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}}, %[[ARG3:.*]]: {{.*}})
- func.func @op_real_dynamic_slice(%arg0: tensor<?xf32>, %arg1: tensor<1xindex>, %arg2: tensor<1xindex>, %arg3: tensor<1xindex>) -> tensor<?xf32> {
--  // CHECK: "vhlo.real_dynamic_slice_v1"(%arg0, %arg1, %arg2, %arg3) : (!vhlo.tensor_v1<?x!vhlo.f32_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>) -> !vhlo.tensor_v1<?x!vhlo.f32_v1>
-+  // CHECK: "vhlo.real_dynamic_slice_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]], %[[ARG3]]) : (!vhlo.tensor_v1<?x!vhlo.f32_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>) -> !vhlo.tensor_v1<?x!vhlo.f32_v1>
-   %0 = "stablehlo.real_dynamic_slice"(%arg0, %arg1, %arg2, %arg3) : (tensor<?xf32>, tensor<1xindex>, tensor<1xindex>, tensor<1xindex>) -> tensor<?xf32>
-   func.return %0 : tensor<?xf32>
- }
- 
- // CHECK-LABEL: "op_real"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_real(%arg0: tensor<complex<f32>>) -> tensor<f32> {
--  // CHECK: "vhlo.real_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.real_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.real"(%arg0) : (tensor<complex<f32>>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_recv"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_recv(%arg0: !stablehlo.token) -> (tensor<f32>, !stablehlo.token) {
--  //      CHECK: "vhlo.recv_v1"(%arg0) <{
-+  //      CHECK: "vhlo.recv_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME:   channel_type = #vhlo.integer_v1<3 : i64>,
-   // CHECK-SAME:   is_host_transfer = #vhlo.bool_v1<true>
-@@ -1590,8 +1713,9 @@
- }
- 
- // CHECK-LABEL: "op_reduce"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_reduce(%arg0: tensor<16xf32>, %arg1: tensor<f32>) -> tensor<f32> {
--  //  CHECK: "vhlo.reduce_v1"(%arg0, %arg1)
-+  //  CHECK: "vhlo.reduce_v1"(%[[ARG0]], %[[ARG1]])
-   //  CHECK:   ^[[BB:bb.*]](%[[ARG1:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>, %[[ARG2:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
-   //  CHECK:     "vhlo.return_v1"(%[[VAL1:.*]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-   //  CHECK: }) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-@@ -1606,8 +1730,9 @@
- }
- 
- // CHECK-LABEL: "op_reduce_precision"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_reduce_precision(%arg0: tensor<f32>) -> tensor<f32> {
--  //      CHECK: "vhlo.reduce_precision_v1"(%arg0) <{
-+  //      CHECK: "vhlo.reduce_precision_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   exponent_bits = #vhlo.integer_v1<8 : i32>
-   // CHECK-SAME:   mantissa_bits = #vhlo.integer_v1<10 : i32>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-@@ -1619,8 +1744,9 @@
- }
- 
- // CHECK-LABEL: "op_reduce_scatter"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_reduce_scatter(%arg0: tensor<16xf32>) -> tensor<16xf32> {
--  //               CHECK: "vhlo.reduce_scatter_v1"(%arg0) <{
-+  //               CHECK: "vhlo.reduce_scatter_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0], [1]]> : tensor<2x1xi64>>,
-   //          CHECK-SAME:   scatter_dimension = #vhlo.integer_v1<0 : i64>
-@@ -1644,8 +1770,9 @@
- }
- 
- // CHECK-LABEL: "op_reduce_window"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_reduce_window(%arg0: tensor<2x17x31x7xf32>, %arg1: tensor<f32>) -> tensor<2x9x16x7xf32> {
--  //               CHECK: "vhlo.reduce_window_v1"(%arg0, %arg1) <{
-+  //               CHECK: "vhlo.reduce_window_v1"(%[[ARG0]], %[[ARG1]]) <{
-   //          CHECK-SAME:   base_dilations = #vhlo.tensor_v1<dense<[1, 2, 2, 1]> : tensor<4xi64>>,
-   // CHECK-SAME{LITERAL}:   padding = #vhlo.tensor_v1<dense<[[0, 0], [2, 0], [0, 2], [0, 0]]> : tensor<4x2xi64>>,
-   //          CHECK-SAME:   window_dilations = #vhlo.tensor_v1<dense<[1, 2, 2, 1]> : tensor<4xi64>>,
-@@ -1671,8 +1798,9 @@
- }
- 
- // CHECK-LABEL: "op_remainder"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_remainder(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.remainder_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.remainder_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.remainder"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
-@@ -1692,16 +1820,18 @@
- }
- 
- // CHECK-LABEL: "op_reshape"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_reshape(%arg0: tensor<16xf32>) -> tensor<4x4xf32> {
--  // CHECK: "vhlo.reshape_v1"(%arg0) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<4x4x!vhlo.f32_v1>
-+  // CHECK: "vhlo.reshape_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<4x4x!vhlo.f32_v1>
-   %0 = "stablehlo.reshape"(%arg0) : (tensor<16xf32>) -> tensor<4x4xf32>
-   func.return %0 : tensor<4x4xf32>
- }
- 
- // CHECK-LABEL: "op_return"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_return(%arg0: tensor<i32>, %arg1: tensor<f32>) -> tensor<f32> {
--  //      CHECK: "vhlo.case_v1"(%arg0) ({
--  // CHECK-NEXT:   "vhlo.return_v1"(%arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-+  //      CHECK: "vhlo.case_v1"(%[[ARG0]]) ({
-+  // CHECK-NEXT:   "vhlo.return_v1"(%[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-   // CHECK-NEXT: }) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.case"(%arg0) ({
-     "stablehlo.return"(%arg1) : (tensor<f32>) -> ()
-@@ -1710,8 +1840,9 @@
- }
- 
- // CHECK-LABEL: "op_reverse"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_reverse(%arg0: tensor<16xf32>) -> tensor<16xf32> {
--  //      CHECK: "vhlo.reverse_v1"(%arg0) <{
-+  //      CHECK: "vhlo.reverse_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   dimensions = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x!vhlo.f32_v1>
-   %0 = "stablehlo.reverse"(%arg0) {
-@@ -1721,8 +1852,9 @@
- }
- 
- // CHECK-LABEL: "op_rng_bit_generator"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_rng_bit_generator(%arg0: tensor<f32>) -> (tensor<f32>, tensor<f32>) {
--  //      CHECK: "vhlo.rng_bit_generator_v1"(%arg0) <{
-+  //      CHECK: "vhlo.rng_bit_generator_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   rng_algorithm = #vhlo<rng_algorithm_v1 PHILOX>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>)
-   %0:2 = "stablehlo.rng_bit_generator"(%arg0) {
-@@ -1732,8 +1864,9 @@
- }
- 
- // CHECK-LABEL: "op_rng"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_rng(%arg0: tensor<f32>, %arg1: tensor<f32>, %arg2: tensor<0xindex>) -> tensor<f32> {
--  //      CHECK: "vhlo.rng_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.rng_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   rng_distribution = #vhlo<rng_distribution_v1 NORMAL>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<0x!vhlo.index_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.rng"(%arg0, %arg1, %arg2) {
-@@ -1743,29 +1876,33 @@
- }
- 
- // CHECK-LABEL: "op_round_nearest_afz"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_round_nearest_afz(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.round_nearest_afz_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.round_nearest_afz_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.round_nearest_afz"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_round_nearest_even"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_round_nearest_even(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.round_nearest_even_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.round_nearest_even_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.round_nearest_even"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_rsqrt"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_rsqrt(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.rsqrt_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.rsqrt_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.rsqrt"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_scatter"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_scatter(%arg0: tensor<200x100x300xf32>, %arg1: tensor<10x2xi32>, %arg2: tensor<10x300xf32>) -> tensor<200x100x300xf32> {
--  //      CHECK: "vhlo.scatter_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.scatter_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   index_vector_dim = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   indices_are_sorted = #vhlo.bool_v1<true>,
-   // CHECK-SAME:   inserted_window_dims = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-@@ -1795,8 +1932,9 @@
- }
- 
- // CHECK-LABEL: "op_select_and_scatter"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_select_and_scatter(%arg0: tensor<10x24x24x64xf32>, %arg1: tensor<12x13x13x66xf32>, %arg2: tensor<f32>) -> tensor<10x24x24x64xf32> {
--  //      CHECK: "vhlo.select_and_scatter_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.select_and_scatter_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   padding = #vhlo.tensor_v1<dense<1> : tensor<4x2xi64>>,
-   // CHECK-SAME:   window_dimensions = #vhlo.tensor_v1<dense<[1, 2, 2, 1]> : tensor<4xi64>>,
-   // CHECK-SAME:   window_strides = #vhlo.tensor_v1<dense<[1, 2, 2, 1]> : tensor<4xi64>>
-@@ -1826,15 +1964,17 @@
- }
- 
- // CHECK-LABEL: "op_select"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_select(%arg0: tensor<i1>, %arg1: tensor<f32>, %arg2: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.select_v1"(%arg0, %arg1, %arg2) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.select_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.select"(%arg0, %arg1, %arg2) : (tensor<i1>, tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_send"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_send(%arg0: tensor<f32>, %arg1: !stablehlo.token) -> !stablehlo.token {
--  //      CHECK: "vhlo.send_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.send_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME:   channel_type = #vhlo.integer_v1<2 : i64>,
-   // CHECK-SAME:   is_host_transfer = #vhlo.bool_v1<true>
-@@ -1847,8 +1987,9 @@
- }
- 
- // CHECK-LABEL: "op_set_dimension_size"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_set_dimension_size(%arg0: tensor<?xf32>, %arg1: tensor<i32>) -> tensor<16xf32> {
--  //      CHECK: "vhlo.set_dimension_size_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.set_dimension_size_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   dimension = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<?x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<16x!vhlo.f32_v1>
-   %0 = "stablehlo.set_dimension_size"(%arg0, %arg1) {
-@@ -1858,43 +1999,49 @@
- }
- 
- // CHECK-LABEL: "op_shift_left"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_shift_left(%arg0: tensor<i32>, %arg1: tensor<i32>) -> tensor<i32> {
--  // CHECK: "vhlo.shift_left_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.i32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-+  // CHECK: "vhlo.shift_left_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-   %0 = "stablehlo.shift_left"(%arg0, %arg1) : (tensor<i32>, tensor<i32>) -> tensor<i32>
-   func.return %0 : tensor<i32>
- }
- 
- // CHECK-LABEL: "op_shift_right_arithmetic"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_shift_right_arithmetic(%arg0: tensor<i32>, %arg1: tensor<i32>) -> tensor<i32> {
--  // CHECK: "vhlo.shift_right_arithmetic_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.i32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-+  // CHECK: "vhlo.shift_right_arithmetic_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-   %0 = "stablehlo.shift_right_arithmetic"(%arg0, %arg1) : (tensor<i32>, tensor<i32>) -> tensor<i32>
-   func.return %0 : tensor<i32>
- }
- 
- // CHECK-LABEL: "op_shift_right_logical"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_shift_right_logical(%arg0: tensor<i32>, %arg1: tensor<i32>) -> tensor<i32> {
--  // CHECK: "vhlo.shift_right_logical_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.i32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-+  // CHECK: "vhlo.shift_right_logical_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-   %0 = "stablehlo.shift_right_logical"(%arg0, %arg1) : (tensor<i32>, tensor<i32>) -> tensor<i32>
-   func.return %0 : tensor<i32>
- }
- 
- // CHECK-LABEL: "op_sign"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_sign(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.sign_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.sign_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.sign"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_sine"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_sine(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.sine_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.sine_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.sine"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_slice"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_slice(%arg0: tensor<16xf32>) -> tensor<4xf32> {
--  //      CHECK: "vhlo.slice_v1"(%arg0) <{
-+  //      CHECK: "vhlo.slice_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   limit_indices = #vhlo.tensor_v1<dense<4> : tensor<1xi64>>,
-   // CHECK-SAME:   start_indices = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>,
-   // CHECK-SAME:   strides = #vhlo.tensor_v1<dense<1> : tensor<1xi64>>
-@@ -1908,8 +2055,9 @@
- }
- 
- // CHECK-LABEL: "op_sort"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_sort(%arg0: tensor<16xf32>) -> tensor<16xf32> {
--  //      CHECK: "vhlo.sort_v1"(%arg0) <{
-+  //      CHECK: "vhlo.sort_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   dimension = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME:   is_stable = #vhlo.bool_v1<true>
-   // CHECK-SAME: }> ({
-@@ -1929,29 +2077,33 @@
- }
- 
- // CHECK-LABEL: "op_sqrt"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_sqrt(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.sqrt_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.sqrt_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.sqrt"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_subtract"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_subtract(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.subtract_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.subtract_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.subtract"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_tanh"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_tanh(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.tanh_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.tanh_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.tanh"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_torch_index_select"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_torch_index_select(%arg0: tensor<5x1x5xf32>, %arg1: tensor<2xi32>) ->  tensor<2x1x5xf32> {
--  //      CHECK: "vhlo.torch_index_select_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.torch_index_select_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   batch_dims = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME:   dim = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<5x1x5x!vhlo.f32_v1>, !vhlo.tensor_v1<2x!vhlo.i32_v1>) -> !vhlo.tensor_v1<2x1x5x!vhlo.f32_v1>
-@@ -1963,8 +2115,9 @@
- }
- 
- // CHECK-LABEL: "op_trace"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_trace(%arg0: tensor<f32>) {
--  //      CHECK: "vhlo.trace_v1"(%arg0) <{
-+  //      CHECK: "vhlo.trace_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   tag = #vhlo.string_v1<"foo">
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-   "stablehlo.trace"(%arg0) {
-@@ -1974,8 +2127,9 @@
- }
- 
- // CHECK-LABEL: "op_transpose"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_transpose(%arg0: tensor<16x8xf32>) ->  tensor<8x16xf32> {
--  //      CHECK: "vhlo.transpose_v1"(%arg0) <{
-+  //      CHECK: "vhlo.transpose_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   permutation = #vhlo.tensor_v1<dense<[1, 0]> : tensor<2xi64>>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<16x8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<8x16x!vhlo.f32_v1>
-   %0 = "stablehlo.transpose"(%arg0) {
-@@ -1985,8 +2139,9 @@
- }
- 
- // CHECK-LABEL: "op_triangular_solve"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_triangular_solve(%arg0: tensor<16x16xf32>, %arg1: tensor<16x16xf32>) ->  tensor<16x16xf32> {
--  //      CHECK: "vhlo.triangular_solve_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.triangular_solve_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   left_side = #vhlo.bool_v1<true>,
-   // CHECK-SAME:   lower = #vhlo.bool_v1<true>,
-   // CHECK-SAME:   transpose_a = #vhlo<transpose_v1 NO_TRANSPOSE>,
-@@ -2002,15 +2157,17 @@
- }
- 
- // CHECK-LABEL: "op_tuple"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_tuple(%arg0: tensor<f32>) -> tuple<tensor<f32>> {
--  // CHECK: "vhlo.tuple_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tuple_v1<!vhlo.tensor_v1<!vhlo.f32_v1>>
-+  // CHECK: "vhlo.tuple_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tuple_v1<!vhlo.tensor_v1<!vhlo.f32_v1>>
-   %0 = "stablehlo.tuple"(%arg0) : (tensor<f32>) -> tuple<tensor<f32>>
-   func.return %0 : tuple<tensor<f32>>
- }
- 
- // CHECK-LABEL: "op_unary_einsum"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_unary_einsum(%arg0: tensor<8x16xf32>) -> tensor<8xf32> {
--  //      CHECK: "vhlo.unary_einsum_v1"(%arg0) <{
-+  //      CHECK: "vhlo.unary_einsum_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   einsum_config = #vhlo.string_v1<"ab->a">
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<8x16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<8x!vhlo.f32_v1>
-   %0 = "stablehlo.unary_einsum"(%arg0) {
-@@ -2020,22 +2177,25 @@
- }
- 
- // CHECK-LABEL: "op_uniform_dequantize"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_uniform_dequantize(%arg0: tensor<!quant.uniform<i8:f32, 34.0:16>>) -> tensor<f32> {
--  // CHECK: "vhlo.uniform_dequantize_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.uniform_dequantize_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.uniform_dequantize"(%arg0) : (tensor<!quant.uniform<i8:f32, 34.0:16>>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_uniform_quantize"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_uniform_quantize(%arg0: tensor<f32>) -> tensor<!quant.uniform<i8:f32, 34.0:16>> {
--  // CHECK: "vhlo.uniform_quantize_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>
-+  // CHECK: "vhlo.uniform_quantize_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>
-   %0 = "stablehlo.uniform_quantize"(%arg0) : (tensor<f32>) -> tensor<!quant.uniform<i8:f32, 34.0:16>>
-   func.return %0 : tensor<!quant.uniform<i8:f32, 34.0:16>>
- }
- 
- // CHECK-LABEL: "op_while"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_while(%arg0: tensor<i1>) -> tensor<i1> {
--  //      CHECK: "vhlo.while_v1"(%arg0) ({
-+  //      CHECK: "vhlo.while_v1"(%[[ARG0]]) ({
-   // CHECK-NEXT:   ^[[BB:bb.*]](%[[ARG1:arg.*]]: !vhlo.tensor_v1<!vhlo.bool_v1>):
-   // CHECK-NEXT:     "vhlo.return_v1"(%[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>) -> ()
-   // CHECK-NEXT:   }, {
-@@ -2053,8 +2213,9 @@
- }
- 
- // CHECK-LABEL: "op_xor"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_xor(%arg0: tensor<i1>, %arg1: tensor<i1>) -> tensor<i1> {
--  // CHECK: "vhlo.xor_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-+  // CHECK: "vhlo.xor_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-   %0 = "stablehlo.xor"(%arg0, %arg1) : (tensor<i1>, tensor<i1>) -> tensor<i1>
-   func.return %0 : tensor<i1>
- }
-@@ -2062,190 +2223,217 @@
- // ============ TYPES ============
- 
- // CHECK-LABEL: "type_i1"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_i1(%arg0: tensor<i1>, %arg1: tensor<i1>) -> tensor<i1> {
--  // CHECK: "vhlo.and_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-+  // CHECK: "vhlo.and_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-   %0 = "stablehlo.and"(%arg0, %arg1) : (tensor<i1>, tensor<i1>) -> tensor<i1>
-   func.return %0 : tensor<i1>
- }
- 
- // CHECK-LABEL: "type_i4"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_i4(%arg0: tensor<i4>, %arg1: tensor<i4>) -> tensor<i4> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.i4_v1>, !vhlo.tensor_v1<!vhlo.i4_v1>) -> !vhlo.tensor_v1<!vhlo.i4_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i4_v1>, !vhlo.tensor_v1<!vhlo.i4_v1>) -> !vhlo.tensor_v1<!vhlo.i4_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<i4>, tensor<i4>) -> tensor<i4>
-   func.return %0 : tensor<i4>
- }
- 
- // CHECK-LABEL: "type_i8"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_i8(%arg0: tensor<i8>, %arg1: tensor<i8>) -> tensor<i8> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.i8_v1>, !vhlo.tensor_v1<!vhlo.i8_v1>) -> !vhlo.tensor_v1<!vhlo.i8_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i8_v1>, !vhlo.tensor_v1<!vhlo.i8_v1>) -> !vhlo.tensor_v1<!vhlo.i8_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<i8>, tensor<i8>) -> tensor<i8>
-   func.return %0 : tensor<i8>
- }
- 
- // CHECK-LABEL: "type_i16"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_i16(%arg0: tensor<i16>, %arg1: tensor<i16>) -> tensor<i16> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.i16_v1>, !vhlo.tensor_v1<!vhlo.i16_v1>) -> !vhlo.tensor_v1<!vhlo.i16_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i16_v1>, !vhlo.tensor_v1<!vhlo.i16_v1>) -> !vhlo.tensor_v1<!vhlo.i16_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<i16>, tensor<i16>) -> tensor<i16>
-   func.return %0 : tensor<i16>
- }
- 
- // CHECK-LABEL: "type_i32"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_i32(%arg0: tensor<i32>, %arg1: tensor<i32>) -> tensor<i32> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.i32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<i32>, tensor<i32>) -> tensor<i32>
-   func.return %0 : tensor<i32>
- }
- 
- // CHECK-LABEL: "type_i64"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_i64(%arg0: tensor<i64>, %arg1: tensor<i64>) -> tensor<i64> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.i64_v1>, !vhlo.tensor_v1<!vhlo.i64_v1>) -> !vhlo.tensor_v1<!vhlo.i64_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i64_v1>, !vhlo.tensor_v1<!vhlo.i64_v1>) -> !vhlo.tensor_v1<!vhlo.i64_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<i64>, tensor<i64>) -> tensor<i64>
-   func.return %0 : tensor<i64>
- }
- 
- // CHECK-LABEL: "type_ui4"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_ui4(%arg0: tensor<ui4>, %arg1: tensor<ui4>) -> tensor<ui4> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.ui4_v1>, !vhlo.tensor_v1<!vhlo.ui4_v1>) -> !vhlo.tensor_v1<!vhlo.ui4_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.ui4_v1>, !vhlo.tensor_v1<!vhlo.ui4_v1>) -> !vhlo.tensor_v1<!vhlo.ui4_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<ui4>, tensor<ui4>) -> tensor<ui4>
-   func.return %0 : tensor<ui4>
- }
- 
- // CHECK-LABEL: "type_ui8"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_ui8(%arg0: tensor<ui8>, %arg1: tensor<ui8>) -> tensor<ui8> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.ui8_v1>, !vhlo.tensor_v1<!vhlo.ui8_v1>) -> !vhlo.tensor_v1<!vhlo.ui8_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.ui8_v1>, !vhlo.tensor_v1<!vhlo.ui8_v1>) -> !vhlo.tensor_v1<!vhlo.ui8_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<ui8>, tensor<ui8>) -> tensor<ui8>
-   func.return %0 : tensor<ui8>
- }
- 
- // CHECK-LABEL: "type_ui16"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_ui16(%arg0: tensor<ui16>, %arg1: tensor<ui16>) -> tensor<ui16> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.ui16_v1>, !vhlo.tensor_v1<!vhlo.ui16_v1>) -> !vhlo.tensor_v1<!vhlo.ui16_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.ui16_v1>, !vhlo.tensor_v1<!vhlo.ui16_v1>) -> !vhlo.tensor_v1<!vhlo.ui16_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<ui16>, tensor<ui16>) -> tensor<ui16>
-   func.return %0 : tensor<ui16>
- }
- 
- // CHECK-LABEL: "type_ui32"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_ui32(%arg0: tensor<ui32>, %arg1: tensor<ui32>) -> tensor<ui32> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.ui32_v1>, !vhlo.tensor_v1<!vhlo.ui32_v1>) -> !vhlo.tensor_v1<!vhlo.ui32_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.ui32_v1>, !vhlo.tensor_v1<!vhlo.ui32_v1>) -> !vhlo.tensor_v1<!vhlo.ui32_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<ui32>, tensor<ui32>) -> tensor<ui32>
-   func.return %0 : tensor<ui32>
- }
- 
- // CHECK-LABEL: "type_ui64"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_ui64(%arg0: tensor<ui64>, %arg1: tensor<ui64>) -> tensor<ui64> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.ui64_v1>, !vhlo.tensor_v1<!vhlo.ui64_v1>) -> !vhlo.tensor_v1<!vhlo.ui64_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.ui64_v1>, !vhlo.tensor_v1<!vhlo.ui64_v1>) -> !vhlo.tensor_v1<!vhlo.ui64_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<ui64>, tensor<ui64>) -> tensor<ui64>
-   func.return %0 : tensor<ui64>
- }
- 
- // CHECK-LABEL: "type_f8E4M3FN"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_f8E4M3FN(%arg0: tensor<f8E4M3FN>, %arg1: tensor<f8E4M3FN>) -> tensor<f8E4M3FN> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f8E4M3FN_v1>, !vhlo.tensor_v1<!vhlo.f8E4M3FN_v1>) -> !vhlo.tensor_v1<!vhlo.f8E4M3FN_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f8E4M3FN_v1>, !vhlo.tensor_v1<!vhlo.f8E4M3FN_v1>) -> !vhlo.tensor_v1<!vhlo.f8E4M3FN_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f8E4M3FN>, tensor<f8E4M3FN>) -> tensor<f8E4M3FN>
-   func.return %0 : tensor<f8E4M3FN>
- }
- 
- // CHECK-LABEL: "type_f8E5M2"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_f8E5M2(%arg0: tensor<f8E5M2>, %arg1: tensor<f8E5M2>) -> tensor<f8E5M2> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f8E5M2_v1>, !vhlo.tensor_v1<!vhlo.f8E5M2_v1>) -> !vhlo.tensor_v1<!vhlo.f8E5M2_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f8E5M2_v1>, !vhlo.tensor_v1<!vhlo.f8E5M2_v1>) -> !vhlo.tensor_v1<!vhlo.f8E5M2_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f8E5M2>, tensor<f8E5M2>) -> tensor<f8E5M2>
-   func.return %0 : tensor<f8E5M2>
- }
- 
- // CHECK-LABEL: "type_f8E4M3FNUZ"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_f8E4M3FNUZ(%arg0: tensor<f8E4M3FNUZ>, %arg1: tensor<f8E4M3FNUZ>) -> tensor<f8E4M3FNUZ> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f8E4M3FNUZ_v1>, !vhlo.tensor_v1<!vhlo.f8E4M3FNUZ_v1>) -> !vhlo.tensor_v1<!vhlo.f8E4M3FNUZ_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f8E4M3FNUZ_v1>, !vhlo.tensor_v1<!vhlo.f8E4M3FNUZ_v1>) -> !vhlo.tensor_v1<!vhlo.f8E4M3FNUZ_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f8E4M3FNUZ>, tensor<f8E4M3FNUZ>) -> tensor<f8E4M3FNUZ>
-   func.return %0 : tensor<f8E4M3FNUZ>
- }
- 
- // CHECK-LABEL: "type_f8E4M3B11FNUZ"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_f8E4M3B11FNUZ(%arg0: tensor<f8E4M3B11FNUZ>, %arg1: tensor<f8E4M3B11FNUZ>) -> tensor<f8E4M3B11FNUZ> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f8E4M3B11FNUZ_v1>, !vhlo.tensor_v1<!vhlo.f8E4M3B11FNUZ_v1>) -> !vhlo.tensor_v1<!vhlo.f8E4M3B11FNUZ_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f8E4M3B11FNUZ_v1>, !vhlo.tensor_v1<!vhlo.f8E4M3B11FNUZ_v1>) -> !vhlo.tensor_v1<!vhlo.f8E4M3B11FNUZ_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f8E4M3B11FNUZ>, tensor<f8E4M3B11FNUZ>) -> tensor<f8E4M3B11FNUZ>
-   func.return %0 : tensor<f8E4M3B11FNUZ>
- }
- 
- // CHECK-LABEL: "type_f8E5M2FNUZ"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_f8E5M2FNUZ(%arg0: tensor<f8E5M2FNUZ>, %arg1: tensor<f8E5M2FNUZ>) -> tensor<f8E5M2FNUZ> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f8E5M2FNUZ_v1>, !vhlo.tensor_v1<!vhlo.f8E5M2FNUZ_v1>) -> !vhlo.tensor_v1<!vhlo.f8E5M2FNUZ_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f8E5M2FNUZ_v1>, !vhlo.tensor_v1<!vhlo.f8E5M2FNUZ_v1>) -> !vhlo.tensor_v1<!vhlo.f8E5M2FNUZ_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f8E5M2FNUZ>, tensor<f8E5M2FNUZ>) -> tensor<f8E5M2FNUZ>
-   func.return %0 : tensor<f8E5M2FNUZ>
- }
- 
- // CHECK-LABEL: "type_bf16"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_bf16(%arg0: tensor<bf16>, %arg1: tensor<bf16>) -> tensor<bf16> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.bf16_v1>, !vhlo.tensor_v1<!vhlo.bf16_v1>) -> !vhlo.tensor_v1<!vhlo.bf16_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.bf16_v1>, !vhlo.tensor_v1<!vhlo.bf16_v1>) -> !vhlo.tensor_v1<!vhlo.bf16_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<bf16>, tensor<bf16>) -> tensor<bf16>
-   func.return %0 : tensor<bf16>
- }
- 
- // CHECK-LABEL: "type_f16"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_f16(%arg0: tensor<f16>, %arg1: tensor<f16>) -> tensor<f16> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f16_v1>, !vhlo.tensor_v1<!vhlo.f16_v1>) -> !vhlo.tensor_v1<!vhlo.f16_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f16_v1>, !vhlo.tensor_v1<!vhlo.f16_v1>) -> !vhlo.tensor_v1<!vhlo.f16_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f16>, tensor<f16>) -> tensor<f16>
-   func.return %0 : tensor<f16>
- }
- 
- // CHECK-LABEL: "type_f32"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_f32(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "type_f64"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_f64(%arg0: tensor<f64>, %arg1: tensor<f64>) -> tensor<f64> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f64_v1>, !vhlo.tensor_v1<!vhlo.f64_v1>) -> !vhlo.tensor_v1<!vhlo.f64_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f64_v1>, !vhlo.tensor_v1<!vhlo.f64_v1>) -> !vhlo.tensor_v1<!vhlo.f64_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f64>, tensor<f64>) -> tensor<f64>
-   func.return %0 : tensor<f64>
- }
- 
- // CHECK-LABEL: "type_complex_f32"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_complex_f32(%arg0: tensor<complex<f32>>, %arg1: tensor<complex<f32>>) -> tensor<complex<f32>> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>, !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>) -> !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>, !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>) -> !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<complex<f32>>, tensor<complex<f32>>) -> tensor<complex<f32>>
-   func.return %0 : tensor<complex<f32>>
- }
- 
- // CHECK-LABEL: "type_complex_f64"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_complex_f64(%arg0: tensor<complex<f64>>, %arg1: tensor<complex<f64>>) -> tensor<complex<f64>> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f64_v1>>, !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f64_v1>>) -> !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f64_v1>>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f64_v1>>, !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f64_v1>>) -> !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f64_v1>>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<complex<f64>>, tensor<complex<f64>>) -> tensor<complex<f64>>
-   func.return %0 : tensor<complex<f64>>
- }
- 
- // CHECK-LABEL: "type_dynamism_ranked"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @type_dynamism_ranked(%arg0: tensor<?xf32>) -> tensor<?xf32> {
--  // CHECK: "vhlo.abs_v1"(%arg0) : (!vhlo.tensor_v1<?x!vhlo.f32_v1>) -> !vhlo.tensor_v1<?x!vhlo.f32_v1>
-+  // CHECK: "vhlo.abs_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<?x!vhlo.f32_v1>) -> !vhlo.tensor_v1<?x!vhlo.f32_v1>
-   %0 = "stablehlo.abs"(%arg0) : (tensor<?xf32>) -> tensor<?xf32>
-   func.return %0 : tensor<?xf32>
- }
- 
- // CHECK-LABEL: "type_quantization"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_quantization(%arg0: tensor<!quant.uniform<i8:f32, 34.0:16>>, %arg1: tensor<!quant.uniform<i8:f32, 34.0:16>>) -> tensor<!quant.uniform<i8:f32, 34.0:16>> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>, !vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>) -> !vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>, !vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>) -> !vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<!quant.uniform<i8:f32, 34.0:16>>, tensor<!quant.uniform<i8:f32, 34.0:16>>) -> tensor<!quant.uniform<i8:f32, 34.0:16>>
-   func.return %0 : tensor<!quant.uniform<i8:f32, 34.0:16>>
- }
- 
- //       CHECK: function_type = #vhlo.type_v1<!vhlo.func_v1<(!vhlo.token_v1) -> !vhlo.token_v1>>
- // CHECK-LABEL: "type_token_callee"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @type_token_callee(%arg0: !stablehlo.token) -> !stablehlo.token {
--  // CHECK: "vhlo.return_v1"(%arg0) : (!vhlo.token_v1) -> ()
-+  // CHECK: "vhlo.return_v1"(%[[ARG0]]) : (!vhlo.token_v1) -> ()
-   return %arg0 : !stablehlo.token
- }
- 
- //       CHECK: function_type = #vhlo.type_v1<!vhlo.func_v1<(!vhlo.token_v1) -> !vhlo.token_v1>>
- // CHECK-LABEL: "type_token_caller"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @type_token_caller(%arg0: !stablehlo.token) -> !stablehlo.token {
--  // CHECK:      "vhlo.call_v1"(%arg0) <{callee = #vhlo.string_v1<"type_token_callee">}
-+  // CHECK:      "vhlo.call_v1"(%[[ARG0]]) <{callee = #vhlo.string_v1<"type_token_callee">}
-   // CHECK-SAME: (!vhlo.token_v1) -> !vhlo.token_v1
-   %0 = func.call @type_token_callee(%arg0) : (!stablehlo.token) -> !stablehlo.token
-   return %0 : !stablehlo.token
- }
- 
- // CHECK-LABEL: "type_tuple"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @type_tuple(%arg0: tuple<tensor<f32>>) -> tuple<!stablehlo.token> {
-   %0 = "stablehlo.custom_call"(%arg0) {
-     call_target_name = "foo"
-diff --ruN a/stablehlo/stablehlo/tests/vhlo/stablehlo_legalize_to_vhlo.0_14_0.mlir b/stablehlo/stablehlo/tests/vhlo/stablehlo_legalize_to_vhlo.0_14_0.mlir
---- stablehlo/stablehlo/tests/vhlo/stablehlo_legalize_to_vhlo.0_14_0.mlir
-+++ stablehlo/stablehlo/tests/vhlo/stablehlo_legalize_to_vhlo.0_14_0.mlir
-@@ -13,6 +13,7 @@
- // ============ ATTRIBUTES ============
- 
- // CHECK-LABEL: "attr_comparison_direction_eq"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_direction_eq(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     // CHECK: comparison_direction = #vhlo<comparison_direction_v1 EQ>
-@@ -22,6 +23,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_direction_ne"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_direction_ne(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     // CHECK: comparison_direction = #vhlo<comparison_direction_v1 NE>
-@@ -31,6 +33,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_direction_ge"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_direction_ge(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     // CHECK: comparison_direction = #vhlo<comparison_direction_v1 GE>
-@@ -40,6 +43,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_direction_gt"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_direction_gt(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     // CHECK: comparison_direction = #vhlo<comparison_direction_v1 GT>
-@@ -49,6 +53,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_direction_le"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_direction_le(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     // CHECK: comparison_direction = #vhlo<comparison_direction_v1 LE>
-@@ -58,6 +63,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_direction_lt"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_direction_lt(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     // CHECK: comparison_direction = #vhlo<comparison_direction_v1 LT>
-@@ -67,6 +73,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_type_notype"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_type_notype(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     comparison_direction = #stablehlo<comparison_direction EQ>
-@@ -76,6 +83,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_type_float"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_type_float(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     comparison_direction = #stablehlo<comparison_direction EQ>,
-@@ -86,6 +94,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_type_totalorder"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_type_totalorder(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     comparison_direction = #stablehlo<comparison_direction EQ>,
-@@ -96,6 +105,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_type_signed"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_type_signed(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     comparison_direction = #stablehlo<comparison_direction EQ>,
-@@ -106,6 +116,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_type_unsigned"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_type_unsigned(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     comparison_direction = #stablehlo<comparison_direction EQ>,
-@@ -118,6 +129,7 @@
- // ConvDimensionNumbers aka #stablehlo.conv is covered below.
- 
- // CHECK-LABEL: "attr_custom_call_api_version_unspecified"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_custom_call_api_version_unspecified(%arg0: tensor<f32>) -> tensor<f32> {
-   %0 = "stablehlo.custom_call"(%arg0) {
-     call_target_name = "foo",
-@@ -128,6 +140,7 @@
- }
- 
- // CHECK-LABEL: "attr_custom_call_api_version_original"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_custom_call_api_version_original(%arg0: tensor<f32>) -> tensor<f32> {
-   %0 = "stablehlo.custom_call"(%arg0) {
-     call_target_name = "foo",
-@@ -138,6 +151,7 @@
- }
- 
- // CHECK-LABEL: "attr_custom_call_api_version_status_returning"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_custom_call_api_version_status_returning(%arg0: tensor<f32>) -> tensor<f32> {
-   %0 = "stablehlo.custom_call"(%arg0) {
-     call_target_name = "foo",
-@@ -148,6 +162,7 @@
- }
- 
- // CHECK-LABEL: "attr_custom_call_api_version_status_returning_unified"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_custom_call_api_version_status_returning_unified(%arg0: tensor<f32>) -> tensor<f32> {
-   %0 = "stablehlo.custom_call"(%arg0) {
-     call_target_name = "foo",
-@@ -160,6 +175,7 @@
- // DotDimensionNumbers aka #stablehlo.dot is covered below.
- 
- // CHECK-LABEL: "attr_fft_type_fft"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_fft_type_fft(%arg0: tensor<16xcomplex<f32>>) -> tensor<16xcomplex<f32>> {
-   %0 = "stablehlo.fft"(%arg0) {
-     // CHECK: fft_type = #vhlo<fft_type_v1 FFT>
-@@ -170,6 +186,7 @@
- }
- 
- // CHECK-LABEL: "attr_fft_type_ifft"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_fft_type_ifft(%arg0: tensor<16xcomplex<f32>>) -> tensor<16xcomplex<f32>> {
-   %0 = "stablehlo.fft"(%arg0) {
-     // CHECK: fft_type = #vhlo<fft_type_v1 IFFT>
-@@ -180,6 +197,7 @@
- }
- 
- // CHECK-LABEL: "attr_fft_type_rfft"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_fft_type_rfft(%arg0: tensor<16xf32>) -> tensor<9xcomplex<f32>> {
-   %0 = "stablehlo.fft"(%arg0) {
-     // CHECK: fft_type = #vhlo<fft_type_v1 RFFT>
-@@ -190,6 +208,7 @@
- }
- 
- // CHECK-LABEL: "attr_fft_type_irfft"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_fft_type_irfft(%arg0: tensor<9xcomplex<f32>>) -> tensor<16xf32> {
-   %0 = "stablehlo.fft"(%arg0) {
-     // CHECK: fft_type = #vhlo<fft_type_v1 IRFFT>
-@@ -202,6 +221,7 @@
- // GatherDimensionNumbers aka #stablehlo.gather is covered below.
- 
- // CHECK-LABEL: "attr_precision_config_default"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_precision_config_default(%arg0: tensor<8x16xf32>, %arg1: tensor<16x8xf32>) -> tensor<8x8xf32> {
-   %0 = "stablehlo.dot"(%arg0, %arg1) {
-     // CHECK: precision_config = #vhlo.array_v1<[#vhlo<precision_v1 DEFAULT>, #vhlo<precision_v1 DEFAULT>]>
-@@ -210,6 +230,7 @@
- }
- 
- // CHECK-LABEL: "attr_precision_config_high"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_precision_config_high(%arg0: tensor<8x16xf32>, %arg1: tensor<16x8xf32>) -> tensor<8x8xf32> {
-   %0 = "stablehlo.dot"(%arg0, %arg1) {
-     // CHECK: precision_config = #vhlo.array_v1<[#vhlo<precision_v1 HIGH>, #vhlo<precision_v1 HIGH>]>
-@@ -219,6 +240,7 @@
- }
- 
- // CHECK-LABEL: "attr_precision_config_highest"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_precision_config_highest(%arg0: tensor<8x16xf32>, %arg1: tensor<16x8xf32>) -> tensor<8x8xf32> {
-   %0 = "stablehlo.dot"(%arg0, %arg1) {
-     // CHECK: precision_config = #vhlo.array_v1<[#vhlo<precision_v1 HIGHEST>, #vhlo<precision_v1 HIGHEST>]>
-@@ -228,6 +250,7 @@
- }
- 
- // CHECK-LABEL: "attr_rng_algorithm_default"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_rng_algorithm_default(%arg0: tensor<f32>) -> (tensor<f32>, tensor<f32>) {
-   %0:2 = "stablehlo.rng_bit_generator"(%arg0) {
-     // CHECK: rng_algorithm = #vhlo<rng_algorithm_v1 DEFAULT>
-@@ -237,6 +260,7 @@
- }
- 
- // CHECK-LABEL: "attr_rng_algorithm_three_fry"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_rng_algorithm_three_fry(%arg0: tensor<f32>) -> (tensor<f32>, tensor<f32>) {
-   %0:2 = "stablehlo.rng_bit_generator"(%arg0) {
-     // CHECK: rng_algorithm = #vhlo<rng_algorithm_v1 THREE_FRY>
-@@ -246,6 +270,7 @@
- }
- 
- // CHECK-LABEL: "attr_rng_algorithm_philox"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_rng_algorithm_philox(%arg0: tensor<f32>) -> (tensor<f32>, tensor<f32>) {
-   %0:2 = "stablehlo.rng_bit_generator"(%arg0) {
-     // CHECK: rng_algorithm = #vhlo<rng_algorithm_v1 PHILOX>
-@@ -255,6 +280,7 @@
- }
- 
- // CHECK-LABEL: "attr_rng_distribution_uniform"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @attr_rng_distribution_uniform(%arg0: tensor<f32>, %arg1: tensor<f32>, %arg2: tensor<0xindex>) -> tensor<f32> {
-   %0 = "stablehlo.rng"(%arg0, %arg1, %arg2) {
-     // CHECK: rng_distribution = #vhlo<rng_distribution_v1 UNIFORM>
-@@ -264,6 +290,7 @@
- }
- 
- // CHECK-LABEL: "attr_rng_distribution_normal"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @attr_rng_distribution_normal(%arg0: tensor<f32>, %arg1: tensor<f32>, %arg2: tensor<0xindex>) -> tensor<f32> {
-   %0 = "stablehlo.rng"(%arg0, %arg1, %arg2) {
-     // CHECK: rng_distribution = #vhlo<rng_distribution_v1 NORMAL>
-@@ -275,6 +302,7 @@
- // ScatterDimensionNumbers aka #stablehlo.scatter is covered below.
- 
- // CHECK-LABEL: "attr_transpose_no_transpose"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_transpose_no_transpose(%arg0: tensor<16x16xf32>, %arg1: tensor<16x16xf32>) ->  tensor<16x16xf32> {
-   %0 = "stablehlo.triangular_solve"(%arg0, %arg1) {
-     left_side = true,
-@@ -287,6 +315,7 @@
- }
- 
- // CHECK-LABEL: "attr_transpose_transpose"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_transpose_transpose(%arg0: tensor<16x16xf32>, %arg1: tensor<16x16xf32>) ->  tensor<16x16xf32> {
-   %0 = "stablehlo.triangular_solve"(%arg0, %arg1) {
-     left_side = true,
-@@ -299,6 +328,7 @@
- }
- 
- // CHECK-LABEL: "attr_transpose_adjoint"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_transpose_adjoint(%arg0: tensor<16x16xf32>, %arg1: tensor<16x16xf32>) ->  tensor<16x16xf32> {
-   %0 = "stablehlo.triangular_solve"(%arg0, %arg1) {
-     left_side = true,
-@@ -313,10 +343,9 @@
- // TypeExtensionsAttr aka #stablehlo.type_extensions is covered below.
- 
- // CHECK-LABEL: "attr_type_extensions_bounds"
--func.func @attr_type_extensions_bounds(
--    %arg0: tensor<?x?xf32, #stablehlo.type_extensions<bounds = [16, ?]>>)
--    -> tensor<?x?xf32, #stablehlo.type_extensions<bounds = [16, ?]>> {
--  // CHECK: "vhlo.return_v1"(%arg0) : (!vhlo.tensor_v1<?x?x!vhlo.f32_v1, #vhlo.type_extensions_v1<bounds = [16, ?]>>) -> ()
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @attr_type_extensions_bounds(%arg0: tensor<?x?xf32, #stablehlo.type_extensions<bounds = [16, ?]>>) -> tensor<?x?xf32, #stablehlo.type_extensions<bounds = [16, ?]>> {
-+  // CHECK: "vhlo.return_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<?x?x!vhlo.f32_v1, #vhlo.type_extensions_v1<bounds = [16, ?]>>) -> ()
-   func.return %arg0 : tensor<?x?xf32, #stablehlo.type_extensions<bounds = [16, ?]>>
- }
- 
-@@ -324,8 +353,9 @@
- // ============ DEFAULTS ============
- 
- // CHECK-LABEL: "default_all_gather"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_all_gather(%arg0: tensor<16x8xf32>) -> tensor<16x16xf32> {
--  //               CHECK: "vhlo.all_gather_v1"(%arg0) <{
-+  //               CHECK: "vhlo.all_gather_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   all_gather_dim = #vhlo.integer_v1<1 : i64>
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0], [1]]> : tensor<2x1xi64>>,
-@@ -339,8 +369,9 @@
- }
- 
- // CHECK-LABEL: "default_all_reduce"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_all_reduce(%arg0: tensor<f32>) -> tensor<f32> {
--  //               CHECK: "vhlo.all_reduce_v1"(%arg0)
-+  //               CHECK: "vhlo.all_reduce_v1"(%[[ARG0]])
-   //          CHECK-SAME: <{
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0], [1]]> : tensor<2x1xi64>>,
-@@ -362,8 +393,9 @@
- }
- 
- // CHECK-LABEL: "default_all_to_all"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_all_to_all(%arg0: tensor<4x16xf32>) -> tensor<16x4xf32> {
--  //               CHECK: "vhlo.all_to_all_v1"(%arg0) <{
-+  //               CHECK: "vhlo.all_to_all_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   //          CHECK-SAME:   concat_dimension = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0, 1, 2, 3]]> : tensor<1x4xi64>>,
-@@ -380,8 +412,9 @@
- }
- 
- // CHECK-LABEL: "default_cholesky"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_cholesky(%arg0: tensor<1x16x16xf32>) -> tensor<1x16x16xf32> {
--  //      CHECK: "vhlo.cholesky_v1"(%arg0) <{
-+  //      CHECK: "vhlo.cholesky_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   lower = #vhlo.bool_v1<false>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<1x16x16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<1x16x16x!vhlo.f32_v1>
-   %0 = "stablehlo.cholesky"(%arg0) : (tensor<1x16x16xf32>) -> tensor<1x16x16xf32>
-@@ -389,8 +422,9 @@
- }
- 
- // CHECK-LABEL: "default_collective_permute"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_collective_permute(%arg0: tensor<16x8xf32>) -> tensor<16x8xf32> {
--  //               CHECK: "vhlo.collective_permute_v1"(%arg0) <{
-+  //               CHECK: "vhlo.collective_permute_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME{LITERAL}:   source_target_pairs = #vhlo.tensor_v1<dense<[[0, 1], [1, 2], [2, 3]]> : tensor<3x2xi64>>
-   //          CHECK-SAME: }> : (!vhlo.tensor_v1<16x8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x8x!vhlo.f32_v1>
-@@ -401,8 +435,9 @@
- }
- 
- // CHECK-LABEL: "default_compare"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @default_compare(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
--  //      CHECK: "vhlo.compare_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.compare_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   compare_type = #vhlo<comparison_type_v1 NOTYPE>,
-   // CHECK-SAME:   comparison_direction = #vhlo<comparison_direction_v1 EQ>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-@@ -413,8 +448,9 @@
- }
- 
- // CHECK-LABEL: "default_convolution"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @default_convolution(%arg0: tensor<1x8x8x207xf32>, %arg1: tensor<3x3x207x16xf32>) -> tensor<1x6x6x16xf32> {
--  //      CHECK: "vhlo.convolution_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.convolution_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   batch_group_count = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   feature_group_count = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   input_batch_dimension = #vhlo.integer_v1<0 : i64>,
-@@ -442,8 +478,9 @@
- }
- 
- // CHECK-LABEL: "default_custom_call"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_custom_call(%arg0: tensor<f32>) -> tensor<f32> {
--  //      CHECK: "vhlo.custom_call_v1"(%arg0) <{
-+  //      CHECK: "vhlo.custom_call_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   api_version = #vhlo<api_version_v1 API_VERSION_ORIGINAL>,
-   // CHECK-SAME:   backend_config = #vhlo.string_v1<"">,
-   // CHECK-SAME:   call_target_name = #vhlo.string_v1<"foo">,
-@@ -460,8 +497,9 @@
- }
- 
- // CHECK-LABEL: "default_dot_general"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @default_dot_general(%arg0: tensor<8x8x16xf32>, %arg1: tensor<8x16x8xf32>) -> tensor<8x8x8xf32> {
--  //      CHECK: "vhlo.dot_general_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.dot_general_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   lhs_batching_dimensions = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>,
-   // CHECK-SAME:   lhs_contracting_dimensions = #vhlo.tensor_v1<dense<2> : tensor<1xi64>>,
-   // CHECK-SAME:   precision_config = #vhlo.array_v1<[#vhlo<precision_v1 DEFAULT>, #vhlo<precision_v1 DEFAULT>]>,
-@@ -480,8 +518,9 @@
- }
- 
- // CHECK-LABEL: "default_dot"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @default_dot(%arg0: tensor<8x16xf32>, %arg1: tensor<16x8xf32>) -> tensor<8x8xf32> {
--  //      CHECK: "vhlo.dot_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.dot_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   precision_config = #vhlo.array_v1<[#vhlo<precision_v1 DEFAULT>, #vhlo<precision_v1 DEFAULT>]>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<8x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<8x8x!vhlo.f32_v1>
-   %0 = "stablehlo.dot"(%arg0, %arg1) : (tensor<8x16xf32>, tensor<16x8xf32>) -> tensor<8x8xf32>
-@@ -489,8 +528,9 @@
- }
- 
- // CHECK-LABEL: "default_dynamic_broadcast_in_dim"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @default_dynamic_broadcast_in_dim(%arg0: tensor<?x?xf32>, %arg1: tensor<2xindex>) -> tensor<?x?xf32> {
--  //      CHECK: "vhlo.dynamic_broadcast_in_dim_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.dynamic_broadcast_in_dim_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   broadcast_dimensions = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-   // CHECK-SAME:   known_expanding_dimensions = #vhlo.tensor_v1<dense<> : tensor<0xi64>>,
-   // CHECK-SAME:   known_nonexpanding_dimensions = #vhlo.tensor_v1<dense<> : tensor<0xi64>>
-@@ -502,8 +542,9 @@
- }
- 
- // CHECK-LABEL: "default_dynamic_conv"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @default_dynamic_conv(%arg0: tensor<1x8x8x207xf32>, %arg1: tensor<3x3x207x16xf32>, %arg2: tensor<4xi32>) -> tensor<1x?x?x16xf32> {
--  //      CHECK: "vhlo.dynamic_conv_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.dynamic_conv_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   batch_group_count = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   feature_group_count = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   input_batch_dimension = #vhlo.integer_v1<0 : i64>,
-@@ -531,8 +572,9 @@
- }
- 
- // CHECK-LABEL: "default_dynamic_gather"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @default_dynamic_gather(%arg0 : tensor<2x4x9xf32>, %arg1 : tensor<1x5x2xi32>, %arg2 : tensor<3xi32>) -> tensor<1x5x8xf32> {
--  //      CHECK: "vhlo.dynamic_gather_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.dynamic_gather_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   collapsed_slice_dims = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-   // CHECK-SAME:   index_vector_dim = #vhlo.integer_v1<2 : i64>,
-   // CHECK-SAME:   indices_are_sorted = #vhlo.bool_v1<false>,
-@@ -558,15 +600,16 @@
-   // CHECK-SAME:   sym_name = #vhlo.string_v1<"default_func">,
-   // CHECK-SAME:   sym_visibility = #vhlo.string_v1<"">
-   // CHECK-SAME: }> ({
--  // CHECK-NEXT:   ^[[BB:bb.*]](%arg0: !vhlo.tensor_v1<!vhlo.f32_v1>):
--  // CHECK-NEXT:     "vhlo.return_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-+  // CHECK-NEXT:   ^[[BB:bb.*]](%[[ARG0:.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
-+  // CHECK-NEXT:     "vhlo.return_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-   // CHECK-NEXT: }) : () -> ()
-   func.return %arg0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "dynamic_gather"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @dynamic_gather(%arg0 : tensor<2x4x9xf32>, %arg1 : tensor<1x5x2xi32>) -> tensor<1x5x1xf32> {
--  //      CHECK: "vhlo.gather_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.gather_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   collapsed_slice_dims = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-   // CHECK-SAME:   index_vector_dim = #vhlo.integer_v1<2 : i64>,
-   // CHECK-SAME:   indices_are_sorted = #vhlo.bool_v1<false>,
-@@ -587,8 +630,9 @@
- }
- 
- // CHECK-LABEL: "default_infeed"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_infeed(%arg0: !stablehlo.token) -> (tensor<f32>, !stablehlo.token) {
--  //               CHECK: "vhlo.infeed_v1"(%arg0) <{
-+  //               CHECK: "vhlo.infeed_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   infeed_config = #vhlo.string_v1<"">,
-   // CHECK-SAME{LITERAL}:   layout = #vhlo.array_v1<[]>
-   //          CHECK-SAME: }> : (!vhlo.token_v1) -> (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.token_v1)
-@@ -597,8 +641,9 @@
- }
- 
- // CHECK-LABEL: "default_outfeed"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @default_outfeed(%arg0: tensor<f32>, %arg1: !stablehlo.token) -> !stablehlo.token {
--  //      CHECK: "vhlo.outfeed_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.outfeed_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   outfeed_config = #vhlo.string_v1<"">
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.token_v1) -> !vhlo.token_v1
-   %0 = "stablehlo.outfeed"(%arg0, %arg1) : (tensor<f32>, !stablehlo.token) -> !stablehlo.token
-@@ -606,8 +651,9 @@
- }
- 
- // CHECK-LABEL: "default_recv"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_recv(%arg0: !stablehlo.token) -> (tensor<f32>, !stablehlo.token) {
--  //      CHECK: "vhlo.recv_v1"(%arg0) <{
-+  //      CHECK: "vhlo.recv_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME:   channel_type = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   is_host_transfer = #vhlo.bool_v1<false>
-@@ -619,8 +665,9 @@
- }
- 
- // CHECK-LABEL: "default_send"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @default_send(%arg0: tensor<f32>, %arg1: !stablehlo.token) -> !stablehlo.token {
--  //      CHECK: "vhlo.send_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.send_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME:   channel_type = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   is_host_transfer = #vhlo.bool_v1<false>
-@@ -632,8 +679,9 @@
- }
- 
- // CHECK-LABEL: "default_reduce_scatter"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_reduce_scatter(%arg0: tensor<16xf32>) -> tensor<16xf32> {
--  //               CHECK: "vhlo.reduce_scatter_v1"(%arg0) <{
-+  //               CHECK: "vhlo.reduce_scatter_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0], [1]]> : tensor<2x1xi64>>,
-   //          CHECK-SAME:   scatter_dimension = #vhlo.integer_v1<0 : i64>
-@@ -655,8 +703,9 @@
- }
- 
- // CHECK-LABEL: "default_reduce_window"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @default_reduce_window(%arg0: tensor<2x17x31x7xf32>, %arg1: tensor<f32>) -> tensor<2x16x30x7xf32> {
--  //               CHECK: "vhlo.reduce_window_v1"(%arg0, %arg1)  <{
-+  //               CHECK: "vhlo.reduce_window_v1"(%[[ARG0]], %[[ARG1]])  <{
-   //          CHECK-SAME:   base_dilations = #vhlo.tensor_v1<dense<1> : tensor<4xi64>>,
-   // CHECK-SAME{LITERAL}:   padding = #vhlo.tensor_v1<dense<0> : tensor<4x2xi64>>,
-   //          CHECK-SAME:   window_dilations = #vhlo.tensor_v1<dense<1> : tensor<4xi64>>,
-@@ -678,8 +727,9 @@
- }
- 
- // CHECK-LABEL: "default_scatter"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @default_scatter(%arg0: tensor<200x100x300xf32>, %arg1: tensor<10x2xi32>, %arg2: tensor<10x300xf32>) -> tensor<200x100x300xf32> {
--  //      CHECK: "vhlo.scatter_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.scatter_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   index_vector_dim = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   indices_are_sorted = #vhlo.bool_v1<false>,
-   // CHECK-SAME:   inserted_window_dims = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-@@ -707,8 +757,9 @@
- }
- 
- // CHECK-LABEL: "default_select_and_scatter"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @default_select_and_scatter(%arg0: tensor<10x24x24x64xf32>, %arg1: tensor<10x23x23x64xf32>, %arg2: tensor<f32>) -> tensor<10x24x24x64xf32> {
--  //      CHECK: "vhlo.select_and_scatter_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.select_and_scatter_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   padding = #vhlo.tensor_v1<dense<0> : tensor<4x2xi64>>,
-   // CHECK-SAME:   window_dimensions = #vhlo.tensor_v1<dense<[1, 2, 2, 1]> : tensor<4xi64>>,
-   // CHECK-SAME:   window_strides = #vhlo.tensor_v1<dense<1> : tensor<4xi64>>
-@@ -736,8 +787,9 @@
- }
- 
- // CHECK-LABEL: "default_sort"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_sort(%arg0: tensor<16xf32>) -> tensor<16xf32> {
--  //      CHECK: "vhlo.sort_v1"(%arg0) <{
-+  //      CHECK: "vhlo.sort_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   dimension = #vhlo.integer_v1<-1 : i64>
-   // CHECK-SAME:   is_stable = #vhlo.bool_v1<false>
-   // CHECK-SAME: }> ({
-@@ -756,29 +808,33 @@
- // ============ OPS ============
- 
- // CHECK-LABEL: "op_abs"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_abs(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.abs_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.abs_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.abs"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_add"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_add(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_after_all"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_after_all(%arg0: !stablehlo.token) -> !stablehlo.token {
--  // CHECK: "vhlo.after_all_v1"(%arg0) : (!vhlo.token_v1) -> !vhlo.token_v1
-+  // CHECK: "vhlo.after_all_v1"(%[[ARG0]]) : (!vhlo.token_v1) -> !vhlo.token_v1
-   %0 = "stablehlo.after_all"(%arg0) : (!stablehlo.token) -> !stablehlo.token
-   func.return %0 : !stablehlo.token
- }
- 
- // CHECK-LABEL: "op_all_gather"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_all_gather(%arg0: tensor<16x8xf32>) -> tensor<16x16xf32> {
--  //               CHECK: "vhlo.all_gather_v1"(%arg0) <{
-+  //               CHECK: "vhlo.all_gather_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   all_gather_dim = #vhlo.integer_v1<1 : i64>
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0], [1]]> : tensor<2x1xi64>>,
-@@ -794,8 +850,9 @@
- }
- 
- // CHECK-LABEL: "op_all_reduce"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_all_reduce(%arg0: tensor<f32>) -> tensor<f32> {
--  //               CHECK: "vhlo.all_reduce_v1"(%arg0) <{
-+  //               CHECK: "vhlo.all_reduce_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0], [1]]> : tensor<2x1xi64>>,
-   //          CHECK-SAME:   use_global_device_ids = #vhlo.bool_v1<true>
-@@ -817,8 +874,9 @@
- }
- 
- // CHECK-LABEL: "op_all_to_all"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_all_to_all(%arg0: tensor<4x16xf32>) -> tensor<16x4xf32> {
--  //               CHECK: "vhlo.all_to_all_v1"(%arg0) <{
-+  //               CHECK: "vhlo.all_to_all_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<1 : i64>,
-   //          CHECK-SAME:   concat_dimension = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0, 1, 2, 3]]> : tensor<1x4xi64>>,
-@@ -836,22 +894,25 @@
- }
- 
- // CHECK-LABEL: "op_and"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_and(%arg0: tensor<i1>, %arg1: tensor<i1>) -> tensor<i1> {
--  // CHECK: "vhlo.and_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-+  // CHECK: "vhlo.and_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-   %0 = "stablehlo.and"(%arg0, %arg1) : (tensor<i1>, tensor<i1>) -> tensor<i1>
-   func.return %0 : tensor<i1>
- }
- 
- // CHECK-LABEL: "op_atan2"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_atan2(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.atan2_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.atan2_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.atan2"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_batch_norm_grad"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}}, %[[ARG3:.*]]: {{.*}}, %[[ARG4:.*]]: {{.*}})
- func.func @op_batch_norm_grad(%arg0: tensor<16x16x16x16xf32>, %arg1: tensor<16xf32>, %arg2: tensor<16xf32>, %arg3: tensor<16xf32>, %arg4: tensor<16x16x16x16xf32>) -> (tensor<16x16x16x16xf32>, tensor<16xf32>, tensor<16xf32>) {
--  //      CHECK: "vhlo.batch_norm_grad_v1"(%arg0, %arg1, %arg2, %arg3, %arg4) <{
-+  //      CHECK: "vhlo.batch_norm_grad_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]], %[[ARG3]], %[[ARG4]]) <{
-   // CHECK-SAME:   epsilon = #vhlo.float_v1<1.000000e-03 : !vhlo.f32_v1>,
-   // CHECK-SAME:   feature_index = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<16x16x16x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x16x16x16x!vhlo.f32_v1>) -> (!vhlo.tensor_v1<16x16x16x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>)
-@@ -863,8 +924,9 @@
- }
- 
- // CHECK-LABEL: "op_batch_norm_inference"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}}, %[[ARG3:.*]]: {{.*}}, %[[ARG4:.*]]: {{.*}})
- func.func @op_batch_norm_inference(%arg0: tensor<16x16x16x16xf32>, %arg1: tensor<16xf32>, %arg2: tensor<16xf32>, %arg3: tensor<16xf32>, %arg4: tensor<16xf32>) -> tensor<16x16x16x16xf32> {
--  //      CHECK: "vhlo.batch_norm_inference_v1"(%arg0, %arg1, %arg2, %arg3, %arg4) <{
-+  //      CHECK: "vhlo.batch_norm_inference_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]], %[[ARG3]], %[[ARG4]]) <{
-   // CHECK-SAME:   epsilon = #vhlo.float_v1<1.000000e-03 : !vhlo.f32_v1>,
-   // CHECK-SAME:   feature_index = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<16x16x16x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x16x16x16x!vhlo.f32_v1>
-@@ -876,8 +938,9 @@
- }
- 
- // CHECK-LABEL: "op_batch_norm_training"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_batch_norm_training(%arg0: tensor<16x16x16x16xf32>, %arg1: tensor<16xf32>, %arg2: tensor<16xf32>) -> (tensor<16x16x16x16xf32>, tensor<16xf32>, tensor<16xf32>) {
--  //      CHECK: "vhlo.batch_norm_training_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.batch_norm_training_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   epsilon = #vhlo.float_v1<1.000000e-03 : !vhlo.f32_v1>,
-   // CHECK-SAME:   feature_index = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<16x16x16x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>) -> (!vhlo.tensor_v1<16x16x16x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>)
-@@ -889,15 +952,17 @@
- }
- 
- // CHECK-LABEL: "op_bitcast_convert"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_bitcast_convert(%arg0: tensor<i32>) -> tensor<f32> {
--  // CHECK: "vhlo.bitcast_convert_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.bitcast_convert_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.bitcast_convert"(%arg0) : (tensor<i32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_broadcast_in_dim"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_broadcast_in_dim(%arg0: tensor<16xf32>) -> tensor<16x16xf32> {
--  //      CHECK: "vhlo.broadcast_in_dim_v1"(%arg0) <{
-+  //      CHECK: "vhlo.broadcast_in_dim_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   broadcast_dimensions = #vhlo.tensor_v1<dense<1> : tensor<1xi64>>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x16x!vhlo.f32_v1>
-   %0 = "stablehlo.broadcast_in_dim"(%arg0) {
-@@ -907,8 +972,9 @@
- }
- 
- // CHECK-LABEL: "op_broadcast"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_broadcast(%arg0: tensor<16xf32>) -> tensor<16x16xf32> {
--  //      CHECK: "vhlo.broadcast_v1"(%arg0) <{
-+  //      CHECK: "vhlo.broadcast_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   broadcast_sizes = #vhlo.tensor_v1<dense<16> : tensor<1xi64>>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x16x!vhlo.f32_v1>
-   %0 = "stablehlo.broadcast"(%arg0) {
-@@ -918,9 +984,10 @@
- }
- 
- // CHECK-LABEL: "op_case"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_case(%arg0: tensor<i32>, %arg1: tensor<f32>) -> tensor<f32> {
--  //      CHECK: "vhlo.case_v1"(%arg0) ({
--  // CHECK-NEXT:   "vhlo.return_v1"(%arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-+  //      CHECK: "vhlo.case_v1"(%[[ARG0]]) ({
-+  // CHECK-NEXT:   "vhlo.return_v1"(%[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-   // CHECK-NEXT: }) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.case"(%arg0) ({
-     "stablehlo.return"(%arg1) : (tensor<f32>) -> ()
-@@ -929,22 +996,25 @@
- }
- 
- // CHECK-LABEL: "op_cbrt"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_cbrt(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.cbrt_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.cbrt_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.cbrt"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_ceil"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_ceil(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.ceil_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.ceil_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.ceil"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_cholesky"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_cholesky(%arg0: tensor<1x16x16xf32>) -> tensor<1x16x16xf32> {
--  //      CHECK: "vhlo.cholesky_v1"(%arg0) <{
-+  //      CHECK: "vhlo.cholesky_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   lower = #vhlo.bool_v1<true>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<1x16x16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<1x16x16x!vhlo.f32_v1>
-   %0 = "stablehlo.cholesky"(%arg0) {
-@@ -954,22 +1024,25 @@
- }
- 
- // CHECK-LABEL: "op_clamp"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_clamp(%arg0: tensor<f32>, %arg1: tensor<f32>, %arg2: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.clamp_v1"(%arg0, %arg1, %arg2) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.clamp_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.clamp"(%arg0, %arg1, %arg2) : (tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_count_leading_zeros"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_count_leading_zeros(%arg0: tensor<i32>) -> tensor<i32> {
--  // CHECK: "vhlo.count_leading_zeros_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-+  // CHECK: "vhlo.count_leading_zeros_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-   %0 = "stablehlo.count_leading_zeros"(%arg0) : (tensor<i32>) -> tensor<i32>
-   func.return %0 : tensor<i32>
- }
- 
- // CHECK-LABEL: "op_collective_permute"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_collective_permute(%arg0: tensor<16x8xf32>) -> tensor<16x8xf32> {
--  //               CHECK: "vhlo.collective_permute_v1"(%arg0) <{
-+  //               CHECK: "vhlo.collective_permute_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME{LITERAL}:   source_target_pairs = #vhlo.tensor_v1<dense<[[0, 1], [1, 2], [2, 3]]> : tensor<3x2xi64>>
-   //          CHECK-SAME: }> : (!vhlo.tensor_v1<16x8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x8x!vhlo.f32_v1>
-@@ -981,8 +1054,9 @@
- }
- 
- // CHECK-LABEL: "op_compare"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_compare(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
--  //      CHECK: "vhlo.compare_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.compare_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   compare_type = #vhlo<comparison_type_v1 TOTALORDER>,
-   // CHECK-SAME:   comparison_direction = #vhlo<comparison_direction_v1 EQ>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-@@ -994,15 +1068,17 @@
- }
- 
- // CHECK-LABEL: "op_complex"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_complex(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<complex<f32>> {
--  // CHECK: "vhlo.complex_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>
-+  // CHECK: "vhlo.complex_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>
-   %0 = "stablehlo.complex"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<complex<f32>>
-   func.return %0 : tensor<complex<f32>>
- }
- 
- // CHECK-LABEL: "op_concatenate"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_concatenate(%arg0: tensor<8xf32>, %arg1: tensor<8xf32>) -> tensor<16xf32> {
--  //      CHECK: "vhlo.concatenate_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.concatenate_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   dimension = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<8x!vhlo.f32_v1>, !vhlo.tensor_v1<8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x!vhlo.f32_v1>
-   %0 = "stablehlo.concatenate"(%arg0, %arg1) {
-@@ -1012,6 +1088,7 @@
- }
- 
- // CHECK-LABEL: "op_constant"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_constant(%arg0: tensor<f32>) -> tensor<f32> {
-   //      CHECK: "vhlo.constant_v1"() <{
-   // CHECK-SAME:   value = #vhlo.tensor_v1<dense<0.000000e+00> : tensor<f32>>
-@@ -1023,15 +1100,17 @@
- }
- 
- // CHECK-LABEL: "op_convert"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_convert(%arg0: tensor<i32>) -> tensor<f32> {
--  // CHECK: "vhlo.convert_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.convert_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.convert"(%arg0) : (tensor<i32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_convolution"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_convolution(%arg0: tensor<1x8x8x207xf32>, %arg1: tensor<3x3x207x16xf32>) -> tensor<1x7x7x16xf32> {
--  //      CHECK: "vhlo.convolution_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.convolution_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   batch_group_count = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   feature_group_count = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   input_batch_dimension = #vhlo.integer_v1<0 : i64>,
-@@ -1065,8 +1144,9 @@
- }
- 
- // CHECK-LABEL: "op_cosine"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_cosine(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.cosine_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.cosine_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.cosine"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
-@@ -1079,8 +1159,9 @@
- }
- 
- // CHECK-LABEL: "op_cross_replica_sum"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_cross_replica_sum(%arg0: tensor<f32>) -> tensor<f32> {
--  //               CHECK: "vhlo.cross-replica-sum_v1"(%arg0) <{
-+  //               CHECK: "vhlo.cross-replica-sum_v1"(%[[ARG0]]) <{
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0], [1]]> : tensor<2x1xi64>>
-   //          CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.cross-replica-sum"(%arg0) {
-@@ -1090,8 +1171,9 @@
- }
- 
- // CHECK-LABEL: "op_custom_call"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_custom_call(%arg0: tensor<f32>) -> tensor<f32> {
--  //      CHECK: "vhlo.custom_call_v1"(%arg0) <{
-+  //      CHECK: "vhlo.custom_call_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   api_version = #vhlo<api_version_v1 API_VERSION_STATUS_RETURNING>,
-   // CHECK-SAME:   backend_config = #vhlo.string_v1<"\08\03\1A\02">,
-   // CHECK-SAME:   call_target_name = #vhlo.string_v1<"foo">,
-@@ -1122,15 +1204,17 @@
- }
- 
- // CHECK-LABEL: "op_divide"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_divide(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.divide_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.divide_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.divide"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_dot_general"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_dot_general(%arg0: tensor<8x8x16xf32>, %arg1: tensor<8x16x8xf32>) -> tensor<8x8x8xf32> {
--  //      CHECK: "vhlo.dot_general_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.dot_general_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   lhs_batching_dimensions = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>,
-   // CHECK-SAME:   lhs_contracting_dimensions = #vhlo.tensor_v1<dense<2> : tensor<1xi64>>,
-   // CHECK-SAME:   precision_config = #vhlo.array_v1<[#vhlo<precision_v1 HIGHEST>, #vhlo<precision_v1 HIGHEST>]>,
-@@ -1150,8 +1234,9 @@
- }
- 
- // CHECK-LABEL: "op_dot"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_dot(%arg0: tensor<8x16xf32>, %arg1: tensor<16x8xf32>) -> tensor<8x8xf32> {
--  //      CHECK: "vhlo.dot_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.dot_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   precision_config = #vhlo.array_v1<[#vhlo<precision_v1 HIGHEST>, #vhlo<precision_v1 HIGHEST>]>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<8x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<8x8x!vhlo.f32_v1>
-   %0 = "stablehlo.dot"(%arg0, %arg1) {
-@@ -1161,8 +1246,9 @@
- }
- 
- // CHECK-LABEL: "op_dynamic_broadcast_in_dim"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_dynamic_broadcast_in_dim(%arg0: tensor<?x?xf32>, %arg1: tensor<2xindex>) -> tensor<?x?xf32> {
--  //      CHECK: "vhlo.dynamic_broadcast_in_dim_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.dynamic_broadcast_in_dim_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   broadcast_dimensions = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-   // CHECK-SAME:   known_expanding_dimensions = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>,
-   // CHECK-SAME:   known_nonexpanding_dimensions = #vhlo.tensor_v1<dense<1> : tensor<1xi64>>
-@@ -1176,8 +1262,9 @@
- }
- 
- // CHECK-LABEL: "op_dynamic_conv"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_dynamic_conv(%arg0: tensor<1x8x8x207xf32>, %arg1: tensor<3x3x207x16xf32>, %arg2: tensor<4xi32>) -> tensor<1x?x?x16xf32> {
--  //      CHECK: "vhlo.dynamic_conv_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.dynamic_conv_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   batch_group_count = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   feature_group_count = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   input_batch_dimension = #vhlo.integer_v1<0 : i64>,
-@@ -1211,8 +1298,9 @@
- }
- 
- // CHECK-LABEL: "op_dynamic_gather"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_dynamic_gather(%arg0 : tensor<2x4x9xf32>, %arg1 : tensor<1x5x2xi32>, %arg2 : tensor<3xi32>) -> tensor<1x5x8xf32> {
--  //      CHECK: "vhlo.dynamic_gather_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.dynamic_gather_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   collapsed_slice_dims = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-   // CHECK-SAME:   index_vector_dim = #vhlo.integer_v1<2 : i64>,
-   // CHECK-SAME:   indices_are_sorted = #vhlo.bool_v1<true>,
-@@ -1232,8 +1320,9 @@
- }
- 
- // CHECK-LABEL: "op_dynamic_iota"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_dynamic_iota(%arg0: tensor<1xindex>) -> tensor<?xf32> {
--  //      CHECK: "vhlo.dynamic_iota_v1"(%arg0) <{
-+  //      CHECK: "vhlo.dynamic_iota_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   iota_dimension = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<1x!vhlo.index_v1>) -> !vhlo.tensor_v1<?x!vhlo.f32_v1>
-   %0 = "stablehlo.dynamic_iota"(%arg0) {
-@@ -1243,22 +1332,25 @@
- }
- 
- // CHECK-LABEL: "op_dynamic_pad"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}}, %[[ARG3:.*]]: {{.*}}, %[[ARG4:.*]]: {{.*}})
- func.func @op_dynamic_pad(%arg0: tensor<?xf32>, %arg1: tensor<f32>, %arg2: tensor<1xindex>, %arg3: tensor<1xindex>, %arg4: tensor<1xindex>) -> tensor<?xf32> {
--  // CHECK: "vhlo.dynamic_pad_v1"(%arg0, %arg1, %arg2, %arg3, %arg4) : (!vhlo.tensor_v1<?x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>) -> !vhlo.tensor_v1<?x!vhlo.f32_v1>
-+  // CHECK: "vhlo.dynamic_pad_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]], %[[ARG3]], %[[ARG4]]) : (!vhlo.tensor_v1<?x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>) -> !vhlo.tensor_v1<?x!vhlo.f32_v1>
-   %0 = "stablehlo.dynamic_pad"(%arg0, %arg1, %arg2, %arg3, %arg4) : (tensor<?xf32>, tensor<f32>, tensor<1xindex>, tensor<1xindex>, tensor<1xindex>) -> tensor<?xf32>
-   func.return %0 : tensor<?xf32>
- }
- 
- // CHECK-LABEL: "op_dynamic_reshape"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_dynamic_reshape(%arg0: tensor<16xf32>, %arg1: tensor<2xindex>) -> tensor<?x?xf32> {
--  // CHECK: "vhlo.dynamic_reshape_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<2x!vhlo.index_v1>) -> !vhlo.tensor_v1<?x?x!vhlo.f32_v1>
-+  // CHECK: "vhlo.dynamic_reshape_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<2x!vhlo.index_v1>) -> !vhlo.tensor_v1<?x?x!vhlo.f32_v1>
-   %0 = "stablehlo.dynamic_reshape"(%arg0, %arg1) : (tensor<16xf32>, tensor<2xindex>) -> tensor<?x?xf32>
-   func.return %0 : tensor<?x?xf32>
- }
- 
- // CHECK-LABEL: "op_dynamic_slice"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_dynamic_slice(%arg0: tensor<16xf32>, %arg1: tensor<i64>) -> tensor<4xf32> {
--  //      CHECK: "vhlo.dynamic_slice_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.dynamic_slice_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   slice_sizes = #vhlo.tensor_v1<dense<4> : tensor<1xi64>>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.i64_v1>) -> !vhlo.tensor_v1<4x!vhlo.f32_v1>
-   %0 = "stablehlo.dynamic_slice"(%arg0, %arg1) {
-@@ -1268,15 +1360,17 @@
- }
- 
- // CHECK-LABEL: "op_dynamic_update_slice"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_dynamic_update_slice(%arg0: tensor<16xf32>, %arg1: tensor<4xf32>, %arg2: tensor<i64>) -> tensor<16xf32> {
--  // CHECK: "vhlo.dynamic_update_slice_v1"(%arg0, %arg1, %arg2) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<4x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.i64_v1>) -> !vhlo.tensor_v1<16x!vhlo.f32_v1>
-+  // CHECK: "vhlo.dynamic_update_slice_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<4x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.i64_v1>) -> !vhlo.tensor_v1<16x!vhlo.f32_v1>
-   %0 = "stablehlo.dynamic_update_slice"(%arg0, %arg1, %arg2) : (tensor<16xf32>, tensor<4xf32>, tensor<i64>) -> tensor<16xf32>
-   func.return %0 : tensor<16xf32>
- }
- 
- // CHECK-LABEL: "op_einsum"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_einsum(%arg0: tensor<8x16xf32>, %arg1: tensor<16x8xf32>) -> tensor<8x8xf32> {
--  //      CHECK: "vhlo.einsum_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.einsum_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   einsum_config = #vhlo.string_v1<"ab,bc->ac">
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<8x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<8x8x!vhlo.f32_v1>
-   %0 = "stablehlo.einsum"(%arg0, %arg1) {
-@@ -1286,22 +1380,25 @@
- }
- 
- // CHECK-LABEL: "op_exponential_minus_one"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_exponential_minus_one(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.exponential_minus_one_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.exponential_minus_one_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.exponential_minus_one"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_exponential"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_exponential(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.exponential_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.exponential_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.exponential"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_fft"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_fft(%arg0: tensor<16xcomplex<f32>>) -> tensor<16xcomplex<f32>> {
--  //      CHECK: "vhlo.fft_v1"(%arg0) <{
-+  //      CHECK: "vhlo.fft_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   fft_length = #vhlo.tensor_v1<dense<16> : tensor<1xi64>>,
-   // CHECK-SAME:   fft_type = #vhlo<fft_type_v1 FFT>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<16x!vhlo.complex_v1<!vhlo.f32_v1>>) -> !vhlo.tensor_v1<16x!vhlo.complex_v1<!vhlo.f32_v1>>
-@@ -1313,8 +1410,9 @@
- }
- 
- // CHECK-LABEL: "op_floor"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_floor(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.floor_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.floor_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.floor"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
-@@ -1327,16 +1425,17 @@
-   // CHECK-SAME:   sym_name = #vhlo.string_v1<"op_func">,
-   // CHECK-SAME:   sym_visibility = #vhlo.string_v1<"private">
-   // CHECK-SAME: }> ({
--  // CHECK-NEXT:   ^[[BB:bb.*]](%arg0: !vhlo.tensor_v1<!vhlo.f32_v1>):
--  // CHECK-NEXT:     "vhlo.return_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-+  // CHECK-NEXT:   ^[[BB:bb.*]](%[[ARG0:.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
-+  // CHECK-NEXT:     "vhlo.return_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-   // CHECK-NEXT: }) : () -> ()
- 
-   func.return %arg0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_gather"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_gather(%arg0 : tensor<2x4x9xf32>, %arg1 : tensor<1x5x2xi32>) -> tensor<1x5x1xf32> {
--  //      CHECK: "vhlo.gather_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.gather_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   collapsed_slice_dims = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-   // CHECK-SAME:   index_vector_dim = #vhlo.integer_v1<2 : i64>,
-   // CHECK-SAME:   indices_are_sorted = #vhlo.bool_v1<true>,
-@@ -1358,8 +1457,9 @@
- }
- 
- // CHECK-LABEL: "op_get_dimension_size"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_get_dimension_size(%arg0: tensor<?xf32>) -> tensor<i32> {
--  //      CHECK: "vhlo.get_dimension_size_v1"(%arg0) <{
-+  //      CHECK: "vhlo.get_dimension_size_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   dimension = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<?x!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-   %0 = "stablehlo.get_dimension_size"(%arg0) {
-@@ -1369,8 +1469,9 @@
- }
- 
- // CHECK-LABEL: "op_get_tuple_element"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_get_tuple_element(%arg0: tuple<tensor<f32>, tensor<i32>>) -> tensor<f32> {
--  //      CHECK: "vhlo.get_tuple_element_v1"(%arg0) <{
-+  //      CHECK: "vhlo.get_tuple_element_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   index = #vhlo.integer_v1<0 : i32>
-   // CHECK-SAME: }> : (!vhlo.tuple_v1<!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.get_tuple_element"(%arg0) {
-@@ -1380,11 +1481,12 @@
- }
- 
- // CHECK-LABEL: "op_if"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_if(%arg0: tensor<i1>, %arg1: tensor<f32>, %arg2: tensor<f32>) -> tensor<f32> {
--  //      CHECK: "vhlo.if_v1"(%arg0) ({
--  // CHECK-NEXT:   "vhlo.return_v1"(%arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-+  //      CHECK: "vhlo.if_v1"(%[[ARG0]]) ({
-+  // CHECK-NEXT:   "vhlo.return_v1"(%[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-   // CHECK-NEXT: }, {
--  // CHECK-NEXT:   "vhlo.return_v1"(%arg2) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-+  // CHECK-NEXT:   "vhlo.return_v1"(%[[ARG2]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-   // CHECK-NEXT: }) : (!vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.if"(%arg0) ({
-     "stablehlo.return"(%arg1) : (tensor<f32>) -> ()
-@@ -1395,15 +1497,17 @@
- }
- 
- // CHECK-LABEL: "op_imag"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_imag(%arg0: tensor<complex<f32>>) -> tensor<f32> {
--  // CHECK: "vhlo.imag_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.imag_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.imag"(%arg0) : (tensor<complex<f32>>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_infeed"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_infeed(%arg0: !stablehlo.token) -> (tensor<f32>, !stablehlo.token) {
--  //               CHECK: "vhlo.infeed_v1"(%arg0) <{
-+  //               CHECK: "vhlo.infeed_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   infeed_config = #vhlo.string_v1<"foo">,
-   // CHECK-SAME{LITERAL}:   layout = #vhlo.array_v1<[#vhlo.array_v1<[]>]>
-   //          CHECK-SAME: }> : (!vhlo.token_v1) -> (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.token_v1)
-@@ -1426,36 +1530,41 @@
- }
- 
- // CHECK-LABEL: "op_is_finite"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_is_finite(%arg0: tensor<f32>) -> tensor<i1> {
--  // CHECK: "vhlo.is_finite_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-+  // CHECK: "vhlo.is_finite_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-   %0 = "stablehlo.is_finite"(%arg0) : (tensor<f32>) -> tensor<i1>
-   func.return %0 : tensor<i1>
- }
- 
- // CHECK-LABEL: "op_log"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_log(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.log_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.log_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.log"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_log_plus_one"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_log_plus_one(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.log_plus_one_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.log_plus_one_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.log_plus_one"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_logistic"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_logistic(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.logistic_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.logistic_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.logistic"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_map"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_map(%arg0: tensor<16xf32>) -> tensor<16xf32> {
--  //      CHECK: "vhlo.map_v1"(%arg0) <{
-+  //      CHECK: "vhlo.map_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   dimensions = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>
-   // CHECK-SAME: }> ({
-   // CHECK-NEXT:   ^[[BB:bb.*]](%[[ARG1:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
-@@ -1473,57 +1582,65 @@
- }
- 
- // CHECK-LABEL: "op_maximum"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_maximum(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.maximum_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.maximum_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.maximum"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_minimum"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_minimum(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.minimum_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.minimum_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.minimum"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_multiply"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_multiply(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.multiply_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.multiply_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.multiply"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_negate"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_negate(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.negate_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.negate_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.negate"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_not"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_not(%arg0: tensor<i1>) -> tensor<i1> {
--  // CHECK: "vhlo.not_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-+  // CHECK: "vhlo.not_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-   %0 = "stablehlo.not"(%arg0) : (tensor<i1>) -> tensor<i1>
-   func.return %0 : tensor<i1>
- }
- 
- // CHECK-LABEL: "op_optimization_barrier"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_optimization_barrier(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.optimization_barrier_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.optimization_barrier_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.optimization_barrier"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_or"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_or(%arg0: tensor<i1>, %arg1: tensor<i1>) -> tensor<i1> {
--  // CHECK: "vhlo.or_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-+  // CHECK: "vhlo.or_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-   %0 = "stablehlo.or"(%arg0, %arg1) : (tensor<i1>, tensor<i1>) -> tensor<i1>
-   func.return %0 : tensor<i1>
- }
- 
- // CHECK-LABEL: "op_outfeed"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_outfeed(%arg0: tensor<f32>, %arg1: !stablehlo.token) -> !stablehlo.token {
--  //      CHECK: "vhlo.outfeed_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.outfeed_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   outfeed_config = #vhlo.string_v1<"foo">
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.token_v1) -> !vhlo.token_v1
-   %0 = "stablehlo.outfeed"(%arg0, %arg1) {
-@@ -1533,8 +1650,9 @@
- }
- 
- // CHECK-LABEL: "op_pad"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_pad(%arg0: tensor<8xf32>, %arg1: tensor<f32>) -> tensor<16xf32> {
--  //      CHECK: "vhlo.pad_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.pad_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   edge_padding_high = #vhlo.tensor_v1<dense<4> : tensor<1xi64>>,
-   // CHECK-SAME:   edge_padding_low = #vhlo.tensor_v1<dense<4> : tensor<1xi64>>,
-   // CHECK-SAME:   interior_padding = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>
-@@ -1548,36 +1666,41 @@
- }
- 
- // CHECK-LABEL: "op_popcnt"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_popcnt(%arg0: tensor<i32>) -> tensor<i32> {
--  // CHECK: "vhlo.popcnt_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-+  // CHECK: "vhlo.popcnt_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-   %0 = "stablehlo.popcnt"(%arg0) : (tensor<i32>) -> tensor<i32>
-   func.return %0 : tensor<i32>
- }
- 
- // CHECK-LABEL: "op_power"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_power(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.power_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.power_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.power"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_real_dynamic_slice"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}}, %[[ARG3:.*]]: {{.*}})
- func.func @op_real_dynamic_slice(%arg0: tensor<?xf32>, %arg1: tensor<1xindex>, %arg2: tensor<1xindex>, %arg3: tensor<1xindex>) -> tensor<?xf32> {
--  // CHECK: "vhlo.real_dynamic_slice_v1"(%arg0, %arg1, %arg2, %arg3) : (!vhlo.tensor_v1<?x!vhlo.f32_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>) -> !vhlo.tensor_v1<?x!vhlo.f32_v1>
-+  // CHECK: "vhlo.real_dynamic_slice_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]], %[[ARG3]]) : (!vhlo.tensor_v1<?x!vhlo.f32_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>) -> !vhlo.tensor_v1<?x!vhlo.f32_v1>
-   %0 = "stablehlo.real_dynamic_slice"(%arg0, %arg1, %arg2, %arg3) : (tensor<?xf32>, tensor<1xindex>, tensor<1xindex>, tensor<1xindex>) -> tensor<?xf32>
-   func.return %0 : tensor<?xf32>
- }
- 
- // CHECK-LABEL: "op_real"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_real(%arg0: tensor<complex<f32>>) -> tensor<f32> {
--  // CHECK: "vhlo.real_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.real_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.real"(%arg0) : (tensor<complex<f32>>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_recv"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_recv(%arg0: !stablehlo.token) -> (tensor<f32>, !stablehlo.token) {
--  //      CHECK: "vhlo.recv_v1"(%arg0) <{
-+  //      CHECK: "vhlo.recv_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME:   channel_type = #vhlo.integer_v1<3 : i64>,
-   // CHECK-SAME:   is_host_transfer = #vhlo.bool_v1<true>
-@@ -1590,8 +1713,9 @@
- }
- 
- // CHECK-LABEL: "op_reduce"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_reduce(%arg0: tensor<16xf32>, %arg1: tensor<f32>) -> tensor<f32> {
--  //  CHECK: "vhlo.reduce_v1"(%arg0, %arg1)
-+  //  CHECK: "vhlo.reduce_v1"(%[[ARG0]], %[[ARG1]])
-   //  CHECK:   ^[[BB:bb.*]](%[[ARG1:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>, %[[ARG2:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
-   //  CHECK:     "vhlo.return_v1"(%[[VAL1:.*]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-   //  CHECK: }) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-@@ -1606,8 +1730,9 @@
- }
- 
- // CHECK-LABEL: "op_reduce_precision"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_reduce_precision(%arg0: tensor<f32>) -> tensor<f32> {
--  //      CHECK: "vhlo.reduce_precision_v1"(%arg0) <{
-+  //      CHECK: "vhlo.reduce_precision_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   exponent_bits = #vhlo.integer_v1<8 : i32>
-   // CHECK-SAME:   mantissa_bits = #vhlo.integer_v1<10 : i32>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-@@ -1619,8 +1744,9 @@
- }
- 
- // CHECK-LABEL: "op_reduce_scatter"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_reduce_scatter(%arg0: tensor<16xf32>) -> tensor<16xf32> {
--  //               CHECK: "vhlo.reduce_scatter_v1"(%arg0) <{
-+  //               CHECK: "vhlo.reduce_scatter_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0], [1]]> : tensor<2x1xi64>>,
-   //          CHECK-SAME:   scatter_dimension = #vhlo.integer_v1<0 : i64>
-@@ -1644,8 +1770,9 @@
- }
- 
- // CHECK-LABEL: "op_reduce_window"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_reduce_window(%arg0: tensor<2x17x31x7xf32>, %arg1: tensor<f32>) -> tensor<2x9x16x7xf32> {
--  //               CHECK: "vhlo.reduce_window_v1"(%arg0, %arg1) <{
-+  //               CHECK: "vhlo.reduce_window_v1"(%[[ARG0]], %[[ARG1]]) <{
-   //          CHECK-SAME:   base_dilations = #vhlo.tensor_v1<dense<[1, 2, 2, 1]> : tensor<4xi64>>,
-   // CHECK-SAME{LITERAL}:   padding = #vhlo.tensor_v1<dense<[[0, 0], [2, 0], [0, 2], [0, 0]]> : tensor<4x2xi64>>,
-   //          CHECK-SAME:   window_dilations = #vhlo.tensor_v1<dense<[1, 2, 2, 1]> : tensor<4xi64>>,
-@@ -1671,8 +1798,9 @@
- }
- 
- // CHECK-LABEL: "op_remainder"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_remainder(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.remainder_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.remainder_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.remainder"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
-@@ -1692,16 +1820,18 @@
- }
- 
- // CHECK-LABEL: "op_reshape"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_reshape(%arg0: tensor<16xf32>) -> tensor<4x4xf32> {
--  // CHECK: "vhlo.reshape_v1"(%arg0) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<4x4x!vhlo.f32_v1>
-+  // CHECK: "vhlo.reshape_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<4x4x!vhlo.f32_v1>
-   %0 = "stablehlo.reshape"(%arg0) : (tensor<16xf32>) -> tensor<4x4xf32>
-   func.return %0 : tensor<4x4xf32>
- }
- 
- // CHECK-LABEL: "op_return"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_return(%arg0: tensor<i32>, %arg1: tensor<f32>) -> tensor<f32> {
--  //      CHECK: "vhlo.case_v1"(%arg0) ({
--  // CHECK-NEXT:   "vhlo.return_v1"(%arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-+  //      CHECK: "vhlo.case_v1"(%[[ARG0]]) ({
-+  // CHECK-NEXT:   "vhlo.return_v1"(%[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-   // CHECK-NEXT: }) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.case"(%arg0) ({
-     "stablehlo.return"(%arg1) : (tensor<f32>) -> ()
-@@ -1710,8 +1840,9 @@
- }
- 
- // CHECK-LABEL: "op_reverse"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_reverse(%arg0: tensor<16xf32>) -> tensor<16xf32> {
--  //      CHECK: "vhlo.reverse_v1"(%arg0) <{
-+  //      CHECK: "vhlo.reverse_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   dimensions = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x!vhlo.f32_v1>
-   %0 = "stablehlo.reverse"(%arg0) {
-@@ -1721,8 +1852,9 @@
- }
- 
- // CHECK-LABEL: "op_rng_bit_generator"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_rng_bit_generator(%arg0: tensor<f32>) -> (tensor<f32>, tensor<f32>) {
--  //      CHECK: "vhlo.rng_bit_generator_v1"(%arg0) <{
-+  //      CHECK: "vhlo.rng_bit_generator_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   rng_algorithm = #vhlo<rng_algorithm_v1 PHILOX>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>)
-   %0:2 = "stablehlo.rng_bit_generator"(%arg0) {
-@@ -1732,8 +1864,9 @@
- }
- 
- // CHECK-LABEL: "op_rng"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_rng(%arg0: tensor<f32>, %arg1: tensor<f32>, %arg2: tensor<0xindex>) -> tensor<f32> {
--  //      CHECK: "vhlo.rng_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.rng_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   rng_distribution = #vhlo<rng_distribution_v1 NORMAL>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<0x!vhlo.index_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.rng"(%arg0, %arg1, %arg2) {
-@@ -1743,29 +1876,33 @@
- }
- 
- // CHECK-LABEL: "op_round_nearest_afz"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_round_nearest_afz(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.round_nearest_afz_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.round_nearest_afz_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.round_nearest_afz"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_round_nearest_even"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_round_nearest_even(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.round_nearest_even_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.round_nearest_even_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.round_nearest_even"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_rsqrt"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_rsqrt(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.rsqrt_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.rsqrt_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.rsqrt"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_scatter"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_scatter(%arg0: tensor<200x100x300xf32>, %arg1: tensor<10x2xi32>, %arg2: tensor<10x300xf32>) -> tensor<200x100x300xf32> {
--  //      CHECK: "vhlo.scatter_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.scatter_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   index_vector_dim = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   indices_are_sorted = #vhlo.bool_v1<true>,
-   // CHECK-SAME:   inserted_window_dims = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-@@ -1795,8 +1932,9 @@
- }
- 
- // CHECK-LABEL: "op_select_and_scatter"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_select_and_scatter(%arg0: tensor<10x24x24x64xf32>, %arg1: tensor<12x13x13x66xf32>, %arg2: tensor<f32>) -> tensor<10x24x24x64xf32> {
--  //      CHECK: "vhlo.select_and_scatter_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.select_and_scatter_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   padding = #vhlo.tensor_v1<dense<1> : tensor<4x2xi64>>,
-   // CHECK-SAME:   window_dimensions = #vhlo.tensor_v1<dense<[1, 2, 2, 1]> : tensor<4xi64>>,
-   // CHECK-SAME:   window_strides = #vhlo.tensor_v1<dense<[1, 2, 2, 1]> : tensor<4xi64>>
-@@ -1826,15 +1964,17 @@
- }
- 
- // CHECK-LABEL: "op_select"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_select(%arg0: tensor<i1>, %arg1: tensor<f32>, %arg2: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.select_v1"(%arg0, %arg1, %arg2) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.select_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.select"(%arg0, %arg1, %arg2) : (tensor<i1>, tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_send"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_send(%arg0: tensor<f32>, %arg1: !stablehlo.token) -> !stablehlo.token {
--  //      CHECK: "vhlo.send_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.send_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME:   channel_type = #vhlo.integer_v1<2 : i64>,
-   // CHECK-SAME:   is_host_transfer = #vhlo.bool_v1<true>
-@@ -1847,8 +1987,9 @@
- }
- 
- // CHECK-LABEL: "op_set_dimension_size"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_set_dimension_size(%arg0: tensor<?xf32>, %arg1: tensor<i32>) -> tensor<16xf32> {
--  //      CHECK: "vhlo.set_dimension_size_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.set_dimension_size_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   dimension = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<?x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<16x!vhlo.f32_v1>
-   %0 = "stablehlo.set_dimension_size"(%arg0, %arg1) {
-@@ -1858,43 +1999,49 @@
- }
- 
- // CHECK-LABEL: "op_shift_left"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_shift_left(%arg0: tensor<i32>, %arg1: tensor<i32>) -> tensor<i32> {
--  // CHECK: "vhlo.shift_left_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.i32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-+  // CHECK: "vhlo.shift_left_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-   %0 = "stablehlo.shift_left"(%arg0, %arg1) : (tensor<i32>, tensor<i32>) -> tensor<i32>
-   func.return %0 : tensor<i32>
- }
- 
- // CHECK-LABEL: "op_shift_right_arithmetic"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_shift_right_arithmetic(%arg0: tensor<i32>, %arg1: tensor<i32>) -> tensor<i32> {
--  // CHECK: "vhlo.shift_right_arithmetic_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.i32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-+  // CHECK: "vhlo.shift_right_arithmetic_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-   %0 = "stablehlo.shift_right_arithmetic"(%arg0, %arg1) : (tensor<i32>, tensor<i32>) -> tensor<i32>
-   func.return %0 : tensor<i32>
- }
- 
- // CHECK-LABEL: "op_shift_right_logical"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_shift_right_logical(%arg0: tensor<i32>, %arg1: tensor<i32>) -> tensor<i32> {
--  // CHECK: "vhlo.shift_right_logical_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.i32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-+  // CHECK: "vhlo.shift_right_logical_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-   %0 = "stablehlo.shift_right_logical"(%arg0, %arg1) : (tensor<i32>, tensor<i32>) -> tensor<i32>
-   func.return %0 : tensor<i32>
- }
- 
- // CHECK-LABEL: "op_sign"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_sign(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.sign_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.sign_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.sign"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_sine"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_sine(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.sine_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.sine_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.sine"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_slice"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_slice(%arg0: tensor<16xf32>) -> tensor<4xf32> {
--  //      CHECK: "vhlo.slice_v1"(%arg0) <{
-+  //      CHECK: "vhlo.slice_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   limit_indices = #vhlo.tensor_v1<dense<4> : tensor<1xi64>>,
-   // CHECK-SAME:   start_indices = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>,
-   // CHECK-SAME:   strides = #vhlo.tensor_v1<dense<1> : tensor<1xi64>>
-@@ -1908,8 +2055,9 @@
- }
- 
- // CHECK-LABEL: "op_sort"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_sort(%arg0: tensor<16xf32>) -> tensor<16xf32> {
--  //      CHECK: "vhlo.sort_v1"(%arg0) <{
-+  //      CHECK: "vhlo.sort_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   dimension = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME:   is_stable = #vhlo.bool_v1<true>
-   // CHECK-SAME: }> ({
-@@ -1929,29 +2077,33 @@
- }
- 
- // CHECK-LABEL: "op_sqrt"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_sqrt(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.sqrt_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.sqrt_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.sqrt"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_subtract"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_subtract(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.subtract_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.subtract_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.subtract"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_tanh"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_tanh(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.tanh_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.tanh_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.tanh"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_torch_index_select"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_torch_index_select(%arg0: tensor<5x1x5xf32>, %arg1: tensor<2xi32>) ->  tensor<2x1x5xf32> {
--  //      CHECK: "vhlo.torch_index_select_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.torch_index_select_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   batch_dims = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME:   dim = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<5x1x5x!vhlo.f32_v1>, !vhlo.tensor_v1<2x!vhlo.i32_v1>) -> !vhlo.tensor_v1<2x1x5x!vhlo.f32_v1>
-@@ -1963,8 +2115,9 @@
- }
- 
- // CHECK-LABEL: "op_trace"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_trace(%arg0: tensor<f32>) {
--  //      CHECK: "vhlo.trace_v1"(%arg0) <{
-+  //      CHECK: "vhlo.trace_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   tag = #vhlo.string_v1<"foo">
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-   "stablehlo.trace"(%arg0) {
-@@ -1974,8 +2127,9 @@
- }
- 
- // CHECK-LABEL: "op_transpose"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_transpose(%arg0: tensor<16x8xf32>) ->  tensor<8x16xf32> {
--  //      CHECK: "vhlo.transpose_v1"(%arg0) <{
-+  //      CHECK: "vhlo.transpose_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   permutation = #vhlo.tensor_v1<dense<[1, 0]> : tensor<2xi64>>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<16x8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<8x16x!vhlo.f32_v1>
-   %0 = "stablehlo.transpose"(%arg0) {
-@@ -1985,8 +2139,9 @@
- }
- 
- // CHECK-LABEL: "op_triangular_solve"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_triangular_solve(%arg0: tensor<16x16xf32>, %arg1: tensor<16x16xf32>) ->  tensor<16x16xf32> {
--  //      CHECK: "vhlo.triangular_solve_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.triangular_solve_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   left_side = #vhlo.bool_v1<true>,
-   // CHECK-SAME:   lower = #vhlo.bool_v1<true>,
-   // CHECK-SAME:   transpose_a = #vhlo<transpose_v1 NO_TRANSPOSE>,
-@@ -2002,15 +2157,17 @@
- }
- 
- // CHECK-LABEL: "op_tuple"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_tuple(%arg0: tensor<f32>) -> tuple<tensor<f32>> {
--  // CHECK: "vhlo.tuple_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tuple_v1<!vhlo.tensor_v1<!vhlo.f32_v1>>
-+  // CHECK: "vhlo.tuple_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tuple_v1<!vhlo.tensor_v1<!vhlo.f32_v1>>
-   %0 = "stablehlo.tuple"(%arg0) : (tensor<f32>) -> tuple<tensor<f32>>
-   func.return %0 : tuple<tensor<f32>>
- }
- 
- // CHECK-LABEL: "op_unary_einsum"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_unary_einsum(%arg0: tensor<8x16xf32>) -> tensor<8xf32> {
--  //      CHECK: "vhlo.unary_einsum_v1"(%arg0) <{
-+  //      CHECK: "vhlo.unary_einsum_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   einsum_config = #vhlo.string_v1<"ab->a">
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<8x16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<8x!vhlo.f32_v1>
-   %0 = "stablehlo.unary_einsum"(%arg0) {
-@@ -2020,22 +2177,25 @@
- }
- 
- // CHECK-LABEL: "op_uniform_dequantize"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_uniform_dequantize(%arg0: tensor<!quant.uniform<i8:f32, 34.0:16>>) -> tensor<f32> {
--  // CHECK: "vhlo.uniform_dequantize_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.uniform_dequantize_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.uniform_dequantize"(%arg0) : (tensor<!quant.uniform<i8:f32, 34.0:16>>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_uniform_quantize"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_uniform_quantize(%arg0: tensor<f32>) -> tensor<!quant.uniform<i8:f32, 34.0:16>> {
--  // CHECK: "vhlo.uniform_quantize_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>
-+  // CHECK: "vhlo.uniform_quantize_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>
-   %0 = "stablehlo.uniform_quantize"(%arg0) : (tensor<f32>) -> tensor<!quant.uniform<i8:f32, 34.0:16>>
-   func.return %0 : tensor<!quant.uniform<i8:f32, 34.0:16>>
- }
- 
- // CHECK-LABEL: "op_while"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_while(%arg0: tensor<i1>) -> tensor<i1> {
--  //      CHECK: "vhlo.while_v1"(%arg0) ({
-+  //      CHECK: "vhlo.while_v1"(%[[ARG0]]) ({
-   // CHECK-NEXT:   ^[[BB:bb.*]](%[[ARG1:arg.*]]: !vhlo.tensor_v1<!vhlo.bool_v1>):
-   // CHECK-NEXT:     "vhlo.return_v1"(%[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>) -> ()
-   // CHECK-NEXT:   }, {
-@@ -2053,8 +2213,9 @@
- }
- 
- // CHECK-LABEL: "op_xor"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_xor(%arg0: tensor<i1>, %arg1: tensor<i1>) -> tensor<i1> {
--  // CHECK: "vhlo.xor_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-+  // CHECK: "vhlo.xor_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-   %0 = "stablehlo.xor"(%arg0, %arg1) : (tensor<i1>, tensor<i1>) -> tensor<i1>
-   func.return %0 : tensor<i1>
- }
-@@ -2062,190 +2223,217 @@
- // ============ TYPES ============
- 
- // CHECK-LABEL: "type_i1"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_i1(%arg0: tensor<i1>, %arg1: tensor<i1>) -> tensor<i1> {
--  // CHECK: "vhlo.and_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-+  // CHECK: "vhlo.and_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-   %0 = "stablehlo.and"(%arg0, %arg1) : (tensor<i1>, tensor<i1>) -> tensor<i1>
-   func.return %0 : tensor<i1>
- }
- 
- // CHECK-LABEL: "type_i4"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_i4(%arg0: tensor<i4>, %arg1: tensor<i4>) -> tensor<i4> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.i4_v1>, !vhlo.tensor_v1<!vhlo.i4_v1>) -> !vhlo.tensor_v1<!vhlo.i4_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i4_v1>, !vhlo.tensor_v1<!vhlo.i4_v1>) -> !vhlo.tensor_v1<!vhlo.i4_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<i4>, tensor<i4>) -> tensor<i4>
-   func.return %0 : tensor<i4>
- }
- 
- // CHECK-LABEL: "type_i8"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_i8(%arg0: tensor<i8>, %arg1: tensor<i8>) -> tensor<i8> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.i8_v1>, !vhlo.tensor_v1<!vhlo.i8_v1>) -> !vhlo.tensor_v1<!vhlo.i8_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i8_v1>, !vhlo.tensor_v1<!vhlo.i8_v1>) -> !vhlo.tensor_v1<!vhlo.i8_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<i8>, tensor<i8>) -> tensor<i8>
-   func.return %0 : tensor<i8>
- }
- 
- // CHECK-LABEL: "type_i16"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_i16(%arg0: tensor<i16>, %arg1: tensor<i16>) -> tensor<i16> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.i16_v1>, !vhlo.tensor_v1<!vhlo.i16_v1>) -> !vhlo.tensor_v1<!vhlo.i16_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i16_v1>, !vhlo.tensor_v1<!vhlo.i16_v1>) -> !vhlo.tensor_v1<!vhlo.i16_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<i16>, tensor<i16>) -> tensor<i16>
-   func.return %0 : tensor<i16>
- }
- 
- // CHECK-LABEL: "type_i32"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_i32(%arg0: tensor<i32>, %arg1: tensor<i32>) -> tensor<i32> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.i32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<i32>, tensor<i32>) -> tensor<i32>
-   func.return %0 : tensor<i32>
- }
- 
- // CHECK-LABEL: "type_i64"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_i64(%arg0: tensor<i64>, %arg1: tensor<i64>) -> tensor<i64> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.i64_v1>, !vhlo.tensor_v1<!vhlo.i64_v1>) -> !vhlo.tensor_v1<!vhlo.i64_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i64_v1>, !vhlo.tensor_v1<!vhlo.i64_v1>) -> !vhlo.tensor_v1<!vhlo.i64_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<i64>, tensor<i64>) -> tensor<i64>
-   func.return %0 : tensor<i64>
- }
- 
- // CHECK-LABEL: "type_ui4"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_ui4(%arg0: tensor<ui4>, %arg1: tensor<ui4>) -> tensor<ui4> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.ui4_v1>, !vhlo.tensor_v1<!vhlo.ui4_v1>) -> !vhlo.tensor_v1<!vhlo.ui4_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.ui4_v1>, !vhlo.tensor_v1<!vhlo.ui4_v1>) -> !vhlo.tensor_v1<!vhlo.ui4_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<ui4>, tensor<ui4>) -> tensor<ui4>
-   func.return %0 : tensor<ui4>
- }
- 
- // CHECK-LABEL: "type_ui8"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_ui8(%arg0: tensor<ui8>, %arg1: tensor<ui8>) -> tensor<ui8> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.ui8_v1>, !vhlo.tensor_v1<!vhlo.ui8_v1>) -> !vhlo.tensor_v1<!vhlo.ui8_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.ui8_v1>, !vhlo.tensor_v1<!vhlo.ui8_v1>) -> !vhlo.tensor_v1<!vhlo.ui8_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<ui8>, tensor<ui8>) -> tensor<ui8>
-   func.return %0 : tensor<ui8>
- }
- 
- // CHECK-LABEL: "type_ui16"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_ui16(%arg0: tensor<ui16>, %arg1: tensor<ui16>) -> tensor<ui16> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.ui16_v1>, !vhlo.tensor_v1<!vhlo.ui16_v1>) -> !vhlo.tensor_v1<!vhlo.ui16_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.ui16_v1>, !vhlo.tensor_v1<!vhlo.ui16_v1>) -> !vhlo.tensor_v1<!vhlo.ui16_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<ui16>, tensor<ui16>) -> tensor<ui16>
-   func.return %0 : tensor<ui16>
- }
- 
- // CHECK-LABEL: "type_ui32"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_ui32(%arg0: tensor<ui32>, %arg1: tensor<ui32>) -> tensor<ui32> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.ui32_v1>, !vhlo.tensor_v1<!vhlo.ui32_v1>) -> !vhlo.tensor_v1<!vhlo.ui32_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.ui32_v1>, !vhlo.tensor_v1<!vhlo.ui32_v1>) -> !vhlo.tensor_v1<!vhlo.ui32_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<ui32>, tensor<ui32>) -> tensor<ui32>
-   func.return %0 : tensor<ui32>
- }
- 
- // CHECK-LABEL: "type_ui64"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_ui64(%arg0: tensor<ui64>, %arg1: tensor<ui64>) -> tensor<ui64> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.ui64_v1>, !vhlo.tensor_v1<!vhlo.ui64_v1>) -> !vhlo.tensor_v1<!vhlo.ui64_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.ui64_v1>, !vhlo.tensor_v1<!vhlo.ui64_v1>) -> !vhlo.tensor_v1<!vhlo.ui64_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<ui64>, tensor<ui64>) -> tensor<ui64>
-   func.return %0 : tensor<ui64>
- }
- 
- // CHECK-LABEL: "type_f8E4M3FN"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_f8E4M3FN(%arg0: tensor<f8E4M3FN>, %arg1: tensor<f8E4M3FN>) -> tensor<f8E4M3FN> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f8E4M3FN_v1>, !vhlo.tensor_v1<!vhlo.f8E4M3FN_v1>) -> !vhlo.tensor_v1<!vhlo.f8E4M3FN_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f8E4M3FN_v1>, !vhlo.tensor_v1<!vhlo.f8E4M3FN_v1>) -> !vhlo.tensor_v1<!vhlo.f8E4M3FN_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f8E4M3FN>, tensor<f8E4M3FN>) -> tensor<f8E4M3FN>
-   func.return %0 : tensor<f8E4M3FN>
- }
- 
- // CHECK-LABEL: "type_f8E5M2"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_f8E5M2(%arg0: tensor<f8E5M2>, %arg1: tensor<f8E5M2>) -> tensor<f8E5M2> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f8E5M2_v1>, !vhlo.tensor_v1<!vhlo.f8E5M2_v1>) -> !vhlo.tensor_v1<!vhlo.f8E5M2_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f8E5M2_v1>, !vhlo.tensor_v1<!vhlo.f8E5M2_v1>) -> !vhlo.tensor_v1<!vhlo.f8E5M2_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f8E5M2>, tensor<f8E5M2>) -> tensor<f8E5M2>
-   func.return %0 : tensor<f8E5M2>
- }
- 
- // CHECK-LABEL: "type_f8E4M3FNUZ"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_f8E4M3FNUZ(%arg0: tensor<f8E4M3FNUZ>, %arg1: tensor<f8E4M3FNUZ>) -> tensor<f8E4M3FNUZ> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f8E4M3FNUZ_v1>, !vhlo.tensor_v1<!vhlo.f8E4M3FNUZ_v1>) -> !vhlo.tensor_v1<!vhlo.f8E4M3FNUZ_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f8E4M3FNUZ_v1>, !vhlo.tensor_v1<!vhlo.f8E4M3FNUZ_v1>) -> !vhlo.tensor_v1<!vhlo.f8E4M3FNUZ_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f8E4M3FNUZ>, tensor<f8E4M3FNUZ>) -> tensor<f8E4M3FNUZ>
-   func.return %0 : tensor<f8E4M3FNUZ>
- }
- 
- // CHECK-LABEL: "type_f8E4M3B11FNUZ"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_f8E4M3B11FNUZ(%arg0: tensor<f8E4M3B11FNUZ>, %arg1: tensor<f8E4M3B11FNUZ>) -> tensor<f8E4M3B11FNUZ> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f8E4M3B11FNUZ_v1>, !vhlo.tensor_v1<!vhlo.f8E4M3B11FNUZ_v1>) -> !vhlo.tensor_v1<!vhlo.f8E4M3B11FNUZ_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f8E4M3B11FNUZ_v1>, !vhlo.tensor_v1<!vhlo.f8E4M3B11FNUZ_v1>) -> !vhlo.tensor_v1<!vhlo.f8E4M3B11FNUZ_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f8E4M3B11FNUZ>, tensor<f8E4M3B11FNUZ>) -> tensor<f8E4M3B11FNUZ>
-   func.return %0 : tensor<f8E4M3B11FNUZ>
- }
- 
- // CHECK-LABEL: "type_f8E5M2FNUZ"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_f8E5M2FNUZ(%arg0: tensor<f8E5M2FNUZ>, %arg1: tensor<f8E5M2FNUZ>) -> tensor<f8E5M2FNUZ> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f8E5M2FNUZ_v1>, !vhlo.tensor_v1<!vhlo.f8E5M2FNUZ_v1>) -> !vhlo.tensor_v1<!vhlo.f8E5M2FNUZ_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f8E5M2FNUZ_v1>, !vhlo.tensor_v1<!vhlo.f8E5M2FNUZ_v1>) -> !vhlo.tensor_v1<!vhlo.f8E5M2FNUZ_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f8E5M2FNUZ>, tensor<f8E5M2FNUZ>) -> tensor<f8E5M2FNUZ>
-   func.return %0 : tensor<f8E5M2FNUZ>
- }
- 
- // CHECK-LABEL: "type_bf16"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_bf16(%arg0: tensor<bf16>, %arg1: tensor<bf16>) -> tensor<bf16> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.bf16_v1>, !vhlo.tensor_v1<!vhlo.bf16_v1>) -> !vhlo.tensor_v1<!vhlo.bf16_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.bf16_v1>, !vhlo.tensor_v1<!vhlo.bf16_v1>) -> !vhlo.tensor_v1<!vhlo.bf16_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<bf16>, tensor<bf16>) -> tensor<bf16>
-   func.return %0 : tensor<bf16>
- }
- 
- // CHECK-LABEL: "type_f16"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_f16(%arg0: tensor<f16>, %arg1: tensor<f16>) -> tensor<f16> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f16_v1>, !vhlo.tensor_v1<!vhlo.f16_v1>) -> !vhlo.tensor_v1<!vhlo.f16_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f16_v1>, !vhlo.tensor_v1<!vhlo.f16_v1>) -> !vhlo.tensor_v1<!vhlo.f16_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f16>, tensor<f16>) -> tensor<f16>
-   func.return %0 : tensor<f16>
- }
- 
- // CHECK-LABEL: "type_f32"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_f32(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "type_f64"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_f64(%arg0: tensor<f64>, %arg1: tensor<f64>) -> tensor<f64> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f64_v1>, !vhlo.tensor_v1<!vhlo.f64_v1>) -> !vhlo.tensor_v1<!vhlo.f64_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f64_v1>, !vhlo.tensor_v1<!vhlo.f64_v1>) -> !vhlo.tensor_v1<!vhlo.f64_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f64>, tensor<f64>) -> tensor<f64>
-   func.return %0 : tensor<f64>
- }
- 
- // CHECK-LABEL: "type_complex_f32"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_complex_f32(%arg0: tensor<complex<f32>>, %arg1: tensor<complex<f32>>) -> tensor<complex<f32>> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>, !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>) -> !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>, !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>) -> !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<complex<f32>>, tensor<complex<f32>>) -> tensor<complex<f32>>
-   func.return %0 : tensor<complex<f32>>
- }
- 
- // CHECK-LABEL: "type_complex_f64"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_complex_f64(%arg0: tensor<complex<f64>>, %arg1: tensor<complex<f64>>) -> tensor<complex<f64>> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f64_v1>>, !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f64_v1>>) -> !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f64_v1>>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f64_v1>>, !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f64_v1>>) -> !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f64_v1>>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<complex<f64>>, tensor<complex<f64>>) -> tensor<complex<f64>>
-   func.return %0 : tensor<complex<f64>>
- }
- 
- // CHECK-LABEL: "type_dynamism_ranked"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @type_dynamism_ranked(%arg0: tensor<?xf32>) -> tensor<?xf32> {
--  // CHECK: "vhlo.abs_v1"(%arg0) : (!vhlo.tensor_v1<?x!vhlo.f32_v1>) -> !vhlo.tensor_v1<?x!vhlo.f32_v1>
-+  // CHECK: "vhlo.abs_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<?x!vhlo.f32_v1>) -> !vhlo.tensor_v1<?x!vhlo.f32_v1>
-   %0 = "stablehlo.abs"(%arg0) : (tensor<?xf32>) -> tensor<?xf32>
-   func.return %0 : tensor<?xf32>
- }
- 
- // CHECK-LABEL: "type_quantization"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_quantization(%arg0: tensor<!quant.uniform<i8:f32, 34.0:16>>, %arg1: tensor<!quant.uniform<i8:f32, 34.0:16>>) -> tensor<!quant.uniform<i8:f32, 34.0:16>> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>, !vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>) -> !vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>, !vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>) -> !vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<!quant.uniform<i8:f32, 34.0:16>>, tensor<!quant.uniform<i8:f32, 34.0:16>>) -> tensor<!quant.uniform<i8:f32, 34.0:16>>
-   func.return %0 : tensor<!quant.uniform<i8:f32, 34.0:16>>
- }
- 
- //       CHECK: function_type = #vhlo.type_v1<!vhlo.func_v1<(!vhlo.token_v1) -> !vhlo.token_v1>>
- // CHECK-LABEL: "type_token_callee"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @type_token_callee(%arg0: !stablehlo.token) -> !stablehlo.token {
--  // CHECK: "vhlo.return_v1"(%arg0) : (!vhlo.token_v1) -> ()
-+  // CHECK: "vhlo.return_v1"(%[[ARG0]]) : (!vhlo.token_v1) -> ()
-   return %arg0 : !stablehlo.token
- }
- 
- //       CHECK: function_type = #vhlo.type_v1<!vhlo.func_v1<(!vhlo.token_v1) -> !vhlo.token_v1>>
- // CHECK-LABEL: "type_token_caller"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @type_token_caller(%arg0: !stablehlo.token) -> !stablehlo.token {
--  // CHECK:      "vhlo.call_v1"(%arg0) <{callee = #vhlo.string_v1<"type_token_callee">}
-+  // CHECK:      "vhlo.call_v1"(%[[ARG0]]) <{callee = #vhlo.string_v1<"type_token_callee">}
-   // CHECK-SAME: (!vhlo.token_v1) -> !vhlo.token_v1
-   %0 = func.call @type_token_callee(%arg0) : (!stablehlo.token) -> !stablehlo.token
-   return %0 : !stablehlo.token
- }
- 
- // CHECK-LABEL: "type_tuple"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @type_tuple(%arg0: tuple<tensor<f32>>) -> tuple<!stablehlo.token> {
-   %0 = "stablehlo.custom_call"(%arg0) {
-     call_target_name = "foo"
-diff --ruN a/stablehlo/stablehlo/tests/vhlo/stablehlo_legalize_to_vhlo.0_15_0.mlir b/stablehlo/stablehlo/tests/vhlo/stablehlo_legalize_to_vhlo.0_15_0.mlir
---- stablehlo/stablehlo/tests/vhlo/stablehlo_legalize_to_vhlo.0_15_0.mlir
-+++ stablehlo/stablehlo/tests/vhlo/stablehlo_legalize_to_vhlo.0_15_0.mlir
-@@ -13,6 +13,7 @@
- // ============ ATTRIBUTES ============
- 
- // CHECK-LABEL: "attr_comparison_direction_eq"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_direction_eq(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     // CHECK: comparison_direction = #vhlo<comparison_direction_v1 EQ>
-@@ -22,6 +23,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_direction_ne"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_direction_ne(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     // CHECK: comparison_direction = #vhlo<comparison_direction_v1 NE>
-@@ -31,6 +33,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_direction_ge"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_direction_ge(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     // CHECK: comparison_direction = #vhlo<comparison_direction_v1 GE>
-@@ -40,6 +43,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_direction_gt"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_direction_gt(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     // CHECK: comparison_direction = #vhlo<comparison_direction_v1 GT>
-@@ -49,6 +53,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_direction_le"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_direction_le(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     // CHECK: comparison_direction = #vhlo<comparison_direction_v1 LE>
-@@ -58,6 +63,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_direction_lt"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_direction_lt(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     // CHECK: comparison_direction = #vhlo<comparison_direction_v1 LT>
-@@ -67,6 +73,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_type_notype"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_type_notype(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     comparison_direction = #stablehlo<comparison_direction EQ>
-@@ -76,6 +83,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_type_float"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_type_float(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     comparison_direction = #stablehlo<comparison_direction EQ>,
-@@ -86,6 +94,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_type_totalorder"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_type_totalorder(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     comparison_direction = #stablehlo<comparison_direction EQ>,
-@@ -96,6 +105,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_type_signed"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_type_signed(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     comparison_direction = #stablehlo<comparison_direction EQ>,
-@@ -106,6 +116,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_type_unsigned"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_type_unsigned(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     comparison_direction = #stablehlo<comparison_direction EQ>,
-@@ -118,6 +129,7 @@
- // ConvDimensionNumbers aka #stablehlo.conv is covered below.
- 
- // CHECK-LABEL: "attr_custom_call_api_version_unspecified"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_custom_call_api_version_unspecified(%arg0: tensor<f32>) -> tensor<f32> {
-   %0 = "stablehlo.custom_call"(%arg0) {
-     call_target_name = "foo",
-@@ -128,6 +140,7 @@
- }
- 
- // CHECK-LABEL: "attr_custom_call_api_version_original"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_custom_call_api_version_original(%arg0: tensor<f32>) -> tensor<f32> {
-   %0 = "stablehlo.custom_call"(%arg0) {
-     call_target_name = "foo",
-@@ -138,6 +151,7 @@
- }
- 
- // CHECK-LABEL: "attr_custom_call_api_version_status_returning"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_custom_call_api_version_status_returning(%arg0: tensor<f32>) -> tensor<f32> {
-   %0 = "stablehlo.custom_call"(%arg0) {
-     call_target_name = "foo",
-@@ -148,6 +162,7 @@
- }
- 
- // CHECK-LABEL: "attr_custom_call_api_version_status_returning_unified"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_custom_call_api_version_status_returning_unified(%arg0: tensor<f32>) -> tensor<f32> {
-   %0 = "stablehlo.custom_call"(%arg0) {
-     call_target_name = "foo",
-@@ -166,6 +181,7 @@
- // DotDimensionNumbers aka #stablehlo.dot is covered below.
- 
- // CHECK-LABEL: "attr_fft_type_fft"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_fft_type_fft(%arg0: tensor<16xcomplex<f32>>) -> tensor<16xcomplex<f32>> {
-   %0 = "stablehlo.fft"(%arg0) {
-     // CHECK: fft_type = #vhlo<fft_type_v1 FFT>
-@@ -176,6 +192,7 @@
- }
- 
- // CHECK-LABEL: "attr_fft_type_ifft"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_fft_type_ifft(%arg0: tensor<16xcomplex<f32>>) -> tensor<16xcomplex<f32>> {
-   %0 = "stablehlo.fft"(%arg0) {
-     // CHECK: fft_type = #vhlo<fft_type_v1 IFFT>
-@@ -186,6 +203,7 @@
- }
- 
- // CHECK-LABEL: "attr_fft_type_rfft"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_fft_type_rfft(%arg0: tensor<16xf32>) -> tensor<9xcomplex<f32>> {
-   %0 = "stablehlo.fft"(%arg0) {
-     // CHECK: fft_type = #vhlo<fft_type_v1 RFFT>
-@@ -196,6 +214,7 @@
- }
- 
- // CHECK-LABEL: "attr_fft_type_irfft"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_fft_type_irfft(%arg0: tensor<9xcomplex<f32>>) -> tensor<16xf32> {
-   %0 = "stablehlo.fft"(%arg0) {
-     // CHECK: fft_type = #vhlo<fft_type_v1 IRFFT>
-@@ -208,6 +227,7 @@
- // GatherDimensionNumbers aka #stablehlo.gather is covered below.
- 
- // CHECK-LABEL: "attr_precision_config_default"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_precision_config_default(%arg0: tensor<8x16xf32>, %arg1: tensor<16x8xf32>) -> tensor<8x8xf32> {
-   %0 = "stablehlo.dot"(%arg0, %arg1) {
-     // CHECK: precision_config = #vhlo.array_v1<[#vhlo<precision_v1 DEFAULT>, #vhlo<precision_v1 DEFAULT>]>
-@@ -216,6 +236,7 @@
- }
- 
- // CHECK-LABEL: "attr_precision_config_high"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_precision_config_high(%arg0: tensor<8x16xf32>, %arg1: tensor<16x8xf32>) -> tensor<8x8xf32> {
-   %0 = "stablehlo.dot"(%arg0, %arg1) {
-     // CHECK: precision_config = #vhlo.array_v1<[#vhlo<precision_v1 HIGH>, #vhlo<precision_v1 HIGH>]>
-@@ -225,6 +246,7 @@
- }
- 
- // CHECK-LABEL: "attr_precision_config_highest"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_precision_config_highest(%arg0: tensor<8x16xf32>, %arg1: tensor<16x8xf32>) -> tensor<8x8xf32> {
-   %0 = "stablehlo.dot"(%arg0, %arg1) {
-     // CHECK: precision_config = #vhlo.array_v1<[#vhlo<precision_v1 HIGHEST>, #vhlo<precision_v1 HIGHEST>]>
-@@ -234,6 +256,7 @@
- }
- 
- // CHECK-LABEL: "attr_rng_algorithm_default"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_rng_algorithm_default(%arg0: tensor<f32>) -> (tensor<f32>, tensor<f32>) {
-   %0:2 = "stablehlo.rng_bit_generator"(%arg0) {
-     // CHECK: rng_algorithm = #vhlo<rng_algorithm_v1 DEFAULT>
-@@ -243,6 +266,7 @@
- }
- 
- // CHECK-LABEL: "attr_rng_algorithm_three_fry"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_rng_algorithm_three_fry(%arg0: tensor<f32>) -> (tensor<f32>, tensor<f32>) {
-   %0:2 = "stablehlo.rng_bit_generator"(%arg0) {
-     // CHECK: rng_algorithm = #vhlo<rng_algorithm_v1 THREE_FRY>
-@@ -252,6 +276,7 @@
- }
- 
- // CHECK-LABEL: "attr_rng_algorithm_philox"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_rng_algorithm_philox(%arg0: tensor<f32>) -> (tensor<f32>, tensor<f32>) {
-   %0:2 = "stablehlo.rng_bit_generator"(%arg0) {
-     // CHECK: rng_algorithm = #vhlo<rng_algorithm_v1 PHILOX>
-@@ -261,6 +286,7 @@
- }
- 
- // CHECK-LABEL: "attr_rng_distribution_uniform"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @attr_rng_distribution_uniform(%arg0: tensor<f32>, %arg1: tensor<f32>, %arg2: tensor<0xindex>) -> tensor<f32> {
-   %0 = "stablehlo.rng"(%arg0, %arg1, %arg2) {
-     // CHECK: rng_distribution = #vhlo<rng_distribution_v1 UNIFORM>
-@@ -270,6 +296,7 @@
- }
- 
- // CHECK-LABEL: "attr_rng_distribution_normal"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @attr_rng_distribution_normal(%arg0: tensor<f32>, %arg1: tensor<f32>, %arg2: tensor<0xindex>) -> tensor<f32> {
-   %0 = "stablehlo.rng"(%arg0, %arg1, %arg2) {
-     // CHECK: rng_distribution = #vhlo<rng_distribution_v1 NORMAL>
-@@ -281,6 +308,7 @@
- // ScatterDimensionNumbers aka #stablehlo.scatter is covered below.
- 
- // CHECK-LABEL: "attr_transpose_no_transpose"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_transpose_no_transpose(%arg0: tensor<16x16xf32>, %arg1: tensor<16x16xf32>) ->  tensor<16x16xf32> {
-   %0 = "stablehlo.triangular_solve"(%arg0, %arg1) {
-     left_side = true,
-@@ -293,6 +321,7 @@
- }
- 
- // CHECK-LABEL: "attr_transpose_transpose"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_transpose_transpose(%arg0: tensor<16x16xf32>, %arg1: tensor<16x16xf32>) ->  tensor<16x16xf32> {
-   %0 = "stablehlo.triangular_solve"(%arg0, %arg1) {
-     left_side = true,
-@@ -305,6 +334,7 @@
- }
- 
- // CHECK-LABEL: "attr_transpose_adjoint"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_transpose_adjoint(%arg0: tensor<16x16xf32>, %arg1: tensor<16x16xf32>) ->  tensor<16x16xf32> {
-   %0 = "stablehlo.triangular_solve"(%arg0, %arg1) {
-     left_side = true,
-@@ -319,10 +349,9 @@
- // TypeExtensionsAttr aka #stablehlo.type_extensions is covered below.
- 
- // CHECK-LABEL: "attr_type_extensions_bounds"
--func.func @attr_type_extensions_bounds(
--    %arg0: tensor<?x?xf32, #stablehlo.type_extensions<bounds = [16, ?]>>)
--    -> tensor<?x?xf32, #stablehlo.type_extensions<bounds = [16, ?]>> {
--  // CHECK: "vhlo.return_v1"(%arg0) : (!vhlo.tensor_v1<?x?x!vhlo.f32_v1, #vhlo.type_extensions_v1<bounds = [16, ?]>>) -> ()
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @attr_type_extensions_bounds(%arg0: tensor<?x?xf32, #stablehlo.type_extensions<bounds = [16, ?]>>) -> tensor<?x?xf32, #stablehlo.type_extensions<bounds = [16, ?]>> {
-+  // CHECK: "vhlo.return_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<?x?x!vhlo.f32_v1, #vhlo.type_extensions_v1<bounds = [16, ?]>>) -> ()
-   func.return %arg0 : tensor<?x?xf32, #stablehlo.type_extensions<bounds = [16, ?]>>
- }
- 
-@@ -330,8 +359,9 @@
- // ============ DEFAULTS ============
- 
- // CHECK-LABEL: "default_all_gather"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_all_gather(%arg0: tensor<16x8xf32>) -> tensor<16x16xf32> {
--  //               CHECK: "vhlo.all_gather_v1"(%arg0) <{
-+  //               CHECK: "vhlo.all_gather_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   all_gather_dim = #vhlo.integer_v1<1 : i64>
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0], [1]]> : tensor<2x1xi64>>,
-@@ -345,8 +375,9 @@
- }
- 
- // CHECK-LABEL: "default_all_reduce"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_all_reduce(%arg0: tensor<f32>) -> tensor<f32> {
--  //               CHECK: "vhlo.all_reduce_v1"(%arg0)
-+  //               CHECK: "vhlo.all_reduce_v1"(%[[ARG0]])
-   //          CHECK-SAME: <{
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0], [1]]> : tensor<2x1xi64>>,
-@@ -368,8 +399,9 @@
- }
- 
- // CHECK-LABEL: "default_all_to_all"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_all_to_all(%arg0: tensor<4x16xf32>) -> tensor<16x4xf32> {
--  //               CHECK: "vhlo.all_to_all_v1"(%arg0) <{
-+  //               CHECK: "vhlo.all_to_all_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   //          CHECK-SAME:   concat_dimension = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0, 1, 2, 3]]> : tensor<1x4xi64>>,
-@@ -386,8 +418,9 @@
- }
- 
- // CHECK-LABEL: "default_cholesky"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_cholesky(%arg0: tensor<1x16x16xf32>) -> tensor<1x16x16xf32> {
--  //      CHECK: "vhlo.cholesky_v1"(%arg0) <{
-+  //      CHECK: "vhlo.cholesky_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   lower = #vhlo.bool_v1<false>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<1x16x16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<1x16x16x!vhlo.f32_v1>
-   %0 = "stablehlo.cholesky"(%arg0) : (tensor<1x16x16xf32>) -> tensor<1x16x16xf32>
-@@ -395,8 +428,9 @@
- }
- 
- // CHECK-LABEL: "default_collective_permute"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_collective_permute(%arg0: tensor<16x8xf32>) -> tensor<16x8xf32> {
--  //               CHECK: "vhlo.collective_permute_v1"(%arg0) <{
-+  //               CHECK: "vhlo.collective_permute_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME{LITERAL}:   source_target_pairs = #vhlo.tensor_v1<dense<[[0, 1], [1, 2], [2, 3]]> : tensor<3x2xi64>>
-   //          CHECK-SAME: }> : (!vhlo.tensor_v1<16x8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x8x!vhlo.f32_v1>
-@@ -407,8 +441,9 @@
- }
- 
- // CHECK-LABEL: "default_compare"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @default_compare(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
--  //      CHECK: "vhlo.compare_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.compare_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   compare_type = #vhlo<comparison_type_v1 NOTYPE>,
-   // CHECK-SAME:   comparison_direction = #vhlo<comparison_direction_v1 EQ>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-@@ -419,8 +454,9 @@
- }
- 
- // CHECK-LABEL: "default_convolution"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @default_convolution(%arg0: tensor<1x8x8x207xf32>, %arg1: tensor<3x3x207x16xf32>) -> tensor<1x6x6x16xf32> {
--  //      CHECK: "vhlo.convolution_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.convolution_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   batch_group_count = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   feature_group_count = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   input_batch_dimension = #vhlo.integer_v1<0 : i64>,
-@@ -448,8 +484,9 @@
- }
- 
- // CHECK-LABEL: "default_custom_call"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_custom_call(%arg0: tensor<f32>) -> tensor<f32> {
--  //      CHECK: "vhlo.custom_call_v1"(%arg0) <{
-+  //      CHECK: "vhlo.custom_call_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   api_version = #vhlo<api_version_v1 API_VERSION_ORIGINAL>,
-   // CHECK-SAME:   backend_config = #vhlo.string_v1<"">,
-   // CHECK-SAME:   call_target_name = #vhlo.string_v1<"foo">,
-@@ -466,8 +503,9 @@
- }
- 
- // CHECK-LABEL: "default_dot_general"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @default_dot_general(%arg0: tensor<8x8x16xf32>, %arg1: tensor<8x16x8xf32>) -> tensor<8x8x8xf32> {
--  //      CHECK: "vhlo.dot_general_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.dot_general_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   lhs_batching_dimensions = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>,
-   // CHECK-SAME:   lhs_contracting_dimensions = #vhlo.tensor_v1<dense<2> : tensor<1xi64>>,
-   // CHECK-SAME:   precision_config = #vhlo.array_v1<[#vhlo<precision_v1 DEFAULT>, #vhlo<precision_v1 DEFAULT>]>,
-@@ -486,8 +524,9 @@
- }
- 
- // CHECK-LABEL: "default_dot"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @default_dot(%arg0: tensor<8x16xf32>, %arg1: tensor<16x8xf32>) -> tensor<8x8xf32> {
--  //      CHECK: "vhlo.dot_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.dot_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   precision_config = #vhlo.array_v1<[#vhlo<precision_v1 DEFAULT>, #vhlo<precision_v1 DEFAULT>]>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<8x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<8x8x!vhlo.f32_v1>
-   %0 = "stablehlo.dot"(%arg0, %arg1) : (tensor<8x16xf32>, tensor<16x8xf32>) -> tensor<8x8xf32>
-@@ -495,8 +534,9 @@
- }
- 
- // CHECK-LABEL: "default_dynamic_broadcast_in_dim"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @default_dynamic_broadcast_in_dim(%arg0: tensor<?x?xf32>, %arg1: tensor<2xindex>) -> tensor<?x?xf32> {
--  //      CHECK: "vhlo.dynamic_broadcast_in_dim_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.dynamic_broadcast_in_dim_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   broadcast_dimensions = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-   // CHECK-SAME:   known_expanding_dimensions = #vhlo.tensor_v1<dense<> : tensor<0xi64>>,
-   // CHECK-SAME:   known_nonexpanding_dimensions = #vhlo.tensor_v1<dense<> : tensor<0xi64>>
-@@ -508,8 +548,9 @@
- }
- 
- // CHECK-LABEL: "default_dynamic_conv"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @default_dynamic_conv(%arg0: tensor<1x8x8x207xf32>, %arg1: tensor<3x3x207x16xf32>, %arg2: tensor<4xi32>) -> tensor<1x?x?x16xf32> {
--  //      CHECK: "vhlo.dynamic_conv_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.dynamic_conv_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   batch_group_count = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   feature_group_count = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   input_batch_dimension = #vhlo.integer_v1<0 : i64>,
-@@ -537,8 +578,9 @@
- }
- 
- // CHECK-LABEL: "default_dynamic_gather"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @default_dynamic_gather(%arg0 : tensor<2x4x9xf32>, %arg1 : tensor<1x5x2xi32>, %arg2 : tensor<3xi32>) -> tensor<1x5x8xf32> {
--  //      CHECK: "vhlo.dynamic_gather_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.dynamic_gather_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   collapsed_slice_dims = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-   // CHECK-SAME:   index_vector_dim = #vhlo.integer_v1<2 : i64>,
-   // CHECK-SAME:   indices_are_sorted = #vhlo.bool_v1<false>,
-@@ -564,15 +606,16 @@
-   // CHECK-SAME:   sym_name = #vhlo.string_v1<"default_func">,
-   // CHECK-SAME:   sym_visibility = #vhlo.string_v1<"">
-   // CHECK-SAME: }> ({
--  // CHECK-NEXT:   ^[[BB:bb.*]](%arg0: !vhlo.tensor_v1<!vhlo.f32_v1>):
--  // CHECK-NEXT:     "vhlo.return_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-+  // CHECK-NEXT:   ^[[BB:bb.*]](%[[ARG0:.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
-+  // CHECK-NEXT:     "vhlo.return_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-   // CHECK-NEXT: }) : () -> ()
-   func.return %arg0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "dynamic_gather"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @dynamic_gather(%arg0 : tensor<2x4x9xf32>, %arg1 : tensor<1x5x2xi32>) -> tensor<1x5x1xf32> {
--  //      CHECK: "vhlo.gather_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.gather_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   collapsed_slice_dims = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-   // CHECK-SAME:   index_vector_dim = #vhlo.integer_v1<2 : i64>,
-   // CHECK-SAME:   indices_are_sorted = #vhlo.bool_v1<false>,
-@@ -593,8 +636,9 @@
- }
- 
- // CHECK-LABEL: "default_infeed"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_infeed(%arg0: !stablehlo.token) -> (tensor<f32>, !stablehlo.token) {
--  //               CHECK: "vhlo.infeed_v1"(%arg0) <{
-+  //               CHECK: "vhlo.infeed_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   infeed_config = #vhlo.string_v1<"">,
-   // CHECK-SAME{LITERAL}:   layout = #vhlo.array_v1<[]>
-   //          CHECK-SAME: }> : (!vhlo.token_v1) -> (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.token_v1)
-@@ -603,8 +647,9 @@
- }
- 
- // CHECK-LABEL: "default_outfeed"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @default_outfeed(%arg0: tensor<f32>, %arg1: !stablehlo.token) -> !stablehlo.token {
--  //      CHECK: "vhlo.outfeed_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.outfeed_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   outfeed_config = #vhlo.string_v1<"">
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.token_v1) -> !vhlo.token_v1
-   %0 = "stablehlo.outfeed"(%arg0, %arg1) : (tensor<f32>, !stablehlo.token) -> !stablehlo.token
-@@ -612,8 +657,9 @@
- }
- 
- // CHECK-LABEL: "default_recv"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_recv(%arg0: !stablehlo.token) -> (tensor<f32>, !stablehlo.token) {
--  //      CHECK: "vhlo.recv_v1"(%arg0) <{
-+  //      CHECK: "vhlo.recv_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME:   channel_type = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   is_host_transfer = #vhlo.bool_v1<false>
-@@ -625,8 +671,9 @@
- }
- 
- // CHECK-LABEL: "default_send"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @default_send(%arg0: tensor<f32>, %arg1: !stablehlo.token) -> !stablehlo.token {
--  //      CHECK: "vhlo.send_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.send_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME:   channel_type = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   is_host_transfer = #vhlo.bool_v1<false>
-@@ -638,8 +685,9 @@
- }
- 
- // CHECK-LABEL: "default_reduce_scatter"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_reduce_scatter(%arg0: tensor<16xf32>) -> tensor<16xf32> {
--  //               CHECK: "vhlo.reduce_scatter_v1"(%arg0) <{
-+  //               CHECK: "vhlo.reduce_scatter_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0], [1]]> : tensor<2x1xi64>>,
-   //          CHECK-SAME:   scatter_dimension = #vhlo.integer_v1<0 : i64>
-@@ -661,8 +709,9 @@
- }
- 
- // CHECK-LABEL: "default_reduce_window"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @default_reduce_window(%arg0: tensor<2x17x31x7xf32>, %arg1: tensor<f32>) -> tensor<2x16x30x7xf32> {
--  //               CHECK: "vhlo.reduce_window_v1"(%arg0, %arg1)  <{
-+  //               CHECK: "vhlo.reduce_window_v1"(%[[ARG0]], %[[ARG1]])  <{
-   //          CHECK-SAME:   base_dilations = #vhlo.tensor_v1<dense<1> : tensor<4xi64>>,
-   // CHECK-SAME{LITERAL}:   padding = #vhlo.tensor_v1<dense<0> : tensor<4x2xi64>>,
-   //          CHECK-SAME:   window_dilations = #vhlo.tensor_v1<dense<1> : tensor<4xi64>>,
-@@ -684,8 +733,9 @@
- }
- 
- // CHECK-LABEL: "default_scatter"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @default_scatter(%arg0: tensor<200x100x300xf32>, %arg1: tensor<10x2xi32>, %arg2: tensor<10x300xf32>) -> tensor<200x100x300xf32> {
--  //      CHECK: "vhlo.scatter_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.scatter_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   index_vector_dim = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   indices_are_sorted = #vhlo.bool_v1<false>,
-   // CHECK-SAME:   inserted_window_dims = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-@@ -713,8 +763,9 @@
- }
- 
- // CHECK-LABEL: "default_select_and_scatter"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @default_select_and_scatter(%arg0: tensor<10x24x24x64xf32>, %arg1: tensor<10x23x23x64xf32>, %arg2: tensor<f32>) -> tensor<10x24x24x64xf32> {
--  //      CHECK: "vhlo.select_and_scatter_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.select_and_scatter_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   padding = #vhlo.tensor_v1<dense<0> : tensor<4x2xi64>>,
-   // CHECK-SAME:   window_dimensions = #vhlo.tensor_v1<dense<[1, 2, 2, 1]> : tensor<4xi64>>,
-   // CHECK-SAME:   window_strides = #vhlo.tensor_v1<dense<1> : tensor<4xi64>>
-@@ -742,8 +793,9 @@
- }
- 
- // CHECK-LABEL: "default_sort"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_sort(%arg0: tensor<16xf32>) -> tensor<16xf32> {
--  //      CHECK: "vhlo.sort_v1"(%arg0) <{
-+  //      CHECK: "vhlo.sort_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   dimension = #vhlo.integer_v1<-1 : i64>
-   // CHECK-SAME:   is_stable = #vhlo.bool_v1<false>
-   // CHECK-SAME: }> ({
-@@ -762,29 +814,33 @@
- // ============ OPS ============
- 
- // CHECK-LABEL: "op_abs"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_abs(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.abs_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.abs_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.abs"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_add"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_add(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_after_all"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_after_all(%arg0: !stablehlo.token) -> !stablehlo.token {
--  // CHECK: "vhlo.after_all_v1"(%arg0) : (!vhlo.token_v1) -> !vhlo.token_v1
-+  // CHECK: "vhlo.after_all_v1"(%[[ARG0]]) : (!vhlo.token_v1) -> !vhlo.token_v1
-   %0 = "stablehlo.after_all"(%arg0) : (!stablehlo.token) -> !stablehlo.token
-   func.return %0 : !stablehlo.token
- }
- 
- // CHECK-LABEL: "op_all_gather"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_all_gather(%arg0: tensor<16x8xf32>) -> tensor<16x16xf32> {
--  //               CHECK: "vhlo.all_gather_v1"(%arg0) <{
-+  //               CHECK: "vhlo.all_gather_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   all_gather_dim = #vhlo.integer_v1<1 : i64>
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0], [1]]> : tensor<2x1xi64>>,
-@@ -800,8 +856,9 @@
- }
- 
- // CHECK-LABEL: "op_all_reduce"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_all_reduce(%arg0: tensor<f32>) -> tensor<f32> {
--  //               CHECK: "vhlo.all_reduce_v1"(%arg0) <{
-+  //               CHECK: "vhlo.all_reduce_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0], [1]]> : tensor<2x1xi64>>,
-   //          CHECK-SAME:   use_global_device_ids = #vhlo.bool_v1<true>
-@@ -823,8 +880,9 @@
- }
- 
- // CHECK-LABEL: "op_all_to_all"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_all_to_all(%arg0: tensor<4x16xf32>) -> tensor<16x4xf32> {
--  //               CHECK: "vhlo.all_to_all_v1"(%arg0) <{
-+  //               CHECK: "vhlo.all_to_all_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<1 : i64>,
-   //          CHECK-SAME:   concat_dimension = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0, 1, 2, 3]]> : tensor<1x4xi64>>,
-@@ -842,22 +900,25 @@
- }
- 
- // CHECK-LABEL: "op_and"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_and(%arg0: tensor<i1>, %arg1: tensor<i1>) -> tensor<i1> {
--  // CHECK: "vhlo.and_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-+  // CHECK: "vhlo.and_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-   %0 = "stablehlo.and"(%arg0, %arg1) : (tensor<i1>, tensor<i1>) -> tensor<i1>
-   func.return %0 : tensor<i1>
- }
- 
- // CHECK-LABEL: "op_atan2"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_atan2(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.atan2_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.atan2_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.atan2"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_batch_norm_grad"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}}, %[[ARG3:.*]]: {{.*}}, %[[ARG4:.*]]: {{.*}})
- func.func @op_batch_norm_grad(%arg0: tensor<16x16x16x16xf32>, %arg1: tensor<16xf32>, %arg2: tensor<16xf32>, %arg3: tensor<16xf32>, %arg4: tensor<16x16x16x16xf32>) -> (tensor<16x16x16x16xf32>, tensor<16xf32>, tensor<16xf32>) {
--  //      CHECK: "vhlo.batch_norm_grad_v1"(%arg0, %arg1, %arg2, %arg3, %arg4) <{
-+  //      CHECK: "vhlo.batch_norm_grad_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]], %[[ARG3]], %[[ARG4]]) <{
-   // CHECK-SAME:   epsilon = #vhlo.float_v1<1.000000e-03 : !vhlo.f32_v1>,
-   // CHECK-SAME:   feature_index = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<16x16x16x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x16x16x16x!vhlo.f32_v1>) -> (!vhlo.tensor_v1<16x16x16x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>)
-@@ -869,8 +930,9 @@
- }
- 
- // CHECK-LABEL: "op_batch_norm_inference"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}}, %[[ARG3:.*]]: {{.*}}, %[[ARG4:.*]]: {{.*}})
- func.func @op_batch_norm_inference(%arg0: tensor<16x16x16x16xf32>, %arg1: tensor<16xf32>, %arg2: tensor<16xf32>, %arg3: tensor<16xf32>, %arg4: tensor<16xf32>) -> tensor<16x16x16x16xf32> {
--  //      CHECK: "vhlo.batch_norm_inference_v1"(%arg0, %arg1, %arg2, %arg3, %arg4) <{
-+  //      CHECK: "vhlo.batch_norm_inference_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]], %[[ARG3]], %[[ARG4]]) <{
-   // CHECK-SAME:   epsilon = #vhlo.float_v1<1.000000e-03 : !vhlo.f32_v1>,
-   // CHECK-SAME:   feature_index = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<16x16x16x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x16x16x16x!vhlo.f32_v1>
-@@ -882,8 +944,9 @@
- }
- 
- // CHECK-LABEL: "op_batch_norm_training"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_batch_norm_training(%arg0: tensor<16x16x16x16xf32>, %arg1: tensor<16xf32>, %arg2: tensor<16xf32>) -> (tensor<16x16x16x16xf32>, tensor<16xf32>, tensor<16xf32>) {
--  //      CHECK: "vhlo.batch_norm_training_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.batch_norm_training_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   epsilon = #vhlo.float_v1<1.000000e-03 : !vhlo.f32_v1>,
-   // CHECK-SAME:   feature_index = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<16x16x16x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>) -> (!vhlo.tensor_v1<16x16x16x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>)
-@@ -895,15 +958,17 @@
- }
- 
- // CHECK-LABEL: "op_bitcast_convert"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_bitcast_convert(%arg0: tensor<i32>) -> tensor<f32> {
--  // CHECK: "vhlo.bitcast_convert_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.bitcast_convert_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.bitcast_convert"(%arg0) : (tensor<i32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_broadcast_in_dim"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_broadcast_in_dim(%arg0: tensor<16xf32>) -> tensor<16x16xf32> {
--  //      CHECK: "vhlo.broadcast_in_dim_v1"(%arg0) <{
-+  //      CHECK: "vhlo.broadcast_in_dim_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   broadcast_dimensions = #vhlo.tensor_v1<dense<1> : tensor<1xi64>>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x16x!vhlo.f32_v1>
-   %0 = "stablehlo.broadcast_in_dim"(%arg0) {
-@@ -913,8 +978,9 @@
- }
- 
- // CHECK-LABEL: "op_broadcast"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_broadcast(%arg0: tensor<16xf32>) -> tensor<16x16xf32> {
--  //      CHECK: "vhlo.broadcast_v1"(%arg0) <{
-+  //      CHECK: "vhlo.broadcast_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   broadcast_sizes = #vhlo.tensor_v1<dense<16> : tensor<1xi64>>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x16x!vhlo.f32_v1>
-   %0 = "stablehlo.broadcast"(%arg0) {
-@@ -924,9 +990,10 @@
- }
- 
- // CHECK-LABEL: "op_case"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_case(%arg0: tensor<i32>, %arg1: tensor<f32>) -> tensor<f32> {
--  //      CHECK: "vhlo.case_v1"(%arg0) ({
--  // CHECK-NEXT:   "vhlo.return_v1"(%arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-+  //      CHECK: "vhlo.case_v1"(%[[ARG0]]) ({
-+  // CHECK-NEXT:   "vhlo.return_v1"(%[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-   // CHECK-NEXT: }) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.case"(%arg0) ({
-     "stablehlo.return"(%arg1) : (tensor<f32>) -> ()
-@@ -935,22 +1002,25 @@
- }
- 
- // CHECK-LABEL: "op_cbrt"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_cbrt(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.cbrt_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.cbrt_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.cbrt"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_ceil"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_ceil(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.ceil_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.ceil_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.ceil"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_cholesky"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_cholesky(%arg0: tensor<1x16x16xf32>) -> tensor<1x16x16xf32> {
--  //      CHECK: "vhlo.cholesky_v1"(%arg0) <{
-+  //      CHECK: "vhlo.cholesky_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   lower = #vhlo.bool_v1<true>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<1x16x16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<1x16x16x!vhlo.f32_v1>
-   %0 = "stablehlo.cholesky"(%arg0) {
-@@ -960,22 +1030,25 @@
- }
- 
- // CHECK-LABEL: "op_clamp"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_clamp(%arg0: tensor<f32>, %arg1: tensor<f32>, %arg2: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.clamp_v1"(%arg0, %arg1, %arg2) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.clamp_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.clamp"(%arg0, %arg1, %arg2) : (tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_count_leading_zeros"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_count_leading_zeros(%arg0: tensor<i32>) -> tensor<i32> {
--  // CHECK: "vhlo.count_leading_zeros_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-+  // CHECK: "vhlo.count_leading_zeros_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-   %0 = "stablehlo.count_leading_zeros"(%arg0) : (tensor<i32>) -> tensor<i32>
-   func.return %0 : tensor<i32>
- }
- 
- // CHECK-LABEL: "op_collective_permute"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_collective_permute(%arg0: tensor<16x8xf32>) -> tensor<16x8xf32> {
--  //               CHECK: "vhlo.collective_permute_v1"(%arg0) <{
-+  //               CHECK: "vhlo.collective_permute_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME{LITERAL}:   source_target_pairs = #vhlo.tensor_v1<dense<[[0, 1], [1, 2], [2, 3]]> : tensor<3x2xi64>>
-   //          CHECK-SAME: }> : (!vhlo.tensor_v1<16x8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x8x!vhlo.f32_v1>
-@@ -987,8 +1060,9 @@
- }
- 
- // CHECK-LABEL: "op_compare"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_compare(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
--  //      CHECK: "vhlo.compare_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.compare_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   compare_type = #vhlo<comparison_type_v1 TOTALORDER>,
-   // CHECK-SAME:   comparison_direction = #vhlo<comparison_direction_v1 EQ>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-@@ -1000,15 +1074,17 @@
- }
- 
- // CHECK-LABEL: "op_complex"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_complex(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<complex<f32>> {
--  // CHECK: "vhlo.complex_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>
-+  // CHECK: "vhlo.complex_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>
-   %0 = "stablehlo.complex"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<complex<f32>>
-   func.return %0 : tensor<complex<f32>>
- }
- 
- // CHECK-LABEL: "op_concatenate"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_concatenate(%arg0: tensor<8xf32>, %arg1: tensor<8xf32>) -> tensor<16xf32> {
--  //      CHECK: "vhlo.concatenate_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.concatenate_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   dimension = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<8x!vhlo.f32_v1>, !vhlo.tensor_v1<8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x!vhlo.f32_v1>
-   %0 = "stablehlo.concatenate"(%arg0, %arg1) {
-@@ -1018,6 +1094,7 @@
- }
- 
- // CHECK-LABEL: "op_constant"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_constant(%arg0: tensor<f32>) -> tensor<f32> {
-   //      CHECK: "vhlo.constant_v1"() <{
-   // CHECK-SAME:   value = #vhlo.tensor_v1<dense<0.000000e+00> : tensor<f32>>
-@@ -1029,15 +1106,17 @@
- }
- 
- // CHECK-LABEL: "op_convert"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_convert(%arg0: tensor<i32>) -> tensor<f32> {
--  // CHECK: "vhlo.convert_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.convert_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.convert"(%arg0) : (tensor<i32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_convolution"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_convolution(%arg0: tensor<1x8x8x207xf32>, %arg1: tensor<3x3x207x16xf32>) -> tensor<1x7x7x16xf32> {
--  //      CHECK: "vhlo.convolution_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.convolution_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   batch_group_count = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   feature_group_count = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   input_batch_dimension = #vhlo.integer_v1<0 : i64>,
-@@ -1071,8 +1150,9 @@
- }
- 
- // CHECK-LABEL: "op_cosine"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_cosine(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.cosine_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.cosine_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.cosine"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
-@@ -1085,8 +1165,9 @@
- }
- 
- // CHECK-LABEL: "op_cross_replica_sum"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_cross_replica_sum(%arg0: tensor<f32>) -> tensor<f32> {
--  //               CHECK: "vhlo.cross-replica-sum_v1"(%arg0) <{
-+  //               CHECK: "vhlo.cross-replica-sum_v1"(%[[ARG0]]) <{
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0], [1]]> : tensor<2x1xi64>>
-   //          CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.cross-replica-sum"(%arg0) {
-@@ -1096,8 +1177,9 @@
- }
- 
- // CHECK-LABEL: "op_custom_call"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_custom_call(%arg0: tensor<f32>) -> tensor<f32> {
--  //      CHECK: "vhlo.custom_call_v1"(%arg0) <{
-+  //      CHECK: "vhlo.custom_call_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   api_version = #vhlo<api_version_v1 API_VERSION_STATUS_RETURNING>,
-   // CHECK-SAME:   backend_config = #vhlo.string_v1<"\08\03\1A\02">,
-   // CHECK-SAME:   call_target_name = #vhlo.string_v1<"foo">,
-@@ -1128,15 +1210,17 @@
- }
- 
- // CHECK-LABEL: "op_divide"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_divide(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.divide_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.divide_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.divide"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_dot_general"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_dot_general(%arg0: tensor<8x8x16xf32>, %arg1: tensor<8x16x8xf32>) -> tensor<8x8x8xf32> {
--  //      CHECK: "vhlo.dot_general_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.dot_general_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   lhs_batching_dimensions = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>,
-   // CHECK-SAME:   lhs_contracting_dimensions = #vhlo.tensor_v1<dense<2> : tensor<1xi64>>,
-   // CHECK-SAME:   precision_config = #vhlo.array_v1<[#vhlo<precision_v1 HIGHEST>, #vhlo<precision_v1 HIGHEST>]>,
-@@ -1156,8 +1240,9 @@
- }
- 
- // CHECK-LABEL: "op_dot"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_dot(%arg0: tensor<8x16xf32>, %arg1: tensor<16x8xf32>) -> tensor<8x8xf32> {
--  //      CHECK: "vhlo.dot_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.dot_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   precision_config = #vhlo.array_v1<[#vhlo<precision_v1 HIGHEST>, #vhlo<precision_v1 HIGHEST>]>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<8x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<8x8x!vhlo.f32_v1>
-   %0 = "stablehlo.dot"(%arg0, %arg1) {
-@@ -1167,8 +1252,9 @@
- }
- 
- // CHECK-LABEL: "op_dynamic_broadcast_in_dim"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_dynamic_broadcast_in_dim(%arg0: tensor<?x?xf32>, %arg1: tensor<2xindex>) -> tensor<?x?xf32> {
--  //      CHECK: "vhlo.dynamic_broadcast_in_dim_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.dynamic_broadcast_in_dim_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   broadcast_dimensions = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-   // CHECK-SAME:   known_expanding_dimensions = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>,
-   // CHECK-SAME:   known_nonexpanding_dimensions = #vhlo.tensor_v1<dense<1> : tensor<1xi64>>
-@@ -1182,8 +1268,9 @@
- }
- 
- // CHECK-LABEL: "op_dynamic_conv"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_dynamic_conv(%arg0: tensor<1x8x8x207xf32>, %arg1: tensor<3x3x207x16xf32>, %arg2: tensor<4xi32>) -> tensor<1x?x?x16xf32> {
--  //      CHECK: "vhlo.dynamic_conv_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.dynamic_conv_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   batch_group_count = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   feature_group_count = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   input_batch_dimension = #vhlo.integer_v1<0 : i64>,
-@@ -1217,8 +1304,9 @@
- }
- 
- // CHECK-LABEL: "op_dynamic_gather"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_dynamic_gather(%arg0 : tensor<2x4x9xf32>, %arg1 : tensor<1x5x2xi32>, %arg2 : tensor<3xi32>) -> tensor<1x5x8xf32> {
--  //      CHECK: "vhlo.dynamic_gather_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.dynamic_gather_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   collapsed_slice_dims = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-   // CHECK-SAME:   index_vector_dim = #vhlo.integer_v1<2 : i64>,
-   // CHECK-SAME:   indices_are_sorted = #vhlo.bool_v1<true>,
-@@ -1238,8 +1326,9 @@
- }
- 
- // CHECK-LABEL: "op_dynamic_iota"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_dynamic_iota(%arg0: tensor<1xindex>) -> tensor<?xf32> {
--  //      CHECK: "vhlo.dynamic_iota_v1"(%arg0) <{
-+  //      CHECK: "vhlo.dynamic_iota_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   iota_dimension = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<1x!vhlo.index_v1>) -> !vhlo.tensor_v1<?x!vhlo.f32_v1>
-   %0 = "stablehlo.dynamic_iota"(%arg0) {
-@@ -1249,22 +1338,25 @@
- }
- 
- // CHECK-LABEL: "op_dynamic_pad"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}}, %[[ARG3:.*]]: {{.*}}, %[[ARG4:.*]]: {{.*}})
- func.func @op_dynamic_pad(%arg0: tensor<?xf32>, %arg1: tensor<f32>, %arg2: tensor<1xindex>, %arg3: tensor<1xindex>, %arg4: tensor<1xindex>) -> tensor<?xf32> {
--  // CHECK: "vhlo.dynamic_pad_v1"(%arg0, %arg1, %arg2, %arg3, %arg4) : (!vhlo.tensor_v1<?x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>) -> !vhlo.tensor_v1<?x!vhlo.f32_v1>
-+  // CHECK: "vhlo.dynamic_pad_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]], %[[ARG3]], %[[ARG4]]) : (!vhlo.tensor_v1<?x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>) -> !vhlo.tensor_v1<?x!vhlo.f32_v1>
-   %0 = "stablehlo.dynamic_pad"(%arg0, %arg1, %arg2, %arg3, %arg4) : (tensor<?xf32>, tensor<f32>, tensor<1xindex>, tensor<1xindex>, tensor<1xindex>) -> tensor<?xf32>
-   func.return %0 : tensor<?xf32>
- }
- 
- // CHECK-LABEL: "op_dynamic_reshape"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_dynamic_reshape(%arg0: tensor<16xf32>, %arg1: tensor<2xindex>) -> tensor<?x?xf32> {
--  // CHECK: "vhlo.dynamic_reshape_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<2x!vhlo.index_v1>) -> !vhlo.tensor_v1<?x?x!vhlo.f32_v1>
-+  // CHECK: "vhlo.dynamic_reshape_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<2x!vhlo.index_v1>) -> !vhlo.tensor_v1<?x?x!vhlo.f32_v1>
-   %0 = "stablehlo.dynamic_reshape"(%arg0, %arg1) : (tensor<16xf32>, tensor<2xindex>) -> tensor<?x?xf32>
-   func.return %0 : tensor<?x?xf32>
- }
- 
- // CHECK-LABEL: "op_dynamic_slice"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_dynamic_slice(%arg0: tensor<16xf32>, %arg1: tensor<i64>) -> tensor<4xf32> {
--  //      CHECK: "vhlo.dynamic_slice_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.dynamic_slice_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   slice_sizes = #vhlo.tensor_v1<dense<4> : tensor<1xi64>>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.i64_v1>) -> !vhlo.tensor_v1<4x!vhlo.f32_v1>
-   %0 = "stablehlo.dynamic_slice"(%arg0, %arg1) {
-@@ -1274,15 +1366,17 @@
- }
- 
- // CHECK-LABEL: "op_dynamic_update_slice"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_dynamic_update_slice(%arg0: tensor<16xf32>, %arg1: tensor<4xf32>, %arg2: tensor<i64>) -> tensor<16xf32> {
--  // CHECK: "vhlo.dynamic_update_slice_v1"(%arg0, %arg1, %arg2) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<4x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.i64_v1>) -> !vhlo.tensor_v1<16x!vhlo.f32_v1>
-+  // CHECK: "vhlo.dynamic_update_slice_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<4x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.i64_v1>) -> !vhlo.tensor_v1<16x!vhlo.f32_v1>
-   %0 = "stablehlo.dynamic_update_slice"(%arg0, %arg1, %arg2) : (tensor<16xf32>, tensor<4xf32>, tensor<i64>) -> tensor<16xf32>
-   func.return %0 : tensor<16xf32>
- }
- 
- // CHECK-LABEL: "op_einsum"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_einsum(%arg0: tensor<8x16xf32>, %arg1: tensor<16x8xf32>) -> tensor<8x8xf32> {
--  //      CHECK: "vhlo.einsum_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.einsum_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   einsum_config = #vhlo.string_v1<"ab,bc->ac">
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<8x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<8x8x!vhlo.f32_v1>
-   %0 = "stablehlo.einsum"(%arg0, %arg1) {
-@@ -1292,22 +1386,25 @@
- }
- 
- // CHECK-LABEL: "op_exponential_minus_one"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_exponential_minus_one(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.exponential_minus_one_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.exponential_minus_one_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.exponential_minus_one"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_exponential"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_exponential(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.exponential_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.exponential_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.exponential"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_fft"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_fft(%arg0: tensor<16xcomplex<f32>>) -> tensor<16xcomplex<f32>> {
--  //      CHECK: "vhlo.fft_v1"(%arg0) <{
-+  //      CHECK: "vhlo.fft_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   fft_length = #vhlo.tensor_v1<dense<16> : tensor<1xi64>>,
-   // CHECK-SAME:   fft_type = #vhlo<fft_type_v1 FFT>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<16x!vhlo.complex_v1<!vhlo.f32_v1>>) -> !vhlo.tensor_v1<16x!vhlo.complex_v1<!vhlo.f32_v1>>
-@@ -1319,8 +1416,9 @@
- }
- 
- // CHECK-LABEL: "op_floor"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_floor(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.floor_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.floor_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.floor"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
-@@ -1333,16 +1431,17 @@
-   // CHECK-SAME:   sym_name = #vhlo.string_v1<"op_func">,
-   // CHECK-SAME:   sym_visibility = #vhlo.string_v1<"private">
-   // CHECK-SAME: }> ({
--  // CHECK-NEXT:   ^[[BB:bb.*]](%arg0: !vhlo.tensor_v1<!vhlo.f32_v1>):
--  // CHECK-NEXT:     "vhlo.return_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-+  // CHECK-NEXT:   ^[[BB:bb.*]](%[[ARG0:.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
-+  // CHECK-NEXT:     "vhlo.return_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-   // CHECK-NEXT: }) : () -> ()
- 
-   func.return %arg0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_gather"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_gather(%arg0 : tensor<2x4x9xf32>, %arg1 : tensor<1x5x2xi32>) -> tensor<1x5x1xf32> {
--  //      CHECK: "vhlo.gather_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.gather_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   collapsed_slice_dims = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-   // CHECK-SAME:   index_vector_dim = #vhlo.integer_v1<2 : i64>,
-   // CHECK-SAME:   indices_are_sorted = #vhlo.bool_v1<true>,
-@@ -1364,8 +1463,9 @@
- }
- 
- // CHECK-LABEL: "op_get_dimension_size"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_get_dimension_size(%arg0: tensor<?xf32>) -> tensor<i32> {
--  //      CHECK: "vhlo.get_dimension_size_v1"(%arg0) <{
-+  //      CHECK: "vhlo.get_dimension_size_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   dimension = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<?x!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-   %0 = "stablehlo.get_dimension_size"(%arg0) {
-@@ -1375,8 +1475,9 @@
- }
- 
- // CHECK-LABEL: "op_get_tuple_element"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_get_tuple_element(%arg0: tuple<tensor<f32>, tensor<i32>>) -> tensor<f32> {
--  //      CHECK: "vhlo.get_tuple_element_v1"(%arg0) <{
-+  //      CHECK: "vhlo.get_tuple_element_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   index = #vhlo.integer_v1<0 : i32>
-   // CHECK-SAME: }> : (!vhlo.tuple_v1<!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.get_tuple_element"(%arg0) {
-@@ -1386,11 +1487,12 @@
- }
- 
- // CHECK-LABEL: "op_if"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_if(%arg0: tensor<i1>, %arg1: tensor<f32>, %arg2: tensor<f32>) -> tensor<f32> {
--  //      CHECK: "vhlo.if_v1"(%arg0) ({
--  // CHECK-NEXT:   "vhlo.return_v1"(%arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-+  //      CHECK: "vhlo.if_v1"(%[[ARG0]]) ({
-+  // CHECK-NEXT:   "vhlo.return_v1"(%[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-   // CHECK-NEXT: }, {
--  // CHECK-NEXT:   "vhlo.return_v1"(%arg2) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-+  // CHECK-NEXT:   "vhlo.return_v1"(%[[ARG2]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-   // CHECK-NEXT: }) : (!vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.if"(%arg0) ({
-     "stablehlo.return"(%arg1) : (tensor<f32>) -> ()
-@@ -1401,15 +1503,17 @@
- }
- 
- // CHECK-LABEL: "op_imag"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_imag(%arg0: tensor<complex<f32>>) -> tensor<f32> {
--  // CHECK: "vhlo.imag_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.imag_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.imag"(%arg0) : (tensor<complex<f32>>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_infeed"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_infeed(%arg0: !stablehlo.token) -> (tensor<f32>, !stablehlo.token) {
--  //               CHECK: "vhlo.infeed_v1"(%arg0) <{
-+  //               CHECK: "vhlo.infeed_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   infeed_config = #vhlo.string_v1<"foo">,
-   // CHECK-SAME{LITERAL}:   layout = #vhlo.array_v1<[#vhlo.array_v1<[]>]>
-   //          CHECK-SAME: }> : (!vhlo.token_v1) -> (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.token_v1)
-@@ -1432,36 +1536,41 @@
- }
- 
- // CHECK-LABEL: "op_is_finite"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_is_finite(%arg0: tensor<f32>) -> tensor<i1> {
--  // CHECK: "vhlo.is_finite_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-+  // CHECK: "vhlo.is_finite_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-   %0 = "stablehlo.is_finite"(%arg0) : (tensor<f32>) -> tensor<i1>
-   func.return %0 : tensor<i1>
- }
- 
- // CHECK-LABEL: "op_log"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_log(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.log_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.log_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.log"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_log_plus_one"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_log_plus_one(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.log_plus_one_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.log_plus_one_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.log_plus_one"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_logistic"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_logistic(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.logistic_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.logistic_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.logistic"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_map"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_map(%arg0: tensor<16xf32>) -> tensor<16xf32> {
--  //      CHECK: "vhlo.map_v1"(%arg0) <{
-+  //      CHECK: "vhlo.map_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   dimensions = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>
-   // CHECK-SAME: }> ({
-   // CHECK-NEXT:   ^[[BB:bb.*]](%[[ARG1:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
-@@ -1479,57 +1588,65 @@
- }
- 
- // CHECK-LABEL: "op_maximum"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_maximum(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.maximum_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.maximum_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.maximum"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_minimum"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_minimum(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.minimum_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.minimum_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.minimum"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_multiply"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_multiply(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.multiply_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.multiply_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.multiply"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_negate"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_negate(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.negate_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.negate_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.negate"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_not"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_not(%arg0: tensor<i1>) -> tensor<i1> {
--  // CHECK: "vhlo.not_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-+  // CHECK: "vhlo.not_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-   %0 = "stablehlo.not"(%arg0) : (tensor<i1>) -> tensor<i1>
-   func.return %0 : tensor<i1>
- }
- 
- // CHECK-LABEL: "op_optimization_barrier"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_optimization_barrier(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.optimization_barrier_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.optimization_barrier_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.optimization_barrier"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_or"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_or(%arg0: tensor<i1>, %arg1: tensor<i1>) -> tensor<i1> {
--  // CHECK: "vhlo.or_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-+  // CHECK: "vhlo.or_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-   %0 = "stablehlo.or"(%arg0, %arg1) : (tensor<i1>, tensor<i1>) -> tensor<i1>
-   func.return %0 : tensor<i1>
- }
- 
- // CHECK-LABEL: "op_outfeed"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_outfeed(%arg0: tensor<f32>, %arg1: !stablehlo.token) -> !stablehlo.token {
--  //      CHECK: "vhlo.outfeed_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.outfeed_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   outfeed_config = #vhlo.string_v1<"foo">
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.token_v1) -> !vhlo.token_v1
-   %0 = "stablehlo.outfeed"(%arg0, %arg1) {
-@@ -1539,8 +1656,9 @@
- }
- 
- // CHECK-LABEL: "op_pad"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_pad(%arg0: tensor<8xf32>, %arg1: tensor<f32>) -> tensor<16xf32> {
--  //      CHECK: "vhlo.pad_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.pad_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   edge_padding_high = #vhlo.tensor_v1<dense<4> : tensor<1xi64>>,
-   // CHECK-SAME:   edge_padding_low = #vhlo.tensor_v1<dense<4> : tensor<1xi64>>,
-   // CHECK-SAME:   interior_padding = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>
-@@ -1554,36 +1672,41 @@
- }
- 
- // CHECK-LABEL: "op_popcnt"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_popcnt(%arg0: tensor<i32>) -> tensor<i32> {
--  // CHECK: "vhlo.popcnt_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-+  // CHECK: "vhlo.popcnt_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-   %0 = "stablehlo.popcnt"(%arg0) : (tensor<i32>) -> tensor<i32>
-   func.return %0 : tensor<i32>
- }
- 
- // CHECK-LABEL: "op_power"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_power(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.power_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.power_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.power"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_real_dynamic_slice"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}}, %[[ARG3:.*]]: {{.*}})
- func.func @op_real_dynamic_slice(%arg0: tensor<?xf32>, %arg1: tensor<1xindex>, %arg2: tensor<1xindex>, %arg3: tensor<1xindex>) -> tensor<?xf32> {
--  // CHECK: "vhlo.real_dynamic_slice_v1"(%arg0, %arg1, %arg2, %arg3) : (!vhlo.tensor_v1<?x!vhlo.f32_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>) -> !vhlo.tensor_v1<?x!vhlo.f32_v1>
-+  // CHECK: "vhlo.real_dynamic_slice_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]], %[[ARG3]]) : (!vhlo.tensor_v1<?x!vhlo.f32_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>) -> !vhlo.tensor_v1<?x!vhlo.f32_v1>
-   %0 = "stablehlo.real_dynamic_slice"(%arg0, %arg1, %arg2, %arg3) : (tensor<?xf32>, tensor<1xindex>, tensor<1xindex>, tensor<1xindex>) -> tensor<?xf32>
-   func.return %0 : tensor<?xf32>
- }
- 
- // CHECK-LABEL: "op_real"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_real(%arg0: tensor<complex<f32>>) -> tensor<f32> {
--  // CHECK: "vhlo.real_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.real_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.real"(%arg0) : (tensor<complex<f32>>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_recv"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_recv(%arg0: !stablehlo.token) -> (tensor<f32>, !stablehlo.token) {
--  //      CHECK: "vhlo.recv_v1"(%arg0) <{
-+  //      CHECK: "vhlo.recv_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME:   channel_type = #vhlo.integer_v1<3 : i64>,
-   // CHECK-SAME:   is_host_transfer = #vhlo.bool_v1<true>
-@@ -1596,8 +1719,9 @@
- }
- 
- // CHECK-LABEL: "op_reduce"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_reduce(%arg0: tensor<16xf32>, %arg1: tensor<f32>) -> tensor<f32> {
--  //  CHECK: "vhlo.reduce_v1"(%arg0, %arg1)
-+  //  CHECK: "vhlo.reduce_v1"(%[[ARG0]], %[[ARG1]])
-   //  CHECK:   ^[[BB:bb.*]](%[[ARG1:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>, %[[ARG2:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
-   //  CHECK:     "vhlo.return_v1"(%[[VAL1:.*]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-   //  CHECK: }) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-@@ -1612,8 +1736,9 @@
- }
- 
- // CHECK-LABEL: "op_reduce_precision"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_reduce_precision(%arg0: tensor<f32>) -> tensor<f32> {
--  //      CHECK: "vhlo.reduce_precision_v1"(%arg0) <{
-+  //      CHECK: "vhlo.reduce_precision_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   exponent_bits = #vhlo.integer_v1<8 : i32>
-   // CHECK-SAME:   mantissa_bits = #vhlo.integer_v1<10 : i32>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-@@ -1625,8 +1750,9 @@
- }
- 
- // CHECK-LABEL: "op_reduce_scatter"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_reduce_scatter(%arg0: tensor<16xf32>) -> tensor<16xf32> {
--  //               CHECK: "vhlo.reduce_scatter_v1"(%arg0) <{
-+  //               CHECK: "vhlo.reduce_scatter_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0], [1]]> : tensor<2x1xi64>>,
-   //          CHECK-SAME:   scatter_dimension = #vhlo.integer_v1<0 : i64>
-@@ -1650,8 +1776,9 @@
- }
- 
- // CHECK-LABEL: "op_reduce_window"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_reduce_window(%arg0: tensor<2x17x31x7xf32>, %arg1: tensor<f32>) -> tensor<2x9x16x7xf32> {
--  //               CHECK: "vhlo.reduce_window_v1"(%arg0, %arg1) <{
-+  //               CHECK: "vhlo.reduce_window_v1"(%[[ARG0]], %[[ARG1]]) <{
-   //          CHECK-SAME:   base_dilations = #vhlo.tensor_v1<dense<[1, 2, 2, 1]> : tensor<4xi64>>,
-   // CHECK-SAME{LITERAL}:   padding = #vhlo.tensor_v1<dense<[[0, 0], [2, 0], [0, 2], [0, 0]]> : tensor<4x2xi64>>,
-   //          CHECK-SAME:   window_dilations = #vhlo.tensor_v1<dense<[1, 2, 2, 1]> : tensor<4xi64>>,
-@@ -1677,8 +1804,9 @@
- }
- 
- // CHECK-LABEL: "op_remainder"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_remainder(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.remainder_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.remainder_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.remainder"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
-@@ -1698,16 +1826,18 @@
- }
- 
- // CHECK-LABEL: "op_reshape"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_reshape(%arg0: tensor<16xf32>) -> tensor<4x4xf32> {
--  // CHECK: "vhlo.reshape_v1"(%arg0) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<4x4x!vhlo.f32_v1>
-+  // CHECK: "vhlo.reshape_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<4x4x!vhlo.f32_v1>
-   %0 = "stablehlo.reshape"(%arg0) : (tensor<16xf32>) -> tensor<4x4xf32>
-   func.return %0 : tensor<4x4xf32>
- }
- 
- // CHECK-LABEL: "op_return"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_return(%arg0: tensor<i32>, %arg1: tensor<f32>) -> tensor<f32> {
--  //      CHECK: "vhlo.case_v1"(%arg0) ({
--  // CHECK-NEXT:   "vhlo.return_v1"(%arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-+  //      CHECK: "vhlo.case_v1"(%[[ARG0]]) ({
-+  // CHECK-NEXT:   "vhlo.return_v1"(%[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-   // CHECK-NEXT: }) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.case"(%arg0) ({
-     "stablehlo.return"(%arg1) : (tensor<f32>) -> ()
-@@ -1716,8 +1846,9 @@
- }
- 
- // CHECK-LABEL: "op_reverse"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_reverse(%arg0: tensor<16xf32>) -> tensor<16xf32> {
--  //      CHECK: "vhlo.reverse_v1"(%arg0) <{
-+  //      CHECK: "vhlo.reverse_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   dimensions = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x!vhlo.f32_v1>
-   %0 = "stablehlo.reverse"(%arg0) {
-@@ -1727,8 +1858,9 @@
- }
- 
- // CHECK-LABEL: "op_rng_bit_generator"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_rng_bit_generator(%arg0: tensor<f32>) -> (tensor<f32>, tensor<f32>) {
--  //      CHECK: "vhlo.rng_bit_generator_v1"(%arg0) <{
-+  //      CHECK: "vhlo.rng_bit_generator_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   rng_algorithm = #vhlo<rng_algorithm_v1 PHILOX>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>)
-   %0:2 = "stablehlo.rng_bit_generator"(%arg0) {
-@@ -1738,8 +1870,9 @@
- }
- 
- // CHECK-LABEL: "op_rng"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_rng(%arg0: tensor<f32>, %arg1: tensor<f32>, %arg2: tensor<0xindex>) -> tensor<f32> {
--  //      CHECK: "vhlo.rng_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.rng_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   rng_distribution = #vhlo<rng_distribution_v1 NORMAL>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<0x!vhlo.index_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.rng"(%arg0, %arg1, %arg2) {
-@@ -1749,29 +1882,33 @@
- }
- 
- // CHECK-LABEL: "op_round_nearest_afz"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_round_nearest_afz(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.round_nearest_afz_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.round_nearest_afz_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.round_nearest_afz"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_round_nearest_even"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_round_nearest_even(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.round_nearest_even_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.round_nearest_even_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.round_nearest_even"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_rsqrt"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_rsqrt(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.rsqrt_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.rsqrt_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.rsqrt"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_scatter"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_scatter(%arg0: tensor<200x100x300xf32>, %arg1: tensor<10x2xi32>, %arg2: tensor<10x300xf32>) -> tensor<200x100x300xf32> {
--  //      CHECK: "vhlo.scatter_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.scatter_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   index_vector_dim = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   indices_are_sorted = #vhlo.bool_v1<true>,
-   // CHECK-SAME:   inserted_window_dims = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-@@ -1801,8 +1938,9 @@
- }
- 
- // CHECK-LABEL: "op_select_and_scatter"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_select_and_scatter(%arg0: tensor<10x24x24x64xf32>, %arg1: tensor<12x13x13x66xf32>, %arg2: tensor<f32>) -> tensor<10x24x24x64xf32> {
--  //      CHECK: "vhlo.select_and_scatter_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.select_and_scatter_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   padding = #vhlo.tensor_v1<dense<1> : tensor<4x2xi64>>,
-   // CHECK-SAME:   window_dimensions = #vhlo.tensor_v1<dense<[1, 2, 2, 1]> : tensor<4xi64>>,
-   // CHECK-SAME:   window_strides = #vhlo.tensor_v1<dense<[1, 2, 2, 1]> : tensor<4xi64>>
-@@ -1832,15 +1970,17 @@
- }
- 
- // CHECK-LABEL: "op_select"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_select(%arg0: tensor<i1>, %arg1: tensor<f32>, %arg2: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.select_v1"(%arg0, %arg1, %arg2) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.select_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.select"(%arg0, %arg1, %arg2) : (tensor<i1>, tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_send"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_send(%arg0: tensor<f32>, %arg1: !stablehlo.token) -> !stablehlo.token {
--  //      CHECK: "vhlo.send_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.send_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME:   channel_type = #vhlo.integer_v1<2 : i64>,
-   // CHECK-SAME:   is_host_transfer = #vhlo.bool_v1<true>
-@@ -1853,8 +1993,9 @@
- }
- 
- // CHECK-LABEL: "op_set_dimension_size"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_set_dimension_size(%arg0: tensor<?xf32>, %arg1: tensor<i32>) -> tensor<16xf32> {
--  //      CHECK: "vhlo.set_dimension_size_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.set_dimension_size_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   dimension = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<?x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<16x!vhlo.f32_v1>
-   %0 = "stablehlo.set_dimension_size"(%arg0, %arg1) {
-@@ -1864,43 +2005,49 @@
- }
- 
- // CHECK-LABEL: "op_shift_left"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_shift_left(%arg0: tensor<i32>, %arg1: tensor<i32>) -> tensor<i32> {
--  // CHECK: "vhlo.shift_left_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.i32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-+  // CHECK: "vhlo.shift_left_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-   %0 = "stablehlo.shift_left"(%arg0, %arg1) : (tensor<i32>, tensor<i32>) -> tensor<i32>
-   func.return %0 : tensor<i32>
- }
- 
- // CHECK-LABEL: "op_shift_right_arithmetic"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_shift_right_arithmetic(%arg0: tensor<i32>, %arg1: tensor<i32>) -> tensor<i32> {
--  // CHECK: "vhlo.shift_right_arithmetic_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.i32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-+  // CHECK: "vhlo.shift_right_arithmetic_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-   %0 = "stablehlo.shift_right_arithmetic"(%arg0, %arg1) : (tensor<i32>, tensor<i32>) -> tensor<i32>
-   func.return %0 : tensor<i32>
- }
- 
- // CHECK-LABEL: "op_shift_right_logical"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_shift_right_logical(%arg0: tensor<i32>, %arg1: tensor<i32>) -> tensor<i32> {
--  // CHECK: "vhlo.shift_right_logical_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.i32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-+  // CHECK: "vhlo.shift_right_logical_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-   %0 = "stablehlo.shift_right_logical"(%arg0, %arg1) : (tensor<i32>, tensor<i32>) -> tensor<i32>
-   func.return %0 : tensor<i32>
- }
- 
- // CHECK-LABEL: "op_sign"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_sign(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.sign_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.sign_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.sign"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_sine"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_sine(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.sine_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.sine_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.sine"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_slice"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_slice(%arg0: tensor<16xf32>) -> tensor<4xf32> {
--  //      CHECK: "vhlo.slice_v1"(%arg0) <{
-+  //      CHECK: "vhlo.slice_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   limit_indices = #vhlo.tensor_v1<dense<4> : tensor<1xi64>>,
-   // CHECK-SAME:   start_indices = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>,
-   // CHECK-SAME:   strides = #vhlo.tensor_v1<dense<1> : tensor<1xi64>>
-@@ -1914,8 +2061,9 @@
- }
- 
- // CHECK-LABEL: "op_sort"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_sort(%arg0: tensor<16xf32>) -> tensor<16xf32> {
--  //      CHECK: "vhlo.sort_v1"(%arg0) <{
-+  //      CHECK: "vhlo.sort_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   dimension = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME:   is_stable = #vhlo.bool_v1<true>
-   // CHECK-SAME: }> ({
-@@ -1935,29 +2083,33 @@
- }
- 
- // CHECK-LABEL: "op_sqrt"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_sqrt(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.sqrt_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.sqrt_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.sqrt"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_subtract"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_subtract(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.subtract_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.subtract_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.subtract"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_tanh"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_tanh(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.tanh_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.tanh_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.tanh"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_torch_index_select"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_torch_index_select(%arg0: tensor<5x1x5xf32>, %arg1: tensor<2xi32>) ->  tensor<2x1x5xf32> {
--  //      CHECK: "vhlo.torch_index_select_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.torch_index_select_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   batch_dims = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME:   dim = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<5x1x5x!vhlo.f32_v1>, !vhlo.tensor_v1<2x!vhlo.i32_v1>) -> !vhlo.tensor_v1<2x1x5x!vhlo.f32_v1>
-@@ -1969,8 +2121,9 @@
- }
- 
- // CHECK-LABEL: "op_trace"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_trace(%arg0: tensor<f32>) {
--  //      CHECK: "vhlo.trace_v1"(%arg0) <{
-+  //      CHECK: "vhlo.trace_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   tag = #vhlo.string_v1<"foo">
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-   "stablehlo.trace"(%arg0) {
-@@ -1980,8 +2133,9 @@
- }
- 
- // CHECK-LABEL: "op_transpose"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_transpose(%arg0: tensor<16x8xf32>) ->  tensor<8x16xf32> {
--  //      CHECK: "vhlo.transpose_v1"(%arg0) <{
-+  //      CHECK: "vhlo.transpose_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   permutation = #vhlo.tensor_v1<dense<[1, 0]> : tensor<2xi64>>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<16x8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<8x16x!vhlo.f32_v1>
-   %0 = "stablehlo.transpose"(%arg0) {
-@@ -1991,8 +2145,9 @@
- }
- 
- // CHECK-LABEL: "op_triangular_solve"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_triangular_solve(%arg0: tensor<16x16xf32>, %arg1: tensor<16x16xf32>) ->  tensor<16x16xf32> {
--  //      CHECK: "vhlo.triangular_solve_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.triangular_solve_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   left_side = #vhlo.bool_v1<true>,
-   // CHECK-SAME:   lower = #vhlo.bool_v1<true>,
-   // CHECK-SAME:   transpose_a = #vhlo<transpose_v1 NO_TRANSPOSE>,
-@@ -2008,15 +2163,17 @@
- }
- 
- // CHECK-LABEL: "op_tuple"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_tuple(%arg0: tensor<f32>) -> tuple<tensor<f32>> {
--  // CHECK: "vhlo.tuple_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tuple_v1<!vhlo.tensor_v1<!vhlo.f32_v1>>
-+  // CHECK: "vhlo.tuple_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tuple_v1<!vhlo.tensor_v1<!vhlo.f32_v1>>
-   %0 = "stablehlo.tuple"(%arg0) : (tensor<f32>) -> tuple<tensor<f32>>
-   func.return %0 : tuple<tensor<f32>>
- }
- 
- // CHECK-LABEL: "op_unary_einsum"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_unary_einsum(%arg0: tensor<8x16xf32>) -> tensor<8xf32> {
--  //      CHECK: "vhlo.unary_einsum_v1"(%arg0) <{
-+  //      CHECK: "vhlo.unary_einsum_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   einsum_config = #vhlo.string_v1<"ab->a">
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<8x16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<8x!vhlo.f32_v1>
-   %0 = "stablehlo.unary_einsum"(%arg0) {
-@@ -2026,22 +2183,25 @@
- }
- 
- // CHECK-LABEL: "op_uniform_dequantize"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_uniform_dequantize(%arg0: tensor<!quant.uniform<i8:f32, 34.0:16>>) -> tensor<f32> {
--  // CHECK: "vhlo.uniform_dequantize_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.uniform_dequantize_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.uniform_dequantize"(%arg0) : (tensor<!quant.uniform<i8:f32, 34.0:16>>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_uniform_quantize"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_uniform_quantize(%arg0: tensor<f32>) -> tensor<!quant.uniform<i8:f32, 34.0:16>> {
--  // CHECK: "vhlo.uniform_quantize_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>
-+  // CHECK: "vhlo.uniform_quantize_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>
-   %0 = "stablehlo.uniform_quantize"(%arg0) : (tensor<f32>) -> tensor<!quant.uniform<i8:f32, 34.0:16>>
-   func.return %0 : tensor<!quant.uniform<i8:f32, 34.0:16>>
- }
- 
- // CHECK-LABEL: "op_while"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_while(%arg0: tensor<i1>) -> tensor<i1> {
--  //      CHECK: "vhlo.while_v1"(%arg0) ({
-+  //      CHECK: "vhlo.while_v1"(%[[ARG0]]) ({
-   // CHECK-NEXT:   ^[[BB:bb.*]](%[[ARG1:arg.*]]: !vhlo.tensor_v1<!vhlo.bool_v1>):
-   // CHECK-NEXT:     "vhlo.return_v1"(%[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>) -> ()
-   // CHECK-NEXT:   }, {
-@@ -2059,8 +2219,9 @@
- }
- 
- // CHECK-LABEL: "op_xor"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_xor(%arg0: tensor<i1>, %arg1: tensor<i1>) -> tensor<i1> {
--  // CHECK: "vhlo.xor_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-+  // CHECK: "vhlo.xor_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-   %0 = "stablehlo.xor"(%arg0, %arg1) : (tensor<i1>, tensor<i1>) -> tensor<i1>
-   func.return %0 : tensor<i1>
- }
-@@ -2068,190 +2229,217 @@
- // ============ TYPES ============
- 
- // CHECK-LABEL: "type_i1"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_i1(%arg0: tensor<i1>, %arg1: tensor<i1>) -> tensor<i1> {
--  // CHECK: "vhlo.and_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-+  // CHECK: "vhlo.and_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-   %0 = "stablehlo.and"(%arg0, %arg1) : (tensor<i1>, tensor<i1>) -> tensor<i1>
-   func.return %0 : tensor<i1>
- }
- 
- // CHECK-LABEL: "type_i4"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_i4(%arg0: tensor<i4>, %arg1: tensor<i4>) -> tensor<i4> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.i4_v1>, !vhlo.tensor_v1<!vhlo.i4_v1>) -> !vhlo.tensor_v1<!vhlo.i4_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i4_v1>, !vhlo.tensor_v1<!vhlo.i4_v1>) -> !vhlo.tensor_v1<!vhlo.i4_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<i4>, tensor<i4>) -> tensor<i4>
-   func.return %0 : tensor<i4>
- }
- 
- // CHECK-LABEL: "type_i8"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_i8(%arg0: tensor<i8>, %arg1: tensor<i8>) -> tensor<i8> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.i8_v1>, !vhlo.tensor_v1<!vhlo.i8_v1>) -> !vhlo.tensor_v1<!vhlo.i8_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i8_v1>, !vhlo.tensor_v1<!vhlo.i8_v1>) -> !vhlo.tensor_v1<!vhlo.i8_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<i8>, tensor<i8>) -> tensor<i8>
-   func.return %0 : tensor<i8>
- }
- 
- // CHECK-LABEL: "type_i16"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_i16(%arg0: tensor<i16>, %arg1: tensor<i16>) -> tensor<i16> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.i16_v1>, !vhlo.tensor_v1<!vhlo.i16_v1>) -> !vhlo.tensor_v1<!vhlo.i16_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i16_v1>, !vhlo.tensor_v1<!vhlo.i16_v1>) -> !vhlo.tensor_v1<!vhlo.i16_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<i16>, tensor<i16>) -> tensor<i16>
-   func.return %0 : tensor<i16>
- }
- 
- // CHECK-LABEL: "type_i32"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_i32(%arg0: tensor<i32>, %arg1: tensor<i32>) -> tensor<i32> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.i32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<i32>, tensor<i32>) -> tensor<i32>
-   func.return %0 : tensor<i32>
- }
- 
- // CHECK-LABEL: "type_i64"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_i64(%arg0: tensor<i64>, %arg1: tensor<i64>) -> tensor<i64> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.i64_v1>, !vhlo.tensor_v1<!vhlo.i64_v1>) -> !vhlo.tensor_v1<!vhlo.i64_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i64_v1>, !vhlo.tensor_v1<!vhlo.i64_v1>) -> !vhlo.tensor_v1<!vhlo.i64_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<i64>, tensor<i64>) -> tensor<i64>
-   func.return %0 : tensor<i64>
- }
- 
- // CHECK-LABEL: "type_ui4"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_ui4(%arg0: tensor<ui4>, %arg1: tensor<ui4>) -> tensor<ui4> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.ui4_v1>, !vhlo.tensor_v1<!vhlo.ui4_v1>) -> !vhlo.tensor_v1<!vhlo.ui4_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.ui4_v1>, !vhlo.tensor_v1<!vhlo.ui4_v1>) -> !vhlo.tensor_v1<!vhlo.ui4_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<ui4>, tensor<ui4>) -> tensor<ui4>
-   func.return %0 : tensor<ui4>
- }
- 
- // CHECK-LABEL: "type_ui8"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_ui8(%arg0: tensor<ui8>, %arg1: tensor<ui8>) -> tensor<ui8> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.ui8_v1>, !vhlo.tensor_v1<!vhlo.ui8_v1>) -> !vhlo.tensor_v1<!vhlo.ui8_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.ui8_v1>, !vhlo.tensor_v1<!vhlo.ui8_v1>) -> !vhlo.tensor_v1<!vhlo.ui8_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<ui8>, tensor<ui8>) -> tensor<ui8>
-   func.return %0 : tensor<ui8>
- }
- 
- // CHECK-LABEL: "type_ui16"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_ui16(%arg0: tensor<ui16>, %arg1: tensor<ui16>) -> tensor<ui16> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.ui16_v1>, !vhlo.tensor_v1<!vhlo.ui16_v1>) -> !vhlo.tensor_v1<!vhlo.ui16_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.ui16_v1>, !vhlo.tensor_v1<!vhlo.ui16_v1>) -> !vhlo.tensor_v1<!vhlo.ui16_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<ui16>, tensor<ui16>) -> tensor<ui16>
-   func.return %0 : tensor<ui16>
- }
- 
- // CHECK-LABEL: "type_ui32"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_ui32(%arg0: tensor<ui32>, %arg1: tensor<ui32>) -> tensor<ui32> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.ui32_v1>, !vhlo.tensor_v1<!vhlo.ui32_v1>) -> !vhlo.tensor_v1<!vhlo.ui32_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.ui32_v1>, !vhlo.tensor_v1<!vhlo.ui32_v1>) -> !vhlo.tensor_v1<!vhlo.ui32_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<ui32>, tensor<ui32>) -> tensor<ui32>
-   func.return %0 : tensor<ui32>
- }
- 
- // CHECK-LABEL: "type_ui64"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_ui64(%arg0: tensor<ui64>, %arg1: tensor<ui64>) -> tensor<ui64> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.ui64_v1>, !vhlo.tensor_v1<!vhlo.ui64_v1>) -> !vhlo.tensor_v1<!vhlo.ui64_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.ui64_v1>, !vhlo.tensor_v1<!vhlo.ui64_v1>) -> !vhlo.tensor_v1<!vhlo.ui64_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<ui64>, tensor<ui64>) -> tensor<ui64>
-   func.return %0 : tensor<ui64>
- }
- 
- // CHECK-LABEL: "type_f8E4M3FN"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_f8E4M3FN(%arg0: tensor<f8E4M3FN>, %arg1: tensor<f8E4M3FN>) -> tensor<f8E4M3FN> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f8E4M3FN_v1>, !vhlo.tensor_v1<!vhlo.f8E4M3FN_v1>) -> !vhlo.tensor_v1<!vhlo.f8E4M3FN_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f8E4M3FN_v1>, !vhlo.tensor_v1<!vhlo.f8E4M3FN_v1>) -> !vhlo.tensor_v1<!vhlo.f8E4M3FN_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f8E4M3FN>, tensor<f8E4M3FN>) -> tensor<f8E4M3FN>
-   func.return %0 : tensor<f8E4M3FN>
- }
- 
- // CHECK-LABEL: "type_f8E5M2"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_f8E5M2(%arg0: tensor<f8E5M2>, %arg1: tensor<f8E5M2>) -> tensor<f8E5M2> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f8E5M2_v1>, !vhlo.tensor_v1<!vhlo.f8E5M2_v1>) -> !vhlo.tensor_v1<!vhlo.f8E5M2_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f8E5M2_v1>, !vhlo.tensor_v1<!vhlo.f8E5M2_v1>) -> !vhlo.tensor_v1<!vhlo.f8E5M2_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f8E5M2>, tensor<f8E5M2>) -> tensor<f8E5M2>
-   func.return %0 : tensor<f8E5M2>
- }
- 
- // CHECK-LABEL: "type_f8E4M3FNUZ"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_f8E4M3FNUZ(%arg0: tensor<f8E4M3FNUZ>, %arg1: tensor<f8E4M3FNUZ>) -> tensor<f8E4M3FNUZ> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f8E4M3FNUZ_v1>, !vhlo.tensor_v1<!vhlo.f8E4M3FNUZ_v1>) -> !vhlo.tensor_v1<!vhlo.f8E4M3FNUZ_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f8E4M3FNUZ_v1>, !vhlo.tensor_v1<!vhlo.f8E4M3FNUZ_v1>) -> !vhlo.tensor_v1<!vhlo.f8E4M3FNUZ_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f8E4M3FNUZ>, tensor<f8E4M3FNUZ>) -> tensor<f8E4M3FNUZ>
-   func.return %0 : tensor<f8E4M3FNUZ>
- }
- 
- // CHECK-LABEL: "type_f8E4M3B11FNUZ"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_f8E4M3B11FNUZ(%arg0: tensor<f8E4M3B11FNUZ>, %arg1: tensor<f8E4M3B11FNUZ>) -> tensor<f8E4M3B11FNUZ> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f8E4M3B11FNUZ_v1>, !vhlo.tensor_v1<!vhlo.f8E4M3B11FNUZ_v1>) -> !vhlo.tensor_v1<!vhlo.f8E4M3B11FNUZ_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f8E4M3B11FNUZ_v1>, !vhlo.tensor_v1<!vhlo.f8E4M3B11FNUZ_v1>) -> !vhlo.tensor_v1<!vhlo.f8E4M3B11FNUZ_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f8E4M3B11FNUZ>, tensor<f8E4M3B11FNUZ>) -> tensor<f8E4M3B11FNUZ>
-   func.return %0 : tensor<f8E4M3B11FNUZ>
- }
- 
- // CHECK-LABEL: "type_f8E5M2FNUZ"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_f8E5M2FNUZ(%arg0: tensor<f8E5M2FNUZ>, %arg1: tensor<f8E5M2FNUZ>) -> tensor<f8E5M2FNUZ> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f8E5M2FNUZ_v1>, !vhlo.tensor_v1<!vhlo.f8E5M2FNUZ_v1>) -> !vhlo.tensor_v1<!vhlo.f8E5M2FNUZ_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f8E5M2FNUZ_v1>, !vhlo.tensor_v1<!vhlo.f8E5M2FNUZ_v1>) -> !vhlo.tensor_v1<!vhlo.f8E5M2FNUZ_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f8E5M2FNUZ>, tensor<f8E5M2FNUZ>) -> tensor<f8E5M2FNUZ>
-   func.return %0 : tensor<f8E5M2FNUZ>
- }
- 
- // CHECK-LABEL: "type_bf16"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_bf16(%arg0: tensor<bf16>, %arg1: tensor<bf16>) -> tensor<bf16> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.bf16_v1>, !vhlo.tensor_v1<!vhlo.bf16_v1>) -> !vhlo.tensor_v1<!vhlo.bf16_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.bf16_v1>, !vhlo.tensor_v1<!vhlo.bf16_v1>) -> !vhlo.tensor_v1<!vhlo.bf16_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<bf16>, tensor<bf16>) -> tensor<bf16>
-   func.return %0 : tensor<bf16>
- }
- 
- // CHECK-LABEL: "type_f16"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_f16(%arg0: tensor<f16>, %arg1: tensor<f16>) -> tensor<f16> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f16_v1>, !vhlo.tensor_v1<!vhlo.f16_v1>) -> !vhlo.tensor_v1<!vhlo.f16_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f16_v1>, !vhlo.tensor_v1<!vhlo.f16_v1>) -> !vhlo.tensor_v1<!vhlo.f16_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f16>, tensor<f16>) -> tensor<f16>
-   func.return %0 : tensor<f16>
- }
- 
- // CHECK-LABEL: "type_f32"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_f32(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "type_f64"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_f64(%arg0: tensor<f64>, %arg1: tensor<f64>) -> tensor<f64> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f64_v1>, !vhlo.tensor_v1<!vhlo.f64_v1>) -> !vhlo.tensor_v1<!vhlo.f64_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f64_v1>, !vhlo.tensor_v1<!vhlo.f64_v1>) -> !vhlo.tensor_v1<!vhlo.f64_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f64>, tensor<f64>) -> tensor<f64>
-   func.return %0 : tensor<f64>
- }
- 
- // CHECK-LABEL: "type_complex_f32"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_complex_f32(%arg0: tensor<complex<f32>>, %arg1: tensor<complex<f32>>) -> tensor<complex<f32>> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>, !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>) -> !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>, !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>) -> !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<complex<f32>>, tensor<complex<f32>>) -> tensor<complex<f32>>
-   func.return %0 : tensor<complex<f32>>
- }
- 
- // CHECK-LABEL: "type_complex_f64"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_complex_f64(%arg0: tensor<complex<f64>>, %arg1: tensor<complex<f64>>) -> tensor<complex<f64>> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f64_v1>>, !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f64_v1>>) -> !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f64_v1>>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f64_v1>>, !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f64_v1>>) -> !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f64_v1>>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<complex<f64>>, tensor<complex<f64>>) -> tensor<complex<f64>>
-   func.return %0 : tensor<complex<f64>>
- }
- 
- // CHECK-LABEL: "type_dynamism_ranked"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @type_dynamism_ranked(%arg0: tensor<?xf32>) -> tensor<?xf32> {
--  // CHECK: "vhlo.abs_v1"(%arg0) : (!vhlo.tensor_v1<?x!vhlo.f32_v1>) -> !vhlo.tensor_v1<?x!vhlo.f32_v1>
-+  // CHECK: "vhlo.abs_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<?x!vhlo.f32_v1>) -> !vhlo.tensor_v1<?x!vhlo.f32_v1>
-   %0 = "stablehlo.abs"(%arg0) : (tensor<?xf32>) -> tensor<?xf32>
-   func.return %0 : tensor<?xf32>
- }
- 
- // CHECK-LABEL: "type_quantization"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_quantization(%arg0: tensor<!quant.uniform<i8:f32, 34.0:16>>, %arg1: tensor<!quant.uniform<i8:f32, 34.0:16>>) -> tensor<!quant.uniform<i8:f32, 34.0:16>> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>, !vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>) -> !vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>, !vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>) -> !vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<!quant.uniform<i8:f32, 34.0:16>>, tensor<!quant.uniform<i8:f32, 34.0:16>>) -> tensor<!quant.uniform<i8:f32, 34.0:16>>
-   func.return %0 : tensor<!quant.uniform<i8:f32, 34.0:16>>
- }
- 
- //       CHECK: function_type = #vhlo.type_v1<!vhlo.func_v1<(!vhlo.token_v1) -> !vhlo.token_v1>>
- // CHECK-LABEL: "type_token_callee"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @type_token_callee(%arg0: !stablehlo.token) -> !stablehlo.token {
--  // CHECK: "vhlo.return_v1"(%arg0) : (!vhlo.token_v1) -> ()
-+  // CHECK: "vhlo.return_v1"(%[[ARG0]]) : (!vhlo.token_v1) -> ()
-   return %arg0 : !stablehlo.token
- }
- 
- //       CHECK: function_type = #vhlo.type_v1<!vhlo.func_v1<(!vhlo.token_v1) -> !vhlo.token_v1>>
- // CHECK-LABEL: "type_token_caller"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @type_token_caller(%arg0: !stablehlo.token) -> !stablehlo.token {
--  // CHECK:      "vhlo.call_v1"(%arg0) <{callee = #vhlo.string_v1<"type_token_callee">}
-+  // CHECK:      "vhlo.call_v1"(%[[ARG0]]) <{callee = #vhlo.string_v1<"type_token_callee">}
-   // CHECK-SAME: (!vhlo.token_v1) -> !vhlo.token_v1
-   %0 = func.call @type_token_callee(%arg0) : (!stablehlo.token) -> !stablehlo.token
-   return %0 : !stablehlo.token
- }
- 
- // CHECK-LABEL: "type_tuple"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @type_tuple(%arg0: tuple<tensor<f32>>) -> tuple<!stablehlo.token> {
-   %0 = "stablehlo.custom_call"(%arg0) {
-     call_target_name = "foo"
-diff --ruN a/stablehlo/stablehlo/tests/vhlo/stablehlo_legalize_to_vhlo.0_16_0.mlir b/stablehlo/stablehlo/tests/vhlo/stablehlo_legalize_to_vhlo.0_16_0.mlir
---- stablehlo/stablehlo/tests/vhlo/stablehlo_legalize_to_vhlo.0_16_0.mlir
-+++ stablehlo/stablehlo/tests/vhlo/stablehlo_legalize_to_vhlo.0_16_0.mlir
-@@ -13,6 +13,7 @@
- // ============ ATTRIBUTES ============
- 
- // CHECK-LABEL: "attr_comparison_direction_eq"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_direction_eq(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     // CHECK: comparison_direction = #vhlo<comparison_direction_v1 EQ>
-@@ -22,6 +23,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_direction_ne"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_direction_ne(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     // CHECK: comparison_direction = #vhlo<comparison_direction_v1 NE>
-@@ -31,6 +33,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_direction_ge"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_direction_ge(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     // CHECK: comparison_direction = #vhlo<comparison_direction_v1 GE>
-@@ -40,6 +43,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_direction_gt"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_direction_gt(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     // CHECK: comparison_direction = #vhlo<comparison_direction_v1 GT>
-@@ -49,6 +53,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_direction_le"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_direction_le(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     // CHECK: comparison_direction = #vhlo<comparison_direction_v1 LE>
-@@ -58,6 +63,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_direction_lt"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_direction_lt(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     // CHECK: comparison_direction = #vhlo<comparison_direction_v1 LT>
-@@ -67,6 +73,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_type_notype"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_type_notype(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     comparison_direction = #stablehlo<comparison_direction EQ>
-@@ -76,6 +83,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_type_float"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_type_float(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     comparison_direction = #stablehlo<comparison_direction EQ>,
-@@ -86,6 +94,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_type_totalorder"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_type_totalorder(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     comparison_direction = #stablehlo<comparison_direction EQ>,
-@@ -96,6 +105,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_type_signed"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_type_signed(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     comparison_direction = #stablehlo<comparison_direction EQ>,
-@@ -106,6 +116,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_type_unsigned"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_type_unsigned(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     comparison_direction = #stablehlo<comparison_direction EQ>,
-@@ -118,6 +129,7 @@
- // ConvDimensionNumbers aka #stablehlo.conv is covered below.
- 
- // CHECK-LABEL: "attr_custom_call_api_version_unspecified"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_custom_call_api_version_unspecified(%arg0: tensor<f32>) -> tensor<f32> {
-   %0 = "stablehlo.custom_call"(%arg0) {
-     call_target_name = "foo",
-@@ -128,6 +140,7 @@
- }
- 
- // CHECK-LABEL: "attr_custom_call_api_version_original"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_custom_call_api_version_original(%arg0: tensor<f32>) -> tensor<f32> {
-   %0 = "stablehlo.custom_call"(%arg0) {
-     call_target_name = "foo",
-@@ -138,6 +151,7 @@
- }
- 
- // CHECK-LABEL: "attr_custom_call_api_version_status_returning"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_custom_call_api_version_status_returning(%arg0: tensor<f32>) -> tensor<f32> {
-   %0 = "stablehlo.custom_call"(%arg0) {
-     call_target_name = "foo",
-@@ -148,6 +162,7 @@
- }
- 
- // CHECK-LABEL: "attr_custom_call_api_version_status_returning_unified"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_custom_call_api_version_status_returning_unified(%arg0: tensor<f32>) -> tensor<f32> {
-   %0 = "stablehlo.custom_call"(%arg0) {
-     call_target_name = "foo",
-@@ -166,6 +181,7 @@
- // DotDimensionNumbers aka #stablehlo.dot is covered below.
- 
- // CHECK-LABEL: "attr_fft_type_fft"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_fft_type_fft(%arg0: tensor<16xcomplex<f32>>) -> tensor<16xcomplex<f32>> {
-   %0 = "stablehlo.fft"(%arg0) {
-     // CHECK: fft_type = #vhlo<fft_type_v1 FFT>
-@@ -176,6 +192,7 @@
- }
- 
- // CHECK-LABEL: "attr_fft_type_ifft"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_fft_type_ifft(%arg0: tensor<16xcomplex<f32>>) -> tensor<16xcomplex<f32>> {
-   %0 = "stablehlo.fft"(%arg0) {
-     // CHECK: fft_type = #vhlo<fft_type_v1 IFFT>
-@@ -186,6 +203,7 @@
- }
- 
- // CHECK-LABEL: "attr_fft_type_rfft"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_fft_type_rfft(%arg0: tensor<16xf32>) -> tensor<9xcomplex<f32>> {
-   %0 = "stablehlo.fft"(%arg0) {
-     // CHECK: fft_type = #vhlo<fft_type_v1 RFFT>
-@@ -196,6 +214,7 @@
- }
- 
- // CHECK-LABEL: "attr_fft_type_irfft"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_fft_type_irfft(%arg0: tensor<9xcomplex<f32>>) -> tensor<16xf32> {
-   %0 = "stablehlo.fft"(%arg0) {
-     // CHECK: fft_type = #vhlo<fft_type_v1 IRFFT>
-@@ -208,6 +227,7 @@
- // GatherDimensionNumbers aka #stablehlo.gather is covered below.
- 
- // CHECK-LABEL: "attr_precision_config_default"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_precision_config_default(%arg0: tensor<8x16xf32>, %arg1: tensor<16x8xf32>) -> tensor<8x8xf32> {
-   %0 = "stablehlo.dot"(%arg0, %arg1) {
-     // CHECK: precision_config = #vhlo.array_v1<[#vhlo<precision_v1 DEFAULT>, #vhlo<precision_v1 DEFAULT>]>
-@@ -216,6 +236,7 @@
- }
- 
- // CHECK-LABEL: "attr_precision_config_high"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_precision_config_high(%arg0: tensor<8x16xf32>, %arg1: tensor<16x8xf32>) -> tensor<8x8xf32> {
-   %0 = "stablehlo.dot"(%arg0, %arg1) {
-     // CHECK: precision_config = #vhlo.array_v1<[#vhlo<precision_v1 HIGH>, #vhlo<precision_v1 HIGH>]>
-@@ -225,6 +246,7 @@
- }
- 
- // CHECK-LABEL: "attr_precision_config_highest"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_precision_config_highest(%arg0: tensor<8x16xf32>, %arg1: tensor<16x8xf32>) -> tensor<8x8xf32> {
-   %0 = "stablehlo.dot"(%arg0, %arg1) {
-     // CHECK: precision_config = #vhlo.array_v1<[#vhlo<precision_v1 HIGHEST>, #vhlo<precision_v1 HIGHEST>]>
-@@ -234,6 +256,7 @@
- }
- 
- // CHECK-LABEL: "attr_rng_algorithm_default"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_rng_algorithm_default(%arg0: tensor<f32>) -> (tensor<f32>, tensor<f32>) {
-   %0:2 = "stablehlo.rng_bit_generator"(%arg0) {
-     // CHECK: rng_algorithm = #vhlo<rng_algorithm_v1 DEFAULT>
-@@ -243,6 +266,7 @@
- }
- 
- // CHECK-LABEL: "attr_rng_algorithm_three_fry"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_rng_algorithm_three_fry(%arg0: tensor<f32>) -> (tensor<f32>, tensor<f32>) {
-   %0:2 = "stablehlo.rng_bit_generator"(%arg0) {
-     // CHECK: rng_algorithm = #vhlo<rng_algorithm_v1 THREE_FRY>
-@@ -252,6 +276,7 @@
- }
- 
- // CHECK-LABEL: "attr_rng_algorithm_philox"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_rng_algorithm_philox(%arg0: tensor<f32>) -> (tensor<f32>, tensor<f32>) {
-   %0:2 = "stablehlo.rng_bit_generator"(%arg0) {
-     // CHECK: rng_algorithm = #vhlo<rng_algorithm_v1 PHILOX>
-@@ -261,6 +286,7 @@
- }
- 
- // CHECK-LABEL: "attr_rng_distribution_uniform"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @attr_rng_distribution_uniform(%arg0: tensor<f32>, %arg1: tensor<f32>, %arg2: tensor<0xindex>) -> tensor<f32> {
-   %0 = "stablehlo.rng"(%arg0, %arg1, %arg2) {
-     // CHECK: rng_distribution = #vhlo<rng_distribution_v1 UNIFORM>
-@@ -270,6 +296,7 @@
- }
- 
- // CHECK-LABEL: "attr_rng_distribution_normal"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @attr_rng_distribution_normal(%arg0: tensor<f32>, %arg1: tensor<f32>, %arg2: tensor<0xindex>) -> tensor<f32> {
-   %0 = "stablehlo.rng"(%arg0, %arg1, %arg2) {
-     // CHECK: rng_distribution = #vhlo<rng_distribution_v1 NORMAL>
-@@ -281,6 +308,7 @@
- // ScatterDimensionNumbers aka #stablehlo.scatter is covered below.
- 
- // CHECK-LABEL: "attr_transpose_no_transpose"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_transpose_no_transpose(%arg0: tensor<16x16xf32>, %arg1: tensor<16x16xf32>) ->  tensor<16x16xf32> {
-   %0 = "stablehlo.triangular_solve"(%arg0, %arg1) {
-     left_side = true,
-@@ -293,6 +321,7 @@
- }
- 
- // CHECK-LABEL: "attr_transpose_transpose"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_transpose_transpose(%arg0: tensor<16x16xf32>, %arg1: tensor<16x16xf32>) ->  tensor<16x16xf32> {
-   %0 = "stablehlo.triangular_solve"(%arg0, %arg1) {
-     left_side = true,
-@@ -305,6 +334,7 @@
- }
- 
- // CHECK-LABEL: "attr_transpose_adjoint"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_transpose_adjoint(%arg0: tensor<16x16xf32>, %arg1: tensor<16x16xf32>) ->  tensor<16x16xf32> {
-   %0 = "stablehlo.triangular_solve"(%arg0, %arg1) {
-     left_side = true,
-@@ -319,10 +349,9 @@
- // TypeExtensionsAttr aka #stablehlo.type_extensions is covered below.
- 
- // CHECK-LABEL: "attr_type_extensions_bounds"
--func.func @attr_type_extensions_bounds(
--    %arg0: tensor<?x?xf32, #stablehlo.type_extensions<bounds = [16, ?]>>)
--    -> tensor<?x?xf32, #stablehlo.type_extensions<bounds = [16, ?]>> {
--  // CHECK: "vhlo.return_v1"(%arg0) : (!vhlo.tensor_v1<?x?x!vhlo.f32_v1, #vhlo.type_extensions_v1<bounds = [16, ?]>>) -> ()
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @attr_type_extensions_bounds(%arg0: tensor<?x?xf32, #stablehlo.type_extensions<bounds = [16, ?]>>) -> tensor<?x?xf32, #stablehlo.type_extensions<bounds = [16, ?]>> {
-+  // CHECK: "vhlo.return_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<?x?x!vhlo.f32_v1, #vhlo.type_extensions_v1<bounds = [16, ?]>>) -> ()
-   func.return %arg0 : tensor<?x?xf32, #stablehlo.type_extensions<bounds = [16, ?]>>
- }
- 
-@@ -330,8 +359,9 @@
- // ============ DEFAULTS ============
- 
- // CHECK-LABEL: "default_all_gather"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_all_gather(%arg0: tensor<16x8xf32>) -> tensor<16x16xf32> {
--  //               CHECK: "vhlo.all_gather_v1"(%arg0) <{
-+  //               CHECK: "vhlo.all_gather_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   all_gather_dim = #vhlo.integer_v1<1 : i64>
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0], [1]]> : tensor<2x1xi64>>,
-@@ -345,8 +375,9 @@
- }
- 
- // CHECK-LABEL: "default_all_reduce"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_all_reduce(%arg0: tensor<f32>) -> tensor<f32> {
--  //               CHECK: "vhlo.all_reduce_v1"(%arg0)
-+  //               CHECK: "vhlo.all_reduce_v1"(%[[ARG0]])
-   //          CHECK-SAME: <{
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0], [1]]> : tensor<2x1xi64>>,
-@@ -368,8 +399,9 @@
- }
- 
- // CHECK-LABEL: "default_all_to_all"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_all_to_all(%arg0: tensor<4x16xf32>) -> tensor<16x4xf32> {
--  //               CHECK: "vhlo.all_to_all_v1"(%arg0) <{
-+  //               CHECK: "vhlo.all_to_all_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   //          CHECK-SAME:   concat_dimension = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0, 1, 2, 3]]> : tensor<1x4xi64>>,
-@@ -386,8 +418,9 @@
- }
- 
- // CHECK-LABEL: "default_cholesky"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_cholesky(%arg0: tensor<1x16x16xf32>) -> tensor<1x16x16xf32> {
--  //      CHECK: "vhlo.cholesky_v1"(%arg0) <{
-+  //      CHECK: "vhlo.cholesky_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   lower = #vhlo.bool_v1<false>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<1x16x16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<1x16x16x!vhlo.f32_v1>
-   %0 = "stablehlo.cholesky"(%arg0) : (tensor<1x16x16xf32>) -> tensor<1x16x16xf32>
-@@ -395,8 +428,9 @@
- }
- 
- // CHECK-LABEL: "default_collective_permute"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_collective_permute(%arg0: tensor<16x8xf32>) -> tensor<16x8xf32> {
--  //               CHECK: "vhlo.collective_permute_v1"(%arg0) <{
-+  //               CHECK: "vhlo.collective_permute_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME{LITERAL}:   source_target_pairs = #vhlo.tensor_v1<dense<[[0, 1], [1, 2], [2, 3]]> : tensor<3x2xi64>>
-   //          CHECK-SAME: }> : (!vhlo.tensor_v1<16x8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x8x!vhlo.f32_v1>
-@@ -407,8 +441,9 @@
- }
- 
- // CHECK-LABEL: "default_collective_broadcast"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_collective_broadcast(%arg0: tensor<16x8xf32>) -> tensor<16x8xf32> {
--  //               CHECK: "vhlo.collective_broadcast_v1"(%arg0) <{
-+  //               CHECK: "vhlo.collective_broadcast_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0, 1]]> : tensor<1x2xi64>>
-   //          CHECK-SAME: }> : (!vhlo.tensor_v1<16x8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x8x!vhlo.f32_v1>
-@@ -419,8 +454,9 @@
- }
- 
- // CHECK-LABEL: "default_compare"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @default_compare(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
--  //      CHECK: "vhlo.compare_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.compare_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   compare_type = #vhlo<comparison_type_v1 NOTYPE>,
-   // CHECK-SAME:   comparison_direction = #vhlo<comparison_direction_v1 EQ>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-@@ -431,8 +467,9 @@
- }
- 
- // CHECK-LABEL: "default_convolution"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @default_convolution(%arg0: tensor<1x8x8x207xf32>, %arg1: tensor<3x3x207x16xf32>) -> tensor<1x6x6x16xf32> {
--  //      CHECK: "vhlo.convolution_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.convolution_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   batch_group_count = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   feature_group_count = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   input_batch_dimension = #vhlo.integer_v1<0 : i64>,
-@@ -460,8 +497,9 @@
- }
- 
- // CHECK-LABEL: "default_custom_call"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_custom_call(%arg0: tensor<f32>) -> tensor<f32> {
--  //      CHECK: "vhlo.custom_call_v1"(%arg0) <{
-+  //      CHECK: "vhlo.custom_call_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   api_version = #vhlo<api_version_v1 API_VERSION_ORIGINAL>,
-   // CHECK-SAME:   backend_config = #vhlo.string_v1<"">,
-   // CHECK-SAME:   call_target_name = #vhlo.string_v1<"foo">,
-@@ -478,8 +516,9 @@
- }
- 
- // CHECK-LABEL: "default_dot_general"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @default_dot_general(%arg0: tensor<8x8x16xf32>, %arg1: tensor<8x16x8xf32>) -> tensor<8x8x8xf32> {
--  //      CHECK: "vhlo.dot_general_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.dot_general_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   lhs_batching_dimensions = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>,
-   // CHECK-SAME:   lhs_contracting_dimensions = #vhlo.tensor_v1<dense<2> : tensor<1xi64>>,
-   // CHECK-SAME:   precision_config = #vhlo.array_v1<[#vhlo<precision_v1 DEFAULT>, #vhlo<precision_v1 DEFAULT>]>,
-@@ -498,8 +537,9 @@
- }
- 
- // CHECK-LABEL: "default_dot"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @default_dot(%arg0: tensor<8x16xf32>, %arg1: tensor<16x8xf32>) -> tensor<8x8xf32> {
--  //      CHECK: "vhlo.dot_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.dot_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   precision_config = #vhlo.array_v1<[#vhlo<precision_v1 DEFAULT>, #vhlo<precision_v1 DEFAULT>]>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<8x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<8x8x!vhlo.f32_v1>
-   %0 = "stablehlo.dot"(%arg0, %arg1) : (tensor<8x16xf32>, tensor<16x8xf32>) -> tensor<8x8xf32>
-@@ -507,8 +547,9 @@
- }
- 
- // CHECK-LABEL: "default_dynamic_broadcast_in_dim"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @default_dynamic_broadcast_in_dim(%arg0: tensor<?x?xf32>, %arg1: tensor<2xindex>) -> tensor<?x?xf32> {
--  //      CHECK: "vhlo.dynamic_broadcast_in_dim_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.dynamic_broadcast_in_dim_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   broadcast_dimensions = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-   // CHECK-SAME:   known_expanding_dimensions = #vhlo.tensor_v1<dense<> : tensor<0xi64>>,
-   // CHECK-SAME:   known_nonexpanding_dimensions = #vhlo.tensor_v1<dense<> : tensor<0xi64>>
-@@ -520,8 +561,9 @@
- }
- 
- // CHECK-LABEL: "default_dynamic_conv"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @default_dynamic_conv(%arg0: tensor<1x8x8x207xf32>, %arg1: tensor<3x3x207x16xf32>, %arg2: tensor<4xi32>) -> tensor<1x?x?x16xf32> {
--  //      CHECK: "vhlo.dynamic_conv_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.dynamic_conv_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   batch_group_count = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   feature_group_count = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   input_batch_dimension = #vhlo.integer_v1<0 : i64>,
-@@ -549,8 +591,9 @@
- }
- 
- // CHECK-LABEL: "default_dynamic_gather"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @default_dynamic_gather(%arg0 : tensor<2x4x9xf32>, %arg1 : tensor<1x5x2xi32>, %arg2 : tensor<3xi32>) -> tensor<1x5x8xf32> {
--  //      CHECK: "vhlo.dynamic_gather_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.dynamic_gather_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   collapsed_slice_dims = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-   // CHECK-SAME:   index_vector_dim = #vhlo.integer_v1<2 : i64>,
-   // CHECK-SAME:   indices_are_sorted = #vhlo.bool_v1<false>,
-@@ -576,15 +619,16 @@
-   // CHECK-SAME:   sym_name = #vhlo.string_v1<"default_func">,
-   // CHECK-SAME:   sym_visibility = #vhlo.string_v1<"">
-   // CHECK-SAME: }> ({
--  // CHECK-NEXT:   ^[[BB:bb.*]](%arg0: !vhlo.tensor_v1<!vhlo.f32_v1>):
--  // CHECK-NEXT:     "vhlo.return_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-+  // CHECK-NEXT:   ^[[BB:bb.*]](%[[ARG0:.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
-+  // CHECK-NEXT:     "vhlo.return_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-   // CHECK-NEXT: }) : () -> ()
-   func.return %arg0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "dynamic_gather"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @dynamic_gather(%arg0 : tensor<2x4x9xf32>, %arg1 : tensor<1x5x2xi32>) -> tensor<1x5x1xf32> {
--  //      CHECK: "vhlo.gather_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.gather_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   collapsed_slice_dims = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-   // CHECK-SAME:   index_vector_dim = #vhlo.integer_v1<2 : i64>,
-   // CHECK-SAME:   indices_are_sorted = #vhlo.bool_v1<false>,
-@@ -605,8 +649,9 @@
- }
- 
- // CHECK-LABEL: "default_infeed"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_infeed(%arg0: !stablehlo.token) -> (tensor<f32>, !stablehlo.token) {
--  //               CHECK: "vhlo.infeed_v1"(%arg0) <{
-+  //               CHECK: "vhlo.infeed_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   infeed_config = #vhlo.string_v1<"">,
-   // CHECK-SAME{LITERAL}:   layout = #vhlo.array_v1<[]>
-   //          CHECK-SAME: }> : (!vhlo.token_v1) -> (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.token_v1)
-@@ -615,8 +660,9 @@
- }
- 
- // CHECK-LABEL: "default_outfeed"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @default_outfeed(%arg0: tensor<f32>, %arg1: !stablehlo.token) -> !stablehlo.token {
--  //      CHECK: "vhlo.outfeed_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.outfeed_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   outfeed_config = #vhlo.string_v1<"">
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.token_v1) -> !vhlo.token_v1
-   %0 = "stablehlo.outfeed"(%arg0, %arg1) : (tensor<f32>, !stablehlo.token) -> !stablehlo.token
-@@ -624,8 +670,9 @@
- }
- 
- // CHECK-LABEL: "default_recv"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_recv(%arg0: !stablehlo.token) -> (tensor<f32>, !stablehlo.token) {
--  //      CHECK: "vhlo.recv_v1"(%arg0) <{
-+  //      CHECK: "vhlo.recv_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME:   channel_type = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   is_host_transfer = #vhlo.bool_v1<false>
-@@ -637,8 +684,9 @@
- }
- 
- // CHECK-LABEL: "default_send"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @default_send(%arg0: tensor<f32>, %arg1: !stablehlo.token) -> !stablehlo.token {
--  //      CHECK: "vhlo.send_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.send_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME:   channel_type = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   is_host_transfer = #vhlo.bool_v1<false>
-@@ -650,8 +698,9 @@
- }
- 
- // CHECK-LABEL: "default_reduce_scatter"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_reduce_scatter(%arg0: tensor<16xf32>) -> tensor<16xf32> {
--  //               CHECK: "vhlo.reduce_scatter_v1"(%arg0) <{
-+  //               CHECK: "vhlo.reduce_scatter_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0], [1]]> : tensor<2x1xi64>>,
-   //          CHECK-SAME:   scatter_dimension = #vhlo.integer_v1<0 : i64>
-@@ -673,8 +722,9 @@
- }
- 
- // CHECK-LABEL: "default_reduce_window"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @default_reduce_window(%arg0: tensor<2x17x31x7xf32>, %arg1: tensor<f32>) -> tensor<2x16x30x7xf32> {
--  //               CHECK: "vhlo.reduce_window_v1"(%arg0, %arg1)  <{
-+  //               CHECK: "vhlo.reduce_window_v1"(%[[ARG0]], %[[ARG1]])  <{
-   //          CHECK-SAME:   base_dilations = #vhlo.tensor_v1<dense<1> : tensor<4xi64>>,
-   // CHECK-SAME{LITERAL}:   padding = #vhlo.tensor_v1<dense<0> : tensor<4x2xi64>>,
-   //          CHECK-SAME:   window_dilations = #vhlo.tensor_v1<dense<1> : tensor<4xi64>>,
-@@ -696,8 +746,9 @@
- }
- 
- // CHECK-LABEL: "default_scatter"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @default_scatter(%arg0: tensor<200x100x300xf32>, %arg1: tensor<10x2xi32>, %arg2: tensor<10x300xf32>) -> tensor<200x100x300xf32> {
--  //      CHECK: "vhlo.scatter_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.scatter_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   index_vector_dim = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   indices_are_sorted = #vhlo.bool_v1<false>,
-   // CHECK-SAME:   inserted_window_dims = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-@@ -725,8 +776,9 @@
- }
- 
- // CHECK-LABEL: "default_select_and_scatter"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @default_select_and_scatter(%arg0: tensor<10x24x24x64xf32>, %arg1: tensor<10x23x23x64xf32>, %arg2: tensor<f32>) -> tensor<10x24x24x64xf32> {
--  //      CHECK: "vhlo.select_and_scatter_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.select_and_scatter_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   padding = #vhlo.tensor_v1<dense<0> : tensor<4x2xi64>>,
-   // CHECK-SAME:   window_dimensions = #vhlo.tensor_v1<dense<[1, 2, 2, 1]> : tensor<4xi64>>,
-   // CHECK-SAME:   window_strides = #vhlo.tensor_v1<dense<1> : tensor<4xi64>>
-@@ -754,8 +806,9 @@
- }
- 
- // CHECK-LABEL: "default_sort"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_sort(%arg0: tensor<16xf32>) -> tensor<16xf32> {
--  //      CHECK: "vhlo.sort_v1"(%arg0) <{
-+  //      CHECK: "vhlo.sort_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   dimension = #vhlo.integer_v1<-1 : i64>
-   // CHECK-SAME:   is_stable = #vhlo.bool_v1<false>
-   // CHECK-SAME: }> ({
-@@ -774,29 +827,33 @@
- // ============ OPS ============
- 
- // CHECK-LABEL: "op_abs"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_abs(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.abs_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.abs_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.abs"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_add"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_add(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_after_all"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_after_all(%arg0: !stablehlo.token) -> !stablehlo.token {
--  // CHECK: "vhlo.after_all_v1"(%arg0) : (!vhlo.token_v1) -> !vhlo.token_v1
-+  // CHECK: "vhlo.after_all_v1"(%[[ARG0]]) : (!vhlo.token_v1) -> !vhlo.token_v1
-   %0 = "stablehlo.after_all"(%arg0) : (!stablehlo.token) -> !stablehlo.token
-   func.return %0 : !stablehlo.token
- }
- 
- // CHECK-LABEL: "op_all_gather"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_all_gather(%arg0: tensor<16x8xf32>) -> tensor<16x16xf32> {
--  //               CHECK: "vhlo.all_gather_v1"(%arg0) <{
-+  //               CHECK: "vhlo.all_gather_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   all_gather_dim = #vhlo.integer_v1<1 : i64>
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0], [1]]> : tensor<2x1xi64>>,
-@@ -812,8 +869,9 @@
- }
- 
- // CHECK-LABEL: "op_all_reduce"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_all_reduce(%arg0: tensor<f32>) -> tensor<f32> {
--  //               CHECK: "vhlo.all_reduce_v1"(%arg0) <{
-+  //               CHECK: "vhlo.all_reduce_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0], [1]]> : tensor<2x1xi64>>,
-   //          CHECK-SAME:   use_global_device_ids = #vhlo.bool_v1<true>
-@@ -835,8 +893,9 @@
- }
- 
- // CHECK-LABEL: "op_all_to_all"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_all_to_all(%arg0: tensor<4x16xf32>) -> tensor<16x4xf32> {
--  //               CHECK: "vhlo.all_to_all_v1"(%arg0) <{
-+  //               CHECK: "vhlo.all_to_all_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<1 : i64>,
-   //          CHECK-SAME:   concat_dimension = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0, 1, 2, 3]]> : tensor<1x4xi64>>,
-@@ -854,22 +913,25 @@
- }
- 
- // CHECK-LABEL: "op_and"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_and(%arg0: tensor<i1>, %arg1: tensor<i1>) -> tensor<i1> {
--  // CHECK: "vhlo.and_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-+  // CHECK: "vhlo.and_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-   %0 = "stablehlo.and"(%arg0, %arg1) : (tensor<i1>, tensor<i1>) -> tensor<i1>
-   func.return %0 : tensor<i1>
- }
- 
- // CHECK-LABEL: "op_atan2"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_atan2(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.atan2_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.atan2_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.atan2"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_batch_norm_grad"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}}, %[[ARG3:.*]]: {{.*}}, %[[ARG4:.*]]: {{.*}})
- func.func @op_batch_norm_grad(%arg0: tensor<16x16x16x16xf32>, %arg1: tensor<16xf32>, %arg2: tensor<16xf32>, %arg3: tensor<16xf32>, %arg4: tensor<16x16x16x16xf32>) -> (tensor<16x16x16x16xf32>, tensor<16xf32>, tensor<16xf32>) {
--  //      CHECK: "vhlo.batch_norm_grad_v1"(%arg0, %arg1, %arg2, %arg3, %arg4) <{
-+  //      CHECK: "vhlo.batch_norm_grad_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]], %[[ARG3]], %[[ARG4]]) <{
-   // CHECK-SAME:   epsilon = #vhlo.float_v1<1.000000e-03 : !vhlo.f32_v1>,
-   // CHECK-SAME:   feature_index = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<16x16x16x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x16x16x16x!vhlo.f32_v1>) -> (!vhlo.tensor_v1<16x16x16x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>)
-@@ -881,8 +943,9 @@
- }
- 
- // CHECK-LABEL: "op_batch_norm_inference"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}}, %[[ARG3:.*]]: {{.*}}, %[[ARG4:.*]]: {{.*}})
- func.func @op_batch_norm_inference(%arg0: tensor<16x16x16x16xf32>, %arg1: tensor<16xf32>, %arg2: tensor<16xf32>, %arg3: tensor<16xf32>, %arg4: tensor<16xf32>) -> tensor<16x16x16x16xf32> {
--  //      CHECK: "vhlo.batch_norm_inference_v1"(%arg0, %arg1, %arg2, %arg3, %arg4) <{
-+  //      CHECK: "vhlo.batch_norm_inference_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]], %[[ARG3]], %[[ARG4]]) <{
-   // CHECK-SAME:   epsilon = #vhlo.float_v1<1.000000e-03 : !vhlo.f32_v1>,
-   // CHECK-SAME:   feature_index = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<16x16x16x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x16x16x16x!vhlo.f32_v1>
-@@ -894,8 +957,9 @@
- }
- 
- // CHECK-LABEL: "op_batch_norm_training"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_batch_norm_training(%arg0: tensor<16x16x16x16xf32>, %arg1: tensor<16xf32>, %arg2: tensor<16xf32>) -> (tensor<16x16x16x16xf32>, tensor<16xf32>, tensor<16xf32>) {
--  //      CHECK: "vhlo.batch_norm_training_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.batch_norm_training_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   epsilon = #vhlo.float_v1<1.000000e-03 : !vhlo.f32_v1>,
-   // CHECK-SAME:   feature_index = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<16x16x16x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>) -> (!vhlo.tensor_v1<16x16x16x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>)
-@@ -907,15 +971,17 @@
- }
- 
- // CHECK-LABEL: "op_bitcast_convert"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_bitcast_convert(%arg0: tensor<i32>) -> tensor<f32> {
--  // CHECK: "vhlo.bitcast_convert_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.bitcast_convert_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.bitcast_convert"(%arg0) : (tensor<i32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_broadcast_in_dim"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_broadcast_in_dim(%arg0: tensor<16xf32>) -> tensor<16x16xf32> {
--  //      CHECK: "vhlo.broadcast_in_dim_v1"(%arg0) <{
-+  //      CHECK: "vhlo.broadcast_in_dim_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   broadcast_dimensions = #vhlo.tensor_v1<dense<1> : tensor<1xi64>>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x16x!vhlo.f32_v1>
-   %0 = "stablehlo.broadcast_in_dim"(%arg0) {
-@@ -925,8 +991,9 @@
- }
- 
- // CHECK-LABEL: "op_broadcast"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_broadcast(%arg0: tensor<16xf32>) -> tensor<16x16xf32> {
--  //      CHECK: "vhlo.broadcast_v1"(%arg0) <{
-+  //      CHECK: "vhlo.broadcast_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   broadcast_sizes = #vhlo.tensor_v1<dense<16> : tensor<1xi64>>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x16x!vhlo.f32_v1>
-   %0 = "stablehlo.broadcast"(%arg0) {
-@@ -936,9 +1003,10 @@
- }
- 
- // CHECK-LABEL: "op_case"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_case(%arg0: tensor<i32>, %arg1: tensor<f32>) -> tensor<f32> {
--  //      CHECK: "vhlo.case_v1"(%arg0) ({
--  // CHECK-NEXT:   "vhlo.return_v1"(%arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-+  //      CHECK: "vhlo.case_v1"(%[[ARG0]]) ({
-+  // CHECK-NEXT:   "vhlo.return_v1"(%[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-   // CHECK-NEXT: }) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.case"(%arg0) ({
-     "stablehlo.return"(%arg1) : (tensor<f32>) -> ()
-@@ -947,22 +1015,25 @@
- }
- 
- // CHECK-LABEL: "op_cbrt"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_cbrt(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.cbrt_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.cbrt_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.cbrt"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_ceil"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_ceil(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.ceil_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.ceil_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.ceil"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_cholesky"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_cholesky(%arg0: tensor<1x16x16xf32>) -> tensor<1x16x16xf32> {
--  //      CHECK: "vhlo.cholesky_v1"(%arg0) <{
-+  //      CHECK: "vhlo.cholesky_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   lower = #vhlo.bool_v1<true>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<1x16x16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<1x16x16x!vhlo.f32_v1>
-   %0 = "stablehlo.cholesky"(%arg0) {
-@@ -972,22 +1043,25 @@
- }
- 
- // CHECK-LABEL: "op_clamp"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_clamp(%arg0: tensor<f32>, %arg1: tensor<f32>, %arg2: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.clamp_v1"(%arg0, %arg1, %arg2) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.clamp_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.clamp"(%arg0, %arg1, %arg2) : (tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_count_leading_zeros"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_count_leading_zeros(%arg0: tensor<i32>) -> tensor<i32> {
--  // CHECK: "vhlo.count_leading_zeros_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-+  // CHECK: "vhlo.count_leading_zeros_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-   %0 = "stablehlo.count_leading_zeros"(%arg0) : (tensor<i32>) -> tensor<i32>
-   func.return %0 : tensor<i32>
- }
- 
- // CHECK-LABEL: "op_collective_permute"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_collective_permute(%arg0: tensor<16x8xf32>) -> tensor<16x8xf32> {
--  //               CHECK: "vhlo.collective_permute_v1"(%arg0) <{
-+  //               CHECK: "vhlo.collective_permute_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME{LITERAL}:   source_target_pairs = #vhlo.tensor_v1<dense<[[0, 1], [1, 2], [2, 3]]> : tensor<3x2xi64>>
-   //          CHECK-SAME: }> : (!vhlo.tensor_v1<16x8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x8x!vhlo.f32_v1>
-@@ -999,8 +1073,9 @@
- }
- 
- // CHECK-LABEL: "op_compare"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_compare(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
--  //      CHECK: "vhlo.compare_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.compare_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   compare_type = #vhlo<comparison_type_v1 TOTALORDER>,
-   // CHECK-SAME:   comparison_direction = #vhlo<comparison_direction_v1 EQ>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-@@ -1012,15 +1087,17 @@
- }
- 
- // CHECK-LABEL: "op_complex"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_complex(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<complex<f32>> {
--  // CHECK: "vhlo.complex_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>
-+  // CHECK: "vhlo.complex_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>
-   %0 = "stablehlo.complex"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<complex<f32>>
-   func.return %0 : tensor<complex<f32>>
- }
- 
- // CHECK-LABEL: "op_concatenate"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_concatenate(%arg0: tensor<8xf32>, %arg1: tensor<8xf32>) -> tensor<16xf32> {
--  //      CHECK: "vhlo.concatenate_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.concatenate_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   dimension = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<8x!vhlo.f32_v1>, !vhlo.tensor_v1<8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x!vhlo.f32_v1>
-   %0 = "stablehlo.concatenate"(%arg0, %arg1) {
-@@ -1030,6 +1107,7 @@
- }
- 
- // CHECK-LABEL: "op_constant"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_constant(%arg0: tensor<f32>) -> tensor<f32> {
-   //      CHECK: "vhlo.constant_v1"() <{
-   // CHECK-SAME:   value = #vhlo.tensor_v1<dense<0.000000e+00> : tensor<f32>>
-@@ -1041,15 +1119,17 @@
- }
- 
- // CHECK-LABEL: "op_convert"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_convert(%arg0: tensor<i32>) -> tensor<f32> {
--  // CHECK: "vhlo.convert_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.convert_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.convert"(%arg0) : (tensor<i32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_convolution"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_convolution(%arg0: tensor<1x8x8x207xf32>, %arg1: tensor<3x3x207x16xf32>) -> tensor<1x7x7x16xf32> {
--  //      CHECK: "vhlo.convolution_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.convolution_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   batch_group_count = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   feature_group_count = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   input_batch_dimension = #vhlo.integer_v1<0 : i64>,
-@@ -1083,8 +1163,9 @@
- }
- 
- // CHECK-LABEL: "op_cosine"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_cosine(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.cosine_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.cosine_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.cosine"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
-@@ -1097,8 +1178,9 @@
- }
- 
- // CHECK-LABEL: "op_cross_replica_sum"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_cross_replica_sum(%arg0: tensor<f32>) -> tensor<f32> {
--  //               CHECK: "vhlo.cross-replica-sum_v1"(%arg0) <{
-+  //               CHECK: "vhlo.cross-replica-sum_v1"(%[[ARG0]]) <{
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0], [1]]> : tensor<2x1xi64>>
-   //          CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.cross-replica-sum"(%arg0) {
-@@ -1108,8 +1190,9 @@
- }
- 
- // CHECK-LABEL: "op_custom_call"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_custom_call(%arg0: tensor<f32>) -> tensor<f32> {
--  //      CHECK: "vhlo.custom_call_v1"(%arg0) <{
-+  //      CHECK: "vhlo.custom_call_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   api_version = #vhlo<api_version_v1 API_VERSION_STATUS_RETURNING>,
-   // CHECK-SAME:   backend_config = #vhlo.string_v1<"\08\03\1A\02">,
-   // CHECK-SAME:   call_target_name = #vhlo.string_v1<"foo">,
-@@ -1140,15 +1223,17 @@
- }
- 
- // CHECK-LABEL: "op_divide"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_divide(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.divide_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.divide_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.divide"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_dot_general"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_dot_general(%arg0: tensor<8x8x16xf32>, %arg1: tensor<8x16x8xf32>) -> tensor<8x8x8xf32> {
--  //      CHECK: "vhlo.dot_general_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.dot_general_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   lhs_batching_dimensions = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>,
-   // CHECK-SAME:   lhs_contracting_dimensions = #vhlo.tensor_v1<dense<2> : tensor<1xi64>>,
-   // CHECK-SAME:   precision_config = #vhlo.array_v1<[#vhlo<precision_v1 HIGHEST>, #vhlo<precision_v1 HIGHEST>]>,
-@@ -1168,8 +1253,9 @@
- }
- 
- // CHECK-LABEL: "op_dot"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_dot(%arg0: tensor<8x16xf32>, %arg1: tensor<16x8xf32>) -> tensor<8x8xf32> {
--  //      CHECK: "vhlo.dot_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.dot_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   precision_config = #vhlo.array_v1<[#vhlo<precision_v1 HIGHEST>, #vhlo<precision_v1 HIGHEST>]>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<8x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<8x8x!vhlo.f32_v1>
-   %0 = "stablehlo.dot"(%arg0, %arg1) {
-@@ -1179,8 +1265,9 @@
- }
- 
- // CHECK-LABEL: "op_dynamic_broadcast_in_dim"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_dynamic_broadcast_in_dim(%arg0: tensor<?x?xf32>, %arg1: tensor<2xindex>) -> tensor<?x?xf32> {
--  //      CHECK: "vhlo.dynamic_broadcast_in_dim_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.dynamic_broadcast_in_dim_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   broadcast_dimensions = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-   // CHECK-SAME:   known_expanding_dimensions = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>,
-   // CHECK-SAME:   known_nonexpanding_dimensions = #vhlo.tensor_v1<dense<1> : tensor<1xi64>>
-@@ -1194,8 +1281,9 @@
- }
- 
- // CHECK-LABEL: "op_dynamic_conv"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_dynamic_conv(%arg0: tensor<1x8x8x207xf32>, %arg1: tensor<3x3x207x16xf32>, %arg2: tensor<4xi32>) -> tensor<1x?x?x16xf32> {
--  //      CHECK: "vhlo.dynamic_conv_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.dynamic_conv_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   batch_group_count = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   feature_group_count = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   input_batch_dimension = #vhlo.integer_v1<0 : i64>,
-@@ -1229,8 +1317,9 @@
- }
- 
- // CHECK-LABEL: "op_dynamic_gather"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_dynamic_gather(%arg0 : tensor<2x4x9xf32>, %arg1 : tensor<1x5x2xi32>, %arg2 : tensor<3xi32>) -> tensor<1x5x8xf32> {
--  //      CHECK: "vhlo.dynamic_gather_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.dynamic_gather_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   collapsed_slice_dims = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-   // CHECK-SAME:   index_vector_dim = #vhlo.integer_v1<2 : i64>,
-   // CHECK-SAME:   indices_are_sorted = #vhlo.bool_v1<true>,
-@@ -1250,8 +1339,9 @@
- }
- 
- // CHECK-LABEL: "op_dynamic_iota"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_dynamic_iota(%arg0: tensor<1xindex>) -> tensor<?xf32> {
--  //      CHECK: "vhlo.dynamic_iota_v1"(%arg0) <{
-+  //      CHECK: "vhlo.dynamic_iota_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   iota_dimension = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<1x!vhlo.index_v1>) -> !vhlo.tensor_v1<?x!vhlo.f32_v1>
-   %0 = "stablehlo.dynamic_iota"(%arg0) {
-@@ -1261,22 +1351,25 @@
- }
- 
- // CHECK-LABEL: "op_dynamic_pad"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}}, %[[ARG3:.*]]: {{.*}}, %[[ARG4:.*]]: {{.*}})
- func.func @op_dynamic_pad(%arg0: tensor<?xf32>, %arg1: tensor<f32>, %arg2: tensor<1xindex>, %arg3: tensor<1xindex>, %arg4: tensor<1xindex>) -> tensor<?xf32> {
--  // CHECK: "vhlo.dynamic_pad_v1"(%arg0, %arg1, %arg2, %arg3, %arg4) : (!vhlo.tensor_v1<?x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>) -> !vhlo.tensor_v1<?x!vhlo.f32_v1>
-+  // CHECK: "vhlo.dynamic_pad_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]], %[[ARG3]], %[[ARG4]]) : (!vhlo.tensor_v1<?x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>) -> !vhlo.tensor_v1<?x!vhlo.f32_v1>
-   %0 = "stablehlo.dynamic_pad"(%arg0, %arg1, %arg2, %arg3, %arg4) : (tensor<?xf32>, tensor<f32>, tensor<1xindex>, tensor<1xindex>, tensor<1xindex>) -> tensor<?xf32>
-   func.return %0 : tensor<?xf32>
- }
- 
- // CHECK-LABEL: "op_dynamic_reshape"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_dynamic_reshape(%arg0: tensor<16xf32>, %arg1: tensor<2xindex>) -> tensor<?x?xf32> {
--  // CHECK: "vhlo.dynamic_reshape_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<2x!vhlo.index_v1>) -> !vhlo.tensor_v1<?x?x!vhlo.f32_v1>
-+  // CHECK: "vhlo.dynamic_reshape_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<2x!vhlo.index_v1>) -> !vhlo.tensor_v1<?x?x!vhlo.f32_v1>
-   %0 = "stablehlo.dynamic_reshape"(%arg0, %arg1) : (tensor<16xf32>, tensor<2xindex>) -> tensor<?x?xf32>
-   func.return %0 : tensor<?x?xf32>
- }
- 
- // CHECK-LABEL: "op_dynamic_slice"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_dynamic_slice(%arg0: tensor<16xf32>, %arg1: tensor<i64>) -> tensor<4xf32> {
--  //      CHECK: "vhlo.dynamic_slice_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.dynamic_slice_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   slice_sizes = #vhlo.tensor_v1<dense<4> : tensor<1xi64>>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.i64_v1>) -> !vhlo.tensor_v1<4x!vhlo.f32_v1>
-   %0 = "stablehlo.dynamic_slice"(%arg0, %arg1) {
-@@ -1286,15 +1379,17 @@
- }
- 
- // CHECK-LABEL: "op_dynamic_update_slice"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_dynamic_update_slice(%arg0: tensor<16xf32>, %arg1: tensor<4xf32>, %arg2: tensor<i64>) -> tensor<16xf32> {
--  // CHECK: "vhlo.dynamic_update_slice_v1"(%arg0, %arg1, %arg2) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<4x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.i64_v1>) -> !vhlo.tensor_v1<16x!vhlo.f32_v1>
-+  // CHECK: "vhlo.dynamic_update_slice_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<4x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.i64_v1>) -> !vhlo.tensor_v1<16x!vhlo.f32_v1>
-   %0 = "stablehlo.dynamic_update_slice"(%arg0, %arg1, %arg2) : (tensor<16xf32>, tensor<4xf32>, tensor<i64>) -> tensor<16xf32>
-   func.return %0 : tensor<16xf32>
- }
- 
- // CHECK-LABEL: "op_einsum"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_einsum(%arg0: tensor<8x16xf32>, %arg1: tensor<16x8xf32>) -> tensor<8x8xf32> {
--  //      CHECK: "vhlo.einsum_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.einsum_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   einsum_config = #vhlo.string_v1<"ab,bc->ac">
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<8x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<8x8x!vhlo.f32_v1>
-   %0 = "stablehlo.einsum"(%arg0, %arg1) {
-@@ -1304,22 +1399,25 @@
- }
- 
- // CHECK-LABEL: "op_exponential_minus_one"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_exponential_minus_one(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.exponential_minus_one_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.exponential_minus_one_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.exponential_minus_one"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_exponential"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_exponential(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.exponential_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.exponential_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.exponential"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_fft"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_fft(%arg0: tensor<16xcomplex<f32>>) -> tensor<16xcomplex<f32>> {
--  //      CHECK: "vhlo.fft_v1"(%arg0) <{
-+  //      CHECK: "vhlo.fft_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   fft_length = #vhlo.tensor_v1<dense<16> : tensor<1xi64>>,
-   // CHECK-SAME:   fft_type = #vhlo<fft_type_v1 FFT>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<16x!vhlo.complex_v1<!vhlo.f32_v1>>) -> !vhlo.tensor_v1<16x!vhlo.complex_v1<!vhlo.f32_v1>>
-@@ -1331,8 +1429,9 @@
- }
- 
- // CHECK-LABEL: "op_floor"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_floor(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.floor_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.floor_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.floor"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
-@@ -1345,16 +1444,17 @@
-   // CHECK-SAME:   sym_name = #vhlo.string_v1<"op_func">,
-   // CHECK-SAME:   sym_visibility = #vhlo.string_v1<"private">
-   // CHECK-SAME: }> ({
--  // CHECK-NEXT:   ^[[BB:bb.*]](%arg0: !vhlo.tensor_v1<!vhlo.f32_v1>):
--  // CHECK-NEXT:     "vhlo.return_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-+  // CHECK-NEXT:   ^[[BB:bb.*]](%[[ARG0:.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
-+  // CHECK-NEXT:     "vhlo.return_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-   // CHECK-NEXT: }) : () -> ()
- 
-   func.return %arg0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_gather"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_gather(%arg0 : tensor<2x4x9xf32>, %arg1 : tensor<1x5x2xi32>) -> tensor<1x5x1xf32> {
--  //      CHECK: "vhlo.gather_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.gather_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   collapsed_slice_dims = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-   // CHECK-SAME:   index_vector_dim = #vhlo.integer_v1<2 : i64>,
-   // CHECK-SAME:   indices_are_sorted = #vhlo.bool_v1<true>,
-@@ -1376,8 +1476,9 @@
- }
- 
- // CHECK-LABEL: "op_get_dimension_size"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_get_dimension_size(%arg0: tensor<?xf32>) -> tensor<i32> {
--  //      CHECK: "vhlo.get_dimension_size_v1"(%arg0) <{
-+  //      CHECK: "vhlo.get_dimension_size_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   dimension = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<?x!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-   %0 = "stablehlo.get_dimension_size"(%arg0) {
-@@ -1387,8 +1488,9 @@
- }
- 
- // CHECK-LABEL: "op_get_tuple_element"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_get_tuple_element(%arg0: tuple<tensor<f32>, tensor<i32>>) -> tensor<f32> {
--  //      CHECK: "vhlo.get_tuple_element_v1"(%arg0) <{
-+  //      CHECK: "vhlo.get_tuple_element_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   index = #vhlo.integer_v1<0 : i32>
-   // CHECK-SAME: }> : (!vhlo.tuple_v1<!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.get_tuple_element"(%arg0) {
-@@ -1398,11 +1500,12 @@
- }
- 
- // CHECK-LABEL: "op_if"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_if(%arg0: tensor<i1>, %arg1: tensor<f32>, %arg2: tensor<f32>) -> tensor<f32> {
--  //      CHECK: "vhlo.if_v1"(%arg0) ({
--  // CHECK-NEXT:   "vhlo.return_v1"(%arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-+  //      CHECK: "vhlo.if_v1"(%[[ARG0]]) ({
-+  // CHECK-NEXT:   "vhlo.return_v1"(%[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-   // CHECK-NEXT: }, {
--  // CHECK-NEXT:   "vhlo.return_v1"(%arg2) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-+  // CHECK-NEXT:   "vhlo.return_v1"(%[[ARG2]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-   // CHECK-NEXT: }) : (!vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.if"(%arg0) ({
-     "stablehlo.return"(%arg1) : (tensor<f32>) -> ()
-@@ -1413,15 +1516,17 @@
- }
- 
- // CHECK-LABEL: "op_imag"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_imag(%arg0: tensor<complex<f32>>) -> tensor<f32> {
--  // CHECK: "vhlo.imag_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.imag_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.imag"(%arg0) : (tensor<complex<f32>>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_infeed"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_infeed(%arg0: !stablehlo.token) -> (tensor<f32>, !stablehlo.token) {
--  //               CHECK: "vhlo.infeed_v1"(%arg0) <{
-+  //               CHECK: "vhlo.infeed_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   infeed_config = #vhlo.string_v1<"foo">,
-   // CHECK-SAME{LITERAL}:   layout = #vhlo.array_v1<[#vhlo.array_v1<[]>]>
-   //          CHECK-SAME: }> : (!vhlo.token_v1) -> (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.token_v1)
-@@ -1444,36 +1549,41 @@
- }
- 
- // CHECK-LABEL: "op_is_finite"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_is_finite(%arg0: tensor<f32>) -> tensor<i1> {
--  // CHECK: "vhlo.is_finite_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-+  // CHECK: "vhlo.is_finite_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-   %0 = "stablehlo.is_finite"(%arg0) : (tensor<f32>) -> tensor<i1>
-   func.return %0 : tensor<i1>
- }
- 
- // CHECK-LABEL: "op_log"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_log(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.log_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.log_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.log"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_log_plus_one"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_log_plus_one(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.log_plus_one_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.log_plus_one_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.log_plus_one"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_logistic"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_logistic(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.logistic_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.logistic_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.logistic"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_map"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_map(%arg0: tensor<16xf32>) -> tensor<16xf32> {
--  //      CHECK: "vhlo.map_v1"(%arg0) <{
-+  //      CHECK: "vhlo.map_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   dimensions = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>
-   // CHECK-SAME: }> ({
-   // CHECK-NEXT:   ^[[BB:bb.*]](%[[ARG1:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
-@@ -1491,57 +1601,65 @@
- }
- 
- // CHECK-LABEL: "op_maximum"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_maximum(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.maximum_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.maximum_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.maximum"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_minimum"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_minimum(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.minimum_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.minimum_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.minimum"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_multiply"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_multiply(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.multiply_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.multiply_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.multiply"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_negate"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_negate(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.negate_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.negate_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.negate"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_not"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_not(%arg0: tensor<i1>) -> tensor<i1> {
--  // CHECK: "vhlo.not_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-+  // CHECK: "vhlo.not_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-   %0 = "stablehlo.not"(%arg0) : (tensor<i1>) -> tensor<i1>
-   func.return %0 : tensor<i1>
- }
- 
- // CHECK-LABEL: "op_optimization_barrier"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_optimization_barrier(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.optimization_barrier_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.optimization_barrier_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.optimization_barrier"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_or"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_or(%arg0: tensor<i1>, %arg1: tensor<i1>) -> tensor<i1> {
--  // CHECK: "vhlo.or_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-+  // CHECK: "vhlo.or_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-   %0 = "stablehlo.or"(%arg0, %arg1) : (tensor<i1>, tensor<i1>) -> tensor<i1>
-   func.return %0 : tensor<i1>
- }
- 
- // CHECK-LABEL: "op_outfeed"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_outfeed(%arg0: tensor<f32>, %arg1: !stablehlo.token) -> !stablehlo.token {
--  //      CHECK: "vhlo.outfeed_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.outfeed_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   outfeed_config = #vhlo.string_v1<"foo">
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.token_v1) -> !vhlo.token_v1
-   %0 = "stablehlo.outfeed"(%arg0, %arg1) {
-@@ -1551,8 +1669,9 @@
- }
- 
- // CHECK-LABEL: "op_pad"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_pad(%arg0: tensor<8xf32>, %arg1: tensor<f32>) -> tensor<16xf32> {
--  //      CHECK: "vhlo.pad_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.pad_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   edge_padding_high = #vhlo.tensor_v1<dense<4> : tensor<1xi64>>,
-   // CHECK-SAME:   edge_padding_low = #vhlo.tensor_v1<dense<4> : tensor<1xi64>>,
-   // CHECK-SAME:   interior_padding = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>
-@@ -1566,36 +1685,41 @@
- }
- 
- // CHECK-LABEL: "op_popcnt"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_popcnt(%arg0: tensor<i32>) -> tensor<i32> {
--  // CHECK: "vhlo.popcnt_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-+  // CHECK: "vhlo.popcnt_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-   %0 = "stablehlo.popcnt"(%arg0) : (tensor<i32>) -> tensor<i32>
-   func.return %0 : tensor<i32>
- }
- 
- // CHECK-LABEL: "op_power"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_power(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.power_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.power_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.power"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_real_dynamic_slice"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}}, %[[ARG3:.*]]: {{.*}})
- func.func @op_real_dynamic_slice(%arg0: tensor<?xf32>, %arg1: tensor<1xindex>, %arg2: tensor<1xindex>, %arg3: tensor<1xindex>) -> tensor<?xf32> {
--  // CHECK: "vhlo.real_dynamic_slice_v1"(%arg0, %arg1, %arg2, %arg3) : (!vhlo.tensor_v1<?x!vhlo.f32_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>) -> !vhlo.tensor_v1<?x!vhlo.f32_v1>
-+  // CHECK: "vhlo.real_dynamic_slice_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]], %[[ARG3]]) : (!vhlo.tensor_v1<?x!vhlo.f32_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>) -> !vhlo.tensor_v1<?x!vhlo.f32_v1>
-   %0 = "stablehlo.real_dynamic_slice"(%arg0, %arg1, %arg2, %arg3) : (tensor<?xf32>, tensor<1xindex>, tensor<1xindex>, tensor<1xindex>) -> tensor<?xf32>
-   func.return %0 : tensor<?xf32>
- }
- 
- // CHECK-LABEL: "op_real"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_real(%arg0: tensor<complex<f32>>) -> tensor<f32> {
--  // CHECK: "vhlo.real_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.real_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.real"(%arg0) : (tensor<complex<f32>>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_recv"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_recv(%arg0: !stablehlo.token) -> (tensor<f32>, !stablehlo.token) {
--  //      CHECK: "vhlo.recv_v1"(%arg0) <{
-+  //      CHECK: "vhlo.recv_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME:   channel_type = #vhlo.integer_v1<3 : i64>,
-   // CHECK-SAME:   is_host_transfer = #vhlo.bool_v1<true>
-@@ -1608,8 +1732,9 @@
- }
- 
- // CHECK-LABEL: "op_reduce"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_reduce(%arg0: tensor<16xf32>, %arg1: tensor<f32>) -> tensor<f32> {
--  //  CHECK: "vhlo.reduce_v1"(%arg0, %arg1)
-+  //  CHECK: "vhlo.reduce_v1"(%[[ARG0]], %[[ARG1]])
-   //  CHECK:   ^[[BB:bb.*]](%[[ARG1:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>, %[[ARG2:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
-   //  CHECK:     "vhlo.return_v1"(%[[VAL1:.*]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-   //  CHECK: }) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-@@ -1624,8 +1749,9 @@
- }
- 
- // CHECK-LABEL: "op_reduce_precision"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_reduce_precision(%arg0: tensor<f32>) -> tensor<f32> {
--  //      CHECK: "vhlo.reduce_precision_v1"(%arg0) <{
-+  //      CHECK: "vhlo.reduce_precision_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   exponent_bits = #vhlo.integer_v1<8 : i32>
-   // CHECK-SAME:   mantissa_bits = #vhlo.integer_v1<10 : i32>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-@@ -1637,8 +1763,9 @@
- }
- 
- // CHECK-LABEL: "op_reduce_scatter"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_reduce_scatter(%arg0: tensor<16xf32>) -> tensor<16xf32> {
--  //               CHECK: "vhlo.reduce_scatter_v1"(%arg0) <{
-+  //               CHECK: "vhlo.reduce_scatter_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0], [1]]> : tensor<2x1xi64>>,
-   //          CHECK-SAME:   scatter_dimension = #vhlo.integer_v1<0 : i64>
-@@ -1662,8 +1789,9 @@
- }
- 
- // CHECK-LABEL: "op_reduce_window"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_reduce_window(%arg0: tensor<2x17x31x7xf32>, %arg1: tensor<f32>) -> tensor<2x9x16x7xf32> {
--  //               CHECK: "vhlo.reduce_window_v1"(%arg0, %arg1) <{
-+  //               CHECK: "vhlo.reduce_window_v1"(%[[ARG0]], %[[ARG1]]) <{
-   //          CHECK-SAME:   base_dilations = #vhlo.tensor_v1<dense<[1, 2, 2, 1]> : tensor<4xi64>>,
-   // CHECK-SAME{LITERAL}:   padding = #vhlo.tensor_v1<dense<[[0, 0], [2, 0], [0, 2], [0, 0]]> : tensor<4x2xi64>>,
-   //          CHECK-SAME:   window_dilations = #vhlo.tensor_v1<dense<[1, 2, 2, 1]> : tensor<4xi64>>,
-@@ -1689,8 +1817,9 @@
- }
- 
- // CHECK-LABEL: "op_remainder"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_remainder(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.remainder_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.remainder_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.remainder"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
-@@ -1710,16 +1839,18 @@
- }
- 
- // CHECK-LABEL: "op_reshape"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_reshape(%arg0: tensor<16xf32>) -> tensor<4x4xf32> {
--  // CHECK: "vhlo.reshape_v1"(%arg0) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<4x4x!vhlo.f32_v1>
-+  // CHECK: "vhlo.reshape_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<4x4x!vhlo.f32_v1>
-   %0 = "stablehlo.reshape"(%arg0) : (tensor<16xf32>) -> tensor<4x4xf32>
-   func.return %0 : tensor<4x4xf32>
- }
- 
- // CHECK-LABEL: "op_return"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_return(%arg0: tensor<i32>, %arg1: tensor<f32>) -> tensor<f32> {
--  //      CHECK: "vhlo.case_v1"(%arg0) ({
--  // CHECK-NEXT:   "vhlo.return_v1"(%arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-+  //      CHECK: "vhlo.case_v1"(%[[ARG0]]) ({
-+  // CHECK-NEXT:   "vhlo.return_v1"(%[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-   // CHECK-NEXT: }) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.case"(%arg0) ({
-     "stablehlo.return"(%arg1) : (tensor<f32>) -> ()
-@@ -1728,8 +1859,9 @@
- }
- 
- // CHECK-LABEL: "op_reverse"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_reverse(%arg0: tensor<16xf32>) -> tensor<16xf32> {
--  //      CHECK: "vhlo.reverse_v1"(%arg0) <{
-+  //      CHECK: "vhlo.reverse_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   dimensions = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x!vhlo.f32_v1>
-   %0 = "stablehlo.reverse"(%arg0) {
-@@ -1739,8 +1871,9 @@
- }
- 
- // CHECK-LABEL: "op_rng_bit_generator"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_rng_bit_generator(%arg0: tensor<f32>) -> (tensor<f32>, tensor<f32>) {
--  //      CHECK: "vhlo.rng_bit_generator_v1"(%arg0) <{
-+  //      CHECK: "vhlo.rng_bit_generator_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   rng_algorithm = #vhlo<rng_algorithm_v1 PHILOX>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>)
-   %0:2 = "stablehlo.rng_bit_generator"(%arg0) {
-@@ -1750,8 +1883,9 @@
- }
- 
- // CHECK-LABEL: "op_rng"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_rng(%arg0: tensor<f32>, %arg1: tensor<f32>, %arg2: tensor<0xindex>) -> tensor<f32> {
--  //      CHECK: "vhlo.rng_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.rng_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   rng_distribution = #vhlo<rng_distribution_v1 NORMAL>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<0x!vhlo.index_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.rng"(%arg0, %arg1, %arg2) {
-@@ -1761,29 +1895,33 @@
- }
- 
- // CHECK-LABEL: "op_round_nearest_afz"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_round_nearest_afz(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.round_nearest_afz_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.round_nearest_afz_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.round_nearest_afz"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_round_nearest_even"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_round_nearest_even(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.round_nearest_even_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.round_nearest_even_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.round_nearest_even"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_rsqrt"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_rsqrt(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.rsqrt_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.rsqrt_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.rsqrt"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_scatter"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_scatter(%arg0: tensor<200x100x300xf32>, %arg1: tensor<10x2xi32>, %arg2: tensor<10x300xf32>) -> tensor<200x100x300xf32> {
--  //      CHECK: "vhlo.scatter_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.scatter_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   index_vector_dim = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   indices_are_sorted = #vhlo.bool_v1<true>,
-   // CHECK-SAME:   inserted_window_dims = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-@@ -1813,8 +1951,9 @@
- }
- 
- // CHECK-LABEL: "op_select_and_scatter"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_select_and_scatter(%arg0: tensor<10x24x24x64xf32>, %arg1: tensor<12x13x13x66xf32>, %arg2: tensor<f32>) -> tensor<10x24x24x64xf32> {
--  //      CHECK: "vhlo.select_and_scatter_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.select_and_scatter_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   padding = #vhlo.tensor_v1<dense<1> : tensor<4x2xi64>>,
-   // CHECK-SAME:   window_dimensions = #vhlo.tensor_v1<dense<[1, 2, 2, 1]> : tensor<4xi64>>,
-   // CHECK-SAME:   window_strides = #vhlo.tensor_v1<dense<[1, 2, 2, 1]> : tensor<4xi64>>
-@@ -1844,15 +1983,17 @@
- }
- 
- // CHECK-LABEL: "op_select"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_select(%arg0: tensor<i1>, %arg1: tensor<f32>, %arg2: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.select_v1"(%arg0, %arg1, %arg2) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.select_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.select"(%arg0, %arg1, %arg2) : (tensor<i1>, tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_send"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_send(%arg0: tensor<f32>, %arg1: !stablehlo.token) -> !stablehlo.token {
--  //      CHECK: "vhlo.send_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.send_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME:   channel_type = #vhlo.integer_v1<2 : i64>,
-   // CHECK-SAME:   is_host_transfer = #vhlo.bool_v1<true>
-@@ -1865,8 +2006,9 @@
- }
- 
- // CHECK-LABEL: "op_set_dimension_size"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_set_dimension_size(%arg0: tensor<?xf32>, %arg1: tensor<i32>) -> tensor<16xf32> {
--  //      CHECK: "vhlo.set_dimension_size_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.set_dimension_size_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   dimension = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<?x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<16x!vhlo.f32_v1>
-   %0 = "stablehlo.set_dimension_size"(%arg0, %arg1) {
-@@ -1876,43 +2018,49 @@
- }
- 
- // CHECK-LABEL: "op_shift_left"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_shift_left(%arg0: tensor<i32>, %arg1: tensor<i32>) -> tensor<i32> {
--  // CHECK: "vhlo.shift_left_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.i32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-+  // CHECK: "vhlo.shift_left_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-   %0 = "stablehlo.shift_left"(%arg0, %arg1) : (tensor<i32>, tensor<i32>) -> tensor<i32>
-   func.return %0 : tensor<i32>
- }
- 
- // CHECK-LABEL: "op_shift_right_arithmetic"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_shift_right_arithmetic(%arg0: tensor<i32>, %arg1: tensor<i32>) -> tensor<i32> {
--  // CHECK: "vhlo.shift_right_arithmetic_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.i32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-+  // CHECK: "vhlo.shift_right_arithmetic_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-   %0 = "stablehlo.shift_right_arithmetic"(%arg0, %arg1) : (tensor<i32>, tensor<i32>) -> tensor<i32>
-   func.return %0 : tensor<i32>
- }
- 
- // CHECK-LABEL: "op_shift_right_logical"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_shift_right_logical(%arg0: tensor<i32>, %arg1: tensor<i32>) -> tensor<i32> {
--  // CHECK: "vhlo.shift_right_logical_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.i32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-+  // CHECK: "vhlo.shift_right_logical_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-   %0 = "stablehlo.shift_right_logical"(%arg0, %arg1) : (tensor<i32>, tensor<i32>) -> tensor<i32>
-   func.return %0 : tensor<i32>
- }
- 
- // CHECK-LABEL: "op_sign"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_sign(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.sign_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.sign_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.sign"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_sine"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_sine(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.sine_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.sine_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.sine"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_slice"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_slice(%arg0: tensor<16xf32>) -> tensor<4xf32> {
--  //      CHECK: "vhlo.slice_v1"(%arg0) <{
-+  //      CHECK: "vhlo.slice_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   limit_indices = #vhlo.tensor_v1<dense<4> : tensor<1xi64>>,
-   // CHECK-SAME:   start_indices = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>,
-   // CHECK-SAME:   strides = #vhlo.tensor_v1<dense<1> : tensor<1xi64>>
-@@ -1926,8 +2074,9 @@
- }
- 
- // CHECK-LABEL: "op_sort"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_sort(%arg0: tensor<16xf32>) -> tensor<16xf32> {
--  //      CHECK: "vhlo.sort_v1"(%arg0) <{
-+  //      CHECK: "vhlo.sort_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   dimension = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME:   is_stable = #vhlo.bool_v1<true>
-   // CHECK-SAME: }> ({
-@@ -1947,29 +2096,33 @@
- }
- 
- // CHECK-LABEL: "op_sqrt"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_sqrt(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.sqrt_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.sqrt_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.sqrt"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_subtract"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_subtract(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.subtract_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.subtract_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.subtract"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_tanh"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_tanh(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.tanh_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.tanh_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.tanh"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_torch_index_select"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_torch_index_select(%arg0: tensor<5x1x5xf32>, %arg1: tensor<2xi32>) ->  tensor<2x1x5xf32> {
--  //      CHECK: "vhlo.torch_index_select_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.torch_index_select_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   batch_dims = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME:   dim = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<5x1x5x!vhlo.f32_v1>, !vhlo.tensor_v1<2x!vhlo.i32_v1>) -> !vhlo.tensor_v1<2x1x5x!vhlo.f32_v1>
-@@ -1981,8 +2134,9 @@
- }
- 
- // CHECK-LABEL: "op_trace"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_trace(%arg0: tensor<f32>) {
--  //      CHECK: "vhlo.trace_v1"(%arg0) <{
-+  //      CHECK: "vhlo.trace_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   tag = #vhlo.string_v1<"foo">
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-   "stablehlo.trace"(%arg0) {
-@@ -1992,8 +2146,9 @@
- }
- 
- // CHECK-LABEL: "op_transpose"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_transpose(%arg0: tensor<16x8xf32>) ->  tensor<8x16xf32> {
--  //      CHECK: "vhlo.transpose_v1"(%arg0) <{
-+  //      CHECK: "vhlo.transpose_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   permutation = #vhlo.tensor_v1<dense<[1, 0]> : tensor<2xi64>>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<16x8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<8x16x!vhlo.f32_v1>
-   %0 = "stablehlo.transpose"(%arg0) {
-@@ -2003,8 +2158,9 @@
- }
- 
- // CHECK-LABEL: "op_triangular_solve"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_triangular_solve(%arg0: tensor<16x16xf32>, %arg1: tensor<16x16xf32>) ->  tensor<16x16xf32> {
--  //      CHECK: "vhlo.triangular_solve_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.triangular_solve_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   left_side = #vhlo.bool_v1<true>,
-   // CHECK-SAME:   lower = #vhlo.bool_v1<true>,
-   // CHECK-SAME:   transpose_a = #vhlo<transpose_v1 NO_TRANSPOSE>,
-@@ -2020,15 +2176,17 @@
- }
- 
- // CHECK-LABEL: "op_tuple"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_tuple(%arg0: tensor<f32>) -> tuple<tensor<f32>> {
--  // CHECK: "vhlo.tuple_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tuple_v1<!vhlo.tensor_v1<!vhlo.f32_v1>>
-+  // CHECK: "vhlo.tuple_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tuple_v1<!vhlo.tensor_v1<!vhlo.f32_v1>>
-   %0 = "stablehlo.tuple"(%arg0) : (tensor<f32>) -> tuple<tensor<f32>>
-   func.return %0 : tuple<tensor<f32>>
- }
- 
- // CHECK-LABEL: "op_unary_einsum"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_unary_einsum(%arg0: tensor<8x16xf32>) -> tensor<8xf32> {
--  //      CHECK: "vhlo.unary_einsum_v1"(%arg0) <{
-+  //      CHECK: "vhlo.unary_einsum_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   einsum_config = #vhlo.string_v1<"ab->a">
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<8x16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<8x!vhlo.f32_v1>
-   %0 = "stablehlo.unary_einsum"(%arg0) {
-@@ -2038,22 +2196,25 @@
- }
- 
- // CHECK-LABEL: "op_uniform_dequantize"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_uniform_dequantize(%arg0: tensor<!quant.uniform<i8:f32, 34.0:16>>) -> tensor<f32> {
--  // CHECK: "vhlo.uniform_dequantize_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.uniform_dequantize_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.uniform_dequantize"(%arg0) : (tensor<!quant.uniform<i8:f32, 34.0:16>>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_uniform_quantize"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_uniform_quantize(%arg0: tensor<f32>) -> tensor<!quant.uniform<i8:f32, 34.0:16>> {
--  // CHECK: "vhlo.uniform_quantize_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>
-+  // CHECK: "vhlo.uniform_quantize_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>
-   %0 = "stablehlo.uniform_quantize"(%arg0) : (tensor<f32>) -> tensor<!quant.uniform<i8:f32, 34.0:16>>
-   func.return %0 : tensor<!quant.uniform<i8:f32, 34.0:16>>
- }
- 
- // CHECK-LABEL: "op_while"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_while(%arg0: tensor<i1>) -> tensor<i1> {
--  //      CHECK: "vhlo.while_v1"(%arg0) ({
-+  //      CHECK: "vhlo.while_v1"(%[[ARG0]]) ({
-   // CHECK-NEXT:   ^[[BB:bb.*]](%[[ARG1:arg.*]]: !vhlo.tensor_v1<!vhlo.bool_v1>):
-   // CHECK-NEXT:     "vhlo.return_v1"(%[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>) -> ()
-   // CHECK-NEXT:   }, {
-@@ -2071,8 +2232,9 @@
- }
- 
- // CHECK-LABEL: "op_xor"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_xor(%arg0: tensor<i1>, %arg1: tensor<i1>) -> tensor<i1> {
--  // CHECK: "vhlo.xor_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-+  // CHECK: "vhlo.xor_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-   %0 = "stablehlo.xor"(%arg0, %arg1) : (tensor<i1>, tensor<i1>) -> tensor<i1>
-   func.return %0 : tensor<i1>
- }
-@@ -2080,190 +2242,217 @@
- // ============ TYPES ============
- 
- // CHECK-LABEL: "type_i1"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_i1(%arg0: tensor<i1>, %arg1: tensor<i1>) -> tensor<i1> {
--  // CHECK: "vhlo.and_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-+  // CHECK: "vhlo.and_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-   %0 = "stablehlo.and"(%arg0, %arg1) : (tensor<i1>, tensor<i1>) -> tensor<i1>
-   func.return %0 : tensor<i1>
- }
- 
- // CHECK-LABEL: "type_i4"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_i4(%arg0: tensor<i4>, %arg1: tensor<i4>) -> tensor<i4> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.i4_v1>, !vhlo.tensor_v1<!vhlo.i4_v1>) -> !vhlo.tensor_v1<!vhlo.i4_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i4_v1>, !vhlo.tensor_v1<!vhlo.i4_v1>) -> !vhlo.tensor_v1<!vhlo.i4_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<i4>, tensor<i4>) -> tensor<i4>
-   func.return %0 : tensor<i4>
- }
- 
- // CHECK-LABEL: "type_i8"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_i8(%arg0: tensor<i8>, %arg1: tensor<i8>) -> tensor<i8> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.i8_v1>, !vhlo.tensor_v1<!vhlo.i8_v1>) -> !vhlo.tensor_v1<!vhlo.i8_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i8_v1>, !vhlo.tensor_v1<!vhlo.i8_v1>) -> !vhlo.tensor_v1<!vhlo.i8_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<i8>, tensor<i8>) -> tensor<i8>
-   func.return %0 : tensor<i8>
- }
- 
- // CHECK-LABEL: "type_i16"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_i16(%arg0: tensor<i16>, %arg1: tensor<i16>) -> tensor<i16> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.i16_v1>, !vhlo.tensor_v1<!vhlo.i16_v1>) -> !vhlo.tensor_v1<!vhlo.i16_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i16_v1>, !vhlo.tensor_v1<!vhlo.i16_v1>) -> !vhlo.tensor_v1<!vhlo.i16_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<i16>, tensor<i16>) -> tensor<i16>
-   func.return %0 : tensor<i16>
- }
- 
- // CHECK-LABEL: "type_i32"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_i32(%arg0: tensor<i32>, %arg1: tensor<i32>) -> tensor<i32> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.i32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<i32>, tensor<i32>) -> tensor<i32>
-   func.return %0 : tensor<i32>
- }
- 
- // CHECK-LABEL: "type_i64"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_i64(%arg0: tensor<i64>, %arg1: tensor<i64>) -> tensor<i64> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.i64_v1>, !vhlo.tensor_v1<!vhlo.i64_v1>) -> !vhlo.tensor_v1<!vhlo.i64_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i64_v1>, !vhlo.tensor_v1<!vhlo.i64_v1>) -> !vhlo.tensor_v1<!vhlo.i64_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<i64>, tensor<i64>) -> tensor<i64>
-   func.return %0 : tensor<i64>
- }
- 
- // CHECK-LABEL: "type_ui4"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_ui4(%arg0: tensor<ui4>, %arg1: tensor<ui4>) -> tensor<ui4> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.ui4_v1>, !vhlo.tensor_v1<!vhlo.ui4_v1>) -> !vhlo.tensor_v1<!vhlo.ui4_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.ui4_v1>, !vhlo.tensor_v1<!vhlo.ui4_v1>) -> !vhlo.tensor_v1<!vhlo.ui4_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<ui4>, tensor<ui4>) -> tensor<ui4>
-   func.return %0 : tensor<ui4>
- }
- 
- // CHECK-LABEL: "type_ui8"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_ui8(%arg0: tensor<ui8>, %arg1: tensor<ui8>) -> tensor<ui8> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.ui8_v1>, !vhlo.tensor_v1<!vhlo.ui8_v1>) -> !vhlo.tensor_v1<!vhlo.ui8_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.ui8_v1>, !vhlo.tensor_v1<!vhlo.ui8_v1>) -> !vhlo.tensor_v1<!vhlo.ui8_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<ui8>, tensor<ui8>) -> tensor<ui8>
-   func.return %0 : tensor<ui8>
- }
- 
- // CHECK-LABEL: "type_ui16"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_ui16(%arg0: tensor<ui16>, %arg1: tensor<ui16>) -> tensor<ui16> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.ui16_v1>, !vhlo.tensor_v1<!vhlo.ui16_v1>) -> !vhlo.tensor_v1<!vhlo.ui16_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.ui16_v1>, !vhlo.tensor_v1<!vhlo.ui16_v1>) -> !vhlo.tensor_v1<!vhlo.ui16_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<ui16>, tensor<ui16>) -> tensor<ui16>
-   func.return %0 : tensor<ui16>
- }
- 
- // CHECK-LABEL: "type_ui32"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_ui32(%arg0: tensor<ui32>, %arg1: tensor<ui32>) -> tensor<ui32> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.ui32_v1>, !vhlo.tensor_v1<!vhlo.ui32_v1>) -> !vhlo.tensor_v1<!vhlo.ui32_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.ui32_v1>, !vhlo.tensor_v1<!vhlo.ui32_v1>) -> !vhlo.tensor_v1<!vhlo.ui32_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<ui32>, tensor<ui32>) -> tensor<ui32>
-   func.return %0 : tensor<ui32>
- }
- 
- // CHECK-LABEL: "type_ui64"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_ui64(%arg0: tensor<ui64>, %arg1: tensor<ui64>) -> tensor<ui64> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.ui64_v1>, !vhlo.tensor_v1<!vhlo.ui64_v1>) -> !vhlo.tensor_v1<!vhlo.ui64_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.ui64_v1>, !vhlo.tensor_v1<!vhlo.ui64_v1>) -> !vhlo.tensor_v1<!vhlo.ui64_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<ui64>, tensor<ui64>) -> tensor<ui64>
-   func.return %0 : tensor<ui64>
- }
- 
- // CHECK-LABEL: "type_f8E4M3FN"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_f8E4M3FN(%arg0: tensor<f8E4M3FN>, %arg1: tensor<f8E4M3FN>) -> tensor<f8E4M3FN> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f8E4M3FN_v1>, !vhlo.tensor_v1<!vhlo.f8E4M3FN_v1>) -> !vhlo.tensor_v1<!vhlo.f8E4M3FN_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f8E4M3FN_v1>, !vhlo.tensor_v1<!vhlo.f8E4M3FN_v1>) -> !vhlo.tensor_v1<!vhlo.f8E4M3FN_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f8E4M3FN>, tensor<f8E4M3FN>) -> tensor<f8E4M3FN>
-   func.return %0 : tensor<f8E4M3FN>
- }
- 
- // CHECK-LABEL: "type_f8E5M2"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_f8E5M2(%arg0: tensor<f8E5M2>, %arg1: tensor<f8E5M2>) -> tensor<f8E5M2> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f8E5M2_v1>, !vhlo.tensor_v1<!vhlo.f8E5M2_v1>) -> !vhlo.tensor_v1<!vhlo.f8E5M2_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f8E5M2_v1>, !vhlo.tensor_v1<!vhlo.f8E5M2_v1>) -> !vhlo.tensor_v1<!vhlo.f8E5M2_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f8E5M2>, tensor<f8E5M2>) -> tensor<f8E5M2>
-   func.return %0 : tensor<f8E5M2>
- }
- 
- // CHECK-LABEL: "type_f8E4M3FNUZ"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_f8E4M3FNUZ(%arg0: tensor<f8E4M3FNUZ>, %arg1: tensor<f8E4M3FNUZ>) -> tensor<f8E4M3FNUZ> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f8E4M3FNUZ_v1>, !vhlo.tensor_v1<!vhlo.f8E4M3FNUZ_v1>) -> !vhlo.tensor_v1<!vhlo.f8E4M3FNUZ_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f8E4M3FNUZ_v1>, !vhlo.tensor_v1<!vhlo.f8E4M3FNUZ_v1>) -> !vhlo.tensor_v1<!vhlo.f8E4M3FNUZ_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f8E4M3FNUZ>, tensor<f8E4M3FNUZ>) -> tensor<f8E4M3FNUZ>
-   func.return %0 : tensor<f8E4M3FNUZ>
- }
- 
- // CHECK-LABEL: "type_f8E4M3B11FNUZ"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_f8E4M3B11FNUZ(%arg0: tensor<f8E4M3B11FNUZ>, %arg1: tensor<f8E4M3B11FNUZ>) -> tensor<f8E4M3B11FNUZ> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f8E4M3B11FNUZ_v1>, !vhlo.tensor_v1<!vhlo.f8E4M3B11FNUZ_v1>) -> !vhlo.tensor_v1<!vhlo.f8E4M3B11FNUZ_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f8E4M3B11FNUZ_v1>, !vhlo.tensor_v1<!vhlo.f8E4M3B11FNUZ_v1>) -> !vhlo.tensor_v1<!vhlo.f8E4M3B11FNUZ_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f8E4M3B11FNUZ>, tensor<f8E4M3B11FNUZ>) -> tensor<f8E4M3B11FNUZ>
-   func.return %0 : tensor<f8E4M3B11FNUZ>
- }
- 
- // CHECK-LABEL: "type_f8E5M2FNUZ"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_f8E5M2FNUZ(%arg0: tensor<f8E5M2FNUZ>, %arg1: tensor<f8E5M2FNUZ>) -> tensor<f8E5M2FNUZ> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f8E5M2FNUZ_v1>, !vhlo.tensor_v1<!vhlo.f8E5M2FNUZ_v1>) -> !vhlo.tensor_v1<!vhlo.f8E5M2FNUZ_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f8E5M2FNUZ_v1>, !vhlo.tensor_v1<!vhlo.f8E5M2FNUZ_v1>) -> !vhlo.tensor_v1<!vhlo.f8E5M2FNUZ_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f8E5M2FNUZ>, tensor<f8E5M2FNUZ>) -> tensor<f8E5M2FNUZ>
-   func.return %0 : tensor<f8E5M2FNUZ>
- }
- 
- // CHECK-LABEL: "type_bf16"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_bf16(%arg0: tensor<bf16>, %arg1: tensor<bf16>) -> tensor<bf16> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.bf16_v1>, !vhlo.tensor_v1<!vhlo.bf16_v1>) -> !vhlo.tensor_v1<!vhlo.bf16_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.bf16_v1>, !vhlo.tensor_v1<!vhlo.bf16_v1>) -> !vhlo.tensor_v1<!vhlo.bf16_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<bf16>, tensor<bf16>) -> tensor<bf16>
-   func.return %0 : tensor<bf16>
- }
- 
- // CHECK-LABEL: "type_f16"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_f16(%arg0: tensor<f16>, %arg1: tensor<f16>) -> tensor<f16> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f16_v1>, !vhlo.tensor_v1<!vhlo.f16_v1>) -> !vhlo.tensor_v1<!vhlo.f16_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f16_v1>, !vhlo.tensor_v1<!vhlo.f16_v1>) -> !vhlo.tensor_v1<!vhlo.f16_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f16>, tensor<f16>) -> tensor<f16>
-   func.return %0 : tensor<f16>
- }
- 
- // CHECK-LABEL: "type_f32"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_f32(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "type_f64"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_f64(%arg0: tensor<f64>, %arg1: tensor<f64>) -> tensor<f64> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f64_v1>, !vhlo.tensor_v1<!vhlo.f64_v1>) -> !vhlo.tensor_v1<!vhlo.f64_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f64_v1>, !vhlo.tensor_v1<!vhlo.f64_v1>) -> !vhlo.tensor_v1<!vhlo.f64_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f64>, tensor<f64>) -> tensor<f64>
-   func.return %0 : tensor<f64>
- }
- 
- // CHECK-LABEL: "type_complex_f32"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_complex_f32(%arg0: tensor<complex<f32>>, %arg1: tensor<complex<f32>>) -> tensor<complex<f32>> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>, !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>) -> !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>, !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>) -> !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<complex<f32>>, tensor<complex<f32>>) -> tensor<complex<f32>>
-   func.return %0 : tensor<complex<f32>>
- }
- 
- // CHECK-LABEL: "type_complex_f64"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_complex_f64(%arg0: tensor<complex<f64>>, %arg1: tensor<complex<f64>>) -> tensor<complex<f64>> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f64_v1>>, !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f64_v1>>) -> !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f64_v1>>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f64_v1>>, !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f64_v1>>) -> !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f64_v1>>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<complex<f64>>, tensor<complex<f64>>) -> tensor<complex<f64>>
-   func.return %0 : tensor<complex<f64>>
- }
- 
- // CHECK-LABEL: "type_dynamism_ranked"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @type_dynamism_ranked(%arg0: tensor<?xf32>) -> tensor<?xf32> {
--  // CHECK: "vhlo.abs_v1"(%arg0) : (!vhlo.tensor_v1<?x!vhlo.f32_v1>) -> !vhlo.tensor_v1<?x!vhlo.f32_v1>
-+  // CHECK: "vhlo.abs_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<?x!vhlo.f32_v1>) -> !vhlo.tensor_v1<?x!vhlo.f32_v1>
-   %0 = "stablehlo.abs"(%arg0) : (tensor<?xf32>) -> tensor<?xf32>
-   func.return %0 : tensor<?xf32>
- }
- 
- // CHECK-LABEL: "type_quantization"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_quantization(%arg0: tensor<!quant.uniform<i8:f32, 34.0:16>>, %arg1: tensor<!quant.uniform<i8:f32, 34.0:16>>) -> tensor<!quant.uniform<i8:f32, 34.0:16>> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>, !vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>) -> !vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>, !vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>) -> !vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<!quant.uniform<i8:f32, 34.0:16>>, tensor<!quant.uniform<i8:f32, 34.0:16>>) -> tensor<!quant.uniform<i8:f32, 34.0:16>>
-   func.return %0 : tensor<!quant.uniform<i8:f32, 34.0:16>>
- }
- 
- //       CHECK: function_type = #vhlo.type_v1<!vhlo.func_v1<(!vhlo.token_v1) -> !vhlo.token_v1>>
- // CHECK-LABEL: "type_token_callee"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @type_token_callee(%arg0: !stablehlo.token) -> !stablehlo.token {
--  // CHECK: "vhlo.return_v1"(%arg0) : (!vhlo.token_v1) -> ()
-+  // CHECK: "vhlo.return_v1"(%[[ARG0]]) : (!vhlo.token_v1) -> ()
-   return %arg0 : !stablehlo.token
- }
- 
- //       CHECK: function_type = #vhlo.type_v1<!vhlo.func_v1<(!vhlo.token_v1) -> !vhlo.token_v1>>
- // CHECK-LABEL: "type_token_caller"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @type_token_caller(%arg0: !stablehlo.token) -> !stablehlo.token {
--  // CHECK:      "vhlo.call_v1"(%arg0) <{callee = #vhlo.string_v1<"type_token_callee">}
-+  // CHECK:      "vhlo.call_v1"(%[[ARG0]]) <{callee = #vhlo.string_v1<"type_token_callee">}
-   // CHECK-SAME: (!vhlo.token_v1) -> !vhlo.token_v1
-   %0 = func.call @type_token_callee(%arg0) : (!stablehlo.token) -> !stablehlo.token
-   return %0 : !stablehlo.token
- }
- 
- // CHECK-LABEL: "type_tuple"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @type_tuple(%arg0: tuple<tensor<f32>>) -> tuple<!stablehlo.token> {
-   %0 = "stablehlo.custom_call"(%arg0) {
-     call_target_name = "foo"
-diff --ruN a/stablehlo/stablehlo/tests/vhlo/stablehlo_legalize_to_vhlo.0_17_0.mlir b/stablehlo/stablehlo/tests/vhlo/stablehlo_legalize_to_vhlo.0_17_0.mlir
---- stablehlo/stablehlo/tests/vhlo/stablehlo_legalize_to_vhlo.0_17_0.mlir
-+++ stablehlo/stablehlo/tests/vhlo/stablehlo_legalize_to_vhlo.0_17_0.mlir
-@@ -13,6 +13,7 @@
- // ============ ATTRIBUTES ============
- 
- // CHECK-LABEL: "attr_comparison_direction_eq"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_direction_eq(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     // CHECK: comparison_direction = #vhlo<comparison_direction_v1 EQ>
-@@ -22,6 +23,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_direction_ne"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_direction_ne(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     // CHECK: comparison_direction = #vhlo<comparison_direction_v1 NE>
-@@ -31,6 +33,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_direction_ge"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_direction_ge(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     // CHECK: comparison_direction = #vhlo<comparison_direction_v1 GE>
-@@ -40,6 +43,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_direction_gt"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_direction_gt(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     // CHECK: comparison_direction = #vhlo<comparison_direction_v1 GT>
-@@ -49,6 +53,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_direction_le"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_direction_le(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     // CHECK: comparison_direction = #vhlo<comparison_direction_v1 LE>
-@@ -58,6 +63,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_direction_lt"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_direction_lt(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     // CHECK: comparison_direction = #vhlo<comparison_direction_v1 LT>
-@@ -67,6 +73,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_type_notype"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_type_notype(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     comparison_direction = #stablehlo<comparison_direction EQ>
-@@ -76,6 +83,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_type_float"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_type_float(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     comparison_direction = #stablehlo<comparison_direction EQ>,
-@@ -86,6 +94,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_type_totalorder"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_type_totalorder(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     comparison_direction = #stablehlo<comparison_direction EQ>,
-@@ -96,6 +105,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_type_signed"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_type_signed(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     comparison_direction = #stablehlo<comparison_direction EQ>,
-@@ -106,6 +116,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_type_unsigned"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_type_unsigned(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     comparison_direction = #stablehlo<comparison_direction EQ>,
-@@ -118,6 +129,7 @@
- // ConvDimensionNumbers aka #stablehlo.conv is covered below.
- 
- // CHECK-LABEL: "attr_custom_call_api_version_unspecified"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_custom_call_api_version_unspecified(%arg0: tensor<f32>) -> tensor<f32> {
-   %0 = "stablehlo.custom_call"(%arg0) {
-     call_target_name = "foo",
-@@ -128,6 +140,7 @@
- }
- 
- // CHECK-LABEL: "attr_custom_call_api_version_original"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_custom_call_api_version_original(%arg0: tensor<f32>) -> tensor<f32> {
-   %0 = "stablehlo.custom_call"(%arg0) {
-     call_target_name = "foo",
-@@ -138,6 +151,7 @@
- }
- 
- // CHECK-LABEL: "attr_custom_call_api_version_status_returning"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_custom_call_api_version_status_returning(%arg0: tensor<f32>) -> tensor<f32> {
-   %0 = "stablehlo.custom_call"(%arg0) {
-     call_target_name = "foo",
-@@ -148,6 +162,7 @@
- }
- 
- // CHECK-LABEL: "attr_custom_call_api_version_status_returning_unified"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_custom_call_api_version_status_returning_unified(%arg0: tensor<f32>) -> tensor<f32> {
-   %0 = "stablehlo.custom_call"(%arg0) {
-     call_target_name = "foo",
-@@ -166,6 +181,7 @@
- // DotDimensionNumbers aka #stablehlo.dot is covered below.
- 
- // CHECK-LABEL: "attr_fft_type_fft"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_fft_type_fft(%arg0: tensor<16xcomplex<f32>>) -> tensor<16xcomplex<f32>> {
-   %0 = "stablehlo.fft"(%arg0) {
-     // CHECK: fft_type = #vhlo<fft_type_v1 FFT>
-@@ -176,6 +192,7 @@
- }
- 
- // CHECK-LABEL: "attr_fft_type_ifft"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_fft_type_ifft(%arg0: tensor<16xcomplex<f32>>) -> tensor<16xcomplex<f32>> {
-   %0 = "stablehlo.fft"(%arg0) {
-     // CHECK: fft_type = #vhlo<fft_type_v1 IFFT>
-@@ -186,6 +203,7 @@
- }
- 
- // CHECK-LABEL: "attr_fft_type_rfft"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_fft_type_rfft(%arg0: tensor<16xf32>) -> tensor<9xcomplex<f32>> {
-   %0 = "stablehlo.fft"(%arg0) {
-     // CHECK: fft_type = #vhlo<fft_type_v1 RFFT>
-@@ -196,6 +214,7 @@
- }
- 
- // CHECK-LABEL: "attr_fft_type_irfft"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_fft_type_irfft(%arg0: tensor<9xcomplex<f32>>) -> tensor<16xf32> {
-   %0 = "stablehlo.fft"(%arg0) {
-     // CHECK: fft_type = #vhlo<fft_type_v1 IRFFT>
-@@ -208,6 +227,7 @@
- // GatherDimensionNumbers aka #stablehlo.gather is covered below.
- 
- // CHECK-LABEL: "attr_precision_config_default"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_precision_config_default(%arg0: tensor<8x16xf32>, %arg1: tensor<16x8xf32>) -> tensor<8x8xf32> {
-   %0 = "stablehlo.dot"(%arg0, %arg1) {
-     // CHECK: precision_config = #vhlo.array_v1<[#vhlo<precision_v1 DEFAULT>, #vhlo<precision_v1 DEFAULT>]>
-@@ -216,6 +236,7 @@
- }
- 
- // CHECK-LABEL: "attr_precision_config_high"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_precision_config_high(%arg0: tensor<8x16xf32>, %arg1: tensor<16x8xf32>) -> tensor<8x8xf32> {
-   %0 = "stablehlo.dot"(%arg0, %arg1) {
-     // CHECK: precision_config = #vhlo.array_v1<[#vhlo<precision_v1 HIGH>, #vhlo<precision_v1 HIGH>]>
-@@ -225,6 +246,7 @@
- }
- 
- // CHECK-LABEL: "attr_precision_config_highest"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_precision_config_highest(%arg0: tensor<8x16xf32>, %arg1: tensor<16x8xf32>) -> tensor<8x8xf32> {
-   %0 = "stablehlo.dot"(%arg0, %arg1) {
-     // CHECK: precision_config = #vhlo.array_v1<[#vhlo<precision_v1 HIGHEST>, #vhlo<precision_v1 HIGHEST>]>
-@@ -234,6 +256,7 @@
- }
- 
- // CHECK-LABEL: "attr_rng_algorithm_default"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_rng_algorithm_default(%arg0: tensor<f32>) -> (tensor<f32>, tensor<f32>) {
-   %0:2 = "stablehlo.rng_bit_generator"(%arg0) {
-     // CHECK: rng_algorithm = #vhlo<rng_algorithm_v1 DEFAULT>
-@@ -243,6 +266,7 @@
- }
- 
- // CHECK-LABEL: "attr_rng_algorithm_three_fry"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_rng_algorithm_three_fry(%arg0: tensor<f32>) -> (tensor<f32>, tensor<f32>) {
-   %0:2 = "stablehlo.rng_bit_generator"(%arg0) {
-     // CHECK: rng_algorithm = #vhlo<rng_algorithm_v1 THREE_FRY>
-@@ -252,6 +276,7 @@
- }
- 
- // CHECK-LABEL: "attr_rng_algorithm_philox"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_rng_algorithm_philox(%arg0: tensor<f32>) -> (tensor<f32>, tensor<f32>) {
-   %0:2 = "stablehlo.rng_bit_generator"(%arg0) {
-     // CHECK: rng_algorithm = #vhlo<rng_algorithm_v1 PHILOX>
-@@ -261,6 +286,7 @@
- }
- 
- // CHECK-LABEL: "attr_rng_distribution_uniform"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @attr_rng_distribution_uniform(%arg0: tensor<f32>, %arg1: tensor<f32>, %arg2: tensor<0xindex>) -> tensor<f32> {
-   %0 = "stablehlo.rng"(%arg0, %arg1, %arg2) {
-     // CHECK: rng_distribution = #vhlo<rng_distribution_v1 UNIFORM>
-@@ -270,6 +296,7 @@
- }
- 
- // CHECK-LABEL: "attr_rng_distribution_normal"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @attr_rng_distribution_normal(%arg0: tensor<f32>, %arg1: tensor<f32>, %arg2: tensor<0xindex>) -> tensor<f32> {
-   %0 = "stablehlo.rng"(%arg0, %arg1, %arg2) {
-     // CHECK: rng_distribution = #vhlo<rng_distribution_v1 NORMAL>
-@@ -281,6 +308,7 @@
- // ScatterDimensionNumbers aka #stablehlo.scatter is covered below.
- 
- // CHECK-LABEL: "attr_transpose_no_transpose"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_transpose_no_transpose(%arg0: tensor<16x16xf32>, %arg1: tensor<16x16xf32>) ->  tensor<16x16xf32> {
-   %0 = "stablehlo.triangular_solve"(%arg0, %arg1) {
-     left_side = true,
-@@ -293,6 +321,7 @@
- }
- 
- // CHECK-LABEL: "attr_transpose_transpose"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_transpose_transpose(%arg0: tensor<16x16xf32>, %arg1: tensor<16x16xf32>) ->  tensor<16x16xf32> {
-   %0 = "stablehlo.triangular_solve"(%arg0, %arg1) {
-     left_side = true,
-@@ -305,6 +334,7 @@
- }
- 
- // CHECK-LABEL: "attr_transpose_adjoint"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_transpose_adjoint(%arg0: tensor<16x16xf32>, %arg1: tensor<16x16xf32>) ->  tensor<16x16xf32> {
-   %0 = "stablehlo.triangular_solve"(%arg0, %arg1) {
-     left_side = true,
-@@ -319,10 +349,9 @@
- // TypeExtensionsAttr aka #stablehlo.type_extensions is covered below.
- 
- // CHECK-LABEL: "attr_type_extensions_bounds"
--func.func @attr_type_extensions_bounds(
--    %arg0: tensor<?x?xf32, #stablehlo.type_extensions<bounds = [16, ?]>>)
--    -> tensor<?x?xf32, #stablehlo.type_extensions<bounds = [16, ?]>> {
--  // CHECK: "vhlo.return_v1"(%arg0) : (!vhlo.tensor_v1<?x?x!vhlo.f32_v1, #vhlo.type_extensions_v1<bounds = [16, ?]>>) -> ()
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @attr_type_extensions_bounds(%arg0: tensor<?x?xf32, #stablehlo.type_extensions<bounds = [16, ?]>>) -> tensor<?x?xf32, #stablehlo.type_extensions<bounds = [16, ?]>> {
-+  // CHECK: "vhlo.return_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<?x?x!vhlo.f32_v1, #vhlo.type_extensions_v1<bounds = [16, ?]>>) -> ()
-   func.return %arg0 : tensor<?x?xf32, #stablehlo.type_extensions<bounds = [16, ?]>>
- }
- 
-@@ -330,8 +359,9 @@
- // ============ DEFAULTS ============
- 
- // CHECK-LABEL: "default_all_gather"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_all_gather(%arg0: tensor<16x8xf32>) -> tensor<16x16xf32> {
--  //               CHECK: "vhlo.all_gather_v1"(%arg0) <{
-+  //               CHECK: "vhlo.all_gather_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   all_gather_dim = #vhlo.integer_v1<1 : i64>
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0], [1]]> : tensor<2x1xi64>>,
-@@ -345,8 +375,9 @@
- }
- 
- // CHECK-LABEL: "default_all_reduce"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_all_reduce(%arg0: tensor<f32>) -> tensor<f32> {
--  //               CHECK: "vhlo.all_reduce_v1"(%arg0)
-+  //               CHECK: "vhlo.all_reduce_v1"(%[[ARG0]])
-   //          CHECK-SAME: <{
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0], [1]]> : tensor<2x1xi64>>,
-@@ -368,8 +399,9 @@
- }
- 
- // CHECK-LABEL: "default_all_to_all"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_all_to_all(%arg0: tensor<4x16xf32>) -> tensor<16x4xf32> {
--  //               CHECK: "vhlo.all_to_all_v1"(%arg0) <{
-+  //               CHECK: "vhlo.all_to_all_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   //          CHECK-SAME:   concat_dimension = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0, 1, 2, 3]]> : tensor<1x4xi64>>,
-@@ -386,8 +418,9 @@
- }
- 
- // CHECK-LABEL: "default_cholesky"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_cholesky(%arg0: tensor<1x16x16xf32>) -> tensor<1x16x16xf32> {
--  //      CHECK: "vhlo.cholesky_v1"(%arg0) <{
-+  //      CHECK: "vhlo.cholesky_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   lower = #vhlo.bool_v1<false>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<1x16x16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<1x16x16x!vhlo.f32_v1>
-   %0 = "stablehlo.cholesky"(%arg0) : (tensor<1x16x16xf32>) -> tensor<1x16x16xf32>
-@@ -395,8 +428,9 @@
- }
- 
- // CHECK-LABEL: "default_collective_permute"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_collective_permute(%arg0: tensor<16x8xf32>) -> tensor<16x8xf32> {
--  //               CHECK: "vhlo.collective_permute_v1"(%arg0) <{
-+  //               CHECK: "vhlo.collective_permute_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME{LITERAL}:   source_target_pairs = #vhlo.tensor_v1<dense<[[0, 1], [1, 2], [2, 3]]> : tensor<3x2xi64>>
-   //          CHECK-SAME: }> : (!vhlo.tensor_v1<16x8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x8x!vhlo.f32_v1>
-@@ -407,8 +441,9 @@
- }
- 
- // CHECK-LABEL: "default_collective_broadcast"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_collective_broadcast(%arg0: tensor<16x8xf32>) -> tensor<16x8xf32> {
--  //               CHECK: "vhlo.collective_broadcast_v1"(%arg0) <{
-+  //               CHECK: "vhlo.collective_broadcast_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0, 1]]> : tensor<1x2xi64>>
-   //          CHECK-SAME: }> : (!vhlo.tensor_v1<16x8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x8x!vhlo.f32_v1>
-@@ -419,8 +454,9 @@
- }
- 
- // CHECK-LABEL: "default_compare"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @default_compare(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
--  //      CHECK: "vhlo.compare_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.compare_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   compare_type = #vhlo<comparison_type_v1 NOTYPE>,
-   // CHECK-SAME:   comparison_direction = #vhlo<comparison_direction_v1 EQ>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-@@ -431,8 +467,9 @@
- }
- 
- // CHECK-LABEL: "default_convolution"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @default_convolution(%arg0: tensor<1x8x8x207xf32>, %arg1: tensor<3x3x207x16xf32>) -> tensor<1x6x6x16xf32> {
--  //      CHECK: "vhlo.convolution_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.convolution_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   batch_group_count = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   feature_group_count = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   input_batch_dimension = #vhlo.integer_v1<0 : i64>,
-@@ -460,8 +497,9 @@
- }
- 
- // CHECK-LABEL: "default_custom_call"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_custom_call(%arg0: tensor<f32>) -> tensor<f32> {
--  //      CHECK: "vhlo.custom_call_v1"(%arg0) <{
-+  //      CHECK: "vhlo.custom_call_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   api_version = #vhlo<api_version_v1 API_VERSION_ORIGINAL>,
-   // CHECK-SAME:   backend_config = #vhlo.string_v1<"">,
-   // CHECK-SAME:   call_target_name = #vhlo.string_v1<"foo">,
-@@ -478,8 +516,9 @@
- }
- 
- // CHECK-LABEL: "default_dot_general"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @default_dot_general(%arg0: tensor<8x8x16xf32>, %arg1: tensor<8x16x8xf32>) -> tensor<8x8x8xf32> {
--  //      CHECK: "vhlo.dot_general_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.dot_general_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   lhs_batching_dimensions = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>,
-   // CHECK-SAME:   lhs_contracting_dimensions = #vhlo.tensor_v1<dense<2> : tensor<1xi64>>,
-   // CHECK-SAME:   precision_config = #vhlo.array_v1<[#vhlo<precision_v1 DEFAULT>, #vhlo<precision_v1 DEFAULT>]>,
-@@ -498,8 +537,9 @@
- }
- 
- // CHECK-LABEL: "default_dot"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @default_dot(%arg0: tensor<8x16xf32>, %arg1: tensor<16x8xf32>) -> tensor<8x8xf32> {
--  //      CHECK: "vhlo.dot_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.dot_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   precision_config = #vhlo.array_v1<[#vhlo<precision_v1 DEFAULT>, #vhlo<precision_v1 DEFAULT>]>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<8x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<8x8x!vhlo.f32_v1>
-   %0 = "stablehlo.dot"(%arg0, %arg1) : (tensor<8x16xf32>, tensor<16x8xf32>) -> tensor<8x8xf32>
-@@ -507,8 +547,9 @@
- }
- 
- // CHECK-LABEL: "default_dynamic_broadcast_in_dim"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @default_dynamic_broadcast_in_dim(%arg0: tensor<?x?xf32>, %arg1: tensor<2xindex>) -> tensor<?x?xf32> {
--  //      CHECK: "vhlo.dynamic_broadcast_in_dim_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.dynamic_broadcast_in_dim_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   broadcast_dimensions = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-   // CHECK-SAME:   known_expanding_dimensions = #vhlo.tensor_v1<dense<> : tensor<0xi64>>,
-   // CHECK-SAME:   known_nonexpanding_dimensions = #vhlo.tensor_v1<dense<> : tensor<0xi64>>
-@@ -520,8 +561,9 @@
- }
- 
- // CHECK-LABEL: "default_dynamic_conv"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @default_dynamic_conv(%arg0: tensor<1x8x8x207xf32>, %arg1: tensor<3x3x207x16xf32>, %arg2: tensor<4xi32>) -> tensor<1x?x?x16xf32> {
--  //      CHECK: "vhlo.dynamic_conv_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.dynamic_conv_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   batch_group_count = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   feature_group_count = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   input_batch_dimension = #vhlo.integer_v1<0 : i64>,
-@@ -549,8 +591,9 @@
- }
- 
- // CHECK-LABEL: "default_dynamic_gather"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @default_dynamic_gather(%arg0 : tensor<2x4x9xf32>, %arg1 : tensor<1x5x2xi32>, %arg2 : tensor<3xi32>) -> tensor<1x5x8xf32> {
--  //      CHECK: "vhlo.dynamic_gather_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.dynamic_gather_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   collapsed_slice_dims = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-   // CHECK-SAME:   index_vector_dim = #vhlo.integer_v1<2 : i64>,
-   // CHECK-SAME:   indices_are_sorted = #vhlo.bool_v1<false>,
-@@ -576,15 +619,16 @@
-   // CHECK-SAME:   sym_name = #vhlo.string_v1<"default_func">,
-   // CHECK-SAME:   sym_visibility = #vhlo.string_v1<"">
-   // CHECK-SAME: }> ({
--  // CHECK-NEXT:   ^[[BB:bb.*]](%arg0: !vhlo.tensor_v1<!vhlo.f32_v1>):
--  // CHECK-NEXT:     "vhlo.return_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-+  // CHECK-NEXT:   ^[[BB:bb.*]](%[[ARG0:.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
-+  // CHECK-NEXT:     "vhlo.return_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-   // CHECK-NEXT: }) : () -> ()
-   func.return %arg0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "dynamic_gather"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @dynamic_gather(%arg0 : tensor<2x4x9xf32>, %arg1 : tensor<1x5x2xi32>) -> tensor<1x5x1xf32> {
--  //      CHECK: "vhlo.gather_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.gather_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   collapsed_slice_dims = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-   // CHECK-SAME:   index_vector_dim = #vhlo.integer_v1<2 : i64>,
-   // CHECK-SAME:   indices_are_sorted = #vhlo.bool_v1<false>,
-@@ -605,8 +649,9 @@
- }
- 
- // CHECK-LABEL: "default_infeed"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_infeed(%arg0: !stablehlo.token) -> (tensor<f32>, !stablehlo.token) {
--  //               CHECK: "vhlo.infeed_v1"(%arg0) <{
-+  //               CHECK: "vhlo.infeed_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   infeed_config = #vhlo.string_v1<"">,
-   // CHECK-SAME{LITERAL}:   layout = #vhlo.array_v1<[]>
-   //          CHECK-SAME: }> : (!vhlo.token_v1) -> (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.token_v1)
-@@ -615,8 +660,9 @@
- }
- 
- // CHECK-LABEL: "default_outfeed"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @default_outfeed(%arg0: tensor<f32>, %arg1: !stablehlo.token) -> !stablehlo.token {
--  //      CHECK: "vhlo.outfeed_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.outfeed_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   outfeed_config = #vhlo.string_v1<"">
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.token_v1) -> !vhlo.token_v1
-   %0 = "stablehlo.outfeed"(%arg0, %arg1) : (tensor<f32>, !stablehlo.token) -> !stablehlo.token
-@@ -624,8 +670,9 @@
- }
- 
- // CHECK-LABEL: "default_recv"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_recv(%arg0: !stablehlo.token) -> (tensor<f32>, !stablehlo.token) {
--  //      CHECK: "vhlo.recv_v1"(%arg0) <{
-+  //      CHECK: "vhlo.recv_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME:   channel_type = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   is_host_transfer = #vhlo.bool_v1<false>
-@@ -637,8 +684,9 @@
- }
- 
- // CHECK-LABEL: "default_send"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @default_send(%arg0: tensor<f32>, %arg1: !stablehlo.token) -> !stablehlo.token {
--  //      CHECK: "vhlo.send_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.send_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME:   channel_type = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   is_host_transfer = #vhlo.bool_v1<false>
-@@ -650,8 +698,9 @@
- }
- 
- // CHECK-LABEL: "default_reduce_scatter"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_reduce_scatter(%arg0: tensor<16xf32>) -> tensor<16xf32> {
--  //               CHECK: "vhlo.reduce_scatter_v1"(%arg0) <{
-+  //               CHECK: "vhlo.reduce_scatter_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0], [1]]> : tensor<2x1xi64>>,
-   //          CHECK-SAME:   scatter_dimension = #vhlo.integer_v1<0 : i64>
-@@ -673,8 +722,9 @@
- }
- 
- // CHECK-LABEL: "default_reduce_window"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @default_reduce_window(%arg0: tensor<2x17x31x7xf32>, %arg1: tensor<f32>) -> tensor<2x16x30x7xf32> {
--  //               CHECK: "vhlo.reduce_window_v1"(%arg0, %arg1)  <{
-+  //               CHECK: "vhlo.reduce_window_v1"(%[[ARG0]], %[[ARG1]])  <{
-   //          CHECK-SAME:   base_dilations = #vhlo.tensor_v1<dense<1> : tensor<4xi64>>,
-   // CHECK-SAME{LITERAL}:   padding = #vhlo.tensor_v1<dense<0> : tensor<4x2xi64>>,
-   //          CHECK-SAME:   window_dilations = #vhlo.tensor_v1<dense<1> : tensor<4xi64>>,
-@@ -696,8 +746,9 @@
- }
- 
- // CHECK-LABEL: "default_scatter"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @default_scatter(%arg0: tensor<200x100x300xf32>, %arg1: tensor<10x2xi32>, %arg2: tensor<10x300xf32>) -> tensor<200x100x300xf32> {
--  //      CHECK: "vhlo.scatter_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.scatter_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   index_vector_dim = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   indices_are_sorted = #vhlo.bool_v1<false>,
-   // CHECK-SAME:   inserted_window_dims = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-@@ -725,8 +776,9 @@
- }
- 
- // CHECK-LABEL: "default_select_and_scatter"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @default_select_and_scatter(%arg0: tensor<10x24x24x64xf32>, %arg1: tensor<10x23x23x64xf32>, %arg2: tensor<f32>) -> tensor<10x24x24x64xf32> {
--  //      CHECK: "vhlo.select_and_scatter_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.select_and_scatter_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   padding = #vhlo.tensor_v1<dense<0> : tensor<4x2xi64>>,
-   // CHECK-SAME:   window_dimensions = #vhlo.tensor_v1<dense<[1, 2, 2, 1]> : tensor<4xi64>>,
-   // CHECK-SAME:   window_strides = #vhlo.tensor_v1<dense<1> : tensor<4xi64>>
-@@ -754,8 +806,9 @@
- }
- 
- // CHECK-LABEL: "default_sort"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_sort(%arg0: tensor<16xf32>) -> tensor<16xf32> {
--  //      CHECK: "vhlo.sort_v1"(%arg0) <{
-+  //      CHECK: "vhlo.sort_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   dimension = #vhlo.integer_v1<-1 : i64>
-   // CHECK-SAME:   is_stable = #vhlo.bool_v1<false>
-   // CHECK-SAME: }> ({
-@@ -774,29 +827,33 @@
- // ============ OPS ============
- 
- // CHECK-LABEL: "op_abs"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_abs(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.abs_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.abs_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.abs"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_add"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_add(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_after_all"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_after_all(%arg0: !stablehlo.token) -> !stablehlo.token {
--  // CHECK: "vhlo.after_all_v1"(%arg0) : (!vhlo.token_v1) -> !vhlo.token_v1
-+  // CHECK: "vhlo.after_all_v1"(%[[ARG0]]) : (!vhlo.token_v1) -> !vhlo.token_v1
-   %0 = "stablehlo.after_all"(%arg0) : (!stablehlo.token) -> !stablehlo.token
-   func.return %0 : !stablehlo.token
- }
- 
- // CHECK-LABEL: "op_all_gather"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_all_gather(%arg0: tensor<16x8xf32>) -> tensor<16x16xf32> {
--  //               CHECK: "vhlo.all_gather_v1"(%arg0) <{
-+  //               CHECK: "vhlo.all_gather_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   all_gather_dim = #vhlo.integer_v1<1 : i64>
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0], [1]]> : tensor<2x1xi64>>,
-@@ -812,8 +869,9 @@
- }
- 
- // CHECK-LABEL: "op_all_reduce"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_all_reduce(%arg0: tensor<f32>) -> tensor<f32> {
--  //               CHECK: "vhlo.all_reduce_v1"(%arg0) <{
-+  //               CHECK: "vhlo.all_reduce_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0], [1]]> : tensor<2x1xi64>>,
-   //          CHECK-SAME:   use_global_device_ids = #vhlo.bool_v1<true>
-@@ -836,7 +894,7 @@
- 
- // CHECK-LABEL: "op_all_reduce_with_promotable_types"
- func.func @op_all_reduce_with_promotable_types(%operand: tensor<f32>) -> tensor<f64> {
--  //  CHECK: "vhlo.all_reduce_v1"(%arg0)
-+  //  CHECK: "vhlo.all_reduce_v1"(%[[ARG0:.*]])
-   //  CHECK:   ^[[BB:bb.*]](%[[ARG1:arg.*]]: !vhlo.tensor_v1<!vhlo.f64_v1>, %[[ARG2:arg.*]]: !vhlo.tensor_v1<!vhlo.f64_v1>):
-   //  CHECK:     "vhlo.return_v1"(%[[VAL1:.*]]) : (!vhlo.tensor_v1<!vhlo.f64_v1>) -> ()
-   //  CHECK: }) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f64_v1>
-@@ -854,8 +912,9 @@
- }
- 
- // CHECK-LABEL: "op_all_to_all"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_all_to_all(%arg0: tensor<4x16xf32>) -> tensor<16x4xf32> {
--  //               CHECK: "vhlo.all_to_all_v1"(%arg0) <{
-+  //               CHECK: "vhlo.all_to_all_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<1 : i64>,
-   //          CHECK-SAME:   concat_dimension = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0, 1, 2, 3]]> : tensor<1x4xi64>>,
-@@ -873,22 +932,25 @@
- }
- 
- // CHECK-LABEL: "op_and"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_and(%arg0: tensor<i1>, %arg1: tensor<i1>) -> tensor<i1> {
--  // CHECK: "vhlo.and_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-+  // CHECK: "vhlo.and_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-   %0 = "stablehlo.and"(%arg0, %arg1) : (tensor<i1>, tensor<i1>) -> tensor<i1>
-   func.return %0 : tensor<i1>
- }
- 
- // CHECK-LABEL: "op_atan2"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_atan2(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.atan2_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.atan2_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.atan2"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_batch_norm_grad"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}}, %[[ARG3:.*]]: {{.*}}, %[[ARG4:.*]]: {{.*}})
- func.func @op_batch_norm_grad(%arg0: tensor<16x16x16x16xf32>, %arg1: tensor<16xf32>, %arg2: tensor<16xf32>, %arg3: tensor<16xf32>, %arg4: tensor<16x16x16x16xf32>) -> (tensor<16x16x16x16xf32>, tensor<16xf32>, tensor<16xf32>) {
--  //      CHECK: "vhlo.batch_norm_grad_v1"(%arg0, %arg1, %arg2, %arg3, %arg4) <{
-+  //      CHECK: "vhlo.batch_norm_grad_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]], %[[ARG3]], %[[ARG4]]) <{
-   // CHECK-SAME:   epsilon = #vhlo.float_v1<1.000000e-03 : !vhlo.f32_v1>,
-   // CHECK-SAME:   feature_index = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<16x16x16x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x16x16x16x!vhlo.f32_v1>) -> (!vhlo.tensor_v1<16x16x16x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>)
-@@ -900,8 +962,9 @@
- }
- 
- // CHECK-LABEL: "op_batch_norm_inference"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}}, %[[ARG3:.*]]: {{.*}}, %[[ARG4:.*]]: {{.*}})
- func.func @op_batch_norm_inference(%arg0: tensor<16x16x16x16xf32>, %arg1: tensor<16xf32>, %arg2: tensor<16xf32>, %arg3: tensor<16xf32>, %arg4: tensor<16xf32>) -> tensor<16x16x16x16xf32> {
--  //      CHECK: "vhlo.batch_norm_inference_v1"(%arg0, %arg1, %arg2, %arg3, %arg4) <{
-+  //      CHECK: "vhlo.batch_norm_inference_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]], %[[ARG3]], %[[ARG4]]) <{
-   // CHECK-SAME:   epsilon = #vhlo.float_v1<1.000000e-03 : !vhlo.f32_v1>,
-   // CHECK-SAME:   feature_index = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<16x16x16x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x16x16x16x!vhlo.f32_v1>
-@@ -913,8 +976,9 @@
- }
- 
- // CHECK-LABEL: "op_batch_norm_training"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_batch_norm_training(%arg0: tensor<16x16x16x16xf32>, %arg1: tensor<16xf32>, %arg2: tensor<16xf32>) -> (tensor<16x16x16x16xf32>, tensor<16xf32>, tensor<16xf32>) {
--  //      CHECK: "vhlo.batch_norm_training_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.batch_norm_training_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   epsilon = #vhlo.float_v1<1.000000e-03 : !vhlo.f32_v1>,
-   // CHECK-SAME:   feature_index = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<16x16x16x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>) -> (!vhlo.tensor_v1<16x16x16x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>)
-@@ -926,15 +990,17 @@
- }
- 
- // CHECK-LABEL: "op_bitcast_convert"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_bitcast_convert(%arg0: tensor<i32>) -> tensor<f32> {
--  // CHECK: "vhlo.bitcast_convert_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.bitcast_convert_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.bitcast_convert"(%arg0) : (tensor<i32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_broadcast_in_dim"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_broadcast_in_dim(%arg0: tensor<16xf32>) -> tensor<16x16xf32> {
--  //      CHECK: "vhlo.broadcast_in_dim_v1"(%arg0) <{
-+  //      CHECK: "vhlo.broadcast_in_dim_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   broadcast_dimensions = #vhlo.tensor_v1<dense<1> : tensor<1xi64>>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x16x!vhlo.f32_v1>
-   %0 = "stablehlo.broadcast_in_dim"(%arg0) {
-@@ -944,8 +1010,9 @@
- }
- 
- // CHECK-LABEL: "op_broadcast"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_broadcast(%arg0: tensor<16xf32>) -> tensor<16x16xf32> {
--  //      CHECK: "vhlo.broadcast_v1"(%arg0) <{
-+  //      CHECK: "vhlo.broadcast_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   broadcast_sizes = #vhlo.tensor_v1<dense<16> : tensor<1xi64>>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x16x!vhlo.f32_v1>
-   %0 = "stablehlo.broadcast"(%arg0) {
-@@ -955,9 +1022,10 @@
- }
- 
- // CHECK-LABEL: "op_case"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_case(%arg0: tensor<i32>, %arg1: tensor<f32>) -> tensor<f32> {
--  //      CHECK: "vhlo.case_v1"(%arg0) ({
--  // CHECK-NEXT:   "vhlo.return_v1"(%arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-+  //      CHECK: "vhlo.case_v1"(%[[ARG0]]) ({
-+  // CHECK-NEXT:   "vhlo.return_v1"(%[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-   // CHECK-NEXT: }) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.case"(%arg0) ({
-     "stablehlo.return"(%arg1) : (tensor<f32>) -> ()
-@@ -966,22 +1034,25 @@
- }
- 
- // CHECK-LABEL: "op_cbrt"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_cbrt(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.cbrt_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.cbrt_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.cbrt"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_ceil"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_ceil(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.ceil_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.ceil_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.ceil"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_cholesky"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_cholesky(%arg0: tensor<1x16x16xf32>) -> tensor<1x16x16xf32> {
--  //      CHECK: "vhlo.cholesky_v1"(%arg0) <{
-+  //      CHECK: "vhlo.cholesky_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   lower = #vhlo.bool_v1<true>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<1x16x16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<1x16x16x!vhlo.f32_v1>
-   %0 = "stablehlo.cholesky"(%arg0) {
-@@ -991,22 +1062,25 @@
- }
- 
- // CHECK-LABEL: "op_clamp"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_clamp(%arg0: tensor<f32>, %arg1: tensor<f32>, %arg2: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.clamp_v1"(%arg0, %arg1, %arg2) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.clamp_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.clamp"(%arg0, %arg1, %arg2) : (tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_count_leading_zeros"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_count_leading_zeros(%arg0: tensor<i32>) -> tensor<i32> {
--  // CHECK: "vhlo.count_leading_zeros_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-+  // CHECK: "vhlo.count_leading_zeros_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-   %0 = "stablehlo.count_leading_zeros"(%arg0) : (tensor<i32>) -> tensor<i32>
-   func.return %0 : tensor<i32>
- }
- 
- // CHECK-LABEL: "op_collective_permute"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_collective_permute(%arg0: tensor<16x8xf32>) -> tensor<16x8xf32> {
--  //               CHECK: "vhlo.collective_permute_v1"(%arg0) <{
-+  //               CHECK: "vhlo.collective_permute_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME{LITERAL}:   source_target_pairs = #vhlo.tensor_v1<dense<[[0, 1], [1, 2], [2, 3]]> : tensor<3x2xi64>>
-   //          CHECK-SAME: }> : (!vhlo.tensor_v1<16x8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x8x!vhlo.f32_v1>
-@@ -1018,8 +1092,9 @@
- }
- 
- // CHECK-LABEL: "op_compare"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_compare(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
--  //      CHECK: "vhlo.compare_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.compare_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   compare_type = #vhlo<comparison_type_v1 TOTALORDER>,
-   // CHECK-SAME:   comparison_direction = #vhlo<comparison_direction_v1 EQ>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-@@ -1031,15 +1106,17 @@
- }
- 
- // CHECK-LABEL: "op_complex"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_complex(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<complex<f32>> {
--  // CHECK: "vhlo.complex_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>
-+  // CHECK: "vhlo.complex_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>
-   %0 = "stablehlo.complex"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<complex<f32>>
-   func.return %0 : tensor<complex<f32>>
- }
- 
- // CHECK-LABEL: "op_concatenate"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_concatenate(%arg0: tensor<8xf32>, %arg1: tensor<8xf32>) -> tensor<16xf32> {
--  //      CHECK: "vhlo.concatenate_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.concatenate_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   dimension = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<8x!vhlo.f32_v1>, !vhlo.tensor_v1<8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x!vhlo.f32_v1>
-   %0 = "stablehlo.concatenate"(%arg0, %arg1) {
-@@ -1049,6 +1126,7 @@
- }
- 
- // CHECK-LABEL: "op_constant"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_constant(%arg0: tensor<f32>) -> tensor<f32> {
-   //      CHECK: "vhlo.constant_v1"() <{
-   // CHECK-SAME:   value = #vhlo.tensor_v1<dense<0.000000e+00> : tensor<f32>>
-@@ -1060,15 +1138,17 @@
- }
- 
- // CHECK-LABEL: "op_convert"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_convert(%arg0: tensor<i32>) -> tensor<f32> {
--  // CHECK: "vhlo.convert_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.convert_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.convert"(%arg0) : (tensor<i32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_convolution"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_convolution(%arg0: tensor<1x8x8x207xf32>, %arg1: tensor<3x3x207x16xf32>) -> tensor<1x7x7x16xf32> {
--  //      CHECK: "vhlo.convolution_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.convolution_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   batch_group_count = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   feature_group_count = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   input_batch_dimension = #vhlo.integer_v1<0 : i64>,
-@@ -1102,8 +1182,9 @@
- }
- 
- // CHECK-LABEL: "op_cosine"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_cosine(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.cosine_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.cosine_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.cosine"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
-@@ -1116,8 +1197,9 @@
- }
- 
- // CHECK-LABEL: "op_cross_replica_sum"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_cross_replica_sum(%arg0: tensor<f32>) -> tensor<f32> {
--  //               CHECK: "vhlo.cross-replica-sum_v1"(%arg0) <{
-+  //               CHECK: "vhlo.cross-replica-sum_v1"(%[[ARG0]]) <{
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0], [1]]> : tensor<2x1xi64>>
-   //          CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.cross-replica-sum"(%arg0) {
-@@ -1127,8 +1209,9 @@
- }
- 
- // CHECK-LABEL: "op_custom_call"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_custom_call(%arg0: tensor<f32>) -> tensor<f32> {
--  //      CHECK: "vhlo.custom_call_v1"(%arg0) <{
-+  //      CHECK: "vhlo.custom_call_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   api_version = #vhlo<api_version_v1 API_VERSION_STATUS_RETURNING>,
-   // CHECK-SAME:   backend_config = #vhlo.string_v1<"\08\03\1A\02">,
-   // CHECK-SAME:   call_target_name = #vhlo.string_v1<"foo">,
-@@ -1159,15 +1242,17 @@
- }
- 
- // CHECK-LABEL: "op_divide"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_divide(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.divide_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.divide_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.divide"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_dot_general"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_dot_general(%arg0: tensor<8x8x16xf32>, %arg1: tensor<8x16x8xf32>) -> tensor<8x8x8xf32> {
--  //      CHECK: "vhlo.dot_general_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.dot_general_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   lhs_batching_dimensions = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>,
-   // CHECK-SAME:   lhs_contracting_dimensions = #vhlo.tensor_v1<dense<2> : tensor<1xi64>>,
-   // CHECK-SAME:   precision_config = #vhlo.array_v1<[#vhlo<precision_v1 HIGHEST>, #vhlo<precision_v1 HIGHEST>]>,
-@@ -1187,8 +1272,9 @@
- }
- 
- // CHECK-LABEL: "op_dot"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_dot(%arg0: tensor<8x16xf32>, %arg1: tensor<16x8xf32>) -> tensor<8x8xf32> {
--  //      CHECK: "vhlo.dot_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.dot_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   precision_config = #vhlo.array_v1<[#vhlo<precision_v1 HIGHEST>, #vhlo<precision_v1 HIGHEST>]>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<8x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<8x8x!vhlo.f32_v1>
-   %0 = "stablehlo.dot"(%arg0, %arg1) {
-@@ -1198,8 +1284,9 @@
- }
- 
- // CHECK-LABEL: "op_dynamic_broadcast_in_dim"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_dynamic_broadcast_in_dim(%arg0: tensor<?x?xf32>, %arg1: tensor<2xindex>) -> tensor<?x?xf32> {
--  //      CHECK: "vhlo.dynamic_broadcast_in_dim_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.dynamic_broadcast_in_dim_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   broadcast_dimensions = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-   // CHECK-SAME:   known_expanding_dimensions = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>,
-   // CHECK-SAME:   known_nonexpanding_dimensions = #vhlo.tensor_v1<dense<1> : tensor<1xi64>>
-@@ -1213,8 +1300,9 @@
- }
- 
- // CHECK-LABEL: "op_dynamic_conv"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_dynamic_conv(%arg0: tensor<1x8x8x207xf32>, %arg1: tensor<3x3x207x16xf32>, %arg2: tensor<4xi32>) -> tensor<1x?x?x16xf32> {
--  //      CHECK: "vhlo.dynamic_conv_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.dynamic_conv_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   batch_group_count = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   feature_group_count = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   input_batch_dimension = #vhlo.integer_v1<0 : i64>,
-@@ -1248,8 +1336,9 @@
- }
- 
- // CHECK-LABEL: "op_dynamic_gather"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_dynamic_gather(%arg0 : tensor<2x4x9xf32>, %arg1 : tensor<1x5x2xi32>, %arg2 : tensor<3xi32>) -> tensor<1x5x8xf32> {
--  //      CHECK: "vhlo.dynamic_gather_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.dynamic_gather_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   collapsed_slice_dims = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-   // CHECK-SAME:   index_vector_dim = #vhlo.integer_v1<2 : i64>,
-   // CHECK-SAME:   indices_are_sorted = #vhlo.bool_v1<true>,
-@@ -1269,8 +1358,9 @@
- }
- 
- // CHECK-LABEL: "op_dynamic_iota"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_dynamic_iota(%arg0: tensor<1xindex>) -> tensor<?xf32> {
--  //      CHECK: "vhlo.dynamic_iota_v1"(%arg0) <{
-+  //      CHECK: "vhlo.dynamic_iota_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   iota_dimension = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<1x!vhlo.index_v1>) -> !vhlo.tensor_v1<?x!vhlo.f32_v1>
-   %0 = "stablehlo.dynamic_iota"(%arg0) {
-@@ -1280,22 +1370,25 @@
- }
- 
- // CHECK-LABEL: "op_dynamic_pad"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}}, %[[ARG3:.*]]: {{.*}}, %[[ARG4:.*]]: {{.*}})
- func.func @op_dynamic_pad(%arg0: tensor<?xf32>, %arg1: tensor<f32>, %arg2: tensor<1xindex>, %arg3: tensor<1xindex>, %arg4: tensor<1xindex>) -> tensor<?xf32> {
--  // CHECK: "vhlo.dynamic_pad_v1"(%arg0, %arg1, %arg2, %arg3, %arg4) : (!vhlo.tensor_v1<?x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>) -> !vhlo.tensor_v1<?x!vhlo.f32_v1>
-+  // CHECK: "vhlo.dynamic_pad_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]], %[[ARG3]], %[[ARG4]]) : (!vhlo.tensor_v1<?x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>) -> !vhlo.tensor_v1<?x!vhlo.f32_v1>
-   %0 = "stablehlo.dynamic_pad"(%arg0, %arg1, %arg2, %arg3, %arg4) : (tensor<?xf32>, tensor<f32>, tensor<1xindex>, tensor<1xindex>, tensor<1xindex>) -> tensor<?xf32>
-   func.return %0 : tensor<?xf32>
- }
- 
- // CHECK-LABEL: "op_dynamic_reshape"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_dynamic_reshape(%arg0: tensor<16xf32>, %arg1: tensor<2xindex>) -> tensor<?x?xf32> {
--  // CHECK: "vhlo.dynamic_reshape_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<2x!vhlo.index_v1>) -> !vhlo.tensor_v1<?x?x!vhlo.f32_v1>
-+  // CHECK: "vhlo.dynamic_reshape_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<2x!vhlo.index_v1>) -> !vhlo.tensor_v1<?x?x!vhlo.f32_v1>
-   %0 = "stablehlo.dynamic_reshape"(%arg0, %arg1) : (tensor<16xf32>, tensor<2xindex>) -> tensor<?x?xf32>
-   func.return %0 : tensor<?x?xf32>
- }
- 
- // CHECK-LABEL: "op_dynamic_slice"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_dynamic_slice(%arg0: tensor<16xf32>, %arg1: tensor<i64>) -> tensor<4xf32> {
--  //      CHECK: "vhlo.dynamic_slice_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.dynamic_slice_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   slice_sizes = #vhlo.tensor_v1<dense<4> : tensor<1xi64>>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.i64_v1>) -> !vhlo.tensor_v1<4x!vhlo.f32_v1>
-   %0 = "stablehlo.dynamic_slice"(%arg0, %arg1) {
-@@ -1305,15 +1398,17 @@
- }
- 
- // CHECK-LABEL: "op_dynamic_update_slice"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_dynamic_update_slice(%arg0: tensor<16xf32>, %arg1: tensor<4xf32>, %arg2: tensor<i64>) -> tensor<16xf32> {
--  // CHECK: "vhlo.dynamic_update_slice_v1"(%arg0, %arg1, %arg2) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<4x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.i64_v1>) -> !vhlo.tensor_v1<16x!vhlo.f32_v1>
-+  // CHECK: "vhlo.dynamic_update_slice_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<4x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.i64_v1>) -> !vhlo.tensor_v1<16x!vhlo.f32_v1>
-   %0 = "stablehlo.dynamic_update_slice"(%arg0, %arg1, %arg2) : (tensor<16xf32>, tensor<4xf32>, tensor<i64>) -> tensor<16xf32>
-   func.return %0 : tensor<16xf32>
- }
- 
- // CHECK-LABEL: "op_einsum"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_einsum(%arg0: tensor<8x16xf32>, %arg1: tensor<16x8xf32>) -> tensor<8x8xf32> {
--  //      CHECK: "vhlo.einsum_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.einsum_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   einsum_config = #vhlo.string_v1<"ab,bc->ac">
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<8x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<8x8x!vhlo.f32_v1>
-   %0 = "stablehlo.einsum"(%arg0, %arg1) {
-@@ -1323,22 +1418,25 @@
- }
- 
- // CHECK-LABEL: "op_exponential_minus_one"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_exponential_minus_one(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.exponential_minus_one_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.exponential_minus_one_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.exponential_minus_one"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_exponential"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_exponential(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.exponential_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.exponential_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.exponential"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_fft"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_fft(%arg0: tensor<16xcomplex<f32>>) -> tensor<16xcomplex<f32>> {
--  //      CHECK: "vhlo.fft_v1"(%arg0) <{
-+  //      CHECK: "vhlo.fft_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   fft_length = #vhlo.tensor_v1<dense<16> : tensor<1xi64>>,
-   // CHECK-SAME:   fft_type = #vhlo<fft_type_v1 FFT>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<16x!vhlo.complex_v1<!vhlo.f32_v1>>) -> !vhlo.tensor_v1<16x!vhlo.complex_v1<!vhlo.f32_v1>>
-@@ -1350,8 +1448,9 @@
- }
- 
- // CHECK-LABEL: "op_floor"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_floor(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.floor_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.floor_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.floor"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
-@@ -1364,16 +1463,17 @@
-   // CHECK-SAME:   sym_name = #vhlo.string_v1<"op_func">,
-   // CHECK-SAME:   sym_visibility = #vhlo.string_v1<"private">
-   // CHECK-SAME: }> ({
--  // CHECK-NEXT:   ^[[BB:bb.*]](%arg0: !vhlo.tensor_v1<!vhlo.f32_v1>):
--  // CHECK-NEXT:     "vhlo.return_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-+  // CHECK-NEXT:   ^[[BB:bb.*]](%[[ARG0:.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
-+  // CHECK-NEXT:     "vhlo.return_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-   // CHECK-NEXT: }) : () -> ()
- 
-   func.return %arg0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_gather"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_gather(%arg0 : tensor<2x4x9xf32>, %arg1 : tensor<1x5x2xi32>) -> tensor<1x5x1xf32> {
--  //      CHECK: "vhlo.gather_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.gather_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   collapsed_slice_dims = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-   // CHECK-SAME:   index_vector_dim = #vhlo.integer_v1<2 : i64>,
-   // CHECK-SAME:   indices_are_sorted = #vhlo.bool_v1<true>,
-@@ -1395,8 +1495,9 @@
- }
- 
- // CHECK-LABEL: "op_get_dimension_size"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_get_dimension_size(%arg0: tensor<?xf32>) -> tensor<i32> {
--  //      CHECK: "vhlo.get_dimension_size_v1"(%arg0) <{
-+  //      CHECK: "vhlo.get_dimension_size_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   dimension = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<?x!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-   %0 = "stablehlo.get_dimension_size"(%arg0) {
-@@ -1406,8 +1507,9 @@
- }
- 
- // CHECK-LABEL: "op_get_tuple_element"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_get_tuple_element(%arg0: tuple<tensor<f32>, tensor<i32>>) -> tensor<f32> {
--  //      CHECK: "vhlo.get_tuple_element_v1"(%arg0) <{
-+  //      CHECK: "vhlo.get_tuple_element_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   index = #vhlo.integer_v1<0 : i32>
-   // CHECK-SAME: }> : (!vhlo.tuple_v1<!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.get_tuple_element"(%arg0) {
-@@ -1417,11 +1519,12 @@
- }
- 
- // CHECK-LABEL: "op_if"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_if(%arg0: tensor<i1>, %arg1: tensor<f32>, %arg2: tensor<f32>) -> tensor<f32> {
--  //      CHECK: "vhlo.if_v1"(%arg0) ({
--  // CHECK-NEXT:   "vhlo.return_v1"(%arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-+  //      CHECK: "vhlo.if_v1"(%[[ARG0]]) ({
-+  // CHECK-NEXT:   "vhlo.return_v1"(%[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-   // CHECK-NEXT: }, {
--  // CHECK-NEXT:   "vhlo.return_v1"(%arg2) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-+  // CHECK-NEXT:   "vhlo.return_v1"(%[[ARG2]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-   // CHECK-NEXT: }) : (!vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.if"(%arg0) ({
-     "stablehlo.return"(%arg1) : (tensor<f32>) -> ()
-@@ -1432,15 +1535,17 @@
- }
- 
- // CHECK-LABEL: "op_imag"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_imag(%arg0: tensor<complex<f32>>) -> tensor<f32> {
--  // CHECK: "vhlo.imag_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.imag_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.imag"(%arg0) : (tensor<complex<f32>>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_infeed"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_infeed(%arg0: !stablehlo.token) -> (tensor<f32>, !stablehlo.token) {
--  //               CHECK: "vhlo.infeed_v1"(%arg0) <{
-+  //               CHECK: "vhlo.infeed_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   infeed_config = #vhlo.string_v1<"foo">,
-   // CHECK-SAME{LITERAL}:   layout = #vhlo.array_v1<[#vhlo.array_v1<[]>]>
-   //          CHECK-SAME: }> : (!vhlo.token_v1) -> (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.token_v1)
-@@ -1463,36 +1568,41 @@
- }
- 
- // CHECK-LABEL: "op_is_finite"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_is_finite(%arg0: tensor<f32>) -> tensor<i1> {
--  // CHECK: "vhlo.is_finite_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-+  // CHECK: "vhlo.is_finite_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-   %0 = "stablehlo.is_finite"(%arg0) : (tensor<f32>) -> tensor<i1>
-   func.return %0 : tensor<i1>
- }
- 
- // CHECK-LABEL: "op_log"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_log(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.log_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.log_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.log"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_log_plus_one"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_log_plus_one(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.log_plus_one_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.log_plus_one_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.log_plus_one"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_logistic"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_logistic(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.logistic_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.logistic_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.logistic"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_map"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_map(%arg0: tensor<16xf32>) -> tensor<16xf32> {
--  //      CHECK: "vhlo.map_v1"(%arg0) <{
-+  //      CHECK: "vhlo.map_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   dimensions = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>
-   // CHECK-SAME: }> ({
-   // CHECK-NEXT:   ^[[BB:bb.*]](%[[ARG1:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
-@@ -1510,57 +1620,65 @@
- }
- 
- // CHECK-LABEL: "op_maximum"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_maximum(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.maximum_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.maximum_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.maximum"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_minimum"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_minimum(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.minimum_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.minimum_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.minimum"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_multiply"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_multiply(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.multiply_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.multiply_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.multiply"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_negate"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_negate(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.negate_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.negate_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.negate"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_not"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_not(%arg0: tensor<i1>) -> tensor<i1> {
--  // CHECK: "vhlo.not_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-+  // CHECK: "vhlo.not_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-   %0 = "stablehlo.not"(%arg0) : (tensor<i1>) -> tensor<i1>
-   func.return %0 : tensor<i1>
- }
- 
- // CHECK-LABEL: "op_optimization_barrier"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_optimization_barrier(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.optimization_barrier_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.optimization_barrier_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.optimization_barrier"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_or"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_or(%arg0: tensor<i1>, %arg1: tensor<i1>) -> tensor<i1> {
--  // CHECK: "vhlo.or_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-+  // CHECK: "vhlo.or_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-   %0 = "stablehlo.or"(%arg0, %arg1) : (tensor<i1>, tensor<i1>) -> tensor<i1>
-   func.return %0 : tensor<i1>
- }
- 
- // CHECK-LABEL: "op_outfeed"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_outfeed(%arg0: tensor<f32>, %arg1: !stablehlo.token) -> !stablehlo.token {
--  //      CHECK: "vhlo.outfeed_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.outfeed_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   outfeed_config = #vhlo.string_v1<"foo">
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.token_v1) -> !vhlo.token_v1
-   %0 = "stablehlo.outfeed"(%arg0, %arg1) {
-@@ -1570,8 +1688,9 @@
- }
- 
- // CHECK-LABEL: "op_pad"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_pad(%arg0: tensor<8xf32>, %arg1: tensor<f32>) -> tensor<16xf32> {
--  //      CHECK: "vhlo.pad_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.pad_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   edge_padding_high = #vhlo.tensor_v1<dense<4> : tensor<1xi64>>,
-   // CHECK-SAME:   edge_padding_low = #vhlo.tensor_v1<dense<4> : tensor<1xi64>>,
-   // CHECK-SAME:   interior_padding = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>
-@@ -1585,36 +1704,41 @@
- }
- 
- // CHECK-LABEL: "op_popcnt"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_popcnt(%arg0: tensor<i32>) -> tensor<i32> {
--  // CHECK: "vhlo.popcnt_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-+  // CHECK: "vhlo.popcnt_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-   %0 = "stablehlo.popcnt"(%arg0) : (tensor<i32>) -> tensor<i32>
-   func.return %0 : tensor<i32>
- }
- 
- // CHECK-LABEL: "op_power"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_power(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.power_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.power_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.power"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_real_dynamic_slice"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}}, %[[ARG3:.*]]: {{.*}})
- func.func @op_real_dynamic_slice(%arg0: tensor<?xf32>, %arg1: tensor<1xindex>, %arg2: tensor<1xindex>, %arg3: tensor<1xindex>) -> tensor<?xf32> {
--  // CHECK: "vhlo.real_dynamic_slice_v1"(%arg0, %arg1, %arg2, %arg3) : (!vhlo.tensor_v1<?x!vhlo.f32_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>) -> !vhlo.tensor_v1<?x!vhlo.f32_v1>
-+  // CHECK: "vhlo.real_dynamic_slice_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]], %[[ARG3]]) : (!vhlo.tensor_v1<?x!vhlo.f32_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>) -> !vhlo.tensor_v1<?x!vhlo.f32_v1>
-   %0 = "stablehlo.real_dynamic_slice"(%arg0, %arg1, %arg2, %arg3) : (tensor<?xf32>, tensor<1xindex>, tensor<1xindex>, tensor<1xindex>) -> tensor<?xf32>
-   func.return %0 : tensor<?xf32>
- }
- 
- // CHECK-LABEL: "op_real"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_real(%arg0: tensor<complex<f32>>) -> tensor<f32> {
--  // CHECK: "vhlo.real_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.real_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.real"(%arg0) : (tensor<complex<f32>>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_recv"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_recv(%arg0: !stablehlo.token) -> (tensor<f32>, !stablehlo.token) {
--  //      CHECK: "vhlo.recv_v1"(%arg0) <{
-+  //      CHECK: "vhlo.recv_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME:   channel_type = #vhlo.integer_v1<3 : i64>,
-   // CHECK-SAME:   is_host_transfer = #vhlo.bool_v1<true>
-@@ -1627,8 +1751,9 @@
- }
- 
- // CHECK-LABEL: "op_reduce"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_reduce(%arg0: tensor<16xf32>, %arg1: tensor<f32>) -> tensor<f32> {
--  //  CHECK: "vhlo.reduce_v1"(%arg0, %arg1)
-+  //  CHECK: "vhlo.reduce_v1"(%[[ARG0]], %[[ARG1]])
-   //  CHECK:   ^[[BB:bb.*]](%[[ARG1:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>, %[[ARG2:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
-   //  CHECK:     "vhlo.return_v1"(%[[VAL1:.*]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-   //  CHECK: }) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-@@ -1643,8 +1768,9 @@
- }
- 
- // CHECK-LABEL: "op_reduce_precision"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_reduce_precision(%arg0: tensor<f32>) -> tensor<f32> {
--  //      CHECK: "vhlo.reduce_precision_v1"(%arg0) <{
-+  //      CHECK: "vhlo.reduce_precision_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   exponent_bits = #vhlo.integer_v1<8 : i32>
-   // CHECK-SAME:   mantissa_bits = #vhlo.integer_v1<10 : i32>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-@@ -1658,7 +1784,7 @@
- // CHECK_lABEL: "op_reduce_with_promotable_types"
- func.func @op_reduce_with_promotable_types(%arg0: tensor<4x4xf32>, %arg1 : tensor<f32>)
-     -> (tensor<4xf64>) {
--  //  CHECK: "vhlo.reduce_v1"(%arg0, %arg1)
-+  //  CHECK: "vhlo.reduce_v1"(%[[ARG0:.*]], %[[ARG1:.*]])
-   //  CHECK:   ^[[BB:bb.*]](%[[ARG1:arg.*]]: !vhlo.tensor_v1<!vhlo.f64_v1>, %[[ARG2:arg.*]]: !vhlo.tensor_v1<!vhlo.f64_v1>):
-   //  CHECK:     "vhlo.return_v1"(%[[VAL1:.*]]) : (!vhlo.tensor_v1<!vhlo.f64_v1>) -> ()
-   //  CHECK: }) : (!vhlo.tensor_v1<4x4x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<4x!vhlo.f64_v1>
-@@ -1673,8 +1799,9 @@
- }
- 
- // CHECK-LABEL: "op_reduce_scatter"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_reduce_scatter(%arg0: tensor<16xf32>) -> tensor<16xf32> {
--  //               CHECK: "vhlo.reduce_scatter_v1"(%arg0) <{
-+  //               CHECK: "vhlo.reduce_scatter_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0], [1]]> : tensor<2x1xi64>>,
-   //          CHECK-SAME:   scatter_dimension = #vhlo.integer_v1<0 : i64>
-@@ -1699,7 +1826,7 @@
- 
- // CHECK_lABEL: "op_reduce_scatter_with_promotable_types"
- func.func @op_reduce_scatter_with_promotable_types(%data: tensor<4x16xf32>) -> tensor<4x4xf64> {
--  //  CHECK: "vhlo.reduce_scatter_v1"(%arg0)
-+  //  CHECK: "vhlo.reduce_scatter_v1"(%[[ARG0:.*]])
-   //  CHECK:   ^[[BB:bb.*]](%[[ARG1:arg.*]]: !vhlo.tensor_v1<!vhlo.f64_v1>, %[[ARG2:arg.*]]: !vhlo.tensor_v1<!vhlo.f64_v1>):
-   //  CHECK:     "vhlo.return_v1"(%[[VAL1:.*]]) : (!vhlo.tensor_v1<!vhlo.f64_v1>) -> ()
-   //  CHECK: }) : (!vhlo.tensor_v1<4x16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<4x4x!vhlo.f64_v1>
-@@ -1716,8 +1843,9 @@
- 
- 
- // CHECK-LABEL: "op_reduce_window"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_reduce_window(%arg0: tensor<2x17x31x7xf32>, %arg1: tensor<f32>) -> tensor<2x9x16x7xf32> {
--  //               CHECK: "vhlo.reduce_window_v1"(%arg0, %arg1) <{
-+  //               CHECK: "vhlo.reduce_window_v1"(%[[ARG0]], %[[ARG1]]) <{
-   //          CHECK-SAME:   base_dilations = #vhlo.tensor_v1<dense<[1, 2, 2, 1]> : tensor<4xi64>>,
-   // CHECK-SAME{LITERAL}:   padding = #vhlo.tensor_v1<dense<[[0, 0], [2, 0], [0, 2], [0, 0]]> : tensor<4x2xi64>>,
-   //          CHECK-SAME:   window_dilations = #vhlo.tensor_v1<dense<[1, 2, 2, 1]> : tensor<4xi64>>,
-@@ -1742,11 +1870,11 @@
-   func.return %0 : tensor<2x9x16x7xf32>
- }
- 
--// CHECK_lABEL: "op_reduce_window_with_promotable_types"
-+// CHECK-LABEL: "op_reduce_window_with_promotable_types"
- func.func @op_reduce_window_with_promotable_types(%arg0: tensor<4x2xf32>,
-     %arg1: tensor<4x2xf32>, %init0: tensor<f32>, %init1: tensor<f32>) ->
-     (tensor<2x2xf64>, tensor<2x2xf32>) {
--  //  CHECK: "vhlo.reduce_window_v1"(%arg0, %arg1, %arg2, %arg3)
-+  //  CHECK: "vhlo.reduce_window_v1"(%[[ARG0:.*]], %[[ARG1:.*]], %[[ARG2:.*]], %[[ARG3:.*]])
-   //  CHECK:   ^[[BB:bb.*]](%[[ARG1:arg.*]]: !vhlo.tensor_v1<!vhlo.f64_v1>, %[[ARG2:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>, %[[ARG3:arg.*]]: !vhlo.tensor_v1<!vhlo.f64_v1>, %[[ARG4:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
-   //  CHECK:     "vhlo.return_v1"(%[[VAL1:.*]], %[[VAL2:.*]]) : (!vhlo.tensor_v1<!vhlo.f64_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-   //  CHECK: }) : (!vhlo.tensor_v1<4x2x!vhlo.f32_v1>, !vhlo.tensor_v1<4x2x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> (!vhlo.tensor_v1<2x2x!vhlo.f64_v1>, !vhlo.tensor_v1<2x2x!vhlo.f32_v1>)
-@@ -1766,8 +1894,9 @@
- }
- 
- // CHECK-LABEL: "op_remainder"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_remainder(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.remainder_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.remainder_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.remainder"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
-@@ -1787,16 +1916,18 @@
- }
- 
- // CHECK-LABEL: "op_reshape"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_reshape(%arg0: tensor<16xf32>) -> tensor<4x4xf32> {
--  // CHECK: "vhlo.reshape_v1"(%arg0) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<4x4x!vhlo.f32_v1>
-+  // CHECK: "vhlo.reshape_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<4x4x!vhlo.f32_v1>
-   %0 = "stablehlo.reshape"(%arg0) : (tensor<16xf32>) -> tensor<4x4xf32>
-   func.return %0 : tensor<4x4xf32>
- }
- 
- // CHECK-LABEL: "op_return"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_return(%arg0: tensor<i32>, %arg1: tensor<f32>) -> tensor<f32> {
--  //      CHECK: "vhlo.case_v1"(%arg0) ({
--  // CHECK-NEXT:   "vhlo.return_v1"(%arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-+  //      CHECK: "vhlo.case_v1"(%[[ARG0]]) ({
-+  // CHECK-NEXT:   "vhlo.return_v1"(%[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-   // CHECK-NEXT: }) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.case"(%arg0) ({
-     "stablehlo.return"(%arg1) : (tensor<f32>) -> ()
-@@ -1805,8 +1936,9 @@
- }
- 
- // CHECK-LABEL: "op_reverse"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_reverse(%arg0: tensor<16xf32>) -> tensor<16xf32> {
--  //      CHECK: "vhlo.reverse_v1"(%arg0) <{
-+  //      CHECK: "vhlo.reverse_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   dimensions = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x!vhlo.f32_v1>
-   %0 = "stablehlo.reverse"(%arg0) {
-@@ -1816,8 +1948,9 @@
- }
- 
- // CHECK-LABEL: "op_rng_bit_generator"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_rng_bit_generator(%arg0: tensor<f32>) -> (tensor<f32>, tensor<f32>) {
--  //      CHECK: "vhlo.rng_bit_generator_v1"(%arg0) <{
-+  //      CHECK: "vhlo.rng_bit_generator_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   rng_algorithm = #vhlo<rng_algorithm_v1 PHILOX>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>)
-   %0:2 = "stablehlo.rng_bit_generator"(%arg0) {
-@@ -1827,8 +1960,9 @@
- }
- 
- // CHECK-LABEL: "op_rng"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_rng(%arg0: tensor<f32>, %arg1: tensor<f32>, %arg2: tensor<0xindex>) -> tensor<f32> {
--  //      CHECK: "vhlo.rng_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.rng_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   rng_distribution = #vhlo<rng_distribution_v1 NORMAL>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<0x!vhlo.index_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.rng"(%arg0, %arg1, %arg2) {
-@@ -1838,29 +1972,33 @@
- }
- 
- // CHECK-LABEL: "op_round_nearest_afz"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_round_nearest_afz(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.round_nearest_afz_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.round_nearest_afz_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.round_nearest_afz"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_round_nearest_even"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_round_nearest_even(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.round_nearest_even_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.round_nearest_even_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.round_nearest_even"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_rsqrt"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_rsqrt(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.rsqrt_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.rsqrt_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.rsqrt"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_scatter"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_scatter(%arg0: tensor<200x100x300xf32>, %arg1: tensor<10x2xi32>, %arg2: tensor<10x300xf32>) -> tensor<200x100x300xf32> {
--  //      CHECK: "vhlo.scatter_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.scatter_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   index_vector_dim = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   indices_are_sorted = #vhlo.bool_v1<true>,
-   // CHECK-SAME:   inserted_window_dims = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-@@ -1893,7 +2031,7 @@
- func.func @op_scatter_with_promotable_types(%input_tensor: tensor<200x100x300xf32>,
-     %scatter_indices: tensor<10x2xi32>, %updates: tensor<10x300xf32>) ->
-       tensor<200x100x300xf64> {
--  //  CHECK: "vhlo.scatter_v1"(%arg0, %arg1, %arg2)
-+  //  CHECK: "vhlo.scatter_v1"(%[[ARG0:.*]], %[[ARG1:.*]], %[[ARG2:.*]])
-   //  CHECK:   ^[[BB:bb.*]](%[[ARG1:arg.*]]: !vhlo.tensor_v1<!vhlo.f64_v1>, %[[ARG2:arg.*]]: !vhlo.tensor_v1<!vhlo.f64_v1>):
-   //  CHECK:     "vhlo.return_v1"(%[[VAL1:.*]]) : (!vhlo.tensor_v1<!vhlo.f64_v1>) -> ()
-   //  CHECK: }) : (!vhlo.tensor_v1<200x100x300x!vhlo.f32_v1>, !vhlo.tensor_v1<10x2x!vhlo.i32_v1>, !vhlo.tensor_v1<10x300x!vhlo.f32_v1>) -> !vhlo.tensor_v1<200x100x300x!vhlo.f64_v1>
-@@ -1916,8 +2054,9 @@
- }
- 
- // CHECK-LABEL: "op_select_and_scatter"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_select_and_scatter(%arg0: tensor<10x24x24x64xf32>, %arg1: tensor<12x13x13x66xf32>, %arg2: tensor<f32>) -> tensor<10x24x24x64xf32> {
--  //      CHECK: "vhlo.select_and_scatter_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.select_and_scatter_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   padding = #vhlo.tensor_v1<dense<1> : tensor<4x2xi64>>,
-   // CHECK-SAME:   window_dimensions = #vhlo.tensor_v1<dense<[1, 2, 2, 1]> : tensor<4xi64>>,
-   // CHECK-SAME:   window_strides = #vhlo.tensor_v1<dense<[1, 2, 2, 1]> : tensor<4xi64>>
-@@ -1947,8 +2086,9 @@
- }
- 
- // CHECK-LABEL: "op_select_and_scatter_with_promotable_types"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_select_and_scatter_with_promotable_types(%arg0: tensor<10x24x24x64xf32>, %arg1: tensor<12x13x13x66xf32>, %arg2: tensor<f32>) -> tensor<10x24x24x64xf64> {
--  // CHECK: "vhlo.select_and_scatter_v1"(%arg0, %arg1, %arg2)
-+  // CHECK: "vhlo.select_and_scatter_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]])
-   // CHECK:   ^[[BB:bb.*]](%[[ARG1:arg.*]]: !vhlo.tensor_v1<!vhlo.f64_v1>, %[[ARG2:arg.*]]: !vhlo.tensor_v1<!vhlo.f64_v1>):
-   // CHECK:     %[[VAL:.*]] = "vhlo.add_v1"(%[[ARG1]], %[[ARG2]]) : (!vhlo.tensor_v1<!vhlo.f64_v1>, !vhlo.tensor_v1<!vhlo.f64_v1>) -> !vhlo.tensor_v1<!vhlo.f64_v1>
-   // CHECK:     "vhlo.return_v1"(%[[VAL]]) : (!vhlo.tensor_v1<!vhlo.f64_v1>) -> ()
-@@ -1970,15 +2110,17 @@
- }
- 
- // CHECK-LABEL: "op_select"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_select(%arg0: tensor<i1>, %arg1: tensor<f32>, %arg2: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.select_v1"(%arg0, %arg1, %arg2) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.select_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.select"(%arg0, %arg1, %arg2) : (tensor<i1>, tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_send"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_send(%arg0: tensor<f32>, %arg1: !stablehlo.token) -> !stablehlo.token {
--  //      CHECK: "vhlo.send_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.send_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME:   channel_type = #vhlo.integer_v1<2 : i64>,
-   // CHECK-SAME:   is_host_transfer = #vhlo.bool_v1<true>
-@@ -1991,8 +2133,9 @@
- }
- 
- // CHECK-LABEL: "op_set_dimension_size"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_set_dimension_size(%arg0: tensor<?xf32>, %arg1: tensor<i32>) -> tensor<16xf32> {
--  //      CHECK: "vhlo.set_dimension_size_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.set_dimension_size_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   dimension = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<?x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<16x!vhlo.f32_v1>
-   %0 = "stablehlo.set_dimension_size"(%arg0, %arg1) {
-@@ -2002,43 +2145,49 @@
- }
- 
- // CHECK-LABEL: "op_shift_left"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_shift_left(%arg0: tensor<i32>, %arg1: tensor<i32>) -> tensor<i32> {
--  // CHECK: "vhlo.shift_left_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.i32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-+  // CHECK: "vhlo.shift_left_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-   %0 = "stablehlo.shift_left"(%arg0, %arg1) : (tensor<i32>, tensor<i32>) -> tensor<i32>
-   func.return %0 : tensor<i32>
- }
- 
- // CHECK-LABEL: "op_shift_right_arithmetic"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_shift_right_arithmetic(%arg0: tensor<i32>, %arg1: tensor<i32>) -> tensor<i32> {
--  // CHECK: "vhlo.shift_right_arithmetic_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.i32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-+  // CHECK: "vhlo.shift_right_arithmetic_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-   %0 = "stablehlo.shift_right_arithmetic"(%arg0, %arg1) : (tensor<i32>, tensor<i32>) -> tensor<i32>
-   func.return %0 : tensor<i32>
- }
- 
- // CHECK-LABEL: "op_shift_right_logical"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_shift_right_logical(%arg0: tensor<i32>, %arg1: tensor<i32>) -> tensor<i32> {
--  // CHECK: "vhlo.shift_right_logical_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.i32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-+  // CHECK: "vhlo.shift_right_logical_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-   %0 = "stablehlo.shift_right_logical"(%arg0, %arg1) : (tensor<i32>, tensor<i32>) -> tensor<i32>
-   func.return %0 : tensor<i32>
- }
- 
- // CHECK-LABEL: "op_sign"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_sign(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.sign_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.sign_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.sign"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_sine"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_sine(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.sine_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.sine_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.sine"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_slice"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_slice(%arg0: tensor<16xf32>) -> tensor<4xf32> {
--  //      CHECK: "vhlo.slice_v1"(%arg0) <{
-+  //      CHECK: "vhlo.slice_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   limit_indices = #vhlo.tensor_v1<dense<4> : tensor<1xi64>>,
-   // CHECK-SAME:   start_indices = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>,
-   // CHECK-SAME:   strides = #vhlo.tensor_v1<dense<1> : tensor<1xi64>>
-@@ -2052,8 +2201,9 @@
- }
- 
- // CHECK-LABEL: "op_sort"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_sort(%arg0: tensor<16xf32>) -> tensor<16xf32> {
--  //      CHECK: "vhlo.sort_v1"(%arg0) <{
-+  //      CHECK: "vhlo.sort_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   dimension = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME:   is_stable = #vhlo.bool_v1<true>
-   // CHECK-SAME: }> ({
-@@ -2073,29 +2223,33 @@
- }
- 
- // CHECK-LABEL: "op_sqrt"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_sqrt(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.sqrt_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.sqrt_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.sqrt"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_subtract"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_subtract(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.subtract_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.subtract_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.subtract"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_tanh"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_tanh(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.tanh_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.tanh_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.tanh"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_torch_index_select"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_torch_index_select(%arg0: tensor<5x1x5xf32>, %arg1: tensor<2xi32>) ->  tensor<2x1x5xf32> {
--  //      CHECK: "vhlo.torch_index_select_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.torch_index_select_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   batch_dims = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME:   dim = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<5x1x5x!vhlo.f32_v1>, !vhlo.tensor_v1<2x!vhlo.i32_v1>) -> !vhlo.tensor_v1<2x1x5x!vhlo.f32_v1>
-@@ -2107,8 +2261,9 @@
- }
- 
- // CHECK-LABEL: "op_trace"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_trace(%arg0: tensor<f32>) {
--  //      CHECK: "vhlo.trace_v1"(%arg0) <{
-+  //      CHECK: "vhlo.trace_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   tag = #vhlo.string_v1<"foo">
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-   "stablehlo.trace"(%arg0) {
-@@ -2118,8 +2273,9 @@
- }
- 
- // CHECK-LABEL: "op_transpose"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_transpose(%arg0: tensor<16x8xf32>) ->  tensor<8x16xf32> {
--  //      CHECK: "vhlo.transpose_v1"(%arg0) <{
-+  //      CHECK: "vhlo.transpose_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   permutation = #vhlo.tensor_v1<dense<[1, 0]> : tensor<2xi64>>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<16x8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<8x16x!vhlo.f32_v1>
-   %0 = "stablehlo.transpose"(%arg0) {
-@@ -2129,8 +2285,9 @@
- }
- 
- // CHECK-LABEL: "op_triangular_solve"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_triangular_solve(%arg0: tensor<16x16xf32>, %arg1: tensor<16x16xf32>) ->  tensor<16x16xf32> {
--  //      CHECK: "vhlo.triangular_solve_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.triangular_solve_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   left_side = #vhlo.bool_v1<true>,
-   // CHECK-SAME:   lower = #vhlo.bool_v1<true>,
-   // CHECK-SAME:   transpose_a = #vhlo<transpose_v1 NO_TRANSPOSE>,
-@@ -2146,15 +2303,17 @@
- }
- 
- // CHECK-LABEL: "op_tuple"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_tuple(%arg0: tensor<f32>) -> tuple<tensor<f32>> {
--  // CHECK: "vhlo.tuple_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tuple_v1<!vhlo.tensor_v1<!vhlo.f32_v1>>
-+  // CHECK: "vhlo.tuple_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tuple_v1<!vhlo.tensor_v1<!vhlo.f32_v1>>
-   %0 = "stablehlo.tuple"(%arg0) : (tensor<f32>) -> tuple<tensor<f32>>
-   func.return %0 : tuple<tensor<f32>>
- }
- 
- // CHECK-LABEL: "op_unary_einsum"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_unary_einsum(%arg0: tensor<8x16xf32>) -> tensor<8xf32> {
--  //      CHECK: "vhlo.unary_einsum_v1"(%arg0) <{
-+  //      CHECK: "vhlo.unary_einsum_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   einsum_config = #vhlo.string_v1<"ab->a">
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<8x16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<8x!vhlo.f32_v1>
-   %0 = "stablehlo.unary_einsum"(%arg0) {
-@@ -2164,22 +2323,25 @@
- }
- 
- // CHECK-LABEL: "op_uniform_dequantize"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_uniform_dequantize(%arg0: tensor<!quant.uniform<i8:f32, 34.0:16>>) -> tensor<f32> {
--  // CHECK: "vhlo.uniform_dequantize_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.uniform_dequantize_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.uniform_dequantize"(%arg0) : (tensor<!quant.uniform<i8:f32, 34.0:16>>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_uniform_quantize"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_uniform_quantize(%arg0: tensor<f32>) -> tensor<!quant.uniform<i8:f32, 34.0:16>> {
--  // CHECK: "vhlo.uniform_quantize_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>
-+  // CHECK: "vhlo.uniform_quantize_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>
-   %0 = "stablehlo.uniform_quantize"(%arg0) : (tensor<f32>) -> tensor<!quant.uniform<i8:f32, 34.0:16>>
-   func.return %0 : tensor<!quant.uniform<i8:f32, 34.0:16>>
- }
- 
- // CHECK-LABEL: "op_while"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_while(%arg0: tensor<i1>) -> tensor<i1> {
--  //      CHECK: "vhlo.while_v1"(%arg0) ({
-+  //      CHECK: "vhlo.while_v1"(%[[ARG0]]) ({
-   // CHECK-NEXT:   ^[[BB:bb.*]](%[[ARG1:arg.*]]: !vhlo.tensor_v1<!vhlo.bool_v1>):
-   // CHECK-NEXT:     "vhlo.return_v1"(%[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>) -> ()
-   // CHECK-NEXT:   }, {
-@@ -2197,8 +2359,9 @@
- }
- 
- // CHECK-LABEL: "op_xor"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_xor(%arg0: tensor<i1>, %arg1: tensor<i1>) -> tensor<i1> {
--  // CHECK: "vhlo.xor_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-+  // CHECK: "vhlo.xor_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-   %0 = "stablehlo.xor"(%arg0, %arg1) : (tensor<i1>, tensor<i1>) -> tensor<i1>
-   func.return %0 : tensor<i1>
- }
-@@ -2206,190 +2369,217 @@
- // ============ TYPES ============
- 
- // CHECK-LABEL: "type_i1"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_i1(%arg0: tensor<i1>, %arg1: tensor<i1>) -> tensor<i1> {
--  // CHECK: "vhlo.and_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-+  // CHECK: "vhlo.and_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-   %0 = "stablehlo.and"(%arg0, %arg1) : (tensor<i1>, tensor<i1>) -> tensor<i1>
-   func.return %0 : tensor<i1>
- }
- 
- // CHECK-LABEL: "type_i4"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_i4(%arg0: tensor<i4>, %arg1: tensor<i4>) -> tensor<i4> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.i4_v1>, !vhlo.tensor_v1<!vhlo.i4_v1>) -> !vhlo.tensor_v1<!vhlo.i4_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i4_v1>, !vhlo.tensor_v1<!vhlo.i4_v1>) -> !vhlo.tensor_v1<!vhlo.i4_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<i4>, tensor<i4>) -> tensor<i4>
-   func.return %0 : tensor<i4>
- }
- 
- // CHECK-LABEL: "type_i8"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_i8(%arg0: tensor<i8>, %arg1: tensor<i8>) -> tensor<i8> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.i8_v1>, !vhlo.tensor_v1<!vhlo.i8_v1>) -> !vhlo.tensor_v1<!vhlo.i8_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i8_v1>, !vhlo.tensor_v1<!vhlo.i8_v1>) -> !vhlo.tensor_v1<!vhlo.i8_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<i8>, tensor<i8>) -> tensor<i8>
-   func.return %0 : tensor<i8>
- }
- 
- // CHECK-LABEL: "type_i16"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_i16(%arg0: tensor<i16>, %arg1: tensor<i16>) -> tensor<i16> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.i16_v1>, !vhlo.tensor_v1<!vhlo.i16_v1>) -> !vhlo.tensor_v1<!vhlo.i16_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i16_v1>, !vhlo.tensor_v1<!vhlo.i16_v1>) -> !vhlo.tensor_v1<!vhlo.i16_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<i16>, tensor<i16>) -> tensor<i16>
-   func.return %0 : tensor<i16>
- }
- 
- // CHECK-LABEL: "type_i32"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_i32(%arg0: tensor<i32>, %arg1: tensor<i32>) -> tensor<i32> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.i32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<i32>, tensor<i32>) -> tensor<i32>
-   func.return %0 : tensor<i32>
- }
- 
- // CHECK-LABEL: "type_i64"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_i64(%arg0: tensor<i64>, %arg1: tensor<i64>) -> tensor<i64> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.i64_v1>, !vhlo.tensor_v1<!vhlo.i64_v1>) -> !vhlo.tensor_v1<!vhlo.i64_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i64_v1>, !vhlo.tensor_v1<!vhlo.i64_v1>) -> !vhlo.tensor_v1<!vhlo.i64_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<i64>, tensor<i64>) -> tensor<i64>
-   func.return %0 : tensor<i64>
- }
- 
- // CHECK-LABEL: "type_ui4"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_ui4(%arg0: tensor<ui4>, %arg1: tensor<ui4>) -> tensor<ui4> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.ui4_v1>, !vhlo.tensor_v1<!vhlo.ui4_v1>) -> !vhlo.tensor_v1<!vhlo.ui4_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.ui4_v1>, !vhlo.tensor_v1<!vhlo.ui4_v1>) -> !vhlo.tensor_v1<!vhlo.ui4_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<ui4>, tensor<ui4>) -> tensor<ui4>
-   func.return %0 : tensor<ui4>
- }
- 
- // CHECK-LABEL: "type_ui8"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_ui8(%arg0: tensor<ui8>, %arg1: tensor<ui8>) -> tensor<ui8> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.ui8_v1>, !vhlo.tensor_v1<!vhlo.ui8_v1>) -> !vhlo.tensor_v1<!vhlo.ui8_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.ui8_v1>, !vhlo.tensor_v1<!vhlo.ui8_v1>) -> !vhlo.tensor_v1<!vhlo.ui8_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<ui8>, tensor<ui8>) -> tensor<ui8>
-   func.return %0 : tensor<ui8>
- }
- 
- // CHECK-LABEL: "type_ui16"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_ui16(%arg0: tensor<ui16>, %arg1: tensor<ui16>) -> tensor<ui16> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.ui16_v1>, !vhlo.tensor_v1<!vhlo.ui16_v1>) -> !vhlo.tensor_v1<!vhlo.ui16_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.ui16_v1>, !vhlo.tensor_v1<!vhlo.ui16_v1>) -> !vhlo.tensor_v1<!vhlo.ui16_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<ui16>, tensor<ui16>) -> tensor<ui16>
-   func.return %0 : tensor<ui16>
- }
- 
- // CHECK-LABEL: "type_ui32"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_ui32(%arg0: tensor<ui32>, %arg1: tensor<ui32>) -> tensor<ui32> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.ui32_v1>, !vhlo.tensor_v1<!vhlo.ui32_v1>) -> !vhlo.tensor_v1<!vhlo.ui32_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.ui32_v1>, !vhlo.tensor_v1<!vhlo.ui32_v1>) -> !vhlo.tensor_v1<!vhlo.ui32_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<ui32>, tensor<ui32>) -> tensor<ui32>
-   func.return %0 : tensor<ui32>
- }
- 
- // CHECK-LABEL: "type_ui64"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_ui64(%arg0: tensor<ui64>, %arg1: tensor<ui64>) -> tensor<ui64> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.ui64_v1>, !vhlo.tensor_v1<!vhlo.ui64_v1>) -> !vhlo.tensor_v1<!vhlo.ui64_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.ui64_v1>, !vhlo.tensor_v1<!vhlo.ui64_v1>) -> !vhlo.tensor_v1<!vhlo.ui64_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<ui64>, tensor<ui64>) -> tensor<ui64>
-   func.return %0 : tensor<ui64>
- }
- 
- // CHECK-LABEL: "type_f8E4M3FN"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_f8E4M3FN(%arg0: tensor<f8E4M3FN>, %arg1: tensor<f8E4M3FN>) -> tensor<f8E4M3FN> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f8E4M3FN_v1>, !vhlo.tensor_v1<!vhlo.f8E4M3FN_v1>) -> !vhlo.tensor_v1<!vhlo.f8E4M3FN_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f8E4M3FN_v1>, !vhlo.tensor_v1<!vhlo.f8E4M3FN_v1>) -> !vhlo.tensor_v1<!vhlo.f8E4M3FN_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f8E4M3FN>, tensor<f8E4M3FN>) -> tensor<f8E4M3FN>
-   func.return %0 : tensor<f8E4M3FN>
- }
- 
- // CHECK-LABEL: "type_f8E5M2"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_f8E5M2(%arg0: tensor<f8E5M2>, %arg1: tensor<f8E5M2>) -> tensor<f8E5M2> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f8E5M2_v1>, !vhlo.tensor_v1<!vhlo.f8E5M2_v1>) -> !vhlo.tensor_v1<!vhlo.f8E5M2_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f8E5M2_v1>, !vhlo.tensor_v1<!vhlo.f8E5M2_v1>) -> !vhlo.tensor_v1<!vhlo.f8E5M2_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f8E5M2>, tensor<f8E5M2>) -> tensor<f8E5M2>
-   func.return %0 : tensor<f8E5M2>
- }
- 
- // CHECK-LABEL: "type_f8E4M3FNUZ"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_f8E4M3FNUZ(%arg0: tensor<f8E4M3FNUZ>, %arg1: tensor<f8E4M3FNUZ>) -> tensor<f8E4M3FNUZ> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f8E4M3FNUZ_v1>, !vhlo.tensor_v1<!vhlo.f8E4M3FNUZ_v1>) -> !vhlo.tensor_v1<!vhlo.f8E4M3FNUZ_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f8E4M3FNUZ_v1>, !vhlo.tensor_v1<!vhlo.f8E4M3FNUZ_v1>) -> !vhlo.tensor_v1<!vhlo.f8E4M3FNUZ_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f8E4M3FNUZ>, tensor<f8E4M3FNUZ>) -> tensor<f8E4M3FNUZ>
-   func.return %0 : tensor<f8E4M3FNUZ>
- }
- 
- // CHECK-LABEL: "type_f8E4M3B11FNUZ"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_f8E4M3B11FNUZ(%arg0: tensor<f8E4M3B11FNUZ>, %arg1: tensor<f8E4M3B11FNUZ>) -> tensor<f8E4M3B11FNUZ> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f8E4M3B11FNUZ_v1>, !vhlo.tensor_v1<!vhlo.f8E4M3B11FNUZ_v1>) -> !vhlo.tensor_v1<!vhlo.f8E4M3B11FNUZ_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f8E4M3B11FNUZ_v1>, !vhlo.tensor_v1<!vhlo.f8E4M3B11FNUZ_v1>) -> !vhlo.tensor_v1<!vhlo.f8E4M3B11FNUZ_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f8E4M3B11FNUZ>, tensor<f8E4M3B11FNUZ>) -> tensor<f8E4M3B11FNUZ>
-   func.return %0 : tensor<f8E4M3B11FNUZ>
- }
- 
- // CHECK-LABEL: "type_f8E5M2FNUZ"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_f8E5M2FNUZ(%arg0: tensor<f8E5M2FNUZ>, %arg1: tensor<f8E5M2FNUZ>) -> tensor<f8E5M2FNUZ> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f8E5M2FNUZ_v1>, !vhlo.tensor_v1<!vhlo.f8E5M2FNUZ_v1>) -> !vhlo.tensor_v1<!vhlo.f8E5M2FNUZ_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f8E5M2FNUZ_v1>, !vhlo.tensor_v1<!vhlo.f8E5M2FNUZ_v1>) -> !vhlo.tensor_v1<!vhlo.f8E5M2FNUZ_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f8E5M2FNUZ>, tensor<f8E5M2FNUZ>) -> tensor<f8E5M2FNUZ>
-   func.return %0 : tensor<f8E5M2FNUZ>
- }
- 
- // CHECK-LABEL: "type_bf16"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_bf16(%arg0: tensor<bf16>, %arg1: tensor<bf16>) -> tensor<bf16> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.bf16_v1>, !vhlo.tensor_v1<!vhlo.bf16_v1>) -> !vhlo.tensor_v1<!vhlo.bf16_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.bf16_v1>, !vhlo.tensor_v1<!vhlo.bf16_v1>) -> !vhlo.tensor_v1<!vhlo.bf16_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<bf16>, tensor<bf16>) -> tensor<bf16>
-   func.return %0 : tensor<bf16>
- }
- 
- // CHECK-LABEL: "type_f16"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_f16(%arg0: tensor<f16>, %arg1: tensor<f16>) -> tensor<f16> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f16_v1>, !vhlo.tensor_v1<!vhlo.f16_v1>) -> !vhlo.tensor_v1<!vhlo.f16_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f16_v1>, !vhlo.tensor_v1<!vhlo.f16_v1>) -> !vhlo.tensor_v1<!vhlo.f16_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f16>, tensor<f16>) -> tensor<f16>
-   func.return %0 : tensor<f16>
- }
- 
- // CHECK-LABEL: "type_f32"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_f32(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "type_f64"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_f64(%arg0: tensor<f64>, %arg1: tensor<f64>) -> tensor<f64> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f64_v1>, !vhlo.tensor_v1<!vhlo.f64_v1>) -> !vhlo.tensor_v1<!vhlo.f64_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f64_v1>, !vhlo.tensor_v1<!vhlo.f64_v1>) -> !vhlo.tensor_v1<!vhlo.f64_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f64>, tensor<f64>) -> tensor<f64>
-   func.return %0 : tensor<f64>
- }
- 
- // CHECK-LABEL: "type_complex_f32"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_complex_f32(%arg0: tensor<complex<f32>>, %arg1: tensor<complex<f32>>) -> tensor<complex<f32>> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>, !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>) -> !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>, !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>) -> !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<complex<f32>>, tensor<complex<f32>>) -> tensor<complex<f32>>
-   func.return %0 : tensor<complex<f32>>
- }
- 
- // CHECK-LABEL: "type_complex_f64"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_complex_f64(%arg0: tensor<complex<f64>>, %arg1: tensor<complex<f64>>) -> tensor<complex<f64>> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f64_v1>>, !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f64_v1>>) -> !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f64_v1>>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f64_v1>>, !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f64_v1>>) -> !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f64_v1>>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<complex<f64>>, tensor<complex<f64>>) -> tensor<complex<f64>>
-   func.return %0 : tensor<complex<f64>>
- }
- 
- // CHECK-LABEL: "type_dynamism_ranked"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @type_dynamism_ranked(%arg0: tensor<?xf32>) -> tensor<?xf32> {
--  // CHECK: "vhlo.abs_v1"(%arg0) : (!vhlo.tensor_v1<?x!vhlo.f32_v1>) -> !vhlo.tensor_v1<?x!vhlo.f32_v1>
-+  // CHECK: "vhlo.abs_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<?x!vhlo.f32_v1>) -> !vhlo.tensor_v1<?x!vhlo.f32_v1>
-   %0 = "stablehlo.abs"(%arg0) : (tensor<?xf32>) -> tensor<?xf32>
-   func.return %0 : tensor<?xf32>
- }
- 
- // CHECK-LABEL: "type_quantization"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_quantization(%arg0: tensor<!quant.uniform<i8:f32, 34.0:16>>, %arg1: tensor<!quant.uniform<i8:f32, 34.0:16>>) -> tensor<!quant.uniform<i8:f32, 34.0:16>> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>, !vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>) -> !vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>, !vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>) -> !vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<!quant.uniform<i8:f32, 34.0:16>>, tensor<!quant.uniform<i8:f32, 34.0:16>>) -> tensor<!quant.uniform<i8:f32, 34.0:16>>
-   func.return %0 : tensor<!quant.uniform<i8:f32, 34.0:16>>
- }
- 
- //       CHECK: function_type = #vhlo.type_v1<!vhlo.func_v1<(!vhlo.token_v1) -> !vhlo.token_v1>>
- // CHECK-LABEL: "type_token_callee"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @type_token_callee(%arg0: !stablehlo.token) -> !stablehlo.token {
--  // CHECK: "vhlo.return_v1"(%arg0) : (!vhlo.token_v1) -> ()
-+  // CHECK: "vhlo.return_v1"(%[[ARG0]]) : (!vhlo.token_v1) -> ()
-   return %arg0 : !stablehlo.token
- }
- 
- //       CHECK: function_type = #vhlo.type_v1<!vhlo.func_v1<(!vhlo.token_v1) -> !vhlo.token_v1>>
- // CHECK-LABEL: "type_token_caller"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @type_token_caller(%arg0: !stablehlo.token) -> !stablehlo.token {
--  // CHECK:      "vhlo.call_v1"(%arg0) <{callee = #vhlo.string_v1<"type_token_callee">}
-+  // CHECK:      "vhlo.call_v1"(%[[ARG0]]) <{callee = #vhlo.string_v1<"type_token_callee">}
-   // CHECK-SAME: (!vhlo.token_v1) -> !vhlo.token_v1
-   %0 = func.call @type_token_callee(%arg0) : (!stablehlo.token) -> !stablehlo.token
-   return %0 : !stablehlo.token
- }
- 
- // CHECK-LABEL: "type_tuple"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @type_tuple(%arg0: tuple<tensor<f32>>) -> tuple<!stablehlo.token> {
-   %0 = "stablehlo.custom_call"(%arg0) {
-     call_target_name = "foo"
-diff --ruN a/stablehlo/stablehlo/tests/vhlo/stablehlo_legalize_to_vhlo.0_18_0.mlir b/stablehlo/stablehlo/tests/vhlo/stablehlo_legalize_to_vhlo.0_18_0.mlir
---- stablehlo/stablehlo/tests/vhlo/stablehlo_legalize_to_vhlo.0_18_0.mlir
-+++ stablehlo/stablehlo/tests/vhlo/stablehlo_legalize_to_vhlo.0_18_0.mlir
-@@ -13,6 +13,7 @@
- // ============ ATTRIBUTES ============
- 
- // CHECK-LABEL: "attr_comparison_direction_eq"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_direction_eq(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     // CHECK: comparison_direction = #vhlo<comparison_direction_v1 EQ>
-@@ -22,6 +23,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_direction_ne"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_direction_ne(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     // CHECK: comparison_direction = #vhlo<comparison_direction_v1 NE>
-@@ -31,6 +33,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_direction_ge"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_direction_ge(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     // CHECK: comparison_direction = #vhlo<comparison_direction_v1 GE>
-@@ -40,6 +43,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_direction_gt"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_direction_gt(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     // CHECK: comparison_direction = #vhlo<comparison_direction_v1 GT>
-@@ -49,6 +53,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_direction_le"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_direction_le(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     // CHECK: comparison_direction = #vhlo<comparison_direction_v1 LE>
-@@ -58,6 +63,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_direction_lt"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_direction_lt(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     // CHECK: comparison_direction = #vhlo<comparison_direction_v1 LT>
-@@ -67,6 +73,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_type_notype"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_type_notype(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     comparison_direction = #stablehlo<comparison_direction EQ>
-@@ -76,6 +83,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_type_float"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_type_float(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     comparison_direction = #stablehlo<comparison_direction EQ>,
-@@ -86,6 +94,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_type_totalorder"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_type_totalorder(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     comparison_direction = #stablehlo<comparison_direction EQ>,
-@@ -96,6 +105,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_type_signed"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_type_signed(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     comparison_direction = #stablehlo<comparison_direction EQ>,
-@@ -106,6 +116,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_type_unsigned"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_type_unsigned(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     comparison_direction = #stablehlo<comparison_direction EQ>,
-@@ -118,6 +129,7 @@
- // ConvDimensionNumbers aka #stablehlo.conv is covered below.
- 
- // CHECK-LABEL: "attr_custom_call_api_version_unspecified"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_custom_call_api_version_unspecified(%arg0: tensor<f32>) -> tensor<f32> {
-   %0 = "stablehlo.custom_call"(%arg0) {
-     call_target_name = "foo",
-@@ -128,6 +140,7 @@
- }
- 
- // CHECK-LABEL: "attr_custom_call_api_version_original"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_custom_call_api_version_original(%arg0: tensor<f32>) -> tensor<f32> {
-   %0 = "stablehlo.custom_call"(%arg0) {
-     call_target_name = "foo",
-@@ -138,6 +151,7 @@
- }
- 
- // CHECK-LABEL: "attr_custom_call_api_version_status_returning"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_custom_call_api_version_status_returning(%arg0: tensor<f32>) -> tensor<f32> {
-   %0 = "stablehlo.custom_call"(%arg0) {
-     call_target_name = "foo",
-@@ -148,6 +162,7 @@
- }
- 
- // CHECK-LABEL: "attr_custom_call_api_version_status_returning_unified"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_custom_call_api_version_status_returning_unified(%arg0: tensor<f32>) -> tensor<f32> {
-   %0 = "stablehlo.custom_call"(%arg0) {
-     call_target_name = "foo",
-@@ -166,6 +181,7 @@
- // DotDimensionNumbers aka #stablehlo.dot is covered below.
- 
- // CHECK-LABEL: "attr_fft_type_fft"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_fft_type_fft(%arg0: tensor<16xcomplex<f32>>) -> tensor<16xcomplex<f32>> {
-   %0 = "stablehlo.fft"(%arg0) {
-     // CHECK: fft_type = #vhlo<fft_type_v1 FFT>
-@@ -176,6 +192,7 @@
- }
- 
- // CHECK-LABEL: "attr_fft_type_ifft"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_fft_type_ifft(%arg0: tensor<16xcomplex<f32>>) -> tensor<16xcomplex<f32>> {
-   %0 = "stablehlo.fft"(%arg0) {
-     // CHECK: fft_type = #vhlo<fft_type_v1 IFFT>
-@@ -186,6 +203,7 @@
- }
- 
- // CHECK-LABEL: "attr_fft_type_rfft"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_fft_type_rfft(%arg0: tensor<16xf32>) -> tensor<9xcomplex<f32>> {
-   %0 = "stablehlo.fft"(%arg0) {
-     // CHECK: fft_type = #vhlo<fft_type_v1 RFFT>
-@@ -196,6 +214,7 @@
- }
- 
- // CHECK-LABEL: "attr_fft_type_irfft"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_fft_type_irfft(%arg0: tensor<9xcomplex<f32>>) -> tensor<16xf32> {
-   %0 = "stablehlo.fft"(%arg0) {
-     // CHECK: fft_type = #vhlo<fft_type_v1 IRFFT>
-@@ -208,6 +227,7 @@
- // GatherDimensionNumbers aka #stablehlo.gather is covered below.
- 
- // CHECK-LABEL: "attr_precision_config_default"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_precision_config_default(%arg0: tensor<8x16xf32>, %arg1: tensor<16x8xf32>) -> tensor<8x8xf32> {
-   %0 = "stablehlo.dot"(%arg0, %arg1) {
-     // CHECK: precision_config = #vhlo.array_v1<[#vhlo<precision_v1 DEFAULT>, #vhlo<precision_v1 DEFAULT>]>
-@@ -216,6 +236,7 @@
- }
- 
- // CHECK-LABEL: "attr_precision_config_high"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_precision_config_high(%arg0: tensor<8x16xf32>, %arg1: tensor<16x8xf32>) -> tensor<8x8xf32> {
-   %0 = "stablehlo.dot"(%arg0, %arg1) {
-     // CHECK: precision_config = #vhlo.array_v1<[#vhlo<precision_v1 HIGH>, #vhlo<precision_v1 HIGH>]>
-@@ -225,6 +246,7 @@
- }
- 
- // CHECK-LABEL: "attr_precision_config_highest"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_precision_config_highest(%arg0: tensor<8x16xf32>, %arg1: tensor<16x8xf32>) -> tensor<8x8xf32> {
-   %0 = "stablehlo.dot"(%arg0, %arg1) {
-     // CHECK: precision_config = #vhlo.array_v1<[#vhlo<precision_v1 HIGHEST>, #vhlo<precision_v1 HIGHEST>]>
-@@ -234,6 +256,7 @@
- }
- 
- // CHECK-LABEL: "attr_rng_algorithm_default"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_rng_algorithm_default(%arg0: tensor<f32>) -> (tensor<f32>, tensor<f32>) {
-   %0:2 = "stablehlo.rng_bit_generator"(%arg0) {
-     // CHECK: rng_algorithm = #vhlo<rng_algorithm_v1 DEFAULT>
-@@ -243,6 +266,7 @@
- }
- 
- // CHECK-LABEL: "attr_rng_algorithm_three_fry"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_rng_algorithm_three_fry(%arg0: tensor<f32>) -> (tensor<f32>, tensor<f32>) {
-   %0:2 = "stablehlo.rng_bit_generator"(%arg0) {
-     // CHECK: rng_algorithm = #vhlo<rng_algorithm_v1 THREE_FRY>
-@@ -252,6 +276,7 @@
- }
- 
- // CHECK-LABEL: "attr_rng_algorithm_philox"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_rng_algorithm_philox(%arg0: tensor<f32>) -> (tensor<f32>, tensor<f32>) {
-   %0:2 = "stablehlo.rng_bit_generator"(%arg0) {
-     // CHECK: rng_algorithm = #vhlo<rng_algorithm_v1 PHILOX>
-@@ -261,6 +286,7 @@
- }
- 
- // CHECK-LABEL: "attr_rng_distribution_uniform"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @attr_rng_distribution_uniform(%arg0: tensor<f32>, %arg1: tensor<f32>, %arg2: tensor<0xindex>) -> tensor<f32> {
-   %0 = "stablehlo.rng"(%arg0, %arg1, %arg2) {
-     // CHECK: rng_distribution = #vhlo<rng_distribution_v1 UNIFORM>
-@@ -270,6 +296,7 @@
- }
- 
- // CHECK-LABEL: "attr_rng_distribution_normal"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @attr_rng_distribution_normal(%arg0: tensor<f32>, %arg1: tensor<f32>, %arg2: tensor<0xindex>) -> tensor<f32> {
-   %0 = "stablehlo.rng"(%arg0, %arg1, %arg2) {
-     // CHECK: rng_distribution = #vhlo<rng_distribution_v1 NORMAL>
-@@ -281,6 +308,7 @@
- // ScatterDimensionNumbers aka #stablehlo.scatter is covered below.
- 
- // CHECK-LABEL: "attr_transpose_no_transpose"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_transpose_no_transpose(%arg0: tensor<16x16xf32>, %arg1: tensor<16x16xf32>) ->  tensor<16x16xf32> {
-   %0 = "stablehlo.triangular_solve"(%arg0, %arg1) {
-     left_side = true,
-@@ -293,6 +321,7 @@
- }
- 
- // CHECK-LABEL: "attr_transpose_transpose"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_transpose_transpose(%arg0: tensor<16x16xf32>, %arg1: tensor<16x16xf32>) ->  tensor<16x16xf32> {
-   %0 = "stablehlo.triangular_solve"(%arg0, %arg1) {
-     left_side = true,
-@@ -305,6 +334,7 @@
- }
- 
- // CHECK-LABEL: "attr_transpose_adjoint"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_transpose_adjoint(%arg0: tensor<16x16xf32>, %arg1: tensor<16x16xf32>) ->  tensor<16x16xf32> {
-   %0 = "stablehlo.triangular_solve"(%arg0, %arg1) {
-     left_side = true,
-@@ -319,10 +349,9 @@
- // TypeExtensionsAttr aka #stablehlo.type_extensions is covered below.
- 
- // CHECK-LABEL: "attr_type_extensions_bounds"
--func.func @attr_type_extensions_bounds(
--    %arg0: tensor<?x?xf32, #stablehlo.type_extensions<bounds = [16, ?]>>)
--    -> tensor<?x?xf32, #stablehlo.type_extensions<bounds = [16, ?]>> {
--  // CHECK: "vhlo.return_v1"(%arg0) : (!vhlo.tensor_v1<?x?x!vhlo.f32_v1, #vhlo.type_extensions_v1<bounds = [16, ?]>>) -> ()
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @attr_type_extensions_bounds(%arg0: tensor<?x?xf32, #stablehlo.type_extensions<bounds = [16, ?]>>) -> tensor<?x?xf32, #stablehlo.type_extensions<bounds = [16, ?]>> {
-+  // CHECK: "vhlo.return_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<?x?x!vhlo.f32_v1, #vhlo.type_extensions_v1<bounds = [16, ?]>>) -> ()
-   func.return %arg0 : tensor<?x?xf32, #stablehlo.type_extensions<bounds = [16, ?]>>
- }
- 
-@@ -330,8 +359,9 @@
- // ============ DEFAULTS ============
- 
- // CHECK-LABEL: "default_all_gather"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_all_gather(%arg0: tensor<16x8xf32>) -> tensor<16x16xf32> {
--  //               CHECK: "vhlo.all_gather_v1"(%arg0) <{
-+  //               CHECK: "vhlo.all_gather_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   all_gather_dim = #vhlo.integer_v1<1 : i64>
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0], [1]]> : tensor<2x1xi64>>,
-@@ -345,8 +375,9 @@
- }
- 
- // CHECK-LABEL: "default_all_reduce"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_all_reduce(%arg0: tensor<f32>) -> tensor<f32> {
--  //               CHECK: "vhlo.all_reduce_v1"(%arg0)
-+  //               CHECK: "vhlo.all_reduce_v1"(%[[ARG0]])
-   //          CHECK-SAME: <{
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0], [1]]> : tensor<2x1xi64>>,
-@@ -368,8 +399,9 @@
- }
- 
- // CHECK-LABEL: "default_all_to_all"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_all_to_all(%arg0: tensor<4x16xf32>) -> tensor<16x4xf32> {
--  //               CHECK: "vhlo.all_to_all_v1"(%arg0) <{
-+  //               CHECK: "vhlo.all_to_all_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   //          CHECK-SAME:   concat_dimension = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0, 1, 2, 3]]> : tensor<1x4xi64>>,
-@@ -386,8 +418,9 @@
- }
- 
- // CHECK-LABEL: "default_cholesky"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_cholesky(%arg0: tensor<1x16x16xf32>) -> tensor<1x16x16xf32> {
--  //      CHECK: "vhlo.cholesky_v1"(%arg0) <{
-+  //      CHECK: "vhlo.cholesky_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   lower = #vhlo.bool_v1<false>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<1x16x16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<1x16x16x!vhlo.f32_v1>
-   %0 = "stablehlo.cholesky"(%arg0) : (tensor<1x16x16xf32>) -> tensor<1x16x16xf32>
-@@ -395,8 +428,9 @@
- }
- 
- // CHECK-LABEL: "default_collective_permute"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_collective_permute(%arg0: tensor<16x8xf32>) -> tensor<16x8xf32> {
--  //               CHECK: "vhlo.collective_permute_v1"(%arg0) <{
-+  //               CHECK: "vhlo.collective_permute_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME{LITERAL}:   source_target_pairs = #vhlo.tensor_v1<dense<[[0, 1], [1, 2], [2, 3]]> : tensor<3x2xi64>>
-   //          CHECK-SAME: }> : (!vhlo.tensor_v1<16x8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x8x!vhlo.f32_v1>
-@@ -407,8 +441,9 @@
- }
- 
- // CHECK-LABEL: "default_collective_broadcast"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_collective_broadcast(%arg0: tensor<16x8xf32>) -> tensor<16x8xf32> {
--  //               CHECK: "vhlo.collective_broadcast_v1"(%arg0) <{
-+  //               CHECK: "vhlo.collective_broadcast_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0, 1]]> : tensor<1x2xi64>>
-   //          CHECK-SAME: }> : (!vhlo.tensor_v1<16x8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x8x!vhlo.f32_v1>
-@@ -419,8 +454,9 @@
- }
- 
- // CHECK-LABEL: "default_compare"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @default_compare(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
--  //      CHECK: "vhlo.compare_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.compare_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   compare_type = #vhlo<comparison_type_v1 NOTYPE>,
-   // CHECK-SAME:   comparison_direction = #vhlo<comparison_direction_v1 EQ>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-@@ -431,8 +467,9 @@
- }
- 
- // CHECK-LABEL: "default_convolution"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @default_convolution(%arg0: tensor<1x8x8x207xf32>, %arg1: tensor<3x3x207x16xf32>) -> tensor<1x6x6x16xf32> {
--  //      CHECK: "vhlo.convolution_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.convolution_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   batch_group_count = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   feature_group_count = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   input_batch_dimension = #vhlo.integer_v1<0 : i64>,
-@@ -460,8 +497,9 @@
- }
- 
- // CHECK-LABEL: "default_custom_call"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_custom_call(%arg0: tensor<f32>) -> tensor<f32> {
--  //      CHECK: "vhlo.custom_call_v1"(%arg0) <{
-+  //      CHECK: "vhlo.custom_call_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   api_version = #vhlo<api_version_v1 API_VERSION_ORIGINAL>,
-   // CHECK-SAME:   backend_config = #vhlo.string_v1<"">,
-   // CHECK-SAME:   call_target_name = #vhlo.string_v1<"foo">,
-@@ -478,8 +516,9 @@
- }
- 
- // CHECK-LABEL: "default_dot_general"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @default_dot_general(%arg0: tensor<8x8x16xf32>, %arg1: tensor<8x16x8xf32>) -> tensor<8x8x8xf32> {
--  //      CHECK: "vhlo.dot_general_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.dot_general_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   lhs_batching_dimensions = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>,
-   // CHECK-SAME:   lhs_contracting_dimensions = #vhlo.tensor_v1<dense<2> : tensor<1xi64>>,
-   // CHECK-SAME:   precision_config = #vhlo.array_v1<[#vhlo<precision_v1 DEFAULT>, #vhlo<precision_v1 DEFAULT>]>,
-@@ -498,8 +537,9 @@
- }
- 
- // CHECK-LABEL: "default_dot"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @default_dot(%arg0: tensor<8x16xf32>, %arg1: tensor<16x8xf32>) -> tensor<8x8xf32> {
--  //      CHECK: "vhlo.dot_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.dot_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   precision_config = #vhlo.array_v1<[#vhlo<precision_v1 DEFAULT>, #vhlo<precision_v1 DEFAULT>]>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<8x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<8x8x!vhlo.f32_v1>
-   %0 = "stablehlo.dot"(%arg0, %arg1) : (tensor<8x16xf32>, tensor<16x8xf32>) -> tensor<8x8xf32>
-@@ -507,8 +547,9 @@
- }
- 
- // CHECK-LABEL: "default_dynamic_broadcast_in_dim"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @default_dynamic_broadcast_in_dim(%arg0: tensor<?x?xf32>, %arg1: tensor<2xindex>) -> tensor<?x?xf32> {
--  //      CHECK: "vhlo.dynamic_broadcast_in_dim_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.dynamic_broadcast_in_dim_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   broadcast_dimensions = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-   // CHECK-SAME:   known_expanding_dimensions = #vhlo.tensor_v1<dense<> : tensor<0xi64>>,
-   // CHECK-SAME:   known_nonexpanding_dimensions = #vhlo.tensor_v1<dense<> : tensor<0xi64>>
-@@ -520,8 +561,9 @@
- }
- 
- // CHECK-LABEL: "default_dynamic_conv"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @default_dynamic_conv(%arg0: tensor<1x8x8x207xf32>, %arg1: tensor<3x3x207x16xf32>, %arg2: tensor<4xi32>) -> tensor<1x?x?x16xf32> {
--  //      CHECK: "vhlo.dynamic_conv_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.dynamic_conv_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   batch_group_count = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   feature_group_count = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   input_batch_dimension = #vhlo.integer_v1<0 : i64>,
-@@ -549,8 +591,9 @@
- }
- 
- // CHECK-LABEL: "default_dynamic_gather"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @default_dynamic_gather(%arg0 : tensor<2x4x9xf32>, %arg1 : tensor<1x5x2xi32>, %arg2 : tensor<3xi32>) -> tensor<1x5x8xf32> {
--  //      CHECK: "vhlo.dynamic_gather_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.dynamic_gather_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   collapsed_slice_dims = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-   // CHECK-SAME:   index_vector_dim = #vhlo.integer_v1<2 : i64>,
-   // CHECK-SAME:   indices_are_sorted = #vhlo.bool_v1<false>,
-@@ -576,15 +619,16 @@
-   // CHECK-SAME:   sym_name = #vhlo.string_v1<"default_func">,
-   // CHECK-SAME:   sym_visibility = #vhlo.string_v1<"">
-   // CHECK-SAME: }> ({
--  // CHECK-NEXT:   ^[[BB:bb.*]](%arg0: !vhlo.tensor_v1<!vhlo.f32_v1>):
--  // CHECK-NEXT:     "vhlo.return_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-+  // CHECK-NEXT:   ^[[BB:bb.*]](%[[ARG0:.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
-+  // CHECK-NEXT:     "vhlo.return_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-   // CHECK-NEXT: }) : () -> ()
-   func.return %arg0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "dynamic_gather"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @dynamic_gather(%arg0 : tensor<2x4x9xf32>, %arg1 : tensor<1x5x2xi32>) -> tensor<1x5x1xf32> {
--  //      CHECK: "vhlo.gather_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.gather_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   collapsed_slice_dims = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-   // CHECK-SAME:   index_vector_dim = #vhlo.integer_v1<2 : i64>,
-   // CHECK-SAME:   indices_are_sorted = #vhlo.bool_v1<false>,
-@@ -605,8 +649,9 @@
- }
- 
- // CHECK-LABEL: "default_infeed"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_infeed(%arg0: !stablehlo.token) -> (tensor<f32>, !stablehlo.token) {
--  //               CHECK: "vhlo.infeed_v1"(%arg0) <{
-+  //               CHECK: "vhlo.infeed_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   infeed_config = #vhlo.string_v1<"">,
-   // CHECK-SAME{LITERAL}:   layout = #vhlo.array_v1<[]>
-   //          CHECK-SAME: }> : (!vhlo.token_v1) -> (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.token_v1)
-@@ -615,8 +660,9 @@
- }
- 
- // CHECK-LABEL: "default_outfeed"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @default_outfeed(%arg0: tensor<f32>, %arg1: !stablehlo.token) -> !stablehlo.token {
--  //      CHECK: "vhlo.outfeed_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.outfeed_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   outfeed_config = #vhlo.string_v1<"">
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.token_v1) -> !vhlo.token_v1
-   %0 = "stablehlo.outfeed"(%arg0, %arg1) : (tensor<f32>, !stablehlo.token) -> !stablehlo.token
-@@ -624,8 +670,9 @@
- }
- 
- // CHECK-LABEL: "default_recv"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_recv(%arg0: !stablehlo.token) -> (tensor<f32>, !stablehlo.token) {
--  //      CHECK: "vhlo.recv_v1"(%arg0) <{
-+  //      CHECK: "vhlo.recv_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME:   channel_type = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   is_host_transfer = #vhlo.bool_v1<false>
-@@ -637,8 +684,9 @@
- }
- 
- // CHECK-LABEL: "default_send"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @default_send(%arg0: tensor<f32>, %arg1: !stablehlo.token) -> !stablehlo.token {
--  //      CHECK: "vhlo.send_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.send_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME:   channel_type = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   is_host_transfer = #vhlo.bool_v1<false>
-@@ -650,8 +698,9 @@
- }
- 
- // CHECK-LABEL: "default_reduce_scatter"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_reduce_scatter(%arg0: tensor<16xf32>) -> tensor<16xf32> {
--  //               CHECK: "vhlo.reduce_scatter_v1"(%arg0) <{
-+  //               CHECK: "vhlo.reduce_scatter_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0], [1]]> : tensor<2x1xi64>>,
-   //          CHECK-SAME:   scatter_dimension = #vhlo.integer_v1<0 : i64>
-@@ -673,8 +722,9 @@
- }
- 
- // CHECK-LABEL: "default_reduce_window"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @default_reduce_window(%arg0: tensor<2x17x31x7xf32>, %arg1: tensor<f32>) -> tensor<2x16x30x7xf32> {
--  //               CHECK: "vhlo.reduce_window_v1"(%arg0, %arg1)  <{
-+  //               CHECK: "vhlo.reduce_window_v1"(%[[ARG0]], %[[ARG1]])  <{
-   //          CHECK-SAME:   base_dilations = #vhlo.tensor_v1<dense<1> : tensor<4xi64>>,
-   // CHECK-SAME{LITERAL}:   padding = #vhlo.tensor_v1<dense<0> : tensor<4x2xi64>>,
-   //          CHECK-SAME:   window_dilations = #vhlo.tensor_v1<dense<1> : tensor<4xi64>>,
-@@ -696,8 +746,9 @@
- }
- 
- // CHECK-LABEL: "default_scatter"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @default_scatter(%arg0: tensor<200x100x300xf32>, %arg1: tensor<10x2xi32>, %arg2: tensor<10x300xf32>) -> tensor<200x100x300xf32> {
--  //      CHECK: "vhlo.scatter_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.scatter_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   index_vector_dim = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   indices_are_sorted = #vhlo.bool_v1<false>,
-   // CHECK-SAME:   inserted_window_dims = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-@@ -725,8 +776,9 @@
- }
- 
- // CHECK-LABEL: "default_select_and_scatter"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @default_select_and_scatter(%arg0: tensor<10x24x24x64xf32>, %arg1: tensor<10x23x23x64xf32>, %arg2: tensor<f32>) -> tensor<10x24x24x64xf32> {
--  //      CHECK: "vhlo.select_and_scatter_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.select_and_scatter_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   padding = #vhlo.tensor_v1<dense<0> : tensor<4x2xi64>>,
-   // CHECK-SAME:   window_dimensions = #vhlo.tensor_v1<dense<[1, 2, 2, 1]> : tensor<4xi64>>,
-   // CHECK-SAME:   window_strides = #vhlo.tensor_v1<dense<1> : tensor<4xi64>>
-@@ -754,8 +806,9 @@
- }
- 
- // CHECK-LABEL: "default_sort"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_sort(%arg0: tensor<16xf32>) -> tensor<16xf32> {
--  //      CHECK: "vhlo.sort_v1"(%arg0) <{
-+  //      CHECK: "vhlo.sort_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   dimension = #vhlo.integer_v1<-1 : i64>
-   // CHECK-SAME:   is_stable = #vhlo.bool_v1<false>
-   // CHECK-SAME: }> ({
-@@ -774,29 +827,33 @@
- // ============ OPS ============
- 
- // CHECK-LABEL: "op_abs"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_abs(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.abs_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.abs_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.abs"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_add"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_add(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_after_all"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_after_all(%arg0: !stablehlo.token) -> !stablehlo.token {
--  // CHECK: "vhlo.after_all_v1"(%arg0) : (!vhlo.token_v1) -> !vhlo.token_v1
-+  // CHECK: "vhlo.after_all_v1"(%[[ARG0]]) : (!vhlo.token_v1) -> !vhlo.token_v1
-   %0 = "stablehlo.after_all"(%arg0) : (!stablehlo.token) -> !stablehlo.token
-   func.return %0 : !stablehlo.token
- }
- 
- // CHECK-LABEL: "op_all_gather"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_all_gather(%arg0: tensor<16x8xf32>) -> tensor<16x16xf32> {
--  //               CHECK: "vhlo.all_gather_v1"(%arg0) <{
-+  //               CHECK: "vhlo.all_gather_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   all_gather_dim = #vhlo.integer_v1<1 : i64>
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0], [1]]> : tensor<2x1xi64>>,
-@@ -812,8 +869,9 @@
- }
- 
- // CHECK-LABEL: "op_all_reduce"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_all_reduce(%arg0: tensor<f32>) -> tensor<f32> {
--  //               CHECK: "vhlo.all_reduce_v1"(%arg0) <{
-+  //               CHECK: "vhlo.all_reduce_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0], [1]]> : tensor<2x1xi64>>,
-   //          CHECK-SAME:   use_global_device_ids = #vhlo.bool_v1<true>
-@@ -836,7 +894,7 @@
- 
- // CHECK-LABEL: "op_all_reduce_with_promotable_types"
- func.func @op_all_reduce_with_promotable_types(%operand: tensor<f32>) -> tensor<f64> {
--  //  CHECK: "vhlo.all_reduce_v1"(%arg0)
-+  //  CHECK: "vhlo.all_reduce_v1"(%[[ARG0:.*]])
-   //  CHECK:   ^[[BB:bb.*]](%[[ARG1:arg.*]]: !vhlo.tensor_v1<!vhlo.f64_v1>, %[[ARG2:arg.*]]: !vhlo.tensor_v1<!vhlo.f64_v1>):
-   //  CHECK:     "vhlo.return_v1"(%[[VAL1:.*]]) : (!vhlo.tensor_v1<!vhlo.f64_v1>) -> ()
-   //  CHECK: }) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f64_v1>
-@@ -854,8 +912,9 @@
- }
- 
- // CHECK-LABEL: "op_all_to_all"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_all_to_all(%arg0: tensor<4x16xf32>) -> tensor<16x4xf32> {
--  //               CHECK: "vhlo.all_to_all_v1"(%arg0) <{
-+  //               CHECK: "vhlo.all_to_all_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<1 : i64>,
-   //          CHECK-SAME:   concat_dimension = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0, 1, 2, 3]]> : tensor<1x4xi64>>,
-@@ -873,22 +932,25 @@
- }
- 
- // CHECK-LABEL: "op_and"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_and(%arg0: tensor<i1>, %arg1: tensor<i1>) -> tensor<i1> {
--  // CHECK: "vhlo.and_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-+  // CHECK: "vhlo.and_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-   %0 = "stablehlo.and"(%arg0, %arg1) : (tensor<i1>, tensor<i1>) -> tensor<i1>
-   func.return %0 : tensor<i1>
- }
- 
- // CHECK-LABEL: "op_atan2"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_atan2(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.atan2_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.atan2_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.atan2"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_batch_norm_grad"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}}, %[[ARG3:.*]]: {{.*}}, %[[ARG4:.*]]: {{.*}})
- func.func @op_batch_norm_grad(%arg0: tensor<16x16x16x16xf32>, %arg1: tensor<16xf32>, %arg2: tensor<16xf32>, %arg3: tensor<16xf32>, %arg4: tensor<16x16x16x16xf32>) -> (tensor<16x16x16x16xf32>, tensor<16xf32>, tensor<16xf32>) {
--  //      CHECK: "vhlo.batch_norm_grad_v1"(%arg0, %arg1, %arg2, %arg3, %arg4) <{
-+  //      CHECK: "vhlo.batch_norm_grad_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]], %[[ARG3]], %[[ARG4]]) <{
-   // CHECK-SAME:   epsilon = #vhlo.float_v1<1.000000e-03 : !vhlo.f32_v1>,
-   // CHECK-SAME:   feature_index = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<16x16x16x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x16x16x16x!vhlo.f32_v1>) -> (!vhlo.tensor_v1<16x16x16x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>)
-@@ -900,8 +962,9 @@
- }
- 
- // CHECK-LABEL: "op_batch_norm_inference"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}}, %[[ARG3:.*]]: {{.*}}, %[[ARG4:.*]]: {{.*}})
- func.func @op_batch_norm_inference(%arg0: tensor<16x16x16x16xf32>, %arg1: tensor<16xf32>, %arg2: tensor<16xf32>, %arg3: tensor<16xf32>, %arg4: tensor<16xf32>) -> tensor<16x16x16x16xf32> {
--  //      CHECK: "vhlo.batch_norm_inference_v1"(%arg0, %arg1, %arg2, %arg3, %arg4) <{
-+  //      CHECK: "vhlo.batch_norm_inference_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]], %[[ARG3]], %[[ARG4]]) <{
-   // CHECK-SAME:   epsilon = #vhlo.float_v1<1.000000e-03 : !vhlo.f32_v1>,
-   // CHECK-SAME:   feature_index = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<16x16x16x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x16x16x16x!vhlo.f32_v1>
-@@ -913,8 +976,9 @@
- }
- 
- // CHECK-LABEL: "op_batch_norm_training"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_batch_norm_training(%arg0: tensor<16x16x16x16xf32>, %arg1: tensor<16xf32>, %arg2: tensor<16xf32>) -> (tensor<16x16x16x16xf32>, tensor<16xf32>, tensor<16xf32>) {
--  //      CHECK: "vhlo.batch_norm_training_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.batch_norm_training_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   epsilon = #vhlo.float_v1<1.000000e-03 : !vhlo.f32_v1>,
-   // CHECK-SAME:   feature_index = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<16x16x16x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>) -> (!vhlo.tensor_v1<16x16x16x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>)
-@@ -926,15 +990,17 @@
- }
- 
- // CHECK-LABEL: "op_bitcast_convert"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_bitcast_convert(%arg0: tensor<i32>) -> tensor<f32> {
--  // CHECK: "vhlo.bitcast_convert_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.bitcast_convert_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.bitcast_convert"(%arg0) : (tensor<i32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_broadcast_in_dim"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_broadcast_in_dim(%arg0: tensor<16xf32>) -> tensor<16x16xf32> {
--  //      CHECK: "vhlo.broadcast_in_dim_v1"(%arg0) <{
-+  //      CHECK: "vhlo.broadcast_in_dim_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   broadcast_dimensions = #vhlo.tensor_v1<dense<1> : tensor<1xi64>>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x16x!vhlo.f32_v1>
-   %0 = "stablehlo.broadcast_in_dim"(%arg0) {
-@@ -944,8 +1010,9 @@
- }
- 
- // CHECK-LABEL: "op_broadcast"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_broadcast(%arg0: tensor<16xf32>) -> tensor<16x16xf32> {
--  //      CHECK: "vhlo.broadcast_v1"(%arg0) <{
-+  //      CHECK: "vhlo.broadcast_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   broadcast_sizes = #vhlo.tensor_v1<dense<16> : tensor<1xi64>>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x16x!vhlo.f32_v1>
-   %0 = "stablehlo.broadcast"(%arg0) {
-@@ -955,9 +1022,10 @@
- }
- 
- // CHECK-LABEL: "op_case"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_case(%arg0: tensor<i32>, %arg1: tensor<f32>) -> tensor<f32> {
--  //      CHECK: "vhlo.case_v1"(%arg0) ({
--  // CHECK-NEXT:   "vhlo.return_v1"(%arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-+  //      CHECK: "vhlo.case_v1"(%[[ARG0]]) ({
-+  // CHECK-NEXT:   "vhlo.return_v1"(%[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-   // CHECK-NEXT: }) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.case"(%arg0) ({
-     "stablehlo.return"(%arg1) : (tensor<f32>) -> ()
-@@ -966,22 +1034,25 @@
- }
- 
- // CHECK-LABEL: "op_cbrt"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_cbrt(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.cbrt_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.cbrt_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.cbrt"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_ceil"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_ceil(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.ceil_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.ceil_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.ceil"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_cholesky"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_cholesky(%arg0: tensor<1x16x16xf32>) -> tensor<1x16x16xf32> {
--  //      CHECK: "vhlo.cholesky_v1"(%arg0) <{
-+  //      CHECK: "vhlo.cholesky_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   lower = #vhlo.bool_v1<true>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<1x16x16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<1x16x16x!vhlo.f32_v1>
-   %0 = "stablehlo.cholesky"(%arg0) {
-@@ -991,22 +1062,25 @@
- }
- 
- // CHECK-LABEL: "op_clamp"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_clamp(%arg0: tensor<f32>, %arg1: tensor<f32>, %arg2: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.clamp_v1"(%arg0, %arg1, %arg2) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.clamp_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.clamp"(%arg0, %arg1, %arg2) : (tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_count_leading_zeros"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_count_leading_zeros(%arg0: tensor<i32>) -> tensor<i32> {
--  // CHECK: "vhlo.count_leading_zeros_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-+  // CHECK: "vhlo.count_leading_zeros_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-   %0 = "stablehlo.count_leading_zeros"(%arg0) : (tensor<i32>) -> tensor<i32>
-   func.return %0 : tensor<i32>
- }
- 
- // CHECK-LABEL: "op_collective_permute"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_collective_permute(%arg0: tensor<16x8xf32>) -> tensor<16x8xf32> {
--  //               CHECK: "vhlo.collective_permute_v1"(%arg0) <{
-+  //               CHECK: "vhlo.collective_permute_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME{LITERAL}:   source_target_pairs = #vhlo.tensor_v1<dense<[[0, 1], [1, 2], [2, 3]]> : tensor<3x2xi64>>
-   //          CHECK-SAME: }> : (!vhlo.tensor_v1<16x8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x8x!vhlo.f32_v1>
-@@ -1018,8 +1092,9 @@
- }
- 
- // CHECK-LABEL: "op_compare"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_compare(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
--  //      CHECK: "vhlo.compare_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.compare_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   compare_type = #vhlo<comparison_type_v1 TOTALORDER>,
-   // CHECK-SAME:   comparison_direction = #vhlo<comparison_direction_v1 EQ>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-@@ -1031,15 +1106,17 @@
- }
- 
- // CHECK-LABEL: "op_complex"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_complex(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<complex<f32>> {
--  // CHECK: "vhlo.complex_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>
-+  // CHECK: "vhlo.complex_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>
-   %0 = "stablehlo.complex"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<complex<f32>>
-   func.return %0 : tensor<complex<f32>>
- }
- 
- // CHECK-LABEL: "op_concatenate"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_concatenate(%arg0: tensor<8xf32>, %arg1: tensor<8xf32>) -> tensor<16xf32> {
--  //      CHECK: "vhlo.concatenate_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.concatenate_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   dimension = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<8x!vhlo.f32_v1>, !vhlo.tensor_v1<8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x!vhlo.f32_v1>
-   %0 = "stablehlo.concatenate"(%arg0, %arg1) {
-@@ -1049,6 +1126,7 @@
- }
- 
- // CHECK-LABEL: "op_constant"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_constant(%arg0: tensor<f32>) -> tensor<f32> {
-   //      CHECK: "vhlo.constant_v1"() <{
-   // CHECK-SAME:   value = #vhlo.tensor_v1<dense<0.000000e+00> : tensor<f32>>
-@@ -1060,15 +1138,17 @@
- }
- 
- // CHECK-LABEL: "op_convert"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_convert(%arg0: tensor<i32>) -> tensor<f32> {
--  // CHECK: "vhlo.convert_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.convert_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.convert"(%arg0) : (tensor<i32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_convolution"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_convolution(%arg0: tensor<1x8x8x207xf32>, %arg1: tensor<3x3x207x16xf32>) -> tensor<1x7x7x16xf32> {
--  //      CHECK: "vhlo.convolution_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.convolution_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   batch_group_count = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   feature_group_count = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   input_batch_dimension = #vhlo.integer_v1<0 : i64>,
-@@ -1102,8 +1182,9 @@
- }
- 
- // CHECK-LABEL: "op_cosine"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_cosine(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.cosine_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.cosine_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.cosine"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
-@@ -1116,8 +1197,9 @@
- }
- 
- // CHECK-LABEL: "op_cross_replica_sum"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_cross_replica_sum(%arg0: tensor<f32>) -> tensor<f32> {
--  //               CHECK: "vhlo.cross-replica-sum_v1"(%arg0) <{
-+  //               CHECK: "vhlo.cross-replica-sum_v1"(%[[ARG0]]) <{
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0], [1]]> : tensor<2x1xi64>>
-   //          CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.cross-replica-sum"(%arg0) {
-@@ -1127,8 +1209,9 @@
- }
- 
- // CHECK-LABEL: "op_custom_call"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_custom_call(%arg0: tensor<f32>) -> tensor<f32> {
--  //      CHECK: "vhlo.custom_call_v1"(%arg0) <{
-+  //      CHECK: "vhlo.custom_call_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   api_version = #vhlo<api_version_v1 API_VERSION_STATUS_RETURNING>,
-   // CHECK-SAME:   backend_config = #vhlo.string_v1<"\08\03\1A\02">,
-   // CHECK-SAME:   call_target_name = #vhlo.string_v1<"foo">,
-@@ -1159,15 +1242,17 @@
- }
- 
- // CHECK-LABEL: "op_divide"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_divide(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.divide_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.divide_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.divide"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_dot_general"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_dot_general(%arg0: tensor<8x8x16xf32>, %arg1: tensor<8x16x8xf32>) -> tensor<8x8x8xf32> {
--  //      CHECK: "vhlo.dot_general_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.dot_general_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   lhs_batching_dimensions = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>,
-   // CHECK-SAME:   lhs_contracting_dimensions = #vhlo.tensor_v1<dense<2> : tensor<1xi64>>,
-   // CHECK-SAME:   precision_config = #vhlo.array_v1<[#vhlo<precision_v1 HIGHEST>, #vhlo<precision_v1 HIGHEST>]>,
-@@ -1187,8 +1272,9 @@
- }
- 
- // CHECK-LABEL: "op_dot"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_dot(%arg0: tensor<8x16xf32>, %arg1: tensor<16x8xf32>) -> tensor<8x8xf32> {
--  //      CHECK: "vhlo.dot_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.dot_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   precision_config = #vhlo.array_v1<[#vhlo<precision_v1 HIGHEST>, #vhlo<precision_v1 HIGHEST>]>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<8x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<8x8x!vhlo.f32_v1>
-   %0 = "stablehlo.dot"(%arg0, %arg1) {
-@@ -1198,8 +1284,9 @@
- }
- 
- // CHECK-LABEL: "op_dynamic_broadcast_in_dim"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_dynamic_broadcast_in_dim(%arg0: tensor<?x?xf32>, %arg1: tensor<2xindex>) -> tensor<?x?xf32> {
--  //      CHECK: "vhlo.dynamic_broadcast_in_dim_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.dynamic_broadcast_in_dim_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   broadcast_dimensions = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-   // CHECK-SAME:   known_expanding_dimensions = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>,
-   // CHECK-SAME:   known_nonexpanding_dimensions = #vhlo.tensor_v1<dense<1> : tensor<1xi64>>
-@@ -1213,8 +1300,9 @@
- }
- 
- // CHECK-LABEL: "op_dynamic_conv"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_dynamic_conv(%arg0: tensor<1x8x8x207xf32>, %arg1: tensor<3x3x207x16xf32>, %arg2: tensor<4xi32>) -> tensor<1x?x?x16xf32> {
--  //      CHECK: "vhlo.dynamic_conv_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.dynamic_conv_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   batch_group_count = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   feature_group_count = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   input_batch_dimension = #vhlo.integer_v1<0 : i64>,
-@@ -1248,8 +1336,9 @@
- }
- 
- // CHECK-LABEL: "op_dynamic_gather"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_dynamic_gather(%arg0 : tensor<2x4x9xf32>, %arg1 : tensor<1x5x2xi32>, %arg2 : tensor<3xi32>) -> tensor<1x5x8xf32> {
--  //      CHECK: "vhlo.dynamic_gather_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.dynamic_gather_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   collapsed_slice_dims = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-   // CHECK-SAME:   index_vector_dim = #vhlo.integer_v1<2 : i64>,
-   // CHECK-SAME:   indices_are_sorted = #vhlo.bool_v1<true>,
-@@ -1269,8 +1358,9 @@
- }
- 
- // CHECK-LABEL: "op_dynamic_iota"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_dynamic_iota(%arg0: tensor<1xindex>) -> tensor<?xf32> {
--  //      CHECK: "vhlo.dynamic_iota_v1"(%arg0) <{
-+  //      CHECK: "vhlo.dynamic_iota_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   iota_dimension = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<1x!vhlo.index_v1>) -> !vhlo.tensor_v1<?x!vhlo.f32_v1>
-   %0 = "stablehlo.dynamic_iota"(%arg0) {
-@@ -1280,22 +1370,25 @@
- }
- 
- // CHECK-LABEL: "op_dynamic_pad"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}}, %[[ARG3:.*]]: {{.*}}, %[[ARG4:.*]]: {{.*}})
- func.func @op_dynamic_pad(%arg0: tensor<?xf32>, %arg1: tensor<f32>, %arg2: tensor<1xindex>, %arg3: tensor<1xindex>, %arg4: tensor<1xindex>) -> tensor<?xf32> {
--  // CHECK: "vhlo.dynamic_pad_v1"(%arg0, %arg1, %arg2, %arg3, %arg4) : (!vhlo.tensor_v1<?x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>) -> !vhlo.tensor_v1<?x!vhlo.f32_v1>
-+  // CHECK: "vhlo.dynamic_pad_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]], %[[ARG3]], %[[ARG4]]) : (!vhlo.tensor_v1<?x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>) -> !vhlo.tensor_v1<?x!vhlo.f32_v1>
-   %0 = "stablehlo.dynamic_pad"(%arg0, %arg1, %arg2, %arg3, %arg4) : (tensor<?xf32>, tensor<f32>, tensor<1xindex>, tensor<1xindex>, tensor<1xindex>) -> tensor<?xf32>
-   func.return %0 : tensor<?xf32>
- }
- 
- // CHECK-LABEL: "op_dynamic_reshape"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_dynamic_reshape(%arg0: tensor<16xf32>, %arg1: tensor<2xindex>) -> tensor<?x?xf32> {
--  // CHECK: "vhlo.dynamic_reshape_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<2x!vhlo.index_v1>) -> !vhlo.tensor_v1<?x?x!vhlo.f32_v1>
-+  // CHECK: "vhlo.dynamic_reshape_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<2x!vhlo.index_v1>) -> !vhlo.tensor_v1<?x?x!vhlo.f32_v1>
-   %0 = "stablehlo.dynamic_reshape"(%arg0, %arg1) : (tensor<16xf32>, tensor<2xindex>) -> tensor<?x?xf32>
-   func.return %0 : tensor<?x?xf32>
- }
- 
- // CHECK-LABEL: "op_dynamic_slice"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_dynamic_slice(%arg0: tensor<16xf32>, %arg1: tensor<i64>) -> tensor<4xf32> {
--  //      CHECK: "vhlo.dynamic_slice_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.dynamic_slice_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   slice_sizes = #vhlo.tensor_v1<dense<4> : tensor<1xi64>>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.i64_v1>) -> !vhlo.tensor_v1<4x!vhlo.f32_v1>
-   %0 = "stablehlo.dynamic_slice"(%arg0, %arg1) {
-@@ -1305,15 +1398,17 @@
- }
- 
- // CHECK-LABEL: "op_dynamic_update_slice"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_dynamic_update_slice(%arg0: tensor<16xf32>, %arg1: tensor<4xf32>, %arg2: tensor<i64>) -> tensor<16xf32> {
--  // CHECK: "vhlo.dynamic_update_slice_v1"(%arg0, %arg1, %arg2) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<4x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.i64_v1>) -> !vhlo.tensor_v1<16x!vhlo.f32_v1>
-+  // CHECK: "vhlo.dynamic_update_slice_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<4x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.i64_v1>) -> !vhlo.tensor_v1<16x!vhlo.f32_v1>
-   %0 = "stablehlo.dynamic_update_slice"(%arg0, %arg1, %arg2) : (tensor<16xf32>, tensor<4xf32>, tensor<i64>) -> tensor<16xf32>
-   func.return %0 : tensor<16xf32>
- }
- 
- // CHECK-LABEL: "op_einsum"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_einsum(%arg0: tensor<8x16xf32>, %arg1: tensor<16x8xf32>) -> tensor<8x8xf32> {
--  //      CHECK: "vhlo.einsum_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.einsum_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   einsum_config = #vhlo.string_v1<"ab,bc->ac">
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<8x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<8x8x!vhlo.f32_v1>
-   %0 = "stablehlo.einsum"(%arg0, %arg1) {
-@@ -1323,22 +1418,25 @@
- }
- 
- // CHECK-LABEL: "op_exponential_minus_one"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_exponential_minus_one(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.exponential_minus_one_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.exponential_minus_one_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.exponential_minus_one"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_exponential"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_exponential(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.exponential_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.exponential_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.exponential"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_fft"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_fft(%arg0: tensor<16xcomplex<f32>>) -> tensor<16xcomplex<f32>> {
--  //      CHECK: "vhlo.fft_v1"(%arg0) <{
-+  //      CHECK: "vhlo.fft_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   fft_length = #vhlo.tensor_v1<dense<16> : tensor<1xi64>>,
-   // CHECK-SAME:   fft_type = #vhlo<fft_type_v1 FFT>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<16x!vhlo.complex_v1<!vhlo.f32_v1>>) -> !vhlo.tensor_v1<16x!vhlo.complex_v1<!vhlo.f32_v1>>
-@@ -1350,8 +1448,9 @@
- }
- 
- // CHECK-LABEL: "op_floor"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_floor(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.floor_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.floor_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.floor"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
-@@ -1364,16 +1463,17 @@
-   // CHECK-SAME:   sym_name = #vhlo.string_v1<"op_func">,
-   // CHECK-SAME:   sym_visibility = #vhlo.string_v1<"private">
-   // CHECK-SAME: }> ({
--  // CHECK-NEXT:   ^[[BB:bb.*]](%arg0: !vhlo.tensor_v1<!vhlo.f32_v1>):
--  // CHECK-NEXT:     "vhlo.return_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-+  // CHECK-NEXT:   ^[[BB:bb.*]](%[[ARG0:.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
-+  // CHECK-NEXT:     "vhlo.return_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-   // CHECK-NEXT: }) : () -> ()
- 
-   func.return %arg0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_gather"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_gather(%arg0 : tensor<2x4x9xf32>, %arg1 : tensor<1x5x2xi32>) -> tensor<1x5x1xf32> {
--  //      CHECK: "vhlo.gather_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.gather_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   collapsed_slice_dims = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-   // CHECK-SAME:   index_vector_dim = #vhlo.integer_v1<2 : i64>,
-   // CHECK-SAME:   indices_are_sorted = #vhlo.bool_v1<true>,
-@@ -1395,8 +1495,9 @@
- }
- 
- // CHECK-LABEL: "op_get_dimension_size"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_get_dimension_size(%arg0: tensor<?xf32>) -> tensor<i32> {
--  //      CHECK: "vhlo.get_dimension_size_v1"(%arg0) <{
-+  //      CHECK: "vhlo.get_dimension_size_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   dimension = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<?x!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-   %0 = "stablehlo.get_dimension_size"(%arg0) {
-@@ -1406,8 +1507,9 @@
- }
- 
- // CHECK-LABEL: "op_get_tuple_element"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_get_tuple_element(%arg0: tuple<tensor<f32>, tensor<i32>>) -> tensor<f32> {
--  //      CHECK: "vhlo.get_tuple_element_v1"(%arg0) <{
-+  //      CHECK: "vhlo.get_tuple_element_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   index = #vhlo.integer_v1<0 : i32>
-   // CHECK-SAME: }> : (!vhlo.tuple_v1<!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.get_tuple_element"(%arg0) {
-@@ -1417,11 +1519,12 @@
- }
- 
- // CHECK-LABEL: "op_if"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_if(%arg0: tensor<i1>, %arg1: tensor<f32>, %arg2: tensor<f32>) -> tensor<f32> {
--  //      CHECK: "vhlo.if_v1"(%arg0) ({
--  // CHECK-NEXT:   "vhlo.return_v1"(%arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-+  //      CHECK: "vhlo.if_v1"(%[[ARG0]]) ({
-+  // CHECK-NEXT:   "vhlo.return_v1"(%[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-   // CHECK-NEXT: }, {
--  // CHECK-NEXT:   "vhlo.return_v1"(%arg2) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-+  // CHECK-NEXT:   "vhlo.return_v1"(%[[ARG2]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-   // CHECK-NEXT: }) : (!vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.if"(%arg0) ({
-     "stablehlo.return"(%arg1) : (tensor<f32>) -> ()
-@@ -1432,15 +1535,17 @@
- }
- 
- // CHECK-LABEL: "op_imag"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_imag(%arg0: tensor<complex<f32>>) -> tensor<f32> {
--  // CHECK: "vhlo.imag_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.imag_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.imag"(%arg0) : (tensor<complex<f32>>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_infeed"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_infeed(%arg0: !stablehlo.token) -> (tensor<f32>, !stablehlo.token) {
--  //               CHECK: "vhlo.infeed_v1"(%arg0) <{
-+  //               CHECK: "vhlo.infeed_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   infeed_config = #vhlo.string_v1<"foo">,
-   // CHECK-SAME{LITERAL}:   layout = #vhlo.array_v1<[#vhlo.array_v1<[]>]>
-   //          CHECK-SAME: }> : (!vhlo.token_v1) -> (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.token_v1)
-@@ -1463,36 +1568,41 @@
- }
- 
- // CHECK-LABEL: "op_is_finite"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_is_finite(%arg0: tensor<f32>) -> tensor<i1> {
--  // CHECK: "vhlo.is_finite_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-+  // CHECK: "vhlo.is_finite_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-   %0 = "stablehlo.is_finite"(%arg0) : (tensor<f32>) -> tensor<i1>
-   func.return %0 : tensor<i1>
- }
- 
- // CHECK-LABEL: "op_log"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_log(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.log_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.log_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.log"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_log_plus_one"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_log_plus_one(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.log_plus_one_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.log_plus_one_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.log_plus_one"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_logistic"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_logistic(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.logistic_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.logistic_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.logistic"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_map"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_map(%arg0: tensor<16xf32>) -> tensor<16xf32> {
--  //      CHECK: "vhlo.map_v1"(%arg0) <{
-+  //      CHECK: "vhlo.map_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   dimensions = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>
-   // CHECK-SAME: }> ({
-   // CHECK-NEXT:   ^[[BB:bb.*]](%[[ARG1:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
-@@ -1510,57 +1620,65 @@
- }
- 
- // CHECK-LABEL: "op_maximum"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_maximum(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.maximum_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.maximum_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.maximum"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_minimum"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_minimum(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.minimum_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.minimum_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.minimum"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_multiply"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_multiply(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.multiply_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.multiply_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.multiply"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_negate"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_negate(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.negate_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.negate_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.negate"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_not"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_not(%arg0: tensor<i1>) -> tensor<i1> {
--  // CHECK: "vhlo.not_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-+  // CHECK: "vhlo.not_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-   %0 = "stablehlo.not"(%arg0) : (tensor<i1>) -> tensor<i1>
-   func.return %0 : tensor<i1>
- }
- 
- // CHECK-LABEL: "op_optimization_barrier"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_optimization_barrier(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.optimization_barrier_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.optimization_barrier_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.optimization_barrier"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_or"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_or(%arg0: tensor<i1>, %arg1: tensor<i1>) -> tensor<i1> {
--  // CHECK: "vhlo.or_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-+  // CHECK: "vhlo.or_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-   %0 = "stablehlo.or"(%arg0, %arg1) : (tensor<i1>, tensor<i1>) -> tensor<i1>
-   func.return %0 : tensor<i1>
- }
- 
- // CHECK-LABEL: "op_outfeed"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_outfeed(%arg0: tensor<f32>, %arg1: !stablehlo.token) -> !stablehlo.token {
--  //      CHECK: "vhlo.outfeed_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.outfeed_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   outfeed_config = #vhlo.string_v1<"foo">
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.token_v1) -> !vhlo.token_v1
-   %0 = "stablehlo.outfeed"(%arg0, %arg1) {
-@@ -1570,8 +1688,9 @@
- }
- 
- // CHECK-LABEL: "op_pad"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_pad(%arg0: tensor<8xf32>, %arg1: tensor<f32>) -> tensor<16xf32> {
--  //      CHECK: "vhlo.pad_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.pad_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   edge_padding_high = #vhlo.tensor_v1<dense<4> : tensor<1xi64>>,
-   // CHECK-SAME:   edge_padding_low = #vhlo.tensor_v1<dense<4> : tensor<1xi64>>,
-   // CHECK-SAME:   interior_padding = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>
-@@ -1585,36 +1704,41 @@
- }
- 
- // CHECK-LABEL: "op_popcnt"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_popcnt(%arg0: tensor<i32>) -> tensor<i32> {
--  // CHECK: "vhlo.popcnt_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-+  // CHECK: "vhlo.popcnt_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-   %0 = "stablehlo.popcnt"(%arg0) : (tensor<i32>) -> tensor<i32>
-   func.return %0 : tensor<i32>
- }
- 
- // CHECK-LABEL: "op_power"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_power(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.power_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.power_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.power"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_real_dynamic_slice"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}}, %[[ARG3:.*]]: {{.*}})
- func.func @op_real_dynamic_slice(%arg0: tensor<?xf32>, %arg1: tensor<1xindex>, %arg2: tensor<1xindex>, %arg3: tensor<1xindex>) -> tensor<?xf32> {
--  // CHECK: "vhlo.real_dynamic_slice_v1"(%arg0, %arg1, %arg2, %arg3) : (!vhlo.tensor_v1<?x!vhlo.f32_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>) -> !vhlo.tensor_v1<?x!vhlo.f32_v1>
-+  // CHECK: "vhlo.real_dynamic_slice_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]], %[[ARG3]]) : (!vhlo.tensor_v1<?x!vhlo.f32_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>) -> !vhlo.tensor_v1<?x!vhlo.f32_v1>
-   %0 = "stablehlo.real_dynamic_slice"(%arg0, %arg1, %arg2, %arg3) : (tensor<?xf32>, tensor<1xindex>, tensor<1xindex>, tensor<1xindex>) -> tensor<?xf32>
-   func.return %0 : tensor<?xf32>
- }
- 
- // CHECK-LABEL: "op_real"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_real(%arg0: tensor<complex<f32>>) -> tensor<f32> {
--  // CHECK: "vhlo.real_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.real_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.real"(%arg0) : (tensor<complex<f32>>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_recv"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_recv(%arg0: !stablehlo.token) -> (tensor<f32>, !stablehlo.token) {
--  //      CHECK: "vhlo.recv_v1"(%arg0) <{
-+  //      CHECK: "vhlo.recv_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME:   channel_type = #vhlo.integer_v1<3 : i64>,
-   // CHECK-SAME:   is_host_transfer = #vhlo.bool_v1<true>
-@@ -1627,8 +1751,9 @@
- }
- 
- // CHECK-LABEL: "op_reduce"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_reduce(%arg0: tensor<16xf32>, %arg1: tensor<f32>) -> tensor<f32> {
--  //  CHECK: "vhlo.reduce_v1"(%arg0, %arg1)
-+  //  CHECK: "vhlo.reduce_v1"(%[[ARG0]], %[[ARG1]])
-   //  CHECK:   ^[[BB:bb.*]](%[[ARG1:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>, %[[ARG2:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
-   //  CHECK:     "vhlo.return_v1"(%[[VAL1:.*]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-   //  CHECK: }) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-@@ -1643,8 +1768,9 @@
- }
- 
- // CHECK-LABEL: "op_reduce_precision"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_reduce_precision(%arg0: tensor<f32>) -> tensor<f32> {
--  //      CHECK: "vhlo.reduce_precision_v1"(%arg0) <{
-+  //      CHECK: "vhlo.reduce_precision_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   exponent_bits = #vhlo.integer_v1<8 : i32>
-   // CHECK-SAME:   mantissa_bits = #vhlo.integer_v1<10 : i32>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-@@ -1658,7 +1784,7 @@
- // CHECK_lABEL: "op_reduce_with_promotable_types"
- func.func @op_reduce_with_promotable_types(%arg0: tensor<4x4xf32>, %arg1 : tensor<f32>)
-     -> (tensor<4xf64>) {
--  //  CHECK: "vhlo.reduce_v1"(%arg0, %arg1)
-+  //  CHECK: "vhlo.reduce_v1"(%[[ARG0:.*]], %[[ARG1:.*]])
-   //  CHECK:   ^[[BB:bb.*]](%[[ARG1:arg.*]]: !vhlo.tensor_v1<!vhlo.f64_v1>, %[[ARG2:arg.*]]: !vhlo.tensor_v1<!vhlo.f64_v1>):
-   //  CHECK:     "vhlo.return_v1"(%[[VAL1:.*]]) : (!vhlo.tensor_v1<!vhlo.f64_v1>) -> ()
-   //  CHECK: }) : (!vhlo.tensor_v1<4x4x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<4x!vhlo.f64_v1>
-@@ -1673,8 +1799,9 @@
- }
- 
- // CHECK-LABEL: "op_reduce_scatter"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_reduce_scatter(%arg0: tensor<16xf32>) -> tensor<16xf32> {
--  //               CHECK: "vhlo.reduce_scatter_v1"(%arg0) <{
-+  //               CHECK: "vhlo.reduce_scatter_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0], [1]]> : tensor<2x1xi64>>,
-   //          CHECK-SAME:   scatter_dimension = #vhlo.integer_v1<0 : i64>
-@@ -1699,7 +1826,7 @@
- 
- // CHECK_lABEL: "op_reduce_scatter_with_promotable_types"
- func.func @op_reduce_scatter_with_promotable_types(%data: tensor<4x16xf32>) -> tensor<4x4xf64> {
--  //  CHECK: "vhlo.reduce_scatter_v1"(%arg0)
-+  //  CHECK: "vhlo.reduce_scatter_v1"(%[[ARG0:.*]])
-   //  CHECK:   ^[[BB:bb.*]](%[[ARG1:arg.*]]: !vhlo.tensor_v1<!vhlo.f64_v1>, %[[ARG2:arg.*]]: !vhlo.tensor_v1<!vhlo.f64_v1>):
-   //  CHECK:     "vhlo.return_v1"(%[[VAL1:.*]]) : (!vhlo.tensor_v1<!vhlo.f64_v1>) -> ()
-   //  CHECK: }) : (!vhlo.tensor_v1<4x16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<4x4x!vhlo.f64_v1>
-@@ -1716,8 +1843,9 @@
- 
- 
- // CHECK-LABEL: "op_reduce_window"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_reduce_window(%arg0: tensor<2x17x31x7xf32>, %arg1: tensor<f32>) -> tensor<2x9x16x7xf32> {
--  //               CHECK: "vhlo.reduce_window_v1"(%arg0, %arg1) <{
-+  //               CHECK: "vhlo.reduce_window_v1"(%[[ARG0]], %[[ARG1]]) <{
-   //          CHECK-SAME:   base_dilations = #vhlo.tensor_v1<dense<[1, 2, 2, 1]> : tensor<4xi64>>,
-   // CHECK-SAME{LITERAL}:   padding = #vhlo.tensor_v1<dense<[[0, 0], [2, 0], [0, 2], [0, 0]]> : tensor<4x2xi64>>,
-   //          CHECK-SAME:   window_dilations = #vhlo.tensor_v1<dense<[1, 2, 2, 1]> : tensor<4xi64>>,
-@@ -1746,7 +1874,7 @@
- func.func @op_reduce_window_with_promotable_types(%arg0: tensor<4x2xf32>,
-     %arg1: tensor<4x2xf32>, %init0: tensor<f32>, %init1: tensor<f32>) ->
-     (tensor<2x2xf64>, tensor<2x2xf32>) {
--  //  CHECK: "vhlo.reduce_window_v1"(%arg0, %arg1, %arg2, %arg3)
-+  //  CHECK: "vhlo.reduce_window_v1"(%[[ARG0:.*]], %[[ARG1:.*]], %[[ARG2:.*]], %[[ARG3:.*]])
-   //  CHECK:   ^[[BB:bb.*]](%[[ARG1:arg.*]]: !vhlo.tensor_v1<!vhlo.f64_v1>, %[[ARG2:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>, %[[ARG3:arg.*]]: !vhlo.tensor_v1<!vhlo.f64_v1>, %[[ARG4:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
-   //  CHECK:     "vhlo.return_v1"(%[[VAL1:.*]], %[[VAL2:.*]]) : (!vhlo.tensor_v1<!vhlo.f64_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-   //  CHECK: }) : (!vhlo.tensor_v1<4x2x!vhlo.f32_v1>, !vhlo.tensor_v1<4x2x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> (!vhlo.tensor_v1<2x2x!vhlo.f64_v1>, !vhlo.tensor_v1<2x2x!vhlo.f32_v1>)
-@@ -1766,8 +1894,9 @@
- }
- 
- // CHECK-LABEL: "op_remainder"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_remainder(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.remainder_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.remainder_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.remainder"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
-@@ -1787,16 +1916,18 @@
- }
- 
- // CHECK-LABEL: "op_reshape"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_reshape(%arg0: tensor<16xf32>) -> tensor<4x4xf32> {
--  // CHECK: "vhlo.reshape_v1"(%arg0) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<4x4x!vhlo.f32_v1>
-+  // CHECK: "vhlo.reshape_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<4x4x!vhlo.f32_v1>
-   %0 = "stablehlo.reshape"(%arg0) : (tensor<16xf32>) -> tensor<4x4xf32>
-   func.return %0 : tensor<4x4xf32>
- }
- 
- // CHECK-LABEL: "op_return"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_return(%arg0: tensor<i32>, %arg1: tensor<f32>) -> tensor<f32> {
--  //      CHECK: "vhlo.case_v1"(%arg0) ({
--  // CHECK-NEXT:   "vhlo.return_v1"(%arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-+  //      CHECK: "vhlo.case_v1"(%[[ARG0]]) ({
-+  // CHECK-NEXT:   "vhlo.return_v1"(%[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-   // CHECK-NEXT: }) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.case"(%arg0) ({
-     "stablehlo.return"(%arg1) : (tensor<f32>) -> ()
-@@ -1805,8 +1936,9 @@
- }
- 
- // CHECK-LABEL: "op_reverse"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_reverse(%arg0: tensor<16xf32>) -> tensor<16xf32> {
--  //      CHECK: "vhlo.reverse_v1"(%arg0) <{
-+  //      CHECK: "vhlo.reverse_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   dimensions = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x!vhlo.f32_v1>
-   %0 = "stablehlo.reverse"(%arg0) {
-@@ -1816,8 +1948,9 @@
- }
- 
- // CHECK-LABEL: "op_rng_bit_generator"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_rng_bit_generator(%arg0: tensor<f32>) -> (tensor<f32>, tensor<f32>) {
--  //      CHECK: "vhlo.rng_bit_generator_v1"(%arg0) <{
-+  //      CHECK: "vhlo.rng_bit_generator_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   rng_algorithm = #vhlo<rng_algorithm_v1 PHILOX>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>)
-   %0:2 = "stablehlo.rng_bit_generator"(%arg0) {
-@@ -1827,8 +1960,9 @@
- }
- 
- // CHECK-LABEL: "op_rng"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_rng(%arg0: tensor<f32>, %arg1: tensor<f32>, %arg2: tensor<0xindex>) -> tensor<f32> {
--  //      CHECK: "vhlo.rng_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.rng_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   rng_distribution = #vhlo<rng_distribution_v1 NORMAL>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<0x!vhlo.index_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.rng"(%arg0, %arg1, %arg2) {
-@@ -1838,29 +1972,33 @@
- }
- 
- // CHECK-LABEL: "op_round_nearest_afz"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_round_nearest_afz(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.round_nearest_afz_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.round_nearest_afz_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.round_nearest_afz"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_round_nearest_even"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_round_nearest_even(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.round_nearest_even_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.round_nearest_even_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.round_nearest_even"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_rsqrt"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_rsqrt(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.rsqrt_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.rsqrt_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.rsqrt"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_scatter"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_scatter(%arg0: tensor<200x100x300xf32>, %arg1: tensor<10x2xi32>, %arg2: tensor<10x300xf32>) -> tensor<200x100x300xf32> {
--  //      CHECK: "vhlo.scatter_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.scatter_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   index_vector_dim = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   indices_are_sorted = #vhlo.bool_v1<true>,
-   // CHECK-SAME:   inserted_window_dims = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-@@ -1893,7 +2031,7 @@
- func.func @op_scatter_with_promotable_types(%input_tensor: tensor<200x100x300xf32>,
-     %scatter_indices: tensor<10x2xi32>, %updates: tensor<10x300xf32>) ->
-       tensor<200x100x300xf64> {
--  //  CHECK: "vhlo.scatter_v1"(%arg0, %arg1, %arg2)
-+  //  CHECK: "vhlo.scatter_v1"(%[[ARG0:.*]], %[[ARG1:.*]], %[[ARG2:.*]])
-   //  CHECK:   ^[[BB:bb.*]](%[[ARG1:arg.*]]: !vhlo.tensor_v1<!vhlo.f64_v1>, %[[ARG2:arg.*]]: !vhlo.tensor_v1<!vhlo.f64_v1>):
-   //  CHECK:     "vhlo.return_v1"(%[[VAL1:.*]]) : (!vhlo.tensor_v1<!vhlo.f64_v1>) -> ()
-   //  CHECK: }) : (!vhlo.tensor_v1<200x100x300x!vhlo.f32_v1>, !vhlo.tensor_v1<10x2x!vhlo.i32_v1>, !vhlo.tensor_v1<10x300x!vhlo.f32_v1>) -> !vhlo.tensor_v1<200x100x300x!vhlo.f64_v1>
-@@ -1916,8 +2054,9 @@
- }
- 
- // CHECK-LABEL: "op_select_and_scatter"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_select_and_scatter(%arg0: tensor<10x24x24x64xf32>, %arg1: tensor<12x13x13x66xf32>, %arg2: tensor<f32>) -> tensor<10x24x24x64xf32> {
--  //      CHECK: "vhlo.select_and_scatter_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.select_and_scatter_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   padding = #vhlo.tensor_v1<dense<1> : tensor<4x2xi64>>,
-   // CHECK-SAME:   window_dimensions = #vhlo.tensor_v1<dense<[1, 2, 2, 1]> : tensor<4xi64>>,
-   // CHECK-SAME:   window_strides = #vhlo.tensor_v1<dense<[1, 2, 2, 1]> : tensor<4xi64>>
-@@ -1947,8 +2086,9 @@
- }
- 
- // CHECK-LABEL: "op_select_and_scatter_with_promotable_types"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_select_and_scatter_with_promotable_types(%arg0: tensor<10x24x24x64xf32>, %arg1: tensor<12x13x13x66xf32>, %arg2: tensor<f32>) -> tensor<10x24x24x64xf64> {
--  // CHECK: "vhlo.select_and_scatter_v1"(%arg0, %arg1, %arg2)
-+  // CHECK: "vhlo.select_and_scatter_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]])
-   // CHECK:   ^[[BB:bb.*]](%[[ARG1:arg.*]]: !vhlo.tensor_v1<!vhlo.f64_v1>, %[[ARG2:arg.*]]: !vhlo.tensor_v1<!vhlo.f64_v1>):
-   // CHECK:     %[[VAL:.*]] = "vhlo.add_v1"(%[[ARG1]], %[[ARG2]]) : (!vhlo.tensor_v1<!vhlo.f64_v1>, !vhlo.tensor_v1<!vhlo.f64_v1>) -> !vhlo.tensor_v1<!vhlo.f64_v1>
-   // CHECK:     "vhlo.return_v1"(%[[VAL]]) : (!vhlo.tensor_v1<!vhlo.f64_v1>) -> ()
-@@ -1970,15 +2110,17 @@
- }
- 
- // CHECK-LABEL: "op_select"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_select(%arg0: tensor<i1>, %arg1: tensor<f32>, %arg2: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.select_v1"(%arg0, %arg1, %arg2) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.select_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.select"(%arg0, %arg1, %arg2) : (tensor<i1>, tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_send"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_send(%arg0: tensor<f32>, %arg1: !stablehlo.token) -> !stablehlo.token {
--  //      CHECK: "vhlo.send_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.send_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME:   channel_type = #vhlo.integer_v1<2 : i64>,
-   // CHECK-SAME:   is_host_transfer = #vhlo.bool_v1<true>
-@@ -1991,8 +2133,9 @@
- }
- 
- // CHECK-LABEL: "op_set_dimension_size"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_set_dimension_size(%arg0: tensor<?xf32>, %arg1: tensor<i32>) -> tensor<16xf32> {
--  //      CHECK: "vhlo.set_dimension_size_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.set_dimension_size_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   dimension = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<?x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<16x!vhlo.f32_v1>
-   %0 = "stablehlo.set_dimension_size"(%arg0, %arg1) {
-@@ -2002,43 +2145,49 @@
- }
- 
- // CHECK-LABEL: "op_shift_left"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_shift_left(%arg0: tensor<i32>, %arg1: tensor<i32>) -> tensor<i32> {
--  // CHECK: "vhlo.shift_left_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.i32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-+  // CHECK: "vhlo.shift_left_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-   %0 = "stablehlo.shift_left"(%arg0, %arg1) : (tensor<i32>, tensor<i32>) -> tensor<i32>
-   func.return %0 : tensor<i32>
- }
- 
- // CHECK-LABEL: "op_shift_right_arithmetic"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_shift_right_arithmetic(%arg0: tensor<i32>, %arg1: tensor<i32>) -> tensor<i32> {
--  // CHECK: "vhlo.shift_right_arithmetic_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.i32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-+  // CHECK: "vhlo.shift_right_arithmetic_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-   %0 = "stablehlo.shift_right_arithmetic"(%arg0, %arg1) : (tensor<i32>, tensor<i32>) -> tensor<i32>
-   func.return %0 : tensor<i32>
- }
- 
- // CHECK-LABEL: "op_shift_right_logical"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_shift_right_logical(%arg0: tensor<i32>, %arg1: tensor<i32>) -> tensor<i32> {
--  // CHECK: "vhlo.shift_right_logical_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.i32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-+  // CHECK: "vhlo.shift_right_logical_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-   %0 = "stablehlo.shift_right_logical"(%arg0, %arg1) : (tensor<i32>, tensor<i32>) -> tensor<i32>
-   func.return %0 : tensor<i32>
- }
- 
- // CHECK-LABEL: "op_sign"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_sign(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.sign_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.sign_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.sign"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_sine"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_sine(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.sine_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.sine_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.sine"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_slice"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_slice(%arg0: tensor<16xf32>) -> tensor<4xf32> {
--  //      CHECK: "vhlo.slice_v1"(%arg0) <{
-+  //      CHECK: "vhlo.slice_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   limit_indices = #vhlo.tensor_v1<dense<4> : tensor<1xi64>>,
-   // CHECK-SAME:   start_indices = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>,
-   // CHECK-SAME:   strides = #vhlo.tensor_v1<dense<1> : tensor<1xi64>>
-@@ -2052,8 +2201,9 @@
- }
- 
- // CHECK-LABEL: "op_sort"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_sort(%arg0: tensor<16xf32>) -> tensor<16xf32> {
--  //      CHECK: "vhlo.sort_v1"(%arg0) <{
-+  //      CHECK: "vhlo.sort_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   dimension = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME:   is_stable = #vhlo.bool_v1<true>
-   // CHECK-SAME: }> ({
-@@ -2073,29 +2223,33 @@
- }
- 
- // CHECK-LABEL: "op_sqrt"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_sqrt(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.sqrt_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.sqrt_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.sqrt"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_subtract"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_subtract(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.subtract_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.subtract_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.subtract"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_tanh"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_tanh(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.tanh_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.tanh_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.tanh"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_torch_index_select"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_torch_index_select(%arg0: tensor<5x1x5xf32>, %arg1: tensor<2xi32>) ->  tensor<2x1x5xf32> {
--  //      CHECK: "vhlo.torch_index_select_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.torch_index_select_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   batch_dims = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME:   dim = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<5x1x5x!vhlo.f32_v1>, !vhlo.tensor_v1<2x!vhlo.i32_v1>) -> !vhlo.tensor_v1<2x1x5x!vhlo.f32_v1>
-@@ -2107,8 +2261,9 @@
- }
- 
- // CHECK-LABEL: "op_trace"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_trace(%arg0: tensor<f32>) {
--  //      CHECK: "vhlo.trace_v1"(%arg0) <{
-+  //      CHECK: "vhlo.trace_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   tag = #vhlo.string_v1<"foo">
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-   "stablehlo.trace"(%arg0) {
-@@ -2118,8 +2273,9 @@
- }
- 
- // CHECK-LABEL: "op_transpose"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_transpose(%arg0: tensor<16x8xf32>) ->  tensor<8x16xf32> {
--  //      CHECK: "vhlo.transpose_v1"(%arg0) <{
-+  //      CHECK: "vhlo.transpose_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   permutation = #vhlo.tensor_v1<dense<[1, 0]> : tensor<2xi64>>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<16x8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<8x16x!vhlo.f32_v1>
-   %0 = "stablehlo.transpose"(%arg0) {
-@@ -2129,8 +2285,9 @@
- }
- 
- // CHECK-LABEL: "op_triangular_solve"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_triangular_solve(%arg0: tensor<16x16xf32>, %arg1: tensor<16x16xf32>) ->  tensor<16x16xf32> {
--  //      CHECK: "vhlo.triangular_solve_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.triangular_solve_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   left_side = #vhlo.bool_v1<true>,
-   // CHECK-SAME:   lower = #vhlo.bool_v1<true>,
-   // CHECK-SAME:   transpose_a = #vhlo<transpose_v1 NO_TRANSPOSE>,
-@@ -2146,15 +2303,17 @@
- }
- 
- // CHECK-LABEL: "op_tuple"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_tuple(%arg0: tensor<f32>) -> tuple<tensor<f32>> {
--  // CHECK: "vhlo.tuple_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tuple_v1<!vhlo.tensor_v1<!vhlo.f32_v1>>
-+  // CHECK: "vhlo.tuple_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tuple_v1<!vhlo.tensor_v1<!vhlo.f32_v1>>
-   %0 = "stablehlo.tuple"(%arg0) : (tensor<f32>) -> tuple<tensor<f32>>
-   func.return %0 : tuple<tensor<f32>>
- }
- 
- // CHECK-LABEL: "op_unary_einsum"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_unary_einsum(%arg0: tensor<8x16xf32>) -> tensor<8xf32> {
--  //      CHECK: "vhlo.unary_einsum_v1"(%arg0) <{
-+  //      CHECK: "vhlo.unary_einsum_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   einsum_config = #vhlo.string_v1<"ab->a">
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<8x16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<8x!vhlo.f32_v1>
-   %0 = "stablehlo.unary_einsum"(%arg0) {
-@@ -2164,22 +2323,25 @@
- }
- 
- // CHECK-LABEL: "op_uniform_dequantize"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_uniform_dequantize(%arg0: tensor<!quant.uniform<i8:f32, 34.0:16>>) -> tensor<f32> {
--  // CHECK: "vhlo.uniform_dequantize_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.uniform_dequantize_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.uniform_dequantize"(%arg0) : (tensor<!quant.uniform<i8:f32, 34.0:16>>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_uniform_quantize"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_uniform_quantize(%arg0: tensor<f32>) -> tensor<!quant.uniform<i8:f32, 34.0:16>> {
--  // CHECK: "vhlo.uniform_quantize_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>
-+  // CHECK: "vhlo.uniform_quantize_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>
-   %0 = "stablehlo.uniform_quantize"(%arg0) : (tensor<f32>) -> tensor<!quant.uniform<i8:f32, 34.0:16>>
-   func.return %0 : tensor<!quant.uniform<i8:f32, 34.0:16>>
- }
- 
- // CHECK-LABEL: "op_while"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_while(%arg0: tensor<i1>) -> tensor<i1> {
--  //      CHECK: "vhlo.while_v1"(%arg0) ({
-+  //      CHECK: "vhlo.while_v1"(%[[ARG0]]) ({
-   // CHECK-NEXT:   ^[[BB:bb.*]](%[[ARG1:arg.*]]: !vhlo.tensor_v1<!vhlo.bool_v1>):
-   // CHECK-NEXT:     "vhlo.return_v1"(%[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>) -> ()
-   // CHECK-NEXT:   }, {
-@@ -2197,8 +2359,9 @@
- }
- 
- // CHECK-LABEL: "op_xor"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_xor(%arg0: tensor<i1>, %arg1: tensor<i1>) -> tensor<i1> {
--  // CHECK: "vhlo.xor_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-+  // CHECK: "vhlo.xor_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-   %0 = "stablehlo.xor"(%arg0, %arg1) : (tensor<i1>, tensor<i1>) -> tensor<i1>
-   func.return %0 : tensor<i1>
- }
-@@ -2206,197 +2369,225 @@
- // ============ TYPES ============
- 
- // CHECK-LABEL: "type_i1"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_i1(%arg0: tensor<i1>, %arg1: tensor<i1>) -> tensor<i1> {
--  // CHECK: "vhlo.and_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-+  // CHECK: "vhlo.and_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-   %0 = "stablehlo.and"(%arg0, %arg1) : (tensor<i1>, tensor<i1>) -> tensor<i1>
-   func.return %0 : tensor<i1>
- }
- 
- // CHECK-LABEL: "type_i4"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_i4(%arg0: tensor<i4>, %arg1: tensor<i4>) -> tensor<i4> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.i4_v1>, !vhlo.tensor_v1<!vhlo.i4_v1>) -> !vhlo.tensor_v1<!vhlo.i4_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i4_v1>, !vhlo.tensor_v1<!vhlo.i4_v1>) -> !vhlo.tensor_v1<!vhlo.i4_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<i4>, tensor<i4>) -> tensor<i4>
-   func.return %0 : tensor<i4>
- }
- 
- // CHECK-LABEL: "type_i8"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_i8(%arg0: tensor<i8>, %arg1: tensor<i8>) -> tensor<i8> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.i8_v1>, !vhlo.tensor_v1<!vhlo.i8_v1>) -> !vhlo.tensor_v1<!vhlo.i8_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i8_v1>, !vhlo.tensor_v1<!vhlo.i8_v1>) -> !vhlo.tensor_v1<!vhlo.i8_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<i8>, tensor<i8>) -> tensor<i8>
-   func.return %0 : tensor<i8>
- }
- 
- // CHECK-LABEL: "type_i16"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_i16(%arg0: tensor<i16>, %arg1: tensor<i16>) -> tensor<i16> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.i16_v1>, !vhlo.tensor_v1<!vhlo.i16_v1>) -> !vhlo.tensor_v1<!vhlo.i16_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i16_v1>, !vhlo.tensor_v1<!vhlo.i16_v1>) -> !vhlo.tensor_v1<!vhlo.i16_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<i16>, tensor<i16>) -> tensor<i16>
-   func.return %0 : tensor<i16>
- }
- 
- // CHECK-LABEL: "type_i32"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_i32(%arg0: tensor<i32>, %arg1: tensor<i32>) -> tensor<i32> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.i32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<i32>, tensor<i32>) -> tensor<i32>
-   func.return %0 : tensor<i32>
- }
- 
- // CHECK-LABEL: "type_i64"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_i64(%arg0: tensor<i64>, %arg1: tensor<i64>) -> tensor<i64> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.i64_v1>, !vhlo.tensor_v1<!vhlo.i64_v1>) -> !vhlo.tensor_v1<!vhlo.i64_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i64_v1>, !vhlo.tensor_v1<!vhlo.i64_v1>) -> !vhlo.tensor_v1<!vhlo.i64_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<i64>, tensor<i64>) -> tensor<i64>
-   func.return %0 : tensor<i64>
- }
- 
- // CHECK-LABEL: "type_ui4"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_ui4(%arg0: tensor<ui4>, %arg1: tensor<ui4>) -> tensor<ui4> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.ui4_v1>, !vhlo.tensor_v1<!vhlo.ui4_v1>) -> !vhlo.tensor_v1<!vhlo.ui4_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.ui4_v1>, !vhlo.tensor_v1<!vhlo.ui4_v1>) -> !vhlo.tensor_v1<!vhlo.ui4_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<ui4>, tensor<ui4>) -> tensor<ui4>
-   func.return %0 : tensor<ui4>
- }
- 
- // CHECK-LABEL: "type_ui8"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_ui8(%arg0: tensor<ui8>, %arg1: tensor<ui8>) -> tensor<ui8> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.ui8_v1>, !vhlo.tensor_v1<!vhlo.ui8_v1>) -> !vhlo.tensor_v1<!vhlo.ui8_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.ui8_v1>, !vhlo.tensor_v1<!vhlo.ui8_v1>) -> !vhlo.tensor_v1<!vhlo.ui8_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<ui8>, tensor<ui8>) -> tensor<ui8>
-   func.return %0 : tensor<ui8>
- }
- 
- // CHECK-LABEL: "type_ui16"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_ui16(%arg0: tensor<ui16>, %arg1: tensor<ui16>) -> tensor<ui16> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.ui16_v1>, !vhlo.tensor_v1<!vhlo.ui16_v1>) -> !vhlo.tensor_v1<!vhlo.ui16_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.ui16_v1>, !vhlo.tensor_v1<!vhlo.ui16_v1>) -> !vhlo.tensor_v1<!vhlo.ui16_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<ui16>, tensor<ui16>) -> tensor<ui16>
-   func.return %0 : tensor<ui16>
- }
- 
- // CHECK-LABEL: "type_ui32"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_ui32(%arg0: tensor<ui32>, %arg1: tensor<ui32>) -> tensor<ui32> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.ui32_v1>, !vhlo.tensor_v1<!vhlo.ui32_v1>) -> !vhlo.tensor_v1<!vhlo.ui32_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.ui32_v1>, !vhlo.tensor_v1<!vhlo.ui32_v1>) -> !vhlo.tensor_v1<!vhlo.ui32_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<ui32>, tensor<ui32>) -> tensor<ui32>
-   func.return %0 : tensor<ui32>
- }
- 
- // CHECK-LABEL: "type_ui64"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_ui64(%arg0: tensor<ui64>, %arg1: tensor<ui64>) -> tensor<ui64> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.ui64_v1>, !vhlo.tensor_v1<!vhlo.ui64_v1>) -> !vhlo.tensor_v1<!vhlo.ui64_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.ui64_v1>, !vhlo.tensor_v1<!vhlo.ui64_v1>) -> !vhlo.tensor_v1<!vhlo.ui64_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<ui64>, tensor<ui64>) -> tensor<ui64>
-   func.return %0 : tensor<ui64>
- }
- 
- // CHECK-LABEL: "type_f8E4M3FN"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_f8E4M3FN(%arg0: tensor<f8E4M3FN>, %arg1: tensor<f8E4M3FN>) -> tensor<f8E4M3FN> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f8E4M3FN_v1>, !vhlo.tensor_v1<!vhlo.f8E4M3FN_v1>) -> !vhlo.tensor_v1<!vhlo.f8E4M3FN_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f8E4M3FN_v1>, !vhlo.tensor_v1<!vhlo.f8E4M3FN_v1>) -> !vhlo.tensor_v1<!vhlo.f8E4M3FN_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f8E4M3FN>, tensor<f8E4M3FN>) -> tensor<f8E4M3FN>
-   func.return %0 : tensor<f8E4M3FN>
- }
- 
- // CHECK-LABEL: "type_f8E5M2"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_f8E5M2(%arg0: tensor<f8E5M2>, %arg1: tensor<f8E5M2>) -> tensor<f8E5M2> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f8E5M2_v1>, !vhlo.tensor_v1<!vhlo.f8E5M2_v1>) -> !vhlo.tensor_v1<!vhlo.f8E5M2_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f8E5M2_v1>, !vhlo.tensor_v1<!vhlo.f8E5M2_v1>) -> !vhlo.tensor_v1<!vhlo.f8E5M2_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f8E5M2>, tensor<f8E5M2>) -> tensor<f8E5M2>
-   func.return %0 : tensor<f8E5M2>
- }
- 
- // CHECK-LABEL: "type_f8E4M3FNUZ"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_f8E4M3FNUZ(%arg0: tensor<f8E4M3FNUZ>, %arg1: tensor<f8E4M3FNUZ>) -> tensor<f8E4M3FNUZ> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f8E4M3FNUZ_v1>, !vhlo.tensor_v1<!vhlo.f8E4M3FNUZ_v1>) -> !vhlo.tensor_v1<!vhlo.f8E4M3FNUZ_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f8E4M3FNUZ_v1>, !vhlo.tensor_v1<!vhlo.f8E4M3FNUZ_v1>) -> !vhlo.tensor_v1<!vhlo.f8E4M3FNUZ_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f8E4M3FNUZ>, tensor<f8E4M3FNUZ>) -> tensor<f8E4M3FNUZ>
-   func.return %0 : tensor<f8E4M3FNUZ>
- }
- 
- // CHECK-LABEL: "type_f8E4M3B11FNUZ"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_f8E4M3B11FNUZ(%arg0: tensor<f8E4M3B11FNUZ>, %arg1: tensor<f8E4M3B11FNUZ>) -> tensor<f8E4M3B11FNUZ> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f8E4M3B11FNUZ_v1>, !vhlo.tensor_v1<!vhlo.f8E4M3B11FNUZ_v1>) -> !vhlo.tensor_v1<!vhlo.f8E4M3B11FNUZ_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f8E4M3B11FNUZ_v1>, !vhlo.tensor_v1<!vhlo.f8E4M3B11FNUZ_v1>) -> !vhlo.tensor_v1<!vhlo.f8E4M3B11FNUZ_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f8E4M3B11FNUZ>, tensor<f8E4M3B11FNUZ>) -> tensor<f8E4M3B11FNUZ>
-   func.return %0 : tensor<f8E4M3B11FNUZ>
- }
- 
- // CHECK-LABEL: "type_f8E5M2FNUZ"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_f8E5M2FNUZ(%arg0: tensor<f8E5M2FNUZ>, %arg1: tensor<f8E5M2FNUZ>) -> tensor<f8E5M2FNUZ> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f8E5M2FNUZ_v1>, !vhlo.tensor_v1<!vhlo.f8E5M2FNUZ_v1>) -> !vhlo.tensor_v1<!vhlo.f8E5M2FNUZ_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f8E5M2FNUZ_v1>, !vhlo.tensor_v1<!vhlo.f8E5M2FNUZ_v1>) -> !vhlo.tensor_v1<!vhlo.f8E5M2FNUZ_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f8E5M2FNUZ>, tensor<f8E5M2FNUZ>) -> tensor<f8E5M2FNUZ>
-   func.return %0 : tensor<f8E5M2FNUZ>
- }
- 
- // CHECK-LABEL: "type_bf16"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_bf16(%arg0: tensor<bf16>, %arg1: tensor<bf16>) -> tensor<bf16> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.bf16_v1>, !vhlo.tensor_v1<!vhlo.bf16_v1>) -> !vhlo.tensor_v1<!vhlo.bf16_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.bf16_v1>, !vhlo.tensor_v1<!vhlo.bf16_v1>) -> !vhlo.tensor_v1<!vhlo.bf16_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<bf16>, tensor<bf16>) -> tensor<bf16>
-   func.return %0 : tensor<bf16>
- }
- 
- // CHECK-LABEL: "type_f16"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_f16(%arg0: tensor<f16>, %arg1: tensor<f16>) -> tensor<f16> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f16_v1>, !vhlo.tensor_v1<!vhlo.f16_v1>) -> !vhlo.tensor_v1<!vhlo.f16_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f16_v1>, !vhlo.tensor_v1<!vhlo.f16_v1>) -> !vhlo.tensor_v1<!vhlo.f16_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f16>, tensor<f16>) -> tensor<f16>
-   func.return %0 : tensor<f16>
- }
- 
- // CHECK-LABEL: "type_f32"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_f32(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "type_f64"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_f64(%arg0: tensor<f64>, %arg1: tensor<f64>) -> tensor<f64> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f64_v1>, !vhlo.tensor_v1<!vhlo.f64_v1>) -> !vhlo.tensor_v1<!vhlo.f64_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f64_v1>, !vhlo.tensor_v1<!vhlo.f64_v1>) -> !vhlo.tensor_v1<!vhlo.f64_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f64>, tensor<f64>) -> tensor<f64>
-   func.return %0 : tensor<f64>
- }
- 
- // CHECK-LABEL: "type_complex_f32"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_complex_f32(%arg0: tensor<complex<f32>>, %arg1: tensor<complex<f32>>) -> tensor<complex<f32>> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>, !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>) -> !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>, !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>) -> !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<complex<f32>>, tensor<complex<f32>>) -> tensor<complex<f32>>
-   func.return %0 : tensor<complex<f32>>
- }
- 
- // CHECK-LABEL: "type_complex_f64"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_complex_f64(%arg0: tensor<complex<f64>>, %arg1: tensor<complex<f64>>) -> tensor<complex<f64>> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f64_v1>>, !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f64_v1>>) -> !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f64_v1>>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f64_v1>>, !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f64_v1>>) -> !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f64_v1>>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<complex<f64>>, tensor<complex<f64>>) -> tensor<complex<f64>>
-   func.return %0 : tensor<complex<f64>>
- }
- 
- // CHECK-LABEL: "type_dynamism_ranked"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @type_dynamism_ranked(%arg0: tensor<?xf32>) -> tensor<?xf32> {
--  // CHECK: "vhlo.abs_v1"(%arg0) : (!vhlo.tensor_v1<?x!vhlo.f32_v1>) -> !vhlo.tensor_v1<?x!vhlo.f32_v1>
-+  // CHECK: "vhlo.abs_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<?x!vhlo.f32_v1>) -> !vhlo.tensor_v1<?x!vhlo.f32_v1>
-   %0 = "stablehlo.abs"(%arg0) : (tensor<?xf32>) -> tensor<?xf32>
-   func.return %0 : tensor<?xf32>
- }
- 
- // CHECK-LABEL: "type_per_tensor_quantization"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_per_tensor_quantization(%arg0: tensor<!quant.uniform<i8:f32, 34.0:16>>, %arg1: tensor<!quant.uniform<i8:f32, 34.0:16>>) -> tensor<!quant.uniform<i8:f32, 34.0:16>> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>, !vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>) -> !vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>, !vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>) -> !vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<!quant.uniform<i8:f32, 34.0:16>>, tensor<!quant.uniform<i8:f32, 34.0:16>>) -> tensor<!quant.uniform<i8:f32, 34.0:16>>
-   func.return %0 : tensor<!quant.uniform<i8:f32, 34.0:16>>
- }
- 
- // CHECK-LABEL: "type_per_axis_quantization"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @type_per_axis_quantization(%arg0: tensor<2x!quant.uniform<i8:f32:0, {34.0:16, 34.0:16}>>) -> tensor<2x!quant.uniform<i8:f32:0, {34.0:16, 34.0:16}>> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg0) : (!vhlo.tensor_v1<2x!vhlo.quant_per_axis_v1<!vhlo.i8_v1:!vhlo.f32_v1, 0, 3.400000e+01, 3.400000e+01, 16, 16, -128:127, 1>>, !vhlo.tensor_v1<2x!vhlo.quant_per_axis_v1<!vhlo.i8_v1:!vhlo.f32_v1, 0, 3.400000e+01, 3.400000e+01, 16, 16, -128:127, 1>>) -> !vhlo.tensor_v1<2x!vhlo.quant_per_axis_v1<!vhlo.i8_v1:!vhlo.f32_v1, 0, 3.400000e+01, 3.400000e+01, 16, 16, -128:127, 1>>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG0]]) : (!vhlo.tensor_v1<2x!vhlo.quant_per_axis_v1<!vhlo.i8_v1:!vhlo.f32_v1, 0, 3.400000e+01, 3.400000e+01, 16, 16, -128:127, 1>>, !vhlo.tensor_v1<2x!vhlo.quant_per_axis_v1<!vhlo.i8_v1:!vhlo.f32_v1, 0, 3.400000e+01, 3.400000e+01, 16, 16, -128:127, 1>>) -> !vhlo.tensor_v1<2x!vhlo.quant_per_axis_v1<!vhlo.i8_v1:!vhlo.f32_v1, 0, 3.400000e+01, 3.400000e+01, 16, 16, -128:127, 1>>
-   %0 = stablehlo.add %arg0, %arg0 : tensor<2x!quant.uniform<i8:f32:0, {34.0:16, 34.0:16}>>
-   func.return %0 : tensor<2x!quant.uniform<i8:f32:0, {34.0:16, 34.0:16}>>
- }
- 
- //       CHECK: function_type = #vhlo.type_v1<!vhlo.func_v1<(!vhlo.token_v1) -> !vhlo.token_v1>>
- // CHECK-LABEL: "type_token_callee"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @type_token_callee(%arg0: !stablehlo.token) -> !stablehlo.token {
--  // CHECK: "vhlo.return_v1"(%arg0) : (!vhlo.token_v1) -> ()
-+  // CHECK: "vhlo.return_v1"(%[[ARG0]]) : (!vhlo.token_v1) -> ()
-   return %arg0 : !stablehlo.token
- }
- 
- //       CHECK: function_type = #vhlo.type_v1<!vhlo.func_v1<(!vhlo.token_v1) -> !vhlo.token_v1>>
- // CHECK-LABEL: "type_token_caller"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @type_token_caller(%arg0: !stablehlo.token) -> !stablehlo.token {
--  // CHECK:      "vhlo.call_v1"(%arg0) <{callee = #vhlo.string_v1<"type_token_callee">}
-+  // CHECK:      "vhlo.call_v1"(%[[ARG0]]) <{callee = #vhlo.string_v1<"type_token_callee">}
-   // CHECK-SAME: (!vhlo.token_v1) -> !vhlo.token_v1
-   %0 = func.call @type_token_callee(%arg0) : (!stablehlo.token) -> !stablehlo.token
-   return %0 : !stablehlo.token
- }
- 
- // CHECK-LABEL: "type_tuple"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @type_tuple(%arg0: tuple<tensor<f32>>) -> tuple<!stablehlo.token> {
-   %0 = "stablehlo.custom_call"(%arg0) {
-     call_target_name = "foo"
-diff --ruN a/stablehlo/stablehlo/tests/vhlo/stablehlo_legalize_to_vhlo.0_19_0.mlir b/stablehlo/stablehlo/tests/vhlo/stablehlo_legalize_to_vhlo.0_19_0.mlir
---- stablehlo/stablehlo/tests/vhlo/stablehlo_legalize_to_vhlo.0_19_0.mlir
-+++ stablehlo/stablehlo/tests/vhlo/stablehlo_legalize_to_vhlo.0_19_0.mlir
-@@ -13,6 +13,7 @@
- // ============ ATTRIBUTES ============
- 
- // CHECK-LABEL: "attr_comparison_direction_eq"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_direction_eq(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     // CHECK: comparison_direction = #vhlo<comparison_direction_v1 EQ>
-@@ -22,6 +23,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_direction_ne"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_direction_ne(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     // CHECK: comparison_direction = #vhlo<comparison_direction_v1 NE>
-@@ -31,6 +33,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_direction_ge"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_direction_ge(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     // CHECK: comparison_direction = #vhlo<comparison_direction_v1 GE>
-@@ -40,6 +43,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_direction_gt"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_direction_gt(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     // CHECK: comparison_direction = #vhlo<comparison_direction_v1 GT>
-@@ -49,6 +53,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_direction_le"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_direction_le(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     // CHECK: comparison_direction = #vhlo<comparison_direction_v1 LE>
-@@ -58,6 +63,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_direction_lt"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_direction_lt(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     // CHECK: comparison_direction = #vhlo<comparison_direction_v1 LT>
-@@ -67,6 +73,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_type_notype"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_type_notype(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     comparison_direction = #stablehlo<comparison_direction EQ>
-@@ -76,6 +83,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_type_float"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_type_float(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     comparison_direction = #stablehlo<comparison_direction EQ>,
-@@ -86,6 +94,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_type_totalorder"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_type_totalorder(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     comparison_direction = #stablehlo<comparison_direction EQ>,
-@@ -96,6 +105,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_type_signed"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_type_signed(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     comparison_direction = #stablehlo<comparison_direction EQ>,
-@@ -106,6 +116,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_type_unsigned"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_type_unsigned(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     comparison_direction = #stablehlo<comparison_direction EQ>,
-@@ -118,6 +129,7 @@
- // ConvDimensionNumbers aka #stablehlo.conv is covered below.
- 
- // CHECK-LABEL: "attr_custom_call_api_version_unspecified"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_custom_call_api_version_unspecified(%arg0: tensor<f32>) -> tensor<f32> {
-   %0 = "stablehlo.custom_call"(%arg0) {
-     call_target_name = "foo",
-@@ -128,6 +140,7 @@
- }
- 
- // CHECK-LABEL: "attr_custom_call_api_version_original"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_custom_call_api_version_original(%arg0: tensor<f32>) -> tensor<f32> {
-   %0 = "stablehlo.custom_call"(%arg0) {
-     call_target_name = "foo",
-@@ -138,6 +151,7 @@
- }
- 
- // CHECK-LABEL: "attr_custom_call_api_version_status_returning"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_custom_call_api_version_status_returning(%arg0: tensor<f32>) -> tensor<f32> {
-   %0 = "stablehlo.custom_call"(%arg0) {
-     call_target_name = "foo",
-@@ -148,6 +162,7 @@
- }
- 
- // CHECK-LABEL: "attr_custom_call_api_version_status_returning_unified"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_custom_call_api_version_status_returning_unified(%arg0: tensor<f32>) -> tensor<f32> {
-   %0 = "stablehlo.custom_call"(%arg0) {
-     call_target_name = "foo",
-@@ -166,6 +181,7 @@
- // DotDimensionNumbers aka #stablehlo.dot is covered below.
- 
- // CHECK-LABEL: "attr_fft_type_fft"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_fft_type_fft(%arg0: tensor<16xcomplex<f32>>) -> tensor<16xcomplex<f32>> {
-   %0 = "stablehlo.fft"(%arg0) {
-     // CHECK: fft_type = #vhlo<fft_type_v1 FFT>
-@@ -176,6 +192,7 @@
- }
- 
- // CHECK-LABEL: "attr_fft_type_ifft"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_fft_type_ifft(%arg0: tensor<16xcomplex<f32>>) -> tensor<16xcomplex<f32>> {
-   %0 = "stablehlo.fft"(%arg0) {
-     // CHECK: fft_type = #vhlo<fft_type_v1 IFFT>
-@@ -186,6 +203,7 @@
- }
- 
- // CHECK-LABEL: "attr_fft_type_rfft"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_fft_type_rfft(%arg0: tensor<16xf32>) -> tensor<9xcomplex<f32>> {
-   %0 = "stablehlo.fft"(%arg0) {
-     // CHECK: fft_type = #vhlo<fft_type_v1 RFFT>
-@@ -196,6 +214,7 @@
- }
- 
- // CHECK-LABEL: "attr_fft_type_irfft"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_fft_type_irfft(%arg0: tensor<9xcomplex<f32>>) -> tensor<16xf32> {
-   %0 = "stablehlo.fft"(%arg0) {
-     // CHECK: fft_type = #vhlo<fft_type_v1 IRFFT>
-@@ -208,6 +227,7 @@
- // GatherDimensionNumbers aka #stablehlo.gather is covered below.
- 
- // CHECK-LABEL: "attr_precision_config_default"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_precision_config_default(%arg0: tensor<8x16xf32>, %arg1: tensor<16x8xf32>) -> tensor<8x8xf32> {
-   %0 = "stablehlo.dot"(%arg0, %arg1) {
-     // CHECK: precision_config = #vhlo.array_v1<[#vhlo<precision_v1 DEFAULT>, #vhlo<precision_v1 DEFAULT>]>
-@@ -216,6 +236,7 @@
- }
- 
- // CHECK-LABEL: "attr_precision_config_high"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_precision_config_high(%arg0: tensor<8x16xf32>, %arg1: tensor<16x8xf32>) -> tensor<8x8xf32> {
-   %0 = "stablehlo.dot"(%arg0, %arg1) {
-     // CHECK: precision_config = #vhlo.array_v1<[#vhlo<precision_v1 HIGH>, #vhlo<precision_v1 HIGH>]>
-@@ -225,6 +246,7 @@
- }
- 
- // CHECK-LABEL: "attr_precision_config_highest"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_precision_config_highest(%arg0: tensor<8x16xf32>, %arg1: tensor<16x8xf32>) -> tensor<8x8xf32> {
-   %0 = "stablehlo.dot"(%arg0, %arg1) {
-     // CHECK: precision_config = #vhlo.array_v1<[#vhlo<precision_v1 HIGHEST>, #vhlo<precision_v1 HIGHEST>]>
-@@ -234,6 +256,7 @@
- }
- 
- // CHECK-LABEL: "attr_rng_algorithm_default"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_rng_algorithm_default(%arg0: tensor<f32>) -> (tensor<f32>, tensor<f32>) {
-   %0:2 = "stablehlo.rng_bit_generator"(%arg0) {
-     // CHECK: rng_algorithm = #vhlo<rng_algorithm_v1 DEFAULT>
-@@ -243,6 +266,7 @@
- }
- 
- // CHECK-LABEL: "attr_rng_algorithm_three_fry"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_rng_algorithm_three_fry(%arg0: tensor<f32>) -> (tensor<f32>, tensor<f32>) {
-   %0:2 = "stablehlo.rng_bit_generator"(%arg0) {
-     // CHECK: rng_algorithm = #vhlo<rng_algorithm_v1 THREE_FRY>
-@@ -252,6 +276,7 @@
- }
- 
- // CHECK-LABEL: "attr_rng_algorithm_philox"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_rng_algorithm_philox(%arg0: tensor<f32>) -> (tensor<f32>, tensor<f32>) {
-   %0:2 = "stablehlo.rng_bit_generator"(%arg0) {
-     // CHECK: rng_algorithm = #vhlo<rng_algorithm_v1 PHILOX>
-@@ -261,6 +286,7 @@
- }
- 
- // CHECK-LABEL: "attr_rng_distribution_uniform"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @attr_rng_distribution_uniform(%arg0: tensor<f32>, %arg1: tensor<f32>, %arg2: tensor<0xindex>) -> tensor<f32> {
-   %0 = "stablehlo.rng"(%arg0, %arg1, %arg2) {
-     // CHECK: rng_distribution = #vhlo<rng_distribution_v1 UNIFORM>
-@@ -270,6 +296,7 @@
- }
- 
- // CHECK-LABEL: "attr_rng_distribution_normal"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @attr_rng_distribution_normal(%arg0: tensor<f32>, %arg1: tensor<f32>, %arg2: tensor<0xindex>) -> tensor<f32> {
-   %0 = "stablehlo.rng"(%arg0, %arg1, %arg2) {
-     // CHECK: rng_distribution = #vhlo<rng_distribution_v1 NORMAL>
-@@ -281,6 +308,7 @@
- // ScatterDimensionNumbers aka #stablehlo.scatter is covered below.
- 
- // CHECK-LABEL: "attr_transpose_no_transpose"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_transpose_no_transpose(%arg0: tensor<16x16xf32>, %arg1: tensor<16x16xf32>) ->  tensor<16x16xf32> {
-   %0 = "stablehlo.triangular_solve"(%arg0, %arg1) {
-     left_side = true,
-@@ -293,6 +321,7 @@
- }
- 
- // CHECK-LABEL: "attr_transpose_transpose"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_transpose_transpose(%arg0: tensor<16x16xf32>, %arg1: tensor<16x16xf32>) ->  tensor<16x16xf32> {
-   %0 = "stablehlo.triangular_solve"(%arg0, %arg1) {
-     left_side = true,
-@@ -305,6 +334,7 @@
- }
- 
- // CHECK-LABEL: "attr_transpose_adjoint"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_transpose_adjoint(%arg0: tensor<16x16xf32>, %arg1: tensor<16x16xf32>) ->  tensor<16x16xf32> {
-   %0 = "stablehlo.triangular_solve"(%arg0, %arg1) {
-     left_side = true,
-@@ -319,10 +349,9 @@
- // TypeExtensionsAttr aka #stablehlo.type_extensions is covered below.
- 
- // CHECK-LABEL: "attr_type_extensions_bounds"
--func.func @attr_type_extensions_bounds(
--    %arg0: tensor<?x?xf32, #stablehlo.type_extensions<bounds = [16, ?]>>)
--    -> tensor<?x?xf32, #stablehlo.type_extensions<bounds = [16, ?]>> {
--  // CHECK: "vhlo.return_v1"(%arg0) : (!vhlo.tensor_v1<?x?x!vhlo.f32_v1, #vhlo.type_extensions_v1<bounds = [16, ?]>>) -> ()
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @attr_type_extensions_bounds(%arg0: tensor<?x?xf32, #stablehlo.type_extensions<bounds = [16, ?]>>) -> tensor<?x?xf32, #stablehlo.type_extensions<bounds = [16, ?]>> {
-+  // CHECK: "vhlo.return_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<?x?x!vhlo.f32_v1, #vhlo.type_extensions_v1<bounds = [16, ?]>>) -> ()
-   func.return %arg0 : tensor<?x?xf32, #stablehlo.type_extensions<bounds = [16, ?]>>
- }
- 
-@@ -330,8 +359,9 @@
- // ============ DEFAULTS ============
- 
- // CHECK-LABEL: "default_all_gather"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_all_gather(%arg0: tensor<16x8xf32>) -> tensor<16x16xf32> {
--  //               CHECK: "vhlo.all_gather_v1"(%arg0) <{
-+  //               CHECK: "vhlo.all_gather_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   all_gather_dim = #vhlo.integer_v1<1 : i64>
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0], [1]]> : tensor<2x1xi64>>,
-@@ -345,8 +375,9 @@
- }
- 
- // CHECK-LABEL: "default_all_reduce"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_all_reduce(%arg0: tensor<f32>) -> tensor<f32> {
--  //               CHECK: "vhlo.all_reduce_v1"(%arg0)
-+  //               CHECK: "vhlo.all_reduce_v1"(%[[ARG0]])
-   //          CHECK-SAME: <{
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0], [1]]> : tensor<2x1xi64>>,
-@@ -368,8 +399,9 @@
- }
- 
- // CHECK-LABEL: "default_all_to_all"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_all_to_all(%arg0: tensor<4x16xf32>) -> tensor<16x4xf32> {
--  //               CHECK: "vhlo.all_to_all_v1"(%arg0) <{
-+  //               CHECK: "vhlo.all_to_all_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   //          CHECK-SAME:   concat_dimension = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0, 1, 2, 3]]> : tensor<1x4xi64>>,
-@@ -386,8 +418,9 @@
- }
- 
- // CHECK-LABEL: "default_cholesky"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_cholesky(%arg0: tensor<1x16x16xf32>) -> tensor<1x16x16xf32> {
--  //      CHECK: "vhlo.cholesky_v1"(%arg0) <{
-+  //      CHECK: "vhlo.cholesky_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   lower = #vhlo.bool_v1<false>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<1x16x16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<1x16x16x!vhlo.f32_v1>
-   %0 = "stablehlo.cholesky"(%arg0) : (tensor<1x16x16xf32>) -> tensor<1x16x16xf32>
-@@ -395,8 +428,9 @@
- }
- 
- // CHECK-LABEL: "default_collective_permute"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_collective_permute(%arg0: tensor<16x8xf32>) -> tensor<16x8xf32> {
--  //               CHECK: "vhlo.collective_permute_v1"(%arg0) <{
-+  //               CHECK: "vhlo.collective_permute_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME{LITERAL}:   source_target_pairs = #vhlo.tensor_v1<dense<[[0, 1], [1, 2], [2, 3]]> : tensor<3x2xi64>>
-   //          CHECK-SAME: }> : (!vhlo.tensor_v1<16x8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x8x!vhlo.f32_v1>
-@@ -407,8 +441,9 @@
- }
- 
- // CHECK-LABEL: "default_collective_broadcast"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_collective_broadcast(%arg0: tensor<16x8xf32>) -> tensor<16x8xf32> {
--  //               CHECK: "vhlo.collective_broadcast_v1"(%arg0) <{
-+  //               CHECK: "vhlo.collective_broadcast_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0, 1]]> : tensor<1x2xi64>>
-   //          CHECK-SAME: }> : (!vhlo.tensor_v1<16x8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x8x!vhlo.f32_v1>
-@@ -419,8 +454,9 @@
- }
- 
- // CHECK-LABEL: "default_compare"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @default_compare(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
--  //      CHECK: "vhlo.compare_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.compare_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   compare_type = #vhlo<comparison_type_v1 NOTYPE>,
-   // CHECK-SAME:   comparison_direction = #vhlo<comparison_direction_v1 EQ>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-@@ -431,8 +467,9 @@
- }
- 
- // CHECK-LABEL: "default_composite"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_composite(%arg0: tensor<f32>) -> tensor<f32> {
--  //               CHECK: "vhlo.composite_v1"(%arg0) <{
-+  //               CHECK: "vhlo.composite_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   composite_attributes = #vhlo.dict_v1<{}>
-   //          CHECK-SAME:   decomposition = #vhlo.string_v1<"composite_target">
-   //          CHECK-SAME:   name = #vhlo.string_v1<"stablehlo.composite_target">
-@@ -446,8 +483,9 @@
- }
- 
- // CHECK-LABEL: "default_convolution"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @default_convolution(%arg0: tensor<1x8x8x207xf32>, %arg1: tensor<3x3x207x16xf32>) -> tensor<1x6x6x16xf32> {
--  //      CHECK: "vhlo.convolution_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.convolution_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   batch_group_count = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   feature_group_count = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   input_batch_dimension = #vhlo.integer_v1<0 : i64>,
-@@ -475,8 +513,9 @@
- }
- 
- // CHECK-LABEL: "default_custom_call"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_custom_call(%arg0: tensor<f32>) -> tensor<f32> {
--  //      CHECK: "vhlo.custom_call_v1"(%arg0) <{
-+  //      CHECK: "vhlo.custom_call_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   api_version = #vhlo<api_version_v1 API_VERSION_ORIGINAL>,
-   // CHECK-SAME:   backend_config = #vhlo.string_v1<"">,
-   // CHECK-SAME:   call_target_name = #vhlo.string_v1<"foo">,
-@@ -493,8 +532,9 @@
- }
- 
- // CHECK-LABEL: "default_dot_general"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @default_dot_general(%arg0: tensor<8x8x16xf32>, %arg1: tensor<8x16x8xf32>) -> tensor<8x8x8xf32> {
--  //      CHECK: "vhlo.dot_general_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.dot_general_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   lhs_batching_dimensions = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>,
-   // CHECK-SAME:   lhs_contracting_dimensions = #vhlo.tensor_v1<dense<2> : tensor<1xi64>>,
-   // CHECK-SAME:   precision_config = #vhlo.array_v1<[#vhlo<precision_v1 DEFAULT>, #vhlo<precision_v1 DEFAULT>]>,
-@@ -513,8 +553,9 @@
- }
- 
- // CHECK-LABEL: "default_dot"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @default_dot(%arg0: tensor<8x16xf32>, %arg1: tensor<16x8xf32>) -> tensor<8x8xf32> {
--  //      CHECK: "vhlo.dot_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.dot_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   precision_config = #vhlo.array_v1<[#vhlo<precision_v1 DEFAULT>, #vhlo<precision_v1 DEFAULT>]>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<8x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<8x8x!vhlo.f32_v1>
-   %0 = "stablehlo.dot"(%arg0, %arg1) : (tensor<8x16xf32>, tensor<16x8xf32>) -> tensor<8x8xf32>
-@@ -522,8 +563,9 @@
- }
- 
- // CHECK-LABEL: "default_dynamic_broadcast_in_dim"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @default_dynamic_broadcast_in_dim(%arg0: tensor<?x?xf32>, %arg1: tensor<2xindex>) -> tensor<?x?xf32> {
--  //      CHECK: "vhlo.dynamic_broadcast_in_dim_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.dynamic_broadcast_in_dim_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   broadcast_dimensions = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-   // CHECK-SAME:   known_expanding_dimensions = #vhlo.tensor_v1<dense<> : tensor<0xi64>>,
-   // CHECK-SAME:   known_nonexpanding_dimensions = #vhlo.tensor_v1<dense<> : tensor<0xi64>>
-@@ -535,8 +577,9 @@
- }
- 
- // CHECK-LABEL: "default_dynamic_conv"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @default_dynamic_conv(%arg0: tensor<1x8x8x207xf32>, %arg1: tensor<3x3x207x16xf32>, %arg2: tensor<4xi32>) -> tensor<1x?x?x16xf32> {
--  //      CHECK: "vhlo.dynamic_conv_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.dynamic_conv_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   batch_group_count = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   feature_group_count = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   input_batch_dimension = #vhlo.integer_v1<0 : i64>,
-@@ -564,8 +607,9 @@
- }
- 
- // CHECK-LABEL: "default_dynamic_gather"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @default_dynamic_gather(%arg0 : tensor<2x4x9xf32>, %arg1 : tensor<1x5x2xi32>, %arg2 : tensor<3xi32>) -> tensor<1x5x8xf32> {
--  //      CHECK: "vhlo.dynamic_gather_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.dynamic_gather_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   collapsed_slice_dims = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-   // CHECK-SAME:   index_vector_dim = #vhlo.integer_v1<2 : i64>,
-   // CHECK-SAME:   indices_are_sorted = #vhlo.bool_v1<false>,
-@@ -591,15 +635,16 @@
-   // CHECK-SAME:   sym_name = #vhlo.string_v1<"default_func">,
-   // CHECK-SAME:   sym_visibility = #vhlo.string_v1<"">
-   // CHECK-SAME: }> ({
--  // CHECK-NEXT:   ^[[BB:bb.*]](%arg0: !vhlo.tensor_v1<!vhlo.f32_v1>):
--  // CHECK-NEXT:     "vhlo.return_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-+  // CHECK-NEXT:   ^[[BB:bb.*]](%[[ARG0:.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
-+  // CHECK-NEXT:     "vhlo.return_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-   // CHECK-NEXT: }) : () -> ()
-   func.return %arg0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "dynamic_gather"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @dynamic_gather(%arg0 : tensor<2x4x9xf32>, %arg1 : tensor<1x5x2xi32>) -> tensor<1x5x1xf32> {
--  //      CHECK: "vhlo.gather_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.gather_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   collapsed_slice_dims = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-   // CHECK-SAME:   index_vector_dim = #vhlo.integer_v1<2 : i64>,
-   // CHECK-SAME:   indices_are_sorted = #vhlo.bool_v1<false>,
-@@ -620,8 +665,9 @@
- }
- 
- // CHECK-LABEL: "default_infeed"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_infeed(%arg0: !stablehlo.token) -> (tensor<f32>, !stablehlo.token) {
--  //               CHECK: "vhlo.infeed_v1"(%arg0) <{
-+  //               CHECK: "vhlo.infeed_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   infeed_config = #vhlo.string_v1<"">,
-   // CHECK-SAME{LITERAL}:   layout = #vhlo.array_v1<[]>
-   //          CHECK-SAME: }> : (!vhlo.token_v1) -> (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.token_v1)
-@@ -630,8 +676,9 @@
- }
- 
- // CHECK-LABEL: "default_outfeed"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @default_outfeed(%arg0: tensor<f32>, %arg1: !stablehlo.token) -> !stablehlo.token {
--  //      CHECK: "vhlo.outfeed_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.outfeed_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   outfeed_config = #vhlo.string_v1<"">
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.token_v1) -> !vhlo.token_v1
-   %0 = "stablehlo.outfeed"(%arg0, %arg1) : (tensor<f32>, !stablehlo.token) -> !stablehlo.token
-@@ -639,8 +686,9 @@
- }
- 
- // CHECK-LABEL: "default_recv"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_recv(%arg0: !stablehlo.token) -> (tensor<f32>, !stablehlo.token) {
--  //      CHECK: "vhlo.recv_v1"(%arg0) <{
-+  //      CHECK: "vhlo.recv_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME:   channel_type = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   is_host_transfer = #vhlo.bool_v1<false>
-@@ -652,8 +700,9 @@
- }
- 
- // CHECK-LABEL: "default_send"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @default_send(%arg0: tensor<f32>, %arg1: !stablehlo.token) -> !stablehlo.token {
--  //      CHECK: "vhlo.send_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.send_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME:   channel_type = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   is_host_transfer = #vhlo.bool_v1<false>
-@@ -665,8 +714,9 @@
- }
- 
- // CHECK-LABEL: "default_reduce_scatter"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_reduce_scatter(%arg0: tensor<16xf32>) -> tensor<16xf32> {
--  //               CHECK: "vhlo.reduce_scatter_v1"(%arg0) <{
-+  //               CHECK: "vhlo.reduce_scatter_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0], [1]]> : tensor<2x1xi64>>,
-   //          CHECK-SAME:   scatter_dimension = #vhlo.integer_v1<0 : i64>
-@@ -688,8 +738,9 @@
- }
- 
- // CHECK-LABEL: "default_reduce_window"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @default_reduce_window(%arg0: tensor<2x17x31x7xf32>, %arg1: tensor<f32>) -> tensor<2x16x30x7xf32> {
--  //               CHECK: "vhlo.reduce_window_v1"(%arg0, %arg1)  <{
-+  //               CHECK: "vhlo.reduce_window_v1"(%[[ARG0]], %[[ARG1]])  <{
-   //          CHECK-SAME:   base_dilations = #vhlo.tensor_v1<dense<1> : tensor<4xi64>>,
-   // CHECK-SAME{LITERAL}:   padding = #vhlo.tensor_v1<dense<0> : tensor<4x2xi64>>,
-   //          CHECK-SAME:   window_dilations = #vhlo.tensor_v1<dense<1> : tensor<4xi64>>,
-@@ -711,8 +762,9 @@
- }
- 
- // CHECK-LABEL: "default_scatter"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @default_scatter(%arg0: tensor<200x100x300xf32>, %arg1: tensor<10x2xi32>, %arg2: tensor<10x300xf32>) -> tensor<200x100x300xf32> {
--  //      CHECK: "vhlo.scatter_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.scatter_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   index_vector_dim = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   indices_are_sorted = #vhlo.bool_v1<false>,
-   // CHECK-SAME:   inserted_window_dims = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-@@ -740,8 +792,9 @@
- }
- 
- // CHECK-LABEL: "default_select_and_scatter"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @default_select_and_scatter(%arg0: tensor<10x24x24x64xf32>, %arg1: tensor<10x23x23x64xf32>, %arg2: tensor<f32>) -> tensor<10x24x24x64xf32> {
--  //      CHECK: "vhlo.select_and_scatter_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.select_and_scatter_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   padding = #vhlo.tensor_v1<dense<0> : tensor<4x2xi64>>,
-   // CHECK-SAME:   window_dimensions = #vhlo.tensor_v1<dense<[1, 2, 2, 1]> : tensor<4xi64>>,
-   // CHECK-SAME:   window_strides = #vhlo.tensor_v1<dense<1> : tensor<4xi64>>
-@@ -769,8 +822,9 @@
- }
- 
- // CHECK-LABEL: "default_sort"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_sort(%arg0: tensor<16xf32>) -> tensor<16xf32> {
--  //      CHECK: "vhlo.sort_v1"(%arg0) <{
-+  //      CHECK: "vhlo.sort_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   dimension = #vhlo.integer_v1<-1 : i64>
-   // CHECK-SAME:   is_stable = #vhlo.bool_v1<false>
-   // CHECK-SAME: }> ({
-@@ -789,29 +843,33 @@
- // ============ OPS ============
- 
- // CHECK-LABEL: "op_abs"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_abs(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.abs_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.abs_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.abs"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_add"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_add(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_after_all"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_after_all(%arg0: !stablehlo.token) -> !stablehlo.token {
--  // CHECK: "vhlo.after_all_v1"(%arg0) : (!vhlo.token_v1) -> !vhlo.token_v1
-+  // CHECK: "vhlo.after_all_v1"(%[[ARG0]]) : (!vhlo.token_v1) -> !vhlo.token_v1
-   %0 = "stablehlo.after_all"(%arg0) : (!stablehlo.token) -> !stablehlo.token
-   func.return %0 : !stablehlo.token
- }
- 
- // CHECK-LABEL: "op_all_gather"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_all_gather(%arg0: tensor<16x8xf32>) -> tensor<16x16xf32> {
--  //               CHECK: "vhlo.all_gather_v1"(%arg0) <{
-+  //               CHECK: "vhlo.all_gather_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   all_gather_dim = #vhlo.integer_v1<1 : i64>
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0], [1]]> : tensor<2x1xi64>>,
-@@ -827,8 +885,9 @@
- }
- 
- // CHECK-LABEL: "op_all_reduce"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_all_reduce(%arg0: tensor<f32>) -> tensor<f32> {
--  //               CHECK: "vhlo.all_reduce_v1"(%arg0) <{
-+  //               CHECK: "vhlo.all_reduce_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0], [1]]> : tensor<2x1xi64>>,
-   //          CHECK-SAME:   use_global_device_ids = #vhlo.bool_v1<true>
-@@ -851,7 +910,7 @@
- 
- // CHECK-LABEL: "op_all_reduce_with_promotable_types"
- func.func @op_all_reduce_with_promotable_types(%operand: tensor<f32>) -> tensor<f64> {
--  //  CHECK: "vhlo.all_reduce_v1"(%arg0)
-+  //  CHECK: "vhlo.all_reduce_v1"(%[[ARG0:.*]])
-   //  CHECK:   ^[[BB:bb.*]](%[[ARG1:arg.*]]: !vhlo.tensor_v1<!vhlo.f64_v1>, %[[ARG2:arg.*]]: !vhlo.tensor_v1<!vhlo.f64_v1>):
-   //  CHECK:     "vhlo.return_v1"(%[[VAL1:.*]]) : (!vhlo.tensor_v1<!vhlo.f64_v1>) -> ()
-   //  CHECK: }) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f64_v1>
-@@ -869,8 +928,9 @@
- }
- 
- // CHECK-LABEL: "op_all_to_all"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_all_to_all(%arg0: tensor<4x16xf32>) -> tensor<16x4xf32> {
--  //               CHECK: "vhlo.all_to_all_v1"(%arg0) <{
-+  //               CHECK: "vhlo.all_to_all_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<1 : i64>,
-   //          CHECK-SAME:   concat_dimension = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0, 1, 2, 3]]> : tensor<1x4xi64>>,
-@@ -888,22 +948,25 @@
- }
- 
- // CHECK-LABEL: "op_and"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_and(%arg0: tensor<i1>, %arg1: tensor<i1>) -> tensor<i1> {
--  // CHECK: "vhlo.and_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-+  // CHECK: "vhlo.and_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-   %0 = "stablehlo.and"(%arg0, %arg1) : (tensor<i1>, tensor<i1>) -> tensor<i1>
-   func.return %0 : tensor<i1>
- }
- 
- // CHECK-LABEL: "op_atan2"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_atan2(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.atan2_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.atan2_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.atan2"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_batch_norm_grad"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}}, %[[ARG3:.*]]: {{.*}}, %[[ARG4:.*]]: {{.*}})
- func.func @op_batch_norm_grad(%arg0: tensor<16x16x16x16xf32>, %arg1: tensor<16xf32>, %arg2: tensor<16xf32>, %arg3: tensor<16xf32>, %arg4: tensor<16x16x16x16xf32>) -> (tensor<16x16x16x16xf32>, tensor<16xf32>, tensor<16xf32>) {
--  //      CHECK: "vhlo.batch_norm_grad_v1"(%arg0, %arg1, %arg2, %arg3, %arg4) <{
-+  //      CHECK: "vhlo.batch_norm_grad_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]], %[[ARG3]], %[[ARG4]]) <{
-   // CHECK-SAME:   epsilon = #vhlo.float_v1<1.000000e-03 : !vhlo.f32_v1>,
-   // CHECK-SAME:   feature_index = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<16x16x16x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x16x16x16x!vhlo.f32_v1>) -> (!vhlo.tensor_v1<16x16x16x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>)
-@@ -915,8 +978,9 @@
- }
- 
- // CHECK-LABEL: "op_batch_norm_inference"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}}, %[[ARG3:.*]]: {{.*}}, %[[ARG4:.*]]: {{.*}})
- func.func @op_batch_norm_inference(%arg0: tensor<16x16x16x16xf32>, %arg1: tensor<16xf32>, %arg2: tensor<16xf32>, %arg3: tensor<16xf32>, %arg4: tensor<16xf32>) -> tensor<16x16x16x16xf32> {
--  //      CHECK: "vhlo.batch_norm_inference_v1"(%arg0, %arg1, %arg2, %arg3, %arg4) <{
-+  //      CHECK: "vhlo.batch_norm_inference_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]], %[[ARG3]], %[[ARG4]]) <{
-   // CHECK-SAME:   epsilon = #vhlo.float_v1<1.000000e-03 : !vhlo.f32_v1>,
-   // CHECK-SAME:   feature_index = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<16x16x16x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x16x16x16x!vhlo.f32_v1>
-@@ -928,8 +992,9 @@
- }
- 
- // CHECK-LABEL: "op_batch_norm_training"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_batch_norm_training(%arg0: tensor<16x16x16x16xf32>, %arg1: tensor<16xf32>, %arg2: tensor<16xf32>) -> (tensor<16x16x16x16xf32>, tensor<16xf32>, tensor<16xf32>) {
--  //      CHECK: "vhlo.batch_norm_training_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.batch_norm_training_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   epsilon = #vhlo.float_v1<1.000000e-03 : !vhlo.f32_v1>,
-   // CHECK-SAME:   feature_index = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<16x16x16x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>) -> (!vhlo.tensor_v1<16x16x16x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>)
-@@ -941,15 +1006,17 @@
- }
- 
- // CHECK-LABEL: "op_bitcast_convert"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_bitcast_convert(%arg0: tensor<i32>) -> tensor<f32> {
--  // CHECK: "vhlo.bitcast_convert_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.bitcast_convert_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.bitcast_convert"(%arg0) : (tensor<i32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_broadcast_in_dim"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_broadcast_in_dim(%arg0: tensor<16xf32>) -> tensor<16x16xf32> {
--  //      CHECK: "vhlo.broadcast_in_dim_v1"(%arg0) <{
-+  //      CHECK: "vhlo.broadcast_in_dim_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   broadcast_dimensions = #vhlo.tensor_v1<dense<1> : tensor<1xi64>>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x16x!vhlo.f32_v1>
-   %0 = "stablehlo.broadcast_in_dim"(%arg0) {
-@@ -959,8 +1026,9 @@
- }
- 
- // CHECK-LABEL: "op_broadcast"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_broadcast(%arg0: tensor<16xf32>) -> tensor<16x16xf32> {
--  //      CHECK: "vhlo.broadcast_v1"(%arg0) <{
-+  //      CHECK: "vhlo.broadcast_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   broadcast_sizes = #vhlo.tensor_v1<dense<16> : tensor<1xi64>>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x16x!vhlo.f32_v1>
-   %0 = "stablehlo.broadcast"(%arg0) {
-@@ -970,9 +1038,10 @@
- }
- 
- // CHECK-LABEL: "op_case"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_case(%arg0: tensor<i32>, %arg1: tensor<f32>) -> tensor<f32> {
--  //      CHECK: "vhlo.case_v1"(%arg0) ({
--  // CHECK-NEXT:   "vhlo.return_v1"(%arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-+  //      CHECK: "vhlo.case_v1"(%[[ARG0]]) ({
-+  // CHECK-NEXT:   "vhlo.return_v1"(%[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-   // CHECK-NEXT: }) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.case"(%arg0) ({
-     "stablehlo.return"(%arg1) : (tensor<f32>) -> ()
-@@ -981,22 +1050,25 @@
- }
- 
- // CHECK-LABEL: "op_cbrt"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_cbrt(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.cbrt_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.cbrt_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.cbrt"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_ceil"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_ceil(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.ceil_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.ceil_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.ceil"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_cholesky"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_cholesky(%arg0: tensor<1x16x16xf32>) -> tensor<1x16x16xf32> {
--  //      CHECK: "vhlo.cholesky_v1"(%arg0) <{
-+  //      CHECK: "vhlo.cholesky_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   lower = #vhlo.bool_v1<true>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<1x16x16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<1x16x16x!vhlo.f32_v1>
-   %0 = "stablehlo.cholesky"(%arg0) {
-@@ -1006,22 +1078,25 @@
- }
- 
- // CHECK-LABEL: "op_clamp"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_clamp(%arg0: tensor<f32>, %arg1: tensor<f32>, %arg2: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.clamp_v1"(%arg0, %arg1, %arg2) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.clamp_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.clamp"(%arg0, %arg1, %arg2) : (tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_count_leading_zeros"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_count_leading_zeros(%arg0: tensor<i32>) -> tensor<i32> {
--  // CHECK: "vhlo.count_leading_zeros_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-+  // CHECK: "vhlo.count_leading_zeros_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-   %0 = "stablehlo.count_leading_zeros"(%arg0) : (tensor<i32>) -> tensor<i32>
-   func.return %0 : tensor<i32>
- }
- 
- // CHECK-LABEL: "op_collective_permute"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_collective_permute(%arg0: tensor<16x8xf32>) -> tensor<16x8xf32> {
--  //               CHECK: "vhlo.collective_permute_v1"(%arg0) <{
-+  //               CHECK: "vhlo.collective_permute_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME{LITERAL}:   source_target_pairs = #vhlo.tensor_v1<dense<[[0, 1], [1, 2], [2, 3]]> : tensor<3x2xi64>>
-   //          CHECK-SAME: }> : (!vhlo.tensor_v1<16x8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x8x!vhlo.f32_v1>
-@@ -1033,8 +1108,9 @@
- }
- 
- // CHECK-LABEL: "op_compare"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_compare(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
--  //      CHECK: "vhlo.compare_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.compare_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   compare_type = #vhlo<comparison_type_v1 TOTALORDER>,
-   // CHECK-SAME:   comparison_direction = #vhlo<comparison_direction_v1 EQ>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-@@ -1046,15 +1122,17 @@
- }
- 
- // CHECK-LABEL: "op_complex"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_complex(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<complex<f32>> {
--  // CHECK: "vhlo.complex_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>
-+  // CHECK: "vhlo.complex_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>
-   %0 = "stablehlo.complex"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<complex<f32>>
-   func.return %0 : tensor<complex<f32>>
- }
- 
- // CHECK-LABEL: "op_composite"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_composite(%arg0: tensor<f32>) -> tensor<f32> {
--  //               CHECK: "vhlo.composite_v1"(%arg0) <{
-+  //               CHECK: "vhlo.composite_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   composite_attributes = #vhlo.dict_v1<{#vhlo.string_v1<"my_int"> = #vhlo.integer_v1<1 : i64>, #vhlo.string_v1<"my_string"> = #vhlo.string_v1<"foo">}>
-   //          CHECK-SAME:   decomposition = #vhlo.string_v1<"composite_target">
-   //          CHECK-SAME:   name = #vhlo.string_v1<"stablehlo.composite_target">
-@@ -1073,8 +1151,9 @@
- }
- 
- // CHECK-LABEL: "op_concatenate"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_concatenate(%arg0: tensor<8xf32>, %arg1: tensor<8xf32>) -> tensor<16xf32> {
--  //      CHECK: "vhlo.concatenate_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.concatenate_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   dimension = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<8x!vhlo.f32_v1>, !vhlo.tensor_v1<8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x!vhlo.f32_v1>
-   %0 = "stablehlo.concatenate"(%arg0, %arg1) {
-@@ -1084,6 +1163,7 @@
- }
- 
- // CHECK-LABEL: "op_constant"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_constant(%arg0: tensor<f32>) -> tensor<f32> {
-   //      CHECK: "vhlo.constant_v1"() <{
-   // CHECK-SAME:   value = #vhlo.tensor_v1<dense<0.000000e+00> : tensor<f32>>
-@@ -1095,15 +1175,17 @@
- }
- 
- // CHECK-LABEL: "op_convert"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_convert(%arg0: tensor<i32>) -> tensor<f32> {
--  // CHECK: "vhlo.convert_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.convert_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.convert"(%arg0) : (tensor<i32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_convolution"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_convolution(%arg0: tensor<1x8x8x207xf32>, %arg1: tensor<3x3x207x16xf32>) -> tensor<1x7x7x16xf32> {
--  //      CHECK: "vhlo.convolution_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.convolution_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   batch_group_count = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   feature_group_count = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   input_batch_dimension = #vhlo.integer_v1<0 : i64>,
-@@ -1137,8 +1219,9 @@
- }
- 
- // CHECK-LABEL: "op_cosine"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_cosine(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.cosine_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.cosine_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.cosine"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
-@@ -1151,8 +1234,9 @@
- }
- 
- // CHECK-LABEL: "op_cross_replica_sum"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_cross_replica_sum(%arg0: tensor<f32>) -> tensor<f32> {
--  //               CHECK: "vhlo.cross-replica-sum_v1"(%arg0) <{
-+  //               CHECK: "vhlo.cross-replica-sum_v1"(%[[ARG0]]) <{
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0], [1]]> : tensor<2x1xi64>>
-   //          CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.cross-replica-sum"(%arg0) {
-@@ -1162,8 +1246,9 @@
- }
- 
- // CHECK-LABEL: "op_custom_call"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_custom_call(%arg0: tensor<f32>) -> tensor<f32> {
--  //      CHECK: "vhlo.custom_call_v1"(%arg0) <{
-+  //      CHECK: "vhlo.custom_call_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   api_version = #vhlo<api_version_v1 API_VERSION_STATUS_RETURNING>,
-   // CHECK-SAME:   backend_config = #vhlo.string_v1<"\08\03\1A\02">,
-   // CHECK-SAME:   call_target_name = #vhlo.string_v1<"foo">,
-@@ -1194,15 +1279,17 @@
- }
- 
- // CHECK-LABEL: "op_divide"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_divide(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.divide_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.divide_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.divide"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_dot_general"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_dot_general(%arg0: tensor<8x8x16xf32>, %arg1: tensor<8x16x8xf32>) -> tensor<8x8x8xf32> {
--  //      CHECK: "vhlo.dot_general_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.dot_general_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   lhs_batching_dimensions = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>,
-   // CHECK-SAME:   lhs_contracting_dimensions = #vhlo.tensor_v1<dense<2> : tensor<1xi64>>,
-   // CHECK-SAME:   precision_config = #vhlo.array_v1<[#vhlo<precision_v1 HIGHEST>, #vhlo<precision_v1 HIGHEST>]>,
-@@ -1222,8 +1309,9 @@
- }
- 
- // CHECK-LABEL: "op_dot"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_dot(%arg0: tensor<8x16xf32>, %arg1: tensor<16x8xf32>) -> tensor<8x8xf32> {
--  //      CHECK: "vhlo.dot_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.dot_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   precision_config = #vhlo.array_v1<[#vhlo<precision_v1 HIGHEST>, #vhlo<precision_v1 HIGHEST>]>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<8x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<8x8x!vhlo.f32_v1>
-   %0 = "stablehlo.dot"(%arg0, %arg1) {
-@@ -1233,8 +1321,9 @@
- }
- 
- // CHECK-LABEL: "op_dynamic_broadcast_in_dim"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_dynamic_broadcast_in_dim(%arg0: tensor<?x?xf32>, %arg1: tensor<2xindex>) -> tensor<?x?xf32> {
--  //      CHECK: "vhlo.dynamic_broadcast_in_dim_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.dynamic_broadcast_in_dim_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   broadcast_dimensions = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-   // CHECK-SAME:   known_expanding_dimensions = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>,
-   // CHECK-SAME:   known_nonexpanding_dimensions = #vhlo.tensor_v1<dense<1> : tensor<1xi64>>
-@@ -1248,8 +1337,9 @@
- }
- 
- // CHECK-LABEL: "op_dynamic_conv"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_dynamic_conv(%arg0: tensor<1x8x8x207xf32>, %arg1: tensor<3x3x207x16xf32>, %arg2: tensor<4xi32>) -> tensor<1x?x?x16xf32> {
--  //      CHECK: "vhlo.dynamic_conv_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.dynamic_conv_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   batch_group_count = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   feature_group_count = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   input_batch_dimension = #vhlo.integer_v1<0 : i64>,
-@@ -1283,8 +1373,9 @@
- }
- 
- // CHECK-LABEL: "op_dynamic_gather"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_dynamic_gather(%arg0 : tensor<2x4x9xf32>, %arg1 : tensor<1x5x2xi32>, %arg2 : tensor<3xi32>) -> tensor<1x5x8xf32> {
--  //      CHECK: "vhlo.dynamic_gather_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.dynamic_gather_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   collapsed_slice_dims = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-   // CHECK-SAME:   index_vector_dim = #vhlo.integer_v1<2 : i64>,
-   // CHECK-SAME:   indices_are_sorted = #vhlo.bool_v1<true>,
-@@ -1304,8 +1395,9 @@
- }
- 
- // CHECK-LABEL: "op_dynamic_iota"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_dynamic_iota(%arg0: tensor<1xindex>) -> tensor<?xf32> {
--  //      CHECK: "vhlo.dynamic_iota_v1"(%arg0) <{
-+  //      CHECK: "vhlo.dynamic_iota_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   iota_dimension = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<1x!vhlo.index_v1>) -> !vhlo.tensor_v1<?x!vhlo.f32_v1>
-   %0 = "stablehlo.dynamic_iota"(%arg0) {
-@@ -1315,22 +1407,25 @@
- }
- 
- // CHECK-LABEL: "op_dynamic_pad"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}}, %[[ARG3:.*]]: {{.*}}, %[[ARG4:.*]]: {{.*}})
- func.func @op_dynamic_pad(%arg0: tensor<?xf32>, %arg1: tensor<f32>, %arg2: tensor<1xindex>, %arg3: tensor<1xindex>, %arg4: tensor<1xindex>) -> tensor<?xf32> {
--  // CHECK: "vhlo.dynamic_pad_v1"(%arg0, %arg1, %arg2, %arg3, %arg4) : (!vhlo.tensor_v1<?x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>) -> !vhlo.tensor_v1<?x!vhlo.f32_v1>
-+  // CHECK: "vhlo.dynamic_pad_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]], %[[ARG3]], %[[ARG4]]) : (!vhlo.tensor_v1<?x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>) -> !vhlo.tensor_v1<?x!vhlo.f32_v1>
-   %0 = "stablehlo.dynamic_pad"(%arg0, %arg1, %arg2, %arg3, %arg4) : (tensor<?xf32>, tensor<f32>, tensor<1xindex>, tensor<1xindex>, tensor<1xindex>) -> tensor<?xf32>
-   func.return %0 : tensor<?xf32>
- }
- 
- // CHECK-LABEL: "op_dynamic_reshape"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_dynamic_reshape(%arg0: tensor<16xf32>, %arg1: tensor<2xindex>) -> tensor<?x?xf32> {
--  // CHECK: "vhlo.dynamic_reshape_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<2x!vhlo.index_v1>) -> !vhlo.tensor_v1<?x?x!vhlo.f32_v1>
-+  // CHECK: "vhlo.dynamic_reshape_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<2x!vhlo.index_v1>) -> !vhlo.tensor_v1<?x?x!vhlo.f32_v1>
-   %0 = "stablehlo.dynamic_reshape"(%arg0, %arg1) : (tensor<16xf32>, tensor<2xindex>) -> tensor<?x?xf32>
-   func.return %0 : tensor<?x?xf32>
- }
- 
- // CHECK-LABEL: "op_dynamic_slice"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_dynamic_slice(%arg0: tensor<16xf32>, %arg1: tensor<i64>) -> tensor<4xf32> {
--  //      CHECK: "vhlo.dynamic_slice_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.dynamic_slice_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   slice_sizes = #vhlo.tensor_v1<dense<4> : tensor<1xi64>>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.i64_v1>) -> !vhlo.tensor_v1<4x!vhlo.f32_v1>
-   %0 = "stablehlo.dynamic_slice"(%arg0, %arg1) {
-@@ -1340,15 +1435,17 @@
- }
- 
- // CHECK-LABEL: "op_dynamic_update_slice"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_dynamic_update_slice(%arg0: tensor<16xf32>, %arg1: tensor<4xf32>, %arg2: tensor<i64>) -> tensor<16xf32> {
--  // CHECK: "vhlo.dynamic_update_slice_v1"(%arg0, %arg1, %arg2) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<4x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.i64_v1>) -> !vhlo.tensor_v1<16x!vhlo.f32_v1>
-+  // CHECK: "vhlo.dynamic_update_slice_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<4x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.i64_v1>) -> !vhlo.tensor_v1<16x!vhlo.f32_v1>
-   %0 = "stablehlo.dynamic_update_slice"(%arg0, %arg1, %arg2) : (tensor<16xf32>, tensor<4xf32>, tensor<i64>) -> tensor<16xf32>
-   func.return %0 : tensor<16xf32>
- }
- 
- // CHECK-LABEL: "op_einsum"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_einsum(%arg0: tensor<8x16xf32>, %arg1: tensor<16x8xf32>) -> tensor<8x8xf32> {
--  //      CHECK: "vhlo.einsum_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.einsum_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   einsum_config = #vhlo.string_v1<"ab,bc->ac">
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<8x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<8x8x!vhlo.f32_v1>
-   %0 = "stablehlo.einsum"(%arg0, %arg1) {
-@@ -1358,22 +1455,25 @@
- }
- 
- // CHECK-LABEL: "op_exponential_minus_one"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_exponential_minus_one(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.exponential_minus_one_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.exponential_minus_one_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.exponential_minus_one"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_exponential"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_exponential(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.exponential_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.exponential_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.exponential"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_fft"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_fft(%arg0: tensor<16xcomplex<f32>>) -> tensor<16xcomplex<f32>> {
--  //      CHECK: "vhlo.fft_v1"(%arg0) <{
-+  //      CHECK: "vhlo.fft_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   fft_length = #vhlo.tensor_v1<dense<16> : tensor<1xi64>>,
-   // CHECK-SAME:   fft_type = #vhlo<fft_type_v1 FFT>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<16x!vhlo.complex_v1<!vhlo.f32_v1>>) -> !vhlo.tensor_v1<16x!vhlo.complex_v1<!vhlo.f32_v1>>
-@@ -1385,8 +1485,9 @@
- }
- 
- // CHECK-LABEL: "op_floor"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_floor(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.floor_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.floor_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.floor"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
-@@ -1399,16 +1500,17 @@
-   // CHECK-SAME:   sym_name = #vhlo.string_v1<"op_func">,
-   // CHECK-SAME:   sym_visibility = #vhlo.string_v1<"private">
-   // CHECK-SAME: }> ({
--  // CHECK-NEXT:   ^[[BB:bb.*]](%arg0: !vhlo.tensor_v1<!vhlo.f32_v1>):
--  // CHECK-NEXT:     "vhlo.return_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-+  // CHECK-NEXT:   ^[[BB:bb.*]](%[[ARG0:.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
-+  // CHECK-NEXT:     "vhlo.return_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-   // CHECK-NEXT: }) : () -> ()
- 
-   func.return %arg0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_gather"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_gather(%arg0 : tensor<2x4x9xf32>, %arg1 : tensor<1x5x2xi32>) -> tensor<1x5x1xf32> {
--  //      CHECK: "vhlo.gather_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.gather_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   collapsed_slice_dims = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-   // CHECK-SAME:   index_vector_dim = #vhlo.integer_v1<2 : i64>,
-   // CHECK-SAME:   indices_are_sorted = #vhlo.bool_v1<true>,
-@@ -1430,8 +1532,9 @@
- }
- 
- // CHECK-LABEL: "op_get_dimension_size"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_get_dimension_size(%arg0: tensor<?xf32>) -> tensor<i32> {
--  //      CHECK: "vhlo.get_dimension_size_v1"(%arg0) <{
-+  //      CHECK: "vhlo.get_dimension_size_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   dimension = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<?x!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-   %0 = "stablehlo.get_dimension_size"(%arg0) {
-@@ -1441,8 +1544,9 @@
- }
- 
- // CHECK-LABEL: "op_get_tuple_element"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_get_tuple_element(%arg0: tuple<tensor<f32>, tensor<i32>>) -> tensor<f32> {
--  //      CHECK: "vhlo.get_tuple_element_v1"(%arg0) <{
-+  //      CHECK: "vhlo.get_tuple_element_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   index = #vhlo.integer_v1<0 : i32>
-   // CHECK-SAME: }> : (!vhlo.tuple_v1<!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.get_tuple_element"(%arg0) {
-@@ -1452,11 +1556,12 @@
- }
- 
- // CHECK-LABEL: "op_if"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_if(%arg0: tensor<i1>, %arg1: tensor<f32>, %arg2: tensor<f32>) -> tensor<f32> {
--  //      CHECK: "vhlo.if_v1"(%arg0) ({
--  // CHECK-NEXT:   "vhlo.return_v1"(%arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-+  //      CHECK: "vhlo.if_v1"(%[[ARG0]]) ({
-+  // CHECK-NEXT:   "vhlo.return_v1"(%[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-   // CHECK-NEXT: }, {
--  // CHECK-NEXT:   "vhlo.return_v1"(%arg2) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-+  // CHECK-NEXT:   "vhlo.return_v1"(%[[ARG2]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-   // CHECK-NEXT: }) : (!vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.if"(%arg0) ({
-     "stablehlo.return"(%arg1) : (tensor<f32>) -> ()
-@@ -1467,15 +1572,17 @@
- }
- 
- // CHECK-LABEL: "op_imag"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_imag(%arg0: tensor<complex<f32>>) -> tensor<f32> {
--  // CHECK: "vhlo.imag_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.imag_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.imag"(%arg0) : (tensor<complex<f32>>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_infeed"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_infeed(%arg0: !stablehlo.token) -> (tensor<f32>, !stablehlo.token) {
--  //               CHECK: "vhlo.infeed_v1"(%arg0) <{
-+  //               CHECK: "vhlo.infeed_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   infeed_config = #vhlo.string_v1<"foo">,
-   // CHECK-SAME{LITERAL}:   layout = #vhlo.array_v1<[#vhlo.array_v1<[]>]>
-   //          CHECK-SAME: }> : (!vhlo.token_v1) -> (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.token_v1)
-@@ -1498,36 +1605,41 @@
- }
- 
- // CHECK-LABEL: "op_is_finite"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_is_finite(%arg0: tensor<f32>) -> tensor<i1> {
--  // CHECK: "vhlo.is_finite_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-+  // CHECK: "vhlo.is_finite_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-   %0 = "stablehlo.is_finite"(%arg0) : (tensor<f32>) -> tensor<i1>
-   func.return %0 : tensor<i1>
- }
- 
- // CHECK-LABEL: "op_log"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_log(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.log_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.log_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.log"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_log_plus_one"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_log_plus_one(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.log_plus_one_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.log_plus_one_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.log_plus_one"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_logistic"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_logistic(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.logistic_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.logistic_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.logistic"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_map"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_map(%arg0: tensor<16xf32>) -> tensor<16xf32> {
--  //      CHECK: "vhlo.map_v1"(%arg0) <{
-+  //      CHECK: "vhlo.map_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   dimensions = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>
-   // CHECK-SAME: }> ({
-   // CHECK-NEXT:   ^[[BB:bb.*]](%[[ARG1:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
-@@ -1545,57 +1657,65 @@
- }
- 
- // CHECK-LABEL: "op_maximum"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_maximum(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.maximum_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.maximum_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.maximum"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_minimum"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_minimum(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.minimum_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.minimum_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.minimum"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_multiply"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_multiply(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.multiply_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.multiply_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.multiply"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_negate"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_negate(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.negate_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.negate_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.negate"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_not"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_not(%arg0: tensor<i1>) -> tensor<i1> {
--  // CHECK: "vhlo.not_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-+  // CHECK: "vhlo.not_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-   %0 = "stablehlo.not"(%arg0) : (tensor<i1>) -> tensor<i1>
-   func.return %0 : tensor<i1>
- }
- 
- // CHECK-LABEL: "op_optimization_barrier"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_optimization_barrier(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.optimization_barrier_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.optimization_barrier_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.optimization_barrier"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_or"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_or(%arg0: tensor<i1>, %arg1: tensor<i1>) -> tensor<i1> {
--  // CHECK: "vhlo.or_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-+  // CHECK: "vhlo.or_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-   %0 = "stablehlo.or"(%arg0, %arg1) : (tensor<i1>, tensor<i1>) -> tensor<i1>
-   func.return %0 : tensor<i1>
- }
- 
- // CHECK-LABEL: "op_outfeed"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_outfeed(%arg0: tensor<f32>, %arg1: !stablehlo.token) -> !stablehlo.token {
--  //      CHECK: "vhlo.outfeed_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.outfeed_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   outfeed_config = #vhlo.string_v1<"foo">
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.token_v1) -> !vhlo.token_v1
-   %0 = "stablehlo.outfeed"(%arg0, %arg1) {
-@@ -1605,8 +1725,9 @@
- }
- 
- // CHECK-LABEL: "op_pad"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_pad(%arg0: tensor<8xf32>, %arg1: tensor<f32>) -> tensor<16xf32> {
--  //      CHECK: "vhlo.pad_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.pad_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   edge_padding_high = #vhlo.tensor_v1<dense<4> : tensor<1xi64>>,
-   // CHECK-SAME:   edge_padding_low = #vhlo.tensor_v1<dense<4> : tensor<1xi64>>,
-   // CHECK-SAME:   interior_padding = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>
-@@ -1620,36 +1741,41 @@
- }
- 
- // CHECK-LABEL: "op_popcnt"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_popcnt(%arg0: tensor<i32>) -> tensor<i32> {
--  // CHECK: "vhlo.popcnt_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-+  // CHECK: "vhlo.popcnt_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-   %0 = "stablehlo.popcnt"(%arg0) : (tensor<i32>) -> tensor<i32>
-   func.return %0 : tensor<i32>
- }
- 
- // CHECK-LABEL: "op_power"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_power(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.power_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.power_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.power"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_real_dynamic_slice"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}}, %[[ARG3:.*]]: {{.*}})
- func.func @op_real_dynamic_slice(%arg0: tensor<?xf32>, %arg1: tensor<1xindex>, %arg2: tensor<1xindex>, %arg3: tensor<1xindex>) -> tensor<?xf32> {
--  // CHECK: "vhlo.real_dynamic_slice_v1"(%arg0, %arg1, %arg2, %arg3) : (!vhlo.tensor_v1<?x!vhlo.f32_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>) -> !vhlo.tensor_v1<?x!vhlo.f32_v1>
-+  // CHECK: "vhlo.real_dynamic_slice_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]], %[[ARG3]]) : (!vhlo.tensor_v1<?x!vhlo.f32_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>) -> !vhlo.tensor_v1<?x!vhlo.f32_v1>
-   %0 = "stablehlo.real_dynamic_slice"(%arg0, %arg1, %arg2, %arg3) : (tensor<?xf32>, tensor<1xindex>, tensor<1xindex>, tensor<1xindex>) -> tensor<?xf32>
-   func.return %0 : tensor<?xf32>
- }
- 
- // CHECK-LABEL: "op_real"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_real(%arg0: tensor<complex<f32>>) -> tensor<f32> {
--  // CHECK: "vhlo.real_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.real_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.real"(%arg0) : (tensor<complex<f32>>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_recv"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_recv(%arg0: !stablehlo.token) -> (tensor<f32>, !stablehlo.token) {
--  //      CHECK: "vhlo.recv_v1"(%arg0) <{
-+  //      CHECK: "vhlo.recv_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME:   channel_type = #vhlo.integer_v1<3 : i64>,
-   // CHECK-SAME:   is_host_transfer = #vhlo.bool_v1<true>
-@@ -1662,8 +1788,9 @@
- }
- 
- // CHECK-LABEL: "op_reduce"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_reduce(%arg0: tensor<16xf32>, %arg1: tensor<f32>) -> tensor<f32> {
--  //  CHECK: "vhlo.reduce_v1"(%arg0, %arg1)
-+  //  CHECK: "vhlo.reduce_v1"(%[[ARG0]], %[[ARG1]])
-   //  CHECK:   ^[[BB:bb.*]](%[[ARG1:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>, %[[ARG2:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
-   //  CHECK:     "vhlo.return_v1"(%[[VAL1:.*]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-   //  CHECK: }) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-@@ -1678,8 +1805,9 @@
- }
- 
- // CHECK-LABEL: "op_reduce_precision"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_reduce_precision(%arg0: tensor<f32>) -> tensor<f32> {
--  //      CHECK: "vhlo.reduce_precision_v1"(%arg0) <{
-+  //      CHECK: "vhlo.reduce_precision_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   exponent_bits = #vhlo.integer_v1<8 : i32>
-   // CHECK-SAME:   mantissa_bits = #vhlo.integer_v1<10 : i32>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-@@ -1690,10 +1818,10 @@
-   func.return %0 : tensor<f32>
- }
- 
--// CHECK_lABEL: "op_reduce_with_promotable_types"
-+// CHECK-LABEL: "op_reduce_with_promotable_types"
- func.func @op_reduce_with_promotable_types(%arg0: tensor<4x4xf32>, %arg1 : tensor<f32>)
-     -> (tensor<4xf64>) {
--  //  CHECK: "vhlo.reduce_v1"(%arg0, %arg1)
-+  //  CHECK: "vhlo.reduce_v1"(%[[ARG0:.*]], %[[ARG1:.*]])
-   //  CHECK:   ^[[BB:bb.*]](%[[ARG1:arg.*]]: !vhlo.tensor_v1<!vhlo.f64_v1>, %[[ARG2:arg.*]]: !vhlo.tensor_v1<!vhlo.f64_v1>):
-   //  CHECK:     "vhlo.return_v1"(%[[VAL1:.*]]) : (!vhlo.tensor_v1<!vhlo.f64_v1>) -> ()
-   //  CHECK: }) : (!vhlo.tensor_v1<4x4x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<4x!vhlo.f64_v1>
-@@ -1708,8 +1836,9 @@
- }
- 
- // CHECK-LABEL: "op_reduce_scatter"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_reduce_scatter(%arg0: tensor<16xf32>) -> tensor<16xf32> {
--  //               CHECK: "vhlo.reduce_scatter_v1"(%arg0) <{
-+  //               CHECK: "vhlo.reduce_scatter_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0], [1]]> : tensor<2x1xi64>>,
-   //          CHECK-SAME:   scatter_dimension = #vhlo.integer_v1<0 : i64>
-@@ -1734,7 +1863,7 @@
- 
- // CHECK_lABEL: "op_reduce_scatter_with_promotable_types"
- func.func @op_reduce_scatter_with_promotable_types(%data: tensor<4x16xf32>) -> tensor<4x4xf64> {
--  //  CHECK: "vhlo.reduce_scatter_v1"(%arg0)
-+  //  CHECK: "vhlo.reduce_scatter_v1"(%[[ARG0:.*]])
-   //  CHECK:   ^[[BB:bb.*]](%[[ARG1:arg.*]]: !vhlo.tensor_v1<!vhlo.f64_v1>, %[[ARG2:arg.*]]: !vhlo.tensor_v1<!vhlo.f64_v1>):
-   //  CHECK:     "vhlo.return_v1"(%[[VAL1:.*]]) : (!vhlo.tensor_v1<!vhlo.f64_v1>) -> ()
-   //  CHECK: }) : (!vhlo.tensor_v1<4x16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<4x4x!vhlo.f64_v1>
-@@ -1751,8 +1880,9 @@
- 
- 
- // CHECK-LABEL: "op_reduce_window"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_reduce_window(%arg0: tensor<2x17x31x7xf32>, %arg1: tensor<f32>) -> tensor<2x9x16x7xf32> {
--  //               CHECK: "vhlo.reduce_window_v1"(%arg0, %arg1) <{
-+  //               CHECK: "vhlo.reduce_window_v1"(%[[ARG0]], %[[ARG1]]) <{
-   //          CHECK-SAME:   base_dilations = #vhlo.tensor_v1<dense<[1, 2, 2, 1]> : tensor<4xi64>>,
-   // CHECK-SAME{LITERAL}:   padding = #vhlo.tensor_v1<dense<[[0, 0], [2, 0], [0, 2], [0, 0]]> : tensor<4x2xi64>>,
-   //          CHECK-SAME:   window_dilations = #vhlo.tensor_v1<dense<[1, 2, 2, 1]> : tensor<4xi64>>,
-@@ -1781,7 +1911,7 @@
- func.func @op_reduce_window_with_promotable_types(%arg0: tensor<4x2xf32>,
-     %arg1: tensor<4x2xf32>, %init0: tensor<f32>, %init1: tensor<f32>) ->
-     (tensor<2x2xf64>, tensor<2x2xf32>) {
--  //  CHECK: "vhlo.reduce_window_v1"(%arg0, %arg1, %arg2, %arg3)
-+  //  CHECK: "vhlo.reduce_window_v1"(%[[ARG0:.*]], %[[ARG1:.*]], %[[ARG2:.*]], %[[ARG3:.*]])
-   //  CHECK:   ^[[BB:bb.*]](%[[ARG1:arg.*]]: !vhlo.tensor_v1<!vhlo.f64_v1>, %[[ARG2:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>, %[[ARG3:arg.*]]: !vhlo.tensor_v1<!vhlo.f64_v1>, %[[ARG4:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
-   //  CHECK:     "vhlo.return_v1"(%[[VAL1:.*]], %[[VAL2:.*]]) : (!vhlo.tensor_v1<!vhlo.f64_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-   //  CHECK: }) : (!vhlo.tensor_v1<4x2x!vhlo.f32_v1>, !vhlo.tensor_v1<4x2x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> (!vhlo.tensor_v1<2x2x!vhlo.f64_v1>, !vhlo.tensor_v1<2x2x!vhlo.f32_v1>)
-@@ -1801,8 +1931,9 @@
- }
- 
- // CHECK-LABEL: "op_remainder"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_remainder(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.remainder_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.remainder_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.remainder"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
-@@ -1822,16 +1953,18 @@
- }
- 
- // CHECK-LABEL: "op_reshape"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_reshape(%arg0: tensor<16xf32>) -> tensor<4x4xf32> {
--  // CHECK: "vhlo.reshape_v1"(%arg0) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<4x4x!vhlo.f32_v1>
-+  // CHECK: "vhlo.reshape_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<4x4x!vhlo.f32_v1>
-   %0 = "stablehlo.reshape"(%arg0) : (tensor<16xf32>) -> tensor<4x4xf32>
-   func.return %0 : tensor<4x4xf32>
- }
- 
- // CHECK-LABEL: "op_return"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_return(%arg0: tensor<i32>, %arg1: tensor<f32>) -> tensor<f32> {
--  //      CHECK: "vhlo.case_v1"(%arg0) ({
--  // CHECK-NEXT:   "vhlo.return_v1"(%arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-+  //      CHECK: "vhlo.case_v1"(%[[ARG0]]) ({
-+  // CHECK-NEXT:   "vhlo.return_v1"(%[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-   // CHECK-NEXT: }) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.case"(%arg0) ({
-     "stablehlo.return"(%arg1) : (tensor<f32>) -> ()
-@@ -1840,8 +1973,9 @@
- }
- 
- // CHECK-LABEL: "op_reverse"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_reverse(%arg0: tensor<16xf32>) -> tensor<16xf32> {
--  //      CHECK: "vhlo.reverse_v1"(%arg0) <{
-+  //      CHECK: "vhlo.reverse_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   dimensions = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x!vhlo.f32_v1>
-   %0 = "stablehlo.reverse"(%arg0) {
-@@ -1851,8 +1985,9 @@
- }
- 
- // CHECK-LABEL: "op_rng_bit_generator"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_rng_bit_generator(%arg0: tensor<f32>) -> (tensor<f32>, tensor<f32>) {
--  //      CHECK: "vhlo.rng_bit_generator_v1"(%arg0) <{
-+  //      CHECK: "vhlo.rng_bit_generator_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   rng_algorithm = #vhlo<rng_algorithm_v1 PHILOX>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>)
-   %0:2 = "stablehlo.rng_bit_generator"(%arg0) {
-@@ -1862,8 +1997,9 @@
- }
- 
- // CHECK-LABEL: "op_rng"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_rng(%arg0: tensor<f32>, %arg1: tensor<f32>, %arg2: tensor<0xindex>) -> tensor<f32> {
--  //      CHECK: "vhlo.rng_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.rng_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   rng_distribution = #vhlo<rng_distribution_v1 NORMAL>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<0x!vhlo.index_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.rng"(%arg0, %arg1, %arg2) {
-@@ -1873,29 +2009,33 @@
- }
- 
- // CHECK-LABEL: "op_round_nearest_afz"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_round_nearest_afz(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.round_nearest_afz_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.round_nearest_afz_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.round_nearest_afz"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_round_nearest_even"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_round_nearest_even(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.round_nearest_even_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.round_nearest_even_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.round_nearest_even"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_rsqrt"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_rsqrt(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.rsqrt_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.rsqrt_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.rsqrt"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_scatter"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_scatter(%arg0: tensor<200x100x300xf32>, %arg1: tensor<10x2xi32>, %arg2: tensor<10x300xf32>) -> tensor<200x100x300xf32> {
--  //      CHECK: "vhlo.scatter_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.scatter_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   index_vector_dim = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   indices_are_sorted = #vhlo.bool_v1<true>,
-   // CHECK-SAME:   inserted_window_dims = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-@@ -1928,7 +2068,7 @@
- func.func @op_scatter_with_promotable_types(%input_tensor: tensor<200x100x300xf32>,
-     %scatter_indices: tensor<10x2xi32>, %updates: tensor<10x300xf32>) ->
-       tensor<200x100x300xf64> {
--  //  CHECK: "vhlo.scatter_v1"(%arg0, %arg1, %arg2)
-+  //  CHECK: "vhlo.scatter_v1"(%[[ARG0:.*]], %[[ARG1:.*]], %[[ARG2:.*]])
-   //  CHECK:   ^[[BB:bb.*]](%[[ARG1:arg.*]]: !vhlo.tensor_v1<!vhlo.f64_v1>, %[[ARG2:arg.*]]: !vhlo.tensor_v1<!vhlo.f64_v1>):
-   //  CHECK:     "vhlo.return_v1"(%[[VAL1:.*]]) : (!vhlo.tensor_v1<!vhlo.f64_v1>) -> ()
-   //  CHECK: }) : (!vhlo.tensor_v1<200x100x300x!vhlo.f32_v1>, !vhlo.tensor_v1<10x2x!vhlo.i32_v1>, !vhlo.tensor_v1<10x300x!vhlo.f32_v1>) -> !vhlo.tensor_v1<200x100x300x!vhlo.f64_v1>
-@@ -1951,8 +2091,9 @@
- }
- 
- // CHECK-LABEL: "op_select_and_scatter"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_select_and_scatter(%arg0: tensor<10x24x24x64xf32>, %arg1: tensor<12x13x13x66xf32>, %arg2: tensor<f32>) -> tensor<10x24x24x64xf32> {
--  //      CHECK: "vhlo.select_and_scatter_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.select_and_scatter_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   padding = #vhlo.tensor_v1<dense<1> : tensor<4x2xi64>>,
-   // CHECK-SAME:   window_dimensions = #vhlo.tensor_v1<dense<[1, 2, 2, 1]> : tensor<4xi64>>,
-   // CHECK-SAME:   window_strides = #vhlo.tensor_v1<dense<[1, 2, 2, 1]> : tensor<4xi64>>
-@@ -1982,8 +2123,9 @@
- }
- 
- // CHECK-LABEL: "op_select_and_scatter_with_promotable_types"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_select_and_scatter_with_promotable_types(%arg0: tensor<10x24x24x64xf32>, %arg1: tensor<12x13x13x66xf32>, %arg2: tensor<f32>) -> tensor<10x24x24x64xf64> {
--  // CHECK: "vhlo.select_and_scatter_v1"(%arg0, %arg1, %arg2)
-+  // CHECK: "vhlo.select_and_scatter_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]])
-   // CHECK:   ^[[BB:bb.*]](%[[ARG1:arg.*]]: !vhlo.tensor_v1<!vhlo.f64_v1>, %[[ARG2:arg.*]]: !vhlo.tensor_v1<!vhlo.f64_v1>):
-   // CHECK:     %[[VAL:.*]] = "vhlo.add_v1"(%[[ARG1]], %[[ARG2]]) : (!vhlo.tensor_v1<!vhlo.f64_v1>, !vhlo.tensor_v1<!vhlo.f64_v1>) -> !vhlo.tensor_v1<!vhlo.f64_v1>
-   // CHECK:     "vhlo.return_v1"(%[[VAL]]) : (!vhlo.tensor_v1<!vhlo.f64_v1>) -> ()
-@@ -2005,15 +2147,17 @@
- }
- 
- // CHECK-LABEL: "op_select"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_select(%arg0: tensor<i1>, %arg1: tensor<f32>, %arg2: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.select_v1"(%arg0, %arg1, %arg2) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.select_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.select"(%arg0, %arg1, %arg2) : (tensor<i1>, tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_send"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_send(%arg0: tensor<f32>, %arg1: !stablehlo.token) -> !stablehlo.token {
--  //      CHECK: "vhlo.send_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.send_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME:   channel_type = #vhlo.integer_v1<2 : i64>,
-   // CHECK-SAME:   is_host_transfer = #vhlo.bool_v1<true>
-@@ -2026,8 +2170,9 @@
- }
- 
- // CHECK-LABEL: "op_set_dimension_size"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_set_dimension_size(%arg0: tensor<?xf32>, %arg1: tensor<i32>) -> tensor<16xf32> {
--  //      CHECK: "vhlo.set_dimension_size_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.set_dimension_size_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   dimension = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<?x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<16x!vhlo.f32_v1>
-   %0 = "stablehlo.set_dimension_size"(%arg0, %arg1) {
-@@ -2037,43 +2182,49 @@
- }
- 
- // CHECK-LABEL: "op_shift_left"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_shift_left(%arg0: tensor<i32>, %arg1: tensor<i32>) -> tensor<i32> {
--  // CHECK: "vhlo.shift_left_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.i32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-+  // CHECK: "vhlo.shift_left_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-   %0 = "stablehlo.shift_left"(%arg0, %arg1) : (tensor<i32>, tensor<i32>) -> tensor<i32>
-   func.return %0 : tensor<i32>
- }
- 
- // CHECK-LABEL: "op_shift_right_arithmetic"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_shift_right_arithmetic(%arg0: tensor<i32>, %arg1: tensor<i32>) -> tensor<i32> {
--  // CHECK: "vhlo.shift_right_arithmetic_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.i32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-+  // CHECK: "vhlo.shift_right_arithmetic_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-   %0 = "stablehlo.shift_right_arithmetic"(%arg0, %arg1) : (tensor<i32>, tensor<i32>) -> tensor<i32>
-   func.return %0 : tensor<i32>
- }
- 
- // CHECK-LABEL: "op_shift_right_logical"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_shift_right_logical(%arg0: tensor<i32>, %arg1: tensor<i32>) -> tensor<i32> {
--  // CHECK: "vhlo.shift_right_logical_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.i32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-+  // CHECK: "vhlo.shift_right_logical_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-   %0 = "stablehlo.shift_right_logical"(%arg0, %arg1) : (tensor<i32>, tensor<i32>) -> tensor<i32>
-   func.return %0 : tensor<i32>
- }
- 
- // CHECK-LABEL: "op_sign"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_sign(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.sign_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.sign_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.sign"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_sine"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_sine(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.sine_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.sine_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.sine"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_slice"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_slice(%arg0: tensor<16xf32>) -> tensor<4xf32> {
--  //      CHECK: "vhlo.slice_v1"(%arg0) <{
-+  //      CHECK: "vhlo.slice_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   limit_indices = #vhlo.tensor_v1<dense<4> : tensor<1xi64>>,
-   // CHECK-SAME:   start_indices = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>,
-   // CHECK-SAME:   strides = #vhlo.tensor_v1<dense<1> : tensor<1xi64>>
-@@ -2087,8 +2238,9 @@
- }
- 
- // CHECK-LABEL: "op_sort"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_sort(%arg0: tensor<16xf32>) -> tensor<16xf32> {
--  //      CHECK: "vhlo.sort_v1"(%arg0) <{
-+  //      CHECK: "vhlo.sort_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   dimension = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME:   is_stable = #vhlo.bool_v1<true>
-   // CHECK-SAME: }> ({
-@@ -2108,29 +2260,33 @@
- }
- 
- // CHECK-LABEL: "op_sqrt"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_sqrt(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.sqrt_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.sqrt_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.sqrt"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_subtract"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_subtract(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.subtract_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.subtract_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.subtract"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_tanh"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_tanh(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.tanh_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.tanh_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.tanh"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_torch_index_select"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_torch_index_select(%arg0: tensor<5x1x5xf32>, %arg1: tensor<2xi32>) ->  tensor<2x1x5xf32> {
--  //      CHECK: "vhlo.torch_index_select_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.torch_index_select_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   batch_dims = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME:   dim = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<5x1x5x!vhlo.f32_v1>, !vhlo.tensor_v1<2x!vhlo.i32_v1>) -> !vhlo.tensor_v1<2x1x5x!vhlo.f32_v1>
-@@ -2142,8 +2298,9 @@
- }
- 
- // CHECK-LABEL: "op_trace"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_trace(%arg0: tensor<f32>) {
--  //      CHECK: "vhlo.trace_v1"(%arg0) <{
-+  //      CHECK: "vhlo.trace_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   tag = #vhlo.string_v1<"foo">
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-   "stablehlo.trace"(%arg0) {
-@@ -2153,8 +2310,9 @@
- }
- 
- // CHECK-LABEL: "op_transpose"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_transpose(%arg0: tensor<16x8xf32>) ->  tensor<8x16xf32> {
--  //      CHECK: "vhlo.transpose_v1"(%arg0) <{
-+  //      CHECK: "vhlo.transpose_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   permutation = #vhlo.tensor_v1<dense<[1, 0]> : tensor<2xi64>>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<16x8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<8x16x!vhlo.f32_v1>
-   %0 = "stablehlo.transpose"(%arg0) {
-@@ -2164,8 +2322,9 @@
- }
- 
- // CHECK-LABEL: "op_triangular_solve"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_triangular_solve(%arg0: tensor<16x16xf32>, %arg1: tensor<16x16xf32>) ->  tensor<16x16xf32> {
--  //      CHECK: "vhlo.triangular_solve_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.triangular_solve_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   left_side = #vhlo.bool_v1<true>,
-   // CHECK-SAME:   lower = #vhlo.bool_v1<true>,
-   // CHECK-SAME:   transpose_a = #vhlo<transpose_v1 NO_TRANSPOSE>,
-@@ -2181,15 +2340,17 @@
- }
- 
- // CHECK-LABEL: "op_tuple"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_tuple(%arg0: tensor<f32>) -> tuple<tensor<f32>> {
--  // CHECK: "vhlo.tuple_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tuple_v1<!vhlo.tensor_v1<!vhlo.f32_v1>>
-+  // CHECK: "vhlo.tuple_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tuple_v1<!vhlo.tensor_v1<!vhlo.f32_v1>>
-   %0 = "stablehlo.tuple"(%arg0) : (tensor<f32>) -> tuple<tensor<f32>>
-   func.return %0 : tuple<tensor<f32>>
- }
- 
- // CHECK-LABEL: "op_unary_einsum"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_unary_einsum(%arg0: tensor<8x16xf32>) -> tensor<8xf32> {
--  //      CHECK: "vhlo.unary_einsum_v1"(%arg0) <{
-+  //      CHECK: "vhlo.unary_einsum_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   einsum_config = #vhlo.string_v1<"ab->a">
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<8x16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<8x!vhlo.f32_v1>
-   %0 = "stablehlo.unary_einsum"(%arg0) {
-@@ -2199,22 +2360,25 @@
- }
- 
- // CHECK-LABEL: "op_uniform_dequantize"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_uniform_dequantize(%arg0: tensor<!quant.uniform<i8:f32, 34.0:16>>) -> tensor<f32> {
--  // CHECK: "vhlo.uniform_dequantize_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.uniform_dequantize_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.uniform_dequantize"(%arg0) : (tensor<!quant.uniform<i8:f32, 34.0:16>>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_uniform_quantize"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_uniform_quantize(%arg0: tensor<f32>) -> tensor<!quant.uniform<i8:f32, 34.0:16>> {
--  // CHECK: "vhlo.uniform_quantize_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>
-+  // CHECK: "vhlo.uniform_quantize_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>
-   %0 = "stablehlo.uniform_quantize"(%arg0) : (tensor<f32>) -> tensor<!quant.uniform<i8:f32, 34.0:16>>
-   func.return %0 : tensor<!quant.uniform<i8:f32, 34.0:16>>
- }
- 
- // CHECK-LABEL: "op_while"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_while(%arg0: tensor<i1>) -> tensor<i1> {
--  //      CHECK: "vhlo.while_v1"(%arg0) ({
-+  //      CHECK: "vhlo.while_v1"(%[[ARG0]]) ({
-   // CHECK-NEXT:   ^[[BB:bb.*]](%[[ARG1:arg.*]]: !vhlo.tensor_v1<!vhlo.bool_v1>):
-   // CHECK-NEXT:     "vhlo.return_v1"(%[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>) -> ()
-   // CHECK-NEXT:   }, {
-@@ -2232,8 +2396,9 @@
- }
- 
- // CHECK-LABEL: "op_xor"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_xor(%arg0: tensor<i1>, %arg1: tensor<i1>) -> tensor<i1> {
--  // CHECK: "vhlo.xor_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-+  // CHECK: "vhlo.xor_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-   %0 = "stablehlo.xor"(%arg0, %arg1) : (tensor<i1>, tensor<i1>) -> tensor<i1>
-   func.return %0 : tensor<i1>
- }
-@@ -2241,197 +2406,225 @@
- // ============ TYPES ============
- 
- // CHECK-LABEL: "type_i1"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_i1(%arg0: tensor<i1>, %arg1: tensor<i1>) -> tensor<i1> {
--  // CHECK: "vhlo.and_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-+  // CHECK: "vhlo.and_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-   %0 = "stablehlo.and"(%arg0, %arg1) : (tensor<i1>, tensor<i1>) -> tensor<i1>
-   func.return %0 : tensor<i1>
- }
- 
- // CHECK-LABEL: "type_i4"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_i4(%arg0: tensor<i4>, %arg1: tensor<i4>) -> tensor<i4> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.i4_v1>, !vhlo.tensor_v1<!vhlo.i4_v1>) -> !vhlo.tensor_v1<!vhlo.i4_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i4_v1>, !vhlo.tensor_v1<!vhlo.i4_v1>) -> !vhlo.tensor_v1<!vhlo.i4_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<i4>, tensor<i4>) -> tensor<i4>
-   func.return %0 : tensor<i4>
- }
- 
- // CHECK-LABEL: "type_i8"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_i8(%arg0: tensor<i8>, %arg1: tensor<i8>) -> tensor<i8> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.i8_v1>, !vhlo.tensor_v1<!vhlo.i8_v1>) -> !vhlo.tensor_v1<!vhlo.i8_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i8_v1>, !vhlo.tensor_v1<!vhlo.i8_v1>) -> !vhlo.tensor_v1<!vhlo.i8_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<i8>, tensor<i8>) -> tensor<i8>
-   func.return %0 : tensor<i8>
- }
- 
- // CHECK-LABEL: "type_i16"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_i16(%arg0: tensor<i16>, %arg1: tensor<i16>) -> tensor<i16> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.i16_v1>, !vhlo.tensor_v1<!vhlo.i16_v1>) -> !vhlo.tensor_v1<!vhlo.i16_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i16_v1>, !vhlo.tensor_v1<!vhlo.i16_v1>) -> !vhlo.tensor_v1<!vhlo.i16_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<i16>, tensor<i16>) -> tensor<i16>
-   func.return %0 : tensor<i16>
- }
- 
- // CHECK-LABEL: "type_i32"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_i32(%arg0: tensor<i32>, %arg1: tensor<i32>) -> tensor<i32> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.i32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<i32>, tensor<i32>) -> tensor<i32>
-   func.return %0 : tensor<i32>
- }
- 
- // CHECK-LABEL: "type_i64"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_i64(%arg0: tensor<i64>, %arg1: tensor<i64>) -> tensor<i64> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.i64_v1>, !vhlo.tensor_v1<!vhlo.i64_v1>) -> !vhlo.tensor_v1<!vhlo.i64_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i64_v1>, !vhlo.tensor_v1<!vhlo.i64_v1>) -> !vhlo.tensor_v1<!vhlo.i64_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<i64>, tensor<i64>) -> tensor<i64>
-   func.return %0 : tensor<i64>
- }
- 
- // CHECK-LABEL: "type_ui4"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_ui4(%arg0: tensor<ui4>, %arg1: tensor<ui4>) -> tensor<ui4> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.ui4_v1>, !vhlo.tensor_v1<!vhlo.ui4_v1>) -> !vhlo.tensor_v1<!vhlo.ui4_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.ui4_v1>, !vhlo.tensor_v1<!vhlo.ui4_v1>) -> !vhlo.tensor_v1<!vhlo.ui4_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<ui4>, tensor<ui4>) -> tensor<ui4>
-   func.return %0 : tensor<ui4>
- }
- 
- // CHECK-LABEL: "type_ui8"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_ui8(%arg0: tensor<ui8>, %arg1: tensor<ui8>) -> tensor<ui8> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.ui8_v1>, !vhlo.tensor_v1<!vhlo.ui8_v1>) -> !vhlo.tensor_v1<!vhlo.ui8_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.ui8_v1>, !vhlo.tensor_v1<!vhlo.ui8_v1>) -> !vhlo.tensor_v1<!vhlo.ui8_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<ui8>, tensor<ui8>) -> tensor<ui8>
-   func.return %0 : tensor<ui8>
- }
- 
- // CHECK-LABEL: "type_ui16"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_ui16(%arg0: tensor<ui16>, %arg1: tensor<ui16>) -> tensor<ui16> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.ui16_v1>, !vhlo.tensor_v1<!vhlo.ui16_v1>) -> !vhlo.tensor_v1<!vhlo.ui16_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.ui16_v1>, !vhlo.tensor_v1<!vhlo.ui16_v1>) -> !vhlo.tensor_v1<!vhlo.ui16_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<ui16>, tensor<ui16>) -> tensor<ui16>
-   func.return %0 : tensor<ui16>
- }
- 
- // CHECK-LABEL: "type_ui32"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_ui32(%arg0: tensor<ui32>, %arg1: tensor<ui32>) -> tensor<ui32> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.ui32_v1>, !vhlo.tensor_v1<!vhlo.ui32_v1>) -> !vhlo.tensor_v1<!vhlo.ui32_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.ui32_v1>, !vhlo.tensor_v1<!vhlo.ui32_v1>) -> !vhlo.tensor_v1<!vhlo.ui32_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<ui32>, tensor<ui32>) -> tensor<ui32>
-   func.return %0 : tensor<ui32>
- }
- 
- // CHECK-LABEL: "type_ui64"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_ui64(%arg0: tensor<ui64>, %arg1: tensor<ui64>) -> tensor<ui64> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.ui64_v1>, !vhlo.tensor_v1<!vhlo.ui64_v1>) -> !vhlo.tensor_v1<!vhlo.ui64_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.ui64_v1>, !vhlo.tensor_v1<!vhlo.ui64_v1>) -> !vhlo.tensor_v1<!vhlo.ui64_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<ui64>, tensor<ui64>) -> tensor<ui64>
-   func.return %0 : tensor<ui64>
- }
- 
- // CHECK-LABEL: "type_f8E4M3FN"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_f8E4M3FN(%arg0: tensor<f8E4M3FN>, %arg1: tensor<f8E4M3FN>) -> tensor<f8E4M3FN> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f8E4M3FN_v1>, !vhlo.tensor_v1<!vhlo.f8E4M3FN_v1>) -> !vhlo.tensor_v1<!vhlo.f8E4M3FN_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f8E4M3FN_v1>, !vhlo.tensor_v1<!vhlo.f8E4M3FN_v1>) -> !vhlo.tensor_v1<!vhlo.f8E4M3FN_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f8E4M3FN>, tensor<f8E4M3FN>) -> tensor<f8E4M3FN>
-   func.return %0 : tensor<f8E4M3FN>
- }
- 
- // CHECK-LABEL: "type_f8E5M2"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_f8E5M2(%arg0: tensor<f8E5M2>, %arg1: tensor<f8E5M2>) -> tensor<f8E5M2> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f8E5M2_v1>, !vhlo.tensor_v1<!vhlo.f8E5M2_v1>) -> !vhlo.tensor_v1<!vhlo.f8E5M2_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f8E5M2_v1>, !vhlo.tensor_v1<!vhlo.f8E5M2_v1>) -> !vhlo.tensor_v1<!vhlo.f8E5M2_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f8E5M2>, tensor<f8E5M2>) -> tensor<f8E5M2>
-   func.return %0 : tensor<f8E5M2>
- }
- 
- // CHECK-LABEL: "type_f8E4M3FNUZ"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_f8E4M3FNUZ(%arg0: tensor<f8E4M3FNUZ>, %arg1: tensor<f8E4M3FNUZ>) -> tensor<f8E4M3FNUZ> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f8E4M3FNUZ_v1>, !vhlo.tensor_v1<!vhlo.f8E4M3FNUZ_v1>) -> !vhlo.tensor_v1<!vhlo.f8E4M3FNUZ_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f8E4M3FNUZ_v1>, !vhlo.tensor_v1<!vhlo.f8E4M3FNUZ_v1>) -> !vhlo.tensor_v1<!vhlo.f8E4M3FNUZ_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f8E4M3FNUZ>, tensor<f8E4M3FNUZ>) -> tensor<f8E4M3FNUZ>
-   func.return %0 : tensor<f8E4M3FNUZ>
- }
- 
- // CHECK-LABEL: "type_f8E4M3B11FNUZ"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_f8E4M3B11FNUZ(%arg0: tensor<f8E4M3B11FNUZ>, %arg1: tensor<f8E4M3B11FNUZ>) -> tensor<f8E4M3B11FNUZ> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f8E4M3B11FNUZ_v1>, !vhlo.tensor_v1<!vhlo.f8E4M3B11FNUZ_v1>) -> !vhlo.tensor_v1<!vhlo.f8E4M3B11FNUZ_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f8E4M3B11FNUZ_v1>, !vhlo.tensor_v1<!vhlo.f8E4M3B11FNUZ_v1>) -> !vhlo.tensor_v1<!vhlo.f8E4M3B11FNUZ_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f8E4M3B11FNUZ>, tensor<f8E4M3B11FNUZ>) -> tensor<f8E4M3B11FNUZ>
-   func.return %0 : tensor<f8E4M3B11FNUZ>
- }
- 
- // CHECK-LABEL: "type_f8E5M2FNUZ"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_f8E5M2FNUZ(%arg0: tensor<f8E5M2FNUZ>, %arg1: tensor<f8E5M2FNUZ>) -> tensor<f8E5M2FNUZ> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f8E5M2FNUZ_v1>, !vhlo.tensor_v1<!vhlo.f8E5M2FNUZ_v1>) -> !vhlo.tensor_v1<!vhlo.f8E5M2FNUZ_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f8E5M2FNUZ_v1>, !vhlo.tensor_v1<!vhlo.f8E5M2FNUZ_v1>) -> !vhlo.tensor_v1<!vhlo.f8E5M2FNUZ_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f8E5M2FNUZ>, tensor<f8E5M2FNUZ>) -> tensor<f8E5M2FNUZ>
-   func.return %0 : tensor<f8E5M2FNUZ>
- }
- 
- // CHECK-LABEL: "type_bf16"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_bf16(%arg0: tensor<bf16>, %arg1: tensor<bf16>) -> tensor<bf16> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.bf16_v1>, !vhlo.tensor_v1<!vhlo.bf16_v1>) -> !vhlo.tensor_v1<!vhlo.bf16_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.bf16_v1>, !vhlo.tensor_v1<!vhlo.bf16_v1>) -> !vhlo.tensor_v1<!vhlo.bf16_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<bf16>, tensor<bf16>) -> tensor<bf16>
-   func.return %0 : tensor<bf16>
- }
- 
- // CHECK-LABEL: "type_f16"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_f16(%arg0: tensor<f16>, %arg1: tensor<f16>) -> tensor<f16> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f16_v1>, !vhlo.tensor_v1<!vhlo.f16_v1>) -> !vhlo.tensor_v1<!vhlo.f16_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f16_v1>, !vhlo.tensor_v1<!vhlo.f16_v1>) -> !vhlo.tensor_v1<!vhlo.f16_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f16>, tensor<f16>) -> tensor<f16>
-   func.return %0 : tensor<f16>
- }
- 
- // CHECK-LABEL: "type_f32"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_f32(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "type_f64"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_f64(%arg0: tensor<f64>, %arg1: tensor<f64>) -> tensor<f64> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f64_v1>, !vhlo.tensor_v1<!vhlo.f64_v1>) -> !vhlo.tensor_v1<!vhlo.f64_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f64_v1>, !vhlo.tensor_v1<!vhlo.f64_v1>) -> !vhlo.tensor_v1<!vhlo.f64_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f64>, tensor<f64>) -> tensor<f64>
-   func.return %0 : tensor<f64>
- }
- 
- // CHECK-LABEL: "type_complex_f32"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_complex_f32(%arg0: tensor<complex<f32>>, %arg1: tensor<complex<f32>>) -> tensor<complex<f32>> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>, !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>) -> !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>, !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>) -> !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<complex<f32>>, tensor<complex<f32>>) -> tensor<complex<f32>>
-   func.return %0 : tensor<complex<f32>>
- }
- 
- // CHECK-LABEL: "type_complex_f64"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_complex_f64(%arg0: tensor<complex<f64>>, %arg1: tensor<complex<f64>>) -> tensor<complex<f64>> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f64_v1>>, !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f64_v1>>) -> !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f64_v1>>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f64_v1>>, !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f64_v1>>) -> !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f64_v1>>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<complex<f64>>, tensor<complex<f64>>) -> tensor<complex<f64>>
-   func.return %0 : tensor<complex<f64>>
- }
- 
- // CHECK-LABEL: "type_dynamism_ranked"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @type_dynamism_ranked(%arg0: tensor<?xf32>) -> tensor<?xf32> {
--  // CHECK: "vhlo.abs_v1"(%arg0) : (!vhlo.tensor_v1<?x!vhlo.f32_v1>) -> !vhlo.tensor_v1<?x!vhlo.f32_v1>
-+  // CHECK: "vhlo.abs_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<?x!vhlo.f32_v1>) -> !vhlo.tensor_v1<?x!vhlo.f32_v1>
-   %0 = "stablehlo.abs"(%arg0) : (tensor<?xf32>) -> tensor<?xf32>
-   func.return %0 : tensor<?xf32>
- }
- 
- // CHECK-LABEL: "type_per_tensor_quantization"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_per_tensor_quantization(%arg0: tensor<!quant.uniform<i8:f32, 34.0:16>>, %arg1: tensor<!quant.uniform<i8:f32, 34.0:16>>) -> tensor<!quant.uniform<i8:f32, 34.0:16>> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>, !vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>) -> !vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>, !vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>) -> !vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<!quant.uniform<i8:f32, 34.0:16>>, tensor<!quant.uniform<i8:f32, 34.0:16>>) -> tensor<!quant.uniform<i8:f32, 34.0:16>>
-   func.return %0 : tensor<!quant.uniform<i8:f32, 34.0:16>>
- }
- 
- // CHECK-LABEL: "type_per_axis_quantization"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @type_per_axis_quantization(%arg0: tensor<2x!quant.uniform<i8:f32:0, {34.0:16, 34.0:16}>>) -> tensor<2x!quant.uniform<i8:f32:0, {34.0:16, 34.0:16}>> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg0) : (!vhlo.tensor_v1<2x!vhlo.quant_per_axis_v1<!vhlo.i8_v1:!vhlo.f32_v1, 0, 3.400000e+01, 3.400000e+01, 16, 16, -128:127, 1>>, !vhlo.tensor_v1<2x!vhlo.quant_per_axis_v1<!vhlo.i8_v1:!vhlo.f32_v1, 0, 3.400000e+01, 3.400000e+01, 16, 16, -128:127, 1>>) -> !vhlo.tensor_v1<2x!vhlo.quant_per_axis_v1<!vhlo.i8_v1:!vhlo.f32_v1, 0, 3.400000e+01, 3.400000e+01, 16, 16, -128:127, 1>>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG0]]) : (!vhlo.tensor_v1<2x!vhlo.quant_per_axis_v1<!vhlo.i8_v1:!vhlo.f32_v1, 0, 3.400000e+01, 3.400000e+01, 16, 16, -128:127, 1>>, !vhlo.tensor_v1<2x!vhlo.quant_per_axis_v1<!vhlo.i8_v1:!vhlo.f32_v1, 0, 3.400000e+01, 3.400000e+01, 16, 16, -128:127, 1>>) -> !vhlo.tensor_v1<2x!vhlo.quant_per_axis_v1<!vhlo.i8_v1:!vhlo.f32_v1, 0, 3.400000e+01, 3.400000e+01, 16, 16, -128:127, 1>>
-   %0 = stablehlo.add %arg0, %arg0 : tensor<2x!quant.uniform<i8:f32:0, {34.0:16, 34.0:16}>>
-   func.return %0 : tensor<2x!quant.uniform<i8:f32:0, {34.0:16, 34.0:16}>>
- }
- 
- //       CHECK: function_type = #vhlo.type_v1<!vhlo.func_v1<(!vhlo.token_v1) -> !vhlo.token_v1>>
- // CHECK-LABEL: "type_token_callee"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @type_token_callee(%arg0: !stablehlo.token) -> !stablehlo.token {
--  // CHECK: "vhlo.return_v1"(%arg0) : (!vhlo.token_v1) -> ()
-+  // CHECK: "vhlo.return_v1"(%[[ARG0]]) : (!vhlo.token_v1) -> ()
-   return %arg0 : !stablehlo.token
- }
- 
- //       CHECK: function_type = #vhlo.type_v1<!vhlo.func_v1<(!vhlo.token_v1) -> !vhlo.token_v1>>
- // CHECK-LABEL: "type_token_caller"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @type_token_caller(%arg0: !stablehlo.token) -> !stablehlo.token {
--  // CHECK:      "vhlo.call_v1"(%arg0) <{callee = #vhlo.string_v1<"type_token_callee">}
-+  // CHECK:      "vhlo.call_v1"(%[[ARG0]]) <{callee = #vhlo.string_v1<"type_token_callee">}
-   // CHECK-SAME: (!vhlo.token_v1) -> !vhlo.token_v1
-   %0 = func.call @type_token_callee(%arg0) : (!stablehlo.token) -> !stablehlo.token
-   return %0 : !stablehlo.token
- }
- 
- // CHECK-LABEL: "type_tuple"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @type_tuple(%arg0: tuple<tensor<f32>>) -> tuple<!stablehlo.token> {
-   %0 = "stablehlo.custom_call"(%arg0) {
-     call_target_name = "foo"
-diff --ruN a/stablehlo/stablehlo/tests/vhlo/stablehlo_legalize_to_vhlo.0_9_0.mlir b/stablehlo/stablehlo/tests/vhlo/stablehlo_legalize_to_vhlo.0_9_0.mlir
---- stablehlo/stablehlo/tests/vhlo/stablehlo_legalize_to_vhlo.0_9_0.mlir
-+++ stablehlo/stablehlo/tests/vhlo/stablehlo_legalize_to_vhlo.0_9_0.mlir
-@@ -13,6 +13,7 @@
- // ============ ATTRIBUTES ============
- 
- // CHECK-LABEL: "attr_comparison_direction_eq"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_direction_eq(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     // CHECK: comparison_direction = #vhlo<comparison_direction_v1 EQ>
-@@ -22,6 +23,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_direction_ne"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_direction_ne(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     // CHECK: comparison_direction = #vhlo<comparison_direction_v1 NE>
-@@ -31,6 +33,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_direction_ge"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_direction_ge(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     // CHECK: comparison_direction = #vhlo<comparison_direction_v1 GE>
-@@ -40,6 +43,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_direction_gt"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_direction_gt(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     // CHECK: comparison_direction = #vhlo<comparison_direction_v1 GT>
-@@ -49,6 +53,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_direction_le"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_direction_le(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     // CHECK: comparison_direction = #vhlo<comparison_direction_v1 LE>
-@@ -58,6 +63,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_direction_lt"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_direction_lt(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     // CHECK: comparison_direction = #vhlo<comparison_direction_v1 LT>
-@@ -67,6 +73,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_type_notype"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_type_notype(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     comparison_direction = #stablehlo<comparison_direction EQ>
-@@ -76,6 +83,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_type_float"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_type_float(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     comparison_direction = #stablehlo<comparison_direction EQ>,
-@@ -86,6 +94,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_type_totalorder"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_type_totalorder(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     comparison_direction = #stablehlo<comparison_direction EQ>,
-@@ -96,6 +105,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_type_signed"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_type_signed(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     comparison_direction = #stablehlo<comparison_direction EQ>,
-@@ -106,6 +116,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_type_unsigned"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_type_unsigned(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     comparison_direction = #stablehlo<comparison_direction EQ>,
-@@ -118,6 +129,7 @@
- // ConvDimensionNumbers aka #stablehlo.conv is covered below.
- 
- // CHECK-LABEL: "attr_custom_call_api_version_unspecified"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_custom_call_api_version_unspecified(%arg0: tensor<f32>) -> tensor<f32> {
-   %0 = "stablehlo.custom_call"(%arg0) {
-     call_target_name = "foo",
-@@ -128,6 +140,7 @@
- }
- 
- // CHECK-LABEL: "attr_custom_call_api_version_original"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_custom_call_api_version_original(%arg0: tensor<f32>) -> tensor<f32> {
-   %0 = "stablehlo.custom_call"(%arg0) {
-     call_target_name = "foo",
-@@ -138,6 +151,7 @@
- }
- 
- // CHECK-LABEL: "attr_custom_call_api_version_status_returning"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_custom_call_api_version_status_returning(%arg0: tensor<f32>) -> tensor<f32> {
-   %0 = "stablehlo.custom_call"(%arg0) {
-     call_target_name = "foo",
-@@ -148,6 +162,7 @@
- }
- 
- // CHECK-LABEL: "attr_custom_call_api_version_status_returning_unified"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_custom_call_api_version_status_returning_unified(%arg0: tensor<f32>) -> tensor<f32> {
-   %0 = "stablehlo.custom_call"(%arg0) {
-     call_target_name = "foo",
-@@ -160,6 +175,7 @@
- // DotDimensionNumbers aka #stablehlo.dot is covered below.
- 
- // CHECK-LABEL: "attr_fft_type_fft"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_fft_type_fft(%arg0: tensor<16xcomplex<f32>>) -> tensor<16xcomplex<f32>> {
-   %0 = "stablehlo.fft"(%arg0) {
-     // CHECK: fft_type = #vhlo<fft_type_v1 FFT>
-@@ -170,6 +186,7 @@
- }
- 
- // CHECK-LABEL: "attr_fft_type_ifft"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_fft_type_ifft(%arg0: tensor<16xcomplex<f32>>) -> tensor<16xcomplex<f32>> {
-   %0 = "stablehlo.fft"(%arg0) {
-     // CHECK: fft_type = #vhlo<fft_type_v1 IFFT>
-@@ -180,6 +197,7 @@
- }
- 
- // CHECK-LABEL: "attr_fft_type_rfft"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_fft_type_rfft(%arg0: tensor<16xf32>) -> tensor<9xcomplex<f32>> {
-   %0 = "stablehlo.fft"(%arg0) {
-     // CHECK: fft_type = #vhlo<fft_type_v1 RFFT>
-@@ -190,6 +208,7 @@
- }
- 
- // CHECK-LABEL: "attr_fft_type_irfft"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_fft_type_irfft(%arg0: tensor<9xcomplex<f32>>) -> tensor<16xf32> {
-   %0 = "stablehlo.fft"(%arg0) {
-     // CHECK: fft_type = #vhlo<fft_type_v1 IRFFT>
-@@ -202,6 +221,7 @@
- // GatherDimensionNumbers aka #stablehlo.gather is covered below.
- 
- // CHECK-LABEL: "attr_precision_config_default"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_precision_config_default(%arg0: tensor<8x16xf32>, %arg1: tensor<16x8xf32>) -> tensor<8x8xf32> {
-   %0 = "stablehlo.dot"(%arg0, %arg1) {
-     // CHECK: precision_config = #vhlo.array_v1<[#vhlo<precision_v1 DEFAULT>, #vhlo<precision_v1 DEFAULT>]>
-@@ -210,6 +230,7 @@
- }
- 
- // CHECK-LABEL: "attr_precision_config_high"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_precision_config_high(%arg0: tensor<8x16xf32>, %arg1: tensor<16x8xf32>) -> tensor<8x8xf32> {
-   %0 = "stablehlo.dot"(%arg0, %arg1) {
-     // CHECK: precision_config = #vhlo.array_v1<[#vhlo<precision_v1 HIGH>, #vhlo<precision_v1 HIGH>]>
-@@ -219,6 +240,7 @@
- }
- 
- // CHECK-LABEL: "attr_precision_config_highest"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_precision_config_highest(%arg0: tensor<8x16xf32>, %arg1: tensor<16x8xf32>) -> tensor<8x8xf32> {
-   %0 = "stablehlo.dot"(%arg0, %arg1) {
-     // CHECK: precision_config = #vhlo.array_v1<[#vhlo<precision_v1 HIGHEST>, #vhlo<precision_v1 HIGHEST>]>
-@@ -228,6 +250,7 @@
- }
- 
- // CHECK-LABEL: "attr_rng_algorithm_default"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_rng_algorithm_default(%arg0: tensor<f32>) -> (tensor<f32>, tensor<f32>) {
-   %0:2 = "stablehlo.rng_bit_generator"(%arg0) {
-     // CHECK: rng_algorithm = #vhlo<rng_algorithm_v1 DEFAULT>
-@@ -237,6 +260,7 @@
- }
- 
- // CHECK-LABEL: "attr_rng_algorithm_three_fry"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_rng_algorithm_three_fry(%arg0: tensor<f32>) -> (tensor<f32>, tensor<f32>) {
-   %0:2 = "stablehlo.rng_bit_generator"(%arg0) {
-     // CHECK: rng_algorithm = #vhlo<rng_algorithm_v1 THREE_FRY>
-@@ -246,6 +270,7 @@
- }
- 
- // CHECK-LABEL: "attr_rng_algorithm_philox"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_rng_algorithm_philox(%arg0: tensor<f32>) -> (tensor<f32>, tensor<f32>) {
-   %0:2 = "stablehlo.rng_bit_generator"(%arg0) {
-     // CHECK: rng_algorithm = #vhlo<rng_algorithm_v1 PHILOX>
-@@ -255,6 +280,7 @@
- }
- 
- // CHECK-LABEL: "attr_rng_distribution_uniform"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @attr_rng_distribution_uniform(%arg0: tensor<f32>, %arg1: tensor<f32>, %arg2: tensor<0xindex>) -> tensor<f32> {
-   %0 = "stablehlo.rng"(%arg0, %arg1, %arg2) {
-     // CHECK: rng_distribution = #vhlo<rng_distribution_v1 UNIFORM>
-@@ -264,6 +290,7 @@
- }
- 
- // CHECK-LABEL: "attr_rng_distribution_normal"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @attr_rng_distribution_normal(%arg0: tensor<f32>, %arg1: tensor<f32>, %arg2: tensor<0xindex>) -> tensor<f32> {
-   %0 = "stablehlo.rng"(%arg0, %arg1, %arg2) {
-     // CHECK: rng_distribution = #vhlo<rng_distribution_v1 NORMAL>
-@@ -275,6 +302,7 @@
- // ScatterDimensionNumbers aka #stablehlo.scatter is covered below.
- 
- // CHECK-LABEL: "attr_transpose_no_transpose"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_transpose_no_transpose(%arg0: tensor<16x16xf32>, %arg1: tensor<16x16xf32>) ->  tensor<16x16xf32> {
-   %0 = "stablehlo.triangular_solve"(%arg0, %arg1) {
-     left_side = true,
-@@ -287,6 +315,7 @@
- }
- 
- // CHECK-LABEL: "attr_transpose_transpose"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_transpose_transpose(%arg0: tensor<16x16xf32>, %arg1: tensor<16x16xf32>) ->  tensor<16x16xf32> {
-   %0 = "stablehlo.triangular_solve"(%arg0, %arg1) {
-     left_side = true,
-@@ -299,6 +328,7 @@
- }
- 
- // CHECK-LABEL: "attr_transpose_adjoint"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_transpose_adjoint(%arg0: tensor<16x16xf32>, %arg1: tensor<16x16xf32>) ->  tensor<16x16xf32> {
-   %0 = "stablehlo.triangular_solve"(%arg0, %arg1) {
-     left_side = true,
-@@ -313,10 +343,9 @@
- // TypeExtensionsAttr aka #stablehlo.type_extensions is covered below.
- 
- // CHECK-LABEL: "attr_type_extensions_bounds"
--func.func @attr_type_extensions_bounds(
--    %arg0: tensor<?x?xf32, #stablehlo.type_extensions<bounds = [16, ?]>>)
--    -> tensor<?x?xf32, #stablehlo.type_extensions<bounds = [16, ?]>> {
--  // CHECK: "vhlo.return_v1"(%arg0) : (!vhlo.tensor_v1<?x?x!vhlo.f32_v1, #vhlo.type_extensions_v1<bounds = [16, ?]>>) -> ()
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @attr_type_extensions_bounds(%arg0: tensor<?x?xf32, #stablehlo.type_extensions<bounds = [16, ?]>>) -> tensor<?x?xf32, #stablehlo.type_extensions<bounds = [16, ?]>> {
-+  // CHECK: "vhlo.return_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<?x?x!vhlo.f32_v1, #vhlo.type_extensions_v1<bounds = [16, ?]>>) -> ()
-   func.return %arg0 : tensor<?x?xf32, #stablehlo.type_extensions<bounds = [16, ?]>>
- }
- 
-@@ -324,8 +353,9 @@
- // ============ DEFAULTS ============
- 
- // CHECK-LABEL: "default_all_gather"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_all_gather(%arg0: tensor<16x8xf32>) -> tensor<16x16xf32> {
--  //               CHECK: "vhlo.all_gather_v1"(%arg0) <{
-+  //               CHECK: "vhlo.all_gather_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   all_gather_dim = #vhlo.integer_v1<1 : i64>
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0], [1]]> : tensor<2x1xi64>>,
-@@ -339,8 +369,9 @@
- }
- 
- // CHECK-LABEL: "default_all_reduce"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_all_reduce(%arg0: tensor<f32>) -> tensor<f32> {
--  //               CHECK: "vhlo.all_reduce_v1"(%arg0)
-+  //               CHECK: "vhlo.all_reduce_v1"(%[[ARG0]])
-   //          CHECK-SAME: <{
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0], [1]]> : tensor<2x1xi64>>,
-@@ -362,8 +393,9 @@
- }
- 
- // CHECK-LABEL: "default_all_to_all"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_all_to_all(%arg0: tensor<4x16xf32>) -> tensor<16x4xf32> {
--  //               CHECK: "vhlo.all_to_all_v1"(%arg0) <{
-+  //               CHECK: "vhlo.all_to_all_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   //          CHECK-SAME:   concat_dimension = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0, 1, 2, 3]]> : tensor<1x4xi64>>,
-@@ -380,8 +412,9 @@
- }
- 
- // CHECK-LABEL: "default_cholesky"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_cholesky(%arg0: tensor<1x16x16xf32>) -> tensor<1x16x16xf32> {
--  //      CHECK: "vhlo.cholesky_v1"(%arg0) <{
-+  //      CHECK: "vhlo.cholesky_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   lower = #vhlo.bool_v1<false>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<1x16x16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<1x16x16x!vhlo.f32_v1>
-   %0 = "stablehlo.cholesky"(%arg0) : (tensor<1x16x16xf32>) -> tensor<1x16x16xf32>
-@@ -389,8 +422,9 @@
- }
- 
- // CHECK-LABEL: "default_collective_permute"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_collective_permute(%arg0: tensor<16x8xf32>) -> tensor<16x8xf32> {
--  //               CHECK: "vhlo.collective_permute_v1"(%arg0) <{
-+  //               CHECK: "vhlo.collective_permute_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME{LITERAL}:   source_target_pairs = #vhlo.tensor_v1<dense<[[0, 1], [1, 2], [2, 3]]> : tensor<3x2xi64>>
-   //          CHECK-SAME: }> : (!vhlo.tensor_v1<16x8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x8x!vhlo.f32_v1>
-@@ -401,8 +435,9 @@
- }
- 
- // CHECK-LABEL: "default_compare"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @default_compare(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
--  //      CHECK: "vhlo.compare_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.compare_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   compare_type = #vhlo<comparison_type_v1 NOTYPE>,
-   // CHECK-SAME:   comparison_direction = #vhlo<comparison_direction_v1 EQ>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-@@ -413,8 +448,9 @@
- }
- 
- // CHECK-LABEL: "default_convolution"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @default_convolution(%arg0: tensor<1x8x8x207xf32>, %arg1: tensor<3x3x207x16xf32>) -> tensor<1x6x6x16xf32> {
--  //      CHECK: "vhlo.convolution_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.convolution_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   batch_group_count = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   feature_group_count = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   input_batch_dimension = #vhlo.integer_v1<0 : i64>,
-@@ -442,8 +478,9 @@
- }
- 
- // CHECK-LABEL: "default_custom_call"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_custom_call(%arg0: tensor<f32>) -> tensor<f32> {
--  //      CHECK: "vhlo.custom_call_v1"(%arg0) <{
-+  //      CHECK: "vhlo.custom_call_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   api_version = #vhlo<api_version_v1 API_VERSION_ORIGINAL>,
-   // CHECK-SAME:   backend_config = #vhlo.string_v1<"">,
-   // CHECK-SAME:   call_target_name = #vhlo.string_v1<"foo">,
-@@ -460,8 +497,9 @@
- }
- 
- // CHECK-LABEL: "default_dot_general"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @default_dot_general(%arg0: tensor<8x8x16xf32>, %arg1: tensor<8x16x8xf32>) -> tensor<8x8x8xf32> {
--  //      CHECK: "vhlo.dot_general_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.dot_general_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   lhs_batching_dimensions = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>,
-   // CHECK-SAME:   lhs_contracting_dimensions = #vhlo.tensor_v1<dense<2> : tensor<1xi64>>,
-   // CHECK-SAME:   precision_config = #vhlo.array_v1<[#vhlo<precision_v1 DEFAULT>, #vhlo<precision_v1 DEFAULT>]>,
-@@ -480,8 +518,9 @@
- }
- 
- // CHECK-LABEL: "default_dot"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @default_dot(%arg0: tensor<8x16xf32>, %arg1: tensor<16x8xf32>) -> tensor<8x8xf32> {
--  //      CHECK: "vhlo.dot_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.dot_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   precision_config = #vhlo.array_v1<[#vhlo<precision_v1 DEFAULT>, #vhlo<precision_v1 DEFAULT>]>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<8x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<8x8x!vhlo.f32_v1>
-   %0 = "stablehlo.dot"(%arg0, %arg1) : (tensor<8x16xf32>, tensor<16x8xf32>) -> tensor<8x8xf32>
-@@ -489,8 +528,9 @@
- }
- 
- // CHECK-LABEL: "default_dynamic_broadcast_in_dim"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @default_dynamic_broadcast_in_dim(%arg0: tensor<?x?xf32>, %arg1: tensor<2xindex>) -> tensor<?x?xf32> {
--  //      CHECK: "vhlo.dynamic_broadcast_in_dim_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.dynamic_broadcast_in_dim_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   broadcast_dimensions = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-   // CHECK-SAME:   known_expanding_dimensions = #vhlo.tensor_v1<dense<> : tensor<0xi64>>,
-   // CHECK-SAME:   known_nonexpanding_dimensions = #vhlo.tensor_v1<dense<> : tensor<0xi64>>
-@@ -502,8 +542,9 @@
- }
- 
- // CHECK-LABEL: "default_dynamic_conv"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @default_dynamic_conv(%arg0: tensor<1x8x8x207xf32>, %arg1: tensor<3x3x207x16xf32>, %arg2: tensor<4xi32>) -> tensor<1x?x?x16xf32> {
--  //      CHECK: "vhlo.dynamic_conv_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.dynamic_conv_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   batch_group_count = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   feature_group_count = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   input_batch_dimension = #vhlo.integer_v1<0 : i64>,
-@@ -531,8 +572,9 @@
- }
- 
- // CHECK-LABEL: "default_dynamic_gather"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @default_dynamic_gather(%arg0 : tensor<2x4x9xf32>, %arg1 : tensor<1x5x2xi32>, %arg2 : tensor<3xi32>) -> tensor<1x5x8xf32> {
--  //      CHECK: "vhlo.dynamic_gather_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.dynamic_gather_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   collapsed_slice_dims = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-   // CHECK-SAME:   index_vector_dim = #vhlo.integer_v1<2 : i64>,
-   // CHECK-SAME:   indices_are_sorted = #vhlo.bool_v1<false>,
-@@ -558,15 +600,16 @@
-   // CHECK-SAME:   sym_name = #vhlo.string_v1<"default_func">,
-   // CHECK-SAME:   sym_visibility = #vhlo.string_v1<"">
-   // CHECK-SAME: }> ({
--  // CHECK-NEXT:   ^[[BB:bb.*]](%arg0: !vhlo.tensor_v1<!vhlo.f32_v1>):
--  // CHECK-NEXT:     "vhlo.return_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-+  // CHECK-NEXT:   ^[[BB:bb.*]](%[[ARG0:.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
-+  // CHECK-NEXT:     "vhlo.return_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-   // CHECK-NEXT: }) : () -> ()
-   func.return %arg0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "dynamic_gather"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @dynamic_gather(%arg0 : tensor<2x4x9xf32>, %arg1 : tensor<1x5x2xi32>) -> tensor<1x5x1xf32> {
--  //      CHECK: "vhlo.gather_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.gather_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   collapsed_slice_dims = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-   // CHECK-SAME:   index_vector_dim = #vhlo.integer_v1<2 : i64>,
-   // CHECK-SAME:   indices_are_sorted = #vhlo.bool_v1<false>,
-@@ -587,8 +630,9 @@
- }
- 
- // CHECK-LABEL: "default_infeed"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_infeed(%arg0: !stablehlo.token) -> (tensor<f32>, !stablehlo.token) {
--  //               CHECK: "vhlo.infeed_v1"(%arg0) <{
-+  //               CHECK: "vhlo.infeed_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   infeed_config = #vhlo.string_v1<"">,
-   // CHECK-SAME{LITERAL}:   layout = #vhlo.array_v1<[]>
-   //          CHECK-SAME: }> : (!vhlo.token_v1) -> (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.token_v1)
-@@ -597,8 +641,9 @@
- }
- 
- // CHECK-LABEL: "default_outfeed"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @default_outfeed(%arg0: tensor<f32>, %arg1: !stablehlo.token) -> !stablehlo.token {
--  //      CHECK: "vhlo.outfeed_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.outfeed_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   outfeed_config = #vhlo.string_v1<"">
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.token_v1) -> !vhlo.token_v1
-   %0 = "stablehlo.outfeed"(%arg0, %arg1) : (tensor<f32>, !stablehlo.token) -> !stablehlo.token
-@@ -606,8 +651,9 @@
- }
- 
- // CHECK-LABEL: "default_recv"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_recv(%arg0: !stablehlo.token) -> (tensor<f32>, !stablehlo.token) {
--  //      CHECK: "vhlo.recv_v1"(%arg0) <{
-+  //      CHECK: "vhlo.recv_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME:   channel_type = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   is_host_transfer = #vhlo.bool_v1<false>
-@@ -619,8 +665,9 @@
- }
- 
- // CHECK-LABEL: "default_send"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @default_send(%arg0: tensor<f32>, %arg1: !stablehlo.token) -> !stablehlo.token {
--  //      CHECK: "vhlo.send_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.send_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME:   channel_type = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   is_host_transfer = #vhlo.bool_v1<false>
-@@ -632,8 +679,9 @@
- }
- 
- // CHECK-LABEL: "default_reduce_scatter"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_reduce_scatter(%arg0: tensor<16xf32>) -> tensor<16xf32> {
--  //               CHECK: "vhlo.reduce_scatter_v1"(%arg0) <{
-+  //               CHECK: "vhlo.reduce_scatter_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0], [1]]> : tensor<2x1xi64>>,
-   //          CHECK-SAME:   scatter_dimension = #vhlo.integer_v1<0 : i64>
-@@ -655,8 +703,9 @@
- }
- 
- // CHECK-LABEL: "default_reduce_window"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @default_reduce_window(%arg0: tensor<2x17x31x7xf32>, %arg1: tensor<f32>) -> tensor<2x16x30x7xf32> {
--  //               CHECK: "vhlo.reduce_window_v1"(%arg0, %arg1)  <{
-+  //               CHECK: "vhlo.reduce_window_v1"(%[[ARG0]], %[[ARG1]])  <{
-   //          CHECK-SAME:   base_dilations = #vhlo.tensor_v1<dense<1> : tensor<4xi64>>,
-   // CHECK-SAME{LITERAL}:   padding = #vhlo.tensor_v1<dense<0> : tensor<4x2xi64>>,
-   //          CHECK-SAME:   window_dilations = #vhlo.tensor_v1<dense<1> : tensor<4xi64>>,
-@@ -678,8 +727,9 @@
- }
- 
- // CHECK-LABEL: "default_scatter"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @default_scatter(%arg0: tensor<200x100x300xf32>, %arg1: tensor<10x2xi32>, %arg2: tensor<10x300xf32>) -> tensor<200x100x300xf32> {
--  //      CHECK: "vhlo.scatter_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.scatter_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   index_vector_dim = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   indices_are_sorted = #vhlo.bool_v1<false>,
-   // CHECK-SAME:   inserted_window_dims = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-@@ -707,8 +757,9 @@
- }
- 
- // CHECK-LABEL: "default_select_and_scatter"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @default_select_and_scatter(%arg0: tensor<10x24x24x64xf32>, %arg1: tensor<10x23x23x64xf32>, %arg2: tensor<f32>) -> tensor<10x24x24x64xf32> {
--  //      CHECK: "vhlo.select_and_scatter_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.select_and_scatter_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   padding = #vhlo.tensor_v1<dense<0> : tensor<4x2xi64>>,
-   // CHECK-SAME:   window_dimensions = #vhlo.tensor_v1<dense<[1, 2, 2, 1]> : tensor<4xi64>>,
-   // CHECK-SAME:   window_strides = #vhlo.tensor_v1<dense<1> : tensor<4xi64>>
-@@ -736,8 +787,9 @@
- }
- 
- // CHECK-LABEL: "default_sort"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_sort(%arg0: tensor<16xf32>) -> tensor<16xf32> {
--  //      CHECK: "vhlo.sort_v1"(%arg0) <{
-+  //      CHECK: "vhlo.sort_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   dimension = #vhlo.integer_v1<-1 : i64>
-   // CHECK-SAME:   is_stable = #vhlo.bool_v1<false>
-   // CHECK-SAME: }> ({
-@@ -756,29 +808,33 @@
- // ============ OPS ============
- 
- // CHECK-LABEL: "op_abs"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_abs(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.abs_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.abs_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.abs"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_add"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_add(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_after_all"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_after_all(%arg0: !stablehlo.token) -> !stablehlo.token {
--  // CHECK: "vhlo.after_all_v1"(%arg0) : (!vhlo.token_v1) -> !vhlo.token_v1
-+  // CHECK: "vhlo.after_all_v1"(%[[ARG0]]) : (!vhlo.token_v1) -> !vhlo.token_v1
-   %0 = "stablehlo.after_all"(%arg0) : (!stablehlo.token) -> !stablehlo.token
-   func.return %0 : !stablehlo.token
- }
- 
- // CHECK-LABEL: "op_all_gather"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_all_gather(%arg0: tensor<16x8xf32>) -> tensor<16x16xf32> {
--  //               CHECK: "vhlo.all_gather_v1"(%arg0) <{
-+  //               CHECK: "vhlo.all_gather_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   all_gather_dim = #vhlo.integer_v1<1 : i64>
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0], [1]]> : tensor<2x1xi64>>,
-@@ -794,8 +850,9 @@
- }
- 
- // CHECK-LABEL: "op_all_reduce"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_all_reduce(%arg0: tensor<f32>) -> tensor<f32> {
--  //               CHECK: "vhlo.all_reduce_v1"(%arg0) <{
-+  //               CHECK: "vhlo.all_reduce_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0], [1]]> : tensor<2x1xi64>>,
-   //          CHECK-SAME:   use_global_device_ids = #vhlo.bool_v1<true>
-@@ -817,8 +874,9 @@
- }
- 
- // CHECK-LABEL: "op_all_to_all"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_all_to_all(%arg0: tensor<4x16xf32>) -> tensor<16x4xf32> {
--  //               CHECK: "vhlo.all_to_all_v1"(%arg0) <{
-+  //               CHECK: "vhlo.all_to_all_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<1 : i64>,
-   //          CHECK-SAME:   concat_dimension = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0, 1, 2, 3]]> : tensor<1x4xi64>>,
-@@ -836,22 +894,25 @@
- }
- 
- // CHECK-LABEL: "op_and"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_and(%arg0: tensor<i1>, %arg1: tensor<i1>) -> tensor<i1> {
--  // CHECK: "vhlo.and_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-+  // CHECK: "vhlo.and_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-   %0 = "stablehlo.and"(%arg0, %arg1) : (tensor<i1>, tensor<i1>) -> tensor<i1>
-   func.return %0 : tensor<i1>
- }
- 
- // CHECK-LABEL: "op_atan2"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_atan2(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.atan2_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.atan2_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.atan2"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_batch_norm_grad"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}}, %[[ARG3:.*]]: {{.*}}, %[[ARG4:.*]]: {{.*}})
- func.func @op_batch_norm_grad(%arg0: tensor<16x16x16x16xf32>, %arg1: tensor<16xf32>, %arg2: tensor<16xf32>, %arg3: tensor<16xf32>, %arg4: tensor<16x16x16x16xf32>) -> (tensor<16x16x16x16xf32>, tensor<16xf32>, tensor<16xf32>) {
--  //      CHECK: "vhlo.batch_norm_grad_v1"(%arg0, %arg1, %arg2, %arg3, %arg4) <{
-+  //      CHECK: "vhlo.batch_norm_grad_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]], %[[ARG3]], %[[ARG4]]) <{
-   // CHECK-SAME:   epsilon = #vhlo.float_v1<1.000000e-03 : !vhlo.f32_v1>,
-   // CHECK-SAME:   feature_index = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<16x16x16x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x16x16x16x!vhlo.f32_v1>) -> (!vhlo.tensor_v1<16x16x16x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>)
-@@ -863,8 +924,9 @@
- }
- 
- // CHECK-LABEL: "op_batch_norm_inference"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}}, %[[ARG3:.*]]: {{.*}}, %[[ARG4:.*]]: {{.*}})
- func.func @op_batch_norm_inference(%arg0: tensor<16x16x16x16xf32>, %arg1: tensor<16xf32>, %arg2: tensor<16xf32>, %arg3: tensor<16xf32>, %arg4: tensor<16xf32>) -> tensor<16x16x16x16xf32> {
--  //      CHECK: "vhlo.batch_norm_inference_v1"(%arg0, %arg1, %arg2, %arg3, %arg4) <{
-+  //      CHECK: "vhlo.batch_norm_inference_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]], %[[ARG3]], %[[ARG4]]) <{
-   // CHECK-SAME:   epsilon = #vhlo.float_v1<1.000000e-03 : !vhlo.f32_v1>,
-   // CHECK-SAME:   feature_index = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<16x16x16x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x16x16x16x!vhlo.f32_v1>
-@@ -876,8 +938,9 @@
- }
- 
- // CHECK-LABEL: "op_batch_norm_training"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_batch_norm_training(%arg0: tensor<16x16x16x16xf32>, %arg1: tensor<16xf32>, %arg2: tensor<16xf32>) -> (tensor<16x16x16x16xf32>, tensor<16xf32>, tensor<16xf32>) {
--  //      CHECK: "vhlo.batch_norm_training_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.batch_norm_training_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   epsilon = #vhlo.float_v1<1.000000e-03 : !vhlo.f32_v1>,
-   // CHECK-SAME:   feature_index = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<16x16x16x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>) -> (!vhlo.tensor_v1<16x16x16x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>)
-@@ -889,15 +952,17 @@
- }
- 
- // CHECK-LABEL: "op_bitcast_convert"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_bitcast_convert(%arg0: tensor<i32>) -> tensor<f32> {
--  // CHECK: "vhlo.bitcast_convert_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.bitcast_convert_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.bitcast_convert"(%arg0) : (tensor<i32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_broadcast_in_dim"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_broadcast_in_dim(%arg0: tensor<16xf32>) -> tensor<16x16xf32> {
--  //      CHECK: "vhlo.broadcast_in_dim_v1"(%arg0) <{
-+  //      CHECK: "vhlo.broadcast_in_dim_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   broadcast_dimensions = #vhlo.tensor_v1<dense<1> : tensor<1xi64>>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x16x!vhlo.f32_v1>
-   %0 = "stablehlo.broadcast_in_dim"(%arg0) {
-@@ -907,8 +972,9 @@
- }
- 
- // CHECK-LABEL: "op_broadcast"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_broadcast(%arg0: tensor<16xf32>) -> tensor<16x16xf32> {
--  //      CHECK: "vhlo.broadcast_v1"(%arg0) <{
-+  //      CHECK: "vhlo.broadcast_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   broadcast_sizes = #vhlo.tensor_v1<dense<16> : tensor<1xi64>>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x16x!vhlo.f32_v1>
-   %0 = "stablehlo.broadcast"(%arg0) {
-@@ -918,9 +984,10 @@
- }
- 
- // CHECK-LABEL: "op_case"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_case(%arg0: tensor<i32>, %arg1: tensor<f32>) -> tensor<f32> {
--  //      CHECK: "vhlo.case_v1"(%arg0) ({
--  // CHECK-NEXT:   "vhlo.return_v1"(%arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-+  //      CHECK: "vhlo.case_v1"(%[[ARG0]]) ({
-+  // CHECK-NEXT:   "vhlo.return_v1"(%[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-   // CHECK-NEXT: }) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.case"(%arg0) ({
-     "stablehlo.return"(%arg1) : (tensor<f32>) -> ()
-@@ -929,22 +996,25 @@
- }
- 
- // CHECK-LABEL: "op_cbrt"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_cbrt(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.cbrt_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.cbrt_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.cbrt"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_ceil"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_ceil(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.ceil_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.ceil_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.ceil"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_cholesky"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_cholesky(%arg0: tensor<1x16x16xf32>) -> tensor<1x16x16xf32> {
--  //      CHECK: "vhlo.cholesky_v1"(%arg0) <{
-+  //      CHECK: "vhlo.cholesky_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   lower = #vhlo.bool_v1<true>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<1x16x16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<1x16x16x!vhlo.f32_v1>
-   %0 = "stablehlo.cholesky"(%arg0) {
-@@ -954,22 +1024,25 @@
- }
- 
- // CHECK-LABEL: "op_clamp"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_clamp(%arg0: tensor<f32>, %arg1: tensor<f32>, %arg2: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.clamp_v1"(%arg0, %arg1, %arg2) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.clamp_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.clamp"(%arg0, %arg1, %arg2) : (tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_count_leading_zeros"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_count_leading_zeros(%arg0: tensor<i32>) -> tensor<i32> {
--  // CHECK: "vhlo.count_leading_zeros_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-+  // CHECK: "vhlo.count_leading_zeros_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-   %0 = "stablehlo.count_leading_zeros"(%arg0) : (tensor<i32>) -> tensor<i32>
-   func.return %0 : tensor<i32>
- }
- 
- // CHECK-LABEL: "op_collective_permute"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_collective_permute(%arg0: tensor<16x8xf32>) -> tensor<16x8xf32> {
--  //               CHECK: "vhlo.collective_permute_v1"(%arg0) <{
-+  //               CHECK: "vhlo.collective_permute_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME{LITERAL}:   source_target_pairs = #vhlo.tensor_v1<dense<[[0, 1], [1, 2], [2, 3]]> : tensor<3x2xi64>>
-   //          CHECK-SAME: }> : (!vhlo.tensor_v1<16x8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x8x!vhlo.f32_v1>
-@@ -981,8 +1054,9 @@
- }
- 
- // CHECK-LABEL: "op_compare"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_compare(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
--  //      CHECK: "vhlo.compare_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.compare_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   compare_type = #vhlo<comparison_type_v1 TOTALORDER>,
-   // CHECK-SAME:   comparison_direction = #vhlo<comparison_direction_v1 EQ>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-@@ -994,15 +1068,17 @@
- }
- 
- // CHECK-LABEL: "op_complex"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_complex(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<complex<f32>> {
--  // CHECK: "vhlo.complex_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>
-+  // CHECK: "vhlo.complex_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>
-   %0 = "stablehlo.complex"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<complex<f32>>
-   func.return %0 : tensor<complex<f32>>
- }
- 
- // CHECK-LABEL: "op_concatenate"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_concatenate(%arg0: tensor<8xf32>, %arg1: tensor<8xf32>) -> tensor<16xf32> {
--  //      CHECK: "vhlo.concatenate_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.concatenate_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   dimension = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<8x!vhlo.f32_v1>, !vhlo.tensor_v1<8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x!vhlo.f32_v1>
-   %0 = "stablehlo.concatenate"(%arg0, %arg1) {
-@@ -1012,6 +1088,7 @@
- }
- 
- // CHECK-LABEL: "op_constant"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_constant(%arg0: tensor<f32>) -> tensor<f32> {
-   //      CHECK: "vhlo.constant_v1"() <{
-   // CHECK-SAME:   value = #vhlo.tensor_v1<dense<0.000000e+00> : tensor<f32>>
-@@ -1023,15 +1100,17 @@
- }
- 
- // CHECK-LABEL: "op_convert"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_convert(%arg0: tensor<i32>) -> tensor<f32> {
--  // CHECK: "vhlo.convert_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.convert_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.convert"(%arg0) : (tensor<i32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_convolution"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_convolution(%arg0: tensor<1x8x8x207xf32>, %arg1: tensor<3x3x207x16xf32>) -> tensor<1x7x7x16xf32> {
--  //      CHECK: "vhlo.convolution_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.convolution_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   batch_group_count = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   feature_group_count = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   input_batch_dimension = #vhlo.integer_v1<0 : i64>,
-@@ -1065,8 +1144,9 @@
- }
- 
- // CHECK-LABEL: "op_cosine"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_cosine(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.cosine_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.cosine_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.cosine"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
-@@ -1079,8 +1159,9 @@
- }
- 
- // CHECK-LABEL: "op_cross_replica_sum"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_cross_replica_sum(%arg0: tensor<f32>) -> tensor<f32> {
--  //               CHECK: "vhlo.cross-replica-sum_v1"(%arg0) <{
-+  //               CHECK: "vhlo.cross-replica-sum_v1"(%[[ARG0]]) <{
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0], [1]]> : tensor<2x1xi64>>
-   //          CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.cross-replica-sum"(%arg0) {
-@@ -1090,8 +1171,9 @@
- }
- 
- // CHECK-LABEL: "op_custom_call"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_custom_call(%arg0: tensor<f32>) -> tensor<f32> {
--  //      CHECK: "vhlo.custom_call_v1"(%arg0) <{
-+  //      CHECK: "vhlo.custom_call_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   api_version = #vhlo<api_version_v1 API_VERSION_STATUS_RETURNING>,
-   // CHECK-SAME:   backend_config = #vhlo.string_v1<"\08\03\1A\02">,
-   // CHECK-SAME:   call_target_name = #vhlo.string_v1<"foo">,
-@@ -1122,15 +1204,17 @@
- }
- 
- // CHECK-LABEL: "op_divide"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_divide(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.divide_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.divide_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.divide"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_dot_general"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_dot_general(%arg0: tensor<8x8x16xf32>, %arg1: tensor<8x16x8xf32>) -> tensor<8x8x8xf32> {
--  //      CHECK: "vhlo.dot_general_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.dot_general_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   lhs_batching_dimensions = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>,
-   // CHECK-SAME:   lhs_contracting_dimensions = #vhlo.tensor_v1<dense<2> : tensor<1xi64>>,
-   // CHECK-SAME:   precision_config = #vhlo.array_v1<[#vhlo<precision_v1 HIGHEST>, #vhlo<precision_v1 HIGHEST>]>,
-@@ -1150,8 +1234,9 @@
- }
- 
- // CHECK-LABEL: "op_dot"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_dot(%arg0: tensor<8x16xf32>, %arg1: tensor<16x8xf32>) -> tensor<8x8xf32> {
--  //      CHECK: "vhlo.dot_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.dot_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   precision_config = #vhlo.array_v1<[#vhlo<precision_v1 HIGHEST>, #vhlo<precision_v1 HIGHEST>]>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<8x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<8x8x!vhlo.f32_v1>
-   %0 = "stablehlo.dot"(%arg0, %arg1) {
-@@ -1161,8 +1246,9 @@
- }
- 
- // CHECK-LABEL: "op_dynamic_broadcast_in_dim"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_dynamic_broadcast_in_dim(%arg0: tensor<?x?xf32>, %arg1: tensor<2xindex>) -> tensor<?x?xf32> {
--  //      CHECK: "vhlo.dynamic_broadcast_in_dim_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.dynamic_broadcast_in_dim_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   broadcast_dimensions = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-   // CHECK-SAME:   known_expanding_dimensions = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>,
-   // CHECK-SAME:   known_nonexpanding_dimensions = #vhlo.tensor_v1<dense<1> : tensor<1xi64>>
-@@ -1176,8 +1262,9 @@
- }
- 
- // CHECK-LABEL: "op_dynamic_conv"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_dynamic_conv(%arg0: tensor<1x8x8x207xf32>, %arg1: tensor<3x3x207x16xf32>, %arg2: tensor<4xi32>) -> tensor<1x?x?x16xf32> {
--  //      CHECK: "vhlo.dynamic_conv_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.dynamic_conv_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   batch_group_count = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   feature_group_count = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   input_batch_dimension = #vhlo.integer_v1<0 : i64>,
-@@ -1211,8 +1298,9 @@
- }
- 
- // CHECK-LABEL: "op_dynamic_gather"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_dynamic_gather(%arg0 : tensor<2x4x9xf32>, %arg1 : tensor<1x5x2xi32>, %arg2 : tensor<3xi32>) -> tensor<1x5x8xf32> {
--  //      CHECK: "vhlo.dynamic_gather_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.dynamic_gather_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   collapsed_slice_dims = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-   // CHECK-SAME:   index_vector_dim = #vhlo.integer_v1<2 : i64>,
-   // CHECK-SAME:   indices_are_sorted = #vhlo.bool_v1<true>,
-@@ -1232,8 +1320,9 @@
- }
- 
- // CHECK-LABEL: "op_dynamic_iota"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_dynamic_iota(%arg0: tensor<1xindex>) -> tensor<?xf32> {
--  //      CHECK: "vhlo.dynamic_iota_v1"(%arg0) <{
-+  //      CHECK: "vhlo.dynamic_iota_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   iota_dimension = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<1x!vhlo.index_v1>) -> !vhlo.tensor_v1<?x!vhlo.f32_v1>
-   %0 = "stablehlo.dynamic_iota"(%arg0) {
-@@ -1243,22 +1332,25 @@
- }
- 
- // CHECK-LABEL: "op_dynamic_pad"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}}, %[[ARG3:.*]]: {{.*}}, %[[ARG4:.*]]: {{.*}})
- func.func @op_dynamic_pad(%arg0: tensor<?xf32>, %arg1: tensor<f32>, %arg2: tensor<1xindex>, %arg3: tensor<1xindex>, %arg4: tensor<1xindex>) -> tensor<?xf32> {
--  // CHECK: "vhlo.dynamic_pad_v1"(%arg0, %arg1, %arg2, %arg3, %arg4) : (!vhlo.tensor_v1<?x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>) -> !vhlo.tensor_v1<?x!vhlo.f32_v1>
-+  // CHECK: "vhlo.dynamic_pad_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]], %[[ARG3]], %[[ARG4]]) : (!vhlo.tensor_v1<?x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>) -> !vhlo.tensor_v1<?x!vhlo.f32_v1>
-   %0 = "stablehlo.dynamic_pad"(%arg0, %arg1, %arg2, %arg3, %arg4) : (tensor<?xf32>, tensor<f32>, tensor<1xindex>, tensor<1xindex>, tensor<1xindex>) -> tensor<?xf32>
-   func.return %0 : tensor<?xf32>
- }
- 
- // CHECK-LABEL: "op_dynamic_reshape"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_dynamic_reshape(%arg0: tensor<16xf32>, %arg1: tensor<2xindex>) -> tensor<?x?xf32> {
--  // CHECK: "vhlo.dynamic_reshape_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<2x!vhlo.index_v1>) -> !vhlo.tensor_v1<?x?x!vhlo.f32_v1>
-+  // CHECK: "vhlo.dynamic_reshape_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<2x!vhlo.index_v1>) -> !vhlo.tensor_v1<?x?x!vhlo.f32_v1>
-   %0 = "stablehlo.dynamic_reshape"(%arg0, %arg1) : (tensor<16xf32>, tensor<2xindex>) -> tensor<?x?xf32>
-   func.return %0 : tensor<?x?xf32>
- }
- 
- // CHECK-LABEL: "op_dynamic_slice"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_dynamic_slice(%arg0: tensor<16xf32>, %arg1: tensor<i64>) -> tensor<4xf32> {
--  //      CHECK: "vhlo.dynamic_slice_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.dynamic_slice_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   slice_sizes = #vhlo.tensor_v1<dense<4> : tensor<1xi64>>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.i64_v1>) -> !vhlo.tensor_v1<4x!vhlo.f32_v1>
-   %0 = "stablehlo.dynamic_slice"(%arg0, %arg1) {
-@@ -1268,15 +1360,17 @@
- }
- 
- // CHECK-LABEL: "op_dynamic_update_slice"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_dynamic_update_slice(%arg0: tensor<16xf32>, %arg1: tensor<4xf32>, %arg2: tensor<i64>) -> tensor<16xf32> {
--  // CHECK: "vhlo.dynamic_update_slice_v1"(%arg0, %arg1, %arg2) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<4x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.i64_v1>) -> !vhlo.tensor_v1<16x!vhlo.f32_v1>
-+  // CHECK: "vhlo.dynamic_update_slice_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<4x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.i64_v1>) -> !vhlo.tensor_v1<16x!vhlo.f32_v1>
-   %0 = "stablehlo.dynamic_update_slice"(%arg0, %arg1, %arg2) : (tensor<16xf32>, tensor<4xf32>, tensor<i64>) -> tensor<16xf32>
-   func.return %0 : tensor<16xf32>
- }
- 
- // CHECK-LABEL: "op_einsum"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_einsum(%arg0: tensor<8x16xf32>, %arg1: tensor<16x8xf32>) -> tensor<8x8xf32> {
--  //      CHECK: "vhlo.einsum_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.einsum_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   einsum_config = #vhlo.string_v1<"ab,bc->ac">
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<8x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<8x8x!vhlo.f32_v1>
-   %0 = "stablehlo.einsum"(%arg0, %arg1) {
-@@ -1286,22 +1380,25 @@
- }
- 
- // CHECK-LABEL: "op_exponential_minus_one"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_exponential_minus_one(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.exponential_minus_one_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.exponential_minus_one_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.exponential_minus_one"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_exponential"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_exponential(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.exponential_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.exponential_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.exponential"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_fft"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_fft(%arg0: tensor<16xcomplex<f32>>) -> tensor<16xcomplex<f32>> {
--  //      CHECK: "vhlo.fft_v1"(%arg0) <{
-+  //      CHECK: "vhlo.fft_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   fft_length = #vhlo.tensor_v1<dense<16> : tensor<1xi64>>,
-   // CHECK-SAME:   fft_type = #vhlo<fft_type_v1 FFT>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<16x!vhlo.complex_v1<!vhlo.f32_v1>>) -> !vhlo.tensor_v1<16x!vhlo.complex_v1<!vhlo.f32_v1>>
-@@ -1313,8 +1410,9 @@
- }
- 
- // CHECK-LABEL: "op_floor"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_floor(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.floor_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.floor_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.floor"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
-@@ -1327,16 +1425,17 @@
-   // CHECK-SAME:   sym_name = #vhlo.string_v1<"op_func">,
-   // CHECK-SAME:   sym_visibility = #vhlo.string_v1<"private">
-   // CHECK-SAME: }> ({
--  // CHECK-NEXT:   ^[[BB:bb.*]](%arg0: !vhlo.tensor_v1<!vhlo.f32_v1>):
--  // CHECK-NEXT:     "vhlo.return_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-+  // CHECK-NEXT:   ^[[BB:bb.*]](%[[ARG0:.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
-+  // CHECK-NEXT:     "vhlo.return_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-   // CHECK-NEXT: }) : () -> ()
- 
-   func.return %arg0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_gather"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_gather(%arg0 : tensor<2x4x9xf32>, %arg1 : tensor<1x5x2xi32>) -> tensor<1x5x1xf32> {
--  //      CHECK: "vhlo.gather_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.gather_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   collapsed_slice_dims = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-   // CHECK-SAME:   index_vector_dim = #vhlo.integer_v1<2 : i64>,
-   // CHECK-SAME:   indices_are_sorted = #vhlo.bool_v1<true>,
-@@ -1358,8 +1457,9 @@
- }
- 
- // CHECK-LABEL: "op_get_dimension_size"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_get_dimension_size(%arg0: tensor<?xf32>) -> tensor<i32> {
--  //      CHECK: "vhlo.get_dimension_size_v1"(%arg0) <{
-+  //      CHECK: "vhlo.get_dimension_size_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   dimension = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<?x!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-   %0 = "stablehlo.get_dimension_size"(%arg0) {
-@@ -1369,8 +1469,9 @@
- }
- 
- // CHECK-LABEL: "op_get_tuple_element"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_get_tuple_element(%arg0: tuple<tensor<f32>, tensor<i32>>) -> tensor<f32> {
--  //      CHECK: "vhlo.get_tuple_element_v1"(%arg0) <{
-+  //      CHECK: "vhlo.get_tuple_element_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   index = #vhlo.integer_v1<0 : i32>
-   // CHECK-SAME: }> : (!vhlo.tuple_v1<!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.get_tuple_element"(%arg0) {
-@@ -1380,11 +1481,12 @@
- }
- 
- // CHECK-LABEL: "op_if"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_if(%arg0: tensor<i1>, %arg1: tensor<f32>, %arg2: tensor<f32>) -> tensor<f32> {
--  //      CHECK: "vhlo.if_v1"(%arg0) ({
--  // CHECK-NEXT:   "vhlo.return_v1"(%arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-+  //      CHECK: "vhlo.if_v1"(%[[ARG0]]) ({
-+  // CHECK-NEXT:   "vhlo.return_v1"(%[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-   // CHECK-NEXT: }, {
--  // CHECK-NEXT:   "vhlo.return_v1"(%arg2) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-+  // CHECK-NEXT:   "vhlo.return_v1"(%[[ARG2]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-   // CHECK-NEXT: }) : (!vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.if"(%arg0) ({
-     "stablehlo.return"(%arg1) : (tensor<f32>) -> ()
-@@ -1395,15 +1497,17 @@
- }
- 
- // CHECK-LABEL: "op_imag"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_imag(%arg0: tensor<complex<f32>>) -> tensor<f32> {
--  // CHECK: "vhlo.imag_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.imag_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.imag"(%arg0) : (tensor<complex<f32>>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_infeed"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_infeed(%arg0: !stablehlo.token) -> (tensor<f32>, !stablehlo.token) {
--  //               CHECK: "vhlo.infeed_v1"(%arg0) <{
-+  //               CHECK: "vhlo.infeed_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   infeed_config = #vhlo.string_v1<"foo">,
-   // CHECK-SAME{LITERAL}:   layout = #vhlo.array_v1<[#vhlo.array_v1<[]>]>
-   //          CHECK-SAME: }> : (!vhlo.token_v1) -> (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.token_v1)
-@@ -1426,36 +1530,41 @@
- }
- 
- // CHECK-LABEL: "op_is_finite"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_is_finite(%arg0: tensor<f32>) -> tensor<i1> {
--  // CHECK: "vhlo.is_finite_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-+  // CHECK: "vhlo.is_finite_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-   %0 = "stablehlo.is_finite"(%arg0) : (tensor<f32>) -> tensor<i1>
-   func.return %0 : tensor<i1>
- }
- 
- // CHECK-LABEL: "op_log"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_log(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.log_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.log_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.log"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_log_plus_one"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_log_plus_one(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.log_plus_one_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.log_plus_one_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.log_plus_one"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_logistic"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_logistic(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.logistic_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.logistic_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.logistic"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_map"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_map(%arg0: tensor<16xf32>) -> tensor<16xf32> {
--  //      CHECK: "vhlo.map_v1"(%arg0) <{
-+  //      CHECK: "vhlo.map_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   dimensions = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>
-   // CHECK-SAME: }> ({
-   // CHECK-NEXT:   ^[[BB:bb.*]](%[[ARG1:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
-@@ -1473,57 +1582,65 @@
- }
- 
- // CHECK-LABEL: "op_maximum"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_maximum(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.maximum_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.maximum_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.maximum"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_minimum"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_minimum(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.minimum_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.minimum_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.minimum"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_multiply"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_multiply(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.multiply_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.multiply_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.multiply"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_negate"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_negate(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.negate_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.negate_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.negate"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_not"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_not(%arg0: tensor<i1>) -> tensor<i1> {
--  // CHECK: "vhlo.not_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-+  // CHECK: "vhlo.not_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-   %0 = "stablehlo.not"(%arg0) : (tensor<i1>) -> tensor<i1>
-   func.return %0 : tensor<i1>
- }
- 
- // CHECK-LABEL: "op_optimization_barrier"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_optimization_barrier(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.optimization_barrier_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.optimization_barrier_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.optimization_barrier"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_or"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_or(%arg0: tensor<i1>, %arg1: tensor<i1>) -> tensor<i1> {
--  // CHECK: "vhlo.or_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-+  // CHECK: "vhlo.or_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-   %0 = "stablehlo.or"(%arg0, %arg1) : (tensor<i1>, tensor<i1>) -> tensor<i1>
-   func.return %0 : tensor<i1>
- }
- 
- // CHECK-LABEL: "op_outfeed"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_outfeed(%arg0: tensor<f32>, %arg1: !stablehlo.token) -> !stablehlo.token {
--  //      CHECK: "vhlo.outfeed_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.outfeed_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   outfeed_config = #vhlo.string_v1<"foo">
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.token_v1) -> !vhlo.token_v1
-   %0 = "stablehlo.outfeed"(%arg0, %arg1) {
-@@ -1533,8 +1650,9 @@
- }
- 
- // CHECK-LABEL: "op_pad"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_pad(%arg0: tensor<8xf32>, %arg1: tensor<f32>) -> tensor<16xf32> {
--  //      CHECK: "vhlo.pad_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.pad_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   edge_padding_high = #vhlo.tensor_v1<dense<4> : tensor<1xi64>>,
-   // CHECK-SAME:   edge_padding_low = #vhlo.tensor_v1<dense<4> : tensor<1xi64>>,
-   // CHECK-SAME:   interior_padding = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>
-@@ -1548,36 +1666,41 @@
- }
- 
- // CHECK-LABEL: "op_popcnt"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_popcnt(%arg0: tensor<i32>) -> tensor<i32> {
--  // CHECK: "vhlo.popcnt_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-+  // CHECK: "vhlo.popcnt_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-   %0 = "stablehlo.popcnt"(%arg0) : (tensor<i32>) -> tensor<i32>
-   func.return %0 : tensor<i32>
- }
- 
- // CHECK-LABEL: "op_power"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_power(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.power_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.power_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.power"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_real_dynamic_slice"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}}, %[[ARG3:.*]]: {{.*}})
- func.func @op_real_dynamic_slice(%arg0: tensor<?xf32>, %arg1: tensor<1xindex>, %arg2: tensor<1xindex>, %arg3: tensor<1xindex>) -> tensor<?xf32> {
--  // CHECK: "vhlo.real_dynamic_slice_v1"(%arg0, %arg1, %arg2, %arg3) : (!vhlo.tensor_v1<?x!vhlo.f32_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>) -> !vhlo.tensor_v1<?x!vhlo.f32_v1>
-+  // CHECK: "vhlo.real_dynamic_slice_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]], %[[ARG3]]) : (!vhlo.tensor_v1<?x!vhlo.f32_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>) -> !vhlo.tensor_v1<?x!vhlo.f32_v1>
-   %0 = "stablehlo.real_dynamic_slice"(%arg0, %arg1, %arg2, %arg3) : (tensor<?xf32>, tensor<1xindex>, tensor<1xindex>, tensor<1xindex>) -> tensor<?xf32>
-   func.return %0 : tensor<?xf32>
- }
- 
- // CHECK-LABEL: "op_real"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_real(%arg0: tensor<complex<f32>>) -> tensor<f32> {
--  // CHECK: "vhlo.real_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.real_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.real"(%arg0) : (tensor<complex<f32>>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_recv"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_recv(%arg0: !stablehlo.token) -> (tensor<f32>, !stablehlo.token) {
--  //      CHECK: "vhlo.recv_v1"(%arg0) <{
-+  //      CHECK: "vhlo.recv_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME:   channel_type = #vhlo.integer_v1<3 : i64>,
-   // CHECK-SAME:   is_host_transfer = #vhlo.bool_v1<true>
-@@ -1590,8 +1713,9 @@
- }
- 
- // CHECK-LABEL: "op_reduce"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_reduce(%arg0: tensor<16xf32>, %arg1: tensor<f32>) -> tensor<f32> {
--  //  CHECK: "vhlo.reduce_v1"(%arg0, %arg1)
-+  //  CHECK: "vhlo.reduce_v1"(%[[ARG0]], %[[ARG1]])
-   //  CHECK:   ^[[BB:bb.*]](%[[ARG1:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>, %[[ARG2:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
-   //  CHECK:     "vhlo.return_v1"(%[[VAL1:.*]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-   //  CHECK: }) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-@@ -1606,8 +1730,9 @@
- }
- 
- // CHECK-LABEL: "op_reduce_precision"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_reduce_precision(%arg0: tensor<f32>) -> tensor<f32> {
--  //      CHECK: "vhlo.reduce_precision_v1"(%arg0) <{
-+  //      CHECK: "vhlo.reduce_precision_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   exponent_bits = #vhlo.integer_v1<8 : i32>
-   // CHECK-SAME:   mantissa_bits = #vhlo.integer_v1<10 : i32>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-@@ -1619,8 +1744,9 @@
- }
- 
- // CHECK-LABEL: "op_reduce_scatter"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_reduce_scatter(%arg0: tensor<16xf32>) -> tensor<16xf32> {
--  //               CHECK: "vhlo.reduce_scatter_v1"(%arg0) <{
-+  //               CHECK: "vhlo.reduce_scatter_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0], [1]]> : tensor<2x1xi64>>,
-   //          CHECK-SAME:   scatter_dimension = #vhlo.integer_v1<0 : i64>
-@@ -1644,8 +1770,9 @@
- }
- 
- // CHECK-LABEL: "op_reduce_window"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_reduce_window(%arg0: tensor<2x17x31x7xf32>, %arg1: tensor<f32>) -> tensor<2x9x16x7xf32> {
--  //               CHECK: "vhlo.reduce_window_v1"(%arg0, %arg1) <{
-+  //               CHECK: "vhlo.reduce_window_v1"(%[[ARG0]], %[[ARG1]]) <{
-   //          CHECK-SAME:   base_dilations = #vhlo.tensor_v1<dense<[1, 2, 2, 1]> : tensor<4xi64>>,
-   // CHECK-SAME{LITERAL}:   padding = #vhlo.tensor_v1<dense<[[0, 0], [2, 0], [0, 2], [0, 0]]> : tensor<4x2xi64>>,
-   //          CHECK-SAME:   window_dilations = #vhlo.tensor_v1<dense<[1, 2, 2, 1]> : tensor<4xi64>>,
-@@ -1671,8 +1798,9 @@
- }
- 
- // CHECK-LABEL: "op_remainder"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_remainder(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.remainder_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.remainder_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.remainder"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
-@@ -1692,16 +1820,18 @@
- }
- 
- // CHECK-LABEL: "op_reshape"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_reshape(%arg0: tensor<16xf32>) -> tensor<4x4xf32> {
--  // CHECK: "vhlo.reshape_v1"(%arg0) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<4x4x!vhlo.f32_v1>
-+  // CHECK: "vhlo.reshape_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<4x4x!vhlo.f32_v1>
-   %0 = "stablehlo.reshape"(%arg0) : (tensor<16xf32>) -> tensor<4x4xf32>
-   func.return %0 : tensor<4x4xf32>
- }
- 
- // CHECK-LABEL: "op_return"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_return(%arg0: tensor<i32>, %arg1: tensor<f32>) -> tensor<f32> {
--  //      CHECK: "vhlo.case_v1"(%arg0) ({
--  // CHECK-NEXT:   "vhlo.return_v1"(%arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-+  //      CHECK: "vhlo.case_v1"(%[[ARG0]]) ({
-+  // CHECK-NEXT:   "vhlo.return_v1"(%[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-   // CHECK-NEXT: }) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.case"(%arg0) ({
-     "stablehlo.return"(%arg1) : (tensor<f32>) -> ()
-@@ -1710,8 +1840,9 @@
- }
- 
- // CHECK-LABEL: "op_reverse"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_reverse(%arg0: tensor<16xf32>) -> tensor<16xf32> {
--  //      CHECK: "vhlo.reverse_v1"(%arg0) <{
-+  //      CHECK: "vhlo.reverse_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   dimensions = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x!vhlo.f32_v1>
-   %0 = "stablehlo.reverse"(%arg0) {
-@@ -1721,8 +1852,9 @@
- }
- 
- // CHECK-LABEL: "op_rng_bit_generator"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_rng_bit_generator(%arg0: tensor<f32>) -> (tensor<f32>, tensor<f32>) {
--  //      CHECK: "vhlo.rng_bit_generator_v1"(%arg0) <{
-+  //      CHECK: "vhlo.rng_bit_generator_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   rng_algorithm = #vhlo<rng_algorithm_v1 PHILOX>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>)
-   %0:2 = "stablehlo.rng_bit_generator"(%arg0) {
-@@ -1732,8 +1864,9 @@
- }
- 
- // CHECK-LABEL: "op_rng"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_rng(%arg0: tensor<f32>, %arg1: tensor<f32>, %arg2: tensor<0xindex>) -> tensor<f32> {
--  //      CHECK: "vhlo.rng_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.rng_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   rng_distribution = #vhlo<rng_distribution_v1 NORMAL>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<0x!vhlo.index_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.rng"(%arg0, %arg1, %arg2) {
-@@ -1743,29 +1876,33 @@
- }
- 
- // CHECK-LABEL: "op_round_nearest_afz"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_round_nearest_afz(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.round_nearest_afz_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.round_nearest_afz_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.round_nearest_afz"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_round_nearest_even"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_round_nearest_even(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.round_nearest_even_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.round_nearest_even_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.round_nearest_even"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_rsqrt"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_rsqrt(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.rsqrt_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.rsqrt_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.rsqrt"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_scatter"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_scatter(%arg0: tensor<200x100x300xf32>, %arg1: tensor<10x2xi32>, %arg2: tensor<10x300xf32>) -> tensor<200x100x300xf32> {
--  //      CHECK: "vhlo.scatter_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.scatter_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   index_vector_dim = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   indices_are_sorted = #vhlo.bool_v1<true>,
-   // CHECK-SAME:   inserted_window_dims = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-@@ -1795,8 +1932,9 @@
- }
- 
- // CHECK-LABEL: "op_select_and_scatter"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_select_and_scatter(%arg0: tensor<10x24x24x64xf32>, %arg1: tensor<12x13x13x66xf32>, %arg2: tensor<f32>) -> tensor<10x24x24x64xf32> {
--  //      CHECK: "vhlo.select_and_scatter_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.select_and_scatter_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   padding = #vhlo.tensor_v1<dense<1> : tensor<4x2xi64>>,
-   // CHECK-SAME:   window_dimensions = #vhlo.tensor_v1<dense<[1, 2, 2, 1]> : tensor<4xi64>>,
-   // CHECK-SAME:   window_strides = #vhlo.tensor_v1<dense<[1, 2, 2, 1]> : tensor<4xi64>>
-@@ -1826,15 +1964,17 @@
- }
- 
- // CHECK-LABEL: "op_select"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_select(%arg0: tensor<i1>, %arg1: tensor<f32>, %arg2: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.select_v1"(%arg0, %arg1, %arg2) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.select_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.select"(%arg0, %arg1, %arg2) : (tensor<i1>, tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_send"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_send(%arg0: tensor<f32>, %arg1: !stablehlo.token) -> !stablehlo.token {
--  //      CHECK: "vhlo.send_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.send_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME:   channel_type = #vhlo.integer_v1<2 : i64>,
-   // CHECK-SAME:   is_host_transfer = #vhlo.bool_v1<true>
-@@ -1847,8 +1987,9 @@
- }
- 
- // CHECK-LABEL: "op_set_dimension_size"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_set_dimension_size(%arg0: tensor<?xf32>, %arg1: tensor<i32>) -> tensor<16xf32> {
--  //      CHECK: "vhlo.set_dimension_size_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.set_dimension_size_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   dimension = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<?x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<16x!vhlo.f32_v1>
-   %0 = "stablehlo.set_dimension_size"(%arg0, %arg1) {
-@@ -1858,43 +1999,49 @@
- }
- 
- // CHECK-LABEL: "op_shift_left"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_shift_left(%arg0: tensor<i32>, %arg1: tensor<i32>) -> tensor<i32> {
--  // CHECK: "vhlo.shift_left_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.i32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-+  // CHECK: "vhlo.shift_left_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-   %0 = "stablehlo.shift_left"(%arg0, %arg1) : (tensor<i32>, tensor<i32>) -> tensor<i32>
-   func.return %0 : tensor<i32>
- }
- 
- // CHECK-LABEL: "op_shift_right_arithmetic"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_shift_right_arithmetic(%arg0: tensor<i32>, %arg1: tensor<i32>) -> tensor<i32> {
--  // CHECK: "vhlo.shift_right_arithmetic_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.i32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-+  // CHECK: "vhlo.shift_right_arithmetic_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-   %0 = "stablehlo.shift_right_arithmetic"(%arg0, %arg1) : (tensor<i32>, tensor<i32>) -> tensor<i32>
-   func.return %0 : tensor<i32>
- }
- 
- // CHECK-LABEL: "op_shift_right_logical"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_shift_right_logical(%arg0: tensor<i32>, %arg1: tensor<i32>) -> tensor<i32> {
--  // CHECK: "vhlo.shift_right_logical_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.i32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-+  // CHECK: "vhlo.shift_right_logical_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-   %0 = "stablehlo.shift_right_logical"(%arg0, %arg1) : (tensor<i32>, tensor<i32>) -> tensor<i32>
-   func.return %0 : tensor<i32>
- }
- 
- // CHECK-LABEL: "op_sign"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_sign(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.sign_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.sign_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.sign"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_sine"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_sine(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.sine_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.sine_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.sine"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_slice"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_slice(%arg0: tensor<16xf32>) -> tensor<4xf32> {
--  //      CHECK: "vhlo.slice_v1"(%arg0) <{
-+  //      CHECK: "vhlo.slice_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   limit_indices = #vhlo.tensor_v1<dense<4> : tensor<1xi64>>,
-   // CHECK-SAME:   start_indices = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>,
-   // CHECK-SAME:   strides = #vhlo.tensor_v1<dense<1> : tensor<1xi64>>
-@@ -1908,8 +2055,9 @@
- }
- 
- // CHECK-LABEL: "op_sort"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_sort(%arg0: tensor<16xf32>) -> tensor<16xf32> {
--  //      CHECK: "vhlo.sort_v1"(%arg0) <{
-+  //      CHECK: "vhlo.sort_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   dimension = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME:   is_stable = #vhlo.bool_v1<true>
-   // CHECK-SAME: }> ({
-@@ -1929,29 +2077,33 @@
- }
- 
- // CHECK-LABEL: "op_sqrt"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_sqrt(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.sqrt_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.sqrt_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.sqrt"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_subtract"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_subtract(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.subtract_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.subtract_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.subtract"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_tanh"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_tanh(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.tanh_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.tanh_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.tanh"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_torch_index_select"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_torch_index_select(%arg0: tensor<5x1x5xf32>, %arg1: tensor<2xi32>) ->  tensor<2x1x5xf32> {
--  //      CHECK: "vhlo.torch_index_select_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.torch_index_select_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   batch_dims = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME:   dim = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<5x1x5x!vhlo.f32_v1>, !vhlo.tensor_v1<2x!vhlo.i32_v1>) -> !vhlo.tensor_v1<2x1x5x!vhlo.f32_v1>
-@@ -1963,8 +2115,9 @@
- }
- 
- // CHECK-LABEL: "op_trace"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_trace(%arg0: tensor<f32>) {
--  //      CHECK: "vhlo.trace_v1"(%arg0) <{
-+  //      CHECK: "vhlo.trace_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   tag = #vhlo.string_v1<"foo">
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-   "stablehlo.trace"(%arg0) {
-@@ -1974,8 +2127,9 @@
- }
- 
- // CHECK-LABEL: "op_transpose"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_transpose(%arg0: tensor<16x8xf32>) ->  tensor<8x16xf32> {
--  //      CHECK: "vhlo.transpose_v1"(%arg0) <{
-+  //      CHECK: "vhlo.transpose_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   permutation = #vhlo.tensor_v1<dense<[1, 0]> : tensor<2xi64>>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<16x8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<8x16x!vhlo.f32_v1>
-   %0 = "stablehlo.transpose"(%arg0) {
-@@ -1985,8 +2139,9 @@
- }
- 
- // CHECK-LABEL: "op_triangular_solve"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_triangular_solve(%arg0: tensor<16x16xf32>, %arg1: tensor<16x16xf32>) ->  tensor<16x16xf32> {
--  //      CHECK: "vhlo.triangular_solve_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.triangular_solve_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   left_side = #vhlo.bool_v1<true>,
-   // CHECK-SAME:   lower = #vhlo.bool_v1<true>,
-   // CHECK-SAME:   transpose_a = #vhlo<transpose_v1 NO_TRANSPOSE>,
-@@ -2002,15 +2157,17 @@
- }
- 
- // CHECK-LABEL: "op_tuple"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_tuple(%arg0: tensor<f32>) -> tuple<tensor<f32>> {
--  // CHECK: "vhlo.tuple_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tuple_v1<!vhlo.tensor_v1<!vhlo.f32_v1>>
-+  // CHECK: "vhlo.tuple_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tuple_v1<!vhlo.tensor_v1<!vhlo.f32_v1>>
-   %0 = "stablehlo.tuple"(%arg0) : (tensor<f32>) -> tuple<tensor<f32>>
-   func.return %0 : tuple<tensor<f32>>
- }
- 
- // CHECK-LABEL: "op_unary_einsum"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_unary_einsum(%arg0: tensor<8x16xf32>) -> tensor<8xf32> {
--  //      CHECK: "vhlo.unary_einsum_v1"(%arg0) <{
-+  //      CHECK: "vhlo.unary_einsum_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   einsum_config = #vhlo.string_v1<"ab->a">
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<8x16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<8x!vhlo.f32_v1>
-   %0 = "stablehlo.unary_einsum"(%arg0) {
-@@ -2020,22 +2177,25 @@
- }
- 
- // CHECK-LABEL: "op_uniform_dequantize"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_uniform_dequantize(%arg0: tensor<!quant.uniform<i8:f32, 34.0:16>>) -> tensor<f32> {
--  // CHECK: "vhlo.uniform_dequantize_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.uniform_dequantize_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.uniform_dequantize"(%arg0) : (tensor<!quant.uniform<i8:f32, 34.0:16>>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_uniform_quantize"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_uniform_quantize(%arg0: tensor<f32>) -> tensor<!quant.uniform<i8:f32, 34.0:16>> {
--  // CHECK: "vhlo.uniform_quantize_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>
-+  // CHECK: "vhlo.uniform_quantize_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>
-   %0 = "stablehlo.uniform_quantize"(%arg0) : (tensor<f32>) -> tensor<!quant.uniform<i8:f32, 34.0:16>>
-   func.return %0 : tensor<!quant.uniform<i8:f32, 34.0:16>>
- }
- 
- // CHECK-LABEL: "op_while"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_while(%arg0: tensor<i1>) -> tensor<i1> {
--  //      CHECK: "vhlo.while_v1"(%arg0) ({
-+  //      CHECK: "vhlo.while_v1"(%[[ARG0]]) ({
-   // CHECK-NEXT:   ^[[BB:bb.*]](%[[ARG1:arg.*]]: !vhlo.tensor_v1<!vhlo.bool_v1>):
-   // CHECK-NEXT:     "vhlo.return_v1"(%[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>) -> ()
-   // CHECK-NEXT:   }, {
-@@ -2053,8 +2213,9 @@
- }
- 
- // CHECK-LABEL: "op_xor"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_xor(%arg0: tensor<i1>, %arg1: tensor<i1>) -> tensor<i1> {
--  // CHECK: "vhlo.xor_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-+  // CHECK: "vhlo.xor_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-   %0 = "stablehlo.xor"(%arg0, %arg1) : (tensor<i1>, tensor<i1>) -> tensor<i1>
-   func.return %0 : tensor<i1>
- }
-@@ -2062,169 +2223,193 @@
- // ============ TYPES ============
- 
- // CHECK-LABEL: "type_i1"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_i1(%arg0: tensor<i1>, %arg1: tensor<i1>) -> tensor<i1> {
--  // CHECK: "vhlo.and_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-+  // CHECK: "vhlo.and_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-   %0 = "stablehlo.and"(%arg0, %arg1) : (tensor<i1>, tensor<i1>) -> tensor<i1>
-   func.return %0 : tensor<i1>
- }
- 
- // CHECK-LABEL: "type_i4"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_i4(%arg0: tensor<i4>, %arg1: tensor<i4>) -> tensor<i4> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.i4_v1>, !vhlo.tensor_v1<!vhlo.i4_v1>) -> !vhlo.tensor_v1<!vhlo.i4_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i4_v1>, !vhlo.tensor_v1<!vhlo.i4_v1>) -> !vhlo.tensor_v1<!vhlo.i4_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<i4>, tensor<i4>) -> tensor<i4>
-   func.return %0 : tensor<i4>
- }
- 
- // CHECK-LABEL: "type_i8"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_i8(%arg0: tensor<i8>, %arg1: tensor<i8>) -> tensor<i8> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.i8_v1>, !vhlo.tensor_v1<!vhlo.i8_v1>) -> !vhlo.tensor_v1<!vhlo.i8_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i8_v1>, !vhlo.tensor_v1<!vhlo.i8_v1>) -> !vhlo.tensor_v1<!vhlo.i8_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<i8>, tensor<i8>) -> tensor<i8>
-   func.return %0 : tensor<i8>
- }
- 
- // CHECK-LABEL: "type_i16"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_i16(%arg0: tensor<i16>, %arg1: tensor<i16>) -> tensor<i16> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.i16_v1>, !vhlo.tensor_v1<!vhlo.i16_v1>) -> !vhlo.tensor_v1<!vhlo.i16_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i16_v1>, !vhlo.tensor_v1<!vhlo.i16_v1>) -> !vhlo.tensor_v1<!vhlo.i16_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<i16>, tensor<i16>) -> tensor<i16>
-   func.return %0 : tensor<i16>
- }
- 
- // CHECK-LABEL: "type_i32"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_i32(%arg0: tensor<i32>, %arg1: tensor<i32>) -> tensor<i32> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.i32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<i32>, tensor<i32>) -> tensor<i32>
-   func.return %0 : tensor<i32>
- }
- 
- // CHECK-LABEL: "type_i64"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_i64(%arg0: tensor<i64>, %arg1: tensor<i64>) -> tensor<i64> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.i64_v1>, !vhlo.tensor_v1<!vhlo.i64_v1>) -> !vhlo.tensor_v1<!vhlo.i64_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i64_v1>, !vhlo.tensor_v1<!vhlo.i64_v1>) -> !vhlo.tensor_v1<!vhlo.i64_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<i64>, tensor<i64>) -> tensor<i64>
-   func.return %0 : tensor<i64>
- }
- 
- // CHECK-LABEL: "type_ui4"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_ui4(%arg0: tensor<ui4>, %arg1: tensor<ui4>) -> tensor<ui4> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.ui4_v1>, !vhlo.tensor_v1<!vhlo.ui4_v1>) -> !vhlo.tensor_v1<!vhlo.ui4_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.ui4_v1>, !vhlo.tensor_v1<!vhlo.ui4_v1>) -> !vhlo.tensor_v1<!vhlo.ui4_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<ui4>, tensor<ui4>) -> tensor<ui4>
-   func.return %0 : tensor<ui4>
- }
- 
- // CHECK-LABEL: "type_ui8"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_ui8(%arg0: tensor<ui8>, %arg1: tensor<ui8>) -> tensor<ui8> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.ui8_v1>, !vhlo.tensor_v1<!vhlo.ui8_v1>) -> !vhlo.tensor_v1<!vhlo.ui8_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.ui8_v1>, !vhlo.tensor_v1<!vhlo.ui8_v1>) -> !vhlo.tensor_v1<!vhlo.ui8_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<ui8>, tensor<ui8>) -> tensor<ui8>
-   func.return %0 : tensor<ui8>
- }
- 
- // CHECK-LABEL: "type_ui16"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_ui16(%arg0: tensor<ui16>, %arg1: tensor<ui16>) -> tensor<ui16> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.ui16_v1>, !vhlo.tensor_v1<!vhlo.ui16_v1>) -> !vhlo.tensor_v1<!vhlo.ui16_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.ui16_v1>, !vhlo.tensor_v1<!vhlo.ui16_v1>) -> !vhlo.tensor_v1<!vhlo.ui16_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<ui16>, tensor<ui16>) -> tensor<ui16>
-   func.return %0 : tensor<ui16>
- }
- 
- // CHECK-LABEL: "type_ui32"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_ui32(%arg0: tensor<ui32>, %arg1: tensor<ui32>) -> tensor<ui32> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.ui32_v1>, !vhlo.tensor_v1<!vhlo.ui32_v1>) -> !vhlo.tensor_v1<!vhlo.ui32_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.ui32_v1>, !vhlo.tensor_v1<!vhlo.ui32_v1>) -> !vhlo.tensor_v1<!vhlo.ui32_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<ui32>, tensor<ui32>) -> tensor<ui32>
-   func.return %0 : tensor<ui32>
- }
- 
- // CHECK-LABEL: "type_ui64"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_ui64(%arg0: tensor<ui64>, %arg1: tensor<ui64>) -> tensor<ui64> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.ui64_v1>, !vhlo.tensor_v1<!vhlo.ui64_v1>) -> !vhlo.tensor_v1<!vhlo.ui64_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.ui64_v1>, !vhlo.tensor_v1<!vhlo.ui64_v1>) -> !vhlo.tensor_v1<!vhlo.ui64_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<ui64>, tensor<ui64>) -> tensor<ui64>
-   func.return %0 : tensor<ui64>
- }
- 
- // CHECK-LABEL: "type_f8E4M3FN"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_f8E4M3FN(%arg0: tensor<f8E4M3FN>, %arg1: tensor<f8E4M3FN>) -> tensor<f8E4M3FN> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f8E4M3FN_v1>, !vhlo.tensor_v1<!vhlo.f8E4M3FN_v1>) -> !vhlo.tensor_v1<!vhlo.f8E4M3FN_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f8E4M3FN_v1>, !vhlo.tensor_v1<!vhlo.f8E4M3FN_v1>) -> !vhlo.tensor_v1<!vhlo.f8E4M3FN_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f8E4M3FN>, tensor<f8E4M3FN>) -> tensor<f8E4M3FN>
-   func.return %0 : tensor<f8E4M3FN>
- }
- 
- // CHECK-LABEL: "type_f8E5M2"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_f8E5M2(%arg0: tensor<f8E5M2>, %arg1: tensor<f8E5M2>) -> tensor<f8E5M2> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f8E5M2_v1>, !vhlo.tensor_v1<!vhlo.f8E5M2_v1>) -> !vhlo.tensor_v1<!vhlo.f8E5M2_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f8E5M2_v1>, !vhlo.tensor_v1<!vhlo.f8E5M2_v1>) -> !vhlo.tensor_v1<!vhlo.f8E5M2_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f8E5M2>, tensor<f8E5M2>) -> tensor<f8E5M2>
-   func.return %0 : tensor<f8E5M2>
- }
- 
- // CHECK-LABEL: "type_bf16"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_bf16(%arg0: tensor<bf16>, %arg1: tensor<bf16>) -> tensor<bf16> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.bf16_v1>, !vhlo.tensor_v1<!vhlo.bf16_v1>) -> !vhlo.tensor_v1<!vhlo.bf16_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.bf16_v1>, !vhlo.tensor_v1<!vhlo.bf16_v1>) -> !vhlo.tensor_v1<!vhlo.bf16_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<bf16>, tensor<bf16>) -> tensor<bf16>
-   func.return %0 : tensor<bf16>
- }
- 
- // CHECK-LABEL: "type_f16"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_f16(%arg0: tensor<f16>, %arg1: tensor<f16>) -> tensor<f16> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f16_v1>, !vhlo.tensor_v1<!vhlo.f16_v1>) -> !vhlo.tensor_v1<!vhlo.f16_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f16_v1>, !vhlo.tensor_v1<!vhlo.f16_v1>) -> !vhlo.tensor_v1<!vhlo.f16_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f16>, tensor<f16>) -> tensor<f16>
-   func.return %0 : tensor<f16>
- }
- 
- // CHECK-LABEL: "type_f32"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_f32(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "type_f64"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_f64(%arg0: tensor<f64>, %arg1: tensor<f64>) -> tensor<f64> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f64_v1>, !vhlo.tensor_v1<!vhlo.f64_v1>) -> !vhlo.tensor_v1<!vhlo.f64_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f64_v1>, !vhlo.tensor_v1<!vhlo.f64_v1>) -> !vhlo.tensor_v1<!vhlo.f64_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f64>, tensor<f64>) -> tensor<f64>
-   func.return %0 : tensor<f64>
- }
- 
- // CHECK-LABEL: "type_complex_f32"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_complex_f32(%arg0: tensor<complex<f32>>, %arg1: tensor<complex<f32>>) -> tensor<complex<f32>> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>, !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>) -> !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>, !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>) -> !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<complex<f32>>, tensor<complex<f32>>) -> tensor<complex<f32>>
-   func.return %0 : tensor<complex<f32>>
- }
- 
- // CHECK-LABEL: "type_complex_f64"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_complex_f64(%arg0: tensor<complex<f64>>, %arg1: tensor<complex<f64>>) -> tensor<complex<f64>> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f64_v1>>, !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f64_v1>>) -> !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f64_v1>>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f64_v1>>, !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f64_v1>>) -> !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f64_v1>>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<complex<f64>>, tensor<complex<f64>>) -> tensor<complex<f64>>
-   func.return %0 : tensor<complex<f64>>
- }
- 
- // CHECK-LABEL: "type_dynamism_ranked"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @type_dynamism_ranked(%arg0: tensor<?xf32>) -> tensor<?xf32> {
--  // CHECK: "vhlo.abs_v1"(%arg0) : (!vhlo.tensor_v1<?x!vhlo.f32_v1>) -> !vhlo.tensor_v1<?x!vhlo.f32_v1>
-+  // CHECK: "vhlo.abs_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<?x!vhlo.f32_v1>) -> !vhlo.tensor_v1<?x!vhlo.f32_v1>
-   %0 = "stablehlo.abs"(%arg0) : (tensor<?xf32>) -> tensor<?xf32>
-   func.return %0 : tensor<?xf32>
- }
- 
- // CHECK-LABEL: "type_quantization"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_quantization(%arg0: tensor<!quant.uniform<i8:f32, 34.0:16>>, %arg1: tensor<!quant.uniform<i8:f32, 34.0:16>>) -> tensor<!quant.uniform<i8:f32, 34.0:16>> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>, !vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>) -> !vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>, !vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>) -> !vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<!quant.uniform<i8:f32, 34.0:16>>, tensor<!quant.uniform<i8:f32, 34.0:16>>) -> tensor<!quant.uniform<i8:f32, 34.0:16>>
-   func.return %0 : tensor<!quant.uniform<i8:f32, 34.0:16>>
- }
- 
- //       CHECK: function_type = #vhlo.type_v1<!vhlo.func_v1<(!vhlo.token_v1) -> !vhlo.token_v1>>
- // CHECK-LABEL: "type_token_callee"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @type_token_callee(%arg0: !stablehlo.token) -> !stablehlo.token {
--  // CHECK: "vhlo.return_v1"(%arg0) : (!vhlo.token_v1) -> ()
-+  // CHECK: "vhlo.return_v1"(%[[ARG0]]) : (!vhlo.token_v1) -> ()
-   return %arg0 : !stablehlo.token
- }
- 
- //       CHECK: function_type = #vhlo.type_v1<!vhlo.func_v1<(!vhlo.token_v1) -> !vhlo.token_v1>>
- // CHECK-LABEL: "type_token_caller"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @type_token_caller(%arg0: !stablehlo.token) -> !stablehlo.token {
--  // CHECK:      "vhlo.call_v1"(%arg0) <{callee = #vhlo.string_v1<"type_token_callee">}
-+  // CHECK:      "vhlo.call_v1"(%[[ARG0]]) <{callee = #vhlo.string_v1<"type_token_callee">}
-   // CHECK-SAME: (!vhlo.token_v1) -> !vhlo.token_v1
-   %0 = func.call @type_token_callee(%arg0) : (!stablehlo.token) -> !stablehlo.token
-   return %0 : !stablehlo.token
- }
- 
- // CHECK-LABEL: "type_tuple"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @type_tuple(%arg0: tuple<tensor<f32>>) -> tuple<!stablehlo.token> {
-   %0 = "stablehlo.custom_call"(%arg0) {
-     call_target_name = "foo"
-diff --ruN a/stablehlo/stablehlo/tests/vhlo/stablehlo_legalize_to_vhlo.mlir b/stablehlo/stablehlo/tests/vhlo/stablehlo_legalize_to_vhlo.mlir
---- stablehlo/stablehlo/tests/vhlo/stablehlo_legalize_to_vhlo.mlir
-+++ stablehlo/stablehlo/tests/vhlo/stablehlo_legalize_to_vhlo.mlir
-@@ -13,6 +13,7 @@
- // ============ ATTRIBUTES ============
- 
- // CHECK-LABEL: "attr_comparison_direction_eq"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_direction_eq(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     // CHECK: comparison_direction = #vhlo<comparison_direction_v1 EQ>
-@@ -22,6 +23,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_direction_ne"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_direction_ne(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     // CHECK: comparison_direction = #vhlo<comparison_direction_v1 NE>
-@@ -31,6 +33,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_direction_ge"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_direction_ge(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     // CHECK: comparison_direction = #vhlo<comparison_direction_v1 GE>
-@@ -40,6 +43,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_direction_gt"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_direction_gt(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     // CHECK: comparison_direction = #vhlo<comparison_direction_v1 GT>
-@@ -49,6 +53,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_direction_le"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_direction_le(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     // CHECK: comparison_direction = #vhlo<comparison_direction_v1 LE>
-@@ -58,6 +63,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_direction_lt"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_direction_lt(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     // CHECK: comparison_direction = #vhlo<comparison_direction_v1 LT>
-@@ -67,6 +73,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_type_notype"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_type_notype(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     comparison_direction = #stablehlo<comparison_direction EQ>
-@@ -76,6 +83,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_type_float"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_type_float(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     comparison_direction = #stablehlo<comparison_direction EQ>,
-@@ -86,6 +94,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_type_totalorder"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_type_totalorder(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     comparison_direction = #stablehlo<comparison_direction EQ>,
-@@ -96,6 +105,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_type_signed"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_type_signed(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     comparison_direction = #stablehlo<comparison_direction EQ>,
-@@ -106,6 +116,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_type_unsigned"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_type_unsigned(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     comparison_direction = #stablehlo<comparison_direction EQ>,
-@@ -118,6 +129,7 @@
- // ConvDimensionNumbers aka #stablehlo.conv is covered below.
- 
- // CHECK-LABEL: "attr_custom_call_api_version_unspecified"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_custom_call_api_version_unspecified(%arg0: tensor<f32>) -> tensor<f32> {
-   %0 = "stablehlo.custom_call"(%arg0) {
-     call_target_name = "foo",
-@@ -128,6 +140,7 @@
- }
- 
- // CHECK-LABEL: "attr_custom_call_api_version_original"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_custom_call_api_version_original(%arg0: tensor<f32>) -> tensor<f32> {
-   %0 = "stablehlo.custom_call"(%arg0) {
-     call_target_name = "foo",
-@@ -138,6 +151,7 @@
- }
- 
- // CHECK-LABEL: "attr_custom_call_api_version_status_returning"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_custom_call_api_version_status_returning(%arg0: tensor<f32>) -> tensor<f32> {
-   %0 = "stablehlo.custom_call"(%arg0) {
-     call_target_name = "foo",
-@@ -148,6 +162,7 @@
- }
- 
- // CHECK-LABEL: "attr_custom_call_api_version_status_returning_unified"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_custom_call_api_version_status_returning_unified(%arg0: tensor<f32>) -> tensor<f32> {
-   %0 = "stablehlo.custom_call"(%arg0) {
-     call_target_name = "foo",
-@@ -166,6 +181,7 @@
- // DotDimensionNumbers aka #stablehlo.dot is covered below.
- 
- // CHECK-LABEL: "attr_fft_type_fft"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_fft_type_fft(%arg0: tensor<16xcomplex<f32>>) -> tensor<16xcomplex<f32>> {
-   %0 = "stablehlo.fft"(%arg0) {
-     // CHECK: fft_type = #vhlo<fft_type_v1 FFT>
-@@ -176,6 +192,7 @@
- }
- 
- // CHECK-LABEL: "attr_fft_type_ifft"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_fft_type_ifft(%arg0: tensor<16xcomplex<f32>>) -> tensor<16xcomplex<f32>> {
-   %0 = "stablehlo.fft"(%arg0) {
-     // CHECK: fft_type = #vhlo<fft_type_v1 IFFT>
-@@ -186,6 +203,7 @@
- }
- 
- // CHECK-LABEL: "attr_fft_type_rfft"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_fft_type_rfft(%arg0: tensor<16xf32>) -> tensor<9xcomplex<f32>> {
-   %0 = "stablehlo.fft"(%arg0) {
-     // CHECK: fft_type = #vhlo<fft_type_v1 RFFT>
-@@ -196,6 +214,7 @@
- }
- 
- // CHECK-LABEL: "attr_fft_type_irfft"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_fft_type_irfft(%arg0: tensor<9xcomplex<f32>>) -> tensor<16xf32> {
-   %0 = "stablehlo.fft"(%arg0) {
-     // CHECK: fft_type = #vhlo<fft_type_v1 IRFFT>
-@@ -208,6 +227,7 @@
- // GatherDimensionNumbers aka #stablehlo.gather is covered below.
- 
- // CHECK-LABEL: "attr_precision_config_default"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_precision_config_default(%arg0: tensor<8x16xf32>, %arg1: tensor<16x8xf32>) -> tensor<8x8xf32> {
-   %0 = "stablehlo.dot"(%arg0, %arg1) {
-     // CHECK: precision_config = #vhlo.array_v1<[#vhlo<precision_v1 DEFAULT>, #vhlo<precision_v1 DEFAULT>]>
-@@ -216,6 +236,7 @@
- }
- 
- // CHECK-LABEL: "attr_precision_config_high"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_precision_config_high(%arg0: tensor<8x16xf32>, %arg1: tensor<16x8xf32>) -> tensor<8x8xf32> {
-   %0 = "stablehlo.dot"(%arg0, %arg1) {
-     // CHECK: precision_config = #vhlo.array_v1<[#vhlo<precision_v1 HIGH>, #vhlo<precision_v1 HIGH>]>
-@@ -225,6 +246,7 @@
- }
- 
- // CHECK-LABEL: "attr_precision_config_highest"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_precision_config_highest(%arg0: tensor<8x16xf32>, %arg1: tensor<16x8xf32>) -> tensor<8x8xf32> {
-   %0 = "stablehlo.dot"(%arg0, %arg1) {
-     // CHECK: precision_config = #vhlo.array_v1<[#vhlo<precision_v1 HIGHEST>, #vhlo<precision_v1 HIGHEST>]>
-@@ -234,6 +256,7 @@
- }
- 
- // CHECK-LABEL: "attr_rng_algorithm_default"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_rng_algorithm_default(%arg0: tensor<f32>) -> (tensor<f32>, tensor<f32>) {
-   %0:2 = "stablehlo.rng_bit_generator"(%arg0) {
-     // CHECK: rng_algorithm = #vhlo<rng_algorithm_v1 DEFAULT>
-@@ -243,6 +266,7 @@
- }
- 
- // CHECK-LABEL: "attr_rng_algorithm_three_fry"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_rng_algorithm_three_fry(%arg0: tensor<f32>) -> (tensor<f32>, tensor<f32>) {
-   %0:2 = "stablehlo.rng_bit_generator"(%arg0) {
-     // CHECK: rng_algorithm = #vhlo<rng_algorithm_v1 THREE_FRY>
-@@ -252,6 +276,7 @@
- }
- 
- // CHECK-LABEL: "attr_rng_algorithm_philox"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_rng_algorithm_philox(%arg0: tensor<f32>) -> (tensor<f32>, tensor<f32>) {
-   %0:2 = "stablehlo.rng_bit_generator"(%arg0) {
-     // CHECK: rng_algorithm = #vhlo<rng_algorithm_v1 PHILOX>
-@@ -261,6 +286,7 @@
- }
- 
- // CHECK-LABEL: "attr_rng_distribution_uniform"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @attr_rng_distribution_uniform(%arg0: tensor<f32>, %arg1: tensor<f32>, %arg2: tensor<0xindex>) -> tensor<f32> {
-   %0 = "stablehlo.rng"(%arg0, %arg1, %arg2) {
-     // CHECK: rng_distribution = #vhlo<rng_distribution_v1 UNIFORM>
-@@ -270,6 +296,7 @@
- }
- 
- // CHECK-LABEL: "attr_rng_distribution_normal"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @attr_rng_distribution_normal(%arg0: tensor<f32>, %arg1: tensor<f32>, %arg2: tensor<0xindex>) -> tensor<f32> {
-   %0 = "stablehlo.rng"(%arg0, %arg1, %arg2) {
-     // CHECK: rng_distribution = #vhlo<rng_distribution_v1 NORMAL>
-@@ -281,6 +308,7 @@
- // ScatterDimensionNumbers aka #stablehlo.scatter is covered below.
- 
- // CHECK-LABEL: "attr_transpose_no_transpose"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_transpose_no_transpose(%arg0: tensor<16x16xf32>, %arg1: tensor<16x16xf32>) ->  tensor<16x16xf32> {
-   %0 = "stablehlo.triangular_solve"(%arg0, %arg1) {
-     left_side = true,
-@@ -293,6 +321,7 @@
- }
- 
- // CHECK-LABEL: "attr_transpose_transpose"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_transpose_transpose(%arg0: tensor<16x16xf32>, %arg1: tensor<16x16xf32>) ->  tensor<16x16xf32> {
-   %0 = "stablehlo.triangular_solve"(%arg0, %arg1) {
-     left_side = true,
-@@ -305,6 +334,7 @@
- }
- 
- // CHECK-LABEL: "attr_transpose_adjoint"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_transpose_adjoint(%arg0: tensor<16x16xf32>, %arg1: tensor<16x16xf32>) ->  tensor<16x16xf32> {
-   %0 = "stablehlo.triangular_solve"(%arg0, %arg1) {
-     left_side = true,
-@@ -319,10 +349,9 @@
- // TypeExtensionsAttr aka #stablehlo.type_extensions is covered below.
- 
- // CHECK-LABEL: "attr_type_extensions_bounds"
--func.func @attr_type_extensions_bounds(
--    %arg0: tensor<?x?xf32, #stablehlo.type_extensions<bounds = [16, ?]>>)
--    -> tensor<?x?xf32, #stablehlo.type_extensions<bounds = [16, ?]>> {
--  // CHECK: "vhlo.return_v1"(%arg0) : (!vhlo.tensor_v1<?x?x!vhlo.f32_v1, #vhlo.type_extensions_v1<bounds = [16, ?]>>) -> ()
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @attr_type_extensions_bounds(%arg0: tensor<?x?xf32, #stablehlo.type_extensions<bounds = [16, ?]>>) -> tensor<?x?xf32, #stablehlo.type_extensions<bounds = [16, ?]>> {
-+  // CHECK: "vhlo.return_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<?x?x!vhlo.f32_v1, #vhlo.type_extensions_v1<bounds = [16, ?]>>) -> ()
-   func.return %arg0 : tensor<?x?xf32, #stablehlo.type_extensions<bounds = [16, ?]>>
- }
- 
-@@ -330,8 +359,9 @@
- // ============ DEFAULTS ============
- 
- // CHECK-LABEL: "default_all_gather"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_all_gather(%arg0: tensor<16x8xf32>) -> tensor<16x16xf32> {
--  //               CHECK: "vhlo.all_gather_v1"(%arg0) <{
-+  //               CHECK: "vhlo.all_gather_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   all_gather_dim = #vhlo.integer_v1<1 : i64>
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0], [1]]> : tensor<2x1xi64>>,
-@@ -345,8 +375,9 @@
- }
- 
- // CHECK-LABEL: "default_all_reduce"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_all_reduce(%arg0: tensor<f32>) -> tensor<f32> {
--  //               CHECK: "vhlo.all_reduce_v1"(%arg0)
-+  //               CHECK: "vhlo.all_reduce_v1"(%[[ARG0]])
-   //          CHECK-SAME: <{
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0], [1]]> : tensor<2x1xi64>>,
-@@ -368,8 +399,9 @@
- }
- 
- // CHECK-LABEL: "default_all_to_all"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_all_to_all(%arg0: tensor<4x16xf32>) -> tensor<16x4xf32> {
--  //               CHECK: "vhlo.all_to_all_v1"(%arg0) <{
-+  //               CHECK: "vhlo.all_to_all_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   //          CHECK-SAME:   concat_dimension = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0, 1, 2, 3]]> : tensor<1x4xi64>>,
-@@ -386,8 +418,9 @@
- }
- 
- // CHECK-LABEL: "default_cholesky"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_cholesky(%arg0: tensor<1x16x16xf32>) -> tensor<1x16x16xf32> {
--  //      CHECK: "vhlo.cholesky_v1"(%arg0) <{
-+  //      CHECK: "vhlo.cholesky_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   lower = #vhlo.bool_v1<false>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<1x16x16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<1x16x16x!vhlo.f32_v1>
-   %0 = "stablehlo.cholesky"(%arg0) : (tensor<1x16x16xf32>) -> tensor<1x16x16xf32>
-@@ -395,8 +428,9 @@
- }
- 
- // CHECK-LABEL: "default_collective_permute"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_collective_permute(%arg0: tensor<16x8xf32>) -> tensor<16x8xf32> {
--  //               CHECK: "vhlo.collective_permute_v1"(%arg0) <{
-+  //               CHECK: "vhlo.collective_permute_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME{LITERAL}:   source_target_pairs = #vhlo.tensor_v1<dense<[[0, 1], [1, 2], [2, 3]]> : tensor<3x2xi64>>
-   //          CHECK-SAME: }> : (!vhlo.tensor_v1<16x8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x8x!vhlo.f32_v1>
-@@ -407,8 +441,9 @@
- }
- 
- // CHECK-LABEL: "default_collective_broadcast"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_collective_broadcast(%arg0: tensor<16x8xf32>) -> tensor<16x8xf32> {
--  //               CHECK: "vhlo.collective_broadcast_v1"(%arg0) <{
-+  //               CHECK: "vhlo.collective_broadcast_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0, 1]]> : tensor<1x2xi64>>
-   //          CHECK-SAME: }> : (!vhlo.tensor_v1<16x8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x8x!vhlo.f32_v1>
-@@ -419,8 +454,9 @@
- }
- 
- // CHECK-LABEL: "default_compare"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @default_compare(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
--  //      CHECK: "vhlo.compare_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.compare_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   compare_type = #vhlo<comparison_type_v1 NOTYPE>,
-   // CHECK-SAME:   comparison_direction = #vhlo<comparison_direction_v1 EQ>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-@@ -431,8 +467,9 @@
- }
- 
- // CHECK-LABEL: "default_composite"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_composite(%arg0: tensor<f32>) -> tensor<f32> {
--  //               CHECK: "vhlo.composite_v1"(%arg0) <{
-+  //               CHECK: "vhlo.composite_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   composite_attributes = #vhlo.dict_v1<{}>
-   //          CHECK-SAME:   decomposition = #vhlo.string_v1<"composite_target">
-   //          CHECK-SAME:   name = #vhlo.string_v1<"stablehlo.composite_target">
-@@ -446,8 +483,9 @@
- }
- 
- // CHECK-LABEL: "default_convolution"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @default_convolution(%arg0: tensor<1x8x8x207xf32>, %arg1: tensor<3x3x207x16xf32>) -> tensor<1x6x6x16xf32> {
--  //      CHECK: "vhlo.convolution_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.convolution_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   batch_group_count = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   feature_group_count = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   input_batch_dimension = #vhlo.integer_v1<0 : i64>,
-@@ -475,8 +513,9 @@
- }
- 
- // CHECK-LABEL: "default_custom_call"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_custom_call(%arg0: tensor<f32>) -> tensor<f32> {
--  //      CHECK: "vhlo.custom_call_v1"(%arg0) <{
-+  //      CHECK: "vhlo.custom_call_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   api_version = #vhlo<api_version_v1 API_VERSION_ORIGINAL>,
-   // CHECK-SAME:   backend_config = #vhlo.string_v1<"">,
-   // CHECK-SAME:   call_target_name = #vhlo.string_v1<"foo">,
-@@ -493,8 +532,9 @@
- }
- 
- // CHECK-LABEL: "default_dot_general"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @default_dot_general(%arg0: tensor<8x8x16xf32>, %arg1: tensor<8x16x8xf32>) -> tensor<8x8x8xf32> {
--  //      CHECK: "vhlo.dot_general_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.dot_general_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   lhs_batching_dimensions = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>,
-   // CHECK-SAME:   lhs_contracting_dimensions = #vhlo.tensor_v1<dense<2> : tensor<1xi64>>,
-   // CHECK-SAME:   precision_config = #vhlo.array_v1<[#vhlo<precision_v1 DEFAULT>, #vhlo<precision_v1 DEFAULT>]>,
-@@ -513,8 +553,9 @@
- }
- 
- // CHECK-LABEL: "default_dot"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @default_dot(%arg0: tensor<8x16xf32>, %arg1: tensor<16x8xf32>) -> tensor<8x8xf32> {
--  //      CHECK: "vhlo.dot_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.dot_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   precision_config = #vhlo.array_v1<[#vhlo<precision_v1 DEFAULT>, #vhlo<precision_v1 DEFAULT>]>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<8x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<8x8x!vhlo.f32_v1>
-   %0 = "stablehlo.dot"(%arg0, %arg1) : (tensor<8x16xf32>, tensor<16x8xf32>) -> tensor<8x8xf32>
-@@ -522,8 +563,9 @@
- }
- 
- // CHECK-LABEL: "default_dynamic_broadcast_in_dim"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @default_dynamic_broadcast_in_dim(%arg0: tensor<?x?xf32>, %arg1: tensor<2xindex>) -> tensor<?x?xf32> {
--  //      CHECK: "vhlo.dynamic_broadcast_in_dim_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.dynamic_broadcast_in_dim_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   broadcast_dimensions = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-   // CHECK-SAME:   known_expanding_dimensions = #vhlo.tensor_v1<dense<> : tensor<0xi64>>,
-   // CHECK-SAME:   known_nonexpanding_dimensions = #vhlo.tensor_v1<dense<> : tensor<0xi64>>
-@@ -535,8 +577,9 @@
- }
- 
- // CHECK-LABEL: "default_dynamic_conv"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @default_dynamic_conv(%arg0: tensor<1x8x8x207xf32>, %arg1: tensor<3x3x207x16xf32>, %arg2: tensor<4xi32>) -> tensor<1x?x?x16xf32> {
--  //      CHECK: "vhlo.dynamic_conv_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.dynamic_conv_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   batch_group_count = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   feature_group_count = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   input_batch_dimension = #vhlo.integer_v1<0 : i64>,
-@@ -564,8 +607,9 @@
- }
- 
- // CHECK-LABEL: "default_dynamic_gather"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @default_dynamic_gather(%arg0 : tensor<2x4x9xf32>, %arg1 : tensor<1x5x2xi32>, %arg2 : tensor<3xi32>) -> tensor<1x5x8xf32> {
--  //      CHECK: "vhlo.dynamic_gather_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.dynamic_gather_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   collapsed_slice_dims = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-   // CHECK-SAME:   index_vector_dim = #vhlo.integer_v1<2 : i64>,
-   // CHECK-SAME:   indices_are_sorted = #vhlo.bool_v1<false>,
-@@ -591,15 +635,16 @@
-   // CHECK-SAME:   sym_name = #vhlo.string_v1<"default_func">,
-   // CHECK-SAME:   sym_visibility = #vhlo.string_v1<"">
-   // CHECK-SAME: }> ({
--  // CHECK-NEXT:   ^[[BB:bb.*]](%arg0: !vhlo.tensor_v1<!vhlo.f32_v1>):
--  // CHECK-NEXT:     "vhlo.return_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-+  // CHECK-NEXT:   ^[[BB:bb.*]](%[[ARG0:.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
-+  // CHECK-NEXT:     "vhlo.return_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-   // CHECK-NEXT: }) : () -> ()
-   func.return %arg0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "dynamic_gather"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @dynamic_gather(%arg0 : tensor<2x4x9xf32>, %arg1 : tensor<1x5x2xi32>) -> tensor<1x5x1xf32> {
--  //      CHECK: "vhlo.gather_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.gather_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   collapsed_slice_dims = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-   // CHECK-SAME:   index_vector_dim = #vhlo.integer_v1<2 : i64>,
-   // CHECK-SAME:   indices_are_sorted = #vhlo.bool_v1<false>,
-@@ -620,8 +665,9 @@
- }
- 
- // CHECK-LABEL: "default_infeed"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_infeed(%arg0: !stablehlo.token) -> (tensor<f32>, !stablehlo.token) {
--  //               CHECK: "vhlo.infeed_v1"(%arg0) <{
-+  //               CHECK: "vhlo.infeed_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   infeed_config = #vhlo.string_v1<"">,
-   // CHECK-SAME{LITERAL}:   layout = #vhlo.array_v1<[]>
-   //          CHECK-SAME: }> : (!vhlo.token_v1) -> (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.token_v1)
-@@ -630,8 +676,9 @@
- }
- 
- // CHECK-LABEL: "default_outfeed"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @default_outfeed(%arg0: tensor<f32>, %arg1: !stablehlo.token) -> !stablehlo.token {
--  //      CHECK: "vhlo.outfeed_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.outfeed_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   outfeed_config = #vhlo.string_v1<"">
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.token_v1) -> !vhlo.token_v1
-   %0 = "stablehlo.outfeed"(%arg0, %arg1) : (tensor<f32>, !stablehlo.token) -> !stablehlo.token
-@@ -639,8 +686,9 @@
- }
- 
- // CHECK-LABEL: "default_recv"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_recv(%arg0: !stablehlo.token) -> (tensor<f32>, !stablehlo.token) {
--  //      CHECK: "vhlo.recv_v1"(%arg0) <{
-+  //      CHECK: "vhlo.recv_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME:   channel_type = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   is_host_transfer = #vhlo.bool_v1<false>
-@@ -652,8 +700,9 @@
- }
- 
- // CHECK-LABEL: "default_send"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @default_send(%arg0: tensor<f32>, %arg1: !stablehlo.token) -> !stablehlo.token {
--  //      CHECK: "vhlo.send_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.send_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME:   channel_type = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   is_host_transfer = #vhlo.bool_v1<false>
-@@ -665,8 +714,9 @@
- }
- 
- // CHECK-LABEL: "default_reduce_scatter"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_reduce_scatter(%arg0: tensor<16xf32>) -> tensor<16xf32> {
--  //               CHECK: "vhlo.reduce_scatter_v1"(%arg0) <{
-+  //               CHECK: "vhlo.reduce_scatter_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0], [1]]> : tensor<2x1xi64>>,
-   //          CHECK-SAME:   scatter_dimension = #vhlo.integer_v1<0 : i64>
-@@ -688,8 +738,9 @@
- }
- 
- // CHECK-LABEL: "default_reduce_window"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @default_reduce_window(%arg0: tensor<2x17x31x7xf32>, %arg1: tensor<f32>) -> tensor<2x16x30x7xf32> {
--  //               CHECK: "vhlo.reduce_window_v1"(%arg0, %arg1)  <{
-+  //               CHECK: "vhlo.reduce_window_v1"(%[[ARG0]], %[[ARG1]])  <{
-   //          CHECK-SAME:   base_dilations = #vhlo.tensor_v1<dense<1> : tensor<4xi64>>,
-   // CHECK-SAME{LITERAL}:   padding = #vhlo.tensor_v1<dense<0> : tensor<4x2xi64>>,
-   //          CHECK-SAME:   window_dilations = #vhlo.tensor_v1<dense<1> : tensor<4xi64>>,
-@@ -711,8 +762,9 @@
- }
- 
- // CHECK-LABEL: "default_scatter"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @default_scatter(%arg0: tensor<200x100x300xf32>, %arg1: tensor<10x2xi32>, %arg2: tensor<10x300xf32>) -> tensor<200x100x300xf32> {
--  //      CHECK: "vhlo.scatter_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.scatter_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   index_vector_dim = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   indices_are_sorted = #vhlo.bool_v1<false>,
-   // CHECK-SAME:   inserted_window_dims = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-@@ -740,8 +792,9 @@
- }
- 
- // CHECK-LABEL: "default_select_and_scatter"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @default_select_and_scatter(%arg0: tensor<10x24x24x64xf32>, %arg1: tensor<10x23x23x64xf32>, %arg2: tensor<f32>) -> tensor<10x24x24x64xf32> {
--  //      CHECK: "vhlo.select_and_scatter_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.select_and_scatter_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   padding = #vhlo.tensor_v1<dense<0> : tensor<4x2xi64>>,
-   // CHECK-SAME:   window_dimensions = #vhlo.tensor_v1<dense<[1, 2, 2, 1]> : tensor<4xi64>>,
-   // CHECK-SAME:   window_strides = #vhlo.tensor_v1<dense<1> : tensor<4xi64>>
-@@ -769,8 +822,9 @@
- }
- 
- // CHECK-LABEL: "default_sort"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_sort(%arg0: tensor<16xf32>) -> tensor<16xf32> {
--  //      CHECK: "vhlo.sort_v1"(%arg0) <{
-+  //      CHECK: "vhlo.sort_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   dimension = #vhlo.integer_v1<-1 : i64>
-   // CHECK-SAME:   is_stable = #vhlo.bool_v1<false>
-   // CHECK-SAME: }> ({
-@@ -789,29 +843,33 @@
- // ============ OPS ============
- 
- // CHECK-LABEL: "op_abs"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_abs(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.abs_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.abs_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.abs"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_add"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_add(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_after_all"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_after_all(%arg0: !stablehlo.token) -> !stablehlo.token {
--  // CHECK: "vhlo.after_all_v1"(%arg0) : (!vhlo.token_v1) -> !vhlo.token_v1
-+  // CHECK: "vhlo.after_all_v1"(%[[ARG0]]) : (!vhlo.token_v1) -> !vhlo.token_v1
-   %0 = "stablehlo.after_all"(%arg0) : (!stablehlo.token) -> !stablehlo.token
-   func.return %0 : !stablehlo.token
- }
- 
- // CHECK-LABEL: "op_all_gather"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_all_gather(%arg0: tensor<16x8xf32>) -> tensor<16x16xf32> {
--  //               CHECK: "vhlo.all_gather_v1"(%arg0) <{
-+  //               CHECK: "vhlo.all_gather_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   all_gather_dim = #vhlo.integer_v1<1 : i64>
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0], [1]]> : tensor<2x1xi64>>,
-@@ -827,8 +885,9 @@
- }
- 
- // CHECK-LABEL: "op_all_reduce"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_all_reduce(%arg0: tensor<f32>) -> tensor<f32> {
--  //               CHECK: "vhlo.all_reduce_v1"(%arg0) <{
-+  //               CHECK: "vhlo.all_reduce_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0], [1]]> : tensor<2x1xi64>>,
-   //          CHECK-SAME:   use_global_device_ids = #vhlo.bool_v1<true>
-@@ -851,7 +910,7 @@
- 
- // CHECK-LABEL: "op_all_reduce_with_promotable_types"
- func.func @op_all_reduce_with_promotable_types(%operand: tensor<f32>) -> tensor<f64> {
--  //  CHECK: "vhlo.all_reduce_v1"(%arg0)
-+  //  CHECK: "vhlo.all_reduce_v1"(%[[ARG0:.*]])
-   //  CHECK:   ^[[BB:bb.*]](%[[ARG1:arg.*]]: !vhlo.tensor_v1<!vhlo.f64_v1>, %[[ARG2:arg.*]]: !vhlo.tensor_v1<!vhlo.f64_v1>):
-   //  CHECK:     "vhlo.return_v1"(%[[VAL1:.*]]) : (!vhlo.tensor_v1<!vhlo.f64_v1>) -> ()
-   //  CHECK: }) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f64_v1>
-@@ -869,8 +928,9 @@
- }
- 
- // CHECK-LABEL: "op_all_to_all"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_all_to_all(%arg0: tensor<4x16xf32>) -> tensor<16x4xf32> {
--  //               CHECK: "vhlo.all_to_all_v1"(%arg0) <{
-+  //               CHECK: "vhlo.all_to_all_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<1 : i64>,
-   //          CHECK-SAME:   concat_dimension = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0, 1, 2, 3]]> : tensor<1x4xi64>>,
-@@ -888,22 +948,25 @@
- }
- 
- // CHECK-LABEL: "op_and"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_and(%arg0: tensor<i1>, %arg1: tensor<i1>) -> tensor<i1> {
--  // CHECK: "vhlo.and_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-+  // CHECK: "vhlo.and_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-   %0 = "stablehlo.and"(%arg0, %arg1) : (tensor<i1>, tensor<i1>) -> tensor<i1>
-   func.return %0 : tensor<i1>
- }
- 
- // CHECK-LABEL: "op_atan2"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_atan2(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.atan2_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.atan2_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.atan2"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_batch_norm_grad"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}}, %[[ARG3:.*]]: {{.*}}, %[[ARG4:.*]]: {{.*}})
- func.func @op_batch_norm_grad(%arg0: tensor<16x16x16x16xf32>, %arg1: tensor<16xf32>, %arg2: tensor<16xf32>, %arg3: tensor<16xf32>, %arg4: tensor<16x16x16x16xf32>) -> (tensor<16x16x16x16xf32>, tensor<16xf32>, tensor<16xf32>) {
--  //      CHECK: "vhlo.batch_norm_grad_v1"(%arg0, %arg1, %arg2, %arg3, %arg4) <{
-+  //      CHECK: "vhlo.batch_norm_grad_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]], %[[ARG3]], %[[ARG4]]) <{
-   // CHECK-SAME:   epsilon = #vhlo.float_v1<1.000000e-03 : !vhlo.f32_v1>,
-   // CHECK-SAME:   feature_index = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<16x16x16x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x16x16x16x!vhlo.f32_v1>) -> (!vhlo.tensor_v1<16x16x16x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>)
-@@ -915,8 +978,9 @@
- }
- 
- // CHECK-LABEL: "op_batch_norm_inference"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}}, %[[ARG3:.*]]: {{.*}}, %[[ARG4:.*]]: {{.*}})
- func.func @op_batch_norm_inference(%arg0: tensor<16x16x16x16xf32>, %arg1: tensor<16xf32>, %arg2: tensor<16xf32>, %arg3: tensor<16xf32>, %arg4: tensor<16xf32>) -> tensor<16x16x16x16xf32> {
--  //      CHECK: "vhlo.batch_norm_inference_v1"(%arg0, %arg1, %arg2, %arg3, %arg4) <{
-+  //      CHECK: "vhlo.batch_norm_inference_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]], %[[ARG3]], %[[ARG4]]) <{
-   // CHECK-SAME:   epsilon = #vhlo.float_v1<1.000000e-03 : !vhlo.f32_v1>,
-   // CHECK-SAME:   feature_index = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<16x16x16x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x16x16x16x!vhlo.f32_v1>
-@@ -928,8 +992,9 @@
- }
- 
- // CHECK-LABEL: "op_batch_norm_training"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_batch_norm_training(%arg0: tensor<16x16x16x16xf32>, %arg1: tensor<16xf32>, %arg2: tensor<16xf32>) -> (tensor<16x16x16x16xf32>, tensor<16xf32>, tensor<16xf32>) {
--  //      CHECK: "vhlo.batch_norm_training_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.batch_norm_training_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   epsilon = #vhlo.float_v1<1.000000e-03 : !vhlo.f32_v1>,
-   // CHECK-SAME:   feature_index = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<16x16x16x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>) -> (!vhlo.tensor_v1<16x16x16x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>)
-@@ -941,15 +1006,17 @@
- }
- 
- // CHECK-LABEL: "op_bitcast_convert"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_bitcast_convert(%arg0: tensor<i32>) -> tensor<f32> {
--  // CHECK: "vhlo.bitcast_convert_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.bitcast_convert_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.bitcast_convert"(%arg0) : (tensor<i32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_broadcast_in_dim"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_broadcast_in_dim(%arg0: tensor<16xf32>) -> tensor<16x16xf32> {
--  //      CHECK: "vhlo.broadcast_in_dim_v1"(%arg0) <{
-+  //      CHECK: "vhlo.broadcast_in_dim_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   broadcast_dimensions = #vhlo.tensor_v1<dense<1> : tensor<1xi64>>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x16x!vhlo.f32_v1>
-   %0 = "stablehlo.broadcast_in_dim"(%arg0) {
-@@ -959,8 +1026,9 @@
- }
- 
- // CHECK-LABEL: "op_broadcast"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_broadcast(%arg0: tensor<16xf32>) -> tensor<16x16xf32> {
--  //      CHECK: "vhlo.broadcast_v1"(%arg0) <{
-+  //      CHECK: "vhlo.broadcast_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   broadcast_sizes = #vhlo.tensor_v1<dense<16> : tensor<1xi64>>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x16x!vhlo.f32_v1>
-   %0 = "stablehlo.broadcast"(%arg0) {
-@@ -970,9 +1038,10 @@
- }
- 
- // CHECK-LABEL: "op_case"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_case(%arg0: tensor<i32>, %arg1: tensor<f32>) -> tensor<f32> {
--  //      CHECK: "vhlo.case_v1"(%arg0) ({
--  // CHECK-NEXT:   "vhlo.return_v1"(%arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-+  //      CHECK: "vhlo.case_v1"(%[[ARG0]]) ({
-+  // CHECK-NEXT:   "vhlo.return_v1"(%[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-   // CHECK-NEXT: }) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.case"(%arg0) ({
-     "stablehlo.return"(%arg1) : (tensor<f32>) -> ()
-@@ -981,22 +1050,25 @@
- }
- 
- // CHECK-LABEL: "op_cbrt"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_cbrt(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.cbrt_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.cbrt_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.cbrt"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_ceil"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_ceil(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.ceil_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.ceil_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.ceil"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_cholesky"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_cholesky(%arg0: tensor<1x16x16xf32>) -> tensor<1x16x16xf32> {
--  //      CHECK: "vhlo.cholesky_v1"(%arg0) <{
-+  //      CHECK: "vhlo.cholesky_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   lower = #vhlo.bool_v1<true>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<1x16x16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<1x16x16x!vhlo.f32_v1>
-   %0 = "stablehlo.cholesky"(%arg0) {
-@@ -1006,22 +1078,25 @@
- }
- 
- // CHECK-LABEL: "op_clamp"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_clamp(%arg0: tensor<f32>, %arg1: tensor<f32>, %arg2: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.clamp_v1"(%arg0, %arg1, %arg2) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.clamp_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.clamp"(%arg0, %arg1, %arg2) : (tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_count_leading_zeros"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_count_leading_zeros(%arg0: tensor<i32>) -> tensor<i32> {
--  // CHECK: "vhlo.count_leading_zeros_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-+  // CHECK: "vhlo.count_leading_zeros_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-   %0 = "stablehlo.count_leading_zeros"(%arg0) : (tensor<i32>) -> tensor<i32>
-   func.return %0 : tensor<i32>
- }
- 
- // CHECK-LABEL: "op_collective_permute"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_collective_permute(%arg0: tensor<16x8xf32>) -> tensor<16x8xf32> {
--  //               CHECK: "vhlo.collective_permute_v1"(%arg0) <{
-+  //               CHECK: "vhlo.collective_permute_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME{LITERAL}:   source_target_pairs = #vhlo.tensor_v1<dense<[[0, 1], [1, 2], [2, 3]]> : tensor<3x2xi64>>
-   //          CHECK-SAME: }> : (!vhlo.tensor_v1<16x8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x8x!vhlo.f32_v1>
-@@ -1033,8 +1108,9 @@
- }
- 
- // CHECK-LABEL: "op_compare"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_compare(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
--  //      CHECK: "vhlo.compare_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.compare_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   compare_type = #vhlo<comparison_type_v1 TOTALORDER>,
-   // CHECK-SAME:   comparison_direction = #vhlo<comparison_direction_v1 EQ>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-@@ -1046,15 +1122,17 @@
- }
- 
- // CHECK-LABEL: "op_complex"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_complex(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<complex<f32>> {
--  // CHECK: "vhlo.complex_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>
-+  // CHECK: "vhlo.complex_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>
-   %0 = "stablehlo.complex"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<complex<f32>>
-   func.return %0 : tensor<complex<f32>>
- }
- 
- // CHECK-LABEL: "op_composite"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_composite(%arg0: tensor<f32>) -> tensor<f32> {
--  //               CHECK: "vhlo.composite_v1"(%arg0) <{
-+  //               CHECK: "vhlo.composite_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   composite_attributes = #vhlo.dict_v1<{#vhlo.string_v1<"my_int"> = #vhlo.integer_v1<1 : i64>, #vhlo.string_v1<"my_string"> = #vhlo.string_v1<"foo">}>
-   //          CHECK-SAME:   decomposition = #vhlo.string_v1<"composite_target">
-   //          CHECK-SAME:   name = #vhlo.string_v1<"stablehlo.composite_target">
-@@ -1073,8 +1151,9 @@
- }
- 
- // CHECK-LABEL: "op_concatenate"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_concatenate(%arg0: tensor<8xf32>, %arg1: tensor<8xf32>) -> tensor<16xf32> {
--  //      CHECK: "vhlo.concatenate_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.concatenate_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   dimension = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<8x!vhlo.f32_v1>, !vhlo.tensor_v1<8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x!vhlo.f32_v1>
-   %0 = "stablehlo.concatenate"(%arg0, %arg1) {
-@@ -1084,6 +1163,7 @@
- }
- 
- // CHECK-LABEL: "op_constant"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_constant(%arg0: tensor<f32>) -> tensor<f32> {
-   //      CHECK: "vhlo.constant_v1"() <{
-   // CHECK-SAME:   value = #vhlo.tensor_v1<dense<0.000000e+00> : tensor<f32>>
-@@ -1095,15 +1175,17 @@
- }
- 
- // CHECK-LABEL: "op_convert"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_convert(%arg0: tensor<i32>) -> tensor<f32> {
--  // CHECK: "vhlo.convert_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.convert_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.convert"(%arg0) : (tensor<i32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_convolution"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_convolution(%arg0: tensor<1x8x8x207xf32>, %arg1: tensor<3x3x207x16xf32>) -> tensor<1x7x7x16xf32> {
--  //      CHECK: "vhlo.convolution_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.convolution_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   batch_group_count = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   feature_group_count = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   input_batch_dimension = #vhlo.integer_v1<0 : i64>,
-@@ -1137,8 +1219,9 @@
- }
- 
- // CHECK-LABEL: "op_cosine"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_cosine(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.cosine_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.cosine_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.cosine"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
-@@ -1151,8 +1234,9 @@
- }
- 
- // CHECK-LABEL: "op_cross_replica_sum"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_cross_replica_sum(%arg0: tensor<f32>) -> tensor<f32> {
--  //               CHECK: "vhlo.cross-replica-sum_v1"(%arg0) <{
-+  //               CHECK: "vhlo.cross-replica-sum_v1"(%[[ARG0]]) <{
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0], [1]]> : tensor<2x1xi64>>
-   //          CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.cross-replica-sum"(%arg0) {
-@@ -1162,8 +1246,9 @@
- }
- 
- // CHECK-LABEL: "op_custom_call"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_custom_call(%arg0: tensor<f32>) -> tensor<f32> {
--  //      CHECK: "vhlo.custom_call_v1"(%arg0) <{
-+  //      CHECK: "vhlo.custom_call_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   api_version = #vhlo<api_version_v1 API_VERSION_STATUS_RETURNING>,
-   // CHECK-SAME:   backend_config = #vhlo.string_v1<"\08\03\1A\02">,
-   // CHECK-SAME:   call_target_name = #vhlo.string_v1<"foo">,
-@@ -1194,15 +1279,17 @@
- }
- 
- // CHECK-LABEL: "op_divide"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_divide(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.divide_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.divide_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.divide"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_dot_general"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_dot_general(%arg0: tensor<8x8x16xf32>, %arg1: tensor<8x16x8xf32>) -> tensor<8x8x8xf32> {
--  //      CHECK: "vhlo.dot_general_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.dot_general_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   lhs_batching_dimensions = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>,
-   // CHECK-SAME:   lhs_contracting_dimensions = #vhlo.tensor_v1<dense<2> : tensor<1xi64>>,
-   // CHECK-SAME:   precision_config = #vhlo.array_v1<[#vhlo<precision_v1 HIGHEST>, #vhlo<precision_v1 HIGHEST>]>,
-@@ -1222,8 +1309,9 @@
- }
- 
- // CHECK-LABEL: "op_dot"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_dot(%arg0: tensor<8x16xf32>, %arg1: tensor<16x8xf32>) -> tensor<8x8xf32> {
--  //      CHECK: "vhlo.dot_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.dot_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   precision_config = #vhlo.array_v1<[#vhlo<precision_v1 HIGHEST>, #vhlo<precision_v1 HIGHEST>]>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<8x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<8x8x!vhlo.f32_v1>
-   %0 = "stablehlo.dot"(%arg0, %arg1) {
-@@ -1233,8 +1321,9 @@
- }
- 
- // CHECK-LABEL: "op_dynamic_broadcast_in_dim"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_dynamic_broadcast_in_dim(%arg0: tensor<?x?xf32>, %arg1: tensor<2xindex>) -> tensor<?x?xf32> {
--  //      CHECK: "vhlo.dynamic_broadcast_in_dim_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.dynamic_broadcast_in_dim_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   broadcast_dimensions = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-   // CHECK-SAME:   known_expanding_dimensions = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>,
-   // CHECK-SAME:   known_nonexpanding_dimensions = #vhlo.tensor_v1<dense<1> : tensor<1xi64>>
-@@ -1248,8 +1337,9 @@
- }
- 
- // CHECK-LABEL: "op_dynamic_conv"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_dynamic_conv(%arg0: tensor<1x8x8x207xf32>, %arg1: tensor<3x3x207x16xf32>, %arg2: tensor<4xi32>) -> tensor<1x?x?x16xf32> {
--  //      CHECK: "vhlo.dynamic_conv_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.dynamic_conv_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   batch_group_count = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   feature_group_count = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   input_batch_dimension = #vhlo.integer_v1<0 : i64>,
-@@ -1283,8 +1373,9 @@
- }
- 
- // CHECK-LABEL: "op_dynamic_gather"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_dynamic_gather(%arg0 : tensor<2x4x9xf32>, %arg1 : tensor<1x5x2xi32>, %arg2 : tensor<3xi32>) -> tensor<1x5x8xf32> {
--  //      CHECK: "vhlo.dynamic_gather_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.dynamic_gather_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   collapsed_slice_dims = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-   // CHECK-SAME:   index_vector_dim = #vhlo.integer_v1<2 : i64>,
-   // CHECK-SAME:   indices_are_sorted = #vhlo.bool_v1<true>,
-@@ -1304,8 +1395,9 @@
- }
- 
- // CHECK-LABEL: "op_dynamic_iota"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_dynamic_iota(%arg0: tensor<1xindex>) -> tensor<?xf32> {
--  //      CHECK: "vhlo.dynamic_iota_v1"(%arg0) <{
-+  //      CHECK: "vhlo.dynamic_iota_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   iota_dimension = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<1x!vhlo.index_v1>) -> !vhlo.tensor_v1<?x!vhlo.f32_v1>
-   %0 = "stablehlo.dynamic_iota"(%arg0) {
-@@ -1315,22 +1407,25 @@
- }
- 
- // CHECK-LABEL: "op_dynamic_pad"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}}, %[[ARG3:.*]]: {{.*}}, %[[ARG4:.*]]: {{.*}})
- func.func @op_dynamic_pad(%arg0: tensor<?xf32>, %arg1: tensor<f32>, %arg2: tensor<1xindex>, %arg3: tensor<1xindex>, %arg4: tensor<1xindex>) -> tensor<?xf32> {
--  // CHECK: "vhlo.dynamic_pad_v1"(%arg0, %arg1, %arg2, %arg3, %arg4) : (!vhlo.tensor_v1<?x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>) -> !vhlo.tensor_v1<?x!vhlo.f32_v1>
-+  // CHECK: "vhlo.dynamic_pad_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]], %[[ARG3]], %[[ARG4]]) : (!vhlo.tensor_v1<?x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>) -> !vhlo.tensor_v1<?x!vhlo.f32_v1>
-   %0 = "stablehlo.dynamic_pad"(%arg0, %arg1, %arg2, %arg3, %arg4) : (tensor<?xf32>, tensor<f32>, tensor<1xindex>, tensor<1xindex>, tensor<1xindex>) -> tensor<?xf32>
-   func.return %0 : tensor<?xf32>
- }
- 
- // CHECK-LABEL: "op_dynamic_reshape"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_dynamic_reshape(%arg0: tensor<16xf32>, %arg1: tensor<2xindex>) -> tensor<?x?xf32> {
--  // CHECK: "vhlo.dynamic_reshape_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<2x!vhlo.index_v1>) -> !vhlo.tensor_v1<?x?x!vhlo.f32_v1>
-+  // CHECK: "vhlo.dynamic_reshape_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<2x!vhlo.index_v1>) -> !vhlo.tensor_v1<?x?x!vhlo.f32_v1>
-   %0 = "stablehlo.dynamic_reshape"(%arg0, %arg1) : (tensor<16xf32>, tensor<2xindex>) -> tensor<?x?xf32>
-   func.return %0 : tensor<?x?xf32>
- }
- 
- // CHECK-LABEL: "op_dynamic_slice"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_dynamic_slice(%arg0: tensor<16xf32>, %arg1: tensor<i64>) -> tensor<4xf32> {
--  //      CHECK: "vhlo.dynamic_slice_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.dynamic_slice_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   slice_sizes = #vhlo.tensor_v1<dense<4> : tensor<1xi64>>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.i64_v1>) -> !vhlo.tensor_v1<4x!vhlo.f32_v1>
-   %0 = "stablehlo.dynamic_slice"(%arg0, %arg1) {
-@@ -1340,15 +1435,17 @@
- }
- 
- // CHECK-LABEL: "op_dynamic_update_slice"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_dynamic_update_slice(%arg0: tensor<16xf32>, %arg1: tensor<4xf32>, %arg2: tensor<i64>) -> tensor<16xf32> {
--  // CHECK: "vhlo.dynamic_update_slice_v1"(%arg0, %arg1, %arg2) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<4x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.i64_v1>) -> !vhlo.tensor_v1<16x!vhlo.f32_v1>
-+  // CHECK: "vhlo.dynamic_update_slice_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<4x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.i64_v1>) -> !vhlo.tensor_v1<16x!vhlo.f32_v1>
-   %0 = "stablehlo.dynamic_update_slice"(%arg0, %arg1, %arg2) : (tensor<16xf32>, tensor<4xf32>, tensor<i64>) -> tensor<16xf32>
-   func.return %0 : tensor<16xf32>
- }
- 
- // CHECK-LABEL: "op_einsum"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_einsum(%arg0: tensor<8x16xf32>, %arg1: tensor<16x8xf32>) -> tensor<8x8xf32> {
--  //      CHECK: "vhlo.einsum_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.einsum_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   einsum_config = #vhlo.string_v1<"ab,bc->ac">
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<8x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<8x8x!vhlo.f32_v1>
-   %0 = "stablehlo.einsum"(%arg0, %arg1) {
-@@ -1358,22 +1455,25 @@
- }
- 
- // CHECK-LABEL: "op_exponential_minus_one"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_exponential_minus_one(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.exponential_minus_one_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.exponential_minus_one_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.exponential_minus_one"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_exponential"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_exponential(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.exponential_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.exponential_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.exponential"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_fft"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_fft(%arg0: tensor<16xcomplex<f32>>) -> tensor<16xcomplex<f32>> {
--  //      CHECK: "vhlo.fft_v1"(%arg0) <{
-+  //      CHECK: "vhlo.fft_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   fft_length = #vhlo.tensor_v1<dense<16> : tensor<1xi64>>,
-   // CHECK-SAME:   fft_type = #vhlo<fft_type_v1 FFT>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<16x!vhlo.complex_v1<!vhlo.f32_v1>>) -> !vhlo.tensor_v1<16x!vhlo.complex_v1<!vhlo.f32_v1>>
-@@ -1385,8 +1485,9 @@
- }
- 
- // CHECK-LABEL: "op_floor"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_floor(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.floor_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.floor_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.floor"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
-@@ -1399,16 +1500,17 @@
-   // CHECK-SAME:   sym_name = #vhlo.string_v1<"op_func">,
-   // CHECK-SAME:   sym_visibility = #vhlo.string_v1<"private">
-   // CHECK-SAME: }> ({
--  // CHECK-NEXT:   ^[[BB:bb.*]](%arg0: !vhlo.tensor_v1<!vhlo.f32_v1>):
--  // CHECK-NEXT:     "vhlo.return_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-+  // CHECK-NEXT:   ^[[BB:bb.*]](%[[ARG0:.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
-+  // CHECK-NEXT:     "vhlo.return_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-   // CHECK-NEXT: }) : () -> ()
- 
-   func.return %arg0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_gather"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_gather(%arg0 : tensor<2x4x9xf32>, %arg1 : tensor<1x5x2xi32>) -> tensor<1x5x1xf32> {
--  //      CHECK: "vhlo.gather_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.gather_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   collapsed_slice_dims = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-   // CHECK-SAME:   index_vector_dim = #vhlo.integer_v1<2 : i64>,
-   // CHECK-SAME:   indices_are_sorted = #vhlo.bool_v1<true>,
-@@ -1430,8 +1532,9 @@
- }
- 
- // CHECK-LABEL: "op_get_dimension_size"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_get_dimension_size(%arg0: tensor<?xf32>) -> tensor<i32> {
--  //      CHECK: "vhlo.get_dimension_size_v1"(%arg0) <{
-+  //      CHECK: "vhlo.get_dimension_size_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   dimension = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<?x!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-   %0 = "stablehlo.get_dimension_size"(%arg0) {
-@@ -1441,8 +1544,9 @@
- }
- 
- // CHECK-LABEL: "op_get_tuple_element"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_get_tuple_element(%arg0: tuple<tensor<f32>, tensor<i32>>) -> tensor<f32> {
--  //      CHECK: "vhlo.get_tuple_element_v1"(%arg0) <{
-+  //      CHECK: "vhlo.get_tuple_element_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   index = #vhlo.integer_v1<0 : i32>
-   // CHECK-SAME: }> : (!vhlo.tuple_v1<!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.get_tuple_element"(%arg0) {
-@@ -1452,11 +1556,12 @@
- }
- 
- // CHECK-LABEL: "op_if"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_if(%arg0: tensor<i1>, %arg1: tensor<f32>, %arg2: tensor<f32>) -> tensor<f32> {
--  //      CHECK: "vhlo.if_v1"(%arg0) ({
--  // CHECK-NEXT:   "vhlo.return_v1"(%arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-+  //      CHECK: "vhlo.if_v1"(%[[ARG0]]) ({
-+  // CHECK-NEXT:   "vhlo.return_v1"(%[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-   // CHECK-NEXT: }, {
--  // CHECK-NEXT:   "vhlo.return_v1"(%arg2) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-+  // CHECK-NEXT:   "vhlo.return_v1"(%[[ARG2]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-   // CHECK-NEXT: }) : (!vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.if"(%arg0) ({
-     "stablehlo.return"(%arg1) : (tensor<f32>) -> ()
-@@ -1467,15 +1572,17 @@
- }
- 
- // CHECK-LABEL: "op_imag"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_imag(%arg0: tensor<complex<f32>>) -> tensor<f32> {
--  // CHECK: "vhlo.imag_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.imag_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.imag"(%arg0) : (tensor<complex<f32>>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_infeed"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_infeed(%arg0: !stablehlo.token) -> (tensor<f32>, !stablehlo.token) {
--  //               CHECK: "vhlo.infeed_v1"(%arg0) <{
-+  //               CHECK: "vhlo.infeed_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   infeed_config = #vhlo.string_v1<"foo">,
-   // CHECK-SAME{LITERAL}:   layout = #vhlo.array_v1<[#vhlo.array_v1<[]>]>
-   //          CHECK-SAME: }> : (!vhlo.token_v1) -> (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.token_v1)
-@@ -1498,36 +1605,41 @@
- }
- 
- // CHECK-LABEL: "op_is_finite"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_is_finite(%arg0: tensor<f32>) -> tensor<i1> {
--  // CHECK: "vhlo.is_finite_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-+  // CHECK: "vhlo.is_finite_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-   %0 = "stablehlo.is_finite"(%arg0) : (tensor<f32>) -> tensor<i1>
-   func.return %0 : tensor<i1>
- }
- 
- // CHECK-LABEL: "op_log"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_log(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.log_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.log_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.log"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_log_plus_one"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_log_plus_one(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.log_plus_one_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.log_plus_one_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.log_plus_one"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_logistic"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_logistic(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.logistic_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.logistic_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.logistic"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_map"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_map(%arg0: tensor<16xf32>) -> tensor<16xf32> {
--  //      CHECK: "vhlo.map_v1"(%arg0) <{
-+  //      CHECK: "vhlo.map_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   dimensions = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>
-   // CHECK-SAME: }> ({
-   // CHECK-NEXT:   ^[[BB:bb.*]](%[[ARG1:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
-@@ -1545,57 +1657,65 @@
- }
- 
- // CHECK-LABEL: "op_maximum"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_maximum(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.maximum_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.maximum_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.maximum"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_minimum"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_minimum(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.minimum_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.minimum_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.minimum"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_multiply"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_multiply(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.multiply_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.multiply_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.multiply"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_negate"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_negate(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.negate_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.negate_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.negate"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_not"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_not(%arg0: tensor<i1>) -> tensor<i1> {
--  // CHECK: "vhlo.not_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-+  // CHECK: "vhlo.not_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-   %0 = "stablehlo.not"(%arg0) : (tensor<i1>) -> tensor<i1>
-   func.return %0 : tensor<i1>
- }
- 
- // CHECK-LABEL: "op_optimization_barrier"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_optimization_barrier(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.optimization_barrier_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.optimization_barrier_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.optimization_barrier"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_or"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_or(%arg0: tensor<i1>, %arg1: tensor<i1>) -> tensor<i1> {
--  // CHECK: "vhlo.or_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-+  // CHECK: "vhlo.or_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-   %0 = "stablehlo.or"(%arg0, %arg1) : (tensor<i1>, tensor<i1>) -> tensor<i1>
-   func.return %0 : tensor<i1>
- }
- 
- // CHECK-LABEL: "op_outfeed"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_outfeed(%arg0: tensor<f32>, %arg1: !stablehlo.token) -> !stablehlo.token {
--  //      CHECK: "vhlo.outfeed_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.outfeed_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   outfeed_config = #vhlo.string_v1<"foo">
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.token_v1) -> !vhlo.token_v1
-   %0 = "stablehlo.outfeed"(%arg0, %arg1) {
-@@ -1605,8 +1725,9 @@
- }
- 
- // CHECK-LABEL: "op_pad"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_pad(%arg0: tensor<8xf32>, %arg1: tensor<f32>) -> tensor<16xf32> {
--  //      CHECK: "vhlo.pad_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.pad_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   edge_padding_high = #vhlo.tensor_v1<dense<4> : tensor<1xi64>>,
-   // CHECK-SAME:   edge_padding_low = #vhlo.tensor_v1<dense<4> : tensor<1xi64>>,
-   // CHECK-SAME:   interior_padding = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>
-@@ -1620,36 +1741,41 @@
- }
- 
- // CHECK-LABEL: "op_popcnt"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_popcnt(%arg0: tensor<i32>) -> tensor<i32> {
--  // CHECK: "vhlo.popcnt_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-+  // CHECK: "vhlo.popcnt_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-   %0 = "stablehlo.popcnt"(%arg0) : (tensor<i32>) -> tensor<i32>
-   func.return %0 : tensor<i32>
- }
- 
- // CHECK-LABEL: "op_power"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_power(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.power_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.power_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.power"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_real_dynamic_slice"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}}, %[[ARG3:.*]]: {{.*}})
- func.func @op_real_dynamic_slice(%arg0: tensor<?xf32>, %arg1: tensor<1xindex>, %arg2: tensor<1xindex>, %arg3: tensor<1xindex>) -> tensor<?xf32> {
--  // CHECK: "vhlo.real_dynamic_slice_v1"(%arg0, %arg1, %arg2, %arg3) : (!vhlo.tensor_v1<?x!vhlo.f32_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>) -> !vhlo.tensor_v1<?x!vhlo.f32_v1>
-+  // CHECK: "vhlo.real_dynamic_slice_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]], %[[ARG3]]) : (!vhlo.tensor_v1<?x!vhlo.f32_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>) -> !vhlo.tensor_v1<?x!vhlo.f32_v1>
-   %0 = "stablehlo.real_dynamic_slice"(%arg0, %arg1, %arg2, %arg3) : (tensor<?xf32>, tensor<1xindex>, tensor<1xindex>, tensor<1xindex>) -> tensor<?xf32>
-   func.return %0 : tensor<?xf32>
- }
- 
- // CHECK-LABEL: "op_real"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_real(%arg0: tensor<complex<f32>>) -> tensor<f32> {
--  // CHECK: "vhlo.real_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.real_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.real"(%arg0) : (tensor<complex<f32>>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_recv"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_recv(%arg0: !stablehlo.token) -> (tensor<f32>, !stablehlo.token) {
--  //      CHECK: "vhlo.recv_v1"(%arg0) <{
-+  //      CHECK: "vhlo.recv_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME:   channel_type = #vhlo.integer_v1<3 : i64>,
-   // CHECK-SAME:   is_host_transfer = #vhlo.bool_v1<true>
-@@ -1662,8 +1788,9 @@
- }
- 
- // CHECK-LABEL: "op_reduce"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_reduce(%arg0: tensor<16xf32>, %arg1: tensor<f32>) -> tensor<f32> {
--  //  CHECK: "vhlo.reduce_v1"(%arg0, %arg1)
-+  //  CHECK: "vhlo.reduce_v1"(%[[ARG0]], %[[ARG1]])
-   //  CHECK:   ^[[BB:bb.*]](%[[ARG1:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>, %[[ARG2:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
-   //  CHECK:     "vhlo.return_v1"(%[[VAL1:.*]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-   //  CHECK: }) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-@@ -1678,8 +1805,9 @@
- }
- 
- // CHECK-LABEL: "op_reduce_precision"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_reduce_precision(%arg0: tensor<f32>) -> tensor<f32> {
--  //      CHECK: "vhlo.reduce_precision_v1"(%arg0) <{
-+  //      CHECK: "vhlo.reduce_precision_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   exponent_bits = #vhlo.integer_v1<8 : i32>
-   // CHECK-SAME:   mantissa_bits = #vhlo.integer_v1<10 : i32>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-@@ -1693,7 +1821,7 @@
- // CHECK_lABEL: "op_reduce_with_promotable_types"
- func.func @op_reduce_with_promotable_types(%arg0: tensor<4x4xf32>, %arg1 : tensor<f32>)
-     -> (tensor<4xf64>) {
--  //  CHECK: "vhlo.reduce_v1"(%arg0, %arg1)
-+  //  CHECK: "vhlo.reduce_v1"(%[[ARG0:.*]], %[[ARG1:.*]])
-   //  CHECK:   ^[[BB:bb.*]](%[[ARG1:arg.*]]: !vhlo.tensor_v1<!vhlo.f64_v1>, %[[ARG2:arg.*]]: !vhlo.tensor_v1<!vhlo.f64_v1>):
-   //  CHECK:     "vhlo.return_v1"(%[[VAL1:.*]]) : (!vhlo.tensor_v1<!vhlo.f64_v1>) -> ()
-   //  CHECK: }) : (!vhlo.tensor_v1<4x4x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<4x!vhlo.f64_v1>
-@@ -1708,8 +1836,9 @@
- }
- 
- // CHECK-LABEL: "op_reduce_scatter"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_reduce_scatter(%arg0: tensor<16xf32>) -> tensor<16xf32> {
--  //               CHECK: "vhlo.reduce_scatter_v1"(%arg0) <{
-+  //               CHECK: "vhlo.reduce_scatter_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0], [1]]> : tensor<2x1xi64>>,
-   //          CHECK-SAME:   scatter_dimension = #vhlo.integer_v1<0 : i64>
-@@ -1734,7 +1863,7 @@
- 
- // CHECK_lABEL: "op_reduce_scatter_with_promotable_types"
- func.func @op_reduce_scatter_with_promotable_types(%data: tensor<4x16xf32>) -> tensor<4x4xf64> {
--  //  CHECK: "vhlo.reduce_scatter_v1"(%arg0)
-+  //  CHECK: "vhlo.reduce_scatter_v1"(%[[ARG0:.*]])
-   //  CHECK:   ^[[BB:bb.*]](%[[ARG1:arg.*]]: !vhlo.tensor_v1<!vhlo.f64_v1>, %[[ARG2:arg.*]]: !vhlo.tensor_v1<!vhlo.f64_v1>):
-   //  CHECK:     "vhlo.return_v1"(%[[VAL1:.*]]) : (!vhlo.tensor_v1<!vhlo.f64_v1>) -> ()
-   //  CHECK: }) : (!vhlo.tensor_v1<4x16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<4x4x!vhlo.f64_v1>
-@@ -1751,8 +1880,9 @@
- 
- 
- // CHECK-LABEL: "op_reduce_window"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_reduce_window(%arg0: tensor<2x17x31x7xf32>, %arg1: tensor<f32>) -> tensor<2x9x16x7xf32> {
--  //               CHECK: "vhlo.reduce_window_v1"(%arg0, %arg1) <{
-+  //               CHECK: "vhlo.reduce_window_v1"(%[[ARG0]], %[[ARG1]]) <{
-   //          CHECK-SAME:   base_dilations = #vhlo.tensor_v1<dense<[1, 2, 2, 1]> : tensor<4xi64>>,
-   // CHECK-SAME{LITERAL}:   padding = #vhlo.tensor_v1<dense<[[0, 0], [2, 0], [0, 2], [0, 0]]> : tensor<4x2xi64>>,
-   //          CHECK-SAME:   window_dilations = #vhlo.tensor_v1<dense<[1, 2, 2, 1]> : tensor<4xi64>>,
-@@ -1777,11 +1907,11 @@
-   func.return %0 : tensor<2x9x16x7xf32>
- }
- 
--// CHECK_lABEL: "op_reduce_window_with_promotable_types"
-+// CHECK-LABEL: "op_reduce_window_with_promotable_types"
- func.func @op_reduce_window_with_promotable_types(%arg0: tensor<4x2xf32>,
-     %arg1: tensor<4x2xf32>, %init0: tensor<f32>, %init1: tensor<f32>) ->
-     (tensor<2x2xf64>, tensor<2x2xf32>) {
--  //  CHECK: "vhlo.reduce_window_v1"(%arg0, %arg1, %arg2, %arg3)
-+  //  CHECK: "vhlo.reduce_window_v1"(%[[ARG0:.*]], %[[ARG1:.*]], %[[ARG2:.*]], %[[ARG3:.*]])
-   //  CHECK:   ^[[BB:bb.*]](%[[ARG1:arg.*]]: !vhlo.tensor_v1<!vhlo.f64_v1>, %[[ARG2:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>, %[[ARG3:arg.*]]: !vhlo.tensor_v1<!vhlo.f64_v1>, %[[ARG4:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
-   //  CHECK:     "vhlo.return_v1"(%[[VAL1:.*]], %[[VAL2:.*]]) : (!vhlo.tensor_v1<!vhlo.f64_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-   //  CHECK: }) : (!vhlo.tensor_v1<4x2x!vhlo.f32_v1>, !vhlo.tensor_v1<4x2x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> (!vhlo.tensor_v1<2x2x!vhlo.f64_v1>, !vhlo.tensor_v1<2x2x!vhlo.f32_v1>)
-@@ -1801,8 +1931,9 @@
- }
- 
- // CHECK-LABEL: "op_remainder"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_remainder(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.remainder_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.remainder_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.remainder"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
-@@ -1822,16 +1953,18 @@
- }
- 
- // CHECK-LABEL: "op_reshape"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_reshape(%arg0: tensor<16xf32>) -> tensor<4x4xf32> {
--  // CHECK: "vhlo.reshape_v1"(%arg0) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<4x4x!vhlo.f32_v1>
-+  // CHECK: "vhlo.reshape_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<4x4x!vhlo.f32_v1>
-   %0 = "stablehlo.reshape"(%arg0) : (tensor<16xf32>) -> tensor<4x4xf32>
-   func.return %0 : tensor<4x4xf32>
- }
- 
- // CHECK-LABEL: "op_return"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_return(%arg0: tensor<i32>, %arg1: tensor<f32>) -> tensor<f32> {
--  //      CHECK: "vhlo.case_v1"(%arg0) ({
--  // CHECK-NEXT:   "vhlo.return_v1"(%arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-+  //      CHECK: "vhlo.case_v1"(%[[ARG0]]) ({
-+  // CHECK-NEXT:   "vhlo.return_v1"(%[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-   // CHECK-NEXT: }) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.case"(%arg0) ({
-     "stablehlo.return"(%arg1) : (tensor<f32>) -> ()
-@@ -1840,8 +1973,9 @@
- }
- 
- // CHECK-LABEL: "op_reverse"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_reverse(%arg0: tensor<16xf32>) -> tensor<16xf32> {
--  //      CHECK: "vhlo.reverse_v1"(%arg0) <{
-+  //      CHECK: "vhlo.reverse_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   dimensions = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x!vhlo.f32_v1>
-   %0 = "stablehlo.reverse"(%arg0) {
-@@ -1851,8 +1985,9 @@
- }
- 
- // CHECK-LABEL: "op_rng_bit_generator"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_rng_bit_generator(%arg0: tensor<f32>) -> (tensor<f32>, tensor<f32>) {
--  //      CHECK: "vhlo.rng_bit_generator_v1"(%arg0) <{
-+  //      CHECK: "vhlo.rng_bit_generator_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   rng_algorithm = #vhlo<rng_algorithm_v1 PHILOX>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>)
-   %0:2 = "stablehlo.rng_bit_generator"(%arg0) {
-@@ -1862,8 +1997,9 @@
- }
- 
- // CHECK-LABEL: "op_rng"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_rng(%arg0: tensor<f32>, %arg1: tensor<f32>, %arg2: tensor<0xindex>) -> tensor<f32> {
--  //      CHECK: "vhlo.rng_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.rng_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   rng_distribution = #vhlo<rng_distribution_v1 NORMAL>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<0x!vhlo.index_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.rng"(%arg0, %arg1, %arg2) {
-@@ -1873,29 +2009,33 @@
- }
- 
- // CHECK-LABEL: "op_round_nearest_afz"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_round_nearest_afz(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.round_nearest_afz_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.round_nearest_afz_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.round_nearest_afz"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_round_nearest_even"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_round_nearest_even(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.round_nearest_even_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.round_nearest_even_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.round_nearest_even"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_rsqrt"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_rsqrt(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.rsqrt_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.rsqrt_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.rsqrt"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_scatter"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_scatter(%arg0: tensor<200x100x300xf32>, %arg1: tensor<10x2xi32>, %arg2: tensor<10x300xf32>) -> tensor<200x100x300xf32> {
--  //      CHECK: "vhlo.scatter_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.scatter_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   index_vector_dim = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   indices_are_sorted = #vhlo.bool_v1<true>,
-   // CHECK-SAME:   inserted_window_dims = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-@@ -1928,7 +2068,7 @@
- func.func @op_scatter_with_promotable_types(%input_tensor: tensor<200x100x300xf32>,
-     %scatter_indices: tensor<10x2xi32>, %updates: tensor<10x300xf32>) ->
-       tensor<200x100x300xf64> {
--  //  CHECK: "vhlo.scatter_v1"(%arg0, %arg1, %arg2)
-+  //  CHECK: "vhlo.scatter_v1"(%[[ARG0:.*]], %[[ARG1:.*]], %[[ARG2:.*]])
-   //  CHECK:   ^[[BB:bb.*]](%[[ARG1:arg.*]]: !vhlo.tensor_v1<!vhlo.f64_v1>, %[[ARG2:arg.*]]: !vhlo.tensor_v1<!vhlo.f64_v1>):
-   //  CHECK:     "vhlo.return_v1"(%[[VAL1:.*]]) : (!vhlo.tensor_v1<!vhlo.f64_v1>) -> ()
-   //  CHECK: }) : (!vhlo.tensor_v1<200x100x300x!vhlo.f32_v1>, !vhlo.tensor_v1<10x2x!vhlo.i32_v1>, !vhlo.tensor_v1<10x300x!vhlo.f32_v1>) -> !vhlo.tensor_v1<200x100x300x!vhlo.f64_v1>
-@@ -1951,8 +2091,9 @@
- }
- 
- // CHECK-LABEL: "op_select_and_scatter"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_select_and_scatter(%arg0: tensor<10x24x24x64xf32>, %arg1: tensor<12x13x13x66xf32>, %arg2: tensor<f32>) -> tensor<10x24x24x64xf32> {
--  //      CHECK: "vhlo.select_and_scatter_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.select_and_scatter_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   padding = #vhlo.tensor_v1<dense<1> : tensor<4x2xi64>>,
-   // CHECK-SAME:   window_dimensions = #vhlo.tensor_v1<dense<[1, 2, 2, 1]> : tensor<4xi64>>,
-   // CHECK-SAME:   window_strides = #vhlo.tensor_v1<dense<[1, 2, 2, 1]> : tensor<4xi64>>
-@@ -1982,8 +2123,9 @@
- }
- 
- // CHECK-LABEL: "op_select_and_scatter_with_promotable_types"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_select_and_scatter_with_promotable_types(%arg0: tensor<10x24x24x64xf32>, %arg1: tensor<12x13x13x66xf32>, %arg2: tensor<f32>) -> tensor<10x24x24x64xf64> {
--  // CHECK: "vhlo.select_and_scatter_v1"(%arg0, %arg1, %arg2)
-+  // CHECK: "vhlo.select_and_scatter_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]])
-   // CHECK:   ^[[BB:bb.*]](%[[ARG1:arg.*]]: !vhlo.tensor_v1<!vhlo.f64_v1>, %[[ARG2:arg.*]]: !vhlo.tensor_v1<!vhlo.f64_v1>):
-   // CHECK:     %[[VAL:.*]] = "vhlo.add_v1"(%[[ARG1]], %[[ARG2]]) : (!vhlo.tensor_v1<!vhlo.f64_v1>, !vhlo.tensor_v1<!vhlo.f64_v1>) -> !vhlo.tensor_v1<!vhlo.f64_v1>
-   // CHECK:     "vhlo.return_v1"(%[[VAL]]) : (!vhlo.tensor_v1<!vhlo.f64_v1>) -> ()
-@@ -2005,15 +2147,17 @@
- }
- 
- // CHECK-LABEL: "op_select"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_select(%arg0: tensor<i1>, %arg1: tensor<f32>, %arg2: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.select_v1"(%arg0, %arg1, %arg2) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.select_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.select"(%arg0, %arg1, %arg2) : (tensor<i1>, tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_send"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_send(%arg0: tensor<f32>, %arg1: !stablehlo.token) -> !stablehlo.token {
--  //      CHECK: "vhlo.send_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.send_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME:   channel_type = #vhlo.integer_v1<2 : i64>,
-   // CHECK-SAME:   is_host_transfer = #vhlo.bool_v1<true>
-@@ -2026,8 +2170,9 @@
- }
- 
- // CHECK-LABEL: "op_set_dimension_size"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_set_dimension_size(%arg0: tensor<?xf32>, %arg1: tensor<i32>) -> tensor<16xf32> {
--  //      CHECK: "vhlo.set_dimension_size_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.set_dimension_size_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   dimension = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<?x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<16x!vhlo.f32_v1>
-   %0 = "stablehlo.set_dimension_size"(%arg0, %arg1) {
-@@ -2037,43 +2182,49 @@
- }
- 
- // CHECK-LABEL: "op_shift_left"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_shift_left(%arg0: tensor<i32>, %arg1: tensor<i32>) -> tensor<i32> {
--  // CHECK: "vhlo.shift_left_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.i32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-+  // CHECK: "vhlo.shift_left_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-   %0 = "stablehlo.shift_left"(%arg0, %arg1) : (tensor<i32>, tensor<i32>) -> tensor<i32>
-   func.return %0 : tensor<i32>
- }
- 
- // CHECK-LABEL: "op_shift_right_arithmetic"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_shift_right_arithmetic(%arg0: tensor<i32>, %arg1: tensor<i32>) -> tensor<i32> {
--  // CHECK: "vhlo.shift_right_arithmetic_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.i32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-+  // CHECK: "vhlo.shift_right_arithmetic_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-   %0 = "stablehlo.shift_right_arithmetic"(%arg0, %arg1) : (tensor<i32>, tensor<i32>) -> tensor<i32>
-   func.return %0 : tensor<i32>
- }
- 
- // CHECK-LABEL: "op_shift_right_logical"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_shift_right_logical(%arg0: tensor<i32>, %arg1: tensor<i32>) -> tensor<i32> {
--  // CHECK: "vhlo.shift_right_logical_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.i32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-+  // CHECK: "vhlo.shift_right_logical_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-   %0 = "stablehlo.shift_right_logical"(%arg0, %arg1) : (tensor<i32>, tensor<i32>) -> tensor<i32>
-   func.return %0 : tensor<i32>
- }
- 
- // CHECK-LABEL: "op_sign"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_sign(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.sign_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.sign_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.sign"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_sine"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_sine(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.sine_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.sine_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.sine"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_slice"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_slice(%arg0: tensor<16xf32>) -> tensor<4xf32> {
--  //      CHECK: "vhlo.slice_v1"(%arg0) <{
-+  //      CHECK: "vhlo.slice_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   limit_indices = #vhlo.tensor_v1<dense<4> : tensor<1xi64>>,
-   // CHECK-SAME:   start_indices = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>,
-   // CHECK-SAME:   strides = #vhlo.tensor_v1<dense<1> : tensor<1xi64>>
-@@ -2087,8 +2238,9 @@
- }
- 
- // CHECK-LABEL: "op_sort"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_sort(%arg0: tensor<16xf32>) -> tensor<16xf32> {
--  //      CHECK: "vhlo.sort_v1"(%arg0) <{
-+  //      CHECK: "vhlo.sort_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   dimension = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME:   is_stable = #vhlo.bool_v1<true>
-   // CHECK-SAME: }> ({
-@@ -2108,29 +2260,33 @@
- }
- 
- // CHECK-LABEL: "op_sqrt"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_sqrt(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.sqrt_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.sqrt_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.sqrt"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_subtract"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_subtract(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.subtract_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.subtract_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.subtract"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_tanh"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_tanh(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.tanh_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.tanh_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.tanh"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_torch_index_select"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_torch_index_select(%arg0: tensor<5x1x5xf32>, %arg1: tensor<2xi32>) ->  tensor<2x1x5xf32> {
--  //      CHECK: "vhlo.torch_index_select_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.torch_index_select_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   batch_dims = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME:   dim = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<5x1x5x!vhlo.f32_v1>, !vhlo.tensor_v1<2x!vhlo.i32_v1>) -> !vhlo.tensor_v1<2x1x5x!vhlo.f32_v1>
-@@ -2142,8 +2298,9 @@
- }
- 
- // CHECK-LABEL: "op_trace"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_trace(%arg0: tensor<f32>) {
--  //      CHECK: "vhlo.trace_v1"(%arg0) <{
-+  //      CHECK: "vhlo.trace_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   tag = #vhlo.string_v1<"foo">
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-   "stablehlo.trace"(%arg0) {
-@@ -2153,8 +2310,9 @@
- }
- 
- // CHECK-LABEL: "op_transpose"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_transpose(%arg0: tensor<16x8xf32>) ->  tensor<8x16xf32> {
--  //      CHECK: "vhlo.transpose_v1"(%arg0) <{
-+  //      CHECK: "vhlo.transpose_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   permutation = #vhlo.tensor_v1<dense<[1, 0]> : tensor<2xi64>>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<16x8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<8x16x!vhlo.f32_v1>
-   %0 = "stablehlo.transpose"(%arg0) {
-@@ -2164,8 +2322,9 @@
- }
- 
- // CHECK-LABEL: "op_triangular_solve"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_triangular_solve(%arg0: tensor<16x16xf32>, %arg1: tensor<16x16xf32>) ->  tensor<16x16xf32> {
--  //      CHECK: "vhlo.triangular_solve_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.triangular_solve_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   left_side = #vhlo.bool_v1<true>,
-   // CHECK-SAME:   lower = #vhlo.bool_v1<true>,
-   // CHECK-SAME:   transpose_a = #vhlo<transpose_v1 NO_TRANSPOSE>,
-@@ -2181,15 +2340,17 @@
- }
- 
- // CHECK-LABEL: "op_tuple"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_tuple(%arg0: tensor<f32>) -> tuple<tensor<f32>> {
--  // CHECK: "vhlo.tuple_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tuple_v1<!vhlo.tensor_v1<!vhlo.f32_v1>>
-+  // CHECK: "vhlo.tuple_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tuple_v1<!vhlo.tensor_v1<!vhlo.f32_v1>>
-   %0 = "stablehlo.tuple"(%arg0) : (tensor<f32>) -> tuple<tensor<f32>>
-   func.return %0 : tuple<tensor<f32>>
- }
- 
- // CHECK-LABEL: "op_unary_einsum"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_unary_einsum(%arg0: tensor<8x16xf32>) -> tensor<8xf32> {
--  //      CHECK: "vhlo.unary_einsum_v1"(%arg0) <{
-+  //      CHECK: "vhlo.unary_einsum_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   einsum_config = #vhlo.string_v1<"ab->a">
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<8x16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<8x!vhlo.f32_v1>
-   %0 = "stablehlo.unary_einsum"(%arg0) {
-@@ -2199,22 +2360,25 @@
- }
- 
- // CHECK-LABEL: "op_uniform_dequantize"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_uniform_dequantize(%arg0: tensor<!quant.uniform<i8:f32, 34.0:16>>) -> tensor<f32> {
--  // CHECK: "vhlo.uniform_dequantize_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.uniform_dequantize_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.uniform_dequantize"(%arg0) : (tensor<!quant.uniform<i8:f32, 34.0:16>>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_uniform_quantize"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_uniform_quantize(%arg0: tensor<f32>) -> tensor<!quant.uniform<i8:f32, 34.0:16>> {
--  // CHECK: "vhlo.uniform_quantize_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>
-+  // CHECK: "vhlo.uniform_quantize_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>
-   %0 = "stablehlo.uniform_quantize"(%arg0) : (tensor<f32>) -> tensor<!quant.uniform<i8:f32, 34.0:16>>
-   func.return %0 : tensor<!quant.uniform<i8:f32, 34.0:16>>
- }
- 
- // CHECK-LABEL: "op_while"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_while(%arg0: tensor<i1>) -> tensor<i1> {
--  //      CHECK: "vhlo.while_v1"(%arg0) ({
-+  //      CHECK: "vhlo.while_v1"(%[[ARG0]]) ({
-   // CHECK-NEXT:   ^[[BB:bb.*]](%[[ARG1:arg.*]]: !vhlo.tensor_v1<!vhlo.bool_v1>):
-   // CHECK-NEXT:     "vhlo.return_v1"(%[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>) -> ()
-   // CHECK-NEXT:   }, {
-@@ -2232,8 +2396,9 @@
- }
- 
- // CHECK-LABEL: "op_xor"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_xor(%arg0: tensor<i1>, %arg1: tensor<i1>) -> tensor<i1> {
--  // CHECK: "vhlo.xor_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-+  // CHECK: "vhlo.xor_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-   %0 = "stablehlo.xor"(%arg0, %arg1) : (tensor<i1>, tensor<i1>) -> tensor<i1>
-   func.return %0 : tensor<i1>
- }
-@@ -2241,197 +2406,225 @@
- // ============ TYPES ============
- 
- // CHECK-LABEL: "type_i1"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_i1(%arg0: tensor<i1>, %arg1: tensor<i1>) -> tensor<i1> {
--  // CHECK: "vhlo.and_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-+  // CHECK: "vhlo.and_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-   %0 = "stablehlo.and"(%arg0, %arg1) : (tensor<i1>, tensor<i1>) -> tensor<i1>
-   func.return %0 : tensor<i1>
- }
- 
- // CHECK-LABEL: "type_i4"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_i4(%arg0: tensor<i4>, %arg1: tensor<i4>) -> tensor<i4> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.i4_v1>, !vhlo.tensor_v1<!vhlo.i4_v1>) -> !vhlo.tensor_v1<!vhlo.i4_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i4_v1>, !vhlo.tensor_v1<!vhlo.i4_v1>) -> !vhlo.tensor_v1<!vhlo.i4_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<i4>, tensor<i4>) -> tensor<i4>
-   func.return %0 : tensor<i4>
- }
- 
- // CHECK-LABEL: "type_i8"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_i8(%arg0: tensor<i8>, %arg1: tensor<i8>) -> tensor<i8> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.i8_v1>, !vhlo.tensor_v1<!vhlo.i8_v1>) -> !vhlo.tensor_v1<!vhlo.i8_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i8_v1>, !vhlo.tensor_v1<!vhlo.i8_v1>) -> !vhlo.tensor_v1<!vhlo.i8_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<i8>, tensor<i8>) -> tensor<i8>
-   func.return %0 : tensor<i8>
- }
- 
- // CHECK-LABEL: "type_i16"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_i16(%arg0: tensor<i16>, %arg1: tensor<i16>) -> tensor<i16> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.i16_v1>, !vhlo.tensor_v1<!vhlo.i16_v1>) -> !vhlo.tensor_v1<!vhlo.i16_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i16_v1>, !vhlo.tensor_v1<!vhlo.i16_v1>) -> !vhlo.tensor_v1<!vhlo.i16_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<i16>, tensor<i16>) -> tensor<i16>
-   func.return %0 : tensor<i16>
- }
- 
- // CHECK-LABEL: "type_i32"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_i32(%arg0: tensor<i32>, %arg1: tensor<i32>) -> tensor<i32> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.i32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<i32>, tensor<i32>) -> tensor<i32>
-   func.return %0 : tensor<i32>
- }
- 
- // CHECK-LABEL: "type_i64"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_i64(%arg0: tensor<i64>, %arg1: tensor<i64>) -> tensor<i64> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.i64_v1>, !vhlo.tensor_v1<!vhlo.i64_v1>) -> !vhlo.tensor_v1<!vhlo.i64_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i64_v1>, !vhlo.tensor_v1<!vhlo.i64_v1>) -> !vhlo.tensor_v1<!vhlo.i64_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<i64>, tensor<i64>) -> tensor<i64>
-   func.return %0 : tensor<i64>
- }
- 
- // CHECK-LABEL: "type_ui4"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_ui4(%arg0: tensor<ui4>, %arg1: tensor<ui4>) -> tensor<ui4> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.ui4_v1>, !vhlo.tensor_v1<!vhlo.ui4_v1>) -> !vhlo.tensor_v1<!vhlo.ui4_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.ui4_v1>, !vhlo.tensor_v1<!vhlo.ui4_v1>) -> !vhlo.tensor_v1<!vhlo.ui4_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<ui4>, tensor<ui4>) -> tensor<ui4>
-   func.return %0 : tensor<ui4>
- }
- 
- // CHECK-LABEL: "type_ui8"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_ui8(%arg0: tensor<ui8>, %arg1: tensor<ui8>) -> tensor<ui8> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.ui8_v1>, !vhlo.tensor_v1<!vhlo.ui8_v1>) -> !vhlo.tensor_v1<!vhlo.ui8_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.ui8_v1>, !vhlo.tensor_v1<!vhlo.ui8_v1>) -> !vhlo.tensor_v1<!vhlo.ui8_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<ui8>, tensor<ui8>) -> tensor<ui8>
-   func.return %0 : tensor<ui8>
- }
- 
- // CHECK-LABEL: "type_ui16"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_ui16(%arg0: tensor<ui16>, %arg1: tensor<ui16>) -> tensor<ui16> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.ui16_v1>, !vhlo.tensor_v1<!vhlo.ui16_v1>) -> !vhlo.tensor_v1<!vhlo.ui16_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.ui16_v1>, !vhlo.tensor_v1<!vhlo.ui16_v1>) -> !vhlo.tensor_v1<!vhlo.ui16_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<ui16>, tensor<ui16>) -> tensor<ui16>
-   func.return %0 : tensor<ui16>
- }
- 
- // CHECK-LABEL: "type_ui32"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_ui32(%arg0: tensor<ui32>, %arg1: tensor<ui32>) -> tensor<ui32> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.ui32_v1>, !vhlo.tensor_v1<!vhlo.ui32_v1>) -> !vhlo.tensor_v1<!vhlo.ui32_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.ui32_v1>, !vhlo.tensor_v1<!vhlo.ui32_v1>) -> !vhlo.tensor_v1<!vhlo.ui32_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<ui32>, tensor<ui32>) -> tensor<ui32>
-   func.return %0 : tensor<ui32>
- }
- 
- // CHECK-LABEL: "type_ui64"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_ui64(%arg0: tensor<ui64>, %arg1: tensor<ui64>) -> tensor<ui64> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.ui64_v1>, !vhlo.tensor_v1<!vhlo.ui64_v1>) -> !vhlo.tensor_v1<!vhlo.ui64_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.ui64_v1>, !vhlo.tensor_v1<!vhlo.ui64_v1>) -> !vhlo.tensor_v1<!vhlo.ui64_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<ui64>, tensor<ui64>) -> tensor<ui64>
-   func.return %0 : tensor<ui64>
- }
- 
- // CHECK-LABEL: "type_f8E4M3FN"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_f8E4M3FN(%arg0: tensor<f8E4M3FN>, %arg1: tensor<f8E4M3FN>) -> tensor<f8E4M3FN> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f8E4M3FN_v1>, !vhlo.tensor_v1<!vhlo.f8E4M3FN_v1>) -> !vhlo.tensor_v1<!vhlo.f8E4M3FN_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f8E4M3FN_v1>, !vhlo.tensor_v1<!vhlo.f8E4M3FN_v1>) -> !vhlo.tensor_v1<!vhlo.f8E4M3FN_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f8E4M3FN>, tensor<f8E4M3FN>) -> tensor<f8E4M3FN>
-   func.return %0 : tensor<f8E4M3FN>
- }
- 
- // CHECK-LABEL: "type_f8E5M2"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_f8E5M2(%arg0: tensor<f8E5M2>, %arg1: tensor<f8E5M2>) -> tensor<f8E5M2> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f8E5M2_v1>, !vhlo.tensor_v1<!vhlo.f8E5M2_v1>) -> !vhlo.tensor_v1<!vhlo.f8E5M2_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f8E5M2_v1>, !vhlo.tensor_v1<!vhlo.f8E5M2_v1>) -> !vhlo.tensor_v1<!vhlo.f8E5M2_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f8E5M2>, tensor<f8E5M2>) -> tensor<f8E5M2>
-   func.return %0 : tensor<f8E5M2>
- }
- 
- // CHECK-LABEL: "type_f8E4M3FNUZ"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_f8E4M3FNUZ(%arg0: tensor<f8E4M3FNUZ>, %arg1: tensor<f8E4M3FNUZ>) -> tensor<f8E4M3FNUZ> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f8E4M3FNUZ_v1>, !vhlo.tensor_v1<!vhlo.f8E4M3FNUZ_v1>) -> !vhlo.tensor_v1<!vhlo.f8E4M3FNUZ_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f8E4M3FNUZ_v1>, !vhlo.tensor_v1<!vhlo.f8E4M3FNUZ_v1>) -> !vhlo.tensor_v1<!vhlo.f8E4M3FNUZ_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f8E4M3FNUZ>, tensor<f8E4M3FNUZ>) -> tensor<f8E4M3FNUZ>
-   func.return %0 : tensor<f8E4M3FNUZ>
- }
- 
- // CHECK-LABEL: "type_f8E4M3B11FNUZ"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_f8E4M3B11FNUZ(%arg0: tensor<f8E4M3B11FNUZ>, %arg1: tensor<f8E4M3B11FNUZ>) -> tensor<f8E4M3B11FNUZ> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f8E4M3B11FNUZ_v1>, !vhlo.tensor_v1<!vhlo.f8E4M3B11FNUZ_v1>) -> !vhlo.tensor_v1<!vhlo.f8E4M3B11FNUZ_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f8E4M3B11FNUZ_v1>, !vhlo.tensor_v1<!vhlo.f8E4M3B11FNUZ_v1>) -> !vhlo.tensor_v1<!vhlo.f8E4M3B11FNUZ_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f8E4M3B11FNUZ>, tensor<f8E4M3B11FNUZ>) -> tensor<f8E4M3B11FNUZ>
-   func.return %0 : tensor<f8E4M3B11FNUZ>
- }
- 
- // CHECK-LABEL: "type_f8E5M2FNUZ"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_f8E5M2FNUZ(%arg0: tensor<f8E5M2FNUZ>, %arg1: tensor<f8E5M2FNUZ>) -> tensor<f8E5M2FNUZ> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f8E5M2FNUZ_v1>, !vhlo.tensor_v1<!vhlo.f8E5M2FNUZ_v1>) -> !vhlo.tensor_v1<!vhlo.f8E5M2FNUZ_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f8E5M2FNUZ_v1>, !vhlo.tensor_v1<!vhlo.f8E5M2FNUZ_v1>) -> !vhlo.tensor_v1<!vhlo.f8E5M2FNUZ_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f8E5M2FNUZ>, tensor<f8E5M2FNUZ>) -> tensor<f8E5M2FNUZ>
-   func.return %0 : tensor<f8E5M2FNUZ>
- }
- 
- // CHECK-LABEL: "type_bf16"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_bf16(%arg0: tensor<bf16>, %arg1: tensor<bf16>) -> tensor<bf16> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.bf16_v1>, !vhlo.tensor_v1<!vhlo.bf16_v1>) -> !vhlo.tensor_v1<!vhlo.bf16_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.bf16_v1>, !vhlo.tensor_v1<!vhlo.bf16_v1>) -> !vhlo.tensor_v1<!vhlo.bf16_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<bf16>, tensor<bf16>) -> tensor<bf16>
-   func.return %0 : tensor<bf16>
- }
- 
- // CHECK-LABEL: "type_f16"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_f16(%arg0: tensor<f16>, %arg1: tensor<f16>) -> tensor<f16> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f16_v1>, !vhlo.tensor_v1<!vhlo.f16_v1>) -> !vhlo.tensor_v1<!vhlo.f16_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f16_v1>, !vhlo.tensor_v1<!vhlo.f16_v1>) -> !vhlo.tensor_v1<!vhlo.f16_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f16>, tensor<f16>) -> tensor<f16>
-   func.return %0 : tensor<f16>
- }
- 
- // CHECK-LABEL: "type_f32"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_f32(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "type_f64"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_f64(%arg0: tensor<f64>, %arg1: tensor<f64>) -> tensor<f64> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f64_v1>, !vhlo.tensor_v1<!vhlo.f64_v1>) -> !vhlo.tensor_v1<!vhlo.f64_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f64_v1>, !vhlo.tensor_v1<!vhlo.f64_v1>) -> !vhlo.tensor_v1<!vhlo.f64_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f64>, tensor<f64>) -> tensor<f64>
-   func.return %0 : tensor<f64>
- }
- 
- // CHECK-LABEL: "type_complex_f32"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_complex_f32(%arg0: tensor<complex<f32>>, %arg1: tensor<complex<f32>>) -> tensor<complex<f32>> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>, !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>) -> !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>, !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>) -> !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<complex<f32>>, tensor<complex<f32>>) -> tensor<complex<f32>>
-   func.return %0 : tensor<complex<f32>>
- }
- 
- // CHECK-LABEL: "type_complex_f64"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_complex_f64(%arg0: tensor<complex<f64>>, %arg1: tensor<complex<f64>>) -> tensor<complex<f64>> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f64_v1>>, !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f64_v1>>) -> !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f64_v1>>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f64_v1>>, !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f64_v1>>) -> !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f64_v1>>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<complex<f64>>, tensor<complex<f64>>) -> tensor<complex<f64>>
-   func.return %0 : tensor<complex<f64>>
- }
- 
- // CHECK-LABEL: "type_dynamism_ranked"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @type_dynamism_ranked(%arg0: tensor<?xf32>) -> tensor<?xf32> {
--  // CHECK: "vhlo.abs_v1"(%arg0) : (!vhlo.tensor_v1<?x!vhlo.f32_v1>) -> !vhlo.tensor_v1<?x!vhlo.f32_v1>
-+  // CHECK: "vhlo.abs_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<?x!vhlo.f32_v1>) -> !vhlo.tensor_v1<?x!vhlo.f32_v1>
-   %0 = "stablehlo.abs"(%arg0) : (tensor<?xf32>) -> tensor<?xf32>
-   func.return %0 : tensor<?xf32>
- }
- 
- // CHECK-LABEL: "type_per_tensor_quantization"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_per_tensor_quantization(%arg0: tensor<!quant.uniform<i8:f32, 34.0:16>>, %arg1: tensor<!quant.uniform<i8:f32, 34.0:16>>) -> tensor<!quant.uniform<i8:f32, 34.0:16>> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>, !vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>) -> !vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>, !vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>) -> !vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<!quant.uniform<i8:f32, 34.0:16>>, tensor<!quant.uniform<i8:f32, 34.0:16>>) -> tensor<!quant.uniform<i8:f32, 34.0:16>>
-   func.return %0 : tensor<!quant.uniform<i8:f32, 34.0:16>>
- }
- 
- // CHECK-LABEL: "type_per_axis_quantization"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @type_per_axis_quantization(%arg0: tensor<2x!quant.uniform<i8:f32:0, {34.0:16, 34.0:16}>>) -> tensor<2x!quant.uniform<i8:f32:0, {34.0:16, 34.0:16}>> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg0) : (!vhlo.tensor_v1<2x!vhlo.quant_per_axis_v1<!vhlo.i8_v1:!vhlo.f32_v1, 0, 3.400000e+01, 3.400000e+01, 16, 16, -128:127, 1>>, !vhlo.tensor_v1<2x!vhlo.quant_per_axis_v1<!vhlo.i8_v1:!vhlo.f32_v1, 0, 3.400000e+01, 3.400000e+01, 16, 16, -128:127, 1>>) -> !vhlo.tensor_v1<2x!vhlo.quant_per_axis_v1<!vhlo.i8_v1:!vhlo.f32_v1, 0, 3.400000e+01, 3.400000e+01, 16, 16, -128:127, 1>>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG0]]) : (!vhlo.tensor_v1<2x!vhlo.quant_per_axis_v1<!vhlo.i8_v1:!vhlo.f32_v1, 0, 3.400000e+01, 3.400000e+01, 16, 16, -128:127, 1>>, !vhlo.tensor_v1<2x!vhlo.quant_per_axis_v1<!vhlo.i8_v1:!vhlo.f32_v1, 0, 3.400000e+01, 3.400000e+01, 16, 16, -128:127, 1>>) -> !vhlo.tensor_v1<2x!vhlo.quant_per_axis_v1<!vhlo.i8_v1:!vhlo.f32_v1, 0, 3.400000e+01, 3.400000e+01, 16, 16, -128:127, 1>>
-   %0 = stablehlo.add %arg0, %arg0 : tensor<2x!quant.uniform<i8:f32:0, {34.0:16, 34.0:16}>>
-   func.return %0 : tensor<2x!quant.uniform<i8:f32:0, {34.0:16, 34.0:16}>>
- }
- 
- //       CHECK: function_type = #vhlo.type_v1<!vhlo.func_v1<(!vhlo.token_v1) -> !vhlo.token_v1>>
- // CHECK-LABEL: "type_token_callee"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @type_token_callee(%arg0: !stablehlo.token) -> !stablehlo.token {
--  // CHECK: "vhlo.return_v1"(%arg0) : (!vhlo.token_v1) -> ()
-+  // CHECK: "vhlo.return_v1"(%[[ARG0]]) : (!vhlo.token_v1) -> ()
-   return %arg0 : !stablehlo.token
- }
- 
- //       CHECK: function_type = #vhlo.type_v1<!vhlo.func_v1<(!vhlo.token_v1) -> !vhlo.token_v1>>
- // CHECK-LABEL: "type_token_caller"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @type_token_caller(%arg0: !stablehlo.token) -> !stablehlo.token {
--  // CHECK:      "vhlo.call_v1"(%arg0) <{callee = #vhlo.string_v1<"type_token_callee">}
-+  // CHECK:      "vhlo.call_v1"(%[[ARG0]]) <{callee = #vhlo.string_v1<"type_token_callee">}
-   // CHECK-SAME: (!vhlo.token_v1) -> !vhlo.token_v1
-   %0 = func.call @type_token_callee(%arg0) : (!stablehlo.token) -> !stablehlo.token
-   return %0 : !stablehlo.token
- }
- 
- // CHECK-LABEL: "type_tuple"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @type_tuple(%arg0: tuple<tensor<f32>>) -> tuple<!stablehlo.token> {
-   %0 = "stablehlo.custom_call"(%arg0) {
-     call_target_name = "foo"
 
diff --git a/third_party/stablehlo/workspace.bzl b/third_party/stablehlo/workspace.bzl
index 6a72c8fa16885c..aaef166d96583c 100644
--- a/third_party/stablehlo/workspace.bzl
+++ b/third_party/stablehlo/workspace.bzl
@@ -4,8 +4,8 @@ load("//third_party:repo.bzl", "tf_http_archive", "tf_mirror_urls")
 
 def repo():
     # LINT.IfChange
-    STABLEHLO_COMMIT = "8ba7728d3fdc3ea882e893ee7e53255c95ee0e5a"
-    STABLEHLO_SHA256 = "1dfc7179dc9200c3ab4ea85edbac4a35393866d8cd8694fcaac00c1d27036408"
+    STABLEHLO_COMMIT = "797bee217e1a041e9aac22cad4db207274596d94"
+    STABLEHLO_SHA256 = "e5619033e131ea2eeb9eab8c8e362f3ba12e111c6b4a15dac789ca216ff22c58"
     # LINT.ThenChange(Google-internal path)
 
     tf_http_archive(
diff --git a/third_party/xla/third_party/stablehlo/temporary.patch b/third_party/xla/third_party/stablehlo/temporary.patch
index f906a856ae69e1..abb329aec579e4 100755
--- a/third_party/xla/third_party/stablehlo/temporary.patch
+++ b/third_party/xla/third_party/stablehlo/temporary.patch
@@ -164,81 +164,6 @@ diff --ruN a/stablehlo/CMakeLists.txt b/stablehlo/CMakeLists.txt
  
  #-------------------------------------------------------------------------------
  # Directory setup
-diff --ruN a/stablehlo/docs/spec.md b/stablehlo/docs/spec.md
---- stablehlo/docs/spec.md
-+++ stablehlo/docs/spec.md
-@@ -2532,10 +2532,10 @@
-         rhs_batching_dimensions, lhs_contracting_dimensions,
-         rhs_contracting_dimensions, precision_config), lhs, rhs, type(result))`.
- 
--For hybrid quantized types, performs `hybrid_dequantize_then_op(
--    lambda lhs, rhs: dot_general(lhs, rhs, lhs_batching_dimensions,
--        rhs_batching_dimensions, lhs_contracting_dimensions,
--        rhs_contracting_dimensions, precision_config), lhs, rhs)`.
-+This only specifies semantics for per-tensor quantization. Per-axis quantization
-+is work in progress ([#1574](https://github.com/openxla/stablehlo/issues/1574)).
-+Also, in the future we may consider adding support for hybrid quantization
-+ ([#1575](https://github.com/openxla/stablehlo/issues/1575)).
- 
- `precision_config` controls the tradeoff between speed and accuracy for
- computations on accelerator backends. This can be one of the following (at the
-@@ -2552,21 +2552,21 @@
- 
- #### Inputs
- 
--| Label | Name                         | Type                                                         | Constraints                                    |
--|-------|------------------------------|--------------------------------------------------------------|------------------------------------------------|
--| (I1)  | `lhs`                        | tensor or per-tensor quantized tensor                        | (C5-C6), (C9-C10), (C12-C14), (C17-C18), (C20) |
--| (I2)  | `rhs`                        | tensor or quantized tensor                                   | (C7-C10), (C12-C20)                            |
--| (I3)  | `lhs_batching_dimensions`    | 1-dimensional tensor constant of type `si64`                 | (C1), (C3), (C5), (C9), (C12)                  |
--| (I4)  | `rhs_batching_dimensions`    | 1-dimensional tensor constant of type `si64`                 | (C1), (C4), (C7), (C9)                         |
--| (I5)  | `lhs_contracting_dimensions` | 1-dimensional tensor constant of type `si64`                 | (C2), (C3), (C6), (C10)                        |
--| (I6)  | `rhs_contracting_dimensions` | 1-dimensional tensor constant of type `si64`                 | (C2), (C4), (C8), (C10), (C16)                 |
--| (I7)  | `precision_config`           | variadic number of enums of `DEFAULT`, `HIGH`, and `HIGHEST` | (C11)                                          |
-+| Label | Name                         | Type                                                         | Constraints                    |
-+|-------|------------------------------|--------------------------------------------------------------|--------------------------------|
-+| (I1)  | `lhs`                        | tensor or per-tensor quantized tensor                        | (C5-C6), (C9-C10), (C12-C16)   |
-+| (I2)  | `rhs`                        | tensor or quantized tensor                                   | (C7-C10), (C12), (C18-C19)     |
-+| (I3)  | `lhs_batching_dimensions`    | 1-dimensional tensor constant of type `si64`                 | (C1), (C3), (C5), (C9), (C12)  |
-+| (I4)  | `rhs_batching_dimensions`    | 1-dimensional tensor constant of type `si64`                 | (C1), (C4), (C7), (C9)         |
-+| (I5)  | `lhs_contracting_dimensions` | 1-dimensional tensor constant of type `si64`                 | (C2), (C3), (C6), (C10)        |
-+| (I6)  | `rhs_contracting_dimensions` | 1-dimensional tensor constant of type `si64`                 | (C2), (C4), (C8), (C10), (C19) |
-+| (I7)  | `precision_config`           | variadic number of enums of `DEFAULT`, `HIGH`, and `HIGHEST` | (C11)                          |
- 
- #### Outputs
- 
- | Name     | Type                       | Constraints                |
- |----------|----------------------------|----------------------------|
--| `result` | tensor or quantized tensor | (C12), (C14), (C18-C20)    |
-+| `result` | tensor or quantized tensor | (C12), (C14), (C16), (C18) |
- 
- #### Constraints
- 
-@@ -2589,17 +2589,14 @@
- * If the operation uses non-quantized tensors:
-   * (C13) `element_type(lhs) = element_type(rhs)`.
- * If the operation uses quantized tensors:
--  * (C14) `is_quantized(lhs) = is_quantized(result) and is_quantized(rhs)`.
--  * (C15) `zero_points(rhs) = 0`.
--  * (C16) If `is_per_axis_quantized(rhs)`, then
-+  * (C14) `is_quantized(lhs) and is_quantized(rhs) and is_quantized(result)`.
-+  * (C15) `storage_type(lhs) = storage_type(rhs)`.
-+  * (C16) `expressed_type(lhs) = expressed_type(rhs) = expressed_type(result)`.
-+  * (C17) `zero_points(rhs) = 0`.
-+  * (C18) If `is_per_tensor_quantized(rhs)`, then
-+    `is_per_tensor_quantized(result)`.
-+  * (C19) If `is_per_axis_quantized(rhs)`, then
-     `quantization_dimension(rhs)` not in `rhs_contracting_dimensions`.
--  * If `is_quantized(lhs)`:
--    * (C17) `storage_type(lhs) = storage_type(rhs)`.
--    * (C18) `expressed_type(lhs) = expressed_type(rhs) = expressed_type(result)`.
--    * (C19) If `is_per_tensor_quantized(rhs)`, then
--      `is_per_tensor_quantized(result)`.
--  * If `!is_quantized(lhs)`:
--    * (C20) `element_type(lhs) = expressed_type(rhs) = element_type(result)`.
- 
- #### Examples
- 
 diff --ruN a/stablehlo/stablehlo/CMakeLists.txt b/stablehlo/stablehlo/CMakeLists.txt
 --- stablehlo/stablehlo/CMakeLists.txt
 +++ stablehlo/stablehlo/CMakeLists.txt
@@ -250,118 +175,10 @@ diff --ruN a/stablehlo/stablehlo/CMakeLists.txt b/stablehlo/stablehlo/CMakeLists
  add_subdirectory(integrations)
  add_subdirectory(reference)
  add_subdirectory(tests)
-diff --ruN a/stablehlo/stablehlo/conversions/linalg/tests/convolution.mlir b/stablehlo/stablehlo/conversions/linalg/tests/convolution.mlir
---- stablehlo/stablehlo/conversions/linalg/tests/convolution.mlir
-+++ stablehlo/stablehlo/conversions/linalg/tests/convolution.mlir
-@@ -356,7 +356,7 @@
- }
- // CHECK-DAG:       %[[CST:.+]] = arith.constant 0.000000e+00 : f32
- // CHECK:       %[[COLLAPSE:.+]] = tensor.collapse_shape %[[FILTER]] {{\[}}[0, 1, 2, 3]] : tensor<2x2x1x6xf32> into tensor<24xf32>
--// CHECK:       %[[EXPAND:.+]] = tensor.expand_shape %[[COLLAPSE]] {{\[}}[0, 1, 2, 3]] : tensor<24xf32> into tensor<2x2x2x3xf32>
-+// CHECK:       %[[EXPAND:.+]] = tensor.expand_shape %[[COLLAPSE]] {{\[}}[0, 1, 2, 3]] output_shape [2, 2, 2, 3] : tensor<24xf32> into tensor<2x2x2x3xf32>
- // CHECK:       %[[INIT:.+]] = tensor.empty() : tensor<2x3x4x2x3xf32>
- // CHECK:       %[[FILL:.+]] = linalg.fill ins(%[[CST]] : f32) outs(%[[INIT]] : tensor<2x3x4x2x3xf32>) -> tensor<2x3x4x2x3xf32>
- // CHECK:       %[[OUT:.+]] = linalg.depthwise_conv_2d_nhwc_hwcm
-diff --ruN a/stablehlo/stablehlo/conversions/linalg/tests/miscellaneous.mlir b/stablehlo/stablehlo/conversions/linalg/tests/miscellaneous.mlir
---- stablehlo/stablehlo/conversions/linalg/tests/miscellaneous.mlir
-+++ stablehlo/stablehlo/conversions/linalg/tests/miscellaneous.mlir
-@@ -865,7 +865,7 @@
-   %0 = "stablehlo.reshape"(%arg0) : (tensor<i32>) -> tensor<1xi32>
-   func.return %0 : tensor<1xi32>
- }
--// CHECK: tensor.expand_shape %{{.*}} [] : tensor<i32> into tensor<1xi32>
-+// CHECK: tensor.expand_shape %{{.*}} [] output_shape [1] : tensor<i32> into tensor<1xi32>
- 
- // -----
- 
-@@ -876,7 +876,7 @@
-   func.return %0 : tensor<1xui32>
- }
- // CHECK:         %[[ARG_SIGNLESS:.*]] = builtin.unrealized_conversion_cast %[[ARG_UNSIGNED]] : tensor<ui32> to tensor<i32>
--// CHECK:         %[[RET_SIGNLESS:.*]] = tensor.expand_shape %[[ARG_SIGNLESS]] [] : tensor<i32> into tensor<1xi32>
-+// CHECK:         %[[RET_SIGNLESS:.*]] = tensor.expand_shape %[[ARG_SIGNLESS]] [] output_shape [1] : tensor<i32> into tensor<1xi32>
- // CHECK:         %[[RET_UNSIGNED:.*]] = builtin.unrealized_conversion_cast %[[RET_SIGNLESS]] : tensor<1xi32> to tensor<1xui32>
- // CHECK:         return %[[RET_UNSIGNED]] : tensor<1xui32>
- 
-@@ -978,7 +978,7 @@
- }
- // CHECK: %[[FLATTEN:.*]] = tensor.collapse_shape %{{.*}} {{\[}}[0, 1]] : tensor<?x?xf32> into tensor<?xf32>
- // CHECK: %[[CAST:.*]] = tensor.cast %[[FLATTEN]] : tensor<?xf32> to tensor<40xf32>
--// CHECK: tensor.expand_shape %[[CAST]] {{\[}}[0, 1, 2]] : tensor<40xf32> into tensor<2x4x5xf32>
-+// CHECK: tensor.expand_shape %[[CAST]] {{\[}}[0, 1, 2]] output_shape [2, 4, 5] : tensor<40xf32> into tensor<2x4x5xf32>
- 
- // -----
- 
-@@ -988,7 +988,7 @@
-   func.return %0 : tensor<1x3xi32>
- }
- // CHECK: %[[CAST:.*]] = tensor.cast %{{.*}} : tensor<?xi32> to tensor<3xi32>
--// CHECK: tensor.expand_shape %[[CAST]] {{\[}}[0, 1]] : tensor<3xi32> into tensor<1x3xi32>
-+// CHECK: tensor.expand_shape %[[CAST]] {{\[}}[0, 1]] output_shape [1, 3] : tensor<3xi32> into tensor<1x3xi32>
- 
- // -----
- 
-diff --ruN a/stablehlo/stablehlo/conversions/linalg/tests/random.mlir b/stablehlo/stablehlo/conversions/linalg/tests/random.mlir
---- stablehlo/stablehlo/conversions/linalg/tests/random.mlir
-+++ stablehlo/stablehlo/conversions/linalg/tests/random.mlir
-@@ -480,8 +480,8 @@
- // CHECK-DAG:   %[[VAL_101:.*]] = arith.xori %[[VAL_100]], %[[VAL_87]] : i32
- 
- // CHECK: linalg.yield %[[YIELDED_1:.*]], %[[YIELDED_2:.*]] : i64, i64
--// CHECK-DAG: %[[VAL_206:.*]] = tensor.expand_shape %[[VAL_207:.*]]#0 {{\[\[}}0, 1]] : tensor<4xi64> into tensor<4x1xi64>
--// CHECK-DAG: %[[VAL_208:.*]] = tensor.expand_shape %[[VAL_207]]#1 {{\[\[}}0, 1]] : tensor<4xi64> into tensor<4x1xi64>
-+// CHECK-DAG: %[[VAL_206:.*]] = tensor.expand_shape %[[VAL_207:.*]]#0 {{\[\[}}0, 1]]
-+// CHECK-DAG: %[[VAL_208:.*]] = tensor.expand_shape %[[VAL_207]]#1 {{\[\[}}0, 1]]
- // CHECK-DAG: %[[VAL_209:.*]] = tensor.empty() : tensor<4x2xi64>
- // CHECK-DAG: %[[VAL_213:.*]] = tensor.insert %[[VAL_30]] into %[[VAL_0]]{{\[}}%[[VAL_19]]] : tensor<2xi64>
- 
-@@ -575,10 +575,10 @@
- // CHECK: %[[COLLAPSE:.+]] = tensor.collapse_shape %[[CONCAT]]
- 
- 
--// CHECK: %[[VAL_213:.*]] = tensor.expand_shape %[[COLLAPSE]] {{\[\[}}0, 1]] : tensor<80xi32> into tensor<80x1xi32>
-+// CHECK: %[[VAL_213:.*]] = tensor.expand_shape %[[COLLAPSE]] {{\[\[}}0, 1]]
- // CHECK: %[[VAL_214:.*]] = tensor.extract_slice %[[VAL_213]][0, 0] [77, 1] [1, 1] : tensor<80x1xi32> to tensor<77x1xi32>
- // CHECK: %[[VAL_215:.*]] = tensor.collapse_shape %[[VAL_214]] {{\[\[}}0, 1]] : tensor<77x1xi32> into tensor<77xi32>
--// CHECK: %[[VAL_216:.*]] = tensor.expand_shape %[[VAL_215]] {{\[\[}}0, 1]] : tensor<77xi32> into tensor<7x11xi32>
-+// CHECK: %[[VAL_216:.*]] = tensor.expand_shape %[[VAL_215]] {{\[\[}}0, 1]]
- // CHECK: %[[VAL_217:.*]] = tensor.insert %[[NEWSTATE]] into %[[ARG0]]{{\[}}%[[C1]]] : tensor<2xi64>
- // CHECK: return %[[VAL_217]], %[[VAL_216]] : tensor<2xi64>, tensor<7x11xi32>
- 
-@@ -616,10 +616,10 @@
- // CHECK-DAG: %[[COLLAPSE:.+]] = tensor.collapse_shape %[[CONCAT]] {{\[\[}}0, 1]] : tensor<8x2xi64> into tensor<16xi64>
- 
- 
--// CHECK-DAG: %[[EXPANDED:.*]] = tensor.expand_shape %[[COLLAPSE]] {{\[\[}}0, 1]] : tensor<16xi64> into tensor<16x1xi64>
-+// CHECK-DAG: %[[EXPANDED:.*]] = tensor.expand_shape %[[COLLAPSE]] {{\[\[}}0, 1]]
- // CHECK-DAG: %[[SLICE:.*]] = tensor.extract_slice %[[EXPANDED]][0, 0] [15, 1] [1, 1] : tensor<16x1xi64> to tensor<15x1xi64>
- // CHECK-DAG: %[[EXPAND_2:.*]] = tensor.collapse_shape %[[SLICE]] {{\[\[}}0, 1]] : tensor<15x1xi64> into tensor<15xi64>
--// CHECK-DAG: %[[RESHAPE:.*]] = tensor.expand_shape %[[EXPAND_2]] {{\[\[}}0, 1]] : tensor<15xi64> into tensor<3x5xi64>
-+// CHECK-DAG: %[[RESHAPE:.*]] = tensor.expand_shape %[[EXPAND_2]] {{\[\[}}0, 1]]
- // CHECK-DAG: %[[INSERTED:.+]] = tensor.insert %[[NEWSTATE]] into %[[ARG0]][%[[C1]]] : tensor<2xi64>
- // CHECK: return %[[INSERTED]], %[[RESHAPE]]
- 
 diff --ruN a/stablehlo/stablehlo/conversions/tosa/tests/binary.mlir b/stablehlo/stablehlo/conversions/tosa/tests/binary.mlir
 --- stablehlo/stablehlo/conversions/tosa/tests/binary.mlir
 +++ stablehlo/stablehlo/conversions/tosa/tests/binary.mlir
-@@ -45,14 +45,14 @@
- 
- // CHECK-LABEL: @divide
- func.func @divide(%arg0 : tensor<10xi32>, %arg1 : tensor<10xi32>) -> tensor<10xi32> {
--  // CHECK: tosa.div
-+  // CHECK: tosa.int_div
-   %0 = "stablehlo.divide"(%arg0, %arg1) : (tensor<10xi32>, tensor<10xi32>) -> tensor<10xi32>
-   return %0 : tensor<10xi32>
- }
- 
- // CHECK-LABEL: @divide_f32
- func.func @divide_f32(%arg0 : tensor<10xf32>, %arg1 : tensor<10xf32>) -> tensor<10xf32> {
--  // tosa.div only supports i32, so this should not legalize.
-+  // tosa.int_div only supports i32, so this should not legalize.
-   // CHECK: stablehlo.divide
-   %0 = "stablehlo.divide"(%arg0, %arg1) : (tensor<10xf32>, tensor<10xf32>) -> tensor<10xf32>
-   return %0 : tensor<10xf32>
-@@ -123,7 +123,7 @@
+@@ -155,7 +155,7 @@
  
  // CHECK-LABEL: @maximum_f64
  func.func @maximum_f64(%arg0 : tensor<10xf64>, %arg1 : tensor<10xf64>) -> tensor<10xf64> {
@@ -383,162 +200,6 @@ diff --ruN a/stablehlo/stablehlo/conversions/tosa/tests/nullary.mlir b/stablehlo
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<10xf64>
    return %0 : tensor<10xf64>
  }
-diff --ruN a/stablehlo/stablehlo/conversions/tosa/transforms/StablehloLegalizeToTosa.pdll b/stablehlo/stablehlo/conversions/tosa/transforms/StablehloLegalizeToTosa.pdll
---- stablehlo/stablehlo/conversions/tosa/transforms/StablehloLegalizeToTosa.pdll
-+++ stablehlo/stablehlo/conversions/tosa/transforms/StablehloLegalizeToTosa.pdll
-@@ -125,7 +125,7 @@
- Pattern =>
-   replace op<stablehlo.divide>(input0 : Value<_: Tosa_Int32Tensor>,
-                           input1 : Value<_: Tosa_Int32Tensor>)
--     with op<tosa.div>(input0, input1);
-+     with op<tosa.int_div>(input0, input1);
- Pattern =>
-   replace op<stablehlo.maximum>(input0 : Value<_: Tosa_Tensor>,
-                            input1 : Value<_: Tosa_Tensor>)
-diff --ruN a/stablehlo/stablehlo/dialect/Base.cpp b/stablehlo/stablehlo/dialect/Base.cpp
---- stablehlo/stablehlo/dialect/Base.cpp
-+++ stablehlo/stablehlo/dialect/Base.cpp
-@@ -651,14 +651,14 @@
- 
-   // quantized_type_c5
-   auto maxPosFiniteNum =
--      APFloat::getLargest(quantizedElementType.getExpressedType()
--                              .cast<FloatType>()
--                              .getFloatSemantics())
-+      APFloat::getLargest(
-+          cast<FloatType>(quantizedElementType.getExpressedType())
-+              .getFloatSemantics())
-           .convertToDouble();
-   auto minPosFiniteNum =
--      APFloat::getSmallest(quantizedElementType.getExpressedType()
--                               .cast<FloatType>()
--                               .getFloatSemantics())
-+      APFloat::getSmallest(
-+          cast<FloatType>(quantizedElementType.getExpressedType())
-+              .getFloatSemantics())
-           .convertToDouble();
-   if (llvm::any_of(scales, [&](double scale) {
-         return scale < minPosFiniteNum || scale > maxPosFiniteNum;
-diff --ruN a/stablehlo/stablehlo/dialect/TypeInference.cpp b/stablehlo/stablehlo/dialect/TypeInference.cpp
---- stablehlo/stablehlo/dialect/TypeInference.cpp
-+++ stablehlo/stablehlo/dialect/TypeInference.cpp
-@@ -171,7 +171,7 @@
- LogicalResult verifyConvolutionDotGeneralCommonQuantizationConstraints(
-     std::optional<Location> location, Type lhsElementType, Type rhsElementType,
-     Type resultElementType) {
--  // convolution_c28 and dot_general_c14
-+  // convolution_c28
-   if (!isa<quant::QuantizedType>(rhsElementType) ||
-       (isa<quant::QuantizedType>(lhsElementType) !=
-        isa<quant::QuantizedType>(resultElementType))) {
-@@ -184,19 +184,19 @@
-   auto rhsQuantType = cast<quant::QuantizedType>(rhsElementType);
-   if (auto lhsQuantType = dyn_cast<quant::QuantizedType>(lhsElementType)) {
-     auto resultQuantType = cast<quant::QuantizedType>(resultElementType);
--    // convolution_c31 and dot_general_c17
-+    // convolution_c31
-     if (lhsQuantType.getStorageType() != rhsQuantType.getStorageType()) {
-       return emitOptionalError(
-           location, "mismatched lhs and rhs quantization storage types");
-     }
--    // convolution_c32 and dot_general_c18
-+    // convolution_c32
-     if (lhsQuantType.getExpressedType() != rhsQuantType.getExpressedType() ||
-         lhsQuantType.getExpressedType() != resultQuantType.getExpressedType()) {
-       return emitOptionalError(
-           location,
-           "mismatched lhs, rhs and result quantization expressed types");
-     }
--    // convolution_c33 and dot_general_c19
-+    // convolution_c33
-     if (isa<quant::UniformQuantizedType>(rhsQuantType) &&
-         !isa<quant::UniformQuantizedType>(resultQuantType)) {
-       return emitOptionalError(
-@@ -204,7 +204,7 @@
-     }
-   } else {
-     Type rhsExpressedType = rhsQuantType.getExpressedType();
--    // convolution_c34 and dot_general_c20
-+    // convolution_c34
-     if (lhsElementType != rhsExpressedType ||
-         lhsElementType != resultElementType) {
-       return emitOptionalError(location,
-@@ -3559,7 +3559,7 @@
-     }
-   }
- 
--  // convolution_c28, convolution_c31 - convolution_c34
-+  // convolution_c31 - convolution_c34
-   return verifyConvolutionDotGeneralCommonQuantizationConstraints(
-       location, lhsElementType, rhsElementType, resultElementType);
- }
-@@ -3626,41 +3626,6 @@
-   return success();
- }
- 
--LogicalResult verifyDotGeneralOpQuantizationConstraints(
--    std::optional<Location> location, Type lhsType, Type rhsType,
--    Type resultType, ArrayRef<int64_t> rhsContractingDimensions) {
--  Type lhsElementType = getElementTypeOrSelf(lhsType);
--  Type rhsElementType = getElementTypeOrSelf(rhsType);
--  Type resultElementType = getElementTypeOrSelf(resultType);
--
--  // dot_general_c15
--  if (auto rhsPerTensorQuantType =
--          dyn_cast<quant::UniformQuantizedType>(rhsElementType)) {
--    if (rhsPerTensorQuantType.getZeroPoint() != 0) {
--      return emitOptionalError(location, "Zero point of rhs should be 0");
--    }
--  } else if (auto rhsPerAxisQuantType =
--                 dyn_cast<quant::UniformQuantizedPerAxisType>(rhsElementType)) {
--    if (llvm::any_of(rhsPerAxisQuantType.getZeroPoints(),
--                     [](int64_t zero_point) { return zero_point != 0; })) {
--      return emitOptionalError(location, "Zero points of rhs should be 0");
--    }
--
--    // dot_general_c16
--    if (llvm::is_contained(rhsContractingDimensions,
--                           rhsPerAxisQuantType.getQuantizedDimension())) {
--      return emitOptionalError(
--          location,
--          "Quantization dimension of rhs should not be in the "
--          "contracting dimension of rhs");
--    }
--  }
--
--  // dot_general_c14, dot_general_c17 - dot_general_c20
--  return verifyConvolutionDotGeneralCommonQuantizationConstraints(
--      location, lhsElementType, rhsElementType, resultElementType);
--}
--
- LogicalResult verifyDotGeneralOp(std::optional<Location> location, Value lhs,
-                                  Value rhs,
-                                  ArrayRef<int64_t> lhsBatchingDimensions,
-@@ -3683,13 +3648,6 @@
-     return emitOptionalError(
-         location, "inferred shape '", dimSizesToString(inferredShape.getDims()),
-         "' ", "is incompatible with return type of operation ", resultType, "");
--
--  Type lhsType = lhs.getType();
--  Type rhsType = rhs.getType();
--  if (anyQuantized<quant::QuantizedType>({lhsType, rhsType, resultType})) {
--    return verifyDotGeneralOpQuantizationConstraints(
--        location, lhsType, rhsType, resultType, rhsContractingDimensions);
--  }
-   return success();
- }
- 
-@@ -3861,8 +3819,8 @@
-   if (SmallVector<int64_t> shape; operandType.hasStaticShape() &&
-                                   matchInts(outputShape, shape).succeeded()) {
-     int64_t operandCount = operandType.getNumElements();
--    int64_t shapeCount = std::accumulate(shape.begin(), shape.end(), 1,
--                                         std::multiplies<int64_t>());
-+    int64_t shapeCount = std::accumulate(
-+        shape.begin(), shape.end(), int64_t{1}, std::multiplies<int64_t>());
-     if (operandCount != shapeCount) {
-       return emitOptionalError(location,
-                                "output_shape is incompatible with input type "
 diff --ruN a/stablehlo/stablehlo/experimental/BUILD.bazel b/stablehlo/stablehlo/experimental/BUILD.bazel
 --- stablehlo/stablehlo/experimental/BUILD.bazel
 +++ stablehlo/stablehlo/experimental/BUILD.bazel
@@ -2928,23052 +2589,4 @@ diff --ruN a/stablehlo/stablehlo/experimental/transforms/StablehloRefineShapes.c
 +}  // namespace experimental
 +}  // namespace stablehlo
 +}  // namespace mlir
-diff --ruN a/stablehlo/stablehlo/integrations/python/tests/stablehlo.py b/stablehlo/stablehlo/integrations/python/tests/stablehlo.py
---- stablehlo/stablehlo/integrations/python/tests/stablehlo.py
-+++ stablehlo/stablehlo/integrations/python/tests/stablehlo.py
-@@ -241,18 +241,18 @@
-   # Formatted as (tensor_type, np_value)
-   # Program runs arg + arg, which is used for expected value
-   tests = [
--    # No numpy types for f8 - skipping fp8 tests
--    ("f16", np.asarray(1, np.float16)),
--    ("f32", np.asarray(2, np.float32)),
--    ("f64", np.asarray(3, np.double)),
--    ("1xi8", np.asarray([4], np.int8)),
--    ("1xi16", np.asarray([5], np.int16)),
--    ("1xi32", np.asarray([-6], np.int32)),
--    # Numpy's uint treated as int by DenseElementsAttr, skipping np.uint tests
--    ("2x2xf16", np.asarray([1, 2, 3, 4], np.float16).reshape(2,2)),
--    ("2x1x2xf16", np.asarray([1, 2, 3, 4], np.float16).reshape(2,1,2)),
--    ("?x?xf16", np.asarray([1, 2, 3, 4], np.float16).reshape(2,2)),
--    ("?x2xf16", np.asarray([1, 2, 3, 4], np.float16).reshape(2,2)),
-+      # No numpy types for f8 - skipping fp8 tests
-+      ("f16", np.asarray(1, np.float16)),
-+      ("f32", np.asarray(2, np.float32)),
-+      ("f64", np.asarray(3, np.double)),
-+      ("1xi8", np.asarray([4], np.int8)),
-+      ("1xi16", np.asarray([5], np.int16)),
-+      ("1xi32", np.asarray([-6], np.int32)),
-+      # Numpy's uint treated as int by DenseElementsAttr, skipping np.uint tests
-+      ("2x2xf16", np.asarray([1, 2, 3, 4], np.float16).reshape(2, 2)),
-+      ("2x1x2xf16", np.asarray([1, 2, 3, 4], np.float16).reshape(2, 1, 2)),
-+      ("?x?xf16", np.asarray([1, 2, 3, 4], np.float16).reshape(2, 2)),
-+      ("?x2xf16", np.asarray([1, 2, 3, 4], np.float16).reshape(2, 2)),
-   ]
-   for test in tests:
-     tensor_type, arg = test
-diff --ruN a/stablehlo/stablehlo/tests/ops_stablehlo_quantized.mlir b/stablehlo/stablehlo/tests/ops_stablehlo_quantized.mlir
---- stablehlo/stablehlo/tests/ops_stablehlo_quantized.mlir
-+++ stablehlo/stablehlo/tests/ops_stablehlo_quantized.mlir
-@@ -1066,146 +1066,6 @@
- 
- // -----
- 
--func.func @dot_general_hybrid_quantized(%arg0: tensor<2x3x4xf32>, %arg1: tensor<2x3x5x!quant.uniform<i8:f32, 1.0:0>>) -> tensor<2x4x5xf32> {
--  %0 = "stablehlo.dot_general"(%arg0, %arg1) {
--    dot_dimension_numbers = #stablehlo.dot<
--      lhs_batching_dimensions = [0],
--      rhs_batching_dimensions = [0],
--      lhs_contracting_dimensions = [1],
--      rhs_contracting_dimensions = [1]
--    >
--  } : (tensor<2x3x4xf32>, tensor<2x3x5x!quant.uniform<i8:f32, 1.0:0>>) -> tensor<2x4x5xf32>
--  func.return %0 : tensor<2x4x5xf32>
--}
--
--// -----
--
--func.func @dot_general_c14(%arg0: tensor<2x3x4xf32>, %arg1: tensor<2x3x5x!quant.uniform<i8:f32, 1.0:0>>) -> tensor<2x4x5x!quant.uniform<i8:f32, 1.0:17>> {
--  // expected-error@+1 {{rhs should be quantized for quantized operations and is_quantized(lhs)=is_quantized(result) should hold}}
--  %0 = "stablehlo.dot_general"(%arg0, %arg1) {
--    dot_dimension_numbers = #stablehlo.dot<
--      lhs_batching_dimensions = [0],
--      rhs_batching_dimensions = [0],
--      lhs_contracting_dimensions = [1],
--      rhs_contracting_dimensions = [1]
--    >
--  } : (tensor<2x3x4xf32>, tensor<2x3x5x!quant.uniform<i8:f32, 1.0:0>>) -> tensor<2x4x5x!quant.uniform<i8:f32, 1.0:17>>
--  func.return %0 : tensor<2x4x5x!quant.uniform<i8:f32, 1.0:17>>
--}
--
--// -----
--
--func.func @dot_general_c15_per_tensor(%arg0: tensor<2x3x4xf32>, %arg1: tensor<2x3x5x!quant.uniform<i8:f32, 1.0:30>>) -> tensor<2x4x5xf32> {
--  // expected-error@+1 {{Zero point of rhs should be 0}}
--  %0 = "stablehlo.dot_general"(%arg0, %arg1) {
--    dot_dimension_numbers = #stablehlo.dot<
--      lhs_batching_dimensions = [0],
--      rhs_batching_dimensions = [0],
--      lhs_contracting_dimensions = [1],
--      rhs_contracting_dimensions = [1]
--    >
--  } : (tensor<2x3x4xf32>, tensor<2x3x5x!quant.uniform<i8:f32, 1.0:30>>) -> tensor<2x4x5xf32>
--  func.return %0 : tensor<2x4x5xf32>
--}
--
--// -----
--
--func.func @dot_general_c15_per_axis(
--  %arg0: tensor<2x3x4x!quant.uniform<i8:f32, 1.0:17>>,
--  %arg1: tensor<2x3x5x!quant.uniform<i8:f32:0, {0.1:10, 0.1:10}>>) -> tensor<2x4x5x!quant.uniform<i8:f32:0, {0.1:-30, 0.1:-30}>> {
--  // expected-error@+1 {{Zero points of rhs should be 0}}
--  %0 = "stablehlo.dot_general"(%arg0, %arg1) {
--    dot_dimension_numbers = #stablehlo.dot<
--      lhs_batching_dimensions = [0],
--      rhs_batching_dimensions = [0],
--      lhs_contracting_dimensions = [1],
--      rhs_contracting_dimensions = [1]
--    >
--  } : (tensor<2x3x4x!quant.uniform<i8:f32, 1.0:17>>,
--  tensor<2x3x5x!quant.uniform<i8:f32:0, {0.1:10, 0.1:10}>>) -> tensor<2x4x5x!quant.uniform<i8:f32:0, {0.1:-30, 0.1:-30}>>
--  func.return %0 : tensor<2x4x5x!quant.uniform<i8:f32:0, {0.1:-30, 0.1:-30}>>
--}
--
--// -----
--
--func.func @dot_general_c16(
--  %arg0: tensor<2x3x4x!quant.uniform<i8:f32, 1.0:17>>,
--  %arg1: tensor<2x3x5x!quant.uniform<i8:f32:0, {0.1:0, 0.1:0}>>) -> tensor<2x4x5x!quant.uniform<i8:f32:0, {0.1:-30, 0.1:-30, 0.1:-30}>> {
--  // expected-error@+1 {{Quantization dimension of rhs should not be in the contracting dimension of rhs}}
--  %0 = "stablehlo.dot_general"(%arg0, %arg1) {
--    dot_dimension_numbers = #stablehlo.dot<
--      lhs_batching_dimensions = [1],
--      rhs_batching_dimensions = [1],
--      lhs_contracting_dimensions = [0],
--      rhs_contracting_dimensions = [0]
--    >
--  } : (tensor<2x3x4x!quant.uniform<i8:f32, 1.0:17>>,
--  tensor<2x3x5x!quant.uniform<i8:f32:0, {0.1:0, 0.1:0}>>) -> tensor<3x4x5x!quant.uniform<i8:f32:0, {0.1:-30, 0.1:-30, 0.1:-30}>>
--  func.return %0 : tensor<3x4x5x!quant.uniform<i8:f32:0, {0.1:-30, 0.1:-30, 0.1:-30}>>
--}
--
--// -----
--
--func.func @dot_general_c17(%arg0: tensor<2x3x4x!quant.uniform<i16:f32, 1.0:17>>, %arg1: tensor<2x3x5x!quant.uniform<i8:f32, 1.0:0>>) -> tensor<2x4x5x!quant.uniform<i8:f32, 1.0:17>> {
--  // expected-error@+1 {{mismatched lhs and rhs quantization storage types}}
--  %0 = "stablehlo.dot_general"(%arg0, %arg1) {
--    dot_dimension_numbers = #stablehlo.dot<
--      lhs_batching_dimensions = [0],
--      rhs_batching_dimensions = [0],
--      lhs_contracting_dimensions = [1],
--      rhs_contracting_dimensions = [1]
--    >
--  } : (tensor<2x3x4x!quant.uniform<i16:f32, 1.0:17>>, tensor<2x3x5x!quant.uniform<i8:f32, 1.0:0>>) -> tensor<2x4x5x!quant.uniform<i8:f32, 1.0:17>>
--  func.return %0 : tensor<2x4x5x!quant.uniform<i8:f32, 1.0:17>>
--}
--
--// -----
--
--func.func @dot_general_c18(%arg0: tensor<2x3x4x!quant.uniform<i8:f32, 1.0:17>>, %arg1: tensor<2x3x5x!quant.uniform<i8:f16, 1.0:0>>) -> tensor<2x4x5x!quant.uniform<i8:f32, 1.0:17>> {
--  // expected-error@+1 {{mismatched lhs, rhs and result quantization expressed types}}
--  %0 = "stablehlo.dot_general"(%arg0, %arg1) {
--    dot_dimension_numbers = #stablehlo.dot<
--      lhs_batching_dimensions = [0],
--      rhs_batching_dimensions = [0],
--      lhs_contracting_dimensions = [1],
--      rhs_contracting_dimensions = [1]
--    >
--  } : (tensor<2x3x4x!quant.uniform<i8:f32, 1.0:17>>, tensor<2x3x5x!quant.uniform<i8:f16, 1.0:0>>) -> tensor<2x4x5x!quant.uniform<i8:f32, 1.0:17>>
--  func.return %0 : tensor<2x4x5x!quant.uniform<i8:f32, 1.0:17>>
--}
--
--// -----
--
--func.func @dot_general_c19(%arg0: tensor<2x3x4x!quant.uniform<i8:f32, 1.0:17>>, %arg1: tensor<2x3x5x!quant.uniform<i8:f32, 1.0:0>>) -> tensor<2x4x5x!quant.uniform<i8:f32:0, {0.1:-30, 0.1:-30}>> {
--  // expected-error@+1 {{mismatched rhs and result quantization granularity}}
--  %0 = "stablehlo.dot_general"(%arg0, %arg1) {
--    dot_dimension_numbers = #stablehlo.dot<
--      lhs_batching_dimensions = [0],
--      rhs_batching_dimensions = [0],
--      lhs_contracting_dimensions = [1],
--      rhs_contracting_dimensions = [1]
--    >
--  } : (tensor<2x3x4x!quant.uniform<i8:f32, 1.0:17>>, tensor<2x3x5x!quant.uniform<i8:f32, 1.0:0>>) -> tensor<2x4x5x!quant.uniform<i8:f32:0, {0.1:-30, 0.1:-30}>>
--  func.return %0 : tensor<2x4x5x!quant.uniform<i8:f32:0, {0.1:-30, 0.1:-30}>>
--}
--
--// -----
--
--func.func @dot_general_c20(%arg0: tensor<2x3x4xf32>, %arg1: tensor<2x3x5x!quant.uniform<i8:f16, 1.0:0>>) -> tensor<2x4x5xf32> {
--  // expected-error@+1 {{mismatched rhs quantization expressed type and lhs and result element type}}
--  %0 = "stablehlo.dot_general"(%arg0, %arg1) {
--    dot_dimension_numbers = #stablehlo.dot<
--      lhs_batching_dimensions = [0],
--      rhs_batching_dimensions = [0],
--      lhs_contracting_dimensions = [1],
--      rhs_contracting_dimensions = [1]
--    >
--  } : (tensor<2x3x4xf32>, tensor<2x3x5x!quant.uniform<i8:f16, 1.0:0>>) -> tensor<2x4x5xf32>
--  func.return %0 : tensor<2x4x5xf32>
--}
--
--// -----
--
- func.func @quantized_element_type_c8(%arg0: tensor<1x2x!quant.uniform<i8<-128:127>:f32, 1.0:300>>) {
-   // expected-error-re@+1 {{operand #0 must be ranked tensor of {{.*}} 4/8/16/32-bit uniform quantized signed integer or 4/8/16/32-bit uniform quantized unsigned integer or 4/8/16/32-bit uniform quantized per axis signed integer or 4/8/16/32-bit uniform quantized per axis unsigned integer values, but got 'tensor<1x2x!quant.uniform<i8:f32, 1.000000e+00:300>>'}}
-   %0 = stablehlo.add %arg0,  %arg0 : tensor<1x2x!quant.uniform<i8<-128:127>:f32, 1.0:300>>
-diff --ruN a/stablehlo/stablehlo/tests/vhlo/stablehlo_legalize_to_vhlo.0_10_0.mlir b/stablehlo/stablehlo/tests/vhlo/stablehlo_legalize_to_vhlo.0_10_0.mlir
---- stablehlo/stablehlo/tests/vhlo/stablehlo_legalize_to_vhlo.0_10_0.mlir
-+++ stablehlo/stablehlo/tests/vhlo/stablehlo_legalize_to_vhlo.0_10_0.mlir
-@@ -13,6 +13,7 @@
- // ============ ATTRIBUTES ============
- 
- // CHECK-LABEL: "attr_comparison_direction_eq"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_direction_eq(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     // CHECK: comparison_direction = #vhlo<comparison_direction_v1 EQ>
-@@ -22,6 +23,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_direction_ne"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_direction_ne(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     // CHECK: comparison_direction = #vhlo<comparison_direction_v1 NE>
-@@ -31,6 +33,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_direction_ge"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_direction_ge(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     // CHECK: comparison_direction = #vhlo<comparison_direction_v1 GE>
-@@ -40,6 +43,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_direction_gt"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_direction_gt(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     // CHECK: comparison_direction = #vhlo<comparison_direction_v1 GT>
-@@ -49,6 +53,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_direction_le"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_direction_le(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     // CHECK: comparison_direction = #vhlo<comparison_direction_v1 LE>
-@@ -58,6 +63,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_direction_lt"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_direction_lt(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     // CHECK: comparison_direction = #vhlo<comparison_direction_v1 LT>
-@@ -67,6 +73,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_type_notype"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_type_notype(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     comparison_direction = #stablehlo<comparison_direction EQ>
-@@ -76,6 +83,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_type_float"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_type_float(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     comparison_direction = #stablehlo<comparison_direction EQ>,
-@@ -86,6 +94,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_type_totalorder"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_type_totalorder(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     comparison_direction = #stablehlo<comparison_direction EQ>,
-@@ -96,6 +105,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_type_signed"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_type_signed(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     comparison_direction = #stablehlo<comparison_direction EQ>,
-@@ -106,6 +116,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_type_unsigned"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_type_unsigned(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     comparison_direction = #stablehlo<comparison_direction EQ>,
-@@ -118,6 +129,7 @@
- // ConvDimensionNumbers aka #stablehlo.conv is covered below.
- 
- // CHECK-LABEL: "attr_custom_call_api_version_unspecified"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_custom_call_api_version_unspecified(%arg0: tensor<f32>) -> tensor<f32> {
-   %0 = "stablehlo.custom_call"(%arg0) {
-     call_target_name = "foo",
-@@ -128,6 +140,7 @@
- }
- 
- // CHECK-LABEL: "attr_custom_call_api_version_original"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_custom_call_api_version_original(%arg0: tensor<f32>) -> tensor<f32> {
-   %0 = "stablehlo.custom_call"(%arg0) {
-     call_target_name = "foo",
-@@ -138,6 +151,7 @@
- }
- 
- // CHECK-LABEL: "attr_custom_call_api_version_status_returning"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_custom_call_api_version_status_returning(%arg0: tensor<f32>) -> tensor<f32> {
-   %0 = "stablehlo.custom_call"(%arg0) {
-     call_target_name = "foo",
-@@ -148,6 +162,7 @@
- }
- 
- // CHECK-LABEL: "attr_custom_call_api_version_status_returning_unified"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_custom_call_api_version_status_returning_unified(%arg0: tensor<f32>) -> tensor<f32> {
-   %0 = "stablehlo.custom_call"(%arg0) {
-     call_target_name = "foo",
-@@ -160,6 +175,7 @@
- // DotDimensionNumbers aka #stablehlo.dot is covered below.
- 
- // CHECK-LABEL: "attr_fft_type_fft"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_fft_type_fft(%arg0: tensor<16xcomplex<f32>>) -> tensor<16xcomplex<f32>> {
-   %0 = "stablehlo.fft"(%arg0) {
-     // CHECK: fft_type = #vhlo<fft_type_v1 FFT>
-@@ -170,6 +186,7 @@
- }
- 
- // CHECK-LABEL: "attr_fft_type_ifft"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_fft_type_ifft(%arg0: tensor<16xcomplex<f32>>) -> tensor<16xcomplex<f32>> {
-   %0 = "stablehlo.fft"(%arg0) {
-     // CHECK: fft_type = #vhlo<fft_type_v1 IFFT>
-@@ -180,6 +197,7 @@
- }
- 
- // CHECK-LABEL: "attr_fft_type_rfft"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_fft_type_rfft(%arg0: tensor<16xf32>) -> tensor<9xcomplex<f32>> {
-   %0 = "stablehlo.fft"(%arg0) {
-     // CHECK: fft_type = #vhlo<fft_type_v1 RFFT>
-@@ -190,6 +208,7 @@
- }
- 
- // CHECK-LABEL: "attr_fft_type_irfft"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_fft_type_irfft(%arg0: tensor<9xcomplex<f32>>) -> tensor<16xf32> {
-   %0 = "stablehlo.fft"(%arg0) {
-     // CHECK: fft_type = #vhlo<fft_type_v1 IRFFT>
-@@ -202,6 +221,7 @@
- // GatherDimensionNumbers aka #stablehlo.gather is covered below.
- 
- // CHECK-LABEL: "attr_precision_config_default"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_precision_config_default(%arg0: tensor<8x16xf32>, %arg1: tensor<16x8xf32>) -> tensor<8x8xf32> {
-   %0 = "stablehlo.dot"(%arg0, %arg1) {
-     // CHECK: precision_config = #vhlo.array_v1<[#vhlo<precision_v1 DEFAULT>, #vhlo<precision_v1 DEFAULT>]>
-@@ -210,6 +230,7 @@
- }
- 
- // CHECK-LABEL: "attr_precision_config_high"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_precision_config_high(%arg0: tensor<8x16xf32>, %arg1: tensor<16x8xf32>) -> tensor<8x8xf32> {
-   %0 = "stablehlo.dot"(%arg0, %arg1) {
-     // CHECK: precision_config = #vhlo.array_v1<[#vhlo<precision_v1 HIGH>, #vhlo<precision_v1 HIGH>]>
-@@ -219,6 +240,7 @@
- }
- 
- // CHECK-LABEL: "attr_precision_config_highest"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_precision_config_highest(%arg0: tensor<8x16xf32>, %arg1: tensor<16x8xf32>) -> tensor<8x8xf32> {
-   %0 = "stablehlo.dot"(%arg0, %arg1) {
-     // CHECK: precision_config = #vhlo.array_v1<[#vhlo<precision_v1 HIGHEST>, #vhlo<precision_v1 HIGHEST>]>
-@@ -228,6 +250,7 @@
- }
- 
- // CHECK-LABEL: "attr_rng_algorithm_default"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_rng_algorithm_default(%arg0: tensor<f32>) -> (tensor<f32>, tensor<f32>) {
-   %0:2 = "stablehlo.rng_bit_generator"(%arg0) {
-     // CHECK: rng_algorithm = #vhlo<rng_algorithm_v1 DEFAULT>
-@@ -237,6 +260,7 @@
- }
- 
- // CHECK-LABEL: "attr_rng_algorithm_three_fry"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_rng_algorithm_three_fry(%arg0: tensor<f32>) -> (tensor<f32>, tensor<f32>) {
-   %0:2 = "stablehlo.rng_bit_generator"(%arg0) {
-     // CHECK: rng_algorithm = #vhlo<rng_algorithm_v1 THREE_FRY>
-@@ -246,6 +270,7 @@
- }
- 
- // CHECK-LABEL: "attr_rng_algorithm_philox"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_rng_algorithm_philox(%arg0: tensor<f32>) -> (tensor<f32>, tensor<f32>) {
-   %0:2 = "stablehlo.rng_bit_generator"(%arg0) {
-     // CHECK: rng_algorithm = #vhlo<rng_algorithm_v1 PHILOX>
-@@ -255,6 +280,7 @@
- }
- 
- // CHECK-LABEL: "attr_rng_distribution_uniform"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @attr_rng_distribution_uniform(%arg0: tensor<f32>, %arg1: tensor<f32>, %arg2: tensor<0xindex>) -> tensor<f32> {
-   %0 = "stablehlo.rng"(%arg0, %arg1, %arg2) {
-     // CHECK: rng_distribution = #vhlo<rng_distribution_v1 UNIFORM>
-@@ -264,6 +290,7 @@
- }
- 
- // CHECK-LABEL: "attr_rng_distribution_normal"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @attr_rng_distribution_normal(%arg0: tensor<f32>, %arg1: tensor<f32>, %arg2: tensor<0xindex>) -> tensor<f32> {
-   %0 = "stablehlo.rng"(%arg0, %arg1, %arg2) {
-     // CHECK: rng_distribution = #vhlo<rng_distribution_v1 NORMAL>
-@@ -275,6 +302,7 @@
- // ScatterDimensionNumbers aka #stablehlo.scatter is covered below.
- 
- // CHECK-LABEL: "attr_transpose_no_transpose"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_transpose_no_transpose(%arg0: tensor<16x16xf32>, %arg1: tensor<16x16xf32>) ->  tensor<16x16xf32> {
-   %0 = "stablehlo.triangular_solve"(%arg0, %arg1) {
-     left_side = true,
-@@ -287,6 +315,7 @@
- }
- 
- // CHECK-LABEL: "attr_transpose_transpose"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_transpose_transpose(%arg0: tensor<16x16xf32>, %arg1: tensor<16x16xf32>) ->  tensor<16x16xf32> {
-   %0 = "stablehlo.triangular_solve"(%arg0, %arg1) {
-     left_side = true,
-@@ -299,6 +328,7 @@
- }
- 
- // CHECK-LABEL: "attr_transpose_adjoint"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_transpose_adjoint(%arg0: tensor<16x16xf32>, %arg1: tensor<16x16xf32>) ->  tensor<16x16xf32> {
-   %0 = "stablehlo.triangular_solve"(%arg0, %arg1) {
-     left_side = true,
-@@ -313,10 +343,9 @@
- // TypeExtensionsAttr aka #stablehlo.type_extensions is covered below.
- 
- // CHECK-LABEL: "attr_type_extensions_bounds"
--func.func @attr_type_extensions_bounds(
--    %arg0: tensor<?x?xf32, #stablehlo.type_extensions<bounds = [16, ?]>>)
--    -> tensor<?x?xf32, #stablehlo.type_extensions<bounds = [16, ?]>> {
--  // CHECK: "vhlo.return_v1"(%arg0) : (!vhlo.tensor_v1<?x?x!vhlo.f32_v1, #vhlo.type_extensions_v1<bounds = [16, ?]>>) -> ()
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @attr_type_extensions_bounds(%arg0: tensor<?x?xf32, #stablehlo.type_extensions<bounds = [16, ?]>>) -> tensor<?x?xf32, #stablehlo.type_extensions<bounds = [16, ?]>> {
-+  // CHECK: "vhlo.return_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<?x?x!vhlo.f32_v1, #vhlo.type_extensions_v1<bounds = [16, ?]>>) -> ()
-   func.return %arg0 : tensor<?x?xf32, #stablehlo.type_extensions<bounds = [16, ?]>>
- }
- 
-@@ -324,8 +353,9 @@
- // ============ DEFAULTS ============
- 
- // CHECK-LABEL: "default_all_gather"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_all_gather(%arg0: tensor<16x8xf32>) -> tensor<16x16xf32> {
--  //               CHECK: "vhlo.all_gather_v1"(%arg0) <{
-+  //               CHECK: "vhlo.all_gather_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   all_gather_dim = #vhlo.integer_v1<1 : i64>
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0], [1]]> : tensor<2x1xi64>>,
-@@ -339,8 +369,9 @@
- }
- 
- // CHECK-LABEL: "default_all_reduce"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_all_reduce(%arg0: tensor<f32>) -> tensor<f32> {
--  //               CHECK: "vhlo.all_reduce_v1"(%arg0)
-+  //               CHECK: "vhlo.all_reduce_v1"(%[[ARG0]])
-   //          CHECK-SAME: <{
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0], [1]]> : tensor<2x1xi64>>,
-@@ -362,8 +393,9 @@
- }
- 
- // CHECK-LABEL: "default_all_to_all"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_all_to_all(%arg0: tensor<4x16xf32>) -> tensor<16x4xf32> {
--  //               CHECK: "vhlo.all_to_all_v1"(%arg0) <{
-+  //               CHECK: "vhlo.all_to_all_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   //          CHECK-SAME:   concat_dimension = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0, 1, 2, 3]]> : tensor<1x4xi64>>,
-@@ -380,8 +412,9 @@
- }
- 
- // CHECK-LABEL: "default_cholesky"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_cholesky(%arg0: tensor<1x16x16xf32>) -> tensor<1x16x16xf32> {
--  //      CHECK: "vhlo.cholesky_v1"(%arg0) <{
-+  //      CHECK: "vhlo.cholesky_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   lower = #vhlo.bool_v1<false>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<1x16x16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<1x16x16x!vhlo.f32_v1>
-   %0 = "stablehlo.cholesky"(%arg0) : (tensor<1x16x16xf32>) -> tensor<1x16x16xf32>
-@@ -389,8 +422,9 @@
- }
- 
- // CHECK-LABEL: "default_collective_permute"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_collective_permute(%arg0: tensor<16x8xf32>) -> tensor<16x8xf32> {
--  //               CHECK: "vhlo.collective_permute_v1"(%arg0) <{
-+  //               CHECK: "vhlo.collective_permute_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME{LITERAL}:   source_target_pairs = #vhlo.tensor_v1<dense<[[0, 1], [1, 2], [2, 3]]> : tensor<3x2xi64>>
-   //          CHECK-SAME: }> : (!vhlo.tensor_v1<16x8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x8x!vhlo.f32_v1>
-@@ -401,8 +435,9 @@
- }
- 
- // CHECK-LABEL: "default_compare"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @default_compare(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
--  //      CHECK: "vhlo.compare_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.compare_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   compare_type = #vhlo<comparison_type_v1 NOTYPE>,
-   // CHECK-SAME:   comparison_direction = #vhlo<comparison_direction_v1 EQ>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-@@ -413,8 +448,9 @@
- }
- 
- // CHECK-LABEL: "default_convolution"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @default_convolution(%arg0: tensor<1x8x8x207xf32>, %arg1: tensor<3x3x207x16xf32>) -> tensor<1x6x6x16xf32> {
--  //      CHECK: "vhlo.convolution_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.convolution_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   batch_group_count = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   feature_group_count = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   input_batch_dimension = #vhlo.integer_v1<0 : i64>,
-@@ -442,8 +478,9 @@
- }
- 
- // CHECK-LABEL: "default_custom_call"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_custom_call(%arg0: tensor<f32>) -> tensor<f32> {
--  //      CHECK: "vhlo.custom_call_v1"(%arg0) <{
-+  //      CHECK: "vhlo.custom_call_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   api_version = #vhlo<api_version_v1 API_VERSION_ORIGINAL>,
-   // CHECK-SAME:   backend_config = #vhlo.string_v1<"">,
-   // CHECK-SAME:   call_target_name = #vhlo.string_v1<"foo">,
-@@ -460,8 +497,9 @@
- }
- 
- // CHECK-LABEL: "default_dot_general"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @default_dot_general(%arg0: tensor<8x8x16xf32>, %arg1: tensor<8x16x8xf32>) -> tensor<8x8x8xf32> {
--  //      CHECK: "vhlo.dot_general_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.dot_general_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   lhs_batching_dimensions = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>,
-   // CHECK-SAME:   lhs_contracting_dimensions = #vhlo.tensor_v1<dense<2> : tensor<1xi64>>,
-   // CHECK-SAME:   precision_config = #vhlo.array_v1<[#vhlo<precision_v1 DEFAULT>, #vhlo<precision_v1 DEFAULT>]>,
-@@ -480,8 +518,9 @@
- }
- 
- // CHECK-LABEL: "default_dot"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @default_dot(%arg0: tensor<8x16xf32>, %arg1: tensor<16x8xf32>) -> tensor<8x8xf32> {
--  //      CHECK: "vhlo.dot_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.dot_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   precision_config = #vhlo.array_v1<[#vhlo<precision_v1 DEFAULT>, #vhlo<precision_v1 DEFAULT>]>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<8x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<8x8x!vhlo.f32_v1>
-   %0 = "stablehlo.dot"(%arg0, %arg1) : (tensor<8x16xf32>, tensor<16x8xf32>) -> tensor<8x8xf32>
-@@ -489,8 +528,9 @@
- }
- 
- // CHECK-LABEL: "default_dynamic_broadcast_in_dim"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @default_dynamic_broadcast_in_dim(%arg0: tensor<?x?xf32>, %arg1: tensor<2xindex>) -> tensor<?x?xf32> {
--  //      CHECK: "vhlo.dynamic_broadcast_in_dim_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.dynamic_broadcast_in_dim_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   broadcast_dimensions = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-   // CHECK-SAME:   known_expanding_dimensions = #vhlo.tensor_v1<dense<> : tensor<0xi64>>,
-   // CHECK-SAME:   known_nonexpanding_dimensions = #vhlo.tensor_v1<dense<> : tensor<0xi64>>
-@@ -502,8 +542,9 @@
- }
- 
- // CHECK-LABEL: "default_dynamic_conv"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @default_dynamic_conv(%arg0: tensor<1x8x8x207xf32>, %arg1: tensor<3x3x207x16xf32>, %arg2: tensor<4xi32>) -> tensor<1x?x?x16xf32> {
--  //      CHECK: "vhlo.dynamic_conv_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.dynamic_conv_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   batch_group_count = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   feature_group_count = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   input_batch_dimension = #vhlo.integer_v1<0 : i64>,
-@@ -531,8 +572,9 @@
- }
- 
- // CHECK-LABEL: "default_dynamic_gather"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @default_dynamic_gather(%arg0 : tensor<2x4x9xf32>, %arg1 : tensor<1x5x2xi32>, %arg2 : tensor<3xi32>) -> tensor<1x5x8xf32> {
--  //      CHECK: "vhlo.dynamic_gather_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.dynamic_gather_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   collapsed_slice_dims = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-   // CHECK-SAME:   index_vector_dim = #vhlo.integer_v1<2 : i64>,
-   // CHECK-SAME:   indices_are_sorted = #vhlo.bool_v1<false>,
-@@ -558,15 +600,16 @@
-   // CHECK-SAME:   sym_name = #vhlo.string_v1<"default_func">,
-   // CHECK-SAME:   sym_visibility = #vhlo.string_v1<"">
-   // CHECK-SAME: }> ({
--  // CHECK-NEXT:   ^[[BB:bb.*]](%arg0: !vhlo.tensor_v1<!vhlo.f32_v1>):
--  // CHECK-NEXT:     "vhlo.return_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-+  // CHECK-NEXT:   ^[[BB:bb.*]](%[[ARG0:.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
-+  // CHECK-NEXT:     "vhlo.return_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-   // CHECK-NEXT: }) : () -> ()
-   func.return %arg0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "dynamic_gather"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @dynamic_gather(%arg0 : tensor<2x4x9xf32>, %arg1 : tensor<1x5x2xi32>) -> tensor<1x5x1xf32> {
--  //      CHECK: "vhlo.gather_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.gather_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   collapsed_slice_dims = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-   // CHECK-SAME:   index_vector_dim = #vhlo.integer_v1<2 : i64>,
-   // CHECK-SAME:   indices_are_sorted = #vhlo.bool_v1<false>,
-@@ -587,8 +630,9 @@
- }
- 
- // CHECK-LABEL: "default_infeed"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_infeed(%arg0: !stablehlo.token) -> (tensor<f32>, !stablehlo.token) {
--  //               CHECK: "vhlo.infeed_v1"(%arg0) <{
-+  //               CHECK: "vhlo.infeed_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   infeed_config = #vhlo.string_v1<"">,
-   // CHECK-SAME{LITERAL}:   layout = #vhlo.array_v1<[]>
-   //          CHECK-SAME: }> : (!vhlo.token_v1) -> (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.token_v1)
-@@ -597,8 +641,9 @@
- }
- 
- // CHECK-LABEL: "default_outfeed"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @default_outfeed(%arg0: tensor<f32>, %arg1: !stablehlo.token) -> !stablehlo.token {
--  //      CHECK: "vhlo.outfeed_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.outfeed_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   outfeed_config = #vhlo.string_v1<"">
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.token_v1) -> !vhlo.token_v1
-   %0 = "stablehlo.outfeed"(%arg0, %arg1) : (tensor<f32>, !stablehlo.token) -> !stablehlo.token
-@@ -606,8 +651,9 @@
- }
- 
- // CHECK-LABEL: "default_recv"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_recv(%arg0: !stablehlo.token) -> (tensor<f32>, !stablehlo.token) {
--  //      CHECK: "vhlo.recv_v1"(%arg0) <{
-+  //      CHECK: "vhlo.recv_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME:   channel_type = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   is_host_transfer = #vhlo.bool_v1<false>
-@@ -619,8 +665,9 @@
- }
- 
- // CHECK-LABEL: "default_send"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @default_send(%arg0: tensor<f32>, %arg1: !stablehlo.token) -> !stablehlo.token {
--  //      CHECK: "vhlo.send_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.send_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME:   channel_type = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   is_host_transfer = #vhlo.bool_v1<false>
-@@ -632,8 +679,9 @@
- }
- 
- // CHECK-LABEL: "default_reduce_scatter"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_reduce_scatter(%arg0: tensor<16xf32>) -> tensor<16xf32> {
--  //               CHECK: "vhlo.reduce_scatter_v1"(%arg0) <{
-+  //               CHECK: "vhlo.reduce_scatter_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0], [1]]> : tensor<2x1xi64>>,
-   //          CHECK-SAME:   scatter_dimension = #vhlo.integer_v1<0 : i64>
-@@ -655,8 +703,9 @@
- }
- 
- // CHECK-LABEL: "default_reduce_window"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @default_reduce_window(%arg0: tensor<2x17x31x7xf32>, %arg1: tensor<f32>) -> tensor<2x16x30x7xf32> {
--  //               CHECK: "vhlo.reduce_window_v1"(%arg0, %arg1)  <{
-+  //               CHECK: "vhlo.reduce_window_v1"(%[[ARG0]], %[[ARG1]])  <{
-   //          CHECK-SAME:   base_dilations = #vhlo.tensor_v1<dense<1> : tensor<4xi64>>,
-   // CHECK-SAME{LITERAL}:   padding = #vhlo.tensor_v1<dense<0> : tensor<4x2xi64>>,
-   //          CHECK-SAME:   window_dilations = #vhlo.tensor_v1<dense<1> : tensor<4xi64>>,
-@@ -678,8 +727,9 @@
- }
- 
- // CHECK-LABEL: "default_scatter"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @default_scatter(%arg0: tensor<200x100x300xf32>, %arg1: tensor<10x2xi32>, %arg2: tensor<10x300xf32>) -> tensor<200x100x300xf32> {
--  //      CHECK: "vhlo.scatter_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.scatter_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   index_vector_dim = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   indices_are_sorted = #vhlo.bool_v1<false>,
-   // CHECK-SAME:   inserted_window_dims = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-@@ -707,8 +757,9 @@
- }
- 
- // CHECK-LABEL: "default_select_and_scatter"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @default_select_and_scatter(%arg0: tensor<10x24x24x64xf32>, %arg1: tensor<10x23x23x64xf32>, %arg2: tensor<f32>) -> tensor<10x24x24x64xf32> {
--  //      CHECK: "vhlo.select_and_scatter_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.select_and_scatter_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   padding = #vhlo.tensor_v1<dense<0> : tensor<4x2xi64>>,
-   // CHECK-SAME:   window_dimensions = #vhlo.tensor_v1<dense<[1, 2, 2, 1]> : tensor<4xi64>>,
-   // CHECK-SAME:   window_strides = #vhlo.tensor_v1<dense<1> : tensor<4xi64>>
-@@ -736,8 +787,9 @@
- }
- 
- // CHECK-LABEL: "default_sort"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_sort(%arg0: tensor<16xf32>) -> tensor<16xf32> {
--  //      CHECK: "vhlo.sort_v1"(%arg0) <{
-+  //      CHECK: "vhlo.sort_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   dimension = #vhlo.integer_v1<-1 : i64>
-   // CHECK-SAME:   is_stable = #vhlo.bool_v1<false>
-   // CHECK-SAME: }> ({
-@@ -756,29 +808,33 @@
- // ============ OPS ============
- 
- // CHECK-LABEL: "op_abs"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_abs(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.abs_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.abs_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.abs"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_add"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_add(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_after_all"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_after_all(%arg0: !stablehlo.token) -> !stablehlo.token {
--  // CHECK: "vhlo.after_all_v1"(%arg0) : (!vhlo.token_v1) -> !vhlo.token_v1
-+  // CHECK: "vhlo.after_all_v1"(%[[ARG0]]) : (!vhlo.token_v1) -> !vhlo.token_v1
-   %0 = "stablehlo.after_all"(%arg0) : (!stablehlo.token) -> !stablehlo.token
-   func.return %0 : !stablehlo.token
- }
- 
- // CHECK-LABEL: "op_all_gather"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_all_gather(%arg0: tensor<16x8xf32>) -> tensor<16x16xf32> {
--  //               CHECK: "vhlo.all_gather_v1"(%arg0) <{
-+  //               CHECK: "vhlo.all_gather_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   all_gather_dim = #vhlo.integer_v1<1 : i64>
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0], [1]]> : tensor<2x1xi64>>,
-@@ -794,8 +850,9 @@
- }
- 
- // CHECK-LABEL: "op_all_reduce"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_all_reduce(%arg0: tensor<f32>) -> tensor<f32> {
--  //               CHECK: "vhlo.all_reduce_v1"(%arg0) <{
-+  //               CHECK: "vhlo.all_reduce_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0], [1]]> : tensor<2x1xi64>>,
-   //          CHECK-SAME:   use_global_device_ids = #vhlo.bool_v1<true>
-@@ -817,8 +874,9 @@
- }
- 
- // CHECK-LABEL: "op_all_to_all"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_all_to_all(%arg0: tensor<4x16xf32>) -> tensor<16x4xf32> {
--  //               CHECK: "vhlo.all_to_all_v1"(%arg0) <{
-+  //               CHECK: "vhlo.all_to_all_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<1 : i64>,
-   //          CHECK-SAME:   concat_dimension = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0, 1, 2, 3]]> : tensor<1x4xi64>>,
-@@ -836,22 +894,25 @@
- }
- 
- // CHECK-LABEL: "op_and"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_and(%arg0: tensor<i1>, %arg1: tensor<i1>) -> tensor<i1> {
--  // CHECK: "vhlo.and_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-+  // CHECK: "vhlo.and_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-   %0 = "stablehlo.and"(%arg0, %arg1) : (tensor<i1>, tensor<i1>) -> tensor<i1>
-   func.return %0 : tensor<i1>
- }
- 
- // CHECK-LABEL: "op_atan2"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_atan2(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.atan2_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.atan2_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.atan2"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_batch_norm_grad"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}}, %[[ARG3:.*]]: {{.*}}, %[[ARG4:.*]]: {{.*}})
- func.func @op_batch_norm_grad(%arg0: tensor<16x16x16x16xf32>, %arg1: tensor<16xf32>, %arg2: tensor<16xf32>, %arg3: tensor<16xf32>, %arg4: tensor<16x16x16x16xf32>) -> (tensor<16x16x16x16xf32>, tensor<16xf32>, tensor<16xf32>) {
--  //      CHECK: "vhlo.batch_norm_grad_v1"(%arg0, %arg1, %arg2, %arg3, %arg4) <{
-+  //      CHECK: "vhlo.batch_norm_grad_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]], %[[ARG3]], %[[ARG4]]) <{
-   // CHECK-SAME:   epsilon = #vhlo.float_v1<1.000000e-03 : !vhlo.f32_v1>,
-   // CHECK-SAME:   feature_index = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<16x16x16x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x16x16x16x!vhlo.f32_v1>) -> (!vhlo.tensor_v1<16x16x16x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>)
-@@ -863,8 +924,9 @@
- }
- 
- // CHECK-LABEL: "op_batch_norm_inference"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}}, %[[ARG3:.*]]: {{.*}}, %[[ARG4:.*]]: {{.*}})
- func.func @op_batch_norm_inference(%arg0: tensor<16x16x16x16xf32>, %arg1: tensor<16xf32>, %arg2: tensor<16xf32>, %arg3: tensor<16xf32>, %arg4: tensor<16xf32>) -> tensor<16x16x16x16xf32> {
--  //      CHECK: "vhlo.batch_norm_inference_v1"(%arg0, %arg1, %arg2, %arg3, %arg4) <{
-+  //      CHECK: "vhlo.batch_norm_inference_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]], %[[ARG3]], %[[ARG4]]) <{
-   // CHECK-SAME:   epsilon = #vhlo.float_v1<1.000000e-03 : !vhlo.f32_v1>,
-   // CHECK-SAME:   feature_index = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<16x16x16x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x16x16x16x!vhlo.f32_v1>
-@@ -876,8 +938,9 @@
- }
- 
- // CHECK-LABEL: "op_batch_norm_training"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_batch_norm_training(%arg0: tensor<16x16x16x16xf32>, %arg1: tensor<16xf32>, %arg2: tensor<16xf32>) -> (tensor<16x16x16x16xf32>, tensor<16xf32>, tensor<16xf32>) {
--  //      CHECK: "vhlo.batch_norm_training_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.batch_norm_training_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   epsilon = #vhlo.float_v1<1.000000e-03 : !vhlo.f32_v1>,
-   // CHECK-SAME:   feature_index = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<16x16x16x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>) -> (!vhlo.tensor_v1<16x16x16x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>)
-@@ -889,15 +952,17 @@
- }
- 
- // CHECK-LABEL: "op_bitcast_convert"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_bitcast_convert(%arg0: tensor<i32>) -> tensor<f32> {
--  // CHECK: "vhlo.bitcast_convert_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.bitcast_convert_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.bitcast_convert"(%arg0) : (tensor<i32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_broadcast_in_dim"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_broadcast_in_dim(%arg0: tensor<16xf32>) -> tensor<16x16xf32> {
--  //      CHECK: "vhlo.broadcast_in_dim_v1"(%arg0) <{
-+  //      CHECK: "vhlo.broadcast_in_dim_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   broadcast_dimensions = #vhlo.tensor_v1<dense<1> : tensor<1xi64>>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x16x!vhlo.f32_v1>
-   %0 = "stablehlo.broadcast_in_dim"(%arg0) {
-@@ -907,8 +972,9 @@
- }
- 
- // CHECK-LABEL: "op_broadcast"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_broadcast(%arg0: tensor<16xf32>) -> tensor<16x16xf32> {
--  //      CHECK: "vhlo.broadcast_v1"(%arg0) <{
-+  //      CHECK: "vhlo.broadcast_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   broadcast_sizes = #vhlo.tensor_v1<dense<16> : tensor<1xi64>>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x16x!vhlo.f32_v1>
-   %0 = "stablehlo.broadcast"(%arg0) {
-@@ -918,9 +984,10 @@
- }
- 
- // CHECK-LABEL: "op_case"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_case(%arg0: tensor<i32>, %arg1: tensor<f32>) -> tensor<f32> {
--  //      CHECK: "vhlo.case_v1"(%arg0) ({
--  // CHECK-NEXT:   "vhlo.return_v1"(%arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-+  //      CHECK: "vhlo.case_v1"(%[[ARG0]]) ({
-+  // CHECK-NEXT:   "vhlo.return_v1"(%[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-   // CHECK-NEXT: }) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.case"(%arg0) ({
-     "stablehlo.return"(%arg1) : (tensor<f32>) -> ()
-@@ -929,22 +996,25 @@
- }
- 
- // CHECK-LABEL: "op_cbrt"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_cbrt(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.cbrt_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.cbrt_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.cbrt"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_ceil"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_ceil(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.ceil_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.ceil_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.ceil"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_cholesky"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_cholesky(%arg0: tensor<1x16x16xf32>) -> tensor<1x16x16xf32> {
--  //      CHECK: "vhlo.cholesky_v1"(%arg0) <{
-+  //      CHECK: "vhlo.cholesky_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   lower = #vhlo.bool_v1<true>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<1x16x16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<1x16x16x!vhlo.f32_v1>
-   %0 = "stablehlo.cholesky"(%arg0) {
-@@ -954,22 +1024,25 @@
- }
- 
- // CHECK-LABEL: "op_clamp"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_clamp(%arg0: tensor<f32>, %arg1: tensor<f32>, %arg2: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.clamp_v1"(%arg0, %arg1, %arg2) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.clamp_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.clamp"(%arg0, %arg1, %arg2) : (tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_count_leading_zeros"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_count_leading_zeros(%arg0: tensor<i32>) -> tensor<i32> {
--  // CHECK: "vhlo.count_leading_zeros_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-+  // CHECK: "vhlo.count_leading_zeros_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-   %0 = "stablehlo.count_leading_zeros"(%arg0) : (tensor<i32>) -> tensor<i32>
-   func.return %0 : tensor<i32>
- }
- 
- // CHECK-LABEL: "op_collective_permute"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_collective_permute(%arg0: tensor<16x8xf32>) -> tensor<16x8xf32> {
--  //               CHECK: "vhlo.collective_permute_v1"(%arg0) <{
-+  //               CHECK: "vhlo.collective_permute_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME{LITERAL}:   source_target_pairs = #vhlo.tensor_v1<dense<[[0, 1], [1, 2], [2, 3]]> : tensor<3x2xi64>>
-   //          CHECK-SAME: }> : (!vhlo.tensor_v1<16x8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x8x!vhlo.f32_v1>
-@@ -981,8 +1054,9 @@
- }
- 
- // CHECK-LABEL: "op_compare"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_compare(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
--  //      CHECK: "vhlo.compare_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.compare_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   compare_type = #vhlo<comparison_type_v1 TOTALORDER>,
-   // CHECK-SAME:   comparison_direction = #vhlo<comparison_direction_v1 EQ>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-@@ -994,15 +1068,17 @@
- }
- 
- // CHECK-LABEL: "op_complex"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_complex(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<complex<f32>> {
--  // CHECK: "vhlo.complex_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>
-+  // CHECK: "vhlo.complex_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>
-   %0 = "stablehlo.complex"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<complex<f32>>
-   func.return %0 : tensor<complex<f32>>
- }
- 
- // CHECK-LABEL: "op_concatenate"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_concatenate(%arg0: tensor<8xf32>, %arg1: tensor<8xf32>) -> tensor<16xf32> {
--  //      CHECK: "vhlo.concatenate_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.concatenate_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   dimension = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<8x!vhlo.f32_v1>, !vhlo.tensor_v1<8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x!vhlo.f32_v1>
-   %0 = "stablehlo.concatenate"(%arg0, %arg1) {
-@@ -1012,6 +1088,7 @@
- }
- 
- // CHECK-LABEL: "op_constant"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_constant(%arg0: tensor<f32>) -> tensor<f32> {
-   //      CHECK: "vhlo.constant_v1"() <{
-   // CHECK-SAME:   value = #vhlo.tensor_v1<dense<0.000000e+00> : tensor<f32>>
-@@ -1023,15 +1100,17 @@
- }
- 
- // CHECK-LABEL: "op_convert"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_convert(%arg0: tensor<i32>) -> tensor<f32> {
--  // CHECK: "vhlo.convert_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.convert_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.convert"(%arg0) : (tensor<i32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_convolution"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_convolution(%arg0: tensor<1x8x8x207xf32>, %arg1: tensor<3x3x207x16xf32>) -> tensor<1x7x7x16xf32> {
--  //      CHECK: "vhlo.convolution_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.convolution_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   batch_group_count = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   feature_group_count = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   input_batch_dimension = #vhlo.integer_v1<0 : i64>,
-@@ -1065,8 +1144,9 @@
- }
- 
- // CHECK-LABEL: "op_cosine"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_cosine(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.cosine_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.cosine_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.cosine"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
-@@ -1079,8 +1159,9 @@
- }
- 
- // CHECK-LABEL: "op_cross_replica_sum"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_cross_replica_sum(%arg0: tensor<f32>) -> tensor<f32> {
--  //               CHECK: "vhlo.cross-replica-sum_v1"(%arg0) <{
-+  //               CHECK: "vhlo.cross-replica-sum_v1"(%[[ARG0]]) <{
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0], [1]]> : tensor<2x1xi64>>
-   //          CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.cross-replica-sum"(%arg0) {
-@@ -1090,8 +1171,9 @@
- }
- 
- // CHECK-LABEL: "op_custom_call"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_custom_call(%arg0: tensor<f32>) -> tensor<f32> {
--  //      CHECK: "vhlo.custom_call_v1"(%arg0) <{
-+  //      CHECK: "vhlo.custom_call_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   api_version = #vhlo<api_version_v1 API_VERSION_STATUS_RETURNING>,
-   // CHECK-SAME:   backend_config = #vhlo.string_v1<"\08\03\1A\02">,
-   // CHECK-SAME:   call_target_name = #vhlo.string_v1<"foo">,
-@@ -1122,15 +1204,17 @@
- }
- 
- // CHECK-LABEL: "op_divide"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_divide(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.divide_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.divide_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.divide"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_dot_general"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_dot_general(%arg0: tensor<8x8x16xf32>, %arg1: tensor<8x16x8xf32>) -> tensor<8x8x8xf32> {
--  //      CHECK: "vhlo.dot_general_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.dot_general_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   lhs_batching_dimensions = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>,
-   // CHECK-SAME:   lhs_contracting_dimensions = #vhlo.tensor_v1<dense<2> : tensor<1xi64>>,
-   // CHECK-SAME:   precision_config = #vhlo.array_v1<[#vhlo<precision_v1 HIGHEST>, #vhlo<precision_v1 HIGHEST>]>,
-@@ -1150,8 +1234,9 @@
- }
- 
- // CHECK-LABEL: "op_dot"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_dot(%arg0: tensor<8x16xf32>, %arg1: tensor<16x8xf32>) -> tensor<8x8xf32> {
--  //      CHECK: "vhlo.dot_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.dot_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   precision_config = #vhlo.array_v1<[#vhlo<precision_v1 HIGHEST>, #vhlo<precision_v1 HIGHEST>]>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<8x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<8x8x!vhlo.f32_v1>
-   %0 = "stablehlo.dot"(%arg0, %arg1) {
-@@ -1161,8 +1246,9 @@
- }
- 
- // CHECK-LABEL: "op_dynamic_broadcast_in_dim"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_dynamic_broadcast_in_dim(%arg0: tensor<?x?xf32>, %arg1: tensor<2xindex>) -> tensor<?x?xf32> {
--  //      CHECK: "vhlo.dynamic_broadcast_in_dim_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.dynamic_broadcast_in_dim_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   broadcast_dimensions = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-   // CHECK-SAME:   known_expanding_dimensions = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>,
-   // CHECK-SAME:   known_nonexpanding_dimensions = #vhlo.tensor_v1<dense<1> : tensor<1xi64>>
-@@ -1176,8 +1262,9 @@
- }
- 
- // CHECK-LABEL: "op_dynamic_conv"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_dynamic_conv(%arg0: tensor<1x8x8x207xf32>, %arg1: tensor<3x3x207x16xf32>, %arg2: tensor<4xi32>) -> tensor<1x?x?x16xf32> {
--  //      CHECK: "vhlo.dynamic_conv_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.dynamic_conv_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   batch_group_count = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   feature_group_count = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   input_batch_dimension = #vhlo.integer_v1<0 : i64>,
-@@ -1211,8 +1298,9 @@
- }
- 
- // CHECK-LABEL: "op_dynamic_gather"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_dynamic_gather(%arg0 : tensor<2x4x9xf32>, %arg1 : tensor<1x5x2xi32>, %arg2 : tensor<3xi32>) -> tensor<1x5x8xf32> {
--  //      CHECK: "vhlo.dynamic_gather_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.dynamic_gather_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   collapsed_slice_dims = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-   // CHECK-SAME:   index_vector_dim = #vhlo.integer_v1<2 : i64>,
-   // CHECK-SAME:   indices_are_sorted = #vhlo.bool_v1<true>,
-@@ -1232,8 +1320,9 @@
- }
- 
- // CHECK-LABEL: "op_dynamic_iota"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_dynamic_iota(%arg0: tensor<1xindex>) -> tensor<?xf32> {
--  //      CHECK: "vhlo.dynamic_iota_v1"(%arg0) <{
-+  //      CHECK: "vhlo.dynamic_iota_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   iota_dimension = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<1x!vhlo.index_v1>) -> !vhlo.tensor_v1<?x!vhlo.f32_v1>
-   %0 = "stablehlo.dynamic_iota"(%arg0) {
-@@ -1243,22 +1332,25 @@
- }
- 
- // CHECK-LABEL: "op_dynamic_pad"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}}, %[[ARG3:.*]]: {{.*}}, %[[ARG4:.*]]: {{.*}})
- func.func @op_dynamic_pad(%arg0: tensor<?xf32>, %arg1: tensor<f32>, %arg2: tensor<1xindex>, %arg3: tensor<1xindex>, %arg4: tensor<1xindex>) -> tensor<?xf32> {
--  // CHECK: "vhlo.dynamic_pad_v1"(%arg0, %arg1, %arg2, %arg3, %arg4) : (!vhlo.tensor_v1<?x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>) -> !vhlo.tensor_v1<?x!vhlo.f32_v1>
-+  // CHECK: "vhlo.dynamic_pad_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]], %[[ARG3]], %[[ARG4]]) : (!vhlo.tensor_v1<?x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>) -> !vhlo.tensor_v1<?x!vhlo.f32_v1>
-   %0 = "stablehlo.dynamic_pad"(%arg0, %arg1, %arg2, %arg3, %arg4) : (tensor<?xf32>, tensor<f32>, tensor<1xindex>, tensor<1xindex>, tensor<1xindex>) -> tensor<?xf32>
-   func.return %0 : tensor<?xf32>
- }
- 
- // CHECK-LABEL: "op_dynamic_reshape"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_dynamic_reshape(%arg0: tensor<16xf32>, %arg1: tensor<2xindex>) -> tensor<?x?xf32> {
--  // CHECK: "vhlo.dynamic_reshape_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<2x!vhlo.index_v1>) -> !vhlo.tensor_v1<?x?x!vhlo.f32_v1>
-+  // CHECK: "vhlo.dynamic_reshape_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<2x!vhlo.index_v1>) -> !vhlo.tensor_v1<?x?x!vhlo.f32_v1>
-   %0 = "stablehlo.dynamic_reshape"(%arg0, %arg1) : (tensor<16xf32>, tensor<2xindex>) -> tensor<?x?xf32>
-   func.return %0 : tensor<?x?xf32>
- }
- 
- // CHECK-LABEL: "op_dynamic_slice"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_dynamic_slice(%arg0: tensor<16xf32>, %arg1: tensor<i64>) -> tensor<4xf32> {
--  //      CHECK: "vhlo.dynamic_slice_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.dynamic_slice_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   slice_sizes = #vhlo.tensor_v1<dense<4> : tensor<1xi64>>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.i64_v1>) -> !vhlo.tensor_v1<4x!vhlo.f32_v1>
-   %0 = "stablehlo.dynamic_slice"(%arg0, %arg1) {
-@@ -1268,15 +1360,17 @@
- }
- 
- // CHECK-LABEL: "op_dynamic_update_slice"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_dynamic_update_slice(%arg0: tensor<16xf32>, %arg1: tensor<4xf32>, %arg2: tensor<i64>) -> tensor<16xf32> {
--  // CHECK: "vhlo.dynamic_update_slice_v1"(%arg0, %arg1, %arg2) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<4x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.i64_v1>) -> !vhlo.tensor_v1<16x!vhlo.f32_v1>
-+  // CHECK: "vhlo.dynamic_update_slice_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<4x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.i64_v1>) -> !vhlo.tensor_v1<16x!vhlo.f32_v1>
-   %0 = "stablehlo.dynamic_update_slice"(%arg0, %arg1, %arg2) : (tensor<16xf32>, tensor<4xf32>, tensor<i64>) -> tensor<16xf32>
-   func.return %0 : tensor<16xf32>
- }
- 
- // CHECK-LABEL: "op_einsum"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_einsum(%arg0: tensor<8x16xf32>, %arg1: tensor<16x8xf32>) -> tensor<8x8xf32> {
--  //      CHECK: "vhlo.einsum_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.einsum_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   einsum_config = #vhlo.string_v1<"ab,bc->ac">
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<8x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<8x8x!vhlo.f32_v1>
-   %0 = "stablehlo.einsum"(%arg0, %arg1) {
-@@ -1286,22 +1380,25 @@
- }
- 
- // CHECK-LABEL: "op_exponential_minus_one"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_exponential_minus_one(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.exponential_minus_one_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.exponential_minus_one_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.exponential_minus_one"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_exponential"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_exponential(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.exponential_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.exponential_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.exponential"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_fft"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_fft(%arg0: tensor<16xcomplex<f32>>) -> tensor<16xcomplex<f32>> {
--  //      CHECK: "vhlo.fft_v1"(%arg0) <{
-+  //      CHECK: "vhlo.fft_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   fft_length = #vhlo.tensor_v1<dense<16> : tensor<1xi64>>,
-   // CHECK-SAME:   fft_type = #vhlo<fft_type_v1 FFT>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<16x!vhlo.complex_v1<!vhlo.f32_v1>>) -> !vhlo.tensor_v1<16x!vhlo.complex_v1<!vhlo.f32_v1>>
-@@ -1313,8 +1410,9 @@
- }
- 
- // CHECK-LABEL: "op_floor"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_floor(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.floor_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.floor_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.floor"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
-@@ -1327,16 +1425,17 @@
-   // CHECK-SAME:   sym_name = #vhlo.string_v1<"op_func">,
-   // CHECK-SAME:   sym_visibility = #vhlo.string_v1<"private">
-   // CHECK-SAME: }> ({
--  // CHECK-NEXT:   ^[[BB:bb.*]](%arg0: !vhlo.tensor_v1<!vhlo.f32_v1>):
--  // CHECK-NEXT:     "vhlo.return_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-+  // CHECK-NEXT:   ^[[BB:bb.*]](%[[ARG0:.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
-+  // CHECK-NEXT:     "vhlo.return_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-   // CHECK-NEXT: }) : () -> ()
- 
-   func.return %arg0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_gather"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_gather(%arg0 : tensor<2x4x9xf32>, %arg1 : tensor<1x5x2xi32>) -> tensor<1x5x1xf32> {
--  //      CHECK: "vhlo.gather_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.gather_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   collapsed_slice_dims = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-   // CHECK-SAME:   index_vector_dim = #vhlo.integer_v1<2 : i64>,
-   // CHECK-SAME:   indices_are_sorted = #vhlo.bool_v1<true>,
-@@ -1358,8 +1457,9 @@
- }
- 
- // CHECK-LABEL: "op_get_dimension_size"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_get_dimension_size(%arg0: tensor<?xf32>) -> tensor<i32> {
--  //      CHECK: "vhlo.get_dimension_size_v1"(%arg0) <{
-+  //      CHECK: "vhlo.get_dimension_size_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   dimension = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<?x!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-   %0 = "stablehlo.get_dimension_size"(%arg0) {
-@@ -1369,8 +1469,9 @@
- }
- 
- // CHECK-LABEL: "op_get_tuple_element"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_get_tuple_element(%arg0: tuple<tensor<f32>, tensor<i32>>) -> tensor<f32> {
--  //      CHECK: "vhlo.get_tuple_element_v1"(%arg0) <{
-+  //      CHECK: "vhlo.get_tuple_element_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   index = #vhlo.integer_v1<0 : i32>
-   // CHECK-SAME: }> : (!vhlo.tuple_v1<!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.get_tuple_element"(%arg0) {
-@@ -1380,11 +1481,12 @@
- }
- 
- // CHECK-LABEL: "op_if"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_if(%arg0: tensor<i1>, %arg1: tensor<f32>, %arg2: tensor<f32>) -> tensor<f32> {
--  //      CHECK: "vhlo.if_v1"(%arg0) ({
--  // CHECK-NEXT:   "vhlo.return_v1"(%arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-+  //      CHECK: "vhlo.if_v1"(%[[ARG0]]) ({
-+  // CHECK-NEXT:   "vhlo.return_v1"(%[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-   // CHECK-NEXT: }, {
--  // CHECK-NEXT:   "vhlo.return_v1"(%arg2) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-+  // CHECK-NEXT:   "vhlo.return_v1"(%[[ARG2]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-   // CHECK-NEXT: }) : (!vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.if"(%arg0) ({
-     "stablehlo.return"(%arg1) : (tensor<f32>) -> ()
-@@ -1395,15 +1497,17 @@
- }
- 
- // CHECK-LABEL: "op_imag"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_imag(%arg0: tensor<complex<f32>>) -> tensor<f32> {
--  // CHECK: "vhlo.imag_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.imag_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.imag"(%arg0) : (tensor<complex<f32>>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_infeed"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_infeed(%arg0: !stablehlo.token) -> (tensor<f32>, !stablehlo.token) {
--  //               CHECK: "vhlo.infeed_v1"(%arg0) <{
-+  //               CHECK: "vhlo.infeed_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   infeed_config = #vhlo.string_v1<"foo">,
-   // CHECK-SAME{LITERAL}:   layout = #vhlo.array_v1<[#vhlo.array_v1<[]>]>
-   //          CHECK-SAME: }> : (!vhlo.token_v1) -> (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.token_v1)
-@@ -1426,36 +1530,41 @@
- }
- 
- // CHECK-LABEL: "op_is_finite"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_is_finite(%arg0: tensor<f32>) -> tensor<i1> {
--  // CHECK: "vhlo.is_finite_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-+  // CHECK: "vhlo.is_finite_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-   %0 = "stablehlo.is_finite"(%arg0) : (tensor<f32>) -> tensor<i1>
-   func.return %0 : tensor<i1>
- }
- 
- // CHECK-LABEL: "op_log"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_log(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.log_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.log_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.log"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_log_plus_one"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_log_plus_one(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.log_plus_one_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.log_plus_one_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.log_plus_one"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_logistic"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_logistic(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.logistic_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.logistic_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.logistic"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_map"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_map(%arg0: tensor<16xf32>) -> tensor<16xf32> {
--  //      CHECK: "vhlo.map_v1"(%arg0) <{
-+  //      CHECK: "vhlo.map_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   dimensions = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>
-   // CHECK-SAME: }> ({
-   // CHECK-NEXT:   ^[[BB:bb.*]](%[[ARG1:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
-@@ -1473,57 +1582,65 @@
- }
- 
- // CHECK-LABEL: "op_maximum"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_maximum(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.maximum_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.maximum_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.maximum"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_minimum"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_minimum(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.minimum_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.minimum_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.minimum"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_multiply"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_multiply(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.multiply_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.multiply_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.multiply"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_negate"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_negate(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.negate_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.negate_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.negate"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_not"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_not(%arg0: tensor<i1>) -> tensor<i1> {
--  // CHECK: "vhlo.not_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-+  // CHECK: "vhlo.not_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-   %0 = "stablehlo.not"(%arg0) : (tensor<i1>) -> tensor<i1>
-   func.return %0 : tensor<i1>
- }
- 
- // CHECK-LABEL: "op_optimization_barrier"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_optimization_barrier(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.optimization_barrier_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.optimization_barrier_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.optimization_barrier"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_or"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_or(%arg0: tensor<i1>, %arg1: tensor<i1>) -> tensor<i1> {
--  // CHECK: "vhlo.or_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-+  // CHECK: "vhlo.or_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-   %0 = "stablehlo.or"(%arg0, %arg1) : (tensor<i1>, tensor<i1>) -> tensor<i1>
-   func.return %0 : tensor<i1>
- }
- 
- // CHECK-LABEL: "op_outfeed"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_outfeed(%arg0: tensor<f32>, %arg1: !stablehlo.token) -> !stablehlo.token {
--  //      CHECK: "vhlo.outfeed_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.outfeed_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   outfeed_config = #vhlo.string_v1<"foo">
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.token_v1) -> !vhlo.token_v1
-   %0 = "stablehlo.outfeed"(%arg0, %arg1) {
-@@ -1533,8 +1650,9 @@
- }
- 
- // CHECK-LABEL: "op_pad"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_pad(%arg0: tensor<8xf32>, %arg1: tensor<f32>) -> tensor<16xf32> {
--  //      CHECK: "vhlo.pad_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.pad_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   edge_padding_high = #vhlo.tensor_v1<dense<4> : tensor<1xi64>>,
-   // CHECK-SAME:   edge_padding_low = #vhlo.tensor_v1<dense<4> : tensor<1xi64>>,
-   // CHECK-SAME:   interior_padding = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>
-@@ -1548,36 +1666,41 @@
- }
- 
- // CHECK-LABEL: "op_popcnt"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_popcnt(%arg0: tensor<i32>) -> tensor<i32> {
--  // CHECK: "vhlo.popcnt_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-+  // CHECK: "vhlo.popcnt_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-   %0 = "stablehlo.popcnt"(%arg0) : (tensor<i32>) -> tensor<i32>
-   func.return %0 : tensor<i32>
- }
- 
- // CHECK-LABEL: "op_power"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_power(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.power_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.power_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.power"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_real_dynamic_slice"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}}, %[[ARG3:.*]]: {{.*}})
- func.func @op_real_dynamic_slice(%arg0: tensor<?xf32>, %arg1: tensor<1xindex>, %arg2: tensor<1xindex>, %arg3: tensor<1xindex>) -> tensor<?xf32> {
--  // CHECK: "vhlo.real_dynamic_slice_v1"(%arg0, %arg1, %arg2, %arg3) : (!vhlo.tensor_v1<?x!vhlo.f32_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>) -> !vhlo.tensor_v1<?x!vhlo.f32_v1>
-+  // CHECK: "vhlo.real_dynamic_slice_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]], %[[ARG3]]) : (!vhlo.tensor_v1<?x!vhlo.f32_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>) -> !vhlo.tensor_v1<?x!vhlo.f32_v1>
-   %0 = "stablehlo.real_dynamic_slice"(%arg0, %arg1, %arg2, %arg3) : (tensor<?xf32>, tensor<1xindex>, tensor<1xindex>, tensor<1xindex>) -> tensor<?xf32>
-   func.return %0 : tensor<?xf32>
- }
- 
- // CHECK-LABEL: "op_real"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_real(%arg0: tensor<complex<f32>>) -> tensor<f32> {
--  // CHECK: "vhlo.real_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.real_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.real"(%arg0) : (tensor<complex<f32>>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_recv"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_recv(%arg0: !stablehlo.token) -> (tensor<f32>, !stablehlo.token) {
--  //      CHECK: "vhlo.recv_v1"(%arg0) <{
-+  //      CHECK: "vhlo.recv_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME:   channel_type = #vhlo.integer_v1<3 : i64>,
-   // CHECK-SAME:   is_host_transfer = #vhlo.bool_v1<true>
-@@ -1590,8 +1713,9 @@
- }
- 
- // CHECK-LABEL: "op_reduce"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_reduce(%arg0: tensor<16xf32>, %arg1: tensor<f32>) -> tensor<f32> {
--  //  CHECK: "vhlo.reduce_v1"(%arg0, %arg1)
-+  //  CHECK: "vhlo.reduce_v1"(%[[ARG0]], %[[ARG1]])
-   //  CHECK:   ^[[BB:bb.*]](%[[ARG1:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>, %[[ARG2:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
-   //  CHECK:     "vhlo.return_v1"(%[[VAL1:.*]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-   //  CHECK: }) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-@@ -1606,8 +1730,9 @@
- }
- 
- // CHECK-LABEL: "op_reduce_precision"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_reduce_precision(%arg0: tensor<f32>) -> tensor<f32> {
--  //      CHECK: "vhlo.reduce_precision_v1"(%arg0) <{
-+  //      CHECK: "vhlo.reduce_precision_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   exponent_bits = #vhlo.integer_v1<8 : i32>
-   // CHECK-SAME:   mantissa_bits = #vhlo.integer_v1<10 : i32>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-@@ -1619,8 +1744,9 @@
- }
- 
- // CHECK-LABEL: "op_reduce_scatter"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_reduce_scatter(%arg0: tensor<16xf32>) -> tensor<16xf32> {
--  //               CHECK: "vhlo.reduce_scatter_v1"(%arg0) <{
-+  //               CHECK: "vhlo.reduce_scatter_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0], [1]]> : tensor<2x1xi64>>,
-   //          CHECK-SAME:   scatter_dimension = #vhlo.integer_v1<0 : i64>
-@@ -1644,8 +1770,9 @@
- }
- 
- // CHECK-LABEL: "op_reduce_window"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_reduce_window(%arg0: tensor<2x17x31x7xf32>, %arg1: tensor<f32>) -> tensor<2x9x16x7xf32> {
--  //               CHECK: "vhlo.reduce_window_v1"(%arg0, %arg1) <{
-+  //               CHECK: "vhlo.reduce_window_v1"(%[[ARG0]], %[[ARG1]]) <{
-   //          CHECK-SAME:   base_dilations = #vhlo.tensor_v1<dense<[1, 2, 2, 1]> : tensor<4xi64>>,
-   // CHECK-SAME{LITERAL}:   padding = #vhlo.tensor_v1<dense<[[0, 0], [2, 0], [0, 2], [0, 0]]> : tensor<4x2xi64>>,
-   //          CHECK-SAME:   window_dilations = #vhlo.tensor_v1<dense<[1, 2, 2, 1]> : tensor<4xi64>>,
-@@ -1671,8 +1798,9 @@
- }
- 
- // CHECK-LABEL: "op_remainder"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_remainder(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.remainder_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.remainder_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.remainder"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
-@@ -1692,16 +1820,18 @@
- }
- 
- // CHECK-LABEL: "op_reshape"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_reshape(%arg0: tensor<16xf32>) -> tensor<4x4xf32> {
--  // CHECK: "vhlo.reshape_v1"(%arg0) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<4x4x!vhlo.f32_v1>
-+  // CHECK: "vhlo.reshape_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<4x4x!vhlo.f32_v1>
-   %0 = "stablehlo.reshape"(%arg0) : (tensor<16xf32>) -> tensor<4x4xf32>
-   func.return %0 : tensor<4x4xf32>
- }
- 
- // CHECK-LABEL: "op_return"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_return(%arg0: tensor<i32>, %arg1: tensor<f32>) -> tensor<f32> {
--  //      CHECK: "vhlo.case_v1"(%arg0) ({
--  // CHECK-NEXT:   "vhlo.return_v1"(%arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-+  //      CHECK: "vhlo.case_v1"(%[[ARG0]]) ({
-+  // CHECK-NEXT:   "vhlo.return_v1"(%[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-   // CHECK-NEXT: }) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.case"(%arg0) ({
-     "stablehlo.return"(%arg1) : (tensor<f32>) -> ()
-@@ -1710,8 +1840,9 @@
- }
- 
- // CHECK-LABEL: "op_reverse"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_reverse(%arg0: tensor<16xf32>) -> tensor<16xf32> {
--  //      CHECK: "vhlo.reverse_v1"(%arg0) <{
-+  //      CHECK: "vhlo.reverse_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   dimensions = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x!vhlo.f32_v1>
-   %0 = "stablehlo.reverse"(%arg0) {
-@@ -1721,8 +1852,9 @@
- }
- 
- // CHECK-LABEL: "op_rng_bit_generator"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_rng_bit_generator(%arg0: tensor<f32>) -> (tensor<f32>, tensor<f32>) {
--  //      CHECK: "vhlo.rng_bit_generator_v1"(%arg0) <{
-+  //      CHECK: "vhlo.rng_bit_generator_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   rng_algorithm = #vhlo<rng_algorithm_v1 PHILOX>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>)
-   %0:2 = "stablehlo.rng_bit_generator"(%arg0) {
-@@ -1732,8 +1864,9 @@
- }
- 
- // CHECK-LABEL: "op_rng"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_rng(%arg0: tensor<f32>, %arg1: tensor<f32>, %arg2: tensor<0xindex>) -> tensor<f32> {
--  //      CHECK: "vhlo.rng_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.rng_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   rng_distribution = #vhlo<rng_distribution_v1 NORMAL>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<0x!vhlo.index_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.rng"(%arg0, %arg1, %arg2) {
-@@ -1743,29 +1876,33 @@
- }
- 
- // CHECK-LABEL: "op_round_nearest_afz"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_round_nearest_afz(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.round_nearest_afz_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.round_nearest_afz_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.round_nearest_afz"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_round_nearest_even"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_round_nearest_even(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.round_nearest_even_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.round_nearest_even_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.round_nearest_even"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_rsqrt"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_rsqrt(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.rsqrt_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.rsqrt_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.rsqrt"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_scatter"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_scatter(%arg0: tensor<200x100x300xf32>, %arg1: tensor<10x2xi32>, %arg2: tensor<10x300xf32>) -> tensor<200x100x300xf32> {
--  //      CHECK: "vhlo.scatter_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.scatter_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   index_vector_dim = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   indices_are_sorted = #vhlo.bool_v1<true>,
-   // CHECK-SAME:   inserted_window_dims = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-@@ -1795,8 +1932,9 @@
- }
- 
- // CHECK-LABEL: "op_select_and_scatter"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_select_and_scatter(%arg0: tensor<10x24x24x64xf32>, %arg1: tensor<12x13x13x66xf32>, %arg2: tensor<f32>) -> tensor<10x24x24x64xf32> {
--  //      CHECK: "vhlo.select_and_scatter_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.select_and_scatter_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   padding = #vhlo.tensor_v1<dense<1> : tensor<4x2xi64>>,
-   // CHECK-SAME:   window_dimensions = #vhlo.tensor_v1<dense<[1, 2, 2, 1]> : tensor<4xi64>>,
-   // CHECK-SAME:   window_strides = #vhlo.tensor_v1<dense<[1, 2, 2, 1]> : tensor<4xi64>>
-@@ -1826,15 +1964,17 @@
- }
- 
- // CHECK-LABEL: "op_select"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_select(%arg0: tensor<i1>, %arg1: tensor<f32>, %arg2: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.select_v1"(%arg0, %arg1, %arg2) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.select_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.select"(%arg0, %arg1, %arg2) : (tensor<i1>, tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_send"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_send(%arg0: tensor<f32>, %arg1: !stablehlo.token) -> !stablehlo.token {
--  //      CHECK: "vhlo.send_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.send_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME:   channel_type = #vhlo.integer_v1<2 : i64>,
-   // CHECK-SAME:   is_host_transfer = #vhlo.bool_v1<true>
-@@ -1847,8 +1987,9 @@
- }
- 
- // CHECK-LABEL: "op_set_dimension_size"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_set_dimension_size(%arg0: tensor<?xf32>, %arg1: tensor<i32>) -> tensor<16xf32> {
--  //      CHECK: "vhlo.set_dimension_size_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.set_dimension_size_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   dimension = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<?x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<16x!vhlo.f32_v1>
-   %0 = "stablehlo.set_dimension_size"(%arg0, %arg1) {
-@@ -1858,43 +1999,49 @@
- }
- 
- // CHECK-LABEL: "op_shift_left"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_shift_left(%arg0: tensor<i32>, %arg1: tensor<i32>) -> tensor<i32> {
--  // CHECK: "vhlo.shift_left_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.i32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-+  // CHECK: "vhlo.shift_left_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-   %0 = "stablehlo.shift_left"(%arg0, %arg1) : (tensor<i32>, tensor<i32>) -> tensor<i32>
-   func.return %0 : tensor<i32>
- }
- 
- // CHECK-LABEL: "op_shift_right_arithmetic"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_shift_right_arithmetic(%arg0: tensor<i32>, %arg1: tensor<i32>) -> tensor<i32> {
--  // CHECK: "vhlo.shift_right_arithmetic_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.i32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-+  // CHECK: "vhlo.shift_right_arithmetic_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-   %0 = "stablehlo.shift_right_arithmetic"(%arg0, %arg1) : (tensor<i32>, tensor<i32>) -> tensor<i32>
-   func.return %0 : tensor<i32>
- }
- 
- // CHECK-LABEL: "op_shift_right_logical"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_shift_right_logical(%arg0: tensor<i32>, %arg1: tensor<i32>) -> tensor<i32> {
--  // CHECK: "vhlo.shift_right_logical_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.i32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-+  // CHECK: "vhlo.shift_right_logical_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-   %0 = "stablehlo.shift_right_logical"(%arg0, %arg1) : (tensor<i32>, tensor<i32>) -> tensor<i32>
-   func.return %0 : tensor<i32>
- }
- 
- // CHECK-LABEL: "op_sign"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_sign(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.sign_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.sign_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.sign"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_sine"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_sine(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.sine_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.sine_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.sine"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_slice"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_slice(%arg0: tensor<16xf32>) -> tensor<4xf32> {
--  //      CHECK: "vhlo.slice_v1"(%arg0) <{
-+  //      CHECK: "vhlo.slice_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   limit_indices = #vhlo.tensor_v1<dense<4> : tensor<1xi64>>,
-   // CHECK-SAME:   start_indices = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>,
-   // CHECK-SAME:   strides = #vhlo.tensor_v1<dense<1> : tensor<1xi64>>
-@@ -1908,8 +2055,9 @@
- }
- 
- // CHECK-LABEL: "op_sort"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_sort(%arg0: tensor<16xf32>) -> tensor<16xf32> {
--  //      CHECK: "vhlo.sort_v1"(%arg0) <{
-+  //      CHECK: "vhlo.sort_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   dimension = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME:   is_stable = #vhlo.bool_v1<true>
-   // CHECK-SAME: }> ({
-@@ -1929,29 +2077,33 @@
- }
- 
- // CHECK-LABEL: "op_sqrt"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_sqrt(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.sqrt_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.sqrt_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.sqrt"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_subtract"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_subtract(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.subtract_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.subtract_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.subtract"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_tanh"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_tanh(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.tanh_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.tanh_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.tanh"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_torch_index_select"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_torch_index_select(%arg0: tensor<5x1x5xf32>, %arg1: tensor<2xi32>) ->  tensor<2x1x5xf32> {
--  //      CHECK: "vhlo.torch_index_select_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.torch_index_select_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   batch_dims = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME:   dim = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<5x1x5x!vhlo.f32_v1>, !vhlo.tensor_v1<2x!vhlo.i32_v1>) -> !vhlo.tensor_v1<2x1x5x!vhlo.f32_v1>
-@@ -1963,8 +2115,9 @@
- }
- 
- // CHECK-LABEL: "op_trace"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_trace(%arg0: tensor<f32>) {
--  //      CHECK: "vhlo.trace_v1"(%arg0) <{
-+  //      CHECK: "vhlo.trace_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   tag = #vhlo.string_v1<"foo">
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-   "stablehlo.trace"(%arg0) {
-@@ -1974,8 +2127,9 @@
- }
- 
- // CHECK-LABEL: "op_transpose"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_transpose(%arg0: tensor<16x8xf32>) ->  tensor<8x16xf32> {
--  //      CHECK: "vhlo.transpose_v1"(%arg0) <{
-+  //      CHECK: "vhlo.transpose_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   permutation = #vhlo.tensor_v1<dense<[1, 0]> : tensor<2xi64>>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<16x8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<8x16x!vhlo.f32_v1>
-   %0 = "stablehlo.transpose"(%arg0) {
-@@ -1985,8 +2139,9 @@
- }
- 
- // CHECK-LABEL: "op_triangular_solve"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_triangular_solve(%arg0: tensor<16x16xf32>, %arg1: tensor<16x16xf32>) ->  tensor<16x16xf32> {
--  //      CHECK: "vhlo.triangular_solve_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.triangular_solve_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   left_side = #vhlo.bool_v1<true>,
-   // CHECK-SAME:   lower = #vhlo.bool_v1<true>,
-   // CHECK-SAME:   transpose_a = #vhlo<transpose_v1 NO_TRANSPOSE>,
-@@ -2002,15 +2157,17 @@
- }
- 
- // CHECK-LABEL: "op_tuple"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_tuple(%arg0: tensor<f32>) -> tuple<tensor<f32>> {
--  // CHECK: "vhlo.tuple_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tuple_v1<!vhlo.tensor_v1<!vhlo.f32_v1>>
-+  // CHECK: "vhlo.tuple_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tuple_v1<!vhlo.tensor_v1<!vhlo.f32_v1>>
-   %0 = "stablehlo.tuple"(%arg0) : (tensor<f32>) -> tuple<tensor<f32>>
-   func.return %0 : tuple<tensor<f32>>
- }
- 
- // CHECK-LABEL: "op_unary_einsum"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_unary_einsum(%arg0: tensor<8x16xf32>) -> tensor<8xf32> {
--  //      CHECK: "vhlo.unary_einsum_v1"(%arg0) <{
-+  //      CHECK: "vhlo.unary_einsum_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   einsum_config = #vhlo.string_v1<"ab->a">
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<8x16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<8x!vhlo.f32_v1>
-   %0 = "stablehlo.unary_einsum"(%arg0) {
-@@ -2020,22 +2177,25 @@
- }
- 
- // CHECK-LABEL: "op_uniform_dequantize"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_uniform_dequantize(%arg0: tensor<!quant.uniform<i8:f32, 34.0:16>>) -> tensor<f32> {
--  // CHECK: "vhlo.uniform_dequantize_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.uniform_dequantize_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.uniform_dequantize"(%arg0) : (tensor<!quant.uniform<i8:f32, 34.0:16>>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_uniform_quantize"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_uniform_quantize(%arg0: tensor<f32>) -> tensor<!quant.uniform<i8:f32, 34.0:16>> {
--  // CHECK: "vhlo.uniform_quantize_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>
-+  // CHECK: "vhlo.uniform_quantize_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>
-   %0 = "stablehlo.uniform_quantize"(%arg0) : (tensor<f32>) -> tensor<!quant.uniform<i8:f32, 34.0:16>>
-   func.return %0 : tensor<!quant.uniform<i8:f32, 34.0:16>>
- }
- 
- // CHECK-LABEL: "op_while"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_while(%arg0: tensor<i1>) -> tensor<i1> {
--  //      CHECK: "vhlo.while_v1"(%arg0) ({
-+  //      CHECK: "vhlo.while_v1"(%[[ARG0]]) ({
-   // CHECK-NEXT:   ^[[BB:bb.*]](%[[ARG1:arg.*]]: !vhlo.tensor_v1<!vhlo.bool_v1>):
-   // CHECK-NEXT:     "vhlo.return_v1"(%[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>) -> ()
-   // CHECK-NEXT:   }, {
-@@ -2053,8 +2213,9 @@
- }
- 
- // CHECK-LABEL: "op_xor"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_xor(%arg0: tensor<i1>, %arg1: tensor<i1>) -> tensor<i1> {
--  // CHECK: "vhlo.xor_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-+  // CHECK: "vhlo.xor_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-   %0 = "stablehlo.xor"(%arg0, %arg1) : (tensor<i1>, tensor<i1>) -> tensor<i1>
-   func.return %0 : tensor<i1>
- }
-@@ -2062,183 +2223,209 @@
- // ============ TYPES ============
- 
- // CHECK-LABEL: "type_i1"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_i1(%arg0: tensor<i1>, %arg1: tensor<i1>) -> tensor<i1> {
--  // CHECK: "vhlo.and_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-+  // CHECK: "vhlo.and_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-   %0 = "stablehlo.and"(%arg0, %arg1) : (tensor<i1>, tensor<i1>) -> tensor<i1>
-   func.return %0 : tensor<i1>
- }
- 
- // CHECK-LABEL: "type_i4"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_i4(%arg0: tensor<i4>, %arg1: tensor<i4>) -> tensor<i4> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.i4_v1>, !vhlo.tensor_v1<!vhlo.i4_v1>) -> !vhlo.tensor_v1<!vhlo.i4_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i4_v1>, !vhlo.tensor_v1<!vhlo.i4_v1>) -> !vhlo.tensor_v1<!vhlo.i4_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<i4>, tensor<i4>) -> tensor<i4>
-   func.return %0 : tensor<i4>
- }
- 
- // CHECK-LABEL: "type_i8"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_i8(%arg0: tensor<i8>, %arg1: tensor<i8>) -> tensor<i8> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.i8_v1>, !vhlo.tensor_v1<!vhlo.i8_v1>) -> !vhlo.tensor_v1<!vhlo.i8_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i8_v1>, !vhlo.tensor_v1<!vhlo.i8_v1>) -> !vhlo.tensor_v1<!vhlo.i8_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<i8>, tensor<i8>) -> tensor<i8>
-   func.return %0 : tensor<i8>
- }
- 
- // CHECK-LABEL: "type_i16"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_i16(%arg0: tensor<i16>, %arg1: tensor<i16>) -> tensor<i16> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.i16_v1>, !vhlo.tensor_v1<!vhlo.i16_v1>) -> !vhlo.tensor_v1<!vhlo.i16_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i16_v1>, !vhlo.tensor_v1<!vhlo.i16_v1>) -> !vhlo.tensor_v1<!vhlo.i16_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<i16>, tensor<i16>) -> tensor<i16>
-   func.return %0 : tensor<i16>
- }
- 
- // CHECK-LABEL: "type_i32"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_i32(%arg0: tensor<i32>, %arg1: tensor<i32>) -> tensor<i32> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.i32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<i32>, tensor<i32>) -> tensor<i32>
-   func.return %0 : tensor<i32>
- }
- 
- // CHECK-LABEL: "type_i64"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_i64(%arg0: tensor<i64>, %arg1: tensor<i64>) -> tensor<i64> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.i64_v1>, !vhlo.tensor_v1<!vhlo.i64_v1>) -> !vhlo.tensor_v1<!vhlo.i64_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i64_v1>, !vhlo.tensor_v1<!vhlo.i64_v1>) -> !vhlo.tensor_v1<!vhlo.i64_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<i64>, tensor<i64>) -> tensor<i64>
-   func.return %0 : tensor<i64>
- }
- 
- // CHECK-LABEL: "type_ui4"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_ui4(%arg0: tensor<ui4>, %arg1: tensor<ui4>) -> tensor<ui4> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.ui4_v1>, !vhlo.tensor_v1<!vhlo.ui4_v1>) -> !vhlo.tensor_v1<!vhlo.ui4_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.ui4_v1>, !vhlo.tensor_v1<!vhlo.ui4_v1>) -> !vhlo.tensor_v1<!vhlo.ui4_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<ui4>, tensor<ui4>) -> tensor<ui4>
-   func.return %0 : tensor<ui4>
- }
- 
- // CHECK-LABEL: "type_ui8"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_ui8(%arg0: tensor<ui8>, %arg1: tensor<ui8>) -> tensor<ui8> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.ui8_v1>, !vhlo.tensor_v1<!vhlo.ui8_v1>) -> !vhlo.tensor_v1<!vhlo.ui8_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.ui8_v1>, !vhlo.tensor_v1<!vhlo.ui8_v1>) -> !vhlo.tensor_v1<!vhlo.ui8_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<ui8>, tensor<ui8>) -> tensor<ui8>
-   func.return %0 : tensor<ui8>
- }
- 
- // CHECK-LABEL: "type_ui16"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_ui16(%arg0: tensor<ui16>, %arg1: tensor<ui16>) -> tensor<ui16> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.ui16_v1>, !vhlo.tensor_v1<!vhlo.ui16_v1>) -> !vhlo.tensor_v1<!vhlo.ui16_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.ui16_v1>, !vhlo.tensor_v1<!vhlo.ui16_v1>) -> !vhlo.tensor_v1<!vhlo.ui16_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<ui16>, tensor<ui16>) -> tensor<ui16>
-   func.return %0 : tensor<ui16>
- }
- 
- // CHECK-LABEL: "type_ui32"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_ui32(%arg0: tensor<ui32>, %arg1: tensor<ui32>) -> tensor<ui32> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.ui32_v1>, !vhlo.tensor_v1<!vhlo.ui32_v1>) -> !vhlo.tensor_v1<!vhlo.ui32_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.ui32_v1>, !vhlo.tensor_v1<!vhlo.ui32_v1>) -> !vhlo.tensor_v1<!vhlo.ui32_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<ui32>, tensor<ui32>) -> tensor<ui32>
-   func.return %0 : tensor<ui32>
- }
- 
- // CHECK-LABEL: "type_ui64"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_ui64(%arg0: tensor<ui64>, %arg1: tensor<ui64>) -> tensor<ui64> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.ui64_v1>, !vhlo.tensor_v1<!vhlo.ui64_v1>) -> !vhlo.tensor_v1<!vhlo.ui64_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.ui64_v1>, !vhlo.tensor_v1<!vhlo.ui64_v1>) -> !vhlo.tensor_v1<!vhlo.ui64_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<ui64>, tensor<ui64>) -> tensor<ui64>
-   func.return %0 : tensor<ui64>
- }
- 
- // CHECK-LABEL: "type_f8E4M3FN"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_f8E4M3FN(%arg0: tensor<f8E4M3FN>, %arg1: tensor<f8E4M3FN>) -> tensor<f8E4M3FN> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f8E4M3FN_v1>, !vhlo.tensor_v1<!vhlo.f8E4M3FN_v1>) -> !vhlo.tensor_v1<!vhlo.f8E4M3FN_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f8E4M3FN_v1>, !vhlo.tensor_v1<!vhlo.f8E4M3FN_v1>) -> !vhlo.tensor_v1<!vhlo.f8E4M3FN_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f8E4M3FN>, tensor<f8E4M3FN>) -> tensor<f8E4M3FN>
-   func.return %0 : tensor<f8E4M3FN>
- }
- 
- // CHECK-LABEL: "type_f8E5M2"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_f8E5M2(%arg0: tensor<f8E5M2>, %arg1: tensor<f8E5M2>) -> tensor<f8E5M2> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f8E5M2_v1>, !vhlo.tensor_v1<!vhlo.f8E5M2_v1>) -> !vhlo.tensor_v1<!vhlo.f8E5M2_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f8E5M2_v1>, !vhlo.tensor_v1<!vhlo.f8E5M2_v1>) -> !vhlo.tensor_v1<!vhlo.f8E5M2_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f8E5M2>, tensor<f8E5M2>) -> tensor<f8E5M2>
-   func.return %0 : tensor<f8E5M2>
- }
- 
- // CHECK-LABEL: "type_f8E4M3FNUZ"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_f8E4M3FNUZ(%arg0: tensor<f8E4M3FNUZ>, %arg1: tensor<f8E4M3FNUZ>) -> tensor<f8E4M3FNUZ> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f8E4M3FNUZ_v1>, !vhlo.tensor_v1<!vhlo.f8E4M3FNUZ_v1>) -> !vhlo.tensor_v1<!vhlo.f8E4M3FNUZ_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f8E4M3FNUZ_v1>, !vhlo.tensor_v1<!vhlo.f8E4M3FNUZ_v1>) -> !vhlo.tensor_v1<!vhlo.f8E4M3FNUZ_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f8E4M3FNUZ>, tensor<f8E4M3FNUZ>) -> tensor<f8E4M3FNUZ>
-   func.return %0 : tensor<f8E4M3FNUZ>
- }
- 
- // CHECK-LABEL: "type_f8E5M2FNUZ"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_f8E5M2FNUZ(%arg0: tensor<f8E5M2FNUZ>, %arg1: tensor<f8E5M2FNUZ>) -> tensor<f8E5M2FNUZ> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f8E5M2FNUZ_v1>, !vhlo.tensor_v1<!vhlo.f8E5M2FNUZ_v1>) -> !vhlo.tensor_v1<!vhlo.f8E5M2FNUZ_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f8E5M2FNUZ_v1>, !vhlo.tensor_v1<!vhlo.f8E5M2FNUZ_v1>) -> !vhlo.tensor_v1<!vhlo.f8E5M2FNUZ_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f8E5M2FNUZ>, tensor<f8E5M2FNUZ>) -> tensor<f8E5M2FNUZ>
-   func.return %0 : tensor<f8E5M2FNUZ>
- }
- 
- // CHECK-LABEL: "type_bf16"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_bf16(%arg0: tensor<bf16>, %arg1: tensor<bf16>) -> tensor<bf16> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.bf16_v1>, !vhlo.tensor_v1<!vhlo.bf16_v1>) -> !vhlo.tensor_v1<!vhlo.bf16_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.bf16_v1>, !vhlo.tensor_v1<!vhlo.bf16_v1>) -> !vhlo.tensor_v1<!vhlo.bf16_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<bf16>, tensor<bf16>) -> tensor<bf16>
-   func.return %0 : tensor<bf16>
- }
- 
- // CHECK-LABEL: "type_f16"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_f16(%arg0: tensor<f16>, %arg1: tensor<f16>) -> tensor<f16> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f16_v1>, !vhlo.tensor_v1<!vhlo.f16_v1>) -> !vhlo.tensor_v1<!vhlo.f16_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f16_v1>, !vhlo.tensor_v1<!vhlo.f16_v1>) -> !vhlo.tensor_v1<!vhlo.f16_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f16>, tensor<f16>) -> tensor<f16>
-   func.return %0 : tensor<f16>
- }
- 
- // CHECK-LABEL: "type_f32"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_f32(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "type_f64"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_f64(%arg0: tensor<f64>, %arg1: tensor<f64>) -> tensor<f64> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f64_v1>, !vhlo.tensor_v1<!vhlo.f64_v1>) -> !vhlo.tensor_v1<!vhlo.f64_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f64_v1>, !vhlo.tensor_v1<!vhlo.f64_v1>) -> !vhlo.tensor_v1<!vhlo.f64_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f64>, tensor<f64>) -> tensor<f64>
-   func.return %0 : tensor<f64>
- }
- 
- // CHECK-LABEL: "type_complex_f32"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_complex_f32(%arg0: tensor<complex<f32>>, %arg1: tensor<complex<f32>>) -> tensor<complex<f32>> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>, !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>) -> !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>, !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>) -> !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<complex<f32>>, tensor<complex<f32>>) -> tensor<complex<f32>>
-   func.return %0 : tensor<complex<f32>>
- }
- 
- // CHECK-LABEL: "type_complex_f64"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_complex_f64(%arg0: tensor<complex<f64>>, %arg1: tensor<complex<f64>>) -> tensor<complex<f64>> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f64_v1>>, !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f64_v1>>) -> !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f64_v1>>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f64_v1>>, !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f64_v1>>) -> !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f64_v1>>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<complex<f64>>, tensor<complex<f64>>) -> tensor<complex<f64>>
-   func.return %0 : tensor<complex<f64>>
- }
- 
- // CHECK-LABEL: "type_dynamism_ranked"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @type_dynamism_ranked(%arg0: tensor<?xf32>) -> tensor<?xf32> {
--  // CHECK: "vhlo.abs_v1"(%arg0) : (!vhlo.tensor_v1<?x!vhlo.f32_v1>) -> !vhlo.tensor_v1<?x!vhlo.f32_v1>
-+  // CHECK: "vhlo.abs_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<?x!vhlo.f32_v1>) -> !vhlo.tensor_v1<?x!vhlo.f32_v1>
-   %0 = "stablehlo.abs"(%arg0) : (tensor<?xf32>) -> tensor<?xf32>
-   func.return %0 : tensor<?xf32>
- }
- 
- // CHECK-LABEL: "type_quantization"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_quantization(%arg0: tensor<!quant.uniform<i8:f32, 34.0:16>>, %arg1: tensor<!quant.uniform<i8:f32, 34.0:16>>) -> tensor<!quant.uniform<i8:f32, 34.0:16>> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>, !vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>) -> !vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>, !vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>) -> !vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<!quant.uniform<i8:f32, 34.0:16>>, tensor<!quant.uniform<i8:f32, 34.0:16>>) -> tensor<!quant.uniform<i8:f32, 34.0:16>>
-   func.return %0 : tensor<!quant.uniform<i8:f32, 34.0:16>>
- }
- 
- //       CHECK: function_type = #vhlo.type_v1<!vhlo.func_v1<(!vhlo.token_v1) -> !vhlo.token_v1>>
- // CHECK-LABEL: "type_token_callee"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @type_token_callee(%arg0: !stablehlo.token) -> !stablehlo.token {
--  // CHECK: "vhlo.return_v1"(%arg0) : (!vhlo.token_v1) -> ()
-+  // CHECK: "vhlo.return_v1"(%[[ARG0]]) : (!vhlo.token_v1) -> ()
-   return %arg0 : !stablehlo.token
- }
- 
- //       CHECK: function_type = #vhlo.type_v1<!vhlo.func_v1<(!vhlo.token_v1) -> !vhlo.token_v1>>
- // CHECK-LABEL: "type_token_caller"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @type_token_caller(%arg0: !stablehlo.token) -> !stablehlo.token {
--  // CHECK:      "vhlo.call_v1"(%arg0) <{callee = #vhlo.string_v1<"type_token_callee">}
-+  // CHECK:      "vhlo.call_v1"(%[[ARG0]]) <{callee = #vhlo.string_v1<"type_token_callee">}
-   // CHECK-SAME: (!vhlo.token_v1) -> !vhlo.token_v1
-   %0 = func.call @type_token_callee(%arg0) : (!stablehlo.token) -> !stablehlo.token
-   return %0 : !stablehlo.token
- }
- 
- // CHECK-LABEL: "type_tuple"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @type_tuple(%arg0: tuple<tensor<f32>>) -> tuple<!stablehlo.token> {
-   %0 = "stablehlo.custom_call"(%arg0) {
-     call_target_name = "foo"
-diff --ruN a/stablehlo/stablehlo/tests/vhlo/stablehlo_legalize_to_vhlo.0_11_0.mlir b/stablehlo/stablehlo/tests/vhlo/stablehlo_legalize_to_vhlo.0_11_0.mlir
---- stablehlo/stablehlo/tests/vhlo/stablehlo_legalize_to_vhlo.0_11_0.mlir
-+++ stablehlo/stablehlo/tests/vhlo/stablehlo_legalize_to_vhlo.0_11_0.mlir
-@@ -13,6 +13,7 @@
- // ============ ATTRIBUTES ============
- 
- // CHECK-LABEL: "attr_comparison_direction_eq"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_direction_eq(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     // CHECK: comparison_direction = #vhlo<comparison_direction_v1 EQ>
-@@ -22,6 +23,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_direction_ne"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_direction_ne(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     // CHECK: comparison_direction = #vhlo<comparison_direction_v1 NE>
-@@ -31,6 +33,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_direction_ge"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_direction_ge(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     // CHECK: comparison_direction = #vhlo<comparison_direction_v1 GE>
-@@ -40,6 +43,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_direction_gt"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_direction_gt(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     // CHECK: comparison_direction = #vhlo<comparison_direction_v1 GT>
-@@ -49,6 +53,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_direction_le"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_direction_le(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     // CHECK: comparison_direction = #vhlo<comparison_direction_v1 LE>
-@@ -58,6 +63,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_direction_lt"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_direction_lt(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     // CHECK: comparison_direction = #vhlo<comparison_direction_v1 LT>
-@@ -67,6 +73,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_type_notype"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_type_notype(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     comparison_direction = #stablehlo<comparison_direction EQ>
-@@ -76,6 +83,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_type_float"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_type_float(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     comparison_direction = #stablehlo<comparison_direction EQ>,
-@@ -86,6 +94,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_type_totalorder"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_type_totalorder(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     comparison_direction = #stablehlo<comparison_direction EQ>,
-@@ -96,6 +105,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_type_signed"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_type_signed(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     comparison_direction = #stablehlo<comparison_direction EQ>,
-@@ -106,6 +116,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_type_unsigned"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_type_unsigned(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     comparison_direction = #stablehlo<comparison_direction EQ>,
-@@ -118,6 +129,7 @@
- // ConvDimensionNumbers aka #stablehlo.conv is covered below.
- 
- // CHECK-LABEL: "attr_custom_call_api_version_unspecified"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_custom_call_api_version_unspecified(%arg0: tensor<f32>) -> tensor<f32> {
-   %0 = "stablehlo.custom_call"(%arg0) {
-     call_target_name = "foo",
-@@ -128,6 +140,7 @@
- }
- 
- // CHECK-LABEL: "attr_custom_call_api_version_original"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_custom_call_api_version_original(%arg0: tensor<f32>) -> tensor<f32> {
-   %0 = "stablehlo.custom_call"(%arg0) {
-     call_target_name = "foo",
-@@ -138,6 +151,7 @@
- }
- 
- // CHECK-LABEL: "attr_custom_call_api_version_status_returning"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_custom_call_api_version_status_returning(%arg0: tensor<f32>) -> tensor<f32> {
-   %0 = "stablehlo.custom_call"(%arg0) {
-     call_target_name = "foo",
-@@ -148,6 +162,7 @@
- }
- 
- // CHECK-LABEL: "attr_custom_call_api_version_status_returning_unified"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_custom_call_api_version_status_returning_unified(%arg0: tensor<f32>) -> tensor<f32> {
-   %0 = "stablehlo.custom_call"(%arg0) {
-     call_target_name = "foo",
-@@ -160,6 +175,7 @@
- // DotDimensionNumbers aka #stablehlo.dot is covered below.
- 
- // CHECK-LABEL: "attr_fft_type_fft"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_fft_type_fft(%arg0: tensor<16xcomplex<f32>>) -> tensor<16xcomplex<f32>> {
-   %0 = "stablehlo.fft"(%arg0) {
-     // CHECK: fft_type = #vhlo<fft_type_v1 FFT>
-@@ -170,6 +186,7 @@
- }
- 
- // CHECK-LABEL: "attr_fft_type_ifft"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_fft_type_ifft(%arg0: tensor<16xcomplex<f32>>) -> tensor<16xcomplex<f32>> {
-   %0 = "stablehlo.fft"(%arg0) {
-     // CHECK: fft_type = #vhlo<fft_type_v1 IFFT>
-@@ -180,6 +197,7 @@
- }
- 
- // CHECK-LABEL: "attr_fft_type_rfft"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_fft_type_rfft(%arg0: tensor<16xf32>) -> tensor<9xcomplex<f32>> {
-   %0 = "stablehlo.fft"(%arg0) {
-     // CHECK: fft_type = #vhlo<fft_type_v1 RFFT>
-@@ -190,6 +208,7 @@
- }
- 
- // CHECK-LABEL: "attr_fft_type_irfft"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_fft_type_irfft(%arg0: tensor<9xcomplex<f32>>) -> tensor<16xf32> {
-   %0 = "stablehlo.fft"(%arg0) {
-     // CHECK: fft_type = #vhlo<fft_type_v1 IRFFT>
-@@ -202,6 +221,7 @@
- // GatherDimensionNumbers aka #stablehlo.gather is covered below.
- 
- // CHECK-LABEL: "attr_precision_config_default"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_precision_config_default(%arg0: tensor<8x16xf32>, %arg1: tensor<16x8xf32>) -> tensor<8x8xf32> {
-   %0 = "stablehlo.dot"(%arg0, %arg1) {
-     // CHECK: precision_config = #vhlo.array_v1<[#vhlo<precision_v1 DEFAULT>, #vhlo<precision_v1 DEFAULT>]>
-@@ -210,6 +230,7 @@
- }
- 
- // CHECK-LABEL: "attr_precision_config_high"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_precision_config_high(%arg0: tensor<8x16xf32>, %arg1: tensor<16x8xf32>) -> tensor<8x8xf32> {
-   %0 = "stablehlo.dot"(%arg0, %arg1) {
-     // CHECK: precision_config = #vhlo.array_v1<[#vhlo<precision_v1 HIGH>, #vhlo<precision_v1 HIGH>]>
-@@ -219,6 +240,7 @@
- }
- 
- // CHECK-LABEL: "attr_precision_config_highest"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_precision_config_highest(%arg0: tensor<8x16xf32>, %arg1: tensor<16x8xf32>) -> tensor<8x8xf32> {
-   %0 = "stablehlo.dot"(%arg0, %arg1) {
-     // CHECK: precision_config = #vhlo.array_v1<[#vhlo<precision_v1 HIGHEST>, #vhlo<precision_v1 HIGHEST>]>
-@@ -228,6 +250,7 @@
- }
- 
- // CHECK-LABEL: "attr_rng_algorithm_default"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_rng_algorithm_default(%arg0: tensor<f32>) -> (tensor<f32>, tensor<f32>) {
-   %0:2 = "stablehlo.rng_bit_generator"(%arg0) {
-     // CHECK: rng_algorithm = #vhlo<rng_algorithm_v1 DEFAULT>
-@@ -237,6 +260,7 @@
- }
- 
- // CHECK-LABEL: "attr_rng_algorithm_three_fry"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_rng_algorithm_three_fry(%arg0: tensor<f32>) -> (tensor<f32>, tensor<f32>) {
-   %0:2 = "stablehlo.rng_bit_generator"(%arg0) {
-     // CHECK: rng_algorithm = #vhlo<rng_algorithm_v1 THREE_FRY>
-@@ -246,6 +270,7 @@
- }
- 
- // CHECK-LABEL: "attr_rng_algorithm_philox"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_rng_algorithm_philox(%arg0: tensor<f32>) -> (tensor<f32>, tensor<f32>) {
-   %0:2 = "stablehlo.rng_bit_generator"(%arg0) {
-     // CHECK: rng_algorithm = #vhlo<rng_algorithm_v1 PHILOX>
-@@ -255,6 +280,7 @@
- }
- 
- // CHECK-LABEL: "attr_rng_distribution_uniform"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @attr_rng_distribution_uniform(%arg0: tensor<f32>, %arg1: tensor<f32>, %arg2: tensor<0xindex>) -> tensor<f32> {
-   %0 = "stablehlo.rng"(%arg0, %arg1, %arg2) {
-     // CHECK: rng_distribution = #vhlo<rng_distribution_v1 UNIFORM>
-@@ -264,6 +290,7 @@
- }
- 
- // CHECK-LABEL: "attr_rng_distribution_normal"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @attr_rng_distribution_normal(%arg0: tensor<f32>, %arg1: tensor<f32>, %arg2: tensor<0xindex>) -> tensor<f32> {
-   %0 = "stablehlo.rng"(%arg0, %arg1, %arg2) {
-     // CHECK: rng_distribution = #vhlo<rng_distribution_v1 NORMAL>
-@@ -275,6 +302,7 @@
- // ScatterDimensionNumbers aka #stablehlo.scatter is covered below.
- 
- // CHECK-LABEL: "attr_transpose_no_transpose"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_transpose_no_transpose(%arg0: tensor<16x16xf32>, %arg1: tensor<16x16xf32>) ->  tensor<16x16xf32> {
-   %0 = "stablehlo.triangular_solve"(%arg0, %arg1) {
-     left_side = true,
-@@ -287,6 +315,7 @@
- }
- 
- // CHECK-LABEL: "attr_transpose_transpose"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_transpose_transpose(%arg0: tensor<16x16xf32>, %arg1: tensor<16x16xf32>) ->  tensor<16x16xf32> {
-   %0 = "stablehlo.triangular_solve"(%arg0, %arg1) {
-     left_side = true,
-@@ -299,6 +328,7 @@
- }
- 
- // CHECK-LABEL: "attr_transpose_adjoint"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_transpose_adjoint(%arg0: tensor<16x16xf32>, %arg1: tensor<16x16xf32>) ->  tensor<16x16xf32> {
-   %0 = "stablehlo.triangular_solve"(%arg0, %arg1) {
-     left_side = true,
-@@ -313,10 +343,9 @@
- // TypeExtensionsAttr aka #stablehlo.type_extensions is covered below.
- 
- // CHECK-LABEL: "attr_type_extensions_bounds"
--func.func @attr_type_extensions_bounds(
--    %arg0: tensor<?x?xf32, #stablehlo.type_extensions<bounds = [16, ?]>>)
--    -> tensor<?x?xf32, #stablehlo.type_extensions<bounds = [16, ?]>> {
--  // CHECK: "vhlo.return_v1"(%arg0) : (!vhlo.tensor_v1<?x?x!vhlo.f32_v1, #vhlo.type_extensions_v1<bounds = [16, ?]>>) -> ()
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @attr_type_extensions_bounds(%arg0: tensor<?x?xf32, #stablehlo.type_extensions<bounds = [16, ?]>>) -> tensor<?x?xf32, #stablehlo.type_extensions<bounds = [16, ?]>> {
-+  // CHECK: "vhlo.return_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<?x?x!vhlo.f32_v1, #vhlo.type_extensions_v1<bounds = [16, ?]>>) -> ()
-   func.return %arg0 : tensor<?x?xf32, #stablehlo.type_extensions<bounds = [16, ?]>>
- }
- 
-@@ -324,8 +353,9 @@
- // ============ DEFAULTS ============
- 
- // CHECK-LABEL: "default_all_gather"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_all_gather(%arg0: tensor<16x8xf32>) -> tensor<16x16xf32> {
--  //               CHECK: "vhlo.all_gather_v1"(%arg0) <{
-+  //               CHECK: "vhlo.all_gather_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   all_gather_dim = #vhlo.integer_v1<1 : i64>
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0], [1]]> : tensor<2x1xi64>>,
-@@ -339,8 +369,9 @@
- }
- 
- // CHECK-LABEL: "default_all_reduce"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_all_reduce(%arg0: tensor<f32>) -> tensor<f32> {
--  //               CHECK: "vhlo.all_reduce_v1"(%arg0)
-+  //               CHECK: "vhlo.all_reduce_v1"(%[[ARG0]])
-   //          CHECK-SAME: <{
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0], [1]]> : tensor<2x1xi64>>,
-@@ -362,8 +393,9 @@
- }
- 
- // CHECK-LABEL: "default_all_to_all"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_all_to_all(%arg0: tensor<4x16xf32>) -> tensor<16x4xf32> {
--  //               CHECK: "vhlo.all_to_all_v1"(%arg0) <{
-+  //               CHECK: "vhlo.all_to_all_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   //          CHECK-SAME:   concat_dimension = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0, 1, 2, 3]]> : tensor<1x4xi64>>,
-@@ -380,8 +412,9 @@
- }
- 
- // CHECK-LABEL: "default_cholesky"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_cholesky(%arg0: tensor<1x16x16xf32>) -> tensor<1x16x16xf32> {
--  //      CHECK: "vhlo.cholesky_v1"(%arg0) <{
-+  //      CHECK: "vhlo.cholesky_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   lower = #vhlo.bool_v1<false>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<1x16x16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<1x16x16x!vhlo.f32_v1>
-   %0 = "stablehlo.cholesky"(%arg0) : (tensor<1x16x16xf32>) -> tensor<1x16x16xf32>
-@@ -389,8 +422,9 @@
- }
- 
- // CHECK-LABEL: "default_collective_permute"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_collective_permute(%arg0: tensor<16x8xf32>) -> tensor<16x8xf32> {
--  //               CHECK: "vhlo.collective_permute_v1"(%arg0) <{
-+  //               CHECK: "vhlo.collective_permute_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME{LITERAL}:   source_target_pairs = #vhlo.tensor_v1<dense<[[0, 1], [1, 2], [2, 3]]> : tensor<3x2xi64>>
-   //          CHECK-SAME: }> : (!vhlo.tensor_v1<16x8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x8x!vhlo.f32_v1>
-@@ -401,8 +435,9 @@
- }
- 
- // CHECK-LABEL: "default_compare"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @default_compare(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
--  //      CHECK: "vhlo.compare_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.compare_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   compare_type = #vhlo<comparison_type_v1 NOTYPE>,
-   // CHECK-SAME:   comparison_direction = #vhlo<comparison_direction_v1 EQ>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-@@ -413,8 +448,9 @@
- }
- 
- // CHECK-LABEL: "default_convolution"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @default_convolution(%arg0: tensor<1x8x8x207xf32>, %arg1: tensor<3x3x207x16xf32>) -> tensor<1x6x6x16xf32> {
--  //      CHECK: "vhlo.convolution_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.convolution_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   batch_group_count = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   feature_group_count = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   input_batch_dimension = #vhlo.integer_v1<0 : i64>,
-@@ -442,8 +478,9 @@
- }
- 
- // CHECK-LABEL: "default_custom_call"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_custom_call(%arg0: tensor<f32>) -> tensor<f32> {
--  //      CHECK: "vhlo.custom_call_v1"(%arg0) <{
-+  //      CHECK: "vhlo.custom_call_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   api_version = #vhlo<api_version_v1 API_VERSION_ORIGINAL>,
-   // CHECK-SAME:   backend_config = #vhlo.string_v1<"">,
-   // CHECK-SAME:   call_target_name = #vhlo.string_v1<"foo">,
-@@ -460,8 +497,9 @@
- }
- 
- // CHECK-LABEL: "default_dot_general"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @default_dot_general(%arg0: tensor<8x8x16xf32>, %arg1: tensor<8x16x8xf32>) -> tensor<8x8x8xf32> {
--  //      CHECK: "vhlo.dot_general_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.dot_general_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   lhs_batching_dimensions = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>,
-   // CHECK-SAME:   lhs_contracting_dimensions = #vhlo.tensor_v1<dense<2> : tensor<1xi64>>,
-   // CHECK-SAME:   precision_config = #vhlo.array_v1<[#vhlo<precision_v1 DEFAULT>, #vhlo<precision_v1 DEFAULT>]>,
-@@ -480,8 +518,9 @@
- }
- 
- // CHECK-LABEL: "default_dot"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @default_dot(%arg0: tensor<8x16xf32>, %arg1: tensor<16x8xf32>) -> tensor<8x8xf32> {
--  //      CHECK: "vhlo.dot_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.dot_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   precision_config = #vhlo.array_v1<[#vhlo<precision_v1 DEFAULT>, #vhlo<precision_v1 DEFAULT>]>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<8x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<8x8x!vhlo.f32_v1>
-   %0 = "stablehlo.dot"(%arg0, %arg1) : (tensor<8x16xf32>, tensor<16x8xf32>) -> tensor<8x8xf32>
-@@ -489,8 +528,9 @@
- }
- 
- // CHECK-LABEL: "default_dynamic_broadcast_in_dim"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @default_dynamic_broadcast_in_dim(%arg0: tensor<?x?xf32>, %arg1: tensor<2xindex>) -> tensor<?x?xf32> {
--  //      CHECK: "vhlo.dynamic_broadcast_in_dim_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.dynamic_broadcast_in_dim_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   broadcast_dimensions = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-   // CHECK-SAME:   known_expanding_dimensions = #vhlo.tensor_v1<dense<> : tensor<0xi64>>,
-   // CHECK-SAME:   known_nonexpanding_dimensions = #vhlo.tensor_v1<dense<> : tensor<0xi64>>
-@@ -502,8 +542,9 @@
- }
- 
- // CHECK-LABEL: "default_dynamic_conv"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @default_dynamic_conv(%arg0: tensor<1x8x8x207xf32>, %arg1: tensor<3x3x207x16xf32>, %arg2: tensor<4xi32>) -> tensor<1x?x?x16xf32> {
--  //      CHECK: "vhlo.dynamic_conv_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.dynamic_conv_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   batch_group_count = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   feature_group_count = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   input_batch_dimension = #vhlo.integer_v1<0 : i64>,
-@@ -531,8 +572,9 @@
- }
- 
- // CHECK-LABEL: "default_dynamic_gather"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @default_dynamic_gather(%arg0 : tensor<2x4x9xf32>, %arg1 : tensor<1x5x2xi32>, %arg2 : tensor<3xi32>) -> tensor<1x5x8xf32> {
--  //      CHECK: "vhlo.dynamic_gather_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.dynamic_gather_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   collapsed_slice_dims = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-   // CHECK-SAME:   index_vector_dim = #vhlo.integer_v1<2 : i64>,
-   // CHECK-SAME:   indices_are_sorted = #vhlo.bool_v1<false>,
-@@ -558,15 +600,16 @@
-   // CHECK-SAME:   sym_name = #vhlo.string_v1<"default_func">,
-   // CHECK-SAME:   sym_visibility = #vhlo.string_v1<"">
-   // CHECK-SAME: }> ({
--  // CHECK-NEXT:   ^[[BB:bb.*]](%arg0: !vhlo.tensor_v1<!vhlo.f32_v1>):
--  // CHECK-NEXT:     "vhlo.return_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-+  // CHECK-NEXT:   ^[[BB:bb.*]](%[[ARG0:.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
-+  // CHECK-NEXT:     "vhlo.return_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-   // CHECK-NEXT: }) : () -> ()
-   func.return %arg0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "dynamic_gather"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @dynamic_gather(%arg0 : tensor<2x4x9xf32>, %arg1 : tensor<1x5x2xi32>) -> tensor<1x5x1xf32> {
--  //      CHECK: "vhlo.gather_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.gather_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   collapsed_slice_dims = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-   // CHECK-SAME:   index_vector_dim = #vhlo.integer_v1<2 : i64>,
-   // CHECK-SAME:   indices_are_sorted = #vhlo.bool_v1<false>,
-@@ -587,8 +630,9 @@
- }
- 
- // CHECK-LABEL: "default_infeed"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_infeed(%arg0: !stablehlo.token) -> (tensor<f32>, !stablehlo.token) {
--  //               CHECK: "vhlo.infeed_v1"(%arg0) <{
-+  //               CHECK: "vhlo.infeed_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   infeed_config = #vhlo.string_v1<"">,
-   // CHECK-SAME{LITERAL}:   layout = #vhlo.array_v1<[]>
-   //          CHECK-SAME: }> : (!vhlo.token_v1) -> (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.token_v1)
-@@ -597,8 +641,9 @@
- }
- 
- // CHECK-LABEL: "default_outfeed"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @default_outfeed(%arg0: tensor<f32>, %arg1: !stablehlo.token) -> !stablehlo.token {
--  //      CHECK: "vhlo.outfeed_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.outfeed_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   outfeed_config = #vhlo.string_v1<"">
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.token_v1) -> !vhlo.token_v1
-   %0 = "stablehlo.outfeed"(%arg0, %arg1) : (tensor<f32>, !stablehlo.token) -> !stablehlo.token
-@@ -606,8 +651,9 @@
- }
- 
- // CHECK-LABEL: "default_recv"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_recv(%arg0: !stablehlo.token) -> (tensor<f32>, !stablehlo.token) {
--  //      CHECK: "vhlo.recv_v1"(%arg0) <{
-+  //      CHECK: "vhlo.recv_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME:   channel_type = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   is_host_transfer = #vhlo.bool_v1<false>
-@@ -619,8 +665,9 @@
- }
- 
- // CHECK-LABEL: "default_send"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @default_send(%arg0: tensor<f32>, %arg1: !stablehlo.token) -> !stablehlo.token {
--  //      CHECK: "vhlo.send_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.send_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME:   channel_type = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   is_host_transfer = #vhlo.bool_v1<false>
-@@ -632,8 +679,9 @@
- }
- 
- // CHECK-LABEL: "default_reduce_scatter"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_reduce_scatter(%arg0: tensor<16xf32>) -> tensor<16xf32> {
--  //               CHECK: "vhlo.reduce_scatter_v1"(%arg0) <{
-+  //               CHECK: "vhlo.reduce_scatter_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0], [1]]> : tensor<2x1xi64>>,
-   //          CHECK-SAME:   scatter_dimension = #vhlo.integer_v1<0 : i64>
-@@ -655,8 +703,9 @@
- }
- 
- // CHECK-LABEL: "default_reduce_window"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @default_reduce_window(%arg0: tensor<2x17x31x7xf32>, %arg1: tensor<f32>) -> tensor<2x16x30x7xf32> {
--  //               CHECK: "vhlo.reduce_window_v1"(%arg0, %arg1)  <{
-+  //               CHECK: "vhlo.reduce_window_v1"(%[[ARG0]], %[[ARG1]])  <{
-   //          CHECK-SAME:   base_dilations = #vhlo.tensor_v1<dense<1> : tensor<4xi64>>,
-   // CHECK-SAME{LITERAL}:   padding = #vhlo.tensor_v1<dense<0> : tensor<4x2xi64>>,
-   //          CHECK-SAME:   window_dilations = #vhlo.tensor_v1<dense<1> : tensor<4xi64>>,
-@@ -678,8 +727,9 @@
- }
- 
- // CHECK-LABEL: "default_scatter"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @default_scatter(%arg0: tensor<200x100x300xf32>, %arg1: tensor<10x2xi32>, %arg2: tensor<10x300xf32>) -> tensor<200x100x300xf32> {
--  //      CHECK: "vhlo.scatter_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.scatter_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   index_vector_dim = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   indices_are_sorted = #vhlo.bool_v1<false>,
-   // CHECK-SAME:   inserted_window_dims = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-@@ -707,8 +757,9 @@
- }
- 
- // CHECK-LABEL: "default_select_and_scatter"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @default_select_and_scatter(%arg0: tensor<10x24x24x64xf32>, %arg1: tensor<10x23x23x64xf32>, %arg2: tensor<f32>) -> tensor<10x24x24x64xf32> {
--  //      CHECK: "vhlo.select_and_scatter_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.select_and_scatter_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   padding = #vhlo.tensor_v1<dense<0> : tensor<4x2xi64>>,
-   // CHECK-SAME:   window_dimensions = #vhlo.tensor_v1<dense<[1, 2, 2, 1]> : tensor<4xi64>>,
-   // CHECK-SAME:   window_strides = #vhlo.tensor_v1<dense<1> : tensor<4xi64>>
-@@ -736,8 +787,9 @@
- }
- 
- // CHECK-LABEL: "default_sort"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_sort(%arg0: tensor<16xf32>) -> tensor<16xf32> {
--  //      CHECK: "vhlo.sort_v1"(%arg0) <{
-+  //      CHECK: "vhlo.sort_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   dimension = #vhlo.integer_v1<-1 : i64>
-   // CHECK-SAME:   is_stable = #vhlo.bool_v1<false>
-   // CHECK-SAME: }> ({
-@@ -756,29 +808,33 @@
- // ============ OPS ============
- 
- // CHECK-LABEL: "op_abs"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_abs(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.abs_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.abs_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.abs"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_add"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_add(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_after_all"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_after_all(%arg0: !stablehlo.token) -> !stablehlo.token {
--  // CHECK: "vhlo.after_all_v1"(%arg0) : (!vhlo.token_v1) -> !vhlo.token_v1
-+  // CHECK: "vhlo.after_all_v1"(%[[ARG0]]) : (!vhlo.token_v1) -> !vhlo.token_v1
-   %0 = "stablehlo.after_all"(%arg0) : (!stablehlo.token) -> !stablehlo.token
-   func.return %0 : !stablehlo.token
- }
- 
- // CHECK-LABEL: "op_all_gather"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_all_gather(%arg0: tensor<16x8xf32>) -> tensor<16x16xf32> {
--  //               CHECK: "vhlo.all_gather_v1"(%arg0) <{
-+  //               CHECK: "vhlo.all_gather_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   all_gather_dim = #vhlo.integer_v1<1 : i64>
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0], [1]]> : tensor<2x1xi64>>,
-@@ -794,8 +850,9 @@
- }
- 
- // CHECK-LABEL: "op_all_reduce"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_all_reduce(%arg0: tensor<f32>) -> tensor<f32> {
--  //               CHECK: "vhlo.all_reduce_v1"(%arg0) <{
-+  //               CHECK: "vhlo.all_reduce_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0], [1]]> : tensor<2x1xi64>>,
-   //          CHECK-SAME:   use_global_device_ids = #vhlo.bool_v1<true>
-@@ -817,8 +874,9 @@
- }
- 
- // CHECK-LABEL: "op_all_to_all"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_all_to_all(%arg0: tensor<4x16xf32>) -> tensor<16x4xf32> {
--  //               CHECK: "vhlo.all_to_all_v1"(%arg0) <{
-+  //               CHECK: "vhlo.all_to_all_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<1 : i64>,
-   //          CHECK-SAME:   concat_dimension = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0, 1, 2, 3]]> : tensor<1x4xi64>>,
-@@ -836,22 +894,25 @@
- }
- 
- // CHECK-LABEL: "op_and"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_and(%arg0: tensor<i1>, %arg1: tensor<i1>) -> tensor<i1> {
--  // CHECK: "vhlo.and_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-+  // CHECK: "vhlo.and_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-   %0 = "stablehlo.and"(%arg0, %arg1) : (tensor<i1>, tensor<i1>) -> tensor<i1>
-   func.return %0 : tensor<i1>
- }
- 
- // CHECK-LABEL: "op_atan2"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_atan2(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.atan2_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.atan2_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.atan2"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_batch_norm_grad"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}}, %[[ARG3:.*]]: {{.*}}, %[[ARG4:.*]]: {{.*}})
- func.func @op_batch_norm_grad(%arg0: tensor<16x16x16x16xf32>, %arg1: tensor<16xf32>, %arg2: tensor<16xf32>, %arg3: tensor<16xf32>, %arg4: tensor<16x16x16x16xf32>) -> (tensor<16x16x16x16xf32>, tensor<16xf32>, tensor<16xf32>) {
--  //      CHECK: "vhlo.batch_norm_grad_v1"(%arg0, %arg1, %arg2, %arg3, %arg4) <{
-+  //      CHECK: "vhlo.batch_norm_grad_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]], %[[ARG3]], %[[ARG4]]) <{
-   // CHECK-SAME:   epsilon = #vhlo.float_v1<1.000000e-03 : !vhlo.f32_v1>,
-   // CHECK-SAME:   feature_index = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<16x16x16x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x16x16x16x!vhlo.f32_v1>) -> (!vhlo.tensor_v1<16x16x16x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>)
-@@ -863,8 +924,9 @@
- }
- 
- // CHECK-LABEL: "op_batch_norm_inference"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}}, %[[ARG3:.*]]: {{.*}}, %[[ARG4:.*]]: {{.*}})
- func.func @op_batch_norm_inference(%arg0: tensor<16x16x16x16xf32>, %arg1: tensor<16xf32>, %arg2: tensor<16xf32>, %arg3: tensor<16xf32>, %arg4: tensor<16xf32>) -> tensor<16x16x16x16xf32> {
--  //      CHECK: "vhlo.batch_norm_inference_v1"(%arg0, %arg1, %arg2, %arg3, %arg4) <{
-+  //      CHECK: "vhlo.batch_norm_inference_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]], %[[ARG3]], %[[ARG4]]) <{
-   // CHECK-SAME:   epsilon = #vhlo.float_v1<1.000000e-03 : !vhlo.f32_v1>,
-   // CHECK-SAME:   feature_index = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<16x16x16x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x16x16x16x!vhlo.f32_v1>
-@@ -876,8 +938,9 @@
- }
- 
- // CHECK-LABEL: "op_batch_norm_training"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_batch_norm_training(%arg0: tensor<16x16x16x16xf32>, %arg1: tensor<16xf32>, %arg2: tensor<16xf32>) -> (tensor<16x16x16x16xf32>, tensor<16xf32>, tensor<16xf32>) {
--  //      CHECK: "vhlo.batch_norm_training_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.batch_norm_training_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   epsilon = #vhlo.float_v1<1.000000e-03 : !vhlo.f32_v1>,
-   // CHECK-SAME:   feature_index = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<16x16x16x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>) -> (!vhlo.tensor_v1<16x16x16x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>)
-@@ -889,15 +952,17 @@
- }
- 
- // CHECK-LABEL: "op_bitcast_convert"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_bitcast_convert(%arg0: tensor<i32>) -> tensor<f32> {
--  // CHECK: "vhlo.bitcast_convert_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.bitcast_convert_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.bitcast_convert"(%arg0) : (tensor<i32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_broadcast_in_dim"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_broadcast_in_dim(%arg0: tensor<16xf32>) -> tensor<16x16xf32> {
--  //      CHECK: "vhlo.broadcast_in_dim_v1"(%arg0) <{
-+  //      CHECK: "vhlo.broadcast_in_dim_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   broadcast_dimensions = #vhlo.tensor_v1<dense<1> : tensor<1xi64>>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x16x!vhlo.f32_v1>
-   %0 = "stablehlo.broadcast_in_dim"(%arg0) {
-@@ -907,8 +972,9 @@
- }
- 
- // CHECK-LABEL: "op_broadcast"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_broadcast(%arg0: tensor<16xf32>) -> tensor<16x16xf32> {
--  //      CHECK: "vhlo.broadcast_v1"(%arg0) <{
-+  //      CHECK: "vhlo.broadcast_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   broadcast_sizes = #vhlo.tensor_v1<dense<16> : tensor<1xi64>>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x16x!vhlo.f32_v1>
-   %0 = "stablehlo.broadcast"(%arg0) {
-@@ -918,9 +984,10 @@
- }
- 
- // CHECK-LABEL: "op_case"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_case(%arg0: tensor<i32>, %arg1: tensor<f32>) -> tensor<f32> {
--  //      CHECK: "vhlo.case_v1"(%arg0) ({
--  // CHECK-NEXT:   "vhlo.return_v1"(%arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-+  //      CHECK: "vhlo.case_v1"(%[[ARG0]]) ({
-+  // CHECK-NEXT:   "vhlo.return_v1"(%[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-   // CHECK-NEXT: }) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.case"(%arg0) ({
-     "stablehlo.return"(%arg1) : (tensor<f32>) -> ()
-@@ -929,22 +996,25 @@
- }
- 
- // CHECK-LABEL: "op_cbrt"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_cbrt(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.cbrt_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.cbrt_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.cbrt"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_ceil"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_ceil(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.ceil_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.ceil_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.ceil"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_cholesky"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_cholesky(%arg0: tensor<1x16x16xf32>) -> tensor<1x16x16xf32> {
--  //      CHECK: "vhlo.cholesky_v1"(%arg0) <{
-+  //      CHECK: "vhlo.cholesky_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   lower = #vhlo.bool_v1<true>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<1x16x16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<1x16x16x!vhlo.f32_v1>
-   %0 = "stablehlo.cholesky"(%arg0) {
-@@ -954,22 +1024,25 @@
- }
- 
- // CHECK-LABEL: "op_clamp"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_clamp(%arg0: tensor<f32>, %arg1: tensor<f32>, %arg2: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.clamp_v1"(%arg0, %arg1, %arg2) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.clamp_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.clamp"(%arg0, %arg1, %arg2) : (tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_count_leading_zeros"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_count_leading_zeros(%arg0: tensor<i32>) -> tensor<i32> {
--  // CHECK: "vhlo.count_leading_zeros_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-+  // CHECK: "vhlo.count_leading_zeros_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-   %0 = "stablehlo.count_leading_zeros"(%arg0) : (tensor<i32>) -> tensor<i32>
-   func.return %0 : tensor<i32>
- }
- 
- // CHECK-LABEL: "op_collective_permute"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_collective_permute(%arg0: tensor<16x8xf32>) -> tensor<16x8xf32> {
--  //               CHECK: "vhlo.collective_permute_v1"(%arg0) <{
-+  //               CHECK: "vhlo.collective_permute_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME{LITERAL}:   source_target_pairs = #vhlo.tensor_v1<dense<[[0, 1], [1, 2], [2, 3]]> : tensor<3x2xi64>>
-   //          CHECK-SAME: }> : (!vhlo.tensor_v1<16x8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x8x!vhlo.f32_v1>
-@@ -981,8 +1054,9 @@
- }
- 
- // CHECK-LABEL: "op_compare"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_compare(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
--  //      CHECK: "vhlo.compare_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.compare_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   compare_type = #vhlo<comparison_type_v1 TOTALORDER>,
-   // CHECK-SAME:   comparison_direction = #vhlo<comparison_direction_v1 EQ>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-@@ -994,15 +1068,17 @@
- }
- 
- // CHECK-LABEL: "op_complex"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_complex(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<complex<f32>> {
--  // CHECK: "vhlo.complex_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>
-+  // CHECK: "vhlo.complex_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>
-   %0 = "stablehlo.complex"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<complex<f32>>
-   func.return %0 : tensor<complex<f32>>
- }
- 
- // CHECK-LABEL: "op_concatenate"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_concatenate(%arg0: tensor<8xf32>, %arg1: tensor<8xf32>) -> tensor<16xf32> {
--  //      CHECK: "vhlo.concatenate_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.concatenate_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   dimension = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<8x!vhlo.f32_v1>, !vhlo.tensor_v1<8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x!vhlo.f32_v1>
-   %0 = "stablehlo.concatenate"(%arg0, %arg1) {
-@@ -1012,6 +1088,7 @@
- }
- 
- // CHECK-LABEL: "op_constant"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_constant(%arg0: tensor<f32>) -> tensor<f32> {
-   //      CHECK: "vhlo.constant_v1"() <{
-   // CHECK-SAME:   value = #vhlo.tensor_v1<dense<0.000000e+00> : tensor<f32>>
-@@ -1023,15 +1100,17 @@
- }
- 
- // CHECK-LABEL: "op_convert"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_convert(%arg0: tensor<i32>) -> tensor<f32> {
--  // CHECK: "vhlo.convert_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.convert_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.convert"(%arg0) : (tensor<i32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_convolution"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_convolution(%arg0: tensor<1x8x8x207xf32>, %arg1: tensor<3x3x207x16xf32>) -> tensor<1x7x7x16xf32> {
--  //      CHECK: "vhlo.convolution_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.convolution_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   batch_group_count = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   feature_group_count = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   input_batch_dimension = #vhlo.integer_v1<0 : i64>,
-@@ -1065,8 +1144,9 @@
- }
- 
- // CHECK-LABEL: "op_cosine"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_cosine(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.cosine_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.cosine_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.cosine"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
-@@ -1079,8 +1159,9 @@
- }
- 
- // CHECK-LABEL: "op_cross_replica_sum"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_cross_replica_sum(%arg0: tensor<f32>) -> tensor<f32> {
--  //               CHECK: "vhlo.cross-replica-sum_v1"(%arg0) <{
-+  //               CHECK: "vhlo.cross-replica-sum_v1"(%[[ARG0]]) <{
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0], [1]]> : tensor<2x1xi64>>
-   //          CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.cross-replica-sum"(%arg0) {
-@@ -1090,8 +1171,9 @@
- }
- 
- // CHECK-LABEL: "op_custom_call"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_custom_call(%arg0: tensor<f32>) -> tensor<f32> {
--  //      CHECK: "vhlo.custom_call_v1"(%arg0) <{
-+  //      CHECK: "vhlo.custom_call_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   api_version = #vhlo<api_version_v1 API_VERSION_STATUS_RETURNING>,
-   // CHECK-SAME:   backend_config = #vhlo.string_v1<"\08\03\1A\02">,
-   // CHECK-SAME:   call_target_name = #vhlo.string_v1<"foo">,
-@@ -1122,15 +1204,17 @@
- }
- 
- // CHECK-LABEL: "op_divide"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_divide(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.divide_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.divide_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.divide"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_dot_general"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_dot_general(%arg0: tensor<8x8x16xf32>, %arg1: tensor<8x16x8xf32>) -> tensor<8x8x8xf32> {
--  //      CHECK: "vhlo.dot_general_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.dot_general_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   lhs_batching_dimensions = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>,
-   // CHECK-SAME:   lhs_contracting_dimensions = #vhlo.tensor_v1<dense<2> : tensor<1xi64>>,
-   // CHECK-SAME:   precision_config = #vhlo.array_v1<[#vhlo<precision_v1 HIGHEST>, #vhlo<precision_v1 HIGHEST>]>,
-@@ -1150,8 +1234,9 @@
- }
- 
- // CHECK-LABEL: "op_dot"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_dot(%arg0: tensor<8x16xf32>, %arg1: tensor<16x8xf32>) -> tensor<8x8xf32> {
--  //      CHECK: "vhlo.dot_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.dot_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   precision_config = #vhlo.array_v1<[#vhlo<precision_v1 HIGHEST>, #vhlo<precision_v1 HIGHEST>]>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<8x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<8x8x!vhlo.f32_v1>
-   %0 = "stablehlo.dot"(%arg0, %arg1) {
-@@ -1161,8 +1246,9 @@
- }
- 
- // CHECK-LABEL: "op_dynamic_broadcast_in_dim"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_dynamic_broadcast_in_dim(%arg0: tensor<?x?xf32>, %arg1: tensor<2xindex>) -> tensor<?x?xf32> {
--  //      CHECK: "vhlo.dynamic_broadcast_in_dim_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.dynamic_broadcast_in_dim_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   broadcast_dimensions = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-   // CHECK-SAME:   known_expanding_dimensions = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>,
-   // CHECK-SAME:   known_nonexpanding_dimensions = #vhlo.tensor_v1<dense<1> : tensor<1xi64>>
-@@ -1176,8 +1262,9 @@
- }
- 
- // CHECK-LABEL: "op_dynamic_conv"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_dynamic_conv(%arg0: tensor<1x8x8x207xf32>, %arg1: tensor<3x3x207x16xf32>, %arg2: tensor<4xi32>) -> tensor<1x?x?x16xf32> {
--  //      CHECK: "vhlo.dynamic_conv_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.dynamic_conv_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   batch_group_count = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   feature_group_count = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   input_batch_dimension = #vhlo.integer_v1<0 : i64>,
-@@ -1211,8 +1298,9 @@
- }
- 
- // CHECK-LABEL: "op_dynamic_gather"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_dynamic_gather(%arg0 : tensor<2x4x9xf32>, %arg1 : tensor<1x5x2xi32>, %arg2 : tensor<3xi32>) -> tensor<1x5x8xf32> {
--  //      CHECK: "vhlo.dynamic_gather_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.dynamic_gather_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   collapsed_slice_dims = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-   // CHECK-SAME:   index_vector_dim = #vhlo.integer_v1<2 : i64>,
-   // CHECK-SAME:   indices_are_sorted = #vhlo.bool_v1<true>,
-@@ -1232,8 +1320,9 @@
- }
- 
- // CHECK-LABEL: "op_dynamic_iota"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_dynamic_iota(%arg0: tensor<1xindex>) -> tensor<?xf32> {
--  //      CHECK: "vhlo.dynamic_iota_v1"(%arg0) <{
-+  //      CHECK: "vhlo.dynamic_iota_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   iota_dimension = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<1x!vhlo.index_v1>) -> !vhlo.tensor_v1<?x!vhlo.f32_v1>
-   %0 = "stablehlo.dynamic_iota"(%arg0) {
-@@ -1243,22 +1332,25 @@
- }
- 
- // CHECK-LABEL: "op_dynamic_pad"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}}, %[[ARG3:.*]]: {{.*}}, %[[ARG4:.*]]: {{.*}})
- func.func @op_dynamic_pad(%arg0: tensor<?xf32>, %arg1: tensor<f32>, %arg2: tensor<1xindex>, %arg3: tensor<1xindex>, %arg4: tensor<1xindex>) -> tensor<?xf32> {
--  // CHECK: "vhlo.dynamic_pad_v1"(%arg0, %arg1, %arg2, %arg3, %arg4) : (!vhlo.tensor_v1<?x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>) -> !vhlo.tensor_v1<?x!vhlo.f32_v1>
-+  // CHECK: "vhlo.dynamic_pad_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]], %[[ARG3]], %[[ARG4]]) : (!vhlo.tensor_v1<?x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>) -> !vhlo.tensor_v1<?x!vhlo.f32_v1>
-   %0 = "stablehlo.dynamic_pad"(%arg0, %arg1, %arg2, %arg3, %arg4) : (tensor<?xf32>, tensor<f32>, tensor<1xindex>, tensor<1xindex>, tensor<1xindex>) -> tensor<?xf32>
-   func.return %0 : tensor<?xf32>
- }
- 
- // CHECK-LABEL: "op_dynamic_reshape"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_dynamic_reshape(%arg0: tensor<16xf32>, %arg1: tensor<2xindex>) -> tensor<?x?xf32> {
--  // CHECK: "vhlo.dynamic_reshape_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<2x!vhlo.index_v1>) -> !vhlo.tensor_v1<?x?x!vhlo.f32_v1>
-+  // CHECK: "vhlo.dynamic_reshape_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<2x!vhlo.index_v1>) -> !vhlo.tensor_v1<?x?x!vhlo.f32_v1>
-   %0 = "stablehlo.dynamic_reshape"(%arg0, %arg1) : (tensor<16xf32>, tensor<2xindex>) -> tensor<?x?xf32>
-   func.return %0 : tensor<?x?xf32>
- }
- 
- // CHECK-LABEL: "op_dynamic_slice"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_dynamic_slice(%arg0: tensor<16xf32>, %arg1: tensor<i64>) -> tensor<4xf32> {
--  //      CHECK: "vhlo.dynamic_slice_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.dynamic_slice_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   slice_sizes = #vhlo.tensor_v1<dense<4> : tensor<1xi64>>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.i64_v1>) -> !vhlo.tensor_v1<4x!vhlo.f32_v1>
-   %0 = "stablehlo.dynamic_slice"(%arg0, %arg1) {
-@@ -1268,15 +1360,17 @@
- }
- 
- // CHECK-LABEL: "op_dynamic_update_slice"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_dynamic_update_slice(%arg0: tensor<16xf32>, %arg1: tensor<4xf32>, %arg2: tensor<i64>) -> tensor<16xf32> {
--  // CHECK: "vhlo.dynamic_update_slice_v1"(%arg0, %arg1, %arg2) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<4x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.i64_v1>) -> !vhlo.tensor_v1<16x!vhlo.f32_v1>
-+  // CHECK: "vhlo.dynamic_update_slice_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<4x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.i64_v1>) -> !vhlo.tensor_v1<16x!vhlo.f32_v1>
-   %0 = "stablehlo.dynamic_update_slice"(%arg0, %arg1, %arg2) : (tensor<16xf32>, tensor<4xf32>, tensor<i64>) -> tensor<16xf32>
-   func.return %0 : tensor<16xf32>
- }
- 
- // CHECK-LABEL: "op_einsum"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_einsum(%arg0: tensor<8x16xf32>, %arg1: tensor<16x8xf32>) -> tensor<8x8xf32> {
--  //      CHECK: "vhlo.einsum_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.einsum_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   einsum_config = #vhlo.string_v1<"ab,bc->ac">
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<8x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<8x8x!vhlo.f32_v1>
-   %0 = "stablehlo.einsum"(%arg0, %arg1) {
-@@ -1286,22 +1380,25 @@
- }
- 
- // CHECK-LABEL: "op_exponential_minus_one"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_exponential_minus_one(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.exponential_minus_one_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.exponential_minus_one_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.exponential_minus_one"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_exponential"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_exponential(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.exponential_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.exponential_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.exponential"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_fft"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_fft(%arg0: tensor<16xcomplex<f32>>) -> tensor<16xcomplex<f32>> {
--  //      CHECK: "vhlo.fft_v1"(%arg0) <{
-+  //      CHECK: "vhlo.fft_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   fft_length = #vhlo.tensor_v1<dense<16> : tensor<1xi64>>,
-   // CHECK-SAME:   fft_type = #vhlo<fft_type_v1 FFT>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<16x!vhlo.complex_v1<!vhlo.f32_v1>>) -> !vhlo.tensor_v1<16x!vhlo.complex_v1<!vhlo.f32_v1>>
-@@ -1313,8 +1410,9 @@
- }
- 
- // CHECK-LABEL: "op_floor"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_floor(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.floor_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.floor_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.floor"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
-@@ -1327,16 +1425,17 @@
-   // CHECK-SAME:   sym_name = #vhlo.string_v1<"op_func">,
-   // CHECK-SAME:   sym_visibility = #vhlo.string_v1<"private">
-   // CHECK-SAME: }> ({
--  // CHECK-NEXT:   ^[[BB:bb.*]](%arg0: !vhlo.tensor_v1<!vhlo.f32_v1>):
--  // CHECK-NEXT:     "vhlo.return_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-+  // CHECK-NEXT:   ^[[BB:bb.*]](%[[ARG0:.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
-+  // CHECK-NEXT:     "vhlo.return_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-   // CHECK-NEXT: }) : () -> ()
- 
-   func.return %arg0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_gather"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_gather(%arg0 : tensor<2x4x9xf32>, %arg1 : tensor<1x5x2xi32>) -> tensor<1x5x1xf32> {
--  //      CHECK: "vhlo.gather_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.gather_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   collapsed_slice_dims = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-   // CHECK-SAME:   index_vector_dim = #vhlo.integer_v1<2 : i64>,
-   // CHECK-SAME:   indices_are_sorted = #vhlo.bool_v1<true>,
-@@ -1358,8 +1457,9 @@
- }
- 
- // CHECK-LABEL: "op_get_dimension_size"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_get_dimension_size(%arg0: tensor<?xf32>) -> tensor<i32> {
--  //      CHECK: "vhlo.get_dimension_size_v1"(%arg0) <{
-+  //      CHECK: "vhlo.get_dimension_size_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   dimension = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<?x!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-   %0 = "stablehlo.get_dimension_size"(%arg0) {
-@@ -1369,8 +1469,9 @@
- }
- 
- // CHECK-LABEL: "op_get_tuple_element"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_get_tuple_element(%arg0: tuple<tensor<f32>, tensor<i32>>) -> tensor<f32> {
--  //      CHECK: "vhlo.get_tuple_element_v1"(%arg0) <{
-+  //      CHECK: "vhlo.get_tuple_element_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   index = #vhlo.integer_v1<0 : i32>
-   // CHECK-SAME: }> : (!vhlo.tuple_v1<!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.get_tuple_element"(%arg0) {
-@@ -1380,11 +1481,12 @@
- }
- 
- // CHECK-LABEL: "op_if"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_if(%arg0: tensor<i1>, %arg1: tensor<f32>, %arg2: tensor<f32>) -> tensor<f32> {
--  //      CHECK: "vhlo.if_v1"(%arg0) ({
--  // CHECK-NEXT:   "vhlo.return_v1"(%arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-+  //      CHECK: "vhlo.if_v1"(%[[ARG0]]) ({
-+  // CHECK-NEXT:   "vhlo.return_v1"(%[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-   // CHECK-NEXT: }, {
--  // CHECK-NEXT:   "vhlo.return_v1"(%arg2) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-+  // CHECK-NEXT:   "vhlo.return_v1"(%[[ARG2]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-   // CHECK-NEXT: }) : (!vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.if"(%arg0) ({
-     "stablehlo.return"(%arg1) : (tensor<f32>) -> ()
-@@ -1395,15 +1497,17 @@
- }
- 
- // CHECK-LABEL: "op_imag"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_imag(%arg0: tensor<complex<f32>>) -> tensor<f32> {
--  // CHECK: "vhlo.imag_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.imag_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.imag"(%arg0) : (tensor<complex<f32>>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_infeed"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_infeed(%arg0: !stablehlo.token) -> (tensor<f32>, !stablehlo.token) {
--  //               CHECK: "vhlo.infeed_v1"(%arg0) <{
-+  //               CHECK: "vhlo.infeed_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   infeed_config = #vhlo.string_v1<"foo">,
-   // CHECK-SAME{LITERAL}:   layout = #vhlo.array_v1<[#vhlo.array_v1<[]>]>
-   //          CHECK-SAME: }> : (!vhlo.token_v1) -> (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.token_v1)
-@@ -1426,36 +1530,41 @@
- }
- 
- // CHECK-LABEL: "op_is_finite"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_is_finite(%arg0: tensor<f32>) -> tensor<i1> {
--  // CHECK: "vhlo.is_finite_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-+  // CHECK: "vhlo.is_finite_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-   %0 = "stablehlo.is_finite"(%arg0) : (tensor<f32>) -> tensor<i1>
-   func.return %0 : tensor<i1>
- }
- 
- // CHECK-LABEL: "op_log"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_log(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.log_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.log_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.log"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_log_plus_one"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_log_plus_one(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.log_plus_one_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.log_plus_one_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.log_plus_one"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_logistic"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_logistic(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.logistic_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.logistic_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.logistic"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_map"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_map(%arg0: tensor<16xf32>) -> tensor<16xf32> {
--  //      CHECK: "vhlo.map_v1"(%arg0) <{
-+  //      CHECK: "vhlo.map_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   dimensions = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>
-   // CHECK-SAME: }> ({
-   // CHECK-NEXT:   ^[[BB:bb.*]](%[[ARG1:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
-@@ -1473,57 +1582,65 @@
- }
- 
- // CHECK-LABEL: "op_maximum"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_maximum(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.maximum_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.maximum_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.maximum"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_minimum"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_minimum(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.minimum_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.minimum_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.minimum"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_multiply"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_multiply(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.multiply_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.multiply_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.multiply"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_negate"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_negate(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.negate_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.negate_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.negate"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_not"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_not(%arg0: tensor<i1>) -> tensor<i1> {
--  // CHECK: "vhlo.not_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-+  // CHECK: "vhlo.not_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-   %0 = "stablehlo.not"(%arg0) : (tensor<i1>) -> tensor<i1>
-   func.return %0 : tensor<i1>
- }
- 
- // CHECK-LABEL: "op_optimization_barrier"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_optimization_barrier(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.optimization_barrier_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.optimization_barrier_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.optimization_barrier"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_or"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_or(%arg0: tensor<i1>, %arg1: tensor<i1>) -> tensor<i1> {
--  // CHECK: "vhlo.or_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-+  // CHECK: "vhlo.or_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-   %0 = "stablehlo.or"(%arg0, %arg1) : (tensor<i1>, tensor<i1>) -> tensor<i1>
-   func.return %0 : tensor<i1>
- }
- 
- // CHECK-LABEL: "op_outfeed"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_outfeed(%arg0: tensor<f32>, %arg1: !stablehlo.token) -> !stablehlo.token {
--  //      CHECK: "vhlo.outfeed_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.outfeed_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   outfeed_config = #vhlo.string_v1<"foo">
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.token_v1) -> !vhlo.token_v1
-   %0 = "stablehlo.outfeed"(%arg0, %arg1) {
-@@ -1533,8 +1650,9 @@
- }
- 
- // CHECK-LABEL: "op_pad"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_pad(%arg0: tensor<8xf32>, %arg1: tensor<f32>) -> tensor<16xf32> {
--  //      CHECK: "vhlo.pad_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.pad_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   edge_padding_high = #vhlo.tensor_v1<dense<4> : tensor<1xi64>>,
-   // CHECK-SAME:   edge_padding_low = #vhlo.tensor_v1<dense<4> : tensor<1xi64>>,
-   // CHECK-SAME:   interior_padding = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>
-@@ -1548,36 +1666,41 @@
- }
- 
- // CHECK-LABEL: "op_popcnt"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_popcnt(%arg0: tensor<i32>) -> tensor<i32> {
--  // CHECK: "vhlo.popcnt_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-+  // CHECK: "vhlo.popcnt_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-   %0 = "stablehlo.popcnt"(%arg0) : (tensor<i32>) -> tensor<i32>
-   func.return %0 : tensor<i32>
- }
- 
- // CHECK-LABEL: "op_power"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_power(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.power_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.power_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.power"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_real_dynamic_slice"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}}, %[[ARG3:.*]]: {{.*}})
- func.func @op_real_dynamic_slice(%arg0: tensor<?xf32>, %arg1: tensor<1xindex>, %arg2: tensor<1xindex>, %arg3: tensor<1xindex>) -> tensor<?xf32> {
--  // CHECK: "vhlo.real_dynamic_slice_v1"(%arg0, %arg1, %arg2, %arg3) : (!vhlo.tensor_v1<?x!vhlo.f32_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>) -> !vhlo.tensor_v1<?x!vhlo.f32_v1>
-+  // CHECK: "vhlo.real_dynamic_slice_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]], %[[ARG3]]) : (!vhlo.tensor_v1<?x!vhlo.f32_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>) -> !vhlo.tensor_v1<?x!vhlo.f32_v1>
-   %0 = "stablehlo.real_dynamic_slice"(%arg0, %arg1, %arg2, %arg3) : (tensor<?xf32>, tensor<1xindex>, tensor<1xindex>, tensor<1xindex>) -> tensor<?xf32>
-   func.return %0 : tensor<?xf32>
- }
- 
- // CHECK-LABEL: "op_real"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_real(%arg0: tensor<complex<f32>>) -> tensor<f32> {
--  // CHECK: "vhlo.real_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.real_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.real"(%arg0) : (tensor<complex<f32>>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_recv"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_recv(%arg0: !stablehlo.token) -> (tensor<f32>, !stablehlo.token) {
--  //      CHECK: "vhlo.recv_v1"(%arg0) <{
-+  //      CHECK: "vhlo.recv_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME:   channel_type = #vhlo.integer_v1<3 : i64>,
-   // CHECK-SAME:   is_host_transfer = #vhlo.bool_v1<true>
-@@ -1590,8 +1713,9 @@
- }
- 
- // CHECK-LABEL: "op_reduce"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_reduce(%arg0: tensor<16xf32>, %arg1: tensor<f32>) -> tensor<f32> {
--  //  CHECK: "vhlo.reduce_v1"(%arg0, %arg1)
-+  //  CHECK: "vhlo.reduce_v1"(%[[ARG0]], %[[ARG1]])
-   //  CHECK:   ^[[BB:bb.*]](%[[ARG1:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>, %[[ARG2:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
-   //  CHECK:     "vhlo.return_v1"(%[[VAL1:.*]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-   //  CHECK: }) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-@@ -1606,8 +1730,9 @@
- }
- 
- // CHECK-LABEL: "op_reduce_precision"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_reduce_precision(%arg0: tensor<f32>) -> tensor<f32> {
--  //      CHECK: "vhlo.reduce_precision_v1"(%arg0) <{
-+  //      CHECK: "vhlo.reduce_precision_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   exponent_bits = #vhlo.integer_v1<8 : i32>
-   // CHECK-SAME:   mantissa_bits = #vhlo.integer_v1<10 : i32>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-@@ -1619,8 +1744,9 @@
- }
- 
- // CHECK-LABEL: "op_reduce_scatter"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_reduce_scatter(%arg0: tensor<16xf32>) -> tensor<16xf32> {
--  //               CHECK: "vhlo.reduce_scatter_v1"(%arg0) <{
-+  //               CHECK: "vhlo.reduce_scatter_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0], [1]]> : tensor<2x1xi64>>,
-   //          CHECK-SAME:   scatter_dimension = #vhlo.integer_v1<0 : i64>
-@@ -1644,8 +1770,9 @@
- }
- 
- // CHECK-LABEL: "op_reduce_window"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_reduce_window(%arg0: tensor<2x17x31x7xf32>, %arg1: tensor<f32>) -> tensor<2x9x16x7xf32> {
--  //               CHECK: "vhlo.reduce_window_v1"(%arg0, %arg1) <{
-+  //               CHECK: "vhlo.reduce_window_v1"(%[[ARG0]], %[[ARG1]]) <{
-   //          CHECK-SAME:   base_dilations = #vhlo.tensor_v1<dense<[1, 2, 2, 1]> : tensor<4xi64>>,
-   // CHECK-SAME{LITERAL}:   padding = #vhlo.tensor_v1<dense<[[0, 0], [2, 0], [0, 2], [0, 0]]> : tensor<4x2xi64>>,
-   //          CHECK-SAME:   window_dilations = #vhlo.tensor_v1<dense<[1, 2, 2, 1]> : tensor<4xi64>>,
-@@ -1671,8 +1798,9 @@
- }
- 
- // CHECK-LABEL: "op_remainder"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_remainder(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.remainder_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.remainder_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.remainder"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
-@@ -1692,16 +1820,18 @@
- }
- 
- // CHECK-LABEL: "op_reshape"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_reshape(%arg0: tensor<16xf32>) -> tensor<4x4xf32> {
--  // CHECK: "vhlo.reshape_v1"(%arg0) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<4x4x!vhlo.f32_v1>
-+  // CHECK: "vhlo.reshape_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<4x4x!vhlo.f32_v1>
-   %0 = "stablehlo.reshape"(%arg0) : (tensor<16xf32>) -> tensor<4x4xf32>
-   func.return %0 : tensor<4x4xf32>
- }
- 
- // CHECK-LABEL: "op_return"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_return(%arg0: tensor<i32>, %arg1: tensor<f32>) -> tensor<f32> {
--  //      CHECK: "vhlo.case_v1"(%arg0) ({
--  // CHECK-NEXT:   "vhlo.return_v1"(%arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-+  //      CHECK: "vhlo.case_v1"(%[[ARG0]]) ({
-+  // CHECK-NEXT:   "vhlo.return_v1"(%[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-   // CHECK-NEXT: }) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.case"(%arg0) ({
-     "stablehlo.return"(%arg1) : (tensor<f32>) -> ()
-@@ -1710,8 +1840,9 @@
- }
- 
- // CHECK-LABEL: "op_reverse"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_reverse(%arg0: tensor<16xf32>) -> tensor<16xf32> {
--  //      CHECK: "vhlo.reverse_v1"(%arg0) <{
-+  //      CHECK: "vhlo.reverse_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   dimensions = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x!vhlo.f32_v1>
-   %0 = "stablehlo.reverse"(%arg0) {
-@@ -1721,8 +1852,9 @@
- }
- 
- // CHECK-LABEL: "op_rng_bit_generator"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_rng_bit_generator(%arg0: tensor<f32>) -> (tensor<f32>, tensor<f32>) {
--  //      CHECK: "vhlo.rng_bit_generator_v1"(%arg0) <{
-+  //      CHECK: "vhlo.rng_bit_generator_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   rng_algorithm = #vhlo<rng_algorithm_v1 PHILOX>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>)
-   %0:2 = "stablehlo.rng_bit_generator"(%arg0) {
-@@ -1732,8 +1864,9 @@
- }
- 
- // CHECK-LABEL: "op_rng"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_rng(%arg0: tensor<f32>, %arg1: tensor<f32>, %arg2: tensor<0xindex>) -> tensor<f32> {
--  //      CHECK: "vhlo.rng_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.rng_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   rng_distribution = #vhlo<rng_distribution_v1 NORMAL>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<0x!vhlo.index_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.rng"(%arg0, %arg1, %arg2) {
-@@ -1743,29 +1876,33 @@
- }
- 
- // CHECK-LABEL: "op_round_nearest_afz"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_round_nearest_afz(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.round_nearest_afz_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.round_nearest_afz_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.round_nearest_afz"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_round_nearest_even"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_round_nearest_even(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.round_nearest_even_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.round_nearest_even_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.round_nearest_even"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_rsqrt"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_rsqrt(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.rsqrt_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.rsqrt_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.rsqrt"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_scatter"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_scatter(%arg0: tensor<200x100x300xf32>, %arg1: tensor<10x2xi32>, %arg2: tensor<10x300xf32>) -> tensor<200x100x300xf32> {
--  //      CHECK: "vhlo.scatter_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.scatter_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   index_vector_dim = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   indices_are_sorted = #vhlo.bool_v1<true>,
-   // CHECK-SAME:   inserted_window_dims = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-@@ -1795,8 +1932,9 @@
- }
- 
- // CHECK-LABEL: "op_select_and_scatter"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_select_and_scatter(%arg0: tensor<10x24x24x64xf32>, %arg1: tensor<12x13x13x66xf32>, %arg2: tensor<f32>) -> tensor<10x24x24x64xf32> {
--  //      CHECK: "vhlo.select_and_scatter_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.select_and_scatter_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   padding = #vhlo.tensor_v1<dense<1> : tensor<4x2xi64>>,
-   // CHECK-SAME:   window_dimensions = #vhlo.tensor_v1<dense<[1, 2, 2, 1]> : tensor<4xi64>>,
-   // CHECK-SAME:   window_strides = #vhlo.tensor_v1<dense<[1, 2, 2, 1]> : tensor<4xi64>>
-@@ -1826,15 +1964,17 @@
- }
- 
- // CHECK-LABEL: "op_select"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_select(%arg0: tensor<i1>, %arg1: tensor<f32>, %arg2: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.select_v1"(%arg0, %arg1, %arg2) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.select_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.select"(%arg0, %arg1, %arg2) : (tensor<i1>, tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_send"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_send(%arg0: tensor<f32>, %arg1: !stablehlo.token) -> !stablehlo.token {
--  //      CHECK: "vhlo.send_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.send_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME:   channel_type = #vhlo.integer_v1<2 : i64>,
-   // CHECK-SAME:   is_host_transfer = #vhlo.bool_v1<true>
-@@ -1847,8 +1987,9 @@
- }
- 
- // CHECK-LABEL: "op_set_dimension_size"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_set_dimension_size(%arg0: tensor<?xf32>, %arg1: tensor<i32>) -> tensor<16xf32> {
--  //      CHECK: "vhlo.set_dimension_size_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.set_dimension_size_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   dimension = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<?x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<16x!vhlo.f32_v1>
-   %0 = "stablehlo.set_dimension_size"(%arg0, %arg1) {
-@@ -1858,43 +1999,49 @@
- }
- 
- // CHECK-LABEL: "op_shift_left"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_shift_left(%arg0: tensor<i32>, %arg1: tensor<i32>) -> tensor<i32> {
--  // CHECK: "vhlo.shift_left_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.i32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-+  // CHECK: "vhlo.shift_left_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-   %0 = "stablehlo.shift_left"(%arg0, %arg1) : (tensor<i32>, tensor<i32>) -> tensor<i32>
-   func.return %0 : tensor<i32>
- }
- 
- // CHECK-LABEL: "op_shift_right_arithmetic"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_shift_right_arithmetic(%arg0: tensor<i32>, %arg1: tensor<i32>) -> tensor<i32> {
--  // CHECK: "vhlo.shift_right_arithmetic_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.i32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-+  // CHECK: "vhlo.shift_right_arithmetic_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-   %0 = "stablehlo.shift_right_arithmetic"(%arg0, %arg1) : (tensor<i32>, tensor<i32>) -> tensor<i32>
-   func.return %0 : tensor<i32>
- }
- 
- // CHECK-LABEL: "op_shift_right_logical"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_shift_right_logical(%arg0: tensor<i32>, %arg1: tensor<i32>) -> tensor<i32> {
--  // CHECK: "vhlo.shift_right_logical_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.i32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-+  // CHECK: "vhlo.shift_right_logical_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-   %0 = "stablehlo.shift_right_logical"(%arg0, %arg1) : (tensor<i32>, tensor<i32>) -> tensor<i32>
-   func.return %0 : tensor<i32>
- }
- 
- // CHECK-LABEL: "op_sign"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_sign(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.sign_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.sign_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.sign"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_sine"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_sine(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.sine_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.sine_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.sine"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_slice"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_slice(%arg0: tensor<16xf32>) -> tensor<4xf32> {
--  //      CHECK: "vhlo.slice_v1"(%arg0) <{
-+  //      CHECK: "vhlo.slice_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   limit_indices = #vhlo.tensor_v1<dense<4> : tensor<1xi64>>,
-   // CHECK-SAME:   start_indices = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>,
-   // CHECK-SAME:   strides = #vhlo.tensor_v1<dense<1> : tensor<1xi64>>
-@@ -1908,8 +2055,9 @@
- }
- 
- // CHECK-LABEL: "op_sort"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_sort(%arg0: tensor<16xf32>) -> tensor<16xf32> {
--  //      CHECK: "vhlo.sort_v1"(%arg0) <{
-+  //      CHECK: "vhlo.sort_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   dimension = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME:   is_stable = #vhlo.bool_v1<true>
-   // CHECK-SAME: }> ({
-@@ -1929,29 +2077,33 @@
- }
- 
- // CHECK-LABEL: "op_sqrt"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_sqrt(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.sqrt_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.sqrt_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.sqrt"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_subtract"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_subtract(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.subtract_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.subtract_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.subtract"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_tanh"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_tanh(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.tanh_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.tanh_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.tanh"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_torch_index_select"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_torch_index_select(%arg0: tensor<5x1x5xf32>, %arg1: tensor<2xi32>) ->  tensor<2x1x5xf32> {
--  //      CHECK: "vhlo.torch_index_select_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.torch_index_select_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   batch_dims = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME:   dim = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<5x1x5x!vhlo.f32_v1>, !vhlo.tensor_v1<2x!vhlo.i32_v1>) -> !vhlo.tensor_v1<2x1x5x!vhlo.f32_v1>
-@@ -1963,8 +2115,9 @@
- }
- 
- // CHECK-LABEL: "op_trace"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_trace(%arg0: tensor<f32>) {
--  //      CHECK: "vhlo.trace_v1"(%arg0) <{
-+  //      CHECK: "vhlo.trace_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   tag = #vhlo.string_v1<"foo">
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-   "stablehlo.trace"(%arg0) {
-@@ -1974,8 +2127,9 @@
- }
- 
- // CHECK-LABEL: "op_transpose"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_transpose(%arg0: tensor<16x8xf32>) ->  tensor<8x16xf32> {
--  //      CHECK: "vhlo.transpose_v1"(%arg0) <{
-+  //      CHECK: "vhlo.transpose_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   permutation = #vhlo.tensor_v1<dense<[1, 0]> : tensor<2xi64>>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<16x8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<8x16x!vhlo.f32_v1>
-   %0 = "stablehlo.transpose"(%arg0) {
-@@ -1985,8 +2139,9 @@
- }
- 
- // CHECK-LABEL: "op_triangular_solve"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_triangular_solve(%arg0: tensor<16x16xf32>, %arg1: tensor<16x16xf32>) ->  tensor<16x16xf32> {
--  //      CHECK: "vhlo.triangular_solve_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.triangular_solve_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   left_side = #vhlo.bool_v1<true>,
-   // CHECK-SAME:   lower = #vhlo.bool_v1<true>,
-   // CHECK-SAME:   transpose_a = #vhlo<transpose_v1 NO_TRANSPOSE>,
-@@ -2002,15 +2157,17 @@
- }
- 
- // CHECK-LABEL: "op_tuple"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_tuple(%arg0: tensor<f32>) -> tuple<tensor<f32>> {
--  // CHECK: "vhlo.tuple_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tuple_v1<!vhlo.tensor_v1<!vhlo.f32_v1>>
-+  // CHECK: "vhlo.tuple_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tuple_v1<!vhlo.tensor_v1<!vhlo.f32_v1>>
-   %0 = "stablehlo.tuple"(%arg0) : (tensor<f32>) -> tuple<tensor<f32>>
-   func.return %0 : tuple<tensor<f32>>
- }
- 
- // CHECK-LABEL: "op_unary_einsum"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_unary_einsum(%arg0: tensor<8x16xf32>) -> tensor<8xf32> {
--  //      CHECK: "vhlo.unary_einsum_v1"(%arg0) <{
-+  //      CHECK: "vhlo.unary_einsum_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   einsum_config = #vhlo.string_v1<"ab->a">
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<8x16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<8x!vhlo.f32_v1>
-   %0 = "stablehlo.unary_einsum"(%arg0) {
-@@ -2020,22 +2177,25 @@
- }
- 
- // CHECK-LABEL: "op_uniform_dequantize"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_uniform_dequantize(%arg0: tensor<!quant.uniform<i8:f32, 34.0:16>>) -> tensor<f32> {
--  // CHECK: "vhlo.uniform_dequantize_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.uniform_dequantize_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.uniform_dequantize"(%arg0) : (tensor<!quant.uniform<i8:f32, 34.0:16>>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_uniform_quantize"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_uniform_quantize(%arg0: tensor<f32>) -> tensor<!quant.uniform<i8:f32, 34.0:16>> {
--  // CHECK: "vhlo.uniform_quantize_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>
-+  // CHECK: "vhlo.uniform_quantize_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>
-   %0 = "stablehlo.uniform_quantize"(%arg0) : (tensor<f32>) -> tensor<!quant.uniform<i8:f32, 34.0:16>>
-   func.return %0 : tensor<!quant.uniform<i8:f32, 34.0:16>>
- }
- 
- // CHECK-LABEL: "op_while"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_while(%arg0: tensor<i1>) -> tensor<i1> {
--  //      CHECK: "vhlo.while_v1"(%arg0) ({
-+  //      CHECK: "vhlo.while_v1"(%[[ARG0]]) ({
-   // CHECK-NEXT:   ^[[BB:bb.*]](%[[ARG1:arg.*]]: !vhlo.tensor_v1<!vhlo.bool_v1>):
-   // CHECK-NEXT:     "vhlo.return_v1"(%[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>) -> ()
-   // CHECK-NEXT:   }, {
-@@ -2053,8 +2213,9 @@
- }
- 
- // CHECK-LABEL: "op_xor"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_xor(%arg0: tensor<i1>, %arg1: tensor<i1>) -> tensor<i1> {
--  // CHECK: "vhlo.xor_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-+  // CHECK: "vhlo.xor_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-   %0 = "stablehlo.xor"(%arg0, %arg1) : (tensor<i1>, tensor<i1>) -> tensor<i1>
-   func.return %0 : tensor<i1>
- }
-@@ -2062,190 +2223,217 @@
- // ============ TYPES ============
- 
- // CHECK-LABEL: "type_i1"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_i1(%arg0: tensor<i1>, %arg1: tensor<i1>) -> tensor<i1> {
--  // CHECK: "vhlo.and_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-+  // CHECK: "vhlo.and_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-   %0 = "stablehlo.and"(%arg0, %arg1) : (tensor<i1>, tensor<i1>) -> tensor<i1>
-   func.return %0 : tensor<i1>
- }
- 
- // CHECK-LABEL: "type_i4"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_i4(%arg0: tensor<i4>, %arg1: tensor<i4>) -> tensor<i4> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.i4_v1>, !vhlo.tensor_v1<!vhlo.i4_v1>) -> !vhlo.tensor_v1<!vhlo.i4_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i4_v1>, !vhlo.tensor_v1<!vhlo.i4_v1>) -> !vhlo.tensor_v1<!vhlo.i4_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<i4>, tensor<i4>) -> tensor<i4>
-   func.return %0 : tensor<i4>
- }
- 
- // CHECK-LABEL: "type_i8"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_i8(%arg0: tensor<i8>, %arg1: tensor<i8>) -> tensor<i8> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.i8_v1>, !vhlo.tensor_v1<!vhlo.i8_v1>) -> !vhlo.tensor_v1<!vhlo.i8_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i8_v1>, !vhlo.tensor_v1<!vhlo.i8_v1>) -> !vhlo.tensor_v1<!vhlo.i8_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<i8>, tensor<i8>) -> tensor<i8>
-   func.return %0 : tensor<i8>
- }
- 
- // CHECK-LABEL: "type_i16"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_i16(%arg0: tensor<i16>, %arg1: tensor<i16>) -> tensor<i16> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.i16_v1>, !vhlo.tensor_v1<!vhlo.i16_v1>) -> !vhlo.tensor_v1<!vhlo.i16_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i16_v1>, !vhlo.tensor_v1<!vhlo.i16_v1>) -> !vhlo.tensor_v1<!vhlo.i16_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<i16>, tensor<i16>) -> tensor<i16>
-   func.return %0 : tensor<i16>
- }
- 
- // CHECK-LABEL: "type_i32"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_i32(%arg0: tensor<i32>, %arg1: tensor<i32>) -> tensor<i32> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.i32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<i32>, tensor<i32>) -> tensor<i32>
-   func.return %0 : tensor<i32>
- }
- 
- // CHECK-LABEL: "type_i64"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_i64(%arg0: tensor<i64>, %arg1: tensor<i64>) -> tensor<i64> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.i64_v1>, !vhlo.tensor_v1<!vhlo.i64_v1>) -> !vhlo.tensor_v1<!vhlo.i64_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i64_v1>, !vhlo.tensor_v1<!vhlo.i64_v1>) -> !vhlo.tensor_v1<!vhlo.i64_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<i64>, tensor<i64>) -> tensor<i64>
-   func.return %0 : tensor<i64>
- }
- 
- // CHECK-LABEL: "type_ui4"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_ui4(%arg0: tensor<ui4>, %arg1: tensor<ui4>) -> tensor<ui4> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.ui4_v1>, !vhlo.tensor_v1<!vhlo.ui4_v1>) -> !vhlo.tensor_v1<!vhlo.ui4_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.ui4_v1>, !vhlo.tensor_v1<!vhlo.ui4_v1>) -> !vhlo.tensor_v1<!vhlo.ui4_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<ui4>, tensor<ui4>) -> tensor<ui4>
-   func.return %0 : tensor<ui4>
- }
- 
- // CHECK-LABEL: "type_ui8"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_ui8(%arg0: tensor<ui8>, %arg1: tensor<ui8>) -> tensor<ui8> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.ui8_v1>, !vhlo.tensor_v1<!vhlo.ui8_v1>) -> !vhlo.tensor_v1<!vhlo.ui8_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.ui8_v1>, !vhlo.tensor_v1<!vhlo.ui8_v1>) -> !vhlo.tensor_v1<!vhlo.ui8_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<ui8>, tensor<ui8>) -> tensor<ui8>
-   func.return %0 : tensor<ui8>
- }
- 
- // CHECK-LABEL: "type_ui16"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_ui16(%arg0: tensor<ui16>, %arg1: tensor<ui16>) -> tensor<ui16> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.ui16_v1>, !vhlo.tensor_v1<!vhlo.ui16_v1>) -> !vhlo.tensor_v1<!vhlo.ui16_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.ui16_v1>, !vhlo.tensor_v1<!vhlo.ui16_v1>) -> !vhlo.tensor_v1<!vhlo.ui16_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<ui16>, tensor<ui16>) -> tensor<ui16>
-   func.return %0 : tensor<ui16>
- }
- 
- // CHECK-LABEL: "type_ui32"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_ui32(%arg0: tensor<ui32>, %arg1: tensor<ui32>) -> tensor<ui32> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.ui32_v1>, !vhlo.tensor_v1<!vhlo.ui32_v1>) -> !vhlo.tensor_v1<!vhlo.ui32_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.ui32_v1>, !vhlo.tensor_v1<!vhlo.ui32_v1>) -> !vhlo.tensor_v1<!vhlo.ui32_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<ui32>, tensor<ui32>) -> tensor<ui32>
-   func.return %0 : tensor<ui32>
- }
- 
- // CHECK-LABEL: "type_ui64"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_ui64(%arg0: tensor<ui64>, %arg1: tensor<ui64>) -> tensor<ui64> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.ui64_v1>, !vhlo.tensor_v1<!vhlo.ui64_v1>) -> !vhlo.tensor_v1<!vhlo.ui64_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.ui64_v1>, !vhlo.tensor_v1<!vhlo.ui64_v1>) -> !vhlo.tensor_v1<!vhlo.ui64_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<ui64>, tensor<ui64>) -> tensor<ui64>
-   func.return %0 : tensor<ui64>
- }
- 
- // CHECK-LABEL: "type_f8E4M3FN"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_f8E4M3FN(%arg0: tensor<f8E4M3FN>, %arg1: tensor<f8E4M3FN>) -> tensor<f8E4M3FN> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f8E4M3FN_v1>, !vhlo.tensor_v1<!vhlo.f8E4M3FN_v1>) -> !vhlo.tensor_v1<!vhlo.f8E4M3FN_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f8E4M3FN_v1>, !vhlo.tensor_v1<!vhlo.f8E4M3FN_v1>) -> !vhlo.tensor_v1<!vhlo.f8E4M3FN_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f8E4M3FN>, tensor<f8E4M3FN>) -> tensor<f8E4M3FN>
-   func.return %0 : tensor<f8E4M3FN>
- }
- 
- // CHECK-LABEL: "type_f8E5M2"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_f8E5M2(%arg0: tensor<f8E5M2>, %arg1: tensor<f8E5M2>) -> tensor<f8E5M2> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f8E5M2_v1>, !vhlo.tensor_v1<!vhlo.f8E5M2_v1>) -> !vhlo.tensor_v1<!vhlo.f8E5M2_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f8E5M2_v1>, !vhlo.tensor_v1<!vhlo.f8E5M2_v1>) -> !vhlo.tensor_v1<!vhlo.f8E5M2_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f8E5M2>, tensor<f8E5M2>) -> tensor<f8E5M2>
-   func.return %0 : tensor<f8E5M2>
- }
- 
- // CHECK-LABEL: "type_f8E4M3FNUZ"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_f8E4M3FNUZ(%arg0: tensor<f8E4M3FNUZ>, %arg1: tensor<f8E4M3FNUZ>) -> tensor<f8E4M3FNUZ> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f8E4M3FNUZ_v1>, !vhlo.tensor_v1<!vhlo.f8E4M3FNUZ_v1>) -> !vhlo.tensor_v1<!vhlo.f8E4M3FNUZ_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f8E4M3FNUZ_v1>, !vhlo.tensor_v1<!vhlo.f8E4M3FNUZ_v1>) -> !vhlo.tensor_v1<!vhlo.f8E4M3FNUZ_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f8E4M3FNUZ>, tensor<f8E4M3FNUZ>) -> tensor<f8E4M3FNUZ>
-   func.return %0 : tensor<f8E4M3FNUZ>
- }
- 
- // CHECK-LABEL: "type_f8E4M3B11FNUZ"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_f8E4M3B11FNUZ(%arg0: tensor<f8E4M3B11FNUZ>, %arg1: tensor<f8E4M3B11FNUZ>) -> tensor<f8E4M3B11FNUZ> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f8E4M3B11FNUZ_v1>, !vhlo.tensor_v1<!vhlo.f8E4M3B11FNUZ_v1>) -> !vhlo.tensor_v1<!vhlo.f8E4M3B11FNUZ_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f8E4M3B11FNUZ_v1>, !vhlo.tensor_v1<!vhlo.f8E4M3B11FNUZ_v1>) -> !vhlo.tensor_v1<!vhlo.f8E4M3B11FNUZ_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f8E4M3B11FNUZ>, tensor<f8E4M3B11FNUZ>) -> tensor<f8E4M3B11FNUZ>
-   func.return %0 : tensor<f8E4M3B11FNUZ>
- }
- 
- // CHECK-LABEL: "type_f8E5M2FNUZ"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_f8E5M2FNUZ(%arg0: tensor<f8E5M2FNUZ>, %arg1: tensor<f8E5M2FNUZ>) -> tensor<f8E5M2FNUZ> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f8E5M2FNUZ_v1>, !vhlo.tensor_v1<!vhlo.f8E5M2FNUZ_v1>) -> !vhlo.tensor_v1<!vhlo.f8E5M2FNUZ_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f8E5M2FNUZ_v1>, !vhlo.tensor_v1<!vhlo.f8E5M2FNUZ_v1>) -> !vhlo.tensor_v1<!vhlo.f8E5M2FNUZ_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f8E5M2FNUZ>, tensor<f8E5M2FNUZ>) -> tensor<f8E5M2FNUZ>
-   func.return %0 : tensor<f8E5M2FNUZ>
- }
- 
- // CHECK-LABEL: "type_bf16"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_bf16(%arg0: tensor<bf16>, %arg1: tensor<bf16>) -> tensor<bf16> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.bf16_v1>, !vhlo.tensor_v1<!vhlo.bf16_v1>) -> !vhlo.tensor_v1<!vhlo.bf16_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.bf16_v1>, !vhlo.tensor_v1<!vhlo.bf16_v1>) -> !vhlo.tensor_v1<!vhlo.bf16_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<bf16>, tensor<bf16>) -> tensor<bf16>
-   func.return %0 : tensor<bf16>
- }
- 
- // CHECK-LABEL: "type_f16"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_f16(%arg0: tensor<f16>, %arg1: tensor<f16>) -> tensor<f16> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f16_v1>, !vhlo.tensor_v1<!vhlo.f16_v1>) -> !vhlo.tensor_v1<!vhlo.f16_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f16_v1>, !vhlo.tensor_v1<!vhlo.f16_v1>) -> !vhlo.tensor_v1<!vhlo.f16_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f16>, tensor<f16>) -> tensor<f16>
-   func.return %0 : tensor<f16>
- }
- 
- // CHECK-LABEL: "type_f32"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_f32(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "type_f64"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_f64(%arg0: tensor<f64>, %arg1: tensor<f64>) -> tensor<f64> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f64_v1>, !vhlo.tensor_v1<!vhlo.f64_v1>) -> !vhlo.tensor_v1<!vhlo.f64_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f64_v1>, !vhlo.tensor_v1<!vhlo.f64_v1>) -> !vhlo.tensor_v1<!vhlo.f64_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f64>, tensor<f64>) -> tensor<f64>
-   func.return %0 : tensor<f64>
- }
- 
- // CHECK-LABEL: "type_complex_f32"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_complex_f32(%arg0: tensor<complex<f32>>, %arg1: tensor<complex<f32>>) -> tensor<complex<f32>> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>, !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>) -> !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>, !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>) -> !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<complex<f32>>, tensor<complex<f32>>) -> tensor<complex<f32>>
-   func.return %0 : tensor<complex<f32>>
- }
- 
- // CHECK-LABEL: "type_complex_f64"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_complex_f64(%arg0: tensor<complex<f64>>, %arg1: tensor<complex<f64>>) -> tensor<complex<f64>> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f64_v1>>, !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f64_v1>>) -> !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f64_v1>>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f64_v1>>, !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f64_v1>>) -> !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f64_v1>>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<complex<f64>>, tensor<complex<f64>>) -> tensor<complex<f64>>
-   func.return %0 : tensor<complex<f64>>
- }
- 
- // CHECK-LABEL: "type_dynamism_ranked"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @type_dynamism_ranked(%arg0: tensor<?xf32>) -> tensor<?xf32> {
--  // CHECK: "vhlo.abs_v1"(%arg0) : (!vhlo.tensor_v1<?x!vhlo.f32_v1>) -> !vhlo.tensor_v1<?x!vhlo.f32_v1>
-+  // CHECK: "vhlo.abs_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<?x!vhlo.f32_v1>) -> !vhlo.tensor_v1<?x!vhlo.f32_v1>
-   %0 = "stablehlo.abs"(%arg0) : (tensor<?xf32>) -> tensor<?xf32>
-   func.return %0 : tensor<?xf32>
- }
- 
- // CHECK-LABEL: "type_quantization"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_quantization(%arg0: tensor<!quant.uniform<i8:f32, 34.0:16>>, %arg1: tensor<!quant.uniform<i8:f32, 34.0:16>>) -> tensor<!quant.uniform<i8:f32, 34.0:16>> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>, !vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>) -> !vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>, !vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>) -> !vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<!quant.uniform<i8:f32, 34.0:16>>, tensor<!quant.uniform<i8:f32, 34.0:16>>) -> tensor<!quant.uniform<i8:f32, 34.0:16>>
-   func.return %0 : tensor<!quant.uniform<i8:f32, 34.0:16>>
- }
- 
- //       CHECK: function_type = #vhlo.type_v1<!vhlo.func_v1<(!vhlo.token_v1) -> !vhlo.token_v1>>
- // CHECK-LABEL: "type_token_callee"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @type_token_callee(%arg0: !stablehlo.token) -> !stablehlo.token {
--  // CHECK: "vhlo.return_v1"(%arg0) : (!vhlo.token_v1) -> ()
-+  // CHECK: "vhlo.return_v1"(%[[ARG0]]) : (!vhlo.token_v1) -> ()
-   return %arg0 : !stablehlo.token
- }
- 
- //       CHECK: function_type = #vhlo.type_v1<!vhlo.func_v1<(!vhlo.token_v1) -> !vhlo.token_v1>>
- // CHECK-LABEL: "type_token_caller"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @type_token_caller(%arg0: !stablehlo.token) -> !stablehlo.token {
--  // CHECK:      "vhlo.call_v1"(%arg0) <{callee = #vhlo.string_v1<"type_token_callee">}
-+  // CHECK:      "vhlo.call_v1"(%[[ARG0]]) <{callee = #vhlo.string_v1<"type_token_callee">}
-   // CHECK-SAME: (!vhlo.token_v1) -> !vhlo.token_v1
-   %0 = func.call @type_token_callee(%arg0) : (!stablehlo.token) -> !stablehlo.token
-   return %0 : !stablehlo.token
- }
- 
- // CHECK-LABEL: "type_tuple"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @type_tuple(%arg0: tuple<tensor<f32>>) -> tuple<!stablehlo.token> {
-   %0 = "stablehlo.custom_call"(%arg0) {
-     call_target_name = "foo"
-diff --ruN a/stablehlo/stablehlo/tests/vhlo/stablehlo_legalize_to_vhlo.0_12_0.mlir b/stablehlo/stablehlo/tests/vhlo/stablehlo_legalize_to_vhlo.0_12_0.mlir
---- stablehlo/stablehlo/tests/vhlo/stablehlo_legalize_to_vhlo.0_12_0.mlir
-+++ stablehlo/stablehlo/tests/vhlo/stablehlo_legalize_to_vhlo.0_12_0.mlir
-@@ -13,6 +13,7 @@
- // ============ ATTRIBUTES ============
- 
- // CHECK-LABEL: "attr_comparison_direction_eq"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_direction_eq(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     // CHECK: comparison_direction = #vhlo<comparison_direction_v1 EQ>
-@@ -22,6 +23,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_direction_ne"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_direction_ne(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     // CHECK: comparison_direction = #vhlo<comparison_direction_v1 NE>
-@@ -31,6 +33,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_direction_ge"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_direction_ge(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     // CHECK: comparison_direction = #vhlo<comparison_direction_v1 GE>
-@@ -40,6 +43,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_direction_gt"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_direction_gt(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     // CHECK: comparison_direction = #vhlo<comparison_direction_v1 GT>
-@@ -49,6 +53,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_direction_le"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_direction_le(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     // CHECK: comparison_direction = #vhlo<comparison_direction_v1 LE>
-@@ -58,6 +63,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_direction_lt"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_direction_lt(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     // CHECK: comparison_direction = #vhlo<comparison_direction_v1 LT>
-@@ -67,6 +73,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_type_notype"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_type_notype(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     comparison_direction = #stablehlo<comparison_direction EQ>
-@@ -76,6 +83,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_type_float"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_type_float(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     comparison_direction = #stablehlo<comparison_direction EQ>,
-@@ -86,6 +94,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_type_totalorder"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_type_totalorder(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     comparison_direction = #stablehlo<comparison_direction EQ>,
-@@ -96,6 +105,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_type_signed"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_type_signed(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     comparison_direction = #stablehlo<comparison_direction EQ>,
-@@ -106,6 +116,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_type_unsigned"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_type_unsigned(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     comparison_direction = #stablehlo<comparison_direction EQ>,
-@@ -118,6 +129,7 @@
- // ConvDimensionNumbers aka #stablehlo.conv is covered below.
- 
- // CHECK-LABEL: "attr_custom_call_api_version_unspecified"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_custom_call_api_version_unspecified(%arg0: tensor<f32>) -> tensor<f32> {
-   %0 = "stablehlo.custom_call"(%arg0) {
-     call_target_name = "foo",
-@@ -128,6 +140,7 @@
- }
- 
- // CHECK-LABEL: "attr_custom_call_api_version_original"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_custom_call_api_version_original(%arg0: tensor<f32>) -> tensor<f32> {
-   %0 = "stablehlo.custom_call"(%arg0) {
-     call_target_name = "foo",
-@@ -138,6 +151,7 @@
- }
- 
- // CHECK-LABEL: "attr_custom_call_api_version_status_returning"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_custom_call_api_version_status_returning(%arg0: tensor<f32>) -> tensor<f32> {
-   %0 = "stablehlo.custom_call"(%arg0) {
-     call_target_name = "foo",
-@@ -148,6 +162,7 @@
- }
- 
- // CHECK-LABEL: "attr_custom_call_api_version_status_returning_unified"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_custom_call_api_version_status_returning_unified(%arg0: tensor<f32>) -> tensor<f32> {
-   %0 = "stablehlo.custom_call"(%arg0) {
-     call_target_name = "foo",
-@@ -160,6 +175,7 @@
- // DotDimensionNumbers aka #stablehlo.dot is covered below.
- 
- // CHECK-LABEL: "attr_fft_type_fft"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_fft_type_fft(%arg0: tensor<16xcomplex<f32>>) -> tensor<16xcomplex<f32>> {
-   %0 = "stablehlo.fft"(%arg0) {
-     // CHECK: fft_type = #vhlo<fft_type_v1 FFT>
-@@ -170,6 +186,7 @@
- }
- 
- // CHECK-LABEL: "attr_fft_type_ifft"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_fft_type_ifft(%arg0: tensor<16xcomplex<f32>>) -> tensor<16xcomplex<f32>> {
-   %0 = "stablehlo.fft"(%arg0) {
-     // CHECK: fft_type = #vhlo<fft_type_v1 IFFT>
-@@ -180,6 +197,7 @@
- }
- 
- // CHECK-LABEL: "attr_fft_type_rfft"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_fft_type_rfft(%arg0: tensor<16xf32>) -> tensor<9xcomplex<f32>> {
-   %0 = "stablehlo.fft"(%arg0) {
-     // CHECK: fft_type = #vhlo<fft_type_v1 RFFT>
-@@ -190,6 +208,7 @@
- }
- 
- // CHECK-LABEL: "attr_fft_type_irfft"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_fft_type_irfft(%arg0: tensor<9xcomplex<f32>>) -> tensor<16xf32> {
-   %0 = "stablehlo.fft"(%arg0) {
-     // CHECK: fft_type = #vhlo<fft_type_v1 IRFFT>
-@@ -202,6 +221,7 @@
- // GatherDimensionNumbers aka #stablehlo.gather is covered below.
- 
- // CHECK-LABEL: "attr_precision_config_default"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_precision_config_default(%arg0: tensor<8x16xf32>, %arg1: tensor<16x8xf32>) -> tensor<8x8xf32> {
-   %0 = "stablehlo.dot"(%arg0, %arg1) {
-     // CHECK: precision_config = #vhlo.array_v1<[#vhlo<precision_v1 DEFAULT>, #vhlo<precision_v1 DEFAULT>]>
-@@ -210,6 +230,7 @@
- }
- 
- // CHECK-LABEL: "attr_precision_config_high"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_precision_config_high(%arg0: tensor<8x16xf32>, %arg1: tensor<16x8xf32>) -> tensor<8x8xf32> {
-   %0 = "stablehlo.dot"(%arg0, %arg1) {
-     // CHECK: precision_config = #vhlo.array_v1<[#vhlo<precision_v1 HIGH>, #vhlo<precision_v1 HIGH>]>
-@@ -219,6 +240,7 @@
- }
- 
- // CHECK-LABEL: "attr_precision_config_highest"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_precision_config_highest(%arg0: tensor<8x16xf32>, %arg1: tensor<16x8xf32>) -> tensor<8x8xf32> {
-   %0 = "stablehlo.dot"(%arg0, %arg1) {
-     // CHECK: precision_config = #vhlo.array_v1<[#vhlo<precision_v1 HIGHEST>, #vhlo<precision_v1 HIGHEST>]>
-@@ -228,6 +250,7 @@
- }
- 
- // CHECK-LABEL: "attr_rng_algorithm_default"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_rng_algorithm_default(%arg0: tensor<f32>) -> (tensor<f32>, tensor<f32>) {
-   %0:2 = "stablehlo.rng_bit_generator"(%arg0) {
-     // CHECK: rng_algorithm = #vhlo<rng_algorithm_v1 DEFAULT>
-@@ -237,6 +260,7 @@
- }
- 
- // CHECK-LABEL: "attr_rng_algorithm_three_fry"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_rng_algorithm_three_fry(%arg0: tensor<f32>) -> (tensor<f32>, tensor<f32>) {
-   %0:2 = "stablehlo.rng_bit_generator"(%arg0) {
-     // CHECK: rng_algorithm = #vhlo<rng_algorithm_v1 THREE_FRY>
-@@ -246,6 +270,7 @@
- }
- 
- // CHECK-LABEL: "attr_rng_algorithm_philox"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_rng_algorithm_philox(%arg0: tensor<f32>) -> (tensor<f32>, tensor<f32>) {
-   %0:2 = "stablehlo.rng_bit_generator"(%arg0) {
-     // CHECK: rng_algorithm = #vhlo<rng_algorithm_v1 PHILOX>
-@@ -255,6 +280,7 @@
- }
- 
- // CHECK-LABEL: "attr_rng_distribution_uniform"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @attr_rng_distribution_uniform(%arg0: tensor<f32>, %arg1: tensor<f32>, %arg2: tensor<0xindex>) -> tensor<f32> {
-   %0 = "stablehlo.rng"(%arg0, %arg1, %arg2) {
-     // CHECK: rng_distribution = #vhlo<rng_distribution_v1 UNIFORM>
-@@ -264,6 +290,7 @@
- }
- 
- // CHECK-LABEL: "attr_rng_distribution_normal"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @attr_rng_distribution_normal(%arg0: tensor<f32>, %arg1: tensor<f32>, %arg2: tensor<0xindex>) -> tensor<f32> {
-   %0 = "stablehlo.rng"(%arg0, %arg1, %arg2) {
-     // CHECK: rng_distribution = #vhlo<rng_distribution_v1 NORMAL>
-@@ -275,6 +302,7 @@
- // ScatterDimensionNumbers aka #stablehlo.scatter is covered below.
- 
- // CHECK-LABEL: "attr_transpose_no_transpose"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_transpose_no_transpose(%arg0: tensor<16x16xf32>, %arg1: tensor<16x16xf32>) ->  tensor<16x16xf32> {
-   %0 = "stablehlo.triangular_solve"(%arg0, %arg1) {
-     left_side = true,
-@@ -287,6 +315,7 @@
- }
- 
- // CHECK-LABEL: "attr_transpose_transpose"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_transpose_transpose(%arg0: tensor<16x16xf32>, %arg1: tensor<16x16xf32>) ->  tensor<16x16xf32> {
-   %0 = "stablehlo.triangular_solve"(%arg0, %arg1) {
-     left_side = true,
-@@ -299,6 +328,7 @@
- }
- 
- // CHECK-LABEL: "attr_transpose_adjoint"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_transpose_adjoint(%arg0: tensor<16x16xf32>, %arg1: tensor<16x16xf32>) ->  tensor<16x16xf32> {
-   %0 = "stablehlo.triangular_solve"(%arg0, %arg1) {
-     left_side = true,
-@@ -313,10 +343,9 @@
- // TypeExtensionsAttr aka #stablehlo.type_extensions is covered below.
- 
- // CHECK-LABEL: "attr_type_extensions_bounds"
--func.func @attr_type_extensions_bounds(
--    %arg0: tensor<?x?xf32, #stablehlo.type_extensions<bounds = [16, ?]>>)
--    -> tensor<?x?xf32, #stablehlo.type_extensions<bounds = [16, ?]>> {
--  // CHECK: "vhlo.return_v1"(%arg0) : (!vhlo.tensor_v1<?x?x!vhlo.f32_v1, #vhlo.type_extensions_v1<bounds = [16, ?]>>) -> ()
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @attr_type_extensions_bounds(%arg0: tensor<?x?xf32, #stablehlo.type_extensions<bounds = [16, ?]>>) -> tensor<?x?xf32, #stablehlo.type_extensions<bounds = [16, ?]>> {
-+  // CHECK: "vhlo.return_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<?x?x!vhlo.f32_v1, #vhlo.type_extensions_v1<bounds = [16, ?]>>) -> ()
-   func.return %arg0 : tensor<?x?xf32, #stablehlo.type_extensions<bounds = [16, ?]>>
- }
- 
-@@ -324,8 +353,9 @@
- // ============ DEFAULTS ============
- 
- // CHECK-LABEL: "default_all_gather"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_all_gather(%arg0: tensor<16x8xf32>) -> tensor<16x16xf32> {
--  //               CHECK: "vhlo.all_gather_v1"(%arg0) <{
-+  //               CHECK: "vhlo.all_gather_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   all_gather_dim = #vhlo.integer_v1<1 : i64>
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0], [1]]> : tensor<2x1xi64>>,
-@@ -339,8 +369,9 @@
- }
- 
- // CHECK-LABEL: "default_all_reduce"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_all_reduce(%arg0: tensor<f32>) -> tensor<f32> {
--  //               CHECK: "vhlo.all_reduce_v1"(%arg0)
-+  //               CHECK: "vhlo.all_reduce_v1"(%[[ARG0]])
-   //          CHECK-SAME: <{
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0], [1]]> : tensor<2x1xi64>>,
-@@ -362,8 +393,9 @@
- }
- 
- // CHECK-LABEL: "default_all_to_all"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_all_to_all(%arg0: tensor<4x16xf32>) -> tensor<16x4xf32> {
--  //               CHECK: "vhlo.all_to_all_v1"(%arg0) <{
-+  //               CHECK: "vhlo.all_to_all_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   //          CHECK-SAME:   concat_dimension = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0, 1, 2, 3]]> : tensor<1x4xi64>>,
-@@ -380,8 +412,9 @@
- }
- 
- // CHECK-LABEL: "default_cholesky"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_cholesky(%arg0: tensor<1x16x16xf32>) -> tensor<1x16x16xf32> {
--  //      CHECK: "vhlo.cholesky_v1"(%arg0) <{
-+  //      CHECK: "vhlo.cholesky_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   lower = #vhlo.bool_v1<false>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<1x16x16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<1x16x16x!vhlo.f32_v1>
-   %0 = "stablehlo.cholesky"(%arg0) : (tensor<1x16x16xf32>) -> tensor<1x16x16xf32>
-@@ -389,8 +422,9 @@
- }
- 
- // CHECK-LABEL: "default_collective_permute"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_collective_permute(%arg0: tensor<16x8xf32>) -> tensor<16x8xf32> {
--  //               CHECK: "vhlo.collective_permute_v1"(%arg0) <{
-+  //               CHECK: "vhlo.collective_permute_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME{LITERAL}:   source_target_pairs = #vhlo.tensor_v1<dense<[[0, 1], [1, 2], [2, 3]]> : tensor<3x2xi64>>
-   //          CHECK-SAME: }> : (!vhlo.tensor_v1<16x8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x8x!vhlo.f32_v1>
-@@ -401,8 +435,9 @@
- }
- 
- // CHECK-LABEL: "default_compare"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @default_compare(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
--  //      CHECK: "vhlo.compare_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.compare_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   compare_type = #vhlo<comparison_type_v1 NOTYPE>,
-   // CHECK-SAME:   comparison_direction = #vhlo<comparison_direction_v1 EQ>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-@@ -413,8 +448,9 @@
- }
- 
- // CHECK-LABEL: "default_convolution"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @default_convolution(%arg0: tensor<1x8x8x207xf32>, %arg1: tensor<3x3x207x16xf32>) -> tensor<1x6x6x16xf32> {
--  //      CHECK: "vhlo.convolution_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.convolution_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   batch_group_count = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   feature_group_count = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   input_batch_dimension = #vhlo.integer_v1<0 : i64>,
-@@ -442,8 +478,9 @@
- }
- 
- // CHECK-LABEL: "default_custom_call"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_custom_call(%arg0: tensor<f32>) -> tensor<f32> {
--  //      CHECK: "vhlo.custom_call_v1"(%arg0) <{
-+  //      CHECK: "vhlo.custom_call_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   api_version = #vhlo<api_version_v1 API_VERSION_ORIGINAL>,
-   // CHECK-SAME:   backend_config = #vhlo.string_v1<"">,
-   // CHECK-SAME:   call_target_name = #vhlo.string_v1<"foo">,
-@@ -460,8 +497,9 @@
- }
- 
- // CHECK-LABEL: "default_dot_general"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @default_dot_general(%arg0: tensor<8x8x16xf32>, %arg1: tensor<8x16x8xf32>) -> tensor<8x8x8xf32> {
--  //      CHECK: "vhlo.dot_general_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.dot_general_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   lhs_batching_dimensions = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>,
-   // CHECK-SAME:   lhs_contracting_dimensions = #vhlo.tensor_v1<dense<2> : tensor<1xi64>>,
-   // CHECK-SAME:   precision_config = #vhlo.array_v1<[#vhlo<precision_v1 DEFAULT>, #vhlo<precision_v1 DEFAULT>]>,
-@@ -480,8 +518,9 @@
- }
- 
- // CHECK-LABEL: "default_dot"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @default_dot(%arg0: tensor<8x16xf32>, %arg1: tensor<16x8xf32>) -> tensor<8x8xf32> {
--  //      CHECK: "vhlo.dot_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.dot_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   precision_config = #vhlo.array_v1<[#vhlo<precision_v1 DEFAULT>, #vhlo<precision_v1 DEFAULT>]>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<8x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<8x8x!vhlo.f32_v1>
-   %0 = "stablehlo.dot"(%arg0, %arg1) : (tensor<8x16xf32>, tensor<16x8xf32>) -> tensor<8x8xf32>
-@@ -489,8 +528,9 @@
- }
- 
- // CHECK-LABEL: "default_dynamic_broadcast_in_dim"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @default_dynamic_broadcast_in_dim(%arg0: tensor<?x?xf32>, %arg1: tensor<2xindex>) -> tensor<?x?xf32> {
--  //      CHECK: "vhlo.dynamic_broadcast_in_dim_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.dynamic_broadcast_in_dim_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   broadcast_dimensions = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-   // CHECK-SAME:   known_expanding_dimensions = #vhlo.tensor_v1<dense<> : tensor<0xi64>>,
-   // CHECK-SAME:   known_nonexpanding_dimensions = #vhlo.tensor_v1<dense<> : tensor<0xi64>>
-@@ -502,8 +542,9 @@
- }
- 
- // CHECK-LABEL: "default_dynamic_conv"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @default_dynamic_conv(%arg0: tensor<1x8x8x207xf32>, %arg1: tensor<3x3x207x16xf32>, %arg2: tensor<4xi32>) -> tensor<1x?x?x16xf32> {
--  //      CHECK: "vhlo.dynamic_conv_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.dynamic_conv_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   batch_group_count = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   feature_group_count = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   input_batch_dimension = #vhlo.integer_v1<0 : i64>,
-@@ -531,8 +572,9 @@
- }
- 
- // CHECK-LABEL: "default_dynamic_gather"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @default_dynamic_gather(%arg0 : tensor<2x4x9xf32>, %arg1 : tensor<1x5x2xi32>, %arg2 : tensor<3xi32>) -> tensor<1x5x8xf32> {
--  //      CHECK: "vhlo.dynamic_gather_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.dynamic_gather_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   collapsed_slice_dims = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-   // CHECK-SAME:   index_vector_dim = #vhlo.integer_v1<2 : i64>,
-   // CHECK-SAME:   indices_are_sorted = #vhlo.bool_v1<false>,
-@@ -558,15 +600,16 @@
-   // CHECK-SAME:   sym_name = #vhlo.string_v1<"default_func">,
-   // CHECK-SAME:   sym_visibility = #vhlo.string_v1<"">
-   // CHECK-SAME: }> ({
--  // CHECK-NEXT:   ^[[BB:bb.*]](%arg0: !vhlo.tensor_v1<!vhlo.f32_v1>):
--  // CHECK-NEXT:     "vhlo.return_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-+  // CHECK-NEXT:   ^[[BB:bb.*]](%[[ARG0:.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
-+  // CHECK-NEXT:     "vhlo.return_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-   // CHECK-NEXT: }) : () -> ()
-   func.return %arg0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "dynamic_gather"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @dynamic_gather(%arg0 : tensor<2x4x9xf32>, %arg1 : tensor<1x5x2xi32>) -> tensor<1x5x1xf32> {
--  //      CHECK: "vhlo.gather_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.gather_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   collapsed_slice_dims = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-   // CHECK-SAME:   index_vector_dim = #vhlo.integer_v1<2 : i64>,
-   // CHECK-SAME:   indices_are_sorted = #vhlo.bool_v1<false>,
-@@ -587,8 +630,9 @@
- }
- 
- // CHECK-LABEL: "default_infeed"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_infeed(%arg0: !stablehlo.token) -> (tensor<f32>, !stablehlo.token) {
--  //               CHECK: "vhlo.infeed_v1"(%arg0) <{
-+  //               CHECK: "vhlo.infeed_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   infeed_config = #vhlo.string_v1<"">,
-   // CHECK-SAME{LITERAL}:   layout = #vhlo.array_v1<[]>
-   //          CHECK-SAME: }> : (!vhlo.token_v1) -> (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.token_v1)
-@@ -597,8 +641,9 @@
- }
- 
- // CHECK-LABEL: "default_outfeed"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @default_outfeed(%arg0: tensor<f32>, %arg1: !stablehlo.token) -> !stablehlo.token {
--  //      CHECK: "vhlo.outfeed_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.outfeed_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   outfeed_config = #vhlo.string_v1<"">
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.token_v1) -> !vhlo.token_v1
-   %0 = "stablehlo.outfeed"(%arg0, %arg1) : (tensor<f32>, !stablehlo.token) -> !stablehlo.token
-@@ -606,8 +651,9 @@
- }
- 
- // CHECK-LABEL: "default_recv"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_recv(%arg0: !stablehlo.token) -> (tensor<f32>, !stablehlo.token) {
--  //      CHECK: "vhlo.recv_v1"(%arg0) <{
-+  //      CHECK: "vhlo.recv_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME:   channel_type = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   is_host_transfer = #vhlo.bool_v1<false>
-@@ -619,8 +665,9 @@
- }
- 
- // CHECK-LABEL: "default_send"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @default_send(%arg0: tensor<f32>, %arg1: !stablehlo.token) -> !stablehlo.token {
--  //      CHECK: "vhlo.send_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.send_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME:   channel_type = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   is_host_transfer = #vhlo.bool_v1<false>
-@@ -632,8 +679,9 @@
- }
- 
- // CHECK-LABEL: "default_reduce_scatter"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_reduce_scatter(%arg0: tensor<16xf32>) -> tensor<16xf32> {
--  //               CHECK: "vhlo.reduce_scatter_v1"(%arg0) <{
-+  //               CHECK: "vhlo.reduce_scatter_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0], [1]]> : tensor<2x1xi64>>,
-   //          CHECK-SAME:   scatter_dimension = #vhlo.integer_v1<0 : i64>
-@@ -655,8 +703,9 @@
- }
- 
- // CHECK-LABEL: "default_reduce_window"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @default_reduce_window(%arg0: tensor<2x17x31x7xf32>, %arg1: tensor<f32>) -> tensor<2x16x30x7xf32> {
--  //               CHECK: "vhlo.reduce_window_v1"(%arg0, %arg1)  <{
-+  //               CHECK: "vhlo.reduce_window_v1"(%[[ARG0]], %[[ARG1]])  <{
-   //          CHECK-SAME:   base_dilations = #vhlo.tensor_v1<dense<1> : tensor<4xi64>>,
-   // CHECK-SAME{LITERAL}:   padding = #vhlo.tensor_v1<dense<0> : tensor<4x2xi64>>,
-   //          CHECK-SAME:   window_dilations = #vhlo.tensor_v1<dense<1> : tensor<4xi64>>,
-@@ -678,8 +727,9 @@
- }
- 
- // CHECK-LABEL: "default_scatter"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @default_scatter(%arg0: tensor<200x100x300xf32>, %arg1: tensor<10x2xi32>, %arg2: tensor<10x300xf32>) -> tensor<200x100x300xf32> {
--  //      CHECK: "vhlo.scatter_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.scatter_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   index_vector_dim = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   indices_are_sorted = #vhlo.bool_v1<false>,
-   // CHECK-SAME:   inserted_window_dims = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-@@ -707,8 +757,9 @@
- }
- 
- // CHECK-LABEL: "default_select_and_scatter"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @default_select_and_scatter(%arg0: tensor<10x24x24x64xf32>, %arg1: tensor<10x23x23x64xf32>, %arg2: tensor<f32>) -> tensor<10x24x24x64xf32> {
--  //      CHECK: "vhlo.select_and_scatter_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.select_and_scatter_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   padding = #vhlo.tensor_v1<dense<0> : tensor<4x2xi64>>,
-   // CHECK-SAME:   window_dimensions = #vhlo.tensor_v1<dense<[1, 2, 2, 1]> : tensor<4xi64>>,
-   // CHECK-SAME:   window_strides = #vhlo.tensor_v1<dense<1> : tensor<4xi64>>
-@@ -736,8 +787,9 @@
- }
- 
- // CHECK-LABEL: "default_sort"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_sort(%arg0: tensor<16xf32>) -> tensor<16xf32> {
--  //      CHECK: "vhlo.sort_v1"(%arg0) <{
-+  //      CHECK: "vhlo.sort_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   dimension = #vhlo.integer_v1<-1 : i64>
-   // CHECK-SAME:   is_stable = #vhlo.bool_v1<false>
-   // CHECK-SAME: }> ({
-@@ -756,29 +808,33 @@
- // ============ OPS ============
- 
- // CHECK-LABEL: "op_abs"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_abs(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.abs_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.abs_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.abs"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_add"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_add(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_after_all"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_after_all(%arg0: !stablehlo.token) -> !stablehlo.token {
--  // CHECK: "vhlo.after_all_v1"(%arg0) : (!vhlo.token_v1) -> !vhlo.token_v1
-+  // CHECK: "vhlo.after_all_v1"(%[[ARG0]]) : (!vhlo.token_v1) -> !vhlo.token_v1
-   %0 = "stablehlo.after_all"(%arg0) : (!stablehlo.token) -> !stablehlo.token
-   func.return %0 : !stablehlo.token
- }
- 
- // CHECK-LABEL: "op_all_gather"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_all_gather(%arg0: tensor<16x8xf32>) -> tensor<16x16xf32> {
--  //               CHECK: "vhlo.all_gather_v1"(%arg0) <{
-+  //               CHECK: "vhlo.all_gather_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   all_gather_dim = #vhlo.integer_v1<1 : i64>
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0], [1]]> : tensor<2x1xi64>>,
-@@ -794,8 +850,9 @@
- }
- 
- // CHECK-LABEL: "op_all_reduce"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_all_reduce(%arg0: tensor<f32>) -> tensor<f32> {
--  //               CHECK: "vhlo.all_reduce_v1"(%arg0) <{
-+  //               CHECK: "vhlo.all_reduce_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0], [1]]> : tensor<2x1xi64>>,
-   //          CHECK-SAME:   use_global_device_ids = #vhlo.bool_v1<true>
-@@ -817,8 +874,9 @@
- }
- 
- // CHECK-LABEL: "op_all_to_all"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_all_to_all(%arg0: tensor<4x16xf32>) -> tensor<16x4xf32> {
--  //               CHECK: "vhlo.all_to_all_v1"(%arg0) <{
-+  //               CHECK: "vhlo.all_to_all_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<1 : i64>,
-   //          CHECK-SAME:   concat_dimension = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0, 1, 2, 3]]> : tensor<1x4xi64>>,
-@@ -836,22 +894,25 @@
- }
- 
- // CHECK-LABEL: "op_and"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_and(%arg0: tensor<i1>, %arg1: tensor<i1>) -> tensor<i1> {
--  // CHECK: "vhlo.and_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-+  // CHECK: "vhlo.and_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-   %0 = "stablehlo.and"(%arg0, %arg1) : (tensor<i1>, tensor<i1>) -> tensor<i1>
-   func.return %0 : tensor<i1>
- }
- 
- // CHECK-LABEL: "op_atan2"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_atan2(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.atan2_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.atan2_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.atan2"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_batch_norm_grad"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}}, %[[ARG3:.*]]: {{.*}}, %[[ARG4:.*]]: {{.*}})
- func.func @op_batch_norm_grad(%arg0: tensor<16x16x16x16xf32>, %arg1: tensor<16xf32>, %arg2: tensor<16xf32>, %arg3: tensor<16xf32>, %arg4: tensor<16x16x16x16xf32>) -> (tensor<16x16x16x16xf32>, tensor<16xf32>, tensor<16xf32>) {
--  //      CHECK: "vhlo.batch_norm_grad_v1"(%arg0, %arg1, %arg2, %arg3, %arg4) <{
-+  //      CHECK: "vhlo.batch_norm_grad_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]], %[[ARG3]], %[[ARG4]]) <{
-   // CHECK-SAME:   epsilon = #vhlo.float_v1<1.000000e-03 : !vhlo.f32_v1>,
-   // CHECK-SAME:   feature_index = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<16x16x16x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x16x16x16x!vhlo.f32_v1>) -> (!vhlo.tensor_v1<16x16x16x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>)
-@@ -863,8 +924,9 @@
- }
- 
- // CHECK-LABEL: "op_batch_norm_inference"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}}, %[[ARG3:.*]]: {{.*}}, %[[ARG4:.*]]: {{.*}})
- func.func @op_batch_norm_inference(%arg0: tensor<16x16x16x16xf32>, %arg1: tensor<16xf32>, %arg2: tensor<16xf32>, %arg3: tensor<16xf32>, %arg4: tensor<16xf32>) -> tensor<16x16x16x16xf32> {
--  //      CHECK: "vhlo.batch_norm_inference_v1"(%arg0, %arg1, %arg2, %arg3, %arg4) <{
-+  //      CHECK: "vhlo.batch_norm_inference_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]], %[[ARG3]], %[[ARG4]]) <{
-   // CHECK-SAME:   epsilon = #vhlo.float_v1<1.000000e-03 : !vhlo.f32_v1>,
-   // CHECK-SAME:   feature_index = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<16x16x16x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x16x16x16x!vhlo.f32_v1>
-@@ -876,8 +938,9 @@
- }
- 
- // CHECK-LABEL: "op_batch_norm_training"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_batch_norm_training(%arg0: tensor<16x16x16x16xf32>, %arg1: tensor<16xf32>, %arg2: tensor<16xf32>) -> (tensor<16x16x16x16xf32>, tensor<16xf32>, tensor<16xf32>) {
--  //      CHECK: "vhlo.batch_norm_training_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.batch_norm_training_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   epsilon = #vhlo.float_v1<1.000000e-03 : !vhlo.f32_v1>,
-   // CHECK-SAME:   feature_index = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<16x16x16x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>) -> (!vhlo.tensor_v1<16x16x16x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>)
-@@ -889,15 +952,17 @@
- }
- 
- // CHECK-LABEL: "op_bitcast_convert"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_bitcast_convert(%arg0: tensor<i32>) -> tensor<f32> {
--  // CHECK: "vhlo.bitcast_convert_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.bitcast_convert_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.bitcast_convert"(%arg0) : (tensor<i32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_broadcast_in_dim"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_broadcast_in_dim(%arg0: tensor<16xf32>) -> tensor<16x16xf32> {
--  //      CHECK: "vhlo.broadcast_in_dim_v1"(%arg0) <{
-+  //      CHECK: "vhlo.broadcast_in_dim_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   broadcast_dimensions = #vhlo.tensor_v1<dense<1> : tensor<1xi64>>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x16x!vhlo.f32_v1>
-   %0 = "stablehlo.broadcast_in_dim"(%arg0) {
-@@ -907,8 +972,9 @@
- }
- 
- // CHECK-LABEL: "op_broadcast"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_broadcast(%arg0: tensor<16xf32>) -> tensor<16x16xf32> {
--  //      CHECK: "vhlo.broadcast_v1"(%arg0) <{
-+  //      CHECK: "vhlo.broadcast_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   broadcast_sizes = #vhlo.tensor_v1<dense<16> : tensor<1xi64>>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x16x!vhlo.f32_v1>
-   %0 = "stablehlo.broadcast"(%arg0) {
-@@ -918,9 +984,10 @@
- }
- 
- // CHECK-LABEL: "op_case"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_case(%arg0: tensor<i32>, %arg1: tensor<f32>) -> tensor<f32> {
--  //      CHECK: "vhlo.case_v1"(%arg0) ({
--  // CHECK-NEXT:   "vhlo.return_v1"(%arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-+  //      CHECK: "vhlo.case_v1"(%[[ARG0]]) ({
-+  // CHECK-NEXT:   "vhlo.return_v1"(%[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-   // CHECK-NEXT: }) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.case"(%arg0) ({
-     "stablehlo.return"(%arg1) : (tensor<f32>) -> ()
-@@ -929,22 +996,25 @@
- }
- 
- // CHECK-LABEL: "op_cbrt"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_cbrt(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.cbrt_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.cbrt_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.cbrt"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_ceil"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_ceil(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.ceil_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.ceil_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.ceil"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_cholesky"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_cholesky(%arg0: tensor<1x16x16xf32>) -> tensor<1x16x16xf32> {
--  //      CHECK: "vhlo.cholesky_v1"(%arg0) <{
-+  //      CHECK: "vhlo.cholesky_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   lower = #vhlo.bool_v1<true>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<1x16x16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<1x16x16x!vhlo.f32_v1>
-   %0 = "stablehlo.cholesky"(%arg0) {
-@@ -954,22 +1024,25 @@
- }
- 
- // CHECK-LABEL: "op_clamp"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_clamp(%arg0: tensor<f32>, %arg1: tensor<f32>, %arg2: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.clamp_v1"(%arg0, %arg1, %arg2) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.clamp_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.clamp"(%arg0, %arg1, %arg2) : (tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_count_leading_zeros"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_count_leading_zeros(%arg0: tensor<i32>) -> tensor<i32> {
--  // CHECK: "vhlo.count_leading_zeros_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-+  // CHECK: "vhlo.count_leading_zeros_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-   %0 = "stablehlo.count_leading_zeros"(%arg0) : (tensor<i32>) -> tensor<i32>
-   func.return %0 : tensor<i32>
- }
- 
- // CHECK-LABEL: "op_collective_permute"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_collective_permute(%arg0: tensor<16x8xf32>) -> tensor<16x8xf32> {
--  //               CHECK: "vhlo.collective_permute_v1"(%arg0) <{
-+  //               CHECK: "vhlo.collective_permute_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME{LITERAL}:   source_target_pairs = #vhlo.tensor_v1<dense<[[0, 1], [1, 2], [2, 3]]> : tensor<3x2xi64>>
-   //          CHECK-SAME: }> : (!vhlo.tensor_v1<16x8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x8x!vhlo.f32_v1>
-@@ -981,8 +1054,9 @@
- }
- 
- // CHECK-LABEL: "op_compare"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_compare(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
--  //      CHECK: "vhlo.compare_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.compare_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   compare_type = #vhlo<comparison_type_v1 TOTALORDER>,
-   // CHECK-SAME:   comparison_direction = #vhlo<comparison_direction_v1 EQ>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-@@ -994,15 +1068,17 @@
- }
- 
- // CHECK-LABEL: "op_complex"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_complex(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<complex<f32>> {
--  // CHECK: "vhlo.complex_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>
-+  // CHECK: "vhlo.complex_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>
-   %0 = "stablehlo.complex"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<complex<f32>>
-   func.return %0 : tensor<complex<f32>>
- }
- 
- // CHECK-LABEL: "op_concatenate"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_concatenate(%arg0: tensor<8xf32>, %arg1: tensor<8xf32>) -> tensor<16xf32> {
--  //      CHECK: "vhlo.concatenate_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.concatenate_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   dimension = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<8x!vhlo.f32_v1>, !vhlo.tensor_v1<8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x!vhlo.f32_v1>
-   %0 = "stablehlo.concatenate"(%arg0, %arg1) {
-@@ -1012,6 +1088,7 @@
- }
- 
- // CHECK-LABEL: "op_constant"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_constant(%arg0: tensor<f32>) -> tensor<f32> {
-   //      CHECK: "vhlo.constant_v1"() <{
-   // CHECK-SAME:   value = #vhlo.tensor_v1<dense<0.000000e+00> : tensor<f32>>
-@@ -1023,15 +1100,17 @@
- }
- 
- // CHECK-LABEL: "op_convert"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_convert(%arg0: tensor<i32>) -> tensor<f32> {
--  // CHECK: "vhlo.convert_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.convert_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.convert"(%arg0) : (tensor<i32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_convolution"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_convolution(%arg0: tensor<1x8x8x207xf32>, %arg1: tensor<3x3x207x16xf32>) -> tensor<1x7x7x16xf32> {
--  //      CHECK: "vhlo.convolution_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.convolution_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   batch_group_count = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   feature_group_count = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   input_batch_dimension = #vhlo.integer_v1<0 : i64>,
-@@ -1065,8 +1144,9 @@
- }
- 
- // CHECK-LABEL: "op_cosine"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_cosine(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.cosine_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.cosine_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.cosine"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
-@@ -1079,8 +1159,9 @@
- }
- 
- // CHECK-LABEL: "op_cross_replica_sum"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_cross_replica_sum(%arg0: tensor<f32>) -> tensor<f32> {
--  //               CHECK: "vhlo.cross-replica-sum_v1"(%arg0) <{
-+  //               CHECK: "vhlo.cross-replica-sum_v1"(%[[ARG0]]) <{
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0], [1]]> : tensor<2x1xi64>>
-   //          CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.cross-replica-sum"(%arg0) {
-@@ -1090,8 +1171,9 @@
- }
- 
- // CHECK-LABEL: "op_custom_call"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_custom_call(%arg0: tensor<f32>) -> tensor<f32> {
--  //      CHECK: "vhlo.custom_call_v1"(%arg0) <{
-+  //      CHECK: "vhlo.custom_call_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   api_version = #vhlo<api_version_v1 API_VERSION_STATUS_RETURNING>,
-   // CHECK-SAME:   backend_config = #vhlo.string_v1<"\08\03\1A\02">,
-   // CHECK-SAME:   call_target_name = #vhlo.string_v1<"foo">,
-@@ -1122,15 +1204,17 @@
- }
- 
- // CHECK-LABEL: "op_divide"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_divide(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.divide_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.divide_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.divide"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_dot_general"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_dot_general(%arg0: tensor<8x8x16xf32>, %arg1: tensor<8x16x8xf32>) -> tensor<8x8x8xf32> {
--  //      CHECK: "vhlo.dot_general_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.dot_general_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   lhs_batching_dimensions = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>,
-   // CHECK-SAME:   lhs_contracting_dimensions = #vhlo.tensor_v1<dense<2> : tensor<1xi64>>,
-   // CHECK-SAME:   precision_config = #vhlo.array_v1<[#vhlo<precision_v1 HIGHEST>, #vhlo<precision_v1 HIGHEST>]>,
-@@ -1150,8 +1234,9 @@
- }
- 
- // CHECK-LABEL: "op_dot"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_dot(%arg0: tensor<8x16xf32>, %arg1: tensor<16x8xf32>) -> tensor<8x8xf32> {
--  //      CHECK: "vhlo.dot_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.dot_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   precision_config = #vhlo.array_v1<[#vhlo<precision_v1 HIGHEST>, #vhlo<precision_v1 HIGHEST>]>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<8x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<8x8x!vhlo.f32_v1>
-   %0 = "stablehlo.dot"(%arg0, %arg1) {
-@@ -1161,8 +1246,9 @@
- }
- 
- // CHECK-LABEL: "op_dynamic_broadcast_in_dim"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_dynamic_broadcast_in_dim(%arg0: tensor<?x?xf32>, %arg1: tensor<2xindex>) -> tensor<?x?xf32> {
--  //      CHECK: "vhlo.dynamic_broadcast_in_dim_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.dynamic_broadcast_in_dim_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   broadcast_dimensions = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-   // CHECK-SAME:   known_expanding_dimensions = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>,
-   // CHECK-SAME:   known_nonexpanding_dimensions = #vhlo.tensor_v1<dense<1> : tensor<1xi64>>
-@@ -1176,8 +1262,9 @@
- }
- 
- // CHECK-LABEL: "op_dynamic_conv"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_dynamic_conv(%arg0: tensor<1x8x8x207xf32>, %arg1: tensor<3x3x207x16xf32>, %arg2: tensor<4xi32>) -> tensor<1x?x?x16xf32> {
--  //      CHECK: "vhlo.dynamic_conv_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.dynamic_conv_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   batch_group_count = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   feature_group_count = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   input_batch_dimension = #vhlo.integer_v1<0 : i64>,
-@@ -1211,8 +1298,9 @@
- }
- 
- // CHECK-LABEL: "op_dynamic_gather"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_dynamic_gather(%arg0 : tensor<2x4x9xf32>, %arg1 : tensor<1x5x2xi32>, %arg2 : tensor<3xi32>) -> tensor<1x5x8xf32> {
--  //      CHECK: "vhlo.dynamic_gather_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.dynamic_gather_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   collapsed_slice_dims = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-   // CHECK-SAME:   index_vector_dim = #vhlo.integer_v1<2 : i64>,
-   // CHECK-SAME:   indices_are_sorted = #vhlo.bool_v1<true>,
-@@ -1232,8 +1320,9 @@
- }
- 
- // CHECK-LABEL: "op_dynamic_iota"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_dynamic_iota(%arg0: tensor<1xindex>) -> tensor<?xf32> {
--  //      CHECK: "vhlo.dynamic_iota_v1"(%arg0) <{
-+  //      CHECK: "vhlo.dynamic_iota_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   iota_dimension = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<1x!vhlo.index_v1>) -> !vhlo.tensor_v1<?x!vhlo.f32_v1>
-   %0 = "stablehlo.dynamic_iota"(%arg0) {
-@@ -1243,22 +1332,25 @@
- }
- 
- // CHECK-LABEL: "op_dynamic_pad"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}}, %[[ARG3:.*]]: {{.*}}, %[[ARG4:.*]]: {{.*}})
- func.func @op_dynamic_pad(%arg0: tensor<?xf32>, %arg1: tensor<f32>, %arg2: tensor<1xindex>, %arg3: tensor<1xindex>, %arg4: tensor<1xindex>) -> tensor<?xf32> {
--  // CHECK: "vhlo.dynamic_pad_v1"(%arg0, %arg1, %arg2, %arg3, %arg4) : (!vhlo.tensor_v1<?x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>) -> !vhlo.tensor_v1<?x!vhlo.f32_v1>
-+  // CHECK: "vhlo.dynamic_pad_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]], %[[ARG3]], %[[ARG4]]) : (!vhlo.tensor_v1<?x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>) -> !vhlo.tensor_v1<?x!vhlo.f32_v1>
-   %0 = "stablehlo.dynamic_pad"(%arg0, %arg1, %arg2, %arg3, %arg4) : (tensor<?xf32>, tensor<f32>, tensor<1xindex>, tensor<1xindex>, tensor<1xindex>) -> tensor<?xf32>
-   func.return %0 : tensor<?xf32>
- }
- 
- // CHECK-LABEL: "op_dynamic_reshape"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_dynamic_reshape(%arg0: tensor<16xf32>, %arg1: tensor<2xindex>) -> tensor<?x?xf32> {
--  // CHECK: "vhlo.dynamic_reshape_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<2x!vhlo.index_v1>) -> !vhlo.tensor_v1<?x?x!vhlo.f32_v1>
-+  // CHECK: "vhlo.dynamic_reshape_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<2x!vhlo.index_v1>) -> !vhlo.tensor_v1<?x?x!vhlo.f32_v1>
-   %0 = "stablehlo.dynamic_reshape"(%arg0, %arg1) : (tensor<16xf32>, tensor<2xindex>) -> tensor<?x?xf32>
-   func.return %0 : tensor<?x?xf32>
- }
- 
- // CHECK-LABEL: "op_dynamic_slice"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_dynamic_slice(%arg0: tensor<16xf32>, %arg1: tensor<i64>) -> tensor<4xf32> {
--  //      CHECK: "vhlo.dynamic_slice_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.dynamic_slice_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   slice_sizes = #vhlo.tensor_v1<dense<4> : tensor<1xi64>>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.i64_v1>) -> !vhlo.tensor_v1<4x!vhlo.f32_v1>
-   %0 = "stablehlo.dynamic_slice"(%arg0, %arg1) {
-@@ -1268,15 +1360,17 @@
- }
- 
- // CHECK-LABEL: "op_dynamic_update_slice"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_dynamic_update_slice(%arg0: tensor<16xf32>, %arg1: tensor<4xf32>, %arg2: tensor<i64>) -> tensor<16xf32> {
--  // CHECK: "vhlo.dynamic_update_slice_v1"(%arg0, %arg1, %arg2) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<4x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.i64_v1>) -> !vhlo.tensor_v1<16x!vhlo.f32_v1>
-+  // CHECK: "vhlo.dynamic_update_slice_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<4x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.i64_v1>) -> !vhlo.tensor_v1<16x!vhlo.f32_v1>
-   %0 = "stablehlo.dynamic_update_slice"(%arg0, %arg1, %arg2) : (tensor<16xf32>, tensor<4xf32>, tensor<i64>) -> tensor<16xf32>
-   func.return %0 : tensor<16xf32>
- }
- 
- // CHECK-LABEL: "op_einsum"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_einsum(%arg0: tensor<8x16xf32>, %arg1: tensor<16x8xf32>) -> tensor<8x8xf32> {
--  //      CHECK: "vhlo.einsum_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.einsum_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   einsum_config = #vhlo.string_v1<"ab,bc->ac">
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<8x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<8x8x!vhlo.f32_v1>
-   %0 = "stablehlo.einsum"(%arg0, %arg1) {
-@@ -1286,22 +1380,25 @@
- }
- 
- // CHECK-LABEL: "op_exponential_minus_one"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_exponential_minus_one(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.exponential_minus_one_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.exponential_minus_one_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.exponential_minus_one"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_exponential"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_exponential(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.exponential_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.exponential_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.exponential"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_fft"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_fft(%arg0: tensor<16xcomplex<f32>>) -> tensor<16xcomplex<f32>> {
--  //      CHECK: "vhlo.fft_v1"(%arg0) <{
-+  //      CHECK: "vhlo.fft_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   fft_length = #vhlo.tensor_v1<dense<16> : tensor<1xi64>>,
-   // CHECK-SAME:   fft_type = #vhlo<fft_type_v1 FFT>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<16x!vhlo.complex_v1<!vhlo.f32_v1>>) -> !vhlo.tensor_v1<16x!vhlo.complex_v1<!vhlo.f32_v1>>
-@@ -1313,8 +1410,9 @@
- }
- 
- // CHECK-LABEL: "op_floor"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_floor(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.floor_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.floor_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.floor"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
-@@ -1327,16 +1425,17 @@
-   // CHECK-SAME:   sym_name = #vhlo.string_v1<"op_func">,
-   // CHECK-SAME:   sym_visibility = #vhlo.string_v1<"private">
-   // CHECK-SAME: }> ({
--  // CHECK-NEXT:   ^[[BB:bb.*]](%arg0: !vhlo.tensor_v1<!vhlo.f32_v1>):
--  // CHECK-NEXT:     "vhlo.return_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-+  // CHECK-NEXT:   ^[[BB:bb.*]](%[[ARG0:.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
-+  // CHECK-NEXT:     "vhlo.return_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-   // CHECK-NEXT: }) : () -> ()
- 
-   func.return %arg0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_gather"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_gather(%arg0 : tensor<2x4x9xf32>, %arg1 : tensor<1x5x2xi32>) -> tensor<1x5x1xf32> {
--  //      CHECK: "vhlo.gather_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.gather_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   collapsed_slice_dims = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-   // CHECK-SAME:   index_vector_dim = #vhlo.integer_v1<2 : i64>,
-   // CHECK-SAME:   indices_are_sorted = #vhlo.bool_v1<true>,
-@@ -1358,8 +1457,9 @@
- }
- 
- // CHECK-LABEL: "op_get_dimension_size"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_get_dimension_size(%arg0: tensor<?xf32>) -> tensor<i32> {
--  //      CHECK: "vhlo.get_dimension_size_v1"(%arg0) <{
-+  //      CHECK: "vhlo.get_dimension_size_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   dimension = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<?x!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-   %0 = "stablehlo.get_dimension_size"(%arg0) {
-@@ -1369,8 +1469,9 @@
- }
- 
- // CHECK-LABEL: "op_get_tuple_element"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_get_tuple_element(%arg0: tuple<tensor<f32>, tensor<i32>>) -> tensor<f32> {
--  //      CHECK: "vhlo.get_tuple_element_v1"(%arg0) <{
-+  //      CHECK: "vhlo.get_tuple_element_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   index = #vhlo.integer_v1<0 : i32>
-   // CHECK-SAME: }> : (!vhlo.tuple_v1<!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.get_tuple_element"(%arg0) {
-@@ -1380,11 +1481,12 @@
- }
- 
- // CHECK-LABEL: "op_if"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_if(%arg0: tensor<i1>, %arg1: tensor<f32>, %arg2: tensor<f32>) -> tensor<f32> {
--  //      CHECK: "vhlo.if_v1"(%arg0) ({
--  // CHECK-NEXT:   "vhlo.return_v1"(%arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-+  //      CHECK: "vhlo.if_v1"(%[[ARG0]]) ({
-+  // CHECK-NEXT:   "vhlo.return_v1"(%[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-   // CHECK-NEXT: }, {
--  // CHECK-NEXT:   "vhlo.return_v1"(%arg2) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-+  // CHECK-NEXT:   "vhlo.return_v1"(%[[ARG2]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-   // CHECK-NEXT: }) : (!vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.if"(%arg0) ({
-     "stablehlo.return"(%arg1) : (tensor<f32>) -> ()
-@@ -1395,15 +1497,17 @@
- }
- 
- // CHECK-LABEL: "op_imag"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_imag(%arg0: tensor<complex<f32>>) -> tensor<f32> {
--  // CHECK: "vhlo.imag_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.imag_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.imag"(%arg0) : (tensor<complex<f32>>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_infeed"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_infeed(%arg0: !stablehlo.token) -> (tensor<f32>, !stablehlo.token) {
--  //               CHECK: "vhlo.infeed_v1"(%arg0) <{
-+  //               CHECK: "vhlo.infeed_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   infeed_config = #vhlo.string_v1<"foo">,
-   // CHECK-SAME{LITERAL}:   layout = #vhlo.array_v1<[#vhlo.array_v1<[]>]>
-   //          CHECK-SAME: }> : (!vhlo.token_v1) -> (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.token_v1)
-@@ -1426,36 +1530,41 @@
- }
- 
- // CHECK-LABEL: "op_is_finite"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_is_finite(%arg0: tensor<f32>) -> tensor<i1> {
--  // CHECK: "vhlo.is_finite_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-+  // CHECK: "vhlo.is_finite_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-   %0 = "stablehlo.is_finite"(%arg0) : (tensor<f32>) -> tensor<i1>
-   func.return %0 : tensor<i1>
- }
- 
- // CHECK-LABEL: "op_log"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_log(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.log_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.log_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.log"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_log_plus_one"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_log_plus_one(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.log_plus_one_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.log_plus_one_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.log_plus_one"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_logistic"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_logistic(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.logistic_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.logistic_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.logistic"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_map"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_map(%arg0: tensor<16xf32>) -> tensor<16xf32> {
--  //      CHECK: "vhlo.map_v1"(%arg0) <{
-+  //      CHECK: "vhlo.map_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   dimensions = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>
-   // CHECK-SAME: }> ({
-   // CHECK-NEXT:   ^[[BB:bb.*]](%[[ARG1:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
-@@ -1473,57 +1582,65 @@
- }
- 
- // CHECK-LABEL: "op_maximum"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_maximum(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.maximum_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.maximum_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.maximum"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_minimum"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_minimum(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.minimum_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.minimum_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.minimum"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_multiply"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_multiply(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.multiply_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.multiply_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.multiply"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_negate"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_negate(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.negate_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.negate_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.negate"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_not"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_not(%arg0: tensor<i1>) -> tensor<i1> {
--  // CHECK: "vhlo.not_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-+  // CHECK: "vhlo.not_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-   %0 = "stablehlo.not"(%arg0) : (tensor<i1>) -> tensor<i1>
-   func.return %0 : tensor<i1>
- }
- 
- // CHECK-LABEL: "op_optimization_barrier"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_optimization_barrier(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.optimization_barrier_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.optimization_barrier_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.optimization_barrier"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_or"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_or(%arg0: tensor<i1>, %arg1: tensor<i1>) -> tensor<i1> {
--  // CHECK: "vhlo.or_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-+  // CHECK: "vhlo.or_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-   %0 = "stablehlo.or"(%arg0, %arg1) : (tensor<i1>, tensor<i1>) -> tensor<i1>
-   func.return %0 : tensor<i1>
- }
- 
- // CHECK-LABEL: "op_outfeed"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_outfeed(%arg0: tensor<f32>, %arg1: !stablehlo.token) -> !stablehlo.token {
--  //      CHECK: "vhlo.outfeed_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.outfeed_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   outfeed_config = #vhlo.string_v1<"foo">
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.token_v1) -> !vhlo.token_v1
-   %0 = "stablehlo.outfeed"(%arg0, %arg1) {
-@@ -1533,8 +1650,9 @@
- }
- 
- // CHECK-LABEL: "op_pad"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_pad(%arg0: tensor<8xf32>, %arg1: tensor<f32>) -> tensor<16xf32> {
--  //      CHECK: "vhlo.pad_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.pad_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   edge_padding_high = #vhlo.tensor_v1<dense<4> : tensor<1xi64>>,
-   // CHECK-SAME:   edge_padding_low = #vhlo.tensor_v1<dense<4> : tensor<1xi64>>,
-   // CHECK-SAME:   interior_padding = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>
-@@ -1548,36 +1666,41 @@
- }
- 
- // CHECK-LABEL: "op_popcnt"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_popcnt(%arg0: tensor<i32>) -> tensor<i32> {
--  // CHECK: "vhlo.popcnt_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-+  // CHECK: "vhlo.popcnt_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-   %0 = "stablehlo.popcnt"(%arg0) : (tensor<i32>) -> tensor<i32>
-   func.return %0 : tensor<i32>
- }
- 
- // CHECK-LABEL: "op_power"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_power(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.power_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.power_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.power"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_real_dynamic_slice"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}}, %[[ARG3:.*]]: {{.*}})
- func.func @op_real_dynamic_slice(%arg0: tensor<?xf32>, %arg1: tensor<1xindex>, %arg2: tensor<1xindex>, %arg3: tensor<1xindex>) -> tensor<?xf32> {
--  // CHECK: "vhlo.real_dynamic_slice_v1"(%arg0, %arg1, %arg2, %arg3) : (!vhlo.tensor_v1<?x!vhlo.f32_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>) -> !vhlo.tensor_v1<?x!vhlo.f32_v1>
-+  // CHECK: "vhlo.real_dynamic_slice_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]], %[[ARG3]]) : (!vhlo.tensor_v1<?x!vhlo.f32_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>) -> !vhlo.tensor_v1<?x!vhlo.f32_v1>
-   %0 = "stablehlo.real_dynamic_slice"(%arg0, %arg1, %arg2, %arg3) : (tensor<?xf32>, tensor<1xindex>, tensor<1xindex>, tensor<1xindex>) -> tensor<?xf32>
-   func.return %0 : tensor<?xf32>
- }
- 
- // CHECK-LABEL: "op_real"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_real(%arg0: tensor<complex<f32>>) -> tensor<f32> {
--  // CHECK: "vhlo.real_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.real_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.real"(%arg0) : (tensor<complex<f32>>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_recv"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_recv(%arg0: !stablehlo.token) -> (tensor<f32>, !stablehlo.token) {
--  //      CHECK: "vhlo.recv_v1"(%arg0) <{
-+  //      CHECK: "vhlo.recv_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME:   channel_type = #vhlo.integer_v1<3 : i64>,
-   // CHECK-SAME:   is_host_transfer = #vhlo.bool_v1<true>
-@@ -1590,8 +1713,9 @@
- }
- 
- // CHECK-LABEL: "op_reduce"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_reduce(%arg0: tensor<16xf32>, %arg1: tensor<f32>) -> tensor<f32> {
--  //  CHECK: "vhlo.reduce_v1"(%arg0, %arg1)
-+  //  CHECK: "vhlo.reduce_v1"(%[[ARG0]], %[[ARG1]])
-   //  CHECK:   ^[[BB:bb.*]](%[[ARG1:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>, %[[ARG2:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
-   //  CHECK:     "vhlo.return_v1"(%[[VAL1:.*]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-   //  CHECK: }) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-@@ -1606,8 +1730,9 @@
- }
- 
- // CHECK-LABEL: "op_reduce_precision"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_reduce_precision(%arg0: tensor<f32>) -> tensor<f32> {
--  //      CHECK: "vhlo.reduce_precision_v1"(%arg0) <{
-+  //      CHECK: "vhlo.reduce_precision_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   exponent_bits = #vhlo.integer_v1<8 : i32>
-   // CHECK-SAME:   mantissa_bits = #vhlo.integer_v1<10 : i32>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-@@ -1619,8 +1744,9 @@
- }
- 
- // CHECK-LABEL: "op_reduce_scatter"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_reduce_scatter(%arg0: tensor<16xf32>) -> tensor<16xf32> {
--  //               CHECK: "vhlo.reduce_scatter_v1"(%arg0) <{
-+  //               CHECK: "vhlo.reduce_scatter_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0], [1]]> : tensor<2x1xi64>>,
-   //          CHECK-SAME:   scatter_dimension = #vhlo.integer_v1<0 : i64>
-@@ -1644,8 +1770,9 @@
- }
- 
- // CHECK-LABEL: "op_reduce_window"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_reduce_window(%arg0: tensor<2x17x31x7xf32>, %arg1: tensor<f32>) -> tensor<2x9x16x7xf32> {
--  //               CHECK: "vhlo.reduce_window_v1"(%arg0, %arg1) <{
-+  //               CHECK: "vhlo.reduce_window_v1"(%[[ARG0]], %[[ARG1]]) <{
-   //          CHECK-SAME:   base_dilations = #vhlo.tensor_v1<dense<[1, 2, 2, 1]> : tensor<4xi64>>,
-   // CHECK-SAME{LITERAL}:   padding = #vhlo.tensor_v1<dense<[[0, 0], [2, 0], [0, 2], [0, 0]]> : tensor<4x2xi64>>,
-   //          CHECK-SAME:   window_dilations = #vhlo.tensor_v1<dense<[1, 2, 2, 1]> : tensor<4xi64>>,
-@@ -1671,8 +1798,9 @@
- }
- 
- // CHECK-LABEL: "op_remainder"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_remainder(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.remainder_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.remainder_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.remainder"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
-@@ -1692,16 +1820,18 @@
- }
- 
- // CHECK-LABEL: "op_reshape"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_reshape(%arg0: tensor<16xf32>) -> tensor<4x4xf32> {
--  // CHECK: "vhlo.reshape_v1"(%arg0) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<4x4x!vhlo.f32_v1>
-+  // CHECK: "vhlo.reshape_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<4x4x!vhlo.f32_v1>
-   %0 = "stablehlo.reshape"(%arg0) : (tensor<16xf32>) -> tensor<4x4xf32>
-   func.return %0 : tensor<4x4xf32>
- }
- 
- // CHECK-LABEL: "op_return"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_return(%arg0: tensor<i32>, %arg1: tensor<f32>) -> tensor<f32> {
--  //      CHECK: "vhlo.case_v1"(%arg0) ({
--  // CHECK-NEXT:   "vhlo.return_v1"(%arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-+  //      CHECK: "vhlo.case_v1"(%[[ARG0]]) ({
-+  // CHECK-NEXT:   "vhlo.return_v1"(%[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-   // CHECK-NEXT: }) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.case"(%arg0) ({
-     "stablehlo.return"(%arg1) : (tensor<f32>) -> ()
-@@ -1710,8 +1840,9 @@
- }
- 
- // CHECK-LABEL: "op_reverse"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_reverse(%arg0: tensor<16xf32>) -> tensor<16xf32> {
--  //      CHECK: "vhlo.reverse_v1"(%arg0) <{
-+  //      CHECK: "vhlo.reverse_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   dimensions = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x!vhlo.f32_v1>
-   %0 = "stablehlo.reverse"(%arg0) {
-@@ -1721,8 +1852,9 @@
- }
- 
- // CHECK-LABEL: "op_rng_bit_generator"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_rng_bit_generator(%arg0: tensor<f32>) -> (tensor<f32>, tensor<f32>) {
--  //      CHECK: "vhlo.rng_bit_generator_v1"(%arg0) <{
-+  //      CHECK: "vhlo.rng_bit_generator_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   rng_algorithm = #vhlo<rng_algorithm_v1 PHILOX>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>)
-   %0:2 = "stablehlo.rng_bit_generator"(%arg0) {
-@@ -1732,8 +1864,9 @@
- }
- 
- // CHECK-LABEL: "op_rng"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_rng(%arg0: tensor<f32>, %arg1: tensor<f32>, %arg2: tensor<0xindex>) -> tensor<f32> {
--  //      CHECK: "vhlo.rng_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.rng_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   rng_distribution = #vhlo<rng_distribution_v1 NORMAL>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<0x!vhlo.index_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.rng"(%arg0, %arg1, %arg2) {
-@@ -1743,29 +1876,33 @@
- }
- 
- // CHECK-LABEL: "op_round_nearest_afz"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_round_nearest_afz(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.round_nearest_afz_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.round_nearest_afz_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.round_nearest_afz"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_round_nearest_even"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_round_nearest_even(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.round_nearest_even_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.round_nearest_even_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.round_nearest_even"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_rsqrt"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_rsqrt(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.rsqrt_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.rsqrt_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.rsqrt"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_scatter"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_scatter(%arg0: tensor<200x100x300xf32>, %arg1: tensor<10x2xi32>, %arg2: tensor<10x300xf32>) -> tensor<200x100x300xf32> {
--  //      CHECK: "vhlo.scatter_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.scatter_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   index_vector_dim = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   indices_are_sorted = #vhlo.bool_v1<true>,
-   // CHECK-SAME:   inserted_window_dims = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-@@ -1795,8 +1932,9 @@
- }
- 
- // CHECK-LABEL: "op_select_and_scatter"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_select_and_scatter(%arg0: tensor<10x24x24x64xf32>, %arg1: tensor<12x13x13x66xf32>, %arg2: tensor<f32>) -> tensor<10x24x24x64xf32> {
--  //      CHECK: "vhlo.select_and_scatter_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.select_and_scatter_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   padding = #vhlo.tensor_v1<dense<1> : tensor<4x2xi64>>,
-   // CHECK-SAME:   window_dimensions = #vhlo.tensor_v1<dense<[1, 2, 2, 1]> : tensor<4xi64>>,
-   // CHECK-SAME:   window_strides = #vhlo.tensor_v1<dense<[1, 2, 2, 1]> : tensor<4xi64>>
-@@ -1826,15 +1964,17 @@
- }
- 
- // CHECK-LABEL: "op_select"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_select(%arg0: tensor<i1>, %arg1: tensor<f32>, %arg2: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.select_v1"(%arg0, %arg1, %arg2) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.select_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.select"(%arg0, %arg1, %arg2) : (tensor<i1>, tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_send"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_send(%arg0: tensor<f32>, %arg1: !stablehlo.token) -> !stablehlo.token {
--  //      CHECK: "vhlo.send_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.send_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME:   channel_type = #vhlo.integer_v1<2 : i64>,
-   // CHECK-SAME:   is_host_transfer = #vhlo.bool_v1<true>
-@@ -1847,8 +1987,9 @@
- }
- 
- // CHECK-LABEL: "op_set_dimension_size"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_set_dimension_size(%arg0: tensor<?xf32>, %arg1: tensor<i32>) -> tensor<16xf32> {
--  //      CHECK: "vhlo.set_dimension_size_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.set_dimension_size_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   dimension = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<?x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<16x!vhlo.f32_v1>
-   %0 = "stablehlo.set_dimension_size"(%arg0, %arg1) {
-@@ -1858,43 +1999,49 @@
- }
- 
- // CHECK-LABEL: "op_shift_left"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_shift_left(%arg0: tensor<i32>, %arg1: tensor<i32>) -> tensor<i32> {
--  // CHECK: "vhlo.shift_left_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.i32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-+  // CHECK: "vhlo.shift_left_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-   %0 = "stablehlo.shift_left"(%arg0, %arg1) : (tensor<i32>, tensor<i32>) -> tensor<i32>
-   func.return %0 : tensor<i32>
- }
- 
- // CHECK-LABEL: "op_shift_right_arithmetic"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_shift_right_arithmetic(%arg0: tensor<i32>, %arg1: tensor<i32>) -> tensor<i32> {
--  // CHECK: "vhlo.shift_right_arithmetic_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.i32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-+  // CHECK: "vhlo.shift_right_arithmetic_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-   %0 = "stablehlo.shift_right_arithmetic"(%arg0, %arg1) : (tensor<i32>, tensor<i32>) -> tensor<i32>
-   func.return %0 : tensor<i32>
- }
- 
- // CHECK-LABEL: "op_shift_right_logical"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_shift_right_logical(%arg0: tensor<i32>, %arg1: tensor<i32>) -> tensor<i32> {
--  // CHECK: "vhlo.shift_right_logical_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.i32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-+  // CHECK: "vhlo.shift_right_logical_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-   %0 = "stablehlo.shift_right_logical"(%arg0, %arg1) : (tensor<i32>, tensor<i32>) -> tensor<i32>
-   func.return %0 : tensor<i32>
- }
- 
- // CHECK-LABEL: "op_sign"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_sign(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.sign_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.sign_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.sign"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_sine"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_sine(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.sine_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.sine_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.sine"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_slice"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_slice(%arg0: tensor<16xf32>) -> tensor<4xf32> {
--  //      CHECK: "vhlo.slice_v1"(%arg0) <{
-+  //      CHECK: "vhlo.slice_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   limit_indices = #vhlo.tensor_v1<dense<4> : tensor<1xi64>>,
-   // CHECK-SAME:   start_indices = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>,
-   // CHECK-SAME:   strides = #vhlo.tensor_v1<dense<1> : tensor<1xi64>>
-@@ -1908,8 +2055,9 @@
- }
- 
- // CHECK-LABEL: "op_sort"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_sort(%arg0: tensor<16xf32>) -> tensor<16xf32> {
--  //      CHECK: "vhlo.sort_v1"(%arg0) <{
-+  //      CHECK: "vhlo.sort_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   dimension = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME:   is_stable = #vhlo.bool_v1<true>
-   // CHECK-SAME: }> ({
-@@ -1929,29 +2077,33 @@
- }
- 
- // CHECK-LABEL: "op_sqrt"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_sqrt(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.sqrt_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.sqrt_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.sqrt"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_subtract"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_subtract(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.subtract_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.subtract_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.subtract"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_tanh"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_tanh(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.tanh_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.tanh_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.tanh"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_torch_index_select"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_torch_index_select(%arg0: tensor<5x1x5xf32>, %arg1: tensor<2xi32>) ->  tensor<2x1x5xf32> {
--  //      CHECK: "vhlo.torch_index_select_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.torch_index_select_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   batch_dims = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME:   dim = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<5x1x5x!vhlo.f32_v1>, !vhlo.tensor_v1<2x!vhlo.i32_v1>) -> !vhlo.tensor_v1<2x1x5x!vhlo.f32_v1>
-@@ -1963,8 +2115,9 @@
- }
- 
- // CHECK-LABEL: "op_trace"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_trace(%arg0: tensor<f32>) {
--  //      CHECK: "vhlo.trace_v1"(%arg0) <{
-+  //      CHECK: "vhlo.trace_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   tag = #vhlo.string_v1<"foo">
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-   "stablehlo.trace"(%arg0) {
-@@ -1974,8 +2127,9 @@
- }
- 
- // CHECK-LABEL: "op_transpose"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_transpose(%arg0: tensor<16x8xf32>) ->  tensor<8x16xf32> {
--  //      CHECK: "vhlo.transpose_v1"(%arg0) <{
-+  //      CHECK: "vhlo.transpose_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   permutation = #vhlo.tensor_v1<dense<[1, 0]> : tensor<2xi64>>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<16x8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<8x16x!vhlo.f32_v1>
-   %0 = "stablehlo.transpose"(%arg0) {
-@@ -1985,8 +2139,9 @@
- }
- 
- // CHECK-LABEL: "op_triangular_solve"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_triangular_solve(%arg0: tensor<16x16xf32>, %arg1: tensor<16x16xf32>) ->  tensor<16x16xf32> {
--  //      CHECK: "vhlo.triangular_solve_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.triangular_solve_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   left_side = #vhlo.bool_v1<true>,
-   // CHECK-SAME:   lower = #vhlo.bool_v1<true>,
-   // CHECK-SAME:   transpose_a = #vhlo<transpose_v1 NO_TRANSPOSE>,
-@@ -2002,15 +2157,17 @@
- }
- 
- // CHECK-LABEL: "op_tuple"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_tuple(%arg0: tensor<f32>) -> tuple<tensor<f32>> {
--  // CHECK: "vhlo.tuple_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tuple_v1<!vhlo.tensor_v1<!vhlo.f32_v1>>
-+  // CHECK: "vhlo.tuple_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tuple_v1<!vhlo.tensor_v1<!vhlo.f32_v1>>
-   %0 = "stablehlo.tuple"(%arg0) : (tensor<f32>) -> tuple<tensor<f32>>
-   func.return %0 : tuple<tensor<f32>>
- }
- 
- // CHECK-LABEL: "op_unary_einsum"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_unary_einsum(%arg0: tensor<8x16xf32>) -> tensor<8xf32> {
--  //      CHECK: "vhlo.unary_einsum_v1"(%arg0) <{
-+  //      CHECK: "vhlo.unary_einsum_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   einsum_config = #vhlo.string_v1<"ab->a">
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<8x16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<8x!vhlo.f32_v1>
-   %0 = "stablehlo.unary_einsum"(%arg0) {
-@@ -2020,22 +2177,25 @@
- }
- 
- // CHECK-LABEL: "op_uniform_dequantize"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_uniform_dequantize(%arg0: tensor<!quant.uniform<i8:f32, 34.0:16>>) -> tensor<f32> {
--  // CHECK: "vhlo.uniform_dequantize_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.uniform_dequantize_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.uniform_dequantize"(%arg0) : (tensor<!quant.uniform<i8:f32, 34.0:16>>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_uniform_quantize"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_uniform_quantize(%arg0: tensor<f32>) -> tensor<!quant.uniform<i8:f32, 34.0:16>> {
--  // CHECK: "vhlo.uniform_quantize_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>
-+  // CHECK: "vhlo.uniform_quantize_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>
-   %0 = "stablehlo.uniform_quantize"(%arg0) : (tensor<f32>) -> tensor<!quant.uniform<i8:f32, 34.0:16>>
-   func.return %0 : tensor<!quant.uniform<i8:f32, 34.0:16>>
- }
- 
- // CHECK-LABEL: "op_while"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_while(%arg0: tensor<i1>) -> tensor<i1> {
--  //      CHECK: "vhlo.while_v1"(%arg0) ({
-+  //      CHECK: "vhlo.while_v1"(%[[ARG0]]) ({
-   // CHECK-NEXT:   ^[[BB:bb.*]](%[[ARG1:arg.*]]: !vhlo.tensor_v1<!vhlo.bool_v1>):
-   // CHECK-NEXT:     "vhlo.return_v1"(%[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>) -> ()
-   // CHECK-NEXT:   }, {
-@@ -2053,8 +2213,9 @@
- }
- 
- // CHECK-LABEL: "op_xor"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_xor(%arg0: tensor<i1>, %arg1: tensor<i1>) -> tensor<i1> {
--  // CHECK: "vhlo.xor_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-+  // CHECK: "vhlo.xor_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-   %0 = "stablehlo.xor"(%arg0, %arg1) : (tensor<i1>, tensor<i1>) -> tensor<i1>
-   func.return %0 : tensor<i1>
- }
-@@ -2062,190 +2223,217 @@
- // ============ TYPES ============
- 
- // CHECK-LABEL: "type_i1"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_i1(%arg0: tensor<i1>, %arg1: tensor<i1>) -> tensor<i1> {
--  // CHECK: "vhlo.and_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-+  // CHECK: "vhlo.and_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-   %0 = "stablehlo.and"(%arg0, %arg1) : (tensor<i1>, tensor<i1>) -> tensor<i1>
-   func.return %0 : tensor<i1>
- }
- 
- // CHECK-LABEL: "type_i4"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_i4(%arg0: tensor<i4>, %arg1: tensor<i4>) -> tensor<i4> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.i4_v1>, !vhlo.tensor_v1<!vhlo.i4_v1>) -> !vhlo.tensor_v1<!vhlo.i4_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i4_v1>, !vhlo.tensor_v1<!vhlo.i4_v1>) -> !vhlo.tensor_v1<!vhlo.i4_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<i4>, tensor<i4>) -> tensor<i4>
-   func.return %0 : tensor<i4>
- }
- 
- // CHECK-LABEL: "type_i8"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_i8(%arg0: tensor<i8>, %arg1: tensor<i8>) -> tensor<i8> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.i8_v1>, !vhlo.tensor_v1<!vhlo.i8_v1>) -> !vhlo.tensor_v1<!vhlo.i8_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i8_v1>, !vhlo.tensor_v1<!vhlo.i8_v1>) -> !vhlo.tensor_v1<!vhlo.i8_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<i8>, tensor<i8>) -> tensor<i8>
-   func.return %0 : tensor<i8>
- }
- 
- // CHECK-LABEL: "type_i16"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_i16(%arg0: tensor<i16>, %arg1: tensor<i16>) -> tensor<i16> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.i16_v1>, !vhlo.tensor_v1<!vhlo.i16_v1>) -> !vhlo.tensor_v1<!vhlo.i16_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i16_v1>, !vhlo.tensor_v1<!vhlo.i16_v1>) -> !vhlo.tensor_v1<!vhlo.i16_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<i16>, tensor<i16>) -> tensor<i16>
-   func.return %0 : tensor<i16>
- }
- 
- // CHECK-LABEL: "type_i32"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_i32(%arg0: tensor<i32>, %arg1: tensor<i32>) -> tensor<i32> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.i32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<i32>, tensor<i32>) -> tensor<i32>
-   func.return %0 : tensor<i32>
- }
- 
- // CHECK-LABEL: "type_i64"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_i64(%arg0: tensor<i64>, %arg1: tensor<i64>) -> tensor<i64> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.i64_v1>, !vhlo.tensor_v1<!vhlo.i64_v1>) -> !vhlo.tensor_v1<!vhlo.i64_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i64_v1>, !vhlo.tensor_v1<!vhlo.i64_v1>) -> !vhlo.tensor_v1<!vhlo.i64_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<i64>, tensor<i64>) -> tensor<i64>
-   func.return %0 : tensor<i64>
- }
- 
- // CHECK-LABEL: "type_ui4"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_ui4(%arg0: tensor<ui4>, %arg1: tensor<ui4>) -> tensor<ui4> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.ui4_v1>, !vhlo.tensor_v1<!vhlo.ui4_v1>) -> !vhlo.tensor_v1<!vhlo.ui4_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.ui4_v1>, !vhlo.tensor_v1<!vhlo.ui4_v1>) -> !vhlo.tensor_v1<!vhlo.ui4_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<ui4>, tensor<ui4>) -> tensor<ui4>
-   func.return %0 : tensor<ui4>
- }
- 
- // CHECK-LABEL: "type_ui8"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_ui8(%arg0: tensor<ui8>, %arg1: tensor<ui8>) -> tensor<ui8> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.ui8_v1>, !vhlo.tensor_v1<!vhlo.ui8_v1>) -> !vhlo.tensor_v1<!vhlo.ui8_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.ui8_v1>, !vhlo.tensor_v1<!vhlo.ui8_v1>) -> !vhlo.tensor_v1<!vhlo.ui8_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<ui8>, tensor<ui8>) -> tensor<ui8>
-   func.return %0 : tensor<ui8>
- }
- 
- // CHECK-LABEL: "type_ui16"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_ui16(%arg0: tensor<ui16>, %arg1: tensor<ui16>) -> tensor<ui16> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.ui16_v1>, !vhlo.tensor_v1<!vhlo.ui16_v1>) -> !vhlo.tensor_v1<!vhlo.ui16_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.ui16_v1>, !vhlo.tensor_v1<!vhlo.ui16_v1>) -> !vhlo.tensor_v1<!vhlo.ui16_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<ui16>, tensor<ui16>) -> tensor<ui16>
-   func.return %0 : tensor<ui16>
- }
- 
- // CHECK-LABEL: "type_ui32"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_ui32(%arg0: tensor<ui32>, %arg1: tensor<ui32>) -> tensor<ui32> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.ui32_v1>, !vhlo.tensor_v1<!vhlo.ui32_v1>) -> !vhlo.tensor_v1<!vhlo.ui32_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.ui32_v1>, !vhlo.tensor_v1<!vhlo.ui32_v1>) -> !vhlo.tensor_v1<!vhlo.ui32_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<ui32>, tensor<ui32>) -> tensor<ui32>
-   func.return %0 : tensor<ui32>
- }
- 
- // CHECK-LABEL: "type_ui64"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_ui64(%arg0: tensor<ui64>, %arg1: tensor<ui64>) -> tensor<ui64> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.ui64_v1>, !vhlo.tensor_v1<!vhlo.ui64_v1>) -> !vhlo.tensor_v1<!vhlo.ui64_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.ui64_v1>, !vhlo.tensor_v1<!vhlo.ui64_v1>) -> !vhlo.tensor_v1<!vhlo.ui64_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<ui64>, tensor<ui64>) -> tensor<ui64>
-   func.return %0 : tensor<ui64>
- }
- 
- // CHECK-LABEL: "type_f8E4M3FN"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_f8E4M3FN(%arg0: tensor<f8E4M3FN>, %arg1: tensor<f8E4M3FN>) -> tensor<f8E4M3FN> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f8E4M3FN_v1>, !vhlo.tensor_v1<!vhlo.f8E4M3FN_v1>) -> !vhlo.tensor_v1<!vhlo.f8E4M3FN_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f8E4M3FN_v1>, !vhlo.tensor_v1<!vhlo.f8E4M3FN_v1>) -> !vhlo.tensor_v1<!vhlo.f8E4M3FN_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f8E4M3FN>, tensor<f8E4M3FN>) -> tensor<f8E4M3FN>
-   func.return %0 : tensor<f8E4M3FN>
- }
- 
- // CHECK-LABEL: "type_f8E5M2"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_f8E5M2(%arg0: tensor<f8E5M2>, %arg1: tensor<f8E5M2>) -> tensor<f8E5M2> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f8E5M2_v1>, !vhlo.tensor_v1<!vhlo.f8E5M2_v1>) -> !vhlo.tensor_v1<!vhlo.f8E5M2_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f8E5M2_v1>, !vhlo.tensor_v1<!vhlo.f8E5M2_v1>) -> !vhlo.tensor_v1<!vhlo.f8E5M2_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f8E5M2>, tensor<f8E5M2>) -> tensor<f8E5M2>
-   func.return %0 : tensor<f8E5M2>
- }
- 
- // CHECK-LABEL: "type_f8E4M3FNUZ"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_f8E4M3FNUZ(%arg0: tensor<f8E4M3FNUZ>, %arg1: tensor<f8E4M3FNUZ>) -> tensor<f8E4M3FNUZ> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f8E4M3FNUZ_v1>, !vhlo.tensor_v1<!vhlo.f8E4M3FNUZ_v1>) -> !vhlo.tensor_v1<!vhlo.f8E4M3FNUZ_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f8E4M3FNUZ_v1>, !vhlo.tensor_v1<!vhlo.f8E4M3FNUZ_v1>) -> !vhlo.tensor_v1<!vhlo.f8E4M3FNUZ_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f8E4M3FNUZ>, tensor<f8E4M3FNUZ>) -> tensor<f8E4M3FNUZ>
-   func.return %0 : tensor<f8E4M3FNUZ>
- }
- 
- // CHECK-LABEL: "type_f8E4M3B11FNUZ"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_f8E4M3B11FNUZ(%arg0: tensor<f8E4M3B11FNUZ>, %arg1: tensor<f8E4M3B11FNUZ>) -> tensor<f8E4M3B11FNUZ> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f8E4M3B11FNUZ_v1>, !vhlo.tensor_v1<!vhlo.f8E4M3B11FNUZ_v1>) -> !vhlo.tensor_v1<!vhlo.f8E4M3B11FNUZ_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f8E4M3B11FNUZ_v1>, !vhlo.tensor_v1<!vhlo.f8E4M3B11FNUZ_v1>) -> !vhlo.tensor_v1<!vhlo.f8E4M3B11FNUZ_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f8E4M3B11FNUZ>, tensor<f8E4M3B11FNUZ>) -> tensor<f8E4M3B11FNUZ>
-   func.return %0 : tensor<f8E4M3B11FNUZ>
- }
- 
- // CHECK-LABEL: "type_f8E5M2FNUZ"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_f8E5M2FNUZ(%arg0: tensor<f8E5M2FNUZ>, %arg1: tensor<f8E5M2FNUZ>) -> tensor<f8E5M2FNUZ> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f8E5M2FNUZ_v1>, !vhlo.tensor_v1<!vhlo.f8E5M2FNUZ_v1>) -> !vhlo.tensor_v1<!vhlo.f8E5M2FNUZ_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f8E5M2FNUZ_v1>, !vhlo.tensor_v1<!vhlo.f8E5M2FNUZ_v1>) -> !vhlo.tensor_v1<!vhlo.f8E5M2FNUZ_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f8E5M2FNUZ>, tensor<f8E5M2FNUZ>) -> tensor<f8E5M2FNUZ>
-   func.return %0 : tensor<f8E5M2FNUZ>
- }
- 
- // CHECK-LABEL: "type_bf16"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_bf16(%arg0: tensor<bf16>, %arg1: tensor<bf16>) -> tensor<bf16> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.bf16_v1>, !vhlo.tensor_v1<!vhlo.bf16_v1>) -> !vhlo.tensor_v1<!vhlo.bf16_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.bf16_v1>, !vhlo.tensor_v1<!vhlo.bf16_v1>) -> !vhlo.tensor_v1<!vhlo.bf16_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<bf16>, tensor<bf16>) -> tensor<bf16>
-   func.return %0 : tensor<bf16>
- }
- 
- // CHECK-LABEL: "type_f16"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_f16(%arg0: tensor<f16>, %arg1: tensor<f16>) -> tensor<f16> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f16_v1>, !vhlo.tensor_v1<!vhlo.f16_v1>) -> !vhlo.tensor_v1<!vhlo.f16_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f16_v1>, !vhlo.tensor_v1<!vhlo.f16_v1>) -> !vhlo.tensor_v1<!vhlo.f16_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f16>, tensor<f16>) -> tensor<f16>
-   func.return %0 : tensor<f16>
- }
- 
- // CHECK-LABEL: "type_f32"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_f32(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "type_f64"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_f64(%arg0: tensor<f64>, %arg1: tensor<f64>) -> tensor<f64> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f64_v1>, !vhlo.tensor_v1<!vhlo.f64_v1>) -> !vhlo.tensor_v1<!vhlo.f64_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f64_v1>, !vhlo.tensor_v1<!vhlo.f64_v1>) -> !vhlo.tensor_v1<!vhlo.f64_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f64>, tensor<f64>) -> tensor<f64>
-   func.return %0 : tensor<f64>
- }
- 
- // CHECK-LABEL: "type_complex_f32"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_complex_f32(%arg0: tensor<complex<f32>>, %arg1: tensor<complex<f32>>) -> tensor<complex<f32>> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>, !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>) -> !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>, !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>) -> !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<complex<f32>>, tensor<complex<f32>>) -> tensor<complex<f32>>
-   func.return %0 : tensor<complex<f32>>
- }
- 
- // CHECK-LABEL: "type_complex_f64"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_complex_f64(%arg0: tensor<complex<f64>>, %arg1: tensor<complex<f64>>) -> tensor<complex<f64>> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f64_v1>>, !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f64_v1>>) -> !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f64_v1>>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f64_v1>>, !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f64_v1>>) -> !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f64_v1>>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<complex<f64>>, tensor<complex<f64>>) -> tensor<complex<f64>>
-   func.return %0 : tensor<complex<f64>>
- }
- 
- // CHECK-LABEL: "type_dynamism_ranked"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @type_dynamism_ranked(%arg0: tensor<?xf32>) -> tensor<?xf32> {
--  // CHECK: "vhlo.abs_v1"(%arg0) : (!vhlo.tensor_v1<?x!vhlo.f32_v1>) -> !vhlo.tensor_v1<?x!vhlo.f32_v1>
-+  // CHECK: "vhlo.abs_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<?x!vhlo.f32_v1>) -> !vhlo.tensor_v1<?x!vhlo.f32_v1>
-   %0 = "stablehlo.abs"(%arg0) : (tensor<?xf32>) -> tensor<?xf32>
-   func.return %0 : tensor<?xf32>
- }
- 
- // CHECK-LABEL: "type_quantization"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_quantization(%arg0: tensor<!quant.uniform<i8:f32, 34.0:16>>, %arg1: tensor<!quant.uniform<i8:f32, 34.0:16>>) -> tensor<!quant.uniform<i8:f32, 34.0:16>> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>, !vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>) -> !vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>, !vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>) -> !vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<!quant.uniform<i8:f32, 34.0:16>>, tensor<!quant.uniform<i8:f32, 34.0:16>>) -> tensor<!quant.uniform<i8:f32, 34.0:16>>
-   func.return %0 : tensor<!quant.uniform<i8:f32, 34.0:16>>
- }
- 
- //       CHECK: function_type = #vhlo.type_v1<!vhlo.func_v1<(!vhlo.token_v1) -> !vhlo.token_v1>>
- // CHECK-LABEL: "type_token_callee"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @type_token_callee(%arg0: !stablehlo.token) -> !stablehlo.token {
--  // CHECK: "vhlo.return_v1"(%arg0) : (!vhlo.token_v1) -> ()
-+  // CHECK: "vhlo.return_v1"(%[[ARG0]]) : (!vhlo.token_v1) -> ()
-   return %arg0 : !stablehlo.token
- }
- 
- //       CHECK: function_type = #vhlo.type_v1<!vhlo.func_v1<(!vhlo.token_v1) -> !vhlo.token_v1>>
- // CHECK-LABEL: "type_token_caller"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @type_token_caller(%arg0: !stablehlo.token) -> !stablehlo.token {
--  // CHECK:      "vhlo.call_v1"(%arg0) <{callee = #vhlo.string_v1<"type_token_callee">}
-+  // CHECK:      "vhlo.call_v1"(%[[ARG0]]) <{callee = #vhlo.string_v1<"type_token_callee">}
-   // CHECK-SAME: (!vhlo.token_v1) -> !vhlo.token_v1
-   %0 = func.call @type_token_callee(%arg0) : (!stablehlo.token) -> !stablehlo.token
-   return %0 : !stablehlo.token
- }
- 
- // CHECK-LABEL: "type_tuple"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @type_tuple(%arg0: tuple<tensor<f32>>) -> tuple<!stablehlo.token> {
-   %0 = "stablehlo.custom_call"(%arg0) {
-     call_target_name = "foo"
-diff --ruN a/stablehlo/stablehlo/tests/vhlo/stablehlo_legalize_to_vhlo.0_13_0.mlir b/stablehlo/stablehlo/tests/vhlo/stablehlo_legalize_to_vhlo.0_13_0.mlir
---- stablehlo/stablehlo/tests/vhlo/stablehlo_legalize_to_vhlo.0_13_0.mlir
-+++ stablehlo/stablehlo/tests/vhlo/stablehlo_legalize_to_vhlo.0_13_0.mlir
-@@ -13,6 +13,7 @@
- // ============ ATTRIBUTES ============
- 
- // CHECK-LABEL: "attr_comparison_direction_eq"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_direction_eq(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     // CHECK: comparison_direction = #vhlo<comparison_direction_v1 EQ>
-@@ -22,6 +23,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_direction_ne"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_direction_ne(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     // CHECK: comparison_direction = #vhlo<comparison_direction_v1 NE>
-@@ -31,6 +33,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_direction_ge"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_direction_ge(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     // CHECK: comparison_direction = #vhlo<comparison_direction_v1 GE>
-@@ -40,6 +43,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_direction_gt"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_direction_gt(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     // CHECK: comparison_direction = #vhlo<comparison_direction_v1 GT>
-@@ -49,6 +53,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_direction_le"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_direction_le(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     // CHECK: comparison_direction = #vhlo<comparison_direction_v1 LE>
-@@ -58,6 +63,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_direction_lt"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_direction_lt(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     // CHECK: comparison_direction = #vhlo<comparison_direction_v1 LT>
-@@ -67,6 +73,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_type_notype"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_type_notype(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     comparison_direction = #stablehlo<comparison_direction EQ>
-@@ -76,6 +83,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_type_float"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_type_float(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     comparison_direction = #stablehlo<comparison_direction EQ>,
-@@ -86,6 +94,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_type_totalorder"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_type_totalorder(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     comparison_direction = #stablehlo<comparison_direction EQ>,
-@@ -96,6 +105,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_type_signed"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_type_signed(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     comparison_direction = #stablehlo<comparison_direction EQ>,
-@@ -106,6 +116,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_type_unsigned"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_type_unsigned(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     comparison_direction = #stablehlo<comparison_direction EQ>,
-@@ -118,6 +129,7 @@
- // ConvDimensionNumbers aka #stablehlo.conv is covered below.
- 
- // CHECK-LABEL: "attr_custom_call_api_version_unspecified"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_custom_call_api_version_unspecified(%arg0: tensor<f32>) -> tensor<f32> {
-   %0 = "stablehlo.custom_call"(%arg0) {
-     call_target_name = "foo",
-@@ -128,6 +140,7 @@
- }
- 
- // CHECK-LABEL: "attr_custom_call_api_version_original"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_custom_call_api_version_original(%arg0: tensor<f32>) -> tensor<f32> {
-   %0 = "stablehlo.custom_call"(%arg0) {
-     call_target_name = "foo",
-@@ -138,6 +151,7 @@
- }
- 
- // CHECK-LABEL: "attr_custom_call_api_version_status_returning"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_custom_call_api_version_status_returning(%arg0: tensor<f32>) -> tensor<f32> {
-   %0 = "stablehlo.custom_call"(%arg0) {
-     call_target_name = "foo",
-@@ -148,6 +162,7 @@
- }
- 
- // CHECK-LABEL: "attr_custom_call_api_version_status_returning_unified"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_custom_call_api_version_status_returning_unified(%arg0: tensor<f32>) -> tensor<f32> {
-   %0 = "stablehlo.custom_call"(%arg0) {
-     call_target_name = "foo",
-@@ -160,6 +175,7 @@
- // DotDimensionNumbers aka #stablehlo.dot is covered below.
- 
- // CHECK-LABEL: "attr_fft_type_fft"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_fft_type_fft(%arg0: tensor<16xcomplex<f32>>) -> tensor<16xcomplex<f32>> {
-   %0 = "stablehlo.fft"(%arg0) {
-     // CHECK: fft_type = #vhlo<fft_type_v1 FFT>
-@@ -170,6 +186,7 @@
- }
- 
- // CHECK-LABEL: "attr_fft_type_ifft"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_fft_type_ifft(%arg0: tensor<16xcomplex<f32>>) -> tensor<16xcomplex<f32>> {
-   %0 = "stablehlo.fft"(%arg0) {
-     // CHECK: fft_type = #vhlo<fft_type_v1 IFFT>
-@@ -180,6 +197,7 @@
- }
- 
- // CHECK-LABEL: "attr_fft_type_rfft"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_fft_type_rfft(%arg0: tensor<16xf32>) -> tensor<9xcomplex<f32>> {
-   %0 = "stablehlo.fft"(%arg0) {
-     // CHECK: fft_type = #vhlo<fft_type_v1 RFFT>
-@@ -190,6 +208,7 @@
- }
- 
- // CHECK-LABEL: "attr_fft_type_irfft"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_fft_type_irfft(%arg0: tensor<9xcomplex<f32>>) -> tensor<16xf32> {
-   %0 = "stablehlo.fft"(%arg0) {
-     // CHECK: fft_type = #vhlo<fft_type_v1 IRFFT>
-@@ -202,6 +221,7 @@
- // GatherDimensionNumbers aka #stablehlo.gather is covered below.
- 
- // CHECK-LABEL: "attr_precision_config_default"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_precision_config_default(%arg0: tensor<8x16xf32>, %arg1: tensor<16x8xf32>) -> tensor<8x8xf32> {
-   %0 = "stablehlo.dot"(%arg0, %arg1) {
-     // CHECK: precision_config = #vhlo.array_v1<[#vhlo<precision_v1 DEFAULT>, #vhlo<precision_v1 DEFAULT>]>
-@@ -210,6 +230,7 @@
- }
- 
- // CHECK-LABEL: "attr_precision_config_high"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_precision_config_high(%arg0: tensor<8x16xf32>, %arg1: tensor<16x8xf32>) -> tensor<8x8xf32> {
-   %0 = "stablehlo.dot"(%arg0, %arg1) {
-     // CHECK: precision_config = #vhlo.array_v1<[#vhlo<precision_v1 HIGH>, #vhlo<precision_v1 HIGH>]>
-@@ -219,6 +240,7 @@
- }
- 
- // CHECK-LABEL: "attr_precision_config_highest"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_precision_config_highest(%arg0: tensor<8x16xf32>, %arg1: tensor<16x8xf32>) -> tensor<8x8xf32> {
-   %0 = "stablehlo.dot"(%arg0, %arg1) {
-     // CHECK: precision_config = #vhlo.array_v1<[#vhlo<precision_v1 HIGHEST>, #vhlo<precision_v1 HIGHEST>]>
-@@ -228,6 +250,7 @@
- }
- 
- // CHECK-LABEL: "attr_rng_algorithm_default"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_rng_algorithm_default(%arg0: tensor<f32>) -> (tensor<f32>, tensor<f32>) {
-   %0:2 = "stablehlo.rng_bit_generator"(%arg0) {
-     // CHECK: rng_algorithm = #vhlo<rng_algorithm_v1 DEFAULT>
-@@ -237,6 +260,7 @@
- }
- 
- // CHECK-LABEL: "attr_rng_algorithm_three_fry"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_rng_algorithm_three_fry(%arg0: tensor<f32>) -> (tensor<f32>, tensor<f32>) {
-   %0:2 = "stablehlo.rng_bit_generator"(%arg0) {
-     // CHECK: rng_algorithm = #vhlo<rng_algorithm_v1 THREE_FRY>
-@@ -246,6 +270,7 @@
- }
- 
- // CHECK-LABEL: "attr_rng_algorithm_philox"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_rng_algorithm_philox(%arg0: tensor<f32>) -> (tensor<f32>, tensor<f32>) {
-   %0:2 = "stablehlo.rng_bit_generator"(%arg0) {
-     // CHECK: rng_algorithm = #vhlo<rng_algorithm_v1 PHILOX>
-@@ -255,6 +280,7 @@
- }
- 
- // CHECK-LABEL: "attr_rng_distribution_uniform"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @attr_rng_distribution_uniform(%arg0: tensor<f32>, %arg1: tensor<f32>, %arg2: tensor<0xindex>) -> tensor<f32> {
-   %0 = "stablehlo.rng"(%arg0, %arg1, %arg2) {
-     // CHECK: rng_distribution = #vhlo<rng_distribution_v1 UNIFORM>
-@@ -264,6 +290,7 @@
- }
- 
- // CHECK-LABEL: "attr_rng_distribution_normal"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @attr_rng_distribution_normal(%arg0: tensor<f32>, %arg1: tensor<f32>, %arg2: tensor<0xindex>) -> tensor<f32> {
-   %0 = "stablehlo.rng"(%arg0, %arg1, %arg2) {
-     // CHECK: rng_distribution = #vhlo<rng_distribution_v1 NORMAL>
-@@ -275,6 +302,7 @@
- // ScatterDimensionNumbers aka #stablehlo.scatter is covered below.
- 
- // CHECK-LABEL: "attr_transpose_no_transpose"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_transpose_no_transpose(%arg0: tensor<16x16xf32>, %arg1: tensor<16x16xf32>) ->  tensor<16x16xf32> {
-   %0 = "stablehlo.triangular_solve"(%arg0, %arg1) {
-     left_side = true,
-@@ -287,6 +315,7 @@
- }
- 
- // CHECK-LABEL: "attr_transpose_transpose"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_transpose_transpose(%arg0: tensor<16x16xf32>, %arg1: tensor<16x16xf32>) ->  tensor<16x16xf32> {
-   %0 = "stablehlo.triangular_solve"(%arg0, %arg1) {
-     left_side = true,
-@@ -299,6 +328,7 @@
- }
- 
- // CHECK-LABEL: "attr_transpose_adjoint"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_transpose_adjoint(%arg0: tensor<16x16xf32>, %arg1: tensor<16x16xf32>) ->  tensor<16x16xf32> {
-   %0 = "stablehlo.triangular_solve"(%arg0, %arg1) {
-     left_side = true,
-@@ -313,10 +343,9 @@
- // TypeExtensionsAttr aka #stablehlo.type_extensions is covered below.
- 
- // CHECK-LABEL: "attr_type_extensions_bounds"
--func.func @attr_type_extensions_bounds(
--    %arg0: tensor<?x?xf32, #stablehlo.type_extensions<bounds = [16, ?]>>)
--    -> tensor<?x?xf32, #stablehlo.type_extensions<bounds = [16, ?]>> {
--  // CHECK: "vhlo.return_v1"(%arg0) : (!vhlo.tensor_v1<?x?x!vhlo.f32_v1, #vhlo.type_extensions_v1<bounds = [16, ?]>>) -> ()
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @attr_type_extensions_bounds(%arg0: tensor<?x?xf32, #stablehlo.type_extensions<bounds = [16, ?]>>) -> tensor<?x?xf32, #stablehlo.type_extensions<bounds = [16, ?]>> {
-+  // CHECK: "vhlo.return_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<?x?x!vhlo.f32_v1, #vhlo.type_extensions_v1<bounds = [16, ?]>>) -> ()
-   func.return %arg0 : tensor<?x?xf32, #stablehlo.type_extensions<bounds = [16, ?]>>
- }
- 
-@@ -324,8 +353,9 @@
- // ============ DEFAULTS ============
- 
- // CHECK-LABEL: "default_all_gather"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_all_gather(%arg0: tensor<16x8xf32>) -> tensor<16x16xf32> {
--  //               CHECK: "vhlo.all_gather_v1"(%arg0) <{
-+  //               CHECK: "vhlo.all_gather_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   all_gather_dim = #vhlo.integer_v1<1 : i64>
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0], [1]]> : tensor<2x1xi64>>,
-@@ -339,8 +369,9 @@
- }
- 
- // CHECK-LABEL: "default_all_reduce"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_all_reduce(%arg0: tensor<f32>) -> tensor<f32> {
--  //               CHECK: "vhlo.all_reduce_v1"(%arg0)
-+  //               CHECK: "vhlo.all_reduce_v1"(%[[ARG0]])
-   //          CHECK-SAME: <{
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0], [1]]> : tensor<2x1xi64>>,
-@@ -362,8 +393,9 @@
- }
- 
- // CHECK-LABEL: "default_all_to_all"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_all_to_all(%arg0: tensor<4x16xf32>) -> tensor<16x4xf32> {
--  //               CHECK: "vhlo.all_to_all_v1"(%arg0) <{
-+  //               CHECK: "vhlo.all_to_all_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   //          CHECK-SAME:   concat_dimension = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0, 1, 2, 3]]> : tensor<1x4xi64>>,
-@@ -380,8 +412,9 @@
- }
- 
- // CHECK-LABEL: "default_cholesky"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_cholesky(%arg0: tensor<1x16x16xf32>) -> tensor<1x16x16xf32> {
--  //      CHECK: "vhlo.cholesky_v1"(%arg0) <{
-+  //      CHECK: "vhlo.cholesky_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   lower = #vhlo.bool_v1<false>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<1x16x16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<1x16x16x!vhlo.f32_v1>
-   %0 = "stablehlo.cholesky"(%arg0) : (tensor<1x16x16xf32>) -> tensor<1x16x16xf32>
-@@ -389,8 +422,9 @@
- }
- 
- // CHECK-LABEL: "default_collective_permute"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_collective_permute(%arg0: tensor<16x8xf32>) -> tensor<16x8xf32> {
--  //               CHECK: "vhlo.collective_permute_v1"(%arg0) <{
-+  //               CHECK: "vhlo.collective_permute_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME{LITERAL}:   source_target_pairs = #vhlo.tensor_v1<dense<[[0, 1], [1, 2], [2, 3]]> : tensor<3x2xi64>>
-   //          CHECK-SAME: }> : (!vhlo.tensor_v1<16x8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x8x!vhlo.f32_v1>
-@@ -401,8 +435,9 @@
- }
- 
- // CHECK-LABEL: "default_compare"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @default_compare(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
--  //      CHECK: "vhlo.compare_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.compare_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   compare_type = #vhlo<comparison_type_v1 NOTYPE>,
-   // CHECK-SAME:   comparison_direction = #vhlo<comparison_direction_v1 EQ>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-@@ -413,8 +448,9 @@
- }
- 
- // CHECK-LABEL: "default_convolution"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @default_convolution(%arg0: tensor<1x8x8x207xf32>, %arg1: tensor<3x3x207x16xf32>) -> tensor<1x6x6x16xf32> {
--  //      CHECK: "vhlo.convolution_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.convolution_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   batch_group_count = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   feature_group_count = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   input_batch_dimension = #vhlo.integer_v1<0 : i64>,
-@@ -442,8 +478,9 @@
- }
- 
- // CHECK-LABEL: "default_custom_call"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_custom_call(%arg0: tensor<f32>) -> tensor<f32> {
--  //      CHECK: "vhlo.custom_call_v1"(%arg0) <{
-+  //      CHECK: "vhlo.custom_call_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   api_version = #vhlo<api_version_v1 API_VERSION_ORIGINAL>,
-   // CHECK-SAME:   backend_config = #vhlo.string_v1<"">,
-   // CHECK-SAME:   call_target_name = #vhlo.string_v1<"foo">,
-@@ -460,8 +497,9 @@
- }
- 
- // CHECK-LABEL: "default_dot_general"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @default_dot_general(%arg0: tensor<8x8x16xf32>, %arg1: tensor<8x16x8xf32>) -> tensor<8x8x8xf32> {
--  //      CHECK: "vhlo.dot_general_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.dot_general_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   lhs_batching_dimensions = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>,
-   // CHECK-SAME:   lhs_contracting_dimensions = #vhlo.tensor_v1<dense<2> : tensor<1xi64>>,
-   // CHECK-SAME:   precision_config = #vhlo.array_v1<[#vhlo<precision_v1 DEFAULT>, #vhlo<precision_v1 DEFAULT>]>,
-@@ -480,8 +518,9 @@
- }
- 
- // CHECK-LABEL: "default_dot"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @default_dot(%arg0: tensor<8x16xf32>, %arg1: tensor<16x8xf32>) -> tensor<8x8xf32> {
--  //      CHECK: "vhlo.dot_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.dot_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   precision_config = #vhlo.array_v1<[#vhlo<precision_v1 DEFAULT>, #vhlo<precision_v1 DEFAULT>]>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<8x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<8x8x!vhlo.f32_v1>
-   %0 = "stablehlo.dot"(%arg0, %arg1) : (tensor<8x16xf32>, tensor<16x8xf32>) -> tensor<8x8xf32>
-@@ -489,8 +528,9 @@
- }
- 
- // CHECK-LABEL: "default_dynamic_broadcast_in_dim"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @default_dynamic_broadcast_in_dim(%arg0: tensor<?x?xf32>, %arg1: tensor<2xindex>) -> tensor<?x?xf32> {
--  //      CHECK: "vhlo.dynamic_broadcast_in_dim_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.dynamic_broadcast_in_dim_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   broadcast_dimensions = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-   // CHECK-SAME:   known_expanding_dimensions = #vhlo.tensor_v1<dense<> : tensor<0xi64>>,
-   // CHECK-SAME:   known_nonexpanding_dimensions = #vhlo.tensor_v1<dense<> : tensor<0xi64>>
-@@ -502,8 +542,9 @@
- }
- 
- // CHECK-LABEL: "default_dynamic_conv"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @default_dynamic_conv(%arg0: tensor<1x8x8x207xf32>, %arg1: tensor<3x3x207x16xf32>, %arg2: tensor<4xi32>) -> tensor<1x?x?x16xf32> {
--  //      CHECK: "vhlo.dynamic_conv_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.dynamic_conv_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   batch_group_count = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   feature_group_count = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   input_batch_dimension = #vhlo.integer_v1<0 : i64>,
-@@ -531,8 +572,9 @@
- }
- 
- // CHECK-LABEL: "default_dynamic_gather"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @default_dynamic_gather(%arg0 : tensor<2x4x9xf32>, %arg1 : tensor<1x5x2xi32>, %arg2 : tensor<3xi32>) -> tensor<1x5x8xf32> {
--  //      CHECK: "vhlo.dynamic_gather_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.dynamic_gather_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   collapsed_slice_dims = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-   // CHECK-SAME:   index_vector_dim = #vhlo.integer_v1<2 : i64>,
-   // CHECK-SAME:   indices_are_sorted = #vhlo.bool_v1<false>,
-@@ -558,15 +600,16 @@
-   // CHECK-SAME:   sym_name = #vhlo.string_v1<"default_func">,
-   // CHECK-SAME:   sym_visibility = #vhlo.string_v1<"">
-   // CHECK-SAME: }> ({
--  // CHECK-NEXT:   ^[[BB:bb.*]](%arg0: !vhlo.tensor_v1<!vhlo.f32_v1>):
--  // CHECK-NEXT:     "vhlo.return_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-+  // CHECK-NEXT:   ^[[BB:bb.*]](%[[ARG0:.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
-+  // CHECK-NEXT:     "vhlo.return_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-   // CHECK-NEXT: }) : () -> ()
-   func.return %arg0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "dynamic_gather"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @dynamic_gather(%arg0 : tensor<2x4x9xf32>, %arg1 : tensor<1x5x2xi32>) -> tensor<1x5x1xf32> {
--  //      CHECK: "vhlo.gather_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.gather_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   collapsed_slice_dims = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-   // CHECK-SAME:   index_vector_dim = #vhlo.integer_v1<2 : i64>,
-   // CHECK-SAME:   indices_are_sorted = #vhlo.bool_v1<false>,
-@@ -587,8 +630,9 @@
- }
- 
- // CHECK-LABEL: "default_infeed"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_infeed(%arg0: !stablehlo.token) -> (tensor<f32>, !stablehlo.token) {
--  //               CHECK: "vhlo.infeed_v1"(%arg0) <{
-+  //               CHECK: "vhlo.infeed_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   infeed_config = #vhlo.string_v1<"">,
-   // CHECK-SAME{LITERAL}:   layout = #vhlo.array_v1<[]>
-   //          CHECK-SAME: }> : (!vhlo.token_v1) -> (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.token_v1)
-@@ -597,8 +641,9 @@
- }
- 
- // CHECK-LABEL: "default_outfeed"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @default_outfeed(%arg0: tensor<f32>, %arg1: !stablehlo.token) -> !stablehlo.token {
--  //      CHECK: "vhlo.outfeed_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.outfeed_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   outfeed_config = #vhlo.string_v1<"">
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.token_v1) -> !vhlo.token_v1
-   %0 = "stablehlo.outfeed"(%arg0, %arg1) : (tensor<f32>, !stablehlo.token) -> !stablehlo.token
-@@ -606,8 +651,9 @@
- }
- 
- // CHECK-LABEL: "default_recv"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_recv(%arg0: !stablehlo.token) -> (tensor<f32>, !stablehlo.token) {
--  //      CHECK: "vhlo.recv_v1"(%arg0) <{
-+  //      CHECK: "vhlo.recv_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME:   channel_type = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   is_host_transfer = #vhlo.bool_v1<false>
-@@ -619,8 +665,9 @@
- }
- 
- // CHECK-LABEL: "default_send"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @default_send(%arg0: tensor<f32>, %arg1: !stablehlo.token) -> !stablehlo.token {
--  //      CHECK: "vhlo.send_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.send_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME:   channel_type = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   is_host_transfer = #vhlo.bool_v1<false>
-@@ -632,8 +679,9 @@
- }
- 
- // CHECK-LABEL: "default_reduce_scatter"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_reduce_scatter(%arg0: tensor<16xf32>) -> tensor<16xf32> {
--  //               CHECK: "vhlo.reduce_scatter_v1"(%arg0) <{
-+  //               CHECK: "vhlo.reduce_scatter_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0], [1]]> : tensor<2x1xi64>>,
-   //          CHECK-SAME:   scatter_dimension = #vhlo.integer_v1<0 : i64>
-@@ -655,8 +703,9 @@
- }
- 
- // CHECK-LABEL: "default_reduce_window"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @default_reduce_window(%arg0: tensor<2x17x31x7xf32>, %arg1: tensor<f32>) -> tensor<2x16x30x7xf32> {
--  //               CHECK: "vhlo.reduce_window_v1"(%arg0, %arg1)  <{
-+  //               CHECK: "vhlo.reduce_window_v1"(%[[ARG0]], %[[ARG1]])  <{
-   //          CHECK-SAME:   base_dilations = #vhlo.tensor_v1<dense<1> : tensor<4xi64>>,
-   // CHECK-SAME{LITERAL}:   padding = #vhlo.tensor_v1<dense<0> : tensor<4x2xi64>>,
-   //          CHECK-SAME:   window_dilations = #vhlo.tensor_v1<dense<1> : tensor<4xi64>>,
-@@ -678,8 +727,9 @@
- }
- 
- // CHECK-LABEL: "default_scatter"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @default_scatter(%arg0: tensor<200x100x300xf32>, %arg1: tensor<10x2xi32>, %arg2: tensor<10x300xf32>) -> tensor<200x100x300xf32> {
--  //      CHECK: "vhlo.scatter_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.scatter_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   index_vector_dim = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   indices_are_sorted = #vhlo.bool_v1<false>,
-   // CHECK-SAME:   inserted_window_dims = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-@@ -707,8 +757,9 @@
- }
- 
- // CHECK-LABEL: "default_select_and_scatter"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @default_select_and_scatter(%arg0: tensor<10x24x24x64xf32>, %arg1: tensor<10x23x23x64xf32>, %arg2: tensor<f32>) -> tensor<10x24x24x64xf32> {
--  //      CHECK: "vhlo.select_and_scatter_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.select_and_scatter_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   padding = #vhlo.tensor_v1<dense<0> : tensor<4x2xi64>>,
-   // CHECK-SAME:   window_dimensions = #vhlo.tensor_v1<dense<[1, 2, 2, 1]> : tensor<4xi64>>,
-   // CHECK-SAME:   window_strides = #vhlo.tensor_v1<dense<1> : tensor<4xi64>>
-@@ -736,8 +787,9 @@
- }
- 
- // CHECK-LABEL: "default_sort"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_sort(%arg0: tensor<16xf32>) -> tensor<16xf32> {
--  //      CHECK: "vhlo.sort_v1"(%arg0) <{
-+  //      CHECK: "vhlo.sort_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   dimension = #vhlo.integer_v1<-1 : i64>
-   // CHECK-SAME:   is_stable = #vhlo.bool_v1<false>
-   // CHECK-SAME: }> ({
-@@ -756,29 +808,33 @@
- // ============ OPS ============
- 
- // CHECK-LABEL: "op_abs"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_abs(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.abs_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.abs_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.abs"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_add"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_add(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_after_all"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_after_all(%arg0: !stablehlo.token) -> !stablehlo.token {
--  // CHECK: "vhlo.after_all_v1"(%arg0) : (!vhlo.token_v1) -> !vhlo.token_v1
-+  // CHECK: "vhlo.after_all_v1"(%[[ARG0]]) : (!vhlo.token_v1) -> !vhlo.token_v1
-   %0 = "stablehlo.after_all"(%arg0) : (!stablehlo.token) -> !stablehlo.token
-   func.return %0 : !stablehlo.token
- }
- 
- // CHECK-LABEL: "op_all_gather"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_all_gather(%arg0: tensor<16x8xf32>) -> tensor<16x16xf32> {
--  //               CHECK: "vhlo.all_gather_v1"(%arg0) <{
-+  //               CHECK: "vhlo.all_gather_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   all_gather_dim = #vhlo.integer_v1<1 : i64>
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0], [1]]> : tensor<2x1xi64>>,
-@@ -794,8 +850,9 @@
- }
- 
- // CHECK-LABEL: "op_all_reduce"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_all_reduce(%arg0: tensor<f32>) -> tensor<f32> {
--  //               CHECK: "vhlo.all_reduce_v1"(%arg0) <{
-+  //               CHECK: "vhlo.all_reduce_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0], [1]]> : tensor<2x1xi64>>,
-   //          CHECK-SAME:   use_global_device_ids = #vhlo.bool_v1<true>
-@@ -817,8 +874,9 @@
- }
- 
- // CHECK-LABEL: "op_all_to_all"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_all_to_all(%arg0: tensor<4x16xf32>) -> tensor<16x4xf32> {
--  //               CHECK: "vhlo.all_to_all_v1"(%arg0) <{
-+  //               CHECK: "vhlo.all_to_all_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<1 : i64>,
-   //          CHECK-SAME:   concat_dimension = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0, 1, 2, 3]]> : tensor<1x4xi64>>,
-@@ -836,22 +894,25 @@
- }
- 
- // CHECK-LABEL: "op_and"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_and(%arg0: tensor<i1>, %arg1: tensor<i1>) -> tensor<i1> {
--  // CHECK: "vhlo.and_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-+  // CHECK: "vhlo.and_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-   %0 = "stablehlo.and"(%arg0, %arg1) : (tensor<i1>, tensor<i1>) -> tensor<i1>
-   func.return %0 : tensor<i1>
- }
- 
- // CHECK-LABEL: "op_atan2"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_atan2(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.atan2_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.atan2_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.atan2"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_batch_norm_grad"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}}, %[[ARG3:.*]]: {{.*}}, %[[ARG4:.*]]: {{.*}})
- func.func @op_batch_norm_grad(%arg0: tensor<16x16x16x16xf32>, %arg1: tensor<16xf32>, %arg2: tensor<16xf32>, %arg3: tensor<16xf32>, %arg4: tensor<16x16x16x16xf32>) -> (tensor<16x16x16x16xf32>, tensor<16xf32>, tensor<16xf32>) {
--  //      CHECK: "vhlo.batch_norm_grad_v1"(%arg0, %arg1, %arg2, %arg3, %arg4) <{
-+  //      CHECK: "vhlo.batch_norm_grad_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]], %[[ARG3]], %[[ARG4]]) <{
-   // CHECK-SAME:   epsilon = #vhlo.float_v1<1.000000e-03 : !vhlo.f32_v1>,
-   // CHECK-SAME:   feature_index = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<16x16x16x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x16x16x16x!vhlo.f32_v1>) -> (!vhlo.tensor_v1<16x16x16x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>)
-@@ -863,8 +924,9 @@
- }
- 
- // CHECK-LABEL: "op_batch_norm_inference"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}}, %[[ARG3:.*]]: {{.*}}, %[[ARG4:.*]]: {{.*}})
- func.func @op_batch_norm_inference(%arg0: tensor<16x16x16x16xf32>, %arg1: tensor<16xf32>, %arg2: tensor<16xf32>, %arg3: tensor<16xf32>, %arg4: tensor<16xf32>) -> tensor<16x16x16x16xf32> {
--  //      CHECK: "vhlo.batch_norm_inference_v1"(%arg0, %arg1, %arg2, %arg3, %arg4) <{
-+  //      CHECK: "vhlo.batch_norm_inference_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]], %[[ARG3]], %[[ARG4]]) <{
-   // CHECK-SAME:   epsilon = #vhlo.float_v1<1.000000e-03 : !vhlo.f32_v1>,
-   // CHECK-SAME:   feature_index = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<16x16x16x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x16x16x16x!vhlo.f32_v1>
-@@ -876,8 +938,9 @@
- }
- 
- // CHECK-LABEL: "op_batch_norm_training"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_batch_norm_training(%arg0: tensor<16x16x16x16xf32>, %arg1: tensor<16xf32>, %arg2: tensor<16xf32>) -> (tensor<16x16x16x16xf32>, tensor<16xf32>, tensor<16xf32>) {
--  //      CHECK: "vhlo.batch_norm_training_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.batch_norm_training_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   epsilon = #vhlo.float_v1<1.000000e-03 : !vhlo.f32_v1>,
-   // CHECK-SAME:   feature_index = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<16x16x16x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>) -> (!vhlo.tensor_v1<16x16x16x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>)
-@@ -889,15 +952,17 @@
- }
- 
- // CHECK-LABEL: "op_bitcast_convert"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_bitcast_convert(%arg0: tensor<i32>) -> tensor<f32> {
--  // CHECK: "vhlo.bitcast_convert_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.bitcast_convert_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.bitcast_convert"(%arg0) : (tensor<i32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_broadcast_in_dim"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_broadcast_in_dim(%arg0: tensor<16xf32>) -> tensor<16x16xf32> {
--  //      CHECK: "vhlo.broadcast_in_dim_v1"(%arg0) <{
-+  //      CHECK: "vhlo.broadcast_in_dim_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   broadcast_dimensions = #vhlo.tensor_v1<dense<1> : tensor<1xi64>>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x16x!vhlo.f32_v1>
-   %0 = "stablehlo.broadcast_in_dim"(%arg0) {
-@@ -907,8 +972,9 @@
- }
- 
- // CHECK-LABEL: "op_broadcast"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_broadcast(%arg0: tensor<16xf32>) -> tensor<16x16xf32> {
--  //      CHECK: "vhlo.broadcast_v1"(%arg0) <{
-+  //      CHECK: "vhlo.broadcast_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   broadcast_sizes = #vhlo.tensor_v1<dense<16> : tensor<1xi64>>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x16x!vhlo.f32_v1>
-   %0 = "stablehlo.broadcast"(%arg0) {
-@@ -918,9 +984,10 @@
- }
- 
- // CHECK-LABEL: "op_case"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_case(%arg0: tensor<i32>, %arg1: tensor<f32>) -> tensor<f32> {
--  //      CHECK: "vhlo.case_v1"(%arg0) ({
--  // CHECK-NEXT:   "vhlo.return_v1"(%arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-+  //      CHECK: "vhlo.case_v1"(%[[ARG0]]) ({
-+  // CHECK-NEXT:   "vhlo.return_v1"(%[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-   // CHECK-NEXT: }) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.case"(%arg0) ({
-     "stablehlo.return"(%arg1) : (tensor<f32>) -> ()
-@@ -929,22 +996,25 @@
- }
- 
- // CHECK-LABEL: "op_cbrt"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_cbrt(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.cbrt_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.cbrt_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.cbrt"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_ceil"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_ceil(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.ceil_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.ceil_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.ceil"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_cholesky"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_cholesky(%arg0: tensor<1x16x16xf32>) -> tensor<1x16x16xf32> {
--  //      CHECK: "vhlo.cholesky_v1"(%arg0) <{
-+  //      CHECK: "vhlo.cholesky_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   lower = #vhlo.bool_v1<true>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<1x16x16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<1x16x16x!vhlo.f32_v1>
-   %0 = "stablehlo.cholesky"(%arg0) {
-@@ -954,22 +1024,25 @@
- }
- 
- // CHECK-LABEL: "op_clamp"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_clamp(%arg0: tensor<f32>, %arg1: tensor<f32>, %arg2: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.clamp_v1"(%arg0, %arg1, %arg2) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.clamp_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.clamp"(%arg0, %arg1, %arg2) : (tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_count_leading_zeros"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_count_leading_zeros(%arg0: tensor<i32>) -> tensor<i32> {
--  // CHECK: "vhlo.count_leading_zeros_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-+  // CHECK: "vhlo.count_leading_zeros_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-   %0 = "stablehlo.count_leading_zeros"(%arg0) : (tensor<i32>) -> tensor<i32>
-   func.return %0 : tensor<i32>
- }
- 
- // CHECK-LABEL: "op_collective_permute"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_collective_permute(%arg0: tensor<16x8xf32>) -> tensor<16x8xf32> {
--  //               CHECK: "vhlo.collective_permute_v1"(%arg0) <{
-+  //               CHECK: "vhlo.collective_permute_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME{LITERAL}:   source_target_pairs = #vhlo.tensor_v1<dense<[[0, 1], [1, 2], [2, 3]]> : tensor<3x2xi64>>
-   //          CHECK-SAME: }> : (!vhlo.tensor_v1<16x8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x8x!vhlo.f32_v1>
-@@ -981,8 +1054,9 @@
- }
- 
- // CHECK-LABEL: "op_compare"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_compare(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
--  //      CHECK: "vhlo.compare_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.compare_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   compare_type = #vhlo<comparison_type_v1 TOTALORDER>,
-   // CHECK-SAME:   comparison_direction = #vhlo<comparison_direction_v1 EQ>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-@@ -994,15 +1068,17 @@
- }
- 
- // CHECK-LABEL: "op_complex"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_complex(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<complex<f32>> {
--  // CHECK: "vhlo.complex_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>
-+  // CHECK: "vhlo.complex_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>
-   %0 = "stablehlo.complex"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<complex<f32>>
-   func.return %0 : tensor<complex<f32>>
- }
- 
- // CHECK-LABEL: "op_concatenate"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_concatenate(%arg0: tensor<8xf32>, %arg1: tensor<8xf32>) -> tensor<16xf32> {
--  //      CHECK: "vhlo.concatenate_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.concatenate_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   dimension = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<8x!vhlo.f32_v1>, !vhlo.tensor_v1<8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x!vhlo.f32_v1>
-   %0 = "stablehlo.concatenate"(%arg0, %arg1) {
-@@ -1012,6 +1088,7 @@
- }
- 
- // CHECK-LABEL: "op_constant"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_constant(%arg0: tensor<f32>) -> tensor<f32> {
-   //      CHECK: "vhlo.constant_v1"() <{
-   // CHECK-SAME:   value = #vhlo.tensor_v1<dense<0.000000e+00> : tensor<f32>>
-@@ -1023,15 +1100,17 @@
- }
- 
- // CHECK-LABEL: "op_convert"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_convert(%arg0: tensor<i32>) -> tensor<f32> {
--  // CHECK: "vhlo.convert_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.convert_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.convert"(%arg0) : (tensor<i32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_convolution"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_convolution(%arg0: tensor<1x8x8x207xf32>, %arg1: tensor<3x3x207x16xf32>) -> tensor<1x7x7x16xf32> {
--  //      CHECK: "vhlo.convolution_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.convolution_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   batch_group_count = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   feature_group_count = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   input_batch_dimension = #vhlo.integer_v1<0 : i64>,
-@@ -1065,8 +1144,9 @@
- }
- 
- // CHECK-LABEL: "op_cosine"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_cosine(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.cosine_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.cosine_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.cosine"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
-@@ -1079,8 +1159,9 @@
- }
- 
- // CHECK-LABEL: "op_cross_replica_sum"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_cross_replica_sum(%arg0: tensor<f32>) -> tensor<f32> {
--  //               CHECK: "vhlo.cross-replica-sum_v1"(%arg0) <{
-+  //               CHECK: "vhlo.cross-replica-sum_v1"(%[[ARG0]]) <{
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0], [1]]> : tensor<2x1xi64>>
-   //          CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.cross-replica-sum"(%arg0) {
-@@ -1090,8 +1171,9 @@
- }
- 
- // CHECK-LABEL: "op_custom_call"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_custom_call(%arg0: tensor<f32>) -> tensor<f32> {
--  //      CHECK: "vhlo.custom_call_v1"(%arg0) <{
-+  //      CHECK: "vhlo.custom_call_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   api_version = #vhlo<api_version_v1 API_VERSION_STATUS_RETURNING>,
-   // CHECK-SAME:   backend_config = #vhlo.string_v1<"\08\03\1A\02">,
-   // CHECK-SAME:   call_target_name = #vhlo.string_v1<"foo">,
-@@ -1122,15 +1204,17 @@
- }
- 
- // CHECK-LABEL: "op_divide"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_divide(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.divide_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.divide_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.divide"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_dot_general"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_dot_general(%arg0: tensor<8x8x16xf32>, %arg1: tensor<8x16x8xf32>) -> tensor<8x8x8xf32> {
--  //      CHECK: "vhlo.dot_general_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.dot_general_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   lhs_batching_dimensions = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>,
-   // CHECK-SAME:   lhs_contracting_dimensions = #vhlo.tensor_v1<dense<2> : tensor<1xi64>>,
-   // CHECK-SAME:   precision_config = #vhlo.array_v1<[#vhlo<precision_v1 HIGHEST>, #vhlo<precision_v1 HIGHEST>]>,
-@@ -1150,8 +1234,9 @@
- }
- 
- // CHECK-LABEL: "op_dot"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_dot(%arg0: tensor<8x16xf32>, %arg1: tensor<16x8xf32>) -> tensor<8x8xf32> {
--  //      CHECK: "vhlo.dot_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.dot_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   precision_config = #vhlo.array_v1<[#vhlo<precision_v1 HIGHEST>, #vhlo<precision_v1 HIGHEST>]>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<8x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<8x8x!vhlo.f32_v1>
-   %0 = "stablehlo.dot"(%arg0, %arg1) {
-@@ -1161,8 +1246,9 @@
- }
- 
- // CHECK-LABEL: "op_dynamic_broadcast_in_dim"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_dynamic_broadcast_in_dim(%arg0: tensor<?x?xf32>, %arg1: tensor<2xindex>) -> tensor<?x?xf32> {
--  //      CHECK: "vhlo.dynamic_broadcast_in_dim_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.dynamic_broadcast_in_dim_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   broadcast_dimensions = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-   // CHECK-SAME:   known_expanding_dimensions = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>,
-   // CHECK-SAME:   known_nonexpanding_dimensions = #vhlo.tensor_v1<dense<1> : tensor<1xi64>>
-@@ -1176,8 +1262,9 @@
- }
- 
- // CHECK-LABEL: "op_dynamic_conv"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_dynamic_conv(%arg0: tensor<1x8x8x207xf32>, %arg1: tensor<3x3x207x16xf32>, %arg2: tensor<4xi32>) -> tensor<1x?x?x16xf32> {
--  //      CHECK: "vhlo.dynamic_conv_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.dynamic_conv_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   batch_group_count = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   feature_group_count = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   input_batch_dimension = #vhlo.integer_v1<0 : i64>,
-@@ -1211,8 +1298,9 @@
- }
- 
- // CHECK-LABEL: "op_dynamic_gather"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_dynamic_gather(%arg0 : tensor<2x4x9xf32>, %arg1 : tensor<1x5x2xi32>, %arg2 : tensor<3xi32>) -> tensor<1x5x8xf32> {
--  //      CHECK: "vhlo.dynamic_gather_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.dynamic_gather_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   collapsed_slice_dims = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-   // CHECK-SAME:   index_vector_dim = #vhlo.integer_v1<2 : i64>,
-   // CHECK-SAME:   indices_are_sorted = #vhlo.bool_v1<true>,
-@@ -1232,8 +1320,9 @@
- }
- 
- // CHECK-LABEL: "op_dynamic_iota"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_dynamic_iota(%arg0: tensor<1xindex>) -> tensor<?xf32> {
--  //      CHECK: "vhlo.dynamic_iota_v1"(%arg0) <{
-+  //      CHECK: "vhlo.dynamic_iota_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   iota_dimension = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<1x!vhlo.index_v1>) -> !vhlo.tensor_v1<?x!vhlo.f32_v1>
-   %0 = "stablehlo.dynamic_iota"(%arg0) {
-@@ -1243,22 +1332,25 @@
- }
- 
- // CHECK-LABEL: "op_dynamic_pad"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}}, %[[ARG3:.*]]: {{.*}}, %[[ARG4:.*]]: {{.*}})
- func.func @op_dynamic_pad(%arg0: tensor<?xf32>, %arg1: tensor<f32>, %arg2: tensor<1xindex>, %arg3: tensor<1xindex>, %arg4: tensor<1xindex>) -> tensor<?xf32> {
--  // CHECK: "vhlo.dynamic_pad_v1"(%arg0, %arg1, %arg2, %arg3, %arg4) : (!vhlo.tensor_v1<?x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>) -> !vhlo.tensor_v1<?x!vhlo.f32_v1>
-+  // CHECK: "vhlo.dynamic_pad_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]], %[[ARG3]], %[[ARG4]]) : (!vhlo.tensor_v1<?x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>) -> !vhlo.tensor_v1<?x!vhlo.f32_v1>
-   %0 = "stablehlo.dynamic_pad"(%arg0, %arg1, %arg2, %arg3, %arg4) : (tensor<?xf32>, tensor<f32>, tensor<1xindex>, tensor<1xindex>, tensor<1xindex>) -> tensor<?xf32>
-   func.return %0 : tensor<?xf32>
- }
- 
- // CHECK-LABEL: "op_dynamic_reshape"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_dynamic_reshape(%arg0: tensor<16xf32>, %arg1: tensor<2xindex>) -> tensor<?x?xf32> {
--  // CHECK: "vhlo.dynamic_reshape_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<2x!vhlo.index_v1>) -> !vhlo.tensor_v1<?x?x!vhlo.f32_v1>
-+  // CHECK: "vhlo.dynamic_reshape_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<2x!vhlo.index_v1>) -> !vhlo.tensor_v1<?x?x!vhlo.f32_v1>
-   %0 = "stablehlo.dynamic_reshape"(%arg0, %arg1) : (tensor<16xf32>, tensor<2xindex>) -> tensor<?x?xf32>
-   func.return %0 : tensor<?x?xf32>
- }
- 
- // CHECK-LABEL: "op_dynamic_slice"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_dynamic_slice(%arg0: tensor<16xf32>, %arg1: tensor<i64>) -> tensor<4xf32> {
--  //      CHECK: "vhlo.dynamic_slice_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.dynamic_slice_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   slice_sizes = #vhlo.tensor_v1<dense<4> : tensor<1xi64>>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.i64_v1>) -> !vhlo.tensor_v1<4x!vhlo.f32_v1>
-   %0 = "stablehlo.dynamic_slice"(%arg0, %arg1) {
-@@ -1268,15 +1360,17 @@
- }
- 
- // CHECK-LABEL: "op_dynamic_update_slice"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_dynamic_update_slice(%arg0: tensor<16xf32>, %arg1: tensor<4xf32>, %arg2: tensor<i64>) -> tensor<16xf32> {
--  // CHECK: "vhlo.dynamic_update_slice_v1"(%arg0, %arg1, %arg2) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<4x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.i64_v1>) -> !vhlo.tensor_v1<16x!vhlo.f32_v1>
-+  // CHECK: "vhlo.dynamic_update_slice_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<4x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.i64_v1>) -> !vhlo.tensor_v1<16x!vhlo.f32_v1>
-   %0 = "stablehlo.dynamic_update_slice"(%arg0, %arg1, %arg2) : (tensor<16xf32>, tensor<4xf32>, tensor<i64>) -> tensor<16xf32>
-   func.return %0 : tensor<16xf32>
- }
- 
- // CHECK-LABEL: "op_einsum"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_einsum(%arg0: tensor<8x16xf32>, %arg1: tensor<16x8xf32>) -> tensor<8x8xf32> {
--  //      CHECK: "vhlo.einsum_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.einsum_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   einsum_config = #vhlo.string_v1<"ab,bc->ac">
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<8x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<8x8x!vhlo.f32_v1>
-   %0 = "stablehlo.einsum"(%arg0, %arg1) {
-@@ -1286,22 +1380,25 @@
- }
- 
- // CHECK-LABEL: "op_exponential_minus_one"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_exponential_minus_one(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.exponential_minus_one_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.exponential_minus_one_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.exponential_minus_one"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_exponential"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_exponential(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.exponential_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.exponential_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.exponential"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_fft"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_fft(%arg0: tensor<16xcomplex<f32>>) -> tensor<16xcomplex<f32>> {
--  //      CHECK: "vhlo.fft_v1"(%arg0) <{
-+  //      CHECK: "vhlo.fft_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   fft_length = #vhlo.tensor_v1<dense<16> : tensor<1xi64>>,
-   // CHECK-SAME:   fft_type = #vhlo<fft_type_v1 FFT>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<16x!vhlo.complex_v1<!vhlo.f32_v1>>) -> !vhlo.tensor_v1<16x!vhlo.complex_v1<!vhlo.f32_v1>>
-@@ -1313,8 +1410,9 @@
- }
- 
- // CHECK-LABEL: "op_floor"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_floor(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.floor_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.floor_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.floor"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
-@@ -1327,16 +1425,17 @@
-   // CHECK-SAME:   sym_name = #vhlo.string_v1<"op_func">,
-   // CHECK-SAME:   sym_visibility = #vhlo.string_v1<"private">
-   // CHECK-SAME: }> ({
--  // CHECK-NEXT:   ^[[BB:bb.*]](%arg0: !vhlo.tensor_v1<!vhlo.f32_v1>):
--  // CHECK-NEXT:     "vhlo.return_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-+  // CHECK-NEXT:   ^[[BB:bb.*]](%[[ARG0:.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
-+  // CHECK-NEXT:     "vhlo.return_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-   // CHECK-NEXT: }) : () -> ()
- 
-   func.return %arg0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_gather"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_gather(%arg0 : tensor<2x4x9xf32>, %arg1 : tensor<1x5x2xi32>) -> tensor<1x5x1xf32> {
--  //      CHECK: "vhlo.gather_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.gather_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   collapsed_slice_dims = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-   // CHECK-SAME:   index_vector_dim = #vhlo.integer_v1<2 : i64>,
-   // CHECK-SAME:   indices_are_sorted = #vhlo.bool_v1<true>,
-@@ -1358,8 +1457,9 @@
- }
- 
- // CHECK-LABEL: "op_get_dimension_size"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_get_dimension_size(%arg0: tensor<?xf32>) -> tensor<i32> {
--  //      CHECK: "vhlo.get_dimension_size_v1"(%arg0) <{
-+  //      CHECK: "vhlo.get_dimension_size_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   dimension = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<?x!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-   %0 = "stablehlo.get_dimension_size"(%arg0) {
-@@ -1369,8 +1469,9 @@
- }
- 
- // CHECK-LABEL: "op_get_tuple_element"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_get_tuple_element(%arg0: tuple<tensor<f32>, tensor<i32>>) -> tensor<f32> {
--  //      CHECK: "vhlo.get_tuple_element_v1"(%arg0) <{
-+  //      CHECK: "vhlo.get_tuple_element_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   index = #vhlo.integer_v1<0 : i32>
-   // CHECK-SAME: }> : (!vhlo.tuple_v1<!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.get_tuple_element"(%arg0) {
-@@ -1380,11 +1481,12 @@
- }
- 
- // CHECK-LABEL: "op_if"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_if(%arg0: tensor<i1>, %arg1: tensor<f32>, %arg2: tensor<f32>) -> tensor<f32> {
--  //      CHECK: "vhlo.if_v1"(%arg0) ({
--  // CHECK-NEXT:   "vhlo.return_v1"(%arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-+  //      CHECK: "vhlo.if_v1"(%[[ARG0]]) ({
-+  // CHECK-NEXT:   "vhlo.return_v1"(%[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-   // CHECK-NEXT: }, {
--  // CHECK-NEXT:   "vhlo.return_v1"(%arg2) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-+  // CHECK-NEXT:   "vhlo.return_v1"(%[[ARG2]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-   // CHECK-NEXT: }) : (!vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.if"(%arg0) ({
-     "stablehlo.return"(%arg1) : (tensor<f32>) -> ()
-@@ -1395,15 +1497,17 @@
- }
- 
- // CHECK-LABEL: "op_imag"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_imag(%arg0: tensor<complex<f32>>) -> tensor<f32> {
--  // CHECK: "vhlo.imag_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.imag_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.imag"(%arg0) : (tensor<complex<f32>>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_infeed"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_infeed(%arg0: !stablehlo.token) -> (tensor<f32>, !stablehlo.token) {
--  //               CHECK: "vhlo.infeed_v1"(%arg0) <{
-+  //               CHECK: "vhlo.infeed_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   infeed_config = #vhlo.string_v1<"foo">,
-   // CHECK-SAME{LITERAL}:   layout = #vhlo.array_v1<[#vhlo.array_v1<[]>]>
-   //          CHECK-SAME: }> : (!vhlo.token_v1) -> (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.token_v1)
-@@ -1426,36 +1530,41 @@
- }
- 
- // CHECK-LABEL: "op_is_finite"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_is_finite(%arg0: tensor<f32>) -> tensor<i1> {
--  // CHECK: "vhlo.is_finite_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-+  // CHECK: "vhlo.is_finite_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-   %0 = "stablehlo.is_finite"(%arg0) : (tensor<f32>) -> tensor<i1>
-   func.return %0 : tensor<i1>
- }
- 
- // CHECK-LABEL: "op_log"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_log(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.log_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.log_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.log"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_log_plus_one"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_log_plus_one(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.log_plus_one_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.log_plus_one_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.log_plus_one"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_logistic"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_logistic(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.logistic_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.logistic_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.logistic"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_map"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_map(%arg0: tensor<16xf32>) -> tensor<16xf32> {
--  //      CHECK: "vhlo.map_v1"(%arg0) <{
-+  //      CHECK: "vhlo.map_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   dimensions = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>
-   // CHECK-SAME: }> ({
-   // CHECK-NEXT:   ^[[BB:bb.*]](%[[ARG1:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
-@@ -1473,57 +1582,65 @@
- }
- 
- // CHECK-LABEL: "op_maximum"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_maximum(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.maximum_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.maximum_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.maximum"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_minimum"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_minimum(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.minimum_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.minimum_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.minimum"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_multiply"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_multiply(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.multiply_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.multiply_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.multiply"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_negate"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_negate(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.negate_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.negate_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.negate"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_not"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_not(%arg0: tensor<i1>) -> tensor<i1> {
--  // CHECK: "vhlo.not_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-+  // CHECK: "vhlo.not_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-   %0 = "stablehlo.not"(%arg0) : (tensor<i1>) -> tensor<i1>
-   func.return %0 : tensor<i1>
- }
- 
- // CHECK-LABEL: "op_optimization_barrier"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_optimization_barrier(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.optimization_barrier_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.optimization_barrier_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.optimization_barrier"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_or"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_or(%arg0: tensor<i1>, %arg1: tensor<i1>) -> tensor<i1> {
--  // CHECK: "vhlo.or_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-+  // CHECK: "vhlo.or_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-   %0 = "stablehlo.or"(%arg0, %arg1) : (tensor<i1>, tensor<i1>) -> tensor<i1>
-   func.return %0 : tensor<i1>
- }
- 
- // CHECK-LABEL: "op_outfeed"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_outfeed(%arg0: tensor<f32>, %arg1: !stablehlo.token) -> !stablehlo.token {
--  //      CHECK: "vhlo.outfeed_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.outfeed_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   outfeed_config = #vhlo.string_v1<"foo">
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.token_v1) -> !vhlo.token_v1
-   %0 = "stablehlo.outfeed"(%arg0, %arg1) {
-@@ -1533,8 +1650,9 @@
- }
- 
- // CHECK-LABEL: "op_pad"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_pad(%arg0: tensor<8xf32>, %arg1: tensor<f32>) -> tensor<16xf32> {
--  //      CHECK: "vhlo.pad_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.pad_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   edge_padding_high = #vhlo.tensor_v1<dense<4> : tensor<1xi64>>,
-   // CHECK-SAME:   edge_padding_low = #vhlo.tensor_v1<dense<4> : tensor<1xi64>>,
-   // CHECK-SAME:   interior_padding = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>
-@@ -1548,36 +1666,41 @@
- }
- 
- // CHECK-LABEL: "op_popcnt"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_popcnt(%arg0: tensor<i32>) -> tensor<i32> {
--  // CHECK: "vhlo.popcnt_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-+  // CHECK: "vhlo.popcnt_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-   %0 = "stablehlo.popcnt"(%arg0) : (tensor<i32>) -> tensor<i32>
-   func.return %0 : tensor<i32>
- }
- 
- // CHECK-LABEL: "op_power"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_power(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.power_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.power_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.power"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_real_dynamic_slice"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}}, %[[ARG3:.*]]: {{.*}})
- func.func @op_real_dynamic_slice(%arg0: tensor<?xf32>, %arg1: tensor<1xindex>, %arg2: tensor<1xindex>, %arg3: tensor<1xindex>) -> tensor<?xf32> {
--  // CHECK: "vhlo.real_dynamic_slice_v1"(%arg0, %arg1, %arg2, %arg3) : (!vhlo.tensor_v1<?x!vhlo.f32_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>) -> !vhlo.tensor_v1<?x!vhlo.f32_v1>
-+  // CHECK: "vhlo.real_dynamic_slice_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]], %[[ARG3]]) : (!vhlo.tensor_v1<?x!vhlo.f32_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>) -> !vhlo.tensor_v1<?x!vhlo.f32_v1>
-   %0 = "stablehlo.real_dynamic_slice"(%arg0, %arg1, %arg2, %arg3) : (tensor<?xf32>, tensor<1xindex>, tensor<1xindex>, tensor<1xindex>) -> tensor<?xf32>
-   func.return %0 : tensor<?xf32>
- }
- 
- // CHECK-LABEL: "op_real"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_real(%arg0: tensor<complex<f32>>) -> tensor<f32> {
--  // CHECK: "vhlo.real_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.real_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.real"(%arg0) : (tensor<complex<f32>>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_recv"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_recv(%arg0: !stablehlo.token) -> (tensor<f32>, !stablehlo.token) {
--  //      CHECK: "vhlo.recv_v1"(%arg0) <{
-+  //      CHECK: "vhlo.recv_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME:   channel_type = #vhlo.integer_v1<3 : i64>,
-   // CHECK-SAME:   is_host_transfer = #vhlo.bool_v1<true>
-@@ -1590,8 +1713,9 @@
- }
- 
- // CHECK-LABEL: "op_reduce"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_reduce(%arg0: tensor<16xf32>, %arg1: tensor<f32>) -> tensor<f32> {
--  //  CHECK: "vhlo.reduce_v1"(%arg0, %arg1)
-+  //  CHECK: "vhlo.reduce_v1"(%[[ARG0]], %[[ARG1]])
-   //  CHECK:   ^[[BB:bb.*]](%[[ARG1:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>, %[[ARG2:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
-   //  CHECK:     "vhlo.return_v1"(%[[VAL1:.*]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-   //  CHECK: }) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-@@ -1606,8 +1730,9 @@
- }
- 
- // CHECK-LABEL: "op_reduce_precision"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_reduce_precision(%arg0: tensor<f32>) -> tensor<f32> {
--  //      CHECK: "vhlo.reduce_precision_v1"(%arg0) <{
-+  //      CHECK: "vhlo.reduce_precision_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   exponent_bits = #vhlo.integer_v1<8 : i32>
-   // CHECK-SAME:   mantissa_bits = #vhlo.integer_v1<10 : i32>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-@@ -1619,8 +1744,9 @@
- }
- 
- // CHECK-LABEL: "op_reduce_scatter"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_reduce_scatter(%arg0: tensor<16xf32>) -> tensor<16xf32> {
--  //               CHECK: "vhlo.reduce_scatter_v1"(%arg0) <{
-+  //               CHECK: "vhlo.reduce_scatter_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0], [1]]> : tensor<2x1xi64>>,
-   //          CHECK-SAME:   scatter_dimension = #vhlo.integer_v1<0 : i64>
-@@ -1644,8 +1770,9 @@
- }
- 
- // CHECK-LABEL: "op_reduce_window"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_reduce_window(%arg0: tensor<2x17x31x7xf32>, %arg1: tensor<f32>) -> tensor<2x9x16x7xf32> {
--  //               CHECK: "vhlo.reduce_window_v1"(%arg0, %arg1) <{
-+  //               CHECK: "vhlo.reduce_window_v1"(%[[ARG0]], %[[ARG1]]) <{
-   //          CHECK-SAME:   base_dilations = #vhlo.tensor_v1<dense<[1, 2, 2, 1]> : tensor<4xi64>>,
-   // CHECK-SAME{LITERAL}:   padding = #vhlo.tensor_v1<dense<[[0, 0], [2, 0], [0, 2], [0, 0]]> : tensor<4x2xi64>>,
-   //          CHECK-SAME:   window_dilations = #vhlo.tensor_v1<dense<[1, 2, 2, 1]> : tensor<4xi64>>,
-@@ -1671,8 +1798,9 @@
- }
- 
- // CHECK-LABEL: "op_remainder"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_remainder(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.remainder_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.remainder_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.remainder"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
-@@ -1692,16 +1820,18 @@
- }
- 
- // CHECK-LABEL: "op_reshape"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_reshape(%arg0: tensor<16xf32>) -> tensor<4x4xf32> {
--  // CHECK: "vhlo.reshape_v1"(%arg0) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<4x4x!vhlo.f32_v1>
-+  // CHECK: "vhlo.reshape_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<4x4x!vhlo.f32_v1>
-   %0 = "stablehlo.reshape"(%arg0) : (tensor<16xf32>) -> tensor<4x4xf32>
-   func.return %0 : tensor<4x4xf32>
- }
- 
- // CHECK-LABEL: "op_return"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_return(%arg0: tensor<i32>, %arg1: tensor<f32>) -> tensor<f32> {
--  //      CHECK: "vhlo.case_v1"(%arg0) ({
--  // CHECK-NEXT:   "vhlo.return_v1"(%arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-+  //      CHECK: "vhlo.case_v1"(%[[ARG0]]) ({
-+  // CHECK-NEXT:   "vhlo.return_v1"(%[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-   // CHECK-NEXT: }) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.case"(%arg0) ({
-     "stablehlo.return"(%arg1) : (tensor<f32>) -> ()
-@@ -1710,8 +1840,9 @@
- }
- 
- // CHECK-LABEL: "op_reverse"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_reverse(%arg0: tensor<16xf32>) -> tensor<16xf32> {
--  //      CHECK: "vhlo.reverse_v1"(%arg0) <{
-+  //      CHECK: "vhlo.reverse_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   dimensions = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x!vhlo.f32_v1>
-   %0 = "stablehlo.reverse"(%arg0) {
-@@ -1721,8 +1852,9 @@
- }
- 
- // CHECK-LABEL: "op_rng_bit_generator"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_rng_bit_generator(%arg0: tensor<f32>) -> (tensor<f32>, tensor<f32>) {
--  //      CHECK: "vhlo.rng_bit_generator_v1"(%arg0) <{
-+  //      CHECK: "vhlo.rng_bit_generator_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   rng_algorithm = #vhlo<rng_algorithm_v1 PHILOX>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>)
-   %0:2 = "stablehlo.rng_bit_generator"(%arg0) {
-@@ -1732,8 +1864,9 @@
- }
- 
- // CHECK-LABEL: "op_rng"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_rng(%arg0: tensor<f32>, %arg1: tensor<f32>, %arg2: tensor<0xindex>) -> tensor<f32> {
--  //      CHECK: "vhlo.rng_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.rng_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   rng_distribution = #vhlo<rng_distribution_v1 NORMAL>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<0x!vhlo.index_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.rng"(%arg0, %arg1, %arg2) {
-@@ -1743,29 +1876,33 @@
- }
- 
- // CHECK-LABEL: "op_round_nearest_afz"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_round_nearest_afz(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.round_nearest_afz_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.round_nearest_afz_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.round_nearest_afz"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_round_nearest_even"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_round_nearest_even(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.round_nearest_even_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.round_nearest_even_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.round_nearest_even"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_rsqrt"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_rsqrt(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.rsqrt_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.rsqrt_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.rsqrt"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_scatter"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_scatter(%arg0: tensor<200x100x300xf32>, %arg1: tensor<10x2xi32>, %arg2: tensor<10x300xf32>) -> tensor<200x100x300xf32> {
--  //      CHECK: "vhlo.scatter_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.scatter_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   index_vector_dim = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   indices_are_sorted = #vhlo.bool_v1<true>,
-   // CHECK-SAME:   inserted_window_dims = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-@@ -1795,8 +1932,9 @@
- }
- 
- // CHECK-LABEL: "op_select_and_scatter"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_select_and_scatter(%arg0: tensor<10x24x24x64xf32>, %arg1: tensor<12x13x13x66xf32>, %arg2: tensor<f32>) -> tensor<10x24x24x64xf32> {
--  //      CHECK: "vhlo.select_and_scatter_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.select_and_scatter_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   padding = #vhlo.tensor_v1<dense<1> : tensor<4x2xi64>>,
-   // CHECK-SAME:   window_dimensions = #vhlo.tensor_v1<dense<[1, 2, 2, 1]> : tensor<4xi64>>,
-   // CHECK-SAME:   window_strides = #vhlo.tensor_v1<dense<[1, 2, 2, 1]> : tensor<4xi64>>
-@@ -1826,15 +1964,17 @@
- }
- 
- // CHECK-LABEL: "op_select"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_select(%arg0: tensor<i1>, %arg1: tensor<f32>, %arg2: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.select_v1"(%arg0, %arg1, %arg2) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.select_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.select"(%arg0, %arg1, %arg2) : (tensor<i1>, tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_send"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_send(%arg0: tensor<f32>, %arg1: !stablehlo.token) -> !stablehlo.token {
--  //      CHECK: "vhlo.send_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.send_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME:   channel_type = #vhlo.integer_v1<2 : i64>,
-   // CHECK-SAME:   is_host_transfer = #vhlo.bool_v1<true>
-@@ -1847,8 +1987,9 @@
- }
- 
- // CHECK-LABEL: "op_set_dimension_size"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_set_dimension_size(%arg0: tensor<?xf32>, %arg1: tensor<i32>) -> tensor<16xf32> {
--  //      CHECK: "vhlo.set_dimension_size_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.set_dimension_size_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   dimension = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<?x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<16x!vhlo.f32_v1>
-   %0 = "stablehlo.set_dimension_size"(%arg0, %arg1) {
-@@ -1858,43 +1999,49 @@
- }
- 
- // CHECK-LABEL: "op_shift_left"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_shift_left(%arg0: tensor<i32>, %arg1: tensor<i32>) -> tensor<i32> {
--  // CHECK: "vhlo.shift_left_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.i32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-+  // CHECK: "vhlo.shift_left_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-   %0 = "stablehlo.shift_left"(%arg0, %arg1) : (tensor<i32>, tensor<i32>) -> tensor<i32>
-   func.return %0 : tensor<i32>
- }
- 
- // CHECK-LABEL: "op_shift_right_arithmetic"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_shift_right_arithmetic(%arg0: tensor<i32>, %arg1: tensor<i32>) -> tensor<i32> {
--  // CHECK: "vhlo.shift_right_arithmetic_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.i32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-+  // CHECK: "vhlo.shift_right_arithmetic_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-   %0 = "stablehlo.shift_right_arithmetic"(%arg0, %arg1) : (tensor<i32>, tensor<i32>) -> tensor<i32>
-   func.return %0 : tensor<i32>
- }
- 
- // CHECK-LABEL: "op_shift_right_logical"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_shift_right_logical(%arg0: tensor<i32>, %arg1: tensor<i32>) -> tensor<i32> {
--  // CHECK: "vhlo.shift_right_logical_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.i32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-+  // CHECK: "vhlo.shift_right_logical_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-   %0 = "stablehlo.shift_right_logical"(%arg0, %arg1) : (tensor<i32>, tensor<i32>) -> tensor<i32>
-   func.return %0 : tensor<i32>
- }
- 
- // CHECK-LABEL: "op_sign"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_sign(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.sign_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.sign_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.sign"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_sine"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_sine(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.sine_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.sine_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.sine"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_slice"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_slice(%arg0: tensor<16xf32>) -> tensor<4xf32> {
--  //      CHECK: "vhlo.slice_v1"(%arg0) <{
-+  //      CHECK: "vhlo.slice_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   limit_indices = #vhlo.tensor_v1<dense<4> : tensor<1xi64>>,
-   // CHECK-SAME:   start_indices = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>,
-   // CHECK-SAME:   strides = #vhlo.tensor_v1<dense<1> : tensor<1xi64>>
-@@ -1908,8 +2055,9 @@
- }
- 
- // CHECK-LABEL: "op_sort"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_sort(%arg0: tensor<16xf32>) -> tensor<16xf32> {
--  //      CHECK: "vhlo.sort_v1"(%arg0) <{
-+  //      CHECK: "vhlo.sort_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   dimension = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME:   is_stable = #vhlo.bool_v1<true>
-   // CHECK-SAME: }> ({
-@@ -1929,29 +2077,33 @@
- }
- 
- // CHECK-LABEL: "op_sqrt"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_sqrt(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.sqrt_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.sqrt_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.sqrt"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_subtract"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_subtract(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.subtract_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.subtract_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.subtract"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_tanh"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_tanh(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.tanh_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.tanh_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.tanh"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_torch_index_select"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_torch_index_select(%arg0: tensor<5x1x5xf32>, %arg1: tensor<2xi32>) ->  tensor<2x1x5xf32> {
--  //      CHECK: "vhlo.torch_index_select_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.torch_index_select_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   batch_dims = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME:   dim = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<5x1x5x!vhlo.f32_v1>, !vhlo.tensor_v1<2x!vhlo.i32_v1>) -> !vhlo.tensor_v1<2x1x5x!vhlo.f32_v1>
-@@ -1963,8 +2115,9 @@
- }
- 
- // CHECK-LABEL: "op_trace"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_trace(%arg0: tensor<f32>) {
--  //      CHECK: "vhlo.trace_v1"(%arg0) <{
-+  //      CHECK: "vhlo.trace_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   tag = #vhlo.string_v1<"foo">
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-   "stablehlo.trace"(%arg0) {
-@@ -1974,8 +2127,9 @@
- }
- 
- // CHECK-LABEL: "op_transpose"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_transpose(%arg0: tensor<16x8xf32>) ->  tensor<8x16xf32> {
--  //      CHECK: "vhlo.transpose_v1"(%arg0) <{
-+  //      CHECK: "vhlo.transpose_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   permutation = #vhlo.tensor_v1<dense<[1, 0]> : tensor<2xi64>>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<16x8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<8x16x!vhlo.f32_v1>
-   %0 = "stablehlo.transpose"(%arg0) {
-@@ -1985,8 +2139,9 @@
- }
- 
- // CHECK-LABEL: "op_triangular_solve"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_triangular_solve(%arg0: tensor<16x16xf32>, %arg1: tensor<16x16xf32>) ->  tensor<16x16xf32> {
--  //      CHECK: "vhlo.triangular_solve_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.triangular_solve_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   left_side = #vhlo.bool_v1<true>,
-   // CHECK-SAME:   lower = #vhlo.bool_v1<true>,
-   // CHECK-SAME:   transpose_a = #vhlo<transpose_v1 NO_TRANSPOSE>,
-@@ -2002,15 +2157,17 @@
- }
- 
- // CHECK-LABEL: "op_tuple"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_tuple(%arg0: tensor<f32>) -> tuple<tensor<f32>> {
--  // CHECK: "vhlo.tuple_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tuple_v1<!vhlo.tensor_v1<!vhlo.f32_v1>>
-+  // CHECK: "vhlo.tuple_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tuple_v1<!vhlo.tensor_v1<!vhlo.f32_v1>>
-   %0 = "stablehlo.tuple"(%arg0) : (tensor<f32>) -> tuple<tensor<f32>>
-   func.return %0 : tuple<tensor<f32>>
- }
- 
- // CHECK-LABEL: "op_unary_einsum"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_unary_einsum(%arg0: tensor<8x16xf32>) -> tensor<8xf32> {
--  //      CHECK: "vhlo.unary_einsum_v1"(%arg0) <{
-+  //      CHECK: "vhlo.unary_einsum_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   einsum_config = #vhlo.string_v1<"ab->a">
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<8x16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<8x!vhlo.f32_v1>
-   %0 = "stablehlo.unary_einsum"(%arg0) {
-@@ -2020,22 +2177,25 @@
- }
- 
- // CHECK-LABEL: "op_uniform_dequantize"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_uniform_dequantize(%arg0: tensor<!quant.uniform<i8:f32, 34.0:16>>) -> tensor<f32> {
--  // CHECK: "vhlo.uniform_dequantize_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.uniform_dequantize_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.uniform_dequantize"(%arg0) : (tensor<!quant.uniform<i8:f32, 34.0:16>>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_uniform_quantize"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_uniform_quantize(%arg0: tensor<f32>) -> tensor<!quant.uniform<i8:f32, 34.0:16>> {
--  // CHECK: "vhlo.uniform_quantize_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>
-+  // CHECK: "vhlo.uniform_quantize_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>
-   %0 = "stablehlo.uniform_quantize"(%arg0) : (tensor<f32>) -> tensor<!quant.uniform<i8:f32, 34.0:16>>
-   func.return %0 : tensor<!quant.uniform<i8:f32, 34.0:16>>
- }
- 
- // CHECK-LABEL: "op_while"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_while(%arg0: tensor<i1>) -> tensor<i1> {
--  //      CHECK: "vhlo.while_v1"(%arg0) ({
-+  //      CHECK: "vhlo.while_v1"(%[[ARG0]]) ({
-   // CHECK-NEXT:   ^[[BB:bb.*]](%[[ARG1:arg.*]]: !vhlo.tensor_v1<!vhlo.bool_v1>):
-   // CHECK-NEXT:     "vhlo.return_v1"(%[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>) -> ()
-   // CHECK-NEXT:   }, {
-@@ -2053,8 +2213,9 @@
- }
- 
- // CHECK-LABEL: "op_xor"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_xor(%arg0: tensor<i1>, %arg1: tensor<i1>) -> tensor<i1> {
--  // CHECK: "vhlo.xor_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-+  // CHECK: "vhlo.xor_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-   %0 = "stablehlo.xor"(%arg0, %arg1) : (tensor<i1>, tensor<i1>) -> tensor<i1>
-   func.return %0 : tensor<i1>
- }
-@@ -2062,190 +2223,217 @@
- // ============ TYPES ============
- 
- // CHECK-LABEL: "type_i1"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_i1(%arg0: tensor<i1>, %arg1: tensor<i1>) -> tensor<i1> {
--  // CHECK: "vhlo.and_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-+  // CHECK: "vhlo.and_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-   %0 = "stablehlo.and"(%arg0, %arg1) : (tensor<i1>, tensor<i1>) -> tensor<i1>
-   func.return %0 : tensor<i1>
- }
- 
- // CHECK-LABEL: "type_i4"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_i4(%arg0: tensor<i4>, %arg1: tensor<i4>) -> tensor<i4> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.i4_v1>, !vhlo.tensor_v1<!vhlo.i4_v1>) -> !vhlo.tensor_v1<!vhlo.i4_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i4_v1>, !vhlo.tensor_v1<!vhlo.i4_v1>) -> !vhlo.tensor_v1<!vhlo.i4_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<i4>, tensor<i4>) -> tensor<i4>
-   func.return %0 : tensor<i4>
- }
- 
- // CHECK-LABEL: "type_i8"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_i8(%arg0: tensor<i8>, %arg1: tensor<i8>) -> tensor<i8> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.i8_v1>, !vhlo.tensor_v1<!vhlo.i8_v1>) -> !vhlo.tensor_v1<!vhlo.i8_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i8_v1>, !vhlo.tensor_v1<!vhlo.i8_v1>) -> !vhlo.tensor_v1<!vhlo.i8_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<i8>, tensor<i8>) -> tensor<i8>
-   func.return %0 : tensor<i8>
- }
- 
- // CHECK-LABEL: "type_i16"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_i16(%arg0: tensor<i16>, %arg1: tensor<i16>) -> tensor<i16> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.i16_v1>, !vhlo.tensor_v1<!vhlo.i16_v1>) -> !vhlo.tensor_v1<!vhlo.i16_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i16_v1>, !vhlo.tensor_v1<!vhlo.i16_v1>) -> !vhlo.tensor_v1<!vhlo.i16_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<i16>, tensor<i16>) -> tensor<i16>
-   func.return %0 : tensor<i16>
- }
- 
- // CHECK-LABEL: "type_i32"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_i32(%arg0: tensor<i32>, %arg1: tensor<i32>) -> tensor<i32> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.i32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<i32>, tensor<i32>) -> tensor<i32>
-   func.return %0 : tensor<i32>
- }
- 
- // CHECK-LABEL: "type_i64"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_i64(%arg0: tensor<i64>, %arg1: tensor<i64>) -> tensor<i64> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.i64_v1>, !vhlo.tensor_v1<!vhlo.i64_v1>) -> !vhlo.tensor_v1<!vhlo.i64_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i64_v1>, !vhlo.tensor_v1<!vhlo.i64_v1>) -> !vhlo.tensor_v1<!vhlo.i64_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<i64>, tensor<i64>) -> tensor<i64>
-   func.return %0 : tensor<i64>
- }
- 
- // CHECK-LABEL: "type_ui4"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_ui4(%arg0: tensor<ui4>, %arg1: tensor<ui4>) -> tensor<ui4> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.ui4_v1>, !vhlo.tensor_v1<!vhlo.ui4_v1>) -> !vhlo.tensor_v1<!vhlo.ui4_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.ui4_v1>, !vhlo.tensor_v1<!vhlo.ui4_v1>) -> !vhlo.tensor_v1<!vhlo.ui4_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<ui4>, tensor<ui4>) -> tensor<ui4>
-   func.return %0 : tensor<ui4>
- }
- 
- // CHECK-LABEL: "type_ui8"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_ui8(%arg0: tensor<ui8>, %arg1: tensor<ui8>) -> tensor<ui8> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.ui8_v1>, !vhlo.tensor_v1<!vhlo.ui8_v1>) -> !vhlo.tensor_v1<!vhlo.ui8_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.ui8_v1>, !vhlo.tensor_v1<!vhlo.ui8_v1>) -> !vhlo.tensor_v1<!vhlo.ui8_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<ui8>, tensor<ui8>) -> tensor<ui8>
-   func.return %0 : tensor<ui8>
- }
- 
- // CHECK-LABEL: "type_ui16"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_ui16(%arg0: tensor<ui16>, %arg1: tensor<ui16>) -> tensor<ui16> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.ui16_v1>, !vhlo.tensor_v1<!vhlo.ui16_v1>) -> !vhlo.tensor_v1<!vhlo.ui16_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.ui16_v1>, !vhlo.tensor_v1<!vhlo.ui16_v1>) -> !vhlo.tensor_v1<!vhlo.ui16_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<ui16>, tensor<ui16>) -> tensor<ui16>
-   func.return %0 : tensor<ui16>
- }
- 
- // CHECK-LABEL: "type_ui32"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_ui32(%arg0: tensor<ui32>, %arg1: tensor<ui32>) -> tensor<ui32> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.ui32_v1>, !vhlo.tensor_v1<!vhlo.ui32_v1>) -> !vhlo.tensor_v1<!vhlo.ui32_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.ui32_v1>, !vhlo.tensor_v1<!vhlo.ui32_v1>) -> !vhlo.tensor_v1<!vhlo.ui32_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<ui32>, tensor<ui32>) -> tensor<ui32>
-   func.return %0 : tensor<ui32>
- }
- 
- // CHECK-LABEL: "type_ui64"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_ui64(%arg0: tensor<ui64>, %arg1: tensor<ui64>) -> tensor<ui64> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.ui64_v1>, !vhlo.tensor_v1<!vhlo.ui64_v1>) -> !vhlo.tensor_v1<!vhlo.ui64_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.ui64_v1>, !vhlo.tensor_v1<!vhlo.ui64_v1>) -> !vhlo.tensor_v1<!vhlo.ui64_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<ui64>, tensor<ui64>) -> tensor<ui64>
-   func.return %0 : tensor<ui64>
- }
- 
- // CHECK-LABEL: "type_f8E4M3FN"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_f8E4M3FN(%arg0: tensor<f8E4M3FN>, %arg1: tensor<f8E4M3FN>) -> tensor<f8E4M3FN> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f8E4M3FN_v1>, !vhlo.tensor_v1<!vhlo.f8E4M3FN_v1>) -> !vhlo.tensor_v1<!vhlo.f8E4M3FN_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f8E4M3FN_v1>, !vhlo.tensor_v1<!vhlo.f8E4M3FN_v1>) -> !vhlo.tensor_v1<!vhlo.f8E4M3FN_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f8E4M3FN>, tensor<f8E4M3FN>) -> tensor<f8E4M3FN>
-   func.return %0 : tensor<f8E4M3FN>
- }
- 
- // CHECK-LABEL: "type_f8E5M2"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_f8E5M2(%arg0: tensor<f8E5M2>, %arg1: tensor<f8E5M2>) -> tensor<f8E5M2> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f8E5M2_v1>, !vhlo.tensor_v1<!vhlo.f8E5M2_v1>) -> !vhlo.tensor_v1<!vhlo.f8E5M2_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f8E5M2_v1>, !vhlo.tensor_v1<!vhlo.f8E5M2_v1>) -> !vhlo.tensor_v1<!vhlo.f8E5M2_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f8E5M2>, tensor<f8E5M2>) -> tensor<f8E5M2>
-   func.return %0 : tensor<f8E5M2>
- }
- 
- // CHECK-LABEL: "type_f8E4M3FNUZ"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_f8E4M3FNUZ(%arg0: tensor<f8E4M3FNUZ>, %arg1: tensor<f8E4M3FNUZ>) -> tensor<f8E4M3FNUZ> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f8E4M3FNUZ_v1>, !vhlo.tensor_v1<!vhlo.f8E4M3FNUZ_v1>) -> !vhlo.tensor_v1<!vhlo.f8E4M3FNUZ_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f8E4M3FNUZ_v1>, !vhlo.tensor_v1<!vhlo.f8E4M3FNUZ_v1>) -> !vhlo.tensor_v1<!vhlo.f8E4M3FNUZ_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f8E4M3FNUZ>, tensor<f8E4M3FNUZ>) -> tensor<f8E4M3FNUZ>
-   func.return %0 : tensor<f8E4M3FNUZ>
- }
- 
- // CHECK-LABEL: "type_f8E4M3B11FNUZ"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_f8E4M3B11FNUZ(%arg0: tensor<f8E4M3B11FNUZ>, %arg1: tensor<f8E4M3B11FNUZ>) -> tensor<f8E4M3B11FNUZ> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f8E4M3B11FNUZ_v1>, !vhlo.tensor_v1<!vhlo.f8E4M3B11FNUZ_v1>) -> !vhlo.tensor_v1<!vhlo.f8E4M3B11FNUZ_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f8E4M3B11FNUZ_v1>, !vhlo.tensor_v1<!vhlo.f8E4M3B11FNUZ_v1>) -> !vhlo.tensor_v1<!vhlo.f8E4M3B11FNUZ_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f8E4M3B11FNUZ>, tensor<f8E4M3B11FNUZ>) -> tensor<f8E4M3B11FNUZ>
-   func.return %0 : tensor<f8E4M3B11FNUZ>
- }
- 
- // CHECK-LABEL: "type_f8E5M2FNUZ"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_f8E5M2FNUZ(%arg0: tensor<f8E5M2FNUZ>, %arg1: tensor<f8E5M2FNUZ>) -> tensor<f8E5M2FNUZ> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f8E5M2FNUZ_v1>, !vhlo.tensor_v1<!vhlo.f8E5M2FNUZ_v1>) -> !vhlo.tensor_v1<!vhlo.f8E5M2FNUZ_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f8E5M2FNUZ_v1>, !vhlo.tensor_v1<!vhlo.f8E5M2FNUZ_v1>) -> !vhlo.tensor_v1<!vhlo.f8E5M2FNUZ_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f8E5M2FNUZ>, tensor<f8E5M2FNUZ>) -> tensor<f8E5M2FNUZ>
-   func.return %0 : tensor<f8E5M2FNUZ>
- }
- 
- // CHECK-LABEL: "type_bf16"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_bf16(%arg0: tensor<bf16>, %arg1: tensor<bf16>) -> tensor<bf16> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.bf16_v1>, !vhlo.tensor_v1<!vhlo.bf16_v1>) -> !vhlo.tensor_v1<!vhlo.bf16_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.bf16_v1>, !vhlo.tensor_v1<!vhlo.bf16_v1>) -> !vhlo.tensor_v1<!vhlo.bf16_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<bf16>, tensor<bf16>) -> tensor<bf16>
-   func.return %0 : tensor<bf16>
- }
- 
- // CHECK-LABEL: "type_f16"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_f16(%arg0: tensor<f16>, %arg1: tensor<f16>) -> tensor<f16> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f16_v1>, !vhlo.tensor_v1<!vhlo.f16_v1>) -> !vhlo.tensor_v1<!vhlo.f16_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f16_v1>, !vhlo.tensor_v1<!vhlo.f16_v1>) -> !vhlo.tensor_v1<!vhlo.f16_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f16>, tensor<f16>) -> tensor<f16>
-   func.return %0 : tensor<f16>
- }
- 
- // CHECK-LABEL: "type_f32"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_f32(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "type_f64"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_f64(%arg0: tensor<f64>, %arg1: tensor<f64>) -> tensor<f64> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f64_v1>, !vhlo.tensor_v1<!vhlo.f64_v1>) -> !vhlo.tensor_v1<!vhlo.f64_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f64_v1>, !vhlo.tensor_v1<!vhlo.f64_v1>) -> !vhlo.tensor_v1<!vhlo.f64_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f64>, tensor<f64>) -> tensor<f64>
-   func.return %0 : tensor<f64>
- }
- 
- // CHECK-LABEL: "type_complex_f32"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_complex_f32(%arg0: tensor<complex<f32>>, %arg1: tensor<complex<f32>>) -> tensor<complex<f32>> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>, !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>) -> !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>, !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>) -> !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<complex<f32>>, tensor<complex<f32>>) -> tensor<complex<f32>>
-   func.return %0 : tensor<complex<f32>>
- }
- 
- // CHECK-LABEL: "type_complex_f64"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_complex_f64(%arg0: tensor<complex<f64>>, %arg1: tensor<complex<f64>>) -> tensor<complex<f64>> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f64_v1>>, !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f64_v1>>) -> !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f64_v1>>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f64_v1>>, !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f64_v1>>) -> !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f64_v1>>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<complex<f64>>, tensor<complex<f64>>) -> tensor<complex<f64>>
-   func.return %0 : tensor<complex<f64>>
- }
- 
- // CHECK-LABEL: "type_dynamism_ranked"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @type_dynamism_ranked(%arg0: tensor<?xf32>) -> tensor<?xf32> {
--  // CHECK: "vhlo.abs_v1"(%arg0) : (!vhlo.tensor_v1<?x!vhlo.f32_v1>) -> !vhlo.tensor_v1<?x!vhlo.f32_v1>
-+  // CHECK: "vhlo.abs_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<?x!vhlo.f32_v1>) -> !vhlo.tensor_v1<?x!vhlo.f32_v1>
-   %0 = "stablehlo.abs"(%arg0) : (tensor<?xf32>) -> tensor<?xf32>
-   func.return %0 : tensor<?xf32>
- }
- 
- // CHECK-LABEL: "type_quantization"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_quantization(%arg0: tensor<!quant.uniform<i8:f32, 34.0:16>>, %arg1: tensor<!quant.uniform<i8:f32, 34.0:16>>) -> tensor<!quant.uniform<i8:f32, 34.0:16>> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>, !vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>) -> !vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>, !vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>) -> !vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<!quant.uniform<i8:f32, 34.0:16>>, tensor<!quant.uniform<i8:f32, 34.0:16>>) -> tensor<!quant.uniform<i8:f32, 34.0:16>>
-   func.return %0 : tensor<!quant.uniform<i8:f32, 34.0:16>>
- }
- 
- //       CHECK: function_type = #vhlo.type_v1<!vhlo.func_v1<(!vhlo.token_v1) -> !vhlo.token_v1>>
- // CHECK-LABEL: "type_token_callee"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @type_token_callee(%arg0: !stablehlo.token) -> !stablehlo.token {
--  // CHECK: "vhlo.return_v1"(%arg0) : (!vhlo.token_v1) -> ()
-+  // CHECK: "vhlo.return_v1"(%[[ARG0]]) : (!vhlo.token_v1) -> ()
-   return %arg0 : !stablehlo.token
- }
- 
- //       CHECK: function_type = #vhlo.type_v1<!vhlo.func_v1<(!vhlo.token_v1) -> !vhlo.token_v1>>
- // CHECK-LABEL: "type_token_caller"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @type_token_caller(%arg0: !stablehlo.token) -> !stablehlo.token {
--  // CHECK:      "vhlo.call_v1"(%arg0) <{callee = #vhlo.string_v1<"type_token_callee">}
-+  // CHECK:      "vhlo.call_v1"(%[[ARG0]]) <{callee = #vhlo.string_v1<"type_token_callee">}
-   // CHECK-SAME: (!vhlo.token_v1) -> !vhlo.token_v1
-   %0 = func.call @type_token_callee(%arg0) : (!stablehlo.token) -> !stablehlo.token
-   return %0 : !stablehlo.token
- }
- 
- // CHECK-LABEL: "type_tuple"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @type_tuple(%arg0: tuple<tensor<f32>>) -> tuple<!stablehlo.token> {
-   %0 = "stablehlo.custom_call"(%arg0) {
-     call_target_name = "foo"
-diff --ruN a/stablehlo/stablehlo/tests/vhlo/stablehlo_legalize_to_vhlo.0_14_0.mlir b/stablehlo/stablehlo/tests/vhlo/stablehlo_legalize_to_vhlo.0_14_0.mlir
---- stablehlo/stablehlo/tests/vhlo/stablehlo_legalize_to_vhlo.0_14_0.mlir
-+++ stablehlo/stablehlo/tests/vhlo/stablehlo_legalize_to_vhlo.0_14_0.mlir
-@@ -13,6 +13,7 @@
- // ============ ATTRIBUTES ============
- 
- // CHECK-LABEL: "attr_comparison_direction_eq"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_direction_eq(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     // CHECK: comparison_direction = #vhlo<comparison_direction_v1 EQ>
-@@ -22,6 +23,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_direction_ne"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_direction_ne(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     // CHECK: comparison_direction = #vhlo<comparison_direction_v1 NE>
-@@ -31,6 +33,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_direction_ge"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_direction_ge(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     // CHECK: comparison_direction = #vhlo<comparison_direction_v1 GE>
-@@ -40,6 +43,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_direction_gt"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_direction_gt(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     // CHECK: comparison_direction = #vhlo<comparison_direction_v1 GT>
-@@ -49,6 +53,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_direction_le"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_direction_le(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     // CHECK: comparison_direction = #vhlo<comparison_direction_v1 LE>
-@@ -58,6 +63,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_direction_lt"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_direction_lt(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     // CHECK: comparison_direction = #vhlo<comparison_direction_v1 LT>
-@@ -67,6 +73,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_type_notype"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_type_notype(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     comparison_direction = #stablehlo<comparison_direction EQ>
-@@ -76,6 +83,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_type_float"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_type_float(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     comparison_direction = #stablehlo<comparison_direction EQ>,
-@@ -86,6 +94,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_type_totalorder"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_type_totalorder(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     comparison_direction = #stablehlo<comparison_direction EQ>,
-@@ -96,6 +105,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_type_signed"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_type_signed(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     comparison_direction = #stablehlo<comparison_direction EQ>,
-@@ -106,6 +116,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_type_unsigned"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_type_unsigned(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     comparison_direction = #stablehlo<comparison_direction EQ>,
-@@ -118,6 +129,7 @@
- // ConvDimensionNumbers aka #stablehlo.conv is covered below.
- 
- // CHECK-LABEL: "attr_custom_call_api_version_unspecified"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_custom_call_api_version_unspecified(%arg0: tensor<f32>) -> tensor<f32> {
-   %0 = "stablehlo.custom_call"(%arg0) {
-     call_target_name = "foo",
-@@ -128,6 +140,7 @@
- }
- 
- // CHECK-LABEL: "attr_custom_call_api_version_original"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_custom_call_api_version_original(%arg0: tensor<f32>) -> tensor<f32> {
-   %0 = "stablehlo.custom_call"(%arg0) {
-     call_target_name = "foo",
-@@ -138,6 +151,7 @@
- }
- 
- // CHECK-LABEL: "attr_custom_call_api_version_status_returning"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_custom_call_api_version_status_returning(%arg0: tensor<f32>) -> tensor<f32> {
-   %0 = "stablehlo.custom_call"(%arg0) {
-     call_target_name = "foo",
-@@ -148,6 +162,7 @@
- }
- 
- // CHECK-LABEL: "attr_custom_call_api_version_status_returning_unified"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_custom_call_api_version_status_returning_unified(%arg0: tensor<f32>) -> tensor<f32> {
-   %0 = "stablehlo.custom_call"(%arg0) {
-     call_target_name = "foo",
-@@ -160,6 +175,7 @@
- // DotDimensionNumbers aka #stablehlo.dot is covered below.
- 
- // CHECK-LABEL: "attr_fft_type_fft"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_fft_type_fft(%arg0: tensor<16xcomplex<f32>>) -> tensor<16xcomplex<f32>> {
-   %0 = "stablehlo.fft"(%arg0) {
-     // CHECK: fft_type = #vhlo<fft_type_v1 FFT>
-@@ -170,6 +186,7 @@
- }
- 
- // CHECK-LABEL: "attr_fft_type_ifft"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_fft_type_ifft(%arg0: tensor<16xcomplex<f32>>) -> tensor<16xcomplex<f32>> {
-   %0 = "stablehlo.fft"(%arg0) {
-     // CHECK: fft_type = #vhlo<fft_type_v1 IFFT>
-@@ -180,6 +197,7 @@
- }
- 
- // CHECK-LABEL: "attr_fft_type_rfft"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_fft_type_rfft(%arg0: tensor<16xf32>) -> tensor<9xcomplex<f32>> {
-   %0 = "stablehlo.fft"(%arg0) {
-     // CHECK: fft_type = #vhlo<fft_type_v1 RFFT>
-@@ -190,6 +208,7 @@
- }
- 
- // CHECK-LABEL: "attr_fft_type_irfft"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_fft_type_irfft(%arg0: tensor<9xcomplex<f32>>) -> tensor<16xf32> {
-   %0 = "stablehlo.fft"(%arg0) {
-     // CHECK: fft_type = #vhlo<fft_type_v1 IRFFT>
-@@ -202,6 +221,7 @@
- // GatherDimensionNumbers aka #stablehlo.gather is covered below.
- 
- // CHECK-LABEL: "attr_precision_config_default"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_precision_config_default(%arg0: tensor<8x16xf32>, %arg1: tensor<16x8xf32>) -> tensor<8x8xf32> {
-   %0 = "stablehlo.dot"(%arg0, %arg1) {
-     // CHECK: precision_config = #vhlo.array_v1<[#vhlo<precision_v1 DEFAULT>, #vhlo<precision_v1 DEFAULT>]>
-@@ -210,6 +230,7 @@
- }
- 
- // CHECK-LABEL: "attr_precision_config_high"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_precision_config_high(%arg0: tensor<8x16xf32>, %arg1: tensor<16x8xf32>) -> tensor<8x8xf32> {
-   %0 = "stablehlo.dot"(%arg0, %arg1) {
-     // CHECK: precision_config = #vhlo.array_v1<[#vhlo<precision_v1 HIGH>, #vhlo<precision_v1 HIGH>]>
-@@ -219,6 +240,7 @@
- }
- 
- // CHECK-LABEL: "attr_precision_config_highest"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_precision_config_highest(%arg0: tensor<8x16xf32>, %arg1: tensor<16x8xf32>) -> tensor<8x8xf32> {
-   %0 = "stablehlo.dot"(%arg0, %arg1) {
-     // CHECK: precision_config = #vhlo.array_v1<[#vhlo<precision_v1 HIGHEST>, #vhlo<precision_v1 HIGHEST>]>
-@@ -228,6 +250,7 @@
- }
- 
- // CHECK-LABEL: "attr_rng_algorithm_default"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_rng_algorithm_default(%arg0: tensor<f32>) -> (tensor<f32>, tensor<f32>) {
-   %0:2 = "stablehlo.rng_bit_generator"(%arg0) {
-     // CHECK: rng_algorithm = #vhlo<rng_algorithm_v1 DEFAULT>
-@@ -237,6 +260,7 @@
- }
- 
- // CHECK-LABEL: "attr_rng_algorithm_three_fry"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_rng_algorithm_three_fry(%arg0: tensor<f32>) -> (tensor<f32>, tensor<f32>) {
-   %0:2 = "stablehlo.rng_bit_generator"(%arg0) {
-     // CHECK: rng_algorithm = #vhlo<rng_algorithm_v1 THREE_FRY>
-@@ -246,6 +270,7 @@
- }
- 
- // CHECK-LABEL: "attr_rng_algorithm_philox"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_rng_algorithm_philox(%arg0: tensor<f32>) -> (tensor<f32>, tensor<f32>) {
-   %0:2 = "stablehlo.rng_bit_generator"(%arg0) {
-     // CHECK: rng_algorithm = #vhlo<rng_algorithm_v1 PHILOX>
-@@ -255,6 +280,7 @@
- }
- 
- // CHECK-LABEL: "attr_rng_distribution_uniform"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @attr_rng_distribution_uniform(%arg0: tensor<f32>, %arg1: tensor<f32>, %arg2: tensor<0xindex>) -> tensor<f32> {
-   %0 = "stablehlo.rng"(%arg0, %arg1, %arg2) {
-     // CHECK: rng_distribution = #vhlo<rng_distribution_v1 UNIFORM>
-@@ -264,6 +290,7 @@
- }
- 
- // CHECK-LABEL: "attr_rng_distribution_normal"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @attr_rng_distribution_normal(%arg0: tensor<f32>, %arg1: tensor<f32>, %arg2: tensor<0xindex>) -> tensor<f32> {
-   %0 = "stablehlo.rng"(%arg0, %arg1, %arg2) {
-     // CHECK: rng_distribution = #vhlo<rng_distribution_v1 NORMAL>
-@@ -275,6 +302,7 @@
- // ScatterDimensionNumbers aka #stablehlo.scatter is covered below.
- 
- // CHECK-LABEL: "attr_transpose_no_transpose"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_transpose_no_transpose(%arg0: tensor<16x16xf32>, %arg1: tensor<16x16xf32>) ->  tensor<16x16xf32> {
-   %0 = "stablehlo.triangular_solve"(%arg0, %arg1) {
-     left_side = true,
-@@ -287,6 +315,7 @@
- }
- 
- // CHECK-LABEL: "attr_transpose_transpose"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_transpose_transpose(%arg0: tensor<16x16xf32>, %arg1: tensor<16x16xf32>) ->  tensor<16x16xf32> {
-   %0 = "stablehlo.triangular_solve"(%arg0, %arg1) {
-     left_side = true,
-@@ -299,6 +328,7 @@
- }
- 
- // CHECK-LABEL: "attr_transpose_adjoint"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_transpose_adjoint(%arg0: tensor<16x16xf32>, %arg1: tensor<16x16xf32>) ->  tensor<16x16xf32> {
-   %0 = "stablehlo.triangular_solve"(%arg0, %arg1) {
-     left_side = true,
-@@ -313,10 +343,9 @@
- // TypeExtensionsAttr aka #stablehlo.type_extensions is covered below.
- 
- // CHECK-LABEL: "attr_type_extensions_bounds"
--func.func @attr_type_extensions_bounds(
--    %arg0: tensor<?x?xf32, #stablehlo.type_extensions<bounds = [16, ?]>>)
--    -> tensor<?x?xf32, #stablehlo.type_extensions<bounds = [16, ?]>> {
--  // CHECK: "vhlo.return_v1"(%arg0) : (!vhlo.tensor_v1<?x?x!vhlo.f32_v1, #vhlo.type_extensions_v1<bounds = [16, ?]>>) -> ()
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @attr_type_extensions_bounds(%arg0: tensor<?x?xf32, #stablehlo.type_extensions<bounds = [16, ?]>>) -> tensor<?x?xf32, #stablehlo.type_extensions<bounds = [16, ?]>> {
-+  // CHECK: "vhlo.return_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<?x?x!vhlo.f32_v1, #vhlo.type_extensions_v1<bounds = [16, ?]>>) -> ()
-   func.return %arg0 : tensor<?x?xf32, #stablehlo.type_extensions<bounds = [16, ?]>>
- }
- 
-@@ -324,8 +353,9 @@
- // ============ DEFAULTS ============
- 
- // CHECK-LABEL: "default_all_gather"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_all_gather(%arg0: tensor<16x8xf32>) -> tensor<16x16xf32> {
--  //               CHECK: "vhlo.all_gather_v1"(%arg0) <{
-+  //               CHECK: "vhlo.all_gather_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   all_gather_dim = #vhlo.integer_v1<1 : i64>
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0], [1]]> : tensor<2x1xi64>>,
-@@ -339,8 +369,9 @@
- }
- 
- // CHECK-LABEL: "default_all_reduce"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_all_reduce(%arg0: tensor<f32>) -> tensor<f32> {
--  //               CHECK: "vhlo.all_reduce_v1"(%arg0)
-+  //               CHECK: "vhlo.all_reduce_v1"(%[[ARG0]])
-   //          CHECK-SAME: <{
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0], [1]]> : tensor<2x1xi64>>,
-@@ -362,8 +393,9 @@
- }
- 
- // CHECK-LABEL: "default_all_to_all"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_all_to_all(%arg0: tensor<4x16xf32>) -> tensor<16x4xf32> {
--  //               CHECK: "vhlo.all_to_all_v1"(%arg0) <{
-+  //               CHECK: "vhlo.all_to_all_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   //          CHECK-SAME:   concat_dimension = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0, 1, 2, 3]]> : tensor<1x4xi64>>,
-@@ -380,8 +412,9 @@
- }
- 
- // CHECK-LABEL: "default_cholesky"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_cholesky(%arg0: tensor<1x16x16xf32>) -> tensor<1x16x16xf32> {
--  //      CHECK: "vhlo.cholesky_v1"(%arg0) <{
-+  //      CHECK: "vhlo.cholesky_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   lower = #vhlo.bool_v1<false>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<1x16x16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<1x16x16x!vhlo.f32_v1>
-   %0 = "stablehlo.cholesky"(%arg0) : (tensor<1x16x16xf32>) -> tensor<1x16x16xf32>
-@@ -389,8 +422,9 @@
- }
- 
- // CHECK-LABEL: "default_collective_permute"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_collective_permute(%arg0: tensor<16x8xf32>) -> tensor<16x8xf32> {
--  //               CHECK: "vhlo.collective_permute_v1"(%arg0) <{
-+  //               CHECK: "vhlo.collective_permute_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME{LITERAL}:   source_target_pairs = #vhlo.tensor_v1<dense<[[0, 1], [1, 2], [2, 3]]> : tensor<3x2xi64>>
-   //          CHECK-SAME: }> : (!vhlo.tensor_v1<16x8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x8x!vhlo.f32_v1>
-@@ -401,8 +435,9 @@
- }
- 
- // CHECK-LABEL: "default_compare"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @default_compare(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
--  //      CHECK: "vhlo.compare_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.compare_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   compare_type = #vhlo<comparison_type_v1 NOTYPE>,
-   // CHECK-SAME:   comparison_direction = #vhlo<comparison_direction_v1 EQ>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-@@ -413,8 +448,9 @@
- }
- 
- // CHECK-LABEL: "default_convolution"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @default_convolution(%arg0: tensor<1x8x8x207xf32>, %arg1: tensor<3x3x207x16xf32>) -> tensor<1x6x6x16xf32> {
--  //      CHECK: "vhlo.convolution_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.convolution_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   batch_group_count = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   feature_group_count = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   input_batch_dimension = #vhlo.integer_v1<0 : i64>,
-@@ -442,8 +478,9 @@
- }
- 
- // CHECK-LABEL: "default_custom_call"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_custom_call(%arg0: tensor<f32>) -> tensor<f32> {
--  //      CHECK: "vhlo.custom_call_v1"(%arg0) <{
-+  //      CHECK: "vhlo.custom_call_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   api_version = #vhlo<api_version_v1 API_VERSION_ORIGINAL>,
-   // CHECK-SAME:   backend_config = #vhlo.string_v1<"">,
-   // CHECK-SAME:   call_target_name = #vhlo.string_v1<"foo">,
-@@ -460,8 +497,9 @@
- }
- 
- // CHECK-LABEL: "default_dot_general"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @default_dot_general(%arg0: tensor<8x8x16xf32>, %arg1: tensor<8x16x8xf32>) -> tensor<8x8x8xf32> {
--  //      CHECK: "vhlo.dot_general_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.dot_general_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   lhs_batching_dimensions = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>,
-   // CHECK-SAME:   lhs_contracting_dimensions = #vhlo.tensor_v1<dense<2> : tensor<1xi64>>,
-   // CHECK-SAME:   precision_config = #vhlo.array_v1<[#vhlo<precision_v1 DEFAULT>, #vhlo<precision_v1 DEFAULT>]>,
-@@ -480,8 +518,9 @@
- }
- 
- // CHECK-LABEL: "default_dot"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @default_dot(%arg0: tensor<8x16xf32>, %arg1: tensor<16x8xf32>) -> tensor<8x8xf32> {
--  //      CHECK: "vhlo.dot_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.dot_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   precision_config = #vhlo.array_v1<[#vhlo<precision_v1 DEFAULT>, #vhlo<precision_v1 DEFAULT>]>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<8x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<8x8x!vhlo.f32_v1>
-   %0 = "stablehlo.dot"(%arg0, %arg1) : (tensor<8x16xf32>, tensor<16x8xf32>) -> tensor<8x8xf32>
-@@ -489,8 +528,9 @@
- }
- 
- // CHECK-LABEL: "default_dynamic_broadcast_in_dim"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @default_dynamic_broadcast_in_dim(%arg0: tensor<?x?xf32>, %arg1: tensor<2xindex>) -> tensor<?x?xf32> {
--  //      CHECK: "vhlo.dynamic_broadcast_in_dim_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.dynamic_broadcast_in_dim_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   broadcast_dimensions = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-   // CHECK-SAME:   known_expanding_dimensions = #vhlo.tensor_v1<dense<> : tensor<0xi64>>,
-   // CHECK-SAME:   known_nonexpanding_dimensions = #vhlo.tensor_v1<dense<> : tensor<0xi64>>
-@@ -502,8 +542,9 @@
- }
- 
- // CHECK-LABEL: "default_dynamic_conv"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @default_dynamic_conv(%arg0: tensor<1x8x8x207xf32>, %arg1: tensor<3x3x207x16xf32>, %arg2: tensor<4xi32>) -> tensor<1x?x?x16xf32> {
--  //      CHECK: "vhlo.dynamic_conv_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.dynamic_conv_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   batch_group_count = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   feature_group_count = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   input_batch_dimension = #vhlo.integer_v1<0 : i64>,
-@@ -531,8 +572,9 @@
- }
- 
- // CHECK-LABEL: "default_dynamic_gather"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @default_dynamic_gather(%arg0 : tensor<2x4x9xf32>, %arg1 : tensor<1x5x2xi32>, %arg2 : tensor<3xi32>) -> tensor<1x5x8xf32> {
--  //      CHECK: "vhlo.dynamic_gather_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.dynamic_gather_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   collapsed_slice_dims = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-   // CHECK-SAME:   index_vector_dim = #vhlo.integer_v1<2 : i64>,
-   // CHECK-SAME:   indices_are_sorted = #vhlo.bool_v1<false>,
-@@ -558,15 +600,16 @@
-   // CHECK-SAME:   sym_name = #vhlo.string_v1<"default_func">,
-   // CHECK-SAME:   sym_visibility = #vhlo.string_v1<"">
-   // CHECK-SAME: }> ({
--  // CHECK-NEXT:   ^[[BB:bb.*]](%arg0: !vhlo.tensor_v1<!vhlo.f32_v1>):
--  // CHECK-NEXT:     "vhlo.return_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-+  // CHECK-NEXT:   ^[[BB:bb.*]](%[[ARG0:.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
-+  // CHECK-NEXT:     "vhlo.return_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-   // CHECK-NEXT: }) : () -> ()
-   func.return %arg0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "dynamic_gather"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @dynamic_gather(%arg0 : tensor<2x4x9xf32>, %arg1 : tensor<1x5x2xi32>) -> tensor<1x5x1xf32> {
--  //      CHECK: "vhlo.gather_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.gather_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   collapsed_slice_dims = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-   // CHECK-SAME:   index_vector_dim = #vhlo.integer_v1<2 : i64>,
-   // CHECK-SAME:   indices_are_sorted = #vhlo.bool_v1<false>,
-@@ -587,8 +630,9 @@
- }
- 
- // CHECK-LABEL: "default_infeed"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_infeed(%arg0: !stablehlo.token) -> (tensor<f32>, !stablehlo.token) {
--  //               CHECK: "vhlo.infeed_v1"(%arg0) <{
-+  //               CHECK: "vhlo.infeed_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   infeed_config = #vhlo.string_v1<"">,
-   // CHECK-SAME{LITERAL}:   layout = #vhlo.array_v1<[]>
-   //          CHECK-SAME: }> : (!vhlo.token_v1) -> (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.token_v1)
-@@ -597,8 +641,9 @@
- }
- 
- // CHECK-LABEL: "default_outfeed"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @default_outfeed(%arg0: tensor<f32>, %arg1: !stablehlo.token) -> !stablehlo.token {
--  //      CHECK: "vhlo.outfeed_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.outfeed_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   outfeed_config = #vhlo.string_v1<"">
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.token_v1) -> !vhlo.token_v1
-   %0 = "stablehlo.outfeed"(%arg0, %arg1) : (tensor<f32>, !stablehlo.token) -> !stablehlo.token
-@@ -606,8 +651,9 @@
- }
- 
- // CHECK-LABEL: "default_recv"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_recv(%arg0: !stablehlo.token) -> (tensor<f32>, !stablehlo.token) {
--  //      CHECK: "vhlo.recv_v1"(%arg0) <{
-+  //      CHECK: "vhlo.recv_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME:   channel_type = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   is_host_transfer = #vhlo.bool_v1<false>
-@@ -619,8 +665,9 @@
- }
- 
- // CHECK-LABEL: "default_send"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @default_send(%arg0: tensor<f32>, %arg1: !stablehlo.token) -> !stablehlo.token {
--  //      CHECK: "vhlo.send_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.send_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME:   channel_type = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   is_host_transfer = #vhlo.bool_v1<false>
-@@ -632,8 +679,9 @@
- }
- 
- // CHECK-LABEL: "default_reduce_scatter"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_reduce_scatter(%arg0: tensor<16xf32>) -> tensor<16xf32> {
--  //               CHECK: "vhlo.reduce_scatter_v1"(%arg0) <{
-+  //               CHECK: "vhlo.reduce_scatter_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0], [1]]> : tensor<2x1xi64>>,
-   //          CHECK-SAME:   scatter_dimension = #vhlo.integer_v1<0 : i64>
-@@ -655,8 +703,9 @@
- }
- 
- // CHECK-LABEL: "default_reduce_window"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @default_reduce_window(%arg0: tensor<2x17x31x7xf32>, %arg1: tensor<f32>) -> tensor<2x16x30x7xf32> {
--  //               CHECK: "vhlo.reduce_window_v1"(%arg0, %arg1)  <{
-+  //               CHECK: "vhlo.reduce_window_v1"(%[[ARG0]], %[[ARG1]])  <{
-   //          CHECK-SAME:   base_dilations = #vhlo.tensor_v1<dense<1> : tensor<4xi64>>,
-   // CHECK-SAME{LITERAL}:   padding = #vhlo.tensor_v1<dense<0> : tensor<4x2xi64>>,
-   //          CHECK-SAME:   window_dilations = #vhlo.tensor_v1<dense<1> : tensor<4xi64>>,
-@@ -678,8 +727,9 @@
- }
- 
- // CHECK-LABEL: "default_scatter"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @default_scatter(%arg0: tensor<200x100x300xf32>, %arg1: tensor<10x2xi32>, %arg2: tensor<10x300xf32>) -> tensor<200x100x300xf32> {
--  //      CHECK: "vhlo.scatter_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.scatter_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   index_vector_dim = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   indices_are_sorted = #vhlo.bool_v1<false>,
-   // CHECK-SAME:   inserted_window_dims = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-@@ -707,8 +757,9 @@
- }
- 
- // CHECK-LABEL: "default_select_and_scatter"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @default_select_and_scatter(%arg0: tensor<10x24x24x64xf32>, %arg1: tensor<10x23x23x64xf32>, %arg2: tensor<f32>) -> tensor<10x24x24x64xf32> {
--  //      CHECK: "vhlo.select_and_scatter_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.select_and_scatter_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   padding = #vhlo.tensor_v1<dense<0> : tensor<4x2xi64>>,
-   // CHECK-SAME:   window_dimensions = #vhlo.tensor_v1<dense<[1, 2, 2, 1]> : tensor<4xi64>>,
-   // CHECK-SAME:   window_strides = #vhlo.tensor_v1<dense<1> : tensor<4xi64>>
-@@ -736,8 +787,9 @@
- }
- 
- // CHECK-LABEL: "default_sort"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_sort(%arg0: tensor<16xf32>) -> tensor<16xf32> {
--  //      CHECK: "vhlo.sort_v1"(%arg0) <{
-+  //      CHECK: "vhlo.sort_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   dimension = #vhlo.integer_v1<-1 : i64>
-   // CHECK-SAME:   is_stable = #vhlo.bool_v1<false>
-   // CHECK-SAME: }> ({
-@@ -756,29 +808,33 @@
- // ============ OPS ============
- 
- // CHECK-LABEL: "op_abs"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_abs(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.abs_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.abs_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.abs"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_add"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_add(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_after_all"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_after_all(%arg0: !stablehlo.token) -> !stablehlo.token {
--  // CHECK: "vhlo.after_all_v1"(%arg0) : (!vhlo.token_v1) -> !vhlo.token_v1
-+  // CHECK: "vhlo.after_all_v1"(%[[ARG0]]) : (!vhlo.token_v1) -> !vhlo.token_v1
-   %0 = "stablehlo.after_all"(%arg0) : (!stablehlo.token) -> !stablehlo.token
-   func.return %0 : !stablehlo.token
- }
- 
- // CHECK-LABEL: "op_all_gather"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_all_gather(%arg0: tensor<16x8xf32>) -> tensor<16x16xf32> {
--  //               CHECK: "vhlo.all_gather_v1"(%arg0) <{
-+  //               CHECK: "vhlo.all_gather_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   all_gather_dim = #vhlo.integer_v1<1 : i64>
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0], [1]]> : tensor<2x1xi64>>,
-@@ -794,8 +850,9 @@
- }
- 
- // CHECK-LABEL: "op_all_reduce"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_all_reduce(%arg0: tensor<f32>) -> tensor<f32> {
--  //               CHECK: "vhlo.all_reduce_v1"(%arg0) <{
-+  //               CHECK: "vhlo.all_reduce_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0], [1]]> : tensor<2x1xi64>>,
-   //          CHECK-SAME:   use_global_device_ids = #vhlo.bool_v1<true>
-@@ -817,8 +874,9 @@
- }
- 
- // CHECK-LABEL: "op_all_to_all"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_all_to_all(%arg0: tensor<4x16xf32>) -> tensor<16x4xf32> {
--  //               CHECK: "vhlo.all_to_all_v1"(%arg0) <{
-+  //               CHECK: "vhlo.all_to_all_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<1 : i64>,
-   //          CHECK-SAME:   concat_dimension = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0, 1, 2, 3]]> : tensor<1x4xi64>>,
-@@ -836,22 +894,25 @@
- }
- 
- // CHECK-LABEL: "op_and"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_and(%arg0: tensor<i1>, %arg1: tensor<i1>) -> tensor<i1> {
--  // CHECK: "vhlo.and_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-+  // CHECK: "vhlo.and_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-   %0 = "stablehlo.and"(%arg0, %arg1) : (tensor<i1>, tensor<i1>) -> tensor<i1>
-   func.return %0 : tensor<i1>
- }
- 
- // CHECK-LABEL: "op_atan2"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_atan2(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.atan2_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.atan2_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.atan2"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_batch_norm_grad"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}}, %[[ARG3:.*]]: {{.*}}, %[[ARG4:.*]]: {{.*}})
- func.func @op_batch_norm_grad(%arg0: tensor<16x16x16x16xf32>, %arg1: tensor<16xf32>, %arg2: tensor<16xf32>, %arg3: tensor<16xf32>, %arg4: tensor<16x16x16x16xf32>) -> (tensor<16x16x16x16xf32>, tensor<16xf32>, tensor<16xf32>) {
--  //      CHECK: "vhlo.batch_norm_grad_v1"(%arg0, %arg1, %arg2, %arg3, %arg4) <{
-+  //      CHECK: "vhlo.batch_norm_grad_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]], %[[ARG3]], %[[ARG4]]) <{
-   // CHECK-SAME:   epsilon = #vhlo.float_v1<1.000000e-03 : !vhlo.f32_v1>,
-   // CHECK-SAME:   feature_index = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<16x16x16x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x16x16x16x!vhlo.f32_v1>) -> (!vhlo.tensor_v1<16x16x16x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>)
-@@ -863,8 +924,9 @@
- }
- 
- // CHECK-LABEL: "op_batch_norm_inference"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}}, %[[ARG3:.*]]: {{.*}}, %[[ARG4:.*]]: {{.*}})
- func.func @op_batch_norm_inference(%arg0: tensor<16x16x16x16xf32>, %arg1: tensor<16xf32>, %arg2: tensor<16xf32>, %arg3: tensor<16xf32>, %arg4: tensor<16xf32>) -> tensor<16x16x16x16xf32> {
--  //      CHECK: "vhlo.batch_norm_inference_v1"(%arg0, %arg1, %arg2, %arg3, %arg4) <{
-+  //      CHECK: "vhlo.batch_norm_inference_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]], %[[ARG3]], %[[ARG4]]) <{
-   // CHECK-SAME:   epsilon = #vhlo.float_v1<1.000000e-03 : !vhlo.f32_v1>,
-   // CHECK-SAME:   feature_index = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<16x16x16x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x16x16x16x!vhlo.f32_v1>
-@@ -876,8 +938,9 @@
- }
- 
- // CHECK-LABEL: "op_batch_norm_training"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_batch_norm_training(%arg0: tensor<16x16x16x16xf32>, %arg1: tensor<16xf32>, %arg2: tensor<16xf32>) -> (tensor<16x16x16x16xf32>, tensor<16xf32>, tensor<16xf32>) {
--  //      CHECK: "vhlo.batch_norm_training_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.batch_norm_training_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   epsilon = #vhlo.float_v1<1.000000e-03 : !vhlo.f32_v1>,
-   // CHECK-SAME:   feature_index = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<16x16x16x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>) -> (!vhlo.tensor_v1<16x16x16x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>)
-@@ -889,15 +952,17 @@
- }
- 
- // CHECK-LABEL: "op_bitcast_convert"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_bitcast_convert(%arg0: tensor<i32>) -> tensor<f32> {
--  // CHECK: "vhlo.bitcast_convert_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.bitcast_convert_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.bitcast_convert"(%arg0) : (tensor<i32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_broadcast_in_dim"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_broadcast_in_dim(%arg0: tensor<16xf32>) -> tensor<16x16xf32> {
--  //      CHECK: "vhlo.broadcast_in_dim_v1"(%arg0) <{
-+  //      CHECK: "vhlo.broadcast_in_dim_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   broadcast_dimensions = #vhlo.tensor_v1<dense<1> : tensor<1xi64>>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x16x!vhlo.f32_v1>
-   %0 = "stablehlo.broadcast_in_dim"(%arg0) {
-@@ -907,8 +972,9 @@
- }
- 
- // CHECK-LABEL: "op_broadcast"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_broadcast(%arg0: tensor<16xf32>) -> tensor<16x16xf32> {
--  //      CHECK: "vhlo.broadcast_v1"(%arg0) <{
-+  //      CHECK: "vhlo.broadcast_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   broadcast_sizes = #vhlo.tensor_v1<dense<16> : tensor<1xi64>>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x16x!vhlo.f32_v1>
-   %0 = "stablehlo.broadcast"(%arg0) {
-@@ -918,9 +984,10 @@
- }
- 
- // CHECK-LABEL: "op_case"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_case(%arg0: tensor<i32>, %arg1: tensor<f32>) -> tensor<f32> {
--  //      CHECK: "vhlo.case_v1"(%arg0) ({
--  // CHECK-NEXT:   "vhlo.return_v1"(%arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-+  //      CHECK: "vhlo.case_v1"(%[[ARG0]]) ({
-+  // CHECK-NEXT:   "vhlo.return_v1"(%[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-   // CHECK-NEXT: }) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.case"(%arg0) ({
-     "stablehlo.return"(%arg1) : (tensor<f32>) -> ()
-@@ -929,22 +996,25 @@
- }
- 
- // CHECK-LABEL: "op_cbrt"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_cbrt(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.cbrt_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.cbrt_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.cbrt"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_ceil"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_ceil(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.ceil_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.ceil_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.ceil"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_cholesky"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_cholesky(%arg0: tensor<1x16x16xf32>) -> tensor<1x16x16xf32> {
--  //      CHECK: "vhlo.cholesky_v1"(%arg0) <{
-+  //      CHECK: "vhlo.cholesky_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   lower = #vhlo.bool_v1<true>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<1x16x16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<1x16x16x!vhlo.f32_v1>
-   %0 = "stablehlo.cholesky"(%arg0) {
-@@ -954,22 +1024,25 @@
- }
- 
- // CHECK-LABEL: "op_clamp"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_clamp(%arg0: tensor<f32>, %arg1: tensor<f32>, %arg2: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.clamp_v1"(%arg0, %arg1, %arg2) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.clamp_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.clamp"(%arg0, %arg1, %arg2) : (tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_count_leading_zeros"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_count_leading_zeros(%arg0: tensor<i32>) -> tensor<i32> {
--  // CHECK: "vhlo.count_leading_zeros_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-+  // CHECK: "vhlo.count_leading_zeros_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-   %0 = "stablehlo.count_leading_zeros"(%arg0) : (tensor<i32>) -> tensor<i32>
-   func.return %0 : tensor<i32>
- }
- 
- // CHECK-LABEL: "op_collective_permute"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_collective_permute(%arg0: tensor<16x8xf32>) -> tensor<16x8xf32> {
--  //               CHECK: "vhlo.collective_permute_v1"(%arg0) <{
-+  //               CHECK: "vhlo.collective_permute_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME{LITERAL}:   source_target_pairs = #vhlo.tensor_v1<dense<[[0, 1], [1, 2], [2, 3]]> : tensor<3x2xi64>>
-   //          CHECK-SAME: }> : (!vhlo.tensor_v1<16x8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x8x!vhlo.f32_v1>
-@@ -981,8 +1054,9 @@
- }
- 
- // CHECK-LABEL: "op_compare"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_compare(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
--  //      CHECK: "vhlo.compare_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.compare_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   compare_type = #vhlo<comparison_type_v1 TOTALORDER>,
-   // CHECK-SAME:   comparison_direction = #vhlo<comparison_direction_v1 EQ>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-@@ -994,15 +1068,17 @@
- }
- 
- // CHECK-LABEL: "op_complex"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_complex(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<complex<f32>> {
--  // CHECK: "vhlo.complex_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>
-+  // CHECK: "vhlo.complex_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>
-   %0 = "stablehlo.complex"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<complex<f32>>
-   func.return %0 : tensor<complex<f32>>
- }
- 
- // CHECK-LABEL: "op_concatenate"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_concatenate(%arg0: tensor<8xf32>, %arg1: tensor<8xf32>) -> tensor<16xf32> {
--  //      CHECK: "vhlo.concatenate_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.concatenate_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   dimension = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<8x!vhlo.f32_v1>, !vhlo.tensor_v1<8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x!vhlo.f32_v1>
-   %0 = "stablehlo.concatenate"(%arg0, %arg1) {
-@@ -1012,6 +1088,7 @@
- }
- 
- // CHECK-LABEL: "op_constant"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_constant(%arg0: tensor<f32>) -> tensor<f32> {
-   //      CHECK: "vhlo.constant_v1"() <{
-   // CHECK-SAME:   value = #vhlo.tensor_v1<dense<0.000000e+00> : tensor<f32>>
-@@ -1023,15 +1100,17 @@
- }
- 
- // CHECK-LABEL: "op_convert"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_convert(%arg0: tensor<i32>) -> tensor<f32> {
--  // CHECK: "vhlo.convert_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.convert_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.convert"(%arg0) : (tensor<i32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_convolution"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_convolution(%arg0: tensor<1x8x8x207xf32>, %arg1: tensor<3x3x207x16xf32>) -> tensor<1x7x7x16xf32> {
--  //      CHECK: "vhlo.convolution_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.convolution_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   batch_group_count = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   feature_group_count = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   input_batch_dimension = #vhlo.integer_v1<0 : i64>,
-@@ -1065,8 +1144,9 @@
- }
- 
- // CHECK-LABEL: "op_cosine"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_cosine(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.cosine_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.cosine_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.cosine"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
-@@ -1079,8 +1159,9 @@
- }
- 
- // CHECK-LABEL: "op_cross_replica_sum"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_cross_replica_sum(%arg0: tensor<f32>) -> tensor<f32> {
--  //               CHECK: "vhlo.cross-replica-sum_v1"(%arg0) <{
-+  //               CHECK: "vhlo.cross-replica-sum_v1"(%[[ARG0]]) <{
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0], [1]]> : tensor<2x1xi64>>
-   //          CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.cross-replica-sum"(%arg0) {
-@@ -1090,8 +1171,9 @@
- }
- 
- // CHECK-LABEL: "op_custom_call"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_custom_call(%arg0: tensor<f32>) -> tensor<f32> {
--  //      CHECK: "vhlo.custom_call_v1"(%arg0) <{
-+  //      CHECK: "vhlo.custom_call_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   api_version = #vhlo<api_version_v1 API_VERSION_STATUS_RETURNING>,
-   // CHECK-SAME:   backend_config = #vhlo.string_v1<"\08\03\1A\02">,
-   // CHECK-SAME:   call_target_name = #vhlo.string_v1<"foo">,
-@@ -1122,15 +1204,17 @@
- }
- 
- // CHECK-LABEL: "op_divide"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_divide(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.divide_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.divide_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.divide"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_dot_general"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_dot_general(%arg0: tensor<8x8x16xf32>, %arg1: tensor<8x16x8xf32>) -> tensor<8x8x8xf32> {
--  //      CHECK: "vhlo.dot_general_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.dot_general_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   lhs_batching_dimensions = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>,
-   // CHECK-SAME:   lhs_contracting_dimensions = #vhlo.tensor_v1<dense<2> : tensor<1xi64>>,
-   // CHECK-SAME:   precision_config = #vhlo.array_v1<[#vhlo<precision_v1 HIGHEST>, #vhlo<precision_v1 HIGHEST>]>,
-@@ -1150,8 +1234,9 @@
- }
- 
- // CHECK-LABEL: "op_dot"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_dot(%arg0: tensor<8x16xf32>, %arg1: tensor<16x8xf32>) -> tensor<8x8xf32> {
--  //      CHECK: "vhlo.dot_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.dot_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   precision_config = #vhlo.array_v1<[#vhlo<precision_v1 HIGHEST>, #vhlo<precision_v1 HIGHEST>]>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<8x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<8x8x!vhlo.f32_v1>
-   %0 = "stablehlo.dot"(%arg0, %arg1) {
-@@ -1161,8 +1246,9 @@
- }
- 
- // CHECK-LABEL: "op_dynamic_broadcast_in_dim"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_dynamic_broadcast_in_dim(%arg0: tensor<?x?xf32>, %arg1: tensor<2xindex>) -> tensor<?x?xf32> {
--  //      CHECK: "vhlo.dynamic_broadcast_in_dim_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.dynamic_broadcast_in_dim_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   broadcast_dimensions = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-   // CHECK-SAME:   known_expanding_dimensions = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>,
-   // CHECK-SAME:   known_nonexpanding_dimensions = #vhlo.tensor_v1<dense<1> : tensor<1xi64>>
-@@ -1176,8 +1262,9 @@
- }
- 
- // CHECK-LABEL: "op_dynamic_conv"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_dynamic_conv(%arg0: tensor<1x8x8x207xf32>, %arg1: tensor<3x3x207x16xf32>, %arg2: tensor<4xi32>) -> tensor<1x?x?x16xf32> {
--  //      CHECK: "vhlo.dynamic_conv_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.dynamic_conv_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   batch_group_count = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   feature_group_count = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   input_batch_dimension = #vhlo.integer_v1<0 : i64>,
-@@ -1211,8 +1298,9 @@
- }
- 
- // CHECK-LABEL: "op_dynamic_gather"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_dynamic_gather(%arg0 : tensor<2x4x9xf32>, %arg1 : tensor<1x5x2xi32>, %arg2 : tensor<3xi32>) -> tensor<1x5x8xf32> {
--  //      CHECK: "vhlo.dynamic_gather_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.dynamic_gather_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   collapsed_slice_dims = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-   // CHECK-SAME:   index_vector_dim = #vhlo.integer_v1<2 : i64>,
-   // CHECK-SAME:   indices_are_sorted = #vhlo.bool_v1<true>,
-@@ -1232,8 +1320,9 @@
- }
- 
- // CHECK-LABEL: "op_dynamic_iota"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_dynamic_iota(%arg0: tensor<1xindex>) -> tensor<?xf32> {
--  //      CHECK: "vhlo.dynamic_iota_v1"(%arg0) <{
-+  //      CHECK: "vhlo.dynamic_iota_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   iota_dimension = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<1x!vhlo.index_v1>) -> !vhlo.tensor_v1<?x!vhlo.f32_v1>
-   %0 = "stablehlo.dynamic_iota"(%arg0) {
-@@ -1243,22 +1332,25 @@
- }
- 
- // CHECK-LABEL: "op_dynamic_pad"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}}, %[[ARG3:.*]]: {{.*}}, %[[ARG4:.*]]: {{.*}})
- func.func @op_dynamic_pad(%arg0: tensor<?xf32>, %arg1: tensor<f32>, %arg2: tensor<1xindex>, %arg3: tensor<1xindex>, %arg4: tensor<1xindex>) -> tensor<?xf32> {
--  // CHECK: "vhlo.dynamic_pad_v1"(%arg0, %arg1, %arg2, %arg3, %arg4) : (!vhlo.tensor_v1<?x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>) -> !vhlo.tensor_v1<?x!vhlo.f32_v1>
-+  // CHECK: "vhlo.dynamic_pad_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]], %[[ARG3]], %[[ARG4]]) : (!vhlo.tensor_v1<?x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>) -> !vhlo.tensor_v1<?x!vhlo.f32_v1>
-   %0 = "stablehlo.dynamic_pad"(%arg0, %arg1, %arg2, %arg3, %arg4) : (tensor<?xf32>, tensor<f32>, tensor<1xindex>, tensor<1xindex>, tensor<1xindex>) -> tensor<?xf32>
-   func.return %0 : tensor<?xf32>
- }
- 
- // CHECK-LABEL: "op_dynamic_reshape"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_dynamic_reshape(%arg0: tensor<16xf32>, %arg1: tensor<2xindex>) -> tensor<?x?xf32> {
--  // CHECK: "vhlo.dynamic_reshape_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<2x!vhlo.index_v1>) -> !vhlo.tensor_v1<?x?x!vhlo.f32_v1>
-+  // CHECK: "vhlo.dynamic_reshape_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<2x!vhlo.index_v1>) -> !vhlo.tensor_v1<?x?x!vhlo.f32_v1>
-   %0 = "stablehlo.dynamic_reshape"(%arg0, %arg1) : (tensor<16xf32>, tensor<2xindex>) -> tensor<?x?xf32>
-   func.return %0 : tensor<?x?xf32>
- }
- 
- // CHECK-LABEL: "op_dynamic_slice"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_dynamic_slice(%arg0: tensor<16xf32>, %arg1: tensor<i64>) -> tensor<4xf32> {
--  //      CHECK: "vhlo.dynamic_slice_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.dynamic_slice_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   slice_sizes = #vhlo.tensor_v1<dense<4> : tensor<1xi64>>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.i64_v1>) -> !vhlo.tensor_v1<4x!vhlo.f32_v1>
-   %0 = "stablehlo.dynamic_slice"(%arg0, %arg1) {
-@@ -1268,15 +1360,17 @@
- }
- 
- // CHECK-LABEL: "op_dynamic_update_slice"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_dynamic_update_slice(%arg0: tensor<16xf32>, %arg1: tensor<4xf32>, %arg2: tensor<i64>) -> tensor<16xf32> {
--  // CHECK: "vhlo.dynamic_update_slice_v1"(%arg0, %arg1, %arg2) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<4x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.i64_v1>) -> !vhlo.tensor_v1<16x!vhlo.f32_v1>
-+  // CHECK: "vhlo.dynamic_update_slice_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<4x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.i64_v1>) -> !vhlo.tensor_v1<16x!vhlo.f32_v1>
-   %0 = "stablehlo.dynamic_update_slice"(%arg0, %arg1, %arg2) : (tensor<16xf32>, tensor<4xf32>, tensor<i64>) -> tensor<16xf32>
-   func.return %0 : tensor<16xf32>
- }
- 
- // CHECK-LABEL: "op_einsum"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_einsum(%arg0: tensor<8x16xf32>, %arg1: tensor<16x8xf32>) -> tensor<8x8xf32> {
--  //      CHECK: "vhlo.einsum_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.einsum_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   einsum_config = #vhlo.string_v1<"ab,bc->ac">
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<8x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<8x8x!vhlo.f32_v1>
-   %0 = "stablehlo.einsum"(%arg0, %arg1) {
-@@ -1286,22 +1380,25 @@
- }
- 
- // CHECK-LABEL: "op_exponential_minus_one"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_exponential_minus_one(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.exponential_minus_one_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.exponential_minus_one_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.exponential_minus_one"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_exponential"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_exponential(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.exponential_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.exponential_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.exponential"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_fft"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_fft(%arg0: tensor<16xcomplex<f32>>) -> tensor<16xcomplex<f32>> {
--  //      CHECK: "vhlo.fft_v1"(%arg0) <{
-+  //      CHECK: "vhlo.fft_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   fft_length = #vhlo.tensor_v1<dense<16> : tensor<1xi64>>,
-   // CHECK-SAME:   fft_type = #vhlo<fft_type_v1 FFT>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<16x!vhlo.complex_v1<!vhlo.f32_v1>>) -> !vhlo.tensor_v1<16x!vhlo.complex_v1<!vhlo.f32_v1>>
-@@ -1313,8 +1410,9 @@
- }
- 
- // CHECK-LABEL: "op_floor"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_floor(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.floor_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.floor_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.floor"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
-@@ -1327,16 +1425,17 @@
-   // CHECK-SAME:   sym_name = #vhlo.string_v1<"op_func">,
-   // CHECK-SAME:   sym_visibility = #vhlo.string_v1<"private">
-   // CHECK-SAME: }> ({
--  // CHECK-NEXT:   ^[[BB:bb.*]](%arg0: !vhlo.tensor_v1<!vhlo.f32_v1>):
--  // CHECK-NEXT:     "vhlo.return_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-+  // CHECK-NEXT:   ^[[BB:bb.*]](%[[ARG0:.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
-+  // CHECK-NEXT:     "vhlo.return_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-   // CHECK-NEXT: }) : () -> ()
- 
-   func.return %arg0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_gather"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_gather(%arg0 : tensor<2x4x9xf32>, %arg1 : tensor<1x5x2xi32>) -> tensor<1x5x1xf32> {
--  //      CHECK: "vhlo.gather_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.gather_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   collapsed_slice_dims = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-   // CHECK-SAME:   index_vector_dim = #vhlo.integer_v1<2 : i64>,
-   // CHECK-SAME:   indices_are_sorted = #vhlo.bool_v1<true>,
-@@ -1358,8 +1457,9 @@
- }
- 
- // CHECK-LABEL: "op_get_dimension_size"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_get_dimension_size(%arg0: tensor<?xf32>) -> tensor<i32> {
--  //      CHECK: "vhlo.get_dimension_size_v1"(%arg0) <{
-+  //      CHECK: "vhlo.get_dimension_size_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   dimension = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<?x!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-   %0 = "stablehlo.get_dimension_size"(%arg0) {
-@@ -1369,8 +1469,9 @@
- }
- 
- // CHECK-LABEL: "op_get_tuple_element"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_get_tuple_element(%arg0: tuple<tensor<f32>, tensor<i32>>) -> tensor<f32> {
--  //      CHECK: "vhlo.get_tuple_element_v1"(%arg0) <{
-+  //      CHECK: "vhlo.get_tuple_element_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   index = #vhlo.integer_v1<0 : i32>
-   // CHECK-SAME: }> : (!vhlo.tuple_v1<!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.get_tuple_element"(%arg0) {
-@@ -1380,11 +1481,12 @@
- }
- 
- // CHECK-LABEL: "op_if"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_if(%arg0: tensor<i1>, %arg1: tensor<f32>, %arg2: tensor<f32>) -> tensor<f32> {
--  //      CHECK: "vhlo.if_v1"(%arg0) ({
--  // CHECK-NEXT:   "vhlo.return_v1"(%arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-+  //      CHECK: "vhlo.if_v1"(%[[ARG0]]) ({
-+  // CHECK-NEXT:   "vhlo.return_v1"(%[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-   // CHECK-NEXT: }, {
--  // CHECK-NEXT:   "vhlo.return_v1"(%arg2) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-+  // CHECK-NEXT:   "vhlo.return_v1"(%[[ARG2]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-   // CHECK-NEXT: }) : (!vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.if"(%arg0) ({
-     "stablehlo.return"(%arg1) : (tensor<f32>) -> ()
-@@ -1395,15 +1497,17 @@
- }
- 
- // CHECK-LABEL: "op_imag"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_imag(%arg0: tensor<complex<f32>>) -> tensor<f32> {
--  // CHECK: "vhlo.imag_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.imag_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.imag"(%arg0) : (tensor<complex<f32>>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_infeed"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_infeed(%arg0: !stablehlo.token) -> (tensor<f32>, !stablehlo.token) {
--  //               CHECK: "vhlo.infeed_v1"(%arg0) <{
-+  //               CHECK: "vhlo.infeed_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   infeed_config = #vhlo.string_v1<"foo">,
-   // CHECK-SAME{LITERAL}:   layout = #vhlo.array_v1<[#vhlo.array_v1<[]>]>
-   //          CHECK-SAME: }> : (!vhlo.token_v1) -> (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.token_v1)
-@@ -1426,36 +1530,41 @@
- }
- 
- // CHECK-LABEL: "op_is_finite"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_is_finite(%arg0: tensor<f32>) -> tensor<i1> {
--  // CHECK: "vhlo.is_finite_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-+  // CHECK: "vhlo.is_finite_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-   %0 = "stablehlo.is_finite"(%arg0) : (tensor<f32>) -> tensor<i1>
-   func.return %0 : tensor<i1>
- }
- 
- // CHECK-LABEL: "op_log"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_log(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.log_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.log_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.log"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_log_plus_one"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_log_plus_one(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.log_plus_one_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.log_plus_one_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.log_plus_one"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_logistic"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_logistic(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.logistic_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.logistic_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.logistic"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_map"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_map(%arg0: tensor<16xf32>) -> tensor<16xf32> {
--  //      CHECK: "vhlo.map_v1"(%arg0) <{
-+  //      CHECK: "vhlo.map_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   dimensions = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>
-   // CHECK-SAME: }> ({
-   // CHECK-NEXT:   ^[[BB:bb.*]](%[[ARG1:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
-@@ -1473,57 +1582,65 @@
- }
- 
- // CHECK-LABEL: "op_maximum"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_maximum(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.maximum_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.maximum_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.maximum"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_minimum"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_minimum(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.minimum_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.minimum_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.minimum"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_multiply"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_multiply(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.multiply_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.multiply_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.multiply"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_negate"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_negate(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.negate_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.negate_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.negate"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_not"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_not(%arg0: tensor<i1>) -> tensor<i1> {
--  // CHECK: "vhlo.not_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-+  // CHECK: "vhlo.not_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-   %0 = "stablehlo.not"(%arg0) : (tensor<i1>) -> tensor<i1>
-   func.return %0 : tensor<i1>
- }
- 
- // CHECK-LABEL: "op_optimization_barrier"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_optimization_barrier(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.optimization_barrier_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.optimization_barrier_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.optimization_barrier"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_or"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_or(%arg0: tensor<i1>, %arg1: tensor<i1>) -> tensor<i1> {
--  // CHECK: "vhlo.or_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-+  // CHECK: "vhlo.or_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-   %0 = "stablehlo.or"(%arg0, %arg1) : (tensor<i1>, tensor<i1>) -> tensor<i1>
-   func.return %0 : tensor<i1>
- }
- 
- // CHECK-LABEL: "op_outfeed"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_outfeed(%arg0: tensor<f32>, %arg1: !stablehlo.token) -> !stablehlo.token {
--  //      CHECK: "vhlo.outfeed_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.outfeed_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   outfeed_config = #vhlo.string_v1<"foo">
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.token_v1) -> !vhlo.token_v1
-   %0 = "stablehlo.outfeed"(%arg0, %arg1) {
-@@ -1533,8 +1650,9 @@
- }
- 
- // CHECK-LABEL: "op_pad"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_pad(%arg0: tensor<8xf32>, %arg1: tensor<f32>) -> tensor<16xf32> {
--  //      CHECK: "vhlo.pad_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.pad_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   edge_padding_high = #vhlo.tensor_v1<dense<4> : tensor<1xi64>>,
-   // CHECK-SAME:   edge_padding_low = #vhlo.tensor_v1<dense<4> : tensor<1xi64>>,
-   // CHECK-SAME:   interior_padding = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>
-@@ -1548,36 +1666,41 @@
- }
- 
- // CHECK-LABEL: "op_popcnt"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_popcnt(%arg0: tensor<i32>) -> tensor<i32> {
--  // CHECK: "vhlo.popcnt_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-+  // CHECK: "vhlo.popcnt_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-   %0 = "stablehlo.popcnt"(%arg0) : (tensor<i32>) -> tensor<i32>
-   func.return %0 : tensor<i32>
- }
- 
- // CHECK-LABEL: "op_power"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_power(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.power_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.power_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.power"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_real_dynamic_slice"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}}, %[[ARG3:.*]]: {{.*}})
- func.func @op_real_dynamic_slice(%arg0: tensor<?xf32>, %arg1: tensor<1xindex>, %arg2: tensor<1xindex>, %arg3: tensor<1xindex>) -> tensor<?xf32> {
--  // CHECK: "vhlo.real_dynamic_slice_v1"(%arg0, %arg1, %arg2, %arg3) : (!vhlo.tensor_v1<?x!vhlo.f32_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>) -> !vhlo.tensor_v1<?x!vhlo.f32_v1>
-+  // CHECK: "vhlo.real_dynamic_slice_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]], %[[ARG3]]) : (!vhlo.tensor_v1<?x!vhlo.f32_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>) -> !vhlo.tensor_v1<?x!vhlo.f32_v1>
-   %0 = "stablehlo.real_dynamic_slice"(%arg0, %arg1, %arg2, %arg3) : (tensor<?xf32>, tensor<1xindex>, tensor<1xindex>, tensor<1xindex>) -> tensor<?xf32>
-   func.return %0 : tensor<?xf32>
- }
- 
- // CHECK-LABEL: "op_real"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_real(%arg0: tensor<complex<f32>>) -> tensor<f32> {
--  // CHECK: "vhlo.real_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.real_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.real"(%arg0) : (tensor<complex<f32>>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_recv"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_recv(%arg0: !stablehlo.token) -> (tensor<f32>, !stablehlo.token) {
--  //      CHECK: "vhlo.recv_v1"(%arg0) <{
-+  //      CHECK: "vhlo.recv_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME:   channel_type = #vhlo.integer_v1<3 : i64>,
-   // CHECK-SAME:   is_host_transfer = #vhlo.bool_v1<true>
-@@ -1590,8 +1713,9 @@
- }
- 
- // CHECK-LABEL: "op_reduce"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_reduce(%arg0: tensor<16xf32>, %arg1: tensor<f32>) -> tensor<f32> {
--  //  CHECK: "vhlo.reduce_v1"(%arg0, %arg1)
-+  //  CHECK: "vhlo.reduce_v1"(%[[ARG0]], %[[ARG1]])
-   //  CHECK:   ^[[BB:bb.*]](%[[ARG1:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>, %[[ARG2:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
-   //  CHECK:     "vhlo.return_v1"(%[[VAL1:.*]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-   //  CHECK: }) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-@@ -1606,8 +1730,9 @@
- }
- 
- // CHECK-LABEL: "op_reduce_precision"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_reduce_precision(%arg0: tensor<f32>) -> tensor<f32> {
--  //      CHECK: "vhlo.reduce_precision_v1"(%arg0) <{
-+  //      CHECK: "vhlo.reduce_precision_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   exponent_bits = #vhlo.integer_v1<8 : i32>
-   // CHECK-SAME:   mantissa_bits = #vhlo.integer_v1<10 : i32>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-@@ -1619,8 +1744,9 @@
- }
- 
- // CHECK-LABEL: "op_reduce_scatter"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_reduce_scatter(%arg0: tensor<16xf32>) -> tensor<16xf32> {
--  //               CHECK: "vhlo.reduce_scatter_v1"(%arg0) <{
-+  //               CHECK: "vhlo.reduce_scatter_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0], [1]]> : tensor<2x1xi64>>,
-   //          CHECK-SAME:   scatter_dimension = #vhlo.integer_v1<0 : i64>
-@@ -1644,8 +1770,9 @@
- }
- 
- // CHECK-LABEL: "op_reduce_window"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_reduce_window(%arg0: tensor<2x17x31x7xf32>, %arg1: tensor<f32>) -> tensor<2x9x16x7xf32> {
--  //               CHECK: "vhlo.reduce_window_v1"(%arg0, %arg1) <{
-+  //               CHECK: "vhlo.reduce_window_v1"(%[[ARG0]], %[[ARG1]]) <{
-   //          CHECK-SAME:   base_dilations = #vhlo.tensor_v1<dense<[1, 2, 2, 1]> : tensor<4xi64>>,
-   // CHECK-SAME{LITERAL}:   padding = #vhlo.tensor_v1<dense<[[0, 0], [2, 0], [0, 2], [0, 0]]> : tensor<4x2xi64>>,
-   //          CHECK-SAME:   window_dilations = #vhlo.tensor_v1<dense<[1, 2, 2, 1]> : tensor<4xi64>>,
-@@ -1671,8 +1798,9 @@
- }
- 
- // CHECK-LABEL: "op_remainder"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_remainder(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.remainder_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.remainder_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.remainder"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
-@@ -1692,16 +1820,18 @@
- }
- 
- // CHECK-LABEL: "op_reshape"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_reshape(%arg0: tensor<16xf32>) -> tensor<4x4xf32> {
--  // CHECK: "vhlo.reshape_v1"(%arg0) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<4x4x!vhlo.f32_v1>
-+  // CHECK: "vhlo.reshape_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<4x4x!vhlo.f32_v1>
-   %0 = "stablehlo.reshape"(%arg0) : (tensor<16xf32>) -> tensor<4x4xf32>
-   func.return %0 : tensor<4x4xf32>
- }
- 
- // CHECK-LABEL: "op_return"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_return(%arg0: tensor<i32>, %arg1: tensor<f32>) -> tensor<f32> {
--  //      CHECK: "vhlo.case_v1"(%arg0) ({
--  // CHECK-NEXT:   "vhlo.return_v1"(%arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-+  //      CHECK: "vhlo.case_v1"(%[[ARG0]]) ({
-+  // CHECK-NEXT:   "vhlo.return_v1"(%[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-   // CHECK-NEXT: }) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.case"(%arg0) ({
-     "stablehlo.return"(%arg1) : (tensor<f32>) -> ()
-@@ -1710,8 +1840,9 @@
- }
- 
- // CHECK-LABEL: "op_reverse"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_reverse(%arg0: tensor<16xf32>) -> tensor<16xf32> {
--  //      CHECK: "vhlo.reverse_v1"(%arg0) <{
-+  //      CHECK: "vhlo.reverse_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   dimensions = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x!vhlo.f32_v1>
-   %0 = "stablehlo.reverse"(%arg0) {
-@@ -1721,8 +1852,9 @@
- }
- 
- // CHECK-LABEL: "op_rng_bit_generator"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_rng_bit_generator(%arg0: tensor<f32>) -> (tensor<f32>, tensor<f32>) {
--  //      CHECK: "vhlo.rng_bit_generator_v1"(%arg0) <{
-+  //      CHECK: "vhlo.rng_bit_generator_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   rng_algorithm = #vhlo<rng_algorithm_v1 PHILOX>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>)
-   %0:2 = "stablehlo.rng_bit_generator"(%arg0) {
-@@ -1732,8 +1864,9 @@
- }
- 
- // CHECK-LABEL: "op_rng"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_rng(%arg0: tensor<f32>, %arg1: tensor<f32>, %arg2: tensor<0xindex>) -> tensor<f32> {
--  //      CHECK: "vhlo.rng_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.rng_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   rng_distribution = #vhlo<rng_distribution_v1 NORMAL>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<0x!vhlo.index_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.rng"(%arg0, %arg1, %arg2) {
-@@ -1743,29 +1876,33 @@
- }
- 
- // CHECK-LABEL: "op_round_nearest_afz"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_round_nearest_afz(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.round_nearest_afz_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.round_nearest_afz_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.round_nearest_afz"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_round_nearest_even"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_round_nearest_even(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.round_nearest_even_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.round_nearest_even_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.round_nearest_even"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_rsqrt"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_rsqrt(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.rsqrt_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.rsqrt_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.rsqrt"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_scatter"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_scatter(%arg0: tensor<200x100x300xf32>, %arg1: tensor<10x2xi32>, %arg2: tensor<10x300xf32>) -> tensor<200x100x300xf32> {
--  //      CHECK: "vhlo.scatter_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.scatter_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   index_vector_dim = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   indices_are_sorted = #vhlo.bool_v1<true>,
-   // CHECK-SAME:   inserted_window_dims = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-@@ -1795,8 +1932,9 @@
- }
- 
- // CHECK-LABEL: "op_select_and_scatter"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_select_and_scatter(%arg0: tensor<10x24x24x64xf32>, %arg1: tensor<12x13x13x66xf32>, %arg2: tensor<f32>) -> tensor<10x24x24x64xf32> {
--  //      CHECK: "vhlo.select_and_scatter_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.select_and_scatter_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   padding = #vhlo.tensor_v1<dense<1> : tensor<4x2xi64>>,
-   // CHECK-SAME:   window_dimensions = #vhlo.tensor_v1<dense<[1, 2, 2, 1]> : tensor<4xi64>>,
-   // CHECK-SAME:   window_strides = #vhlo.tensor_v1<dense<[1, 2, 2, 1]> : tensor<4xi64>>
-@@ -1826,15 +1964,17 @@
- }
- 
- // CHECK-LABEL: "op_select"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_select(%arg0: tensor<i1>, %arg1: tensor<f32>, %arg2: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.select_v1"(%arg0, %arg1, %arg2) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.select_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.select"(%arg0, %arg1, %arg2) : (tensor<i1>, tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_send"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_send(%arg0: tensor<f32>, %arg1: !stablehlo.token) -> !stablehlo.token {
--  //      CHECK: "vhlo.send_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.send_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME:   channel_type = #vhlo.integer_v1<2 : i64>,
-   // CHECK-SAME:   is_host_transfer = #vhlo.bool_v1<true>
-@@ -1847,8 +1987,9 @@
- }
- 
- // CHECK-LABEL: "op_set_dimension_size"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_set_dimension_size(%arg0: tensor<?xf32>, %arg1: tensor<i32>) -> tensor<16xf32> {
--  //      CHECK: "vhlo.set_dimension_size_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.set_dimension_size_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   dimension = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<?x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<16x!vhlo.f32_v1>
-   %0 = "stablehlo.set_dimension_size"(%arg0, %arg1) {
-@@ -1858,43 +1999,49 @@
- }
- 
- // CHECK-LABEL: "op_shift_left"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_shift_left(%arg0: tensor<i32>, %arg1: tensor<i32>) -> tensor<i32> {
--  // CHECK: "vhlo.shift_left_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.i32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-+  // CHECK: "vhlo.shift_left_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-   %0 = "stablehlo.shift_left"(%arg0, %arg1) : (tensor<i32>, tensor<i32>) -> tensor<i32>
-   func.return %0 : tensor<i32>
- }
- 
- // CHECK-LABEL: "op_shift_right_arithmetic"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_shift_right_arithmetic(%arg0: tensor<i32>, %arg1: tensor<i32>) -> tensor<i32> {
--  // CHECK: "vhlo.shift_right_arithmetic_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.i32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-+  // CHECK: "vhlo.shift_right_arithmetic_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-   %0 = "stablehlo.shift_right_arithmetic"(%arg0, %arg1) : (tensor<i32>, tensor<i32>) -> tensor<i32>
-   func.return %0 : tensor<i32>
- }
- 
- // CHECK-LABEL: "op_shift_right_logical"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_shift_right_logical(%arg0: tensor<i32>, %arg1: tensor<i32>) -> tensor<i32> {
--  // CHECK: "vhlo.shift_right_logical_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.i32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-+  // CHECK: "vhlo.shift_right_logical_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-   %0 = "stablehlo.shift_right_logical"(%arg0, %arg1) : (tensor<i32>, tensor<i32>) -> tensor<i32>
-   func.return %0 : tensor<i32>
- }
- 
- // CHECK-LABEL: "op_sign"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_sign(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.sign_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.sign_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.sign"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_sine"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_sine(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.sine_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.sine_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.sine"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_slice"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_slice(%arg0: tensor<16xf32>) -> tensor<4xf32> {
--  //      CHECK: "vhlo.slice_v1"(%arg0) <{
-+  //      CHECK: "vhlo.slice_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   limit_indices = #vhlo.tensor_v1<dense<4> : tensor<1xi64>>,
-   // CHECK-SAME:   start_indices = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>,
-   // CHECK-SAME:   strides = #vhlo.tensor_v1<dense<1> : tensor<1xi64>>
-@@ -1908,8 +2055,9 @@
- }
- 
- // CHECK-LABEL: "op_sort"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_sort(%arg0: tensor<16xf32>) -> tensor<16xf32> {
--  //      CHECK: "vhlo.sort_v1"(%arg0) <{
-+  //      CHECK: "vhlo.sort_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   dimension = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME:   is_stable = #vhlo.bool_v1<true>
-   // CHECK-SAME: }> ({
-@@ -1929,29 +2077,33 @@
- }
- 
- // CHECK-LABEL: "op_sqrt"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_sqrt(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.sqrt_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.sqrt_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.sqrt"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_subtract"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_subtract(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.subtract_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.subtract_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.subtract"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_tanh"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_tanh(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.tanh_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.tanh_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.tanh"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_torch_index_select"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_torch_index_select(%arg0: tensor<5x1x5xf32>, %arg1: tensor<2xi32>) ->  tensor<2x1x5xf32> {
--  //      CHECK: "vhlo.torch_index_select_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.torch_index_select_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   batch_dims = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME:   dim = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<5x1x5x!vhlo.f32_v1>, !vhlo.tensor_v1<2x!vhlo.i32_v1>) -> !vhlo.tensor_v1<2x1x5x!vhlo.f32_v1>
-@@ -1963,8 +2115,9 @@
- }
- 
- // CHECK-LABEL: "op_trace"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_trace(%arg0: tensor<f32>) {
--  //      CHECK: "vhlo.trace_v1"(%arg0) <{
-+  //      CHECK: "vhlo.trace_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   tag = #vhlo.string_v1<"foo">
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-   "stablehlo.trace"(%arg0) {
-@@ -1974,8 +2127,9 @@
- }
- 
- // CHECK-LABEL: "op_transpose"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_transpose(%arg0: tensor<16x8xf32>) ->  tensor<8x16xf32> {
--  //      CHECK: "vhlo.transpose_v1"(%arg0) <{
-+  //      CHECK: "vhlo.transpose_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   permutation = #vhlo.tensor_v1<dense<[1, 0]> : tensor<2xi64>>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<16x8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<8x16x!vhlo.f32_v1>
-   %0 = "stablehlo.transpose"(%arg0) {
-@@ -1985,8 +2139,9 @@
- }
- 
- // CHECK-LABEL: "op_triangular_solve"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_triangular_solve(%arg0: tensor<16x16xf32>, %arg1: tensor<16x16xf32>) ->  tensor<16x16xf32> {
--  //      CHECK: "vhlo.triangular_solve_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.triangular_solve_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   left_side = #vhlo.bool_v1<true>,
-   // CHECK-SAME:   lower = #vhlo.bool_v1<true>,
-   // CHECK-SAME:   transpose_a = #vhlo<transpose_v1 NO_TRANSPOSE>,
-@@ -2002,15 +2157,17 @@
- }
- 
- // CHECK-LABEL: "op_tuple"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_tuple(%arg0: tensor<f32>) -> tuple<tensor<f32>> {
--  // CHECK: "vhlo.tuple_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tuple_v1<!vhlo.tensor_v1<!vhlo.f32_v1>>
-+  // CHECK: "vhlo.tuple_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tuple_v1<!vhlo.tensor_v1<!vhlo.f32_v1>>
-   %0 = "stablehlo.tuple"(%arg0) : (tensor<f32>) -> tuple<tensor<f32>>
-   func.return %0 : tuple<tensor<f32>>
- }
- 
- // CHECK-LABEL: "op_unary_einsum"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_unary_einsum(%arg0: tensor<8x16xf32>) -> tensor<8xf32> {
--  //      CHECK: "vhlo.unary_einsum_v1"(%arg0) <{
-+  //      CHECK: "vhlo.unary_einsum_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   einsum_config = #vhlo.string_v1<"ab->a">
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<8x16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<8x!vhlo.f32_v1>
-   %0 = "stablehlo.unary_einsum"(%arg0) {
-@@ -2020,22 +2177,25 @@
- }
- 
- // CHECK-LABEL: "op_uniform_dequantize"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_uniform_dequantize(%arg0: tensor<!quant.uniform<i8:f32, 34.0:16>>) -> tensor<f32> {
--  // CHECK: "vhlo.uniform_dequantize_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.uniform_dequantize_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.uniform_dequantize"(%arg0) : (tensor<!quant.uniform<i8:f32, 34.0:16>>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_uniform_quantize"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_uniform_quantize(%arg0: tensor<f32>) -> tensor<!quant.uniform<i8:f32, 34.0:16>> {
--  // CHECK: "vhlo.uniform_quantize_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>
-+  // CHECK: "vhlo.uniform_quantize_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>
-   %0 = "stablehlo.uniform_quantize"(%arg0) : (tensor<f32>) -> tensor<!quant.uniform<i8:f32, 34.0:16>>
-   func.return %0 : tensor<!quant.uniform<i8:f32, 34.0:16>>
- }
- 
- // CHECK-LABEL: "op_while"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_while(%arg0: tensor<i1>) -> tensor<i1> {
--  //      CHECK: "vhlo.while_v1"(%arg0) ({
-+  //      CHECK: "vhlo.while_v1"(%[[ARG0]]) ({
-   // CHECK-NEXT:   ^[[BB:bb.*]](%[[ARG1:arg.*]]: !vhlo.tensor_v1<!vhlo.bool_v1>):
-   // CHECK-NEXT:     "vhlo.return_v1"(%[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>) -> ()
-   // CHECK-NEXT:   }, {
-@@ -2053,8 +2213,9 @@
- }
- 
- // CHECK-LABEL: "op_xor"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_xor(%arg0: tensor<i1>, %arg1: tensor<i1>) -> tensor<i1> {
--  // CHECK: "vhlo.xor_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-+  // CHECK: "vhlo.xor_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-   %0 = "stablehlo.xor"(%arg0, %arg1) : (tensor<i1>, tensor<i1>) -> tensor<i1>
-   func.return %0 : tensor<i1>
- }
-@@ -2062,190 +2223,217 @@
- // ============ TYPES ============
- 
- // CHECK-LABEL: "type_i1"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_i1(%arg0: tensor<i1>, %arg1: tensor<i1>) -> tensor<i1> {
--  // CHECK: "vhlo.and_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-+  // CHECK: "vhlo.and_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-   %0 = "stablehlo.and"(%arg0, %arg1) : (tensor<i1>, tensor<i1>) -> tensor<i1>
-   func.return %0 : tensor<i1>
- }
- 
- // CHECK-LABEL: "type_i4"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_i4(%arg0: tensor<i4>, %arg1: tensor<i4>) -> tensor<i4> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.i4_v1>, !vhlo.tensor_v1<!vhlo.i4_v1>) -> !vhlo.tensor_v1<!vhlo.i4_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i4_v1>, !vhlo.tensor_v1<!vhlo.i4_v1>) -> !vhlo.tensor_v1<!vhlo.i4_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<i4>, tensor<i4>) -> tensor<i4>
-   func.return %0 : tensor<i4>
- }
- 
- // CHECK-LABEL: "type_i8"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_i8(%arg0: tensor<i8>, %arg1: tensor<i8>) -> tensor<i8> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.i8_v1>, !vhlo.tensor_v1<!vhlo.i8_v1>) -> !vhlo.tensor_v1<!vhlo.i8_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i8_v1>, !vhlo.tensor_v1<!vhlo.i8_v1>) -> !vhlo.tensor_v1<!vhlo.i8_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<i8>, tensor<i8>) -> tensor<i8>
-   func.return %0 : tensor<i8>
- }
- 
- // CHECK-LABEL: "type_i16"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_i16(%arg0: tensor<i16>, %arg1: tensor<i16>) -> tensor<i16> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.i16_v1>, !vhlo.tensor_v1<!vhlo.i16_v1>) -> !vhlo.tensor_v1<!vhlo.i16_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i16_v1>, !vhlo.tensor_v1<!vhlo.i16_v1>) -> !vhlo.tensor_v1<!vhlo.i16_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<i16>, tensor<i16>) -> tensor<i16>
-   func.return %0 : tensor<i16>
- }
- 
- // CHECK-LABEL: "type_i32"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_i32(%arg0: tensor<i32>, %arg1: tensor<i32>) -> tensor<i32> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.i32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<i32>, tensor<i32>) -> tensor<i32>
-   func.return %0 : tensor<i32>
- }
- 
- // CHECK-LABEL: "type_i64"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_i64(%arg0: tensor<i64>, %arg1: tensor<i64>) -> tensor<i64> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.i64_v1>, !vhlo.tensor_v1<!vhlo.i64_v1>) -> !vhlo.tensor_v1<!vhlo.i64_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i64_v1>, !vhlo.tensor_v1<!vhlo.i64_v1>) -> !vhlo.tensor_v1<!vhlo.i64_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<i64>, tensor<i64>) -> tensor<i64>
-   func.return %0 : tensor<i64>
- }
- 
- // CHECK-LABEL: "type_ui4"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_ui4(%arg0: tensor<ui4>, %arg1: tensor<ui4>) -> tensor<ui4> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.ui4_v1>, !vhlo.tensor_v1<!vhlo.ui4_v1>) -> !vhlo.tensor_v1<!vhlo.ui4_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.ui4_v1>, !vhlo.tensor_v1<!vhlo.ui4_v1>) -> !vhlo.tensor_v1<!vhlo.ui4_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<ui4>, tensor<ui4>) -> tensor<ui4>
-   func.return %0 : tensor<ui4>
- }
- 
- // CHECK-LABEL: "type_ui8"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_ui8(%arg0: tensor<ui8>, %arg1: tensor<ui8>) -> tensor<ui8> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.ui8_v1>, !vhlo.tensor_v1<!vhlo.ui8_v1>) -> !vhlo.tensor_v1<!vhlo.ui8_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.ui8_v1>, !vhlo.tensor_v1<!vhlo.ui8_v1>) -> !vhlo.tensor_v1<!vhlo.ui8_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<ui8>, tensor<ui8>) -> tensor<ui8>
-   func.return %0 : tensor<ui8>
- }
- 
- // CHECK-LABEL: "type_ui16"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_ui16(%arg0: tensor<ui16>, %arg1: tensor<ui16>) -> tensor<ui16> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.ui16_v1>, !vhlo.tensor_v1<!vhlo.ui16_v1>) -> !vhlo.tensor_v1<!vhlo.ui16_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.ui16_v1>, !vhlo.tensor_v1<!vhlo.ui16_v1>) -> !vhlo.tensor_v1<!vhlo.ui16_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<ui16>, tensor<ui16>) -> tensor<ui16>
-   func.return %0 : tensor<ui16>
- }
- 
- // CHECK-LABEL: "type_ui32"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_ui32(%arg0: tensor<ui32>, %arg1: tensor<ui32>) -> tensor<ui32> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.ui32_v1>, !vhlo.tensor_v1<!vhlo.ui32_v1>) -> !vhlo.tensor_v1<!vhlo.ui32_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.ui32_v1>, !vhlo.tensor_v1<!vhlo.ui32_v1>) -> !vhlo.tensor_v1<!vhlo.ui32_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<ui32>, tensor<ui32>) -> tensor<ui32>
-   func.return %0 : tensor<ui32>
- }
- 
- // CHECK-LABEL: "type_ui64"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_ui64(%arg0: tensor<ui64>, %arg1: tensor<ui64>) -> tensor<ui64> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.ui64_v1>, !vhlo.tensor_v1<!vhlo.ui64_v1>) -> !vhlo.tensor_v1<!vhlo.ui64_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.ui64_v1>, !vhlo.tensor_v1<!vhlo.ui64_v1>) -> !vhlo.tensor_v1<!vhlo.ui64_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<ui64>, tensor<ui64>) -> tensor<ui64>
-   func.return %0 : tensor<ui64>
- }
- 
- // CHECK-LABEL: "type_f8E4M3FN"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_f8E4M3FN(%arg0: tensor<f8E4M3FN>, %arg1: tensor<f8E4M3FN>) -> tensor<f8E4M3FN> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f8E4M3FN_v1>, !vhlo.tensor_v1<!vhlo.f8E4M3FN_v1>) -> !vhlo.tensor_v1<!vhlo.f8E4M3FN_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f8E4M3FN_v1>, !vhlo.tensor_v1<!vhlo.f8E4M3FN_v1>) -> !vhlo.tensor_v1<!vhlo.f8E4M3FN_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f8E4M3FN>, tensor<f8E4M3FN>) -> tensor<f8E4M3FN>
-   func.return %0 : tensor<f8E4M3FN>
- }
- 
- // CHECK-LABEL: "type_f8E5M2"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_f8E5M2(%arg0: tensor<f8E5M2>, %arg1: tensor<f8E5M2>) -> tensor<f8E5M2> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f8E5M2_v1>, !vhlo.tensor_v1<!vhlo.f8E5M2_v1>) -> !vhlo.tensor_v1<!vhlo.f8E5M2_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f8E5M2_v1>, !vhlo.tensor_v1<!vhlo.f8E5M2_v1>) -> !vhlo.tensor_v1<!vhlo.f8E5M2_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f8E5M2>, tensor<f8E5M2>) -> tensor<f8E5M2>
-   func.return %0 : tensor<f8E5M2>
- }
- 
- // CHECK-LABEL: "type_f8E4M3FNUZ"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_f8E4M3FNUZ(%arg0: tensor<f8E4M3FNUZ>, %arg1: tensor<f8E4M3FNUZ>) -> tensor<f8E4M3FNUZ> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f8E4M3FNUZ_v1>, !vhlo.tensor_v1<!vhlo.f8E4M3FNUZ_v1>) -> !vhlo.tensor_v1<!vhlo.f8E4M3FNUZ_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f8E4M3FNUZ_v1>, !vhlo.tensor_v1<!vhlo.f8E4M3FNUZ_v1>) -> !vhlo.tensor_v1<!vhlo.f8E4M3FNUZ_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f8E4M3FNUZ>, tensor<f8E4M3FNUZ>) -> tensor<f8E4M3FNUZ>
-   func.return %0 : tensor<f8E4M3FNUZ>
- }
- 
- // CHECK-LABEL: "type_f8E4M3B11FNUZ"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_f8E4M3B11FNUZ(%arg0: tensor<f8E4M3B11FNUZ>, %arg1: tensor<f8E4M3B11FNUZ>) -> tensor<f8E4M3B11FNUZ> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f8E4M3B11FNUZ_v1>, !vhlo.tensor_v1<!vhlo.f8E4M3B11FNUZ_v1>) -> !vhlo.tensor_v1<!vhlo.f8E4M3B11FNUZ_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f8E4M3B11FNUZ_v1>, !vhlo.tensor_v1<!vhlo.f8E4M3B11FNUZ_v1>) -> !vhlo.tensor_v1<!vhlo.f8E4M3B11FNUZ_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f8E4M3B11FNUZ>, tensor<f8E4M3B11FNUZ>) -> tensor<f8E4M3B11FNUZ>
-   func.return %0 : tensor<f8E4M3B11FNUZ>
- }
- 
- // CHECK-LABEL: "type_f8E5M2FNUZ"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_f8E5M2FNUZ(%arg0: tensor<f8E5M2FNUZ>, %arg1: tensor<f8E5M2FNUZ>) -> tensor<f8E5M2FNUZ> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f8E5M2FNUZ_v1>, !vhlo.tensor_v1<!vhlo.f8E5M2FNUZ_v1>) -> !vhlo.tensor_v1<!vhlo.f8E5M2FNUZ_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f8E5M2FNUZ_v1>, !vhlo.tensor_v1<!vhlo.f8E5M2FNUZ_v1>) -> !vhlo.tensor_v1<!vhlo.f8E5M2FNUZ_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f8E5M2FNUZ>, tensor<f8E5M2FNUZ>) -> tensor<f8E5M2FNUZ>
-   func.return %0 : tensor<f8E5M2FNUZ>
- }
- 
- // CHECK-LABEL: "type_bf16"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_bf16(%arg0: tensor<bf16>, %arg1: tensor<bf16>) -> tensor<bf16> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.bf16_v1>, !vhlo.tensor_v1<!vhlo.bf16_v1>) -> !vhlo.tensor_v1<!vhlo.bf16_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.bf16_v1>, !vhlo.tensor_v1<!vhlo.bf16_v1>) -> !vhlo.tensor_v1<!vhlo.bf16_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<bf16>, tensor<bf16>) -> tensor<bf16>
-   func.return %0 : tensor<bf16>
- }
- 
- // CHECK-LABEL: "type_f16"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_f16(%arg0: tensor<f16>, %arg1: tensor<f16>) -> tensor<f16> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f16_v1>, !vhlo.tensor_v1<!vhlo.f16_v1>) -> !vhlo.tensor_v1<!vhlo.f16_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f16_v1>, !vhlo.tensor_v1<!vhlo.f16_v1>) -> !vhlo.tensor_v1<!vhlo.f16_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f16>, tensor<f16>) -> tensor<f16>
-   func.return %0 : tensor<f16>
- }
- 
- // CHECK-LABEL: "type_f32"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_f32(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "type_f64"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_f64(%arg0: tensor<f64>, %arg1: tensor<f64>) -> tensor<f64> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f64_v1>, !vhlo.tensor_v1<!vhlo.f64_v1>) -> !vhlo.tensor_v1<!vhlo.f64_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f64_v1>, !vhlo.tensor_v1<!vhlo.f64_v1>) -> !vhlo.tensor_v1<!vhlo.f64_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f64>, tensor<f64>) -> tensor<f64>
-   func.return %0 : tensor<f64>
- }
- 
- // CHECK-LABEL: "type_complex_f32"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_complex_f32(%arg0: tensor<complex<f32>>, %arg1: tensor<complex<f32>>) -> tensor<complex<f32>> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>, !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>) -> !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>, !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>) -> !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<complex<f32>>, tensor<complex<f32>>) -> tensor<complex<f32>>
-   func.return %0 : tensor<complex<f32>>
- }
- 
- // CHECK-LABEL: "type_complex_f64"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_complex_f64(%arg0: tensor<complex<f64>>, %arg1: tensor<complex<f64>>) -> tensor<complex<f64>> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f64_v1>>, !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f64_v1>>) -> !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f64_v1>>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f64_v1>>, !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f64_v1>>) -> !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f64_v1>>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<complex<f64>>, tensor<complex<f64>>) -> tensor<complex<f64>>
-   func.return %0 : tensor<complex<f64>>
- }
- 
- // CHECK-LABEL: "type_dynamism_ranked"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @type_dynamism_ranked(%arg0: tensor<?xf32>) -> tensor<?xf32> {
--  // CHECK: "vhlo.abs_v1"(%arg0) : (!vhlo.tensor_v1<?x!vhlo.f32_v1>) -> !vhlo.tensor_v1<?x!vhlo.f32_v1>
-+  // CHECK: "vhlo.abs_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<?x!vhlo.f32_v1>) -> !vhlo.tensor_v1<?x!vhlo.f32_v1>
-   %0 = "stablehlo.abs"(%arg0) : (tensor<?xf32>) -> tensor<?xf32>
-   func.return %0 : tensor<?xf32>
- }
- 
- // CHECK-LABEL: "type_quantization"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_quantization(%arg0: tensor<!quant.uniform<i8:f32, 34.0:16>>, %arg1: tensor<!quant.uniform<i8:f32, 34.0:16>>) -> tensor<!quant.uniform<i8:f32, 34.0:16>> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>, !vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>) -> !vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>, !vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>) -> !vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<!quant.uniform<i8:f32, 34.0:16>>, tensor<!quant.uniform<i8:f32, 34.0:16>>) -> tensor<!quant.uniform<i8:f32, 34.0:16>>
-   func.return %0 : tensor<!quant.uniform<i8:f32, 34.0:16>>
- }
- 
- //       CHECK: function_type = #vhlo.type_v1<!vhlo.func_v1<(!vhlo.token_v1) -> !vhlo.token_v1>>
- // CHECK-LABEL: "type_token_callee"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @type_token_callee(%arg0: !stablehlo.token) -> !stablehlo.token {
--  // CHECK: "vhlo.return_v1"(%arg0) : (!vhlo.token_v1) -> ()
-+  // CHECK: "vhlo.return_v1"(%[[ARG0]]) : (!vhlo.token_v1) -> ()
-   return %arg0 : !stablehlo.token
- }
- 
- //       CHECK: function_type = #vhlo.type_v1<!vhlo.func_v1<(!vhlo.token_v1) -> !vhlo.token_v1>>
- // CHECK-LABEL: "type_token_caller"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @type_token_caller(%arg0: !stablehlo.token) -> !stablehlo.token {
--  // CHECK:      "vhlo.call_v1"(%arg0) <{callee = #vhlo.string_v1<"type_token_callee">}
-+  // CHECK:      "vhlo.call_v1"(%[[ARG0]]) <{callee = #vhlo.string_v1<"type_token_callee">}
-   // CHECK-SAME: (!vhlo.token_v1) -> !vhlo.token_v1
-   %0 = func.call @type_token_callee(%arg0) : (!stablehlo.token) -> !stablehlo.token
-   return %0 : !stablehlo.token
- }
- 
- // CHECK-LABEL: "type_tuple"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @type_tuple(%arg0: tuple<tensor<f32>>) -> tuple<!stablehlo.token> {
-   %0 = "stablehlo.custom_call"(%arg0) {
-     call_target_name = "foo"
-diff --ruN a/stablehlo/stablehlo/tests/vhlo/stablehlo_legalize_to_vhlo.0_15_0.mlir b/stablehlo/stablehlo/tests/vhlo/stablehlo_legalize_to_vhlo.0_15_0.mlir
---- stablehlo/stablehlo/tests/vhlo/stablehlo_legalize_to_vhlo.0_15_0.mlir
-+++ stablehlo/stablehlo/tests/vhlo/stablehlo_legalize_to_vhlo.0_15_0.mlir
-@@ -13,6 +13,7 @@
- // ============ ATTRIBUTES ============
- 
- // CHECK-LABEL: "attr_comparison_direction_eq"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_direction_eq(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     // CHECK: comparison_direction = #vhlo<comparison_direction_v1 EQ>
-@@ -22,6 +23,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_direction_ne"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_direction_ne(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     // CHECK: comparison_direction = #vhlo<comparison_direction_v1 NE>
-@@ -31,6 +33,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_direction_ge"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_direction_ge(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     // CHECK: comparison_direction = #vhlo<comparison_direction_v1 GE>
-@@ -40,6 +43,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_direction_gt"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_direction_gt(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     // CHECK: comparison_direction = #vhlo<comparison_direction_v1 GT>
-@@ -49,6 +53,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_direction_le"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_direction_le(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     // CHECK: comparison_direction = #vhlo<comparison_direction_v1 LE>
-@@ -58,6 +63,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_direction_lt"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_direction_lt(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     // CHECK: comparison_direction = #vhlo<comparison_direction_v1 LT>
-@@ -67,6 +73,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_type_notype"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_type_notype(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     comparison_direction = #stablehlo<comparison_direction EQ>
-@@ -76,6 +83,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_type_float"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_type_float(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     comparison_direction = #stablehlo<comparison_direction EQ>,
-@@ -86,6 +94,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_type_totalorder"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_type_totalorder(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     comparison_direction = #stablehlo<comparison_direction EQ>,
-@@ -96,6 +105,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_type_signed"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_type_signed(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     comparison_direction = #stablehlo<comparison_direction EQ>,
-@@ -106,6 +116,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_type_unsigned"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_type_unsigned(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     comparison_direction = #stablehlo<comparison_direction EQ>,
-@@ -118,6 +129,7 @@
- // ConvDimensionNumbers aka #stablehlo.conv is covered below.
- 
- // CHECK-LABEL: "attr_custom_call_api_version_unspecified"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_custom_call_api_version_unspecified(%arg0: tensor<f32>) -> tensor<f32> {
-   %0 = "stablehlo.custom_call"(%arg0) {
-     call_target_name = "foo",
-@@ -128,6 +140,7 @@
- }
- 
- // CHECK-LABEL: "attr_custom_call_api_version_original"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_custom_call_api_version_original(%arg0: tensor<f32>) -> tensor<f32> {
-   %0 = "stablehlo.custom_call"(%arg0) {
-     call_target_name = "foo",
-@@ -138,6 +151,7 @@
- }
- 
- // CHECK-LABEL: "attr_custom_call_api_version_status_returning"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_custom_call_api_version_status_returning(%arg0: tensor<f32>) -> tensor<f32> {
-   %0 = "stablehlo.custom_call"(%arg0) {
-     call_target_name = "foo",
-@@ -148,6 +162,7 @@
- }
- 
- // CHECK-LABEL: "attr_custom_call_api_version_status_returning_unified"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_custom_call_api_version_status_returning_unified(%arg0: tensor<f32>) -> tensor<f32> {
-   %0 = "stablehlo.custom_call"(%arg0) {
-     call_target_name = "foo",
-@@ -166,6 +181,7 @@
- // DotDimensionNumbers aka #stablehlo.dot is covered below.
- 
- // CHECK-LABEL: "attr_fft_type_fft"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_fft_type_fft(%arg0: tensor<16xcomplex<f32>>) -> tensor<16xcomplex<f32>> {
-   %0 = "stablehlo.fft"(%arg0) {
-     // CHECK: fft_type = #vhlo<fft_type_v1 FFT>
-@@ -176,6 +192,7 @@
- }
- 
- // CHECK-LABEL: "attr_fft_type_ifft"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_fft_type_ifft(%arg0: tensor<16xcomplex<f32>>) -> tensor<16xcomplex<f32>> {
-   %0 = "stablehlo.fft"(%arg0) {
-     // CHECK: fft_type = #vhlo<fft_type_v1 IFFT>
-@@ -186,6 +203,7 @@
- }
- 
- // CHECK-LABEL: "attr_fft_type_rfft"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_fft_type_rfft(%arg0: tensor<16xf32>) -> tensor<9xcomplex<f32>> {
-   %0 = "stablehlo.fft"(%arg0) {
-     // CHECK: fft_type = #vhlo<fft_type_v1 RFFT>
-@@ -196,6 +214,7 @@
- }
- 
- // CHECK-LABEL: "attr_fft_type_irfft"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_fft_type_irfft(%arg0: tensor<9xcomplex<f32>>) -> tensor<16xf32> {
-   %0 = "stablehlo.fft"(%arg0) {
-     // CHECK: fft_type = #vhlo<fft_type_v1 IRFFT>
-@@ -208,6 +227,7 @@
- // GatherDimensionNumbers aka #stablehlo.gather is covered below.
- 
- // CHECK-LABEL: "attr_precision_config_default"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_precision_config_default(%arg0: tensor<8x16xf32>, %arg1: tensor<16x8xf32>) -> tensor<8x8xf32> {
-   %0 = "stablehlo.dot"(%arg0, %arg1) {
-     // CHECK: precision_config = #vhlo.array_v1<[#vhlo<precision_v1 DEFAULT>, #vhlo<precision_v1 DEFAULT>]>
-@@ -216,6 +236,7 @@
- }
- 
- // CHECK-LABEL: "attr_precision_config_high"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_precision_config_high(%arg0: tensor<8x16xf32>, %arg1: tensor<16x8xf32>) -> tensor<8x8xf32> {
-   %0 = "stablehlo.dot"(%arg0, %arg1) {
-     // CHECK: precision_config = #vhlo.array_v1<[#vhlo<precision_v1 HIGH>, #vhlo<precision_v1 HIGH>]>
-@@ -225,6 +246,7 @@
- }
- 
- // CHECK-LABEL: "attr_precision_config_highest"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_precision_config_highest(%arg0: tensor<8x16xf32>, %arg1: tensor<16x8xf32>) -> tensor<8x8xf32> {
-   %0 = "stablehlo.dot"(%arg0, %arg1) {
-     // CHECK: precision_config = #vhlo.array_v1<[#vhlo<precision_v1 HIGHEST>, #vhlo<precision_v1 HIGHEST>]>
-@@ -234,6 +256,7 @@
- }
- 
- // CHECK-LABEL: "attr_rng_algorithm_default"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_rng_algorithm_default(%arg0: tensor<f32>) -> (tensor<f32>, tensor<f32>) {
-   %0:2 = "stablehlo.rng_bit_generator"(%arg0) {
-     // CHECK: rng_algorithm = #vhlo<rng_algorithm_v1 DEFAULT>
-@@ -243,6 +266,7 @@
- }
- 
- // CHECK-LABEL: "attr_rng_algorithm_three_fry"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_rng_algorithm_three_fry(%arg0: tensor<f32>) -> (tensor<f32>, tensor<f32>) {
-   %0:2 = "stablehlo.rng_bit_generator"(%arg0) {
-     // CHECK: rng_algorithm = #vhlo<rng_algorithm_v1 THREE_FRY>
-@@ -252,6 +276,7 @@
- }
- 
- // CHECK-LABEL: "attr_rng_algorithm_philox"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_rng_algorithm_philox(%arg0: tensor<f32>) -> (tensor<f32>, tensor<f32>) {
-   %0:2 = "stablehlo.rng_bit_generator"(%arg0) {
-     // CHECK: rng_algorithm = #vhlo<rng_algorithm_v1 PHILOX>
-@@ -261,6 +286,7 @@
- }
- 
- // CHECK-LABEL: "attr_rng_distribution_uniform"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @attr_rng_distribution_uniform(%arg0: tensor<f32>, %arg1: tensor<f32>, %arg2: tensor<0xindex>) -> tensor<f32> {
-   %0 = "stablehlo.rng"(%arg0, %arg1, %arg2) {
-     // CHECK: rng_distribution = #vhlo<rng_distribution_v1 UNIFORM>
-@@ -270,6 +296,7 @@
- }
- 
- // CHECK-LABEL: "attr_rng_distribution_normal"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @attr_rng_distribution_normal(%arg0: tensor<f32>, %arg1: tensor<f32>, %arg2: tensor<0xindex>) -> tensor<f32> {
-   %0 = "stablehlo.rng"(%arg0, %arg1, %arg2) {
-     // CHECK: rng_distribution = #vhlo<rng_distribution_v1 NORMAL>
-@@ -281,6 +308,7 @@
- // ScatterDimensionNumbers aka #stablehlo.scatter is covered below.
- 
- // CHECK-LABEL: "attr_transpose_no_transpose"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_transpose_no_transpose(%arg0: tensor<16x16xf32>, %arg1: tensor<16x16xf32>) ->  tensor<16x16xf32> {
-   %0 = "stablehlo.triangular_solve"(%arg0, %arg1) {
-     left_side = true,
-@@ -293,6 +321,7 @@
- }
- 
- // CHECK-LABEL: "attr_transpose_transpose"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_transpose_transpose(%arg0: tensor<16x16xf32>, %arg1: tensor<16x16xf32>) ->  tensor<16x16xf32> {
-   %0 = "stablehlo.triangular_solve"(%arg0, %arg1) {
-     left_side = true,
-@@ -305,6 +334,7 @@
- }
- 
- // CHECK-LABEL: "attr_transpose_adjoint"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_transpose_adjoint(%arg0: tensor<16x16xf32>, %arg1: tensor<16x16xf32>) ->  tensor<16x16xf32> {
-   %0 = "stablehlo.triangular_solve"(%arg0, %arg1) {
-     left_side = true,
-@@ -319,10 +349,9 @@
- // TypeExtensionsAttr aka #stablehlo.type_extensions is covered below.
- 
- // CHECK-LABEL: "attr_type_extensions_bounds"
--func.func @attr_type_extensions_bounds(
--    %arg0: tensor<?x?xf32, #stablehlo.type_extensions<bounds = [16, ?]>>)
--    -> tensor<?x?xf32, #stablehlo.type_extensions<bounds = [16, ?]>> {
--  // CHECK: "vhlo.return_v1"(%arg0) : (!vhlo.tensor_v1<?x?x!vhlo.f32_v1, #vhlo.type_extensions_v1<bounds = [16, ?]>>) -> ()
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @attr_type_extensions_bounds(%arg0: tensor<?x?xf32, #stablehlo.type_extensions<bounds = [16, ?]>>) -> tensor<?x?xf32, #stablehlo.type_extensions<bounds = [16, ?]>> {
-+  // CHECK: "vhlo.return_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<?x?x!vhlo.f32_v1, #vhlo.type_extensions_v1<bounds = [16, ?]>>) -> ()
-   func.return %arg0 : tensor<?x?xf32, #stablehlo.type_extensions<bounds = [16, ?]>>
- }
- 
-@@ -330,8 +359,9 @@
- // ============ DEFAULTS ============
- 
- // CHECK-LABEL: "default_all_gather"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_all_gather(%arg0: tensor<16x8xf32>) -> tensor<16x16xf32> {
--  //               CHECK: "vhlo.all_gather_v1"(%arg0) <{
-+  //               CHECK: "vhlo.all_gather_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   all_gather_dim = #vhlo.integer_v1<1 : i64>
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0], [1]]> : tensor<2x1xi64>>,
-@@ -345,8 +375,9 @@
- }
- 
- // CHECK-LABEL: "default_all_reduce"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_all_reduce(%arg0: tensor<f32>) -> tensor<f32> {
--  //               CHECK: "vhlo.all_reduce_v1"(%arg0)
-+  //               CHECK: "vhlo.all_reduce_v1"(%[[ARG0]])
-   //          CHECK-SAME: <{
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0], [1]]> : tensor<2x1xi64>>,
-@@ -368,8 +399,9 @@
- }
- 
- // CHECK-LABEL: "default_all_to_all"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_all_to_all(%arg0: tensor<4x16xf32>) -> tensor<16x4xf32> {
--  //               CHECK: "vhlo.all_to_all_v1"(%arg0) <{
-+  //               CHECK: "vhlo.all_to_all_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   //          CHECK-SAME:   concat_dimension = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0, 1, 2, 3]]> : tensor<1x4xi64>>,
-@@ -386,8 +418,9 @@
- }
- 
- // CHECK-LABEL: "default_cholesky"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_cholesky(%arg0: tensor<1x16x16xf32>) -> tensor<1x16x16xf32> {
--  //      CHECK: "vhlo.cholesky_v1"(%arg0) <{
-+  //      CHECK: "vhlo.cholesky_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   lower = #vhlo.bool_v1<false>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<1x16x16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<1x16x16x!vhlo.f32_v1>
-   %0 = "stablehlo.cholesky"(%arg0) : (tensor<1x16x16xf32>) -> tensor<1x16x16xf32>
-@@ -395,8 +428,9 @@
- }
- 
- // CHECK-LABEL: "default_collective_permute"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_collective_permute(%arg0: tensor<16x8xf32>) -> tensor<16x8xf32> {
--  //               CHECK: "vhlo.collective_permute_v1"(%arg0) <{
-+  //               CHECK: "vhlo.collective_permute_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME{LITERAL}:   source_target_pairs = #vhlo.tensor_v1<dense<[[0, 1], [1, 2], [2, 3]]> : tensor<3x2xi64>>
-   //          CHECK-SAME: }> : (!vhlo.tensor_v1<16x8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x8x!vhlo.f32_v1>
-@@ -407,8 +441,9 @@
- }
- 
- // CHECK-LABEL: "default_compare"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @default_compare(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
--  //      CHECK: "vhlo.compare_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.compare_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   compare_type = #vhlo<comparison_type_v1 NOTYPE>,
-   // CHECK-SAME:   comparison_direction = #vhlo<comparison_direction_v1 EQ>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-@@ -419,8 +454,9 @@
- }
- 
- // CHECK-LABEL: "default_convolution"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @default_convolution(%arg0: tensor<1x8x8x207xf32>, %arg1: tensor<3x3x207x16xf32>) -> tensor<1x6x6x16xf32> {
--  //      CHECK: "vhlo.convolution_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.convolution_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   batch_group_count = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   feature_group_count = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   input_batch_dimension = #vhlo.integer_v1<0 : i64>,
-@@ -448,8 +484,9 @@
- }
- 
- // CHECK-LABEL: "default_custom_call"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_custom_call(%arg0: tensor<f32>) -> tensor<f32> {
--  //      CHECK: "vhlo.custom_call_v1"(%arg0) <{
-+  //      CHECK: "vhlo.custom_call_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   api_version = #vhlo<api_version_v1 API_VERSION_ORIGINAL>,
-   // CHECK-SAME:   backend_config = #vhlo.string_v1<"">,
-   // CHECK-SAME:   call_target_name = #vhlo.string_v1<"foo">,
-@@ -466,8 +503,9 @@
- }
- 
- // CHECK-LABEL: "default_dot_general"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @default_dot_general(%arg0: tensor<8x8x16xf32>, %arg1: tensor<8x16x8xf32>) -> tensor<8x8x8xf32> {
--  //      CHECK: "vhlo.dot_general_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.dot_general_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   lhs_batching_dimensions = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>,
-   // CHECK-SAME:   lhs_contracting_dimensions = #vhlo.tensor_v1<dense<2> : tensor<1xi64>>,
-   // CHECK-SAME:   precision_config = #vhlo.array_v1<[#vhlo<precision_v1 DEFAULT>, #vhlo<precision_v1 DEFAULT>]>,
-@@ -486,8 +524,9 @@
- }
- 
- // CHECK-LABEL: "default_dot"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @default_dot(%arg0: tensor<8x16xf32>, %arg1: tensor<16x8xf32>) -> tensor<8x8xf32> {
--  //      CHECK: "vhlo.dot_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.dot_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   precision_config = #vhlo.array_v1<[#vhlo<precision_v1 DEFAULT>, #vhlo<precision_v1 DEFAULT>]>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<8x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<8x8x!vhlo.f32_v1>
-   %0 = "stablehlo.dot"(%arg0, %arg1) : (tensor<8x16xf32>, tensor<16x8xf32>) -> tensor<8x8xf32>
-@@ -495,8 +534,9 @@
- }
- 
- // CHECK-LABEL: "default_dynamic_broadcast_in_dim"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @default_dynamic_broadcast_in_dim(%arg0: tensor<?x?xf32>, %arg1: tensor<2xindex>) -> tensor<?x?xf32> {
--  //      CHECK: "vhlo.dynamic_broadcast_in_dim_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.dynamic_broadcast_in_dim_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   broadcast_dimensions = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-   // CHECK-SAME:   known_expanding_dimensions = #vhlo.tensor_v1<dense<> : tensor<0xi64>>,
-   // CHECK-SAME:   known_nonexpanding_dimensions = #vhlo.tensor_v1<dense<> : tensor<0xi64>>
-@@ -508,8 +548,9 @@
- }
- 
- // CHECK-LABEL: "default_dynamic_conv"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @default_dynamic_conv(%arg0: tensor<1x8x8x207xf32>, %arg1: tensor<3x3x207x16xf32>, %arg2: tensor<4xi32>) -> tensor<1x?x?x16xf32> {
--  //      CHECK: "vhlo.dynamic_conv_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.dynamic_conv_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   batch_group_count = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   feature_group_count = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   input_batch_dimension = #vhlo.integer_v1<0 : i64>,
-@@ -537,8 +578,9 @@
- }
- 
- // CHECK-LABEL: "default_dynamic_gather"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @default_dynamic_gather(%arg0 : tensor<2x4x9xf32>, %arg1 : tensor<1x5x2xi32>, %arg2 : tensor<3xi32>) -> tensor<1x5x8xf32> {
--  //      CHECK: "vhlo.dynamic_gather_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.dynamic_gather_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   collapsed_slice_dims = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-   // CHECK-SAME:   index_vector_dim = #vhlo.integer_v1<2 : i64>,
-   // CHECK-SAME:   indices_are_sorted = #vhlo.bool_v1<false>,
-@@ -564,15 +606,16 @@
-   // CHECK-SAME:   sym_name = #vhlo.string_v1<"default_func">,
-   // CHECK-SAME:   sym_visibility = #vhlo.string_v1<"">
-   // CHECK-SAME: }> ({
--  // CHECK-NEXT:   ^[[BB:bb.*]](%arg0: !vhlo.tensor_v1<!vhlo.f32_v1>):
--  // CHECK-NEXT:     "vhlo.return_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-+  // CHECK-NEXT:   ^[[BB:bb.*]](%[[ARG0:.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
-+  // CHECK-NEXT:     "vhlo.return_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-   // CHECK-NEXT: }) : () -> ()
-   func.return %arg0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "dynamic_gather"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @dynamic_gather(%arg0 : tensor<2x4x9xf32>, %arg1 : tensor<1x5x2xi32>) -> tensor<1x5x1xf32> {
--  //      CHECK: "vhlo.gather_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.gather_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   collapsed_slice_dims = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-   // CHECK-SAME:   index_vector_dim = #vhlo.integer_v1<2 : i64>,
-   // CHECK-SAME:   indices_are_sorted = #vhlo.bool_v1<false>,
-@@ -593,8 +636,9 @@
- }
- 
- // CHECK-LABEL: "default_infeed"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_infeed(%arg0: !stablehlo.token) -> (tensor<f32>, !stablehlo.token) {
--  //               CHECK: "vhlo.infeed_v1"(%arg0) <{
-+  //               CHECK: "vhlo.infeed_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   infeed_config = #vhlo.string_v1<"">,
-   // CHECK-SAME{LITERAL}:   layout = #vhlo.array_v1<[]>
-   //          CHECK-SAME: }> : (!vhlo.token_v1) -> (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.token_v1)
-@@ -603,8 +647,9 @@
- }
- 
- // CHECK-LABEL: "default_outfeed"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @default_outfeed(%arg0: tensor<f32>, %arg1: !stablehlo.token) -> !stablehlo.token {
--  //      CHECK: "vhlo.outfeed_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.outfeed_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   outfeed_config = #vhlo.string_v1<"">
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.token_v1) -> !vhlo.token_v1
-   %0 = "stablehlo.outfeed"(%arg0, %arg1) : (tensor<f32>, !stablehlo.token) -> !stablehlo.token
-@@ -612,8 +657,9 @@
- }
- 
- // CHECK-LABEL: "default_recv"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_recv(%arg0: !stablehlo.token) -> (tensor<f32>, !stablehlo.token) {
--  //      CHECK: "vhlo.recv_v1"(%arg0) <{
-+  //      CHECK: "vhlo.recv_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME:   channel_type = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   is_host_transfer = #vhlo.bool_v1<false>
-@@ -625,8 +671,9 @@
- }
- 
- // CHECK-LABEL: "default_send"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @default_send(%arg0: tensor<f32>, %arg1: !stablehlo.token) -> !stablehlo.token {
--  //      CHECK: "vhlo.send_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.send_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME:   channel_type = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   is_host_transfer = #vhlo.bool_v1<false>
-@@ -638,8 +685,9 @@
- }
- 
- // CHECK-LABEL: "default_reduce_scatter"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_reduce_scatter(%arg0: tensor<16xf32>) -> tensor<16xf32> {
--  //               CHECK: "vhlo.reduce_scatter_v1"(%arg0) <{
-+  //               CHECK: "vhlo.reduce_scatter_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0], [1]]> : tensor<2x1xi64>>,
-   //          CHECK-SAME:   scatter_dimension = #vhlo.integer_v1<0 : i64>
-@@ -661,8 +709,9 @@
- }
- 
- // CHECK-LABEL: "default_reduce_window"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @default_reduce_window(%arg0: tensor<2x17x31x7xf32>, %arg1: tensor<f32>) -> tensor<2x16x30x7xf32> {
--  //               CHECK: "vhlo.reduce_window_v1"(%arg0, %arg1)  <{
-+  //               CHECK: "vhlo.reduce_window_v1"(%[[ARG0]], %[[ARG1]])  <{
-   //          CHECK-SAME:   base_dilations = #vhlo.tensor_v1<dense<1> : tensor<4xi64>>,
-   // CHECK-SAME{LITERAL}:   padding = #vhlo.tensor_v1<dense<0> : tensor<4x2xi64>>,
-   //          CHECK-SAME:   window_dilations = #vhlo.tensor_v1<dense<1> : tensor<4xi64>>,
-@@ -684,8 +733,9 @@
- }
- 
- // CHECK-LABEL: "default_scatter"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @default_scatter(%arg0: tensor<200x100x300xf32>, %arg1: tensor<10x2xi32>, %arg2: tensor<10x300xf32>) -> tensor<200x100x300xf32> {
--  //      CHECK: "vhlo.scatter_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.scatter_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   index_vector_dim = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   indices_are_sorted = #vhlo.bool_v1<false>,
-   // CHECK-SAME:   inserted_window_dims = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-@@ -713,8 +763,9 @@
- }
- 
- // CHECK-LABEL: "default_select_and_scatter"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @default_select_and_scatter(%arg0: tensor<10x24x24x64xf32>, %arg1: tensor<10x23x23x64xf32>, %arg2: tensor<f32>) -> tensor<10x24x24x64xf32> {
--  //      CHECK: "vhlo.select_and_scatter_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.select_and_scatter_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   padding = #vhlo.tensor_v1<dense<0> : tensor<4x2xi64>>,
-   // CHECK-SAME:   window_dimensions = #vhlo.tensor_v1<dense<[1, 2, 2, 1]> : tensor<4xi64>>,
-   // CHECK-SAME:   window_strides = #vhlo.tensor_v1<dense<1> : tensor<4xi64>>
-@@ -742,8 +793,9 @@
- }
- 
- // CHECK-LABEL: "default_sort"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_sort(%arg0: tensor<16xf32>) -> tensor<16xf32> {
--  //      CHECK: "vhlo.sort_v1"(%arg0) <{
-+  //      CHECK: "vhlo.sort_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   dimension = #vhlo.integer_v1<-1 : i64>
-   // CHECK-SAME:   is_stable = #vhlo.bool_v1<false>
-   // CHECK-SAME: }> ({
-@@ -762,29 +814,33 @@
- // ============ OPS ============
- 
- // CHECK-LABEL: "op_abs"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_abs(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.abs_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.abs_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.abs"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_add"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_add(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_after_all"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_after_all(%arg0: !stablehlo.token) -> !stablehlo.token {
--  // CHECK: "vhlo.after_all_v1"(%arg0) : (!vhlo.token_v1) -> !vhlo.token_v1
-+  // CHECK: "vhlo.after_all_v1"(%[[ARG0]]) : (!vhlo.token_v1) -> !vhlo.token_v1
-   %0 = "stablehlo.after_all"(%arg0) : (!stablehlo.token) -> !stablehlo.token
-   func.return %0 : !stablehlo.token
- }
- 
- // CHECK-LABEL: "op_all_gather"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_all_gather(%arg0: tensor<16x8xf32>) -> tensor<16x16xf32> {
--  //               CHECK: "vhlo.all_gather_v1"(%arg0) <{
-+  //               CHECK: "vhlo.all_gather_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   all_gather_dim = #vhlo.integer_v1<1 : i64>
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0], [1]]> : tensor<2x1xi64>>,
-@@ -800,8 +856,9 @@
- }
- 
- // CHECK-LABEL: "op_all_reduce"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_all_reduce(%arg0: tensor<f32>) -> tensor<f32> {
--  //               CHECK: "vhlo.all_reduce_v1"(%arg0) <{
-+  //               CHECK: "vhlo.all_reduce_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0], [1]]> : tensor<2x1xi64>>,
-   //          CHECK-SAME:   use_global_device_ids = #vhlo.bool_v1<true>
-@@ -823,8 +880,9 @@
- }
- 
- // CHECK-LABEL: "op_all_to_all"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_all_to_all(%arg0: tensor<4x16xf32>) -> tensor<16x4xf32> {
--  //               CHECK: "vhlo.all_to_all_v1"(%arg0) <{
-+  //               CHECK: "vhlo.all_to_all_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<1 : i64>,
-   //          CHECK-SAME:   concat_dimension = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0, 1, 2, 3]]> : tensor<1x4xi64>>,
-@@ -842,22 +900,25 @@
- }
- 
- // CHECK-LABEL: "op_and"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_and(%arg0: tensor<i1>, %arg1: tensor<i1>) -> tensor<i1> {
--  // CHECK: "vhlo.and_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-+  // CHECK: "vhlo.and_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-   %0 = "stablehlo.and"(%arg0, %arg1) : (tensor<i1>, tensor<i1>) -> tensor<i1>
-   func.return %0 : tensor<i1>
- }
- 
- // CHECK-LABEL: "op_atan2"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_atan2(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.atan2_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.atan2_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.atan2"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_batch_norm_grad"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}}, %[[ARG3:.*]]: {{.*}}, %[[ARG4:.*]]: {{.*}})
- func.func @op_batch_norm_grad(%arg0: tensor<16x16x16x16xf32>, %arg1: tensor<16xf32>, %arg2: tensor<16xf32>, %arg3: tensor<16xf32>, %arg4: tensor<16x16x16x16xf32>) -> (tensor<16x16x16x16xf32>, tensor<16xf32>, tensor<16xf32>) {
--  //      CHECK: "vhlo.batch_norm_grad_v1"(%arg0, %arg1, %arg2, %arg3, %arg4) <{
-+  //      CHECK: "vhlo.batch_norm_grad_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]], %[[ARG3]], %[[ARG4]]) <{
-   // CHECK-SAME:   epsilon = #vhlo.float_v1<1.000000e-03 : !vhlo.f32_v1>,
-   // CHECK-SAME:   feature_index = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<16x16x16x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x16x16x16x!vhlo.f32_v1>) -> (!vhlo.tensor_v1<16x16x16x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>)
-@@ -869,8 +930,9 @@
- }
- 
- // CHECK-LABEL: "op_batch_norm_inference"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}}, %[[ARG3:.*]]: {{.*}}, %[[ARG4:.*]]: {{.*}})
- func.func @op_batch_norm_inference(%arg0: tensor<16x16x16x16xf32>, %arg1: tensor<16xf32>, %arg2: tensor<16xf32>, %arg3: tensor<16xf32>, %arg4: tensor<16xf32>) -> tensor<16x16x16x16xf32> {
--  //      CHECK: "vhlo.batch_norm_inference_v1"(%arg0, %arg1, %arg2, %arg3, %arg4) <{
-+  //      CHECK: "vhlo.batch_norm_inference_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]], %[[ARG3]], %[[ARG4]]) <{
-   // CHECK-SAME:   epsilon = #vhlo.float_v1<1.000000e-03 : !vhlo.f32_v1>,
-   // CHECK-SAME:   feature_index = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<16x16x16x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x16x16x16x!vhlo.f32_v1>
-@@ -882,8 +944,9 @@
- }
- 
- // CHECK-LABEL: "op_batch_norm_training"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_batch_norm_training(%arg0: tensor<16x16x16x16xf32>, %arg1: tensor<16xf32>, %arg2: tensor<16xf32>) -> (tensor<16x16x16x16xf32>, tensor<16xf32>, tensor<16xf32>) {
--  //      CHECK: "vhlo.batch_norm_training_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.batch_norm_training_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   epsilon = #vhlo.float_v1<1.000000e-03 : !vhlo.f32_v1>,
-   // CHECK-SAME:   feature_index = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<16x16x16x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>) -> (!vhlo.tensor_v1<16x16x16x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>)
-@@ -895,15 +958,17 @@
- }
- 
- // CHECK-LABEL: "op_bitcast_convert"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_bitcast_convert(%arg0: tensor<i32>) -> tensor<f32> {
--  // CHECK: "vhlo.bitcast_convert_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.bitcast_convert_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.bitcast_convert"(%arg0) : (tensor<i32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_broadcast_in_dim"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_broadcast_in_dim(%arg0: tensor<16xf32>) -> tensor<16x16xf32> {
--  //      CHECK: "vhlo.broadcast_in_dim_v1"(%arg0) <{
-+  //      CHECK: "vhlo.broadcast_in_dim_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   broadcast_dimensions = #vhlo.tensor_v1<dense<1> : tensor<1xi64>>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x16x!vhlo.f32_v1>
-   %0 = "stablehlo.broadcast_in_dim"(%arg0) {
-@@ -913,8 +978,9 @@
- }
- 
- // CHECK-LABEL: "op_broadcast"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_broadcast(%arg0: tensor<16xf32>) -> tensor<16x16xf32> {
--  //      CHECK: "vhlo.broadcast_v1"(%arg0) <{
-+  //      CHECK: "vhlo.broadcast_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   broadcast_sizes = #vhlo.tensor_v1<dense<16> : tensor<1xi64>>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x16x!vhlo.f32_v1>
-   %0 = "stablehlo.broadcast"(%arg0) {
-@@ -924,9 +990,10 @@
- }
- 
- // CHECK-LABEL: "op_case"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_case(%arg0: tensor<i32>, %arg1: tensor<f32>) -> tensor<f32> {
--  //      CHECK: "vhlo.case_v1"(%arg0) ({
--  // CHECK-NEXT:   "vhlo.return_v1"(%arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-+  //      CHECK: "vhlo.case_v1"(%[[ARG0]]) ({
-+  // CHECK-NEXT:   "vhlo.return_v1"(%[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-   // CHECK-NEXT: }) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.case"(%arg0) ({
-     "stablehlo.return"(%arg1) : (tensor<f32>) -> ()
-@@ -935,22 +1002,25 @@
- }
- 
- // CHECK-LABEL: "op_cbrt"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_cbrt(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.cbrt_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.cbrt_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.cbrt"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_ceil"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_ceil(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.ceil_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.ceil_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.ceil"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_cholesky"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_cholesky(%arg0: tensor<1x16x16xf32>) -> tensor<1x16x16xf32> {
--  //      CHECK: "vhlo.cholesky_v1"(%arg0) <{
-+  //      CHECK: "vhlo.cholesky_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   lower = #vhlo.bool_v1<true>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<1x16x16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<1x16x16x!vhlo.f32_v1>
-   %0 = "stablehlo.cholesky"(%arg0) {
-@@ -960,22 +1030,25 @@
- }
- 
- // CHECK-LABEL: "op_clamp"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_clamp(%arg0: tensor<f32>, %arg1: tensor<f32>, %arg2: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.clamp_v1"(%arg0, %arg1, %arg2) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.clamp_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.clamp"(%arg0, %arg1, %arg2) : (tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_count_leading_zeros"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_count_leading_zeros(%arg0: tensor<i32>) -> tensor<i32> {
--  // CHECK: "vhlo.count_leading_zeros_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-+  // CHECK: "vhlo.count_leading_zeros_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-   %0 = "stablehlo.count_leading_zeros"(%arg0) : (tensor<i32>) -> tensor<i32>
-   func.return %0 : tensor<i32>
- }
- 
- // CHECK-LABEL: "op_collective_permute"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_collective_permute(%arg0: tensor<16x8xf32>) -> tensor<16x8xf32> {
--  //               CHECK: "vhlo.collective_permute_v1"(%arg0) <{
-+  //               CHECK: "vhlo.collective_permute_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME{LITERAL}:   source_target_pairs = #vhlo.tensor_v1<dense<[[0, 1], [1, 2], [2, 3]]> : tensor<3x2xi64>>
-   //          CHECK-SAME: }> : (!vhlo.tensor_v1<16x8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x8x!vhlo.f32_v1>
-@@ -987,8 +1060,9 @@
- }
- 
- // CHECK-LABEL: "op_compare"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_compare(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
--  //      CHECK: "vhlo.compare_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.compare_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   compare_type = #vhlo<comparison_type_v1 TOTALORDER>,
-   // CHECK-SAME:   comparison_direction = #vhlo<comparison_direction_v1 EQ>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-@@ -1000,15 +1074,17 @@
- }
- 
- // CHECK-LABEL: "op_complex"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_complex(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<complex<f32>> {
--  // CHECK: "vhlo.complex_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>
-+  // CHECK: "vhlo.complex_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>
-   %0 = "stablehlo.complex"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<complex<f32>>
-   func.return %0 : tensor<complex<f32>>
- }
- 
- // CHECK-LABEL: "op_concatenate"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_concatenate(%arg0: tensor<8xf32>, %arg1: tensor<8xf32>) -> tensor<16xf32> {
--  //      CHECK: "vhlo.concatenate_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.concatenate_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   dimension = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<8x!vhlo.f32_v1>, !vhlo.tensor_v1<8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x!vhlo.f32_v1>
-   %0 = "stablehlo.concatenate"(%arg0, %arg1) {
-@@ -1018,6 +1094,7 @@
- }
- 
- // CHECK-LABEL: "op_constant"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_constant(%arg0: tensor<f32>) -> tensor<f32> {
-   //      CHECK: "vhlo.constant_v1"() <{
-   // CHECK-SAME:   value = #vhlo.tensor_v1<dense<0.000000e+00> : tensor<f32>>
-@@ -1029,15 +1106,17 @@
- }
- 
- // CHECK-LABEL: "op_convert"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_convert(%arg0: tensor<i32>) -> tensor<f32> {
--  // CHECK: "vhlo.convert_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.convert_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.convert"(%arg0) : (tensor<i32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_convolution"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_convolution(%arg0: tensor<1x8x8x207xf32>, %arg1: tensor<3x3x207x16xf32>) -> tensor<1x7x7x16xf32> {
--  //      CHECK: "vhlo.convolution_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.convolution_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   batch_group_count = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   feature_group_count = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   input_batch_dimension = #vhlo.integer_v1<0 : i64>,
-@@ -1071,8 +1150,9 @@
- }
- 
- // CHECK-LABEL: "op_cosine"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_cosine(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.cosine_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.cosine_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.cosine"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
-@@ -1085,8 +1165,9 @@
- }
- 
- // CHECK-LABEL: "op_cross_replica_sum"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_cross_replica_sum(%arg0: tensor<f32>) -> tensor<f32> {
--  //               CHECK: "vhlo.cross-replica-sum_v1"(%arg0) <{
-+  //               CHECK: "vhlo.cross-replica-sum_v1"(%[[ARG0]]) <{
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0], [1]]> : tensor<2x1xi64>>
-   //          CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.cross-replica-sum"(%arg0) {
-@@ -1096,8 +1177,9 @@
- }
- 
- // CHECK-LABEL: "op_custom_call"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_custom_call(%arg0: tensor<f32>) -> tensor<f32> {
--  //      CHECK: "vhlo.custom_call_v1"(%arg0) <{
-+  //      CHECK: "vhlo.custom_call_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   api_version = #vhlo<api_version_v1 API_VERSION_STATUS_RETURNING>,
-   // CHECK-SAME:   backend_config = #vhlo.string_v1<"\08\03\1A\02">,
-   // CHECK-SAME:   call_target_name = #vhlo.string_v1<"foo">,
-@@ -1128,15 +1210,17 @@
- }
- 
- // CHECK-LABEL: "op_divide"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_divide(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.divide_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.divide_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.divide"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_dot_general"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_dot_general(%arg0: tensor<8x8x16xf32>, %arg1: tensor<8x16x8xf32>) -> tensor<8x8x8xf32> {
--  //      CHECK: "vhlo.dot_general_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.dot_general_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   lhs_batching_dimensions = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>,
-   // CHECK-SAME:   lhs_contracting_dimensions = #vhlo.tensor_v1<dense<2> : tensor<1xi64>>,
-   // CHECK-SAME:   precision_config = #vhlo.array_v1<[#vhlo<precision_v1 HIGHEST>, #vhlo<precision_v1 HIGHEST>]>,
-@@ -1156,8 +1240,9 @@
- }
- 
- // CHECK-LABEL: "op_dot"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_dot(%arg0: tensor<8x16xf32>, %arg1: tensor<16x8xf32>) -> tensor<8x8xf32> {
--  //      CHECK: "vhlo.dot_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.dot_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   precision_config = #vhlo.array_v1<[#vhlo<precision_v1 HIGHEST>, #vhlo<precision_v1 HIGHEST>]>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<8x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<8x8x!vhlo.f32_v1>
-   %0 = "stablehlo.dot"(%arg0, %arg1) {
-@@ -1167,8 +1252,9 @@
- }
- 
- // CHECK-LABEL: "op_dynamic_broadcast_in_dim"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_dynamic_broadcast_in_dim(%arg0: tensor<?x?xf32>, %arg1: tensor<2xindex>) -> tensor<?x?xf32> {
--  //      CHECK: "vhlo.dynamic_broadcast_in_dim_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.dynamic_broadcast_in_dim_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   broadcast_dimensions = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-   // CHECK-SAME:   known_expanding_dimensions = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>,
-   // CHECK-SAME:   known_nonexpanding_dimensions = #vhlo.tensor_v1<dense<1> : tensor<1xi64>>
-@@ -1182,8 +1268,9 @@
- }
- 
- // CHECK-LABEL: "op_dynamic_conv"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_dynamic_conv(%arg0: tensor<1x8x8x207xf32>, %arg1: tensor<3x3x207x16xf32>, %arg2: tensor<4xi32>) -> tensor<1x?x?x16xf32> {
--  //      CHECK: "vhlo.dynamic_conv_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.dynamic_conv_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   batch_group_count = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   feature_group_count = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   input_batch_dimension = #vhlo.integer_v1<0 : i64>,
-@@ -1217,8 +1304,9 @@
- }
- 
- // CHECK-LABEL: "op_dynamic_gather"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_dynamic_gather(%arg0 : tensor<2x4x9xf32>, %arg1 : tensor<1x5x2xi32>, %arg2 : tensor<3xi32>) -> tensor<1x5x8xf32> {
--  //      CHECK: "vhlo.dynamic_gather_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.dynamic_gather_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   collapsed_slice_dims = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-   // CHECK-SAME:   index_vector_dim = #vhlo.integer_v1<2 : i64>,
-   // CHECK-SAME:   indices_are_sorted = #vhlo.bool_v1<true>,
-@@ -1238,8 +1326,9 @@
- }
- 
- // CHECK-LABEL: "op_dynamic_iota"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_dynamic_iota(%arg0: tensor<1xindex>) -> tensor<?xf32> {
--  //      CHECK: "vhlo.dynamic_iota_v1"(%arg0) <{
-+  //      CHECK: "vhlo.dynamic_iota_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   iota_dimension = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<1x!vhlo.index_v1>) -> !vhlo.tensor_v1<?x!vhlo.f32_v1>
-   %0 = "stablehlo.dynamic_iota"(%arg0) {
-@@ -1249,22 +1338,25 @@
- }
- 
- // CHECK-LABEL: "op_dynamic_pad"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}}, %[[ARG3:.*]]: {{.*}}, %[[ARG4:.*]]: {{.*}})
- func.func @op_dynamic_pad(%arg0: tensor<?xf32>, %arg1: tensor<f32>, %arg2: tensor<1xindex>, %arg3: tensor<1xindex>, %arg4: tensor<1xindex>) -> tensor<?xf32> {
--  // CHECK: "vhlo.dynamic_pad_v1"(%arg0, %arg1, %arg2, %arg3, %arg4) : (!vhlo.tensor_v1<?x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>) -> !vhlo.tensor_v1<?x!vhlo.f32_v1>
-+  // CHECK: "vhlo.dynamic_pad_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]], %[[ARG3]], %[[ARG4]]) : (!vhlo.tensor_v1<?x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>) -> !vhlo.tensor_v1<?x!vhlo.f32_v1>
-   %0 = "stablehlo.dynamic_pad"(%arg0, %arg1, %arg2, %arg3, %arg4) : (tensor<?xf32>, tensor<f32>, tensor<1xindex>, tensor<1xindex>, tensor<1xindex>) -> tensor<?xf32>
-   func.return %0 : tensor<?xf32>
- }
- 
- // CHECK-LABEL: "op_dynamic_reshape"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_dynamic_reshape(%arg0: tensor<16xf32>, %arg1: tensor<2xindex>) -> tensor<?x?xf32> {
--  // CHECK: "vhlo.dynamic_reshape_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<2x!vhlo.index_v1>) -> !vhlo.tensor_v1<?x?x!vhlo.f32_v1>
-+  // CHECK: "vhlo.dynamic_reshape_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<2x!vhlo.index_v1>) -> !vhlo.tensor_v1<?x?x!vhlo.f32_v1>
-   %0 = "stablehlo.dynamic_reshape"(%arg0, %arg1) : (tensor<16xf32>, tensor<2xindex>) -> tensor<?x?xf32>
-   func.return %0 : tensor<?x?xf32>
- }
- 
- // CHECK-LABEL: "op_dynamic_slice"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_dynamic_slice(%arg0: tensor<16xf32>, %arg1: tensor<i64>) -> tensor<4xf32> {
--  //      CHECK: "vhlo.dynamic_slice_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.dynamic_slice_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   slice_sizes = #vhlo.tensor_v1<dense<4> : tensor<1xi64>>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.i64_v1>) -> !vhlo.tensor_v1<4x!vhlo.f32_v1>
-   %0 = "stablehlo.dynamic_slice"(%arg0, %arg1) {
-@@ -1274,15 +1366,17 @@
- }
- 
- // CHECK-LABEL: "op_dynamic_update_slice"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_dynamic_update_slice(%arg0: tensor<16xf32>, %arg1: tensor<4xf32>, %arg2: tensor<i64>) -> tensor<16xf32> {
--  // CHECK: "vhlo.dynamic_update_slice_v1"(%arg0, %arg1, %arg2) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<4x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.i64_v1>) -> !vhlo.tensor_v1<16x!vhlo.f32_v1>
-+  // CHECK: "vhlo.dynamic_update_slice_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<4x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.i64_v1>) -> !vhlo.tensor_v1<16x!vhlo.f32_v1>
-   %0 = "stablehlo.dynamic_update_slice"(%arg0, %arg1, %arg2) : (tensor<16xf32>, tensor<4xf32>, tensor<i64>) -> tensor<16xf32>
-   func.return %0 : tensor<16xf32>
- }
- 
- // CHECK-LABEL: "op_einsum"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_einsum(%arg0: tensor<8x16xf32>, %arg1: tensor<16x8xf32>) -> tensor<8x8xf32> {
--  //      CHECK: "vhlo.einsum_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.einsum_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   einsum_config = #vhlo.string_v1<"ab,bc->ac">
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<8x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<8x8x!vhlo.f32_v1>
-   %0 = "stablehlo.einsum"(%arg0, %arg1) {
-@@ -1292,22 +1386,25 @@
- }
- 
- // CHECK-LABEL: "op_exponential_minus_one"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_exponential_minus_one(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.exponential_minus_one_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.exponential_minus_one_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.exponential_minus_one"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_exponential"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_exponential(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.exponential_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.exponential_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.exponential"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_fft"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_fft(%arg0: tensor<16xcomplex<f32>>) -> tensor<16xcomplex<f32>> {
--  //      CHECK: "vhlo.fft_v1"(%arg0) <{
-+  //      CHECK: "vhlo.fft_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   fft_length = #vhlo.tensor_v1<dense<16> : tensor<1xi64>>,
-   // CHECK-SAME:   fft_type = #vhlo<fft_type_v1 FFT>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<16x!vhlo.complex_v1<!vhlo.f32_v1>>) -> !vhlo.tensor_v1<16x!vhlo.complex_v1<!vhlo.f32_v1>>
-@@ -1319,8 +1416,9 @@
- }
- 
- // CHECK-LABEL: "op_floor"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_floor(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.floor_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.floor_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.floor"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
-@@ -1333,16 +1431,17 @@
-   // CHECK-SAME:   sym_name = #vhlo.string_v1<"op_func">,
-   // CHECK-SAME:   sym_visibility = #vhlo.string_v1<"private">
-   // CHECK-SAME: }> ({
--  // CHECK-NEXT:   ^[[BB:bb.*]](%arg0: !vhlo.tensor_v1<!vhlo.f32_v1>):
--  // CHECK-NEXT:     "vhlo.return_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-+  // CHECK-NEXT:   ^[[BB:bb.*]](%[[ARG0:.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
-+  // CHECK-NEXT:     "vhlo.return_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-   // CHECK-NEXT: }) : () -> ()
- 
-   func.return %arg0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_gather"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_gather(%arg0 : tensor<2x4x9xf32>, %arg1 : tensor<1x5x2xi32>) -> tensor<1x5x1xf32> {
--  //      CHECK: "vhlo.gather_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.gather_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   collapsed_slice_dims = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-   // CHECK-SAME:   index_vector_dim = #vhlo.integer_v1<2 : i64>,
-   // CHECK-SAME:   indices_are_sorted = #vhlo.bool_v1<true>,
-@@ -1364,8 +1463,9 @@
- }
- 
- // CHECK-LABEL: "op_get_dimension_size"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_get_dimension_size(%arg0: tensor<?xf32>) -> tensor<i32> {
--  //      CHECK: "vhlo.get_dimension_size_v1"(%arg0) <{
-+  //      CHECK: "vhlo.get_dimension_size_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   dimension = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<?x!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-   %0 = "stablehlo.get_dimension_size"(%arg0) {
-@@ -1375,8 +1475,9 @@
- }
- 
- // CHECK-LABEL: "op_get_tuple_element"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_get_tuple_element(%arg0: tuple<tensor<f32>, tensor<i32>>) -> tensor<f32> {
--  //      CHECK: "vhlo.get_tuple_element_v1"(%arg0) <{
-+  //      CHECK: "vhlo.get_tuple_element_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   index = #vhlo.integer_v1<0 : i32>
-   // CHECK-SAME: }> : (!vhlo.tuple_v1<!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.get_tuple_element"(%arg0) {
-@@ -1386,11 +1487,12 @@
- }
- 
- // CHECK-LABEL: "op_if"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_if(%arg0: tensor<i1>, %arg1: tensor<f32>, %arg2: tensor<f32>) -> tensor<f32> {
--  //      CHECK: "vhlo.if_v1"(%arg0) ({
--  // CHECK-NEXT:   "vhlo.return_v1"(%arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-+  //      CHECK: "vhlo.if_v1"(%[[ARG0]]) ({
-+  // CHECK-NEXT:   "vhlo.return_v1"(%[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-   // CHECK-NEXT: }, {
--  // CHECK-NEXT:   "vhlo.return_v1"(%arg2) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-+  // CHECK-NEXT:   "vhlo.return_v1"(%[[ARG2]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-   // CHECK-NEXT: }) : (!vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.if"(%arg0) ({
-     "stablehlo.return"(%arg1) : (tensor<f32>) -> ()
-@@ -1401,15 +1503,17 @@
- }
- 
- // CHECK-LABEL: "op_imag"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_imag(%arg0: tensor<complex<f32>>) -> tensor<f32> {
--  // CHECK: "vhlo.imag_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.imag_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.imag"(%arg0) : (tensor<complex<f32>>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_infeed"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_infeed(%arg0: !stablehlo.token) -> (tensor<f32>, !stablehlo.token) {
--  //               CHECK: "vhlo.infeed_v1"(%arg0) <{
-+  //               CHECK: "vhlo.infeed_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   infeed_config = #vhlo.string_v1<"foo">,
-   // CHECK-SAME{LITERAL}:   layout = #vhlo.array_v1<[#vhlo.array_v1<[]>]>
-   //          CHECK-SAME: }> : (!vhlo.token_v1) -> (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.token_v1)
-@@ -1432,36 +1536,41 @@
- }
- 
- // CHECK-LABEL: "op_is_finite"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_is_finite(%arg0: tensor<f32>) -> tensor<i1> {
--  // CHECK: "vhlo.is_finite_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-+  // CHECK: "vhlo.is_finite_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-   %0 = "stablehlo.is_finite"(%arg0) : (tensor<f32>) -> tensor<i1>
-   func.return %0 : tensor<i1>
- }
- 
- // CHECK-LABEL: "op_log"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_log(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.log_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.log_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.log"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_log_plus_one"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_log_plus_one(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.log_plus_one_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.log_plus_one_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.log_plus_one"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_logistic"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_logistic(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.logistic_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.logistic_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.logistic"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_map"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_map(%arg0: tensor<16xf32>) -> tensor<16xf32> {
--  //      CHECK: "vhlo.map_v1"(%arg0) <{
-+  //      CHECK: "vhlo.map_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   dimensions = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>
-   // CHECK-SAME: }> ({
-   // CHECK-NEXT:   ^[[BB:bb.*]](%[[ARG1:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
-@@ -1479,57 +1588,65 @@
- }
- 
- // CHECK-LABEL: "op_maximum"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_maximum(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.maximum_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.maximum_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.maximum"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_minimum"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_minimum(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.minimum_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.minimum_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.minimum"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_multiply"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_multiply(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.multiply_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.multiply_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.multiply"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_negate"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_negate(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.negate_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.negate_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.negate"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_not"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_not(%arg0: tensor<i1>) -> tensor<i1> {
--  // CHECK: "vhlo.not_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-+  // CHECK: "vhlo.not_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-   %0 = "stablehlo.not"(%arg0) : (tensor<i1>) -> tensor<i1>
-   func.return %0 : tensor<i1>
- }
- 
- // CHECK-LABEL: "op_optimization_barrier"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_optimization_barrier(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.optimization_barrier_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.optimization_barrier_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.optimization_barrier"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_or"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_or(%arg0: tensor<i1>, %arg1: tensor<i1>) -> tensor<i1> {
--  // CHECK: "vhlo.or_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-+  // CHECK: "vhlo.or_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-   %0 = "stablehlo.or"(%arg0, %arg1) : (tensor<i1>, tensor<i1>) -> tensor<i1>
-   func.return %0 : tensor<i1>
- }
- 
- // CHECK-LABEL: "op_outfeed"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_outfeed(%arg0: tensor<f32>, %arg1: !stablehlo.token) -> !stablehlo.token {
--  //      CHECK: "vhlo.outfeed_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.outfeed_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   outfeed_config = #vhlo.string_v1<"foo">
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.token_v1) -> !vhlo.token_v1
-   %0 = "stablehlo.outfeed"(%arg0, %arg1) {
-@@ -1539,8 +1656,9 @@
- }
- 
- // CHECK-LABEL: "op_pad"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_pad(%arg0: tensor<8xf32>, %arg1: tensor<f32>) -> tensor<16xf32> {
--  //      CHECK: "vhlo.pad_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.pad_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   edge_padding_high = #vhlo.tensor_v1<dense<4> : tensor<1xi64>>,
-   // CHECK-SAME:   edge_padding_low = #vhlo.tensor_v1<dense<4> : tensor<1xi64>>,
-   // CHECK-SAME:   interior_padding = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>
-@@ -1554,36 +1672,41 @@
- }
- 
- // CHECK-LABEL: "op_popcnt"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_popcnt(%arg0: tensor<i32>) -> tensor<i32> {
--  // CHECK: "vhlo.popcnt_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-+  // CHECK: "vhlo.popcnt_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-   %0 = "stablehlo.popcnt"(%arg0) : (tensor<i32>) -> tensor<i32>
-   func.return %0 : tensor<i32>
- }
- 
- // CHECK-LABEL: "op_power"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_power(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.power_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.power_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.power"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_real_dynamic_slice"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}}, %[[ARG3:.*]]: {{.*}})
- func.func @op_real_dynamic_slice(%arg0: tensor<?xf32>, %arg1: tensor<1xindex>, %arg2: tensor<1xindex>, %arg3: tensor<1xindex>) -> tensor<?xf32> {
--  // CHECK: "vhlo.real_dynamic_slice_v1"(%arg0, %arg1, %arg2, %arg3) : (!vhlo.tensor_v1<?x!vhlo.f32_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>) -> !vhlo.tensor_v1<?x!vhlo.f32_v1>
-+  // CHECK: "vhlo.real_dynamic_slice_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]], %[[ARG3]]) : (!vhlo.tensor_v1<?x!vhlo.f32_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>) -> !vhlo.tensor_v1<?x!vhlo.f32_v1>
-   %0 = "stablehlo.real_dynamic_slice"(%arg0, %arg1, %arg2, %arg3) : (tensor<?xf32>, tensor<1xindex>, tensor<1xindex>, tensor<1xindex>) -> tensor<?xf32>
-   func.return %0 : tensor<?xf32>
- }
- 
- // CHECK-LABEL: "op_real"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_real(%arg0: tensor<complex<f32>>) -> tensor<f32> {
--  // CHECK: "vhlo.real_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.real_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.real"(%arg0) : (tensor<complex<f32>>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_recv"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_recv(%arg0: !stablehlo.token) -> (tensor<f32>, !stablehlo.token) {
--  //      CHECK: "vhlo.recv_v1"(%arg0) <{
-+  //      CHECK: "vhlo.recv_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME:   channel_type = #vhlo.integer_v1<3 : i64>,
-   // CHECK-SAME:   is_host_transfer = #vhlo.bool_v1<true>
-@@ -1596,8 +1719,9 @@
- }
- 
- // CHECK-LABEL: "op_reduce"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_reduce(%arg0: tensor<16xf32>, %arg1: tensor<f32>) -> tensor<f32> {
--  //  CHECK: "vhlo.reduce_v1"(%arg0, %arg1)
-+  //  CHECK: "vhlo.reduce_v1"(%[[ARG0]], %[[ARG1]])
-   //  CHECK:   ^[[BB:bb.*]](%[[ARG1:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>, %[[ARG2:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
-   //  CHECK:     "vhlo.return_v1"(%[[VAL1:.*]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-   //  CHECK: }) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-@@ -1612,8 +1736,9 @@
- }
- 
- // CHECK-LABEL: "op_reduce_precision"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_reduce_precision(%arg0: tensor<f32>) -> tensor<f32> {
--  //      CHECK: "vhlo.reduce_precision_v1"(%arg0) <{
-+  //      CHECK: "vhlo.reduce_precision_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   exponent_bits = #vhlo.integer_v1<8 : i32>
-   // CHECK-SAME:   mantissa_bits = #vhlo.integer_v1<10 : i32>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-@@ -1625,8 +1750,9 @@
- }
- 
- // CHECK-LABEL: "op_reduce_scatter"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_reduce_scatter(%arg0: tensor<16xf32>) -> tensor<16xf32> {
--  //               CHECK: "vhlo.reduce_scatter_v1"(%arg0) <{
-+  //               CHECK: "vhlo.reduce_scatter_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0], [1]]> : tensor<2x1xi64>>,
-   //          CHECK-SAME:   scatter_dimension = #vhlo.integer_v1<0 : i64>
-@@ -1650,8 +1776,9 @@
- }
- 
- // CHECK-LABEL: "op_reduce_window"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_reduce_window(%arg0: tensor<2x17x31x7xf32>, %arg1: tensor<f32>) -> tensor<2x9x16x7xf32> {
--  //               CHECK: "vhlo.reduce_window_v1"(%arg0, %arg1) <{
-+  //               CHECK: "vhlo.reduce_window_v1"(%[[ARG0]], %[[ARG1]]) <{
-   //          CHECK-SAME:   base_dilations = #vhlo.tensor_v1<dense<[1, 2, 2, 1]> : tensor<4xi64>>,
-   // CHECK-SAME{LITERAL}:   padding = #vhlo.tensor_v1<dense<[[0, 0], [2, 0], [0, 2], [0, 0]]> : tensor<4x2xi64>>,
-   //          CHECK-SAME:   window_dilations = #vhlo.tensor_v1<dense<[1, 2, 2, 1]> : tensor<4xi64>>,
-@@ -1677,8 +1804,9 @@
- }
- 
- // CHECK-LABEL: "op_remainder"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_remainder(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.remainder_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.remainder_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.remainder"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
-@@ -1698,16 +1826,18 @@
- }
- 
- // CHECK-LABEL: "op_reshape"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_reshape(%arg0: tensor<16xf32>) -> tensor<4x4xf32> {
--  // CHECK: "vhlo.reshape_v1"(%arg0) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<4x4x!vhlo.f32_v1>
-+  // CHECK: "vhlo.reshape_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<4x4x!vhlo.f32_v1>
-   %0 = "stablehlo.reshape"(%arg0) : (tensor<16xf32>) -> tensor<4x4xf32>
-   func.return %0 : tensor<4x4xf32>
- }
- 
- // CHECK-LABEL: "op_return"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_return(%arg0: tensor<i32>, %arg1: tensor<f32>) -> tensor<f32> {
--  //      CHECK: "vhlo.case_v1"(%arg0) ({
--  // CHECK-NEXT:   "vhlo.return_v1"(%arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-+  //      CHECK: "vhlo.case_v1"(%[[ARG0]]) ({
-+  // CHECK-NEXT:   "vhlo.return_v1"(%[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-   // CHECK-NEXT: }) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.case"(%arg0) ({
-     "stablehlo.return"(%arg1) : (tensor<f32>) -> ()
-@@ -1716,8 +1846,9 @@
- }
- 
- // CHECK-LABEL: "op_reverse"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_reverse(%arg0: tensor<16xf32>) -> tensor<16xf32> {
--  //      CHECK: "vhlo.reverse_v1"(%arg0) <{
-+  //      CHECK: "vhlo.reverse_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   dimensions = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x!vhlo.f32_v1>
-   %0 = "stablehlo.reverse"(%arg0) {
-@@ -1727,8 +1858,9 @@
- }
- 
- // CHECK-LABEL: "op_rng_bit_generator"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_rng_bit_generator(%arg0: tensor<f32>) -> (tensor<f32>, tensor<f32>) {
--  //      CHECK: "vhlo.rng_bit_generator_v1"(%arg0) <{
-+  //      CHECK: "vhlo.rng_bit_generator_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   rng_algorithm = #vhlo<rng_algorithm_v1 PHILOX>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>)
-   %0:2 = "stablehlo.rng_bit_generator"(%arg0) {
-@@ -1738,8 +1870,9 @@
- }
- 
- // CHECK-LABEL: "op_rng"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_rng(%arg0: tensor<f32>, %arg1: tensor<f32>, %arg2: tensor<0xindex>) -> tensor<f32> {
--  //      CHECK: "vhlo.rng_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.rng_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   rng_distribution = #vhlo<rng_distribution_v1 NORMAL>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<0x!vhlo.index_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.rng"(%arg0, %arg1, %arg2) {
-@@ -1749,29 +1882,33 @@
- }
- 
- // CHECK-LABEL: "op_round_nearest_afz"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_round_nearest_afz(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.round_nearest_afz_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.round_nearest_afz_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.round_nearest_afz"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_round_nearest_even"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_round_nearest_even(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.round_nearest_even_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.round_nearest_even_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.round_nearest_even"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_rsqrt"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_rsqrt(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.rsqrt_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.rsqrt_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.rsqrt"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_scatter"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_scatter(%arg0: tensor<200x100x300xf32>, %arg1: tensor<10x2xi32>, %arg2: tensor<10x300xf32>) -> tensor<200x100x300xf32> {
--  //      CHECK: "vhlo.scatter_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.scatter_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   index_vector_dim = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   indices_are_sorted = #vhlo.bool_v1<true>,
-   // CHECK-SAME:   inserted_window_dims = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-@@ -1801,8 +1938,9 @@
- }
- 
- // CHECK-LABEL: "op_select_and_scatter"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_select_and_scatter(%arg0: tensor<10x24x24x64xf32>, %arg1: tensor<12x13x13x66xf32>, %arg2: tensor<f32>) -> tensor<10x24x24x64xf32> {
--  //      CHECK: "vhlo.select_and_scatter_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.select_and_scatter_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   padding = #vhlo.tensor_v1<dense<1> : tensor<4x2xi64>>,
-   // CHECK-SAME:   window_dimensions = #vhlo.tensor_v1<dense<[1, 2, 2, 1]> : tensor<4xi64>>,
-   // CHECK-SAME:   window_strides = #vhlo.tensor_v1<dense<[1, 2, 2, 1]> : tensor<4xi64>>
-@@ -1832,15 +1970,17 @@
- }
- 
- // CHECK-LABEL: "op_select"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_select(%arg0: tensor<i1>, %arg1: tensor<f32>, %arg2: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.select_v1"(%arg0, %arg1, %arg2) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.select_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.select"(%arg0, %arg1, %arg2) : (tensor<i1>, tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_send"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_send(%arg0: tensor<f32>, %arg1: !stablehlo.token) -> !stablehlo.token {
--  //      CHECK: "vhlo.send_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.send_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME:   channel_type = #vhlo.integer_v1<2 : i64>,
-   // CHECK-SAME:   is_host_transfer = #vhlo.bool_v1<true>
-@@ -1853,8 +1993,9 @@
- }
- 
- // CHECK-LABEL: "op_set_dimension_size"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_set_dimension_size(%arg0: tensor<?xf32>, %arg1: tensor<i32>) -> tensor<16xf32> {
--  //      CHECK: "vhlo.set_dimension_size_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.set_dimension_size_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   dimension = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<?x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<16x!vhlo.f32_v1>
-   %0 = "stablehlo.set_dimension_size"(%arg0, %arg1) {
-@@ -1864,43 +2005,49 @@
- }
- 
- // CHECK-LABEL: "op_shift_left"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_shift_left(%arg0: tensor<i32>, %arg1: tensor<i32>) -> tensor<i32> {
--  // CHECK: "vhlo.shift_left_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.i32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-+  // CHECK: "vhlo.shift_left_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-   %0 = "stablehlo.shift_left"(%arg0, %arg1) : (tensor<i32>, tensor<i32>) -> tensor<i32>
-   func.return %0 : tensor<i32>
- }
- 
- // CHECK-LABEL: "op_shift_right_arithmetic"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_shift_right_arithmetic(%arg0: tensor<i32>, %arg1: tensor<i32>) -> tensor<i32> {
--  // CHECK: "vhlo.shift_right_arithmetic_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.i32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-+  // CHECK: "vhlo.shift_right_arithmetic_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-   %0 = "stablehlo.shift_right_arithmetic"(%arg0, %arg1) : (tensor<i32>, tensor<i32>) -> tensor<i32>
-   func.return %0 : tensor<i32>
- }
- 
- // CHECK-LABEL: "op_shift_right_logical"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_shift_right_logical(%arg0: tensor<i32>, %arg1: tensor<i32>) -> tensor<i32> {
--  // CHECK: "vhlo.shift_right_logical_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.i32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-+  // CHECK: "vhlo.shift_right_logical_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-   %0 = "stablehlo.shift_right_logical"(%arg0, %arg1) : (tensor<i32>, tensor<i32>) -> tensor<i32>
-   func.return %0 : tensor<i32>
- }
- 
- // CHECK-LABEL: "op_sign"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_sign(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.sign_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.sign_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.sign"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_sine"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_sine(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.sine_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.sine_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.sine"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_slice"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_slice(%arg0: tensor<16xf32>) -> tensor<4xf32> {
--  //      CHECK: "vhlo.slice_v1"(%arg0) <{
-+  //      CHECK: "vhlo.slice_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   limit_indices = #vhlo.tensor_v1<dense<4> : tensor<1xi64>>,
-   // CHECK-SAME:   start_indices = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>,
-   // CHECK-SAME:   strides = #vhlo.tensor_v1<dense<1> : tensor<1xi64>>
-@@ -1914,8 +2061,9 @@
- }
- 
- // CHECK-LABEL: "op_sort"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_sort(%arg0: tensor<16xf32>) -> tensor<16xf32> {
--  //      CHECK: "vhlo.sort_v1"(%arg0) <{
-+  //      CHECK: "vhlo.sort_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   dimension = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME:   is_stable = #vhlo.bool_v1<true>
-   // CHECK-SAME: }> ({
-@@ -1935,29 +2083,33 @@
- }
- 
- // CHECK-LABEL: "op_sqrt"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_sqrt(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.sqrt_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.sqrt_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.sqrt"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_subtract"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_subtract(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.subtract_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.subtract_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.subtract"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_tanh"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_tanh(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.tanh_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.tanh_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.tanh"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_torch_index_select"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_torch_index_select(%arg0: tensor<5x1x5xf32>, %arg1: tensor<2xi32>) ->  tensor<2x1x5xf32> {
--  //      CHECK: "vhlo.torch_index_select_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.torch_index_select_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   batch_dims = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME:   dim = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<5x1x5x!vhlo.f32_v1>, !vhlo.tensor_v1<2x!vhlo.i32_v1>) -> !vhlo.tensor_v1<2x1x5x!vhlo.f32_v1>
-@@ -1969,8 +2121,9 @@
- }
- 
- // CHECK-LABEL: "op_trace"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_trace(%arg0: tensor<f32>) {
--  //      CHECK: "vhlo.trace_v1"(%arg0) <{
-+  //      CHECK: "vhlo.trace_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   tag = #vhlo.string_v1<"foo">
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-   "stablehlo.trace"(%arg0) {
-@@ -1980,8 +2133,9 @@
- }
- 
- // CHECK-LABEL: "op_transpose"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_transpose(%arg0: tensor<16x8xf32>) ->  tensor<8x16xf32> {
--  //      CHECK: "vhlo.transpose_v1"(%arg0) <{
-+  //      CHECK: "vhlo.transpose_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   permutation = #vhlo.tensor_v1<dense<[1, 0]> : tensor<2xi64>>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<16x8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<8x16x!vhlo.f32_v1>
-   %0 = "stablehlo.transpose"(%arg0) {
-@@ -1991,8 +2145,9 @@
- }
- 
- // CHECK-LABEL: "op_triangular_solve"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_triangular_solve(%arg0: tensor<16x16xf32>, %arg1: tensor<16x16xf32>) ->  tensor<16x16xf32> {
--  //      CHECK: "vhlo.triangular_solve_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.triangular_solve_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   left_side = #vhlo.bool_v1<true>,
-   // CHECK-SAME:   lower = #vhlo.bool_v1<true>,
-   // CHECK-SAME:   transpose_a = #vhlo<transpose_v1 NO_TRANSPOSE>,
-@@ -2008,15 +2163,17 @@
- }
- 
- // CHECK-LABEL: "op_tuple"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_tuple(%arg0: tensor<f32>) -> tuple<tensor<f32>> {
--  // CHECK: "vhlo.tuple_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tuple_v1<!vhlo.tensor_v1<!vhlo.f32_v1>>
-+  // CHECK: "vhlo.tuple_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tuple_v1<!vhlo.tensor_v1<!vhlo.f32_v1>>
-   %0 = "stablehlo.tuple"(%arg0) : (tensor<f32>) -> tuple<tensor<f32>>
-   func.return %0 : tuple<tensor<f32>>
- }
- 
- // CHECK-LABEL: "op_unary_einsum"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_unary_einsum(%arg0: tensor<8x16xf32>) -> tensor<8xf32> {
--  //      CHECK: "vhlo.unary_einsum_v1"(%arg0) <{
-+  //      CHECK: "vhlo.unary_einsum_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   einsum_config = #vhlo.string_v1<"ab->a">
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<8x16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<8x!vhlo.f32_v1>
-   %0 = "stablehlo.unary_einsum"(%arg0) {
-@@ -2026,22 +2183,25 @@
- }
- 
- // CHECK-LABEL: "op_uniform_dequantize"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_uniform_dequantize(%arg0: tensor<!quant.uniform<i8:f32, 34.0:16>>) -> tensor<f32> {
--  // CHECK: "vhlo.uniform_dequantize_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.uniform_dequantize_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.uniform_dequantize"(%arg0) : (tensor<!quant.uniform<i8:f32, 34.0:16>>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_uniform_quantize"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_uniform_quantize(%arg0: tensor<f32>) -> tensor<!quant.uniform<i8:f32, 34.0:16>> {
--  // CHECK: "vhlo.uniform_quantize_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>
-+  // CHECK: "vhlo.uniform_quantize_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>
-   %0 = "stablehlo.uniform_quantize"(%arg0) : (tensor<f32>) -> tensor<!quant.uniform<i8:f32, 34.0:16>>
-   func.return %0 : tensor<!quant.uniform<i8:f32, 34.0:16>>
- }
- 
- // CHECK-LABEL: "op_while"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_while(%arg0: tensor<i1>) -> tensor<i1> {
--  //      CHECK: "vhlo.while_v1"(%arg0) ({
-+  //      CHECK: "vhlo.while_v1"(%[[ARG0]]) ({
-   // CHECK-NEXT:   ^[[BB:bb.*]](%[[ARG1:arg.*]]: !vhlo.tensor_v1<!vhlo.bool_v1>):
-   // CHECK-NEXT:     "vhlo.return_v1"(%[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>) -> ()
-   // CHECK-NEXT:   }, {
-@@ -2059,8 +2219,9 @@
- }
- 
- // CHECK-LABEL: "op_xor"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_xor(%arg0: tensor<i1>, %arg1: tensor<i1>) -> tensor<i1> {
--  // CHECK: "vhlo.xor_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-+  // CHECK: "vhlo.xor_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-   %0 = "stablehlo.xor"(%arg0, %arg1) : (tensor<i1>, tensor<i1>) -> tensor<i1>
-   func.return %0 : tensor<i1>
- }
-@@ -2068,190 +2229,217 @@
- // ============ TYPES ============
- 
- // CHECK-LABEL: "type_i1"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_i1(%arg0: tensor<i1>, %arg1: tensor<i1>) -> tensor<i1> {
--  // CHECK: "vhlo.and_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-+  // CHECK: "vhlo.and_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-   %0 = "stablehlo.and"(%arg0, %arg1) : (tensor<i1>, tensor<i1>) -> tensor<i1>
-   func.return %0 : tensor<i1>
- }
- 
- // CHECK-LABEL: "type_i4"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_i4(%arg0: tensor<i4>, %arg1: tensor<i4>) -> tensor<i4> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.i4_v1>, !vhlo.tensor_v1<!vhlo.i4_v1>) -> !vhlo.tensor_v1<!vhlo.i4_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i4_v1>, !vhlo.tensor_v1<!vhlo.i4_v1>) -> !vhlo.tensor_v1<!vhlo.i4_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<i4>, tensor<i4>) -> tensor<i4>
-   func.return %0 : tensor<i4>
- }
- 
- // CHECK-LABEL: "type_i8"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_i8(%arg0: tensor<i8>, %arg1: tensor<i8>) -> tensor<i8> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.i8_v1>, !vhlo.tensor_v1<!vhlo.i8_v1>) -> !vhlo.tensor_v1<!vhlo.i8_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i8_v1>, !vhlo.tensor_v1<!vhlo.i8_v1>) -> !vhlo.tensor_v1<!vhlo.i8_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<i8>, tensor<i8>) -> tensor<i8>
-   func.return %0 : tensor<i8>
- }
- 
- // CHECK-LABEL: "type_i16"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_i16(%arg0: tensor<i16>, %arg1: tensor<i16>) -> tensor<i16> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.i16_v1>, !vhlo.tensor_v1<!vhlo.i16_v1>) -> !vhlo.tensor_v1<!vhlo.i16_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i16_v1>, !vhlo.tensor_v1<!vhlo.i16_v1>) -> !vhlo.tensor_v1<!vhlo.i16_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<i16>, tensor<i16>) -> tensor<i16>
-   func.return %0 : tensor<i16>
- }
- 
- // CHECK-LABEL: "type_i32"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_i32(%arg0: tensor<i32>, %arg1: tensor<i32>) -> tensor<i32> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.i32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<i32>, tensor<i32>) -> tensor<i32>
-   func.return %0 : tensor<i32>
- }
- 
- // CHECK-LABEL: "type_i64"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_i64(%arg0: tensor<i64>, %arg1: tensor<i64>) -> tensor<i64> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.i64_v1>, !vhlo.tensor_v1<!vhlo.i64_v1>) -> !vhlo.tensor_v1<!vhlo.i64_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i64_v1>, !vhlo.tensor_v1<!vhlo.i64_v1>) -> !vhlo.tensor_v1<!vhlo.i64_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<i64>, tensor<i64>) -> tensor<i64>
-   func.return %0 : tensor<i64>
- }
- 
- // CHECK-LABEL: "type_ui4"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_ui4(%arg0: tensor<ui4>, %arg1: tensor<ui4>) -> tensor<ui4> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.ui4_v1>, !vhlo.tensor_v1<!vhlo.ui4_v1>) -> !vhlo.tensor_v1<!vhlo.ui4_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.ui4_v1>, !vhlo.tensor_v1<!vhlo.ui4_v1>) -> !vhlo.tensor_v1<!vhlo.ui4_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<ui4>, tensor<ui4>) -> tensor<ui4>
-   func.return %0 : tensor<ui4>
- }
- 
- // CHECK-LABEL: "type_ui8"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_ui8(%arg0: tensor<ui8>, %arg1: tensor<ui8>) -> tensor<ui8> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.ui8_v1>, !vhlo.tensor_v1<!vhlo.ui8_v1>) -> !vhlo.tensor_v1<!vhlo.ui8_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.ui8_v1>, !vhlo.tensor_v1<!vhlo.ui8_v1>) -> !vhlo.tensor_v1<!vhlo.ui8_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<ui8>, tensor<ui8>) -> tensor<ui8>
-   func.return %0 : tensor<ui8>
- }
- 
- // CHECK-LABEL: "type_ui16"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_ui16(%arg0: tensor<ui16>, %arg1: tensor<ui16>) -> tensor<ui16> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.ui16_v1>, !vhlo.tensor_v1<!vhlo.ui16_v1>) -> !vhlo.tensor_v1<!vhlo.ui16_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.ui16_v1>, !vhlo.tensor_v1<!vhlo.ui16_v1>) -> !vhlo.tensor_v1<!vhlo.ui16_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<ui16>, tensor<ui16>) -> tensor<ui16>
-   func.return %0 : tensor<ui16>
- }
- 
- // CHECK-LABEL: "type_ui32"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_ui32(%arg0: tensor<ui32>, %arg1: tensor<ui32>) -> tensor<ui32> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.ui32_v1>, !vhlo.tensor_v1<!vhlo.ui32_v1>) -> !vhlo.tensor_v1<!vhlo.ui32_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.ui32_v1>, !vhlo.tensor_v1<!vhlo.ui32_v1>) -> !vhlo.tensor_v1<!vhlo.ui32_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<ui32>, tensor<ui32>) -> tensor<ui32>
-   func.return %0 : tensor<ui32>
- }
- 
- // CHECK-LABEL: "type_ui64"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_ui64(%arg0: tensor<ui64>, %arg1: tensor<ui64>) -> tensor<ui64> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.ui64_v1>, !vhlo.tensor_v1<!vhlo.ui64_v1>) -> !vhlo.tensor_v1<!vhlo.ui64_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.ui64_v1>, !vhlo.tensor_v1<!vhlo.ui64_v1>) -> !vhlo.tensor_v1<!vhlo.ui64_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<ui64>, tensor<ui64>) -> tensor<ui64>
-   func.return %0 : tensor<ui64>
- }
- 
- // CHECK-LABEL: "type_f8E4M3FN"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_f8E4M3FN(%arg0: tensor<f8E4M3FN>, %arg1: tensor<f8E4M3FN>) -> tensor<f8E4M3FN> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f8E4M3FN_v1>, !vhlo.tensor_v1<!vhlo.f8E4M3FN_v1>) -> !vhlo.tensor_v1<!vhlo.f8E4M3FN_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f8E4M3FN_v1>, !vhlo.tensor_v1<!vhlo.f8E4M3FN_v1>) -> !vhlo.tensor_v1<!vhlo.f8E4M3FN_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f8E4M3FN>, tensor<f8E4M3FN>) -> tensor<f8E4M3FN>
-   func.return %0 : tensor<f8E4M3FN>
- }
- 
- // CHECK-LABEL: "type_f8E5M2"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_f8E5M2(%arg0: tensor<f8E5M2>, %arg1: tensor<f8E5M2>) -> tensor<f8E5M2> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f8E5M2_v1>, !vhlo.tensor_v1<!vhlo.f8E5M2_v1>) -> !vhlo.tensor_v1<!vhlo.f8E5M2_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f8E5M2_v1>, !vhlo.tensor_v1<!vhlo.f8E5M2_v1>) -> !vhlo.tensor_v1<!vhlo.f8E5M2_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f8E5M2>, tensor<f8E5M2>) -> tensor<f8E5M2>
-   func.return %0 : tensor<f8E5M2>
- }
- 
- // CHECK-LABEL: "type_f8E4M3FNUZ"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_f8E4M3FNUZ(%arg0: tensor<f8E4M3FNUZ>, %arg1: tensor<f8E4M3FNUZ>) -> tensor<f8E4M3FNUZ> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f8E4M3FNUZ_v1>, !vhlo.tensor_v1<!vhlo.f8E4M3FNUZ_v1>) -> !vhlo.tensor_v1<!vhlo.f8E4M3FNUZ_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f8E4M3FNUZ_v1>, !vhlo.tensor_v1<!vhlo.f8E4M3FNUZ_v1>) -> !vhlo.tensor_v1<!vhlo.f8E4M3FNUZ_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f8E4M3FNUZ>, tensor<f8E4M3FNUZ>) -> tensor<f8E4M3FNUZ>
-   func.return %0 : tensor<f8E4M3FNUZ>
- }
- 
- // CHECK-LABEL: "type_f8E4M3B11FNUZ"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_f8E4M3B11FNUZ(%arg0: tensor<f8E4M3B11FNUZ>, %arg1: tensor<f8E4M3B11FNUZ>) -> tensor<f8E4M3B11FNUZ> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f8E4M3B11FNUZ_v1>, !vhlo.tensor_v1<!vhlo.f8E4M3B11FNUZ_v1>) -> !vhlo.tensor_v1<!vhlo.f8E4M3B11FNUZ_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f8E4M3B11FNUZ_v1>, !vhlo.tensor_v1<!vhlo.f8E4M3B11FNUZ_v1>) -> !vhlo.tensor_v1<!vhlo.f8E4M3B11FNUZ_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f8E4M3B11FNUZ>, tensor<f8E4M3B11FNUZ>) -> tensor<f8E4M3B11FNUZ>
-   func.return %0 : tensor<f8E4M3B11FNUZ>
- }
- 
- // CHECK-LABEL: "type_f8E5M2FNUZ"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_f8E5M2FNUZ(%arg0: tensor<f8E5M2FNUZ>, %arg1: tensor<f8E5M2FNUZ>) -> tensor<f8E5M2FNUZ> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f8E5M2FNUZ_v1>, !vhlo.tensor_v1<!vhlo.f8E5M2FNUZ_v1>) -> !vhlo.tensor_v1<!vhlo.f8E5M2FNUZ_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f8E5M2FNUZ_v1>, !vhlo.tensor_v1<!vhlo.f8E5M2FNUZ_v1>) -> !vhlo.tensor_v1<!vhlo.f8E5M2FNUZ_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f8E5M2FNUZ>, tensor<f8E5M2FNUZ>) -> tensor<f8E5M2FNUZ>
-   func.return %0 : tensor<f8E5M2FNUZ>
- }
- 
- // CHECK-LABEL: "type_bf16"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_bf16(%arg0: tensor<bf16>, %arg1: tensor<bf16>) -> tensor<bf16> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.bf16_v1>, !vhlo.tensor_v1<!vhlo.bf16_v1>) -> !vhlo.tensor_v1<!vhlo.bf16_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.bf16_v1>, !vhlo.tensor_v1<!vhlo.bf16_v1>) -> !vhlo.tensor_v1<!vhlo.bf16_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<bf16>, tensor<bf16>) -> tensor<bf16>
-   func.return %0 : tensor<bf16>
- }
- 
- // CHECK-LABEL: "type_f16"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_f16(%arg0: tensor<f16>, %arg1: tensor<f16>) -> tensor<f16> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f16_v1>, !vhlo.tensor_v1<!vhlo.f16_v1>) -> !vhlo.tensor_v1<!vhlo.f16_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f16_v1>, !vhlo.tensor_v1<!vhlo.f16_v1>) -> !vhlo.tensor_v1<!vhlo.f16_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f16>, tensor<f16>) -> tensor<f16>
-   func.return %0 : tensor<f16>
- }
- 
- // CHECK-LABEL: "type_f32"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_f32(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "type_f64"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_f64(%arg0: tensor<f64>, %arg1: tensor<f64>) -> tensor<f64> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f64_v1>, !vhlo.tensor_v1<!vhlo.f64_v1>) -> !vhlo.tensor_v1<!vhlo.f64_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f64_v1>, !vhlo.tensor_v1<!vhlo.f64_v1>) -> !vhlo.tensor_v1<!vhlo.f64_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f64>, tensor<f64>) -> tensor<f64>
-   func.return %0 : tensor<f64>
- }
- 
- // CHECK-LABEL: "type_complex_f32"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_complex_f32(%arg0: tensor<complex<f32>>, %arg1: tensor<complex<f32>>) -> tensor<complex<f32>> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>, !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>) -> !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>, !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>) -> !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<complex<f32>>, tensor<complex<f32>>) -> tensor<complex<f32>>
-   func.return %0 : tensor<complex<f32>>
- }
- 
- // CHECK-LABEL: "type_complex_f64"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_complex_f64(%arg0: tensor<complex<f64>>, %arg1: tensor<complex<f64>>) -> tensor<complex<f64>> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f64_v1>>, !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f64_v1>>) -> !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f64_v1>>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f64_v1>>, !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f64_v1>>) -> !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f64_v1>>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<complex<f64>>, tensor<complex<f64>>) -> tensor<complex<f64>>
-   func.return %0 : tensor<complex<f64>>
- }
- 
- // CHECK-LABEL: "type_dynamism_ranked"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @type_dynamism_ranked(%arg0: tensor<?xf32>) -> tensor<?xf32> {
--  // CHECK: "vhlo.abs_v1"(%arg0) : (!vhlo.tensor_v1<?x!vhlo.f32_v1>) -> !vhlo.tensor_v1<?x!vhlo.f32_v1>
-+  // CHECK: "vhlo.abs_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<?x!vhlo.f32_v1>) -> !vhlo.tensor_v1<?x!vhlo.f32_v1>
-   %0 = "stablehlo.abs"(%arg0) : (tensor<?xf32>) -> tensor<?xf32>
-   func.return %0 : tensor<?xf32>
- }
- 
- // CHECK-LABEL: "type_quantization"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_quantization(%arg0: tensor<!quant.uniform<i8:f32, 34.0:16>>, %arg1: tensor<!quant.uniform<i8:f32, 34.0:16>>) -> tensor<!quant.uniform<i8:f32, 34.0:16>> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>, !vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>) -> !vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>, !vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>) -> !vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<!quant.uniform<i8:f32, 34.0:16>>, tensor<!quant.uniform<i8:f32, 34.0:16>>) -> tensor<!quant.uniform<i8:f32, 34.0:16>>
-   func.return %0 : tensor<!quant.uniform<i8:f32, 34.0:16>>
- }
- 
- //       CHECK: function_type = #vhlo.type_v1<!vhlo.func_v1<(!vhlo.token_v1) -> !vhlo.token_v1>>
- // CHECK-LABEL: "type_token_callee"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @type_token_callee(%arg0: !stablehlo.token) -> !stablehlo.token {
--  // CHECK: "vhlo.return_v1"(%arg0) : (!vhlo.token_v1) -> ()
-+  // CHECK: "vhlo.return_v1"(%[[ARG0]]) : (!vhlo.token_v1) -> ()
-   return %arg0 : !stablehlo.token
- }
- 
- //       CHECK: function_type = #vhlo.type_v1<!vhlo.func_v1<(!vhlo.token_v1) -> !vhlo.token_v1>>
- // CHECK-LABEL: "type_token_caller"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @type_token_caller(%arg0: !stablehlo.token) -> !stablehlo.token {
--  // CHECK:      "vhlo.call_v1"(%arg0) <{callee = #vhlo.string_v1<"type_token_callee">}
-+  // CHECK:      "vhlo.call_v1"(%[[ARG0]]) <{callee = #vhlo.string_v1<"type_token_callee">}
-   // CHECK-SAME: (!vhlo.token_v1) -> !vhlo.token_v1
-   %0 = func.call @type_token_callee(%arg0) : (!stablehlo.token) -> !stablehlo.token
-   return %0 : !stablehlo.token
- }
- 
- // CHECK-LABEL: "type_tuple"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @type_tuple(%arg0: tuple<tensor<f32>>) -> tuple<!stablehlo.token> {
-   %0 = "stablehlo.custom_call"(%arg0) {
-     call_target_name = "foo"
-diff --ruN a/stablehlo/stablehlo/tests/vhlo/stablehlo_legalize_to_vhlo.0_16_0.mlir b/stablehlo/stablehlo/tests/vhlo/stablehlo_legalize_to_vhlo.0_16_0.mlir
---- stablehlo/stablehlo/tests/vhlo/stablehlo_legalize_to_vhlo.0_16_0.mlir
-+++ stablehlo/stablehlo/tests/vhlo/stablehlo_legalize_to_vhlo.0_16_0.mlir
-@@ -13,6 +13,7 @@
- // ============ ATTRIBUTES ============
- 
- // CHECK-LABEL: "attr_comparison_direction_eq"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_direction_eq(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     // CHECK: comparison_direction = #vhlo<comparison_direction_v1 EQ>
-@@ -22,6 +23,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_direction_ne"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_direction_ne(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     // CHECK: comparison_direction = #vhlo<comparison_direction_v1 NE>
-@@ -31,6 +33,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_direction_ge"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_direction_ge(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     // CHECK: comparison_direction = #vhlo<comparison_direction_v1 GE>
-@@ -40,6 +43,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_direction_gt"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_direction_gt(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     // CHECK: comparison_direction = #vhlo<comparison_direction_v1 GT>
-@@ -49,6 +53,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_direction_le"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_direction_le(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     // CHECK: comparison_direction = #vhlo<comparison_direction_v1 LE>
-@@ -58,6 +63,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_direction_lt"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_direction_lt(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     // CHECK: comparison_direction = #vhlo<comparison_direction_v1 LT>
-@@ -67,6 +73,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_type_notype"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_type_notype(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     comparison_direction = #stablehlo<comparison_direction EQ>
-@@ -76,6 +83,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_type_float"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_type_float(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     comparison_direction = #stablehlo<comparison_direction EQ>,
-@@ -86,6 +94,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_type_totalorder"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_type_totalorder(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     comparison_direction = #stablehlo<comparison_direction EQ>,
-@@ -96,6 +105,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_type_signed"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_type_signed(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     comparison_direction = #stablehlo<comparison_direction EQ>,
-@@ -106,6 +116,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_type_unsigned"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_type_unsigned(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     comparison_direction = #stablehlo<comparison_direction EQ>,
-@@ -118,6 +129,7 @@
- // ConvDimensionNumbers aka #stablehlo.conv is covered below.
- 
- // CHECK-LABEL: "attr_custom_call_api_version_unspecified"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_custom_call_api_version_unspecified(%arg0: tensor<f32>) -> tensor<f32> {
-   %0 = "stablehlo.custom_call"(%arg0) {
-     call_target_name = "foo",
-@@ -128,6 +140,7 @@
- }
- 
- // CHECK-LABEL: "attr_custom_call_api_version_original"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_custom_call_api_version_original(%arg0: tensor<f32>) -> tensor<f32> {
-   %0 = "stablehlo.custom_call"(%arg0) {
-     call_target_name = "foo",
-@@ -138,6 +151,7 @@
- }
- 
- // CHECK-LABEL: "attr_custom_call_api_version_status_returning"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_custom_call_api_version_status_returning(%arg0: tensor<f32>) -> tensor<f32> {
-   %0 = "stablehlo.custom_call"(%arg0) {
-     call_target_name = "foo",
-@@ -148,6 +162,7 @@
- }
- 
- // CHECK-LABEL: "attr_custom_call_api_version_status_returning_unified"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_custom_call_api_version_status_returning_unified(%arg0: tensor<f32>) -> tensor<f32> {
-   %0 = "stablehlo.custom_call"(%arg0) {
-     call_target_name = "foo",
-@@ -166,6 +181,7 @@
- // DotDimensionNumbers aka #stablehlo.dot is covered below.
- 
- // CHECK-LABEL: "attr_fft_type_fft"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_fft_type_fft(%arg0: tensor<16xcomplex<f32>>) -> tensor<16xcomplex<f32>> {
-   %0 = "stablehlo.fft"(%arg0) {
-     // CHECK: fft_type = #vhlo<fft_type_v1 FFT>
-@@ -176,6 +192,7 @@
- }
- 
- // CHECK-LABEL: "attr_fft_type_ifft"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_fft_type_ifft(%arg0: tensor<16xcomplex<f32>>) -> tensor<16xcomplex<f32>> {
-   %0 = "stablehlo.fft"(%arg0) {
-     // CHECK: fft_type = #vhlo<fft_type_v1 IFFT>
-@@ -186,6 +203,7 @@
- }
- 
- // CHECK-LABEL: "attr_fft_type_rfft"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_fft_type_rfft(%arg0: tensor<16xf32>) -> tensor<9xcomplex<f32>> {
-   %0 = "stablehlo.fft"(%arg0) {
-     // CHECK: fft_type = #vhlo<fft_type_v1 RFFT>
-@@ -196,6 +214,7 @@
- }
- 
- // CHECK-LABEL: "attr_fft_type_irfft"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_fft_type_irfft(%arg0: tensor<9xcomplex<f32>>) -> tensor<16xf32> {
-   %0 = "stablehlo.fft"(%arg0) {
-     // CHECK: fft_type = #vhlo<fft_type_v1 IRFFT>
-@@ -208,6 +227,7 @@
- // GatherDimensionNumbers aka #stablehlo.gather is covered below.
- 
- // CHECK-LABEL: "attr_precision_config_default"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_precision_config_default(%arg0: tensor<8x16xf32>, %arg1: tensor<16x8xf32>) -> tensor<8x8xf32> {
-   %0 = "stablehlo.dot"(%arg0, %arg1) {
-     // CHECK: precision_config = #vhlo.array_v1<[#vhlo<precision_v1 DEFAULT>, #vhlo<precision_v1 DEFAULT>]>
-@@ -216,6 +236,7 @@
- }
- 
- // CHECK-LABEL: "attr_precision_config_high"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_precision_config_high(%arg0: tensor<8x16xf32>, %arg1: tensor<16x8xf32>) -> tensor<8x8xf32> {
-   %0 = "stablehlo.dot"(%arg0, %arg1) {
-     // CHECK: precision_config = #vhlo.array_v1<[#vhlo<precision_v1 HIGH>, #vhlo<precision_v1 HIGH>]>
-@@ -225,6 +246,7 @@
- }
- 
- // CHECK-LABEL: "attr_precision_config_highest"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_precision_config_highest(%arg0: tensor<8x16xf32>, %arg1: tensor<16x8xf32>) -> tensor<8x8xf32> {
-   %0 = "stablehlo.dot"(%arg0, %arg1) {
-     // CHECK: precision_config = #vhlo.array_v1<[#vhlo<precision_v1 HIGHEST>, #vhlo<precision_v1 HIGHEST>]>
-@@ -234,6 +256,7 @@
- }
- 
- // CHECK-LABEL: "attr_rng_algorithm_default"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_rng_algorithm_default(%arg0: tensor<f32>) -> (tensor<f32>, tensor<f32>) {
-   %0:2 = "stablehlo.rng_bit_generator"(%arg0) {
-     // CHECK: rng_algorithm = #vhlo<rng_algorithm_v1 DEFAULT>
-@@ -243,6 +266,7 @@
- }
- 
- // CHECK-LABEL: "attr_rng_algorithm_three_fry"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_rng_algorithm_three_fry(%arg0: tensor<f32>) -> (tensor<f32>, tensor<f32>) {
-   %0:2 = "stablehlo.rng_bit_generator"(%arg0) {
-     // CHECK: rng_algorithm = #vhlo<rng_algorithm_v1 THREE_FRY>
-@@ -252,6 +276,7 @@
- }
- 
- // CHECK-LABEL: "attr_rng_algorithm_philox"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_rng_algorithm_philox(%arg0: tensor<f32>) -> (tensor<f32>, tensor<f32>) {
-   %0:2 = "stablehlo.rng_bit_generator"(%arg0) {
-     // CHECK: rng_algorithm = #vhlo<rng_algorithm_v1 PHILOX>
-@@ -261,6 +286,7 @@
- }
- 
- // CHECK-LABEL: "attr_rng_distribution_uniform"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @attr_rng_distribution_uniform(%arg0: tensor<f32>, %arg1: tensor<f32>, %arg2: tensor<0xindex>) -> tensor<f32> {
-   %0 = "stablehlo.rng"(%arg0, %arg1, %arg2) {
-     // CHECK: rng_distribution = #vhlo<rng_distribution_v1 UNIFORM>
-@@ -270,6 +296,7 @@
- }
- 
- // CHECK-LABEL: "attr_rng_distribution_normal"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @attr_rng_distribution_normal(%arg0: tensor<f32>, %arg1: tensor<f32>, %arg2: tensor<0xindex>) -> tensor<f32> {
-   %0 = "stablehlo.rng"(%arg0, %arg1, %arg2) {
-     // CHECK: rng_distribution = #vhlo<rng_distribution_v1 NORMAL>
-@@ -281,6 +308,7 @@
- // ScatterDimensionNumbers aka #stablehlo.scatter is covered below.
- 
- // CHECK-LABEL: "attr_transpose_no_transpose"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_transpose_no_transpose(%arg0: tensor<16x16xf32>, %arg1: tensor<16x16xf32>) ->  tensor<16x16xf32> {
-   %0 = "stablehlo.triangular_solve"(%arg0, %arg1) {
-     left_side = true,
-@@ -293,6 +321,7 @@
- }
- 
- // CHECK-LABEL: "attr_transpose_transpose"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_transpose_transpose(%arg0: tensor<16x16xf32>, %arg1: tensor<16x16xf32>) ->  tensor<16x16xf32> {
-   %0 = "stablehlo.triangular_solve"(%arg0, %arg1) {
-     left_side = true,
-@@ -305,6 +334,7 @@
- }
- 
- // CHECK-LABEL: "attr_transpose_adjoint"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_transpose_adjoint(%arg0: tensor<16x16xf32>, %arg1: tensor<16x16xf32>) ->  tensor<16x16xf32> {
-   %0 = "stablehlo.triangular_solve"(%arg0, %arg1) {
-     left_side = true,
-@@ -319,10 +349,9 @@
- // TypeExtensionsAttr aka #stablehlo.type_extensions is covered below.
- 
- // CHECK-LABEL: "attr_type_extensions_bounds"
--func.func @attr_type_extensions_bounds(
--    %arg0: tensor<?x?xf32, #stablehlo.type_extensions<bounds = [16, ?]>>)
--    -> tensor<?x?xf32, #stablehlo.type_extensions<bounds = [16, ?]>> {
--  // CHECK: "vhlo.return_v1"(%arg0) : (!vhlo.tensor_v1<?x?x!vhlo.f32_v1, #vhlo.type_extensions_v1<bounds = [16, ?]>>) -> ()
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @attr_type_extensions_bounds(%arg0: tensor<?x?xf32, #stablehlo.type_extensions<bounds = [16, ?]>>) -> tensor<?x?xf32, #stablehlo.type_extensions<bounds = [16, ?]>> {
-+  // CHECK: "vhlo.return_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<?x?x!vhlo.f32_v1, #vhlo.type_extensions_v1<bounds = [16, ?]>>) -> ()
-   func.return %arg0 : tensor<?x?xf32, #stablehlo.type_extensions<bounds = [16, ?]>>
- }
- 
-@@ -330,8 +359,9 @@
- // ============ DEFAULTS ============
- 
- // CHECK-LABEL: "default_all_gather"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_all_gather(%arg0: tensor<16x8xf32>) -> tensor<16x16xf32> {
--  //               CHECK: "vhlo.all_gather_v1"(%arg0) <{
-+  //               CHECK: "vhlo.all_gather_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   all_gather_dim = #vhlo.integer_v1<1 : i64>
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0], [1]]> : tensor<2x1xi64>>,
-@@ -345,8 +375,9 @@
- }
- 
- // CHECK-LABEL: "default_all_reduce"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_all_reduce(%arg0: tensor<f32>) -> tensor<f32> {
--  //               CHECK: "vhlo.all_reduce_v1"(%arg0)
-+  //               CHECK: "vhlo.all_reduce_v1"(%[[ARG0]])
-   //          CHECK-SAME: <{
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0], [1]]> : tensor<2x1xi64>>,
-@@ -368,8 +399,9 @@
- }
- 
- // CHECK-LABEL: "default_all_to_all"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_all_to_all(%arg0: tensor<4x16xf32>) -> tensor<16x4xf32> {
--  //               CHECK: "vhlo.all_to_all_v1"(%arg0) <{
-+  //               CHECK: "vhlo.all_to_all_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   //          CHECK-SAME:   concat_dimension = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0, 1, 2, 3]]> : tensor<1x4xi64>>,
-@@ -386,8 +418,9 @@
- }
- 
- // CHECK-LABEL: "default_cholesky"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_cholesky(%arg0: tensor<1x16x16xf32>) -> tensor<1x16x16xf32> {
--  //      CHECK: "vhlo.cholesky_v1"(%arg0) <{
-+  //      CHECK: "vhlo.cholesky_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   lower = #vhlo.bool_v1<false>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<1x16x16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<1x16x16x!vhlo.f32_v1>
-   %0 = "stablehlo.cholesky"(%arg0) : (tensor<1x16x16xf32>) -> tensor<1x16x16xf32>
-@@ -395,8 +428,9 @@
- }
- 
- // CHECK-LABEL: "default_collective_permute"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_collective_permute(%arg0: tensor<16x8xf32>) -> tensor<16x8xf32> {
--  //               CHECK: "vhlo.collective_permute_v1"(%arg0) <{
-+  //               CHECK: "vhlo.collective_permute_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME{LITERAL}:   source_target_pairs = #vhlo.tensor_v1<dense<[[0, 1], [1, 2], [2, 3]]> : tensor<3x2xi64>>
-   //          CHECK-SAME: }> : (!vhlo.tensor_v1<16x8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x8x!vhlo.f32_v1>
-@@ -407,8 +441,9 @@
- }
- 
- // CHECK-LABEL: "default_collective_broadcast"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_collective_broadcast(%arg0: tensor<16x8xf32>) -> tensor<16x8xf32> {
--  //               CHECK: "vhlo.collective_broadcast_v1"(%arg0) <{
-+  //               CHECK: "vhlo.collective_broadcast_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0, 1]]> : tensor<1x2xi64>>
-   //          CHECK-SAME: }> : (!vhlo.tensor_v1<16x8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x8x!vhlo.f32_v1>
-@@ -419,8 +454,9 @@
- }
- 
- // CHECK-LABEL: "default_compare"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @default_compare(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
--  //      CHECK: "vhlo.compare_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.compare_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   compare_type = #vhlo<comparison_type_v1 NOTYPE>,
-   // CHECK-SAME:   comparison_direction = #vhlo<comparison_direction_v1 EQ>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-@@ -431,8 +467,9 @@
- }
- 
- // CHECK-LABEL: "default_convolution"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @default_convolution(%arg0: tensor<1x8x8x207xf32>, %arg1: tensor<3x3x207x16xf32>) -> tensor<1x6x6x16xf32> {
--  //      CHECK: "vhlo.convolution_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.convolution_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   batch_group_count = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   feature_group_count = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   input_batch_dimension = #vhlo.integer_v1<0 : i64>,
-@@ -460,8 +497,9 @@
- }
- 
- // CHECK-LABEL: "default_custom_call"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_custom_call(%arg0: tensor<f32>) -> tensor<f32> {
--  //      CHECK: "vhlo.custom_call_v1"(%arg0) <{
-+  //      CHECK: "vhlo.custom_call_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   api_version = #vhlo<api_version_v1 API_VERSION_ORIGINAL>,
-   // CHECK-SAME:   backend_config = #vhlo.string_v1<"">,
-   // CHECK-SAME:   call_target_name = #vhlo.string_v1<"foo">,
-@@ -478,8 +516,9 @@
- }
- 
- // CHECK-LABEL: "default_dot_general"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @default_dot_general(%arg0: tensor<8x8x16xf32>, %arg1: tensor<8x16x8xf32>) -> tensor<8x8x8xf32> {
--  //      CHECK: "vhlo.dot_general_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.dot_general_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   lhs_batching_dimensions = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>,
-   // CHECK-SAME:   lhs_contracting_dimensions = #vhlo.tensor_v1<dense<2> : tensor<1xi64>>,
-   // CHECK-SAME:   precision_config = #vhlo.array_v1<[#vhlo<precision_v1 DEFAULT>, #vhlo<precision_v1 DEFAULT>]>,
-@@ -498,8 +537,9 @@
- }
- 
- // CHECK-LABEL: "default_dot"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @default_dot(%arg0: tensor<8x16xf32>, %arg1: tensor<16x8xf32>) -> tensor<8x8xf32> {
--  //      CHECK: "vhlo.dot_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.dot_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   precision_config = #vhlo.array_v1<[#vhlo<precision_v1 DEFAULT>, #vhlo<precision_v1 DEFAULT>]>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<8x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<8x8x!vhlo.f32_v1>
-   %0 = "stablehlo.dot"(%arg0, %arg1) : (tensor<8x16xf32>, tensor<16x8xf32>) -> tensor<8x8xf32>
-@@ -507,8 +547,9 @@
- }
- 
- // CHECK-LABEL: "default_dynamic_broadcast_in_dim"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @default_dynamic_broadcast_in_dim(%arg0: tensor<?x?xf32>, %arg1: tensor<2xindex>) -> tensor<?x?xf32> {
--  //      CHECK: "vhlo.dynamic_broadcast_in_dim_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.dynamic_broadcast_in_dim_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   broadcast_dimensions = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-   // CHECK-SAME:   known_expanding_dimensions = #vhlo.tensor_v1<dense<> : tensor<0xi64>>,
-   // CHECK-SAME:   known_nonexpanding_dimensions = #vhlo.tensor_v1<dense<> : tensor<0xi64>>
-@@ -520,8 +561,9 @@
- }
- 
- // CHECK-LABEL: "default_dynamic_conv"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @default_dynamic_conv(%arg0: tensor<1x8x8x207xf32>, %arg1: tensor<3x3x207x16xf32>, %arg2: tensor<4xi32>) -> tensor<1x?x?x16xf32> {
--  //      CHECK: "vhlo.dynamic_conv_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.dynamic_conv_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   batch_group_count = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   feature_group_count = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   input_batch_dimension = #vhlo.integer_v1<0 : i64>,
-@@ -549,8 +591,9 @@
- }
- 
- // CHECK-LABEL: "default_dynamic_gather"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @default_dynamic_gather(%arg0 : tensor<2x4x9xf32>, %arg1 : tensor<1x5x2xi32>, %arg2 : tensor<3xi32>) -> tensor<1x5x8xf32> {
--  //      CHECK: "vhlo.dynamic_gather_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.dynamic_gather_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   collapsed_slice_dims = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-   // CHECK-SAME:   index_vector_dim = #vhlo.integer_v1<2 : i64>,
-   // CHECK-SAME:   indices_are_sorted = #vhlo.bool_v1<false>,
-@@ -576,15 +619,16 @@
-   // CHECK-SAME:   sym_name = #vhlo.string_v1<"default_func">,
-   // CHECK-SAME:   sym_visibility = #vhlo.string_v1<"">
-   // CHECK-SAME: }> ({
--  // CHECK-NEXT:   ^[[BB:bb.*]](%arg0: !vhlo.tensor_v1<!vhlo.f32_v1>):
--  // CHECK-NEXT:     "vhlo.return_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-+  // CHECK-NEXT:   ^[[BB:bb.*]](%[[ARG0:.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
-+  // CHECK-NEXT:     "vhlo.return_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-   // CHECK-NEXT: }) : () -> ()
-   func.return %arg0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "dynamic_gather"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @dynamic_gather(%arg0 : tensor<2x4x9xf32>, %arg1 : tensor<1x5x2xi32>) -> tensor<1x5x1xf32> {
--  //      CHECK: "vhlo.gather_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.gather_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   collapsed_slice_dims = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-   // CHECK-SAME:   index_vector_dim = #vhlo.integer_v1<2 : i64>,
-   // CHECK-SAME:   indices_are_sorted = #vhlo.bool_v1<false>,
-@@ -605,8 +649,9 @@
- }
- 
- // CHECK-LABEL: "default_infeed"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_infeed(%arg0: !stablehlo.token) -> (tensor<f32>, !stablehlo.token) {
--  //               CHECK: "vhlo.infeed_v1"(%arg0) <{
-+  //               CHECK: "vhlo.infeed_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   infeed_config = #vhlo.string_v1<"">,
-   // CHECK-SAME{LITERAL}:   layout = #vhlo.array_v1<[]>
-   //          CHECK-SAME: }> : (!vhlo.token_v1) -> (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.token_v1)
-@@ -615,8 +660,9 @@
- }
- 
- // CHECK-LABEL: "default_outfeed"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @default_outfeed(%arg0: tensor<f32>, %arg1: !stablehlo.token) -> !stablehlo.token {
--  //      CHECK: "vhlo.outfeed_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.outfeed_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   outfeed_config = #vhlo.string_v1<"">
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.token_v1) -> !vhlo.token_v1
-   %0 = "stablehlo.outfeed"(%arg0, %arg1) : (tensor<f32>, !stablehlo.token) -> !stablehlo.token
-@@ -624,8 +670,9 @@
- }
- 
- // CHECK-LABEL: "default_recv"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_recv(%arg0: !stablehlo.token) -> (tensor<f32>, !stablehlo.token) {
--  //      CHECK: "vhlo.recv_v1"(%arg0) <{
-+  //      CHECK: "vhlo.recv_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME:   channel_type = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   is_host_transfer = #vhlo.bool_v1<false>
-@@ -637,8 +684,9 @@
- }
- 
- // CHECK-LABEL: "default_send"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @default_send(%arg0: tensor<f32>, %arg1: !stablehlo.token) -> !stablehlo.token {
--  //      CHECK: "vhlo.send_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.send_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME:   channel_type = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   is_host_transfer = #vhlo.bool_v1<false>
-@@ -650,8 +698,9 @@
- }
- 
- // CHECK-LABEL: "default_reduce_scatter"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_reduce_scatter(%arg0: tensor<16xf32>) -> tensor<16xf32> {
--  //               CHECK: "vhlo.reduce_scatter_v1"(%arg0) <{
-+  //               CHECK: "vhlo.reduce_scatter_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0], [1]]> : tensor<2x1xi64>>,
-   //          CHECK-SAME:   scatter_dimension = #vhlo.integer_v1<0 : i64>
-@@ -673,8 +722,9 @@
- }
- 
- // CHECK-LABEL: "default_reduce_window"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @default_reduce_window(%arg0: tensor<2x17x31x7xf32>, %arg1: tensor<f32>) -> tensor<2x16x30x7xf32> {
--  //               CHECK: "vhlo.reduce_window_v1"(%arg0, %arg1)  <{
-+  //               CHECK: "vhlo.reduce_window_v1"(%[[ARG0]], %[[ARG1]])  <{
-   //          CHECK-SAME:   base_dilations = #vhlo.tensor_v1<dense<1> : tensor<4xi64>>,
-   // CHECK-SAME{LITERAL}:   padding = #vhlo.tensor_v1<dense<0> : tensor<4x2xi64>>,
-   //          CHECK-SAME:   window_dilations = #vhlo.tensor_v1<dense<1> : tensor<4xi64>>,
-@@ -696,8 +746,9 @@
- }
- 
- // CHECK-LABEL: "default_scatter"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @default_scatter(%arg0: tensor<200x100x300xf32>, %arg1: tensor<10x2xi32>, %arg2: tensor<10x300xf32>) -> tensor<200x100x300xf32> {
--  //      CHECK: "vhlo.scatter_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.scatter_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   index_vector_dim = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   indices_are_sorted = #vhlo.bool_v1<false>,
-   // CHECK-SAME:   inserted_window_dims = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-@@ -725,8 +776,9 @@
- }
- 
- // CHECK-LABEL: "default_select_and_scatter"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @default_select_and_scatter(%arg0: tensor<10x24x24x64xf32>, %arg1: tensor<10x23x23x64xf32>, %arg2: tensor<f32>) -> tensor<10x24x24x64xf32> {
--  //      CHECK: "vhlo.select_and_scatter_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.select_and_scatter_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   padding = #vhlo.tensor_v1<dense<0> : tensor<4x2xi64>>,
-   // CHECK-SAME:   window_dimensions = #vhlo.tensor_v1<dense<[1, 2, 2, 1]> : tensor<4xi64>>,
-   // CHECK-SAME:   window_strides = #vhlo.tensor_v1<dense<1> : tensor<4xi64>>
-@@ -754,8 +806,9 @@
- }
- 
- // CHECK-LABEL: "default_sort"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_sort(%arg0: tensor<16xf32>) -> tensor<16xf32> {
--  //      CHECK: "vhlo.sort_v1"(%arg0) <{
-+  //      CHECK: "vhlo.sort_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   dimension = #vhlo.integer_v1<-1 : i64>
-   // CHECK-SAME:   is_stable = #vhlo.bool_v1<false>
-   // CHECK-SAME: }> ({
-@@ -774,29 +827,33 @@
- // ============ OPS ============
- 
- // CHECK-LABEL: "op_abs"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_abs(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.abs_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.abs_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.abs"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_add"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_add(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_after_all"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_after_all(%arg0: !stablehlo.token) -> !stablehlo.token {
--  // CHECK: "vhlo.after_all_v1"(%arg0) : (!vhlo.token_v1) -> !vhlo.token_v1
-+  // CHECK: "vhlo.after_all_v1"(%[[ARG0]]) : (!vhlo.token_v1) -> !vhlo.token_v1
-   %0 = "stablehlo.after_all"(%arg0) : (!stablehlo.token) -> !stablehlo.token
-   func.return %0 : !stablehlo.token
- }
- 
- // CHECK-LABEL: "op_all_gather"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_all_gather(%arg0: tensor<16x8xf32>) -> tensor<16x16xf32> {
--  //               CHECK: "vhlo.all_gather_v1"(%arg0) <{
-+  //               CHECK: "vhlo.all_gather_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   all_gather_dim = #vhlo.integer_v1<1 : i64>
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0], [1]]> : tensor<2x1xi64>>,
-@@ -812,8 +869,9 @@
- }
- 
- // CHECK-LABEL: "op_all_reduce"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_all_reduce(%arg0: tensor<f32>) -> tensor<f32> {
--  //               CHECK: "vhlo.all_reduce_v1"(%arg0) <{
-+  //               CHECK: "vhlo.all_reduce_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0], [1]]> : tensor<2x1xi64>>,
-   //          CHECK-SAME:   use_global_device_ids = #vhlo.bool_v1<true>
-@@ -835,8 +893,9 @@
- }
- 
- // CHECK-LABEL: "op_all_to_all"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_all_to_all(%arg0: tensor<4x16xf32>) -> tensor<16x4xf32> {
--  //               CHECK: "vhlo.all_to_all_v1"(%arg0) <{
-+  //               CHECK: "vhlo.all_to_all_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<1 : i64>,
-   //          CHECK-SAME:   concat_dimension = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0, 1, 2, 3]]> : tensor<1x4xi64>>,
-@@ -854,22 +913,25 @@
- }
- 
- // CHECK-LABEL: "op_and"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_and(%arg0: tensor<i1>, %arg1: tensor<i1>) -> tensor<i1> {
--  // CHECK: "vhlo.and_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-+  // CHECK: "vhlo.and_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-   %0 = "stablehlo.and"(%arg0, %arg1) : (tensor<i1>, tensor<i1>) -> tensor<i1>
-   func.return %0 : tensor<i1>
- }
- 
- // CHECK-LABEL: "op_atan2"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_atan2(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.atan2_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.atan2_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.atan2"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_batch_norm_grad"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}}, %[[ARG3:.*]]: {{.*}}, %[[ARG4:.*]]: {{.*}})
- func.func @op_batch_norm_grad(%arg0: tensor<16x16x16x16xf32>, %arg1: tensor<16xf32>, %arg2: tensor<16xf32>, %arg3: tensor<16xf32>, %arg4: tensor<16x16x16x16xf32>) -> (tensor<16x16x16x16xf32>, tensor<16xf32>, tensor<16xf32>) {
--  //      CHECK: "vhlo.batch_norm_grad_v1"(%arg0, %arg1, %arg2, %arg3, %arg4) <{
-+  //      CHECK: "vhlo.batch_norm_grad_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]], %[[ARG3]], %[[ARG4]]) <{
-   // CHECK-SAME:   epsilon = #vhlo.float_v1<1.000000e-03 : !vhlo.f32_v1>,
-   // CHECK-SAME:   feature_index = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<16x16x16x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x16x16x16x!vhlo.f32_v1>) -> (!vhlo.tensor_v1<16x16x16x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>)
-@@ -881,8 +943,9 @@
- }
- 
- // CHECK-LABEL: "op_batch_norm_inference"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}}, %[[ARG3:.*]]: {{.*}}, %[[ARG4:.*]]: {{.*}})
- func.func @op_batch_norm_inference(%arg0: tensor<16x16x16x16xf32>, %arg1: tensor<16xf32>, %arg2: tensor<16xf32>, %arg3: tensor<16xf32>, %arg4: tensor<16xf32>) -> tensor<16x16x16x16xf32> {
--  //      CHECK: "vhlo.batch_norm_inference_v1"(%arg0, %arg1, %arg2, %arg3, %arg4) <{
-+  //      CHECK: "vhlo.batch_norm_inference_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]], %[[ARG3]], %[[ARG4]]) <{
-   // CHECK-SAME:   epsilon = #vhlo.float_v1<1.000000e-03 : !vhlo.f32_v1>,
-   // CHECK-SAME:   feature_index = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<16x16x16x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x16x16x16x!vhlo.f32_v1>
-@@ -894,8 +957,9 @@
- }
- 
- // CHECK-LABEL: "op_batch_norm_training"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_batch_norm_training(%arg0: tensor<16x16x16x16xf32>, %arg1: tensor<16xf32>, %arg2: tensor<16xf32>) -> (tensor<16x16x16x16xf32>, tensor<16xf32>, tensor<16xf32>) {
--  //      CHECK: "vhlo.batch_norm_training_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.batch_norm_training_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   epsilon = #vhlo.float_v1<1.000000e-03 : !vhlo.f32_v1>,
-   // CHECK-SAME:   feature_index = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<16x16x16x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>) -> (!vhlo.tensor_v1<16x16x16x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>)
-@@ -907,15 +971,17 @@
- }
- 
- // CHECK-LABEL: "op_bitcast_convert"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_bitcast_convert(%arg0: tensor<i32>) -> tensor<f32> {
--  // CHECK: "vhlo.bitcast_convert_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.bitcast_convert_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.bitcast_convert"(%arg0) : (tensor<i32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_broadcast_in_dim"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_broadcast_in_dim(%arg0: tensor<16xf32>) -> tensor<16x16xf32> {
--  //      CHECK: "vhlo.broadcast_in_dim_v1"(%arg0) <{
-+  //      CHECK: "vhlo.broadcast_in_dim_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   broadcast_dimensions = #vhlo.tensor_v1<dense<1> : tensor<1xi64>>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x16x!vhlo.f32_v1>
-   %0 = "stablehlo.broadcast_in_dim"(%arg0) {
-@@ -925,8 +991,9 @@
- }
- 
- // CHECK-LABEL: "op_broadcast"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_broadcast(%arg0: tensor<16xf32>) -> tensor<16x16xf32> {
--  //      CHECK: "vhlo.broadcast_v1"(%arg0) <{
-+  //      CHECK: "vhlo.broadcast_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   broadcast_sizes = #vhlo.tensor_v1<dense<16> : tensor<1xi64>>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x16x!vhlo.f32_v1>
-   %0 = "stablehlo.broadcast"(%arg0) {
-@@ -936,9 +1003,10 @@
- }
- 
- // CHECK-LABEL: "op_case"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_case(%arg0: tensor<i32>, %arg1: tensor<f32>) -> tensor<f32> {
--  //      CHECK: "vhlo.case_v1"(%arg0) ({
--  // CHECK-NEXT:   "vhlo.return_v1"(%arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-+  //      CHECK: "vhlo.case_v1"(%[[ARG0]]) ({
-+  // CHECK-NEXT:   "vhlo.return_v1"(%[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-   // CHECK-NEXT: }) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.case"(%arg0) ({
-     "stablehlo.return"(%arg1) : (tensor<f32>) -> ()
-@@ -947,22 +1015,25 @@
- }
- 
- // CHECK-LABEL: "op_cbrt"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_cbrt(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.cbrt_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.cbrt_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.cbrt"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_ceil"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_ceil(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.ceil_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.ceil_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.ceil"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_cholesky"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_cholesky(%arg0: tensor<1x16x16xf32>) -> tensor<1x16x16xf32> {
--  //      CHECK: "vhlo.cholesky_v1"(%arg0) <{
-+  //      CHECK: "vhlo.cholesky_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   lower = #vhlo.bool_v1<true>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<1x16x16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<1x16x16x!vhlo.f32_v1>
-   %0 = "stablehlo.cholesky"(%arg0) {
-@@ -972,22 +1043,25 @@
- }
- 
- // CHECK-LABEL: "op_clamp"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_clamp(%arg0: tensor<f32>, %arg1: tensor<f32>, %arg2: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.clamp_v1"(%arg0, %arg1, %arg2) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.clamp_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.clamp"(%arg0, %arg1, %arg2) : (tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_count_leading_zeros"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_count_leading_zeros(%arg0: tensor<i32>) -> tensor<i32> {
--  // CHECK: "vhlo.count_leading_zeros_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-+  // CHECK: "vhlo.count_leading_zeros_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-   %0 = "stablehlo.count_leading_zeros"(%arg0) : (tensor<i32>) -> tensor<i32>
-   func.return %0 : tensor<i32>
- }
- 
- // CHECK-LABEL: "op_collective_permute"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_collective_permute(%arg0: tensor<16x8xf32>) -> tensor<16x8xf32> {
--  //               CHECK: "vhlo.collective_permute_v1"(%arg0) <{
-+  //               CHECK: "vhlo.collective_permute_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME{LITERAL}:   source_target_pairs = #vhlo.tensor_v1<dense<[[0, 1], [1, 2], [2, 3]]> : tensor<3x2xi64>>
-   //          CHECK-SAME: }> : (!vhlo.tensor_v1<16x8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x8x!vhlo.f32_v1>
-@@ -999,8 +1073,9 @@
- }
- 
- // CHECK-LABEL: "op_compare"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_compare(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
--  //      CHECK: "vhlo.compare_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.compare_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   compare_type = #vhlo<comparison_type_v1 TOTALORDER>,
-   // CHECK-SAME:   comparison_direction = #vhlo<comparison_direction_v1 EQ>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-@@ -1012,15 +1087,17 @@
- }
- 
- // CHECK-LABEL: "op_complex"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_complex(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<complex<f32>> {
--  // CHECK: "vhlo.complex_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>
-+  // CHECK: "vhlo.complex_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>
-   %0 = "stablehlo.complex"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<complex<f32>>
-   func.return %0 : tensor<complex<f32>>
- }
- 
- // CHECK-LABEL: "op_concatenate"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_concatenate(%arg0: tensor<8xf32>, %arg1: tensor<8xf32>) -> tensor<16xf32> {
--  //      CHECK: "vhlo.concatenate_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.concatenate_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   dimension = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<8x!vhlo.f32_v1>, !vhlo.tensor_v1<8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x!vhlo.f32_v1>
-   %0 = "stablehlo.concatenate"(%arg0, %arg1) {
-@@ -1030,6 +1107,7 @@
- }
- 
- // CHECK-LABEL: "op_constant"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_constant(%arg0: tensor<f32>) -> tensor<f32> {
-   //      CHECK: "vhlo.constant_v1"() <{
-   // CHECK-SAME:   value = #vhlo.tensor_v1<dense<0.000000e+00> : tensor<f32>>
-@@ -1041,15 +1119,17 @@
- }
- 
- // CHECK-LABEL: "op_convert"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_convert(%arg0: tensor<i32>) -> tensor<f32> {
--  // CHECK: "vhlo.convert_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.convert_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.convert"(%arg0) : (tensor<i32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_convolution"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_convolution(%arg0: tensor<1x8x8x207xf32>, %arg1: tensor<3x3x207x16xf32>) -> tensor<1x7x7x16xf32> {
--  //      CHECK: "vhlo.convolution_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.convolution_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   batch_group_count = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   feature_group_count = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   input_batch_dimension = #vhlo.integer_v1<0 : i64>,
-@@ -1083,8 +1163,9 @@
- }
- 
- // CHECK-LABEL: "op_cosine"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_cosine(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.cosine_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.cosine_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.cosine"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
-@@ -1097,8 +1178,9 @@
- }
- 
- // CHECK-LABEL: "op_cross_replica_sum"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_cross_replica_sum(%arg0: tensor<f32>) -> tensor<f32> {
--  //               CHECK: "vhlo.cross-replica-sum_v1"(%arg0) <{
-+  //               CHECK: "vhlo.cross-replica-sum_v1"(%[[ARG0]]) <{
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0], [1]]> : tensor<2x1xi64>>
-   //          CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.cross-replica-sum"(%arg0) {
-@@ -1108,8 +1190,9 @@
- }
- 
- // CHECK-LABEL: "op_custom_call"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_custom_call(%arg0: tensor<f32>) -> tensor<f32> {
--  //      CHECK: "vhlo.custom_call_v1"(%arg0) <{
-+  //      CHECK: "vhlo.custom_call_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   api_version = #vhlo<api_version_v1 API_VERSION_STATUS_RETURNING>,
-   // CHECK-SAME:   backend_config = #vhlo.string_v1<"\08\03\1A\02">,
-   // CHECK-SAME:   call_target_name = #vhlo.string_v1<"foo">,
-@@ -1140,15 +1223,17 @@
- }
- 
- // CHECK-LABEL: "op_divide"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_divide(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.divide_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.divide_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.divide"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_dot_general"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_dot_general(%arg0: tensor<8x8x16xf32>, %arg1: tensor<8x16x8xf32>) -> tensor<8x8x8xf32> {
--  //      CHECK: "vhlo.dot_general_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.dot_general_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   lhs_batching_dimensions = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>,
-   // CHECK-SAME:   lhs_contracting_dimensions = #vhlo.tensor_v1<dense<2> : tensor<1xi64>>,
-   // CHECK-SAME:   precision_config = #vhlo.array_v1<[#vhlo<precision_v1 HIGHEST>, #vhlo<precision_v1 HIGHEST>]>,
-@@ -1168,8 +1253,9 @@
- }
- 
- // CHECK-LABEL: "op_dot"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_dot(%arg0: tensor<8x16xf32>, %arg1: tensor<16x8xf32>) -> tensor<8x8xf32> {
--  //      CHECK: "vhlo.dot_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.dot_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   precision_config = #vhlo.array_v1<[#vhlo<precision_v1 HIGHEST>, #vhlo<precision_v1 HIGHEST>]>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<8x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<8x8x!vhlo.f32_v1>
-   %0 = "stablehlo.dot"(%arg0, %arg1) {
-@@ -1179,8 +1265,9 @@
- }
- 
- // CHECK-LABEL: "op_dynamic_broadcast_in_dim"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_dynamic_broadcast_in_dim(%arg0: tensor<?x?xf32>, %arg1: tensor<2xindex>) -> tensor<?x?xf32> {
--  //      CHECK: "vhlo.dynamic_broadcast_in_dim_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.dynamic_broadcast_in_dim_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   broadcast_dimensions = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-   // CHECK-SAME:   known_expanding_dimensions = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>,
-   // CHECK-SAME:   known_nonexpanding_dimensions = #vhlo.tensor_v1<dense<1> : tensor<1xi64>>
-@@ -1194,8 +1281,9 @@
- }
- 
- // CHECK-LABEL: "op_dynamic_conv"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_dynamic_conv(%arg0: tensor<1x8x8x207xf32>, %arg1: tensor<3x3x207x16xf32>, %arg2: tensor<4xi32>) -> tensor<1x?x?x16xf32> {
--  //      CHECK: "vhlo.dynamic_conv_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.dynamic_conv_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   batch_group_count = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   feature_group_count = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   input_batch_dimension = #vhlo.integer_v1<0 : i64>,
-@@ -1229,8 +1317,9 @@
- }
- 
- // CHECK-LABEL: "op_dynamic_gather"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_dynamic_gather(%arg0 : tensor<2x4x9xf32>, %arg1 : tensor<1x5x2xi32>, %arg2 : tensor<3xi32>) -> tensor<1x5x8xf32> {
--  //      CHECK: "vhlo.dynamic_gather_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.dynamic_gather_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   collapsed_slice_dims = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-   // CHECK-SAME:   index_vector_dim = #vhlo.integer_v1<2 : i64>,
-   // CHECK-SAME:   indices_are_sorted = #vhlo.bool_v1<true>,
-@@ -1250,8 +1339,9 @@
- }
- 
- // CHECK-LABEL: "op_dynamic_iota"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_dynamic_iota(%arg0: tensor<1xindex>) -> tensor<?xf32> {
--  //      CHECK: "vhlo.dynamic_iota_v1"(%arg0) <{
-+  //      CHECK: "vhlo.dynamic_iota_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   iota_dimension = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<1x!vhlo.index_v1>) -> !vhlo.tensor_v1<?x!vhlo.f32_v1>
-   %0 = "stablehlo.dynamic_iota"(%arg0) {
-@@ -1261,22 +1351,25 @@
- }
- 
- // CHECK-LABEL: "op_dynamic_pad"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}}, %[[ARG3:.*]]: {{.*}}, %[[ARG4:.*]]: {{.*}})
- func.func @op_dynamic_pad(%arg0: tensor<?xf32>, %arg1: tensor<f32>, %arg2: tensor<1xindex>, %arg3: tensor<1xindex>, %arg4: tensor<1xindex>) -> tensor<?xf32> {
--  // CHECK: "vhlo.dynamic_pad_v1"(%arg0, %arg1, %arg2, %arg3, %arg4) : (!vhlo.tensor_v1<?x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>) -> !vhlo.tensor_v1<?x!vhlo.f32_v1>
-+  // CHECK: "vhlo.dynamic_pad_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]], %[[ARG3]], %[[ARG4]]) : (!vhlo.tensor_v1<?x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>) -> !vhlo.tensor_v1<?x!vhlo.f32_v1>
-   %0 = "stablehlo.dynamic_pad"(%arg0, %arg1, %arg2, %arg3, %arg4) : (tensor<?xf32>, tensor<f32>, tensor<1xindex>, tensor<1xindex>, tensor<1xindex>) -> tensor<?xf32>
-   func.return %0 : tensor<?xf32>
- }
- 
- // CHECK-LABEL: "op_dynamic_reshape"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_dynamic_reshape(%arg0: tensor<16xf32>, %arg1: tensor<2xindex>) -> tensor<?x?xf32> {
--  // CHECK: "vhlo.dynamic_reshape_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<2x!vhlo.index_v1>) -> !vhlo.tensor_v1<?x?x!vhlo.f32_v1>
-+  // CHECK: "vhlo.dynamic_reshape_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<2x!vhlo.index_v1>) -> !vhlo.tensor_v1<?x?x!vhlo.f32_v1>
-   %0 = "stablehlo.dynamic_reshape"(%arg0, %arg1) : (tensor<16xf32>, tensor<2xindex>) -> tensor<?x?xf32>
-   func.return %0 : tensor<?x?xf32>
- }
- 
- // CHECK-LABEL: "op_dynamic_slice"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_dynamic_slice(%arg0: tensor<16xf32>, %arg1: tensor<i64>) -> tensor<4xf32> {
--  //      CHECK: "vhlo.dynamic_slice_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.dynamic_slice_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   slice_sizes = #vhlo.tensor_v1<dense<4> : tensor<1xi64>>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.i64_v1>) -> !vhlo.tensor_v1<4x!vhlo.f32_v1>
-   %0 = "stablehlo.dynamic_slice"(%arg0, %arg1) {
-@@ -1286,15 +1379,17 @@
- }
- 
- // CHECK-LABEL: "op_dynamic_update_slice"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_dynamic_update_slice(%arg0: tensor<16xf32>, %arg1: tensor<4xf32>, %arg2: tensor<i64>) -> tensor<16xf32> {
--  // CHECK: "vhlo.dynamic_update_slice_v1"(%arg0, %arg1, %arg2) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<4x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.i64_v1>) -> !vhlo.tensor_v1<16x!vhlo.f32_v1>
-+  // CHECK: "vhlo.dynamic_update_slice_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<4x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.i64_v1>) -> !vhlo.tensor_v1<16x!vhlo.f32_v1>
-   %0 = "stablehlo.dynamic_update_slice"(%arg0, %arg1, %arg2) : (tensor<16xf32>, tensor<4xf32>, tensor<i64>) -> tensor<16xf32>
-   func.return %0 : tensor<16xf32>
- }
- 
- // CHECK-LABEL: "op_einsum"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_einsum(%arg0: tensor<8x16xf32>, %arg1: tensor<16x8xf32>) -> tensor<8x8xf32> {
--  //      CHECK: "vhlo.einsum_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.einsum_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   einsum_config = #vhlo.string_v1<"ab,bc->ac">
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<8x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<8x8x!vhlo.f32_v1>
-   %0 = "stablehlo.einsum"(%arg0, %arg1) {
-@@ -1304,22 +1399,25 @@
- }
- 
- // CHECK-LABEL: "op_exponential_minus_one"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_exponential_minus_one(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.exponential_minus_one_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.exponential_minus_one_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.exponential_minus_one"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_exponential"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_exponential(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.exponential_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.exponential_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.exponential"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_fft"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_fft(%arg0: tensor<16xcomplex<f32>>) -> tensor<16xcomplex<f32>> {
--  //      CHECK: "vhlo.fft_v1"(%arg0) <{
-+  //      CHECK: "vhlo.fft_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   fft_length = #vhlo.tensor_v1<dense<16> : tensor<1xi64>>,
-   // CHECK-SAME:   fft_type = #vhlo<fft_type_v1 FFT>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<16x!vhlo.complex_v1<!vhlo.f32_v1>>) -> !vhlo.tensor_v1<16x!vhlo.complex_v1<!vhlo.f32_v1>>
-@@ -1331,8 +1429,9 @@
- }
- 
- // CHECK-LABEL: "op_floor"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_floor(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.floor_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.floor_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.floor"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
-@@ -1345,16 +1444,17 @@
-   // CHECK-SAME:   sym_name = #vhlo.string_v1<"op_func">,
-   // CHECK-SAME:   sym_visibility = #vhlo.string_v1<"private">
-   // CHECK-SAME: }> ({
--  // CHECK-NEXT:   ^[[BB:bb.*]](%arg0: !vhlo.tensor_v1<!vhlo.f32_v1>):
--  // CHECK-NEXT:     "vhlo.return_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-+  // CHECK-NEXT:   ^[[BB:bb.*]](%[[ARG0:.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
-+  // CHECK-NEXT:     "vhlo.return_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-   // CHECK-NEXT: }) : () -> ()
- 
-   func.return %arg0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_gather"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_gather(%arg0 : tensor<2x4x9xf32>, %arg1 : tensor<1x5x2xi32>) -> tensor<1x5x1xf32> {
--  //      CHECK: "vhlo.gather_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.gather_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   collapsed_slice_dims = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-   // CHECK-SAME:   index_vector_dim = #vhlo.integer_v1<2 : i64>,
-   // CHECK-SAME:   indices_are_sorted = #vhlo.bool_v1<true>,
-@@ -1376,8 +1476,9 @@
- }
- 
- // CHECK-LABEL: "op_get_dimension_size"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_get_dimension_size(%arg0: tensor<?xf32>) -> tensor<i32> {
--  //      CHECK: "vhlo.get_dimension_size_v1"(%arg0) <{
-+  //      CHECK: "vhlo.get_dimension_size_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   dimension = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<?x!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-   %0 = "stablehlo.get_dimension_size"(%arg0) {
-@@ -1387,8 +1488,9 @@
- }
- 
- // CHECK-LABEL: "op_get_tuple_element"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_get_tuple_element(%arg0: tuple<tensor<f32>, tensor<i32>>) -> tensor<f32> {
--  //      CHECK: "vhlo.get_tuple_element_v1"(%arg0) <{
-+  //      CHECK: "vhlo.get_tuple_element_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   index = #vhlo.integer_v1<0 : i32>
-   // CHECK-SAME: }> : (!vhlo.tuple_v1<!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.get_tuple_element"(%arg0) {
-@@ -1398,11 +1500,12 @@
- }
- 
- // CHECK-LABEL: "op_if"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_if(%arg0: tensor<i1>, %arg1: tensor<f32>, %arg2: tensor<f32>) -> tensor<f32> {
--  //      CHECK: "vhlo.if_v1"(%arg0) ({
--  // CHECK-NEXT:   "vhlo.return_v1"(%arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-+  //      CHECK: "vhlo.if_v1"(%[[ARG0]]) ({
-+  // CHECK-NEXT:   "vhlo.return_v1"(%[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-   // CHECK-NEXT: }, {
--  // CHECK-NEXT:   "vhlo.return_v1"(%arg2) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-+  // CHECK-NEXT:   "vhlo.return_v1"(%[[ARG2]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-   // CHECK-NEXT: }) : (!vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.if"(%arg0) ({
-     "stablehlo.return"(%arg1) : (tensor<f32>) -> ()
-@@ -1413,15 +1516,17 @@
- }
- 
- // CHECK-LABEL: "op_imag"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_imag(%arg0: tensor<complex<f32>>) -> tensor<f32> {
--  // CHECK: "vhlo.imag_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.imag_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.imag"(%arg0) : (tensor<complex<f32>>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_infeed"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_infeed(%arg0: !stablehlo.token) -> (tensor<f32>, !stablehlo.token) {
--  //               CHECK: "vhlo.infeed_v1"(%arg0) <{
-+  //               CHECK: "vhlo.infeed_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   infeed_config = #vhlo.string_v1<"foo">,
-   // CHECK-SAME{LITERAL}:   layout = #vhlo.array_v1<[#vhlo.array_v1<[]>]>
-   //          CHECK-SAME: }> : (!vhlo.token_v1) -> (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.token_v1)
-@@ -1444,36 +1549,41 @@
- }
- 
- // CHECK-LABEL: "op_is_finite"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_is_finite(%arg0: tensor<f32>) -> tensor<i1> {
--  // CHECK: "vhlo.is_finite_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-+  // CHECK: "vhlo.is_finite_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-   %0 = "stablehlo.is_finite"(%arg0) : (tensor<f32>) -> tensor<i1>
-   func.return %0 : tensor<i1>
- }
- 
- // CHECK-LABEL: "op_log"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_log(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.log_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.log_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.log"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_log_plus_one"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_log_plus_one(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.log_plus_one_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.log_plus_one_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.log_plus_one"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_logistic"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_logistic(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.logistic_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.logistic_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.logistic"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_map"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_map(%arg0: tensor<16xf32>) -> tensor<16xf32> {
--  //      CHECK: "vhlo.map_v1"(%arg0) <{
-+  //      CHECK: "vhlo.map_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   dimensions = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>
-   // CHECK-SAME: }> ({
-   // CHECK-NEXT:   ^[[BB:bb.*]](%[[ARG1:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
-@@ -1491,57 +1601,65 @@
- }
- 
- // CHECK-LABEL: "op_maximum"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_maximum(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.maximum_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.maximum_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.maximum"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_minimum"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_minimum(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.minimum_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.minimum_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.minimum"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_multiply"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_multiply(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.multiply_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.multiply_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.multiply"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_negate"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_negate(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.negate_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.negate_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.negate"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_not"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_not(%arg0: tensor<i1>) -> tensor<i1> {
--  // CHECK: "vhlo.not_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-+  // CHECK: "vhlo.not_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-   %0 = "stablehlo.not"(%arg0) : (tensor<i1>) -> tensor<i1>
-   func.return %0 : tensor<i1>
- }
- 
- // CHECK-LABEL: "op_optimization_barrier"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_optimization_barrier(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.optimization_barrier_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.optimization_barrier_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.optimization_barrier"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_or"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_or(%arg0: tensor<i1>, %arg1: tensor<i1>) -> tensor<i1> {
--  // CHECK: "vhlo.or_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-+  // CHECK: "vhlo.or_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-   %0 = "stablehlo.or"(%arg0, %arg1) : (tensor<i1>, tensor<i1>) -> tensor<i1>
-   func.return %0 : tensor<i1>
- }
- 
- // CHECK-LABEL: "op_outfeed"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_outfeed(%arg0: tensor<f32>, %arg1: !stablehlo.token) -> !stablehlo.token {
--  //      CHECK: "vhlo.outfeed_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.outfeed_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   outfeed_config = #vhlo.string_v1<"foo">
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.token_v1) -> !vhlo.token_v1
-   %0 = "stablehlo.outfeed"(%arg0, %arg1) {
-@@ -1551,8 +1669,9 @@
- }
- 
- // CHECK-LABEL: "op_pad"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_pad(%arg0: tensor<8xf32>, %arg1: tensor<f32>) -> tensor<16xf32> {
--  //      CHECK: "vhlo.pad_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.pad_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   edge_padding_high = #vhlo.tensor_v1<dense<4> : tensor<1xi64>>,
-   // CHECK-SAME:   edge_padding_low = #vhlo.tensor_v1<dense<4> : tensor<1xi64>>,
-   // CHECK-SAME:   interior_padding = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>
-@@ -1566,36 +1685,41 @@
- }
- 
- // CHECK-LABEL: "op_popcnt"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_popcnt(%arg0: tensor<i32>) -> tensor<i32> {
--  // CHECK: "vhlo.popcnt_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-+  // CHECK: "vhlo.popcnt_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-   %0 = "stablehlo.popcnt"(%arg0) : (tensor<i32>) -> tensor<i32>
-   func.return %0 : tensor<i32>
- }
- 
- // CHECK-LABEL: "op_power"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_power(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.power_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.power_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.power"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_real_dynamic_slice"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}}, %[[ARG3:.*]]: {{.*}})
- func.func @op_real_dynamic_slice(%arg0: tensor<?xf32>, %arg1: tensor<1xindex>, %arg2: tensor<1xindex>, %arg3: tensor<1xindex>) -> tensor<?xf32> {
--  // CHECK: "vhlo.real_dynamic_slice_v1"(%arg0, %arg1, %arg2, %arg3) : (!vhlo.tensor_v1<?x!vhlo.f32_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>) -> !vhlo.tensor_v1<?x!vhlo.f32_v1>
-+  // CHECK: "vhlo.real_dynamic_slice_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]], %[[ARG3]]) : (!vhlo.tensor_v1<?x!vhlo.f32_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>) -> !vhlo.tensor_v1<?x!vhlo.f32_v1>
-   %0 = "stablehlo.real_dynamic_slice"(%arg0, %arg1, %arg2, %arg3) : (tensor<?xf32>, tensor<1xindex>, tensor<1xindex>, tensor<1xindex>) -> tensor<?xf32>
-   func.return %0 : tensor<?xf32>
- }
- 
- // CHECK-LABEL: "op_real"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_real(%arg0: tensor<complex<f32>>) -> tensor<f32> {
--  // CHECK: "vhlo.real_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.real_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.real"(%arg0) : (tensor<complex<f32>>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_recv"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_recv(%arg0: !stablehlo.token) -> (tensor<f32>, !stablehlo.token) {
--  //      CHECK: "vhlo.recv_v1"(%arg0) <{
-+  //      CHECK: "vhlo.recv_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME:   channel_type = #vhlo.integer_v1<3 : i64>,
-   // CHECK-SAME:   is_host_transfer = #vhlo.bool_v1<true>
-@@ -1608,8 +1732,9 @@
- }
- 
- // CHECK-LABEL: "op_reduce"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_reduce(%arg0: tensor<16xf32>, %arg1: tensor<f32>) -> tensor<f32> {
--  //  CHECK: "vhlo.reduce_v1"(%arg0, %arg1)
-+  //  CHECK: "vhlo.reduce_v1"(%[[ARG0]], %[[ARG1]])
-   //  CHECK:   ^[[BB:bb.*]](%[[ARG1:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>, %[[ARG2:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
-   //  CHECK:     "vhlo.return_v1"(%[[VAL1:.*]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-   //  CHECK: }) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-@@ -1624,8 +1749,9 @@
- }
- 
- // CHECK-LABEL: "op_reduce_precision"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_reduce_precision(%arg0: tensor<f32>) -> tensor<f32> {
--  //      CHECK: "vhlo.reduce_precision_v1"(%arg0) <{
-+  //      CHECK: "vhlo.reduce_precision_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   exponent_bits = #vhlo.integer_v1<8 : i32>
-   // CHECK-SAME:   mantissa_bits = #vhlo.integer_v1<10 : i32>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-@@ -1637,8 +1763,9 @@
- }
- 
- // CHECK-LABEL: "op_reduce_scatter"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_reduce_scatter(%arg0: tensor<16xf32>) -> tensor<16xf32> {
--  //               CHECK: "vhlo.reduce_scatter_v1"(%arg0) <{
-+  //               CHECK: "vhlo.reduce_scatter_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0], [1]]> : tensor<2x1xi64>>,
-   //          CHECK-SAME:   scatter_dimension = #vhlo.integer_v1<0 : i64>
-@@ -1662,8 +1789,9 @@
- }
- 
- // CHECK-LABEL: "op_reduce_window"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_reduce_window(%arg0: tensor<2x17x31x7xf32>, %arg1: tensor<f32>) -> tensor<2x9x16x7xf32> {
--  //               CHECK: "vhlo.reduce_window_v1"(%arg0, %arg1) <{
-+  //               CHECK: "vhlo.reduce_window_v1"(%[[ARG0]], %[[ARG1]]) <{
-   //          CHECK-SAME:   base_dilations = #vhlo.tensor_v1<dense<[1, 2, 2, 1]> : tensor<4xi64>>,
-   // CHECK-SAME{LITERAL}:   padding = #vhlo.tensor_v1<dense<[[0, 0], [2, 0], [0, 2], [0, 0]]> : tensor<4x2xi64>>,
-   //          CHECK-SAME:   window_dilations = #vhlo.tensor_v1<dense<[1, 2, 2, 1]> : tensor<4xi64>>,
-@@ -1689,8 +1817,9 @@
- }
- 
- // CHECK-LABEL: "op_remainder"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_remainder(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.remainder_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.remainder_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.remainder"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
-@@ -1710,16 +1839,18 @@
- }
- 
- // CHECK-LABEL: "op_reshape"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_reshape(%arg0: tensor<16xf32>) -> tensor<4x4xf32> {
--  // CHECK: "vhlo.reshape_v1"(%arg0) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<4x4x!vhlo.f32_v1>
-+  // CHECK: "vhlo.reshape_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<4x4x!vhlo.f32_v1>
-   %0 = "stablehlo.reshape"(%arg0) : (tensor<16xf32>) -> tensor<4x4xf32>
-   func.return %0 : tensor<4x4xf32>
- }
- 
- // CHECK-LABEL: "op_return"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_return(%arg0: tensor<i32>, %arg1: tensor<f32>) -> tensor<f32> {
--  //      CHECK: "vhlo.case_v1"(%arg0) ({
--  // CHECK-NEXT:   "vhlo.return_v1"(%arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-+  //      CHECK: "vhlo.case_v1"(%[[ARG0]]) ({
-+  // CHECK-NEXT:   "vhlo.return_v1"(%[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-   // CHECK-NEXT: }) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.case"(%arg0) ({
-     "stablehlo.return"(%arg1) : (tensor<f32>) -> ()
-@@ -1728,8 +1859,9 @@
- }
- 
- // CHECK-LABEL: "op_reverse"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_reverse(%arg0: tensor<16xf32>) -> tensor<16xf32> {
--  //      CHECK: "vhlo.reverse_v1"(%arg0) <{
-+  //      CHECK: "vhlo.reverse_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   dimensions = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x!vhlo.f32_v1>
-   %0 = "stablehlo.reverse"(%arg0) {
-@@ -1739,8 +1871,9 @@
- }
- 
- // CHECK-LABEL: "op_rng_bit_generator"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_rng_bit_generator(%arg0: tensor<f32>) -> (tensor<f32>, tensor<f32>) {
--  //      CHECK: "vhlo.rng_bit_generator_v1"(%arg0) <{
-+  //      CHECK: "vhlo.rng_bit_generator_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   rng_algorithm = #vhlo<rng_algorithm_v1 PHILOX>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>)
-   %0:2 = "stablehlo.rng_bit_generator"(%arg0) {
-@@ -1750,8 +1883,9 @@
- }
- 
- // CHECK-LABEL: "op_rng"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_rng(%arg0: tensor<f32>, %arg1: tensor<f32>, %arg2: tensor<0xindex>) -> tensor<f32> {
--  //      CHECK: "vhlo.rng_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.rng_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   rng_distribution = #vhlo<rng_distribution_v1 NORMAL>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<0x!vhlo.index_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.rng"(%arg0, %arg1, %arg2) {
-@@ -1761,29 +1895,33 @@
- }
- 
- // CHECK-LABEL: "op_round_nearest_afz"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_round_nearest_afz(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.round_nearest_afz_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.round_nearest_afz_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.round_nearest_afz"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_round_nearest_even"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_round_nearest_even(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.round_nearest_even_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.round_nearest_even_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.round_nearest_even"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_rsqrt"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_rsqrt(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.rsqrt_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.rsqrt_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.rsqrt"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_scatter"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_scatter(%arg0: tensor<200x100x300xf32>, %arg1: tensor<10x2xi32>, %arg2: tensor<10x300xf32>) -> tensor<200x100x300xf32> {
--  //      CHECK: "vhlo.scatter_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.scatter_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   index_vector_dim = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   indices_are_sorted = #vhlo.bool_v1<true>,
-   // CHECK-SAME:   inserted_window_dims = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-@@ -1813,8 +1951,9 @@
- }
- 
- // CHECK-LABEL: "op_select_and_scatter"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_select_and_scatter(%arg0: tensor<10x24x24x64xf32>, %arg1: tensor<12x13x13x66xf32>, %arg2: tensor<f32>) -> tensor<10x24x24x64xf32> {
--  //      CHECK: "vhlo.select_and_scatter_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.select_and_scatter_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   padding = #vhlo.tensor_v1<dense<1> : tensor<4x2xi64>>,
-   // CHECK-SAME:   window_dimensions = #vhlo.tensor_v1<dense<[1, 2, 2, 1]> : tensor<4xi64>>,
-   // CHECK-SAME:   window_strides = #vhlo.tensor_v1<dense<[1, 2, 2, 1]> : tensor<4xi64>>
-@@ -1844,15 +1983,17 @@
- }
- 
- // CHECK-LABEL: "op_select"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_select(%arg0: tensor<i1>, %arg1: tensor<f32>, %arg2: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.select_v1"(%arg0, %arg1, %arg2) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.select_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.select"(%arg0, %arg1, %arg2) : (tensor<i1>, tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_send"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_send(%arg0: tensor<f32>, %arg1: !stablehlo.token) -> !stablehlo.token {
--  //      CHECK: "vhlo.send_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.send_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME:   channel_type = #vhlo.integer_v1<2 : i64>,
-   // CHECK-SAME:   is_host_transfer = #vhlo.bool_v1<true>
-@@ -1865,8 +2006,9 @@
- }
- 
- // CHECK-LABEL: "op_set_dimension_size"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_set_dimension_size(%arg0: tensor<?xf32>, %arg1: tensor<i32>) -> tensor<16xf32> {
--  //      CHECK: "vhlo.set_dimension_size_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.set_dimension_size_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   dimension = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<?x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<16x!vhlo.f32_v1>
-   %0 = "stablehlo.set_dimension_size"(%arg0, %arg1) {
-@@ -1876,43 +2018,49 @@
- }
- 
- // CHECK-LABEL: "op_shift_left"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_shift_left(%arg0: tensor<i32>, %arg1: tensor<i32>) -> tensor<i32> {
--  // CHECK: "vhlo.shift_left_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.i32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-+  // CHECK: "vhlo.shift_left_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-   %0 = "stablehlo.shift_left"(%arg0, %arg1) : (tensor<i32>, tensor<i32>) -> tensor<i32>
-   func.return %0 : tensor<i32>
- }
- 
- // CHECK-LABEL: "op_shift_right_arithmetic"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_shift_right_arithmetic(%arg0: tensor<i32>, %arg1: tensor<i32>) -> tensor<i32> {
--  // CHECK: "vhlo.shift_right_arithmetic_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.i32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-+  // CHECK: "vhlo.shift_right_arithmetic_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-   %0 = "stablehlo.shift_right_arithmetic"(%arg0, %arg1) : (tensor<i32>, tensor<i32>) -> tensor<i32>
-   func.return %0 : tensor<i32>
- }
- 
- // CHECK-LABEL: "op_shift_right_logical"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_shift_right_logical(%arg0: tensor<i32>, %arg1: tensor<i32>) -> tensor<i32> {
--  // CHECK: "vhlo.shift_right_logical_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.i32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-+  // CHECK: "vhlo.shift_right_logical_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-   %0 = "stablehlo.shift_right_logical"(%arg0, %arg1) : (tensor<i32>, tensor<i32>) -> tensor<i32>
-   func.return %0 : tensor<i32>
- }
- 
- // CHECK-LABEL: "op_sign"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_sign(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.sign_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.sign_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.sign"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_sine"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_sine(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.sine_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.sine_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.sine"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_slice"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_slice(%arg0: tensor<16xf32>) -> tensor<4xf32> {
--  //      CHECK: "vhlo.slice_v1"(%arg0) <{
-+  //      CHECK: "vhlo.slice_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   limit_indices = #vhlo.tensor_v1<dense<4> : tensor<1xi64>>,
-   // CHECK-SAME:   start_indices = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>,
-   // CHECK-SAME:   strides = #vhlo.tensor_v1<dense<1> : tensor<1xi64>>
-@@ -1926,8 +2074,9 @@
- }
- 
- // CHECK-LABEL: "op_sort"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_sort(%arg0: tensor<16xf32>) -> tensor<16xf32> {
--  //      CHECK: "vhlo.sort_v1"(%arg0) <{
-+  //      CHECK: "vhlo.sort_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   dimension = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME:   is_stable = #vhlo.bool_v1<true>
-   // CHECK-SAME: }> ({
-@@ -1947,29 +2096,33 @@
- }
- 
- // CHECK-LABEL: "op_sqrt"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_sqrt(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.sqrt_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.sqrt_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.sqrt"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_subtract"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_subtract(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.subtract_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.subtract_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.subtract"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_tanh"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_tanh(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.tanh_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.tanh_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.tanh"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_torch_index_select"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_torch_index_select(%arg0: tensor<5x1x5xf32>, %arg1: tensor<2xi32>) ->  tensor<2x1x5xf32> {
--  //      CHECK: "vhlo.torch_index_select_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.torch_index_select_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   batch_dims = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME:   dim = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<5x1x5x!vhlo.f32_v1>, !vhlo.tensor_v1<2x!vhlo.i32_v1>) -> !vhlo.tensor_v1<2x1x5x!vhlo.f32_v1>
-@@ -1981,8 +2134,9 @@
- }
- 
- // CHECK-LABEL: "op_trace"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_trace(%arg0: tensor<f32>) {
--  //      CHECK: "vhlo.trace_v1"(%arg0) <{
-+  //      CHECK: "vhlo.trace_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   tag = #vhlo.string_v1<"foo">
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-   "stablehlo.trace"(%arg0) {
-@@ -1992,8 +2146,9 @@
- }
- 
- // CHECK-LABEL: "op_transpose"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_transpose(%arg0: tensor<16x8xf32>) ->  tensor<8x16xf32> {
--  //      CHECK: "vhlo.transpose_v1"(%arg0) <{
-+  //      CHECK: "vhlo.transpose_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   permutation = #vhlo.tensor_v1<dense<[1, 0]> : tensor<2xi64>>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<16x8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<8x16x!vhlo.f32_v1>
-   %0 = "stablehlo.transpose"(%arg0) {
-@@ -2003,8 +2158,9 @@
- }
- 
- // CHECK-LABEL: "op_triangular_solve"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_triangular_solve(%arg0: tensor<16x16xf32>, %arg1: tensor<16x16xf32>) ->  tensor<16x16xf32> {
--  //      CHECK: "vhlo.triangular_solve_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.triangular_solve_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   left_side = #vhlo.bool_v1<true>,
-   // CHECK-SAME:   lower = #vhlo.bool_v1<true>,
-   // CHECK-SAME:   transpose_a = #vhlo<transpose_v1 NO_TRANSPOSE>,
-@@ -2020,15 +2176,17 @@
- }
- 
- // CHECK-LABEL: "op_tuple"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_tuple(%arg0: tensor<f32>) -> tuple<tensor<f32>> {
--  // CHECK: "vhlo.tuple_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tuple_v1<!vhlo.tensor_v1<!vhlo.f32_v1>>
-+  // CHECK: "vhlo.tuple_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tuple_v1<!vhlo.tensor_v1<!vhlo.f32_v1>>
-   %0 = "stablehlo.tuple"(%arg0) : (tensor<f32>) -> tuple<tensor<f32>>
-   func.return %0 : tuple<tensor<f32>>
- }
- 
- // CHECK-LABEL: "op_unary_einsum"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_unary_einsum(%arg0: tensor<8x16xf32>) -> tensor<8xf32> {
--  //      CHECK: "vhlo.unary_einsum_v1"(%arg0) <{
-+  //      CHECK: "vhlo.unary_einsum_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   einsum_config = #vhlo.string_v1<"ab->a">
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<8x16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<8x!vhlo.f32_v1>
-   %0 = "stablehlo.unary_einsum"(%arg0) {
-@@ -2038,22 +2196,25 @@
- }
- 
- // CHECK-LABEL: "op_uniform_dequantize"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_uniform_dequantize(%arg0: tensor<!quant.uniform<i8:f32, 34.0:16>>) -> tensor<f32> {
--  // CHECK: "vhlo.uniform_dequantize_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.uniform_dequantize_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.uniform_dequantize"(%arg0) : (tensor<!quant.uniform<i8:f32, 34.0:16>>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_uniform_quantize"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_uniform_quantize(%arg0: tensor<f32>) -> tensor<!quant.uniform<i8:f32, 34.0:16>> {
--  // CHECK: "vhlo.uniform_quantize_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>
-+  // CHECK: "vhlo.uniform_quantize_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>
-   %0 = "stablehlo.uniform_quantize"(%arg0) : (tensor<f32>) -> tensor<!quant.uniform<i8:f32, 34.0:16>>
-   func.return %0 : tensor<!quant.uniform<i8:f32, 34.0:16>>
- }
- 
- // CHECK-LABEL: "op_while"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_while(%arg0: tensor<i1>) -> tensor<i1> {
--  //      CHECK: "vhlo.while_v1"(%arg0) ({
-+  //      CHECK: "vhlo.while_v1"(%[[ARG0]]) ({
-   // CHECK-NEXT:   ^[[BB:bb.*]](%[[ARG1:arg.*]]: !vhlo.tensor_v1<!vhlo.bool_v1>):
-   // CHECK-NEXT:     "vhlo.return_v1"(%[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>) -> ()
-   // CHECK-NEXT:   }, {
-@@ -2071,8 +2232,9 @@
- }
- 
- // CHECK-LABEL: "op_xor"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_xor(%arg0: tensor<i1>, %arg1: tensor<i1>) -> tensor<i1> {
--  // CHECK: "vhlo.xor_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-+  // CHECK: "vhlo.xor_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-   %0 = "stablehlo.xor"(%arg0, %arg1) : (tensor<i1>, tensor<i1>) -> tensor<i1>
-   func.return %0 : tensor<i1>
- }
-@@ -2080,190 +2242,217 @@
- // ============ TYPES ============
- 
- // CHECK-LABEL: "type_i1"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_i1(%arg0: tensor<i1>, %arg1: tensor<i1>) -> tensor<i1> {
--  // CHECK: "vhlo.and_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-+  // CHECK: "vhlo.and_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-   %0 = "stablehlo.and"(%arg0, %arg1) : (tensor<i1>, tensor<i1>) -> tensor<i1>
-   func.return %0 : tensor<i1>
- }
- 
- // CHECK-LABEL: "type_i4"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_i4(%arg0: tensor<i4>, %arg1: tensor<i4>) -> tensor<i4> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.i4_v1>, !vhlo.tensor_v1<!vhlo.i4_v1>) -> !vhlo.tensor_v1<!vhlo.i4_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i4_v1>, !vhlo.tensor_v1<!vhlo.i4_v1>) -> !vhlo.tensor_v1<!vhlo.i4_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<i4>, tensor<i4>) -> tensor<i4>
-   func.return %0 : tensor<i4>
- }
- 
- // CHECK-LABEL: "type_i8"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_i8(%arg0: tensor<i8>, %arg1: tensor<i8>) -> tensor<i8> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.i8_v1>, !vhlo.tensor_v1<!vhlo.i8_v1>) -> !vhlo.tensor_v1<!vhlo.i8_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i8_v1>, !vhlo.tensor_v1<!vhlo.i8_v1>) -> !vhlo.tensor_v1<!vhlo.i8_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<i8>, tensor<i8>) -> tensor<i8>
-   func.return %0 : tensor<i8>
- }
- 
- // CHECK-LABEL: "type_i16"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_i16(%arg0: tensor<i16>, %arg1: tensor<i16>) -> tensor<i16> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.i16_v1>, !vhlo.tensor_v1<!vhlo.i16_v1>) -> !vhlo.tensor_v1<!vhlo.i16_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i16_v1>, !vhlo.tensor_v1<!vhlo.i16_v1>) -> !vhlo.tensor_v1<!vhlo.i16_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<i16>, tensor<i16>) -> tensor<i16>
-   func.return %0 : tensor<i16>
- }
- 
- // CHECK-LABEL: "type_i32"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_i32(%arg0: tensor<i32>, %arg1: tensor<i32>) -> tensor<i32> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.i32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<i32>, tensor<i32>) -> tensor<i32>
-   func.return %0 : tensor<i32>
- }
- 
- // CHECK-LABEL: "type_i64"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_i64(%arg0: tensor<i64>, %arg1: tensor<i64>) -> tensor<i64> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.i64_v1>, !vhlo.tensor_v1<!vhlo.i64_v1>) -> !vhlo.tensor_v1<!vhlo.i64_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i64_v1>, !vhlo.tensor_v1<!vhlo.i64_v1>) -> !vhlo.tensor_v1<!vhlo.i64_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<i64>, tensor<i64>) -> tensor<i64>
-   func.return %0 : tensor<i64>
- }
- 
- // CHECK-LABEL: "type_ui4"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_ui4(%arg0: tensor<ui4>, %arg1: tensor<ui4>) -> tensor<ui4> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.ui4_v1>, !vhlo.tensor_v1<!vhlo.ui4_v1>) -> !vhlo.tensor_v1<!vhlo.ui4_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.ui4_v1>, !vhlo.tensor_v1<!vhlo.ui4_v1>) -> !vhlo.tensor_v1<!vhlo.ui4_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<ui4>, tensor<ui4>) -> tensor<ui4>
-   func.return %0 : tensor<ui4>
- }
- 
- // CHECK-LABEL: "type_ui8"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_ui8(%arg0: tensor<ui8>, %arg1: tensor<ui8>) -> tensor<ui8> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.ui8_v1>, !vhlo.tensor_v1<!vhlo.ui8_v1>) -> !vhlo.tensor_v1<!vhlo.ui8_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.ui8_v1>, !vhlo.tensor_v1<!vhlo.ui8_v1>) -> !vhlo.tensor_v1<!vhlo.ui8_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<ui8>, tensor<ui8>) -> tensor<ui8>
-   func.return %0 : tensor<ui8>
- }
- 
- // CHECK-LABEL: "type_ui16"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_ui16(%arg0: tensor<ui16>, %arg1: tensor<ui16>) -> tensor<ui16> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.ui16_v1>, !vhlo.tensor_v1<!vhlo.ui16_v1>) -> !vhlo.tensor_v1<!vhlo.ui16_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.ui16_v1>, !vhlo.tensor_v1<!vhlo.ui16_v1>) -> !vhlo.tensor_v1<!vhlo.ui16_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<ui16>, tensor<ui16>) -> tensor<ui16>
-   func.return %0 : tensor<ui16>
- }
- 
- // CHECK-LABEL: "type_ui32"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_ui32(%arg0: tensor<ui32>, %arg1: tensor<ui32>) -> tensor<ui32> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.ui32_v1>, !vhlo.tensor_v1<!vhlo.ui32_v1>) -> !vhlo.tensor_v1<!vhlo.ui32_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.ui32_v1>, !vhlo.tensor_v1<!vhlo.ui32_v1>) -> !vhlo.tensor_v1<!vhlo.ui32_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<ui32>, tensor<ui32>) -> tensor<ui32>
-   func.return %0 : tensor<ui32>
- }
- 
- // CHECK-LABEL: "type_ui64"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_ui64(%arg0: tensor<ui64>, %arg1: tensor<ui64>) -> tensor<ui64> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.ui64_v1>, !vhlo.tensor_v1<!vhlo.ui64_v1>) -> !vhlo.tensor_v1<!vhlo.ui64_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.ui64_v1>, !vhlo.tensor_v1<!vhlo.ui64_v1>) -> !vhlo.tensor_v1<!vhlo.ui64_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<ui64>, tensor<ui64>) -> tensor<ui64>
-   func.return %0 : tensor<ui64>
- }
- 
- // CHECK-LABEL: "type_f8E4M3FN"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_f8E4M3FN(%arg0: tensor<f8E4M3FN>, %arg1: tensor<f8E4M3FN>) -> tensor<f8E4M3FN> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f8E4M3FN_v1>, !vhlo.tensor_v1<!vhlo.f8E4M3FN_v1>) -> !vhlo.tensor_v1<!vhlo.f8E4M3FN_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f8E4M3FN_v1>, !vhlo.tensor_v1<!vhlo.f8E4M3FN_v1>) -> !vhlo.tensor_v1<!vhlo.f8E4M3FN_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f8E4M3FN>, tensor<f8E4M3FN>) -> tensor<f8E4M3FN>
-   func.return %0 : tensor<f8E4M3FN>
- }
- 
- // CHECK-LABEL: "type_f8E5M2"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_f8E5M2(%arg0: tensor<f8E5M2>, %arg1: tensor<f8E5M2>) -> tensor<f8E5M2> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f8E5M2_v1>, !vhlo.tensor_v1<!vhlo.f8E5M2_v1>) -> !vhlo.tensor_v1<!vhlo.f8E5M2_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f8E5M2_v1>, !vhlo.tensor_v1<!vhlo.f8E5M2_v1>) -> !vhlo.tensor_v1<!vhlo.f8E5M2_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f8E5M2>, tensor<f8E5M2>) -> tensor<f8E5M2>
-   func.return %0 : tensor<f8E5M2>
- }
- 
- // CHECK-LABEL: "type_f8E4M3FNUZ"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_f8E4M3FNUZ(%arg0: tensor<f8E4M3FNUZ>, %arg1: tensor<f8E4M3FNUZ>) -> tensor<f8E4M3FNUZ> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f8E4M3FNUZ_v1>, !vhlo.tensor_v1<!vhlo.f8E4M3FNUZ_v1>) -> !vhlo.tensor_v1<!vhlo.f8E4M3FNUZ_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f8E4M3FNUZ_v1>, !vhlo.tensor_v1<!vhlo.f8E4M3FNUZ_v1>) -> !vhlo.tensor_v1<!vhlo.f8E4M3FNUZ_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f8E4M3FNUZ>, tensor<f8E4M3FNUZ>) -> tensor<f8E4M3FNUZ>
-   func.return %0 : tensor<f8E4M3FNUZ>
- }
- 
- // CHECK-LABEL: "type_f8E4M3B11FNUZ"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_f8E4M3B11FNUZ(%arg0: tensor<f8E4M3B11FNUZ>, %arg1: tensor<f8E4M3B11FNUZ>) -> tensor<f8E4M3B11FNUZ> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f8E4M3B11FNUZ_v1>, !vhlo.tensor_v1<!vhlo.f8E4M3B11FNUZ_v1>) -> !vhlo.tensor_v1<!vhlo.f8E4M3B11FNUZ_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f8E4M3B11FNUZ_v1>, !vhlo.tensor_v1<!vhlo.f8E4M3B11FNUZ_v1>) -> !vhlo.tensor_v1<!vhlo.f8E4M3B11FNUZ_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f8E4M3B11FNUZ>, tensor<f8E4M3B11FNUZ>) -> tensor<f8E4M3B11FNUZ>
-   func.return %0 : tensor<f8E4M3B11FNUZ>
- }
- 
- // CHECK-LABEL: "type_f8E5M2FNUZ"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_f8E5M2FNUZ(%arg0: tensor<f8E5M2FNUZ>, %arg1: tensor<f8E5M2FNUZ>) -> tensor<f8E5M2FNUZ> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f8E5M2FNUZ_v1>, !vhlo.tensor_v1<!vhlo.f8E5M2FNUZ_v1>) -> !vhlo.tensor_v1<!vhlo.f8E5M2FNUZ_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f8E5M2FNUZ_v1>, !vhlo.tensor_v1<!vhlo.f8E5M2FNUZ_v1>) -> !vhlo.tensor_v1<!vhlo.f8E5M2FNUZ_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f8E5M2FNUZ>, tensor<f8E5M2FNUZ>) -> tensor<f8E5M2FNUZ>
-   func.return %0 : tensor<f8E5M2FNUZ>
- }
- 
- // CHECK-LABEL: "type_bf16"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_bf16(%arg0: tensor<bf16>, %arg1: tensor<bf16>) -> tensor<bf16> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.bf16_v1>, !vhlo.tensor_v1<!vhlo.bf16_v1>) -> !vhlo.tensor_v1<!vhlo.bf16_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.bf16_v1>, !vhlo.tensor_v1<!vhlo.bf16_v1>) -> !vhlo.tensor_v1<!vhlo.bf16_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<bf16>, tensor<bf16>) -> tensor<bf16>
-   func.return %0 : tensor<bf16>
- }
- 
- // CHECK-LABEL: "type_f16"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_f16(%arg0: tensor<f16>, %arg1: tensor<f16>) -> tensor<f16> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f16_v1>, !vhlo.tensor_v1<!vhlo.f16_v1>) -> !vhlo.tensor_v1<!vhlo.f16_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f16_v1>, !vhlo.tensor_v1<!vhlo.f16_v1>) -> !vhlo.tensor_v1<!vhlo.f16_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f16>, tensor<f16>) -> tensor<f16>
-   func.return %0 : tensor<f16>
- }
- 
- // CHECK-LABEL: "type_f32"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_f32(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "type_f64"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_f64(%arg0: tensor<f64>, %arg1: tensor<f64>) -> tensor<f64> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f64_v1>, !vhlo.tensor_v1<!vhlo.f64_v1>) -> !vhlo.tensor_v1<!vhlo.f64_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f64_v1>, !vhlo.tensor_v1<!vhlo.f64_v1>) -> !vhlo.tensor_v1<!vhlo.f64_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f64>, tensor<f64>) -> tensor<f64>
-   func.return %0 : tensor<f64>
- }
- 
- // CHECK-LABEL: "type_complex_f32"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_complex_f32(%arg0: tensor<complex<f32>>, %arg1: tensor<complex<f32>>) -> tensor<complex<f32>> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>, !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>) -> !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>, !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>) -> !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<complex<f32>>, tensor<complex<f32>>) -> tensor<complex<f32>>
-   func.return %0 : tensor<complex<f32>>
- }
- 
- // CHECK-LABEL: "type_complex_f64"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_complex_f64(%arg0: tensor<complex<f64>>, %arg1: tensor<complex<f64>>) -> tensor<complex<f64>> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f64_v1>>, !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f64_v1>>) -> !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f64_v1>>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f64_v1>>, !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f64_v1>>) -> !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f64_v1>>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<complex<f64>>, tensor<complex<f64>>) -> tensor<complex<f64>>
-   func.return %0 : tensor<complex<f64>>
- }
- 
- // CHECK-LABEL: "type_dynamism_ranked"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @type_dynamism_ranked(%arg0: tensor<?xf32>) -> tensor<?xf32> {
--  // CHECK: "vhlo.abs_v1"(%arg0) : (!vhlo.tensor_v1<?x!vhlo.f32_v1>) -> !vhlo.tensor_v1<?x!vhlo.f32_v1>
-+  // CHECK: "vhlo.abs_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<?x!vhlo.f32_v1>) -> !vhlo.tensor_v1<?x!vhlo.f32_v1>
-   %0 = "stablehlo.abs"(%arg0) : (tensor<?xf32>) -> tensor<?xf32>
-   func.return %0 : tensor<?xf32>
- }
- 
- // CHECK-LABEL: "type_quantization"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_quantization(%arg0: tensor<!quant.uniform<i8:f32, 34.0:16>>, %arg1: tensor<!quant.uniform<i8:f32, 34.0:16>>) -> tensor<!quant.uniform<i8:f32, 34.0:16>> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>, !vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>) -> !vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>, !vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>) -> !vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<!quant.uniform<i8:f32, 34.0:16>>, tensor<!quant.uniform<i8:f32, 34.0:16>>) -> tensor<!quant.uniform<i8:f32, 34.0:16>>
-   func.return %0 : tensor<!quant.uniform<i8:f32, 34.0:16>>
- }
- 
- //       CHECK: function_type = #vhlo.type_v1<!vhlo.func_v1<(!vhlo.token_v1) -> !vhlo.token_v1>>
- // CHECK-LABEL: "type_token_callee"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @type_token_callee(%arg0: !stablehlo.token) -> !stablehlo.token {
--  // CHECK: "vhlo.return_v1"(%arg0) : (!vhlo.token_v1) -> ()
-+  // CHECK: "vhlo.return_v1"(%[[ARG0]]) : (!vhlo.token_v1) -> ()
-   return %arg0 : !stablehlo.token
- }
- 
- //       CHECK: function_type = #vhlo.type_v1<!vhlo.func_v1<(!vhlo.token_v1) -> !vhlo.token_v1>>
- // CHECK-LABEL: "type_token_caller"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @type_token_caller(%arg0: !stablehlo.token) -> !stablehlo.token {
--  // CHECK:      "vhlo.call_v1"(%arg0) <{callee = #vhlo.string_v1<"type_token_callee">}
-+  // CHECK:      "vhlo.call_v1"(%[[ARG0]]) <{callee = #vhlo.string_v1<"type_token_callee">}
-   // CHECK-SAME: (!vhlo.token_v1) -> !vhlo.token_v1
-   %0 = func.call @type_token_callee(%arg0) : (!stablehlo.token) -> !stablehlo.token
-   return %0 : !stablehlo.token
- }
- 
- // CHECK-LABEL: "type_tuple"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @type_tuple(%arg0: tuple<tensor<f32>>) -> tuple<!stablehlo.token> {
-   %0 = "stablehlo.custom_call"(%arg0) {
-     call_target_name = "foo"
-diff --ruN a/stablehlo/stablehlo/tests/vhlo/stablehlo_legalize_to_vhlo.0_17_0.mlir b/stablehlo/stablehlo/tests/vhlo/stablehlo_legalize_to_vhlo.0_17_0.mlir
---- stablehlo/stablehlo/tests/vhlo/stablehlo_legalize_to_vhlo.0_17_0.mlir
-+++ stablehlo/stablehlo/tests/vhlo/stablehlo_legalize_to_vhlo.0_17_0.mlir
-@@ -13,6 +13,7 @@
- // ============ ATTRIBUTES ============
- 
- // CHECK-LABEL: "attr_comparison_direction_eq"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_direction_eq(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     // CHECK: comparison_direction = #vhlo<comparison_direction_v1 EQ>
-@@ -22,6 +23,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_direction_ne"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_direction_ne(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     // CHECK: comparison_direction = #vhlo<comparison_direction_v1 NE>
-@@ -31,6 +33,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_direction_ge"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_direction_ge(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     // CHECK: comparison_direction = #vhlo<comparison_direction_v1 GE>
-@@ -40,6 +43,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_direction_gt"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_direction_gt(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     // CHECK: comparison_direction = #vhlo<comparison_direction_v1 GT>
-@@ -49,6 +53,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_direction_le"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_direction_le(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     // CHECK: comparison_direction = #vhlo<comparison_direction_v1 LE>
-@@ -58,6 +63,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_direction_lt"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_direction_lt(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     // CHECK: comparison_direction = #vhlo<comparison_direction_v1 LT>
-@@ -67,6 +73,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_type_notype"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_type_notype(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     comparison_direction = #stablehlo<comparison_direction EQ>
-@@ -76,6 +83,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_type_float"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_type_float(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     comparison_direction = #stablehlo<comparison_direction EQ>,
-@@ -86,6 +94,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_type_totalorder"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_type_totalorder(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     comparison_direction = #stablehlo<comparison_direction EQ>,
-@@ -96,6 +105,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_type_signed"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_type_signed(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     comparison_direction = #stablehlo<comparison_direction EQ>,
-@@ -106,6 +116,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_type_unsigned"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_type_unsigned(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     comparison_direction = #stablehlo<comparison_direction EQ>,
-@@ -118,6 +129,7 @@
- // ConvDimensionNumbers aka #stablehlo.conv is covered below.
- 
- // CHECK-LABEL: "attr_custom_call_api_version_unspecified"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_custom_call_api_version_unspecified(%arg0: tensor<f32>) -> tensor<f32> {
-   %0 = "stablehlo.custom_call"(%arg0) {
-     call_target_name = "foo",
-@@ -128,6 +140,7 @@
- }
- 
- // CHECK-LABEL: "attr_custom_call_api_version_original"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_custom_call_api_version_original(%arg0: tensor<f32>) -> tensor<f32> {
-   %0 = "stablehlo.custom_call"(%arg0) {
-     call_target_name = "foo",
-@@ -138,6 +151,7 @@
- }
- 
- // CHECK-LABEL: "attr_custom_call_api_version_status_returning"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_custom_call_api_version_status_returning(%arg0: tensor<f32>) -> tensor<f32> {
-   %0 = "stablehlo.custom_call"(%arg0) {
-     call_target_name = "foo",
-@@ -148,6 +162,7 @@
- }
- 
- // CHECK-LABEL: "attr_custom_call_api_version_status_returning_unified"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_custom_call_api_version_status_returning_unified(%arg0: tensor<f32>) -> tensor<f32> {
-   %0 = "stablehlo.custom_call"(%arg0) {
-     call_target_name = "foo",
-@@ -166,6 +181,7 @@
- // DotDimensionNumbers aka #stablehlo.dot is covered below.
- 
- // CHECK-LABEL: "attr_fft_type_fft"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_fft_type_fft(%arg0: tensor<16xcomplex<f32>>) -> tensor<16xcomplex<f32>> {
-   %0 = "stablehlo.fft"(%arg0) {
-     // CHECK: fft_type = #vhlo<fft_type_v1 FFT>
-@@ -176,6 +192,7 @@
- }
- 
- // CHECK-LABEL: "attr_fft_type_ifft"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_fft_type_ifft(%arg0: tensor<16xcomplex<f32>>) -> tensor<16xcomplex<f32>> {
-   %0 = "stablehlo.fft"(%arg0) {
-     // CHECK: fft_type = #vhlo<fft_type_v1 IFFT>
-@@ -186,6 +203,7 @@
- }
- 
- // CHECK-LABEL: "attr_fft_type_rfft"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_fft_type_rfft(%arg0: tensor<16xf32>) -> tensor<9xcomplex<f32>> {
-   %0 = "stablehlo.fft"(%arg0) {
-     // CHECK: fft_type = #vhlo<fft_type_v1 RFFT>
-@@ -196,6 +214,7 @@
- }
- 
- // CHECK-LABEL: "attr_fft_type_irfft"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_fft_type_irfft(%arg0: tensor<9xcomplex<f32>>) -> tensor<16xf32> {
-   %0 = "stablehlo.fft"(%arg0) {
-     // CHECK: fft_type = #vhlo<fft_type_v1 IRFFT>
-@@ -208,6 +227,7 @@
- // GatherDimensionNumbers aka #stablehlo.gather is covered below.
- 
- // CHECK-LABEL: "attr_precision_config_default"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_precision_config_default(%arg0: tensor<8x16xf32>, %arg1: tensor<16x8xf32>) -> tensor<8x8xf32> {
-   %0 = "stablehlo.dot"(%arg0, %arg1) {
-     // CHECK: precision_config = #vhlo.array_v1<[#vhlo<precision_v1 DEFAULT>, #vhlo<precision_v1 DEFAULT>]>
-@@ -216,6 +236,7 @@
- }
- 
- // CHECK-LABEL: "attr_precision_config_high"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_precision_config_high(%arg0: tensor<8x16xf32>, %arg1: tensor<16x8xf32>) -> tensor<8x8xf32> {
-   %0 = "stablehlo.dot"(%arg0, %arg1) {
-     // CHECK: precision_config = #vhlo.array_v1<[#vhlo<precision_v1 HIGH>, #vhlo<precision_v1 HIGH>]>
-@@ -225,6 +246,7 @@
- }
- 
- // CHECK-LABEL: "attr_precision_config_highest"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_precision_config_highest(%arg0: tensor<8x16xf32>, %arg1: tensor<16x8xf32>) -> tensor<8x8xf32> {
-   %0 = "stablehlo.dot"(%arg0, %arg1) {
-     // CHECK: precision_config = #vhlo.array_v1<[#vhlo<precision_v1 HIGHEST>, #vhlo<precision_v1 HIGHEST>]>
-@@ -234,6 +256,7 @@
- }
- 
- // CHECK-LABEL: "attr_rng_algorithm_default"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_rng_algorithm_default(%arg0: tensor<f32>) -> (tensor<f32>, tensor<f32>) {
-   %0:2 = "stablehlo.rng_bit_generator"(%arg0) {
-     // CHECK: rng_algorithm = #vhlo<rng_algorithm_v1 DEFAULT>
-@@ -243,6 +266,7 @@
- }
- 
- // CHECK-LABEL: "attr_rng_algorithm_three_fry"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_rng_algorithm_three_fry(%arg0: tensor<f32>) -> (tensor<f32>, tensor<f32>) {
-   %0:2 = "stablehlo.rng_bit_generator"(%arg0) {
-     // CHECK: rng_algorithm = #vhlo<rng_algorithm_v1 THREE_FRY>
-@@ -252,6 +276,7 @@
- }
- 
- // CHECK-LABEL: "attr_rng_algorithm_philox"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_rng_algorithm_philox(%arg0: tensor<f32>) -> (tensor<f32>, tensor<f32>) {
-   %0:2 = "stablehlo.rng_bit_generator"(%arg0) {
-     // CHECK: rng_algorithm = #vhlo<rng_algorithm_v1 PHILOX>
-@@ -261,6 +286,7 @@
- }
- 
- // CHECK-LABEL: "attr_rng_distribution_uniform"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @attr_rng_distribution_uniform(%arg0: tensor<f32>, %arg1: tensor<f32>, %arg2: tensor<0xindex>) -> tensor<f32> {
-   %0 = "stablehlo.rng"(%arg0, %arg1, %arg2) {
-     // CHECK: rng_distribution = #vhlo<rng_distribution_v1 UNIFORM>
-@@ -270,6 +296,7 @@
- }
- 
- // CHECK-LABEL: "attr_rng_distribution_normal"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @attr_rng_distribution_normal(%arg0: tensor<f32>, %arg1: tensor<f32>, %arg2: tensor<0xindex>) -> tensor<f32> {
-   %0 = "stablehlo.rng"(%arg0, %arg1, %arg2) {
-     // CHECK: rng_distribution = #vhlo<rng_distribution_v1 NORMAL>
-@@ -281,6 +308,7 @@
- // ScatterDimensionNumbers aka #stablehlo.scatter is covered below.
- 
- // CHECK-LABEL: "attr_transpose_no_transpose"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_transpose_no_transpose(%arg0: tensor<16x16xf32>, %arg1: tensor<16x16xf32>) ->  tensor<16x16xf32> {
-   %0 = "stablehlo.triangular_solve"(%arg0, %arg1) {
-     left_side = true,
-@@ -293,6 +321,7 @@
- }
- 
- // CHECK-LABEL: "attr_transpose_transpose"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_transpose_transpose(%arg0: tensor<16x16xf32>, %arg1: tensor<16x16xf32>) ->  tensor<16x16xf32> {
-   %0 = "stablehlo.triangular_solve"(%arg0, %arg1) {
-     left_side = true,
-@@ -305,6 +334,7 @@
- }
- 
- // CHECK-LABEL: "attr_transpose_adjoint"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_transpose_adjoint(%arg0: tensor<16x16xf32>, %arg1: tensor<16x16xf32>) ->  tensor<16x16xf32> {
-   %0 = "stablehlo.triangular_solve"(%arg0, %arg1) {
-     left_side = true,
-@@ -319,10 +349,9 @@
- // TypeExtensionsAttr aka #stablehlo.type_extensions is covered below.
- 
- // CHECK-LABEL: "attr_type_extensions_bounds"
--func.func @attr_type_extensions_bounds(
--    %arg0: tensor<?x?xf32, #stablehlo.type_extensions<bounds = [16, ?]>>)
--    -> tensor<?x?xf32, #stablehlo.type_extensions<bounds = [16, ?]>> {
--  // CHECK: "vhlo.return_v1"(%arg0) : (!vhlo.tensor_v1<?x?x!vhlo.f32_v1, #vhlo.type_extensions_v1<bounds = [16, ?]>>) -> ()
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @attr_type_extensions_bounds(%arg0: tensor<?x?xf32, #stablehlo.type_extensions<bounds = [16, ?]>>) -> tensor<?x?xf32, #stablehlo.type_extensions<bounds = [16, ?]>> {
-+  // CHECK: "vhlo.return_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<?x?x!vhlo.f32_v1, #vhlo.type_extensions_v1<bounds = [16, ?]>>) -> ()
-   func.return %arg0 : tensor<?x?xf32, #stablehlo.type_extensions<bounds = [16, ?]>>
- }
- 
-@@ -330,8 +359,9 @@
- // ============ DEFAULTS ============
- 
- // CHECK-LABEL: "default_all_gather"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_all_gather(%arg0: tensor<16x8xf32>) -> tensor<16x16xf32> {
--  //               CHECK: "vhlo.all_gather_v1"(%arg0) <{
-+  //               CHECK: "vhlo.all_gather_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   all_gather_dim = #vhlo.integer_v1<1 : i64>
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0], [1]]> : tensor<2x1xi64>>,
-@@ -345,8 +375,9 @@
- }
- 
- // CHECK-LABEL: "default_all_reduce"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_all_reduce(%arg0: tensor<f32>) -> tensor<f32> {
--  //               CHECK: "vhlo.all_reduce_v1"(%arg0)
-+  //               CHECK: "vhlo.all_reduce_v1"(%[[ARG0]])
-   //          CHECK-SAME: <{
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0], [1]]> : tensor<2x1xi64>>,
-@@ -368,8 +399,9 @@
- }
- 
- // CHECK-LABEL: "default_all_to_all"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_all_to_all(%arg0: tensor<4x16xf32>) -> tensor<16x4xf32> {
--  //               CHECK: "vhlo.all_to_all_v1"(%arg0) <{
-+  //               CHECK: "vhlo.all_to_all_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   //          CHECK-SAME:   concat_dimension = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0, 1, 2, 3]]> : tensor<1x4xi64>>,
-@@ -386,8 +418,9 @@
- }
- 
- // CHECK-LABEL: "default_cholesky"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_cholesky(%arg0: tensor<1x16x16xf32>) -> tensor<1x16x16xf32> {
--  //      CHECK: "vhlo.cholesky_v1"(%arg0) <{
-+  //      CHECK: "vhlo.cholesky_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   lower = #vhlo.bool_v1<false>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<1x16x16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<1x16x16x!vhlo.f32_v1>
-   %0 = "stablehlo.cholesky"(%arg0) : (tensor<1x16x16xf32>) -> tensor<1x16x16xf32>
-@@ -395,8 +428,9 @@
- }
- 
- // CHECK-LABEL: "default_collective_permute"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_collective_permute(%arg0: tensor<16x8xf32>) -> tensor<16x8xf32> {
--  //               CHECK: "vhlo.collective_permute_v1"(%arg0) <{
-+  //               CHECK: "vhlo.collective_permute_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME{LITERAL}:   source_target_pairs = #vhlo.tensor_v1<dense<[[0, 1], [1, 2], [2, 3]]> : tensor<3x2xi64>>
-   //          CHECK-SAME: }> : (!vhlo.tensor_v1<16x8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x8x!vhlo.f32_v1>
-@@ -407,8 +441,9 @@
- }
- 
- // CHECK-LABEL: "default_collective_broadcast"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_collective_broadcast(%arg0: tensor<16x8xf32>) -> tensor<16x8xf32> {
--  //               CHECK: "vhlo.collective_broadcast_v1"(%arg0) <{
-+  //               CHECK: "vhlo.collective_broadcast_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0, 1]]> : tensor<1x2xi64>>
-   //          CHECK-SAME: }> : (!vhlo.tensor_v1<16x8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x8x!vhlo.f32_v1>
-@@ -419,8 +454,9 @@
- }
- 
- // CHECK-LABEL: "default_compare"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @default_compare(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
--  //      CHECK: "vhlo.compare_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.compare_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   compare_type = #vhlo<comparison_type_v1 NOTYPE>,
-   // CHECK-SAME:   comparison_direction = #vhlo<comparison_direction_v1 EQ>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-@@ -431,8 +467,9 @@
- }
- 
- // CHECK-LABEL: "default_convolution"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @default_convolution(%arg0: tensor<1x8x8x207xf32>, %arg1: tensor<3x3x207x16xf32>) -> tensor<1x6x6x16xf32> {
--  //      CHECK: "vhlo.convolution_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.convolution_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   batch_group_count = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   feature_group_count = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   input_batch_dimension = #vhlo.integer_v1<0 : i64>,
-@@ -460,8 +497,9 @@
- }
- 
- // CHECK-LABEL: "default_custom_call"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_custom_call(%arg0: tensor<f32>) -> tensor<f32> {
--  //      CHECK: "vhlo.custom_call_v1"(%arg0) <{
-+  //      CHECK: "vhlo.custom_call_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   api_version = #vhlo<api_version_v1 API_VERSION_ORIGINAL>,
-   // CHECK-SAME:   backend_config = #vhlo.string_v1<"">,
-   // CHECK-SAME:   call_target_name = #vhlo.string_v1<"foo">,
-@@ -478,8 +516,9 @@
- }
- 
- // CHECK-LABEL: "default_dot_general"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @default_dot_general(%arg0: tensor<8x8x16xf32>, %arg1: tensor<8x16x8xf32>) -> tensor<8x8x8xf32> {
--  //      CHECK: "vhlo.dot_general_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.dot_general_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   lhs_batching_dimensions = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>,
-   // CHECK-SAME:   lhs_contracting_dimensions = #vhlo.tensor_v1<dense<2> : tensor<1xi64>>,
-   // CHECK-SAME:   precision_config = #vhlo.array_v1<[#vhlo<precision_v1 DEFAULT>, #vhlo<precision_v1 DEFAULT>]>,
-@@ -498,8 +537,9 @@
- }
- 
- // CHECK-LABEL: "default_dot"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @default_dot(%arg0: tensor<8x16xf32>, %arg1: tensor<16x8xf32>) -> tensor<8x8xf32> {
--  //      CHECK: "vhlo.dot_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.dot_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   precision_config = #vhlo.array_v1<[#vhlo<precision_v1 DEFAULT>, #vhlo<precision_v1 DEFAULT>]>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<8x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<8x8x!vhlo.f32_v1>
-   %0 = "stablehlo.dot"(%arg0, %arg1) : (tensor<8x16xf32>, tensor<16x8xf32>) -> tensor<8x8xf32>
-@@ -507,8 +547,9 @@
- }
- 
- // CHECK-LABEL: "default_dynamic_broadcast_in_dim"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @default_dynamic_broadcast_in_dim(%arg0: tensor<?x?xf32>, %arg1: tensor<2xindex>) -> tensor<?x?xf32> {
--  //      CHECK: "vhlo.dynamic_broadcast_in_dim_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.dynamic_broadcast_in_dim_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   broadcast_dimensions = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-   // CHECK-SAME:   known_expanding_dimensions = #vhlo.tensor_v1<dense<> : tensor<0xi64>>,
-   // CHECK-SAME:   known_nonexpanding_dimensions = #vhlo.tensor_v1<dense<> : tensor<0xi64>>
-@@ -520,8 +561,9 @@
- }
- 
- // CHECK-LABEL: "default_dynamic_conv"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @default_dynamic_conv(%arg0: tensor<1x8x8x207xf32>, %arg1: tensor<3x3x207x16xf32>, %arg2: tensor<4xi32>) -> tensor<1x?x?x16xf32> {
--  //      CHECK: "vhlo.dynamic_conv_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.dynamic_conv_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   batch_group_count = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   feature_group_count = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   input_batch_dimension = #vhlo.integer_v1<0 : i64>,
-@@ -549,8 +591,9 @@
- }
- 
- // CHECK-LABEL: "default_dynamic_gather"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @default_dynamic_gather(%arg0 : tensor<2x4x9xf32>, %arg1 : tensor<1x5x2xi32>, %arg2 : tensor<3xi32>) -> tensor<1x5x8xf32> {
--  //      CHECK: "vhlo.dynamic_gather_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.dynamic_gather_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   collapsed_slice_dims = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-   // CHECK-SAME:   index_vector_dim = #vhlo.integer_v1<2 : i64>,
-   // CHECK-SAME:   indices_are_sorted = #vhlo.bool_v1<false>,
-@@ -576,15 +619,16 @@
-   // CHECK-SAME:   sym_name = #vhlo.string_v1<"default_func">,
-   // CHECK-SAME:   sym_visibility = #vhlo.string_v1<"">
-   // CHECK-SAME: }> ({
--  // CHECK-NEXT:   ^[[BB:bb.*]](%arg0: !vhlo.tensor_v1<!vhlo.f32_v1>):
--  // CHECK-NEXT:     "vhlo.return_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-+  // CHECK-NEXT:   ^[[BB:bb.*]](%[[ARG0:.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
-+  // CHECK-NEXT:     "vhlo.return_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-   // CHECK-NEXT: }) : () -> ()
-   func.return %arg0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "dynamic_gather"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @dynamic_gather(%arg0 : tensor<2x4x9xf32>, %arg1 : tensor<1x5x2xi32>) -> tensor<1x5x1xf32> {
--  //      CHECK: "vhlo.gather_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.gather_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   collapsed_slice_dims = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-   // CHECK-SAME:   index_vector_dim = #vhlo.integer_v1<2 : i64>,
-   // CHECK-SAME:   indices_are_sorted = #vhlo.bool_v1<false>,
-@@ -605,8 +649,9 @@
- }
- 
- // CHECK-LABEL: "default_infeed"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_infeed(%arg0: !stablehlo.token) -> (tensor<f32>, !stablehlo.token) {
--  //               CHECK: "vhlo.infeed_v1"(%arg0) <{
-+  //               CHECK: "vhlo.infeed_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   infeed_config = #vhlo.string_v1<"">,
-   // CHECK-SAME{LITERAL}:   layout = #vhlo.array_v1<[]>
-   //          CHECK-SAME: }> : (!vhlo.token_v1) -> (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.token_v1)
-@@ -615,8 +660,9 @@
- }
- 
- // CHECK-LABEL: "default_outfeed"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @default_outfeed(%arg0: tensor<f32>, %arg1: !stablehlo.token) -> !stablehlo.token {
--  //      CHECK: "vhlo.outfeed_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.outfeed_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   outfeed_config = #vhlo.string_v1<"">
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.token_v1) -> !vhlo.token_v1
-   %0 = "stablehlo.outfeed"(%arg0, %arg1) : (tensor<f32>, !stablehlo.token) -> !stablehlo.token
-@@ -624,8 +670,9 @@
- }
- 
- // CHECK-LABEL: "default_recv"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_recv(%arg0: !stablehlo.token) -> (tensor<f32>, !stablehlo.token) {
--  //      CHECK: "vhlo.recv_v1"(%arg0) <{
-+  //      CHECK: "vhlo.recv_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME:   channel_type = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   is_host_transfer = #vhlo.bool_v1<false>
-@@ -637,8 +684,9 @@
- }
- 
- // CHECK-LABEL: "default_send"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @default_send(%arg0: tensor<f32>, %arg1: !stablehlo.token) -> !stablehlo.token {
--  //      CHECK: "vhlo.send_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.send_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME:   channel_type = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   is_host_transfer = #vhlo.bool_v1<false>
-@@ -650,8 +698,9 @@
- }
- 
- // CHECK-LABEL: "default_reduce_scatter"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_reduce_scatter(%arg0: tensor<16xf32>) -> tensor<16xf32> {
--  //               CHECK: "vhlo.reduce_scatter_v1"(%arg0) <{
-+  //               CHECK: "vhlo.reduce_scatter_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0], [1]]> : tensor<2x1xi64>>,
-   //          CHECK-SAME:   scatter_dimension = #vhlo.integer_v1<0 : i64>
-@@ -673,8 +722,9 @@
- }
- 
- // CHECK-LABEL: "default_reduce_window"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @default_reduce_window(%arg0: tensor<2x17x31x7xf32>, %arg1: tensor<f32>) -> tensor<2x16x30x7xf32> {
--  //               CHECK: "vhlo.reduce_window_v1"(%arg0, %arg1)  <{
-+  //               CHECK: "vhlo.reduce_window_v1"(%[[ARG0]], %[[ARG1]])  <{
-   //          CHECK-SAME:   base_dilations = #vhlo.tensor_v1<dense<1> : tensor<4xi64>>,
-   // CHECK-SAME{LITERAL}:   padding = #vhlo.tensor_v1<dense<0> : tensor<4x2xi64>>,
-   //          CHECK-SAME:   window_dilations = #vhlo.tensor_v1<dense<1> : tensor<4xi64>>,
-@@ -696,8 +746,9 @@
- }
- 
- // CHECK-LABEL: "default_scatter"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @default_scatter(%arg0: tensor<200x100x300xf32>, %arg1: tensor<10x2xi32>, %arg2: tensor<10x300xf32>) -> tensor<200x100x300xf32> {
--  //      CHECK: "vhlo.scatter_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.scatter_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   index_vector_dim = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   indices_are_sorted = #vhlo.bool_v1<false>,
-   // CHECK-SAME:   inserted_window_dims = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-@@ -725,8 +776,9 @@
- }
- 
- // CHECK-LABEL: "default_select_and_scatter"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @default_select_and_scatter(%arg0: tensor<10x24x24x64xf32>, %arg1: tensor<10x23x23x64xf32>, %arg2: tensor<f32>) -> tensor<10x24x24x64xf32> {
--  //      CHECK: "vhlo.select_and_scatter_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.select_and_scatter_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   padding = #vhlo.tensor_v1<dense<0> : tensor<4x2xi64>>,
-   // CHECK-SAME:   window_dimensions = #vhlo.tensor_v1<dense<[1, 2, 2, 1]> : tensor<4xi64>>,
-   // CHECK-SAME:   window_strides = #vhlo.tensor_v1<dense<1> : tensor<4xi64>>
-@@ -754,8 +806,9 @@
- }
- 
- // CHECK-LABEL: "default_sort"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_sort(%arg0: tensor<16xf32>) -> tensor<16xf32> {
--  //      CHECK: "vhlo.sort_v1"(%arg0) <{
-+  //      CHECK: "vhlo.sort_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   dimension = #vhlo.integer_v1<-1 : i64>
-   // CHECK-SAME:   is_stable = #vhlo.bool_v1<false>
-   // CHECK-SAME: }> ({
-@@ -774,29 +827,33 @@
- // ============ OPS ============
- 
- // CHECK-LABEL: "op_abs"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_abs(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.abs_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.abs_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.abs"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_add"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_add(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_after_all"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_after_all(%arg0: !stablehlo.token) -> !stablehlo.token {
--  // CHECK: "vhlo.after_all_v1"(%arg0) : (!vhlo.token_v1) -> !vhlo.token_v1
-+  // CHECK: "vhlo.after_all_v1"(%[[ARG0]]) : (!vhlo.token_v1) -> !vhlo.token_v1
-   %0 = "stablehlo.after_all"(%arg0) : (!stablehlo.token) -> !stablehlo.token
-   func.return %0 : !stablehlo.token
- }
- 
- // CHECK-LABEL: "op_all_gather"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_all_gather(%arg0: tensor<16x8xf32>) -> tensor<16x16xf32> {
--  //               CHECK: "vhlo.all_gather_v1"(%arg0) <{
-+  //               CHECK: "vhlo.all_gather_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   all_gather_dim = #vhlo.integer_v1<1 : i64>
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0], [1]]> : tensor<2x1xi64>>,
-@@ -812,8 +869,9 @@
- }
- 
- // CHECK-LABEL: "op_all_reduce"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_all_reduce(%arg0: tensor<f32>) -> tensor<f32> {
--  //               CHECK: "vhlo.all_reduce_v1"(%arg0) <{
-+  //               CHECK: "vhlo.all_reduce_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0], [1]]> : tensor<2x1xi64>>,
-   //          CHECK-SAME:   use_global_device_ids = #vhlo.bool_v1<true>
-@@ -836,7 +894,7 @@
- 
- // CHECK-LABEL: "op_all_reduce_with_promotable_types"
- func.func @op_all_reduce_with_promotable_types(%operand: tensor<f32>) -> tensor<f64> {
--  //  CHECK: "vhlo.all_reduce_v1"(%arg0)
-+  //  CHECK: "vhlo.all_reduce_v1"(%[[ARG0:.*]])
-   //  CHECK:   ^[[BB:bb.*]](%[[ARG1:arg.*]]: !vhlo.tensor_v1<!vhlo.f64_v1>, %[[ARG2:arg.*]]: !vhlo.tensor_v1<!vhlo.f64_v1>):
-   //  CHECK:     "vhlo.return_v1"(%[[VAL1:.*]]) : (!vhlo.tensor_v1<!vhlo.f64_v1>) -> ()
-   //  CHECK: }) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f64_v1>
-@@ -854,8 +912,9 @@
- }
- 
- // CHECK-LABEL: "op_all_to_all"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_all_to_all(%arg0: tensor<4x16xf32>) -> tensor<16x4xf32> {
--  //               CHECK: "vhlo.all_to_all_v1"(%arg0) <{
-+  //               CHECK: "vhlo.all_to_all_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<1 : i64>,
-   //          CHECK-SAME:   concat_dimension = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0, 1, 2, 3]]> : tensor<1x4xi64>>,
-@@ -873,22 +932,25 @@
- }
- 
- // CHECK-LABEL: "op_and"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_and(%arg0: tensor<i1>, %arg1: tensor<i1>) -> tensor<i1> {
--  // CHECK: "vhlo.and_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-+  // CHECK: "vhlo.and_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-   %0 = "stablehlo.and"(%arg0, %arg1) : (tensor<i1>, tensor<i1>) -> tensor<i1>
-   func.return %0 : tensor<i1>
- }
- 
- // CHECK-LABEL: "op_atan2"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_atan2(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.atan2_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.atan2_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.atan2"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_batch_norm_grad"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}}, %[[ARG3:.*]]: {{.*}}, %[[ARG4:.*]]: {{.*}})
- func.func @op_batch_norm_grad(%arg0: tensor<16x16x16x16xf32>, %arg1: tensor<16xf32>, %arg2: tensor<16xf32>, %arg3: tensor<16xf32>, %arg4: tensor<16x16x16x16xf32>) -> (tensor<16x16x16x16xf32>, tensor<16xf32>, tensor<16xf32>) {
--  //      CHECK: "vhlo.batch_norm_grad_v1"(%arg0, %arg1, %arg2, %arg3, %arg4) <{
-+  //      CHECK: "vhlo.batch_norm_grad_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]], %[[ARG3]], %[[ARG4]]) <{
-   // CHECK-SAME:   epsilon = #vhlo.float_v1<1.000000e-03 : !vhlo.f32_v1>,
-   // CHECK-SAME:   feature_index = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<16x16x16x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x16x16x16x!vhlo.f32_v1>) -> (!vhlo.tensor_v1<16x16x16x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>)
-@@ -900,8 +962,9 @@
- }
- 
- // CHECK-LABEL: "op_batch_norm_inference"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}}, %[[ARG3:.*]]: {{.*}}, %[[ARG4:.*]]: {{.*}})
- func.func @op_batch_norm_inference(%arg0: tensor<16x16x16x16xf32>, %arg1: tensor<16xf32>, %arg2: tensor<16xf32>, %arg3: tensor<16xf32>, %arg4: tensor<16xf32>) -> tensor<16x16x16x16xf32> {
--  //      CHECK: "vhlo.batch_norm_inference_v1"(%arg0, %arg1, %arg2, %arg3, %arg4) <{
-+  //      CHECK: "vhlo.batch_norm_inference_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]], %[[ARG3]], %[[ARG4]]) <{
-   // CHECK-SAME:   epsilon = #vhlo.float_v1<1.000000e-03 : !vhlo.f32_v1>,
-   // CHECK-SAME:   feature_index = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<16x16x16x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x16x16x16x!vhlo.f32_v1>
-@@ -913,8 +976,9 @@
- }
- 
- // CHECK-LABEL: "op_batch_norm_training"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_batch_norm_training(%arg0: tensor<16x16x16x16xf32>, %arg1: tensor<16xf32>, %arg2: tensor<16xf32>) -> (tensor<16x16x16x16xf32>, tensor<16xf32>, tensor<16xf32>) {
--  //      CHECK: "vhlo.batch_norm_training_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.batch_norm_training_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   epsilon = #vhlo.float_v1<1.000000e-03 : !vhlo.f32_v1>,
-   // CHECK-SAME:   feature_index = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<16x16x16x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>) -> (!vhlo.tensor_v1<16x16x16x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>)
-@@ -926,15 +990,17 @@
- }
- 
- // CHECK-LABEL: "op_bitcast_convert"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_bitcast_convert(%arg0: tensor<i32>) -> tensor<f32> {
--  // CHECK: "vhlo.bitcast_convert_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.bitcast_convert_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.bitcast_convert"(%arg0) : (tensor<i32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_broadcast_in_dim"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_broadcast_in_dim(%arg0: tensor<16xf32>) -> tensor<16x16xf32> {
--  //      CHECK: "vhlo.broadcast_in_dim_v1"(%arg0) <{
-+  //      CHECK: "vhlo.broadcast_in_dim_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   broadcast_dimensions = #vhlo.tensor_v1<dense<1> : tensor<1xi64>>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x16x!vhlo.f32_v1>
-   %0 = "stablehlo.broadcast_in_dim"(%arg0) {
-@@ -944,8 +1010,9 @@
- }
- 
- // CHECK-LABEL: "op_broadcast"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_broadcast(%arg0: tensor<16xf32>) -> tensor<16x16xf32> {
--  //      CHECK: "vhlo.broadcast_v1"(%arg0) <{
-+  //      CHECK: "vhlo.broadcast_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   broadcast_sizes = #vhlo.tensor_v1<dense<16> : tensor<1xi64>>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x16x!vhlo.f32_v1>
-   %0 = "stablehlo.broadcast"(%arg0) {
-@@ -955,9 +1022,10 @@
- }
- 
- // CHECK-LABEL: "op_case"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_case(%arg0: tensor<i32>, %arg1: tensor<f32>) -> tensor<f32> {
--  //      CHECK: "vhlo.case_v1"(%arg0) ({
--  // CHECK-NEXT:   "vhlo.return_v1"(%arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-+  //      CHECK: "vhlo.case_v1"(%[[ARG0]]) ({
-+  // CHECK-NEXT:   "vhlo.return_v1"(%[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-   // CHECK-NEXT: }) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.case"(%arg0) ({
-     "stablehlo.return"(%arg1) : (tensor<f32>) -> ()
-@@ -966,22 +1034,25 @@
- }
- 
- // CHECK-LABEL: "op_cbrt"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_cbrt(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.cbrt_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.cbrt_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.cbrt"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_ceil"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_ceil(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.ceil_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.ceil_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.ceil"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_cholesky"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_cholesky(%arg0: tensor<1x16x16xf32>) -> tensor<1x16x16xf32> {
--  //      CHECK: "vhlo.cholesky_v1"(%arg0) <{
-+  //      CHECK: "vhlo.cholesky_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   lower = #vhlo.bool_v1<true>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<1x16x16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<1x16x16x!vhlo.f32_v1>
-   %0 = "stablehlo.cholesky"(%arg0) {
-@@ -991,22 +1062,25 @@
- }
- 
- // CHECK-LABEL: "op_clamp"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_clamp(%arg0: tensor<f32>, %arg1: tensor<f32>, %arg2: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.clamp_v1"(%arg0, %arg1, %arg2) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.clamp_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.clamp"(%arg0, %arg1, %arg2) : (tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_count_leading_zeros"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_count_leading_zeros(%arg0: tensor<i32>) -> tensor<i32> {
--  // CHECK: "vhlo.count_leading_zeros_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-+  // CHECK: "vhlo.count_leading_zeros_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-   %0 = "stablehlo.count_leading_zeros"(%arg0) : (tensor<i32>) -> tensor<i32>
-   func.return %0 : tensor<i32>
- }
- 
- // CHECK-LABEL: "op_collective_permute"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_collective_permute(%arg0: tensor<16x8xf32>) -> tensor<16x8xf32> {
--  //               CHECK: "vhlo.collective_permute_v1"(%arg0) <{
-+  //               CHECK: "vhlo.collective_permute_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME{LITERAL}:   source_target_pairs = #vhlo.tensor_v1<dense<[[0, 1], [1, 2], [2, 3]]> : tensor<3x2xi64>>
-   //          CHECK-SAME: }> : (!vhlo.tensor_v1<16x8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x8x!vhlo.f32_v1>
-@@ -1018,8 +1092,9 @@
- }
- 
- // CHECK-LABEL: "op_compare"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_compare(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
--  //      CHECK: "vhlo.compare_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.compare_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   compare_type = #vhlo<comparison_type_v1 TOTALORDER>,
-   // CHECK-SAME:   comparison_direction = #vhlo<comparison_direction_v1 EQ>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-@@ -1031,15 +1106,17 @@
- }
- 
- // CHECK-LABEL: "op_complex"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_complex(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<complex<f32>> {
--  // CHECK: "vhlo.complex_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>
-+  // CHECK: "vhlo.complex_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>
-   %0 = "stablehlo.complex"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<complex<f32>>
-   func.return %0 : tensor<complex<f32>>
- }
- 
- // CHECK-LABEL: "op_concatenate"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_concatenate(%arg0: tensor<8xf32>, %arg1: tensor<8xf32>) -> tensor<16xf32> {
--  //      CHECK: "vhlo.concatenate_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.concatenate_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   dimension = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<8x!vhlo.f32_v1>, !vhlo.tensor_v1<8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x!vhlo.f32_v1>
-   %0 = "stablehlo.concatenate"(%arg0, %arg1) {
-@@ -1049,6 +1126,7 @@
- }
- 
- // CHECK-LABEL: "op_constant"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_constant(%arg0: tensor<f32>) -> tensor<f32> {
-   //      CHECK: "vhlo.constant_v1"() <{
-   // CHECK-SAME:   value = #vhlo.tensor_v1<dense<0.000000e+00> : tensor<f32>>
-@@ -1060,15 +1138,17 @@
- }
- 
- // CHECK-LABEL: "op_convert"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_convert(%arg0: tensor<i32>) -> tensor<f32> {
--  // CHECK: "vhlo.convert_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.convert_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.convert"(%arg0) : (tensor<i32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_convolution"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_convolution(%arg0: tensor<1x8x8x207xf32>, %arg1: tensor<3x3x207x16xf32>) -> tensor<1x7x7x16xf32> {
--  //      CHECK: "vhlo.convolution_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.convolution_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   batch_group_count = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   feature_group_count = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   input_batch_dimension = #vhlo.integer_v1<0 : i64>,
-@@ -1102,8 +1182,9 @@
- }
- 
- // CHECK-LABEL: "op_cosine"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_cosine(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.cosine_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.cosine_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.cosine"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
-@@ -1116,8 +1197,9 @@
- }
- 
- // CHECK-LABEL: "op_cross_replica_sum"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_cross_replica_sum(%arg0: tensor<f32>) -> tensor<f32> {
--  //               CHECK: "vhlo.cross-replica-sum_v1"(%arg0) <{
-+  //               CHECK: "vhlo.cross-replica-sum_v1"(%[[ARG0]]) <{
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0], [1]]> : tensor<2x1xi64>>
-   //          CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.cross-replica-sum"(%arg0) {
-@@ -1127,8 +1209,9 @@
- }
- 
- // CHECK-LABEL: "op_custom_call"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_custom_call(%arg0: tensor<f32>) -> tensor<f32> {
--  //      CHECK: "vhlo.custom_call_v1"(%arg0) <{
-+  //      CHECK: "vhlo.custom_call_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   api_version = #vhlo<api_version_v1 API_VERSION_STATUS_RETURNING>,
-   // CHECK-SAME:   backend_config = #vhlo.string_v1<"\08\03\1A\02">,
-   // CHECK-SAME:   call_target_name = #vhlo.string_v1<"foo">,
-@@ -1159,15 +1242,17 @@
- }
- 
- // CHECK-LABEL: "op_divide"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_divide(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.divide_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.divide_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.divide"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_dot_general"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_dot_general(%arg0: tensor<8x8x16xf32>, %arg1: tensor<8x16x8xf32>) -> tensor<8x8x8xf32> {
--  //      CHECK: "vhlo.dot_general_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.dot_general_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   lhs_batching_dimensions = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>,
-   // CHECK-SAME:   lhs_contracting_dimensions = #vhlo.tensor_v1<dense<2> : tensor<1xi64>>,
-   // CHECK-SAME:   precision_config = #vhlo.array_v1<[#vhlo<precision_v1 HIGHEST>, #vhlo<precision_v1 HIGHEST>]>,
-@@ -1187,8 +1272,9 @@
- }
- 
- // CHECK-LABEL: "op_dot"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_dot(%arg0: tensor<8x16xf32>, %arg1: tensor<16x8xf32>) -> tensor<8x8xf32> {
--  //      CHECK: "vhlo.dot_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.dot_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   precision_config = #vhlo.array_v1<[#vhlo<precision_v1 HIGHEST>, #vhlo<precision_v1 HIGHEST>]>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<8x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<8x8x!vhlo.f32_v1>
-   %0 = "stablehlo.dot"(%arg0, %arg1) {
-@@ -1198,8 +1284,9 @@
- }
- 
- // CHECK-LABEL: "op_dynamic_broadcast_in_dim"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_dynamic_broadcast_in_dim(%arg0: tensor<?x?xf32>, %arg1: tensor<2xindex>) -> tensor<?x?xf32> {
--  //      CHECK: "vhlo.dynamic_broadcast_in_dim_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.dynamic_broadcast_in_dim_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   broadcast_dimensions = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-   // CHECK-SAME:   known_expanding_dimensions = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>,
-   // CHECK-SAME:   known_nonexpanding_dimensions = #vhlo.tensor_v1<dense<1> : tensor<1xi64>>
-@@ -1213,8 +1300,9 @@
- }
- 
- // CHECK-LABEL: "op_dynamic_conv"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_dynamic_conv(%arg0: tensor<1x8x8x207xf32>, %arg1: tensor<3x3x207x16xf32>, %arg2: tensor<4xi32>) -> tensor<1x?x?x16xf32> {
--  //      CHECK: "vhlo.dynamic_conv_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.dynamic_conv_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   batch_group_count = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   feature_group_count = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   input_batch_dimension = #vhlo.integer_v1<0 : i64>,
-@@ -1248,8 +1336,9 @@
- }
- 
- // CHECK-LABEL: "op_dynamic_gather"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_dynamic_gather(%arg0 : tensor<2x4x9xf32>, %arg1 : tensor<1x5x2xi32>, %arg2 : tensor<3xi32>) -> tensor<1x5x8xf32> {
--  //      CHECK: "vhlo.dynamic_gather_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.dynamic_gather_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   collapsed_slice_dims = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-   // CHECK-SAME:   index_vector_dim = #vhlo.integer_v1<2 : i64>,
-   // CHECK-SAME:   indices_are_sorted = #vhlo.bool_v1<true>,
-@@ -1269,8 +1358,9 @@
- }
- 
- // CHECK-LABEL: "op_dynamic_iota"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_dynamic_iota(%arg0: tensor<1xindex>) -> tensor<?xf32> {
--  //      CHECK: "vhlo.dynamic_iota_v1"(%arg0) <{
-+  //      CHECK: "vhlo.dynamic_iota_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   iota_dimension = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<1x!vhlo.index_v1>) -> !vhlo.tensor_v1<?x!vhlo.f32_v1>
-   %0 = "stablehlo.dynamic_iota"(%arg0) {
-@@ -1280,22 +1370,25 @@
- }
- 
- // CHECK-LABEL: "op_dynamic_pad"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}}, %[[ARG3:.*]]: {{.*}}, %[[ARG4:.*]]: {{.*}})
- func.func @op_dynamic_pad(%arg0: tensor<?xf32>, %arg1: tensor<f32>, %arg2: tensor<1xindex>, %arg3: tensor<1xindex>, %arg4: tensor<1xindex>) -> tensor<?xf32> {
--  // CHECK: "vhlo.dynamic_pad_v1"(%arg0, %arg1, %arg2, %arg3, %arg4) : (!vhlo.tensor_v1<?x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>) -> !vhlo.tensor_v1<?x!vhlo.f32_v1>
-+  // CHECK: "vhlo.dynamic_pad_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]], %[[ARG3]], %[[ARG4]]) : (!vhlo.tensor_v1<?x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>) -> !vhlo.tensor_v1<?x!vhlo.f32_v1>
-   %0 = "stablehlo.dynamic_pad"(%arg0, %arg1, %arg2, %arg3, %arg4) : (tensor<?xf32>, tensor<f32>, tensor<1xindex>, tensor<1xindex>, tensor<1xindex>) -> tensor<?xf32>
-   func.return %0 : tensor<?xf32>
- }
- 
- // CHECK-LABEL: "op_dynamic_reshape"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_dynamic_reshape(%arg0: tensor<16xf32>, %arg1: tensor<2xindex>) -> tensor<?x?xf32> {
--  // CHECK: "vhlo.dynamic_reshape_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<2x!vhlo.index_v1>) -> !vhlo.tensor_v1<?x?x!vhlo.f32_v1>
-+  // CHECK: "vhlo.dynamic_reshape_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<2x!vhlo.index_v1>) -> !vhlo.tensor_v1<?x?x!vhlo.f32_v1>
-   %0 = "stablehlo.dynamic_reshape"(%arg0, %arg1) : (tensor<16xf32>, tensor<2xindex>) -> tensor<?x?xf32>
-   func.return %0 : tensor<?x?xf32>
- }
- 
- // CHECK-LABEL: "op_dynamic_slice"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_dynamic_slice(%arg0: tensor<16xf32>, %arg1: tensor<i64>) -> tensor<4xf32> {
--  //      CHECK: "vhlo.dynamic_slice_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.dynamic_slice_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   slice_sizes = #vhlo.tensor_v1<dense<4> : tensor<1xi64>>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.i64_v1>) -> !vhlo.tensor_v1<4x!vhlo.f32_v1>
-   %0 = "stablehlo.dynamic_slice"(%arg0, %arg1) {
-@@ -1305,15 +1398,17 @@
- }
- 
- // CHECK-LABEL: "op_dynamic_update_slice"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_dynamic_update_slice(%arg0: tensor<16xf32>, %arg1: tensor<4xf32>, %arg2: tensor<i64>) -> tensor<16xf32> {
--  // CHECK: "vhlo.dynamic_update_slice_v1"(%arg0, %arg1, %arg2) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<4x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.i64_v1>) -> !vhlo.tensor_v1<16x!vhlo.f32_v1>
-+  // CHECK: "vhlo.dynamic_update_slice_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<4x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.i64_v1>) -> !vhlo.tensor_v1<16x!vhlo.f32_v1>
-   %0 = "stablehlo.dynamic_update_slice"(%arg0, %arg1, %arg2) : (tensor<16xf32>, tensor<4xf32>, tensor<i64>) -> tensor<16xf32>
-   func.return %0 : tensor<16xf32>
- }
- 
- // CHECK-LABEL: "op_einsum"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_einsum(%arg0: tensor<8x16xf32>, %arg1: tensor<16x8xf32>) -> tensor<8x8xf32> {
--  //      CHECK: "vhlo.einsum_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.einsum_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   einsum_config = #vhlo.string_v1<"ab,bc->ac">
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<8x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<8x8x!vhlo.f32_v1>
-   %0 = "stablehlo.einsum"(%arg0, %arg1) {
-@@ -1323,22 +1418,25 @@
- }
- 
- // CHECK-LABEL: "op_exponential_minus_one"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_exponential_minus_one(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.exponential_minus_one_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.exponential_minus_one_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.exponential_minus_one"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_exponential"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_exponential(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.exponential_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.exponential_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.exponential"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_fft"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_fft(%arg0: tensor<16xcomplex<f32>>) -> tensor<16xcomplex<f32>> {
--  //      CHECK: "vhlo.fft_v1"(%arg0) <{
-+  //      CHECK: "vhlo.fft_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   fft_length = #vhlo.tensor_v1<dense<16> : tensor<1xi64>>,
-   // CHECK-SAME:   fft_type = #vhlo<fft_type_v1 FFT>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<16x!vhlo.complex_v1<!vhlo.f32_v1>>) -> !vhlo.tensor_v1<16x!vhlo.complex_v1<!vhlo.f32_v1>>
-@@ -1350,8 +1448,9 @@
- }
- 
- // CHECK-LABEL: "op_floor"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_floor(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.floor_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.floor_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.floor"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
-@@ -1364,16 +1463,17 @@
-   // CHECK-SAME:   sym_name = #vhlo.string_v1<"op_func">,
-   // CHECK-SAME:   sym_visibility = #vhlo.string_v1<"private">
-   // CHECK-SAME: }> ({
--  // CHECK-NEXT:   ^[[BB:bb.*]](%arg0: !vhlo.tensor_v1<!vhlo.f32_v1>):
--  // CHECK-NEXT:     "vhlo.return_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-+  // CHECK-NEXT:   ^[[BB:bb.*]](%[[ARG0:.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
-+  // CHECK-NEXT:     "vhlo.return_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-   // CHECK-NEXT: }) : () -> ()
- 
-   func.return %arg0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_gather"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_gather(%arg0 : tensor<2x4x9xf32>, %arg1 : tensor<1x5x2xi32>) -> tensor<1x5x1xf32> {
--  //      CHECK: "vhlo.gather_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.gather_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   collapsed_slice_dims = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-   // CHECK-SAME:   index_vector_dim = #vhlo.integer_v1<2 : i64>,
-   // CHECK-SAME:   indices_are_sorted = #vhlo.bool_v1<true>,
-@@ -1395,8 +1495,9 @@
- }
- 
- // CHECK-LABEL: "op_get_dimension_size"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_get_dimension_size(%arg0: tensor<?xf32>) -> tensor<i32> {
--  //      CHECK: "vhlo.get_dimension_size_v1"(%arg0) <{
-+  //      CHECK: "vhlo.get_dimension_size_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   dimension = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<?x!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-   %0 = "stablehlo.get_dimension_size"(%arg0) {
-@@ -1406,8 +1507,9 @@
- }
- 
- // CHECK-LABEL: "op_get_tuple_element"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_get_tuple_element(%arg0: tuple<tensor<f32>, tensor<i32>>) -> tensor<f32> {
--  //      CHECK: "vhlo.get_tuple_element_v1"(%arg0) <{
-+  //      CHECK: "vhlo.get_tuple_element_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   index = #vhlo.integer_v1<0 : i32>
-   // CHECK-SAME: }> : (!vhlo.tuple_v1<!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.get_tuple_element"(%arg0) {
-@@ -1417,11 +1519,12 @@
- }
- 
- // CHECK-LABEL: "op_if"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_if(%arg0: tensor<i1>, %arg1: tensor<f32>, %arg2: tensor<f32>) -> tensor<f32> {
--  //      CHECK: "vhlo.if_v1"(%arg0) ({
--  // CHECK-NEXT:   "vhlo.return_v1"(%arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-+  //      CHECK: "vhlo.if_v1"(%[[ARG0]]) ({
-+  // CHECK-NEXT:   "vhlo.return_v1"(%[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-   // CHECK-NEXT: }, {
--  // CHECK-NEXT:   "vhlo.return_v1"(%arg2) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-+  // CHECK-NEXT:   "vhlo.return_v1"(%[[ARG2]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-   // CHECK-NEXT: }) : (!vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.if"(%arg0) ({
-     "stablehlo.return"(%arg1) : (tensor<f32>) -> ()
-@@ -1432,15 +1535,17 @@
- }
- 
- // CHECK-LABEL: "op_imag"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_imag(%arg0: tensor<complex<f32>>) -> tensor<f32> {
--  // CHECK: "vhlo.imag_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.imag_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.imag"(%arg0) : (tensor<complex<f32>>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_infeed"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_infeed(%arg0: !stablehlo.token) -> (tensor<f32>, !stablehlo.token) {
--  //               CHECK: "vhlo.infeed_v1"(%arg0) <{
-+  //               CHECK: "vhlo.infeed_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   infeed_config = #vhlo.string_v1<"foo">,
-   // CHECK-SAME{LITERAL}:   layout = #vhlo.array_v1<[#vhlo.array_v1<[]>]>
-   //          CHECK-SAME: }> : (!vhlo.token_v1) -> (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.token_v1)
-@@ -1463,36 +1568,41 @@
- }
- 
- // CHECK-LABEL: "op_is_finite"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_is_finite(%arg0: tensor<f32>) -> tensor<i1> {
--  // CHECK: "vhlo.is_finite_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-+  // CHECK: "vhlo.is_finite_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-   %0 = "stablehlo.is_finite"(%arg0) : (tensor<f32>) -> tensor<i1>
-   func.return %0 : tensor<i1>
- }
- 
- // CHECK-LABEL: "op_log"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_log(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.log_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.log_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.log"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_log_plus_one"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_log_plus_one(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.log_plus_one_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.log_plus_one_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.log_plus_one"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_logistic"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_logistic(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.logistic_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.logistic_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.logistic"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_map"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_map(%arg0: tensor<16xf32>) -> tensor<16xf32> {
--  //      CHECK: "vhlo.map_v1"(%arg0) <{
-+  //      CHECK: "vhlo.map_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   dimensions = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>
-   // CHECK-SAME: }> ({
-   // CHECK-NEXT:   ^[[BB:bb.*]](%[[ARG1:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
-@@ -1510,57 +1620,65 @@
- }
- 
- // CHECK-LABEL: "op_maximum"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_maximum(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.maximum_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.maximum_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.maximum"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_minimum"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_minimum(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.minimum_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.minimum_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.minimum"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_multiply"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_multiply(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.multiply_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.multiply_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.multiply"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_negate"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_negate(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.negate_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.negate_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.negate"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_not"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_not(%arg0: tensor<i1>) -> tensor<i1> {
--  // CHECK: "vhlo.not_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-+  // CHECK: "vhlo.not_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-   %0 = "stablehlo.not"(%arg0) : (tensor<i1>) -> tensor<i1>
-   func.return %0 : tensor<i1>
- }
- 
- // CHECK-LABEL: "op_optimization_barrier"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_optimization_barrier(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.optimization_barrier_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.optimization_barrier_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.optimization_barrier"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_or"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_or(%arg0: tensor<i1>, %arg1: tensor<i1>) -> tensor<i1> {
--  // CHECK: "vhlo.or_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-+  // CHECK: "vhlo.or_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-   %0 = "stablehlo.or"(%arg0, %arg1) : (tensor<i1>, tensor<i1>) -> tensor<i1>
-   func.return %0 : tensor<i1>
- }
- 
- // CHECK-LABEL: "op_outfeed"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_outfeed(%arg0: tensor<f32>, %arg1: !stablehlo.token) -> !stablehlo.token {
--  //      CHECK: "vhlo.outfeed_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.outfeed_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   outfeed_config = #vhlo.string_v1<"foo">
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.token_v1) -> !vhlo.token_v1
-   %0 = "stablehlo.outfeed"(%arg0, %arg1) {
-@@ -1570,8 +1688,9 @@
- }
- 
- // CHECK-LABEL: "op_pad"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_pad(%arg0: tensor<8xf32>, %arg1: tensor<f32>) -> tensor<16xf32> {
--  //      CHECK: "vhlo.pad_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.pad_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   edge_padding_high = #vhlo.tensor_v1<dense<4> : tensor<1xi64>>,
-   // CHECK-SAME:   edge_padding_low = #vhlo.tensor_v1<dense<4> : tensor<1xi64>>,
-   // CHECK-SAME:   interior_padding = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>
-@@ -1585,36 +1704,41 @@
- }
- 
- // CHECK-LABEL: "op_popcnt"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_popcnt(%arg0: tensor<i32>) -> tensor<i32> {
--  // CHECK: "vhlo.popcnt_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-+  // CHECK: "vhlo.popcnt_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-   %0 = "stablehlo.popcnt"(%arg0) : (tensor<i32>) -> tensor<i32>
-   func.return %0 : tensor<i32>
- }
- 
- // CHECK-LABEL: "op_power"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_power(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.power_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.power_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.power"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_real_dynamic_slice"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}}, %[[ARG3:.*]]: {{.*}})
- func.func @op_real_dynamic_slice(%arg0: tensor<?xf32>, %arg1: tensor<1xindex>, %arg2: tensor<1xindex>, %arg3: tensor<1xindex>) -> tensor<?xf32> {
--  // CHECK: "vhlo.real_dynamic_slice_v1"(%arg0, %arg1, %arg2, %arg3) : (!vhlo.tensor_v1<?x!vhlo.f32_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>) -> !vhlo.tensor_v1<?x!vhlo.f32_v1>
-+  // CHECK: "vhlo.real_dynamic_slice_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]], %[[ARG3]]) : (!vhlo.tensor_v1<?x!vhlo.f32_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>) -> !vhlo.tensor_v1<?x!vhlo.f32_v1>
-   %0 = "stablehlo.real_dynamic_slice"(%arg0, %arg1, %arg2, %arg3) : (tensor<?xf32>, tensor<1xindex>, tensor<1xindex>, tensor<1xindex>) -> tensor<?xf32>
-   func.return %0 : tensor<?xf32>
- }
- 
- // CHECK-LABEL: "op_real"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_real(%arg0: tensor<complex<f32>>) -> tensor<f32> {
--  // CHECK: "vhlo.real_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.real_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.real"(%arg0) : (tensor<complex<f32>>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_recv"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_recv(%arg0: !stablehlo.token) -> (tensor<f32>, !stablehlo.token) {
--  //      CHECK: "vhlo.recv_v1"(%arg0) <{
-+  //      CHECK: "vhlo.recv_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME:   channel_type = #vhlo.integer_v1<3 : i64>,
-   // CHECK-SAME:   is_host_transfer = #vhlo.bool_v1<true>
-@@ -1627,8 +1751,9 @@
- }
- 
- // CHECK-LABEL: "op_reduce"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_reduce(%arg0: tensor<16xf32>, %arg1: tensor<f32>) -> tensor<f32> {
--  //  CHECK: "vhlo.reduce_v1"(%arg0, %arg1)
-+  //  CHECK: "vhlo.reduce_v1"(%[[ARG0]], %[[ARG1]])
-   //  CHECK:   ^[[BB:bb.*]](%[[ARG1:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>, %[[ARG2:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
-   //  CHECK:     "vhlo.return_v1"(%[[VAL1:.*]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-   //  CHECK: }) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-@@ -1643,8 +1768,9 @@
- }
- 
- // CHECK-LABEL: "op_reduce_precision"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_reduce_precision(%arg0: tensor<f32>) -> tensor<f32> {
--  //      CHECK: "vhlo.reduce_precision_v1"(%arg0) <{
-+  //      CHECK: "vhlo.reduce_precision_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   exponent_bits = #vhlo.integer_v1<8 : i32>
-   // CHECK-SAME:   mantissa_bits = #vhlo.integer_v1<10 : i32>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-@@ -1658,7 +1784,7 @@
- // CHECK_lABEL: "op_reduce_with_promotable_types"
- func.func @op_reduce_with_promotable_types(%arg0: tensor<4x4xf32>, %arg1 : tensor<f32>)
-     -> (tensor<4xf64>) {
--  //  CHECK: "vhlo.reduce_v1"(%arg0, %arg1)
-+  //  CHECK: "vhlo.reduce_v1"(%[[ARG0:.*]], %[[ARG1:.*]])
-   //  CHECK:   ^[[BB:bb.*]](%[[ARG1:arg.*]]: !vhlo.tensor_v1<!vhlo.f64_v1>, %[[ARG2:arg.*]]: !vhlo.tensor_v1<!vhlo.f64_v1>):
-   //  CHECK:     "vhlo.return_v1"(%[[VAL1:.*]]) : (!vhlo.tensor_v1<!vhlo.f64_v1>) -> ()
-   //  CHECK: }) : (!vhlo.tensor_v1<4x4x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<4x!vhlo.f64_v1>
-@@ -1673,8 +1799,9 @@
- }
- 
- // CHECK-LABEL: "op_reduce_scatter"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_reduce_scatter(%arg0: tensor<16xf32>) -> tensor<16xf32> {
--  //               CHECK: "vhlo.reduce_scatter_v1"(%arg0) <{
-+  //               CHECK: "vhlo.reduce_scatter_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0], [1]]> : tensor<2x1xi64>>,
-   //          CHECK-SAME:   scatter_dimension = #vhlo.integer_v1<0 : i64>
-@@ -1699,7 +1826,7 @@
- 
- // CHECK_lABEL: "op_reduce_scatter_with_promotable_types"
- func.func @op_reduce_scatter_with_promotable_types(%data: tensor<4x16xf32>) -> tensor<4x4xf64> {
--  //  CHECK: "vhlo.reduce_scatter_v1"(%arg0)
-+  //  CHECK: "vhlo.reduce_scatter_v1"(%[[ARG0:.*]])
-   //  CHECK:   ^[[BB:bb.*]](%[[ARG1:arg.*]]: !vhlo.tensor_v1<!vhlo.f64_v1>, %[[ARG2:arg.*]]: !vhlo.tensor_v1<!vhlo.f64_v1>):
-   //  CHECK:     "vhlo.return_v1"(%[[VAL1:.*]]) : (!vhlo.tensor_v1<!vhlo.f64_v1>) -> ()
-   //  CHECK: }) : (!vhlo.tensor_v1<4x16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<4x4x!vhlo.f64_v1>
-@@ -1716,8 +1843,9 @@
- 
- 
- // CHECK-LABEL: "op_reduce_window"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_reduce_window(%arg0: tensor<2x17x31x7xf32>, %arg1: tensor<f32>) -> tensor<2x9x16x7xf32> {
--  //               CHECK: "vhlo.reduce_window_v1"(%arg0, %arg1) <{
-+  //               CHECK: "vhlo.reduce_window_v1"(%[[ARG0]], %[[ARG1]]) <{
-   //          CHECK-SAME:   base_dilations = #vhlo.tensor_v1<dense<[1, 2, 2, 1]> : tensor<4xi64>>,
-   // CHECK-SAME{LITERAL}:   padding = #vhlo.tensor_v1<dense<[[0, 0], [2, 0], [0, 2], [0, 0]]> : tensor<4x2xi64>>,
-   //          CHECK-SAME:   window_dilations = #vhlo.tensor_v1<dense<[1, 2, 2, 1]> : tensor<4xi64>>,
-@@ -1742,11 +1870,11 @@
-   func.return %0 : tensor<2x9x16x7xf32>
- }
- 
--// CHECK_lABEL: "op_reduce_window_with_promotable_types"
-+// CHECK-LABEL: "op_reduce_window_with_promotable_types"
- func.func @op_reduce_window_with_promotable_types(%arg0: tensor<4x2xf32>,
-     %arg1: tensor<4x2xf32>, %init0: tensor<f32>, %init1: tensor<f32>) ->
-     (tensor<2x2xf64>, tensor<2x2xf32>) {
--  //  CHECK: "vhlo.reduce_window_v1"(%arg0, %arg1, %arg2, %arg3)
-+  //  CHECK: "vhlo.reduce_window_v1"(%[[ARG0:.*]], %[[ARG1:.*]], %[[ARG2:.*]], %[[ARG3:.*]])
-   //  CHECK:   ^[[BB:bb.*]](%[[ARG1:arg.*]]: !vhlo.tensor_v1<!vhlo.f64_v1>, %[[ARG2:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>, %[[ARG3:arg.*]]: !vhlo.tensor_v1<!vhlo.f64_v1>, %[[ARG4:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
-   //  CHECK:     "vhlo.return_v1"(%[[VAL1:.*]], %[[VAL2:.*]]) : (!vhlo.tensor_v1<!vhlo.f64_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-   //  CHECK: }) : (!vhlo.tensor_v1<4x2x!vhlo.f32_v1>, !vhlo.tensor_v1<4x2x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> (!vhlo.tensor_v1<2x2x!vhlo.f64_v1>, !vhlo.tensor_v1<2x2x!vhlo.f32_v1>)
-@@ -1766,8 +1894,9 @@
- }
- 
- // CHECK-LABEL: "op_remainder"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_remainder(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.remainder_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.remainder_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.remainder"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
-@@ -1787,16 +1916,18 @@
- }
- 
- // CHECK-LABEL: "op_reshape"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_reshape(%arg0: tensor<16xf32>) -> tensor<4x4xf32> {
--  // CHECK: "vhlo.reshape_v1"(%arg0) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<4x4x!vhlo.f32_v1>
-+  // CHECK: "vhlo.reshape_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<4x4x!vhlo.f32_v1>
-   %0 = "stablehlo.reshape"(%arg0) : (tensor<16xf32>) -> tensor<4x4xf32>
-   func.return %0 : tensor<4x4xf32>
- }
- 
- // CHECK-LABEL: "op_return"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_return(%arg0: tensor<i32>, %arg1: tensor<f32>) -> tensor<f32> {
--  //      CHECK: "vhlo.case_v1"(%arg0) ({
--  // CHECK-NEXT:   "vhlo.return_v1"(%arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-+  //      CHECK: "vhlo.case_v1"(%[[ARG0]]) ({
-+  // CHECK-NEXT:   "vhlo.return_v1"(%[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-   // CHECK-NEXT: }) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.case"(%arg0) ({
-     "stablehlo.return"(%arg1) : (tensor<f32>) -> ()
-@@ -1805,8 +1936,9 @@
- }
- 
- // CHECK-LABEL: "op_reverse"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_reverse(%arg0: tensor<16xf32>) -> tensor<16xf32> {
--  //      CHECK: "vhlo.reverse_v1"(%arg0) <{
-+  //      CHECK: "vhlo.reverse_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   dimensions = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x!vhlo.f32_v1>
-   %0 = "stablehlo.reverse"(%arg0) {
-@@ -1816,8 +1948,9 @@
- }
- 
- // CHECK-LABEL: "op_rng_bit_generator"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_rng_bit_generator(%arg0: tensor<f32>) -> (tensor<f32>, tensor<f32>) {
--  //      CHECK: "vhlo.rng_bit_generator_v1"(%arg0) <{
-+  //      CHECK: "vhlo.rng_bit_generator_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   rng_algorithm = #vhlo<rng_algorithm_v1 PHILOX>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>)
-   %0:2 = "stablehlo.rng_bit_generator"(%arg0) {
-@@ -1827,8 +1960,9 @@
- }
- 
- // CHECK-LABEL: "op_rng"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_rng(%arg0: tensor<f32>, %arg1: tensor<f32>, %arg2: tensor<0xindex>) -> tensor<f32> {
--  //      CHECK: "vhlo.rng_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.rng_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   rng_distribution = #vhlo<rng_distribution_v1 NORMAL>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<0x!vhlo.index_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.rng"(%arg0, %arg1, %arg2) {
-@@ -1838,29 +1972,33 @@
- }
- 
- // CHECK-LABEL: "op_round_nearest_afz"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_round_nearest_afz(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.round_nearest_afz_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.round_nearest_afz_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.round_nearest_afz"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_round_nearest_even"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_round_nearest_even(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.round_nearest_even_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.round_nearest_even_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.round_nearest_even"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_rsqrt"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_rsqrt(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.rsqrt_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.rsqrt_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.rsqrt"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_scatter"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_scatter(%arg0: tensor<200x100x300xf32>, %arg1: tensor<10x2xi32>, %arg2: tensor<10x300xf32>) -> tensor<200x100x300xf32> {
--  //      CHECK: "vhlo.scatter_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.scatter_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   index_vector_dim = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   indices_are_sorted = #vhlo.bool_v1<true>,
-   // CHECK-SAME:   inserted_window_dims = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-@@ -1893,7 +2031,7 @@
- func.func @op_scatter_with_promotable_types(%input_tensor: tensor<200x100x300xf32>,
-     %scatter_indices: tensor<10x2xi32>, %updates: tensor<10x300xf32>) ->
-       tensor<200x100x300xf64> {
--  //  CHECK: "vhlo.scatter_v1"(%arg0, %arg1, %arg2)
-+  //  CHECK: "vhlo.scatter_v1"(%[[ARG0:.*]], %[[ARG1:.*]], %[[ARG2:.*]])
-   //  CHECK:   ^[[BB:bb.*]](%[[ARG1:arg.*]]: !vhlo.tensor_v1<!vhlo.f64_v1>, %[[ARG2:arg.*]]: !vhlo.tensor_v1<!vhlo.f64_v1>):
-   //  CHECK:     "vhlo.return_v1"(%[[VAL1:.*]]) : (!vhlo.tensor_v1<!vhlo.f64_v1>) -> ()
-   //  CHECK: }) : (!vhlo.tensor_v1<200x100x300x!vhlo.f32_v1>, !vhlo.tensor_v1<10x2x!vhlo.i32_v1>, !vhlo.tensor_v1<10x300x!vhlo.f32_v1>) -> !vhlo.tensor_v1<200x100x300x!vhlo.f64_v1>
-@@ -1916,8 +2054,9 @@
- }
- 
- // CHECK-LABEL: "op_select_and_scatter"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_select_and_scatter(%arg0: tensor<10x24x24x64xf32>, %arg1: tensor<12x13x13x66xf32>, %arg2: tensor<f32>) -> tensor<10x24x24x64xf32> {
--  //      CHECK: "vhlo.select_and_scatter_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.select_and_scatter_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   padding = #vhlo.tensor_v1<dense<1> : tensor<4x2xi64>>,
-   // CHECK-SAME:   window_dimensions = #vhlo.tensor_v1<dense<[1, 2, 2, 1]> : tensor<4xi64>>,
-   // CHECK-SAME:   window_strides = #vhlo.tensor_v1<dense<[1, 2, 2, 1]> : tensor<4xi64>>
-@@ -1947,8 +2086,9 @@
- }
- 
- // CHECK-LABEL: "op_select_and_scatter_with_promotable_types"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_select_and_scatter_with_promotable_types(%arg0: tensor<10x24x24x64xf32>, %arg1: tensor<12x13x13x66xf32>, %arg2: tensor<f32>) -> tensor<10x24x24x64xf64> {
--  // CHECK: "vhlo.select_and_scatter_v1"(%arg0, %arg1, %arg2)
-+  // CHECK: "vhlo.select_and_scatter_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]])
-   // CHECK:   ^[[BB:bb.*]](%[[ARG1:arg.*]]: !vhlo.tensor_v1<!vhlo.f64_v1>, %[[ARG2:arg.*]]: !vhlo.tensor_v1<!vhlo.f64_v1>):
-   // CHECK:     %[[VAL:.*]] = "vhlo.add_v1"(%[[ARG1]], %[[ARG2]]) : (!vhlo.tensor_v1<!vhlo.f64_v1>, !vhlo.tensor_v1<!vhlo.f64_v1>) -> !vhlo.tensor_v1<!vhlo.f64_v1>
-   // CHECK:     "vhlo.return_v1"(%[[VAL]]) : (!vhlo.tensor_v1<!vhlo.f64_v1>) -> ()
-@@ -1970,15 +2110,17 @@
- }
- 
- // CHECK-LABEL: "op_select"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_select(%arg0: tensor<i1>, %arg1: tensor<f32>, %arg2: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.select_v1"(%arg0, %arg1, %arg2) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.select_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.select"(%arg0, %arg1, %arg2) : (tensor<i1>, tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_send"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_send(%arg0: tensor<f32>, %arg1: !stablehlo.token) -> !stablehlo.token {
--  //      CHECK: "vhlo.send_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.send_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME:   channel_type = #vhlo.integer_v1<2 : i64>,
-   // CHECK-SAME:   is_host_transfer = #vhlo.bool_v1<true>
-@@ -1991,8 +2133,9 @@
- }
- 
- // CHECK-LABEL: "op_set_dimension_size"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_set_dimension_size(%arg0: tensor<?xf32>, %arg1: tensor<i32>) -> tensor<16xf32> {
--  //      CHECK: "vhlo.set_dimension_size_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.set_dimension_size_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   dimension = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<?x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<16x!vhlo.f32_v1>
-   %0 = "stablehlo.set_dimension_size"(%arg0, %arg1) {
-@@ -2002,43 +2145,49 @@
- }
- 
- // CHECK-LABEL: "op_shift_left"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_shift_left(%arg0: tensor<i32>, %arg1: tensor<i32>) -> tensor<i32> {
--  // CHECK: "vhlo.shift_left_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.i32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-+  // CHECK: "vhlo.shift_left_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-   %0 = "stablehlo.shift_left"(%arg0, %arg1) : (tensor<i32>, tensor<i32>) -> tensor<i32>
-   func.return %0 : tensor<i32>
- }
- 
- // CHECK-LABEL: "op_shift_right_arithmetic"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_shift_right_arithmetic(%arg0: tensor<i32>, %arg1: tensor<i32>) -> tensor<i32> {
--  // CHECK: "vhlo.shift_right_arithmetic_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.i32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-+  // CHECK: "vhlo.shift_right_arithmetic_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-   %0 = "stablehlo.shift_right_arithmetic"(%arg0, %arg1) : (tensor<i32>, tensor<i32>) -> tensor<i32>
-   func.return %0 : tensor<i32>
- }
- 
- // CHECK-LABEL: "op_shift_right_logical"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_shift_right_logical(%arg0: tensor<i32>, %arg1: tensor<i32>) -> tensor<i32> {
--  // CHECK: "vhlo.shift_right_logical_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.i32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-+  // CHECK: "vhlo.shift_right_logical_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-   %0 = "stablehlo.shift_right_logical"(%arg0, %arg1) : (tensor<i32>, tensor<i32>) -> tensor<i32>
-   func.return %0 : tensor<i32>
- }
- 
- // CHECK-LABEL: "op_sign"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_sign(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.sign_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.sign_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.sign"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_sine"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_sine(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.sine_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.sine_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.sine"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_slice"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_slice(%arg0: tensor<16xf32>) -> tensor<4xf32> {
--  //      CHECK: "vhlo.slice_v1"(%arg0) <{
-+  //      CHECK: "vhlo.slice_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   limit_indices = #vhlo.tensor_v1<dense<4> : tensor<1xi64>>,
-   // CHECK-SAME:   start_indices = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>,
-   // CHECK-SAME:   strides = #vhlo.tensor_v1<dense<1> : tensor<1xi64>>
-@@ -2052,8 +2201,9 @@
- }
- 
- // CHECK-LABEL: "op_sort"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_sort(%arg0: tensor<16xf32>) -> tensor<16xf32> {
--  //      CHECK: "vhlo.sort_v1"(%arg0) <{
-+  //      CHECK: "vhlo.sort_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   dimension = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME:   is_stable = #vhlo.bool_v1<true>
-   // CHECK-SAME: }> ({
-@@ -2073,29 +2223,33 @@
- }
- 
- // CHECK-LABEL: "op_sqrt"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_sqrt(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.sqrt_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.sqrt_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.sqrt"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_subtract"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_subtract(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.subtract_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.subtract_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.subtract"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_tanh"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_tanh(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.tanh_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.tanh_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.tanh"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_torch_index_select"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_torch_index_select(%arg0: tensor<5x1x5xf32>, %arg1: tensor<2xi32>) ->  tensor<2x1x5xf32> {
--  //      CHECK: "vhlo.torch_index_select_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.torch_index_select_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   batch_dims = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME:   dim = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<5x1x5x!vhlo.f32_v1>, !vhlo.tensor_v1<2x!vhlo.i32_v1>) -> !vhlo.tensor_v1<2x1x5x!vhlo.f32_v1>
-@@ -2107,8 +2261,9 @@
- }
- 
- // CHECK-LABEL: "op_trace"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_trace(%arg0: tensor<f32>) {
--  //      CHECK: "vhlo.trace_v1"(%arg0) <{
-+  //      CHECK: "vhlo.trace_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   tag = #vhlo.string_v1<"foo">
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-   "stablehlo.trace"(%arg0) {
-@@ -2118,8 +2273,9 @@
- }
- 
- // CHECK-LABEL: "op_transpose"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_transpose(%arg0: tensor<16x8xf32>) ->  tensor<8x16xf32> {
--  //      CHECK: "vhlo.transpose_v1"(%arg0) <{
-+  //      CHECK: "vhlo.transpose_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   permutation = #vhlo.tensor_v1<dense<[1, 0]> : tensor<2xi64>>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<16x8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<8x16x!vhlo.f32_v1>
-   %0 = "stablehlo.transpose"(%arg0) {
-@@ -2129,8 +2285,9 @@
- }
- 
- // CHECK-LABEL: "op_triangular_solve"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_triangular_solve(%arg0: tensor<16x16xf32>, %arg1: tensor<16x16xf32>) ->  tensor<16x16xf32> {
--  //      CHECK: "vhlo.triangular_solve_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.triangular_solve_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   left_side = #vhlo.bool_v1<true>,
-   // CHECK-SAME:   lower = #vhlo.bool_v1<true>,
-   // CHECK-SAME:   transpose_a = #vhlo<transpose_v1 NO_TRANSPOSE>,
-@@ -2146,15 +2303,17 @@
- }
- 
- // CHECK-LABEL: "op_tuple"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_tuple(%arg0: tensor<f32>) -> tuple<tensor<f32>> {
--  // CHECK: "vhlo.tuple_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tuple_v1<!vhlo.tensor_v1<!vhlo.f32_v1>>
-+  // CHECK: "vhlo.tuple_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tuple_v1<!vhlo.tensor_v1<!vhlo.f32_v1>>
-   %0 = "stablehlo.tuple"(%arg0) : (tensor<f32>) -> tuple<tensor<f32>>
-   func.return %0 : tuple<tensor<f32>>
- }
- 
- // CHECK-LABEL: "op_unary_einsum"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_unary_einsum(%arg0: tensor<8x16xf32>) -> tensor<8xf32> {
--  //      CHECK: "vhlo.unary_einsum_v1"(%arg0) <{
-+  //      CHECK: "vhlo.unary_einsum_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   einsum_config = #vhlo.string_v1<"ab->a">
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<8x16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<8x!vhlo.f32_v1>
-   %0 = "stablehlo.unary_einsum"(%arg0) {
-@@ -2164,22 +2323,25 @@
- }
- 
- // CHECK-LABEL: "op_uniform_dequantize"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_uniform_dequantize(%arg0: tensor<!quant.uniform<i8:f32, 34.0:16>>) -> tensor<f32> {
--  // CHECK: "vhlo.uniform_dequantize_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.uniform_dequantize_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.uniform_dequantize"(%arg0) : (tensor<!quant.uniform<i8:f32, 34.0:16>>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_uniform_quantize"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_uniform_quantize(%arg0: tensor<f32>) -> tensor<!quant.uniform<i8:f32, 34.0:16>> {
--  // CHECK: "vhlo.uniform_quantize_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>
-+  // CHECK: "vhlo.uniform_quantize_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>
-   %0 = "stablehlo.uniform_quantize"(%arg0) : (tensor<f32>) -> tensor<!quant.uniform<i8:f32, 34.0:16>>
-   func.return %0 : tensor<!quant.uniform<i8:f32, 34.0:16>>
- }
- 
- // CHECK-LABEL: "op_while"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_while(%arg0: tensor<i1>) -> tensor<i1> {
--  //      CHECK: "vhlo.while_v1"(%arg0) ({
-+  //      CHECK: "vhlo.while_v1"(%[[ARG0]]) ({
-   // CHECK-NEXT:   ^[[BB:bb.*]](%[[ARG1:arg.*]]: !vhlo.tensor_v1<!vhlo.bool_v1>):
-   // CHECK-NEXT:     "vhlo.return_v1"(%[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>) -> ()
-   // CHECK-NEXT:   }, {
-@@ -2197,8 +2359,9 @@
- }
- 
- // CHECK-LABEL: "op_xor"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_xor(%arg0: tensor<i1>, %arg1: tensor<i1>) -> tensor<i1> {
--  // CHECK: "vhlo.xor_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-+  // CHECK: "vhlo.xor_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-   %0 = "stablehlo.xor"(%arg0, %arg1) : (tensor<i1>, tensor<i1>) -> tensor<i1>
-   func.return %0 : tensor<i1>
- }
-@@ -2206,190 +2369,217 @@
- // ============ TYPES ============
- 
- // CHECK-LABEL: "type_i1"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_i1(%arg0: tensor<i1>, %arg1: tensor<i1>) -> tensor<i1> {
--  // CHECK: "vhlo.and_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-+  // CHECK: "vhlo.and_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-   %0 = "stablehlo.and"(%arg0, %arg1) : (tensor<i1>, tensor<i1>) -> tensor<i1>
-   func.return %0 : tensor<i1>
- }
- 
- // CHECK-LABEL: "type_i4"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_i4(%arg0: tensor<i4>, %arg1: tensor<i4>) -> tensor<i4> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.i4_v1>, !vhlo.tensor_v1<!vhlo.i4_v1>) -> !vhlo.tensor_v1<!vhlo.i4_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i4_v1>, !vhlo.tensor_v1<!vhlo.i4_v1>) -> !vhlo.tensor_v1<!vhlo.i4_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<i4>, tensor<i4>) -> tensor<i4>
-   func.return %0 : tensor<i4>
- }
- 
- // CHECK-LABEL: "type_i8"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_i8(%arg0: tensor<i8>, %arg1: tensor<i8>) -> tensor<i8> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.i8_v1>, !vhlo.tensor_v1<!vhlo.i8_v1>) -> !vhlo.tensor_v1<!vhlo.i8_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i8_v1>, !vhlo.tensor_v1<!vhlo.i8_v1>) -> !vhlo.tensor_v1<!vhlo.i8_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<i8>, tensor<i8>) -> tensor<i8>
-   func.return %0 : tensor<i8>
- }
- 
- // CHECK-LABEL: "type_i16"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_i16(%arg0: tensor<i16>, %arg1: tensor<i16>) -> tensor<i16> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.i16_v1>, !vhlo.tensor_v1<!vhlo.i16_v1>) -> !vhlo.tensor_v1<!vhlo.i16_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i16_v1>, !vhlo.tensor_v1<!vhlo.i16_v1>) -> !vhlo.tensor_v1<!vhlo.i16_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<i16>, tensor<i16>) -> tensor<i16>
-   func.return %0 : tensor<i16>
- }
- 
- // CHECK-LABEL: "type_i32"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_i32(%arg0: tensor<i32>, %arg1: tensor<i32>) -> tensor<i32> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.i32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<i32>, tensor<i32>) -> tensor<i32>
-   func.return %0 : tensor<i32>
- }
- 
- // CHECK-LABEL: "type_i64"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_i64(%arg0: tensor<i64>, %arg1: tensor<i64>) -> tensor<i64> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.i64_v1>, !vhlo.tensor_v1<!vhlo.i64_v1>) -> !vhlo.tensor_v1<!vhlo.i64_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i64_v1>, !vhlo.tensor_v1<!vhlo.i64_v1>) -> !vhlo.tensor_v1<!vhlo.i64_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<i64>, tensor<i64>) -> tensor<i64>
-   func.return %0 : tensor<i64>
- }
- 
- // CHECK-LABEL: "type_ui4"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_ui4(%arg0: tensor<ui4>, %arg1: tensor<ui4>) -> tensor<ui4> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.ui4_v1>, !vhlo.tensor_v1<!vhlo.ui4_v1>) -> !vhlo.tensor_v1<!vhlo.ui4_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.ui4_v1>, !vhlo.tensor_v1<!vhlo.ui4_v1>) -> !vhlo.tensor_v1<!vhlo.ui4_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<ui4>, tensor<ui4>) -> tensor<ui4>
-   func.return %0 : tensor<ui4>
- }
- 
- // CHECK-LABEL: "type_ui8"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_ui8(%arg0: tensor<ui8>, %arg1: tensor<ui8>) -> tensor<ui8> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.ui8_v1>, !vhlo.tensor_v1<!vhlo.ui8_v1>) -> !vhlo.tensor_v1<!vhlo.ui8_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.ui8_v1>, !vhlo.tensor_v1<!vhlo.ui8_v1>) -> !vhlo.tensor_v1<!vhlo.ui8_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<ui8>, tensor<ui8>) -> tensor<ui8>
-   func.return %0 : tensor<ui8>
- }
- 
- // CHECK-LABEL: "type_ui16"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_ui16(%arg0: tensor<ui16>, %arg1: tensor<ui16>) -> tensor<ui16> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.ui16_v1>, !vhlo.tensor_v1<!vhlo.ui16_v1>) -> !vhlo.tensor_v1<!vhlo.ui16_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.ui16_v1>, !vhlo.tensor_v1<!vhlo.ui16_v1>) -> !vhlo.tensor_v1<!vhlo.ui16_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<ui16>, tensor<ui16>) -> tensor<ui16>
-   func.return %0 : tensor<ui16>
- }
- 
- // CHECK-LABEL: "type_ui32"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_ui32(%arg0: tensor<ui32>, %arg1: tensor<ui32>) -> tensor<ui32> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.ui32_v1>, !vhlo.tensor_v1<!vhlo.ui32_v1>) -> !vhlo.tensor_v1<!vhlo.ui32_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.ui32_v1>, !vhlo.tensor_v1<!vhlo.ui32_v1>) -> !vhlo.tensor_v1<!vhlo.ui32_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<ui32>, tensor<ui32>) -> tensor<ui32>
-   func.return %0 : tensor<ui32>
- }
- 
- // CHECK-LABEL: "type_ui64"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_ui64(%arg0: tensor<ui64>, %arg1: tensor<ui64>) -> tensor<ui64> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.ui64_v1>, !vhlo.tensor_v1<!vhlo.ui64_v1>) -> !vhlo.tensor_v1<!vhlo.ui64_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.ui64_v1>, !vhlo.tensor_v1<!vhlo.ui64_v1>) -> !vhlo.tensor_v1<!vhlo.ui64_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<ui64>, tensor<ui64>) -> tensor<ui64>
-   func.return %0 : tensor<ui64>
- }
- 
- // CHECK-LABEL: "type_f8E4M3FN"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_f8E4M3FN(%arg0: tensor<f8E4M3FN>, %arg1: tensor<f8E4M3FN>) -> tensor<f8E4M3FN> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f8E4M3FN_v1>, !vhlo.tensor_v1<!vhlo.f8E4M3FN_v1>) -> !vhlo.tensor_v1<!vhlo.f8E4M3FN_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f8E4M3FN_v1>, !vhlo.tensor_v1<!vhlo.f8E4M3FN_v1>) -> !vhlo.tensor_v1<!vhlo.f8E4M3FN_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f8E4M3FN>, tensor<f8E4M3FN>) -> tensor<f8E4M3FN>
-   func.return %0 : tensor<f8E4M3FN>
- }
- 
- // CHECK-LABEL: "type_f8E5M2"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_f8E5M2(%arg0: tensor<f8E5M2>, %arg1: tensor<f8E5M2>) -> tensor<f8E5M2> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f8E5M2_v1>, !vhlo.tensor_v1<!vhlo.f8E5M2_v1>) -> !vhlo.tensor_v1<!vhlo.f8E5M2_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f8E5M2_v1>, !vhlo.tensor_v1<!vhlo.f8E5M2_v1>) -> !vhlo.tensor_v1<!vhlo.f8E5M2_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f8E5M2>, tensor<f8E5M2>) -> tensor<f8E5M2>
-   func.return %0 : tensor<f8E5M2>
- }
- 
- // CHECK-LABEL: "type_f8E4M3FNUZ"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_f8E4M3FNUZ(%arg0: tensor<f8E4M3FNUZ>, %arg1: tensor<f8E4M3FNUZ>) -> tensor<f8E4M3FNUZ> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f8E4M3FNUZ_v1>, !vhlo.tensor_v1<!vhlo.f8E4M3FNUZ_v1>) -> !vhlo.tensor_v1<!vhlo.f8E4M3FNUZ_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f8E4M3FNUZ_v1>, !vhlo.tensor_v1<!vhlo.f8E4M3FNUZ_v1>) -> !vhlo.tensor_v1<!vhlo.f8E4M3FNUZ_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f8E4M3FNUZ>, tensor<f8E4M3FNUZ>) -> tensor<f8E4M3FNUZ>
-   func.return %0 : tensor<f8E4M3FNUZ>
- }
- 
- // CHECK-LABEL: "type_f8E4M3B11FNUZ"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_f8E4M3B11FNUZ(%arg0: tensor<f8E4M3B11FNUZ>, %arg1: tensor<f8E4M3B11FNUZ>) -> tensor<f8E4M3B11FNUZ> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f8E4M3B11FNUZ_v1>, !vhlo.tensor_v1<!vhlo.f8E4M3B11FNUZ_v1>) -> !vhlo.tensor_v1<!vhlo.f8E4M3B11FNUZ_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f8E4M3B11FNUZ_v1>, !vhlo.tensor_v1<!vhlo.f8E4M3B11FNUZ_v1>) -> !vhlo.tensor_v1<!vhlo.f8E4M3B11FNUZ_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f8E4M3B11FNUZ>, tensor<f8E4M3B11FNUZ>) -> tensor<f8E4M3B11FNUZ>
-   func.return %0 : tensor<f8E4M3B11FNUZ>
- }
- 
- // CHECK-LABEL: "type_f8E5M2FNUZ"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_f8E5M2FNUZ(%arg0: tensor<f8E5M2FNUZ>, %arg1: tensor<f8E5M2FNUZ>) -> tensor<f8E5M2FNUZ> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f8E5M2FNUZ_v1>, !vhlo.tensor_v1<!vhlo.f8E5M2FNUZ_v1>) -> !vhlo.tensor_v1<!vhlo.f8E5M2FNUZ_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f8E5M2FNUZ_v1>, !vhlo.tensor_v1<!vhlo.f8E5M2FNUZ_v1>) -> !vhlo.tensor_v1<!vhlo.f8E5M2FNUZ_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f8E5M2FNUZ>, tensor<f8E5M2FNUZ>) -> tensor<f8E5M2FNUZ>
-   func.return %0 : tensor<f8E5M2FNUZ>
- }
- 
- // CHECK-LABEL: "type_bf16"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_bf16(%arg0: tensor<bf16>, %arg1: tensor<bf16>) -> tensor<bf16> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.bf16_v1>, !vhlo.tensor_v1<!vhlo.bf16_v1>) -> !vhlo.tensor_v1<!vhlo.bf16_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.bf16_v1>, !vhlo.tensor_v1<!vhlo.bf16_v1>) -> !vhlo.tensor_v1<!vhlo.bf16_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<bf16>, tensor<bf16>) -> tensor<bf16>
-   func.return %0 : tensor<bf16>
- }
- 
- // CHECK-LABEL: "type_f16"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_f16(%arg0: tensor<f16>, %arg1: tensor<f16>) -> tensor<f16> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f16_v1>, !vhlo.tensor_v1<!vhlo.f16_v1>) -> !vhlo.tensor_v1<!vhlo.f16_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f16_v1>, !vhlo.tensor_v1<!vhlo.f16_v1>) -> !vhlo.tensor_v1<!vhlo.f16_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f16>, tensor<f16>) -> tensor<f16>
-   func.return %0 : tensor<f16>
- }
- 
- // CHECK-LABEL: "type_f32"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_f32(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "type_f64"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_f64(%arg0: tensor<f64>, %arg1: tensor<f64>) -> tensor<f64> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f64_v1>, !vhlo.tensor_v1<!vhlo.f64_v1>) -> !vhlo.tensor_v1<!vhlo.f64_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f64_v1>, !vhlo.tensor_v1<!vhlo.f64_v1>) -> !vhlo.tensor_v1<!vhlo.f64_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f64>, tensor<f64>) -> tensor<f64>
-   func.return %0 : tensor<f64>
- }
- 
- // CHECK-LABEL: "type_complex_f32"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_complex_f32(%arg0: tensor<complex<f32>>, %arg1: tensor<complex<f32>>) -> tensor<complex<f32>> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>, !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>) -> !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>, !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>) -> !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<complex<f32>>, tensor<complex<f32>>) -> tensor<complex<f32>>
-   func.return %0 : tensor<complex<f32>>
- }
- 
- // CHECK-LABEL: "type_complex_f64"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_complex_f64(%arg0: tensor<complex<f64>>, %arg1: tensor<complex<f64>>) -> tensor<complex<f64>> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f64_v1>>, !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f64_v1>>) -> !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f64_v1>>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f64_v1>>, !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f64_v1>>) -> !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f64_v1>>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<complex<f64>>, tensor<complex<f64>>) -> tensor<complex<f64>>
-   func.return %0 : tensor<complex<f64>>
- }
- 
- // CHECK-LABEL: "type_dynamism_ranked"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @type_dynamism_ranked(%arg0: tensor<?xf32>) -> tensor<?xf32> {
--  // CHECK: "vhlo.abs_v1"(%arg0) : (!vhlo.tensor_v1<?x!vhlo.f32_v1>) -> !vhlo.tensor_v1<?x!vhlo.f32_v1>
-+  // CHECK: "vhlo.abs_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<?x!vhlo.f32_v1>) -> !vhlo.tensor_v1<?x!vhlo.f32_v1>
-   %0 = "stablehlo.abs"(%arg0) : (tensor<?xf32>) -> tensor<?xf32>
-   func.return %0 : tensor<?xf32>
- }
- 
- // CHECK-LABEL: "type_quantization"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_quantization(%arg0: tensor<!quant.uniform<i8:f32, 34.0:16>>, %arg1: tensor<!quant.uniform<i8:f32, 34.0:16>>) -> tensor<!quant.uniform<i8:f32, 34.0:16>> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>, !vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>) -> !vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>, !vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>) -> !vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<!quant.uniform<i8:f32, 34.0:16>>, tensor<!quant.uniform<i8:f32, 34.0:16>>) -> tensor<!quant.uniform<i8:f32, 34.0:16>>
-   func.return %0 : tensor<!quant.uniform<i8:f32, 34.0:16>>
- }
- 
- //       CHECK: function_type = #vhlo.type_v1<!vhlo.func_v1<(!vhlo.token_v1) -> !vhlo.token_v1>>
- // CHECK-LABEL: "type_token_callee"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @type_token_callee(%arg0: !stablehlo.token) -> !stablehlo.token {
--  // CHECK: "vhlo.return_v1"(%arg0) : (!vhlo.token_v1) -> ()
-+  // CHECK: "vhlo.return_v1"(%[[ARG0]]) : (!vhlo.token_v1) -> ()
-   return %arg0 : !stablehlo.token
- }
- 
- //       CHECK: function_type = #vhlo.type_v1<!vhlo.func_v1<(!vhlo.token_v1) -> !vhlo.token_v1>>
- // CHECK-LABEL: "type_token_caller"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @type_token_caller(%arg0: !stablehlo.token) -> !stablehlo.token {
--  // CHECK:      "vhlo.call_v1"(%arg0) <{callee = #vhlo.string_v1<"type_token_callee">}
-+  // CHECK:      "vhlo.call_v1"(%[[ARG0]]) <{callee = #vhlo.string_v1<"type_token_callee">}
-   // CHECK-SAME: (!vhlo.token_v1) -> !vhlo.token_v1
-   %0 = func.call @type_token_callee(%arg0) : (!stablehlo.token) -> !stablehlo.token
-   return %0 : !stablehlo.token
- }
- 
- // CHECK-LABEL: "type_tuple"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @type_tuple(%arg0: tuple<tensor<f32>>) -> tuple<!stablehlo.token> {
-   %0 = "stablehlo.custom_call"(%arg0) {
-     call_target_name = "foo"
-diff --ruN a/stablehlo/stablehlo/tests/vhlo/stablehlo_legalize_to_vhlo.0_18_0.mlir b/stablehlo/stablehlo/tests/vhlo/stablehlo_legalize_to_vhlo.0_18_0.mlir
---- stablehlo/stablehlo/tests/vhlo/stablehlo_legalize_to_vhlo.0_18_0.mlir
-+++ stablehlo/stablehlo/tests/vhlo/stablehlo_legalize_to_vhlo.0_18_0.mlir
-@@ -13,6 +13,7 @@
- // ============ ATTRIBUTES ============
- 
- // CHECK-LABEL: "attr_comparison_direction_eq"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_direction_eq(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     // CHECK: comparison_direction = #vhlo<comparison_direction_v1 EQ>
-@@ -22,6 +23,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_direction_ne"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_direction_ne(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     // CHECK: comparison_direction = #vhlo<comparison_direction_v1 NE>
-@@ -31,6 +33,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_direction_ge"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_direction_ge(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     // CHECK: comparison_direction = #vhlo<comparison_direction_v1 GE>
-@@ -40,6 +43,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_direction_gt"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_direction_gt(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     // CHECK: comparison_direction = #vhlo<comparison_direction_v1 GT>
-@@ -49,6 +53,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_direction_le"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_direction_le(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     // CHECK: comparison_direction = #vhlo<comparison_direction_v1 LE>
-@@ -58,6 +63,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_direction_lt"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_direction_lt(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     // CHECK: comparison_direction = #vhlo<comparison_direction_v1 LT>
-@@ -67,6 +73,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_type_notype"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_type_notype(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     comparison_direction = #stablehlo<comparison_direction EQ>
-@@ -76,6 +83,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_type_float"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_type_float(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     comparison_direction = #stablehlo<comparison_direction EQ>,
-@@ -86,6 +94,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_type_totalorder"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_type_totalorder(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     comparison_direction = #stablehlo<comparison_direction EQ>,
-@@ -96,6 +105,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_type_signed"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_type_signed(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     comparison_direction = #stablehlo<comparison_direction EQ>,
-@@ -106,6 +116,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_type_unsigned"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_type_unsigned(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     comparison_direction = #stablehlo<comparison_direction EQ>,
-@@ -118,6 +129,7 @@
- // ConvDimensionNumbers aka #stablehlo.conv is covered below.
- 
- // CHECK-LABEL: "attr_custom_call_api_version_unspecified"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_custom_call_api_version_unspecified(%arg0: tensor<f32>) -> tensor<f32> {
-   %0 = "stablehlo.custom_call"(%arg0) {
-     call_target_name = "foo",
-@@ -128,6 +140,7 @@
- }
- 
- // CHECK-LABEL: "attr_custom_call_api_version_original"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_custom_call_api_version_original(%arg0: tensor<f32>) -> tensor<f32> {
-   %0 = "stablehlo.custom_call"(%arg0) {
-     call_target_name = "foo",
-@@ -138,6 +151,7 @@
- }
- 
- // CHECK-LABEL: "attr_custom_call_api_version_status_returning"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_custom_call_api_version_status_returning(%arg0: tensor<f32>) -> tensor<f32> {
-   %0 = "stablehlo.custom_call"(%arg0) {
-     call_target_name = "foo",
-@@ -148,6 +162,7 @@
- }
- 
- // CHECK-LABEL: "attr_custom_call_api_version_status_returning_unified"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_custom_call_api_version_status_returning_unified(%arg0: tensor<f32>) -> tensor<f32> {
-   %0 = "stablehlo.custom_call"(%arg0) {
-     call_target_name = "foo",
-@@ -166,6 +181,7 @@
- // DotDimensionNumbers aka #stablehlo.dot is covered below.
- 
- // CHECK-LABEL: "attr_fft_type_fft"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_fft_type_fft(%arg0: tensor<16xcomplex<f32>>) -> tensor<16xcomplex<f32>> {
-   %0 = "stablehlo.fft"(%arg0) {
-     // CHECK: fft_type = #vhlo<fft_type_v1 FFT>
-@@ -176,6 +192,7 @@
- }
- 
- // CHECK-LABEL: "attr_fft_type_ifft"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_fft_type_ifft(%arg0: tensor<16xcomplex<f32>>) -> tensor<16xcomplex<f32>> {
-   %0 = "stablehlo.fft"(%arg0) {
-     // CHECK: fft_type = #vhlo<fft_type_v1 IFFT>
-@@ -186,6 +203,7 @@
- }
- 
- // CHECK-LABEL: "attr_fft_type_rfft"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_fft_type_rfft(%arg0: tensor<16xf32>) -> tensor<9xcomplex<f32>> {
-   %0 = "stablehlo.fft"(%arg0) {
-     // CHECK: fft_type = #vhlo<fft_type_v1 RFFT>
-@@ -196,6 +214,7 @@
- }
- 
- // CHECK-LABEL: "attr_fft_type_irfft"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_fft_type_irfft(%arg0: tensor<9xcomplex<f32>>) -> tensor<16xf32> {
-   %0 = "stablehlo.fft"(%arg0) {
-     // CHECK: fft_type = #vhlo<fft_type_v1 IRFFT>
-@@ -208,6 +227,7 @@
- // GatherDimensionNumbers aka #stablehlo.gather is covered below.
- 
- // CHECK-LABEL: "attr_precision_config_default"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_precision_config_default(%arg0: tensor<8x16xf32>, %arg1: tensor<16x8xf32>) -> tensor<8x8xf32> {
-   %0 = "stablehlo.dot"(%arg0, %arg1) {
-     // CHECK: precision_config = #vhlo.array_v1<[#vhlo<precision_v1 DEFAULT>, #vhlo<precision_v1 DEFAULT>]>
-@@ -216,6 +236,7 @@
- }
- 
- // CHECK-LABEL: "attr_precision_config_high"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_precision_config_high(%arg0: tensor<8x16xf32>, %arg1: tensor<16x8xf32>) -> tensor<8x8xf32> {
-   %0 = "stablehlo.dot"(%arg0, %arg1) {
-     // CHECK: precision_config = #vhlo.array_v1<[#vhlo<precision_v1 HIGH>, #vhlo<precision_v1 HIGH>]>
-@@ -225,6 +246,7 @@
- }
- 
- // CHECK-LABEL: "attr_precision_config_highest"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_precision_config_highest(%arg0: tensor<8x16xf32>, %arg1: tensor<16x8xf32>) -> tensor<8x8xf32> {
-   %0 = "stablehlo.dot"(%arg0, %arg1) {
-     // CHECK: precision_config = #vhlo.array_v1<[#vhlo<precision_v1 HIGHEST>, #vhlo<precision_v1 HIGHEST>]>
-@@ -234,6 +256,7 @@
- }
- 
- // CHECK-LABEL: "attr_rng_algorithm_default"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_rng_algorithm_default(%arg0: tensor<f32>) -> (tensor<f32>, tensor<f32>) {
-   %0:2 = "stablehlo.rng_bit_generator"(%arg0) {
-     // CHECK: rng_algorithm = #vhlo<rng_algorithm_v1 DEFAULT>
-@@ -243,6 +266,7 @@
- }
- 
- // CHECK-LABEL: "attr_rng_algorithm_three_fry"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_rng_algorithm_three_fry(%arg0: tensor<f32>) -> (tensor<f32>, tensor<f32>) {
-   %0:2 = "stablehlo.rng_bit_generator"(%arg0) {
-     // CHECK: rng_algorithm = #vhlo<rng_algorithm_v1 THREE_FRY>
-@@ -252,6 +276,7 @@
- }
- 
- // CHECK-LABEL: "attr_rng_algorithm_philox"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_rng_algorithm_philox(%arg0: tensor<f32>) -> (tensor<f32>, tensor<f32>) {
-   %0:2 = "stablehlo.rng_bit_generator"(%arg0) {
-     // CHECK: rng_algorithm = #vhlo<rng_algorithm_v1 PHILOX>
-@@ -261,6 +286,7 @@
- }
- 
- // CHECK-LABEL: "attr_rng_distribution_uniform"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @attr_rng_distribution_uniform(%arg0: tensor<f32>, %arg1: tensor<f32>, %arg2: tensor<0xindex>) -> tensor<f32> {
-   %0 = "stablehlo.rng"(%arg0, %arg1, %arg2) {
-     // CHECK: rng_distribution = #vhlo<rng_distribution_v1 UNIFORM>
-@@ -270,6 +296,7 @@
- }
- 
- // CHECK-LABEL: "attr_rng_distribution_normal"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @attr_rng_distribution_normal(%arg0: tensor<f32>, %arg1: tensor<f32>, %arg2: tensor<0xindex>) -> tensor<f32> {
-   %0 = "stablehlo.rng"(%arg0, %arg1, %arg2) {
-     // CHECK: rng_distribution = #vhlo<rng_distribution_v1 NORMAL>
-@@ -281,6 +308,7 @@
- // ScatterDimensionNumbers aka #stablehlo.scatter is covered below.
- 
- // CHECK-LABEL: "attr_transpose_no_transpose"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_transpose_no_transpose(%arg0: tensor<16x16xf32>, %arg1: tensor<16x16xf32>) ->  tensor<16x16xf32> {
-   %0 = "stablehlo.triangular_solve"(%arg0, %arg1) {
-     left_side = true,
-@@ -293,6 +321,7 @@
- }
- 
- // CHECK-LABEL: "attr_transpose_transpose"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_transpose_transpose(%arg0: tensor<16x16xf32>, %arg1: tensor<16x16xf32>) ->  tensor<16x16xf32> {
-   %0 = "stablehlo.triangular_solve"(%arg0, %arg1) {
-     left_side = true,
-@@ -305,6 +334,7 @@
- }
- 
- // CHECK-LABEL: "attr_transpose_adjoint"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_transpose_adjoint(%arg0: tensor<16x16xf32>, %arg1: tensor<16x16xf32>) ->  tensor<16x16xf32> {
-   %0 = "stablehlo.triangular_solve"(%arg0, %arg1) {
-     left_side = true,
-@@ -319,10 +349,9 @@
- // TypeExtensionsAttr aka #stablehlo.type_extensions is covered below.
- 
- // CHECK-LABEL: "attr_type_extensions_bounds"
--func.func @attr_type_extensions_bounds(
--    %arg0: tensor<?x?xf32, #stablehlo.type_extensions<bounds = [16, ?]>>)
--    -> tensor<?x?xf32, #stablehlo.type_extensions<bounds = [16, ?]>> {
--  // CHECK: "vhlo.return_v1"(%arg0) : (!vhlo.tensor_v1<?x?x!vhlo.f32_v1, #vhlo.type_extensions_v1<bounds = [16, ?]>>) -> ()
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @attr_type_extensions_bounds(%arg0: tensor<?x?xf32, #stablehlo.type_extensions<bounds = [16, ?]>>) -> tensor<?x?xf32, #stablehlo.type_extensions<bounds = [16, ?]>> {
-+  // CHECK: "vhlo.return_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<?x?x!vhlo.f32_v1, #vhlo.type_extensions_v1<bounds = [16, ?]>>) -> ()
-   func.return %arg0 : tensor<?x?xf32, #stablehlo.type_extensions<bounds = [16, ?]>>
- }
- 
-@@ -330,8 +359,9 @@
- // ============ DEFAULTS ============
- 
- // CHECK-LABEL: "default_all_gather"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_all_gather(%arg0: tensor<16x8xf32>) -> tensor<16x16xf32> {
--  //               CHECK: "vhlo.all_gather_v1"(%arg0) <{
-+  //               CHECK: "vhlo.all_gather_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   all_gather_dim = #vhlo.integer_v1<1 : i64>
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0], [1]]> : tensor<2x1xi64>>,
-@@ -345,8 +375,9 @@
- }
- 
- // CHECK-LABEL: "default_all_reduce"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_all_reduce(%arg0: tensor<f32>) -> tensor<f32> {
--  //               CHECK: "vhlo.all_reduce_v1"(%arg0)
-+  //               CHECK: "vhlo.all_reduce_v1"(%[[ARG0]])
-   //          CHECK-SAME: <{
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0], [1]]> : tensor<2x1xi64>>,
-@@ -368,8 +399,9 @@
- }
- 
- // CHECK-LABEL: "default_all_to_all"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_all_to_all(%arg0: tensor<4x16xf32>) -> tensor<16x4xf32> {
--  //               CHECK: "vhlo.all_to_all_v1"(%arg0) <{
-+  //               CHECK: "vhlo.all_to_all_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   //          CHECK-SAME:   concat_dimension = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0, 1, 2, 3]]> : tensor<1x4xi64>>,
-@@ -386,8 +418,9 @@
- }
- 
- // CHECK-LABEL: "default_cholesky"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_cholesky(%arg0: tensor<1x16x16xf32>) -> tensor<1x16x16xf32> {
--  //      CHECK: "vhlo.cholesky_v1"(%arg0) <{
-+  //      CHECK: "vhlo.cholesky_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   lower = #vhlo.bool_v1<false>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<1x16x16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<1x16x16x!vhlo.f32_v1>
-   %0 = "stablehlo.cholesky"(%arg0) : (tensor<1x16x16xf32>) -> tensor<1x16x16xf32>
-@@ -395,8 +428,9 @@
- }
- 
- // CHECK-LABEL: "default_collective_permute"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_collective_permute(%arg0: tensor<16x8xf32>) -> tensor<16x8xf32> {
--  //               CHECK: "vhlo.collective_permute_v1"(%arg0) <{
-+  //               CHECK: "vhlo.collective_permute_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME{LITERAL}:   source_target_pairs = #vhlo.tensor_v1<dense<[[0, 1], [1, 2], [2, 3]]> : tensor<3x2xi64>>
-   //          CHECK-SAME: }> : (!vhlo.tensor_v1<16x8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x8x!vhlo.f32_v1>
-@@ -407,8 +441,9 @@
- }
- 
- // CHECK-LABEL: "default_collective_broadcast"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_collective_broadcast(%arg0: tensor<16x8xf32>) -> tensor<16x8xf32> {
--  //               CHECK: "vhlo.collective_broadcast_v1"(%arg0) <{
-+  //               CHECK: "vhlo.collective_broadcast_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0, 1]]> : tensor<1x2xi64>>
-   //          CHECK-SAME: }> : (!vhlo.tensor_v1<16x8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x8x!vhlo.f32_v1>
-@@ -419,8 +454,9 @@
- }
- 
- // CHECK-LABEL: "default_compare"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @default_compare(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
--  //      CHECK: "vhlo.compare_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.compare_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   compare_type = #vhlo<comparison_type_v1 NOTYPE>,
-   // CHECK-SAME:   comparison_direction = #vhlo<comparison_direction_v1 EQ>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-@@ -431,8 +467,9 @@
- }
- 
- // CHECK-LABEL: "default_convolution"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @default_convolution(%arg0: tensor<1x8x8x207xf32>, %arg1: tensor<3x3x207x16xf32>) -> tensor<1x6x6x16xf32> {
--  //      CHECK: "vhlo.convolution_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.convolution_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   batch_group_count = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   feature_group_count = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   input_batch_dimension = #vhlo.integer_v1<0 : i64>,
-@@ -460,8 +497,9 @@
- }
- 
- // CHECK-LABEL: "default_custom_call"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_custom_call(%arg0: tensor<f32>) -> tensor<f32> {
--  //      CHECK: "vhlo.custom_call_v1"(%arg0) <{
-+  //      CHECK: "vhlo.custom_call_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   api_version = #vhlo<api_version_v1 API_VERSION_ORIGINAL>,
-   // CHECK-SAME:   backend_config = #vhlo.string_v1<"">,
-   // CHECK-SAME:   call_target_name = #vhlo.string_v1<"foo">,
-@@ -478,8 +516,9 @@
- }
- 
- // CHECK-LABEL: "default_dot_general"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @default_dot_general(%arg0: tensor<8x8x16xf32>, %arg1: tensor<8x16x8xf32>) -> tensor<8x8x8xf32> {
--  //      CHECK: "vhlo.dot_general_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.dot_general_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   lhs_batching_dimensions = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>,
-   // CHECK-SAME:   lhs_contracting_dimensions = #vhlo.tensor_v1<dense<2> : tensor<1xi64>>,
-   // CHECK-SAME:   precision_config = #vhlo.array_v1<[#vhlo<precision_v1 DEFAULT>, #vhlo<precision_v1 DEFAULT>]>,
-@@ -498,8 +537,9 @@
- }
- 
- // CHECK-LABEL: "default_dot"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @default_dot(%arg0: tensor<8x16xf32>, %arg1: tensor<16x8xf32>) -> tensor<8x8xf32> {
--  //      CHECK: "vhlo.dot_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.dot_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   precision_config = #vhlo.array_v1<[#vhlo<precision_v1 DEFAULT>, #vhlo<precision_v1 DEFAULT>]>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<8x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<8x8x!vhlo.f32_v1>
-   %0 = "stablehlo.dot"(%arg0, %arg1) : (tensor<8x16xf32>, tensor<16x8xf32>) -> tensor<8x8xf32>
-@@ -507,8 +547,9 @@
- }
- 
- // CHECK-LABEL: "default_dynamic_broadcast_in_dim"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @default_dynamic_broadcast_in_dim(%arg0: tensor<?x?xf32>, %arg1: tensor<2xindex>) -> tensor<?x?xf32> {
--  //      CHECK: "vhlo.dynamic_broadcast_in_dim_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.dynamic_broadcast_in_dim_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   broadcast_dimensions = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-   // CHECK-SAME:   known_expanding_dimensions = #vhlo.tensor_v1<dense<> : tensor<0xi64>>,
-   // CHECK-SAME:   known_nonexpanding_dimensions = #vhlo.tensor_v1<dense<> : tensor<0xi64>>
-@@ -520,8 +561,9 @@
- }
- 
- // CHECK-LABEL: "default_dynamic_conv"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @default_dynamic_conv(%arg0: tensor<1x8x8x207xf32>, %arg1: tensor<3x3x207x16xf32>, %arg2: tensor<4xi32>) -> tensor<1x?x?x16xf32> {
--  //      CHECK: "vhlo.dynamic_conv_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.dynamic_conv_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   batch_group_count = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   feature_group_count = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   input_batch_dimension = #vhlo.integer_v1<0 : i64>,
-@@ -549,8 +591,9 @@
- }
- 
- // CHECK-LABEL: "default_dynamic_gather"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @default_dynamic_gather(%arg0 : tensor<2x4x9xf32>, %arg1 : tensor<1x5x2xi32>, %arg2 : tensor<3xi32>) -> tensor<1x5x8xf32> {
--  //      CHECK: "vhlo.dynamic_gather_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.dynamic_gather_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   collapsed_slice_dims = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-   // CHECK-SAME:   index_vector_dim = #vhlo.integer_v1<2 : i64>,
-   // CHECK-SAME:   indices_are_sorted = #vhlo.bool_v1<false>,
-@@ -576,15 +619,16 @@
-   // CHECK-SAME:   sym_name = #vhlo.string_v1<"default_func">,
-   // CHECK-SAME:   sym_visibility = #vhlo.string_v1<"">
-   // CHECK-SAME: }> ({
--  // CHECK-NEXT:   ^[[BB:bb.*]](%arg0: !vhlo.tensor_v1<!vhlo.f32_v1>):
--  // CHECK-NEXT:     "vhlo.return_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-+  // CHECK-NEXT:   ^[[BB:bb.*]](%[[ARG0:.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
-+  // CHECK-NEXT:     "vhlo.return_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-   // CHECK-NEXT: }) : () -> ()
-   func.return %arg0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "dynamic_gather"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @dynamic_gather(%arg0 : tensor<2x4x9xf32>, %arg1 : tensor<1x5x2xi32>) -> tensor<1x5x1xf32> {
--  //      CHECK: "vhlo.gather_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.gather_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   collapsed_slice_dims = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-   // CHECK-SAME:   index_vector_dim = #vhlo.integer_v1<2 : i64>,
-   // CHECK-SAME:   indices_are_sorted = #vhlo.bool_v1<false>,
-@@ -605,8 +649,9 @@
- }
- 
- // CHECK-LABEL: "default_infeed"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_infeed(%arg0: !stablehlo.token) -> (tensor<f32>, !stablehlo.token) {
--  //               CHECK: "vhlo.infeed_v1"(%arg0) <{
-+  //               CHECK: "vhlo.infeed_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   infeed_config = #vhlo.string_v1<"">,
-   // CHECK-SAME{LITERAL}:   layout = #vhlo.array_v1<[]>
-   //          CHECK-SAME: }> : (!vhlo.token_v1) -> (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.token_v1)
-@@ -615,8 +660,9 @@
- }
- 
- // CHECK-LABEL: "default_outfeed"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @default_outfeed(%arg0: tensor<f32>, %arg1: !stablehlo.token) -> !stablehlo.token {
--  //      CHECK: "vhlo.outfeed_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.outfeed_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   outfeed_config = #vhlo.string_v1<"">
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.token_v1) -> !vhlo.token_v1
-   %0 = "stablehlo.outfeed"(%arg0, %arg1) : (tensor<f32>, !stablehlo.token) -> !stablehlo.token
-@@ -624,8 +670,9 @@
- }
- 
- // CHECK-LABEL: "default_recv"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_recv(%arg0: !stablehlo.token) -> (tensor<f32>, !stablehlo.token) {
--  //      CHECK: "vhlo.recv_v1"(%arg0) <{
-+  //      CHECK: "vhlo.recv_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME:   channel_type = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   is_host_transfer = #vhlo.bool_v1<false>
-@@ -637,8 +684,9 @@
- }
- 
- // CHECK-LABEL: "default_send"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @default_send(%arg0: tensor<f32>, %arg1: !stablehlo.token) -> !stablehlo.token {
--  //      CHECK: "vhlo.send_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.send_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME:   channel_type = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   is_host_transfer = #vhlo.bool_v1<false>
-@@ -650,8 +698,9 @@
- }
- 
- // CHECK-LABEL: "default_reduce_scatter"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_reduce_scatter(%arg0: tensor<16xf32>) -> tensor<16xf32> {
--  //               CHECK: "vhlo.reduce_scatter_v1"(%arg0) <{
-+  //               CHECK: "vhlo.reduce_scatter_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0], [1]]> : tensor<2x1xi64>>,
-   //          CHECK-SAME:   scatter_dimension = #vhlo.integer_v1<0 : i64>
-@@ -673,8 +722,9 @@
- }
- 
- // CHECK-LABEL: "default_reduce_window"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @default_reduce_window(%arg0: tensor<2x17x31x7xf32>, %arg1: tensor<f32>) -> tensor<2x16x30x7xf32> {
--  //               CHECK: "vhlo.reduce_window_v1"(%arg0, %arg1)  <{
-+  //               CHECK: "vhlo.reduce_window_v1"(%[[ARG0]], %[[ARG1]])  <{
-   //          CHECK-SAME:   base_dilations = #vhlo.tensor_v1<dense<1> : tensor<4xi64>>,
-   // CHECK-SAME{LITERAL}:   padding = #vhlo.tensor_v1<dense<0> : tensor<4x2xi64>>,
-   //          CHECK-SAME:   window_dilations = #vhlo.tensor_v1<dense<1> : tensor<4xi64>>,
-@@ -696,8 +746,9 @@
- }
- 
- // CHECK-LABEL: "default_scatter"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @default_scatter(%arg0: tensor<200x100x300xf32>, %arg1: tensor<10x2xi32>, %arg2: tensor<10x300xf32>) -> tensor<200x100x300xf32> {
--  //      CHECK: "vhlo.scatter_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.scatter_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   index_vector_dim = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   indices_are_sorted = #vhlo.bool_v1<false>,
-   // CHECK-SAME:   inserted_window_dims = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-@@ -725,8 +776,9 @@
- }
- 
- // CHECK-LABEL: "default_select_and_scatter"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @default_select_and_scatter(%arg0: tensor<10x24x24x64xf32>, %arg1: tensor<10x23x23x64xf32>, %arg2: tensor<f32>) -> tensor<10x24x24x64xf32> {
--  //      CHECK: "vhlo.select_and_scatter_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.select_and_scatter_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   padding = #vhlo.tensor_v1<dense<0> : tensor<4x2xi64>>,
-   // CHECK-SAME:   window_dimensions = #vhlo.tensor_v1<dense<[1, 2, 2, 1]> : tensor<4xi64>>,
-   // CHECK-SAME:   window_strides = #vhlo.tensor_v1<dense<1> : tensor<4xi64>>
-@@ -754,8 +806,9 @@
- }
- 
- // CHECK-LABEL: "default_sort"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_sort(%arg0: tensor<16xf32>) -> tensor<16xf32> {
--  //      CHECK: "vhlo.sort_v1"(%arg0) <{
-+  //      CHECK: "vhlo.sort_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   dimension = #vhlo.integer_v1<-1 : i64>
-   // CHECK-SAME:   is_stable = #vhlo.bool_v1<false>
-   // CHECK-SAME: }> ({
-@@ -774,29 +827,33 @@
- // ============ OPS ============
- 
- // CHECK-LABEL: "op_abs"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_abs(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.abs_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.abs_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.abs"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_add"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_add(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_after_all"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_after_all(%arg0: !stablehlo.token) -> !stablehlo.token {
--  // CHECK: "vhlo.after_all_v1"(%arg0) : (!vhlo.token_v1) -> !vhlo.token_v1
-+  // CHECK: "vhlo.after_all_v1"(%[[ARG0]]) : (!vhlo.token_v1) -> !vhlo.token_v1
-   %0 = "stablehlo.after_all"(%arg0) : (!stablehlo.token) -> !stablehlo.token
-   func.return %0 : !stablehlo.token
- }
- 
- // CHECK-LABEL: "op_all_gather"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_all_gather(%arg0: tensor<16x8xf32>) -> tensor<16x16xf32> {
--  //               CHECK: "vhlo.all_gather_v1"(%arg0) <{
-+  //               CHECK: "vhlo.all_gather_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   all_gather_dim = #vhlo.integer_v1<1 : i64>
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0], [1]]> : tensor<2x1xi64>>,
-@@ -812,8 +869,9 @@
- }
- 
- // CHECK-LABEL: "op_all_reduce"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_all_reduce(%arg0: tensor<f32>) -> tensor<f32> {
--  //               CHECK: "vhlo.all_reduce_v1"(%arg0) <{
-+  //               CHECK: "vhlo.all_reduce_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0], [1]]> : tensor<2x1xi64>>,
-   //          CHECK-SAME:   use_global_device_ids = #vhlo.bool_v1<true>
-@@ -836,7 +894,7 @@
- 
- // CHECK-LABEL: "op_all_reduce_with_promotable_types"
- func.func @op_all_reduce_with_promotable_types(%operand: tensor<f32>) -> tensor<f64> {
--  //  CHECK: "vhlo.all_reduce_v1"(%arg0)
-+  //  CHECK: "vhlo.all_reduce_v1"(%[[ARG0:.*]])
-   //  CHECK:   ^[[BB:bb.*]](%[[ARG1:arg.*]]: !vhlo.tensor_v1<!vhlo.f64_v1>, %[[ARG2:arg.*]]: !vhlo.tensor_v1<!vhlo.f64_v1>):
-   //  CHECK:     "vhlo.return_v1"(%[[VAL1:.*]]) : (!vhlo.tensor_v1<!vhlo.f64_v1>) -> ()
-   //  CHECK: }) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f64_v1>
-@@ -854,8 +912,9 @@
- }
- 
- // CHECK-LABEL: "op_all_to_all"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_all_to_all(%arg0: tensor<4x16xf32>) -> tensor<16x4xf32> {
--  //               CHECK: "vhlo.all_to_all_v1"(%arg0) <{
-+  //               CHECK: "vhlo.all_to_all_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<1 : i64>,
-   //          CHECK-SAME:   concat_dimension = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0, 1, 2, 3]]> : tensor<1x4xi64>>,
-@@ -873,22 +932,25 @@
- }
- 
- // CHECK-LABEL: "op_and"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_and(%arg0: tensor<i1>, %arg1: tensor<i1>) -> tensor<i1> {
--  // CHECK: "vhlo.and_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-+  // CHECK: "vhlo.and_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-   %0 = "stablehlo.and"(%arg0, %arg1) : (tensor<i1>, tensor<i1>) -> tensor<i1>
-   func.return %0 : tensor<i1>
- }
- 
- // CHECK-LABEL: "op_atan2"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_atan2(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.atan2_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.atan2_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.atan2"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_batch_norm_grad"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}}, %[[ARG3:.*]]: {{.*}}, %[[ARG4:.*]]: {{.*}})
- func.func @op_batch_norm_grad(%arg0: tensor<16x16x16x16xf32>, %arg1: tensor<16xf32>, %arg2: tensor<16xf32>, %arg3: tensor<16xf32>, %arg4: tensor<16x16x16x16xf32>) -> (tensor<16x16x16x16xf32>, tensor<16xf32>, tensor<16xf32>) {
--  //      CHECK: "vhlo.batch_norm_grad_v1"(%arg0, %arg1, %arg2, %arg3, %arg4) <{
-+  //      CHECK: "vhlo.batch_norm_grad_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]], %[[ARG3]], %[[ARG4]]) <{
-   // CHECK-SAME:   epsilon = #vhlo.float_v1<1.000000e-03 : !vhlo.f32_v1>,
-   // CHECK-SAME:   feature_index = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<16x16x16x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x16x16x16x!vhlo.f32_v1>) -> (!vhlo.tensor_v1<16x16x16x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>)
-@@ -900,8 +962,9 @@
- }
- 
- // CHECK-LABEL: "op_batch_norm_inference"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}}, %[[ARG3:.*]]: {{.*}}, %[[ARG4:.*]]: {{.*}})
- func.func @op_batch_norm_inference(%arg0: tensor<16x16x16x16xf32>, %arg1: tensor<16xf32>, %arg2: tensor<16xf32>, %arg3: tensor<16xf32>, %arg4: tensor<16xf32>) -> tensor<16x16x16x16xf32> {
--  //      CHECK: "vhlo.batch_norm_inference_v1"(%arg0, %arg1, %arg2, %arg3, %arg4) <{
-+  //      CHECK: "vhlo.batch_norm_inference_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]], %[[ARG3]], %[[ARG4]]) <{
-   // CHECK-SAME:   epsilon = #vhlo.float_v1<1.000000e-03 : !vhlo.f32_v1>,
-   // CHECK-SAME:   feature_index = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<16x16x16x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x16x16x16x!vhlo.f32_v1>
-@@ -913,8 +976,9 @@
- }
- 
- // CHECK-LABEL: "op_batch_norm_training"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_batch_norm_training(%arg0: tensor<16x16x16x16xf32>, %arg1: tensor<16xf32>, %arg2: tensor<16xf32>) -> (tensor<16x16x16x16xf32>, tensor<16xf32>, tensor<16xf32>) {
--  //      CHECK: "vhlo.batch_norm_training_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.batch_norm_training_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   epsilon = #vhlo.float_v1<1.000000e-03 : !vhlo.f32_v1>,
-   // CHECK-SAME:   feature_index = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<16x16x16x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>) -> (!vhlo.tensor_v1<16x16x16x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>)
-@@ -926,15 +990,17 @@
- }
- 
- // CHECK-LABEL: "op_bitcast_convert"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_bitcast_convert(%arg0: tensor<i32>) -> tensor<f32> {
--  // CHECK: "vhlo.bitcast_convert_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.bitcast_convert_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.bitcast_convert"(%arg0) : (tensor<i32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_broadcast_in_dim"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_broadcast_in_dim(%arg0: tensor<16xf32>) -> tensor<16x16xf32> {
--  //      CHECK: "vhlo.broadcast_in_dim_v1"(%arg0) <{
-+  //      CHECK: "vhlo.broadcast_in_dim_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   broadcast_dimensions = #vhlo.tensor_v1<dense<1> : tensor<1xi64>>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x16x!vhlo.f32_v1>
-   %0 = "stablehlo.broadcast_in_dim"(%arg0) {
-@@ -944,8 +1010,9 @@
- }
- 
- // CHECK-LABEL: "op_broadcast"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_broadcast(%arg0: tensor<16xf32>) -> tensor<16x16xf32> {
--  //      CHECK: "vhlo.broadcast_v1"(%arg0) <{
-+  //      CHECK: "vhlo.broadcast_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   broadcast_sizes = #vhlo.tensor_v1<dense<16> : tensor<1xi64>>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x16x!vhlo.f32_v1>
-   %0 = "stablehlo.broadcast"(%arg0) {
-@@ -955,9 +1022,10 @@
- }
- 
- // CHECK-LABEL: "op_case"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_case(%arg0: tensor<i32>, %arg1: tensor<f32>) -> tensor<f32> {
--  //      CHECK: "vhlo.case_v1"(%arg0) ({
--  // CHECK-NEXT:   "vhlo.return_v1"(%arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-+  //      CHECK: "vhlo.case_v1"(%[[ARG0]]) ({
-+  // CHECK-NEXT:   "vhlo.return_v1"(%[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-   // CHECK-NEXT: }) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.case"(%arg0) ({
-     "stablehlo.return"(%arg1) : (tensor<f32>) -> ()
-@@ -966,22 +1034,25 @@
- }
- 
- // CHECK-LABEL: "op_cbrt"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_cbrt(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.cbrt_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.cbrt_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.cbrt"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_ceil"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_ceil(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.ceil_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.ceil_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.ceil"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_cholesky"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_cholesky(%arg0: tensor<1x16x16xf32>) -> tensor<1x16x16xf32> {
--  //      CHECK: "vhlo.cholesky_v1"(%arg0) <{
-+  //      CHECK: "vhlo.cholesky_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   lower = #vhlo.bool_v1<true>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<1x16x16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<1x16x16x!vhlo.f32_v1>
-   %0 = "stablehlo.cholesky"(%arg0) {
-@@ -991,22 +1062,25 @@
- }
- 
- // CHECK-LABEL: "op_clamp"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_clamp(%arg0: tensor<f32>, %arg1: tensor<f32>, %arg2: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.clamp_v1"(%arg0, %arg1, %arg2) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.clamp_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.clamp"(%arg0, %arg1, %arg2) : (tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_count_leading_zeros"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_count_leading_zeros(%arg0: tensor<i32>) -> tensor<i32> {
--  // CHECK: "vhlo.count_leading_zeros_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-+  // CHECK: "vhlo.count_leading_zeros_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-   %0 = "stablehlo.count_leading_zeros"(%arg0) : (tensor<i32>) -> tensor<i32>
-   func.return %0 : tensor<i32>
- }
- 
- // CHECK-LABEL: "op_collective_permute"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_collective_permute(%arg0: tensor<16x8xf32>) -> tensor<16x8xf32> {
--  //               CHECK: "vhlo.collective_permute_v1"(%arg0) <{
-+  //               CHECK: "vhlo.collective_permute_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME{LITERAL}:   source_target_pairs = #vhlo.tensor_v1<dense<[[0, 1], [1, 2], [2, 3]]> : tensor<3x2xi64>>
-   //          CHECK-SAME: }> : (!vhlo.tensor_v1<16x8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x8x!vhlo.f32_v1>
-@@ -1018,8 +1092,9 @@
- }
- 
- // CHECK-LABEL: "op_compare"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_compare(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
--  //      CHECK: "vhlo.compare_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.compare_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   compare_type = #vhlo<comparison_type_v1 TOTALORDER>,
-   // CHECK-SAME:   comparison_direction = #vhlo<comparison_direction_v1 EQ>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-@@ -1031,15 +1106,17 @@
- }
- 
- // CHECK-LABEL: "op_complex"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_complex(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<complex<f32>> {
--  // CHECK: "vhlo.complex_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>
-+  // CHECK: "vhlo.complex_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>
-   %0 = "stablehlo.complex"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<complex<f32>>
-   func.return %0 : tensor<complex<f32>>
- }
- 
- // CHECK-LABEL: "op_concatenate"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_concatenate(%arg0: tensor<8xf32>, %arg1: tensor<8xf32>) -> tensor<16xf32> {
--  //      CHECK: "vhlo.concatenate_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.concatenate_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   dimension = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<8x!vhlo.f32_v1>, !vhlo.tensor_v1<8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x!vhlo.f32_v1>
-   %0 = "stablehlo.concatenate"(%arg0, %arg1) {
-@@ -1049,6 +1126,7 @@
- }
- 
- // CHECK-LABEL: "op_constant"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_constant(%arg0: tensor<f32>) -> tensor<f32> {
-   //      CHECK: "vhlo.constant_v1"() <{
-   // CHECK-SAME:   value = #vhlo.tensor_v1<dense<0.000000e+00> : tensor<f32>>
-@@ -1060,15 +1138,17 @@
- }
- 
- // CHECK-LABEL: "op_convert"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_convert(%arg0: tensor<i32>) -> tensor<f32> {
--  // CHECK: "vhlo.convert_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.convert_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.convert"(%arg0) : (tensor<i32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_convolution"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_convolution(%arg0: tensor<1x8x8x207xf32>, %arg1: tensor<3x3x207x16xf32>) -> tensor<1x7x7x16xf32> {
--  //      CHECK: "vhlo.convolution_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.convolution_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   batch_group_count = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   feature_group_count = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   input_batch_dimension = #vhlo.integer_v1<0 : i64>,
-@@ -1102,8 +1182,9 @@
- }
- 
- // CHECK-LABEL: "op_cosine"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_cosine(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.cosine_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.cosine_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.cosine"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
-@@ -1116,8 +1197,9 @@
- }
- 
- // CHECK-LABEL: "op_cross_replica_sum"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_cross_replica_sum(%arg0: tensor<f32>) -> tensor<f32> {
--  //               CHECK: "vhlo.cross-replica-sum_v1"(%arg0) <{
-+  //               CHECK: "vhlo.cross-replica-sum_v1"(%[[ARG0]]) <{
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0], [1]]> : tensor<2x1xi64>>
-   //          CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.cross-replica-sum"(%arg0) {
-@@ -1127,8 +1209,9 @@
- }
- 
- // CHECK-LABEL: "op_custom_call"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_custom_call(%arg0: tensor<f32>) -> tensor<f32> {
--  //      CHECK: "vhlo.custom_call_v1"(%arg0) <{
-+  //      CHECK: "vhlo.custom_call_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   api_version = #vhlo<api_version_v1 API_VERSION_STATUS_RETURNING>,
-   // CHECK-SAME:   backend_config = #vhlo.string_v1<"\08\03\1A\02">,
-   // CHECK-SAME:   call_target_name = #vhlo.string_v1<"foo">,
-@@ -1159,15 +1242,17 @@
- }
- 
- // CHECK-LABEL: "op_divide"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_divide(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.divide_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.divide_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.divide"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_dot_general"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_dot_general(%arg0: tensor<8x8x16xf32>, %arg1: tensor<8x16x8xf32>) -> tensor<8x8x8xf32> {
--  //      CHECK: "vhlo.dot_general_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.dot_general_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   lhs_batching_dimensions = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>,
-   // CHECK-SAME:   lhs_contracting_dimensions = #vhlo.tensor_v1<dense<2> : tensor<1xi64>>,
-   // CHECK-SAME:   precision_config = #vhlo.array_v1<[#vhlo<precision_v1 HIGHEST>, #vhlo<precision_v1 HIGHEST>]>,
-@@ -1187,8 +1272,9 @@
- }
- 
- // CHECK-LABEL: "op_dot"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_dot(%arg0: tensor<8x16xf32>, %arg1: tensor<16x8xf32>) -> tensor<8x8xf32> {
--  //      CHECK: "vhlo.dot_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.dot_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   precision_config = #vhlo.array_v1<[#vhlo<precision_v1 HIGHEST>, #vhlo<precision_v1 HIGHEST>]>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<8x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<8x8x!vhlo.f32_v1>
-   %0 = "stablehlo.dot"(%arg0, %arg1) {
-@@ -1198,8 +1284,9 @@
- }
- 
- // CHECK-LABEL: "op_dynamic_broadcast_in_dim"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_dynamic_broadcast_in_dim(%arg0: tensor<?x?xf32>, %arg1: tensor<2xindex>) -> tensor<?x?xf32> {
--  //      CHECK: "vhlo.dynamic_broadcast_in_dim_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.dynamic_broadcast_in_dim_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   broadcast_dimensions = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-   // CHECK-SAME:   known_expanding_dimensions = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>,
-   // CHECK-SAME:   known_nonexpanding_dimensions = #vhlo.tensor_v1<dense<1> : tensor<1xi64>>
-@@ -1213,8 +1300,9 @@
- }
- 
- // CHECK-LABEL: "op_dynamic_conv"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_dynamic_conv(%arg0: tensor<1x8x8x207xf32>, %arg1: tensor<3x3x207x16xf32>, %arg2: tensor<4xi32>) -> tensor<1x?x?x16xf32> {
--  //      CHECK: "vhlo.dynamic_conv_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.dynamic_conv_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   batch_group_count = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   feature_group_count = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   input_batch_dimension = #vhlo.integer_v1<0 : i64>,
-@@ -1248,8 +1336,9 @@
- }
- 
- // CHECK-LABEL: "op_dynamic_gather"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_dynamic_gather(%arg0 : tensor<2x4x9xf32>, %arg1 : tensor<1x5x2xi32>, %arg2 : tensor<3xi32>) -> tensor<1x5x8xf32> {
--  //      CHECK: "vhlo.dynamic_gather_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.dynamic_gather_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   collapsed_slice_dims = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-   // CHECK-SAME:   index_vector_dim = #vhlo.integer_v1<2 : i64>,
-   // CHECK-SAME:   indices_are_sorted = #vhlo.bool_v1<true>,
-@@ -1269,8 +1358,9 @@
- }
- 
- // CHECK-LABEL: "op_dynamic_iota"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_dynamic_iota(%arg0: tensor<1xindex>) -> tensor<?xf32> {
--  //      CHECK: "vhlo.dynamic_iota_v1"(%arg0) <{
-+  //      CHECK: "vhlo.dynamic_iota_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   iota_dimension = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<1x!vhlo.index_v1>) -> !vhlo.tensor_v1<?x!vhlo.f32_v1>
-   %0 = "stablehlo.dynamic_iota"(%arg0) {
-@@ -1280,22 +1370,25 @@
- }
- 
- // CHECK-LABEL: "op_dynamic_pad"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}}, %[[ARG3:.*]]: {{.*}}, %[[ARG4:.*]]: {{.*}})
- func.func @op_dynamic_pad(%arg0: tensor<?xf32>, %arg1: tensor<f32>, %arg2: tensor<1xindex>, %arg3: tensor<1xindex>, %arg4: tensor<1xindex>) -> tensor<?xf32> {
--  // CHECK: "vhlo.dynamic_pad_v1"(%arg0, %arg1, %arg2, %arg3, %arg4) : (!vhlo.tensor_v1<?x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>) -> !vhlo.tensor_v1<?x!vhlo.f32_v1>
-+  // CHECK: "vhlo.dynamic_pad_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]], %[[ARG3]], %[[ARG4]]) : (!vhlo.tensor_v1<?x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>) -> !vhlo.tensor_v1<?x!vhlo.f32_v1>
-   %0 = "stablehlo.dynamic_pad"(%arg0, %arg1, %arg2, %arg3, %arg4) : (tensor<?xf32>, tensor<f32>, tensor<1xindex>, tensor<1xindex>, tensor<1xindex>) -> tensor<?xf32>
-   func.return %0 : tensor<?xf32>
- }
- 
- // CHECK-LABEL: "op_dynamic_reshape"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_dynamic_reshape(%arg0: tensor<16xf32>, %arg1: tensor<2xindex>) -> tensor<?x?xf32> {
--  // CHECK: "vhlo.dynamic_reshape_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<2x!vhlo.index_v1>) -> !vhlo.tensor_v1<?x?x!vhlo.f32_v1>
-+  // CHECK: "vhlo.dynamic_reshape_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<2x!vhlo.index_v1>) -> !vhlo.tensor_v1<?x?x!vhlo.f32_v1>
-   %0 = "stablehlo.dynamic_reshape"(%arg0, %arg1) : (tensor<16xf32>, tensor<2xindex>) -> tensor<?x?xf32>
-   func.return %0 : tensor<?x?xf32>
- }
- 
- // CHECK-LABEL: "op_dynamic_slice"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_dynamic_slice(%arg0: tensor<16xf32>, %arg1: tensor<i64>) -> tensor<4xf32> {
--  //      CHECK: "vhlo.dynamic_slice_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.dynamic_slice_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   slice_sizes = #vhlo.tensor_v1<dense<4> : tensor<1xi64>>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.i64_v1>) -> !vhlo.tensor_v1<4x!vhlo.f32_v1>
-   %0 = "stablehlo.dynamic_slice"(%arg0, %arg1) {
-@@ -1305,15 +1398,17 @@
- }
- 
- // CHECK-LABEL: "op_dynamic_update_slice"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_dynamic_update_slice(%arg0: tensor<16xf32>, %arg1: tensor<4xf32>, %arg2: tensor<i64>) -> tensor<16xf32> {
--  // CHECK: "vhlo.dynamic_update_slice_v1"(%arg0, %arg1, %arg2) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<4x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.i64_v1>) -> !vhlo.tensor_v1<16x!vhlo.f32_v1>
-+  // CHECK: "vhlo.dynamic_update_slice_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<4x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.i64_v1>) -> !vhlo.tensor_v1<16x!vhlo.f32_v1>
-   %0 = "stablehlo.dynamic_update_slice"(%arg0, %arg1, %arg2) : (tensor<16xf32>, tensor<4xf32>, tensor<i64>) -> tensor<16xf32>
-   func.return %0 : tensor<16xf32>
- }
- 
- // CHECK-LABEL: "op_einsum"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_einsum(%arg0: tensor<8x16xf32>, %arg1: tensor<16x8xf32>) -> tensor<8x8xf32> {
--  //      CHECK: "vhlo.einsum_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.einsum_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   einsum_config = #vhlo.string_v1<"ab,bc->ac">
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<8x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<8x8x!vhlo.f32_v1>
-   %0 = "stablehlo.einsum"(%arg0, %arg1) {
-@@ -1323,22 +1418,25 @@
- }
- 
- // CHECK-LABEL: "op_exponential_minus_one"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_exponential_minus_one(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.exponential_minus_one_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.exponential_minus_one_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.exponential_minus_one"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_exponential"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_exponential(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.exponential_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.exponential_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.exponential"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_fft"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_fft(%arg0: tensor<16xcomplex<f32>>) -> tensor<16xcomplex<f32>> {
--  //      CHECK: "vhlo.fft_v1"(%arg0) <{
-+  //      CHECK: "vhlo.fft_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   fft_length = #vhlo.tensor_v1<dense<16> : tensor<1xi64>>,
-   // CHECK-SAME:   fft_type = #vhlo<fft_type_v1 FFT>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<16x!vhlo.complex_v1<!vhlo.f32_v1>>) -> !vhlo.tensor_v1<16x!vhlo.complex_v1<!vhlo.f32_v1>>
-@@ -1350,8 +1448,9 @@
- }
- 
- // CHECK-LABEL: "op_floor"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_floor(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.floor_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.floor_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.floor"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
-@@ -1364,16 +1463,17 @@
-   // CHECK-SAME:   sym_name = #vhlo.string_v1<"op_func">,
-   // CHECK-SAME:   sym_visibility = #vhlo.string_v1<"private">
-   // CHECK-SAME: }> ({
--  // CHECK-NEXT:   ^[[BB:bb.*]](%arg0: !vhlo.tensor_v1<!vhlo.f32_v1>):
--  // CHECK-NEXT:     "vhlo.return_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-+  // CHECK-NEXT:   ^[[BB:bb.*]](%[[ARG0:.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
-+  // CHECK-NEXT:     "vhlo.return_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-   // CHECK-NEXT: }) : () -> ()
- 
-   func.return %arg0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_gather"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_gather(%arg0 : tensor<2x4x9xf32>, %arg1 : tensor<1x5x2xi32>) -> tensor<1x5x1xf32> {
--  //      CHECK: "vhlo.gather_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.gather_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   collapsed_slice_dims = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-   // CHECK-SAME:   index_vector_dim = #vhlo.integer_v1<2 : i64>,
-   // CHECK-SAME:   indices_are_sorted = #vhlo.bool_v1<true>,
-@@ -1395,8 +1495,9 @@
- }
- 
- // CHECK-LABEL: "op_get_dimension_size"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_get_dimension_size(%arg0: tensor<?xf32>) -> tensor<i32> {
--  //      CHECK: "vhlo.get_dimension_size_v1"(%arg0) <{
-+  //      CHECK: "vhlo.get_dimension_size_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   dimension = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<?x!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-   %0 = "stablehlo.get_dimension_size"(%arg0) {
-@@ -1406,8 +1507,9 @@
- }
- 
- // CHECK-LABEL: "op_get_tuple_element"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_get_tuple_element(%arg0: tuple<tensor<f32>, tensor<i32>>) -> tensor<f32> {
--  //      CHECK: "vhlo.get_tuple_element_v1"(%arg0) <{
-+  //      CHECK: "vhlo.get_tuple_element_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   index = #vhlo.integer_v1<0 : i32>
-   // CHECK-SAME: }> : (!vhlo.tuple_v1<!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.get_tuple_element"(%arg0) {
-@@ -1417,11 +1519,12 @@
- }
- 
- // CHECK-LABEL: "op_if"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_if(%arg0: tensor<i1>, %arg1: tensor<f32>, %arg2: tensor<f32>) -> tensor<f32> {
--  //      CHECK: "vhlo.if_v1"(%arg0) ({
--  // CHECK-NEXT:   "vhlo.return_v1"(%arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-+  //      CHECK: "vhlo.if_v1"(%[[ARG0]]) ({
-+  // CHECK-NEXT:   "vhlo.return_v1"(%[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-   // CHECK-NEXT: }, {
--  // CHECK-NEXT:   "vhlo.return_v1"(%arg2) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-+  // CHECK-NEXT:   "vhlo.return_v1"(%[[ARG2]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-   // CHECK-NEXT: }) : (!vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.if"(%arg0) ({
-     "stablehlo.return"(%arg1) : (tensor<f32>) -> ()
-@@ -1432,15 +1535,17 @@
- }
- 
- // CHECK-LABEL: "op_imag"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_imag(%arg0: tensor<complex<f32>>) -> tensor<f32> {
--  // CHECK: "vhlo.imag_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.imag_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.imag"(%arg0) : (tensor<complex<f32>>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_infeed"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_infeed(%arg0: !stablehlo.token) -> (tensor<f32>, !stablehlo.token) {
--  //               CHECK: "vhlo.infeed_v1"(%arg0) <{
-+  //               CHECK: "vhlo.infeed_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   infeed_config = #vhlo.string_v1<"foo">,
-   // CHECK-SAME{LITERAL}:   layout = #vhlo.array_v1<[#vhlo.array_v1<[]>]>
-   //          CHECK-SAME: }> : (!vhlo.token_v1) -> (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.token_v1)
-@@ -1463,36 +1568,41 @@
- }
- 
- // CHECK-LABEL: "op_is_finite"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_is_finite(%arg0: tensor<f32>) -> tensor<i1> {
--  // CHECK: "vhlo.is_finite_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-+  // CHECK: "vhlo.is_finite_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-   %0 = "stablehlo.is_finite"(%arg0) : (tensor<f32>) -> tensor<i1>
-   func.return %0 : tensor<i1>
- }
- 
- // CHECK-LABEL: "op_log"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_log(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.log_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.log_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.log"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_log_plus_one"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_log_plus_one(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.log_plus_one_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.log_plus_one_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.log_plus_one"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_logistic"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_logistic(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.logistic_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.logistic_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.logistic"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_map"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_map(%arg0: tensor<16xf32>) -> tensor<16xf32> {
--  //      CHECK: "vhlo.map_v1"(%arg0) <{
-+  //      CHECK: "vhlo.map_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   dimensions = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>
-   // CHECK-SAME: }> ({
-   // CHECK-NEXT:   ^[[BB:bb.*]](%[[ARG1:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
-@@ -1510,57 +1620,65 @@
- }
- 
- // CHECK-LABEL: "op_maximum"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_maximum(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.maximum_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.maximum_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.maximum"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_minimum"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_minimum(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.minimum_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.minimum_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.minimum"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_multiply"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_multiply(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.multiply_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.multiply_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.multiply"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_negate"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_negate(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.negate_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.negate_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.negate"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_not"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_not(%arg0: tensor<i1>) -> tensor<i1> {
--  // CHECK: "vhlo.not_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-+  // CHECK: "vhlo.not_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-   %0 = "stablehlo.not"(%arg0) : (tensor<i1>) -> tensor<i1>
-   func.return %0 : tensor<i1>
- }
- 
- // CHECK-LABEL: "op_optimization_barrier"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_optimization_barrier(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.optimization_barrier_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.optimization_barrier_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.optimization_barrier"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_or"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_or(%arg0: tensor<i1>, %arg1: tensor<i1>) -> tensor<i1> {
--  // CHECK: "vhlo.or_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-+  // CHECK: "vhlo.or_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-   %0 = "stablehlo.or"(%arg0, %arg1) : (tensor<i1>, tensor<i1>) -> tensor<i1>
-   func.return %0 : tensor<i1>
- }
- 
- // CHECK-LABEL: "op_outfeed"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_outfeed(%arg0: tensor<f32>, %arg1: !stablehlo.token) -> !stablehlo.token {
--  //      CHECK: "vhlo.outfeed_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.outfeed_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   outfeed_config = #vhlo.string_v1<"foo">
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.token_v1) -> !vhlo.token_v1
-   %0 = "stablehlo.outfeed"(%arg0, %arg1) {
-@@ -1570,8 +1688,9 @@
- }
- 
- // CHECK-LABEL: "op_pad"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_pad(%arg0: tensor<8xf32>, %arg1: tensor<f32>) -> tensor<16xf32> {
--  //      CHECK: "vhlo.pad_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.pad_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   edge_padding_high = #vhlo.tensor_v1<dense<4> : tensor<1xi64>>,
-   // CHECK-SAME:   edge_padding_low = #vhlo.tensor_v1<dense<4> : tensor<1xi64>>,
-   // CHECK-SAME:   interior_padding = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>
-@@ -1585,36 +1704,41 @@
- }
- 
- // CHECK-LABEL: "op_popcnt"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_popcnt(%arg0: tensor<i32>) -> tensor<i32> {
--  // CHECK: "vhlo.popcnt_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-+  // CHECK: "vhlo.popcnt_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-   %0 = "stablehlo.popcnt"(%arg0) : (tensor<i32>) -> tensor<i32>
-   func.return %0 : tensor<i32>
- }
- 
- // CHECK-LABEL: "op_power"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_power(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.power_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.power_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.power"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_real_dynamic_slice"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}}, %[[ARG3:.*]]: {{.*}})
- func.func @op_real_dynamic_slice(%arg0: tensor<?xf32>, %arg1: tensor<1xindex>, %arg2: tensor<1xindex>, %arg3: tensor<1xindex>) -> tensor<?xf32> {
--  // CHECK: "vhlo.real_dynamic_slice_v1"(%arg0, %arg1, %arg2, %arg3) : (!vhlo.tensor_v1<?x!vhlo.f32_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>) -> !vhlo.tensor_v1<?x!vhlo.f32_v1>
-+  // CHECK: "vhlo.real_dynamic_slice_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]], %[[ARG3]]) : (!vhlo.tensor_v1<?x!vhlo.f32_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>) -> !vhlo.tensor_v1<?x!vhlo.f32_v1>
-   %0 = "stablehlo.real_dynamic_slice"(%arg0, %arg1, %arg2, %arg3) : (tensor<?xf32>, tensor<1xindex>, tensor<1xindex>, tensor<1xindex>) -> tensor<?xf32>
-   func.return %0 : tensor<?xf32>
- }
- 
- // CHECK-LABEL: "op_real"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_real(%arg0: tensor<complex<f32>>) -> tensor<f32> {
--  // CHECK: "vhlo.real_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.real_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.real"(%arg0) : (tensor<complex<f32>>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_recv"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_recv(%arg0: !stablehlo.token) -> (tensor<f32>, !stablehlo.token) {
--  //      CHECK: "vhlo.recv_v1"(%arg0) <{
-+  //      CHECK: "vhlo.recv_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME:   channel_type = #vhlo.integer_v1<3 : i64>,
-   // CHECK-SAME:   is_host_transfer = #vhlo.bool_v1<true>
-@@ -1627,8 +1751,9 @@
- }
- 
- // CHECK-LABEL: "op_reduce"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_reduce(%arg0: tensor<16xf32>, %arg1: tensor<f32>) -> tensor<f32> {
--  //  CHECK: "vhlo.reduce_v1"(%arg0, %arg1)
-+  //  CHECK: "vhlo.reduce_v1"(%[[ARG0]], %[[ARG1]])
-   //  CHECK:   ^[[BB:bb.*]](%[[ARG1:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>, %[[ARG2:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
-   //  CHECK:     "vhlo.return_v1"(%[[VAL1:.*]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-   //  CHECK: }) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-@@ -1643,8 +1768,9 @@
- }
- 
- // CHECK-LABEL: "op_reduce_precision"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_reduce_precision(%arg0: tensor<f32>) -> tensor<f32> {
--  //      CHECK: "vhlo.reduce_precision_v1"(%arg0) <{
-+  //      CHECK: "vhlo.reduce_precision_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   exponent_bits = #vhlo.integer_v1<8 : i32>
-   // CHECK-SAME:   mantissa_bits = #vhlo.integer_v1<10 : i32>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-@@ -1658,7 +1784,7 @@
- // CHECK_lABEL: "op_reduce_with_promotable_types"
- func.func @op_reduce_with_promotable_types(%arg0: tensor<4x4xf32>, %arg1 : tensor<f32>)
-     -> (tensor<4xf64>) {
--  //  CHECK: "vhlo.reduce_v1"(%arg0, %arg1)
-+  //  CHECK: "vhlo.reduce_v1"(%[[ARG0:.*]], %[[ARG1:.*]])
-   //  CHECK:   ^[[BB:bb.*]](%[[ARG1:arg.*]]: !vhlo.tensor_v1<!vhlo.f64_v1>, %[[ARG2:arg.*]]: !vhlo.tensor_v1<!vhlo.f64_v1>):
-   //  CHECK:     "vhlo.return_v1"(%[[VAL1:.*]]) : (!vhlo.tensor_v1<!vhlo.f64_v1>) -> ()
-   //  CHECK: }) : (!vhlo.tensor_v1<4x4x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<4x!vhlo.f64_v1>
-@@ -1673,8 +1799,9 @@
- }
- 
- // CHECK-LABEL: "op_reduce_scatter"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_reduce_scatter(%arg0: tensor<16xf32>) -> tensor<16xf32> {
--  //               CHECK: "vhlo.reduce_scatter_v1"(%arg0) <{
-+  //               CHECK: "vhlo.reduce_scatter_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0], [1]]> : tensor<2x1xi64>>,
-   //          CHECK-SAME:   scatter_dimension = #vhlo.integer_v1<0 : i64>
-@@ -1699,7 +1826,7 @@
- 
- // CHECK_lABEL: "op_reduce_scatter_with_promotable_types"
- func.func @op_reduce_scatter_with_promotable_types(%data: tensor<4x16xf32>) -> tensor<4x4xf64> {
--  //  CHECK: "vhlo.reduce_scatter_v1"(%arg0)
-+  //  CHECK: "vhlo.reduce_scatter_v1"(%[[ARG0:.*]])
-   //  CHECK:   ^[[BB:bb.*]](%[[ARG1:arg.*]]: !vhlo.tensor_v1<!vhlo.f64_v1>, %[[ARG2:arg.*]]: !vhlo.tensor_v1<!vhlo.f64_v1>):
-   //  CHECK:     "vhlo.return_v1"(%[[VAL1:.*]]) : (!vhlo.tensor_v1<!vhlo.f64_v1>) -> ()
-   //  CHECK: }) : (!vhlo.tensor_v1<4x16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<4x4x!vhlo.f64_v1>
-@@ -1716,8 +1843,9 @@
- 
- 
- // CHECK-LABEL: "op_reduce_window"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_reduce_window(%arg0: tensor<2x17x31x7xf32>, %arg1: tensor<f32>) -> tensor<2x9x16x7xf32> {
--  //               CHECK: "vhlo.reduce_window_v1"(%arg0, %arg1) <{
-+  //               CHECK: "vhlo.reduce_window_v1"(%[[ARG0]], %[[ARG1]]) <{
-   //          CHECK-SAME:   base_dilations = #vhlo.tensor_v1<dense<[1, 2, 2, 1]> : tensor<4xi64>>,
-   // CHECK-SAME{LITERAL}:   padding = #vhlo.tensor_v1<dense<[[0, 0], [2, 0], [0, 2], [0, 0]]> : tensor<4x2xi64>>,
-   //          CHECK-SAME:   window_dilations = #vhlo.tensor_v1<dense<[1, 2, 2, 1]> : tensor<4xi64>>,
-@@ -1746,7 +1874,7 @@
- func.func @op_reduce_window_with_promotable_types(%arg0: tensor<4x2xf32>,
-     %arg1: tensor<4x2xf32>, %init0: tensor<f32>, %init1: tensor<f32>) ->
-     (tensor<2x2xf64>, tensor<2x2xf32>) {
--  //  CHECK: "vhlo.reduce_window_v1"(%arg0, %arg1, %arg2, %arg3)
-+  //  CHECK: "vhlo.reduce_window_v1"(%[[ARG0:.*]], %[[ARG1:.*]], %[[ARG2:.*]], %[[ARG3:.*]])
-   //  CHECK:   ^[[BB:bb.*]](%[[ARG1:arg.*]]: !vhlo.tensor_v1<!vhlo.f64_v1>, %[[ARG2:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>, %[[ARG3:arg.*]]: !vhlo.tensor_v1<!vhlo.f64_v1>, %[[ARG4:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
-   //  CHECK:     "vhlo.return_v1"(%[[VAL1:.*]], %[[VAL2:.*]]) : (!vhlo.tensor_v1<!vhlo.f64_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-   //  CHECK: }) : (!vhlo.tensor_v1<4x2x!vhlo.f32_v1>, !vhlo.tensor_v1<4x2x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> (!vhlo.tensor_v1<2x2x!vhlo.f64_v1>, !vhlo.tensor_v1<2x2x!vhlo.f32_v1>)
-@@ -1766,8 +1894,9 @@
- }
- 
- // CHECK-LABEL: "op_remainder"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_remainder(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.remainder_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.remainder_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.remainder"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
-@@ -1787,16 +1916,18 @@
- }
- 
- // CHECK-LABEL: "op_reshape"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_reshape(%arg0: tensor<16xf32>) -> tensor<4x4xf32> {
--  // CHECK: "vhlo.reshape_v1"(%arg0) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<4x4x!vhlo.f32_v1>
-+  // CHECK: "vhlo.reshape_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<4x4x!vhlo.f32_v1>
-   %0 = "stablehlo.reshape"(%arg0) : (tensor<16xf32>) -> tensor<4x4xf32>
-   func.return %0 : tensor<4x4xf32>
- }
- 
- // CHECK-LABEL: "op_return"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_return(%arg0: tensor<i32>, %arg1: tensor<f32>) -> tensor<f32> {
--  //      CHECK: "vhlo.case_v1"(%arg0) ({
--  // CHECK-NEXT:   "vhlo.return_v1"(%arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-+  //      CHECK: "vhlo.case_v1"(%[[ARG0]]) ({
-+  // CHECK-NEXT:   "vhlo.return_v1"(%[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-   // CHECK-NEXT: }) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.case"(%arg0) ({
-     "stablehlo.return"(%arg1) : (tensor<f32>) -> ()
-@@ -1805,8 +1936,9 @@
- }
- 
- // CHECK-LABEL: "op_reverse"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_reverse(%arg0: tensor<16xf32>) -> tensor<16xf32> {
--  //      CHECK: "vhlo.reverse_v1"(%arg0) <{
-+  //      CHECK: "vhlo.reverse_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   dimensions = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x!vhlo.f32_v1>
-   %0 = "stablehlo.reverse"(%arg0) {
-@@ -1816,8 +1948,9 @@
- }
- 
- // CHECK-LABEL: "op_rng_bit_generator"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_rng_bit_generator(%arg0: tensor<f32>) -> (tensor<f32>, tensor<f32>) {
--  //      CHECK: "vhlo.rng_bit_generator_v1"(%arg0) <{
-+  //      CHECK: "vhlo.rng_bit_generator_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   rng_algorithm = #vhlo<rng_algorithm_v1 PHILOX>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>)
-   %0:2 = "stablehlo.rng_bit_generator"(%arg0) {
-@@ -1827,8 +1960,9 @@
- }
- 
- // CHECK-LABEL: "op_rng"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_rng(%arg0: tensor<f32>, %arg1: tensor<f32>, %arg2: tensor<0xindex>) -> tensor<f32> {
--  //      CHECK: "vhlo.rng_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.rng_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   rng_distribution = #vhlo<rng_distribution_v1 NORMAL>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<0x!vhlo.index_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.rng"(%arg0, %arg1, %arg2) {
-@@ -1838,29 +1972,33 @@
- }
- 
- // CHECK-LABEL: "op_round_nearest_afz"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_round_nearest_afz(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.round_nearest_afz_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.round_nearest_afz_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.round_nearest_afz"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_round_nearest_even"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_round_nearest_even(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.round_nearest_even_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.round_nearest_even_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.round_nearest_even"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_rsqrt"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_rsqrt(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.rsqrt_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.rsqrt_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.rsqrt"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_scatter"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_scatter(%arg0: tensor<200x100x300xf32>, %arg1: tensor<10x2xi32>, %arg2: tensor<10x300xf32>) -> tensor<200x100x300xf32> {
--  //      CHECK: "vhlo.scatter_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.scatter_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   index_vector_dim = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   indices_are_sorted = #vhlo.bool_v1<true>,
-   // CHECK-SAME:   inserted_window_dims = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-@@ -1893,7 +2031,7 @@
- func.func @op_scatter_with_promotable_types(%input_tensor: tensor<200x100x300xf32>,
-     %scatter_indices: tensor<10x2xi32>, %updates: tensor<10x300xf32>) ->
-       tensor<200x100x300xf64> {
--  //  CHECK: "vhlo.scatter_v1"(%arg0, %arg1, %arg2)
-+  //  CHECK: "vhlo.scatter_v1"(%[[ARG0:.*]], %[[ARG1:.*]], %[[ARG2:.*]])
-   //  CHECK:   ^[[BB:bb.*]](%[[ARG1:arg.*]]: !vhlo.tensor_v1<!vhlo.f64_v1>, %[[ARG2:arg.*]]: !vhlo.tensor_v1<!vhlo.f64_v1>):
-   //  CHECK:     "vhlo.return_v1"(%[[VAL1:.*]]) : (!vhlo.tensor_v1<!vhlo.f64_v1>) -> ()
-   //  CHECK: }) : (!vhlo.tensor_v1<200x100x300x!vhlo.f32_v1>, !vhlo.tensor_v1<10x2x!vhlo.i32_v1>, !vhlo.tensor_v1<10x300x!vhlo.f32_v1>) -> !vhlo.tensor_v1<200x100x300x!vhlo.f64_v1>
-@@ -1916,8 +2054,9 @@
- }
- 
- // CHECK-LABEL: "op_select_and_scatter"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_select_and_scatter(%arg0: tensor<10x24x24x64xf32>, %arg1: tensor<12x13x13x66xf32>, %arg2: tensor<f32>) -> tensor<10x24x24x64xf32> {
--  //      CHECK: "vhlo.select_and_scatter_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.select_and_scatter_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   padding = #vhlo.tensor_v1<dense<1> : tensor<4x2xi64>>,
-   // CHECK-SAME:   window_dimensions = #vhlo.tensor_v1<dense<[1, 2, 2, 1]> : tensor<4xi64>>,
-   // CHECK-SAME:   window_strides = #vhlo.tensor_v1<dense<[1, 2, 2, 1]> : tensor<4xi64>>
-@@ -1947,8 +2086,9 @@
- }
- 
- // CHECK-LABEL: "op_select_and_scatter_with_promotable_types"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_select_and_scatter_with_promotable_types(%arg0: tensor<10x24x24x64xf32>, %arg1: tensor<12x13x13x66xf32>, %arg2: tensor<f32>) -> tensor<10x24x24x64xf64> {
--  // CHECK: "vhlo.select_and_scatter_v1"(%arg0, %arg1, %arg2)
-+  // CHECK: "vhlo.select_and_scatter_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]])
-   // CHECK:   ^[[BB:bb.*]](%[[ARG1:arg.*]]: !vhlo.tensor_v1<!vhlo.f64_v1>, %[[ARG2:arg.*]]: !vhlo.tensor_v1<!vhlo.f64_v1>):
-   // CHECK:     %[[VAL:.*]] = "vhlo.add_v1"(%[[ARG1]], %[[ARG2]]) : (!vhlo.tensor_v1<!vhlo.f64_v1>, !vhlo.tensor_v1<!vhlo.f64_v1>) -> !vhlo.tensor_v1<!vhlo.f64_v1>
-   // CHECK:     "vhlo.return_v1"(%[[VAL]]) : (!vhlo.tensor_v1<!vhlo.f64_v1>) -> ()
-@@ -1970,15 +2110,17 @@
- }
- 
- // CHECK-LABEL: "op_select"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_select(%arg0: tensor<i1>, %arg1: tensor<f32>, %arg2: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.select_v1"(%arg0, %arg1, %arg2) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.select_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.select"(%arg0, %arg1, %arg2) : (tensor<i1>, tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_send"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_send(%arg0: tensor<f32>, %arg1: !stablehlo.token) -> !stablehlo.token {
--  //      CHECK: "vhlo.send_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.send_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME:   channel_type = #vhlo.integer_v1<2 : i64>,
-   // CHECK-SAME:   is_host_transfer = #vhlo.bool_v1<true>
-@@ -1991,8 +2133,9 @@
- }
- 
- // CHECK-LABEL: "op_set_dimension_size"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_set_dimension_size(%arg0: tensor<?xf32>, %arg1: tensor<i32>) -> tensor<16xf32> {
--  //      CHECK: "vhlo.set_dimension_size_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.set_dimension_size_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   dimension = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<?x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<16x!vhlo.f32_v1>
-   %0 = "stablehlo.set_dimension_size"(%arg0, %arg1) {
-@@ -2002,43 +2145,49 @@
- }
- 
- // CHECK-LABEL: "op_shift_left"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_shift_left(%arg0: tensor<i32>, %arg1: tensor<i32>) -> tensor<i32> {
--  // CHECK: "vhlo.shift_left_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.i32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-+  // CHECK: "vhlo.shift_left_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-   %0 = "stablehlo.shift_left"(%arg0, %arg1) : (tensor<i32>, tensor<i32>) -> tensor<i32>
-   func.return %0 : tensor<i32>
- }
- 
- // CHECK-LABEL: "op_shift_right_arithmetic"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_shift_right_arithmetic(%arg0: tensor<i32>, %arg1: tensor<i32>) -> tensor<i32> {
--  // CHECK: "vhlo.shift_right_arithmetic_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.i32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-+  // CHECK: "vhlo.shift_right_arithmetic_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-   %0 = "stablehlo.shift_right_arithmetic"(%arg0, %arg1) : (tensor<i32>, tensor<i32>) -> tensor<i32>
-   func.return %0 : tensor<i32>
- }
- 
- // CHECK-LABEL: "op_shift_right_logical"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_shift_right_logical(%arg0: tensor<i32>, %arg1: tensor<i32>) -> tensor<i32> {
--  // CHECK: "vhlo.shift_right_logical_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.i32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-+  // CHECK: "vhlo.shift_right_logical_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-   %0 = "stablehlo.shift_right_logical"(%arg0, %arg1) : (tensor<i32>, tensor<i32>) -> tensor<i32>
-   func.return %0 : tensor<i32>
- }
- 
- // CHECK-LABEL: "op_sign"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_sign(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.sign_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.sign_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.sign"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_sine"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_sine(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.sine_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.sine_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.sine"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_slice"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_slice(%arg0: tensor<16xf32>) -> tensor<4xf32> {
--  //      CHECK: "vhlo.slice_v1"(%arg0) <{
-+  //      CHECK: "vhlo.slice_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   limit_indices = #vhlo.tensor_v1<dense<4> : tensor<1xi64>>,
-   // CHECK-SAME:   start_indices = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>,
-   // CHECK-SAME:   strides = #vhlo.tensor_v1<dense<1> : tensor<1xi64>>
-@@ -2052,8 +2201,9 @@
- }
- 
- // CHECK-LABEL: "op_sort"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_sort(%arg0: tensor<16xf32>) -> tensor<16xf32> {
--  //      CHECK: "vhlo.sort_v1"(%arg0) <{
-+  //      CHECK: "vhlo.sort_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   dimension = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME:   is_stable = #vhlo.bool_v1<true>
-   // CHECK-SAME: }> ({
-@@ -2073,29 +2223,33 @@
- }
- 
- // CHECK-LABEL: "op_sqrt"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_sqrt(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.sqrt_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.sqrt_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.sqrt"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_subtract"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_subtract(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.subtract_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.subtract_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.subtract"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_tanh"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_tanh(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.tanh_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.tanh_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.tanh"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_torch_index_select"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_torch_index_select(%arg0: tensor<5x1x5xf32>, %arg1: tensor<2xi32>) ->  tensor<2x1x5xf32> {
--  //      CHECK: "vhlo.torch_index_select_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.torch_index_select_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   batch_dims = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME:   dim = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<5x1x5x!vhlo.f32_v1>, !vhlo.tensor_v1<2x!vhlo.i32_v1>) -> !vhlo.tensor_v1<2x1x5x!vhlo.f32_v1>
-@@ -2107,8 +2261,9 @@
- }
- 
- // CHECK-LABEL: "op_trace"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_trace(%arg0: tensor<f32>) {
--  //      CHECK: "vhlo.trace_v1"(%arg0) <{
-+  //      CHECK: "vhlo.trace_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   tag = #vhlo.string_v1<"foo">
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-   "stablehlo.trace"(%arg0) {
-@@ -2118,8 +2273,9 @@
- }
- 
- // CHECK-LABEL: "op_transpose"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_transpose(%arg0: tensor<16x8xf32>) ->  tensor<8x16xf32> {
--  //      CHECK: "vhlo.transpose_v1"(%arg0) <{
-+  //      CHECK: "vhlo.transpose_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   permutation = #vhlo.tensor_v1<dense<[1, 0]> : tensor<2xi64>>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<16x8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<8x16x!vhlo.f32_v1>
-   %0 = "stablehlo.transpose"(%arg0) {
-@@ -2129,8 +2285,9 @@
- }
- 
- // CHECK-LABEL: "op_triangular_solve"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_triangular_solve(%arg0: tensor<16x16xf32>, %arg1: tensor<16x16xf32>) ->  tensor<16x16xf32> {
--  //      CHECK: "vhlo.triangular_solve_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.triangular_solve_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   left_side = #vhlo.bool_v1<true>,
-   // CHECK-SAME:   lower = #vhlo.bool_v1<true>,
-   // CHECK-SAME:   transpose_a = #vhlo<transpose_v1 NO_TRANSPOSE>,
-@@ -2146,15 +2303,17 @@
- }
- 
- // CHECK-LABEL: "op_tuple"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_tuple(%arg0: tensor<f32>) -> tuple<tensor<f32>> {
--  // CHECK: "vhlo.tuple_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tuple_v1<!vhlo.tensor_v1<!vhlo.f32_v1>>
-+  // CHECK: "vhlo.tuple_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tuple_v1<!vhlo.tensor_v1<!vhlo.f32_v1>>
-   %0 = "stablehlo.tuple"(%arg0) : (tensor<f32>) -> tuple<tensor<f32>>
-   func.return %0 : tuple<tensor<f32>>
- }
- 
- // CHECK-LABEL: "op_unary_einsum"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_unary_einsum(%arg0: tensor<8x16xf32>) -> tensor<8xf32> {
--  //      CHECK: "vhlo.unary_einsum_v1"(%arg0) <{
-+  //      CHECK: "vhlo.unary_einsum_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   einsum_config = #vhlo.string_v1<"ab->a">
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<8x16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<8x!vhlo.f32_v1>
-   %0 = "stablehlo.unary_einsum"(%arg0) {
-@@ -2164,22 +2323,25 @@
- }
- 
- // CHECK-LABEL: "op_uniform_dequantize"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_uniform_dequantize(%arg0: tensor<!quant.uniform<i8:f32, 34.0:16>>) -> tensor<f32> {
--  // CHECK: "vhlo.uniform_dequantize_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.uniform_dequantize_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.uniform_dequantize"(%arg0) : (tensor<!quant.uniform<i8:f32, 34.0:16>>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_uniform_quantize"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_uniform_quantize(%arg0: tensor<f32>) -> tensor<!quant.uniform<i8:f32, 34.0:16>> {
--  // CHECK: "vhlo.uniform_quantize_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>
-+  // CHECK: "vhlo.uniform_quantize_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>
-   %0 = "stablehlo.uniform_quantize"(%arg0) : (tensor<f32>) -> tensor<!quant.uniform<i8:f32, 34.0:16>>
-   func.return %0 : tensor<!quant.uniform<i8:f32, 34.0:16>>
- }
- 
- // CHECK-LABEL: "op_while"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_while(%arg0: tensor<i1>) -> tensor<i1> {
--  //      CHECK: "vhlo.while_v1"(%arg0) ({
-+  //      CHECK: "vhlo.while_v1"(%[[ARG0]]) ({
-   // CHECK-NEXT:   ^[[BB:bb.*]](%[[ARG1:arg.*]]: !vhlo.tensor_v1<!vhlo.bool_v1>):
-   // CHECK-NEXT:     "vhlo.return_v1"(%[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>) -> ()
-   // CHECK-NEXT:   }, {
-@@ -2197,8 +2359,9 @@
- }
- 
- // CHECK-LABEL: "op_xor"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_xor(%arg0: tensor<i1>, %arg1: tensor<i1>) -> tensor<i1> {
--  // CHECK: "vhlo.xor_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-+  // CHECK: "vhlo.xor_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-   %0 = "stablehlo.xor"(%arg0, %arg1) : (tensor<i1>, tensor<i1>) -> tensor<i1>
-   func.return %0 : tensor<i1>
- }
-@@ -2206,197 +2369,225 @@
- // ============ TYPES ============
- 
- // CHECK-LABEL: "type_i1"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_i1(%arg0: tensor<i1>, %arg1: tensor<i1>) -> tensor<i1> {
--  // CHECK: "vhlo.and_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-+  // CHECK: "vhlo.and_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-   %0 = "stablehlo.and"(%arg0, %arg1) : (tensor<i1>, tensor<i1>) -> tensor<i1>
-   func.return %0 : tensor<i1>
- }
- 
- // CHECK-LABEL: "type_i4"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_i4(%arg0: tensor<i4>, %arg1: tensor<i4>) -> tensor<i4> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.i4_v1>, !vhlo.tensor_v1<!vhlo.i4_v1>) -> !vhlo.tensor_v1<!vhlo.i4_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i4_v1>, !vhlo.tensor_v1<!vhlo.i4_v1>) -> !vhlo.tensor_v1<!vhlo.i4_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<i4>, tensor<i4>) -> tensor<i4>
-   func.return %0 : tensor<i4>
- }
- 
- // CHECK-LABEL: "type_i8"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_i8(%arg0: tensor<i8>, %arg1: tensor<i8>) -> tensor<i8> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.i8_v1>, !vhlo.tensor_v1<!vhlo.i8_v1>) -> !vhlo.tensor_v1<!vhlo.i8_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i8_v1>, !vhlo.tensor_v1<!vhlo.i8_v1>) -> !vhlo.tensor_v1<!vhlo.i8_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<i8>, tensor<i8>) -> tensor<i8>
-   func.return %0 : tensor<i8>
- }
- 
- // CHECK-LABEL: "type_i16"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_i16(%arg0: tensor<i16>, %arg1: tensor<i16>) -> tensor<i16> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.i16_v1>, !vhlo.tensor_v1<!vhlo.i16_v1>) -> !vhlo.tensor_v1<!vhlo.i16_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i16_v1>, !vhlo.tensor_v1<!vhlo.i16_v1>) -> !vhlo.tensor_v1<!vhlo.i16_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<i16>, tensor<i16>) -> tensor<i16>
-   func.return %0 : tensor<i16>
- }
- 
- // CHECK-LABEL: "type_i32"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_i32(%arg0: tensor<i32>, %arg1: tensor<i32>) -> tensor<i32> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.i32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<i32>, tensor<i32>) -> tensor<i32>
-   func.return %0 : tensor<i32>
- }
- 
- // CHECK-LABEL: "type_i64"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_i64(%arg0: tensor<i64>, %arg1: tensor<i64>) -> tensor<i64> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.i64_v1>, !vhlo.tensor_v1<!vhlo.i64_v1>) -> !vhlo.tensor_v1<!vhlo.i64_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i64_v1>, !vhlo.tensor_v1<!vhlo.i64_v1>) -> !vhlo.tensor_v1<!vhlo.i64_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<i64>, tensor<i64>) -> tensor<i64>
-   func.return %0 : tensor<i64>
- }
- 
- // CHECK-LABEL: "type_ui4"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_ui4(%arg0: tensor<ui4>, %arg1: tensor<ui4>) -> tensor<ui4> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.ui4_v1>, !vhlo.tensor_v1<!vhlo.ui4_v1>) -> !vhlo.tensor_v1<!vhlo.ui4_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.ui4_v1>, !vhlo.tensor_v1<!vhlo.ui4_v1>) -> !vhlo.tensor_v1<!vhlo.ui4_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<ui4>, tensor<ui4>) -> tensor<ui4>
-   func.return %0 : tensor<ui4>
- }
- 
- // CHECK-LABEL: "type_ui8"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_ui8(%arg0: tensor<ui8>, %arg1: tensor<ui8>) -> tensor<ui8> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.ui8_v1>, !vhlo.tensor_v1<!vhlo.ui8_v1>) -> !vhlo.tensor_v1<!vhlo.ui8_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.ui8_v1>, !vhlo.tensor_v1<!vhlo.ui8_v1>) -> !vhlo.tensor_v1<!vhlo.ui8_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<ui8>, tensor<ui8>) -> tensor<ui8>
-   func.return %0 : tensor<ui8>
- }
- 
- // CHECK-LABEL: "type_ui16"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_ui16(%arg0: tensor<ui16>, %arg1: tensor<ui16>) -> tensor<ui16> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.ui16_v1>, !vhlo.tensor_v1<!vhlo.ui16_v1>) -> !vhlo.tensor_v1<!vhlo.ui16_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.ui16_v1>, !vhlo.tensor_v1<!vhlo.ui16_v1>) -> !vhlo.tensor_v1<!vhlo.ui16_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<ui16>, tensor<ui16>) -> tensor<ui16>
-   func.return %0 : tensor<ui16>
- }
- 
- // CHECK-LABEL: "type_ui32"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_ui32(%arg0: tensor<ui32>, %arg1: tensor<ui32>) -> tensor<ui32> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.ui32_v1>, !vhlo.tensor_v1<!vhlo.ui32_v1>) -> !vhlo.tensor_v1<!vhlo.ui32_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.ui32_v1>, !vhlo.tensor_v1<!vhlo.ui32_v1>) -> !vhlo.tensor_v1<!vhlo.ui32_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<ui32>, tensor<ui32>) -> tensor<ui32>
-   func.return %0 : tensor<ui32>
- }
- 
- // CHECK-LABEL: "type_ui64"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_ui64(%arg0: tensor<ui64>, %arg1: tensor<ui64>) -> tensor<ui64> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.ui64_v1>, !vhlo.tensor_v1<!vhlo.ui64_v1>) -> !vhlo.tensor_v1<!vhlo.ui64_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.ui64_v1>, !vhlo.tensor_v1<!vhlo.ui64_v1>) -> !vhlo.tensor_v1<!vhlo.ui64_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<ui64>, tensor<ui64>) -> tensor<ui64>
-   func.return %0 : tensor<ui64>
- }
- 
- // CHECK-LABEL: "type_f8E4M3FN"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_f8E4M3FN(%arg0: tensor<f8E4M3FN>, %arg1: tensor<f8E4M3FN>) -> tensor<f8E4M3FN> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f8E4M3FN_v1>, !vhlo.tensor_v1<!vhlo.f8E4M3FN_v1>) -> !vhlo.tensor_v1<!vhlo.f8E4M3FN_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f8E4M3FN_v1>, !vhlo.tensor_v1<!vhlo.f8E4M3FN_v1>) -> !vhlo.tensor_v1<!vhlo.f8E4M3FN_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f8E4M3FN>, tensor<f8E4M3FN>) -> tensor<f8E4M3FN>
-   func.return %0 : tensor<f8E4M3FN>
- }
- 
- // CHECK-LABEL: "type_f8E5M2"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_f8E5M2(%arg0: tensor<f8E5M2>, %arg1: tensor<f8E5M2>) -> tensor<f8E5M2> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f8E5M2_v1>, !vhlo.tensor_v1<!vhlo.f8E5M2_v1>) -> !vhlo.tensor_v1<!vhlo.f8E5M2_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f8E5M2_v1>, !vhlo.tensor_v1<!vhlo.f8E5M2_v1>) -> !vhlo.tensor_v1<!vhlo.f8E5M2_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f8E5M2>, tensor<f8E5M2>) -> tensor<f8E5M2>
-   func.return %0 : tensor<f8E5M2>
- }
- 
- // CHECK-LABEL: "type_f8E4M3FNUZ"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_f8E4M3FNUZ(%arg0: tensor<f8E4M3FNUZ>, %arg1: tensor<f8E4M3FNUZ>) -> tensor<f8E4M3FNUZ> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f8E4M3FNUZ_v1>, !vhlo.tensor_v1<!vhlo.f8E4M3FNUZ_v1>) -> !vhlo.tensor_v1<!vhlo.f8E4M3FNUZ_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f8E4M3FNUZ_v1>, !vhlo.tensor_v1<!vhlo.f8E4M3FNUZ_v1>) -> !vhlo.tensor_v1<!vhlo.f8E4M3FNUZ_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f8E4M3FNUZ>, tensor<f8E4M3FNUZ>) -> tensor<f8E4M3FNUZ>
-   func.return %0 : tensor<f8E4M3FNUZ>
- }
- 
- // CHECK-LABEL: "type_f8E4M3B11FNUZ"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_f8E4M3B11FNUZ(%arg0: tensor<f8E4M3B11FNUZ>, %arg1: tensor<f8E4M3B11FNUZ>) -> tensor<f8E4M3B11FNUZ> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f8E4M3B11FNUZ_v1>, !vhlo.tensor_v1<!vhlo.f8E4M3B11FNUZ_v1>) -> !vhlo.tensor_v1<!vhlo.f8E4M3B11FNUZ_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f8E4M3B11FNUZ_v1>, !vhlo.tensor_v1<!vhlo.f8E4M3B11FNUZ_v1>) -> !vhlo.tensor_v1<!vhlo.f8E4M3B11FNUZ_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f8E4M3B11FNUZ>, tensor<f8E4M3B11FNUZ>) -> tensor<f8E4M3B11FNUZ>
-   func.return %0 : tensor<f8E4M3B11FNUZ>
- }
- 
- // CHECK-LABEL: "type_f8E5M2FNUZ"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_f8E5M2FNUZ(%arg0: tensor<f8E5M2FNUZ>, %arg1: tensor<f8E5M2FNUZ>) -> tensor<f8E5M2FNUZ> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f8E5M2FNUZ_v1>, !vhlo.tensor_v1<!vhlo.f8E5M2FNUZ_v1>) -> !vhlo.tensor_v1<!vhlo.f8E5M2FNUZ_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f8E5M2FNUZ_v1>, !vhlo.tensor_v1<!vhlo.f8E5M2FNUZ_v1>) -> !vhlo.tensor_v1<!vhlo.f8E5M2FNUZ_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f8E5M2FNUZ>, tensor<f8E5M2FNUZ>) -> tensor<f8E5M2FNUZ>
-   func.return %0 : tensor<f8E5M2FNUZ>
- }
- 
- // CHECK-LABEL: "type_bf16"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_bf16(%arg0: tensor<bf16>, %arg1: tensor<bf16>) -> tensor<bf16> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.bf16_v1>, !vhlo.tensor_v1<!vhlo.bf16_v1>) -> !vhlo.tensor_v1<!vhlo.bf16_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.bf16_v1>, !vhlo.tensor_v1<!vhlo.bf16_v1>) -> !vhlo.tensor_v1<!vhlo.bf16_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<bf16>, tensor<bf16>) -> tensor<bf16>
-   func.return %0 : tensor<bf16>
- }
- 
- // CHECK-LABEL: "type_f16"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_f16(%arg0: tensor<f16>, %arg1: tensor<f16>) -> tensor<f16> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f16_v1>, !vhlo.tensor_v1<!vhlo.f16_v1>) -> !vhlo.tensor_v1<!vhlo.f16_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f16_v1>, !vhlo.tensor_v1<!vhlo.f16_v1>) -> !vhlo.tensor_v1<!vhlo.f16_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f16>, tensor<f16>) -> tensor<f16>
-   func.return %0 : tensor<f16>
- }
- 
- // CHECK-LABEL: "type_f32"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_f32(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "type_f64"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_f64(%arg0: tensor<f64>, %arg1: tensor<f64>) -> tensor<f64> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f64_v1>, !vhlo.tensor_v1<!vhlo.f64_v1>) -> !vhlo.tensor_v1<!vhlo.f64_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f64_v1>, !vhlo.tensor_v1<!vhlo.f64_v1>) -> !vhlo.tensor_v1<!vhlo.f64_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f64>, tensor<f64>) -> tensor<f64>
-   func.return %0 : tensor<f64>
- }
- 
- // CHECK-LABEL: "type_complex_f32"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_complex_f32(%arg0: tensor<complex<f32>>, %arg1: tensor<complex<f32>>) -> tensor<complex<f32>> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>, !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>) -> !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>, !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>) -> !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<complex<f32>>, tensor<complex<f32>>) -> tensor<complex<f32>>
-   func.return %0 : tensor<complex<f32>>
- }
- 
- // CHECK-LABEL: "type_complex_f64"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_complex_f64(%arg0: tensor<complex<f64>>, %arg1: tensor<complex<f64>>) -> tensor<complex<f64>> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f64_v1>>, !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f64_v1>>) -> !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f64_v1>>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f64_v1>>, !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f64_v1>>) -> !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f64_v1>>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<complex<f64>>, tensor<complex<f64>>) -> tensor<complex<f64>>
-   func.return %0 : tensor<complex<f64>>
- }
- 
- // CHECK-LABEL: "type_dynamism_ranked"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @type_dynamism_ranked(%arg0: tensor<?xf32>) -> tensor<?xf32> {
--  // CHECK: "vhlo.abs_v1"(%arg0) : (!vhlo.tensor_v1<?x!vhlo.f32_v1>) -> !vhlo.tensor_v1<?x!vhlo.f32_v1>
-+  // CHECK: "vhlo.abs_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<?x!vhlo.f32_v1>) -> !vhlo.tensor_v1<?x!vhlo.f32_v1>
-   %0 = "stablehlo.abs"(%arg0) : (tensor<?xf32>) -> tensor<?xf32>
-   func.return %0 : tensor<?xf32>
- }
- 
- // CHECK-LABEL: "type_per_tensor_quantization"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_per_tensor_quantization(%arg0: tensor<!quant.uniform<i8:f32, 34.0:16>>, %arg1: tensor<!quant.uniform<i8:f32, 34.0:16>>) -> tensor<!quant.uniform<i8:f32, 34.0:16>> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>, !vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>) -> !vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>, !vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>) -> !vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<!quant.uniform<i8:f32, 34.0:16>>, tensor<!quant.uniform<i8:f32, 34.0:16>>) -> tensor<!quant.uniform<i8:f32, 34.0:16>>
-   func.return %0 : tensor<!quant.uniform<i8:f32, 34.0:16>>
- }
- 
- // CHECK-LABEL: "type_per_axis_quantization"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @type_per_axis_quantization(%arg0: tensor<2x!quant.uniform<i8:f32:0, {34.0:16, 34.0:16}>>) -> tensor<2x!quant.uniform<i8:f32:0, {34.0:16, 34.0:16}>> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg0) : (!vhlo.tensor_v1<2x!vhlo.quant_per_axis_v1<!vhlo.i8_v1:!vhlo.f32_v1, 0, 3.400000e+01, 3.400000e+01, 16, 16, -128:127, 1>>, !vhlo.tensor_v1<2x!vhlo.quant_per_axis_v1<!vhlo.i8_v1:!vhlo.f32_v1, 0, 3.400000e+01, 3.400000e+01, 16, 16, -128:127, 1>>) -> !vhlo.tensor_v1<2x!vhlo.quant_per_axis_v1<!vhlo.i8_v1:!vhlo.f32_v1, 0, 3.400000e+01, 3.400000e+01, 16, 16, -128:127, 1>>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG0]]) : (!vhlo.tensor_v1<2x!vhlo.quant_per_axis_v1<!vhlo.i8_v1:!vhlo.f32_v1, 0, 3.400000e+01, 3.400000e+01, 16, 16, -128:127, 1>>, !vhlo.tensor_v1<2x!vhlo.quant_per_axis_v1<!vhlo.i8_v1:!vhlo.f32_v1, 0, 3.400000e+01, 3.400000e+01, 16, 16, -128:127, 1>>) -> !vhlo.tensor_v1<2x!vhlo.quant_per_axis_v1<!vhlo.i8_v1:!vhlo.f32_v1, 0, 3.400000e+01, 3.400000e+01, 16, 16, -128:127, 1>>
-   %0 = stablehlo.add %arg0, %arg0 : tensor<2x!quant.uniform<i8:f32:0, {34.0:16, 34.0:16}>>
-   func.return %0 : tensor<2x!quant.uniform<i8:f32:0, {34.0:16, 34.0:16}>>
- }
- 
- //       CHECK: function_type = #vhlo.type_v1<!vhlo.func_v1<(!vhlo.token_v1) -> !vhlo.token_v1>>
- // CHECK-LABEL: "type_token_callee"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @type_token_callee(%arg0: !stablehlo.token) -> !stablehlo.token {
--  // CHECK: "vhlo.return_v1"(%arg0) : (!vhlo.token_v1) -> ()
-+  // CHECK: "vhlo.return_v1"(%[[ARG0]]) : (!vhlo.token_v1) -> ()
-   return %arg0 : !stablehlo.token
- }
- 
- //       CHECK: function_type = #vhlo.type_v1<!vhlo.func_v1<(!vhlo.token_v1) -> !vhlo.token_v1>>
- // CHECK-LABEL: "type_token_caller"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @type_token_caller(%arg0: !stablehlo.token) -> !stablehlo.token {
--  // CHECK:      "vhlo.call_v1"(%arg0) <{callee = #vhlo.string_v1<"type_token_callee">}
-+  // CHECK:      "vhlo.call_v1"(%[[ARG0]]) <{callee = #vhlo.string_v1<"type_token_callee">}
-   // CHECK-SAME: (!vhlo.token_v1) -> !vhlo.token_v1
-   %0 = func.call @type_token_callee(%arg0) : (!stablehlo.token) -> !stablehlo.token
-   return %0 : !stablehlo.token
- }
- 
- // CHECK-LABEL: "type_tuple"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @type_tuple(%arg0: tuple<tensor<f32>>) -> tuple<!stablehlo.token> {
-   %0 = "stablehlo.custom_call"(%arg0) {
-     call_target_name = "foo"
-diff --ruN a/stablehlo/stablehlo/tests/vhlo/stablehlo_legalize_to_vhlo.0_19_0.mlir b/stablehlo/stablehlo/tests/vhlo/stablehlo_legalize_to_vhlo.0_19_0.mlir
---- stablehlo/stablehlo/tests/vhlo/stablehlo_legalize_to_vhlo.0_19_0.mlir
-+++ stablehlo/stablehlo/tests/vhlo/stablehlo_legalize_to_vhlo.0_19_0.mlir
-@@ -13,6 +13,7 @@
- // ============ ATTRIBUTES ============
- 
- // CHECK-LABEL: "attr_comparison_direction_eq"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_direction_eq(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     // CHECK: comparison_direction = #vhlo<comparison_direction_v1 EQ>
-@@ -22,6 +23,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_direction_ne"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_direction_ne(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     // CHECK: comparison_direction = #vhlo<comparison_direction_v1 NE>
-@@ -31,6 +33,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_direction_ge"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_direction_ge(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     // CHECK: comparison_direction = #vhlo<comparison_direction_v1 GE>
-@@ -40,6 +43,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_direction_gt"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_direction_gt(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     // CHECK: comparison_direction = #vhlo<comparison_direction_v1 GT>
-@@ -49,6 +53,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_direction_le"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_direction_le(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     // CHECK: comparison_direction = #vhlo<comparison_direction_v1 LE>
-@@ -58,6 +63,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_direction_lt"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_direction_lt(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     // CHECK: comparison_direction = #vhlo<comparison_direction_v1 LT>
-@@ -67,6 +73,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_type_notype"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_type_notype(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     comparison_direction = #stablehlo<comparison_direction EQ>
-@@ -76,6 +83,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_type_float"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_type_float(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     comparison_direction = #stablehlo<comparison_direction EQ>,
-@@ -86,6 +94,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_type_totalorder"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_type_totalorder(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     comparison_direction = #stablehlo<comparison_direction EQ>,
-@@ -96,6 +105,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_type_signed"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_type_signed(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     comparison_direction = #stablehlo<comparison_direction EQ>,
-@@ -106,6 +116,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_type_unsigned"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_type_unsigned(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     comparison_direction = #stablehlo<comparison_direction EQ>,
-@@ -118,6 +129,7 @@
- // ConvDimensionNumbers aka #stablehlo.conv is covered below.
- 
- // CHECK-LABEL: "attr_custom_call_api_version_unspecified"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_custom_call_api_version_unspecified(%arg0: tensor<f32>) -> tensor<f32> {
-   %0 = "stablehlo.custom_call"(%arg0) {
-     call_target_name = "foo",
-@@ -128,6 +140,7 @@
- }
- 
- // CHECK-LABEL: "attr_custom_call_api_version_original"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_custom_call_api_version_original(%arg0: tensor<f32>) -> tensor<f32> {
-   %0 = "stablehlo.custom_call"(%arg0) {
-     call_target_name = "foo",
-@@ -138,6 +151,7 @@
- }
- 
- // CHECK-LABEL: "attr_custom_call_api_version_status_returning"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_custom_call_api_version_status_returning(%arg0: tensor<f32>) -> tensor<f32> {
-   %0 = "stablehlo.custom_call"(%arg0) {
-     call_target_name = "foo",
-@@ -148,6 +162,7 @@
- }
- 
- // CHECK-LABEL: "attr_custom_call_api_version_status_returning_unified"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_custom_call_api_version_status_returning_unified(%arg0: tensor<f32>) -> tensor<f32> {
-   %0 = "stablehlo.custom_call"(%arg0) {
-     call_target_name = "foo",
-@@ -166,6 +181,7 @@
- // DotDimensionNumbers aka #stablehlo.dot is covered below.
- 
- // CHECK-LABEL: "attr_fft_type_fft"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_fft_type_fft(%arg0: tensor<16xcomplex<f32>>) -> tensor<16xcomplex<f32>> {
-   %0 = "stablehlo.fft"(%arg0) {
-     // CHECK: fft_type = #vhlo<fft_type_v1 FFT>
-@@ -176,6 +192,7 @@
- }
- 
- // CHECK-LABEL: "attr_fft_type_ifft"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_fft_type_ifft(%arg0: tensor<16xcomplex<f32>>) -> tensor<16xcomplex<f32>> {
-   %0 = "stablehlo.fft"(%arg0) {
-     // CHECK: fft_type = #vhlo<fft_type_v1 IFFT>
-@@ -186,6 +203,7 @@
- }
- 
- // CHECK-LABEL: "attr_fft_type_rfft"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_fft_type_rfft(%arg0: tensor<16xf32>) -> tensor<9xcomplex<f32>> {
-   %0 = "stablehlo.fft"(%arg0) {
-     // CHECK: fft_type = #vhlo<fft_type_v1 RFFT>
-@@ -196,6 +214,7 @@
- }
- 
- // CHECK-LABEL: "attr_fft_type_irfft"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_fft_type_irfft(%arg0: tensor<9xcomplex<f32>>) -> tensor<16xf32> {
-   %0 = "stablehlo.fft"(%arg0) {
-     // CHECK: fft_type = #vhlo<fft_type_v1 IRFFT>
-@@ -208,6 +227,7 @@
- // GatherDimensionNumbers aka #stablehlo.gather is covered below.
- 
- // CHECK-LABEL: "attr_precision_config_default"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_precision_config_default(%arg0: tensor<8x16xf32>, %arg1: tensor<16x8xf32>) -> tensor<8x8xf32> {
-   %0 = "stablehlo.dot"(%arg0, %arg1) {
-     // CHECK: precision_config = #vhlo.array_v1<[#vhlo<precision_v1 DEFAULT>, #vhlo<precision_v1 DEFAULT>]>
-@@ -216,6 +236,7 @@
- }
- 
- // CHECK-LABEL: "attr_precision_config_high"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_precision_config_high(%arg0: tensor<8x16xf32>, %arg1: tensor<16x8xf32>) -> tensor<8x8xf32> {
-   %0 = "stablehlo.dot"(%arg0, %arg1) {
-     // CHECK: precision_config = #vhlo.array_v1<[#vhlo<precision_v1 HIGH>, #vhlo<precision_v1 HIGH>]>
-@@ -225,6 +246,7 @@
- }
- 
- // CHECK-LABEL: "attr_precision_config_highest"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_precision_config_highest(%arg0: tensor<8x16xf32>, %arg1: tensor<16x8xf32>) -> tensor<8x8xf32> {
-   %0 = "stablehlo.dot"(%arg0, %arg1) {
-     // CHECK: precision_config = #vhlo.array_v1<[#vhlo<precision_v1 HIGHEST>, #vhlo<precision_v1 HIGHEST>]>
-@@ -234,6 +256,7 @@
- }
- 
- // CHECK-LABEL: "attr_rng_algorithm_default"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_rng_algorithm_default(%arg0: tensor<f32>) -> (tensor<f32>, tensor<f32>) {
-   %0:2 = "stablehlo.rng_bit_generator"(%arg0) {
-     // CHECK: rng_algorithm = #vhlo<rng_algorithm_v1 DEFAULT>
-@@ -243,6 +266,7 @@
- }
- 
- // CHECK-LABEL: "attr_rng_algorithm_three_fry"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_rng_algorithm_three_fry(%arg0: tensor<f32>) -> (tensor<f32>, tensor<f32>) {
-   %0:2 = "stablehlo.rng_bit_generator"(%arg0) {
-     // CHECK: rng_algorithm = #vhlo<rng_algorithm_v1 THREE_FRY>
-@@ -252,6 +276,7 @@
- }
- 
- // CHECK-LABEL: "attr_rng_algorithm_philox"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_rng_algorithm_philox(%arg0: tensor<f32>) -> (tensor<f32>, tensor<f32>) {
-   %0:2 = "stablehlo.rng_bit_generator"(%arg0) {
-     // CHECK: rng_algorithm = #vhlo<rng_algorithm_v1 PHILOX>
-@@ -261,6 +286,7 @@
- }
- 
- // CHECK-LABEL: "attr_rng_distribution_uniform"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @attr_rng_distribution_uniform(%arg0: tensor<f32>, %arg1: tensor<f32>, %arg2: tensor<0xindex>) -> tensor<f32> {
-   %0 = "stablehlo.rng"(%arg0, %arg1, %arg2) {
-     // CHECK: rng_distribution = #vhlo<rng_distribution_v1 UNIFORM>
-@@ -270,6 +296,7 @@
- }
- 
- // CHECK-LABEL: "attr_rng_distribution_normal"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @attr_rng_distribution_normal(%arg0: tensor<f32>, %arg1: tensor<f32>, %arg2: tensor<0xindex>) -> tensor<f32> {
-   %0 = "stablehlo.rng"(%arg0, %arg1, %arg2) {
-     // CHECK: rng_distribution = #vhlo<rng_distribution_v1 NORMAL>
-@@ -281,6 +308,7 @@
- // ScatterDimensionNumbers aka #stablehlo.scatter is covered below.
- 
- // CHECK-LABEL: "attr_transpose_no_transpose"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_transpose_no_transpose(%arg0: tensor<16x16xf32>, %arg1: tensor<16x16xf32>) ->  tensor<16x16xf32> {
-   %0 = "stablehlo.triangular_solve"(%arg0, %arg1) {
-     left_side = true,
-@@ -293,6 +321,7 @@
- }
- 
- // CHECK-LABEL: "attr_transpose_transpose"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_transpose_transpose(%arg0: tensor<16x16xf32>, %arg1: tensor<16x16xf32>) ->  tensor<16x16xf32> {
-   %0 = "stablehlo.triangular_solve"(%arg0, %arg1) {
-     left_side = true,
-@@ -305,6 +334,7 @@
- }
- 
- // CHECK-LABEL: "attr_transpose_adjoint"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_transpose_adjoint(%arg0: tensor<16x16xf32>, %arg1: tensor<16x16xf32>) ->  tensor<16x16xf32> {
-   %0 = "stablehlo.triangular_solve"(%arg0, %arg1) {
-     left_side = true,
-@@ -319,10 +349,9 @@
- // TypeExtensionsAttr aka #stablehlo.type_extensions is covered below.
- 
- // CHECK-LABEL: "attr_type_extensions_bounds"
--func.func @attr_type_extensions_bounds(
--    %arg0: tensor<?x?xf32, #stablehlo.type_extensions<bounds = [16, ?]>>)
--    -> tensor<?x?xf32, #stablehlo.type_extensions<bounds = [16, ?]>> {
--  // CHECK: "vhlo.return_v1"(%arg0) : (!vhlo.tensor_v1<?x?x!vhlo.f32_v1, #vhlo.type_extensions_v1<bounds = [16, ?]>>) -> ()
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @attr_type_extensions_bounds(%arg0: tensor<?x?xf32, #stablehlo.type_extensions<bounds = [16, ?]>>) -> tensor<?x?xf32, #stablehlo.type_extensions<bounds = [16, ?]>> {
-+  // CHECK: "vhlo.return_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<?x?x!vhlo.f32_v1, #vhlo.type_extensions_v1<bounds = [16, ?]>>) -> ()
-   func.return %arg0 : tensor<?x?xf32, #stablehlo.type_extensions<bounds = [16, ?]>>
- }
- 
-@@ -330,8 +359,9 @@
- // ============ DEFAULTS ============
- 
- // CHECK-LABEL: "default_all_gather"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_all_gather(%arg0: tensor<16x8xf32>) -> tensor<16x16xf32> {
--  //               CHECK: "vhlo.all_gather_v1"(%arg0) <{
-+  //               CHECK: "vhlo.all_gather_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   all_gather_dim = #vhlo.integer_v1<1 : i64>
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0], [1]]> : tensor<2x1xi64>>,
-@@ -345,8 +375,9 @@
- }
- 
- // CHECK-LABEL: "default_all_reduce"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_all_reduce(%arg0: tensor<f32>) -> tensor<f32> {
--  //               CHECK: "vhlo.all_reduce_v1"(%arg0)
-+  //               CHECK: "vhlo.all_reduce_v1"(%[[ARG0]])
-   //          CHECK-SAME: <{
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0], [1]]> : tensor<2x1xi64>>,
-@@ -368,8 +399,9 @@
- }
- 
- // CHECK-LABEL: "default_all_to_all"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_all_to_all(%arg0: tensor<4x16xf32>) -> tensor<16x4xf32> {
--  //               CHECK: "vhlo.all_to_all_v1"(%arg0) <{
-+  //               CHECK: "vhlo.all_to_all_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   //          CHECK-SAME:   concat_dimension = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0, 1, 2, 3]]> : tensor<1x4xi64>>,
-@@ -386,8 +418,9 @@
- }
- 
- // CHECK-LABEL: "default_cholesky"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_cholesky(%arg0: tensor<1x16x16xf32>) -> tensor<1x16x16xf32> {
--  //      CHECK: "vhlo.cholesky_v1"(%arg0) <{
-+  //      CHECK: "vhlo.cholesky_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   lower = #vhlo.bool_v1<false>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<1x16x16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<1x16x16x!vhlo.f32_v1>
-   %0 = "stablehlo.cholesky"(%arg0) : (tensor<1x16x16xf32>) -> tensor<1x16x16xf32>
-@@ -395,8 +428,9 @@
- }
- 
- // CHECK-LABEL: "default_collective_permute"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_collective_permute(%arg0: tensor<16x8xf32>) -> tensor<16x8xf32> {
--  //               CHECK: "vhlo.collective_permute_v1"(%arg0) <{
-+  //               CHECK: "vhlo.collective_permute_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME{LITERAL}:   source_target_pairs = #vhlo.tensor_v1<dense<[[0, 1], [1, 2], [2, 3]]> : tensor<3x2xi64>>
-   //          CHECK-SAME: }> : (!vhlo.tensor_v1<16x8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x8x!vhlo.f32_v1>
-@@ -407,8 +441,9 @@
- }
- 
- // CHECK-LABEL: "default_collective_broadcast"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_collective_broadcast(%arg0: tensor<16x8xf32>) -> tensor<16x8xf32> {
--  //               CHECK: "vhlo.collective_broadcast_v1"(%arg0) <{
-+  //               CHECK: "vhlo.collective_broadcast_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0, 1]]> : tensor<1x2xi64>>
-   //          CHECK-SAME: }> : (!vhlo.tensor_v1<16x8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x8x!vhlo.f32_v1>
-@@ -419,8 +454,9 @@
- }
- 
- // CHECK-LABEL: "default_compare"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @default_compare(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
--  //      CHECK: "vhlo.compare_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.compare_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   compare_type = #vhlo<comparison_type_v1 NOTYPE>,
-   // CHECK-SAME:   comparison_direction = #vhlo<comparison_direction_v1 EQ>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-@@ -431,8 +467,9 @@
- }
- 
- // CHECK-LABEL: "default_composite"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_composite(%arg0: tensor<f32>) -> tensor<f32> {
--  //               CHECK: "vhlo.composite_v1"(%arg0) <{
-+  //               CHECK: "vhlo.composite_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   composite_attributes = #vhlo.dict_v1<{}>
-   //          CHECK-SAME:   decomposition = #vhlo.string_v1<"composite_target">
-   //          CHECK-SAME:   name = #vhlo.string_v1<"stablehlo.composite_target">
-@@ -446,8 +483,9 @@
- }
- 
- // CHECK-LABEL: "default_convolution"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @default_convolution(%arg0: tensor<1x8x8x207xf32>, %arg1: tensor<3x3x207x16xf32>) -> tensor<1x6x6x16xf32> {
--  //      CHECK: "vhlo.convolution_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.convolution_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   batch_group_count = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   feature_group_count = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   input_batch_dimension = #vhlo.integer_v1<0 : i64>,
-@@ -475,8 +513,9 @@
- }
- 
- // CHECK-LABEL: "default_custom_call"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_custom_call(%arg0: tensor<f32>) -> tensor<f32> {
--  //      CHECK: "vhlo.custom_call_v1"(%arg0) <{
-+  //      CHECK: "vhlo.custom_call_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   api_version = #vhlo<api_version_v1 API_VERSION_ORIGINAL>,
-   // CHECK-SAME:   backend_config = #vhlo.string_v1<"">,
-   // CHECK-SAME:   call_target_name = #vhlo.string_v1<"foo">,
-@@ -493,8 +532,9 @@
- }
- 
- // CHECK-LABEL: "default_dot_general"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @default_dot_general(%arg0: tensor<8x8x16xf32>, %arg1: tensor<8x16x8xf32>) -> tensor<8x8x8xf32> {
--  //      CHECK: "vhlo.dot_general_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.dot_general_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   lhs_batching_dimensions = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>,
-   // CHECK-SAME:   lhs_contracting_dimensions = #vhlo.tensor_v1<dense<2> : tensor<1xi64>>,
-   // CHECK-SAME:   precision_config = #vhlo.array_v1<[#vhlo<precision_v1 DEFAULT>, #vhlo<precision_v1 DEFAULT>]>,
-@@ -513,8 +553,9 @@
- }
- 
- // CHECK-LABEL: "default_dot"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @default_dot(%arg0: tensor<8x16xf32>, %arg1: tensor<16x8xf32>) -> tensor<8x8xf32> {
--  //      CHECK: "vhlo.dot_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.dot_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   precision_config = #vhlo.array_v1<[#vhlo<precision_v1 DEFAULT>, #vhlo<precision_v1 DEFAULT>]>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<8x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<8x8x!vhlo.f32_v1>
-   %0 = "stablehlo.dot"(%arg0, %arg1) : (tensor<8x16xf32>, tensor<16x8xf32>) -> tensor<8x8xf32>
-@@ -522,8 +563,9 @@
- }
- 
- // CHECK-LABEL: "default_dynamic_broadcast_in_dim"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @default_dynamic_broadcast_in_dim(%arg0: tensor<?x?xf32>, %arg1: tensor<2xindex>) -> tensor<?x?xf32> {
--  //      CHECK: "vhlo.dynamic_broadcast_in_dim_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.dynamic_broadcast_in_dim_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   broadcast_dimensions = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-   // CHECK-SAME:   known_expanding_dimensions = #vhlo.tensor_v1<dense<> : tensor<0xi64>>,
-   // CHECK-SAME:   known_nonexpanding_dimensions = #vhlo.tensor_v1<dense<> : tensor<0xi64>>
-@@ -535,8 +577,9 @@
- }
- 
- // CHECK-LABEL: "default_dynamic_conv"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @default_dynamic_conv(%arg0: tensor<1x8x8x207xf32>, %arg1: tensor<3x3x207x16xf32>, %arg2: tensor<4xi32>) -> tensor<1x?x?x16xf32> {
--  //      CHECK: "vhlo.dynamic_conv_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.dynamic_conv_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   batch_group_count = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   feature_group_count = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   input_batch_dimension = #vhlo.integer_v1<0 : i64>,
-@@ -564,8 +607,9 @@
- }
- 
- // CHECK-LABEL: "default_dynamic_gather"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @default_dynamic_gather(%arg0 : tensor<2x4x9xf32>, %arg1 : tensor<1x5x2xi32>, %arg2 : tensor<3xi32>) -> tensor<1x5x8xf32> {
--  //      CHECK: "vhlo.dynamic_gather_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.dynamic_gather_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   collapsed_slice_dims = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-   // CHECK-SAME:   index_vector_dim = #vhlo.integer_v1<2 : i64>,
-   // CHECK-SAME:   indices_are_sorted = #vhlo.bool_v1<false>,
-@@ -591,15 +635,16 @@
-   // CHECK-SAME:   sym_name = #vhlo.string_v1<"default_func">,
-   // CHECK-SAME:   sym_visibility = #vhlo.string_v1<"">
-   // CHECK-SAME: }> ({
--  // CHECK-NEXT:   ^[[BB:bb.*]](%arg0: !vhlo.tensor_v1<!vhlo.f32_v1>):
--  // CHECK-NEXT:     "vhlo.return_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-+  // CHECK-NEXT:   ^[[BB:bb.*]](%[[ARG0:.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
-+  // CHECK-NEXT:     "vhlo.return_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-   // CHECK-NEXT: }) : () -> ()
-   func.return %arg0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "dynamic_gather"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @dynamic_gather(%arg0 : tensor<2x4x9xf32>, %arg1 : tensor<1x5x2xi32>) -> tensor<1x5x1xf32> {
--  //      CHECK: "vhlo.gather_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.gather_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   collapsed_slice_dims = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-   // CHECK-SAME:   index_vector_dim = #vhlo.integer_v1<2 : i64>,
-   // CHECK-SAME:   indices_are_sorted = #vhlo.bool_v1<false>,
-@@ -620,8 +665,9 @@
- }
- 
- // CHECK-LABEL: "default_infeed"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_infeed(%arg0: !stablehlo.token) -> (tensor<f32>, !stablehlo.token) {
--  //               CHECK: "vhlo.infeed_v1"(%arg0) <{
-+  //               CHECK: "vhlo.infeed_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   infeed_config = #vhlo.string_v1<"">,
-   // CHECK-SAME{LITERAL}:   layout = #vhlo.array_v1<[]>
-   //          CHECK-SAME: }> : (!vhlo.token_v1) -> (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.token_v1)
-@@ -630,8 +676,9 @@
- }
- 
- // CHECK-LABEL: "default_outfeed"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @default_outfeed(%arg0: tensor<f32>, %arg1: !stablehlo.token) -> !stablehlo.token {
--  //      CHECK: "vhlo.outfeed_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.outfeed_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   outfeed_config = #vhlo.string_v1<"">
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.token_v1) -> !vhlo.token_v1
-   %0 = "stablehlo.outfeed"(%arg0, %arg1) : (tensor<f32>, !stablehlo.token) -> !stablehlo.token
-@@ -639,8 +686,9 @@
- }
- 
- // CHECK-LABEL: "default_recv"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_recv(%arg0: !stablehlo.token) -> (tensor<f32>, !stablehlo.token) {
--  //      CHECK: "vhlo.recv_v1"(%arg0) <{
-+  //      CHECK: "vhlo.recv_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME:   channel_type = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   is_host_transfer = #vhlo.bool_v1<false>
-@@ -652,8 +700,9 @@
- }
- 
- // CHECK-LABEL: "default_send"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @default_send(%arg0: tensor<f32>, %arg1: !stablehlo.token) -> !stablehlo.token {
--  //      CHECK: "vhlo.send_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.send_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME:   channel_type = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   is_host_transfer = #vhlo.bool_v1<false>
-@@ -665,8 +714,9 @@
- }
- 
- // CHECK-LABEL: "default_reduce_scatter"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_reduce_scatter(%arg0: tensor<16xf32>) -> tensor<16xf32> {
--  //               CHECK: "vhlo.reduce_scatter_v1"(%arg0) <{
-+  //               CHECK: "vhlo.reduce_scatter_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0], [1]]> : tensor<2x1xi64>>,
-   //          CHECK-SAME:   scatter_dimension = #vhlo.integer_v1<0 : i64>
-@@ -688,8 +738,9 @@
- }
- 
- // CHECK-LABEL: "default_reduce_window"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @default_reduce_window(%arg0: tensor<2x17x31x7xf32>, %arg1: tensor<f32>) -> tensor<2x16x30x7xf32> {
--  //               CHECK: "vhlo.reduce_window_v1"(%arg0, %arg1)  <{
-+  //               CHECK: "vhlo.reduce_window_v1"(%[[ARG0]], %[[ARG1]])  <{
-   //          CHECK-SAME:   base_dilations = #vhlo.tensor_v1<dense<1> : tensor<4xi64>>,
-   // CHECK-SAME{LITERAL}:   padding = #vhlo.tensor_v1<dense<0> : tensor<4x2xi64>>,
-   //          CHECK-SAME:   window_dilations = #vhlo.tensor_v1<dense<1> : tensor<4xi64>>,
-@@ -711,8 +762,9 @@
- }
- 
- // CHECK-LABEL: "default_scatter"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @default_scatter(%arg0: tensor<200x100x300xf32>, %arg1: tensor<10x2xi32>, %arg2: tensor<10x300xf32>) -> tensor<200x100x300xf32> {
--  //      CHECK: "vhlo.scatter_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.scatter_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   index_vector_dim = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   indices_are_sorted = #vhlo.bool_v1<false>,
-   // CHECK-SAME:   inserted_window_dims = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-@@ -740,8 +792,9 @@
- }
- 
- // CHECK-LABEL: "default_select_and_scatter"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @default_select_and_scatter(%arg0: tensor<10x24x24x64xf32>, %arg1: tensor<10x23x23x64xf32>, %arg2: tensor<f32>) -> tensor<10x24x24x64xf32> {
--  //      CHECK: "vhlo.select_and_scatter_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.select_and_scatter_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   padding = #vhlo.tensor_v1<dense<0> : tensor<4x2xi64>>,
-   // CHECK-SAME:   window_dimensions = #vhlo.tensor_v1<dense<[1, 2, 2, 1]> : tensor<4xi64>>,
-   // CHECK-SAME:   window_strides = #vhlo.tensor_v1<dense<1> : tensor<4xi64>>
-@@ -769,8 +822,9 @@
- }
- 
- // CHECK-LABEL: "default_sort"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_sort(%arg0: tensor<16xf32>) -> tensor<16xf32> {
--  //      CHECK: "vhlo.sort_v1"(%arg0) <{
-+  //      CHECK: "vhlo.sort_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   dimension = #vhlo.integer_v1<-1 : i64>
-   // CHECK-SAME:   is_stable = #vhlo.bool_v1<false>
-   // CHECK-SAME: }> ({
-@@ -789,29 +843,33 @@
- // ============ OPS ============
- 
- // CHECK-LABEL: "op_abs"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_abs(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.abs_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.abs_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.abs"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_add"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_add(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_after_all"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_after_all(%arg0: !stablehlo.token) -> !stablehlo.token {
--  // CHECK: "vhlo.after_all_v1"(%arg0) : (!vhlo.token_v1) -> !vhlo.token_v1
-+  // CHECK: "vhlo.after_all_v1"(%[[ARG0]]) : (!vhlo.token_v1) -> !vhlo.token_v1
-   %0 = "stablehlo.after_all"(%arg0) : (!stablehlo.token) -> !stablehlo.token
-   func.return %0 : !stablehlo.token
- }
- 
- // CHECK-LABEL: "op_all_gather"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_all_gather(%arg0: tensor<16x8xf32>) -> tensor<16x16xf32> {
--  //               CHECK: "vhlo.all_gather_v1"(%arg0) <{
-+  //               CHECK: "vhlo.all_gather_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   all_gather_dim = #vhlo.integer_v1<1 : i64>
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0], [1]]> : tensor<2x1xi64>>,
-@@ -827,8 +885,9 @@
- }
- 
- // CHECK-LABEL: "op_all_reduce"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_all_reduce(%arg0: tensor<f32>) -> tensor<f32> {
--  //               CHECK: "vhlo.all_reduce_v1"(%arg0) <{
-+  //               CHECK: "vhlo.all_reduce_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0], [1]]> : tensor<2x1xi64>>,
-   //          CHECK-SAME:   use_global_device_ids = #vhlo.bool_v1<true>
-@@ -851,7 +910,7 @@
- 
- // CHECK-LABEL: "op_all_reduce_with_promotable_types"
- func.func @op_all_reduce_with_promotable_types(%operand: tensor<f32>) -> tensor<f64> {
--  //  CHECK: "vhlo.all_reduce_v1"(%arg0)
-+  //  CHECK: "vhlo.all_reduce_v1"(%[[ARG0:.*]])
-   //  CHECK:   ^[[BB:bb.*]](%[[ARG1:arg.*]]: !vhlo.tensor_v1<!vhlo.f64_v1>, %[[ARG2:arg.*]]: !vhlo.tensor_v1<!vhlo.f64_v1>):
-   //  CHECK:     "vhlo.return_v1"(%[[VAL1:.*]]) : (!vhlo.tensor_v1<!vhlo.f64_v1>) -> ()
-   //  CHECK: }) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f64_v1>
-@@ -869,8 +928,9 @@
- }
- 
- // CHECK-LABEL: "op_all_to_all"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_all_to_all(%arg0: tensor<4x16xf32>) -> tensor<16x4xf32> {
--  //               CHECK: "vhlo.all_to_all_v1"(%arg0) <{
-+  //               CHECK: "vhlo.all_to_all_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<1 : i64>,
-   //          CHECK-SAME:   concat_dimension = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0, 1, 2, 3]]> : tensor<1x4xi64>>,
-@@ -888,22 +948,25 @@
- }
- 
- // CHECK-LABEL: "op_and"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_and(%arg0: tensor<i1>, %arg1: tensor<i1>) -> tensor<i1> {
--  // CHECK: "vhlo.and_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-+  // CHECK: "vhlo.and_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-   %0 = "stablehlo.and"(%arg0, %arg1) : (tensor<i1>, tensor<i1>) -> tensor<i1>
-   func.return %0 : tensor<i1>
- }
- 
- // CHECK-LABEL: "op_atan2"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_atan2(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.atan2_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.atan2_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.atan2"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_batch_norm_grad"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}}, %[[ARG3:.*]]: {{.*}}, %[[ARG4:.*]]: {{.*}})
- func.func @op_batch_norm_grad(%arg0: tensor<16x16x16x16xf32>, %arg1: tensor<16xf32>, %arg2: tensor<16xf32>, %arg3: tensor<16xf32>, %arg4: tensor<16x16x16x16xf32>) -> (tensor<16x16x16x16xf32>, tensor<16xf32>, tensor<16xf32>) {
--  //      CHECK: "vhlo.batch_norm_grad_v1"(%arg0, %arg1, %arg2, %arg3, %arg4) <{
-+  //      CHECK: "vhlo.batch_norm_grad_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]], %[[ARG3]], %[[ARG4]]) <{
-   // CHECK-SAME:   epsilon = #vhlo.float_v1<1.000000e-03 : !vhlo.f32_v1>,
-   // CHECK-SAME:   feature_index = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<16x16x16x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x16x16x16x!vhlo.f32_v1>) -> (!vhlo.tensor_v1<16x16x16x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>)
-@@ -915,8 +978,9 @@
- }
- 
- // CHECK-LABEL: "op_batch_norm_inference"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}}, %[[ARG3:.*]]: {{.*}}, %[[ARG4:.*]]: {{.*}})
- func.func @op_batch_norm_inference(%arg0: tensor<16x16x16x16xf32>, %arg1: tensor<16xf32>, %arg2: tensor<16xf32>, %arg3: tensor<16xf32>, %arg4: tensor<16xf32>) -> tensor<16x16x16x16xf32> {
--  //      CHECK: "vhlo.batch_norm_inference_v1"(%arg0, %arg1, %arg2, %arg3, %arg4) <{
-+  //      CHECK: "vhlo.batch_norm_inference_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]], %[[ARG3]], %[[ARG4]]) <{
-   // CHECK-SAME:   epsilon = #vhlo.float_v1<1.000000e-03 : !vhlo.f32_v1>,
-   // CHECK-SAME:   feature_index = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<16x16x16x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x16x16x16x!vhlo.f32_v1>
-@@ -928,8 +992,9 @@
- }
- 
- // CHECK-LABEL: "op_batch_norm_training"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_batch_norm_training(%arg0: tensor<16x16x16x16xf32>, %arg1: tensor<16xf32>, %arg2: tensor<16xf32>) -> (tensor<16x16x16x16xf32>, tensor<16xf32>, tensor<16xf32>) {
--  //      CHECK: "vhlo.batch_norm_training_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.batch_norm_training_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   epsilon = #vhlo.float_v1<1.000000e-03 : !vhlo.f32_v1>,
-   // CHECK-SAME:   feature_index = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<16x16x16x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>) -> (!vhlo.tensor_v1<16x16x16x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>)
-@@ -941,15 +1006,17 @@
- }
- 
- // CHECK-LABEL: "op_bitcast_convert"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_bitcast_convert(%arg0: tensor<i32>) -> tensor<f32> {
--  // CHECK: "vhlo.bitcast_convert_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.bitcast_convert_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.bitcast_convert"(%arg0) : (tensor<i32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_broadcast_in_dim"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_broadcast_in_dim(%arg0: tensor<16xf32>) -> tensor<16x16xf32> {
--  //      CHECK: "vhlo.broadcast_in_dim_v1"(%arg0) <{
-+  //      CHECK: "vhlo.broadcast_in_dim_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   broadcast_dimensions = #vhlo.tensor_v1<dense<1> : tensor<1xi64>>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x16x!vhlo.f32_v1>
-   %0 = "stablehlo.broadcast_in_dim"(%arg0) {
-@@ -959,8 +1026,9 @@
- }
- 
- // CHECK-LABEL: "op_broadcast"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_broadcast(%arg0: tensor<16xf32>) -> tensor<16x16xf32> {
--  //      CHECK: "vhlo.broadcast_v1"(%arg0) <{
-+  //      CHECK: "vhlo.broadcast_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   broadcast_sizes = #vhlo.tensor_v1<dense<16> : tensor<1xi64>>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x16x!vhlo.f32_v1>
-   %0 = "stablehlo.broadcast"(%arg0) {
-@@ -970,9 +1038,10 @@
- }
- 
- // CHECK-LABEL: "op_case"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_case(%arg0: tensor<i32>, %arg1: tensor<f32>) -> tensor<f32> {
--  //      CHECK: "vhlo.case_v1"(%arg0) ({
--  // CHECK-NEXT:   "vhlo.return_v1"(%arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-+  //      CHECK: "vhlo.case_v1"(%[[ARG0]]) ({
-+  // CHECK-NEXT:   "vhlo.return_v1"(%[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-   // CHECK-NEXT: }) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.case"(%arg0) ({
-     "stablehlo.return"(%arg1) : (tensor<f32>) -> ()
-@@ -981,22 +1050,25 @@
- }
- 
- // CHECK-LABEL: "op_cbrt"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_cbrt(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.cbrt_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.cbrt_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.cbrt"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_ceil"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_ceil(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.ceil_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.ceil_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.ceil"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_cholesky"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_cholesky(%arg0: tensor<1x16x16xf32>) -> tensor<1x16x16xf32> {
--  //      CHECK: "vhlo.cholesky_v1"(%arg0) <{
-+  //      CHECK: "vhlo.cholesky_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   lower = #vhlo.bool_v1<true>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<1x16x16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<1x16x16x!vhlo.f32_v1>
-   %0 = "stablehlo.cholesky"(%arg0) {
-@@ -1006,22 +1078,25 @@
- }
- 
- // CHECK-LABEL: "op_clamp"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_clamp(%arg0: tensor<f32>, %arg1: tensor<f32>, %arg2: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.clamp_v1"(%arg0, %arg1, %arg2) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.clamp_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.clamp"(%arg0, %arg1, %arg2) : (tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_count_leading_zeros"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_count_leading_zeros(%arg0: tensor<i32>) -> tensor<i32> {
--  // CHECK: "vhlo.count_leading_zeros_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-+  // CHECK: "vhlo.count_leading_zeros_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-   %0 = "stablehlo.count_leading_zeros"(%arg0) : (tensor<i32>) -> tensor<i32>
-   func.return %0 : tensor<i32>
- }
- 
- // CHECK-LABEL: "op_collective_permute"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_collective_permute(%arg0: tensor<16x8xf32>) -> tensor<16x8xf32> {
--  //               CHECK: "vhlo.collective_permute_v1"(%arg0) <{
-+  //               CHECK: "vhlo.collective_permute_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME{LITERAL}:   source_target_pairs = #vhlo.tensor_v1<dense<[[0, 1], [1, 2], [2, 3]]> : tensor<3x2xi64>>
-   //          CHECK-SAME: }> : (!vhlo.tensor_v1<16x8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x8x!vhlo.f32_v1>
-@@ -1033,8 +1108,9 @@
- }
- 
- // CHECK-LABEL: "op_compare"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_compare(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
--  //      CHECK: "vhlo.compare_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.compare_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   compare_type = #vhlo<comparison_type_v1 TOTALORDER>,
-   // CHECK-SAME:   comparison_direction = #vhlo<comparison_direction_v1 EQ>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-@@ -1046,15 +1122,17 @@
- }
- 
- // CHECK-LABEL: "op_complex"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_complex(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<complex<f32>> {
--  // CHECK: "vhlo.complex_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>
-+  // CHECK: "vhlo.complex_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>
-   %0 = "stablehlo.complex"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<complex<f32>>
-   func.return %0 : tensor<complex<f32>>
- }
- 
- // CHECK-LABEL: "op_composite"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_composite(%arg0: tensor<f32>) -> tensor<f32> {
--  //               CHECK: "vhlo.composite_v1"(%arg0) <{
-+  //               CHECK: "vhlo.composite_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   composite_attributes = #vhlo.dict_v1<{#vhlo.string_v1<"my_int"> = #vhlo.integer_v1<1 : i64>, #vhlo.string_v1<"my_string"> = #vhlo.string_v1<"foo">}>
-   //          CHECK-SAME:   decomposition = #vhlo.string_v1<"composite_target">
-   //          CHECK-SAME:   name = #vhlo.string_v1<"stablehlo.composite_target">
-@@ -1073,8 +1151,9 @@
- }
- 
- // CHECK-LABEL: "op_concatenate"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_concatenate(%arg0: tensor<8xf32>, %arg1: tensor<8xf32>) -> tensor<16xf32> {
--  //      CHECK: "vhlo.concatenate_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.concatenate_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   dimension = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<8x!vhlo.f32_v1>, !vhlo.tensor_v1<8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x!vhlo.f32_v1>
-   %0 = "stablehlo.concatenate"(%arg0, %arg1) {
-@@ -1084,6 +1163,7 @@
- }
- 
- // CHECK-LABEL: "op_constant"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_constant(%arg0: tensor<f32>) -> tensor<f32> {
-   //      CHECK: "vhlo.constant_v1"() <{
-   // CHECK-SAME:   value = #vhlo.tensor_v1<dense<0.000000e+00> : tensor<f32>>
-@@ -1095,15 +1175,17 @@
- }
- 
- // CHECK-LABEL: "op_convert"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_convert(%arg0: tensor<i32>) -> tensor<f32> {
--  // CHECK: "vhlo.convert_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.convert_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.convert"(%arg0) : (tensor<i32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_convolution"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_convolution(%arg0: tensor<1x8x8x207xf32>, %arg1: tensor<3x3x207x16xf32>) -> tensor<1x7x7x16xf32> {
--  //      CHECK: "vhlo.convolution_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.convolution_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   batch_group_count = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   feature_group_count = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   input_batch_dimension = #vhlo.integer_v1<0 : i64>,
-@@ -1137,8 +1219,9 @@
- }
- 
- // CHECK-LABEL: "op_cosine"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_cosine(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.cosine_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.cosine_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.cosine"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
-@@ -1151,8 +1234,9 @@
- }
- 
- // CHECK-LABEL: "op_cross_replica_sum"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_cross_replica_sum(%arg0: tensor<f32>) -> tensor<f32> {
--  //               CHECK: "vhlo.cross-replica-sum_v1"(%arg0) <{
-+  //               CHECK: "vhlo.cross-replica-sum_v1"(%[[ARG0]]) <{
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0], [1]]> : tensor<2x1xi64>>
-   //          CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.cross-replica-sum"(%arg0) {
-@@ -1162,8 +1246,9 @@
- }
- 
- // CHECK-LABEL: "op_custom_call"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_custom_call(%arg0: tensor<f32>) -> tensor<f32> {
--  //      CHECK: "vhlo.custom_call_v1"(%arg0) <{
-+  //      CHECK: "vhlo.custom_call_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   api_version = #vhlo<api_version_v1 API_VERSION_STATUS_RETURNING>,
-   // CHECK-SAME:   backend_config = #vhlo.string_v1<"\08\03\1A\02">,
-   // CHECK-SAME:   call_target_name = #vhlo.string_v1<"foo">,
-@@ -1194,15 +1279,17 @@
- }
- 
- // CHECK-LABEL: "op_divide"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_divide(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.divide_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.divide_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.divide"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_dot_general"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_dot_general(%arg0: tensor<8x8x16xf32>, %arg1: tensor<8x16x8xf32>) -> tensor<8x8x8xf32> {
--  //      CHECK: "vhlo.dot_general_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.dot_general_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   lhs_batching_dimensions = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>,
-   // CHECK-SAME:   lhs_contracting_dimensions = #vhlo.tensor_v1<dense<2> : tensor<1xi64>>,
-   // CHECK-SAME:   precision_config = #vhlo.array_v1<[#vhlo<precision_v1 HIGHEST>, #vhlo<precision_v1 HIGHEST>]>,
-@@ -1222,8 +1309,9 @@
- }
- 
- // CHECK-LABEL: "op_dot"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_dot(%arg0: tensor<8x16xf32>, %arg1: tensor<16x8xf32>) -> tensor<8x8xf32> {
--  //      CHECK: "vhlo.dot_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.dot_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   precision_config = #vhlo.array_v1<[#vhlo<precision_v1 HIGHEST>, #vhlo<precision_v1 HIGHEST>]>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<8x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<8x8x!vhlo.f32_v1>
-   %0 = "stablehlo.dot"(%arg0, %arg1) {
-@@ -1233,8 +1321,9 @@
- }
- 
- // CHECK-LABEL: "op_dynamic_broadcast_in_dim"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_dynamic_broadcast_in_dim(%arg0: tensor<?x?xf32>, %arg1: tensor<2xindex>) -> tensor<?x?xf32> {
--  //      CHECK: "vhlo.dynamic_broadcast_in_dim_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.dynamic_broadcast_in_dim_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   broadcast_dimensions = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-   // CHECK-SAME:   known_expanding_dimensions = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>,
-   // CHECK-SAME:   known_nonexpanding_dimensions = #vhlo.tensor_v1<dense<1> : tensor<1xi64>>
-@@ -1248,8 +1337,9 @@
- }
- 
- // CHECK-LABEL: "op_dynamic_conv"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_dynamic_conv(%arg0: tensor<1x8x8x207xf32>, %arg1: tensor<3x3x207x16xf32>, %arg2: tensor<4xi32>) -> tensor<1x?x?x16xf32> {
--  //      CHECK: "vhlo.dynamic_conv_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.dynamic_conv_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   batch_group_count = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   feature_group_count = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   input_batch_dimension = #vhlo.integer_v1<0 : i64>,
-@@ -1283,8 +1373,9 @@
- }
- 
- // CHECK-LABEL: "op_dynamic_gather"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_dynamic_gather(%arg0 : tensor<2x4x9xf32>, %arg1 : tensor<1x5x2xi32>, %arg2 : tensor<3xi32>) -> tensor<1x5x8xf32> {
--  //      CHECK: "vhlo.dynamic_gather_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.dynamic_gather_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   collapsed_slice_dims = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-   // CHECK-SAME:   index_vector_dim = #vhlo.integer_v1<2 : i64>,
-   // CHECK-SAME:   indices_are_sorted = #vhlo.bool_v1<true>,
-@@ -1304,8 +1395,9 @@
- }
- 
- // CHECK-LABEL: "op_dynamic_iota"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_dynamic_iota(%arg0: tensor<1xindex>) -> tensor<?xf32> {
--  //      CHECK: "vhlo.dynamic_iota_v1"(%arg0) <{
-+  //      CHECK: "vhlo.dynamic_iota_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   iota_dimension = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<1x!vhlo.index_v1>) -> !vhlo.tensor_v1<?x!vhlo.f32_v1>
-   %0 = "stablehlo.dynamic_iota"(%arg0) {
-@@ -1315,22 +1407,25 @@
- }
- 
- // CHECK-LABEL: "op_dynamic_pad"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}}, %[[ARG3:.*]]: {{.*}}, %[[ARG4:.*]]: {{.*}})
- func.func @op_dynamic_pad(%arg0: tensor<?xf32>, %arg1: tensor<f32>, %arg2: tensor<1xindex>, %arg3: tensor<1xindex>, %arg4: tensor<1xindex>) -> tensor<?xf32> {
--  // CHECK: "vhlo.dynamic_pad_v1"(%arg0, %arg1, %arg2, %arg3, %arg4) : (!vhlo.tensor_v1<?x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>) -> !vhlo.tensor_v1<?x!vhlo.f32_v1>
-+  // CHECK: "vhlo.dynamic_pad_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]], %[[ARG3]], %[[ARG4]]) : (!vhlo.tensor_v1<?x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>) -> !vhlo.tensor_v1<?x!vhlo.f32_v1>
-   %0 = "stablehlo.dynamic_pad"(%arg0, %arg1, %arg2, %arg3, %arg4) : (tensor<?xf32>, tensor<f32>, tensor<1xindex>, tensor<1xindex>, tensor<1xindex>) -> tensor<?xf32>
-   func.return %0 : tensor<?xf32>
- }
- 
- // CHECK-LABEL: "op_dynamic_reshape"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_dynamic_reshape(%arg0: tensor<16xf32>, %arg1: tensor<2xindex>) -> tensor<?x?xf32> {
--  // CHECK: "vhlo.dynamic_reshape_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<2x!vhlo.index_v1>) -> !vhlo.tensor_v1<?x?x!vhlo.f32_v1>
-+  // CHECK: "vhlo.dynamic_reshape_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<2x!vhlo.index_v1>) -> !vhlo.tensor_v1<?x?x!vhlo.f32_v1>
-   %0 = "stablehlo.dynamic_reshape"(%arg0, %arg1) : (tensor<16xf32>, tensor<2xindex>) -> tensor<?x?xf32>
-   func.return %0 : tensor<?x?xf32>
- }
- 
- // CHECK-LABEL: "op_dynamic_slice"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_dynamic_slice(%arg0: tensor<16xf32>, %arg1: tensor<i64>) -> tensor<4xf32> {
--  //      CHECK: "vhlo.dynamic_slice_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.dynamic_slice_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   slice_sizes = #vhlo.tensor_v1<dense<4> : tensor<1xi64>>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.i64_v1>) -> !vhlo.tensor_v1<4x!vhlo.f32_v1>
-   %0 = "stablehlo.dynamic_slice"(%arg0, %arg1) {
-@@ -1340,15 +1435,17 @@
- }
- 
- // CHECK-LABEL: "op_dynamic_update_slice"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_dynamic_update_slice(%arg0: tensor<16xf32>, %arg1: tensor<4xf32>, %arg2: tensor<i64>) -> tensor<16xf32> {
--  // CHECK: "vhlo.dynamic_update_slice_v1"(%arg0, %arg1, %arg2) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<4x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.i64_v1>) -> !vhlo.tensor_v1<16x!vhlo.f32_v1>
-+  // CHECK: "vhlo.dynamic_update_slice_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<4x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.i64_v1>) -> !vhlo.tensor_v1<16x!vhlo.f32_v1>
-   %0 = "stablehlo.dynamic_update_slice"(%arg0, %arg1, %arg2) : (tensor<16xf32>, tensor<4xf32>, tensor<i64>) -> tensor<16xf32>
-   func.return %0 : tensor<16xf32>
- }
- 
- // CHECK-LABEL: "op_einsum"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_einsum(%arg0: tensor<8x16xf32>, %arg1: tensor<16x8xf32>) -> tensor<8x8xf32> {
--  //      CHECK: "vhlo.einsum_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.einsum_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   einsum_config = #vhlo.string_v1<"ab,bc->ac">
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<8x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<8x8x!vhlo.f32_v1>
-   %0 = "stablehlo.einsum"(%arg0, %arg1) {
-@@ -1358,22 +1455,25 @@
- }
- 
- // CHECK-LABEL: "op_exponential_minus_one"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_exponential_minus_one(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.exponential_minus_one_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.exponential_minus_one_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.exponential_minus_one"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_exponential"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_exponential(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.exponential_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.exponential_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.exponential"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_fft"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_fft(%arg0: tensor<16xcomplex<f32>>) -> tensor<16xcomplex<f32>> {
--  //      CHECK: "vhlo.fft_v1"(%arg0) <{
-+  //      CHECK: "vhlo.fft_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   fft_length = #vhlo.tensor_v1<dense<16> : tensor<1xi64>>,
-   // CHECK-SAME:   fft_type = #vhlo<fft_type_v1 FFT>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<16x!vhlo.complex_v1<!vhlo.f32_v1>>) -> !vhlo.tensor_v1<16x!vhlo.complex_v1<!vhlo.f32_v1>>
-@@ -1385,8 +1485,9 @@
- }
- 
- // CHECK-LABEL: "op_floor"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_floor(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.floor_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.floor_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.floor"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
-@@ -1399,16 +1500,17 @@
-   // CHECK-SAME:   sym_name = #vhlo.string_v1<"op_func">,
-   // CHECK-SAME:   sym_visibility = #vhlo.string_v1<"private">
-   // CHECK-SAME: }> ({
--  // CHECK-NEXT:   ^[[BB:bb.*]](%arg0: !vhlo.tensor_v1<!vhlo.f32_v1>):
--  // CHECK-NEXT:     "vhlo.return_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-+  // CHECK-NEXT:   ^[[BB:bb.*]](%[[ARG0:.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
-+  // CHECK-NEXT:     "vhlo.return_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-   // CHECK-NEXT: }) : () -> ()
- 
-   func.return %arg0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_gather"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_gather(%arg0 : tensor<2x4x9xf32>, %arg1 : tensor<1x5x2xi32>) -> tensor<1x5x1xf32> {
--  //      CHECK: "vhlo.gather_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.gather_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   collapsed_slice_dims = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-   // CHECK-SAME:   index_vector_dim = #vhlo.integer_v1<2 : i64>,
-   // CHECK-SAME:   indices_are_sorted = #vhlo.bool_v1<true>,
-@@ -1430,8 +1532,9 @@
- }
- 
- // CHECK-LABEL: "op_get_dimension_size"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_get_dimension_size(%arg0: tensor<?xf32>) -> tensor<i32> {
--  //      CHECK: "vhlo.get_dimension_size_v1"(%arg0) <{
-+  //      CHECK: "vhlo.get_dimension_size_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   dimension = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<?x!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-   %0 = "stablehlo.get_dimension_size"(%arg0) {
-@@ -1441,8 +1544,9 @@
- }
- 
- // CHECK-LABEL: "op_get_tuple_element"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_get_tuple_element(%arg0: tuple<tensor<f32>, tensor<i32>>) -> tensor<f32> {
--  //      CHECK: "vhlo.get_tuple_element_v1"(%arg0) <{
-+  //      CHECK: "vhlo.get_tuple_element_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   index = #vhlo.integer_v1<0 : i32>
-   // CHECK-SAME: }> : (!vhlo.tuple_v1<!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.get_tuple_element"(%arg0) {
-@@ -1452,11 +1556,12 @@
- }
- 
- // CHECK-LABEL: "op_if"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_if(%arg0: tensor<i1>, %arg1: tensor<f32>, %arg2: tensor<f32>) -> tensor<f32> {
--  //      CHECK: "vhlo.if_v1"(%arg0) ({
--  // CHECK-NEXT:   "vhlo.return_v1"(%arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-+  //      CHECK: "vhlo.if_v1"(%[[ARG0]]) ({
-+  // CHECK-NEXT:   "vhlo.return_v1"(%[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-   // CHECK-NEXT: }, {
--  // CHECK-NEXT:   "vhlo.return_v1"(%arg2) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-+  // CHECK-NEXT:   "vhlo.return_v1"(%[[ARG2]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-   // CHECK-NEXT: }) : (!vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.if"(%arg0) ({
-     "stablehlo.return"(%arg1) : (tensor<f32>) -> ()
-@@ -1467,15 +1572,17 @@
- }
- 
- // CHECK-LABEL: "op_imag"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_imag(%arg0: tensor<complex<f32>>) -> tensor<f32> {
--  // CHECK: "vhlo.imag_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.imag_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.imag"(%arg0) : (tensor<complex<f32>>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_infeed"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_infeed(%arg0: !stablehlo.token) -> (tensor<f32>, !stablehlo.token) {
--  //               CHECK: "vhlo.infeed_v1"(%arg0) <{
-+  //               CHECK: "vhlo.infeed_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   infeed_config = #vhlo.string_v1<"foo">,
-   // CHECK-SAME{LITERAL}:   layout = #vhlo.array_v1<[#vhlo.array_v1<[]>]>
-   //          CHECK-SAME: }> : (!vhlo.token_v1) -> (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.token_v1)
-@@ -1498,36 +1605,41 @@
- }
- 
- // CHECK-LABEL: "op_is_finite"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_is_finite(%arg0: tensor<f32>) -> tensor<i1> {
--  // CHECK: "vhlo.is_finite_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-+  // CHECK: "vhlo.is_finite_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-   %0 = "stablehlo.is_finite"(%arg0) : (tensor<f32>) -> tensor<i1>
-   func.return %0 : tensor<i1>
- }
- 
- // CHECK-LABEL: "op_log"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_log(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.log_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.log_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.log"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_log_plus_one"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_log_plus_one(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.log_plus_one_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.log_plus_one_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.log_plus_one"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_logistic"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_logistic(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.logistic_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.logistic_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.logistic"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_map"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_map(%arg0: tensor<16xf32>) -> tensor<16xf32> {
--  //      CHECK: "vhlo.map_v1"(%arg0) <{
-+  //      CHECK: "vhlo.map_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   dimensions = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>
-   // CHECK-SAME: }> ({
-   // CHECK-NEXT:   ^[[BB:bb.*]](%[[ARG1:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
-@@ -1545,57 +1657,65 @@
- }
- 
- // CHECK-LABEL: "op_maximum"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_maximum(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.maximum_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.maximum_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.maximum"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_minimum"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_minimum(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.minimum_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.minimum_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.minimum"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_multiply"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_multiply(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.multiply_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.multiply_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.multiply"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_negate"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_negate(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.negate_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.negate_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.negate"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_not"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_not(%arg0: tensor<i1>) -> tensor<i1> {
--  // CHECK: "vhlo.not_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-+  // CHECK: "vhlo.not_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-   %0 = "stablehlo.not"(%arg0) : (tensor<i1>) -> tensor<i1>
-   func.return %0 : tensor<i1>
- }
- 
- // CHECK-LABEL: "op_optimization_barrier"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_optimization_barrier(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.optimization_barrier_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.optimization_barrier_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.optimization_barrier"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_or"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_or(%arg0: tensor<i1>, %arg1: tensor<i1>) -> tensor<i1> {
--  // CHECK: "vhlo.or_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-+  // CHECK: "vhlo.or_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-   %0 = "stablehlo.or"(%arg0, %arg1) : (tensor<i1>, tensor<i1>) -> tensor<i1>
-   func.return %0 : tensor<i1>
- }
- 
- // CHECK-LABEL: "op_outfeed"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_outfeed(%arg0: tensor<f32>, %arg1: !stablehlo.token) -> !stablehlo.token {
--  //      CHECK: "vhlo.outfeed_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.outfeed_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   outfeed_config = #vhlo.string_v1<"foo">
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.token_v1) -> !vhlo.token_v1
-   %0 = "stablehlo.outfeed"(%arg0, %arg1) {
-@@ -1605,8 +1725,9 @@
- }
- 
- // CHECK-LABEL: "op_pad"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_pad(%arg0: tensor<8xf32>, %arg1: tensor<f32>) -> tensor<16xf32> {
--  //      CHECK: "vhlo.pad_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.pad_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   edge_padding_high = #vhlo.tensor_v1<dense<4> : tensor<1xi64>>,
-   // CHECK-SAME:   edge_padding_low = #vhlo.tensor_v1<dense<4> : tensor<1xi64>>,
-   // CHECK-SAME:   interior_padding = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>
-@@ -1620,36 +1741,41 @@
- }
- 
- // CHECK-LABEL: "op_popcnt"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_popcnt(%arg0: tensor<i32>) -> tensor<i32> {
--  // CHECK: "vhlo.popcnt_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-+  // CHECK: "vhlo.popcnt_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-   %0 = "stablehlo.popcnt"(%arg0) : (tensor<i32>) -> tensor<i32>
-   func.return %0 : tensor<i32>
- }
- 
- // CHECK-LABEL: "op_power"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_power(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.power_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.power_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.power"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_real_dynamic_slice"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}}, %[[ARG3:.*]]: {{.*}})
- func.func @op_real_dynamic_slice(%arg0: tensor<?xf32>, %arg1: tensor<1xindex>, %arg2: tensor<1xindex>, %arg3: tensor<1xindex>) -> tensor<?xf32> {
--  // CHECK: "vhlo.real_dynamic_slice_v1"(%arg0, %arg1, %arg2, %arg3) : (!vhlo.tensor_v1<?x!vhlo.f32_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>) -> !vhlo.tensor_v1<?x!vhlo.f32_v1>
-+  // CHECK: "vhlo.real_dynamic_slice_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]], %[[ARG3]]) : (!vhlo.tensor_v1<?x!vhlo.f32_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>) -> !vhlo.tensor_v1<?x!vhlo.f32_v1>
-   %0 = "stablehlo.real_dynamic_slice"(%arg0, %arg1, %arg2, %arg3) : (tensor<?xf32>, tensor<1xindex>, tensor<1xindex>, tensor<1xindex>) -> tensor<?xf32>
-   func.return %0 : tensor<?xf32>
- }
- 
- // CHECK-LABEL: "op_real"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_real(%arg0: tensor<complex<f32>>) -> tensor<f32> {
--  // CHECK: "vhlo.real_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.real_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.real"(%arg0) : (tensor<complex<f32>>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_recv"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_recv(%arg0: !stablehlo.token) -> (tensor<f32>, !stablehlo.token) {
--  //      CHECK: "vhlo.recv_v1"(%arg0) <{
-+  //      CHECK: "vhlo.recv_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME:   channel_type = #vhlo.integer_v1<3 : i64>,
-   // CHECK-SAME:   is_host_transfer = #vhlo.bool_v1<true>
-@@ -1662,8 +1788,9 @@
- }
- 
- // CHECK-LABEL: "op_reduce"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_reduce(%arg0: tensor<16xf32>, %arg1: tensor<f32>) -> tensor<f32> {
--  //  CHECK: "vhlo.reduce_v1"(%arg0, %arg1)
-+  //  CHECK: "vhlo.reduce_v1"(%[[ARG0]], %[[ARG1]])
-   //  CHECK:   ^[[BB:bb.*]](%[[ARG1:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>, %[[ARG2:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
-   //  CHECK:     "vhlo.return_v1"(%[[VAL1:.*]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-   //  CHECK: }) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-@@ -1678,8 +1805,9 @@
- }
- 
- // CHECK-LABEL: "op_reduce_precision"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_reduce_precision(%arg0: tensor<f32>) -> tensor<f32> {
--  //      CHECK: "vhlo.reduce_precision_v1"(%arg0) <{
-+  //      CHECK: "vhlo.reduce_precision_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   exponent_bits = #vhlo.integer_v1<8 : i32>
-   // CHECK-SAME:   mantissa_bits = #vhlo.integer_v1<10 : i32>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-@@ -1690,10 +1818,10 @@
-   func.return %0 : tensor<f32>
- }
- 
--// CHECK_lABEL: "op_reduce_with_promotable_types"
-+// CHECK-LABEL: "op_reduce_with_promotable_types"
- func.func @op_reduce_with_promotable_types(%arg0: tensor<4x4xf32>, %arg1 : tensor<f32>)
-     -> (tensor<4xf64>) {
--  //  CHECK: "vhlo.reduce_v1"(%arg0, %arg1)
-+  //  CHECK: "vhlo.reduce_v1"(%[[ARG0:.*]], %[[ARG1:.*]])
-   //  CHECK:   ^[[BB:bb.*]](%[[ARG1:arg.*]]: !vhlo.tensor_v1<!vhlo.f64_v1>, %[[ARG2:arg.*]]: !vhlo.tensor_v1<!vhlo.f64_v1>):
-   //  CHECK:     "vhlo.return_v1"(%[[VAL1:.*]]) : (!vhlo.tensor_v1<!vhlo.f64_v1>) -> ()
-   //  CHECK: }) : (!vhlo.tensor_v1<4x4x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<4x!vhlo.f64_v1>
-@@ -1708,8 +1836,9 @@
- }
- 
- // CHECK-LABEL: "op_reduce_scatter"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_reduce_scatter(%arg0: tensor<16xf32>) -> tensor<16xf32> {
--  //               CHECK: "vhlo.reduce_scatter_v1"(%arg0) <{
-+  //               CHECK: "vhlo.reduce_scatter_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0], [1]]> : tensor<2x1xi64>>,
-   //          CHECK-SAME:   scatter_dimension = #vhlo.integer_v1<0 : i64>
-@@ -1734,7 +1863,7 @@
- 
- // CHECK_lABEL: "op_reduce_scatter_with_promotable_types"
- func.func @op_reduce_scatter_with_promotable_types(%data: tensor<4x16xf32>) -> tensor<4x4xf64> {
--  //  CHECK: "vhlo.reduce_scatter_v1"(%arg0)
-+  //  CHECK: "vhlo.reduce_scatter_v1"(%[[ARG0:.*]])
-   //  CHECK:   ^[[BB:bb.*]](%[[ARG1:arg.*]]: !vhlo.tensor_v1<!vhlo.f64_v1>, %[[ARG2:arg.*]]: !vhlo.tensor_v1<!vhlo.f64_v1>):
-   //  CHECK:     "vhlo.return_v1"(%[[VAL1:.*]]) : (!vhlo.tensor_v1<!vhlo.f64_v1>) -> ()
-   //  CHECK: }) : (!vhlo.tensor_v1<4x16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<4x4x!vhlo.f64_v1>
-@@ -1751,8 +1880,9 @@
- 
- 
- // CHECK-LABEL: "op_reduce_window"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_reduce_window(%arg0: tensor<2x17x31x7xf32>, %arg1: tensor<f32>) -> tensor<2x9x16x7xf32> {
--  //               CHECK: "vhlo.reduce_window_v1"(%arg0, %arg1) <{
-+  //               CHECK: "vhlo.reduce_window_v1"(%[[ARG0]], %[[ARG1]]) <{
-   //          CHECK-SAME:   base_dilations = #vhlo.tensor_v1<dense<[1, 2, 2, 1]> : tensor<4xi64>>,
-   // CHECK-SAME{LITERAL}:   padding = #vhlo.tensor_v1<dense<[[0, 0], [2, 0], [0, 2], [0, 0]]> : tensor<4x2xi64>>,
-   //          CHECK-SAME:   window_dilations = #vhlo.tensor_v1<dense<[1, 2, 2, 1]> : tensor<4xi64>>,
-@@ -1781,7 +1911,7 @@
- func.func @op_reduce_window_with_promotable_types(%arg0: tensor<4x2xf32>,
-     %arg1: tensor<4x2xf32>, %init0: tensor<f32>, %init1: tensor<f32>) ->
-     (tensor<2x2xf64>, tensor<2x2xf32>) {
--  //  CHECK: "vhlo.reduce_window_v1"(%arg0, %arg1, %arg2, %arg3)
-+  //  CHECK: "vhlo.reduce_window_v1"(%[[ARG0:.*]], %[[ARG1:.*]], %[[ARG2:.*]], %[[ARG3:.*]])
-   //  CHECK:   ^[[BB:bb.*]](%[[ARG1:arg.*]]: !vhlo.tensor_v1<!vhlo.f64_v1>, %[[ARG2:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>, %[[ARG3:arg.*]]: !vhlo.tensor_v1<!vhlo.f64_v1>, %[[ARG4:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
-   //  CHECK:     "vhlo.return_v1"(%[[VAL1:.*]], %[[VAL2:.*]]) : (!vhlo.tensor_v1<!vhlo.f64_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-   //  CHECK: }) : (!vhlo.tensor_v1<4x2x!vhlo.f32_v1>, !vhlo.tensor_v1<4x2x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> (!vhlo.tensor_v1<2x2x!vhlo.f64_v1>, !vhlo.tensor_v1<2x2x!vhlo.f32_v1>)
-@@ -1801,8 +1931,9 @@
- }
- 
- // CHECK-LABEL: "op_remainder"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_remainder(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.remainder_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.remainder_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.remainder"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
-@@ -1822,16 +1953,18 @@
- }
- 
- // CHECK-LABEL: "op_reshape"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_reshape(%arg0: tensor<16xf32>) -> tensor<4x4xf32> {
--  // CHECK: "vhlo.reshape_v1"(%arg0) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<4x4x!vhlo.f32_v1>
-+  // CHECK: "vhlo.reshape_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<4x4x!vhlo.f32_v1>
-   %0 = "stablehlo.reshape"(%arg0) : (tensor<16xf32>) -> tensor<4x4xf32>
-   func.return %0 : tensor<4x4xf32>
- }
- 
- // CHECK-LABEL: "op_return"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_return(%arg0: tensor<i32>, %arg1: tensor<f32>) -> tensor<f32> {
--  //      CHECK: "vhlo.case_v1"(%arg0) ({
--  // CHECK-NEXT:   "vhlo.return_v1"(%arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-+  //      CHECK: "vhlo.case_v1"(%[[ARG0]]) ({
-+  // CHECK-NEXT:   "vhlo.return_v1"(%[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-   // CHECK-NEXT: }) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.case"(%arg0) ({
-     "stablehlo.return"(%arg1) : (tensor<f32>) -> ()
-@@ -1840,8 +1973,9 @@
- }
- 
- // CHECK-LABEL: "op_reverse"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_reverse(%arg0: tensor<16xf32>) -> tensor<16xf32> {
--  //      CHECK: "vhlo.reverse_v1"(%arg0) <{
-+  //      CHECK: "vhlo.reverse_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   dimensions = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x!vhlo.f32_v1>
-   %0 = "stablehlo.reverse"(%arg0) {
-@@ -1851,8 +1985,9 @@
- }
- 
- // CHECK-LABEL: "op_rng_bit_generator"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_rng_bit_generator(%arg0: tensor<f32>) -> (tensor<f32>, tensor<f32>) {
--  //      CHECK: "vhlo.rng_bit_generator_v1"(%arg0) <{
-+  //      CHECK: "vhlo.rng_bit_generator_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   rng_algorithm = #vhlo<rng_algorithm_v1 PHILOX>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>)
-   %0:2 = "stablehlo.rng_bit_generator"(%arg0) {
-@@ -1862,8 +1997,9 @@
- }
- 
- // CHECK-LABEL: "op_rng"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_rng(%arg0: tensor<f32>, %arg1: tensor<f32>, %arg2: tensor<0xindex>) -> tensor<f32> {
--  //      CHECK: "vhlo.rng_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.rng_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   rng_distribution = #vhlo<rng_distribution_v1 NORMAL>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<0x!vhlo.index_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.rng"(%arg0, %arg1, %arg2) {
-@@ -1873,29 +2009,33 @@
- }
- 
- // CHECK-LABEL: "op_round_nearest_afz"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_round_nearest_afz(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.round_nearest_afz_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.round_nearest_afz_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.round_nearest_afz"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_round_nearest_even"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_round_nearest_even(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.round_nearest_even_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.round_nearest_even_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.round_nearest_even"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_rsqrt"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_rsqrt(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.rsqrt_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.rsqrt_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.rsqrt"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_scatter"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_scatter(%arg0: tensor<200x100x300xf32>, %arg1: tensor<10x2xi32>, %arg2: tensor<10x300xf32>) -> tensor<200x100x300xf32> {
--  //      CHECK: "vhlo.scatter_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.scatter_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   index_vector_dim = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   indices_are_sorted = #vhlo.bool_v1<true>,
-   // CHECK-SAME:   inserted_window_dims = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-@@ -1928,7 +2068,7 @@
- func.func @op_scatter_with_promotable_types(%input_tensor: tensor<200x100x300xf32>,
-     %scatter_indices: tensor<10x2xi32>, %updates: tensor<10x300xf32>) ->
-       tensor<200x100x300xf64> {
--  //  CHECK: "vhlo.scatter_v1"(%arg0, %arg1, %arg2)
-+  //  CHECK: "vhlo.scatter_v1"(%[[ARG0:.*]], %[[ARG1:.*]], %[[ARG2:.*]])
-   //  CHECK:   ^[[BB:bb.*]](%[[ARG1:arg.*]]: !vhlo.tensor_v1<!vhlo.f64_v1>, %[[ARG2:arg.*]]: !vhlo.tensor_v1<!vhlo.f64_v1>):
-   //  CHECK:     "vhlo.return_v1"(%[[VAL1:.*]]) : (!vhlo.tensor_v1<!vhlo.f64_v1>) -> ()
-   //  CHECK: }) : (!vhlo.tensor_v1<200x100x300x!vhlo.f32_v1>, !vhlo.tensor_v1<10x2x!vhlo.i32_v1>, !vhlo.tensor_v1<10x300x!vhlo.f32_v1>) -> !vhlo.tensor_v1<200x100x300x!vhlo.f64_v1>
-@@ -1951,8 +2091,9 @@
- }
- 
- // CHECK-LABEL: "op_select_and_scatter"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_select_and_scatter(%arg0: tensor<10x24x24x64xf32>, %arg1: tensor<12x13x13x66xf32>, %arg2: tensor<f32>) -> tensor<10x24x24x64xf32> {
--  //      CHECK: "vhlo.select_and_scatter_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.select_and_scatter_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   padding = #vhlo.tensor_v1<dense<1> : tensor<4x2xi64>>,
-   // CHECK-SAME:   window_dimensions = #vhlo.tensor_v1<dense<[1, 2, 2, 1]> : tensor<4xi64>>,
-   // CHECK-SAME:   window_strides = #vhlo.tensor_v1<dense<[1, 2, 2, 1]> : tensor<4xi64>>
-@@ -1982,8 +2123,9 @@
- }
- 
- // CHECK-LABEL: "op_select_and_scatter_with_promotable_types"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_select_and_scatter_with_promotable_types(%arg0: tensor<10x24x24x64xf32>, %arg1: tensor<12x13x13x66xf32>, %arg2: tensor<f32>) -> tensor<10x24x24x64xf64> {
--  // CHECK: "vhlo.select_and_scatter_v1"(%arg0, %arg1, %arg2)
-+  // CHECK: "vhlo.select_and_scatter_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]])
-   // CHECK:   ^[[BB:bb.*]](%[[ARG1:arg.*]]: !vhlo.tensor_v1<!vhlo.f64_v1>, %[[ARG2:arg.*]]: !vhlo.tensor_v1<!vhlo.f64_v1>):
-   // CHECK:     %[[VAL:.*]] = "vhlo.add_v1"(%[[ARG1]], %[[ARG2]]) : (!vhlo.tensor_v1<!vhlo.f64_v1>, !vhlo.tensor_v1<!vhlo.f64_v1>) -> !vhlo.tensor_v1<!vhlo.f64_v1>
-   // CHECK:     "vhlo.return_v1"(%[[VAL]]) : (!vhlo.tensor_v1<!vhlo.f64_v1>) -> ()
-@@ -2005,15 +2147,17 @@
- }
- 
- // CHECK-LABEL: "op_select"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_select(%arg0: tensor<i1>, %arg1: tensor<f32>, %arg2: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.select_v1"(%arg0, %arg1, %arg2) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.select_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.select"(%arg0, %arg1, %arg2) : (tensor<i1>, tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_send"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_send(%arg0: tensor<f32>, %arg1: !stablehlo.token) -> !stablehlo.token {
--  //      CHECK: "vhlo.send_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.send_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME:   channel_type = #vhlo.integer_v1<2 : i64>,
-   // CHECK-SAME:   is_host_transfer = #vhlo.bool_v1<true>
-@@ -2026,8 +2170,9 @@
- }
- 
- // CHECK-LABEL: "op_set_dimension_size"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_set_dimension_size(%arg0: tensor<?xf32>, %arg1: tensor<i32>) -> tensor<16xf32> {
--  //      CHECK: "vhlo.set_dimension_size_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.set_dimension_size_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   dimension = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<?x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<16x!vhlo.f32_v1>
-   %0 = "stablehlo.set_dimension_size"(%arg0, %arg1) {
-@@ -2037,43 +2182,49 @@
- }
- 
- // CHECK-LABEL: "op_shift_left"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_shift_left(%arg0: tensor<i32>, %arg1: tensor<i32>) -> tensor<i32> {
--  // CHECK: "vhlo.shift_left_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.i32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-+  // CHECK: "vhlo.shift_left_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-   %0 = "stablehlo.shift_left"(%arg0, %arg1) : (tensor<i32>, tensor<i32>) -> tensor<i32>
-   func.return %0 : tensor<i32>
- }
- 
- // CHECK-LABEL: "op_shift_right_arithmetic"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_shift_right_arithmetic(%arg0: tensor<i32>, %arg1: tensor<i32>) -> tensor<i32> {
--  // CHECK: "vhlo.shift_right_arithmetic_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.i32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-+  // CHECK: "vhlo.shift_right_arithmetic_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-   %0 = "stablehlo.shift_right_arithmetic"(%arg0, %arg1) : (tensor<i32>, tensor<i32>) -> tensor<i32>
-   func.return %0 : tensor<i32>
- }
- 
- // CHECK-LABEL: "op_shift_right_logical"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_shift_right_logical(%arg0: tensor<i32>, %arg1: tensor<i32>) -> tensor<i32> {
--  // CHECK: "vhlo.shift_right_logical_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.i32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-+  // CHECK: "vhlo.shift_right_logical_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-   %0 = "stablehlo.shift_right_logical"(%arg0, %arg1) : (tensor<i32>, tensor<i32>) -> tensor<i32>
-   func.return %0 : tensor<i32>
- }
- 
- // CHECK-LABEL: "op_sign"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_sign(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.sign_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.sign_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.sign"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_sine"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_sine(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.sine_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.sine_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.sine"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_slice"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_slice(%arg0: tensor<16xf32>) -> tensor<4xf32> {
--  //      CHECK: "vhlo.slice_v1"(%arg0) <{
-+  //      CHECK: "vhlo.slice_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   limit_indices = #vhlo.tensor_v1<dense<4> : tensor<1xi64>>,
-   // CHECK-SAME:   start_indices = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>,
-   // CHECK-SAME:   strides = #vhlo.tensor_v1<dense<1> : tensor<1xi64>>
-@@ -2087,8 +2238,9 @@
- }
- 
- // CHECK-LABEL: "op_sort"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_sort(%arg0: tensor<16xf32>) -> tensor<16xf32> {
--  //      CHECK: "vhlo.sort_v1"(%arg0) <{
-+  //      CHECK: "vhlo.sort_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   dimension = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME:   is_stable = #vhlo.bool_v1<true>
-   // CHECK-SAME: }> ({
-@@ -2108,29 +2260,33 @@
- }
- 
- // CHECK-LABEL: "op_sqrt"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_sqrt(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.sqrt_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.sqrt_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.sqrt"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_subtract"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_subtract(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.subtract_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.subtract_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.subtract"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_tanh"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_tanh(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.tanh_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.tanh_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.tanh"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_torch_index_select"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_torch_index_select(%arg0: tensor<5x1x5xf32>, %arg1: tensor<2xi32>) ->  tensor<2x1x5xf32> {
--  //      CHECK: "vhlo.torch_index_select_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.torch_index_select_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   batch_dims = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME:   dim = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<5x1x5x!vhlo.f32_v1>, !vhlo.tensor_v1<2x!vhlo.i32_v1>) -> !vhlo.tensor_v1<2x1x5x!vhlo.f32_v1>
-@@ -2142,8 +2298,9 @@
- }
- 
- // CHECK-LABEL: "op_trace"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_trace(%arg0: tensor<f32>) {
--  //      CHECK: "vhlo.trace_v1"(%arg0) <{
-+  //      CHECK: "vhlo.trace_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   tag = #vhlo.string_v1<"foo">
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-   "stablehlo.trace"(%arg0) {
-@@ -2153,8 +2310,9 @@
- }
- 
- // CHECK-LABEL: "op_transpose"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_transpose(%arg0: tensor<16x8xf32>) ->  tensor<8x16xf32> {
--  //      CHECK: "vhlo.transpose_v1"(%arg0) <{
-+  //      CHECK: "vhlo.transpose_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   permutation = #vhlo.tensor_v1<dense<[1, 0]> : tensor<2xi64>>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<16x8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<8x16x!vhlo.f32_v1>
-   %0 = "stablehlo.transpose"(%arg0) {
-@@ -2164,8 +2322,9 @@
- }
- 
- // CHECK-LABEL: "op_triangular_solve"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_triangular_solve(%arg0: tensor<16x16xf32>, %arg1: tensor<16x16xf32>) ->  tensor<16x16xf32> {
--  //      CHECK: "vhlo.triangular_solve_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.triangular_solve_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   left_side = #vhlo.bool_v1<true>,
-   // CHECK-SAME:   lower = #vhlo.bool_v1<true>,
-   // CHECK-SAME:   transpose_a = #vhlo<transpose_v1 NO_TRANSPOSE>,
-@@ -2181,15 +2340,17 @@
- }
- 
- // CHECK-LABEL: "op_tuple"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_tuple(%arg0: tensor<f32>) -> tuple<tensor<f32>> {
--  // CHECK: "vhlo.tuple_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tuple_v1<!vhlo.tensor_v1<!vhlo.f32_v1>>
-+  // CHECK: "vhlo.tuple_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tuple_v1<!vhlo.tensor_v1<!vhlo.f32_v1>>
-   %0 = "stablehlo.tuple"(%arg0) : (tensor<f32>) -> tuple<tensor<f32>>
-   func.return %0 : tuple<tensor<f32>>
- }
- 
- // CHECK-LABEL: "op_unary_einsum"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_unary_einsum(%arg0: tensor<8x16xf32>) -> tensor<8xf32> {
--  //      CHECK: "vhlo.unary_einsum_v1"(%arg0) <{
-+  //      CHECK: "vhlo.unary_einsum_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   einsum_config = #vhlo.string_v1<"ab->a">
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<8x16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<8x!vhlo.f32_v1>
-   %0 = "stablehlo.unary_einsum"(%arg0) {
-@@ -2199,22 +2360,25 @@
- }
- 
- // CHECK-LABEL: "op_uniform_dequantize"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_uniform_dequantize(%arg0: tensor<!quant.uniform<i8:f32, 34.0:16>>) -> tensor<f32> {
--  // CHECK: "vhlo.uniform_dequantize_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.uniform_dequantize_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.uniform_dequantize"(%arg0) : (tensor<!quant.uniform<i8:f32, 34.0:16>>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_uniform_quantize"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_uniform_quantize(%arg0: tensor<f32>) -> tensor<!quant.uniform<i8:f32, 34.0:16>> {
--  // CHECK: "vhlo.uniform_quantize_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>
-+  // CHECK: "vhlo.uniform_quantize_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>
-   %0 = "stablehlo.uniform_quantize"(%arg0) : (tensor<f32>) -> tensor<!quant.uniform<i8:f32, 34.0:16>>
-   func.return %0 : tensor<!quant.uniform<i8:f32, 34.0:16>>
- }
- 
- // CHECK-LABEL: "op_while"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_while(%arg0: tensor<i1>) -> tensor<i1> {
--  //      CHECK: "vhlo.while_v1"(%arg0) ({
-+  //      CHECK: "vhlo.while_v1"(%[[ARG0]]) ({
-   // CHECK-NEXT:   ^[[BB:bb.*]](%[[ARG1:arg.*]]: !vhlo.tensor_v1<!vhlo.bool_v1>):
-   // CHECK-NEXT:     "vhlo.return_v1"(%[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>) -> ()
-   // CHECK-NEXT:   }, {
-@@ -2232,8 +2396,9 @@
- }
- 
- // CHECK-LABEL: "op_xor"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_xor(%arg0: tensor<i1>, %arg1: tensor<i1>) -> tensor<i1> {
--  // CHECK: "vhlo.xor_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-+  // CHECK: "vhlo.xor_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-   %0 = "stablehlo.xor"(%arg0, %arg1) : (tensor<i1>, tensor<i1>) -> tensor<i1>
-   func.return %0 : tensor<i1>
- }
-@@ -2241,197 +2406,225 @@
- // ============ TYPES ============
- 
- // CHECK-LABEL: "type_i1"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_i1(%arg0: tensor<i1>, %arg1: tensor<i1>) -> tensor<i1> {
--  // CHECK: "vhlo.and_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-+  // CHECK: "vhlo.and_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-   %0 = "stablehlo.and"(%arg0, %arg1) : (tensor<i1>, tensor<i1>) -> tensor<i1>
-   func.return %0 : tensor<i1>
- }
- 
- // CHECK-LABEL: "type_i4"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_i4(%arg0: tensor<i4>, %arg1: tensor<i4>) -> tensor<i4> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.i4_v1>, !vhlo.tensor_v1<!vhlo.i4_v1>) -> !vhlo.tensor_v1<!vhlo.i4_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i4_v1>, !vhlo.tensor_v1<!vhlo.i4_v1>) -> !vhlo.tensor_v1<!vhlo.i4_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<i4>, tensor<i4>) -> tensor<i4>
-   func.return %0 : tensor<i4>
- }
- 
- // CHECK-LABEL: "type_i8"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_i8(%arg0: tensor<i8>, %arg1: tensor<i8>) -> tensor<i8> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.i8_v1>, !vhlo.tensor_v1<!vhlo.i8_v1>) -> !vhlo.tensor_v1<!vhlo.i8_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i8_v1>, !vhlo.tensor_v1<!vhlo.i8_v1>) -> !vhlo.tensor_v1<!vhlo.i8_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<i8>, tensor<i8>) -> tensor<i8>
-   func.return %0 : tensor<i8>
- }
- 
- // CHECK-LABEL: "type_i16"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_i16(%arg0: tensor<i16>, %arg1: tensor<i16>) -> tensor<i16> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.i16_v1>, !vhlo.tensor_v1<!vhlo.i16_v1>) -> !vhlo.tensor_v1<!vhlo.i16_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i16_v1>, !vhlo.tensor_v1<!vhlo.i16_v1>) -> !vhlo.tensor_v1<!vhlo.i16_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<i16>, tensor<i16>) -> tensor<i16>
-   func.return %0 : tensor<i16>
- }
- 
- // CHECK-LABEL: "type_i32"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_i32(%arg0: tensor<i32>, %arg1: tensor<i32>) -> tensor<i32> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.i32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<i32>, tensor<i32>) -> tensor<i32>
-   func.return %0 : tensor<i32>
- }
- 
- // CHECK-LABEL: "type_i64"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_i64(%arg0: tensor<i64>, %arg1: tensor<i64>) -> tensor<i64> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.i64_v1>, !vhlo.tensor_v1<!vhlo.i64_v1>) -> !vhlo.tensor_v1<!vhlo.i64_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i64_v1>, !vhlo.tensor_v1<!vhlo.i64_v1>) -> !vhlo.tensor_v1<!vhlo.i64_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<i64>, tensor<i64>) -> tensor<i64>
-   func.return %0 : tensor<i64>
- }
- 
- // CHECK-LABEL: "type_ui4"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_ui4(%arg0: tensor<ui4>, %arg1: tensor<ui4>) -> tensor<ui4> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.ui4_v1>, !vhlo.tensor_v1<!vhlo.ui4_v1>) -> !vhlo.tensor_v1<!vhlo.ui4_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.ui4_v1>, !vhlo.tensor_v1<!vhlo.ui4_v1>) -> !vhlo.tensor_v1<!vhlo.ui4_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<ui4>, tensor<ui4>) -> tensor<ui4>
-   func.return %0 : tensor<ui4>
- }
- 
- // CHECK-LABEL: "type_ui8"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_ui8(%arg0: tensor<ui8>, %arg1: tensor<ui8>) -> tensor<ui8> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.ui8_v1>, !vhlo.tensor_v1<!vhlo.ui8_v1>) -> !vhlo.tensor_v1<!vhlo.ui8_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.ui8_v1>, !vhlo.tensor_v1<!vhlo.ui8_v1>) -> !vhlo.tensor_v1<!vhlo.ui8_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<ui8>, tensor<ui8>) -> tensor<ui8>
-   func.return %0 : tensor<ui8>
- }
- 
- // CHECK-LABEL: "type_ui16"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_ui16(%arg0: tensor<ui16>, %arg1: tensor<ui16>) -> tensor<ui16> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.ui16_v1>, !vhlo.tensor_v1<!vhlo.ui16_v1>) -> !vhlo.tensor_v1<!vhlo.ui16_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.ui16_v1>, !vhlo.tensor_v1<!vhlo.ui16_v1>) -> !vhlo.tensor_v1<!vhlo.ui16_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<ui16>, tensor<ui16>) -> tensor<ui16>
-   func.return %0 : tensor<ui16>
- }
- 
- // CHECK-LABEL: "type_ui32"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_ui32(%arg0: tensor<ui32>, %arg1: tensor<ui32>) -> tensor<ui32> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.ui32_v1>, !vhlo.tensor_v1<!vhlo.ui32_v1>) -> !vhlo.tensor_v1<!vhlo.ui32_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.ui32_v1>, !vhlo.tensor_v1<!vhlo.ui32_v1>) -> !vhlo.tensor_v1<!vhlo.ui32_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<ui32>, tensor<ui32>) -> tensor<ui32>
-   func.return %0 : tensor<ui32>
- }
- 
- // CHECK-LABEL: "type_ui64"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_ui64(%arg0: tensor<ui64>, %arg1: tensor<ui64>) -> tensor<ui64> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.ui64_v1>, !vhlo.tensor_v1<!vhlo.ui64_v1>) -> !vhlo.tensor_v1<!vhlo.ui64_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.ui64_v1>, !vhlo.tensor_v1<!vhlo.ui64_v1>) -> !vhlo.tensor_v1<!vhlo.ui64_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<ui64>, tensor<ui64>) -> tensor<ui64>
-   func.return %0 : tensor<ui64>
- }
- 
- // CHECK-LABEL: "type_f8E4M3FN"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_f8E4M3FN(%arg0: tensor<f8E4M3FN>, %arg1: tensor<f8E4M3FN>) -> tensor<f8E4M3FN> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f8E4M3FN_v1>, !vhlo.tensor_v1<!vhlo.f8E4M3FN_v1>) -> !vhlo.tensor_v1<!vhlo.f8E4M3FN_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f8E4M3FN_v1>, !vhlo.tensor_v1<!vhlo.f8E4M3FN_v1>) -> !vhlo.tensor_v1<!vhlo.f8E4M3FN_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f8E4M3FN>, tensor<f8E4M3FN>) -> tensor<f8E4M3FN>
-   func.return %0 : tensor<f8E4M3FN>
- }
- 
- // CHECK-LABEL: "type_f8E5M2"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_f8E5M2(%arg0: tensor<f8E5M2>, %arg1: tensor<f8E5M2>) -> tensor<f8E5M2> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f8E5M2_v1>, !vhlo.tensor_v1<!vhlo.f8E5M2_v1>) -> !vhlo.tensor_v1<!vhlo.f8E5M2_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f8E5M2_v1>, !vhlo.tensor_v1<!vhlo.f8E5M2_v1>) -> !vhlo.tensor_v1<!vhlo.f8E5M2_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f8E5M2>, tensor<f8E5M2>) -> tensor<f8E5M2>
-   func.return %0 : tensor<f8E5M2>
- }
- 
- // CHECK-LABEL: "type_f8E4M3FNUZ"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_f8E4M3FNUZ(%arg0: tensor<f8E4M3FNUZ>, %arg1: tensor<f8E4M3FNUZ>) -> tensor<f8E4M3FNUZ> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f8E4M3FNUZ_v1>, !vhlo.tensor_v1<!vhlo.f8E4M3FNUZ_v1>) -> !vhlo.tensor_v1<!vhlo.f8E4M3FNUZ_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f8E4M3FNUZ_v1>, !vhlo.tensor_v1<!vhlo.f8E4M3FNUZ_v1>) -> !vhlo.tensor_v1<!vhlo.f8E4M3FNUZ_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f8E4M3FNUZ>, tensor<f8E4M3FNUZ>) -> tensor<f8E4M3FNUZ>
-   func.return %0 : tensor<f8E4M3FNUZ>
- }
- 
- // CHECK-LABEL: "type_f8E4M3B11FNUZ"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_f8E4M3B11FNUZ(%arg0: tensor<f8E4M3B11FNUZ>, %arg1: tensor<f8E4M3B11FNUZ>) -> tensor<f8E4M3B11FNUZ> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f8E4M3B11FNUZ_v1>, !vhlo.tensor_v1<!vhlo.f8E4M3B11FNUZ_v1>) -> !vhlo.tensor_v1<!vhlo.f8E4M3B11FNUZ_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f8E4M3B11FNUZ_v1>, !vhlo.tensor_v1<!vhlo.f8E4M3B11FNUZ_v1>) -> !vhlo.tensor_v1<!vhlo.f8E4M3B11FNUZ_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f8E4M3B11FNUZ>, tensor<f8E4M3B11FNUZ>) -> tensor<f8E4M3B11FNUZ>
-   func.return %0 : tensor<f8E4M3B11FNUZ>
- }
- 
- // CHECK-LABEL: "type_f8E5M2FNUZ"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_f8E5M2FNUZ(%arg0: tensor<f8E5M2FNUZ>, %arg1: tensor<f8E5M2FNUZ>) -> tensor<f8E5M2FNUZ> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f8E5M2FNUZ_v1>, !vhlo.tensor_v1<!vhlo.f8E5M2FNUZ_v1>) -> !vhlo.tensor_v1<!vhlo.f8E5M2FNUZ_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f8E5M2FNUZ_v1>, !vhlo.tensor_v1<!vhlo.f8E5M2FNUZ_v1>) -> !vhlo.tensor_v1<!vhlo.f8E5M2FNUZ_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f8E5M2FNUZ>, tensor<f8E5M2FNUZ>) -> tensor<f8E5M2FNUZ>
-   func.return %0 : tensor<f8E5M2FNUZ>
- }
- 
- // CHECK-LABEL: "type_bf16"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_bf16(%arg0: tensor<bf16>, %arg1: tensor<bf16>) -> tensor<bf16> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.bf16_v1>, !vhlo.tensor_v1<!vhlo.bf16_v1>) -> !vhlo.tensor_v1<!vhlo.bf16_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.bf16_v1>, !vhlo.tensor_v1<!vhlo.bf16_v1>) -> !vhlo.tensor_v1<!vhlo.bf16_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<bf16>, tensor<bf16>) -> tensor<bf16>
-   func.return %0 : tensor<bf16>
- }
- 
- // CHECK-LABEL: "type_f16"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_f16(%arg0: tensor<f16>, %arg1: tensor<f16>) -> tensor<f16> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f16_v1>, !vhlo.tensor_v1<!vhlo.f16_v1>) -> !vhlo.tensor_v1<!vhlo.f16_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f16_v1>, !vhlo.tensor_v1<!vhlo.f16_v1>) -> !vhlo.tensor_v1<!vhlo.f16_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f16>, tensor<f16>) -> tensor<f16>
-   func.return %0 : tensor<f16>
- }
- 
- // CHECK-LABEL: "type_f32"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_f32(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "type_f64"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_f64(%arg0: tensor<f64>, %arg1: tensor<f64>) -> tensor<f64> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f64_v1>, !vhlo.tensor_v1<!vhlo.f64_v1>) -> !vhlo.tensor_v1<!vhlo.f64_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f64_v1>, !vhlo.tensor_v1<!vhlo.f64_v1>) -> !vhlo.tensor_v1<!vhlo.f64_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f64>, tensor<f64>) -> tensor<f64>
-   func.return %0 : tensor<f64>
- }
- 
- // CHECK-LABEL: "type_complex_f32"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_complex_f32(%arg0: tensor<complex<f32>>, %arg1: tensor<complex<f32>>) -> tensor<complex<f32>> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>, !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>) -> !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>, !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>) -> !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<complex<f32>>, tensor<complex<f32>>) -> tensor<complex<f32>>
-   func.return %0 : tensor<complex<f32>>
- }
- 
- // CHECK-LABEL: "type_complex_f64"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_complex_f64(%arg0: tensor<complex<f64>>, %arg1: tensor<complex<f64>>) -> tensor<complex<f64>> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f64_v1>>, !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f64_v1>>) -> !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f64_v1>>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f64_v1>>, !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f64_v1>>) -> !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f64_v1>>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<complex<f64>>, tensor<complex<f64>>) -> tensor<complex<f64>>
-   func.return %0 : tensor<complex<f64>>
- }
- 
- // CHECK-LABEL: "type_dynamism_ranked"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @type_dynamism_ranked(%arg0: tensor<?xf32>) -> tensor<?xf32> {
--  // CHECK: "vhlo.abs_v1"(%arg0) : (!vhlo.tensor_v1<?x!vhlo.f32_v1>) -> !vhlo.tensor_v1<?x!vhlo.f32_v1>
-+  // CHECK: "vhlo.abs_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<?x!vhlo.f32_v1>) -> !vhlo.tensor_v1<?x!vhlo.f32_v1>
-   %0 = "stablehlo.abs"(%arg0) : (tensor<?xf32>) -> tensor<?xf32>
-   func.return %0 : tensor<?xf32>
- }
- 
- // CHECK-LABEL: "type_per_tensor_quantization"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_per_tensor_quantization(%arg0: tensor<!quant.uniform<i8:f32, 34.0:16>>, %arg1: tensor<!quant.uniform<i8:f32, 34.0:16>>) -> tensor<!quant.uniform<i8:f32, 34.0:16>> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>, !vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>) -> !vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>, !vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>) -> !vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<!quant.uniform<i8:f32, 34.0:16>>, tensor<!quant.uniform<i8:f32, 34.0:16>>) -> tensor<!quant.uniform<i8:f32, 34.0:16>>
-   func.return %0 : tensor<!quant.uniform<i8:f32, 34.0:16>>
- }
- 
- // CHECK-LABEL: "type_per_axis_quantization"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @type_per_axis_quantization(%arg0: tensor<2x!quant.uniform<i8:f32:0, {34.0:16, 34.0:16}>>) -> tensor<2x!quant.uniform<i8:f32:0, {34.0:16, 34.0:16}>> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg0) : (!vhlo.tensor_v1<2x!vhlo.quant_per_axis_v1<!vhlo.i8_v1:!vhlo.f32_v1, 0, 3.400000e+01, 3.400000e+01, 16, 16, -128:127, 1>>, !vhlo.tensor_v1<2x!vhlo.quant_per_axis_v1<!vhlo.i8_v1:!vhlo.f32_v1, 0, 3.400000e+01, 3.400000e+01, 16, 16, -128:127, 1>>) -> !vhlo.tensor_v1<2x!vhlo.quant_per_axis_v1<!vhlo.i8_v1:!vhlo.f32_v1, 0, 3.400000e+01, 3.400000e+01, 16, 16, -128:127, 1>>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG0]]) : (!vhlo.tensor_v1<2x!vhlo.quant_per_axis_v1<!vhlo.i8_v1:!vhlo.f32_v1, 0, 3.400000e+01, 3.400000e+01, 16, 16, -128:127, 1>>, !vhlo.tensor_v1<2x!vhlo.quant_per_axis_v1<!vhlo.i8_v1:!vhlo.f32_v1, 0, 3.400000e+01, 3.400000e+01, 16, 16, -128:127, 1>>) -> !vhlo.tensor_v1<2x!vhlo.quant_per_axis_v1<!vhlo.i8_v1:!vhlo.f32_v1, 0, 3.400000e+01, 3.400000e+01, 16, 16, -128:127, 1>>
-   %0 = stablehlo.add %arg0, %arg0 : tensor<2x!quant.uniform<i8:f32:0, {34.0:16, 34.0:16}>>
-   func.return %0 : tensor<2x!quant.uniform<i8:f32:0, {34.0:16, 34.0:16}>>
- }
- 
- //       CHECK: function_type = #vhlo.type_v1<!vhlo.func_v1<(!vhlo.token_v1) -> !vhlo.token_v1>>
- // CHECK-LABEL: "type_token_callee"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @type_token_callee(%arg0: !stablehlo.token) -> !stablehlo.token {
--  // CHECK: "vhlo.return_v1"(%arg0) : (!vhlo.token_v1) -> ()
-+  // CHECK: "vhlo.return_v1"(%[[ARG0]]) : (!vhlo.token_v1) -> ()
-   return %arg0 : !stablehlo.token
- }
- 
- //       CHECK: function_type = #vhlo.type_v1<!vhlo.func_v1<(!vhlo.token_v1) -> !vhlo.token_v1>>
- // CHECK-LABEL: "type_token_caller"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @type_token_caller(%arg0: !stablehlo.token) -> !stablehlo.token {
--  // CHECK:      "vhlo.call_v1"(%arg0) <{callee = #vhlo.string_v1<"type_token_callee">}
-+  // CHECK:      "vhlo.call_v1"(%[[ARG0]]) <{callee = #vhlo.string_v1<"type_token_callee">}
-   // CHECK-SAME: (!vhlo.token_v1) -> !vhlo.token_v1
-   %0 = func.call @type_token_callee(%arg0) : (!stablehlo.token) -> !stablehlo.token
-   return %0 : !stablehlo.token
- }
- 
- // CHECK-LABEL: "type_tuple"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @type_tuple(%arg0: tuple<tensor<f32>>) -> tuple<!stablehlo.token> {
-   %0 = "stablehlo.custom_call"(%arg0) {
-     call_target_name = "foo"
-diff --ruN a/stablehlo/stablehlo/tests/vhlo/stablehlo_legalize_to_vhlo.0_9_0.mlir b/stablehlo/stablehlo/tests/vhlo/stablehlo_legalize_to_vhlo.0_9_0.mlir
---- stablehlo/stablehlo/tests/vhlo/stablehlo_legalize_to_vhlo.0_9_0.mlir
-+++ stablehlo/stablehlo/tests/vhlo/stablehlo_legalize_to_vhlo.0_9_0.mlir
-@@ -13,6 +13,7 @@
- // ============ ATTRIBUTES ============
- 
- // CHECK-LABEL: "attr_comparison_direction_eq"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_direction_eq(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     // CHECK: comparison_direction = #vhlo<comparison_direction_v1 EQ>
-@@ -22,6 +23,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_direction_ne"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_direction_ne(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     // CHECK: comparison_direction = #vhlo<comparison_direction_v1 NE>
-@@ -31,6 +33,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_direction_ge"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_direction_ge(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     // CHECK: comparison_direction = #vhlo<comparison_direction_v1 GE>
-@@ -40,6 +43,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_direction_gt"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_direction_gt(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     // CHECK: comparison_direction = #vhlo<comparison_direction_v1 GT>
-@@ -49,6 +53,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_direction_le"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_direction_le(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     // CHECK: comparison_direction = #vhlo<comparison_direction_v1 LE>
-@@ -58,6 +63,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_direction_lt"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_direction_lt(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     // CHECK: comparison_direction = #vhlo<comparison_direction_v1 LT>
-@@ -67,6 +73,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_type_notype"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_type_notype(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     comparison_direction = #stablehlo<comparison_direction EQ>
-@@ -76,6 +83,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_type_float"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_type_float(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     comparison_direction = #stablehlo<comparison_direction EQ>,
-@@ -86,6 +94,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_type_totalorder"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_type_totalorder(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     comparison_direction = #stablehlo<comparison_direction EQ>,
-@@ -96,6 +105,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_type_signed"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_type_signed(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     comparison_direction = #stablehlo<comparison_direction EQ>,
-@@ -106,6 +116,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_type_unsigned"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_type_unsigned(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     comparison_direction = #stablehlo<comparison_direction EQ>,
-@@ -118,6 +129,7 @@
- // ConvDimensionNumbers aka #stablehlo.conv is covered below.
- 
- // CHECK-LABEL: "attr_custom_call_api_version_unspecified"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_custom_call_api_version_unspecified(%arg0: tensor<f32>) -> tensor<f32> {
-   %0 = "stablehlo.custom_call"(%arg0) {
-     call_target_name = "foo",
-@@ -128,6 +140,7 @@
- }
- 
- // CHECK-LABEL: "attr_custom_call_api_version_original"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_custom_call_api_version_original(%arg0: tensor<f32>) -> tensor<f32> {
-   %0 = "stablehlo.custom_call"(%arg0) {
-     call_target_name = "foo",
-@@ -138,6 +151,7 @@
- }
- 
- // CHECK-LABEL: "attr_custom_call_api_version_status_returning"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_custom_call_api_version_status_returning(%arg0: tensor<f32>) -> tensor<f32> {
-   %0 = "stablehlo.custom_call"(%arg0) {
-     call_target_name = "foo",
-@@ -148,6 +162,7 @@
- }
- 
- // CHECK-LABEL: "attr_custom_call_api_version_status_returning_unified"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_custom_call_api_version_status_returning_unified(%arg0: tensor<f32>) -> tensor<f32> {
-   %0 = "stablehlo.custom_call"(%arg0) {
-     call_target_name = "foo",
-@@ -160,6 +175,7 @@
- // DotDimensionNumbers aka #stablehlo.dot is covered below.
- 
- // CHECK-LABEL: "attr_fft_type_fft"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_fft_type_fft(%arg0: tensor<16xcomplex<f32>>) -> tensor<16xcomplex<f32>> {
-   %0 = "stablehlo.fft"(%arg0) {
-     // CHECK: fft_type = #vhlo<fft_type_v1 FFT>
-@@ -170,6 +186,7 @@
- }
- 
- // CHECK-LABEL: "attr_fft_type_ifft"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_fft_type_ifft(%arg0: tensor<16xcomplex<f32>>) -> tensor<16xcomplex<f32>> {
-   %0 = "stablehlo.fft"(%arg0) {
-     // CHECK: fft_type = #vhlo<fft_type_v1 IFFT>
-@@ -180,6 +197,7 @@
- }
- 
- // CHECK-LABEL: "attr_fft_type_rfft"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_fft_type_rfft(%arg0: tensor<16xf32>) -> tensor<9xcomplex<f32>> {
-   %0 = "stablehlo.fft"(%arg0) {
-     // CHECK: fft_type = #vhlo<fft_type_v1 RFFT>
-@@ -190,6 +208,7 @@
- }
- 
- // CHECK-LABEL: "attr_fft_type_irfft"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_fft_type_irfft(%arg0: tensor<9xcomplex<f32>>) -> tensor<16xf32> {
-   %0 = "stablehlo.fft"(%arg0) {
-     // CHECK: fft_type = #vhlo<fft_type_v1 IRFFT>
-@@ -202,6 +221,7 @@
- // GatherDimensionNumbers aka #stablehlo.gather is covered below.
- 
- // CHECK-LABEL: "attr_precision_config_default"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_precision_config_default(%arg0: tensor<8x16xf32>, %arg1: tensor<16x8xf32>) -> tensor<8x8xf32> {
-   %0 = "stablehlo.dot"(%arg0, %arg1) {
-     // CHECK: precision_config = #vhlo.array_v1<[#vhlo<precision_v1 DEFAULT>, #vhlo<precision_v1 DEFAULT>]>
-@@ -210,6 +230,7 @@
- }
- 
- // CHECK-LABEL: "attr_precision_config_high"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_precision_config_high(%arg0: tensor<8x16xf32>, %arg1: tensor<16x8xf32>) -> tensor<8x8xf32> {
-   %0 = "stablehlo.dot"(%arg0, %arg1) {
-     // CHECK: precision_config = #vhlo.array_v1<[#vhlo<precision_v1 HIGH>, #vhlo<precision_v1 HIGH>]>
-@@ -219,6 +240,7 @@
- }
- 
- // CHECK-LABEL: "attr_precision_config_highest"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_precision_config_highest(%arg0: tensor<8x16xf32>, %arg1: tensor<16x8xf32>) -> tensor<8x8xf32> {
-   %0 = "stablehlo.dot"(%arg0, %arg1) {
-     // CHECK: precision_config = #vhlo.array_v1<[#vhlo<precision_v1 HIGHEST>, #vhlo<precision_v1 HIGHEST>]>
-@@ -228,6 +250,7 @@
- }
- 
- // CHECK-LABEL: "attr_rng_algorithm_default"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_rng_algorithm_default(%arg0: tensor<f32>) -> (tensor<f32>, tensor<f32>) {
-   %0:2 = "stablehlo.rng_bit_generator"(%arg0) {
-     // CHECK: rng_algorithm = #vhlo<rng_algorithm_v1 DEFAULT>
-@@ -237,6 +260,7 @@
- }
- 
- // CHECK-LABEL: "attr_rng_algorithm_three_fry"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_rng_algorithm_three_fry(%arg0: tensor<f32>) -> (tensor<f32>, tensor<f32>) {
-   %0:2 = "stablehlo.rng_bit_generator"(%arg0) {
-     // CHECK: rng_algorithm = #vhlo<rng_algorithm_v1 THREE_FRY>
-@@ -246,6 +270,7 @@
- }
- 
- // CHECK-LABEL: "attr_rng_algorithm_philox"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_rng_algorithm_philox(%arg0: tensor<f32>) -> (tensor<f32>, tensor<f32>) {
-   %0:2 = "stablehlo.rng_bit_generator"(%arg0) {
-     // CHECK: rng_algorithm = #vhlo<rng_algorithm_v1 PHILOX>
-@@ -255,6 +280,7 @@
- }
- 
- // CHECK-LABEL: "attr_rng_distribution_uniform"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @attr_rng_distribution_uniform(%arg0: tensor<f32>, %arg1: tensor<f32>, %arg2: tensor<0xindex>) -> tensor<f32> {
-   %0 = "stablehlo.rng"(%arg0, %arg1, %arg2) {
-     // CHECK: rng_distribution = #vhlo<rng_distribution_v1 UNIFORM>
-@@ -264,6 +290,7 @@
- }
- 
- // CHECK-LABEL: "attr_rng_distribution_normal"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @attr_rng_distribution_normal(%arg0: tensor<f32>, %arg1: tensor<f32>, %arg2: tensor<0xindex>) -> tensor<f32> {
-   %0 = "stablehlo.rng"(%arg0, %arg1, %arg2) {
-     // CHECK: rng_distribution = #vhlo<rng_distribution_v1 NORMAL>
-@@ -275,6 +302,7 @@
- // ScatterDimensionNumbers aka #stablehlo.scatter is covered below.
- 
- // CHECK-LABEL: "attr_transpose_no_transpose"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_transpose_no_transpose(%arg0: tensor<16x16xf32>, %arg1: tensor<16x16xf32>) ->  tensor<16x16xf32> {
-   %0 = "stablehlo.triangular_solve"(%arg0, %arg1) {
-     left_side = true,
-@@ -287,6 +315,7 @@
- }
- 
- // CHECK-LABEL: "attr_transpose_transpose"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_transpose_transpose(%arg0: tensor<16x16xf32>, %arg1: tensor<16x16xf32>) ->  tensor<16x16xf32> {
-   %0 = "stablehlo.triangular_solve"(%arg0, %arg1) {
-     left_side = true,
-@@ -299,6 +328,7 @@
- }
- 
- // CHECK-LABEL: "attr_transpose_adjoint"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_transpose_adjoint(%arg0: tensor<16x16xf32>, %arg1: tensor<16x16xf32>) ->  tensor<16x16xf32> {
-   %0 = "stablehlo.triangular_solve"(%arg0, %arg1) {
-     left_side = true,
-@@ -313,10 +343,9 @@
- // TypeExtensionsAttr aka #stablehlo.type_extensions is covered below.
- 
- // CHECK-LABEL: "attr_type_extensions_bounds"
--func.func @attr_type_extensions_bounds(
--    %arg0: tensor<?x?xf32, #stablehlo.type_extensions<bounds = [16, ?]>>)
--    -> tensor<?x?xf32, #stablehlo.type_extensions<bounds = [16, ?]>> {
--  // CHECK: "vhlo.return_v1"(%arg0) : (!vhlo.tensor_v1<?x?x!vhlo.f32_v1, #vhlo.type_extensions_v1<bounds = [16, ?]>>) -> ()
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @attr_type_extensions_bounds(%arg0: tensor<?x?xf32, #stablehlo.type_extensions<bounds = [16, ?]>>) -> tensor<?x?xf32, #stablehlo.type_extensions<bounds = [16, ?]>> {
-+  // CHECK: "vhlo.return_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<?x?x!vhlo.f32_v1, #vhlo.type_extensions_v1<bounds = [16, ?]>>) -> ()
-   func.return %arg0 : tensor<?x?xf32, #stablehlo.type_extensions<bounds = [16, ?]>>
- }
- 
-@@ -324,8 +353,9 @@
- // ============ DEFAULTS ============
- 
- // CHECK-LABEL: "default_all_gather"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_all_gather(%arg0: tensor<16x8xf32>) -> tensor<16x16xf32> {
--  //               CHECK: "vhlo.all_gather_v1"(%arg0) <{
-+  //               CHECK: "vhlo.all_gather_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   all_gather_dim = #vhlo.integer_v1<1 : i64>
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0], [1]]> : tensor<2x1xi64>>,
-@@ -339,8 +369,9 @@
- }
- 
- // CHECK-LABEL: "default_all_reduce"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_all_reduce(%arg0: tensor<f32>) -> tensor<f32> {
--  //               CHECK: "vhlo.all_reduce_v1"(%arg0)
-+  //               CHECK: "vhlo.all_reduce_v1"(%[[ARG0]])
-   //          CHECK-SAME: <{
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0], [1]]> : tensor<2x1xi64>>,
-@@ -362,8 +393,9 @@
- }
- 
- // CHECK-LABEL: "default_all_to_all"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_all_to_all(%arg0: tensor<4x16xf32>) -> tensor<16x4xf32> {
--  //               CHECK: "vhlo.all_to_all_v1"(%arg0) <{
-+  //               CHECK: "vhlo.all_to_all_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   //          CHECK-SAME:   concat_dimension = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0, 1, 2, 3]]> : tensor<1x4xi64>>,
-@@ -380,8 +412,9 @@
- }
- 
- // CHECK-LABEL: "default_cholesky"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_cholesky(%arg0: tensor<1x16x16xf32>) -> tensor<1x16x16xf32> {
--  //      CHECK: "vhlo.cholesky_v1"(%arg0) <{
-+  //      CHECK: "vhlo.cholesky_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   lower = #vhlo.bool_v1<false>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<1x16x16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<1x16x16x!vhlo.f32_v1>
-   %0 = "stablehlo.cholesky"(%arg0) : (tensor<1x16x16xf32>) -> tensor<1x16x16xf32>
-@@ -389,8 +422,9 @@
- }
- 
- // CHECK-LABEL: "default_collective_permute"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_collective_permute(%arg0: tensor<16x8xf32>) -> tensor<16x8xf32> {
--  //               CHECK: "vhlo.collective_permute_v1"(%arg0) <{
-+  //               CHECK: "vhlo.collective_permute_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME{LITERAL}:   source_target_pairs = #vhlo.tensor_v1<dense<[[0, 1], [1, 2], [2, 3]]> : tensor<3x2xi64>>
-   //          CHECK-SAME: }> : (!vhlo.tensor_v1<16x8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x8x!vhlo.f32_v1>
-@@ -401,8 +435,9 @@
- }
- 
- // CHECK-LABEL: "default_compare"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @default_compare(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
--  //      CHECK: "vhlo.compare_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.compare_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   compare_type = #vhlo<comparison_type_v1 NOTYPE>,
-   // CHECK-SAME:   comparison_direction = #vhlo<comparison_direction_v1 EQ>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-@@ -413,8 +448,9 @@
- }
- 
- // CHECK-LABEL: "default_convolution"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @default_convolution(%arg0: tensor<1x8x8x207xf32>, %arg1: tensor<3x3x207x16xf32>) -> tensor<1x6x6x16xf32> {
--  //      CHECK: "vhlo.convolution_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.convolution_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   batch_group_count = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   feature_group_count = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   input_batch_dimension = #vhlo.integer_v1<0 : i64>,
-@@ -442,8 +478,9 @@
- }
- 
- // CHECK-LABEL: "default_custom_call"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_custom_call(%arg0: tensor<f32>) -> tensor<f32> {
--  //      CHECK: "vhlo.custom_call_v1"(%arg0) <{
-+  //      CHECK: "vhlo.custom_call_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   api_version = #vhlo<api_version_v1 API_VERSION_ORIGINAL>,
-   // CHECK-SAME:   backend_config = #vhlo.string_v1<"">,
-   // CHECK-SAME:   call_target_name = #vhlo.string_v1<"foo">,
-@@ -460,8 +497,9 @@
- }
- 
- // CHECK-LABEL: "default_dot_general"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @default_dot_general(%arg0: tensor<8x8x16xf32>, %arg1: tensor<8x16x8xf32>) -> tensor<8x8x8xf32> {
--  //      CHECK: "vhlo.dot_general_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.dot_general_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   lhs_batching_dimensions = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>,
-   // CHECK-SAME:   lhs_contracting_dimensions = #vhlo.tensor_v1<dense<2> : tensor<1xi64>>,
-   // CHECK-SAME:   precision_config = #vhlo.array_v1<[#vhlo<precision_v1 DEFAULT>, #vhlo<precision_v1 DEFAULT>]>,
-@@ -480,8 +518,9 @@
- }
- 
- // CHECK-LABEL: "default_dot"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @default_dot(%arg0: tensor<8x16xf32>, %arg1: tensor<16x8xf32>) -> tensor<8x8xf32> {
--  //      CHECK: "vhlo.dot_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.dot_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   precision_config = #vhlo.array_v1<[#vhlo<precision_v1 DEFAULT>, #vhlo<precision_v1 DEFAULT>]>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<8x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<8x8x!vhlo.f32_v1>
-   %0 = "stablehlo.dot"(%arg0, %arg1) : (tensor<8x16xf32>, tensor<16x8xf32>) -> tensor<8x8xf32>
-@@ -489,8 +528,9 @@
- }
- 
- // CHECK-LABEL: "default_dynamic_broadcast_in_dim"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @default_dynamic_broadcast_in_dim(%arg0: tensor<?x?xf32>, %arg1: tensor<2xindex>) -> tensor<?x?xf32> {
--  //      CHECK: "vhlo.dynamic_broadcast_in_dim_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.dynamic_broadcast_in_dim_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   broadcast_dimensions = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-   // CHECK-SAME:   known_expanding_dimensions = #vhlo.tensor_v1<dense<> : tensor<0xi64>>,
-   // CHECK-SAME:   known_nonexpanding_dimensions = #vhlo.tensor_v1<dense<> : tensor<0xi64>>
-@@ -502,8 +542,9 @@
- }
- 
- // CHECK-LABEL: "default_dynamic_conv"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @default_dynamic_conv(%arg0: tensor<1x8x8x207xf32>, %arg1: tensor<3x3x207x16xf32>, %arg2: tensor<4xi32>) -> tensor<1x?x?x16xf32> {
--  //      CHECK: "vhlo.dynamic_conv_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.dynamic_conv_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   batch_group_count = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   feature_group_count = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   input_batch_dimension = #vhlo.integer_v1<0 : i64>,
-@@ -531,8 +572,9 @@
- }
- 
- // CHECK-LABEL: "default_dynamic_gather"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @default_dynamic_gather(%arg0 : tensor<2x4x9xf32>, %arg1 : tensor<1x5x2xi32>, %arg2 : tensor<3xi32>) -> tensor<1x5x8xf32> {
--  //      CHECK: "vhlo.dynamic_gather_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.dynamic_gather_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   collapsed_slice_dims = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-   // CHECK-SAME:   index_vector_dim = #vhlo.integer_v1<2 : i64>,
-   // CHECK-SAME:   indices_are_sorted = #vhlo.bool_v1<false>,
-@@ -558,15 +600,16 @@
-   // CHECK-SAME:   sym_name = #vhlo.string_v1<"default_func">,
-   // CHECK-SAME:   sym_visibility = #vhlo.string_v1<"">
-   // CHECK-SAME: }> ({
--  // CHECK-NEXT:   ^[[BB:bb.*]](%arg0: !vhlo.tensor_v1<!vhlo.f32_v1>):
--  // CHECK-NEXT:     "vhlo.return_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-+  // CHECK-NEXT:   ^[[BB:bb.*]](%[[ARG0:.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
-+  // CHECK-NEXT:     "vhlo.return_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-   // CHECK-NEXT: }) : () -> ()
-   func.return %arg0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "dynamic_gather"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @dynamic_gather(%arg0 : tensor<2x4x9xf32>, %arg1 : tensor<1x5x2xi32>) -> tensor<1x5x1xf32> {
--  //      CHECK: "vhlo.gather_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.gather_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   collapsed_slice_dims = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-   // CHECK-SAME:   index_vector_dim = #vhlo.integer_v1<2 : i64>,
-   // CHECK-SAME:   indices_are_sorted = #vhlo.bool_v1<false>,
-@@ -587,8 +630,9 @@
- }
- 
- // CHECK-LABEL: "default_infeed"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_infeed(%arg0: !stablehlo.token) -> (tensor<f32>, !stablehlo.token) {
--  //               CHECK: "vhlo.infeed_v1"(%arg0) <{
-+  //               CHECK: "vhlo.infeed_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   infeed_config = #vhlo.string_v1<"">,
-   // CHECK-SAME{LITERAL}:   layout = #vhlo.array_v1<[]>
-   //          CHECK-SAME: }> : (!vhlo.token_v1) -> (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.token_v1)
-@@ -597,8 +641,9 @@
- }
- 
- // CHECK-LABEL: "default_outfeed"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @default_outfeed(%arg0: tensor<f32>, %arg1: !stablehlo.token) -> !stablehlo.token {
--  //      CHECK: "vhlo.outfeed_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.outfeed_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   outfeed_config = #vhlo.string_v1<"">
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.token_v1) -> !vhlo.token_v1
-   %0 = "stablehlo.outfeed"(%arg0, %arg1) : (tensor<f32>, !stablehlo.token) -> !stablehlo.token
-@@ -606,8 +651,9 @@
- }
- 
- // CHECK-LABEL: "default_recv"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_recv(%arg0: !stablehlo.token) -> (tensor<f32>, !stablehlo.token) {
--  //      CHECK: "vhlo.recv_v1"(%arg0) <{
-+  //      CHECK: "vhlo.recv_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME:   channel_type = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   is_host_transfer = #vhlo.bool_v1<false>
-@@ -619,8 +665,9 @@
- }
- 
- // CHECK-LABEL: "default_send"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @default_send(%arg0: tensor<f32>, %arg1: !stablehlo.token) -> !stablehlo.token {
--  //      CHECK: "vhlo.send_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.send_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME:   channel_type = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   is_host_transfer = #vhlo.bool_v1<false>
-@@ -632,8 +679,9 @@
- }
- 
- // CHECK-LABEL: "default_reduce_scatter"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_reduce_scatter(%arg0: tensor<16xf32>) -> tensor<16xf32> {
--  //               CHECK: "vhlo.reduce_scatter_v1"(%arg0) <{
-+  //               CHECK: "vhlo.reduce_scatter_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0], [1]]> : tensor<2x1xi64>>,
-   //          CHECK-SAME:   scatter_dimension = #vhlo.integer_v1<0 : i64>
-@@ -655,8 +703,9 @@
- }
- 
- // CHECK-LABEL: "default_reduce_window"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @default_reduce_window(%arg0: tensor<2x17x31x7xf32>, %arg1: tensor<f32>) -> tensor<2x16x30x7xf32> {
--  //               CHECK: "vhlo.reduce_window_v1"(%arg0, %arg1)  <{
-+  //               CHECK: "vhlo.reduce_window_v1"(%[[ARG0]], %[[ARG1]])  <{
-   //          CHECK-SAME:   base_dilations = #vhlo.tensor_v1<dense<1> : tensor<4xi64>>,
-   // CHECK-SAME{LITERAL}:   padding = #vhlo.tensor_v1<dense<0> : tensor<4x2xi64>>,
-   //          CHECK-SAME:   window_dilations = #vhlo.tensor_v1<dense<1> : tensor<4xi64>>,
-@@ -678,8 +727,9 @@
- }
- 
- // CHECK-LABEL: "default_scatter"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @default_scatter(%arg0: tensor<200x100x300xf32>, %arg1: tensor<10x2xi32>, %arg2: tensor<10x300xf32>) -> tensor<200x100x300xf32> {
--  //      CHECK: "vhlo.scatter_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.scatter_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   index_vector_dim = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   indices_are_sorted = #vhlo.bool_v1<false>,
-   // CHECK-SAME:   inserted_window_dims = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-@@ -707,8 +757,9 @@
- }
- 
- // CHECK-LABEL: "default_select_and_scatter"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @default_select_and_scatter(%arg0: tensor<10x24x24x64xf32>, %arg1: tensor<10x23x23x64xf32>, %arg2: tensor<f32>) -> tensor<10x24x24x64xf32> {
--  //      CHECK: "vhlo.select_and_scatter_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.select_and_scatter_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   padding = #vhlo.tensor_v1<dense<0> : tensor<4x2xi64>>,
-   // CHECK-SAME:   window_dimensions = #vhlo.tensor_v1<dense<[1, 2, 2, 1]> : tensor<4xi64>>,
-   // CHECK-SAME:   window_strides = #vhlo.tensor_v1<dense<1> : tensor<4xi64>>
-@@ -736,8 +787,9 @@
- }
- 
- // CHECK-LABEL: "default_sort"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_sort(%arg0: tensor<16xf32>) -> tensor<16xf32> {
--  //      CHECK: "vhlo.sort_v1"(%arg0) <{
-+  //      CHECK: "vhlo.sort_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   dimension = #vhlo.integer_v1<-1 : i64>
-   // CHECK-SAME:   is_stable = #vhlo.bool_v1<false>
-   // CHECK-SAME: }> ({
-@@ -756,29 +808,33 @@
- // ============ OPS ============
- 
- // CHECK-LABEL: "op_abs"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_abs(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.abs_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.abs_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.abs"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_add"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_add(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_after_all"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_after_all(%arg0: !stablehlo.token) -> !stablehlo.token {
--  // CHECK: "vhlo.after_all_v1"(%arg0) : (!vhlo.token_v1) -> !vhlo.token_v1
-+  // CHECK: "vhlo.after_all_v1"(%[[ARG0]]) : (!vhlo.token_v1) -> !vhlo.token_v1
-   %0 = "stablehlo.after_all"(%arg0) : (!stablehlo.token) -> !stablehlo.token
-   func.return %0 : !stablehlo.token
- }
- 
- // CHECK-LABEL: "op_all_gather"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_all_gather(%arg0: tensor<16x8xf32>) -> tensor<16x16xf32> {
--  //               CHECK: "vhlo.all_gather_v1"(%arg0) <{
-+  //               CHECK: "vhlo.all_gather_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   all_gather_dim = #vhlo.integer_v1<1 : i64>
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0], [1]]> : tensor<2x1xi64>>,
-@@ -794,8 +850,9 @@
- }
- 
- // CHECK-LABEL: "op_all_reduce"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_all_reduce(%arg0: tensor<f32>) -> tensor<f32> {
--  //               CHECK: "vhlo.all_reduce_v1"(%arg0) <{
-+  //               CHECK: "vhlo.all_reduce_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0], [1]]> : tensor<2x1xi64>>,
-   //          CHECK-SAME:   use_global_device_ids = #vhlo.bool_v1<true>
-@@ -817,8 +874,9 @@
- }
- 
- // CHECK-LABEL: "op_all_to_all"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_all_to_all(%arg0: tensor<4x16xf32>) -> tensor<16x4xf32> {
--  //               CHECK: "vhlo.all_to_all_v1"(%arg0) <{
-+  //               CHECK: "vhlo.all_to_all_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<1 : i64>,
-   //          CHECK-SAME:   concat_dimension = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0, 1, 2, 3]]> : tensor<1x4xi64>>,
-@@ -836,22 +894,25 @@
- }
- 
- // CHECK-LABEL: "op_and"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_and(%arg0: tensor<i1>, %arg1: tensor<i1>) -> tensor<i1> {
--  // CHECK: "vhlo.and_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-+  // CHECK: "vhlo.and_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-   %0 = "stablehlo.and"(%arg0, %arg1) : (tensor<i1>, tensor<i1>) -> tensor<i1>
-   func.return %0 : tensor<i1>
- }
- 
- // CHECK-LABEL: "op_atan2"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_atan2(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.atan2_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.atan2_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.atan2"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_batch_norm_grad"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}}, %[[ARG3:.*]]: {{.*}}, %[[ARG4:.*]]: {{.*}})
- func.func @op_batch_norm_grad(%arg0: tensor<16x16x16x16xf32>, %arg1: tensor<16xf32>, %arg2: tensor<16xf32>, %arg3: tensor<16xf32>, %arg4: tensor<16x16x16x16xf32>) -> (tensor<16x16x16x16xf32>, tensor<16xf32>, tensor<16xf32>) {
--  //      CHECK: "vhlo.batch_norm_grad_v1"(%arg0, %arg1, %arg2, %arg3, %arg4) <{
-+  //      CHECK: "vhlo.batch_norm_grad_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]], %[[ARG3]], %[[ARG4]]) <{
-   // CHECK-SAME:   epsilon = #vhlo.float_v1<1.000000e-03 : !vhlo.f32_v1>,
-   // CHECK-SAME:   feature_index = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<16x16x16x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x16x16x16x!vhlo.f32_v1>) -> (!vhlo.tensor_v1<16x16x16x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>)
-@@ -863,8 +924,9 @@
- }
- 
- // CHECK-LABEL: "op_batch_norm_inference"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}}, %[[ARG3:.*]]: {{.*}}, %[[ARG4:.*]]: {{.*}})
- func.func @op_batch_norm_inference(%arg0: tensor<16x16x16x16xf32>, %arg1: tensor<16xf32>, %arg2: tensor<16xf32>, %arg3: tensor<16xf32>, %arg4: tensor<16xf32>) -> tensor<16x16x16x16xf32> {
--  //      CHECK: "vhlo.batch_norm_inference_v1"(%arg0, %arg1, %arg2, %arg3, %arg4) <{
-+  //      CHECK: "vhlo.batch_norm_inference_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]], %[[ARG3]], %[[ARG4]]) <{
-   // CHECK-SAME:   epsilon = #vhlo.float_v1<1.000000e-03 : !vhlo.f32_v1>,
-   // CHECK-SAME:   feature_index = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<16x16x16x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x16x16x16x!vhlo.f32_v1>
-@@ -876,8 +938,9 @@
- }
- 
- // CHECK-LABEL: "op_batch_norm_training"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_batch_norm_training(%arg0: tensor<16x16x16x16xf32>, %arg1: tensor<16xf32>, %arg2: tensor<16xf32>) -> (tensor<16x16x16x16xf32>, tensor<16xf32>, tensor<16xf32>) {
--  //      CHECK: "vhlo.batch_norm_training_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.batch_norm_training_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   epsilon = #vhlo.float_v1<1.000000e-03 : !vhlo.f32_v1>,
-   // CHECK-SAME:   feature_index = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<16x16x16x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>) -> (!vhlo.tensor_v1<16x16x16x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>)
-@@ -889,15 +952,17 @@
- }
- 
- // CHECK-LABEL: "op_bitcast_convert"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_bitcast_convert(%arg0: tensor<i32>) -> tensor<f32> {
--  // CHECK: "vhlo.bitcast_convert_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.bitcast_convert_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.bitcast_convert"(%arg0) : (tensor<i32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_broadcast_in_dim"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_broadcast_in_dim(%arg0: tensor<16xf32>) -> tensor<16x16xf32> {
--  //      CHECK: "vhlo.broadcast_in_dim_v1"(%arg0) <{
-+  //      CHECK: "vhlo.broadcast_in_dim_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   broadcast_dimensions = #vhlo.tensor_v1<dense<1> : tensor<1xi64>>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x16x!vhlo.f32_v1>
-   %0 = "stablehlo.broadcast_in_dim"(%arg0) {
-@@ -907,8 +972,9 @@
- }
- 
- // CHECK-LABEL: "op_broadcast"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_broadcast(%arg0: tensor<16xf32>) -> tensor<16x16xf32> {
--  //      CHECK: "vhlo.broadcast_v1"(%arg0) <{
-+  //      CHECK: "vhlo.broadcast_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   broadcast_sizes = #vhlo.tensor_v1<dense<16> : tensor<1xi64>>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x16x!vhlo.f32_v1>
-   %0 = "stablehlo.broadcast"(%arg0) {
-@@ -918,9 +984,10 @@
- }
- 
- // CHECK-LABEL: "op_case"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_case(%arg0: tensor<i32>, %arg1: tensor<f32>) -> tensor<f32> {
--  //      CHECK: "vhlo.case_v1"(%arg0) ({
--  // CHECK-NEXT:   "vhlo.return_v1"(%arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-+  //      CHECK: "vhlo.case_v1"(%[[ARG0]]) ({
-+  // CHECK-NEXT:   "vhlo.return_v1"(%[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-   // CHECK-NEXT: }) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.case"(%arg0) ({
-     "stablehlo.return"(%arg1) : (tensor<f32>) -> ()
-@@ -929,22 +996,25 @@
- }
- 
- // CHECK-LABEL: "op_cbrt"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_cbrt(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.cbrt_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.cbrt_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.cbrt"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_ceil"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_ceil(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.ceil_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.ceil_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.ceil"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_cholesky"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_cholesky(%arg0: tensor<1x16x16xf32>) -> tensor<1x16x16xf32> {
--  //      CHECK: "vhlo.cholesky_v1"(%arg0) <{
-+  //      CHECK: "vhlo.cholesky_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   lower = #vhlo.bool_v1<true>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<1x16x16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<1x16x16x!vhlo.f32_v1>
-   %0 = "stablehlo.cholesky"(%arg0) {
-@@ -954,22 +1024,25 @@
- }
- 
- // CHECK-LABEL: "op_clamp"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_clamp(%arg0: tensor<f32>, %arg1: tensor<f32>, %arg2: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.clamp_v1"(%arg0, %arg1, %arg2) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.clamp_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.clamp"(%arg0, %arg1, %arg2) : (tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_count_leading_zeros"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_count_leading_zeros(%arg0: tensor<i32>) -> tensor<i32> {
--  // CHECK: "vhlo.count_leading_zeros_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-+  // CHECK: "vhlo.count_leading_zeros_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-   %0 = "stablehlo.count_leading_zeros"(%arg0) : (tensor<i32>) -> tensor<i32>
-   func.return %0 : tensor<i32>
- }
- 
- // CHECK-LABEL: "op_collective_permute"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_collective_permute(%arg0: tensor<16x8xf32>) -> tensor<16x8xf32> {
--  //               CHECK: "vhlo.collective_permute_v1"(%arg0) <{
-+  //               CHECK: "vhlo.collective_permute_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME{LITERAL}:   source_target_pairs = #vhlo.tensor_v1<dense<[[0, 1], [1, 2], [2, 3]]> : tensor<3x2xi64>>
-   //          CHECK-SAME: }> : (!vhlo.tensor_v1<16x8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x8x!vhlo.f32_v1>
-@@ -981,8 +1054,9 @@
- }
- 
- // CHECK-LABEL: "op_compare"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_compare(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
--  //      CHECK: "vhlo.compare_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.compare_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   compare_type = #vhlo<comparison_type_v1 TOTALORDER>,
-   // CHECK-SAME:   comparison_direction = #vhlo<comparison_direction_v1 EQ>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-@@ -994,15 +1068,17 @@
- }
- 
- // CHECK-LABEL: "op_complex"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_complex(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<complex<f32>> {
--  // CHECK: "vhlo.complex_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>
-+  // CHECK: "vhlo.complex_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>
-   %0 = "stablehlo.complex"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<complex<f32>>
-   func.return %0 : tensor<complex<f32>>
- }
- 
- // CHECK-LABEL: "op_concatenate"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_concatenate(%arg0: tensor<8xf32>, %arg1: tensor<8xf32>) -> tensor<16xf32> {
--  //      CHECK: "vhlo.concatenate_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.concatenate_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   dimension = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<8x!vhlo.f32_v1>, !vhlo.tensor_v1<8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x!vhlo.f32_v1>
-   %0 = "stablehlo.concatenate"(%arg0, %arg1) {
-@@ -1012,6 +1088,7 @@
- }
- 
- // CHECK-LABEL: "op_constant"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_constant(%arg0: tensor<f32>) -> tensor<f32> {
-   //      CHECK: "vhlo.constant_v1"() <{
-   // CHECK-SAME:   value = #vhlo.tensor_v1<dense<0.000000e+00> : tensor<f32>>
-@@ -1023,15 +1100,17 @@
- }
- 
- // CHECK-LABEL: "op_convert"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_convert(%arg0: tensor<i32>) -> tensor<f32> {
--  // CHECK: "vhlo.convert_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.convert_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.convert"(%arg0) : (tensor<i32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_convolution"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_convolution(%arg0: tensor<1x8x8x207xf32>, %arg1: tensor<3x3x207x16xf32>) -> tensor<1x7x7x16xf32> {
--  //      CHECK: "vhlo.convolution_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.convolution_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   batch_group_count = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   feature_group_count = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   input_batch_dimension = #vhlo.integer_v1<0 : i64>,
-@@ -1065,8 +1144,9 @@
- }
- 
- // CHECK-LABEL: "op_cosine"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_cosine(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.cosine_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.cosine_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.cosine"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
-@@ -1079,8 +1159,9 @@
- }
- 
- // CHECK-LABEL: "op_cross_replica_sum"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_cross_replica_sum(%arg0: tensor<f32>) -> tensor<f32> {
--  //               CHECK: "vhlo.cross-replica-sum_v1"(%arg0) <{
-+  //               CHECK: "vhlo.cross-replica-sum_v1"(%[[ARG0]]) <{
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0], [1]]> : tensor<2x1xi64>>
-   //          CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.cross-replica-sum"(%arg0) {
-@@ -1090,8 +1171,9 @@
- }
- 
- // CHECK-LABEL: "op_custom_call"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_custom_call(%arg0: tensor<f32>) -> tensor<f32> {
--  //      CHECK: "vhlo.custom_call_v1"(%arg0) <{
-+  //      CHECK: "vhlo.custom_call_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   api_version = #vhlo<api_version_v1 API_VERSION_STATUS_RETURNING>,
-   // CHECK-SAME:   backend_config = #vhlo.string_v1<"\08\03\1A\02">,
-   // CHECK-SAME:   call_target_name = #vhlo.string_v1<"foo">,
-@@ -1122,15 +1204,17 @@
- }
- 
- // CHECK-LABEL: "op_divide"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_divide(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.divide_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.divide_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.divide"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_dot_general"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_dot_general(%arg0: tensor<8x8x16xf32>, %arg1: tensor<8x16x8xf32>) -> tensor<8x8x8xf32> {
--  //      CHECK: "vhlo.dot_general_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.dot_general_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   lhs_batching_dimensions = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>,
-   // CHECK-SAME:   lhs_contracting_dimensions = #vhlo.tensor_v1<dense<2> : tensor<1xi64>>,
-   // CHECK-SAME:   precision_config = #vhlo.array_v1<[#vhlo<precision_v1 HIGHEST>, #vhlo<precision_v1 HIGHEST>]>,
-@@ -1150,8 +1234,9 @@
- }
- 
- // CHECK-LABEL: "op_dot"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_dot(%arg0: tensor<8x16xf32>, %arg1: tensor<16x8xf32>) -> tensor<8x8xf32> {
--  //      CHECK: "vhlo.dot_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.dot_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   precision_config = #vhlo.array_v1<[#vhlo<precision_v1 HIGHEST>, #vhlo<precision_v1 HIGHEST>]>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<8x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<8x8x!vhlo.f32_v1>
-   %0 = "stablehlo.dot"(%arg0, %arg1) {
-@@ -1161,8 +1246,9 @@
- }
- 
- // CHECK-LABEL: "op_dynamic_broadcast_in_dim"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_dynamic_broadcast_in_dim(%arg0: tensor<?x?xf32>, %arg1: tensor<2xindex>) -> tensor<?x?xf32> {
--  //      CHECK: "vhlo.dynamic_broadcast_in_dim_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.dynamic_broadcast_in_dim_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   broadcast_dimensions = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-   // CHECK-SAME:   known_expanding_dimensions = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>,
-   // CHECK-SAME:   known_nonexpanding_dimensions = #vhlo.tensor_v1<dense<1> : tensor<1xi64>>
-@@ -1176,8 +1262,9 @@
- }
- 
- // CHECK-LABEL: "op_dynamic_conv"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_dynamic_conv(%arg0: tensor<1x8x8x207xf32>, %arg1: tensor<3x3x207x16xf32>, %arg2: tensor<4xi32>) -> tensor<1x?x?x16xf32> {
--  //      CHECK: "vhlo.dynamic_conv_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.dynamic_conv_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   batch_group_count = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   feature_group_count = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   input_batch_dimension = #vhlo.integer_v1<0 : i64>,
-@@ -1211,8 +1298,9 @@
- }
- 
- // CHECK-LABEL: "op_dynamic_gather"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_dynamic_gather(%arg0 : tensor<2x4x9xf32>, %arg1 : tensor<1x5x2xi32>, %arg2 : tensor<3xi32>) -> tensor<1x5x8xf32> {
--  //      CHECK: "vhlo.dynamic_gather_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.dynamic_gather_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   collapsed_slice_dims = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-   // CHECK-SAME:   index_vector_dim = #vhlo.integer_v1<2 : i64>,
-   // CHECK-SAME:   indices_are_sorted = #vhlo.bool_v1<true>,
-@@ -1232,8 +1320,9 @@
- }
- 
- // CHECK-LABEL: "op_dynamic_iota"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_dynamic_iota(%arg0: tensor<1xindex>) -> tensor<?xf32> {
--  //      CHECK: "vhlo.dynamic_iota_v1"(%arg0) <{
-+  //      CHECK: "vhlo.dynamic_iota_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   iota_dimension = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<1x!vhlo.index_v1>) -> !vhlo.tensor_v1<?x!vhlo.f32_v1>
-   %0 = "stablehlo.dynamic_iota"(%arg0) {
-@@ -1243,22 +1332,25 @@
- }
- 
- // CHECK-LABEL: "op_dynamic_pad"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}}, %[[ARG3:.*]]: {{.*}}, %[[ARG4:.*]]: {{.*}})
- func.func @op_dynamic_pad(%arg0: tensor<?xf32>, %arg1: tensor<f32>, %arg2: tensor<1xindex>, %arg3: tensor<1xindex>, %arg4: tensor<1xindex>) -> tensor<?xf32> {
--  // CHECK: "vhlo.dynamic_pad_v1"(%arg0, %arg1, %arg2, %arg3, %arg4) : (!vhlo.tensor_v1<?x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>) -> !vhlo.tensor_v1<?x!vhlo.f32_v1>
-+  // CHECK: "vhlo.dynamic_pad_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]], %[[ARG3]], %[[ARG4]]) : (!vhlo.tensor_v1<?x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>) -> !vhlo.tensor_v1<?x!vhlo.f32_v1>
-   %0 = "stablehlo.dynamic_pad"(%arg0, %arg1, %arg2, %arg3, %arg4) : (tensor<?xf32>, tensor<f32>, tensor<1xindex>, tensor<1xindex>, tensor<1xindex>) -> tensor<?xf32>
-   func.return %0 : tensor<?xf32>
- }
- 
- // CHECK-LABEL: "op_dynamic_reshape"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_dynamic_reshape(%arg0: tensor<16xf32>, %arg1: tensor<2xindex>) -> tensor<?x?xf32> {
--  // CHECK: "vhlo.dynamic_reshape_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<2x!vhlo.index_v1>) -> !vhlo.tensor_v1<?x?x!vhlo.f32_v1>
-+  // CHECK: "vhlo.dynamic_reshape_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<2x!vhlo.index_v1>) -> !vhlo.tensor_v1<?x?x!vhlo.f32_v1>
-   %0 = "stablehlo.dynamic_reshape"(%arg0, %arg1) : (tensor<16xf32>, tensor<2xindex>) -> tensor<?x?xf32>
-   func.return %0 : tensor<?x?xf32>
- }
- 
- // CHECK-LABEL: "op_dynamic_slice"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_dynamic_slice(%arg0: tensor<16xf32>, %arg1: tensor<i64>) -> tensor<4xf32> {
--  //      CHECK: "vhlo.dynamic_slice_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.dynamic_slice_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   slice_sizes = #vhlo.tensor_v1<dense<4> : tensor<1xi64>>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.i64_v1>) -> !vhlo.tensor_v1<4x!vhlo.f32_v1>
-   %0 = "stablehlo.dynamic_slice"(%arg0, %arg1) {
-@@ -1268,15 +1360,17 @@
- }
- 
- // CHECK-LABEL: "op_dynamic_update_slice"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_dynamic_update_slice(%arg0: tensor<16xf32>, %arg1: tensor<4xf32>, %arg2: tensor<i64>) -> tensor<16xf32> {
--  // CHECK: "vhlo.dynamic_update_slice_v1"(%arg0, %arg1, %arg2) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<4x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.i64_v1>) -> !vhlo.tensor_v1<16x!vhlo.f32_v1>
-+  // CHECK: "vhlo.dynamic_update_slice_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<4x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.i64_v1>) -> !vhlo.tensor_v1<16x!vhlo.f32_v1>
-   %0 = "stablehlo.dynamic_update_slice"(%arg0, %arg1, %arg2) : (tensor<16xf32>, tensor<4xf32>, tensor<i64>) -> tensor<16xf32>
-   func.return %0 : tensor<16xf32>
- }
- 
- // CHECK-LABEL: "op_einsum"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_einsum(%arg0: tensor<8x16xf32>, %arg1: tensor<16x8xf32>) -> tensor<8x8xf32> {
--  //      CHECK: "vhlo.einsum_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.einsum_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   einsum_config = #vhlo.string_v1<"ab,bc->ac">
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<8x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<8x8x!vhlo.f32_v1>
-   %0 = "stablehlo.einsum"(%arg0, %arg1) {
-@@ -1286,22 +1380,25 @@
- }
- 
- // CHECK-LABEL: "op_exponential_minus_one"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_exponential_minus_one(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.exponential_minus_one_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.exponential_minus_one_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.exponential_minus_one"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_exponential"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_exponential(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.exponential_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.exponential_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.exponential"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_fft"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_fft(%arg0: tensor<16xcomplex<f32>>) -> tensor<16xcomplex<f32>> {
--  //      CHECK: "vhlo.fft_v1"(%arg0) <{
-+  //      CHECK: "vhlo.fft_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   fft_length = #vhlo.tensor_v1<dense<16> : tensor<1xi64>>,
-   // CHECK-SAME:   fft_type = #vhlo<fft_type_v1 FFT>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<16x!vhlo.complex_v1<!vhlo.f32_v1>>) -> !vhlo.tensor_v1<16x!vhlo.complex_v1<!vhlo.f32_v1>>
-@@ -1313,8 +1410,9 @@
- }
- 
- // CHECK-LABEL: "op_floor"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_floor(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.floor_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.floor_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.floor"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
-@@ -1327,16 +1425,17 @@
-   // CHECK-SAME:   sym_name = #vhlo.string_v1<"op_func">,
-   // CHECK-SAME:   sym_visibility = #vhlo.string_v1<"private">
-   // CHECK-SAME: }> ({
--  // CHECK-NEXT:   ^[[BB:bb.*]](%arg0: !vhlo.tensor_v1<!vhlo.f32_v1>):
--  // CHECK-NEXT:     "vhlo.return_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-+  // CHECK-NEXT:   ^[[BB:bb.*]](%[[ARG0:.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
-+  // CHECK-NEXT:     "vhlo.return_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-   // CHECK-NEXT: }) : () -> ()
- 
-   func.return %arg0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_gather"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_gather(%arg0 : tensor<2x4x9xf32>, %arg1 : tensor<1x5x2xi32>) -> tensor<1x5x1xf32> {
--  //      CHECK: "vhlo.gather_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.gather_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   collapsed_slice_dims = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-   // CHECK-SAME:   index_vector_dim = #vhlo.integer_v1<2 : i64>,
-   // CHECK-SAME:   indices_are_sorted = #vhlo.bool_v1<true>,
-@@ -1358,8 +1457,9 @@
- }
- 
- // CHECK-LABEL: "op_get_dimension_size"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_get_dimension_size(%arg0: tensor<?xf32>) -> tensor<i32> {
--  //      CHECK: "vhlo.get_dimension_size_v1"(%arg0) <{
-+  //      CHECK: "vhlo.get_dimension_size_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   dimension = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<?x!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-   %0 = "stablehlo.get_dimension_size"(%arg0) {
-@@ -1369,8 +1469,9 @@
- }
- 
- // CHECK-LABEL: "op_get_tuple_element"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_get_tuple_element(%arg0: tuple<tensor<f32>, tensor<i32>>) -> tensor<f32> {
--  //      CHECK: "vhlo.get_tuple_element_v1"(%arg0) <{
-+  //      CHECK: "vhlo.get_tuple_element_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   index = #vhlo.integer_v1<0 : i32>
-   // CHECK-SAME: }> : (!vhlo.tuple_v1<!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.get_tuple_element"(%arg0) {
-@@ -1380,11 +1481,12 @@
- }
- 
- // CHECK-LABEL: "op_if"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_if(%arg0: tensor<i1>, %arg1: tensor<f32>, %arg2: tensor<f32>) -> tensor<f32> {
--  //      CHECK: "vhlo.if_v1"(%arg0) ({
--  // CHECK-NEXT:   "vhlo.return_v1"(%arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-+  //      CHECK: "vhlo.if_v1"(%[[ARG0]]) ({
-+  // CHECK-NEXT:   "vhlo.return_v1"(%[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-   // CHECK-NEXT: }, {
--  // CHECK-NEXT:   "vhlo.return_v1"(%arg2) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-+  // CHECK-NEXT:   "vhlo.return_v1"(%[[ARG2]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-   // CHECK-NEXT: }) : (!vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.if"(%arg0) ({
-     "stablehlo.return"(%arg1) : (tensor<f32>) -> ()
-@@ -1395,15 +1497,17 @@
- }
- 
- // CHECK-LABEL: "op_imag"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_imag(%arg0: tensor<complex<f32>>) -> tensor<f32> {
--  // CHECK: "vhlo.imag_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.imag_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.imag"(%arg0) : (tensor<complex<f32>>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_infeed"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_infeed(%arg0: !stablehlo.token) -> (tensor<f32>, !stablehlo.token) {
--  //               CHECK: "vhlo.infeed_v1"(%arg0) <{
-+  //               CHECK: "vhlo.infeed_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   infeed_config = #vhlo.string_v1<"foo">,
-   // CHECK-SAME{LITERAL}:   layout = #vhlo.array_v1<[#vhlo.array_v1<[]>]>
-   //          CHECK-SAME: }> : (!vhlo.token_v1) -> (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.token_v1)
-@@ -1426,36 +1530,41 @@
- }
- 
- // CHECK-LABEL: "op_is_finite"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_is_finite(%arg0: tensor<f32>) -> tensor<i1> {
--  // CHECK: "vhlo.is_finite_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-+  // CHECK: "vhlo.is_finite_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-   %0 = "stablehlo.is_finite"(%arg0) : (tensor<f32>) -> tensor<i1>
-   func.return %0 : tensor<i1>
- }
- 
- // CHECK-LABEL: "op_log"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_log(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.log_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.log_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.log"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_log_plus_one"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_log_plus_one(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.log_plus_one_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.log_plus_one_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.log_plus_one"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_logistic"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_logistic(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.logistic_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.logistic_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.logistic"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_map"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_map(%arg0: tensor<16xf32>) -> tensor<16xf32> {
--  //      CHECK: "vhlo.map_v1"(%arg0) <{
-+  //      CHECK: "vhlo.map_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   dimensions = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>
-   // CHECK-SAME: }> ({
-   // CHECK-NEXT:   ^[[BB:bb.*]](%[[ARG1:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
-@@ -1473,57 +1582,65 @@
- }
- 
- // CHECK-LABEL: "op_maximum"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_maximum(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.maximum_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.maximum_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.maximum"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_minimum"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_minimum(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.minimum_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.minimum_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.minimum"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_multiply"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_multiply(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.multiply_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.multiply_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.multiply"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_negate"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_negate(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.negate_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.negate_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.negate"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_not"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_not(%arg0: tensor<i1>) -> tensor<i1> {
--  // CHECK: "vhlo.not_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-+  // CHECK: "vhlo.not_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-   %0 = "stablehlo.not"(%arg0) : (tensor<i1>) -> tensor<i1>
-   func.return %0 : tensor<i1>
- }
- 
- // CHECK-LABEL: "op_optimization_barrier"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_optimization_barrier(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.optimization_barrier_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.optimization_barrier_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.optimization_barrier"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_or"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_or(%arg0: tensor<i1>, %arg1: tensor<i1>) -> tensor<i1> {
--  // CHECK: "vhlo.or_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-+  // CHECK: "vhlo.or_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-   %0 = "stablehlo.or"(%arg0, %arg1) : (tensor<i1>, tensor<i1>) -> tensor<i1>
-   func.return %0 : tensor<i1>
- }
- 
- // CHECK-LABEL: "op_outfeed"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_outfeed(%arg0: tensor<f32>, %arg1: !stablehlo.token) -> !stablehlo.token {
--  //      CHECK: "vhlo.outfeed_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.outfeed_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   outfeed_config = #vhlo.string_v1<"foo">
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.token_v1) -> !vhlo.token_v1
-   %0 = "stablehlo.outfeed"(%arg0, %arg1) {
-@@ -1533,8 +1650,9 @@
- }
- 
- // CHECK-LABEL: "op_pad"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_pad(%arg0: tensor<8xf32>, %arg1: tensor<f32>) -> tensor<16xf32> {
--  //      CHECK: "vhlo.pad_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.pad_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   edge_padding_high = #vhlo.tensor_v1<dense<4> : tensor<1xi64>>,
-   // CHECK-SAME:   edge_padding_low = #vhlo.tensor_v1<dense<4> : tensor<1xi64>>,
-   // CHECK-SAME:   interior_padding = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>
-@@ -1548,36 +1666,41 @@
- }
- 
- // CHECK-LABEL: "op_popcnt"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_popcnt(%arg0: tensor<i32>) -> tensor<i32> {
--  // CHECK: "vhlo.popcnt_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-+  // CHECK: "vhlo.popcnt_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-   %0 = "stablehlo.popcnt"(%arg0) : (tensor<i32>) -> tensor<i32>
-   func.return %0 : tensor<i32>
- }
- 
- // CHECK-LABEL: "op_power"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_power(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.power_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.power_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.power"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_real_dynamic_slice"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}}, %[[ARG3:.*]]: {{.*}})
- func.func @op_real_dynamic_slice(%arg0: tensor<?xf32>, %arg1: tensor<1xindex>, %arg2: tensor<1xindex>, %arg3: tensor<1xindex>) -> tensor<?xf32> {
--  // CHECK: "vhlo.real_dynamic_slice_v1"(%arg0, %arg1, %arg2, %arg3) : (!vhlo.tensor_v1<?x!vhlo.f32_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>) -> !vhlo.tensor_v1<?x!vhlo.f32_v1>
-+  // CHECK: "vhlo.real_dynamic_slice_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]], %[[ARG3]]) : (!vhlo.tensor_v1<?x!vhlo.f32_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>) -> !vhlo.tensor_v1<?x!vhlo.f32_v1>
-   %0 = "stablehlo.real_dynamic_slice"(%arg0, %arg1, %arg2, %arg3) : (tensor<?xf32>, tensor<1xindex>, tensor<1xindex>, tensor<1xindex>) -> tensor<?xf32>
-   func.return %0 : tensor<?xf32>
- }
- 
- // CHECK-LABEL: "op_real"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_real(%arg0: tensor<complex<f32>>) -> tensor<f32> {
--  // CHECK: "vhlo.real_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.real_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.real"(%arg0) : (tensor<complex<f32>>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_recv"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_recv(%arg0: !stablehlo.token) -> (tensor<f32>, !stablehlo.token) {
--  //      CHECK: "vhlo.recv_v1"(%arg0) <{
-+  //      CHECK: "vhlo.recv_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME:   channel_type = #vhlo.integer_v1<3 : i64>,
-   // CHECK-SAME:   is_host_transfer = #vhlo.bool_v1<true>
-@@ -1590,8 +1713,9 @@
- }
- 
- // CHECK-LABEL: "op_reduce"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_reduce(%arg0: tensor<16xf32>, %arg1: tensor<f32>) -> tensor<f32> {
--  //  CHECK: "vhlo.reduce_v1"(%arg0, %arg1)
-+  //  CHECK: "vhlo.reduce_v1"(%[[ARG0]], %[[ARG1]])
-   //  CHECK:   ^[[BB:bb.*]](%[[ARG1:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>, %[[ARG2:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
-   //  CHECK:     "vhlo.return_v1"(%[[VAL1:.*]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-   //  CHECK: }) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-@@ -1606,8 +1730,9 @@
- }
- 
- // CHECK-LABEL: "op_reduce_precision"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_reduce_precision(%arg0: tensor<f32>) -> tensor<f32> {
--  //      CHECK: "vhlo.reduce_precision_v1"(%arg0) <{
-+  //      CHECK: "vhlo.reduce_precision_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   exponent_bits = #vhlo.integer_v1<8 : i32>
-   // CHECK-SAME:   mantissa_bits = #vhlo.integer_v1<10 : i32>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-@@ -1619,8 +1744,9 @@
- }
- 
- // CHECK-LABEL: "op_reduce_scatter"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_reduce_scatter(%arg0: tensor<16xf32>) -> tensor<16xf32> {
--  //               CHECK: "vhlo.reduce_scatter_v1"(%arg0) <{
-+  //               CHECK: "vhlo.reduce_scatter_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0], [1]]> : tensor<2x1xi64>>,
-   //          CHECK-SAME:   scatter_dimension = #vhlo.integer_v1<0 : i64>
-@@ -1644,8 +1770,9 @@
- }
- 
- // CHECK-LABEL: "op_reduce_window"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_reduce_window(%arg0: tensor<2x17x31x7xf32>, %arg1: tensor<f32>) -> tensor<2x9x16x7xf32> {
--  //               CHECK: "vhlo.reduce_window_v1"(%arg0, %arg1) <{
-+  //               CHECK: "vhlo.reduce_window_v1"(%[[ARG0]], %[[ARG1]]) <{
-   //          CHECK-SAME:   base_dilations = #vhlo.tensor_v1<dense<[1, 2, 2, 1]> : tensor<4xi64>>,
-   // CHECK-SAME{LITERAL}:   padding = #vhlo.tensor_v1<dense<[[0, 0], [2, 0], [0, 2], [0, 0]]> : tensor<4x2xi64>>,
-   //          CHECK-SAME:   window_dilations = #vhlo.tensor_v1<dense<[1, 2, 2, 1]> : tensor<4xi64>>,
-@@ -1671,8 +1798,9 @@
- }
- 
- // CHECK-LABEL: "op_remainder"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_remainder(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.remainder_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.remainder_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.remainder"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
-@@ -1692,16 +1820,18 @@
- }
- 
- // CHECK-LABEL: "op_reshape"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_reshape(%arg0: tensor<16xf32>) -> tensor<4x4xf32> {
--  // CHECK: "vhlo.reshape_v1"(%arg0) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<4x4x!vhlo.f32_v1>
-+  // CHECK: "vhlo.reshape_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<4x4x!vhlo.f32_v1>
-   %0 = "stablehlo.reshape"(%arg0) : (tensor<16xf32>) -> tensor<4x4xf32>
-   func.return %0 : tensor<4x4xf32>
- }
- 
- // CHECK-LABEL: "op_return"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_return(%arg0: tensor<i32>, %arg1: tensor<f32>) -> tensor<f32> {
--  //      CHECK: "vhlo.case_v1"(%arg0) ({
--  // CHECK-NEXT:   "vhlo.return_v1"(%arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-+  //      CHECK: "vhlo.case_v1"(%[[ARG0]]) ({
-+  // CHECK-NEXT:   "vhlo.return_v1"(%[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-   // CHECK-NEXT: }) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.case"(%arg0) ({
-     "stablehlo.return"(%arg1) : (tensor<f32>) -> ()
-@@ -1710,8 +1840,9 @@
- }
- 
- // CHECK-LABEL: "op_reverse"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_reverse(%arg0: tensor<16xf32>) -> tensor<16xf32> {
--  //      CHECK: "vhlo.reverse_v1"(%arg0) <{
-+  //      CHECK: "vhlo.reverse_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   dimensions = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x!vhlo.f32_v1>
-   %0 = "stablehlo.reverse"(%arg0) {
-@@ -1721,8 +1852,9 @@
- }
- 
- // CHECK-LABEL: "op_rng_bit_generator"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_rng_bit_generator(%arg0: tensor<f32>) -> (tensor<f32>, tensor<f32>) {
--  //      CHECK: "vhlo.rng_bit_generator_v1"(%arg0) <{
-+  //      CHECK: "vhlo.rng_bit_generator_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   rng_algorithm = #vhlo<rng_algorithm_v1 PHILOX>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>)
-   %0:2 = "stablehlo.rng_bit_generator"(%arg0) {
-@@ -1732,8 +1864,9 @@
- }
- 
- // CHECK-LABEL: "op_rng"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_rng(%arg0: tensor<f32>, %arg1: tensor<f32>, %arg2: tensor<0xindex>) -> tensor<f32> {
--  //      CHECK: "vhlo.rng_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.rng_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   rng_distribution = #vhlo<rng_distribution_v1 NORMAL>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<0x!vhlo.index_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.rng"(%arg0, %arg1, %arg2) {
-@@ -1743,29 +1876,33 @@
- }
- 
- // CHECK-LABEL: "op_round_nearest_afz"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_round_nearest_afz(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.round_nearest_afz_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.round_nearest_afz_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.round_nearest_afz"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_round_nearest_even"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_round_nearest_even(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.round_nearest_even_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.round_nearest_even_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.round_nearest_even"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_rsqrt"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_rsqrt(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.rsqrt_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.rsqrt_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.rsqrt"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_scatter"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_scatter(%arg0: tensor<200x100x300xf32>, %arg1: tensor<10x2xi32>, %arg2: tensor<10x300xf32>) -> tensor<200x100x300xf32> {
--  //      CHECK: "vhlo.scatter_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.scatter_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   index_vector_dim = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   indices_are_sorted = #vhlo.bool_v1<true>,
-   // CHECK-SAME:   inserted_window_dims = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-@@ -1795,8 +1932,9 @@
- }
- 
- // CHECK-LABEL: "op_select_and_scatter"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_select_and_scatter(%arg0: tensor<10x24x24x64xf32>, %arg1: tensor<12x13x13x66xf32>, %arg2: tensor<f32>) -> tensor<10x24x24x64xf32> {
--  //      CHECK: "vhlo.select_and_scatter_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.select_and_scatter_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   padding = #vhlo.tensor_v1<dense<1> : tensor<4x2xi64>>,
-   // CHECK-SAME:   window_dimensions = #vhlo.tensor_v1<dense<[1, 2, 2, 1]> : tensor<4xi64>>,
-   // CHECK-SAME:   window_strides = #vhlo.tensor_v1<dense<[1, 2, 2, 1]> : tensor<4xi64>>
-@@ -1826,15 +1964,17 @@
- }
- 
- // CHECK-LABEL: "op_select"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_select(%arg0: tensor<i1>, %arg1: tensor<f32>, %arg2: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.select_v1"(%arg0, %arg1, %arg2) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.select_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.select"(%arg0, %arg1, %arg2) : (tensor<i1>, tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_send"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_send(%arg0: tensor<f32>, %arg1: !stablehlo.token) -> !stablehlo.token {
--  //      CHECK: "vhlo.send_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.send_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME:   channel_type = #vhlo.integer_v1<2 : i64>,
-   // CHECK-SAME:   is_host_transfer = #vhlo.bool_v1<true>
-@@ -1847,8 +1987,9 @@
- }
- 
- // CHECK-LABEL: "op_set_dimension_size"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_set_dimension_size(%arg0: tensor<?xf32>, %arg1: tensor<i32>) -> tensor<16xf32> {
--  //      CHECK: "vhlo.set_dimension_size_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.set_dimension_size_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   dimension = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<?x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<16x!vhlo.f32_v1>
-   %0 = "stablehlo.set_dimension_size"(%arg0, %arg1) {
-@@ -1858,43 +1999,49 @@
- }
- 
- // CHECK-LABEL: "op_shift_left"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_shift_left(%arg0: tensor<i32>, %arg1: tensor<i32>) -> tensor<i32> {
--  // CHECK: "vhlo.shift_left_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.i32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-+  // CHECK: "vhlo.shift_left_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-   %0 = "stablehlo.shift_left"(%arg0, %arg1) : (tensor<i32>, tensor<i32>) -> tensor<i32>
-   func.return %0 : tensor<i32>
- }
- 
- // CHECK-LABEL: "op_shift_right_arithmetic"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_shift_right_arithmetic(%arg0: tensor<i32>, %arg1: tensor<i32>) -> tensor<i32> {
--  // CHECK: "vhlo.shift_right_arithmetic_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.i32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-+  // CHECK: "vhlo.shift_right_arithmetic_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-   %0 = "stablehlo.shift_right_arithmetic"(%arg0, %arg1) : (tensor<i32>, tensor<i32>) -> tensor<i32>
-   func.return %0 : tensor<i32>
- }
- 
- // CHECK-LABEL: "op_shift_right_logical"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_shift_right_logical(%arg0: tensor<i32>, %arg1: tensor<i32>) -> tensor<i32> {
--  // CHECK: "vhlo.shift_right_logical_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.i32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-+  // CHECK: "vhlo.shift_right_logical_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-   %0 = "stablehlo.shift_right_logical"(%arg0, %arg1) : (tensor<i32>, tensor<i32>) -> tensor<i32>
-   func.return %0 : tensor<i32>
- }
- 
- // CHECK-LABEL: "op_sign"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_sign(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.sign_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.sign_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.sign"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_sine"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_sine(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.sine_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.sine_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.sine"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_slice"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_slice(%arg0: tensor<16xf32>) -> tensor<4xf32> {
--  //      CHECK: "vhlo.slice_v1"(%arg0) <{
-+  //      CHECK: "vhlo.slice_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   limit_indices = #vhlo.tensor_v1<dense<4> : tensor<1xi64>>,
-   // CHECK-SAME:   start_indices = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>,
-   // CHECK-SAME:   strides = #vhlo.tensor_v1<dense<1> : tensor<1xi64>>
-@@ -1908,8 +2055,9 @@
- }
- 
- // CHECK-LABEL: "op_sort"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_sort(%arg0: tensor<16xf32>) -> tensor<16xf32> {
--  //      CHECK: "vhlo.sort_v1"(%arg0) <{
-+  //      CHECK: "vhlo.sort_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   dimension = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME:   is_stable = #vhlo.bool_v1<true>
-   // CHECK-SAME: }> ({
-@@ -1929,29 +2077,33 @@
- }
- 
- // CHECK-LABEL: "op_sqrt"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_sqrt(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.sqrt_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.sqrt_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.sqrt"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_subtract"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_subtract(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.subtract_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.subtract_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.subtract"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_tanh"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_tanh(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.tanh_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.tanh_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.tanh"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_torch_index_select"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_torch_index_select(%arg0: tensor<5x1x5xf32>, %arg1: tensor<2xi32>) ->  tensor<2x1x5xf32> {
--  //      CHECK: "vhlo.torch_index_select_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.torch_index_select_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   batch_dims = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME:   dim = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<5x1x5x!vhlo.f32_v1>, !vhlo.tensor_v1<2x!vhlo.i32_v1>) -> !vhlo.tensor_v1<2x1x5x!vhlo.f32_v1>
-@@ -1963,8 +2115,9 @@
- }
- 
- // CHECK-LABEL: "op_trace"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_trace(%arg0: tensor<f32>) {
--  //      CHECK: "vhlo.trace_v1"(%arg0) <{
-+  //      CHECK: "vhlo.trace_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   tag = #vhlo.string_v1<"foo">
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-   "stablehlo.trace"(%arg0) {
-@@ -1974,8 +2127,9 @@
- }
- 
- // CHECK-LABEL: "op_transpose"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_transpose(%arg0: tensor<16x8xf32>) ->  tensor<8x16xf32> {
--  //      CHECK: "vhlo.transpose_v1"(%arg0) <{
-+  //      CHECK: "vhlo.transpose_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   permutation = #vhlo.tensor_v1<dense<[1, 0]> : tensor<2xi64>>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<16x8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<8x16x!vhlo.f32_v1>
-   %0 = "stablehlo.transpose"(%arg0) {
-@@ -1985,8 +2139,9 @@
- }
- 
- // CHECK-LABEL: "op_triangular_solve"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_triangular_solve(%arg0: tensor<16x16xf32>, %arg1: tensor<16x16xf32>) ->  tensor<16x16xf32> {
--  //      CHECK: "vhlo.triangular_solve_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.triangular_solve_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   left_side = #vhlo.bool_v1<true>,
-   // CHECK-SAME:   lower = #vhlo.bool_v1<true>,
-   // CHECK-SAME:   transpose_a = #vhlo<transpose_v1 NO_TRANSPOSE>,
-@@ -2002,15 +2157,17 @@
- }
- 
- // CHECK-LABEL: "op_tuple"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_tuple(%arg0: tensor<f32>) -> tuple<tensor<f32>> {
--  // CHECK: "vhlo.tuple_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tuple_v1<!vhlo.tensor_v1<!vhlo.f32_v1>>
-+  // CHECK: "vhlo.tuple_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tuple_v1<!vhlo.tensor_v1<!vhlo.f32_v1>>
-   %0 = "stablehlo.tuple"(%arg0) : (tensor<f32>) -> tuple<tensor<f32>>
-   func.return %0 : tuple<tensor<f32>>
- }
- 
- // CHECK-LABEL: "op_unary_einsum"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_unary_einsum(%arg0: tensor<8x16xf32>) -> tensor<8xf32> {
--  //      CHECK: "vhlo.unary_einsum_v1"(%arg0) <{
-+  //      CHECK: "vhlo.unary_einsum_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   einsum_config = #vhlo.string_v1<"ab->a">
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<8x16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<8x!vhlo.f32_v1>
-   %0 = "stablehlo.unary_einsum"(%arg0) {
-@@ -2020,22 +2177,25 @@
- }
- 
- // CHECK-LABEL: "op_uniform_dequantize"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_uniform_dequantize(%arg0: tensor<!quant.uniform<i8:f32, 34.0:16>>) -> tensor<f32> {
--  // CHECK: "vhlo.uniform_dequantize_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.uniform_dequantize_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.uniform_dequantize"(%arg0) : (tensor<!quant.uniform<i8:f32, 34.0:16>>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_uniform_quantize"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_uniform_quantize(%arg0: tensor<f32>) -> tensor<!quant.uniform<i8:f32, 34.0:16>> {
--  // CHECK: "vhlo.uniform_quantize_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>
-+  // CHECK: "vhlo.uniform_quantize_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>
-   %0 = "stablehlo.uniform_quantize"(%arg0) : (tensor<f32>) -> tensor<!quant.uniform<i8:f32, 34.0:16>>
-   func.return %0 : tensor<!quant.uniform<i8:f32, 34.0:16>>
- }
- 
- // CHECK-LABEL: "op_while"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_while(%arg0: tensor<i1>) -> tensor<i1> {
--  //      CHECK: "vhlo.while_v1"(%arg0) ({
-+  //      CHECK: "vhlo.while_v1"(%[[ARG0]]) ({
-   // CHECK-NEXT:   ^[[BB:bb.*]](%[[ARG1:arg.*]]: !vhlo.tensor_v1<!vhlo.bool_v1>):
-   // CHECK-NEXT:     "vhlo.return_v1"(%[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>) -> ()
-   // CHECK-NEXT:   }, {
-@@ -2053,8 +2213,9 @@
- }
- 
- // CHECK-LABEL: "op_xor"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_xor(%arg0: tensor<i1>, %arg1: tensor<i1>) -> tensor<i1> {
--  // CHECK: "vhlo.xor_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-+  // CHECK: "vhlo.xor_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-   %0 = "stablehlo.xor"(%arg0, %arg1) : (tensor<i1>, tensor<i1>) -> tensor<i1>
-   func.return %0 : tensor<i1>
- }
-@@ -2062,169 +2223,193 @@
- // ============ TYPES ============
- 
- // CHECK-LABEL: "type_i1"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_i1(%arg0: tensor<i1>, %arg1: tensor<i1>) -> tensor<i1> {
--  // CHECK: "vhlo.and_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-+  // CHECK: "vhlo.and_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-   %0 = "stablehlo.and"(%arg0, %arg1) : (tensor<i1>, tensor<i1>) -> tensor<i1>
-   func.return %0 : tensor<i1>
- }
- 
- // CHECK-LABEL: "type_i4"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_i4(%arg0: tensor<i4>, %arg1: tensor<i4>) -> tensor<i4> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.i4_v1>, !vhlo.tensor_v1<!vhlo.i4_v1>) -> !vhlo.tensor_v1<!vhlo.i4_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i4_v1>, !vhlo.tensor_v1<!vhlo.i4_v1>) -> !vhlo.tensor_v1<!vhlo.i4_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<i4>, tensor<i4>) -> tensor<i4>
-   func.return %0 : tensor<i4>
- }
- 
- // CHECK-LABEL: "type_i8"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_i8(%arg0: tensor<i8>, %arg1: tensor<i8>) -> tensor<i8> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.i8_v1>, !vhlo.tensor_v1<!vhlo.i8_v1>) -> !vhlo.tensor_v1<!vhlo.i8_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i8_v1>, !vhlo.tensor_v1<!vhlo.i8_v1>) -> !vhlo.tensor_v1<!vhlo.i8_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<i8>, tensor<i8>) -> tensor<i8>
-   func.return %0 : tensor<i8>
- }
- 
- // CHECK-LABEL: "type_i16"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_i16(%arg0: tensor<i16>, %arg1: tensor<i16>) -> tensor<i16> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.i16_v1>, !vhlo.tensor_v1<!vhlo.i16_v1>) -> !vhlo.tensor_v1<!vhlo.i16_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i16_v1>, !vhlo.tensor_v1<!vhlo.i16_v1>) -> !vhlo.tensor_v1<!vhlo.i16_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<i16>, tensor<i16>) -> tensor<i16>
-   func.return %0 : tensor<i16>
- }
- 
- // CHECK-LABEL: "type_i32"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_i32(%arg0: tensor<i32>, %arg1: tensor<i32>) -> tensor<i32> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.i32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<i32>, tensor<i32>) -> tensor<i32>
-   func.return %0 : tensor<i32>
- }
- 
- // CHECK-LABEL: "type_i64"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_i64(%arg0: tensor<i64>, %arg1: tensor<i64>) -> tensor<i64> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.i64_v1>, !vhlo.tensor_v1<!vhlo.i64_v1>) -> !vhlo.tensor_v1<!vhlo.i64_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i64_v1>, !vhlo.tensor_v1<!vhlo.i64_v1>) -> !vhlo.tensor_v1<!vhlo.i64_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<i64>, tensor<i64>) -> tensor<i64>
-   func.return %0 : tensor<i64>
- }
- 
- // CHECK-LABEL: "type_ui4"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_ui4(%arg0: tensor<ui4>, %arg1: tensor<ui4>) -> tensor<ui4> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.ui4_v1>, !vhlo.tensor_v1<!vhlo.ui4_v1>) -> !vhlo.tensor_v1<!vhlo.ui4_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.ui4_v1>, !vhlo.tensor_v1<!vhlo.ui4_v1>) -> !vhlo.tensor_v1<!vhlo.ui4_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<ui4>, tensor<ui4>) -> tensor<ui4>
-   func.return %0 : tensor<ui4>
- }
- 
- // CHECK-LABEL: "type_ui8"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_ui8(%arg0: tensor<ui8>, %arg1: tensor<ui8>) -> tensor<ui8> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.ui8_v1>, !vhlo.tensor_v1<!vhlo.ui8_v1>) -> !vhlo.tensor_v1<!vhlo.ui8_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.ui8_v1>, !vhlo.tensor_v1<!vhlo.ui8_v1>) -> !vhlo.tensor_v1<!vhlo.ui8_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<ui8>, tensor<ui8>) -> tensor<ui8>
-   func.return %0 : tensor<ui8>
- }
- 
- // CHECK-LABEL: "type_ui16"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_ui16(%arg0: tensor<ui16>, %arg1: tensor<ui16>) -> tensor<ui16> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.ui16_v1>, !vhlo.tensor_v1<!vhlo.ui16_v1>) -> !vhlo.tensor_v1<!vhlo.ui16_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.ui16_v1>, !vhlo.tensor_v1<!vhlo.ui16_v1>) -> !vhlo.tensor_v1<!vhlo.ui16_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<ui16>, tensor<ui16>) -> tensor<ui16>
-   func.return %0 : tensor<ui16>
- }
- 
- // CHECK-LABEL: "type_ui32"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_ui32(%arg0: tensor<ui32>, %arg1: tensor<ui32>) -> tensor<ui32> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.ui32_v1>, !vhlo.tensor_v1<!vhlo.ui32_v1>) -> !vhlo.tensor_v1<!vhlo.ui32_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.ui32_v1>, !vhlo.tensor_v1<!vhlo.ui32_v1>) -> !vhlo.tensor_v1<!vhlo.ui32_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<ui32>, tensor<ui32>) -> tensor<ui32>
-   func.return %0 : tensor<ui32>
- }
- 
- // CHECK-LABEL: "type_ui64"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_ui64(%arg0: tensor<ui64>, %arg1: tensor<ui64>) -> tensor<ui64> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.ui64_v1>, !vhlo.tensor_v1<!vhlo.ui64_v1>) -> !vhlo.tensor_v1<!vhlo.ui64_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.ui64_v1>, !vhlo.tensor_v1<!vhlo.ui64_v1>) -> !vhlo.tensor_v1<!vhlo.ui64_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<ui64>, tensor<ui64>) -> tensor<ui64>
-   func.return %0 : tensor<ui64>
- }
- 
- // CHECK-LABEL: "type_f8E4M3FN"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_f8E4M3FN(%arg0: tensor<f8E4M3FN>, %arg1: tensor<f8E4M3FN>) -> tensor<f8E4M3FN> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f8E4M3FN_v1>, !vhlo.tensor_v1<!vhlo.f8E4M3FN_v1>) -> !vhlo.tensor_v1<!vhlo.f8E4M3FN_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f8E4M3FN_v1>, !vhlo.tensor_v1<!vhlo.f8E4M3FN_v1>) -> !vhlo.tensor_v1<!vhlo.f8E4M3FN_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f8E4M3FN>, tensor<f8E4M3FN>) -> tensor<f8E4M3FN>
-   func.return %0 : tensor<f8E4M3FN>
- }
- 
- // CHECK-LABEL: "type_f8E5M2"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_f8E5M2(%arg0: tensor<f8E5M2>, %arg1: tensor<f8E5M2>) -> tensor<f8E5M2> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f8E5M2_v1>, !vhlo.tensor_v1<!vhlo.f8E5M2_v1>) -> !vhlo.tensor_v1<!vhlo.f8E5M2_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f8E5M2_v1>, !vhlo.tensor_v1<!vhlo.f8E5M2_v1>) -> !vhlo.tensor_v1<!vhlo.f8E5M2_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f8E5M2>, tensor<f8E5M2>) -> tensor<f8E5M2>
-   func.return %0 : tensor<f8E5M2>
- }
- 
- // CHECK-LABEL: "type_bf16"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_bf16(%arg0: tensor<bf16>, %arg1: tensor<bf16>) -> tensor<bf16> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.bf16_v1>, !vhlo.tensor_v1<!vhlo.bf16_v1>) -> !vhlo.tensor_v1<!vhlo.bf16_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.bf16_v1>, !vhlo.tensor_v1<!vhlo.bf16_v1>) -> !vhlo.tensor_v1<!vhlo.bf16_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<bf16>, tensor<bf16>) -> tensor<bf16>
-   func.return %0 : tensor<bf16>
- }
- 
- // CHECK-LABEL: "type_f16"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_f16(%arg0: tensor<f16>, %arg1: tensor<f16>) -> tensor<f16> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f16_v1>, !vhlo.tensor_v1<!vhlo.f16_v1>) -> !vhlo.tensor_v1<!vhlo.f16_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f16_v1>, !vhlo.tensor_v1<!vhlo.f16_v1>) -> !vhlo.tensor_v1<!vhlo.f16_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f16>, tensor<f16>) -> tensor<f16>
-   func.return %0 : tensor<f16>
- }
- 
- // CHECK-LABEL: "type_f32"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_f32(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "type_f64"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_f64(%arg0: tensor<f64>, %arg1: tensor<f64>) -> tensor<f64> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f64_v1>, !vhlo.tensor_v1<!vhlo.f64_v1>) -> !vhlo.tensor_v1<!vhlo.f64_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f64_v1>, !vhlo.tensor_v1<!vhlo.f64_v1>) -> !vhlo.tensor_v1<!vhlo.f64_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f64>, tensor<f64>) -> tensor<f64>
-   func.return %0 : tensor<f64>
- }
- 
- // CHECK-LABEL: "type_complex_f32"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_complex_f32(%arg0: tensor<complex<f32>>, %arg1: tensor<complex<f32>>) -> tensor<complex<f32>> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>, !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>) -> !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>, !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>) -> !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<complex<f32>>, tensor<complex<f32>>) -> tensor<complex<f32>>
-   func.return %0 : tensor<complex<f32>>
- }
- 
- // CHECK-LABEL: "type_complex_f64"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_complex_f64(%arg0: tensor<complex<f64>>, %arg1: tensor<complex<f64>>) -> tensor<complex<f64>> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f64_v1>>, !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f64_v1>>) -> !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f64_v1>>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f64_v1>>, !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f64_v1>>) -> !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f64_v1>>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<complex<f64>>, tensor<complex<f64>>) -> tensor<complex<f64>>
-   func.return %0 : tensor<complex<f64>>
- }
- 
- // CHECK-LABEL: "type_dynamism_ranked"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @type_dynamism_ranked(%arg0: tensor<?xf32>) -> tensor<?xf32> {
--  // CHECK: "vhlo.abs_v1"(%arg0) : (!vhlo.tensor_v1<?x!vhlo.f32_v1>) -> !vhlo.tensor_v1<?x!vhlo.f32_v1>
-+  // CHECK: "vhlo.abs_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<?x!vhlo.f32_v1>) -> !vhlo.tensor_v1<?x!vhlo.f32_v1>
-   %0 = "stablehlo.abs"(%arg0) : (tensor<?xf32>) -> tensor<?xf32>
-   func.return %0 : tensor<?xf32>
- }
- 
- // CHECK-LABEL: "type_quantization"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_quantization(%arg0: tensor<!quant.uniform<i8:f32, 34.0:16>>, %arg1: tensor<!quant.uniform<i8:f32, 34.0:16>>) -> tensor<!quant.uniform<i8:f32, 34.0:16>> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>, !vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>) -> !vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>, !vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>) -> !vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<!quant.uniform<i8:f32, 34.0:16>>, tensor<!quant.uniform<i8:f32, 34.0:16>>) -> tensor<!quant.uniform<i8:f32, 34.0:16>>
-   func.return %0 : tensor<!quant.uniform<i8:f32, 34.0:16>>
- }
- 
- //       CHECK: function_type = #vhlo.type_v1<!vhlo.func_v1<(!vhlo.token_v1) -> !vhlo.token_v1>>
- // CHECK-LABEL: "type_token_callee"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @type_token_callee(%arg0: !stablehlo.token) -> !stablehlo.token {
--  // CHECK: "vhlo.return_v1"(%arg0) : (!vhlo.token_v1) -> ()
-+  // CHECK: "vhlo.return_v1"(%[[ARG0]]) : (!vhlo.token_v1) -> ()
-   return %arg0 : !stablehlo.token
- }
- 
- //       CHECK: function_type = #vhlo.type_v1<!vhlo.func_v1<(!vhlo.token_v1) -> !vhlo.token_v1>>
- // CHECK-LABEL: "type_token_caller"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @type_token_caller(%arg0: !stablehlo.token) -> !stablehlo.token {
--  // CHECK:      "vhlo.call_v1"(%arg0) <{callee = #vhlo.string_v1<"type_token_callee">}
-+  // CHECK:      "vhlo.call_v1"(%[[ARG0]]) <{callee = #vhlo.string_v1<"type_token_callee">}
-   // CHECK-SAME: (!vhlo.token_v1) -> !vhlo.token_v1
-   %0 = func.call @type_token_callee(%arg0) : (!stablehlo.token) -> !stablehlo.token
-   return %0 : !stablehlo.token
- }
- 
- // CHECK-LABEL: "type_tuple"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @type_tuple(%arg0: tuple<tensor<f32>>) -> tuple<!stablehlo.token> {
-   %0 = "stablehlo.custom_call"(%arg0) {
-     call_target_name = "foo"
-diff --ruN a/stablehlo/stablehlo/tests/vhlo/stablehlo_legalize_to_vhlo.mlir b/stablehlo/stablehlo/tests/vhlo/stablehlo_legalize_to_vhlo.mlir
---- stablehlo/stablehlo/tests/vhlo/stablehlo_legalize_to_vhlo.mlir
-+++ stablehlo/stablehlo/tests/vhlo/stablehlo_legalize_to_vhlo.mlir
-@@ -13,6 +13,7 @@
- // ============ ATTRIBUTES ============
- 
- // CHECK-LABEL: "attr_comparison_direction_eq"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_direction_eq(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     // CHECK: comparison_direction = #vhlo<comparison_direction_v1 EQ>
-@@ -22,6 +23,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_direction_ne"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_direction_ne(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     // CHECK: comparison_direction = #vhlo<comparison_direction_v1 NE>
-@@ -31,6 +33,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_direction_ge"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_direction_ge(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     // CHECK: comparison_direction = #vhlo<comparison_direction_v1 GE>
-@@ -40,6 +43,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_direction_gt"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_direction_gt(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     // CHECK: comparison_direction = #vhlo<comparison_direction_v1 GT>
-@@ -49,6 +53,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_direction_le"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_direction_le(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     // CHECK: comparison_direction = #vhlo<comparison_direction_v1 LE>
-@@ -58,6 +63,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_direction_lt"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_direction_lt(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     // CHECK: comparison_direction = #vhlo<comparison_direction_v1 LT>
-@@ -67,6 +73,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_type_notype"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_type_notype(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     comparison_direction = #stablehlo<comparison_direction EQ>
-@@ -76,6 +83,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_type_float"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_type_float(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     comparison_direction = #stablehlo<comparison_direction EQ>,
-@@ -86,6 +94,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_type_totalorder"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_type_totalorder(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     comparison_direction = #stablehlo<comparison_direction EQ>,
-@@ -96,6 +105,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_type_signed"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_type_signed(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     comparison_direction = #stablehlo<comparison_direction EQ>,
-@@ -106,6 +116,7 @@
- }
- 
- // CHECK-LABEL: "attr_comparison_type_unsigned"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_comparison_type_unsigned(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-   %0 = "stablehlo.compare"(%arg0, %arg1) {
-     comparison_direction = #stablehlo<comparison_direction EQ>,
-@@ -118,6 +129,7 @@
- // ConvDimensionNumbers aka #stablehlo.conv is covered below.
- 
- // CHECK-LABEL: "attr_custom_call_api_version_unspecified"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_custom_call_api_version_unspecified(%arg0: tensor<f32>) -> tensor<f32> {
-   %0 = "stablehlo.custom_call"(%arg0) {
-     call_target_name = "foo",
-@@ -128,6 +140,7 @@
- }
- 
- // CHECK-LABEL: "attr_custom_call_api_version_original"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_custom_call_api_version_original(%arg0: tensor<f32>) -> tensor<f32> {
-   %0 = "stablehlo.custom_call"(%arg0) {
-     call_target_name = "foo",
-@@ -138,6 +151,7 @@
- }
- 
- // CHECK-LABEL: "attr_custom_call_api_version_status_returning"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_custom_call_api_version_status_returning(%arg0: tensor<f32>) -> tensor<f32> {
-   %0 = "stablehlo.custom_call"(%arg0) {
-     call_target_name = "foo",
-@@ -148,6 +162,7 @@
- }
- 
- // CHECK-LABEL: "attr_custom_call_api_version_status_returning_unified"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_custom_call_api_version_status_returning_unified(%arg0: tensor<f32>) -> tensor<f32> {
-   %0 = "stablehlo.custom_call"(%arg0) {
-     call_target_name = "foo",
-@@ -166,6 +181,7 @@
- // DotDimensionNumbers aka #stablehlo.dot is covered below.
- 
- // CHECK-LABEL: "attr_fft_type_fft"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_fft_type_fft(%arg0: tensor<16xcomplex<f32>>) -> tensor<16xcomplex<f32>> {
-   %0 = "stablehlo.fft"(%arg0) {
-     // CHECK: fft_type = #vhlo<fft_type_v1 FFT>
-@@ -176,6 +192,7 @@
- }
- 
- // CHECK-LABEL: "attr_fft_type_ifft"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_fft_type_ifft(%arg0: tensor<16xcomplex<f32>>) -> tensor<16xcomplex<f32>> {
-   %0 = "stablehlo.fft"(%arg0) {
-     // CHECK: fft_type = #vhlo<fft_type_v1 IFFT>
-@@ -186,6 +203,7 @@
- }
- 
- // CHECK-LABEL: "attr_fft_type_rfft"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_fft_type_rfft(%arg0: tensor<16xf32>) -> tensor<9xcomplex<f32>> {
-   %0 = "stablehlo.fft"(%arg0) {
-     // CHECK: fft_type = #vhlo<fft_type_v1 RFFT>
-@@ -196,6 +214,7 @@
- }
- 
- // CHECK-LABEL: "attr_fft_type_irfft"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_fft_type_irfft(%arg0: tensor<9xcomplex<f32>>) -> tensor<16xf32> {
-   %0 = "stablehlo.fft"(%arg0) {
-     // CHECK: fft_type = #vhlo<fft_type_v1 IRFFT>
-@@ -208,6 +227,7 @@
- // GatherDimensionNumbers aka #stablehlo.gather is covered below.
- 
- // CHECK-LABEL: "attr_precision_config_default"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_precision_config_default(%arg0: tensor<8x16xf32>, %arg1: tensor<16x8xf32>) -> tensor<8x8xf32> {
-   %0 = "stablehlo.dot"(%arg0, %arg1) {
-     // CHECK: precision_config = #vhlo.array_v1<[#vhlo<precision_v1 DEFAULT>, #vhlo<precision_v1 DEFAULT>]>
-@@ -216,6 +236,7 @@
- }
- 
- // CHECK-LABEL: "attr_precision_config_high"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_precision_config_high(%arg0: tensor<8x16xf32>, %arg1: tensor<16x8xf32>) -> tensor<8x8xf32> {
-   %0 = "stablehlo.dot"(%arg0, %arg1) {
-     // CHECK: precision_config = #vhlo.array_v1<[#vhlo<precision_v1 HIGH>, #vhlo<precision_v1 HIGH>]>
-@@ -225,6 +246,7 @@
- }
- 
- // CHECK-LABEL: "attr_precision_config_highest"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_precision_config_highest(%arg0: tensor<8x16xf32>, %arg1: tensor<16x8xf32>) -> tensor<8x8xf32> {
-   %0 = "stablehlo.dot"(%arg0, %arg1) {
-     // CHECK: precision_config = #vhlo.array_v1<[#vhlo<precision_v1 HIGHEST>, #vhlo<precision_v1 HIGHEST>]>
-@@ -234,6 +256,7 @@
- }
- 
- // CHECK-LABEL: "attr_rng_algorithm_default"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_rng_algorithm_default(%arg0: tensor<f32>) -> (tensor<f32>, tensor<f32>) {
-   %0:2 = "stablehlo.rng_bit_generator"(%arg0) {
-     // CHECK: rng_algorithm = #vhlo<rng_algorithm_v1 DEFAULT>
-@@ -243,6 +266,7 @@
- }
- 
- // CHECK-LABEL: "attr_rng_algorithm_three_fry"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_rng_algorithm_three_fry(%arg0: tensor<f32>) -> (tensor<f32>, tensor<f32>) {
-   %0:2 = "stablehlo.rng_bit_generator"(%arg0) {
-     // CHECK: rng_algorithm = #vhlo<rng_algorithm_v1 THREE_FRY>
-@@ -252,6 +276,7 @@
- }
- 
- // CHECK-LABEL: "attr_rng_algorithm_philox"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @attr_rng_algorithm_philox(%arg0: tensor<f32>) -> (tensor<f32>, tensor<f32>) {
-   %0:2 = "stablehlo.rng_bit_generator"(%arg0) {
-     // CHECK: rng_algorithm = #vhlo<rng_algorithm_v1 PHILOX>
-@@ -261,6 +286,7 @@
- }
- 
- // CHECK-LABEL: "attr_rng_distribution_uniform"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @attr_rng_distribution_uniform(%arg0: tensor<f32>, %arg1: tensor<f32>, %arg2: tensor<0xindex>) -> tensor<f32> {
-   %0 = "stablehlo.rng"(%arg0, %arg1, %arg2) {
-     // CHECK: rng_distribution = #vhlo<rng_distribution_v1 UNIFORM>
-@@ -270,6 +296,7 @@
- }
- 
- // CHECK-LABEL: "attr_rng_distribution_normal"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @attr_rng_distribution_normal(%arg0: tensor<f32>, %arg1: tensor<f32>, %arg2: tensor<0xindex>) -> tensor<f32> {
-   %0 = "stablehlo.rng"(%arg0, %arg1, %arg2) {
-     // CHECK: rng_distribution = #vhlo<rng_distribution_v1 NORMAL>
-@@ -281,6 +308,7 @@
- // ScatterDimensionNumbers aka #stablehlo.scatter is covered below.
- 
- // CHECK-LABEL: "attr_transpose_no_transpose"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_transpose_no_transpose(%arg0: tensor<16x16xf32>, %arg1: tensor<16x16xf32>) ->  tensor<16x16xf32> {
-   %0 = "stablehlo.triangular_solve"(%arg0, %arg1) {
-     left_side = true,
-@@ -293,6 +321,7 @@
- }
- 
- // CHECK-LABEL: "attr_transpose_transpose"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_transpose_transpose(%arg0: tensor<16x16xf32>, %arg1: tensor<16x16xf32>) ->  tensor<16x16xf32> {
-   %0 = "stablehlo.triangular_solve"(%arg0, %arg1) {
-     left_side = true,
-@@ -305,6 +334,7 @@
- }
- 
- // CHECK-LABEL: "attr_transpose_adjoint"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @attr_transpose_adjoint(%arg0: tensor<16x16xf32>, %arg1: tensor<16x16xf32>) ->  tensor<16x16xf32> {
-   %0 = "stablehlo.triangular_solve"(%arg0, %arg1) {
-     left_side = true,
-@@ -319,10 +349,9 @@
- // TypeExtensionsAttr aka #stablehlo.type_extensions is covered below.
- 
- // CHECK-LABEL: "attr_type_extensions_bounds"
--func.func @attr_type_extensions_bounds(
--    %arg0: tensor<?x?xf32, #stablehlo.type_extensions<bounds = [16, ?]>>)
--    -> tensor<?x?xf32, #stablehlo.type_extensions<bounds = [16, ?]>> {
--  // CHECK: "vhlo.return_v1"(%arg0) : (!vhlo.tensor_v1<?x?x!vhlo.f32_v1, #vhlo.type_extensions_v1<bounds = [16, ?]>>) -> ()
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @attr_type_extensions_bounds(%arg0: tensor<?x?xf32, #stablehlo.type_extensions<bounds = [16, ?]>>) -> tensor<?x?xf32, #stablehlo.type_extensions<bounds = [16, ?]>> {
-+  // CHECK: "vhlo.return_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<?x?x!vhlo.f32_v1, #vhlo.type_extensions_v1<bounds = [16, ?]>>) -> ()
-   func.return %arg0 : tensor<?x?xf32, #stablehlo.type_extensions<bounds = [16, ?]>>
- }
- 
-@@ -330,8 +359,9 @@
- // ============ DEFAULTS ============
- 
- // CHECK-LABEL: "default_all_gather"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_all_gather(%arg0: tensor<16x8xf32>) -> tensor<16x16xf32> {
--  //               CHECK: "vhlo.all_gather_v1"(%arg0) <{
-+  //               CHECK: "vhlo.all_gather_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   all_gather_dim = #vhlo.integer_v1<1 : i64>
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0], [1]]> : tensor<2x1xi64>>,
-@@ -345,8 +375,9 @@
- }
- 
- // CHECK-LABEL: "default_all_reduce"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_all_reduce(%arg0: tensor<f32>) -> tensor<f32> {
--  //               CHECK: "vhlo.all_reduce_v1"(%arg0)
-+  //               CHECK: "vhlo.all_reduce_v1"(%[[ARG0]])
-   //          CHECK-SAME: <{
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0], [1]]> : tensor<2x1xi64>>,
-@@ -368,8 +399,9 @@
- }
- 
- // CHECK-LABEL: "default_all_to_all"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_all_to_all(%arg0: tensor<4x16xf32>) -> tensor<16x4xf32> {
--  //               CHECK: "vhlo.all_to_all_v1"(%arg0) <{
-+  //               CHECK: "vhlo.all_to_all_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   //          CHECK-SAME:   concat_dimension = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0, 1, 2, 3]]> : tensor<1x4xi64>>,
-@@ -386,8 +418,9 @@
- }
- 
- // CHECK-LABEL: "default_cholesky"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_cholesky(%arg0: tensor<1x16x16xf32>) -> tensor<1x16x16xf32> {
--  //      CHECK: "vhlo.cholesky_v1"(%arg0) <{
-+  //      CHECK: "vhlo.cholesky_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   lower = #vhlo.bool_v1<false>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<1x16x16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<1x16x16x!vhlo.f32_v1>
-   %0 = "stablehlo.cholesky"(%arg0) : (tensor<1x16x16xf32>) -> tensor<1x16x16xf32>
-@@ -395,8 +428,9 @@
- }
- 
- // CHECK-LABEL: "default_collective_permute"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_collective_permute(%arg0: tensor<16x8xf32>) -> tensor<16x8xf32> {
--  //               CHECK: "vhlo.collective_permute_v1"(%arg0) <{
-+  //               CHECK: "vhlo.collective_permute_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME{LITERAL}:   source_target_pairs = #vhlo.tensor_v1<dense<[[0, 1], [1, 2], [2, 3]]> : tensor<3x2xi64>>
-   //          CHECK-SAME: }> : (!vhlo.tensor_v1<16x8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x8x!vhlo.f32_v1>
-@@ -407,8 +441,9 @@
- }
- 
- // CHECK-LABEL: "default_collective_broadcast"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_collective_broadcast(%arg0: tensor<16x8xf32>) -> tensor<16x8xf32> {
--  //               CHECK: "vhlo.collective_broadcast_v1"(%arg0) <{
-+  //               CHECK: "vhlo.collective_broadcast_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0, 1]]> : tensor<1x2xi64>>
-   //          CHECK-SAME: }> : (!vhlo.tensor_v1<16x8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x8x!vhlo.f32_v1>
-@@ -419,8 +454,9 @@
- }
- 
- // CHECK-LABEL: "default_compare"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @default_compare(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
--  //      CHECK: "vhlo.compare_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.compare_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   compare_type = #vhlo<comparison_type_v1 NOTYPE>,
-   // CHECK-SAME:   comparison_direction = #vhlo<comparison_direction_v1 EQ>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-@@ -431,8 +467,9 @@
- }
- 
- // CHECK-LABEL: "default_composite"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_composite(%arg0: tensor<f32>) -> tensor<f32> {
--  //               CHECK: "vhlo.composite_v1"(%arg0) <{
-+  //               CHECK: "vhlo.composite_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   composite_attributes = #vhlo.dict_v1<{}>
-   //          CHECK-SAME:   decomposition = #vhlo.string_v1<"composite_target">
-   //          CHECK-SAME:   name = #vhlo.string_v1<"stablehlo.composite_target">
-@@ -446,8 +483,9 @@
- }
- 
- // CHECK-LABEL: "default_convolution"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @default_convolution(%arg0: tensor<1x8x8x207xf32>, %arg1: tensor<3x3x207x16xf32>) -> tensor<1x6x6x16xf32> {
--  //      CHECK: "vhlo.convolution_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.convolution_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   batch_group_count = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   feature_group_count = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   input_batch_dimension = #vhlo.integer_v1<0 : i64>,
-@@ -475,8 +513,9 @@
- }
- 
- // CHECK-LABEL: "default_custom_call"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_custom_call(%arg0: tensor<f32>) -> tensor<f32> {
--  //      CHECK: "vhlo.custom_call_v1"(%arg0) <{
-+  //      CHECK: "vhlo.custom_call_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   api_version = #vhlo<api_version_v1 API_VERSION_ORIGINAL>,
-   // CHECK-SAME:   backend_config = #vhlo.string_v1<"">,
-   // CHECK-SAME:   call_target_name = #vhlo.string_v1<"foo">,
-@@ -493,8 +532,9 @@
- }
- 
- // CHECK-LABEL: "default_dot_general"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @default_dot_general(%arg0: tensor<8x8x16xf32>, %arg1: tensor<8x16x8xf32>) -> tensor<8x8x8xf32> {
--  //      CHECK: "vhlo.dot_general_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.dot_general_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   lhs_batching_dimensions = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>,
-   // CHECK-SAME:   lhs_contracting_dimensions = #vhlo.tensor_v1<dense<2> : tensor<1xi64>>,
-   // CHECK-SAME:   precision_config = #vhlo.array_v1<[#vhlo<precision_v1 DEFAULT>, #vhlo<precision_v1 DEFAULT>]>,
-@@ -513,8 +553,9 @@
- }
- 
- // CHECK-LABEL: "default_dot"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @default_dot(%arg0: tensor<8x16xf32>, %arg1: tensor<16x8xf32>) -> tensor<8x8xf32> {
--  //      CHECK: "vhlo.dot_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.dot_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   precision_config = #vhlo.array_v1<[#vhlo<precision_v1 DEFAULT>, #vhlo<precision_v1 DEFAULT>]>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<8x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<8x8x!vhlo.f32_v1>
-   %0 = "stablehlo.dot"(%arg0, %arg1) : (tensor<8x16xf32>, tensor<16x8xf32>) -> tensor<8x8xf32>
-@@ -522,8 +563,9 @@
- }
- 
- // CHECK-LABEL: "default_dynamic_broadcast_in_dim"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @default_dynamic_broadcast_in_dim(%arg0: tensor<?x?xf32>, %arg1: tensor<2xindex>) -> tensor<?x?xf32> {
--  //      CHECK: "vhlo.dynamic_broadcast_in_dim_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.dynamic_broadcast_in_dim_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   broadcast_dimensions = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-   // CHECK-SAME:   known_expanding_dimensions = #vhlo.tensor_v1<dense<> : tensor<0xi64>>,
-   // CHECK-SAME:   known_nonexpanding_dimensions = #vhlo.tensor_v1<dense<> : tensor<0xi64>>
-@@ -535,8 +577,9 @@
- }
- 
- // CHECK-LABEL: "default_dynamic_conv"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @default_dynamic_conv(%arg0: tensor<1x8x8x207xf32>, %arg1: tensor<3x3x207x16xf32>, %arg2: tensor<4xi32>) -> tensor<1x?x?x16xf32> {
--  //      CHECK: "vhlo.dynamic_conv_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.dynamic_conv_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   batch_group_count = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   feature_group_count = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   input_batch_dimension = #vhlo.integer_v1<0 : i64>,
-@@ -564,8 +607,9 @@
- }
- 
- // CHECK-LABEL: "default_dynamic_gather"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @default_dynamic_gather(%arg0 : tensor<2x4x9xf32>, %arg1 : tensor<1x5x2xi32>, %arg2 : tensor<3xi32>) -> tensor<1x5x8xf32> {
--  //      CHECK: "vhlo.dynamic_gather_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.dynamic_gather_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   collapsed_slice_dims = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-   // CHECK-SAME:   index_vector_dim = #vhlo.integer_v1<2 : i64>,
-   // CHECK-SAME:   indices_are_sorted = #vhlo.bool_v1<false>,
-@@ -591,15 +635,16 @@
-   // CHECK-SAME:   sym_name = #vhlo.string_v1<"default_func">,
-   // CHECK-SAME:   sym_visibility = #vhlo.string_v1<"">
-   // CHECK-SAME: }> ({
--  // CHECK-NEXT:   ^[[BB:bb.*]](%arg0: !vhlo.tensor_v1<!vhlo.f32_v1>):
--  // CHECK-NEXT:     "vhlo.return_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-+  // CHECK-NEXT:   ^[[BB:bb.*]](%[[ARG0:.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
-+  // CHECK-NEXT:     "vhlo.return_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-   // CHECK-NEXT: }) : () -> ()
-   func.return %arg0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "dynamic_gather"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @dynamic_gather(%arg0 : tensor<2x4x9xf32>, %arg1 : tensor<1x5x2xi32>) -> tensor<1x5x1xf32> {
--  //      CHECK: "vhlo.gather_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.gather_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   collapsed_slice_dims = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-   // CHECK-SAME:   index_vector_dim = #vhlo.integer_v1<2 : i64>,
-   // CHECK-SAME:   indices_are_sorted = #vhlo.bool_v1<false>,
-@@ -620,8 +665,9 @@
- }
- 
- // CHECK-LABEL: "default_infeed"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_infeed(%arg0: !stablehlo.token) -> (tensor<f32>, !stablehlo.token) {
--  //               CHECK: "vhlo.infeed_v1"(%arg0) <{
-+  //               CHECK: "vhlo.infeed_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   infeed_config = #vhlo.string_v1<"">,
-   // CHECK-SAME{LITERAL}:   layout = #vhlo.array_v1<[]>
-   //          CHECK-SAME: }> : (!vhlo.token_v1) -> (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.token_v1)
-@@ -630,8 +676,9 @@
- }
- 
- // CHECK-LABEL: "default_outfeed"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @default_outfeed(%arg0: tensor<f32>, %arg1: !stablehlo.token) -> !stablehlo.token {
--  //      CHECK: "vhlo.outfeed_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.outfeed_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   outfeed_config = #vhlo.string_v1<"">
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.token_v1) -> !vhlo.token_v1
-   %0 = "stablehlo.outfeed"(%arg0, %arg1) : (tensor<f32>, !stablehlo.token) -> !stablehlo.token
-@@ -639,8 +686,9 @@
- }
- 
- // CHECK-LABEL: "default_recv"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_recv(%arg0: !stablehlo.token) -> (tensor<f32>, !stablehlo.token) {
--  //      CHECK: "vhlo.recv_v1"(%arg0) <{
-+  //      CHECK: "vhlo.recv_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME:   channel_type = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   is_host_transfer = #vhlo.bool_v1<false>
-@@ -652,8 +700,9 @@
- }
- 
- // CHECK-LABEL: "default_send"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @default_send(%arg0: tensor<f32>, %arg1: !stablehlo.token) -> !stablehlo.token {
--  //      CHECK: "vhlo.send_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.send_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME:   channel_type = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   is_host_transfer = #vhlo.bool_v1<false>
-@@ -665,8 +714,9 @@
- }
- 
- // CHECK-LABEL: "default_reduce_scatter"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_reduce_scatter(%arg0: tensor<16xf32>) -> tensor<16xf32> {
--  //               CHECK: "vhlo.reduce_scatter_v1"(%arg0) <{
-+  //               CHECK: "vhlo.reduce_scatter_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0], [1]]> : tensor<2x1xi64>>,
-   //          CHECK-SAME:   scatter_dimension = #vhlo.integer_v1<0 : i64>
-@@ -688,8 +738,9 @@
- }
- 
- // CHECK-LABEL: "default_reduce_window"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @default_reduce_window(%arg0: tensor<2x17x31x7xf32>, %arg1: tensor<f32>) -> tensor<2x16x30x7xf32> {
--  //               CHECK: "vhlo.reduce_window_v1"(%arg0, %arg1)  <{
-+  //               CHECK: "vhlo.reduce_window_v1"(%[[ARG0]], %[[ARG1]])  <{
-   //          CHECK-SAME:   base_dilations = #vhlo.tensor_v1<dense<1> : tensor<4xi64>>,
-   // CHECK-SAME{LITERAL}:   padding = #vhlo.tensor_v1<dense<0> : tensor<4x2xi64>>,
-   //          CHECK-SAME:   window_dilations = #vhlo.tensor_v1<dense<1> : tensor<4xi64>>,
-@@ -711,8 +762,9 @@
- }
- 
- // CHECK-LABEL: "default_scatter"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @default_scatter(%arg0: tensor<200x100x300xf32>, %arg1: tensor<10x2xi32>, %arg2: tensor<10x300xf32>) -> tensor<200x100x300xf32> {
--  //      CHECK: "vhlo.scatter_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.scatter_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   index_vector_dim = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   indices_are_sorted = #vhlo.bool_v1<false>,
-   // CHECK-SAME:   inserted_window_dims = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-@@ -740,8 +792,9 @@
- }
- 
- // CHECK-LABEL: "default_select_and_scatter"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @default_select_and_scatter(%arg0: tensor<10x24x24x64xf32>, %arg1: tensor<10x23x23x64xf32>, %arg2: tensor<f32>) -> tensor<10x24x24x64xf32> {
--  //      CHECK: "vhlo.select_and_scatter_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.select_and_scatter_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   padding = #vhlo.tensor_v1<dense<0> : tensor<4x2xi64>>,
-   // CHECK-SAME:   window_dimensions = #vhlo.tensor_v1<dense<[1, 2, 2, 1]> : tensor<4xi64>>,
-   // CHECK-SAME:   window_strides = #vhlo.tensor_v1<dense<1> : tensor<4xi64>>
-@@ -769,8 +822,9 @@
- }
- 
- // CHECK-LABEL: "default_sort"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @default_sort(%arg0: tensor<16xf32>) -> tensor<16xf32> {
--  //      CHECK: "vhlo.sort_v1"(%arg0) <{
-+  //      CHECK: "vhlo.sort_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   dimension = #vhlo.integer_v1<-1 : i64>
-   // CHECK-SAME:   is_stable = #vhlo.bool_v1<false>
-   // CHECK-SAME: }> ({
-@@ -789,29 +843,33 @@
- // ============ OPS ============
- 
- // CHECK-LABEL: "op_abs"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_abs(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.abs_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.abs_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.abs"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_add"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_add(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_after_all"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_after_all(%arg0: !stablehlo.token) -> !stablehlo.token {
--  // CHECK: "vhlo.after_all_v1"(%arg0) : (!vhlo.token_v1) -> !vhlo.token_v1
-+  // CHECK: "vhlo.after_all_v1"(%[[ARG0]]) : (!vhlo.token_v1) -> !vhlo.token_v1
-   %0 = "stablehlo.after_all"(%arg0) : (!stablehlo.token) -> !stablehlo.token
-   func.return %0 : !stablehlo.token
- }
- 
- // CHECK-LABEL: "op_all_gather"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_all_gather(%arg0: tensor<16x8xf32>) -> tensor<16x16xf32> {
--  //               CHECK: "vhlo.all_gather_v1"(%arg0) <{
-+  //               CHECK: "vhlo.all_gather_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   all_gather_dim = #vhlo.integer_v1<1 : i64>
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0], [1]]> : tensor<2x1xi64>>,
-@@ -827,8 +885,9 @@
- }
- 
- // CHECK-LABEL: "op_all_reduce"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_all_reduce(%arg0: tensor<f32>) -> tensor<f32> {
--  //               CHECK: "vhlo.all_reduce_v1"(%arg0) <{
-+  //               CHECK: "vhlo.all_reduce_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0], [1]]> : tensor<2x1xi64>>,
-   //          CHECK-SAME:   use_global_device_ids = #vhlo.bool_v1<true>
-@@ -851,7 +910,7 @@
- 
- // CHECK-LABEL: "op_all_reduce_with_promotable_types"
- func.func @op_all_reduce_with_promotable_types(%operand: tensor<f32>) -> tensor<f64> {
--  //  CHECK: "vhlo.all_reduce_v1"(%arg0)
-+  //  CHECK: "vhlo.all_reduce_v1"(%[[ARG0:.*]])
-   //  CHECK:   ^[[BB:bb.*]](%[[ARG1:arg.*]]: !vhlo.tensor_v1<!vhlo.f64_v1>, %[[ARG2:arg.*]]: !vhlo.tensor_v1<!vhlo.f64_v1>):
-   //  CHECK:     "vhlo.return_v1"(%[[VAL1:.*]]) : (!vhlo.tensor_v1<!vhlo.f64_v1>) -> ()
-   //  CHECK: }) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f64_v1>
-@@ -869,8 +928,9 @@
- }
- 
- // CHECK-LABEL: "op_all_to_all"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_all_to_all(%arg0: tensor<4x16xf32>) -> tensor<16x4xf32> {
--  //               CHECK: "vhlo.all_to_all_v1"(%arg0) <{
-+  //               CHECK: "vhlo.all_to_all_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<1 : i64>,
-   //          CHECK-SAME:   concat_dimension = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0, 1, 2, 3]]> : tensor<1x4xi64>>,
-@@ -888,22 +948,25 @@
- }
- 
- // CHECK-LABEL: "op_and"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_and(%arg0: tensor<i1>, %arg1: tensor<i1>) -> tensor<i1> {
--  // CHECK: "vhlo.and_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-+  // CHECK: "vhlo.and_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-   %0 = "stablehlo.and"(%arg0, %arg1) : (tensor<i1>, tensor<i1>) -> tensor<i1>
-   func.return %0 : tensor<i1>
- }
- 
- // CHECK-LABEL: "op_atan2"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_atan2(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.atan2_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.atan2_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.atan2"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_batch_norm_grad"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}}, %[[ARG3:.*]]: {{.*}}, %[[ARG4:.*]]: {{.*}})
- func.func @op_batch_norm_grad(%arg0: tensor<16x16x16x16xf32>, %arg1: tensor<16xf32>, %arg2: tensor<16xf32>, %arg3: tensor<16xf32>, %arg4: tensor<16x16x16x16xf32>) -> (tensor<16x16x16x16xf32>, tensor<16xf32>, tensor<16xf32>) {
--  //      CHECK: "vhlo.batch_norm_grad_v1"(%arg0, %arg1, %arg2, %arg3, %arg4) <{
-+  //      CHECK: "vhlo.batch_norm_grad_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]], %[[ARG3]], %[[ARG4]]) <{
-   // CHECK-SAME:   epsilon = #vhlo.float_v1<1.000000e-03 : !vhlo.f32_v1>,
-   // CHECK-SAME:   feature_index = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<16x16x16x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x16x16x16x!vhlo.f32_v1>) -> (!vhlo.tensor_v1<16x16x16x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>)
-@@ -915,8 +978,9 @@
- }
- 
- // CHECK-LABEL: "op_batch_norm_inference"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}}, %[[ARG3:.*]]: {{.*}}, %[[ARG4:.*]]: {{.*}})
- func.func @op_batch_norm_inference(%arg0: tensor<16x16x16x16xf32>, %arg1: tensor<16xf32>, %arg2: tensor<16xf32>, %arg3: tensor<16xf32>, %arg4: tensor<16xf32>) -> tensor<16x16x16x16xf32> {
--  //      CHECK: "vhlo.batch_norm_inference_v1"(%arg0, %arg1, %arg2, %arg3, %arg4) <{
-+  //      CHECK: "vhlo.batch_norm_inference_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]], %[[ARG3]], %[[ARG4]]) <{
-   // CHECK-SAME:   epsilon = #vhlo.float_v1<1.000000e-03 : !vhlo.f32_v1>,
-   // CHECK-SAME:   feature_index = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<16x16x16x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x16x16x16x!vhlo.f32_v1>
-@@ -928,8 +992,9 @@
- }
- 
- // CHECK-LABEL: "op_batch_norm_training"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_batch_norm_training(%arg0: tensor<16x16x16x16xf32>, %arg1: tensor<16xf32>, %arg2: tensor<16xf32>) -> (tensor<16x16x16x16xf32>, tensor<16xf32>, tensor<16xf32>) {
--  //      CHECK: "vhlo.batch_norm_training_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.batch_norm_training_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   epsilon = #vhlo.float_v1<1.000000e-03 : !vhlo.f32_v1>,
-   // CHECK-SAME:   feature_index = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<16x16x16x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>) -> (!vhlo.tensor_v1<16x16x16x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>)
-@@ -941,15 +1006,17 @@
- }
- 
- // CHECK-LABEL: "op_bitcast_convert"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_bitcast_convert(%arg0: tensor<i32>) -> tensor<f32> {
--  // CHECK: "vhlo.bitcast_convert_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.bitcast_convert_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.bitcast_convert"(%arg0) : (tensor<i32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_broadcast_in_dim"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_broadcast_in_dim(%arg0: tensor<16xf32>) -> tensor<16x16xf32> {
--  //      CHECK: "vhlo.broadcast_in_dim_v1"(%arg0) <{
-+  //      CHECK: "vhlo.broadcast_in_dim_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   broadcast_dimensions = #vhlo.tensor_v1<dense<1> : tensor<1xi64>>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x16x!vhlo.f32_v1>
-   %0 = "stablehlo.broadcast_in_dim"(%arg0) {
-@@ -959,8 +1026,9 @@
- }
- 
- // CHECK-LABEL: "op_broadcast"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_broadcast(%arg0: tensor<16xf32>) -> tensor<16x16xf32> {
--  //      CHECK: "vhlo.broadcast_v1"(%arg0) <{
-+  //      CHECK: "vhlo.broadcast_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   broadcast_sizes = #vhlo.tensor_v1<dense<16> : tensor<1xi64>>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x16x!vhlo.f32_v1>
-   %0 = "stablehlo.broadcast"(%arg0) {
-@@ -970,9 +1038,10 @@
- }
- 
- // CHECK-LABEL: "op_case"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_case(%arg0: tensor<i32>, %arg1: tensor<f32>) -> tensor<f32> {
--  //      CHECK: "vhlo.case_v1"(%arg0) ({
--  // CHECK-NEXT:   "vhlo.return_v1"(%arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-+  //      CHECK: "vhlo.case_v1"(%[[ARG0]]) ({
-+  // CHECK-NEXT:   "vhlo.return_v1"(%[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-   // CHECK-NEXT: }) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.case"(%arg0) ({
-     "stablehlo.return"(%arg1) : (tensor<f32>) -> ()
-@@ -981,22 +1050,25 @@
- }
- 
- // CHECK-LABEL: "op_cbrt"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_cbrt(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.cbrt_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.cbrt_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.cbrt"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_ceil"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_ceil(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.ceil_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.ceil_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.ceil"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_cholesky"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_cholesky(%arg0: tensor<1x16x16xf32>) -> tensor<1x16x16xf32> {
--  //      CHECK: "vhlo.cholesky_v1"(%arg0) <{
-+  //      CHECK: "vhlo.cholesky_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   lower = #vhlo.bool_v1<true>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<1x16x16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<1x16x16x!vhlo.f32_v1>
-   %0 = "stablehlo.cholesky"(%arg0) {
-@@ -1006,22 +1078,25 @@
- }
- 
- // CHECK-LABEL: "op_clamp"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_clamp(%arg0: tensor<f32>, %arg1: tensor<f32>, %arg2: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.clamp_v1"(%arg0, %arg1, %arg2) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.clamp_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.clamp"(%arg0, %arg1, %arg2) : (tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_count_leading_zeros"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_count_leading_zeros(%arg0: tensor<i32>) -> tensor<i32> {
--  // CHECK: "vhlo.count_leading_zeros_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-+  // CHECK: "vhlo.count_leading_zeros_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-   %0 = "stablehlo.count_leading_zeros"(%arg0) : (tensor<i32>) -> tensor<i32>
-   func.return %0 : tensor<i32>
- }
- 
- // CHECK-LABEL: "op_collective_permute"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_collective_permute(%arg0: tensor<16x8xf32>) -> tensor<16x8xf32> {
--  //               CHECK: "vhlo.collective_permute_v1"(%arg0) <{
-+  //               CHECK: "vhlo.collective_permute_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME{LITERAL}:   source_target_pairs = #vhlo.tensor_v1<dense<[[0, 1], [1, 2], [2, 3]]> : tensor<3x2xi64>>
-   //          CHECK-SAME: }> : (!vhlo.tensor_v1<16x8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x8x!vhlo.f32_v1>
-@@ -1033,8 +1108,9 @@
- }
- 
- // CHECK-LABEL: "op_compare"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_compare(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
--  //      CHECK: "vhlo.compare_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.compare_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   compare_type = #vhlo<comparison_type_v1 TOTALORDER>,
-   // CHECK-SAME:   comparison_direction = #vhlo<comparison_direction_v1 EQ>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-@@ -1046,15 +1122,17 @@
- }
- 
- // CHECK-LABEL: "op_complex"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_complex(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<complex<f32>> {
--  // CHECK: "vhlo.complex_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>
-+  // CHECK: "vhlo.complex_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>
-   %0 = "stablehlo.complex"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<complex<f32>>
-   func.return %0 : tensor<complex<f32>>
- }
- 
- // CHECK-LABEL: "op_composite"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_composite(%arg0: tensor<f32>) -> tensor<f32> {
--  //               CHECK: "vhlo.composite_v1"(%arg0) <{
-+  //               CHECK: "vhlo.composite_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   composite_attributes = #vhlo.dict_v1<{#vhlo.string_v1<"my_int"> = #vhlo.integer_v1<1 : i64>, #vhlo.string_v1<"my_string"> = #vhlo.string_v1<"foo">}>
-   //          CHECK-SAME:   decomposition = #vhlo.string_v1<"composite_target">
-   //          CHECK-SAME:   name = #vhlo.string_v1<"stablehlo.composite_target">
-@@ -1073,8 +1151,9 @@
- }
- 
- // CHECK-LABEL: "op_concatenate"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_concatenate(%arg0: tensor<8xf32>, %arg1: tensor<8xf32>) -> tensor<16xf32> {
--  //      CHECK: "vhlo.concatenate_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.concatenate_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   dimension = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<8x!vhlo.f32_v1>, !vhlo.tensor_v1<8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x!vhlo.f32_v1>
-   %0 = "stablehlo.concatenate"(%arg0, %arg1) {
-@@ -1084,6 +1163,7 @@
- }
- 
- // CHECK-LABEL: "op_constant"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_constant(%arg0: tensor<f32>) -> tensor<f32> {
-   //      CHECK: "vhlo.constant_v1"() <{
-   // CHECK-SAME:   value = #vhlo.tensor_v1<dense<0.000000e+00> : tensor<f32>>
-@@ -1095,15 +1175,17 @@
- }
- 
- // CHECK-LABEL: "op_convert"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_convert(%arg0: tensor<i32>) -> tensor<f32> {
--  // CHECK: "vhlo.convert_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.convert_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.convert"(%arg0) : (tensor<i32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_convolution"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_convolution(%arg0: tensor<1x8x8x207xf32>, %arg1: tensor<3x3x207x16xf32>) -> tensor<1x7x7x16xf32> {
--  //      CHECK: "vhlo.convolution_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.convolution_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   batch_group_count = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   feature_group_count = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   input_batch_dimension = #vhlo.integer_v1<0 : i64>,
-@@ -1137,8 +1219,9 @@
- }
- 
- // CHECK-LABEL: "op_cosine"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_cosine(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.cosine_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.cosine_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.cosine"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
-@@ -1151,8 +1234,9 @@
- }
- 
- // CHECK-LABEL: "op_cross_replica_sum"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_cross_replica_sum(%arg0: tensor<f32>) -> tensor<f32> {
--  //               CHECK: "vhlo.cross-replica-sum_v1"(%arg0) <{
-+  //               CHECK: "vhlo.cross-replica-sum_v1"(%[[ARG0]]) <{
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0], [1]]> : tensor<2x1xi64>>
-   //          CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.cross-replica-sum"(%arg0) {
-@@ -1162,8 +1246,9 @@
- }
- 
- // CHECK-LABEL: "op_custom_call"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_custom_call(%arg0: tensor<f32>) -> tensor<f32> {
--  //      CHECK: "vhlo.custom_call_v1"(%arg0) <{
-+  //      CHECK: "vhlo.custom_call_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   api_version = #vhlo<api_version_v1 API_VERSION_STATUS_RETURNING>,
-   // CHECK-SAME:   backend_config = #vhlo.string_v1<"\08\03\1A\02">,
-   // CHECK-SAME:   call_target_name = #vhlo.string_v1<"foo">,
-@@ -1194,15 +1279,17 @@
- }
- 
- // CHECK-LABEL: "op_divide"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_divide(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.divide_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.divide_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.divide"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_dot_general"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_dot_general(%arg0: tensor<8x8x16xf32>, %arg1: tensor<8x16x8xf32>) -> tensor<8x8x8xf32> {
--  //      CHECK: "vhlo.dot_general_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.dot_general_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   lhs_batching_dimensions = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>,
-   // CHECK-SAME:   lhs_contracting_dimensions = #vhlo.tensor_v1<dense<2> : tensor<1xi64>>,
-   // CHECK-SAME:   precision_config = #vhlo.array_v1<[#vhlo<precision_v1 HIGHEST>, #vhlo<precision_v1 HIGHEST>]>,
-@@ -1222,8 +1309,9 @@
- }
- 
- // CHECK-LABEL: "op_dot"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_dot(%arg0: tensor<8x16xf32>, %arg1: tensor<16x8xf32>) -> tensor<8x8xf32> {
--  //      CHECK: "vhlo.dot_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.dot_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   precision_config = #vhlo.array_v1<[#vhlo<precision_v1 HIGHEST>, #vhlo<precision_v1 HIGHEST>]>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<8x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<8x8x!vhlo.f32_v1>
-   %0 = "stablehlo.dot"(%arg0, %arg1) {
-@@ -1233,8 +1321,9 @@
- }
- 
- // CHECK-LABEL: "op_dynamic_broadcast_in_dim"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_dynamic_broadcast_in_dim(%arg0: tensor<?x?xf32>, %arg1: tensor<2xindex>) -> tensor<?x?xf32> {
--  //      CHECK: "vhlo.dynamic_broadcast_in_dim_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.dynamic_broadcast_in_dim_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   broadcast_dimensions = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-   // CHECK-SAME:   known_expanding_dimensions = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>,
-   // CHECK-SAME:   known_nonexpanding_dimensions = #vhlo.tensor_v1<dense<1> : tensor<1xi64>>
-@@ -1248,8 +1337,9 @@
- }
- 
- // CHECK-LABEL: "op_dynamic_conv"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_dynamic_conv(%arg0: tensor<1x8x8x207xf32>, %arg1: tensor<3x3x207x16xf32>, %arg2: tensor<4xi32>) -> tensor<1x?x?x16xf32> {
--  //      CHECK: "vhlo.dynamic_conv_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.dynamic_conv_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   batch_group_count = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   feature_group_count = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   input_batch_dimension = #vhlo.integer_v1<0 : i64>,
-@@ -1283,8 +1373,9 @@
- }
- 
- // CHECK-LABEL: "op_dynamic_gather"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_dynamic_gather(%arg0 : tensor<2x4x9xf32>, %arg1 : tensor<1x5x2xi32>, %arg2 : tensor<3xi32>) -> tensor<1x5x8xf32> {
--  //      CHECK: "vhlo.dynamic_gather_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.dynamic_gather_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   collapsed_slice_dims = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-   // CHECK-SAME:   index_vector_dim = #vhlo.integer_v1<2 : i64>,
-   // CHECK-SAME:   indices_are_sorted = #vhlo.bool_v1<true>,
-@@ -1304,8 +1395,9 @@
- }
- 
- // CHECK-LABEL: "op_dynamic_iota"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_dynamic_iota(%arg0: tensor<1xindex>) -> tensor<?xf32> {
--  //      CHECK: "vhlo.dynamic_iota_v1"(%arg0) <{
-+  //      CHECK: "vhlo.dynamic_iota_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   iota_dimension = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<1x!vhlo.index_v1>) -> !vhlo.tensor_v1<?x!vhlo.f32_v1>
-   %0 = "stablehlo.dynamic_iota"(%arg0) {
-@@ -1315,22 +1407,25 @@
- }
- 
- // CHECK-LABEL: "op_dynamic_pad"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}}, %[[ARG3:.*]]: {{.*}}, %[[ARG4:.*]]: {{.*}})
- func.func @op_dynamic_pad(%arg0: tensor<?xf32>, %arg1: tensor<f32>, %arg2: tensor<1xindex>, %arg3: tensor<1xindex>, %arg4: tensor<1xindex>) -> tensor<?xf32> {
--  // CHECK: "vhlo.dynamic_pad_v1"(%arg0, %arg1, %arg2, %arg3, %arg4) : (!vhlo.tensor_v1<?x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>) -> !vhlo.tensor_v1<?x!vhlo.f32_v1>
-+  // CHECK: "vhlo.dynamic_pad_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]], %[[ARG3]], %[[ARG4]]) : (!vhlo.tensor_v1<?x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>) -> !vhlo.tensor_v1<?x!vhlo.f32_v1>
-   %0 = "stablehlo.dynamic_pad"(%arg0, %arg1, %arg2, %arg3, %arg4) : (tensor<?xf32>, tensor<f32>, tensor<1xindex>, tensor<1xindex>, tensor<1xindex>) -> tensor<?xf32>
-   func.return %0 : tensor<?xf32>
- }
- 
- // CHECK-LABEL: "op_dynamic_reshape"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_dynamic_reshape(%arg0: tensor<16xf32>, %arg1: tensor<2xindex>) -> tensor<?x?xf32> {
--  // CHECK: "vhlo.dynamic_reshape_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<2x!vhlo.index_v1>) -> !vhlo.tensor_v1<?x?x!vhlo.f32_v1>
-+  // CHECK: "vhlo.dynamic_reshape_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<2x!vhlo.index_v1>) -> !vhlo.tensor_v1<?x?x!vhlo.f32_v1>
-   %0 = "stablehlo.dynamic_reshape"(%arg0, %arg1) : (tensor<16xf32>, tensor<2xindex>) -> tensor<?x?xf32>
-   func.return %0 : tensor<?x?xf32>
- }
- 
- // CHECK-LABEL: "op_dynamic_slice"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_dynamic_slice(%arg0: tensor<16xf32>, %arg1: tensor<i64>) -> tensor<4xf32> {
--  //      CHECK: "vhlo.dynamic_slice_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.dynamic_slice_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   slice_sizes = #vhlo.tensor_v1<dense<4> : tensor<1xi64>>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.i64_v1>) -> !vhlo.tensor_v1<4x!vhlo.f32_v1>
-   %0 = "stablehlo.dynamic_slice"(%arg0, %arg1) {
-@@ -1340,15 +1435,17 @@
- }
- 
- // CHECK-LABEL: "op_dynamic_update_slice"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_dynamic_update_slice(%arg0: tensor<16xf32>, %arg1: tensor<4xf32>, %arg2: tensor<i64>) -> tensor<16xf32> {
--  // CHECK: "vhlo.dynamic_update_slice_v1"(%arg0, %arg1, %arg2) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<4x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.i64_v1>) -> !vhlo.tensor_v1<16x!vhlo.f32_v1>
-+  // CHECK: "vhlo.dynamic_update_slice_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<4x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.i64_v1>) -> !vhlo.tensor_v1<16x!vhlo.f32_v1>
-   %0 = "stablehlo.dynamic_update_slice"(%arg0, %arg1, %arg2) : (tensor<16xf32>, tensor<4xf32>, tensor<i64>) -> tensor<16xf32>
-   func.return %0 : tensor<16xf32>
- }
- 
- // CHECK-LABEL: "op_einsum"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_einsum(%arg0: tensor<8x16xf32>, %arg1: tensor<16x8xf32>) -> tensor<8x8xf32> {
--  //      CHECK: "vhlo.einsum_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.einsum_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   einsum_config = #vhlo.string_v1<"ab,bc->ac">
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<8x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<8x8x!vhlo.f32_v1>
-   %0 = "stablehlo.einsum"(%arg0, %arg1) {
-@@ -1358,22 +1455,25 @@
- }
- 
- // CHECK-LABEL: "op_exponential_minus_one"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_exponential_minus_one(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.exponential_minus_one_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.exponential_minus_one_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.exponential_minus_one"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_exponential"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_exponential(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.exponential_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.exponential_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.exponential"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_fft"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_fft(%arg0: tensor<16xcomplex<f32>>) -> tensor<16xcomplex<f32>> {
--  //      CHECK: "vhlo.fft_v1"(%arg0) <{
-+  //      CHECK: "vhlo.fft_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   fft_length = #vhlo.tensor_v1<dense<16> : tensor<1xi64>>,
-   // CHECK-SAME:   fft_type = #vhlo<fft_type_v1 FFT>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<16x!vhlo.complex_v1<!vhlo.f32_v1>>) -> !vhlo.tensor_v1<16x!vhlo.complex_v1<!vhlo.f32_v1>>
-@@ -1385,8 +1485,9 @@
- }
- 
- // CHECK-LABEL: "op_floor"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_floor(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.floor_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.floor_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.floor"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
-@@ -1399,16 +1500,17 @@
-   // CHECK-SAME:   sym_name = #vhlo.string_v1<"op_func">,
-   // CHECK-SAME:   sym_visibility = #vhlo.string_v1<"private">
-   // CHECK-SAME: }> ({
--  // CHECK-NEXT:   ^[[BB:bb.*]](%arg0: !vhlo.tensor_v1<!vhlo.f32_v1>):
--  // CHECK-NEXT:     "vhlo.return_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-+  // CHECK-NEXT:   ^[[BB:bb.*]](%[[ARG0:.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
-+  // CHECK-NEXT:     "vhlo.return_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-   // CHECK-NEXT: }) : () -> ()
- 
-   func.return %arg0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_gather"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_gather(%arg0 : tensor<2x4x9xf32>, %arg1 : tensor<1x5x2xi32>) -> tensor<1x5x1xf32> {
--  //      CHECK: "vhlo.gather_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.gather_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   collapsed_slice_dims = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-   // CHECK-SAME:   index_vector_dim = #vhlo.integer_v1<2 : i64>,
-   // CHECK-SAME:   indices_are_sorted = #vhlo.bool_v1<true>,
-@@ -1430,8 +1532,9 @@
- }
- 
- // CHECK-LABEL: "op_get_dimension_size"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_get_dimension_size(%arg0: tensor<?xf32>) -> tensor<i32> {
--  //      CHECK: "vhlo.get_dimension_size_v1"(%arg0) <{
-+  //      CHECK: "vhlo.get_dimension_size_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   dimension = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<?x!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-   %0 = "stablehlo.get_dimension_size"(%arg0) {
-@@ -1441,8 +1544,9 @@
- }
- 
- // CHECK-LABEL: "op_get_tuple_element"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_get_tuple_element(%arg0: tuple<tensor<f32>, tensor<i32>>) -> tensor<f32> {
--  //      CHECK: "vhlo.get_tuple_element_v1"(%arg0) <{
-+  //      CHECK: "vhlo.get_tuple_element_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   index = #vhlo.integer_v1<0 : i32>
-   // CHECK-SAME: }> : (!vhlo.tuple_v1<!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.get_tuple_element"(%arg0) {
-@@ -1452,11 +1556,12 @@
- }
- 
- // CHECK-LABEL: "op_if"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_if(%arg0: tensor<i1>, %arg1: tensor<f32>, %arg2: tensor<f32>) -> tensor<f32> {
--  //      CHECK: "vhlo.if_v1"(%arg0) ({
--  // CHECK-NEXT:   "vhlo.return_v1"(%arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-+  //      CHECK: "vhlo.if_v1"(%[[ARG0]]) ({
-+  // CHECK-NEXT:   "vhlo.return_v1"(%[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-   // CHECK-NEXT: }, {
--  // CHECK-NEXT:   "vhlo.return_v1"(%arg2) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-+  // CHECK-NEXT:   "vhlo.return_v1"(%[[ARG2]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-   // CHECK-NEXT: }) : (!vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.if"(%arg0) ({
-     "stablehlo.return"(%arg1) : (tensor<f32>) -> ()
-@@ -1467,15 +1572,17 @@
- }
- 
- // CHECK-LABEL: "op_imag"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_imag(%arg0: tensor<complex<f32>>) -> tensor<f32> {
--  // CHECK: "vhlo.imag_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.imag_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.imag"(%arg0) : (tensor<complex<f32>>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_infeed"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_infeed(%arg0: !stablehlo.token) -> (tensor<f32>, !stablehlo.token) {
--  //               CHECK: "vhlo.infeed_v1"(%arg0) <{
-+  //               CHECK: "vhlo.infeed_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   infeed_config = #vhlo.string_v1<"foo">,
-   // CHECK-SAME{LITERAL}:   layout = #vhlo.array_v1<[#vhlo.array_v1<[]>]>
-   //          CHECK-SAME: }> : (!vhlo.token_v1) -> (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.token_v1)
-@@ -1498,36 +1605,41 @@
- }
- 
- // CHECK-LABEL: "op_is_finite"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_is_finite(%arg0: tensor<f32>) -> tensor<i1> {
--  // CHECK: "vhlo.is_finite_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-+  // CHECK: "vhlo.is_finite_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-   %0 = "stablehlo.is_finite"(%arg0) : (tensor<f32>) -> tensor<i1>
-   func.return %0 : tensor<i1>
- }
- 
- // CHECK-LABEL: "op_log"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_log(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.log_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.log_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.log"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_log_plus_one"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_log_plus_one(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.log_plus_one_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.log_plus_one_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.log_plus_one"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_logistic"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_logistic(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.logistic_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.logistic_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.logistic"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_map"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_map(%arg0: tensor<16xf32>) -> tensor<16xf32> {
--  //      CHECK: "vhlo.map_v1"(%arg0) <{
-+  //      CHECK: "vhlo.map_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   dimensions = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>
-   // CHECK-SAME: }> ({
-   // CHECK-NEXT:   ^[[BB:bb.*]](%[[ARG1:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
-@@ -1545,57 +1657,65 @@
- }
- 
- // CHECK-LABEL: "op_maximum"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_maximum(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.maximum_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.maximum_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.maximum"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_minimum"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_minimum(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.minimum_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.minimum_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.minimum"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_multiply"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_multiply(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.multiply_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.multiply_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.multiply"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_negate"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_negate(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.negate_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.negate_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.negate"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_not"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_not(%arg0: tensor<i1>) -> tensor<i1> {
--  // CHECK: "vhlo.not_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-+  // CHECK: "vhlo.not_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-   %0 = "stablehlo.not"(%arg0) : (tensor<i1>) -> tensor<i1>
-   func.return %0 : tensor<i1>
- }
- 
- // CHECK-LABEL: "op_optimization_barrier"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_optimization_barrier(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.optimization_barrier_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.optimization_barrier_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.optimization_barrier"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_or"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_or(%arg0: tensor<i1>, %arg1: tensor<i1>) -> tensor<i1> {
--  // CHECK: "vhlo.or_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-+  // CHECK: "vhlo.or_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-   %0 = "stablehlo.or"(%arg0, %arg1) : (tensor<i1>, tensor<i1>) -> tensor<i1>
-   func.return %0 : tensor<i1>
- }
- 
- // CHECK-LABEL: "op_outfeed"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_outfeed(%arg0: tensor<f32>, %arg1: !stablehlo.token) -> !stablehlo.token {
--  //      CHECK: "vhlo.outfeed_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.outfeed_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   outfeed_config = #vhlo.string_v1<"foo">
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.token_v1) -> !vhlo.token_v1
-   %0 = "stablehlo.outfeed"(%arg0, %arg1) {
-@@ -1605,8 +1725,9 @@
- }
- 
- // CHECK-LABEL: "op_pad"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_pad(%arg0: tensor<8xf32>, %arg1: tensor<f32>) -> tensor<16xf32> {
--  //      CHECK: "vhlo.pad_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.pad_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   edge_padding_high = #vhlo.tensor_v1<dense<4> : tensor<1xi64>>,
-   // CHECK-SAME:   edge_padding_low = #vhlo.tensor_v1<dense<4> : tensor<1xi64>>,
-   // CHECK-SAME:   interior_padding = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>
-@@ -1620,36 +1741,41 @@
- }
- 
- // CHECK-LABEL: "op_popcnt"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_popcnt(%arg0: tensor<i32>) -> tensor<i32> {
--  // CHECK: "vhlo.popcnt_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-+  // CHECK: "vhlo.popcnt_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-   %0 = "stablehlo.popcnt"(%arg0) : (tensor<i32>) -> tensor<i32>
-   func.return %0 : tensor<i32>
- }
- 
- // CHECK-LABEL: "op_power"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_power(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.power_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.power_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.power"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_real_dynamic_slice"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}}, %[[ARG3:.*]]: {{.*}})
- func.func @op_real_dynamic_slice(%arg0: tensor<?xf32>, %arg1: tensor<1xindex>, %arg2: tensor<1xindex>, %arg3: tensor<1xindex>) -> tensor<?xf32> {
--  // CHECK: "vhlo.real_dynamic_slice_v1"(%arg0, %arg1, %arg2, %arg3) : (!vhlo.tensor_v1<?x!vhlo.f32_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>) -> !vhlo.tensor_v1<?x!vhlo.f32_v1>
-+  // CHECK: "vhlo.real_dynamic_slice_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]], %[[ARG3]]) : (!vhlo.tensor_v1<?x!vhlo.f32_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>) -> !vhlo.tensor_v1<?x!vhlo.f32_v1>
-   %0 = "stablehlo.real_dynamic_slice"(%arg0, %arg1, %arg2, %arg3) : (tensor<?xf32>, tensor<1xindex>, tensor<1xindex>, tensor<1xindex>) -> tensor<?xf32>
-   func.return %0 : tensor<?xf32>
- }
- 
- // CHECK-LABEL: "op_real"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_real(%arg0: tensor<complex<f32>>) -> tensor<f32> {
--  // CHECK: "vhlo.real_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.real_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.real"(%arg0) : (tensor<complex<f32>>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_recv"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_recv(%arg0: !stablehlo.token) -> (tensor<f32>, !stablehlo.token) {
--  //      CHECK: "vhlo.recv_v1"(%arg0) <{
-+  //      CHECK: "vhlo.recv_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME:   channel_type = #vhlo.integer_v1<3 : i64>,
-   // CHECK-SAME:   is_host_transfer = #vhlo.bool_v1<true>
-@@ -1662,8 +1788,9 @@
- }
- 
- // CHECK-LABEL: "op_reduce"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_reduce(%arg0: tensor<16xf32>, %arg1: tensor<f32>) -> tensor<f32> {
--  //  CHECK: "vhlo.reduce_v1"(%arg0, %arg1)
-+  //  CHECK: "vhlo.reduce_v1"(%[[ARG0]], %[[ARG1]])
-   //  CHECK:   ^[[BB:bb.*]](%[[ARG1:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>, %[[ARG2:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
-   //  CHECK:     "vhlo.return_v1"(%[[VAL1:.*]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-   //  CHECK: }) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-@@ -1678,8 +1805,9 @@
- }
- 
- // CHECK-LABEL: "op_reduce_precision"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_reduce_precision(%arg0: tensor<f32>) -> tensor<f32> {
--  //      CHECK: "vhlo.reduce_precision_v1"(%arg0) <{
-+  //      CHECK: "vhlo.reduce_precision_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   exponent_bits = #vhlo.integer_v1<8 : i32>
-   // CHECK-SAME:   mantissa_bits = #vhlo.integer_v1<10 : i32>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-@@ -1693,7 +1821,7 @@
- // CHECK_lABEL: "op_reduce_with_promotable_types"
- func.func @op_reduce_with_promotable_types(%arg0: tensor<4x4xf32>, %arg1 : tensor<f32>)
-     -> (tensor<4xf64>) {
--  //  CHECK: "vhlo.reduce_v1"(%arg0, %arg1)
-+  //  CHECK: "vhlo.reduce_v1"(%[[ARG0:.*]], %[[ARG1:.*]])
-   //  CHECK:   ^[[BB:bb.*]](%[[ARG1:arg.*]]: !vhlo.tensor_v1<!vhlo.f64_v1>, %[[ARG2:arg.*]]: !vhlo.tensor_v1<!vhlo.f64_v1>):
-   //  CHECK:     "vhlo.return_v1"(%[[VAL1:.*]]) : (!vhlo.tensor_v1<!vhlo.f64_v1>) -> ()
-   //  CHECK: }) : (!vhlo.tensor_v1<4x4x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<4x!vhlo.f64_v1>
-@@ -1708,8 +1836,9 @@
- }
- 
- // CHECK-LABEL: "op_reduce_scatter"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_reduce_scatter(%arg0: tensor<16xf32>) -> tensor<16xf32> {
--  //               CHECK: "vhlo.reduce_scatter_v1"(%arg0) <{
-+  //               CHECK: "vhlo.reduce_scatter_v1"(%[[ARG0]]) <{
-   //          CHECK-SAME:   channel_id = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0], [1]]> : tensor<2x1xi64>>,
-   //          CHECK-SAME:   scatter_dimension = #vhlo.integer_v1<0 : i64>
-@@ -1734,7 +1863,7 @@
- 
- // CHECK_lABEL: "op_reduce_scatter_with_promotable_types"
- func.func @op_reduce_scatter_with_promotable_types(%data: tensor<4x16xf32>) -> tensor<4x4xf64> {
--  //  CHECK: "vhlo.reduce_scatter_v1"(%arg0)
-+  //  CHECK: "vhlo.reduce_scatter_v1"(%[[ARG0:.*]])
-   //  CHECK:   ^[[BB:bb.*]](%[[ARG1:arg.*]]: !vhlo.tensor_v1<!vhlo.f64_v1>, %[[ARG2:arg.*]]: !vhlo.tensor_v1<!vhlo.f64_v1>):
-   //  CHECK:     "vhlo.return_v1"(%[[VAL1:.*]]) : (!vhlo.tensor_v1<!vhlo.f64_v1>) -> ()
-   //  CHECK: }) : (!vhlo.tensor_v1<4x16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<4x4x!vhlo.f64_v1>
-@@ -1751,8 +1880,9 @@
- 
- 
- // CHECK-LABEL: "op_reduce_window"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_reduce_window(%arg0: tensor<2x17x31x7xf32>, %arg1: tensor<f32>) -> tensor<2x9x16x7xf32> {
--  //               CHECK: "vhlo.reduce_window_v1"(%arg0, %arg1) <{
-+  //               CHECK: "vhlo.reduce_window_v1"(%[[ARG0]], %[[ARG1]]) <{
-   //          CHECK-SAME:   base_dilations = #vhlo.tensor_v1<dense<[1, 2, 2, 1]> : tensor<4xi64>>,
-   // CHECK-SAME{LITERAL}:   padding = #vhlo.tensor_v1<dense<[[0, 0], [2, 0], [0, 2], [0, 0]]> : tensor<4x2xi64>>,
-   //          CHECK-SAME:   window_dilations = #vhlo.tensor_v1<dense<[1, 2, 2, 1]> : tensor<4xi64>>,
-@@ -1777,11 +1907,11 @@
-   func.return %0 : tensor<2x9x16x7xf32>
- }
- 
--// CHECK_lABEL: "op_reduce_window_with_promotable_types"
-+// CHECK-LABEL: "op_reduce_window_with_promotable_types"
- func.func @op_reduce_window_with_promotable_types(%arg0: tensor<4x2xf32>,
-     %arg1: tensor<4x2xf32>, %init0: tensor<f32>, %init1: tensor<f32>) ->
-     (tensor<2x2xf64>, tensor<2x2xf32>) {
--  //  CHECK: "vhlo.reduce_window_v1"(%arg0, %arg1, %arg2, %arg3)
-+  //  CHECK: "vhlo.reduce_window_v1"(%[[ARG0:.*]], %[[ARG1:.*]], %[[ARG2:.*]], %[[ARG3:.*]])
-   //  CHECK:   ^[[BB:bb.*]](%[[ARG1:arg.*]]: !vhlo.tensor_v1<!vhlo.f64_v1>, %[[ARG2:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>, %[[ARG3:arg.*]]: !vhlo.tensor_v1<!vhlo.f64_v1>, %[[ARG4:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
-   //  CHECK:     "vhlo.return_v1"(%[[VAL1:.*]], %[[VAL2:.*]]) : (!vhlo.tensor_v1<!vhlo.f64_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-   //  CHECK: }) : (!vhlo.tensor_v1<4x2x!vhlo.f32_v1>, !vhlo.tensor_v1<4x2x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> (!vhlo.tensor_v1<2x2x!vhlo.f64_v1>, !vhlo.tensor_v1<2x2x!vhlo.f32_v1>)
-@@ -1801,8 +1931,9 @@
- }
- 
- // CHECK-LABEL: "op_remainder"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_remainder(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.remainder_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.remainder_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.remainder"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
-@@ -1822,16 +1953,18 @@
- }
- 
- // CHECK-LABEL: "op_reshape"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_reshape(%arg0: tensor<16xf32>) -> tensor<4x4xf32> {
--  // CHECK: "vhlo.reshape_v1"(%arg0) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<4x4x!vhlo.f32_v1>
-+  // CHECK: "vhlo.reshape_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<4x4x!vhlo.f32_v1>
-   %0 = "stablehlo.reshape"(%arg0) : (tensor<16xf32>) -> tensor<4x4xf32>
-   func.return %0 : tensor<4x4xf32>
- }
- 
- // CHECK-LABEL: "op_return"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_return(%arg0: tensor<i32>, %arg1: tensor<f32>) -> tensor<f32> {
--  //      CHECK: "vhlo.case_v1"(%arg0) ({
--  // CHECK-NEXT:   "vhlo.return_v1"(%arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-+  //      CHECK: "vhlo.case_v1"(%[[ARG0]]) ({
-+  // CHECK-NEXT:   "vhlo.return_v1"(%[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-   // CHECK-NEXT: }) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.case"(%arg0) ({
-     "stablehlo.return"(%arg1) : (tensor<f32>) -> ()
-@@ -1840,8 +1973,9 @@
- }
- 
- // CHECK-LABEL: "op_reverse"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_reverse(%arg0: tensor<16xf32>) -> tensor<16xf32> {
--  //      CHECK: "vhlo.reverse_v1"(%arg0) <{
-+  //      CHECK: "vhlo.reverse_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   dimensions = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x!vhlo.f32_v1>
-   %0 = "stablehlo.reverse"(%arg0) {
-@@ -1851,8 +1985,9 @@
- }
- 
- // CHECK-LABEL: "op_rng_bit_generator"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_rng_bit_generator(%arg0: tensor<f32>) -> (tensor<f32>, tensor<f32>) {
--  //      CHECK: "vhlo.rng_bit_generator_v1"(%arg0) <{
-+  //      CHECK: "vhlo.rng_bit_generator_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   rng_algorithm = #vhlo<rng_algorithm_v1 PHILOX>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>)
-   %0:2 = "stablehlo.rng_bit_generator"(%arg0) {
-@@ -1862,8 +1997,9 @@
- }
- 
- // CHECK-LABEL: "op_rng"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_rng(%arg0: tensor<f32>, %arg1: tensor<f32>, %arg2: tensor<0xindex>) -> tensor<f32> {
--  //      CHECK: "vhlo.rng_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.rng_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   rng_distribution = #vhlo<rng_distribution_v1 NORMAL>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<0x!vhlo.index_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.rng"(%arg0, %arg1, %arg2) {
-@@ -1873,29 +2009,33 @@
- }
- 
- // CHECK-LABEL: "op_round_nearest_afz"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_round_nearest_afz(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.round_nearest_afz_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.round_nearest_afz_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.round_nearest_afz"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_round_nearest_even"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_round_nearest_even(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.round_nearest_even_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.round_nearest_even_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.round_nearest_even"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_rsqrt"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_rsqrt(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.rsqrt_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.rsqrt_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.rsqrt"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_scatter"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_scatter(%arg0: tensor<200x100x300xf32>, %arg1: tensor<10x2xi32>, %arg2: tensor<10x300xf32>) -> tensor<200x100x300xf32> {
--  //      CHECK: "vhlo.scatter_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.scatter_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   index_vector_dim = #vhlo.integer_v1<1 : i64>,
-   // CHECK-SAME:   indices_are_sorted = #vhlo.bool_v1<true>,
-   // CHECK-SAME:   inserted_window_dims = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-@@ -1928,7 +2068,7 @@
- func.func @op_scatter_with_promotable_types(%input_tensor: tensor<200x100x300xf32>,
-     %scatter_indices: tensor<10x2xi32>, %updates: tensor<10x300xf32>) ->
-       tensor<200x100x300xf64> {
--  //  CHECK: "vhlo.scatter_v1"(%arg0, %arg1, %arg2)
-+  //  CHECK: "vhlo.scatter_v1"(%[[ARG0:.*]], %[[ARG1:.*]], %[[ARG2:.*]])
-   //  CHECK:   ^[[BB:bb.*]](%[[ARG1:arg.*]]: !vhlo.tensor_v1<!vhlo.f64_v1>, %[[ARG2:arg.*]]: !vhlo.tensor_v1<!vhlo.f64_v1>):
-   //  CHECK:     "vhlo.return_v1"(%[[VAL1:.*]]) : (!vhlo.tensor_v1<!vhlo.f64_v1>) -> ()
-   //  CHECK: }) : (!vhlo.tensor_v1<200x100x300x!vhlo.f32_v1>, !vhlo.tensor_v1<10x2x!vhlo.i32_v1>, !vhlo.tensor_v1<10x300x!vhlo.f32_v1>) -> !vhlo.tensor_v1<200x100x300x!vhlo.f64_v1>
-@@ -1951,8 +2091,9 @@
- }
- 
- // CHECK-LABEL: "op_select_and_scatter"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_select_and_scatter(%arg0: tensor<10x24x24x64xf32>, %arg1: tensor<12x13x13x66xf32>, %arg2: tensor<f32>) -> tensor<10x24x24x64xf32> {
--  //      CHECK: "vhlo.select_and_scatter_v1"(%arg0, %arg1, %arg2) <{
-+  //      CHECK: "vhlo.select_and_scatter_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-   // CHECK-SAME:   padding = #vhlo.tensor_v1<dense<1> : tensor<4x2xi64>>,
-   // CHECK-SAME:   window_dimensions = #vhlo.tensor_v1<dense<[1, 2, 2, 1]> : tensor<4xi64>>,
-   // CHECK-SAME:   window_strides = #vhlo.tensor_v1<dense<[1, 2, 2, 1]> : tensor<4xi64>>
-@@ -1982,8 +2123,9 @@
- }
- 
- // CHECK-LABEL: "op_select_and_scatter_with_promotable_types"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_select_and_scatter_with_promotable_types(%arg0: tensor<10x24x24x64xf32>, %arg1: tensor<12x13x13x66xf32>, %arg2: tensor<f32>) -> tensor<10x24x24x64xf64> {
--  // CHECK: "vhlo.select_and_scatter_v1"(%arg0, %arg1, %arg2)
-+  // CHECK: "vhlo.select_and_scatter_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]])
-   // CHECK:   ^[[BB:bb.*]](%[[ARG1:arg.*]]: !vhlo.tensor_v1<!vhlo.f64_v1>, %[[ARG2:arg.*]]: !vhlo.tensor_v1<!vhlo.f64_v1>):
-   // CHECK:     %[[VAL:.*]] = "vhlo.add_v1"(%[[ARG1]], %[[ARG2]]) : (!vhlo.tensor_v1<!vhlo.f64_v1>, !vhlo.tensor_v1<!vhlo.f64_v1>) -> !vhlo.tensor_v1<!vhlo.f64_v1>
-   // CHECK:     "vhlo.return_v1"(%[[VAL]]) : (!vhlo.tensor_v1<!vhlo.f64_v1>) -> ()
-@@ -2005,15 +2147,17 @@
- }
- 
- // CHECK-LABEL: "op_select"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
- func.func @op_select(%arg0: tensor<i1>, %arg1: tensor<f32>, %arg2: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.select_v1"(%arg0, %arg1, %arg2) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.select_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.select"(%arg0, %arg1, %arg2) : (tensor<i1>, tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_send"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_send(%arg0: tensor<f32>, %arg1: !stablehlo.token) -> !stablehlo.token {
--  //      CHECK: "vhlo.send_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.send_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-   // CHECK-SAME:   channel_type = #vhlo.integer_v1<2 : i64>,
-   // CHECK-SAME:   is_host_transfer = #vhlo.bool_v1<true>
-@@ -2026,8 +2170,9 @@
- }
- 
- // CHECK-LABEL: "op_set_dimension_size"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_set_dimension_size(%arg0: tensor<?xf32>, %arg1: tensor<i32>) -> tensor<16xf32> {
--  //      CHECK: "vhlo.set_dimension_size_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.set_dimension_size_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   dimension = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<?x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<16x!vhlo.f32_v1>
-   %0 = "stablehlo.set_dimension_size"(%arg0, %arg1) {
-@@ -2037,43 +2182,49 @@
- }
- 
- // CHECK-LABEL: "op_shift_left"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_shift_left(%arg0: tensor<i32>, %arg1: tensor<i32>) -> tensor<i32> {
--  // CHECK: "vhlo.shift_left_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.i32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-+  // CHECK: "vhlo.shift_left_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-   %0 = "stablehlo.shift_left"(%arg0, %arg1) : (tensor<i32>, tensor<i32>) -> tensor<i32>
-   func.return %0 : tensor<i32>
- }
- 
- // CHECK-LABEL: "op_shift_right_arithmetic"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_shift_right_arithmetic(%arg0: tensor<i32>, %arg1: tensor<i32>) -> tensor<i32> {
--  // CHECK: "vhlo.shift_right_arithmetic_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.i32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-+  // CHECK: "vhlo.shift_right_arithmetic_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-   %0 = "stablehlo.shift_right_arithmetic"(%arg0, %arg1) : (tensor<i32>, tensor<i32>) -> tensor<i32>
-   func.return %0 : tensor<i32>
- }
- 
- // CHECK-LABEL: "op_shift_right_logical"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_shift_right_logical(%arg0: tensor<i32>, %arg1: tensor<i32>) -> tensor<i32> {
--  // CHECK: "vhlo.shift_right_logical_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.i32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-+  // CHECK: "vhlo.shift_right_logical_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-   %0 = "stablehlo.shift_right_logical"(%arg0, %arg1) : (tensor<i32>, tensor<i32>) -> tensor<i32>
-   func.return %0 : tensor<i32>
- }
- 
- // CHECK-LABEL: "op_sign"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_sign(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.sign_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.sign_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.sign"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_sine"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_sine(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.sine_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.sine_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.sine"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_slice"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_slice(%arg0: tensor<16xf32>) -> tensor<4xf32> {
--  //      CHECK: "vhlo.slice_v1"(%arg0) <{
-+  //      CHECK: "vhlo.slice_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   limit_indices = #vhlo.tensor_v1<dense<4> : tensor<1xi64>>,
-   // CHECK-SAME:   start_indices = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>,
-   // CHECK-SAME:   strides = #vhlo.tensor_v1<dense<1> : tensor<1xi64>>
-@@ -2087,8 +2238,9 @@
- }
- 
- // CHECK-LABEL: "op_sort"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_sort(%arg0: tensor<16xf32>) -> tensor<16xf32> {
--  //      CHECK: "vhlo.sort_v1"(%arg0) <{
-+  //      CHECK: "vhlo.sort_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   dimension = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME:   is_stable = #vhlo.bool_v1<true>
-   // CHECK-SAME: }> ({
-@@ -2108,29 +2260,33 @@
- }
- 
- // CHECK-LABEL: "op_sqrt"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_sqrt(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.sqrt_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.sqrt_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.sqrt"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_subtract"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_subtract(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.subtract_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.subtract_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.subtract"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_tanh"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_tanh(%arg0: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.tanh_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.tanh_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.tanh"(%arg0) : (tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_torch_index_select"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_torch_index_select(%arg0: tensor<5x1x5xf32>, %arg1: tensor<2xi32>) ->  tensor<2x1x5xf32> {
--  //      CHECK: "vhlo.torch_index_select_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.torch_index_select_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   batch_dims = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME:   dim = #vhlo.integer_v1<0 : i64>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<5x1x5x!vhlo.f32_v1>, !vhlo.tensor_v1<2x!vhlo.i32_v1>) -> !vhlo.tensor_v1<2x1x5x!vhlo.f32_v1>
-@@ -2142,8 +2298,9 @@
- }
- 
- // CHECK-LABEL: "op_trace"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_trace(%arg0: tensor<f32>) {
--  //      CHECK: "vhlo.trace_v1"(%arg0) <{
-+  //      CHECK: "vhlo.trace_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   tag = #vhlo.string_v1<"foo">
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-   "stablehlo.trace"(%arg0) {
-@@ -2153,8 +2310,9 @@
- }
- 
- // CHECK-LABEL: "op_transpose"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_transpose(%arg0: tensor<16x8xf32>) ->  tensor<8x16xf32> {
--  //      CHECK: "vhlo.transpose_v1"(%arg0) <{
-+  //      CHECK: "vhlo.transpose_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   permutation = #vhlo.tensor_v1<dense<[1, 0]> : tensor<2xi64>>
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<16x8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<8x16x!vhlo.f32_v1>
-   %0 = "stablehlo.transpose"(%arg0) {
-@@ -2164,8 +2322,9 @@
- }
- 
- // CHECK-LABEL: "op_triangular_solve"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_triangular_solve(%arg0: tensor<16x16xf32>, %arg1: tensor<16x16xf32>) ->  tensor<16x16xf32> {
--  //      CHECK: "vhlo.triangular_solve_v1"(%arg0, %arg1) <{
-+  //      CHECK: "vhlo.triangular_solve_v1"(%[[ARG0]], %[[ARG1]]) <{
-   // CHECK-SAME:   left_side = #vhlo.bool_v1<true>,
-   // CHECK-SAME:   lower = #vhlo.bool_v1<true>,
-   // CHECK-SAME:   transpose_a = #vhlo<transpose_v1 NO_TRANSPOSE>,
-@@ -2181,15 +2340,17 @@
- }
- 
- // CHECK-LABEL: "op_tuple"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_tuple(%arg0: tensor<f32>) -> tuple<tensor<f32>> {
--  // CHECK: "vhlo.tuple_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tuple_v1<!vhlo.tensor_v1<!vhlo.f32_v1>>
-+  // CHECK: "vhlo.tuple_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tuple_v1<!vhlo.tensor_v1<!vhlo.f32_v1>>
-   %0 = "stablehlo.tuple"(%arg0) : (tensor<f32>) -> tuple<tensor<f32>>
-   func.return %0 : tuple<tensor<f32>>
- }
- 
- // CHECK-LABEL: "op_unary_einsum"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_unary_einsum(%arg0: tensor<8x16xf32>) -> tensor<8xf32> {
--  //      CHECK: "vhlo.unary_einsum_v1"(%arg0) <{
-+  //      CHECK: "vhlo.unary_einsum_v1"(%[[ARG0]]) <{
-   // CHECK-SAME:   einsum_config = #vhlo.string_v1<"ab->a">
-   // CHECK-SAME: }> : (!vhlo.tensor_v1<8x16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<8x!vhlo.f32_v1>
-   %0 = "stablehlo.unary_einsum"(%arg0) {
-@@ -2199,22 +2360,25 @@
- }
- 
- // CHECK-LABEL: "op_uniform_dequantize"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_uniform_dequantize(%arg0: tensor<!quant.uniform<i8:f32, 34.0:16>>) -> tensor<f32> {
--  // CHECK: "vhlo.uniform_dequantize_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.uniform_dequantize_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.uniform_dequantize"(%arg0) : (tensor<!quant.uniform<i8:f32, 34.0:16>>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "op_uniform_quantize"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_uniform_quantize(%arg0: tensor<f32>) -> tensor<!quant.uniform<i8:f32, 34.0:16>> {
--  // CHECK: "vhlo.uniform_quantize_v1"(%arg0) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>
-+  // CHECK: "vhlo.uniform_quantize_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>
-   %0 = "stablehlo.uniform_quantize"(%arg0) : (tensor<f32>) -> tensor<!quant.uniform<i8:f32, 34.0:16>>
-   func.return %0 : tensor<!quant.uniform<i8:f32, 34.0:16>>
- }
- 
- // CHECK-LABEL: "op_while"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_while(%arg0: tensor<i1>) -> tensor<i1> {
--  //      CHECK: "vhlo.while_v1"(%arg0) ({
-+  //      CHECK: "vhlo.while_v1"(%[[ARG0]]) ({
-   // CHECK-NEXT:   ^[[BB:bb.*]](%[[ARG1:arg.*]]: !vhlo.tensor_v1<!vhlo.bool_v1>):
-   // CHECK-NEXT:     "vhlo.return_v1"(%[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>) -> ()
-   // CHECK-NEXT:   }, {
-@@ -2232,8 +2396,9 @@
- }
- 
- // CHECK-LABEL: "op_xor"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @op_xor(%arg0: tensor<i1>, %arg1: tensor<i1>) -> tensor<i1> {
--  // CHECK: "vhlo.xor_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-+  // CHECK: "vhlo.xor_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-   %0 = "stablehlo.xor"(%arg0, %arg1) : (tensor<i1>, tensor<i1>) -> tensor<i1>
-   func.return %0 : tensor<i1>
- }
-@@ -2241,197 +2406,225 @@
- // ============ TYPES ============
- 
- // CHECK-LABEL: "type_i1"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_i1(%arg0: tensor<i1>, %arg1: tensor<i1>) -> tensor<i1> {
--  // CHECK: "vhlo.and_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-+  // CHECK: "vhlo.and_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-   %0 = "stablehlo.and"(%arg0, %arg1) : (tensor<i1>, tensor<i1>) -> tensor<i1>
-   func.return %0 : tensor<i1>
- }
- 
- // CHECK-LABEL: "type_i4"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_i4(%arg0: tensor<i4>, %arg1: tensor<i4>) -> tensor<i4> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.i4_v1>, !vhlo.tensor_v1<!vhlo.i4_v1>) -> !vhlo.tensor_v1<!vhlo.i4_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i4_v1>, !vhlo.tensor_v1<!vhlo.i4_v1>) -> !vhlo.tensor_v1<!vhlo.i4_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<i4>, tensor<i4>) -> tensor<i4>
-   func.return %0 : tensor<i4>
- }
- 
- // CHECK-LABEL: "type_i8"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_i8(%arg0: tensor<i8>, %arg1: tensor<i8>) -> tensor<i8> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.i8_v1>, !vhlo.tensor_v1<!vhlo.i8_v1>) -> !vhlo.tensor_v1<!vhlo.i8_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i8_v1>, !vhlo.tensor_v1<!vhlo.i8_v1>) -> !vhlo.tensor_v1<!vhlo.i8_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<i8>, tensor<i8>) -> tensor<i8>
-   func.return %0 : tensor<i8>
- }
- 
- // CHECK-LABEL: "type_i16"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_i16(%arg0: tensor<i16>, %arg1: tensor<i16>) -> tensor<i16> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.i16_v1>, !vhlo.tensor_v1<!vhlo.i16_v1>) -> !vhlo.tensor_v1<!vhlo.i16_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i16_v1>, !vhlo.tensor_v1<!vhlo.i16_v1>) -> !vhlo.tensor_v1<!vhlo.i16_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<i16>, tensor<i16>) -> tensor<i16>
-   func.return %0 : tensor<i16>
- }
- 
- // CHECK-LABEL: "type_i32"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_i32(%arg0: tensor<i32>, %arg1: tensor<i32>) -> tensor<i32> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.i32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<i32>, tensor<i32>) -> tensor<i32>
-   func.return %0 : tensor<i32>
- }
- 
- // CHECK-LABEL: "type_i64"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_i64(%arg0: tensor<i64>, %arg1: tensor<i64>) -> tensor<i64> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.i64_v1>, !vhlo.tensor_v1<!vhlo.i64_v1>) -> !vhlo.tensor_v1<!vhlo.i64_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i64_v1>, !vhlo.tensor_v1<!vhlo.i64_v1>) -> !vhlo.tensor_v1<!vhlo.i64_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<i64>, tensor<i64>) -> tensor<i64>
-   func.return %0 : tensor<i64>
- }
- 
- // CHECK-LABEL: "type_ui4"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_ui4(%arg0: tensor<ui4>, %arg1: tensor<ui4>) -> tensor<ui4> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.ui4_v1>, !vhlo.tensor_v1<!vhlo.ui4_v1>) -> !vhlo.tensor_v1<!vhlo.ui4_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.ui4_v1>, !vhlo.tensor_v1<!vhlo.ui4_v1>) -> !vhlo.tensor_v1<!vhlo.ui4_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<ui4>, tensor<ui4>) -> tensor<ui4>
-   func.return %0 : tensor<ui4>
- }
- 
- // CHECK-LABEL: "type_ui8"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_ui8(%arg0: tensor<ui8>, %arg1: tensor<ui8>) -> tensor<ui8> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.ui8_v1>, !vhlo.tensor_v1<!vhlo.ui8_v1>) -> !vhlo.tensor_v1<!vhlo.ui8_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.ui8_v1>, !vhlo.tensor_v1<!vhlo.ui8_v1>) -> !vhlo.tensor_v1<!vhlo.ui8_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<ui8>, tensor<ui8>) -> tensor<ui8>
-   func.return %0 : tensor<ui8>
- }
- 
- // CHECK-LABEL: "type_ui16"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_ui16(%arg0: tensor<ui16>, %arg1: tensor<ui16>) -> tensor<ui16> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.ui16_v1>, !vhlo.tensor_v1<!vhlo.ui16_v1>) -> !vhlo.tensor_v1<!vhlo.ui16_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.ui16_v1>, !vhlo.tensor_v1<!vhlo.ui16_v1>) -> !vhlo.tensor_v1<!vhlo.ui16_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<ui16>, tensor<ui16>) -> tensor<ui16>
-   func.return %0 : tensor<ui16>
- }
- 
- // CHECK-LABEL: "type_ui32"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_ui32(%arg0: tensor<ui32>, %arg1: tensor<ui32>) -> tensor<ui32> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.ui32_v1>, !vhlo.tensor_v1<!vhlo.ui32_v1>) -> !vhlo.tensor_v1<!vhlo.ui32_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.ui32_v1>, !vhlo.tensor_v1<!vhlo.ui32_v1>) -> !vhlo.tensor_v1<!vhlo.ui32_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<ui32>, tensor<ui32>) -> tensor<ui32>
-   func.return %0 : tensor<ui32>
- }
- 
- // CHECK-LABEL: "type_ui64"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_ui64(%arg0: tensor<ui64>, %arg1: tensor<ui64>) -> tensor<ui64> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.ui64_v1>, !vhlo.tensor_v1<!vhlo.ui64_v1>) -> !vhlo.tensor_v1<!vhlo.ui64_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.ui64_v1>, !vhlo.tensor_v1<!vhlo.ui64_v1>) -> !vhlo.tensor_v1<!vhlo.ui64_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<ui64>, tensor<ui64>) -> tensor<ui64>
-   func.return %0 : tensor<ui64>
- }
- 
- // CHECK-LABEL: "type_f8E4M3FN"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_f8E4M3FN(%arg0: tensor<f8E4M3FN>, %arg1: tensor<f8E4M3FN>) -> tensor<f8E4M3FN> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f8E4M3FN_v1>, !vhlo.tensor_v1<!vhlo.f8E4M3FN_v1>) -> !vhlo.tensor_v1<!vhlo.f8E4M3FN_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f8E4M3FN_v1>, !vhlo.tensor_v1<!vhlo.f8E4M3FN_v1>) -> !vhlo.tensor_v1<!vhlo.f8E4M3FN_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f8E4M3FN>, tensor<f8E4M3FN>) -> tensor<f8E4M3FN>
-   func.return %0 : tensor<f8E4M3FN>
- }
- 
- // CHECK-LABEL: "type_f8E5M2"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_f8E5M2(%arg0: tensor<f8E5M2>, %arg1: tensor<f8E5M2>) -> tensor<f8E5M2> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f8E5M2_v1>, !vhlo.tensor_v1<!vhlo.f8E5M2_v1>) -> !vhlo.tensor_v1<!vhlo.f8E5M2_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f8E5M2_v1>, !vhlo.tensor_v1<!vhlo.f8E5M2_v1>) -> !vhlo.tensor_v1<!vhlo.f8E5M2_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f8E5M2>, tensor<f8E5M2>) -> tensor<f8E5M2>
-   func.return %0 : tensor<f8E5M2>
- }
- 
- // CHECK-LABEL: "type_f8E4M3FNUZ"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_f8E4M3FNUZ(%arg0: tensor<f8E4M3FNUZ>, %arg1: tensor<f8E4M3FNUZ>) -> tensor<f8E4M3FNUZ> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f8E4M3FNUZ_v1>, !vhlo.tensor_v1<!vhlo.f8E4M3FNUZ_v1>) -> !vhlo.tensor_v1<!vhlo.f8E4M3FNUZ_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f8E4M3FNUZ_v1>, !vhlo.tensor_v1<!vhlo.f8E4M3FNUZ_v1>) -> !vhlo.tensor_v1<!vhlo.f8E4M3FNUZ_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f8E4M3FNUZ>, tensor<f8E4M3FNUZ>) -> tensor<f8E4M3FNUZ>
-   func.return %0 : tensor<f8E4M3FNUZ>
- }
- 
- // CHECK-LABEL: "type_f8E4M3B11FNUZ"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_f8E4M3B11FNUZ(%arg0: tensor<f8E4M3B11FNUZ>, %arg1: tensor<f8E4M3B11FNUZ>) -> tensor<f8E4M3B11FNUZ> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f8E4M3B11FNUZ_v1>, !vhlo.tensor_v1<!vhlo.f8E4M3B11FNUZ_v1>) -> !vhlo.tensor_v1<!vhlo.f8E4M3B11FNUZ_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f8E4M3B11FNUZ_v1>, !vhlo.tensor_v1<!vhlo.f8E4M3B11FNUZ_v1>) -> !vhlo.tensor_v1<!vhlo.f8E4M3B11FNUZ_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f8E4M3B11FNUZ>, tensor<f8E4M3B11FNUZ>) -> tensor<f8E4M3B11FNUZ>
-   func.return %0 : tensor<f8E4M3B11FNUZ>
- }
- 
- // CHECK-LABEL: "type_f8E5M2FNUZ"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_f8E5M2FNUZ(%arg0: tensor<f8E5M2FNUZ>, %arg1: tensor<f8E5M2FNUZ>) -> tensor<f8E5M2FNUZ> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f8E5M2FNUZ_v1>, !vhlo.tensor_v1<!vhlo.f8E5M2FNUZ_v1>) -> !vhlo.tensor_v1<!vhlo.f8E5M2FNUZ_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f8E5M2FNUZ_v1>, !vhlo.tensor_v1<!vhlo.f8E5M2FNUZ_v1>) -> !vhlo.tensor_v1<!vhlo.f8E5M2FNUZ_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f8E5M2FNUZ>, tensor<f8E5M2FNUZ>) -> tensor<f8E5M2FNUZ>
-   func.return %0 : tensor<f8E5M2FNUZ>
- }
- 
- // CHECK-LABEL: "type_bf16"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_bf16(%arg0: tensor<bf16>, %arg1: tensor<bf16>) -> tensor<bf16> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.bf16_v1>, !vhlo.tensor_v1<!vhlo.bf16_v1>) -> !vhlo.tensor_v1<!vhlo.bf16_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.bf16_v1>, !vhlo.tensor_v1<!vhlo.bf16_v1>) -> !vhlo.tensor_v1<!vhlo.bf16_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<bf16>, tensor<bf16>) -> tensor<bf16>
-   func.return %0 : tensor<bf16>
- }
- 
- // CHECK-LABEL: "type_f16"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_f16(%arg0: tensor<f16>, %arg1: tensor<f16>) -> tensor<f16> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f16_v1>, !vhlo.tensor_v1<!vhlo.f16_v1>) -> !vhlo.tensor_v1<!vhlo.f16_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f16_v1>, !vhlo.tensor_v1<!vhlo.f16_v1>) -> !vhlo.tensor_v1<!vhlo.f16_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f16>, tensor<f16>) -> tensor<f16>
-   func.return %0 : tensor<f16>
- }
- 
- // CHECK-LABEL: "type_f32"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_f32(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-   func.return %0 : tensor<f32>
- }
- 
- // CHECK-LABEL: "type_f64"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_f64(%arg0: tensor<f64>, %arg1: tensor<f64>) -> tensor<f64> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f64_v1>, !vhlo.tensor_v1<!vhlo.f64_v1>) -> !vhlo.tensor_v1<!vhlo.f64_v1>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f64_v1>, !vhlo.tensor_v1<!vhlo.f64_v1>) -> !vhlo.tensor_v1<!vhlo.f64_v1>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f64>, tensor<f64>) -> tensor<f64>
-   func.return %0 : tensor<f64>
- }
- 
- // CHECK-LABEL: "type_complex_f32"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_complex_f32(%arg0: tensor<complex<f32>>, %arg1: tensor<complex<f32>>) -> tensor<complex<f32>> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>, !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>) -> !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>, !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>) -> !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<complex<f32>>, tensor<complex<f32>>) -> tensor<complex<f32>>
-   func.return %0 : tensor<complex<f32>>
- }
- 
- // CHECK-LABEL: "type_complex_f64"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_complex_f64(%arg0: tensor<complex<f64>>, %arg1: tensor<complex<f64>>) -> tensor<complex<f64>> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f64_v1>>, !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f64_v1>>) -> !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f64_v1>>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f64_v1>>, !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f64_v1>>) -> !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f64_v1>>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<complex<f64>>, tensor<complex<f64>>) -> tensor<complex<f64>>
-   func.return %0 : tensor<complex<f64>>
- }
- 
- // CHECK-LABEL: "type_dynamism_ranked"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @type_dynamism_ranked(%arg0: tensor<?xf32>) -> tensor<?xf32> {
--  // CHECK: "vhlo.abs_v1"(%arg0) : (!vhlo.tensor_v1<?x!vhlo.f32_v1>) -> !vhlo.tensor_v1<?x!vhlo.f32_v1>
-+  // CHECK: "vhlo.abs_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<?x!vhlo.f32_v1>) -> !vhlo.tensor_v1<?x!vhlo.f32_v1>
-   %0 = "stablehlo.abs"(%arg0) : (tensor<?xf32>) -> tensor<?xf32>
-   func.return %0 : tensor<?xf32>
- }
- 
- // CHECK-LABEL: "type_per_tensor_quantization"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @type_per_tensor_quantization(%arg0: tensor<!quant.uniform<i8:f32, 34.0:16>>, %arg1: tensor<!quant.uniform<i8:f32, 34.0:16>>) -> tensor<!quant.uniform<i8:f32, 34.0:16>> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>, !vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>) -> !vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>, !vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>) -> !vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<!quant.uniform<i8:f32, 34.0:16>>, tensor<!quant.uniform<i8:f32, 34.0:16>>) -> tensor<!quant.uniform<i8:f32, 34.0:16>>
-   func.return %0 : tensor<!quant.uniform<i8:f32, 34.0:16>>
- }
- 
- // CHECK-LABEL: "type_per_axis_quantization"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @type_per_axis_quantization(%arg0: tensor<2x!quant.uniform<i8:f32:0, {34.0:16, 34.0:16}>>) -> tensor<2x!quant.uniform<i8:f32:0, {34.0:16, 34.0:16}>> {
--  // CHECK: "vhlo.add_v1"(%arg0, %arg0) : (!vhlo.tensor_v1<2x!vhlo.quant_per_axis_v1<!vhlo.i8_v1:!vhlo.f32_v1, 0, 3.400000e+01, 3.400000e+01, 16, 16, -128:127, 1>>, !vhlo.tensor_v1<2x!vhlo.quant_per_axis_v1<!vhlo.i8_v1:!vhlo.f32_v1, 0, 3.400000e+01, 3.400000e+01, 16, 16, -128:127, 1>>) -> !vhlo.tensor_v1<2x!vhlo.quant_per_axis_v1<!vhlo.i8_v1:!vhlo.f32_v1, 0, 3.400000e+01, 3.400000e+01, 16, 16, -128:127, 1>>
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG0]]) : (!vhlo.tensor_v1<2x!vhlo.quant_per_axis_v1<!vhlo.i8_v1:!vhlo.f32_v1, 0, 3.400000e+01, 3.400000e+01, 16, 16, -128:127, 1>>, !vhlo.tensor_v1<2x!vhlo.quant_per_axis_v1<!vhlo.i8_v1:!vhlo.f32_v1, 0, 3.400000e+01, 3.400000e+01, 16, 16, -128:127, 1>>) -> !vhlo.tensor_v1<2x!vhlo.quant_per_axis_v1<!vhlo.i8_v1:!vhlo.f32_v1, 0, 3.400000e+01, 3.400000e+01, 16, 16, -128:127, 1>>
-   %0 = stablehlo.add %arg0, %arg0 : tensor<2x!quant.uniform<i8:f32:0, {34.0:16, 34.0:16}>>
-   func.return %0 : tensor<2x!quant.uniform<i8:f32:0, {34.0:16, 34.0:16}>>
- }
- 
- //       CHECK: function_type = #vhlo.type_v1<!vhlo.func_v1<(!vhlo.token_v1) -> !vhlo.token_v1>>
- // CHECK-LABEL: "type_token_callee"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @type_token_callee(%arg0: !stablehlo.token) -> !stablehlo.token {
--  // CHECK: "vhlo.return_v1"(%arg0) : (!vhlo.token_v1) -> ()
-+  // CHECK: "vhlo.return_v1"(%[[ARG0]]) : (!vhlo.token_v1) -> ()
-   return %arg0 : !stablehlo.token
- }
- 
- //       CHECK: function_type = #vhlo.type_v1<!vhlo.func_v1<(!vhlo.token_v1) -> !vhlo.token_v1>>
- // CHECK-LABEL: "type_token_caller"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @type_token_caller(%arg0: !stablehlo.token) -> !stablehlo.token {
--  // CHECK:      "vhlo.call_v1"(%arg0) <{callee = #vhlo.string_v1<"type_token_callee">}
-+  // CHECK:      "vhlo.call_v1"(%[[ARG0]]) <{callee = #vhlo.string_v1<"type_token_callee">}
-   // CHECK-SAME: (!vhlo.token_v1) -> !vhlo.token_v1
-   %0 = func.call @type_token_callee(%arg0) : (!stablehlo.token) -> !stablehlo.token
-   return %0 : !stablehlo.token
- }
- 
- // CHECK-LABEL: "type_tuple"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @type_tuple(%arg0: tuple<tensor<f32>>) -> tuple<!stablehlo.token> {
-   %0 = "stablehlo.custom_call"(%arg0) {
-     call_target_name = "foo"
 
diff --git a/third_party/xla/third_party/stablehlo/workspace.bzl b/third_party/xla/third_party/stablehlo/workspace.bzl
index 6a72c8fa16885c..aaef166d96583c 100644
--- a/third_party/xla/third_party/stablehlo/workspace.bzl
+++ b/third_party/xla/third_party/stablehlo/workspace.bzl
@@ -4,8 +4,8 @@ load("//third_party:repo.bzl", "tf_http_archive", "tf_mirror_urls")
 
 def repo():
     # LINT.IfChange
-    STABLEHLO_COMMIT = "8ba7728d3fdc3ea882e893ee7e53255c95ee0e5a"
-    STABLEHLO_SHA256 = "1dfc7179dc9200c3ab4ea85edbac4a35393866d8cd8694fcaac00c1d27036408"
+    STABLEHLO_COMMIT = "797bee217e1a041e9aac22cad4db207274596d94"
+    STABLEHLO_SHA256 = "e5619033e131ea2eeb9eab8c8e362f3ba12e111c6b4a15dac789ca216ff22c58"
     # LINT.ThenChange(Google-internal path)
 
     tf_http_archive(
diff --git a/third_party/xla/xla/mlir_hlo/mhlo/IR/hlo_base.td b/third_party/xla/xla/mlir_hlo/mhlo/IR/hlo_base.td
index 15d8dcc8cf5f5f..992aca563edf97 100644
--- a/third_party/xla/xla/mlir_hlo/mhlo/IR/hlo_base.td
+++ b/third_party/xla/xla/mlir_hlo/mhlo/IR/hlo_base.td
@@ -131,4 +131,6 @@ defvar MHLO_StaticShapeIntOrFpTensor = HLO_StaticShapeIntOrFpTensor;
 
 defvar MHLO_StaticShapeIntFpOrComplexTensor = HLO_StaticShapeIntFpOrComplexTensor;
 
+defvar MHLO_Static1DIntTensor = HLO_Static1DIntTensor;
+
 #endif // MLIR_HLO_DIALECT_MHLO_IR_HLO_BASE
diff --git a/third_party/xla/xla/mlir_hlo/mhlo/IR/hlo_ops.td b/third_party/xla/xla/mlir_hlo/mhlo/IR/hlo_ops.td
index 75db5312d336c1..cfba735634a366 100644
--- a/third_party/xla/xla/mlir_hlo/mhlo/IR/hlo_ops.td
+++ b/third_party/xla/xla/mlir_hlo/mhlo/IR/hlo_ops.td
@@ -120,12 +120,12 @@ def MHLO_IotaOp : MHLO_Op<"iota", [Pure]> {
 def MHLO_DynamicIotaOp: MHLO_ShapedInterfaceOp<"dynamic_iota", [Pure]> {
   let summary = "DynamicIota operation";
   let description = [{
-    This operation is a work in progress, so it is not yet included in
-    the specification: https://github.com/openxla/stablehlo/issues/8.
+    This operation is functionally identical to
+    [iota](https://github.com/openxla/stablehlo/blob/main/docs/spec.md#iota)
+    op, but the result shape is specified dynamically via `output_shape`.
 
-    Informally, this operation does the same thing as IotaOp except that the
-    result shape is specified dynamically via `output_shape`:
-    https://github.com/openxla/stablehlo/blob/main/docs/spec.md#iota
+    See:
+    https://github.com/openxla/stablehlo/blob/main/docs/spec.md#dynamic_iota
 
     Example:
     ```mlir
@@ -2126,19 +2126,30 @@ def MHLO_DynamicBroadcastInDimOp : MHLO_ShapedInterfaceOp<
     "dynamic_broadcast_in_dim", [Pure]> {
   let summary = "DynamicBroadcastInDim operation";
   let description = [{
-    This operation is a work in progress, so it is not yet included in
-    the specification: https://github.com/openxla/stablehlo/issues/8.
-
-    Informally, this operation does the same thing as BroadcastInDimOp except
-    that the result shape is specified dynamically via `output_dimensions`:
-    https://github.com/openxla/stablehlo/blob/main/docs/spec.md#broadcast_in_dim
+    This operation is functionally identical to
+    [broadcast_in_dim](https://github.com/openxla/stablehlo/blob/main/docs/spec.md#broadcast_in_dim)
+    op, but the result shape is specified dynamically via `output_dimensions`.
 
     It also accepts optional attributes to express static knowledge about the
     expanding behavior of dimensions. If not specified, all dimensions are
     assumed to be possibly expanding. The sets of dimensions that are known to
     be expanding and the set of dimensions that are known to be non-expanding
     must be disjoint and they must be a subset of the operand's dimensions.
+
+    See: https://github.com/openxla/stablehlo/blob/main/docs/spec.md#dynamic_broadcast_in_dim
+
+    Example:
+    ```mlir
+    %operand = mhlo.constant dense<[[1, 2, 3]]> : tensor<1x3xi64>
+    %output_dimensions = mhlo.constant dense<[2, 3, 2]> : tensor<3xi64>
+    %result = "mhlo.dynamic_broadcast_in_dim"(%operand, %output_dimensions) {
+      broadcast_dimensions = array<i64: 2, 1>,
+      known_expanding_dimensions = array<i64: 0>,
+      known_non_expanding_dimensions = array<i64: 1>
+    } : (tensor<1x3xi64>, tensor<3xi64>) -> tensor<2x3x2xi64>
+    ```
   }];
+
   let arguments = (ins
     MHLO_Tensor:$operand,
     MHLO_DimensionTensor:$output_dimensions,
@@ -2854,16 +2865,17 @@ def MHLO_ReshapeOp: MHLO_Op<"reshape",
 def MHLO_DynamicReshapeOp: MHLO_ShapedInterfaceOp<"dynamic_reshape", [Pure]> {
   let summary = "DynamicReshape operation";
   let description = [{
-    This operation is a work in progress, so it is not yet included in
-    the specification: https://github.com/openxla/stablehlo/issues/8.
+    This operation is functionally identical to
+    [reshape](https://github.com/openxla/stablehlo/blob/main/docs/spec.md#reshape)
+    op, but the result shape is specified dynamically via `output_shape`.
 
-    Informally, this operation does the same thing as ReshapeOp except that the
-    result shape is specified dynamically via `output_shape`:
-    https://github.com/openxla/stablehlo/blob/main/docs/spec.md#reshape
+    See:
+    https://github.com/openxla/stablehlo/blob/main/docs/spec.md#dynamic_reshape
 
     Example:
     ```mlir
-    %0 = mhlo.dynamic_reshape %arg0, %shape : (tensor<?xf32>, tensor<2xindex>) -> tensor<?x?xf32>
+    %output_shape = mhlo.constant dense<[3, 2]> : tensor<2xi64>
+    %result = mhlo.dynamic_reshape %operand, %output_shape : (tensor<2x3xi64>, tensor<2xi64>) -> tensor<3x2xi64>
     ```
   }];
 
@@ -3795,7 +3807,7 @@ def MHLO_DynamicGatherOp: MHLO_Op<"dynamic_gather",
   let arguments = (ins
     MHLO_Tensor:$operand,
     MHLO_IntTensor:$start_indices,
-    MHLO_IntTensor:$slice_sizes,
+    MHLO_Static1DIntTensor:$slice_sizes,
     MHLO_GatherDimensionNumbers:$dimension_numbers,
     DefaultValuedOptionalAttr<BoolAttr, "false">:$indices_are_sorted
   );
diff --git a/third_party/xla/xla/mlir_hlo/mhlo/transforms/map_stablehlo_to_hlo_op.h b/third_party/xla/xla/mlir_hlo/mhlo/transforms/map_stablehlo_to_hlo_op.h
index 031ce586e80a5c..80c6c0b0bd1ea7 100644
--- a/third_party/xla/xla/mlir_hlo/mhlo/transforms/map_stablehlo_to_hlo_op.h
+++ b/third_party/xla/xla/mlir_hlo/mhlo/transforms/map_stablehlo_to_hlo_op.h
@@ -150,7 +150,6 @@ MAP_STABLEHLO_TO_HLO(SqrtOp)
 MAP_STABLEHLO_TO_HLO(SubtractOp)
 MAP_STABLEHLO_TO_HLO(TanhOp)
 MAP_STABLEHLO_TO_HLO(TorchIndexSelectOp)
-MAP_STABLEHLO_TO_HLO(TraceOp)
 MAP_STABLEHLO_TO_HLO(TransposeOp)
 MAP_STABLEHLO_TO_HLO(TriangularSolveOp)
 MAP_STABLEHLO_TO_HLO(TupleOp)
diff --git a/third_party/xla/xla/mlir_hlo/tests/Dialect/mhlo/hlo-legalize-to-stablehlo.mlir b/third_party/xla/xla/mlir_hlo/tests/Dialect/mhlo/hlo-legalize-to-stablehlo.mlir
index b624a4341dba83..32f555f88cc1cc 100644
--- a/third_party/xla/xla/mlir_hlo/tests/Dialect/mhlo/hlo-legalize-to-stablehlo.mlir
+++ b/third_party/xla/xla/mlir_hlo/tests/Dialect/mhlo/hlo-legalize-to-stablehlo.mlir
@@ -821,21 +821,19 @@ func.func @op_dynamic_broadcast_in_dim(%arg0: tensor<?xf32>, %arg1: tensor<2xind
 }
 
 // CHECK-LABEL: "op_dynamic_conv"
-func.func @op_dynamic_conv(%arg0: tensor<1x8x8x207xf32>, %arg1: tensor<3x3x207x16xf32>, %arg2: tensor<4xi32>) -> tensor<1x?x?x16xf32> {
+func.func @op_dynamic_conv(%arg0: tensor<1x8x8x207xf32>, %arg1: tensor<3x3x207x16xf32>, %arg2: tensor<2x2xi32>) -> tensor<1x?x?x16xf32> {
   //      CHECK: "stablehlo.dynamic_conv"([[ARG0:%arg[0-9]+]], [[ARG1:%arg[0-9]+]], [[ARG2:%arg[0-9]+]]) <{
   // CHECK-SAME:   batch_group_count = 1 : i64,
   // CHECK-SAME:   dimension_numbers = #stablehlo.conv<[b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f]>,
   // CHECK-SAME:   feature_group_count = 1 : i64,
   // CHECK-SAME:   lhs_dilation = array<i64: 1, 1>,
-  // CHECK-SAME:   padding = dense<1> : tensor<2x2xi64>,
   // CHECK-SAME:   precision_config = [#stablehlo<precision DEFAULT>, #stablehlo<precision DEFAULT>],
   // CHECK-SAME:   rhs_dilation = array<i64: 1, 1>,
   // CHECK-SAME:   window_reversal = array<i1: false, false>,
   // CHECK-SAME:   window_strides = array<i64: 1, 1>
-  // CHECK-SAME: }> : (tensor<1x8x8x207xf32>, tensor<3x3x207x16xf32>, tensor<4xi32>) -> tensor<1x?x?x16xf32>
+  // CHECK-SAME: }> : (tensor<1x8x8x207xf32>, tensor<3x3x207x16xf32>, tensor<2x2xi32>) -> tensor<1x?x?x16xf32>
   %0 = "mhlo.dynamic_conv"(%arg0, %arg1, %arg2) {
     window_strides = dense<1> : tensor<2xi64>,
-    padding = dense<1> : tensor<2x2xi64>,
     lhs_dilation = dense<1> : tensor<2xi64>,
     rhs_dilation = dense<1> : tensor<2xi64>,
     window_reversal = dense<false> : tensor<2xi1>,
@@ -843,7 +841,7 @@ func.func @op_dynamic_conv(%arg0: tensor<1x8x8x207xf32>, %arg1: tensor<3x3x207x1
     feature_group_count = 1 : i64,
     batch_group_count = 1 : i64,
     precision_config = [#mhlo<precision DEFAULT>, #mhlo<precision DEFAULT>]
-  } : (tensor<1x8x8x207xf32>, tensor<3x3x207x16xf32>, tensor<4xi32>) -> tensor<1x?x?x16xf32>
+  } : (tensor<1x8x8x207xf32>, tensor<3x3x207x16xf32>, tensor<2x2xi32>) -> tensor<1x?x?x16xf32>
   func.return %0 : tensor<1x?x?x16xf32>
 }
 
@@ -1604,17 +1602,6 @@ func.func @op_torch_index_select(%arg0: tensor<5x1x5xf32>, %arg1: tensor<2xi32>)
   func.return %0 : tensor<2x1x5xf32>
 }
 
-// CHECK-LABEL: "op_trace"
-func.func @op_trace(%arg0: tensor<f32>) {
-  //      CHECK: "stablehlo.trace"([[ARG0:%arg[0-9]+]]) <{
-  // CHECK-SAME:   tag = "foo"
-  // CHECK-SAME: }> : (tensor<f32>) -> ()
-  "mhlo.trace"(%arg0) {
-    tag = "foo"
-  } : (tensor<f32>) -> ()
-  func.return
-}
-
 // CHECK-LABEL: "op_transpose"
 func.func @op_transpose(%arg0: tensor<16x8xf32>) ->  tensor<8x16xf32> {
   //      CHECK: "stablehlo.transpose"([[ARG0:%arg[0-9]+]]) <{
diff --git a/third_party/xla/xla/mlir_hlo/tests/Dialect/mhlo/mhlo-quant-legalize-to-int.mlir b/third_party/xla/xla/mlir_hlo/tests/Dialect/mhlo/mhlo-quant-legalize-to-int.mlir
index 7d9455a5535ac9..6b0404881d9e2b 100644
--- a/third_party/xla/xla/mlir_hlo/tests/Dialect/mhlo/mhlo-quant-legalize-to-int.mlir
+++ b/third_party/xla/xla/mlir_hlo/tests/Dialect/mhlo/mhlo-quant-legalize-to-int.mlir
@@ -697,7 +697,7 @@ func.func @dot_dynamic_batch_dim(
 // CHECK-LABEL: func @dot_general
 func.func @dot_general(
     %arg0: tensor<2x5x6x!quant.uniform<i8:f32, 2.000000e+00:3>>,
-    %arg1: tensor<6x8x2x!quant.uniform<i8:f32, 1.000000e+00:5>>
+    %arg1: tensor<6x8x2x!quant.uniform<i8:f32, 1.000000e+00:0>>
   ) -> tensor<2x5x8x!quant.uniform<i32:f32, 4.000000e+00:7>> {
   // CHECK: %[[DOT_RES:.*]] = "mhlo.dot_general"
   // CHECK-SAME: lhs_batching_dimensions = [0]
@@ -705,22 +705,7 @@ func.func @dot_general(
   // CHECK-SAME: lhs_contracting_dimensions = [2]
   // CHECK-SAME: rhs_contracting_dimensions = [0]
 
-  // Zero point offset contribution from LHS tensor * RHS ZP.
-
-  // CHECK: %[[LHS_I32:.*]] = mhlo.convert %[[LHS:.*]] : (tensor<2x5x6xi8>)
-  // CHECK-SAME: -> tensor<2x5x6xi32>
-  // CHECK: %[[LHS_REDUCE_INIT:.*]] = mhlo.constant dense<0> : tensor<i32>
-  // CHECK: %[[LHS_REDUCE:.*]] = mhlo.reduce(%[[LHS_I32]] init: %[[LHS_REDUCE_INIT]])
-  // CHECK-SAME: applies mhlo.add across dimensions = [2]
-  // CHECK-SAME: (tensor<2x5x6xi32>, tensor<i32>)
-  // CHECK-SAME: -> tensor<2x5xi32>
-  // CHECK: %[[RHS_ZP:.*]] = mhlo.constant dense<5> : tensor<i32>
-  // CHECK: %[[LHS_ZP_CONTRIB:.*]] = chlo.broadcast_multiply
-  // CHECK-SAME: %[[LHS_REDUCE]], %[[RHS_ZP]] :
-  // CHECK-SAME: (tensor<2x5xi32>, tensor<i32>) -> tensor<2x5xi32>
-  // CHECK: %[[LHS_ZP_BCAST:.*]] = "mhlo.broadcast_in_dim"(%[[LHS_ZP_CONTRIB]])
-  // CHECK-SAME: broadcast_dimensions = dense<[0, 1]>
-  // CHECK-SAME: (tensor<2x5xi32>) -> tensor<2x5x8xi32>
+  // Zero point offset contribution from LHS tensor * RHS ZP is 0 and skipped.
 
   // Zero point offset contribution from RHS tensor * LHS ZP.
 
@@ -738,13 +723,8 @@ func.func @dot_general(
   // CHECK: %[[RHS_ZP_BCAST:.*]] = "mhlo.broadcast_in_dim"(%[[RHS_ZP_CONTRIB]])
   // CHECK-SAME: broadcast_dimensions = dense<[2, 0]>
   // CHECK-SAME: (tensor<8x2xi32>) -> tensor<2x5x8xi32>
-  // CHECK: %[[ZP_TOTAL_1:.*]] = mhlo.add %[[LHS_ZP_BCAST]], %[[RHS_ZP_BCAST]]
-
-  // Zero point offset contribution from LHS ZP * RHS ZP.
 
-  // CHECK: %[[ZPS:.*]] = mhlo.constant dense<90> : tensor<i32>
-  // CHECK: %[[ZP_TOTAL_2:.*]] = chlo.broadcast_subtract %[[ZP_TOTAL_1]], %[[ZPS]]
-  // CHECK-SAME: (tensor<2x5x8xi32>, tensor<i32>) -> tensor<2x5x8xi32>
+  // Zero point offset contribution from LHS ZP * RHS ZP is 0 and skipped.
 
   // Combine dot result with zero point offset and output final result.
 
@@ -756,17 +736,17 @@ func.func @dot_general(
   // CHECK: %[[RES_INT:.*]] = mhlo.convert %[[RES_FP_1]]
   // CHECK-SAME: (tensor<2x5x8xf32>) -> tensor<2x5x8xi32>
 
-  // CHECK: %[[ZP_TOTAL_3:.*]] = mhlo.convert %[[ZP_TOTAL_2]]
+  // CHECK: %[[ZP_TOTAL_1:.*]] = mhlo.convert %[[RHS_ZP_BCAST]]
   // CHECK-SAME: (tensor<2x5x8xi32>) -> tensor<2x5x8xf32>
-  // CHECK: %[[ZP_TOTAL_4:.*]] = chlo.broadcast_multiply
-  // CHECK-SAME: %[[ZP_TOTAL_3:.*]], %[[COMBINED_SCALE]]
-  // CHECK: %[[ZP_TOTAL_5:.*]] = mhlo.convert %[[ZP_TOTAL_4]]
+  // CHECK: %[[ZP_TOTAL_2:.*]] = chlo.broadcast_multiply
+  // CHECK-SAME: %[[ZP_TOTAL_1:.*]], %[[COMBINED_SCALE]]
+  // CHECK: %[[ZP_TOTAL_3:.*]] = mhlo.convert %[[ZP_TOTAL_2]]
   // CHECK-SAME: (tensor<2x5x8xf32>) -> tensor<2x5x8xi32>
 
   // CHECK: %[[RES_ZP:.*]] = mhlo.constant dense<7> : tensor<i32>
-  // CHECK: %[[ZP_TOTAL_6:.*]] = chlo.broadcast_subtract %[[RES_ZP]], %[[ZP_TOTAL_5]]
+  // CHECK: %[[ZP_TOTAL_4:.*]] = chlo.broadcast_subtract %[[RES_ZP]], %[[ZP_TOTAL_3]]
   // CHECK-SAME: (tensor<i32>, tensor<2x5x8xi32>) -> tensor<2x5x8xi32>
-  // CHECK: chlo.broadcast_add %[[RES_INT]], %[[ZP_TOTAL_6]]
+  // CHECK: chlo.broadcast_add %[[RES_INT]], %[[ZP_TOTAL_4]]
 
   %0 = "mhlo.dot_general" (%arg0, %arg1) {
     dot_dimension_numbers = #mhlo.dot<
@@ -776,7 +756,7 @@ func.func @dot_general(
       rhs_contracting_dimensions = [0]
     >} : (
       tensor<2x5x6x!quant.uniform<i8:f32, 2.000000e+00:3>>,
-      tensor<6x8x2x!quant.uniform<i8:f32, 1.000000e+00:5>>
+      tensor<6x8x2x!quant.uniform<i8:f32, 1.000000e+00:0>>
     ) -> tensor<2x5x8x!quant.uniform<i32:f32, 4.000000e+00:7>>
   return %0 : tensor<2x5x8x!quant.uniform<i32:f32, 4.000000e+00:7>>
 }
@@ -786,7 +766,7 @@ func.func @dot_general(
 // CHECK-LABEL: func @dot_general_combined_scale_1
 func.func @dot_general_combined_scale_1(
     %arg0: tensor<2x5x6x!quant.uniform<i8:f32, 2.000000e+00:3>>,
-    %arg1: tensor<6x8x2x!quant.uniform<i8:f32, 3.000000e+00:5>>
+    %arg1: tensor<6x8x2x!quant.uniform<i8:f32, 3.000000e+00:0>>
   ) -> tensor<2x5x8x!quant.uniform<i32:f32, 6.000000e+00:7>> {
   // CHECK: %[[DOT_RES:.*]] = "mhlo.dot_general"
   // CHECK-SAME: lhs_batching_dimensions = [0]
@@ -794,22 +774,7 @@ func.func @dot_general_combined_scale_1(
   // CHECK-SAME: lhs_contracting_dimensions = [2]
   // CHECK-SAME: rhs_contracting_dimensions = [0]
 
-  // Zero point offset contribution from LHS tensor * RHS ZP.
-
-  // CHECK: %[[LHS_I32:.*]] = mhlo.convert %[[LHS:.*]] : (tensor<2x5x6xi8>)
-  // CHECK-SAME: -> tensor<2x5x6xi32>
-  // CHECK: %[[LHS_REDUCE_INIT:.*]] = mhlo.constant dense<0> : tensor<i32>
-  // CHECK: %[[LHS_REDUCE:.*]] = mhlo.reduce(%[[LHS_I32]] init: %[[LHS_REDUCE_INIT]])
-  // CHECK-SAME: applies mhlo.add across dimensions = [2]
-  // CHECK-SAME: (tensor<2x5x6xi32>, tensor<i32>)
-  // CHECK-SAME: -> tensor<2x5xi32>
-  // CHECK: %[[RHS_ZP:.*]] = mhlo.constant dense<5> : tensor<i32>
-  // CHECK: %[[LHS_ZP_CONTRIB:.*]] = chlo.broadcast_multiply
-  // CHECK-SAME: %[[LHS_REDUCE]], %[[RHS_ZP]] :
-  // CHECK-SAME: (tensor<2x5xi32>, tensor<i32>) -> tensor<2x5xi32>
-  // CHECK: %[[LHS_ZP_BCAST:.*]] = "mhlo.broadcast_in_dim"(%[[LHS_ZP_CONTRIB]])
-  // CHECK-SAME: broadcast_dimensions = dense<[0, 1]>
-  // CHECK-SAME: (tensor<2x5xi32>) -> tensor<2x5x8xi32>
+  // Zero point offset contribution from LHS tensor * RHS ZP is 0 and skipped.
 
   // Zero point offset contribution from RHS tensor * LHS ZP.
 
@@ -827,21 +792,11 @@ func.func @dot_general_combined_scale_1(
   // CHECK: %[[RHS_ZP_BCAST:.*]] = "mhlo.broadcast_in_dim"(%[[RHS_ZP_CONTRIB]])
   // CHECK-SAME: broadcast_dimensions = dense<[2, 0]>
   // CHECK-SAME: (tensor<8x2xi32>) -> tensor<2x5x8xi32>
-  // CHECK: %[[ZP_TOTAL_1:.*]] = mhlo.add %[[LHS_ZP_BCAST]], %[[RHS_ZP_BCAST]]
-
-  // Zero point offset contribution from LHS ZP * RHS ZP.
-
-  // CHECK: %[[ZPS:.*]] = mhlo.constant dense<90> : tensor<i32>
-  // CHECK: %[[ZP_TOTAL_2:.*]] = chlo.broadcast_subtract %[[ZP_TOTAL_1]], %[[ZPS]]
-  // CHECK-SAME: (tensor<2x5x8xi32>, tensor<i32>) -> tensor<2x5x8xi32>
-
-  // Combine dot result with zero point offset and output final result.
-  // Do not multiply by combined scale since it is 1.0 and thus no-op.
 
   // CHECK: %[[RES_ZP:.*]] = mhlo.constant dense<7> : tensor<i32>
-  // CHECK: %[[ZP_TOTAL_3:.*]] = chlo.broadcast_subtract %[[RES_ZP]], %[[ZP_TOTAL_2]]
+  // CHECK: %[[ZP_TOTAL_1:.*]] = chlo.broadcast_subtract %[[RES_ZP]], %[[RHS_ZP_BCAST]]
   // CHECK-SAME: (tensor<i32>, tensor<2x5x8xi32>) -> tensor<2x5x8xi32>
-  // CHECK: chlo.broadcast_add %[[DOT_RES]], %[[ZP_TOTAL_3]]
+  // CHECK: chlo.broadcast_add %[[DOT_RES]], %[[ZP_TOTAL_1]]
 
   %0 = "mhlo.dot_general" (%arg0, %arg1) {
     dot_dimension_numbers = #mhlo.dot<
@@ -851,7 +806,7 @@ func.func @dot_general_combined_scale_1(
       rhs_contracting_dimensions = [0]
     >} : (
       tensor<2x5x6x!quant.uniform<i8:f32, 2.000000e+00:3>>,
-      tensor<6x8x2x!quant.uniform<i8:f32, 3.000000e+00:5>>
+      tensor<6x8x2x!quant.uniform<i8:f32, 3.000000e+00:0>>
     ) -> tensor<2x5x8x!quant.uniform<i32:f32, 6.000000e+00:7>>
   return %0 : tensor<2x5x8x!quant.uniform<i32:f32, 6.000000e+00:7>>
 }
@@ -861,7 +816,7 @@ func.func @dot_general_combined_scale_1(
 // CHECK-LABEL: func @dot_general_multiple_batching_dims
 func.func @dot_general_multiple_batching_dims(
     %arg0: tensor<2x5x3x7x6x!quant.uniform<i8:f32, 2.000000e+00:3>>,
-    %arg1: tensor<6x2x7x8x3x!quant.uniform<i8:f32, 1.000000e+00:5>>
+    %arg1: tensor<6x2x7x8x3x!quant.uniform<i8:f32, 1.000000e+00:0>>
   ) -> tensor<2x3x5x8x!quant.uniform<i32:f32, 4.000000e+00:7>> {
   // CHECK: %[[DOT_RES:.*]] = "mhlo.dot_general"
   // CHECK-SAME: lhs_batching_dimensions = [0, 2]
@@ -869,22 +824,6 @@ func.func @dot_general_multiple_batching_dims(
   // CHECK-SAME: lhs_contracting_dimensions = [4, 3]
   // CHECK-SAME: rhs_contracting_dimensions = [0, 2]>}
 
-  // Zero point offset contribution from LHS tensor * RHS ZP.
-
-  // CHECK: %[[LHS_I32:.*]] = mhlo.convert %[[LHS:.*]] : (tensor<2x5x3x7x6xi8>)
-  // CHECK-SAME: -> tensor<2x5x3x7x6xi32>
-  // CHECK: %[[LHS_REDUCE_INIT:.*]] = mhlo.constant dense<0> : tensor<i32>
-  // CHECK: %[[LHS_REDUCE:.*]] = mhlo.reduce(%[[LHS_I32]] init: %[[LHS_REDUCE_INIT]])
-  // CHECK-SAME: applies mhlo.add across dimensions = [4, 3]
-  // CHECK-SAME: (tensor<2x5x3x7x6xi32>, tensor<i32>)
-  // CHECK-SAME: -> tensor<2x5x3xi32>
-  // CHECK: %[[RHS_ZP:.*]] = mhlo.constant dense<5> : tensor<i32>
-  // CHECK: %[[LHS_ZP_CONTRIB:.*]] = chlo.broadcast_multiply
-  // CHECK-SAME: %[[LHS_REDUCE]], %[[RHS_ZP]] :
-  // CHECK-SAME: (tensor<2x5x3xi32>, tensor<i32>) -> tensor<2x5x3xi32>
-  // CHECK: %[[LHS_ZP_BCAST:.*]] = "mhlo.broadcast_in_dim"(%[[LHS_ZP_CONTRIB]])
-  // CHECK-SAME: broadcast_dimensions = dense<[0, 2, 1]>
-  // CHECK-SAME: (tensor<2x5x3xi32>) -> tensor<2x3x5x8xi32>
 
   // Zero point offset contribution from RHS tensor * LHS ZP.
 
@@ -902,13 +841,7 @@ func.func @dot_general_multiple_batching_dims(
   // CHECK: %[[RHS_ZP_BCAST:.*]] = "mhlo.broadcast_in_dim"(%[[RHS_ZP_CONTRIB]])
   // CHECK-SAME: broadcast_dimensions = dense<[0, 3, 1]>
   // CHECK-SAME: (tensor<2x8x3xi32>) -> tensor<2x3x5x8xi32>
-  // CHECK: %[[ZP_TOTAL_1:.*]] = mhlo.add %[[LHS_ZP_BCAST]], %[[RHS_ZP_BCAST]]
 
-  // Zero point offset contribution from LHS ZP * RHS ZP.
-
-  // CHECK: %[[ZPS:.*]] = mhlo.constant dense<630> : tensor<i32>
-  // CHECK: %[[ZP_TOTAL_2:.*]] = chlo.broadcast_subtract %[[ZP_TOTAL_1]], %[[ZPS]]
-  // CHECK-SAME: (tensor<2x3x5x8xi32>, tensor<i32>) -> tensor<2x3x5x8xi32>
 
   // Combine dot result with zero point offset and output final result.
 
@@ -920,17 +853,17 @@ func.func @dot_general_multiple_batching_dims(
   // CHECK: %[[RES_INT:.*]] = mhlo.convert %[[RES_FP_1]]
   // CHECK-SAME: (tensor<2x3x5x8xf32>) -> tensor<2x3x5x8xi32>
 
-  // CHECK: %[[ZP_TOTAL_3:.*]] = mhlo.convert %[[ZP_TOTAL_2]]
+  // CHECK: %[[ZP_TOTAL_1:.*]] = mhlo.convert %[[RHS_ZP_BCAST]]
   // CHECK-SAME: (tensor<2x3x5x8xi32>) -> tensor<2x3x5x8xf32>
-  // CHECK: %[[ZP_TOTAL_4:.*]] = chlo.broadcast_multiply
-  // CHECK-SAME: %[[ZP_TOTAL_3:.*]], %[[COMBINED_SCALE]]
-  // CHECK: %[[ZP_TOTAL_5:.*]] = mhlo.convert %[[ZP_TOTAL_4]]
+  // CHECK: %[[ZP_TOTAL_2:.*]] = chlo.broadcast_multiply
+  // CHECK-SAME: %[[ZP_TOTAL_1:.*]], %[[COMBINED_SCALE]]
+  // CHECK: %[[ZP_TOTAL_3:.*]] = mhlo.convert %[[ZP_TOTAL_2]]
   // CHECK-SAME: (tensor<2x3x5x8xf32>) -> tensor<2x3x5x8xi32>
 
   // CHECK: %[[RES_ZP:.*]] = mhlo.constant dense<7> : tensor<i32>
-  // CHECK: %[[ZP_TOTAL_6:.*]] = chlo.broadcast_subtract %[[RES_ZP]], %[[ZP_TOTAL_5]]
+  // CHECK: %[[ZP_TOTAL_4:.*]] = chlo.broadcast_subtract %[[RES_ZP]], %[[ZP_TOTAL_3]]
   // CHECK-SAME: (tensor<i32>, tensor<2x3x5x8xi32>) -> tensor<2x3x5x8xi32>
-  // CHECK: chlo.broadcast_add %[[RES_INT]], %[[ZP_TOTAL_6]]
+  // CHECK: chlo.broadcast_add %[[RES_INT]], %[[ZP_TOTAL_4]]
 
   %0 = "mhlo.dot_general" (%arg0, %arg1) {
     dot_dimension_numbers = #mhlo.dot<
@@ -940,7 +873,7 @@ func.func @dot_general_multiple_batching_dims(
       rhs_contracting_dimensions = [0, 2]
     >} : (
       tensor<2x5x3x7x6x!quant.uniform<i8:f32, 2.000000e+00:3>>,
-      tensor<6x2x7x8x3x!quant.uniform<i8:f32, 1.000000e+00:5>>
+      tensor<6x2x7x8x3x!quant.uniform<i8:f32, 1.000000e+00:0>>
     ) -> tensor<2x3x5x8x!quant.uniform<i32:f32, 4.000000e+00:7>>
   return %0 : tensor<2x3x5x8x!quant.uniform<i32:f32, 4.000000e+00:7>>
 }
@@ -1058,7 +991,7 @@ func.func @dot_general_zero_zp(
 // CHECK-LABEL: func @dot_general_multiple_dynamic_dims
 func.func @dot_general_multiple_dynamic_dims(
     %arg0: tensor<?x?x3x?x6x!quant.uniform<i8:f32, 2.000000e+00:3>>,
-    %arg1: tensor<6x?x?x8x3x!quant.uniform<i8:f32, 1.000000e+00:5>>
+    %arg1: tensor<6x?x?x8x3x!quant.uniform<i8:f32, 1.000000e+00:0>>
   ) -> tensor<?x3x?x8x!quant.uniform<i32:f32, 4.000000e+00:7>> {
   // CHECK: %[[DOT_RES:.*]] = "mhlo.dot_general"
   // CHECK-SAME: lhs_batching_dimensions = [0, 2]
@@ -1068,17 +1001,17 @@ func.func @dot_general_multiple_dynamic_dims(
 
   // Zero point offset contribution from LHS tensor * RHS ZP.
 
-  // CHECK: %[[LHS_I32:.*]] = mhlo.convert %[[LHS:.*]] : (tensor<?x?x3x?x6xi8>)
-  // CHECK-SAME: -> tensor<?x?x3x?x6xi32>
-  // CHECK: %[[LHS_REDUCE_INIT:.*]] = mhlo.constant dense<0> : tensor<i32>
-  // CHECK: %[[LHS_REDUCE:.*]] = mhlo.reduce(%[[LHS_I32]] init: %[[LHS_REDUCE_INIT]])
-  // CHECK-SAME: applies mhlo.add across dimensions = [4, 3]
-  // CHECK-SAME: (tensor<?x?x3x?x6xi32>, tensor<i32>)
-  // CHECK-SAME: -> tensor<?x?x3xi32>
-  // CHECK: %[[RHS_ZP:.*]] = mhlo.constant dense<5> : tensor<i32>
-  // CHECK: %[[LHS_ZP_CONTRIB:.*]] = chlo.broadcast_multiply
-  // CHECK-SAME: %[[LHS_REDUCE]], %[[RHS_ZP]] :
-  // CHECK-SAME: (tensor<?x?x3xi32>, tensor<i32>) -> tensor<?x?x3xi32>
+  // CHECK: %[[RHS_I32:.*]] = mhlo.convert %[[RHS:.*]] : (tensor<6x?x?x8x3xi8>)
+  // CHECK-SAME: -> tensor<6x?x?x8x3xi32>
+  // CHECK: %[[RHS_REDUCE_INIT:.*]] = mhlo.constant dense<0> : tensor<i32>
+  // CHECK: %[[RHS_REDUCE:.*]] = mhlo.reduce(%[[RHS_I32]] init: %[[RHS_REDUCE_INIT]])
+  // CHECK-SAME: applies mhlo.add across dimensions = [0, 2]
+  // CHECK-SAME: (tensor<6x?x?x8x3xi32>, tensor<i32>)
+  // CHECK-SAME: -> tensor<?x8x3xi32>
+  // CHECK: %[[RHS_ZP:.*]] = mhlo.constant dense<3> : tensor<i32>
+  // CHECK: %[[RHS_ZP_CONTRIB:.*]] = chlo.broadcast_multiply
+  // CHECK-SAME: %[[RHS_REDUCE]], %[[RHS_ZP]] :
+  // CHECK-SAME: (tensor<?x8x3xi32>, tensor<i32>) -> tensor<?x8x3xi32>
 
   // Calculate output dynamic dims.
   // CHECK: %[[DIM_1_1:.*]] = "mhlo.get_dimension_size"(%[[DOT_RES]])
@@ -1094,40 +1027,10 @@ func.func @dot_general_multiple_dynamic_dims(
   // CHECK: %[[OUTPUT_DIMS:.*]] = "mhlo.concatenate"
   // CHECK-SAME: %[[DIM_1]], %[[DIM_2]], %[[DIM_3]], %[[DIM_4]]
 
-  // CHECK: %[[LHS_ZP_BCAST:.*]] = "mhlo.dynamic_broadcast_in_dim"
-  // CHECK-SAME: (%[[LHS_ZP_CONTRIB]], %[[OUTPUT_DIMS]])
-  // CHECK-SAME: broadcast_dimensions = dense<[0, 2, 1]>
-  // CHECK-SAME: (tensor<?x?x3xi32>, tensor<4xi64>) -> tensor<?x3x?x8xi32>
-
-  // Zero point offset contribution from RHS tensor * LHS ZP.
-
-  // CHECK: %[[RHS_I32:.*]] = mhlo.convert %[[RHS:.*]] : (tensor<6x?x?x8x3xi8>)
-  // CHECK-SAME: -> tensor<6x?x?x8x3xi32>
-  // CHECK: %[[RHS_REDUCE_INIT:.*]] = mhlo.constant dense<0> : tensor<i32>
-  // CHECK: %[[RHS_REDUCE:.*]] = mhlo.reduce(%[[RHS_I32]] init: %[[RHS_REDUCE_INIT]])
-  // CHECK-SAME: applies mhlo.add across dimensions = [0, 2]
-  // CHECK-SAME: (tensor<6x?x?x8x3xi32>, tensor<i32>)
-  // CHECK-SAME: -> tensor<?x8x3xi32>
-  // CHECK: %[[RHS_ZP:.*]] = mhlo.constant dense<3> : tensor<i32>
-  // CHECK: %[[RHS_ZP_CONTRIB:.*]] = chlo.broadcast_multiply
-  // CHECK-SAME: %[[RHS_REDUCE]], %[[RHS_ZP]] :
-  // CHECK-SAME: (tensor<?x8x3xi32>, tensor<i32>) -> tensor<?x8x3xi32>
-
   // CHECK: %[[RHS_ZP_BCAST:.*]] = "mhlo.dynamic_broadcast_in_dim"
   // CHECK-SAME: (%[[RHS_ZP_CONTRIB]], %[[OUTPUT_DIMS]])
   // CHECK-SAME: broadcast_dimensions = dense<[0, 3, 1]>
   // CHECK-SAME: (tensor<?x8x3xi32>, tensor<4xi64>) -> tensor<?x3x?x8xi32>
-  // CHECK: %[[ZP_TOTAL_1:.*]] = mhlo.add %[[LHS_ZP_BCAST]], %[[RHS_ZP_BCAST]]
-
-  // Zero point offset contribution from LHS ZP * RHS ZP.
-
-  // CHECK: %[[ZPS_INIT:.*]] = mhlo.constant dense<1> : tensor<i32>
-  // CHECK: %[[DYN_DIM:.*]] = "mhlo.get_dimension_size"(%[[RHS]])
-  // CHECK: %[[ZPS_1:.*]] = mhlo.multiply %[[ZPS_INIT]], %[[DYN_DIM]]
-  // CHECK: %[[STATIC_DIM:.*]] = mhlo.constant dense<90> : tensor<i32>
-  // CHECK: %[[ZPS:.*]] = mhlo.multiply %[[STATIC_DIM]], %[[ZPS_1]]
-  // CHECK: %[[ZP_TOTAL_2:.*]] = chlo.broadcast_subtract %[[ZP_TOTAL_1]], %[[ZPS]]
-  // CHECK-SAME: (tensor<?x3x?x8xi32>, tensor<i32>) -> tensor<?x3x?x8xi32>
 
   // Combine dot result with zero point offset and output final result.
 
@@ -1139,17 +1042,17 @@ func.func @dot_general_multiple_dynamic_dims(
   // CHECK: %[[RES_INT:.*]] = mhlo.convert %[[RES_FP_1]]
   // CHECK-SAME: (tensor<?x3x?x8xf32>) -> tensor<?x3x?x8xi32>
 
-  // CHECK: %[[ZP_TOTAL_3:.*]] = mhlo.convert %[[ZP_TOTAL_2]]
+  // CHECK: %[[ZP_TOTAL_1:.*]] = mhlo.convert %[[RHS_ZP_BCAST]]
   // CHECK-SAME: (tensor<?x3x?x8xi32>) -> tensor<?x3x?x8xf32>
-  // CHECK: %[[ZP_TOTAL_4:.*]] = chlo.broadcast_multiply
-  // CHECK-SAME: %[[ZP_TOTAL_3:.*]], %[[COMBINED_SCALE]]
-  // CHECK: %[[ZP_TOTAL_5:.*]] = mhlo.convert %[[ZP_TOTAL_4]]
+  // CHECK: %[[ZP_TOTAL_2:.*]] = chlo.broadcast_multiply
+  // CHECK-SAME: %[[ZP_TOTAL_1:.*]], %[[COMBINED_SCALE]]
+  // CHECK: %[[ZP_TOTAL_3:.*]] = mhlo.convert %[[ZP_TOTAL_2]]
   // CHECK-SAME: (tensor<?x3x?x8xf32>) -> tensor<?x3x?x8xi32>
 
   // CHECK: %[[RES_ZP:.*]] = mhlo.constant dense<7> : tensor<i32>
-  // CHECK: %[[ZP_TOTAL_6:.*]] = chlo.broadcast_subtract %[[RES_ZP]], %[[ZP_TOTAL_5]]
+  // CHECK: %[[ZP_TOTAL_4:.*]] = chlo.broadcast_subtract %[[RES_ZP]], %[[ZP_TOTAL_3]]
   // CHECK-SAME: (tensor<i32>, tensor<?x3x?x8xi32>) -> tensor<?x3x?x8xi32>
-  // CHECK: chlo.broadcast_add %[[RES_INT]], %[[ZP_TOTAL_6]]
+  // CHECK: chlo.broadcast_add %[[RES_INT]], %[[ZP_TOTAL_4]]
 
   %0 = "mhlo.dot_general" (%arg0, %arg1) {
     dot_dimension_numbers = #mhlo.dot<
@@ -1159,7 +1062,7 @@ func.func @dot_general_multiple_dynamic_dims(
       rhs_contracting_dimensions = [0, 2]
     >} : (
       tensor<?x?x3x?x6x!quant.uniform<i8:f32, 2.000000e+00:3>>,
-      tensor<6x?x?x8x3x!quant.uniform<i8:f32, 1.000000e+00:5>>
+      tensor<6x?x?x8x3x!quant.uniform<i8:f32, 1.000000e+00:0>>
     ) -> tensor<?x3x?x8x!quant.uniform<i32:f32, 4.000000e+00:7>>
   return %0 : tensor<?x3x?x8x!quant.uniform<i32:f32, 4.000000e+00:7>>
 }
@@ -1811,14 +1714,13 @@ func.func @dot_general_hybrid_per_channel(
 // CHECK-SAME: %[[ARG1:.*]]: tensor<2x2xi8>
 func.func @dot_general_hybrid_per_channel_asymmetric(
     %arg0: tensor<3x2xf32>,
-    %arg1: tensor<2x2x!quant.uniform<i8<-127:127>:f32:1, {3.000000e+00:10, 4.000000e+00:20}>>
+    %arg1: tensor<2x2x!quant.uniform<i8<-127:127>:f32:1, {3.000000e+00:0, 4.000000e+00:0}>>
   ) -> tensor<3x2xf32> {
   // CHECK-DAG: %[[BARRIER:.*]] = mhlo.optimization_barrier %[[ARG1]] : tensor<2x2xi8>
   // CHECK-DAG: %[[SCALES:.*]] = mhlo.constant dense<[3.000000e+00, 4.000000e+00]> : tensor<2xf32>
-  // CHECK-DAG: %[[ZPS:.*]] = mhlo.constant dense<[1.000000e+01, 2.000000e+01]> : tensor<2xf32>
+  // CHECK-DAG: %[[ZPS:.*]] = mhlo.constant dense<0.000000e+00> : tensor<2xf32>
   // CHECK-DAG: %[[CONVERT:.*]] = mhlo.convert %[[BARRIER]] : (tensor<2x2xi8>) -> tensor<2x2xf32>
-  // CHECK: %[[SUB:.*]] = chlo.broadcast_subtract %[[CONVERT]], %[[ZPS]] {broadcast_dimensions = array<i64: 1>} : (tensor<2x2xf32>, tensor<2xf32>) -> tensor<2x2xf32>
-  // CHECK: %[[MUL:.*]] = chlo.broadcast_multiply %[[SUB]], %[[SCALES]] {broadcast_dimensions = array<i64: 1>} : (tensor<2x2xf32>, tensor<2xf32>) -> tensor<2x2xf32>
+  // CHECK: %[[MUL:.*]] = chlo.broadcast_multiply %[[CONVERT]], %[[SCALES]] {broadcast_dimensions = array<i64: 1>} : (tensor<2x2xf32>, tensor<2xf32>) -> tensor<2x2xf32>
   // CHECK: %[[DOT:.*]] = "mhlo.dot_general"(%[[ARG0]], %[[MUL]])
   // CHECK-SAME: (tensor<3x2xf32>, tensor<2x2xf32>) -> tensor<3x2xf32>
   // CHECK: return %[[DOT]]
@@ -1827,7 +1729,7 @@ func.func @dot_general_hybrid_per_channel_asymmetric(
       dot_dimension_numbers = #mhlo.dot<lhs_contracting_dimensions = [1],
       rhs_contracting_dimensions = [0]>} : (
     tensor<3x2xf32>,
-    tensor<2x2x!quant.uniform<i8<-127:127>:f32:1, {3.000000e+00:10, 4.000000e+00:20}>>
+    tensor<2x2x!quant.uniform<i8<-127:127>:f32:1, {3.000000e+00:0, 4.000000e+00:0}>>
   ) -> tensor<3x2xf32>
   return %0 : tensor<3x2xf32>
 }
@@ -1989,9 +1891,9 @@ func.func @conv2d_hybrid_result_not_float(
 
 // -----
 
-func.func @dot_general_hybrid_result_not_float(
-    %arg0: tensor<2x5x6xf32>,
-    %arg1: tensor<6x8x2x!quant.uniform<i8:f32, 1.000000e+00:5>>) {
+func.func @dot_general_non_hybrid_result_not_float(
+  %arg0: tensor<2x5x6x!quant.uniform<i8:f32, 1.000000e+00:0>>,
+  %arg1: tensor<6x8x2x!quant.uniform<i8:f32:2, {1.000000e+00:0, 1.000000e+00:0}>>) {
   // expected-error@+2 {{Invalid input/output type for Dot/Convolution op}}
   // expected-error@+1 {{failed to legalize operation 'mhlo.dot_general' that was explicitly marked illegal}}
   %0 = "mhlo.dot_general" (%arg0, %arg1) {
@@ -2001,8 +1903,8 @@ func.func @dot_general_hybrid_result_not_float(
       lhs_contracting_dimensions = [2],
       rhs_contracting_dimensions = [0]
     >} : (
-      tensor<2x5x6xf32>,
-      tensor<6x8x2x!quant.uniform<i8:f32, 1.000000e+00:5>>
+      tensor<2x5x6x!quant.uniform<i8:f32, 1.000000e+00:0>>,
+      tensor<6x8x2x!quant.uniform<i8:f32:2, {1.000000e+00:0, 1.000000e+00:0}>>
     ) -> tensor<2x5x8x!quant.uniform<i8:f32, 4.000000e+00:7>>
   return
 }
diff --git a/third_party/xla/xla/mlir_hlo/tests/Dialect/mhlo/ops.mlir b/third_party/xla/xla/mlir_hlo/tests/Dialect/mhlo/ops.mlir
index 3f6e71bee0b7ff..b4bdfba82540b2 100644
--- a/third_party/xla/xla/mlir_hlo/tests/Dialect/mhlo/ops.mlir
+++ b/third_party/xla/xla/mlir_hlo/tests/Dialect/mhlo/ops.mlir
@@ -3449,7 +3449,7 @@ func.func @compatible_shapes(%arg0: tensor<?xf32>, %shape: tensor<2xindex>) -> t
 // -----
 
 func.func @incompatible_shapes(%arg0: tensor<?xf32>, %shape: tensor<2xindex>) -> tensor<?xf32> {
-  // expected-error @+1 {{output should have a rank equal to the number of elements in output_shape}}
+  // expected-error @+1 {{result should have a rank equal to the number of elements in output_shape}}
   %0 = "mhlo.dynamic_reshape"(%arg0, %shape) : (tensor<?xf32>, tensor<2xindex>) -> tensor<?xf32>
   func.return %0 : tensor<?xf32>
 }
@@ -4061,7 +4061,7 @@ func.func @dynamic_gather(%operand : tensor<2x4x9xi32>, %start_indices : tensor<
 
 // -----
 
-func.func @dynamic_gather(%operand : tensor<?x?x?xi32>, %start_indices : tensor<?x?x?xi32>, %slice_sizes : tensor<?xi32>) -> tensor<?x?x?xi32> {
+func.func @dynamic_gather(%operand : tensor<?x?x?xi32>, %start_indices : tensor<?x?x?xi32>, %slice_sizes : tensor<3xi32>) -> tensor<?x?x?xi32> {
   %res = "mhlo.dynamic_gather"(%operand, %start_indices, %slice_sizes) {
     dimension_numbers = #mhlo.gather<
       collapsed_slice_dims = [0, 1],
@@ -4070,13 +4070,13 @@ func.func @dynamic_gather(%operand : tensor<?x?x?xi32>, %start_indices : tensor<
       start_index_map = [0, 1]
     >,
     indices_are_sorted = false
-  } : (tensor<?x?x?xi32>, tensor<?x?x?xi32>, tensor<?xi32>) -> tensor<?x?x?xi32>
+  } : (tensor<?x?x?xi32>, tensor<?x?x?xi32>, tensor<3xi32>) -> tensor<?x?x?xi32>
   func.return %res : tensor<?x?x?xi32>
 }
 
 // -----
 
-func.func @dynamic_gather(%operand : tensor<2x4x9xi32>, %start_indices : tensor<?x?x?xi32>, %slice_sizes : tensor<?xi32>) -> tensor<?x?x?xi32> {
+func.func @dynamic_gather(%operand : tensor<2x4x9xi32>, %start_indices : tensor<?x?x?xi32>, %slice_sizes : tensor<3xi32>) -> tensor<?x?x?xi32> {
   %res = "mhlo.dynamic_gather"(%operand, %start_indices, %slice_sizes) {
     dimension_numbers = #mhlo.gather<
       collapsed_slice_dims = [0, 1],
@@ -4085,14 +4085,31 @@ func.func @dynamic_gather(%operand : tensor<2x4x9xi32>, %start_indices : tensor<
       start_index_map = [0, 1]
     >,
     indices_are_sorted = false
-  } : (tensor<2x4x9xi32>, tensor<?x?x?xi32>, tensor<?xi32>) -> tensor<?x?x?xi32>
+  } : (tensor<2x4x9xi32>, tensor<?x?x?xi32>, tensor<3xi32>) -> tensor<?x?x?xi32>
   func.return %res : tensor<?x?x?xi32>
 }
 
 // -----
 
-func.func @dynamic_gather(%operand : tensor<?x?x?xi32>, %start_indices : tensor<?x?x?xi32>, %slice_sizes : tensor<?xi32>) -> tensor<?x?x?xi32> {
-  // @expected-error@+2 {{'mhlo.dynamic_gather' op failed to infer returned types}}
+func.func @dynamic_gather_c1(%operand : tensor<2x4x9xi32>, %start_indices : tensor<1x5x2xi32>, %slice_sizes : tensor<3xi32>) -> tensor<1x5x8xi32> {
+  // expected-error@+2 {{failed to infer returned types}}
+  // expected-error@+1 {{offset_dims size (2) plus collapse_slice_dims size (2) is not equal to operand rank (3)}}
+  %res = "mhlo.dynamic_gather"(%operand, %start_indices, %slice_sizes) {
+    dimension_numbers = #mhlo.gather<
+      offset_dims = [1, 2],
+      collapsed_slice_dims = [0, 1],
+      start_index_map = [0, 1],
+      index_vector_dim = 2
+    >,
+    indices_are_sorted = false
+  } : (tensor<2x4x9xi32>, tensor<1x5x2xi32>, tensor<3xi32>) -> tensor<1x5x8xi32>
+  func.return %res : tensor<1x5x8xi32>
+}
+
+// -----
+
+func.func @dynamic_gather_c2(%operand : tensor<?x?x?xi32>, %start_indices : tensor<?x?x?xi32>, %slice_sizes : tensor<3xi32>) -> tensor<?xi32> {
+  // expected-error@+2 {{failed to infer returned types}}
   // expected-error@+1 {{index_vector_dim 4 is out of bounds for start indices with rank 3}}
   %res = "mhlo.dynamic_gather"(%operand, %start_indices, %slice_sizes) {
     dimension_numbers = #mhlo.gather<
@@ -4102,65 +4119,221 @@ func.func @dynamic_gather(%operand : tensor<?x?x?xi32>, %start_indices : tensor<
       start_index_map = [0, 1]
     >,
     indices_are_sorted = false
-  } : (tensor<?x?x?xi32>, tensor<?x?x?xi32>, tensor<?xi32>) -> tensor<?x?x?xi32>
+  } : (tensor<?x?x?xi32>, tensor<?x?x?xi32>, tensor<3xi32>) -> tensor<?xi32>
+  func.return %res : tensor<?xi32>
+}
+
+// -----
+
+func.func @dynamic_gather_c3(%operand : tensor<?x?x?xi32>, %start_indices : tensor<?x?x2xi32>, %slice_sizes : tensor<3xi32>) -> tensor<?xi32> {
+  // expected-error@+2 {{failed to infer returned types}}
+  // expected-error@+1 {{start_index_map size (1) is not equal to size of index dimension (2) of start_indices (2)}}
+  %res = "mhlo.dynamic_gather"(%operand, %start_indices, %slice_sizes) {
+    dimension_numbers = #mhlo.gather<
+      collapsed_slice_dims = [0, 1],
+      index_vector_dim = 2,
+      offset_dims = [2],
+      start_index_map = [0]
+    >,
+    indices_are_sorted = false
+  } : (tensor<?x?x?xi32>, tensor<?x?x2xi32>, tensor<3xi32>) -> tensor<?xi32>
+  func.return %res : tensor<?xi32>
+}
+
+// -----
+
+func.func @dynamic_gather_c4(%operand : tensor<2x4x9xi32>, %start_indices : tensor<?x?x?xi32>, %slice_sizes : tensor<3xi32>) -> tensor<?x?x?xi32> {
+  // expected-error@+2 {{failed to infer returned types}}
+  // expected-error@+1 {{expects offset_dims to be sorted, got: [2, 1]}}
+  %res = "mhlo.dynamic_gather"(%operand, %start_indices, %slice_sizes) {
+    dimension_numbers = #mhlo.gather<
+      collapsed_slice_dims = [0, 1],
+      index_vector_dim = 2,
+      offset_dims = [2, 1],
+      start_index_map = [0, 1]
+    >,
+    indices_are_sorted = false
+  } : (tensor<2x4x9xi32>, tensor<?x?x?xi32>, tensor<3xi32>) -> tensor<?x?x?xi32>
   func.return %res : tensor<?x?x?xi32>
 }
 
 // -----
 
-func.func @dynamic_gather(%operand : tensor<?x?x?xi32>, %start_indices : tensor<?x?x?xi32>, %slice_sizes : tensor<?xi32>) -> tensor<?x?x?xi32> {
-  // @expected-error@+2 {{'mhlo.dynamic_gather' op failed to infer returned types}}
-  // expected-error@+1 {{offset_dims size (2) plus collapse_slice_dims size (2) is not equal to operand rank (3)}}
+func.func @dynamic_gather_c5(%operand : tensor<2x4x9xi32>, %start_indices : tensor<?x?x?xi32>, %slice_sizes : tensor<3xi32>) -> tensor<?x?x?xi32> {
+  // expected-error@+2 {{failed to infer returned types}}
+  // expected-error@+1 {{offset_dims[0]: -1 is out of bounds for implied result rank 3}}
   %res = "mhlo.dynamic_gather"(%operand, %start_indices, %slice_sizes) {
     dimension_numbers = #mhlo.gather<
       collapsed_slice_dims = [0, 1],
       index_vector_dim = 2,
-      offset_dims = [1, 2],
+      offset_dims = [-1],
       start_index_map = [0, 1]
     >,
     indices_are_sorted = false
-  } : (tensor<?x?x?xi32>, tensor<?x?x?xi32>, tensor<?xi32>) -> tensor<?x?x?xi32>
+  } : (tensor<2x4x9xi32>, tensor<?x?x?xi32>, tensor<3xi32>) -> tensor<?x?x?xi32>
   func.return %res : tensor<?x?x?xi32>
 }
 
 // -----
 
-func.func @dynamic_gather(%operand : tensor<?x?x?xi32>, %start_indices : tensor<?x?x2xi32>, %slice_sizes : tensor<?xi32>) -> tensor<?x?x?xi32> {
-  // @expected-error@+2 {{'mhlo.dynamic_gather' op failed to infer returned types}}
-  // expected-error@+1 {{start_index_map size (1) is not equal to size of index dimension (2) of start_indices (2)}}
+func.func @dynamic_gather_c5(%operand : tensor<2x4x9xi32>, %start_indices : tensor<?x?x?xi32>, %slice_sizes : tensor<3xi32>) -> tensor<?x?x?xi32> {
+  // expected-error@+2 {{failed to infer returned types}}
+  // expected-error@+1 {{offset_dims[0]: 3 is out of bounds for implied result rank 3}}
   %res = "mhlo.dynamic_gather"(%operand, %start_indices, %slice_sizes) {
     dimension_numbers = #mhlo.gather<
       collapsed_slice_dims = [0, 1],
       index_vector_dim = 2,
+      offset_dims = [3],
+      start_index_map = [0, 1]
+    >,
+    indices_are_sorted = false
+  } : (tensor<2x4x9xi32>, tensor<?x?x?xi32>, tensor<3xi32>) -> tensor<?x?x?xi32>
+  func.return %res : tensor<?x?x?xi32>
+}
+
+// -----
+
+func.func @dynamic_gather_c6(%operand : tensor<2x4x9xi32>, %start_indices : tensor<?x?x?xi32>, %slice_sizes : tensor<3xi32>) -> tensor<?x?x?xi32> {
+  // expected-error@+2 {{failed to infer returned types}}
+  // expected-error@+1 {{expects collapsed_slice_dims to be sorted, got: [1, 0]}}
+  %res = "mhlo.dynamic_gather"(%operand, %start_indices, %slice_sizes) {
+    dimension_numbers = #mhlo.gather<
+      collapsed_slice_dims = [1, 0],
+      index_vector_dim = 2,
       offset_dims = [2],
-      start_index_map = [0]
+      start_index_map = [0, 1]
     >,
     indices_are_sorted = false
-  } : (tensor<?x?x?xi32>, tensor<?x?x2xi32>, tensor<?xi32>) -> tensor<?x?x?xi32>
+  } : (tensor<2x4x9xi32>, tensor<?x?x?xi32>, tensor<3xi32>) -> tensor<?x?x?xi32>
   func.return %res : tensor<?x?x?xi32>
 }
 
 // -----
 
-func.func @dynamic_gather(%operand : tensor<?x?x?xi32>, %start_indices : tensor<?x?x?xi32>, %slice_sizes : tensor<?x?xi32>) -> tensor<?xi32> {
-  // @expected-error@+2 {{'mhlo.dynamic_gather' op failed to infer returned types}}
-  // expected-error@+1 {{slice_sizes.rank != 1}}
+func.func @dynamic_gather_c6(%operand : tensor<2x4x9xi32>, %start_indices : tensor<?x?x?xi32>, %slice_sizes : tensor<3xi32>) -> tensor<?x?x?xi32> {
+  // expected-error@+2 {{failed to infer returned types}}
+  // expected-error@+1 {{expects collapsed_slice_dims to not repeat, got: [1, 1]}}
   %res = "mhlo.dynamic_gather"(%operand, %start_indices, %slice_sizes) {
     dimension_numbers = #mhlo.gather<
-      collapsed_slice_dims = [0, 1],
+      collapsed_slice_dims = [1, 1],
       index_vector_dim = 2,
       offset_dims = [2],
       start_index_map = [0, 1]
     >,
     indices_are_sorted = false
-  } : (tensor<?x?x?xi32>, tensor<?x?x?xi32>, tensor<?x?xi32>) -> tensor<?x?x?xi32>
+  } : (tensor<2x4x9xi32>, tensor<?x?x?xi32>, tensor<3xi32>) -> tensor<?x?x?xi32>
   func.return %res : tensor<?x?x?xi32>
 }
 
 // -----
 
-func.func @dynamic_gather(%operand : tensor<?x?x?xi32>, %start_indices : tensor<?x?x?xi32>, %slice_sizes : tensor<2xi32>) -> tensor<?x?x?xi32> {
-  // @expected-error@+2 {{'mhlo.dynamic_gather' op failed to infer returned types}}
+func.func @dynamic_gather_c7(%operand : tensor<2x4x9xi32>, %start_indices : tensor<?x?x?xi32>, %slice_sizes : tensor<3xi32>) -> tensor<?x?x?xi32> {
+  // expected-error@+2 {{failed to infer returned types}}
+  // expected-error@+1 {{collapsed dimension -1 is out of bounds for slice_sizes.size (3)}}
+  %res = "mhlo.dynamic_gather"(%operand, %start_indices, %slice_sizes) {
+    dimension_numbers = #mhlo.gather<
+      collapsed_slice_dims = [-1, 1],
+      index_vector_dim = 2,
+      offset_dims = [2],
+      start_index_map = [0, 1]
+    >,
+    indices_are_sorted = false
+  } : (tensor<2x4x9xi32>, tensor<?x?x?xi32>, tensor<3xi32>) -> tensor<?x?x?xi32>
+  func.return %res : tensor<?x?x?xi32>
+}
+
+// -----
+
+func.func @dynamic_gather_c7(%operand : tensor<2x4x9xi32>, %start_indices : tensor<?x?x?xi32>, %slice_sizes : tensor<3xi32>) -> tensor<?x?x?xi32> {
+  // expected-error@+2 {{failed to infer returned types}}
+  // expected-error@+1 {{collapsed dimension 17 is out of bounds for slice_sizes.size (3)}}
+  %res = "mhlo.dynamic_gather"(%operand, %start_indices, %slice_sizes) {
+    dimension_numbers = #mhlo.gather<
+      collapsed_slice_dims = [0, 17],
+      index_vector_dim = 2,
+      offset_dims = [2],
+      start_index_map = [0, 1]
+    >,
+    indices_are_sorted = false
+  } : (tensor<2x4x9xi32>, tensor<?x?x?xi32>, tensor<3xi32>) -> tensor<?x?x?xi32>
+  func.return %res : tensor<?x?x?xi32>
+}
+
+// -----
+
+func.func @dynamic_gather_c8(%operand : tensor<?x?x?xi32>, %start_indices : tensor<?x?x?xi32>) -> tensor<?x?x?xi32> {
+  %slize_sizes = mhlo.constant dense<[1,1,8]> : tensor<3xi32>
+  // expected-error@+2 {{failed to infer returned types}}
+  // expected-error@+1 {{slice_sizes collapsed dimension 2 should <= 1 but got 8}}
+  %res = "mhlo.dynamic_gather"(%operand, %start_indices, %slize_sizes) {
+    dimension_numbers = #mhlo.gather<
+      offset_dims = [2],
+      collapsed_slice_dims = [0, 2],
+      start_index_map = [0, 1],
+      index_vector_dim = 2
+    >,
+    indices_are_sorted = false
+  } : (tensor<?x?x?xi32>, tensor<?x?x?xi32>, tensor<3xi32>) -> tensor<?x?x?xi32>
+  func.return %res : tensor<?x?x?xi32>
+}
+
+// -----
+
+func.func @dynamic_gather_c9(%operand : tensor<2x4x9xi32>, %start_indices : tensor<1x5x2xi32>, %slize_sizes : tensor<3xi32>) -> tensor<1x5x8xi32> {
+  // expected-error@+2 {{failed to infer returned types}}
+  // expected-error@+1 {{expects start_index_map to not repeat, got: [0, 0]}}
+  %res = "mhlo.dynamic_gather"(%operand, %start_indices, %slize_sizes) {
+    dimension_numbers = #mhlo.gather<
+      offset_dims = [2],
+      collapsed_slice_dims = [0, 1],
+      start_index_map = [0, 0],
+      index_vector_dim = 2
+    >,
+    indices_are_sorted = false
+  } : (tensor<2x4x9xi32>, tensor<1x5x2xi32>, tensor<3xi32>) -> tensor<1x5x8xi32>
+  func.return %res : tensor<1x5x8xi32>
+}
+
+// -----
+
+func.func @dynamic_gather_c10(%operand : tensor<2x4x9xi32>, %start_indices : tensor<1x5x2xi32>, %slize_sizes : tensor<3xi32>) -> tensor<1x5x8xi32> {
+  // expected-error@+2 {{failed to infer returned types}}
+  // expected-error@+1 {{start_index_map[0]: -2 is out of bounds for operand rank 3}}
+  %res = "mhlo.dynamic_gather"(%operand, %start_indices, %slize_sizes) {
+    dimension_numbers = #mhlo.gather<
+      offset_dims = [2],
+      collapsed_slice_dims = [0, 1],
+      start_index_map = [-2, -1],
+      index_vector_dim = 2
+    >,
+    slice_sizes = array<i64: 1, 1, 8>,
+    indices_are_sorted = false
+  } : (tensor<2x4x9xi32>, tensor<1x5x2xi32>, tensor<3xi32>) -> tensor<1x5x8xi32>
+  func.return %res : tensor<1x5x8xi32>
+}
+
+// -----
+
+func.func @dynamic_gather_c10(%operand : tensor<2x4x9xi32>, %start_indices : tensor<1x5x2xi32>, %slize_sizes : tensor<3xi32>) -> tensor<1x5x8xi32> {
+  // expected-error@+2 {{failed to infer returned types}}
+  // expected-error@+1 {{start_index_map[1]: 3 is out of bounds for operand rank 3}}
+  %res = "mhlo.dynamic_gather"(%operand, %start_indices, %slize_sizes) {
+    dimension_numbers = #mhlo.gather<
+      offset_dims = [2],
+      collapsed_slice_dims = [0, 1],
+      start_index_map = [0, 3],
+      index_vector_dim = 2
+    >,
+    slice_sizes = array<i64: 1, 1, 8>,
+    indices_are_sorted = false
+  } : (tensor<2x4x9xi32>, tensor<1x5x2xi32>, tensor<3xi32>) -> tensor<1x5x8xi32>
+  func.return %res : tensor<1x5x8xi32>
+}
+
+// -----
+
+func.func @dynamic_gather_c11(%operand : tensor<?x?x?xi32>, %start_indices : tensor<?x?x?xi32>, %slice_sizes : tensor<2xi32>) -> tensor<?x?x?xi32> {
+  // expected-error@+2 {{failed to infer returned types}}
   // expected-error@+1 {{slice_sizes size (2) not equal to (implied) operand rank (3)}}
   %res = "mhlo.dynamic_gather"(%operand, %start_indices, %slice_sizes) {
     dimension_numbers = #mhlo.gather<
@@ -4176,8 +4349,45 @@ func.func @dynamic_gather(%operand : tensor<?x?x?xi32>, %start_indices : tensor<
 
 // -----
 
-func.func @dynamic_gather(%operand : tensor<2x4x9xi32>, %start_indices : tensor<1x5x2xi32>, %slice_sizes : tensor<3xi32>) -> tensor<3xi32> {
-  // @expected-error@+2 {{'mhlo.dynamic_gather' op failed to infer returned types}}
+func.func @dynamic_gather_c12(%operand : tensor<?x?x2xi32>, %start_indices : tensor<?x?x?xi32>) -> tensor<?x?x?xi32> {
+  %slice_sizes = mhlo.constant dense<[1,1,-1]> : tensor<3xi32>
+  // expected-error@+2 {{failed to infer returned types}}
+  // expected-error@+1 {{slice size (-1) is out of bounds for operand dimension (2) at index 2}}
+  %res = "mhlo.dynamic_gather"(%operand, %start_indices, %slice_sizes) {
+    dimension_numbers = #mhlo.gather<
+      offset_dims = [2],
+      collapsed_slice_dims = [0, 1],
+      start_index_map = [0, 1],
+      index_vector_dim = 2
+    >,
+    indices_are_sorted = false
+  } : (tensor<?x?x2xi32>, tensor<?x?x?xi32>, tensor<3xi32>) -> tensor<?x?x?xi32>
+  func.return %res : tensor<?x?x?xi32>
+}
+
+// -----
+
+func.func @dynamic_gather_c12(%operand : tensor<?x?x2xi32>, %start_indices : tensor<?x?x?xi32>) -> tensor<?x?x?xi32> {
+  %slice_sizes = mhlo.constant dense<[1,1,8]> : tensor<3xi32>
+  // expected-error@+2 {{failed to infer returned types}}
+  // expected-error@+1 {{slice size (8) is out of bounds for operand dimension (2) at index 2}}
+  %res = "mhlo.dynamic_gather"(%operand, %start_indices, %slice_sizes) {
+    dimension_numbers = #mhlo.gather<
+      offset_dims = [2],
+      collapsed_slice_dims = [0, 1],
+      start_index_map = [0, 1],
+      index_vector_dim = 2
+    >,
+    slice_sizes = array<i64: 1, 1, 8>,
+    indices_are_sorted = false
+  } : (tensor<?x?x2xi32>, tensor<?x?x?xi32>, tensor<3xi32>) -> tensor<?x?x?xi32>
+  func.return %res : tensor<?x?x?xi32>
+}
+
+// -----
+
+func.func @dynamic_gather_c13(%operand : tensor<2x4x9xi32>, %start_indices : tensor<1x5x2xi32>, %slice_sizes : tensor<3xi32>) -> tensor<3xi32> {
+  // expected-error@+2 {{failed to infer returned types}}
   // expected-error@+1 {{inferred type(s) 'tensor<1x5x?xi32>' are incompatible with return type(s) of operation 'tensor<3xi32>'}}
   %res = "mhlo.dynamic_gather"(%operand, %start_indices, %slice_sizes) {
     dimension_numbers = #mhlo.gather<
@@ -4193,8 +4403,8 @@ func.func @dynamic_gather(%operand : tensor<2x4x9xi32>, %start_indices : tensor<
 
 // -----
 
-func.func @dynamic_gather(%operand : tensor<2x4x9xi32>, %start_indices : tensor<1x5x2xi32>, %slice_sizes : tensor<?xi32>) -> tensor<3xi32> {
-  // @expected-error@+2 {{'mhlo.dynamic_gather' op failed to infer returned types}}
+func.func @dynamic_gather_c13(%operand : tensor<2x4x9xi32>, %start_indices : tensor<1x5x2xi32>, %slice_sizes : tensor<3xi32>) -> tensor<3xi32> {
+  // expected-error@+2 {{failed to infer returned types}}
   // expected-error@+1 {{inferred type(s) 'tensor<1x5x?xi32>' are incompatible with return type(s) of operation 'tensor<3xi32>'}}
   %res = "mhlo.dynamic_gather"(%operand, %start_indices, %slice_sizes) {
     dimension_numbers = #mhlo.gather<
@@ -4204,14 +4414,14 @@ func.func @dynamic_gather(%operand : tensor<2x4x9xi32>, %start_indices : tensor<
       start_index_map = [0, 1]
     >,
     indices_are_sorted = false
-  } : (tensor<2x4x9xi32>, tensor<1x5x2xi32>, tensor<?xi32>) -> tensor<3xi32>
+  } : (tensor<2x4x9xi32>, tensor<1x5x2xi32>, tensor<3xi32>) -> tensor<3xi32>
   func.return %res : tensor<3xi32>
 }
 
 // -----
 
-func.func @dynamic_gather(%operand : tensor<?x?x?xi32>, %start_indices : tensor<?x?x?xi32>, %slice_sizes : tensor<3xi32>) -> tensor<3xi32> {
-  // @expected-error@+2 {{'mhlo.dynamic_gather' op failed to infer returned types}}
+func.func @dynamic_gather_c13(%operand : tensor<?x?x?xi32>, %start_indices : tensor<?x?x?xi32>, %slice_sizes : tensor<3xi32>) -> tensor<3xi32> {
+  // expected-error@+2 {{failed to infer returned types}}
   // expected-error@+1 {{inferred type(s) 'tensor<?x?x?xi32>' are incompatible with return type(s) of operation 'tensor<3xi32>'}}
   %res = "mhlo.dynamic_gather"(%operand, %start_indices, %slice_sizes) {
     dimension_numbers = #mhlo.gather<
@@ -4227,8 +4437,8 @@ func.func @dynamic_gather(%operand : tensor<?x?x?xi32>, %start_indices : tensor<
 
 // -----
 
-func.func @dynamic_gather(%operand : tensor<?x?x?xi32>, %start_indices : tensor<?x?x?xi32>, %slice_sizes : tensor<?xi32>) -> tensor<?xi32> {
-  // @expected-error@+2 {{'mhlo.dynamic_gather' op failed to infer returned types}}
+func.func @dynamic_gather_c13(%operand : tensor<?x?x?xi32>, %start_indices : tensor<?x?x?xi32>, %slice_sizes : tensor<3xi32>) -> tensor<?xi32> {
+  // expected-error@+2 {{failed to infer returned types}}
   // expected-error@+1 {{inferred type(s) 'tensor<?x?x?xi32>' are incompatible with return type(s) of operation 'tensor<?xi32>'}}
   %res = "mhlo.dynamic_gather"(%operand, %start_indices, %slice_sizes) {
     dimension_numbers = #mhlo.gather<
@@ -4238,7 +4448,7 @@ func.func @dynamic_gather(%operand : tensor<?x?x?xi32>, %start_indices : tensor<
       start_index_map = [0, 1]
     >,
     indices_are_sorted = false
-  } : (tensor<?x?x?xi32>, tensor<?x?x?xi32>, tensor<?xi32>) -> tensor<?xi32>
+  } : (tensor<?x?x?xi32>, tensor<?x?x?xi32>, tensor<3xi32>) -> tensor<?xi32>
   func.return %res : tensor<?xi32>
 }
 
@@ -5532,7 +5742,7 @@ func.func @quantized_dot_i4(%arg0: tensor<2x2x!quant.uniform<i4:f32, 2.0:1>>, %a
 // -----
 
 // CHECK-LABEL: func @quantized_dot_general
-func.func @quantized_dot_general(%arg0: tensor<2x16x32x!quant.uniform<i8:f32, 2.0:15>>, %arg1: tensor<2x32x32x!quant.uniform<i8:f32, 5.0:20>>) -> tensor<2x16x32x!quant.uniform<i8:f32, 10.0:50>> {
+func.func @quantized_dot_general(%arg0: tensor<2x16x32x!quant.uniform<i8:f32, 2.0:15>>, %arg1: tensor<2x32x32x!quant.uniform<i8:f32, 5.0:0>>) -> tensor<2x16x32x!quant.uniform<i8:f32, 10.0:50>> {
   %0 = "mhlo.dot_general"(%arg0, %arg1) {
     dot_dimension_numbers = #mhlo.dot<
       lhs_batching_dimensions = [0],
@@ -5541,7 +5751,7 @@ func.func @quantized_dot_general(%arg0: tensor<2x16x32x!quant.uniform<i8:f32, 2.
       rhs_contracting_dimensions = [1]
     >,
     precision_config = [#mhlo<precision DEFAULT>, #mhlo<precision DEFAULT>]}
-    : (tensor<2x16x32x!quant.uniform<i8:f32, 2.0:15>>, tensor<2x32x32x!quant.uniform<i8:f32, 5.0:20>>) -> tensor<2x16x32x!quant.uniform<i8:f32, 10.0:50>>
+    : (tensor<2x16x32x!quant.uniform<i8:f32, 2.0:15>>, tensor<2x32x32x!quant.uniform<i8:f32, 5.0:0>>) -> tensor<2x16x32x!quant.uniform<i8:f32, 10.0:50>>
   func.return %0 : tensor<2x16x32x!quant.uniform<i8:f32, 10.0:50>>
 }
 
@@ -5779,19 +5989,6 @@ func.func @pad_dynamic(%arg0: tensor<?x48x48x32xf32>) -> tensor<?x48x48x48xf32>
 
 // -----
 
-func.func @pad_i2(%arg0: tensor<1x2x3xf16>, %arg1: tensor<2xf16>) -> tensor<2x4x7xf16> {
-  // @expected-error@+2 {{'mhlo.pad' op failed to infer returned types}}
-  // expected-error@+1 {{padding value type should be a rank-0 tensor, is rank 1}}
-  %0 = "mhlo.pad"(%arg0, %arg1) {
-    edge_padding_low = dense<[0, 1, 2]> : tensor<3xi64>,
-    edge_padding_high = dense<[1, 1, 0]> : tensor<3xi64>,
-    interior_padding = dense<[0, 0, 1]> : tensor<3xi64>
-  } : (tensor<1x2x3xf16>, tensor<2xf16>) -> tensor<2x4x7xf16>
-  func.return %0 : tensor<2x4x7xf16>
-}
-
-// -----
-
 func.func @pad_i3(%arg0: tensor<1x2x3xf16>, %arg1: tensor<f16>) -> tensor<2x4x7xf16> {
   // @expected-error@+2 {{'mhlo.pad' op failed to infer returned types}}
   // expected-error@+1 {{edge_padding_low has rank 0 instead of required rank 1}}
diff --git a/third_party/xla/xla/mlir_hlo/tests/Dialect/mhlo/stablehlo-legalize-to-hlo.mlir b/third_party/xla/xla/mlir_hlo/tests/Dialect/mhlo/stablehlo-legalize-to-hlo.mlir
index 963e2a46459e44..7107b7d615e7b8 100644
--- a/third_party/xla/xla/mlir_hlo/tests/Dialect/mhlo/stablehlo-legalize-to-hlo.mlir
+++ b/third_party/xla/xla/mlir_hlo/tests/Dialect/mhlo/stablehlo-legalize-to-hlo.mlir
@@ -806,21 +806,19 @@ func.func @op_dynamic_broadcast_in_dim(%arg0: tensor<?xf32>, %arg1: tensor<2xind
 }
 
 // CHECK-LABEL: "op_dynamic_conv"
-func.func @op_dynamic_conv(%arg0: tensor<1x8x8x207xf32>, %arg1: tensor<3x3x207x16xf32>, %arg2: tensor<4xi32>) -> tensor<1x?x?x16xf32> {
+func.func @op_dynamic_conv(%arg0: tensor<1x8x8x207xf32>, %arg1: tensor<3x3x207x16xf32>, %arg2: tensor<2x2xi32>) -> tensor<1x?x?x16xf32> {
   //      CHECK: "mhlo.dynamic_conv"([[ARG0:%arg[0-9]+]], [[ARG1:%arg[0-9]+]], [[ARG2:%arg[0-9]+]]) <{
   // CHECK-SAME:   batch_group_count = 1 : i64,
   // CHECK-SAME:   dimension_numbers = #mhlo.conv<[b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f]>,
   // CHECK-SAME:   feature_group_count = 1 : i64,
   // CHECK-SAME:   lhs_dilation = dense<1> : tensor<2xi64>,
-  // CHECK-SAME:   padding = dense<1> : tensor<2x2xi64>,
   // CHECK-SAME:   precision_config = [#mhlo<precision DEFAULT>, #mhlo<precision DEFAULT>],
   // CHECK-SAME:   rhs_dilation = dense<1> : tensor<2xi64>,
   // CHECK-SAME:   window_reversal = dense<false> : tensor<2xi1>,
   // CHECK-SAME:   window_strides = dense<1> : tensor<2xi64>
-  // CHECK-SAME: }> : (tensor<1x8x8x207xf32>, tensor<3x3x207x16xf32>, tensor<4xi32>) -> tensor<1x?x?x16xf32>
+  // CHECK-SAME: }> : (tensor<1x8x8x207xf32>, tensor<3x3x207x16xf32>, tensor<2x2xi32>) -> tensor<1x?x?x16xf32>
   %0 = "stablehlo.dynamic_conv"(%arg0, %arg1, %arg2) {
     window_strides = array<i64: 1, 1>,
-    padding = dense<1> : tensor<2x2xi64>,
     lhs_dilation = array<i64: 1, 1>,
     rhs_dilation = array<i64: 1, 1>,
     window_reversal = array<i1: false, false>,
@@ -828,7 +826,7 @@ func.func @op_dynamic_conv(%arg0: tensor<1x8x8x207xf32>, %arg1: tensor<3x3x207x1
     feature_group_count = 1 : i64,
     batch_group_count = 1 : i64,
     precision_config = [#stablehlo<precision DEFAULT>, #stablehlo<precision DEFAULT>]
-  } : (tensor<1x8x8x207xf32>, tensor<3x3x207x16xf32>, tensor<4xi32>) -> tensor<1x?x?x16xf32>
+  } : (tensor<1x8x8x207xf32>, tensor<3x3x207x16xf32>, tensor<2x2xi32>) -> tensor<1x?x?x16xf32>
   func.return %0 : tensor<1x?x?x16xf32>
 }
 
@@ -1579,17 +1577,6 @@ func.func @op_torch_index_select(%arg0: tensor<5x1x5xf32>, %arg1: tensor<2xi32>)
   func.return %0 : tensor<2x1x5xf32>
 }
 
-// CHECK-LABEL: "op_trace"
-func.func @op_trace(%arg0: tensor<f32>) {
-  //      CHECK: "mhlo.trace"([[ARG0:%arg[0-9]+]]) <{
-  // CHECK-SAME:   tag = "foo"
-  // CHECK-SAME: }> : (tensor<f32>) -> ()
-  "stablehlo.trace"(%arg0) {
-    tag = "foo"
-  } : (tensor<f32>) -> ()
-  func.return
-}
-
 // CHECK-LABEL: "op_transpose"
 func.func @op_transpose(%arg0: tensor<16x8xf32>) ->  tensor<8x16xf32> {
   //      CHECK: "mhlo.transpose"([[ARG0:%arg[0-9]+]]) <{

From 86eff97b8c84bacca4b8c02b998cb6c4b2b6d9ef Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Tue, 14 May 2024 10:22:16 -0700
Subject: [PATCH 119/478] Support converting strings to unsigned integer types
 in `tf.strings.to_number`.

PiperOrigin-RevId: 633628312
---
 tensorflow/core/kernels/string_to_number_op.cc | 4 ++++
 tensorflow/core/ops/parsing_ops.cc             | 2 +-
 tensorflow/python/ops/string_ops.py            | 2 +-
 3 files changed, 6 insertions(+), 2 deletions(-)

diff --git a/tensorflow/core/kernels/string_to_number_op.cc b/tensorflow/core/kernels/string_to_number_op.cc
index d0ee2a96e4572a..7d1553874d9a7a 100644
--- a/tensorflow/core/kernels/string_to_number_op.cc
+++ b/tensorflow/core/kernels/string_to_number_op.cc
@@ -16,6 +16,8 @@ limitations under the License.
 // See docs in ../ops/parse_ops.cc.
 
 #include <errno.h>
+
+#include <cstdint>
 #include <string>
 
 #include "tensorflow/core/framework/kernel_def_builder.h"
@@ -68,6 +70,8 @@ REGISTER(float);
 REGISTER(double);
 REGISTER(int32);
 REGISTER(int64_t);
+REGISTER(uint32_t);
+REGISTER(uint64_t);
 #undef REGISTER
 
 }  // namespace tensorflow
diff --git a/tensorflow/core/ops/parsing_ops.cc b/tensorflow/core/ops/parsing_ops.cc
index d765455b6f1864..a3e801a87099c7 100644
--- a/tensorflow/core/ops/parsing_ops.cc
+++ b/tensorflow/core/ops/parsing_ops.cc
@@ -498,7 +498,7 @@ REGISTER_OP("DecodeCSV")
 REGISTER_OP("StringToNumber")
     .Input("string_tensor: string")
     .Output("output: out_type")
-    .Attr("out_type: {float, double, int32, int64} = DT_FLOAT")
+    .Attr("out_type: {float, double, int32, int64, uint32, uint64} = DT_FLOAT")
     .SetShapeFn(shape_inference::UnchangedShape);
 
 }  // namespace tensorflow
diff --git a/tensorflow/python/ops/string_ops.py b/tensorflow/python/ops/string_ops.py
index 2e3bb68e0c20f2..b6c0e46c5f1028 100644
--- a/tensorflow/python/ops/string_ops.py
+++ b/tensorflow/python/ops/string_ops.py
@@ -478,7 +478,7 @@ def string_to_number(input, out_type=dtypes.float32, name=None):
   Args:
     input: A `Tensor` of type `string`.
     out_type: An optional `tf.DType` from: `tf.float32, tf.float64, tf.int32,
-      tf.int64`. Defaults to `tf.float32`.
+      tf.int64, tf.uint32, tf.uint64`. Defaults to `tf.float32`.
       The numeric type to interpret each string in `string_tensor` as.
     name: A name for the operation (optional).
 

From 7fb5f26eaa627d5e280e3c52ad1d3464764e716d Mon Sep 17 00:00:00 2001
From: Kyle Lucke <klucke@google.com>
Date: Tue, 14 May 2024 10:25:34 -0700
Subject: [PATCH 120/478] Add a typedef to allow users to easily create the
 correct Factory type to create a TypedKernel.

PiperOrigin-RevId: 633629571
---
 third_party/xla/xla/stream_executor/kernel.h | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/third_party/xla/xla/stream_executor/kernel.h b/third_party/xla/xla/stream_executor/kernel.h
index f90ad3fc50fbf5..677e8db1a37b35 100644
--- a/third_party/xla/xla/stream_executor/kernel.h
+++ b/third_party/xla/xla/stream_executor/kernel.h
@@ -338,6 +338,9 @@ class TypedKernel {
 
   operator bool() const { return static_cast<bool>(kernel_); }  // NOLINT
 
+  // Type of factory used to create a TypedKernel.
+  using FactoryType = TypedKernelFactory<Params...>;
+
  private:
   friend class TypedKernelFactory<Params...>;
   explicit TypedKernel(std::unique_ptr<Kernel> kernel)

From f6cc0f900c36d8f4634b07202cd6d46a2e4ffaac Mon Sep 17 00:00:00 2001
From: Vlad Sytchenko <vsytch@google.com>
Date: Tue, 14 May 2024 10:33:13 -0700
Subject: [PATCH 121/478] [XLA] Account for private stack allocations in
 HloRematerialization

Buffer assignment allocates a single stack for all asynchronous computations, which persists for the entire duration of the program. We need to account for this by adjusting the memory limit.

PiperOrigin-RevId: 633632062
---
 .../xla/xla/service/gpu/gpu_compiler.cc       |  3 +-
 .../xla/xla/service/hlo_rematerialization.cc  | 74 +++++++++++++--
 .../xla/xla/service/hlo_rematerialization.h   | 25 +++--
 .../xla/service/hlo_rematerialization_test.cc | 91 ++++++++++++++++++-
 4 files changed, 169 insertions(+), 24 deletions(-)

diff --git a/third_party/xla/xla/service/gpu/gpu_compiler.cc b/third_party/xla/xla/service/gpu/gpu_compiler.cc
index 55edb1e4b00336..6d293297d2ffad 100644
--- a/third_party/xla/xla/service/gpu/gpu_compiler.cc
+++ b/third_party/xla/xla/service/gpu/gpu_compiler.cc
@@ -2186,8 +2186,7 @@ absl::Status GpuCompiler::RunPostSchedulingPipelines(
         // Assume 75% of the total device memory is available for XLA.
         /*memory_limit_bytes=*/scheduler_mem_limit,
         /*block_size_limit=*/1, /*block_rematerialization_factor=*/1,
-        /*min_remat_size=*/0, /*compact_shape_function=*/nullptr,
-        /*host_memory_offload_config=*/std::nullopt);
+        /*min_remat_size=*/0, /*compact_shape_function=*/nullptr);
     HloRematerialization::RematerializationSizes sizes;
     pipeline.AddPass<HloRematerialization>(options, sizes);
     pipeline.AddPass<StreamAttributeAnnotator>();
diff --git a/third_party/xla/xla/service/hlo_rematerialization.cc b/third_party/xla/xla/service/hlo_rematerialization.cc
index 0dbacab5e7bbf9..a948cde80aa7d2 100644
--- a/third_party/xla/xla/service/hlo_rematerialization.cc
+++ b/third_party/xla/xla/service/hlo_rematerialization.cc
@@ -22,8 +22,8 @@ limitations under the License.
 #include <memory>
 #include <optional>
 #include <set>
-#include <sstream>
 #include <string>
+#include <string_view>
 #include <tuple>
 #include <utility>
 #include <vector>
@@ -56,7 +56,6 @@ limitations under the License.
 #include "xla/shape_util.h"
 #include "xla/status.h"
 #include "xla/status_macros.h"
-#include "xla/statusor.h"
 #include "xla/util.h"
 #include "tsl/platform/errors.h"
 
@@ -2761,7 +2760,9 @@ absl::StatusOr<bool> HloRematerialization::RematerializeComputation(
       // in the callee computations.
       for (HloComputation* called_computation :
            callsite->called_computations()) {
-        if (!ContainsKey(rematerialized_computations_, called_computation)) {
+        if (!ContainsKey(rematerialized_computations_, called_computation) &&
+            HloInstruction::IsThreadIncluded(
+                called_computation->execution_thread(), execution_threads)) {
           // Memory limit for the subcomputation is the memory limit less the
           // amount of memory used at this point in the computation.
           int64_t subcomputation_memory_limit_bytes = std::max<int64_t>(
@@ -2862,15 +2863,67 @@ absl::StatusOr<bool> HloRematerialization::Run(
         module_output_size += options_.hlo_cost_analysis.GetShapeSize(subshape);
       });
 
-  const int64_t adjusted_memory_limit_bytes =
+  int64_t adjusted_memory_limit_bytes =
       std::max<int64_t>(0, options_.memory_limit_bytes - module_output_size);
   VLOG(1) << "Adjusted memory limit accounting for output ("
           << HumanReadableNumBytes(module_output_size)
           << "): " << HumanReadableNumBytes(adjusted_memory_limit_bytes);
 
+  call_graph_ = CallGraph::Build(module);
+
+  // Buffer assignment allocates a single stack for all asynchronous
+  // computations of the same thread, which persists for the entire duration of
+  // the program. We need to account for this by adjusting the memory limit.
+  int64_t total_async_peak_memory = 0;
+  if (!options_.async_computation_parallelism.empty()) {
+    // We cannot compute memory usage for both the main and asynchronous threads
+    // at the same time, as that will cause the asynchronous callee usage to be
+    // added to the main thread callers usage. The callee's memory is
+    // preallocated, so the caller doesn't pay for it.
+    absl::flat_hash_set<std::string_view> async_threads;
+    for (const auto& [computation, _] :
+         options_.async_computation_parallelism) {
+      async_threads.insert(computation->execution_thread());
+    }
+    TF_RETURN_IF_ERROR(call_graph_->VisitNodes(
+        [this, module, &async_threads](const CallGraphNode& node) -> Status {
+          auto callee_thread = node.computation()->execution_thread();
+          if (node.context() == CallContext::kControlFlow &&
+              HloInstruction::IsThreadIncluded(callee_thread, async_threads)) {
+            TF_ASSIGN_OR_RETURN(computation_peak_memory_[node.computation()],
+                                ComputePeakMemory(node.computation(),
+                                                  module->schedule().sequence(
+                                                      node.computation()),
+                                                  {callee_thread}));
+          }
+          return OkStatus();
+        },
+        /*visit_unreachable_nodes=*/false));
+
+    int64_t async_peak_memory = 0;
+    // Only consider asynchronous computations invoked from the main thread.
+    for (const auto [entry_computation, parallel_threads] :
+         options_.async_computation_parallelism) {
+      const int64_t peak_memory =
+          computation_peak_memory_.at(entry_computation);
+      // Adjust memory usage for parallel execution of the same computation
+      // on different devices.
+      const int64_t parallel_peak_memory = peak_memory * parallel_threads;
+      async_peak_memory = std::max(async_peak_memory, parallel_peak_memory);
+    }
+    adjusted_memory_limit_bytes =
+        std::max<int64_t>(0, adjusted_memory_limit_bytes - async_peak_memory);
+    total_async_peak_memory += async_peak_memory;
+    VLOG(1) << "Adjusted memory limit accounting for async computations ("
+            << HumanReadableNumBytes(async_peak_memory)
+            << "): " << HumanReadableNumBytes(adjusted_memory_limit_bytes);
+
+    // Reset back to a clean state, since we don't expect to utilize the
+    // async computation memory usage anymore.
+    computation_peak_memory_.clear();
+  }
   // Compute peak memory usage of all computations in the module called in a
   // sequential context.
-  call_graph_ = CallGraph::Build(module);
   TF_RETURN_IF_ERROR(call_graph_->VisitNodes(
       [this, module, &execution_threads](const CallGraphNode& node) -> Status {
         if (node.context() == CallContext::kControlFlow &&
@@ -2887,12 +2940,13 @@ absl::StatusOr<bool> HloRematerialization::Run(
       /*visit_unreachable_nodes=*/false));
 
   // The peak memory usage of the module equals the peak memory use of the entry
-  // computation plus the output size of the computation. This is because the
-  // peak memory for a computation does not include the output as this is
-  // typically accounted for in the caller.
+  // computation plus the output size of the computation plus memory use of
+  // asynchronous computations. This is because the peak memory for a
+  // computation does not include the output as this is typically accounted for
+  // in the caller.
   const int64_t before_peak_memory =
       computation_peak_memory_.at(module->entry_computation()) +
-      module_output_size;
+      module_output_size + total_async_peak_memory;
   VLOG(1) << "Peak memory usage of module (before): "
           << HumanReadableNumBytes(before_peak_memory);
 
@@ -2928,7 +2982,7 @@ absl::StatusOr<bool> HloRematerialization::Run(
           << net_instructions_added_ << " net instructions added";
   const int64_t current_peak_memory =
       computation_peak_memory_.at(module->entry_computation()) +
-      module_output_size;
+      module_output_size + total_async_peak_memory;
   VLOG(1) << "Peak memory usage of module now "
           << HumanReadableNumBytes(current_peak_memory) << " ("
           << current_peak_memory << " bytes), was "
diff --git a/third_party/xla/xla/service/hlo_rematerialization.h b/third_party/xla/xla/service/hlo_rematerialization.h
index 4fae1cc2ddbbd0..e5bb5cc8aa8d2e 100644
--- a/third_party/xla/xla/service/hlo_rematerialization.h
+++ b/third_party/xla/xla/service/hlo_rematerialization.h
@@ -90,13 +90,15 @@ class HloRematerialization : public HloModulePass {
   static Shape DefaultCompactShapeFunction(const Shape& shape) { return shape; }
 
   struct Options {
-    explicit Options(
-        HloCostAnalysis& hlo_cost_analysis,
-        const RematerializationModeConfig& remat_mode_config,
-        int64_t memory_limit_bytes, int block_size_limit,
-        int block_rematerialization_factor, int64_t min_remat_size,
-        CompactShapeFunction compact_shape_function,
-        std::optional<HostMemoryOffloadConfig> host_memory_offload_config)
+    explicit Options(HloCostAnalysis& hlo_cost_analysis,
+                     const RematerializationModeConfig& remat_mode_config,
+                     int64_t memory_limit_bytes, int block_size_limit,
+                     int block_rematerialization_factor, int64_t min_remat_size,
+                     CompactShapeFunction compact_shape_function,
+                     std::optional<HostMemoryOffloadConfig>
+                         host_memory_offload_config = std::nullopt,
+                     absl::flat_hash_map<HloComputation*, int64_t>
+                         async_computation_parallelism = {})
         : hlo_cost_analysis(hlo_cost_analysis),
           remat_mode_config(remat_mode_config),
           memory_limit_bytes(memory_limit_bytes),
@@ -106,7 +108,8 @@ class HloRematerialization : public HloModulePass {
           compact_shape_function(compact_shape_function == nullptr
                                      ? DefaultCompactShapeFunction
                                      : std::move(compact_shape_function)),
-          host_memory_offload_config(host_memory_offload_config) {}
+          host_memory_offload_config(host_memory_offload_config),
+          async_computation_parallelism(async_computation_parallelism) {}
 
     // The cost model used for decisions during rematerialization for host
     // memory offload. It is also used for getting Shape size.
@@ -133,7 +136,7 @@ class HloRematerialization : public HloModulePass {
     // return for potentially reduced memory consumption.
     int block_rematerialization_factor;
 
-    // The minimim size, in bytes, of a tensor to be considered for
+    // The minimum size, in bytes, of a tensor to be considered for
     // rematerialization. All tensors smaller than this size will be skipped
     // over.
     int64_t min_remat_size;
@@ -143,6 +146,10 @@ class HloRematerialization : public HloModulePass {
     CompactShapeFunction compact_shape_function;
 
     std::optional<HostMemoryOffloadConfig> host_memory_offload_config;
+
+    // Collection of async entry computations and their number of parallel
+    // invocations.
+    absl::flat_hash_map<HloComputation*, int64_t> async_computation_parallelism;
   };
 
   explicit HloRematerialization(Options options, RematerializationSizes& sizes)
diff --git a/third_party/xla/xla/service/hlo_rematerialization_test.cc b/third_party/xla/xla/service/hlo_rematerialization_test.cc
index 780c881aa94bf1..b30cf8293e48e9 100644
--- a/third_party/xla/xla/service/hlo_rematerialization_test.cc
+++ b/third_party/xla/xla/service/hlo_rematerialization_test.cc
@@ -47,6 +47,88 @@ namespace op = xla::testing::opcode_matchers;
 
 using ::testing::_;
 
+class AsyncRematerializationTest : public RematerializationTestBase {
+ protected:
+  absl::StatusOr<bool> RunHloRematerialization(
+      int64_t memory_limit_bytes, HloModule* module,
+      const absl::flat_hash_map<HloComputation*, int64_t>&
+          async_computation_parallelism,
+      int64_t min_remat_size = 0) {
+    TF_EXPECT_OK(verifier().Run(module).status());
+    if (!module->has_schedule()) {
+      HloMemoryScheduler scheduler(
+          [](const BufferValue& buffer) { return ByteSizeOf(buffer.shape()); },
+          ComputationSchedulerToModuleScheduler(DefaultMemoryScheduler));
+      TF_EXPECT_OK(scheduler.Run(module).status());
+    }
+    HloRematerialization::RematerializationModeConfig config(
+        /*recompute=*/true, /*compress=*/true, /*host_offload=*/false);
+    auto shape_size_func = [](const Shape& shape) { return ByteSizeOf(shape); };
+    HloCostAnalysis cost_analysis(shape_size_func);
+    HloRematerialization::Options options(
+        cost_analysis, config, memory_limit_bytes,
+        /*block_size_limit=*/1, /*block_rematerialization_factor=*/1,
+        min_remat_size, /*compact_shape_function=*/nullptr,
+        /*host_memory_offload_config=*/std::nullopt,
+        /*async_computation_parallelism=*/async_computation_parallelism);
+    HloRematerialization::RematerializationSizes sizes;
+    HloRematerialization remat(options, sizes);
+    return remat.Run(module, {HloInstruction::kMainExecutionThread});
+  }
+
+  static constexpr int64_t kNumParallelThreads = 16;
+};
+
+TEST_F(AsyncRematerializationTest, AsyncComputation) {
+  constexpr std::string_view hlo = R"(
+HloModule async, is_scheduled=true
+
+%offload_computation {
+  %param = f32[1]{0} parameter(0)
+  %reshape = f32[] reshape(f32[1]{0} %param)
+  %broadcast = f32[1024]{0} broadcast(f32[] %reshape), dimensions={}
+  %negate = f32[1024]{0} negate(f32[1024]{0} %broadcast)
+  %concatenate = f32[2048]{0} concatenate(f32[1024]{0} %negate, f32[1024]{0} %negate), dimensions={0}
+  %slice = f32[1]{0} slice(f32[2048]{0} %concatenate), slice={[0:1]}
+  %concatenate.1 = f32[1025]{0} concatenate(f32[1024]{0} %broadcast, f32[1]{0} %slice), dimensions={0}
+  ROOT %slice.1 = f32[1]{0} slice(f32[1025]{0} %concatenate.1), slice={[0:1]}
+}
+
+%main_computation {
+  %param = f32[1]{0} parameter(0)
+  %reshape = f32[] reshape(f32[1]{0} %param)
+  %broadcast = f32[1024]{0} broadcast(f32[] %reshape), dimensions={}
+  %negate = f32[1024]{0} negate(f32[1024]{0} %broadcast)
+  %concatenate = f32[2048]{0} concatenate(f32[1024]{0} %negate, f32[1024]{0} %negate), dimensions={0}
+  %slice = f32[1]{0} slice(f32[2048]{0} %concatenate), slice={[0:1]}
+  %concatenate.1 = f32[1025]{0} concatenate(f32[1024]{0} %broadcast, f32[1]{0} %slice), dimensions={0}
+  ROOT %slice.1 = f32[1]{0} slice(f32[1025]{0} %concatenate.1), slice={[0:1]}
+}
+
+ENTRY %main {
+  %param = f32[1]{0} parameter(0)
+  %call-start = ((f32[1]{0}), f32[1]{0}, s32[]) call-start(f32[1]{0} %param), to_apply=%offload_computation, async_execution_thread="offload"
+  %call-done = f32[1]{0} call-done(((f32[1]{0}), f32[1]{0}, s32[]) %call-start)
+  ROOT %call = f32[1]{0} call(f32[1]{0} %call-done), to_apply=%main_computation
+}
+)";
+
+  TF_ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo));
+
+  HloInstruction* call_start = FindInstruction(module.get(), "call-start");
+  // Computation requires 16KB without rematerialization, but uses only 12KB
+  // with rematerialization so pick a memory limit between these values (14KB).
+  // Asynchronous computation will run on 16 devices and we do not rematerialize
+  // it, so it will reserve 16 * 16Kb from the memory limit.
+  TF_ASSERT_OK_AND_ASSIGN(
+      bool changed,
+      RunHloRematerialization(
+          /*memory_limit_bytes=*/kNumParallelThreads * 16 * 1024 + 14 * 1024,
+          module.get(),
+          {{call_start->async_wrapped_computation(), kNumParallelThreads}}));
+  EXPECT_TRUE(changed);
+}
+
 // Inherits methods to create rematerializable computations. See
 // RematerializationTestBase for more.
 class RecomputeAndCompressHloRematerializationTest
@@ -80,7 +162,8 @@ class RecomputeAndCompressHloRematerializationTest
         cost_analysis, config, memory_limit_bytes,
         /*block_size_limit=*/1, /*block_rematerialization_factor=*/1,
         min_remat_size, /*compact_shape_function=*/nullptr,
-        /*host_memory_offload_config=*/std::nullopt);
+        /*host_memory_offload_config=*/std::nullopt,
+        /*async_threads=*/{});
     HloRematerialization::RematerializationSizes sizes;
     HloRematerialization remat(options, sizes);
     absl::StatusOr<bool> result = remat.Run(module);
@@ -1105,7 +1188,8 @@ class CompressingRematerializationTest : public RematerializationTestBase {
         cost_analysis, config, memory_limit_bytes,
         /*block_size_limit=*/1, /*block_rematerialization_factor=*/1,
         min_remat_size, ChooseCompactLayoutForShape,
-        /*host_memory_offload_config=*/std::nullopt);
+        /*host_memory_offload_config=*/std::nullopt,
+        /*async_threads=*/{});
     HloRematerialization::RematerializationSizes sizes;
     HloRematerialization remat(options, sizes);
     return remat.Run(module);
@@ -1314,7 +1398,8 @@ class OffloadingRematerializationTest : public RematerializationTestBase {
         cost_analysis, config, memory_limit_bytes,
         /*block_size_limit=*/1, /*block_rematerialization_factor=*/1,
         min_remat_size, /*compact_shape_function=*/nullptr,
-        host_memory_offload_config);
+        host_memory_offload_config,
+        /*async_threads=*/{});
     HloRematerialization::RematerializationSizes sizes;
     HloRematerialization remat(options, sizes);
     return remat.Run(module);

From b10ad05b3b2f78d1c0c96ec7194f58bbb0be54e9 Mon Sep 17 00:00:00 2001
From: Goran Flegar <gflegar@google.com>
Date: Tue, 14 May 2024 10:52:12 -0700
Subject: [PATCH 122/478] Add a flag to override the GEMM autotuner with a
 specified textproto configuration.

This makes it easier to debug issues with Triton, since [persisted autotuning](https://openxla.org/xla/persisted_autotuning) does not offer a workflow that would be efficient enough:

- `--xla_gpu_dump_autotune_results_to=` doesn't work if the autotuner crashes, so we would need to construct the cache manually from the logs; furthermore, the cache only applies to a specific HW, so we can't use it to override the autotuner on a different GPU
- `--xla_gpu_load_autotune_results_from=` reads a file so we need to copy the file into the source tree, and modify the BUILD file to include it into the build; with this new flag, we can just copy the flag from the VLOG without having to fiddle around with any files

PiperOrigin-RevId: 633638219
---
 third_party/xla/xla/debug_options_flags.cc    |  8 ++++
 third_party/xla/xla/service/gpu/BUILD         |  1 +
 .../xla/service/gpu/gemm_fusion_autotuner.cc  | 43 +++++++++++++++++--
 third_party/xla/xla/xla.proto                 |  4 +-
 4 files changed, 51 insertions(+), 5 deletions(-)

diff --git a/third_party/xla/xla/debug_options_flags.cc b/third_party/xla/xla/debug_options_flags.cc
index 831a26b2c96793..aa95aa8a3f6120 100644
--- a/third_party/xla/xla/debug_options_flags.cc
+++ b/third_party/xla/xla/debug_options_flags.cc
@@ -1472,6 +1472,14 @@ void MakeDebugOptionsFlags(std::vector<tsl::Flag>* flag_list,
       "Dumps autotuned GEMM fusions to the directory specified by "
       "xla_dump_to or stdout. Each fusion is dumped only once, as an optimized "
       "HLO."));
+  flag_list->push_back(tsl::Flag(
+      "xla_gpu_override_gemm_autotuner",
+      string_setter_for(&DebugOptions::set_xla_gpu_override_gemm_autotuner),
+      debug_options->xla_gpu_override_gemm_autotuner(),
+      "Overrides the GEMM autotuner to use the specified "
+      "(AutotuneResult::TritonGemmKey) textproto configuration for all Triton "
+      "GEMM fusions. (You can get such textprotos from the debug logs of the "
+      "GEMM autotuner.) "));
   flag_list->push_back(tsl::Flag(
       "xla_gpu_copy_insertion_use_region_analysis",
       bool_setter_for(
diff --git a/third_party/xla/xla/service/gpu/BUILD b/third_party/xla/xla/service/gpu/BUILD
index 91776ad61dee01..8b841bcf9be6ea 100644
--- a/third_party/xla/xla/service/gpu/BUILD
+++ b/third_party/xla/xla/service/gpu/BUILD
@@ -758,6 +758,7 @@ cc_library(
         "@local_tsl//tsl/platform:blocking_counter",
         "@local_tsl//tsl/platform:env",
         "@local_tsl//tsl/platform:errors",
+        "@local_tsl//tsl/platform:protobuf",
         "@local_tsl//tsl/platform:status",
         "@local_tsl//tsl/platform:statusor",
         "@local_tsl//tsl/profiler/lib:scoped_annotation",
diff --git a/third_party/xla/xla/service/gpu/gemm_fusion_autotuner.cc b/third_party/xla/xla/service/gpu/gemm_fusion_autotuner.cc
index d5b34fb877add3..a2a1cc1c063493 100644
--- a/third_party/xla/xla/service/gpu/gemm_fusion_autotuner.cc
+++ b/third_party/xla/xla/service/gpu/gemm_fusion_autotuner.cc
@@ -90,6 +90,7 @@ limitations under the License.
 #include "tsl/platform/blocking_counter.h"
 #include "tsl/platform/errors.h"
 #include "tsl/platform/path.h"
+#include "tsl/platform/protobuf.h"
 #include "tsl/platform/status.h"
 #include "tsl/platform/statusor.h"
 #include "tsl/platform/threadpool.h"
@@ -501,6 +502,17 @@ absl::Status DumpAutotunedFusion(const AutotuneConfig& autotune_config,
   return absl::OkStatus();
 }
 
+std::string Serialize(const Config& config) {
+  if (auto triton_config = std::get_if<TritonGemmConfig>(&config)) {
+    tsl::protobuf::TextFormat::Printer printer;
+    printer.SetSingleLineMode(true);
+    std::string result;
+    printer.PrintToString(triton_config->ToProto(), &result);
+    return result;
+  }
+  return GemmFusionAutotunerImpl::ToString(config);
+}
+
 }  // anonymous namespace
 
 // Methods required for sorting the configs.
@@ -757,12 +769,18 @@ GemmFusionAutotunerImpl::CompileAll(
       const HloFusionInstruction* fusion = key_value.first;
       const std::vector<Config>& gemm_config_set = key_value.second;
 
-      VLOG(10) << "Compiling the fusion: " << fusion->name();
+      VLOG(10) << "Compiling fusion: " << fusion->name();
       VLOG(10) << "Dumping fusion computation: "
                << fusion->called_computation()->ToString();
       for (const Config& config : gemm_config_set) {
         thread_pool_->Schedule([&, fusion] {
-          VLOG(10) << "Trying configuration: " << ToString(config);
+          VLOG(10) << "Trying configuration forceable through: "
+                      "--xla_gpu_override_gemm_autotuner='"
+                   << Serialize(config) << "'";
+          VLOG(10) << "WARNING: you are running in multithreaded-mode, the "
+                      "last configuration printed out might not be the one "
+                      "causing issues! Use "
+                      "--xla_gpu_force_compilation_parallelism=1 to fix.";
           absl::StatusOr<bool> has_executable =
               compile(fusion, config, gemm_config_set.size() > 1);
           TF_CHECK_OK(has_executable.status())
@@ -790,11 +808,13 @@ GemmFusionAutotunerImpl::CompileAll(
       const HloFusionInstruction* fusion = key_value.first;
       const auto& gemm_config_set = key_value.second;
 
-      VLOG(10) << "Compiling the fusion: " << fusion->name();
+      VLOG(10) << "Compiling fusion: " << fusion->name();
       VLOG(10) << "Dumping fusion computation: "
                << fusion->called_computation()->ToString();
       for (const Config& config : gemm_config_set) {
-        VLOG(10) << "Trying configuration: " << ToString(config);
+        VLOG(10) << "Trying configuration forceable through: "
+                    "--xla_gpu_override_gemm_autotuner='"
+                 << Serialize(config) << "'";
         TF_ASSIGN_OR_RETURN(
             bool has_executable,
             compile(fusion, config, gemm_config_set.size() > 1));
@@ -1123,6 +1143,21 @@ absl::StatusOr<bool> GemmFusionAutotuner::Run(
           tsl::proto_utils::ToDurationProto(absl::ZeroDuration());
       AutotunerUtil::AddResult(key, res);
     }
+  } else if (!debug_options.xla_gpu_override_gemm_autotuner().empty()) {
+    // TODO(gflegar): support overriding with non-Triton configs (cuBLAS, cuDNN)
+    AutotuneResult::TritonGemmKey gemm_key;
+    CHECK(tsl::protobuf::TextFormat::ParseFromString(
+        debug_options.xla_gpu_override_gemm_autotuner(), &gemm_key));
+    VLOG(1) << "Overriding GEMM autotuner with the following config: "
+            << gemm_key.DebugString();
+    for (const auto& [fusion, unused] : gemm_config_sets) {
+      const AutotuneCacheKey key = AutotunerUtil::GetKey(fusion, config_);
+      AutotuneResult res;
+      *res.mutable_triton() = gemm_key;
+      *res.mutable_run_time() =
+          tsl::proto_utils::ToDurationProto(absl::ZeroDuration());
+      AutotunerUtil::AddResult(key, res);
+    }
   } else if (!config_.IsDeviceless()) {
     TF_ASSIGN_OR_RETURN(std::optional<AutotunerCompileUtil> opt_compile_util,
                         AutotunerCompileUtil::Create(config_, debug_options));
diff --git a/third_party/xla/xla/xla.proto b/third_party/xla/xla/xla.proto
index b5044afe4d6292..56ddedc2d6416b 100644
--- a/third_party/xla/xla/xla.proto
+++ b/third_party/xla/xla/xla.proto
@@ -633,6 +633,8 @@ message DebugOptions {
 
   bool xla_gpu_dump_autotuned_gemm_fusions = 232;
 
+  string xla_gpu_override_gemm_autotuner = 295;
+
   bool xla_gpu_copy_insertion_use_region_analysis = 236;
 
   // If true, each fusion instruction will have a cost model runtime estimate in
@@ -767,7 +769,7 @@ message DebugOptions {
   // Base length to rewrite the reduce window to, no rewrite if set to 0.
   int64 xla_reduce_window_rewrite_base_length = 293;
 
-  // Next id: 295
+  // Next id: 296
 
   // Extra options to pass to the compilation backend (e.g. LLVM); specific
   // interpretation of these values is left to the backend.

From d8d9f2e9ca26e531169f9a55012d66f90a81c01a Mon Sep 17 00:00:00 2001
From: Arturo Schmidt <arturoschmidt@google.com>
Date: Tue, 14 May 2024 11:15:31 -0700
Subject: [PATCH 123/478] Remove unused translate_tf_dialect.

PiperOrigin-RevId: 633647120
---
 tensorflow/compiler/mlir/BUILD                |  1 -
 tensorflow/compiler/mlir/tensorflow/BUILD     |  1 -
 .../mlir2graphdef/simple_tf_dialect_op.mlir   | 32 --------
 .../compiler/mlir/tensorflow/translate/BUILD  | 16 ----
 .../translate/translate_tf_dialect_op.cc      | 75 -------------------
 5 files changed, 125 deletions(-)
 delete mode 100644 tensorflow/compiler/mlir/tensorflow/tests/mlir2graphdef/simple_tf_dialect_op.mlir
 delete mode 100644 tensorflow/compiler/mlir/tensorflow/translate/translate_tf_dialect_op.cc

diff --git a/tensorflow/compiler/mlir/BUILD b/tensorflow/compiler/mlir/BUILD
index 810e4277e8c58f..46d5e7e9fb9005 100644
--- a/tensorflow/compiler/mlir/BUILD
+++ b/tensorflow/compiler/mlir/BUILD
@@ -226,7 +226,6 @@ tf_cc_binary(
         "//tensorflow/compiler/mlir/tensorflow:translate_cl_options",
         "//tensorflow/compiler/mlir/tensorflow:translate_lib",
         "//tensorflow/compiler/mlir/tensorflow:translate_registration",
-        "//tensorflow/compiler/mlir/tensorflow:translate_tf_dialect_op",
         "//tensorflow/core:lib",
         "//tensorflow/core:tensorflow",
         "@com_google_absl//absl/strings",
diff --git a/tensorflow/compiler/mlir/tensorflow/BUILD b/tensorflow/compiler/mlir/tensorflow/BUILD
index 0be47767d4cabe..85f8644fc2d88f 100644
--- a/tensorflow/compiler/mlir/tensorflow/BUILD
+++ b/tensorflow/compiler/mlir/tensorflow/BUILD
@@ -1664,7 +1664,6 @@ aliased_targets = [
     "export_graphdef",
     "import_model",
     "export_tf_dialect_op",
-    "translate_tf_dialect_op",
     "mlir_roundtrip_flags",
     "mlir_import_options",
     "translate_lib",
diff --git a/tensorflow/compiler/mlir/tensorflow/tests/mlir2graphdef/simple_tf_dialect_op.mlir b/tensorflow/compiler/mlir/tensorflow/tests/mlir2graphdef/simple_tf_dialect_op.mlir
deleted file mode 100644
index 780406e0c16127..00000000000000
--- a/tensorflow/compiler/mlir/tensorflow/tests/mlir2graphdef/simple_tf_dialect_op.mlir
+++ /dev/null
@@ -1,32 +0,0 @@
-// RUN: tf-mlir-translate -test-only-mlir-to-tf-nodedef %s -o - | FileCheck %s
-
-func.func @main() {
-^bb0:
-  // CHECK: name: "node_name"
-  // CHECK-NEXT: op: "Const"
-  // CHECK-NEXT: attr {
-  // CHECK:      key: "dtype"
-  // CHECK-NEXT:   value {
-  // CHECK-NEXT:     type: DT_INT32
-  // CHECK-NEXT:   }
-  // CHECK-NEXT: }
-  // CHECK-NEXT: attr {
-  // CHECK-NEXT:   key: "value"
-  // CHECK-NEXT:   value {
-  // CHECK-NEXT:     tensor {
-  // CHECK-NEXT:       dtype: DT_INT32
-  // CHECK-NEXT:       tensor_shape {
-  // CHECK-NEXT:         dim {
-  // CHECK-NEXT:           size: 2
-  // CHECK-NEXT:         }
-  // CHECK-NEXT:       }
-  // CHECK-NEXT:       tensor_content: "\200\000\000\000\200\000\000\000"
-  // CHECK:      experimental_debug_info {
-  // CHECK-NEXT:   original_node_names: "n1"
-  // CHECK-NEXT:   original_func_names: "f1"
-  // CHECK-NEXT: }
-  %0 = "tf.Const"() {value = #tf_type<tensor_proto : "0x746674656E736F722464747970653A2044545F494E5433320A74656E736F725F7368617065207B0A202064696D207B0A2020202073697A653A20320A20207D0A7D0A74656E736F725F636F6E74656E743A20225C3230305C3030305C3030305C3030305C3230305C3030305C3030305C303030220A"> : tensor<2xi32>} : () -> (tensor<2xi32>) loc(fused[callsite("n1@f1" at callsite("node_name" at "file_loc"))])
-  func.return
-}
-
-
diff --git a/tensorflow/compiler/mlir/tensorflow/translate/BUILD b/tensorflow/compiler/mlir/tensorflow/translate/BUILD
index 620e4d7783609e..f0280340dddf62 100644
--- a/tensorflow/compiler/mlir/tensorflow/translate/BUILD
+++ b/tensorflow/compiler/mlir/tensorflow/translate/BUILD
@@ -112,22 +112,6 @@ cc_library(
     ],
 )
 
-cc_library(
-    name = "translate_tf_dialect_op",
-    srcs = ["translate_tf_dialect_op.cc"],
-    deps = [
-        ":export_tf_dialect_op",
-        "//tensorflow/compiler/mlir/tensorflow",
-        "@llvm-project//llvm:Support",
-        "@llvm-project//mlir:FuncDialect",
-        "@llvm-project//mlir:IR",
-        "@llvm-project//mlir:Support",
-        "@llvm-project//mlir:TranslateLib",
-        "@local_tsl//tsl/platform:protobuf",
-    ],
-    alwayslink = 1,
-)
-
 cc_library(
     name = "mlir_roundtrip_flags",
     srcs = ["mlir_roundtrip_flags.cc"],
diff --git a/tensorflow/compiler/mlir/tensorflow/translate/translate_tf_dialect_op.cc b/tensorflow/compiler/mlir/tensorflow/translate/translate_tf_dialect_op.cc
deleted file mode 100644
index 856db032e501ae..00000000000000
--- a/tensorflow/compiler/mlir/tensorflow/translate/translate_tf_dialect_op.cc
+++ /dev/null
@@ -1,75 +0,0 @@
-/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
-
-Licensed under the Apache License, Version 2.0 (the "License");
-you may not use this file except in compliance with the License.
-You may obtain a copy of the License at
-
-    http://www.apache.org/licenses/LICENSE-2.0
-
-Unless required by applicable law or agreed to in writing, software
-distributed under the License is distributed on an "AS IS" BASIS,
-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.
-==============================================================================*/
-
-#include "llvm/ADT/STLExtras.h"
-#include "llvm/Support/ToolOutputFile.h"
-#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
-#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
-#include "mlir/IR/Location.h"  // from @llvm-project
-#include "mlir/IR/MLIRContext.h"  // from @llvm-project
-#include "mlir/Tools/mlir-translate/Translation.h"  // from @llvm-project
-#include "tensorflow/compiler/mlir/tensorflow/dialect_registration.h"
-#include "tensorflow/compiler/mlir/tensorflow/translate/export_tf_dialect_op.h"
-#include "tsl/platform/protobuf.h"
-
-namespace mlir {
-static mlir::Operation* ExtractOnlyOp(mlir::ModuleOp module) {
-  mlir::func::FuncOp fn = module.lookupSymbol<mlir::func::FuncOp>("main");
-  if (!fn) return nullptr;
-
-  if (!llvm::hasSingleElement(fn)) return nullptr;
-
-  // Here, modules with exactly two operations in the only basic block are
-  // supported. The last operation should be a terminator operation and the
-  // other operation is the operation of interest.
-  auto& block = fn.front();
-  if (block.getOperations().size() != 2) return nullptr;
-  if (!block.back().hasTrait<OpTrait::IsTerminator>()) return nullptr;
-
-  return &block.front();
-}
-
-static LogicalResult MlirToTfNodeDef(ModuleOp module,
-                                     llvm::raw_ostream& output) {
-  auto* context = module.getContext();
-
-  Operation* op = ExtractOnlyOp(module);
-  if (!op) {
-    emitError(UnknownLoc::get(context),
-              "modules with exactly one op other than terminator in a "
-              "'main' function's "
-              "only block are supported");
-    return failure();
-  }
-
-  auto node_def_or = tensorflow::ConvertTFDialectOpToNodeDef(
-      op, "node_name", /*ignore_unregistered_attrs=*/false);
-  if (!node_def_or.ok()) {
-    op->emitError("failed to convert to TF NodeDef:")
-        << node_def_or.status().ToString();
-    return failure();
-  }
-
-  output << tsl::LegacyUnredactedDebugString(*node_def_or.value());
-  return success();
-}
-
-// Test only translation to convert a simple MLIR module with a single TF
-// dialect op to NodeDef.
-static TranslateFromMLIRRegistration translate_from_mlir_registration(
-    "test-only-mlir-to-tf-nodedef", "test-only-mlir-to-tf-nodedef",
-    MlirToTfNodeDef, mlir::RegisterAllTensorFlowDialects);
-
-}  // namespace mlir

From fe35120d7bb513623e46dbb2db41aa49bf00e36a Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Tue, 14 May 2024 11:22:48 -0700
Subject: [PATCH 124/478] Update ops-related pbtxt files.

PiperOrigin-RevId: 633649526
---
 .../ops/compat/ops_history_v2/Abort.pbtxt     |  2 +-
 .../core/ops/compat/ops_history_v2/Abs.pbtxt  |  2 +-
 .../compat/ops_history_v2/AccumulateNV2.pbtxt |  2 +-
 .../AccumulatorApplyGradient.pbtxt            |  2 +-
 .../AccumulatorNumAccumulated.pbtxt           |  2 +-
 .../AccumulatorSetGlobalStep.pbtxt            |  2 +-
 .../AccumulatorTakeGradient.pbtxt             |  2 +-
 .../core/ops/compat/ops_history_v2/Acos.pbtxt |  2 +-
 .../ops/compat/ops_history_v2/Acosh.pbtxt     |  2 +-
 .../core/ops/compat/ops_history_v2/Add.pbtxt  |  2 +-
 .../AddManySparseToTensorsMap.pbtxt           |  2 +-
 .../core/ops/compat/ops_history_v2/AddN.pbtxt |  2 +-
 .../AddSparseToTensorsMap.pbtxt               |  2 +-
 .../ops/compat/ops_history_v2/AddV2.pbtxt     |  2 +-
 .../ops_history_v2/AdjustContrast.pbtxt       |  2 +-
 .../ops_history_v2/AdjustContrastv2.pbtxt     |  2 +-
 .../ops/compat/ops_history_v2/AdjustHue.pbtxt |  2 +-
 .../ops_history_v2/AdjustSaturation.pbtxt     |  2 +-
 .../core/ops/compat/ops_history_v2/All.pbtxt  |  2 +-
 .../ops_history_v2/AllCandidateSampler.pbtxt  |  2 +-
 .../ops/compat/ops_history_v2/AllToAll.pbtxt  |  2 +-
 .../ops/compat/ops_history_v2/Angle.pbtxt     |  2 +-
 .../ops_history_v2/AnonymousHashTable.pbtxt   |  2 +-
 .../ops_history_v2/AnonymousIterator.pbtxt    |  2 +-
 .../ops_history_v2/AnonymousIteratorV2.pbtxt  |  2 +-
 .../ops_history_v2/AnonymousIteratorV3.pbtxt  |  2 +-
 .../ops_history_v2/AnonymousMemoryCache.pbtxt |  2 +-
 .../AnonymousMultiDeviceIterator.pbtxt        |  2 +-
 .../AnonymousMultiDeviceIteratorV3.pbtxt      |  2 +-
 .../AnonymousMutableDenseHashTable.pbtxt      |  2 +-
 .../AnonymousMutableHashTable.pbtxt           |  2 +-
 .../AnonymousMutableHashTableOfTensors.pbtxt  |  2 +-
 .../AnonymousRandomSeedGenerator.pbtxt        |  2 +-
 .../AnonymousSeedGenerator.pbtxt              |  2 +-
 .../core/ops/compat/ops_history_v2/Any.pbtxt  |  2 +-
 .../compat/ops_history_v2/ApplyAdaMax.pbtxt   |  2 +-
 .../compat/ops_history_v2/ApplyAdadelta.pbtxt |  2 +-
 .../compat/ops_history_v2/ApplyAdagrad.pbtxt  |  2 +-
 .../ops_history_v2/ApplyAdagradDA.pbtxt       |  2 +-
 .../ops_history_v2/ApplyAdagradV2.pbtxt       |  2 +-
 .../ops/compat/ops_history_v2/ApplyAdam.pbtxt |  2 +-
 .../compat/ops_history_v2/ApplyAddSign.pbtxt  |  2 +-
 .../ops_history_v2/ApplyCenteredRMSProp.pbtxt |  2 +-
 .../ops/compat/ops_history_v2/ApplyFtrl.pbtxt |  2 +-
 .../compat/ops_history_v2/ApplyFtrlV2.pbtxt   |  2 +-
 .../ops_history_v2/ApplyGradientDescent.pbtxt |  2 +-
 .../compat/ops_history_v2/ApplyMomentum.pbtxt |  2 +-
 .../ops_history_v2/ApplyPowerSign.pbtxt       |  2 +-
 .../ops_history_v2/ApplyProximalAdagrad.pbtxt |  2 +-
 .../ApplyProximalGradientDescent.pbtxt        |  2 +-
 .../compat/ops_history_v2/ApplyRMSProp.pbtxt  |  2 +-
 .../compat/ops_history_v2/ApproxTopK.pbtxt    |  2 +-
 .../ops_history_v2/ApproximateEqual.pbtxt     |  2 +-
 .../ops/compat/ops_history_v2/ArgMax.pbtxt    |  2 +-
 .../ops/compat/ops_history_v2/ArgMin.pbtxt    |  2 +-
 .../ops/compat/ops_history_v2/AsString.pbtxt  |  2 +-
 .../core/ops/compat/ops_history_v2/Asin.pbtxt |  2 +-
 .../ops/compat/ops_history_v2/Asinh.pbtxt     |  2 +-
 .../ops/compat/ops_history_v2/Assert.pbtxt    |  2 +-
 .../AssertCardinalityDataset.pbtxt            |  2 +-
 .../ops_history_v2/AssertNextDataset.pbtxt    |  2 +-
 .../ops_history_v2/AssertPrevDataset.pbtxt    |  2 +-
 .../ops/compat/ops_history_v2/Assign.pbtxt    |  2 +-
 .../ops/compat/ops_history_v2/AssignAdd.pbtxt |  2 +-
 .../ops_history_v2/AssignAddVariableOp.pbtxt  |  2 +-
 .../ops/compat/ops_history_v2/AssignSub.pbtxt |  2 +-
 .../ops_history_v2/AssignSubVariableOp.pbtxt  |  2 +-
 .../ops_history_v2/AssignVariableOp.pbtxt     |  2 +-
 .../AssignVariableXlaConcatND.pbtxt           |  2 +-
 .../core/ops/compat/ops_history_v2/Atan.pbtxt |  2 +-
 .../ops/compat/ops_history_v2/Atan2.pbtxt     |  2 +-
 .../ops/compat/ops_history_v2/Atanh.pbtxt     |  2 +-
 .../ops_history_v2/AudioSpectrogram.pbtxt     |  2 +-
 .../compat/ops_history_v2/AudioSummary.pbtxt  |  2 +-
 .../ops_history_v2/AudioSummaryV2.pbtxt       |  2 +-
 .../ops_history_v2/AutoShardDataset.pbtxt     |  2 +-
 .../ops/compat/ops_history_v2/AvgPool.pbtxt   |  2 +-
 .../ops/compat/ops_history_v2/AvgPool3D.pbtxt |  2 +-
 .../compat/ops_history_v2/AvgPool3DGrad.pbtxt |  2 +-
 .../compat/ops_history_v2/AvgPoolGrad.pbtxt   |  2 +-
 .../BandedTriangularSolve.pbtxt               |  2 +-
 .../ops/compat/ops_history_v2/Barrier.pbtxt   |  2 +-
 .../compat/ops_history_v2/BarrierClose.pbtxt  |  2 +-
 .../BarrierIncompleteSize.pbtxt               |  2 +-
 .../ops_history_v2/BarrierInsertMany.pbtxt    |  2 +-
 .../ops_history_v2/BarrierReadySize.pbtxt     |  2 +-
 .../ops_history_v2/BarrierTakeMany.pbtxt      |  2 +-
 .../ops/compat/ops_history_v2/Batch.pbtxt     |  2 +-
 .../compat/ops_history_v2/BatchCholesky.pbtxt |  2 +-
 .../ops_history_v2/BatchCholeskyGrad.pbtxt    |  2 +-
 .../compat/ops_history_v2/BatchDataset.pbtxt  |  2 +-
 .../ops_history_v2/BatchDatasetV2.pbtxt       |  2 +-
 .../ops/compat/ops_history_v2/BatchFFT.pbtxt  |  2 +-
 .../compat/ops_history_v2/BatchFFT2D.pbtxt    |  2 +-
 .../compat/ops_history_v2/BatchFFT3D.pbtxt    |  2 +-
 .../compat/ops_history_v2/BatchFunction.pbtxt |  2 +-
 .../ops/compat/ops_history_v2/BatchIFFT.pbtxt |  2 +-
 .../compat/ops_history_v2/BatchIFFT2D.pbtxt   |  2 +-
 .../compat/ops_history_v2/BatchIFFT3D.pbtxt   |  2 +-
 .../compat/ops_history_v2/BatchMatMul.pbtxt   |  2 +-
 .../compat/ops_history_v2/BatchMatMulV2.pbtxt |  2 +-
 .../compat/ops_history_v2/BatchMatMulV3.pbtxt |  2 +-
 .../ops_history_v2/BatchMatrixBandPart.pbtxt  |  2 +-
 .../BatchMatrixDeterminant.pbtxt              |  2 +-
 .../ops_history_v2/BatchMatrixDiag.pbtxt      |  2 +-
 .../ops_history_v2/BatchMatrixDiagPart.pbtxt  |  2 +-
 .../ops_history_v2/BatchMatrixInverse.pbtxt   |  2 +-
 .../ops_history_v2/BatchMatrixSetDiag.pbtxt   |  2 +-
 .../ops_history_v2/BatchMatrixSolve.pbtxt     |  2 +-
 .../ops_history_v2/BatchMatrixSolveLs.pbtxt   |  2 +-
 .../BatchMatrixTriangularSolve.pbtxt          |  2 +-
 .../BatchNormWithGlobalNormalization.pbtxt    |  2 +-
 ...BatchNormWithGlobalNormalizationGrad.pbtxt |  2 +-
 .../ops_history_v2/BatchSelfAdjointEig.pbtxt  |  2 +-
 .../BatchSelfAdjointEigV2.pbtxt               |  2 +-
 .../ops/compat/ops_history_v2/BatchSvd.pbtxt  |  2 +-
 .../compat/ops_history_v2/BatchToSpace.pbtxt  |  2 +-
 .../ops_history_v2/BatchToSpaceND.pbtxt       |  2 +-
 .../ops/compat/ops_history_v2/BesselI0.pbtxt  |  2 +-
 .../ops/compat/ops_history_v2/BesselI0e.pbtxt |  2 +-
 .../ops/compat/ops_history_v2/BesselI1.pbtxt  |  2 +-
 .../ops/compat/ops_history_v2/BesselI1e.pbtxt |  2 +-
 .../ops/compat/ops_history_v2/BesselJ0.pbtxt  |  2 +-
 .../ops/compat/ops_history_v2/BesselJ1.pbtxt  |  2 +-
 .../ops/compat/ops_history_v2/BesselK0.pbtxt  |  2 +-
 .../ops/compat/ops_history_v2/BesselK0e.pbtxt |  2 +-
 .../ops/compat/ops_history_v2/BesselK1.pbtxt  |  2 +-
 .../ops/compat/ops_history_v2/BesselK1e.pbtxt |  2 +-
 .../ops/compat/ops_history_v2/BesselY0.pbtxt  |  2 +-
 .../ops/compat/ops_history_v2/BesselY1.pbtxt  |  2 +-
 .../ops/compat/ops_history_v2/Betainc.pbtxt   |  2 +-
 .../ops/compat/ops_history_v2/BiasAdd.pbtxt   |  2 +-
 .../compat/ops_history_v2/BiasAddGrad.pbtxt   |  2 +-
 .../ops/compat/ops_history_v2/BiasAddV1.pbtxt |  2 +-
 .../ops/compat/ops_history_v2/Bincount.pbtxt  |  2 +-
 .../ops/compat/ops_history_v2/Bitcast.pbtxt   |  2 +-
 .../compat/ops_history_v2/BitwiseAnd.pbtxt    |  2 +-
 .../ops/compat/ops_history_v2/BitwiseOr.pbtxt |  2 +-
 .../compat/ops_history_v2/BitwiseXor.pbtxt    |  2 +-
 .../ops/compat/ops_history_v2/BlockLSTM.pbtxt |  2 +-
 .../compat/ops_history_v2/BlockLSTMGrad.pbtxt |  2 +-
 .../ops_history_v2/BlockLSTMGradV2.pbtxt      |  2 +-
 .../compat/ops_history_v2/BlockLSTMV2.pbtxt   |  2 +-
 .../BoostedTreesAggregateStats.pbtxt          |  2 +-
 .../BoostedTreesBucketize.pbtxt               |  2 +-
 ...oostedTreesCalculateBestFeatureSplit.pbtxt |  2 +-
 ...stedTreesCalculateBestFeatureSplitV2.pbtxt |  2 +-
 ...tedTreesCalculateBestGainsPerFeature.pbtxt |  2 +-
 .../BoostedTreesCenterBias.pbtxt              |  2 +-
 .../BoostedTreesCreateEnsemble.pbtxt          |  2 +-
 ...tedTreesCreateQuantileStreamResource.pbtxt |  2 +-
 .../BoostedTreesDeserializeEnsemble.pbtxt     |  2 +-
 ...BoostedTreesEnsembleResourceHandleOp.pbtxt |  2 +-
 .../BoostedTreesExampleDebugOutputs.pbtxt     |  2 +-
 .../BoostedTreesFlushQuantileSummaries.pbtxt  |  2 +-
 .../BoostedTreesGetEnsembleStates.pbtxt       |  2 +-
 .../BoostedTreesMakeQuantileSummaries.pbtxt   |  2 +-
 .../BoostedTreesMakeStatsSummary.pbtxt        |  2 +-
 .../ops_history_v2/BoostedTreesPredict.pbtxt  |  2 +-
 ...esQuantileStreamResourceAddSummaries.pbtxt |  2 +-
 ...eesQuantileStreamResourceDeserialize.pbtxt |  2 +-
 ...stedTreesQuantileStreamResourceFlush.pbtxt |  2 +-
 ...ileStreamResourceGetBucketBoundaries.pbtxt |  2 +-
 ...dTreesQuantileStreamResourceHandleOp.pbtxt |  2 +-
 .../BoostedTreesSerializeEnsemble.pbtxt       |  2 +-
 .../BoostedTreesSparseAggregateStats.pbtxt    |  2 +-
 ...TreesSparseCalculateBestFeatureSplit.pbtxt |  2 +-
 .../BoostedTreesTrainingPredict.pbtxt         |  2 +-
 .../BoostedTreesUpdateEnsemble.pbtxt          |  2 +-
 .../BoostedTreesUpdateEnsembleV2.pbtxt        |  2 +-
 .../compat/ops_history_v2/BroadcastArgs.pbtxt |  2 +-
 .../BroadcastGradientArgs.pbtxt               |  2 +-
 .../compat/ops_history_v2/BroadcastTo.pbtxt   |  2 +-
 .../ops/compat/ops_history_v2/Bucketize.pbtxt |  2 +-
 .../BytesProducedStatsDataset.pbtxt           |  2 +-
 .../CSRSparseMatrixComponents.pbtxt           |  2 +-
 .../CSRSparseMatrixToDense.pbtxt              |  2 +-
 .../CSRSparseMatrixToSparseTensor.pbtxt       |  2 +-
 .../compat/ops_history_v2/CSVDataset.pbtxt    |  2 +-
 .../compat/ops_history_v2/CSVDatasetV2.pbtxt  |  2 +-
 .../ops_history_v2/CTCBeamSearchDecoder.pbtxt |  2 +-
 .../ops_history_v2/CTCGreedyDecoder.pbtxt     |  2 +-
 .../ops/compat/ops_history_v2/CTCLoss.pbtxt   |  2 +-
 .../ops/compat/ops_history_v2/CTCLossV2.pbtxt |  2 +-
 .../compat/ops_history_v2/CacheDataset.pbtxt  |  2 +-
 .../ops_history_v2/CacheDatasetV2.pbtxt       |  2 +-
 .../core/ops/compat/ops_history_v2/Case.pbtxt |  2 +-
 .../core/ops/compat/ops_history_v2/Cast.pbtxt |  2 +-
 .../core/ops/compat/ops_history_v2/Ceil.pbtxt |  2 +-
 .../compat/ops_history_v2/CheckNumerics.pbtxt |  2 +-
 .../ops_history_v2/CheckNumericsV2.pbtxt      |  2 +-
 .../ops/compat/ops_history_v2/Cholesky.pbtxt  |  2 +-
 .../compat/ops_history_v2/CholeskyGrad.pbtxt  |  2 +-
 .../ChooseFastestBranchDataset.pbtxt          |  2 +-
 .../ops_history_v2/ChooseFastestDataset.pbtxt |  2 +-
 .../compat/ops_history_v2/ClipByValue.pbtxt   |  2 +-
 .../ops_history_v2/CloseSummaryWriter.pbtxt   |  2 +-
 .../CollateTPUEmbeddingMemory.pbtxt           |  2 +-
 .../ops_history_v2/CollectiveAllToAllV2.pbtxt |  2 +-
 .../ops_history_v2/CollectiveAllToAllV3.pbtxt |  2 +-
 .../CollectiveAssignGroupV2.pbtxt             |  2 +-
 .../ops_history_v2/CollectiveBcastRecv.pbtxt  |  2 +-
 .../CollectiveBcastRecvV2.pbtxt               |  2 +-
 .../ops_history_v2/CollectiveBcastSend.pbtxt  |  2 +-
 .../CollectiveBcastSendV2.pbtxt               |  2 +-
 .../ops_history_v2/CollectiveGather.pbtxt     |  2 +-
 .../ops_history_v2/CollectiveGatherV2.pbtxt   |  2 +-
 .../CollectiveInitializeCommunicator.pbtxt    |  2 +-
 .../ops_history_v2/CollectivePermute.pbtxt    |  2 +-
 .../ops_history_v2/CollectiveReduce.pbtxt     |  2 +-
 .../CollectiveReduceScatterV2.pbtxt           |  2 +-
 .../ops_history_v2/CollectiveReduceV2.pbtxt   |  2 +-
 .../ops_history_v2/CollectiveReduceV3.pbtxt   |  2 +-
 .../CombinedNonMaxSuppression.pbtxt           |  2 +-
 .../ops/compat/ops_history_v2/Complex.pbtxt   |  2 +-
 .../compat/ops_history_v2/ComplexAbs.pbtxt    |  2 +-
 ...CompositeTensorVariantFromComponents.pbtxt |  2 +-
 .../CompositeTensorVariantToComponents.pbtxt  |  2 +-
 .../ops_history_v2/CompressElement.pbtxt      |  2 +-
 .../ComputeAccidentalHits.pbtxt               |  2 +-
 .../ops_history_v2/ComputeBatchSize.pbtxt     |  2 +-
 .../ops_history_v2/ComputeDedupDataSize.pbtxt |  2 +-
 .../ComputeDedupDataSizeV2.pbtxt              |  2 +-
 .../ComputeDedupDataTupleMask.pbtxt           |  2 +-
 .../ComputeDedupDataTupleMaskV2.pbtxt         |  2 +-
 .../ops/compat/ops_history_v2/Concat.pbtxt    |  2 +-
 .../compat/ops_history_v2/ConcatOffset.pbtxt  |  2 +-
 .../ops/compat/ops_history_v2/ConcatV2.pbtxt  |  2 +-
 .../ops_history_v2/ConcatenateDataset.pbtxt   |  2 +-
 .../ConditionalAccumulator.pbtxt              |  2 +-
 .../ConfigureAndInitializeGlobalTPU.pbtxt     |  2 +-
 .../ConfigureDistributedTPU.pbtxt             |  2 +-
 .../ConfigureTPUEmbedding.pbtxt               |  2 +-
 .../ConfigureTPUEmbeddingHost.pbtxt           |  2 +-
 .../ConfigureTPUEmbeddingMemory.pbtxt         |  2 +-
 .../core/ops/compat/ops_history_v2/Conj.pbtxt |  2 +-
 .../ops_history_v2/ConjugateTranspose.pbtxt   |  2 +-
 .../ConnectTPUEmbeddingHosts.pbtxt            |  2 +-
 .../ops/compat/ops_history_v2/Const.pbtxt     |  2 +-
 .../ops_history_v2/ConsumeMutexLock.pbtxt     |  2 +-
 .../ops_history_v2/ControlTrigger.pbtxt       |  2 +-
 .../core/ops/compat/ops_history_v2/Conv.pbtxt |  2 +-
 .../ops/compat/ops_history_v2/Conv2D.pbtxt    |  2 +-
 .../ops_history_v2/Conv2DBackpropFilter.pbtxt |  2 +-
 .../Conv2DBackpropFilterV2.pbtxt              |  2 +-
 .../ops_history_v2/Conv2DBackpropInput.pbtxt  |  2 +-
 .../Conv2DBackpropInputV2.pbtxt               |  2 +-
 .../ops/compat/ops_history_v2/Conv3D.pbtxt    |  2 +-
 .../ops_history_v2/Conv3DBackpropFilter.pbtxt |  2 +-
 .../Conv3DBackpropFilterV2.pbtxt              |  2 +-
 .../ops_history_v2/Conv3DBackpropInput.pbtxt  |  2 +-
 .../Conv3DBackpropInputV2.pbtxt               |  2 +-
 .../ops_history_v2/ConvertToCooTensor.pbtxt   |  2 +-
 .../ConvertToListOfSparseCoreCooTensors.pbtxt |  2 +-
 ...nvertToSparseCoreCsrWrappedCooTensor.pbtxt |  2 +-
 .../core/ops/compat/ops_history_v2/Copy.pbtxt |  2 +-
 .../ops/compat/ops_history_v2/CopyHost.pbtxt  |  2 +-
 .../compat/ops_history_v2/CopyToMesh.pbtxt    |  2 +-
 .../ops_history_v2/CopyToMeshGrad.pbtxt       |  2 +-
 .../core/ops/compat/ops_history_v2/Cos.pbtxt  |  2 +-
 .../core/ops/compat/ops_history_v2/Cosh.pbtxt |  2 +-
 .../ops/compat/ops_history_v2/CountUpTo.pbtxt |  2 +-
 .../CreateSummaryDbWriter.pbtxt               |  2 +-
 .../CreateSummaryFileWriter.pbtxt             |  2 +-
 .../compat/ops_history_v2/CropAndResize.pbtxt |  2 +-
 .../CropAndResizeGradBoxes.pbtxt              |  2 +-
 .../CropAndResizeGradImage.pbtxt              |  2 +-
 .../ops/compat/ops_history_v2/Cross.pbtxt     |  2 +-
 .../ops_history_v2/CrossReplicaSum.pbtxt      |  2 +-
 .../ops/compat/ops_history_v2/CudnnRNN.pbtxt  |  2 +-
 .../ops_history_v2/CudnnRNNBackprop.pbtxt     |  2 +-
 .../ops_history_v2/CudnnRNNBackpropV2.pbtxt   |  2 +-
 .../ops_history_v2/CudnnRNNBackpropV3.pbtxt   |  2 +-
 .../CudnnRNNCanonicalToParams.pbtxt           |  2 +-
 .../CudnnRNNCanonicalToParamsV2.pbtxt         |  2 +-
 .../ops_history_v2/CudnnRNNParamsSize.pbtxt   |  2 +-
 .../CudnnRNNParamsToCanonical.pbtxt           |  2 +-
 .../CudnnRNNParamsToCanonicalV2.pbtxt         |  2 +-
 .../compat/ops_history_v2/CudnnRNNV2.pbtxt    |  2 +-
 .../compat/ops_history_v2/CudnnRNNV3.pbtxt    |  2 +-
 .../ops/compat/ops_history_v2/Cumprod.pbtxt   |  2 +-
 .../ops/compat/ops_history_v2/Cumsum.pbtxt    |  2 +-
 .../ops_history_v2/CumulativeLogsumexp.pbtxt  |  2 +-
 .../ops_history_v2/DTensorRestoreV2.pbtxt     |  2 +-
 .../DTensorSetGlobalTPUArray.pbtxt            |  2 +-
 .../ops_history_v2/DataFormatDimMap.pbtxt     |  2 +-
 .../ops_history_v2/DataFormatVecPermute.pbtxt |  2 +-
 .../ops_history_v2/DataServiceDataset.pbtxt   |  2 +-
 .../ops_history_v2/DataServiceDatasetV2.pbtxt |  2 +-
 .../ops_history_v2/DataServiceDatasetV3.pbtxt |  2 +-
 .../ops_history_v2/DataServiceDatasetV4.pbtxt |  2 +-
 .../ops_history_v2/DatasetCardinality.pbtxt   |  2 +-
 .../ops_history_v2/DatasetFingerprint.pbtxt   |  2 +-
 .../ops_history_v2/DatasetFromGraph.pbtxt     |  2 +-
 .../ops_history_v2/DatasetToGraph.pbtxt       |  2 +-
 .../ops_history_v2/DatasetToGraphV2.pbtxt     |  2 +-
 .../DatasetToSingleElement.pbtxt              |  2 +-
 .../ops_history_v2/DatasetToTFRecord.pbtxt    |  2 +-
 .../ops/compat/ops_history_v2/Dawsn.pbtxt     |  2 +-
 .../DebugGradientIdentity.pbtxt               |  2 +-
 .../DebugGradientRefIdentity.pbtxt            |  2 +-
 .../compat/ops_history_v2/DebugIdentity.pbtxt |  2 +-
 .../ops_history_v2/DebugIdentityV2.pbtxt      |  2 +-
 .../ops_history_v2/DebugIdentityV3.pbtxt      |  2 +-
 .../compat/ops_history_v2/DebugNanCount.pbtxt |  2 +-
 .../ops_history_v2/DebugNumericSummary.pbtxt  |  2 +-
 .../DebugNumericSummaryV2.pbtxt               |  2 +-
 .../ops_history_v2/DecodeAndCropJpeg.pbtxt    |  2 +-
 .../compat/ops_history_v2/DecodeBase64.pbtxt  |  2 +-
 .../ops/compat/ops_history_v2/DecodeBmp.pbtxt |  2 +-
 .../ops/compat/ops_history_v2/DecodeCSV.pbtxt |  2 +-
 .../ops_history_v2/DecodeCompressed.pbtxt     |  2 +-
 .../ops/compat/ops_history_v2/DecodeGif.pbtxt |  2 +-
 .../compat/ops_history_v2/DecodeImage.pbtxt   |  2 +-
 .../ops_history_v2/DecodeJSONExample.pbtxt    |  2 +-
 .../compat/ops_history_v2/DecodeJpeg.pbtxt    |  2 +-
 .../ops_history_v2/DecodePaddedRaw.pbtxt      |  2 +-
 .../ops/compat/ops_history_v2/DecodePng.pbtxt |  2 +-
 .../compat/ops_history_v2/DecodeProtoV2.pbtxt |  2 +-
 .../ops/compat/ops_history_v2/DecodeRaw.pbtxt |  2 +-
 .../ops/compat/ops_history_v2/DecodeWav.pbtxt |  2 +-
 .../ops/compat/ops_history_v2/DeepCopy.pbtxt  |  2 +-
 .../ops_history_v2/DeleteIterator.pbtxt       |  2 +-
 .../ops_history_v2/DeleteMemoryCache.pbtxt    |  2 +-
 .../DeleteMultiDeviceIterator.pbtxt           |  2 +-
 .../DeleteRandomSeedGenerator.pbtxt           |  2 +-
 .../ops_history_v2/DeleteSeedGenerator.pbtxt  |  2 +-
 .../ops_history_v2/DeleteSessionTensor.pbtxt  |  2 +-
 .../compat/ops_history_v2/DenseBincount.pbtxt |  2 +-
 .../DenseCountSparseOutput.pbtxt              |  2 +-
 .../DenseToCSRSparseMatrix.pbtxt              |  2 +-
 .../DenseToDenseSetOperation.pbtxt            |  2 +-
 .../DenseToSparseBatchDataset.pbtxt           |  2 +-
 .../DenseToSparseSetOperation.pbtxt           |  2 +-
 .../compat/ops_history_v2/DepthToSpace.pbtxt  |  2 +-
 .../DepthwiseConv2dNative.pbtxt               |  2 +-
 .../DepthwiseConv2dNativeBackpropFilter.pbtxt |  2 +-
 .../DepthwiseConv2dNativeBackpropInput.pbtxt  |  2 +-
 .../compat/ops_history_v2/Dequantize.pbtxt    |  2 +-
 .../ops_history_v2/DeserializeIterator.pbtxt  |  2 +-
 .../DeserializeManySparse.pbtxt               |  2 +-
 .../ops_history_v2/DeserializeSparse.pbtxt    |  2 +-
 .../ops_history_v2/DestroyResourceOp.pbtxt    |  2 +-
 .../DestroyTemporaryVariable.pbtxt            |  2 +-
 .../compat/ops_history_v2/DeviceIndex.pbtxt   |  2 +-
 .../core/ops/compat/ops_history_v2/Diag.pbtxt |  2 +-
 .../ops/compat/ops_history_v2/DiagPart.pbtxt  |  2 +-
 .../ops/compat/ops_history_v2/Digamma.pbtxt   |  2 +-
 .../compat/ops_history_v2/Dilation2D.pbtxt    |  2 +-
 .../Dilation2DBackpropFilter.pbtxt            |  2 +-
 .../Dilation2DBackpropInput.pbtxt             |  2 +-
 .../DirectedInterleaveDataset.pbtxt           |  2 +-
 .../ops_history_v2/DisableCopyOnRead.pbtxt    |  2 +-
 .../ops_history_v2/DistributedSave.pbtxt      |  2 +-
 .../core/ops/compat/ops_history_v2/Div.pbtxt  |  2 +-
 .../ops/compat/ops_history_v2/DivNoNan.pbtxt  |  2 +-
 .../ops_history_v2/DrawBoundingBoxes.pbtxt    |  2 +-
 .../ops_history_v2/DrawBoundingBoxesV2.pbtxt  |  2 +-
 .../DummyIterationCounter.pbtxt               |  2 +-
 .../ops_history_v2/DummyMemoryCache.pbtxt     |  2 +-
 .../ops_history_v2/DummySeedGenerator.pbtxt   |  2 +-
 ...ueueTPUEmbeddingArbitraryTensorBatch.pbtxt |  2 +-
 ...EnqueueTPUEmbeddingRaggedTensorBatch.pbtxt |  2 +-
 .../ops_history_v2/DynamicPartition.pbtxt     |  2 +-
 .../compat/ops_history_v2/DynamicStitch.pbtxt |  2 +-
 .../compat/ops_history_v2/EagerPyFunc.pbtxt   |  2 +-
 .../compat/ops_history_v2/EditDistance.pbtxt  |  2 +-
 .../core/ops/compat/ops_history_v2/Eig.pbtxt  |  2 +-
 .../ops/compat/ops_history_v2/Einsum.pbtxt    |  2 +-
 .../core/ops/compat/ops_history_v2/Elu.pbtxt  |  2 +-
 .../ops/compat/ops_history_v2/EluGrad.pbtxt   |  2 +-
 .../ops/compat/ops_history_v2/Empty.pbtxt     |  2 +-
 .../ops_history_v2/EmptyTensorList.pbtxt      |  2 +-
 .../ops_history_v2/EmptyTensorMap.pbtxt       |  2 +-
 .../compat/ops_history_v2/EncodeBase64.pbtxt  |  2 +-
 .../compat/ops_history_v2/EncodeJpeg.pbtxt    |  2 +-
 .../EncodeJpegVariableQuality.pbtxt           |  2 +-
 .../ops/compat/ops_history_v2/EncodePng.pbtxt |  2 +-
 .../compat/ops_history_v2/EncodeProto.pbtxt   |  2 +-
 .../ops/compat/ops_history_v2/EncodeWav.pbtxt |  2 +-
 ...ueueTPUEmbeddingArbitraryTensorBatch.pbtxt |  2 +-
 .../EnqueueTPUEmbeddingBatch.pbtxt            |  2 +-
 .../EnqueueTPUEmbeddingIntegerBatch.pbtxt     |  2 +-
 ...EnqueueTPUEmbeddingRaggedTensorBatch.pbtxt |  2 +-
 .../EnqueueTPUEmbeddingSparseBatch.pbtxt      |  2 +-
 ...EnqueueTPUEmbeddingSparseTensorBatch.pbtxt |  2 +-
 .../compat/ops_history_v2/EnsureShape.pbtxt   |  2 +-
 .../ops/compat/ops_history_v2/Enter.pbtxt     |  2 +-
 .../ops/compat/ops_history_v2/Equal.pbtxt     |  2 +-
 .../core/ops/compat/ops_history_v2/Erf.pbtxt  |  2 +-
 .../core/ops/compat/ops_history_v2/Erfc.pbtxt |  2 +-
 .../ops/compat/ops_history_v2/Erfinv.pbtxt    |  2 +-
 .../compat/ops_history_v2/EuclideanNorm.pbtxt |  2 +-
 .../ExecuteTPUEmbeddingPartitioner.pbtxt      |  2 +-
 .../core/ops/compat/ops_history_v2/Exit.pbtxt |  2 +-
 .../core/ops/compat/ops_history_v2/Exp.pbtxt  |  2 +-
 .../compat/ops_history_v2/ExpandDims.pbtxt    |  2 +-
 .../ExperimentalAssertNextDataset.pbtxt       |  2 +-
 .../ExperimentalAutoShardDataset.pbtxt        |  2 +-
 ...xperimentalBytesProducedStatsDataset.pbtxt |  2 +-
 .../ExperimentalCSVDataset.pbtxt              |  2 +-
 .../ExperimentalChooseFastestDataset.pbtxt    |  2 +-
 .../ExperimentalDatasetCardinality.pbtxt      |  2 +-
 .../ExperimentalDatasetToTFRecord.pbtxt       |  2 +-
 ...xperimentalDenseToSparseBatchDataset.pbtxt |  2 +-
 ...xperimentalDirectedInterleaveDataset.pbtxt |  2 +-
 .../ExperimentalGroupByReducerDataset.pbtxt   |  2 +-
 .../ExperimentalGroupByWindowDataset.pbtxt    |  2 +-
 .../ExperimentalIgnoreErrorsDataset.pbtxt     |  2 +-
 .../ExperimentalIteratorGetDevice.pbtxt       |  2 +-
 .../ExperimentalLMDBDataset.pbtxt             |  2 +-
 .../ExperimentalLatencyStatsDataset.pbtxt     |  2 +-
 .../ExperimentalMapAndBatchDataset.pbtxt      |  2 +-
 .../ExperimentalMapDataset.pbtxt              |  2 +-
 .../ExperimentalMatchingFilesDataset.pbtxt    |  2 +-
 ...rimentalMaxIntraOpParallelismDataset.pbtxt |  2 +-
 .../ExperimentalNonSerializableDataset.pbtxt  |  2 +-
 ...xperimentalParallelInterleaveDataset.pbtxt |  2 +-
 .../ExperimentalParseExampleDataset.pbtxt     |  2 +-
 ...ExperimentalPrivateThreadPoolDataset.pbtxt |  2 +-
 .../ExperimentalRandomDataset.pbtxt           |  2 +-
 .../ExperimentalRebatchDataset.pbtxt          |  2 +-
 .../ExperimentalScanDataset.pbtxt             |  2 +-
 ...xperimentalSetStatsAggregatorDataset.pbtxt |  2 +-
 .../ExperimentalSleepDataset.pbtxt            |  2 +-
 .../ExperimentalSlidingWindowDataset.pbtxt    |  2 +-
 .../ExperimentalSqlDataset.pbtxt              |  2 +-
 .../ExperimentalStatsAggregatorHandle.pbtxt   |  2 +-
 .../ExperimentalStatsAggregatorSummary.pbtxt  |  2 +-
 .../ExperimentalTakeWhileDataset.pbtxt        |  2 +-
 .../ExperimentalThreadPoolDataset.pbtxt       |  2 +-
 .../ExperimentalThreadPoolHandle.pbtxt        |  2 +-
 .../ExperimentalUnbatchDataset.pbtxt          |  2 +-
 .../ExperimentalUniqueDataset.pbtxt           |  2 +-
 .../ops/compat/ops_history_v2/Expint.pbtxt    |  2 +-
 .../ops/compat/ops_history_v2/Expm1.pbtxt     |  2 +-
 .../ops_history_v2/ExtractGlimpse.pbtxt       |  2 +-
 .../ops_history_v2/ExtractGlimpseV2.pbtxt     |  2 +-
 .../ops_history_v2/ExtractImagePatches.pbtxt  |  2 +-
 .../ops_history_v2/ExtractJpegShape.pbtxt     |  2 +-
 .../ops_history_v2/ExtractVolumePatches.pbtxt |  2 +-
 .../core/ops/compat/ops_history_v2/FFT.pbtxt  |  2 +-
 .../ops/compat/ops_history_v2/FFT2D.pbtxt     |  2 +-
 .../ops/compat/ops_history_v2/FFT3D.pbtxt     |  2 +-
 .../ops/compat/ops_history_v2/FFTND.pbtxt     |  2 +-
 .../ops/compat/ops_history_v2/FIFOQueue.pbtxt |  2 +-
 .../compat/ops_history_v2/FIFOQueueV2.pbtxt   |  2 +-
 .../core/ops/compat/ops_history_v2/Fact.pbtxt |  2 +-
 .../ops/compat/ops_history_v2/FakeParam.pbtxt |  2 +-
 .../FakeQuantWithMinMaxArgs.pbtxt             |  2 +-
 .../FakeQuantWithMinMaxArgsGradient.pbtxt     |  2 +-
 .../FakeQuantWithMinMaxVars.pbtxt             |  2 +-
 .../FakeQuantWithMinMaxVarsGradient.pbtxt     |  2 +-
 .../FakeQuantWithMinMaxVarsPerChannel.pbtxt   |  2 +-
 ...uantWithMinMaxVarsPerChannelGradient.pbtxt |  2 +-
 .../ops/compat/ops_history_v2/FakeQueue.pbtxt |  2 +-
 .../FileSystemSetConfiguration.pbtxt          |  2 +-
 .../core/ops/compat/ops_history_v2/Fill.pbtxt |  2 +-
 .../FilterByLastComponentDataset.pbtxt        |  2 +-
 .../compat/ops_history_v2/FilterDataset.pbtxt |  2 +-
 .../ops_history_v2/FinalizeDataset.pbtxt      |  2 +-
 .../ops_history_v2/FinalizeTPUEmbedding.pbtxt |  2 +-
 .../FinalizeTPUEmbeddingV2.pbtxt              |  2 +-
 .../compat/ops_history_v2/Fingerprint.pbtxt   |  2 +-
 .../FixedLengthRecordDataset.pbtxt            |  2 +-
 .../FixedLengthRecordDatasetV2.pbtxt          |  2 +-
 .../FixedLengthRecordReader.pbtxt             |  2 +-
 .../FixedLengthRecordReaderV2.pbtxt           |  2 +-
 .../FixedUnigramCandidateSampler.pbtxt        |  2 +-
 .../ops_history_v2/FlatMapDataset.pbtxt       |  2 +-
 .../ops/compat/ops_history_v2/Floor.pbtxt     |  2 +-
 .../ops/compat/ops_history_v2/FloorDiv.pbtxt  |  2 +-
 .../ops/compat/ops_history_v2/FloorMod.pbtxt  |  2 +-
 .../ops_history_v2/FlushSummaryWriter.pbtxt   |  2 +-
 .../core/ops/compat/ops_history_v2/For.pbtxt  |  2 +-
 .../ops_history_v2/FractionalAvgPool.pbtxt    |  2 +-
 .../FractionalAvgPoolGrad.pbtxt               |  2 +-
 .../ops_history_v2/FractionalMaxPool.pbtxt    |  2 +-
 .../FractionalMaxPoolGrad.pbtxt               |  2 +-
 .../compat/ops_history_v2/FresnelCos.pbtxt    |  2 +-
 .../compat/ops_history_v2/FresnelSin.pbtxt    |  2 +-
 .../ops_history_v2/FusedBatchNorm.pbtxt       |  2 +-
 .../ops_history_v2/FusedBatchNormGrad.pbtxt   |  2 +-
 .../ops_history_v2/FusedBatchNormGradV2.pbtxt |  2 +-
 .../ops_history_v2/FusedBatchNormGradV3.pbtxt |  2 +-
 .../ops_history_v2/FusedBatchNormV2.pbtxt     |  2 +-
 .../ops_history_v2/FusedBatchNormV3.pbtxt     |  2 +-
 .../ops_history_v2/FusedPadConv2D.pbtxt       |  2 +-
 .../FusedResizeAndPadConv2D.pbtxt             |  2 +-
 .../compat/ops_history_v2/GRUBlockCell.pbtxt  |  2 +-
 .../ops_history_v2/GRUBlockCellGrad.pbtxt     |  2 +-
 .../ops/compat/ops_history_v2/Gather.pbtxt    |  2 +-
 .../ops/compat/ops_history_v2/GatherNd.pbtxt  |  2 +-
 .../ops/compat/ops_history_v2/GatherV2.pbtxt  |  2 +-
 .../GenerateBoundingBoxProposals.pbtxt        |  2 +-
 .../GenerateVocabRemapping.pbtxt              |  2 +-
 .../ops_history_v2/GeneratorDataset.pbtxt     |  2 +-
 .../ops_history_v2/GetElementAtIndex.pbtxt    |  2 +-
 ...etMinibatchSplitsWithPhysicalReplica.pbtxt |  2 +-
 ...tMinibatchesInCsrWithPhysicalReplica.pbtxt |  2 +-
 .../compat/ops_history_v2/GetOptions.pbtxt    |  2 +-
 .../ops_history_v2/GetSessionHandle.pbtxt     |  2 +-
 .../ops_history_v2/GetSessionHandleV2.pbtxt   |  2 +-
 .../ops_history_v2/GetSessionTensor.pbtxt     |  2 +-
 ...tStatsFromListOfSparseCoreCooTensors.pbtxt |  2 +-
 .../compat/ops_history_v2/GetTpuTaskId.pbtxt  |  1 -
 .../compat/ops_history_v2/GlobalIterId.pbtxt  |  2 +-
 .../ops_history_v2/GlobalShuffleDataset.pbtxt |  2 +-
 .../ops/compat/ops_history_v2/Greater.pbtxt   |  2 +-
 .../compat/ops_history_v2/GreaterEqual.pbtxt  |  2 +-
 .../GroupByReducerDataset.pbtxt               |  2 +-
 .../ops_history_v2/GroupByWindowDataset.pbtxt |  2 +-
 .../ops_history_v2/GuaranteeConst.pbtxt       |  2 +-
 .../ops/compat/ops_history_v2/HSVToRGB.pbtxt  |  2 +-
 .../ops/compat/ops_history_v2/HashTable.pbtxt |  2 +-
 .../compat/ops_history_v2/HashTableV2.pbtxt   |  2 +-
 .../ops_history_v2/HistogramFixedWidth.pbtxt  |  2 +-
 .../ops_history_v2/HistogramSummary.pbtxt     |  2 +-
 .../ops/compat/ops_history_v2/HostConst.pbtxt |  2 +-
 .../core/ops/compat/ops_history_v2/IFFT.pbtxt |  2 +-
 .../ops/compat/ops_history_v2/IFFT2D.pbtxt    |  2 +-
 .../ops/compat/ops_history_v2/IFFT3D.pbtxt    |  2 +-
 .../ops/compat/ops_history_v2/IFFTND.pbtxt    |  2 +-
 .../ops/compat/ops_history_v2/IRFFT.pbtxt     |  2 +-
 .../ops/compat/ops_history_v2/IRFFT2D.pbtxt   |  2 +-
 .../ops/compat/ops_history_v2/IRFFT3D.pbtxt   |  2 +-
 .../ops/compat/ops_history_v2/IRFFTND.pbtxt   |  2 +-
 .../ops/compat/ops_history_v2/Identity.pbtxt  |  2 +-
 .../ops/compat/ops_history_v2/IdentityN.pbtxt |  2 +-
 .../ops_history_v2/IdentityReader.pbtxt       |  2 +-
 .../ops_history_v2/IdentityReaderV2.pbtxt     |  2 +-
 .../core/ops/compat/ops_history_v2/If.pbtxt   |  2 +-
 .../ops/compat/ops_history_v2/Igamma.pbtxt    |  2 +-
 .../compat/ops_history_v2/IgammaGradA.pbtxt   |  2 +-
 .../ops/compat/ops_history_v2/Igammac.pbtxt   |  2 +-
 .../ops_history_v2/IgnoreErrorsDataset.pbtxt  |  2 +-
 .../core/ops/compat/ops_history_v2/Imag.pbtxt |  2 +-
 .../ImageProjectiveTransformV2.pbtxt          |  2 +-
 .../ImageProjectiveTransformV3.pbtxt          |  2 +-
 .../compat/ops_history_v2/ImageSummary.pbtxt  |  2 +-
 .../ops_history_v2/ImmutableConst.pbtxt       |  2 +-
 .../compat/ops_history_v2/ImportEvent.pbtxt   |  2 +-
 .../ops/compat/ops_history_v2/InTopK.pbtxt    |  2 +-
 .../ops/compat/ops_history_v2/InTopKV2.pbtxt  |  2 +-
 .../ops_history_v2/IndexFlatMapDataset.pbtxt  |  2 +-
 .../compat/ops_history_v2/InfeedDequeue.pbtxt |  2 +-
 .../ops_history_v2/InfeedDequeueTuple.pbtxt   |  2 +-
 .../compat/ops_history_v2/InfeedEnqueue.pbtxt |  2 +-
 .../InfeedEnqueuePrelinearizedBuffer.pbtxt    |  2 +-
 .../ops_history_v2/InfeedEnqueueTuple.pbtxt   |  2 +-
 .../ops_history_v2/InitializeTable.pbtxt      |  2 +-
 .../InitializeTableFromDataset.pbtxt          |  2 +-
 .../InitializeTableFromTextFile.pbtxt         |  2 +-
 .../InitializeTableFromTextFileV2.pbtxt       |  2 +-
 .../ops_history_v2/InitializeTableV2.pbtxt    |  2 +-
 .../compat/ops_history_v2/InplaceAdd.pbtxt    |  2 +-
 .../compat/ops_history_v2/InplaceSub.pbtxt    |  2 +-
 .../compat/ops_history_v2/InplaceUpdate.pbtxt |  2 +-
 .../ops_history_v2/InterleaveDataset.pbtxt    |  2 +-
 .../core/ops/compat/ops_history_v2/Inv.pbtxt  |  2 +-
 .../ops/compat/ops_history_v2/InvGrad.pbtxt   |  2 +-
 .../ops/compat/ops_history_v2/Invert.pbtxt    |  2 +-
 .../ops_history_v2/InvertPermutation.pbtxt    |  2 +-
 .../IsBoostedTreesEnsembleInitialized.pbtxt   |  2 +-
 ...eesQuantileStreamResourceInitialized.pbtxt |  2 +-
 .../ops/compat/ops_history_v2/IsFinite.pbtxt  |  2 +-
 .../ops/compat/ops_history_v2/IsInf.pbtxt     |  2 +-
 .../ops/compat/ops_history_v2/IsNan.pbtxt     |  2 +-
 .../IsTPUEmbeddingInitialized.pbtxt           |  2 +-
 .../IsVariableInitialized.pbtxt               |  2 +-
 .../ops_history_v2/IsotonicRegression.pbtxt   |  2 +-
 .../ops/compat/ops_history_v2/Iterator.pbtxt  |  2 +-
 .../IteratorFromStringHandle.pbtxt            |  2 +-
 .../IteratorFromStringHandleV2.pbtxt          |  2 +-
 .../ops_history_v2/IteratorGetDevice.pbtxt    |  2 +-
 .../IteratorGetModelProto.pbtxt               |  2 +-
 .../ops_history_v2/IteratorGetNext.pbtxt      |  2 +-
 .../IteratorGetNextAsOptional.pbtxt           |  2 +-
 .../ops_history_v2/IteratorGetNextSync.pbtxt  |  2 +-
 .../IteratorToStringHandle.pbtxt              |  2 +-
 .../compat/ops_history_v2/IteratorV2.pbtxt    |  2 +-
 .../KMC2ChainInitialization.pbtxt             |  2 +-
 .../KmeansPlusPlusInitialization.pbtxt        |  2 +-
 .../ops_history_v2/KthOrderStatistic.pbtxt    |  2 +-
 .../ops/compat/ops_history_v2/L2Loss.pbtxt    |  2 +-
 .../compat/ops_history_v2/LMDBDataset.pbtxt   |  2 +-
 .../compat/ops_history_v2/LMDBReader.pbtxt    |  2 +-
 .../core/ops/compat/ops_history_v2/LRN.pbtxt  |  2 +-
 .../ops/compat/ops_history_v2/LRNGrad.pbtxt   |  2 +-
 .../compat/ops_history_v2/LSTMBlockCell.pbtxt |  2 +-
 .../ops_history_v2/LSTMBlockCellGrad.pbtxt    |  2 +-
 .../ops_history_v2/LatencyStatsDataset.pbtxt  |  2 +-
 .../ops/compat/ops_history_v2/LeakyRelu.pbtxt |  2 +-
 .../compat/ops_history_v2/LeakyReluGrad.pbtxt |  2 +-
 .../LearnedUnigramCandidateSampler.pbtxt      |  2 +-
 .../ops/compat/ops_history_v2/LeftShift.pbtxt |  2 +-
 .../LegacyParallelInterleaveDatasetV2.pbtxt   |  2 +-
 .../core/ops/compat/ops_history_v2/Less.pbtxt |  2 +-
 .../ops/compat/ops_history_v2/LessEqual.pbtxt |  2 +-
 .../ops/compat/ops_history_v2/Lgamma.pbtxt    |  2 +-
 .../ops/compat/ops_history_v2/LinSpace.pbtxt  |  2 +-
 .../compat/ops_history_v2/ListDataset.pbtxt   |  2 +-
 .../ops/compat/ops_history_v2/ListDiff.pbtxt  |  2 +-
 .../ListSnapshotChunksDataset.pbtxt           |  2 +-
 .../LoadAllTPUEmbeddingParameters.pbtxt       |  2 +-
 .../ops_history_v2/LoadAndRemapMatrix.pbtxt   |  2 +-
 .../compat/ops_history_v2/LoadDataset.pbtxt   |  2 +-
 .../LoadTPUEmbeddingADAMParameters.pbtxt      |  2 +-
 .../LoadTPUEmbeddingAdadeltaParameters.pbtxt  |  2 +-
 ...PUEmbeddingAdagradMomentumParameters.pbtxt |  2 +-
 .../LoadTPUEmbeddingAdagradParameters.pbtxt   |  2 +-
 ...PUEmbeddingCenteredRMSPropParameters.pbtxt |  2 +-
 .../LoadTPUEmbeddingFTRLParameters.pbtxt      |  2 +-
 ...mbeddingFrequencyEstimatorParameters.pbtxt |  2 +-
 ...PUEmbeddingMDLAdagradLightParameters.pbtxt |  2 +-
 .../LoadTPUEmbeddingMomentumParameters.pbtxt  |  2 +-
 ...PUEmbeddingProximalAdagradParameters.pbtxt |  2 +-
 ...adTPUEmbeddingProximalYogiParameters.pbtxt |  2 +-
 .../LoadTPUEmbeddingRMSPropParameters.pbtxt   |  2 +-
 ...gStochasticGradientDescentParameters.pbtxt |  2 +-
 .../core/ops/compat/ops_history_v2/Log.pbtxt  |  2 +-
 .../ops/compat/ops_history_v2/Log1p.pbtxt     |  2 +-
 .../ops_history_v2/LogMatrixDeterminant.pbtxt |  2 +-
 .../compat/ops_history_v2/LogSoftmax.pbtxt    |  2 +-
 .../LogUniformCandidateSampler.pbtxt          |  2 +-
 .../compat/ops_history_v2/LogicalAnd.pbtxt    |  2 +-
 .../compat/ops_history_v2/LogicalNot.pbtxt    |  2 +-
 .../ops/compat/ops_history_v2/LogicalOr.pbtxt |  2 +-
 .../ops_history_v2/LookupTableExport.pbtxt    |  2 +-
 .../ops_history_v2/LookupTableExportV2.pbtxt  |  2 +-
 .../ops_history_v2/LookupTableFind.pbtxt      |  2 +-
 .../ops_history_v2/LookupTableFindV2.pbtxt    |  2 +-
 .../ops_history_v2/LookupTableImport.pbtxt    |  2 +-
 .../ops_history_v2/LookupTableImportV2.pbtxt  |  2 +-
 .../ops_history_v2/LookupTableInsert.pbtxt    |  2 +-
 .../ops_history_v2/LookupTableInsertV2.pbtxt  |  2 +-
 .../ops_history_v2/LookupTableRemoveV2.pbtxt  |  2 +-
 .../ops_history_v2/LookupTableSize.pbtxt      |  2 +-
 .../ops_history_v2/LookupTableSizeV2.pbtxt    |  2 +-
 .../ops/compat/ops_history_v2/LoopCond.pbtxt  |  2 +-
 .../compat/ops_history_v2/LowerBound.pbtxt    |  2 +-
 .../core/ops/compat/ops_history_v2/Lu.pbtxt   |  2 +-
 .../compat/ops_history_v2/MakeIterator.pbtxt  |  2 +-
 .../compat/ops_history_v2/MakeUnique.pbtxt    |  2 +-
 .../ops_history_v2/MapAndBatchDataset.pbtxt   |  2 +-
 .../ops/compat/ops_history_v2/MapClear.pbtxt  |  2 +-
 .../compat/ops_history_v2/MapDataset.pbtxt    |  2 +-
 .../ops/compat/ops_history_v2/MapDefun.pbtxt  |  2 +-
 .../ops_history_v2/MapIncompleteSize.pbtxt    |  2 +-
 .../ops/compat/ops_history_v2/MapPeek.pbtxt   |  2 +-
 .../ops/compat/ops_history_v2/MapSize.pbtxt   |  2 +-
 .../ops/compat/ops_history_v2/MapStage.pbtxt  |  2 +-
 .../compat/ops_history_v2/MapUnstage.pbtxt    |  2 +-
 .../ops_history_v2/MapUnstageNoKey.pbtxt      |  2 +-
 .../ops/compat/ops_history_v2/MatMul.pbtxt    |  2 +-
 .../compat/ops_history_v2/MatchingFiles.pbtxt |  2 +-
 .../ops_history_v2/MatchingFilesDataset.pbtxt |  2 +-
 .../ops_history_v2/MatrixBandPart.pbtxt       |  2 +-
 .../ops_history_v2/MatrixDeterminant.pbtxt    |  2 +-
 .../compat/ops_history_v2/MatrixDiag.pbtxt    |  2 +-
 .../ops_history_v2/MatrixDiagPart.pbtxt       |  2 +-
 .../ops_history_v2/MatrixDiagPartV2.pbtxt     |  2 +-
 .../ops_history_v2/MatrixDiagPartV3.pbtxt     |  2 +-
 .../compat/ops_history_v2/MatrixDiagV2.pbtxt  |  2 +-
 .../compat/ops_history_v2/MatrixDiagV3.pbtxt  |  2 +-
 .../ops_history_v2/MatrixExponential.pbtxt    |  2 +-
 .../compat/ops_history_v2/MatrixInverse.pbtxt |  2 +-
 .../ops_history_v2/MatrixLogarithm.pbtxt      |  2 +-
 .../compat/ops_history_v2/MatrixSetDiag.pbtxt |  2 +-
 .../ops_history_v2/MatrixSetDiagV2.pbtxt      |  2 +-
 .../ops_history_v2/MatrixSetDiagV3.pbtxt      |  2 +-
 .../compat/ops_history_v2/MatrixSolve.pbtxt   |  2 +-
 .../compat/ops_history_v2/MatrixSolveLs.pbtxt |  2 +-
 .../ops_history_v2/MatrixSquareRoot.pbtxt     |  2 +-
 .../MatrixTriangularSolve.pbtxt               |  2 +-
 .../core/ops/compat/ops_history_v2/Max.pbtxt  |  2 +-
 .../MaxIntraOpParallelismDataset.pbtxt        |  2 +-
 .../ops/compat/ops_history_v2/MaxPool.pbtxt   |  2 +-
 .../ops/compat/ops_history_v2/MaxPool3D.pbtxt |  2 +-
 .../compat/ops_history_v2/MaxPool3DGrad.pbtxt |  2 +-
 .../ops_history_v2/MaxPool3DGradGrad.pbtxt    |  2 +-
 .../compat/ops_history_v2/MaxPoolGrad.pbtxt   |  2 +-
 .../ops_history_v2/MaxPoolGradGrad.pbtxt      |  2 +-
 .../ops_history_v2/MaxPoolGradGradV2.pbtxt    |  2 +-
 .../MaxPoolGradGradWithArgmax.pbtxt           |  2 +-
 .../compat/ops_history_v2/MaxPoolGradV2.pbtxt |  2 +-
 .../MaxPoolGradWithArgmax.pbtxt               |  2 +-
 .../ops/compat/ops_history_v2/MaxPoolV2.pbtxt |  2 +-
 .../ops_history_v2/MaxPoolWithArgmax.pbtxt    |  2 +-
 .../ops/compat/ops_history_v2/Maximum.pbtxt   |  2 +-
 .../core/ops/compat/ops_history_v2/Mean.pbtxt |  2 +-
 .../ops/compat/ops_history_v2/Merge.pbtxt     |  2 +-
 .../ops_history_v2/MergeDedupData.pbtxt       |  2 +-
 .../compat/ops_history_v2/MergeSummary.pbtxt  |  2 +-
 .../ops_history_v2/MergeV2Checkpoints.pbtxt   |  2 +-
 .../core/ops/compat/ops_history_v2/Mfcc.pbtxt |  2 +-
 .../core/ops/compat/ops_history_v2/Min.pbtxt  |  2 +-
 .../ops/compat/ops_history_v2/Minimum.pbtxt   |  2 +-
 .../ops/compat/ops_history_v2/MirrorPad.pbtxt |  2 +-
 .../compat/ops_history_v2/MirrorPadGrad.pbtxt |  2 +-
 .../ops_history_v2/MlirPassthroughOp.pbtxt    |  2 +-
 .../core/ops/compat/ops_history_v2/Mod.pbtxt  |  2 +-
 .../compat/ops_history_v2/ModelDataset.pbtxt  |  2 +-
 .../core/ops/compat/ops_history_v2/Mul.pbtxt  |  2 +-
 .../ops/compat/ops_history_v2/MulNoNan.pbtxt  |  2 +-
 .../ops_history_v2/MultiDeviceIterator.pbtxt  |  2 +-
 .../MultiDeviceIteratorFromStringHandle.pbtxt |  2 +-
 .../MultiDeviceIteratorGetNextFromShard.pbtxt |  2 +-
 .../MultiDeviceIteratorInit.pbtxt             |  2 +-
 .../MultiDeviceIteratorToStringHandle.pbtxt   |  2 +-
 .../compat/ops_history_v2/Multinomial.pbtxt   |  2 +-
 .../MutableDenseHashTable.pbtxt               |  2 +-
 .../MutableDenseHashTableV2.pbtxt             |  2 +-
 .../ops_history_v2/MutableHashTable.pbtxt     |  2 +-
 .../MutableHashTableOfTensors.pbtxt           |  2 +-
 .../MutableHashTableOfTensorsV2.pbtxt         |  2 +-
 .../ops_history_v2/MutableHashTableV2.pbtxt   |  2 +-
 .../ops/compat/ops_history_v2/MutexLock.pbtxt |  2 +-
 .../ops/compat/ops_history_v2/MutexV2.pbtxt   |  2 +-
 .../compat/ops_history_v2/NcclAllReduce.pbtxt |  2 +-
 .../compat/ops_history_v2/NcclBroadcast.pbtxt |  2 +-
 .../compat/ops_history_v2/NcclReduce.pbtxt    |  2 +-
 .../ops/compat/ops_history_v2/Ndtri.pbtxt     |  2 +-
 .../ops_history_v2/NearestNeighbors.pbtxt     |  2 +-
 .../core/ops/compat/ops_history_v2/Neg.pbtxt  |  2 +-
 .../ops/compat/ops_history_v2/NegTrain.pbtxt  |  2 +-
 .../ops/compat/ops_history_v2/NextAfter.pbtxt |  2 +-
 .../compat/ops_history_v2/NextIteration.pbtxt |  2 +-
 .../core/ops/compat/ops_history_v2/NoOp.pbtxt |  2 +-
 .../ops_history_v2/NonDeterministicInts.pbtxt |  2 +-
 .../ops_history_v2/NonMaxSuppression.pbtxt    |  2 +-
 .../ops_history_v2/NonMaxSuppressionV2.pbtxt  |  2 +-
 .../ops_history_v2/NonMaxSuppressionV3.pbtxt  |  2 +-
 .../ops_history_v2/NonMaxSuppressionV4.pbtxt  |  2 +-
 .../ops_history_v2/NonMaxSuppressionV5.pbtxt  |  2 +-
 .../NonMaxSuppressionWithOverlaps.pbtxt       |  2 +-
 .../NonSerializableDataset.pbtxt              |  2 +-
 .../ops/compat/ops_history_v2/NotEqual.pbtxt  |  2 +-
 .../compat/ops_history_v2/NthElement.pbtxt    |  2 +-
 .../ops/compat/ops_history_v2/OneHot.pbtxt    |  2 +-
 .../ops_history_v2/OneShotIterator.pbtxt      |  2 +-
 .../ops/compat/ops_history_v2/OnesLike.pbtxt  |  2 +-
 .../ops_history_v2/OptimizeDataset.pbtxt      |  2 +-
 .../ops_history_v2/OptimizeDatasetV2.pbtxt    |  2 +-
 .../ops_history_v2/OptionalFromValue.pbtxt    |  2 +-
 .../ops_history_v2/OptionalGetValue.pbtxt     |  2 +-
 .../ops_history_v2/OptionalHasValue.pbtxt     |  2 +-
 .../compat/ops_history_v2/OptionalNone.pbtxt  |  2 +-
 .../ops_history_v2/OptionsDataset.pbtxt       |  2 +-
 .../ops_history_v2/OrderedMapClear.pbtxt      |  2 +-
 .../OrderedMapIncompleteSize.pbtxt            |  2 +-
 .../ops_history_v2/OrderedMapPeek.pbtxt       |  2 +-
 .../ops_history_v2/OrderedMapSize.pbtxt       |  2 +-
 .../ops_history_v2/OrderedMapStage.pbtxt      |  2 +-
 .../ops_history_v2/OrderedMapUnstage.pbtxt    |  2 +-
 .../OrderedMapUnstageNoKey.pbtxt              |  2 +-
 .../ops_history_v2/OutfeedDequeue.pbtxt       |  2 +-
 .../ops_history_v2/OutfeedDequeueTuple.pbtxt  |  2 +-
 .../OutfeedDequeueTupleV2.pbtxt               |  2 +-
 .../ops_history_v2/OutfeedDequeueV2.pbtxt     |  2 +-
 .../ops_history_v2/OutfeedEnqueue.pbtxt       |  2 +-
 .../ops_history_v2/OutfeedEnqueueTuple.pbtxt  |  2 +-
 .../core/ops/compat/ops_history_v2/Pack.pbtxt |  2 +-
 .../core/ops/compat/ops_history_v2/Pad.pbtxt  |  2 +-
 .../ops/compat/ops_history_v2/PadV2.pbtxt     |  2 +-
 .../ops_history_v2/PaddedBatchDataset.pbtxt   |  2 +-
 .../ops_history_v2/PaddedBatchDatasetV2.pbtxt |  2 +-
 .../ops_history_v2/PaddingFIFOQueue.pbtxt     |  2 +-
 .../ops_history_v2/PaddingFIFOQueueV2.pbtxt   |  2 +-
 .../ops_history_v2/ParallelBatchDataset.pbtxt |  2 +-
 .../ops_history_v2/ParallelConcat.pbtxt       |  2 +-
 .../ParallelDynamicStitch.pbtxt               |  2 +-
 .../ParallelFilterDataset.pbtxt               |  2 +-
 .../ParallelInterleaveDataset.pbtxt           |  2 +-
 .../ParallelInterleaveDatasetV2.pbtxt         |  2 +-
 .../ParallelInterleaveDatasetV3.pbtxt         |  2 +-
 .../ParallelInterleaveDatasetV4.pbtxt         |  2 +-
 .../ops_history_v2/ParallelMapDataset.pbtxt   |  2 +-
 .../ops_history_v2/ParallelMapDatasetV2.pbtxt |  2 +-
 .../ParameterizedTruncatedNormal.pbtxt        |  2 +-
 .../compat/ops_history_v2/ParseExample.pbtxt  |  2 +-
 .../ops_history_v2/ParseExampleDataset.pbtxt  |  2 +-
 .../ParseExampleDatasetV2.pbtxt               |  2 +-
 .../ops_history_v2/ParseExampleV2.pbtxt       |  2 +-
 .../ops_history_v2/ParseSequenceExample.pbtxt |  2 +-
 .../ParseSequenceExampleV2.pbtxt              |  2 +-
 .../ops_history_v2/ParseSingleExample.pbtxt   |  2 +-
 .../ParseSingleSequenceExample.pbtxt          |  2 +-
 .../compat/ops_history_v2/ParseTensor.pbtxt   |  2 +-
 .../ops_history_v2/PartitionedCall.pbtxt      |  2 +-
 .../compat/ops_history_v2/Placeholder.pbtxt   |  2 +-
 .../compat/ops_history_v2/PlaceholderV2.pbtxt |  2 +-
 .../PlaceholderWithDefault.pbtxt              |  2 +-
 .../ops/compat/ops_history_v2/Polygamma.pbtxt |  2 +-
 .../ops_history_v2/PopulationCount.pbtxt      |  2 +-
 .../core/ops/compat/ops_history_v2/Pow.pbtxt  |  2 +-
 .../ops_history_v2/PrefetchDataset.pbtxt      |  2 +-
 .../compat/ops_history_v2/Prelinearize.pbtxt  |  2 +-
 .../ops_history_v2/PrelinearizeTuple.pbtxt    |  2 +-
 .../ops_history_v2/PreventGradient.pbtxt      |  2 +-
 .../ops/compat/ops_history_v2/Print.pbtxt     |  2 +-
 .../ops/compat/ops_history_v2/PrintV2.pbtxt   |  2 +-
 .../compat/ops_history_v2/PriorityQueue.pbtxt |  2 +-
 .../ops_history_v2/PriorityQueueV2.pbtxt      |  2 +-
 .../PrivateThreadPoolDataset.pbtxt            |  2 +-
 .../core/ops/compat/ops_history_v2/Prod.pbtxt |  2 +-
 .../ops/compat/ops_history_v2/PyFunc.pbtxt    |  2 +-
 .../ops_history_v2/PyFuncStateless.pbtxt      |  2 +-
 .../core/ops/compat/ops_history_v2/Qr.pbtxt   |  2 +-
 .../QuantizeAndDequantize.pbtxt               |  2 +-
 .../QuantizeAndDequantizeV2.pbtxt             |  2 +-
 .../QuantizeAndDequantizeV3.pbtxt             |  2 +-
 .../QuantizeAndDequantizeV4.pbtxt             |  2 +-
 .../QuantizeAndDequantizeV4Grad.pbtxt         |  2 +-
 .../QuantizeDownAndShrinkRange.pbtxt          |  2 +-
 .../compat/ops_history_v2/QuantizeV2.pbtxt    |  2 +-
 .../compat/ops_history_v2/QuantizedAdd.pbtxt  |  2 +-
 .../ops_history_v2/QuantizedAvgPool.pbtxt     |  2 +-
 ...izedBatchNormWithGlobalNormalization.pbtxt |  2 +-
 .../ops_history_v2/QuantizedBiasAdd.pbtxt     |  2 +-
 .../ops_history_v2/QuantizedConcat.pbtxt      |  2 +-
 .../ops_history_v2/QuantizedConv2D.pbtxt      |  2 +-
 .../QuantizedConv2DAndRelu.pbtxt              |  2 +-
 .../QuantizedConv2DAndReluAndRequantize.pbtxt |  2 +-
 .../QuantizedConv2DAndRequantize.pbtxt        |  2 +-
 .../QuantizedConv2DPerChannel.pbtxt           |  2 +-
 .../QuantizedConv2DWithBias.pbtxt             |  2 +-
 .../QuantizedConv2DWithBiasAndRelu.pbtxt      |  2 +-
 ...edConv2DWithBiasAndReluAndRequantize.pbtxt |  2 +-
 ...QuantizedConv2DWithBiasAndRequantize.pbtxt |  2 +-
 ...ithBiasSignedSumAndReluAndRequantize.pbtxt |  2 +-
 .../QuantizedConv2DWithBiasSumAndRelu.pbtxt   |  2 +-
 ...onv2DWithBiasSumAndReluAndRequantize.pbtxt |  2 +-
 .../QuantizedDepthwiseConv2D.pbtxt            |  2 +-
 .../QuantizedDepthwiseConv2DWithBias.pbtxt    |  2 +-
 ...ntizedDepthwiseConv2DWithBiasAndRelu.pbtxt |  2 +-
 ...seConv2DWithBiasAndReluAndRequantize.pbtxt |  2 +-
 .../QuantizedInstanceNorm.pbtxt               |  2 +-
 .../ops_history_v2/QuantizedMatMul.pbtxt      |  2 +-
 .../QuantizedMatMulWithBias.pbtxt             |  2 +-
 ...QuantizedMatMulWithBiasAndDequantize.pbtxt |  2 +-
 .../QuantizedMatMulWithBiasAndRelu.pbtxt      |  2 +-
 ...edMatMulWithBiasAndReluAndRequantize.pbtxt |  2 +-
 ...QuantizedMatMulWithBiasAndRequantize.pbtxt |  2 +-
 .../ops_history_v2/QuantizedMaxPool.pbtxt     |  2 +-
 .../compat/ops_history_v2/QuantizedMul.pbtxt  |  2 +-
 .../compat/ops_history_v2/QuantizedRelu.pbtxt |  2 +-
 .../ops_history_v2/QuantizedRelu6.pbtxt       |  2 +-
 .../ops_history_v2/QuantizedReluX.pbtxt       |  2 +-
 .../ops_history_v2/QuantizedReshape.pbtxt     |  2 +-
 .../QuantizedResizeBilinear.pbtxt             |  2 +-
 .../compat/ops_history_v2/QueueClose.pbtxt    |  2 +-
 .../compat/ops_history_v2/QueueCloseV2.pbtxt  |  2 +-
 .../compat/ops_history_v2/QueueDequeue.pbtxt  |  2 +-
 .../ops_history_v2/QueueDequeueMany.pbtxt     |  2 +-
 .../ops_history_v2/QueueDequeueManyV2.pbtxt   |  2 +-
 .../ops_history_v2/QueueDequeueUpTo.pbtxt     |  2 +-
 .../ops_history_v2/QueueDequeueUpToV2.pbtxt   |  2 +-
 .../ops_history_v2/QueueDequeueV2.pbtxt       |  2 +-
 .../compat/ops_history_v2/QueueEnqueue.pbtxt  |  2 +-
 .../ops_history_v2/QueueEnqueueMany.pbtxt     |  2 +-
 .../ops_history_v2/QueueEnqueueManyV2.pbtxt   |  2 +-
 .../ops_history_v2/QueueEnqueueV2.pbtxt       |  2 +-
 .../compat/ops_history_v2/QueueIsClosed.pbtxt |  2 +-
 .../ops_history_v2/QueueIsClosedV2.pbtxt      |  2 +-
 .../ops/compat/ops_history_v2/QueueSize.pbtxt |  2 +-
 .../compat/ops_history_v2/QueueSizeV2.pbtxt   |  2 +-
 .../core/ops/compat/ops_history_v2/RFFT.pbtxt |  2 +-
 .../ops/compat/ops_history_v2/RFFT2D.pbtxt    |  2 +-
 .../ops/compat/ops_history_v2/RFFT3D.pbtxt    |  2 +-
 .../ops/compat/ops_history_v2/RFFTND.pbtxt    |  2 +-
 .../ops/compat/ops_history_v2/RGBToHSV.pbtxt  |  2 +-
 .../ops_history_v2/RaggedBincount.pbtxt       |  2 +-
 .../RaggedCountSparseOutput.pbtxt             |  2 +-
 .../compat/ops_history_v2/RaggedCross.pbtxt   |  2 +-
 .../ops_history_v2/RaggedFillEmptyRows.pbtxt  |  2 +-
 .../RaggedFillEmptyRowsGrad.pbtxt             |  2 +-
 .../compat/ops_history_v2/RaggedGather.pbtxt  |  2 +-
 .../compat/ops_history_v2/RaggedRange.pbtxt   |  2 +-
 .../RaggedTensorFromVariant.pbtxt             |  2 +-
 .../ops_history_v2/RaggedTensorToSparse.pbtxt |  2 +-
 .../ops_history_v2/RaggedTensorToTensor.pbtxt |  2 +-
 .../RaggedTensorToVariant.pbtxt               |  2 +-
 .../RaggedTensorToVariantGradient.pbtxt       |  2 +-
 .../compat/ops_history_v2/RandomCrop.pbtxt    |  2 +-
 .../compat/ops_history_v2/RandomDataset.pbtxt |  2 +-
 .../ops_history_v2/RandomDatasetV2.pbtxt      |  2 +-
 .../compat/ops_history_v2/RandomGamma.pbtxt   |  2 +-
 .../ops_history_v2/RandomGammaGrad.pbtxt      |  2 +-
 .../ops_history_v2/RandomIndexShuffle.pbtxt   |  2 +-
 .../compat/ops_history_v2/RandomPoisson.pbtxt |  2 +-
 .../ops_history_v2/RandomPoissonV2.pbtxt      |  2 +-
 .../compat/ops_history_v2/RandomShuffle.pbtxt |  2 +-
 .../ops_history_v2/RandomShuffleQueue.pbtxt   |  2 +-
 .../ops_history_v2/RandomShuffleQueueV2.pbtxt |  2 +-
 .../ops_history_v2/RandomStandardNormal.pbtxt |  2 +-
 .../compat/ops_history_v2/RandomUniform.pbtxt |  2 +-
 .../ops_history_v2/RandomUniformInt.pbtxt     |  2 +-
 .../ops/compat/ops_history_v2/Range.pbtxt     |  2 +-
 .../compat/ops_history_v2/RangeDataset.pbtxt  |  2 +-
 .../core/ops/compat/ops_history_v2/Rank.pbtxt |  2 +-
 .../ops/compat/ops_history_v2/ReadFile.pbtxt  |  2 +-
 .../ops_history_v2/ReadVariableOp.pbtxt       |  2 +-
 .../ReadVariableXlaSplitND.pbtxt              |  2 +-
 .../ReaderNumRecordsProduced.pbtxt            |  2 +-
 .../ReaderNumRecordsProducedV2.pbtxt          |  2 +-
 .../ReaderNumWorkUnitsCompleted.pbtxt         |  2 +-
 .../ReaderNumWorkUnitsCompletedV2.pbtxt       |  2 +-
 .../compat/ops_history_v2/ReaderRead.pbtxt    |  2 +-
 .../ops_history_v2/ReaderReadUpTo.pbtxt       |  2 +-
 .../ops_history_v2/ReaderReadUpToV2.pbtxt     |  2 +-
 .../compat/ops_history_v2/ReaderReadV2.pbtxt  |  2 +-
 .../compat/ops_history_v2/ReaderReset.pbtxt   |  2 +-
 .../compat/ops_history_v2/ReaderResetV2.pbtxt |  2 +-
 .../ops_history_v2/ReaderRestoreState.pbtxt   |  2 +-
 .../ops_history_v2/ReaderRestoreStateV2.pbtxt |  2 +-
 .../ops_history_v2/ReaderSerializeState.pbtxt |  2 +-
 .../ReaderSerializeStateV2.pbtxt              |  2 +-
 .../core/ops/compat/ops_history_v2/Real.pbtxt |  2 +-
 .../ops/compat/ops_history_v2/RealDiv.pbtxt   |  2 +-
 .../ops_history_v2/RebatchDataset.pbtxt       |  2 +-
 .../ops_history_v2/RebatchDatasetV2.pbtxt     |  2 +-
 .../compat/ops_history_v2/Reciprocal.pbtxt    |  2 +-
 .../ops_history_v2/ReciprocalGrad.pbtxt       |  2 +-
 .../compat/ops_history_v2/RecordInput.pbtxt   |  2 +-
 .../core/ops/compat/ops_history_v2/Recv.pbtxt |  2 +-
 .../RecvTPUEmbeddingActivations.pbtxt         |  2 +-
 .../compat/ops_history_v2/ReduceDataset.pbtxt |  2 +-
 .../compat/ops_history_v2/ReduceJoin.pbtxt    |  2 +-
 .../ops/compat/ops_history_v2/RefEnter.pbtxt  |  2 +-
 .../ops/compat/ops_history_v2/RefExit.pbtxt   |  2 +-
 .../compat/ops_history_v2/RefIdentity.pbtxt   |  2 +-
 .../ops/compat/ops_history_v2/RefMerge.pbtxt  |  2 +-
 .../ops_history_v2/RefNextIteration.pbtxt     |  2 +-
 .../ops/compat/ops_history_v2/RefSelect.pbtxt |  2 +-
 .../ops/compat/ops_history_v2/RefSwitch.pbtxt |  2 +-
 .../ops_history_v2/RegexFullMatch.pbtxt       |  2 +-
 .../compat/ops_history_v2/RegexReplace.pbtxt  |  2 +-
 .../ops_history_v2/RegisterDataset.pbtxt      |  2 +-
 .../ops_history_v2/RegisterDatasetV2.pbtxt    |  2 +-
 .../ops/compat/ops_history_v2/Relayout.pbtxt  |  2 +-
 .../compat/ops_history_v2/RelayoutLike.pbtxt  |  2 +-
 .../core/ops/compat/ops_history_v2/Relu.pbtxt |  2 +-
 .../ops/compat/ops_history_v2/Relu6.pbtxt     |  2 +-
 .../ops/compat/ops_history_v2/Relu6Grad.pbtxt |  2 +-
 .../ops/compat/ops_history_v2/ReluGrad.pbtxt  |  2 +-
 .../compat/ops_history_v2/RemoteCall.pbtxt    |  2 +-
 .../compat/ops_history_v2/RepeatDataset.pbtxt |  2 +-
 .../ops_history_v2/RequantizationRange.pbtxt  |  2 +-
 .../RequantizationRangePerChannel.pbtxt       |  2 +-
 .../compat/ops_history_v2/Requantize.pbtxt    |  2 +-
 .../ops_history_v2/RequantizePerChannel.pbtxt |  2 +-
 .../ops/compat/ops_history_v2/Reshape.pbtxt   |  2 +-
 .../compat/ops_history_v2/ResizeArea.pbtxt    |  2 +-
 .../compat/ops_history_v2/ResizeBicubic.pbtxt |  2 +-
 .../ops_history_v2/ResizeBicubicGrad.pbtxt    |  2 +-
 .../ops_history_v2/ResizeBilinear.pbtxt       |  2 +-
 .../ops_history_v2/ResizeBilinearGrad.pbtxt   |  2 +-
 .../ResizeNearestNeighbor.pbtxt               |  2 +-
 .../ResizeNearestNeighborGrad.pbtxt           |  2 +-
 .../ResourceAccumulatorApplyGradient.pbtxt    |  2 +-
 .../ResourceAccumulatorNumAccumulated.pbtxt   |  2 +-
 .../ResourceAccumulatorSetGlobalStep.pbtxt    |  2 +-
 .../ResourceAccumulatorTakeGradient.pbtxt     |  2 +-
 .../ops_history_v2/ResourceApplyAdaMax.pbtxt  |  2 +-
 .../ResourceApplyAdadelta.pbtxt               |  2 +-
 .../ops_history_v2/ResourceApplyAdagrad.pbtxt |  2 +-
 .../ResourceApplyAdagradDA.pbtxt              |  2 +-
 .../ResourceApplyAdagradV2.pbtxt              |  2 +-
 .../ops_history_v2/ResourceApplyAdam.pbtxt    |  2 +-
 .../ResourceApplyAdamWithAmsgrad.pbtxt        |  2 +-
 .../ops_history_v2/ResourceApplyAddSign.pbtxt |  2 +-
 .../ResourceApplyCenteredRMSProp.pbtxt        |  2 +-
 .../ops_history_v2/ResourceApplyFtrl.pbtxt    |  2 +-
 .../ops_history_v2/ResourceApplyFtrlV2.pbtxt  |  2 +-
 .../ResourceApplyGradientDescent.pbtxt        |  2 +-
 .../ResourceApplyKerasMomentum.pbtxt          |  2 +-
 .../ResourceApplyMomentum.pbtxt               |  2 +-
 .../ResourceApplyPowerSign.pbtxt              |  2 +-
 .../ResourceApplyProximalAdagrad.pbtxt        |  2 +-
 ...ResourceApplyProximalGradientDescent.pbtxt |  2 +-
 .../ops_history_v2/ResourceApplyRMSProp.pbtxt |  2 +-
 .../ResourceConditionalAccumulator.pbtxt      |  2 +-
 .../ops_history_v2/ResourceCountUpTo.pbtxt    |  2 +-
 .../ops_history_v2/ResourceGather.pbtxt       |  2 +-
 .../ops_history_v2/ResourceGatherNd.pbtxt     |  2 +-
 .../ops_history_v2/ResourceScatterAdd.pbtxt   |  2 +-
 .../ops_history_v2/ResourceScatterDiv.pbtxt   |  2 +-
 .../ops_history_v2/ResourceScatterMax.pbtxt   |  2 +-
 .../ops_history_v2/ResourceScatterMin.pbtxt   |  2 +-
 .../ops_history_v2/ResourceScatterMul.pbtxt   |  2 +-
 .../ops_history_v2/ResourceScatterNdAdd.pbtxt |  2 +-
 .../ops_history_v2/ResourceScatterNdMax.pbtxt |  2 +-
 .../ops_history_v2/ResourceScatterNdMin.pbtxt |  2 +-
 .../ops_history_v2/ResourceScatterNdSub.pbtxt |  2 +-
 .../ResourceScatterNdUpdate.pbtxt             |  2 +-
 .../ops_history_v2/ResourceScatterSub.pbtxt   |  2 +-
 .../ResourceScatterUpdate.pbtxt               |  2 +-
 .../ResourceSparseApplyAdadelta.pbtxt         |  2 +-
 .../ResourceSparseApplyAdagrad.pbtxt          |  2 +-
 .../ResourceSparseApplyAdagradDA.pbtxt        |  2 +-
 .../ResourceSparseApplyAdagradV2.pbtxt        |  2 +-
 .../ResourceSparseApplyCenteredRMSProp.pbtxt  |  2 +-
 .../ResourceSparseApplyFtrl.pbtxt             |  2 +-
 .../ResourceSparseApplyFtrlV2.pbtxt           |  2 +-
 .../ResourceSparseApplyKerasMomentum.pbtxt    |  2 +-
 .../ResourceSparseApplyMomentum.pbtxt         |  2 +-
 .../ResourceSparseApplyProximalAdagrad.pbtxt  |  2 +-
 ...ceSparseApplyProximalGradientDescent.pbtxt |  2 +-
 .../ResourceSparseApplyRMSProp.pbtxt          |  2 +-
 .../ResourceStridedSliceAssign.pbtxt          |  2 +-
 .../ops/compat/ops_history_v2/Restore.pbtxt   |  2 +-
 .../compat/ops_history_v2/RestoreSlice.pbtxt  |  2 +-
 .../ops/compat/ops_history_v2/RestoreV2.pbtxt |  2 +-
 .../RetrieveAllTPUEmbeddingParameters.pbtxt   |  2 +-
 .../RetrieveTPUEmbeddingADAMParameters.pbtxt  |  2 +-
 ...trieveTPUEmbeddingAdadeltaParameters.pbtxt |  2 +-
 ...PUEmbeddingAdagradMomentumParameters.pbtxt |  2 +-
 ...etrieveTPUEmbeddingAdagradParameters.pbtxt |  2 +-
 ...PUEmbeddingCenteredRMSPropParameters.pbtxt |  2 +-
 .../RetrieveTPUEmbeddingFTRLParameters.pbtxt  |  2 +-
 ...mbeddingFrequencyEstimatorParameters.pbtxt |  2 +-
 ...PUEmbeddingMDLAdagradLightParameters.pbtxt |  2 +-
 ...trieveTPUEmbeddingMomentumParameters.pbtxt |  2 +-
 ...PUEmbeddingProximalAdagradParameters.pbtxt |  2 +-
 ...veTPUEmbeddingProximalYogiParameters.pbtxt |  2 +-
 ...etrieveTPUEmbeddingRMSPropParameters.pbtxt |  2 +-
 ...gStochasticGradientDescentParameters.pbtxt |  2 +-
 .../ops/compat/ops_history_v2/Reverse.pbtxt   |  2 +-
 .../ops_history_v2/ReverseSequence.pbtxt      |  2 +-
 .../ops/compat/ops_history_v2/ReverseV2.pbtxt |  2 +-
 .../ops_history_v2/RewriteDataset.pbtxt       |  2 +-
 .../compat/ops_history_v2/RightShift.pbtxt    |  2 +-
 .../core/ops/compat/ops_history_v2/Rint.pbtxt |  2 +-
 .../ops/compat/ops_history_v2/RiscAbs.pbtxt   |  2 +-
 .../ops/compat/ops_history_v2/RiscAdd.pbtxt   |  2 +-
 .../ops_history_v2/RiscBinaryArithmetic.pbtxt |  2 +-
 .../ops_history_v2/RiscBinaryComparison.pbtxt |  2 +-
 .../compat/ops_history_v2/RiscBitcast.pbtxt   |  2 +-
 .../compat/ops_history_v2/RiscBroadcast.pbtxt |  2 +-
 .../ops/compat/ops_history_v2/RiscCast.pbtxt  |  2 +-
 .../ops/compat/ops_history_v2/RiscCeil.pbtxt  |  2 +-
 .../compat/ops_history_v2/RiscCholesky.pbtxt  |  2 +-
 .../compat/ops_history_v2/RiscConcat.pbtxt    |  2 +-
 .../compat/ops_history_v2/RiscCondition.pbtxt |  2 +-
 .../ops/compat/ops_history_v2/RiscConv.pbtxt  |  2 +-
 .../ops/compat/ops_history_v2/RiscCos.pbtxt   |  2 +-
 .../ops/compat/ops_history_v2/RiscDiv.pbtxt   |  2 +-
 .../ops/compat/ops_history_v2/RiscDot.pbtxt   |  2 +-
 .../ops/compat/ops_history_v2/RiscExp.pbtxt   |  2 +-
 .../ops/compat/ops_history_v2/RiscFft.pbtxt   |  2 +-
 .../ops/compat/ops_history_v2/RiscFloor.pbtxt |  2 +-
 .../compat/ops_history_v2/RiscGather.pbtxt    |  2 +-
 .../ops/compat/ops_history_v2/RiscImag.pbtxt  |  2 +-
 .../compat/ops_history_v2/RiscIsFinite.pbtxt  |  2 +-
 .../ops/compat/ops_history_v2/RiscLog.pbtxt   |  2 +-
 .../ops_history_v2/RiscLogicalAnd.pbtxt       |  2 +-
 .../ops_history_v2/RiscLogicalNot.pbtxt       |  2 +-
 .../compat/ops_history_v2/RiscLogicalOr.pbtxt |  2 +-
 .../ops/compat/ops_history_v2/RiscMax.pbtxt   |  2 +-
 .../ops/compat/ops_history_v2/RiscMin.pbtxt   |  2 +-
 .../ops/compat/ops_history_v2/RiscMul.pbtxt   |  2 +-
 .../ops/compat/ops_history_v2/RiscNeg.pbtxt   |  2 +-
 .../ops/compat/ops_history_v2/RiscPad.pbtxt   |  2 +-
 .../ops/compat/ops_history_v2/RiscPool.pbtxt  |  2 +-
 .../ops/compat/ops_history_v2/RiscPow.pbtxt   |  2 +-
 .../ops_history_v2/RiscRandomUniform.pbtxt    |  2 +-
 .../ops/compat/ops_history_v2/RiscReal.pbtxt  |  2 +-
 .../compat/ops_history_v2/RiscReduce.pbtxt    |  2 +-
 .../ops/compat/ops_history_v2/RiscRem.pbtxt   |  2 +-
 .../compat/ops_history_v2/RiscReshape.pbtxt   |  2 +-
 .../compat/ops_history_v2/RiscReverse.pbtxt   |  2 +-
 .../compat/ops_history_v2/RiscScatter.pbtxt   |  2 +-
 .../ops/compat/ops_history_v2/RiscShape.pbtxt |  2 +-
 .../ops/compat/ops_history_v2/RiscSign.pbtxt  |  2 +-
 .../ops/compat/ops_history_v2/RiscSlice.pbtxt |  2 +-
 .../ops/compat/ops_history_v2/RiscSort.pbtxt  |  2 +-
 .../compat/ops_history_v2/RiscSqueeze.pbtxt   |  2 +-
 .../ops/compat/ops_history_v2/RiscSub.pbtxt   |  2 +-
 .../compat/ops_history_v2/RiscTranspose.pbtxt |  2 +-
 .../ops_history_v2/RiscTriangularSolve.pbtxt  |  2 +-
 .../ops/compat/ops_history_v2/RiscUnary.pbtxt |  2 +-
 .../ops/compat/ops_history_v2/RiscWhile.pbtxt |  2 +-
 .../ops_history_v2/RngReadAndSkip.pbtxt       |  2 +-
 .../ops/compat/ops_history_v2/RngSkip.pbtxt   |  2 +-
 .../core/ops/compat/ops_history_v2/Roll.pbtxt |  2 +-
 .../ops/compat/ops_history_v2/Round.pbtxt     |  2 +-
 .../ops/compat/ops_history_v2/Rsqrt.pbtxt     |  2 +-
 .../ops/compat/ops_history_v2/RsqrtGrad.pbtxt |  2 +-
 .../SampleDistortedBoundingBox.pbtxt          |  2 +-
 .../SampleDistortedBoundingBoxV2.pbtxt        |  2 +-
 .../ops_history_v2/SamplingDataset.pbtxt      |  2 +-
 .../core/ops/compat/ops_history_v2/Save.pbtxt |  2 +-
 .../compat/ops_history_v2/SaveDataset.pbtxt   |  2 +-
 .../compat/ops_history_v2/SaveDatasetV2.pbtxt |  2 +-
 .../compat/ops_history_v2/SaveSlices.pbtxt    |  2 +-
 .../ops/compat/ops_history_v2/SaveV2.pbtxt    |  2 +-
 .../compat/ops_history_v2/ScalarSummary.pbtxt |  2 +-
 .../ops_history_v2/ScaleAndTranslate.pbtxt    |  2 +-
 .../ScaleAndTranslateGrad.pbtxt               |  2 +-
 .../compat/ops_history_v2/ScanDataset.pbtxt   |  2 +-
 .../compat/ops_history_v2/ScatterAdd.pbtxt    |  2 +-
 .../compat/ops_history_v2/ScatterDiv.pbtxt    |  2 +-
 .../compat/ops_history_v2/ScatterMax.pbtxt    |  2 +-
 .../compat/ops_history_v2/ScatterMin.pbtxt    |  2 +-
 .../compat/ops_history_v2/ScatterMul.pbtxt    |  2 +-
 .../ops/compat/ops_history_v2/ScatterNd.pbtxt |  2 +-
 .../compat/ops_history_v2/ScatterNdAdd.pbtxt  |  2 +-
 .../compat/ops_history_v2/ScatterNdMax.pbtxt  |  2 +-
 .../compat/ops_history_v2/ScatterNdMin.pbtxt  |  2 +-
 .../ScatterNdNonAliasingAdd.pbtxt             |  2 +-
 .../compat/ops_history_v2/ScatterNdSub.pbtxt  |  2 +-
 .../ops_history_v2/ScatterNdUpdate.pbtxt      |  2 +-
 .../compat/ops_history_v2/ScatterSub.pbtxt    |  2 +-
 .../compat/ops_history_v2/ScatterUpdate.pbtxt |  2 +-
 .../compat/ops_history_v2/SdcaFprint.pbtxt    |  2 +-
 .../compat/ops_history_v2/SdcaOptimizer.pbtxt |  2 +-
 .../ops_history_v2/SdcaOptimizerV2.pbtxt      |  2 +-
 .../compat/ops_history_v2/SdcaShrinkL1.pbtxt  |  2 +-
 .../compat/ops_history_v2/SegmentMax.pbtxt    |  2 +-
 .../compat/ops_history_v2/SegmentMaxV2.pbtxt  |  2 +-
 .../compat/ops_history_v2/SegmentMean.pbtxt   |  2 +-
 .../compat/ops_history_v2/SegmentMin.pbtxt    |  2 +-
 .../compat/ops_history_v2/SegmentMinV2.pbtxt  |  2 +-
 .../compat/ops_history_v2/SegmentProd.pbtxt   |  2 +-
 .../compat/ops_history_v2/SegmentProdV2.pbtxt |  2 +-
 .../compat/ops_history_v2/SegmentSum.pbtxt    |  2 +-
 .../compat/ops_history_v2/SegmentSumV2.pbtxt  |  2 +-
 .../ops/compat/ops_history_v2/Select.pbtxt    |  2 +-
 .../ops/compat/ops_history_v2/SelectV2.pbtxt  |  2 +-
 .../ops_history_v2/SelfAdjointEig.pbtxt       |  2 +-
 .../ops_history_v2/SelfAdjointEigV2.pbtxt     |  2 +-
 .../core/ops/compat/ops_history_v2/Selu.pbtxt |  2 +-
 .../ops/compat/ops_history_v2/SeluGrad.pbtxt  |  2 +-
 .../core/ops/compat/ops_history_v2/Send.pbtxt |  2 +-
 .../SendTPUEmbeddingGradients.pbtxt           |  2 +-
 .../ops_history_v2/SerializeIterator.pbtxt    |  2 +-
 .../ops_history_v2/SerializeManySparse.pbtxt  |  2 +-
 .../ops_history_v2/SerializeSparse.pbtxt      |  2 +-
 .../ops_history_v2/SerializeTensor.pbtxt      |  2 +-
 .../ops/compat/ops_history_v2/SetSize.pbtxt   |  2 +-
 .../SetStatsAggregatorDataset.pbtxt           |  2 +-
 .../ops/compat/ops_history_v2/Shape.pbtxt     |  2 +-
 .../ops/compat/ops_history_v2/ShapeN.pbtxt    |  2 +-
 .../compat/ops_history_v2/ShardDataset.pbtxt  |  2 +-
 .../ops_history_v2/ShardedFilename.pbtxt      |  2 +-
 .../ops_history_v2/ShardedFilespec.pbtxt      |  2 +-
 .../ShuffleAndRepeatDataset.pbtxt             |  2 +-
 .../ShuffleAndRepeatDatasetV2.pbtxt           |  2 +-
 .../ops_history_v2/ShuffleDataset.pbtxt       |  2 +-
 .../ops_history_v2/ShuffleDatasetV2.pbtxt     |  2 +-
 .../ops_history_v2/ShuffleDatasetV3.pbtxt     |  2 +-
 .../ShutdownDistributedTPU.pbtxt              |  2 +-
 .../ops_history_v2/ShutdownTPUSystem.pbtxt    |  2 +-
 .../ops/compat/ops_history_v2/Sigmoid.pbtxt   |  2 +-
 .../compat/ops_history_v2/SigmoidGrad.pbtxt   |  2 +-
 .../core/ops/compat/ops_history_v2/Sign.pbtxt |  2 +-
 .../core/ops/compat/ops_history_v2/Sin.pbtxt  |  2 +-
 .../core/ops/compat/ops_history_v2/Sinh.pbtxt |  2 +-
 .../core/ops/compat/ops_history_v2/Size.pbtxt |  2 +-
 .../compat/ops_history_v2/SkipDataset.pbtxt   |  2 +-
 .../ops/compat/ops_history_v2/Skipgram.pbtxt  |  2 +-
 .../compat/ops_history_v2/SleepDataset.pbtxt  |  2 +-
 .../ops/compat/ops_history_v2/Slice.pbtxt     |  2 +-
 .../ops_history_v2/SlidingWindowDataset.pbtxt |  2 +-
 .../ops/compat/ops_history_v2/Snapshot.pbtxt  |  2 +-
 .../ops_history_v2/SnapshotChunkDataset.pbtxt |  2 +-
 .../ops_history_v2/SnapshotDataset.pbtxt      |  2 +-
 .../SnapshotDatasetReader.pbtxt               |  2 +-
 .../ops_history_v2/SnapshotDatasetV2.pbtxt    |  2 +-
 .../SnapshotNestedDatasetReader.pbtxt         |  2 +-
 .../compat/ops_history_v2/SobolSample.pbtxt   |  2 +-
 .../ops/compat/ops_history_v2/Softmax.pbtxt   |  2 +-
 .../SoftmaxCrossEntropyWithLogits.pbtxt       |  2 +-
 .../ops/compat/ops_history_v2/Softplus.pbtxt  |  2 +-
 .../compat/ops_history_v2/SoftplusGrad.pbtxt  |  2 +-
 .../ops/compat/ops_history_v2/Softsign.pbtxt  |  2 +-
 .../compat/ops_history_v2/SoftsignGrad.pbtxt  |  2 +-
 .../SortListOfSparseCoreCooTensors.pbtxt      |  2 +-
 .../compat/ops_history_v2/SpaceToBatch.pbtxt  |  2 +-
 .../ops_history_v2/SpaceToBatchND.pbtxt       |  2 +-
 .../compat/ops_history_v2/SpaceToDepth.pbtxt  |  2 +-
 .../SparseAccumulatorApplyGradient.pbtxt      |  2 +-
 .../SparseAccumulatorTakeGradient.pbtxt       |  2 +-
 .../ops/compat/ops_history_v2/SparseAdd.pbtxt |  2 +-
 .../compat/ops_history_v2/SparseAddGrad.pbtxt |  2 +-
 .../ops_history_v2/SparseApplyAdadelta.pbtxt  |  2 +-
 .../ops_history_v2/SparseApplyAdagrad.pbtxt   |  2 +-
 .../ops_history_v2/SparseApplyAdagradDA.pbtxt |  2 +-
 .../ops_history_v2/SparseApplyAdagradV2.pbtxt |  2 +-
 .../SparseApplyCenteredRMSProp.pbtxt          |  2 +-
 .../ops_history_v2/SparseApplyFtrl.pbtxt      |  2 +-
 .../ops_history_v2/SparseApplyFtrlV2.pbtxt    |  2 +-
 .../ops_history_v2/SparseApplyMomentum.pbtxt  |  2 +-
 .../SparseApplyProximalAdagrad.pbtxt          |  2 +-
 .../SparseApplyProximalGradientDescent.pbtxt  |  2 +-
 .../ops_history_v2/SparseApplyRMSProp.pbtxt   |  2 +-
 .../ops_history_v2/SparseBincount.pbtxt       |  2 +-
 .../compat/ops_history_v2/SparseConcat.pbtxt  |  2 +-
 .../SparseConditionalAccumulator.pbtxt        |  2 +-
 .../SparseCountSparseOutput.pbtxt             |  2 +-
 .../compat/ops_history_v2/SparseCross.pbtxt   |  2 +-
 .../ops_history_v2/SparseCrossHashed.pbtxt    |  2 +-
 .../compat/ops_history_v2/SparseCrossV2.pbtxt |  2 +-
 .../ops_history_v2/SparseDenseCwiseAdd.pbtxt  |  2 +-
 .../ops_history_v2/SparseDenseCwiseDiv.pbtxt  |  2 +-
 .../ops_history_v2/SparseDenseCwiseMul.pbtxt  |  2 +-
 .../ops_history_v2/SparseFillEmptyRows.pbtxt  |  2 +-
 .../SparseFillEmptyRowsGrad.pbtxt             |  2 +-
 .../compat/ops_history_v2/SparseMatMul.pbtxt  |  2 +-
 .../ops_history_v2/SparseMatrixAdd.pbtxt      |  2 +-
 .../ops_history_v2/SparseMatrixMatMul.pbtxt   |  2 +-
 .../ops_history_v2/SparseMatrixMul.pbtxt      |  2 +-
 .../ops_history_v2/SparseMatrixNNZ.pbtxt      |  2 +-
 .../SparseMatrixOrderingAMD.pbtxt             |  2 +-
 .../ops_history_v2/SparseMatrixSoftmax.pbtxt  |  2 +-
 .../SparseMatrixSoftmaxGrad.pbtxt             |  2 +-
 .../SparseMatrixSparseCholesky.pbtxt          |  2 +-
 .../SparseMatrixSparseMatMul.pbtxt            |  2 +-
 .../SparseMatrixTranspose.pbtxt               |  2 +-
 .../ops_history_v2/SparseMatrixZeros.pbtxt    |  2 +-
 .../ops_history_v2/SparseReduceMax.pbtxt      |  2 +-
 .../SparseReduceMaxSparse.pbtxt               |  2 +-
 .../ops_history_v2/SparseReduceSum.pbtxt      |  2 +-
 .../SparseReduceSumSparse.pbtxt               |  2 +-
 .../compat/ops_history_v2/SparseReorder.pbtxt |  2 +-
 .../compat/ops_history_v2/SparseReshape.pbtxt |  2 +-
 .../ops_history_v2/SparseSegmentMean.pbtxt    |  2 +-
 .../SparseSegmentMeanGrad.pbtxt               |  2 +-
 .../SparseSegmentMeanGradV2.pbtxt             |  2 +-
 .../SparseSegmentMeanWithNumSegments.pbtxt    |  2 +-
 .../ops_history_v2/SparseSegmentSqrtN.pbtxt   |  2 +-
 .../SparseSegmentSqrtNGrad.pbtxt              |  2 +-
 .../SparseSegmentSqrtNGradV2.pbtxt            |  2 +-
 .../SparseSegmentSqrtNWithNumSegments.pbtxt   |  2 +-
 .../ops_history_v2/SparseSegmentSum.pbtxt     |  2 +-
 .../ops_history_v2/SparseSegmentSumGrad.pbtxt |  2 +-
 .../SparseSegmentSumGradV2.pbtxt              |  2 +-
 .../SparseSegmentSumWithNumSegments.pbtxt     |  2 +-
 .../compat/ops_history_v2/SparseSlice.pbtxt   |  2 +-
 .../ops_history_v2/SparseSliceGrad.pbtxt      |  2 +-
 .../compat/ops_history_v2/SparseSoftmax.pbtxt |  2 +-
 .../SparseSoftmaxCrossEntropyWithLogits.pbtxt |  2 +-
 .../ops_history_v2/SparseSparseMaximum.pbtxt  |  2 +-
 .../ops_history_v2/SparseSparseMinimum.pbtxt  |  2 +-
 .../compat/ops_history_v2/SparseSplit.pbtxt   |  2 +-
 .../ops_history_v2/SparseTensorDenseAdd.pbtxt |  2 +-
 .../SparseTensorDenseMatMul.pbtxt             |  2 +-
 .../SparseTensorSliceDataset.pbtxt            |  2 +-
 .../SparseTensorToCSRSparseMatrix.pbtxt       |  2 +-
 .../compat/ops_history_v2/SparseToDense.pbtxt |  2 +-
 .../SparseToSparseSetOperation.pbtxt          |  2 +-
 .../ops/compat/ops_history_v2/Spence.pbtxt    |  2 +-
 .../ops/compat/ops_history_v2/Split.pbtxt     |  2 +-
 .../ops_history_v2/SplitDedupData.pbtxt       |  2 +-
 .../ops/compat/ops_history_v2/SplitV.pbtxt    |  2 +-
 .../compat/ops_history_v2/SqlDataset.pbtxt    |  2 +-
 .../core/ops/compat/ops_history_v2/Sqrt.pbtxt |  2 +-
 .../ops/compat/ops_history_v2/SqrtGrad.pbtxt  |  2 +-
 .../ops/compat/ops_history_v2/Square.pbtxt    |  2 +-
 .../ops_history_v2/SquaredDifference.pbtxt    |  2 +-
 .../ops/compat/ops_history_v2/Squeeze.pbtxt   |  2 +-
 .../ops/compat/ops_history_v2/Stack.pbtxt     |  2 +-
 .../compat/ops_history_v2/StackClose.pbtxt    |  2 +-
 .../compat/ops_history_v2/StackCloseV2.pbtxt  |  2 +-
 .../ops/compat/ops_history_v2/StackPop.pbtxt  |  2 +-
 .../compat/ops_history_v2/StackPopV2.pbtxt    |  2 +-
 .../ops/compat/ops_history_v2/StackPush.pbtxt |  2 +-
 .../compat/ops_history_v2/StackPushV2.pbtxt   |  2 +-
 .../ops/compat/ops_history_v2/StackV2.pbtxt   |  2 +-
 .../ops/compat/ops_history_v2/Stage.pbtxt     |  2 +-
 .../compat/ops_history_v2/StageClear.pbtxt    |  2 +-
 .../ops/compat/ops_history_v2/StagePeek.pbtxt |  2 +-
 .../ops/compat/ops_history_v2/StageSize.pbtxt |  2 +-
 .../StatefulPartitionedCall.pbtxt             |  2 +-
 .../StatefulRandomBinomial.pbtxt              |  2 +-
 .../StatefulStandardNormal.pbtxt              |  2 +-
 .../StatefulStandardNormalV2.pbtxt            |  2 +-
 .../StatefulTruncatedNormal.pbtxt             |  2 +-
 .../ops_history_v2/StatefulUniform.pbtxt      |  2 +-
 .../StatefulUniformFullInt.pbtxt              |  2 +-
 .../ops_history_v2/StatefulUniformInt.pbtxt   |  2 +-
 .../compat/ops_history_v2/StatelessCase.pbtxt |  2 +-
 .../compat/ops_history_v2/StatelessIf.pbtxt   |  2 +-
 .../ops_history_v2/StatelessMultinomial.pbtxt |  2 +-
 ...tatelessParameterizedTruncatedNormal.pbtxt |  2 +-
 .../StatelessRandomBinomial.pbtxt             |  2 +-
 .../StatelessRandomGammaV2.pbtxt              |  2 +-
 .../StatelessRandomGammaV3.pbtxt              |  2 +-
 .../StatelessRandomGetAlg.pbtxt               |  2 +-
 .../StatelessRandomGetKeyCounter.pbtxt        |  2 +-
 .../StatelessRandomGetKeyCounterAlg.pbtxt     |  2 +-
 .../StatelessRandomNormal.pbtxt               |  2 +-
 .../StatelessRandomNormalV2.pbtxt             |  2 +-
 .../StatelessRandomPoisson.pbtxt              |  2 +-
 .../StatelessRandomUniform.pbtxt              |  2 +-
 .../StatelessRandomUniformFullInt.pbtxt       |  2 +-
 .../StatelessRandomUniformFullIntV2.pbtxt     |  2 +-
 .../StatelessRandomUniformInt.pbtxt           |  2 +-
 .../StatelessRandomUniformIntV2.pbtxt         |  2 +-
 .../StatelessRandomUniformV2.pbtxt            |  2 +-
 .../StatelessSampleDistortedBoundingBox.pbtxt |  2 +-
 .../ops_history_v2/StatelessShuffle.pbtxt     |  2 +-
 .../StatelessTruncatedNormal.pbtxt            |  2 +-
 .../StatelessTruncatedNormalV2.pbtxt          |  2 +-
 .../ops_history_v2/StatelessWhile.pbtxt       |  2 +-
 .../ops_history_v2/StaticRegexFullMatch.pbtxt |  2 +-
 .../ops_history_v2/StaticRegexReplace.pbtxt   |  2 +-
 .../StatsAggregatorHandle.pbtxt               |  2 +-
 .../StatsAggregatorHandleV2.pbtxt             |  2 +-
 .../StatsAggregatorSetSummaryWriter.pbtxt     |  2 +-
 .../StatsAggregatorSummary.pbtxt              |  2 +-
 .../ops_history_v2/StochasticCastToInt.pbtxt  |  2 +-
 .../compat/ops_history_v2/StopGradient.pbtxt  |  2 +-
 .../StoreMinibatchStatisticsInFdo.pbtxt       |  2 +-
 .../compat/ops_history_v2/StridedSlice.pbtxt  |  2 +-
 .../ops_history_v2/StridedSliceAssign.pbtxt   |  2 +-
 .../ops_history_v2/StridedSliceGrad.pbtxt     |  2 +-
 .../compat/ops_history_v2/StringFormat.pbtxt  |  2 +-
 .../compat/ops_history_v2/StringJoin.pbtxt    |  2 +-
 .../compat/ops_history_v2/StringLength.pbtxt  |  2 +-
 .../compat/ops_history_v2/StringLower.pbtxt   |  2 +-
 .../compat/ops_history_v2/StringNGrams.pbtxt  |  2 +-
 .../compat/ops_history_v2/StringSplit.pbtxt   |  2 +-
 .../compat/ops_history_v2/StringSplitV2.pbtxt |  2 +-
 .../compat/ops_history_v2/StringStrip.pbtxt   |  2 +-
 .../ops_history_v2/StringToHashBucket.pbtxt   |  2 +-
 .../StringToHashBucketFast.pbtxt              |  2 +-
 .../StringToHashBucketStrong.pbtxt            |  2 +-
 .../ops_history_v2/StringToNumber.pbtxt       | 30 ++++++++++++++++++-
 .../compat/ops_history_v2/StringUpper.pbtxt   |  2 +-
 .../core/ops/compat/ops_history_v2/Sub.pbtxt  |  2 +-
 .../ops/compat/ops_history_v2/Substr.pbtxt    |  2 +-
 .../core/ops/compat/ops_history_v2/Sum.pbtxt  |  2 +-
 .../compat/ops_history_v2/SummaryWriter.pbtxt |  2 +-
 .../core/ops/compat/ops_history_v2/Svd.pbtxt  |  2 +-
 .../ops/compat/ops_history_v2/Switch.pbtxt    |  2 +-
 .../ops_history_v2/SymbolicGradient.pbtxt     |  2 +-
 .../compat/ops_history_v2/SyncDevice.pbtxt    |  2 +-
 .../ops_history_v2/TFRecordDataset.pbtxt      |  2 +-
 .../ops_history_v2/TFRecordDatasetV2.pbtxt    |  2 +-
 .../ops_history_v2/TFRecordReader.pbtxt       |  2 +-
 .../ops_history_v2/TFRecordReaderV2.pbtxt     |  2 +-
 .../TPUAnnotateTensorsWithDynamicShape.pbtxt  |  2 +-
 .../ops_history_v2/TPUCompilationResult.pbtxt |  2 +-
 .../compat/ops_history_v2/TPUCompile.pbtxt    |  2 +-
 .../TPUCompileSucceededAssert.pbtxt           |  2 +-
 .../TPUCopyWithDynamicShape.pbtxt             |  2 +-
 .../TPUEmbeddingActivations.pbtxt             |  2 +-
 .../compat/ops_history_v2/TPUExecute.pbtxt    |  2 +-
 .../TPUExecuteAndUpdateVariables.pbtxt        |  2 +-
 .../ops_history_v2/TPUOrdinalSelector.pbtxt   |  2 +-
 .../ops_history_v2/TPUPartitionedCall.pbtxt   |  2 +-
 .../ops_history_v2/TPUPartitionedInput.pbtxt  |  2 +-
 .../TPUPartitionedInputV2.pbtxt               |  2 +-
 .../ops_history_v2/TPUPartitionedOutput.pbtxt |  2 +-
 .../TPUPartitionedOutputV2.pbtxt              |  2 +-
 .../ops_history_v2/TPUReplicateMetadata.pbtxt |  2 +-
 .../ops_history_v2/TPUReplicatedInput.pbtxt   |  2 +-
 .../ops_history_v2/TPUReplicatedOutput.pbtxt  |  2 +-
 .../ops_history_v2/TPUReshardVariables.pbtxt  |  2 +-
 .../compat/ops_history_v2/TPURoundRobin.pbtxt |  2 +-
 .../compat/ops_history_v2/TakeDataset.pbtxt   |  2 +-
 .../TakeManySparseFromTensorsMap.pbtxt        |  2 +-
 .../ops_history_v2/TakeWhileDataset.pbtxt     |  2 +-
 .../core/ops/compat/ops_history_v2/Tan.pbtxt  |  2 +-
 .../core/ops/compat/ops_history_v2/Tanh.pbtxt |  2 +-
 .../ops/compat/ops_history_v2/TanhGrad.pbtxt  |  2 +-
 .../ops_history_v2/TemporaryVariable.pbtxt    |  2 +-
 .../compat/ops_history_v2/TensorArray.pbtxt   |  2 +-
 .../ops_history_v2/TensorArrayClose.pbtxt     |  2 +-
 .../ops_history_v2/TensorArrayCloseV2.pbtxt   |  2 +-
 .../ops_history_v2/TensorArrayCloseV3.pbtxt   |  2 +-
 .../ops_history_v2/TensorArrayConcat.pbtxt    |  2 +-
 .../ops_history_v2/TensorArrayConcatV2.pbtxt  |  2 +-
 .../ops_history_v2/TensorArrayConcatV3.pbtxt  |  2 +-
 .../ops_history_v2/TensorArrayGather.pbtxt    |  2 +-
 .../ops_history_v2/TensorArrayGatherV2.pbtxt  |  2 +-
 .../ops_history_v2/TensorArrayGatherV3.pbtxt  |  2 +-
 .../ops_history_v2/TensorArrayGrad.pbtxt      |  2 +-
 .../ops_history_v2/TensorArrayGradV2.pbtxt    |  2 +-
 .../ops_history_v2/TensorArrayGradV3.pbtxt    |  2 +-
 .../TensorArrayGradWithShape.pbtxt            |  2 +-
 .../ops_history_v2/TensorArrayPack.pbtxt      |  2 +-
 .../ops_history_v2/TensorArrayRead.pbtxt      |  2 +-
 .../ops_history_v2/TensorArrayReadV2.pbtxt    |  2 +-
 .../ops_history_v2/TensorArrayReadV3.pbtxt    |  2 +-
 .../ops_history_v2/TensorArrayScatter.pbtxt   |  2 +-
 .../ops_history_v2/TensorArrayScatterV2.pbtxt |  2 +-
 .../ops_history_v2/TensorArrayScatterV3.pbtxt |  2 +-
 .../ops_history_v2/TensorArraySize.pbtxt      |  2 +-
 .../ops_history_v2/TensorArraySizeV2.pbtxt    |  2 +-
 .../ops_history_v2/TensorArraySizeV3.pbtxt    |  2 +-
 .../ops_history_v2/TensorArraySplit.pbtxt     |  2 +-
 .../ops_history_v2/TensorArraySplitV2.pbtxt   |  2 +-
 .../ops_history_v2/TensorArraySplitV3.pbtxt   |  2 +-
 .../ops_history_v2/TensorArrayUnpack.pbtxt    |  2 +-
 .../compat/ops_history_v2/TensorArrayV2.pbtxt |  2 +-
 .../compat/ops_history_v2/TensorArrayV3.pbtxt |  2 +-
 .../ops_history_v2/TensorArrayWrite.pbtxt     |  2 +-
 .../ops_history_v2/TensorArrayWriteV2.pbtxt   |  2 +-
 .../ops_history_v2/TensorArrayWriteV3.pbtxt   |  2 +-
 .../compat/ops_history_v2/TensorDataset.pbtxt |  2 +-
 .../ops_history_v2/TensorListConcat.pbtxt     |  2 +-
 .../TensorListConcatLists.pbtxt               |  2 +-
 .../ops_history_v2/TensorListConcatV2.pbtxt   |  2 +-
 .../TensorListElementShape.pbtxt              |  2 +-
 .../ops_history_v2/TensorListFromTensor.pbtxt |  2 +-
 .../ops_history_v2/TensorListGather.pbtxt     |  2 +-
 .../ops_history_v2/TensorListGetItem.pbtxt    |  2 +-
 .../ops_history_v2/TensorListLength.pbtxt     |  2 +-
 .../ops_history_v2/TensorListPopBack.pbtxt    |  2 +-
 .../ops_history_v2/TensorListPushBack.pbtxt   |  2 +-
 .../TensorListPushBackBatch.pbtxt             |  2 +-
 .../ops_history_v2/TensorListReserve.pbtxt    |  2 +-
 .../ops_history_v2/TensorListResize.pbtxt     |  2 +-
 .../ops_history_v2/TensorListScatter.pbtxt    |  2 +-
 .../TensorListScatterIntoExistingList.pbtxt   |  2 +-
 .../ops_history_v2/TensorListScatterV2.pbtxt  |  2 +-
 .../ops_history_v2/TensorListSetItem.pbtxt    |  2 +-
 .../ops_history_v2/TensorListSplit.pbtxt      |  2 +-
 .../ops_history_v2/TensorListStack.pbtxt      |  2 +-
 .../ops_history_v2/TensorMapErase.pbtxt       |  2 +-
 .../ops_history_v2/TensorMapHasKey.pbtxt      |  2 +-
 .../ops_history_v2/TensorMapInsert.pbtxt      |  2 +-
 .../ops_history_v2/TensorMapLookup.pbtxt      |  2 +-
 .../compat/ops_history_v2/TensorMapSize.pbtxt |  2 +-
 .../ops_history_v2/TensorMapStackKeys.pbtxt   |  2 +-
 .../ops_history_v2/TensorScatterAdd.pbtxt     |  2 +-
 .../ops_history_v2/TensorScatterMax.pbtxt     |  2 +-
 .../ops_history_v2/TensorScatterMin.pbtxt     |  2 +-
 .../ops_history_v2/TensorScatterSub.pbtxt     |  2 +-
 .../ops_history_v2/TensorScatterUpdate.pbtxt  |  2 +-
 .../ops_history_v2/TensorSliceDataset.pbtxt   |  2 +-
 .../TensorStridedSliceUpdate.pbtxt            |  2 +-
 .../compat/ops_history_v2/TensorSummary.pbtxt |  2 +-
 .../ops_history_v2/TensorSummaryV2.pbtxt      |  2 +-
 .../ops_history_v2/TextLineDataset.pbtxt      |  2 +-
 .../ops_history_v2/TextLineReader.pbtxt       |  2 +-
 .../ops_history_v2/TextLineReaderV2.pbtxt     |  2 +-
 .../ops_history_v2/ThreadPoolDataset.pbtxt    |  2 +-
 .../ops_history_v2/ThreadPoolHandle.pbtxt     |  2 +-
 .../ThreadUnsafeUnigramCandidateSampler.pbtxt |  2 +-
 .../core/ops/compat/ops_history_v2/Tile.pbtxt |  2 +-
 .../ops/compat/ops_history_v2/TileGrad.pbtxt  |  2 +-
 .../ops/compat/ops_history_v2/Timestamp.pbtxt |  2 +-
 .../ops/compat/ops_history_v2/ToBool.pbtxt    |  2 +-
 .../core/ops/compat/ops_history_v2/TopK.pbtxt |  2 +-
 .../compat/ops_history_v2/TopKUnique.pbtxt    |  2 +-
 .../ops/compat/ops_history_v2/TopKV2.pbtxt    |  2 +-
 .../ops_history_v2/TopKWithUnique.pbtxt       |  2 +-
 .../ops_history_v2/TpuHandleToProtoKey.pbtxt  |  2 +-
 .../ops/compat/ops_history_v2/Transpose.pbtxt |  2 +-
 .../ops_history_v2/TridiagonalMatMul.pbtxt    |  2 +-
 .../ops_history_v2/TridiagonalSolve.pbtxt     |  2 +-
 .../compat/ops_history_v2/TruncateDiv.pbtxt   |  2 +-
 .../compat/ops_history_v2/TruncateMod.pbtxt   |  2 +-
 .../ops_history_v2/TruncatedNormal.pbtxt      |  2 +-
 .../ops/compat/ops_history_v2/Unbatch.pbtxt   |  2 +-
 .../ops_history_v2/UnbatchDataset.pbtxt       |  2 +-
 .../compat/ops_history_v2/UnbatchGrad.pbtxt   |  2 +-
 .../ops_history_v2/UncompressElement.pbtxt    |  2 +-
 .../compat/ops_history_v2/UnicodeDecode.pbtxt |  2 +-
 .../UnicodeDecodeWithOffsets.pbtxt            |  2 +-
 .../compat/ops_history_v2/UnicodeEncode.pbtxt |  2 +-
 .../compat/ops_history_v2/UnicodeScript.pbtxt |  2 +-
 .../ops_history_v2/UnicodeTranscode.pbtxt     |  2 +-
 .../UniformCandidateSampler.pbtxt             |  2 +-
 .../ops_history_v2/UniformDequantize.pbtxt    |  2 +-
 .../ops_history_v2/UniformQuantize.pbtxt      |  2 +-
 .../ops_history_v2/UniformQuantizedAdd.pbtxt  |  2 +-
 .../UniformQuantizedClipByValue.pbtxt         |  2 +-
 .../UniformQuantizedConvolution.pbtxt         |  2 +-
 .../UniformQuantizedConvolutionHybrid.pbtxt   |  2 +-
 .../ops_history_v2/UniformQuantizedDot.pbtxt  |  2 +-
 .../UniformQuantizedDotHybrid.pbtxt           |  2 +-
 .../ops_history_v2/UniformRequantize.pbtxt    |  2 +-
 .../ops/compat/ops_history_v2/Unique.pbtxt    |  2 +-
 .../compat/ops_history_v2/UniqueDataset.pbtxt |  2 +-
 .../ops/compat/ops_history_v2/UniqueV2.pbtxt  |  2 +-
 .../ops_history_v2/UniqueWithCounts.pbtxt     |  2 +-
 .../ops_history_v2/UniqueWithCountsV2.pbtxt   |  2 +-
 .../ops/compat/ops_history_v2/Unpack.pbtxt    |  2 +-
 .../compat/ops_history_v2/UnravelIndex.pbtxt  |  2 +-
 .../ops_history_v2/UnsortedSegmentJoin.pbtxt  |  2 +-
 .../ops_history_v2/UnsortedSegmentMax.pbtxt   |  2 +-
 .../ops_history_v2/UnsortedSegmentMin.pbtxt   |  2 +-
 .../ops_history_v2/UnsortedSegmentProd.pbtxt  |  2 +-
 .../ops_history_v2/UnsortedSegmentSum.pbtxt   |  2 +-
 .../ops/compat/ops_history_v2/Unstage.pbtxt   |  2 +-
 .../ops_history_v2/UnwrapDatasetVariant.pbtxt |  2 +-
 .../compat/ops_history_v2/UpperBound.pbtxt    |  2 +-
 .../compat/ops_history_v2/VarHandleOp.pbtxt   |  2 +-
 .../ops_history_v2/VarIsInitializedOp.pbtxt   |  2 +-
 .../ops/compat/ops_history_v2/Variable.pbtxt  |  2 +-
 .../compat/ops_history_v2/VariableShape.pbtxt |  2 +-
 .../compat/ops_history_v2/VariableV2.pbtxt    |  2 +-
 .../WeightedFlatMapDataset.pbtxt              |  2 +-
 .../ops/compat/ops_history_v2/Where.pbtxt     |  2 +-
 .../ops/compat/ops_history_v2/While.pbtxt     |  2 +-
 .../ops_history_v2/WholeFileReader.pbtxt      |  2 +-
 .../ops_history_v2/WholeFileReaderV2.pbtxt    |  2 +-
 .../compat/ops_history_v2/WindowDataset.pbtxt |  2 +-
 .../ops/compat/ops_history_v2/WindowOp.pbtxt  |  2 +-
 .../ops_history_v2/WorkerHeartbeat.pbtxt      |  2 +-
 .../ops_history_v2/WrapDatasetVariant.pbtxt   |  2 +-
 .../ops_history_v2/WriteAudioSummary.pbtxt    |  2 +-
 .../ops/compat/ops_history_v2/WriteFile.pbtxt |  2 +-
 .../ops_history_v2/WriteGraphSummary.pbtxt    |  2 +-
 .../WriteHistogramSummary.pbtxt               |  2 +-
 .../ops_history_v2/WriteImageSummary.pbtxt    |  2 +-
 .../ops_history_v2/WriteRawProtoSummary.pbtxt |  2 +-
 .../ops_history_v2/WriteScalarSummary.pbtxt   |  2 +-
 .../compat/ops_history_v2/WriteSummary.pbtxt  |  2 +-
 .../ops/compat/ops_history_v2/Xdivy.pbtxt     |  2 +-
 .../compat/ops_history_v2/XlaConcatND.pbtxt   |  2 +-
 .../ops_history_v2/XlaHostCompute.pbtxt       |  2 +-
 .../ops_history_v2/XlaRecvFromHost.pbtxt      |  2 +-
 .../XlaRecvTPUEmbeddingActivations.pbtxt      |  2 +-
 .../XlaRecvTPUEmbeddingActivationsV2.pbtxt    |  2 +-
 ...XlaRecvTPUEmbeddingDeduplicationData.pbtxt |  2 +-
 ...aRecvTPUEmbeddingDeduplicationDataV2.pbtxt |  2 +-
 .../XlaSendTPUEmbeddingGradients.pbtxt        |  2 +-
 .../XlaSendTPUEmbeddingGradientsV2.pbtxt      |  2 +-
 .../compat/ops_history_v2/XlaSendToHost.pbtxt |  2 +-
 .../ops_history_v2/XlaSparseCoreAdagrad.pbtxt |  2 +-
 .../XlaSparseCoreAdagradMomentum.pbtxt        |  2 +-
 .../ops_history_v2/XlaSparseCoreAdam.pbtxt    |  2 +-
 .../ops_history_v2/XlaSparseCoreFtrl.pbtxt    |  2 +-
 .../ops_history_v2/XlaSparseCoreSgd.pbtxt     |  2 +-
 .../ops_history_v2/XlaSparseDenseMatmul.pbtxt |  2 +-
 ...enseMatmulGradWithAdagradAndCsrInput.pbtxt |  2 +-
 ...ulGradWithAdagradAndStaticBufferSize.pbtxt |  2 +-
 ...ulGradWithAdagradMomentumAndCsrInput.pbtxt |  2 +-
 ...thAdagradMomentumAndStaticBufferSize.pbtxt |  2 +-
 ...seDenseMatmulGradWithAdamAndCsrInput.pbtxt |  2 +-
 ...atmulGradWithAdamAndStaticBufferSize.pbtxt |  2 +-
 ...seDenseMatmulGradWithFtrlAndCsrInput.pbtxt |  2 +-
 ...atmulGradWithFtrlAndStaticBufferSize.pbtxt |  2 +-
 ...rseDenseMatmulGradWithSgdAndCsrInput.pbtxt |  2 +-
 ...MatmulGradWithSgdAndStaticBufferSize.pbtxt |  2 +-
 .../XlaSparseDenseMatmulWithCsrInput.pbtxt    |  2 +-
 ...parseDenseMatmulWithStaticBufferSize.pbtxt |  2 +-
 .../compat/ops_history_v2/XlaSplitND.pbtxt    |  2 +-
 .../ops/compat/ops_history_v2/Xlog1py.pbtxt   |  2 +-
 .../ops/compat/ops_history_v2/Xlogy.pbtxt     |  2 +-
 .../ops/compat/ops_history_v2/ZerosLike.pbtxt |  2 +-
 .../core/ops/compat/ops_history_v2/Zeta.pbtxt |  2 +-
 .../compat/ops_history_v2/ZipDataset.pbtxt    |  2 +-
 tensorflow/core/ops/ops.pbtxt                 |  4 ++-
 1550 files changed, 1579 insertions(+), 1550 deletions(-)

diff --git a/tensorflow/core/ops/compat/ops_history_v2/Abort.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/Abort.pbtxt
index 5d88f788a8cda1..4752385d6ecf6d 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/Abort.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/Abort.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "Abort"
   attr {
     name: "error_msg"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/Abs.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/Abs.pbtxt
index e901bc5794f682..5f44f9c6dca631 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/Abs.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/Abs.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "Abs"
   input_arg {
     name: "x"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/AccumulateNV2.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/AccumulateNV2.pbtxt
index 7a44ffd9fda37d..3e94aa154434f3 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/AccumulateNV2.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/AccumulateNV2.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "AccumulateNV2"
   input_arg {
     name: "inputs"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/AccumulatorApplyGradient.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/AccumulatorApplyGradient.pbtxt
index e920df37858235..dd1c9870a63327 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/AccumulatorApplyGradient.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/AccumulatorApplyGradient.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "AccumulatorApplyGradient"
   input_arg {
     name: "handle"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/AccumulatorNumAccumulated.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/AccumulatorNumAccumulated.pbtxt
index 131af0c1b5b24a..f378509e1e02f9 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/AccumulatorNumAccumulated.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/AccumulatorNumAccumulated.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "AccumulatorNumAccumulated"
   input_arg {
     name: "handle"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/AccumulatorSetGlobalStep.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/AccumulatorSetGlobalStep.pbtxt
index 0f50ac1852b834..9b4170df332a6a 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/AccumulatorSetGlobalStep.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/AccumulatorSetGlobalStep.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "AccumulatorSetGlobalStep"
   input_arg {
     name: "handle"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/AccumulatorTakeGradient.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/AccumulatorTakeGradient.pbtxt
index de3a79650b8f9d..1e28a68455fbf7 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/AccumulatorTakeGradient.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/AccumulatorTakeGradient.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "AccumulatorTakeGradient"
   input_arg {
     name: "handle"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/Acos.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/Acos.pbtxt
index 577976c6e4c579..504486147b2d30 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/Acos.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/Acos.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "Acos"
   input_arg {
     name: "x"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/Acosh.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/Acosh.pbtxt
index 136b9bfb209241..e53c8177f7f2f7 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/Acosh.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/Acosh.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "Acosh"
   input_arg {
     name: "x"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/Add.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/Add.pbtxt
index d1f87827922fac..ce30e6d5544d04 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/Add.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/Add.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "Add"
   input_arg {
     name: "x"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/AddManySparseToTensorsMap.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/AddManySparseToTensorsMap.pbtxt
index 433a7f40b042e9..c1433ccbaf9e3b 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/AddManySparseToTensorsMap.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/AddManySparseToTensorsMap.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "AddManySparseToTensorsMap"
   input_arg {
     name: "sparse_indices"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/AddN.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/AddN.pbtxt
index 106bb1a9a9b7ae..8935c304edc5a8 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/AddN.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/AddN.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "AddN"
   input_arg {
     name: "inputs"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/AddSparseToTensorsMap.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/AddSparseToTensorsMap.pbtxt
index 8c226e1965ce41..8a4c020d067b18 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/AddSparseToTensorsMap.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/AddSparseToTensorsMap.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "AddSparseToTensorsMap"
   input_arg {
     name: "sparse_indices"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/AddV2.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/AddV2.pbtxt
index ad79c179db20ab..a2be10b4abe51e 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/AddV2.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/AddV2.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "AddV2"
   input_arg {
     name: "x"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/AdjustContrast.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/AdjustContrast.pbtxt
index b2dc5e78e4b4e6..e51900d718bcc6 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/AdjustContrast.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/AdjustContrast.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "AdjustContrast"
   input_arg {
     name: "images"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/AdjustContrastv2.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/AdjustContrastv2.pbtxt
index 08fc84a5ab0cb0..6869f269dadf87 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/AdjustContrastv2.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/AdjustContrastv2.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "AdjustContrastv2"
   input_arg {
     name: "images"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/AdjustHue.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/AdjustHue.pbtxt
index 6cce51abdc8b28..9a6c72d3d8f515 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/AdjustHue.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/AdjustHue.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "AdjustHue"
   input_arg {
     name: "images"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/AdjustSaturation.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/AdjustSaturation.pbtxt
index 4bc90aec904be0..918ea188d8523b 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/AdjustSaturation.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/AdjustSaturation.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "AdjustSaturation"
   input_arg {
     name: "images"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/All.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/All.pbtxt
index 0afd8d468145d4..c0bc8f4beae4fc 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/All.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/All.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "All"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/AllCandidateSampler.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/AllCandidateSampler.pbtxt
index ff93f20c348920..e452850c261223 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/AllCandidateSampler.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/AllCandidateSampler.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "AllCandidateSampler"
   input_arg {
     name: "true_classes"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/AllToAll.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/AllToAll.pbtxt
index 23796980cd62fc..005d16aec0d9e6 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/AllToAll.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/AllToAll.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "AllToAll"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/Angle.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/Angle.pbtxt
index 5a1ff8b86c421c..ce28927f2b8118 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/Angle.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/Angle.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "Angle"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/AnonymousHashTable.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/AnonymousHashTable.pbtxt
index 1a62f991251223..15826399d357e4 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/AnonymousHashTable.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/AnonymousHashTable.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "AnonymousHashTable"
   output_arg {
     name: "table_handle"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/AnonymousIterator.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/AnonymousIterator.pbtxt
index 8094c8d79d6511..bf8f8fc2ed49d8 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/AnonymousIterator.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/AnonymousIterator.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "AnonymousIterator"
   output_arg {
     name: "handle"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/AnonymousIteratorV2.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/AnonymousIteratorV2.pbtxt
index 6b6dad7212bfe9..e7dca69e3e041c 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/AnonymousIteratorV2.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/AnonymousIteratorV2.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "AnonymousIteratorV2"
   output_arg {
     name: "handle"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/AnonymousIteratorV3.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/AnonymousIteratorV3.pbtxt
index 328434c6042ce2..15e975a4453c2e 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/AnonymousIteratorV3.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/AnonymousIteratorV3.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "AnonymousIteratorV3"
   output_arg {
     name: "handle"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/AnonymousMemoryCache.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/AnonymousMemoryCache.pbtxt
index b3ab7638e496a9..7f15df3e956212 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/AnonymousMemoryCache.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/AnonymousMemoryCache.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "AnonymousMemoryCache"
   output_arg {
     name: "handle"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/AnonymousMultiDeviceIterator.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/AnonymousMultiDeviceIterator.pbtxt
index a19b45630b5d2f..b8afaa363d6a66 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/AnonymousMultiDeviceIterator.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/AnonymousMultiDeviceIterator.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "AnonymousMultiDeviceIterator"
   output_arg {
     name: "handle"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/AnonymousMultiDeviceIteratorV3.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/AnonymousMultiDeviceIteratorV3.pbtxt
index 20d70a6acc9855..c503bf5c9d3e83 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/AnonymousMultiDeviceIteratorV3.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/AnonymousMultiDeviceIteratorV3.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "AnonymousMultiDeviceIteratorV3"
   output_arg {
     name: "handle"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/AnonymousMutableDenseHashTable.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/AnonymousMutableDenseHashTable.pbtxt
index f14a62b4faa484..6026fd263d36ff 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/AnonymousMutableDenseHashTable.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/AnonymousMutableDenseHashTable.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "AnonymousMutableDenseHashTable"
   input_arg {
     name: "empty_key"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/AnonymousMutableHashTable.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/AnonymousMutableHashTable.pbtxt
index b375400621b8ee..6f34858f4d3bd3 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/AnonymousMutableHashTable.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/AnonymousMutableHashTable.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "AnonymousMutableHashTable"
   output_arg {
     name: "table_handle"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/AnonymousMutableHashTableOfTensors.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/AnonymousMutableHashTableOfTensors.pbtxt
index 692d1963adffc8..21b3744a98b973 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/AnonymousMutableHashTableOfTensors.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/AnonymousMutableHashTableOfTensors.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "AnonymousMutableHashTableOfTensors"
   output_arg {
     name: "table_handle"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/AnonymousRandomSeedGenerator.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/AnonymousRandomSeedGenerator.pbtxt
index e805c04147a283..da2558b596ac6b 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/AnonymousRandomSeedGenerator.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/AnonymousRandomSeedGenerator.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "AnonymousRandomSeedGenerator"
   input_arg {
     name: "seed"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/AnonymousSeedGenerator.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/AnonymousSeedGenerator.pbtxt
index c2b9eee9b6a976..370b0460e22f3c 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/AnonymousSeedGenerator.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/AnonymousSeedGenerator.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "AnonymousSeedGenerator"
   input_arg {
     name: "seed"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/Any.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/Any.pbtxt
index 9b5d6350b83cfc..da020906a6d358 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/Any.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/Any.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "Any"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ApplyAdaMax.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ApplyAdaMax.pbtxt
index fe6e128a18034c..06f6134c329f5a 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ApplyAdaMax.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ApplyAdaMax.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ApplyAdaMax"
   input_arg {
     name: "var"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ApplyAdadelta.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ApplyAdadelta.pbtxt
index 211c8d4a64efca..477fd7e824e958 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ApplyAdadelta.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ApplyAdadelta.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ApplyAdadelta"
   input_arg {
     name: "var"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ApplyAdagrad.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ApplyAdagrad.pbtxt
index 5ede2bc76ea499..3da12f57396533 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ApplyAdagrad.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ApplyAdagrad.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ApplyAdagrad"
   input_arg {
     name: "var"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ApplyAdagradDA.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ApplyAdagradDA.pbtxt
index 3d00cb6f01bc9f..6795b12ea72d2f 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ApplyAdagradDA.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ApplyAdagradDA.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ApplyAdagradDA"
   input_arg {
     name: "var"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ApplyAdagradV2.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ApplyAdagradV2.pbtxt
index f7673fe102988f..12f70b41e8c237 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ApplyAdagradV2.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ApplyAdagradV2.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ApplyAdagradV2"
   input_arg {
     name: "var"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ApplyAdam.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ApplyAdam.pbtxt
index fdcf81e0d03a74..ae512d525f0a76 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ApplyAdam.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ApplyAdam.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ApplyAdam"
   input_arg {
     name: "var"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ApplyAddSign.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ApplyAddSign.pbtxt
index 9e485fcfec65c2..f489857b9940d1 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ApplyAddSign.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ApplyAddSign.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ApplyAddSign"
   input_arg {
     name: "var"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ApplyCenteredRMSProp.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ApplyCenteredRMSProp.pbtxt
index 3e50dc38c7c12f..908b913f567d1c 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ApplyCenteredRMSProp.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ApplyCenteredRMSProp.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ApplyCenteredRMSProp"
   input_arg {
     name: "var"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ApplyFtrl.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ApplyFtrl.pbtxt
index 6a67647bad7fae..2025575439504e 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ApplyFtrl.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ApplyFtrl.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ApplyFtrl"
   input_arg {
     name: "var"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ApplyFtrlV2.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ApplyFtrlV2.pbtxt
index 201d61b2fa8373..d2d1af5d40ccb6 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ApplyFtrlV2.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ApplyFtrlV2.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ApplyFtrlV2"
   input_arg {
     name: "var"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ApplyGradientDescent.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ApplyGradientDescent.pbtxt
index 25fb5723ebdaf9..8451d74ea1c4ba 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ApplyGradientDescent.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ApplyGradientDescent.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ApplyGradientDescent"
   input_arg {
     name: "var"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ApplyMomentum.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ApplyMomentum.pbtxt
index 289c6ea6151f89..f54384d4fbd58f 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ApplyMomentum.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ApplyMomentum.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ApplyMomentum"
   input_arg {
     name: "var"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ApplyPowerSign.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ApplyPowerSign.pbtxt
index fb3838c02a2e16..f2f8f79d5bad9b 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ApplyPowerSign.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ApplyPowerSign.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ApplyPowerSign"
   input_arg {
     name: "var"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ApplyProximalAdagrad.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ApplyProximalAdagrad.pbtxt
index c25959517fcd7b..632d4d5ff91e2f 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ApplyProximalAdagrad.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ApplyProximalAdagrad.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ApplyProximalAdagrad"
   input_arg {
     name: "var"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ApplyProximalGradientDescent.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ApplyProximalGradientDescent.pbtxt
index 3482b511cffd59..6a9c4e9ae6d27c 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ApplyProximalGradientDescent.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ApplyProximalGradientDescent.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ApplyProximalGradientDescent"
   input_arg {
     name: "var"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ApplyRMSProp.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ApplyRMSProp.pbtxt
index 24fe49118c943d..77a1e75f1f959f 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ApplyRMSProp.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ApplyRMSProp.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ApplyRMSProp"
   input_arg {
     name: "var"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ApproxTopK.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ApproxTopK.pbtxt
index a7a32594df6aa1..b97ee1259579ee 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ApproxTopK.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ApproxTopK.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ApproxTopK"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ApproximateEqual.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ApproximateEqual.pbtxt
index be01e038fdef3a..0b62abc45a378d 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ApproximateEqual.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ApproximateEqual.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ApproximateEqual"
   input_arg {
     name: "x"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ArgMax.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ArgMax.pbtxt
index 90679987916abd..0cfef9a64eb544 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ArgMax.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ArgMax.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ArgMax"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ArgMin.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ArgMin.pbtxt
index d2113ec185ec1d..4b83dd746cc3c5 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ArgMin.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ArgMin.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ArgMin"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/AsString.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/AsString.pbtxt
index b2b275470159db..7b985f6126d4a3 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/AsString.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/AsString.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "AsString"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/Asin.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/Asin.pbtxt
index 690d2a556e4be8..652e0ea057b672 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/Asin.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/Asin.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "Asin"
   input_arg {
     name: "x"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/Asinh.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/Asinh.pbtxt
index 3d78d9de59399d..7f31ec1236ce91 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/Asinh.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/Asinh.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "Asinh"
   input_arg {
     name: "x"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/Assert.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/Assert.pbtxt
index 343b2bd0b8da08..a891ca8c601309 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/Assert.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/Assert.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "Assert"
   input_arg {
     name: "condition"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/AssertCardinalityDataset.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/AssertCardinalityDataset.pbtxt
index 147978a3363667..edf77a307ece00 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/AssertCardinalityDataset.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/AssertCardinalityDataset.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "AssertCardinalityDataset"
   input_arg {
     name: "input_dataset"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/AssertNextDataset.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/AssertNextDataset.pbtxt
index 0bdc9b4f9b278e..f1ebd27d543b14 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/AssertNextDataset.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/AssertNextDataset.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "AssertNextDataset"
   input_arg {
     name: "input_dataset"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/AssertPrevDataset.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/AssertPrevDataset.pbtxt
index 23bced136aa919..62899074c77fd8 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/AssertPrevDataset.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/AssertPrevDataset.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "AssertPrevDataset"
   input_arg {
     name: "input_dataset"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/Assign.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/Assign.pbtxt
index 6760fe93d26221..9255e12f1a8f11 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/Assign.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/Assign.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "Assign"
   input_arg {
     name: "ref"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/AssignAdd.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/AssignAdd.pbtxt
index eea9bd9ad00ac5..a631b1b807a5ba 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/AssignAdd.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/AssignAdd.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "AssignAdd"
   input_arg {
     name: "ref"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/AssignAddVariableOp.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/AssignAddVariableOp.pbtxt
index 91adfade5b6bc3..c3a8b74a0daaf8 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/AssignAddVariableOp.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/AssignAddVariableOp.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "AssignAddVariableOp"
   input_arg {
     name: "resource"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/AssignSub.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/AssignSub.pbtxt
index 3e13c2c8c52e17..0337c992553583 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/AssignSub.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/AssignSub.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "AssignSub"
   input_arg {
     name: "ref"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/AssignSubVariableOp.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/AssignSubVariableOp.pbtxt
index a93f5576589e29..a5c9a567d077ae 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/AssignSubVariableOp.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/AssignSubVariableOp.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "AssignSubVariableOp"
   input_arg {
     name: "resource"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/AssignVariableOp.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/AssignVariableOp.pbtxt
index fa84589b303260..e7c5290c53ecdb 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/AssignVariableOp.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/AssignVariableOp.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "AssignVariableOp"
   input_arg {
     name: "resource"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/AssignVariableXlaConcatND.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/AssignVariableXlaConcatND.pbtxt
index 8a17f4cadc3417..7c9d5df2c9c8c8 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/AssignVariableXlaConcatND.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/AssignVariableXlaConcatND.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "AssignVariableXlaConcatND"
   input_arg {
     name: "resource"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/Atan.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/Atan.pbtxt
index 78fd00eea4d39b..b063bf81719aed 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/Atan.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/Atan.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "Atan"
   input_arg {
     name: "x"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/Atan2.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/Atan2.pbtxt
index 9ae51baed83ca7..e58675db4c19d0 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/Atan2.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/Atan2.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "Atan2"
   input_arg {
     name: "y"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/Atanh.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/Atanh.pbtxt
index 76aaad73963159..28d417a08544e9 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/Atanh.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/Atanh.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "Atanh"
   input_arg {
     name: "x"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/AudioSpectrogram.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/AudioSpectrogram.pbtxt
index 96bf31789fcec6..dbc2a2280dee3c 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/AudioSpectrogram.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/AudioSpectrogram.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "AudioSpectrogram"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/AudioSummary.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/AudioSummary.pbtxt
index 0a33e5ca25463e..4b1830595e07e4 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/AudioSummary.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/AudioSummary.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "AudioSummary"
   input_arg {
     name: "tag"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/AudioSummaryV2.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/AudioSummaryV2.pbtxt
index e92012709af98d..313c044aaeb506 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/AudioSummaryV2.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/AudioSummaryV2.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "AudioSummaryV2"
   input_arg {
     name: "tag"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/AutoShardDataset.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/AutoShardDataset.pbtxt
index 6ab98040d710eb..465b757c8e967b 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/AutoShardDataset.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/AutoShardDataset.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "AutoShardDataset"
   input_arg {
     name: "input_dataset"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/AvgPool.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/AvgPool.pbtxt
index f7472845fbcbd4..8e7db139a9a3c7 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/AvgPool.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/AvgPool.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "AvgPool"
   input_arg {
     name: "value"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/AvgPool3D.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/AvgPool3D.pbtxt
index 619b1b4c4eaa70..f3f60cbc1f18d2 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/AvgPool3D.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/AvgPool3D.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "AvgPool3D"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/AvgPool3DGrad.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/AvgPool3DGrad.pbtxt
index 3c74eac2382c39..67fef9572878db 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/AvgPool3DGrad.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/AvgPool3DGrad.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "AvgPool3DGrad"
   input_arg {
     name: "orig_input_shape"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/AvgPoolGrad.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/AvgPoolGrad.pbtxt
index f3952fab895c1f..6c72effaffaa43 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/AvgPoolGrad.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/AvgPoolGrad.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "AvgPoolGrad"
   input_arg {
     name: "orig_input_shape"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/BandedTriangularSolve.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/BandedTriangularSolve.pbtxt
index c69a9650010a38..5cf85a62392a2f 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/BandedTriangularSolve.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/BandedTriangularSolve.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "BandedTriangularSolve"
   input_arg {
     name: "matrix"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/Barrier.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/Barrier.pbtxt
index bdc8f8f53ba666..9391157b888851 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/Barrier.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/Barrier.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "Barrier"
   output_arg {
     name: "handle"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/BarrierClose.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/BarrierClose.pbtxt
index 4d121c65796ef0..69230484813264 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/BarrierClose.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/BarrierClose.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "BarrierClose"
   input_arg {
     name: "handle"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/BarrierIncompleteSize.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/BarrierIncompleteSize.pbtxt
index cb040bc1db3c0f..0d17c183684932 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/BarrierIncompleteSize.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/BarrierIncompleteSize.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "BarrierIncompleteSize"
   input_arg {
     name: "handle"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/BarrierInsertMany.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/BarrierInsertMany.pbtxt
index db48fd4b4554d9..86b64f603ebb0a 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/BarrierInsertMany.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/BarrierInsertMany.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "BarrierInsertMany"
   input_arg {
     name: "handle"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/BarrierReadySize.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/BarrierReadySize.pbtxt
index 8d214c6e6fedf0..e7b063005593aa 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/BarrierReadySize.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/BarrierReadySize.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "BarrierReadySize"
   input_arg {
     name: "handle"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/BarrierTakeMany.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/BarrierTakeMany.pbtxt
index 71e75d20ca7a08..e324042930451c 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/BarrierTakeMany.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/BarrierTakeMany.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "BarrierTakeMany"
   input_arg {
     name: "handle"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/Batch.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/Batch.pbtxt
index 24fc21b2f288db..d3ee8d8513d434 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/Batch.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/Batch.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "Batch"
   input_arg {
     name: "in_tensors"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/BatchCholesky.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/BatchCholesky.pbtxt
index a8d4223751ce30..5d38acc7c2d563 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/BatchCholesky.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/BatchCholesky.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "BatchCholesky"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/BatchCholeskyGrad.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/BatchCholeskyGrad.pbtxt
index 1beef0ebc8dda0..286ae3a81169d6 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/BatchCholeskyGrad.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/BatchCholeskyGrad.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "BatchCholeskyGrad"
   input_arg {
     name: "l"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/BatchDataset.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/BatchDataset.pbtxt
index da9db473d53bbb..39467ae1bb6d33 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/BatchDataset.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/BatchDataset.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "BatchDataset"
   input_arg {
     name: "input_dataset"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/BatchDatasetV2.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/BatchDatasetV2.pbtxt
index 4e95dbe7edd433..a3dc3afed0f53a 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/BatchDatasetV2.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/BatchDatasetV2.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "BatchDatasetV2"
   input_arg {
     name: "input_dataset"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/BatchFFT.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/BatchFFT.pbtxt
index 872afa3d2c76d1..4fe86a392f079b 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/BatchFFT.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/BatchFFT.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "BatchFFT"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/BatchFFT2D.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/BatchFFT2D.pbtxt
index 1bd0127c2a50ee..b52a6bdca44512 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/BatchFFT2D.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/BatchFFT2D.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "BatchFFT2D"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/BatchFFT3D.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/BatchFFT3D.pbtxt
index c1d39ed70f7671..7f19cf13c10573 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/BatchFFT3D.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/BatchFFT3D.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "BatchFFT3D"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/BatchFunction.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/BatchFunction.pbtxt
index e35c2c8618eddd..cf5e5896d084ba 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/BatchFunction.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/BatchFunction.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "BatchFunction"
   input_arg {
     name: "in_tensors"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/BatchIFFT.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/BatchIFFT.pbtxt
index 256417563cbc66..09d7b4ad7863a0 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/BatchIFFT.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/BatchIFFT.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "BatchIFFT"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/BatchIFFT2D.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/BatchIFFT2D.pbtxt
index 1c21807dededf1..23cc9cc51df5fc 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/BatchIFFT2D.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/BatchIFFT2D.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "BatchIFFT2D"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/BatchIFFT3D.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/BatchIFFT3D.pbtxt
index 3d6ddd79d52095..10a78fab914335 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/BatchIFFT3D.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/BatchIFFT3D.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "BatchIFFT3D"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/BatchMatMul.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/BatchMatMul.pbtxt
index 9d7ac3ca8e2a33..8bd778087f1467 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/BatchMatMul.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/BatchMatMul.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "BatchMatMul"
   input_arg {
     name: "x"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/BatchMatMulV2.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/BatchMatMulV2.pbtxt
index 4769d8220f53e1..ed724ddc49db94 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/BatchMatMulV2.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/BatchMatMulV2.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "BatchMatMulV2"
   input_arg {
     name: "x"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/BatchMatMulV3.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/BatchMatMulV3.pbtxt
index 1bcfdb937064ca..052a39e77c85ca 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/BatchMatMulV3.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/BatchMatMulV3.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "BatchMatMulV3"
   input_arg {
     name: "x"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/BatchMatrixBandPart.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/BatchMatrixBandPart.pbtxt
index 167fe7b69d484e..413681e612999d 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/BatchMatrixBandPart.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/BatchMatrixBandPart.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "BatchMatrixBandPart"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/BatchMatrixDeterminant.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/BatchMatrixDeterminant.pbtxt
index 02e8c5dc93ec83..4bc6081aa4482a 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/BatchMatrixDeterminant.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/BatchMatrixDeterminant.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "BatchMatrixDeterminant"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/BatchMatrixDiag.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/BatchMatrixDiag.pbtxt
index 4b06bc66e2a6dd..6104bef9340001 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/BatchMatrixDiag.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/BatchMatrixDiag.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "BatchMatrixDiag"
   input_arg {
     name: "diagonal"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/BatchMatrixDiagPart.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/BatchMatrixDiagPart.pbtxt
index 611f21ce1a1b67..9bd200f8cf5384 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/BatchMatrixDiagPart.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/BatchMatrixDiagPart.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "BatchMatrixDiagPart"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/BatchMatrixInverse.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/BatchMatrixInverse.pbtxt
index 7e413ef3b05da5..03a694d973a1b5 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/BatchMatrixInverse.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/BatchMatrixInverse.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "BatchMatrixInverse"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/BatchMatrixSetDiag.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/BatchMatrixSetDiag.pbtxt
index 377fe7cea09d92..f459184a0a3a43 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/BatchMatrixSetDiag.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/BatchMatrixSetDiag.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "BatchMatrixSetDiag"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/BatchMatrixSolve.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/BatchMatrixSolve.pbtxt
index 020873dea15c47..909502e91ab546 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/BatchMatrixSolve.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/BatchMatrixSolve.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "BatchMatrixSolve"
   input_arg {
     name: "matrix"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/BatchMatrixSolveLs.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/BatchMatrixSolveLs.pbtxt
index f6bc34b54f7b77..8c9d24efc7a404 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/BatchMatrixSolveLs.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/BatchMatrixSolveLs.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "BatchMatrixSolveLs"
   input_arg {
     name: "matrix"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/BatchMatrixTriangularSolve.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/BatchMatrixTriangularSolve.pbtxt
index f259cceac49a17..406fa62171f511 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/BatchMatrixTriangularSolve.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/BatchMatrixTriangularSolve.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "BatchMatrixTriangularSolve"
   input_arg {
     name: "matrix"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/BatchNormWithGlobalNormalization.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/BatchNormWithGlobalNormalization.pbtxt
index 846475ea1bfbbf..b9959a7d2dba82 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/BatchNormWithGlobalNormalization.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/BatchNormWithGlobalNormalization.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "BatchNormWithGlobalNormalization"
   input_arg {
     name: "t"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/BatchNormWithGlobalNormalizationGrad.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/BatchNormWithGlobalNormalizationGrad.pbtxt
index c8b1b878780d0c..170d512aec6406 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/BatchNormWithGlobalNormalizationGrad.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/BatchNormWithGlobalNormalizationGrad.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "BatchNormWithGlobalNormalizationGrad"
   input_arg {
     name: "t"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/BatchSelfAdjointEig.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/BatchSelfAdjointEig.pbtxt
index e124892ed2d677..42ba04199f7c81 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/BatchSelfAdjointEig.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/BatchSelfAdjointEig.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "BatchSelfAdjointEig"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/BatchSelfAdjointEigV2.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/BatchSelfAdjointEigV2.pbtxt
index bc9128925886cb..df3996ea2378c1 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/BatchSelfAdjointEigV2.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/BatchSelfAdjointEigV2.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "BatchSelfAdjointEigV2"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/BatchSvd.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/BatchSvd.pbtxt
index 7080ed06437879..0595ffcd2a6eb6 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/BatchSvd.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/BatchSvd.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "BatchSvd"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/BatchToSpace.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/BatchToSpace.pbtxt
index a96fdc58dfcb32..ac089e5ca76cf0 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/BatchToSpace.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/BatchToSpace.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "BatchToSpace"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/BatchToSpaceND.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/BatchToSpaceND.pbtxt
index 18560c832264a4..464beb31614de7 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/BatchToSpaceND.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/BatchToSpaceND.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "BatchToSpaceND"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/BesselI0.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/BesselI0.pbtxt
index b45aa841a7be67..78d524c916c861 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/BesselI0.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/BesselI0.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "BesselI0"
   input_arg {
     name: "x"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/BesselI0e.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/BesselI0e.pbtxt
index 8696576a64aa23..299cf82535aa06 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/BesselI0e.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/BesselI0e.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "BesselI0e"
   input_arg {
     name: "x"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/BesselI1.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/BesselI1.pbtxt
index 600fdb5771c3c1..e756c4655ddfdf 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/BesselI1.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/BesselI1.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "BesselI1"
   input_arg {
     name: "x"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/BesselI1e.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/BesselI1e.pbtxt
index d7284658a8bbb8..a9c8d0eb0e5a1b 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/BesselI1e.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/BesselI1e.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "BesselI1e"
   input_arg {
     name: "x"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/BesselJ0.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/BesselJ0.pbtxt
index 73ee83477debbf..35e14e5fdf173e 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/BesselJ0.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/BesselJ0.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "BesselJ0"
   input_arg {
     name: "x"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/BesselJ1.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/BesselJ1.pbtxt
index de8e56a83c2739..ef8814ea8f723a 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/BesselJ1.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/BesselJ1.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "BesselJ1"
   input_arg {
     name: "x"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/BesselK0.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/BesselK0.pbtxt
index 6bd13898bf74fa..ebb364d0371e52 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/BesselK0.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/BesselK0.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "BesselK0"
   input_arg {
     name: "x"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/BesselK0e.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/BesselK0e.pbtxt
index 4b1125eb3eac38..e3e680c9549023 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/BesselK0e.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/BesselK0e.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "BesselK0e"
   input_arg {
     name: "x"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/BesselK1.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/BesselK1.pbtxt
index 04aef9f8fe0e5b..f7ca7c2f6e27b1 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/BesselK1.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/BesselK1.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "BesselK1"
   input_arg {
     name: "x"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/BesselK1e.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/BesselK1e.pbtxt
index 461db7a9222ea1..96fe68d7b7f313 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/BesselK1e.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/BesselK1e.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "BesselK1e"
   input_arg {
     name: "x"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/BesselY0.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/BesselY0.pbtxt
index 6d97aba4578264..cd62af3477370a 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/BesselY0.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/BesselY0.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "BesselY0"
   input_arg {
     name: "x"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/BesselY1.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/BesselY1.pbtxt
index 6aa5123957547c..06f4c08eaf6932 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/BesselY1.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/BesselY1.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "BesselY1"
   input_arg {
     name: "x"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/Betainc.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/Betainc.pbtxt
index 330e5456a7afa6..b1523bff9e9807 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/Betainc.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/Betainc.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "Betainc"
   input_arg {
     name: "a"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/BiasAdd.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/BiasAdd.pbtxt
index 15e7dad6982866..2eba00671d375c 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/BiasAdd.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/BiasAdd.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "BiasAdd"
   input_arg {
     name: "value"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/BiasAddGrad.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/BiasAddGrad.pbtxt
index ea11e9ee5dda20..f85a2b9f6a5c01 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/BiasAddGrad.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/BiasAddGrad.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "BiasAddGrad"
   input_arg {
     name: "out_backprop"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/BiasAddV1.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/BiasAddV1.pbtxt
index b1f6b0cc1fc478..b35e45165b8d63 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/BiasAddV1.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/BiasAddV1.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "BiasAddV1"
   input_arg {
     name: "value"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/Bincount.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/Bincount.pbtxt
index 1e8d56cd7b2dc9..12135bbd54a063 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/Bincount.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/Bincount.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "Bincount"
   input_arg {
     name: "arr"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/Bitcast.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/Bitcast.pbtxt
index e3c5814e29872b..993a0c6da9ec65 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/Bitcast.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/Bitcast.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "Bitcast"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/BitwiseAnd.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/BitwiseAnd.pbtxt
index 1d8c1eb88c94e3..4b90e0e3de225a 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/BitwiseAnd.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/BitwiseAnd.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "BitwiseAnd"
   input_arg {
     name: "x"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/BitwiseOr.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/BitwiseOr.pbtxt
index 681b469fc95130..393a506f339896 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/BitwiseOr.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/BitwiseOr.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "BitwiseOr"
   input_arg {
     name: "x"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/BitwiseXor.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/BitwiseXor.pbtxt
index b8d801443cff93..c72b23fc432331 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/BitwiseXor.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/BitwiseXor.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "BitwiseXor"
   input_arg {
     name: "x"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/BlockLSTM.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/BlockLSTM.pbtxt
index 2c4ad7866c77ea..63180f534f17b0 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/BlockLSTM.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/BlockLSTM.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "BlockLSTM"
   input_arg {
     name: "seq_len_max"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/BlockLSTMGrad.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/BlockLSTMGrad.pbtxt
index a03cda1c387f96..e7b6458bc8adb9 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/BlockLSTMGrad.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/BlockLSTMGrad.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "BlockLSTMGrad"
   input_arg {
     name: "seq_len_max"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/BlockLSTMGradV2.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/BlockLSTMGradV2.pbtxt
index 11b3e4c48f42ea..ed0bd6b245645a 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/BlockLSTMGradV2.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/BlockLSTMGradV2.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "BlockLSTMGradV2"
   input_arg {
     name: "seq_len_max"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/BlockLSTMV2.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/BlockLSTMV2.pbtxt
index 35df5226c632c4..5fce517277de83 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/BlockLSTMV2.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/BlockLSTMV2.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "BlockLSTMV2"
   input_arg {
     name: "seq_len_max"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/BoostedTreesAggregateStats.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/BoostedTreesAggregateStats.pbtxt
index 45fe3a867da0b5..72994094399599 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/BoostedTreesAggregateStats.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/BoostedTreesAggregateStats.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "BoostedTreesAggregateStats"
   input_arg {
     name: "node_ids"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/BoostedTreesBucketize.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/BoostedTreesBucketize.pbtxt
index 61a170763f1dd8..5f277d3e0db3ae 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/BoostedTreesBucketize.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/BoostedTreesBucketize.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "BoostedTreesBucketize"
   input_arg {
     name: "float_values"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/BoostedTreesCalculateBestFeatureSplit.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/BoostedTreesCalculateBestFeatureSplit.pbtxt
index f885b3317006fe..50f35695e006fb 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/BoostedTreesCalculateBestFeatureSplit.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/BoostedTreesCalculateBestFeatureSplit.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "BoostedTreesCalculateBestFeatureSplit"
   input_arg {
     name: "node_id_range"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/BoostedTreesCalculateBestFeatureSplitV2.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/BoostedTreesCalculateBestFeatureSplitV2.pbtxt
index 38ddb60c44ff2c..e900ed9c6748ca 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/BoostedTreesCalculateBestFeatureSplitV2.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/BoostedTreesCalculateBestFeatureSplitV2.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "BoostedTreesCalculateBestFeatureSplitV2"
   input_arg {
     name: "node_id_range"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/BoostedTreesCalculateBestGainsPerFeature.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/BoostedTreesCalculateBestGainsPerFeature.pbtxt
index e5f863eb4ac8ef..f100db7b38671c 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/BoostedTreesCalculateBestGainsPerFeature.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/BoostedTreesCalculateBestGainsPerFeature.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "BoostedTreesCalculateBestGainsPerFeature"
   input_arg {
     name: "node_id_range"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/BoostedTreesCenterBias.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/BoostedTreesCenterBias.pbtxt
index 78eba4fcb4ba17..5c2fb9b5c54747 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/BoostedTreesCenterBias.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/BoostedTreesCenterBias.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "BoostedTreesCenterBias"
   input_arg {
     name: "tree_ensemble_handle"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/BoostedTreesCreateEnsemble.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/BoostedTreesCreateEnsemble.pbtxt
index 8816746475cf45..cea6d23f91ff7b 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/BoostedTreesCreateEnsemble.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/BoostedTreesCreateEnsemble.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "BoostedTreesCreateEnsemble"
   input_arg {
     name: "tree_ensemble_handle"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/BoostedTreesCreateQuantileStreamResource.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/BoostedTreesCreateQuantileStreamResource.pbtxt
index d14e8e40b3a966..3d0d64adcd58fc 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/BoostedTreesCreateQuantileStreamResource.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/BoostedTreesCreateQuantileStreamResource.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "BoostedTreesCreateQuantileStreamResource"
   input_arg {
     name: "quantile_stream_resource_handle"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/BoostedTreesDeserializeEnsemble.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/BoostedTreesDeserializeEnsemble.pbtxt
index 18aeb311841f89..b6d55ea0544c06 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/BoostedTreesDeserializeEnsemble.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/BoostedTreesDeserializeEnsemble.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "BoostedTreesDeserializeEnsemble"
   input_arg {
     name: "tree_ensemble_handle"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/BoostedTreesEnsembleResourceHandleOp.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/BoostedTreesEnsembleResourceHandleOp.pbtxt
index 6c79d9a75101cf..00573c1b95a7ec 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/BoostedTreesEnsembleResourceHandleOp.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/BoostedTreesEnsembleResourceHandleOp.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "BoostedTreesEnsembleResourceHandleOp"
   output_arg {
     name: "resource"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/BoostedTreesExampleDebugOutputs.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/BoostedTreesExampleDebugOutputs.pbtxt
index 8e8f1aa7f332c7..066be042842aef 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/BoostedTreesExampleDebugOutputs.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/BoostedTreesExampleDebugOutputs.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "BoostedTreesExampleDebugOutputs"
   input_arg {
     name: "tree_ensemble_handle"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/BoostedTreesFlushQuantileSummaries.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/BoostedTreesFlushQuantileSummaries.pbtxt
index d2b84d1e01c5ae..ae35e1023d4fae 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/BoostedTreesFlushQuantileSummaries.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/BoostedTreesFlushQuantileSummaries.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "BoostedTreesFlushQuantileSummaries"
   input_arg {
     name: "quantile_stream_resource_handle"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/BoostedTreesGetEnsembleStates.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/BoostedTreesGetEnsembleStates.pbtxt
index a6bc2b2749409e..1959384a36b4a5 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/BoostedTreesGetEnsembleStates.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/BoostedTreesGetEnsembleStates.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "BoostedTreesGetEnsembleStates"
   input_arg {
     name: "tree_ensemble_handle"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/BoostedTreesMakeQuantileSummaries.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/BoostedTreesMakeQuantileSummaries.pbtxt
index 14eacf93d81cdd..bbefa8b8711b1a 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/BoostedTreesMakeQuantileSummaries.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/BoostedTreesMakeQuantileSummaries.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "BoostedTreesMakeQuantileSummaries"
   input_arg {
     name: "float_values"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/BoostedTreesMakeStatsSummary.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/BoostedTreesMakeStatsSummary.pbtxt
index 0772fc259310aa..49a82d2ba0c772 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/BoostedTreesMakeStatsSummary.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/BoostedTreesMakeStatsSummary.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "BoostedTreesMakeStatsSummary"
   input_arg {
     name: "node_ids"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/BoostedTreesPredict.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/BoostedTreesPredict.pbtxt
index 7f66b1248fe293..7f176cdc901666 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/BoostedTreesPredict.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/BoostedTreesPredict.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "BoostedTreesPredict"
   input_arg {
     name: "tree_ensemble_handle"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/BoostedTreesQuantileStreamResourceAddSummaries.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/BoostedTreesQuantileStreamResourceAddSummaries.pbtxt
index ba12d9cdccfd88..97e875f4ea0bd8 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/BoostedTreesQuantileStreamResourceAddSummaries.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/BoostedTreesQuantileStreamResourceAddSummaries.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "BoostedTreesQuantileStreamResourceAddSummaries"
   input_arg {
     name: "quantile_stream_resource_handle"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/BoostedTreesQuantileStreamResourceDeserialize.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/BoostedTreesQuantileStreamResourceDeserialize.pbtxt
index a08188e2dfbcc2..c3f01fef2a49a9 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/BoostedTreesQuantileStreamResourceDeserialize.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/BoostedTreesQuantileStreamResourceDeserialize.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "BoostedTreesQuantileStreamResourceDeserialize"
   input_arg {
     name: "quantile_stream_resource_handle"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/BoostedTreesQuantileStreamResourceFlush.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/BoostedTreesQuantileStreamResourceFlush.pbtxt
index ca7affea02c882..fc2613a0d372e9 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/BoostedTreesQuantileStreamResourceFlush.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/BoostedTreesQuantileStreamResourceFlush.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "BoostedTreesQuantileStreamResourceFlush"
   input_arg {
     name: "quantile_stream_resource_handle"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/BoostedTreesQuantileStreamResourceGetBucketBoundaries.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/BoostedTreesQuantileStreamResourceGetBucketBoundaries.pbtxt
index b359fce14edefc..b2aa8dd4e72321 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/BoostedTreesQuantileStreamResourceGetBucketBoundaries.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/BoostedTreesQuantileStreamResourceGetBucketBoundaries.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "BoostedTreesQuantileStreamResourceGetBucketBoundaries"
   input_arg {
     name: "quantile_stream_resource_handle"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/BoostedTreesQuantileStreamResourceHandleOp.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/BoostedTreesQuantileStreamResourceHandleOp.pbtxt
index c106a233c27023..ca40a0aa1ede35 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/BoostedTreesQuantileStreamResourceHandleOp.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/BoostedTreesQuantileStreamResourceHandleOp.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "BoostedTreesQuantileStreamResourceHandleOp"
   output_arg {
     name: "resource"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/BoostedTreesSerializeEnsemble.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/BoostedTreesSerializeEnsemble.pbtxt
index 88a5330954209c..29d19f03117634 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/BoostedTreesSerializeEnsemble.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/BoostedTreesSerializeEnsemble.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "BoostedTreesSerializeEnsemble"
   input_arg {
     name: "tree_ensemble_handle"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/BoostedTreesSparseAggregateStats.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/BoostedTreesSparseAggregateStats.pbtxt
index b482c304d3dac7..9260634bf1edb5 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/BoostedTreesSparseAggregateStats.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/BoostedTreesSparseAggregateStats.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "BoostedTreesSparseAggregateStats"
   input_arg {
     name: "node_ids"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/BoostedTreesSparseCalculateBestFeatureSplit.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/BoostedTreesSparseCalculateBestFeatureSplit.pbtxt
index cd590bfbc3838b..86f7a5ffd2118c 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/BoostedTreesSparseCalculateBestFeatureSplit.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/BoostedTreesSparseCalculateBestFeatureSplit.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "BoostedTreesSparseCalculateBestFeatureSplit"
   input_arg {
     name: "node_id_range"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/BoostedTreesTrainingPredict.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/BoostedTreesTrainingPredict.pbtxt
index 8b8b1053ff2c64..615f52c656de08 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/BoostedTreesTrainingPredict.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/BoostedTreesTrainingPredict.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "BoostedTreesTrainingPredict"
   input_arg {
     name: "tree_ensemble_handle"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/BoostedTreesUpdateEnsemble.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/BoostedTreesUpdateEnsemble.pbtxt
index 9d50a5d4797069..9cd779e314a38c 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/BoostedTreesUpdateEnsemble.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/BoostedTreesUpdateEnsemble.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "BoostedTreesUpdateEnsemble"
   input_arg {
     name: "tree_ensemble_handle"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/BoostedTreesUpdateEnsembleV2.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/BoostedTreesUpdateEnsembleV2.pbtxt
index f1dc44c6f7ed06..2a573217f3ec25 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/BoostedTreesUpdateEnsembleV2.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/BoostedTreesUpdateEnsembleV2.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "BoostedTreesUpdateEnsembleV2"
   input_arg {
     name: "tree_ensemble_handle"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/BroadcastArgs.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/BroadcastArgs.pbtxt
index 328d3fdb0da8ed..e6dc3990634597 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/BroadcastArgs.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/BroadcastArgs.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "BroadcastArgs"
   input_arg {
     name: "s0"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/BroadcastGradientArgs.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/BroadcastGradientArgs.pbtxt
index a0e831191f46cf..2e1d739f9881ac 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/BroadcastGradientArgs.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/BroadcastGradientArgs.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "BroadcastGradientArgs"
   input_arg {
     name: "s0"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/BroadcastTo.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/BroadcastTo.pbtxt
index 350e630e9df0fe..4d29f9ebeb8d22 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/BroadcastTo.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/BroadcastTo.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "BroadcastTo"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/Bucketize.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/Bucketize.pbtxt
index 76174ffde85193..abe818e666cafd 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/Bucketize.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/Bucketize.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "Bucketize"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/BytesProducedStatsDataset.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/BytesProducedStatsDataset.pbtxt
index 2e71b250953aaf..7cbfbbd146c83f 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/BytesProducedStatsDataset.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/BytesProducedStatsDataset.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "BytesProducedStatsDataset"
   input_arg {
     name: "input_dataset"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/CSRSparseMatrixComponents.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/CSRSparseMatrixComponents.pbtxt
index 050ebb37abfa96..614097be7efe9e 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/CSRSparseMatrixComponents.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/CSRSparseMatrixComponents.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "CSRSparseMatrixComponents"
   input_arg {
     name: "csr_sparse_matrix"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/CSRSparseMatrixToDense.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/CSRSparseMatrixToDense.pbtxt
index be8df4e2115271..ed0cab0fcca8f4 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/CSRSparseMatrixToDense.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/CSRSparseMatrixToDense.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "CSRSparseMatrixToDense"
   input_arg {
     name: "sparse_input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/CSRSparseMatrixToSparseTensor.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/CSRSparseMatrixToSparseTensor.pbtxt
index dba86fe65b7ed8..39870fe0ca893c 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/CSRSparseMatrixToSparseTensor.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/CSRSparseMatrixToSparseTensor.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "CSRSparseMatrixToSparseTensor"
   input_arg {
     name: "sparse_matrix"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/CSVDataset.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/CSVDataset.pbtxt
index 02655c49a13501..56e1a03ab62995 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/CSVDataset.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/CSVDataset.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "CSVDataset"
   input_arg {
     name: "filenames"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/CSVDatasetV2.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/CSVDatasetV2.pbtxt
index 0d3909d2bff660..c540e8407e8480 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/CSVDatasetV2.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/CSVDatasetV2.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "CSVDatasetV2"
   input_arg {
     name: "filenames"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/CTCBeamSearchDecoder.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/CTCBeamSearchDecoder.pbtxt
index 6bcd60420e5fdb..5bd1968c832d29 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/CTCBeamSearchDecoder.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/CTCBeamSearchDecoder.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "CTCBeamSearchDecoder"
   input_arg {
     name: "inputs"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/CTCGreedyDecoder.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/CTCGreedyDecoder.pbtxt
index ab7f9af6e8c943..65e36d169a3bd9 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/CTCGreedyDecoder.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/CTCGreedyDecoder.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "CTCGreedyDecoder"
   input_arg {
     name: "inputs"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/CTCLoss.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/CTCLoss.pbtxt
index 54ced032292324..1d7e041dd9111e 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/CTCLoss.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/CTCLoss.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "CTCLoss"
   input_arg {
     name: "inputs"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/CTCLossV2.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/CTCLossV2.pbtxt
index e9631882f6595a..5a68abaa4e5ad8 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/CTCLossV2.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/CTCLossV2.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "CTCLossV2"
   input_arg {
     name: "inputs"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/CacheDataset.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/CacheDataset.pbtxt
index 3acc3ce899e6c4..8b8ec246c5be80 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/CacheDataset.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/CacheDataset.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "CacheDataset"
   input_arg {
     name: "input_dataset"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/CacheDatasetV2.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/CacheDatasetV2.pbtxt
index 78d13d28238562..43fe482cba6e0a 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/CacheDatasetV2.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/CacheDatasetV2.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "CacheDatasetV2"
   input_arg {
     name: "input_dataset"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/Case.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/Case.pbtxt
index 5a2cb6a6dfafbb..39cfc3f723630d 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/Case.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/Case.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "Case"
   input_arg {
     name: "branch_index"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/Cast.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/Cast.pbtxt
index 695048c5c775e2..581a8e4a8cc8a7 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/Cast.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/Cast.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "Cast"
   input_arg {
     name: "x"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/Ceil.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/Ceil.pbtxt
index ad3771e9a84fdb..cdec08500072df 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/Ceil.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/Ceil.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "Ceil"
   input_arg {
     name: "x"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/CheckNumerics.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/CheckNumerics.pbtxt
index 64bafc119602a5..9e63b170b231db 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/CheckNumerics.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/CheckNumerics.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "CheckNumerics"
   input_arg {
     name: "tensor"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/CheckNumericsV2.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/CheckNumericsV2.pbtxt
index 362728e3dc53b1..ba06f6a3524e48 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/CheckNumericsV2.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/CheckNumericsV2.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "CheckNumericsV2"
   input_arg {
     name: "tensor"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/Cholesky.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/Cholesky.pbtxt
index bd26758db797c8..e3cee5fcf89aa8 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/Cholesky.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/Cholesky.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "Cholesky"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/CholeskyGrad.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/CholeskyGrad.pbtxt
index e526acd1f84513..0f7c7efd88c46f 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/CholeskyGrad.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/CholeskyGrad.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "CholeskyGrad"
   input_arg {
     name: "l"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ChooseFastestBranchDataset.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ChooseFastestBranchDataset.pbtxt
index ab120b02bcfafd..4850496a3954b3 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ChooseFastestBranchDataset.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ChooseFastestBranchDataset.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ChooseFastestBranchDataset"
   input_arg {
     name: "input_dataset"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ChooseFastestDataset.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ChooseFastestDataset.pbtxt
index f27c03e6e881d4..476e834edde27e 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ChooseFastestDataset.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ChooseFastestDataset.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ChooseFastestDataset"
   input_arg {
     name: "input_datasets"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ClipByValue.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ClipByValue.pbtxt
index 821ca0c2425f7d..8652a168e85187 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ClipByValue.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ClipByValue.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ClipByValue"
   input_arg {
     name: "t"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/CloseSummaryWriter.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/CloseSummaryWriter.pbtxt
index 7a4dea11575c66..f67e1aaff2470c 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/CloseSummaryWriter.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/CloseSummaryWriter.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "CloseSummaryWriter"
   input_arg {
     name: "writer"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/CollateTPUEmbeddingMemory.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/CollateTPUEmbeddingMemory.pbtxt
index 61f6f6fcccabaf..16c232371ad944 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/CollateTPUEmbeddingMemory.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/CollateTPUEmbeddingMemory.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "CollateTPUEmbeddingMemory"
   input_arg {
     name: "memory_configs"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/CollectiveAllToAllV2.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/CollectiveAllToAllV2.pbtxt
index db1520d1af3f19..85e76487c019f1 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/CollectiveAllToAllV2.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/CollectiveAllToAllV2.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "CollectiveAllToAllV2"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/CollectiveAllToAllV3.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/CollectiveAllToAllV3.pbtxt
index ce690300d651bb..8e08f0bbbf21bd 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/CollectiveAllToAllV3.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/CollectiveAllToAllV3.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "CollectiveAllToAllV3"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/CollectiveAssignGroupV2.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/CollectiveAssignGroupV2.pbtxt
index 888884e1dd1a87..891eab6162c961 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/CollectiveAssignGroupV2.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/CollectiveAssignGroupV2.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "CollectiveAssignGroupV2"
   input_arg {
     name: "group_assignment"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/CollectiveBcastRecv.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/CollectiveBcastRecv.pbtxt
index aee2ceeb9dd728..bb2f18d7fd2637 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/CollectiveBcastRecv.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/CollectiveBcastRecv.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "CollectiveBcastRecv"
   output_arg {
     name: "data"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/CollectiveBcastRecvV2.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/CollectiveBcastRecvV2.pbtxt
index 7da7797030f61b..35a48eec2582f2 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/CollectiveBcastRecvV2.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/CollectiveBcastRecvV2.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "CollectiveBcastRecvV2"
   input_arg {
     name: "group_size"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/CollectiveBcastSend.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/CollectiveBcastSend.pbtxt
index 5ac9ecf25b7da5..50719b2f193899 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/CollectiveBcastSend.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/CollectiveBcastSend.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "CollectiveBcastSend"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/CollectiveBcastSendV2.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/CollectiveBcastSendV2.pbtxt
index f5f058fca6ec59..3fa3a7d269d76b 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/CollectiveBcastSendV2.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/CollectiveBcastSendV2.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "CollectiveBcastSendV2"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/CollectiveGather.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/CollectiveGather.pbtxt
index 1c5a520a41daf3..7a4610a33f6072 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/CollectiveGather.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/CollectiveGather.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "CollectiveGather"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/CollectiveGatherV2.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/CollectiveGatherV2.pbtxt
index 3bd706437fafed..d85b0612c3dddb 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/CollectiveGatherV2.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/CollectiveGatherV2.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "CollectiveGatherV2"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/CollectiveInitializeCommunicator.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/CollectiveInitializeCommunicator.pbtxt
index 4ac5ab48cdba99..9b9340a118a35e 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/CollectiveInitializeCommunicator.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/CollectiveInitializeCommunicator.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "CollectiveInitializeCommunicator"
   input_arg {
     name: "group_key"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/CollectivePermute.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/CollectivePermute.pbtxt
index 9a4ba13ea1917b..b3d33d2447e8e7 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/CollectivePermute.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/CollectivePermute.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "CollectivePermute"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/CollectiveReduce.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/CollectiveReduce.pbtxt
index 303ce3ebd83abd..a93a8505dd3209 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/CollectiveReduce.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/CollectiveReduce.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "CollectiveReduce"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/CollectiveReduceScatterV2.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/CollectiveReduceScatterV2.pbtxt
index 49551a88bcf242..f64772c2d270f8 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/CollectiveReduceScatterV2.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/CollectiveReduceScatterV2.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "CollectiveReduceScatterV2"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/CollectiveReduceV2.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/CollectiveReduceV2.pbtxt
index 37e60303ec3657..b42db7ed9ae53c 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/CollectiveReduceV2.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/CollectiveReduceV2.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "CollectiveReduceV2"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/CollectiveReduceV3.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/CollectiveReduceV3.pbtxt
index 4459e5580d2264..53d1697add4c55 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/CollectiveReduceV3.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/CollectiveReduceV3.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "CollectiveReduceV3"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/CombinedNonMaxSuppression.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/CombinedNonMaxSuppression.pbtxt
index 983b747b1bac53..55e27122e9d862 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/CombinedNonMaxSuppression.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/CombinedNonMaxSuppression.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "CombinedNonMaxSuppression"
   input_arg {
     name: "boxes"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/Complex.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/Complex.pbtxt
index c22e2b8045e10a..5d17643c89473c 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/Complex.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/Complex.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "Complex"
   input_arg {
     name: "real"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ComplexAbs.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ComplexAbs.pbtxt
index 60f95372d7c985..6e7cfc1266aa15 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ComplexAbs.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ComplexAbs.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ComplexAbs"
   input_arg {
     name: "x"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/CompositeTensorVariantFromComponents.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/CompositeTensorVariantFromComponents.pbtxt
index f6a900fd692d55..ed0660bc02e05f 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/CompositeTensorVariantFromComponents.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/CompositeTensorVariantFromComponents.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "CompositeTensorVariantFromComponents"
   input_arg {
     name: "components"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/CompositeTensorVariantToComponents.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/CompositeTensorVariantToComponents.pbtxt
index b877c9535bd31d..fa8c5a4d95b316 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/CompositeTensorVariantToComponents.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/CompositeTensorVariantToComponents.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "CompositeTensorVariantToComponents"
   input_arg {
     name: "encoded"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/CompressElement.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/CompressElement.pbtxt
index 5da73190c445c8..07d8cb461afcbe 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/CompressElement.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/CompressElement.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "CompressElement"
   input_arg {
     name: "components"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ComputeAccidentalHits.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ComputeAccidentalHits.pbtxt
index 6d8dfad02d5845..0bac269ba6b9f2 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ComputeAccidentalHits.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ComputeAccidentalHits.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ComputeAccidentalHits"
   input_arg {
     name: "true_classes"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ComputeBatchSize.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ComputeBatchSize.pbtxt
index 67dfb6ef08def0..13ab4eef4d0a61 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ComputeBatchSize.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ComputeBatchSize.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ComputeBatchSize"
   input_arg {
     name: "input_dataset"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ComputeDedupDataSize.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ComputeDedupDataSize.pbtxt
index b9f68e23d3d77b..0d4ec98a96f357 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ComputeDedupDataSize.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ComputeDedupDataSize.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ComputeDedupDataSize"
   output_arg {
     name: "num_elements"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ComputeDedupDataSizeV2.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ComputeDedupDataSizeV2.pbtxt
index 7851f740a6b679..2493251c1fddc6 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ComputeDedupDataSizeV2.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ComputeDedupDataSizeV2.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ComputeDedupDataSizeV2"
   output_arg {
     name: "num_elements"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ComputeDedupDataTupleMask.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ComputeDedupDataTupleMask.pbtxt
index d04dc21c9546d6..ba91c01af1ba78 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ComputeDedupDataTupleMask.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ComputeDedupDataTupleMask.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ComputeDedupDataTupleMask"
   output_arg {
     name: "output_shape"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ComputeDedupDataTupleMaskV2.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ComputeDedupDataTupleMaskV2.pbtxt
index f1e3e21c75ad1b..87deca2a2daecd 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ComputeDedupDataTupleMaskV2.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ComputeDedupDataTupleMaskV2.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ComputeDedupDataTupleMaskV2"
   output_arg {
     name: "output_shape"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/Concat.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/Concat.pbtxt
index 878098fbb5586a..21ff0fda4338f0 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/Concat.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/Concat.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "Concat"
   input_arg {
     name: "concat_dim"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ConcatOffset.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ConcatOffset.pbtxt
index a2d85962951864..f2fbb0062a1f11 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ConcatOffset.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ConcatOffset.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ConcatOffset"
   input_arg {
     name: "concat_dim"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ConcatV2.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ConcatV2.pbtxt
index d14b7b8d507d4e..d11dc14a9fffba 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ConcatV2.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ConcatV2.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ConcatV2"
   input_arg {
     name: "values"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ConcatenateDataset.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ConcatenateDataset.pbtxt
index f3d827b100dd35..3bdf420d07b14b 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ConcatenateDataset.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ConcatenateDataset.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ConcatenateDataset"
   input_arg {
     name: "input_dataset"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ConditionalAccumulator.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ConditionalAccumulator.pbtxt
index fbb1e42c13cb5b..678e9736ff978e 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ConditionalAccumulator.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ConditionalAccumulator.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ConditionalAccumulator"
   output_arg {
     name: "handle"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ConfigureAndInitializeGlobalTPU.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ConfigureAndInitializeGlobalTPU.pbtxt
index 0ea06f390bcb5f..96ba7c2527d732 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ConfigureAndInitializeGlobalTPU.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ConfigureAndInitializeGlobalTPU.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ConfigureAndInitializeGlobalTPU"
   output_arg {
     name: "output"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ConfigureDistributedTPU.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ConfigureDistributedTPU.pbtxt
index 790c463a57b7f2..72f944042dcaaf 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ConfigureDistributedTPU.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ConfigureDistributedTPU.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ConfigureDistributedTPU"
   output_arg {
     name: "topology"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ConfigureTPUEmbedding.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ConfigureTPUEmbedding.pbtxt
index 6949705a6a0183..6e61f8870c30b7 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ConfigureTPUEmbedding.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ConfigureTPUEmbedding.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ConfigureTPUEmbedding"
   attr {
     name: "config"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ConfigureTPUEmbeddingHost.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ConfigureTPUEmbeddingHost.pbtxt
index e7c07085a59269..2e99447437a6e9 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ConfigureTPUEmbeddingHost.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ConfigureTPUEmbeddingHost.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ConfigureTPUEmbeddingHost"
   input_arg {
     name: "common_config"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ConfigureTPUEmbeddingMemory.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ConfigureTPUEmbeddingMemory.pbtxt
index 138d1486a5bcaf..ab1b0b8511e4c0 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ConfigureTPUEmbeddingMemory.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ConfigureTPUEmbeddingMemory.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ConfigureTPUEmbeddingMemory"
   input_arg {
     name: "common_config"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/Conj.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/Conj.pbtxt
index 74a543143d071a..6e98e166726cd9 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/Conj.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/Conj.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "Conj"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ConjugateTranspose.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ConjugateTranspose.pbtxt
index 16047a1a25d6fd..417a2a53cd0e7d 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ConjugateTranspose.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ConjugateTranspose.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ConjugateTranspose"
   input_arg {
     name: "x"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ConnectTPUEmbeddingHosts.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ConnectTPUEmbeddingHosts.pbtxt
index 78b2395bae806d..af1e2d5e0e596f 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ConnectTPUEmbeddingHosts.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ConnectTPUEmbeddingHosts.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ConnectTPUEmbeddingHosts"
   input_arg {
     name: "network_configs"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/Const.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/Const.pbtxt
index ea0c1c0864951d..6512d22e435f73 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/Const.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/Const.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "Const"
   output_arg {
     name: "output"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ConsumeMutexLock.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ConsumeMutexLock.pbtxt
index 3a36b864063ee6..4340267edf935b 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ConsumeMutexLock.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ConsumeMutexLock.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ConsumeMutexLock"
   input_arg {
     name: "mutex_lock"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ControlTrigger.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ControlTrigger.pbtxt
index 6c77cdfebf25dd..2fe84fee80152f 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ControlTrigger.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ControlTrigger.pbtxt
@@ -1,3 +1,3 @@
-op 	 {
+op {
   name: "ControlTrigger"
 }
diff --git a/tensorflow/core/ops/compat/ops_history_v2/Conv.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/Conv.pbtxt
index 1651b7a2150dd7..c626db885b594f 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/Conv.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/Conv.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "Conv"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/Conv2D.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/Conv2D.pbtxt
index 1e7916c736761c..a95e3d23aea137 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/Conv2D.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/Conv2D.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "Conv2D"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/Conv2DBackpropFilter.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/Conv2DBackpropFilter.pbtxt
index a7bb1cbfc63127..1c656127f7e76a 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/Conv2DBackpropFilter.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/Conv2DBackpropFilter.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "Conv2DBackpropFilter"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/Conv2DBackpropFilterV2.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/Conv2DBackpropFilterV2.pbtxt
index c56a507acddd69..e9bf463cc70237 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/Conv2DBackpropFilterV2.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/Conv2DBackpropFilterV2.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "Conv2DBackpropFilterV2"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/Conv2DBackpropInput.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/Conv2DBackpropInput.pbtxt
index 0c46e05bbe485e..04d32885860ea6 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/Conv2DBackpropInput.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/Conv2DBackpropInput.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "Conv2DBackpropInput"
   input_arg {
     name: "input_sizes"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/Conv2DBackpropInputV2.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/Conv2DBackpropInputV2.pbtxt
index 82148176cadab3..414a06604d1d71 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/Conv2DBackpropInputV2.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/Conv2DBackpropInputV2.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "Conv2DBackpropInputV2"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/Conv3D.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/Conv3D.pbtxt
index f62fda0c7c17c6..a04a66081fff37 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/Conv3D.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/Conv3D.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "Conv3D"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/Conv3DBackpropFilter.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/Conv3DBackpropFilter.pbtxt
index 1c43d8de15d2f2..3091cfdcba7629 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/Conv3DBackpropFilter.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/Conv3DBackpropFilter.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "Conv3DBackpropFilter"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/Conv3DBackpropFilterV2.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/Conv3DBackpropFilterV2.pbtxt
index fa96e2ceea7d65..2494eba6e8a2b6 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/Conv3DBackpropFilterV2.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/Conv3DBackpropFilterV2.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "Conv3DBackpropFilterV2"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/Conv3DBackpropInput.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/Conv3DBackpropInput.pbtxt
index 0a2429303607fa..7fa3a5548d7628 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/Conv3DBackpropInput.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/Conv3DBackpropInput.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "Conv3DBackpropInput"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/Conv3DBackpropInputV2.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/Conv3DBackpropInputV2.pbtxt
index 0bdc6e2d7d3d30..e01b33dc4a87ec 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/Conv3DBackpropInputV2.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/Conv3DBackpropInputV2.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "Conv3DBackpropInputV2"
   input_arg {
     name: "input_sizes"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ConvertToCooTensor.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ConvertToCooTensor.pbtxt
index 31aeb0dfd962c4..65577c03134142 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ConvertToCooTensor.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ConvertToCooTensor.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ConvertToCooTensor"
   input_arg {
     name: "indices_or_row_splits"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ConvertToListOfSparseCoreCooTensors.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ConvertToListOfSparseCoreCooTensors.pbtxt
index ad7dc09892cd22..137b4eca0acfb6 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ConvertToListOfSparseCoreCooTensors.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ConvertToListOfSparseCoreCooTensors.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ConvertToListOfSparseCoreCooTensors"
   input_arg {
     name: "indices_or_row_splits"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ConvertToSparseCoreCsrWrappedCooTensor.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ConvertToSparseCoreCsrWrappedCooTensor.pbtxt
index dc53acec41aed5..defd0e3b53b1c0 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ConvertToSparseCoreCsrWrappedCooTensor.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ConvertToSparseCoreCsrWrappedCooTensor.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ConvertToSparseCoreCsrWrappedCooTensor"
   input_arg {
     name: "sorted_row_ids_list"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/Copy.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/Copy.pbtxt
index 693fbfe081ac3f..258aecc2947457 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/Copy.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/Copy.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "Copy"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/CopyHost.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/CopyHost.pbtxt
index 89d49896724450..07eb864f460c25 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/CopyHost.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/CopyHost.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "CopyHost"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/CopyToMesh.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/CopyToMesh.pbtxt
index 3d1b4b1bffb059..50e0a66e784a74 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/CopyToMesh.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/CopyToMesh.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "CopyToMesh"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/CopyToMeshGrad.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/CopyToMeshGrad.pbtxt
index c64c8dc3790bf8..e75ffe9bc3eb37 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/CopyToMeshGrad.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/CopyToMeshGrad.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "CopyToMeshGrad"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/Cos.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/Cos.pbtxt
index 00cdee090bb89b..52b7c1e795560c 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/Cos.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/Cos.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "Cos"
   input_arg {
     name: "x"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/Cosh.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/Cosh.pbtxt
index 67b5ea0e591614..7a29316305061a 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/Cosh.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/Cosh.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "Cosh"
   input_arg {
     name: "x"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/CountUpTo.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/CountUpTo.pbtxt
index e9ff14eb74db60..05726df8c11e4c 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/CountUpTo.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/CountUpTo.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "CountUpTo"
   input_arg {
     name: "ref"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/CreateSummaryDbWriter.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/CreateSummaryDbWriter.pbtxt
index 1b8760c6bc65c9..7a5f844bbb6ed8 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/CreateSummaryDbWriter.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/CreateSummaryDbWriter.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "CreateSummaryDbWriter"
   input_arg {
     name: "writer"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/CreateSummaryFileWriter.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/CreateSummaryFileWriter.pbtxt
index b2edd27ffa5695..61106e9fc851eb 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/CreateSummaryFileWriter.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/CreateSummaryFileWriter.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "CreateSummaryFileWriter"
   input_arg {
     name: "writer"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/CropAndResize.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/CropAndResize.pbtxt
index cb5c951b82f088..57b02c6638874e 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/CropAndResize.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/CropAndResize.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "CropAndResize"
   input_arg {
     name: "image"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/CropAndResizeGradBoxes.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/CropAndResizeGradBoxes.pbtxt
index 6d6be12cf1146f..d3f62e3ba4cea4 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/CropAndResizeGradBoxes.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/CropAndResizeGradBoxes.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "CropAndResizeGradBoxes"
   input_arg {
     name: "grads"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/CropAndResizeGradImage.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/CropAndResizeGradImage.pbtxt
index a3d0713d1f34ac..6ae744f428a506 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/CropAndResizeGradImage.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/CropAndResizeGradImage.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "CropAndResizeGradImage"
   input_arg {
     name: "grads"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/Cross.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/Cross.pbtxt
index 12ba58289127fb..b80215fcbc16ee 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/Cross.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/Cross.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "Cross"
   input_arg {
     name: "a"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/CrossReplicaSum.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/CrossReplicaSum.pbtxt
index 13c1bb582fccf9..947028f2c539a9 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/CrossReplicaSum.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/CrossReplicaSum.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "CrossReplicaSum"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/CudnnRNN.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/CudnnRNN.pbtxt
index 4a7753b4988a00..181b2c30b153d5 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/CudnnRNN.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/CudnnRNN.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "CudnnRNN"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/CudnnRNNBackprop.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/CudnnRNNBackprop.pbtxt
index bc95978eb0da08..fbab0e98b82484 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/CudnnRNNBackprop.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/CudnnRNNBackprop.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "CudnnRNNBackprop"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/CudnnRNNBackpropV2.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/CudnnRNNBackpropV2.pbtxt
index c76d1583656abd..3d2b03657a5e88 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/CudnnRNNBackpropV2.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/CudnnRNNBackpropV2.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "CudnnRNNBackpropV2"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/CudnnRNNBackpropV3.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/CudnnRNNBackpropV3.pbtxt
index b901c64dbcbd3f..19caef08548618 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/CudnnRNNBackpropV3.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/CudnnRNNBackpropV3.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "CudnnRNNBackpropV3"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/CudnnRNNCanonicalToParams.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/CudnnRNNCanonicalToParams.pbtxt
index 3b46f61998c093..b8162cb8123909 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/CudnnRNNCanonicalToParams.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/CudnnRNNCanonicalToParams.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "CudnnRNNCanonicalToParams"
   input_arg {
     name: "num_layers"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/CudnnRNNCanonicalToParamsV2.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/CudnnRNNCanonicalToParamsV2.pbtxt
index 9c51d58cdbe96c..25b6fa3b1bb323 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/CudnnRNNCanonicalToParamsV2.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/CudnnRNNCanonicalToParamsV2.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "CudnnRNNCanonicalToParamsV2"
   input_arg {
     name: "num_layers"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/CudnnRNNParamsSize.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/CudnnRNNParamsSize.pbtxt
index d88ab4c5499f7a..d67825a0b67366 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/CudnnRNNParamsSize.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/CudnnRNNParamsSize.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "CudnnRNNParamsSize"
   input_arg {
     name: "num_layers"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/CudnnRNNParamsToCanonical.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/CudnnRNNParamsToCanonical.pbtxt
index 4cce256653d512..0c4917bdf6b554 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/CudnnRNNParamsToCanonical.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/CudnnRNNParamsToCanonical.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "CudnnRNNParamsToCanonical"
   input_arg {
     name: "num_layers"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/CudnnRNNParamsToCanonicalV2.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/CudnnRNNParamsToCanonicalV2.pbtxt
index 70fe8004c48bdd..40c3cd58536000 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/CudnnRNNParamsToCanonicalV2.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/CudnnRNNParamsToCanonicalV2.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "CudnnRNNParamsToCanonicalV2"
   input_arg {
     name: "num_layers"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/CudnnRNNV2.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/CudnnRNNV2.pbtxt
index caf4d17262626c..fb05fea4beb8c7 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/CudnnRNNV2.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/CudnnRNNV2.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "CudnnRNNV2"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/CudnnRNNV3.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/CudnnRNNV3.pbtxt
index 91535bbe7adab8..8e71b147b727af 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/CudnnRNNV3.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/CudnnRNNV3.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "CudnnRNNV3"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/Cumprod.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/Cumprod.pbtxt
index 9ae7fcfd322abb..2294de2efe238c 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/Cumprod.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/Cumprod.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "Cumprod"
   input_arg {
     name: "x"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/Cumsum.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/Cumsum.pbtxt
index a5817e2d44679b..4bf125962944d0 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/Cumsum.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/Cumsum.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "Cumsum"
   input_arg {
     name: "x"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/CumulativeLogsumexp.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/CumulativeLogsumexp.pbtxt
index e6e47e700f278f..21d25d772f478e 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/CumulativeLogsumexp.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/CumulativeLogsumexp.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "CumulativeLogsumexp"
   input_arg {
     name: "x"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/DTensorRestoreV2.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/DTensorRestoreV2.pbtxt
index 8b96f73f15b9b1..3e082064faff93 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/DTensorRestoreV2.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/DTensorRestoreV2.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "DTensorRestoreV2"
   input_arg {
     name: "prefix"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/DTensorSetGlobalTPUArray.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/DTensorSetGlobalTPUArray.pbtxt
index f466f7f4a0f78a..b7cf1cfa9c5904 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/DTensorSetGlobalTPUArray.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/DTensorSetGlobalTPUArray.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "DTensorSetGlobalTPUArray"
   input_arg {
     name: "topology"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/DataFormatDimMap.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/DataFormatDimMap.pbtxt
index 0550a275713d1d..a01806bad6afa3 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/DataFormatDimMap.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/DataFormatDimMap.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "DataFormatDimMap"
   input_arg {
     name: "x"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/DataFormatVecPermute.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/DataFormatVecPermute.pbtxt
index 1304fbbc4bacd8..e439414d6d3503 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/DataFormatVecPermute.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/DataFormatVecPermute.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "DataFormatVecPermute"
   input_arg {
     name: "x"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/DataServiceDataset.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/DataServiceDataset.pbtxt
index f15d2af6c6b9b9..449ae514cbb89c 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/DataServiceDataset.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/DataServiceDataset.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "DataServiceDataset"
   input_arg {
     name: "dataset_id"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/DataServiceDatasetV2.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/DataServiceDatasetV2.pbtxt
index cc213b8a2c0f84..1b3547e62b6ad8 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/DataServiceDatasetV2.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/DataServiceDatasetV2.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "DataServiceDatasetV2"
   input_arg {
     name: "dataset_id"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/DataServiceDatasetV3.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/DataServiceDatasetV3.pbtxt
index 8860baa354c08b..0151a0ed46c148 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/DataServiceDatasetV3.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/DataServiceDatasetV3.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "DataServiceDatasetV3"
   input_arg {
     name: "dataset_id"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/DataServiceDatasetV4.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/DataServiceDatasetV4.pbtxt
index ca1f1709f98197..1ea88a869b037c 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/DataServiceDatasetV4.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/DataServiceDatasetV4.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "DataServiceDatasetV4"
   input_arg {
     name: "dataset_id"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/DatasetCardinality.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/DatasetCardinality.pbtxt
index 8e6e5f85f83e4d..435bf6774e03cd 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/DatasetCardinality.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/DatasetCardinality.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "DatasetCardinality"
   input_arg {
     name: "input_dataset"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/DatasetFingerprint.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/DatasetFingerprint.pbtxt
index 11fcf992a5dd7b..db1c39a335720d 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/DatasetFingerprint.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/DatasetFingerprint.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "DatasetFingerprint"
   input_arg {
     name: "input_dataset"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/DatasetFromGraph.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/DatasetFromGraph.pbtxt
index b7afa744de1fb4..54c8ed8b014816 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/DatasetFromGraph.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/DatasetFromGraph.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "DatasetFromGraph"
   input_arg {
     name: "graph_def"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/DatasetToGraph.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/DatasetToGraph.pbtxt
index 1cf25d1f8bf874..5bc711b6b66080 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/DatasetToGraph.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/DatasetToGraph.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "DatasetToGraph"
   input_arg {
     name: "input_dataset"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/DatasetToGraphV2.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/DatasetToGraphV2.pbtxt
index 8fd09dfe6ba036..9f0500c8bdf214 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/DatasetToGraphV2.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/DatasetToGraphV2.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "DatasetToGraphV2"
   input_arg {
     name: "input_dataset"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/DatasetToSingleElement.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/DatasetToSingleElement.pbtxt
index 1531425688c05b..8633fe65f4dcc1 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/DatasetToSingleElement.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/DatasetToSingleElement.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "DatasetToSingleElement"
   input_arg {
     name: "dataset"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/DatasetToTFRecord.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/DatasetToTFRecord.pbtxt
index e6547a87e970ec..e13f88e1c25d3b 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/DatasetToTFRecord.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/DatasetToTFRecord.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "DatasetToTFRecord"
   input_arg {
     name: "input_dataset"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/Dawsn.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/Dawsn.pbtxt
index 7b69d80753e825..a2f8cba9d58b43 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/Dawsn.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/Dawsn.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "Dawsn"
   input_arg {
     name: "x"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/DebugGradientIdentity.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/DebugGradientIdentity.pbtxt
index 9c94f17ae44f02..e1b425730d8cd4 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/DebugGradientIdentity.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/DebugGradientIdentity.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "DebugGradientIdentity"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/DebugGradientRefIdentity.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/DebugGradientRefIdentity.pbtxt
index 1cd87c9d5335b1..f75b7784ac0c05 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/DebugGradientRefIdentity.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/DebugGradientRefIdentity.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "DebugGradientRefIdentity"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/DebugIdentity.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/DebugIdentity.pbtxt
index f971fb59d398e2..50f97d847d4f3d 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/DebugIdentity.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/DebugIdentity.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "DebugIdentity"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/DebugIdentityV2.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/DebugIdentityV2.pbtxt
index c6a523b2038bd3..ea92aaa4943cb1 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/DebugIdentityV2.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/DebugIdentityV2.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "DebugIdentityV2"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/DebugIdentityV3.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/DebugIdentityV3.pbtxt
index ddcfd4f6a874a0..73d76b888bc0d8 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/DebugIdentityV3.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/DebugIdentityV3.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "DebugIdentityV3"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/DebugNanCount.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/DebugNanCount.pbtxt
index cf635ac53aa6e1..82ae073497fd90 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/DebugNanCount.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/DebugNanCount.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "DebugNanCount"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/DebugNumericSummary.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/DebugNumericSummary.pbtxt
index 965ef30ab24a90..d108b54e1cdea9 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/DebugNumericSummary.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/DebugNumericSummary.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "DebugNumericSummary"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/DebugNumericSummaryV2.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/DebugNumericSummaryV2.pbtxt
index f58db91ce09782..58c11f58af0bf5 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/DebugNumericSummaryV2.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/DebugNumericSummaryV2.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "DebugNumericSummaryV2"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/DecodeAndCropJpeg.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/DecodeAndCropJpeg.pbtxt
index ee8f8536e500b8..d1d767f4714b32 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/DecodeAndCropJpeg.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/DecodeAndCropJpeg.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "DecodeAndCropJpeg"
   input_arg {
     name: "contents"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/DecodeBase64.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/DecodeBase64.pbtxt
index 700152f16eccb5..cc7d61ce119cf4 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/DecodeBase64.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/DecodeBase64.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "DecodeBase64"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/DecodeBmp.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/DecodeBmp.pbtxt
index d8ab837cc10e35..40ac5f0b20bd00 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/DecodeBmp.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/DecodeBmp.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "DecodeBmp"
   input_arg {
     name: "contents"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/DecodeCSV.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/DecodeCSV.pbtxt
index e60342992df29d..f4fee2a95a5ff3 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/DecodeCSV.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/DecodeCSV.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "DecodeCSV"
   input_arg {
     name: "records"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/DecodeCompressed.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/DecodeCompressed.pbtxt
index ec0e15b92b74ac..8a345fffa8f409 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/DecodeCompressed.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/DecodeCompressed.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "DecodeCompressed"
   input_arg {
     name: "bytes"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/DecodeGif.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/DecodeGif.pbtxt
index c0e732cd5f5d8d..89b21b376e4035 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/DecodeGif.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/DecodeGif.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "DecodeGif"
   input_arg {
     name: "contents"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/DecodeImage.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/DecodeImage.pbtxt
index 6771daae228047..066ffd1091d0ee 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/DecodeImage.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/DecodeImage.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "DecodeImage"
   input_arg {
     name: "contents"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/DecodeJSONExample.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/DecodeJSONExample.pbtxt
index 348bf9b1cda17a..ec37ae546c5b7d 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/DecodeJSONExample.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/DecodeJSONExample.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "DecodeJSONExample"
   input_arg {
     name: "json_examples"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/DecodeJpeg.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/DecodeJpeg.pbtxt
index ad09dea051a347..9a4b4e4443c530 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/DecodeJpeg.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/DecodeJpeg.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "DecodeJpeg"
   input_arg {
     name: "contents"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/DecodePaddedRaw.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/DecodePaddedRaw.pbtxt
index d0bd867a9c3893..d4c3dab8607e8a 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/DecodePaddedRaw.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/DecodePaddedRaw.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "DecodePaddedRaw"
   input_arg {
     name: "input_bytes"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/DecodePng.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/DecodePng.pbtxt
index 68c77685db7118..dd7bd024365408 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/DecodePng.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/DecodePng.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "DecodePng"
   input_arg {
     name: "contents"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/DecodeProtoV2.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/DecodeProtoV2.pbtxt
index e6a52119199998..ae72d29d4e72f9 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/DecodeProtoV2.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/DecodeProtoV2.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "DecodeProtoV2"
   input_arg {
     name: "bytes"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/DecodeRaw.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/DecodeRaw.pbtxt
index 5a3d5a475ffce7..bc8ce80563a202 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/DecodeRaw.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/DecodeRaw.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "DecodeRaw"
   input_arg {
     name: "bytes"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/DecodeWav.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/DecodeWav.pbtxt
index 698e67694d1224..8eba7b95ad01f0 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/DecodeWav.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/DecodeWav.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "DecodeWav"
   input_arg {
     name: "contents"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/DeepCopy.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/DeepCopy.pbtxt
index 197842671a3467..e673960be5dc04 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/DeepCopy.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/DeepCopy.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "DeepCopy"
   input_arg {
     name: "x"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/DeleteIterator.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/DeleteIterator.pbtxt
index afd3cf0c3fdc15..3050ea922611dd 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/DeleteIterator.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/DeleteIterator.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "DeleteIterator"
   input_arg {
     name: "handle"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/DeleteMemoryCache.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/DeleteMemoryCache.pbtxt
index 21de51004ebdda..821293ba6a7587 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/DeleteMemoryCache.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/DeleteMemoryCache.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "DeleteMemoryCache"
   input_arg {
     name: "handle"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/DeleteMultiDeviceIterator.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/DeleteMultiDeviceIterator.pbtxt
index 4bba829418c3e4..b4ae640cec2abc 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/DeleteMultiDeviceIterator.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/DeleteMultiDeviceIterator.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "DeleteMultiDeviceIterator"
   input_arg {
     name: "multi_device_iterator"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/DeleteRandomSeedGenerator.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/DeleteRandomSeedGenerator.pbtxt
index 9272e55a0d7b63..0c0d2d1057422b 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/DeleteRandomSeedGenerator.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/DeleteRandomSeedGenerator.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "DeleteRandomSeedGenerator"
   input_arg {
     name: "handle"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/DeleteSeedGenerator.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/DeleteSeedGenerator.pbtxt
index 125eae81544aa3..e588b1e9e7e85e 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/DeleteSeedGenerator.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/DeleteSeedGenerator.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "DeleteSeedGenerator"
   input_arg {
     name: "handle"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/DeleteSessionTensor.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/DeleteSessionTensor.pbtxt
index 6b4c55134baf66..def4c105535d95 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/DeleteSessionTensor.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/DeleteSessionTensor.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "DeleteSessionTensor"
   input_arg {
     name: "handle"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/DenseBincount.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/DenseBincount.pbtxt
index c96abfc5ab012c..9bab6854e406a9 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/DenseBincount.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/DenseBincount.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "DenseBincount"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/DenseCountSparseOutput.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/DenseCountSparseOutput.pbtxt
index cb500a1f4c9931..be566eab9f438d 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/DenseCountSparseOutput.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/DenseCountSparseOutput.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "DenseCountSparseOutput"
   input_arg {
     name: "values"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/DenseToCSRSparseMatrix.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/DenseToCSRSparseMatrix.pbtxt
index 159fe06324c252..c8b2f66380791a 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/DenseToCSRSparseMatrix.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/DenseToCSRSparseMatrix.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "DenseToCSRSparseMatrix"
   input_arg {
     name: "dense_input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/DenseToDenseSetOperation.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/DenseToDenseSetOperation.pbtxt
index 74a3d2aaa02b5b..5188a82414ef18 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/DenseToDenseSetOperation.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/DenseToDenseSetOperation.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "DenseToDenseSetOperation"
   input_arg {
     name: "set1"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/DenseToSparseBatchDataset.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/DenseToSparseBatchDataset.pbtxt
index 16838ea7d82424..cb972cca8aa42f 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/DenseToSparseBatchDataset.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/DenseToSparseBatchDataset.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "DenseToSparseBatchDataset"
   input_arg {
     name: "input_dataset"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/DenseToSparseSetOperation.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/DenseToSparseSetOperation.pbtxt
index e79d8f5e213011..71c9c37798f7b1 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/DenseToSparseSetOperation.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/DenseToSparseSetOperation.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "DenseToSparseSetOperation"
   input_arg {
     name: "set1"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/DepthToSpace.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/DepthToSpace.pbtxt
index f5dd95b60c97d0..422fe7ff53f2a4 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/DepthToSpace.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/DepthToSpace.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "DepthToSpace"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/DepthwiseConv2dNative.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/DepthwiseConv2dNative.pbtxt
index 90eae83c303cc0..805bf2d379aa03 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/DepthwiseConv2dNative.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/DepthwiseConv2dNative.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "DepthwiseConv2dNative"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/DepthwiseConv2dNativeBackpropFilter.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/DepthwiseConv2dNativeBackpropFilter.pbtxt
index 199b883851c5ab..119933bf0ae7de 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/DepthwiseConv2dNativeBackpropFilter.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/DepthwiseConv2dNativeBackpropFilter.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "DepthwiseConv2dNativeBackpropFilter"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/DepthwiseConv2dNativeBackpropInput.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/DepthwiseConv2dNativeBackpropInput.pbtxt
index 99edc181b7d635..8ad511910587f6 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/DepthwiseConv2dNativeBackpropInput.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/DepthwiseConv2dNativeBackpropInput.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "DepthwiseConv2dNativeBackpropInput"
   input_arg {
     name: "input_sizes"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/Dequantize.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/Dequantize.pbtxt
index 8db5d571979aea..1a39e1869bfb90 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/Dequantize.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/Dequantize.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "Dequantize"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/DeserializeIterator.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/DeserializeIterator.pbtxt
index 2d531bd93df362..1ae290e93c0f37 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/DeserializeIterator.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/DeserializeIterator.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "DeserializeIterator"
   input_arg {
     name: "resource_handle"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/DeserializeManySparse.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/DeserializeManySparse.pbtxt
index 0ebb833a4b6383..f0e75d96d94d7e 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/DeserializeManySparse.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/DeserializeManySparse.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "DeserializeManySparse"
   input_arg {
     name: "serialized_sparse"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/DeserializeSparse.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/DeserializeSparse.pbtxt
index c327f01280ec25..c23a9b58a62ed6 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/DeserializeSparse.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/DeserializeSparse.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "DeserializeSparse"
   input_arg {
     name: "serialized_sparse"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/DestroyResourceOp.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/DestroyResourceOp.pbtxt
index 355227a484d32b..aa16c5ad5235a5 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/DestroyResourceOp.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/DestroyResourceOp.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "DestroyResourceOp"
   input_arg {
     name: "resource"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/DestroyTemporaryVariable.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/DestroyTemporaryVariable.pbtxt
index 773e399783efe7..7e073b2f20b040 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/DestroyTemporaryVariable.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/DestroyTemporaryVariable.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "DestroyTemporaryVariable"
   input_arg {
     name: "ref"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/DeviceIndex.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/DeviceIndex.pbtxt
index cfb79b60e3e44d..c513889e7b1813 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/DeviceIndex.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/DeviceIndex.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "DeviceIndex"
   output_arg {
     name: "index"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/Diag.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/Diag.pbtxt
index bae6dbfc7d2e76..92cb2071cf9788 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/Diag.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/Diag.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "Diag"
   input_arg {
     name: "diagonal"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/DiagPart.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/DiagPart.pbtxt
index 00e64aa39669e2..aec8c871407ccb 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/DiagPart.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/DiagPart.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "DiagPart"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/Digamma.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/Digamma.pbtxt
index d29ad8723b06de..0c294e54f216ca 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/Digamma.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/Digamma.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "Digamma"
   input_arg {
     name: "x"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/Dilation2D.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/Dilation2D.pbtxt
index 30f06dc42062a5..1db8503014ae72 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/Dilation2D.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/Dilation2D.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "Dilation2D"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/Dilation2DBackpropFilter.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/Dilation2DBackpropFilter.pbtxt
index 4dcc4b6f248e45..5a5a9f1dbb39cd 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/Dilation2DBackpropFilter.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/Dilation2DBackpropFilter.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "Dilation2DBackpropFilter"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/Dilation2DBackpropInput.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/Dilation2DBackpropInput.pbtxt
index dc2e964457a8bf..8944211d86d45a 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/Dilation2DBackpropInput.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/Dilation2DBackpropInput.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "Dilation2DBackpropInput"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/DirectedInterleaveDataset.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/DirectedInterleaveDataset.pbtxt
index 9aaadbf0a6c014..61a9462fac3b59 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/DirectedInterleaveDataset.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/DirectedInterleaveDataset.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "DirectedInterleaveDataset"
   input_arg {
     name: "selector_input_dataset"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/DisableCopyOnRead.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/DisableCopyOnRead.pbtxt
index 8a91880bab9c5b..61d189ae526cfa 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/DisableCopyOnRead.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/DisableCopyOnRead.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "DisableCopyOnRead"
   input_arg {
     name: "resource"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/DistributedSave.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/DistributedSave.pbtxt
index 3cb284d442a94e..221820fddb1962 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/DistributedSave.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/DistributedSave.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "DistributedSave"
   input_arg {
     name: "dataset"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/Div.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/Div.pbtxt
index 3ec0077535986c..fdc955f19ebe17 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/Div.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/Div.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "Div"
   input_arg {
     name: "x"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/DivNoNan.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/DivNoNan.pbtxt
index 40fc208137bcbd..ee9025f52f25a7 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/DivNoNan.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/DivNoNan.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "DivNoNan"
   input_arg {
     name: "x"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/DrawBoundingBoxes.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/DrawBoundingBoxes.pbtxt
index 35952f3512e107..729817314e9c39 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/DrawBoundingBoxes.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/DrawBoundingBoxes.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "DrawBoundingBoxes"
   input_arg {
     name: "images"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/DrawBoundingBoxesV2.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/DrawBoundingBoxesV2.pbtxt
index 85ab33cee4f34c..0a561796ca0b61 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/DrawBoundingBoxesV2.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/DrawBoundingBoxesV2.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "DrawBoundingBoxesV2"
   input_arg {
     name: "images"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/DummyIterationCounter.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/DummyIterationCounter.pbtxt
index 2f6bf602b173d3..b1df20cae731ca 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/DummyIterationCounter.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/DummyIterationCounter.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "DummyIterationCounter"
   output_arg {
     name: "handle"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/DummyMemoryCache.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/DummyMemoryCache.pbtxt
index bbce6cafdc3958..63901e2585ea15 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/DummyMemoryCache.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/DummyMemoryCache.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "DummyMemoryCache"
   output_arg {
     name: "handle"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/DummySeedGenerator.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/DummySeedGenerator.pbtxt
index f08c0e07723d46..585bc7c7528344 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/DummySeedGenerator.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/DummySeedGenerator.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "DummySeedGenerator"
   output_arg {
     name: "handle"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/DynamicEnqueueTPUEmbeddingArbitraryTensorBatch.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/DynamicEnqueueTPUEmbeddingArbitraryTensorBatch.pbtxt
index 5ac83512996ed4..b494ece1a77218 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/DynamicEnqueueTPUEmbeddingArbitraryTensorBatch.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/DynamicEnqueueTPUEmbeddingArbitraryTensorBatch.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "DynamicEnqueueTPUEmbeddingArbitraryTensorBatch"
   input_arg {
     name: "sample_indices_or_row_lengths"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/DynamicEnqueueTPUEmbeddingRaggedTensorBatch.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/DynamicEnqueueTPUEmbeddingRaggedTensorBatch.pbtxt
index 506a023aa23583..46adf791fe5977 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/DynamicEnqueueTPUEmbeddingRaggedTensorBatch.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/DynamicEnqueueTPUEmbeddingRaggedTensorBatch.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "DynamicEnqueueTPUEmbeddingRaggedTensorBatch"
   input_arg {
     name: "sample_splits"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/DynamicPartition.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/DynamicPartition.pbtxt
index c497964b636bc3..3565bd6f754540 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/DynamicPartition.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/DynamicPartition.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "DynamicPartition"
   input_arg {
     name: "data"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/DynamicStitch.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/DynamicStitch.pbtxt
index 76226f65312d0a..aba8346995ccd3 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/DynamicStitch.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/DynamicStitch.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "DynamicStitch"
   input_arg {
     name: "indices"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/EagerPyFunc.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/EagerPyFunc.pbtxt
index 302e420778b6f0..56c12e3845c6ff 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/EagerPyFunc.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/EagerPyFunc.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "EagerPyFunc"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/EditDistance.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/EditDistance.pbtxt
index 78fbc3f1ac4d9e..aba098b7020cdb 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/EditDistance.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/EditDistance.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "EditDistance"
   input_arg {
     name: "hypothesis_indices"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/Eig.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/Eig.pbtxt
index 3872e50f11ed76..d95892a88367df 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/Eig.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/Eig.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "Eig"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/Einsum.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/Einsum.pbtxt
index 8a365bb22bf4d9..3855daa079bd40 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/Einsum.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/Einsum.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "Einsum"
   input_arg {
     name: "inputs"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/Elu.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/Elu.pbtxt
index 3c9ea1633638b9..4b8a81527569a9 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/Elu.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/Elu.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "Elu"
   input_arg {
     name: "features"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/EluGrad.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/EluGrad.pbtxt
index 2b2e8b974c18f7..cfbc9f99e314f3 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/EluGrad.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/EluGrad.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "EluGrad"
   input_arg {
     name: "gradients"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/Empty.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/Empty.pbtxt
index 760cab96732ccc..147854bb88cf23 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/Empty.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/Empty.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "Empty"
   input_arg {
     name: "shape"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/EmptyTensorList.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/EmptyTensorList.pbtxt
index 46015e651e219f..d15fa1ad47048c 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/EmptyTensorList.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/EmptyTensorList.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "EmptyTensorList"
   input_arg {
     name: "element_shape"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/EmptyTensorMap.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/EmptyTensorMap.pbtxt
index ee1fc166627b41..25327b4e1e8bfb 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/EmptyTensorMap.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/EmptyTensorMap.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "EmptyTensorMap"
   output_arg {
     name: "handle"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/EncodeBase64.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/EncodeBase64.pbtxt
index 001c3d0d9fb36d..6e5241d0fd7033 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/EncodeBase64.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/EncodeBase64.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "EncodeBase64"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/EncodeJpeg.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/EncodeJpeg.pbtxt
index 165a02aed32de1..9f3c3453e33e32 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/EncodeJpeg.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/EncodeJpeg.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "EncodeJpeg"
   input_arg {
     name: "image"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/EncodeJpegVariableQuality.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/EncodeJpegVariableQuality.pbtxt
index 97f93d28b288f0..94c41ea4e5d574 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/EncodeJpegVariableQuality.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/EncodeJpegVariableQuality.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "EncodeJpegVariableQuality"
   input_arg {
     name: "images"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/EncodePng.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/EncodePng.pbtxt
index 7b7d161e8cba6e..7d2cbd85225125 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/EncodePng.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/EncodePng.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "EncodePng"
   input_arg {
     name: "image"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/EncodeProto.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/EncodeProto.pbtxt
index 2cea90e4183ddf..e619618946a6c0 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/EncodeProto.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/EncodeProto.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "EncodeProto"
   input_arg {
     name: "sizes"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/EncodeWav.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/EncodeWav.pbtxt
index e468d229ff4024..b013362a47d9e4 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/EncodeWav.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/EncodeWav.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "EncodeWav"
   input_arg {
     name: "audio"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/EnqueueTPUEmbeddingArbitraryTensorBatch.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/EnqueueTPUEmbeddingArbitraryTensorBatch.pbtxt
index f8d3f09b3cd32a..efb854a44baf5e 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/EnqueueTPUEmbeddingArbitraryTensorBatch.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/EnqueueTPUEmbeddingArbitraryTensorBatch.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "EnqueueTPUEmbeddingArbitraryTensorBatch"
   input_arg {
     name: "sample_indices_or_row_lengths"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/EnqueueTPUEmbeddingBatch.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/EnqueueTPUEmbeddingBatch.pbtxt
index b1be42b17eed44..a09c5d87f44fda 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/EnqueueTPUEmbeddingBatch.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/EnqueueTPUEmbeddingBatch.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "EnqueueTPUEmbeddingBatch"
   input_arg {
     name: "batch"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/EnqueueTPUEmbeddingIntegerBatch.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/EnqueueTPUEmbeddingIntegerBatch.pbtxt
index 37b123d9791a2a..26d63b6e49dc5e 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/EnqueueTPUEmbeddingIntegerBatch.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/EnqueueTPUEmbeddingIntegerBatch.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "EnqueueTPUEmbeddingIntegerBatch"
   input_arg {
     name: "batch"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/EnqueueTPUEmbeddingRaggedTensorBatch.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/EnqueueTPUEmbeddingRaggedTensorBatch.pbtxt
index 5e1a15cbc95fcc..327bd4fb387700 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/EnqueueTPUEmbeddingRaggedTensorBatch.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/EnqueueTPUEmbeddingRaggedTensorBatch.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "EnqueueTPUEmbeddingRaggedTensorBatch"
   input_arg {
     name: "sample_splits"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/EnqueueTPUEmbeddingSparseBatch.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/EnqueueTPUEmbeddingSparseBatch.pbtxt
index 03a19c2feed572..64b8cb5178c5a6 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/EnqueueTPUEmbeddingSparseBatch.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/EnqueueTPUEmbeddingSparseBatch.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "EnqueueTPUEmbeddingSparseBatch"
   input_arg {
     name: "sample_indices"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/EnqueueTPUEmbeddingSparseTensorBatch.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/EnqueueTPUEmbeddingSparseTensorBatch.pbtxt
index 19521bb9213393..ab1c9d264eca1e 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/EnqueueTPUEmbeddingSparseTensorBatch.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/EnqueueTPUEmbeddingSparseTensorBatch.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "EnqueueTPUEmbeddingSparseTensorBatch"
   input_arg {
     name: "sample_indices"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/EnsureShape.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/EnsureShape.pbtxt
index 336e2aed15cff6..24fa5589131df3 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/EnsureShape.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/EnsureShape.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "EnsureShape"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/Enter.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/Enter.pbtxt
index d0a555da70c264..d39d15f34dbf6f 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/Enter.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/Enter.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "Enter"
   input_arg {
     name: "data"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/Equal.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/Equal.pbtxt
index 8ad87805a6c558..a50cbdfcfeb1ac 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/Equal.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/Equal.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "Equal"
   input_arg {
     name: "x"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/Erf.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/Erf.pbtxt
index a53d8e380ef17c..680b736fa3e4a3 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/Erf.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/Erf.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "Erf"
   input_arg {
     name: "x"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/Erfc.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/Erfc.pbtxt
index 589ca5a917c652..2fcfc68f04fbf2 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/Erfc.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/Erfc.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "Erfc"
   input_arg {
     name: "x"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/Erfinv.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/Erfinv.pbtxt
index 3b51f3b83c9d3b..78443dada86bc4 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/Erfinv.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/Erfinv.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "Erfinv"
   input_arg {
     name: "x"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/EuclideanNorm.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/EuclideanNorm.pbtxt
index b4916ba2c7ee8c..b88f521b57199f 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/EuclideanNorm.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/EuclideanNorm.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "EuclideanNorm"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ExecuteTPUEmbeddingPartitioner.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ExecuteTPUEmbeddingPartitioner.pbtxt
index c2a3d248d1f844..37025570305f98 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ExecuteTPUEmbeddingPartitioner.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ExecuteTPUEmbeddingPartitioner.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ExecuteTPUEmbeddingPartitioner"
   output_arg {
     name: "common_config"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/Exit.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/Exit.pbtxt
index a722f2b7d5ced9..56a1371bec6b3f 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/Exit.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/Exit.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "Exit"
   input_arg {
     name: "data"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/Exp.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/Exp.pbtxt
index a237e834372291..7afeb677a29a54 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/Exp.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/Exp.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "Exp"
   input_arg {
     name: "x"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ExpandDims.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ExpandDims.pbtxt
index 9bd506742ce3d7..c7bb353162c8c8 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ExpandDims.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ExpandDims.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ExpandDims"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ExperimentalAssertNextDataset.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ExperimentalAssertNextDataset.pbtxt
index d33df8039e8b66..937afb4e0ae48f 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ExperimentalAssertNextDataset.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ExperimentalAssertNextDataset.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ExperimentalAssertNextDataset"
   input_arg {
     name: "input_dataset"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ExperimentalAutoShardDataset.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ExperimentalAutoShardDataset.pbtxt
index 7a609f0300a365..d8d8d7ffc96e4e 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ExperimentalAutoShardDataset.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ExperimentalAutoShardDataset.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ExperimentalAutoShardDataset"
   input_arg {
     name: "input_dataset"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ExperimentalBytesProducedStatsDataset.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ExperimentalBytesProducedStatsDataset.pbtxt
index 509dd811947653..10555bb3ebfbf5 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ExperimentalBytesProducedStatsDataset.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ExperimentalBytesProducedStatsDataset.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ExperimentalBytesProducedStatsDataset"
   input_arg {
     name: "input_dataset"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ExperimentalCSVDataset.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ExperimentalCSVDataset.pbtxt
index 85d8950a845328..a618d55fcb6289 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ExperimentalCSVDataset.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ExperimentalCSVDataset.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ExperimentalCSVDataset"
   input_arg {
     name: "filenames"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ExperimentalChooseFastestDataset.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ExperimentalChooseFastestDataset.pbtxt
index 7d52752fffe21d..2c04d58db4c0a4 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ExperimentalChooseFastestDataset.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ExperimentalChooseFastestDataset.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ExperimentalChooseFastestDataset"
   input_arg {
     name: "input_datasets"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ExperimentalDatasetCardinality.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ExperimentalDatasetCardinality.pbtxt
index d0e7b0934d0dd6..f6ba3657864e1a 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ExperimentalDatasetCardinality.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ExperimentalDatasetCardinality.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ExperimentalDatasetCardinality"
   input_arg {
     name: "input_dataset"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ExperimentalDatasetToTFRecord.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ExperimentalDatasetToTFRecord.pbtxt
index 76b1f8c1d0036f..0d0e46c8b39ade 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ExperimentalDatasetToTFRecord.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ExperimentalDatasetToTFRecord.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ExperimentalDatasetToTFRecord"
   input_arg {
     name: "input_dataset"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ExperimentalDenseToSparseBatchDataset.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ExperimentalDenseToSparseBatchDataset.pbtxt
index 94afefd024b797..c322ef95777609 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ExperimentalDenseToSparseBatchDataset.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ExperimentalDenseToSparseBatchDataset.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ExperimentalDenseToSparseBatchDataset"
   input_arg {
     name: "input_dataset"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ExperimentalDirectedInterleaveDataset.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ExperimentalDirectedInterleaveDataset.pbtxt
index c2214754670c61..2a877497ff29fd 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ExperimentalDirectedInterleaveDataset.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ExperimentalDirectedInterleaveDataset.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ExperimentalDirectedInterleaveDataset"
   input_arg {
     name: "selector_input_dataset"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ExperimentalGroupByReducerDataset.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ExperimentalGroupByReducerDataset.pbtxt
index 94d88c86c5fd0c..5e2fd15a22c908 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ExperimentalGroupByReducerDataset.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ExperimentalGroupByReducerDataset.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ExperimentalGroupByReducerDataset"
   input_arg {
     name: "input_dataset"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ExperimentalGroupByWindowDataset.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ExperimentalGroupByWindowDataset.pbtxt
index 25ede182ce533a..35f9c3c1a9547d 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ExperimentalGroupByWindowDataset.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ExperimentalGroupByWindowDataset.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ExperimentalGroupByWindowDataset"
   input_arg {
     name: "input_dataset"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ExperimentalIgnoreErrorsDataset.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ExperimentalIgnoreErrorsDataset.pbtxt
index 5180b62d1e57b5..8fc4e7d2c8adee 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ExperimentalIgnoreErrorsDataset.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ExperimentalIgnoreErrorsDataset.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ExperimentalIgnoreErrorsDataset"
   input_arg {
     name: "input_dataset"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ExperimentalIteratorGetDevice.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ExperimentalIteratorGetDevice.pbtxt
index d65b5f8bb6c4a0..8e1e10240f9ea7 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ExperimentalIteratorGetDevice.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ExperimentalIteratorGetDevice.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ExperimentalIteratorGetDevice"
   input_arg {
     name: "resource"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ExperimentalLMDBDataset.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ExperimentalLMDBDataset.pbtxt
index 05c9db9bd03516..e8b8694a947d95 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ExperimentalLMDBDataset.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ExperimentalLMDBDataset.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ExperimentalLMDBDataset"
   input_arg {
     name: "filenames"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ExperimentalLatencyStatsDataset.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ExperimentalLatencyStatsDataset.pbtxt
index aa401c14384320..29ba38a500c2b7 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ExperimentalLatencyStatsDataset.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ExperimentalLatencyStatsDataset.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ExperimentalLatencyStatsDataset"
   input_arg {
     name: "input_dataset"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ExperimentalMapAndBatchDataset.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ExperimentalMapAndBatchDataset.pbtxt
index 519fa71185c471..7799116408c237 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ExperimentalMapAndBatchDataset.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ExperimentalMapAndBatchDataset.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ExperimentalMapAndBatchDataset"
   input_arg {
     name: "input_dataset"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ExperimentalMapDataset.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ExperimentalMapDataset.pbtxt
index 75fef0d63ddb9f..401331b716073d 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ExperimentalMapDataset.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ExperimentalMapDataset.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ExperimentalMapDataset"
   input_arg {
     name: "input_dataset"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ExperimentalMatchingFilesDataset.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ExperimentalMatchingFilesDataset.pbtxt
index 29a30d9b257a07..45ef522bf7fdb4 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ExperimentalMatchingFilesDataset.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ExperimentalMatchingFilesDataset.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ExperimentalMatchingFilesDataset"
   input_arg {
     name: "patterns"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ExperimentalMaxIntraOpParallelismDataset.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ExperimentalMaxIntraOpParallelismDataset.pbtxt
index 004b1dd4f9067d..109f3906b31852 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ExperimentalMaxIntraOpParallelismDataset.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ExperimentalMaxIntraOpParallelismDataset.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ExperimentalMaxIntraOpParallelismDataset"
   input_arg {
     name: "input_dataset"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ExperimentalNonSerializableDataset.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ExperimentalNonSerializableDataset.pbtxt
index 72b5a2ef5b2d32..b0c45ac19941ce 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ExperimentalNonSerializableDataset.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ExperimentalNonSerializableDataset.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ExperimentalNonSerializableDataset"
   input_arg {
     name: "input_dataset"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ExperimentalParallelInterleaveDataset.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ExperimentalParallelInterleaveDataset.pbtxt
index 5d9a463a93ab4f..a90031c9dad06a 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ExperimentalParallelInterleaveDataset.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ExperimentalParallelInterleaveDataset.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ExperimentalParallelInterleaveDataset"
   input_arg {
     name: "input_dataset"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ExperimentalParseExampleDataset.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ExperimentalParseExampleDataset.pbtxt
index 3e150e33913691..44701f67286318 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ExperimentalParseExampleDataset.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ExperimentalParseExampleDataset.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ExperimentalParseExampleDataset"
   input_arg {
     name: "input_dataset"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ExperimentalPrivateThreadPoolDataset.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ExperimentalPrivateThreadPoolDataset.pbtxt
index e81bde0383e847..3098c30ba69d11 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ExperimentalPrivateThreadPoolDataset.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ExperimentalPrivateThreadPoolDataset.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ExperimentalPrivateThreadPoolDataset"
   input_arg {
     name: "input_dataset"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ExperimentalRandomDataset.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ExperimentalRandomDataset.pbtxt
index 04d661ee71297c..c3276a46df5b41 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ExperimentalRandomDataset.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ExperimentalRandomDataset.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ExperimentalRandomDataset"
   input_arg {
     name: "seed"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ExperimentalRebatchDataset.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ExperimentalRebatchDataset.pbtxt
index 2656cc09900abd..2c024741713d48 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ExperimentalRebatchDataset.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ExperimentalRebatchDataset.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ExperimentalRebatchDataset"
   input_arg {
     name: "input_dataset"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ExperimentalScanDataset.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ExperimentalScanDataset.pbtxt
index cc4fc8c0e9c249..39d42061ef58d6 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ExperimentalScanDataset.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ExperimentalScanDataset.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ExperimentalScanDataset"
   input_arg {
     name: "input_dataset"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ExperimentalSetStatsAggregatorDataset.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ExperimentalSetStatsAggregatorDataset.pbtxt
index 4a6d59b28c2e80..291597bf11b8dc 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ExperimentalSetStatsAggregatorDataset.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ExperimentalSetStatsAggregatorDataset.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ExperimentalSetStatsAggregatorDataset"
   input_arg {
     name: "input_dataset"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ExperimentalSleepDataset.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ExperimentalSleepDataset.pbtxt
index 06dbf87163fdae..806d764d9cdf12 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ExperimentalSleepDataset.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ExperimentalSleepDataset.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ExperimentalSleepDataset"
   input_arg {
     name: "input_dataset"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ExperimentalSlidingWindowDataset.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ExperimentalSlidingWindowDataset.pbtxt
index 6d81c4c97f833b..ab18f4e214c578 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ExperimentalSlidingWindowDataset.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ExperimentalSlidingWindowDataset.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ExperimentalSlidingWindowDataset"
   input_arg {
     name: "input_dataset"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ExperimentalSqlDataset.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ExperimentalSqlDataset.pbtxt
index 634a65912fa64a..f56ce488df0aba 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ExperimentalSqlDataset.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ExperimentalSqlDataset.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ExperimentalSqlDataset"
   input_arg {
     name: "driver_name"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ExperimentalStatsAggregatorHandle.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ExperimentalStatsAggregatorHandle.pbtxt
index a3493ec933b7c0..b00cadbca09498 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ExperimentalStatsAggregatorHandle.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ExperimentalStatsAggregatorHandle.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ExperimentalStatsAggregatorHandle"
   output_arg {
     name: "handle"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ExperimentalStatsAggregatorSummary.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ExperimentalStatsAggregatorSummary.pbtxt
index 87f63b282fe5f9..7886f7a6cb305e 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ExperimentalStatsAggregatorSummary.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ExperimentalStatsAggregatorSummary.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ExperimentalStatsAggregatorSummary"
   input_arg {
     name: "iterator"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ExperimentalTakeWhileDataset.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ExperimentalTakeWhileDataset.pbtxt
index db64c7eb605f44..7c9b4f86adbbe4 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ExperimentalTakeWhileDataset.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ExperimentalTakeWhileDataset.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ExperimentalTakeWhileDataset"
   input_arg {
     name: "input_dataset"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ExperimentalThreadPoolDataset.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ExperimentalThreadPoolDataset.pbtxt
index dc50cd1b975c71..da23c415fd24f0 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ExperimentalThreadPoolDataset.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ExperimentalThreadPoolDataset.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ExperimentalThreadPoolDataset"
   input_arg {
     name: "input_dataset"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ExperimentalThreadPoolHandle.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ExperimentalThreadPoolHandle.pbtxt
index 47dc873ee476fb..8b230f90470f29 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ExperimentalThreadPoolHandle.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ExperimentalThreadPoolHandle.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ExperimentalThreadPoolHandle"
   output_arg {
     name: "handle"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ExperimentalUnbatchDataset.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ExperimentalUnbatchDataset.pbtxt
index 42819a367ee74e..83f3a39f5e9244 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ExperimentalUnbatchDataset.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ExperimentalUnbatchDataset.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ExperimentalUnbatchDataset"
   input_arg {
     name: "input_dataset"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ExperimentalUniqueDataset.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ExperimentalUniqueDataset.pbtxt
index 4aa39a71b2d8bc..95668c930d7269 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ExperimentalUniqueDataset.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ExperimentalUniqueDataset.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ExperimentalUniqueDataset"
   input_arg {
     name: "input_dataset"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/Expint.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/Expint.pbtxt
index afb722b09ac2c3..3080bf15de9170 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/Expint.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/Expint.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "Expint"
   input_arg {
     name: "x"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/Expm1.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/Expm1.pbtxt
index 6ea265efc65ffe..b09aac454d000e 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/Expm1.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/Expm1.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "Expm1"
   input_arg {
     name: "x"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ExtractGlimpse.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ExtractGlimpse.pbtxt
index 1b305eb38867d2..597a77a3f3b8fa 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ExtractGlimpse.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ExtractGlimpse.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ExtractGlimpse"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ExtractGlimpseV2.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ExtractGlimpseV2.pbtxt
index 8dababaccb9423..08725f4504ce01 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ExtractGlimpseV2.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ExtractGlimpseV2.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ExtractGlimpseV2"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ExtractImagePatches.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ExtractImagePatches.pbtxt
index ada7acbbf57863..dee8034d6c7076 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ExtractImagePatches.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ExtractImagePatches.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ExtractImagePatches"
   input_arg {
     name: "images"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ExtractJpegShape.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ExtractJpegShape.pbtxt
index ef8068e1e91f00..ac3d34ca234fea 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ExtractJpegShape.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ExtractJpegShape.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ExtractJpegShape"
   input_arg {
     name: "contents"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ExtractVolumePatches.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ExtractVolumePatches.pbtxt
index 96d79b25ea52eb..09cc21a38b47e6 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ExtractVolumePatches.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ExtractVolumePatches.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ExtractVolumePatches"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/FFT.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/FFT.pbtxt
index 1bc74e426a247f..e986f323936f51 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/FFT.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/FFT.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "FFT"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/FFT2D.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/FFT2D.pbtxt
index 38c837f8e91a87..adb1c253867f22 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/FFT2D.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/FFT2D.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "FFT2D"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/FFT3D.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/FFT3D.pbtxt
index df20d1970b2bd3..9266d6db4a688f 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/FFT3D.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/FFT3D.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "FFT3D"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/FFTND.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/FFTND.pbtxt
index 1cb40fa334b2f2..20afcc3c5b466d 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/FFTND.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/FFTND.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "FFTND"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/FIFOQueue.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/FIFOQueue.pbtxt
index d3d57195e72efd..c3321a8c6e4782 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/FIFOQueue.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/FIFOQueue.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "FIFOQueue"
   output_arg {
     name: "handle"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/FIFOQueueV2.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/FIFOQueueV2.pbtxt
index 6f9c48db9ebcb7..9b1c8404d0ddac 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/FIFOQueueV2.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/FIFOQueueV2.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "FIFOQueueV2"
   output_arg {
     name: "handle"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/Fact.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/Fact.pbtxt
index 426124e73678e1..90a0ad8dd00112 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/Fact.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/Fact.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "Fact"
   output_arg {
     name: "fact"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/FakeParam.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/FakeParam.pbtxt
index 7b286c4697e1d3..dc2a7c5ea46991 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/FakeParam.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/FakeParam.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "FakeParam"
   output_arg {
     name: "output"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/FakeQuantWithMinMaxArgs.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/FakeQuantWithMinMaxArgs.pbtxt
index 5faa52d9136ca7..2d8eac83c59be4 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/FakeQuantWithMinMaxArgs.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/FakeQuantWithMinMaxArgs.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "FakeQuantWithMinMaxArgs"
   input_arg {
     name: "inputs"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/FakeQuantWithMinMaxArgsGradient.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/FakeQuantWithMinMaxArgsGradient.pbtxt
index 02850fbbc9b4d1..5d02f59da1e444 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/FakeQuantWithMinMaxArgsGradient.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/FakeQuantWithMinMaxArgsGradient.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "FakeQuantWithMinMaxArgsGradient"
   input_arg {
     name: "gradients"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/FakeQuantWithMinMaxVars.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/FakeQuantWithMinMaxVars.pbtxt
index 8fc979749d448a..233f5cc2f66134 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/FakeQuantWithMinMaxVars.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/FakeQuantWithMinMaxVars.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "FakeQuantWithMinMaxVars"
   input_arg {
     name: "inputs"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/FakeQuantWithMinMaxVarsGradient.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/FakeQuantWithMinMaxVarsGradient.pbtxt
index 595d2688c7b99d..cf8ed6f8b7e18d 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/FakeQuantWithMinMaxVarsGradient.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/FakeQuantWithMinMaxVarsGradient.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "FakeQuantWithMinMaxVarsGradient"
   input_arg {
     name: "gradients"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/FakeQuantWithMinMaxVarsPerChannel.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/FakeQuantWithMinMaxVarsPerChannel.pbtxt
index 7300a1daf33c8a..551ae79cd94ce9 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/FakeQuantWithMinMaxVarsPerChannel.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/FakeQuantWithMinMaxVarsPerChannel.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "FakeQuantWithMinMaxVarsPerChannel"
   input_arg {
     name: "inputs"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/FakeQuantWithMinMaxVarsPerChannelGradient.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/FakeQuantWithMinMaxVarsPerChannelGradient.pbtxt
index 80e037e94e46f3..a787e251c60a6e 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/FakeQuantWithMinMaxVarsPerChannelGradient.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/FakeQuantWithMinMaxVarsPerChannelGradient.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "FakeQuantWithMinMaxVarsPerChannelGradient"
   input_arg {
     name: "gradients"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/FakeQueue.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/FakeQueue.pbtxt
index 1a45d2509e62a4..5e4cb62d941eb9 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/FakeQueue.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/FakeQueue.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "FakeQueue"
   input_arg {
     name: "resource"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/FileSystemSetConfiguration.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/FileSystemSetConfiguration.pbtxt
index 06b193366fea68..95d26fc09818f5 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/FileSystemSetConfiguration.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/FileSystemSetConfiguration.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "FileSystemSetConfiguration"
   input_arg {
     name: "scheme"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/Fill.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/Fill.pbtxt
index e0859caefb7b41..543ae42239b4c3 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/Fill.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/Fill.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "Fill"
   input_arg {
     name: "dims"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/FilterByLastComponentDataset.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/FilterByLastComponentDataset.pbtxt
index 804e6f5dadc754..cf9bbc586524a9 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/FilterByLastComponentDataset.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/FilterByLastComponentDataset.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "FilterByLastComponentDataset"
   input_arg {
     name: "input_dataset"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/FilterDataset.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/FilterDataset.pbtxt
index 193fa2eb091518..aad48d7aed4f62 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/FilterDataset.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/FilterDataset.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "FilterDataset"
   input_arg {
     name: "input_dataset"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/FinalizeDataset.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/FinalizeDataset.pbtxt
index cc8c78b502d528..38e49288d662e1 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/FinalizeDataset.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/FinalizeDataset.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "FinalizeDataset"
   input_arg {
     name: "input_dataset"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/FinalizeTPUEmbedding.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/FinalizeTPUEmbedding.pbtxt
index 00bc279a92250a..bd367fcdd12451 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/FinalizeTPUEmbedding.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/FinalizeTPUEmbedding.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "FinalizeTPUEmbedding"
   input_arg {
     name: "common_config"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/FinalizeTPUEmbeddingV2.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/FinalizeTPUEmbeddingV2.pbtxt
index fd57e3726f6040..63c69eaff3aaba 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/FinalizeTPUEmbeddingV2.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/FinalizeTPUEmbeddingV2.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "FinalizeTPUEmbeddingV2"
   input_arg {
     name: "common_config"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/Fingerprint.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/Fingerprint.pbtxt
index b00a23b0f22f8c..3a5585701ba722 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/Fingerprint.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/Fingerprint.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "Fingerprint"
   input_arg {
     name: "data"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/FixedLengthRecordDataset.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/FixedLengthRecordDataset.pbtxt
index 04aff9c6ec69fa..63522d5f47c331 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/FixedLengthRecordDataset.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/FixedLengthRecordDataset.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "FixedLengthRecordDataset"
   input_arg {
     name: "filenames"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/FixedLengthRecordDatasetV2.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/FixedLengthRecordDatasetV2.pbtxt
index e27a084317ac5a..88d744513572c6 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/FixedLengthRecordDatasetV2.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/FixedLengthRecordDatasetV2.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "FixedLengthRecordDatasetV2"
   input_arg {
     name: "filenames"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/FixedLengthRecordReader.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/FixedLengthRecordReader.pbtxt
index 9f3aeb4fa8f374..75b6018f249c13 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/FixedLengthRecordReader.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/FixedLengthRecordReader.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "FixedLengthRecordReader"
   output_arg {
     name: "reader_handle"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/FixedLengthRecordReaderV2.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/FixedLengthRecordReaderV2.pbtxt
index 1bd086e2d92964..b16e5225240580 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/FixedLengthRecordReaderV2.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/FixedLengthRecordReaderV2.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "FixedLengthRecordReaderV2"
   output_arg {
     name: "reader_handle"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/FixedUnigramCandidateSampler.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/FixedUnigramCandidateSampler.pbtxt
index 264c48f683c7ab..a7911344886c97 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/FixedUnigramCandidateSampler.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/FixedUnigramCandidateSampler.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "FixedUnigramCandidateSampler"
   input_arg {
     name: "true_classes"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/FlatMapDataset.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/FlatMapDataset.pbtxt
index f0515e584c7b92..dcf1a7ae71c41f 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/FlatMapDataset.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/FlatMapDataset.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "FlatMapDataset"
   input_arg {
     name: "input_dataset"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/Floor.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/Floor.pbtxt
index 0653fc9864ca3f..27e405e22de85e 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/Floor.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/Floor.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "Floor"
   input_arg {
     name: "x"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/FloorDiv.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/FloorDiv.pbtxt
index 5b6aa1c29f34e1..dcaff127c7d9d0 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/FloorDiv.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/FloorDiv.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "FloorDiv"
   input_arg {
     name: "x"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/FloorMod.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/FloorMod.pbtxt
index 2f912df3b49de2..bbf48e9b826570 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/FloorMod.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/FloorMod.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "FloorMod"
   input_arg {
     name: "x"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/FlushSummaryWriter.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/FlushSummaryWriter.pbtxt
index c7c659b79f6d5c..f928d4abe99a55 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/FlushSummaryWriter.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/FlushSummaryWriter.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "FlushSummaryWriter"
   input_arg {
     name: "writer"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/For.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/For.pbtxt
index d0cdc427284e7f..139990f3994b84 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/For.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/For.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "For"
   input_arg {
     name: "start"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/FractionalAvgPool.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/FractionalAvgPool.pbtxt
index d503b56852ad5d..5fc527b066de60 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/FractionalAvgPool.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/FractionalAvgPool.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "FractionalAvgPool"
   input_arg {
     name: "value"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/FractionalAvgPoolGrad.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/FractionalAvgPoolGrad.pbtxt
index 56028951959e56..cceb2fe903ab05 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/FractionalAvgPoolGrad.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/FractionalAvgPoolGrad.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "FractionalAvgPoolGrad"
   input_arg {
     name: "orig_input_tensor_shape"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/FractionalMaxPool.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/FractionalMaxPool.pbtxt
index 4360d2356fc979..a11b4ef05f1bc9 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/FractionalMaxPool.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/FractionalMaxPool.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "FractionalMaxPool"
   input_arg {
     name: "value"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/FractionalMaxPoolGrad.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/FractionalMaxPoolGrad.pbtxt
index 517e9117186358..711e98a5df1479 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/FractionalMaxPoolGrad.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/FractionalMaxPoolGrad.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "FractionalMaxPoolGrad"
   input_arg {
     name: "orig_input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/FresnelCos.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/FresnelCos.pbtxt
index 19d0bf180e584b..7be5bbcb2ff8d5 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/FresnelCos.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/FresnelCos.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "FresnelCos"
   input_arg {
     name: "x"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/FresnelSin.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/FresnelSin.pbtxt
index a05ab41d50ac8e..c8c91ba6a68b8d 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/FresnelSin.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/FresnelSin.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "FresnelSin"
   input_arg {
     name: "x"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/FusedBatchNorm.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/FusedBatchNorm.pbtxt
index f2480179c00962..e5ac169b31ef96 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/FusedBatchNorm.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/FusedBatchNorm.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "FusedBatchNorm"
   input_arg {
     name: "x"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/FusedBatchNormGrad.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/FusedBatchNormGrad.pbtxt
index 9aabfafdd18be8..bff7eecf0ce852 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/FusedBatchNormGrad.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/FusedBatchNormGrad.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "FusedBatchNormGrad"
   input_arg {
     name: "y_backprop"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/FusedBatchNormGradV2.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/FusedBatchNormGradV2.pbtxt
index 78e2509276d1e5..dea20af8afcc6e 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/FusedBatchNormGradV2.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/FusedBatchNormGradV2.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "FusedBatchNormGradV2"
   input_arg {
     name: "y_backprop"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/FusedBatchNormGradV3.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/FusedBatchNormGradV3.pbtxt
index 2d2d3234b91df4..aa05a575bfeec0 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/FusedBatchNormGradV3.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/FusedBatchNormGradV3.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "FusedBatchNormGradV3"
   input_arg {
     name: "y_backprop"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/FusedBatchNormV2.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/FusedBatchNormV2.pbtxt
index 1f8fe62eb8b782..99f482fc721d6d 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/FusedBatchNormV2.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/FusedBatchNormV2.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "FusedBatchNormV2"
   input_arg {
     name: "x"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/FusedBatchNormV3.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/FusedBatchNormV3.pbtxt
index 1afdb5c8875269..b1f608dbe3659d 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/FusedBatchNormV3.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/FusedBatchNormV3.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "FusedBatchNormV3"
   input_arg {
     name: "x"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/FusedPadConv2D.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/FusedPadConv2D.pbtxt
index b4142aea4fbbf9..1a89c018c4e059 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/FusedPadConv2D.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/FusedPadConv2D.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "FusedPadConv2D"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/FusedResizeAndPadConv2D.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/FusedResizeAndPadConv2D.pbtxt
index b5870f695e766a..cfc716fa1d695c 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/FusedResizeAndPadConv2D.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/FusedResizeAndPadConv2D.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "FusedResizeAndPadConv2D"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/GRUBlockCell.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/GRUBlockCell.pbtxt
index 7c4519f6d7191f..7c0dd9d5fcff40 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/GRUBlockCell.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/GRUBlockCell.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "GRUBlockCell"
   input_arg {
     name: "x"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/GRUBlockCellGrad.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/GRUBlockCellGrad.pbtxt
index a8b3b505af37d7..723bcbd0b6f49c 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/GRUBlockCellGrad.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/GRUBlockCellGrad.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "GRUBlockCellGrad"
   input_arg {
     name: "x"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/Gather.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/Gather.pbtxt
index 37b2ae11c7eeb4..264a8366bb82b8 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/Gather.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/Gather.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "Gather"
   input_arg {
     name: "params"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/GatherNd.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/GatherNd.pbtxt
index 78376f9baff7b4..5ec2fd92ae606e 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/GatherNd.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/GatherNd.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "GatherNd"
   input_arg {
     name: "params"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/GatherV2.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/GatherV2.pbtxt
index 607cb850411e0d..891e82e5fb3ed6 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/GatherV2.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/GatherV2.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "GatherV2"
   input_arg {
     name: "params"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/GenerateBoundingBoxProposals.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/GenerateBoundingBoxProposals.pbtxt
index 6cb47345d8ccb7..adbc9d4c8e00f3 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/GenerateBoundingBoxProposals.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/GenerateBoundingBoxProposals.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "GenerateBoundingBoxProposals"
   input_arg {
     name: "scores"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/GenerateVocabRemapping.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/GenerateVocabRemapping.pbtxt
index adb2f799c542d0..a095253dbb20b4 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/GenerateVocabRemapping.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/GenerateVocabRemapping.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "GenerateVocabRemapping"
   input_arg {
     name: "new_vocab_file"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/GeneratorDataset.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/GeneratorDataset.pbtxt
index b9ac804e012177..9f8da9c542648e 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/GeneratorDataset.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/GeneratorDataset.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "GeneratorDataset"
   input_arg {
     name: "init_func_other_args"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/GetElementAtIndex.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/GetElementAtIndex.pbtxt
index 22b27a5530c826..82121301bd9fbb 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/GetElementAtIndex.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/GetElementAtIndex.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "GetElementAtIndex"
   input_arg {
     name: "dataset"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/GetMinibatchSplitsWithPhysicalReplica.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/GetMinibatchSplitsWithPhysicalReplica.pbtxt
index 764a8982073313..85bc30bf6dc098 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/GetMinibatchSplitsWithPhysicalReplica.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/GetMinibatchSplitsWithPhysicalReplica.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "GetMinibatchSplitsWithPhysicalReplica"
   input_arg {
     name: "program_key"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/GetMinibatchesInCsrWithPhysicalReplica.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/GetMinibatchesInCsrWithPhysicalReplica.pbtxt
index 45c3b531422fc9..79632c44322921 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/GetMinibatchesInCsrWithPhysicalReplica.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/GetMinibatchesInCsrWithPhysicalReplica.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "GetMinibatchesInCsrWithPhysicalReplica"
   input_arg {
     name: "program_key"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/GetOptions.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/GetOptions.pbtxt
index 1562e89627eddd..f8f161238d5630 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/GetOptions.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/GetOptions.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "GetOptions"
   input_arg {
     name: "input_dataset"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/GetSessionHandle.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/GetSessionHandle.pbtxt
index 90e66c7bd59690..e5345ec6f0ea1f 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/GetSessionHandle.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/GetSessionHandle.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "GetSessionHandle"
   input_arg {
     name: "value"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/GetSessionHandleV2.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/GetSessionHandleV2.pbtxt
index 55dcf277a9e465..60405234b168b4 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/GetSessionHandleV2.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/GetSessionHandleV2.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "GetSessionHandleV2"
   input_arg {
     name: "value"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/GetSessionTensor.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/GetSessionTensor.pbtxt
index fef3d08c0a7b9f..5c4cf8af9c55dc 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/GetSessionTensor.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/GetSessionTensor.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "GetSessionTensor"
   input_arg {
     name: "handle"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/GetStatsFromListOfSparseCoreCooTensors.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/GetStatsFromListOfSparseCoreCooTensors.pbtxt
index 8a9eb07d6259e5..f2ea7cfeb6ed88 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/GetStatsFromListOfSparseCoreCooTensors.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/GetStatsFromListOfSparseCoreCooTensors.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "GetStatsFromListOfSparseCoreCooTensors"
   input_arg {
     name: "row_ids_list"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/GetTpuTaskId.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/GetTpuTaskId.pbtxt
index dab64d5df42503..a5c46e67ad3e0e 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/GetTpuTaskId.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/GetTpuTaskId.pbtxt
@@ -4,5 +4,4 @@ op {
     name: "tpu_task_id"
     type: DT_INT32
   }
-  is_stateful: false
 }
diff --git a/tensorflow/core/ops/compat/ops_history_v2/GlobalIterId.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/GlobalIterId.pbtxt
index 5fa2302622c9ac..9dcfcb49a6d0cb 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/GlobalIterId.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/GlobalIterId.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "GlobalIterId"
   output_arg {
     name: "iter_id"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/GlobalShuffleDataset.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/GlobalShuffleDataset.pbtxt
index 131281a80ec590..0bce0beea54070 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/GlobalShuffleDataset.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/GlobalShuffleDataset.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "GlobalShuffleDataset"
   input_arg {
     name: "input_dataset"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/Greater.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/Greater.pbtxt
index bbf5f46aab7d53..8860e3c0c1097c 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/Greater.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/Greater.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "Greater"
   input_arg {
     name: "x"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/GreaterEqual.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/GreaterEqual.pbtxt
index db2bbd30ff2ce3..5bcdd3789c789a 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/GreaterEqual.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/GreaterEqual.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "GreaterEqual"
   input_arg {
     name: "x"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/GroupByReducerDataset.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/GroupByReducerDataset.pbtxt
index d550546eb64bd1..320e628f8aabba 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/GroupByReducerDataset.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/GroupByReducerDataset.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "GroupByReducerDataset"
   input_arg {
     name: "input_dataset"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/GroupByWindowDataset.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/GroupByWindowDataset.pbtxt
index da3bfd2542624d..0de0de53d0e7d3 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/GroupByWindowDataset.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/GroupByWindowDataset.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "GroupByWindowDataset"
   input_arg {
     name: "input_dataset"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/GuaranteeConst.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/GuaranteeConst.pbtxt
index 9ec864b74c024b..71d47e3758090e 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/GuaranteeConst.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/GuaranteeConst.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "GuaranteeConst"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/HSVToRGB.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/HSVToRGB.pbtxt
index 6169e00a62058a..2b209cc6547ab6 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/HSVToRGB.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/HSVToRGB.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "HSVToRGB"
   input_arg {
     name: "images"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/HashTable.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/HashTable.pbtxt
index eb3a57246ed277..83afe2b9448103 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/HashTable.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/HashTable.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "HashTable"
   output_arg {
     name: "table_handle"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/HashTableV2.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/HashTableV2.pbtxt
index ea6383d9f392b7..24a9bc7176d8ae 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/HashTableV2.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/HashTableV2.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "HashTableV2"
   output_arg {
     name: "table_handle"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/HistogramFixedWidth.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/HistogramFixedWidth.pbtxt
index 201df5b3467a1c..f39eabe4f72506 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/HistogramFixedWidth.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/HistogramFixedWidth.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "HistogramFixedWidth"
   input_arg {
     name: "values"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/HistogramSummary.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/HistogramSummary.pbtxt
index 6b5c4c9df1303a..0c46f397972bf6 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/HistogramSummary.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/HistogramSummary.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "HistogramSummary"
   input_arg {
     name: "tag"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/HostConst.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/HostConst.pbtxt
index 81975b57d5e37d..6dd4c1757073fd 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/HostConst.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/HostConst.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "HostConst"
   output_arg {
     name: "output"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/IFFT.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/IFFT.pbtxt
index a4c712b70cdf89..8571a132950cb1 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/IFFT.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/IFFT.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "IFFT"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/IFFT2D.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/IFFT2D.pbtxt
index f079604a4bec8c..0b208d46939354 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/IFFT2D.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/IFFT2D.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "IFFT2D"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/IFFT3D.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/IFFT3D.pbtxt
index bba5431e81392a..8b9667f882c67c 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/IFFT3D.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/IFFT3D.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "IFFT3D"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/IFFTND.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/IFFTND.pbtxt
index 72956229f0cde9..0dfb1fc9c51904 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/IFFTND.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/IFFTND.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "IFFTND"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/IRFFT.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/IRFFT.pbtxt
index 1a42324399eeb1..8ac3dfc979d8ef 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/IRFFT.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/IRFFT.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "IRFFT"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/IRFFT2D.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/IRFFT2D.pbtxt
index e348e4e8cbdf39..5d1f872d605e32 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/IRFFT2D.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/IRFFT2D.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "IRFFT2D"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/IRFFT3D.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/IRFFT3D.pbtxt
index 3a6fa0da90e102..b69417ee1a7e57 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/IRFFT3D.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/IRFFT3D.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "IRFFT3D"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/IRFFTND.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/IRFFTND.pbtxt
index c1d1da300962db..175092d5aa8158 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/IRFFTND.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/IRFFTND.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "IRFFTND"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/Identity.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/Identity.pbtxt
index 7646ed8ddb0964..f3ca3dbd24324a 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/Identity.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/Identity.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "Identity"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/IdentityN.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/IdentityN.pbtxt
index aedd30ebafdddb..61c3b63279003c 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/IdentityN.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/IdentityN.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "IdentityN"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/IdentityReader.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/IdentityReader.pbtxt
index 9e1aa0e11e219d..3330154b4d60ca 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/IdentityReader.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/IdentityReader.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "IdentityReader"
   output_arg {
     name: "reader_handle"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/IdentityReaderV2.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/IdentityReaderV2.pbtxt
index 27c72561ec1364..f37e9cedab961a 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/IdentityReaderV2.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/IdentityReaderV2.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "IdentityReaderV2"
   output_arg {
     name: "reader_handle"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/If.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/If.pbtxt
index 595f6002939acb..7ccb12afa61896 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/If.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/If.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "If"
   input_arg {
     name: "cond"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/Igamma.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/Igamma.pbtxt
index 821fa48275501f..d356f5ec5816c6 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/Igamma.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/Igamma.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "Igamma"
   input_arg {
     name: "a"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/IgammaGradA.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/IgammaGradA.pbtxt
index 739385e8c4da4f..964067de5dcb69 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/IgammaGradA.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/IgammaGradA.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "IgammaGradA"
   input_arg {
     name: "a"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/Igammac.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/Igammac.pbtxt
index 1cfa801aeeb09a..cdf44f684bbdfc 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/Igammac.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/Igammac.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "Igammac"
   input_arg {
     name: "a"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/IgnoreErrorsDataset.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/IgnoreErrorsDataset.pbtxt
index 94fed463e57f7c..32af3bb466528e 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/IgnoreErrorsDataset.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/IgnoreErrorsDataset.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "IgnoreErrorsDataset"
   input_arg {
     name: "input_dataset"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/Imag.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/Imag.pbtxt
index f8c3ce95533bf1..1444b0c60b05f6 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/Imag.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/Imag.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "Imag"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ImageProjectiveTransformV2.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ImageProjectiveTransformV2.pbtxt
index bd58faf8e16445..891b8b1cb7a6cc 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ImageProjectiveTransformV2.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ImageProjectiveTransformV2.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ImageProjectiveTransformV2"
   input_arg {
     name: "images"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ImageProjectiveTransformV3.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ImageProjectiveTransformV3.pbtxt
index 92887b2cc18151..115ee4269e7490 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ImageProjectiveTransformV3.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ImageProjectiveTransformV3.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ImageProjectiveTransformV3"
   input_arg {
     name: "images"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ImageSummary.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ImageSummary.pbtxt
index cbe4a0123f2eda..fafd7173195e3d 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ImageSummary.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ImageSummary.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ImageSummary"
   input_arg {
     name: "tag"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ImmutableConst.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ImmutableConst.pbtxt
index 51bb4050b9f65c..ba1180951f7083 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ImmutableConst.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ImmutableConst.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ImmutableConst"
   output_arg {
     name: "tensor"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ImportEvent.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ImportEvent.pbtxt
index d80a4f171a4090..7be31dd0ae5844 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ImportEvent.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ImportEvent.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ImportEvent"
   input_arg {
     name: "writer"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/InTopK.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/InTopK.pbtxt
index 881ead0a3631b5..6acd3b62e91d28 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/InTopK.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/InTopK.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "InTopK"
   input_arg {
     name: "predictions"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/InTopKV2.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/InTopKV2.pbtxt
index 003655ca052f60..a6ca2b83a45a37 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/InTopKV2.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/InTopKV2.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "InTopKV2"
   input_arg {
     name: "predictions"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/IndexFlatMapDataset.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/IndexFlatMapDataset.pbtxt
index 348e21901d1743..e28bead11f8c03 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/IndexFlatMapDataset.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/IndexFlatMapDataset.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "IndexFlatMapDataset"
   input_arg {
     name: "input_dataset"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/InfeedDequeue.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/InfeedDequeue.pbtxt
index 595fbefc51366c..a48d840da663f3 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/InfeedDequeue.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/InfeedDequeue.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "InfeedDequeue"
   output_arg {
     name: "output"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/InfeedDequeueTuple.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/InfeedDequeueTuple.pbtxt
index e3c4cdc4a0f5bb..dc6ab2b0b66476 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/InfeedDequeueTuple.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/InfeedDequeueTuple.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "InfeedDequeueTuple"
   output_arg {
     name: "outputs"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/InfeedEnqueue.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/InfeedEnqueue.pbtxt
index f62d220e778769..759b91401e9de0 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/InfeedEnqueue.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/InfeedEnqueue.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "InfeedEnqueue"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/InfeedEnqueuePrelinearizedBuffer.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/InfeedEnqueuePrelinearizedBuffer.pbtxt
index e6dfae7dcea10c..d281b700bd4aad 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/InfeedEnqueuePrelinearizedBuffer.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/InfeedEnqueuePrelinearizedBuffer.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "InfeedEnqueuePrelinearizedBuffer"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/InfeedEnqueueTuple.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/InfeedEnqueueTuple.pbtxt
index f1339fc123d0b6..459c5d9218fd6a 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/InfeedEnqueueTuple.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/InfeedEnqueueTuple.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "InfeedEnqueueTuple"
   input_arg {
     name: "inputs"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/InitializeTable.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/InitializeTable.pbtxt
index 571557348711b2..35a46a99c24987 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/InitializeTable.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/InitializeTable.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "InitializeTable"
   input_arg {
     name: "table_handle"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/InitializeTableFromDataset.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/InitializeTableFromDataset.pbtxt
index 69dc6753d84f1a..fe0ec4d4b176c2 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/InitializeTableFromDataset.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/InitializeTableFromDataset.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "InitializeTableFromDataset"
   input_arg {
     name: "table_handle"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/InitializeTableFromTextFile.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/InitializeTableFromTextFile.pbtxt
index 08c63cbe8d4ced..77be4cadf98245 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/InitializeTableFromTextFile.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/InitializeTableFromTextFile.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "InitializeTableFromTextFile"
   input_arg {
     name: "table_handle"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/InitializeTableFromTextFileV2.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/InitializeTableFromTextFileV2.pbtxt
index 69b41a04dc866e..6593434f2a6f4d 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/InitializeTableFromTextFileV2.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/InitializeTableFromTextFileV2.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "InitializeTableFromTextFileV2"
   input_arg {
     name: "table_handle"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/InitializeTableV2.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/InitializeTableV2.pbtxt
index 6e7aeb5f95946e..62c565902faf69 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/InitializeTableV2.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/InitializeTableV2.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "InitializeTableV2"
   input_arg {
     name: "table_handle"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/InplaceAdd.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/InplaceAdd.pbtxt
index 31799cddf4fee9..7c6685770b7f3f 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/InplaceAdd.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/InplaceAdd.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "InplaceAdd"
   input_arg {
     name: "x"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/InplaceSub.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/InplaceSub.pbtxt
index 28754d0755e050..42d6c14a586c49 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/InplaceSub.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/InplaceSub.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "InplaceSub"
   input_arg {
     name: "x"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/InplaceUpdate.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/InplaceUpdate.pbtxt
index 448bff4d381669..94b7f24aecc2ca 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/InplaceUpdate.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/InplaceUpdate.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "InplaceUpdate"
   input_arg {
     name: "x"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/InterleaveDataset.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/InterleaveDataset.pbtxt
index 9a2505d5540d19..124a84cf82fc55 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/InterleaveDataset.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/InterleaveDataset.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "InterleaveDataset"
   input_arg {
     name: "input_dataset"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/Inv.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/Inv.pbtxt
index 30ed71735243f7..0c191790030e8d 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/Inv.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/Inv.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "Inv"
   input_arg {
     name: "x"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/InvGrad.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/InvGrad.pbtxt
index 1c1d4971bf43f5..af882a90b23f58 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/InvGrad.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/InvGrad.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "InvGrad"
   input_arg {
     name: "x"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/Invert.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/Invert.pbtxt
index 6915b9d38fd471..cd9c81231790c4 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/Invert.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/Invert.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "Invert"
   input_arg {
     name: "x"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/InvertPermutation.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/InvertPermutation.pbtxt
index 74c7d64b10ba81..fa028961e3220c 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/InvertPermutation.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/InvertPermutation.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "InvertPermutation"
   input_arg {
     name: "x"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/IsBoostedTreesEnsembleInitialized.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/IsBoostedTreesEnsembleInitialized.pbtxt
index 4086fdf15d3ffa..1b19fef0df2c93 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/IsBoostedTreesEnsembleInitialized.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/IsBoostedTreesEnsembleInitialized.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "IsBoostedTreesEnsembleInitialized"
   input_arg {
     name: "tree_ensemble_handle"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/IsBoostedTreesQuantileStreamResourceInitialized.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/IsBoostedTreesQuantileStreamResourceInitialized.pbtxt
index 868a34f393b315..359e0e9ba5798c 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/IsBoostedTreesQuantileStreamResourceInitialized.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/IsBoostedTreesQuantileStreamResourceInitialized.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "IsBoostedTreesQuantileStreamResourceInitialized"
   input_arg {
     name: "quantile_stream_resource_handle"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/IsFinite.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/IsFinite.pbtxt
index d3051d645ba2b1..8410dce0cb011f 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/IsFinite.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/IsFinite.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "IsFinite"
   input_arg {
     name: "x"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/IsInf.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/IsInf.pbtxt
index 20604d6e930955..1ce6c74691e3d6 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/IsInf.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/IsInf.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "IsInf"
   input_arg {
     name: "x"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/IsNan.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/IsNan.pbtxt
index bc6c66435d00e0..826f2fff6c507d 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/IsNan.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/IsNan.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "IsNan"
   input_arg {
     name: "x"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/IsTPUEmbeddingInitialized.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/IsTPUEmbeddingInitialized.pbtxt
index 6a6a893bfddcd4..bbfe80cc97ac21 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/IsTPUEmbeddingInitialized.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/IsTPUEmbeddingInitialized.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "IsTPUEmbeddingInitialized"
   output_arg {
     name: "is_tpu_embedding_initialized"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/IsVariableInitialized.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/IsVariableInitialized.pbtxt
index c61e3263a6aeb9..03496db8d32030 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/IsVariableInitialized.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/IsVariableInitialized.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "IsVariableInitialized"
   input_arg {
     name: "ref"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/IsotonicRegression.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/IsotonicRegression.pbtxt
index 3c9050f67f2607..abe6fb4bbd849a 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/IsotonicRegression.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/IsotonicRegression.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "IsotonicRegression"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/Iterator.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/Iterator.pbtxt
index a9b8b66ee19046..76b9fdef4ed26e 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/Iterator.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/Iterator.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "Iterator"
   output_arg {
     name: "handle"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/IteratorFromStringHandle.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/IteratorFromStringHandle.pbtxt
index e4964182ad4259..ebd34378194d2d 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/IteratorFromStringHandle.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/IteratorFromStringHandle.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "IteratorFromStringHandle"
   input_arg {
     name: "string_handle"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/IteratorFromStringHandleV2.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/IteratorFromStringHandleV2.pbtxt
index 3faf956c0b6030..624c47394db730 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/IteratorFromStringHandleV2.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/IteratorFromStringHandleV2.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "IteratorFromStringHandleV2"
   input_arg {
     name: "string_handle"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/IteratorGetDevice.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/IteratorGetDevice.pbtxt
index 02020169c71f71..8d379c1557b2ab 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/IteratorGetDevice.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/IteratorGetDevice.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "IteratorGetDevice"
   input_arg {
     name: "resource"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/IteratorGetModelProto.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/IteratorGetModelProto.pbtxt
index ee47f9d7d43634..b1343becfa9bd6 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/IteratorGetModelProto.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/IteratorGetModelProto.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "IteratorGetModelProto"
   input_arg {
     name: "iterator"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/IteratorGetNext.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/IteratorGetNext.pbtxt
index daddc26f37be32..f204011ed431b7 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/IteratorGetNext.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/IteratorGetNext.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "IteratorGetNext"
   input_arg {
     name: "iterator"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/IteratorGetNextAsOptional.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/IteratorGetNextAsOptional.pbtxt
index 5cbd8314b9f0e7..c1a532f165c6f0 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/IteratorGetNextAsOptional.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/IteratorGetNextAsOptional.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "IteratorGetNextAsOptional"
   input_arg {
     name: "iterator"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/IteratorGetNextSync.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/IteratorGetNextSync.pbtxt
index c7f29d2773cfad..e1a7351d2da357 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/IteratorGetNextSync.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/IteratorGetNextSync.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "IteratorGetNextSync"
   input_arg {
     name: "iterator"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/IteratorToStringHandle.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/IteratorToStringHandle.pbtxt
index 47fa0764034917..87f2dffc941b41 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/IteratorToStringHandle.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/IteratorToStringHandle.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "IteratorToStringHandle"
   input_arg {
     name: "resource_handle"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/IteratorV2.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/IteratorV2.pbtxt
index 244fe9d0a49c62..6f7ab705485da4 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/IteratorV2.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/IteratorV2.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "IteratorV2"
   output_arg {
     name: "handle"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/KMC2ChainInitialization.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/KMC2ChainInitialization.pbtxt
index f695595658365b..e9640975b0e685 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/KMC2ChainInitialization.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/KMC2ChainInitialization.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "KMC2ChainInitialization"
   input_arg {
     name: "distances"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/KmeansPlusPlusInitialization.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/KmeansPlusPlusInitialization.pbtxt
index ea5914ba22f17c..27ab4b34885bc9 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/KmeansPlusPlusInitialization.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/KmeansPlusPlusInitialization.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "KmeansPlusPlusInitialization"
   input_arg {
     name: "points"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/KthOrderStatistic.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/KthOrderStatistic.pbtxt
index 3b0ca9d2d07479..8e5b79cec0442b 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/KthOrderStatistic.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/KthOrderStatistic.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "KthOrderStatistic"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/L2Loss.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/L2Loss.pbtxt
index 1e72b270d39aee..90e8619d09f2e7 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/L2Loss.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/L2Loss.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "L2Loss"
   input_arg {
     name: "t"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/LMDBDataset.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/LMDBDataset.pbtxt
index 6609f2d64fd84f..9ba1bd98191f8e 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/LMDBDataset.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/LMDBDataset.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "LMDBDataset"
   input_arg {
     name: "filenames"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/LMDBReader.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/LMDBReader.pbtxt
index b6f52544c52965..967c74bb72c778 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/LMDBReader.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/LMDBReader.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "LMDBReader"
   output_arg {
     name: "reader_handle"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/LRN.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/LRN.pbtxt
index aba0c94025dd81..75880682c31830 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/LRN.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/LRN.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "LRN"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/LRNGrad.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/LRNGrad.pbtxt
index 65a6c221aaeff0..37db775eaa236b 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/LRNGrad.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/LRNGrad.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "LRNGrad"
   input_arg {
     name: "input_grads"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/LSTMBlockCell.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/LSTMBlockCell.pbtxt
index 88643e69b1c1ae..f1071f7fc51699 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/LSTMBlockCell.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/LSTMBlockCell.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "LSTMBlockCell"
   input_arg {
     name: "x"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/LSTMBlockCellGrad.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/LSTMBlockCellGrad.pbtxt
index 1851e5bd6e6bfb..b20d47c5c01ef5 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/LSTMBlockCellGrad.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/LSTMBlockCellGrad.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "LSTMBlockCellGrad"
   input_arg {
     name: "x"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/LatencyStatsDataset.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/LatencyStatsDataset.pbtxt
index cdb1716e83521e..546bcdcbd1233e 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/LatencyStatsDataset.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/LatencyStatsDataset.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "LatencyStatsDataset"
   input_arg {
     name: "input_dataset"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/LeakyRelu.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/LeakyRelu.pbtxt
index 4cbb4fdf427b38..c0358f96a87f10 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/LeakyRelu.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/LeakyRelu.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "LeakyRelu"
   input_arg {
     name: "features"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/LeakyReluGrad.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/LeakyReluGrad.pbtxt
index 524638322a3f62..786872202c456d 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/LeakyReluGrad.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/LeakyReluGrad.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "LeakyReluGrad"
   input_arg {
     name: "gradients"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/LearnedUnigramCandidateSampler.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/LearnedUnigramCandidateSampler.pbtxt
index cce62fa5d0cc4b..71466c56726990 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/LearnedUnigramCandidateSampler.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/LearnedUnigramCandidateSampler.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "LearnedUnigramCandidateSampler"
   input_arg {
     name: "true_classes"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/LeftShift.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/LeftShift.pbtxt
index 94088ea60d6d12..c3f56bee3bb47a 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/LeftShift.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/LeftShift.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "LeftShift"
   input_arg {
     name: "x"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/LegacyParallelInterleaveDatasetV2.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/LegacyParallelInterleaveDatasetV2.pbtxt
index 6d4d712d004182..49f6a5574721c8 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/LegacyParallelInterleaveDatasetV2.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/LegacyParallelInterleaveDatasetV2.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "LegacyParallelInterleaveDatasetV2"
   input_arg {
     name: "input_dataset"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/Less.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/Less.pbtxt
index 818d92f302140e..e4f12455aa5ae3 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/Less.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/Less.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "Less"
   input_arg {
     name: "x"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/LessEqual.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/LessEqual.pbtxt
index 477d8bf1d92fdd..9162a684069fd8 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/LessEqual.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/LessEqual.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "LessEqual"
   input_arg {
     name: "x"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/Lgamma.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/Lgamma.pbtxt
index f43959c6afe82a..fcb0241217bdfa 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/Lgamma.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/Lgamma.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "Lgamma"
   input_arg {
     name: "x"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/LinSpace.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/LinSpace.pbtxt
index ad58235a544266..044ba244206f96 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/LinSpace.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/LinSpace.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "LinSpace"
   input_arg {
     name: "start"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ListDataset.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ListDataset.pbtxt
index c095f6d1194705..2180d6d82a474c 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ListDataset.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ListDataset.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ListDataset"
   input_arg {
     name: "tensors"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ListDiff.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ListDiff.pbtxt
index 090baff95aefa0..39c3ee8606ccfd 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ListDiff.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ListDiff.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ListDiff"
   input_arg {
     name: "x"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ListSnapshotChunksDataset.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ListSnapshotChunksDataset.pbtxt
index be35470141fb28..1b667052f0e069 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ListSnapshotChunksDataset.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ListSnapshotChunksDataset.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ListSnapshotChunksDataset"
   input_arg {
     name: "snapshot_path"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/LoadAllTPUEmbeddingParameters.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/LoadAllTPUEmbeddingParameters.pbtxt
index 68a146348b2758..5ac0c1cfb2b345 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/LoadAllTPUEmbeddingParameters.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/LoadAllTPUEmbeddingParameters.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "LoadAllTPUEmbeddingParameters"
   input_arg {
     name: "parameters"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/LoadAndRemapMatrix.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/LoadAndRemapMatrix.pbtxt
index be16dec44ee721..54b4a68f2b344b 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/LoadAndRemapMatrix.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/LoadAndRemapMatrix.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "LoadAndRemapMatrix"
   input_arg {
     name: "ckpt_path"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/LoadDataset.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/LoadDataset.pbtxt
index 41e306f86429a9..c46f54ef3c53a1 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/LoadDataset.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/LoadDataset.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "LoadDataset"
   input_arg {
     name: "path"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/LoadTPUEmbeddingADAMParameters.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/LoadTPUEmbeddingADAMParameters.pbtxt
index 38aec474b28889..5294493f6d1f30 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/LoadTPUEmbeddingADAMParameters.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/LoadTPUEmbeddingADAMParameters.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "LoadTPUEmbeddingADAMParameters"
   input_arg {
     name: "parameters"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/LoadTPUEmbeddingAdadeltaParameters.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/LoadTPUEmbeddingAdadeltaParameters.pbtxt
index c8d55510e0c983..c093049ee0d2f3 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/LoadTPUEmbeddingAdadeltaParameters.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/LoadTPUEmbeddingAdadeltaParameters.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "LoadTPUEmbeddingAdadeltaParameters"
   input_arg {
     name: "parameters"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/LoadTPUEmbeddingAdagradMomentumParameters.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/LoadTPUEmbeddingAdagradMomentumParameters.pbtxt
index be33451441efe9..82f2cb2161fc7a 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/LoadTPUEmbeddingAdagradMomentumParameters.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/LoadTPUEmbeddingAdagradMomentumParameters.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "LoadTPUEmbeddingAdagradMomentumParameters"
   input_arg {
     name: "parameters"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/LoadTPUEmbeddingAdagradParameters.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/LoadTPUEmbeddingAdagradParameters.pbtxt
index b4325420e97e5d..485ec861a1eae4 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/LoadTPUEmbeddingAdagradParameters.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/LoadTPUEmbeddingAdagradParameters.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "LoadTPUEmbeddingAdagradParameters"
   input_arg {
     name: "parameters"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/LoadTPUEmbeddingCenteredRMSPropParameters.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/LoadTPUEmbeddingCenteredRMSPropParameters.pbtxt
index 42b53f9cf2e60c..72c02df3763afb 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/LoadTPUEmbeddingCenteredRMSPropParameters.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/LoadTPUEmbeddingCenteredRMSPropParameters.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "LoadTPUEmbeddingCenteredRMSPropParameters"
   input_arg {
     name: "parameters"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/LoadTPUEmbeddingFTRLParameters.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/LoadTPUEmbeddingFTRLParameters.pbtxt
index 0bceec9c7897de..a429697249129a 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/LoadTPUEmbeddingFTRLParameters.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/LoadTPUEmbeddingFTRLParameters.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "LoadTPUEmbeddingFTRLParameters"
   input_arg {
     name: "parameters"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/LoadTPUEmbeddingFrequencyEstimatorParameters.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/LoadTPUEmbeddingFrequencyEstimatorParameters.pbtxt
index 2485f0c96ca4f5..ab6af2875d5dc1 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/LoadTPUEmbeddingFrequencyEstimatorParameters.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/LoadTPUEmbeddingFrequencyEstimatorParameters.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "LoadTPUEmbeddingFrequencyEstimatorParameters"
   input_arg {
     name: "parameters"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/LoadTPUEmbeddingMDLAdagradLightParameters.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/LoadTPUEmbeddingMDLAdagradLightParameters.pbtxt
index 2140a2b84b2b59..a175817a6e931d 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/LoadTPUEmbeddingMDLAdagradLightParameters.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/LoadTPUEmbeddingMDLAdagradLightParameters.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "LoadTPUEmbeddingMDLAdagradLightParameters"
   input_arg {
     name: "parameters"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/LoadTPUEmbeddingMomentumParameters.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/LoadTPUEmbeddingMomentumParameters.pbtxt
index 3da833e5595b2a..0f135f7c6f0e25 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/LoadTPUEmbeddingMomentumParameters.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/LoadTPUEmbeddingMomentumParameters.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "LoadTPUEmbeddingMomentumParameters"
   input_arg {
     name: "parameters"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/LoadTPUEmbeddingProximalAdagradParameters.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/LoadTPUEmbeddingProximalAdagradParameters.pbtxt
index c272880303b6ad..10d611a38b2e78 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/LoadTPUEmbeddingProximalAdagradParameters.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/LoadTPUEmbeddingProximalAdagradParameters.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "LoadTPUEmbeddingProximalAdagradParameters"
   input_arg {
     name: "parameters"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/LoadTPUEmbeddingProximalYogiParameters.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/LoadTPUEmbeddingProximalYogiParameters.pbtxt
index 7feae76245ed14..8f51ed094ae8e1 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/LoadTPUEmbeddingProximalYogiParameters.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/LoadTPUEmbeddingProximalYogiParameters.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "LoadTPUEmbeddingProximalYogiParameters"
   input_arg {
     name: "parameters"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/LoadTPUEmbeddingRMSPropParameters.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/LoadTPUEmbeddingRMSPropParameters.pbtxt
index 711f9e56dfe055..cfec4be6aa7d1c 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/LoadTPUEmbeddingRMSPropParameters.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/LoadTPUEmbeddingRMSPropParameters.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "LoadTPUEmbeddingRMSPropParameters"
   input_arg {
     name: "parameters"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/LoadTPUEmbeddingStochasticGradientDescentParameters.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/LoadTPUEmbeddingStochasticGradientDescentParameters.pbtxt
index a61ee8b1e4bc4e..48b965b4ea3691 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/LoadTPUEmbeddingStochasticGradientDescentParameters.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/LoadTPUEmbeddingStochasticGradientDescentParameters.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "LoadTPUEmbeddingStochasticGradientDescentParameters"
   input_arg {
     name: "parameters"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/Log.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/Log.pbtxt
index 7e5b08a05a1282..a16862c0735f93 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/Log.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/Log.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "Log"
   input_arg {
     name: "x"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/Log1p.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/Log1p.pbtxt
index b1f686eebe6982..1f8ba12957e726 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/Log1p.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/Log1p.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "Log1p"
   input_arg {
     name: "x"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/LogMatrixDeterminant.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/LogMatrixDeterminant.pbtxt
index aee2342ce070aa..3807cdda42595e 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/LogMatrixDeterminant.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/LogMatrixDeterminant.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "LogMatrixDeterminant"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/LogSoftmax.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/LogSoftmax.pbtxt
index 59748cdb10238b..92d2727bbe7e23 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/LogSoftmax.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/LogSoftmax.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "LogSoftmax"
   input_arg {
     name: "logits"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/LogUniformCandidateSampler.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/LogUniformCandidateSampler.pbtxt
index 1066189f03b1b3..9ec45571159093 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/LogUniformCandidateSampler.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/LogUniformCandidateSampler.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "LogUniformCandidateSampler"
   input_arg {
     name: "true_classes"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/LogicalAnd.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/LogicalAnd.pbtxt
index bdad772540320e..b10b115df4f23f 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/LogicalAnd.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/LogicalAnd.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "LogicalAnd"
   input_arg {
     name: "x"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/LogicalNot.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/LogicalNot.pbtxt
index cba489e4022795..5cf13ad8399c7b 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/LogicalNot.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/LogicalNot.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "LogicalNot"
   input_arg {
     name: "x"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/LogicalOr.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/LogicalOr.pbtxt
index 9306d3d95b357a..635a66d8ba0634 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/LogicalOr.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/LogicalOr.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "LogicalOr"
   input_arg {
     name: "x"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/LookupTableExport.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/LookupTableExport.pbtxt
index f083b589d1cd21..6c56cdeb1de748 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/LookupTableExport.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/LookupTableExport.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "LookupTableExport"
   input_arg {
     name: "table_handle"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/LookupTableExportV2.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/LookupTableExportV2.pbtxt
index 264b3867d11698..b86fd3a32ad47b 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/LookupTableExportV2.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/LookupTableExportV2.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "LookupTableExportV2"
   input_arg {
     name: "table_handle"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/LookupTableFind.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/LookupTableFind.pbtxt
index e9e33182da56de..5923b502abcb15 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/LookupTableFind.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/LookupTableFind.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "LookupTableFind"
   input_arg {
     name: "table_handle"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/LookupTableFindV2.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/LookupTableFindV2.pbtxt
index 4fc069f0e48d0d..53cbafbeee082f 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/LookupTableFindV2.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/LookupTableFindV2.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "LookupTableFindV2"
   input_arg {
     name: "table_handle"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/LookupTableImport.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/LookupTableImport.pbtxt
index c3184809ed3962..73b53a5b6f5784 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/LookupTableImport.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/LookupTableImport.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "LookupTableImport"
   input_arg {
     name: "table_handle"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/LookupTableImportV2.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/LookupTableImportV2.pbtxt
index aa684ec6b5bcb2..41c03b83c71ba1 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/LookupTableImportV2.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/LookupTableImportV2.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "LookupTableImportV2"
   input_arg {
     name: "table_handle"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/LookupTableInsert.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/LookupTableInsert.pbtxt
index b33cb4837a7bfc..b96cb478887922 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/LookupTableInsert.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/LookupTableInsert.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "LookupTableInsert"
   input_arg {
     name: "table_handle"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/LookupTableInsertV2.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/LookupTableInsertV2.pbtxt
index c2e5b1a3cb448c..19d7d49b8600a3 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/LookupTableInsertV2.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/LookupTableInsertV2.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "LookupTableInsertV2"
   input_arg {
     name: "table_handle"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/LookupTableRemoveV2.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/LookupTableRemoveV2.pbtxt
index 325b4316da3810..d7fe0bb4dd4e40 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/LookupTableRemoveV2.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/LookupTableRemoveV2.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "LookupTableRemoveV2"
   input_arg {
     name: "table_handle"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/LookupTableSize.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/LookupTableSize.pbtxt
index a878f30b9ea764..0d4bf61189fa33 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/LookupTableSize.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/LookupTableSize.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "LookupTableSize"
   input_arg {
     name: "table_handle"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/LookupTableSizeV2.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/LookupTableSizeV2.pbtxt
index 6fb6f47f945262..511beedff01e3b 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/LookupTableSizeV2.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/LookupTableSizeV2.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "LookupTableSizeV2"
   input_arg {
     name: "table_handle"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/LoopCond.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/LoopCond.pbtxt
index 3be4701ef48460..7111fff007b5f2 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/LoopCond.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/LoopCond.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "LoopCond"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/LowerBound.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/LowerBound.pbtxt
index 934c879c0907dc..b7d1dee7797707 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/LowerBound.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/LowerBound.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "LowerBound"
   input_arg {
     name: "sorted_inputs"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/Lu.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/Lu.pbtxt
index a5140b8af7a87e..59c28e09e704d4 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/Lu.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/Lu.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "Lu"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/MakeIterator.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/MakeIterator.pbtxt
index 84e12f49dba0e6..b11c2b9e1dd1b1 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/MakeIterator.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/MakeIterator.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "MakeIterator"
   input_arg {
     name: "dataset"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/MakeUnique.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/MakeUnique.pbtxt
index a6dbebb1e588a9..685f52d66eaead 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/MakeUnique.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/MakeUnique.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "MakeUnique"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/MapAndBatchDataset.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/MapAndBatchDataset.pbtxt
index d4e5d9a29a79f2..8e7b6a32493801 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/MapAndBatchDataset.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/MapAndBatchDataset.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "MapAndBatchDataset"
   input_arg {
     name: "input_dataset"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/MapClear.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/MapClear.pbtxt
index 261e8ba9b6f5b0..22c5e5fcad020b 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/MapClear.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/MapClear.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "MapClear"
   attr {
     name: "capacity"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/MapDataset.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/MapDataset.pbtxt
index 6840222b360148..b01b535e48d6fd 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/MapDataset.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/MapDataset.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "MapDataset"
   input_arg {
     name: "input_dataset"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/MapDefun.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/MapDefun.pbtxt
index 243d9866d0bbe0..7cb9d19231cfd9 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/MapDefun.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/MapDefun.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "MapDefun"
   input_arg {
     name: "arguments"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/MapIncompleteSize.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/MapIncompleteSize.pbtxt
index b6c12953967540..ca9c629887fa58 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/MapIncompleteSize.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/MapIncompleteSize.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "MapIncompleteSize"
   output_arg {
     name: "size"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/MapPeek.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/MapPeek.pbtxt
index 13973572d1cf49..4a61cb9e40f47c 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/MapPeek.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/MapPeek.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "MapPeek"
   input_arg {
     name: "key"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/MapSize.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/MapSize.pbtxt
index 5bffb06eb1ea59..6828f8fbb09a3f 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/MapSize.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/MapSize.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "MapSize"
   output_arg {
     name: "size"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/MapStage.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/MapStage.pbtxt
index c34737c065777f..4ad2131a1f1844 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/MapStage.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/MapStage.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "MapStage"
   input_arg {
     name: "key"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/MapUnstage.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/MapUnstage.pbtxt
index f288ab5f51cdb2..9901130961c8a1 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/MapUnstage.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/MapUnstage.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "MapUnstage"
   input_arg {
     name: "key"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/MapUnstageNoKey.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/MapUnstageNoKey.pbtxt
index ce11b85c5063cb..ee4cca51346065 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/MapUnstageNoKey.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/MapUnstageNoKey.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "MapUnstageNoKey"
   input_arg {
     name: "indices"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/MatMul.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/MatMul.pbtxt
index 8f79fa11000f7f..42a2c794353672 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/MatMul.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/MatMul.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "MatMul"
   input_arg {
     name: "a"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/MatchingFiles.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/MatchingFiles.pbtxt
index e374694de6f10a..3f8af5f3226375 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/MatchingFiles.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/MatchingFiles.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "MatchingFiles"
   input_arg {
     name: "pattern"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/MatchingFilesDataset.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/MatchingFilesDataset.pbtxt
index 916f7e1e60a1c5..bf43730d6dc87f 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/MatchingFilesDataset.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/MatchingFilesDataset.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "MatchingFilesDataset"
   input_arg {
     name: "patterns"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/MatrixBandPart.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/MatrixBandPart.pbtxt
index c6ab4921aa4f70..c25aa9615c51d1 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/MatrixBandPart.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/MatrixBandPart.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "MatrixBandPart"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/MatrixDeterminant.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/MatrixDeterminant.pbtxt
index 791c7f7a5afc57..4dd524d4894ed4 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/MatrixDeterminant.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/MatrixDeterminant.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "MatrixDeterminant"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/MatrixDiag.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/MatrixDiag.pbtxt
index 299678fef206a3..9b0ddb0285d0cd 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/MatrixDiag.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/MatrixDiag.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "MatrixDiag"
   input_arg {
     name: "diagonal"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/MatrixDiagPart.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/MatrixDiagPart.pbtxt
index 4952fc6b2f2acc..efb1e18fccba39 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/MatrixDiagPart.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/MatrixDiagPart.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "MatrixDiagPart"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/MatrixDiagPartV2.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/MatrixDiagPartV2.pbtxt
index 3325a0d4069336..f709c6d5eb556c 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/MatrixDiagPartV2.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/MatrixDiagPartV2.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "MatrixDiagPartV2"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/MatrixDiagPartV3.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/MatrixDiagPartV3.pbtxt
index c7df214d46ed30..75a1307f279142 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/MatrixDiagPartV3.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/MatrixDiagPartV3.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "MatrixDiagPartV3"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/MatrixDiagV2.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/MatrixDiagV2.pbtxt
index 4fafe31b76cc93..3f6aa1e6a72fa5 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/MatrixDiagV2.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/MatrixDiagV2.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "MatrixDiagV2"
   input_arg {
     name: "diagonal"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/MatrixDiagV3.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/MatrixDiagV3.pbtxt
index c8829be1bc41f8..793efcec3b6425 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/MatrixDiagV3.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/MatrixDiagV3.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "MatrixDiagV3"
   input_arg {
     name: "diagonal"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/MatrixExponential.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/MatrixExponential.pbtxt
index d173fbe0515836..008291a4cafd64 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/MatrixExponential.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/MatrixExponential.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "MatrixExponential"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/MatrixInverse.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/MatrixInverse.pbtxt
index 6b3befc4dd031f..81d35ad1d082ae 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/MatrixInverse.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/MatrixInverse.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "MatrixInverse"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/MatrixLogarithm.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/MatrixLogarithm.pbtxt
index c3cc81e38c6cb0..0a87e5905d0123 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/MatrixLogarithm.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/MatrixLogarithm.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "MatrixLogarithm"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/MatrixSetDiag.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/MatrixSetDiag.pbtxt
index 0dcfca6102cc1b..e8c08f8d295192 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/MatrixSetDiag.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/MatrixSetDiag.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "MatrixSetDiag"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/MatrixSetDiagV2.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/MatrixSetDiagV2.pbtxt
index c11740fbd69a34..1147220c00c774 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/MatrixSetDiagV2.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/MatrixSetDiagV2.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "MatrixSetDiagV2"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/MatrixSetDiagV3.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/MatrixSetDiagV3.pbtxt
index 1490839cf0b3f6..d5a6af8b119291 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/MatrixSetDiagV3.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/MatrixSetDiagV3.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "MatrixSetDiagV3"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/MatrixSolve.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/MatrixSolve.pbtxt
index 50cf802d9ff263..2a28fa0adb0377 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/MatrixSolve.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/MatrixSolve.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "MatrixSolve"
   input_arg {
     name: "matrix"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/MatrixSolveLs.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/MatrixSolveLs.pbtxt
index 082064ef822021..5df48fc28e3b29 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/MatrixSolveLs.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/MatrixSolveLs.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "MatrixSolveLs"
   input_arg {
     name: "matrix"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/MatrixSquareRoot.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/MatrixSquareRoot.pbtxt
index b7e690181cad3e..32ff859e8f9826 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/MatrixSquareRoot.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/MatrixSquareRoot.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "MatrixSquareRoot"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/MatrixTriangularSolve.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/MatrixTriangularSolve.pbtxt
index e925e1189322b4..915e582d00a0e1 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/MatrixTriangularSolve.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/MatrixTriangularSolve.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "MatrixTriangularSolve"
   input_arg {
     name: "matrix"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/Max.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/Max.pbtxt
index 91262d8e435aee..bf147acf0f405f 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/Max.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/Max.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "Max"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/MaxIntraOpParallelismDataset.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/MaxIntraOpParallelismDataset.pbtxt
index 946f2f18f84161..85547917acc6e3 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/MaxIntraOpParallelismDataset.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/MaxIntraOpParallelismDataset.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "MaxIntraOpParallelismDataset"
   input_arg {
     name: "input_dataset"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/MaxPool.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/MaxPool.pbtxt
index 2fd8174576359a..f4fd1cccf29f7d 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/MaxPool.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/MaxPool.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "MaxPool"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/MaxPool3D.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/MaxPool3D.pbtxt
index 928c2c0aeaf32b..7af4fca0e93e15 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/MaxPool3D.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/MaxPool3D.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "MaxPool3D"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/MaxPool3DGrad.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/MaxPool3DGrad.pbtxt
index 8e8f10e4fe80e7..77edcb4c89887a 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/MaxPool3DGrad.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/MaxPool3DGrad.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "MaxPool3DGrad"
   input_arg {
     name: "orig_input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/MaxPool3DGradGrad.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/MaxPool3DGradGrad.pbtxt
index 5012c312b1bf1f..55d26c13c9cab9 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/MaxPool3DGradGrad.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/MaxPool3DGradGrad.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "MaxPool3DGradGrad"
   input_arg {
     name: "orig_input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/MaxPoolGrad.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/MaxPoolGrad.pbtxt
index 2a36894500bdd8..131a3633cf98f7 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/MaxPoolGrad.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/MaxPoolGrad.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "MaxPoolGrad"
   input_arg {
     name: "orig_input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/MaxPoolGradGrad.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/MaxPoolGradGrad.pbtxt
index 9589f707f4c481..9b1f4de08ea069 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/MaxPoolGradGrad.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/MaxPoolGradGrad.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "MaxPoolGradGrad"
   input_arg {
     name: "orig_input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/MaxPoolGradGradV2.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/MaxPoolGradGradV2.pbtxt
index c716043ae9ac56..fba1ab57dc6b28 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/MaxPoolGradGradV2.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/MaxPoolGradGradV2.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "MaxPoolGradGradV2"
   input_arg {
     name: "orig_input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/MaxPoolGradGradWithArgmax.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/MaxPoolGradGradWithArgmax.pbtxt
index dc96f707c66e8a..3c3cdbb90d1287 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/MaxPoolGradGradWithArgmax.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/MaxPoolGradGradWithArgmax.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "MaxPoolGradGradWithArgmax"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/MaxPoolGradV2.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/MaxPoolGradV2.pbtxt
index 42d63ab79b1e60..7e38cf840dd0e0 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/MaxPoolGradV2.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/MaxPoolGradV2.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "MaxPoolGradV2"
   input_arg {
     name: "orig_input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/MaxPoolGradWithArgmax.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/MaxPoolGradWithArgmax.pbtxt
index 1d30e8703369a6..7c3ab4a0cd05c6 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/MaxPoolGradWithArgmax.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/MaxPoolGradWithArgmax.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "MaxPoolGradWithArgmax"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/MaxPoolV2.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/MaxPoolV2.pbtxt
index 82dc586d6d487d..3ef7da8d9d9848 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/MaxPoolV2.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/MaxPoolV2.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "MaxPoolV2"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/MaxPoolWithArgmax.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/MaxPoolWithArgmax.pbtxt
index 0a4a4df9c8135d..d33bbd2f70769e 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/MaxPoolWithArgmax.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/MaxPoolWithArgmax.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "MaxPoolWithArgmax"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/Maximum.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/Maximum.pbtxt
index 32bbcb645478ef..c6ee10114c63f0 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/Maximum.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/Maximum.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "Maximum"
   input_arg {
     name: "x"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/Mean.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/Mean.pbtxt
index 10769852ec7d10..e0b5f145616867 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/Mean.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/Mean.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "Mean"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/Merge.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/Merge.pbtxt
index 23610f501a46a0..d08f9cc55e94f0 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/Merge.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/Merge.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "Merge"
   input_arg {
     name: "inputs"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/MergeDedupData.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/MergeDedupData.pbtxt
index 147341c8f374af..a5bcb48c34a0f1 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/MergeDedupData.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/MergeDedupData.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "MergeDedupData"
   input_arg {
     name: "integer_tensor"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/MergeSummary.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/MergeSummary.pbtxt
index ffe893e53bfa8a..d9b14d4511093e 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/MergeSummary.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/MergeSummary.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "MergeSummary"
   input_arg {
     name: "inputs"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/MergeV2Checkpoints.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/MergeV2Checkpoints.pbtxt
index 96b092e55660ac..2a6b60b4e531c9 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/MergeV2Checkpoints.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/MergeV2Checkpoints.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "MergeV2Checkpoints"
   input_arg {
     name: "checkpoint_prefixes"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/Mfcc.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/Mfcc.pbtxt
index 8f4577250c9f3d..4c22eb8c69fe03 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/Mfcc.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/Mfcc.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "Mfcc"
   input_arg {
     name: "spectrogram"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/Min.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/Min.pbtxt
index bb0cd548c8a13a..4959b5e8d583b9 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/Min.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/Min.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "Min"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/Minimum.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/Minimum.pbtxt
index 756451b0b748d9..01cc483ba54dcb 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/Minimum.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/Minimum.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "Minimum"
   input_arg {
     name: "x"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/MirrorPad.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/MirrorPad.pbtxt
index 0642f97fafe52a..bf64a6ca5040b1 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/MirrorPad.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/MirrorPad.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "MirrorPad"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/MirrorPadGrad.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/MirrorPadGrad.pbtxt
index d1503a556c6127..b544cfbe72e7c4 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/MirrorPadGrad.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/MirrorPadGrad.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "MirrorPadGrad"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/MlirPassthroughOp.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/MlirPassthroughOp.pbtxt
index 4f7ddeb99acaf0..5990eb2850281d 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/MlirPassthroughOp.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/MlirPassthroughOp.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "MlirPassthroughOp"
   input_arg {
     name: "inputs"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/Mod.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/Mod.pbtxt
index 37a302eaa25de1..6c39ed683f6bee 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/Mod.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/Mod.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "Mod"
   input_arg {
     name: "x"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ModelDataset.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ModelDataset.pbtxt
index d3228c40e1d4dc..14ac940f6a7764 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ModelDataset.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ModelDataset.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ModelDataset"
   input_arg {
     name: "input_dataset"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/Mul.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/Mul.pbtxt
index d13b3adae90bd6..ef592669ff62cd 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/Mul.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/Mul.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "Mul"
   input_arg {
     name: "x"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/MulNoNan.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/MulNoNan.pbtxt
index ba06844a75700e..ca5c92fb15bfde 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/MulNoNan.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/MulNoNan.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "MulNoNan"
   input_arg {
     name: "x"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/MultiDeviceIterator.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/MultiDeviceIterator.pbtxt
index d7067ebf6eab43..d85c553f186f3b 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/MultiDeviceIterator.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/MultiDeviceIterator.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "MultiDeviceIterator"
   output_arg {
     name: "handle"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/MultiDeviceIteratorFromStringHandle.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/MultiDeviceIteratorFromStringHandle.pbtxt
index 74acb5eb9ce0c9..384b1477dbc599 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/MultiDeviceIteratorFromStringHandle.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/MultiDeviceIteratorFromStringHandle.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "MultiDeviceIteratorFromStringHandle"
   input_arg {
     name: "string_handle"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/MultiDeviceIteratorGetNextFromShard.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/MultiDeviceIteratorGetNextFromShard.pbtxt
index 8cf4fb6c28514c..2e007c25b2337d 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/MultiDeviceIteratorGetNextFromShard.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/MultiDeviceIteratorGetNextFromShard.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "MultiDeviceIteratorGetNextFromShard"
   input_arg {
     name: "multi_device_iterator"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/MultiDeviceIteratorInit.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/MultiDeviceIteratorInit.pbtxt
index a2829e514c5beb..a011997186af14 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/MultiDeviceIteratorInit.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/MultiDeviceIteratorInit.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "MultiDeviceIteratorInit"
   input_arg {
     name: "dataset"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/MultiDeviceIteratorToStringHandle.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/MultiDeviceIteratorToStringHandle.pbtxt
index 6f332755692ba1..d7780d79687166 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/MultiDeviceIteratorToStringHandle.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/MultiDeviceIteratorToStringHandle.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "MultiDeviceIteratorToStringHandle"
   input_arg {
     name: "multi_device_iterator"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/Multinomial.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/Multinomial.pbtxt
index ee55af8dddac95..c258fa6e7ed4c8 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/Multinomial.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/Multinomial.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "Multinomial"
   input_arg {
     name: "logits"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/MutableDenseHashTable.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/MutableDenseHashTable.pbtxt
index 4ebb385704164c..eecaeb2d4fb570 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/MutableDenseHashTable.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/MutableDenseHashTable.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "MutableDenseHashTable"
   input_arg {
     name: "empty_key"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/MutableDenseHashTableV2.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/MutableDenseHashTableV2.pbtxt
index 861d25c0c1b3f4..739079ced16cb6 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/MutableDenseHashTableV2.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/MutableDenseHashTableV2.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "MutableDenseHashTableV2"
   input_arg {
     name: "empty_key"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/MutableHashTable.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/MutableHashTable.pbtxt
index c1f2e216da5b83..a8ecc34cb18bd0 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/MutableHashTable.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/MutableHashTable.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "MutableHashTable"
   output_arg {
     name: "table_handle"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/MutableHashTableOfTensors.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/MutableHashTableOfTensors.pbtxt
index 7b71665fb96804..bdec2ff5939b39 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/MutableHashTableOfTensors.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/MutableHashTableOfTensors.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "MutableHashTableOfTensors"
   output_arg {
     name: "table_handle"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/MutableHashTableOfTensorsV2.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/MutableHashTableOfTensorsV2.pbtxt
index d32880aba77e28..dc46d075df3d94 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/MutableHashTableOfTensorsV2.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/MutableHashTableOfTensorsV2.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "MutableHashTableOfTensorsV2"
   output_arg {
     name: "table_handle"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/MutableHashTableV2.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/MutableHashTableV2.pbtxt
index eb690ee06e1aa9..610214dfa76e66 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/MutableHashTableV2.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/MutableHashTableV2.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "MutableHashTableV2"
   output_arg {
     name: "table_handle"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/MutexLock.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/MutexLock.pbtxt
index 16f1ad6e0a1049..6b5747cd4e383e 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/MutexLock.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/MutexLock.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "MutexLock"
   input_arg {
     name: "mutex"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/MutexV2.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/MutexV2.pbtxt
index ef04c075cd8d6e..b20f9b1e7996b3 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/MutexV2.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/MutexV2.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "MutexV2"
   output_arg {
     name: "resource"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/NcclAllReduce.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/NcclAllReduce.pbtxt
index 59ab7c672eb6cc..80f91edef1d237 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/NcclAllReduce.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/NcclAllReduce.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "NcclAllReduce"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/NcclBroadcast.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/NcclBroadcast.pbtxt
index 46a2cab258fbf0..02a5487d1acf6f 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/NcclBroadcast.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/NcclBroadcast.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "NcclBroadcast"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/NcclReduce.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/NcclReduce.pbtxt
index 141189947600f3..507f92cff2c87e 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/NcclReduce.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/NcclReduce.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "NcclReduce"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/Ndtri.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/Ndtri.pbtxt
index c1549d77e8e841..a7a923f780db40 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/Ndtri.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/Ndtri.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "Ndtri"
   input_arg {
     name: "x"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/NearestNeighbors.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/NearestNeighbors.pbtxt
index 42b7b7c906bd85..5d1e5ed57659ac 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/NearestNeighbors.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/NearestNeighbors.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "NearestNeighbors"
   input_arg {
     name: "points"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/Neg.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/Neg.pbtxt
index 22a09654a9b7c3..864d0257fe4b2d 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/Neg.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/Neg.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "Neg"
   input_arg {
     name: "x"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/NegTrain.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/NegTrain.pbtxt
index d97f107378d3b9..f12529fd6328b8 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/NegTrain.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/NegTrain.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "NegTrain"
   input_arg {
     name: "w_in"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/NextAfter.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/NextAfter.pbtxt
index 1cbacb54fe2bec..70e4afe6c77db5 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/NextAfter.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/NextAfter.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "NextAfter"
   input_arg {
     name: "x1"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/NextIteration.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/NextIteration.pbtxt
index 6044c37879b4d2..7186fc0b684029 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/NextIteration.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/NextIteration.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "NextIteration"
   input_arg {
     name: "data"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/NoOp.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/NoOp.pbtxt
index 6c06e127e1bd6b..8f0370633fc1ce 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/NoOp.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/NoOp.pbtxt
@@ -1,3 +1,3 @@
-op 	 {
+op {
   name: "NoOp"
 }
diff --git a/tensorflow/core/ops/compat/ops_history_v2/NonDeterministicInts.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/NonDeterministicInts.pbtxt
index 470bdf155b510a..3fa5aa4a605c7d 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/NonDeterministicInts.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/NonDeterministicInts.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "NonDeterministicInts"
   input_arg {
     name: "shape"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/NonMaxSuppression.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/NonMaxSuppression.pbtxt
index a63ddf63dbac33..ded8b3728f4613 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/NonMaxSuppression.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/NonMaxSuppression.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "NonMaxSuppression"
   input_arg {
     name: "boxes"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/NonMaxSuppressionV2.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/NonMaxSuppressionV2.pbtxt
index 0ff1ffcb5ae632..90c23bc0457dcc 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/NonMaxSuppressionV2.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/NonMaxSuppressionV2.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "NonMaxSuppressionV2"
   input_arg {
     name: "boxes"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/NonMaxSuppressionV3.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/NonMaxSuppressionV3.pbtxt
index bf9aac7c28f2c5..daeffd841b8e00 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/NonMaxSuppressionV3.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/NonMaxSuppressionV3.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "NonMaxSuppressionV3"
   input_arg {
     name: "boxes"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/NonMaxSuppressionV4.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/NonMaxSuppressionV4.pbtxt
index ee1f325a9bd8d6..07ca92fef71df1 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/NonMaxSuppressionV4.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/NonMaxSuppressionV4.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "NonMaxSuppressionV4"
   input_arg {
     name: "boxes"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/NonMaxSuppressionV5.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/NonMaxSuppressionV5.pbtxt
index fbfe05a13141a1..cabec767a15617 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/NonMaxSuppressionV5.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/NonMaxSuppressionV5.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "NonMaxSuppressionV5"
   input_arg {
     name: "boxes"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/NonMaxSuppressionWithOverlaps.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/NonMaxSuppressionWithOverlaps.pbtxt
index a150d688af1b69..d89eeee4a4a34b 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/NonMaxSuppressionWithOverlaps.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/NonMaxSuppressionWithOverlaps.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "NonMaxSuppressionWithOverlaps"
   input_arg {
     name: "overlaps"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/NonSerializableDataset.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/NonSerializableDataset.pbtxt
index b47c4157f9dbde..5fbd4bdb656c57 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/NonSerializableDataset.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/NonSerializableDataset.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "NonSerializableDataset"
   input_arg {
     name: "input_dataset"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/NotEqual.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/NotEqual.pbtxt
index 738fc0de0114d8..099ef75d622d83 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/NotEqual.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/NotEqual.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "NotEqual"
   input_arg {
     name: "x"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/NthElement.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/NthElement.pbtxt
index a1c7a12ba00126..c9e797273df119 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/NthElement.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/NthElement.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "NthElement"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/OneHot.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/OneHot.pbtxt
index 83163e8685a19e..0c1cfb6d686a95 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/OneHot.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/OneHot.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "OneHot"
   input_arg {
     name: "indices"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/OneShotIterator.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/OneShotIterator.pbtxt
index 362e5b271b172e..a2969bcc0e36c9 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/OneShotIterator.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/OneShotIterator.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "OneShotIterator"
   output_arg {
     name: "handle"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/OnesLike.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/OnesLike.pbtxt
index 6827c87d832660..d4609438139b35 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/OnesLike.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/OnesLike.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "OnesLike"
   input_arg {
     name: "x"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/OptimizeDataset.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/OptimizeDataset.pbtxt
index 4a2b3f7420fd6c..5b40b213eb5dc8 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/OptimizeDataset.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/OptimizeDataset.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "OptimizeDataset"
   input_arg {
     name: "input_dataset"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/OptimizeDatasetV2.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/OptimizeDatasetV2.pbtxt
index d647c7bc1f0d58..7e1537ebe57a3a 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/OptimizeDatasetV2.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/OptimizeDatasetV2.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "OptimizeDatasetV2"
   input_arg {
     name: "input_dataset"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/OptionalFromValue.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/OptionalFromValue.pbtxt
index 3a9510a674c426..ddfc9c43d348c3 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/OptionalFromValue.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/OptionalFromValue.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "OptionalFromValue"
   input_arg {
     name: "components"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/OptionalGetValue.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/OptionalGetValue.pbtxt
index e5f3ab817ab570..e7364a1014afe8 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/OptionalGetValue.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/OptionalGetValue.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "OptionalGetValue"
   input_arg {
     name: "optional"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/OptionalHasValue.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/OptionalHasValue.pbtxt
index e744d908129750..da76333cecbf70 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/OptionalHasValue.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/OptionalHasValue.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "OptionalHasValue"
   input_arg {
     name: "optional"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/OptionalNone.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/OptionalNone.pbtxt
index a051d978224361..c47d6a745481e9 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/OptionalNone.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/OptionalNone.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "OptionalNone"
   output_arg {
     name: "optional"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/OptionsDataset.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/OptionsDataset.pbtxt
index 6e3c0a9fdb22a7..fc63e5ee2951cf 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/OptionsDataset.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/OptionsDataset.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "OptionsDataset"
   input_arg {
     name: "input_dataset"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/OrderedMapClear.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/OrderedMapClear.pbtxt
index 2f77ea5e73bbb4..726e26e6172f05 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/OrderedMapClear.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/OrderedMapClear.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "OrderedMapClear"
   attr {
     name: "capacity"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/OrderedMapIncompleteSize.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/OrderedMapIncompleteSize.pbtxt
index a6439051fe23a0..9a9572a51beea8 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/OrderedMapIncompleteSize.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/OrderedMapIncompleteSize.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "OrderedMapIncompleteSize"
   output_arg {
     name: "size"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/OrderedMapPeek.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/OrderedMapPeek.pbtxt
index 8e61ac23f9bb73..0d9fd20fe077dd 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/OrderedMapPeek.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/OrderedMapPeek.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "OrderedMapPeek"
   input_arg {
     name: "key"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/OrderedMapSize.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/OrderedMapSize.pbtxt
index 46777f1d7e7e8e..ea07d7e4215fb8 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/OrderedMapSize.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/OrderedMapSize.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "OrderedMapSize"
   output_arg {
     name: "size"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/OrderedMapStage.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/OrderedMapStage.pbtxt
index badbf44b424839..76af456ed8372b 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/OrderedMapStage.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/OrderedMapStage.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "OrderedMapStage"
   input_arg {
     name: "key"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/OrderedMapUnstage.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/OrderedMapUnstage.pbtxt
index f3d4dced87d04e..c09b4be94f4908 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/OrderedMapUnstage.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/OrderedMapUnstage.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "OrderedMapUnstage"
   input_arg {
     name: "key"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/OrderedMapUnstageNoKey.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/OrderedMapUnstageNoKey.pbtxt
index 295f8258d599c4..bc3e8c7da30200 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/OrderedMapUnstageNoKey.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/OrderedMapUnstageNoKey.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "OrderedMapUnstageNoKey"
   input_arg {
     name: "indices"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/OutfeedDequeue.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/OutfeedDequeue.pbtxt
index b802b78283c0ee..29dc8b5c5879bd 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/OutfeedDequeue.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/OutfeedDequeue.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "OutfeedDequeue"
   output_arg {
     name: "output"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/OutfeedDequeueTuple.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/OutfeedDequeueTuple.pbtxt
index 580babfb502d44..3e0d31078b3ec6 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/OutfeedDequeueTuple.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/OutfeedDequeueTuple.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "OutfeedDequeueTuple"
   output_arg {
     name: "outputs"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/OutfeedDequeueTupleV2.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/OutfeedDequeueTupleV2.pbtxt
index e8d13591681564..744744b4545cad 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/OutfeedDequeueTupleV2.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/OutfeedDequeueTupleV2.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "OutfeedDequeueTupleV2"
   input_arg {
     name: "device_ordinal"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/OutfeedDequeueV2.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/OutfeedDequeueV2.pbtxt
index 5863aae467fed3..c5ca1f31a676cb 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/OutfeedDequeueV2.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/OutfeedDequeueV2.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "OutfeedDequeueV2"
   input_arg {
     name: "device_ordinal"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/OutfeedEnqueue.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/OutfeedEnqueue.pbtxt
index 4836db6b12b31f..d8c16f4d62978c 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/OutfeedEnqueue.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/OutfeedEnqueue.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "OutfeedEnqueue"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/OutfeedEnqueueTuple.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/OutfeedEnqueueTuple.pbtxt
index e7464383175911..0bf1a5ba4804eb 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/OutfeedEnqueueTuple.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/OutfeedEnqueueTuple.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "OutfeedEnqueueTuple"
   input_arg {
     name: "inputs"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/Pack.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/Pack.pbtxt
index 12eecb122fbe6a..65eb67509d3fd8 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/Pack.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/Pack.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "Pack"
   input_arg {
     name: "values"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/Pad.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/Pad.pbtxt
index 5d1b80e2f3976c..1c7b9c7b457b3c 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/Pad.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/Pad.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "Pad"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/PadV2.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/PadV2.pbtxt
index aa1a0fb109be72..463cb71f207ed0 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/PadV2.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/PadV2.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "PadV2"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/PaddedBatchDataset.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/PaddedBatchDataset.pbtxt
index f67cb9f4e4e5af..a118fc102f10a2 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/PaddedBatchDataset.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/PaddedBatchDataset.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "PaddedBatchDataset"
   input_arg {
     name: "input_dataset"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/PaddedBatchDatasetV2.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/PaddedBatchDatasetV2.pbtxt
index 1142cbf9eb5bed..4ae5a66624a65d 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/PaddedBatchDatasetV2.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/PaddedBatchDatasetV2.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "PaddedBatchDatasetV2"
   input_arg {
     name: "input_dataset"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/PaddingFIFOQueue.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/PaddingFIFOQueue.pbtxt
index e03ca8fecfec5f..f5eca52ba1927e 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/PaddingFIFOQueue.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/PaddingFIFOQueue.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "PaddingFIFOQueue"
   output_arg {
     name: "handle"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/PaddingFIFOQueueV2.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/PaddingFIFOQueueV2.pbtxt
index 12404b10303679..c398f9ee3a8a14 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/PaddingFIFOQueueV2.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/PaddingFIFOQueueV2.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "PaddingFIFOQueueV2"
   output_arg {
     name: "handle"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ParallelBatchDataset.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ParallelBatchDataset.pbtxt
index 023d855a4fa2fd..5c160cae2ddad7 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ParallelBatchDataset.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ParallelBatchDataset.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ParallelBatchDataset"
   input_arg {
     name: "input_dataset"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ParallelConcat.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ParallelConcat.pbtxt
index 154f48829d78d4..b0d1cc3918574a 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ParallelConcat.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ParallelConcat.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ParallelConcat"
   input_arg {
     name: "values"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ParallelDynamicStitch.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ParallelDynamicStitch.pbtxt
index 8566fc274f031d..9ab18a1ba5e69e 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ParallelDynamicStitch.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ParallelDynamicStitch.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ParallelDynamicStitch"
   input_arg {
     name: "indices"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ParallelFilterDataset.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ParallelFilterDataset.pbtxt
index e33694113b5105..1c895e2767636e 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ParallelFilterDataset.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ParallelFilterDataset.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ParallelFilterDataset"
   input_arg {
     name: "input_dataset"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ParallelInterleaveDataset.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ParallelInterleaveDataset.pbtxt
index 730f9b76b490be..f278cb0efc6dfb 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ParallelInterleaveDataset.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ParallelInterleaveDataset.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ParallelInterleaveDataset"
   input_arg {
     name: "input_dataset"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ParallelInterleaveDatasetV2.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ParallelInterleaveDatasetV2.pbtxt
index 507463da15bb97..110573b42ed39f 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ParallelInterleaveDatasetV2.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ParallelInterleaveDatasetV2.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ParallelInterleaveDatasetV2"
   input_arg {
     name: "input_dataset"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ParallelInterleaveDatasetV3.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ParallelInterleaveDatasetV3.pbtxt
index 3f73fc38abc0d0..096460fb1efcdb 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ParallelInterleaveDatasetV3.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ParallelInterleaveDatasetV3.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ParallelInterleaveDatasetV3"
   input_arg {
     name: "input_dataset"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ParallelInterleaveDatasetV4.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ParallelInterleaveDatasetV4.pbtxt
index 65d63780827683..94f9ae0f6eebd1 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ParallelInterleaveDatasetV4.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ParallelInterleaveDatasetV4.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ParallelInterleaveDatasetV4"
   input_arg {
     name: "input_dataset"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ParallelMapDataset.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ParallelMapDataset.pbtxt
index 20ac8a4dcc3d51..991e4192983c3d 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ParallelMapDataset.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ParallelMapDataset.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ParallelMapDataset"
   input_arg {
     name: "input_dataset"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ParallelMapDatasetV2.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ParallelMapDatasetV2.pbtxt
index 87012c3c49f3d6..55e73b740adefd 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ParallelMapDatasetV2.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ParallelMapDatasetV2.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ParallelMapDatasetV2"
   input_arg {
     name: "input_dataset"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ParameterizedTruncatedNormal.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ParameterizedTruncatedNormal.pbtxt
index 8bca662ec52d00..1f96da6f7886e6 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ParameterizedTruncatedNormal.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ParameterizedTruncatedNormal.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ParameterizedTruncatedNormal"
   input_arg {
     name: "shape"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ParseExample.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ParseExample.pbtxt
index b1020d356343db..a1e35bde86b0a1 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ParseExample.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ParseExample.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ParseExample"
   input_arg {
     name: "serialized"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ParseExampleDataset.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ParseExampleDataset.pbtxt
index 4e12e94a6a28c8..4dc9ac1efb6cd3 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ParseExampleDataset.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ParseExampleDataset.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ParseExampleDataset"
   input_arg {
     name: "input_dataset"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ParseExampleDatasetV2.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ParseExampleDatasetV2.pbtxt
index 8af77a0bd59095..59632a160b121c 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ParseExampleDatasetV2.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ParseExampleDatasetV2.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ParseExampleDatasetV2"
   input_arg {
     name: "input_dataset"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ParseExampleV2.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ParseExampleV2.pbtxt
index cc2706b046e368..0d7e1d37e23ac9 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ParseExampleV2.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ParseExampleV2.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ParseExampleV2"
   input_arg {
     name: "serialized"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ParseSequenceExample.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ParseSequenceExample.pbtxt
index 8060fe90a3553e..03ac5be8d26160 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ParseSequenceExample.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ParseSequenceExample.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ParseSequenceExample"
   input_arg {
     name: "serialized"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ParseSequenceExampleV2.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ParseSequenceExampleV2.pbtxt
index 07dfc9d7467e14..7ba59734893b0b 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ParseSequenceExampleV2.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ParseSequenceExampleV2.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ParseSequenceExampleV2"
   input_arg {
     name: "serialized"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ParseSingleExample.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ParseSingleExample.pbtxt
index 0ea9857ac83b6f..aaa69af4f62ac6 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ParseSingleExample.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ParseSingleExample.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ParseSingleExample"
   input_arg {
     name: "serialized"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ParseSingleSequenceExample.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ParseSingleSequenceExample.pbtxt
index c0f76a20b44160..a0f52dbdd1f406 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ParseSingleSequenceExample.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ParseSingleSequenceExample.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ParseSingleSequenceExample"
   input_arg {
     name: "serialized"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ParseTensor.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ParseTensor.pbtxt
index 20232f99165c8d..63d1f1292eed08 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ParseTensor.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ParseTensor.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ParseTensor"
   input_arg {
     name: "serialized"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/PartitionedCall.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/PartitionedCall.pbtxt
index 59752220042a63..b51bd1de9fce0d 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/PartitionedCall.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/PartitionedCall.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "PartitionedCall"
   input_arg {
     name: "args"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/Placeholder.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/Placeholder.pbtxt
index ec0fdcf19f7ebd..7c0f57a94e177d 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/Placeholder.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/Placeholder.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "Placeholder"
   output_arg {
     name: "output"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/PlaceholderV2.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/PlaceholderV2.pbtxt
index e9d5b4840cb8e0..b2cd20b238fa5e 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/PlaceholderV2.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/PlaceholderV2.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "PlaceholderV2"
   output_arg {
     name: "output"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/PlaceholderWithDefault.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/PlaceholderWithDefault.pbtxt
index 3d02d762312073..79a2ffb4492a6a 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/PlaceholderWithDefault.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/PlaceholderWithDefault.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "PlaceholderWithDefault"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/Polygamma.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/Polygamma.pbtxt
index 3f20bc10cd2332..6bf0d9ba4cf9cc 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/Polygamma.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/Polygamma.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "Polygamma"
   input_arg {
     name: "a"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/PopulationCount.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/PopulationCount.pbtxt
index 97e98373418383..d66c1ac00ae122 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/PopulationCount.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/PopulationCount.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "PopulationCount"
   input_arg {
     name: "x"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/Pow.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/Pow.pbtxt
index ad30b536cd7742..b1cc1b8f479801 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/Pow.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/Pow.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "Pow"
   input_arg {
     name: "x"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/PrefetchDataset.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/PrefetchDataset.pbtxt
index 92926f6523d1ff..81953a0e75c0c9 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/PrefetchDataset.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/PrefetchDataset.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "PrefetchDataset"
   input_arg {
     name: "input_dataset"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/Prelinearize.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/Prelinearize.pbtxt
index f74c9b381dce7d..b5ed810c25a426 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/Prelinearize.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/Prelinearize.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "Prelinearize"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/PrelinearizeTuple.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/PrelinearizeTuple.pbtxt
index 92de7f372a74d5..bb1ae7d3e2f828 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/PrelinearizeTuple.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/PrelinearizeTuple.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "PrelinearizeTuple"
   input_arg {
     name: "inputs"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/PreventGradient.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/PreventGradient.pbtxt
index 19a95b09d5e645..1649fc808aa935 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/PreventGradient.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/PreventGradient.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "PreventGradient"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/Print.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/Print.pbtxt
index 1966093e81c6ad..fbbb514b177737 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/Print.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/Print.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "Print"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/PrintV2.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/PrintV2.pbtxt
index 8f748465b810bd..c5942f0a6145a5 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/PrintV2.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/PrintV2.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "PrintV2"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/PriorityQueue.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/PriorityQueue.pbtxt
index af0f36d238d9e6..b44d83dfb20036 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/PriorityQueue.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/PriorityQueue.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "PriorityQueue"
   output_arg {
     name: "handle"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/PriorityQueueV2.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/PriorityQueueV2.pbtxt
index ab426238988baf..a4e7c750b65a14 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/PriorityQueueV2.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/PriorityQueueV2.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "PriorityQueueV2"
   output_arg {
     name: "handle"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/PrivateThreadPoolDataset.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/PrivateThreadPoolDataset.pbtxt
index 483f4b88fd93f7..c16c1eb164728d 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/PrivateThreadPoolDataset.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/PrivateThreadPoolDataset.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "PrivateThreadPoolDataset"
   input_arg {
     name: "input_dataset"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/Prod.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/Prod.pbtxt
index d3c71eacd78cf9..fe9126bca16ed1 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/Prod.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/Prod.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "Prod"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/PyFunc.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/PyFunc.pbtxt
index de5661de7e2150..987f028051ea0d 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/PyFunc.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/PyFunc.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "PyFunc"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/PyFuncStateless.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/PyFuncStateless.pbtxt
index fa2ac3e446b1a7..2a587d53d937c8 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/PyFuncStateless.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/PyFuncStateless.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "PyFuncStateless"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/Qr.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/Qr.pbtxt
index 251f8ff7f7f0d6..8319528f8a016e 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/Qr.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/Qr.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "Qr"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/QuantizeAndDequantize.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/QuantizeAndDequantize.pbtxt
index 30d0b750a8df0a..fb662f7057c2f0 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/QuantizeAndDequantize.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/QuantizeAndDequantize.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "QuantizeAndDequantize"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/QuantizeAndDequantizeV2.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/QuantizeAndDequantizeV2.pbtxt
index d1fa4fba7fd4e5..d946e753cf7bbf 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/QuantizeAndDequantizeV2.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/QuantizeAndDequantizeV2.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "QuantizeAndDequantizeV2"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/QuantizeAndDequantizeV3.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/QuantizeAndDequantizeV3.pbtxt
index 8d6a0317190cfe..cb7762a9deda20 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/QuantizeAndDequantizeV3.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/QuantizeAndDequantizeV3.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "QuantizeAndDequantizeV3"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/QuantizeAndDequantizeV4.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/QuantizeAndDequantizeV4.pbtxt
index fca137de6cd4bc..2a49131faaf2d8 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/QuantizeAndDequantizeV4.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/QuantizeAndDequantizeV4.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "QuantizeAndDequantizeV4"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/QuantizeAndDequantizeV4Grad.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/QuantizeAndDequantizeV4Grad.pbtxt
index 71df7854231502..0bbe87452b145d 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/QuantizeAndDequantizeV4Grad.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/QuantizeAndDequantizeV4Grad.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "QuantizeAndDequantizeV4Grad"
   input_arg {
     name: "gradients"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/QuantizeDownAndShrinkRange.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/QuantizeDownAndShrinkRange.pbtxt
index ea6737ee390894..42783d3a14ebb4 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/QuantizeDownAndShrinkRange.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/QuantizeDownAndShrinkRange.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "QuantizeDownAndShrinkRange"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/QuantizeV2.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/QuantizeV2.pbtxt
index d164717d5aaf39..37cd1384176d6e 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/QuantizeV2.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/QuantizeV2.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "QuantizeV2"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/QuantizedAdd.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/QuantizedAdd.pbtxt
index be3f511f03aec3..4532bc23d50a2b 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/QuantizedAdd.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/QuantizedAdd.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "QuantizedAdd"
   input_arg {
     name: "x"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/QuantizedAvgPool.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/QuantizedAvgPool.pbtxt
index 44e51eb39509f7..0ae3390d303c11 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/QuantizedAvgPool.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/QuantizedAvgPool.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "QuantizedAvgPool"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/QuantizedBatchNormWithGlobalNormalization.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/QuantizedBatchNormWithGlobalNormalization.pbtxt
index be4b9356265362..832b8ba577577d 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/QuantizedBatchNormWithGlobalNormalization.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/QuantizedBatchNormWithGlobalNormalization.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "QuantizedBatchNormWithGlobalNormalization"
   input_arg {
     name: "t"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/QuantizedBiasAdd.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/QuantizedBiasAdd.pbtxt
index 3c94fb54f0a497..b479c2c54e4d26 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/QuantizedBiasAdd.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/QuantizedBiasAdd.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "QuantizedBiasAdd"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/QuantizedConcat.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/QuantizedConcat.pbtxt
index 2752ebfca09e38..449f588ac8f498 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/QuantizedConcat.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/QuantizedConcat.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "QuantizedConcat"
   input_arg {
     name: "concat_dim"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/QuantizedConv2D.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/QuantizedConv2D.pbtxt
index 3680e575ace2f3..b1cf1c8d334182 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/QuantizedConv2D.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/QuantizedConv2D.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "QuantizedConv2D"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/QuantizedConv2DAndRelu.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/QuantizedConv2DAndRelu.pbtxt
index c9f1696d1ecf67..229e4c436dd622 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/QuantizedConv2DAndRelu.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/QuantizedConv2DAndRelu.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "QuantizedConv2DAndRelu"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/QuantizedConv2DAndReluAndRequantize.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/QuantizedConv2DAndReluAndRequantize.pbtxt
index 1cbc11d524e167..bc566896f53dfe 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/QuantizedConv2DAndReluAndRequantize.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/QuantizedConv2DAndReluAndRequantize.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "QuantizedConv2DAndReluAndRequantize"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/QuantizedConv2DAndRequantize.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/QuantizedConv2DAndRequantize.pbtxt
index d3166c425a984f..5d26709f14e31c 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/QuantizedConv2DAndRequantize.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/QuantizedConv2DAndRequantize.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "QuantizedConv2DAndRequantize"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/QuantizedConv2DPerChannel.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/QuantizedConv2DPerChannel.pbtxt
index 06c3d5f115fab3..93640944477061 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/QuantizedConv2DPerChannel.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/QuantizedConv2DPerChannel.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "QuantizedConv2DPerChannel"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/QuantizedConv2DWithBias.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/QuantizedConv2DWithBias.pbtxt
index f059cf786d14f0..8372a882260457 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/QuantizedConv2DWithBias.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/QuantizedConv2DWithBias.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "QuantizedConv2DWithBias"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/QuantizedConv2DWithBiasAndRelu.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/QuantizedConv2DWithBiasAndRelu.pbtxt
index 4eef8cb4c4f3c2..af0ce39a844546 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/QuantizedConv2DWithBiasAndRelu.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/QuantizedConv2DWithBiasAndRelu.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "QuantizedConv2DWithBiasAndRelu"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/QuantizedConv2DWithBiasAndReluAndRequantize.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/QuantizedConv2DWithBiasAndReluAndRequantize.pbtxt
index 174a196a036974..599f19e666dcdc 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/QuantizedConv2DWithBiasAndReluAndRequantize.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/QuantizedConv2DWithBiasAndReluAndRequantize.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "QuantizedConv2DWithBiasAndReluAndRequantize"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/QuantizedConv2DWithBiasAndRequantize.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/QuantizedConv2DWithBiasAndRequantize.pbtxt
index 96ed2852b68046..8cf8fbb2eae658 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/QuantizedConv2DWithBiasAndRequantize.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/QuantizedConv2DWithBiasAndRequantize.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "QuantizedConv2DWithBiasAndRequantize"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/QuantizedConv2DWithBiasSignedSumAndReluAndRequantize.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/QuantizedConv2DWithBiasSignedSumAndReluAndRequantize.pbtxt
index 37f884f1cb0e88..e46786a9a74a86 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/QuantizedConv2DWithBiasSignedSumAndReluAndRequantize.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/QuantizedConv2DWithBiasSignedSumAndReluAndRequantize.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "QuantizedConv2DWithBiasSignedSumAndReluAndRequantize"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/QuantizedConv2DWithBiasSumAndRelu.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/QuantizedConv2DWithBiasSumAndRelu.pbtxt
index ed100f68a9b58a..d74439b670ee37 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/QuantizedConv2DWithBiasSumAndRelu.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/QuantizedConv2DWithBiasSumAndRelu.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "QuantizedConv2DWithBiasSumAndRelu"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/QuantizedConv2DWithBiasSumAndReluAndRequantize.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/QuantizedConv2DWithBiasSumAndReluAndRequantize.pbtxt
index d4a3c49b46c5c7..70c2366a19f2bd 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/QuantizedConv2DWithBiasSumAndReluAndRequantize.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/QuantizedConv2DWithBiasSumAndReluAndRequantize.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "QuantizedConv2DWithBiasSumAndReluAndRequantize"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/QuantizedDepthwiseConv2D.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/QuantizedDepthwiseConv2D.pbtxt
index 6b4163c799cdfa..f88bba239bf07b 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/QuantizedDepthwiseConv2D.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/QuantizedDepthwiseConv2D.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "QuantizedDepthwiseConv2D"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/QuantizedDepthwiseConv2DWithBias.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/QuantizedDepthwiseConv2DWithBias.pbtxt
index c9e584543f3767..4faf839b3b6651 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/QuantizedDepthwiseConv2DWithBias.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/QuantizedDepthwiseConv2DWithBias.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "QuantizedDepthwiseConv2DWithBias"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/QuantizedDepthwiseConv2DWithBiasAndRelu.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/QuantizedDepthwiseConv2DWithBiasAndRelu.pbtxt
index a4bb2fcd372ded..6a4b514e23ffec 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/QuantizedDepthwiseConv2DWithBiasAndRelu.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/QuantizedDepthwiseConv2DWithBiasAndRelu.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "QuantizedDepthwiseConv2DWithBiasAndRelu"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/QuantizedDepthwiseConv2DWithBiasAndReluAndRequantize.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/QuantizedDepthwiseConv2DWithBiasAndReluAndRequantize.pbtxt
index c24ba3165a8aca..3de6c3f9e28f47 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/QuantizedDepthwiseConv2DWithBiasAndReluAndRequantize.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/QuantizedDepthwiseConv2DWithBiasAndReluAndRequantize.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "QuantizedDepthwiseConv2DWithBiasAndReluAndRequantize"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/QuantizedInstanceNorm.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/QuantizedInstanceNorm.pbtxt
index 720358f1bd01d7..98136d82ebda45 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/QuantizedInstanceNorm.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/QuantizedInstanceNorm.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "QuantizedInstanceNorm"
   input_arg {
     name: "x"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/QuantizedMatMul.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/QuantizedMatMul.pbtxt
index d1332abc5169ca..7e4707a316f80d 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/QuantizedMatMul.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/QuantizedMatMul.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "QuantizedMatMul"
   input_arg {
     name: "a"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/QuantizedMatMulWithBias.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/QuantizedMatMulWithBias.pbtxt
index 55c80678a6eb90..a59adb7f78c6d7 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/QuantizedMatMulWithBias.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/QuantizedMatMulWithBias.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "QuantizedMatMulWithBias"
   input_arg {
     name: "a"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/QuantizedMatMulWithBiasAndDequantize.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/QuantizedMatMulWithBiasAndDequantize.pbtxt
index 6c6370a1ac0191..04ecfdbd855f83 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/QuantizedMatMulWithBiasAndDequantize.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/QuantizedMatMulWithBiasAndDequantize.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "QuantizedMatMulWithBiasAndDequantize"
   input_arg {
     name: "a"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/QuantizedMatMulWithBiasAndRelu.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/QuantizedMatMulWithBiasAndRelu.pbtxt
index ecdeb1d05d1c90..cd0acb9d721657 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/QuantizedMatMulWithBiasAndRelu.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/QuantizedMatMulWithBiasAndRelu.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "QuantizedMatMulWithBiasAndRelu"
   input_arg {
     name: "a"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/QuantizedMatMulWithBiasAndReluAndRequantize.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/QuantizedMatMulWithBiasAndReluAndRequantize.pbtxt
index 140ed993453766..b591d3fb37c868 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/QuantizedMatMulWithBiasAndReluAndRequantize.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/QuantizedMatMulWithBiasAndReluAndRequantize.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "QuantizedMatMulWithBiasAndReluAndRequantize"
   input_arg {
     name: "a"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/QuantizedMatMulWithBiasAndRequantize.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/QuantizedMatMulWithBiasAndRequantize.pbtxt
index 16b180475e46e2..1aab9762eea036 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/QuantizedMatMulWithBiasAndRequantize.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/QuantizedMatMulWithBiasAndRequantize.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "QuantizedMatMulWithBiasAndRequantize"
   input_arg {
     name: "a"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/QuantizedMaxPool.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/QuantizedMaxPool.pbtxt
index 9164c0fcfc6160..47d6ac80518ef9 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/QuantizedMaxPool.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/QuantizedMaxPool.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "QuantizedMaxPool"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/QuantizedMul.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/QuantizedMul.pbtxt
index 97c025672c48b2..795ab1341d9c67 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/QuantizedMul.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/QuantizedMul.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "QuantizedMul"
   input_arg {
     name: "x"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/QuantizedRelu.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/QuantizedRelu.pbtxt
index 77fb5c2eeb9041..724d8b3946bdb1 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/QuantizedRelu.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/QuantizedRelu.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "QuantizedRelu"
   input_arg {
     name: "features"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/QuantizedRelu6.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/QuantizedRelu6.pbtxt
index b76f1159ec4f16..0f389d5eae9337 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/QuantizedRelu6.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/QuantizedRelu6.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "QuantizedRelu6"
   input_arg {
     name: "features"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/QuantizedReluX.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/QuantizedReluX.pbtxt
index 772ac59eba26ff..9ee6f0d2e274b9 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/QuantizedReluX.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/QuantizedReluX.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "QuantizedReluX"
   input_arg {
     name: "features"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/QuantizedReshape.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/QuantizedReshape.pbtxt
index fa7b1fee5cf5f1..f54db98943cc14 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/QuantizedReshape.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/QuantizedReshape.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "QuantizedReshape"
   input_arg {
     name: "tensor"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/QuantizedResizeBilinear.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/QuantizedResizeBilinear.pbtxt
index aca635c5016a1a..bee577ed23b411 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/QuantizedResizeBilinear.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/QuantizedResizeBilinear.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "QuantizedResizeBilinear"
   input_arg {
     name: "images"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/QueueClose.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/QueueClose.pbtxt
index 0a4a305cc97568..582eeccd6a263a 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/QueueClose.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/QueueClose.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "QueueClose"
   input_arg {
     name: "handle"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/QueueCloseV2.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/QueueCloseV2.pbtxt
index ab9c620c7fb7b4..e0544c13e654b4 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/QueueCloseV2.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/QueueCloseV2.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "QueueCloseV2"
   input_arg {
     name: "handle"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/QueueDequeue.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/QueueDequeue.pbtxt
index f4a3b6fe195c42..f06745f20fe456 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/QueueDequeue.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/QueueDequeue.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "QueueDequeue"
   input_arg {
     name: "handle"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/QueueDequeueMany.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/QueueDequeueMany.pbtxt
index 986897a5bcd27f..374ecfb18a87c3 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/QueueDequeueMany.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/QueueDequeueMany.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "QueueDequeueMany"
   input_arg {
     name: "handle"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/QueueDequeueManyV2.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/QueueDequeueManyV2.pbtxt
index 355b7905840281..f3ebc6c7e59288 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/QueueDequeueManyV2.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/QueueDequeueManyV2.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "QueueDequeueManyV2"
   input_arg {
     name: "handle"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/QueueDequeueUpTo.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/QueueDequeueUpTo.pbtxt
index 6e006006a784f5..6fa30ac810a0e5 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/QueueDequeueUpTo.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/QueueDequeueUpTo.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "QueueDequeueUpTo"
   input_arg {
     name: "handle"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/QueueDequeueUpToV2.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/QueueDequeueUpToV2.pbtxt
index 5901d313926bdf..2016cc7f04a0ad 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/QueueDequeueUpToV2.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/QueueDequeueUpToV2.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "QueueDequeueUpToV2"
   input_arg {
     name: "handle"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/QueueDequeueV2.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/QueueDequeueV2.pbtxt
index 44a893ea7c3a00..e338ccbd355a76 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/QueueDequeueV2.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/QueueDequeueV2.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "QueueDequeueV2"
   input_arg {
     name: "handle"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/QueueEnqueue.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/QueueEnqueue.pbtxt
index 689046ebae4e2d..fb94d288f2400e 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/QueueEnqueue.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/QueueEnqueue.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "QueueEnqueue"
   input_arg {
     name: "handle"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/QueueEnqueueMany.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/QueueEnqueueMany.pbtxt
index f3cbf429f28639..2d9582430725b9 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/QueueEnqueueMany.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/QueueEnqueueMany.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "QueueEnqueueMany"
   input_arg {
     name: "handle"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/QueueEnqueueManyV2.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/QueueEnqueueManyV2.pbtxt
index 159b3012075455..c327d27e2f11f6 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/QueueEnqueueManyV2.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/QueueEnqueueManyV2.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "QueueEnqueueManyV2"
   input_arg {
     name: "handle"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/QueueEnqueueV2.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/QueueEnqueueV2.pbtxt
index e6bf061965d84b..da8cdd3cd67c48 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/QueueEnqueueV2.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/QueueEnqueueV2.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "QueueEnqueueV2"
   input_arg {
     name: "handle"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/QueueIsClosed.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/QueueIsClosed.pbtxt
index 42426552b51808..11a421b27c85e3 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/QueueIsClosed.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/QueueIsClosed.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "QueueIsClosed"
   input_arg {
     name: "handle"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/QueueIsClosedV2.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/QueueIsClosedV2.pbtxt
index efe540401c9591..7cf1fde1bcee16 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/QueueIsClosedV2.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/QueueIsClosedV2.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "QueueIsClosedV2"
   input_arg {
     name: "handle"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/QueueSize.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/QueueSize.pbtxt
index 4dfe7aff427356..d2a49624f20ded 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/QueueSize.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/QueueSize.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "QueueSize"
   input_arg {
     name: "handle"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/QueueSizeV2.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/QueueSizeV2.pbtxt
index c7e11d7cf326e4..46eb229a3fa7e6 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/QueueSizeV2.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/QueueSizeV2.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "QueueSizeV2"
   input_arg {
     name: "handle"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/RFFT.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/RFFT.pbtxt
index 0879b0cfe5f2ae..02456ea217b791 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/RFFT.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/RFFT.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "RFFT"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/RFFT2D.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/RFFT2D.pbtxt
index 79ed06b41d4dea..f3676f45cde226 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/RFFT2D.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/RFFT2D.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "RFFT2D"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/RFFT3D.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/RFFT3D.pbtxt
index 754d372cf560ce..6475cd47316c19 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/RFFT3D.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/RFFT3D.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "RFFT3D"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/RFFTND.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/RFFTND.pbtxt
index 7f24e19efc06bf..8f75bea6baef82 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/RFFTND.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/RFFTND.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "RFFTND"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/RGBToHSV.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/RGBToHSV.pbtxt
index 22b6995682fb33..9ed50d337d0346 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/RGBToHSV.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/RGBToHSV.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "RGBToHSV"
   input_arg {
     name: "images"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/RaggedBincount.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/RaggedBincount.pbtxt
index 3fda312a6a4653..4f5fb24109cad4 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/RaggedBincount.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/RaggedBincount.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "RaggedBincount"
   input_arg {
     name: "splits"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/RaggedCountSparseOutput.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/RaggedCountSparseOutput.pbtxt
index 644b0eca3ad379..aa1a4e07aafaa2 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/RaggedCountSparseOutput.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/RaggedCountSparseOutput.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "RaggedCountSparseOutput"
   input_arg {
     name: "splits"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/RaggedCross.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/RaggedCross.pbtxt
index 2407315da2ee8a..0e9fe0adcfba2c 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/RaggedCross.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/RaggedCross.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "RaggedCross"
   input_arg {
     name: "ragged_values"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/RaggedFillEmptyRows.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/RaggedFillEmptyRows.pbtxt
index 54ba7be3cb1fdc..4587abc0f66e9a 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/RaggedFillEmptyRows.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/RaggedFillEmptyRows.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "RaggedFillEmptyRows"
   input_arg {
     name: "value_rowids"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/RaggedFillEmptyRowsGrad.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/RaggedFillEmptyRowsGrad.pbtxt
index 04f8b3e6f3a821..ea9546d504068f 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/RaggedFillEmptyRowsGrad.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/RaggedFillEmptyRowsGrad.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "RaggedFillEmptyRowsGrad"
   input_arg {
     name: "reverse_index_map"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/RaggedGather.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/RaggedGather.pbtxt
index 5e0e69e365a36c..afa14e8e8a8b20 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/RaggedGather.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/RaggedGather.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "RaggedGather"
   input_arg {
     name: "params_nested_splits"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/RaggedRange.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/RaggedRange.pbtxt
index e895110b6372a1..866c9b472d906d 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/RaggedRange.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/RaggedRange.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "RaggedRange"
   input_arg {
     name: "starts"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/RaggedTensorFromVariant.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/RaggedTensorFromVariant.pbtxt
index 50d8a8dd1b86db..5c750d0ce7e567 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/RaggedTensorFromVariant.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/RaggedTensorFromVariant.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "RaggedTensorFromVariant"
   input_arg {
     name: "encoded_ragged"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/RaggedTensorToSparse.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/RaggedTensorToSparse.pbtxt
index bcbb0d25bc94ff..f9172b4cf3772a 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/RaggedTensorToSparse.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/RaggedTensorToSparse.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "RaggedTensorToSparse"
   input_arg {
     name: "rt_nested_splits"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/RaggedTensorToTensor.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/RaggedTensorToTensor.pbtxt
index 8ecd7fb2c86ec3..60fceb565ebc03 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/RaggedTensorToTensor.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/RaggedTensorToTensor.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "RaggedTensorToTensor"
   input_arg {
     name: "shape"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/RaggedTensorToVariant.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/RaggedTensorToVariant.pbtxt
index 93d6b27fd053a1..f875a2a2ca7fae 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/RaggedTensorToVariant.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/RaggedTensorToVariant.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "RaggedTensorToVariant"
   input_arg {
     name: "rt_nested_splits"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/RaggedTensorToVariantGradient.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/RaggedTensorToVariantGradient.pbtxt
index fc6ede9bd741bf..45f2fcefe04210 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/RaggedTensorToVariantGradient.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/RaggedTensorToVariantGradient.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "RaggedTensorToVariantGradient"
   input_arg {
     name: "encoded_ragged_grad"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/RandomCrop.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/RandomCrop.pbtxt
index 9ca7d2f4b5b3bd..a5353cf58d5d41 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/RandomCrop.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/RandomCrop.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "RandomCrop"
   input_arg {
     name: "image"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/RandomDataset.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/RandomDataset.pbtxt
index 99fb23493250ca..d646d19b2e5ef2 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/RandomDataset.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/RandomDataset.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "RandomDataset"
   input_arg {
     name: "seed"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/RandomDatasetV2.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/RandomDatasetV2.pbtxt
index d561045e0b738d..3cf01f7644e9f1 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/RandomDatasetV2.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/RandomDatasetV2.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "RandomDatasetV2"
   input_arg {
     name: "seed"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/RandomGamma.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/RandomGamma.pbtxt
index 7c655027ce1832..2f38a20f8f0a2e 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/RandomGamma.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/RandomGamma.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "RandomGamma"
   input_arg {
     name: "shape"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/RandomGammaGrad.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/RandomGammaGrad.pbtxt
index f0c1b5033a5ddd..1e1c0723f6cbfe 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/RandomGammaGrad.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/RandomGammaGrad.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "RandomGammaGrad"
   input_arg {
     name: "alpha"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/RandomIndexShuffle.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/RandomIndexShuffle.pbtxt
index 06954c6cc639a3..22e6b88af0a921 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/RandomIndexShuffle.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/RandomIndexShuffle.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "RandomIndexShuffle"
   input_arg {
     name: "index"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/RandomPoisson.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/RandomPoisson.pbtxt
index 33d8920e36dcaa..5499e8d678c590 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/RandomPoisson.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/RandomPoisson.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "RandomPoisson"
   input_arg {
     name: "shape"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/RandomPoissonV2.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/RandomPoissonV2.pbtxt
index 1330833972ef23..6c3d9827a35cf8 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/RandomPoissonV2.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/RandomPoissonV2.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "RandomPoissonV2"
   input_arg {
     name: "shape"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/RandomShuffle.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/RandomShuffle.pbtxt
index 189e7d770a0652..ddd1a8d3f2b2c1 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/RandomShuffle.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/RandomShuffle.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "RandomShuffle"
   input_arg {
     name: "value"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/RandomShuffleQueue.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/RandomShuffleQueue.pbtxt
index 5793974674e561..550acae8d5be11 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/RandomShuffleQueue.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/RandomShuffleQueue.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "RandomShuffleQueue"
   output_arg {
     name: "handle"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/RandomShuffleQueueV2.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/RandomShuffleQueueV2.pbtxt
index b73d70541fafab..7d9807c4e9564b 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/RandomShuffleQueueV2.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/RandomShuffleQueueV2.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "RandomShuffleQueueV2"
   output_arg {
     name: "handle"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/RandomStandardNormal.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/RandomStandardNormal.pbtxt
index c693a9f7ba2e17..71fe5e5ef326b8 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/RandomStandardNormal.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/RandomStandardNormal.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "RandomStandardNormal"
   input_arg {
     name: "shape"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/RandomUniform.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/RandomUniform.pbtxt
index 63c637f0de1309..449a9ef973929c 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/RandomUniform.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/RandomUniform.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "RandomUniform"
   input_arg {
     name: "shape"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/RandomUniformInt.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/RandomUniformInt.pbtxt
index 35ad67cb371c9e..3b89715afcad98 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/RandomUniformInt.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/RandomUniformInt.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "RandomUniformInt"
   input_arg {
     name: "shape"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/Range.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/Range.pbtxt
index 6c2023eb68e78c..306116f59bfd42 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/Range.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/Range.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "Range"
   input_arg {
     name: "start"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/RangeDataset.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/RangeDataset.pbtxt
index 91fe787a7be594..9121cf0567e518 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/RangeDataset.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/RangeDataset.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "RangeDataset"
   input_arg {
     name: "start"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/Rank.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/Rank.pbtxt
index d44aeb9ec5fe4e..c12fd9a0abf07b 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/Rank.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/Rank.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "Rank"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ReadFile.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ReadFile.pbtxt
index 8bf1beeb43a473..ce1985ec3c5cba 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ReadFile.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ReadFile.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ReadFile"
   input_arg {
     name: "filename"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ReadVariableOp.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ReadVariableOp.pbtxt
index 2123cfedd6cb68..5459632d58351e 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ReadVariableOp.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ReadVariableOp.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ReadVariableOp"
   input_arg {
     name: "resource"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ReadVariableXlaSplitND.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ReadVariableXlaSplitND.pbtxt
index be5364ceb1f1b5..f5216da5c96aac 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ReadVariableXlaSplitND.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ReadVariableXlaSplitND.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ReadVariableXlaSplitND"
   input_arg {
     name: "resource"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ReaderNumRecordsProduced.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ReaderNumRecordsProduced.pbtxt
index 4514d3aac61611..50b1ea00da8f29 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ReaderNumRecordsProduced.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ReaderNumRecordsProduced.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ReaderNumRecordsProduced"
   input_arg {
     name: "reader_handle"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ReaderNumRecordsProducedV2.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ReaderNumRecordsProducedV2.pbtxt
index bff4305f3bfbb8..f560f01d443b08 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ReaderNumRecordsProducedV2.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ReaderNumRecordsProducedV2.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ReaderNumRecordsProducedV2"
   input_arg {
     name: "reader_handle"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ReaderNumWorkUnitsCompleted.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ReaderNumWorkUnitsCompleted.pbtxt
index 4ef2c1950f9507..b1e361e0119f22 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ReaderNumWorkUnitsCompleted.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ReaderNumWorkUnitsCompleted.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ReaderNumWorkUnitsCompleted"
   input_arg {
     name: "reader_handle"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ReaderNumWorkUnitsCompletedV2.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ReaderNumWorkUnitsCompletedV2.pbtxt
index 8f136052d08810..ee4c93e19bab64 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ReaderNumWorkUnitsCompletedV2.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ReaderNumWorkUnitsCompletedV2.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ReaderNumWorkUnitsCompletedV2"
   input_arg {
     name: "reader_handle"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ReaderRead.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ReaderRead.pbtxt
index fefc032a3460d9..b2a933892c4226 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ReaderRead.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ReaderRead.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ReaderRead"
   input_arg {
     name: "reader_handle"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ReaderReadUpTo.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ReaderReadUpTo.pbtxt
index 148095de5cae3a..e3bb64ec391d86 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ReaderReadUpTo.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ReaderReadUpTo.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ReaderReadUpTo"
   input_arg {
     name: "reader_handle"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ReaderReadUpToV2.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ReaderReadUpToV2.pbtxt
index f007588ebb608b..2ad62b16c9e928 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ReaderReadUpToV2.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ReaderReadUpToV2.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ReaderReadUpToV2"
   input_arg {
     name: "reader_handle"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ReaderReadV2.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ReaderReadV2.pbtxt
index 4e74f1b5b75080..3a1573147dcbe4 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ReaderReadV2.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ReaderReadV2.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ReaderReadV2"
   input_arg {
     name: "reader_handle"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ReaderReset.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ReaderReset.pbtxt
index 4379c400833e0f..9607f83c47016d 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ReaderReset.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ReaderReset.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ReaderReset"
   input_arg {
     name: "reader_handle"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ReaderResetV2.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ReaderResetV2.pbtxt
index eecd28bb726e79..56f862ae795184 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ReaderResetV2.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ReaderResetV2.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ReaderResetV2"
   input_arg {
     name: "reader_handle"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ReaderRestoreState.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ReaderRestoreState.pbtxt
index 1c5f71fe0f3521..717a5c34c8de93 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ReaderRestoreState.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ReaderRestoreState.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ReaderRestoreState"
   input_arg {
     name: "reader_handle"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ReaderRestoreStateV2.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ReaderRestoreStateV2.pbtxt
index 1b10902e081b7b..f75b04fc59d4e0 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ReaderRestoreStateV2.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ReaderRestoreStateV2.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ReaderRestoreStateV2"
   input_arg {
     name: "reader_handle"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ReaderSerializeState.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ReaderSerializeState.pbtxt
index d3a95d06f33347..2f708cb8926695 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ReaderSerializeState.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ReaderSerializeState.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ReaderSerializeState"
   input_arg {
     name: "reader_handle"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ReaderSerializeStateV2.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ReaderSerializeStateV2.pbtxt
index 4946352509bcb4..c4ade1409fbacd 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ReaderSerializeStateV2.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ReaderSerializeStateV2.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ReaderSerializeStateV2"
   input_arg {
     name: "reader_handle"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/Real.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/Real.pbtxt
index 0f326a3de147ed..d7e783ebe72f15 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/Real.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/Real.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "Real"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/RealDiv.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/RealDiv.pbtxt
index a5a78f14cc7c5a..6f725e22c64473 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/RealDiv.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/RealDiv.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "RealDiv"
   input_arg {
     name: "x"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/RebatchDataset.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/RebatchDataset.pbtxt
index bc70fcfadadb8e..176f94a3329c2b 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/RebatchDataset.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/RebatchDataset.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "RebatchDataset"
   input_arg {
     name: "input_dataset"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/RebatchDatasetV2.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/RebatchDatasetV2.pbtxt
index fda65d7bfaba18..fd853b65b2685c 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/RebatchDatasetV2.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/RebatchDatasetV2.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "RebatchDatasetV2"
   input_arg {
     name: "input_dataset"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/Reciprocal.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/Reciprocal.pbtxt
index 35cab85ee2a898..7e03554871a4eb 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/Reciprocal.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/Reciprocal.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "Reciprocal"
   input_arg {
     name: "x"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ReciprocalGrad.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ReciprocalGrad.pbtxt
index eea1bbe1c8d667..8884c796da5e6e 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ReciprocalGrad.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ReciprocalGrad.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ReciprocalGrad"
   input_arg {
     name: "x"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/RecordInput.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/RecordInput.pbtxt
index c353a17b318002..a72374420ee28e 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/RecordInput.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/RecordInput.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "RecordInput"
   output_arg {
     name: "records"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/Recv.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/Recv.pbtxt
index 87c7eb65ca87e8..e6717e57f8dedf 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/Recv.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/Recv.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "Recv"
   output_arg {
     name: "tensor"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/RecvTPUEmbeddingActivations.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/RecvTPUEmbeddingActivations.pbtxt
index 4cb9bd42ec045c..0fec828421f91b 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/RecvTPUEmbeddingActivations.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/RecvTPUEmbeddingActivations.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "RecvTPUEmbeddingActivations"
   output_arg {
     name: "outputs"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ReduceDataset.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ReduceDataset.pbtxt
index c98f1045e53af2..a1a17bc20c534d 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ReduceDataset.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ReduceDataset.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ReduceDataset"
   input_arg {
     name: "input_dataset"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ReduceJoin.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ReduceJoin.pbtxt
index 3603cf38c0b098..28880fcfb06a99 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ReduceJoin.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ReduceJoin.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ReduceJoin"
   input_arg {
     name: "inputs"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/RefEnter.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/RefEnter.pbtxt
index 8c7921571d9d24..9af599d1b338bb 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/RefEnter.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/RefEnter.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "RefEnter"
   input_arg {
     name: "data"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/RefExit.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/RefExit.pbtxt
index b13adfac00995a..1f9e84e7fade8b 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/RefExit.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/RefExit.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "RefExit"
   input_arg {
     name: "data"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/RefIdentity.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/RefIdentity.pbtxt
index 4fa1105b6952d0..d2293fdf467a7b 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/RefIdentity.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/RefIdentity.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "RefIdentity"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/RefMerge.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/RefMerge.pbtxt
index 31913cc794708d..fc4794d2f2f9d7 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/RefMerge.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/RefMerge.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "RefMerge"
   input_arg {
     name: "inputs"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/RefNextIteration.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/RefNextIteration.pbtxt
index 453bce7c335929..d447a3a87b23e4 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/RefNextIteration.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/RefNextIteration.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "RefNextIteration"
   input_arg {
     name: "data"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/RefSelect.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/RefSelect.pbtxt
index 688a34a5886fd4..aa2645f9ff17a3 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/RefSelect.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/RefSelect.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "RefSelect"
   input_arg {
     name: "index"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/RefSwitch.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/RefSwitch.pbtxt
index fba45dc4f98e00..6d12be2e5a8fa9 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/RefSwitch.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/RefSwitch.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "RefSwitch"
   input_arg {
     name: "data"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/RegexFullMatch.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/RegexFullMatch.pbtxt
index 7b337477adb60d..f2c0b7b99f1c71 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/RegexFullMatch.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/RegexFullMatch.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "RegexFullMatch"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/RegexReplace.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/RegexReplace.pbtxt
index d8e7b8c143f4be..591773ce37416d 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/RegexReplace.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/RegexReplace.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "RegexReplace"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/RegisterDataset.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/RegisterDataset.pbtxt
index 3cac6a9d2a0b6d..a144d96163e746 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/RegisterDataset.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/RegisterDataset.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "RegisterDataset"
   input_arg {
     name: "dataset"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/RegisterDatasetV2.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/RegisterDatasetV2.pbtxt
index 3031d65c39b18b..1ebced2af8d75f 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/RegisterDatasetV2.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/RegisterDatasetV2.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "RegisterDatasetV2"
   input_arg {
     name: "dataset"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/Relayout.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/Relayout.pbtxt
index 118574bf66fa05..da96a32cf86245 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/Relayout.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/Relayout.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "Relayout"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/RelayoutLike.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/RelayoutLike.pbtxt
index df5b2f89a73424..7a5af4f919e1c9 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/RelayoutLike.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/RelayoutLike.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "RelayoutLike"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/Relu.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/Relu.pbtxt
index 42ae6688c797ef..703fbbeff56f91 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/Relu.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/Relu.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "Relu"
   input_arg {
     name: "features"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/Relu6.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/Relu6.pbtxt
index 110c799682733e..311c3297a16411 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/Relu6.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/Relu6.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "Relu6"
   input_arg {
     name: "features"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/Relu6Grad.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/Relu6Grad.pbtxt
index 5455dc384cca0c..618e13a2297acc 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/Relu6Grad.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/Relu6Grad.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "Relu6Grad"
   input_arg {
     name: "gradients"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ReluGrad.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ReluGrad.pbtxt
index 89c486883f26c7..b14f23bb30bddc 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ReluGrad.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ReluGrad.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ReluGrad"
   input_arg {
     name: "gradients"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/RemoteCall.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/RemoteCall.pbtxt
index efe22174170af6..c6bc594510340c 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/RemoteCall.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/RemoteCall.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "RemoteCall"
   input_arg {
     name: "target"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/RepeatDataset.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/RepeatDataset.pbtxt
index f8f2cd8da0f50c..b7914feb4dfbc7 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/RepeatDataset.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/RepeatDataset.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "RepeatDataset"
   input_arg {
     name: "input_dataset"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/RequantizationRange.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/RequantizationRange.pbtxt
index 9ff7aa997d1a0c..6a489081eda2da 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/RequantizationRange.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/RequantizationRange.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "RequantizationRange"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/RequantizationRangePerChannel.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/RequantizationRangePerChannel.pbtxt
index f0352b164b215b..b621afb7a80b52 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/RequantizationRangePerChannel.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/RequantizationRangePerChannel.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "RequantizationRangePerChannel"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/Requantize.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/Requantize.pbtxt
index a142594170dc41..c04d32f19aeb16 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/Requantize.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/Requantize.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "Requantize"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/RequantizePerChannel.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/RequantizePerChannel.pbtxt
index 1e5b406fc0fb1e..3ed03fe12de763 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/RequantizePerChannel.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/RequantizePerChannel.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "RequantizePerChannel"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/Reshape.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/Reshape.pbtxt
index fcbd49a8ce805f..e422ffa2470633 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/Reshape.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/Reshape.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "Reshape"
   input_arg {
     name: "tensor"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ResizeArea.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ResizeArea.pbtxt
index 315069f44fa499..ec861564fc4cc6 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ResizeArea.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ResizeArea.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ResizeArea"
   input_arg {
     name: "images"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ResizeBicubic.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ResizeBicubic.pbtxt
index 2de5ed2e280f54..5e479a9432b876 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ResizeBicubic.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ResizeBicubic.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ResizeBicubic"
   input_arg {
     name: "images"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ResizeBicubicGrad.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ResizeBicubicGrad.pbtxt
index b085c1d6f46f67..6de227d7e0fd8f 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ResizeBicubicGrad.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ResizeBicubicGrad.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ResizeBicubicGrad"
   input_arg {
     name: "grads"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ResizeBilinear.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ResizeBilinear.pbtxt
index ea43150a420c04..b991a8ad63ddc8 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ResizeBilinear.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ResizeBilinear.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ResizeBilinear"
   input_arg {
     name: "images"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ResizeBilinearGrad.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ResizeBilinearGrad.pbtxt
index a78a0b699d7072..79d1605fd1c4a8 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ResizeBilinearGrad.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ResizeBilinearGrad.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ResizeBilinearGrad"
   input_arg {
     name: "grads"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ResizeNearestNeighbor.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ResizeNearestNeighbor.pbtxt
index a61928b06f38ff..18d87c3583c856 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ResizeNearestNeighbor.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ResizeNearestNeighbor.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ResizeNearestNeighbor"
   input_arg {
     name: "images"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ResizeNearestNeighborGrad.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ResizeNearestNeighborGrad.pbtxt
index 5bf5aa96eccc2b..001525b07337d6 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ResizeNearestNeighborGrad.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ResizeNearestNeighborGrad.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ResizeNearestNeighborGrad"
   input_arg {
     name: "grads"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ResourceAccumulatorApplyGradient.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ResourceAccumulatorApplyGradient.pbtxt
index 05ec4d8a404634..a1e03efd365a02 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ResourceAccumulatorApplyGradient.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ResourceAccumulatorApplyGradient.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ResourceAccumulatorApplyGradient"
   input_arg {
     name: "handle"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ResourceAccumulatorNumAccumulated.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ResourceAccumulatorNumAccumulated.pbtxt
index 5773bd5cc85b18..398171da210bf4 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ResourceAccumulatorNumAccumulated.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ResourceAccumulatorNumAccumulated.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ResourceAccumulatorNumAccumulated"
   input_arg {
     name: "handle"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ResourceAccumulatorSetGlobalStep.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ResourceAccumulatorSetGlobalStep.pbtxt
index 902748a29f5836..3e9c5a29882859 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ResourceAccumulatorSetGlobalStep.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ResourceAccumulatorSetGlobalStep.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ResourceAccumulatorSetGlobalStep"
   input_arg {
     name: "handle"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ResourceAccumulatorTakeGradient.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ResourceAccumulatorTakeGradient.pbtxt
index 5f37567706f7f0..d690fd198e62bd 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ResourceAccumulatorTakeGradient.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ResourceAccumulatorTakeGradient.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ResourceAccumulatorTakeGradient"
   input_arg {
     name: "handle"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ResourceApplyAdaMax.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ResourceApplyAdaMax.pbtxt
index ef87929c4f5eb6..3f307b70b7cac2 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ResourceApplyAdaMax.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ResourceApplyAdaMax.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ResourceApplyAdaMax"
   input_arg {
     name: "var"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ResourceApplyAdadelta.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ResourceApplyAdadelta.pbtxt
index 2a671e082a687f..f4c58b9860a033 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ResourceApplyAdadelta.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ResourceApplyAdadelta.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ResourceApplyAdadelta"
   input_arg {
     name: "var"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ResourceApplyAdagrad.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ResourceApplyAdagrad.pbtxt
index 67881c44ac854d..3c8e9d071a98b3 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ResourceApplyAdagrad.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ResourceApplyAdagrad.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ResourceApplyAdagrad"
   input_arg {
     name: "var"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ResourceApplyAdagradDA.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ResourceApplyAdagradDA.pbtxt
index 1bab7daf7aea46..996bc0b1a5cd71 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ResourceApplyAdagradDA.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ResourceApplyAdagradDA.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ResourceApplyAdagradDA"
   input_arg {
     name: "var"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ResourceApplyAdagradV2.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ResourceApplyAdagradV2.pbtxt
index 66aaa456bedf82..09168ff5c4e03b 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ResourceApplyAdagradV2.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ResourceApplyAdagradV2.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ResourceApplyAdagradV2"
   input_arg {
     name: "var"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ResourceApplyAdam.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ResourceApplyAdam.pbtxt
index e76225de60701e..88cafe521d6800 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ResourceApplyAdam.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ResourceApplyAdam.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ResourceApplyAdam"
   input_arg {
     name: "var"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ResourceApplyAdamWithAmsgrad.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ResourceApplyAdamWithAmsgrad.pbtxt
index ab8696a3263a89..ccd84e52d5a820 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ResourceApplyAdamWithAmsgrad.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ResourceApplyAdamWithAmsgrad.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ResourceApplyAdamWithAmsgrad"
   input_arg {
     name: "var"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ResourceApplyAddSign.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ResourceApplyAddSign.pbtxt
index 44f5b7d88e0c07..4653229b0aebf1 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ResourceApplyAddSign.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ResourceApplyAddSign.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ResourceApplyAddSign"
   input_arg {
     name: "var"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ResourceApplyCenteredRMSProp.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ResourceApplyCenteredRMSProp.pbtxt
index b7c69b1c832fff..f6fa5b2a351164 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ResourceApplyCenteredRMSProp.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ResourceApplyCenteredRMSProp.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ResourceApplyCenteredRMSProp"
   input_arg {
     name: "var"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ResourceApplyFtrl.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ResourceApplyFtrl.pbtxt
index f94944686d5a14..f917b0fc22fa6f 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ResourceApplyFtrl.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ResourceApplyFtrl.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ResourceApplyFtrl"
   input_arg {
     name: "var"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ResourceApplyFtrlV2.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ResourceApplyFtrlV2.pbtxt
index 597ce4ab164aa0..d5586b7f73de74 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ResourceApplyFtrlV2.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ResourceApplyFtrlV2.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ResourceApplyFtrlV2"
   input_arg {
     name: "var"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ResourceApplyGradientDescent.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ResourceApplyGradientDescent.pbtxt
index 6bd3170e2c6bde..83639c23ea1d7d 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ResourceApplyGradientDescent.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ResourceApplyGradientDescent.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ResourceApplyGradientDescent"
   input_arg {
     name: "var"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ResourceApplyKerasMomentum.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ResourceApplyKerasMomentum.pbtxt
index 12e22d4167d7dc..737020cd744516 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ResourceApplyKerasMomentum.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ResourceApplyKerasMomentum.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ResourceApplyKerasMomentum"
   input_arg {
     name: "var"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ResourceApplyMomentum.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ResourceApplyMomentum.pbtxt
index 20de47ab2895d0..a6d0dd9243836e 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ResourceApplyMomentum.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ResourceApplyMomentum.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ResourceApplyMomentum"
   input_arg {
     name: "var"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ResourceApplyPowerSign.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ResourceApplyPowerSign.pbtxt
index 261463a5512584..5a7e103c707c04 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ResourceApplyPowerSign.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ResourceApplyPowerSign.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ResourceApplyPowerSign"
   input_arg {
     name: "var"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ResourceApplyProximalAdagrad.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ResourceApplyProximalAdagrad.pbtxt
index 2c6007597ca3b1..c3adae4e5fc900 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ResourceApplyProximalAdagrad.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ResourceApplyProximalAdagrad.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ResourceApplyProximalAdagrad"
   input_arg {
     name: "var"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ResourceApplyProximalGradientDescent.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ResourceApplyProximalGradientDescent.pbtxt
index dbe02a88ff079d..9c55be37dfc8da 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ResourceApplyProximalGradientDescent.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ResourceApplyProximalGradientDescent.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ResourceApplyProximalGradientDescent"
   input_arg {
     name: "var"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ResourceApplyRMSProp.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ResourceApplyRMSProp.pbtxt
index 90f24a83fb81e7..b0be4acdf19643 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ResourceApplyRMSProp.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ResourceApplyRMSProp.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ResourceApplyRMSProp"
   input_arg {
     name: "var"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ResourceConditionalAccumulator.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ResourceConditionalAccumulator.pbtxt
index 389486faef4e8e..cf3c150ccb9984 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ResourceConditionalAccumulator.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ResourceConditionalAccumulator.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ResourceConditionalAccumulator"
   output_arg {
     name: "handle"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ResourceCountUpTo.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ResourceCountUpTo.pbtxt
index 3f07aa17a613e4..352935c2c167a4 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ResourceCountUpTo.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ResourceCountUpTo.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ResourceCountUpTo"
   input_arg {
     name: "resource"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ResourceGather.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ResourceGather.pbtxt
index 47a841492de645..9aa33d994c6cea 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ResourceGather.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ResourceGather.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ResourceGather"
   input_arg {
     name: "resource"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ResourceGatherNd.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ResourceGatherNd.pbtxt
index f8df557edbbdbb..04794f402cde00 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ResourceGatherNd.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ResourceGatherNd.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ResourceGatherNd"
   input_arg {
     name: "resource"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ResourceScatterAdd.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ResourceScatterAdd.pbtxt
index 5b10cf3f16c0ce..a061da9db360bc 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ResourceScatterAdd.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ResourceScatterAdd.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ResourceScatterAdd"
   input_arg {
     name: "resource"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ResourceScatterDiv.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ResourceScatterDiv.pbtxt
index 0b94ef0dec43bd..d5b683c107ef47 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ResourceScatterDiv.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ResourceScatterDiv.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ResourceScatterDiv"
   input_arg {
     name: "resource"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ResourceScatterMax.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ResourceScatterMax.pbtxt
index d6e97844047d95..d7f1e9c8a37e69 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ResourceScatterMax.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ResourceScatterMax.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ResourceScatterMax"
   input_arg {
     name: "resource"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ResourceScatterMin.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ResourceScatterMin.pbtxt
index d012a861218190..617f12b35c83a7 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ResourceScatterMin.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ResourceScatterMin.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ResourceScatterMin"
   input_arg {
     name: "resource"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ResourceScatterMul.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ResourceScatterMul.pbtxt
index aa859ad4d252fe..313b2ed03c038b 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ResourceScatterMul.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ResourceScatterMul.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ResourceScatterMul"
   input_arg {
     name: "resource"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ResourceScatterNdAdd.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ResourceScatterNdAdd.pbtxt
index 5cc51a5559e4a0..507b30eff9a2d6 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ResourceScatterNdAdd.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ResourceScatterNdAdd.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ResourceScatterNdAdd"
   input_arg {
     name: "ref"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ResourceScatterNdMax.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ResourceScatterNdMax.pbtxt
index 9933b54896edff..8a40ac6e842a03 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ResourceScatterNdMax.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ResourceScatterNdMax.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ResourceScatterNdMax"
   input_arg {
     name: "ref"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ResourceScatterNdMin.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ResourceScatterNdMin.pbtxt
index bd0e7589184074..7567b43d39f025 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ResourceScatterNdMin.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ResourceScatterNdMin.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ResourceScatterNdMin"
   input_arg {
     name: "ref"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ResourceScatterNdSub.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ResourceScatterNdSub.pbtxt
index bd5785e83d4bf8..9d1a74daa931f8 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ResourceScatterNdSub.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ResourceScatterNdSub.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ResourceScatterNdSub"
   input_arg {
     name: "ref"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ResourceScatterNdUpdate.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ResourceScatterNdUpdate.pbtxt
index 0b6ae92e036d84..4305163fc92280 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ResourceScatterNdUpdate.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ResourceScatterNdUpdate.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ResourceScatterNdUpdate"
   input_arg {
     name: "ref"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ResourceScatterSub.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ResourceScatterSub.pbtxt
index 32257603e3aeef..3a267f9f63b289 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ResourceScatterSub.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ResourceScatterSub.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ResourceScatterSub"
   input_arg {
     name: "resource"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ResourceScatterUpdate.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ResourceScatterUpdate.pbtxt
index 320a7b0285e79d..55101b84c33871 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ResourceScatterUpdate.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ResourceScatterUpdate.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ResourceScatterUpdate"
   input_arg {
     name: "resource"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ResourceSparseApplyAdadelta.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ResourceSparseApplyAdadelta.pbtxt
index 21aa0addab9259..562e667d584b04 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ResourceSparseApplyAdadelta.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ResourceSparseApplyAdadelta.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ResourceSparseApplyAdadelta"
   input_arg {
     name: "var"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ResourceSparseApplyAdagrad.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ResourceSparseApplyAdagrad.pbtxt
index 7cc9e1f5cd2927..4588dc668bd9b4 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ResourceSparseApplyAdagrad.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ResourceSparseApplyAdagrad.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ResourceSparseApplyAdagrad"
   input_arg {
     name: "var"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ResourceSparseApplyAdagradDA.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ResourceSparseApplyAdagradDA.pbtxt
index c28ddc89e06058..e9ef6077dcedf7 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ResourceSparseApplyAdagradDA.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ResourceSparseApplyAdagradDA.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ResourceSparseApplyAdagradDA"
   input_arg {
     name: "var"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ResourceSparseApplyAdagradV2.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ResourceSparseApplyAdagradV2.pbtxt
index a14ba1dc5d2d55..e0724d11142c5a 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ResourceSparseApplyAdagradV2.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ResourceSparseApplyAdagradV2.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ResourceSparseApplyAdagradV2"
   input_arg {
     name: "var"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ResourceSparseApplyCenteredRMSProp.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ResourceSparseApplyCenteredRMSProp.pbtxt
index 855c9982f3f6f5..6d02f81629efda 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ResourceSparseApplyCenteredRMSProp.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ResourceSparseApplyCenteredRMSProp.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ResourceSparseApplyCenteredRMSProp"
   input_arg {
     name: "var"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ResourceSparseApplyFtrl.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ResourceSparseApplyFtrl.pbtxt
index 7bb28e90106812..8ac38c09e514b4 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ResourceSparseApplyFtrl.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ResourceSparseApplyFtrl.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ResourceSparseApplyFtrl"
   input_arg {
     name: "var"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ResourceSparseApplyFtrlV2.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ResourceSparseApplyFtrlV2.pbtxt
index c43d4a60f27998..cff68beaf1cbe7 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ResourceSparseApplyFtrlV2.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ResourceSparseApplyFtrlV2.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ResourceSparseApplyFtrlV2"
   input_arg {
     name: "var"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ResourceSparseApplyKerasMomentum.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ResourceSparseApplyKerasMomentum.pbtxt
index 96b22b6cf0bf4a..648d85a4989da1 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ResourceSparseApplyKerasMomentum.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ResourceSparseApplyKerasMomentum.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ResourceSparseApplyKerasMomentum"
   input_arg {
     name: "var"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ResourceSparseApplyMomentum.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ResourceSparseApplyMomentum.pbtxt
index 03b05570cb5ccd..ec3e5a96db592b 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ResourceSparseApplyMomentum.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ResourceSparseApplyMomentum.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ResourceSparseApplyMomentum"
   input_arg {
     name: "var"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ResourceSparseApplyProximalAdagrad.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ResourceSparseApplyProximalAdagrad.pbtxt
index 0093fb2c50f763..59ae118cd9522d 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ResourceSparseApplyProximalAdagrad.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ResourceSparseApplyProximalAdagrad.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ResourceSparseApplyProximalAdagrad"
   input_arg {
     name: "var"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ResourceSparseApplyProximalGradientDescent.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ResourceSparseApplyProximalGradientDescent.pbtxt
index 0844c8d434921c..246f442d27756e 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ResourceSparseApplyProximalGradientDescent.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ResourceSparseApplyProximalGradientDescent.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ResourceSparseApplyProximalGradientDescent"
   input_arg {
     name: "var"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ResourceSparseApplyRMSProp.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ResourceSparseApplyRMSProp.pbtxt
index 87803f3e89366f..8735a6fcacc18e 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ResourceSparseApplyRMSProp.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ResourceSparseApplyRMSProp.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ResourceSparseApplyRMSProp"
   input_arg {
     name: "var"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ResourceStridedSliceAssign.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ResourceStridedSliceAssign.pbtxt
index 430195fab5c0d0..867f205958e76c 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ResourceStridedSliceAssign.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ResourceStridedSliceAssign.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ResourceStridedSliceAssign"
   input_arg {
     name: "ref"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/Restore.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/Restore.pbtxt
index 269b27c0b71f59..1db029076032ac 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/Restore.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/Restore.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "Restore"
   input_arg {
     name: "file_pattern"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/RestoreSlice.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/RestoreSlice.pbtxt
index 74303c89f05152..03d2aa3bbf56ff 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/RestoreSlice.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/RestoreSlice.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "RestoreSlice"
   input_arg {
     name: "file_pattern"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/RestoreV2.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/RestoreV2.pbtxt
index 99319760e597e7..a88db314bc4906 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/RestoreV2.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/RestoreV2.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "RestoreV2"
   input_arg {
     name: "prefix"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/RetrieveAllTPUEmbeddingParameters.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/RetrieveAllTPUEmbeddingParameters.pbtxt
index ea305f05a813bb..226bed39734bf1 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/RetrieveAllTPUEmbeddingParameters.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/RetrieveAllTPUEmbeddingParameters.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "RetrieveAllTPUEmbeddingParameters"
   output_arg {
     name: "parameters"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/RetrieveTPUEmbeddingADAMParameters.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/RetrieveTPUEmbeddingADAMParameters.pbtxt
index adfc5f50337be2..dfef40371c8098 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/RetrieveTPUEmbeddingADAMParameters.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/RetrieveTPUEmbeddingADAMParameters.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "RetrieveTPUEmbeddingADAMParameters"
   output_arg {
     name: "parameters"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/RetrieveTPUEmbeddingAdadeltaParameters.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/RetrieveTPUEmbeddingAdadeltaParameters.pbtxt
index 3c514f34bde42c..f887714c78b20e 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/RetrieveTPUEmbeddingAdadeltaParameters.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/RetrieveTPUEmbeddingAdadeltaParameters.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "RetrieveTPUEmbeddingAdadeltaParameters"
   output_arg {
     name: "parameters"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/RetrieveTPUEmbeddingAdagradMomentumParameters.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/RetrieveTPUEmbeddingAdagradMomentumParameters.pbtxt
index 747aecc018ccc4..9a84dc7b0cc57f 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/RetrieveTPUEmbeddingAdagradMomentumParameters.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/RetrieveTPUEmbeddingAdagradMomentumParameters.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "RetrieveTPUEmbeddingAdagradMomentumParameters"
   output_arg {
     name: "parameters"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/RetrieveTPUEmbeddingAdagradParameters.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/RetrieveTPUEmbeddingAdagradParameters.pbtxt
index cc9f172411d370..27e1aa441b60ae 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/RetrieveTPUEmbeddingAdagradParameters.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/RetrieveTPUEmbeddingAdagradParameters.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "RetrieveTPUEmbeddingAdagradParameters"
   output_arg {
     name: "parameters"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/RetrieveTPUEmbeddingCenteredRMSPropParameters.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/RetrieveTPUEmbeddingCenteredRMSPropParameters.pbtxt
index 6ce83985ca788f..afdcbd5cb5e7c9 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/RetrieveTPUEmbeddingCenteredRMSPropParameters.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/RetrieveTPUEmbeddingCenteredRMSPropParameters.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "RetrieveTPUEmbeddingCenteredRMSPropParameters"
   output_arg {
     name: "parameters"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/RetrieveTPUEmbeddingFTRLParameters.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/RetrieveTPUEmbeddingFTRLParameters.pbtxt
index d0d5fd74a5f8b1..161fde82db383b 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/RetrieveTPUEmbeddingFTRLParameters.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/RetrieveTPUEmbeddingFTRLParameters.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "RetrieveTPUEmbeddingFTRLParameters"
   output_arg {
     name: "parameters"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/RetrieveTPUEmbeddingFrequencyEstimatorParameters.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/RetrieveTPUEmbeddingFrequencyEstimatorParameters.pbtxt
index 1125ad0e594c68..633bf51cab0d68 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/RetrieveTPUEmbeddingFrequencyEstimatorParameters.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/RetrieveTPUEmbeddingFrequencyEstimatorParameters.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "RetrieveTPUEmbeddingFrequencyEstimatorParameters"
   output_arg {
     name: "parameters"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/RetrieveTPUEmbeddingMDLAdagradLightParameters.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/RetrieveTPUEmbeddingMDLAdagradLightParameters.pbtxt
index c2e4047eacfc4a..bf9d96bf37dea8 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/RetrieveTPUEmbeddingMDLAdagradLightParameters.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/RetrieveTPUEmbeddingMDLAdagradLightParameters.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "RetrieveTPUEmbeddingMDLAdagradLightParameters"
   output_arg {
     name: "parameters"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/RetrieveTPUEmbeddingMomentumParameters.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/RetrieveTPUEmbeddingMomentumParameters.pbtxt
index 6db5898c1e3a9a..3b8f98b755312f 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/RetrieveTPUEmbeddingMomentumParameters.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/RetrieveTPUEmbeddingMomentumParameters.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "RetrieveTPUEmbeddingMomentumParameters"
   output_arg {
     name: "parameters"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/RetrieveTPUEmbeddingProximalAdagradParameters.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/RetrieveTPUEmbeddingProximalAdagradParameters.pbtxt
index 5cf6dbfe8d767a..bd7c8288ff2af4 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/RetrieveTPUEmbeddingProximalAdagradParameters.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/RetrieveTPUEmbeddingProximalAdagradParameters.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "RetrieveTPUEmbeddingProximalAdagradParameters"
   output_arg {
     name: "parameters"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/RetrieveTPUEmbeddingProximalYogiParameters.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/RetrieveTPUEmbeddingProximalYogiParameters.pbtxt
index dedd0ee1826811..a74c0caa0862ad 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/RetrieveTPUEmbeddingProximalYogiParameters.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/RetrieveTPUEmbeddingProximalYogiParameters.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "RetrieveTPUEmbeddingProximalYogiParameters"
   output_arg {
     name: "parameters"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/RetrieveTPUEmbeddingRMSPropParameters.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/RetrieveTPUEmbeddingRMSPropParameters.pbtxt
index d4028e2d164b75..137d3bd497dd76 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/RetrieveTPUEmbeddingRMSPropParameters.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/RetrieveTPUEmbeddingRMSPropParameters.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "RetrieveTPUEmbeddingRMSPropParameters"
   output_arg {
     name: "parameters"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/RetrieveTPUEmbeddingStochasticGradientDescentParameters.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/RetrieveTPUEmbeddingStochasticGradientDescentParameters.pbtxt
index 44b3a038b1ca5e..da65f5ae98e542 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/RetrieveTPUEmbeddingStochasticGradientDescentParameters.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/RetrieveTPUEmbeddingStochasticGradientDescentParameters.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "RetrieveTPUEmbeddingStochasticGradientDescentParameters"
   output_arg {
     name: "parameters"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/Reverse.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/Reverse.pbtxt
index 9eced34f0a50d9..aecf71f072f0cd 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/Reverse.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/Reverse.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "Reverse"
   input_arg {
     name: "tensor"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ReverseSequence.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ReverseSequence.pbtxt
index c3119a2fe49707..74d3601e1f4aba 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ReverseSequence.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ReverseSequence.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ReverseSequence"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ReverseV2.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ReverseV2.pbtxt
index fc9a980a8712bc..8cbc98e30a5390 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ReverseV2.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ReverseV2.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ReverseV2"
   input_arg {
     name: "tensor"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/RewriteDataset.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/RewriteDataset.pbtxt
index d831b7cac3f549..3e3e43b46fed65 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/RewriteDataset.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/RewriteDataset.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "RewriteDataset"
   input_arg {
     name: "input_dataset"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/RightShift.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/RightShift.pbtxt
index 76bcdc4bd20b38..97257a019725a2 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/RightShift.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/RightShift.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "RightShift"
   input_arg {
     name: "x"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/Rint.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/Rint.pbtxt
index f128a08a47a130..feed3bca0b3ef1 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/Rint.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/Rint.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "Rint"
   input_arg {
     name: "x"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/RiscAbs.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/RiscAbs.pbtxt
index a1d64124f38587..2b3111bf580739 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/RiscAbs.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/RiscAbs.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "RiscAbs"
   input_arg {
     name: "x"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/RiscAdd.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/RiscAdd.pbtxt
index d98094d312b9ad..86cc81e964c101 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/RiscAdd.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/RiscAdd.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "RiscAdd"
   input_arg {
     name: "x"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/RiscBinaryArithmetic.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/RiscBinaryArithmetic.pbtxt
index f2919e4c26cbf9..9d5f080ed02465 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/RiscBinaryArithmetic.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/RiscBinaryArithmetic.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "RiscBinaryArithmetic"
   input_arg {
     name: "x"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/RiscBinaryComparison.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/RiscBinaryComparison.pbtxt
index 88731716d93946..d131476f8fe944 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/RiscBinaryComparison.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/RiscBinaryComparison.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "RiscBinaryComparison"
   input_arg {
     name: "x"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/RiscBitcast.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/RiscBitcast.pbtxt
index 738659831a958d..1d37369adec753 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/RiscBitcast.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/RiscBitcast.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "RiscBitcast"
   input_arg {
     name: "x"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/RiscBroadcast.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/RiscBroadcast.pbtxt
index 61318f39439e77..e81e8413dfb88b 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/RiscBroadcast.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/RiscBroadcast.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "RiscBroadcast"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/RiscCast.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/RiscCast.pbtxt
index 2b82bc3ddc9939..344d0496b27962 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/RiscCast.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/RiscCast.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "RiscCast"
   input_arg {
     name: "x"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/RiscCeil.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/RiscCeil.pbtxt
index 904aca9f50b86d..ff1fefc0115839 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/RiscCeil.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/RiscCeil.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "RiscCeil"
   input_arg {
     name: "x"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/RiscCholesky.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/RiscCholesky.pbtxt
index 3b8d59b6efac76..c6b24d107be7f3 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/RiscCholesky.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/RiscCholesky.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "RiscCholesky"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/RiscConcat.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/RiscConcat.pbtxt
index 0b062957805f5f..889de347165b26 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/RiscConcat.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/RiscConcat.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "RiscConcat"
   input_arg {
     name: "values"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/RiscCondition.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/RiscCondition.pbtxt
index a348796ece3143..859814c65905b2 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/RiscCondition.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/RiscCondition.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "RiscCondition"
   input_arg {
     name: "pred"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/RiscConv.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/RiscConv.pbtxt
index 06103b7729dab5..e9d326be8f286f 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/RiscConv.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/RiscConv.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "RiscConv"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/RiscCos.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/RiscCos.pbtxt
index d89ad75fc9fa25..98cebc12c5a9ee 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/RiscCos.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/RiscCos.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "RiscCos"
   input_arg {
     name: "x"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/RiscDiv.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/RiscDiv.pbtxt
index d2196f83d64b61..98aaf5ae97b8c5 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/RiscDiv.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/RiscDiv.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "RiscDiv"
   input_arg {
     name: "x"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/RiscDot.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/RiscDot.pbtxt
index 55101191acb75f..0b3ff5c11e3646 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/RiscDot.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/RiscDot.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "RiscDot"
   input_arg {
     name: "a"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/RiscExp.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/RiscExp.pbtxt
index 9775069c7f693e..4386db96ae9960 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/RiscExp.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/RiscExp.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "RiscExp"
   input_arg {
     name: "x"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/RiscFft.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/RiscFft.pbtxt
index de6960750f34ef..605cd7edd6e04e 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/RiscFft.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/RiscFft.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "RiscFft"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/RiscFloor.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/RiscFloor.pbtxt
index 71b73dca0fe261..55941ca352c2ec 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/RiscFloor.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/RiscFloor.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "RiscFloor"
   input_arg {
     name: "x"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/RiscGather.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/RiscGather.pbtxt
index 87681965ec7f3d..18d4ba3fdc90d4 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/RiscGather.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/RiscGather.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "RiscGather"
   input_arg {
     name: "params"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/RiscImag.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/RiscImag.pbtxt
index c6e17d5296d93e..555ea9d071fe6a 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/RiscImag.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/RiscImag.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "RiscImag"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/RiscIsFinite.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/RiscIsFinite.pbtxt
index 1b6574899d378d..19a4ae6617c0cb 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/RiscIsFinite.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/RiscIsFinite.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "RiscIsFinite"
   input_arg {
     name: "x"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/RiscLog.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/RiscLog.pbtxt
index e52bc8fbe8ad02..23bbef3b07a3b0 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/RiscLog.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/RiscLog.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "RiscLog"
   input_arg {
     name: "x"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/RiscLogicalAnd.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/RiscLogicalAnd.pbtxt
index 04777105ffae7f..8bd4410a056174 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/RiscLogicalAnd.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/RiscLogicalAnd.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "RiscLogicalAnd"
   input_arg {
     name: "x"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/RiscLogicalNot.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/RiscLogicalNot.pbtxt
index d743f6c5f935bf..3496ef02e435a8 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/RiscLogicalNot.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/RiscLogicalNot.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "RiscLogicalNot"
   input_arg {
     name: "x"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/RiscLogicalOr.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/RiscLogicalOr.pbtxt
index 5550a34f7d9bc1..3cf31921d5aa0c 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/RiscLogicalOr.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/RiscLogicalOr.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "RiscLogicalOr"
   input_arg {
     name: "x"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/RiscMax.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/RiscMax.pbtxt
index 9ecd924242589f..11c4517d7566b3 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/RiscMax.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/RiscMax.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "RiscMax"
   input_arg {
     name: "x"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/RiscMin.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/RiscMin.pbtxt
index 225a66f5f952de..7ac92ec5ca4a12 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/RiscMin.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/RiscMin.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "RiscMin"
   input_arg {
     name: "x"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/RiscMul.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/RiscMul.pbtxt
index 787d58e56f074d..e55fca7a910539 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/RiscMul.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/RiscMul.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "RiscMul"
   input_arg {
     name: "x"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/RiscNeg.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/RiscNeg.pbtxt
index 2d4a2b3425a37c..429d94153c2490 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/RiscNeg.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/RiscNeg.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "RiscNeg"
   input_arg {
     name: "x"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/RiscPad.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/RiscPad.pbtxt
index 3707792d018c3c..13ea65b0bd3974 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/RiscPad.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/RiscPad.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "RiscPad"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/RiscPool.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/RiscPool.pbtxt
index 766557c3b401f4..57847ff00d0cb4 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/RiscPool.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/RiscPool.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "RiscPool"
   input_arg {
     name: "value"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/RiscPow.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/RiscPow.pbtxt
index 89d58894ce23f3..150c846fb0396f 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/RiscPow.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/RiscPow.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "RiscPow"
   input_arg {
     name: "x"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/RiscRandomUniform.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/RiscRandomUniform.pbtxt
index 29261856fe30a6..2d3cd00a70c0ea 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/RiscRandomUniform.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/RiscRandomUniform.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "RiscRandomUniform"
   input_arg {
     name: "shape"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/RiscReal.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/RiscReal.pbtxt
index 5c299538003771..cf624262303984 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/RiscReal.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/RiscReal.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "RiscReal"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/RiscReduce.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/RiscReduce.pbtxt
index fc067edf50fa64..1dff780022c33e 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/RiscReduce.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/RiscReduce.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "RiscReduce"
   input_arg {
     name: "tensor"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/RiscRem.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/RiscRem.pbtxt
index 5392ef47912f15..ffa8ff1bc61e82 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/RiscRem.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/RiscRem.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "RiscRem"
   input_arg {
     name: "x"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/RiscReshape.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/RiscReshape.pbtxt
index a4508f5347d382..f686036497b850 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/RiscReshape.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/RiscReshape.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "RiscReshape"
   input_arg {
     name: "tensor"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/RiscReverse.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/RiscReverse.pbtxt
index 3ee303ebc4ea2c..60dec4d9b9d64b 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/RiscReverse.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/RiscReverse.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "RiscReverse"
   input_arg {
     name: "tensor"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/RiscScatter.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/RiscScatter.pbtxt
index d37788f996fb13..5def9d11dc486b 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/RiscScatter.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/RiscScatter.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "RiscScatter"
   input_arg {
     name: "indices"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/RiscShape.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/RiscShape.pbtxt
index a9e95b87a8abd1..615e2b7e8bbba9 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/RiscShape.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/RiscShape.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "RiscShape"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/RiscSign.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/RiscSign.pbtxt
index c3478d94634773..db7894468abf7b 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/RiscSign.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/RiscSign.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "RiscSign"
   input_arg {
     name: "x"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/RiscSlice.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/RiscSlice.pbtxt
index afdc888c886ef2..b09072c5ae7f08 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/RiscSlice.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/RiscSlice.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "RiscSlice"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/RiscSort.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/RiscSort.pbtxt
index 60ca56bc40cab6..c49a6951bdbeff 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/RiscSort.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/RiscSort.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "RiscSort"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/RiscSqueeze.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/RiscSqueeze.pbtxt
index 0c70b1088f6aba..bf4e7123ddabc0 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/RiscSqueeze.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/RiscSqueeze.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "RiscSqueeze"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/RiscSub.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/RiscSub.pbtxt
index 8a83934419e223..2590f9fa34cce6 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/RiscSub.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/RiscSub.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "RiscSub"
   input_arg {
     name: "x"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/RiscTranspose.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/RiscTranspose.pbtxt
index eb7dce89a9aa35..856b0d67ffb98c 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/RiscTranspose.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/RiscTranspose.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "RiscTranspose"
   input_arg {
     name: "x"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/RiscTriangularSolve.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/RiscTriangularSolve.pbtxt
index c2095295f40a00..5b8518fc9bedd7 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/RiscTriangularSolve.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/RiscTriangularSolve.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "RiscTriangularSolve"
   input_arg {
     name: "matrix"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/RiscUnary.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/RiscUnary.pbtxt
index d825acd378f0b8..0a7af35478abb0 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/RiscUnary.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/RiscUnary.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "RiscUnary"
   input_arg {
     name: "x"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/RiscWhile.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/RiscWhile.pbtxt
index 9e4695be194b3f..8bb4745a1e135c 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/RiscWhile.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/RiscWhile.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "RiscWhile"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/RngReadAndSkip.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/RngReadAndSkip.pbtxt
index d369d4dc9ae400..e64f5dd6d2680c 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/RngReadAndSkip.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/RngReadAndSkip.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "RngReadAndSkip"
   input_arg {
     name: "resource"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/RngSkip.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/RngSkip.pbtxt
index 6ef7c4b9ae69ea..dc3e9b948b7ed6 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/RngSkip.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/RngSkip.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "RngSkip"
   input_arg {
     name: "resource"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/Roll.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/Roll.pbtxt
index cfb7c101757c7b..ac81404fece17e 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/Roll.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/Roll.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "Roll"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/Round.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/Round.pbtxt
index bb883d57b9a84e..c5685dc61439b1 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/Round.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/Round.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "Round"
   input_arg {
     name: "x"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/Rsqrt.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/Rsqrt.pbtxt
index 87e8420b1e96f6..6d066c9e00cee9 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/Rsqrt.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/Rsqrt.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "Rsqrt"
   input_arg {
     name: "x"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/RsqrtGrad.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/RsqrtGrad.pbtxt
index 9298cb8a73e1dd..4509b1af361e3d 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/RsqrtGrad.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/RsqrtGrad.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "RsqrtGrad"
   input_arg {
     name: "x"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/SampleDistortedBoundingBox.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/SampleDistortedBoundingBox.pbtxt
index 527d0b2f5000c6..95b4a2ddd5f95f 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/SampleDistortedBoundingBox.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/SampleDistortedBoundingBox.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "SampleDistortedBoundingBox"
   input_arg {
     name: "image_size"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/SampleDistortedBoundingBoxV2.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/SampleDistortedBoundingBoxV2.pbtxt
index 3600eaf126b48b..d857ee0a68795d 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/SampleDistortedBoundingBoxV2.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/SampleDistortedBoundingBoxV2.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "SampleDistortedBoundingBoxV2"
   input_arg {
     name: "image_size"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/SamplingDataset.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/SamplingDataset.pbtxt
index 549b71b08c77fd..160a9e9bb16588 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/SamplingDataset.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/SamplingDataset.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "SamplingDataset"
   input_arg {
     name: "input_dataset"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/Save.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/Save.pbtxt
index c815380c490ad8..c632730bd91e7e 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/Save.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/Save.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "Save"
   input_arg {
     name: "filename"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/SaveDataset.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/SaveDataset.pbtxt
index 8b5b8fcbd1dbc7..18f99fbe3591ee 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/SaveDataset.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/SaveDataset.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "SaveDataset"
   input_arg {
     name: "input_dataset"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/SaveDatasetV2.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/SaveDatasetV2.pbtxt
index e20303ad3b8d21..9407de7812bb12 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/SaveDatasetV2.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/SaveDatasetV2.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "SaveDatasetV2"
   input_arg {
     name: "input_dataset"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/SaveSlices.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/SaveSlices.pbtxt
index 4415c13dbb8f69..306d67bd688456 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/SaveSlices.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/SaveSlices.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "SaveSlices"
   input_arg {
     name: "filename"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/SaveV2.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/SaveV2.pbtxt
index c91f4d3159d0d3..d9bae4c8b8e12b 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/SaveV2.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/SaveV2.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "SaveV2"
   input_arg {
     name: "prefix"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ScalarSummary.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ScalarSummary.pbtxt
index 48fa50e9382251..bf4948076abe9c 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ScalarSummary.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ScalarSummary.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ScalarSummary"
   input_arg {
     name: "tags"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ScaleAndTranslate.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ScaleAndTranslate.pbtxt
index 4c3b0727ac80b0..516cca34539b9b 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ScaleAndTranslate.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ScaleAndTranslate.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ScaleAndTranslate"
   input_arg {
     name: "images"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ScaleAndTranslateGrad.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ScaleAndTranslateGrad.pbtxt
index 647f70653585b0..8eaa03c3933199 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ScaleAndTranslateGrad.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ScaleAndTranslateGrad.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ScaleAndTranslateGrad"
   input_arg {
     name: "grads"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ScanDataset.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ScanDataset.pbtxt
index 7563a34c8c7df4..25de8c51a7a388 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ScanDataset.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ScanDataset.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ScanDataset"
   input_arg {
     name: "input_dataset"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ScatterAdd.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ScatterAdd.pbtxt
index 70f2fc7f5ae9ae..0e47d683e05717 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ScatterAdd.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ScatterAdd.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ScatterAdd"
   input_arg {
     name: "ref"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ScatterDiv.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ScatterDiv.pbtxt
index c34776538354d3..2e75f22b0ed101 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ScatterDiv.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ScatterDiv.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ScatterDiv"
   input_arg {
     name: "ref"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ScatterMax.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ScatterMax.pbtxt
index 7ad8c98ea588fe..fe176e143b1874 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ScatterMax.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ScatterMax.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ScatterMax"
   input_arg {
     name: "ref"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ScatterMin.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ScatterMin.pbtxt
index e9dc8087a99349..7099d89f366c6a 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ScatterMin.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ScatterMin.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ScatterMin"
   input_arg {
     name: "ref"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ScatterMul.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ScatterMul.pbtxt
index aa6f863e0bca77..149540b04aecc1 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ScatterMul.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ScatterMul.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ScatterMul"
   input_arg {
     name: "ref"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ScatterNd.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ScatterNd.pbtxt
index 06877c844cbd08..75749c807b7368 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ScatterNd.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ScatterNd.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ScatterNd"
   input_arg {
     name: "indices"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ScatterNdAdd.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ScatterNdAdd.pbtxt
index 409f7d35e5f1e8..d093276f490609 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ScatterNdAdd.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ScatterNdAdd.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ScatterNdAdd"
   input_arg {
     name: "ref"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ScatterNdMax.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ScatterNdMax.pbtxt
index 1425d91fc0563a..2af2bf50069c0c 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ScatterNdMax.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ScatterNdMax.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ScatterNdMax"
   input_arg {
     name: "ref"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ScatterNdMin.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ScatterNdMin.pbtxt
index 996fd4036e1cb5..c2a91a0ed313c8 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ScatterNdMin.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ScatterNdMin.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ScatterNdMin"
   input_arg {
     name: "ref"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ScatterNdNonAliasingAdd.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ScatterNdNonAliasingAdd.pbtxt
index 5cf7d9f91f111b..eb8f48444b62ed 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ScatterNdNonAliasingAdd.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ScatterNdNonAliasingAdd.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ScatterNdNonAliasingAdd"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ScatterNdSub.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ScatterNdSub.pbtxt
index 6e85f0669efa5e..c9c73604d8ad20 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ScatterNdSub.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ScatterNdSub.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ScatterNdSub"
   input_arg {
     name: "ref"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ScatterNdUpdate.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ScatterNdUpdate.pbtxt
index c3f7a023e18568..73def71d094156 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ScatterNdUpdate.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ScatterNdUpdate.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ScatterNdUpdate"
   input_arg {
     name: "ref"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ScatterSub.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ScatterSub.pbtxt
index bf168e222834a8..9665d98a94e2d7 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ScatterSub.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ScatterSub.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ScatterSub"
   input_arg {
     name: "ref"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ScatterUpdate.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ScatterUpdate.pbtxt
index fc91f285899ac2..2f292734f17fb6 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ScatterUpdate.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ScatterUpdate.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ScatterUpdate"
   input_arg {
     name: "ref"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/SdcaFprint.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/SdcaFprint.pbtxt
index 53ee0176bcf390..979c0016b3400d 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/SdcaFprint.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/SdcaFprint.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "SdcaFprint"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/SdcaOptimizer.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/SdcaOptimizer.pbtxt
index d13dd131607794..3746f9504ac61c 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/SdcaOptimizer.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/SdcaOptimizer.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "SdcaOptimizer"
   input_arg {
     name: "sparse_example_indices"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/SdcaOptimizerV2.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/SdcaOptimizerV2.pbtxt
index 2052c0c4a26703..cb16c8f55a1677 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/SdcaOptimizerV2.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/SdcaOptimizerV2.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "SdcaOptimizerV2"
   input_arg {
     name: "sparse_example_indices"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/SdcaShrinkL1.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/SdcaShrinkL1.pbtxt
index 8bc28bda66d735..23d9fdd793a2e9 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/SdcaShrinkL1.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/SdcaShrinkL1.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "SdcaShrinkL1"
   input_arg {
     name: "weights"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/SegmentMax.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/SegmentMax.pbtxt
index c137b5ce06d672..a1d5968f9dd7bc 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/SegmentMax.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/SegmentMax.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "SegmentMax"
   input_arg {
     name: "data"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/SegmentMaxV2.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/SegmentMaxV2.pbtxt
index a2186b6d58c8cb..f61b9101812e8a 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/SegmentMaxV2.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/SegmentMaxV2.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "SegmentMaxV2"
   input_arg {
     name: "data"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/SegmentMean.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/SegmentMean.pbtxt
index a5a4a423220e35..b3c5ead77d8510 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/SegmentMean.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/SegmentMean.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "SegmentMean"
   input_arg {
     name: "data"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/SegmentMin.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/SegmentMin.pbtxt
index feae22be352cfd..bf87e8294e9421 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/SegmentMin.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/SegmentMin.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "SegmentMin"
   input_arg {
     name: "data"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/SegmentMinV2.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/SegmentMinV2.pbtxt
index 16fd710dac8b47..fe2e5396ff0dde 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/SegmentMinV2.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/SegmentMinV2.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "SegmentMinV2"
   input_arg {
     name: "data"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/SegmentProd.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/SegmentProd.pbtxt
index 4726236438d1a9..8fbc39235bf0f7 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/SegmentProd.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/SegmentProd.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "SegmentProd"
   input_arg {
     name: "data"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/SegmentProdV2.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/SegmentProdV2.pbtxt
index a66300ef3e4230..9868f42e95da01 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/SegmentProdV2.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/SegmentProdV2.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "SegmentProdV2"
   input_arg {
     name: "data"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/SegmentSum.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/SegmentSum.pbtxt
index 9f033b0db192b4..69289d5b7444b8 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/SegmentSum.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/SegmentSum.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "SegmentSum"
   input_arg {
     name: "data"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/SegmentSumV2.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/SegmentSumV2.pbtxt
index f0bc5dce91d3b8..e11efffc4be916 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/SegmentSumV2.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/SegmentSumV2.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "SegmentSumV2"
   input_arg {
     name: "data"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/Select.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/Select.pbtxt
index c26b378e476155..38d00af197a978 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/Select.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/Select.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "Select"
   input_arg {
     name: "condition"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/SelectV2.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/SelectV2.pbtxt
index 0536e625bc1f15..a7c59f0d2a1778 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/SelectV2.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/SelectV2.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "SelectV2"
   input_arg {
     name: "condition"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/SelfAdjointEig.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/SelfAdjointEig.pbtxt
index c9bcf0815924a1..3657cc1bdc365d 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/SelfAdjointEig.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/SelfAdjointEig.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "SelfAdjointEig"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/SelfAdjointEigV2.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/SelfAdjointEigV2.pbtxt
index 2e4b645639d6a6..8fbbfc961d7788 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/SelfAdjointEigV2.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/SelfAdjointEigV2.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "SelfAdjointEigV2"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/Selu.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/Selu.pbtxt
index d24219cff310f3..2acf579a5ca9a7 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/Selu.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/Selu.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "Selu"
   input_arg {
     name: "features"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/SeluGrad.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/SeluGrad.pbtxt
index d55750d1e75dd1..f96c7cbc158564 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/SeluGrad.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/SeluGrad.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "SeluGrad"
   input_arg {
     name: "gradients"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/Send.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/Send.pbtxt
index c2f8bd93e93c5c..73835060ddbcb2 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/Send.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/Send.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "Send"
   input_arg {
     name: "tensor"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/SendTPUEmbeddingGradients.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/SendTPUEmbeddingGradients.pbtxt
index 0794e1c7e5c448..f6c486fe542bef 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/SendTPUEmbeddingGradients.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/SendTPUEmbeddingGradients.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "SendTPUEmbeddingGradients"
   input_arg {
     name: "inputs"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/SerializeIterator.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/SerializeIterator.pbtxt
index ba098cea894104..262c9def883fca 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/SerializeIterator.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/SerializeIterator.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "SerializeIterator"
   input_arg {
     name: "resource_handle"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/SerializeManySparse.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/SerializeManySparse.pbtxt
index cb0c4c86419eec..9e741634385c16 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/SerializeManySparse.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/SerializeManySparse.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "SerializeManySparse"
   input_arg {
     name: "sparse_indices"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/SerializeSparse.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/SerializeSparse.pbtxt
index b9a18bcf35107c..5040d77fc2f407 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/SerializeSparse.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/SerializeSparse.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "SerializeSparse"
   input_arg {
     name: "sparse_indices"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/SerializeTensor.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/SerializeTensor.pbtxt
index c739dc6a8330f3..4d7b5cf5766745 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/SerializeTensor.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/SerializeTensor.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "SerializeTensor"
   input_arg {
     name: "tensor"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/SetSize.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/SetSize.pbtxt
index 38c8b19d19dee3..185e5e09734e96 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/SetSize.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/SetSize.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "SetSize"
   input_arg {
     name: "set_indices"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/SetStatsAggregatorDataset.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/SetStatsAggregatorDataset.pbtxt
index cf7b57075cfa14..fa2dfd389adb8c 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/SetStatsAggregatorDataset.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/SetStatsAggregatorDataset.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "SetStatsAggregatorDataset"
   input_arg {
     name: "input_dataset"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/Shape.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/Shape.pbtxt
index c679caa24aad13..371bb9ef2fde71 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/Shape.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/Shape.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "Shape"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ShapeN.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ShapeN.pbtxt
index 04be1dd59c613c..15e9f11c0f8d3b 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ShapeN.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ShapeN.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ShapeN"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ShardDataset.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ShardDataset.pbtxt
index 2c08879ca68dd0..c23a4d3d2e3f98 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ShardDataset.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ShardDataset.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ShardDataset"
   input_arg {
     name: "input_dataset"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ShardedFilename.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ShardedFilename.pbtxt
index df6d834f1f8c1e..cf46ffdbd78d54 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ShardedFilename.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ShardedFilename.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ShardedFilename"
   input_arg {
     name: "basename"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ShardedFilespec.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ShardedFilespec.pbtxt
index 7c2e0806b5ec58..7d1badcf09e83b 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ShardedFilespec.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ShardedFilespec.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ShardedFilespec"
   input_arg {
     name: "basename"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ShuffleAndRepeatDataset.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ShuffleAndRepeatDataset.pbtxt
index ca111a93648d73..195d66b2ab8956 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ShuffleAndRepeatDataset.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ShuffleAndRepeatDataset.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ShuffleAndRepeatDataset"
   input_arg {
     name: "input_dataset"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ShuffleAndRepeatDatasetV2.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ShuffleAndRepeatDatasetV2.pbtxt
index 2ac2e838c3db86..1d22404cf064e0 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ShuffleAndRepeatDatasetV2.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ShuffleAndRepeatDatasetV2.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ShuffleAndRepeatDatasetV2"
   input_arg {
     name: "input_dataset"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ShuffleDataset.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ShuffleDataset.pbtxt
index e1b5ff3283e57a..35c0aa70c11696 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ShuffleDataset.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ShuffleDataset.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ShuffleDataset"
   input_arg {
     name: "input_dataset"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ShuffleDatasetV2.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ShuffleDatasetV2.pbtxt
index 5dce75878aa05f..9ec7fa282d6307 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ShuffleDatasetV2.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ShuffleDatasetV2.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ShuffleDatasetV2"
   input_arg {
     name: "input_dataset"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ShuffleDatasetV3.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ShuffleDatasetV3.pbtxt
index ea30455a9f4acd..e037b818d4ffe1 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ShuffleDatasetV3.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ShuffleDatasetV3.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ShuffleDatasetV3"
   input_arg {
     name: "input_dataset"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ShutdownDistributedTPU.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ShutdownDistributedTPU.pbtxt
index 5784b4418c3112..9e60b7f507554b 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ShutdownDistributedTPU.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ShutdownDistributedTPU.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ShutdownDistributedTPU"
   is_stateful: true
 }
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ShutdownTPUSystem.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ShutdownTPUSystem.pbtxt
index df9e4f58f57dcf..ddddeddc63a18e 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ShutdownTPUSystem.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ShutdownTPUSystem.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ShutdownTPUSystem"
   output_arg {
     name: "success"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/Sigmoid.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/Sigmoid.pbtxt
index 16f433a337c378..dee59f6fa02f41 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/Sigmoid.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/Sigmoid.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "Sigmoid"
   input_arg {
     name: "x"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/SigmoidGrad.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/SigmoidGrad.pbtxt
index 20c59f8f9311e6..788c3385098097 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/SigmoidGrad.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/SigmoidGrad.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "SigmoidGrad"
   input_arg {
     name: "x"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/Sign.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/Sign.pbtxt
index ee388e8ae2a4bf..07cb519996650f 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/Sign.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/Sign.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "Sign"
   input_arg {
     name: "x"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/Sin.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/Sin.pbtxt
index 6a3398c8da6e69..f6122e6e30b1f6 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/Sin.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/Sin.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "Sin"
   input_arg {
     name: "x"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/Sinh.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/Sinh.pbtxt
index 0ff80863b11fda..7225234c7edcdf 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/Sinh.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/Sinh.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "Sinh"
   input_arg {
     name: "x"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/Size.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/Size.pbtxt
index 2c5e61f1496822..db039e4254ced1 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/Size.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/Size.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "Size"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/SkipDataset.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/SkipDataset.pbtxt
index e46a083e7dac16..07e0cf257f87ce 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/SkipDataset.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/SkipDataset.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "SkipDataset"
   input_arg {
     name: "input_dataset"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/Skipgram.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/Skipgram.pbtxt
index 3734477985b239..d31bc826301db8 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/Skipgram.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/Skipgram.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "Skipgram"
   output_arg {
     name: "vocab_word"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/SleepDataset.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/SleepDataset.pbtxt
index c9cdd0432da162..0a1d637995e146 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/SleepDataset.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/SleepDataset.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "SleepDataset"
   input_arg {
     name: "input_dataset"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/Slice.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/Slice.pbtxt
index 0d66369807d6ae..ced3fb6e0f0886 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/Slice.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/Slice.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "Slice"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/SlidingWindowDataset.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/SlidingWindowDataset.pbtxt
index d3361f9d45b76e..ab63899bf4fb51 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/SlidingWindowDataset.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/SlidingWindowDataset.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "SlidingWindowDataset"
   input_arg {
     name: "input_dataset"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/Snapshot.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/Snapshot.pbtxt
index c1fb2d69ed9253..aea213f7c50b1f 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/Snapshot.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/Snapshot.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "Snapshot"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/SnapshotChunkDataset.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/SnapshotChunkDataset.pbtxt
index 339bcd9980923d..e20fbcefee50e0 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/SnapshotChunkDataset.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/SnapshotChunkDataset.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "SnapshotChunkDataset"
   input_arg {
     name: "chunk_file"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/SnapshotDataset.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/SnapshotDataset.pbtxt
index b535c43a80371c..6d9002761ae02f 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/SnapshotDataset.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/SnapshotDataset.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "SnapshotDataset"
   input_arg {
     name: "input_dataset"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/SnapshotDatasetReader.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/SnapshotDatasetReader.pbtxt
index 7a9354125604db..d59d8edf2fc492 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/SnapshotDatasetReader.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/SnapshotDatasetReader.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "SnapshotDatasetReader"
   input_arg {
     name: "shard_dir"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/SnapshotDatasetV2.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/SnapshotDatasetV2.pbtxt
index e8636e3537e451..c9e244ed0e9099 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/SnapshotDatasetV2.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/SnapshotDatasetV2.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "SnapshotDatasetV2"
   input_arg {
     name: "input_dataset"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/SnapshotNestedDatasetReader.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/SnapshotNestedDatasetReader.pbtxt
index 60728b657a6dd4..078460bcb23930 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/SnapshotNestedDatasetReader.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/SnapshotNestedDatasetReader.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "SnapshotNestedDatasetReader"
   input_arg {
     name: "inputs"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/SobolSample.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/SobolSample.pbtxt
index 16377290c6194d..4fe7c45282a15f 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/SobolSample.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/SobolSample.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "SobolSample"
   input_arg {
     name: "dim"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/Softmax.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/Softmax.pbtxt
index 886e1dcb7168c8..03f499777eef54 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/Softmax.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/Softmax.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "Softmax"
   input_arg {
     name: "logits"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/SoftmaxCrossEntropyWithLogits.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/SoftmaxCrossEntropyWithLogits.pbtxt
index 4b258297d20a27..8ac8052e30ffac 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/SoftmaxCrossEntropyWithLogits.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/SoftmaxCrossEntropyWithLogits.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "SoftmaxCrossEntropyWithLogits"
   input_arg {
     name: "features"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/Softplus.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/Softplus.pbtxt
index 66a82cdfd19bb1..3757e8d75039a8 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/Softplus.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/Softplus.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "Softplus"
   input_arg {
     name: "features"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/SoftplusGrad.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/SoftplusGrad.pbtxt
index 20ed580575492b..331b1abbf371f0 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/SoftplusGrad.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/SoftplusGrad.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "SoftplusGrad"
   input_arg {
     name: "gradients"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/Softsign.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/Softsign.pbtxt
index 99ac45e08ba465..c83bc9929ab275 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/Softsign.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/Softsign.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "Softsign"
   input_arg {
     name: "features"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/SoftsignGrad.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/SoftsignGrad.pbtxt
index da12fc0333ef5a..5411f9b5187758 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/SoftsignGrad.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/SoftsignGrad.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "SoftsignGrad"
   input_arg {
     name: "gradients"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/SortListOfSparseCoreCooTensors.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/SortListOfSparseCoreCooTensors.pbtxt
index 209998f61f0bd8..b86af96551b108 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/SortListOfSparseCoreCooTensors.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/SortListOfSparseCoreCooTensors.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "SortListOfSparseCoreCooTensors"
   input_arg {
     name: "row_ids_list"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/SpaceToBatch.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/SpaceToBatch.pbtxt
index 6a3a33227e6b4f..155e1b3a985e44 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/SpaceToBatch.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/SpaceToBatch.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "SpaceToBatch"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/SpaceToBatchND.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/SpaceToBatchND.pbtxt
index 3b0379dbec0ae4..c38026e5cde09c 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/SpaceToBatchND.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/SpaceToBatchND.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "SpaceToBatchND"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/SpaceToDepth.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/SpaceToDepth.pbtxt
index 3abc14c711a049..c7dd03ea1041fc 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/SpaceToDepth.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/SpaceToDepth.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "SpaceToDepth"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/SparseAccumulatorApplyGradient.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/SparseAccumulatorApplyGradient.pbtxt
index cbd19b930643ff..7973ae8b558156 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/SparseAccumulatorApplyGradient.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/SparseAccumulatorApplyGradient.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "SparseAccumulatorApplyGradient"
   input_arg {
     name: "handle"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/SparseAccumulatorTakeGradient.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/SparseAccumulatorTakeGradient.pbtxt
index 5a35297ec764e3..1aaa40667ff84a 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/SparseAccumulatorTakeGradient.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/SparseAccumulatorTakeGradient.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "SparseAccumulatorTakeGradient"
   input_arg {
     name: "handle"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/SparseAdd.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/SparseAdd.pbtxt
index 84b00a504d9f9d..06122b12980987 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/SparseAdd.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/SparseAdd.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "SparseAdd"
   input_arg {
     name: "a_indices"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/SparseAddGrad.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/SparseAddGrad.pbtxt
index 96192d0094101c..5c5e9aa0118f89 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/SparseAddGrad.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/SparseAddGrad.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "SparseAddGrad"
   input_arg {
     name: "backprop_val_grad"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/SparseApplyAdadelta.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/SparseApplyAdadelta.pbtxt
index abff5aab28cbc4..5cd7caf5d819d7 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/SparseApplyAdadelta.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/SparseApplyAdadelta.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "SparseApplyAdadelta"
   input_arg {
     name: "var"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/SparseApplyAdagrad.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/SparseApplyAdagrad.pbtxt
index 248c28c0b8e9e7..3e2e873c9ab59d 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/SparseApplyAdagrad.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/SparseApplyAdagrad.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "SparseApplyAdagrad"
   input_arg {
     name: "var"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/SparseApplyAdagradDA.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/SparseApplyAdagradDA.pbtxt
index 8dcc79dc49aa9c..e47c75ddc6cfa1 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/SparseApplyAdagradDA.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/SparseApplyAdagradDA.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "SparseApplyAdagradDA"
   input_arg {
     name: "var"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/SparseApplyAdagradV2.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/SparseApplyAdagradV2.pbtxt
index 00a3ca86f43359..2b5360f68b6ae1 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/SparseApplyAdagradV2.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/SparseApplyAdagradV2.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "SparseApplyAdagradV2"
   input_arg {
     name: "var"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/SparseApplyCenteredRMSProp.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/SparseApplyCenteredRMSProp.pbtxt
index a5ae87d42c445a..8f958ce67f413a 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/SparseApplyCenteredRMSProp.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/SparseApplyCenteredRMSProp.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "SparseApplyCenteredRMSProp"
   input_arg {
     name: "var"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/SparseApplyFtrl.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/SparseApplyFtrl.pbtxt
index d43fe26ffc1a11..17f289a6bbba7f 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/SparseApplyFtrl.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/SparseApplyFtrl.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "SparseApplyFtrl"
   input_arg {
     name: "var"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/SparseApplyFtrlV2.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/SparseApplyFtrlV2.pbtxt
index 4ced4fe22d596a..1d7b1cab4169df 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/SparseApplyFtrlV2.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/SparseApplyFtrlV2.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "SparseApplyFtrlV2"
   input_arg {
     name: "var"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/SparseApplyMomentum.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/SparseApplyMomentum.pbtxt
index 17e60ae80df5eb..ffabce2fb7af25 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/SparseApplyMomentum.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/SparseApplyMomentum.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "SparseApplyMomentum"
   input_arg {
     name: "var"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/SparseApplyProximalAdagrad.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/SparseApplyProximalAdagrad.pbtxt
index 5fb249592fba4f..c2a7440e3dcacb 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/SparseApplyProximalAdagrad.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/SparseApplyProximalAdagrad.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "SparseApplyProximalAdagrad"
   input_arg {
     name: "var"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/SparseApplyProximalGradientDescent.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/SparseApplyProximalGradientDescent.pbtxt
index f04e6553369cc9..525119034a7653 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/SparseApplyProximalGradientDescent.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/SparseApplyProximalGradientDescent.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "SparseApplyProximalGradientDescent"
   input_arg {
     name: "var"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/SparseApplyRMSProp.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/SparseApplyRMSProp.pbtxt
index aa27af72a44866..f5f7725c0e95e8 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/SparseApplyRMSProp.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/SparseApplyRMSProp.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "SparseApplyRMSProp"
   input_arg {
     name: "var"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/SparseBincount.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/SparseBincount.pbtxt
index 9739aa7f2588de..9bbc5132845f1f 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/SparseBincount.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/SparseBincount.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "SparseBincount"
   input_arg {
     name: "indices"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/SparseConcat.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/SparseConcat.pbtxt
index 640957dd1750a5..ac291f4acbacce 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/SparseConcat.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/SparseConcat.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "SparseConcat"
   input_arg {
     name: "indices"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/SparseConditionalAccumulator.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/SparseConditionalAccumulator.pbtxt
index 26bc9fa77b61d6..59def3f130ef21 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/SparseConditionalAccumulator.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/SparseConditionalAccumulator.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "SparseConditionalAccumulator"
   output_arg {
     name: "handle"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/SparseCountSparseOutput.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/SparseCountSparseOutput.pbtxt
index d80e465d0205ba..ed79733f97fcdf 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/SparseCountSparseOutput.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/SparseCountSparseOutput.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "SparseCountSparseOutput"
   input_arg {
     name: "indices"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/SparseCross.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/SparseCross.pbtxt
index 81bd4101693f0e..f25372f5808567 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/SparseCross.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/SparseCross.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "SparseCross"
   input_arg {
     name: "indices"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/SparseCrossHashed.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/SparseCrossHashed.pbtxt
index b0be6ee5dbcc64..73002a92f24850 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/SparseCrossHashed.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/SparseCrossHashed.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "SparseCrossHashed"
   input_arg {
     name: "indices"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/SparseCrossV2.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/SparseCrossV2.pbtxt
index e2a3e7053512c3..206542e4713902 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/SparseCrossV2.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/SparseCrossV2.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "SparseCrossV2"
   input_arg {
     name: "indices"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/SparseDenseCwiseAdd.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/SparseDenseCwiseAdd.pbtxt
index 474457187b36bf..ca71405a4c9f61 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/SparseDenseCwiseAdd.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/SparseDenseCwiseAdd.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "SparseDenseCwiseAdd"
   input_arg {
     name: "sp_indices"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/SparseDenseCwiseDiv.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/SparseDenseCwiseDiv.pbtxt
index d91c4c89585002..fe4b629ad1cf5b 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/SparseDenseCwiseDiv.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/SparseDenseCwiseDiv.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "SparseDenseCwiseDiv"
   input_arg {
     name: "sp_indices"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/SparseDenseCwiseMul.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/SparseDenseCwiseMul.pbtxt
index f6fd9e956c0884..80f7245d66d1bb 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/SparseDenseCwiseMul.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/SparseDenseCwiseMul.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "SparseDenseCwiseMul"
   input_arg {
     name: "sp_indices"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/SparseFillEmptyRows.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/SparseFillEmptyRows.pbtxt
index a93278ac57811e..d99257aa7103af 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/SparseFillEmptyRows.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/SparseFillEmptyRows.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "SparseFillEmptyRows"
   input_arg {
     name: "indices"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/SparseFillEmptyRowsGrad.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/SparseFillEmptyRowsGrad.pbtxt
index af97bf3745301b..87f1c5c4e2d1e3 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/SparseFillEmptyRowsGrad.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/SparseFillEmptyRowsGrad.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "SparseFillEmptyRowsGrad"
   input_arg {
     name: "reverse_index_map"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/SparseMatMul.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/SparseMatMul.pbtxt
index e3a0ab25bb5f8d..d1eaa6a5edcb98 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/SparseMatMul.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/SparseMatMul.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "SparseMatMul"
   input_arg {
     name: "a"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/SparseMatrixAdd.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/SparseMatrixAdd.pbtxt
index 1a87d1beedee98..3a9efffdae686c 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/SparseMatrixAdd.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/SparseMatrixAdd.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "SparseMatrixAdd"
   input_arg {
     name: "a"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/SparseMatrixMatMul.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/SparseMatrixMatMul.pbtxt
index db0ded9e1f7cae..a38613747717c9 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/SparseMatrixMatMul.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/SparseMatrixMatMul.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "SparseMatrixMatMul"
   input_arg {
     name: "a"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/SparseMatrixMul.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/SparseMatrixMul.pbtxt
index 1ed54fe50f9f58..649992a0298912 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/SparseMatrixMul.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/SparseMatrixMul.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "SparseMatrixMul"
   input_arg {
     name: "a"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/SparseMatrixNNZ.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/SparseMatrixNNZ.pbtxt
index 2073ae629c54cd..40363327c68c83 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/SparseMatrixNNZ.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/SparseMatrixNNZ.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "SparseMatrixNNZ"
   input_arg {
     name: "sparse_matrix"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/SparseMatrixOrderingAMD.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/SparseMatrixOrderingAMD.pbtxt
index cfa3dda19d4b26..b851bde928900d 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/SparseMatrixOrderingAMD.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/SparseMatrixOrderingAMD.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "SparseMatrixOrderingAMD"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/SparseMatrixSoftmax.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/SparseMatrixSoftmax.pbtxt
index 63579485eae114..d994082a1b3e9e 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/SparseMatrixSoftmax.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/SparseMatrixSoftmax.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "SparseMatrixSoftmax"
   input_arg {
     name: "logits"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/SparseMatrixSoftmaxGrad.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/SparseMatrixSoftmaxGrad.pbtxt
index 133ffbbfc1047c..3bb68d7797d134 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/SparseMatrixSoftmaxGrad.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/SparseMatrixSoftmaxGrad.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "SparseMatrixSoftmaxGrad"
   input_arg {
     name: "softmax"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/SparseMatrixSparseCholesky.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/SparseMatrixSparseCholesky.pbtxt
index f14d9834fa04ef..1f74136d3a480e 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/SparseMatrixSparseCholesky.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/SparseMatrixSparseCholesky.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "SparseMatrixSparseCholesky"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/SparseMatrixSparseMatMul.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/SparseMatrixSparseMatMul.pbtxt
index 51afd8da1a2ec7..3726d6bd8c0a4d 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/SparseMatrixSparseMatMul.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/SparseMatrixSparseMatMul.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "SparseMatrixSparseMatMul"
   input_arg {
     name: "a"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/SparseMatrixTranspose.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/SparseMatrixTranspose.pbtxt
index ea071da4e77177..be6bcd6dc1e496 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/SparseMatrixTranspose.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/SparseMatrixTranspose.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "SparseMatrixTranspose"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/SparseMatrixZeros.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/SparseMatrixZeros.pbtxt
index 87b09ffd306f2c..6d743e24a0ad4c 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/SparseMatrixZeros.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/SparseMatrixZeros.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "SparseMatrixZeros"
   input_arg {
     name: "dense_shape"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/SparseReduceMax.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/SparseReduceMax.pbtxt
index 86a4d1e0046906..4df1254af7460d 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/SparseReduceMax.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/SparseReduceMax.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "SparseReduceMax"
   input_arg {
     name: "input_indices"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/SparseReduceMaxSparse.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/SparseReduceMaxSparse.pbtxt
index c44461c9d090be..81896440ab75a9 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/SparseReduceMaxSparse.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/SparseReduceMaxSparse.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "SparseReduceMaxSparse"
   input_arg {
     name: "input_indices"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/SparseReduceSum.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/SparseReduceSum.pbtxt
index d0e5258dad8801..9f5a2d2e2be133 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/SparseReduceSum.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/SparseReduceSum.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "SparseReduceSum"
   input_arg {
     name: "input_indices"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/SparseReduceSumSparse.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/SparseReduceSumSparse.pbtxt
index 12a18fbda1c045..b554e7ca6a00f6 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/SparseReduceSumSparse.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/SparseReduceSumSparse.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "SparseReduceSumSparse"
   input_arg {
     name: "input_indices"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/SparseReorder.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/SparseReorder.pbtxt
index 9a9bd24e406c00..5c5ad9078385e6 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/SparseReorder.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/SparseReorder.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "SparseReorder"
   input_arg {
     name: "input_indices"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/SparseReshape.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/SparseReshape.pbtxt
index ab9f157b09f23f..934b5010a1e415 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/SparseReshape.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/SparseReshape.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "SparseReshape"
   input_arg {
     name: "input_indices"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/SparseSegmentMean.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/SparseSegmentMean.pbtxt
index 6f56f3f5364548..0891aa6b0175ec 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/SparseSegmentMean.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/SparseSegmentMean.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "SparseSegmentMean"
   input_arg {
     name: "data"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/SparseSegmentMeanGrad.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/SparseSegmentMeanGrad.pbtxt
index a75fd742bf4e91..54e44b3e62d64b 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/SparseSegmentMeanGrad.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/SparseSegmentMeanGrad.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "SparseSegmentMeanGrad"
   input_arg {
     name: "grad"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/SparseSegmentMeanGradV2.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/SparseSegmentMeanGradV2.pbtxt
index b02a4838190c33..e1282b858a90ec 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/SparseSegmentMeanGradV2.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/SparseSegmentMeanGradV2.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "SparseSegmentMeanGradV2"
   input_arg {
     name: "grad"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/SparseSegmentMeanWithNumSegments.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/SparseSegmentMeanWithNumSegments.pbtxt
index 930da74b96b061..a27a2efc560465 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/SparseSegmentMeanWithNumSegments.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/SparseSegmentMeanWithNumSegments.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "SparseSegmentMeanWithNumSegments"
   input_arg {
     name: "data"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/SparseSegmentSqrtN.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/SparseSegmentSqrtN.pbtxt
index feec752c94d180..b7a209c48d9ee6 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/SparseSegmentSqrtN.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/SparseSegmentSqrtN.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "SparseSegmentSqrtN"
   input_arg {
     name: "data"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/SparseSegmentSqrtNGrad.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/SparseSegmentSqrtNGrad.pbtxt
index 18397fbb9cda45..d9a3c930a122de 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/SparseSegmentSqrtNGrad.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/SparseSegmentSqrtNGrad.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "SparseSegmentSqrtNGrad"
   input_arg {
     name: "grad"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/SparseSegmentSqrtNGradV2.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/SparseSegmentSqrtNGradV2.pbtxt
index 1d8a3b4c921189..9436be8beb8e17 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/SparseSegmentSqrtNGradV2.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/SparseSegmentSqrtNGradV2.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "SparseSegmentSqrtNGradV2"
   input_arg {
     name: "grad"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/SparseSegmentSqrtNWithNumSegments.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/SparseSegmentSqrtNWithNumSegments.pbtxt
index 5aa63ca35fcace..b67358819acdc6 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/SparseSegmentSqrtNWithNumSegments.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/SparseSegmentSqrtNWithNumSegments.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "SparseSegmentSqrtNWithNumSegments"
   input_arg {
     name: "data"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/SparseSegmentSum.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/SparseSegmentSum.pbtxt
index db9042011a6f8d..5962058393cf14 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/SparseSegmentSum.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/SparseSegmentSum.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "SparseSegmentSum"
   input_arg {
     name: "data"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/SparseSegmentSumGrad.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/SparseSegmentSumGrad.pbtxt
index e441e7f0a60f69..ce9742f898ef6b 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/SparseSegmentSumGrad.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/SparseSegmentSumGrad.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "SparseSegmentSumGrad"
   input_arg {
     name: "grad"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/SparseSegmentSumGradV2.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/SparseSegmentSumGradV2.pbtxt
index 4643278f47fad1..3baf45183e9464 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/SparseSegmentSumGradV2.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/SparseSegmentSumGradV2.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "SparseSegmentSumGradV2"
   input_arg {
     name: "grad"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/SparseSegmentSumWithNumSegments.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/SparseSegmentSumWithNumSegments.pbtxt
index 55674b28bbae69..4b8183aeeeabcd 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/SparseSegmentSumWithNumSegments.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/SparseSegmentSumWithNumSegments.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "SparseSegmentSumWithNumSegments"
   input_arg {
     name: "data"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/SparseSlice.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/SparseSlice.pbtxt
index 60d914313b3da7..a6434cbfa71d2c 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/SparseSlice.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/SparseSlice.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "SparseSlice"
   input_arg {
     name: "indices"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/SparseSliceGrad.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/SparseSliceGrad.pbtxt
index d3d6693044a564..77134f37194b5a 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/SparseSliceGrad.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/SparseSliceGrad.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "SparseSliceGrad"
   input_arg {
     name: "backprop_val_grad"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/SparseSoftmax.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/SparseSoftmax.pbtxt
index 003ea791ed61f0..5e447367f3be6d 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/SparseSoftmax.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/SparseSoftmax.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "SparseSoftmax"
   input_arg {
     name: "sp_indices"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/SparseSoftmaxCrossEntropyWithLogits.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/SparseSoftmaxCrossEntropyWithLogits.pbtxt
index 884008f4ea9937..57d8f4c4662535 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/SparseSoftmaxCrossEntropyWithLogits.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/SparseSoftmaxCrossEntropyWithLogits.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "SparseSoftmaxCrossEntropyWithLogits"
   input_arg {
     name: "features"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/SparseSparseMaximum.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/SparseSparseMaximum.pbtxt
index c1bb9cb237fb35..bdd017c2252867 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/SparseSparseMaximum.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/SparseSparseMaximum.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "SparseSparseMaximum"
   input_arg {
     name: "a_indices"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/SparseSparseMinimum.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/SparseSparseMinimum.pbtxt
index 7d1569221cce76..b600e95b3ae3bb 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/SparseSparseMinimum.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/SparseSparseMinimum.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "SparseSparseMinimum"
   input_arg {
     name: "a_indices"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/SparseSplit.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/SparseSplit.pbtxt
index af8deac1a17908..997b2b21abdd44 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/SparseSplit.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/SparseSplit.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "SparseSplit"
   input_arg {
     name: "split_dim"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/SparseTensorDenseAdd.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/SparseTensorDenseAdd.pbtxt
index c1b647cef05e96..8e95e67610204f 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/SparseTensorDenseAdd.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/SparseTensorDenseAdd.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "SparseTensorDenseAdd"
   input_arg {
     name: "a_indices"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/SparseTensorDenseMatMul.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/SparseTensorDenseMatMul.pbtxt
index 8fa6733bfc1b91..ce66c5306a5c2a 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/SparseTensorDenseMatMul.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/SparseTensorDenseMatMul.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "SparseTensorDenseMatMul"
   input_arg {
     name: "a_indices"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/SparseTensorSliceDataset.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/SparseTensorSliceDataset.pbtxt
index bc7b81876d1c4e..af26fd8c180a3f 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/SparseTensorSliceDataset.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/SparseTensorSliceDataset.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "SparseTensorSliceDataset"
   input_arg {
     name: "indices"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/SparseTensorToCSRSparseMatrix.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/SparseTensorToCSRSparseMatrix.pbtxt
index a1215c2000ae1e..b45376fcd00d3c 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/SparseTensorToCSRSparseMatrix.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/SparseTensorToCSRSparseMatrix.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "SparseTensorToCSRSparseMatrix"
   input_arg {
     name: "indices"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/SparseToDense.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/SparseToDense.pbtxt
index 7e693267c06d2a..351603424e6ccd 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/SparseToDense.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/SparseToDense.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "SparseToDense"
   input_arg {
     name: "sparse_indices"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/SparseToSparseSetOperation.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/SparseToSparseSetOperation.pbtxt
index bac3763b0510a7..a7775a2f24a465 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/SparseToSparseSetOperation.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/SparseToSparseSetOperation.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "SparseToSparseSetOperation"
   input_arg {
     name: "set1_indices"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/Spence.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/Spence.pbtxt
index fefe73924fda23..7032cac3dce437 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/Spence.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/Spence.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "Spence"
   input_arg {
     name: "x"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/Split.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/Split.pbtxt
index 3d69a43f2e71ec..49428f7e5ce590 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/Split.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/Split.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "Split"
   input_arg {
     name: "split_dim"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/SplitDedupData.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/SplitDedupData.pbtxt
index 53d3ccdf1c6a6f..ada9d479888933 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/SplitDedupData.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/SplitDedupData.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "SplitDedupData"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/SplitV.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/SplitV.pbtxt
index 577d8a604fac7b..706d1c3c81c729 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/SplitV.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/SplitV.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "SplitV"
   input_arg {
     name: "value"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/SqlDataset.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/SqlDataset.pbtxt
index fb233f139edaa7..68af0ac17eff32 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/SqlDataset.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/SqlDataset.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "SqlDataset"
   input_arg {
     name: "driver_name"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/Sqrt.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/Sqrt.pbtxt
index 371520ff6ae795..3c566b98b0dbff 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/Sqrt.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/Sqrt.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "Sqrt"
   input_arg {
     name: "x"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/SqrtGrad.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/SqrtGrad.pbtxt
index 055baa35e98f93..d738e2023ffc24 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/SqrtGrad.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/SqrtGrad.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "SqrtGrad"
   input_arg {
     name: "x"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/Square.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/Square.pbtxt
index 52c0e31486d601..7501f29f1d95e4 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/Square.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/Square.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "Square"
   input_arg {
     name: "x"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/SquaredDifference.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/SquaredDifference.pbtxt
index a680e0087f531d..29ea33c95e2d2a 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/SquaredDifference.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/SquaredDifference.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "SquaredDifference"
   input_arg {
     name: "x"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/Squeeze.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/Squeeze.pbtxt
index d3733f99e07910..54335545f4807f 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/Squeeze.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/Squeeze.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "Squeeze"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/Stack.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/Stack.pbtxt
index e4398a4783de8f..e8e459cfe2c08d 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/Stack.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/Stack.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "Stack"
   output_arg {
     name: "handle"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/StackClose.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/StackClose.pbtxt
index dd2b8efd9c95ab..8c916ab52a2c69 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/StackClose.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/StackClose.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "StackClose"
   input_arg {
     name: "handle"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/StackCloseV2.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/StackCloseV2.pbtxt
index 66ddab1d90bed8..18c5934b0d7961 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/StackCloseV2.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/StackCloseV2.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "StackCloseV2"
   input_arg {
     name: "handle"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/StackPop.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/StackPop.pbtxt
index 53965ff252d508..80e3ef79d09c61 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/StackPop.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/StackPop.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "StackPop"
   input_arg {
     name: "handle"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/StackPopV2.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/StackPopV2.pbtxt
index 77300ed42e46f9..438d52b8ea5625 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/StackPopV2.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/StackPopV2.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "StackPopV2"
   input_arg {
     name: "handle"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/StackPush.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/StackPush.pbtxt
index f2052fe41c1ab4..44fae0ce455f0c 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/StackPush.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/StackPush.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "StackPush"
   input_arg {
     name: "handle"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/StackPushV2.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/StackPushV2.pbtxt
index 6b3f8b5655b2a3..7149b4fda435c6 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/StackPushV2.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/StackPushV2.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "StackPushV2"
   input_arg {
     name: "handle"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/StackV2.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/StackV2.pbtxt
index 5d934a7b8bb00b..606361dd26fdb1 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/StackV2.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/StackV2.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "StackV2"
   input_arg {
     name: "max_size"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/Stage.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/Stage.pbtxt
index d7b0fbc1b9a4b7..8a64d696118f7f 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/Stage.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/Stage.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "Stage"
   input_arg {
     name: "values"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/StageClear.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/StageClear.pbtxt
index 7ce3219c586d58..1f43cdb901967d 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/StageClear.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/StageClear.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "StageClear"
   attr {
     name: "capacity"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/StagePeek.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/StagePeek.pbtxt
index e9e6c68ed56344..a7397c488167fe 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/StagePeek.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/StagePeek.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "StagePeek"
   input_arg {
     name: "index"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/StageSize.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/StageSize.pbtxt
index b9bb218b6909aa..6f22fd3d032706 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/StageSize.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/StageSize.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "StageSize"
   output_arg {
     name: "size"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/StatefulPartitionedCall.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/StatefulPartitionedCall.pbtxt
index a069806185b334..7d411c155ab2b9 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/StatefulPartitionedCall.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/StatefulPartitionedCall.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "StatefulPartitionedCall"
   input_arg {
     name: "args"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/StatefulRandomBinomial.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/StatefulRandomBinomial.pbtxt
index 798854feef481b..97eb7d4e8f0fa0 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/StatefulRandomBinomial.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/StatefulRandomBinomial.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "StatefulRandomBinomial"
   input_arg {
     name: "resource"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/StatefulStandardNormal.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/StatefulStandardNormal.pbtxt
index 4738b843a0b92a..44ef92c5a8f7ed 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/StatefulStandardNormal.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/StatefulStandardNormal.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "StatefulStandardNormal"
   input_arg {
     name: "resource"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/StatefulStandardNormalV2.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/StatefulStandardNormalV2.pbtxt
index 76e31871a6bbae..1b99b2320b0264 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/StatefulStandardNormalV2.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/StatefulStandardNormalV2.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "StatefulStandardNormalV2"
   input_arg {
     name: "resource"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/StatefulTruncatedNormal.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/StatefulTruncatedNormal.pbtxt
index 8ef450390d5367..e74de4f0fceb55 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/StatefulTruncatedNormal.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/StatefulTruncatedNormal.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "StatefulTruncatedNormal"
   input_arg {
     name: "resource"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/StatefulUniform.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/StatefulUniform.pbtxt
index fef8635b75cb94..fd2b87c6e45988 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/StatefulUniform.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/StatefulUniform.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "StatefulUniform"
   input_arg {
     name: "resource"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/StatefulUniformFullInt.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/StatefulUniformFullInt.pbtxt
index 0536cdf62ce4cc..35ab70e0f3a1e2 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/StatefulUniformFullInt.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/StatefulUniformFullInt.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "StatefulUniformFullInt"
   input_arg {
     name: "resource"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/StatefulUniformInt.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/StatefulUniformInt.pbtxt
index cd71c24e2762c1..06f62faaace7b6 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/StatefulUniformInt.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/StatefulUniformInt.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "StatefulUniformInt"
   input_arg {
     name: "resource"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/StatelessCase.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/StatelessCase.pbtxt
index 367c9ddfdb9602..174c00e5c8ab4a 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/StatelessCase.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/StatelessCase.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "StatelessCase"
   input_arg {
     name: "branch_index"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/StatelessIf.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/StatelessIf.pbtxt
index a38480703ef14a..6eda6df052a58b 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/StatelessIf.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/StatelessIf.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "StatelessIf"
   input_arg {
     name: "cond"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/StatelessMultinomial.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/StatelessMultinomial.pbtxt
index 3a06bef6ef7747..16dac7dfd152bb 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/StatelessMultinomial.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/StatelessMultinomial.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "StatelessMultinomial"
   input_arg {
     name: "logits"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/StatelessParameterizedTruncatedNormal.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/StatelessParameterizedTruncatedNormal.pbtxt
index 86509705c4f188..598125677b114a 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/StatelessParameterizedTruncatedNormal.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/StatelessParameterizedTruncatedNormal.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "StatelessParameterizedTruncatedNormal"
   input_arg {
     name: "shape"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/StatelessRandomBinomial.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/StatelessRandomBinomial.pbtxt
index da4c351d221a53..1ba7e5c119b147 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/StatelessRandomBinomial.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/StatelessRandomBinomial.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "StatelessRandomBinomial"
   input_arg {
     name: "shape"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/StatelessRandomGammaV2.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/StatelessRandomGammaV2.pbtxt
index 56d573cf2ec6d9..b3988b1f407d96 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/StatelessRandomGammaV2.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/StatelessRandomGammaV2.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "StatelessRandomGammaV2"
   input_arg {
     name: "shape"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/StatelessRandomGammaV3.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/StatelessRandomGammaV3.pbtxt
index d2fbd60387e24d..00d9da80c101e1 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/StatelessRandomGammaV3.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/StatelessRandomGammaV3.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "StatelessRandomGammaV3"
   input_arg {
     name: "shape"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/StatelessRandomGetAlg.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/StatelessRandomGetAlg.pbtxt
index ff50f6fad96a1c..522fce59c88a05 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/StatelessRandomGetAlg.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/StatelessRandomGetAlg.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "StatelessRandomGetAlg"
   output_arg {
     name: "alg"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/StatelessRandomGetKeyCounter.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/StatelessRandomGetKeyCounter.pbtxt
index 7532f9b2f6ffca..0d9c547181bb61 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/StatelessRandomGetKeyCounter.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/StatelessRandomGetKeyCounter.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "StatelessRandomGetKeyCounter"
   input_arg {
     name: "seed"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/StatelessRandomGetKeyCounterAlg.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/StatelessRandomGetKeyCounterAlg.pbtxt
index 6ed78cbec62348..4897ee80bb82d5 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/StatelessRandomGetKeyCounterAlg.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/StatelessRandomGetKeyCounterAlg.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "StatelessRandomGetKeyCounterAlg"
   input_arg {
     name: "seed"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/StatelessRandomNormal.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/StatelessRandomNormal.pbtxt
index 0de87cb2569ce2..804d904c148234 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/StatelessRandomNormal.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/StatelessRandomNormal.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "StatelessRandomNormal"
   input_arg {
     name: "shape"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/StatelessRandomNormalV2.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/StatelessRandomNormalV2.pbtxt
index e5a48f3116056e..dac945afe5354e 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/StatelessRandomNormalV2.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/StatelessRandomNormalV2.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "StatelessRandomNormalV2"
   input_arg {
     name: "shape"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/StatelessRandomPoisson.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/StatelessRandomPoisson.pbtxt
index 575f54b2e22617..525b933d8b005b 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/StatelessRandomPoisson.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/StatelessRandomPoisson.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "StatelessRandomPoisson"
   input_arg {
     name: "shape"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/StatelessRandomUniform.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/StatelessRandomUniform.pbtxt
index d27528775e2460..22a5b25466b90f 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/StatelessRandomUniform.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/StatelessRandomUniform.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "StatelessRandomUniform"
   input_arg {
     name: "shape"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/StatelessRandomUniformFullInt.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/StatelessRandomUniformFullInt.pbtxt
index ede922b448a4de..a9d652634d7cb7 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/StatelessRandomUniformFullInt.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/StatelessRandomUniformFullInt.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "StatelessRandomUniformFullInt"
   input_arg {
     name: "shape"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/StatelessRandomUniformFullIntV2.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/StatelessRandomUniformFullIntV2.pbtxt
index 5a5bd12f262c5d..d4511c5447bda5 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/StatelessRandomUniformFullIntV2.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/StatelessRandomUniformFullIntV2.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "StatelessRandomUniformFullIntV2"
   input_arg {
     name: "shape"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/StatelessRandomUniformInt.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/StatelessRandomUniformInt.pbtxt
index 13014cca798e16..834a6fd5ad3983 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/StatelessRandomUniformInt.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/StatelessRandomUniformInt.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "StatelessRandomUniformInt"
   input_arg {
     name: "shape"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/StatelessRandomUniformIntV2.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/StatelessRandomUniformIntV2.pbtxt
index 96bc39656a4b17..be4ed6072852be 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/StatelessRandomUniformIntV2.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/StatelessRandomUniformIntV2.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "StatelessRandomUniformIntV2"
   input_arg {
     name: "shape"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/StatelessRandomUniformV2.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/StatelessRandomUniformV2.pbtxt
index 9434be69beff1f..f66ee72bd4af30 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/StatelessRandomUniformV2.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/StatelessRandomUniformV2.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "StatelessRandomUniformV2"
   input_arg {
     name: "shape"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/StatelessSampleDistortedBoundingBox.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/StatelessSampleDistortedBoundingBox.pbtxt
index 6c322d848ac079..6858a110cf46f5 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/StatelessSampleDistortedBoundingBox.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/StatelessSampleDistortedBoundingBox.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "StatelessSampleDistortedBoundingBox"
   input_arg {
     name: "image_size"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/StatelessShuffle.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/StatelessShuffle.pbtxt
index 06b6eb02eac58d..eab3565990135e 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/StatelessShuffle.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/StatelessShuffle.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "StatelessShuffle"
   input_arg {
     name: "value"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/StatelessTruncatedNormal.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/StatelessTruncatedNormal.pbtxt
index a301c48afbe058..c8c8d850341f5f 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/StatelessTruncatedNormal.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/StatelessTruncatedNormal.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "StatelessTruncatedNormal"
   input_arg {
     name: "shape"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/StatelessTruncatedNormalV2.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/StatelessTruncatedNormalV2.pbtxt
index f074ea7423b525..23f886f104dd24 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/StatelessTruncatedNormalV2.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/StatelessTruncatedNormalV2.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "StatelessTruncatedNormalV2"
   input_arg {
     name: "shape"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/StatelessWhile.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/StatelessWhile.pbtxt
index dc0b6353373c65..28579edbde5e67 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/StatelessWhile.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/StatelessWhile.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "StatelessWhile"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/StaticRegexFullMatch.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/StaticRegexFullMatch.pbtxt
index d053ad6e1db6f6..be6078c102232d 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/StaticRegexFullMatch.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/StaticRegexFullMatch.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "StaticRegexFullMatch"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/StaticRegexReplace.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/StaticRegexReplace.pbtxt
index e570f0de09ff55..fe3eb69a1a0044 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/StaticRegexReplace.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/StaticRegexReplace.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "StaticRegexReplace"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/StatsAggregatorHandle.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/StatsAggregatorHandle.pbtxt
index 45080da97423ca..2d55e00492fddc 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/StatsAggregatorHandle.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/StatsAggregatorHandle.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "StatsAggregatorHandle"
   output_arg {
     name: "handle"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/StatsAggregatorHandleV2.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/StatsAggregatorHandleV2.pbtxt
index 0fca6c398316bf..7dc361e958d794 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/StatsAggregatorHandleV2.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/StatsAggregatorHandleV2.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "StatsAggregatorHandleV2"
   output_arg {
     name: "handle"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/StatsAggregatorSetSummaryWriter.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/StatsAggregatorSetSummaryWriter.pbtxt
index 55d5207beec293..24730ade1494c6 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/StatsAggregatorSetSummaryWriter.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/StatsAggregatorSetSummaryWriter.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "StatsAggregatorSetSummaryWriter"
   input_arg {
     name: "stats_aggregator"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/StatsAggregatorSummary.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/StatsAggregatorSummary.pbtxt
index 92d1f2ba507923..a0702a11168ff4 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/StatsAggregatorSummary.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/StatsAggregatorSummary.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "StatsAggregatorSummary"
   input_arg {
     name: "iterator"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/StochasticCastToInt.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/StochasticCastToInt.pbtxt
index 66896c5d662940..930525b4364525 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/StochasticCastToInt.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/StochasticCastToInt.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "StochasticCastToInt"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/StopGradient.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/StopGradient.pbtxt
index 06500717b6034a..26f7c677ab8e8a 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/StopGradient.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/StopGradient.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "StopGradient"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/StoreMinibatchStatisticsInFdo.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/StoreMinibatchStatisticsInFdo.pbtxt
index 2250ba0eee369c..22766cb4409917 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/StoreMinibatchStatisticsInFdo.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/StoreMinibatchStatisticsInFdo.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "StoreMinibatchStatisticsInFdo"
   input_arg {
     name: "program_key"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/StridedSlice.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/StridedSlice.pbtxt
index 60186d9449d593..2c60bcb73c757f 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/StridedSlice.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/StridedSlice.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "StridedSlice"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/StridedSliceAssign.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/StridedSliceAssign.pbtxt
index ba3eaed57e270e..8393dc7272c59c 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/StridedSliceAssign.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/StridedSliceAssign.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "StridedSliceAssign"
   input_arg {
     name: "ref"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/StridedSliceGrad.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/StridedSliceGrad.pbtxt
index 92bad979c0f579..14f6a464020a5a 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/StridedSliceGrad.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/StridedSliceGrad.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "StridedSliceGrad"
   input_arg {
     name: "shape"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/StringFormat.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/StringFormat.pbtxt
index 23c48e47c2761d..bea32908608e6f 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/StringFormat.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/StringFormat.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "StringFormat"
   input_arg {
     name: "inputs"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/StringJoin.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/StringJoin.pbtxt
index a45262ee65ed29..790cb7b5b6c8d6 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/StringJoin.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/StringJoin.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "StringJoin"
   input_arg {
     name: "inputs"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/StringLength.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/StringLength.pbtxt
index 169adeecbed33f..5bdf993f907a2a 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/StringLength.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/StringLength.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "StringLength"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/StringLower.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/StringLower.pbtxt
index 5f7dbf3b2bb283..1c886146d0560d 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/StringLower.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/StringLower.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "StringLower"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/StringNGrams.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/StringNGrams.pbtxt
index 4281b96aee5007..025fc052819bf2 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/StringNGrams.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/StringNGrams.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "StringNGrams"
   input_arg {
     name: "data"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/StringSplit.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/StringSplit.pbtxt
index 1832d334e79837..35e8594235e170 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/StringSplit.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/StringSplit.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "StringSplit"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/StringSplitV2.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/StringSplitV2.pbtxt
index 03ba25fd70168b..fbdf8e06f372c1 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/StringSplitV2.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/StringSplitV2.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "StringSplitV2"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/StringStrip.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/StringStrip.pbtxt
index 153bd63f4e2313..3fff999e93789b 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/StringStrip.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/StringStrip.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "StringStrip"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/StringToHashBucket.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/StringToHashBucket.pbtxt
index 6cbdd4c7bf8c28..7147a40a12d5af 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/StringToHashBucket.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/StringToHashBucket.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "StringToHashBucket"
   input_arg {
     name: "string_tensor"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/StringToHashBucketFast.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/StringToHashBucketFast.pbtxt
index a07a00a41eecee..8ef1227faae9b7 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/StringToHashBucketFast.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/StringToHashBucketFast.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "StringToHashBucketFast"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/StringToHashBucketStrong.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/StringToHashBucketStrong.pbtxt
index de121b287b3aa9..2dbd9920711fe4 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/StringToHashBucketStrong.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/StringToHashBucketStrong.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "StringToHashBucketStrong"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/StringToNumber.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/StringToNumber.pbtxt
index 833aaa61533c1f..fb09d67516e7fd 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/StringToNumber.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/StringToNumber.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "StringToNumber"
   input_arg {
     name: "string_tensor"
@@ -48,3 +48,31 @@ op {
     }
   }
 }
+op {
+  name: "StringToNumber"
+  input_arg {
+    name: "string_tensor"
+    type: DT_STRING
+  }
+  output_arg {
+    name: "output"
+    type_attr: "out_type"
+  }
+  attr {
+    name: "out_type"
+    type: "type"
+    default_value {
+      type: DT_FLOAT
+    }
+    allowed_values {
+      list {
+        type: DT_FLOAT
+        type: DT_DOUBLE
+        type: DT_INT32
+        type: DT_INT64
+        type: DT_UINT32
+        type: DT_UINT64
+      }
+    }
+  }
+}
diff --git a/tensorflow/core/ops/compat/ops_history_v2/StringUpper.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/StringUpper.pbtxt
index 69c6133ad0501b..8df4881554cefe 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/StringUpper.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/StringUpper.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "StringUpper"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/Sub.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/Sub.pbtxt
index 44761cb1a3bf95..4d89817f561874 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/Sub.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/Sub.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "Sub"
   input_arg {
     name: "x"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/Substr.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/Substr.pbtxt
index ebfc6279bbe635..a5c1d2c0ae0aa8 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/Substr.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/Substr.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "Substr"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/Sum.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/Sum.pbtxt
index 4a71be5f59cfe9..fb7ce8600c4ecd 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/Sum.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/Sum.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "Sum"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/SummaryWriter.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/SummaryWriter.pbtxt
index 199ded68ff3bbf..a6fd9170f2a121 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/SummaryWriter.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/SummaryWriter.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "SummaryWriter"
   output_arg {
     name: "writer"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/Svd.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/Svd.pbtxt
index 3750a39144739e..48003906cc05ed 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/Svd.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/Svd.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "Svd"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/Switch.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/Switch.pbtxt
index 2ff607548010d7..0856f3459b3c02 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/Switch.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/Switch.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "Switch"
   input_arg {
     name: "data"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/SymbolicGradient.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/SymbolicGradient.pbtxt
index 5193954ea7eebd..aae5457863e176 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/SymbolicGradient.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/SymbolicGradient.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "SymbolicGradient"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/SyncDevice.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/SyncDevice.pbtxt
index cc6bf6b6f0c0b0..e55c5f4ade1d4b 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/SyncDevice.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/SyncDevice.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "SyncDevice"
   is_stateful: true
 }
diff --git a/tensorflow/core/ops/compat/ops_history_v2/TFRecordDataset.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/TFRecordDataset.pbtxt
index b68eea8fbd264c..6caa343ebc05b8 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/TFRecordDataset.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/TFRecordDataset.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "TFRecordDataset"
   input_arg {
     name: "filenames"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/TFRecordDatasetV2.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/TFRecordDatasetV2.pbtxt
index af944b699a0e9d..d3b89247493e91 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/TFRecordDatasetV2.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/TFRecordDatasetV2.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "TFRecordDatasetV2"
   input_arg {
     name: "filenames"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/TFRecordReader.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/TFRecordReader.pbtxt
index 0f223c57ec5ed3..684c21ea45e8f5 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/TFRecordReader.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/TFRecordReader.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "TFRecordReader"
   output_arg {
     name: "reader_handle"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/TFRecordReaderV2.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/TFRecordReaderV2.pbtxt
index 0f9f02ce01d876..bcdb4764d378da 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/TFRecordReaderV2.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/TFRecordReaderV2.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "TFRecordReaderV2"
   output_arg {
     name: "reader_handle"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/TPUAnnotateTensorsWithDynamicShape.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/TPUAnnotateTensorsWithDynamicShape.pbtxt
index 09d484e3194e76..eb1f07856a9ef6 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/TPUAnnotateTensorsWithDynamicShape.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/TPUAnnotateTensorsWithDynamicShape.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "TPUAnnotateTensorsWithDynamicShape"
   input_arg {
     name: "tensors"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/TPUCompilationResult.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/TPUCompilationResult.pbtxt
index bdaf1f9e51a8c9..04a95cc089fd4e 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/TPUCompilationResult.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/TPUCompilationResult.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "TPUCompilationResult"
   output_arg {
     name: "output"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/TPUCompile.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/TPUCompile.pbtxt
index 7328f9214e3be6..be95091c809227 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/TPUCompile.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/TPUCompile.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "TPUCompile"
   input_arg {
     name: "dynamic_shapes"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/TPUCompileSucceededAssert.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/TPUCompileSucceededAssert.pbtxt
index a39d0e72843b10..bc1b3c153f10ec 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/TPUCompileSucceededAssert.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/TPUCompileSucceededAssert.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "TPUCompileSucceededAssert"
   input_arg {
     name: "compilation_status"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/TPUCopyWithDynamicShape.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/TPUCopyWithDynamicShape.pbtxt
index 8b897ff34cca06..1e8386d91a5760 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/TPUCopyWithDynamicShape.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/TPUCopyWithDynamicShape.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "TPUCopyWithDynamicShape"
   input_arg {
     name: "tensors"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/TPUEmbeddingActivations.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/TPUEmbeddingActivations.pbtxt
index 0bd460f3a0bf50..3975077297a6fa 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/TPUEmbeddingActivations.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/TPUEmbeddingActivations.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "TPUEmbeddingActivations"
   input_arg {
     name: "embedding_variable"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/TPUExecute.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/TPUExecute.pbtxt
index a231036be50e8b..97a2c2a1f6673d 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/TPUExecute.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/TPUExecute.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "TPUExecute"
   input_arg {
     name: "args"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/TPUExecuteAndUpdateVariables.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/TPUExecuteAndUpdateVariables.pbtxt
index 2d41c28beff465..5a611f9202d83d 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/TPUExecuteAndUpdateVariables.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/TPUExecuteAndUpdateVariables.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "TPUExecuteAndUpdateVariables"
   input_arg {
     name: "args"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/TPUOrdinalSelector.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/TPUOrdinalSelector.pbtxt
index de910326bd6a79..3fb272504068ce 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/TPUOrdinalSelector.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/TPUOrdinalSelector.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "TPUOrdinalSelector"
   output_arg {
     name: "device_ordinals"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/TPUPartitionedCall.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/TPUPartitionedCall.pbtxt
index cc0652a3790b31..1ec9cb3e43400f 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/TPUPartitionedCall.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/TPUPartitionedCall.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "TPUPartitionedCall"
   input_arg {
     name: "args"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/TPUPartitionedInput.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/TPUPartitionedInput.pbtxt
index b5eca4ed94d216..aab0574d99eb95 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/TPUPartitionedInput.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/TPUPartitionedInput.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "TPUPartitionedInput"
   input_arg {
     name: "inputs"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/TPUPartitionedInputV2.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/TPUPartitionedInputV2.pbtxt
index 52ec50cc1bf7bd..fe0ef345a1c628 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/TPUPartitionedInputV2.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/TPUPartitionedInputV2.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "TPUPartitionedInputV2"
   input_arg {
     name: "inputs"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/TPUPartitionedOutput.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/TPUPartitionedOutput.pbtxt
index ad5122c3f1b732..38a85e319644ac 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/TPUPartitionedOutput.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/TPUPartitionedOutput.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "TPUPartitionedOutput"
   input_arg {
     name: "inputs"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/TPUPartitionedOutputV2.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/TPUPartitionedOutputV2.pbtxt
index 83b7375fa2f1bf..3e7a6039ad8ca0 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/TPUPartitionedOutputV2.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/TPUPartitionedOutputV2.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "TPUPartitionedOutputV2"
   input_arg {
     name: "inputs"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/TPUReplicateMetadata.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/TPUReplicateMetadata.pbtxt
index 9ea81b7f929a1f..9742ad0d8d40c9 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/TPUReplicateMetadata.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/TPUReplicateMetadata.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "TPUReplicateMetadata"
   attr {
     name: "num_replicas"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/TPUReplicatedInput.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/TPUReplicatedInput.pbtxt
index f5a8ddcf284366..b549b570c13777 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/TPUReplicatedInput.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/TPUReplicatedInput.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "TPUReplicatedInput"
   input_arg {
     name: "inputs"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/TPUReplicatedOutput.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/TPUReplicatedOutput.pbtxt
index f7e9600cf4fb99..70b7d0ae71aadc 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/TPUReplicatedOutput.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/TPUReplicatedOutput.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "TPUReplicatedOutput"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/TPUReshardVariables.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/TPUReshardVariables.pbtxt
index 0bc7b4611fc7e8..fecd05d06eb1b9 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/TPUReshardVariables.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/TPUReshardVariables.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "TPUReshardVariables"
   input_arg {
     name: "vars"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/TPURoundRobin.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/TPURoundRobin.pbtxt
index 1405bc8cd6b167..7630e0cac746a0 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/TPURoundRobin.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/TPURoundRobin.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "TPURoundRobin"
   output_arg {
     name: "device_ordinal"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/TakeDataset.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/TakeDataset.pbtxt
index 4d9da96372ae09..8ced9c67054cc1 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/TakeDataset.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/TakeDataset.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "TakeDataset"
   input_arg {
     name: "input_dataset"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/TakeManySparseFromTensorsMap.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/TakeManySparseFromTensorsMap.pbtxt
index 7b852a26ddec81..0e3ca630eb178e 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/TakeManySparseFromTensorsMap.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/TakeManySparseFromTensorsMap.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "TakeManySparseFromTensorsMap"
   input_arg {
     name: "sparse_handles"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/TakeWhileDataset.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/TakeWhileDataset.pbtxt
index 7586b070a73be4..bfde2664966ef2 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/TakeWhileDataset.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/TakeWhileDataset.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "TakeWhileDataset"
   input_arg {
     name: "input_dataset"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/Tan.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/Tan.pbtxt
index 45442d6d712d71..a78f07b8d21382 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/Tan.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/Tan.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "Tan"
   input_arg {
     name: "x"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/Tanh.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/Tanh.pbtxt
index 9b7ccfe0bc9503..1672b0dc825c79 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/Tanh.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/Tanh.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "Tanh"
   input_arg {
     name: "x"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/TanhGrad.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/TanhGrad.pbtxt
index 41a827121b8c56..67d28f8ad7e0b4 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/TanhGrad.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/TanhGrad.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "TanhGrad"
   input_arg {
     name: "x"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/TemporaryVariable.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/TemporaryVariable.pbtxt
index d5c19d9f1ef34d..191354ec959700 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/TemporaryVariable.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/TemporaryVariable.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "TemporaryVariable"
   output_arg {
     name: "ref"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/TensorArray.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/TensorArray.pbtxt
index 67aa5f61327ac7..74b1a54976cc8b 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/TensorArray.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/TensorArray.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "TensorArray"
   input_arg {
     name: "size"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/TensorArrayClose.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/TensorArrayClose.pbtxt
index 44b06bcbf10e95..63c0100942005b 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/TensorArrayClose.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/TensorArrayClose.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "TensorArrayClose"
   input_arg {
     name: "handle"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/TensorArrayCloseV2.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/TensorArrayCloseV2.pbtxt
index 5298e3d85742b8..b0fb5804f1a2cf 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/TensorArrayCloseV2.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/TensorArrayCloseV2.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "TensorArrayCloseV2"
   input_arg {
     name: "handle"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/TensorArrayCloseV3.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/TensorArrayCloseV3.pbtxt
index 63af4407f62ae6..c5d1c2b1f244ce 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/TensorArrayCloseV3.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/TensorArrayCloseV3.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "TensorArrayCloseV3"
   input_arg {
     name: "handle"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/TensorArrayConcat.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/TensorArrayConcat.pbtxt
index 1111e79677e141..e2c59abd687402 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/TensorArrayConcat.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/TensorArrayConcat.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "TensorArrayConcat"
   input_arg {
     name: "handle"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/TensorArrayConcatV2.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/TensorArrayConcatV2.pbtxt
index b08f04af049c80..72376bd561c910 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/TensorArrayConcatV2.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/TensorArrayConcatV2.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "TensorArrayConcatV2"
   input_arg {
     name: "handle"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/TensorArrayConcatV3.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/TensorArrayConcatV3.pbtxt
index 70a7939d25c9d2..91e575ca87f6b3 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/TensorArrayConcatV3.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/TensorArrayConcatV3.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "TensorArrayConcatV3"
   input_arg {
     name: "handle"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/TensorArrayGather.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/TensorArrayGather.pbtxt
index 30b54dbce7b8b4..a8ded38550bc8b 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/TensorArrayGather.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/TensorArrayGather.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "TensorArrayGather"
   input_arg {
     name: "handle"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/TensorArrayGatherV2.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/TensorArrayGatherV2.pbtxt
index 3025ec43a0da9c..f72968388437f6 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/TensorArrayGatherV2.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/TensorArrayGatherV2.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "TensorArrayGatherV2"
   input_arg {
     name: "handle"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/TensorArrayGatherV3.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/TensorArrayGatherV3.pbtxt
index bf80504e78a1fb..c87538a40d26ab 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/TensorArrayGatherV3.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/TensorArrayGatherV3.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "TensorArrayGatherV3"
   input_arg {
     name: "handle"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/TensorArrayGrad.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/TensorArrayGrad.pbtxt
index 326cb594c3b5af..422154510dbcdf 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/TensorArrayGrad.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/TensorArrayGrad.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "TensorArrayGrad"
   input_arg {
     name: "handle"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/TensorArrayGradV2.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/TensorArrayGradV2.pbtxt
index 092dd8435055b7..d989c4071435ee 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/TensorArrayGradV2.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/TensorArrayGradV2.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "TensorArrayGradV2"
   input_arg {
     name: "handle"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/TensorArrayGradV3.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/TensorArrayGradV3.pbtxt
index 749e282ae92a2d..53e20429ec0e6f 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/TensorArrayGradV3.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/TensorArrayGradV3.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "TensorArrayGradV3"
   input_arg {
     name: "handle"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/TensorArrayGradWithShape.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/TensorArrayGradWithShape.pbtxt
index 100c3e7c78a4ef..1ce739062eb4c2 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/TensorArrayGradWithShape.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/TensorArrayGradWithShape.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "TensorArrayGradWithShape"
   input_arg {
     name: "handle"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/TensorArrayPack.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/TensorArrayPack.pbtxt
index 7b4ca8e9f4d39e..f608e453cee31b 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/TensorArrayPack.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/TensorArrayPack.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "TensorArrayPack"
   input_arg {
     name: "handle"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/TensorArrayRead.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/TensorArrayRead.pbtxt
index c3398036b5e4b5..62660bec758965 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/TensorArrayRead.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/TensorArrayRead.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "TensorArrayRead"
   input_arg {
     name: "handle"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/TensorArrayReadV2.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/TensorArrayReadV2.pbtxt
index 5620c0ade70a6c..cd0a2a32a8c06b 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/TensorArrayReadV2.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/TensorArrayReadV2.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "TensorArrayReadV2"
   input_arg {
     name: "handle"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/TensorArrayReadV3.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/TensorArrayReadV3.pbtxt
index bb53ce3649d56c..59e66fc84d5cdb 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/TensorArrayReadV3.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/TensorArrayReadV3.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "TensorArrayReadV3"
   input_arg {
     name: "handle"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/TensorArrayScatter.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/TensorArrayScatter.pbtxt
index 37c69a9b398e2e..b2017163f5c072 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/TensorArrayScatter.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/TensorArrayScatter.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "TensorArrayScatter"
   input_arg {
     name: "handle"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/TensorArrayScatterV2.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/TensorArrayScatterV2.pbtxt
index 798222d0d00d64..1eacf2d9acfa3c 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/TensorArrayScatterV2.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/TensorArrayScatterV2.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "TensorArrayScatterV2"
   input_arg {
     name: "handle"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/TensorArrayScatterV3.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/TensorArrayScatterV3.pbtxt
index c9ecffbc9692c1..5053ed60b1130e 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/TensorArrayScatterV3.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/TensorArrayScatterV3.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "TensorArrayScatterV3"
   input_arg {
     name: "handle"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/TensorArraySize.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/TensorArraySize.pbtxt
index 690f2118b6269f..7f6ce9510a0bff 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/TensorArraySize.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/TensorArraySize.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "TensorArraySize"
   input_arg {
     name: "handle"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/TensorArraySizeV2.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/TensorArraySizeV2.pbtxt
index bf446335854293..8ee9eda30bd21f 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/TensorArraySizeV2.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/TensorArraySizeV2.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "TensorArraySizeV2"
   input_arg {
     name: "handle"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/TensorArraySizeV3.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/TensorArraySizeV3.pbtxt
index c156e6af69ea9b..8932b0dcf2ddf7 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/TensorArraySizeV3.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/TensorArraySizeV3.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "TensorArraySizeV3"
   input_arg {
     name: "handle"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/TensorArraySplit.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/TensorArraySplit.pbtxt
index 4a693b29ff9cba..06bf8bfc3595fd 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/TensorArraySplit.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/TensorArraySplit.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "TensorArraySplit"
   input_arg {
     name: "handle"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/TensorArraySplitV2.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/TensorArraySplitV2.pbtxt
index fa9e02adb9326b..b45ea7a6108a66 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/TensorArraySplitV2.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/TensorArraySplitV2.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "TensorArraySplitV2"
   input_arg {
     name: "handle"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/TensorArraySplitV3.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/TensorArraySplitV3.pbtxt
index 2eed16078de27e..c072c0c65fc008 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/TensorArraySplitV3.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/TensorArraySplitV3.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "TensorArraySplitV3"
   input_arg {
     name: "handle"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/TensorArrayUnpack.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/TensorArrayUnpack.pbtxt
index 6fd145fbf55fac..81e5abec891b9b 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/TensorArrayUnpack.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/TensorArrayUnpack.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "TensorArrayUnpack"
   input_arg {
     name: "handle"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/TensorArrayV2.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/TensorArrayV2.pbtxt
index 3c01113a63e8c4..1293e1999c2031 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/TensorArrayV2.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/TensorArrayV2.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "TensorArrayV2"
   input_arg {
     name: "size"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/TensorArrayV3.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/TensorArrayV3.pbtxt
index ef70cea090839b..906e407de181e0 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/TensorArrayV3.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/TensorArrayV3.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "TensorArrayV3"
   input_arg {
     name: "size"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/TensorArrayWrite.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/TensorArrayWrite.pbtxt
index af5ad923f1d8c1..8f1a94c36b3878 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/TensorArrayWrite.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/TensorArrayWrite.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "TensorArrayWrite"
   input_arg {
     name: "handle"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/TensorArrayWriteV2.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/TensorArrayWriteV2.pbtxt
index ba33173b1dc971..fa0c1a679f7dfa 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/TensorArrayWriteV2.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/TensorArrayWriteV2.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "TensorArrayWriteV2"
   input_arg {
     name: "handle"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/TensorArrayWriteV3.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/TensorArrayWriteV3.pbtxt
index 44670bec744cf5..45327d42b1be1b 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/TensorArrayWriteV3.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/TensorArrayWriteV3.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "TensorArrayWriteV3"
   input_arg {
     name: "handle"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/TensorDataset.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/TensorDataset.pbtxt
index a70bea1c6975d5..9e71deef2c597c 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/TensorDataset.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/TensorDataset.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "TensorDataset"
   input_arg {
     name: "components"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/TensorListConcat.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/TensorListConcat.pbtxt
index 8b715d4864011e..010be2e120b46e 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/TensorListConcat.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/TensorListConcat.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "TensorListConcat"
   input_arg {
     name: "input_handle"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/TensorListConcatLists.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/TensorListConcatLists.pbtxt
index e532be3918c367..57dd05a90feb7d 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/TensorListConcatLists.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/TensorListConcatLists.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "TensorListConcatLists"
   input_arg {
     name: "input_a"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/TensorListConcatV2.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/TensorListConcatV2.pbtxt
index 6051430fd552bf..0bb9546d155d29 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/TensorListConcatV2.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/TensorListConcatV2.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "TensorListConcatV2"
   input_arg {
     name: "input_handle"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/TensorListElementShape.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/TensorListElementShape.pbtxt
index e4143fdbe9dd06..26b982f6cfd9c3 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/TensorListElementShape.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/TensorListElementShape.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "TensorListElementShape"
   input_arg {
     name: "input_handle"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/TensorListFromTensor.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/TensorListFromTensor.pbtxt
index f3111b9756db4e..c2be2938c8607e 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/TensorListFromTensor.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/TensorListFromTensor.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "TensorListFromTensor"
   input_arg {
     name: "tensor"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/TensorListGather.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/TensorListGather.pbtxt
index 271fc2db2c0161..43b4773a4f7e0f 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/TensorListGather.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/TensorListGather.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "TensorListGather"
   input_arg {
     name: "input_handle"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/TensorListGetItem.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/TensorListGetItem.pbtxt
index ff087f92a0721a..fa124bc94971fe 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/TensorListGetItem.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/TensorListGetItem.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "TensorListGetItem"
   input_arg {
     name: "input_handle"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/TensorListLength.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/TensorListLength.pbtxt
index 8b95320f200740..b4ea660dca13de 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/TensorListLength.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/TensorListLength.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "TensorListLength"
   input_arg {
     name: "input_handle"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/TensorListPopBack.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/TensorListPopBack.pbtxt
index 61dd1111cb64f5..35aa68e07584cc 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/TensorListPopBack.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/TensorListPopBack.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "TensorListPopBack"
   input_arg {
     name: "input_handle"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/TensorListPushBack.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/TensorListPushBack.pbtxt
index 1ae20365fb9bab..8175cfe350dab0 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/TensorListPushBack.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/TensorListPushBack.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "TensorListPushBack"
   input_arg {
     name: "input_handle"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/TensorListPushBackBatch.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/TensorListPushBackBatch.pbtxt
index 6e805caed1bcc2..29b878e527a4d5 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/TensorListPushBackBatch.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/TensorListPushBackBatch.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "TensorListPushBackBatch"
   input_arg {
     name: "input_handles"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/TensorListReserve.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/TensorListReserve.pbtxt
index bef49b63477ebe..98ade8cb4a9fd0 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/TensorListReserve.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/TensorListReserve.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "TensorListReserve"
   input_arg {
     name: "element_shape"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/TensorListResize.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/TensorListResize.pbtxt
index 50e070018e2a0b..b322d89d47aa68 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/TensorListResize.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/TensorListResize.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "TensorListResize"
   input_arg {
     name: "input_handle"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/TensorListScatter.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/TensorListScatter.pbtxt
index 31fa2452bb11d9..daa2f4130ab06b 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/TensorListScatter.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/TensorListScatter.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "TensorListScatter"
   input_arg {
     name: "tensor"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/TensorListScatterIntoExistingList.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/TensorListScatterIntoExistingList.pbtxt
index 311157650998f6..4427bab8a358c2 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/TensorListScatterIntoExistingList.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/TensorListScatterIntoExistingList.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "TensorListScatterIntoExistingList"
   input_arg {
     name: "input_handle"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/TensorListScatterV2.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/TensorListScatterV2.pbtxt
index 29cf4aba78680f..de588984614839 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/TensorListScatterV2.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/TensorListScatterV2.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "TensorListScatterV2"
   input_arg {
     name: "tensor"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/TensorListSetItem.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/TensorListSetItem.pbtxt
index d16b98d59bdc54..e2510ca98e0b8f 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/TensorListSetItem.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/TensorListSetItem.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "TensorListSetItem"
   input_arg {
     name: "input_handle"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/TensorListSplit.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/TensorListSplit.pbtxt
index 2ed8f207fed4c3..ff83247addf89b 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/TensorListSplit.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/TensorListSplit.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "TensorListSplit"
   input_arg {
     name: "tensor"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/TensorListStack.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/TensorListStack.pbtxt
index 8623e34a934845..5a8e7bcd81b9ef 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/TensorListStack.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/TensorListStack.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "TensorListStack"
   input_arg {
     name: "input_handle"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/TensorMapErase.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/TensorMapErase.pbtxt
index 009fd49ee6288a..854e7311eab331 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/TensorMapErase.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/TensorMapErase.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "TensorMapErase"
   input_arg {
     name: "input_handle"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/TensorMapHasKey.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/TensorMapHasKey.pbtxt
index e2856e3ecaea56..a095c36d7c26f5 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/TensorMapHasKey.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/TensorMapHasKey.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "TensorMapHasKey"
   input_arg {
     name: "input_handle"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/TensorMapInsert.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/TensorMapInsert.pbtxt
index 492d3538de0ed5..10061ea1cde6dc 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/TensorMapInsert.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/TensorMapInsert.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "TensorMapInsert"
   input_arg {
     name: "input_handle"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/TensorMapLookup.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/TensorMapLookup.pbtxt
index 9f1097226fc7b3..b48fda8ac4623f 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/TensorMapLookup.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/TensorMapLookup.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "TensorMapLookup"
   input_arg {
     name: "input_handle"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/TensorMapSize.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/TensorMapSize.pbtxt
index a94f2c515ece36..dd8ade84414f56 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/TensorMapSize.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/TensorMapSize.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "TensorMapSize"
   input_arg {
     name: "input_handle"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/TensorMapStackKeys.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/TensorMapStackKeys.pbtxt
index 10aef43aca1343..c3befaa320a385 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/TensorMapStackKeys.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/TensorMapStackKeys.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "TensorMapStackKeys"
   input_arg {
     name: "input_handle"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/TensorScatterAdd.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/TensorScatterAdd.pbtxt
index 71f0e9fb625334..5fb5b8cb0dd693 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/TensorScatterAdd.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/TensorScatterAdd.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "TensorScatterAdd"
   input_arg {
     name: "tensor"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/TensorScatterMax.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/TensorScatterMax.pbtxt
index 3abba955f62158..84a05c9bd5dba9 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/TensorScatterMax.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/TensorScatterMax.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "TensorScatterMax"
   input_arg {
     name: "tensor"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/TensorScatterMin.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/TensorScatterMin.pbtxt
index 047619805f2aff..d1ae6117921ee0 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/TensorScatterMin.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/TensorScatterMin.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "TensorScatterMin"
   input_arg {
     name: "tensor"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/TensorScatterSub.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/TensorScatterSub.pbtxt
index d66d823ef317c8..81920523ae9bbc 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/TensorScatterSub.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/TensorScatterSub.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "TensorScatterSub"
   input_arg {
     name: "tensor"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/TensorScatterUpdate.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/TensorScatterUpdate.pbtxt
index 4bc0747881f943..1e8281cc247958 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/TensorScatterUpdate.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/TensorScatterUpdate.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "TensorScatterUpdate"
   input_arg {
     name: "tensor"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/TensorSliceDataset.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/TensorSliceDataset.pbtxt
index 7be6e03ca4a19c..3810a6b4023657 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/TensorSliceDataset.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/TensorSliceDataset.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "TensorSliceDataset"
   input_arg {
     name: "components"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/TensorStridedSliceUpdate.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/TensorStridedSliceUpdate.pbtxt
index 67dae5ab4dd346..3854eeed137057 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/TensorStridedSliceUpdate.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/TensorStridedSliceUpdate.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "TensorStridedSliceUpdate"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/TensorSummary.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/TensorSummary.pbtxt
index 5486c2ef84aaed..bf4114aeef398c 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/TensorSummary.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/TensorSummary.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "TensorSummary"
   input_arg {
     name: "tensor"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/TensorSummaryV2.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/TensorSummaryV2.pbtxt
index 107c53939f3b29..39092b078161f0 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/TensorSummaryV2.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/TensorSummaryV2.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "TensorSummaryV2"
   input_arg {
     name: "tag"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/TextLineDataset.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/TextLineDataset.pbtxt
index 5f539c525b0048..c51a34124f2d1e 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/TextLineDataset.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/TextLineDataset.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "TextLineDataset"
   input_arg {
     name: "filenames"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/TextLineReader.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/TextLineReader.pbtxt
index 2e0924bb51786e..baf1ef10d91047 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/TextLineReader.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/TextLineReader.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "TextLineReader"
   output_arg {
     name: "reader_handle"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/TextLineReaderV2.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/TextLineReaderV2.pbtxt
index d30d08b5fbf150..c669951acdf6b6 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/TextLineReaderV2.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/TextLineReaderV2.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "TextLineReaderV2"
   output_arg {
     name: "reader_handle"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ThreadPoolDataset.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ThreadPoolDataset.pbtxt
index b8c817fc13876a..8e185af579f2be 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ThreadPoolDataset.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ThreadPoolDataset.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ThreadPoolDataset"
   input_arg {
     name: "input_dataset"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ThreadPoolHandle.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ThreadPoolHandle.pbtxt
index 4fac8fb83b79f1..e2518b1439d732 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ThreadPoolHandle.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ThreadPoolHandle.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ThreadPoolHandle"
   output_arg {
     name: "handle"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ThreadUnsafeUnigramCandidateSampler.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ThreadUnsafeUnigramCandidateSampler.pbtxt
index 0e99e93edcc732..89106aab220583 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ThreadUnsafeUnigramCandidateSampler.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ThreadUnsafeUnigramCandidateSampler.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ThreadUnsafeUnigramCandidateSampler"
   input_arg {
     name: "true_classes"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/Tile.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/Tile.pbtxt
index f3a139e2265abd..67de1e5201698c 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/Tile.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/Tile.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "Tile"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/TileGrad.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/TileGrad.pbtxt
index f068a74b1c3d3d..f710e1c470f254 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/TileGrad.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/TileGrad.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "TileGrad"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/Timestamp.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/Timestamp.pbtxt
index 18f5f4d2d23f03..6e51504d17653e 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/Timestamp.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/Timestamp.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "Timestamp"
   output_arg {
     name: "ts"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ToBool.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ToBool.pbtxt
index 2903fe7f0b2e7f..6e02fdbb52cdeb 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ToBool.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ToBool.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ToBool"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/TopK.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/TopK.pbtxt
index 8ad9a3c23bb9c7..71c98b7fd8120e 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/TopK.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/TopK.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "TopK"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/TopKUnique.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/TopKUnique.pbtxt
index 23b15af2aef16e..12463385bcc816 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/TopKUnique.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/TopKUnique.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "TopKUnique"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/TopKV2.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/TopKV2.pbtxt
index 5da8c479b9a8b3..22908f661e62f9 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/TopKV2.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/TopKV2.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "TopKV2"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/TopKWithUnique.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/TopKWithUnique.pbtxt
index cfb4e9da928507..5e3216fa554877 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/TopKWithUnique.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/TopKWithUnique.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "TopKWithUnique"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/TpuHandleToProtoKey.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/TpuHandleToProtoKey.pbtxt
index 3d2c41d36b4d15..1f8d36887b9733 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/TpuHandleToProtoKey.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/TpuHandleToProtoKey.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "TpuHandleToProtoKey"
   input_arg {
     name: "uid"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/Transpose.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/Transpose.pbtxt
index e83ec349f19b72..fa4fb6d58937c4 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/Transpose.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/Transpose.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "Transpose"
   input_arg {
     name: "x"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/TridiagonalMatMul.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/TridiagonalMatMul.pbtxt
index 98f70443f67a82..117d68b48d29b6 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/TridiagonalMatMul.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/TridiagonalMatMul.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "TridiagonalMatMul"
   input_arg {
     name: "superdiag"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/TridiagonalSolve.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/TridiagonalSolve.pbtxt
index f86be1a7508c37..d824c1cb76c35e 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/TridiagonalSolve.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/TridiagonalSolve.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "TridiagonalSolve"
   input_arg {
     name: "diagonals"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/TruncateDiv.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/TruncateDiv.pbtxt
index 74a4db11dc5efc..0b6c414f616163 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/TruncateDiv.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/TruncateDiv.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "TruncateDiv"
   input_arg {
     name: "x"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/TruncateMod.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/TruncateMod.pbtxt
index 72517db5294d89..70ce81b35c4209 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/TruncateMod.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/TruncateMod.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "TruncateMod"
   input_arg {
     name: "x"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/TruncatedNormal.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/TruncatedNormal.pbtxt
index c066b3283ec0ad..018d657985d5e8 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/TruncatedNormal.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/TruncatedNormal.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "TruncatedNormal"
   input_arg {
     name: "shape"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/Unbatch.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/Unbatch.pbtxt
index 624522ee54ced0..3934b1823ff052 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/Unbatch.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/Unbatch.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "Unbatch"
   input_arg {
     name: "batched_tensor"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/UnbatchDataset.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/UnbatchDataset.pbtxt
index fa3075b92f41ee..84479c117206d7 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/UnbatchDataset.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/UnbatchDataset.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "UnbatchDataset"
   input_arg {
     name: "input_dataset"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/UnbatchGrad.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/UnbatchGrad.pbtxt
index f2619637143ee8..97240f0be53a0f 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/UnbatchGrad.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/UnbatchGrad.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "UnbatchGrad"
   input_arg {
     name: "original_input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/UncompressElement.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/UncompressElement.pbtxt
index 04674945d82c97..68406e0e4bc755 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/UncompressElement.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/UncompressElement.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "UncompressElement"
   input_arg {
     name: "compressed"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/UnicodeDecode.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/UnicodeDecode.pbtxt
index a8aac23ae733e7..fa036b31ef6c38 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/UnicodeDecode.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/UnicodeDecode.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "UnicodeDecode"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/UnicodeDecodeWithOffsets.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/UnicodeDecodeWithOffsets.pbtxt
index 05a35cc1ad2c7f..29d274738da829 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/UnicodeDecodeWithOffsets.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/UnicodeDecodeWithOffsets.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "UnicodeDecodeWithOffsets"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/UnicodeEncode.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/UnicodeEncode.pbtxt
index de0b916b6ab2a2..31a7a5b838820d 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/UnicodeEncode.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/UnicodeEncode.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "UnicodeEncode"
   input_arg {
     name: "input_values"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/UnicodeScript.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/UnicodeScript.pbtxt
index b9d7590e7c2ff2..60877b544480fe 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/UnicodeScript.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/UnicodeScript.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "UnicodeScript"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/UnicodeTranscode.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/UnicodeTranscode.pbtxt
index 494d7a9d6d8373..5cab73782ce8ec 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/UnicodeTranscode.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/UnicodeTranscode.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "UnicodeTranscode"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/UniformCandidateSampler.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/UniformCandidateSampler.pbtxt
index affff4ad02d2dc..bea963f908ee14 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/UniformCandidateSampler.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/UniformCandidateSampler.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "UniformCandidateSampler"
   input_arg {
     name: "true_classes"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/UniformDequantize.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/UniformDequantize.pbtxt
index 68557e82c563e1..7653370635d18c 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/UniformDequantize.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/UniformDequantize.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "UniformDequantize"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/UniformQuantize.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/UniformQuantize.pbtxt
index 3e7d42c86198c8..900da8e5be53c3 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/UniformQuantize.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/UniformQuantize.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "UniformQuantize"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/UniformQuantizedAdd.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/UniformQuantizedAdd.pbtxt
index 480e70a6693739..1adac5ae59e790 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/UniformQuantizedAdd.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/UniformQuantizedAdd.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "UniformQuantizedAdd"
   input_arg {
     name: "lhs"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/UniformQuantizedClipByValue.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/UniformQuantizedClipByValue.pbtxt
index a9d0a79fded26e..8ab01d6c42d9a6 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/UniformQuantizedClipByValue.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/UniformQuantizedClipByValue.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "UniformQuantizedClipByValue"
   input_arg {
     name: "operand"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/UniformQuantizedConvolution.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/UniformQuantizedConvolution.pbtxt
index 68c4746c5210ae..4ebbe6c80690e8 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/UniformQuantizedConvolution.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/UniformQuantizedConvolution.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "UniformQuantizedConvolution"
   input_arg {
     name: "lhs"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/UniformQuantizedConvolutionHybrid.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/UniformQuantizedConvolutionHybrid.pbtxt
index 9a9f8fc085fccb..23096f391d9b6e 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/UniformQuantizedConvolutionHybrid.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/UniformQuantizedConvolutionHybrid.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "UniformQuantizedConvolutionHybrid"
   input_arg {
     name: "lhs"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/UniformQuantizedDot.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/UniformQuantizedDot.pbtxt
index 9b3030bb3bd8a6..159dabb9798621 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/UniformQuantizedDot.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/UniformQuantizedDot.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "UniformQuantizedDot"
   input_arg {
     name: "lhs"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/UniformQuantizedDotHybrid.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/UniformQuantizedDotHybrid.pbtxt
index 38c8ed55f5978e..2cabf91667e386 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/UniformQuantizedDotHybrid.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/UniformQuantizedDotHybrid.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "UniformQuantizedDotHybrid"
   input_arg {
     name: "lhs"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/UniformRequantize.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/UniformRequantize.pbtxt
index bfc3d47945b686..af2d7387c5d695 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/UniformRequantize.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/UniformRequantize.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "UniformRequantize"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/Unique.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/Unique.pbtxt
index 9d91d9ef4c6164..be389ba1482be8 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/Unique.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/Unique.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "Unique"
   input_arg {
     name: "x"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/UniqueDataset.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/UniqueDataset.pbtxt
index 58fc7e07a35822..281ba7c1bd0619 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/UniqueDataset.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/UniqueDataset.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "UniqueDataset"
   input_arg {
     name: "input_dataset"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/UniqueV2.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/UniqueV2.pbtxt
index e76969551065c0..83113e14232fb2 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/UniqueV2.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/UniqueV2.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "UniqueV2"
   input_arg {
     name: "x"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/UniqueWithCounts.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/UniqueWithCounts.pbtxt
index 4d5682561cf926..c386059943a70e 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/UniqueWithCounts.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/UniqueWithCounts.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "UniqueWithCounts"
   input_arg {
     name: "x"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/UniqueWithCountsV2.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/UniqueWithCountsV2.pbtxt
index d54bba74d97327..85a12b70007320 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/UniqueWithCountsV2.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/UniqueWithCountsV2.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "UniqueWithCountsV2"
   input_arg {
     name: "x"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/Unpack.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/Unpack.pbtxt
index 6f24385e25bce6..cc5fd918d4c694 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/Unpack.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/Unpack.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "Unpack"
   input_arg {
     name: "value"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/UnravelIndex.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/UnravelIndex.pbtxt
index 36e66f2d2b0f4b..df2c2bc8469451 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/UnravelIndex.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/UnravelIndex.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "UnravelIndex"
   input_arg {
     name: "indices"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/UnsortedSegmentJoin.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/UnsortedSegmentJoin.pbtxt
index 8c95acc7ab9dba..dcbb91bc2f13c4 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/UnsortedSegmentJoin.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/UnsortedSegmentJoin.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "UnsortedSegmentJoin"
   input_arg {
     name: "inputs"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/UnsortedSegmentMax.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/UnsortedSegmentMax.pbtxt
index fc9ca18d6cfeed..ee8578f289bdc0 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/UnsortedSegmentMax.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/UnsortedSegmentMax.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "UnsortedSegmentMax"
   input_arg {
     name: "data"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/UnsortedSegmentMin.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/UnsortedSegmentMin.pbtxt
index 33e3635173d589..6a8e5ba6d1fbe6 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/UnsortedSegmentMin.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/UnsortedSegmentMin.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "UnsortedSegmentMin"
   input_arg {
     name: "data"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/UnsortedSegmentProd.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/UnsortedSegmentProd.pbtxt
index e1543a33441e81..d100cde127e9f1 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/UnsortedSegmentProd.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/UnsortedSegmentProd.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "UnsortedSegmentProd"
   input_arg {
     name: "data"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/UnsortedSegmentSum.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/UnsortedSegmentSum.pbtxt
index ff0bf2d7a8cf8a..28ddbd3bdec499 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/UnsortedSegmentSum.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/UnsortedSegmentSum.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "UnsortedSegmentSum"
   input_arg {
     name: "data"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/Unstage.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/Unstage.pbtxt
index af51a8ce0b8ca6..4bcfd02758ced6 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/Unstage.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/Unstage.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "Unstage"
   output_arg {
     name: "values"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/UnwrapDatasetVariant.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/UnwrapDatasetVariant.pbtxt
index fc01d23c863703..10e23a97750a70 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/UnwrapDatasetVariant.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/UnwrapDatasetVariant.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "UnwrapDatasetVariant"
   input_arg {
     name: "input_handle"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/UpperBound.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/UpperBound.pbtxt
index 5cfd0d536c5238..d1b3fa060c6942 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/UpperBound.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/UpperBound.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "UpperBound"
   input_arg {
     name: "sorted_inputs"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/VarHandleOp.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/VarHandleOp.pbtxt
index acca5ff60992e7..80bc633b5ac613 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/VarHandleOp.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/VarHandleOp.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "VarHandleOp"
   output_arg {
     name: "resource"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/VarIsInitializedOp.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/VarIsInitializedOp.pbtxt
index 975983e81af8ff..395360158262b3 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/VarIsInitializedOp.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/VarIsInitializedOp.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "VarIsInitializedOp"
   input_arg {
     name: "resource"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/Variable.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/Variable.pbtxt
index 7ad65f1fe6a3f6..943c24def5944d 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/Variable.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/Variable.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "Variable"
   output_arg {
     name: "ref"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/VariableShape.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/VariableShape.pbtxt
index e0bb7d05dafe25..570b4f241aaa95 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/VariableShape.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/VariableShape.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "VariableShape"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/VariableV2.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/VariableV2.pbtxt
index dabd46dcfb467c..c27112f15887b1 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/VariableV2.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/VariableV2.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "VariableV2"
   output_arg {
     name: "ref"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/WeightedFlatMapDataset.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/WeightedFlatMapDataset.pbtxt
index 5c50229e4c5b07..98c853bb2f3e64 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/WeightedFlatMapDataset.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/WeightedFlatMapDataset.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "WeightedFlatMapDataset"
   input_arg {
     name: "input_datasets"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/Where.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/Where.pbtxt
index b23dca17037197..8e64cd2419e2c1 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/Where.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/Where.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "Where"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/While.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/While.pbtxt
index 8e609867f8d94f..807461b00984bd 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/While.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/While.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "While"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/WholeFileReader.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/WholeFileReader.pbtxt
index b1513138650504..729d76503e53dd 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/WholeFileReader.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/WholeFileReader.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "WholeFileReader"
   output_arg {
     name: "reader_handle"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/WholeFileReaderV2.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/WholeFileReaderV2.pbtxt
index f451cf41c57bc3..2430494342d709 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/WholeFileReaderV2.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/WholeFileReaderV2.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "WholeFileReaderV2"
   output_arg {
     name: "reader_handle"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/WindowDataset.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/WindowDataset.pbtxt
index a5fe3aabbfe45e..43784faad6dc62 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/WindowDataset.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/WindowDataset.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "WindowDataset"
   input_arg {
     name: "input_dataset"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/WindowOp.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/WindowOp.pbtxt
index 60e080a0472b0e..336a13805eaf6a 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/WindowOp.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/WindowOp.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "WindowOp"
   input_arg {
     name: "inputs"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/WorkerHeartbeat.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/WorkerHeartbeat.pbtxt
index bfadb835956750..ae5c7b8caaad6a 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/WorkerHeartbeat.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/WorkerHeartbeat.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "WorkerHeartbeat"
   input_arg {
     name: "request"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/WrapDatasetVariant.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/WrapDatasetVariant.pbtxt
index 247262ec65ce5d..0b1e4363bd20e6 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/WrapDatasetVariant.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/WrapDatasetVariant.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "WrapDatasetVariant"
   input_arg {
     name: "input_handle"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/WriteAudioSummary.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/WriteAudioSummary.pbtxt
index 6ec7f394c44bb2..8cc81eba8ff3a2 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/WriteAudioSummary.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/WriteAudioSummary.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "WriteAudioSummary"
   input_arg {
     name: "writer"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/WriteFile.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/WriteFile.pbtxt
index 55a4c60b10173b..6a15b39873d560 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/WriteFile.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/WriteFile.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "WriteFile"
   input_arg {
     name: "filename"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/WriteGraphSummary.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/WriteGraphSummary.pbtxt
index 7851117301a1b3..2957e224f59514 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/WriteGraphSummary.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/WriteGraphSummary.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "WriteGraphSummary"
   input_arg {
     name: "writer"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/WriteHistogramSummary.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/WriteHistogramSummary.pbtxt
index 869b72fc133101..492d573056823b 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/WriteHistogramSummary.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/WriteHistogramSummary.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "WriteHistogramSummary"
   input_arg {
     name: "writer"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/WriteImageSummary.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/WriteImageSummary.pbtxt
index 45e3243c6bcf40..1cfc5ca69a7f40 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/WriteImageSummary.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/WriteImageSummary.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "WriteImageSummary"
   input_arg {
     name: "writer"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/WriteRawProtoSummary.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/WriteRawProtoSummary.pbtxt
index 6c571b9f2c7f7a..82ac51a137894b 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/WriteRawProtoSummary.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/WriteRawProtoSummary.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "WriteRawProtoSummary"
   input_arg {
     name: "writer"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/WriteScalarSummary.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/WriteScalarSummary.pbtxt
index e40411aabdc256..0f359a85dce91b 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/WriteScalarSummary.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/WriteScalarSummary.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "WriteScalarSummary"
   input_arg {
     name: "writer"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/WriteSummary.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/WriteSummary.pbtxt
index f6f887199db2e1..a641ece08df095 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/WriteSummary.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/WriteSummary.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "WriteSummary"
   input_arg {
     name: "writer"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/Xdivy.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/Xdivy.pbtxt
index 6536322552738d..898987f947cd3a 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/Xdivy.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/Xdivy.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "Xdivy"
   input_arg {
     name: "x"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/XlaConcatND.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/XlaConcatND.pbtxt
index 7d320b8f5c4544..c571497bcd5535 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/XlaConcatND.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/XlaConcatND.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "XlaConcatND"
   input_arg {
     name: "inputs"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/XlaHostCompute.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/XlaHostCompute.pbtxt
index 787e744ed8350a..9675bda59a86dd 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/XlaHostCompute.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/XlaHostCompute.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "XlaHostCompute"
   input_arg {
     name: "inputs"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/XlaRecvFromHost.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/XlaRecvFromHost.pbtxt
index 10c8f79678c334..d3760ea79b1272 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/XlaRecvFromHost.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/XlaRecvFromHost.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "XlaRecvFromHost"
   output_arg {
     name: "output"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/XlaRecvTPUEmbeddingActivations.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/XlaRecvTPUEmbeddingActivations.pbtxt
index 208eea2857258f..b624b26a628bc2 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/XlaRecvTPUEmbeddingActivations.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/XlaRecvTPUEmbeddingActivations.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "XlaRecvTPUEmbeddingActivations"
   input_arg {
     name: "deduplication_data"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/XlaRecvTPUEmbeddingActivationsV2.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/XlaRecvTPUEmbeddingActivationsV2.pbtxt
index c5abbb2d5f1a4e..2e8fb4d4f2530c 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/XlaRecvTPUEmbeddingActivationsV2.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/XlaRecvTPUEmbeddingActivationsV2.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "XlaRecvTPUEmbeddingActivationsV2"
   input_arg {
     name: "deduplication_data"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/XlaRecvTPUEmbeddingDeduplicationData.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/XlaRecvTPUEmbeddingDeduplicationData.pbtxt
index c931212779893b..3c3b92f9d7bff7 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/XlaRecvTPUEmbeddingDeduplicationData.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/XlaRecvTPUEmbeddingDeduplicationData.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "XlaRecvTPUEmbeddingDeduplicationData"
   output_arg {
     name: "output"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/XlaRecvTPUEmbeddingDeduplicationDataV2.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/XlaRecvTPUEmbeddingDeduplicationDataV2.pbtxt
index 71632c6d871452..d97710b91e46fb 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/XlaRecvTPUEmbeddingDeduplicationDataV2.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/XlaRecvTPUEmbeddingDeduplicationDataV2.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "XlaRecvTPUEmbeddingDeduplicationDataV2"
   output_arg {
     name: "output"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/XlaSendTPUEmbeddingGradients.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/XlaSendTPUEmbeddingGradients.pbtxt
index 6da5fb2f5395cb..77f6547229554f 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/XlaSendTPUEmbeddingGradients.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/XlaSendTPUEmbeddingGradients.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "XlaSendTPUEmbeddingGradients"
   input_arg {
     name: "gradients"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/XlaSendTPUEmbeddingGradientsV2.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/XlaSendTPUEmbeddingGradientsV2.pbtxt
index f52b83abc14719..b416d0ad1a8f0c 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/XlaSendTPUEmbeddingGradientsV2.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/XlaSendTPUEmbeddingGradientsV2.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "XlaSendTPUEmbeddingGradientsV2"
   input_arg {
     name: "gradients"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/XlaSendToHost.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/XlaSendToHost.pbtxt
index 90cc16e20a2e95..f2dfeaf444491e 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/XlaSendToHost.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/XlaSendToHost.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "XlaSendToHost"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/XlaSparseCoreAdagrad.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/XlaSparseCoreAdagrad.pbtxt
index 9cf626b404b1e7..bc20baf287c8b1 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/XlaSparseCoreAdagrad.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/XlaSparseCoreAdagrad.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "XlaSparseCoreAdagrad"
   input_arg {
     name: "indices"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/XlaSparseCoreAdagradMomentum.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/XlaSparseCoreAdagradMomentum.pbtxt
index b64460485d5ac9..5e2b17ab4c9238 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/XlaSparseCoreAdagradMomentum.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/XlaSparseCoreAdagradMomentum.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "XlaSparseCoreAdagradMomentum"
   input_arg {
     name: "indices"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/XlaSparseCoreAdam.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/XlaSparseCoreAdam.pbtxt
index 38af8afcc1d10d..625c164bb20ed3 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/XlaSparseCoreAdam.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/XlaSparseCoreAdam.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "XlaSparseCoreAdam"
   input_arg {
     name: "embedding_table"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/XlaSparseCoreFtrl.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/XlaSparseCoreFtrl.pbtxt
index afbf9e024d7041..b65b707befee31 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/XlaSparseCoreFtrl.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/XlaSparseCoreFtrl.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "XlaSparseCoreFtrl"
   input_arg {
     name: "embedding_table"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/XlaSparseCoreSgd.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/XlaSparseCoreSgd.pbtxt
index 7f507c7d722106..677ce14ddf039d 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/XlaSparseCoreSgd.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/XlaSparseCoreSgd.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "XlaSparseCoreSgd"
   input_arg {
     name: "indices"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/XlaSparseDenseMatmul.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/XlaSparseDenseMatmul.pbtxt
index 5ecf0c2472748a..90aa2cf22849b3 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/XlaSparseDenseMatmul.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/XlaSparseDenseMatmul.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "XlaSparseDenseMatmul"
   input_arg {
     name: "row_ids"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/XlaSparseDenseMatmulGradWithAdagradAndCsrInput.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/XlaSparseDenseMatmulGradWithAdagradAndCsrInput.pbtxt
index e13cbfcff32417..fbf266b10e35fa 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/XlaSparseDenseMatmulGradWithAdagradAndCsrInput.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/XlaSparseDenseMatmulGradWithAdagradAndCsrInput.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "XlaSparseDenseMatmulGradWithAdagradAndCsrInput"
   input_arg {
     name: "row_pointers"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/XlaSparseDenseMatmulGradWithAdagradAndStaticBufferSize.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/XlaSparseDenseMatmulGradWithAdagradAndStaticBufferSize.pbtxt
index 182bfdbf3bcab2..359a038ea9b6b7 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/XlaSparseDenseMatmulGradWithAdagradAndStaticBufferSize.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/XlaSparseDenseMatmulGradWithAdagradAndStaticBufferSize.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "XlaSparseDenseMatmulGradWithAdagradAndStaticBufferSize"
   input_arg {
     name: "row_pointers"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/XlaSparseDenseMatmulGradWithAdagradMomentumAndCsrInput.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/XlaSparseDenseMatmulGradWithAdagradMomentumAndCsrInput.pbtxt
index e6f2eed0c1d75b..5150a4f23b598f 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/XlaSparseDenseMatmulGradWithAdagradMomentumAndCsrInput.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/XlaSparseDenseMatmulGradWithAdagradMomentumAndCsrInput.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "XlaSparseDenseMatmulGradWithAdagradMomentumAndCsrInput"
   input_arg {
     name: "row_pointers"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/XlaSparseDenseMatmulGradWithAdagradMomentumAndStaticBufferSize.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/XlaSparseDenseMatmulGradWithAdagradMomentumAndStaticBufferSize.pbtxt
index 3bd492cc55a70c..4fd6fa9bb5a5b2 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/XlaSparseDenseMatmulGradWithAdagradMomentumAndStaticBufferSize.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/XlaSparseDenseMatmulGradWithAdagradMomentumAndStaticBufferSize.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "XlaSparseDenseMatmulGradWithAdagradMomentumAndStaticBufferSize"
   input_arg {
     name: "row_pointers"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/XlaSparseDenseMatmulGradWithAdamAndCsrInput.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/XlaSparseDenseMatmulGradWithAdamAndCsrInput.pbtxt
index 202e6f4f4f7b09..aaa27b25954a9e 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/XlaSparseDenseMatmulGradWithAdamAndCsrInput.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/XlaSparseDenseMatmulGradWithAdamAndCsrInput.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "XlaSparseDenseMatmulGradWithAdamAndCsrInput"
   input_arg {
     name: "row_pointers"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/XlaSparseDenseMatmulGradWithAdamAndStaticBufferSize.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/XlaSparseDenseMatmulGradWithAdamAndStaticBufferSize.pbtxt
index f058e30f800527..5024f72b5c66cb 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/XlaSparseDenseMatmulGradWithAdamAndStaticBufferSize.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/XlaSparseDenseMatmulGradWithAdamAndStaticBufferSize.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "XlaSparseDenseMatmulGradWithAdamAndStaticBufferSize"
   input_arg {
     name: "row_pointers"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/XlaSparseDenseMatmulGradWithFtrlAndCsrInput.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/XlaSparseDenseMatmulGradWithFtrlAndCsrInput.pbtxt
index 96121a6bab883d..261f25bebfd7df 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/XlaSparseDenseMatmulGradWithFtrlAndCsrInput.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/XlaSparseDenseMatmulGradWithFtrlAndCsrInput.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "XlaSparseDenseMatmulGradWithFtrlAndCsrInput"
   input_arg {
     name: "row_pointers"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/XlaSparseDenseMatmulGradWithFtrlAndStaticBufferSize.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/XlaSparseDenseMatmulGradWithFtrlAndStaticBufferSize.pbtxt
index 23c6d671dd7914..f2f57f2f744d7b 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/XlaSparseDenseMatmulGradWithFtrlAndStaticBufferSize.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/XlaSparseDenseMatmulGradWithFtrlAndStaticBufferSize.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "XlaSparseDenseMatmulGradWithFtrlAndStaticBufferSize"
   input_arg {
     name: "row_pointers"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/XlaSparseDenseMatmulGradWithSgdAndCsrInput.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/XlaSparseDenseMatmulGradWithSgdAndCsrInput.pbtxt
index 3ad518fadb629e..9446a6fa98c515 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/XlaSparseDenseMatmulGradWithSgdAndCsrInput.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/XlaSparseDenseMatmulGradWithSgdAndCsrInput.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "XlaSparseDenseMatmulGradWithSgdAndCsrInput"
   input_arg {
     name: "row_pointers"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/XlaSparseDenseMatmulGradWithSgdAndStaticBufferSize.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/XlaSparseDenseMatmulGradWithSgdAndStaticBufferSize.pbtxt
index 46a1fe8d6e2f6a..dbb06c95f6d643 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/XlaSparseDenseMatmulGradWithSgdAndStaticBufferSize.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/XlaSparseDenseMatmulGradWithSgdAndStaticBufferSize.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "XlaSparseDenseMatmulGradWithSgdAndStaticBufferSize"
   input_arg {
     name: "row_pointers"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/XlaSparseDenseMatmulWithCsrInput.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/XlaSparseDenseMatmulWithCsrInput.pbtxt
index 1aa1743718a32a..2b4bc1dcba74ac 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/XlaSparseDenseMatmulWithCsrInput.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/XlaSparseDenseMatmulWithCsrInput.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "XlaSparseDenseMatmulWithCsrInput"
   input_arg {
     name: "row_pointers"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/XlaSparseDenseMatmulWithStaticBufferSize.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/XlaSparseDenseMatmulWithStaticBufferSize.pbtxt
index 85888d026c595c..471ded1635244a 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/XlaSparseDenseMatmulWithStaticBufferSize.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/XlaSparseDenseMatmulWithStaticBufferSize.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "XlaSparseDenseMatmulWithStaticBufferSize"
   input_arg {
     name: "row_pointers"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/XlaSplitND.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/XlaSplitND.pbtxt
index d4d5391b1340df..353b6e166d0f2e 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/XlaSplitND.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/XlaSplitND.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "XlaSplitND"
   input_arg {
     name: "input"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/Xlog1py.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/Xlog1py.pbtxt
index f37a09eea17228..9fe466c2fb8165 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/Xlog1py.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/Xlog1py.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "Xlog1py"
   input_arg {
     name: "x"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/Xlogy.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/Xlogy.pbtxt
index 8da356ec49fe8c..8e7df823ff73dc 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/Xlogy.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/Xlogy.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "Xlogy"
   input_arg {
     name: "x"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ZerosLike.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ZerosLike.pbtxt
index c1200273821a51..5bb8d0ab3781d5 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ZerosLike.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ZerosLike.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ZerosLike"
   input_arg {
     name: "x"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/Zeta.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/Zeta.pbtxt
index 7f86bd29b404cd..c391bd1f22c6a1 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/Zeta.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/Zeta.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "Zeta"
   input_arg {
     name: "x"
diff --git a/tensorflow/core/ops/compat/ops_history_v2/ZipDataset.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/ZipDataset.pbtxt
index 60efabb2d58f62..5dd34535a666d3 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/ZipDataset.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/ZipDataset.pbtxt
@@ -1,4 +1,4 @@
-op 	 {
+op {
   name: "ZipDataset"
   input_arg {
     name: "input_datasets"
diff --git a/tensorflow/core/ops/ops.pbtxt b/tensorflow/core/ops/ops.pbtxt
index c8c55e08ca0d32..c0a6aa926c51e0 100644
--- a/tensorflow/core/ops/ops.pbtxt
+++ b/tensorflow/core/ops/ops.pbtxt
@@ -1,4 +1,4 @@
-go/debugonly    
+go/debugproto  
 op {
   name: "Abort"
   attr {
@@ -60061,6 +60061,8 @@ op {
         type: DT_DOUBLE
         type: DT_INT32
         type: DT_INT64
+        type: DT_UINT32
+        type: DT_UINT64
       }
     }
   }

From 52911fba4c7e2852dc36a601fc3918263f552e2c Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Tue, 14 May 2024 11:44:50 -0700
Subject: [PATCH 125/478] Update visibility in
 third_party/tensorflow/python/platform/BUILD.

PiperOrigin-RevId: 633657239
---
 tensorflow/python/platform/BUILD | 1 +
 1 file changed, 1 insertion(+)

diff --git a/tensorflow/python/platform/BUILD b/tensorflow/python/platform/BUILD
index df26d8b84ec5a3..0ca7e7bfae738f 100644
--- a/tensorflow/python/platform/BUILD
+++ b/tensorflow/python/platform/BUILD
@@ -15,6 +15,7 @@ visibility = [
     # copybara:uncomment "//learning/brain/contrib/eager/numlib/benchmarks/kumamon:__subpackages__",
     # copybara:uncomment "//learning/brain/mobile/lite/tooling/model_analyzer:__subpackages__",
     # copybara:uncomment "//tensorflow_serving/model_servers:__subpackages__",
+    # copybara:uncomment "//third_party/odml/model_customization/quantization:__subpackages__",
 ]
 
 package(

From f6a2b933193c44bf382e47e099122deccc4d2f8c Mon Sep 17 00:00:00 2001
From: "Dimitar (Mitko) Asenov" <dasenov@google.com>
Date: Tue, 14 May 2024 11:54:14 -0700
Subject: [PATCH 126/478] [XLA:GPU] Remove the deprecated xla flag
 `--xla_gpu_simplify_all_fp_conversions`.

Use `--xla_allow_excess_precision` instead.

PiperOrigin-RevId: 633659944
---
 third_party/xla/xla/debug_options_flags.cc      | 6 ------
 third_party/xla/xla/service/gpu/gpu_compiler.cc | 6 ++----
 third_party/xla/xla/xla.proto                   | 3 +--
 3 files changed, 3 insertions(+), 12 deletions(-)

diff --git a/third_party/xla/xla/debug_options_flags.cc b/third_party/xla/xla/debug_options_flags.cc
index aa95aa8a3f6120..aef3c9ef09de2c 100644
--- a/third_party/xla/xla/debug_options_flags.cc
+++ b/third_party/xla/xla/debug_options_flags.cc
@@ -147,7 +147,6 @@ DebugOptions DefaultDebugOptionsIgnoringFlags() {
   opts.set_xla_gpu_redzone_padding_bytes(8 * 1024 * 1024);
   opts.set_xla_gpu_shape_checks(DebugOptions::RUNTIME);
   opts.set_xla_gpu_normalize_layouts(true);
-  opts.set_xla_gpu_simplify_all_fp_conversions(true);
   opts.set_xla_dump_latency_hiding_schedule(false);
   opts.set_xla_gpu_enable_latency_hiding_scheduler(false);
   opts.set_xla_gpu_lhs_enable_gpu_async_tracker(true);
@@ -1229,11 +1228,6 @@ void MakeDebugOptionsFlags(std::vector<tsl::Flag>* flag_list,
       "Amount of padding the redzone allocator will put on one side of each "
       "buffer it allocates. (So the buffer's total size will be increased by "
       "2x this value.)"));
-  flag_list->push_back(tsl::Flag(
-      "xla_gpu_simplify_all_fp_conversions",
-      bool_setter_for(&DebugOptions::set_xla_gpu_simplify_all_fp_conversions),
-      debug_options->xla_gpu_simplify_all_fp_conversions(),
-      "Allows any chain of floating-point conversions to be simplified."));
   flag_list->push_back(tsl::Flag(
       "xla_gpu_shape_checks", setter_for_xla_gpu_shape_checks,
       DebugOptions::ShapeChecks_Name(debug_options->xla_gpu_shape_checks()),
diff --git a/third_party/xla/xla/service/gpu/gpu_compiler.cc b/third_party/xla/xla/service/gpu/gpu_compiler.cc
index 6d293297d2ffad..542c2ae0524d6f 100644
--- a/third_party/xla/xla/service/gpu/gpu_compiler.cc
+++ b/third_party/xla/xla/service/gpu/gpu_compiler.cc
@@ -1324,8 +1324,7 @@ absl::Status GpuCompiler::OptimizeHloPostLayoutAssignment(
     sub_pipeline.AddPass<FloatNormalization>(&f8e5m2fnuz_support);
     sub_pipeline.AddPass<FloatNormalization>(&f8e4m3fnuz_support);
     // Remove `f32 -> bf16 -> f32` casts inserted by bf16 normalization.
-    if (debug_options.xla_allow_excess_precision() &&
-        debug_options.xla_gpu_simplify_all_fp_conversions()) {
+    if (debug_options.xla_allow_excess_precision()) {
       sub_pipeline.AddPass<SimplifyFPConversions>();
     }
   };
@@ -1466,8 +1465,7 @@ absl::Status GpuCompiler::OptimizeHloPostLayoutAssignment(
   pipeline.AddPass<HloPassFix<GpuAlgebraicSimplifier>>(simplifier_options,
                                                        gpu_version);
 
-  if (debug_options.xla_allow_excess_precision() &&
-      debug_options.xla_gpu_simplify_all_fp_conversions()) {
+  if (debug_options.xla_allow_excess_precision()) {
     // This pass cleans up chains of compiler-generated converts
     // (i.e. f32 -> bf16 -> f32) that have been produced by the algebraic
     // simplifier by rearranging ops (i.e. by pushing broadcasts towards the
diff --git a/third_party/xla/xla/xla.proto b/third_party/xla/xla/xla.proto
index 56ddedc2d6416b..4e0abf59b50906 100644
--- a/third_party/xla/xla/xla.proto
+++ b/third_party/xla/xla/xla.proto
@@ -519,8 +519,7 @@ message DebugOptions {
   // scratch), so this can be multiplied by quite a lot.
   int64 xla_gpu_redzone_padding_bytes = 228;
 
-  // Deprecated. Use xla_allow_excess_precision instead.
-  bool xla_gpu_simplify_all_fp_conversions = 168 [deprecated = true];
+  reserved 168;  // Was xla_allow_excess_precision.
 
   // An experimental option to force all layouts present in the
   // after-optimizations HLO to be descending, e.g.

From e3e649bf7b7e48ca39fa49065e7b7db3b90a3f8a Mon Sep 17 00:00:00 2001
From: Deqiang Chen <deqiangc@google.com>
Date: Tue, 14 May 2024 12:01:51 -0700
Subject: [PATCH 127/478] Reduce messages from sharding_utils.cc

PiperOrigin-RevId: 633662509
---
 tensorflow/core/tfrt/ifrt/sharding_utils.cc | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/tensorflow/core/tfrt/ifrt/sharding_utils.cc b/tensorflow/core/tfrt/ifrt/sharding_utils.cc
index 03cc5f953edb38..13e0c2941f4201 100644
--- a/tensorflow/core/tfrt/ifrt/sharding_utils.cc
+++ b/tensorflow/core/tfrt/ifrt/sharding_utils.cc
@@ -174,8 +174,8 @@ SplitAndCreateArraysFromHostBuffer(
                   kImmutableUntilTransferCompletes,
               [tensor, slice_idx]() {
                 // Keep tensor alive
-                LOG(INFO) << "Done with host buffer for slice " << slice_idx
-                          << " at " << tensor.data();
+                VLOG(2) << "Done with host buffer for slice " << slice_idx
+                        << " at " << tensor.data();
               }));
       arrays.push_back(std::move(array));
       device_iter++;

From b4a18a16059f0940bb9f57b520dff678aacf19ac Mon Sep 17 00:00:00 2001
From: Rohit Upadhyaya <rohitju@google.com>
Date: Tue, 14 May 2024 12:17:34 -0700
Subject: [PATCH 128/478] Cleanup change: tf::Status is an alias for
 absl::Status. Add tests

PiperOrigin-RevId: 633667291
---
 tensorflow/core/tfrt/utils/BUILD              |  3 ++-
 tensorflow/core/tfrt/utils/error_util.cc      |  6 ------
 tensorflow/core/tfrt/utils/error_util.h       |  2 --
 tensorflow/core/tfrt/utils/error_util_test.cc | 10 ++++++++++
 4 files changed, 12 insertions(+), 9 deletions(-)

diff --git a/tensorflow/core/tfrt/utils/BUILD b/tensorflow/core/tfrt/utils/BUILD
index 7517e07928f87b..7c5cf7a46579a5 100644
--- a/tensorflow/core/tfrt/utils/BUILD
+++ b/tensorflow/core/tfrt/utils/BUILD
@@ -155,9 +155,10 @@ tf_cc_test(
     deps = [
         ":error_util",
         "//tensorflow/core/platform:status",
+        "@com_google_absl//absl/status",
         "@com_google_googletest//:gtest_main",
+        "@local_xla//xla/tsl/concurrency:async_value",
         "@tf_runtime//:support",
-        "@tf_runtime//cpp_tests:common",
     ],
 )
 
diff --git a/tensorflow/core/tfrt/utils/error_util.cc b/tensorflow/core/tfrt/utils/error_util.cc
index 1f8fb6a6770570..2530b98f051041 100644
--- a/tensorflow/core/tfrt/utils/error_util.cc
+++ b/tensorflow/core/tfrt/utils/error_util.cc
@@ -47,10 +47,4 @@ tensorflow::Status ToTfStatus(const tfrt::AsyncValue* av) {
   return absl::OkStatus();
 }
 
-absl::Status AbslStatusFromTfStatus(tensorflow::Status status) {
-  if (status.ok()) return absl::OkStatus();
-  return absl::Status(static_cast<absl::StatusCode>(status.code()),
-                      status.message());
-}
-
 }  // namespace tfrt
diff --git a/tensorflow/core/tfrt/utils/error_util.h b/tensorflow/core/tfrt/utils/error_util.h
index e694931f82e2e5..ee7bcd81dd913f 100644
--- a/tensorflow/core/tfrt/utils/error_util.h
+++ b/tensorflow/core/tfrt/utils/error_util.h
@@ -76,8 +76,6 @@ inline llvm::Error MakeStatusError(tensorflow::Status status) {
   return MakeStringError(MakeStatusString(status));
 }
 
-absl::Status AbslStatusFromTfStatus(tensorflow::Status status);
-
 }  // namespace tfrt
 
 #endif  // TENSORFLOW_CORE_TFRT_UTILS_ERROR_UTIL_H_
diff --git a/tensorflow/core/tfrt/utils/error_util_test.cc b/tensorflow/core/tfrt/utils/error_util_test.cc
index 07c65905825a53..06edb63c897af4 100644
--- a/tensorflow/core/tfrt/utils/error_util_test.cc
+++ b/tensorflow/core/tfrt/utils/error_util_test.cc
@@ -16,6 +16,8 @@ limitations under the License.
 
 #include <gmock/gmock.h>
 #include <gtest/gtest.h>
+#include "absl/status/status.h"
+#include "xla/tsl/concurrency/async_value_ref.h"
 #include "tensorflow/core/platform/status.h"
 #include "tfrt/support/error_util.h"  // from @tf_runtime
 
@@ -38,5 +40,13 @@ TEST(ErrorUtilTest, UnsupportedErrorConversion) {
             tfrt::ErrorCode::kUnknown);
 }
 
+TEST(ErrorUtilTest, ToTfStatusError) {
+  auto error_av =
+      tsl::MakeErrorAsyncValueRef(absl::UnauthenticatedError("test_error"));
+  auto status = ToTfStatus(error_av.get());
+  EXPECT_EQ(status.code(), absl::StatusCode::kUnauthenticated);
+  EXPECT_EQ(status.message(), "test_error");
+}
+
 }  // namespace
 }  // namespace tfrt

From d7a56fffe080617331aaf3f1b5d172f62babf0d5 Mon Sep 17 00:00:00 2001
From: Kuy Mainwaring <kuym@google.com>
Date: Tue, 14 May 2024 12:28:42 -0700
Subject: [PATCH 129/478] [XLA:GPU] Clang-tidy fixes for
 xla/service/gpu/runtime/fft_thunk.cc

PiperOrigin-RevId: 633670507
---
 third_party/xla/xla/service/gpu/runtime/BUILD     |  8 +++-----
 .../xla/xla/service/gpu/runtime/fft_thunk.cc      | 15 ++++++++++++++-
 2 files changed, 17 insertions(+), 6 deletions(-)

diff --git a/third_party/xla/xla/service/gpu/runtime/BUILD b/third_party/xla/xla/service/gpu/runtime/BUILD
index 4175990422e7bb..e53035e558ee9d 100644
--- a/third_party/xla/xla/service/gpu/runtime/BUILD
+++ b/third_party/xla/xla/service/gpu/runtime/BUILD
@@ -613,25 +613,23 @@ cc_library(
     srcs = ["fft_thunk.cc"],
     hdrs = ["fft_thunk.h"],
     deps = [
+        ":thunk",
         "//xla:shape_util",
+        "//xla:status_macros",
         "//xla:types",
         "//xla:util",
         "//xla:xla_data_proto_cc",
-        "//xla/hlo/ir:hlo",
         "//xla/service:buffer_assignment",
-        "//xla/service/gpu:buffer_allocations",
-        "//xla/service/gpu/runtime:thunk",
         "//xla/stream_executor",
         "@com_google_absl//absl/base:core_headers",
         "@com_google_absl//absl/container:flat_hash_map",
         "@com_google_absl//absl/status",
         "@com_google_absl//absl/status:statusor",
-        "@com_google_absl//absl/strings",
         "@com_google_absl//absl/strings:str_format",
         "@com_google_absl//absl/synchronization",
         "@com_google_absl//absl/types:span",
         "@local_tsl//tsl/platform:logging",
-        "@local_tsl//tsl/platform:status",
+        "@local_tsl//tsl/platform:statusor",
     ],
 )
 
diff --git a/third_party/xla/xla/service/gpu/runtime/fft_thunk.cc b/third_party/xla/xla/service/gpu/runtime/fft_thunk.cc
index 728c36752aeed5..7d620522146acf 100644
--- a/third_party/xla/xla/service/gpu/runtime/fft_thunk.cc
+++ b/third_party/xla/xla/service/gpu/runtime/fft_thunk.cc
@@ -15,17 +15,30 @@ limitations under the License.
 
 #include "xla/service/gpu/runtime/fft_thunk.h"
 
+#include <cstdint>
+#include <memory>
 #include <string>
 
 #include "absl/status/status.h"
 #include "absl/status/statusor.h"
-#include "absl/strings/str_cat.h"
 #include "absl/strings/str_format.h"
+#include "absl/synchronization/mutex.h"
+#include "absl/types/span.h"
+#include "xla/service/buffer_assignment.h"
+#include "xla/service/gpu/runtime/thunk.h"
+#include "xla/shape.h"
+#include "xla/shape_util.h"
+#include "xla/status_macros.h"
+#include "xla/stream_executor/blas.h"
+#include "xla/stream_executor/device_memory.h"
+#include "xla/stream_executor/device_memory_allocator.h"
+#include "xla/stream_executor/fft.h"
 #include "xla/stream_executor/scratch_allocator.h"
 #include "xla/stream_executor/stream_executor.h"
 #include "xla/types.h"
 #include "xla/util.h"
 #include "tsl/platform/logging.h"
+#include "tsl/platform/statusor.h"
 
 namespace xla {
 namespace gpu {

From 3d1bd4049fc906071088c1d560b75ea77e898a62 Mon Sep 17 00:00:00 2001
From: Gunhyun Park <gunhyun@google.com>
Date: Tue, 14 May 2024 12:31:06 -0700
Subject: [PATCH 130/478] Add unbounded dynamism test for AllToAllOp.

PiperOrigin-RevId: 633671113
---
 third_party/xla/xla/client/xla_builder.cc     |  71 ++++++++++-
 third_party/xla/xla/client/xla_builder.h      |   5 +-
 .../xla/xla/client/xla_builder_test.cc        | 120 ++++++++++++++++++
 .../xla/xla/service/shape_inference.cc        |  30 ++++-
 .../xla/xla/service/shape_inference_test.cc   |  26 ++++
 5 files changed, 241 insertions(+), 11 deletions(-)

diff --git a/third_party/xla/xla/client/xla_builder.cc b/third_party/xla/xla/client/xla_builder.cc
index ecc6b86dfcb575..9d32c48ea9afad 100644
--- a/third_party/xla/xla/client/xla_builder.cc
+++ b/third_party/xla/xla/client/xla_builder.cc
@@ -15,6 +15,7 @@ limitations under the License.
 
 #include "xla/client/xla_builder.h"
 
+#include <algorithm>
 #include <cstddef>
 #include <cstdint>
 #include <functional>
@@ -3789,15 +3790,55 @@ XlaOp XlaBuilder::AllToAllArray(
       return all_to_all;
     }
     DimensionVector sizes;
+    const bool is_unbounded = operand_shape->is_unbounded_dynamic();
+    std::vector<XlaOp> dynamic_sizes;
+    auto GetR1DimensionSizeOrConstant = [&](XlaOp operand,
+                                            int64_t dimension) -> XlaOp {
+      if (operand_shape->is_unbounded_dynamic_dimension(dimension)) {
+        return Reshape(GetDimensionSize(operand, dimension), {1});
+      }
+      return ConstantR1<int32_t>(
+          this, {static_cast<int32_t>(operand_shape->dimensions(dimension))});
+    };
+    XlaOp r1_split_count =
+        ConstantR1<int32_t>(this, {static_cast<int32_t>(split_count)});
     for (int64_t i = 0; i < operand_shape->rank(); ++i) {
       if (i != split_dimension) {
         sizes.push_back(operand_shape->dimensions(i));
+        if (is_unbounded) {
+          dynamic_sizes.push_back(GetR1DimensionSizeOrConstant(operand, i));
+        }
         continue;
       }
       sizes.push_back(split_count);
-      sizes.push_back(operand_shape->dimensions(i) / split_count);
+      sizes.push_back(operand_shape->is_unbounded_dynamic_dimension(i)
+                          ? Shape::kUnboundedSize
+                          : operand_shape->dimensions(i) / split_count);
+
+      if (is_unbounded) {
+        dynamic_sizes.push_back(r1_split_count);
+        dynamic_sizes.push_back(
+            operand_shape->is_unbounded_dynamic_dimension(i)
+                ? Div(GetR1DimensionSizeOrConstant(operand, i), r1_split_count)
+                : ConstantR1<int32_t>(this,
+                                      {static_cast<int32_t>(sizes.back())}));
+      }
+    }
+
+    if (is_unbounded) {
+      std::vector<bool> dynamic_dimensions;
+      std::transform(
+          sizes.begin(), sizes.end(), std::back_inserter(dynamic_dimensions),
+          [](int64_t size) { return size == Shape::kUnboundedSize; });
+      TF_ASSIGN_OR_RETURN(
+          const Shape shape,
+          ShapeUtil::MakeValidatedShape(all_to_all_shape.element_type(), sizes,
+                                        dynamic_dimensions));
+      all_to_all =
+          MhloDynamicReshape(all_to_all, ConcatInDim(dynamic_sizes, 0), shape);
+    } else {
+      all_to_all = Reshape(all_to_all, sizes);
     }
-    all_to_all = Reshape(all_to_all, sizes);
 
     std::vector<int64_t> permutation;
     const auto rank = operand_shape->rank();
@@ -3810,6 +3851,21 @@ XlaOp XlaBuilder::AllToAllArray(
       permutation.push_back(dim_after_reshape);
     }
     all_to_all = Transpose(all_to_all, permutation);
+
+    if (is_unbounded) {
+      std::vector<XlaOp> new_dimensions;
+      for (int64_t i = 0; i < operand_shape->rank(); ++i) {
+        new_dimensions.push_back(GetR1DimensionSizeOrConstant(operand, i));
+      }
+      new_dimensions[split_dimension] =
+          Div(new_dimensions[split_dimension], r1_split_count);
+      new_dimensions[concat_dimension] =
+          Mul(new_dimensions[concat_dimension], r1_split_count);
+
+      return MhloDynamicReshape(all_to_all, ConcatInDim(new_dimensions, 0),
+                                all_to_all_shape);
+    }
+
     return Reshape(all_to_all_shape, all_to_all);
   });
 }
@@ -3865,6 +3921,13 @@ XlaOp XlaBuilder::AllToAllTuple(
     const std::optional<ChannelHandle>& channel_id) {
   return ReportErrorOrReturn([&]() -> absl::StatusOr<XlaOp> {
     TF_ASSIGN_OR_RETURN(const Shape* operand_shape, GetShapePtr(operand));
+    if (operand_shape->is_unbounded_dynamic() ||
+        split_dimension == Shape::kUnboundedSize ||
+        concat_dimension == Shape::kUnboundedSize ||
+        split_count == Shape::kUnboundedSize) {
+      return InvalidArgument(
+          "AllToAllTuple does not support unbounded dynamic shapes");
+    }
 
     // The HloInstruction for AllToAll currently only handles the data
     // communication: it accepts N already split parts and scatters them to N
@@ -3890,14 +3953,14 @@ XlaOp XlaBuilder::AllToAllTuple(
     }
 
     // Handle data communication.
-    XlaOp alltoall =
+    XlaOp all_to_all =
         this->AllToAllTuple(slices, replica_groups, layout, channel_id);
 
     // Concat the N received parts.
     std::vector<XlaOp> received;
     received.reserve(split_count);
     for (int i = 0; i < split_count; i++) {
-      received.push_back(this->GetTupleElement(alltoall, i));
+      received.push_back(this->GetTupleElement(all_to_all, i));
     }
     return this->ConcatInDim(received, concat_dimension);
   });
diff --git a/third_party/xla/xla/client/xla_builder.h b/third_party/xla/xla/client/xla_builder.h
index 571904f7cd7995..84fa7588c69016 100644
--- a/third_party/xla/xla/client/xla_builder.h
+++ b/third_party/xla/xla/client/xla_builder.h
@@ -2564,7 +2564,10 @@ XlaOp ReduceScatter(
     const std::optional<Layout>& layout = std::nullopt,
     std::optional<bool> use_global_device_ids = std::nullopt);
 
-// Enqueues an operation that do an Alltoall of the operand cross cores.
+// Enqueues an operation that do an AllToAll of the operand cross cores.
+// This involves AllToAll, followed by Reshape, Transpose, and another Reshape
+// to get proper codegen. See implementation for additional details.
+//
 // An optional `layout` can be specified to force the layout of the instruction.
 // This is used to guarantee the same layout for a group of AllToAll ops
 // compiled separately.
diff --git a/third_party/xla/xla/client/xla_builder_test.cc b/third_party/xla/xla/client/xla_builder_test.cc
index 3aa8c1ca00c5aa..c1c9640aa841f1 100644
--- a/third_party/xla/xla/client/xla_builder_test.cc
+++ b/third_party/xla/xla/client/xla_builder_test.cc
@@ -1999,6 +1999,126 @@ TEST(XlaBuilderTest, UnboundedAllReduce) {
               GmockMatch(m::Op().WithShapeEqualTo(&expected)));
 }
 
+TEST(XlaBuilderTest, UnboundedAllToAllDynamicSplitDimension) {
+  XlaBuilder b(TestName());
+  TF_ASSERT_OK_AND_ASSIGN(const Shape operand, ParseShape("f32[?, 15]"));
+  TF_ASSERT_OK_AND_ASSIGN(const Shape expected, ParseShape("f32[?, 45]"));
+  AllToAll(/*operand=*/Parameter(&b, 0, operand, "operand"),
+           /*split_dimension=*/0,
+           /*concat_dimension=*/1,
+           /*split_count=*/3,
+           /*replica_groups=*/{});
+  TF_ASSERT_OK_AND_ASSIGN(const std::unique_ptr<HloModule> module,
+                          BuildHloModule(b));
+  EXPECT_THAT(GetRoot(*module),
+              GmockMatch(m::Op().WithShapeEqualTo(&expected)));
+}
+
+TEST(XlaBuilderTest, UnboundedAllToAllDynamicConcatDimension) {
+  XlaBuilder b(TestName());
+  TF_ASSERT_OK_AND_ASSIGN(const Shape operand, ParseShape("f32[?, 15]"));
+  TF_ASSERT_OK_AND_ASSIGN(const Shape expected, ParseShape("f32[?, 5]"));
+  AllToAll(/*operand=*/Parameter(&b, 0, operand, "operand"),
+           /*split_dimension=*/1,
+           /*concat_dimension=*/0,
+           /*split_count=*/3,
+           /*replica_groups=*/{});
+  TF_ASSERT_OK_AND_ASSIGN(const std::unique_ptr<HloModule> module,
+                          BuildHloModule(b));
+  EXPECT_THAT(GetRoot(*module),
+              GmockMatch(m::Op().WithShapeEqualTo(&expected)));
+}
+
+TEST(XlaBuilderTest, UnboundedAllToAllDynamicSplitAndConcatDimensionEqual) {
+  XlaBuilder b(TestName());
+  TF_ASSERT_OK_AND_ASSIGN(const Shape operand, ParseShape("f32[?, 15]"));
+  TF_ASSERT_OK_AND_ASSIGN(const Shape expected, ParseShape("f32[?, 15]"));
+  AllToAll(/*operand=*/Parameter(&b, 0, operand, "operand"),
+           /*split_dimension=*/0,
+           /*concat_dimension=*/0,
+           /*split_count=*/3,
+           /*replica_groups=*/{});
+  TF_ASSERT_OK_AND_ASSIGN(const std::unique_ptr<HloModule> module,
+                          BuildHloModule(b));
+  EXPECT_THAT(GetRoot(*module),
+              GmockMatch(m::Op().WithShapeEqualTo(&expected)));
+}
+
+TEST(XlaBuilderTest, UnboundedAllToAllFullyDynamic) {
+  XlaBuilder b(TestName());
+  TF_ASSERT_OK_AND_ASSIGN(const Shape operand, ParseShape("f32[?, ?]"));
+  TF_ASSERT_OK_AND_ASSIGN(const Shape expected, ParseShape("f32[?, ?]"));
+  AllToAll(/*operand=*/Parameter(&b, 0, operand, "operand"),
+           /*split_dimension=*/0,
+           /*concat_dimension=*/1,
+           /*split_count=*/3,
+           /*replica_groups=*/{});
+  TF_ASSERT_OK_AND_ASSIGN(const std::unique_ptr<HloModule> module,
+                          BuildHloModule(b));
+  EXPECT_THAT(GetRoot(*module),
+              GmockMatch(m::Op().WithShapeEqualTo(&expected)));
+}
+
+TEST(XlaBuilderTest, UnboundedAllToAllTupleVariadicUnsupported) {
+  XlaBuilder b(TestName());
+  TF_ASSERT_OK_AND_ASSIGN(const Shape operand, ParseShape("f32[?, 15]{1,0}"));
+  b.ReportErrorOrReturn(
+      AllToAllTuple(/*operands=*/{Parameter(&b, 0, operand, "operand0"),
+                                  Parameter(&b, 1, operand, "operand1")},
+                    /*replica_groups=*/{}));
+  EXPECT_THAT(
+      BuildHloModule(b),
+      StatusIs(_,
+               HasSubstr(
+                   "AllToAllTuple does not support unbounded dynamic shapes")));
+}
+
+TEST(XlaBuilderTest, UnboundedAllToAllTupleUnsupported) {
+  XlaBuilder b(TestName());
+  TF_ASSERT_OK_AND_ASSIGN(const Shape operand, ParseShape("f32[?, 15]{1,0}"));
+  b.ReportErrorOrReturn(
+      AllToAllTuple(/*operand=*/Parameter(&b, 0, operand, "operand"),
+                    /*split_dimension=*/0,
+                    /*concat_dimension=*/1,
+                    /*split_count=*/3,
+                    /*replica_groups=*/{}));
+  EXPECT_THAT(
+      BuildHloModule(b),
+      StatusIs(_,
+               HasSubstr(
+                   "AllToAllTuple does not support unbounded dynamic shapes")));
+}
+
+TEST(XlaBuilderTest, BoundedAllToAllTupleUnsupported) {
+  XlaBuilder b(TestName());
+  TF_ASSERT_OK_AND_ASSIGN(const Shape operand, ParseShape("f32[3, <=15]{1,0}"));
+  b.ReportErrorOrReturn(
+      AllToAllTuple(/*operand=*/Parameter(&b, 0, operand, "operand"),
+                    /*split_dimension=*/0,
+                    /*concat_dimension=*/1,
+                    /*split_count=*/3,
+                    /*replica_groups=*/{}));
+  EXPECT_THAT(
+      BuildHloModule(b),
+      StatusIs(_,
+               HasSubstr("AllToAll does not support bounded dynamic shapes")));
+}
+
+TEST(XlaBuilderTest, BoundedAllToAllUnsupported) {
+  XlaBuilder b(TestName());
+  TF_ASSERT_OK_AND_ASSIGN(const Shape operand, ParseShape("f32[3, <=15]{1,0}"));
+  b.ReportErrorOrReturn(
+      AllToAllTuple(/*operand=*/Parameter(&b, 0, operand, "operand"),
+                    /*split_dimension=*/0,
+                    /*concat_dimension=*/1,
+                    /*split_count=*/3,
+                    /*replica_groups=*/{}));
+  EXPECT_THAT(
+      BuildHloModule(b),
+      StatusIs(_,
+               HasSubstr("AllToAll does not support bounded dynamic shapes")));
+}
+
 TEST(XlaBuilderTest, UnboundedAnd) {
   XlaBuilder b(TestName());
   TF_ASSERT_OK_AND_ASSIGN(const Shape lhs,
diff --git a/third_party/xla/xla/service/shape_inference.cc b/third_party/xla/xla/service/shape_inference.cc
index a6ba2645d35273..10d2f5d3f75b44 100644
--- a/third_party/xla/xla/service/shape_inference.cc
+++ b/third_party/xla/xla/service/shape_inference.cc
@@ -2492,6 +2492,8 @@ ShapeInference::InferScalarBroadcastShape(absl::Span<const Shape> shapes) {
     const Shape& shape, int64_t split_dimension, int64_t concat_dimension,
     int64_t split_count) {
   TF_RET_CHECK(split_count > 0);
+  TF_RET_CHECK(!shape.is_bounded_dynamic())
+      << "AllToAll does not support bounded dynamic shapes";
   if (split_dimension >= shape.rank() || split_dimension < 0) {
     return InvalidArgument(
         "AllToAll split_dimension %d is out-of-bounds in shape %s.",
@@ -2502,25 +2504,41 @@ ShapeInference::InferScalarBroadcastShape(absl::Span<const Shape> shapes) {
         "AllToAll concat_dimension %d is out-of-bounds in shape %s.",
         concat_dimension, ShapeUtil::HumanString(shape));
   }
-  if (shape.dimensions(split_dimension) % split_count != 0) {
+  int64_t split_dimension_size = shape.dimensions(split_dimension);
+  if (!IsUnboundedDynamicSize(split_dimension_size) &&
+      split_dimension_size % split_count != 0) {
     return InvalidArgument(
         "AllToAll split dimension size %d must be dividable by split_count "
         "%d.",
-        shape.dimensions(split_dimension), split_count);
+        split_dimension_size, split_count);
   }
   std::vector<int64_t> new_dimensions(shape.dimensions().begin(),
                                       shape.dimensions().end());
-  new_dimensions[split_dimension] /= split_count;
-  new_dimensions[concat_dimension] *= split_count;
-  return ShapeUtil::MakeShape(shape.element_type(), new_dimensions);
+  new_dimensions[split_dimension] =
+      IsUnboundedDynamicSize(new_dimensions[split_dimension])
+          ? Shape::kUnboundedSize
+          : new_dimensions[split_dimension] / split_count;
+  new_dimensions[concat_dimension] =
+      IsUnboundedDynamicSize(new_dimensions[concat_dimension])
+          ? Shape::kUnboundedSize
+          : new_dimensions[concat_dimension] * split_count;
+
+  const std::vector<bool> dynamic_dimensions(shape.dynamic_dimensions().begin(),
+                                             shape.dynamic_dimensions().end());
+  return ShapeUtil::MakeShape(shape.element_type(), new_dimensions,
+                              dynamic_dimensions);
 }
 
 /* static */ absl::StatusOr<Shape> ShapeInference::InferAllToAllTupleShape(
     absl::Span<const Shape* const> operand_shapes) {
-  // An Alltoall HLO instruction receives N operands (with the same shape) and
+  // An AllToAll HLO instruction receives N operands (with the same shape) and
   // returns a tuple that contains N array shapes.
   TF_RET_CHECK(!operand_shapes.empty());
   for (int i = 0; i < operand_shapes.size(); i++) {
+    if (operand_shapes[i]->is_unbounded_dynamic()) {
+      return InvalidArgument(
+          "AllToAllTuple does not support unbounded dynamic shapes");
+    }
     if (!Shape::Equal().IgnoreMemorySpaceInLayout()(*operand_shapes[0],
                                                     *operand_shapes[i])) {
       return InvalidArgument(
diff --git a/third_party/xla/xla/service/shape_inference_test.cc b/third_party/xla/xla/service/shape_inference_test.cc
index 3afa14fa302eb8..2ec1968a94b7c2 100644
--- a/third_party/xla/xla/service/shape_inference_test.cc
+++ b/third_party/xla/xla/service/shape_inference_test.cc
@@ -4056,6 +4056,32 @@ TEST_F(ShapeInferenceTest, UnboundedAllReduce) {
       << " expected: " << ShapeUtil::HumanString(expected);
 }
 
+TEST_F(ShapeInferenceTest, UnboundedAllToAll) {
+  TF_ASSERT_OK_AND_ASSIGN(const Shape operand, ParseShape("f32[?, 10]"));
+  TF_ASSERT_OK_AND_ASSIGN(const Shape expected, ParseShape("f32[?, 10]"));
+  TF_ASSERT_OK_AND_ASSIGN(
+      const Shape inferred_shape,
+      ShapeInference::InferAllToAllShape(/*shape=*/operand,
+                                         /*split_dimension=*/0,
+                                         /*concat_dimension=*/0,
+                                         /*split_count=*/3));
+  EXPECT_TRUE(ShapeUtil::Equal(inferred_shape, expected))
+      << "inferred: " << ShapeUtil::HumanString(inferred_shape)
+      << " expected: " << ShapeUtil::HumanString(expected);
+}
+
+TEST_F(ShapeInferenceTest, UnboundedAllToAllTupleUnsupported) {
+  TF_ASSERT_OK_AND_ASSIGN(const Shape operand, ParseShape("f32[?, 10]"));
+  TF_ASSERT_OK_AND_ASSIGN(const Shape expected,
+                          ParseShape("(f32[?, 10], f32[?, 10])"));
+  const absl::StatusOr<Shape> inferred_shape =
+      ShapeInference::InferAllToAllTupleShape(
+          /*operand_shapes=*/{&operand, &operand});
+  EXPECT_THAT(
+      inferred_shape.status().message(),
+      HasSubstr("AllToAllTuple does not support unbounded dynamic shapes"));
+}
+
 TEST_P(UnboundedLogicalOpShapeInferenceTest, UnboundedAnd) {
   TF_ASSERT_OK_AND_ASSIGN(const Shape lhs, ParseShape(GetParam().lhs));
   TF_ASSERT_OK_AND_ASSIGN(const Shape rhs, ParseShape(GetParam().rhs));

From ec95a1e4bb77931881f2881f4846ec73212e2c4e Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Tue, 14 May 2024 12:37:00 -0700
Subject: [PATCH 131/478] Create basic server coverage and model tests.

PiperOrigin-RevId: 633672834
---
 tensorflow/compiler/mlir/quantization/stablehlo/python/BUILD | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/tensorflow/compiler/mlir/quantization/stablehlo/python/BUILD b/tensorflow/compiler/mlir/quantization/stablehlo/python/BUILD
index a4bdf7da09f8ee..0999d37da524c2 100644
--- a/tensorflow/compiler/mlir/quantization/stablehlo/python/BUILD
+++ b/tensorflow/compiler/mlir/quantization/stablehlo/python/BUILD
@@ -30,6 +30,7 @@ package(
 pytype_strict_library(
     name = "quantization",
     srcs = ["quantization.py"],
+    visibility = ["//visibility:public"],
     deps = [
         ":pywrap_quantization",
         "//tensorflow/compiler/mlir/quantization/stablehlo:quantization_config_proto_py",
@@ -45,6 +46,10 @@ pytype_strict_library(
 #     testonly = 1,
 #     srcs = ["integration_test/quantize_model_test_base.py"],
 #     tags = ["no_pip"],
+#     visibility = [
+#         "//learning/brain/mlir/quantization/stablehlo:__subpackages__",
+#         "//tensorflow/compiler/mlir/quantization:__subpackages__",
+#     ],
 #     deps = [
 #         "//third_party/py/mlir:ir",
 #         "//third_party/py/mlir:stablehlo_dialect",

From e621b7eb3ce7a806c5427ad68acf0b880871286f Mon Sep 17 00:00:00 2001
From: "Dimitar (Mitko) Asenov" <dasenov@google.com>
Date: Tue, 14 May 2024 12:52:34 -0700
Subject: [PATCH 132/478] [XLA] Fix incorrect comment.

PiperOrigin-RevId: 633677008
---
 third_party/xla/xla/xla.proto | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/third_party/xla/xla/xla.proto b/third_party/xla/xla/xla.proto
index 4e0abf59b50906..391ddd053ffcad 100644
--- a/third_party/xla/xla/xla.proto
+++ b/third_party/xla/xla/xla.proto
@@ -519,7 +519,7 @@ message DebugOptions {
   // scratch), so this can be multiplied by quite a lot.
   int64 xla_gpu_redzone_padding_bytes = 228;
 
-  reserved 168;  // Was xla_allow_excess_precision.
+  reserved 168;  // Was xla_gpu_simplify_all_fp_conversions.
 
   // An experimental option to force all layouts present in the
   // after-optimizations HLO to be descending, e.g.

From 39db3ef2065adeaf6ab48fd400a1bd4817717725 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Tue, 14 May 2024 13:24:17 -0700
Subject: [PATCH 133/478] Update ops-related pbtxt files.

PiperOrigin-RevId: 633686615
---
 tensorflow/core/ops/ops.pbtxt | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorflow/core/ops/ops.pbtxt b/tensorflow/core/ops/ops.pbtxt
index c0a6aa926c51e0..4ac4c1beb21194 100644
--- a/tensorflow/core/ops/ops.pbtxt
+++ b/tensorflow/core/ops/ops.pbtxt
@@ -1,4 +1,4 @@
-go/debugproto  
+go/nodeserialize    
 op {
   name: "Abort"
   attr {

From a1750554dc91ccf3c1585dbc90bcb1b2856855b7 Mon Sep 17 00:00:00 2001
From: Ilia Sergachev <isergachev@nvidia.com>
Date: Tue, 14 May 2024 13:33:54 -0700
Subject: [PATCH 134/478] PR #12465: [GPU] Fix handling of
 xla_gpu_require_complete_aot_autotune_results flag.

Imported from GitHub PR https://github.com/openxla/xla/pull/12465

Copybara import of the project:

--
4a050dce154b924b98134294f3f15ffce35bdcbb by Ilia Sergachev <isergachev@nvidia.com>:

[GPU] Fix handling of xla_gpu_require_complete_aot_autotune_results flag.

Merging this change closes #12465

PiperOrigin-RevId: 633690035
---
 third_party/xla/xla/debug_options_flags.cc | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/third_party/xla/xla/debug_options_flags.cc b/third_party/xla/xla/debug_options_flags.cc
index aef3c9ef09de2c..8b782dbbe5dc39 100644
--- a/third_party/xla/xla/debug_options_flags.cc
+++ b/third_party/xla/xla/debug_options_flags.cc
@@ -249,6 +249,8 @@ DebugOptions DefaultDebugOptionsIgnoringFlags() {
 
   opts.set_xla_reduce_window_rewrite_base_length(32);
 
+  opts.set_xla_gpu_require_complete_aot_autotune_results(false);
+
   return opts;
 }
 
@@ -1427,7 +1429,7 @@ void MakeDebugOptionsFlags(std::vector<tsl::Flag>* flag_list,
       "xla_gpu_require_complete_aot_autotune_results",
       bool_setter_for(
           &DebugOptions::set_xla_gpu_require_complete_aot_autotune_results),
-      debug_options->xla_gpu_multi_streamed_windowed_einsum(),
+      debug_options->xla_gpu_require_complete_aot_autotune_results(),
       "Whether to require complete AOT autotuning results."));
   flag_list->push_back(tsl::Flag(
       "xla_gpu_auto_spmd_partitioning_memory_budget_gb",

From 9bd24f556207dec3e7253c59e764c6a71d16fe4d Mon Sep 17 00:00:00 2001
From: David Dunleavy <ddunleavy@google.com>
Date: Tue, 14 May 2024 14:07:59 -0700
Subject: [PATCH 135/478] Remove `gpu_any` tag in favor of
 `requires-gpu-nvidia`

This also enables v100 tests.

We now use the same sets of tags for gpu tests between internal and external, modulo the extra tags for local execution of gpu tests.

PiperOrigin-RevId: 633700962
---
 third_party/xla/.kokoro/linux/build.sh                     | 2 +-
 .../tsl/tsl/platform/default/build_config_root.bzl         | 2 +-
 third_party/xla/xla/service/gpu/BUILD                      | 1 +
 third_party/xla/xla/service/gpu/tests/BUILD                | 2 +-
 third_party/xla/xla/tools/multihost_hlo_runner/BUILD       | 7 +++----
 third_party/xla/xla/xla.bzl                                | 2 +-
 6 files changed, 8 insertions(+), 8 deletions(-)

diff --git a/third_party/xla/.kokoro/linux/build.sh b/third_party/xla/.kokoro/linux/build.sh
index 81251bd3f44a6c..f9610594e5646d 100644
--- a/third_party/xla/.kokoro/linux/build.sh
+++ b/third_party/xla/.kokoro/linux/build.sh
@@ -54,7 +54,7 @@ RBE_FLAGS=""
 TARGET_FILTERS="-@local_tsl//tsl/platform:subprocess_test -@local_tsl//tsl/platform/cloud:google_auth_provider_test -@local_tsl//tsl/platform/cloud:oauth_client_test"
 
 if is_linux_gpu_job ; then
-    TAGS_FILTER="$TAGS_FILTER,gpu_any,requires-gpu-nvidia,-no_gpu"
+    TAGS_FILTER="$TAGS_FILTER,requires-gpu-nvidia,-no_gpu"
 
     # We are currently running XLA presubmits on machines with NVIDIA T4 GPUs,
     # which have a compute compatibility of 7.5. Se we filter out all the tests
diff --git a/third_party/xla/third_party/tsl/tsl/platform/default/build_config_root.bzl b/third_party/xla/third_party/tsl/tsl/platform/default/build_config_root.bzl
index a900565143027d..142641b16d2fa3 100644
--- a/third_party/xla/third_party/tsl/tsl/platform/default/build_config_root.bzl
+++ b/third_party/xla/third_party/tsl/tsl/platform/default/build_config_root.bzl
@@ -12,7 +12,7 @@ GPU_TEST_PROPERTIES = {
 }
 
 def tf_gpu_tests_tags():
-    return ["gpu"] + gpu_test_tags()
+    return ["requires-gpu-nvidia", "gpu"] + gpu_test_tags()
 
 # terminology changes: saving tf_cuda_* for compatibility
 def tf_cuda_tests_tags():
diff --git a/third_party/xla/xla/service/gpu/BUILD b/third_party/xla/xla/service/gpu/BUILD
index 8b841bcf9be6ea..3d6c90b8119d96 100644
--- a/third_party/xla/xla/service/gpu/BUILD
+++ b/third_party/xla/xla/service/gpu/BUILD
@@ -5648,6 +5648,7 @@ xla_test(
         "gpu_a100",
     ],
     tags = [
+        "no_oss",  # Needs fix for `ConvertGenerator`
         "nomac",
     ],
     deps = [
diff --git a/third_party/xla/xla/service/gpu/tests/BUILD b/third_party/xla/xla/service/gpu/tests/BUILD
index 1c3dfbf7311549..257e3ddf87973f 100644
--- a/third_party/xla/xla/service/gpu/tests/BUILD
+++ b/third_party/xla/xla/service/gpu/tests/BUILD
@@ -820,7 +820,7 @@ lit_test_suite(
         "//xla/tools/hlo_opt:gpu_specs/p100.txtpb",
         "//xla/tools/hlo_opt:gpu_specs/v100.txtpb",
     ],
-    default_tags = tf_cuda_tests_tags() + ["gpu_any"],
+    default_tags = tf_cuda_tests_tags(),
     tags_override = {
         "element_wise_row_vectorization.hlo": ["no_rocm"],
         "scatter_bf16.hlo": ["no_rocm"],
diff --git a/third_party/xla/xla/tools/multihost_hlo_runner/BUILD b/third_party/xla/xla/tools/multihost_hlo_runner/BUILD
index eb9b29fa9a3364..8014fbe8ed3f0d 100644
--- a/third_party/xla/xla/tools/multihost_hlo_runner/BUILD
+++ b/third_party/xla/xla/tools/multihost_hlo_runner/BUILD
@@ -1,6 +1,7 @@
 load("@bazel_skylib//rules:build_test.bzl", "build_test")
 load("@local_config_cuda//cuda:build_defs.bzl", "if_cuda")
 load("@local_config_rocm//rocm:build_defs.bzl", "if_rocm")
+load("@local_tsl//tsl/platform:build_config_root.bzl", "tf_gpu_tests_tags")
 load("@local_tsl//tsl/platform:rules_cc.bzl", "cc_library")
 load("//xla:xla.bzl", "xla_cc_binary")
 load("//xla/tests:build_defs.bzl", "xla_test")
@@ -16,7 +17,7 @@ build_test(
     name = "hlo_runner_main_build_test",
     tags = [
         "cpu",
-        "gpu_any",
+        "gpu",  # TODO(ddunleavy): this configuration of tags will only build on GPU in OSS.
     ],
     targets = [
         ":hlo_runner_main",
@@ -28,11 +29,9 @@ xla_cc_binary(
     testonly = True,
     srcs = ["hlo_runner_main.cc"],
     tags = [
-        "gpu",
-        "gpu_any",
         "noasan",  # Exceeds linker limit.
         "nomac",
-    ],
+    ] + tf_gpu_tests_tags(),
     deps = [
         ":functional_hlo_runner",
         ":hlo_runner_flags",
diff --git a/third_party/xla/xla/xla.bzl b/third_party/xla/xla/xla.bzl
index 93508a19769711..2143baaf6b5d41 100644
--- a/third_party/xla/xla/xla.bzl
+++ b/third_party/xla/xla/xla.bzl
@@ -73,7 +73,7 @@ def xla_cc_binary(deps = [], copts = tsl_copts(), **kwargs):
 def xla_cc_test(name, deps = [], use_gpu = False, **kwargs):
     # Need to do it this way so that `tf_exec_properties` can read tags.
     _tags = kwargs.get("tags", [])
-    kwargs["tags"] = _tags + tf_gpu_tests_tags() + ["gpu_any"] if use_gpu else _tags
+    kwargs["tags"] = _tags + tf_gpu_tests_tags() if use_gpu else _tags
 
     native.cc_test(
         name = name,

From a6b2dfb52b593d8e765aa1fcc664c802e112771a Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Tue, 14 May 2024 14:18:46 -0700
Subject: [PATCH 136/478] Fix build error for
 third_party/tensorflow/core/tfrt/gpu.

PiperOrigin-RevId: 633704451
---
 tensorflow/core/tfrt/gpu/kernel/BUILD              | 1 +
 tensorflow/core/tfrt/gpu/kernel/gpu_runner_test.cc | 4 +++-
 2 files changed, 4 insertions(+), 1 deletion(-)

diff --git a/tensorflow/core/tfrt/gpu/kernel/BUILD b/tensorflow/core/tfrt/gpu/kernel/BUILD
index 4c9046644c93c2..2b79330c050d82 100644
--- a/tensorflow/core/tfrt/gpu/kernel/BUILD
+++ b/tensorflow/core/tfrt/gpu/kernel/BUILD
@@ -102,6 +102,7 @@ tf_cuda_cc_test(
         "//tensorflow/core/tfrt/common:pjrt_util",
         "//tensorflow/core/tfrt/fallback:fallback_state",
         "@com_google_googletest//:gtest_main",
+        "@local_tsl//tsl/framework:serving_device_selector_policies",
         "@tf_runtime//:hostcontext",
         "@tf_runtime//:tensor",
     ],
diff --git a/tensorflow/core/tfrt/gpu/kernel/gpu_runner_test.cc b/tensorflow/core/tfrt/gpu/kernel/gpu_runner_test.cc
index e2b999cb23ec06..7371f62df1f20b 100644
--- a/tensorflow/core/tfrt/gpu/kernel/gpu_runner_test.cc
+++ b/tensorflow/core/tfrt/gpu/kernel/gpu_runner_test.cc
@@ -15,6 +15,8 @@ limitations under the License.
 #if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
 #include "tensorflow/core/tfrt/gpu/kernel/gpu_runner.h"
 
+#include <memory>
+
 #include <gmock/gmock.h>
 #include <gtest/gtest.h>
 #include "tensorflow/cc/framework/scope.h"
@@ -145,7 +147,7 @@ class GpuRunnerTest : public ::testing::Test {
     exec_ctx_ = std::make_unique<tfrt::ExecutionContext>(std::move(*req_ctx));
 
     // Create a gpu runner.
-    auto policy = std::make_unique<RoundRobinPolicy>();
+    auto policy = std::make_unique<tsl::RoundRobinPolicy>();
     serving_device_selector_ = std::make_unique<GpuServingDeviceSelector>(
         kNumVirtualGpuDevices, std::move(policy));
     gpu_runner_ = std::make_unique<GpuRunner>(serving_device_selector_.get());

From edfcdee26c09216b3ef5e6fadbeb7828f3c7bcc2 Mon Sep 17 00:00:00 2001
From: Kuy Mainwaring <kuym@google.com>
Date: Tue, 14 May 2024 14:29:23 -0700
Subject: [PATCH 137/478] [XLA:GPU] Clang-tidy fixes for
 xla/service/gpu/runtime/sequential_thunk.cc

PiperOrigin-RevId: 633708120
---
 third_party/xla/xla/service/gpu/runtime/sequential_thunk.cc | 2 --
 1 file changed, 2 deletions(-)

diff --git a/third_party/xla/xla/service/gpu/runtime/sequential_thunk.cc b/third_party/xla/xla/service/gpu/runtime/sequential_thunk.cc
index 143ad94a29071d..c58f31b5207a7a 100644
--- a/third_party/xla/xla/service/gpu/runtime/sequential_thunk.cc
+++ b/third_party/xla/xla/service/gpu/runtime/sequential_thunk.cc
@@ -28,8 +28,6 @@ limitations under the License.
 namespace xla {
 namespace gpu {
 
-using ::tsl::profiler::ScopedAnnotation;
-
 SequentialThunk::SequentialThunk(ThunkInfo thunk_info, ThunkSequence thunks)
     : Thunk(Kind::kSequential, thunk_info), thunks_(std::move(thunks)) {}
 

From f8e34052d7c6a44cb6fdb593f1d16e4d6db30610 Mon Sep 17 00:00:00 2001
From: Bryan Massoth <bmassoth@google.com>
Date: Tue, 14 May 2024 15:00:17 -0700
Subject: [PATCH 138/478] Fix NPE bug in XContextStatsAccessor::GetStat() when
 stats_metadata_ is nullptr. Refactor preprocess_xplane to enforce nullability
 as part of type.

PiperOrigin-RevId: 633717014
---
 .../third_party/tsl/tsl/profiler/utils/BUILD  |   1 +
 .../tsl/profiler/utils/preprocess_xplane.cc   |  18 +-
 .../tsl/profiler/utils/preprocess_xplane.h    | 218 +++++++++---------
 .../profiler/utils/preprocess_xplane_test.cc  |  15 ++
 4 files changed, 137 insertions(+), 115 deletions(-)

diff --git a/third_party/xla/third_party/tsl/tsl/profiler/utils/BUILD b/third_party/xla/third_party/tsl/tsl/profiler/utils/BUILD
index f0145697048df9..26358c8ec80db5 100644
--- a/third_party/xla/third_party/tsl/tsl/profiler/utils/BUILD
+++ b/third_party/xla/third_party/tsl/tsl/profiler/utils/BUILD
@@ -446,6 +446,7 @@ tsl_cc_test(
         "//tsl/platform:test",
         "//tsl/platform:test_main",
         "//tsl/profiler/lib:connected_traceme",
+        "//tsl/profiler/protobuf:xplane_proto_cc",
         "@com_google_absl//absl/container:flat_hash_map",
         "@com_google_absl//absl/hash",
     ],
diff --git a/third_party/xla/third_party/tsl/tsl/profiler/utils/preprocess_xplane.cc b/third_party/xla/third_party/tsl/tsl/profiler/utils/preprocess_xplane.cc
index 7e8e7f3a431c97..3d06a05609a118 100644
--- a/third_party/xla/third_party/tsl/tsl/profiler/utils/preprocess_xplane.cc
+++ b/third_party/xla/third_party/tsl/tsl/profiler/utils/preprocess_xplane.cc
@@ -38,16 +38,16 @@ using ::tsl::profiler::XPlane;
 using ::tsl::profiler::XPlaneBuilder;
 using ::tsl::profiler::XSpace;
 
-void MutateXPlane(XPlane* plane,
+void MutateXPlane(XPlane& plane,
                   const std::vector<std::unique_ptr<XplaneEventMutatorFactory>>&
                       mutator_factories) {
-  XPlaneBuilder plane_builder(plane);
+  XPlaneBuilder plane_builder(&plane);
 
   absl::flat_hash_map<int64_t, std::vector<std::unique_ptr<XplaneEventMutator>>>
       mutators_from_event_metadata_id;
   std::vector<std::unique_ptr<XplaneEventMutator>> line_mutators;
   for (const auto& mutator_factory : mutator_factories) {
-    auto mutators = mutator_factory->CreateMutators(&plane_builder);
+    auto mutators = mutator_factory->CreateMutators(plane_builder);
     for (auto& mutator : mutators) {
       if (mutator->event_metadata()) {
         auto id = mutator->event_metadata()->id();
@@ -63,7 +63,7 @@ void MutateXPlane(XPlane* plane,
 
   plane_builder.ForEachLine([&](XLineBuilder line_builder) {
     for (const auto& mutator : line_mutators) {
-      mutator->MutateEventsInLine(&line_builder);
+      mutator->MutateEventsInLine(line_builder);
     }
     if (mutators_from_event_metadata_id.empty()) return;
     line_builder.ForEachEvent([&](XEventBuilder event_builder) {
@@ -71,7 +71,7 @@ void MutateXPlane(XPlane* plane,
           mutators_from_event_metadata_id.find(event_builder.MetadataId());
       if (event_mutators != mutators_from_event_metadata_id.end()) {
         for (const auto& mutator : event_mutators->second) {
-          mutator->Mutate(&event_builder);
+          mutator->Mutate(event_builder);
         }
       }
     });
@@ -150,14 +150,18 @@ CreateMutatorFactories() {
 }  // namespace
 
 void PreprocessXPlane(XPlane* plane) {
+  if (plane == nullptr) return;
+
   auto mutator_factories = CreateMutatorFactories();
-  MutateXPlane(plane, mutator_factories);
+  MutateXPlane(*plane, mutator_factories);
 }
 
 void PreprocessXSpace(XSpace* space) {
+  if (space == nullptr) return;
+
   auto mutator_factories = CreateMutatorFactories();
   for (XPlane& plane : *space->mutable_planes()) {
-    MutateXPlane(&plane, mutator_factories);
+    MutateXPlane(plane, mutator_factories);
   }
 }
 
diff --git a/third_party/xla/third_party/tsl/tsl/profiler/utils/preprocess_xplane.h b/third_party/xla/third_party/tsl/tsl/profiler/utils/preprocess_xplane.h
index 2433cd825cc842..724abb30429968 100644
--- a/third_party/xla/third_party/tsl/tsl/profiler/utils/preprocess_xplane.h
+++ b/third_party/xla/third_party/tsl/tsl/profiler/utils/preprocess_xplane.h
@@ -52,9 +52,9 @@ class XplaneEventMutator {
   virtual ~XplaneEventMutator() = default;
 
   // Mutate event by event specified by the event_metadata.
-  virtual void Mutate(XEventBuilder* builder) = 0;
+  virtual void Mutate(XEventBuilder& builder) = 0;
   // Mutate line by line if event_metadata() return nullptr.
-  virtual void MutateEventsInLine(XLineBuilder* line) = 0;
+  virtual void MutateEventsInLine(XLineBuilder& line) = 0;
 
   const XEventMetadata* event_metadata() const { return event_metadata_; }
 
@@ -70,7 +70,7 @@ class XplaneEventMutatorFactory {
   virtual ~XplaneEventMutatorFactory() = default;
 
   virtual std::vector<std::unique_ptr<XplaneEventMutator>> CreateMutators(
-      XPlaneBuilder* xplane) const = 0;
+      XPlaneBuilder& xplane) const = 0;
 
  protected:
   XplaneEventMutatorFactory() = default;
@@ -84,21 +84,20 @@ class XplaneRootEventMutatorFactory : public XplaneEventMutatorFactory {
  public:
   static std::unique_ptr<XplaneEventMutatorFactory> CreateFactory(
       HostEventType event_type, int64_t root_level) {
-    std::unique_ptr<XplaneEventMutatorFactory> base;
-    base.reset(new XplaneRootEventMutatorFactory(event_type, root_level));
-    return base;
+    return absl::WrapUnique(
+        new XplaneRootEventMutatorFactory(event_type, root_level));
   }
 
   std::vector<std::unique_ptr<XplaneEventMutator>> CreateMutators(
-      XPlaneBuilder* xplane) const override {
+      XPlaneBuilder& xplane) const override {
     std::vector<std::unique_ptr<XplaneEventMutator>> mutators;
-    XEventMetadata* event_metadata =
-        xplane->GetEventMetadata(GetHostEventTypeStr(event_type_));
-    if (event_metadata == nullptr) return {};
-    XStatMetadata* root_metadata =
-        xplane->GetOrCreateStatMetadata(GetStatTypeStr(StatType::kIsRoot));
-    mutators.emplace_back(std::make_unique<XplaneRootEventMutator>(
-        event_metadata, root_metadata, root_level_));
+    if (auto* event_metadata =
+            xplane.GetEventMetadata(GetHostEventTypeStr(event_type_))) {
+      XStatMetadata* root_metadata =
+          xplane.GetOrCreateStatMetadata(GetStatTypeStr(StatType::kIsRoot));
+      mutators.emplace_back(std::make_unique<XplaneRootEventMutator>(
+          event_metadata, *root_metadata, root_level_));
+    }
     return mutators;
   }
 
@@ -110,20 +109,20 @@ class XplaneRootEventMutatorFactory : public XplaneEventMutatorFactory {
   class XplaneRootEventMutator : public XplaneEventMutator {
    public:
     XplaneRootEventMutator(XEventMetadata* event_metadata,
-                           XStatMetadata* root_stats_metadata,
+                           XStatMetadata& root_stats_metadata,
                            int64_t root_level)
         : XplaneEventMutator(event_metadata),
           root_stats_metadata_(root_stats_metadata),
           root_level_(root_level) {}
-    void Mutate(XEventBuilder* event_builder) override {
-      event_builder->SetOrAddStatValue(*root_stats_metadata_, root_level_);
+    void Mutate(XEventBuilder& event_builder) override {
+      event_builder.SetOrAddStatValue(root_stats_metadata_, root_level_);
     }
-    void MutateEventsInLine(XLineBuilder* line) override {
+    void MutateEventsInLine(XLineBuilder& line) override {
       CHECK(false);  // Crash OK
     }
 
    private:
-    XStatMetadata* root_stats_metadata_;
+    XStatMetadata& root_stats_metadata_;
     int64_t root_level_;
   };
 
@@ -136,18 +135,19 @@ class XContextStatsAccessor {
  public:
   using value_type = StatValueType;
 
-  bool Initialize(XPlaneBuilder* xplane) {
-    stats_metadata_ = xplane->GetStatMetadata(GetStatTypeStr(kStatId));
+  bool Initialize(XPlaneBuilder& xplane) {
+    stats_metadata_ = xplane.GetStatMetadata(GetStatTypeStr(kStatId));
     return stats_metadata_;
   }
 
-  std::optional<StatValueType> GetStat(XEventBuilder* event_builder) {
-    auto* stat = event_builder->GetStat(*stats_metadata_);
+  std::optional<StatValueType> GetStat(XEventBuilder& event_builder) {
+    if (stats_metadata_ == nullptr) return std::nullopt;
+    auto* stat = event_builder.GetStat(*stats_metadata_);
     if (stat == nullptr) return std::nullopt;
     if constexpr (std::is_integral_v<StatValueType>) {
-      return event_builder->IntOrUintValue(*stat);
+      return event_builder.IntOrUintValue(*stat);
     } else {
-      return event_builder->StrOrRefValue(*stat);
+      return event_builder.StrOrRefValue(*stat);
     }
   }
 
@@ -160,19 +160,19 @@ class XContextStatsAccessorWithDefault {
  public:
   using value_type = StatValueType;
 
-  bool Initialize(XPlaneBuilder* xplane) {
-    stats_metadata_ = xplane->GetStatMetadata(GetStatTypeStr(kStatId));
+  bool Initialize(XPlaneBuilder& xplane) {
+    stats_metadata_ = xplane.GetStatMetadata(GetStatTypeStr(kStatId));
     return true;  // Always return true, even stat_metadata doesn't exist.
   }
 
-  std::optional<StatValueType> GetStat(XEventBuilder* event_builder) {
+  std::optional<StatValueType> GetStat(XEventBuilder& event_builder) {
     if (stats_metadata_ == nullptr) return kDefaultValue;
-    auto* stat = event_builder->GetStat(*stats_metadata_);
+    auto* stat = event_builder.GetStat(*stats_metadata_);
     if (stat == nullptr) return kDefaultValue;
     if constexpr (std::is_integral_v<StatValueType>) {
-      return event_builder->IntOrUintValue(*stat);
+      return event_builder.IntOrUintValue(*stat);
     } else {
-      return event_builder->StrOrRefValue(*stat);
+      return event_builder.StrOrRefValue(*stat);
     }
   }
 
@@ -222,40 +222,40 @@ class XplaneConnectedEventMutatorFactory : public XplaneEventMutatorFactory {
   using StatsAccessors = std::tuple<StatsAccessorTypes...>;
 
   std::vector<std::unique_ptr<XplaneEventMutator>> CreateMutators(
-      XPlaneBuilder* xplane) const override {
+      XPlaneBuilder& xplane) const override {
     // Check if all stats exist in current plane.
     StatsAccessors stats_accessors;
     bool all_required_stats_exist = true;
     auto check_stats_meta = [&all_required_stats_exist,
-                             xplane](auto&& accessor) {
-      if (all_required_stats_exist == false) return;
-      if (!accessor.Initialize(xplane)) all_required_stats_exist = false;
+                             &xplane](auto&& accessor) {
+      all_required_stats_exist =
+          all_required_stats_exist && accessor.Initialize(xplane);
     };
     for_each(stats_accessors, check_stats_meta);
     if (!all_required_stats_exist) return {};
 
     XEventMetadata* producer_event_metadata =
-        xplane->GetEventMetadata(GetHostEventTypeStr(producer_event));
+        xplane.GetEventMetadata(GetHostEventTypeStr(producer_event));
     XEventMetadata* consumer_event_metadata =
-        xplane->GetEventMetadata(GetHostEventTypeStr(consumer_event));
+        xplane.GetEventMetadata(GetHostEventTypeStr(consumer_event));
 
     std::vector<std::unique_ptr<XplaneEventMutator>> mutators;
     if (producer_event_metadata) {
-      XStatMetadata* context_type_metadata = xplane->GetOrCreateStatMetadata(
+      XStatMetadata* context_type_metadata = xplane.GetOrCreateStatMetadata(
           GetStatTypeStr(StatType::kProducerType));
-      XStatMetadata* context_id_metadata = xplane->GetOrCreateStatMetadata(
-          GetStatTypeStr(StatType::kProducerId));
+      XStatMetadata* context_id_metadata =
+          xplane.GetOrCreateStatMetadata(GetStatTypeStr(StatType::kProducerId));
       mutators.emplace_back(std::make_unique<XplaneConnectedEventMutator>(
-          producer_event_metadata, context_type_metadata, context_id_metadata,
+          producer_event_metadata, *context_type_metadata, *context_id_metadata,
           stats_accessors));
     }
     if (consumer_event_metadata) {
-      XStatMetadata* context_type_metadata = xplane->GetOrCreateStatMetadata(
+      XStatMetadata* context_type_metadata = xplane.GetOrCreateStatMetadata(
           GetStatTypeStr(StatType::kConsumerType));
-      XStatMetadata* context_id_metadata = xplane->GetOrCreateStatMetadata(
-          GetStatTypeStr(StatType::kConsumerId));
+      XStatMetadata* context_id_metadata =
+          xplane.GetOrCreateStatMetadata(GetStatTypeStr(StatType::kConsumerId));
       mutators.emplace_back(std::make_unique<XplaneConnectedEventMutator>(
-          consumer_event_metadata, context_type_metadata, context_id_metadata,
+          consumer_event_metadata, *context_type_metadata, *context_id_metadata,
           stats_accessors));
     }
     return mutators;
@@ -267,19 +267,19 @@ class XplaneConnectedEventMutatorFactory : public XplaneEventMutatorFactory {
   class XplaneConnectedEventMutator : public XplaneEventMutator {
    public:
     XplaneConnectedEventMutator(XEventMetadata* event_metadata,
-                                XStatMetadata* context_type_metadata,
-                                XStatMetadata* context_id_metadata,
+                                XStatMetadata& context_type_metadata,
+                                XStatMetadata& context_id_metadata,
                                 const StatsAccessors& accessors)
         : XplaneEventMutator(event_metadata),
           context_type_metadata_(context_type_metadata),
           context_id_metadata_(context_id_metadata),
           accessors_(accessors) {}
 
-    void Mutate(XEventBuilder* event_builder) override {
+    void Mutate(XEventBuilder& event_builder) override {
       bool all_required_stats_exist = true;
       std::vector<std::variant<absl::string_view, uint64_t>> required_stats;
       auto check_stats_meta = [&all_required_stats_exist, &required_stats,
-                               event_builder](auto&& accessor) {
+                               &event_builder](auto&& accessor) {
         if (all_required_stats_exist == false) return;
         auto stats_data = accessor.GetStat(event_builder);
         if (!stats_data) {
@@ -299,18 +299,18 @@ class XplaneConnectedEventMutatorFactory : public XplaneEventMutatorFactory {
         context_id =
             absl::HashOf(producer_event, consumer_event, required_stats);
       }
-      event_builder->SetOrAddStatValue(*context_type_metadata_,
-                                       static_cast<int64_t>(context_type));
-      event_builder->SetOrAddStatValue(*context_id_metadata_, context_id);
+      event_builder.SetOrAddStatValue(context_type_metadata_,
+                                      static_cast<int64_t>(context_type));
+      event_builder.SetOrAddStatValue(context_id_metadata_, context_id);
     }
 
-    void MutateEventsInLine(XLineBuilder* line) override {
+    void MutateEventsInLine(XLineBuilder& line) override {
       CHECK(false);  // Crash OK
     }
 
    private:
-    XStatMetadata* context_type_metadata_;
-    XStatMetadata* context_id_metadata_;
+    XStatMetadata& context_type_metadata_;
+    XStatMetadata& context_id_metadata_;
     StatsAccessors accessors_;
   };
 };
@@ -323,17 +323,18 @@ class HostRunIdMutatorFactory : public XplaneEventMutatorFactory {
   }
 
   std::vector<std::unique_ptr<XplaneEventMutator>> CreateMutators(
-      XPlaneBuilder* xplane) const override {
+      XPlaneBuilder& xplane) const override {
     std::vector<std::unique_ptr<XplaneEventMutator>> mutators;
-    XEventMetadata* event_metadata =
-        xplane->GetEventMetadata(GetHostEventTypeStr(event_type));
-    if (event_metadata == nullptr) return {};
-    XContextStatsAccessor<int64_t, StatType::kRunId> run_id_stats_accessor;
-    run_id_stats_accessor.Initialize(xplane);
-    XStatMetadata* run_id_metadata =
-        xplane->GetOrCreateStatMetadata(GetStatTypeStr(StatType::kRunId));
-    mutators.emplace_back(std::make_unique<HostRunIdMutator>(
-        event_metadata, run_id_stats_accessor, run_id_metadata));
+    if (auto* event_metadata =
+            xplane.GetEventMetadata(GetHostEventTypeStr(event_type))) {
+      XContextStatsAccessor<int64_t, StatType::kRunId> run_id_stats_accessor;
+      if (run_id_stats_accessor.Initialize(xplane)) {
+        XStatMetadata* run_id_metadata =
+            xplane.GetOrCreateStatMetadata(GetStatTypeStr(StatType::kRunId));
+        mutators.emplace_back(std::make_unique<HostRunIdMutator>(
+            event_metadata, run_id_stats_accessor, *run_id_metadata));
+      }
+    }
     return mutators;
   }
 
@@ -344,25 +345,25 @@ class HostRunIdMutatorFactory : public XplaneEventMutatorFactory {
     HostRunIdMutator(
         XEventMetadata* event_metadata,
         XContextStatsAccessor<int64_t, StatType::kRunId> run_id_stats_accessor,
-        XStatMetadata* run_id_metadata)
+        XStatMetadata& run_id_metadata)
         : XplaneEventMutator(event_metadata),
           run_id_stats_accessor_(run_id_stats_accessor),
           run_id_metadata_(run_id_metadata) {}
 
-    void Mutate(XEventBuilder* event_builder) override {
+    void Mutate(XEventBuilder& event_builder) override {
       auto run_id = run_id_stats_accessor_.GetStat(event_builder);
       if (!run_id) return;
       int64_t fixed_run_id = ((uint64_t)run_id.value() & kRunIdMask);
-      event_builder->SetOrAddStatValue(*run_id_metadata_, fixed_run_id);
+      event_builder.SetOrAddStatValue(run_id_metadata_, fixed_run_id);
     }
 
-    void MutateEventsInLine(XLineBuilder* line) override {
+    void MutateEventsInLine(XLineBuilder& line) override {
       CHECK(false);  // Crash OK
     }
 
    private:
     XContextStatsAccessor<int64_t, StatType::kRunId> run_id_stats_accessor_;
-    XStatMetadata* run_id_metadata_;
+    XStatMetadata& run_id_metadata_;
   };
 };
 
@@ -377,27 +378,28 @@ class TpuModuleLineMutatorFactory : public XplaneEventMutatorFactory {
   }
 
   std::vector<std::unique_ptr<XplaneEventMutator>> CreateMutators(
-      XPlaneBuilder* xplane) const override {
+      XPlaneBuilder& xplane) const override {
     std::vector<std::unique_ptr<XplaneEventMutator>> mutators;
-    if (absl::StartsWith(xplane->Name(), kTpuPlanePrefix) &&
-        GetTensorCoreId(xplane->Name()).has_value()) {
-      if (auto device_ordinal = ParseDeviceOrdinal(xplane->Name())) {
-        XStatMetadata* context_type_metadata = xplane->GetOrCreateStatMetadata(
+    if (absl::StartsWith(xplane.Name(), kTpuPlanePrefix) &&
+        GetTensorCoreId(xplane.Name()).has_value()) {
+      if (auto device_ordinal = ParseDeviceOrdinal(xplane.Name())) {
+        XStatMetadata* context_type_metadata = xplane.GetOrCreateStatMetadata(
             GetStatTypeStr(StatType::kConsumerType));
-        XStatMetadata* context_id_metadata = xplane->GetOrCreateStatMetadata(
+        XStatMetadata* context_id_metadata = xplane.GetOrCreateStatMetadata(
             GetStatTypeStr(StatType::kConsumerId));
         XContextStatsAccessor<uint64_t, StatType::kQueueId>
             queue_id_stats_accessor;
         XContextStatsAccessor<uint64_t, StatType::kRunId> run_id_stats_accessor;
         XContextStatsAccessorWithDefault<uint64_t, StatType::kCoreType, 0ULL>
             core_type_stats_accessor;
-        queue_id_stats_accessor.Initialize(xplane);
-        run_id_stats_accessor.Initialize(xplane);
-        core_type_stats_accessor.Initialize(xplane);
-        mutators.emplace_back(std::make_unique<TpuModuleLineMutator>(
-            *device_ordinal, context_type_metadata, context_id_metadata,
-            queue_id_stats_accessor, run_id_stats_accessor,
-            core_type_stats_accessor));
+        if (queue_id_stats_accessor.Initialize(xplane) &&
+            run_id_stats_accessor.Initialize(xplane) &&
+            core_type_stats_accessor.Initialize(xplane)) {
+          mutators.emplace_back(std::make_unique<TpuModuleLineMutator>(
+              *device_ordinal, *context_type_metadata, *context_id_metadata,
+              queue_id_stats_accessor, run_id_stats_accessor,
+              core_type_stats_accessor));
+        }
       }
     }
     return mutators;
@@ -409,8 +411,8 @@ class TpuModuleLineMutatorFactory : public XplaneEventMutatorFactory {
   class TpuModuleLineMutator : public XplaneEventMutator {
    public:
     TpuModuleLineMutator(
-        uint32_t device_ordinal, XStatMetadata* context_type_metadata,
-        XStatMetadata* context_id_metadata,
+        uint32_t device_ordinal, XStatMetadata& context_type_metadata,
+        XStatMetadata& context_id_metadata,
         XContextStatsAccessor<uint64_t, StatType::kQueueId>
             queue_id_stats_accessor,
         XContextStatsAccessor<uint64_t, StatType::kRunId> run_id_stats_accessor,
@@ -424,16 +426,16 @@ class TpuModuleLineMutatorFactory : public XplaneEventMutatorFactory {
           run_id_stats_accessor_(run_id_stats_accessor),
           core_type_stats_accessor_(core_type_stats_accessor) {}
 
-    void Mutate(XEventBuilder* event_builder) override {
+    void Mutate(XEventBuilder& event_builder) override {
       CHECK(false);  // Crash OK
     }
 
-    void MutateEventsInLine(XLineBuilder* line) override {
-      if (line->Name() != kXlaModuleLineName) return;
-      line->ForEachEvent([&](XEventBuilder event) {
-        auto run_id = run_id_stats_accessor_.GetStat(&event);
-        auto queue_id = queue_id_stats_accessor_.GetStat(&event);
-        auto core_type = core_type_stats_accessor_.GetStat(&event);
+    void MutateEventsInLine(XLineBuilder& line) override {
+      if (line.Name() != kXlaModuleLineName) return;
+      line.ForEachEvent([&](XEventBuilder event) {
+        auto run_id = run_id_stats_accessor_.GetStat(event);
+        auto queue_id = queue_id_stats_accessor_.GetStat(event);
+        auto core_type = core_type_stats_accessor_.GetStat(event);
         if (!run_id || !queue_id) return;
         // The order of tuple <device_ordinal, queue_id, run_id> need to be
         // consistent with other kTpuLaunch types.
@@ -444,16 +446,16 @@ class TpuModuleLineMutatorFactory : public XplaneEventMutatorFactory {
         required_stats.emplace_back(*run_id);
         required_stats.emplace_back(static_cast<uint64_t>(*core_type));
         int64_t context_id = absl::HashOf(required_stats);
-        event.SetOrAddStatValue(*context_type_metadata_,
+        event.SetOrAddStatValue(context_type_metadata_,
                                 static_cast<int64_t>(ContextType::kTpuLaunch));
-        event.SetOrAddStatValue(*context_id_metadata_, context_id);
+        event.SetOrAddStatValue(context_id_metadata_, context_id);
       });
     }
 
    private:
     uint64_t device_ordinal_;
-    XStatMetadata* context_type_metadata_;
-    XStatMetadata* context_id_metadata_;
+    XStatMetadata& context_type_metadata_;
+    XStatMetadata& context_id_metadata_;
     XContextStatsAccessor<uint64_t, StatType::kQueueId>
         queue_id_stats_accessor_;
     XContextStatsAccessor<uint64_t, StatType::kRunId> run_id_stats_accessor_;
@@ -473,7 +475,7 @@ class ThreadpoolLineMutatorFactory : public XplaneEventMutatorFactory {
   }
 
   std::vector<std::unique_ptr<XplaneEventMutator>> CreateMutators(
-      XPlaneBuilder* xplane) const override {
+      XPlaneBuilder& xplane) const override {
     std::vector<std::unique_ptr<XplaneEventMutator>> mutators;
     mutators.emplace_back(std::make_unique<ThreadpoolLineMutator>(xplane));
     return mutators;
@@ -484,25 +486,25 @@ class ThreadpoolLineMutatorFactory : public XplaneEventMutatorFactory {
 
   class ThreadpoolLineMutator : public XplaneEventMutator {
    public:
-    explicit ThreadpoolLineMutator(XPlaneBuilder* xplane)
+    explicit ThreadpoolLineMutator(XPlaneBuilder& xplane)
         : XplaneEventMutator(nullptr), xplane_(xplane) {
       start_region_metadata_ =
-          xplane_->GetEventMetadata(kThreadpoolListenerStartRegion);
+          xplane_.GetEventMetadata(kThreadpoolListenerStartRegion);
       stop_region_metadata_ =
-          xplane_->GetEventMetadata(kThreadpoolListenerStopRegion);
+          xplane_.GetEventMetadata(kThreadpoolListenerStopRegion);
       thread_pool_metadata_ =
-          xplane_->GetOrCreateEventMetadata(kThreadpoolListenerRegion);
-      consumer_ = xplane_->GetOrCreateStatMetadata(
+          xplane_.GetOrCreateEventMetadata(kThreadpoolListenerRegion);
+      consumer_ = xplane_.GetOrCreateStatMetadata(
           GetStatTypeStr(StatType::kConsumerId));
-      consumer_type_ = xplane_->GetOrCreateStatMetadata(
+      consumer_type_ = xplane_.GetOrCreateStatMetadata(
           GetStatTypeStr(StatType::kConsumerType));
     }
 
-    void Mutate(XEventBuilder* event_builder) override {
+    void Mutate(XEventBuilder& event_builder) override {
       CHECK(false);  // Crash OK
     }
 
-    void MutateEventsInLine(XLineBuilder* line) override {
+    void MutateEventsInLine(XLineBuilder& line) override {
       if (start_region_metadata_ == nullptr ||
           stop_region_metadata_ == nullptr) {
         // Skip mutations for xplanes that do not have region markers. These
@@ -519,7 +521,7 @@ class ThreadpoolLineMutatorFactory : public XplaneEventMutatorFactory {
       };
 
       std::vector<EventMetadata> event_metadata;
-      line->ForEachEvent([&](const XEventBuilder& event) {
+      line.ForEachEvent([&](const XEventBuilder& event) {
         if (event.MetadataId() == start_region_metadata_->id()) {
           auto consumer_id = event.GetStat(*consumer_);
           if (!consumer_id) return;
@@ -535,7 +537,7 @@ class ThreadpoolLineMutatorFactory : public XplaneEventMutatorFactory {
         }
       });
       for (const auto& event_metadata : event_metadata) {
-        XEventBuilder region = line->AddEvent(*thread_pool_metadata_);
+        XEventBuilder region = line.AddEvent(*thread_pool_metadata_);
         region.SetTimestampPs(event_metadata.start_region_timestamp_ps);
         region.SetEndTimestampPs(event_metadata.end_region_timestamp_ps);
         region.SetOrAddStatValue(*consumer_, event_metadata.region_id);
@@ -548,7 +550,7 @@ class ThreadpoolLineMutatorFactory : public XplaneEventMutatorFactory {
    private:
     XStatMetadata* consumer_;
     XStatMetadata* consumer_type_;
-    XPlaneBuilder* xplane_;
+    XPlaneBuilder& xplane_;
     XEventMetadata* start_region_metadata_;
     XEventMetadata* stop_region_metadata_;
     XEventMetadata* thread_pool_metadata_;
diff --git a/third_party/xla/third_party/tsl/tsl/profiler/utils/preprocess_xplane_test.cc b/third_party/xla/third_party/tsl/tsl/profiler/utils/preprocess_xplane_test.cc
index 5912c00c397853..9712893645090c 100644
--- a/third_party/xla/third_party/tsl/tsl/profiler/utils/preprocess_xplane_test.cc
+++ b/third_party/xla/third_party/tsl/tsl/profiler/utils/preprocess_xplane_test.cc
@@ -16,12 +16,14 @@ limitations under the License.
 #include "tsl/profiler/utils/preprocess_xplane.h"
 
 #include <cstdint>
+#include <memory>
 #include <optional>
 
 #include "absl/container/flat_hash_map.h"
 #include "absl/hash/hash.h"
 #include "tsl/platform/test.h"
 #include "tsl/profiler/lib/connected_traceme.h"
+#include "tsl/profiler/protobuf/xplane.pb.h"
 #include "tsl/profiler/utils/tf_xplane_visitor.h"
 #include "tsl/profiler/utils/xplane_builder.h"
 #include "tsl/profiler/utils/xplane_schema.h"
@@ -290,6 +292,19 @@ TEST(PreprocessXPlane, ThreadPoolPreprocessorTest) {
   EXPECT_TRUE(new_event_added);
 }
 
+TEST(PreprocessXPlane, XContextStatsAccessorNPETest) {
+  auto xplane = std::make_unique<XPlane>();
+  XPlaneBuilder xplane_builder(xplane.get());
+  XLine xline;
+  XLineBuilder xline_builder(&xline, &xplane_builder);
+  XEvent xevent;
+  XEventBuilder xevent_builder(&xline, &xplane_builder, &xevent);
+  XContextStatsAccessor<int64_t, StatType::kRunId> run_id_accessor;
+
+  ASSERT_FALSE(run_id_accessor.Initialize(xplane_builder));
+  EXPECT_EQ(run_id_accessor.GetStat(xevent_builder), std::nullopt);
+}
+
 }  // namespace
 }  // namespace profiler
 }  // namespace tsl

From c4f3a1d9180968265f779513ca4ac2292799cf3a Mon Sep 17 00:00:00 2001
From: Kyle Lucke <klucke@google.com>
Date: Tue, 14 May 2024 15:18:47 -0700
Subject: [PATCH 139/478] Remove StreamExecutorInterface::AllocateStream in
 favor of making CreateStream do the necessary work.

PiperOrigin-RevId: 633722475
---
 .../stream_executor/stream_executor.cc        | 14 +++----------
 .../xla/xla/backends/interpreter/executor.h   |  2 --
 .../xla/stream_executor/cuda/cuda_executor.cc | 20 +++++++++----------
 .../xla/stream_executor/gpu/gpu_executor.h    |  2 --
 .../xla/xla/stream_executor/gpu/gpu_stream.h  |  3 +--
 .../xla/stream_executor/host/host_executor.cc |  6 +-----
 .../xla/stream_executor/host/host_executor.h  |  1 -
 .../stream_executor/mock_stream_executor.h    |  1 -
 .../xla/stream_executor/rocm/rocm_executor.cc | 20 +++++++++----------
 third_party/xla/xla/stream_executor/stream.cc | 11 ----------
 third_party/xla/xla/stream_executor/stream.h  |  5 -----
 .../stream_executor_interface.h               |  4 ----
 .../xla/stream_executor/tpu/tpu_executor.cc   |  6 ------
 .../xla/stream_executor/tpu/tpu_executor.h    |  2 --
 .../tpu/tpu_executor_init_fns.inc             |  1 -
 15 files changed, 25 insertions(+), 73 deletions(-)

diff --git a/tensorflow/c/experimental/stream_executor/stream_executor.cc b/tensorflow/c/experimental/stream_executor/stream_executor.cc
index c524dd1d27ac89..6e706366d47127 100644
--- a/tensorflow/c/experimental/stream_executor/stream_executor.cc
+++ b/tensorflow/c/experimental/stream_executor/stream_executor.cc
@@ -438,14 +438,6 @@ class CStreamExecutor : public StreamExecutor {
         stream_executor_->get_event_status(&device_, event_handle);
     return SEEventStatusToEventStatus(event_status);
   }
-  bool AllocateStream(Stream* stream) override {
-    DCHECK(stream != nullptr);
-    absl::Status status =
-        static_cast<CStream*>(stream->implementation())->Create();
-    // TODO(annarev): update AllocateStream to return status instead
-    // (similar to AllocateEvent).
-    return status.ok();
-  }
   void DeallocateStream(Stream* stream) override {
     static_cast<CStream*>(stream->implementation())->Destroy();
   }
@@ -568,9 +560,9 @@ class CStreamExecutor : public StreamExecutor {
   absl::StatusOr<std::unique_ptr<Stream>> CreateStream(
       std::optional<std::variant<StreamPriority, int>> priority =
           std::nullopt) override {
-    auto stream = std::make_unique<Stream>(
-        this, std::make_unique<CStream>(&device_, stream_executor_));
-    TF_RETURN_IF_ERROR(stream->Initialize(priority));
+    auto c_stream = std::make_unique<CStream>(&device_, stream_executor_);
+    TF_RETURN_IF_ERROR(c_stream->Create());
+    auto stream = std::make_unique<Stream>(this, std::move(c_stream));
     return std::move(stream);
   }
 
diff --git a/third_party/xla/xla/backends/interpreter/executor.h b/third_party/xla/xla/backends/interpreter/executor.h
index 9f72e99dd0a69d..7ca85d7c22048f 100644
--- a/third_party/xla/xla/backends/interpreter/executor.h
+++ b/third_party/xla/xla/backends/interpreter/executor.h
@@ -132,7 +132,6 @@ class XlaInterpreterExecutor : public StreamExecutor {
     return Event::Status::kError;
   }
 
-  bool AllocateStream(Stream *stream) override { return true; }
   void DeallocateStream(Stream *stream) override {}
   bool CreateStreamDependency(Stream *dependent, Stream *other) override;
 
@@ -167,7 +166,6 @@ class XlaInterpreterExecutor : public StreamExecutor {
           std::nullopt) override {
     auto stream =
         std::make_unique<Stream>(this, std::make_unique<host::HostStream>());
-    TF_RETURN_IF_ERROR(stream->Initialize(priority));
     return std::move(stream);
   }
 
diff --git a/third_party/xla/xla/stream_executor/cuda/cuda_executor.cc b/third_party/xla/xla/stream_executor/cuda/cuda_executor.cc
index edcd40c16aba7f..596a98fa028e14 100644
--- a/third_party/xla/xla/stream_executor/cuda/cuda_executor.cc
+++ b/third_party/xla/xla/stream_executor/cuda/cuda_executor.cc
@@ -794,13 +794,6 @@ Event::Status GpuExecutor::PollForEventStatus(Event* event) {
   return AsGpuEvent(event)->PollForStatus();
 }
 
-bool GpuExecutor::AllocateStream(Stream* stream) {
-  absl::MutexLock l(&alive_gpu_streams_mu_);
-  bool out = AsGpuStream(stream)->Init();
-  alive_gpu_streams_[stream->platform_specific_handle().stream] = stream;
-  return out;
-}
-
 void GpuExecutor::DeallocateStream(Stream* stream) {
   {
     absl::MutexLock lock(&mu_);
@@ -969,9 +962,16 @@ absl::StatusOr<std::unique_ptr<Stream>> GpuExecutor::CreateStream(
       gpu_stream->SetPriority(std::get<int>(*priority));
     }
   }
-  auto stream = std::make_unique<Stream>(this, std::move(gpu_stream));
-  TF_RETURN_IF_ERROR(stream->Initialize(priority));
-  return std::move(stream);
+  absl::MutexLock l(&alive_gpu_streams_mu_);
+  bool init_worked = gpu_stream->Init();
+  if (init_worked) {
+    auto platform_specific_stream = gpu_stream->platform_specific_stream();
+    auto stream = std::make_unique<Stream>(this, std::move(gpu_stream));
+    alive_gpu_streams_[platform_specific_stream] = stream.get();
+    return std::move(stream);
+  } else {
+    return absl::InvalidArgumentError("Failed to initialize gpu stream");
+  }
 }
 
 absl::StatusOr<std::unique_ptr<Kernel>> GpuExecutor::CreateKernel() {
diff --git a/third_party/xla/xla/stream_executor/gpu/gpu_executor.h b/third_party/xla/xla/stream_executor/gpu/gpu_executor.h
index ab4b97db67c8f3..4fda798c54078f 100644
--- a/third_party/xla/xla/stream_executor/gpu/gpu_executor.h
+++ b/third_party/xla/xla/stream_executor/gpu/gpu_executor.h
@@ -229,8 +229,6 @@ class GpuExecutor : public StreamExecutor {
   bool HostCallback(Stream* stream,
                     absl::AnyInvocable<absl::Status() &&> callback) override;
 
-  bool AllocateStream(Stream* stream) override;
-
   void DeallocateStream(Stream* stream) override;
 
   bool CreateStreamDependency(Stream* dependent, Stream* other) override;
diff --git a/third_party/xla/xla/stream_executor/gpu/gpu_stream.h b/third_party/xla/xla/stream_executor/gpu/gpu_stream.h
index ebd143953e54b5..a06839b5c1c79b 100644
--- a/third_party/xla/xla/stream_executor/gpu/gpu_stream.h
+++ b/third_party/xla/xla/stream_executor/gpu/gpu_stream.h
@@ -45,8 +45,7 @@ class GpuStream : public StreamInterface {
 
   void* platform_specific_stream() override { return gpu_stream_; }
 
-  // Explicitly initialize the CUDA resources associated with this stream, used
-  // by StreamExecutor::AllocateStream().
+  // Explicitly initialize the CUDA resources associated with this stream.
   bool Init();
 
   void SetPriority(StreamPriority priority) override {
diff --git a/third_party/xla/xla/stream_executor/host/host_executor.cc b/third_party/xla/xla/stream_executor/host/host_executor.cc
index 3d59b1c69ea587..8a9c812bb7ecdc 100644
--- a/third_party/xla/xla/stream_executor/host/host_executor.cc
+++ b/third_party/xla/xla/stream_executor/host/host_executor.cc
@@ -212,8 +212,6 @@ bool HostExecutor::HostCallback(
   return true;
 }
 
-bool HostExecutor::AllocateStream(Stream* stream) { return true; }
-
 void HostExecutor::DeallocateStream(Stream* stream) {}
 
 bool HostExecutor::CreateStreamDependency(Stream* dependent, Stream* other) {
@@ -304,9 +302,7 @@ HostExecutor::CreateDeviceDescription(int device_ordinal) {
 
 absl::StatusOr<std::unique_ptr<Stream>> HostExecutor::CreateStream(
     std::optional<std::variant<StreamPriority, int>> priority) {
-  auto stream = std::make_unique<Stream>(this, std::make_unique<HostStream>());
-  TF_RETURN_IF_ERROR(stream->Initialize(priority));
-  return std::move(stream);
+  return std::make_unique<Stream>(this, std::make_unique<HostStream>());
 }
 
 }  // namespace host
diff --git a/third_party/xla/xla/stream_executor/host/host_executor.h b/third_party/xla/xla/stream_executor/host/host_executor.h
index 40cb0ccea89330..c57a078e6f6b51 100644
--- a/third_party/xla/xla/stream_executor/host/host_executor.h
+++ b/third_party/xla/xla/stream_executor/host/host_executor.h
@@ -111,7 +111,6 @@ class HostExecutor : public StreamExecutor {
   absl::Status WaitForEvent(Stream* stream, Event* event) override;
   Event::Status PollForEventStatus(Event* event) override;
 
-  bool AllocateStream(Stream* stream) override;
   void DeallocateStream(Stream* stream) override;
   bool CreateStreamDependency(Stream* dependent, Stream* other) override;
 
diff --git a/third_party/xla/xla/stream_executor/mock_stream_executor.h b/third_party/xla/xla/stream_executor/mock_stream_executor.h
index 12a0c2794baf7b..89464874408753 100644
--- a/third_party/xla/xla/stream_executor/mock_stream_executor.h
+++ b/third_party/xla/xla/stream_executor/mock_stream_executor.h
@@ -142,7 +142,6 @@ class MockStreamExecutor : public StreamExecutorInterface {
   MOCK_METHOD(absl::Status, WaitForEventOnExternalStream,
               (std::intptr_t stream, Event* event), (override));
   MOCK_METHOD(Event::Status, PollForEventStatus, (Event * event), (override));
-  MOCK_METHOD(bool, AllocateStream, (Stream * stream), (override));
   MOCK_METHOD(void, DeallocateStream, (Stream * stream), (override));
   MOCK_METHOD(bool, CreateStreamDependency, (Stream * dependent, Stream* other),
               (override));
diff --git a/third_party/xla/xla/stream_executor/rocm/rocm_executor.cc b/third_party/xla/xla/stream_executor/rocm/rocm_executor.cc
index 7f714d765354dc..226cf3afe231b7 100644
--- a/third_party/xla/xla/stream_executor/rocm/rocm_executor.cc
+++ b/third_party/xla/xla/stream_executor/rocm/rocm_executor.cc
@@ -698,13 +698,6 @@ Event::Status GpuExecutor::PollForEventStatus(Event* event) {
   return AsGpuEvent(event)->PollForStatus();
 }
 
-bool GpuExecutor::AllocateStream(Stream* stream) {
-  absl::MutexLock l(&alive_gpu_streams_mu_);
-  bool out = AsGpuStream(stream)->Init();
-  alive_gpu_streams_[stream->platform_specific_handle().stream] = stream;
-  return out;
-}
-
 void GpuExecutor::DeallocateStream(Stream* stream) {
   {
     absl::MutexLock lock(&mu_);
@@ -875,9 +868,16 @@ absl::StatusOr<std::unique_ptr<Stream>> GpuExecutor::CreateStream(
       gpu_stream->SetPriority(std::get<int>(*priority));
     }
   }
-  auto stream = std::make_unique<Stream>(this, std::move(gpu_stream));
-  TF_RETURN_IF_ERROR(stream->Initialize(priority));
-  return std::move(stream);
+  absl::MutexLock l(&alive_gpu_streams_mu_);
+  bool init_worked = gpu_stream->Init();
+  if (init_worked) {
+    auto platform_specific_stream = gpu_stream->platform_specific_stream();
+    auto stream = std::make_unique<Stream>(this, std::move(gpu_stream));
+    alive_gpu_streams_[platform_specific_stream] = stream.get();
+    return std::move(stream);
+  } else {
+    return absl::InvalidArgumentError("Failed to initialize GPU stream");
+  }
 }
 
 absl::StatusOr<std::unique_ptr<Kernel>> GpuExecutor::CreateKernel() {
diff --git a/third_party/xla/xla/stream_executor/stream.cc b/third_party/xla/xla/stream_executor/stream.cc
index 4fc078b0ae2c14..dbb0ca213cced9 100644
--- a/third_party/xla/xla/stream_executor/stream.cc
+++ b/third_party/xla/xla/stream_executor/stream.cc
@@ -47,17 +47,6 @@ Stream::Stream(StreamExecutor *parent,
       implementation_(std::move(implementation)),
       status_(absl::OkStatus()) {}
 
-absl::Status Stream::Initialize(
-    std::optional<std::variant<StreamPriority, int>> priority) {
-  absl::MutexLock lock(&mu_);
-  if (parent_->AllocateStream(this)) {
-    // Successful initialization!
-    return absl::OkStatus();
-  }
-
-  return absl::InternalError("failed to allocate stream during initialization");
-}
-
 Stream::~Stream() {
   // Ensure the stream is completed.
   auto status = BlockHostUntilDone();
diff --git a/third_party/xla/xla/stream_executor/stream.h b/third_party/xla/xla/stream_executor/stream.h
index cc0d5d94ae0524..f6bf0a49335d82 100644
--- a/third_party/xla/xla/stream_executor/stream.h
+++ b/third_party/xla/xla/stream_executor/stream.h
@@ -110,11 +110,6 @@ class Stream {
   // execution status.
   absl::Status RefreshStatus() TF_LOCKS_EXCLUDED(mu_);
 
-  // Initialize the stream. This must be performed before entraining any other
-  // operations.
-  absl::Status Initialize(
-      std::optional<std::variant<StreamPriority, int>> priority = std::nullopt);
-
   // Get or create a sub-stream from this stream. If there is any sub-stream in
   // the pool that can be reused then just return this sub-stream.  Otherwise
   // create a new sub-stream.
diff --git a/third_party/xla/xla/stream_executor/stream_executor_interface.h b/third_party/xla/xla/stream_executor/stream_executor_interface.h
index 9097ea4f064c91..65eff81b0a4d25 100644
--- a/third_party/xla/xla/stream_executor/stream_executor_interface.h
+++ b/third_party/xla/xla/stream_executor/stream_executor_interface.h
@@ -282,10 +282,6 @@ class StreamExecutorInterface {
   // Requests the current status of the event from the underlying platform.
   virtual Event::Status PollForEventStatus(Event* event) = 0;
 
-  // Allocates stream resources on the underlying platform and initializes its
-  // internals.
-  virtual bool AllocateStream(Stream* stream) = 0;
-
   // Deallocates stream resources on the underlying platform.
   virtual void DeallocateStream(Stream* stream) = 0;
 
diff --git a/third_party/xla/xla/stream_executor/tpu/tpu_executor.cc b/third_party/xla/xla/stream_executor/tpu/tpu_executor.cc
index c36d09f5c238f0..c59c0c45db1b82 100644
--- a/third_party/xla/xla/stream_executor/tpu/tpu_executor.cc
+++ b/third_party/xla/xla/stream_executor/tpu/tpu_executor.cc
@@ -81,11 +81,6 @@ tensorflow::tpu::TpuCoreLocationExternal TpuExecutor::GetCoreLocationExternal()
       ExecutorApiFn()->TpuExecutor_GetCoreLocationFn(executor_));
 }
 
-bool TpuExecutor::AllocateStream(Stream* stream) {
-  return ExecutorApiFn()->TpuExecutor_AllocateStreamFn(
-      executor_, get_stream(stream->implementation()));
-}
-
 void TpuExecutor::DeallocateStream(Stream* stream) {
   ExecutorApiFn()->TpuExecutor_DeallocateStreamFn(
       executor_, get_stream(stream->implementation()));
@@ -142,7 +137,6 @@ absl::StatusOr<std::unique_ptr<Stream>> TpuExecutor::CreateStream(
   stream_map()[ptr.get()] = tpu_stream;
   tpu_platform().mutex().Unlock();
   auto stream = std::make_unique<Stream>(this, std::move(ptr));
-  TF_RETURN_IF_ERROR(stream->Initialize(priority));
   return std::move(stream);
 }
 
diff --git a/third_party/xla/xla/stream_executor/tpu/tpu_executor.h b/third_party/xla/xla/stream_executor/tpu/tpu_executor.h
index 2c38b2ed0b464c..74b17765ef5b92 100644
--- a/third_party/xla/xla/stream_executor/tpu/tpu_executor.h
+++ b/third_party/xla/xla/stream_executor/tpu/tpu_executor.h
@@ -70,8 +70,6 @@ class TpuExecutor : public tensorflow::tpu::TpuExecutorInterface {
 
   absl::Status AllocateEvent(Event* event) override;
 
-  bool AllocateStream(Stream* stream) override;
-
   absl::Status BlockHostUntilDone(Stream* stream) override;
 
   StatusOr<std::unique_ptr<DeviceDescription>> CreateDeviceDescription()
diff --git a/third_party/xla/xla/stream_executor/tpu/tpu_executor_init_fns.inc b/third_party/xla/xla/stream_executor/tpu/tpu_executor_init_fns.inc
index 3221fe01c0076d..4b5f8b64890efd 100644
--- a/third_party/xla/xla/stream_executor/tpu/tpu_executor_init_fns.inc
+++ b/third_party/xla/xla/stream_executor/tpu/tpu_executor_init_fns.inc
@@ -22,7 +22,6 @@ absl::Status SetExecutorStructFn(
   TFTPU_SET_FN(executor_fn, TpuExecutor_Deallocate);
   TFTPU_SET_FN(executor_fn, TpuExecutor_GetAllocatorStats);
   TFTPU_SET_FN(executor_fn, TpuExecutor_DeviceMemoryUsage);
-  TFTPU_SET_FN(executor_fn, TpuExecutor_AllocateStream);
   TFTPU_SET_FN(executor_fn, TpuExecutor_DeallocateStream);
   TFTPU_SET_FN(executor_fn, TpuExecutor_CreateStreamDependency);
   TFTPU_SET_FN(executor_fn, TpuExecutor_GetStatus);

From 5699e32d0b2f7bd93e4a891eeee05c3f52315eae Mon Sep 17 00:00:00 2001
From: Kuy Mainwaring <kuym@google.com>
Date: Tue, 14 May 2024 15:19:41 -0700
Subject: [PATCH 140/478] [XLA:GPU] Clang-tidy fixes for gpu_norm_runner.h,
 in_place_dynamic_update_slice_mlir.h, tiling_util.h and indexing_analysis.h

PiperOrigin-RevId: 633722707
---
 third_party/xla/xla/service/gpu/fusions/BUILD  |  2 ++
 .../in_place_dynamic_update_slice_mlir.h       |  3 +--
 .../xla/xla/service/gpu/fusions/tiling_util.h  |  7 +++++++
 .../xla/xla/service/gpu/gpu_norm_runner.h      |  2 +-
 .../xla/service/gpu/model/indexing_analysis.h  | 18 +++++++++---------
 5 files changed, 20 insertions(+), 12 deletions(-)

diff --git a/third_party/xla/xla/service/gpu/fusions/BUILD b/third_party/xla/xla/service/gpu/fusions/BUILD
index 1b1d9c8a5ef0e5..17b48188bed3ea 100644
--- a/third_party/xla/xla/service/gpu/fusions/BUILD
+++ b/third_party/xla/xla/service/gpu/fusions/BUILD
@@ -585,8 +585,10 @@ cc_library(
         "//xla/service/llvm_ir:kernel_support_library",
         "//xla/service/llvm_ir:llvm_loop",
         "//xla/service/llvm_ir:llvm_util",
+        "@com_google_absl//absl/container:inlined_vector",
         "@com_google_absl//absl/log:check",
         "@com_google_absl//absl/strings",
+        "@com_google_absl//absl/strings:str_format",
         "@com_google_absl//absl/types:span",
         "@llvm-project//llvm:Support",
         "@llvm-project//llvm:ir_headers",
diff --git a/third_party/xla/xla/service/gpu/fusions/in_place_dynamic_update_slice_mlir.h b/third_party/xla/xla/service/gpu/fusions/in_place_dynamic_update_slice_mlir.h
index 9d1514c131524d..efcfce8b2c6ede 100644
--- a/third_party/xla/xla/service/gpu/fusions/in_place_dynamic_update_slice_mlir.h
+++ b/third_party/xla/xla/service/gpu/fusions/in_place_dynamic_update_slice_mlir.h
@@ -21,7 +21,6 @@ limitations under the License.
 
 #include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
 #include "mlir/IR/MLIRContext.h"  // from @llvm-project
-#include "mlir/Interfaces/DataLayoutInterfaces.h"  // from @llvm-project
 #include "xla/hlo/ir/hlo_instruction.h"
 #include "xla/hlo/ir/hlo_instructions.h"
 #include "xla/service/gpu/fusions/mlir/computation_partitioner.h"
@@ -62,7 +61,7 @@ class MlirInPlaceDynamicUpdateSliceFusion : public MlirFusionEmitterBase {
 
   std::optional<IndexingMap> ComputeThreadIdToInputIndexing(
       int64_t root_index, int64_t hero_operand_index,
-      mlir::MLIRContext* indexing_context) const override;
+      mlir::MLIRContext* mlir_context) const override;
 
  protected:
   absl::Status EmitEntryFunction(
diff --git a/third_party/xla/xla/service/gpu/fusions/tiling_util.h b/third_party/xla/xla/service/gpu/fusions/tiling_util.h
index f06ae8ccab4280..977dbdaf986f72 100644
--- a/third_party/xla/xla/service/gpu/fusions/tiling_util.h
+++ b/third_party/xla/xla/service/gpu/fusions/tiling_util.h
@@ -19,9 +19,16 @@ limitations under the License.
 #include <functional>
 #include <string>
 
+#include "absl/container/inlined_vector.h"
 #include "absl/log/check.h"
+#include "absl/strings/str_format.h"
+#include "absl/strings/str_join.h"
 #include "absl/types/span.h"
+#include "llvm/IR/IRBuilder.h"
+#include "llvm/IR/Type.h"
+#include "llvm/IR/Value.h"
 #include "xla/service/llvm_ir/ir_array.h"
+#include "xla/shape.h"
 #include "xla/shape_util.h"
 #include "xla/util.h"
 
diff --git a/third_party/xla/xla/service/gpu/gpu_norm_runner.h b/third_party/xla/xla/service/gpu/gpu_norm_runner.h
index 8461671e86d037..7c9da554a16fa3 100644
--- a/third_party/xla/xla/service/gpu/gpu_norm_runner.h
+++ b/third_party/xla/xla/service/gpu/gpu_norm_runner.h
@@ -184,7 +184,7 @@ absl::Status RunGpuNorm(const GpuNormConfig& conv_config,
                         const se::DeviceMemoryBase& y_or_dx_buffer,
                         std::optional<se::DeviceMemoryBase> bias_buffer,
                         std::optional<se::DeviceMemoryBase> dy_buffer,
-                        std::optional<se::DeviceMemoryBase> exepctation_buffer,
+                        std::optional<se::DeviceMemoryBase> expectation_buffer,
                         std::optional<se::DeviceMemoryBase> norm_factor_buffer,
                         std::optional<se::DeviceMemoryBase> dscale_buffer,
                         std::optional<se::DeviceMemoryBase> dbias_buffer,
diff --git a/third_party/xla/xla/service/gpu/model/indexing_analysis.h b/third_party/xla/xla/service/gpu/model/indexing_analysis.h
index 5408d1e5dd9708..76fc8ab39e476c 100644
--- a/third_party/xla/xla/service/gpu/model/indexing_analysis.h
+++ b/third_party/xla/xla/service/gpu/model/indexing_analysis.h
@@ -104,7 +104,7 @@ HloInstructionIndexing ComputeInputToOutputIndexing(const HloInstruction* instr,
 // fusion does not make much sense, but they are created sometimes.
 IndexingMap ComputeEpilogueInputToOutputIndexing(
     HloInstructionAdaptor epilogue_parent, HloInstructionAdaptor epilogue_root,
-    mlir::MLIRContext* ctx);
+    mlir::MLIRContext* mlir_context);
 
 using GroupedByOpIndexingMap =
     absl::flat_hash_map<const HloInstruction*, IndexingMapSet>;
@@ -132,32 +132,32 @@ bool FuseProducerConsumerOutputToInputIndexing(
 // Equivalent to linearizing the input_shape index and then delinearizing it
 // to output_shape.
 IndexingMap GetBitcastMap(const Shape& input_shape, const Shape& output_shape,
-                          mlir::MLIRContext* ctx);
+                          mlir::MLIRContext* mlir_context);
 
 // Creates an indexing map from the physical layout of the tensor to its logical
 // layout.
-IndexingMap GetIndexingMapFromPhysicalLayoutToLogical(const Shape& shape,
-                                                      mlir::MLIRContext* ctx);
+IndexingMap GetIndexingMapFromPhysicalLayoutToLogical(
+    const Shape& shape, mlir::MLIRContext* mlir_context);
 
 // Creates an indexing map from the logical layout of the tensor to its physical
 // layout.
-IndexingMap GetIndexingMapFromLogicalToPhysicalLayout(const Shape& shape,
-                                                      mlir::MLIRContext* ctx);
+IndexingMap GetIndexingMapFromLogicalToPhysicalLayout(
+    const Shape& shape, mlir::MLIRContext* mlir_context);
 
 // Creates an indexing map from thread and block IDs to elements of the tiled
 // shape. Uses the same convention as KernelFusionInterface: dimensions 0 to 2
 // are thread indices (currently only 0 is used), dimensions 3 to 5 are block
 // indices (currently only 3 is used).
 mlir::AffineMap GetBlockOffsetsForTiling(const Tiling& tiling,
-                                         mlir::MLIRContext* ctx);
+                                         mlir::MLIRContext* mlir_context);
 mlir::AffineMap GetThreadOffsetsForTiling(const Tiling& tiling,
-                                          mlir::MLIRContext* ctx);
+                                          mlir::MLIRContext* mlir_context);
 
 // Convenience functions for the two functions above
 // (`GetBlockOffsestsForTiling` + `GetThreadOffsetsForTiling`). Also sets up
 // the ranges of dimensions and symbols.
 IndexingMap GetIndexingMapForTiling(const Tiling& tiling,
-                                    mlir::MLIRContext* ctx);
+                                    mlir::MLIRContext* mlir_context);
 IndexingMap GetIndexingMapForTiling(mlir::AffineMap block_offsets,
                                     mlir::AffineMap thread_offsets,
                                     int64_t threads_per_block,

From 8d5affaa0e8d971becc90b253ff5b4e6bf69908f Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Tue, 14 May 2024 15:24:32 -0700
Subject: [PATCH 141/478] Update ops-related pbtxt files.

PiperOrigin-RevId: 633723992
---
 tensorflow/core/ops/ops.pbtxt | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorflow/core/ops/ops.pbtxt b/tensorflow/core/ops/ops.pbtxt
index 4ac4c1beb21194..cdc0191b8ad6c9 100644
--- a/tensorflow/core/ops/ops.pbtxt
+++ b/tensorflow/core/ops/ops.pbtxt
@@ -1,4 +1,4 @@
-go/nodeserialize    
+go/debugproto   
 op {
   name: "Abort"
   attr {

From d17f7cea03029b5e8b7848ea1345edcff09d8f2b Mon Sep 17 00:00:00 2001
From: Kyle Lucke <klucke@google.com>
Date: Tue, 14 May 2024 15:31:17 -0700
Subject: [PATCH 142/478] Eliminate deprecated Create methods, and use the
 KernelFactory methods everywhere.

PiperOrigin-RevId: 633726272
---
 third_party/xla/xla/service/gpu/BUILD         |  3 +
 .../xla/xla/service/gpu/buffer_comparator.cc  | 10 +--
 third_party/xla/xla/service/gpu/kernels/BUILD |  1 +
 .../xla/service/gpu/kernels/topk_kernel.cc    |  8 +-
 .../xla/service/gpu/make_batch_pointers.cc    |  3 +-
 .../xla/service/gpu/stream_executor_util.cc   |  6 +-
 third_party/xla/xla/stream_executor/BUILD     |  1 +
 third_party/xla/xla/stream_executor/gpu/BUILD |  3 +
 .../stream_executor/gpu/gpu_command_buffer.cc | 29 ++++---
 .../gpu/gpu_timer_kernel_cuda.cu.cc           |  7 +-
 .../gpu/redzone_allocator_kernel_cuda.cc      |  7 +-
 .../xla/xla/stream_executor/host/BUILD        |  1 +
 .../stream_executor/host/host_kernel_test.cc  |  3 +-
 third_party/xla/xla/stream_executor/kernel.cc |  7 --
 third_party/xla/xla/stream_executor/kernel.h  | 78 -------------------
 .../xla/xla/stream_executor/kernel_test.cc    |  3 +-
 .../stream_executor/typed_kernel_factory.h    |  6 +-
 17 files changed, 59 insertions(+), 117 deletions(-)

diff --git a/third_party/xla/xla/service/gpu/BUILD b/third_party/xla/xla/service/gpu/BUILD
index 3d6c90b8119d96..c19905c49698c8 100644
--- a/third_party/xla/xla/service/gpu/BUILD
+++ b/third_party/xla/xla/service/gpu/BUILD
@@ -4367,6 +4367,7 @@ cc_library(
         "//xla/stream_executor",
         "//xla/stream_executor:kernel_factory",
         "//xla/stream_executor:launch_dim",
+        "//xla/stream_executor:typed_kernel_factory",
         "//xla/tsl/util:env_var",
         "//xla/tsl/util/proto:proto_utils",
         "@com_google_absl//absl/algorithm:container",
@@ -4476,6 +4477,7 @@ cc_library(
         "//xla/service:hlo_module_config",
         "//xla/stream_executor",
         "//xla/stream_executor:device_memory_handle",
+        "//xla/stream_executor:typed_kernel_factory",
         "//xla/stream_executor/gpu:asm_compiler",
         "@local_tsl//tsl/platform:errors",
         "@local_tsl//tsl/platform:logging",
@@ -5329,6 +5331,7 @@ cc_library(
         "//xla:util",
         "//xla/stream_executor",
         "//xla/stream_executor:device_memory",
+        "//xla/stream_executor:typed_kernel_factory",
         "//xla/stream_executor/gpu:gpu_stream_header",
         "@com_google_absl//absl/status",
         "@local_tsl//tsl/platform:errors",
diff --git a/third_party/xla/xla/service/gpu/buffer_comparator.cc b/third_party/xla/xla/service/gpu/buffer_comparator.cc
index ab9141b3e042f5..f4f5112add357b 100644
--- a/third_party/xla/xla/service/gpu/buffer_comparator.cc
+++ b/third_party/xla/xla/service/gpu/buffer_comparator.cc
@@ -31,6 +31,7 @@ limitations under the License.
 #include "xla/stream_executor/device_memory_handle.h"
 #include "xla/stream_executor/kernel.h"
 #include "xla/stream_executor/stream_executor.h"
+#include "xla/stream_executor/typed_kernel_factory.h"
 #include "xla/util.h"
 #include "tsl/platform/errors.h"
 #include "tsl/platform/logging.h"
@@ -75,11 +76,10 @@ static absl::StatusOr<bool> DeviceCompare(se::Stream* stream,
 
   TF_ASSIGN_OR_RETURN(
       ComparisonKernelT<ElementT> comparison_kernel,
-      (se::TypedKernel<se::DeviceMemory<ElementT>, se::DeviceMemory<ElementT>,
-                       float, uint64_t,
-                       se::DeviceMemory<uint64_t>>::Create(executor,
-                                                           kernel_name,
-                                                           kernel_symbol)));
+      (se::TypedKernelFactory<
+          se::DeviceMemory<ElementT>, se::DeviceMemory<ElementT>, float,
+          uint64_t, se::DeviceMemory<uint64_t>>::Create(executor, kernel_name,
+                                                        kernel_symbol)));
 
   const se::DeviceDescription& gpu_device_info =
       executor->GetDeviceDescription();
diff --git a/third_party/xla/xla/service/gpu/kernels/BUILD b/third_party/xla/xla/service/gpu/kernels/BUILD
index fc9c7ec54b2aff..a7add705572576 100644
--- a/third_party/xla/xla/service/gpu/kernels/BUILD
+++ b/third_party/xla/xla/service/gpu/kernels/BUILD
@@ -143,6 +143,7 @@ cc_library(
         "//xla:xla_data_proto_cc",
         "//xla/stream_executor",  # build_cleaner: keep
         "//xla/stream_executor:platform",
+        "//xla/stream_executor:typed_kernel_factory",
         "//xla/stream_executor/gpu:gpu_stream_header",
         "//xla/stream_executor/gpu:gpu_types_header",
         "@com_google_absl//absl/numeric:bits",
diff --git a/third_party/xla/xla/service/gpu/kernels/topk_kernel.cc b/third_party/xla/xla/service/gpu/kernels/topk_kernel.cc
index 3ce37ef5978b13..1595d823b41fd8 100644
--- a/third_party/xla/xla/service/gpu/kernels/topk_kernel.cc
+++ b/third_party/xla/xla/service/gpu/kernels/topk_kernel.cc
@@ -32,6 +32,7 @@ limitations under the License.
 #include "xla/stream_executor/kernel.h"
 #include "xla/stream_executor/launch_dim.h"
 #include "xla/stream_executor/stream.h"
+#include "xla/stream_executor/typed_kernel_factory.h"
 #include "xla/types.h"
 #include "xla/xla_data.pb.h"
 #include "tsl/platform/errors.h"
@@ -84,9 +85,10 @@ absl::Status TypedTopK(se::Stream* stream, se::DeviceMemoryBase data,
   TF_ASSIGN_OR_RETURN(void* kernel_symbol, GetKernel<T>(num_elements, k));
   TF_ASSIGN_OR_RETURN(
       auto kernel,
-      (se::TypedKernel<se::DeviceMemory<T>, size_t, se::DeviceMemory<T>,
-                       se::DeviceMemory<uint32_t>,
-                       size_t>::Create(executor, "topk", kernel_symbol)));
+      (se::TypedKernelFactory<se::DeviceMemory<T>, size_t, se::DeviceMemory<T>,
+                              se::DeviceMemory<uint32_t>,
+                              size_t>::Create(executor, "topk",
+                                              kernel_symbol)));
 
   TF_RETURN_IF_ERROR(stream->ThenLaunch(
       se::ThreadDim(num_threads, 1, 1), se::BlockDim(batch_size, 1, 1),
diff --git a/third_party/xla/xla/service/gpu/make_batch_pointers.cc b/third_party/xla/xla/service/gpu/make_batch_pointers.cc
index e7788d652dcc40..f0de161a35e3ce 100644
--- a/third_party/xla/xla/service/gpu/make_batch_pointers.cc
+++ b/third_party/xla/xla/service/gpu/make_batch_pointers.cc
@@ -23,6 +23,7 @@ limitations under the License.
 #include "xla/stream_executor/kernel.h"
 #include "xla/stream_executor/launch_dim.h"
 #include "xla/stream_executor/stream.h"
+#include "xla/stream_executor/typed_kernel_factory.h"
 #include "xla/util.h"
 #include "tsl/platform/errors.h"
 #include "tsl/platform/statusor.h"
@@ -60,7 +61,7 @@ absl::Status MakeBatchPointers(se::Stream* stream,
 
   TF_ASSIGN_OR_RETURN(
       auto kernel,
-      (se::TypedKernel<
+      (se::TypedKernelFactory<
           se::DeviceMemoryBase, size_t, size_t,
           se::DeviceMemoryBase>::Create(executor, "make_batch_pointers",
                                         make_batch_pointers::kernel())));
diff --git a/third_party/xla/xla/service/gpu/stream_executor_util.cc b/third_party/xla/xla/service/gpu/stream_executor_util.cc
index 8760faa725bcef..93c0a63bb14566 100644
--- a/third_party/xla/xla/service/gpu/stream_executor_util.cc
+++ b/third_party/xla/xla/service/gpu/stream_executor_util.cc
@@ -57,6 +57,7 @@ limitations under the License.
 #include "xla/stream_executor/launch_dim.h"
 #include "xla/stream_executor/platform.h"
 #include "xla/stream_executor/stream.h"
+#include "xla/stream_executor/typed_kernel_factory.h"
 #include "xla/tsl/util/env_var.h"
 #include "xla/tsl/util/proto/proto_utils.h"
 #include "xla/util.h"
@@ -489,8 +490,9 @@ static void InitializeTypedBuffer(se::Stream* stream,
   // Repeat the host_buffer_size elements at the start of `buf` to the end
   CHECK_EQ(elements_to_fill, buffer.size() / sizeof(T) - host_buffer_size);
   se::StreamExecutor* executor = stream->parent();
-  auto kernel = se::TypedKernel<se::DeviceMemoryBase, int64_t, int64_t>::Create(
-      executor, "RepeatBufferKernel", repeat_buffer_kernel::kernel());
+  auto kernel =
+      se::TypedKernelFactory<se::DeviceMemoryBase, int64_t, int64_t>::Create(
+          executor, "RepeatBufferKernel", repeat_buffer_kernel::kernel());
   if (!kernel.ok()) {
     LOG(FATAL) << "Could not create RepeatBufferKernel: " << kernel.status();
   }
diff --git a/third_party/xla/xla/stream_executor/BUILD b/third_party/xla/xla/stream_executor/BUILD
index 3899b903553264..d41de1e0ff7032 100644
--- a/third_party/xla/xla/stream_executor/BUILD
+++ b/third_party/xla/xla/stream_executor/BUILD
@@ -752,6 +752,7 @@ xla_cc_test(
     deps = [
         ":device_memory",
         ":stream_executor",
+        ":typed_kernel_factory",
         "//xla/stream_executor/host:host_platform",
         "@local_tsl//tsl/platform:test",
         "@local_tsl//tsl/platform:test_benchmark",
diff --git a/third_party/xla/xla/stream_executor/gpu/BUILD b/third_party/xla/xla/stream_executor/gpu/BUILD
index 8e70a01a43a6ed..619a5bb034ab76 100644
--- a/third_party/xla/xla/stream_executor/gpu/BUILD
+++ b/third_party/xla/xla/stream_executor/gpu/BUILD
@@ -152,6 +152,7 @@ gpu_only_cc_library(
         ":gpu_types_header",
         "//xla/stream_executor",
         "//xla/stream_executor:stream_executor_interface",
+        "//xla/stream_executor:typed_kernel_factory",
         "@com_google_absl//absl/container:flat_hash_map",
         "@com_google_absl//absl/container:inlined_vector",
         "@com_google_absl//absl/functional:any_invocable",
@@ -343,6 +344,7 @@ gpu_kernel_library(
         ":gpu_semaphore",
         ":gpu_stream",
         "//xla/stream_executor",
+        "//xla/stream_executor:typed_kernel_factory",
         "@com_google_absl//absl/status:statusor",
     ],
 )
@@ -484,6 +486,7 @@ gpu_kernel_library(
         "//xla/stream_executor",
         "//xla/stream_executor:device_memory",
         "//xla/stream_executor:stream_executor_headers",
+        "//xla/stream_executor:typed_kernel_factory",
         "//xla/stream_executor/cuda:cuda_asm_compiler",
         "//xla/stream_executor/cuda:cuda_driver",
         "@com_google_absl//absl/base",
diff --git a/third_party/xla/xla/stream_executor/gpu/gpu_command_buffer.cc b/third_party/xla/xla/stream_executor/gpu/gpu_command_buffer.cc
index c0788cf21a9e58..bcfc6a3790507f 100644
--- a/third_party/xla/xla/stream_executor/gpu/gpu_command_buffer.cc
+++ b/third_party/xla/xla/stream_executor/gpu/gpu_command_buffer.cc
@@ -46,6 +46,7 @@ limitations under the License.
 #include "xla/stream_executor/kernel_spec.h"
 #include "xla/stream_executor/launch_dim.h"
 #include "xla/stream_executor/stream_executor.h"
+#include "xla/stream_executor/typed_kernel_factory.h"
 #include "tsl/platform/env.h"
 #include "tsl/platform/errors.h"
 #include "tsl/platform/logging.h"
@@ -217,8 +218,9 @@ GpuCommandBuffer::GetSetIfConditionKernel() {
   if (!set_if_condition_kernel_) {
     MultiKernelLoaderSpec spec(/*arity=*/2);
     spec.AddCudaPtxInMemory(gpu::GetSetIfConditionKernel(), "set_if_condition");
-    TF_ASSIGN_OR_RETURN(set_if_condition_kernel_,
-                        SetIfConditionKernel::Create(parent_, spec));
+    TF_ASSIGN_OR_RETURN(
+        set_if_condition_kernel_,
+        SetIfConditionKernel::FactoryType::Create(parent_, spec));
   }
   return &set_if_condition_kernel_;
 }
@@ -229,8 +231,9 @@ GpuCommandBuffer::GetSetIfElseConditionKernel() {
     MultiKernelLoaderSpec spec(/*arity=*/3);
     spec.AddCudaPtxInMemory(gpu::GetSetIfElseConditionKernel(),
                             "set_if_else_condition");
-    TF_ASSIGN_OR_RETURN(set_if_else_condition_kernel_,
-                        SetIfElseConditionKernel::Create(parent_, spec));
+    TF_ASSIGN_OR_RETURN(
+        set_if_else_condition_kernel_,
+        SetIfElseConditionKernel::FactoryType::Create(parent_, spec));
   }
   return &set_if_else_condition_kernel_;
 }
@@ -241,8 +244,9 @@ GpuCommandBuffer::GetSetCaseConditionKernel() {
     MultiKernelLoaderSpec spec(/*arity=*/10);
     spec.AddCudaPtxInMemory(gpu::GetSetCaseConditionKernel(),
                             "set_case_condition");
-    TF_ASSIGN_OR_RETURN(set_case_condition_kernel_,
-                        SetCaseConditionKernel::Create(parent_, spec));
+    TF_ASSIGN_OR_RETURN(
+        set_case_condition_kernel_,
+        SetCaseConditionKernel::FactoryType::Create(parent_, spec));
   }
   return &set_case_condition_kernel_;
 }
@@ -253,8 +257,9 @@ GpuCommandBuffer::GetSetForConditionKernel() {
     MultiKernelLoaderSpec spec(/*arity=*/3);
     spec.AddCudaPtxInMemory(gpu::GetSetForConditionKernel(),
                             "set_for_condition");
-    TF_ASSIGN_OR_RETURN(set_for_condition_kernel_,
-                        SetForConditionKernel::Create(parent_, spec));
+    TF_ASSIGN_OR_RETURN(
+        set_for_condition_kernel_,
+        SetForConditionKernel::FactoryType::Create(parent_, spec));
   }
   return &set_for_condition_kernel_;
 }
@@ -265,8 +270,9 @@ GpuCommandBuffer::GetSetWhileConditionKernel() {
     MultiKernelLoaderSpec spec(/*arity=*/2);
     spec.AddCudaPtxInMemory(gpu::GetSetWhileConditionKernel(),
                             "set_while_condition");
-    TF_ASSIGN_OR_RETURN(set_while_condition_kernel_,
-                        SetWhileConditionKernel::Create(parent_, spec));
+    TF_ASSIGN_OR_RETURN(
+        set_while_condition_kernel_,
+        SetWhileConditionKernel::FactoryType::Create(parent_, spec));
   }
   return &set_while_condition_kernel_;
 }
@@ -277,7 +283,8 @@ GpuCommandBuffer::GetNoOpKernel() {
   if (!noop_kernel_) {
     MultiKernelLoaderSpec spec(/*arity=*/0);
     spec.AddCudaPtxInMemory(gpu::kNoOpKernel, "noop");
-    TF_ASSIGN_OR_RETURN(noop_kernel_, NoOpKernel::Create(parent_, spec));
+    TF_ASSIGN_OR_RETURN(noop_kernel_,
+                        NoOpKernel::FactoryType::Create(parent_, spec));
   }
   return &noop_kernel_;
 #else
diff --git a/third_party/xla/xla/stream_executor/gpu/gpu_timer_kernel_cuda.cu.cc b/third_party/xla/xla/stream_executor/gpu/gpu_timer_kernel_cuda.cu.cc
index b4af6019234d04..8a320c583547c6 100644
--- a/third_party/xla/xla/stream_executor/gpu/gpu_timer_kernel_cuda.cu.cc
+++ b/third_party/xla/xla/stream_executor/gpu/gpu_timer_kernel_cuda.cu.cc
@@ -19,6 +19,7 @@ limitations under the License.
 #include "xla/stream_executor/gpu/gpu_executor.h"
 #include "xla/stream_executor/gpu/gpu_semaphore.h"
 #include "xla/stream_executor/gpu/gpu_timer_kernel.h"
+#include "xla/stream_executor/typed_kernel_factory.h"
 
 namespace stream_executor::gpu {
 namespace {
@@ -60,8 +61,10 @@ absl::StatusOr<GpuSemaphore> LaunchDelayKernel(Stream* stream) {
   // multiple GpuTimer objects.
   TF_ASSIGN_OR_RETURN(
       auto kernel,
-      (TypedKernel<DeviceMemory<GpuSemaphoreState>, GpuSemaphoreState>::Create(
-          executor, "DelayKernel", reinterpret_cast<void*>(DelayKernel))));
+      (TypedKernelFactory<DeviceMemory<GpuSemaphoreState>,
+                          GpuSemaphoreState>::Create(executor, "DelayKernel",
+                                                     reinterpret_cast<void*>(
+                                                         DelayKernel))));
   // Launch a delay kernel into this stream, which will spin until
   // GetElapsedDuration() is called, the timer is destroyed, or the timeout
   // in the kernel is reached.
diff --git a/third_party/xla/xla/stream_executor/gpu/redzone_allocator_kernel_cuda.cc b/third_party/xla/xla/stream_executor/gpu/redzone_allocator_kernel_cuda.cc
index 2432c1e5da79d2..6e6e5ad1077590 100644
--- a/third_party/xla/xla/stream_executor/gpu/redzone_allocator_kernel_cuda.cc
+++ b/third_party/xla/xla/stream_executor/gpu/redzone_allocator_kernel_cuda.cc
@@ -34,6 +34,7 @@ limitations under the License.
 #include "xla/stream_executor/gpu/redzone_allocator_kernel.h"
 #include "xla/stream_executor/kernel.h"
 #include "xla/stream_executor/stream_executor_pimpl.h"
+#include "xla/stream_executor/typed_kernel_factory.h"
 #include "tsl/platform/statusor.h"
 
 namespace stream_executor {
@@ -54,9 +55,9 @@ static absl::StatusOr<TypedKernel<Args...>*> LoadKernelOrGetPtr(
 
   auto it = kernel_ptr_cache.find(kernel_ptr_cache_key);
   if (it == kernel_ptr_cache.end()) {
-    TF_ASSIGN_OR_RETURN(
-        TypedKernel<Args...> loaded,
-        (TypedKernel<Args...>::Create(executor, kernel_name, ptx, cubin_data)));
+    TF_ASSIGN_OR_RETURN(TypedKernel<Args...> loaded,
+                        (TypedKernelFactory<Args...>::Create(
+                            executor, kernel_name, ptx, cubin_data)));
     it =
         kernel_ptr_cache.emplace(kernel_ptr_cache_key, std::move(loaded)).first;
   }
diff --git a/third_party/xla/xla/stream_executor/host/BUILD b/third_party/xla/xla/stream_executor/host/BUILD
index da6b0d35ba7d5b..6ad6b619df78c7 100644
--- a/third_party/xla/xla/stream_executor/host/BUILD
+++ b/third_party/xla/xla/stream_executor/host/BUILD
@@ -158,6 +158,7 @@ xla_cc_test(
         ":host_platform",
         "//xla/stream_executor",
         "//xla/stream_executor:device_memory",
+        "//xla/stream_executor:kernel_factory",
         "@com_google_absl//absl/types:span",
         "@local_tsl//tsl/lib/core:status_test_util",
         "@local_tsl//tsl/platform:statusor",
diff --git a/third_party/xla/xla/stream_executor/host/host_kernel_test.cc b/third_party/xla/xla/stream_executor/host/host_kernel_test.cc
index 2dcb432cf81bf4..cebe2b3220c7fd 100644
--- a/third_party/xla/xla/stream_executor/host/host_kernel_test.cc
+++ b/third_party/xla/xla/stream_executor/host/host_kernel_test.cc
@@ -23,6 +23,7 @@ limitations under the License.
 #include "absl/types/span.h"
 #include "xla/stream_executor/device_memory.h"
 #include "xla/stream_executor/host/host_kernel_c_api.h"
+#include "xla/stream_executor/kernel_factory.h"
 #include "xla/stream_executor/kernel_spec.h"
 #include "xla/stream_executor/launch_dim.h"
 #include "xla/stream_executor/platform.h"
@@ -116,7 +117,7 @@ TEST(HostKernelTest, LlvmAddition) {
   auto executor = NewStreamExecutor();
   auto eg = executor.get();
   EXPECT_NE(eg, nullptr);
-  TF_ASSERT_OK_AND_ASSIGN(auto add, HostKernel::Create(eg, spec));
+  TF_ASSERT_OK_AND_ASSIGN(auto add, KernelFactory::Create(eg, spec));
 
   // TODO(tsilytskyi): implement Launch part
   // TF_ASSERT_OK(executor->Launch(ThreadDim(4), args));
diff --git a/third_party/xla/xla/stream_executor/kernel.cc b/third_party/xla/xla/stream_executor/kernel.cc
index ba8f0bbda43b22..81767a7fe5025b 100644
--- a/third_party/xla/xla/stream_executor/kernel.cc
+++ b/third_party/xla/xla/stream_executor/kernel.cc
@@ -52,13 +52,6 @@ void KernelMetadata::set_shared_memory_bytes(int shared_memory_bytes) {
 // Kernel
 //===----------------------------------------------------------------------===//
 
-absl::StatusOr<std::unique_ptr<Kernel>> Kernel::Create(
-    StreamExecutorInterface *executor, const MultiKernelLoaderSpec &spec) {
-  TF_ASSIGN_OR_RETURN(auto kernel, executor->CreateKernel());
-  TF_RETURN_IF_ERROR(executor->GetKernel(spec, kernel.get()));
-  return kernel;
-}
-
 void Kernel::set_name(absl::string_view name) {
   name_ = std::string(name);
 
diff --git a/third_party/xla/xla/stream_executor/kernel.h b/third_party/xla/xla/stream_executor/kernel.h
index 677e8db1a37b35..98de009e1b8a16 100644
--- a/third_party/xla/xla/stream_executor/kernel.h
+++ b/third_party/xla/xla/stream_executor/kernel.h
@@ -229,14 +229,6 @@ class Kernel {
       std::function<absl::StatusOr<std::unique_ptr<KernelArgsPackedArrayBase>>(
           const Kernel &kernel, const KernelArgs &args)>;
 
-  // TODO(b/323534971): Kernel constructor should be moved to StreamExecutor or
-  // a dedicated KernelFactory accessible via StreamExecutor.
-
-  // Creates kernel on a given executor from a given kernel specification.
-  ABSL_DEPRECATED("Use KernelFactory::Create instead.")
-  static absl::StatusOr<std::unique_ptr<Kernel>> Create(
-      StreamExecutorInterface *executor, const MultiKernelLoaderSpec &spec);
-
   Kernel() = default;
   virtual ~Kernel() = default;
 
@@ -293,41 +285,6 @@ class TypedKernel {
  public:
   static constexpr size_t kNumberOfParameters = sizeof...(Params);
 
-  // Creates a typed kernel on a given executor from a kernel specification.
-  ABSL_DEPRECATED("Use TypedKernelFactory::Create instead.")
-  static absl::StatusOr<TypedKernel> Create(StreamExecutorInterface *executor,
-                                            const MultiKernelLoaderSpec &spec) {
-    TF_ASSIGN_OR_RETURN(std::unique_ptr<Kernel> kernel,
-                        Kernel::Create(executor, spec));
-    return TypedKernel(std::move(kernel));
-  }
-
-  // Creates a kernel which can be launched with `stream.ThenLaunch(...)` from a
-  // PTX (and optional CUBIN), such that the types of the arguments provided for
-  // launch would have to match types of the arguments provided at creation
-  // time. The canonical storage for both ptx and cubin_data should outlive the
-  // lifetime of the kernel.
-  ABSL_DEPRECATED("Use TypedKernelFactory::Create instead.")
-  static absl::StatusOr<TypedKernel> Create(
-      StreamExecutorInterface *executor, absl::string_view kernel_name,
-      absl::string_view ptx, absl::Span<const uint8_t> cubin_data);
-
-  // Creates a kernel which can be launched with `stream.ThenLaunch(...)` from
-  // an in-process symbol pointer.
-  ABSL_DEPRECATED("Use TypedKernelFactory::Create instead.")
-  static absl::StatusOr<TypedKernel> Create(StreamExecutorInterface *executor,
-                                            absl::string_view kernel_name,
-                                            void *symbol);
-
-  // Creates a kernel which can be launched with `stream.ThenLaunch(...)` from
-  // an LLVM IR.
-  ABSL_DEPRECATED("Use TypedKernelFactory::Create instead.")
-  static absl::StatusOr<TypedKernel> Create(StreamExecutorInterface *executor,
-                                            absl::string_view ir,
-                                            absl::string_view entrypoint,
-                                            absl::string_view kernel_name,
-                                            absl::Span<std::string> options);
-
   TypedKernel() = default;
 
   Kernel &operator*() { return *kernel_; }
@@ -754,41 +711,6 @@ std::unique_ptr<KernelArgsPackedArrayBase> PackKernelArgs(
   return std::make_unique<PackedArgs>(std::forward<Args>(args)..., shmem_bytes);
 }
 
-template <typename... Args>
-inline absl::StatusOr<TypedKernel<Args...>> TypedKernel<Args...>::Create(
-    StreamExecutorInterface *executor, absl::string_view kernel_name,
-    absl::string_view ptx, absl::Span<const uint8_t> cubin_data) {
-  MultiKernelLoaderSpec loader_spec(TypedKernel<Args...>::kNumberOfParameters);
-  loader_spec.AddCudaPtxInMemory(ptx, kernel_name);
-
-  if (!cubin_data.empty()) {
-    loader_spec.AddCudaCubinInMemory(cubin_data, kernel_name);
-  }
-
-  return TypedKernel<Args...>::Create(executor, loader_spec);
-}
-
-template <typename... Args>
-inline absl::StatusOr<TypedKernel<Args...>> TypedKernel<Args...>::Create(
-    StreamExecutorInterface *executor, absl::string_view kernel_name,
-    void *symbol) {
-  MultiKernelLoaderSpec loader_spec(TypedKernel<Args...>::kNumberOfParameters);
-  loader_spec.AddInProcessSymbol(symbol, kernel_name);
-
-  return TypedKernel<Args...>::Create(executor, loader_spec);
-}
-
-template <typename... Args>
-inline absl::StatusOr<TypedKernel<Args...>> TypedKernel<Args...>::Create(
-    StreamExecutorInterface *executor, absl::string_view ir,
-    absl::string_view entrypoint, absl::string_view kernel_name,
-    absl::Span<std::string> options) {
-  MultiKernelLoaderSpec loader_spec(TypedKernel<Args...>::kNumberOfParameters);
-  loader_spec.AddLlvmHostKernel(ir, entrypoint, kernel_name, options);
-
-  return TypedKernel<Args...>::Create(executor, loader_spec);
-}
-
 }  // namespace stream_executor
 
 #endif  // XLA_STREAM_EXECUTOR_KERNEL_H_
diff --git a/third_party/xla/xla/stream_executor/kernel_test.cc b/third_party/xla/xla/stream_executor/kernel_test.cc
index 205c559170e729..cf63e5b0a55281 100644
--- a/third_party/xla/xla/stream_executor/kernel_test.cc
+++ b/third_party/xla/xla/stream_executor/kernel_test.cc
@@ -26,6 +26,7 @@ limitations under the License.
 #include "xla/stream_executor/platform.h"
 #include "xla/stream_executor/platform_manager.h"
 #include "xla/stream_executor/stream_executor.h"
+#include "xla/stream_executor/typed_kernel_factory.h"
 #include "tsl/platform/test.h"
 #include "tsl/platform/test_benchmark.h"
 
@@ -124,7 +125,7 @@ TEST(KernelTest, FailToCreateTypedKernelFromEmptySpec) {
   MultiKernelLoaderSpec empty_spec(/*arity=*/0);
 
   auto executor = NewStreamExecutor();
-  auto kernel = TypedKernel<>::Create(executor.get(), empty_spec);
+  auto kernel = TypedKernelFactory<>::Create(executor.get(), empty_spec);
   EXPECT_FALSE(kernel.ok());
 }
 
diff --git a/third_party/xla/xla/stream_executor/typed_kernel_factory.h b/third_party/xla/xla/stream_executor/typed_kernel_factory.h
index ac908b13b08e0c..95b9eb6c505ca5 100644
--- a/third_party/xla/xla/stream_executor/typed_kernel_factory.h
+++ b/third_party/xla/xla/stream_executor/typed_kernel_factory.h
@@ -61,7 +61,7 @@ class TypedKernelFactory {
       loader_spec.AddCudaCubinInMemory(cubin_data, kernel_name);
     }
 
-    return TypedKernel<Params...>::Create(executor, loader_spec);
+    return Create(executor, loader_spec);
   }
 
   // Creates a kernel which can be launched with `stream.ThenLaunch(...)` from
@@ -73,7 +73,7 @@ class TypedKernelFactory {
         TypedKernel<Params...>::kNumberOfParameters);
     loader_spec.AddInProcessSymbol(symbol, kernel_name);
 
-    return TypedKernel<Params...>::Create(executor, loader_spec);
+    return Create(executor, loader_spec);
   }
 
   // Creates a kernel which can be launched with `stream.ThenLaunch(...)` from
@@ -86,7 +86,7 @@ class TypedKernelFactory {
         TypedKernel<Params...>::kNumberOfParameters);
     loader_spec.AddLlvmHostKernel(ir, entrypoint, kernel_name, options);
 
-    return TypedKernel<Params...>::Create(executor, loader_spec);
+    return Create(executor, loader_spec);
   }
 };
 

From 2fca8835decdae0e261fbc98c2600ceed36a4d13 Mon Sep 17 00:00:00 2001
From: Deqiang Chen <deqiangc@google.com>
Date: Tue, 14 May 2024 16:03:02 -0700
Subject: [PATCH 143/478] tf_host_callback in tfrt/ifrt use DeviceMgr instead
 of StaticDeviceMgr for better generality and not owning the DeviceMgr since
 that can be owned/created in fallback_request

PiperOrigin-RevId: 633734973
---
 tensorflow/core/tfrt/fallback/fallback_state.h       |  1 +
 tensorflow/core/tfrt/ifrt/ifrt_model_context.h       | 10 ++++------
 tensorflow/core/tfrt/ifrt/ifrt_serving_executable.cc |  6 +++---
 tensorflow/core/tfrt/ifrt/ifrt_serving_executable.h  |  6 +++---
 tensorflow/core/tfrt/ifrt/tf_host_callback.cc        |  2 +-
 tensorflow/core/tfrt/ifrt/tf_host_callback.h         |  2 +-
 tensorflow/core/tfrt/runtime/runtime.h               |  8 ++++++++
 tensorflow/core/tfrt/saved_model/BUILD               |  1 +
 tensorflow/core/tfrt/saved_model/saved_model.cc      |  2 ++
 tensorflow/core/tfrt/tfrt_session/tfrt_session.cc    |  2 ++
 10 files changed, 26 insertions(+), 14 deletions(-)

diff --git a/tensorflow/core/tfrt/fallback/fallback_state.h b/tensorflow/core/tfrt/fallback/fallback_state.h
index ae5158018b9b67..cf293b1b406a28 100644
--- a/tensorflow/core/tfrt/fallback/fallback_state.h
+++ b/tensorflow/core/tfrt/fallback/fallback_state.h
@@ -62,6 +62,7 @@ class FallbackState {
   const SessionOptions &session_options() const { return session_options_; }
 
   const DeviceMgr &device_manager() const { return device_manager_; }
+  DeviceMgr &device_manager() { return device_manager_; }
 
   const DeviceSet &device_set() const { return device_set_; }
 
diff --git a/tensorflow/core/tfrt/ifrt/ifrt_model_context.h b/tensorflow/core/tfrt/ifrt/ifrt_model_context.h
index 9bf38b8c92cb74..20af882d634a7a 100644
--- a/tensorflow/core/tfrt/ifrt/ifrt_model_context.h
+++ b/tensorflow/core/tfrt/ifrt/ifrt_model_context.h
@@ -65,12 +65,12 @@ class IfrtModelContext {
       std::shared_ptr<xla::ifrt::Client> client,
       IfrtServingCoreSelector* ifrt_serving_core_selector,
       const tsl::thread::ThreadPool* thread_pool,
-      std::unique_ptr<tensorflow::StaticDeviceMgr> device_mgr,
+      tensorflow::DeviceMgr* device_mgr,
       tensorflow::XlaHelpers::ShapeRepresentationFn shape_representation_fn)
       : client_(std::move(client)),
         ifrt_serving_core_selector_(ifrt_serving_core_selector),
         thread_pool_(*thread_pool),
-        device_mgr_(std::move(device_mgr)),
+        device_mgr_(device_mgr),
         shape_representation_fn_(shape_representation_fn) {}
 
   void RegisterHandle(ServingExecutableRegistry::Handle handle) {
@@ -100,9 +100,7 @@ class IfrtModelContext {
     return restore_tensor_registry_;
   }
 
-  tensorflow::StaticDeviceMgr* GetDeviceMgr() const {
-    return device_mgr_.get();
-  }
+  tensorflow::DeviceMgr* GetDeviceMgr() const { return device_mgr_; }
   IfrtServingCoreSelector* GetIfrtServingCoreSelector() const {
     return ifrt_serving_core_selector_;
   }
@@ -127,7 +125,7 @@ class IfrtModelContext {
   IfrtServingCoreSelector* ifrt_serving_core_selector_;  // May be nullptr
   const tsl::thread::ThreadPool& thread_pool_;
 
-  std::unique_ptr<tensorflow::StaticDeviceMgr> device_mgr_;
+  tensorflow::DeviceMgr* device_mgr_;  // Not owned.
   tensorflow::XlaHelpers::ShapeRepresentationFn shape_representation_fn_ =
       tensorflow::IdentityShapeRepresentationFn();
 
diff --git a/tensorflow/core/tfrt/ifrt/ifrt_serving_executable.cc b/tensorflow/core/tfrt/ifrt/ifrt_serving_executable.cc
index b65b9c68eb1f78..822df677a97240 100644
--- a/tensorflow/core/tfrt/ifrt/ifrt_serving_executable.cc
+++ b/tensorflow/core/tfrt/ifrt/ifrt_serving_executable.cc
@@ -165,7 +165,7 @@ IfrtServingExecutable::Create(
     IfrtLoadedVariableRegistry* ifrt_loaded_variable_registry,
     const IfrtRestoreTensorRegistry* ifrt_restore,
     tfrt::ConcurrentWorkQueue* checkpoint_loader_queue,
-    tensorflow::StaticDeviceMgr* device_mgr,
+    tensorflow::DeviceMgr* device_mgr,
     tensorflow::XlaHelpers::ShapeRepresentationFn shape_representation_fn,
     IfrtServingCoreSelector* ifrt_serving_core_selector) {
   TF_ASSIGN_OR_RETURN(
@@ -237,7 +237,7 @@ GroupHostCallbackByKey(const Tf2HloResult& tf2hlo_result) {
 // TODO: shape propagation in module
 absl::StatusOr<xla::HostCallback> BuildHostCallback(
     absl::string_view key, const HostCallbackBuilderInfo& builder_info,
-    mlir::ModuleOp module, tensorflow::StaticDeviceMgr* device_mgr,
+    mlir::ModuleOp module, tensorflow::DeviceMgr* device_mgr,
     std::vector<std::unique_ptr<TfHostCallback>>& tf_host_callbacks) {
   VLOG(2) << "BuildHostCallback for key: " << key;
 
@@ -310,7 +310,7 @@ absl::StatusOr<xla::HostCallback> BuildHostCallback(
 
 absl::StatusOr<std::vector<xla::HostCallback>> BuildHostCallbacks(
     const Tf2HloResult& tf2hlo_result, mlir::ModuleOp module,
-    tensorflow::StaticDeviceMgr* device_mgr,
+    tensorflow::DeviceMgr* device_mgr,
     std::vector<std::unique_ptr<TfHostCallback>>& tf_host_callbacks) {
   TF_ASSIGN_OR_RETURN(auto host_callback_maps,
                       GroupHostCallbackByKey(tf2hlo_result));
diff --git a/tensorflow/core/tfrt/ifrt/ifrt_serving_executable.h b/tensorflow/core/tfrt/ifrt/ifrt_serving_executable.h
index 5a336689073cd6..1d224d9345c9dd 100644
--- a/tensorflow/core/tfrt/ifrt/ifrt_serving_executable.h
+++ b/tensorflow/core/tfrt/ifrt/ifrt_serving_executable.h
@@ -70,7 +70,7 @@ class IfrtServingExecutable {
       IfrtLoadedVariableRegistry* ifrt_loaded_variable_registry,
       const IfrtRestoreTensorRegistry* ifrt_restore,
       tfrt::ConcurrentWorkQueue* checkpoint_loader_queue,
-      tensorflow::StaticDeviceMgr* device_mgr,
+      tensorflow::DeviceMgr* device_mgr,
       tensorflow::XlaHelpers::ShapeRepresentationFn shape_representation_fn,
       IfrtServingCoreSelector* ifrt_serving_core_selector);
 
@@ -140,7 +140,7 @@ class IfrtServingExecutable {
       IfrtLoadedVariableRegistry* ifrt_loaded_variable_registry,
       const IfrtRestoreTensorRegistry* ifrt_restore_tensor_registry,
       tfrt::ConcurrentWorkQueue* checkpoint_loader_queue,
-      tensorflow::StaticDeviceMgr* device_mgr,
+      tensorflow::DeviceMgr* device_mgr,
       tensorflow::XlaHelpers::ShapeRepresentationFn shape_representation_fn,
       IfrtServingCoreSelector* ifrt_serving_core_selector,
       tensorflow::tpu::TPUCompileMetadataProto original_compile_metadata)
@@ -176,7 +176,7 @@ class IfrtServingExecutable {
   IfrtLoadedVariableRegistry& ifrt_loaded_variable_registry_;
   const IfrtRestoreTensorRegistry& ifrt_restore_tensor_registry_;
   tfrt::ConcurrentWorkQueue* checkpoint_loader_queue_;
-  tensorflow::StaticDeviceMgr* device_mgr_;  // Not owned. For host callback.
+  tensorflow::DeviceMgr* device_mgr_;  // Not owned. For host callback.
   tensorflow::XlaHelpers::ShapeRepresentationFn shape_representation_fn_;
   IfrtServingCoreSelector* ifrt_serving_core_selector_;
 
diff --git a/tensorflow/core/tfrt/ifrt/tf_host_callback.cc b/tensorflow/core/tfrt/ifrt/tf_host_callback.cc
index 084d1ea1a3ec7a..5c5a48f4fc52b4 100644
--- a/tensorflow/core/tfrt/ifrt/tf_host_callback.cc
+++ b/tensorflow/core/tfrt/ifrt/tf_host_callback.cc
@@ -127,7 +127,7 @@ absl::StatusOr<std::unique_ptr<TfHostCallback>> TfHostCallback::Create(
     absl::string_view entry_function_name,
     absl::Span<const DtypeAndShape> operand_type_and_shapes,
     absl::Span<const DtypeAndShape> result_type_and_shapes,
-    tensorflow::StaticDeviceMgr* device_mgr) {
+    tensorflow::DeviceMgr* device_mgr) {
   tensorflow::SessionOptions options;
   // Explicitly disable non-CPU devices to avoid triggering TPU device
   // initialization inside TF.
diff --git a/tensorflow/core/tfrt/ifrt/tf_host_callback.h b/tensorflow/core/tfrt/ifrt/tf_host_callback.h
index ddb00050e8df58..a78b0e5d0aecea 100644
--- a/tensorflow/core/tfrt/ifrt/tf_host_callback.h
+++ b/tensorflow/core/tfrt/ifrt/tf_host_callback.h
@@ -45,7 +45,7 @@ class TfHostCallback {
       absl::string_view entry_function_name,
       absl::Span<const DtypeAndShape> operand_type_and_shapes,
       absl::Span<const DtypeAndShape> result_type_and_shapes,
-      tensorflow::StaticDeviceMgr* device_mgr);
+      tensorflow::DeviceMgr* device_mgr);
 
   // The host callback function takes two pointer arrays, each element of which
   // points to allocated host buffer in host layout according to corresponding
diff --git a/tensorflow/core/tfrt/runtime/runtime.h b/tensorflow/core/tfrt/runtime/runtime.h
index 830210d1426dab..1a6925c1535cc9 100644
--- a/tensorflow/core/tfrt/runtime/runtime.h
+++ b/tensorflow/core/tfrt/runtime/runtime.h
@@ -27,6 +27,8 @@ limitations under the License.
 #include "absl/log/log.h"
 #include "absl/status/status.h"
 #include "absl/strings/string_view.h"
+#include "tensorflow/core/common_runtime/device_mgr.h"
+#include "tensorflow/core/framework/device.h"
 #include "tensorflow/core/framework/function.h"
 #include "tensorflow/core/graph/graph.h"
 #include "tensorflow/core/platform/statusor.h"
@@ -79,6 +81,11 @@ class ModelRuntimeContext {
     flib_def_ = flib_def;
   }
 
+  tensorflow::DeviceMgr* device_mgr() const { return device_mgr_; }
+  void set_device_mgr(tensorflow::DeviceMgr* device_mgr) {
+    device_mgr_ = device_mgr;
+  }
+
   bool is_local_session() const { return is_local_session_; }
 
   void set_is_local_session(bool is_local_session) {
@@ -104,6 +111,7 @@ class ModelRuntimeContext {
   const GraphDef* graph_def_ = nullptr;
   const CallableOptions* callable_options_ = nullptr;
   tfrt::ResourceContext* resource_context_ = nullptr;
+  tensorflow::DeviceMgr* device_mgr_ = nullptr;
 
   FunctionLibraryDefinition* flib_def_ = nullptr;
 
diff --git a/tensorflow/core/tfrt/saved_model/BUILD b/tensorflow/core/tfrt/saved_model/BUILD
index 3e49778bed8416..9e85c14baef362 100644
--- a/tensorflow/core/tfrt/saved_model/BUILD
+++ b/tensorflow/core/tfrt/saved_model/BUILD
@@ -118,6 +118,7 @@ cc_library(
         "//tensorflow/core:core_cpu_base",
         "//tensorflow/core:framework",
         "//tensorflow/core:lib",
+        "//tensorflow/core/common_runtime:device_mgr",
         "//tensorflow/core/framework:function_proto_cc",
         "//tensorflow/core/framework:graph_proto_cc",
         "//tensorflow/core/framework:tensor_proto_cc",
diff --git a/tensorflow/core/tfrt/saved_model/saved_model.cc b/tensorflow/core/tfrt/saved_model/saved_model.cc
index 85bc3fabc94fad..13ff33bb9e086c 100644
--- a/tensorflow/core/tfrt/saved_model/saved_model.cc
+++ b/tensorflow/core/tfrt/saved_model/saved_model.cc
@@ -577,6 +577,8 @@ absl::StatusOr<std::unique_ptr<SavedModel>> SavedModelImpl::LoadSavedModel(
         CombineSignatureDefs(meta_graph_def.signature_def());
     model_context.set_graph_def(&meta_graph_def.graph_def());
     model_context.set_callable_options(&callable_options);
+    model_context.set_device_mgr(&fallback_state->device_manager());
+
     TF_RETURN_IF_ERROR(
         options.graph_execution_options.runtime->CreateRuntimeResources(
             model_context));
diff --git a/tensorflow/core/tfrt/tfrt_session/tfrt_session.cc b/tensorflow/core/tfrt/tfrt_session/tfrt_session.cc
index 8116c4aab1fce3..67475b0b23ed3e 100644
--- a/tensorflow/core/tfrt/tfrt_session/tfrt_session.cc
+++ b/tensorflow/core/tfrt/tfrt_session/tfrt_session.cc
@@ -35,6 +35,7 @@ limitations under the License.
 #include "Eigen/ThreadPool"  // from @eigen_archive
 #include "llvm/ADT/STLExtras.h"
 #include "tensorflow/compiler/mlir/tfrt/translate/tfrt_compile_options.h"
+#include "tensorflow/core/common_runtime/device_mgr.h"
 #include "tensorflow/core/common_runtime/local_session_selection.h"
 #include "tensorflow/core/common_runtime/process_util.h"
 #include "tensorflow/core/common_runtime/session_factory.h"
@@ -218,6 +219,7 @@ class TfrtSession : public tensorflow::Session {
         &options, /*export_dir=*/"unknown_export_dir", resource_context.get());
     // TODO(b/334641254): Offer a Session option that prunes the graph_def.
     model_context.set_graph_def(&graph);
+    model_context.set_device_mgr(&fallback_state->device_manager());
     // In the multi-host case, this prevents local Sessions from running
     // global resource creation functions.
     model_context.set_is_local_session(

From bde6f752c22ac8463971fa9b3292f4c95f26e5bb Mon Sep 17 00:00:00 2001
From: David Dunleavy <ddunleavy@google.com>
Date: Tue, 14 May 2024 16:15:59 -0700
Subject: [PATCH 144/478] Remove references to unused tags in XLA build script

PiperOrigin-RevId: 633738393
---
 third_party/xla/.kokoro/linux/build.sh | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/third_party/xla/.kokoro/linux/build.sh b/third_party/xla/.kokoro/linux/build.sh
index f9610594e5646d..f7e572c5046281 100644
--- a/third_party/xla/.kokoro/linux/build.sh
+++ b/third_party/xla/.kokoro/linux/build.sh
@@ -48,13 +48,13 @@ docker run --name xla -w /tf/xla -itd --rm \
     "$DOCKER_IMAGE" \
     bash
 
-TAGS_FILTER="-no_oss,-oss_excluded,-oss_serial"
+TAGS_FILTER="-no_oss"
 ADDITIONAL_FLAGS=""
 RBE_FLAGS=""
 TARGET_FILTERS="-@local_tsl//tsl/platform:subprocess_test -@local_tsl//tsl/platform/cloud:google_auth_provider_test -@local_tsl//tsl/platform/cloud:oauth_client_test"
 
 if is_linux_gpu_job ; then
-    TAGS_FILTER="$TAGS_FILTER,requires-gpu-nvidia,-no_gpu"
+    TAGS_FILTER="$TAGS_FILTER,requires-gpu-nvidia"
 
     # We are currently running XLA presubmits on machines with NVIDIA T4 GPUs,
     # which have a compute compatibility of 7.5. Se we filter out all the tests

From c7bb3fbf20089e8ee17405650cca4c022d491717 Mon Sep 17 00:00:00 2001
From: Kyle Lucke <klucke@google.com>
Date: Tue, 14 May 2024 16:21:09 -0700
Subject: [PATCH 145/478] Clean unnecessary include and definition from
 kernel.h.

PiperOrigin-RevId: 633739642
---
 third_party/xla/xla/stream_executor/BUILD                  | 2 +-
 third_party/xla/xla/stream_executor/gpu/BUILD              | 1 +
 .../xla/xla/stream_executor/gpu/stream_search_test.cc      | 1 +
 third_party/xla/xla/stream_executor/kernel.cc              | 7 -------
 third_party/xla/xla/stream_executor/kernel.h               | 2 --
 third_party/xla/xla/stream_executor/stream.cc              | 1 +
 6 files changed, 4 insertions(+), 10 deletions(-)

diff --git a/third_party/xla/xla/stream_executor/BUILD b/third_party/xla/xla/stream_executor/BUILD
index d41de1e0ff7032..7137b54991ad52 100644
--- a/third_party/xla/xla/stream_executor/BUILD
+++ b/third_party/xla/xla/stream_executor/BUILD
@@ -605,8 +605,8 @@ cc_library(
     deps = [
         ":device_memory",
         ":kernel_spec",
+        ":launch_dim",
         ":platform",
-        ":stream_executor_headers",
         "@com_google_absl//absl/base:core_headers",
         "@com_google_absl//absl/container:inlined_vector",
         "@com_google_absl//absl/meta:type_traits",
diff --git a/third_party/xla/xla/stream_executor/gpu/BUILD b/third_party/xla/xla/stream_executor/gpu/BUILD
index 619a5bb034ab76..b2bd00c972d26b 100644
--- a/third_party/xla/xla/stream_executor/gpu/BUILD
+++ b/third_party/xla/xla/stream_executor/gpu/BUILD
@@ -753,6 +753,7 @@ xla_cc_test(
         "//xla/stream_executor",
         "//xla/stream_executor/host:host_platform",
         "@com_google_absl//absl/status:statusor",
+        "@local_tsl//tsl/platform:statusor",
         "@local_tsl//tsl/platform:test",
         "@local_tsl//tsl/platform:test_main",
     ] + if_cuda([
diff --git a/third_party/xla/xla/stream_executor/gpu/stream_search_test.cc b/third_party/xla/xla/stream_executor/gpu/stream_search_test.cc
index 9f91c63ae0d972..c0f66159400039 100644
--- a/third_party/xla/xla/stream_executor/gpu/stream_search_test.cc
+++ b/third_party/xla/xla/stream_executor/gpu/stream_search_test.cc
@@ -17,6 +17,7 @@ limitations under the License.
 #include "xla/stream_executor/platform.h"
 #include "xla/stream_executor/platform_manager.h"
 #include "xla/stream_executor/stream_executor.h"
+#include "tsl/platform/statusor.h"
 #include "tsl/platform/test.h"
 
 namespace stream_executor {
diff --git a/third_party/xla/xla/stream_executor/kernel.cc b/third_party/xla/xla/stream_executor/kernel.cc
index 81767a7fe5025b..8f40daee24f284 100644
--- a/third_party/xla/xla/stream_executor/kernel.cc
+++ b/third_party/xla/xla/stream_executor/kernel.cc
@@ -16,19 +16,12 @@ limitations under the License.
 #include "xla/stream_executor/kernel.h"
 
 #include <cstdint>
-#include <memory>
 #include <optional>
 #include <string>
 
-#include "absl/status/statusor.h"
 #include "absl/strings/string_view.h"
 #include "absl/strings/strip.h"
-#include "xla/stream_executor/kernel_spec.h"
-#include "xla/stream_executor/platform.h"
-#include "xla/stream_executor/stream_executor.h"
 #include "tsl/platform/demangle.h"
-#include "tsl/platform/errors.h"
-#include "tsl/platform/statusor.h"
 
 namespace stream_executor {
 
diff --git a/third_party/xla/xla/stream_executor/kernel.h b/third_party/xla/xla/stream_executor/kernel.h
index 98de009e1b8a16..0f51a6ff447eb2 100644
--- a/third_party/xla/xla/stream_executor/kernel.h
+++ b/third_party/xla/xla/stream_executor/kernel.h
@@ -92,12 +92,10 @@ limitations under the License.
 #include "xla/stream_executor/kernel_spec.h"
 #include "xla/stream_executor/launch_dim.h"
 #include "tsl/platform/logging.h"
-#include "tsl/platform/statusor.h"
 
 namespace stream_executor {
 
 class Kernel;
-class StreamExecutorInterface;
 
 //===----------------------------------------------------------------------===//
 // Kernel cache config
diff --git a/third_party/xla/xla/stream_executor/stream.cc b/third_party/xla/xla/stream_executor/stream.cc
index dbb0ca213cced9..0c611032332b06 100644
--- a/third_party/xla/xla/stream_executor/stream.cc
+++ b/third_party/xla/xla/stream_executor/stream.cc
@@ -38,6 +38,7 @@ limitations under the License.
 #include "tsl/platform/errors.h"
 #include "tsl/platform/logging.h"
 #include "tsl/platform/stacktrace.h"
+#include "tsl/platform/statusor.h"
 
 namespace stream_executor {
 

From 244464154460be6e75e53ca8f69c965d256beb73 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Tue, 14 May 2024 16:46:46 -0700
Subject: [PATCH 146/478] Expands the test coverage of GetReducedIntervals(),
 since this is now the preferred way of interacting with the Memory Term
 Reducer.

PiperOrigin-RevId: 633746142
---
 .../auto_sharding_memory_test.cc              | 42 +++++++++++++++++++
 1 file changed, 42 insertions(+)

diff --git a/third_party/xla/xla/hlo/experimental/auto_sharding/auto_sharding_memory_test.cc b/third_party/xla/xla/hlo/experimental/auto_sharding/auto_sharding_memory_test.cc
index 4128d683c1dc7d..cfa7b4d2af23ee 100644
--- a/third_party/xla/xla/hlo/experimental/auto_sharding/auto_sharding_memory_test.cc
+++ b/third_party/xla/xla/hlo/experimental/auto_sharding/auto_sharding_memory_test.cc
@@ -68,10 +68,13 @@ TEST(AutoShardingMemoryTest, WithoutOverlap) {
        {   1},
        {   1},
        {   1}};
+  const std::vector<std::pair<int64_t, int64_t>> expected_reduced_intervals =
+      {{0, 2}, {3, 5}};
   const std::vector<absl::btree_set<int64_t>> expected_reduced_groups = {};
   const std::pair<int64_t, int64_t> expected_num_terms = {6, 6};
   EXPECT_EQ(num_terms, expected_num_terms);
   EXPECT_EQ(reducer.GetReducedLive(), expected_reduced_live);
+  EXPECT_EQ(reducer.GetReducedIntervals(), expected_reduced_intervals);
   EXPECT_EQ(reducer.GetReducedGroups(), expected_reduced_groups);
 }
 
@@ -100,11 +103,14 @@ TEST(AutoShardingMemoryTest, PartialOverlap) {
        {      2},
        {      2},
        {   1   }};
+  const std::vector<std::pair<int64_t, int64_t>> expected_reduced_intervals =
+      {{0, 0}, {5, 5}, {1, 4}};
   const std::vector<absl::btree_set<int64_t>> expected_reduced_groups =
       {{0, 1}};
   const std::pair<int64_t, int64_t> expected_num_terms = {10, 8};
   EXPECT_EQ(num_terms, expected_num_terms);
   EXPECT_EQ(reducer.GetReducedLive(), expected_reduced_live);
+  EXPECT_EQ(reducer.GetReducedIntervals(), expected_reduced_intervals);
   EXPECT_EQ(reducer.GetReducedGroups(), expected_reduced_groups);
 }
 
@@ -133,11 +139,14 @@ TEST(AutoShardingMemoryTest, PartialOverlapReversed) {
        {      2},
        {      2},
        {0      }};
+  const std::vector<std::pair<int64_t, int64_t>> expected_reduced_intervals =
+      {{5, 5}, {0, 0}, {1, 4}};
   const std::vector<absl::btree_set<int64_t>> expected_reduced_groups =
       {{0, 1}};
   const std::pair<int64_t, int64_t> expected_num_terms = {10, 8};
   EXPECT_EQ(num_terms, expected_num_terms);
   EXPECT_EQ(reducer.GetReducedLive(), expected_reduced_live);
+  EXPECT_EQ(reducer.GetReducedIntervals(), expected_reduced_intervals);
   EXPECT_EQ(reducer.GetReducedGroups(), expected_reduced_groups);
 }
 
@@ -165,10 +174,13 @@ TEST(AutoShardingMemoryTest, DoesNotSplitPrimitive) {
        {0, 1},
        {0, 1},
        {0   }};
+  const std::vector<std::pair<int64_t, int64_t>> expected_reduced_intervals =
+      {{0, 5}, {1, 4}};
   const std::vector<absl::btree_set<int64_t>> expected_reduced_groups = {};
   const std::pair<int64_t, int64_t> expected_num_terms = {10, 10};
   EXPECT_EQ(num_terms, expected_num_terms);
   EXPECT_EQ(reducer.GetReducedLive(), expected_reduced_live);
+  EXPECT_EQ(reducer.GetReducedIntervals(), expected_reduced_intervals);
   EXPECT_EQ(reducer.GetReducedGroups(), expected_reduced_groups);
 }
 
@@ -197,11 +209,14 @@ TEST(AutoShardingMemoryTest, OnePrimitiveVanishes) {
        {   2},
        {   2},
        {   2}};
+  const std::vector<std::pair<int64_t, int64_t>> expected_reduced_intervals =
+      {{0, 0}, {6, 0}, {1, 5}};
   const std::vector<absl::btree_set<int64_t>> expected_reduced_groups =
       {{0, 1}};
   const std::pair<int64_t, int64_t> expected_num_terms = {11, 8};
   EXPECT_EQ(num_terms, expected_num_terms);
   EXPECT_EQ(reducer.GetReducedLive(), expected_reduced_live);
+  EXPECT_EQ(reducer.GetReducedIntervals(), expected_reduced_intervals);
   EXPECT_EQ(reducer.GetReducedGroups(), expected_reduced_groups);
 }
 
@@ -230,11 +245,14 @@ TEST(AutoShardingMemoryTest, BothPrimitivesVanish) {
        {2},
        {2},
        {2}};
+  const std::vector<std::pair<int64_t, int64_t>> expected_reduced_intervals =
+      {{6, -1}, {6, -1}, {0, 5}};
   const std::vector<absl::btree_set<int64_t>> expected_reduced_groups =
       {{0, 1}};
   const std::pair<int64_t, int64_t> expected_num_terms = {12, 8};
   EXPECT_EQ(num_terms, expected_num_terms);
   EXPECT_EQ(reducer.GetReducedLive(), expected_reduced_live);
+  EXPECT_EQ(reducer.GetReducedIntervals(), expected_reduced_intervals);
   EXPECT_EQ(reducer.GetReducedGroups(), expected_reduced_groups);
 }
 
@@ -270,11 +288,14 @@ TEST(AutoShardingMemoryTest, OneGroupingPreventsAnother) {
        {1, 2   },
        {1, 2   },
        {   2   }};
+  const std::vector<std::pair<int64_t, int64_t>> expected_reduced_intervals =
+      {{5, -1}, {5, 7}, {4, 8}, {0, 4}};
   const std::vector<absl::btree_set<int64_t>> expected_reduced_groups =
       {{0, 1}};
   const std::pair<int64_t, int64_t> expected_num_terms = {18, 15};
   EXPECT_EQ(num_terms, expected_num_terms);
   EXPECT_EQ(reducer.GetReducedLive(), expected_reduced_live);
+  EXPECT_EQ(reducer.GetReducedIntervals(), expected_reduced_intervals);
   EXPECT_EQ(reducer.GetReducedGroups(), expected_reduced_groups);
 }
 
@@ -304,11 +325,14 @@ TEST(AutoShardingMemoryTest, TwoGroups) {
        {4},
        {4},
        {4}};
+  const std::vector<std::pair<int64_t, int64_t>> expected_reduced_intervals =
+      {{6, 2}, {3, -1}, {6, 2}, {0, 2}, {3, 5}};
   const std::vector<absl::btree_set<int64_t>> expected_reduced_groups =
       {{0, 1}, {0, 2}};
   const std::pair<int64_t, int64_t> expected_num_terms = {12, 10};
   EXPECT_EQ(num_terms, expected_num_terms);
   EXPECT_EQ(reducer.GetReducedLive(), expected_reduced_live);
+  EXPECT_EQ(reducer.GetReducedIntervals(), expected_reduced_intervals);
   EXPECT_EQ(reducer.GetReducedGroups(), expected_reduced_groups);
 }
 
@@ -343,11 +367,14 @@ TEST(AutoShardingMemoryTest, TwoGroupsMutuallyExclusive) {
        {      5},
        {      5},
        {   3   }};
+  const std::vector<std::pair<int64_t, int64_t>> expected_reduced_intervals =
+      {{0, 0}, {4, 0}, {7, 3}, {7, 7}, {1, 3}, {4, 6}};
   const std::vector<absl::btree_set<int64_t>> expected_reduced_groups =
       {{0, 1}, {2, 3}};
   const std::pair<int64_t, int64_t> expected_num_terms = {14, 12};
   EXPECT_EQ(num_terms, expected_num_terms);
   EXPECT_EQ(reducer.GetReducedLive(), expected_reduced_live);
+  EXPECT_EQ(reducer.GetReducedIntervals(), expected_reduced_intervals);
   EXPECT_EQ(reducer.GetReducedGroups(), expected_reduced_groups);
 }
 
@@ -367,10 +394,13 @@ TEST(AutoShardingMemoryTest, MergingPrimitivesWouldNotReduceTerms) {
   const std::vector<std::vector<int64_t>> expected_reduced_live =
       {{0, 1},
        {0, 1}};
+  const std::vector<std::pair<int64_t, int64_t>> expected_reduced_intervals =
+      {{0, 1}, {0, 1}};
   const std::vector<absl::btree_set<int64_t>> expected_reduced_groups = {};
   const std::pair<int64_t, int64_t> expected_num_terms = {4, 4};
   EXPECT_EQ(num_terms, expected_num_terms);
   EXPECT_EQ(reducer.GetReducedLive(), expected_reduced_live);
+  EXPECT_EQ(reducer.GetReducedIntervals(), expected_reduced_intervals);
   EXPECT_EQ(reducer.GetReducedGroups(), expected_reduced_groups);
 }
 
@@ -400,11 +430,14 @@ TEST(AutoShardingMemoryTest, AllPrimitivesVanish) {
        {3},
        {3},
        {3}};
+  const std::vector<std::pair<int64_t, int64_t>> expected_reduced_intervals =
+      {{6, -1}, {6, -1}, {6, -1}, {0, 5}};
   const std::vector<absl::btree_set<int64_t>> expected_reduced_groups =
       {{0, 1, 2}};
   const std::pair<int64_t, int64_t> expected_num_terms = {18, 9};
   EXPECT_EQ(num_terms, expected_num_terms);
   EXPECT_EQ(reducer.GetReducedLive(), expected_reduced_live);
+  EXPECT_EQ(reducer.GetReducedIntervals(), expected_reduced_intervals);
   EXPECT_EQ(reducer.GetReducedGroups(), expected_reduced_groups);
 }
 
@@ -442,11 +475,14 @@ TEST(AutoShardingMemoryTest, MergingGroupsWouldNotReduceTerms) {
        {4, 5},
        {   5},
        {   5}};
+  const std::vector<std::pair<int64_t, int64_t>> expected_reduced_intervals =
+      {{7, -1}, {7, -1}, {9, 2}, {9, 2}, {0, 6}, {3, 8}};
   const std::vector<absl::btree_set<int64_t>> expected_reduced_groups =
       {{0, 1}, {2, 3}};
   const std::pair<int64_t, int64_t> expected_num_terms = {26, 17};
   EXPECT_EQ(num_terms, expected_num_terms);
   EXPECT_EQ(reducer.GetReducedLive(), expected_reduced_live);
+  EXPECT_EQ(reducer.GetReducedIntervals(), expected_reduced_intervals);
   EXPECT_EQ(reducer.GetReducedGroups(), expected_reduced_groups);
 }
 
@@ -493,11 +529,14 @@ TEST(AutoShardingMemoryTest, ExampleFromDocumentation) {
        {      5},
        {      5},
        {   3   }};
+  const std::vector<std::pair<int64_t, int64_t>> expected_reduced_intervals =
+      {{0, 0}, {10, 0}, {13, 4}, {13, 13}, {1, 4}, {10, 12}, {5, 9}};
   const std::vector<absl::btree_set<int64_t>> expected_reduced_groups =
       {{0, 1}, {2, 3}, {0, 1, 2, 3}};
   const std::pair<int64_t, int64_t> expected_num_terms = {36, 22};
   EXPECT_EQ(num_terms, expected_num_terms);
   EXPECT_EQ(reducer.GetReducedLive(), expected_reduced_live);
+  EXPECT_EQ(reducer.GetReducedIntervals(), expected_reduced_intervals);
   EXPECT_EQ(reducer.GetReducedGroups(), expected_reduced_groups);
 }
 
@@ -523,11 +562,14 @@ TEST(AutoShardingMemoryTest, MergesWithRightmost) {
        {      3},
        {      3},
        {1, 2   }};
+  const std::vector<std::pair<int64_t, int64_t>> expected_reduced_intervals =
+      {{3, -1}, {3, 3}, {3, 3}, {0, 2}};
   const std::vector<absl::btree_set<int64_t>> expected_reduced_groups =
       {{0, 2}};
   const std::pair<int64_t, int64_t> expected_num_terms = {8, 7};
   EXPECT_EQ(num_terms, expected_num_terms);
   EXPECT_EQ(reducer.GetReducedLive(), expected_reduced_live);
+  EXPECT_EQ(reducer.GetReducedIntervals(), expected_reduced_intervals);
   EXPECT_EQ(reducer.GetReducedGroups(), expected_reduced_groups);
 }
 

From 5938029cd2670e6aa29953de9aea7a88fb598fa0 Mon Sep 17 00:00:00 2001
From: Kuy Mainwaring <kuym@google.com>
Date: Tue, 14 May 2024 16:52:32 -0700
Subject: [PATCH 147/478] [XLA:GPU] Clang-tidy cleanup for xla/service/gpu/...
 for const-correctness

PiperOrigin-RevId: 633747545
---
 third_party/xla/xla/service/gpu/cublas_pad_for_gemms.cc      | 2 +-
 third_party/xla/xla/service/gpu/cudnn_fused_mha_rewriter.cc  | 5 +++--
 .../xla/xla/service/gpu/cudnn_fused_mha_transpose_fusion.cc  | 2 --
 third_party/xla/xla/service/gpu/gpu_conv_rewriter.cc         | 2 +-
 .../xla/service/gpu/runtime/nccl_collective_permute_thunk.cc | 4 ++--
 5 files changed, 7 insertions(+), 8 deletions(-)

diff --git a/third_party/xla/xla/service/gpu/cublas_pad_for_gemms.cc b/third_party/xla/xla/service/gpu/cublas_pad_for_gemms.cc
index 050f219d12b6c8..a90981f1c932fb 100644
--- a/third_party/xla/xla/service/gpu/cublas_pad_for_gemms.cc
+++ b/third_party/xla/xla/service/gpu/cublas_pad_for_gemms.cc
@@ -133,7 +133,7 @@ namespace {
 // We need this check because PadForGemm works in the assumption that
 // the dot instruction is canonicalized.
 bool CheckCanonical(HloDotInstruction* dot) {
-  auto dimension_numbers = dot->dot_dimension_numbers();
+  const auto& dimension_numbers = dot->dot_dimension_numbers();
 
   if (dimension_numbers.lhs_batch_dimensions_size() + 2 !=
           dot->operand(0)->shape().rank() ||
diff --git a/third_party/xla/xla/service/gpu/cudnn_fused_mha_rewriter.cc b/third_party/xla/xla/service/gpu/cudnn_fused_mha_rewriter.cc
index b25cfe8d7cbe55..0dea3331bc11cd 100644
--- a/third_party/xla/xla/service/gpu/cudnn_fused_mha_rewriter.cc
+++ b/third_party/xla/xla/service/gpu/cudnn_fused_mha_rewriter.cc
@@ -1344,7 +1344,7 @@ absl::StatusOr<HloInstruction*> FuseFwdMultiHeadedAttentionBlock(
         bcast_dimensions.push_back(dim - starting_index);
       }
 
-      Shape bcast_shape = bmm_1->shape();
+      const Shape& bcast_shape = bmm_1->shape();
       bias = comp->AddInstruction(HloInstruction::CreateBroadcast(
           bcast_shape, original_bias, bcast_dimensions));
     }
@@ -1402,7 +1402,8 @@ absl::StatusOr<bool> FuseBwdMultiHeadedAttentionBlock(
 
   TF_ASSIGN_OR_RETURN(GpuBackendConfig gpu_config,
                       fwd_fmha_call->backend_config<GpuBackendConfig>());
-  CudnnfMHABackendConfig fwd_config = gpu_config.cudnn_fmha_backend_config();
+  const CudnnfMHABackendConfig& fwd_config =
+      gpu_config.cudnn_fmha_backend_config();
   bool is_causal_mask =
       fwd_config.mask_type() == CudnnfMHABackendConfig::CAUSAL;
   CudnnfMHABackendConfig bwd_fmha_config;
diff --git a/third_party/xla/xla/service/gpu/cudnn_fused_mha_transpose_fusion.cc b/third_party/xla/xla/service/gpu/cudnn_fused_mha_transpose_fusion.cc
index 81074a29638b9e..665cc0bf824383 100644
--- a/third_party/xla/xla/service/gpu/cudnn_fused_mha_transpose_fusion.cc
+++ b/third_party/xla/xla/service/gpu/cudnn_fused_mha_transpose_fusion.cc
@@ -463,8 +463,6 @@ absl::StatusOr<bool> FusePrologueTransposeWithcuDNNFMHA(HloComputation* comp) {
       // should have the same layout of O
       TF_ASSIGN_OR_RETURN(auto gpu_config,
                           fmha->backend_config<GpuBackendConfig>());
-      const CudnnfMHABackendConfig config =
-          gpu_config.cudnn_fmha_backend_config();
       if (changed && VLOG_IS_ON(2)) {
         VLOG(2) << "After CudnnFusedMHATransposeFusion Arg 4: \n"
                 << comp->parent()->ToString();
diff --git a/third_party/xla/xla/service/gpu/gpu_conv_rewriter.cc b/third_party/xla/xla/service/gpu/gpu_conv_rewriter.cc
index 8ba1af16338870..b8dbf25879db5a 100644
--- a/third_party/xla/xla/service/gpu/gpu_conv_rewriter.cc
+++ b/third_party/xla/xla/service/gpu/gpu_conv_rewriter.cc
@@ -65,7 +65,7 @@ bool MaybeConv1dToConv2d(HloInstruction* conv) {
   if (reshape_degenerate.has_value() &&
       reshape_degenerate->deleted_dimensions.empty() &&
       reshape_degenerate->inserted_dimensions.size() == 1) {
-    auto dnums = conv->convolution_dimension_numbers();
+    const auto& dnums = conv->convolution_dimension_numbers();
     for (auto dim : dnums.kernel_spatial_dimensions()) {
       if (dim == reshape_degenerate->inserted_dimensions[0]) {
         return true;
diff --git a/third_party/xla/xla/service/gpu/runtime/nccl_collective_permute_thunk.cc b/third_party/xla/xla/service/gpu/runtime/nccl_collective_permute_thunk.cc
index 804c30b42e3690..02a8a583d754ef 100644
--- a/third_party/xla/xla/service/gpu/runtime/nccl_collective_permute_thunk.cc
+++ b/third_party/xla/xla/service/gpu/runtime/nccl_collective_permute_thunk.cc
@@ -89,7 +89,7 @@ NcclCollectivePermuteStartThunk::NcclCollectivePermuteStartThunk(
     replica_group.add_replica_ids(i);
   }
 
-  const std::vector<std::pair<int64_t, int64_t>> source_target_pairs =
+  const std::vector<std::pair<int64_t, int64_t>>& source_target_pairs =
       instr->source_target_pairs();
 
   for (const std::pair<int64_t, int64_t>& source_target : source_target_pairs) {
@@ -110,7 +110,7 @@ NcclCollectivePermuteStartThunk::NcclCollectivePermuteStartThunk(
     int64_t partition_count) {
   // The collective permute is degenerate if all source-target pairs are
   // identity, and all the IDs appear in the list.
-  const std::vector<std::pair<int64_t, int64_t>> source_target_pairs =
+  const std::vector<std::pair<int64_t, int64_t>>& source_target_pairs =
       instr->source_target_pairs();
   // Each ID can appear only once as a source and as a target. So if all pairs
   // are identity, all IDs must appear in the list is the size == number of

From 5f64c63ba91479f6343be97ee4c60cb7eadfb4dd Mon Sep 17 00:00:00 2001
From: Fangrui Song <maskray@google.com>
Date: Tue, 14 May 2024 16:57:21 -0700
Subject: [PATCH 148/478] Integrate LLVM at llvm/llvm-project@f89b1b8a6806

Updates LLVM usage to match
[f89b1b8a6806](https://github.com/llvm/llvm-project/commit/f89b1b8a6806)

PiperOrigin-RevId: 633748740
---
 third_party/llvm/workspace.bzl | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/third_party/llvm/workspace.bzl b/third_party/llvm/workspace.bzl
index 8d2c0c7e5a2254..a7d6c7b299393d 100644
--- a/third_party/llvm/workspace.bzl
+++ b/third_party/llvm/workspace.bzl
@@ -4,8 +4,8 @@ load("//third_party:repo.bzl", "tf_http_archive")
 
 def repo(name):
     """Imports LLVM."""
-    LLVM_COMMIT = "e6d3a4212d20b49a8e63f11fedea79cccf261479"
-    LLVM_SHA256 = "a9581601d91836d182180e35a698550b9c8257eacead9e5cc7ec956128200fce"
+    LLVM_COMMIT = "f89b1b8a68065c4b880417abb0563bce21399b52"
+    LLVM_SHA256 = "eba609768a9fbce79c7096d156fc96e999296cc61aca9ca736df7a9c5236215c"
 
     tf_http_archive(
         name = name,

From 161d89a2856b920645467a7243843461db121195 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Tue, 14 May 2024 17:19:42 -0700
Subject: [PATCH 149/478] Update ops-related pbtxt files.

PiperOrigin-RevId: 633754444
---
 tensorflow/core/ops/ops.pbtxt | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorflow/core/ops/ops.pbtxt b/tensorflow/core/ops/ops.pbtxt
index cdc0191b8ad6c9..4b2857a8994764 100644
--- a/tensorflow/core/ops/ops.pbtxt
+++ b/tensorflow/core/ops/ops.pbtxt
@@ -1,4 +1,4 @@
-go/debugproto   
+go/debugonly  
 op {
   name: "Abort"
   attr {

From 102e08c2c67113289b4ac1d60fd8c16ce2f401b9 Mon Sep 17 00:00:00 2001
From: Yang Chen <yangchen@google.com>
Date: Tue, 14 May 2024 17:52:44 -0700
Subject: [PATCH 150/478] #tf-data Support random access for the unbatch
 dataset.

PiperOrigin-RevId: 633761597
---
 .../core/kernels/data/experimental/BUILD      |  3 +-
 .../data/experimental/unbatch_dataset_op.cc   | 28 +++++++++++++++++
 tensorflow/python/data/kernel_tests/BUILD     |  1 +
 .../python/data/kernel_tests/unbatch_test.py  | 30 +++++++++++++++++++
 4 files changed, 61 insertions(+), 1 deletion(-)

diff --git a/tensorflow/core/kernels/data/experimental/BUILD b/tensorflow/core/kernels/data/experimental/BUILD
index f24cff66f1ef98..5adbe7abe25ee8 100644
--- a/tensorflow/core/kernels/data/experimental/BUILD
+++ b/tensorflow/core/kernels/data/experimental/BUILD
@@ -922,8 +922,9 @@ tf_kernel_library(
         "//tensorflow/core:framework_internal",
         "//tensorflow/core:lib",
         "//tensorflow/core:lib_internal",
-        "//tensorflow/core/data:dataset_utils",
         "//tensorflow/core/framework:dataset_options_proto_cc",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/strings",
         "@local_tsl//tsl/platform:errors",
         "@local_tsl//tsl/platform:mutex",
         "@local_tsl//tsl/platform:status",
diff --git a/tensorflow/core/kernels/data/experimental/unbatch_dataset_op.cc b/tensorflow/core/kernels/data/experimental/unbatch_dataset_op.cc
index 4b4433ca27b27d..3cc332024328e8 100644
--- a/tensorflow/core/kernels/data/experimental/unbatch_dataset_op.cc
+++ b/tensorflow/core/kernels/data/experimental/unbatch_dataset_op.cc
@@ -19,6 +19,8 @@ limitations under the License.
 #include <utility>
 #include <vector>
 
+#include "absl/status/status.h"
+#include "absl/strings/str_cat.h"
 #include "tensorflow/core/framework/dataset.h"
 #include "tensorflow/core/framework/dataset_options.pb.h"
 #include "tensorflow/core/framework/model.h"
@@ -122,6 +124,32 @@ class UnbatchDatasetOp : public UnaryDatasetOpKernel {
       return input_->CheckExternalState();
     }
 
+    absl::Status Get(OpKernelContext* ctx, int64_t index,
+                     std::vector<Tensor>* out_tensors) const override {
+      TF_RETURN_IF_ERROR(CheckRandomAccessCompatible(index));
+      if (batch_size_ <= 0) {
+        return absl::FailedPreconditionError(absl::StrCat(
+            "Random access for the `unbatch` dataset requires a known batch "
+            "size. Got ",
+            batch_size_, "."));
+      }
+
+      const int64_t input_index = index / batch_size_;
+      const int64_t input_offset = index % batch_size_;
+      std::vector<Tensor> input_tensors;
+      TF_RETURN_IF_ERROR(input_->Get(ctx, input_index, &input_tensors));
+      for (int64_t i = 0; i < input_tensors.size(); ++i) {
+        const DataType& dtype = input_tensors[i].dtype();
+        TensorShape shape = input_tensors[i].shape();
+        shape.RemoveDim(0);
+
+        out_tensors->emplace_back(ctx->get_allocator({}), dtype, shape);
+        TF_RETURN_IF_ERROR(batch_util::MaybeMoveSliceToElement(
+            &input_tensors[i], &out_tensors->back(), input_offset));
+      }
+      return absl::OkStatus();
+    }
+
    protected:
     Status AsGraphDefInternal(SerializationContext* ctx,
                               DatasetGraphDefBuilder* b,
diff --git a/tensorflow/python/data/kernel_tests/BUILD b/tensorflow/python/data/kernel_tests/BUILD
index dc0d785798622b..da147e0f39435e 100644
--- a/tensorflow/python/data/kernel_tests/BUILD
+++ b/tensorflow/python/data/kernel_tests/BUILD
@@ -1360,6 +1360,7 @@ tf_py_strict_test(
     deps = [
         ":checkpoint_test_base",
         ":test_base",
+        "//tensorflow/python/data/experimental/ops:random_access",
         "//tensorflow/python/data/ops:dataset_ops",
         "//tensorflow/python/data/ops:options",
         "//tensorflow/python/framework:combinations",
diff --git a/tensorflow/python/data/kernel_tests/unbatch_test.py b/tensorflow/python/data/kernel_tests/unbatch_test.py
index d86075cd2b73bf..a65c6a2f2f393d 100644
--- a/tensorflow/python/data/kernel_tests/unbatch_test.py
+++ b/tensorflow/python/data/kernel_tests/unbatch_test.py
@@ -16,6 +16,7 @@
 from absl.testing import parameterized
 import numpy as np
 
+from tensorflow.python.data.experimental.ops import random_access
 from tensorflow.python.data.kernel_tests import checkpoint_test_base
 from tensorflow.python.data.kernel_tests import test_base
 from tensorflow.python.data.ops import dataset_ops
@@ -262,5 +263,34 @@ def test(self, verify_fn, symbolic_checkpoint):
         num_outputs)
 
 
+class UnbatchRandomAccessTest(test_base.DatasetTestBase,
+                              parameterized.TestCase):
+  @combinations.generate(test_base.default_test_combinations())
+  def test(self):
+    dataset = dataset_ops.Dataset.range(10)
+    dataset = dataset.batch(4, drop_remainder=True)
+    dataset = dataset.unbatch()
+    for i in range(8):
+      self.assertEqual(self.evaluate(random_access.at(dataset, i)), i)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testNotDropRemainder(self):
+    dataset = dataset_ops.Dataset.range(10)
+    dataset = dataset.batch(4, drop_remainder=False)
+    dataset = dataset.unbatch()
+    with self.assertRaises(errors.FailedPreconditionError):
+      self.evaluate(random_access.at(dataset, 0))
+
+  @combinations.generate(
+      combinations.times(test_base.default_test_combinations(),
+                         combinations.combine(index=[-1, 100])))
+  def testInvalidIndex(self, index):
+    dataset = dataset_ops.Dataset.range(10)
+    dataset = dataset.batch(4, drop_remainder=True)
+    dataset = dataset.unbatch()
+    with self.assertRaises(errors.OutOfRangeError):
+      self.evaluate(random_access.at(dataset, index=index))
+
+
 if __name__ == "__main__":
   test.main()

From f6b72954734f8304bfb83228bd8406a3ba3394f4 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Tue, 14 May 2024 19:17:56 -0700
Subject: [PATCH 151/478] Update ops-related pbtxt files.

PiperOrigin-RevId: 633777654
---
 tensorflow/core/ops/ops.pbtxt | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorflow/core/ops/ops.pbtxt b/tensorflow/core/ops/ops.pbtxt
index 4b2857a8994764..9dc272e0a4f6ed 100644
--- a/tensorflow/core/ops/ops.pbtxt
+++ b/tensorflow/core/ops/ops.pbtxt
@@ -1,4 +1,4 @@
-go/debugonly  
+go/debugproto    
 op {
   name: "Abort"
   attr {

From 4f02320c5b379ff7a9f4cc9b665f1937a27d9a47 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Tue, 14 May 2024 21:17:42 -0700
Subject: [PATCH 152/478] Update ops-related pbtxt files.

PiperOrigin-RevId: 633801239
---
 tensorflow/core/ops/ops.pbtxt | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorflow/core/ops/ops.pbtxt b/tensorflow/core/ops/ops.pbtxt
index 9dc272e0a4f6ed..4b2857a8994764 100644
--- a/tensorflow/core/ops/ops.pbtxt
+++ b/tensorflow/core/ops/ops.pbtxt
@@ -1,4 +1,4 @@
-go/debugproto    
+go/debugonly  
 op {
   name: "Abort"
   attr {

From f34cf56711a673edbf082d42a318626815a40a86 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Tue, 14 May 2024 22:04:15 -0700
Subject: [PATCH 153/478] Automated Code Change

PiperOrigin-RevId: 633809834
---
 third_party/xla/xla/BUILD                  | 5 ++++-
 third_party/xla/xla/frontend_attributes.cc | 3 +++
 2 files changed, 7 insertions(+), 1 deletion(-)

diff --git a/third_party/xla/xla/BUILD b/third_party/xla/xla/BUILD
index 2219ec3a37e78c..304ed5a4d2b031 100644
--- a/third_party/xla/xla/BUILD
+++ b/third_party/xla/xla/BUILD
@@ -224,7 +224,10 @@ cc_library(
         "frontend_attributes.h",
     ],
     visibility = internal_visibility([":friends"]),
-    deps = ["//xla/hlo/ir:hlo"],
+    deps = [
+        ":xla_data_proto_cc",
+        "//xla/hlo/ir:hlo",
+    ],
 )
 
 cc_library(
diff --git a/third_party/xla/xla/frontend_attributes.cc b/third_party/xla/xla/frontend_attributes.cc
index 8831040f89c15c..53ee0a4d1f1643 100644
--- a/third_party/xla/xla/frontend_attributes.cc
+++ b/third_party/xla/xla/frontend_attributes.cc
@@ -14,6 +14,9 @@ limitations under the License.
 ==============================================================================*/
 #include "xla/frontend_attributes.h"
 
+#include "xla/hlo/ir/hlo_instruction.h"
+#include "xla/xla_data.pb.h"
+
 namespace xla {
 
 void SetDisjointReadWriteRegionsAttr(HloInstruction* instruction) {

From bf42cf1551e1cfd0226ae9db1e79818f720eb317 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Tue, 14 May 2024 22:41:38 -0700
Subject: [PATCH 154/478] Automated Code Change

PiperOrigin-RevId: 633817088
---
 tensorflow/tools/benchmark/BUILD               |  8 +++++++-
 tensorflow/tools/benchmark/benchmark_model.cc  | 17 +++++++++++++----
 tensorflow/tools/benchmark/benchmark_model.h   |  4 ++++
 .../tools/benchmark/benchmark_model_test.cc    | 18 +++++++++++++-----
 4 files changed, 37 insertions(+), 10 deletions(-)

diff --git a/tensorflow/tools/benchmark/BUILD b/tensorflow/tools/benchmark/BUILD
index dbac33541f0872..cdfb36c3a3eae2 100644
--- a/tensorflow/tools/benchmark/BUILD
+++ b/tensorflow/tools/benchmark/BUILD
@@ -40,7 +40,11 @@ cc_library(
             "//tensorflow/core:tensorflow",
             "//tensorflow/core:test",
         ],
-    }),
+    }) + [
+        "//tensorflow/core/platform:numbers",
+        "//tensorflow/core/util:stats_calculator_portable",
+        "@local_tsl//tsl/platform:status",
+    ],
 )
 
 tf_cc_test(
@@ -50,9 +54,11 @@ tf_cc_test(
     deps = [
         ":benchmark_model_lib",
         "//tensorflow/cc:cc_ops",
+        "//tensorflow/cc:scope",
         "//tensorflow/core:core_cpu",
         "//tensorflow/core:framework",
         "//tensorflow/core:lib",
+        "//tensorflow/core:protos_all_cc",
         "//tensorflow/core:test",
         "//tensorflow/core:test_main",
         "//tensorflow/core:testlib",
diff --git a/tensorflow/tools/benchmark/benchmark_model.cc b/tensorflow/tools/benchmark/benchmark_model.cc
index f0a97c2170a3fc..b687d725dccf06 100644
--- a/tensorflow/tools/benchmark/benchmark_model.cc
+++ b/tensorflow/tools/benchmark/benchmark_model.cc
@@ -26,25 +26,34 @@ limitations under the License.
 #include <unordered_set>
 #include <vector>
 
-#include "tensorflow/core/common_runtime/graph_constructor.h"
 #include "tensorflow/core/framework/graph.pb.h"
 #include "tensorflow/core/framework/node_def.pb.h"
+#include "tensorflow/core/framework/numeric_types.h"
 #include "tensorflow/core/framework/step_stats.pb.h"
 #include "tensorflow/core/framework/tensor.h"
-#include "tensorflow/core/graph/algorithm.h"
+#include "tensorflow/core/framework/tensor_shape.h"
+#include "tensorflow/core/framework/types.h"
+#include "tensorflow/core/framework/types.pb.h"
 #include "tensorflow/core/graph/graph.h"
-#include "tensorflow/core/lib/strings/numbers.h"
 #include "tensorflow/core/lib/strings/str_util.h"
 #include "tensorflow/core/lib/strings/strcat.h"
 #include "tensorflow/core/platform/env.h"
 #include "tensorflow/core/platform/init_main.h"
 #include "tensorflow/core/platform/logging.h"
-#include "tensorflow/core/platform/platform.h"
+#include "tensorflow/core/platform/numbers.h"
+#include "tensorflow/core/platform/status.h"
+#include "tensorflow/core/platform/tstring.h"
 #include "tensorflow/core/platform/types.h"
+#include "tensorflow/core/protobuf/config.pb.h"
 #include "tensorflow/core/public/session.h"
+#include "tensorflow/core/public/session_options.h"
 #include "tensorflow/core/util/command_line_flags.h"
 #include "tensorflow/core/util/reporter.h"
 #include "tensorflow/core/util/stat_summarizer.h"
+#include "tensorflow/core/util/stat_summarizer_options.h"
+#include "tensorflow/core/util/stats_calculator.h"
+#include "tsl/platform/errors.h"
+#include "tsl/platform/status.h"
 
 namespace tensorflow {
 namespace benchmark_model {
diff --git a/tensorflow/tools/benchmark/benchmark_model.h b/tensorflow/tools/benchmark/benchmark_model.h
index 8211605ace6dea..e983ea4167d740 100644
--- a/tensorflow/tools/benchmark/benchmark_model.h
+++ b/tensorflow/tools/benchmark/benchmark_model.h
@@ -16,6 +16,10 @@ limitations under the License.
 #ifndef TENSORFLOW_TOOLS_BENCHMARK_BENCHMARK_MODEL_H_
 #define TENSORFLOW_TOOLS_BENCHMARK_BENCHMARK_MODEL_H_
 
+#include "tensorflow/core/framework/tensor_shape.h"
+#include "tensorflow/core/framework/types.pb.h"
+#include "tensorflow/core/platform/status.h"
+#include "tensorflow/core/platform/types.h"
 #include "tensorflow/core/public/session.h"
 #include "tensorflow/core/util/stat_summarizer.h"
 
diff --git a/tensorflow/tools/benchmark/benchmark_model_test.cc b/tensorflow/tools/benchmark/benchmark_model_test.cc
index c2a953ec91f7a7..c6e42840c6a689 100644
--- a/tensorflow/tools/benchmark/benchmark_model_test.cc
+++ b/tensorflow/tools/benchmark/benchmark_model_test.cc
@@ -15,13 +15,21 @@ limitations under the License.
 
 #include "tensorflow/tools/benchmark/benchmark_model.h"
 
-#include "tensorflow/cc/ops/standard_ops.h"
+#include "tensorflow/cc/framework/scope.h"
+#include "tensorflow/cc/ops/array_ops.h"
+#include "tensorflow/cc/ops/math_ops.h"
+#include "tensorflow/core/framework/graph.pb.h"
+#include "tensorflow/core/framework/tensor.h"
+#include "tensorflow/core/framework/tensor_shape.h"
 #include "tensorflow/core/framework/tensor_testutil.h"
-#include "tensorflow/core/graph/graph_def_builder.h"
-#include "tensorflow/core/lib/core/status_test_util.h"
-#include "tensorflow/core/lib/io/path.h"
+#include "tensorflow/core/framework/types.pb.h"
+#include "tensorflow/core/platform/env.h"
+#include "tensorflow/core/platform/path.h"
 #include "tensorflow/core/platform/test.h"
-#include "tensorflow/core/platform/test_benchmark.h"
+#include "tensorflow/core/platform/types.h"
+#include "tensorflow/core/public/session.h"
+#include "tensorflow/core/util/stat_summarizer.h"
+#include "tsl/lib/core/status_test_util.h"
 
 namespace tensorflow {
 namespace {

From 17de42e5c78c8944a1643101187efd461b07cf0d Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Tue, 14 May 2024 23:17:37 -0700
Subject: [PATCH 155/478] Update ops-related pbtxt files.

PiperOrigin-RevId: 633823815
---
 tensorflow/core/ops/ops.pbtxt | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorflow/core/ops/ops.pbtxt b/tensorflow/core/ops/ops.pbtxt
index 4b2857a8994764..846207488b92f3 100644
--- a/tensorflow/core/ops/ops.pbtxt
+++ b/tensorflow/core/ops/ops.pbtxt
@@ -1,4 +1,4 @@
-go/debugonly  
+go/debugstr   
 op {
   name: "Abort"
   attr {

From 2206fdfecea8507224029b670fea938b44d54395 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Tue, 14 May 2024 23:31:50 -0700
Subject: [PATCH 156/478] Automated Code Change

PiperOrigin-RevId: 633826568
---
 third_party/xla/xla/tsl/concurrency/concurrent_vector_test.cc | 1 +
 1 file changed, 1 insertion(+)

diff --git a/third_party/xla/xla/tsl/concurrency/concurrent_vector_test.cc b/third_party/xla/xla/tsl/concurrency/concurrent_vector_test.cc
index 9d384ad8b80d08..5106909ce06146 100644
--- a/third_party/xla/xla/tsl/concurrency/concurrent_vector_test.cc
+++ b/third_party/xla/xla/tsl/concurrency/concurrent_vector_test.cc
@@ -18,6 +18,7 @@ limitations under the License.
 #include <algorithm>
 #include <vector>
 
+#include "tsl/platform/env.h"
 #include "tsl/platform/test.h"
 #include "tsl/platform/threadpool.h"
 

From 30dc4d79b93c1af722d07c7e73251e2af407eab4 Mon Sep 17 00:00:00 2001
From: Dateng Lin <datenglin@google.com>
Date: Tue, 14 May 2024 23:47:32 -0700
Subject: [PATCH 157/478] Fixed the test failure.

PiperOrigin-RevId: 633829720
---
 tensorflow/core/tfrt/ifrt/ifrt_model_context.h | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorflow/core/tfrt/ifrt/ifrt_model_context.h b/tensorflow/core/tfrt/ifrt/ifrt_model_context.h
index 20af882d634a7a..76813a3a80937a 100644
--- a/tensorflow/core/tfrt/ifrt/ifrt_model_context.h
+++ b/tensorflow/core/tfrt/ifrt/ifrt_model_context.h
@@ -125,7 +125,7 @@ class IfrtModelContext {
   IfrtServingCoreSelector* ifrt_serving_core_selector_;  // May be nullptr
   const tsl::thread::ThreadPool& thread_pool_;
 
-  tensorflow::DeviceMgr* device_mgr_;  // Not owned.
+  tensorflow::DeviceMgr* device_mgr_ = nullptr;  // Not owned.
   tensorflow::XlaHelpers::ShapeRepresentationFn shape_representation_fn_ =
       tensorflow::IdentityShapeRepresentationFn();
 

From 6c46714176e0cb907de346758d5ffd19bf0382af Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Wed, 15 May 2024 00:13:28 -0700
Subject: [PATCH 158/478] Automated Code Change

PiperOrigin-RevId: 633835586
---
 tensorflow/core/platform/file_system_test.cc | 18 +++++++++---------
 1 file changed, 9 insertions(+), 9 deletions(-)

diff --git a/tensorflow/core/platform/file_system_test.cc b/tensorflow/core/platform/file_system_test.cc
index 2ecc93f11ab507..1c4d978f13ca6c 100644
--- a/tensorflow/core/platform/file_system_test.cc
+++ b/tensorflow/core/platform/file_system_test.cc
@@ -38,7 +38,7 @@ class InterPlanetaryFileSystem : public NullFileSystem {
     string parsed_path;
     ParsePath(fname, &parsed_path);
     if (BodyExists(parsed_path)) {
-      return OkStatus();
+      return absl::OkStatus();
     }
     return Status(absl::StatusCode::kNotFound, "File does not exist");
   }
@@ -58,13 +58,13 @@ class InterPlanetaryFileSystem : public NullFileSystem {
       return Status(absl::StatusCode::kInvalidArgument, "Bad dirname");
     }
     if (split_path.empty()) {
-      return OkStatus();
+      return absl::OkStatus();
     }
     if (split_path.size() == 1) {
       celestial_bodies_[""].insert(parsed_path);
       celestial_bodies_.insert(
           std::pair<string, std::set<string>>(parsed_path, {}));
-      return OkStatus();
+      return absl::OkStatus();
     }
     if (split_path.size() == 2) {
       if (!BodyExists(split_path[0])) {
@@ -74,7 +74,7 @@ class InterPlanetaryFileSystem : public NullFileSystem {
       celestial_bodies_[split_path[0]].insert(split_path[1]);
       celestial_bodies_.insert(
           std::pair<string, std::set<string>>(parsed_path, {}));
-      return OkStatus();
+      return absl::OkStatus();
     }
     if (split_path.size() == 3) {
       const string& parent_path = this->JoinPath(split_path[0], split_path[1]);
@@ -85,7 +85,7 @@ class InterPlanetaryFileSystem : public NullFileSystem {
       celestial_bodies_[parent_path].insert(split_path[2]);
       celestial_bodies_.insert(
           std::pair<string, std::set<string>>(parsed_path, {}));
-      return OkStatus();
+      return absl::OkStatus();
     }
     return Status(absl::StatusCode::kFailedPrecondition, "Failed to create");
   }
@@ -102,7 +102,7 @@ class InterPlanetaryFileSystem : public NullFileSystem {
       return Status(absl::StatusCode::kFailedPrecondition, "Not a dir");
     }
     if (celestial_bodies_.find(parsed_path) != celestial_bodies_.end()) {
-      return OkStatus();
+      return absl::OkStatus();
     }
     return Status(absl::StatusCode::kFailedPrecondition, "Not a dir");
   }
@@ -114,7 +114,7 @@ class InterPlanetaryFileSystem : public NullFileSystem {
     ParsePath(dir, &parsed_path);
     result->insert(result->begin(), celestial_bodies_[parsed_path].begin(),
                    celestial_bodies_[parsed_path].end());
-    return OkStatus();
+    return absl::OkStatus();
   }
 
  private:
@@ -279,7 +279,7 @@ class TestFileSystem : public NullFileSystem {
   // Only allow for a single root directory.
   Status IsDirectory(const string& dirname, TransactionToken* token) override {
     if (dirname == "." || dirname.empty()) {
-      return OkStatus();
+      return absl::OkStatus();
     }
     return Status(absl::StatusCode::kFailedPrecondition, "Not a dir");
   }
@@ -290,7 +290,7 @@ class TestFileSystem : public NullFileSystem {
     if (dir == "." || dir.empty()) {
       result->push_back("test");
     }
-    return OkStatus();
+    return absl::OkStatus();
   }
 };
 

From c42c1478f0e2e36861f3dda1d320d136a74e1ba6 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Wed, 15 May 2024 00:14:36 -0700
Subject: [PATCH 159/478] Use min subgroup size when 16 is not supported

PiperOrigin-RevId: 633835830
---
 tensorflow/lite/delegates/gpu/common/BUILD           |  1 +
 tensorflow/lite/delegates/gpu/common/gpu_info.cc     | 12 ++++++++++++
 tensorflow/lite/delegates/gpu/common/gpu_info.h      |  2 ++
 tensorflow/lite/delegates/gpu/common/tasks/BUILD     |  1 +
 .../lite/delegates/gpu/common/tasks/conv_generic.cc  | 10 +++++++---
 5 files changed, 23 insertions(+), 3 deletions(-)

diff --git a/tensorflow/lite/delegates/gpu/common/BUILD b/tensorflow/lite/delegates/gpu/common/BUILD
index 0a55ad05a76968..195124f269c7c8 100644
--- a/tensorflow/lite/delegates/gpu/common/BUILD
+++ b/tensorflow/lite/delegates/gpu/common/BUILD
@@ -44,6 +44,7 @@ cc_library(
     deps = [
         ":data_type",
         "@com_google_absl//absl/container:flat_hash_set",
+        "@com_google_absl//absl/status",
         "@com_google_absl//absl/strings",
     ],
 )
diff --git a/tensorflow/lite/delegates/gpu/common/gpu_info.cc b/tensorflow/lite/delegates/gpu/common/gpu_info.cc
index 944202cb3d8561..2627adda13c6bd 100644
--- a/tensorflow/lite/delegates/gpu/common/gpu_info.cc
+++ b/tensorflow/lite/delegates/gpu/common/gpu_info.cc
@@ -21,6 +21,7 @@ limitations under the License.
 #include <utility>
 #include <vector>
 
+#include "absl/status/status.h"
 #include "absl/strings/ascii.h"
 
 namespace tflite {
@@ -981,6 +982,17 @@ bool GpuInfo::SupportsSubGroupWithSize(int sub_group_size) const {
   return false;
 }
 
+absl::Status GpuInfo::GetMinSubGroupSize(int& min_sub_group_size) const {
+  auto begin = supported_subgroup_sizes.begin();
+  auto end = supported_subgroup_sizes.end();
+  auto min = std::min_element(begin, end);
+  if (min == end) {
+    return absl::InternalError("No supported subgroup sizes");
+  }
+  min_sub_group_size = *min;
+  return absl::OkStatus();
+}
+
 bool GpuInfo::SupportsFloatImage2D(DataType data_type, int channels) const {
   if (IsApiOpenCl()) {
     return opencl_info.supported_images_2d.SupportsImage2D(data_type, channels);
diff --git a/tensorflow/lite/delegates/gpu/common/gpu_info.h b/tensorflow/lite/delegates/gpu/common/gpu_info.h
index 15bdb065f15361..9849e2405d32e2 100644
--- a/tensorflow/lite/delegates/gpu/common/gpu_info.h
+++ b/tensorflow/lite/delegates/gpu/common/gpu_info.h
@@ -21,6 +21,7 @@ limitations under the License.
 #include <vector>
 
 #include "absl/container/flat_hash_set.h"
+#include "absl/status/status.h"
 #include "absl/strings/match.h"
 #include "tensorflow/lite/delegates/gpu/common/data_type.h"
 
@@ -524,6 +525,7 @@ struct GpuInfo {
   // returns true if device have fixed wave size equal to 32
   bool IsWaveSizeEqualTo32() const;
   bool SupportsSubGroupWithSize(int sub_group_size) const;
+  absl::Status GetMinSubGroupSize(int& min_sub_group_size) const;
 
   bool SupportsFloatImage2D(DataType data_type, int channels) const;
   bool SupportsExtension(const std::string& extension) const;
diff --git a/tensorflow/lite/delegates/gpu/common/tasks/BUILD b/tensorflow/lite/delegates/gpu/common/tasks/BUILD
index 01ad443db54ae2..0b9662471277e4 100644
--- a/tensorflow/lite/delegates/gpu/common/tasks/BUILD
+++ b/tensorflow/lite/delegates/gpu/common/tasks/BUILD
@@ -143,6 +143,7 @@ cc_library(
     hdrs = ["conv_generic.h"],
     deps = [
         "//tensorflow/lite/delegates/gpu/common:data_type",
+        "//tensorflow/lite/delegates/gpu/common:gpu_info",
         "//tensorflow/lite/delegates/gpu/common:operations",
         "//tensorflow/lite/delegates/gpu/common:shape",
         "//tensorflow/lite/delegates/gpu/common:status",
diff --git a/tensorflow/lite/delegates/gpu/common/tasks/conv_generic.cc b/tensorflow/lite/delegates/gpu/common/tasks/conv_generic.cc
index c2a1bd8f539512..1e5cc2f14b9e82 100644
--- a/tensorflow/lite/delegates/gpu/common/tasks/conv_generic.cc
+++ b/tensorflow/lite/delegates/gpu/common/tasks/conv_generic.cc
@@ -1791,16 +1791,20 @@ ConvGeneric::ConvParams ConvGeneric::GuessBestParams(
         const int kSubGroupSize = 16;
         const bool supports_subgroup_size_control =
             gpu_info.SupportsExtension("cl_intel_required_subgroup_size");
+        int min_subgroup_size;
+        auto min_subgroup_size_status =
+            gpu_info.GetMinSubGroupSize(min_subgroup_size);
         if (supports_subgroup_size_control &&
             gpu_info.SupportsSubGroupWithSize(kSubGroupSize)) {
           conv_params.weights_upload_type =
               WeightsUploadType::PRIVATE_MEM_SIMD_BROADCAST;
           conv_params.simd_size = kSubGroupSize;
-        } else if (gpu_info.opencl_info.IsCLVK()) {
-          // It will work because of specific driver using subgroup size 16
+        } else if (supports_subgroup_size_control &&
+                   min_subgroup_size_status.ok()) {
           conv_params.weights_upload_type =
               WeightsUploadType::PRIVATE_MEM_SIMD_BROADCAST;
-          conv_params.simd_size = 16;
+          conv_params.simd_size = min_subgroup_size;
+          work_group_size_ = int3(min_subgroup_size, 1, 1);
         } else {
           // no support of subgroup size control
           // only smallest subgroup size (8) can be used safely, otherwise

From 32f6db8a0b0f4801445e9007d7ffc94a54b8ac14 Mon Sep 17 00:00:00 2001
From: Siqiao Wu <siqiaowu@google.com>
Date: Wed, 15 May 2024 00:16:34 -0700
Subject: [PATCH 160/478] Add core selector support for TFRT+IFRT serving on
 tensorflow serving

PiperOrigin-RevId: 633836270
---
 tensorflow/core/tfrt/ifrt/ifrt_config.proto | 8 ++++++++
 1 file changed, 8 insertions(+)

diff --git a/tensorflow/core/tfrt/ifrt/ifrt_config.proto b/tensorflow/core/tfrt/ifrt/ifrt_config.proto
index 38c0f6ec0e0902..f11d61049ce666 100644
--- a/tensorflow/core/tfrt/ifrt/ifrt_config.proto
+++ b/tensorflow/core/tfrt/ifrt/ifrt_config.proto
@@ -10,3 +10,11 @@ message VariableDeviceShardingConfigProto {
   xla.OpSharding sharding = 1;
   repeated int32 device_ids = 2;
 }
+
+enum IfrtServingCoreSelectionPolicy {
+  // Default policy to select the soonest-to-finish core defined at
+  // http://shortn/_XQhI1ASAu0.
+  IFRT_SERVING_CORE_SELECTION_POLICY_DEFAULT = 0;
+  // Policy that round robin with local ordinal http://shortn/_7BtVe4dkp5.
+  IFRT_SERVING_CORE_SELECTION_POLICY_LOCAL_ROUND_ROBIN = 1;
+}

From 7087cf81034c4d4c1a2d87f052b9aeec88b738a5 Mon Sep 17 00:00:00 2001
From: Henning Becker <hebecker@google.com>
Date: Wed, 15 May 2024 00:24:47 -0700
Subject: [PATCH 161/478] [XLA:GPU] Remove unused sleep kernel

Nothing is using the sleep kernel, so let's remove it

PiperOrigin-RevId: 633838040
---
 third_party/xla/xla/service/gpu/BUILD         |  7 ----
 .../xla/xla/service/gpu/sleep_kernel.cu.cc    | 32 -------------------
 .../xla/xla/service/gpu/sleep_kernel.h        | 26 ---------------
 3 files changed, 65 deletions(-)
 delete mode 100644 third_party/xla/xla/service/gpu/sleep_kernel.cu.cc
 delete mode 100644 third_party/xla/xla/service/gpu/sleep_kernel.h

diff --git a/third_party/xla/xla/service/gpu/BUILD b/third_party/xla/xla/service/gpu/BUILD
index c19905c49698c8..9206be42cc5ddc 100644
--- a/third_party/xla/xla/service/gpu/BUILD
+++ b/third_party/xla/xla/service/gpu/BUILD
@@ -874,13 +874,6 @@ cc_library(
     ],
 )
 
-cuda_library(
-    name = "sleep_kernel",
-    srcs = if_cuda_is_configured(["sleep_kernel.cu.cc"]),
-    hdrs = if_cuda_is_configured(["sleep_kernel.h"]),
-    deps = ["@local_config_cuda//cuda:cuda_headers"],
-)
-
 # TODO(b/244780257): Remove this config.
 bool_flag(
     name = "enable_xlir",
diff --git a/third_party/xla/xla/service/gpu/sleep_kernel.cu.cc b/third_party/xla/xla/service/gpu/sleep_kernel.cu.cc
deleted file mode 100644
index 8f37d47e677347..00000000000000
--- a/third_party/xla/xla/service/gpu/sleep_kernel.cu.cc
+++ /dev/null
@@ -1,32 +0,0 @@
-/* Copyright 2023 The OpenXLA Authors.
-
-Licensed under the Apache License, Version 2.0 (the "License");
-you may not use this file except in compliance with the License.
-You may obtain a copy of the License at
-
-    http://www.apache.org/licenses/LICENSE-2.0
-
-Unless required by applicable law or agreed to in writing, software
-distributed under the License is distributed on an "AS IS" BASIS,
-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.
-==============================================================================*/
-
-#include "xla/service/gpu/sleep_kernel.h"
-
-namespace xla::gpu {
-namespace {
-
-// Use busy waiting instead of __nanosleep() to make the code more portable
-// (__nanosleep requires __CUDA_ARCH__ >= 700)
-__global__ void sleep(int64_t num_clocks) {
-  int64_t start = clock64();
-  while (clock64() - start < num_clocks) continue;
-}
-
-}  // namespace
-
-void* GetSleepKernel() { return reinterpret_cast<void*>(&sleep); }
-
-}  // namespace xla::gpu
diff --git a/third_party/xla/xla/service/gpu/sleep_kernel.h b/third_party/xla/xla/service/gpu/sleep_kernel.h
deleted file mode 100644
index 3e040b10860e65..00000000000000
--- a/third_party/xla/xla/service/gpu/sleep_kernel.h
+++ /dev/null
@@ -1,26 +0,0 @@
-/* Copyright 2023 The OpenXLA Authors.
-
-Licensed under the Apache License, Version 2.0 (the "License");
-you may not use this file except in compliance with the License.
-You may obtain a copy of the License at
-
-    http://www.apache.org/licenses/LICENSE-2.0
-
-Unless required by applicable law or agreed to in writing, software
-distributed under the License is distributed on an "AS IS" BASIS,
-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.
-==============================================================================*/
-
-#ifndef XLA_SERVICE_GPU_SLEEP_KERNEL_H_
-#define XLA_SERVICE_GPU_SLEEP_KERNEL_H_
-
-namespace xla::gpu {
-
-// Returns a pointer to CUDA kernel that does sleep operation on device.
-void* GetSleepKernel();
-
-}  // namespace xla::gpu
-
-#endif  // XLA_SERVICE_GPU_SLEEP_KERNEL_H_

From a8a6dd4b922ee56b136b014a27ab0694b427eef0 Mon Sep 17 00:00:00 2001
From: lingzhi98 <103185827+lingzhi98@users.noreply.github.com>
Date: Wed, 15 May 2024 01:13:45 -0700
Subject: [PATCH 162/478] PR #12328: Make shared cache read/write logic more
 clearly for transpose mlir emitter

Imported from GitHub PR https://github.com/openxla/xla/pull/12328

Current transpose mlir emitter allocate shared cache with shape 32x1x32 for transpose 2-1-0. But the read indices of shared cache are {0, y, x} as [this line](https://github.com/openxla/xla/blob/main/xla/service/gpu/fusions/transpose_mlir.cc#L190) shows, which is not compatible with 32x1x32 shape. What's strange is that transpose 2-1-0 can run successfully using transpose mlir emitter. I find the reason is that lower tensor pass use [linear index](https://github.com/openxla/xla/blob/main/xla/service/gpu/fusions/mlir/lower_tensors.cc#L148) to access shared cache, which is lucky to get right result. For example, the strides of 32x1x32 are {32, 32, 1}, and the linear index of {0, y ,x} is 0 * 32 + y * 32 + 32.
I am not sure if it is as expected or just mistake. If reviewer think no need of this PR, feel free to close.
Copybara import of the project:

--
bfb21798ee518dc11293a5683669add619a38e53 by Zhou, Lingzhi <lingzhi.zhou@intel.com>:

make shared cache read/write logic more clearly for transpose mlir emitter

--
0c9033334835bc8a14310e5ee059489cea7b5309 by Zhou, Lingzhi <lingzhi.zhou@intel.com>:

refactor

--
5554110835fc18207fb466587c1aeb20c3a542fe by Zhou, Lingzhi <lingzhi.zhou@intel.com>:

pad shared cache

--
8c17818baa1e2477952df15e412a6463f73106ab by Zhou, Lingzhi <lingzhi.zhou@intel.com>:

include missing file

Merging this change closes #12328

PiperOrigin-RevId: 633848774
---
 .../xla/service/gpu/fusions/transpose_mlir.cc | 32 +++++++++----
 .../gpu/fusions/transpose_mlir_test.cc        | 46 ++++++++++++++++++-
 2 files changed, 66 insertions(+), 12 deletions(-)

diff --git a/third_party/xla/xla/service/gpu/fusions/transpose_mlir.cc b/third_party/xla/xla/service/gpu/fusions/transpose_mlir.cc
index 2f602a75b185f4..79111485562508 100644
--- a/third_party/xla/xla/service/gpu/fusions/transpose_mlir.cc
+++ b/third_party/xla/xla/service/gpu/fusions/transpose_mlir.cc
@@ -180,13 +180,20 @@ IndexingMap GetSharedMemoryWriteIndexingMap(
   AffineExpr th_x = mlir::getAffineDimExpr(0, mlir_context);
   SmallVector<AffineExpr, 3> tile_sizes(3);
   mlir::bindSymbolsList(mlir_context, llvm::MutableArrayRef(tile_sizes));
+  SmallVector<AffineExpr, 3> shared_memory_indices = {
+      th_x.floorDiv(32) + 4 * tile_sizes[loop_dim], th_x % 32};
+  for (auto [index, range_val] :
+       llvm::enumerate(thread_id_indexing.GetRangeVars())) {
+    if (range_val.range.NumElements() == 1) {
+      shared_memory_indices.insert(shared_memory_indices.begin() + index, c0);
+      break;
+    }
+  }
 
   IndexingMap shmem_write_indexing{
-      AffineMap::get(
-          thread_id_indexing.GetDimensionCount(),
-          thread_id_indexing.GetSymbolCount(),
-          {c0, th_x.floorDiv(32) + 4 * tile_sizes[loop_dim], th_x % 32},
-          mlir_context),
+      AffineMap::get(thread_id_indexing.GetDimensionCount(),
+                     thread_id_indexing.GetSymbolCount(), shared_memory_indices,
+                     mlir_context),
       thread_id_indexing.GetDimVars(), thread_id_indexing.GetRangeVars(),
       thread_id_indexing.GetRTVars(), thread_id_indexing.GetConstraints()};
   return shmem_write_indexing;
@@ -195,10 +202,12 @@ IndexingMap GetSharedMemoryWriteIndexingMap(
 // Returns an indexing map with block_x, block_y, block_z set to 0 and swapped
 // 2nd and 3rd results.
 IndexingMap GetSharedMemoryReadIndexingMap(
-    const IndexingMap& thread_id_indexing, int loop_dim) {
-  IndexingMap write_indexing =
-      GetSharedMemoryWriteIndexingMap(thread_id_indexing, loop_dim);
-  return IndexingMap{write_indexing.GetAffineMap().getSubMap({0, 2, 1}),
+    const IndexingMap& thread_id_indexing, Vector3 permutation) {
+  IndexingMap write_indexing = GetSharedMemoryWriteIndexingMap(
+      thread_id_indexing, /*loop_dim=*/permutation[2]);
+  llvm::SmallVector<unsigned, 3> positions;
+  absl::c_copy(permutation, std::back_inserter(positions));
+  return IndexingMap{write_indexing.GetAffineMap().getSubMap(positions),
                      write_indexing.GetDimVars(), write_indexing.GetRangeVars(),
                      write_indexing.GetRTVars(),
                      write_indexing.GetConstraints()};
@@ -212,6 +221,9 @@ MlirTransposeFusion::WriteResult MlirTransposeFusion::EmitWriteToShMemMlir(
     ValueRange output_args) const {
   std::vector<int64_t> shmem_tensor_size(tiling_.GetBlockTileSize().begin(),
                                          tiling_.GetBlockTileSize().end());
+  // Avoid bank conflict.
+  ++shmem_tensor_size.back();
+
   MLIRContext* ctx = builder.getContext();
 
   WriteResult write_result;
@@ -305,7 +317,7 @@ void MlirTransposeFusion::EmitReadFromShMemMlir(
   auto* mlir_context = builder.getContext();
   auto output_indexing = *ComputeThreadIdToOutputIndexing(0, mlir_context);
   auto shmem_output_indexing =
-      GetSharedMemoryReadIndexingMap(output_indexing, permutation_[2]);
+      GetSharedMemoryReadIndexingMap(output_indexing, permutation_);
   auto result_tensors = EmitThreadLoopNest(
       builder, written.updated_outputs, output_indexing,
       [&](ValueRange output_tensors, ValueRange dim_values,
diff --git a/third_party/xla/xla/service/gpu/fusions/transpose_mlir_test.cc b/third_party/xla/xla/service/gpu/fusions/transpose_mlir_test.cc
index 086ebf8f2ad911..b2343d916537cf 100644
--- a/third_party/xla/xla/service/gpu/fusions/transpose_mlir_test.cc
+++ b/third_party/xla/xla/service/gpu/fusions/transpose_mlir_test.cc
@@ -171,7 +171,49 @@ TEST_F(MlirTransposeFusionTest, FusedTranspose021) {
     // CHECK-DAG:  %[[C1:.*]] = arith.constant 1 : index
     // CHECK-DAG:  %[[C8:.*]] = arith.constant 8 : index
 
-    // CHECK:      %[[SHMEM:.*]] = xla_gpu.allocate_shared : tensor<1x32x32xf32>
+    // CHECK:      %[[SHMEM:.*]] = xla_gpu.allocate_shared : tensor<1x32x33xf32>
+    // CHECK:      %[[SHMEM_WITH_VALS:.*]] = scf.for
+    // CHECK-SAME:     %[[C0]] to %[[C8]] step %[[C1]]
+    // CHECK-SAME:     iter_args(%[[SHMEM_:.*]] = %[[SHMEM]])
+    // CHECK:        %[[EXP:.*]] = xla_gpu.pure_call @fused_computation_exp
+    // CHECK:        tensor.insert %[[EXP]] into %[[SHMEM_]]
+
+    // CHECK:      %[[SYNC:.*]] = xla_gpu.sync_threads %[[SHMEM_WITH_VALS]]
+
+    // CHECK:      scf.for
+    // CHECK-SAME:    %[[C0]] to %[[C8]] step %[[C1]]
+    // CHECK-SAME:    iter_args(%[[OUT_:.*]] = %[[OUT]])
+    // CHECK:       %[[ABS:.*]] = xla_gpu.pure_call @fused_computation__epilogue__
+    // CHECK:       tensor.insert %[[ABS]] into %[[OUT_]]
+  )"));
+  EXPECT_TRUE(RunAndCompareNoHloPasses(kHloString, ErrorSpec{1e-3}));
+}
+
+TEST_F(MlirTransposeFusionTest, FusedTranspose210) {
+  auto kHloString = R"(
+    HloModule Transpose
+
+    %fused_computation {
+      %p0 = f32[20,160,170] parameter(0)
+      %exp = f32[20,160,170] exponential(%p0)
+      %transpose = f32[170,160,20] transpose(%exp), dimensions={2,1,0}
+      ROOT %abs = f32[170,160,20] abs(%transpose)
+    }
+    ENTRY main {
+      %param = f32[20,160,170] parameter(0)
+      ROOT %fusion = f32[170,160,20] fusion(%param), kind=kInput,
+        calls=%fused_computation
+    }
+  )";
+  TF_ASSERT_OK(EmitAndCheckIR(kHloString, R"(
+    // CHECK-LABEL: func.func @fused_computation(
+    // CHECK-SAME:   }, %[[OUT:.*]]: tensor<170x160x20xf32>
+    //
+    // CHECK-DAG:  %[[C0:.*]] = arith.constant 0 : index
+    // CHECK-DAG:  %[[C1:.*]] = arith.constant 1 : index
+    // CHECK-DAG:  %[[C8:.*]] = arith.constant 8 : index
+
+    // CHECK:      %[[SHMEM:.*]] = xla_gpu.allocate_shared : tensor<32x1x33xf32>
     // CHECK:      %[[SHMEM_WITH_VALS:.*]] = scf.for
     // CHECK-SAME:     %[[C0]] to %[[C8]] step %[[C1]]
     // CHECK-SAME:     iter_args(%[[SHMEM_:.*]] = %[[SHMEM]])
@@ -212,7 +254,7 @@ TEST_F(MlirTransposeFusionTest, Transpose021_Parameter) {
     // CHECK-DAG:  %[[C1:.*]] = arith.constant 1 : index
     // CHECK-DAG:  %[[C8:.*]] = arith.constant 8 : index
 
-    // CHECK:      %[[SHMEM:.*]] = xla_gpu.allocate_shared : tensor<1x32x32xf32>
+    // CHECK:      %[[SHMEM:.*]] = xla_gpu.allocate_shared : tensor<1x32x33xf32>
     // CHECK:      %[[SHMEM_WITH_VALS:.*]] = scf.for
     // CHECK-SAME:     %[[C0]] to %[[C8]] step %[[C1]]
     // CHECK-SAME:     iter_args(%[[SHMEM_:.*]] = %[[SHMEM]])

From 42aeac11da3458b463937dc382882ac459d8edea Mon Sep 17 00:00:00 2001
From: Harsha H S <hsharsha@users.noreply.github.com>
Date: Wed, 15 May 2024 01:26:40 -0700
Subject: [PATCH 163/478] PR #12463: [ROCm] Fix build break in
 `xla/service/gpu/runtime/convolution_thunk.cc`

Imported from GitHub PR https://github.com/openxla/xla/pull/12463

Fix build break introduced in db8f9c0cf2bf08b171aec2cbe3da2fcb074a6207
Copybara import of the project:

--
42d5839ffb47362a7f6789ba98ef50b31189dccc by Harsha HS <harsha.havanurshamsundara@amd.com>:

[ROCm] Fix build break in `xla/service/gpu/runtime/convolution_thunk.cc`

--
fd95e342a4e8de6c8348df3576510584a0020150 by Harsha HS <harsha.havanurshamsundara@amd.com>:

Gaurd header using TENSORFLOW_USE_ROCM

Merging this change closes #12463

PiperOrigin-RevId: 633851634
---
 third_party/xla/xla/service/gpu/runtime/BUILD                | 2 +-
 third_party/xla/xla/service/gpu/runtime/convolution_thunk.cc | 3 +++
 2 files changed, 4 insertions(+), 1 deletion(-)

diff --git a/third_party/xla/xla/service/gpu/runtime/BUILD b/third_party/xla/xla/service/gpu/runtime/BUILD
index e53035e558ee9d..4327ce5583e62b 100644
--- a/third_party/xla/xla/service/gpu/runtime/BUILD
+++ b/third_party/xla/xla/service/gpu/runtime/BUILD
@@ -527,7 +527,7 @@ cc_library(
         "@com_google_absl//absl/synchronization",
         "@com_google_absl//absl/types:span",
         "@local_tsl//tsl/platform:errors",
-    ],
+    ] + if_rocm_is_configured(["//xla/service/gpu:stream_executor_util"]),
 )
 
 cc_library(
diff --git a/third_party/xla/xla/service/gpu/runtime/convolution_thunk.cc b/third_party/xla/xla/service/gpu/runtime/convolution_thunk.cc
index b7bba0fab3b418..32e81855e7ffb8 100644
--- a/third_party/xla/xla/service/gpu/runtime/convolution_thunk.cc
+++ b/third_party/xla/xla/service/gpu/runtime/convolution_thunk.cc
@@ -27,6 +27,9 @@ limitations under the License.
 #include "absl/synchronization/mutex.h"
 #include "absl/types/span.h"
 #include "xla/service/buffer_assignment.h"
+#if TENSORFLOW_USE_ROCM
+#include "xla/service/gpu/stream_executor_util.h"
+#endif  // TENSORFLOW_USE_ROCM
 #include "xla/service/gpu/gpu_conv_runner.h"
 #include "xla/service/gpu/runtime/thunk.h"
 #include "xla/stream_executor/device_memory.h"

From e9215bfcc11aa5fe4115b4d4379894f0597720e6 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Wed, 15 May 2024 01:27:19 -0700
Subject: [PATCH 164/478] Update ops-related pbtxt files.

PiperOrigin-RevId: 633851744
---
 tensorflow/core/ops/ops.pbtxt | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorflow/core/ops/ops.pbtxt b/tensorflow/core/ops/ops.pbtxt
index 846207488b92f3..a15dada136af38 100644
--- a/tensorflow/core/ops/ops.pbtxt
+++ b/tensorflow/core/ops/ops.pbtxt
@@ -1,4 +1,4 @@
-go/debugstr   
+go/nodeserialize  
 op {
   name: "Abort"
   attr {

From f55a2ac80498820b4e7ae15c81e21f847b133a36 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Wed, 15 May 2024 01:53:14 -0700
Subject: [PATCH 165/478] Automated Code Change

PiperOrigin-RevId: 633857603
---
 tensorflow/compiler/mlir/tensorflow/BUILD           |  2 ++
 .../tensorflow/analysis/resource_alias_analysis.cc  | 13 ++++++++++---
 .../tensorflow/analysis/resource_alias_analysis.h   |  7 +++++++
 .../mlir/tensorflow/analysis/resource_dataflow.h    |  2 ++
 .../analysis/resource_value_typed_analyzer.cc       |  9 +++++++++
 .../analysis/resource_value_typed_analyzer.h        |  5 +++++
 .../tensorflow/analysis/side_effect_analysis.cc     | 11 +++++++++--
 .../mlir/tensorflow/analysis/side_effect_analysis.h |  6 ++++++
 8 files changed, 50 insertions(+), 5 deletions(-)

diff --git a/tensorflow/compiler/mlir/tensorflow/BUILD b/tensorflow/compiler/mlir/tensorflow/BUILD
index 85f8644fc2d88f..b138c2d3efd598 100644
--- a/tensorflow/compiler/mlir/tensorflow/BUILD
+++ b/tensorflow/compiler/mlir/tensorflow/BUILD
@@ -676,10 +676,12 @@ cc_library(
         ":tensorflow",
         ":tensorflow_op_interfaces",
         ":tensorflow_side_effects",
+        ":tensorflow_traits",
         ":tensorflow_types",
         "@com_google_absl//absl/container:flat_hash_map",
         "@com_google_absl//absl/container:flat_hash_set",
         "@com_google_absl//absl/container:node_hash_map",
+        "@com_google_absl//absl/log",
         "@llvm-project//llvm:Support",
         "@llvm-project//mlir:Analysis",
         "@llvm-project//mlir:FuncDialect",
diff --git a/tensorflow/compiler/mlir/tensorflow/analysis/resource_alias_analysis.cc b/tensorflow/compiler/mlir/tensorflow/analysis/resource_alias_analysis.cc
index cab89bb10b5fb9..267bc48d17e06d 100644
--- a/tensorflow/compiler/mlir/tensorflow/analysis/resource_alias_analysis.cc
+++ b/tensorflow/compiler/mlir/tensorflow/analysis/resource_alias_analysis.cc
@@ -20,29 +20,36 @@ limitations under the License.
 #include <optional>
 #include <utility>
 
+#include "absl/log/log.h"
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/SCCIterator.h"
 #include "llvm/ADT/STLExtras.h"
+#include "llvm/ADT/SetVector.h"
+#include "llvm/ADT/SmallSet.h"
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/ADT/iterator_range.h"
 #include "llvm/Support/Casting.h"
 #include "mlir/Analysis/CallGraph.h"  // from @llvm-project
 #include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
-#include "mlir/IR/Attributes.h"  // from @llvm-project
 #include "mlir/IR/Block.h"  // from @llvm-project
+#include "mlir/IR/BuiltinAttributes.h"  // from @llvm-project
 #include "mlir/IR/BuiltinOps.h"  // from @llvm-project
-#include "mlir/IR/BuiltinTypes.h"  // from @llvm-project
+#include "mlir/IR/OpDefinition.h"  // from @llvm-project
 #include "mlir/IR/Operation.h"  // from @llvm-project
 #include "mlir/IR/Value.h"  // from @llvm-project
+#include "mlir/IR/ValueRange.h"  // from @llvm-project
 #include "mlir/IR/Visitors.h"  // from @llvm-project
 #include "mlir/Interfaces/CallInterfaces.h"  // from @llvm-project
+#include "mlir/Interfaces/SideEffectInterfaces.h"  // from @llvm-project
 #include "mlir/Support/LLVM.h"  // from @llvm-project
-#include "mlir/Support/LogicalResult.h"  // from @llvm-project
+#include "mlir/Support/TypeID.h"  // from @llvm-project
 #include "tensorflow/compiler/mlir/tensorflow/ir/tf_device.h"
 #include "tensorflow/compiler/mlir/tensorflow/ir/tf_executor.h"
 #include "tensorflow/compiler/mlir/tensorflow/ir/tf_op_interfaces.h"
 #include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
+#include "tensorflow/compiler/mlir/tensorflow/ir/tf_traits.h"
+#include "tensorflow/compiler/mlir/tensorflow/ir/tf_types.h"
 
 namespace mlir {
 namespace TF {
diff --git a/tensorflow/compiler/mlir/tensorflow/analysis/resource_alias_analysis.h b/tensorflow/compiler/mlir/tensorflow/analysis/resource_alias_analysis.h
index 7afec29bc5df75..c49852c1864763 100644
--- a/tensorflow/compiler/mlir/tensorflow/analysis/resource_alias_analysis.h
+++ b/tensorflow/compiler/mlir/tensorflow/analysis/resource_alias_analysis.h
@@ -21,14 +21,21 @@ limitations under the License.
 #include <memory>
 
 #include "llvm/ADT/ArrayRef.h"
+#include "llvm/ADT/DenseMap.h"
+#include "llvm/ADT/DenseSet.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/SetVector.h"
 #include "llvm/ADT/SmallSet.h"
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/ADT/StringMap.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
+#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
 #include "mlir/IR/Operation.h"  // from @llvm-project
 #include "mlir/IR/Region.h"  // from @llvm-project
+#include "mlir/IR/SymbolTable.h"  // from @llvm-project
 #include "mlir/IR/TypeUtilities.h"  // from @llvm-project
+#include "mlir/IR/Value.h"  // from @llvm-project
+#include "mlir/Support/TypeID.h"  // from @llvm-project
 #include "tensorflow/compiler/mlir/tensorflow/analysis/per_function_aggregate_analysis.h"
 #include "tensorflow/compiler/mlir/tensorflow/ir/tf_types.h"
 
diff --git a/tensorflow/compiler/mlir/tensorflow/analysis/resource_dataflow.h b/tensorflow/compiler/mlir/tensorflow/analysis/resource_dataflow.h
index 0cf3611af1d20c..1e68ac41d25b54 100644
--- a/tensorflow/compiler/mlir/tensorflow/analysis/resource_dataflow.h
+++ b/tensorflow/compiler/mlir/tensorflow/analysis/resource_dataflow.h
@@ -25,9 +25,11 @@ limitations under the License.
 #include "llvm/Support/Debug.h"
 #include "mlir/Analysis/DataFlow/DeadCodeAnalysis.h"  // from @llvm-project
 #include "mlir/Analysis/DataFlow/SparseAnalysis.h"  // from @llvm-project
+#include "mlir/Analysis/DataFlowFramework.h"  // from @llvm-project
 #include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
 #include "mlir/IR/Builders.h"  // from @llvm-project
 #include "mlir/IR/BuiltinOps.h"  // from @llvm-project
+#include "mlir/IR/MLIRContext.h"  // from @llvm-project
 #include "mlir/IR/SymbolTable.h"  // from @llvm-project
 #include "mlir/IR/Value.h"  // from @llvm-project
 #include "mlir/Pass/Pass.h"  // from @llvm-project
diff --git a/tensorflow/compiler/mlir/tensorflow/analysis/resource_value_typed_analyzer.cc b/tensorflow/compiler/mlir/tensorflow/analysis/resource_value_typed_analyzer.cc
index abace5111184ff..372446641382ac 100644
--- a/tensorflow/compiler/mlir/tensorflow/analysis/resource_value_typed_analyzer.cc
+++ b/tensorflow/compiler/mlir/tensorflow/analysis/resource_value_typed_analyzer.cc
@@ -16,9 +16,18 @@ limitations under the License.
 
 #include <tuple>
 
+#include "llvm/ADT/STLExtras.h"
 #include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
 #include "mlir/IR/BuiltinAttributes.h"  // from @llvm-project
+#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
+#include "mlir/IR/BuiltinTypes.h"  // from @llvm-project
+#include "mlir/IR/Region.h"  // from @llvm-project
+#include "mlir/IR/SymbolTable.h"  // from @llvm-project
+#include "mlir/IR/Types.h"  // from @llvm-project
+#include "mlir/IR/Value.h"  // from @llvm-project
+#include "mlir/Interfaces/CallInterfaces.h"  // from @llvm-project
 #include "mlir/Support/LLVM.h"  // from @llvm-project
+#include "mlir/Support/LogicalResult.h"  // from @llvm-project
 #include "mlir/Transforms/RegionUtils.h"  // from @llvm-project
 #include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
 #include "tensorflow/compiler/mlir/tensorflow/ir/tf_saved_model.h"
diff --git a/tensorflow/compiler/mlir/tensorflow/analysis/resource_value_typed_analyzer.h b/tensorflow/compiler/mlir/tensorflow/analysis/resource_value_typed_analyzer.h
index 9817b290c4cbdb..738d8c1df3d395 100644
--- a/tensorflow/compiler/mlir/tensorflow/analysis/resource_value_typed_analyzer.h
+++ b/tensorflow/compiler/mlir/tensorflow/analysis/resource_value_typed_analyzer.h
@@ -22,6 +22,11 @@ limitations under the License.
 #include "llvm/ADT/StringRef.h"
 #include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
 #include "mlir/IR/BuiltinOps.h"  // from @llvm-project
+#include "mlir/IR/Operation.h"  // from @llvm-project
+#include "mlir/IR/Region.h"  // from @llvm-project
+#include "mlir/IR/Value.h"  // from @llvm-project
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+#include "mlir/Support/LogicalResult.h"  // from @llvm-project
 #include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
 
 namespace mlir {
diff --git a/tensorflow/compiler/mlir/tensorflow/analysis/side_effect_analysis.cc b/tensorflow/compiler/mlir/tensorflow/analysis/side_effect_analysis.cc
index df0138a20a0c74..179b3979348161 100644
--- a/tensorflow/compiler/mlir/tensorflow/analysis/side_effect_analysis.cc
+++ b/tensorflow/compiler/mlir/tensorflow/analysis/side_effect_analysis.cc
@@ -26,25 +26,32 @@ limitations under the License.
 #include "absl/container/flat_hash_map.h"
 #include "absl/container/flat_hash_set.h"
 #include "absl/container/node_hash_map.h"
+#include "absl/log/log.h"
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/DenseSet.h"
 #include "llvm/ADT/STLExtras.h"
+#include "llvm/ADT/STLFunctionalExtras.h"
+#include "llvm/ADT/SmallSet.h"
 #include "llvm/ADT/SmallVector.h"
-#include "llvm/ADT/iterator_range.h"
 #include "llvm/Support/Casting.h"
 #include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
-#include "mlir/IR/Attributes.h"  // from @llvm-project
 #include "mlir/IR/Block.h"  // from @llvm-project
 #include "mlir/IR/Builders.h"  // from @llvm-project
 #include "mlir/IR/BuiltinAttributes.h"  // from @llvm-project
 #include "mlir/IR/BuiltinOps.h"  // from @llvm-project
 #include "mlir/IR/Operation.h"  // from @llvm-project
+#include "mlir/IR/SymbolTable.h"  // from @llvm-project
+#include "mlir/IR/TypeUtilities.h"  // from @llvm-project
 #include "mlir/IR/Value.h"  // from @llvm-project
 #include "mlir/IR/Visitors.h"  // from @llvm-project
+#include "mlir/Interfaces/CallInterfaces.h"  // from @llvm-project
 #include "mlir/Interfaces/SideEffectInterfaces.h"  // from @llvm-project
 #include "mlir/Support/DebugStringHelper.h"  // from @llvm-project
 #include "mlir/Support/LLVM.h"  // from @llvm-project
+#include "mlir/Support/TypeID.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/tensorflow/analysis/resource_alias_analysis.h"
 #include "tensorflow/compiler/mlir/tensorflow/ir/tf_device.h"
+#include "tensorflow/compiler/mlir/tensorflow/ir/tf_dialect.h"
 #include "tensorflow/compiler/mlir/tensorflow/ir/tf_executor.h"
 #include "tensorflow/compiler/mlir/tensorflow/ir/tf_op_interfaces.h"
 #include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
diff --git a/tensorflow/compiler/mlir/tensorflow/analysis/side_effect_analysis.h b/tensorflow/compiler/mlir/tensorflow/analysis/side_effect_analysis.h
index 97fcd30d36d02f..feb90de18857b2 100644
--- a/tensorflow/compiler/mlir/tensorflow/analysis/side_effect_analysis.h
+++ b/tensorflow/compiler/mlir/tensorflow/analysis/side_effect_analysis.h
@@ -23,12 +23,18 @@ limitations under the License.
 
 #include "absl/container/flat_hash_map.h"
 #include "absl/container/flat_hash_set.h"
+#include "llvm/ADT/DenseMap.h"
+#include "llvm/ADT/STLFunctionalExtras.h"
 #include "llvm/ADT/SetVector.h"
+#include "llvm/ADT/SmallPtrSet.h"
 #include "llvm/ADT/SmallSet.h"
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/ADT/StringMap.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
+#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
 #include "mlir/IR/Operation.h"  // from @llvm-project
 #include "mlir/IR/Region.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/tensorflow/analysis/per_function_aggregate_analysis.h"
 #include "tensorflow/compiler/mlir/tensorflow/analysis/resource_alias_analysis.h"
 
 namespace mlir {

From 8f18c6a99aea641ddd355b3ca8bc9f61f32cee4a Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Wed, 15 May 2024 02:03:30 -0700
Subject: [PATCH 166/478] compat: Update forward compatibility horizon to
 2024-05-15

PiperOrigin-RevId: 633859960
---
 tensorflow/python/compat/compat.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorflow/python/compat/compat.py b/tensorflow/python/compat/compat.py
index f372d9e16c9363..8416badea05697 100644
--- a/tensorflow/python/compat/compat.py
+++ b/tensorflow/python/compat/compat.py
@@ -29,7 +29,7 @@
 # This value changes every day with an automatic CL. It can be modified in code
 # via `forward_compatibility_horizon()` or with the environment variable
 # TF_FORWARD_COMPATIBILITY_DELTA_DAYS, which is added to the compatibility date.
-_FORWARD_COMPATIBILITY_HORIZON = datetime.date(2024, 5, 14)
+_FORWARD_COMPATIBILITY_HORIZON = datetime.date(2024, 5, 15)
 _FORWARD_COMPATIBILITY_DELTA_DAYS_VAR_NAME = "TF_FORWARD_COMPATIBILITY_DELTA_DAYS"
 _FORWARD_COMPATIBILITY_DATE_NUMBER = None
 

From 55afb206dafb17a1e4f31b4fd0030c1f5ca625b8 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Wed, 15 May 2024 02:03:37 -0700
Subject: [PATCH 167/478] Update GraphDef version to 1863.

PiperOrigin-RevId: 633859985
---
 tensorflow/core/public/version.h | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorflow/core/public/version.h b/tensorflow/core/public/version.h
index e6d5b27b16cb2c..a60f2dd95fe679 100644
--- a/tensorflow/core/public/version.h
+++ b/tensorflow/core/public/version.h
@@ -108,7 +108,7 @@ limitations under the License.
 
 #define TF_GRAPH_DEF_VERSION_MIN_PRODUCER 0
 #define TF_GRAPH_DEF_VERSION_MIN_CONSUMER 0
-#define TF_GRAPH_DEF_VERSION 1862  // Updated: 2024/5/14
+#define TF_GRAPH_DEF_VERSION 1863  // Updated: 2024/5/15
 
 // Checkpoint compatibility versions (the versions field in SavedSliceMeta).
 //

From a77a2470aed0f56b11db21f15153fcadfaee3f15 Mon Sep 17 00:00:00 2001
From: Greg Olechwierowicz <olechwierowicz@google.com>
Date: Wed, 15 May 2024 02:04:45 -0700
Subject: [PATCH 168/478] [XLA:GPU] Migrate while loop extraction to utility
 function.

PiperOrigin-RevId: 633860340
---
 third_party/xla/xla/service/gpu/BUILD         |   1 +
 .../gpu/double_buffer_loop_unrolling_test.cc  | 153 +++++++-----------
 2 files changed, 58 insertions(+), 96 deletions(-)

diff --git a/third_party/xla/xla/service/gpu/BUILD b/third_party/xla/xla/service/gpu/BUILD
index 9206be42cc5ddc..3883a34d7c4b6e 100644
--- a/third_party/xla/xla/service/gpu/BUILD
+++ b/third_party/xla/xla/service/gpu/BUILD
@@ -5838,6 +5838,7 @@ xla_cc_test(
         "//xla:xla_data_proto_cc",
         "//xla:xla_proto_cc",
         "//xla/hlo/ir:hlo",
+        "//xla/hlo/utils:hlo_query",
         "//xla/service:tuple_simplifier",
         "//xla/tests:hlo_test_base",
         "//xla/tests:xla_internal_test_main",
diff --git a/third_party/xla/xla/service/gpu/double_buffer_loop_unrolling_test.cc b/third_party/xla/xla/service/gpu/double_buffer_loop_unrolling_test.cc
index b7ff9a24de7e2e..2d2f33f160db87 100644
--- a/third_party/xla/xla/service/gpu/double_buffer_loop_unrolling_test.cc
+++ b/third_party/xla/xla/service/gpu/double_buffer_loop_unrolling_test.cc
@@ -24,6 +24,7 @@ limitations under the License.
 #include "xla/hlo/ir/hlo_instruction.h"
 #include "xla/hlo/ir/hlo_module.h"
 #include "xla/hlo/ir/hlo_opcode.h"
+#include "xla/hlo/utils/hlo_query.h"
 #include "xla/service/tuple_simplifier.h"
 #include "xla/test.h"
 #include "xla/tests/hlo_test_base.h"
@@ -110,12 +111,8 @@ ENTRY main {
   EXPECT_TRUE(changed);
   TF_ASSERT_OK_AND_ASSIGN(changed, tuple_simp.Run(module.get()));
   EXPECT_TRUE(changed);
-  HloInstruction* while_instruction;
-  for (auto instr : module->entry_computation()->instructions()) {
-    if (instr->opcode() == HloOpcode::kWhile) {
-      while_instruction = instr;
-    }
-  }
+  HloInstruction* while_instruction = hlo_query::GetFirstInstructionWithOpcode(
+      *module->entry_computation(), HloOpcode::kWhile);
   TF_ASSERT_OK_AND_ASSIGN(
       WhileLoopBackendConfig config,
       while_instruction->backend_config<WhileLoopBackendConfig>());
@@ -243,12 +240,8 @@ ENTRY main {
   TF_ASSERT_OK_AND_ASSIGN(changed, tuple_simp.Run(module.get()));
   EXPECT_TRUE(changed);
 
-  HloInstruction* while_instruction;
-  for (auto instr : module->entry_computation()->instructions()) {
-    if (instr->opcode() == HloOpcode::kWhile) {
-      while_instruction = instr;
-    }
-  }
+  HloInstruction* while_instruction = hlo_query::GetFirstInstructionWithOpcode(
+      *module->entry_computation(), HloOpcode::kWhile);
   TF_ASSERT_OK_AND_ASSIGN(
       WhileLoopBackendConfig config,
       while_instruction->backend_config<WhileLoopBackendConfig>());
@@ -314,12 +307,8 @@ ENTRY main {
 
   // We expect that for the while loop, no further copy needs to be added to the
   // module.
-  HloInstruction* while_instruction;
-  for (auto instr : module->entry_computation()->instructions()) {
-    if (instr->opcode() == HloOpcode::kWhile) {
-      while_instruction = instr;
-    }
-  }
+  HloInstruction* while_instruction = hlo_query::GetFirstInstructionWithOpcode(
+      *module->entry_computation(), HloOpcode::kWhile);
   TF_ASSERT_OK_AND_ASSIGN(
       WhileLoopBackendConfig config,
       while_instruction->backend_config<WhileLoopBackendConfig>());
@@ -377,12 +366,8 @@ ENTRY main {
   EXPECT_THAT(double_buffer.Run(module.get()), IsOkAndHolds(true));
   EXPECT_THAT(tuple_simp.Run(module.get()), IsOkAndHolds(true));
 
-  HloInstruction* while_instruction;
-  for (auto instr : module->entry_computation()->instructions()) {
-    if (instr->opcode() == HloOpcode::kWhile) {
-      while_instruction = instr;
-    }
-  }
+  HloInstruction* while_instruction = hlo_query::GetFirstInstructionWithOpcode(
+      *module->entry_computation(), HloOpcode::kWhile);
   TF_ASSERT_OK_AND_ASSIGN(
       WhileLoopBackendConfig config,
       while_instruction->backend_config<WhileLoopBackendConfig>());
@@ -449,12 +434,8 @@ ENTRY main {
   EXPECT_THAT(double_buffer.Run(module.get()), IsOkAndHolds(true));
   EXPECT_THAT(tuple_simp.Run(module.get()), IsOkAndHolds(true));
 
-  HloInstruction* while_instruction;
-  for (auto instr : module->entry_computation()->instructions()) {
-    if (instr->opcode() == HloOpcode::kWhile) {
-      while_instruction = instr;
-    }
-  }
+  HloInstruction* while_instruction = hlo_query::GetFirstInstructionWithOpcode(
+      *module->entry_computation(), HloOpcode::kWhile);
   TF_ASSERT_OK_AND_ASSIGN(
       WhileLoopBackendConfig config,
       while_instruction->backend_config<WhileLoopBackendConfig>());
@@ -468,14 +449,14 @@ ENTRY main {
                               HloOpcode::kAllReduceStart),
             2);
   absl::flat_hash_set<int64_t> channel_ids;
-  for (HloInstruction* ar : while_instruction->while_body()->instructions()) {
-    if (ar->opcode() == HloOpcode::kAllReduceStart) {
-      // We expect that after unrolling, all-reduces should not have any control
-      // deps.
-      EXPECT_EQ(ar->control_predecessors().size(), 0);
-      channel_ids.insert(*(ar->channel_id()));
-    }
-  }
+  hlo_query::ForEachInstructionWithOpcode(
+      *while_instruction->while_body(), HloOpcode::kAllReduceStart,
+      [&channel_ids](HloInstruction* ar) {
+        // We expect that after unrolling, all-reduces should not have any
+        // control deps.
+        EXPECT_EQ(ar->control_predecessors().size(), 0);
+        channel_ids.insert(*(ar->channel_id()));
+      });
   // we expect that all 2 all-reduces will have different channel ids.
   EXPECT_EQ(channel_ids.size(), 2);
 }
@@ -523,12 +504,8 @@ ENTRY main {
   EXPECT_THAT(double_buffer.Run(module.get()), IsOkAndHolds(true));
   EXPECT_THAT(tuple_simp.Run(module.get()), IsOkAndHolds(true));
 
-  HloInstruction* while_instruction;
-  for (auto instr : module->entry_computation()->instructions()) {
-    if (instr->opcode() == HloOpcode::kWhile) {
-      while_instruction = instr;
-    }
-  }
+  HloInstruction* while_instruction = hlo_query::GetFirstInstructionWithOpcode(
+      *module->entry_computation(), HloOpcode::kWhile);
   TF_ASSERT_OK_AND_ASSIGN(
       WhileLoopBackendConfig config,
       while_instruction->backend_config<WhileLoopBackendConfig>());
@@ -540,14 +517,14 @@ ENTRY main {
                               HloOpcode::kAllReduceStart),
             10);
   absl::flat_hash_set<int64_t> channel_ids;
-  for (HloInstruction* ar : while_instruction->while_body()->instructions()) {
-    if (ar->opcode() == HloOpcode::kAllReduceStart) {
-      // We expect that after unrolling, all-reduces should not have any control
-      // deps.
-      EXPECT_EQ(ar->control_predecessors().size(), 0);
-      channel_ids.insert(*(ar->channel_id()));
-    }
-  }
+  hlo_query::ForEachInstructionWithOpcode(
+      *while_instruction->while_body(), HloOpcode::kAllReduceStart,
+      [&channel_ids](HloInstruction* ar) {
+        // We expect that after unrolling, all-reduces should not have any
+        // control deps.
+        EXPECT_EQ(ar->control_predecessors().size(), 0);
+        channel_ids.insert(*(ar->channel_id()));
+      });
   // we expect that all 10 all-reduces will have different channel ids.
   EXPECT_EQ(channel_ids.size(), 10);
 }
@@ -597,15 +574,13 @@ ENTRY main {
 
   absl::flat_hash_set<const HloComputation*> while_loops_callees;
 
-  for (const HloComputation* computation : module->computations()) {
-    for (const HloInstruction* instr : computation->instructions()) {
-      if (instr->opcode() == HloOpcode::kWhile) {
+  hlo_query::ForEachInstructionWithOpcode(
+      *module, HloOpcode::kWhile,
+      [&while_loops_callees](HloInstruction* instr) {
         EXPECT_TRUE(
             while_loops_callees.insert(instr->while_condition()).second);
         EXPECT_TRUE(while_loops_callees.insert(instr->while_body()).second);
-      }
-    }
-  }
+      });
 
   // We expect that the nested while loop has been duplicated, along with its
   // associated computations.
@@ -656,15 +631,13 @@ ENTRY main {
 
   absl::flat_hash_set<const HloComputation*> while_loops_callees;
 
-  for (const HloComputation* computation : module->computations()) {
-    for (const HloInstruction* instr : computation->instructions()) {
-      if (instr->opcode() == HloOpcode::kWhile) {
+  hlo_query::ForEachInstructionWithOpcode(
+      *module, HloOpcode::kWhile,
+      [&while_loops_callees](HloInstruction* instr) {
         EXPECT_TRUE(
             while_loops_callees.insert(instr->while_condition()).second);
         EXPECT_TRUE(while_loops_callees.insert(instr->while_body()).second);
-      }
-    }
-  }
+      });
 
   // We expect that the nested while loop has been duplicated, along with its
   // associated computations.
@@ -716,25 +689,23 @@ ENTRY main {
 
   absl::flat_hash_set<const HloComputation*> while_loops_callees;
 
-  for (const HloComputation* computation : module->computations()) {
-    for (const HloInstruction* instr : computation->instructions()) {
-      if (instr->opcode() == HloOpcode::kWhile) {
+  hlo_query::ForEachInstructionWithOpcode(
+      *module, HloOpcode::kWhile,
+      [&while_loops_callees](HloInstruction* instr) {
         EXPECT_TRUE(
             while_loops_callees.insert(instr->while_condition()).second);
         EXPECT_TRUE(while_loops_callees.insert(instr->while_body()).second);
-      }
-    }
-  }
-
-  for (const HloInstruction* instr :
-       module->entry_computation()->instructions()) {
-    if (instr->opcode() == HloOpcode::kWhile) {
-      TF_ASSERT_OK_AND_ASSIGN(WhileLoopBackendConfig config,
-                              instr->backend_config<WhileLoopBackendConfig>());
-      int64_t exact_trip_count = config.known_trip_count().n();
-      EXPECT_EQ(exact_trip_count, 1);
-    }
-  }
+      });
+
+  hlo_query::ForEachInstructionWithOpcode(
+      *module->entry_computation(), HloOpcode::kWhile,
+      [](HloInstruction* instr) {
+        TF_ASSERT_OK_AND_ASSIGN(
+            WhileLoopBackendConfig config,
+            instr->backend_config<WhileLoopBackendConfig>());
+        int64_t exact_trip_count = config.known_trip_count().n();
+        EXPECT_EQ(exact_trip_count, 1);
+      });
 
   // We expect that the nested while loop has been fully duplicated 10
   // times. The one outer while loop still remains so that's 11 while
@@ -780,19 +751,14 @@ ENTRY main {
   EXPECT_THAT(double_buffer.Run(module.get()), IsOkAndHolds(true));
 
   int64_t num_whiles = 0;
-  for (const HloComputation* computation : module->computations()) {
-    for (const HloInstruction* instr : computation->instructions()) {
-      if (instr->opcode() == HloOpcode::kWhile) {
-        // All loops in the module should be unrolled now and have trip count
-        // of 5.
+  hlo_query::ForEachInstructionWithOpcode(
+      *module, HloOpcode::kWhile, [&num_whiles](HloInstruction* instr) {
         EXPECT_EQ(instr->backend_config<WhileLoopBackendConfig>()
                       ->known_trip_count()
                       .n(),
                   5);
         ++num_whiles;
-      }
-    }
-  }
+      });
   // We expect the number of while loops to be 4 in total after unrolling.
   EXPECT_EQ(num_whiles, 4);
 }
@@ -835,19 +801,14 @@ ENTRY main {
   EXPECT_THAT(double_buffer.Run(module.get()), IsOkAndHolds(true));
 
   int64_t num_whiles = 0;
-  // TODO(olechwierowicz): We should have an abstraction in HloModule which lets
-  // us to iterate over instrucitons of certain kind.
-  for (const HloComputation* computation : module->computations()) {
-    for (const HloInstruction* instr : computation->instructions()) {
-      if (instr->opcode() == HloOpcode::kWhile) {
+  hlo_query::ForEachInstructionWithOpcode(
+      *module, HloOpcode::kWhile, [&num_whiles](HloInstruction* instr) {
         EXPECT_EQ(instr->backend_config<WhileLoopBackendConfig>()
                       ->known_trip_count()
                       .n(),
                   1);
         ++num_whiles;
-      }
-    }
-  }
+      });
   EXPECT_EQ(num_whiles, 12);
 }
 

From a77e969118f85ae5a7ae738e0f423d01c85291f3 Mon Sep 17 00:00:00 2001
From: Henning Becker <hebecker@google.com>
Date: Wed, 15 May 2024 03:11:15 -0700
Subject: [PATCH 169/478] Fix ROCm build after TypedKernel change

PiperOrigin-RevId: 633874819
---
 third_party/xla/xla/stream_executor/gpu/BUILD         |  1 +
 .../gpu/redzone_allocator_kernel_rocm.cu.cc           | 11 ++++++-----
 2 files changed, 7 insertions(+), 5 deletions(-)

diff --git a/third_party/xla/xla/stream_executor/gpu/BUILD b/third_party/xla/xla/stream_executor/gpu/BUILD
index b2bd00c972d26b..86e594ce824f3a 100644
--- a/third_party/xla/xla/stream_executor/gpu/BUILD
+++ b/third_party/xla/xla/stream_executor/gpu/BUILD
@@ -515,6 +515,7 @@ gpu_kernel_library(
         "//xla/stream_executor",
         "//xla/stream_executor:device_memory",
         "//xla/stream_executor:stream_executor_headers",
+        "//xla/stream_executor:typed_kernel_factory",
         "@com_google_absl//absl/status:statusor",
         "@local_config_rocm//rocm:rocm_headers",
         "@local_tsl//tsl/platform:statusor",
diff --git a/third_party/xla/xla/stream_executor/gpu/redzone_allocator_kernel_rocm.cu.cc b/third_party/xla/xla/stream_executor/gpu/redzone_allocator_kernel_rocm.cu.cc
index e7958b4ba4e631..c1e7200cddbff7 100644
--- a/third_party/xla/xla/stream_executor/gpu/redzone_allocator_kernel_rocm.cu.cc
+++ b/third_party/xla/xla/stream_executor/gpu/redzone_allocator_kernel_rocm.cu.cc
@@ -21,6 +21,7 @@ limitations under the License.
 #include "xla/stream_executor/kernel.h"
 #include "xla/stream_executor/stream_executor.h"
 #include "xla/stream_executor/stream_executor_pimpl.h"
+#include "xla/stream_executor/typed_kernel_factory.h"
 #include "tsl/platform/statusor.h"
 
 namespace {
@@ -38,11 +39,11 @@ namespace stream_executor {
 
 absl::StatusOr<const ComparisonKernel*> GetComparisonKernel(
     StreamExecutor* executor, GpuAsmOpts /*gpu_asm_opts*/) {
-  static auto kernel =
-      TypedKernel<DeviceMemory<uint8>, uint8, uint64_t,
-                  DeviceMemory<uint64_t>>::Create(executor, "redzone_checker",
-                                                  reinterpret_cast<void*>(
-                                                      redzone_checker_kernel));
+  static auto kernel = TypedKernelFactory<
+      DeviceMemory<uint8>, uint8, uint64_t,
+      DeviceMemory<uint64_t>>::Create(executor, "redzone_checker",
+                                      reinterpret_cast<void*>(
+                                          redzone_checker_kernel));
 
   if (!kernel.ok()) return kernel.status();
   return &kernel.value();

From 12ec42e0c07fa97d009dcd050b4a009d95f51640 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Wed, 15 May 2024 03:19:26 -0700
Subject: [PATCH 170/478] Update ops-related pbtxt files.

PiperOrigin-RevId: 633876272
---
 tensorflow/core/ops/ops.pbtxt | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorflow/core/ops/ops.pbtxt b/tensorflow/core/ops/ops.pbtxt
index a15dada136af38..2073f2e00ecd81 100644
--- a/tensorflow/core/ops/ops.pbtxt
+++ b/tensorflow/core/ops/ops.pbtxt
@@ -1,4 +1,4 @@
-go/nodeserialize  
+go/debugonly    
 op {
   name: "Abort"
   attr {

From e3f99f4f953f870a4695718b280fdf232088288b Mon Sep 17 00:00:00 2001
From: Artsiom Ablavatski <artsiom@google.com>
Date: Wed, 15 May 2024 03:25:10 -0700
Subject: [PATCH 171/478] Introduce reference implementation of DRQ
 (dynamically quantized) TransposeConv with per-channel quantization

PiperOrigin-RevId: 633877445
---
 .../internal/reference/transpose_conv.h       |  97 ++++++++++
 tensorflow/lite/kernels/transpose_conv.cc     | 179 +++++++++++++++++-
 .../lite/kernels/transpose_conv_test.cc       |  96 ++++++++++
 3 files changed, 364 insertions(+), 8 deletions(-)

diff --git a/tensorflow/lite/kernels/internal/reference/transpose_conv.h b/tensorflow/lite/kernels/internal/reference/transpose_conv.h
index 8a51e0fa5e9742..744ed0f826b335 100644
--- a/tensorflow/lite/kernels/internal/reference/transpose_conv.h
+++ b/tensorflow/lite/kernels/internal/reference/transpose_conv.h
@@ -219,6 +219,103 @@ inline void TransposeConv(
   }
 }
 
+inline void HybridTransposeConv(
+    const ConvParams& params, float* scaling_factors_ptr,
+    const RuntimeShape& input_shape, const int8_t* input_data,
+    const RuntimeShape& filter_shape, const int8_t* filter_data,
+    const RuntimeShape& bias_shape, const float* bias_data,
+    const RuntimeShape& output_shape, float* output_data,
+    const float* per_channel_scale, int32_t* input_offset) {
+  const int stride_width = params.stride_width;
+  const int stride_height = params.stride_height;
+  const int pad_width = params.padding_values.width;
+  const int pad_height = params.padding_values.height;
+  TFLITE_DCHECK_EQ(input_shape.DimensionsCount(), 4);
+  TFLITE_DCHECK_EQ(filter_shape.DimensionsCount(), 4);
+  TFLITE_DCHECK_EQ(output_shape.DimensionsCount(), 4);
+
+  const int batches = MatchingDim(input_shape, 0, output_shape, 0);
+  const int input_depth = MatchingDim(input_shape, 3, filter_shape, 3);
+  const int output_depth = MatchingDim(filter_shape, 0, output_shape, 3);
+  const int input_height = input_shape.Dims(1);
+  const int input_width = input_shape.Dims(2);
+  const int filter_height = filter_shape.Dims(1);
+  const int filter_width = filter_shape.Dims(2);
+  const int output_height = output_shape.Dims(1);
+  const int output_width = output_shape.Dims(2);
+  const float output_activation_min = params.float_activation_min;
+  const float output_activation_max = params.float_activation_max;
+  if (bias_data) {
+    TFLITE_DCHECK_EQ(bias_shape.FlatSize(), output_depth);
+  }
+
+  // Although transpose convolution simplifies to convolution with transposed
+  // weights for strides of 1, non-unitary striding complicates matters. To
+  // keep this reference implementation as clear as possible, we use a
+  // "scatter" access pattern, where we loop through all the input elements,
+  // computing their influence on the output, rather than looping through the
+  // output elements in the typical "gather" access pattern of a conv. We
+  // therefore must initialize the output array to zero.
+  const int num_elements = output_shape.FlatSize();
+  for (int i = 0; i < num_elements; i++) {
+    output_data[i] = 0.0f;
+  }
+
+  // Loop through input elements one at a time.
+  for (int batch = 0; batch < batches; ++batch) {
+    const float scaling_factor = scaling_factors_ptr[batch];
+    for (int in_y = 0; in_y < input_height; ++in_y) {
+      for (int in_x = 0; in_x < input_width; ++in_x) {
+        for (int in_channel = 0; in_channel < input_depth; ++in_channel) {
+          // Loop through the output elements it will influence
+          const int out_x_origin = (in_x * stride_width) - pad_width;
+          const int out_y_origin = (in_y * stride_height) - pad_height;
+          for (int filter_y = 0; filter_y < filter_height; ++filter_y) {
+            for (int filter_x = 0; filter_x < filter_width; ++filter_x) {
+              for (int out_channel = 0; out_channel < output_depth;
+                   ++out_channel) {
+                // Compute output element location
+                const int out_x = out_x_origin + filter_x;
+                const int out_y = out_y_origin + filter_y;
+                // We cannot accumulate out of bounds
+                if ((out_x >= 0) && (out_x < output_width) && (out_y >= 0) &&
+                    (out_y < output_height)) {
+                  int32_t input_value = input_data[Offset(
+                      input_shape, batch, in_y, in_x, in_channel)];
+                  int32_t filter_value =
+                      filter_data[Offset(filter_shape, out_channel, filter_y,
+                                         filter_x, in_channel)];
+                  int32_t acc =
+                      (input_value - input_offset[batch]) * filter_value;
+                  output_data[Offset(output_shape, batch, out_y, out_x,
+                                     out_channel)] +=
+                      acc * per_channel_scale[out_channel] * scaling_factor;
+                }
+              }
+            }
+          }
+        }
+      }
+    }
+  }
+
+  for (int batch = 0; batch < batches; ++batch) {
+    for (int out_y = 0; out_y < output_height; ++out_y) {
+      for (int out_x = 0; out_x < output_width; ++out_x) {
+        for (int out_channel = 0; out_channel < output_depth; ++out_channel) {
+          float acc = output_data[Offset(output_shape, batch, out_y, out_x,
+                                         out_channel)];
+          if (bias_data) acc += bias_data[out_channel];
+
+          output_data[Offset(output_shape, batch, out_y, out_x, out_channel)] =
+              ActivationFunctionWithMinMax(acc, output_activation_min,
+                                           output_activation_max);
+        }
+      }
+    }
+  }
+}
+
 }  // namespace reference_ops
 }  // namespace tflite
 
diff --git a/tensorflow/lite/kernels/transpose_conv.cc b/tensorflow/lite/kernels/transpose_conv.cc
index 93c6df28890c9c..7d7c9a410ef451 100644
--- a/tensorflow/lite/kernels/transpose_conv.cc
+++ b/tensorflow/lite/kernels/transpose_conv.cc
@@ -26,8 +26,10 @@ limitations under the License.
 #include "tensorflow/lite/kernels/internal/optimized/integer_ops/transpose_conv.h"
 #include "tensorflow/lite/kernels/internal/optimized/optimized_ops.h"
 // NOLINTNEXTLINE - This header file shouldn't go to the top.
+#include "tensorflow/lite/kernels/internal/portable_tensor_utils.h"
 #include "tensorflow/lite/kernels/internal/reference/integer_ops/transpose_conv.h"
 #include "tensorflow/lite/kernels/internal/reference/reference_ops.h"
+#include "tensorflow/lite/kernels/internal/reference/transpose_conv.h"
 #include "tensorflow/lite/kernels/internal/tensor.h"
 #include "tensorflow/lite/kernels/internal/tensor_ctypes.h"
 #include "tensorflow/lite/kernels/internal/types.h"
@@ -59,6 +61,9 @@ struct OpData {
   int col2im_id = kTensorNotAllocated;
   int transposed_weights_id = kTensorNotAllocated;
   int scratch_tensor_id = kTensorNotAllocated;
+  int input_quantized_id = kTensorNotAllocated;
+  int scaling_factors_id = kTensorNotAllocated;
+  int input_offset_id = kTensorNotAllocated;
 
   // col2im is the temporary tensor allocated and used in optimized path for
   // storing col2im data:gemm result for input_matrix x filter_matrix.
@@ -73,6 +78,11 @@ struct OpData {
   // results.
   int32_t scratch_tensor_index;
 
+  // Indexes are used for hybrid (dynamic range quantization) path.
+  int32_t input_quantized_index;
+  int32_t scaling_factors_index;
+  int32_t input_offset_index;
+
   TfLitePaddingValues padding;
   // The scaling factor from input to output (aka the 'real multiplier') can
   // be represented as a fixed point multiplier plus a left shift.
@@ -160,6 +170,32 @@ static TfLiteStatus AllocateTemporaryTensorsIfRequired(TfLiteContext* context,
     ++temporaries_count;
   }
 
+  if (input_type == kTfLiteFloat32 && weights_type == kTfLiteInt8) {
+    // Allocate tensor to store the on-the-fly quantized inputs.
+    data->input_quantized_index = temporaries_count;
+    if (data->input_quantized_id == kTensorNotAllocated) {
+      TF_LITE_ENSURE_OK(
+          context, context->AddTensors(context, 1, &data->input_quantized_id));
+    }
+    ++temporaries_count;
+
+    // Allocate tensor to store the quantization params computed during
+    // on-the-fly input quantization.
+    data->scaling_factors_index = temporaries_count;
+    if (data->scaling_factors_id == kTensorNotAllocated) {
+      TF_LITE_ENSURE_OK(
+          context, context->AddTensors(context, 1, &data->scaling_factors_id));
+    }
+    ++temporaries_count;
+
+    data->input_offset_index = temporaries_count;
+    if (data->input_offset_id == kTensorNotAllocated) {
+      TF_LITE_ENSURE_OK(
+          context, context->AddTensors(context, 1, &data->input_offset_id));
+    }
+    ++temporaries_count;
+  }
+
   TfLiteIntArrayFree(node->temporaries);
   node->temporaries = TfLiteIntArrayCreate(temporaries_count);
 
@@ -308,8 +344,6 @@ TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
                                   bias->type == params->quantized_bias_type);
       data->quantized_bias_type = params->quantized_bias_type;
     }
-  } else {
-    TF_LITE_ENSURE_TYPES_EQ(context, weights->type, input->type);
   }
   TF_LITE_ENSURE_TYPES_EQ(context, output->type, input->type);
   // Ensure that weights and inputs have the same channel dimension.
@@ -406,6 +440,69 @@ TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
         data->per_channel_output_shift.data(), channels_out));
   }
 
+  if (input->type == kTfLiteFloat32 && weights->type == kTfLiteInt8) {
+    node->temporaries->data[data->input_quantized_index] =
+        data->input_quantized_id;
+    TfLiteTensor* input_quantized;
+    TF_LITE_ENSURE_OK(
+        context, GetTemporarySafe(context, node, data->input_quantized_index,
+                                  &input_quantized));
+    input_quantized->type = kTfLiteInt8;
+    input_quantized->allocation_type = kTfLiteArenaRw;
+    if (!TfLiteIntArrayEqual(input_quantized->dims, input->dims)) {
+      TfLiteIntArray* input_quantized_size = TfLiteIntArrayCopy(input->dims);
+      TF_LITE_ENSURE_OK(context, context->ResizeTensor(context, input_quantized,
+                                                       input_quantized_size));
+    }
+
+    node->temporaries->data[data->scaling_factors_index] =
+        data->scaling_factors_id;
+    TfLiteTensor* scaling_factors;
+    TF_LITE_ENSURE_OK(
+        context, GetTemporarySafe(context, node, data->scaling_factors_index,
+                                  &scaling_factors));
+    scaling_factors->type = kTfLiteFloat32;
+    scaling_factors->allocation_type = kTfLiteArenaRw;
+    // Only one scale factor per batch is typically necessary. See optimized
+    // implementation for why we need to allocate for the height of the inputs
+    // flattened to 2D.
+    const int channels_in = weights->dims->data[3];
+    TF_LITE_ENSURE(context, channels_in != 0);
+    const int height = NumElements(input) / channels_in;
+    int scaling_dims[1] = {height};
+    if (!TfLiteIntArrayEqualsArray(scaling_factors->dims, 1, scaling_dims)) {
+      TfLiteIntArray* scaling_factors_size = TfLiteIntArrayCreate(1);
+      scaling_factors_size->data[0] = height;
+      TF_LITE_ENSURE_OK(context, context->ResizeTensor(context, scaling_factors,
+                                                       scaling_factors_size));
+    }
+
+    const auto* affine_quantization =
+        reinterpret_cast<TfLiteAffineQuantization*>(
+            weights->quantization.params);
+    TF_LITE_ENSURE(context, affine_quantization);
+    TF_LITE_ENSURE(context, affine_quantization->scale);
+    TF_LITE_ENSURE_EQ(
+        context, affine_quantization->scale->size,
+        weights->dims->data[affine_quantization->quantized_dimension]);
+    node->temporaries->data[data->input_offset_index] = data->input_offset_id;
+    TfLiteTensor* input_offsets;
+    TF_LITE_ENSURE_OK(context,
+                      GetTemporarySafe(context, node, data->input_offset_index,
+                                       &input_offsets));
+    input_offsets->type = kTfLiteInt32;
+    input_offsets->allocation_type = kTfLiteArenaRw;
+    // See above comment for the need to allocate for height of inputs.
+    TF_LITE_ENSURE(context, channels_in != 0);
+    const int input_offset_dims[1] = {height};
+    if (!TfLiteIntArrayEqualsArray(input_offsets->dims, 1, input_offset_dims)) {
+      TfLiteIntArray* input_offsets_size = TfLiteIntArrayCreate(1);
+      input_offsets_size->data[0] = input_offset_dims[0];
+      TF_LITE_ENSURE_OK(context, context->ResizeTensor(context, input_offsets,
+                                                       input_offsets_size));
+    }
+  }
+
   return kTfLiteOk;
 }
 
@@ -617,6 +714,67 @@ void EvalQuantizedPerChannel16x8(
   }
 }
 
+TfLiteStatus EvalHybrid(TfLiteContext* context, TfLiteNode* node,
+                        const TfLiteTransposeConvParams* params, OpData* data,
+                        const TfLiteTensor* input, const TfLiteTensor* weights,
+                        const TfLiteTensor* bias, TfLiteTensor* output) {
+  float output_activation_min, output_activation_max;
+  CalculateActivationRange(params->activation, &output_activation_min,
+                           &output_activation_max);
+
+  const int batch_size = SizeOfDimension(input, 0);
+  TF_LITE_ENSURE(context, batch_size != 0);
+  const int input_size = NumElements(input) / batch_size;
+  TfLiteTensor* quantized_input_tensor;
+  TF_LITE_ENSURE_OK(context,
+                    GetTemporarySafe(context, node, data->input_quantized_index,
+                                     &quantized_input_tensor));
+  int8_t* quantized_input_ptr_batch =
+      GetTensorData<int8_t>(quantized_input_tensor);
+  TfLiteTensor* scaling_factors_tensor;
+  TF_LITE_ENSURE_OK(context,
+                    GetTemporarySafe(context, node, data->scaling_factors_index,
+                                     &scaling_factors_tensor));
+  float* scaling_factors_ptr = GetTensorData<float>(scaling_factors_tensor);
+  TfLiteTensor* input_offset_tensor;
+  TF_LITE_ENSURE_OK(context,
+                    GetTemporarySafe(context, node, data->input_offset_index,
+                                     &input_offset_tensor));
+  int32_t* input_offset_ptr = GetTensorData<int32_t>(input_offset_tensor);
+
+  for (int b = 0; b < batch_size; ++b) {
+    const int offset = b * input_size;
+    tensor_utils::AsymmetricQuantizeFloats(
+        GetTensorData<float>(input) + offset, input_size,
+        quantized_input_ptr_batch + offset, &scaling_factors_ptr[b],
+        &input_offset_ptr[b]);
+  }
+
+  const auto* affine_quantization =
+      reinterpret_cast<TfLiteAffineQuantization*>(weights->quantization.params);
+
+  tflite::ConvParams op_params;
+  op_params.padding_type = PaddingType::kSame;
+  op_params.padding_values.width = data->padding.width;
+  op_params.padding_values.height = data->padding.height;
+  op_params.padding_values.width_offset = data->padding.width_offset;
+  op_params.padding_values.height_offset = data->padding.height_offset;
+  op_params.stride_width = params->stride_width;
+  op_params.stride_height = params->stride_height;
+  op_params.float_activation_min = output_activation_min;
+  op_params.float_activation_max = output_activation_max;
+
+  reference_ops::HybridTransposeConv(
+      op_params, scaling_factors_ptr, GetTensorShape(input),
+      quantized_input_ptr_batch, GetTensorShape(weights),
+      GetTensorData<int8_t>(weights), GetTensorShape(bias),
+      GetTensorData<float>(bias), GetTensorShape(output),
+      GetTensorData<float>(output), affine_quantization->scale->data,
+      input_offset_ptr);
+
+  return kTfLiteOk;
+}
+
 template <KernelType kernel_type>
 TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) {
   // Retrieve tensors (All should be allocated by now)
@@ -677,14 +835,19 @@ TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) {
   // Currently support float32, uint8, int8, int16.
   switch (input->type) {
     case kTfLiteFloat32: {
-      // Only for GenericOptimized path, we use transposed weights.
-      if (data->weights_are_transposed) {
-        if (!IsConstantTensor(weights)) {
-          ResizeAndTransposeWeights(context, weights, transposed_weights);
+      if (weights->type == kTfLiteInt8) {
+        TF_LITE_ENSURE_OK(context, EvalHybrid(context, node, params, data,
+                                              input, weights, bias, output));
+      } else {
+        // Only for GenericOptimized path, we use transposed weights.
+        if (data->weights_are_transposed) {
+          if (!IsConstantTensor(weights)) {
+            ResizeAndTransposeWeights(context, weights, transposed_weights);
+          }
         }
+        EvalFloat<kernel_type>(context, params, data, input, weights, bias,
+                               transposed_weights, col2im, output);
       }
-      EvalFloat<kernel_type>(context, params, data, input, weights, bias,
-                             transposed_weights, col2im, output);
       break;
     }
     case kTfLiteUInt8: {
diff --git a/tensorflow/lite/kernels/transpose_conv_test.cc b/tensorflow/lite/kernels/transpose_conv_test.cc
index 84dfd29dfb9945..bbe7d3f022c39b 100644
--- a/tensorflow/lite/kernels/transpose_conv_test.cc
+++ b/tensorflow/lite/kernels/transpose_conv_test.cc
@@ -1251,6 +1251,102 @@ TEST_P(TransposeConvOpTest, SimpleBiasTestQuantizedPerChannel16x8Bias64) {
   EXPECT_THAT(model.GetOutputShape(), ElementsAreArray({1, 2, 3, 2}));
 }
 
+class HybridTransposeConvOpModel
+    : public BaseTransposeConvBiasOpModel<float, int8_t> {
+ public:
+  using BaseTransposeConvBiasOpModel::BaseTransposeConvBiasOpModel;
+
+  void SetFilter(std::initializer_list<float> f) {
+    PerChannelSymmetricQuantizeAndPopulate(filter_, f);
+  }
+
+  void SetBias(std::initializer_list<float> b) { PopulateTensor(bias_, b); }
+
+  std::vector<float> GetOutput() { return ExtractVector<float>(output_); }
+};
+
+TEST_P(TransposeConvOpTest, SimpleTestHybridInt8) {
+  const std::initializer_list<float> filter_data = {1, 2, 3, 4, 5, 6, 7, 8, 9};
+  const std::initializer_list<int8_t> const_filter_data = {14, 28, 42,  56, 71,
+                                                           85, 99, 113, 127};
+  HybridTransposeConvOpModel model(
+      /*registration=*/GetRegistration(), /*output_shape_data=*/{1, 4, 4, 1},
+      /*filter=*/
+      {TensorType_INT8, {1, 3, 3, 1}, 0, 0, 0, 0, true, {9.0 / 127}, {0}, 0},
+      /*filter_data=*/const_filter_data,
+      /*input=*/{TensorType_FLOAT32, {1, 4, 4, 1}},
+      /*output=*/{TensorType_FLOAT32, {}},
+      /*padding=*/Padding_SAME, /*stride_w=*/1, /*stride_h=*/1,
+      /*fused_activation=*/ActivationFunctionType_NONE,
+      /*test_type=*/GetTestType(),
+      /*version=*/3,
+      /*bias_type=*/TensorType_FLOAT32);
+  model.SetInput({1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16});
+  if (GetTestType() == TestType::kDynamic) {
+    model.SetFilter(filter_data);
+  }
+
+  model.SetBias({1});
+  ASSERT_EQ(model.Invoke(), kTfLiteOk);
+
+  // The values are taken from float model "SimpleTest".
+  EXPECT_THAT(model.GetOutput(), ElementsAreArray(ArrayFloatNear(
+                                     {30, 63, 84, 76, 100, 193, 238, 199, 208,
+                                      373, 417.5, 331, 263.7, 447, 486, 366.5},
+                                     0.19)));
+
+  // GetOutputShape() should always be same as model.SetOutputShape(...);
+  EXPECT_THAT(model.GetOutputShape(), ElementsAreArray({1, 4, 4, 1}));
+}
+
+TEST_P(TransposeConvOpTest, SimpleTestHybridInt8MultiChannel) {
+  const std::initializer_list<float> filter_data = {
+      1, 3, 5, 7, 9, 11, 13, 15, 17, 2, 4, 6, 8, 10, 12, 14, 16, 18};
+  const std::initializer_list<int8_t> const_filter_data = {
+      7,  22, 37, 52, 67, 82, 97, 112, 127,
+      14, 28, 42, 56, 71, 85, 99, 113, 127};
+  HybridTransposeConvOpModel model(
+      /*registration=*/GetRegistration(), /*output_shape_data=*/{1, 5, 5, 2},
+      /*filter=*/
+      {TensorType_INT8,
+       {2, 3, 3, 1},
+       0,
+       0,
+       0,
+       0,
+       true,
+       {17.0 / 127, 18.0 / 127},
+       {0, 0},
+       0},
+      /*filter_data=*/const_filter_data,
+      /*input=*/{TensorType_FLOAT32, {1, 2, 2, 1}},
+      /*output=*/{TensorType_FLOAT32, {}},
+      /*padding=*/Padding_VALID, /*stride_w=*/2, /*stride_h=*/2,
+      /*fused_activation=*/ActivationFunctionType_NONE,
+      /*test_type=*/GetTestType(),
+      /*version=*/3,
+      /*bias_type=*/TensorType_FLOAT32);
+
+  model.SetInput({1, 2, 3, 4});
+  if (GetTestType() == TestType::kDynamic) {
+    model.SetFilter(filter_data);
+  }
+  model.SetBias({3, 4});
+
+  ASSERT_EQ(model.Invoke(), kTfLiteOk);
+
+  // The values are taken from float model "MultiChannelBiasTest".
+  EXPECT_THAT(
+      model.GetOutput(),
+      ElementsAreArray(ArrayFloatNear(
+          {4,    6,  6,    8,    10, 14,  9,   12, 13,   16, 10, 12,   12,
+           14,   28, 32,   21,   24, 25,  28,  19, 24,   27, 32, 64.5, 76,
+           44.5, 52, 56.5, 63.5, 24, 28,  30,  34, 63.5, 72, 39, 44,   47,
+           52,   42, 46,   48,   52, 106, 114, 63, 68,   71, 76},
+          0.26)));
+  EXPECT_THAT(model.GetOutputShape(), ElementsAreArray({1, 5, 5, 2}));
+}
+
 INSTANTIATE_TEST_SUITE_P(
     TransposeConvOpTest, TransposeConvOpTest,
     ::testing::Combine(

From 382d00d7b4e250586c617f0c0b59644867c4e149 Mon Sep 17 00:00:00 2001
From: gaurides <gauri1.deshpande@intel.com>
Date: Wed, 15 May 2024 03:40:54 -0700
Subject: [PATCH 172/478] PR #12436: [XLA:CPU][oneDNN] Enable Matmul + BiasAdd
 + Elu / Tanh / Relu6 fusions for onednn

Imported from GitHub PR https://github.com/openxla/xla/pull/12436

This PR enables MatMul  fusions with different activations viz. Elu, Tanh and Relu6.
Copybara import of the project:

--
9895b022aa66ee0f1040d0f6e557382c9dd77b99 by Gauri1 Deshpande <gauri1.deshpande@intel.com>:

Enable Matmul + BiasAdd + Elu / Tanh / Relu6 fusions for onednn

--
c277c70da3a1bbe05d0d637ec2d4e9ce220f7f3f by Gauri1 Deshpande <gauri1.deshpande@intel.com>:

address review comments - sort name, remove default layout and remove default attributes

Merging this change closes #12436

PiperOrigin-RevId: 633881207
---
 third_party/xla/xla/service/cpu/BUILD         |  11 +
 .../xla/xla/service/cpu/backend_config.proto  |   2 +
 .../xla/xla/service/cpu/onednn_matmul.cc      |   6 +
 .../xla/service/cpu/onednn_matmul_rewriter.cc | 122 +++++---
 .../xla/service/cpu/onednn_ops_rewriter.cc    |  21 +-
 .../xla/service/cpu/onednn_pattern_utils.h    |  60 ++++
 third_party/xla/xla/service/pattern_matcher.h |   1 +
 .../xla/xla/tests/onednn_matmul_test.cc       | 275 ++++++++++++++++++
 8 files changed, 451 insertions(+), 47 deletions(-)
 create mode 100644 third_party/xla/xla/service/cpu/onednn_pattern_utils.h

diff --git a/third_party/xla/xla/service/cpu/BUILD b/third_party/xla/xla/service/cpu/BUILD
index 1d0defbdcced10..3f2bdc0f211410 100644
--- a/third_party/xla/xla/service/cpu/BUILD
+++ b/third_party/xla/xla/service/cpu/BUILD
@@ -1658,6 +1658,15 @@ cc_library(
     ] + mkl_deps(),
 )
 
+cc_library(
+    name = "onednn_pattern_utils",
+    hdrs = ["onednn_pattern_utils.h"],
+    visibility = ["//visibility:public"],
+    deps = [
+        ":onednn_util",
+    ] + mkl_deps(),
+)
+
 cc_library(
     name = "onednn_matmul_rewriter",
     srcs = ["onednn_matmul_rewriter.cc"],
@@ -1671,6 +1680,7 @@ cc_library(
         ":backend_config_proto_cc",
         ":onednn_matmul",
         ":onednn_memory_util",
+        ":onednn_pattern_utils",
         ":onednn_util",
         "//xla:executable_run_options",
         "//xla:shape_util",
@@ -1699,6 +1709,7 @@ cc_library(
     deps = [
         ":backend_config_proto_cc",
         ":onednn_memory_util",
+        ":onednn_pattern_utils",
         ":onednn_util",
         "//xla:status_macros",
         "//xla:xla_data_proto_cc",
diff --git a/third_party/xla/xla/service/cpu/backend_config.proto b/third_party/xla/xla/service/cpu/backend_config.proto
index d5ce59285e55c9..b63bcd5c376014 100644
--- a/third_party/xla/xla/service/cpu/backend_config.proto
+++ b/third_party/xla/xla/service/cpu/backend_config.proto
@@ -27,6 +27,8 @@ message OneDnnMatMulConfig {
     GELU_TANH = 5;
     BINARY_ADD = 6;
     LINEAR = 7;
+    ELU = 8;
+    RELU6 = 9;
   }
   repeated FusionKind fused_ops = 3;
   bool bias_broadcast = 4;
diff --git a/third_party/xla/xla/service/cpu/onednn_matmul.cc b/third_party/xla/xla/service/cpu/onednn_matmul.cc
index 04e0bd47a2fb00..4a6484ec07d405 100644
--- a/third_party/xla/xla/service/cpu/onednn_matmul.cc
+++ b/third_party/xla/xla/service/cpu/onednn_matmul.cc
@@ -138,6 +138,9 @@ std::unique_ptr<matmul::primitive_desc> CreateMatMulPrimDesc(
       case OneDnnMatMulConfig::GELU_ERF:
         post_ops.append_eltwise(dnnl::algorithm::eltwise_gelu_erf, 0.f, 0.f);
         break;
+      case OneDnnMatMulConfig::RELU6:
+        post_ops.append_eltwise(dnnl::algorithm::eltwise_clip_v2, 0.f, 6.0f);
+        break;
       case OneDnnMatMulConfig::BIAS: {
         bias_md = fused_mds.at(fused_operand_idx);
         // Extend bias rank to match result rank.
@@ -156,6 +159,9 @@ std::unique_ptr<matmul::primitive_desc> CreateMatMulPrimDesc(
         }
         fused_operand_idx++;
       } break;
+      case OneDnnMatMulConfig::ELU:
+        post_ops.append_eltwise(dnnl::algorithm::eltwise_elu, 1.0f, 0.0f);
+        break;
       case OneDnnMatMulConfig::BINARY_ADD: {
         auto binary_md = fused_mds.at(fused_operand_idx);
         if (fused_operands_ref) {
diff --git a/third_party/xla/xla/service/cpu/onednn_matmul_rewriter.cc b/third_party/xla/xla/service/cpu/onednn_matmul_rewriter.cc
index cc97076a509af6..a792a5ba52130d 100644
--- a/third_party/xla/xla/service/cpu/onednn_matmul_rewriter.cc
+++ b/third_party/xla/xla/service/cpu/onednn_matmul_rewriter.cc
@@ -28,6 +28,7 @@ limitations under the License.
 #include "xla/service/cpu/backend_config.pb.h"
 #include "xla/service/cpu/onednn_matmul.h"
 #include "xla/service/cpu/onednn_memory_util.h"
+#include "xla/service/cpu/onednn_pattern_utils.h"
 #include "xla/service/cpu/onednn_util.h"
 #include "xla/service/hlo_cost_analysis.h"
 #include "xla/service/pattern_matcher.h"
@@ -40,6 +41,7 @@ namespace cpu {
 
 namespace {
 namespace m = match;
+namespace pu = ::xla::cpu::onednn_pattern_utils_internal;
 
 inline Status ValidateDotDimensionNumbers(
     const DotDimensionNumbers& dim_numbers) {
@@ -63,26 +65,6 @@ inline bool CompatibleElementType(const HloInstruction* instr) {
   return element_type == BF16 || element_type == F32 || element_type == F16;
 }
 
-// Type conversion from and to any of BF16 and FP32.
-// TODO(intel-tf): Support more types when enabled.
-template <typename Pattern>
-inline auto SupportedConvert(Pattern pattern) {
-  auto supported_convert = [](const HloInstruction* instr) -> bool {
-    return CompatibleElementType(instr) &&
-           CompatibleElementType(instr->operand(0));
-  };
-  return m::Convert(pattern).WithPredicate(supported_convert);
-}
-
-template <typename Pattern>
-inline auto SupportedConvert(HloInstruction** convert, Pattern pattern) {
-  auto supported_convert = [](const HloInstruction* instr) -> bool {
-    return CompatibleElementType(instr) &&
-           CompatibleElementType(instr->operand(0));
-  };
-  return m::Convert(convert, pattern).WithPredicate(supported_convert);
-}
-
 inline bool IsRowMajor(const Shape& shape) {
   return LayoutUtil::IsMonotonicWithDim0Major(shape.layout());
 }
@@ -117,9 +99,9 @@ auto ElementwiseSafeIntermediates(HloInstruction** instr,
       m::Slice(instr, pattern.WithOneUser()),
       m::Bitcast(instr, pattern.WithOneUser()),
       m::Reshape(instr, pattern.WithOneUser()),
-      SupportedConvert(instr, pattern.WithOneUser()),
-      SupportedConvert(instr, BitcastWithReshapeSemantics(
-                                  optional_bitcast, pattern.WithOneUser())),
+      pu::SupportedConvert(instr, pattern.WithOneUser()),
+      pu::SupportedConvert(instr, BitcastWithReshapeSemantics(
+                                      optional_bitcast, pattern.WithOneUser())),
       pattern);
 }
 
@@ -140,6 +122,10 @@ inline auto BcastConstScalar(double value) {
   return BcastConstScalar(nullptr, value);
 }
 
+inline auto BcastConvertConstScalar(double value) {
+  return m::Broadcast(pu::OptionalConvert(m::ConstantScalar(value)));
+}
+
 inline bool IsBatchDot(const HloInstruction& instr) {
   if (auto* dot_instr = DynCast<HloDotInstruction>(&instr)) {
     return dot_instr->dot_dimension_numbers().lhs_batch_dimensions_size() > 0;
@@ -368,7 +354,7 @@ inline auto OptionalConvertAndBitcast(HloInstruction** optional_convert,
   //   3. pattern-root -> bitcast
   //   4. pattern-root
   auto common = m::AnyOf<HloInstruction>(
-      SupportedConvert(optional_convert, std::move(pattern).WithOneUser())
+      pu::SupportedConvert(optional_convert, std::move(pattern).WithOneUser())
           .WithOperand(0, m::Op().WithElementType(PrimitiveType::BF16))
           .WithElementType(PrimitiveType::F32),
       std::move(pattern).WithOneUser());
@@ -562,7 +548,8 @@ class OneDnnMatMulRewriteVisitor : public DfsHloRewriteVisitor {
       }
 
       // Validate addend for fusion.
-      if (CompatibleElementType(addend) && IsOperandFusible(addend, dot)) {
+      if (IsSupportedType(addend->shape().element_type()) &&
+          IsOperandFusible(addend, dot)) {
         new_operands.push_back(addend);
       } else {
         return OkStatus();
@@ -647,6 +634,71 @@ class OneDnnMatMulRewriteVisitor : public DfsHloRewriteVisitor {
     return OkStatus();
   }
 
+  auto ELUActivation(HloInstruction* instr, HloInstruction** src) {
+    //  Reference: tensorflow/compiler/tf2xla/kernels/elu_op.cc
+    //  const auto zero = ScalarLike(x, 0);
+    //  const auto pred = Gt(x, zero);
+    //  const auto expm1 = Expm1(x);
+    //  return Select(pred, x, expm1);
+    auto pattern = m::Select(
+        m::Gt(pu::OptionalConvert(m::Op(src)), BcastConvertConstScalar(0)),
+        m::Op(src),
+        pu::OptionalConvert(m::Expm1(pu::OptionalConvert(m::Op(src)))));
+    return Match(instr, pattern);
+  }
+
+  Status HandleSelect(HloInstruction* instr) override {
+    HloInstruction* matmul_call;
+    HloInstruction* intermediate_instr = nullptr;
+    HloInstruction* optional_bitcast = nullptr;
+    HloInstruction* src;
+    // Attempt to elide ELU subgraph and fuse ELU activation into GEMM,
+    // including when slicing or bitcasting is applied to the result.
+    if (ELUActivation(instr, &src)) {
+      if (Match(src, ElementwiseSafeIntermediates(
+                         &intermediate_instr, &optional_bitcast,
+                         OneDnnMatmulInstr(&matmul_call)))) {
+        return FuseActivation(OneDnnMatMulConfig::ELU, instr, matmul_call,
+                              intermediate_instr);
+      }
+    }
+    return OkStatus();
+  }
+
+  Status HandleTanh(HloInstruction* instr) override {
+    HloInstruction* matmul_call;
+    HloInstruction* intermediate_instr = nullptr;
+    HloInstruction* optional_bitcast = nullptr;
+    // Attempt to elide Tanh and fuse Tanh activation into GEMM, including
+    // when slicing or bitcasting is applied to the result.
+    if (Match(instr, m::Tanh(ElementwiseSafeIntermediates(
+                                 &intermediate_instr, &optional_bitcast,
+                                 OneDnnMatmulInstr(&matmul_call))
+                                 .WithOneUser()))) {
+      return FuseActivation(OneDnnMatMulConfig::TANH, instr, matmul_call,
+                            intermediate_instr);
+    }
+    return OkStatus();
+  }
+
+  Status HandleClamp(HloInstruction* instr) override {
+    HloInstruction* matmul_call;
+    HloInstruction* intermediate_instr = nullptr;
+    HloInstruction* optional_bitcast = nullptr;
+    // Attempt to elide RELU6 and fuse RELU6 activation into GEMM, including
+    // when slicing or bitcasting is applied to the result.
+    if (Match(instr, m::Clamp(BcastConstScalar(0),
+                              ElementwiseSafeIntermediates(
+                                  &intermediate_instr, &optional_bitcast,
+                                  OneDnnMatmulInstr(&matmul_call))
+                                  .WithOneUser(),
+                              BcastConstScalar(6)))) {
+      return FuseActivation(OneDnnMatMulConfig::RELU6, instr, matmul_call,
+                            intermediate_instr);
+    }
+    return OkStatus();
+  }
+
   Status HandleMultiply(HloInstruction* instr) override {
     HloInstruction* matmul_call;
     HloInstruction* intermediate_instr = nullptr;
@@ -664,16 +716,16 @@ class OneDnnMatMulRewriteVisitor : public DfsHloRewriteVisitor {
 
     HloInstruction *dot, *constant;
     HloInstruction* optional_convert = nullptr;
-    auto pattern =
-        m::Op(&instr)
-            .WithOpcode(HloOpcode::kMultiply)
-            .WithBinaryOperandsAnyOrder(
-                m::AnyOf<HloInstruction>(
-                    SupportedConvert(&optional_convert, OneDnnMatmulInstr(&dot))
-                        .WithElementType(PrimitiveType::F32),
-                    OneDnnMatmulInstr(&dot))
-                    .WithOneUser(),
-                m::Broadcast(m::Constant(&constant)));
+    auto pattern = m::Op(&instr)
+                       .WithOpcode(HloOpcode::kMultiply)
+                       .WithBinaryOperandsAnyOrder(
+                           m::AnyOf<HloInstruction>(
+                               pu::SupportedConvert(&optional_convert,
+                                                    OneDnnMatmulInstr(&dot))
+                                   .WithElementType(PrimitiveType::F32),
+                               OneDnnMatmulInstr(&dot))
+                               .WithOneUser(),
+                           m::Broadcast(m::Constant(&constant)));
 
     if (Match(instr, pattern)) {
       std::vector<HloInstruction*> new_operands;
diff --git a/third_party/xla/xla/service/cpu/onednn_ops_rewriter.cc b/third_party/xla/xla/service/cpu/onednn_ops_rewriter.cc
index 058355223cd5ee..06c2137a782d50 100644
--- a/third_party/xla/xla/service/cpu/onednn_ops_rewriter.cc
+++ b/third_party/xla/xla/service/cpu/onednn_ops_rewriter.cc
@@ -19,6 +19,7 @@ limitations under the License.
 #include "xla/hlo/ir/hlo_instruction.h"
 #include "xla/service/cpu/backend_config.pb.h"
 #include "xla/service/cpu/onednn_memory_util.h"
+#include "xla/service/cpu/onednn_pattern_utils.h"
 #include "xla/service/cpu/onednn_util.h"
 #include "xla/service/pattern_matcher.h"
 #include "xla/status_macros.h"
@@ -28,11 +29,7 @@ namespace cpu {
 
 namespace {
 namespace m = match;
-
-template <typename Pattern>
-auto OptionalConvert(Pattern pattern) {
-  return m::AnyOf<HloInstruction>(m::Convert(pattern), std::move(pattern));
-}
+namespace pu = ::xla::cpu::onednn_pattern_utils_internal;
 
 inline auto OneDnnConvertibleInstr(HloInstruction** instr) {
   return m::AnyOf<HloInstruction>(m::CustomCall(instr, {"__onednn$layernorm"}),
@@ -101,9 +98,9 @@ std::optional<HloInstruction*> MatchSoftmax(HloInstruction* instr) {
   if (!Match(instr,
              m::Divide(
                  m::Exp(&left_exponential, m::Op()),
-                 m::Broadcast(m::Reshape(
-                     m::Broadcast(OptionalConvert(m::Reshape(OptionalConvert(
-                         m::Reduce(OptionalConvert(
+                 m::Broadcast(m::Reshape(m::Broadcast(
+                     pu::OptionalConvert(m::Reshape(pu::OptionalConvert(
+                         m::Reduce(pu::OptionalConvert(
                                        m::Exp(&right_exponential, m::Op())),
                                    m::Op())
                              .WithPredicate([](const HloInstruction* reduce) {
@@ -256,7 +253,7 @@ bool MatchFlaxLayerNorm(HloInstruction* instr, HloInstruction** src,
           .WithBinaryOperandsAnyOrder(
               m::Op(&hinge).WithOneUser(),
               m::Subtract(
-                  OptionalConvert(m::Op(&prod_s)),
+                  pu::OptionalConvert(m::Op(&prod_s)),
                   m::Broadcast(
                       m::Reshape(
                           m::Broadcast(m::Reshape(m::Op(&div_red).WithOpcode(
@@ -326,8 +323,8 @@ bool MatchFlaxLayerNorm(HloInstruction* instr, HloInstruction** src,
   auto div_red_mul_src =
       m::Divide()
           .WithOperand(0, m::Reduce(m::Multiply().WithBinaryOperandsAnyOrder(
-                                        OptionalConvert(m::Op(&mul_in0)),
-                                        OptionalConvert(m::Op(&mul_in1))),
+                                        pu::OptionalConvert(m::Op(&mul_in0)),
+                                        pu::OptionalConvert(m::Op(&mul_in1))),
                                     m::Constant())
                               .WithPredicate([](const HloInstruction* reduce) {
                                 HloComputation* reducer = reduce->to_apply();
@@ -348,7 +345,7 @@ bool MatchFlaxLayerNorm(HloInstruction* instr, HloInstruction** src,
       m::Divide()
           .WithOperand(
               0,
-              m::Reduce(OptionalConvert(m::Op(&reduce_in0)), m::Constant())
+              m::Reduce(pu::OptionalConvert(m::Op(&reduce_in0)), m::Constant())
                   .WithPredicate([](const HloInstruction* reduce) {
                     HloComputation* reducer = reduce->to_apply();
                     return (reducer->root_instruction()->opcode() ==
diff --git a/third_party/xla/xla/service/cpu/onednn_pattern_utils.h b/third_party/xla/xla/service/cpu/onednn_pattern_utils.h
new file mode 100644
index 00000000000000..6635d79d115938
--- /dev/null
+++ b/third_party/xla/xla/service/cpu/onednn_pattern_utils.h
@@ -0,0 +1,60 @@
+/* Copyright 2024 The OpenXLA Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+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.
+==============================================================================*/
+
+#ifndef XLA_SERVICE_CPU_ONEDNN_PATTERN_UTILS_H_
+#define XLA_SERVICE_CPU_ONEDNN_PATTERN_UTILS_H_
+#if defined(INTEL_MKL) && defined(ENABLE_ONEDNN_V3)
+
+#include "xla/hlo/ir/hlo_instruction.h"
+#include "xla/hlo/ir/hlo_instructions.h"
+#include "xla/service/cpu/onednn_util.h"
+#include "xla/service/pattern_matcher.h"
+
+namespace xla {
+namespace cpu {
+
+namespace onednn_pattern_utils_internal {
+namespace m = match;
+
+template <typename Pattern>
+auto OptionalConvert(Pattern pattern) {
+  return m::AnyOf<HloInstruction>(m::Convert(pattern), std::move(pattern));
+}
+
+// Type conversion from and to any of BF16 and FP32.
+// TODO(intel-tf): Support more types when enabled.
+template <typename Pattern>
+inline auto SupportedConvert(Pattern pattern) {
+  auto supported_convert = [](const HloInstruction* instr) -> bool {
+    return IsSupportedType(instr->shape().element_type()) &&
+           IsSupportedType(instr->operand(0)->shape().element_type());
+  };
+  return m::Convert(pattern).WithPredicate(supported_convert);
+}
+
+template <typename Pattern>
+inline auto SupportedConvert(HloInstruction** convert, Pattern pattern) {
+  auto supported_convert = [](const HloInstruction* instr) -> bool {
+    return IsSupportedType(instr->shape().element_type()) &&
+           IsSupportedType(instr->operand(0)->shape().element_type());
+  };
+  return m::Convert(convert, pattern).WithPredicate(supported_convert);
+}
+}  // namespace onednn_pattern_utils_internal
+}  // namespace cpu
+}  // namespace xla
+
+#endif  // INTEL_MKL && ENABLE_ONEDNN_V3
+#endif  // XLA_SERVICE_CPU_ONEDNN_PATTERN_UTILS_H_
diff --git a/third_party/xla/xla/service/pattern_matcher.h b/third_party/xla/xla/service/pattern_matcher.h
index 477a63e2a6f91d..b17c53a9baf699 100644
--- a/third_party/xla/xla/service/pattern_matcher.h
+++ b/third_party/xla/xla/service/pattern_matcher.h
@@ -2687,6 +2687,7 @@ XLA_UNOP_PATTERN(CollectivePermuteStart)
 XLA_UNOP_PATTERN(CollectivePermuteDone)
 XLA_UNOP_PATTERN(Domain)
 XLA_UNOP_PATTERN(Exp)
+XLA_UNOP_PATTERN(Expm1)
 XLA_UNOP_PATTERN(Fft)
 XLA_UNOP_PATTERN(Floor)
 XLA_UNOP_PATTERN(GetTupleElement)
diff --git a/third_party/xla/xla/tests/onednn_matmul_test.cc b/third_party/xla/xla/tests/onednn_matmul_test.cc
index 22089f177e3666..2248b05f5ddcdf 100644
--- a/third_party/xla/xla/tests/onednn_matmul_test.cc
+++ b/third_party/xla/xla/tests/onednn_matmul_test.cc
@@ -81,6 +81,34 @@ class MatmulTest : public HloTestBase {
     ; CHECK-DAG:   }
     ; CHECK:     }
     )";
+
+  const char* fused_matmul_bias_elu_rewrite_str_ = R"(
+    ; CHECK:     custom_call_target="__onednn$matmul",
+    ; CHECK:       backend_config={
+    ; CHECK-DAG:     "outer_dimension_partitions":[],
+    ; CHECK-DAG:     "onednn_matmul_config":{
+    ; CHECK-DAG:       "fused_ops":["BIAS","ELU"]
+    ; CHECK-DAG:   }
+    ; CHECK:     }
+    )";
+  const char* fused_matmul_bias_tanh_rewrite_str_ = R"(
+    ; CHECK:     custom_call_target="__onednn$matmul",
+    ; CHECK:       backend_config={
+    ; CHECK-DAG:     "outer_dimension_partitions":[],
+    ; CHECK-DAG:     "onednn_matmul_config":{
+    ; CHECK-DAG:       "fused_ops":["BIAS","TANH"]
+    ; CHECK-DAG:   }
+    ; CHECK:     }
+    )";
+  const char* fused_matmul_bias_relu6_rewrite_str_ = R"(
+    ; CHECK:     custom_call_target="__onednn$matmul",
+    ; CHECK:       backend_config={
+    ; CHECK-DAG:     "outer_dimension_partitions":[],
+    ; CHECK-DAG:     "onednn_matmul_config":{
+    ; CHECK-DAG:       "fused_ops":["BIAS","RELU6"]
+    ; CHECK-DAG:   }
+    ; CHECK:     }
+    )";
 };
 
 TEST_F(MatmulTest, SimpleTestF32) {
@@ -923,6 +951,117 @@ TEST_F(MatmulTest, TestF32ConstantWeights) {
   )");
 }
 
+// MM + BiasAdd + Elu : FP32
+TEST_F(MatmulTest, BiasAddELUFusion_F32) {
+  const char* matmul_module_str = R"(
+  HloModule matmul.test.f32
+
+  ENTRY matmul.test.f32 {
+    arg0.1 = f32[1024,1024] parameter(0)
+    arg1.2 = f32[1024,1024] parameter(1)
+    dot.3 = f32[1024,1024] dot(arg1.2, arg0.1), lhs_contracting_dims={1}, rhs_contracting_dims={0}
+    arg2.4 = f32[1024] parameter(2)
+    broadcast.5 = f32[1024,1024] broadcast(arg2.4), dimensions={1}
+    add.6 = f32[1024,1024] add(dot.3, broadcast.5)
+    constant.7 = f32[] constant(0)
+    broadcast.8 = f32[1024,1024] broadcast(constant.7), dimensions={}
+    compare.9 = pred[1024,1024] compare(add.6, broadcast.8), direction=GT
+    exponential-minus-one.10 = f32[1024,1024] exponential-minus-one(add.6)
+    ROOT select.11 = f32[1024,1024] select(compare.9, add.6, exponential-minus-one.10)
+  })";
+
+  EXPECT_TRUE(RunAndCompare(matmul_module_str, ErrorSpec{1e-4, 1e-4}));
+  MatchOptimizedHlo(matmul_module_str, fused_matmul_bias_elu_rewrite_str_);
+}
+
+// MM + BiasAdd + Elu : BF16
+TEST_F(MatmulTest, BiasAddELUFusion_BF16) {
+  if (!IsSupportedType(PrimitiveType::BF16)) {
+    GTEST_SKIP() << "CPU does not support BF16.";
+  }
+  const char* matmul_module_str = R"(
+  HloModule matmul.test.bf16
+  ENTRY matmul.test.bf16 {
+    arg0.1 = f32[1024,512] parameter(0)
+    convert.2 = bf16[1024,512] convert(arg0.1)
+    arg1.3 = f32[256,512] parameter(1)
+    convert.4 = bf16[256,512] convert(arg1.3)
+    dot.5 = bf16[1024,256] dot(convert.2, convert.4), lhs_contracting_dims={1}, rhs_contracting_dims={1}
+    convert.6 = f32[1024,256] convert(dot.5)
+    arg2.7 = f32[256] parameter(2)
+    broadcast.8 = f32[1024,256] broadcast(arg2.7), dimensions={1}
+    add.9 = f32[1024,256] add(convert.6, broadcast.8)
+    constant.10 = f32[] constant(0)
+    broadcast.11 = f32[1024,256] broadcast(constant.10), dimensions={}
+    compare.12 = pred[1024,256] compare(add.9, broadcast.11), direction=GT
+    convert.13 = bf16[1024,256] convert(add.9)
+    exponential-minus-one.14 = f32[1024,256] exponential-minus-one(add.9)
+    convert.15 = bf16[1024,256] convert(exponential-minus-one.14)
+    select.16 = bf16[1024,256] select(compare.12, convert.13, convert.15)
+    ROOT convert.17 = f32[1024,256] convert(select.16)
+  })";
+
+  EXPECT_TRUE(RunAndCompare(matmul_module_str, ErrorSpec{1e-2, 1e-2}));
+  MatchOptimizedHlo(matmul_module_str, fused_matmul_bias_elu_rewrite_str_);
+}
+
+// MM + BiasAdd + Elu : F16
+TEST_F(MatmulTest, BiasAddELUFusion_F16) {
+  if (!IsSupportedType(PrimitiveType::F16)) {
+    GTEST_SKIP() << "CPU does not support F16.";
+  }
+  const char* matmul_module_str = R"(
+  HloModule matmul.test.f16
+
+  ENTRY matmul.test.f16 {
+    arg0.1 = f16[1024,1024] parameter(0)
+    arg1.2 = f16[1024,1024] parameter(1)
+    dot.3 = f16[1024,1024] dot(arg1.2, arg0.1), lhs_contracting_dims={1}, rhs_contracting_dims={0}
+    arg2.4 = f16[1024] parameter(2)
+    broadcast.5 = f16[1024,1024] broadcast(arg2.4), dimensions={1}
+    add.6 = f16[1024,1024] add(dot.3, broadcast.5)
+    constant.7 = f16[] constant(0)
+    broadcast.8 = f16[1024,1024] broadcast(constant.7), dimensions={}
+    compare.9 = pred[1024,1024] compare(add.6, broadcast.8), direction=GT
+    exponential-minus-one.10 = f16[1024,1024] exponential-minus-one(add.6)
+    ROOT select.11 = f16[1024,1024] select(compare.9, add.6, exponential-minus-one.10)
+  })";
+
+  EXPECT_TRUE(RunAndCompare(matmul_module_str, ErrorSpec{1e-2, 1e-2}));
+  MatchOptimizedHlo(matmul_module_str, fused_matmul_bias_elu_rewrite_str_);
+}
+
+// MM + BiasAdd + Elu + MM : FP16_2
+TEST_F(MatmulTest, BiasAddELUFusion_F16_2) {
+  if (!IsSupportedType(PrimitiveType::F16)) {
+    GTEST_SKIP() << "CPU does not support F16.";
+  }
+  const char* matmul_module_str = R"(
+  HloModule matmul.test.f16
+
+  ENTRY matmul.test.f16 {
+    arg0.1 = f32[1024,1024] parameter(0)
+    convert.2 = f16[1024,1024] convert(arg0.1)
+    arg1.3 = f32[1024,1024] parameter(2)
+    convert.4 = f16[1024,1024] convert(arg1.3)
+    dot.5 = f16[1024,1024] dot(convert.2, convert.4), lhs_contracting_dims={1}, rhs_contracting_dims={0}
+    arg2.6 = f32[1024] parameter(1)
+    convert.7 = f16[1024] convert(arg2.6)
+    broadcast.8 = f16[1024,1024] broadcast(convert.7), dimensions={1}
+    add.9 = f16[1024,1024] add(dot.5, broadcast.8)
+    constant.10 = f16[] constant(0)
+    broadcast.11 = f16[1024,1024] broadcast(constant.10), dimensions={}
+    compare.12 = pred[1024,1024] compare(add.9, broadcast.11), direction=GT
+    exponential-minus-one.13 = f16[1024,1024] exponential-minus-one(add.9)
+    select.14 = f16[1024,1024] select(compare.12, add.9, exponential-minus-one.13)
+    dot.15 = f16[1024,1024] dot(select.14, convert.4), lhs_contracting_dims={1}, rhs_contracting_dims={0}
+    ROOT convert.16 = f32[1024,1024] convert(dot.15)
+  })";
+
+  EXPECT_TRUE(RunAndCompare(matmul_module_str, ErrorSpec{1e-2, 1e-2}));
+  MatchOptimizedHlo(matmul_module_str, fused_matmul_bias_elu_rewrite_str_);
+}
+
 TEST_F(MatmulTest, SimpleTestBF16Gemv1) {
   if (!IsSupportedType(PrimitiveType::BF16)) {
     GTEST_SKIP() << "CPU does not support BF16.";
@@ -1010,6 +1149,142 @@ TEST_F(MatmulTest, SimpleTestF32WithMulAndAddFusion) {
     )");
 }
 
+TEST_F(MatmulTest, BiasAddTanhFusionTest_F32) {
+  const char* matmul_module_str = R"(
+  HloModule matmul.bias.tanh.test.f32
+  ENTRY matmul.bias.tanh.test.f32 {
+    arg.0 = f32[32,32,40,30] parameter(0)
+    arg.1 = f32[32,32,30,40] parameter(1)
+    dot.2 = f32[32,32,40,40] dot(arg.0, arg.1), lhs_batch_dims={0,1}, lhs_contracting_dims={3}, rhs_batch_dims={0,1}, rhs_contracting_dims={2}
+    const.3 = f32[40] constant(15)
+    bcast.4 = f32[32,32,40,40] broadcast(const.3), dimensions={3}
+    add.5 = f32[32,32,40,40] add(dot.2, bcast.4)
+    tanh.6 = f32[32,32,40,40] tanh(add.5)
+    tuple.7 = (f32[32,32,40,40]) tuple(tanh.6)
+    ROOT get-tuple-element.8 = f32[32,32,40,40] get-tuple-element(tuple.7), index=0
+  })";
+
+  EXPECT_TRUE(RunAndCompare(matmul_module_str, ErrorSpec{1e-4, 1e-4}));
+  MatchOptimizedHlo(matmul_module_str, fused_matmul_bias_tanh_rewrite_str_);
+}
+
+TEST_F(MatmulTest, BiasAddTanhFusionTest_BF16) {
+  if (!IsSupportedType(PrimitiveType::BF16)) {
+    GTEST_SKIP() << "CPU does not support BF16.";
+  }
+  const char* matmul_module_str = R"(
+  HloModule matmul.bias.tanh.test.f32
+  ENTRY matmul.bias.tanh.test.f32 {
+    arg0.1 = f32[1024,512] parameter(0)
+    convert.2 = bf16[1024,512] convert(arg0.1)
+    arg1.3 = f32[256,512] parameter(1)
+    convert.4 = bf16[256,512] convert(arg1.3)
+    dot.5 = bf16[1024,256] dot(convert.2, convert.4), lhs_contracting_dims={1}, rhs_contracting_dims={1}
+    convert.6 = f32[1024,256] convert(dot.5)
+    arg2.7 = f32[256] parameter(2)
+    broadcast.8 = f32[1024,256] broadcast(arg2.7), dimensions={1}
+    add.9 = f32[1024,256] add(convert.6, broadcast.8)
+    ROOT tanh.10 = f32[1024,256] tanh(add.9)
+  })";
+  EXPECT_TRUE(RunAndCompare(matmul_module_str, ErrorSpec{1e-2, 1e-2}));
+  MatchOptimizedHlo(matmul_module_str, fused_matmul_bias_tanh_rewrite_str_);
+}
+
+TEST_F(MatmulTest, BiasAddTanhFusionTest_F16) {
+  if (!IsSupportedType(PrimitiveType::F16)) {
+    GTEST_SKIP() << "CPU does not support F16.";
+  }
+  const char* matmul_module_str = R"(
+  HloModule matmul.bias.tanh.test.f16
+  ENTRY matmul.bias.tanh.test.f16 {
+    arg0.1 = f16[1024,1024] parameter(0)
+    arg1.2 = f16[1024,1024] parameter(1)
+    dot.3 = f16[1024,1024] dot(arg1.2, arg0.1), lhs_contracting_dims={1}, rhs_contracting_dims={0}
+    arg2.4 = f16[1024] parameter(2)
+    broadcast.5 = f16[1024,1024] broadcast(arg2.4), dimensions={1}
+    add.6 = f16[1024,1024] add(dot.3, broadcast.5)
+    ROOT tanh.7 = f16[1024,1024] tanh(add.6)
+  })";
+
+  EXPECT_TRUE(RunAndCompare(matmul_module_str, ErrorSpec{1e-4, 1e-4}));
+  MatchOptimizedHlo(matmul_module_str, fused_matmul_bias_tanh_rewrite_str_);
+}
+
+// Test MM + BiasAdd + Relu6 fusion : F32
+TEST_F(MatmulTest, BiasAddRelu6Fusion_F32) {
+  const char* matmul_module_str = R"(
+  HloModule matmul.bias.relu6.test.f32
+  ENTRY matmul.bias.relu6.test.f32 {
+    constant.1 = f32[] constant(0)
+    broadcast.2 = f32[1024,1024] broadcast(constant.1), dimensions={}
+    arg1.3 = f32[1024,1024] parameter(1)
+    arg2.4 = f32[1024,1024] parameter(0)
+    dot.5 = f32[1024,1024] dot(arg1.3, arg2.4), lhs_contracting_dims={1}, rhs_contracting_dims={0}
+    arg3.6 = f32[1024] parameter(2)
+    broadcast.7 = f32[1024,1024] broadcast(arg3.6), dimensions={1}
+    add.8 = f32[1024,1024] add(dot.5, broadcast.7)
+    constant.9 = f32[] constant(6)
+    broadcast.10 = f32[1024,1024] broadcast(constant.9), dimensions={}
+    ROOT clamp.11 = f32[1024,1024] clamp(broadcast.2, add.8, broadcast.10)
+  })";
+
+  EXPECT_TRUE(RunAndCompare(matmul_module_str, ErrorSpec{1e-2, 1e-2}));
+  MatchOptimizedHlo(matmul_module_str, fused_matmul_bias_relu6_rewrite_str_);
+}
+
+// Test MM + BiasAdd + Relu6 fusion : BF16
+TEST_F(MatmulTest, BiasAddRelu6Fusion_BF16) {
+  if (!IsSupportedType(PrimitiveType::BF16)) {
+    GTEST_SKIP() << "CPU does not support BF16.";
+  }
+  const char* matmul_module_str = R"(
+  HloModule matmul.bias.relu6.test.bf16
+  ENTRY matmul.bias.relu6.test.bf16 {
+    constant.1 = f32[] constant(0)
+    broadcast.2 = f32[1024,256] broadcast(constant.1), dimensions={}
+    arg0.3 = f32[1024,512] parameter(0)
+    convert.4 = bf16[1024,512] convert(arg0.3)
+    arg1.5 = f32[256,512] parameter(1)
+    convert.6 = bf16[256,512] convert(arg1.5)
+    dot.7 = bf16[1024,256] dot(convert.4, convert.6), lhs_contracting_dims={1}, rhs_contracting_dims={1}
+    convert.8 = f32[1024,256] convert(dot.7)
+    arg2.9 = f32[256] parameter(2)
+    broadcast.10 = f32[1024,256] broadcast(arg2.9), dimensions={1}
+    add.11 = f32[1024,256] add(convert.8, broadcast.10)
+    constant.12 = f32[] constant(6)
+    broadcast.13 = f32[1024,256] broadcast(constant.12), dimensions={}
+    ROOT clamp.14 = f32[1024,256] clamp(broadcast.2, add.11, broadcast.13)
+  })";
+
+  EXPECT_TRUE(RunAndCompare(matmul_module_str, ErrorSpec{1e-2, 1e-2}));
+  MatchOptimizedHlo(matmul_module_str, fused_matmul_bias_relu6_rewrite_str_);
+}
+
+// Test MM + BiasAdd + Relu6 fusion : F16
+TEST_F(MatmulTest, BiasAddRelu6Fusion_F16) {
+  if (!IsSupportedType(PrimitiveType::F16)) {
+    GTEST_SKIP() << "CPU does not support F16.";
+  }
+  const char* matmul_module_str = R"(
+  HloModule matmul.bias.relu6.test.f16
+  ENTRY matmul.bias.relu6.test.f16 {
+    constant.1 = f16[] constant(0)
+    broadcast.2 = f16[1024,1024] broadcast(constant.7), dimensions={}
+    arg0.3 = f16[1024,1024] parameter(0)
+    arg1.4 = f16[1024,1024] parameter(1)
+    dot.5 = f16[1024,1024] dot(arg1.4, arg0.3), lhs_contracting_dims={1}, rhs_contracting_dims={0}
+    arg2.6 = f16[1024] parameter(2)
+    broadcast.7 = f16[1024,1024] broadcast(arg2.6), dimensions={1}
+    add.8 = f16[1024,1024] add(dot.5, broadcast.7)
+    constant.9 = f16[] constant(6)
+    broadcast.10 = f16[1024,1024] broadcast(constant.9), dimensions={}
+    ROOT clamp.11 = f16[1024,1024] clamp(broadcast.2, add.8, broadcast.10)
+  })";
+
+  EXPECT_TRUE(RunAndCompare(matmul_module_str, ErrorSpec{1e-4, 1e-4}));
+  MatchOptimizedHlo(matmul_module_str, fused_matmul_bias_relu6_rewrite_str_);
+}
+
 TEST_F(MatmulTest, SimpleTestBF16WithMulAndAddFusion) {
   if (!IsSupportedType(PrimitiveType::BF16)) {
     GTEST_SKIP() << "CPU does not support BF16.";

From 99ddf8b8239d4955de73917d5cfa179e648628a2 Mon Sep 17 00:00:00 2001
From: Mohammed Anany <manany@google.com>
Date: Wed, 15 May 2024 04:31:17 -0700
Subject: [PATCH 173/478] [Triton] Bring back a test to show that TF32 is not
 supported for 8-bit or less types with F32.

PiperOrigin-RevId: 633894011
---
 .../xla/service/gpu/ir_emitter_triton_test.cc | 40 +++++++++++++++++++
 1 file changed, 40 insertions(+)

diff --git a/third_party/xla/xla/service/gpu/ir_emitter_triton_test.cc b/third_party/xla/xla/service/gpu/ir_emitter_triton_test.cc
index 972b5cb57c8b6c..a9a0ed4ba69e9d 100644
--- a/third_party/xla/xla/service/gpu/ir_emitter_triton_test.cc
+++ b/third_party/xla/xla/service/gpu/ir_emitter_triton_test.cc
@@ -5204,6 +5204,46 @@ ENTRY entry {
               "Triton support is only enabled for Ampere GPUs and up.")));
 }
 
+// This test could be modified to allow TF32 once this bug is fixed.
+// TODO(b/320659359) Allow TF32 for 8-bit or less types with F32.
+TEST_F(TritonFilecheckTest, NoTF32For8BitOrLessWithF32) {
+  const std::string hlo_text = R"(
+HloModule t
+
+triton_dot {
+  parameter_0 = s32[11,24]{1,0} parameter(0)
+  broadcast.1747 = s32[11,24,128]{2,1,0} broadcast(parameter_0),
+  dimensions={0,1} parameter_1 = s32[11,24,128]{2,1,0} parameter(1)
+  compare.49 = pred[11,24,128]{2,1,0} compare(broadcast.1747, parameter_1),
+      direction=EQ bitcast.4717 = pred[264,128]{1,0} bitcast(compare.49)
+  convert.142 = f32[264,128]{1,0} convert(bitcast.4717)
+  parameter_2 = f32[128,8]{1,0} parameter(2)
+  ROOT dot.381 = f32[264,8]{1,0} dot(convert.142, parameter_2),
+      lhs_contracting_dims={1}, rhs_contracting_dims={0}
+}
+
+ENTRY e {
+  p0 = s32[11,24]{1,0} parameter(0)
+  p1 = s32[11,24,128]{2,1,0} parameter(1)
+  p2 = f32[128,8]{1,0} parameter(2)
+  ROOT _ = f32[264,8] fusion(p0, p1, p2), kind=kCustom, calls=triton_dot,
+    backend_config={"fusion_backend_config": {kind: "__triton_gemm",
+      triton_gemm_config:
+        {"block_m":32,"block_n":16,"block_k":128,
+         "split_k":1,"num_stages":1,"num_warps":4,
+         "num_ctas":1}}}
+})";
+
+  TritonGemmConfig config(32, 16, 128, 1, 1, 4);
+  ASSERT_OK(
+      CreateTritonIrAndFileCheck(hlo_text, config, EmitMatMul, "triton_dot", R"(
+CHECK:      tt.dot
+CHECK-NOT:  inputPrecision = tf32
+  )"));
+
+  EXPECT_TRUE(RunAndCompare(hlo_text, ErrorSpec{/*aabs=*/1e-3, /*arel=*/1e-3}));
+}
+
 }  // namespace
 }  // namespace gpu
 }  // namespace xla

From d1d3cc601529cf4a4a580c42eacf951adfd3ba8b Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Wed, 15 May 2024 05:14:14 -0700
Subject: [PATCH 174/478] Automated Code Change

PiperOrigin-RevId: 633909924
---
 .../lite/tools/evaluation/stages/utils/image_metrics_test.cc     | 1 -
 1 file changed, 1 deletion(-)

diff --git a/tensorflow/lite/tools/evaluation/stages/utils/image_metrics_test.cc b/tensorflow/lite/tools/evaluation/stages/utils/image_metrics_test.cc
index 11d71d663fff22..5e95165081781c 100644
--- a/tensorflow/lite/tools/evaluation/stages/utils/image_metrics_test.cc
+++ b/tensorflow/lite/tools/evaluation/stages/utils/image_metrics_test.cc
@@ -20,7 +20,6 @@ limitations under the License.
 #include <cmath>
 #include <cstdlib>
 
-#include <gmock/gmock.h>
 #include <gtest/gtest.h>
 
 namespace tflite {

From 6895571f33ff1dc7275ca6abe211f559766b1bb7 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Wed, 15 May 2024 05:17:36 -0700
Subject: [PATCH 175/478] Update ops-related pbtxt files.

PiperOrigin-RevId: 633910995
---
 tensorflow/core/ops/ops.pbtxt | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorflow/core/ops/ops.pbtxt b/tensorflow/core/ops/ops.pbtxt
index 2073f2e00ecd81..0751aea36d13e7 100644
--- a/tensorflow/core/ops/ops.pbtxt
+++ b/tensorflow/core/ops/ops.pbtxt
@@ -1,4 +1,4 @@
-go/debugonly    
+go/debugonly   
 op {
   name: "Abort"
   attr {

From 3c7d9c657e9f0e122cda70b178f56ef29aedf71d Mon Sep 17 00:00:00 2001
From: Johannes Reifferscheid <jreiffers@google.com>
Date: Wed, 15 May 2024 05:57:51 -0700
Subject: [PATCH 176/478] Revive MLIR interpreter.

This is needed for debugging MLIR emitters. So far, the new GPU ops aren't
supported. This just reverts the deletion, with a few changes:

- removal of depreacted dialects (gml_st, thlo, lhlo, deallocation, xla CPU, xla rt)
- adjustments for upstream changes (expand_shape, vector combiners, scf.parallel)
- removal of mhlo ComputeReshapeShape
- replacement of tsl::Status with absl::Status
- move it out of mlir_hlo and adjust to TF style guide (...mostly)

The tooling still needs to be hooked up to the GPU compiler.

Reverts changelist 568496965

PiperOrigin-RevId: 633922803
---
 .../xla/xla/mlir/tools/mlir_bisect/BUILD      |  66 ++
 .../xla/xla/mlir/tools/mlir_bisect/README.md  |  85 ++
 .../xla/mlir/tools/mlir_bisect/bisect_lib.cc  |  87 ++
 .../xla/mlir/tools/mlir_bisect/bisect_lib.h   |  96 ++
 .../xla/mlir/tools/mlir_bisect/mlir_bisect.cc | 360 +++++++
 .../xla/mlir/tools/mlir_bisect/rewrites/BUILD |  26 +
 .../mlir/tools/mlir_bisect/rewrites/func.cc   |  86 ++
 .../tools/mlir_bisect/rewrites/general.cc     | 194 ++++
 .../mlir/tools/mlir_bisect/rewrites/scf.cc    | 144 +++
 .../tools/mlir_bisect/rewrites/tests/BUILD    |  13 +
 .../tests/erase-op-without-results.mlir       |  12 +
 .../rewrites/tests/inline-scf-while.mlir      |  40 +
 .../tests/reduce-scf-forall-bounds.mlir       |  16 +
 .../tests/replace-op-with-constant.mlir       |  26 +
 .../rewrites/tests/replace-op-with-value.mlir |  16 +
 .../tests/replace-operand-with-constant.mlir  |  28 +
 ...eturn-operands-of-terminator-operands.mlir |  15 +
 .../rewrites/tests/truncate-function.mlir     |  31 +
 .../xla/mlir/tools/mlir_bisect/test_passes.cc |  48 +
 .../xla/mlir/tools/mlir_bisect/test_passes.h  |  29 +
 .../xla/mlir/tools/mlir_bisect/tests/BUILD    |  17 +
 .../mlir/tools/mlir_bisect/tests/bisect.mlir  |  46 +
 .../mlir/tools/mlir_bisect/tests/no-bug.mlir  |  10 +
 .../tools/mlir_bisect/tests/snapshot.mlir     |  12 +
 .../tools/mlir_bisect/tests/snapshot.mlir.pb  | Bin 0 -> 68 bytes
 .../xla/xla/mlir/tools/mlir_interpreter/BUILD |  40 +
 .../tools/mlir_interpreter/dialects/BUILD     |  67 ++
 .../tools/mlir_interpreter/dialects/affine.cc |  53 ++
 .../tools/mlir_interpreter/dialects/arith.cc  | 311 +++++++
 .../dialects/bufferization.cc                 |  70 ++
 .../mlir_interpreter/dialects/builtin.cc      |  64 ++
 .../mlir_interpreter/dialects/comparators.h   | 105 +++
 .../mlir_interpreter/dialects/complex.cc      |  68 ++
 .../mlir_interpreter/dialects/cwise_math.h    | 242 +++++
 .../tools/mlir_interpreter/dialects/func.cc   | 122 +++
 .../tools/mlir_interpreter/dialects/linalg.cc | 338 +++++++
 .../tools/mlir_interpreter/dialects/math.cc   |  64 ++
 .../tools/mlir_interpreter/dialects/memref.cc | 253 +++++
 .../tools/mlir_interpreter/dialects/mhlo.cc   | 871 +++++++++++++++++
 .../dialects/mhlo_binary_cwise.cc             |  45 +
 .../dialects/mhlo_unary_cwise.cc              |  84 ++
 .../tools/mlir_interpreter/dialects/scf.cc    | 222 +++++
 .../tools/mlir_interpreter/dialects/tensor.cc | 262 ++++++
 .../mlir_interpreter/dialects/tests/BUILD     |  13 +
 .../dialects/tests/affine/apply.mlir          |  63 ++
 .../dialects/tests/affine/minmax.mlir         |  36 +
 .../dialects/tests/arith/bitcast.mlir         |  21 +
 .../dialects/tests/arith/cmpf.mlir            | 129 +++
 .../dialects/tests/arith/cmpi.mlir            | 147 +++
 .../dialects/tests/arith/constant.mlir        |  37 +
 .../dialects/tests/arith/extf.mlir            |  11 +
 .../dialects/tests/arith/fptosi.mlir          |  21 +
 .../dialects/tests/arith/index_cast.mlir      |  28 +
 .../dialects/tests/arith/int_math.mlir        | 111 +++
 .../dialects/tests/arith/minmax.mlir          |  25 +
 .../dialects/tests/arith/negf.mlir            |  21 +
 .../dialects/tests/arith/remf.mlir            |  12 +
 .../dialects/tests/arith/select.mlir          |  52 ++
 .../dialects/tests/arith/sitofp.mlir          |  31 +
 .../dialects/tests/arith/uitofp.mlir          |  31 +
 .../dialects/tests/arith/vector_math.mlir     |  12 +
 .../tests/bufferization/alloc_tensor.mlir     |  30 +
 .../dialects/tests/bufferization/clone.mlir   |  14 +
 .../tests/bufferization/to_memref.mlir        |  10 +
 .../tests/bufferization/to_tensor.mlir        |  11 +
 .../builtin/unrealized_conversion_cast.mlir   |  21 +
 .../dialects/tests/complex/complex.mlir       | 186 ++++
 .../dialects/tests/func/call.mlir             |  48 +
 .../dialects/tests/linalg/broadcast.mlir      |  30 +
 .../dialects/tests/linalg/dot.mlir            |  14 +
 .../dialects/tests/linalg/fill.mlir           |  24 +
 .../dialects/tests/linalg/generic.mlir        | 113 +++
 .../dialects/tests/linalg/map.mlir            |  74 ++
 .../dialects/tests/linalg/matmul.mlir         |  41 +
 .../dialects/tests/linalg/reduce.mlir         |  57 ++
 .../dialects/tests/linalg/transpose.mlir      |  27 +
 .../dialects/tests/linalg/vecmat.mlir         |  14 +
 .../dialects/tests/math/math.mlir             | 252 +++++
 .../dialects/tests/memref/alloc.mlir          |  57 ++
 .../dialects/tests/memref/collapse_shape.mlir |  33 +
 .../dialects/tests/memref/copy.mlir           |  39 +
 .../dialects/tests/memref/dim.mlir            |  12 +
 .../dialects/tests/memref/expand_shape.mlir   |  52 ++
 .../dialects/tests/memref/get_global.mlir     |  12 +
 .../dialects/tests/memref/invalid.mlir        |  77 ++
 .../dialects/tests/memref/load.mlir           |  12 +
 .../dialects/tests/memref/subview.mlir        | 120 +++
 .../dialects/tests/mhlo/bitcast_convert.mlir  |  11 +
 .../dialects/tests/mhlo/broadcast_in_dim.mlir |  20 +
 .../dialects/tests/mhlo/case.mlir             |  17 +
 .../dialects/tests/mhlo/clamp.mlir            |  27 +
 .../dialects/tests/mhlo/compare.mlir          | 143 +++
 .../dialects/tests/mhlo/complex_math.mlir     | 100 ++
 .../dialects/tests/mhlo/concatenate.mlir      |  37 +
 .../dialects/tests/mhlo/constant.mlir         |  25 +
 .../dialects/tests/mhlo/convert.mlir          |  21 +
 .../dialects/tests/mhlo/dot.mlir              |  37 +
 .../dialects/tests/mhlo/dot_general.mlir      |  73 ++
 .../dialects/tests/mhlo/dynamic_slice.mlir    |  32 +
 .../tests/mhlo/dynamic_update_slice.mlir      |  34 +
 .../dialects/tests/mhlo/float_math.mlir       | 199 ++++
 .../dialects/tests/mhlo/gather.mlir           |  58 ++
 .../dialects/tests/mhlo/int_math.mlir         | 358 +++++++
 .../dialects/tests/mhlo/iota.mlir             |  30 +
 .../dialects/tests/mhlo/pad.mlir              |  56 ++
 .../dialects/tests/mhlo/reduce.mlir           |  17 +
 .../dialects/tests/mhlo/reshape.mlir          |  34 +
 .../dialects/tests/mhlo/scatter.mlir          |  55 ++
 .../dialects/tests/mhlo/select.mlir           |  14 +
 .../dialects/tests/mhlo/slice.mlir            |  16 +
 .../dialects/tests/mhlo/sort.mlir             |  25 +
 .../dialects/tests/mhlo/subtract.mlir         |  10 +
 .../dialects/tests/mhlo/transpose.mlir        |  28 +
 .../dialects/tests/mhlo/tuple.mlir            |  30 +
 .../dialects/tests/mhlo/while.mlir            |  25 +
 .../dialects/tests/scf/for.mlir               |  82 ++
 .../dialects/tests/scf/forall.mlir            |  62 ++
 .../dialects/tests/scf/if.mlir                |  69 ++
 .../dialects/tests/scf/parallel.mlir          |  42 +
 .../dialects/tests/scf/while.mlir             |  45 +
 .../dialects/tests/tensor/collapse_shape.mlir |  42 +
 .../dialects/tests/tensor/dim.mlir            |  12 +
 .../dialects/tests/tensor/empty.mlir          |  21 +
 .../dialects/tests/tensor/expand_shape.mlir   |  31 +
 .../dialects/tests/tensor/extract.mlir        |  13 +
 .../dialects/tests/tensor/extract_slice.mlir  |  62 ++
 .../dialects/tests/tensor/from_elements.mlir  |  25 +
 .../dialects/tests/tensor/generate.mlir       |  29 +
 .../dialects/tests/tensor/insert.mlir         |  14 +
 .../dialects/tests/tensor/insert_slice.mlir   |  25 +
 .../dialects/tests/tensor/pad.mlir            |  38 +
 .../dialects/tests/vector/bitcast.mlir        |  32 +
 .../dialects/tests/vector/broadcast.mlir      |  51 +
 .../dialects/tests/vector/compressstore.mlir  |  16 +
 .../dialects/tests/vector/constant_mask.mlir  |  14 +
 .../dialects/tests/vector/contract.mlir       | 141 +++
 .../dialects/tests/vector/create_mask.mlir    |  16 +
 .../dialects/tests/vector/expandload.mlir     |  19 +
 .../dialects/tests/vector/extract.mlir        |  52 ++
 .../tests/vector/extract_strided_slice.mlir   |  18 +
 .../dialects/tests/vector/extractelement.mlir |  22 +
 .../dialects/tests/vector/flat_transpose.mlir |  23 +
 .../dialects/tests/vector/fma.mlir            |  13 +
 .../dialects/tests/vector/gather.mlir         |  50 +
 .../dialects/tests/vector/insert.mlir         |  57 ++
 .../tests/vector/insert_strided_slice.mlir    |  17 +
 .../dialects/tests/vector/insertelement.mlir  |  24 +
 .../dialects/tests/vector/invalid.mlir        |  39 +
 .../dialects/tests/vector/load.mlir           |  27 +
 .../dialects/tests/vector/maskedload.mlir     |  19 +
 .../dialects/tests/vector/maskedstore.mlir    |  18 +
 .../tests/vector/multi_reduction.mlir         |  46 +
 .../dialects/tests/vector/outerproduct.mlir   | 155 ++++
 .../dialects/tests/vector/reduction.mlir      | 235 +++++
 .../dialects/tests/vector/shape_cast.mlir     |  23 +
 .../dialects/tests/vector/shuffle.mlir        |  34 +
 .../dialects/tests/vector/splat.mlir          |  11 +
 .../dialects/tests/vector/store.mlir          |  39 +
 .../dialects/tests/vector/transfer_read.mlir  | 118 +++
 .../dialects/tests/vector/transfer_write.mlir |  91 ++
 .../dialects/tests/vector/transpose.mlir      |  28 +
 .../dialects/tests/vector/type_cast.mlir      |  11 +
 .../dialects/tests/vector/vscale.mlir         |  12 +
 .../tools/mlir_interpreter/dialects/util.cc   | 180 ++++
 .../tools/mlir_interpreter/dialects/util.h    |  85 ++
 .../tools/mlir_interpreter/dialects/vector.cc | 875 ++++++++++++++++++
 .../tools/mlir_interpreter/framework/BUILD    |  42 +
 .../mlir_interpreter/framework/interpreter.cc | 150 +++
 .../mlir_interpreter/framework/interpreter.h  | 186 ++++
 .../framework/interpreter_value.cc            | 382 ++++++++
 .../framework/interpreter_value.h             | 227 +++++
 .../framework/interpreter_value_util.h        | 190 ++++
 .../framework/registration.cc                 | 122 +++
 .../mlir_interpreter/framework/registration.h | 229 +++++
 .../framework/tensor_or_memref.cc             | 162 ++++
 .../framework/tensor_or_memref.h              | 364 ++++++++
 .../mlir_interpreter/framework/tests/BUILD    |  25 +
 .../framework/tests/interpreter_value_test.cc | 238 +++++
 .../framework/tests/tensor_or_memref_test.cc  | 104 +++
 .../mlir_interpreter_runner.cc                | 138 +++
 .../xla/xla/mlir/tools/mlir_replay/BUILD      |  63 ++
 .../xla/xla/mlir/tools/mlir_replay/README.md  |  48 +
 .../xla/mlir/tools/mlir_replay/mlir_replay.cc | 236 +++++
 .../mlir/tools/mlir_replay/mlir_replay_lib.cc | 254 +++++
 .../mlir/tools/mlir_replay/mlir_replay_lib.h  |  40 +
 .../xla/mlir/tools/mlir_replay/public/BUILD   |  56 ++
 .../mlir_replay/public/compiler_trace.proto   |  31 +
 .../mlir_replay/public/execution_trace.proto  |  72 ++
 .../public/execution_trace_utils.cc           | 449 +++++++++
 .../public/execution_trace_utils.h            |  76 ++
 .../public/execution_trace_utils_test.cc      | 139 +++
 191 files changed, 15743 insertions(+)
 create mode 100644 third_party/xla/xla/mlir/tools/mlir_bisect/BUILD
 create mode 100644 third_party/xla/xla/mlir/tools/mlir_bisect/README.md
 create mode 100644 third_party/xla/xla/mlir/tools/mlir_bisect/bisect_lib.cc
 create mode 100644 third_party/xla/xla/mlir/tools/mlir_bisect/bisect_lib.h
 create mode 100644 third_party/xla/xla/mlir/tools/mlir_bisect/mlir_bisect.cc
 create mode 100644 third_party/xla/xla/mlir/tools/mlir_bisect/rewrites/BUILD
 create mode 100644 third_party/xla/xla/mlir/tools/mlir_bisect/rewrites/func.cc
 create mode 100644 third_party/xla/xla/mlir/tools/mlir_bisect/rewrites/general.cc
 create mode 100644 third_party/xla/xla/mlir/tools/mlir_bisect/rewrites/scf.cc
 create mode 100644 third_party/xla/xla/mlir/tools/mlir_bisect/rewrites/tests/BUILD
 create mode 100644 third_party/xla/xla/mlir/tools/mlir_bisect/rewrites/tests/erase-op-without-results.mlir
 create mode 100644 third_party/xla/xla/mlir/tools/mlir_bisect/rewrites/tests/inline-scf-while.mlir
 create mode 100644 third_party/xla/xla/mlir/tools/mlir_bisect/rewrites/tests/reduce-scf-forall-bounds.mlir
 create mode 100644 third_party/xla/xla/mlir/tools/mlir_bisect/rewrites/tests/replace-op-with-constant.mlir
 create mode 100644 third_party/xla/xla/mlir/tools/mlir_bisect/rewrites/tests/replace-op-with-value.mlir
 create mode 100644 third_party/xla/xla/mlir/tools/mlir_bisect/rewrites/tests/replace-operand-with-constant.mlir
 create mode 100644 third_party/xla/xla/mlir/tools/mlir_bisect/rewrites/tests/return-operands-of-terminator-operands.mlir
 create mode 100644 third_party/xla/xla/mlir/tools/mlir_bisect/rewrites/tests/truncate-function.mlir
 create mode 100644 third_party/xla/xla/mlir/tools/mlir_bisect/test_passes.cc
 create mode 100644 third_party/xla/xla/mlir/tools/mlir_bisect/test_passes.h
 create mode 100644 third_party/xla/xla/mlir/tools/mlir_bisect/tests/BUILD
 create mode 100644 third_party/xla/xla/mlir/tools/mlir_bisect/tests/bisect.mlir
 create mode 100644 third_party/xla/xla/mlir/tools/mlir_bisect/tests/no-bug.mlir
 create mode 100644 third_party/xla/xla/mlir/tools/mlir_bisect/tests/snapshot.mlir
 create mode 100644 third_party/xla/xla/mlir/tools/mlir_bisect/tests/snapshot.mlir.pb
 create mode 100644 third_party/xla/xla/mlir/tools/mlir_interpreter/BUILD
 create mode 100644 third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/BUILD
 create mode 100644 third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/affine.cc
 create mode 100644 third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/arith.cc
 create mode 100644 third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/bufferization.cc
 create mode 100644 third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/builtin.cc
 create mode 100644 third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/comparators.h
 create mode 100644 third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/complex.cc
 create mode 100644 third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/cwise_math.h
 create mode 100644 third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/func.cc
 create mode 100644 third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/linalg.cc
 create mode 100644 third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/math.cc
 create mode 100644 third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/memref.cc
 create mode 100644 third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/mhlo.cc
 create mode 100644 third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/mhlo_binary_cwise.cc
 create mode 100644 third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/mhlo_unary_cwise.cc
 create mode 100644 third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/scf.cc
 create mode 100644 third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tensor.cc
 create mode 100644 third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/BUILD
 create mode 100644 third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/affine/apply.mlir
 create mode 100644 third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/affine/minmax.mlir
 create mode 100644 third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/arith/bitcast.mlir
 create mode 100644 third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/arith/cmpf.mlir
 create mode 100644 third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/arith/cmpi.mlir
 create mode 100644 third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/arith/constant.mlir
 create mode 100644 third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/arith/extf.mlir
 create mode 100644 third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/arith/fptosi.mlir
 create mode 100644 third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/arith/index_cast.mlir
 create mode 100644 third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/arith/int_math.mlir
 create mode 100644 third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/arith/minmax.mlir
 create mode 100644 third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/arith/negf.mlir
 create mode 100644 third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/arith/remf.mlir
 create mode 100644 third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/arith/select.mlir
 create mode 100644 third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/arith/sitofp.mlir
 create mode 100644 third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/arith/uitofp.mlir
 create mode 100644 third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/arith/vector_math.mlir
 create mode 100644 third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/bufferization/alloc_tensor.mlir
 create mode 100644 third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/bufferization/clone.mlir
 create mode 100644 third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/bufferization/to_memref.mlir
 create mode 100644 third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/bufferization/to_tensor.mlir
 create mode 100644 third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/builtin/unrealized_conversion_cast.mlir
 create mode 100644 third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/complex/complex.mlir
 create mode 100644 third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/func/call.mlir
 create mode 100644 third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/linalg/broadcast.mlir
 create mode 100644 third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/linalg/dot.mlir
 create mode 100644 third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/linalg/fill.mlir
 create mode 100644 third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/linalg/generic.mlir
 create mode 100644 third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/linalg/map.mlir
 create mode 100644 third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/linalg/matmul.mlir
 create mode 100644 third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/linalg/reduce.mlir
 create mode 100644 third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/linalg/transpose.mlir
 create mode 100644 third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/linalg/vecmat.mlir
 create mode 100644 third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/math/math.mlir
 create mode 100644 third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/memref/alloc.mlir
 create mode 100644 third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/memref/collapse_shape.mlir
 create mode 100644 third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/memref/copy.mlir
 create mode 100644 third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/memref/dim.mlir
 create mode 100644 third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/memref/expand_shape.mlir
 create mode 100644 third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/memref/get_global.mlir
 create mode 100644 third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/memref/invalid.mlir
 create mode 100644 third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/memref/load.mlir
 create mode 100644 third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/memref/subview.mlir
 create mode 100644 third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/mhlo/bitcast_convert.mlir
 create mode 100644 third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/mhlo/broadcast_in_dim.mlir
 create mode 100644 third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/mhlo/case.mlir
 create mode 100644 third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/mhlo/clamp.mlir
 create mode 100644 third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/mhlo/compare.mlir
 create mode 100644 third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/mhlo/complex_math.mlir
 create mode 100644 third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/mhlo/concatenate.mlir
 create mode 100644 third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/mhlo/constant.mlir
 create mode 100644 third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/mhlo/convert.mlir
 create mode 100644 third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/mhlo/dot.mlir
 create mode 100644 third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/mhlo/dot_general.mlir
 create mode 100644 third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/mhlo/dynamic_slice.mlir
 create mode 100644 third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/mhlo/dynamic_update_slice.mlir
 create mode 100644 third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/mhlo/float_math.mlir
 create mode 100644 third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/mhlo/gather.mlir
 create mode 100644 third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/mhlo/int_math.mlir
 create mode 100644 third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/mhlo/iota.mlir
 create mode 100644 third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/mhlo/pad.mlir
 create mode 100644 third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/mhlo/reduce.mlir
 create mode 100644 third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/mhlo/reshape.mlir
 create mode 100644 third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/mhlo/scatter.mlir
 create mode 100644 third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/mhlo/select.mlir
 create mode 100644 third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/mhlo/slice.mlir
 create mode 100644 third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/mhlo/sort.mlir
 create mode 100644 third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/mhlo/subtract.mlir
 create mode 100644 third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/mhlo/transpose.mlir
 create mode 100644 third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/mhlo/tuple.mlir
 create mode 100644 third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/mhlo/while.mlir
 create mode 100644 third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/scf/for.mlir
 create mode 100644 third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/scf/forall.mlir
 create mode 100644 third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/scf/if.mlir
 create mode 100644 third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/scf/parallel.mlir
 create mode 100644 third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/scf/while.mlir
 create mode 100644 third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/tensor/collapse_shape.mlir
 create mode 100644 third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/tensor/dim.mlir
 create mode 100644 third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/tensor/empty.mlir
 create mode 100644 third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/tensor/expand_shape.mlir
 create mode 100644 third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/tensor/extract.mlir
 create mode 100644 third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/tensor/extract_slice.mlir
 create mode 100644 third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/tensor/from_elements.mlir
 create mode 100644 third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/tensor/generate.mlir
 create mode 100644 third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/tensor/insert.mlir
 create mode 100644 third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/tensor/insert_slice.mlir
 create mode 100644 third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/tensor/pad.mlir
 create mode 100644 third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/vector/bitcast.mlir
 create mode 100644 third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/vector/broadcast.mlir
 create mode 100644 third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/vector/compressstore.mlir
 create mode 100644 third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/vector/constant_mask.mlir
 create mode 100644 third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/vector/contract.mlir
 create mode 100644 third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/vector/create_mask.mlir
 create mode 100644 third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/vector/expandload.mlir
 create mode 100644 third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/vector/extract.mlir
 create mode 100644 third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/vector/extract_strided_slice.mlir
 create mode 100644 third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/vector/extractelement.mlir
 create mode 100644 third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/vector/flat_transpose.mlir
 create mode 100644 third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/vector/fma.mlir
 create mode 100644 third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/vector/gather.mlir
 create mode 100644 third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/vector/insert.mlir
 create mode 100644 third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/vector/insert_strided_slice.mlir
 create mode 100644 third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/vector/insertelement.mlir
 create mode 100644 third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/vector/invalid.mlir
 create mode 100644 third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/vector/load.mlir
 create mode 100644 third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/vector/maskedload.mlir
 create mode 100644 third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/vector/maskedstore.mlir
 create mode 100644 third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/vector/multi_reduction.mlir
 create mode 100644 third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/vector/outerproduct.mlir
 create mode 100644 third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/vector/reduction.mlir
 create mode 100644 third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/vector/shape_cast.mlir
 create mode 100644 third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/vector/shuffle.mlir
 create mode 100644 third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/vector/splat.mlir
 create mode 100644 third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/vector/store.mlir
 create mode 100644 third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/vector/transfer_read.mlir
 create mode 100644 third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/vector/transfer_write.mlir
 create mode 100644 third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/vector/transpose.mlir
 create mode 100644 third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/vector/type_cast.mlir
 create mode 100644 third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/vector/vscale.mlir
 create mode 100644 third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/util.cc
 create mode 100644 third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/util.h
 create mode 100644 third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/vector.cc
 create mode 100644 third_party/xla/xla/mlir/tools/mlir_interpreter/framework/BUILD
 create mode 100644 third_party/xla/xla/mlir/tools/mlir_interpreter/framework/interpreter.cc
 create mode 100644 third_party/xla/xla/mlir/tools/mlir_interpreter/framework/interpreter.h
 create mode 100644 third_party/xla/xla/mlir/tools/mlir_interpreter/framework/interpreter_value.cc
 create mode 100644 third_party/xla/xla/mlir/tools/mlir_interpreter/framework/interpreter_value.h
 create mode 100644 third_party/xla/xla/mlir/tools/mlir_interpreter/framework/interpreter_value_util.h
 create mode 100644 third_party/xla/xla/mlir/tools/mlir_interpreter/framework/registration.cc
 create mode 100644 third_party/xla/xla/mlir/tools/mlir_interpreter/framework/registration.h
 create mode 100644 third_party/xla/xla/mlir/tools/mlir_interpreter/framework/tensor_or_memref.cc
 create mode 100644 third_party/xla/xla/mlir/tools/mlir_interpreter/framework/tensor_or_memref.h
 create mode 100644 third_party/xla/xla/mlir/tools/mlir_interpreter/framework/tests/BUILD
 create mode 100644 third_party/xla/xla/mlir/tools/mlir_interpreter/framework/tests/interpreter_value_test.cc
 create mode 100644 third_party/xla/xla/mlir/tools/mlir_interpreter/framework/tests/tensor_or_memref_test.cc
 create mode 100644 third_party/xla/xla/mlir/tools/mlir_interpreter/mlir_interpreter_runner.cc
 create mode 100644 third_party/xla/xla/mlir/tools/mlir_replay/BUILD
 create mode 100644 third_party/xla/xla/mlir/tools/mlir_replay/README.md
 create mode 100644 third_party/xla/xla/mlir/tools/mlir_replay/mlir_replay.cc
 create mode 100644 third_party/xla/xla/mlir/tools/mlir_replay/mlir_replay_lib.cc
 create mode 100644 third_party/xla/xla/mlir/tools/mlir_replay/mlir_replay_lib.h
 create mode 100644 third_party/xla/xla/mlir/tools/mlir_replay/public/BUILD
 create mode 100644 third_party/xla/xla/mlir/tools/mlir_replay/public/compiler_trace.proto
 create mode 100644 third_party/xla/xla/mlir/tools/mlir_replay/public/execution_trace.proto
 create mode 100644 third_party/xla/xla/mlir/tools/mlir_replay/public/execution_trace_utils.cc
 create mode 100644 third_party/xla/xla/mlir/tools/mlir_replay/public/execution_trace_utils.h
 create mode 100644 third_party/xla/xla/mlir/tools/mlir_replay/public/execution_trace_utils_test.cc

diff --git a/third_party/xla/xla/mlir/tools/mlir_bisect/BUILD b/third_party/xla/xla/mlir/tools/mlir_bisect/BUILD
new file mode 100644
index 00000000000000..c394e91ad24568
--- /dev/null
+++ b/third_party/xla/xla/mlir/tools/mlir_bisect/BUILD
@@ -0,0 +1,66 @@
+load("@bazel_skylib//rules:build_test.bzl", "build_test")
+load("@local_tsl//tsl/platform:rules_cc.bzl", "cc_library")
+load("//xla:xla.bzl", "xla_cc_binary")
+
+# copybara:uncomment package(default_applicable_licenses = ["//tensorflow:license"])
+
+build_test(
+    name = "mlir-bisect_build_test",
+    targets = [
+        ":mlir-bisect",
+    ],
+)
+
+xla_cc_binary(
+    name = "mlir-bisect",
+    testonly = True,
+    srcs = ["mlir_bisect.cc"],
+    visibility = ["//visibility:public"],
+    deps = [
+        ":bisect_lib",
+        "//xla:literal",
+        "//xla/mlir/tools/mlir_bisect/rewrites",
+        "//xla/mlir/tools/mlir_interpreter/dialects",
+        "//xla/mlir/tools/mlir_interpreter/framework",
+        "//xla/mlir/tools/mlir_replay/public:execution_trace_utils",
+        "//xla/mlir_hlo:hlo_dialect_registration",
+        "//xla/mlir_hlo:mhlo_passes",
+        "//xla/service:hlo_proto_cc",
+        "@com_google_absl//absl/log:check",
+        "@llvm-project//llvm:Support",
+        "@llvm-project//mlir:AffineUtils",
+        "@llvm-project//mlir:AllPassesAndDialects",
+        "@llvm-project//mlir:ArithDialect",
+        "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:ParseUtilities",
+        "@llvm-project//mlir:Parser",
+        "@llvm-project//mlir:Pass",
+        "@llvm-project//mlir:Support",
+        "@llvm-project//mlir:Transforms",
+        "@local_tsl//tsl/platform:env",
+        "@local_tsl//tsl/platform:platform_port",
+    ],
+)
+
+cc_library(
+    name = "bisect_lib",
+    srcs = [
+        "bisect_lib.cc",
+        "test_passes.cc",
+    ],
+    hdrs = [
+        "bisect_lib.h",
+        "test_passes.h",
+    ],
+    visibility = ["//visibility:public"],
+    deps = [
+        "//xla/mlir/tools/mlir_replay/public:execution_trace_proto_cc",
+        "//xla/mlir/tools/mlir_replay/public:execution_trace_proto_cc_impl",
+        "//xla/mlir/tools/mlir_replay/public:execution_trace_utils",
+        "@llvm-project//llvm:Support",
+        "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:LinalgDialect",
+        "@llvm-project//mlir:Pass",
+        "@llvm-project//mlir:Support",
+    ],
+)
diff --git a/third_party/xla/xla/mlir/tools/mlir_bisect/README.md b/third_party/xla/xla/mlir/tools/mlir_bisect/README.md
new file mode 100644
index 00000000000000..570e92b6e53802
--- /dev/null
+++ b/third_party/xla/xla/mlir/tools/mlir_bisect/README.md
@@ -0,0 +1,85 @@
+# MLIR HLO mlir_bisect
+
+This is a test case reduction tool, similar in purpose to `mlir-reduce`, but
+specific to the `mlir-interpreter` infrastructure. In particular, reductions can
+depend on concrete values encountered during execution, and reductions can (and
+usually do) generate multiple candidates.
+
+For example, the `ReplaceOpWithConstant` reduction will attempt to replace each
+op with each of its results. If the op is in a loop, each execution will be a
+candidate for replacement.
+
+## Using this tool
+
+1.  Run a JAX test with snapshots enabled:
+
+    ```
+    bazel test some-jax-test
+      --test_env=XLA_FLAGS="--xla_cpu_use_xla_runtime --xla_dump_to=/tmp/dump
+      --xla_dump_hlo_snapshots" --test_filter=SomeSpecific.Test
+      --test_sharding_strategy=disabled --test_strategy=local
+    ```
+
+1.  Figure out the culprit module and pass (sorry, no automation yet):
+
+    ```
+    bazel run tensorflow/compiler/xla/mlir/tools/mlir_replay:mlir_replay -- \
+      --mlir-compilation-trace=/tmp/dump/module_0000.jit__something.mlir-trace.pb \
+      --hlo-snapshot=/tmp/dump/module_0000.jit__something.snapshot.0.pb \
+      --print-changes-only \
+      --execution-trace-dir=/tmp/execution
+    ```
+
+    You should see a pass after which results change. You'll want to use the
+    .mlir file in `/tmp/execution` corresponding to the pass *before* that with
+    the bisect tool.
+
+    Note: If the failing pass is bufferization, you may have to use an earlier
+    snapshot, e.g. before EmptyTensorToAllocTensor.
+1.  Run bisect:
+
+    ```
+    bazel run tensorflow/compiler/xla/mlir/tools/mlir_bisect:mlir-bisect -- \
+      --hlo-snapshot=/tmp/dump/module_0000.jit_something.snapshot.0.pb \
+      --pass-pipeline="builtin.module(empty-tensor-to-alloc-tensor,one-shot-bufferize{allow-return-allocs bufferize-function-boundaries create-deallocs=0})" \
+      /tmp/execution/0052.ScalarizationPass.mlir
+    ```
+
+## Adding a reduction
+
+To add a reduction, create a function that generates the candidates and register
+it:
+
+```
+SmallVector<OwningOpRef<ModuleOp>>
+FrobulateAndDefenestrate(BisectState&, dialect::SomeOp some_op) {
+  auto [cloned_module_1, cloned_op_1] = CloneModuleFor(some_op);
+  Frobulate(cloned_op_1);
+
+  auto [cloned_module_2, cloned_op_2] = CloneModuleFor(some_op);
+  Defenestrate(cloned_op_2);
+
+  return {cloned_module_1, cloned_module_2};
+}
+
+REGISTER_MLIR_REDUCE_STRATEGY(FrobulateAndDefenestrate);
+```
+
+Then, add a test for the strategy. Make sure your strategy is linked into
+mlir-bisect and has `alwayslink` set.
+
+```
+// RUN: mlir-bisect %s --debug-strategy=FrobulateAndDefenestrate | FileCheck %s
+
+func.func @main() {
+  dialect.some_op()
+}
+
+// CHECK: func @main()
+// CHECK-NEXT: frobulated
+
+// CHECK: func @main()
+// CHECK-NEXT: defenestrated
+```
+
+`--debug-strategy` will print all candidates generated by the given strategy.
diff --git a/third_party/xla/xla/mlir/tools/mlir_bisect/bisect_lib.cc b/third_party/xla/xla/mlir/tools/mlir_bisect/bisect_lib.cc
new file mode 100644
index 00000000000000..c76e226ff7a6fa
--- /dev/null
+++ b/third_party/xla/xla/mlir/tools/mlir_bisect/bisect_lib.cc
@@ -0,0 +1,87 @@
+/* Copyright 2023 The OpenXLA Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+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.
+==============================================================================*/
+
+#include "xla/mlir/tools/mlir_bisect/bisect_lib.h"
+
+#include <cassert>
+#include <functional>
+#include <iterator>
+#include <utility>
+
+#include "llvm/ADT/STLExtras.h"
+#include "llvm/Support/Casting.h"
+#include "llvm/Support/ErrorHandling.h"
+#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
+#include "mlir/IR/Operation.h"  // from @llvm-project
+#include "mlir/IR/OwningOpRef.h"  // from @llvm-project
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+
+namespace mlir {
+namespace bisect {
+
+Operation* FindInClone(Operation* op, ModuleOp clone) {
+  if (llvm::isa<ModuleOp>(op)) {
+    return clone;
+  }
+
+  auto* parent_clone = FindInClone(op->getParentOp(), clone);
+  auto cloned_ops =
+      parent_clone->getRegions()[op->getParentRegion()->getRegionNumber()]
+          .getOps();
+  for (auto [original_op, cloned_op] :
+       llvm::zip(op->getParentRegion()->getOps(), cloned_ops)) {
+    if (&original_op == op) {
+      return &cloned_op;
+    }
+  }
+
+  llvm_unreachable("Op not found in clone.");
+}
+
+std::pair<OwningOpRef<ModuleOp>, Operation*> CloneModuleFor(Operation* op) {
+  auto module = op->getParentOfType<ModuleOp>().clone();
+  return {OwningOpRef<ModuleOp>{module}, FindInClone(op, module)};
+}
+
+namespace detail {
+
+DenseMap<StringRef, std::function<CandidateVector(BisectState&, Operation*)>>&
+GetStrategies() {
+  static auto* strategies =
+      new DenseMap<StringRef,
+                   std::function<CandidateVector(BisectState&, Operation*)>>();
+  return *strategies;
+}
+
+void RegisterReduceStrategy(
+    StringRef name,
+    std::function<CandidateVector(BisectState&, Operation*)> fn) {
+  GetStrategies()[name] = std::move(fn);
+}
+
+CandidateVector GetCandidates(
+    const std::function<CandidateVector(BisectState&, Operation*)>& strategy,
+    BisectState& state, ModuleOp op) {
+  assert(strategy && "GetCandidates was passed a null strategy");
+  CandidateVector result;
+  op.lookupSymbol("main")->walk([&](Operation* subOp) {
+    llvm::move(strategy(state, subOp), std::back_inserter(result));
+  });
+  return result;
+}
+
+}  // namespace detail
+}  // namespace bisect
+}  // namespace mlir
diff --git a/third_party/xla/xla/mlir/tools/mlir_bisect/bisect_lib.h b/third_party/xla/xla/mlir/tools/mlir_bisect/bisect_lib.h
new file mode 100644
index 00000000000000..a4784263313de5
--- /dev/null
+++ b/third_party/xla/xla/mlir/tools/mlir_bisect/bisect_lib.h
@@ -0,0 +1,96 @@
+/* Copyright 2023 The OpenXLA Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+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.
+==============================================================================*/
+
+#ifndef XLA_MLIR_TOOLS_MLIR_BISECT_BISECT_LIB_H_
+#define XLA_MLIR_TOOLS_MLIR_BISECT_BISECT_LIB_H_
+
+#include <functional>
+#include <tuple>
+#include <utility>
+
+#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+#include "xla/mlir/tools/mlir_replay/public/execution_trace.pb.h"
+#include "xla/mlir/tools/mlir_replay/public/execution_trace_utils.h"
+
+#define REGISTER_MLIR_REDUCE_STRATEGY(name)                      \
+  static int name##_init = []() {                                \
+    ::mlir::bisect::detail::RegisterReduceStrategy(#name, name); \
+    return 1;                                                    \
+  }();
+
+namespace mlir {
+namespace bisect {
+
+class BisectState {
+ public:
+  void SetTrace(mlir::interpreter::ExecutionTrace trace) {
+    trace_ = std::move(trace);
+  }
+
+  // Returns all executions of the given op.
+  llvm::SmallVector<const interpreter::InstructionTrace*> GetExecutions(
+      mlir::Operation* op) const {
+    return interpreter::FindOpExecutionsInTrace(trace_, op);
+  }
+
+ private:
+  mlir::interpreter::ExecutionTrace trace_;
+};
+
+std::pair<OwningOpRef<ModuleOp>, Operation*> CloneModuleFor(Operation* op);
+Operation* FindInClone(Operation* op, ModuleOp clone);
+
+template <typename Op>
+std::pair<OwningOpRef<ModuleOp>, Op> CloneModuleFor(Op op) {
+  auto [module, op_clone] = CloneModuleFor(op.getOperation());
+  return {std::move(module), llvm::cast<Op>(op_clone)};
+}
+
+namespace detail {
+
+using CandidateVector = SmallVector<std::function<OwningOpRef<ModuleOp>()>>;
+
+CandidateVector GetCandidates(
+    const std::function<CandidateVector(BisectState&, Operation*)>& strategy,
+    BisectState& state, ModuleOp op);
+
+DenseMap<StringRef, std::function<CandidateVector(BisectState&, Operation*)>>&
+GetStrategies();
+
+// Registers a strategy that applies to all ops.
+void RegisterReduceStrategy(
+    StringRef name,
+    std::function<CandidateVector(BisectState&, Operation*)> fn);
+
+// Registers a strategy that applies to specific ops.
+template <typename Op>
+void RegisterReduceStrategy(StringRef name,
+                            CandidateVector (*fn)(BisectState&, Op)) {
+  RegisterReduceStrategy(
+      name, [fn](BisectState& state, Operation* op) -> CandidateVector {
+        if (auto cast = llvm::dyn_cast<Op>(op)) {
+          return fn(state, cast);
+        }
+        return {};
+      });
+}
+
+}  // namespace detail
+
+}  // namespace bisect
+}  // namespace mlir
+
+#endif  // XLA_MLIR_TOOLS_MLIR_BISECT_BISECT_LIB_H_
diff --git a/third_party/xla/xla/mlir/tools/mlir_bisect/mlir_bisect.cc b/third_party/xla/xla/mlir/tools/mlir_bisect/mlir_bisect.cc
new file mode 100644
index 00000000000000..ca1a699444b452
--- /dev/null
+++ b/third_party/xla/xla/mlir/tools/mlir_bisect/mlir_bisect.cc
@@ -0,0 +1,360 @@
+/* Copyright 2023 The OpenXLA Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+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.
+==============================================================================*/
+
+#include <algorithm>
+#include <cstdint>
+#include <iterator>
+#include <memory>
+#include <optional>
+#include <string>
+#include <utility>
+
+#include "absl/log/check.h"
+#include "llvm/ADT/STLExtras.h"
+#include "llvm/ADT/SmallVector.h"
+#include "llvm/ADT/Twine.h"
+#include "llvm/Support/Casting.h"
+#include "llvm/Support/CommandLine.h"
+#include "llvm/Support/SourceMgr.h"
+#include "llvm/Support/raw_ostream.h"
+#include "mlir/Dialect/Affine/Utils.h"  // from @llvm-project
+#include "mlir/Dialect/Arith/IR/Arith.h"  // from @llvm-project
+#include "mlir/IR/Builders.h"  // from @llvm-project
+#include "mlir/IR/BuiltinAttributeInterfaces.h"  // from @llvm-project
+#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
+#include "mlir/IR/DialectRegistry.h"  // from @llvm-project
+#include "mlir/IR/MLIRContext.h"  // from @llvm-project
+#include "mlir/IR/OwningOpRef.h"  // from @llvm-project
+#include "mlir/IR/SymbolTable.h"  // from @llvm-project
+#include "mlir/IR/Verifier.h"  // from @llvm-project
+#include "mlir/InitAllDialects.h"  // from @llvm-project
+#include "mlir/InitAllPasses.h"  // from @llvm-project
+#include "mlir/Parser/Parser.h"  // from @llvm-project
+#include "mlir/Pass/PassManager.h"  // from @llvm-project
+#include "mlir/Pass/PassRegistry.h"  // from @llvm-project
+#include "mlir/Support/FileUtilities.h"  // from @llvm-project
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+#include "mlir/Support/LogicalResult.h"  // from @llvm-project
+#include "mlir/Tools/ParseUtilities.h"  // from @llvm-project
+#include "mlir/Transforms/Passes.h"  // from @llvm-project
+#include "xla/literal.h"
+#include "xla/mlir/tools/mlir_bisect/bisect_lib.h"
+#include "xla/mlir/tools/mlir_bisect/test_passes.h"
+#include "xla/mlir/tools/mlir_interpreter/framework/interpreter.h"
+#include "xla/mlir/tools/mlir_replay/public/execution_trace_utils.h"
+#include "xla/mlir_hlo/mhlo/IR/register.h"
+#include "xla/mlir_hlo/mhlo/transforms/passes.h"
+#include "xla/service/hlo.pb.h"
+#include "tsl/platform/env.h"
+#include "tsl/platform/init_main.h"
+
+struct Options {
+  llvm::cl::opt<std::string> input_filename{llvm::cl::Positional,
+                                            llvm::cl::desc("<input file>"),
+                                            llvm::cl::init("-")};
+  llvm::cl::opt<std::string> hlo_snapshot{
+      "hlo-snapshot",
+      llvm::cl::desc(
+          "If set, get argument values from the given snapshot. If not set, "
+          "the input function must not have any arguments."),
+      llvm::cl::init("")};
+  llvm::cl::opt<std::string> debug_strategy{
+      "debug-strategy",
+      llvm::cl::desc("If set, print all reductions for the given strategy and "
+                     "exit. For testing."),
+      llvm::cl::init("")};
+  llvm::cl::opt<std::string> expected_error{
+      "expected-error",
+      llvm::cl::desc("If set, expect the given error message after applying "
+                     "the pass instead of a successful execution."),
+      llvm::cl::init("")};
+  llvm::cl::opt<int64_t> max_steps_per_run{
+      "max-steps-per-run",
+      llvm::cl::desc("Maximum number of steps to execute for each attempt."),
+      llvm::cl::init(100000)};
+  mlir::PassPipelineCLParser pass_pipeline{"", "Passes to run"};
+  llvm::cl::opt<bool> canonicalize{
+      "enable-canonicalization",
+      llvm::cl::desc("If set, canonicalize candidates before trying them. Set "
+                     "to false if you're bisecting --canonicalize."),
+      llvm::cl::init(true)};
+};
+
+namespace mlir {
+namespace bisect {
+namespace {
+
+OwningOpRef<ModuleOp> ParseMlirInput(llvm::StringRef inputFilename,
+                                     MLIRContext* context) {
+  std::string error_message;
+  auto file = mlir::openInputFile(inputFilename, &error_message);
+  if (!file) {
+    llvm::errs() << error_message << "\n";
+    return {};
+  }
+
+  auto source_mgr = std::make_shared<llvm::SourceMgr>();
+  source_mgr->AddNewSourceBuffer(std::move(file), SMLoc());
+  return parseSourceFile<ModuleOp>(source_mgr, context);
+}
+
+LogicalResult RunPipeline(ModuleOp module, const Options& options) {
+  if (!options.pass_pipeline.hasAnyOccurrences()) {
+    return mlir::success();
+  }
+
+  auto error_handler = [&](const llvm::Twine& msg) {
+    llvm::errs() << msg << "\n";
+    return failure();
+  };
+  PassManager pm(module.getContext());
+  if (failed(options.pass_pipeline.addToPipeline(pm, error_handler)) ||
+      failed(pm.run(module))) {
+    llvm::errs() << "pipeline failed\n";
+    return failure();
+  }
+  return success();
+}
+
+LogicalResult Run(mlir::Operation* module, interpreter::ExecutionTrace* trace,
+                  const Options& options) {
+  SymbolTable symbol_table{module};
+  interpreter::ExecutionTraceListener tracer(trace);
+  interpreter::InterpreterOptions interpreter_options;
+  interpreter_options.listener = &tracer;
+  interpreter_options.max_steps = options.max_steps_per_run;
+  auto results_before_pass = interpreter::RunInterpreter(
+      symbol_table, llvm::cast<func::FuncOp>(symbol_table.lookup("main")), {},
+      interpreter_options);
+
+  if (!results_before_pass.ok()) {
+    llvm::errs() << "Interpreter failed\n";
+    return failure();
+  }
+
+  if (!options.debug_strategy.empty()) {
+    return success();
+  }
+
+  OwningOpRef<ModuleOp> clone(llvm::cast<ModuleOp>(module).clone());
+  if (!succeeded(RunPipeline(*clone, options))) {
+    return failure();
+  }
+
+  SymbolTable symbol_table_after{*clone};
+  interpreter_options.listener = nullptr;
+  bool found_expected_error = false;
+  if (!options.expected_error.empty()) {
+    auto original_handler = std::move(interpreter_options.error_handler);
+    interpreter_options.error_handler = [&](llvm::StringRef failure) {
+      found_expected_error |=
+          failure.find(options.expected_error) != std::string::npos;
+      original_handler(failure);
+    };
+  }
+
+  auto results_after_pass = interpreter::RunInterpreter(
+      symbol_table_after,
+      llvm::cast<func::FuncOp>(symbol_table_after.lookup("main")), {},
+      std::move(interpreter_options));
+
+  if (!results_after_pass.ok()) {
+    if (found_expected_error) {
+      return success();
+    }
+    llvm::errs() << "Interpreter failed\n";
+    return failure();
+  } else if (!options.expected_error.empty()) {
+    llvm::errs() << "Expected error not seen\n";
+    return failure();
+  }
+
+  // If the results are the same, the bug is no longer present.
+  if (*results_before_pass == *results_after_pass) {
+    return failure();
+  }
+
+  llvm::errs() << "results before:\n";
+  for (auto& result : *results_before_pass) {
+    llvm::errs() << "  " << result.ToString() << "\n";
+  }
+  llvm::errs() << "\nresults after:\n";
+  for (auto& result : *results_after_pass) {
+    llvm::errs() << "  " << result.ToString() << "\n";
+  }
+
+  return success();
+}
+
+LogicalResult Canonicalize(ModuleOp module) {
+  PassManager pm(module.getContext());
+  pm.addPass(createCanonicalizerPass());
+  return pm.run(module.getOperation());
+}
+
+OwningOpRef<ModuleOp> ReduceModule(OwningOpRef<ModuleOp> module,
+                                   BisectState& state, const Options& options) {
+  auto strategies = llvm::to_vector(mlir::bisect::detail::GetStrategies());
+
+  auto apply_step = [&]() -> std::optional<OwningOpRef<ModuleOp>> {
+    for (auto it = strategies.begin(); it != strategies.end(); ++it) {
+      for (auto& candidate_fn :
+           detail::GetCandidates(it->second, state, *module)) {
+        auto candidate = candidate_fn();
+        if (!candidate || !mlir::verify(*candidate).succeeded()) {
+          continue;
+        }
+        if (options.canonicalize && !Canonicalize(*candidate).succeeded()) {
+          continue;
+        }
+
+        interpreter::ExecutionTrace trace;
+        // Verify that the candidate is still buggy.
+        if (!Run(*candidate, &trace, options).succeeded()) {
+          continue;
+        }
+
+        // Print the new buggy module.
+        llvm::outs() << "module after " << it->first << ":\n"
+                     << *candidate << "\n\n";
+
+        // Update the trace.
+        state.SetTrace(std::move(trace));
+
+        // Move strategies to the end.
+        decltype(strategies) new_strategies;
+        std::copy(it + 1, strategies.end(), std::back_inserter(new_strategies));
+        std::copy(strategies.begin(), it + 1,
+                  std::back_inserter(new_strategies));
+        strategies = std::move(new_strategies);
+        return {candidate.release()};
+      }
+    }
+    return std::nullopt;
+  };
+
+  while (auto new_module = apply_step()) {
+    module = std::move(*new_module);
+  }
+  return module;
+}
+
+void ReplaceArgsWithConstants(ModuleOp module,
+                              const xla::HloSnapshot& snapshot) {
+  auto main = llvm::cast<func::FuncOp>(module.lookupSymbol("main"));
+  OpBuilder b(main.getBody());
+  for (auto [arg, bbarg] :
+       llvm::zip(snapshot.arguments(), main.getBody().getArguments())) {
+    auto attr = interpreter::ValueToAttribute(
+        *interpreter::LiteralToValue(*xla::Literal::CreateFromProto(arg)),
+        bbarg.getType());
+    CHECK_EQ(attr.size(), 1) << "unsupported argument";
+
+    auto constant = b.create<arith::ConstantOp>(
+        main.getLoc(), bbarg.getType(), llvm::cast<TypedAttr>(attr.front()));
+    bbarg.replaceAllUsesWith(constant);
+  }
+
+  // The remaining ops are output args, so we replace them with allocs.
+  for (auto arg :
+       main.getBody().getArguments().drop_front(snapshot.arguments().size())) {
+    CHECK(llvm::isa<MemRefType>(arg.getType())) << "unsupported argument";
+    arg.replaceAllUsesWith(b.create<memref::AllocOp>(
+        module.getLoc(), llvm::cast<MemRefType>(arg.getType())));
+  }
+  while (main.getBody().getNumArguments() > 0) {
+    main.getBody().eraseArgument(0);
+  }
+  main.setFunctionType(FunctionType::get(main.getContext(), /*inputs=*/{},
+                                         main.getFunctionType().getResults()));
+  main.setArgAttrsAttr(b.getArrayAttr({}));
+}
+
+}  // namespace
+}  // namespace bisect
+}  // namespace mlir
+
+int main(int argc, char* argv[]) {
+  llvm::errs().tie(&llvm::outs());
+  llvm::outs().tie(&llvm::errs());
+  int dummy_argc = 1;
+  tsl::port::InitMain("", &dummy_argc, &argv);
+
+  Options options;
+  llvm::cl::ParseCommandLineOptions(argc, argv, "MLIR bisect tool\n");
+
+  mlir::DialectRegistry registry;
+  mlir::registerAllDialects(registry);
+  mlir::registerAllPasses();
+  mlir::bisect::test::RegisterTestPasses();
+  mlir::mhlo::registerAllMhloPasses();
+  mlir::mhlo::registerAllMhloDialects(registry);
+
+  registry.insert<mlir::arith::ArithDialect>();
+
+  mlir::MLIRContext context(registry);
+  context.getOrLoadDialect<mlir::arith::ArithDialect>();
+  auto module = mlir::bisect::ParseMlirInput(options.input_filename, &context);
+
+  if (!options.hlo_snapshot.empty()) {
+    xla::HloSnapshot snapshot;
+    CHECK_OK(tsl::ReadBinaryProto(tsl::Env::Default(), options.hlo_snapshot,
+                                  &snapshot));
+    mlir::bisect::ReplaceArgsWithConstants(*module, snapshot);
+  }
+
+  if (options.debug_strategy.empty()) {
+    llvm::outs() << "initial module:\n" << *module << "\n";
+  }
+
+  mlir::interpreter::ExecutionTrace trace;
+  if (!mlir::bisect::Run(*module, &trace, options).succeeded()) {
+    llvm::outs() << "Did not find bug in initial module\n";
+    if (options.pass_pipeline.hasAnyOccurrences() &&
+        mlir::succeeded(mlir::bisect::RunPipeline(*module, options))) {
+      llvm::outs() << "Module after running pipeline:\n" << *module << "\n";
+    }
+    return 1;
+  }
+
+  mlir::bisect::BisectState state;
+  state.SetTrace(std::move(trace));
+  if (!options.debug_strategy.empty()) {
+    bool some_failed = false;
+    for (auto& candidate : mlir::bisect::detail::GetCandidates(
+             mlir::bisect::detail::GetStrategies()[options.debug_strategy],
+             state, *module)) {
+      auto new_module = candidate();
+      if (!new_module) {
+        continue;
+      }
+      llvm::outs() << *new_module << "\n\n";
+      if (!mlir::verify(*new_module).succeeded()) {
+        some_failed = true;
+        llvm::errs() << "verification failed\n";
+      }
+    }
+    return some_failed ? 1 : 0;
+  }
+
+  module = mlir::bisect::ReduceModule(std::move(module), state, options);
+
+  llvm::outs() << "Final module:\n" << *module << "\n";
+  if (options.pass_pipeline.hasAnyOccurrences() &&
+      mlir::succeeded(mlir::bisect::RunPipeline(*module, options))) {
+    llvm::outs() << "Final module after running pipeline:\n" << *module << "\n";
+  }
+  return 0;
+}
diff --git a/third_party/xla/xla/mlir/tools/mlir_bisect/rewrites/BUILD b/third_party/xla/xla/mlir/tools/mlir_bisect/rewrites/BUILD
new file mode 100644
index 00000000000000..18620fa5b880ec
--- /dev/null
+++ b/third_party/xla/xla/mlir/tools/mlir_bisect/rewrites/BUILD
@@ -0,0 +1,26 @@
+load("@local_tsl//tsl/platform:rules_cc.bzl", "cc_library")
+
+# copybara:uncomment package(default_applicable_licenses = ["//tensorflow:license"])
+
+cc_library(
+    name = "rewrites",
+    srcs = [
+        "func.cc",
+        "general.cc",
+        "scf.cc",
+    ],
+    visibility = ["//visibility:public"],
+    deps = [
+        "//xla/mlir/tools/mlir_bisect:bisect_lib",
+        "//xla/mlir/tools/mlir_replay/public:execution_trace_utils",
+        "@llvm-project//llvm:Support",
+        "@llvm-project//mlir:ArithDialect",
+        "@llvm-project//mlir:DialectUtils",
+        "@llvm-project//mlir:FuncDialect",
+        "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:SCFDialect",
+        "@llvm-project//mlir:SideEffectInterfaces",
+        "@llvm-project//mlir:Support",
+    ],
+    alwayslink = 1,
+)
diff --git a/third_party/xla/xla/mlir/tools/mlir_bisect/rewrites/func.cc b/third_party/xla/xla/mlir/tools/mlir_bisect/rewrites/func.cc
new file mode 100644
index 00000000000000..3715f36b825692
--- /dev/null
+++ b/third_party/xla/xla/mlir/tools/mlir_bisect/rewrites/func.cc
@@ -0,0 +1,86 @@
+/* Copyright 2023 The OpenXLA Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+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.
+==============================================================================*/
+
+#include <functional>
+#include <iterator>
+#include <utility>
+
+#include "llvm/ADT/STLExtras.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
+#include "mlir/IR/BuiltinAttributes.h"  // from @llvm-project
+#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
+#include "mlir/IR/OwningOpRef.h"  // from @llvm-project
+#include "mlir/IR/Value.h"  // from @llvm-project
+#include "mlir/IR/ValueRange.h"  // from @llvm-project
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+#include "xla/mlir/tools/mlir_bisect/bisect_lib.h"
+
+namespace mlir {
+namespace bisect {
+namespace {
+
+void SetReturnValues(func::FuncOp func, ValueRange values) {
+  // We only operate on functions without arguments.
+  func.setFunctionType(mlir::FunctionType::get(func.getContext(), /*inputs=*/{},
+                                               values.getTypes()));
+  func.getBody().getBlocks().front().getTerminator()->setOperands(values);
+}
+
+SmallVector<std::function<OwningOpRef<ModuleOp>()>> TruncateFunction(
+    BisectState&, func::FuncOp func) {
+  SmallVector<std::function<OwningOpRef<ModuleOp>()>> result;
+  for (auto& ret : func.getBody().getBlocks().front().without_terminator()) {
+    if (func.getBody().getBlocks().front().getTerminator()->getOperands() ==
+        ret.getResults()) {
+      continue;
+    }
+    auto fun = [r = &ret]() -> OwningOpRef<ModuleOp> {
+      auto [module, ret_clone] = CloneModuleFor(r);
+      SetReturnValues(ret_clone->getParentOfType<func::FuncOp>(),
+                      ret_clone->getResults());
+      return std::move(module);
+    };
+    result.push_back(fun);
+  }
+  return result;
+}
+
+SmallVector<std::function<OwningOpRef<ModuleOp>()>>
+ReturnOperandsOfTerminatorOperands(BisectState&, func::FuncOp func) {
+  SmallVector<std::function<OwningOpRef<ModuleOp>()>> result;
+  result.push_back([func]() -> OwningOpRef<ModuleOp> {
+    auto [module, func_clone] = CloneModuleFor(func);
+    auto* terminator = func_clone.getBody().getBlocks().front().getTerminator();
+    SmallVector<Value> new_operands;
+    for (auto operand : terminator->getOperands()) {
+      if (operand.getDefiningOp()) {
+        llvm::copy(operand.getDefiningOp()->getOperands(),
+                   std::back_inserter(new_operands));
+      } else {
+        return nullptr;
+      }
+    }
+    SetReturnValues(func_clone, new_operands);
+    return std::move(module);
+  });
+  return result;
+}
+
+REGISTER_MLIR_REDUCE_STRATEGY(TruncateFunction);
+REGISTER_MLIR_REDUCE_STRATEGY(ReturnOperandsOfTerminatorOperands);
+
+}  // namespace
+}  // namespace bisect
+}  // namespace mlir
diff --git a/third_party/xla/xla/mlir/tools/mlir_bisect/rewrites/general.cc b/third_party/xla/xla/mlir/tools/mlir_bisect/rewrites/general.cc
new file mode 100644
index 00000000000000..e54e81aebeebb8
--- /dev/null
+++ b/third_party/xla/xla/mlir/tools/mlir_bisect/rewrites/general.cc
@@ -0,0 +1,194 @@
+/* Copyright 2023 The OpenXLA Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+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.
+==============================================================================*/
+
+#include <cassert>
+#include <cstdint>
+#include <functional>
+#include <utility>
+
+#include "llvm/ADT/STLExtras.h"
+#include "llvm/ADT/SmallVector.h"
+#include "llvm/Support/Casting.h"
+#include "mlir/Dialect/Arith/IR/Arith.h"  // from @llvm-project
+#include "mlir/IR/Builders.h"  // from @llvm-project
+#include "mlir/IR/BuiltinAttributeInterfaces.h"  // from @llvm-project
+#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
+#include "mlir/IR/OpDefinition.h"  // from @llvm-project
+#include "mlir/IR/Operation.h"  // from @llvm-project
+#include "mlir/IR/OwningOpRef.h"  // from @llvm-project
+#include "mlir/IR/Value.h"  // from @llvm-project
+#include "mlir/Interfaces/SideEffectInterfaces.h"  // from @llvm-project
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+#include "xla/mlir/tools/mlir_bisect/bisect_lib.h"
+#include "xla/mlir/tools/mlir_replay/public/execution_trace_utils.h"
+
+namespace mlir {
+namespace bisect {
+namespace {
+
+bool IsTerminator(Operation* op) {
+  return op->hasTrait<OpTrait::IsTerminator>();
+}
+
+bool IsTopLevelOp(Operation* op) {
+  return !op->getBlock()->back().mightHaveTrait<OpTrait::IsTerminator>();
+}
+
+SmallVector<std::function<OwningOpRef<ModuleOp>()>> EraseOpWithoutResults(
+    BisectState& state, Operation* op) {
+  // Only erase ops with results if they're unused.
+  if (op->getNumResults() > 0 && !op->use_empty()) {
+    return {};
+  }
+
+  // Don't erase entire functions, constants, terminators.
+  if (IsTopLevelOp(op) || IsTerminator(op)) {
+    return {};
+  }
+
+  SmallVector<std::function<OwningOpRef<ModuleOp>()>> ret;
+  ret.push_back([op]() {
+    auto [module, cloned_op] = CloneModuleFor(op);
+    cloned_op->erase();
+    return std::move(module);
+  });
+  return ret;
+}
+
+llvm::SmallVector<std::function<OwningOpRef<ModuleOp>()>> ReplaceOpWithConstant(
+    BisectState& state, Operation* op) {
+  llvm::SmallVector<std::function<OwningOpRef<ModuleOp>()>> result;
+  if (op->hasTrait<OpTrait::ConstantLike>() || IsTopLevelOp(op) ||
+      IsTerminator(op) || op->use_empty() || op->getNumResults() == 0) {
+    return result;
+  }
+
+  auto mii = llvm::dyn_cast<MemoryEffectOpInterface>(op);
+  if (mii && mii.hasEffect<MemoryEffects::Allocate>()) {
+    // Don't replace allocs with constants.
+    return result;
+  }
+
+  // Ops that are never executed won't be replaced here, but we have other
+  // strategies that get rid of them (e.g. deleting the entire region).
+  for (auto* execution : state.GetExecutions(op)) {
+    assert(execution->results_size() == op->getNumResults() &&
+           "unexpected number of results");
+
+    result.push_back([execution, op]() -> OwningOpRef<ModuleOp> {
+      auto [module_clone, op_clone] = CloneModuleFor(op);
+      SmallVector<Value> results;
+      OpBuilder b(op_clone);
+      for (int64_t i = 0; i < op->getNumResults(); ++i) {
+        auto type = op->getResultTypes()[i];
+        auto value = *interpreter::TracedValueToValue(
+            execution->results(static_cast<int>(i)));
+        auto attribute = interpreter::ValueToAttribute(value, type);
+        // We don't currently support tuples.
+        if (attribute.size() != 1) {
+          return nullptr;
+        }
+        op_clone->getResults()[i].replaceAllUsesWith(
+            b.create<arith::ConstantOp>(
+                op_clone->getLoc(), type,
+                llvm::cast<TypedAttr>(attribute.front())));
+      }
+      return std::move(module_clone);
+    });
+  }
+  return result;
+}
+
+llvm::SmallVector<std::function<OwningOpRef<ModuleOp>()>>
+ReplaceOperandWithConstant(BisectState& state, Operation* op) {
+  llvm::SmallVector<std::function<OwningOpRef<ModuleOp>()>> result;
+  if (IsTopLevelOp(op) || op->getNumOperands() == 0) {
+    return result;
+  }
+
+  for (auto* execution : state.GetExecutions(op)) {
+    for (int64_t i = 0; i < op->getNumOperands(); ++i) {
+      auto operand = op->getOperand(i);
+      if (operand.getDefiningOp() &&
+          operand.getDefiningOp()->hasTrait<OpTrait::ConstantLike>()) {
+        continue;
+      }
+      result.push_back([execution, i, op]() -> OwningOpRef<ModuleOp> {
+        auto type = op->getOperandTypes()[i];
+        auto value = *interpreter::TracedValueToValue(
+            execution->args(static_cast<int>(i)));
+        auto attribute = interpreter::ValueToAttribute(value, type);
+        if (attribute.size() != 1) {
+          return nullptr;
+        }
+        auto [module_clone, op_clone] = CloneModuleFor(op);
+        OpBuilder b(op_clone);
+        op_clone->setOperand(i, b.create<arith::ConstantOp>(
+                                    op_clone->getLoc(), type,
+                                    llvm::cast<TypedAttr>(attribute.front())));
+        return std::move(module_clone);
+      });
+    }
+  }
+  return result;
+}
+
+// Replaces an op's result with some other value with the same type defined
+// previously in the same region.
+llvm::SmallVector<std::function<OwningOpRef<ModuleOp>()>> ReplaceOpWithValue(
+    BisectState&, Operation* op) {
+  llvm::SmallVector<std::function<OwningOpRef<ModuleOp>()>> ret;
+  if (op->hasTrait<OpTrait::ConstantLike>() || IsTopLevelOp(op) ||
+      IsTerminator(op)) {
+    return ret;
+  }
+
+  // TODO(jreiffers): Consider bbargs.
+  llvm::DenseMap<mlir::Type, llvm::SmallVector<std::pair<Operation*, int64_t>>>
+      candidates_by_type;
+  for (auto* pred = op->getPrevNode(); pred != nullptr;
+       pred = pred->getPrevNode()) {
+    for (auto [index, result] : llvm::enumerate(pred->getResults())) {
+      candidates_by_type[result.getType()].emplace_back(pred, index);
+    }
+  }
+
+  for (auto [index, result] : llvm::enumerate(op->getResults())) {
+    if (result.use_empty()) {
+      continue;
+    }
+
+    for (auto [new_result_op, new_result_index] :
+         candidates_by_type[result.getType()]) {
+      ret.push_back(
+          [op, i = index, j = new_result_index, result_op = new_result_op]() {
+            auto [module_clone, op_clone] = CloneModuleFor(op);
+            op_clone->getResults()[i].replaceAllUsesWith(
+                FindInClone(result_op, module_clone.get())->getResults()[j]);
+            return std::move(module_clone);
+          });
+    }
+  }
+  return ret;
+}
+
+REGISTER_MLIR_REDUCE_STRATEGY(EraseOpWithoutResults);
+REGISTER_MLIR_REDUCE_STRATEGY(ReplaceOpWithConstant);
+REGISTER_MLIR_REDUCE_STRATEGY(ReplaceOpWithValue);
+REGISTER_MLIR_REDUCE_STRATEGY(ReplaceOperandWithConstant);
+
+}  // namespace
+}  // namespace bisect
+}  // namespace mlir
diff --git a/third_party/xla/xla/mlir/tools/mlir_bisect/rewrites/scf.cc b/third_party/xla/xla/mlir/tools/mlir_bisect/rewrites/scf.cc
new file mode 100644
index 00000000000000..3445393ff9798b
--- /dev/null
+++ b/third_party/xla/xla/mlir/tools/mlir_bisect/rewrites/scf.cc
@@ -0,0 +1,144 @@
+/* Copyright 2023 The OpenXLA Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+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.
+==============================================================================*/
+
+#include "mlir/Dialect/SCF/IR/SCF.h"  // from @llvm-project
+
+#include <cstdint>     // NOLINT
+#include <functional>  // NOLINT
+#include <utility>     // NOLINT
+
+#include "mlir/Dialect/Utils/StaticValueUtils.h"  // from @llvm-project
+#include "mlir/IR/Builders.h"  // from @llvm-project
+#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
+#include "mlir/IR/OpDefinition.h"  // from @llvm-project
+#include "mlir/IR/OwningOpRef.h"  // from @llvm-project
+#include "mlir/IR/PatternMatch.h"  // from @llvm-project
+#include "mlir/IR/Value.h"  // from @llvm-project
+#include "mlir/IR/ValueRange.h"  // from @llvm-project
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+#include "xla/mlir/tools/mlir_bisect/bisect_lib.h"
+
+namespace mlir {
+namespace bisect {
+namespace {
+
+constexpr int64_t kMaxWhileIterations = 1;
+
+// Rewrites a while loop to execute its body a fixed number of times. The
+// condition is executed, but its result is ignored.
+// For ease of implementation, this generates scf.execute_region ops. These are
+// subsequently canonicalized away.
+llvm::SmallVector<std::function<OwningOpRef<ModuleOp>()>> InlineScfWhile(
+    BisectState&, scf::WhileOp while_op) {
+  llvm::SmallVector<std::function<OwningOpRef<ModuleOp>()>> result;
+  for (int64_t num_executions = 0; num_executions <= kMaxWhileIterations;
+       ++num_executions) {
+    using ::mlir::scf::ExecuteRegionOp;
+
+    result.push_back([while_op, num_executions]() -> OwningOpRef<ModuleOp> {
+      auto [module, op] = CloneModuleFor(while_op);
+      OpBuilder b(op);
+      llvm::SmallVector<scf::ExecuteRegionOp> regions;
+
+      auto wrap_region_in_execute = [&,
+                                     loc = op.getLoc()](mlir::Region& region) {
+        regions
+            .emplace_back(b.create<ExecuteRegionOp>(
+                loc,
+                region.getBlocks().front().getTerminator()->getOperandTypes(),
+                mlir::ValueRange{}))
+            .getRegion()
+            .takeBody(region);
+      };
+
+      wrap_region_in_execute(op.getBefore());
+      // Replace the condition terminator with a yield terminator.
+      {
+        auto& before_block = regions[0].getRegion().getBlocks().front();
+        OpBuilder before_builder(before_block.getTerminator());
+        IRRewriter before_rewriter(before_builder);
+        before_rewriter.replaceOpWithNewOp<scf::YieldOp>(
+            before_block.getTerminator(),
+            before_block.getTerminator()->getOperands());
+      }
+
+      // Clone the execute region ops the requested number of times.
+      if (num_executions > 0) {
+        wrap_region_in_execute(op.getAfter());
+        for (int64_t i = 0; i < num_executions - 1; ++i) {
+          b.insert(regions.emplace_back(regions[0].clone()));
+          b.insert(regions.emplace_back(regions[1].clone()));
+        }
+        b.insert(regions.emplace_back(regions[0].clone()));
+      }
+
+      // Rewire region arguments and erase them.
+      for (int64_t i = 0; i < regions.size(); ++i) {
+        auto args = i == 0 ? ValueRange{op.getOperands()}
+                           : ValueRange{regions[i - 1].getResults()};
+        bool is_after_region = (i & 1) == 1;
+        auto& region = regions[i].getRegion();
+        for (int64_t arg = static_cast<int64_t>(region.getNumArguments()) - 1;
+             arg >= 0; --arg) {
+          region.getArgument(arg).replaceAllUsesWith(
+              args[is_after_region ? arg + 1 : arg]);
+          region.eraseArgument(arg);
+        }
+      }
+      op->replaceAllUsesWith(regions.back().getResults().drop_front(1));
+      op->erase();
+      return std::move(module);
+    });
+  }
+  return result;
+}
+
+SmallVector<std::function<OwningOpRef<ModuleOp>()>> ReduceScfForallBounds(
+    BisectState&, scf::ForallOp forall_op) {
+  SmallVector<OpFoldResult> new_upper_bound{forall_op.getMixedUpperBound()};
+  OpBuilder b(forall_op);
+  bool any_replaced = false;
+  for (auto& ub : new_upper_bound) {
+    auto constant_or = mlir::getConstantIntValue(ub);
+    if (!constant_or.has_value()) {
+      continue;
+    }
+    any_replaced = true;
+    ub = b.getIndexAttr(*constant_or - 1);
+  }
+  SmallVector<std::function<OwningOpRef<ModuleOp>()>> result;
+  if (!any_replaced) {
+    return result;
+  }
+  result.push_back([=]() -> OwningOpRef<ModuleOp> {
+    auto [module, op] = CloneModuleFor(forall_op);
+    OpBuilder b(op);
+    SmallVector<Value> dynamic_upper_bound;
+    SmallVector<int64_t> static_upper_bound;
+    dispatchIndexOpFoldResults(new_upper_bound, dynamic_upper_bound,
+                               static_upper_bound);
+    op.getDynamicUpperBoundMutable().assign(dynamic_upper_bound);
+    op.setStaticUpperBound(static_upper_bound);
+    return std::move(module);
+  });
+  return result;
+}
+
+REGISTER_MLIR_REDUCE_STRATEGY(ReduceScfForallBounds);
+REGISTER_MLIR_REDUCE_STRATEGY(InlineScfWhile);
+
+}  // namespace
+}  // namespace bisect
+}  // namespace mlir
diff --git a/third_party/xla/xla/mlir/tools/mlir_bisect/rewrites/tests/BUILD b/third_party/xla/xla/mlir/tools/mlir_bisect/rewrites/tests/BUILD
new file mode 100644
index 00000000000000..3cca2e53bb6945
--- /dev/null
+++ b/third_party/xla/xla/mlir/tools/mlir_bisect/rewrites/tests/BUILD
@@ -0,0 +1,13 @@
+load("//xla:lit.bzl", "lit_test_suite")
+
+# copybara:uncomment package(default_applicable_licenses = ["//tensorflow:license"])
+
+lit_test_suite(
+    name = "all_tests",
+    srcs = glob(["*.mlir"]),
+    cfg = "//xla:lit.cfg.py",
+    tools = [
+        "//xla/mlir/tools/mlir_bisect:mlir-bisect",
+        "@llvm-project//llvm:FileCheck",
+    ],
+)
diff --git a/third_party/xla/xla/mlir/tools/mlir_bisect/rewrites/tests/erase-op-without-results.mlir b/third_party/xla/xla/mlir/tools/mlir_bisect/rewrites/tests/erase-op-without-results.mlir
new file mode 100644
index 00000000000000..e918e112fe46f3
--- /dev/null
+++ b/third_party/xla/xla/mlir/tools/mlir_bisect/rewrites/tests/erase-op-without-results.mlir
@@ -0,0 +1,12 @@
+// RUN: mlir-bisect %s --debug-strategy=EraseOpWithoutResults | FileCheck %s
+
+func.func @main() -> memref<i32> {
+  %a = arith.constant 1 : i32
+  %b = memref.alloc() : memref<i32>
+  memref.store %a, %b[] : memref<i32>
+  func.return %b : memref<i32>
+}
+
+//      CHECK: func.func @main()
+//      CHECK:   %[[ALLOC:.*]] = memref.alloc
+// CHECK-NEXT:   return %[[ALLOC]]
diff --git a/third_party/xla/xla/mlir/tools/mlir_bisect/rewrites/tests/inline-scf-while.mlir b/third_party/xla/xla/mlir/tools/mlir_bisect/rewrites/tests/inline-scf-while.mlir
new file mode 100644
index 00000000000000..6c9deddbc37cb5
--- /dev/null
+++ b/third_party/xla/xla/mlir/tools/mlir_bisect/rewrites/tests/inline-scf-while.mlir
@@ -0,0 +1,40 @@
+// RUN: mlir-bisect %s --debug-strategy=InlineScfWhile | FileCheck %s
+
+func.func @main() -> i64 {
+  %c0 = arith.constant 0 : i64
+  %c1 = arith.constant 1 : i64
+  %c4 = arith.constant 4 : i64
+  %alloc = memref.alloc() : memref<i64>
+  memref.store %c0, %alloc[] : memref<i64>
+  %ret = scf.while(%arg0 = %c0): (i64) -> (i64) {
+    %cond = arith.cmpi slt, %arg0, %c4 : i64
+    scf.condition(%cond) %arg0 : i64
+  } do {
+  ^bb0(%arg1: i64):
+    %add = arith.addi %arg1, %c1 : i64
+    scf.yield %add : i64
+  }
+  return %ret : i64
+}
+
+//     CHECK: func @main
+// CHECK-DAG:   %[[C0:.*]] = arith.constant 0
+// CHECK-DAG:   %[[C4:.*]] = arith.constant 4
+//     CHECK:   %[[RET:.*]]:2 = scf.execute_region
+//     CHECK:     arith.cmpi slt, %[[C0]], %[[C4]]
+//     CHECK:     yield {{.*}}, %[[C0]]
+//     CHECK:   return %[[RET]]#1
+
+//     CHECK: func @main
+// CHECK-DAG:   %[[C0:.*]] = arith.constant 0
+// CHECK-DAG:   %[[C1:.*]] = arith.constant 1
+//     CHECK:   %[[BEFORE0:.*]]:2 = scf.execute_region
+//     CHECK:     arith.cmpi
+//     CHECK:     yield {{.*}}, %[[C0]]
+//     CHECK:   %[[AFTER:.*]] = scf.execute_region
+//     CHECK:     %[[ADD:.*]] = arith.addi %[[BEFORE0]]#1, %[[C1]]
+//     CHECK:     yield %[[ADD]]
+//     CHECK:   %[[BEFORE1:.*]]:2 = scf.execute_region
+//     CHECK:     arith.cmpi
+//     CHECK:     yield {{.*}}, %[[AFTER]]
+//     CHECK:   return %[[BEFORE1]]#1
\ No newline at end of file
diff --git a/third_party/xla/xla/mlir/tools/mlir_bisect/rewrites/tests/reduce-scf-forall-bounds.mlir b/third_party/xla/xla/mlir/tools/mlir_bisect/rewrites/tests/reduce-scf-forall-bounds.mlir
new file mode 100644
index 00000000000000..61f289d3c5cd6c
--- /dev/null
+++ b/third_party/xla/xla/mlir/tools/mlir_bisect/rewrites/tests/reduce-scf-forall-bounds.mlir
@@ -0,0 +1,16 @@
+// RUN: mlir-bisect %s --debug-strategy=ReduceScfForallBounds | FileCheck %s
+
+func.func @main() -> tensor<8xindex> {
+  %init = tensor.empty() : tensor<8xindex>
+  %iota = scf.forall (%i) = (0) to (8) step (1)
+      shared_outs (%init_ = %init) -> (tensor<8xindex>) {
+    %tensor = tensor.from_elements %i : tensor<1xindex>
+    scf.forall.in_parallel {
+      tensor.parallel_insert_slice %tensor into %init_[%i] [1] [1]
+        : tensor<1xindex> into tensor<8xindex>
+    }
+  }
+  func.return %iota : tensor<8xindex>
+}
+// CHECK: func @main()
+// CHECK:   scf.forall ({{.*}}) in (7)
diff --git a/third_party/xla/xla/mlir/tools/mlir_bisect/rewrites/tests/replace-op-with-constant.mlir b/third_party/xla/xla/mlir/tools/mlir_bisect/rewrites/tests/replace-op-with-constant.mlir
new file mode 100644
index 00000000000000..171472ad733642
--- /dev/null
+++ b/third_party/xla/xla/mlir/tools/mlir_bisect/rewrites/tests/replace-op-with-constant.mlir
@@ -0,0 +1,26 @@
+// RUN: mlir-bisect %s --debug-strategy=ReplaceOpWithConstant | FileCheck %s
+
+func.func @main() -> tensor<2xi32> {
+  %a = arith.constant dense<3> : tensor<2xi32>
+  %b = arith.constant dense<2> : tensor<2xi32>
+  %c = mhlo.add %a, %b : tensor<2xi32>
+  %d = mhlo.multiply %b, %c : tensor<2xi32>
+  func.return %d : tensor<2xi32>
+}
+
+//      CHECK: func.func @main()
+// CHECK-NEXT:   arith.constant dense<3>
+// CHECK-NEXT:   arith.constant dense<2>
+// CHECK-NEXT:   arith.constant dense<5>
+// CHECK-NEXT:   %[[ADD:.*]] = mhlo.add
+//  CHECK-NOT:   %[[ADD]]
+// CHECK-NEXT:   mhlo.multiply
+// CHECK-NEXT:   return
+
+//      CHECK: func.func @main()
+// CHECK-NEXT:   arith.constant dense<3>
+// CHECK-NEXT:   arith.constant dense<2>
+// CHECK-NEXT:   mhlo.add
+// CHECK-NEXT:   %[[D:.*]] = arith.constant dense<10>
+// CHECK-NEXT:   mhlo.multiply
+// CHECK-NEXT:   return %[[D]]
diff --git a/third_party/xla/xla/mlir/tools/mlir_bisect/rewrites/tests/replace-op-with-value.mlir b/third_party/xla/xla/mlir/tools/mlir_bisect/rewrites/tests/replace-op-with-value.mlir
new file mode 100644
index 00000000000000..f89f647f14ddc6
--- /dev/null
+++ b/third_party/xla/xla/mlir/tools/mlir_bisect/rewrites/tests/replace-op-with-value.mlir
@@ -0,0 +1,16 @@
+// RUN: mlir-bisect %s --debug-strategy=ReplaceOpWithValue | FileCheck %s
+
+func.func @main() -> (memref<i32>, memref<i32>) {
+  %a = memref.alloc() : memref<i32>
+  %b = memref.alloc() : memref<i32>
+  %c0 = arith.constant 0 : i32
+  memref.store %c0, %b[] : memref<i32>
+  return %a, %b : memref<i32>, memref<i32>
+}
+
+//      CHECK: func @main()
+//      CHECK:   %[[ALLOC:.*]] = memref.alloc()
+// CHECK-NEXT:   memref.alloc
+// CHECK-NEXT:   constant
+// CHECK-NEXT:   memref.store {{.*}}, %[[ALLOC]]
+// CHECK-NEXT:   return %[[ALLOC]], %[[ALLOC]]
diff --git a/third_party/xla/xla/mlir/tools/mlir_bisect/rewrites/tests/replace-operand-with-constant.mlir b/third_party/xla/xla/mlir/tools/mlir_bisect/rewrites/tests/replace-operand-with-constant.mlir
new file mode 100644
index 00000000000000..7619a8a500c5e4
--- /dev/null
+++ b/third_party/xla/xla/mlir/tools/mlir_bisect/rewrites/tests/replace-operand-with-constant.mlir
@@ -0,0 +1,28 @@
+// RUN: mlir-bisect %s --debug-strategy=ReplaceOperandWithConstant | FileCheck %s
+
+func.func @main() -> (tensor<2xi32>, tensor<2xi32>) {
+  %a = arith.constant dense<3> : tensor<2xi32>
+  %b = arith.constant dense<2> : tensor<2xi32>
+  %c = mhlo.add %a, %b : tensor<2xi32>
+  %d = mhlo.multiply %b, %c : tensor<2xi32>
+  func.return %c, %d : tensor<2xi32>, tensor<2xi32>
+}
+
+// CHECK: func @main()
+// CHECK:   %[[C2:.*]] = arith.constant dense<2>
+// CHECK:   %[[ADD:.*]] = mhlo.add
+// CHECK:   %[[C5:.*]] = arith.constant dense<5>
+// CHECK:   %[[MUL:.*]] = mhlo.multiply %[[C2]], %[[C5]] : tensor<2xi32>
+// CHECK:   return %[[ADD]], %[[MUL]]
+
+// CHECK: func @main()
+// CHECK:   mhlo.add
+// CHECK:   %[[MUL:.*]] = mhlo.multiply %cst_0, %0 : tensor<2xi32>
+// CHECK:   %[[C5:.*]] = arith.constant dense<5>
+// CHECK:   return %[[C5]], %[[MUL]]
+
+// CHECK: func @main()
+// CHECK:   %[[ADD:.*]] = mhlo.add
+// CHECK:   mhlo.multiply
+// CHECK:   %[[C10:.*]] = arith.constant dense<10>
+// CHECK:   return %[[ADD]], %[[C10]]
diff --git a/third_party/xla/xla/mlir/tools/mlir_bisect/rewrites/tests/return-operands-of-terminator-operands.mlir b/third_party/xla/xla/mlir/tools/mlir_bisect/rewrites/tests/return-operands-of-terminator-operands.mlir
new file mode 100644
index 00000000000000..8584e2a0008fa0
--- /dev/null
+++ b/third_party/xla/xla/mlir/tools/mlir_bisect/rewrites/tests/return-operands-of-terminator-operands.mlir
@@ -0,0 +1,15 @@
+// RUN: mlir-bisect %s --debug-strategy=ReturnOperandsOfTerminatorOperands | FileCheck %s
+
+func.func @main() -> tensor<2xi32> {
+  %a = arith.constant dense<3> : tensor<2xi32>
+  %b = arith.constant dense<2> : tensor<2xi32>
+  %c = mhlo.add %a, %b : tensor<2xi32>
+  %d = mhlo.multiply %b, %c : tensor<2xi32>
+  func.return %d : tensor<2xi32>
+}
+
+// CHECK: @main
+// CHECK:   %[[C2:.*]] = arith.constant dense<2>
+// CHECK:   %[[ADD:.*]] = mhlo.add
+// CHECK:   mhlo.multiply
+// CHECK:   return %[[C2]], %[[ADD]]
\ No newline at end of file
diff --git a/third_party/xla/xla/mlir/tools/mlir_bisect/rewrites/tests/truncate-function.mlir b/third_party/xla/xla/mlir/tools/mlir_bisect/rewrites/tests/truncate-function.mlir
new file mode 100644
index 00000000000000..af06778bd47c54
--- /dev/null
+++ b/third_party/xla/xla/mlir/tools/mlir_bisect/rewrites/tests/truncate-function.mlir
@@ -0,0 +1,31 @@
+// RUN: mlir-bisect %s --debug-strategy=TruncateFunction | FileCheck %s
+
+// Function to prevent constant folding below.
+func.func private @cst() -> tensor<2xi32> {
+  %cst = arith.constant dense<2> : tensor<2xi32>
+  return %cst : tensor<2xi32>
+}
+
+func.func @main() -> tensor<2xi32> {
+  %a = arith.constant dense<1> : tensor<2xi32>
+  %b = func.call @cst() : () -> tensor<2xi32>
+  %c = mhlo.add %a, %b : tensor<2xi32>
+  %d = mhlo.multiply %b, %c : tensor<2xi32>
+  func.return %d : tensor<2xi32>
+}
+
+//     CHECK: func @main()
+//     CHECK:   %[[A:.*]] = arith.constant dense<1>
+//     CHECK:   return %[[A]]
+
+//     CHECK: func @main()
+//     CHECK:   %[[B:.*]] = call @cst()
+//     CHECK:   return %[[B]]
+
+//     CHECK: func @main()
+//     CHECK:   %[[A:.*]] = arith.constant dense<1>
+//     CHECK:   %[[B:.*]] = call @cst()
+//     CHECK:   %[[ADD:.*]] = mhlo.add
+// CHECK-DAG:   %[[A]]
+// CHECK-DAG:   %[[B]]
+//     CHECK:   return %[[ADD]]
diff --git a/third_party/xla/xla/mlir/tools/mlir_bisect/test_passes.cc b/third_party/xla/xla/mlir/tools/mlir_bisect/test_passes.cc
new file mode 100644
index 00000000000000..a48b6a0b8cc2e9
--- /dev/null
+++ b/third_party/xla/xla/mlir/tools/mlir_bisect/test_passes.cc
@@ -0,0 +1,48 @@
+/* Copyright 2023 The OpenXLA Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+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.
+==============================================================================*/
+
+#include "xla/mlir/tools/mlir_bisect/test_passes.h"
+
+#include "mlir/Dialect/Linalg/IR/Linalg.h"  // from @llvm-project
+#include "mlir/Pass/Pass.h"  // from @llvm-project
+
+namespace mlir {
+namespace bisect {
+namespace test {
+namespace {
+
+struct BreakLinalgTransposePass
+    : public PassWrapper<BreakLinalgTransposePass, OperationPass<ModuleOp>> {
+  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(BreakLinalgTransposePass)
+
+  StringRef getArgument() const final { return "test-break-linalg-transpose"; }
+  StringRef getDescription() const final { return "breaks linalg transpose"; }
+  BreakLinalgTransposePass() = default;
+
+  void runOnOperation() override {
+    getOperation().walk([](linalg::TransposeOp op) {
+      auto permutation = llvm::to_vector(op.getPermutation());
+      std::swap(permutation[0], permutation[1]);
+      op.setPermutation(permutation);
+    });
+  }
+};
+}  // namespace
+
+void RegisterTestPasses() { PassRegistration<BreakLinalgTransposePass>(); }
+
+}  // namespace test
+}  // namespace bisect
+}  // namespace mlir
diff --git a/third_party/xla/xla/mlir/tools/mlir_bisect/test_passes.h b/third_party/xla/xla/mlir/tools/mlir_bisect/test_passes.h
new file mode 100644
index 00000000000000..b90f7e6b86529d
--- /dev/null
+++ b/third_party/xla/xla/mlir/tools/mlir_bisect/test_passes.h
@@ -0,0 +1,29 @@
+/* Copyright 2023 The OpenXLA Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+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.
+==============================================================================*/
+
+#ifndef XLA_MLIR_TOOLS_MLIR_BISECT_TEST_PASSES_H_
+#define XLA_MLIR_TOOLS_MLIR_BISECT_TEST_PASSES_H_
+
+namespace mlir {
+namespace bisect {
+namespace test {
+
+void RegisterTestPasses();
+
+}
+}  // namespace bisect
+}  // namespace mlir
+
+#endif  // XLA_MLIR_TOOLS_MLIR_BISECT_TEST_PASSES_H_
diff --git a/third_party/xla/xla/mlir/tools/mlir_bisect/tests/BUILD b/third_party/xla/xla/mlir/tools/mlir_bisect/tests/BUILD
new file mode 100644
index 00000000000000..5fb8b35f77d82a
--- /dev/null
+++ b/third_party/xla/xla/mlir/tools/mlir_bisect/tests/BUILD
@@ -0,0 +1,17 @@
+load("//xla:lit.bzl", "lit_test_suite")
+
+# copybara:uncomment package(default_applicable_licenses = ["//tensorflow:license"])
+
+lit_test_suite(
+    name = "all_tests",
+    srcs = glob(["*.mlir"]),
+    cfg = "//xla:lit.cfg.py",
+    data = [
+        "snapshot.mlir.pb",
+    ],
+    tools = [
+        "//xla/mlir/tools/mlir_bisect:mlir-bisect",
+        "@llvm-project//llvm:FileCheck",
+        "@llvm-project//llvm:not",
+    ],
+)
diff --git a/third_party/xla/xla/mlir/tools/mlir_bisect/tests/bisect.mlir b/third_party/xla/xla/mlir/tools/mlir_bisect/tests/bisect.mlir
new file mode 100644
index 00000000000000..ca839d982c416a
--- /dev/null
+++ b/third_party/xla/xla/mlir/tools/mlir_bisect/tests/bisect.mlir
@@ -0,0 +1,46 @@
+// RUN: mlir-bisect %s \
+// RUN: --pass-pipeline="builtin.module(test-break-linalg-transpose)" \
+// RUN: --max-steps-per-run=200 \
+// RUN: | FileCheck %s
+
+func.func @main() -> (memref<2x2xindex>, memref<2x2xindex>) {
+  %a = memref.alloc() : memref<2x2xindex>
+  %b = memref.alloc() : memref<2x2xindex>
+  %c = memref.alloc() : memref<2x2xindex>
+  %c0 = arith.constant 0 : index
+  %c1 = arith.constant 1 : index
+  %c2 = arith.constant 2 : index
+  %c3 = arith.constant 3 : index
+  scf.for %i = %c0 to %c2 step %c1 {
+    scf.for %j = %c0 to %c2 step %c1 {
+      memref.store %i, %a[%i, %j] : memref<2x2xindex>
+      memref.store %j, %b[%i, %j] : memref<2x2xindex>
+    }
+  }
+
+  %i = scf.while: () -> (index) {
+    %value = memref.load %a[%c0, %c0] : memref<2x2xindex>
+    %cond = arith.cmpi slt, %value, %c3 : index
+    scf.condition(%cond) %value : index
+  } do {
+  ^bb0(%_: index):
+    %value = memref.load %a[%c0, %c0] : memref<2x2xindex>
+    %add = arith.addi %value, %c1 : index
+    memref.store %add, %a[%c0, %c0] : memref<2x2xindex>
+    linalg.transpose ins(%b : memref<2x2xindex>) outs(%c : memref<2x2xindex>)
+      permutation = [1, 0]
+    memref.copy %c, %b : memref<2x2xindex> to memref<2x2xindex>
+    scf.yield
+  }
+
+  return %a, %b : memref<2x2xindex>, memref<2x2xindex>
+}
+
+//     CHECK: Final module
+//     CHECK: func @main() -> memref<2x2xindex> {
+// CHECK-NOT: scf.while
+// CHECK-NOT: scf.for
+//     CHECK: linalg.transpose {{.*}} permutation = [1, 0]
+
+//     CHECK: Final module after running pipeline
+//     CHECK: linalg.transpose {{.*}} permutation = [0, 1]
diff --git a/third_party/xla/xla/mlir/tools/mlir_bisect/tests/no-bug.mlir b/third_party/xla/xla/mlir/tools/mlir_bisect/tests/no-bug.mlir
new file mode 100644
index 00000000000000..df343f3bf8b09f
--- /dev/null
+++ b/third_party/xla/xla/mlir/tools/mlir_bisect/tests/no-bug.mlir
@@ -0,0 +1,10 @@
+// RUN: not mlir-bisect %s \
+// RUN: --pass-pipeline="builtin.module(test-break-linalg-transpose)" \
+// RUN: | FileCheck %s
+
+func.func @main() -> memref<2x2xindex> {
+  %a = memref.alloc() : memref<2x2xindex>
+  return %a : memref<2x2xindex>
+}
+
+// CHECK: Did not find bug in initial module
diff --git a/third_party/xla/xla/mlir/tools/mlir_bisect/tests/snapshot.mlir b/third_party/xla/xla/mlir/tools/mlir_bisect/tests/snapshot.mlir
new file mode 100644
index 00000000000000..916ca47ab0fd8e
--- /dev/null
+++ b/third_party/xla/xla/mlir/tools/mlir_bisect/tests/snapshot.mlir
@@ -0,0 +1,12 @@
+// RUN: not mlir-bisect %s --hlo-snapshot=%s.pb \
+// RUN: --pass-pipeline="builtin.module(test-break-linalg-transpose)" \
+// RUN: | FileCheck %s
+
+func.func @main(%a: tensor<3x1xi32>, %b: tensor<3x1xi32>) -> tensor<3x1xi32> {
+  return %a : tensor<3x1xi32>
+}
+
+// CHECK: initial module
+// CHECK: func @main() -> tensor<3x1xi32> {
+// CHECK{LITERAL}: arith.constant dense<[[2], [-4], [5]]> : tensor<3x1xi32>
+// CHECK{LITERAL}: arith.constant dense<[[0], [7], [-5]]> : tensor<3x1xi32>
diff --git a/third_party/xla/xla/mlir/tools/mlir_bisect/tests/snapshot.mlir.pb b/third_party/xla/xla/mlir/tools/mlir_bisect/tests/snapshot.mlir.pb
new file mode 100644
index 0000000000000000000000000000000000000000..ee3c8f759494db153cd7114783124b1cb7fb5da0
GIT binary patch
literal 68
scmWeq;1UpEkz!(I)MDXcVq`F4Vqj3>VfynQ3K&_1u&Q8S{|#3H00LzcG5`Po

literal 0
HcmV?d00001

diff --git a/third_party/xla/xla/mlir/tools/mlir_interpreter/BUILD b/third_party/xla/xla/mlir/tools/mlir_interpreter/BUILD
new file mode 100644
index 00000000000000..87d94287d3285a
--- /dev/null
+++ b/third_party/xla/xla/mlir/tools/mlir_interpreter/BUILD
@@ -0,0 +1,40 @@
+load("@bazel_skylib//rules:build_test.bzl", "build_test")
+load("//xla:xla.bzl", "xla_cc_binary")
+
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
+    default_visibility = [":friends"],
+    licenses = ["notice"],
+)
+
+package_group(
+    name = "friends",
+    includes = [
+        "//xla:friends",
+    ],
+)
+
+build_test(
+    name = "mlir-interpreter-runner_build_test",
+    targets = [
+        ":mlir-interpreter-runner",
+    ],
+)
+
+xla_cc_binary(
+    name = "mlir-interpreter-runner",
+    testonly = True,
+    srcs = ["mlir_interpreter_runner.cc"],
+    visibility = ["//visibility:public"],
+    deps = [
+        "//xla/mlir/tools/mlir_interpreter/dialects",
+        "//xla/mlir/tools/mlir_interpreter/framework",
+        "//xla/mlir_hlo:hlo_dialect_registration",
+        "//xla/mlir_hlo:mhlo_passes",
+        "@llvm-project//llvm:Support",
+        "@llvm-project//mlir:AllPassesAndDialects",
+        "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:ParseUtilities",
+        "@llvm-project//mlir:Support",
+    ],
+)
diff --git a/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/BUILD b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/BUILD
new file mode 100644
index 00000000000000..e6f5d71d0e19f6
--- /dev/null
+++ b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/BUILD
@@ -0,0 +1,67 @@
+load("@local_tsl//tsl/platform:rules_cc.bzl", "cc_library")
+
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
+    default_visibility = [":friends"],
+    licenses = ["notice"],
+)
+
+package_group(
+    name = "friends",
+    includes = [
+        "//xla:friends",
+    ],
+)
+
+cc_library(
+    name = "dialects",
+    srcs = glob(
+        [
+            "*.cc",
+        ],
+        exclude = ["util.cc"],
+    ),
+    deps = [
+        ":dialect_utils",
+        "//xla/mlir/tools/mlir_interpreter/framework",
+        "//xla/mlir_hlo",
+        "@llvm-project//llvm:Support",
+        "@llvm-project//mlir:AffineDialect",
+        "@llvm-project//mlir:AffineUtils",
+        "@llvm-project//mlir:ArithDialect",
+        "@llvm-project//mlir:BufferizationDialect",
+        "@llvm-project//mlir:ComplexDialect",
+        "@llvm-project//mlir:FuncDialect",
+        "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:LinalgDialect",
+        "@llvm-project//mlir:MemRefDialect",
+        "@llvm-project//mlir:SCFDialect",
+        "@llvm-project//mlir:Support",
+        "@llvm-project//mlir:TensorDialect",
+        "@llvm-project//mlir:VectorDialect",
+    ],
+    alwayslink = 1,
+)
+
+cc_library(
+    name = "dialect_utils",
+    srcs = [
+        "util.cc",
+    ],
+    hdrs = [
+        "comparators.h",
+        "cwise_math.h",
+        "util.h",
+    ],
+    deps = [
+        "//xla/mlir/tools/mlir_interpreter/framework",
+        "@com_google_absl//absl/strings",
+        "@llvm-project//llvm:Support",
+        "@llvm-project//mlir:ArithDialect",
+        "@llvm-project//mlir:FuncDialect",
+        "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:Support",
+        "@llvm-project//mlir:TensorDialect",
+        "@llvm-project//mlir:ViewLikeInterface",
+    ],
+)
diff --git a/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/affine.cc b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/affine.cc
new file mode 100644
index 00000000000000..3b952f5c4e54db
--- /dev/null
+++ b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/affine.cc
@@ -0,0 +1,53 @@
+/* Copyright 2022 The OpenXLA Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+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.
+==============================================================================*/
+
+#include <algorithm>
+#include <cstdint>
+
+#include "llvm/ADT/SmallVector.h"
+#include "mlir/Dialect/Affine/IR/AffineOps.h"  // from @llvm-project
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+#include "xla/mlir/tools/mlir_interpreter/dialects/util.h"
+#include "xla/mlir/tools/mlir_interpreter/framework/interpreter.h"
+#include "xla/mlir/tools/mlir_interpreter/framework/registration.h"
+
+namespace mlir {
+namespace interpreter {
+namespace {
+
+llvm::SmallVector<int64_t> Apply(InterpreterState&, affine::AffineApplyOp op,
+                                 ArrayRef<int64_t> operands) {
+  return EvalAffineMap(op.getAffineMap(), operands);
+}
+
+int64_t Min(InterpreterState&, affine::AffineMinOp op,
+            ArrayRef<int64_t> operands) {
+  auto results = EvalAffineMap(op.getAffineMap(), operands);
+  return *std::min_element(results.begin(), results.end());
+}
+
+int64_t Max(InterpreterState&, affine::AffineMaxOp op,
+            ArrayRef<int64_t> operands) {
+  auto results = EvalAffineMap(op.getAffineMap(), operands);
+  return *std::max_element(results.begin(), results.end());
+}
+
+REGISTER_MLIR_INTERPRETER_OP(Apply);
+REGISTER_MLIR_INTERPRETER_OP(Max);
+REGISTER_MLIR_INTERPRETER_OP(Min);
+
+}  // namespace
+}  // namespace interpreter
+}  // namespace mlir
diff --git a/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/arith.cc b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/arith.cc
new file mode 100644
index 00000000000000..17fd94c3019141
--- /dev/null
+++ b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/arith.cc
@@ -0,0 +1,311 @@
+/* Copyright 2022 The OpenXLA Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+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.
+==============================================================================*/
+
+#include "mlir/Dialect/Arith/IR/Arith.h"  // from @llvm-project
+
+#include <variant>  // NOLINT
+
+#include "llvm/ADT/SmallVector.h"
+#include "llvm/Support/ErrorHandling.h"
+#include "mlir/IR/Builders.h"  // from @llvm-project
+#include "mlir/IR/BuiltinAttributes.h"  // from @llvm-project
+#include "mlir/IR/BuiltinTypeInterfaces.h"  // from @llvm-project
+#include "mlir/IR/Types.h"  // from @llvm-project
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+#include "xla/mlir/tools/mlir_interpreter/dialects/comparators.h"
+#include "xla/mlir/tools/mlir_interpreter/dialects/cwise_math.h"
+#include "xla/mlir/tools/mlir_interpreter/framework/interpreter.h"
+#include "xla/mlir/tools/mlir_interpreter/framework/interpreter_value.h"
+#include "xla/mlir/tools/mlir_interpreter/framework/interpreter_value_util.h"
+#include "xla/mlir/tools/mlir_interpreter/framework/registration.h"
+#include "xla/mlir/tools/mlir_interpreter/framework/tensor_or_memref.h"
+
+namespace mlir {
+namespace interpreter {
+namespace {
+
+InterpreterValue Bitcast(InterpreterState&, arith::BitcastOp op,
+                         const InterpreterValue& in) {
+  Type ty = op->getResultTypes()[0];
+  auto shaped_ty = dyn_cast<ShapedType>(ty);
+  auto result = DispatchScalarType(ty, [&](auto dummy) -> InterpreterValue {
+    TensorOrMemref<decltype(dummy)> result;
+    result.view = {};
+    if (shaped_ty) {
+      result.buffer = in.Clone().Buffer();
+    } else {
+      result.buffer = in.AsUnitTensor().Buffer();
+    }
+    return {result};
+  });
+  if (!shaped_ty) {
+    return result.ExtractElement({});
+  }
+  auto& outView = result.View();
+  outView.strides = BufferView::GetDefaultStrides(shaped_ty.getShape());
+  outView.sizes = llvm::to_vector(shaped_ty.getShape());
+  return result;
+}
+
+InterpreterValue Constant(InterpreterState&, arith::ConstantOp constant) {
+  auto ty = constant->getResultTypes()[0];
+  auto shaped_ty = dyn_cast<ShapedType>(ty);
+  auto elem_ty = shaped_ty ? shaped_ty.getElementType() : ty;
+  return DispatchScalarType(elem_ty, [&](auto dummy) -> InterpreterValue {
+    using T = decltype(dummy);
+    if (shaped_ty) {
+      auto values = cast<DenseElementsAttr>(constant.getValue()).getValues<T>();
+      auto result = TensorOrMemref<T>::Empty(shaped_ty.getShape());
+      auto valueIt = values.begin();
+      result.view.is_vector = isa<VectorType>(shaped_ty);
+      for (const auto& index : result.view.Indices(true)) {
+        result.at(index) = *valueIt;
+        ++valueIt;
+      }
+      return {result};
+    }
+
+    auto value = constant.getValue();
+    if (auto integer = value.dyn_cast<IntegerAttr>()) {
+      return {static_cast<T>(integer.getInt())};
+    }
+    if (auto floatValue = value.dyn_cast<FloatAttr>()) {
+      return {static_cast<T>(floatValue.getValueAsDouble())};
+    }
+
+    llvm_unreachable("unsupported constant type");
+  });
+}
+
+template <typename Op>
+InterpreterValue IntCast(InterpreterState&, Op op,
+                         const InterpreterValue& arg) {
+  if (arg.IsTensor()) {
+    return DispatchScalarType(
+        op->getResultTypes()[0], [&](auto dummy) -> InterpreterValue {
+          auto result = TensorOrMemref<decltype(dummy)>::EmptyLike(arg.View());
+          for (const auto& index : result.view.Indices()) {
+            result.at(index) =
+                static_cast<decltype(dummy)>(arg.ExtractElement(index).AsInt());
+          }
+          return {result};
+        });
+  }
+
+  return DispatchScalarType(
+      op->getResultTypes()[0], [&](auto dummy) -> InterpreterValue {
+        return {static_cast<decltype(dummy)>(arg.AsInt())};
+      });
+}
+
+template <typename Op>
+InterpreterValue FloatCast(InterpreterState&, Op op,
+                           const InterpreterValue& arg) {
+  if (arg.IsTensor()) {
+    return DispatchScalarType(
+        op->getResultTypes()[0], [&](auto dummy) -> InterpreterValue {
+          auto result = TensorOrMemref<decltype(dummy)>::EmptyLike(arg.View());
+          for (const auto& index : result.view.Indices()) {
+            result.at(index) = static_cast<decltype(dummy)>(
+                arg.ExtractElement(index).AsDouble());
+          }
+          return {result};
+        });
+  }
+
+  return DispatchScalarType(
+      op->getResultTypes()[0], [&](auto dummy) -> InterpreterValue {
+        return {static_cast<decltype(dummy)>(arg.AsDouble())};
+      });
+}
+
+llvm::SmallVector<InterpreterValue> UiToFP(
+    MutableArrayRef<InterpreterValue> args, mlir::Operation* op,
+    InterpreterState&) {
+  if (args[0].IsTensor()) {
+    auto ty = op->getResultTypes()[0].cast<ShapedType>();
+    return {DispatchScalarType(
+        ty.getElementType(), [&](auto dummy) -> InterpreterValue {
+          auto result =
+              TensorOrMemref<decltype(dummy)>::EmptyLike(args[0].View());
+          for (const auto& index : result.view.Indices()) {
+            result.at(index) = static_cast<decltype(dummy)>(
+                args[0].ExtractElement(index).AsUInt());
+          }
+          return {result};
+        })};
+  }
+
+  return {DispatchScalarType(
+      op->getResultTypes()[0], [&](auto dummy) -> InterpreterValue {
+        return {static_cast<decltype(dummy)>(args[0].AsUInt())};
+      })};
+}
+
+InterpreterValue CmpI(InterpreterState&, arith::CmpIOp compare,
+                      const InterpreterValue& lhs,
+                      const InterpreterValue& rhs) {
+  switch (compare.getPredicate()) {
+    case arith::CmpIPredicate::eq:
+      return ApplyCwiseBinaryMap<Foeq>(lhs, rhs);
+    case arith::CmpIPredicate::ne:
+      return ApplyCwiseBinaryMap<Fone>(lhs, rhs);
+    case arith::CmpIPredicate::slt:
+      return ApplyCwiseBinaryMap<Folt>(lhs, rhs);
+    case arith::CmpIPredicate::sle:
+      return ApplyCwiseBinaryMap<Fole>(lhs, rhs);
+    case arith::CmpIPredicate::sgt:
+      return ApplyCwiseBinaryMap<Fogt>(lhs, rhs);
+    case arith::CmpIPredicate::sge:
+      return ApplyCwiseBinaryMap<Foge>(lhs, rhs);
+    case arith::CmpIPredicate::ult:
+      return ApplyCwiseBinaryMap<Iult>(lhs, rhs);
+    case arith::CmpIPredicate::ule:
+      return ApplyCwiseBinaryMap<Iule>(lhs, rhs);
+    case arith::CmpIPredicate::ugt:
+      return ApplyCwiseBinaryMap<Iugt>(lhs, rhs);
+    case arith::CmpIPredicate::uge:
+      return ApplyCwiseBinaryMap<Iuge>(lhs, rhs);
+  }
+}
+
+template <bool value>
+struct ConstFunctor : CwiseAll {
+  template <typename T>
+  static bool Apply(T, T) {
+    return value;
+  }
+};
+
+InterpreterValue CmpF(InterpreterState&, arith::CmpFOp compare,
+                      const InterpreterValue& lhs,
+                      const InterpreterValue& rhs) {
+  switch (compare.getPredicate()) {
+    case arith::CmpFPredicate::AlwaysFalse:
+      return ApplyCwiseBinaryMap<ConstFunctor<false>>(lhs, rhs);
+    case arith::CmpFPredicate::OEQ:
+      return ApplyCwiseBinaryMap<Foeq>(lhs, rhs);
+    case arith::CmpFPredicate::OGT:
+      return ApplyCwiseBinaryMap<Fogt>(lhs, rhs);
+    case arith::CmpFPredicate::OGE:
+      return ApplyCwiseBinaryMap<Foge>(lhs, rhs);
+    case arith::CmpFPredicate::OLT:
+      return ApplyCwiseBinaryMap<Folt>(lhs, rhs);
+    case arith::CmpFPredicate::OLE:
+      return ApplyCwiseBinaryMap<Fole>(lhs, rhs);
+    case arith::CmpFPredicate::ONE:
+      return ApplyCwiseBinaryMap<Fone>(lhs, rhs);
+    case arith::CmpFPredicate::ORD:
+      return ApplyCwiseBinaryMap<Ford>(lhs, rhs);
+    case arith::CmpFPredicate::UEQ:
+      return ApplyCwiseBinaryMap<Fueq>(lhs, rhs);
+    case arith::CmpFPredicate::UGT:
+      return ApplyCwiseBinaryMap<Fugt>(lhs, rhs);
+    case arith::CmpFPredicate::UGE:
+      return ApplyCwiseBinaryMap<Fuge>(lhs, rhs);
+    case arith::CmpFPredicate::ULT:
+      return ApplyCwiseBinaryMap<Fult>(lhs, rhs);
+    case arith::CmpFPredicate::ULE:
+      return ApplyCwiseBinaryMap<Fule>(lhs, rhs);
+    case arith::CmpFPredicate::UNE:
+      return ApplyCwiseBinaryMap<Fune>(lhs, rhs);
+    case arith::CmpFPredicate::UNO:
+      return ApplyCwiseBinaryMap<Funo>(lhs, rhs);
+    case arith::CmpFPredicate::AlwaysTrue:
+      return ApplyCwiseBinaryMap<ConstFunctor<true>>(lhs, rhs);
+  }
+}
+
+InterpreterValue Select(InterpreterState& state, arith::SelectOp,
+                        const InterpreterValue& cond,
+                        const InterpreterValue& trueValue,
+                        const InterpreterValue& falseValue) {
+  if (std::holds_alternative<bool>(cond.storage)) {
+    return std::get<bool>(cond.storage) ? trueValue : falseValue;
+  }
+
+  if (!cond.IsTensor() && !cond.View().is_vector) {
+    state.AddFailure("select requires a scalar or vector argument");
+    return {};
+  }
+
+  auto ret = trueValue.Clone();
+  for (const auto& index : cond.View().Indices(/*include_vector_dims=*/true)) {
+    if (cond.ExtractElement(index).AsInt() == 0) {
+      ret.InsertElement(index, falseValue.ExtractElement(index));
+    }
+  }
+  return ret;
+}
+
+template <typename R>
+struct ExtFFunctor : CwiseFloat {
+  template <typename A>
+  static R Apply(A v) {
+    return v;
+  }
+};
+
+InterpreterValue ExtF(InterpreterState&, arith::ExtFOp op,
+                      const InterpreterValue& in) {
+  return DispatchScalarType(
+      op->getResultTypes()[0], [&](auto dummy) -> InterpreterValue {
+        return ApplyCwiseMap<ExtFFunctor<decltype(dummy)>>(in);
+      });
+}
+
+REGISTER_MLIR_INTERPRETER_OP("arith.addf", ApplyCwiseBinaryMap<Plus>);
+REGISTER_MLIR_INTERPRETER_OP("arith.andi", ApplyCwiseBinaryMap<BitAnd>);
+REGISTER_MLIR_INTERPRETER_OP("arith.divf", ApplyCwiseBinaryMap<Divide>);
+REGISTER_MLIR_INTERPRETER_OP("arith.extui", UiToFP);
+REGISTER_MLIR_INTERPRETER_OP("arith.maxf", ApplyCwiseBinaryMap<Max>);
+REGISTER_MLIR_INTERPRETER_OP("arith.minf", ApplyCwiseBinaryMap<Min>);
+REGISTER_MLIR_INTERPRETER_OP("arith.mulf", ApplyCwiseBinaryMap<Multiply>);
+REGISTER_MLIR_INTERPRETER_OP("arith.negf", ApplyCwiseMap<Neg>);
+REGISTER_MLIR_INTERPRETER_OP("arith.ori", ApplyCwiseBinaryMap<BitOr>);
+REGISTER_MLIR_INTERPRETER_OP("arith.remf", ApplyCwiseBinaryMap<Remainder>);
+REGISTER_MLIR_INTERPRETER_OP("arith.subf", ApplyCwiseBinaryMap<Minus>);
+REGISTER_MLIR_INTERPRETER_OP("arith.uitofp", UiToFP);
+REGISTER_MLIR_INTERPRETER_OP("arith.xori", ApplyCwiseBinaryMap<BitXor>);
+REGISTER_MLIR_INTERPRETER_OP("arith.shrui",
+                             ApplyCwiseBinaryMap<ShiftRightLogical>);
+REGISTER_MLIR_INTERPRETER_OP("arith.shrsi",
+                             ApplyCwiseBinaryMap<ShiftRightArith>);
+REGISTER_MLIR_INTERPRETER_OP("arith.shli", ApplyCwiseBinaryMap<ShiftLeft>);
+
+// The float implementations support ints too.
+REGISTER_MLIR_INTERPRETER_OP("arith.addi", "arith.addf");
+REGISTER_MLIR_INTERPRETER_OP("arith.divsi", "arith.divf");
+REGISTER_MLIR_INTERPRETER_OP("arith.maxsi", "arith.maxf");
+REGISTER_MLIR_INTERPRETER_OP("arith.minsi", "arith.minf");
+REGISTER_MLIR_INTERPRETER_OP("arith.muli", "arith.mulf");
+REGISTER_MLIR_INTERPRETER_OP("arith.remsi", "arith.remf");
+REGISTER_MLIR_INTERPRETER_OP("arith.subi", "arith.subf");
+
+REGISTER_MLIR_INTERPRETER_OP(Bitcast);
+REGISTER_MLIR_INTERPRETER_OP(CmpF);
+REGISTER_MLIR_INTERPRETER_OP(CmpI);
+REGISTER_MLIR_INTERPRETER_OP(Constant);
+REGISTER_MLIR_INTERPRETER_OP(ExtF);
+REGISTER_MLIR_INTERPRETER_OP(FloatCast<arith::FPToSIOp>);
+REGISTER_MLIR_INTERPRETER_OP(IntCast<arith::ExtSIOp>);
+REGISTER_MLIR_INTERPRETER_OP(IntCast<arith::IndexCastOp>);
+REGISTER_MLIR_INTERPRETER_OP(IntCast<arith::SIToFPOp>);
+REGISTER_MLIR_INTERPRETER_OP(IntCast<arith::TruncIOp>);
+REGISTER_MLIR_INTERPRETER_OP(Select);
+
+}  // namespace
+}  // namespace interpreter
+}  // namespace mlir
diff --git a/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/bufferization.cc b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/bufferization.cc
new file mode 100644
index 00000000000000..1561d54b4b31eb
--- /dev/null
+++ b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/bufferization.cc
@@ -0,0 +1,70 @@
+/* Copyright 2022 The OpenXLA Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+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.
+==============================================================================*/
+
+#include "mlir/Dialect/Bufferization/IR/Bufferization.h"  // from @llvm-project
+
+#include <algorithm>  // NOLINT
+#include <optional>   // NOLINT
+
+#include "xla/mlir/tools/mlir_interpreter/dialects/util.h"
+#include "xla/mlir/tools/mlir_interpreter/framework/interpreter.h"
+#include "xla/mlir/tools/mlir_interpreter/framework/interpreter_value.h"
+#include "xla/mlir/tools/mlir_interpreter/framework/registration.h"
+
+namespace mlir {
+namespace interpreter {
+namespace {
+
+InterpreterValue ToTensor(InterpreterState&, bufferization::ToTensorOp,
+                          const InterpreterValue& in) {
+  return in.Clone();
+}
+
+InterpreterValue ToMemref(InterpreterState&, bufferization::ToMemrefOp,
+                          const InterpreterValue& in) {
+  return in;
+}
+
+InterpreterValue AllocTensor(
+    InterpreterState&, bufferization::AllocTensorOp alloc,
+    ArrayRef<int64_t> dynamic_sizes, std::optional<InterpreterValue> copy,
+    const std::optional<InterpreterValue>& /*sizeHint*/) {
+  auto ty = alloc->getResultTypes().front().cast<mlir::ShapedType>();
+  auto shape = ReplaceDynamicVals(ty.getShape(), dynamic_sizes);
+
+  if (copy) {
+    return copy->Clone();
+  }
+  return InterpreterValue::MakeTensor(ty.getElementType(), shape);
+}
+
+InterpreterValue Clone(InterpreterState& state, bufferization::CloneOp,
+                       const InterpreterValue& in) {
+  if (auto* stats = state.GetOptions().stats) {
+    stats->heap_size += in.Buffer()->GetByteSize();
+    stats->peak_heap_size = std::max(stats->peak_heap_size, stats->heap_size);
+    ++stats->num_allocations;
+  }
+  return in.Clone();
+}
+
+REGISTER_MLIR_INTERPRETER_OP(AllocTensor);
+REGISTER_MLIR_INTERPRETER_OP(Clone);
+REGISTER_MLIR_INTERPRETER_OP(ToMemref);
+REGISTER_MLIR_INTERPRETER_OP(ToTensor);
+
+}  // namespace
+}  // namespace interpreter
+}  // namespace mlir
diff --git a/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/builtin.cc b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/builtin.cc
new file mode 100644
index 00000000000000..48b8008dc07cfc
--- /dev/null
+++ b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/builtin.cc
@@ -0,0 +1,64 @@
+/* Copyright 2022 The OpenXLA Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+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.
+==============================================================================*/
+
+#include "llvm/ADT/SmallVector.h"
+#include "llvm/Support/Casting.h"
+#include "llvm/Support/ErrorHandling.h"
+#include "llvm/Support/raw_ostream.h"
+#include "mlir/IR/BuiltinTypeInterfaces.h"  // from @llvm-project
+#include "mlir/IR/Operation.h"  // from @llvm-project
+#include "mlir/IR/TypeUtilities.h"  // from @llvm-project
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+#include "xla/mlir/tools/mlir_interpreter/framework/interpreter.h"
+#include "xla/mlir/tools/mlir_interpreter/framework/interpreter_value.h"
+#include "xla/mlir/tools/mlir_interpreter/framework/registration.h"
+#include "xla/mlir/tools/mlir_interpreter/framework/tensor_or_memref.h"
+
+namespace mlir {
+namespace interpreter {
+namespace {
+
+llvm::SmallVector<InterpreterValue> UnrealizedConversionCast(
+    MutableArrayRef<InterpreterValue> args, mlir::Operation* op,
+    InterpreterState&) {
+  auto result_ty = op->getResultTypes()[0];
+  auto operand_ty = op->getOperandTypes()[0];
+  if (result_ty == operand_ty) {
+    return {args[0]};
+  }
+
+  if (auto r = llvm::dyn_cast<ShapedType>(result_ty)) {
+    if (auto o = llvm::dyn_cast<ShapedType>(operand_ty)) {
+      if (verifyCompatibleShapes({o, r}).succeeded()) {
+        return {DispatchScalarType(r, [&](auto dummy) -> InterpreterValue {
+          TensorOrMemref<decltype(dummy)> result;
+          result.view = args[0].View();
+          result.buffer = args[0].Buffer();
+          return {result};
+        })};
+      }
+    }
+  }
+
+  llvm::errs() << "Unimplemented cast: " << *op << "\n";
+  llvm_unreachable("unimplemented cast");
+}
+
+REGISTER_MLIR_INTERPRETER_OP("builtin.unrealized_conversion_cast",
+                             UnrealizedConversionCast);
+
+}  // namespace
+}  // namespace interpreter
+}  // namespace mlir
diff --git a/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/comparators.h b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/comparators.h
new file mode 100644
index 00000000000000..c26b75af6ae660
--- /dev/null
+++ b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/comparators.h
@@ -0,0 +1,105 @@
+/* Copyright 2022 The OpenXLA Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+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.
+==============================================================================*/
+
+#ifndef XLA_MLIR_TOOLS_MLIR_INTERPRETER_DIALECTS_COMPARATORS_H_
+#define XLA_MLIR_TOOLS_MLIR_INTERPRETER_DIALECTS_COMPARATORS_H_
+
+#include <complex>
+#include <cstdint>
+#include <type_traits>
+
+#include "llvm/Support/ErrorHandling.h"
+#include "xla/mlir/tools/mlir_interpreter/framework/interpreter_value_util.h"
+
+namespace mlir {
+namespace interpreter {
+
+// Despite the name, this works on integers and complex too.
+template <int64_t v, bool r, bool nan_result>
+struct FloatCompare : CwiseAll {
+  template <typename T>
+  static bool Apply(T a, T b) {
+    if (isnan(a) || isnan(b)) return nan_result;
+    if constexpr (v == 0) {
+      // For complex eq/ne.
+      return (a == b) == r;
+    } else if constexpr (std::is_floating_point_v<T> || std::is_integral_v<T>) {
+      auto cmp = a > b ? 1 : (a < b ? -1 : 0);
+      return (cmp == v) == r;
+    } else {
+      llvm_unreachable("operation not supported for this type");
+    }
+  }
+
+  template <typename T>
+  static bool isnan(T a) {
+    return std::isnan(a);
+  }
+  template <typename T>
+  static bool isnan(std::complex<T> a) {
+    return std::isnan(std::real(a)) || std::isnan(std::imag(a));
+  }
+};
+
+using Foeq = FloatCompare<0, true, false>;
+using Foge = FloatCompare<-1, false, false>;
+using Fogt = FloatCompare<1, true, false>;
+using Fole = FloatCompare<1, false, false>;
+using Folt = FloatCompare<-1, true, false>;
+using Fone = FloatCompare<0, false, false>;
+using Ford = FloatCompare<99, false, false>;
+using Fueq = FloatCompare<0, true, true>;
+using Fuge = FloatCompare<-1, false, true>;
+using Fugt = FloatCompare<1, true, true>;
+using Fule = FloatCompare<1, false, true>;
+using Fult = FloatCompare<-1, true, true>;
+using Fune = FloatCompare<0, false, true>;
+using Funo = FloatCompare<99, true, true>;
+
+template <int64_t v, bool r>
+struct UnsignedCompare : CwiseInt {
+  template <typename T>
+  static bool Apply(T a, T b) {
+    using U = std::make_unsigned_t<T>;
+    auto a_u = static_cast<U>(a);
+    auto b_u = static_cast<U>(b);
+    auto cmp = a_u > b_u ? 1 : (a_u < b_u ? -1 : 0);
+    return (cmp == v) == r;
+  }
+};
+
+using Iuge = UnsignedCompare<-1, false>;
+using Iule = UnsignedCompare<1, false>;
+using Iugt = UnsignedCompare<1, true>;
+using Iult = UnsignedCompare<-1, true>;
+
+struct Iumax {
+  template <typename T>
+  static T apply(T a, T b) {
+    return Iuge::Apply(a, b) ? a : b;
+  }
+};
+
+struct Iumin {
+  template <typename T>
+  static T apply(T a, T b) {
+    return Iule::Apply(a, b) ? a : b;
+  }
+};
+
+}  // namespace interpreter
+}  // namespace mlir
+
+#endif  // XLA_MLIR_TOOLS_MLIR_INTERPRETER_DIALECTS_COMPARATORS_H_
diff --git a/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/complex.cc b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/complex.cc
new file mode 100644
index 00000000000000..35ab806eef8ca3
--- /dev/null
+++ b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/complex.cc
@@ -0,0 +1,68 @@
+/* Copyright 2023 The OpenXLA Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+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.
+==============================================================================*/
+
+#include "mlir/Dialect/Complex/IR/Complex.h"  // from @llvm-project
+
+#include "llvm/ADT/SmallVector.h"
+#include "llvm/Support/ErrorHandling.h"
+#include "mlir/IR/BuiltinAttributes.h"  // from @llvm-project
+#include "xla/mlir/tools/mlir_interpreter/dialects/cwise_math.h"
+#include "xla/mlir/tools/mlir_interpreter/framework/interpreter.h"
+#include "xla/mlir/tools/mlir_interpreter/framework/interpreter_value.h"
+#include "xla/mlir/tools/mlir_interpreter/framework/interpreter_value_util.h"
+#include "xla/mlir/tools/mlir_interpreter/framework/registration.h"
+
+namespace mlir {
+namespace interpreter {
+namespace {
+
+InterpreterValue Constant(InterpreterState&, complex::ConstantOp constant) {
+  auto ty = constant->getResultTypes()[0];
+  return DispatchScalarType(ty, [&](auto dummy) -> InterpreterValue {
+    if constexpr (is_complex_v<decltype(dummy)>) {
+      using T = typename decltype(dummy)::value_type;
+      auto values =
+          llvm::to_vector(constant.getValue().getAsValueRange<FloatAttr>());
+      return {decltype(dummy){static_cast<T>(values[0].convertToDouble()),
+                              static_cast<T>(values[1].convertToDouble())}};
+    } else {
+      llvm_unreachable("invalid constant");
+    }
+  });
+}
+
+REGISTER_MLIR_INTERPRETER_OP("complex.abs", "math.absf");
+REGISTER_MLIR_INTERPRETER_OP("complex.add", "arith.addf");
+REGISTER_MLIR_INTERPRETER_OP("complex.cos", ApplyCwiseMap<Cos>);
+REGISTER_MLIR_INTERPRETER_OP("complex.create", ApplyCwiseBinaryMap<Complex>);
+REGISTER_MLIR_INTERPRETER_OP("complex.div", ApplyCwiseBinaryMap<Divide>);
+REGISTER_MLIR_INTERPRETER_OP("complex.exp", ApplyCwiseMap<Exp>);
+REGISTER_MLIR_INTERPRETER_OP("complex.expm1", ApplyCwiseMap<ExpM1>);
+REGISTER_MLIR_INTERPRETER_OP("complex.im", ApplyCwiseMap<Imag>);
+REGISTER_MLIR_INTERPRETER_OP("complex.log", ApplyCwiseMap<Log>);
+REGISTER_MLIR_INTERPRETER_OP("complex.log1p", ApplyCwiseMap<Log1P>);
+REGISTER_MLIR_INTERPRETER_OP("complex.mul", ApplyCwiseBinaryMap<Multiply>);
+REGISTER_MLIR_INTERPRETER_OP("complex.neg", ApplyCwiseMap<Neg>);
+REGISTER_MLIR_INTERPRETER_OP("complex.pow", ApplyCwiseBinaryMap<Power>);
+REGISTER_MLIR_INTERPRETER_OP("complex.re", ApplyCwiseMap<Real>);
+REGISTER_MLIR_INTERPRETER_OP("complex.rsqrt", ApplyCwiseMap<RSqrt>);
+REGISTER_MLIR_INTERPRETER_OP("complex.sin", ApplyCwiseMap<Sin>);
+REGISTER_MLIR_INTERPRETER_OP("complex.sqrt", ApplyCwiseMap<Sqrt>);
+REGISTER_MLIR_INTERPRETER_OP("complex.tanh", ApplyCwiseMap<TanH>);
+REGISTER_MLIR_INTERPRETER_OP(Constant);
+
+}  // namespace
+}  // namespace interpreter
+}  // namespace mlir
diff --git a/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/cwise_math.h b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/cwise_math.h
new file mode 100644
index 00000000000000..6f36f8a65a8142
--- /dev/null
+++ b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/cwise_math.h
@@ -0,0 +1,242 @@
+/* Copyright 2023 The OpenXLA Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+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.
+==============================================================================*/
+
+#ifndef XLA_MLIR_TOOLS_MLIR_INTERPRETER_DIALECTS_CWISE_MATH_H_
+#define XLA_MLIR_TOOLS_MLIR_INTERPRETER_DIALECTS_CWISE_MATH_H_
+
+#include <climits>
+#include <complex>
+#include <cstdint>
+#include <functional>
+#include <type_traits>
+
+#include "xla/mlir/tools/mlir_interpreter/framework/interpreter_value_util.h"
+
+namespace mlir {
+namespace interpreter {
+
+struct ATan2 : CwiseReal {
+  template <typename T>
+  static T Apply(T a, T b) {
+    return std::atan2(a, b);
+  }
+};
+
+struct Clz : CwiseInt {
+  template <typename T>
+  static T Apply(T a) {
+    if (!a) {
+      // Return something well-defined for zeroes.
+      return sizeof(T{}) * CHAR_BIT;
+    }
+    return __builtin_clzl(
+               static_cast<uint64_t>(static_cast<std::make_unsigned_t<T>>(a))) -
+           (sizeof(uint64_t) - sizeof(T{})) * CHAR_BIT;
+  }
+};
+
+struct Ctz : CwiseInt {
+  template <typename T>
+  static T Apply(T a) {
+    if (!a) {
+      // Return something well-defined for zeroes.
+      return sizeof(T{}) * CHAR_BIT;
+    }
+    return __builtin_ctzl(static_cast<uint64_t>(a));
+  }
+};
+
+struct Complex : CwiseFloat {
+  template <typename T>
+  static std::complex<T> Apply(T a, T b) {
+    return {a, b};
+  }
+};
+
+struct Max : CwiseReal {
+  template <typename T>
+  static T Apply(T a, T b) {
+    return std::max(a, b);
+  }
+};
+
+struct Min : CwiseReal {
+  template <typename T>
+  static T Apply(T a, T b) {
+    return std::min(a, b);
+  }
+};
+
+struct Power : CwiseArith {
+  template <typename T>
+  static T Apply(T a, T b) {
+    if constexpr (std::is_integral_v<T>) {
+      if constexpr (std::is_signed_v<T>) {
+        if (b < 0) {
+          return a == 1 ? 1 : 0;
+        }
+      }
+      T result = 1;
+      while (b > 0) {
+        if (b & 1) result *= a;
+        b >>= 1;
+        if (b) {
+          a *= a;
+        }
+      }
+      return result;
+    } else {
+      return std::pow(a, b);
+    }
+  }
+};
+
+struct Remainder : CwiseReal {
+  template <typename T>
+  static T Apply(T a, T b) {
+    if constexpr (std::is_integral_v<T>) {
+      return a % b;
+    } else {
+      return std::fmod(a, b);
+    }
+  }
+};
+
+struct ShiftRightArith : CwiseInt {
+  template <typename T>
+  static T Apply(T a, T b) {
+    return b >= sizeof(T) * CHAR_BIT ? 0 : (a >> b);
+  }
+};
+
+struct ShiftRightLogical : CwiseInt {
+  template <typename T>
+  static T Apply(T a, T b) {
+    return b >= sizeof(T) * CHAR_BIT
+               ? 0
+               : static_cast<std::make_unsigned_t<T>>(a) >> b;
+  }
+};
+
+struct ShiftLeft : CwiseInt {
+  template <typename T>
+  static T Apply(T a, T b) {
+    return b >= sizeof(T) * CHAR_BIT ? 0 : (a << b);
+  }
+};
+
+namespace detail {
+template <template <typename T> class F, typename trait>
+struct Wrap : trait {
+  template <typename T>
+  static T Apply(T a, T b) {
+    return F<T>{}(a, b);
+  }
+};
+}  // namespace detail
+
+using Plus = detail::Wrap<std::plus, CwiseArith>;
+using Divide = detail::Wrap<std::divides, CwiseArith>;
+using Multiply = detail::Wrap<std::multiplies, CwiseArith>;
+using Minus = detail::Wrap<std::minus, CwiseAll>;
+using BitAnd = detail::Wrap<std::bit_and, CwiseIntegral>;
+using BitOr = detail::Wrap<std::bit_or, CwiseIntegral>;
+using BitXor = detail::Wrap<std::bit_xor, CwiseIntegral>;
+
+struct RSqrt : CwiseNonIntegral {
+  template <typename T>
+  static T Apply(T a) {
+    return static_cast<T>(T{1} / std::sqrt(a));
+  }
+};
+
+#define DEFINE_WRAPPER(name, std_fun, trait) \
+  struct name : trait {                      \
+    template <typename T>                    \
+    static auto Apply(T a) {                 \
+      return std_fun(a);                     \
+    }                                        \
+  };
+
+DEFINE_WRAPPER(ATan, std::atan, CwiseNonIntegral);
+DEFINE_WRAPPER(Abs, std::abs, CwiseSignedOrComplex);
+DEFINE_WRAPPER(Cbrt, std::cbrt, CwiseFloat);
+DEFINE_WRAPPER(Ceil, std::ceil, CwiseFloat);
+DEFINE_WRAPPER(Cos, std::cos, CwiseNonIntegral);
+DEFINE_WRAPPER(Erf, std::erf, CwiseFloat);
+DEFINE_WRAPPER(Exp, std::exp, CwiseNonIntegral);
+DEFINE_WRAPPER(Exp2, std::exp2, CwiseFloat);
+DEFINE_WRAPPER(Floor, std::floor, CwiseFloat);
+DEFINE_WRAPPER(Imag, std::imag, CwiseComplex);
+DEFINE_WRAPPER(IsFinite, std::isfinite, CwiseFloat);
+DEFINE_WRAPPER(Log, std::log, CwiseNonIntegral);
+DEFINE_WRAPPER(Log10, std::log10, CwiseNonIntegral);
+DEFINE_WRAPPER(Log2, std::log2, CwiseFloat);
+DEFINE_WRAPPER(NearbyInt, std::nearbyint, CwiseFloat);
+DEFINE_WRAPPER(Neg, std::negate<T>{}, CwiseSignedOrComplex);
+DEFINE_WRAPPER(Real, std::real, CwiseComplex);
+DEFINE_WRAPPER(Round, std::round, CwiseFloat);
+DEFINE_WRAPPER(Sin, std::sin, CwiseNonIntegral);
+DEFINE_WRAPPER(Sqrt, std::sqrt, CwiseNonIntegral);
+DEFINE_WRAPPER(Tan, std::tan, CwiseNonIntegral);
+DEFINE_WRAPPER(TanH, std::tanh, CwiseNonIntegral);
+DEFINE_WRAPPER(Trunc, std::trunc, CwiseFloat);
+
+#undef DEFINE_WRAPPER
+
+struct ExpM1 : CwiseNonIntegral {
+  template <typename T>
+  static T Apply(T a) {
+    if constexpr (std::is_floating_point_v<T>) {
+      return std::expm1(a);
+    } else {
+      auto r = std::real(a);
+      auto i = std::imag(a);
+      auto s = std::sin(i / 2);
+      auto real = std::expm1(r) * std::cos(i) - 2 * s * s;
+      auto imag = std::exp(r) * std::sin(i);
+      return {real, imag};
+    }
+  }
+};
+
+struct Log1P : CwiseNonIntegral {
+  template <typename T>
+  static T Apply(T a) {
+    if constexpr (std::is_floating_point_v<T>) {
+      return std::log1p(a);
+    } else {
+      auto r = std::real(a);
+      auto i = std::imag(a);
+      auto l = std::hypot(r + 1, i);
+      auto real = std::log(l);
+      auto imag = std::atan2(i, r + 1);
+      return {real, imag};
+    }
+  }
+};
+
+struct PopCount : CwiseInt {
+  template <typename T>
+  static T Apply(T a) {
+    return __builtin_popcountl(
+        static_cast<uint64_t>(static_cast<std::make_unsigned_t<T>>(a)));
+  }
+};
+
+}  // namespace interpreter
+}  // namespace mlir
+
+#endif  // XLA_MLIR_TOOLS_MLIR_INTERPRETER_DIALECTS_CWISE_MATH_H_
diff --git a/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/func.cc b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/func.cc
new file mode 100644
index 00000000000000..b20d8dfe3607e4
--- /dev/null
+++ b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/func.cc
@@ -0,0 +1,122 @@
+/* Copyright 2023 The OpenXLA Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+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.
+==============================================================================*/
+
+#include <dlfcn.h>
+
+#include <tuple>
+#include <type_traits>
+
+#include "llvm/Support/Casting.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
+#include "mlir/IR/Operation.h"  // from @llvm-project
+#include "mlir/IR/Types.h"  // from @llvm-project
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+#include "xla/mlir/tools/mlir_interpreter/dialects/util.h"
+#include "xla/mlir/tools/mlir_interpreter/framework/interpreter.h"
+#include "xla/mlir/tools/mlir_interpreter/framework/interpreter_value.h"
+#include "xla/mlir/tools/mlir_interpreter/framework/registration.h"
+
+namespace mlir {
+namespace interpreter {
+namespace {
+
+template <typename T>
+bool TypeMatches(mlir::Type type) {
+  if constexpr (std::is_same_v<T, float>) {
+    return type.isF32();
+  } else if constexpr (std::is_same_v<T, double>) {
+    return type.isF64();
+  } else {
+    return false;
+  }
+}
+
+template <typename Dummy>
+bool TypesMatch(ArrayRef<mlir::Type> types) {
+  return types.empty();
+}
+
+template <typename Dummy, typename T, typename... R>
+bool TypesMatch(ArrayRef<mlir::Type> types) {
+  if (types.empty() || !TypeMatches<T>(types.front())) return false;
+  return TypesMatch<Dummy, R...>(types.drop_front());
+}
+
+template <int n, typename... Args>
+using Arg = std::tuple_element_t<n, std::tuple<Args...>>;
+
+template <typename Ret, typename... Args>
+bool TryCall(void* sym, func::FuncOp callee,
+             MutableArrayRef<InterpreterValue> args, InterpreterValue& ret) {
+  if (args.size() != callee.getNumArguments() || callee.getNumResults() != 1) {
+    return false;
+  }
+
+  if (!TypeMatches<Ret>(callee.getResultTypes()[0])) {
+    return false;
+  }
+  if (!TypesMatch<void, Args...>(callee.getArgumentTypes())) {
+    return false;
+  }
+
+  static_assert(sizeof...(Args) <= 2);
+  using FnType = Ret (*)(Args...);
+  auto fn = reinterpret_cast<FnType>(sym);
+  constexpr int n = sizeof...(Args);
+
+  if constexpr (n == 1) {
+    ret = {fn(std::get<Arg<0, Args...>>(args[0].storage))};
+  } else {
+    static_assert(n == 2);
+    ret = {fn(std::get<Arg<0, Args...>>(args[0].storage),
+              std::get<Arg<1, Args...>>(args[1].storage))};
+  }
+  return true;
+}
+
+llvm::SmallVector<InterpreterValue> Call(MutableArrayRef<InterpreterValue> args,
+                                         mlir::Operation* op,
+                                         InterpreterState& state) {
+  auto call = llvm::cast<func::CallOp>(op);
+  auto callee =
+      llvm::cast<func::FuncOp>(state.GetSymbols().lookup(call.getCallee()));
+  if (callee->getRegion(0).hasOneBlock()) {
+    return Interpret(state, callee.getRegion(), args);
+  }
+
+  void* sym = dlsym(RTLD_DEFAULT, callee.getSymName().str().c_str());
+  if (sym == nullptr) {
+    state.AddFailure("callee not found");
+    return {};
+  }
+
+  InterpreterValue result;
+  if (TryCall<float, float>(sym, callee, args, result) ||
+      TryCall<float, float, float>(sym, callee, args, result) ||
+      TryCall<double, double>(sym, callee, args, result) ||
+      TryCall<double, double, double>(sym, callee, args, result)) {
+    return {result};
+  }
+
+  state.AddFailure("unsupported call target");
+  return {};
+}
+
+REGISTER_MLIR_INTERPRETER_OP("func.call", Call);
+REGISTER_MLIR_INTERPRETER_OP("func.return", NoOpTerminator);
+
+}  // namespace
+}  // namespace interpreter
+}  // namespace mlir
diff --git a/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/linalg.cc b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/linalg.cc
new file mode 100644
index 00000000000000..f65cbaf4ef18de
--- /dev/null
+++ b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/linalg.cc
@@ -0,0 +1,338 @@
+/* Copyright 2022 The OpenXLA Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+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.
+==============================================================================*/
+
+#include "mlir/Dialect/Linalg/IR/Linalg.h"  // from @llvm-project
+
+// clang-format erroneously puts the Linalg header above.
+#include <algorithm>   // NOLINT
+#include <cstdint>     // NOLINT
+#include <functional>  // NOLINT
+#include <memory>      // NOLINT
+#include <utility>     // NOLINT
+
+#include "llvm/ADT/ArrayRef.h"
+#include "llvm/ADT/STLExtras.h"
+#include "llvm/ADT/SmallVector.h"
+#include "mlir/IR/BuiltinTypes.h"  // from @llvm-project
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+#include "xla/mlir/tools/mlir_interpreter/dialects/util.h"
+#include "xla/mlir/tools/mlir_interpreter/framework/interpreter.h"
+#include "xla/mlir/tools/mlir_interpreter/framework/interpreter_value.h"
+#include "xla/mlir/tools/mlir_interpreter/framework/registration.h"
+#include "xla/mlir/tools/mlir_interpreter/framework/tensor_or_memref.h"
+
+namespace mlir {
+namespace interpreter {
+namespace {
+
+class IterationIndexSideChannel : public InterpreterSideChannel {
+ public:
+  explicit IterationIndexSideChannel(ArrayRef<int64_t> indices)
+      : indices_(indices) {}
+  ArrayRef<int64_t> GetIndices() const { return indices_; }
+
+ private:
+  ArrayRef<int64_t> indices_;
+};
+
+llvm::SmallVector<InterpreterValue> Broadcast(InterpreterState&,
+                                              linalg::BroadcastOp broadcast,
+                                              const InterpreterValue& input,
+                                              InterpreterValue init) {
+  auto broadcast_dims = llvm::to_vector(broadcast.getDimensions());
+  llvm::sort(broadcast_dims);
+
+  auto out_shape = init.View().sizes;
+  auto out = input;
+  auto& out_view = out.View();
+  for (int64_t dim : broadcast_dims) {
+    out_view.sizes.insert(out_view.sizes.begin() + dim, out_shape[dim]);
+    out_view.strides.insert(out_view.strides.begin() + dim, 0);
+  }
+
+  if (broadcast.getNumResults() == 1) {
+    return {std::move(out)};
+  }
+
+  init.Fill([&](llvm::ArrayRef<int64_t> indices) {
+    return out.ExtractElement(indices);
+  });
+  return {};
+}
+
+llvm::SmallVector<InterpreterValue> Generic(
+    InterpreterState& state, linalg::GenericOp generic,
+    MutableArrayRef<InterpreterValue> inputs,
+    MutableArrayRef<InterpreterValue> outputs_ref) {
+  SmallVector<int64_t> shapes;
+  for (auto& value : llvm::concat<InterpreterValue>(inputs, outputs_ref)) {
+    if (value.IsTensor() /* (or memref) */) {
+      llvm::append_range(
+          shapes, ArrayRef<int64_t>(value.View().sizes)
+                      .drop_back(value.View().num_vector_dims.value_or(0)));
+    }
+  }
+  auto ranges = generic.getShapesToLoopsMap().compose(shapes);
+
+  llvm::SmallVector<InterpreterValue> outputs;
+  for (int64_t output = 0; output < outputs_ref.size(); ++output) {
+    outputs.push_back(
+        GetInitOperand(generic.getOutputs(), output, outputs_ref));
+  }
+
+  llvm::SmallVector<int64_t> ivs(ranges.size());
+  InterpreterScope scope(state);
+  scope.SetSideChannel(std::make_shared<IterationIndexSideChannel>(ivs));
+
+  auto indexing_maps = generic.getIndexingMapsArray();
+  auto output_maps =
+      ArrayRef<AffineMap>(indexing_maps).drop_front(inputs.size());
+  std::function<void(int64_t)> run;
+  run = [&](int64_t loopIndex) {
+    // Abort recursion if we encountered some error previously.
+    if (state.HasFailure()) {
+      return;
+    }
+
+    if (loopIndex < ranges.size()) {
+      for (int64_t index = 0; index < ranges[loopIndex]; ++index) {
+        ivs[loopIndex] = index;
+        run(loopIndex + 1);
+      }
+    } else {
+      llvm::SmallVector<InterpreterValue> bbargs;
+      // Build bbargs: 1. inputs, 2. outputs.
+      for (auto [input, map] : llvm::zip(inputs, indexing_maps)) {
+        auto indices = EvalAffineMap(map, ivs);
+        bbargs.push_back(input.ExtractElement(indices));
+      }
+      llvm::SmallVector<llvm::SmallVector<int64_t>> outputIndices;
+      for (auto [output, map] : llvm::zip(outputs, output_maps)) {
+        auto& indices = outputIndices.emplace_back(EvalAffineMap(map, ivs));
+        bbargs.push_back(output.ExtractElement(indices));
+      }
+      // Evaluate region.
+      auto yielded = Interpret(state, generic.getRegion(), bbargs);
+      if (state.HasFailure()) {
+        return;
+      }
+      // Insert yielded values in the outputs.
+      for (auto [output, indices, yield] :
+           llvm::zip(outputs, outputIndices, yielded)) {
+        output.InsertElement(indices, yield);
+      }
+    }
+  };
+  run(0);
+
+  if (generic.getNumResults() == 0) return {};
+  return outputs;
+}
+
+llvm::SmallVector<InterpreterValue> Map(InterpreterState& state,
+                                        linalg::MapOp op,
+                                        ArrayRef<InterpreterValue> inputs,
+                                        const InterpreterValue& init) {
+  InterpreterValue output =
+      isa<TensorType>(op.getInit().getType()) ? init.Clone() : init;
+
+  InterpreterScope scope(state);
+  SmallVector<int64_t> ivs(output.View().Rank());
+  scope.SetSideChannel(std::make_shared<IterationIndexSideChannel>(ivs));
+  for (const auto& indices : output.View().Indices()) {
+    std::copy(indices.begin(), indices.end(), ivs.begin());
+    llvm::SmallVector<InterpreterValue> args;
+    for (auto& input : inputs) {
+      args.push_back(input.ExtractElement(indices));
+    }
+    auto yielded = Interpret(state, op.getRegion(), args);
+    if (state.HasFailure()) {
+      break;
+    }
+    output.InsertElement(indices, yielded[0]);
+  }
+
+  if (op.getNumResults() == 0) {
+    return {};
+  }
+  return {std::move(output)};
+}
+
+llvm::SmallVector<InterpreterValue> Reduce(InterpreterState& state,
+                                           linalg::ReduceOp reduce,
+                                           ArrayRef<InterpreterValue> ins,
+                                           ArrayRef<InterpreterValue> inits) {
+  auto dims = reduce.getDimensions();
+  SmallVector<InterpreterValue> output;
+  for (auto [ty, init] : llvm::zip(reduce.getInits().getTypes(), inits)) {
+    output.push_back(isa<TensorType>(ty) ? init.Clone() : init);
+  }
+  for (const auto& index : ins[0].View().Indices()) {
+    auto dst_index = index;
+    for (int64_t dim : llvm::reverse(dims)) {
+      dst_index.erase(dst_index.begin() + dim);
+    }
+
+    SmallVector<InterpreterValue> args;
+    for (auto& in : ins) {
+      args.push_back(in.ExtractElement(index));
+    }
+    for (auto& out : output) {
+      args.push_back(out.ExtractElement(dst_index));
+    }
+    auto newValues = Interpret(state, reduce.getRegion(), args);
+    if (state.HasFailure()) {
+      return {};
+    }
+
+    for (auto [out, value] : llvm::zip(output, newValues)) {
+      out.InsertElement(dst_index, value);
+    }
+  }
+  if (reduce->getNumResults() == 0) {
+    return {};
+  }
+  return output;
+}
+
+llvm::SmallVector<InterpreterValue> Fill(InterpreterState&, linalg::FillOp op,
+                                         const InterpreterValue& value,
+                                         const InterpreterValue& init) {
+  // TODO(jreiffers): Support variadic Fill.
+  InterpreterValue output = GetInitOperand(op.getOutputs(), 0, {init});
+  output.Fill([&](llvm::ArrayRef<int64_t>) { return value; });
+  if (op.getNumResults() == 0) {
+    return {};
+  }
+  return {std::move(output)};
+}
+
+int64_t Index(InterpreterState& state, linalg::IndexOp index) {
+  return state.GetTopScope()
+      ->GetSideChannel<IterationIndexSideChannel>()
+      ->GetIndices()[index.getDim()];
+}
+
+SmallVector<InterpreterValue> Matmul(InterpreterState& state,
+                                     linalg::MatmulOp matmul,
+                                     ArrayRef<InterpreterValue> inputs,
+                                     const InterpreterValue& init) {
+  if (inputs.size() != 2) {
+    state.AddFailure("Invalid matmul");
+    return {};
+  }
+  const auto& lhs = inputs[0];
+  const auto& rhs = inputs[1];
+  auto ty = matmul.getOutputs()[0].getType();
+  auto result = isa<TensorType>(ty) ? init.Clone() : init;
+  DispatchScalarType(ty, [&](auto dummy) {
+    using TT = TensorOrMemref<decltype(dummy)>;
+    auto lhs_tensor = std::get<TT>(lhs.storage);
+    auto rhs_tensor = std::get<TT>(rhs.storage);
+    auto result_tensor = std::get<TT>(result.storage);
+    for (int64_t i = 0; i < result_tensor.view.sizes[0]; ++i) {
+      for (int64_t j = 0; j < result_tensor.view.sizes[1]; ++j) {
+        for (int64_t k = 0; k < lhs_tensor.view.sizes[1]; ++k) {
+          result_tensor.at({i, j}) +=
+              lhs_tensor.at({i, k}) * rhs_tensor.at({k, j});
+        }
+      }
+    }
+  });
+
+  if (matmul.getNumResults() == 0) {
+    return {};
+  }
+  return {result};
+}
+
+SmallVector<InterpreterValue> Transpose(InterpreterState&,
+                                        linalg::TransposeOp transpose,
+                                        const InterpreterValue& input,
+                                        InterpreterValue init) {
+  auto transposed = TransposeImpl(input, transpose.getPermutation());
+  if (transpose.getNumResults() == 1) {
+    return {transposed};
+  }
+
+  init.Fill([&](auto index) { return transposed.ExtractElement(index); });
+  return {};
+}
+
+SmallVector<InterpreterValue> Dot(InterpreterState&, linalg::DotOp op,
+                                  ArrayRef<InterpreterValue> inputs,
+                                  InterpreterValue acc) {
+  const auto& lhs = inputs[0];
+  const auto& rhs = inputs[1];
+  if (op.getOutputs()[0].getType().isa<TensorType>()) {
+    acc = acc.Clone();
+  }
+  DispatchScalarType(op.getOutputs()[0].getType(), [&](auto dummy) {
+    using TT = TensorOrMemref<decltype(dummy)>;
+    auto lhs_tensor = std::get<TT>(lhs.storage);
+    auto rhs_tensor = std::get<TT>(rhs.storage);
+    auto result_tensor = std::get<TT>(acc.storage);
+    for (int64_t k = 0; k < lhs_tensor.view.sizes[0]; ++k) {
+      result_tensor.at({}) += lhs_tensor.at(k) * rhs_tensor.at(k);
+    }
+  });
+
+  if (op.getNumResults() == 0) {
+    return {};
+  }
+  return {std::move(acc)};
+}
+
+SmallVector<InterpreterValue> Vecmat(InterpreterState&, linalg::VecmatOp op,
+                                     ArrayRef<InterpreterValue> inputs,
+                                     InterpreterValue acc) {
+  const auto& lhs = inputs[0];
+  const auto& rhs = inputs[1];
+  if (op.getOutputs()[0].getType().isa<TensorType>()) {
+    acc = acc.Clone();
+  }
+  DispatchScalarType(op.getOutputs()[0].getType(), [&](auto dummy) {
+    using TT = TensorOrMemref<decltype(dummy)>;
+    auto lhs_tensor = std::get<TT>(lhs.storage);
+    auto rhs_tensor = std::get<TT>(rhs.storage);
+    auto result_tensor = std::get<TT>(acc.storage);
+    for (int64_t j = 0; j < result_tensor.view.sizes[0]; ++j) {
+      for (int64_t k = 0; k < lhs_tensor.view.sizes[0]; ++k) {
+        result_tensor.at(j) += lhs_tensor.at(k) * rhs_tensor.at({k, j});
+      }
+    }
+  });
+
+  if (op.getNumResults() == 0) {
+    return {};
+  }
+  return {std::move(acc)};
+}
+
+REGISTER_MLIR_INTERPRETER_OP("linalg.yield", NoOpTerminator);
+REGISTER_MLIR_INTERPRETER_OP(Broadcast);
+REGISTER_MLIR_INTERPRETER_OP(Dot);
+REGISTER_MLIR_INTERPRETER_OP(Fill);
+REGISTER_MLIR_INTERPRETER_OP(Generic);
+REGISTER_MLIR_INTERPRETER_OP(Index);
+REGISTER_MLIR_INTERPRETER_OP(Map);
+REGISTER_MLIR_INTERPRETER_OP(Matmul);
+REGISTER_MLIR_INTERPRETER_OP(Reduce);
+REGISTER_MLIR_INTERPRETER_OP(Transpose);
+REGISTER_MLIR_INTERPRETER_OP(Vecmat);
+
+}  // namespace
+}  // namespace interpreter
+}  // namespace mlir
diff --git a/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/math.cc b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/math.cc
new file mode 100644
index 00000000000000..02bfded3404e8d
--- /dev/null
+++ b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/math.cc
@@ -0,0 +1,64 @@
+/* Copyright 2022 The OpenXLA Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+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.
+==============================================================================*/
+
+#include "xla/mlir/tools/mlir_interpreter/dialects/cwise_math.h"
+#include "xla/mlir/tools/mlir_interpreter/framework/interpreter_value_util.h"
+#include "xla/mlir/tools/mlir_interpreter/framework/registration.h"
+
+namespace mlir {
+namespace interpreter {
+namespace {
+
+struct CopySign : CwiseFloat {
+  template <typename T>
+  static T Apply(T a, T b) {
+    return std::copysign(a, b);
+  }
+};
+
+REGISTER_MLIR_INTERPRETER_OP("math.absf", ApplyCwiseMap<Abs>);
+REGISTER_MLIR_INTERPRETER_OP("math.absi", ApplyCwiseMap<Abs>);
+REGISTER_MLIR_INTERPRETER_OP("math.atan", ApplyCwiseMap<ATan>);
+REGISTER_MLIR_INTERPRETER_OP("math.atan2", ApplyCwiseBinaryMap<ATan2>);
+REGISTER_MLIR_INTERPRETER_OP("math.cbrt", ApplyCwiseMap<Cbrt>);
+REGISTER_MLIR_INTERPRETER_OP("math.ceil", ApplyCwiseMap<Ceil>);
+REGISTER_MLIR_INTERPRETER_OP("math.copysign", ApplyCwiseBinaryMap<CopySign>);
+REGISTER_MLIR_INTERPRETER_OP("math.cos", ApplyCwiseMap<Cos>);
+REGISTER_MLIR_INTERPRETER_OP("math.ctlz", ApplyCwiseMap<Clz>);
+REGISTER_MLIR_INTERPRETER_OP("math.ctpop", ApplyCwiseMap<PopCount>);
+REGISTER_MLIR_INTERPRETER_OP("math.cttz", ApplyCwiseMap<Ctz>);
+REGISTER_MLIR_INTERPRETER_OP("math.erf", ApplyCwiseMap<Erf>);
+REGISTER_MLIR_INTERPRETER_OP("math.exp", ApplyCwiseMap<Exp>);
+REGISTER_MLIR_INTERPRETER_OP("math.exp2", ApplyCwiseMap<Exp2>);
+REGISTER_MLIR_INTERPRETER_OP("math.expm1", ApplyCwiseMap<ExpM1>);
+REGISTER_MLIR_INTERPRETER_OP("math.floor", ApplyCwiseMap<Floor>);
+REGISTER_MLIR_INTERPRETER_OP("math.ipowi", ApplyCwiseBinaryMap<Power>);
+REGISTER_MLIR_INTERPRETER_OP("math.log", ApplyCwiseMap<Log>);
+REGISTER_MLIR_INTERPRETER_OP("math.log10", ApplyCwiseMap<Log10>);
+REGISTER_MLIR_INTERPRETER_OP("math.log1p", ApplyCwiseMap<Log1P>);
+REGISTER_MLIR_INTERPRETER_OP("math.log2", ApplyCwiseMap<Log2>);
+REGISTER_MLIR_INTERPRETER_OP("math.powf", ApplyCwiseBinaryMap<Power>);
+REGISTER_MLIR_INTERPRETER_OP("math.round", ApplyCwiseMap<Round>);
+REGISTER_MLIR_INTERPRETER_OP("math.roundeven", ApplyCwiseMap<NearbyInt>);
+REGISTER_MLIR_INTERPRETER_OP("math.rsqrt", ApplyCwiseMap<RSqrt>);
+REGISTER_MLIR_INTERPRETER_OP("math.sin", ApplyCwiseMap<Sin>);
+REGISTER_MLIR_INTERPRETER_OP("math.sqrt", ApplyCwiseMap<Sqrt>);
+REGISTER_MLIR_INTERPRETER_OP("math.tan", ApplyCwiseMap<Tan>);
+REGISTER_MLIR_INTERPRETER_OP("math.tanh", ApplyCwiseMap<TanH>);
+REGISTER_MLIR_INTERPRETER_OP("math.trunc", ApplyCwiseMap<Trunc>);
+
+}  // namespace
+}  // namespace interpreter
+}  // namespace mlir
diff --git a/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/memref.cc b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/memref.cc
new file mode 100644
index 00000000000000..0fd938054c1b9b
--- /dev/null
+++ b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/memref.cc
@@ -0,0 +1,253 @@
+/* Copyright 2022 The OpenXLA Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+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.
+==============================================================================*/
+
+#include "mlir/Dialect/MemRef/IR/MemRef.h"  // from @llvm-project
+
+// clang-format erroneously puts the MemRef header above.
+#include <algorithm>  // NOLINT
+#include <cassert>    // NOLINT
+#include <cstdint>    // NOLINT
+#include <utility>    // NOLINT
+
+#include "llvm/ADT/ArrayRef.h"
+#include "llvm/ADT/SmallVector.h"
+#include "llvm/Support/Casting.h"
+#include "mlir/IR/BuiltinTypeInterfaces.h"  // from @llvm-project
+#include "mlir/IR/BuiltinTypes.h"  // from @llvm-project
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+#include "xla/mlir/tools/mlir_interpreter/dialects/util.h"
+#include "xla/mlir/tools/mlir_interpreter/framework/interpreter.h"
+#include "xla/mlir/tools/mlir_interpreter/framework/interpreter_value.h"
+#include "xla/mlir/tools/mlir_interpreter/framework/registration.h"
+#include "xla/mlir/tools/mlir_interpreter/framework/tensor_or_memref.h"
+
+namespace mlir {
+namespace interpreter {
+namespace {
+
+InterpreterValue Load(InterpreterState& state, memref::LoadOp,
+                      const InterpreterValue& memref,
+                      ArrayRef<int64_t> indices) {
+  if (!memref.Buffer()) {
+    state.AddFailure("null pointer dereference.");
+    return {};
+  }
+  if (!memref.View().InBounds(indices)) {
+    state.AddFailure("array index out of bounds");
+    return {};
+  }
+  return memref.ExtractElement(indices);
+}
+
+void Store(InterpreterState& state, memref::StoreOp,
+           const InterpreterValue& value, InterpreterValue memref,
+           ArrayRef<int64_t> indices) {
+  if (memref.View().InBounds(indices)) {
+    memref.InsertElement(indices, value);
+  } else {
+    state.AddFailure("array index out of bounds");
+  }
+}
+
+// TODO(jreiffers): Support symbol operands.
+InterpreterValue Alloc(InterpreterState& state, memref::AllocOp alloc,
+                       ArrayRef<int64_t> dynamic_sizes) {
+  auto ty = cast<ShapedType>(alloc->getResultTypes().front());
+  auto shape = ReplaceDynamicVals(ty.getShape(), dynamic_sizes);
+  auto result =
+      InterpreterValue::MakeTensor(ty.getElementType(), std::move(shape));
+  if (auto* stats = state.GetOptions().stats) {
+    stats->heap_size += result.Buffer()->GetByteSize();
+    stats->peak_heap_size = std::max(stats->peak_heap_size, stats->heap_size);
+    ++stats->num_allocations;
+  }
+  result.Buffer()->SetAllocatedBy(alloc);
+  return result;
+}
+
+InterpreterValue AllocA(InterpreterState&, memref::AllocaOp alloc,
+                        ArrayRef<int64_t> dynamic_sizes) {
+  auto ty = cast<ShapedType>(alloc->getResultTypes().front());
+  auto shape = ReplaceDynamicVals(ty.getShape(), dynamic_sizes);
+  auto result =
+      InterpreterValue::MakeTensor(ty.getElementType(), std::move(shape));
+  result.Buffer()->SetIsAlloca();
+  result.Buffer()->SetAllocatedBy(alloc);
+  return result;
+}
+
+void Dealloc(InterpreterState& state, memref::DeallocOp op,
+             InterpreterValue memref) {
+  if (!memref.Buffer()) {
+    state.AddFailure("attempting to deallocate null pointer.");
+    return;
+  }
+  auto buffer = memref.Buffer();
+  const auto& view = memref.View();
+  if (auto* stats = state.GetOptions().stats) {
+    stats->heap_size -= buffer->GetByteSize();
+    ++stats->num_deallocations;
+  }
+  if (view.GetNumElements() * memref.GetByteSizeOfElement() !=
+      buffer->GetByteSize()) {
+    state.AddFailure("Attempting to deallocate a subview");
+  } else if (!state.GetOptions().disable_deallocations) {
+    buffer->Deallocate(op);
+  }
+}
+
+void Copy(InterpreterState&, memref::CopyOp, InterpreterValue source,
+          InterpreterValue dest) {
+  dest.Fill([&](llvm::ArrayRef<int64_t> indices) {
+    return source.ExtractElement(indices);
+  });
+}
+
+InterpreterValue Subview(InterpreterState& state, memref::SubViewOp subview,
+                         const InterpreterValue& memref,
+                         ArrayRef<int64_t> dynamic_offsets,
+                         ArrayRef<int64_t> dynamic_sizes,
+                         ArrayRef<int64_t> dynamic_strides) {
+  auto v = ExtractOffsetsSizesStrides(dynamic_offsets, dynamic_sizes,
+                                      dynamic_strides, subview);
+  auto out = memref;
+  auto& out_view = out.View();
+  if (!out_view.Subview(v.offsets, v.sizes, v.strides).succeeded()) {
+    state.AddFailure("subview out of bounds");
+    return {};
+  }
+
+  if (subview.getResult().getType().getRank() == out_view.Rank()) {
+    return out;
+  }
+
+  auto shape = subview.getResult().getType().getShape();
+  // TODO(jreiffers): Check why subview.getDroppedDims() yields the wrong shape
+  // here for 1x2x2x3 (-> 1x2x1x3) -> 1x2x3 (claiming 0 is dropped).
+  int64_t dim = 0;
+  while (dim < out_view.Rank() && dim < shape.size()) {
+    if (shape[dim] != 1 && out_view.sizes[dim] == 1) {
+      out_view.sizes.erase(out_view.sizes.begin() + dim);
+      out_view.strides.erase(out_view.strides.begin() + dim);
+    } else {
+      assert((shape[dim] < 0 || out_view.sizes[dim] == shape[dim]) &&
+             "expected static size to match");
+      ++dim;
+    }
+  }
+  while (dim < out_view.Rank()) {
+    assert(out_view.sizes.back() == 1 && "expected remaining dims to be 1");
+    out_view.sizes.pop_back();
+    out_view.strides.pop_back();
+  }
+  return out;
+}
+
+llvm::SmallVector<InterpreterValue> CollapseShape(
+    InterpreterState& state, memref::CollapseShapeOp collapse,
+    const InterpreterValue& memref) {
+  const BufferView& inputView = memref.View();
+  InterpreterValue out = memref;
+  auto& out_view = out.View();
+  out_view.sizes.clear();
+  out_view.strides.clear();
+
+  for (const auto& group : collapse.getReassociationIndices()) {
+    if (auto stride = inputView.GetCollapsedStride(group)) {
+      out_view.strides.push_back(*stride);
+      int64_t& size = out_view.sizes.emplace_back(1);
+      for (int64_t dim : group) size *= inputView.sizes[dim];
+    } else {
+      state.AddFailure("cannot collapse dimensions without a common stride");
+      return {};
+    }
+  }
+
+  return {out};
+}
+
+InterpreterValue Cast(InterpreterState&, memref::CastOp,
+                      InterpreterValue memref) {
+  return memref;
+}
+
+// TODO(jreiffers): Implement full expand_shape support.
+InterpreterValue ExpandShape(InterpreterState& state, memref::ExpandShapeOp op,
+                             InterpreterValue memref) {
+  BufferView inputView = memref.View();
+  auto outTy = cast<MemRefType>(op->getResultTypes()[0]);
+  if (outTy.getNumDynamicDims() > 0) {
+    state.AddFailure("dynamic dimensions unsupported.");
+    return {};
+  }
+
+  InterpreterValue out = memref;
+  auto& out_view = out.View();
+  out_view.strides.clear();
+  out_view.sizes = llvm::to_vector(outTy.getShape());
+  int64_t dummy;
+  if (!getStridesAndOffset(outTy, out_view.strides, dummy).succeeded()) {
+    if (inputView.strides != BufferView::GetDefaultStrides(inputView.sizes)) {
+      state.AddFailure("unsupported strides");
+      return {};
+    }
+    out_view.strides = BufferView::GetDefaultStrides(out_view.sizes);
+  }
+
+  return out;
+}
+
+InterpreterValue GetGlobal(InterpreterState& state,
+                           memref::GetGlobalOp getGlobal) {
+  auto global = llvm::cast<memref::GlobalOp>(
+      state.GetSymbols().lookup(getGlobal.getName()));
+
+  auto value = global.getConstantInitValue();
+  assert(value && "mutable globals are not implemented");
+
+  auto ty = getGlobal->getResultTypes()[0].cast<ShapedType>();
+  return DispatchScalarType(ty, [&](auto dummy) -> InterpreterValue {
+    auto values = value.getValues<decltype(dummy)>();
+    auto result = TensorOrMemref<decltype(dummy)>::Empty(ty.getShape());
+    auto valueIt = values.begin();
+    for (const auto& index : result.view.Indices()) {
+      result.at(index) = *valueIt;
+      ++valueIt;
+    }
+    return {result};
+  });
+}
+
+int64_t Dim(InterpreterState& state, memref::DimOp,
+            const InterpreterValue& memref, int64_t dim) {
+  return DimImpl(memref, dim, state);
+}
+
+REGISTER_MLIR_INTERPRETER_OP(Alloc);
+REGISTER_MLIR_INTERPRETER_OP(AllocA);
+REGISTER_MLIR_INTERPRETER_OP(Cast);
+REGISTER_MLIR_INTERPRETER_OP(CollapseShape);
+REGISTER_MLIR_INTERPRETER_OP(Copy);
+REGISTER_MLIR_INTERPRETER_OP(Dealloc);
+REGISTER_MLIR_INTERPRETER_OP(Dim);
+REGISTER_MLIR_INTERPRETER_OP(ExpandShape);
+REGISTER_MLIR_INTERPRETER_OP(GetGlobal);
+REGISTER_MLIR_INTERPRETER_OP(Load);
+REGISTER_MLIR_INTERPRETER_OP(Store);
+REGISTER_MLIR_INTERPRETER_OP(Subview);
+
+}  // namespace
+}  // namespace interpreter
+}  // namespace mlir
diff --git a/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/mhlo.cc b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/mhlo.cc
new file mode 100644
index 00000000000000..bf77f451c84e11
--- /dev/null
+++ b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/mhlo.cc
@@ -0,0 +1,871 @@
+/* Copyright 2022 The OpenXLA Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+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.
+==============================================================================*/
+
+// NOTE: These ops aim to match the StableHLO spec and the behavior of the XLA
+// compiler. For op semantics and a reference implementation, check the
+// StableHLO repository (the spec is here:
+// https://github.com/openxla/stablehlo/blob/main/docs/spec.md).
+
+#include <algorithm>
+#include <cassert>
+#include <cstdint>
+#include <functional>
+#include <memory>
+#include <numeric>
+#include <type_traits>
+#include <utility>
+
+#include "llvm/ADT/STLExtras.h"
+#include "llvm/ADT/SmallVector.h"
+#include "llvm/Support/Casting.h"
+#include "llvm/Support/ErrorHandling.h"
+#include "mlir/IR/BuiltinAttributes.h"  // from @llvm-project
+#include "mlir/IR/BuiltinTypeInterfaces.h"  // from @llvm-project
+#include "mlir/IR/Operation.h"  // from @llvm-project
+#include "mlir/IR/Types.h"  // from @llvm-project
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+#include "xla/mlir/tools/mlir_interpreter/dialects/comparators.h"
+#include "xla/mlir/tools/mlir_interpreter/dialects/cwise_math.h"
+#include "xla/mlir/tools/mlir_interpreter/dialects/util.h"
+#include "xla/mlir/tools/mlir_interpreter/framework/interpreter.h"
+#include "xla/mlir/tools/mlir_interpreter/framework/interpreter_value.h"
+#include "xla/mlir/tools/mlir_interpreter/framework/interpreter_value_util.h"
+#include "xla/mlir/tools/mlir_interpreter/framework/registration.h"
+#include "xla/mlir/tools/mlir_interpreter/framework/tensor_or_memref.h"
+#include "xla/mlir_hlo/mhlo/IR/hlo_ops.h"
+
+namespace mlir {
+namespace interpreter {
+namespace {
+
+InterpreterValue MakeTuple(MutableArrayRef<InterpreterValue> values) {
+  Tuple result;
+  for (auto& value : values) {
+    result.values.push_back(
+        std::make_shared<InterpreterValue>(std::move(value)));
+  }
+  return {result};
+}
+
+InterpreterValue GetTupleElement(InterpreterState&, mhlo::GetTupleElementOp get,
+                                 const InterpreterValue& tuple) {
+  return *std::get<Tuple>(tuple.storage).values[get.getIndex()];
+}
+
+InterpreterValue BitcastConvert(InterpreterState&, mhlo::BitcastConvertOp op,
+                                const InterpreterValue& in) {
+  ShapedType ty = cast<ShapedType>(op->getResultTypes()[0]);
+  auto result = DispatchScalarType(ty, [&](auto dummy) -> InterpreterValue {
+    TensorOrMemref<decltype(dummy)> result;
+    result.view = {};
+    result.buffer = in.Clone().Buffer();
+    return {result};
+  });
+  auto& out_view = result.View();
+  out_view.strides = BufferView::GetDefaultStrides(ty.getShape());
+  out_view.sizes = llvm::to_vector(ty.getShape());
+  return result;
+}
+
+InterpreterValue BroadcastInDim(InterpreterState&,
+                                mhlo::BroadcastInDimOp broadcast,
+                                InterpreterValue in) {
+  auto boradcast_dims = broadcast.getBroadcastDimensions().getValues<int64_t>();
+  const auto& in_sizes = in.View().sizes;
+  auto out = in.TypedAlike(
+      cast<ShapedType>(broadcast->getResultTypes()[0]).getShape());
+  // TODO(jreiffers): Skip the copy.
+  out.Fill([&](llvm::ArrayRef<int64_t> out_indices) {
+    llvm::SmallVector<int64_t> in_indices;
+    for (auto [inDim, out_dim] : llvm::enumerate(boradcast_dims)) {
+      in_indices.push_back(in_sizes[inDim] == 1 ? 0 : out_indices[out_dim]);
+    }
+    return in.ExtractElement(in_indices);
+  });
+  return out;
+}
+
+InterpreterValue Clamp(InterpreterState&, mhlo::ClampOp,
+                       const InterpreterValue& lb, const InterpreterValue& arg,
+                       const InterpreterValue& ub) {
+  auto result = arg.Clone();
+  for (const auto& index : arg.View().Indices()) {
+    auto lb_scalar = lb.IsTensor() ? lb.ExtractElement(index) : lb;
+    auto ub_scalar = ub.IsTensor() ? ub.ExtractElement(index) : ub;
+    assert(arg.IsTensor() && "clamp only bcasts scalar bounds");
+    auto arg_scalar = arg.ExtractElement(index);
+    auto result_scalar = ApplyCwiseBinaryMap<Min>(
+        ApplyCwiseBinaryMap<Max>(arg_scalar, lb_scalar), ub_scalar);
+    result.InsertElement(index, result_scalar);
+  }
+  return result;
+}
+
+InterpreterValue Concatenate(InterpreterState&, mhlo::ConcatenateOp concat,
+                             ArrayRef<InterpreterValue> vals) {
+  uint64_t dim = concat.getDimension();
+  auto sizes = vals[0].View().sizes;
+  sizes[dim] = 0;
+  for (const auto& val : vals) {
+    sizes[dim] += val.View().sizes[dim];
+  }
+
+  auto result = vals[0].TypedAlike(sizes);
+  int64_t offset = 0;
+  for (const auto& val : vals) {
+    for (auto index : val.View().Indices()) {
+      auto item = val.ExtractElement(index);
+      index[dim] += offset;
+      result.InsertElement(index, item);
+    }
+    offset += val.View().sizes[dim];
+  }
+  return result;
+}
+
+InterpreterValue Reshape(InterpreterState&, mhlo::ReshapeOp reshape,
+                         const InterpreterValue& in) {
+  auto ty = reshape->getResultTypes()[0].cast<mlir::ShapedType>();
+  return ReshapeTensor(in, ty.getShape());
+}
+
+llvm::SmallVector<int64_t> ClampStarts(ArrayRef<int64_t> starts,
+                                       ArrayRef<int64_t> slice_sizes,
+                                       ArrayRef<int64_t> tensor_sizes) {
+  llvm::SmallVector<int64_t> result;
+  for (auto [start, slice_size, tensor_size] :
+       llvm::zip(starts, slice_sizes, tensor_sizes)) {
+    result.push_back(
+        std::max(int64_t{0}, std::min(tensor_size - slice_size, start)));
+  }
+  return result;
+}
+
+InterpreterValue DynamicSlice(InterpreterState&, mhlo::DynamicSliceOp slice,
+                              const InterpreterValue& in,
+                              ArrayRef<int64_t> starts) {
+  auto result = in.TypedAlike(
+      llvm::to_vector(slice.getSliceSizes().getValues<int64_t>()));
+  auto clamed_starts =
+      ClampStarts(starts, result.View().sizes, in.View().sizes);
+  // TODO(jreiffers): Skip the copy.
+  result.Fill([&](llvm::ArrayRef<int64_t> out_indices) {
+    llvm::SmallVector<int64_t> in_indices;
+    for (auto [start, index] : llvm::zip(clamed_starts, out_indices)) {
+      in_indices.push_back(start + index);
+    }
+    return in.ExtractElement(in_indices);
+  });
+  return result;
+}
+
+InterpreterValue DynamicUpdateSlice(InterpreterState&,
+                                    mhlo::DynamicUpdateSliceOp,
+                                    const InterpreterValue& in,
+                                    const InterpreterValue& updates,
+                                    ArrayRef<int64_t> starts) {
+  auto result = in.Clone();
+  auto clamed_starts =
+      ClampStarts(starts, updates.View().sizes, result.View().sizes);
+  for (auto in_indices : updates.View().Indices()) {
+    llvm::SmallVector<int64_t> out_indices;
+    for (auto [start, index] : llvm::zip(clamed_starts, in_indices)) {
+      out_indices.push_back(start + index);
+    }
+    result.InsertElement(out_indices, updates.ExtractElement(in_indices));
+  }
+  return result;
+}
+
+InterpreterValue Slice(InterpreterState&, mhlo::SliceOp slice,
+                       InterpreterValue in) {
+  auto starts = slice.getStartIndices().getValues<int64_t>();
+  auto limits = slice.getLimitIndices().getValues<int64_t>();
+  auto strides = slice.getStrides().getValues<int64_t>();
+
+  llvm::SmallVector<int64_t> sizes;
+  for (auto [start, limit, stride] : llvm::zip(starts, limits, strides)) {
+    sizes.push_back(((limit - start) + (stride - 1)) / stride);
+  }
+  // TODO(jreiffers): Skip the copy.
+  auto result = in.TypedAlike(sizes);
+  result.Fill([&](llvm::ArrayRef<int64_t> out_indices) {
+    llvm::SmallVector<int64_t> in_indices;
+    for (auto [start, stride, index] :
+         llvm::zip(starts, strides, out_indices)) {
+      in_indices.push_back(start + stride * index);
+    }
+    return in.ExtractElement(in_indices);
+  });
+  return result;
+}
+
+llvm::SmallVector<InterpreterValue> Constant(InterpreterState&,
+                                             mhlo::ConstantOp constant) {
+  auto ty = constant->getResultTypes()[0].cast<ShapedType>();
+  return {DispatchScalarType(ty, [&](auto dummy) -> InterpreterValue {
+    if (ty.getElementType().isUnsignedInteger()) {
+      if constexpr (!std::is_same_v<decltype(dummy), bool> &&
+                    std::is_integral_v<decltype(dummy)>) {
+        auto values =
+            constant.getValue()
+                .getValues<
+                    typename std::make_unsigned<decltype(dummy)>::type>();
+        auto result = TensorOrMemref<decltype(dummy)>::Empty(ty.getShape());
+        auto value_it = values.begin();
+        for (const auto& index : result.view.Indices()) {
+          result.at(index) = *value_it;
+          ++value_it;
+        }
+        return {result};
+      } else {
+        llvm_unreachable("invalid input");
+      }
+    } else {
+      auto values = constant.getValue().getValues<decltype(dummy)>();
+      auto result = TensorOrMemref<decltype(dummy)>::Empty(ty.getShape());
+      auto value_it = values.begin();
+      for (const auto& index : result.view.Indices()) {
+        result.at(index) = *value_it;
+        ++value_it;
+      }
+      return {result};
+    }
+  })};
+}
+
+InterpreterValue Pad(InterpreterState&, mhlo::PadOp pad, InterpreterValue arg,
+                     InterpreterValue padding_value) {
+  padding_value = padding_value.ExtractElement({});
+  auto his = pad.getEdgePaddingHigh().getValues<int64_t>();
+  auto los = pad.getEdgePaddingLow().getValues<int64_t>();
+  auto ins = pad.getInteriorPadding().getValues<int64_t>();
+
+  llvm::SmallVector<int64_t> sizes;
+  for (auto [size, lo, in, hi] : llvm::zip(arg.View().sizes, los, ins, his)) {
+    sizes.push_back(size + lo + hi + (size - 1) * in);
+  }
+
+  auto result = arg.TypedAlike(sizes);
+  result.Fill([&](llvm::ArrayRef<int64_t>) { return padding_value; });
+
+  for (const auto& in_indices : arg.View().Indices()) {
+    llvm::SmallVector<int64_t> out_indices;
+    for (auto [in_index, in, lo] : llvm::zip(in_indices, ins, los)) {
+      out_indices.push_back(in_index * (in + 1) + lo);
+    }
+    if (result.View().InBounds(out_indices)) {
+      result.InsertElement(out_indices, arg.ExtractElement(in_indices));
+    }
+  }
+
+  return result;
+}
+
+InterpreterValue Compare(InterpreterState&, mhlo::CompareOp compare,
+                         const InterpreterValue& lhs,
+                         const InterpreterValue& rhs) {
+  switch (compare.getComparisonDirection()) {
+    case mlir::mhlo::ComparisonDirection::EQ:
+      return ApplyCwiseBinaryMap<Foeq>(lhs, rhs);
+    case mlir::mhlo::ComparisonDirection::NE:
+      return ApplyCwiseBinaryMap<Fune>(lhs, rhs);
+    case mlir::mhlo::ComparisonDirection::GE:
+      return ApplyCwiseBinaryMap<Foge>(lhs, rhs);
+    case mlir::mhlo::ComparisonDirection::GT:
+      return ApplyCwiseBinaryMap<Fogt>(lhs, rhs);
+    case mlir::mhlo::ComparisonDirection::LE:
+      return ApplyCwiseBinaryMap<Fole>(lhs, rhs);
+    case mlir::mhlo::ComparisonDirection::LT:
+      return ApplyCwiseBinaryMap<Folt>(lhs, rhs);
+  }
+}
+
+llvm::SmallVector<InterpreterValue> Gather(
+    MutableArrayRef<InterpreterValue> args, mlir::Operation* op,
+    InterpreterState& state) {
+  auto gather = llvm::dyn_cast<mhlo::GatherOp>(op);
+  auto dynamic_gather = llvm::dyn_cast<mhlo::DynamicGatherOp>(op);
+
+  if (!gather && !dynamic_gather) {
+    state.AddFailure("invalid gather");
+    return {};
+  }
+
+  const auto& dims = gather ? gather.getDimensionNumbers()
+                            : dynamic_gather.getDimensionNumbers();
+
+  auto index_vector_dim = dims.getIndexVectorDim();
+  auto start_index_map = dims.getStartIndexMap();
+  auto offset_dims = dims.getOffsetDims();
+  auto collapsed_slice_dims = dims.getCollapsedSliceDims();
+  auto slice_sizes =
+      gather ? llvm::to_vector(gather.getSliceSizes().getValues<int64_t>())
+             : llvm::to_vector(llvm::map_range(
+                   args[2].View().Indices(), [&](const auto& indices) {
+                     return args[2].ExtractElement(indices).AsInt();
+                   }));
+
+  auto& operand = args[0];
+  auto& start_indices = args[1];
+  const auto& operand_view = operand.View();
+  int64_t operand_rank = operand_view.Rank();
+
+  // Make a fake BufferView for the start indices.
+  BufferView start_indices_view = start_indices.View();
+  auto output_rank =
+      static_cast<int64_t>(start_indices_view.Rank() + offset_dims.size());
+  if (index_vector_dim < start_indices_view.Rank()) {
+    --output_rank;
+    start_indices_view.sizes[index_vector_dim] = 1;
+  }
+
+  SmallVector<int64_t> batch_dims;
+  for (int64_t i = 0; i < output_rank; ++i) {
+    if (!llvm::is_contained(offset_dims, i)) {
+      batch_dims.push_back(i);
+    }
+  }
+
+  // Make a fake BufferView for the slice indices.
+  BufferView slice_indices_view{0, SmallVector<int64_t>{slice_sizes}, {}};
+
+  SmallVector<int64_t> non_collapsed_slice_dims;
+  for (int64_t i = 0; i < operand_rank; ++i) {
+    if (!llvm::is_contained(collapsed_slice_dims, i)) {
+      non_collapsed_slice_dims.push_back(i);
+    }
+  }
+
+  SmallVector<int64_t> output_sizes(output_rank);
+  for (auto [output_dim, slice_dim] :
+       llvm::zip(offset_dims, non_collapsed_slice_dims)) {
+    output_sizes[output_dim] = slice_sizes[slice_dim];
+  }
+  for (auto [batch_index, output_dim] : llvm::enumerate(batch_dims)) {
+    if (batch_index >= index_vector_dim) {
+      ++batch_index;
+    }
+    output_sizes[output_dim] = start_indices_view.sizes[batch_index];
+  }
+
+  auto output = operand.TypedAlike(output_sizes);
+  for (auto start_indices_index : start_indices_view.Indices()) {
+    SmallVector<int64_t> operand_base_indices(operand_rank);
+    for (auto [i, dim] : llvm::enumerate(start_index_map)) {
+      if (index_vector_dim < start_indices_view.Rank()) {
+        start_indices_index[index_vector_dim] = static_cast<int64_t>(i);
+      }
+      operand_base_indices[dim] = std::max<int64_t>(
+          0, std::min(start_indices.ExtractElement(start_indices_index).AsInt(),
+                      operand_view.sizes[dim] - slice_sizes[dim]));
+    }
+
+    for (const auto& slice_indices : slice_indices_view.Indices()) {
+      SmallVector<int64_t> operand_indices;
+      for (int64_t i = 0; i < operand_rank; ++i) {
+        operand_indices.push_back(operand_base_indices[i] + slice_indices[i]);
+      }
+
+      SmallVector<int64_t> output_indices(output_rank);
+      for (auto [output_dim, slice_dim] :
+           llvm::zip(offset_dims, non_collapsed_slice_dims)) {
+        output_indices[output_dim] = slice_indices[slice_dim];
+      }
+      for (auto [batch_index, output_dim] : llvm::enumerate(batch_dims)) {
+        output_indices[output_dim] =
+            start_indices_index[batch_index >= index_vector_dim
+                                    ? batch_index + 1
+                                    : batch_index];
+      }
+
+      auto value = operand.ExtractElement(operand_indices);
+      output.InsertElement(output_indices, value);
+    }
+  }
+
+  return {output};
+}
+
+llvm::SmallVector<InterpreterValue> Scatter(
+    InterpreterState& state, mhlo::ScatterOp scatter,
+    ArrayRef<InterpreterValue> n_inputs, InterpreterValue scatter_indices,
+    ArrayRef<InterpreterValue> n_updates) {
+  const auto& dims = scatter.getScatterDimensionNumbers();
+  auto index_vector_dim = dims.getIndexVectorDim();
+  auto scatter_dims_to_operand_dims = dims.getScatterDimsToOperandDims();
+  auto inserted_window_dims = dims.getInsertedWindowDims();
+  auto update_window_dims = dims.getUpdateWindowDims();
+
+  auto input_view = n_inputs.front().View();
+  int64_t operand_rank = input_view.Rank();
+  int64_t updates_rank = n_updates.front().View().Rank();
+  int64_t indices_rank = scatter_indices.View().Rank();
+
+  llvm::SmallVector<int64_t> batch_dims;
+  for (int64_t dim = 0; dim < operand_rank; ++dim) {
+    if (!llvm::is_contained(inserted_window_dims, dim)) {
+      batch_dims.push_back(dim);
+    }
+  }
+
+  llvm::SmallVector<int64_t> update_scatter_dims;
+  for (int64_t dim = 0; dim < updates_rank; ++dim) {
+    if (!llvm::is_contained(update_window_dims, dim)) {
+      update_scatter_dims.push_back(dim);
+    }
+  }
+
+  llvm::SmallVector<InterpreterValue> n_results;
+  for (auto& inputs : n_inputs) {
+    n_results.push_back(inputs.Clone());
+  }
+
+  for (auto [inputs, updates, results] :
+       llvm::zip(n_results, n_updates, n_results)) {
+    const auto& updates_view = updates.View();
+    for (const auto& update_indices : updates_view.Indices()) {
+      llvm::SmallVector<int64_t> input_indices(operand_rank);
+      llvm::SmallVector<int64_t> max_indices(operand_rank);
+      llvm::SmallVector<int64_t> min_indices(operand_rank);
+      llvm::SmallVector<int64_t> scatter_indices_index(indices_rank);
+
+      for (auto [index, dim] : llvm::enumerate(update_scatter_dims)) {
+        scatter_indices_index[index >= index_vector_dim ? index + 1 : index] +=
+            update_indices[dim];
+      }
+
+      for (auto [update_dim, operand_dim] :
+           llvm::zip(update_window_dims, batch_dims)) {
+        input_indices[operand_dim] = update_indices[update_dim];
+        max_indices[operand_dim] = updates_view.sizes[update_dim] - 1;
+      }
+
+      for (auto [index, dim] : llvm::enumerate(scatter_dims_to_operand_dims)) {
+        if (index_vector_dim < indices_rank) {
+          scatter_indices_index[index_vector_dim] = static_cast<int64_t>(index);
+        }
+
+        int64_t scatter_index =
+            scatter_indices.ExtractElement(scatter_indices_index).AsInt();
+        input_indices[dim] += scatter_index;
+        min_indices[dim] += scatter_index;
+        max_indices[dim] += scatter_index;
+      }
+
+      if (!input_view.InBounds(min_indices) ||
+          !input_view.InBounds(max_indices)) {
+        continue;
+      }
+
+      auto current_value = inputs.ExtractElement(input_indices).AsUnitTensor();
+      auto update = updates.ExtractElement(update_indices).AsUnitTensor();
+
+      auto result = Interpret(state, scatter.getUpdateComputation(),
+                              {current_value, update});
+      if (state.HasFailure()) {
+        return n_results;
+      }
+      inputs.InsertElement(input_indices, result.front().ExtractElement({}));
+    }
+  }
+
+  return n_results;
+}
+
+InterpreterValue Select(InterpreterState&, mhlo::SelectOp,
+                        TensorOrMemref<bool> condition,
+                        const InterpreterValue& true_vals,
+                        const InterpreterValue& false_vals) {
+  // TODO(jreiffers): Support implicit broadcasting.
+  auto result = true_vals.Clone();
+  for (const auto& indices : condition.view.Indices()) {
+    if (!condition.at(indices)) {
+      result.InsertElement(indices, false_vals.ExtractElement(indices));
+    }
+  }
+  return result;
+}
+
+InterpreterValue Transpose(InterpreterState&, mhlo::TransposeOp transpose,
+                           const InterpreterValue& tensor) {
+  auto permutation = transpose.getPermutation().getValues<int64_t>();
+  return TransposeImpl(tensor, llvm::to_vector(permutation));
+}
+
+InterpreterValue Iota(InterpreterState&, mhlo::IotaOp iota) {
+  auto dim = iota.getIotaDimension();
+  auto ty = iota->getResultTypes()[0].cast<ShapedType>();
+  return DispatchScalarType(ty, [&](auto dummy) -> InterpreterValue {
+    auto result = TensorOrMemref<decltype(dummy)>::Empty({ty.getShape()[dim]});
+    for (const auto& index : result.view.Indices()) {
+      result.at(index) = index[0];
+    }
+    result.view.sizes = SmallVector<int64_t>(ty.getShape());
+    result.view.strides = SmallVector<int64_t>(result.view.sizes.size());
+    result.view.strides[dim] = 1;
+    return {result};
+  });
+}
+
+template <typename R>
+struct Caster {
+  template <typename A>
+  constexpr static bool SupportedType() {
+    return std::is_convertible_v<A, R>;
+  }
+
+  template <typename A>
+  static R Apply(A v) {
+    return v;
+  }
+};
+
+InterpreterValue Convert(InterpreterState&, mhlo::ConvertOp op,
+                         InterpreterValue in) {
+  return DispatchScalarType(op->getResultTypes()[0],
+                            [&](auto dummy) -> InterpreterValue {
+                              return ApplyCwiseMap<Caster<decltype(dummy)>>(in);
+                            });
+}
+
+llvm::SmallVector<InterpreterValue> While(
+    InterpreterState& state, mhlo::WhileOp whileop,
+    SmallVector<InterpreterValue> loop_vars) {
+  auto cond = Interpret(state, whileop.getRegion(0), loop_vars);
+  while (!state.HasFailure() &&
+         std::get<TensorOrMemref<bool>>(cond.front().storage).at({})) {
+    loop_vars = Interpret(state, whileop.getRegion(1), loop_vars);
+    if (state.HasFailure()) {
+      break;
+    }
+    cond = Interpret(state, whileop.getRegion(0), loop_vars);
+  }
+  return loop_vars;
+}
+
+InterpreterValue Copy(InterpreterState&, mhlo::CopyOp op,
+                      const InterpreterValue& tensor) {
+  auto layout = op->getAttr("result_layout");
+  if (!layout) {
+    return tensor.Clone();
+  }
+  return tensor.Clone(
+      llvm::to_vector(cast<DenseIntElementsAttr>(layout).getValues<int64_t>()));
+}
+
+llvm::SmallVector<InterpreterValue> Fusion(InterpreterState& state,
+                                           mhlo::FusionOp fusion,
+                                           ArrayRef<InterpreterValue> args) {
+  return Interpret(state, fusion.getRegion(), args);
+}
+
+llvm::SmallVector<InterpreterValue> Reduce(InterpreterState& state,
+                                           mhlo::ReduceOp reduce,
+                                           ArrayRef<InterpreterValue> operands,
+                                           ArrayRef<InterpreterValue> inits) {
+  auto dims = reduce.getDimensions().getValues<int64_t>();
+
+  auto out_sizes = operands.front().View().sizes;
+  for (int64_t dim : llvm::reverse(dims)) {
+    out_sizes.erase(out_sizes.begin() + dim);
+  }
+
+  SmallVector<InterpreterValue> results;
+  for (auto [operand, init] : llvm::zip(operands, inits)) {
+    results.push_back(operand.TypedAlike(out_sizes));
+    auto init_scalar = init.ExtractElement({});
+    results.back().Fill([&](llvm::ArrayRef<int64_t>) { return init_scalar; });
+  }
+
+  for (const auto& index : operands[0].View().Indices()) {
+    auto dst_index = index;
+    for (int64_t dim : llvm::reverse(dims)) {
+      dst_index.erase(dst_index.begin() + dim);
+    }
+
+    SmallVector<InterpreterValue> args;
+    for (auto& result : results) {
+      args.push_back(result.ExtractElement(dst_index).AsUnitTensor());
+    }
+    for (auto& operand : operands) {
+      args.push_back(operand.ExtractElement(index).AsUnitTensor());
+    }
+    auto new_values = Interpret(state, reduce.getRegion(), args);
+    if (state.HasFailure()) {
+      return results;
+    }
+
+    for (auto [result, value] : llvm::zip(results, new_values)) {
+      result.InsertElement(dst_index, value.ExtractElement({}));
+    }
+  }
+  return results;
+}
+
+SmallVector<InterpreterValue> Sort(InterpreterState& state, mhlo::SortOp op,
+                                   ArrayRef<InterpreterValue> inputs) {
+  const auto& shape = inputs.front().View().sizes;
+  uint64_t dim =
+      op.getDimension() < shape.size() ? op.getDimension() : shape.size() - 1;
+  SmallVector<int64_t> indices(shape[dim]);
+
+  auto iter_view = inputs.front().View();
+  iter_view.sizes[dim] = 1;
+  SmallVector<InterpreterValue> results;
+  for (const auto& input : inputs) {
+    results.push_back(input.TypedAlike(input.View().sizes));
+  }
+  for (const auto& offset : iter_view.Indices()) {
+    std::iota(indices.begin(), indices.end(), 0);
+    auto comparator = [&](int64_t a, int64_t b) {
+      if (state.HasFailure()) {
+        return false;
+      }
+      SmallVector<InterpreterValue> args;
+      auto index_a = offset;
+      index_a[dim] = a;
+      auto index_b = offset;
+      index_b[dim] = b;
+      for (const auto& input : inputs) {
+        args.push_back(input.ExtractElement(index_a).AsUnitTensor());
+        args.push_back(input.ExtractElement(index_b).AsUnitTensor());
+      }
+      auto result = Interpret(state, op.getComparator(), args);
+      return !state.HasFailure() && result[0].ExtractElement({}).AsInt() != 0;
+    };
+
+    if (op.getIsStable()) {
+      std::stable_sort(indices.begin(), indices.end(), comparator);
+    } else {
+      std::sort(indices.begin(), indices.end(), comparator);
+    }
+
+    auto copy = offset;
+    for (auto [dst, src] : llvm::enumerate(indices)) {
+      for (auto [input, result] : llvm::zip(inputs, results)) {
+        copy[dim] = src;
+        auto elem = input.ExtractElement(copy);
+        copy[dim] = static_cast<int64_t>(dst);
+        result.InsertElement(copy, elem);
+      }
+    }
+  }
+  return results;
+}
+
+SmallVector<InterpreterValue> Case(InterpreterState& state, mhlo::CaseOp op,
+                                   int64_t index) {
+  if (index < 0 || index >= op->getNumResults()) {
+    index = op->getNumRegions() - 1;
+  }
+  return Interpret(state, op->getRegion(index), {});
+}
+
+InterpreterValue DotGeneralImpl(InterpreterValue& lhs, InterpreterValue& rhs,
+                                ArrayRef<int64_t> lhs_contracting,
+                                ArrayRef<int64_t> rhs_contracting,
+                                ArrayRef<int64_t> lhs_batch,
+                                ArrayRef<int64_t> rhs_batch,
+                                mlir::Type element_ty) {
+  auto& lhsv = lhs.View();
+  auto& rhsv = rhs.View();
+  SmallVector<int64_t> dimensions;
+  SmallVector<int64_t> lhs_non_batch;
+  SmallVector<int64_t> rhs_non_batch;
+  auto nbatch = static_cast<int64_t>(lhs_batch.size());
+  for (int64_t lhs_dim : lhs_batch) {
+    dimensions.push_back(lhsv.sizes[lhs_dim]);
+  }
+  for (int64_t i = 0; i < lhsv.Rank(); i++) {
+    if (!llvm::is_contained(lhs_contracting, i) &&
+        !llvm::is_contained(lhs_batch, i)) {
+      dimensions.push_back(lhsv.sizes[i]);
+      lhs_non_batch.push_back(i);
+    }
+  }
+  for (int64_t i = 0; i < rhs.View().Rank(); i++) {
+    if (!llvm::is_contained(rhs_contracting, i) &&
+        !llvm::is_contained(rhs_batch, i)) {
+      dimensions.push_back(rhsv.sizes[i]);
+      rhs_non_batch.push_back(i);
+    }
+  }
+
+  SmallVector<int64_t> lhs_index(lhsv.Rank());
+  SmallVector<int64_t> rhs_index(rhsv.Rank());
+  SmallVector<int64_t> output_index(dimensions.size());
+  auto output = lhs.TypedAlike(dimensions);
+
+  std::function<void(int64_t)> apply_batch, apply_lhs, apply_rhs,
+      apply_contract;
+
+  apply_batch = [&](int64_t dim) {
+    if (dim >= nbatch) {
+      apply_lhs(0);
+      return;
+    }
+    for (int64_t i = 0; i < dimensions[dim]; ++i) {
+      lhs_index[lhs_batch[dim]] = i;
+      rhs_index[rhs_batch[dim]] = i;
+      output_index[dim] = i;
+      apply_batch(dim + 1);
+    }
+  };
+
+  apply_lhs = [&](int64_t dim) {
+    if (dim >= lhs_non_batch.size()) {
+      apply_rhs(0);
+      return;
+    }
+    int64_t out_dim = nbatch + dim;
+    for (int64_t i = 0; i < dimensions[out_dim]; ++i) {
+      lhs_index[lhs_non_batch[dim]] = i;
+      output_index[out_dim] = i;
+      apply_lhs(dim + 1);
+    }
+  };
+
+  apply_rhs = [&](int64_t dim) {
+    if (dim >= rhs_non_batch.size()) {
+      apply_contract(0);
+      return;
+    }
+    auto out_dim = static_cast<int64_t>(nbatch + lhs_non_batch.size() + dim);
+    for (int64_t i = 0; i < dimensions[out_dim]; ++i) {
+      rhs_index[rhs_non_batch[dim]] = i;
+      output_index[out_dim] = i;
+      apply_rhs(dim + 1);
+    }
+  };
+
+  apply_contract = [&](int64_t dim) {
+    if (dim >= lhs_contracting.size()) {
+      DispatchScalarType(element_ty, [&](auto dummy) {
+        using T = TensorOrMemref<decltype(dummy)>;
+        std::get<T>(output.storage).at(output_index) +=
+            std::get<T>(lhs.storage).at(lhs_index) *
+            std::get<T>(rhs.storage).at(rhs_index);
+      });
+      return;
+    }
+    for (int64_t i = 0; i < lhsv.sizes[lhs_contracting[dim]]; ++i) {
+      lhs_index[lhs_contracting[dim]] = i;
+      rhs_index[rhs_contracting[dim]] = i;
+      apply_contract(dim + 1);
+    }
+  };
+
+  apply_batch(0);
+  return output;
+}
+
+// TODO(jreiffers): Unify this with DotGeneral.
+InterpreterValue Dot(InterpreterState& state, mhlo::DotOp op,
+                     const InterpreterValue& lhs, const InterpreterValue& rhs) {
+  auto ty = cast<ShapedType>(op->getResultTypes()[0]);
+  auto result = lhs.TypedAlike(ty.getShape());
+
+  if (lhs.View().Rank() == 1 && rhs.View().Rank() == 1) {
+    DispatchScalarType(ty, [&](auto dummy) {
+      using T = decltype(dummy);
+      using TT = TensorOrMemref<T>;
+      auto lhs_tensor = std::get<TT>(lhs.storage);
+      auto rhs_tensor = std::get<TT>(rhs.storage);
+
+      T product = 0;
+      for (int64_t i = 0; i < lhs_tensor.view.sizes[0]; ++i) {
+        product += lhs_tensor.at(i) * rhs_tensor.at(i);
+      }
+      std::get<TT>(result.storage).at({}) += product;
+    });
+  } else if (lhs.View().Rank() == 2 && rhs.View().Rank() == 1) {
+    DispatchScalarType(ty, [&](auto dummy) {
+      using TT = TensorOrMemref<decltype(dummy)>;
+      auto lhs_tensor = std::get<TT>(lhs.storage);
+      auto rhs_tensor = std::get<TT>(rhs.storage);
+      auto result_tensor = std::get<TT>(result.storage);
+      for (int64_t i = 0; i < result_tensor.view.sizes[0]; ++i) {
+        for (int64_t j = 0; j < rhs_tensor.view.sizes[0]; ++j) {
+          result_tensor.at(i) += lhs_tensor.at({i, j}) * rhs_tensor.at(j);
+        }
+      }
+    });
+  } else if (lhs.View().Rank() == 2 && rhs.View().Rank() == 2) {
+    DispatchScalarType(ty, [&](auto dummy) {
+      using TT = TensorOrMemref<decltype(dummy)>;
+      auto lhs_tensor = std::get<TT>(lhs.storage);
+      auto rhs_tensor = std::get<TT>(rhs.storage);
+      auto result_tensor = std::get<TT>(result.storage);
+      for (int64_t i = 0; i < result_tensor.view.sizes[0]; ++i) {
+        for (int64_t j = 0; j < result_tensor.view.sizes[1]; ++j) {
+          for (int64_t k = 0; k < lhs_tensor.view.sizes[1]; ++k) {
+            result_tensor.at({i, j}) +=
+                lhs_tensor.at({i, k}) * rhs_tensor.at({k, j});
+          }
+        }
+      }
+    });
+  } else {
+    state.AddFailure("unsupported dot");
+  }
+
+  return result;
+}
+
+InterpreterValue DotGeneral(InterpreterState&, mhlo::DotGeneralOp op,
+                            InterpreterValue lhs, InterpreterValue rhs) {
+  const auto& dims = op.getDotDimensionNumbers();
+  return DotGeneralImpl(
+      lhs, rhs, dims.getLhsContractingDimensions(),
+      dims.getRhsContractingDimensions(), dims.getLhsBatchingDimensions(),
+      dims.getRhsBatchingDimensions(),
+      op->getResultTypes()[0].cast<ShapedType>().getElementType());
+}
+
+// TODO(jreiffers): Migrate remaining ops to the safer signature.
+REGISTER_MLIR_INTERPRETER_OP("mhlo.dynamic_gather", Gather);
+REGISTER_MLIR_INTERPRETER_OP("mhlo.gather", Gather);
+REGISTER_MLIR_INTERPRETER_OP("mhlo.return", NoOpTerminator);
+REGISTER_MLIR_INTERPRETER_OP("mhlo.tuple", MakeTuple);
+REGISTER_MLIR_INTERPRETER_OP(BitcastConvert);
+REGISTER_MLIR_INTERPRETER_OP(BroadcastInDim);
+REGISTER_MLIR_INTERPRETER_OP(Clamp);
+REGISTER_MLIR_INTERPRETER_OP(Compare);
+REGISTER_MLIR_INTERPRETER_OP(Concatenate);
+REGISTER_MLIR_INTERPRETER_OP(Constant);
+REGISTER_MLIR_INTERPRETER_OP(Convert);
+REGISTER_MLIR_INTERPRETER_OP(Copy);
+REGISTER_MLIR_INTERPRETER_OP(Dot);
+REGISTER_MLIR_INTERPRETER_OP(DotGeneral);
+REGISTER_MLIR_INTERPRETER_OP(DynamicSlice);
+REGISTER_MLIR_INTERPRETER_OP(DynamicUpdateSlice);
+REGISTER_MLIR_INTERPRETER_OP(Fusion);
+REGISTER_MLIR_INTERPRETER_OP(GetTupleElement);
+REGISTER_MLIR_INTERPRETER_OP(Iota);
+REGISTER_MLIR_INTERPRETER_OP(Case);
+REGISTER_MLIR_INTERPRETER_OP(While);
+REGISTER_MLIR_INTERPRETER_OP(Pad);
+REGISTER_MLIR_INTERPRETER_OP(Reduce);
+REGISTER_MLIR_INTERPRETER_OP(Reshape);
+REGISTER_MLIR_INTERPRETER_OP(Scatter);
+REGISTER_MLIR_INTERPRETER_OP(Select);
+REGISTER_MLIR_INTERPRETER_OP(Slice);
+REGISTER_MLIR_INTERPRETER_OP(Sort);
+REGISTER_MLIR_INTERPRETER_OP(Transpose);
+
+}  // namespace
+}  // namespace interpreter
+}  // namespace mlir
diff --git a/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/mhlo_binary_cwise.cc b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/mhlo_binary_cwise.cc
new file mode 100644
index 00000000000000..00c4a38aa2286f
--- /dev/null
+++ b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/mhlo_binary_cwise.cc
@@ -0,0 +1,45 @@
+/* Copyright 2023 The OpenXLA Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+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.
+==============================================================================*/
+
+#include "xla/mlir/tools/mlir_interpreter/dialects/cwise_math.h"
+#include "xla/mlir/tools/mlir_interpreter/framework/interpreter_value_util.h"
+#include "xla/mlir/tools/mlir_interpreter/framework/registration.h"
+
+namespace mlir {
+namespace interpreter {
+namespace {
+
+REGISTER_MLIR_INTERPRETER_OP("mhlo.add", ApplyCwiseBinaryMap<Plus>);
+REGISTER_MLIR_INTERPRETER_OP("mhlo.and", ApplyCwiseBinaryMap<BitAnd>);
+REGISTER_MLIR_INTERPRETER_OP("mhlo.atan2", ApplyCwiseBinaryMap<ATan2>);
+REGISTER_MLIR_INTERPRETER_OP("mhlo.complex", ApplyCwiseBinaryMap<Complex>);
+REGISTER_MLIR_INTERPRETER_OP("mhlo.divide", ApplyCwiseBinaryMap<Divide>);
+REGISTER_MLIR_INTERPRETER_OP("mhlo.maximum", ApplyCwiseBinaryMap<Max>);
+REGISTER_MLIR_INTERPRETER_OP("mhlo.minimum", ApplyCwiseBinaryMap<Min>);
+REGISTER_MLIR_INTERPRETER_OP("mhlo.multiply", ApplyCwiseBinaryMap<Multiply>);
+REGISTER_MLIR_INTERPRETER_OP("mhlo.or", ApplyCwiseBinaryMap<BitOr>);
+REGISTER_MLIR_INTERPRETER_OP("mhlo.power", ApplyCwiseBinaryMap<Power>);
+REGISTER_MLIR_INTERPRETER_OP("mhlo.remainder", ApplyCwiseBinaryMap<Remainder>);
+REGISTER_MLIR_INTERPRETER_OP("mhlo.shift_left", ApplyCwiseBinaryMap<ShiftLeft>);
+REGISTER_MLIR_INTERPRETER_OP("mhlo.shift_right_arithmetic",
+                             ApplyCwiseBinaryMap<ShiftRightArith>);
+REGISTER_MLIR_INTERPRETER_OP("mhlo.shift_right_logical",
+                             ApplyCwiseBinaryMap<ShiftRightLogical>);
+REGISTER_MLIR_INTERPRETER_OP("mhlo.subtract", ApplyCwiseBinaryMap<Minus>);
+REGISTER_MLIR_INTERPRETER_OP("mhlo.xor", ApplyCwiseBinaryMap<BitXor>);
+
+}  // namespace
+}  // namespace interpreter
+}  // namespace mlir
diff --git a/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/mhlo_unary_cwise.cc b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/mhlo_unary_cwise.cc
new file mode 100644
index 00000000000000..e8d09418369589
--- /dev/null
+++ b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/mhlo_unary_cwise.cc
@@ -0,0 +1,84 @@
+/* Copyright 2023 The OpenXLA Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+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.
+==============================================================================*/
+
+#include <complex>
+
+#include "xla/mlir/tools/mlir_interpreter/dialects/cwise_math.h"
+#include "xla/mlir/tools/mlir_interpreter/framework/interpreter_value_util.h"
+#include "xla/mlir/tools/mlir_interpreter/framework/registration.h"
+
+namespace mlir {
+namespace interpreter {
+namespace {
+
+struct Logistic : CwiseNonIntegral {
+  template <typename T>
+  static T Apply(T a) {
+    if (std::real(a) < 0) {
+      T e = std::exp(a);
+      return e / (e + T{1});
+    }
+    return T{1} / (std::exp(-a) + T{1});
+  }
+};
+
+// Note: this can't be replaced with std::bit_not, which always returns true for
+// bools.
+struct Not : CwiseIntegral {
+  static bool Apply(bool a) { return !a; }
+
+  template <typename T>
+  static T Apply(T a) {
+    return ~a;
+  }
+};
+
+struct Sign : CwiseSigned {
+  template <typename T>
+  static T Apply(T a) {
+    return std::copysign(T{1}, a);
+  }
+};
+
+REGISTER_MLIR_INTERPRETER_OP("mhlo.abs", ApplyCwiseMap<Abs>);
+REGISTER_MLIR_INTERPRETER_OP("mhlo.cbrt", ApplyCwiseMap<Cbrt>);
+REGISTER_MLIR_INTERPRETER_OP("mhlo.ceil", ApplyCwiseMap<Ceil>);
+REGISTER_MLIR_INTERPRETER_OP("mhlo.cosine", ApplyCwiseMap<Cos>);
+REGISTER_MLIR_INTERPRETER_OP("mhlo.count_leading_zeros", ApplyCwiseMap<Clz>);
+REGISTER_MLIR_INTERPRETER_OP("mhlo.exponential", ApplyCwiseMap<Exp>);
+REGISTER_MLIR_INTERPRETER_OP("mhlo.exponential_minus_one",
+                             ApplyCwiseMap<ExpM1>);
+REGISTER_MLIR_INTERPRETER_OP("mhlo.floor", ApplyCwiseMap<Floor>);
+REGISTER_MLIR_INTERPRETER_OP("mhlo.imag", ApplyCwiseMap<Imag>);
+REGISTER_MLIR_INTERPRETER_OP("mhlo.is_finite", ApplyCwiseMap<IsFinite>);
+REGISTER_MLIR_INTERPRETER_OP("mhlo.log", ApplyCwiseMap<Log>);
+REGISTER_MLIR_INTERPRETER_OP("mhlo.log_plus_one", ApplyCwiseMap<Log1P>);
+REGISTER_MLIR_INTERPRETER_OP("mhlo.logistic", ApplyCwiseMap<Logistic>);
+REGISTER_MLIR_INTERPRETER_OP("mhlo.negate", ApplyCwiseMap<Neg>);
+REGISTER_MLIR_INTERPRETER_OP("mhlo.not", ApplyCwiseMap<Not>);
+REGISTER_MLIR_INTERPRETER_OP("mhlo.popcnt", ApplyCwiseMap<PopCount>);
+REGISTER_MLIR_INTERPRETER_OP("mhlo.real", ApplyCwiseMap<Real>);
+REGISTER_MLIR_INTERPRETER_OP("mhlo.round_nearest_afz", ApplyCwiseMap<Round>);
+REGISTER_MLIR_INTERPRETER_OP("mhlo.round_nearest_even",
+                             ApplyCwiseMap<NearbyInt>);
+REGISTER_MLIR_INTERPRETER_OP("mhlo.rsqrt", ApplyCwiseMap<RSqrt>);
+REGISTER_MLIR_INTERPRETER_OP("mhlo.sign", ApplyCwiseMap<Sign>);
+REGISTER_MLIR_INTERPRETER_OP("mhlo.sine", ApplyCwiseMap<Sin>);
+REGISTER_MLIR_INTERPRETER_OP("mhlo.sqrt", ApplyCwiseMap<Sqrt>);
+REGISTER_MLIR_INTERPRETER_OP("mhlo.tanh", ApplyCwiseMap<TanH>);
+
+}  // namespace
+}  // namespace interpreter
+}  // namespace mlir
diff --git a/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/scf.cc b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/scf.cc
new file mode 100644
index 00000000000000..60137e379441dd
--- /dev/null
+++ b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/scf.cc
@@ -0,0 +1,222 @@
+/* Copyright 2022 The OpenXLA Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+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.
+==============================================================================*/
+
+#include "mlir/Dialect/SCF/IR/SCF.h"  // from @llvm-project
+
+#include <cassert>   // NOLINT
+#include <cstdint>   // NOLINT
+#include <iterator>  // NOLINT
+#include <memory>    // NOLINT
+#include <utility>   // NOLINT
+
+#include "llvm/ADT/ArrayRef.h"
+#include "llvm/ADT/STLExtras.h"
+#include "llvm/ADT/SmallVector.h"
+#include "mlir/Dialect/Affine/LoopUtils.h"  // from @llvm-project
+#include "mlir/IR/BuiltinTypes.h"  // from @llvm-project
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+#include "xla/mlir/tools/mlir_interpreter/dialects/util.h"
+#include "xla/mlir/tools/mlir_interpreter/framework/interpreter.h"
+#include "xla/mlir/tools/mlir_interpreter/framework/interpreter_value.h"
+#include "xla/mlir/tools/mlir_interpreter/framework/registration.h"
+#include "xla/mlir/tools/mlir_interpreter/framework/tensor_or_memref.h"
+
+namespace mlir {
+namespace interpreter {
+namespace {
+
+class ParallelSideChannel : public InterpreterSideChannel {
+ public:
+  explicit ParallelSideChannel(llvm::SmallVector<InterpreterValue>& results)
+      : results_(results) {}
+
+  InterpreterValue& result(int index) const { return results_[index]; }
+
+ private:
+  SmallVector<InterpreterValue>& results_;
+};
+
+llvm::SmallVector<InterpreterValue> For(InterpreterState& state, scf::ForOp op,
+                                        int64_t lb, int64_t ub, int64_t step,
+                                        ArrayRef<InterpreterValue> inits) {
+  llvm::SmallVector<InterpreterValue> results;
+  for (int64_t i = 0; i < inits.size(); ++i) {
+    results.push_back(GetInitOperand(op.getInitArgs(), i, inits));
+  }
+
+  auto& region = op->getRegion(0);
+  for (; lb < ub; lb += step) {
+    SmallVector<InterpreterValue> inputs;
+    DispatchScalarType(op.getLowerBound().getType(), [&](auto dummy) {
+      inputs.push_back(InterpreterValue{static_cast<decltype(dummy)>(lb)});
+    });
+    llvm::copy(results, std::back_inserter(inputs));
+    results = Interpret(state, region, inputs);
+    if (state.HasFailure()) {
+      break;
+    }
+  }
+  return results;
+}
+
+llvm::SmallVector<InterpreterValue> ForAll(
+    InterpreterState& state, scf::ForallOp op,
+    ArrayRef<InterpreterValue> dynamic_lower_bounds,
+    ArrayRef<InterpreterValue> dynamic_upper_bounds,
+    ArrayRef<InterpreterValue> dynamic_steps,
+    ArrayRef<InterpreterValue> inits) {
+  bool is_bufferized = op.getNumResults() == 0;
+
+  // Clone any tensors that are passed in `shared_outs`.
+  SmallVector<InterpreterValue> outs = llvm::to_vector(inits);
+  for (auto [type, out] : llvm::zip(op.getOutputs().getTypes(), outs)) {
+    if (isa<TensorType>(type)) {
+      out = out.Clone();
+    }
+  }
+
+  auto lbs = ReplaceDynamicVals(op.getStaticLowerBound(), dynamic_lower_bounds);
+  auto ubs = ReplaceDynamicVals(op.getStaticUpperBound(), dynamic_upper_bounds);
+  auto steps = ReplaceDynamicVals(op.getStaticStep(), dynamic_steps);
+
+  SmallVector<int64_t> iter_sizes;
+  for (auto [lb, ub, step] : llvm::zip(lbs, ubs, steps)) {
+    if (step == 0) {
+      state.AddFailure("invalid step");
+      return {};
+    }
+    iter_sizes.push_back((ub - lb + (step - 1)) / step);
+  }
+
+  // Make a fake buffer view to abuse its index iterator.
+  BufferView view{0, std::move(iter_sizes), {}};
+  for (const auto& indices : view.Indices()) {
+    SmallVector<InterpreterValue> args;
+    for (auto [i, lb, step] : llvm::zip(indices, lbs, steps)) {
+      args.push_back(InterpreterValue{i * step + lb});
+    }
+    llvm::copy(outs, std::back_inserter(args));
+
+    auto yielded = Interpret(state, op->getRegion(0), args);
+    if (state.HasFailure()) {
+      break;
+    }
+    assert(yielded.empty() && "forall loop shouldn't have yielded anything");
+  }
+
+  if (is_bufferized) {
+    return {};
+  }
+  return outs;
+}
+
+void InParallel(InterpreterState& state, scf::InParallelOp op) {
+  Interpret(state, op.getRegion(), {});
+}
+
+llvm::SmallVector<InterpreterValue> If(InterpreterState& state, scf::IfOp op,
+                                       bool condition) {
+  if (condition) {
+    return Interpret(state, op.getThenRegion(), {});
+  }
+  if (op.getElseRegion().hasOneBlock()) {
+    return Interpret(state, op.getElseRegion(), {});
+  }
+  return {};
+}
+
+llvm::SmallVector<InterpreterValue> Parallel(InterpreterState& state,
+                                             scf::ParallelOp parallel,
+                                             ArrayRef<int64_t> lbs,
+                                             ArrayRef<int64_t> ubs,
+                                             ArrayRef<int64_t> steps,
+                                             ArrayRef<InterpreterValue> inits) {
+  llvm::SmallVector<InterpreterValue> results;
+  for (int64_t i = 0; i < parallel.getNumReductions(); ++i) {
+    results.push_back(GetInitOperand(parallel.getInitVals(), i, inits));
+  }
+
+  BufferView iter;
+  for (auto [lb, ub, step] : llvm::zip(lbs, ubs, steps)) {
+    iter.sizes.push_back((ub - lb + (step - 1)) / step);
+  }
+
+  // Make the results available to reduce ops.
+  state.GetTopScope()->SetSideChannel(
+      std::make_shared<ParallelSideChannel>(results));
+  for (const auto& indices : iter.Indices()) {
+    SmallVector<InterpreterValue> iter_args;
+    for (auto [i, lb, step] : llvm::zip(indices, lbs, steps)) {
+      iter_args.push_back(InterpreterValue{i * step + lb});
+    }
+
+    // Execute the region. It has no results.
+    Interpret(state, parallel.getRegion(), iter_args);
+  }
+
+  return results;
+}
+
+void Reduce(InterpreterState& state, scf::ReduceOp reduce,
+            ArrayRef<InterpreterValue> operands) {
+  if (operands.size() != reduce.getNumRegions()) {
+    state.AddFailure("reduce op has wrong number of operands");
+    return;
+  }
+  for (int i = 0; i < reduce.getNumRegions(); ++i) {
+    auto& accumulator =
+        state.GetTopScope()->GetSideChannel<ParallelSideChannel>()->result(i);
+    auto results =
+        Interpret(state, reduce.getRegion(i), {accumulator, operands[i]});
+    if (state.HasFailure()) {
+      return;
+    }
+    accumulator = results.front();
+  }
+}
+
+llvm::SmallVector<InterpreterValue> While(
+    InterpreterState& state, scf::WhileOp op,
+    MutableArrayRef<InterpreterValue> inits) {
+  auto loop_vars = Interpret(state, op.getBefore(), inits);
+  while (!state.HasFailure() && std::get<bool>(loop_vars.front().storage)) {
+    loop_vars = Interpret(state, op.getAfter(),
+                          ArrayRef<InterpreterValue>(loop_vars).drop_front());
+    if (state.HasFailure()) {
+      break;
+    }
+    loop_vars = Interpret(state, op.getBefore(), loop_vars);
+  }
+  if (state.HasFailure()) {
+    return {};
+  }
+  loop_vars.erase(loop_vars.begin());
+  return loop_vars;
+}
+
+REGISTER_MLIR_INTERPRETER_OP("scf.condition", NoOpTerminator);
+REGISTER_MLIR_INTERPRETER_OP("scf.reduce.return", NoOpTerminator);
+REGISTER_MLIR_INTERPRETER_OP("scf.yield", NoOpTerminator);
+REGISTER_MLIR_INTERPRETER_OP(ForAll);
+REGISTER_MLIR_INTERPRETER_OP(InParallel);
+REGISTER_MLIR_INTERPRETER_OP(Parallel);
+REGISTER_MLIR_INTERPRETER_OP(Reduce);
+REGISTER_MLIR_INTERPRETER_OP(For);
+REGISTER_MLIR_INTERPRETER_OP(If);
+REGISTER_MLIR_INTERPRETER_OP(While);
+
+}  // namespace
+}  // namespace interpreter
+}  // namespace mlir
diff --git a/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tensor.cc b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tensor.cc
new file mode 100644
index 00000000000000..0abed8dc48b573
--- /dev/null
+++ b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tensor.cc
@@ -0,0 +1,262 @@
+/* Copyright 2022 The OpenXLA Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+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.
+==============================================================================*/
+
+#include "mlir/Dialect/Tensor/IR/Tensor.h"  // from @llvm-project
+
+#include <cassert>   // NOLINT
+#include <cstdint>   // NOLINT
+#include <optional>  // NOLINT
+
+#include "llvm/ADT/STLExtras.h"
+#include "llvm/ADT/SmallVector.h"
+#include "mlir/IR/BuiltinTypeInterfaces.h"  // from @llvm-project
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+#include "xla/mlir/tools/mlir_interpreter/dialects/util.h"
+#include "xla/mlir/tools/mlir_interpreter/framework/interpreter.h"
+#include "xla/mlir/tools/mlir_interpreter/framework/interpreter_value.h"
+#include "xla/mlir/tools/mlir_interpreter/framework/interpreter_value_util.h"
+#include "xla/mlir/tools/mlir_interpreter/framework/registration.h"
+
+namespace mlir {
+namespace interpreter {
+namespace {
+
+int64_t Dim(InterpreterState& state, tensor::DimOp,
+            const InterpreterValue& tensor, int64_t dim) {
+  return DimImpl(tensor, dim, state);
+}
+
+InterpreterValue Empty(InterpreterState&, tensor::EmptyOp op,
+                       ArrayRef<int64_t> dynamic_sizes) {
+  auto ty = op->getResultTypes().front().cast<mlir::ShapedType>();
+  auto shape = ReplaceDynamicVals(ty.getShape(), dynamic_sizes);
+  return InterpreterValue::MakeTensor(ty.getElementType(), shape);
+}
+
+InterpreterValue Extract(InterpreterState& state, tensor::ExtractOp,
+                         InterpreterValue tensor, ArrayRef<int64_t> indices) {
+  if (!tensor.View().InBounds(indices)) {
+    state.AddFailure("array index out of bounds");
+    return {};
+  }
+  return tensor.ExtractElement(indices);
+}
+
+InterpreterValue FromElements(InterpreterState&, tensor::FromElementsOp op,
+                              MutableArrayRef<InterpreterValue> elements) {
+  auto ty = op->getResultTypes().front().cast<mlir::ShapedType>();
+  auto result = InterpreterValue::MakeTensor(ty.getElementType(),
+                                             llvm::to_vector(ty.getShape()));
+  for (auto [index, element] : llvm::zip(result.View().Indices(), elements)) {
+    result.InsertElement(index, element);
+  }
+  return result;
+}
+
+llvm::SmallVector<InterpreterValue> CollapseShape(
+    InterpreterState&, tensor::CollapseShapeOp op,
+    const InterpreterValue& tensor) {
+  SmallVector<int64_t> sizes;
+  for (const auto& indices : op.getReassociationIndices()) {
+    int64_t size = 1;
+    for (auto dim : indices) {
+      size *= tensor.View().sizes[dim];
+    }
+    sizes.push_back(size);
+  }
+  return {ReshapeTensor(tensor, sizes)};
+}
+
+llvm::SmallVector<InterpreterValue> ExpandShape(
+    InterpreterState&, tensor::ExpandShapeOp op,
+    const InterpreterValue& tensor) {
+  auto ty = cast<ShapedType>(op->getResultTypes().front());
+  auto sizes = llvm::to_vector(ty.getShape());
+  for (const auto& [src_index, dst_indices] :
+       llvm::enumerate(op.getReassociationIndices())) {
+    int64_t size = tensor.View().sizes[src_index];
+    std::optional<int64_t> dyn_index = std::nullopt;
+    for (auto dim : dst_indices) {
+      if (sizes[dim] < 0) {
+        dyn_index = dim;
+      } else {
+        size /= sizes[dim];
+      }
+    }
+    if (dyn_index) {
+      sizes[*dyn_index] = size;
+    }
+  }
+  return {ReshapeTensor(tensor, sizes)};
+}
+
+llvm::SmallVector<InterpreterValue> ExtractSlice(
+    InterpreterState&, tensor::ExtractSliceOp extract, InterpreterValue tensor,
+    ArrayRef<int64_t> dynamic_offsets, ArrayRef<int64_t> dynamic_sizes,
+    ArrayRef<int64_t> dynamic_strides) {
+  auto v = ExtractOffsetsSizesStrides(dynamic_offsets, dynamic_sizes,
+                                      dynamic_strides, extract);
+  int64_t rank = v.offsets.size();
+  auto out = tensor.TypedAlike(v.sizes);
+  out.Fill([&](llvm::ArrayRef<int64_t> indices) {
+    llvm::SmallVector<int64_t> src_indices;
+    for (int64_t i = 0; i < rank; ++i) {
+      src_indices.push_back(indices[i] * v.strides[i] + v.offsets[i]);
+    }
+    return tensor.ExtractElement(src_indices);
+  });
+
+  int64_t num_dropped = 0;
+  auto& out_view = out.View();
+
+  // TODO(jreiffers): Figure out why getDroppedDims fails here when there's
+  // no rank reduction and the output has a dynamic shape.
+  int64_t dim = 0;
+  const auto& result_sizes = extract.getResultType().getShape();
+  const auto& static_sizes = extract.getStaticSizes();
+  while (dim < out_view.Rank()) {
+    if (static_sizes[num_dropped + dim] == 1 &&
+        (dim >= result_sizes.size() || result_sizes[dim] != 1)) {
+      out_view.sizes.erase(out_view.sizes.begin() + dim);
+      out_view.strides.erase(out_view.strides.begin() + dim);
+      ++num_dropped;
+    } else {
+      ++dim;
+    }
+  }
+
+  return {out};
+}
+
+template <typename Op>
+llvm::SmallVector<InterpreterValue> InsertSlice(
+    InterpreterState&, Op insert, InterpreterValue src, InterpreterValue dest,
+    ArrayRef<int64_t> dynamic_offsets, ArrayRef<int64_t> dynamic_sizes,
+    ArrayRef<int64_t> dynamic_strides) {
+  // parallel_insert_slice actually writes to its destination.
+  if (insert->getNumResults() == 1) {
+    dest = dest.Clone();
+  }
+  auto v = ExtractOffsetsSizesStrides(dynamic_offsets, dynamic_sizes,
+                                      dynamic_strides, insert);
+
+  auto static_sizes = insert.getStaticSizes();
+  llvm::SmallVector<int64_t> inserted_dims;
+  auto& src_view = src.View();
+  auto* src_size_it = src_view.sizes.begin();
+  for (auto [dim, size] : llvm::enumerate(static_sizes)) {
+    if (src_size_it == src_view.sizes.end() ||
+        (*src_size_it != size && size >= 0)) {
+      assert(size == 1 && "Can only insert unit dims");
+      inserted_dims.push_back(dim);
+    } else {
+      ++src_size_it;
+    }
+  }
+
+  for (const auto& src_indices : src_view.Indices()) {
+    llvm::SmallVector<int64_t> src_with_inserted_dims = src_indices;
+    for (int64_t dim : inserted_dims) {
+      src_with_inserted_dims.insert(src_with_inserted_dims.begin() + dim, 0);
+    }
+    llvm::SmallVector<int64_t> dst_indices;
+    for (auto [src_index, stride, offset] :
+         llvm::zip(src_with_inserted_dims, v.strides, v.offsets)) {
+      dst_indices.push_back(src_index * stride + offset);
+    }
+    dest.InsertElement(dst_indices, src.ExtractElement(src_indices));
+  }
+  if (insert->getNumResults() == 1) {
+    return {dest};
+  }
+  return {};
+}
+
+InterpreterValue Generate(InterpreterState& state, tensor::GenerateOp generate,
+                          ArrayRef<int64_t> dynamic_sizes) {
+  auto ty = generate->getResultTypes().front().cast<ShapedType>();
+  auto sizes = ReplaceDynamicVals(ty.getShape(), dynamic_sizes);
+
+  auto result = InterpreterValue::MakeTensor(ty.getElementType(), sizes);
+  result.Fill([&](ArrayRef<int64_t> indices) -> InterpreterValue {
+    auto values = Interpret(state, generate.getRegion(),
+                            PackInterpreterValues<int64_t>(indices));
+    if (state.HasFailure()) {
+      return {result.ExtractElement(indices)};
+    }
+    return values.front();
+  });
+  return {result};
+}
+
+InterpreterValue Insert(InterpreterState& state, tensor::InsertOp,
+                        const InterpreterValue& value,
+                        const InterpreterValue& tensor,
+                        ArrayRef<int64_t> indices) {
+  auto result = tensor.Clone();
+  if (result.View().InBounds(indices)) {
+    result.InsertElement(indices, value);
+  } else {
+    state.AddFailure("array index out of bounds");
+  }
+  return result;
+}
+
+InterpreterValue Pad(InterpreterState& state, tensor::PadOp pad,
+                     InterpreterValue tensor, ArrayRef<int64_t> dynamic_lows,
+                     ArrayRef<int64_t> dynamic_highs) {
+  auto lows = ReplaceDynamicVals(pad.getStaticLow(), dynamic_lows);
+  auto highs = ReplaceDynamicVals(pad.getStaticHigh(), dynamic_highs);
+
+  auto& view = tensor.View();
+  llvm::SmallVector<int64_t> result_sizes;
+  for (auto [size, low, high] : llvm::zip(view.sizes, lows, highs)) {
+    result_sizes.push_back(size + low + high);
+  }
+
+  auto result = tensor.TypedAlike(result_sizes);
+  result.Fill([&](llvm::ArrayRef<int64_t> out_index) -> InterpreterValue {
+    llvm::SmallVector<int64_t> in_index;
+    for (auto [index, low] : llvm::zip(out_index, lows)) {
+      in_index.push_back(index - low);
+    }
+    if (view.InBounds(in_index)) {
+      return tensor.ExtractElement(in_index);
+    }
+    return Interpret(state, pad.getRegion(), PackInterpreterValues(out_index))
+        .front();
+  });
+  return result;
+}
+
+REGISTER_MLIR_INTERPRETER_OP("tensor.cast",
+                             "builtin.unrealized_conversion_cast");
+REGISTER_MLIR_INTERPRETER_OP("tensor.yield", NoOpTerminator);
+REGISTER_MLIR_INTERPRETER_OP(CollapseShape);
+REGISTER_MLIR_INTERPRETER_OP(Dim);
+REGISTER_MLIR_INTERPRETER_OP(Empty);
+REGISTER_MLIR_INTERPRETER_OP(ExpandShape);
+REGISTER_MLIR_INTERPRETER_OP(Extract);
+REGISTER_MLIR_INTERPRETER_OP(ExtractSlice);
+REGISTER_MLIR_INTERPRETER_OP(FromElements);
+REGISTER_MLIR_INTERPRETER_OP(Generate);
+REGISTER_MLIR_INTERPRETER_OP(Insert);
+REGISTER_MLIR_INTERPRETER_OP(InsertSlice<tensor::InsertSliceOp>);
+REGISTER_MLIR_INTERPRETER_OP(InsertSlice<tensor::ParallelInsertSliceOp>);
+REGISTER_MLIR_INTERPRETER_OP(Pad);
+
+}  // namespace
+}  // namespace interpreter
+}  // namespace mlir
diff --git a/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/BUILD b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/BUILD
new file mode 100644
index 00000000000000..c146c9ea69a970
--- /dev/null
+++ b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/BUILD
@@ -0,0 +1,13 @@
+load("//xla:lit.bzl", "lit_test_suite")
+
+# copybara:uncomment package(default_applicable_licenses = ["//tensorflow:license"])
+
+lit_test_suite(
+    name = "all_tests",
+    srcs = glob(["**/*.mlir"]),
+    cfg = "//xla:lit.cfg.py",
+    tools = [
+        "//xla/mlir/tools/mlir_interpreter:mlir-interpreter-runner",
+        "@llvm-project//llvm:FileCheck",
+    ],
+)
diff --git a/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/affine/apply.mlir b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/affine/apply.mlir
new file mode 100644
index 00000000000000..43feca045e927f
--- /dev/null
+++ b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/affine/apply.mlir
@@ -0,0 +1,63 @@
+// RUN: mlir-interpreter-runner %s -run-all | FileCheck %s
+
+func.func @ceildiv() -> index {
+  %c25 = arith.constant 25 : index
+  %ret = affine.apply affine_map<(d0)[] -> ((d0 ceildiv 8) * 8)>(%c25)[]
+  return %ret : index
+}
+
+// CHECK-LABEL: @ceildiv
+// CHECK: Results
+// CHECK-NEXT: 32
+
+func.func @floordiv() -> index {
+  %c25 = arith.constant 25 : index
+  %ret = affine.apply affine_map<(d0)[] -> ((d0 floordiv 8) * 8)>(%c25)[]
+  return %ret : index
+}
+
+// CHECK-LABEL: @floordiv
+// CHECK: Results
+// CHECK-NEXT: 24
+
+func.func @add() -> index {
+  %c100 = arith.constant 100 : index
+  %c42 = arith.constant 42 : index
+  %ret = affine.apply affine_map<(d0, d1)[] -> (d0 + d1)>(%c100, %c42)[]
+  return %ret : index
+}
+
+// CHECK-LABEL: @add
+// CHECK: Results
+// CHECK-NEXT: 142
+
+func.func @mod() -> index {
+  %c99 = arith.constant 99 : index
+  %ret = affine.apply affine_map<(d0)[] -> (d0 mod 10)>(%c99)[]
+  return %ret : index
+}
+
+// CHECK-LABEL: @mod
+// CHECK: Results
+// CHECK-NEXT: 9
+
+func.func @mul() -> index {
+  %c100 = arith.constant 100 : index
+  %ret = affine.apply affine_map<(d0)[] -> (d0 * 42)>(%c100)[]
+  return %ret : index
+}
+
+// CHECK-LABEL: @mul
+// CHECK: Results
+// CHECK-NEXT: 4200
+
+func.func @symbol() -> index {
+  %c1 = arith.constant 1 : index
+  %c20 = arith.constant 20 : index
+  %ret = affine.apply affine_map<(d0)[s0] -> (d0 + s0)>(%c1)[%c20]
+  return %ret : index
+}
+
+// CHECK-LABEL: @symbol
+// CHECK: Results
+// CHECK-NEXT: 21
diff --git a/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/affine/minmax.mlir b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/affine/minmax.mlir
new file mode 100644
index 00000000000000..cf38c45f3608ab
--- /dev/null
+++ b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/affine/minmax.mlir
@@ -0,0 +1,36 @@
+// RUN: mlir-interpreter-runner %s -run-all | FileCheck %s
+
+#map = affine_map<(d0)[] -> (1000, d0 + 512, d0*100)>
+
+func.func @min() -> (index, index, index) {
+  %c1 = arith.constant 1 : index
+  %c8 = arith.constant 8 : index
+  %c500 = arith.constant 500 : index
+
+  %0 = affine.min #map (%c1)[]
+  %1 = affine.min #map (%c8)[]
+  %2 = affine.min #map (%c500)[]
+
+  return %0, %1, %2 : index, index, index
+}
+
+// CHECK-LABEL: @min
+// CHECK-NEXT: Results
+// CHECK-NEXT: 100
+// CHECK-NEXT: 520
+// CHECK-NEXT: 1000
+
+func.func @max() -> (index, index) {
+  %c1 = arith.constant 1 : index
+  %c11 = arith.constant 11 : index
+
+  %0 = affine.max #map (%c1)[]
+  %1 = affine.max #map (%c11)[]
+
+  return %0, %1 : index, index
+}
+
+// CHECK-LABEL: @max
+// CHECK-NEXT: Results
+// CHECK-NEXT: 1000
+// CHECK-NEXT: 1100
diff --git a/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/arith/bitcast.mlir b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/arith/bitcast.mlir
new file mode 100644
index 00000000000000..5dc30c1fbbba20
--- /dev/null
+++ b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/arith/bitcast.mlir
@@ -0,0 +1,21 @@
+// RUN: mlir-interpreter-runner %s -run-all | FileCheck %s
+
+func.func @bitcast_scalar() -> i32 {
+  %c15 = arith.constant 1.5 : f32
+  %ret = arith.bitcast %c15 : f32 to i32
+  return %ret : i32
+}
+
+// CHECK-LABEL: @bitcast_scalar
+// CHECK-NEXT: Results
+// CHECK-NEXT: 1069547520
+
+func.func @bitcast_vector() -> vector<2xf32> {
+  %c15 = arith.constant dense<[15, 25]> : vector<2xi32>
+  %ret = arith.bitcast %c15 : vector<2xi32> to vector<2xf32>
+  return %ret : vector<2xf32>
+}
+
+// CHECK-LABEL: @bitcast_vector
+// CHECK-NEXT: Results
+// CHECK-NEXT: [2.101948e-44, 3.503246e-44]
diff --git a/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/arith/cmpf.mlir b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/arith/cmpf.mlir
new file mode 100644
index 00000000000000..e416f353b18781
--- /dev/null
+++ b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/arith/cmpf.mlir
@@ -0,0 +1,129 @@
+// RUN: mlir-interpreter-runner %s -run-all | FileCheck %s
+
+func.func private @compare(%a: f32, %b: f32) -> tensor<16xi1> {
+  %false = arith.cmpf false, %a, %b : f32
+  %oeq = arith.cmpf oeq, %a, %b : f32
+  %ogt = arith.cmpf ogt, %a, %b : f32
+  %oge = arith.cmpf oge, %a, %b : f32
+
+  %olt = arith.cmpf olt, %a, %b : f32
+  %ole = arith.cmpf ole, %a, %b : f32
+  %one = arith.cmpf one, %a, %b : f32
+  %ord = arith.cmpf ord, %a, %b : f32
+
+  %ueq = arith.cmpf ueq, %a, %b : f32
+  %ugt = arith.cmpf ugt, %a, %b : f32
+  %uge = arith.cmpf uge, %a, %b : f32
+  %ult = arith.cmpf ult, %a, %b : f32
+
+  %ule = arith.cmpf ule, %a, %b : f32
+  %une = arith.cmpf une, %a, %b : f32
+  %uno = arith.cmpf uno, %a, %b : f32
+  %true = arith.cmpf true, %a, %b : f32
+
+  %ret = tensor.from_elements
+    %false, %oeq, %ogt, %oge,
+    %olt, %ole, %one, %ord,
+    %ueq, %ugt, %uge, %ult,
+    %ule, %une, %uno, %true : tensor<16xi1>
+  func.return %ret : tensor<16xi1>
+}
+
+func.func @nan_vs_one() -> tensor<16xi1> {
+  %nan = arith.constant 0x7fc00000 : f32
+  %one = arith.constant 1.0 : f32
+
+  %ret = func.call @compare(%nan, %one) : (f32, f32) -> tensor<16xi1>
+  func.return %ret : tensor<16xi1>
+}
+
+// CHECK-LABEL: @nan_vs_one
+// CHECK-NEXT: Results
+// CHECK-NEXT: [false, false, false, false,
+// CHECK-SAME:  false, false, false, false,
+// CHECK-SAME:  true, true, true, true,
+// CHECK-SAME:  true, true, true, true]
+
+func.func @one_vs_nan() -> tensor<16xi1> {
+  %nan = arith.constant 0x7fc00000 : f32
+  %one = arith.constant 1.0 : f32
+
+  %ret = func.call @compare(%one, %nan) : (f32, f32) -> tensor<16xi1>
+  func.return %ret : tensor<16xi1>
+}
+
+// CHECK-LABEL: @one_vs_nan
+// CHECK-NEXT: Results
+// CHECK-NEXT: [false, false, false, false,
+// CHECK-SAME:  false, false, false, false,
+// CHECK-SAME:  true, true, true, true,
+// CHECK-SAME:  true, true, true, true]
+
+func.func @nan_vs_nan() -> tensor<16xi1> {
+  %nan = arith.constant 0x7fc00000 : f32
+
+  %ret = func.call @compare(%nan, %nan) : (f32, f32) -> tensor<16xi1>
+  func.return %ret : tensor<16xi1>
+}
+
+// CHECK-LABEL: @nan_vs_nan
+// CHECK-NEXT: Results
+// CHECK-NEXT: [false, false, false, false,
+// CHECK-SAME:  false, false, false, false,
+// CHECK-SAME:  true, true, true, true,
+// CHECK-SAME:  true, true, true, true]
+
+func.func @one_vs_one() -> tensor<16xi1> {
+  %one = arith.constant 1.0 : f32
+
+  %ret = func.call @compare(%one, %one) : (f32, f32) -> tensor<16xi1>
+  func.return %ret : tensor<16xi1>
+}
+
+// CHECK-LABEL: @one_vs_one
+// CHECK-NEXT: Results
+// CHECK-NEXT: [false, true, false, true,
+// CHECK-SAME:  false, true, false, true,
+// CHECK-SAME:  true, false, true, false,
+// CHECK-SAME:  true, false, false, true]
+
+func.func @one_vs_two() -> tensor<16xi1> {
+  %one = arith.constant 1.0 : f32
+  %two = arith.constant 2.0 : f32
+
+  %ret = func.call @compare(%one, %two) : (f32, f32) -> tensor<16xi1>
+  func.return %ret : tensor<16xi1>
+}
+
+// CHECK-LABEL: @one_vs_two
+// CHECK-NEXT: Results
+// CHECK-NEXT: [false, false, false, false,
+// CHECK-SAME:  true, true, true, true,
+// CHECK-SAME:  false, false, false, true,
+// CHECK-SAME:  true, true, false, true]
+
+func.func @two_vs_one() -> tensor<16xi1> {
+  %one = arith.constant 1.0 : f32
+  %two = arith.constant 2.0 : f32
+
+  %ret = func.call @compare(%two, %one) : (f32, f32) -> tensor<16xi1>
+  func.return %ret : tensor<16xi1>
+}
+
+// CHECK-LABEL: @two_vs_one
+// CHECK-NEXT: Results
+// CHECK-NEXT: [false, false, true, true,
+// CHECK-SAME:  false, false, true, true,
+// CHECK-SAME:  false, true, true, false,
+// CHECK-SAME:  false, true, false, true]
+
+func.func @vector() -> vector<4xi1> {
+  %0 = arith.constant dense<[1.0, 2.0, 3.0, 4.0]> : vector<4xf32>
+  %1 = arith.constant dense<[10.0, 2.0, 30.0, 4.0]> : vector<4xf32>
+  %ret = arith.cmpf oeq, %0, %1 : vector<4xf32>
+  return %ret : vector<4xi1>
+}
+
+// CHECK-LABEL: @vector
+// CHECK-NEXT: Results
+// CHECK-NEXT: vector<4xi1>: [false, true, false, true]
diff --git a/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/arith/cmpi.mlir b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/arith/cmpi.mlir
new file mode 100644
index 00000000000000..dbca4f718d5ccd
--- /dev/null
+++ b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/arith/cmpi.mlir
@@ -0,0 +1,147 @@
+// RUN: mlir-interpreter-runner %s -run-all | FileCheck %s
+
+func.func @eq() -> (i1, i1) {
+  %c1 = arith.constant 1 : index
+  %c2 = arith.constant 2 : index
+  %e1 = arith.cmpi eq, %c1, %c2 : index
+  %e2 = arith.cmpi eq, %c1, %c1 : index
+  return %e1, %e2 : i1, i1
+}
+
+// CHECK-LABEL: @eq
+// CHECK-NEXT: Results
+// CHECK-NEXT{LITERAL}: false
+// CHECK-NEXT{LITERAL}: true
+
+func.func @ne() -> (i1, i1) {
+  %c1 = arith.constant 1 : index
+  %c2 = arith.constant 2 : index
+  %e1 = arith.cmpi ne, %c1, %c2 : index
+  %e2 = arith.cmpi ne, %c1, %c1 : index
+  return %e1, %e2 : i1, i1
+}
+
+// CHECK-LABEL: @ne
+// CHECK-NEXT: Results
+// CHECK-NEXT{LITERAL}: true
+// CHECK-NEXT{LITERAL}: false
+
+func.func @slt() -> (i1, i1, i1) {
+  %c-1 = arith.constant -1 : index
+  %c1 = arith.constant 1 : index
+  %e1 = arith.cmpi slt, %c-1, %c1 : index
+  %e2 = arith.cmpi slt, %c1, %c-1 : index
+  %e3 = arith.cmpi slt, %c1, %c1 : index
+  return %e1, %e2, %e3 : i1, i1, i1
+}
+
+// CHECK-LABEL: @slt
+// CHECK-NEXT: Results
+// CHECK-NEXT{LITERAL}: true
+// CHECK-NEXT{LITERAL}: false
+// CHECK-NEXT{LITERAL}: false
+
+func.func @sle() -> (i1, i1, i1) {
+  %c-1 = arith.constant -1 : index
+  %c1 = arith.constant 1 : index
+  %e1 = arith.cmpi sle, %c-1, %c1 : index
+  %e2 = arith.cmpi sle, %c1, %c-1 : index
+  %e3 = arith.cmpi sle, %c1, %c1 : index
+  return %e1, %e2, %e3 : i1, i1, i1
+}
+
+// CHECK-LABEL: @sle
+// CHECK-NEXT: Results
+// CHECK-NEXT{LITERAL}: true
+// CHECK-NEXT{LITERAL}: false
+// CHECK-NEXT{LITERAL}: true
+
+func.func @sgt() -> (i1, i1, i1) {
+  %c-1 = arith.constant -1 : index
+  %c1 = arith.constant 1 : index
+  %e1 = arith.cmpi sgt, %c-1, %c1 : index
+  %e2 = arith.cmpi sgt, %c1, %c-1 : index
+  %e3 = arith.cmpi sgt, %c1, %c1 : index
+  return %e1, %e2, %e3 : i1, i1, i1
+}
+
+// CHECK-LABEL: @sgt
+// CHECK-NEXT: Results
+// CHECK-NEXT{LITERAL}: false
+// CHECK-NEXT{LITERAL}: true
+// CHECK-NEXT{LITERAL}: false
+
+func.func @sge() -> (i1, i1, i1) {
+  %c-1 = arith.constant -1 : index
+  %c1 = arith.constant 1 : index
+  %e1 = arith.cmpi sge, %c-1, %c1 : index
+  %e2 = arith.cmpi sge, %c1, %c-1 : index
+  %e3 = arith.cmpi sge, %c1, %c1 : index
+  return %e1, %e2, %e3 : i1, i1, i1
+}
+
+// CHECK-LABEL: @sge
+// CHECK-NEXT: Results
+// CHECK-NEXT{LITERAL}: false
+// CHECK-NEXT{LITERAL}: true
+// CHECK-NEXT{LITERAL}: true
+
+func.func @ult() -> (i1, i1, i1) {
+  %c-1 = arith.constant -1 : index
+  %c1 = arith.constant 1 : index
+  %e1 = arith.cmpi ult, %c-1, %c1 : index
+  %e2 = arith.cmpi ult, %c1, %c-1 : index
+  %e3 = arith.cmpi ult, %c1, %c1 : index
+  return %e1, %e2, %e3 : i1, i1, i1
+}
+
+// CHECK-LABEL: @ult
+// CHECK-NEXT: Results
+// CHECK-NEXT{LITERAL}: false
+// CHECK-NEXT{LITERAL}: true
+// CHECK-NEXT{LITERAL}: false
+
+func.func @ule() -> (i1, i1, i1) {
+  %c-1 = arith.constant -1 : index
+  %c1 = arith.constant 1 : index
+  %e1 = arith.cmpi ule, %c-1, %c1 : index
+  %e2 = arith.cmpi ule, %c1, %c-1 : index
+  %e3 = arith.cmpi ule, %c1, %c1 : index
+  return %e1, %e2, %e3 : i1, i1, i1
+}
+
+// CHECK-LABEL: @ule
+// CHECK-NEXT: Results
+// CHECK-NEXT{LITERAL}: false
+// CHECK-NEXT{LITERAL}: true
+// CHECK-NEXT{LITERAL}: true
+
+func.func @ugt() -> (i1, i1, i1) {
+  %c-1 = arith.constant -1 : index
+  %c1 = arith.constant 1 : index
+  %e1 = arith.cmpi ugt, %c-1, %c1 : index
+  %e2 = arith.cmpi ugt, %c1, %c-1 : index
+  %e3 = arith.cmpi ugt, %c1, %c1 : index
+  return %e1, %e2, %e3 : i1, i1, i1
+}
+
+// CHECK-LABEL: @ugt
+// CHECK-NEXT: Results
+// CHECK-NEXT{LITERAL}: true
+// CHECK-NEXT{LITERAL}: false
+// CHECK-NEXT{LITERAL}: false
+
+func.func @uge() -> (i1, i1, i1) {
+  %c-1 = arith.constant -1 : index
+  %c1 = arith.constant 1 : index
+  %e1 = arith.cmpi uge, %c-1, %c1 : index
+  %e2 = arith.cmpi uge, %c1, %c-1 : index
+  %e3 = arith.cmpi uge, %c1, %c1 : index
+  return %e1, %e2, %e3 : i1, i1, i1
+}
+
+// CHECK-LABEL: @uge
+// CHECK-NEXT: Results
+// CHECK-NEXT{LITERAL}: true
+// CHECK-NEXT{LITERAL}: false
+// CHECK-NEXT{LITERAL}: true
diff --git a/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/arith/constant.mlir b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/arith/constant.mlir
new file mode 100644
index 00000000000000..b378f008bada34
--- /dev/null
+++ b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/arith/constant.mlir
@@ -0,0 +1,37 @@
+// RUN: mlir-interpreter-runner %s -run-all | FileCheck %s
+
+func.func @tensor() -> tensor<2xi16> {
+  %cst = arith.constant dense<[42, 43]> : tensor<2xi16>
+  return %cst : tensor<2xi16>
+}
+
+// CHECK-LABEL: @tensor
+// CHECK-NEXT: Results
+// CHECK-NEXT: [42, 43]
+
+func.func @tensor_splat() -> tensor<2xi32> {
+  %cst = arith.constant dense<42> : tensor<2xi32>
+  return %cst : tensor<2xi32>
+}
+
+// CHECK-LABEL: @tensor_splat
+// CHECK-NEXT: Results
+// CHECK-NEXT: [42, 42]
+
+func.func @scalar() -> i1 {
+  %cst = arith.constant true
+  return %cst : i1
+}
+
+// CHECK-LABEL: @scalar
+// CHECK-NEXT: Results
+// CHECK-NEXT: true
+
+func.func @vector() -> vector<2x3xi32> {
+  %cst = arith.constant dense<[[1, 2, 3], [4, 5, 6]]> : vector<2x3xi32>
+  return %cst : vector<2x3xi32>
+}
+
+// CHECK-LABEL: @vector
+// CHECK-NEXT: Results
+// CHECK-NEXT{LITERAL}: vector<2x3xi32>: [[1, 2, 3], [4, 5, 6]]
diff --git a/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/arith/extf.mlir b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/arith/extf.mlir
new file mode 100644
index 00000000000000..f6124ea07c3706
--- /dev/null
+++ b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/arith/extf.mlir
@@ -0,0 +1,11 @@
+// RUN: mlir-interpreter-runner %s -run-all | FileCheck %s
+
+func.func @extf() -> f64 {
+  %c1 = arith.constant 1.0 : f32
+  %ext = arith.extf %c1 : f32 to f64
+  return %ext : f64
+}
+
+// CHECK-LABEL: @extf
+// CHECK-NEXT: Results
+// CHECK-NEXT: f64: 1.000000e+00
diff --git a/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/arith/fptosi.mlir b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/arith/fptosi.mlir
new file mode 100644
index 00000000000000..46592e49f2ed33
--- /dev/null
+++ b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/arith/fptosi.mlir
@@ -0,0 +1,21 @@
+// RUN: mlir-interpreter-runner %s -run-all | FileCheck %s
+
+func.func @i16() -> i16 {
+  %c-1 = arith.constant -1.4 : f32
+  %r = arith.fptosi %c-1 : f32 to i16
+  return %r : i16
+}
+
+// CHECK-LABEL: @i16
+// CHECK-NEXT: Results
+// CHECK-NEXT{LITERAL}: -1
+
+func.func @vector() -> vector<2xi32> {
+  %c = arith.constant dense<[-1.1, -0.9]> : vector<2xf32>
+  %r = arith.fptosi %c : vector<2xf32> to vector<2xi32>
+  return %r : vector<2xi32>
+}
+
+// CHECK-LABEL: @vector
+// CHECK-NEXT: Results
+// CHECK-NEXT{LITERAL}: vector<2xi32>: [-1, 0]
diff --git a/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/arith/index_cast.mlir b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/arith/index_cast.mlir
new file mode 100644
index 00000000000000..d47438e1992062
--- /dev/null
+++ b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/arith/index_cast.mlir
@@ -0,0 +1,28 @@
+// RUN: mlir-interpreter-runner %s -run-all | FileCheck %s
+
+func.func @i32() -> (index) {
+  %c1 = arith.constant 42 : i32
+  %index = arith.index_cast %c1 : i32 to index
+  return %index : index
+}
+
+// CHECK-LABEL: @i32
+// CHECK{LITERAL}: 42
+
+func.func @i64() -> (index) {
+  %c1 = arith.constant 43 : i64
+  %index = arith.index_cast %c1 : i64 to index
+  return %index : index
+}
+
+// CHECK-LABEL: @i64
+// CHECK{LITERAL}: 43
+
+func.func @narrowing() -> (i32) {
+  %c1 = arith.constant 0x100000001 : index
+  %i32 = arith.index_cast %c1 : index to i32
+  return %i32 : i32
+}
+
+// CHECK-LABEL: @narrowing
+// CHECK{LITERAL}: i32: 1
diff --git a/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/arith/int_math.mlir b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/arith/int_math.mlir
new file mode 100644
index 00000000000000..a1c32c64730355
--- /dev/null
+++ b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/arith/int_math.mlir
@@ -0,0 +1,111 @@
+// RUN: mlir-interpreter-runner %s -run-all | FileCheck %s
+
+func.func @addi() -> i32 {
+  %c1 = arith.constant 1 : i32
+  %c2 = arith.constant 2 : i32
+  %ret = arith.addi %c1, %c2 : i32
+  return %ret : i32
+}
+
+// CHECK-LABEL: @addi
+// CHECK-NEXT: Results
+// CHECK-NEXT: i32: 3
+
+func.func @muli() -> i32 {
+  %c3 = arith.constant 3 : i32
+  %c5 = arith.constant 5 : i32
+  %ret = arith.muli %c3, %c5 : i32
+  return %ret : i32
+}
+
+// CHECK-LABEL: @muli
+// CHECK-NEXT: Results
+// CHECK-NEXT: i32: 15
+
+func.func @divsi() -> i32 {
+  %c10 = arith.constant 10 : i32
+  %c-2 = arith.constant -2 : i32
+  %ret = arith.divsi %c10, %c-2 : i32
+  return %ret : i32
+}
+
+// CHECK-LABEL: @divsi
+// CHECK-NEXT: Results
+// CHECK-NEXT: i32: -5
+
+func.func @subi() -> i32 {
+  %c10 = arith.constant 10 : i32
+  %c3 = arith.constant 3 : i32
+  %ret = arith.subi %c10, %c3 : i32
+  return %ret : i32
+}
+
+// CHECK-LABEL: @subi
+// CHECK-NEXT: Results
+// CHECK-NEXT: i32: 7
+
+func.func @andi() -> i32 {
+  %c63 = arith.constant 63 : i32
+  %c131 = arith.constant 131 : i32
+  %ret = arith.andi %c63, %c131 : i32
+  return %ret : i32
+}
+
+// CHECK-LABEL: @andi
+// CHECK-NEXT: Results
+// CHECK-NEXT: i32: 3
+
+func.func @ori() -> i32 {
+  %c3 = arith.constant 3 : i32
+  %c10 = arith.constant 10 : i32
+  %ret = arith.ori %c3, %c10 : i32
+  return %ret : i32
+}
+
+// CHECK-LABEL: @ori
+// CHECK-NEXT: Results
+// CHECK-NEXT: i32: 11
+
+func.func @shrui() -> i32 {
+  %c20 = arith.constant 20 : i32
+  %c-1 = arith.constant -1 : i32
+  %ret = arith.shrui %c-1, %c20 : i32
+  return %ret : i32
+}
+
+// CHECK-LABEL: @shrui
+// CHECK-NEXT: Results
+// CHECK-NEXT: 4095
+
+func.func @shli() -> i32 {
+  %c2 = arith.constant 2 : i32
+  %c42 = arith.constant 42 : i32
+  %ret = arith.shli %c42, %c2 : i32
+  return %ret : i32
+}
+
+// CHECK-LABEL: @shli
+// CHECK-NEXT: Results
+// CHECK-NEXT: 168
+
+func.func @shrsi() -> i32 {
+  %c3 = arith.constant 3 : i32
+  %c-1023 = arith.constant -1023 : i32
+  %ret = arith.shrsi %c-1023, %c3 : i32
+  return %ret : i32
+}
+
+// CHECK-LABEL: @shrsi
+// CHECK-NEXT: Results
+// CHECK-NEXT: -128
+
+func.func @remsi() -> i32 {
+  %c7 = arith.constant 7 : i32
+  %c-1023 = arith.constant -1023 : i32
+  %ret = arith.remsi %c-1023, %c7 : i32
+  return %ret : i32
+}
+
+// CHECK-LABEL: @remsi
+// CHECK-NEXT: Results
+// CHECK-NEXT: -1
diff --git a/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/arith/minmax.mlir b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/arith/minmax.mlir
new file mode 100644
index 00000000000000..6d11378b5d7d31
--- /dev/null
+++ b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/arith/minmax.mlir
@@ -0,0 +1,25 @@
+// RUN: mlir-interpreter-runner %s -run-all | FileCheck %s
+
+func.func @i32() -> (i32, i32) {
+  %c-1 = arith.constant -1 : i32
+  %c1 = arith.constant 1 : i32
+  %r1 = arith.minsi %c-1, %c1 : i32
+  %r2 = arith.maxsi %c-1, %c1 : i32
+  return %r1, %r2 : i32, i32
+}
+
+// CHECK-LABEL: @i32
+// CHECK{LITERAL}: -1
+// CHECK-NEXT{LITERAL}: 1
+
+func.func @i64() -> (i64, i64) {
+  %c-1 = arith.constant -1 : i64
+  %c1 = arith.constant 1000000000000 : i64
+  %r1 = arith.minsi %c-1, %c1 : i64
+  %r2 = arith.maxsi %c-1, %c1 : i64
+  return %r1, %r2 : i64, i64
+}
+
+// CHECK-LABEL: @i64
+// CHECK{LITERAL}: -1
+// CHECK-NEXT{LITERAL}: 1000000000000
diff --git a/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/arith/negf.mlir b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/arith/negf.mlir
new file mode 100644
index 00000000000000..8b45b85a0611ba
--- /dev/null
+++ b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/arith/negf.mlir
@@ -0,0 +1,21 @@
+// RUN: mlir-interpreter-runner %s -run-all | FileCheck %s
+
+func.func @negf32() -> f32 {
+  %c = arith.constant -1.5 : f32
+  %ret = arith.negf %c : f32
+  return %ret : f32
+}
+
+// CHECK-LABEL: @negf32
+// CHECK-NEXT: Results
+// CHECK-NEXT: f32: 1.500000e+00
+
+func.func @negf64() -> f64 {
+  %c = arith.constant 3.5 : f64
+  %ret = arith.negf %c : f64
+  return %ret : f64
+}
+
+// CHECK-LABEL: @negf64
+// CHECK-NEXT: Results
+// CHECK-NEXT: f64: -3.500000e+00
diff --git a/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/arith/remf.mlir b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/arith/remf.mlir
new file mode 100644
index 00000000000000..3dc3033646cabd
--- /dev/null
+++ b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/arith/remf.mlir
@@ -0,0 +1,12 @@
+// RUN: mlir-interpreter-runner %s -run-all | FileCheck %s
+
+func.func @remf() -> f32 {
+  %a = arith.constant 3.5 : f32
+  %b = arith.constant 2.25 : f32
+  %ret = arith.remf %a, %b : f32
+  return %ret : f32
+}
+
+// CHECK-LABEL: @remf
+// CHECK-NEXT: Results
+// CHECK-NEXT: f32: 1.250000e+00
diff --git a/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/arith/select.mlir b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/arith/select.mlir
new file mode 100644
index 00000000000000..103842d4079e92
--- /dev/null
+++ b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/arith/select.mlir
@@ -0,0 +1,52 @@
+// RUN: mlir-interpreter-runner %s -run-all | FileCheck %s
+
+func.func @select() -> (i32, i32) {
+  %c-1 = arith.constant -1 : i32
+  %c1 = arith.constant 1 : i32
+  %true = arith.constant true
+  %false = arith.constant false
+  %r1 = arith.select %true, %c-1, %c1 : i32
+  %r2 = arith.select %false, %c-1, %c1 : i32
+  return %r1, %r2 : i32, i32
+}
+
+// CHECK-LABEL: @select
+// CHECK-NEXT: Results
+// CHECK-NEXT{LITERAL}: -1
+// CHECK-NEXT{LITERAL}: 1
+
+func.func @vector() -> vector<4xi32> {
+  %a = arith.constant dense<[1, 2, 3, 4]> : vector<4xi32>
+  %b = arith.constant dense<[10, 20, 30, 40]> : vector<4xi32>
+  %c = arith.constant dense<[true, false, true, false]> : vector<4xi1>
+  %r = arith.select %c, %a, %b : vector<4xi1>, vector<4xi32>
+  return %r : vector<4xi32>
+}
+
+// CHECK-LABEL: @vector
+// CHECK-NEXT: Results
+// CHECK-NEXT{LITERAL}: vector<4xi32>: [1, 20, 3, 40]
+
+func.func @scalar_vector() -> vector<4xi32> {
+  %a = arith.constant dense<[1, 2, 3, 4]> : vector<4xi32>
+  %b = arith.constant dense<[10, 20, 30, 40]> : vector<4xi32>
+  %c = arith.constant false
+  %r = arith.select %c, %a, %b : vector<4xi32>
+  return %r : vector<4xi32>
+}
+
+// CHECK-LABEL: @scalar_vector
+// CHECK-NEXT: Results
+// CHECK-NEXT{LITERAL}: vector<4xi32>: [10, 20, 30, 40]
+
+func.func @tensor() -> tensor<4xi32> {
+  %a = arith.constant dense<[1, 2, 3, 4]> : tensor<4xi32>
+  %b = arith.constant dense<[10, 20, 30, 40]> : tensor<4xi32>
+  %c = arith.constant dense<[true, false, true, false]> : tensor<4xi1>
+  %r = arith.select %c, %a, %b : tensor<4xi1>, tensor<4xi32>
+  return %r : tensor<4xi32>
+}
+
+// CHECK-LABEL: @tensor
+// CHECK-NEXT: Results
+// CHECK-NEXT{LITERAL}: TensorOrMemref<4xi32>: [1, 20, 3, 40]
diff --git a/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/arith/sitofp.mlir b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/arith/sitofp.mlir
new file mode 100644
index 00000000000000..b2a18564d0fcd3
--- /dev/null
+++ b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/arith/sitofp.mlir
@@ -0,0 +1,31 @@
+// RUN: mlir-interpreter-runner %s -run-all | FileCheck %s
+
+func.func @i16() -> f32 {
+  %c-1 = arith.constant -1 : i16
+  %r = arith.sitofp %c-1 : i16 to f32
+  return %r : f32
+}
+
+// CHECK-LABEL: @i16
+// CHECK-NEXT: Results
+// CHECK-NEXT{LITERAL}: -1.000000e+00
+
+func.func @i1() -> f64 {
+  %true = arith.constant true
+  %r = arith.sitofp %true : i1 to f64
+  return %r : f64
+}
+
+// CHECK-LABEL: @i1
+// CHECK-NEXT: Results
+// CHECK-NEXT{LITERAL}: 1.000000e+00
+
+func.func @vector() -> vector<1xf32> {
+  %c-1 = arith.constant dense<-1> : vector<1xi8>
+  %r = arith.sitofp %c-1 : vector<1xi8> to vector<1xf32>
+  return %r : vector<1xf32>
+}
+
+// CHECK-LABEL: @vector
+// CHECK-NEXT: Results
+// CHECK-NEXT{LITERAL}: vector<1xf32>: [-1.000000e+00]
diff --git a/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/arith/uitofp.mlir b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/arith/uitofp.mlir
new file mode 100644
index 00000000000000..990e0196041858
--- /dev/null
+++ b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/arith/uitofp.mlir
@@ -0,0 +1,31 @@
+// RUN: mlir-interpreter-runner %s -run-all | FileCheck %s
+
+func.func @i16() -> f32 {
+  %c-1 = arith.constant -1 : i16
+  %r = arith.uitofp %c-1 : i16 to f32
+  return %r : f32
+}
+
+// CHECK-LABEL: @i16
+// CHECK-NEXT: Results
+// CHECK-NEXT{LITERAL}: 6.553500e+04
+
+func.func @i1() -> f64 {
+  %true = arith.constant true
+  %r = arith.uitofp %true : i1 to f64
+  return %r : f64
+}
+
+// CHECK-LABEL: @i1
+// CHECK-NEXT: Results
+// CHECK-NEXT{LITERAL}: 1.000000e+00
+
+func.func @vector() -> vector<1xf32> {
+  %c-1 = arith.constant dense<-1> : vector<1xi8>
+  %r = arith.uitofp %c-1 : vector<1xi8> to vector<1xf32>
+  return %r : vector<1xf32>
+}
+
+// CHECK-LABEL: @vector
+// CHECK-NEXT: Results
+// CHECK-NEXT{LITERAL}: vector<1xf32>: [2.550000e+02]
diff --git a/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/arith/vector_math.mlir b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/arith/vector_math.mlir
new file mode 100644
index 00000000000000..4ffc4d0ed81961
--- /dev/null
+++ b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/arith/vector_math.mlir
@@ -0,0 +1,12 @@
+// RUN: mlir-interpreter-runner %s -run-all | FileCheck %s
+
+func.func @addi() -> vector<2xi32> {
+  %c1 = arith.constant dense<[1, 2]> : vector<2xi32>
+  %c2 = arith.constant dense<[3, 4]> : vector<2xi32>
+  %ret = arith.addi %c1, %c2 : vector<2xi32>
+  return %ret : vector<2xi32>
+}
+
+// CHECK-LABEL: @addi
+// CHECK-NEXT: Results
+// CHECK-NEXT: vector<2xi32>: [4, 6]
diff --git a/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/bufferization/alloc_tensor.mlir b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/bufferization/alloc_tensor.mlir
new file mode 100644
index 00000000000000..8d47f63656a05f
--- /dev/null
+++ b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/bufferization/alloc_tensor.mlir
@@ -0,0 +1,30 @@
+// RUN: mlir-interpreter-runner %s -run-all | FileCheck %s
+
+func.func @static() -> tensor<1x2x3xi32> {
+  %t = bufferization.alloc_tensor() : tensor<1x2x3xi32>
+  return %t : tensor<1x2x3xi32>
+}
+
+// CHECK-LABEL: @static
+// CHECK-NEXT: Results
+// CHECK-NEXT{LITERAL}: [[[0, 0, 0], [0, 0, 0]]]
+
+func.func @dynamic() -> tensor<?x1xi32> {
+  %c4 = arith.constant 4 : index
+  %t = bufferization.alloc_tensor(%c4) : tensor<?x1xi32>
+  return %t : tensor<?x1xi32>
+}
+
+// CHECK-LABEL: @dynamic
+// CHECK-NEXT: Results
+// CHECK-NEXT{LITERAL}: [[0], [0], [0], [0]]
+
+func.func @copy() -> tensor<i32> {
+  %c = arith.constant dense<123> : tensor<i32>
+  %t = bufferization.alloc_tensor() copy(%c) : tensor<i32>
+  return %t : tensor<i32>
+}
+
+// CHECK-LABEL: @copy
+// CHECK-NEXT: Results
+// CHECK-NEXT: 123
diff --git a/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/bufferization/clone.mlir b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/bufferization/clone.mlir
new file mode 100644
index 00000000000000..09670eb52c81d6
--- /dev/null
+++ b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/bufferization/clone.mlir
@@ -0,0 +1,14 @@
+// RUN: mlir-interpreter-runner %s -run-all | FileCheck %s
+
+func.func @clone() -> (memref<i32>, memref<i32>) {
+  %a = arith.constant dense<1> : memref<i32>
+  %b = bufferization.clone %a : memref<i32> to memref<i32>
+  %c = arith.constant 2 : i32
+  memref.store %c, %b[] : memref<i32>
+  return %a, %b : memref<i32>, memref<i32>
+}
+
+// CHECK-LABEL: @clone
+// CHECK-NEXT: Results
+// CHECK-NEXT: TensorOrMemref<i32>: 1
+// CHECK-NEXT: TensorOrMemref<i32>: 2
diff --git a/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/bufferization/to_memref.mlir b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/bufferization/to_memref.mlir
new file mode 100644
index 00000000000000..2ed938c36c6ece
--- /dev/null
+++ b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/bufferization/to_memref.mlir
@@ -0,0 +1,10 @@
+// RUN: mlir-interpreter-runner %s -run-all | FileCheck %s
+
+func.func @memref() -> memref<2xi16> {
+  %cst = arith.constant dense<[42, 43]> : tensor<2xi16>
+  %memref = bufferization.to_memref %cst : memref<2xi16>
+  return %memref : memref<2xi16>
+}
+
+// CHECK-LABEL: @memref
+// CHECK{LITERAL}: [42, 43]
diff --git a/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/bufferization/to_tensor.mlir b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/bufferization/to_tensor.mlir
new file mode 100644
index 00000000000000..8bc19b4d895bff
--- /dev/null
+++ b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/bufferization/to_tensor.mlir
@@ -0,0 +1,11 @@
+// RUN: mlir-interpreter-runner %s -run-all | FileCheck %s
+
+func.func @tensor() -> tensor<2xi16> {
+  %cst = arith.constant dense<[43, 44]> : tensor<2xi16>
+  %memref = bufferization.to_memref %cst : memref<2xi16>
+  %tensor = bufferization.to_tensor %memref : memref<2xi16>
+  return %tensor : tensor<2xi16>
+}
+
+// CHECK-LABEL: @tensor
+// CHECK{LITERAL}: [43, 44]
diff --git a/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/builtin/unrealized_conversion_cast.mlir b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/builtin/unrealized_conversion_cast.mlir
new file mode 100644
index 00000000000000..55c481677eb4ad
--- /dev/null
+++ b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/builtin/unrealized_conversion_cast.mlir
@@ -0,0 +1,21 @@
+// RUN: mlir-interpreter-runner %s -run-all | FileCheck %s
+
+func.func @no_op_cast() -> i32 {
+  %cst = arith.constant 42 : i32
+  %cast = builtin.unrealized_conversion_cast %cst : i32 to i32
+  return %cast : i32
+}
+
+// CHECK-LABEL: @no_op_cast
+// CHECK-NEXT: Results
+// CHECK{LITERAL}: 42
+
+func.func @cast_to_dynamic() -> tensor<?xi32> {
+  %cst = arith.constant dense<[0, 1, 2]> : tensor<3xi32>
+  %cast = builtin.unrealized_conversion_cast %cst : tensor<3xi32> to tensor<?xi32>
+  return %cast : tensor<?xi32>
+}
+
+// CHECK-LABEL: @cast_to_dynamic
+// CHECK-NEXT: Results
+// CHECK{LITERAL}: [0, 1, 2]
diff --git a/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/complex/complex.mlir b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/complex/complex.mlir
new file mode 100644
index 00000000000000..590ce8fb2916a7
--- /dev/null
+++ b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/complex/complex.mlir
@@ -0,0 +1,186 @@
+// RUN: mlir-interpreter-runner %s -run-all | FileCheck %s
+
+func.func @abs() -> f32 {
+  %c = complex.constant [1.0 : f32, 2.0 : f32] : complex<f32>
+  %ret = complex.abs %c : complex<f32>
+  return %ret : f32
+}
+
+// CHECK-LABEL: @abs
+// CHECK-NEXT: Results
+// CHECK-NEXT: 2.236068e+00
+
+func.func @add() -> complex<f32> {
+  %a = complex.constant [1.0 : f32, 2.0 : f32] : complex<f32>
+  %b = complex.constant [10.0 : f32, 200.0 : f32] : complex<f32>
+  %ret = complex.add %a, %b : complex<f32>
+  return %ret : complex<f32>
+}
+
+// CHECK-LABEL: @add
+// CHECK-NEXT: Results
+// CHECK-NEXT: 1.100000e+01+2.020000e+02i
+
+func.func @cos() -> complex<f32> {
+  %a = complex.constant [1.0 : f32, 2.0 : f32] : complex<f32>
+  %ret = complex.cos %a : complex<f32>
+  return %ret : complex<f32>
+}
+
+// CHECK-LABEL: @cos
+// CHECK-NEXT: Results
+// CHECK-NEXT: 2.032723e+00-3.051898e+00i
+
+func.func @create() -> complex<f32> {
+  %c1 = arith.constant 1.0 : f32
+  %c2 = arith.constant 2.0 : f32
+  %ret = complex.create %c1, %c2 : complex<f32>
+  return %ret : complex<f32>
+}
+
+// CHECK-LABEL: @create
+// CHECK-NEXT: Results
+// CHECK-NEXT: 1.000000e+00+2.000000e+00i
+
+func.func @div() -> complex<f32> {
+  %a = complex.constant [1.0 : f32, 2.0 : f32] : complex<f32>
+  %b = complex.constant [3.0 : f32, 4.0 : f32] : complex<f32>
+  %ret = complex.div %a, %b : complex<f32>
+  return %ret : complex<f32>
+}
+
+// CHECK-LABEL: @div
+// CHECK-NEXT: Results
+// CHECK-NEXT: 4.400000e-01+8.000000e-02i
+
+func.func @exp() -> complex<f32> {
+  %a = complex.constant [1.0 : f32, 2.0 : f32] : complex<f32>
+  %ret = complex.exp %a : complex<f32>
+  return %ret : complex<f32>
+}
+
+// CHECK-LABEL: @exp
+// CHECK-NEXT: Results
+// CHECK-NEXT: -1.131204e+00+2.471727e+00i
+
+func.func @expm1() -> complex<f32> {
+  %a = complex.constant [1.0e-06 : f32, 1.0e-06 : f32] : complex<f32>
+  %ret = complex.expm1 %a : complex<f32>
+  return %ret : complex<f32>
+}
+
+// CHECK-LABEL: @expm1
+// CHECK-NEXT: Results
+// CHECK-NEXT: 1.000000e-06+1.000001e-06i
+
+func.func @im() -> f32 {
+  %a = complex.constant [1.0 : f32, 2.0 : f32] : complex<f32>
+  %ret = complex.im %a : complex<f32>
+  return %ret : f32
+}
+
+// CHECK-LABEL: @im
+// CHECK-NEXT: Results
+// CHECK-NEXT: 2.000000e+00
+
+func.func @log() -> complex<f32> {
+  %a = complex.constant [1.0 : f32, 2.0 : f32] : complex<f32>
+  %ret = complex.log %a : complex<f32>
+  return %ret : complex<f32>
+}
+
+// CHECK-LABEL: @log
+// CHECK-NEXT: Results
+// CHECK-NEXT: 8.047190e-01+1.107149e+00i
+
+func.func @log1p() -> complex<f32> {
+  %a = complex.constant [1.0e-07 : f32, 1.0e-20 : f32] : complex<f32>
+  %ret = complex.log1p %a : complex<f32>
+  return %ret : complex<f32>
+}
+
+// CHECK-LABEL: @log1p
+// CHECK-NEXT: Results
+// CHECK-NEXT: 1.192093e-07+9.999999e-21i
+
+func.func @mul() -> complex<f32> {
+  %a = complex.constant [1.0 : f32, 2.0 : f32] : complex<f32>
+  %b = complex.constant [3.0 : f32, 4.0 : f32] : complex<f32>
+  %ret = complex.mul %a, %b : complex<f32>
+  return %ret : complex<f32>
+}
+
+// CHECK-LABEL: @mul
+// CHECK-NEXT: Results
+// CHECK-NEXT: -5.000000e+00+1.000000e+01i
+
+func.func @neg() -> complex<f32> {
+  %a = complex.constant [1.0 : f32, 2.0 : f32] : complex<f32>
+  %ret = complex.neg %a : complex<f32>
+  return %ret : complex<f32>
+}
+
+// CHECK-LABEL: @neg
+// CHECK-NEXT: Results
+// CHECK-NEXT: -1.000000e+00-2.000000e+00i
+
+func.func @pow() -> complex<f32> {
+  %a = complex.constant [1.0 : f32, 2.0 : f32] : complex<f32>
+  %b = complex.constant [3.0 : f32, 4.0 : f32] : complex<f32>
+  %ret = complex.pow %a, %b : complex<f32>
+  return %ret : complex<f32>
+}
+
+// CHECK-LABEL: @pow
+// CHECK-NEXT: Results
+// CHECK-NEXT: 1.290096e-01+3.392413e-02i
+
+func.func @re() -> f32 {
+  %a = complex.constant [1.0 : f32, 2.0 : f32] : complex<f32>
+  %ret = complex.re %a : complex<f32>
+  return %ret : f32
+}
+
+// CHECK-LABEL: @re
+// CHECK-NEXT: Results
+// CHECK-NEXT: 1.000000e+00
+
+func.func @rsqrt() -> complex<f32> {
+  %a = complex.constant [1.0 : f32, 2.0 : f32] : complex<f32>
+  %ret = complex.rsqrt %a : complex<f32>
+  return %ret : complex<f32>
+}
+
+// CHECK-LABEL: @rsqrt
+// CHECK-NEXT: Results
+// CHECK-NEXT: 5.688645e-01-3.515776e-01i
+
+func.func @sin() -> complex<f64> {
+  %a = complex.constant [1.0 : f64, 2.0 : f64] : complex<f64>
+  %ret = complex.sin %a : complex<f64>
+  return %ret : complex<f64>
+}
+
+// CHECK-LABEL: @sin
+// CHECK-NEXT: Results
+// CHECK-NEXT: 3.165779e+00+1.959601e+00i
+
+func.func @sqrt() -> complex<f32> {
+  %a = complex.constant [1.0 : f32, 2.0 : f32] : complex<f32>
+  %ret = complex.sqrt %a : complex<f32>
+  return %ret : complex<f32>
+}
+
+// CHECK-LABEL: @sqrt
+// CHECK-NEXT: Results
+// CHECK-NEXT: 1.27202{{.*}}e+00+7.86151{{.*}}e-01i
+
+func.func @tanh() -> complex<f32> {
+  %a = complex.constant [1.0 : f32, 2.0 : f32] : complex<f32>
+  %ret = complex.tanh %a : complex<f32>
+  return %ret : complex<f32>
+}
+
+// CHECK-LABEL: @tanh
+// CHECK-NEXT: Results
+// CHECK-NEXT: 1.166736e+00-2.434582e-01i
diff --git a/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/func/call.mlir b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/func/call.mlir
new file mode 100644
index 00000000000000..b23f9d7423b05b
--- /dev/null
+++ b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/func/call.mlir
@@ -0,0 +1,48 @@
+// RUN: mlir-interpreter-runner %s -run-all | FileCheck %s
+
+func.func @cbrtf_caller() -> f32 {
+  %c-27 = arith.constant -27.0 : f32
+  %ret = func.call @cbrtf(%c-27) : (f32) -> f32
+  func.return %ret : f32
+}
+
+// CHECK-LABEL: @cbrtf_caller
+// CHECK-NEXT: Results
+// CHECK-NEXT: -3.000000e+00
+
+func.func @cbrt_caller() -> f64 {
+  %c-8 = arith.constant -8.0 : f64
+  %ret = func.call @cbrt(%c-8) : (f64) -> f64
+  func.return %ret : f64
+}
+
+// CHECK-LABEL: @cbrt_caller
+// CHECK-NEXT: Results
+// CHECK-NEXT: -2.000000e+00
+
+func.func @atan2f_caller() -> f32 {
+  %c1 = arith.constant 1.0 : f32
+  %c2 = arith.constant 2.0 : f32
+  %ret = func.call @atan2f(%c1, %c2) : (f32, f32) -> f32
+  func.return %ret : f32
+}
+
+// CHECK-LABEL: @atan2f_caller
+// CHECK-NEXT: Results
+// CHECK-NEXT: 4.636476e-01
+
+func.func @atan2_caller() -> f64 {
+  %c2 = arith.constant 2.0 : f64
+  %c3 = arith.constant 3.0 : f64
+  %ret = func.call @atan2(%c2, %c3) : (f64, f64) -> f64
+  func.return %ret : f64
+}
+
+// CHECK-LABEL: @atan2_caller
+// CHECK-NEXT: Results
+// CHECK-NEXT: 5.880026e-01
+
+func.func private @cbrtf(%a: f32) -> f32
+func.func private @cbrt(%a: f64) -> f64
+func.func private @atan2f(%a: f32, %b: f32) -> f32
+func.func private @atan2(%a: f64, %b: f64) -> f64
diff --git a/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/linalg/broadcast.mlir b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/linalg/broadcast.mlir
new file mode 100644
index 00000000000000..56fa8396a6391a
--- /dev/null
+++ b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/linalg/broadcast.mlir
@@ -0,0 +1,30 @@
+// RUN: mlir-interpreter-runner %s -run-all | FileCheck %s
+
+func.func @broadcast() -> (tensor<2x3xi32>, tensor<2x3xi32>) {
+  %v = arith.constant dense<[1,2]> : tensor<2xi32>
+  %init = tensor.empty() : tensor<2x3xi32>
+  %bcast = linalg.broadcast
+      ins(%v: tensor<2xi32>)
+      outs(%init: tensor<2x3xi32>)
+      dimensions = [1]
+  func.return %init, %bcast : tensor<2x3xi32>, tensor<2x3xi32>
+}
+
+// CHECK-LABEL: @broadcast
+// CHECK-NEXT: Results
+// CHECK-NEXT{LITERAL}: [[0, 0, 0], [0, 0, 0]]
+// CHECK-NEXT{LITERAL}: [[1, 1, 1], [2, 2, 2]]
+
+func.func @bufferized() -> memref<2x3xi32> {
+  %v = arith.constant dense<[1,2]> : memref<2xi32>
+  %alloc = memref.alloc() : memref<2x3xi32>
+  linalg.broadcast
+      ins(%v: memref<2xi32>)
+      outs(%alloc: memref<2x3xi32>)
+      dimensions = [1]
+  func.return %alloc : memref<2x3xi32>
+}
+
+// CHECK-LABEL: @bufferized
+// CHECK-NEXT: Results
+// CHECK-NEXT{LITERAL}: [[1, 1, 1], [2, 2, 2]]
diff --git a/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/linalg/dot.mlir b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/linalg/dot.mlir
new file mode 100644
index 00000000000000..895f51e0b0dd93
--- /dev/null
+++ b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/linalg/dot.mlir
@@ -0,0 +1,14 @@
+// RUN: mlir-interpreter-runner %s -run-all | FileCheck %s
+
+func.func @dot() -> tensor<i32> {
+  %lhs = arith.constant dense<[1, 2]> : tensor<2xi32>
+  %rhs = arith.constant dense<[3, 4]> : tensor<2xi32>
+  %init = tensor.empty() : tensor<i32>
+  %ret = linalg.dot ins(%lhs, %rhs: tensor<2xi32>, tensor<2xi32>)
+                    outs(%init: tensor<i32>) -> tensor<i32>
+  return %ret : tensor<i32>
+}
+
+// CHECK-LABEL: @dot
+// CHECK-NEXT: Results
+// CHECK-NEXT{LITERAL}: 11
diff --git a/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/linalg/fill.mlir b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/linalg/fill.mlir
new file mode 100644
index 00000000000000..20a91a7768df6b
--- /dev/null
+++ b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/linalg/fill.mlir
@@ -0,0 +1,24 @@
+// RUN: mlir-interpreter-runner %s -run-all | FileCheck %s
+
+func.func @fill() -> (tensor<2xi32>, tensor<2xi32>) {
+  %c42 = arith.constant 42 : i32
+  %init = tensor.empty() : tensor<2xi32>
+  %fill = linalg.fill ins(%c42 : i32) outs(%init : tensor<2xi32>) -> tensor<2xi32>
+  func.return %init, %fill : tensor<2xi32>, tensor<2xi32>
+}
+
+// CHECK-LABEL: @fill
+// CHECK-NEXT: Results
+// CHECK-NEXT{LITERAL}: [0, 0]
+// CHECK-NEXT{LITERAL}: [42, 42]
+
+func.func @bufferized() -> memref<2xi32> {
+  %c42 = arith.constant 42 : i32
+  %alloc = memref.alloc() : memref<2xi32>
+  linalg.fill ins(%c42 : i32) outs(%alloc : memref<2xi32>)
+  func.return %alloc : memref<2xi32>
+}
+
+// CHECK-LABEL: @bufferized
+// CHECK-NEXT: Results
+// CHECK-NEXT{LITERAL}: [42, 42]
diff --git a/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/linalg/generic.mlir b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/linalg/generic.mlir
new file mode 100644
index 00000000000000..c27c765bf3337e
--- /dev/null
+++ b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/linalg/generic.mlir
@@ -0,0 +1,113 @@
+// RUN: mlir-interpreter-runner %s -run-all | FileCheck %s
+
+#matmul_trait = {
+  indexing_maps = [
+    affine_map<(m, n, k) -> (m, k)>,
+    affine_map<(m, n, k) -> (k, n)>,
+    affine_map<(m, n, k) -> (m, n)>
+  ],
+  iterator_types = ["parallel", "parallel", "reduction"]
+}
+
+func.func @matmul() -> (tensor<2x2xi32>, tensor<2x2xi32>) {
+  %lhs = arith.constant dense<[[1, 2, 3], [4, 5, 6]]> : tensor<2x3xi32>
+  %rhs = arith.constant dense<[[1, 2], [3, 4], [5, 6]]> : tensor<3x2xi32>
+  %init = tensor.empty() : tensor<2x2xi32>
+  %ret = linalg.generic #matmul_trait
+    ins(%lhs, %rhs : tensor<2x3xi32>, tensor<3x2xi32>)
+    outs(%init : tensor<2x2xi32>) {
+    ^bb(%a: i32, %b: i32, %c: i32):
+      %d = arith.muli %a, %b: i32
+      %e = arith.addi %c, %d: i32
+      linalg.yield %e : i32
+    } -> tensor<2x2xi32>
+  return %ret, %init : tensor<2x2xi32>, tensor<2x2xi32>
+}
+
+// CHECK-LABEL: @matmul
+// CHECK-NEXT: Results
+// CHECK-NEXT{LITERAL}: [[22, 28], [49, 64]]
+// CHECK-NEXT{LITERAL}: [[0, 0], [0, 0]]
+
+func.func @bufferized() -> memref<2x2xi32> {
+  %lhs = arith.constant dense<[[1, 2, 3], [4, 5, 6]]> : memref<2x3xi32>
+  %rhs = arith.constant dense<[[2, 1], [3, 4], [5, 6]]> : memref<3x2xi32>
+  %alloc = memref.alloc() : memref<2x2xi32>
+  linalg.generic #matmul_trait
+    ins(%lhs, %rhs : memref<2x3xi32>, memref<3x2xi32>)
+    outs(%alloc : memref<2x2xi32>) {
+    ^bb(%a: i32, %b: i32, %c: i32):
+      %d = arith.muli %a, %b: i32
+      %e = arith.addi %c, %d: i32
+      linalg.yield %e : i32
+    }
+  return %alloc : memref<2x2xi32>
+}
+
+// CHECK-LABEL: @bufferized
+// CHECK-NEXT: Results
+// CHECK-NEXT{LITERAL}: [[23, 27], [53, 60]]
+
+#map = affine_map<(d0) -> (d0)>
+
+func.func @vector() -> tensor<4xvector<2xi32>> {
+  %c0 = arith.constant 0 : index
+  %c1 = arith.constant 1 : index
+  %lhs = arith.constant dense<[1, 2, 3, 4]> : tensor<4xi32>
+  %rhs = arith.constant dense<[5, 6, 7, 8]> : tensor<4xi32>
+  %init = tensor.empty() : tensor<4xvector<2xi32>>
+  %ret = linalg.generic {
+      indexing_maps = [#map, #map, #map],
+      iterator_types = ["parallel"]
+    }
+    ins(%lhs, %rhs : tensor<4xi32>, tensor<4xi32>)
+    outs(%init : tensor<4xvector<2xi32>>) {
+    ^bb(%a: i32, %b: i32, %c: vector<2xi32>):
+      %d = vector.insertelement %a, %c[%c0 : index] : vector<2xi32>
+      %e = vector.insertelement %b, %d[%c1 : index] : vector<2xi32>
+      linalg.yield %e : vector<2xi32>
+    } -> tensor<4xvector<2xi32>>
+  return %ret : tensor<4xvector<2xi32>>
+}
+
+// CHECK-LABEL: @vector
+// CHECK-NEXT: Results
+// CHECK-NEXT{LITERAL}: TensorOrMemref<4xvector<2xi32>>: [[1, 5], [2, 6], [3, 7], [4, 8]]
+
+func.func @matmul_dynamic() -> tensor<2x2xi32> {
+  %lhs = arith.constant dense<[[1, 2, 3], [4, 5, 6]]> : tensor<2x3xi32>
+  %rhs = arith.constant dense<[[1, 2], [3, 4], [5, 6]]> : tensor<3x2xi32>
+  %lhs_cast = tensor.cast %lhs : tensor<2x3xi32> to tensor<2x?xi32>
+  %rhs_cast = tensor.cast %rhs : tensor<3x2xi32> to tensor<?x2xi32>
+  %init = tensor.empty() : tensor<2x2xi32>
+  %ret = linalg.generic #matmul_trait
+    ins(%lhs_cast, %rhs_cast : tensor<2x?xi32>, tensor<?x2xi32>)
+    outs(%init : tensor<2x2xi32>) {
+    ^bb(%a: i32, %b: i32, %c: i32):
+      %d = arith.muli %a, %b: i32
+      %e = arith.addi %c, %d: i32
+      linalg.yield %e : i32
+    } -> tensor<2x2xi32>
+  return %ret : tensor<2x2xi32>
+}
+
+// CHECK-LABEL: @matmul_dynamic
+// CHECK-NEXT: Results
+// CHECK-NEXT{LITERAL}: [[22, 28], [49, 64]]
+
+func.func @dynamic_generic_w_cst() -> tensor<4x?xf64> {
+  %cst_1 =  arith.constant 123.456 : f64
+  %extracted_slice_ = arith.constant dense<[[1.1, 2.2], [3.3, 4.4], [5.5, 6.6], [7.7, 8.8]]> : tensor<4x2xf64>
+  %extracted_slice = tensor.cast %extracted_slice_ : tensor<4x2xf64> to tensor<4x?xf64>
+  %6 = linalg.generic { indexing_maps = [affine_map<(d0, d1) -> ()>,
+      affine_map<(d0, d1) -> (d0, d1)>], iterator_types = ["parallel",
+      "parallel"]} ins(%cst_1 : f64) outs(%extracted_slice : tensor<4x?xf64>) {
+  ^bb0(%in: f64, %out: f64):
+    linalg.yield %in : f64
+  } -> tensor<4x?xf64>
+  return %6 : tensor<4x?xf64>
+}
+
+// CHECK-LABEL: @dynamic_generic_w_cst
+// CHECK-NEXT: Results
+// CHECK-NEXT{LITERAL}: TensorOrMemref<4x2xf64>: [[1.234560e+02, 1.234560e+02], [1.234560e+02, 1.234560e+02], [1.234560e+02, 1.234560e+02], [1.234560e+02, 1.234560e+02]]
diff --git a/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/linalg/map.mlir b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/linalg/map.mlir
new file mode 100644
index 00000000000000..4fef179d7989d6
--- /dev/null
+++ b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/linalg/map.mlir
@@ -0,0 +1,74 @@
+// RUN: mlir-interpreter-runner %s -run-all | FileCheck %s
+
+func.func @no_inputs() -> tensor<4xf32> {
+  %init = arith.constant dense<[1.0,2.0,3.0,4.0]> : tensor<4xf32>
+  %zero = linalg.map outs(%init:tensor<4xf32>)() {
+    %0 = arith.constant 0.0: f32
+    linalg.yield %0: f32
+  }
+  func.return %zero : tensor<4xf32>
+}
+
+// CHECK-LABEL: @no_inputs
+// CHECK-NEXT: Results
+// CHECK-NEXT{LITERAL}: [0.000000e+00, 0.000000e+00, 0.000000e+00, 0.000000e+00]
+
+func.func @binary() -> tensor<4xi32> {
+  %init = tensor.empty() : tensor<4xi32>
+  %lhs = arith.constant dense<[0, 1, 2, 3]> : tensor<4xi32>
+  %rhs = arith.constant dense<[10, 20, 30, 40]> : tensor<4xi32>
+  %add = linalg.map ins(%lhs, %rhs: tensor<4xi32>, tensor<4xi32>)
+                    outs(%init: tensor<4xi32>)
+    (%lhs_elem: i32, %rhs_elem: i32) {
+      %0 = arith.addi %lhs_elem, %rhs_elem: i32
+      linalg.yield %0: i32
+    }
+  func.return %add : tensor<4xi32>
+}
+
+// CHECK-LABEL: @binary
+// CHECK-NEXT: Results
+// CHECK-NEXT{LITERAL}: [10, 21, 32, 43]
+
+func.func @memref() -> memref<4xi32> {
+  %alloc = memref.alloc() : memref<4xi32>
+  %lhs = arith.constant dense<[0, 1, 2, 3]> : memref<4xi32>
+  %rhs = arith.constant dense<[10, 20, 30, 40]> : memref<4xi32>
+  linalg.map ins(%lhs, %rhs: memref<4xi32>, memref<4xi32>)
+             outs(%alloc: memref<4xi32>)
+    (%lhs_elem: i32, %rhs_elem: i32) {
+      %0 = arith.muli %lhs_elem, %rhs_elem: i32
+      linalg.yield %0: i32
+    }
+  func.return %alloc : memref<4xi32>
+}
+
+// CHECK-LABEL: @memref
+// CHECK-NEXT: Results
+// CHECK-NEXT{LITERAL}: [0, 20, 60, 120]
+
+func.func @index() -> memref<4xindex> {
+  %alloc = memref.alloc() : memref<4xindex>
+  linalg.map outs(%alloc: memref<4xindex>)() {
+    %0 = linalg.index 0 : index
+    linalg.yield %0: index
+  }
+  func.return %alloc : memref<4xindex>
+}
+
+// CHECK-LABEL: @index
+// CHECK-NEXT: Results
+// CHECK-NEXT{LITERAL}: [0, 1, 2, 3]
+
+func.func @vector() -> memref<4xvector<2xindex>> {
+  %c = arith.constant dense<42> : vector<2xindex>
+  %alloc = memref.alloc() : memref<4xvector<2xindex>>
+  linalg.map outs(%alloc: memref<4xvector<2xindex>>)() {
+    linalg.yield %c: vector<2xindex>
+  }
+  func.return %alloc : memref<4xvector<2xindex>>
+}
+
+// CHECK-LABEL: @vector
+// CHECK-NEXT: Results
+// CHECK-NEXT{LITERAL}: [[42, 42], [42, 42], [42, 42], [42, 42]]
diff --git a/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/linalg/matmul.mlir b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/linalg/matmul.mlir
new file mode 100644
index 00000000000000..131d2db123792b
--- /dev/null
+++ b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/linalg/matmul.mlir
@@ -0,0 +1,41 @@
+// RUN: mlir-interpreter-runner %s -run-all | FileCheck %s
+
+func.func @m1x1() -> (tensor<1x1xi32>, tensor<1x1xi32>) {
+  %a = arith.constant dense<1> : tensor<1x1xi32>
+  %b = arith.constant dense<2> : tensor<1x1xi32>
+  %c = arith.constant dense<3> : tensor<1x1xi32>
+  %ret = linalg.matmul ins(%a, %b : tensor<1x1xi32>, tensor<1x1xi32>)
+                       outs(%c : tensor<1x1xi32>) -> tensor<1x1xi32>
+  return %c, %ret : tensor<1x1xi32>, tensor<1x1xi32>
+}
+
+// CHECK-LABEL: @m1x1
+// CHECK-NEXT: Results
+// CHECK-NEXT{LITERAL}: [[3]]
+// CHECK-NEXT{LITERAL}: [[5]]
+
+func.func @m2x2() -> tensor<2x2xi32> {
+  %a = arith.constant dense<[[1, 2], [3, 4]]> : tensor<2x2xi32>
+  %b = arith.constant dense<[[4, 5], [6, 7]]> : tensor<2x2xi32>
+  %c = tensor.empty() : tensor<2x2xi32>
+  %ret = linalg.matmul ins(%a, %b : tensor<2x2xi32>, tensor<2x2xi32>)
+                       outs(%c : tensor<2x2xi32>) -> tensor<2x2xi32>
+  return %ret : tensor<2x2xi32>
+}
+
+// CHECK-LABEL: @m2x2
+// CHECK-NEXT: Results
+// CHECK-NEXT{LITERAL}: [[16, 19], [36, 43]]
+
+func.func @m1x1_bufferized() -> memref<1x1xi32> {
+  %a = arith.constant dense<1> : memref<1x1xi32>
+  %b = arith.constant dense<2> : memref<1x1xi32>
+  %c = arith.constant dense<3> : memref<1x1xi32>
+  linalg.matmul ins(%a, %b : memref<1x1xi32>, memref<1x1xi32>)
+                outs(%c : memref<1x1xi32>)
+  return %c : memref<1x1xi32>
+}
+
+// CHECK-LABEL: @m1x1_bufferized
+// CHECK-NEXT: Results
+// CHECK-NEXT{LITERAL}: [[5]]
diff --git a/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/linalg/reduce.mlir b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/linalg/reduce.mlir
new file mode 100644
index 00000000000000..9b2e2d760369e2
--- /dev/null
+++ b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/linalg/reduce.mlir
@@ -0,0 +1,57 @@
+// RUN: mlir-interpreter-runner %s -run-all | FileCheck %s
+
+func.func @reduce() -> tensor<2xi32> {
+  %v = arith.constant dense<[[1,2,3,4], [5,6,7,8]]> : tensor<2x4xi32>
+  %init = arith.constant dense<[9, 10]> : tensor<2xi32>
+  %ret = linalg.reduce ins(%v : tensor<2x4xi32>)
+                       outs(%init: tensor<2xi32>)
+                       dimensions = [1]
+                       (%in: i32, %out: i32) {
+                         %sum = arith.addi %in, %out : i32
+                         linalg.yield %sum: i32
+                       }
+  func.return %ret : tensor<2xi32>
+}
+
+// CHECK-LABEL: @reduce
+// CHECK-NEXT: Results
+// CHECK-NEXT: [19, 36]
+
+func.func @variadic() -> (tensor<2xi32>, tensor<2xf32>) {
+  %v = arith.constant dense<[[1,2,3,4], [5,6,7,8]]> : tensor<2x4xi32>
+  %w = arith.constant dense<[[1.0,2.0,3.0,4.0], [5.0,6.0,7.0,8.0]]> : tensor<2x4xf32>
+  %init = arith.constant dense<[9, 10]> : tensor<2xi32>
+  %init2 = arith.constant dense<[9.0, 10.0]> : tensor<2xf32>
+  %ret, %retf = linalg.reduce ins(%v, %w : tensor<2x4xi32>, tensor<2x4xf32>)
+                              outs(%init, %init2: tensor<2xi32>, tensor<2xf32>)
+                              dimensions = [1]
+    (%in: i32, %inf: f32, %out: i32, %outf: f32) {
+      %sum = arith.addi %in, %out : i32
+      %sumf = arith.addf %inf, %outf : f32
+      linalg.yield %sum, %sumf: i32, f32
+    }
+  func.return %ret, %retf : tensor<2xi32>, tensor<2xf32>
+}
+
+// CHECK-LABEL: @variadic
+// CHECK-NEXT: Results
+// CHECK-NEXT: [19, 36]
+// CHECK-NEXT: [1.900000e+01, 3.600000e+01]
+
+func.func @bufferized() -> memref<2xi32> {
+  %v = arith.constant dense<[[1,2,3,4], [5,6,7,8]]> : memref<2x4xi32>
+  %init = arith.constant dense<[9, 10]> : memref<2xi32>
+  linalg.reduce ins(%v : memref<2x4xi32>)
+                outs(%init: memref<2xi32>)
+                dimensions = [1]
+                (%in: i32, %out: i32) {
+                  %sum = arith.addi %in, %out : i32
+                  linalg.yield %sum: i32
+                }
+  func.return %init : memref<2xi32>
+}
+
+// CHECK-LABEL: @bufferized
+// CHECK-NEXT: Results
+// CHECK-NEXT: [19, 36]
+
diff --git a/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/linalg/transpose.mlir b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/linalg/transpose.mlir
new file mode 100644
index 00000000000000..8a60a6ef83ce8e
--- /dev/null
+++ b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/linalg/transpose.mlir
@@ -0,0 +1,27 @@
+// RUN: mlir-interpreter-runner %s -run-all | FileCheck %s
+
+func.func @transpose() -> tensor<2x3xi32> {
+  %a = arith.constant dense<[[1, 2], [3, 4], [5, 6]]> : tensor<3x2xi32>
+  %b = tensor.empty() : tensor<2x3xi32>
+  %ret = linalg.transpose ins(%a : tensor<3x2xi32>)
+                          outs(%b : tensor<2x3xi32>)
+                          permutation = [1, 0]
+  return %ret : tensor<2x3xi32>
+}
+
+// CHECK-LABEL: @transpose
+// CHECK-NEXT: Results
+// CHECK-NEXT{LITERAL}: [[1, 3, 5], [2, 4, 6]]
+
+func.func @transpose_bufferized() -> memref<2x3xi32> {
+  %a = arith.constant dense<[[1, 2], [3, 4], [5, 6]]> : memref<3x2xi32>
+  %b = memref.alloc() : memref<2x3xi32>
+  linalg.transpose ins(%a : memref<3x2xi32>)
+                   outs(%b : memref<2x3xi32>)
+                   permutation = [1, 0]
+  return %b : memref<2x3xi32>
+}
+
+// CHECK-LABEL: @transpose_bufferized
+// CHECK-NEXT: Results
+// CHECK-NEXT{LITERAL}: [[1, 3, 5], [2, 4, 6]]
diff --git a/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/linalg/vecmat.mlir b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/linalg/vecmat.mlir
new file mode 100644
index 00000000000000..2a4280f1581b2c
--- /dev/null
+++ b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/linalg/vecmat.mlir
@@ -0,0 +1,14 @@
+// RUN: mlir-interpreter-runner %s -run-all | FileCheck %s
+
+func.func @vecmat() -> tensor<2xi32> {
+  %lhs = arith.constant dense<[4, 5]> : tensor<2xi32>
+  %rhs = arith.constant dense<[[1, 2], [3, 4]]> : tensor<2x2xi32>
+  %init = tensor.empty() : tensor<2xi32>
+  %ret = linalg.vecmat ins(%lhs, %rhs: tensor<2xi32>, tensor<2x2xi32>)
+                       outs(%init: tensor<2xi32>) -> tensor<2xi32>
+  return %ret : tensor<2xi32>
+}
+
+// CHECK-LABEL: @vecmat
+// CHECK-NEXT: Results
+// CHECK-NEXT{LITERAL}: [19, 28]
diff --git a/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/math/math.mlir b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/math/math.mlir
new file mode 100644
index 00000000000000..3e99a449c8cef6
--- /dev/null
+++ b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/math/math.mlir
@@ -0,0 +1,252 @@
+// RUN: mlir-interpreter-runner %s -run-all | FileCheck %s
+
+func.func @absi() -> i32 {
+  %a = arith.constant -1 : i32
+  %ret = math.absi %a : i32
+  return %ret : i32
+}
+
+// CHECK-LABEL: @absi
+// CHECK-NEXT: Results
+// CHECK-NEXT: 1
+
+func.func @absf() -> (f32, f32) {
+  %a = arith.constant 1.0 : f32
+  %b = arith.constant -1.0 : f32
+  %c = math.absf %a : f32
+  %d = math.absf %b : f32
+  return %c, %d : f32, f32
+}
+
+// CHECK-LABEL: @absf
+// CHECK-NEXT: Results
+// CHECK-NEXT: 1.000000e+00
+// CHECK-NEXT: 1.000000e+00
+
+func.func @atan() -> f32 {
+  %c1 = arith.constant 1.0 : f32
+  %ret = math.atan %c1 : f32
+  return %ret : f32
+}
+
+// CHECK-LABEL: @atan
+// CHECK-NEXT: Results
+// CHECK-NEXT: 7.853982e-01
+
+func.func @atan2() -> f32 {
+  %c10 = arith.constant 10.0 : f32
+  %c1 = arith.constant 1.0 : f32
+  %ret = math.atan2 %c10, %c1 : f32
+  return %ret : f32
+}
+
+// CHECK-LABEL: @atan2
+// CHECK-NEXT: Results
+// CHECK-NEXT: 1.471128e+00
+
+func.func @cbrt() -> f32 {
+  %c-27 = arith.constant -27.0 : f32
+  %ret = math.cbrt %c-27 : f32
+  return %ret : f32
+}
+
+// CHECK-LABEL: @cbrt
+// CHECK-NEXT: Results
+// CHECK-NEXT: -3.000000e+00
+
+func.func @ceil() -> f32 {
+  %a = arith.constant 1.234 : f32
+  %ret = math.ceil %a : f32
+  return %ret : f32
+}
+
+// CHECK-LABEL: @ceil
+// CHECK-NEXT: Results
+// CHECK-NEXT: 2.000000e+00
+
+func.func @ctlz() -> i8 {
+  %c1 = arith.constant 1 : i8
+  %ret = math.ctlz %c1 : i8
+  return %ret : i8
+}
+
+// CHECK-LABEL: @ctlz
+// CHECK-NEXT: Results
+// CHECK-NEXT: 7
+
+func.func @ctpop() -> i8 {
+  %c1 = arith.constant 17 : i8
+  %ret = math.ctpop %c1 : i8
+  return %ret : i8
+}
+
+// CHECK-LABEL: @ctpop
+// CHECK-NEXT: Results
+// CHECK-NEXT: i8: 2
+
+func.func @ctz() -> i8 {
+  %c1 = arith.constant 8 : i8
+  %ret = math.cttz %c1 : i8
+  return %ret : i8
+}
+
+// CHECK-LABEL: @ctz
+// CHECK-NEXT: Results
+// CHECK-NEXT: 3
+
+func.func @copysign() -> (f32, f32) {
+  %a = arith.constant 1.5 : f32
+  %b = arith.constant -10.0 : f32
+  %c = math.copysign %a, %b : f32
+  %d = math.copysign %a, %a : f32
+  return %c, %d : f32, f32
+}
+
+// CHECK-LABEL: @copysign
+// CHECK-NEXT: Results
+// CHECK-NEXT: -1.500000e+00
+// CHECK-NEXT: 1.500000e+00
+
+func.func @cos() -> f32 {
+  %a = arith.constant 0.0 : f32
+  %b = math.cos %a : f32
+  return %b : f32
+}
+
+// CHECK-LABEL: @cos
+// CHECK-NEXT: Results
+// CHECK-NEXT: 1.000000e+00
+
+func.func @erf() -> f32 {
+  %a = arith.constant 1.0 : f32
+  %b = math.erf %a : f32
+  return %b : f32
+}
+
+// CHECK-LABEL: @erf
+// CHECK-NEXT: Results
+// CHECK-NEXT: 8.427008e-01
+
+func.func @exp() -> f32 {
+  %a = arith.constant 1.0 : f32
+  %b = math.exp %a : f32
+  return %b : f32
+}
+
+// CHECK-LABEL: @exp
+// CHECK-NEXT: Results
+// CHECK-NEXT: 2.718282e+00
+
+func.func @exp2() -> f32 {
+  %a = arith.constant 3.0 : f32
+  %b = math.exp2 %a : f32
+  return %b : f32
+}
+
+// CHECK-LABEL: @exp2
+// CHECK-NEXT: Results
+// CHECK-NEXT: 8.000000e+00
+
+func.func @expm1() -> f32 {
+  %a = arith.constant 1.0e-06 : f32
+  %b = math.expm1 %a : f32
+  return %b : f32
+}
+
+// CHECK-LABEL: @expm1
+// CHECK-NEXT: Results
+// CHECK-NEXT: 1.000000e-06
+
+func.func @floor() -> (f32, f32) {
+  %a = arith.constant 1.5 : f32
+  %b = arith.constant -10.5 : f32
+  %c = math.floor %a : f32
+  %d = math.floor %b : f32
+  return %c, %d : f32, f32
+}
+
+// CHECK-LABEL: @floor
+// CHECK-NEXT: Results
+// CHECK-NEXT: 1.000000e+00
+// CHECK-NEXT: -1.100000e+01
+
+func.func @ipowi() -> i32 {
+  %a = arith.constant 2 : i32
+  %b = arith.constant 3 : i32
+  %c = math.ipowi %a, %b : i32
+  return %c : i32
+}
+
+// CHECK-LABEL: @ipowi
+// CHECK-NEXT: Results
+// CHECK-NEXT: i32: 8
+
+func.func @log() -> f32 {
+  %a = arith.constant 1.0 : f32
+  %ret = math.log %a : f32
+  return %ret : f32
+}
+
+// CHECK-LABEL: @log
+// CHECK-NEXT: Results
+// CHECK-NEXT: 0.000000e+00
+
+func.func @log10() -> f32 {
+  %a = arith.constant 1.0e12 : f32
+  %ret = math.log10 %a : f32
+  return %ret : f32
+}
+
+// CHECK-LABEL: @log
+// CHECK-NEXT: Results
+// CHECK-NEXT: 1.200000e+01
+
+func.func @log1p() -> f32 {
+  %a = arith.constant 1.0e-10 : f32
+  %ret = math.log1p %a : f32
+  return %ret : f32
+}
+
+// CHECK-LABEL: @log1p
+// CHECK-NEXT: Results
+// CHECK-NEXT: 1.000000e-10
+
+func.func @log2() -> f32 {
+  %a = arith.constant 65536.0 : f32
+  %ret = math.log2 %a : f32
+  return %ret : f32
+}
+
+// CHECK-LABEL: @log2
+// CHECK-NEXT: Results
+// CHECK-NEXT: 1.600000e+01
+
+func.func @powf() -> f32 {
+  %a = arith.constant 2.0 : f32
+  %b = arith.constant 3.0 : f32
+  %c = math.powf %a, %b : f32
+  return %c : f32
+}
+
+// CHECK-LABEL: @powf
+// CHECK-NEXT: Results
+// CHECK-NEXT: 8.000000e+00
+
+//REGISTER_MLIR_INTERPRETER_OP("math.round", applyCwiseMap<Round>);
+//REGISTER_MLIR_INTERPRETER_OP("math.roundeven",  applyCwiseMap<NearbyInt>);
+//REGISTER_MLIR_INTERPRETER_OP("math.rsqrt", applyCwiseMap<RSqrt>);
+
+func.func @sin() -> f32 {
+  %a = arith.constant 0.0 : f32
+  %b = math.sin %a : f32
+  return %b : f32
+}
+
+// CHECK-LABEL: @sin
+// CHECK-NEXT: Results
+// CHECK-NEXT: 0.000000e+00
+
+//REGISTER_MLIR_INTERPRETER_OP("math.sqrt", applyCwiseMap<Sqrt>);
+//REGISTER_MLIR_INTERPRETER_OP("math.tan", applyCwiseMap<Tan>);
+//REGISTER_MLIR_INTERPRETER_OP("math.tanh", applyCwiseMap<TanH>);
+//REGISTER_MLIR_INTERPRETER_OP("math.trunc", applyCwiseMap<Trunc>);
diff --git a/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/memref/alloc.mlir b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/memref/alloc.mlir
new file mode 100644
index 00000000000000..d6026b5b913a1b
--- /dev/null
+++ b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/memref/alloc.mlir
@@ -0,0 +1,57 @@
+// RUN: mlir-interpreter-runner %s -run-all | FileCheck %s
+
+func.func @alloc() -> memref<2x3xi32> {
+  %ret = memref.alloc() : memref<2x3xi32>
+  return %ret : memref<2x3xi32>
+}
+
+// CHECK-LABEL: @alloc
+// CHECK-NEXT: Results
+// CHECK-NEXT{LITERAL}: <2x3xi32>: [[0, 0, 0], [0, 0, 0]]
+
+func.func @alloc_unit() -> memref<i32> {
+  %ret = memref.alloc() : memref<i32>
+  return %ret : memref<i32>
+}
+
+// CHECK-LABEL: @alloc_unit
+// CHECK-NEXT: Results
+// CHECK-NEXT{LITERAL}: <i32>: 0
+
+func.func @alloc_unit_vector() -> memref<vector<i32>> {
+  %ret = memref.alloc() : memref<vector<i32>>
+  return %ret : memref<vector<i32>>
+}
+
+// CHECK-LABEL: @alloc_unit_vector
+// CHECK-NEXT: Results
+// CHECK-NEXT{LITERAL}: <vector<i32>>: 0
+
+func.func @alloc_vector() -> memref<2xvector<3xi32>> {
+  %ret = memref.alloc() : memref<2xvector<3xi32>>
+  return %ret : memref<2xvector<3xi32>>
+}
+
+// CHECK-LABEL: @alloc_vector
+// CHECK-NEXT: Results
+// CHECK-NEXT{LITERAL}: <2xvector<3xi32>>: [[0, 0, 0], [0, 0, 0]]
+
+func.func @alloc_dynamic() -> memref<?x3xi32> {
+  %c2 = arith.constant 2 : index
+  %ret = memref.alloc(%c2) : memref<?x3xi32>
+  return %ret : memref<?x3xi32>
+}
+
+// CHECK-LABEL: @alloc_dynamic
+// CHECK-NEXT: Results
+// CHECK-NEXT{LITERAL}: <2x3xi32>: [[0, 0, 0], [0, 0, 0]]
+
+func.func @dealloc() -> memref<i32> {
+  %a = memref.alloc() : memref<i32>
+  memref.dealloc %a : memref<i32>
+  return %a : memref<i32>
+}
+
+// CHECK-LABEL: @dealloc
+// CHECK-NEXT: Results
+// CHECK-NEXT: TensorOrMemref<i32>: <<deallocated>>
diff --git a/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/memref/collapse_shape.mlir b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/memref/collapse_shape.mlir
new file mode 100644
index 00000000000000..027b59755bed8c
--- /dev/null
+++ b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/memref/collapse_shape.mlir
@@ -0,0 +1,33 @@
+// RUN: mlir-interpreter-runner %s -run-all | FileCheck %s
+
+func.func @collapse_shape()
+    -> (memref<1x2x3xi32>, memref<1x6xi32>, memref<2x3xi32>, memref<6xi32>) {
+  %cst = arith.constant dense<[[[1, 2, 3], [4, 5, 6]]]> : memref<1x2x3xi32>
+  %collapse1 = memref.collapse_shape %cst [[0], [1, 2]]
+      : memref<1x2x3xi32> into memref<1x6xi32>
+  %collapse2 = memref.collapse_shape %cst [[0, 1], [2]]
+      : memref<1x2x3xi32> into memref<2x3xi32>
+  %collapse3 = memref.collapse_shape %cst [[0, 1, 2]]
+      : memref<1x2x3xi32> into memref<6xi32>
+  return %cst, %collapse1, %collapse2, %collapse3
+      : memref<1x2x3xi32>, memref<1x6xi32>, memref<2x3xi32>, memref<6xi32>
+}
+
+// CHECK-LABEL: @collapse_shape
+// CHECK-NEXT: Results
+// CHECK-NEXT{LITERAL}: [[[1, 2, 3], [4, 5, 6]]]
+// CHECK-NEXT{LITERAL}: [[1, 2, 3, 4, 5, 6]]
+// CHECK-NEXT{LITERAL}: [[1, 2, 3], [4, 5, 6]]
+// CHECK-NEXT{LITERAL}: [1, 2, 3, 4, 5, 6]
+
+func.func @zero_dim()
+    -> (memref<6x0xi32>) {
+  %cst = arith.constant dense<> : memref<1x2x3x0xi32>
+  %collapse = memref.collapse_shape %cst [[0, 1, 2], [3]]
+      : memref<1x2x3x0xi32> into memref<6x0xi32>
+  return %collapse : memref<6x0xi32>
+}
+
+// CHECK-LABEL: @zero_dim
+// CHECK-NEXT: Results
+// CHECK-NEXT: TensorOrMemref<6x0xi32>
diff --git a/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/memref/copy.mlir b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/memref/copy.mlir
new file mode 100644
index 00000000000000..09815f5e8cf61a
--- /dev/null
+++ b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/memref/copy.mlir
@@ -0,0 +1,39 @@
+// RUN: mlir-interpreter-runner %s -run-all | FileCheck %s
+
+func.func @subview() -> memref<4x4xi32, strided<[4, 1], offset: 0>> {
+  %c1 = arith.constant 1 : index
+  %c2 = arith.constant 2 : index
+
+  %0 = memref.alloc() : memref<4x4xi32, strided<[4, 1], offset: 0>>
+  %1 = memref.subview %0[%c1, %c1][%c2, %c2][%c1, %c1]
+    : memref<4x4xi32, strided<[4, 1], offset: 0>> to
+      memref<?x?xi32, strided<[?, ?], offset: ?>>
+
+  %cst = arith.constant dense<[[1, 2], [3, 4]]> : memref<2x2xi32>
+  memref.copy %cst, %1 : memref<2x2xi32> to memref<?x?xi32, strided<[?, ?], offset: ?>>
+
+  return %0 : memref<4x4xi32, strided<[4, 1], offset: 0>>
+}
+
+// CHECK-LABEL: @subview
+// CHECK-NEXT: Results
+// CHECK-NEXT{LITERAL}: [[0, 0, 0, 0], [0, 1, 2, 0], [0, 3, 4, 0], [0, 0, 0, 0]]
+
+func.func @strided() -> memref<4x4xi32> {
+  %c1 = arith.constant 1 : index
+  %c2 = arith.constant 2 : index
+
+  %0 = memref.alloc() : memref<4x4xi32>
+  %1 = memref.subview %0[%c1, %c1][%c2, %c2][%c2, %c2]
+    : memref<4x4xi32> to memref<?x?xi32, strided<[?, ?], offset: ?>>
+
+  %cst = arith.constant dense<[[1, 2], [3, 4]]> : memref<2x2xi32>
+  memref.copy %cst, %1
+    : memref<2x2xi32> to memref<?x?xi32, strided<[?, ?], offset: ?>>
+
+  return %0 : memref<4x4xi32>
+}
+
+// CHECK-LABEL: @strided
+// CHECK-NEXT: Results
+// CHECK-NEXT{LITERAL}: [[0, 0, 0, 0], [0, 1, 0, 2], [0, 0, 0, 0], [0, 3, 0, 4]]
diff --git a/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/memref/dim.mlir b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/memref/dim.mlir
new file mode 100644
index 00000000000000..2740b7df65de6d
--- /dev/null
+++ b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/memref/dim.mlir
@@ -0,0 +1,12 @@
+// RUN: mlir-interpreter-runner %s -run-all | FileCheck %s
+
+func.func @dim() -> index {
+  %alloc = memref.alloc() {alignment = 64 : i64} : memref<10x50xf32>
+  %c1 = arith.constant 1 : index
+  %dim = memref.dim %alloc, %c1 : memref<10x50xf32>
+  return %dim : index
+}
+
+// CHECK-LABEL: @dim
+// CHECK-NEXT: Results
+// CHECK-NEXT: i64: 50
diff --git a/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/memref/expand_shape.mlir b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/memref/expand_shape.mlir
new file mode 100644
index 00000000000000..7f9cdc7e27f49f
--- /dev/null
+++ b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/memref/expand_shape.mlir
@@ -0,0 +1,52 @@
+// RUN: mlir-interpreter-runner %s -run-all | FileCheck %s
+
+func.func @expand_shape()
+    -> (memref<1x2x3xi32>, memref<1x2x3xi32>, memref<2x1x3xi32>,
+        memref<1x2x3xi32>) {
+  %cst1 = arith.constant dense<[[1, 2, 3, 4, 5, 6]]> : memref<1x6xi32>
+  %cst2 = arith.constant dense<[[1, 2, 3], [4, 5, 6]]> : memref<2x3xi32>
+  %cst3 = arith.constant dense<[1, 2, 3, 4, 5, 6]> : memref<6xi32>
+  %expand1 = memref.expand_shape %cst1 [[0], [1, 2]] output_shape [1, 2, 3]
+      : memref<1x6xi32> into memref<1x2x3xi32>
+  %expand2 = memref.expand_shape %cst2 [[0, 1], [2]] output_shape [1, 2, 3]
+      : memref<2x3xi32> into memref<1x2x3xi32>
+  %expand3 = memref.expand_shape %cst2 [[0, 1], [2]] output_shape [2, 1, 3]
+      : memref<2x3xi32> into memref<2x1x3xi32>
+  %expand4 = memref.expand_shape %cst3 [[0, 1, 2]] output_shape [1, 2, 3]
+      : memref<6xi32> into memref<1x2x3xi32>
+  return %expand1, %expand2, %expand3, %expand4
+      : memref<1x2x3xi32>, memref<1x2x3xi32>, memref<2x1x3xi32>,
+        memref<1x2x3xi32>
+}
+
+// CHECK-LABEL: @expand_shape
+// CHECK-NEXT: Results
+// CHECK-NEXT{LITERAL}: [[[1, 2, 3], [4, 5, 6]]]
+// CHECK-NEXT{LITERAL}: [[[1, 2, 3], [4, 5, 6]]]
+// CHECK-NEXT{LITERAL}: [[[1, 2, 3]], [[4, 5, 6]]]
+// CHECK-NEXT{LITERAL}: [[[1, 2, 3], [4, 5, 6]]]
+
+func.func @zero_rank()
+    -> (memref<1x1xi32>) {
+  %cst = arith.constant dense<1> : memref<i32>
+  %expand = memref.expand_shape %cst [] output_shape [1, 1]
+      : memref<i32> into memref<1x1xi32>
+  return %expand : memref<1x1xi32>
+}
+
+// CHECK-LABEL: @zero_rank
+// CHECK-NEXT: Results
+// CHECK-NEXT{LITERAL}: [[1]]
+
+func.func @split_dim() -> memref<3x2x2xi32, strided<[4, 2, 1], offset: 0>> {
+  %cst = arith.constant dense<[[1, 2, 3, 4], [5, 6, 7, 8], [9, 10, 11, 12]]>
+    : memref<3x4xi32, strided<[4, 1], offset: 0>>
+  %ret = memref.expand_shape %cst [[0], [1, 2]] output_shape [3, 2, 2]
+    : memref<3x4xi32, strided<[4, 1], offset: 0>> into
+      memref<3x2x2xi32, strided<[4, 2, 1], offset: 0>>
+  return %ret : memref<3x2x2xi32, strided<[4, 2, 1], offset: 0>>
+}
+
+// CHECK-LABEL: @split_dim
+// CHECK-NEXT: Results
+// CHECK-NEXT{LITERAL}: [[[1, 2], [3, 4]], [[5, 6], [7, 8]], [[9, 10], [11, 12]]]
diff --git a/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/memref/get_global.mlir b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/memref/get_global.mlir
new file mode 100644
index 00000000000000..a8b5de8930cc83
--- /dev/null
+++ b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/memref/get_global.mlir
@@ -0,0 +1,12 @@
+// RUN: mlir-interpreter-runner %s -run-all | FileCheck %s
+
+memref.global "private" constant @cst : memref<2xi16> = dense<[1, 2]>
+
+func.func @get_global() -> memref<2xi16> {
+  %0 = memref.get_global @cst : memref<2xi16>
+  return %0 : memref<2xi16>
+}
+
+// CHECK-LABEL: @get_global
+// CHECK-NEXT: Results
+// CHECK-NEXT{LITERAL}: [1, 2]
diff --git a/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/memref/invalid.mlir b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/memref/invalid.mlir
new file mode 100644
index 00000000000000..a533a9664c235d
--- /dev/null
+++ b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/memref/invalid.mlir
@@ -0,0 +1,77 @@
+// RUN: (! mlir-interpreter-runner %s -run-all 2>&1) | FileCheck %s
+
+func.func @out_of_bounds_load() -> i32 {
+  %c3 = arith.constant 3 : index
+  %cst = arith.constant dense<[1, 2]> : memref<2xi32>
+  %ret = memref.load %cst[%c3] : memref<2xi32>
+  return %ret : i32
+}
+
+// CHECK-LABEL: @out_of_bounds_load
+// CHECK-NEXT: array index out of bounds
+
+func.func @out_of_bounds_store() {
+  %c3 = arith.constant 3 : index
+  %v = arith.constant 32 : i32
+  %cst = arith.constant dense<[1, 2]> : memref<2xi32>
+  memref.store %v, %cst[%c3] : memref<2xi32>
+  return
+}
+
+// CHECK-LABEL: @out_of_bounds_store
+// CHECK-NEXT: array index out of bounds
+
+func.func @out_of_bounds_subview() -> memref<?xi32, strided<[?], offset: ?>> {
+  %c1 = arith.constant 1 : index
+  %c2 = arith.constant 2 : index
+  %v = arith.constant 32 : i32
+  %cst = arith.constant dense<[1, 2]> : memref<2xi32>
+  %subview = memref.subview %cst[%c1][%c2][%c1]
+    : memref<2xi32> to memref<?xi32, strided<[?], offset: ?>>
+  return %subview : memref<?xi32, strided<[?], offset: ?>>
+}
+
+// CHECK-LABEL: @out_of_bounds_subview
+// CHECK-NEXT: subview out of bounds
+
+func.func @collapse_shape_no_common_stride()
+    -> (memref<1x2x3xi32>, memref<1x6xi32>) {
+  %a = arith.constant dense<[[[[0, 1, 2], [3, 4, 5]],
+                              [[6, 7, 8], [9, 10,11]]]]>
+    : memref<1x2x2x3xi32>
+  %b = memref.subview %a[0, 0, 0, 0][1, 2, 1, 3][1, 1, 1, 1]
+    : memref<1x2x2x3xi32> to memref<1x2x3xi32>
+  %c = memref.collapse_shape %b [[0], [1, 2]]
+    : memref<1x2x3xi32> into memref<1x6xi32>
+  return %b, %c : memref<1x2x3xi32>, memref<1x6xi32>
+}
+
+// CHECK-LABEL: @collapse_shape_no_common_stride
+// CHECK-NEXT: cannot collapse dimensions without a common stride
+
+func.func @double_free() {
+  %a = memref.alloc() : memref<i32>
+  memref.dealloc %a : memref<i32>
+  memref.dealloc %a : memref<i32>
+  return
+}
+
+// CHECK-LABEL: @double_free
+// CHECK-NEXT: Interpreter failure: double-free
+// CHECK-NEXT: Note: allocated by %alloc = memref.alloc() : memref<i32>
+// CHECK-NEXT: Note: previously freed by memref.dealloc %alloc : memref<i32>
+// CHECK-NEXT{2}: Encountered failure while executing memref.dealloc %alloc : memref<i32>
+
+func.func @use_after_free() {
+  %a = memref.alloc() : memref<i32>
+  memref.dealloc %a : memref<i32>
+  %b = arith.constant 1 : i32
+  memref.store %b, %a[] : memref<i32>
+  return
+}
+
+// CHECK-LABEL: @use_after_free
+// CHECK-NEXT: Interpreter failure: use-after-free
+// CHECK-NEXT: Note: allocated by %alloc = memref.alloc() : memref<i32>
+// CHECK-NEXT: Note: previously freed by memref.dealloc %alloc : memref<i32>
+// CHECK-NEXT{2}: Encountered failure while executing memref.store %c1_i32, %alloc[] : memref<i32>
diff --git a/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/memref/load.mlir b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/memref/load.mlir
new file mode 100644
index 00000000000000..17d935319622e1
--- /dev/null
+++ b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/memref/load.mlir
@@ -0,0 +1,12 @@
+// RUN: mlir-interpreter-runner %s -run-all | FileCheck %s
+
+func.func @load() -> i32 {
+  %c1 = arith.constant 1 : index
+  %cst = arith.constant dense<[[1, 2], [3, 4]]> : memref<2x2xi32>
+  %ret = memref.load %cst[%c1, %c1] : memref<2x2xi32>
+  return %ret : i32
+}
+
+// CHECK-LABEL: @load
+// CHECK-NEXT: Results
+// CHECK-NEXT: 4
diff --git a/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/memref/subview.mlir b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/memref/subview.mlir
new file mode 100644
index 00000000000000..441548ffe2c84d
--- /dev/null
+++ b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/memref/subview.mlir
@@ -0,0 +1,120 @@
+// RUN: mlir-interpreter-runner %s -run-all | FileCheck %s
+
+func.func @subview() -> (memref<4x4xi32, strided<[4, 1], offset: 0>>,
+                         memref<?x?xi32, strided<[?, ?], offset: ?>>) {
+  %c0 = arith.constant 0 : index
+  %c1 = arith.constant 1 : index
+  %c2 = arith.constant 2 : index
+
+  %0 = memref.alloc() : memref<4x4xi32, strided<[4, 1], offset: 0>>
+  %1 = memref.subview %0[%c1, %c1][%c2, %c2][%c1, %c1]
+    : memref<4x4xi32, strided<[4, 1], offset: 0>> to
+      memref<?x?xi32, strided<[?, ?], offset: ?>>
+
+  %i1 = arith.constant 1 : i32
+  %i2 = arith.constant 2 : i32
+  %i3 = arith.constant 3 : i32
+  %i4 = arith.constant 4 : i32
+
+  memref.store %i1, %1[%c0, %c0] : memref<?x?xi32, strided<[?, ?], offset: ?>>
+  memref.store %i2, %1[%c0, %c1] : memref<?x?xi32, strided<[?, ?], offset: ?>>
+  memref.store %i3, %1[%c1, %c0] : memref<?x?xi32, strided<[?, ?], offset: ?>>
+  memref.store %i4, %1[%c1, %c1] : memref<?x?xi32, strided<[?, ?], offset: ?>>
+
+  return %0, %1 : memref<4x4xi32, strided<[4, 1], offset: 0>>,
+                  memref<?x?xi32, strided<[?, ?], offset: ?>>
+}
+
+// CHECK-LABEL: @subview
+// CHECK-NEXT: Results
+// CHECK-NEXT{LITERAL}: [[0, 0, 0, 0], [0, 1, 2, 0], [0, 3, 4, 0], [0, 0, 0, 0]]
+// CHECK-NEXT{LITERAL}: [[1, 2], [3, 4]]
+
+func.func @strided() -> memref<?x?xi32, strided<[?, ?], offset: ?>> {
+  %0 = arith.constant dense<[[1, 2, 3, 4, 5], [6, 7, 8, 9, 10]]> : memref<2x5xi32>
+
+  %c0 = arith.constant 0 : index
+  %c1 = arith.constant 1 : index
+  %c2 = arith.constant 2 : index
+
+  %1 = memref.subview %0[%c0, %c1][%c2, %c2][%c1, %c2]
+    : memref<2x5xi32> to
+      memref<?x?xi32, strided<[?, ?], offset: ?>>
+
+  return %1 : memref<?x?xi32, strided<[?, ?], offset: ?>>
+}
+
+// CHECK-LABEL: @strided
+// CHECK-NEXT: Results
+// CHECK-NEXT{LITERAL}: [[2, 4], [7, 9]]
+
+func.func @subview_of_subview() -> (memref<?xi32, strided<[?], offset: ?>>,
+                                     memref<?xi32, strided<[?], offset: ?>>) {
+  %0 = arith.constant dense<[1, 2, 3, 4, 5, 6, 7, 8, 9, 10]> : memref<10xi32>
+
+  %c0 = arith.constant 0 : index
+  %c1 = arith.constant 1 : index
+  %c2 = arith.constant 2 : index
+  %c4 = arith.constant 4 : index
+
+  %1 = memref.subview %0[%c1][%c4][%c2]
+    : memref<10xi32> to memref<?xi32, strided<[?], offset: ?>>
+  %2 = memref.subview %1[%c1][%c2][%c2]
+    : memref<?xi32, strided<[?], offset: ?>> to
+      memref<?xi32, strided<[?], offset: ?>>
+
+  return %1, %2 : memref<?xi32, strided<[?], offset: ?>>,
+                  memref<?xi32, strided<[?], offset: ?>>
+}
+
+// CHECK-LABEL: @subview_of_subview
+// CHECK-NEXT: Results
+// CHECK-NEXT{LITERAL}: [2, 4, 6, 8]
+// CHECK-NEXT{LITERAL}: [4, 8]
+
+func.func @negative_stride() -> memref<?xi32, strided<[?], offset: ?>> {
+  %0 = arith.constant dense<[1, 2, 3, 4]> : memref<4xi32>
+  %c-1 = arith.constant -1 : index
+  %c3 = arith.constant 3 : index
+  %c4 = arith.constant 4 : index
+  %1 = memref.subview %0[%c3][%c4][%c-1]
+    : memref<4xi32> to memref<?xi32, strided<[?], offset: ?>>
+  return %1 : memref<?xi32, strided<[?], offset: ?>>
+}
+
+// CHECK-LABEL: @negative_stride
+// CHECK-NEXT: Results
+// CHECK-NEXT: [4, 3, 2, 1]
+
+func.func @negative_stride_of_subview() -> (memref<?xi32, strided<[?], offset: ?>>, memref<?xi32, strided<[?], offset: ?>>) {
+  %0 = arith.constant dense<[1, 2, 3, 4, 5, 6, 7, 8, 9, 10]> : memref<10xi32>
+  %c-2 = arith.constant -2 : index
+  %c1 = arith.constant 1 : index
+  %c2 = arith.constant 2 : index
+  %c3 = arith.constant 3 : index
+  %c4 = arith.constant 4 : index
+  %1 = memref.subview %0[%c1][%c4][%c2]
+    : memref<10xi32> to memref<?xi32, strided<[?], offset: ?>>
+  %2 = memref.subview %1[%c3][%c2][%c-2]
+    : memref<?xi32, strided<[?], offset: ?>> to
+      memref<?xi32, strided<[?], offset: ?>>
+  return %1, %2 : memref<?xi32, strided<[?], offset: ?>>, memref<?xi32, strided<[?], offset: ?>>
+}
+
+// CHECK-LABEL: @negative_stride_of_subview
+// CHECK-NEXT: Results
+// CHECK-NEXT{LITERAL}: [2, 4, 6, 8]
+// CHECK-NEXT{LITERAL}: [8, 4]
+
+func.func @rank_reduce_middle() -> memref<1x2x3xi32> {
+  %a = arith.constant dense<[[[[0, 1, 2], [3, 4, 5]],
+                              [[6, 7, 8], [9, 10,11]]]]>
+    : memref<1x2x2x3xi32>
+  %b = memref.subview %a[0, 0, 0, 0][1, 2, 1, 3][1, 1, 1, 1]
+    : memref<1x2x2x3xi32> to memref<1x2x3xi32>
+  return %b : memref<1x2x3xi32>
+}
+
+// CHECK-LABEL: @rank_reduce_middle
+// CHECK-NEXT: Results
+// CHECK-NEXT{LITERAL}: [[[0, 1, 2], [6, 7, 8]]]
diff --git a/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/mhlo/bitcast_convert.mlir b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/mhlo/bitcast_convert.mlir
new file mode 100644
index 00000000000000..f628ebadd21abc
--- /dev/null
+++ b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/mhlo/bitcast_convert.mlir
@@ -0,0 +1,11 @@
+// RUN: mlir-interpreter-runner %s -run-all | FileCheck %s
+
+func.func @bitcast_convert() -> tensor<3xf32> {
+  %c = arith.constant dense<[10000000,20000000,30000000]> : tensor<3xi32>
+  %ret = mhlo.bitcast_convert %c : (tensor<3xi32>) -> tensor<3xf32>
+  return %ret : tensor<3xf32>
+}
+
+// CHECK-LABEL: @bitcast_convert
+// CHECK-NEXT: Results
+// CHECK-NEXT: [1.401298e-38, 3.254205e-38, 7.411627e-38]
diff --git a/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/mhlo/broadcast_in_dim.mlir b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/mhlo/broadcast_in_dim.mlir
new file mode 100644
index 00000000000000..6f834dfb2fecea
--- /dev/null
+++ b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/mhlo/broadcast_in_dim.mlir
@@ -0,0 +1,20 @@
+// RUN: mlir-interpreter-runner %s -run-all | FileCheck %s
+
+func.func @broadcast() -> tensor<2x3xui16> {
+  %0 = mhlo.constant dense<[1, 2]> : tensor<2xui16>
+  %1 = "mhlo.broadcast_in_dim"(%0) {
+    broadcast_dimensions = dense<0> : tensor<1xi64>} : (tensor<2xui16>) -> tensor<2x3xui16>
+  return %1 : tensor<2x3xui16>
+}
+
+// CHECK{LITERAL}: [[1, 1, 1], [2, 2, 2]]
+
+func.func @zero_rank() -> tensor<1x2x3xi32> {
+  %in = mhlo.constant dense<1> : tensor<i32>
+  %0 = "mhlo.broadcast_in_dim"(%in) {
+    broadcast_dimensions = dense<[]> : tensor<0xi64>
+  } : (tensor<i32>) -> tensor<1x2x3xi32>
+  func.return %0 : tensor<1x2x3xi32>
+}
+
+// CHECK{LITERAL}: [[[1, 1, 1], [1, 1, 1]]]
diff --git a/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/mhlo/case.mlir b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/mhlo/case.mlir
new file mode 100644
index 00000000000000..17d7f8fecd3641
--- /dev/null
+++ b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/mhlo/case.mlir
@@ -0,0 +1,17 @@
+// RUN: mlir-interpreter-runner %s -run-all | FileCheck %s
+
+func.func @case() -> tensor<i32> {
+  %c1 = mhlo.constant dense<1> : tensor<i32>
+  %c2 = mhlo.constant dense<2> : tensor<i32>
+  %c3 = mhlo.constant dense<3> : tensor<i32>
+  %ret = "mhlo.case"(%c1) ({
+    "mhlo.return"(%c2) : (tensor<i32>) -> ()
+  }, {
+    "mhlo.return"(%c3) : (tensor<i32>) -> ()
+  }) : (tensor<i32>) -> tensor<i32>
+  func.return %ret : tensor<i32>
+}
+
+// CHECK-LABEL: @case
+// CHECK-NEXT: Results
+// CHECK-NEXT: <i32>: 3
diff --git a/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/mhlo/clamp.mlir b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/mhlo/clamp.mlir
new file mode 100644
index 00000000000000..f7a455a78e2293
--- /dev/null
+++ b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/mhlo/clamp.mlir
@@ -0,0 +1,27 @@
+// RUN: mlir-interpreter-runner %s -run-all | FileCheck %s
+
+func.func @clamp() -> tensor<2x2xi32> {
+  %lb = mhlo.constant dense<[[1, 7], [1, 7]]> : tensor<2x2xi32>
+  %arg = mhlo.constant dense<[[4, 5], [6, 9]]> : tensor<2x2xi32>
+  %ub = mhlo.constant dense<[[5, 9], [3, 6]]> : tensor<2x2xi32>
+  %clamp = "mhlo.clamp"(%lb, %arg, %ub)
+      : (tensor<2x2xi32>, tensor<2x2xi32>, tensor<2x2xi32>) -> tensor<2x2xi32>
+  return %clamp : tensor<2x2xi32>
+}
+
+// CHECK-LABEL:         @clamp
+// CHECK-NEXT:          Results
+// CHECK-NEXT{LITERAL}: [[4, 7], [3, 6]]
+
+func.func @clamp_f32() -> tensor<2x2xf32> {
+  %lb = mhlo.constant dense<[[1.1, 7.1], [1.1, 7.1]]> : tensor<2x2xf32>
+  %arg = mhlo.constant dense<[[4.1, 5.1], [6.1, 9.1]]> : tensor<2x2xf32>
+  %ub = mhlo.constant dense<[[5.1, 9.1], [3.1, 6.1]]> : tensor<2x2xf32>
+  %clamp = "mhlo.clamp"(%lb, %arg, %ub)
+      : (tensor<2x2xf32>, tensor<2x2xf32>, tensor<2x2xf32>) -> tensor<2x2xf32>
+  return %clamp : tensor<2x2xf32>
+}
+
+// CHECK-LABEL:         @clamp
+// CHECK-NEXT:          Results
+// CHECK-NEXT{LITERAL}: [[4.100000e+00, 7.100000e+00], [3.100000e+00, 6.100000e+00]]
diff --git a/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/mhlo/compare.mlir b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/mhlo/compare.mlir
new file mode 100644
index 00000000000000..f9d7eb99cbeac7
--- /dev/null
+++ b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/mhlo/compare.mlir
@@ -0,0 +1,143 @@
+// RUN: mlir-interpreter-runner %s -run-all | FileCheck %s
+
+func.func @eq() -> (tensor<i1>, tensor<i1>) {
+  %c1 = arith.constant dense<1> : tensor<i32>
+  %c2 = arith.constant dense<2> : tensor<i32>
+  %0 = mhlo.compare EQ, %c1, %c2 : (tensor<i32>, tensor<i32>) -> tensor<i1>
+  %1 = mhlo.compare EQ, %c1, %c1 : (tensor<i32>, tensor<i32>) -> tensor<i1>
+  return %0, %1 : tensor<i1>, tensor<i1>
+}
+
+// CHECK-LABEL: @eq
+// CHECK-NEXT: Results
+// CHECK-NEXT: false
+// CHECK-NEXT: true
+
+func.func @ne() -> (tensor<i1>, tensor<i1>) {
+  %c1 = arith.constant dense<1> : tensor<i32>
+  %c2 = arith.constant dense<2> : tensor<i32>
+  %0 = mhlo.compare NE, %c1, %c2 : (tensor<i32>, tensor<i32>) -> tensor<i1>
+  %1 = mhlo.compare NE, %c1, %c1 : (tensor<i32>, tensor<i32>) -> tensor<i1>
+  return %0, %1 : tensor<i1>, tensor<i1>
+}
+
+// CHECK-LABEL: @ne
+// CHECK-NEXT: Results
+// CHECK-NEXT: true
+// CHECK-NEXT: false
+
+func.func @ge() -> (tensor<i1>, tensor<i1>, tensor<i1>) {
+  %c1 = arith.constant dense<1> : tensor<i32>
+  %c2 = arith.constant dense<2> : tensor<i32>
+  %0 = mhlo.compare GE, %c1, %c2 : (tensor<i32>, tensor<i32>) -> tensor<i1>
+  %1 = mhlo.compare GE, %c1, %c1 : (tensor<i32>, tensor<i32>) -> tensor<i1>
+  %2 = mhlo.compare GE, %c2, %c1 : (tensor<i32>, tensor<i32>) -> tensor<i1>
+  return %0, %1, %2 : tensor<i1>, tensor<i1>, tensor<i1>
+}
+
+// CHECK-LABEL: @ge
+// CHECK-NEXT: Results
+// CHECK-NEXT: false
+// CHECK-NEXT: true
+// CHECK-NEXT: true
+
+func.func @gt() -> (tensor<i1>, tensor<i1>, tensor<i1>) {
+  %c1 = arith.constant dense<1> : tensor<i32>
+  %c2 = arith.constant dense<2> : tensor<i32>
+  %0 = mhlo.compare GT, %c1, %c2 : (tensor<i32>, tensor<i32>) -> tensor<i1>
+  %1 = mhlo.compare GT, %c1, %c1 : (tensor<i32>, tensor<i32>) -> tensor<i1>
+  %2 = mhlo.compare GT, %c2, %c1 : (tensor<i32>, tensor<i32>) -> tensor<i1>
+  return %0, %1, %2 : tensor<i1>, tensor<i1>, tensor<i1>
+}
+
+// CHECK-LABEL: @gt
+// CHECK-NEXT: Results
+// CHECK-NEXT: false
+// CHECK-NEXT: false
+// CHECK-NEXT: true
+
+func.func @le() -> (tensor<i1>, tensor<i1>, tensor<i1>) {
+  %c1 = arith.constant dense<1> : tensor<i32>
+  %c2 = arith.constant dense<2> : tensor<i32>
+  %0 = mhlo.compare LE, %c1, %c2 : (tensor<i32>, tensor<i32>) -> tensor<i1>
+  %1 = mhlo.compare LE, %c1, %c1 : (tensor<i32>, tensor<i32>) -> tensor<i1>
+  %2 = mhlo.compare LE, %c2, %c1 : (tensor<i32>, tensor<i32>) -> tensor<i1>
+  return %0, %1, %2 : tensor<i1>, tensor<i1>, tensor<i1>
+}
+
+// CHECK-LABEL: @le
+// CHECK-NEXT: Results
+// CHECK-NEXT: true
+// CHECK-NEXT: true
+// CHECK-NEXT: false
+
+func.func @lt() -> (tensor<i1>, tensor<i1>, tensor<i1>) {
+  %c1 = arith.constant dense<1> : tensor<i32>
+  %c2 = arith.constant dense<2> : tensor<i32>
+  %0 = mhlo.compare LT, %c1, %c2 : (tensor<i32>, tensor<i32>) -> tensor<i1>
+  %1 = mhlo.compare LT, %c1, %c1 : (tensor<i32>, tensor<i32>) -> tensor<i1>
+  %2 = mhlo.compare LT, %c2, %c1 : (tensor<i32>, tensor<i32>) -> tensor<i1>
+  return %0, %1, %2 : tensor<i1>, tensor<i1>, tensor<i1>
+}
+
+// CHECK-LABEL: @lt
+// CHECK-NEXT: Results
+// CHECK-NEXT: true
+// CHECK-NEXT: false
+// CHECK-NEXT: false
+
+func.func @complex_eq() -> (tensor<i1>, tensor<i1>) {
+  %c1 = arith.constant dense<(1.0, 1.0)> : tensor<complex<f32>>
+  %c2 = arith.constant dense<(1.0, 2.0)> : tensor<complex<f32>>
+  %0 = mhlo.compare EQ, %c1, %c2
+    : (tensor<complex<f32>>, tensor<complex<f32>>) -> tensor<i1>
+  %1 = mhlo.compare EQ, %c1, %c1
+    : (tensor<complex<f32>>, tensor<complex<f32>>) -> tensor<i1>
+  return %0, %1 : tensor<i1>, tensor<i1>
+}
+
+// CHECK-LABEL: @complex_eq
+// CHECK-NEXT: Results
+// CHECK-NEXT: false
+// CHECK-NEXT: true
+
+func.func @complex_nan_compare() -> (tensor<i1>, tensor<i1>) {
+  %nan = arith.constant dense<0x7FC00000> : tensor<f32>
+  %c1 = arith.constant dense<1.0> : tensor<f32>
+  %c = mhlo.complex %c1, %nan : tensor<complex<f32>>
+  %0 = mhlo.compare EQ, %c, %c
+    : (tensor<complex<f32>>, tensor<complex<f32>>) -> tensor<i1>
+  %1 = mhlo.compare NE, %c, %c
+    : (tensor<complex<f32>>, tensor<complex<f32>>) -> tensor<i1>
+  return %0, %1 : tensor<i1>, tensor<i1>
+}
+
+// CHECK-LABEL: @complex_nan_compare
+// CHECK-NEXT: Results
+// CHECK-NEXT: false
+// CHECK-NEXT: true
+
+func.func @float_eq() -> (tensor<i1>, tensor<i1>) {
+  %c1 = arith.constant dense<1.0> : tensor<f32>
+  %c2 = arith.constant dense<2.0> : tensor<f32>
+  %0 = mhlo.compare EQ, %c1, %c2 : (tensor<f32>, tensor<f32>) -> tensor<i1>
+  %1 = mhlo.compare EQ, %c1, %c1 : (tensor<f32>, tensor<f32>) -> tensor<i1>
+  return %0, %1 : tensor<i1>, tensor<i1>
+}
+
+// CHECK-LABEL: @float_eq
+// CHECK-NEXT: Results
+// CHECK-NEXT: false
+// CHECK-NEXT: true
+
+func.func @float_nan_compare() -> (tensor<i1>, tensor<i1>) {
+  %nan = arith.constant dense<0x7FC00000> : tensor<f32>
+  %0 = mhlo.compare EQ, %nan, %nan : (tensor<f32>, tensor<f32>) -> tensor<i1>
+  %1 = mhlo.compare NE, %nan, %nan : (tensor<f32>, tensor<f32>) -> tensor<i1>
+  return %0, %1 : tensor<i1>, tensor<i1>
+}
+
+// CHECK-LABEL: @float_nan_compare
+// CHECK-NEXT: Results
+// CHECK-NEXT: false
+// CHECK-NEXT: true
\ No newline at end of file
diff --git a/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/mhlo/complex_math.mlir b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/mhlo/complex_math.mlir
new file mode 100644
index 00000000000000..c31953be62158e
--- /dev/null
+++ b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/mhlo/complex_math.mlir
@@ -0,0 +1,100 @@
+// RUN: mlir-interpreter-runner %s -run-all | FileCheck %s
+
+func.func @cos() -> tensor<complex<f64>> {
+  %c = mhlo.constant dense<(1.0, 1.0)> : tensor<complex<f64>>
+  %cos = mhlo.cosine %c : tensor<complex<f64>>
+  return %cos : tensor<complex<f64>>
+}
+
+// CHECK-LABEL: @cos
+// CHECK-NEXT: Results
+// CHECK-NEXT: 8.337300e-01-9.888977e-01i
+
+func.func @expm1() -> tensor<complex<f32>> {
+  // import numpy as np  -- not jax.numpy
+  // np.expm1(np.array(1e-6 + 1e-6j, dtype=np.complex64))
+  // Don't run this with jax.numpy, it returns 9.536743e-07+0.j.
+  %c = mhlo.constant dense<(1.0e-06, 1.0e-06)> : tensor<complex<f32>>
+  %expm1 = mhlo.exponential_minus_one %c : tensor<complex<f32>>
+  return %expm1 : tensor<complex<f32>>
+}
+
+// CHECK-LABEL: @expm1
+// CHECK-NEXT: Results
+// CHECK-NEXT: 1.000000e-06+1.000001e-06i
+
+func.func @expm1d() -> tensor<complex<f64>> {
+  %c = mhlo.constant dense<(1.0e-50, 1.0e-50)> : tensor<complex<f64>>
+  %expm1 = mhlo.exponential_minus_one %c : tensor<complex<f64>>
+  return %expm1 : tensor<complex<f64>>
+}
+
+// CHECK-LABEL: @expm1d
+// CHECK-NEXT: Results
+// CHECK-NEXT: 1.000000e-50+1.000000e-50i
+
+func.func @imag() -> tensor<f64> {
+  %c = mhlo.constant dense<(1.0, 2.0)> : tensor<complex<f64>>
+  %imag = mhlo.imag %c : (tensor<complex<f64>>) -> tensor<f64>
+  return %imag : tensor<f64>
+}
+
+// CHECK-LABEL: @imag
+// CHECK-NEXT: Results
+// CHECK-NEXT: 2.000000e+00
+
+func.func @log1pf() -> tensor<complex<f32>> {
+  %c = mhlo.constant dense<(1.0e-07, 1.0e-20)> : tensor<complex<f32>>
+  %cos = mhlo.log_plus_one %c : tensor<complex<f32>>
+  return %cos : tensor<complex<f32>>
+}
+
+// CHECK-LABEL: @log1p
+// CHECK-NEXT: Results
+// The accuracy of this is rather poor, but it matches numpy.
+// CHECK-NEXT: 1.192093e-07+9.999999e-21i
+
+func.func @log1pd() -> tensor<complex<f64>> {
+  %c = mhlo.constant dense<(1.0e-07, 1.0e-20)> : tensor<complex<f64>>
+  %cos = mhlo.log_plus_one %c : tensor<complex<f64>>
+  return %cos : tensor<complex<f64>>
+}
+
+// CHECK-LABEL: @log1pd
+// CHECK-NEXT: Results
+// CHECK-NEXT: 1.000000e-07+9.999999e-21i
+
+func.func @logistic() -> tensor<5xcomplex<f32>> {
+  %c = mhlo.constant dense<[(-1.0e5, 0.0), (-1.0, 0.0),
+      (0.0, 0.0), (1.0, 0.0), (1.0e5, 0.0)]> : tensor<5xcomplex<f32>>
+  %ret = mhlo.logistic %c : tensor<5xcomplex<f32>>
+  return %ret : tensor<5xcomplex<f32>>
+}
+
+// CHECK-LABEL: @logistic
+// CHECK-NEXT: Results
+// CHECK-NEXT: [0.000000e+00+0.000000e+00i,
+// CHECK-SAME:  2.689414e-01+0.000000e+00i,
+// CHECK-SAME:  5.000000e-01+0.000000e+00i,
+// CHECK-SAME:  7.310586e-01+0.000000e+00i,
+// CHECK-SAME:  1.000000e+00+0.000000e+00i]
+
+func.func @real() -> tensor<f64> {
+  %c = mhlo.constant dense<(1.0, 2.0)> : tensor<complex<f64>>
+  %real = mhlo.real %c : (tensor<complex<f64>>) -> tensor<f64>
+  return %real : tensor<f64>
+}
+
+// CHECK-LABEL: @real
+// CHECK-NEXT: Results
+// CHECK-NEXT: 1.000000e+00
+
+func.func @sin() -> tensor<complex<f64>> {
+  %c = mhlo.constant dense<(1.0, 1.0)> : tensor<complex<f64>>
+  %cos = mhlo.sine %c : tensor<complex<f64>>
+  return %cos : tensor<complex<f64>>
+}
+
+// CHECK-LABEL: @sin
+// CHECK-NEXT: Results
+// CHECK-NEXT: 1.298458e+00+6.349639e-01i
diff --git a/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/mhlo/concatenate.mlir b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/mhlo/concatenate.mlir
new file mode 100644
index 00000000000000..fb88f0c0da1fe3
--- /dev/null
+++ b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/mhlo/concatenate.mlir
@@ -0,0 +1,37 @@
+// RUN: mlir-interpreter-runner %s -run-all | FileCheck %s
+
+func.func @concat_1d()  -> tensor<3xi32> {
+  %a = mhlo.constant dense<[1]> : tensor<1xi32>
+  %b = mhlo.constant dense<[2, 3]> : tensor<2xi32>
+  %0 = "mhlo.concatenate"(%a, %b) { dimension = 0 : i64 }
+     : (tensor<1xi32>, tensor<2xi32>) -> tensor<3xi32>
+  func.return %0 : tensor<3xi32>
+}
+
+// CHECK-LABEL: @concat_1d
+// CHECK-NEXT: Results
+// CHECK-NEXT: [1, 2, 3]
+
+func.func @concat_dim0() -> tensor<4x2xi32> {
+  %a = mhlo.constant dense<1> : tensor<2x2xi32>
+  %b = mhlo.constant dense<2> : tensor<2x2xi32>
+  %0 = "mhlo.concatenate"(%a, %b) { dimension = 0 : i64 }
+     : (tensor<2x2xi32>, tensor<2x2xi32>) -> tensor<4x2xi32>
+  func.return %0 : tensor<4x2xi32>
+}
+
+// CHECK-LABEL: @concat_dim0
+// CHECK-NEXT: Results
+// CHECK-NEXT{LITERAL}: [[1, 1], [1, 1], [2, 2], [2, 2]]
+
+func.func @concat_dim1() -> tensor<2x4xi32> {
+  %a = mhlo.constant dense<1> : tensor<2x2xi32>
+  %b = mhlo.constant dense<2> : tensor<2x2xi32>
+  %0 = "mhlo.concatenate"(%a, %b) { dimension = 1 : i64 }
+     : (tensor<2x2xi32>, tensor<2x2xi32>) -> tensor<2x4xi32>
+  func.return %0 : tensor<2x4xi32>
+}
+
+// CHECK-LABEL: @concat_dim1
+// CHECK-NEXT: Results
+// CHECK-NEXT{LITERAL}: [[1, 1, 2, 2], [1, 1, 2, 2]]
diff --git a/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/mhlo/constant.mlir b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/mhlo/constant.mlir
new file mode 100644
index 00000000000000..77547aff85ba20
--- /dev/null
+++ b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/mhlo/constant.mlir
@@ -0,0 +1,25 @@
+// RUN: mlir-interpreter-runner %s -run-all | FileCheck %s
+
+func.func @main() -> (tensor<2x3xi32>, tensor<2x3xf32>, tensor<0x0x3xi16>, tensor<f64>) {
+  %i32 = mhlo.constant dense<[[0, 1, 2], [3, 4, 5]]> : tensor<2x3xi32>
+  %f32 = mhlo.constant dense<[[0.0, 0.1, 0.2], [0.3, 0.4, 0.5]]> : tensor<2x3xf32>
+  %empty = mhlo.constant dense<> : tensor<0x0x3xi16>
+  %scalar = mhlo.constant dense<3.14> : tensor<f64>
+  return %i32, %f32, %empty, %scalar : tensor<2x3xi32>, tensor<2x3xf32>, tensor<0x0x3xi16>, tensor<f64>
+}
+
+// CHECK-LABEL: @main
+// CHECK-NEXT: Results
+// CHECK-NEXT{LITERAL}: [[0, 1, 2], [3, 4, 5]]
+// CHECK-NEXT{LITERAL}: [[0.000000e+00, 1.000000e-01, 2.000000e-01], [3.000000e-01, 4.000000e-01, 5.000000e-01]]
+// CHECK-NEXT{LITERAL}: []
+// CHECK-NEXT{LITERAL}: 3.140000e+00
+
+func.func @ui8() -> tensor<ui8> {
+  %v = mhlo.constant dense<123> : tensor<ui8>
+  return %v : tensor<ui8>
+}
+
+// CHECK-LABEL: @ui8
+// CHECK-NEXT: Results
+// CHECK-NEXT: 123
\ No newline at end of file
diff --git a/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/mhlo/convert.mlir b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/mhlo/convert.mlir
new file mode 100644
index 00000000000000..fe897b2856a3b6
--- /dev/null
+++ b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/mhlo/convert.mlir
@@ -0,0 +1,21 @@
+// RUN: mlir-interpreter-runner %s -run-all | FileCheck %s
+
+func.func @convert_i1_to_f32() -> tensor<2xf32> {
+  %input = mhlo.constant dense<[true, false]> : tensor<2xi1>
+  %result = "mhlo.convert"(%input) : (tensor<2xi1>) -> tensor<2xf32>
+  func.return %result : tensor<2xf32>
+}
+
+// CHECK-LABEL: @convert_i1_to_f32
+// CHECK-NEXT: Results
+// CHECK-NEXT: [1.000000e+00, 0.000000e+00]
+
+func.func @convert_f32_to_i16() -> tensor<2xi16> {
+  %input = mhlo.constant dense<[1.4, 2.55]> : tensor<2xf32>
+  %result = "mhlo.convert"(%input) : (tensor<2xf32>) -> tensor<2xi16>
+  func.return %result : tensor<2xi16>
+}
+
+// CHECK-LABEL: @convert_f32_to_i16
+// CHECK-NEXT: Results
+// CHECK-NEXT: [1, 2]
diff --git a/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/mhlo/dot.mlir b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/mhlo/dot.mlir
new file mode 100644
index 00000000000000..66ed3956c0cbae
--- /dev/null
+++ b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/mhlo/dot.mlir
@@ -0,0 +1,37 @@
+// RUN: mlir-interpreter-runner %s -run-all | FileCheck %s
+
+func.func @dot_2d() -> tensor<2x2xi32> {
+  %lhs = mhlo.constant dense<[[1, 2], [3, 4]]> : tensor<2x2xi32>
+  %rhs = mhlo.constant dense<[[4, 5], [6, 7]]> : tensor<2x2xi32>
+  %dot = "mhlo.dot"(%lhs, %rhs)
+    : (tensor<2x2xi32>, tensor<2x2xi32>) -> tensor<2x2xi32>
+  return %dot : tensor<2x2xi32>
+}
+
+// CHECK-LABEL: @dot_2d
+// CHECK-NEXT: Results
+// CHECK-NEXT{LITERAL}: [[16, 19], [36, 43]]
+
+func.func @dot_2d_1d() -> tensor<2xi32> {
+  %lhs = mhlo.constant dense<[[1, 2], [3, 4]]> : tensor<2x2xi32>
+  %rhs = mhlo.constant dense<[4, 5]> : tensor<2xi32>
+  %dot = "mhlo.dot"(%lhs, %rhs)
+    : (tensor<2x2xi32>, tensor<2xi32>) -> tensor<2xi32>
+  return %dot : tensor<2xi32>
+}
+
+// CHECK-LABEL: @dot_2d_1d
+// CHECK-NEXT: Results
+// CHECK-NEXT{LITERAL}: [14, 32]
+
+func.func @dot_1d_1d() -> tensor<i32> {
+  %lhs = mhlo.constant dense<[1, 2]> : tensor<2xi32>
+  %rhs = mhlo.constant dense<[4, 5]> : tensor<2xi32>
+  %dot = "mhlo.dot"(%lhs, %rhs)
+    : (tensor<2xi32>, tensor<2xi32>) -> tensor<i32>
+  return %dot : tensor<i32>
+}
+
+// CHECK-LABEL: @dot_1d_1d
+// CHECK-NEXT: Results
+// CHECK-NEXT{LITERAL}: 14
\ No newline at end of file
diff --git a/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/mhlo/dot_general.mlir b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/mhlo/dot_general.mlir
new file mode 100644
index 00000000000000..1f83ad5ff63d62
--- /dev/null
+++ b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/mhlo/dot_general.mlir
@@ -0,0 +1,73 @@
+// RUN: mlir-interpreter-runner %s -run-all | FileCheck %s
+
+func.func @dot_general_2d() -> tensor<2x2xi32> {
+  %lhs = mhlo.constant dense<[[1, 2], [3, 4]]> : tensor<2x2xi32>
+  %rhs = mhlo.constant dense<[[4, 5], [6, 7]]> : tensor<2x2xi32>
+  %dot = "mhlo.dot_general"(%lhs, %rhs) {
+    dot_dimension_numbers = #mhlo.dot<
+      lhs_contracting_dimensions = [1],
+      rhs_contracting_dimensions = [0],
+      lhs_batching_dimensions = [],
+      rhs_batching_dimensions = []
+    >
+  } : (tensor<2x2xi32>, tensor<2x2xi32>) -> tensor<2x2xi32>
+  return %dot : tensor<2x2xi32>
+}
+
+// CHECK-LABEL: @dot_general_2d
+// CHECK-NEXT: Results
+// CHECK-NEXT{LITERAL}: [[16, 19], [36, 43]]
+
+func.func @dot_general_2d_2() -> tensor<2x2xi32> {
+  %lhs = mhlo.constant dense<[[1, 2], [3, 4]]> : tensor<2x2xi32>
+  %rhs = mhlo.constant dense<[[4, 5], [6, 7]]> : tensor<2x2xi32>
+  %dot = "mhlo.dot_general"(%lhs, %rhs) {
+    dot_dimension_numbers = #mhlo.dot<
+      lhs_contracting_dimensions = [0],
+      rhs_contracting_dimensions = [1],
+      lhs_batching_dimensions = [],
+      rhs_batching_dimensions = []
+    >
+  } : (tensor<2x2xi32>, tensor<2x2xi32>) -> tensor<2x2xi32>
+  return %dot : tensor<2x2xi32>
+}
+
+// CHECK-LABEL: @dot_general_2d_2
+// CHECK-NEXT: Results
+// CHECK-NEXT{LITERAL}: [[19, 27], [28, 40]]
+
+func.func @dot_general_2d_1d() -> tensor<2xi32> {
+  %lhs = mhlo.constant dense<[[1, 2], [3, 4]]> : tensor<2x2xi32>
+  %rhs = mhlo.constant dense<[4, 5]> : tensor<2xi32>
+  %dot = "mhlo.dot_general"(%lhs, %rhs) {
+    dot_dimension_numbers = #mhlo.dot<
+      lhs_contracting_dimensions = [0],
+      rhs_contracting_dimensions = [0],
+      lhs_batching_dimensions = [],
+      rhs_batching_dimensions = []
+    >
+  } : (tensor<2x2xi32>, tensor<2xi32>) -> tensor<2xi32>
+  return %dot : tensor<2xi32>
+}
+
+// CHECK-LABEL: @dot_general_2d_1d
+// CHECK-NEXT: Results
+// CHECK-NEXT{LITERAL}: [19, 28]
+
+func.func @dot_general_batch_only() -> tensor<2x2xi32> {
+  %lhs = mhlo.constant dense<[[1, 2], [3, 4]]> : tensor<2x2xi32>
+  %rhs = mhlo.constant dense<[4, 5]> : tensor<2xi32>
+  %dot = "mhlo.dot_general"(%lhs, %rhs) {
+    dot_dimension_numbers = #mhlo.dot<
+      lhs_contracting_dimensions = [],
+      rhs_contracting_dimensions = [],
+      lhs_batching_dimensions = [1],
+      rhs_batching_dimensions = [0]
+    >
+  } : (tensor<2x2xi32>, tensor<2xi32>) -> tensor<2x2xi32>
+  return %dot : tensor<2x2xi32>
+}
+
+// CHECK-LABEL: @dot_general_batch_only
+// CHECK-NEXT: Results
+// CHECK-NEXT{LITERAL}: [[4, 12], [10, 20]]
\ No newline at end of file
diff --git a/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/mhlo/dynamic_slice.mlir b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/mhlo/dynamic_slice.mlir
new file mode 100644
index 00000000000000..8708694d3a8504
--- /dev/null
+++ b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/mhlo/dynamic_slice.mlir
@@ -0,0 +1,32 @@
+// RUN: mlir-interpreter-runner %s -run-all | FileCheck %s
+
+func.func @dynamic_slice() -> tensor<1x2xi32> {
+  %cst = mhlo.constant dense<[[1, 2, 3, 4], [5, 6, 7, 8], [9, 10, 11, 12]]>
+    : tensor<3x4xi32>
+  %s0 = mhlo.constant dense<1> : tensor<i32>
+  %s1 = mhlo.constant dense<0> : tensor<i32>
+  %0 = "mhlo.dynamic_slice"(%cst, %s0, %s1) {
+    slice_sizes = dense<[1, 2]> : tensor<2xi64>
+  } : (tensor<3x4xi32>, tensor<i32>, tensor<i32>) -> tensor<1x2xi32>
+  func.return %0 : tensor<1x2xi32>
+}
+
+// CHECK-LABEL: @dynamic_slice
+// CHECK-NEXT: Results
+// CHECK-NEXT{LITERAL}: [[5, 6]]
+
+func.func @clamp_starts() -> tensor<2x3xi32> {
+
+  %cst = mhlo.constant dense<[[1, 2, 3, 4], [5, 6, 7, 8], [9, 10, 11, 12]]>
+    : tensor<3x4xi32>
+  %s0 = mhlo.constant dense<-10> : tensor<i32>
+  %s1 = mhlo.constant dense<10> : tensor<i32>
+  %0 = "mhlo.dynamic_slice"(%cst, %s0, %s1) {
+    slice_sizes = dense<[2, 3]> : tensor<2xi64>
+  } : (tensor<3x4xi32>, tensor<i32>, tensor<i32>) -> tensor<2x3xi32>
+  func.return %0 : tensor<2x3xi32>
+}
+
+// CHECK-LABEL: @clamp_starts
+// CHECK-NEXT: Results
+// CHECK-NEXT{LITERAL}: [[2, 3, 4], [6, 7, 8]]
diff --git a/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/mhlo/dynamic_update_slice.mlir b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/mhlo/dynamic_update_slice.mlir
new file mode 100644
index 00000000000000..481389e156e927
--- /dev/null
+++ b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/mhlo/dynamic_update_slice.mlir
@@ -0,0 +1,34 @@
+// RUN: mlir-interpreter-runner %s -run-all | FileCheck %s
+
+func.func @dynamic_update_slice() -> tensor<3x4xi32> {
+  %cst = mhlo.constant dense<[[1, 2, 3, 4], [5, 6, 7, 8], [9, 10, 11, 12]]>
+    : tensor<3x4xi32>
+  %v = mhlo.constant dense<[[13, 14]]> : tensor<1x2xi32>
+  %s0 = mhlo.constant dense<1> : tensor<i32>
+  %s1 = mhlo.constant dense<0> : tensor<i32>
+  %0 = "mhlo.dynamic_update_slice"(%cst, %v, %s0, %s1)
+    : (tensor<3x4xi32>, tensor<1x2xi32>, tensor<i32>, tensor<i32>)
+      -> tensor<3x4xi32>
+  func.return %0 : tensor<3x4xi32>
+}
+
+// CHECK-LABEL: @dynamic_update_slice
+// CHECK-NEXT: Results
+// CHECK-NEXT{LITERAL}: [[1, 2, 3, 4], [13, 14, 7, 8], [9, 10, 11, 12]]
+
+func.func @clamp_starts() -> tensor<3x4xi32> {
+  %cst = mhlo.constant dense<[[1, 2, 3, 4], [5, 6, 7, 8], [9, 10, 11, 12]]>
+    : tensor<3x4xi32>
+  %v = mhlo.constant dense<[[13, 14]]> : tensor<1x2xi32>
+  %s0 = mhlo.constant dense<-10> : tensor<i32>
+  %s1 = mhlo.constant dense<10> : tensor<i32>
+  %0 = "mhlo.dynamic_update_slice"(%cst, %v, %s0, %s1) {
+    slice_sizes = dense<[2, 3]> : tensor<2xi64>
+  } : (tensor<3x4xi32>, tensor<1x2xi32>, tensor<i32>, tensor<i32>)
+      -> tensor<3x4xi32>
+  func.return %0 : tensor<3x4xi32>
+}
+
+// CHECK-LABEL: @clamp_starts
+// CHECK-NEXT: Results
+// CHECK-NEXT{LITERAL}: [[1, 2, 13, 14], [5, 6, 7, 8], [9, 10, 11, 12]]
diff --git a/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/mhlo/float_math.mlir b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/mhlo/float_math.mlir
new file mode 100644
index 00000000000000..c713e4bf46394d
--- /dev/null
+++ b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/mhlo/float_math.mlir
@@ -0,0 +1,199 @@
+// RUN: mlir-interpreter-runner %s -run-all | FileCheck %s
+
+func.func @atan2() -> tensor<1xf32> {
+  %c10 = mhlo.constant dense<10.0> : tensor<1xf32>
+  %c1 = mhlo.constant dense<1.0> : tensor<1xf32>
+  %ret = mhlo.atan2 %c10, %c1 : tensor<1xf32>
+  return %ret : tensor<1xf32>
+}
+
+// CHECK-LABEL: @atan2
+// CHECK-NEXT: Results
+// CHECK-NEXT: [1.471128e+00]
+
+func.func @cbrt() -> tensor<1xf32> {
+  %c-27 = mhlo.constant dense<-27.0> : tensor<1xf32>
+  %ret = mhlo.cbrt %c-27 : tensor<1xf32>
+  return %ret : tensor<1xf32>
+}
+
+// CHECK-LABEL: @cbrt
+// CHECK-NEXT: Results
+// CHECK-NEXT: [-3.000000e+00]
+
+func.func @ceil() -> tensor<1xf32> {
+  %c = mhlo.constant dense<0.123> : tensor<1xf32>
+  %ret = mhlo.ceil %c : tensor<1xf32>
+  return %ret : tensor<1xf32>
+}
+
+// CHECK-LABEL: @ceil
+// CHECK-NEXT: Results
+// CHECK-NEXT: [1.000000e+00]
+
+func.func @complex() -> tensor<complex<f32>> {
+  %c1 = mhlo.constant dense<1.0> : tensor<f32>
+  %c2 = mhlo.constant dense<2.0> : tensor<f32>
+  %ret = mhlo.complex %c1, %c2 : tensor<complex<f32>>
+  return %ret : tensor<complex<f32>>
+}
+
+// CHECK-LABEL: @complex
+// CHECK-NEXT: Results
+// CHECK-NEXT: <complex<f32>>: 1.000000e+00+2.000000e+00i
+
+func.func @exp() -> tensor<1xf32> {
+  %c = mhlo.constant dense<0.0> : tensor<1xf32>
+  %ret = mhlo.exponential %c : tensor<1xf32>
+  return %ret : tensor<1xf32>
+}
+
+// CHECK-LABEL: @exp
+// CHECK-NEXT: Results
+// CHECK-NEXT: [1.000000e+00]
+
+func.func @floor() -> tensor<1xf32> {
+  %c = mhlo.constant dense<3.123> : tensor<1xf32>
+  %ret = mhlo.floor %c : tensor<1xf32>
+  return %ret : tensor<1xf32>
+}
+
+// CHECK-LABEL: @floor
+// CHECK-NEXT: Results
+// CHECK-NEXT: [3.000000e+00]
+
+func.func @is_finite() -> tensor<3xi1> {
+  %c = mhlo.constant dense<[0x7FC00000, 2.0, 0x7F800000]> : tensor<3xf32>
+  %is_finite = mhlo.is_finite %c : (tensor<3xf32>) -> tensor<3xi1>
+  return %is_finite : tensor<3xi1>
+}
+
+// CHECK-LABEL: @is_finite
+// CHECK-NEXT: Results
+// CHECK-NEXT: [false, true, false]
+
+func.func @log() -> tensor<1xf32> {
+  %c = mhlo.constant dense<1.0> : tensor<1xf32>
+  %ret = mhlo.log %c : tensor<1xf32>
+  return %ret : tensor<1xf32>
+}
+
+// CHECK-LABEL: @log
+// CHECK-NEXT: Results
+// CHECK-NEXT: [0.000000e+00]
+
+func.func @logistic() -> tensor<5xf32> {
+  %c = mhlo.constant dense<[-1.0e5, -1.0, 0.0, 1.0, 1.0e5]> : tensor<5xf32>
+  %ret = mhlo.logistic %c : tensor<5xf32>
+  return %ret : tensor<5xf32>
+}
+
+// CHECK-LABEL: @logistic
+// CHECK-NEXT: Results
+// CHECK-NEXT: [0.000000e+00, 2.689414e-01, 5.000000e-01, 7.310586e-01, 1.000000e+00]
+
+func.func @maximum() -> tensor<f32> {
+  %c1 = mhlo.constant dense<1.0> : tensor<f32>
+  %c2 = mhlo.constant dense<2.0> : tensor<f32>
+  %ret = mhlo.maximum %c1, %c2 : tensor<f32>
+  return %ret : tensor<f32>
+}
+
+// CHECK-LABEL: @maximum
+// CHECK-NEXT: Results
+// CHECK-NEXT: 2.000000e+00
+
+func.func @minimum() -> tensor<f32> {
+  %c1 = mhlo.constant dense<1.0> : tensor<f32>
+  %c2 = mhlo.constant dense<2.0> : tensor<f32>
+  %ret = mhlo.minimum %c1, %c2 : tensor<f32>
+  return %ret : tensor<f32>
+}
+
+// CHECK-LABEL: @minimum
+// CHECK-NEXT: Results
+// CHECK-NEXT: 1.000000e+00
+
+func.func @pow() -> tensor<f32> {
+  %c2 = mhlo.constant dense<2.0> : tensor<f32>
+  %c5 = mhlo.constant dense<5.0> : tensor<f32>
+  %pow = mhlo.power %c2, %c5 : tensor<f32>
+  return %pow : tensor<f32>
+}
+
+// CHECK-LABEL: @pow
+// CHECK-NEXT: Results
+// CHECK-NEXT: 3.200000e+01
+
+func.func @rem() -> tensor<8xf32> {
+  %0 = mhlo.constant dense<[-2.5, 2.25, -10.0, 6.0, 3.0, 3.0, -1.0, -8.0]> : tensor<8xf32>
+  %1 = mhlo.constant dense<[10.0, 1.0, 10.0, -6.0, 2.0, -2.0, 7.0, -4.0]> : tensor<8xf32>
+  %2 = mhlo.remainder %0, %1 : tensor<8xf32>
+  func.return %2 : tensor<8xf32>
+}
+
+// CHECK-LABEL: @rem
+// CHECK-NEXT: Results
+// CHECK-NEXT: [-2.500000e+00, 2.500000e-01, -0.000000e+00, 0.000000e+00, 1.000000e+00, 1.000000e+00, -1.000000e+00, -0.000000e+00]
+
+func.func @rem_inf() -> (tensor<1xf32>, tensor<1xf32>) {
+  %0 = mhlo.constant dense<[1.0]> : tensor<1xf32>
+  %1 = mhlo.constant dense<[0x7F800000]> : tensor<1xf32>
+  %2 = mhlo.remainder %0, %1 : tensor<1xf32>
+  func.return %2, %1 : tensor<1xf32>, tensor<1xf32>
+}
+
+// CHECK-LABEL: @rem_inf
+// CHECK-NEXT: Results
+// CHECK-NEXT: [1.000000e+00]
+// CHECK-NEXT: [INF]
+
+func.func @round_nearest_afz() -> tensor<4xf32> {
+  %c = mhlo.constant dense<[-1.5, -0.5, 0.5, 1.5]> : tensor<4xf32>
+  %ret = mhlo.round_nearest_afz %c : tensor<4xf32>
+  return %ret : tensor<4xf32>
+}
+
+// CHECK-LABEL: @round_nearest_afz
+// CHECK-NEXT: Results
+// CHECK-NEXT: [-2.000000e+00, -1.000000e+00, 1.000000e+00, 2.000000e+00]
+
+func.func @round_nearest_even() -> tensor<5xf32> {
+  %c = mhlo.constant dense<[-1.5, -0.5, 0.5, 0.6, 1.5]> : tensor<5xf32>
+  %ret = mhlo.round_nearest_even %c : tensor<5xf32>
+  return %ret : tensor<5xf32>
+}
+
+// CHECK-LABEL: @round_nearest_even
+// CHECK-NEXT: Results
+// CHECK-NEXT: [-2.000000e+00, -0.000000e+00, 0.000000e+00, 1.000000e+00, 2.000000e+00]
+
+func.func @rsqrt() -> tensor<1xf32> {
+  %c4 = mhlo.constant dense<4.0> : tensor<1xf32>
+  %ret = mhlo.rsqrt %c4 : tensor<1xf32>
+  return %ret : tensor<1xf32>
+}
+
+// CHECK-LABEL: @rsqrt
+// CHECK-NEXT: Results
+// CHECK-NEXT: [5.000000e-01]
+
+func.func @sign() -> tensor<3xf32> {
+  %c = mhlo.constant dense<[-1.0, 2.0, 0x7F800000]> : tensor<3xf32>
+  %ret = mhlo.sign %c : tensor<3xf32>
+  return %ret : tensor<3xf32>
+}
+
+// CHECK-LABEL: @sign
+// CHECK-NEXT: Results
+// CHECK-NEXT: [-1.000000e+00, 1.000000e+00, 1.000000e+00]
+
+func.func @tanh() -> tensor<1xf32> {
+  %c = mhlo.constant dense<[1.0]> : tensor<1xf32>
+  %ret = mhlo.tanh %c : tensor<1xf32>
+  return %ret : tensor<1xf32>
+}
+
+// CHECK-LABEL: @tanh
+// CHECK-NEXT: Results
+// CHECK-NEXT: [7.615942e-01]
diff --git a/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/mhlo/gather.mlir b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/mhlo/gather.mlir
new file mode 100644
index 00000000000000..c6992ba6d527dd
--- /dev/null
+++ b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/mhlo/gather.mlir
@@ -0,0 +1,58 @@
+// RUN: mlir-interpreter-runner %s -run-all | FileCheck %s
+
+func.func @bounds_check() -> tensor<4x3xi32> {
+  %operand = mhlo.constant dense<[0, 1, 2, 3, 4, 5, 6, 7, 8, 9]> : tensor<10xi32>
+  %indices = mhlo.constant dense<[[1], [8], [-3]]> : tensor<3x1xi32>
+  %gather = "mhlo.gather"(%operand, %indices) {
+    dimension_numbers = #mhlo.gather<
+      collapsed_slice_dims = [],
+      index_vector_dim = 1,
+      offset_dims = [0],
+      start_index_map = [0]
+    >,
+    slice_sizes = dense<[4]> : tensor<1xi64>
+  } : (tensor<10xi32>, tensor<3x1xi32>) -> tensor<4x3xi32>
+  return %gather : tensor<4x3xi32>
+}
+
+// CHECK-LABEL: @bounds_check
+// CHECK-NEXT: Results
+// CHECK-NEXT{LITERAL}: [[1, 6, 0], [2, 7, 1], [3, 8, 2], [4, 9, 3]]
+
+func.func @gather_2d() -> tensor<4x2xi32> {
+  // operand = np.arange(1, 16).reshape([5, 3, 1])
+  // indices = np.array([[0, 0], [1, 0], [4, 3], [-1, -1]])
+  // lax.gather(operand, indices, lax.GatherDimensionNumbers(offset_dims=(1,),
+  //   collapsed_slice_dims=(1,2), start_index_map=(0,1,)), slice_sizes=[2,1,1]))
+  %operand = arith.constant dense<[
+    [[1], [2], [3]],
+    [[4], [5], [6]],
+    [[7], [8], [9]],
+    [[10], [11], [12]],
+    [[13], [14], [15]]
+  ]> : tensor<5x3x1xi32>
+
+  %indices = arith.constant dense<[
+    [0, 0],
+    [1, 0],
+    [4, 3],
+    [-1, -1]
+  ]> : tensor<4x2xi64>
+
+  %0 = "mhlo.gather"(%operand, %indices) {
+    dimension_numbers = #mhlo.gather<
+      collapsed_slice_dims = [1, 2],
+      index_vector_dim = 1,
+      offset_dims = [1],
+      start_index_map = [0, 1]
+    >,
+    indices_are_sorted = false,
+    slice_sizes = dense<[2, 1, 1]> : tensor<3xi64>
+  } : (tensor<5x3x1xi32>, tensor<4x2xi64>) -> tensor<4x2xi32>
+
+  func.return %0 : tensor<4x2xi32>
+}
+
+// CHECK-LABEL: @gather_2d
+// CHECK-NEXT: Results
+// CHECK-NEXT{LITERAL}: [[1, 4], [4, 7], [12, 15], [1, 4]]
diff --git a/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/mhlo/int_math.mlir b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/mhlo/int_math.mlir
new file mode 100644
index 00000000000000..999f6f6996e8b6
--- /dev/null
+++ b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/mhlo/int_math.mlir
@@ -0,0 +1,358 @@
+// RUN: mlir-interpreter-runner %s -run-all | FileCheck %s
+
+func.func @abs() -> tensor<1xi32> {
+  %c-3 = mhlo.constant dense<-3> : tensor<1xi32>
+  %ret = mhlo.abs %c-3 : tensor<1xi32>
+  return %ret : tensor<1xi32>
+}
+
+// CHECK-LABEL: @abs
+// CHECK-NEXT: Results
+// CHECK-NEXT: [3]
+
+func.func @add_tensor() -> tensor<2x3xi32> {
+  %lhs = mhlo.constant dense<[[0, 1, 2], [3, 4, 5]]> : tensor<2x3xi32>
+  %rhs = mhlo.constant dense<[[10, 20, 30], [40, 50, 60]]> : tensor<2x3xi32>
+  %result = mhlo.add %lhs, %rhs : tensor<2x3xi32>
+  return %result : tensor<2x3xi32>
+}
+
+// CHECK-LABEL: @add_tensor
+// CHECK-NEXT: Results
+// CHECK-NEXT{LITERAL}: [[10, 21, 32], [43, 54, 65]]
+
+func.func @add_scalar_i8() -> tensor<i8> {
+  %lhs = mhlo.constant dense<40> : tensor<i8>
+  %rhs = mhlo.constant dense<2> : tensor<i8>
+  %result = mhlo.add %lhs, %rhs : tensor<i8>
+  return %result : tensor<i8>
+}
+
+// CHECK-LABEL: @add_scalar_i8
+// CHECK-NEXT: Results
+// CHECK-NEXT: <i8>: 42
+
+func.func @add_scalar_i16() -> tensor<i16> {
+  %lhs = mhlo.constant dense<40> : tensor<i16>
+  %rhs = mhlo.constant dense<2> : tensor<i16>
+  %result = mhlo.add %lhs, %rhs : tensor<i16>
+  return %result : tensor<i16>
+}
+
+// CHECK-LABEL: @add_scalar_i16
+// CHECK-NEXT: Results
+// CHECK-NEXT: <i16>: 42
+
+func.func @add_scalar_i32() -> tensor<i32> {
+  %lhs = mhlo.constant dense<40> : tensor<i32>
+  %rhs = mhlo.constant dense<2> : tensor<i32>
+  %result = mhlo.add %lhs, %rhs : tensor<i32>
+  return %result : tensor<i32>
+}
+
+// CHECK-LABEL: @add_scalar_i32
+// CHECK-NEXT: Results
+// CHECK-NEXT: <i32>: 42
+
+func.func @and() -> tensor<1xi32> {
+  %c63 = mhlo.constant dense<63> : tensor<1xi32>
+  %c131 = mhlo.constant dense<131> : tensor<1xi32>
+  %ret = mhlo.and %c63, %c131 : tensor<1xi32>
+  return %ret : tensor<1xi32>
+}
+
+// CHECK-LABEL: @and
+// CHECK-NEXT: Results
+// CHECK-NEXT: [3]
+
+func.func @and_i16() -> tensor<i16> {
+  %c63 = mhlo.constant dense<63> : tensor<i16>
+  %c131 = mhlo.constant dense<131> : tensor<i16>
+  %ret = mhlo.and %c63, %c131 : tensor<i16>
+  return %ret : tensor<i16>
+}
+
+// CHECK-LABEL: @and_i16
+// CHECK-NEXT: Results
+// CHECK-NEXT: <i16>: 3
+
+func.func @clz_negative()
+    -> (tensor<1xi8>, tensor<1xi16>, tensor<1xi32>, tensor<1xi64>) {
+  %c-1_8 = mhlo.constant dense<-1> : tensor<1xi8>
+  %c-2_16 = mhlo.constant dense<-2> : tensor<1xi16>
+  %c-4_32 = mhlo.constant dense<-4> : tensor<1xi32>
+  %c-8_64 = mhlo.constant dense<-8> : tensor<1xi64>
+  %clz_8 = mhlo.count_leading_zeros %c-1_8 : tensor<1xi8>
+  %clz_16 = mhlo.count_leading_zeros %c-2_16 : tensor<1xi16>
+  %clz_32 = mhlo.count_leading_zeros %c-4_32 : tensor<1xi32>
+  %clz_64 = mhlo.count_leading_zeros %c-8_64 : tensor<1xi64>
+  return %clz_8, %clz_16, %clz_32, %clz_64
+    : tensor<1xi8>, tensor<1xi16>, tensor<1xi32>, tensor<1xi64>
+}
+
+// CHECK-LABEL: @clz_negative
+// CHECK-NEXT: Results
+// CHECK-NEXT: [0]
+// CHECK-NEXT: [0]
+// CHECK-NEXT: [0]
+// CHECK-NEXT: [0]
+
+func.func @clz_signed()
+    -> (tensor<1xi8>, tensor<1xi16>, tensor<1xi32>, tensor<1xi64>) {
+  %c1_8 = mhlo.constant dense<1> : tensor<1xi8>
+  %c2_16 = mhlo.constant dense<2> : tensor<1xi16>
+  %c4_32 = mhlo.constant dense<4> : tensor<1xi32>
+  %c8_64 = mhlo.constant dense<8> : tensor<1xi64>
+  %clz_8 = mhlo.count_leading_zeros %c1_8 : tensor<1xi8>
+  %clz_16 = mhlo.count_leading_zeros %c2_16 : tensor<1xi16>
+  %clz_32 = mhlo.count_leading_zeros %c4_32 : tensor<1xi32>
+  %clz_64 = mhlo.count_leading_zeros %c8_64 : tensor<1xi64>
+  return %clz_8, %clz_16, %clz_32, %clz_64
+    : tensor<1xi8>, tensor<1xi16>, tensor<1xi32>, tensor<1xi64>
+}
+
+// CHECK-LABEL: @clz_signed
+// CHECK-NEXT: Results
+// CHECK-NEXT: [7]
+// CHECK-NEXT: [14]
+// CHECK-NEXT: [29]
+// CHECK-NEXT: [60]
+
+func.func @clz_unsigned()
+    -> (tensor<1xui8>, tensor<1xui16>, tensor<1xui32>, tensor<1xui64>) {
+  %c1_8 = mhlo.constant dense<1> : tensor<1xui8>
+  %c2_16 = mhlo.constant dense<2> : tensor<1xui16>
+  %c4_32 = mhlo.constant dense<4> : tensor<1xui32>
+  %c8_64 = mhlo.constant dense<8> : tensor<1xui64>
+  %clz_8 = mhlo.count_leading_zeros %c1_8 : tensor<1xui8>
+  %clz_16 = mhlo.count_leading_zeros %c2_16 : tensor<1xui16>
+  %clz_32 = mhlo.count_leading_zeros %c4_32 : tensor<1xui32>
+  %clz_64 = mhlo.count_leading_zeros %c8_64 : tensor<1xui64>
+  return %clz_8, %clz_16, %clz_32, %clz_64
+    : tensor<1xui8>, tensor<1xui16>, tensor<1xui32>, tensor<1xui64>
+}
+
+// CHECK-LABEL: @clz_unsigned
+// CHECK-NEXT: Results
+// CHECK-NEXT: [7]
+// CHECK-NEXT: [14]
+// CHECK-NEXT: [29]
+// CHECK-NEXT: [60]
+
+func.func @divide() -> tensor<1xi32> {
+  %c-10 = mhlo.constant dense<-10> : tensor<1xi32>
+  %c-2 = mhlo.constant dense<-2> : tensor<1xi32>
+  %ret = mhlo.divide %c-10, %c-2 : tensor<1xi32>
+  return %ret : tensor<1xi32>
+}
+
+// CHECK-LABEL: @divide
+// CHECK-NEXT: Results
+// CHECK-NEXT: [5]
+
+func.func @subtract() -> tensor<1xi32> {
+  %c10 = mhlo.constant dense<10> : tensor<1xi32>
+  %c3 = mhlo.constant dense<3> : tensor<1xi32>
+  %ret = mhlo.subtract %c10, %c3 : tensor<1xi32>
+  return %ret : tensor<1xi32>
+}
+
+// CHECK-LABEL: @subtract
+// CHECK-NEXT: Results
+// CHECK-NEXT: [7]
+
+func.func @or() -> tensor<1xi32> {
+  %c3 = mhlo.constant dense<3> : tensor<1xi32>
+  %c10 = mhlo.constant dense<10> : tensor<1xi32>
+  %ret = mhlo.or %c3, %c10 : tensor<1xi32>
+  return %ret : tensor<1xi32>
+}
+
+// CHECK-LABEL: @or
+// CHECK-NEXT: Results
+// CHECK-NEXT: [11]
+
+func.func @max_scalar() -> tensor<i32> {
+  %lhs = mhlo.constant dense<40> : tensor<i32>
+  %rhs = mhlo.constant dense<2> : tensor<i32>
+  %result = mhlo.maximum %lhs, %rhs : tensor<i32>
+  return %result : tensor<i32>
+}
+
+// CHECK-LABEL: @max_scalar
+// CHECK-NEXT: Results
+// CHECK-NEXT: 40
+
+func.func @multiply() -> tensor<1xi32> {
+  %c3 = mhlo.constant dense<3> : tensor<1xi32>
+  %c-5 = mhlo.constant dense<-5> : tensor<1xi32>
+  %ret = mhlo.multiply %c3, %c-5 : tensor<1xi32>
+  return %ret : tensor<1xi32>
+}
+
+// CHECK-LABEL: @multiply
+// CHECK-NEXT: Results
+// CHECK-NEXT: [-15]
+
+func.func @multiply_scalar_i16() -> tensor<i16> {
+  %lhs = mhlo.constant dense<40> : tensor<i16>
+  %rhs = mhlo.constant dense<2> : tensor<i16>
+  %result = mhlo.multiply %lhs, %rhs : tensor<i16>
+  return %result : tensor<i16>
+}
+
+// CHECK-LABEL: @multiply_scalar_i16
+// CHECK-NEXT: Results
+// CHECK-NEXT: <i16>: 80
+
+func.func @multiply_scalar_ui16() -> tensor<ui16> {
+  %lhs = mhlo.constant dense<40> : tensor<ui16>
+  %rhs = mhlo.constant dense<2> : tensor<ui16>
+  %result = mhlo.multiply %lhs, %rhs : tensor<ui16>
+  return %result : tensor<ui16>
+}
+
+// CHECK-LABEL: @multiply_scalar_ui16
+// CHECK-NEXT: Results
+// CHECK-NEXT: <ui16>: 80
+
+func.func @not_i1() -> tensor<2xi1> {
+  %cst = mhlo.constant dense<[false, true]> : tensor<2xi1>
+  %not = mhlo.not %cst : tensor<2xi1>
+  return %not : tensor<2xi1>
+}
+
+// CHECK-LABEL: @not_i1
+// CHECK-NEXT: Results
+// CHECK-NEXT: [true, false]
+
+func.func @not_ui16() -> tensor<ui16> {
+  %cst = mhlo.constant dense<1> : tensor<ui16>
+  %not = mhlo.not %cst : tensor<ui16>
+  return %not : tensor<ui16>
+}
+
+// CHECK-LABEL: @not_ui16
+// CHECK-NEXT: Results
+// CHECK-NEXT: 65534
+
+func.func @popcnt_negative()
+    -> (tensor<1xi8>, tensor<1xi16>, tensor<1xi32>, tensor<1xi64>) {
+  %c-1_8 = mhlo.constant dense<-1> : tensor<1xi8>
+  %c-2_16 = mhlo.constant dense<-2> : tensor<1xi16>
+  %c-4_32 = mhlo.constant dense<-4> : tensor<1xi32>
+  %c-8_64 = mhlo.constant dense<-8> : tensor<1xi64>
+  %pop_8 = mhlo.popcnt %c-1_8 : tensor<1xi8>
+  %pop_16 = mhlo.popcnt %c-2_16 : tensor<1xi16>
+  %pop_32 = mhlo.popcnt %c-4_32 : tensor<1xi32>
+  %pop_64 = mhlo.popcnt %c-8_64 : tensor<1xi64>
+  return %pop_8, %pop_16, %pop_32, %pop_64
+    : tensor<1xi8>, tensor<1xi16>, tensor<1xi32>, tensor<1xi64>
+}
+
+// CHECK-LABEL: @popcnt_negative
+// CHECK-NEXT: Results
+// CHECK-NEXT: [8]
+// CHECK-NEXT: [15]
+// CHECK-NEXT: [30]
+// CHECK-NEXT: [61]
+
+func.func @pow() -> tensor<ui16> {
+  %c2 = mhlo.constant dense<2> : tensor<ui16>
+  %c5 = mhlo.constant dense<5> : tensor<ui16>
+  %pow = mhlo.power %c2, %c5 : tensor<ui16>
+  return %pow : tensor<ui16>
+}
+
+// CHECK-LABEL: @pow
+// CHECK-NEXT: Results
+// CHECK-NEXT: <ui16>: 32
+
+func.func @pow_negative() -> tensor<3xi64> {
+  %c = mhlo.constant dense<[0, 1, 2]> : tensor<3xi64>
+  %c2 = mhlo.constant dense<-1> : tensor<3xi64>
+  %pow = mhlo.power %c, %c2 : tensor<3xi64>
+  return %pow : tensor<3xi64>
+}
+
+// CHECK-LABEL: @pow_negative
+// CHECK-NEXT: Results
+// CHECK-NEXT: [0, 1, 0]
+
+func.func @pow_non_double() -> tensor<i64> {
+  %c3 = mhlo.constant dense<3> : tensor<i64>
+  %c35 = mhlo.constant dense<35> : tensor<i64>
+  // The result of this operation cannot be represented by a double.
+  %pow = mhlo.power %c3, %c35 : tensor<i64>
+  return %pow : tensor<i64>
+}
+
+// CHECK-LABEL: @pow_non_double
+// CHECK-NEXT: Results
+// CHECK-NEXT: 50031545098999707
+
+func.func @rem() -> tensor<4xi32> {
+  %0 = mhlo.constant dense<[5, 66, 5, -1]> : tensor<4xi32>
+  %1 = mhlo.constant dense<[3, 5, 1, -2]> : tensor<4xi32>
+  %2 = mhlo.remainder %0, %1 : tensor<4xi32>
+  func.return %2 : tensor<4xi32>
+}
+
+// CHECK-LABEL: @rem
+// CHECK-NEXT: Results
+// CHECK-NEXT: [2, 1, 0, -1]
+
+func.func @shift_left() -> tensor<2xi32> {
+  %0 = mhlo.constant dense<[3, 7]> : tensor<2xi32>
+  %1 = mhlo.constant dense<[4, 31]> : tensor<2xi32>
+  %ret = mhlo.shift_left %0, %1 : tensor<2xi32>
+  func.return %ret : tensor<2xi32>
+}
+
+// CHECK-LABEL: @shift_left
+// CHECK-NEXT: Results
+// CHECK-NEXT: [48, -2147483648]
+
+func.func @shift_right_arith() -> tensor<2xi32> {
+  %0 = mhlo.constant dense<[100, -100]> : tensor<2xi32>
+  %1 = mhlo.constant dense<[2, 2]> : tensor<2xi32>
+  %ret = mhlo.shift_right_arithmetic %0, %1 : tensor<2xi32>
+  func.return %ret : tensor<2xi32>
+}
+
+// CHECK-LABEL: @shift_right_arith
+// CHECK-NEXT: Results
+// CHECK-NEXT: [25, -25]
+
+func.func @shift_right_arith_ui32() -> tensor<ui32> {
+  %0 = mhlo.constant dense<100> : tensor<ui32>
+  %1 = mhlo.constant dense<2> : tensor<ui32>
+  %ret = mhlo.shift_right_arithmetic %0, %1 : tensor<ui32>
+  func.return %ret : tensor<ui32>
+}
+
+// CHECK-LABEL: @shift_right_arith_ui32
+// CHECK-NEXT: Results
+// CHECK-NEXT: 25
+
+func.func @shift_right_logical() -> tensor<2xi32> {
+  %0 = mhlo.constant dense<[100, -100]> : tensor<2xi32>
+  %1 = mhlo.constant dense<[2, 2]> : tensor<2xi32>
+  %ret = mhlo.shift_right_logical %0, %1 : tensor<2xi32>
+  func.return %ret : tensor<2xi32>
+}
+
+// CHECK-LABEL: @shift_right_logical
+// CHECK-NEXT: Results
+// CHECK-NEXT: [25, 1073741799]
+
+func.func @shift_right_logical_ui32() -> tensor<ui32> {
+  %0 = mhlo.constant dense<100> : tensor<ui32>
+  %1 = mhlo.constant dense<2> : tensor<ui32>
+  %ret = mhlo.shift_right_logical %0, %1 : tensor<ui32>
+  func.return %ret : tensor<ui32>
+}
+
+// CHECK-LABEL: @shift_right_logical_ui32
+// CHECK-NEXT: Results
+// CHECK-NEXT: 25
diff --git a/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/mhlo/iota.mlir b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/mhlo/iota.mlir
new file mode 100644
index 00000000000000..e2d12d998e0bb5
--- /dev/null
+++ b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/mhlo/iota.mlir
@@ -0,0 +1,30 @@
+// RUN: mlir-interpreter-runner %s -run-all | FileCheck %s
+
+func.func @iota_f32() -> tensor<1x2x3x4xf32> {
+  %result = "mhlo.iota"() {
+    iota_dimension = 2 : i64
+  } : () -> tensor<1x2x3x4xf32>
+  func.return %result : tensor<1x2x3x4xf32>
+}
+
+// CHECK-LABEL: @iota_f32
+// CHECK-NEXT: Results
+// CHECK-NEXT{LITERAL}: [[[[0.000000e+00, 0.000000e+00, 0.000000e+00, 0.000000e+00],
+// CHECK{LITERAL}:         [1.000000e+00, 1.000000e+00, 1.000000e+00, 1.000000e+00],
+// CHECK{LITERAL}:         [2.000000e+00, 2.000000e+00, 2.000000e+00, 2.000000e+00]]]]
+
+func.func @iota_i32() -> tensor<1x2x3x4xi32> {
+  %result = "mhlo.iota"() {
+    iota_dimension = 3 : i64
+  } : () -> tensor<1x2x3x4xi32>
+  func.return %result : tensor<1x2x3x4xi32>
+}
+
+// CHECK-LABEL: @iota_i32
+// CHECK-NEXT: Results
+// CHECK-NEXT{LITERAL}: [[[[0, 1, 2, 3],
+// CHECK{LITERAL}          [0, 1, 2, 3],
+// CHECK{LITERAL}          [0, 1, 2, 3]],
+// CHECK{LITERAL}         [[0, 1, 2, 3],
+// CHECK{LITERAL}          [0, 1, 2, 3],
+// CHECK{LITERAL}          [0, 1, 2, 3]]]]
diff --git a/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/mhlo/pad.mlir b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/mhlo/pad.mlir
new file mode 100644
index 00000000000000..1f563b2d5ec57f
--- /dev/null
+++ b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/mhlo/pad.mlir
@@ -0,0 +1,56 @@
+// RUN: mlir-interpreter-runner %s -run-all | FileCheck %s
+
+func.func @static_pad() -> tensor<2x4x7xi32> {
+  // lax.pad(np.array([[[1,2,3],[4,5,6]]]), 42,
+  //                  [(0, 1, 0), (1, 1, 0), (2, 0, 1)])
+  %cst = mhlo.constant dense<[[[1,2,3],[4,5,6]]]> : tensor<1x2x3xi32>
+  %pad_value = mhlo.constant dense<42> : tensor<i32>
+  %0 = "mhlo.pad"(%cst, %pad_value) {
+    edge_padding_low = dense<[0, 1, 2]> : tensor<3xi64>,
+    edge_padding_high = dense<[1, 1, 0]> : tensor<3xi64>,
+    interior_padding = dense<[0, 0, 1]> : tensor<3xi64>
+  } : (tensor<1x2x3xi32>, tensor<i32>) -> tensor<2x4x7xi32>
+  func.return %0 : tensor<2x4x7xi32>
+}
+
+// CHECK-LABEL: @static_pad
+// CHECK-NEXT: Results
+// CHECK-NEXT{LITERAL}: [[[42, 42, 42, 42, 42, 42, 42],
+// CHECK-SAME{LITERAL}:   [42, 42, 1, 42, 2, 42, 3]
+// CHECK-SAME{LITERAL}:   [42, 42, 4, 42, 5, 42, 6],
+// CHECK-SAME{LITERAL}:   [42, 42, 42, 42, 42, 42, 42]],
+// CHECK-SAME{LITERAL}:  [[42, 42, 42, 42, 42, 42, 42],
+// CHECK-SAME{LITERAL}:   [42, 42, 42, 42, 42, 42, 42],
+// CHECK-SAME{LITERAL}:   [42, 42, 42, 42, 42, 42, 42],
+// CHECK-SAME{LITERAL}:   [42, 42, 42, 42, 42, 42, 42]]]
+
+func.func @dynamic_pad() -> tensor<?x4xi32> {
+  %c1 = arith.constant 1 : index
+  %empty = tensor.empty(%c1) : tensor<?x2xi32>
+  %pad_value = mhlo.constant dense<42> : tensor<i32>
+  %0 = "mhlo.pad"(%empty, %pad_value) {
+    edge_padding_low = dense<[0, 1]> : tensor<2xi64>,
+    edge_padding_high = dense<[1, 1]> : tensor<2xi64>,
+    interior_padding = dense<[0, 0]> : tensor<2xi64>
+  } : (tensor<?x2xi32>, tensor<i32>) -> tensor<?x4xi32>
+  func.return %0 : tensor<?x4xi32>
+}
+
+// CHECK-LABEL: @dynamic_pad
+// CHECK-NEXT: Results
+// CHECK-NEXT{LITERAL}: [[42, 0, 0, 42], [42, 42, 42, 42]]
+
+func.func @negative_pad() -> tensor<10xi32> {
+  %empty = arith.constant dense<[1,2,3,4,5,6,7]> : tensor<7xi32>
+  %pad_value = mhlo.constant dense<42> : tensor<i32>
+  %0 = "mhlo.pad"(%empty, %pad_value) {
+    edge_padding_low = dense<[-2]> : tensor<1xi64>,
+    edge_padding_high = dense<[-1]> : tensor<1xi64>,
+    interior_padding = dense<[1]> : tensor<1xi64>
+  } : (tensor<7xi32>, tensor<i32>) -> tensor<10xi32>
+  func.return %0 : tensor<10xi32>
+}
+
+// CHECK-LABEL: @negative_pad
+// CHECK-NEXT: Results
+// CHECK-NEXT{LITERAL}: [2, 42, 3, 42, 4, 42, 5, 42, 6, 42]
diff --git a/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/mhlo/reduce.mlir b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/mhlo/reduce.mlir
new file mode 100644
index 00000000000000..f371769413c2c5
--- /dev/null
+++ b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/mhlo/reduce.mlir
@@ -0,0 +1,17 @@
+// RUN: mlir-interpreter-runner %s -run-all | FileCheck %s
+
+func.func @reduce() -> tensor<3xi32> {
+  %cst = mhlo.constant dense<[[1, 2, 3], [4, 5, 6]]> : tensor<2x3xi32>
+  %init = mhlo.constant dense<1> : tensor<i32>
+  %reduce = mhlo.reduce(%cst init: %init) across dimensions = [0]
+      : (tensor<2x3xi32>, tensor<i32>) -> tensor<3xi32>
+    reducer(%arg0: tensor<i32>, %arg1: tensor<i32>)  {
+      %0 = mhlo.add %arg0, %arg1 : tensor<i32>
+      mhlo.return %0 : tensor<i32>
+    }
+  return %reduce : tensor<3xi32>
+}
+
+// CHECK-LABEL: @reduce
+// CHECK-NEXT: Results
+// CHECK-NEXT: [6, 8, 10]
diff --git a/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/mhlo/reshape.mlir b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/mhlo/reshape.mlir
new file mode 100644
index 00000000000000..117489375b00e7
--- /dev/null
+++ b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/mhlo/reshape.mlir
@@ -0,0 +1,34 @@
+// RUN: mlir-interpreter-runner %s -run-all | FileCheck %s
+
+func.func @reshape() -> tensor<2x4x2xi32> {
+  %cst = mhlo.constant dense<
+    [[0, 1, 2, 3], [4, 5, 6, 7],
+     [8, 9, 10, 11], [12, 13, 14, 15]]> : tensor<4x4xi32>
+  %reshape = mhlo.reshape %cst : (tensor<4x4xi32>) -> tensor<2x4x2xi32>
+  func.return %reshape : tensor<2x4x2xi32>
+}
+
+// CHECK-LABEL: @reshape
+// CHECK-NEXT: Results
+// CHECK-NEXT{LITERAL}: <2x4x2xi32>: [[[0, 1], [2, 3], [4, 5], [6, 7]],
+// CHECK-SAME{LITERAL}:               [[8, 9], [10, 11], [12, 13], [14, 15]]]
+
+func.func @reshape_0d_1d() -> tensor<1xi32> {
+  %cst = mhlo.constant dense<42> : tensor<i32>
+  %reshape = mhlo.reshape %cst : (tensor<i32>) -> tensor<1xi32>
+  func.return %reshape : tensor<1xi32>
+}
+
+// CHECK-LABEL: @reshape_0d_1d
+// CHECK-NEXT: Results
+// CHECK-NEXT: [42]
+
+func.func @reshape_1d_0d() -> tensor<i32> {
+  %cst = mhlo.constant dense<42> : tensor<1xi32>
+  %reshape = mhlo.reshape %cst : (tensor<1xi32>) -> tensor<i32>
+  func.return %reshape : tensor<i32>
+}
+
+// CHECK-LABEL: @reshape_1d_0d
+// CHECK-NEXT: Results
+// CHECK-NEXT: 42
diff --git a/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/mhlo/scatter.mlir b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/mhlo/scatter.mlir
new file mode 100644
index 00000000000000..85bb6225f9ba54
--- /dev/null
+++ b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/mhlo/scatter.mlir
@@ -0,0 +1,55 @@
+// RUN: mlir-interpreter-runner %s -run-all | FileCheck %s
+
+func.func @bounds_check() -> tensor<10xi32> {
+  // operand = np.zeros([8], dtype=np.int32)
+  // indices = np.array([[1], [8], [-1]])
+  // updates = np.array([[4, 5, 6], [6, 7, 8], [8, 9, 10]])
+  // lax.scatter_add(operand, indices, updates,
+  //   dimension_numbers=lax.ScatterDimensionNumbers(
+  //      update_window_dims=(0,), inserted_window_dims=(),
+  //      scatter_dims_to_operand_dims=(0,)))
+  %operand = mhlo.constant dense<0> : tensor<10xi32>
+  %indices = mhlo.constant dense<[[1], [8], [-1]]> : tensor<3x1xi32>
+  %updates = mhlo.constant dense<[[4, 5, 6], [6, 7, 8], [8, 9, 10]]> : tensor<3x3xi32>
+  %scatter = "mhlo.scatter"(%operand, %indices, %updates) ({
+  ^bb0(%lhs: tensor<i32>, %rhs: tensor<i32>):
+    %add = mhlo.add %lhs, %rhs : tensor<i32>
+    "mhlo.return"(%add) : (tensor<i32>) -> ()
+  }) {
+    scatter_dimension_numbers = #mhlo.scatter<
+      update_window_dims = [0],
+      inserted_window_dims = [],
+      index_vector_dim = 1,
+      scatter_dims_to_operand_dims = [0]
+    >
+  } : (tensor<10xi32>, tensor<3x1xi32>, tensor<3x3xi32>) -> tensor<10xi32>
+  return %scatter : tensor<10xi32>
+}
+
+// CHECK-LABEL: @bounds_check
+// CHECK-NEXT: Results
+// CHECK-NEXT{LITERAL}: [0, 4, 6, 8, 0, 0, 0, 0, 0, 0]
+
+func.func @update_last_element() -> tensor<2xi32> {
+  %operand = mhlo.constant dense<[1, 1]> : tensor<2xi32>
+  %indices = mhlo.constant dense<[[1]]> : tensor<1x1xi32>
+  %updates = mhlo.constant dense<[[0]]> : tensor<1x1xi32>
+
+  %scatter = "mhlo.scatter"(%operand, %indices, %updates) ({
+    ^bb0(%arg3: tensor<i32>, %arg4: tensor<i32>):
+      "mhlo.return"(%arg4) : (tensor<i32>) -> ()
+    }) {
+      indices_are_sorted = false,
+      scatter_dimension_numbers = #mhlo.scatter<
+        update_window_dims = [1],
+        scatter_dims_to_operand_dims = [0],
+        index_vector_dim = 1
+      >,
+      unique_indices = false
+    } : (tensor<2xi32>, tensor<1x1xi32>, tensor<1x1xi32>) -> tensor<2xi32>
+  return %scatter : tensor<2xi32>
+}
+
+// CHECK-LABEL: @update_last_element
+// CHECK-NEXT: Results
+// CHECK-NEXT{LITERAL}: [1, 0]
diff --git a/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/mhlo/select.mlir b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/mhlo/select.mlir
new file mode 100644
index 00000000000000..5959b3cf105b89
--- /dev/null
+++ b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/mhlo/select.mlir
@@ -0,0 +1,14 @@
+// RUN: mlir-interpreter-runner %s -run-all | FileCheck %s
+
+func.func @reshape() -> tensor<2xi32> {
+  %cst = mhlo.constant dense<[true, false]> : tensor<2xi1>
+  %a = mhlo.constant dense<[1, 2]> : tensor<2xi32>
+  %b = mhlo.constant dense<[3, 4]> : tensor<2xi32>
+  %ret = "mhlo.select"(%cst, %a, %b) :
+    (tensor<2xi1>, tensor<2xi32>, tensor<2xi32>) -> tensor<2xi32>
+  return %ret : tensor<2xi32>
+}
+
+// CHECK-LABEL: @reshape
+// CHECK-NEXT: Results
+// CHECK-NEXT: [1, 4]
diff --git a/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/mhlo/slice.mlir b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/mhlo/slice.mlir
new file mode 100644
index 00000000000000..c5bb53f8740cc0
--- /dev/null
+++ b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/mhlo/slice.mlir
@@ -0,0 +1,16 @@
+// RUN: mlir-interpreter-runner %s -run-all | FileCheck %s
+
+func.func @slice() -> tensor<1x2xi32> {
+  %cst = mhlo.constant dense<[[1, 2, 3, 4], [5, 6, 7, 8], [9, 10, 11, 12]]>
+    : tensor<3x4xi32>
+  %0 = "mhlo.slice"(%cst) {
+    start_indices = dense<[1, 0]> : tensor<2xi64>,
+    limit_indices = dense<[2, 4]> : tensor<2xi64>,
+    strides = dense<[1, 2]> : tensor<2xi64>
+  } : (tensor<3x4xi32>) -> tensor<1x2xi32>
+  func.return %0 : tensor<1x2xi32>
+}
+
+// CHECK-LABEL: @slice
+// CHECK-NEXT: Results
+// CHECK-NEXT{LITERAL}: [[5, 7]]
diff --git a/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/mhlo/sort.mlir b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/mhlo/sort.mlir
new file mode 100644
index 00000000000000..fa163718e91fbf
--- /dev/null
+++ b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/mhlo/sort.mlir
@@ -0,0 +1,25 @@
+// RUN: mlir-interpreter-runner %s -run-all | FileCheck %s
+
+func.func @sort() -> (tensor<2x5xf32>, tensor<2x5xi32>) {
+  %input0 = arith.constant dense<[
+    [4.0, 2.0, 1.0, 5.0, 3.0],
+    [6.0, 9.0, 8.0, 7.0, 10.0]
+  ]> : tensor<2x5xf32>
+
+  %input1 = arith.constant dense<[
+    [1, 2, 3, 4,  5],
+    [6, 7, 8, 9, 10]
+  ]> : tensor<2x5xi32>
+  %0, %1 = "mhlo.sort"(%input0, %input1) ({
+  ^bb0(%arg0: tensor<f32>, %arg1: tensor<f32>, %arg2: tensor<i32>, %arg3: tensor<i32>):
+    %7 = "mhlo.compare"(%arg0, %arg1) {comparison_direction = #mhlo<comparison_direction LT>} : (tensor<f32>, tensor<f32>) -> tensor<i1>
+    "mhlo.return"(%7) : (tensor<i1>) -> ()
+  }) {dimension = 1 : i64, is_stable = true} : (tensor<2x5xf32>, tensor<2x5xi32>) -> (tensor<2x5xf32>, tensor<2x5xi32>)
+
+  return %0, %1 : tensor<2x5xf32>, tensor<2x5xi32>
+}
+
+// CHECK-LABEL: @sort
+// CHECK-NEXT: Results
+// CHECK-NEXT{LITERAL}: [[1.000000e+00, 2.000000e+00, 3.000000e+00, 4.000000e+00, 5.000000e+00], [6.000000e+00, 7.000000e+00, 8.000000e+00, 9.000000e+00, 1.000000e+01]]
+// CHECK-NEXT{LITERAL}: [[3, 2, 5, 1, 4], [6, 9, 8, 7, 10]]
diff --git a/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/mhlo/subtract.mlir b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/mhlo/subtract.mlir
new file mode 100644
index 00000000000000..334e8421863ef2
--- /dev/null
+++ b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/mhlo/subtract.mlir
@@ -0,0 +1,10 @@
+// RUN: mlir-interpreter-runner %s | FileCheck %s
+
+func.func @main() -> tensor<2x3xi64> {
+  %lhs = mhlo.constant dense<[[0, 1, 2], [3, 4, 5]]> : tensor<2x3xi64>
+  %rhs = mhlo.constant dense<[[10, 20, 30], [40, 50, 60]]> : tensor<2x3xi64>
+  %result = mhlo.subtract %lhs, %rhs : tensor<2x3xi64>
+  return %result : tensor<2x3xi64>
+}
+
+// CHECK{LITERAL}: [[-10, -19, -28], [-37, -46, -55]]
\ No newline at end of file
diff --git a/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/mhlo/transpose.mlir b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/mhlo/transpose.mlir
new file mode 100644
index 00000000000000..09f9822aa8daa8
--- /dev/null
+++ b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/mhlo/transpose.mlir
@@ -0,0 +1,28 @@
+// RUN: mlir-interpreter-runner %s -run-all | FileCheck %s
+
+func.func @transpose() -> tensor<2x1x4x3xi32> {
+  %0 = mhlo.constant dense<[[[
+      [000, 001, 002, 003],
+      [010, 011, 012, 013],
+      [020, 021, 022, 023]
+    ],
+    [
+      [100, 101, 102, 103],
+      [110, 111, 112, 113],
+      [120, 121, 122, 123]
+    ]]]> : tensor<1x2x3x4xi32>
+  %1 = "mhlo.transpose"(%0) {permutation = dense<[1, 0, 3, 2]> : tensor<4xi64>}
+    : (tensor<1x2x3x4xi32>) -> tensor<2x1x4x3xi32>
+  return %1 : tensor<2x1x4x3xi32>
+}
+
+// CHECK-LABEL: @transpose
+// CHECK-NEXT: Results
+// CHECK-NEXT{LITERAL}: [[[[0, 10, 20],
+// CHECK{LITERAL}:         [1, 11, 21],
+// CHECK{LITERAL}:         [2, 12, 22],
+// CHECK{LITERAL}:         [3, 13, 23]]],
+// CHECK{LITERAL}:       [[[100, 110, 120],
+// CHECK{LITERAL}:         [101, 111, 121],
+// CHECK{LITERAL}:         [102, 112, 122],
+// CHECK{LITERAL}:         [103, 113, 123]]]]
diff --git a/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/mhlo/tuple.mlir b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/mhlo/tuple.mlir
new file mode 100644
index 00000000000000..02cb16c72e8b5d
--- /dev/null
+++ b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/mhlo/tuple.mlir
@@ -0,0 +1,30 @@
+// RUN: mlir-interpreter-runner %s -run-all | FileCheck %s
+
+func.func @tuple() -> tuple<tensor<2xi1>, tensor<i32>> {
+  %cst = mhlo.constant dense<[true, false]> : tensor<2xi1>
+  %c1 = mhlo.constant dense<1> : tensor<i32>
+  %ret = "mhlo.tuple"(%cst, %c1)
+    : (tensor<2xi1>, tensor<i32>) -> tuple<tensor<2xi1>, tensor<i32>>
+  return %ret : tuple<tensor<2xi1>, tensor<i32>>
+}
+
+// CHECK-LABEL: @tuple
+// CHECK-NEXT: Results
+// CHECK-NEXT: (TensorOrMemref<2xi1>: [true, false], TensorOrMemref<i32>: 1)
+
+func.func @get_tuple_element() -> (tensor<2xi1>, tensor<i32>) {
+  %cst = mhlo.constant dense<[true, false]> : tensor<2xi1>
+  %c42 = mhlo.constant dense<42> : tensor<i32>
+  %tuple = "mhlo.tuple"(%cst, %c42)
+    : (tensor<2xi1>, tensor<i32>) -> tuple<tensor<2xi1>, tensor<i32>>
+  %r0 = "mhlo.get_tuple_element"(%tuple) {index = 0 : i32}
+    : (tuple<tensor<2xi1>, tensor<i32>>) -> tensor<2xi1>
+  %r1 = "mhlo.get_tuple_element"(%tuple) {index = 1 : i32}
+    : (tuple<tensor<2xi1>, tensor<i32>>) -> tensor<i32>
+  return %r0, %r1 : tensor<2xi1>, tensor<i32>
+}
+
+// CHECK-LABEL: @get_tuple_element
+// CHECK-NEXT: Results
+// CHECK-NEXT: [true, false]
+// CHECK-NEXT: 42
diff --git a/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/mhlo/while.mlir b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/mhlo/while.mlir
new file mode 100644
index 00000000000000..4525174d7af4be
--- /dev/null
+++ b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/mhlo/while.mlir
@@ -0,0 +1,25 @@
+// RUN: mlir-interpreter-runner %s -run-all | FileCheck %s
+
+func.func @while() -> (tensor<i32>, tensor<i32>) {
+  %c0 = mhlo.constant dense<0> : tensor<i32>
+  %c1 = mhlo.constant dense<1> : tensor<i32>
+  %c10 = mhlo.constant dense<10> : tensor<i32>
+  %3:2 = "mhlo.while"(%c0, %c1) ({
+    ^bb0(%arg0: tensor<i32>, %arg1: tensor<i32>):
+      %4 = "mhlo.compare"(%arg0, %c10) {
+        comparison_direction = #mhlo<comparison_direction LT>
+      } : (tensor<i32>, tensor<i32>) -> tensor<i1>
+      "mhlo.return"(%4) : (tensor<i1>) -> ()
+    },  {
+    ^bb0(%arg0: tensor<i32>, %arg1: tensor<i32>):
+      %5 = mhlo.add %arg0, %c1 : tensor<i32>
+      %6 = mhlo.add %arg1, %arg1 : tensor<i32>
+      "mhlo.return"(%5, %6) : (tensor<i32>, tensor<i32>) -> ()
+    }) : (tensor<i32>, tensor<i32>) -> (tensor<i32>, tensor<i32>)
+  func.return %3#0, %3#1 : tensor<i32>, tensor<i32>
+}
+
+// CHECK-LABEL: @while
+// CHECK-NEXT: Results
+// CHECK-NEXT: TensorOrMemref<i32>: 10
+// CHECK-NEXT: TensorOrMemref<i32>: 1024
diff --git a/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/scf/for.mlir b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/scf/for.mlir
new file mode 100644
index 00000000000000..bad5951f4b4ca9
--- /dev/null
+++ b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/scf/for.mlir
@@ -0,0 +1,82 @@
+
+// RUN: mlir-interpreter-runner %s -run-all | FileCheck %s
+
+func.func @for() -> memref<4xi64> {
+  %c0 = arith.constant 0 : index
+  %c2 = arith.constant 2 : index
+  %c4 = arith.constant 4 : index
+  %alloc = memref.alloc() : memref<4xi64>
+  scf.for %i = %c0 to %c4 step %c2 {
+    %1 = arith.index_cast %i: index to i64
+    memref.store %1, %alloc[%i]: memref<4xi64>
+  }
+  return %alloc : memref<4xi64>
+}
+
+// CHECK-LABEL: @for
+// CHECK: Results
+// CHECK-NEXT{LITERAL}: [0, 0, 2, 0]
+
+func.func @nested() -> memref<2x2xindex> {
+  %c0 = arith.constant 0 : index
+  %c1 = arith.constant 1 : index
+  %c2 = arith.constant 2 : index
+  %alloc = memref.alloc() : memref<2x2xindex>
+  scf.for %i = %c0 to %c2 step %c1 {
+    scf.for %j = %c0 to %c2 step %c1 {
+      memref.store %c1, %alloc[%i, %j]: memref<2x2xindex>
+    }
+  }
+  return %alloc : memref<2x2xindex>
+}
+
+// CHECK-LABEL: @nested
+// CHECK-NEXT: Results
+// CHECK-NEXT{LITERAL}: [[1, 1], [1, 1]]
+
+func.func @iter_arg() -> index {
+  %c0 = arith.constant 0 : index
+  %c4 = arith.constant 4 : index
+  %c1 = arith.constant 1 : index
+  %sum = scf.for %i = %c0 to %c4 step %c1 iter_args(%x = %c1) -> index {
+    %sum = arith.addi %i, %x : index
+    scf.yield %sum : index
+  }
+  return %sum : index
+}
+
+// CHECK-LABEL: @iter_arg
+// CHECK-NEXT: Results
+// CHECK-NEXT{LITERAL}: i64: 7
+
+func.func @int32() -> memref<4xi32> {
+  %c0 = arith.constant 0 : i32
+  %c1 = arith.constant 1 : i32
+  %c4 = arith.constant 4 : i32
+  %alloc = memref.alloc() : memref<4xi32>
+  scf.for %i = %c0 to %c4 step %c1 : i32 {
+    %index = arith.index_cast %i : i32 to index
+    memref.store %i, %alloc[%index]: memref<4xi32>
+  }
+  return %alloc : memref<4xi32>
+}
+
+// CHECK-LABEL: int32
+// CHECK-NEXT: Results
+// CHECK-NEXT{LITERAL}: <4xi32>: [0, 1, 2, 3]
+
+func.func @int16() -> memref<4xi16> {
+  %c0 = arith.constant 0 : i16
+  %c1 = arith.constant 1 : i16
+  %c4 = arith.constant 4 : i16
+  %alloc = memref.alloc() : memref<4xi16>
+  scf.for %i = %c0 to %c4 step %c1 : i16 {
+    %index = arith.index_cast %i : i16 to index
+    memref.store %i, %alloc[%index]: memref<4xi16>
+  }
+  return %alloc : memref<4xi16>
+}
+
+// CHECK-LABEL: int16
+// CHECK-NEXT: Results
+// CHECK-NEXT{LITERAL}: <4xi16>: [0, 1, 2, 3]
diff --git a/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/scf/forall.mlir b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/scf/forall.mlir
new file mode 100644
index 00000000000000..29bec5e463ea92
--- /dev/null
+++ b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/scf/forall.mlir
@@ -0,0 +1,62 @@
+
+// RUN: mlir-interpreter-runner %s -run-all | FileCheck %s
+
+func.func @add() -> (tensor<3xi32>, tensor<3xi32>) {
+  %c3 = arith.constant 3 : index
+  %in = arith.constant dense<[1, 2, 3]> : tensor<3xi32>
+  %out = arith.constant dense<[4, 5, 6]> : tensor<3xi32>
+
+  %result = scf.forall (%i) in (%c3) shared_outs(%o = %out) -> tensor<3xi32> {
+    %addend = tensor.extract_slice %in[%i][1][1] : tensor<3xi32> to tensor<1xi32>
+    %augend = tensor.extract_slice %out[%i][1][1] : tensor<3xi32> to tensor<1xi32>
+    %sum = mhlo.add %augend, %addend : tensor<1xi32>
+    scf.forall.in_parallel {
+      tensor.parallel_insert_slice %sum into %o[%i][1][1]
+        : tensor<1xi32> into tensor<3xi32>
+    }
+  }
+
+  return %out, %result : tensor<3xi32>, tensor<3xi32>
+}
+
+// CHECK-LABEL: @add
+// CHECK-NEXT: Results
+// CHECK-NEXT: [4, 5, 6]
+// CHECK-NEXT: [5, 7, 9]
+
+func.func @bufferized_add() -> memref<3xi32> {
+  %c3 = arith.constant 3 : index
+  %in = arith.constant dense<[1, 2, 3]> : tensor<3xi32>
+  %out = arith.constant dense<[4, 5, 6]> : memref<3xi32>
+
+  scf.forall (%i) in (%c3) {
+    %addend = tensor.extract %in[%i] : tensor<3xi32>
+    %augend = memref.load %out[%i] : memref<3xi32>
+    %sum = arith.addi %augend, %addend : i32
+    memref.store %sum, %out[%i] : memref<3xi32>
+  }
+
+  return %out : memref<3xi32>
+}
+
+// CHECK-LABEL: @bufferized_add
+// CHECK-NEXT: Results
+// CHECK-NEXT: [5, 7, 9]
+
+func.func @step() -> memref<4xi32> {
+  %c2 = arith.constant 2 : index
+  %c4 = arith.constant 4 : index
+
+  %c42 = arith.constant 42 : i32
+  %out = arith.constant dense<[1, 1, 1, 1]> : memref<4xi32>
+
+  scf.forall (%i) = (0) to (%c4) step (%c2) {
+    memref.store %c42, %out[%i] : memref<4xi32>
+  }
+
+  return %out : memref<4xi32>
+}
+
+// CHECK-LABEL: @step
+// CHECK-NEXT: Results
+// CHECK-NEXT: [42, 1, 42, 1]
diff --git a/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/scf/if.mlir b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/scf/if.mlir
new file mode 100644
index 00000000000000..52d899d13328a6
--- /dev/null
+++ b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/scf/if.mlir
@@ -0,0 +1,69 @@
+// RUN: mlir-interpreter-runner %s -run-all | FileCheck %s
+
+func.func @true() -> i64 {
+  %c0 = arith.constant 0 : i64
+  %c1 = arith.constant 1 : i64
+  %true = arith.constant true
+  %ret = scf.if %true -> i64 {
+    scf.yield %c0 : i64
+  } else {
+    scf.yield %c1 : i64
+  }
+  return %ret : i64
+}
+
+// CHECK-LABEL: @true
+// CHECK-NEXT: Results
+// CHECK-NEXT: i64: 0
+
+func.func @false() -> i64 {
+  %c2 = arith.constant 2 : i64
+  %c3 = arith.constant 3 : i64
+  %false = arith.constant false
+  %ret = scf.if %false -> i64 {
+    scf.yield %c2 : i64
+  } else {
+    scf.yield %c3 : i64
+  }
+  return %ret : i64
+}
+
+// CHECK-LABEL: @false
+// CHECK-NEXT: Results
+// CHECK-NEXT: i64: 3
+
+func.func @side_effect() -> memref<i64> {
+  %alloc = memref.alloc() : memref<i64>
+  %true = arith.constant true
+  %c124 = arith.constant 124 : i64
+  %c125 = arith.constant 125 : i64
+  scf.if %true {
+    memref.store %c124, %alloc[] : memref<i64>
+    scf.yield
+  } else {
+    memref.store %c125, %alloc[] : memref<i64>
+    scf.yield
+  }
+  return %alloc : memref<i64>
+}
+
+// CHECK-LABEL: @side_effect
+// CHECK-NEXT: Results
+// CHECK-NEXT: <i64>: 124
+
+func.func @side_effect_not_executed() -> memref<i64> {
+  %alloc = memref.alloc() : memref<i64>
+  %false = arith.constant false
+  %c126 = arith.constant 126 : i64
+  memref.store %c126, %alloc[] : memref<i64>
+  %c127 = arith.constant 127 : i64
+  scf.if %false {
+    memref.store %c127, %alloc[] : memref<i64>
+    scf.yield
+  }
+  return %alloc : memref<i64>
+}
+
+// CHECK-LABEL: @side_effect_not_executed
+// CHECK-NEXT: Results
+// CHECK-NEXT: <i64>: 126
diff --git a/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/scf/parallel.mlir b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/scf/parallel.mlir
new file mode 100644
index 00000000000000..72007d2b2966fd
--- /dev/null
+++ b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/scf/parallel.mlir
@@ -0,0 +1,42 @@
+// RUN: mlir-interpreter-runner %s -run-all | FileCheck %s
+
+func.func @parallel() -> memref<4x4xi32> {
+  %ret = memref.alloc() : memref<4x4xi32>
+  %c0 = arith.constant 0 : index
+  %c1 = arith.constant 1 : index
+  %c2 = arith.constant 2 : index
+  %c4 = arith.constant 4 : index
+  %c9 = arith.constant 9 : i32
+  scf.parallel (%i, %j) = (%c0, %c1) to (%c4, %c4) step (%c1, %c2) {
+    memref.store %c9, %ret[%i, %j] : memref<4x4xi32>
+  }
+  return %ret : memref<4x4xi32>
+}
+
+// CHECK-LABEL: @parallel
+// CHECK-NEXT: Results
+// CHECK-NEXT{LITERAL}: [[0, 9, 0, 9], [0, 9, 0, 9], [0, 9, 0, 9], [0, 9, 0, 9]]
+
+func.func @reduce_2() -> (index, index) {
+  %c1 = arith.constant 1 : index
+  %c6 = arith.constant 6 : index
+  %ret:2 = scf.parallel (%i) = (%c1) to (%c6) step (%c1)
+             init (%c1, %c1) -> (index, index) {
+    scf.reduce (%i, %i : index, index) {
+      ^bb0(%lhs: index, %rhs: index):
+        %ret = arith.muli %lhs, %rhs : index
+        scf.reduce.return %ret : index
+    }, {
+      ^bb0(%lhs: index, %rhs: index):
+        %ret = arith.addi %lhs, %rhs : index
+        scf.reduce.return %ret : index
+    }
+  }
+  return %ret#0, %ret#1 : index, index
+}
+
+// CHECK-LABEL: @reduce_2
+// CHECK-NEXT: Results
+// CHECK-NEXT{LITERAL}: i64: 120
+// CHECK-NEXT{LITERAL}: i64: 16
+
diff --git a/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/scf/while.mlir b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/scf/while.mlir
new file mode 100644
index 00000000000000..1be3814f40365f
--- /dev/null
+++ b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/scf/while.mlir
@@ -0,0 +1,45 @@
+// RUN: mlir-interpreter-runner %s -run-all | FileCheck %s
+
+func.func @while_empty() -> memref<i64> {
+  %c0 = arith.constant 0 : i64
+  %c1 = arith.constant 1 : i64
+  %c4 = arith.constant 4 : i64
+  %alloc = memref.alloc() : memref<i64>
+  memref.store %c0, %alloc[] : memref<i64>
+  scf.while: () -> () {
+    %value = memref.load %alloc[] : memref<i64>
+    %cond = arith.cmpi slt, %value, %c4 : i64
+    scf.condition(%cond)
+  } do {
+    %value = memref.load %alloc[] : memref<i64>
+    %add = arith.addi %value, %c1 : i64
+    memref.store %add, %alloc[] : memref<i64>
+    scf.yield
+  }
+  return %alloc : memref<i64>
+}
+
+// CHECK-LABEL: @while_empty
+// CHECK-NEXT: Results
+// CHECK-NEXT{LITERAL}: TensorOrMemref<i64>: 4
+
+func.func @while_var() -> i64 {
+  %c0 = arith.constant 0 : i64
+  %c1 = arith.constant 1 : i64
+  %c4 = arith.constant 4 : i64
+  %alloc = memref.alloc() : memref<i64>
+  memref.store %c0, %alloc[] : memref<i64>
+  %ret = scf.while(%arg0 = %c0): (i64) -> (i64) {
+    %cond = arith.cmpi slt, %arg0, %c4 : i64
+    scf.condition(%cond) %arg0 : i64
+  } do {
+  ^bb0(%arg1: i64):
+    %add = arith.addi %arg1, %c1 : i64
+    scf.yield %add : i64
+  }
+  return %ret : i64
+}
+
+// CHECK-LABEL: @while_var
+// CHECK-NEXT: Results
+// CHECK-NEXT{LITERAL}: i64: 4
diff --git a/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/tensor/collapse_shape.mlir b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/tensor/collapse_shape.mlir
new file mode 100644
index 00000000000000..1201a2a342e73a
--- /dev/null
+++ b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/tensor/collapse_shape.mlir
@@ -0,0 +1,42 @@
+// RUN: mlir-interpreter-runner %s -run-all | FileCheck %s
+
+func.func @collapse_shape()
+    -> (tensor<1x6xi32>, tensor<2x3xi32>, tensor<6xi32>) {
+  %cst = arith.constant dense<[[[1, 2, 3], [4, 5, 6]]]> : tensor<1x2x3xi32>
+  %collapse1 = tensor.collapse_shape %cst [[0], [1, 2]]
+      : tensor<1x2x3xi32> into tensor<1x6xi32>
+  %collapse2 = tensor.collapse_shape %cst [[0, 1], [2]]
+      : tensor<1x2x3xi32> into tensor<2x3xi32>
+  %collapse3 = tensor.collapse_shape %cst [[0, 1, 2]]
+      : tensor<1x2x3xi32> into tensor<6xi32>
+  return %collapse1, %collapse2, %collapse3
+      : tensor<1x6xi32>, tensor<2x3xi32>, tensor<6xi32>
+}
+
+// CHECK-LABEL: @collapse_shape
+// CHECK-NEXT: Results
+// CHECK-NEXT{LITERAL}: [[1, 2, 3, 4, 5, 6]]
+// CHECK-NEXT{LITERAL}: [[1, 2, 3], [4, 5, 6]]
+// CHECK-NEXT{LITERAL}: [1, 2, 3, 4, 5, 6]
+
+func.func @to_unit() -> tensor<i32> {
+  %cst = arith.constant dense<42> : tensor<1x1x1x1xi32>
+  %collapse = tensor.collapse_shape %cst []
+    : tensor<1x1x1x1xi32> into tensor<i32>
+  return %collapse : tensor<i32>
+}
+
+// CHECK-LABEL: @to_unit
+// CHECK-NEXT: Results
+// CHECK-NEXT: 42
+
+func.func @dynamic() -> tensor<?xi32> {
+  %cst = arith.constant dense<42> : tensor<2x3xi32>
+  %cast = tensor.cast %cst : tensor<2x3xi32> to tensor<?x3xi32>
+  %collapse = tensor.collapse_shape %cast [[0, 1]] : tensor<?x3xi32> into tensor<?xi32>
+  return %collapse : tensor<?xi32>
+}
+
+// CHECK-LABEL: @dynamic
+// CHECK-NEXT: Results
+// CHECK-NEXT: [42, 42, 42, 42, 42, 42]
diff --git a/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/tensor/dim.mlir b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/tensor/dim.mlir
new file mode 100644
index 00000000000000..6e38150587ed27
--- /dev/null
+++ b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/tensor/dim.mlir
@@ -0,0 +1,12 @@
+// RUN: mlir-interpreter-runner %s -run-all | FileCheck %s
+
+func.func @dim() -> index {
+  %it = tensor.empty() : tensor<2x3x4xi32>
+  %c1 = arith.constant 1 : index
+  %dim = tensor.dim %it, %c1 : tensor<2x3x4xi32>
+  return %dim : index
+}
+
+// CHECK-LABEL: @dim
+// CHECK-NEXT: Results
+// CHECK-NEXT: i64: 3
diff --git a/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/tensor/empty.mlir b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/tensor/empty.mlir
new file mode 100644
index 00000000000000..1fdd75afc7e227
--- /dev/null
+++ b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/tensor/empty.mlir
@@ -0,0 +1,21 @@
+// RUN: mlir-interpreter-runner %s -run-all | FileCheck %s
+
+func.func @static() -> tensor<2x3xi32> {
+  %ret = tensor.empty() : tensor<2x3xi32>
+  return %ret : tensor<2x3xi32>
+}
+
+// CHECK-LABEL: @static
+// CHECK-NEXT: Results
+// CHECK-NEXT{LITERAL}: [[0, 0, 0], [0, 0, 0]]
+
+func.func @dynamic() -> tensor<2x?x3x?xi32> {
+  %c5 = arith.constant 5 : index
+  %c7 = arith.constant 7 : index
+  %ret = tensor.empty(%c5, %c7) : tensor<2x?x3x?xi32>
+  return %ret : tensor<2x?x3x?xi32>
+}
+
+// CHECK-LABEL: @dynamic
+// CHECK-NEXT: Results
+// CHECK-NEXT{LITERAL}: TensorOrMemref<2x5x3x7xi32>
diff --git a/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/tensor/expand_shape.mlir b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/tensor/expand_shape.mlir
new file mode 100644
index 00000000000000..77948509ae3334
--- /dev/null
+++ b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/tensor/expand_shape.mlir
@@ -0,0 +1,31 @@
+// RUN: mlir-interpreter-runner %s -run-all | FileCheck %s
+
+func.func @expand_shape()
+    -> (tensor<1x6xi32>, tensor<2x3xi32>, tensor<1x2x3xi32>) {
+  %cst = arith.constant dense<[1, 2, 3, 4, 5, 6]> : tensor<6xi32>
+  %expand1 = tensor.expand_shape %cst [[0, 1]] output_shape [1, 6]
+      : tensor<6xi32> into tensor<1x6xi32>
+  %expand2 = tensor.expand_shape %cst [[0, 1]] output_shape [2, 3]
+      : tensor<6xi32> into tensor<2x3xi32>
+  %expand3 = tensor.expand_shape %cst [[0, 1, 2]] output_shape [1, 2, 3]
+      : tensor<6xi32> into tensor<1x2x3xi32>
+  return %expand1, %expand2, %expand3
+      : tensor<1x6xi32>, tensor<2x3xi32>, tensor<1x2x3xi32>
+}
+
+// CHECK-LABEL: @expand_shape
+// CHECK-NEXT: Results
+// CHECK-NEXT{LITERAL}: [[1, 2, 3, 4, 5, 6]]
+// CHECK-NEXT{LITERAL}: [[1, 2, 3], [4, 5, 6]]
+// CHECK-NEXT{LITERAL}: [[[1, 2, 3], [4, 5, 6]]]
+
+func.func @from_unit() -> tensor<1x1xi32> {
+  %cst = arith.constant dense<42> : tensor<i32>
+  %expand = tensor.expand_shape %cst [] output_shape [1, 1]
+    : tensor<i32> into tensor<1x1xi32>
+  return %expand : tensor<1x1xi32>
+}
+
+// CHECK-LABEL: @from_unit
+// CHECK-NEXT: Results
+// CHECK-NEXT{LITERAL}: TensorOrMemref<1x1xi32>: [[42]]
diff --git a/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/tensor/extract.mlir b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/tensor/extract.mlir
new file mode 100644
index 00000000000000..9209eea17d9a50
--- /dev/null
+++ b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/tensor/extract.mlir
@@ -0,0 +1,13 @@
+// RUN: mlir-interpreter-runner %s -run-all | FileCheck %s
+
+func.func @extract() -> i32 {
+  %c0 = arith.constant 0 : index
+  %c1 = arith.constant 1 : index
+  %cst = arith.constant dense<[[[1, 2, 3], [4, 5, 6]]]> : tensor<1x2x3xi32>
+  %ret = tensor.extract %cst[%c0, %c1, %c1] : tensor<1x2x3xi32>
+  return %ret : i32
+}
+
+// CHECK-LABEL: @extract
+// CHECK-NEXT: Results
+// CHECK-NEXT: i32: 5
diff --git a/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/tensor/extract_slice.mlir b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/tensor/extract_slice.mlir
new file mode 100644
index 00000000000000..81721e12b2abf7
--- /dev/null
+++ b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/tensor/extract_slice.mlir
@@ -0,0 +1,62 @@
+// RUN: mlir-interpreter-runner %s -run-all | FileCheck %s
+
+func.func @extract() -> tensor<1x2x1xi32> {
+  %cst = arith.constant dense<[[[1, 2, 3], [4, 5, 6]]]> : tensor<1x2x3xi32>
+  %ret = tensor.extract_slice %cst[0, 0, 0][1, 2, 1][1, 1, 1] : tensor<1x2x3xi32> to tensor<1x2x1xi32>
+  return %ret : tensor<1x2x1xi32>
+}
+
+// CHECK-LABEL: @extract
+// CHECK-NEXT: Results
+// CHECK-NEXT{LITERAL}: [[[1], [4]]]
+
+func.func @rank_reduction() -> tensor<2xi32> {
+  %cst = arith.constant dense<[[[1, 2, 3], [4, 5, 6]]]> : tensor<1x2x3xi32>
+  %ret = tensor.extract_slice %cst[0, 0, 0][1, 2, 1][1, 1, 1] : tensor<1x2x3xi32> to tensor<2xi32>
+  return %ret : tensor<2xi32>
+}
+
+// CHECK-LABEL: @rank_reduction
+// CHECK-NEXT: Results
+// CHECK-NEXT{LITERAL}: [1, 4]
+
+func.func @rank_reduction_to_dynamic() -> tensor<?xi32> {
+  %cst = arith.constant dense<[[[1, 2, 3], [4, 5, 6]]]> : tensor<1x2x3xi32>
+  %c2 = arith.constant 2 : index
+  %ret = tensor.extract_slice %cst[0, 0, 0][1, %c2, 1][1, 1, 1] : tensor<1x2x3xi32> to tensor<?xi32>
+  return %ret : tensor<?xi32>
+}
+
+// CHECK-LABEL: @rank_reduction_to_dynamic
+// CHECK-NEXT: Results
+// CHECK-NEXT{LITERAL}: [1, 4]
+
+func.func @no_rank_reduction_to_dynamic() -> tensor<1x?xi32> {
+  %cst = arith.constant dense<[[0], [1], [2], [3], [4]]> : tensor<5x1xi32>
+  %c0 = arith.constant 0 : index
+  %c1 = arith.constant 1 : index
+  %c2 = arith.constant 2 : index
+  %ret = tensor.extract_slice %cst[%c2, %c0] [1, %c1] [1, 1]
+    : tensor<5x1xi32> to tensor<1x?xi32>
+  return %ret : tensor<1x?xi32>
+}
+
+// CHECK-LABEL: @no_rank_reduction_to_dynamic
+// CHECK-NEXT: Results
+// CHECK-NEXT{LITERAL}: [[2]]
+
+func.func @extract_from_extract() -> tensor<1x1xi32> {
+  %cst = arith.constant dense<[[0], [1], [2], [3], [4]]> : tensor<5x1xi32>
+  %c0 = arith.constant 0 : index
+  %c1 = arith.constant 1 : index
+  %c2 = arith.constant 2 : index
+  %extracted = tensor.extract_slice %cst[%c2, %c0] [1, %c1] [1, 1]
+    : tensor<5x1xi32> to tensor<1x?xi32>
+  %ret = tensor.extract_slice %extracted[%c0, %c0] [1, 1] [1, 1]
+    : tensor<1x?xi32> to tensor<1x1xi32>
+  return %ret : tensor<1x1xi32>
+}
+
+// CHECK-LABEL: @extract_from_extract
+// CHECK-NEXT: Results
+// CHECK-NEXT{LITERAL}: [[2]]
diff --git a/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/tensor/from_elements.mlir b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/tensor/from_elements.mlir
new file mode 100644
index 00000000000000..e8bda72ada6dd4
--- /dev/null
+++ b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/tensor/from_elements.mlir
@@ -0,0 +1,25 @@
+// RUN: mlir-interpreter-runner %s -run-all | FileCheck %s
+
+func.func @from_elements() -> tensor<2x3xindex> {
+  %c0 = arith.constant 0 : index
+  %c1 = arith.constant 1 : index
+  %c2 = arith.constant 2 : index
+  %c3 = arith.constant 3 : index
+  %c4 = arith.constant 4 : index
+  %c5 = arith.constant 5 : index
+  %ret = tensor.from_elements %c0, %c1, %c2, %c3, %c4, %c5 :  tensor<2x3xindex>
+  return %ret : tensor<2x3xindex>
+}
+
+// CHECK-LABEL: @from_elements
+// CHECK-NEXT: Results
+// CHECK-NEXT{LITERAL}: [[0, 1, 2], [3, 4, 5]]
+
+func.func @empty() -> tensor<0xindex> {
+  %ret = tensor.from_elements : tensor<0xindex>
+  return %ret : tensor<0xindex>
+}
+
+// CHECK-LABEL: @empty
+// CHECK-NEXT: Results
+// CHECK-NEXT{LITERAL}: []
diff --git a/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/tensor/generate.mlir b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/tensor/generate.mlir
new file mode 100644
index 00000000000000..9ae2a4d43de1fa
--- /dev/null
+++ b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/tensor/generate.mlir
@@ -0,0 +1,29 @@
+// RUN: mlir-interpreter-runner %s -run-all | FileCheck %s
+
+func.func @generate() -> tensor<?xindex> {
+  %size = arith.constant 5 : index
+  %iota = tensor.generate %size {
+    ^bb0(%i : index):
+      tensor.yield %i : index
+    } : tensor<?xindex>
+  return %iota : tensor<?xindex>
+}
+
+// CHECK-LABEL: @generate
+// CHECK-NEXT: Results
+// CHECK-NEXT{LITERAL}: [0, 1, 2, 3, 4]
+
+func.func @generate_2d() -> tensor<?x?xindex> {
+  %c5 = arith.constant 5 : index
+  %c2 = arith.constant 2 : index
+  %iota = tensor.generate %c5, %c2 {
+    ^bb0(%i : index, %j : index):
+      %sum = arith.addi %i, %j : index
+      tensor.yield %sum : index
+    } : tensor<?x?xindex>
+  return %iota : tensor<?x?xindex>
+}
+
+// CHECK-LABEL: @generate_2d
+// CHECK-NEXT: Results
+// CHECK-NEXT{LITERAL}: [[0, 1], [1, 2], [2, 3], [3, 4], [4, 5]]
diff --git a/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/tensor/insert.mlir b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/tensor/insert.mlir
new file mode 100644
index 00000000000000..47d5677694d89c
--- /dev/null
+++ b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/tensor/insert.mlir
@@ -0,0 +1,14 @@
+// RUN: mlir-interpreter-runner %s -run-all | FileCheck %s
+
+func.func @insert() -> tensor<1x2x3xi32> {
+  %c0 = arith.constant 0 : index
+  %c1 = arith.constant 1 : index
+  %c7 = arith.constant 7 : i32
+  %cst = arith.constant dense<[[[1, 2, 3], [4, 5, 6]]]> : tensor<1x2x3xi32>
+  %ret = tensor.insert %c7 into %cst[%c0, %c1, %c1] : tensor<1x2x3xi32>
+  return %ret : tensor<1x2x3xi32>
+}
+
+// CHECK-LABEL: @insert
+// CHECK-NEXT: Results
+// CHECK-NEXT{LITERAL}: [[[1, 2, 3], [4, 7, 6]]]
diff --git a/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/tensor/insert_slice.mlir b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/tensor/insert_slice.mlir
new file mode 100644
index 00000000000000..ec5383d5c2852a
--- /dev/null
+++ b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/tensor/insert_slice.mlir
@@ -0,0 +1,25 @@
+// RUN: mlir-interpreter-runner %s -run-all | FileCheck %s
+
+func.func @insert() -> tensor<1x3x3xi32> {
+  %cst = arith.constant dense<[[[1, 2, 3], [4, 5, 6], [7, 8, 9]]]> : tensor<1x3x3xi32>
+  %cst_2 = arith.constant dense<[[[10], [11]]]> : tensor<1x2x1xi32>
+  %ret = tensor.insert_slice %cst_2 into %cst[0, 1, 1][1, 2, 1][1, 1, 1]
+    : tensor<1x2x1xi32> into tensor<1x3x3xi32>
+  return %ret : tensor<1x3x3xi32>
+}
+
+// CHECK-LABEL: @insert
+// CHECK-NEXT: Results
+// CHECK-NEXT{LITERAL}: [[[1, 2, 3], [4, 10, 6], [7, 11, 9]]]
+
+func.func @rank_increase() -> tensor<1x3x3xi32> {
+  %cst = arith.constant dense<[[[1, 2, 3], [4, 5, 6], [7, 8, 9]]]> : tensor<1x3x3xi32>
+  %cst_2 = arith.constant dense<[10, 11]> : tensor<2xi32>
+  %ret = tensor.insert_slice %cst_2 into %cst[0, 1, 1][1, 2, 1][1, 1, 1]
+    : tensor<2xi32> into tensor<1x3x3xi32>
+  return %ret : tensor<1x3x3xi32>
+}
+
+// CHECK-LABEL: @rank_increase
+// CHECK-NEXT: Results
+// CHECK-NEXT{LITERAL}: [[[1, 2, 3], [4, 10, 6], [7, 11, 9]]]
diff --git a/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/tensor/pad.mlir b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/tensor/pad.mlir
new file mode 100644
index 00000000000000..a2a6470b6e3f17
--- /dev/null
+++ b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/tensor/pad.mlir
@@ -0,0 +1,38 @@
+// RUN: mlir-interpreter-runner %s -run-all | FileCheck %s
+
+func.func @pad() -> tensor<1x?x?xi32> {
+  %it = arith.constant dense<[[[1, 2, 3], [2, 3, 4]]]> : tensor<1x2x3xi32>
+  %offset = arith.constant 2 : index
+  %c0 = arith.constant 0 : index
+  %out = tensor.pad %it low[%c0, %offset, 0] high[0, %c0, %offset]  {
+    ^bb0(%a: index, %b: index, %c: index):
+      %c5 = arith.constant 5 : i32
+      tensor.yield %c5 : i32
+    } : tensor<1x2x3xi32> to tensor<1x?x?xi32>
+  return %out : tensor<1x?x?xi32>
+}
+
+// CHECK-LABEL: @pad
+// CHECK-NEXT: Results
+// CHECK-NEXT{LITERAL}: [[[5, 5, 5, 5, 5],
+// CHECK-SAME{LITERAL}:   [5, 5, 5, 5, 5],
+// CHECK-SAME{LITERAL}:   [1, 2, 3, 5, 5],
+// CHECK-SAME{LITERAL}:   [2, 3, 4, 5, 5]]]
+
+func.func @pad_args() -> tensor<3x3xindex> {
+  %it = arith.constant dense<[[999]]> : tensor<1x1xindex>
+  %out = tensor.pad %it low[1, 1] high[1, 1]  {
+    ^bb0(%a: index, %b: index):
+      %c10 = arith.constant 10 : index
+      %mul = arith.muli %a, %c10 : index
+      %ret = arith.addi %mul, %b : index
+      tensor.yield %ret : index
+    } : tensor<1x1xindex> to tensor<3x3xindex>
+  return %out : tensor<3x3xindex>
+}
+
+// CHECK-LABEL: @pad_args
+// CHECK-NEXT: Results
+// CHECK-NEXT{LITERAL}: [[0, 1, 2],
+// CHECK-SAME{LITERAL}:  [10, 999, 12],
+// CHECK-SAME{LITERAL}:  [20, 21, 22]]
\ No newline at end of file
diff --git a/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/vector/bitcast.mlir b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/vector/bitcast.mlir
new file mode 100644
index 00000000000000..242b025324e934
--- /dev/null
+++ b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/vector/bitcast.mlir
@@ -0,0 +1,32 @@
+// RUN: mlir-interpreter-runner %s -run-all | FileCheck %s
+
+func.func @bitcast1d() -> vector<4xi8> {
+  %c = arith.constant dense<-2> : vector<1xi32>
+  %b = vector.bitcast %c : vector<1xi32> to vector<4xi8>
+  return %b : vector<4xi8>
+}
+
+// CHECK-LABEL: @bitcast1d
+// CHECK-NEXT: Results
+// CHECK-NEXT: vector<4xi8>: [-2, -1, -1, -1]
+
+func.func @bitcast2d() -> vector<2x1xi64> {
+  %c = arith.constant dense<[[0, 1], [2, 3]]> : vector<2x2xi32>
+  %b = vector.bitcast %c : vector<2x2xi32> to vector<2x1xi64>
+  return %b : vector<2x1xi64>
+}
+
+// CHECK-LABEL: @bitcast2d
+// CHECK-NEXT: Results
+// CHECK-NEXT{LITERAL}: vector<2x1xi64>: [[4294967296], [12884901890]]
+
+func.func @transpose_bitcast() -> vector<2x1xi64> {
+  %c = arith.constant dense<[[0, 1], [2, 3]]> : vector<2x2xi32>
+  %d = vector.transpose %c, [1, 0] : vector<2x2xi32> to vector<2x2xi32>
+  %b = vector.bitcast %d : vector<2x2xi32> to vector<2x1xi64>
+  return %b : vector<2x1xi64>
+}
+
+// CHECK-LABEL: @transpose_bitcast
+// CHECK-NEXT: Results
+// CHECK-NEXT{LITERAL}: vector<2x1xi64>: [[8589934592], [12884901889]]
diff --git a/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/vector/broadcast.mlir b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/vector/broadcast.mlir
new file mode 100644
index 00000000000000..1e30f7aa9e603f
--- /dev/null
+++ b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/vector/broadcast.mlir
@@ -0,0 +1,51 @@
+// RUN: mlir-interpreter-runner %s -run-all | FileCheck %s
+
+func.func @scalar_to_unit() -> vector<i32> {
+  %c1 = arith.constant 1 : i32
+  %b = vector.broadcast %c1 : i32 to vector<i32>
+  return %b : vector<i32>
+}
+
+// CHECK-LABEL: @scalar_to_unit
+// CHECK-NEXT: Results
+// CHECK-NEXT: vector<i32>: 1
+
+func.func @scalar_to_2d() -> vector<2x3xi32> {
+  %c1 = arith.constant 1 : i32
+  %b = vector.broadcast %c1 : i32 to vector<2x3xi32>
+  return %b : vector<2x3xi32>
+}
+
+// CHECK-LABEL: @scalar_to_2d
+// CHECK-NEXT: Results
+// CHECK-NEXT{LITERAL}: vector<2x3xi32>: [[1, 1, 1], [1, 1, 1]]
+
+func.func @unit_to_unit() -> vector<i32> {
+  %c1 = arith.constant dense<1> : vector<i32>
+  %b = vector.broadcast %c1 : vector<i32> to vector<i32>
+  return %b : vector<i32>
+}
+
+// CHECK-LABEL: @unit_to_unit
+// CHECK-NEXT: Results
+// CHECK-NEXT: vector<i32>: 1
+
+func.func @stretch() -> vector<3xi32> {
+  %c1 = arith.constant dense<1> : vector<1xi32>
+  %b = vector.broadcast %c1 : vector<1xi32> to vector<3xi32>
+  return %b : vector<3xi32>
+}
+
+// CHECK-LABEL: @stretch
+// CHECK-NEXT: Results
+// CHECK-NEXT: vector<3xi32>: [1, 1, 1]
+
+func.func @stretch_and_broadcast() -> vector<2x1x2x2xi32> {
+  %c1 = arith.constant dense<[[1], [2]]> : vector<2x1xi32>
+  %b = vector.broadcast %c1 : vector<2x1xi32> to vector<2x1x2x2xi32>
+  return %b : vector<2x1x2x2xi32>
+}
+
+// CHECK-LABEL: @stretch_and_broadcast
+// CHECK-NEXT: Results
+// CHECK-NEXT{LITERAL}: vector<2x1x2x2xi32>: [[[[1, 1], [2, 2]]], [[[1, 1], [2, 2]]]]
diff --git a/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/vector/compressstore.mlir b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/vector/compressstore.mlir
new file mode 100644
index 00000000000000..079633094723c0
--- /dev/null
+++ b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/vector/compressstore.mlir
@@ -0,0 +1,16 @@
+// RUN: mlir-interpreter-runner %s -run-all | FileCheck %s
+
+func.func @compressstore() -> memref<3x4xi32> {
+  %alloc = memref.alloc() : memref<3x4xi32>
+  %c = arith.constant dense<[1,2,3]> : vector<3xi32>
+  %m = arith.constant dense<[true,false,true]> : vector<3xi1>
+  %c1 = arith.constant 1 : index
+  %c2 = arith.constant 2 : index
+  vector.compressstore %alloc[%c1, %c2], %m, %c
+    : memref<3x4xi32>, vector<3xi1>, vector<3xi32>
+  return %alloc : memref<3x4xi32>
+}
+
+// CHECK-LABEL: @compressstore
+// CHECK-NEXT: Results
+// CHECK-NEXT{LITERAL}: [[0, 0, 0, 0], [0, 0, 1, 3], [0, 0, 0, 0]]
diff --git a/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/vector/constant_mask.mlir b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/vector/constant_mask.mlir
new file mode 100644
index 00000000000000..a3e3e0d03cc959
--- /dev/null
+++ b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/vector/constant_mask.mlir
@@ -0,0 +1,14 @@
+// RUN: mlir-interpreter-runner %s -run-all | FileCheck %s
+
+func.func @constant_mask() -> vector<4x3xi1> {
+  %1 = vector.constant_mask [3, 2] : vector<4x3xi1>
+  return %1 : vector<4x3xi1>
+}
+
+// CHECK-LABEL: @constant_mask
+// CHECK-NEXT: Results
+// CHECK-NEXT{LITERAL}: vector<4x3xi1>:
+// CHECK-SAME{LITERAL}: [[true, true, false],
+// CHECK-SAME{LITERAL}:  [true, true, false],
+// CHECK-SAME{LITERAL}:  [true, true, false],
+// CHECK-SAME{LITERAL}:  [false, false, false]]
diff --git a/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/vector/contract.mlir b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/vector/contract.mlir
new file mode 100644
index 00000000000000..f7f26fed00d8d0
--- /dev/null
+++ b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/vector/contract.mlir
@@ -0,0 +1,141 @@
+// RUN: mlir-interpreter-runner %s -run-all | FileCheck %s
+
+// Adapted from
+// mlir/test/Integration/Dialect/Vector/CPU/test-contraction.mlir
+
+#dotp_accesses = [
+  affine_map<(i) -> (i)>,
+  affine_map<(i) -> (i)>,
+  affine_map<(i) -> ()>
+]
+#dotp_trait = {
+  indexing_maps = #dotp_accesses,
+  iterator_types = ["reduction"]
+}
+
+#matvec_accesses = [
+  affine_map<(i, j) -> (i, j)>,
+  affine_map<(i, j) -> (j)>,
+  affine_map<(i, j) -> (i)>
+]
+#matvec_trait = {
+  indexing_maps = #matvec_accesses,
+  iterator_types = ["parallel", "reduction"]
+}
+
+#mattransvec_accesses = [
+  affine_map<(i, j) -> (j, i)>,
+  affine_map<(i, j) -> (j)>,
+  affine_map<(i, j) -> (i)>
+]
+#mattransvec_trait = {
+  indexing_maps = #mattransvec_accesses,
+  iterator_types = ["parallel", "reduction"]
+}
+
+#matmat_accesses = [
+  affine_map<(i, j, k) -> (i, k)>,
+  affine_map<(i, j, k) -> (k, j)>,
+  affine_map<(i, j, k) -> (i, j)>
+]
+#matmat_trait = {
+  indexing_maps = #matmat_accesses,
+  iterator_types = ["parallel", "parallel", "reduction"]
+}
+
+#column_major_matmat_accesses = [
+  affine_map<(i, j, k) -> (k, j)>,
+  affine_map<(i, j, k) -> (i, k)>,
+  affine_map<(i, j, k) -> (j, i)>
+]
+#column_major_matmat_trait = {
+  indexing_maps = #column_major_matmat_accesses,
+  iterator_types = ["parallel", "parallel", "reduction"]
+}
+
+func.func @dot_products() -> (i32, i32) {
+  %c0 = arith.constant 0 : i32
+  %c1 = arith.constant 1 : i32
+  %a = arith.constant dense<[1, 2]> : vector<2xi32>
+  %b = arith.constant dense<[3, 4]> : vector<2xi32>
+  // Contraction: dot-product a x b
+  %dp1 = vector.contract #dotp_trait %a, %b, %c0
+    : vector<2xi32>, vector<2xi32> into i32
+  %dp2 = vector.contract #dotp_trait %a, %b, %c1
+    : vector<2xi32>, vector<2xi32> into i32
+  return %dp1, %dp2 : i32, i32
+}
+
+// CHECK-LABEL: @dot_product
+// CHECK-NEXT: Results
+// CHECK-NEXT: i32: 11
+// CHECK-NEXT: i32: 12
+
+func.func @matrix_vector() -> (vector<2xi32>, vector<2xi32>) {
+  %z1 = arith.constant dense<0> : vector<2xi32>
+  %a = arith.constant dense<[1, 2]> : vector<2xi32>
+  %c = arith.constant dense<[5, 6]> : vector<2xi32>
+  %A = arith.constant dense<[[1, 2], [3, 4]]> : vector<2x2xi32>
+  %mv1 = vector.contract #matvec_trait %A, %c, %z1
+    : vector<2x2xi32>, vector<2xi32> into vector<2xi32>
+  %mv2 = vector.contract #matvec_trait %A, %c, %a
+    : vector<2x2xi32>, vector<2xi32> into vector<2xi32>
+  return %mv1, %mv2 : vector<2xi32>, vector<2xi32>
+}
+
+// CHECK-LABEL: @matrix_vector
+// CHECK-NEXT: Results
+// CHECK-NEXT: [17, 39]
+// CHECK-NEXT: [18, 41]
+
+func.func @matrix_trans_vector() -> (vector<2xi32>, vector<2xi32>) {
+  %z1 = arith.constant dense<0> : vector<2xi32>
+  %a = arith.constant dense<[1, 2]> : vector<2xi32>
+  %c = arith.constant dense<[5, 6]> : vector<2xi32>
+  %A = arith.constant dense<[[1, 2], [3, 4]]> : vector<2x2xi32>
+  %mv1 = vector.contract #mattransvec_trait %A, %c, %z1
+    : vector<2x2xi32>, vector<2xi32> into vector<2xi32>
+  %mv2 = vector.contract #mattransvec_trait %A, %c, %a
+    : vector<2x2xi32>, vector<2xi32> into vector<2xi32>
+  return %mv1, %mv2 : vector<2xi32>, vector<2xi32>
+}
+
+// CHECK-LABEL: @matrix_trans_vector
+// CHECK-NEXT: Results
+// CHECK-NEXT: [23, 34]
+// CHECK-NEXT: [24, 36]
+
+func.func @matrix_matrix() -> (vector<2x2xi32>, vector<2x2xi32>) {
+  %z2 = arith.constant dense<0> : vector<2x2xi32>
+  %A = arith.constant dense<[[1, 2], [3, 4]]> : vector<2x2xi32>
+  %B = arith.constant dense<[[5, 6], [7, 8]]> : vector<2x2xi32>
+  %mm1 = vector.contract #matmat_trait %A, %B, %z2
+    : vector<2x2xi32>, vector<2x2xi32> into vector<2x2xi32>
+  %mm2 = vector.contract #matmat_trait %A, %B, %A
+    : vector<2x2xi32>, vector<2x2xi32> into vector<2x2xi32>
+  return %mm1, %mm2 : vector<2x2xi32>, vector<2x2xi32>
+}
+
+// CHECK-LABEL: @matrix_matrix
+// CHECK-NEXT: Results
+// CHECK-NEXT{LITERAL}: [[19, 22], [43, 50]]
+// CHECK-NEXT{LITERAL}: [[20, 24], [46, 54]]
+
+func.func @matrix_matrix_column_major() -> (vector<2x2xi32>, vector<2x2xi32>) {
+  %z2 = arith.constant dense<0> : vector<2x2xi32>
+  %A = arith.constant dense<[[1, 2], [3, 4]]> : vector<2x2xi32>
+  %B = arith.constant dense<[[5, 6], [7, 8]]> : vector<2x2xi32>
+  %llvm_matrix_column_major_mm0 =
+    vector.contract #column_major_matmat_trait %A, %B, %z2
+      : vector<2x2xi32>, vector<2x2xi32> into vector<2x2xi32>
+  %llvm_matrix_column_major_mm1 =
+    vector.contract #column_major_matmat_trait %A, %B, %A
+      : vector<2x2xi32>, vector<2x2xi32> into vector<2x2xi32>
+  return %llvm_matrix_column_major_mm0, %llvm_matrix_column_major_mm1 :
+    vector<2x2xi32>, vector<2x2xi32>
+}
+
+// CHECK-LABEL: @matrix_matrix_column_major
+// CHECK-NEXT: Results
+// CHECK-NEXT{LITERAL}: [[23, 31], [34, 46]]
+// CHECK-NEXT{LITERAL}: [[24, 33], [37, 50]]
diff --git a/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/vector/create_mask.mlir b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/vector/create_mask.mlir
new file mode 100644
index 00000000000000..2616eddb273fb2
--- /dev/null
+++ b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/vector/create_mask.mlir
@@ -0,0 +1,16 @@
+// RUN: mlir-interpreter-runner %s -run-all | FileCheck %s
+
+func.func @create_mask() -> vector<4x3xi1> {
+  %c2 = arith.constant 2 : index
+  %c3 = arith.constant 3 : index
+  %1 = vector.create_mask %c3, %c2 : vector<4x3xi1>
+  return %1 : vector<4x3xi1>
+}
+
+// CHECK-LABEL: @create_mask
+// CHECK-NEXT: Results
+// CHECK-NEXT{LITERAL}: vector<4x3xi1>:
+// CHECK-SAME{LITERAL}: [[true, true, false],
+// CHECK-SAME{LITERAL}:  [true, true, false],
+// CHECK-SAME{LITERAL}:  [true, true, false],
+// CHECK-SAME{LITERAL}:  [false, false, false]]
diff --git a/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/vector/expandload.mlir b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/vector/expandload.mlir
new file mode 100644
index 00000000000000..35a4bc27358128
--- /dev/null
+++ b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/vector/expandload.mlir
@@ -0,0 +1,19 @@
+// RUN: mlir-interpreter-runner %s -run-all | FileCheck %s
+
+func.func @expandload() -> (vector<4xi32>, vector<4xi32>) {
+  %passthrough = arith.constant dense<[1,2,3,4]> : vector<4xi32>
+  %mask = arith.constant dense<[true, false, true, false]> : vector<4xi1>
+  %memref = arith.constant dense<[[10,11,12,13,14],
+                                  [15,16,17,18,19]]> : memref<2x5xi32>
+  %c1 = arith.constant 1 : index
+  %c2 = arith.constant 2 : index
+  %ret = vector.expandload %memref[%c1, %c2], %mask, %passthrough
+    : memref<2x5xi32>, vector<4xi1>, vector<4xi32> into vector<4xi32>
+
+  return %passthrough, %ret : vector<4xi32>, vector<4xi32>
+}
+
+// CHECK-LABEL: @expandload
+// CHECK-NEXT: Results
+// CHECK-NEXT: vector<4xi32>: [1, 2, 3, 4]
+// CHECK-NEXT: vector<4xi32>: [17, 2, 18, 4]
diff --git a/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/vector/extract.mlir b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/vector/extract.mlir
new file mode 100644
index 00000000000000..6186cf62257302
--- /dev/null
+++ b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/vector/extract.mlir
@@ -0,0 +1,52 @@
+// RUN: mlir-interpreter-runner %s -run-all | FileCheck %s
+
+func.func @extract_1d_1d() -> vector<2xi32> {
+  %c = arith.constant dense<[1, 2]> : vector<2xi32>
+  %i = vector.extract %c[] : vector<2xi32> from vector<2xi32>
+  return %i : vector<2xi32>
+}
+
+// CHECK-LABEL: @extract_1d_1d
+// CHECK-NEXT: Results
+// CHECK-NEXT: vector<2xi32>: [1, 2]
+
+func.func @extract_1d_0d() -> i32 {
+  %c = arith.constant dense<[1, 2]> : vector<2xi32>
+  %i = vector.extract %c[1] : i32 from vector<2xi32>
+  return %i : i32
+}
+
+// CHECK-LABEL: @extract_1d_0d
+// CHECK-NEXT: Results
+// CHECK-NEXT: i32: 2
+
+func.func @extract_2d_0d() -> i32 {
+  %c = arith.constant dense<[[1, 2], [3, 4]]> : vector<2x2xi32>
+  %i = vector.extract %c[0, 1] : i32 from vector<2x2xi32>
+  return %i : i32
+}
+
+// CHECK-LABEL: @extract_2d_0d
+// CHECK-NEXT: Results
+// CHECK-NEXT: i32: 2
+
+func.func @extract_2d_1d() -> vector<3xi32> {
+  %c = arith.constant dense<[[1, 2, 3], [4, 5, 6]]> : vector<2x3xi32>
+  %i = vector.extract %c[0] : vector<3xi32> from vector<2x3xi32>
+  return %i : vector<3xi32>
+}
+
+// CHECK-LABEL: @extract_2d_1d
+// CHECK-NEXT: Results
+// CHECK-NEXT: vector<3xi32>: [1, 2, 3]
+
+func.func @extract_2d_2d() -> vector<2x2xi32> {
+  %c = arith.constant dense<[[1, 2], [3, 4]]> : vector<2x2xi32>
+  %i = vector.extract %c[] : vector<2x2xi32> from vector<2x2xi32>
+  return %i : vector<2x2xi32>
+}
+
+// CHECK-LABEL: @extract_2d_2d
+// CHECK-NEXT: Results
+// CHECK-NEXT{LITERAL}: vector<2x2xi32>: [[1, 2], [3, 4]]
+
diff --git a/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/vector/extract_strided_slice.mlir b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/vector/extract_strided_slice.mlir
new file mode 100644
index 00000000000000..eae5ae9e5e483d
--- /dev/null
+++ b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/vector/extract_strided_slice.mlir
@@ -0,0 +1,18 @@
+// RUN: mlir-interpreter-runner %s -run-all | FileCheck %s
+
+func.func @extract_strided_slice() -> vector<2x3xi32> {
+  %c = arith.constant dense<[[1,2,3,4],
+                             [5,6,7,8],
+                             [9,10,11,12]]> : vector<3x4xi32>
+  %o = vector.extract_strided_slice %c {
+    offsets = [0, 1],
+    sizes = [2, 3],
+    // TODO(jreiffers): Test non-unit strides when supported by verifier.
+    strides = [1, 1]
+  } : vector<3x4xi32> to vector<2x3xi32>
+  return %o : vector<2x3xi32>
+}
+
+// CHECK-LABEL: @extract_strided_slice
+// CHECK-NEXT: Results
+// CHECK-NEXT{LITERAL}: [[2, 3, 4], [6, 7, 8]]
diff --git a/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/vector/extractelement.mlir b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/vector/extractelement.mlir
new file mode 100644
index 00000000000000..2e4ffc708d64ea
--- /dev/null
+++ b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/vector/extractelement.mlir
@@ -0,0 +1,22 @@
+// RUN: mlir-interpreter-runner %s -run-all | FileCheck %s
+
+func.func @extract0d() -> i32 {
+  %c = arith.constant dense<1> : vector<i32>
+  %i = vector.extractelement %c[] : vector<i32>
+  return %i : i32
+}
+
+// CHECK-LABEL: @extract0d
+// CHECK-NEXT: Results
+// CHECK-NEXT: i32: 1
+
+func.func @extract1d() -> i32 {
+  %c = arith.constant dense<[1,2]> : vector<2xi32>
+  %c1 = arith.constant 1 : index
+  %i = vector.extractelement %c[%c1 : index] : vector<2xi32>
+  return %i : i32
+}
+
+// CHECK-LABEL: @extract1d
+// CHECK-NEXT: Results
+// CHECK-NEXT: i32: 2
\ No newline at end of file
diff --git a/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/vector/flat_transpose.mlir b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/vector/flat_transpose.mlir
new file mode 100644
index 00000000000000..b67a84a4cce49b
--- /dev/null
+++ b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/vector/flat_transpose.mlir
@@ -0,0 +1,23 @@
+// RUN: mlir-interpreter-runner %s -run-all | FileCheck %s
+
+func.func @flattranspose_2x3() -> vector<6xi32> {
+  %c = arith.constant dense<[0, 1, 2, 3, 4, 5]> : vector<6xi32>
+  %ret = vector.flat_transpose %c { columns = 2: i32, rows = 3: i32 }
+    : vector<6xi32> -> vector<6xi32>
+  return %ret : vector<6xi32>
+}
+
+// CHECK-LABEL: @flattranspose_2x3
+// CHECK-NEXT: Results
+// CHECK-NEXT: [0, 3, 1, 4, 2, 5]
+
+func.func @flattranspose_3x2() -> vector<6xi32> {
+  %c = arith.constant dense<[0, 1, 2, 3, 4, 5]> : vector<6xi32>
+  %ret = vector.flat_transpose %c { columns = 3: i32, rows = 2: i32 }
+    : vector<6xi32> -> vector<6xi32>
+  return %ret : vector<6xi32>
+}
+
+// CHECK-LABEL: @flattranspose_3x2
+// CHECK-NEXT: Results
+// CHECK-NEXT: [0, 2, 4, 1, 3, 5]
diff --git a/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/vector/fma.mlir b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/vector/fma.mlir
new file mode 100644
index 00000000000000..696799d5c99c6c
--- /dev/null
+++ b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/vector/fma.mlir
@@ -0,0 +1,13 @@
+// RUN: mlir-interpreter-runner %s -run-all | FileCheck %s
+
+func.func @fma() -> vector<2xf32> {
+  %a = arith.constant dense<[1.0,2.0]> : vector<2xf32>
+  %b = arith.constant dense<[3.0,4.0]> : vector<2xf32>
+  %c = arith.constant dense<[5.0,6.0]> : vector<2xf32>
+  %r = vector.fma %a, %b, %c : vector<2xf32>
+  return %r : vector<2xf32>
+}
+
+// CHECK-LABEL: @fma
+// CHECK-NEXT: Results
+// CHECK-NEXT: vector<2xf32>: [8.000000e+00, 1.400000e+01]
diff --git a/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/vector/gather.mlir b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/vector/gather.mlir
new file mode 100644
index 00000000000000..c92a5589b1ecb4
--- /dev/null
+++ b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/vector/gather.mlir
@@ -0,0 +1,50 @@
+// RUN: mlir-interpreter-runner %s -run-all | FileCheck %s
+
+// Adapted from mlir/test/Integration/Dialect/Vector/CPU/test-gather.mlir
+
+func.func private @gather8(%base: memref<10xi32>, %indices: vector<8xi32>,
+              %mask: vector<8xi1>, %pass_thru: vector<8xi32>) -> vector<8xi32> {
+  %c0 = arith.constant 0: index
+  %g = vector.gather %base[%c0][%indices], %mask, %pass_thru
+    : memref<10xi32>, vector<8xi32>, vector<8xi1>, vector<8xi32> into vector<8xi32>
+  return %g : vector<8xi32>
+}
+
+func.func @gather() ->
+    (vector<8xi32>, vector<8xi32>, vector<8xi32>, vector<8xi32>, vector<8xi32>) {
+  %A = arith.constant dense<[0, 1, 2, 3, 4, 5, 6, 7, 8, 9]> : memref<10xi32>
+  %idx = arith.constant dense<[0, 6, 1, 3, 5, 4, 9, 2]> : vector<8xi32>
+  %pass = arith.constant dense<-7> : vector<8xi32>
+  %none = vector.constant_mask [0] : vector<8xi1>
+  %all = vector.constant_mask [8] : vector<8xi1>
+  %some = vector.constant_mask [4] : vector<8xi1>
+  %true = arith.constant true
+  %more = vector.insert %true, %some[7] : i1 into vector<8xi1>
+
+  %g1 = call @gather8(%A, %idx, %all, %pass)
+    : (memref<10xi32>, vector<8xi32>, vector<8xi1>, vector<8xi32>)
+    -> (vector<8xi32>)
+  %g2 = call @gather8(%A, %idx, %none, %pass)
+    : (memref<10xi32>, vector<8xi32>, vector<8xi1>, vector<8xi32>)
+    -> (vector<8xi32>)
+  %g3 = call @gather8(%A, %idx, %some, %pass)
+    : (memref<10xi32>, vector<8xi32>, vector<8xi1>, vector<8xi32>)
+    -> (vector<8xi32>)
+  %g4 = call @gather8(%A, %idx, %more, %pass)
+    : (memref<10xi32>, vector<8xi32>, vector<8xi1>, vector<8xi32>)
+    -> (vector<8xi32>)
+  %g5 = call @gather8(%A, %idx, %all, %pass)
+    : (memref<10xi32>, vector<8xi32>, vector<8xi1>, vector<8xi32>)
+    -> (vector<8xi32>)
+
+  return %g1, %g2, %g3, %g4, %g5
+    : vector<8xi32>, vector<8xi32>, vector<8xi32>, vector<8xi32>, vector<8xi32>
+}
+
+// CHECK-LABEL: @gather
+// CHECK-NEXT: Results
+// CHECK-NEXT: [0, 6, 1, 3, 5, 4, 9, 2]
+// CHECK-NEXT: [-7, -7, -7, -7, -7, -7, -7, -7]
+// CHECK-NEXT: [0, 6, 1, 3, -7, -7, -7, -7]
+// CHECK-NEXT: [0, 6, 1, 3, -7, -7, -7, 2]
+// CHECK-NEXT: [0, 6, 1, 3, 5, 4, 9, 2]
diff --git a/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/vector/insert.mlir b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/vector/insert.mlir
new file mode 100644
index 00000000000000..97019cf98e4775
--- /dev/null
+++ b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/vector/insert.mlir
@@ -0,0 +1,57 @@
+// RUN: mlir-interpreter-runner %s -run-all | FileCheck %s
+
+func.func @insert_1d_1d() -> vector<2xi32> {
+  %c = arith.constant dense<[3, 4]> : vector<2xi32>
+  %d = arith.constant dense<[1, 2]> : vector<2xi32>
+  %i = vector.insert %c, %d[] : vector<2xi32> into vector<2xi32>
+  return %i : vector<2xi32>
+}
+
+// CHECK-LABEL: @insert_1d_1d
+// CHECK-NEXT: Results
+// CHECK-NEXT: vector<2xi32>: [3, 4]
+
+func.func @insert_1d_0d() -> vector<2xi32> {
+  %c = arith.constant 42 : i32
+  %d = arith.constant dense<[1, 2]> : vector<2xi32>
+  %i = vector.insert %c, %d[1] : i32 into vector<2xi32>
+  return %i : vector<2xi32>
+}
+
+// CHECK-LABEL: @insert_1d_0d
+// CHECK-NEXT: Results
+// CHECK-NEXT: vector<2xi32>: [1, 42]
+
+func.func @insert_2d_0d() -> vector<2x2xi32> {
+  %c = arith.constant 42 : i32
+  %d = arith.constant dense<[[1, 2], [3, 4]]> : vector<2x2xi32>
+  %i = vector.insert %c, %d[0, 1] : i32 into vector<2x2xi32>
+  return %i : vector<2x2xi32>
+}
+
+// CHECK-LABEL: @insert_2d_0d
+// CHECK-NEXT: Results
+// CHECK-NEXT{LITERAL}: vector<2x2xi32>: [[1, 42], [3, 4]]
+
+func.func @insert_2d_1d() -> vector<2x3xi32> {
+  %c = arith.constant dense<[42, 43, 44]> : vector<3xi32>
+  %d = arith.constant dense<[[1, 2, 3], [4, 5, 6]]> : vector<2x3xi32>
+  %i = vector.insert %c, %d[0] : vector<3xi32> into vector<2x3xi32>
+  return %i : vector<2x3xi32>
+}
+
+// CHECK-LABEL: @insert_2d_1d
+// CHECK-NEXT: Results
+// CHECK-NEXT{LITERAL}: vector<2x3xi32>: [[42, 43, 44], [4, 5, 6]]
+
+func.func @insert_2d_2d() -> vector<2x2xi32> {
+  %c = arith.constant dense<[[1, 2], [3, 4]]> : vector<2x2xi32>
+  %d = arith.constant dense<0> : vector<2x2xi32>
+  %i = vector.insert %c, %d[] : vector<2x2xi32> into vector<2x2xi32>
+  return %i : vector<2x2xi32>
+}
+
+// CHECK-LABEL: @insert_2d_2d
+// CHECK-NEXT: Results
+// CHECK-NEXT{LITERAL}: vector<2x2xi32>: [[1, 2], [3, 4]]
+
diff --git a/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/vector/insert_strided_slice.mlir b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/vector/insert_strided_slice.mlir
new file mode 100644
index 00000000000000..140cc716bfdf0a
--- /dev/null
+++ b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/vector/insert_strided_slice.mlir
@@ -0,0 +1,17 @@
+// RUN: mlir-interpreter-runner %s -run-all | FileCheck %s
+
+func.func @insert_strided_slice() -> (vector<3x4xi32>, vector<3x4xi32>) {
+  %v = arith.constant dense<[[2, 3, 4], [6, 7, 8]]> : vector<2x3xi32>
+  %c = arith.constant dense<0> : vector<3x4xi32>
+  %o = vector.insert_strided_slice %v, %c {
+    offsets = [0, 1],
+    // TODO(jreiffers): Test non-unit strides when supported by verifier.
+    strides = [1, 1]
+  } : vector<2x3xi32> into vector<3x4xi32>
+  return %c, %o : vector<3x4xi32>, vector<3x4xi32>
+}
+
+// CHECK-LABEL: @insert_strided_slice
+// CHECK-NEXT: Results
+// CHECK-NEXT{LITERAL}: [[0, 0, 0, 0], [0, 0, 0, 0], [0, 0, 0, 0]]
+// CHECK-NEXT{LITERAL}: [[0, 2, 3, 4], [0, 6, 7, 8], [0, 0, 0, 0]]
diff --git a/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/vector/insertelement.mlir b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/vector/insertelement.mlir
new file mode 100644
index 00000000000000..4befb8c9598486
--- /dev/null
+++ b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/vector/insertelement.mlir
@@ -0,0 +1,24 @@
+// RUN: mlir-interpreter-runner %s -run-all | FileCheck %s
+
+func.func @insert0d() -> vector<i32> {
+  %c = arith.constant dense<1> : vector<i32>
+  %v = arith.constant 42 : i32
+  %i = vector.insertelement %v, %c[] : vector<i32>
+  return %i : vector<i32>
+}
+
+// CHECK-LABEL: @insert0d
+// CHECK-NEXT: Results
+// CHECK-NEXT: vector<i32>: 42
+
+func.func @insert1d() -> vector<2xi32> {
+  %c = arith.constant dense<1> : vector<2xi32>
+  %c1 = arith.constant 1 : index
+  %v = arith.constant 42 : i32
+  %i = vector.insertelement %v, %c[%c1 : index] : vector<2xi32>
+  return %i : vector<2xi32>
+}
+
+// CHECK-LABEL: @insert1d
+// CHECK-NEXT: Results
+// CHECK-NEXT: vector<2xi32>: [1, 42]
diff --git a/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/vector/invalid.mlir b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/vector/invalid.mlir
new file mode 100644
index 00000000000000..7a08cf50923f51
--- /dev/null
+++ b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/vector/invalid.mlir
@@ -0,0 +1,39 @@
+// RUN: (! mlir-interpreter-runner %s -run-all 2>&1) | FileCheck %s
+
+func.func @write_4_at_3_inbounds() {
+  %a = memref.alloc() : memref<5xi32>
+  %base = arith.constant 3 : index
+  %f = arith.constant dense<[1, 2, 3, 4]> : vector<4xi32>
+  vector.transfer_write %f, %a[%base]
+    {permutation_map = affine_map<(d0) -> (d0)>, in_bounds = [true]}
+    : vector<4xi32>, memref<5xi32>
+  return
+}
+
+// CHECK-LABEL: @write_4_at_3_inbounds
+// CHECK-NEXT: index out of bounds
+
+func.func @transfer_read_2d_1d_oob()-> vector<2xi32> {
+  %a = arith.constant dense<[[0, 1, 2, 3], [4, 5, 6, 7]]> : memref<2x4xi32>
+  %c0 = arith.constant 0 : index
+  %c2 = arith.constant 2 : index
+  %c-42 = arith.constant -42: i32
+  %f = vector.transfer_read %a[%c2, %c0], %c-42
+      : memref<2x4xi32>, vector<2xi32>
+  return %f : vector<2xi32>
+}
+
+// CHECK-LABEL: @transfer_read_2d_1d_oob
+// CHECK-NEXT: index out of bounds
+
+func.func @store_vector_memref() -> memref<1x2xvector<2xi32>> {
+  %m = memref.alloc() : memref<1x2xvector<2xi32>>
+  %c0 = arith.constant 0 : index
+  %c1 = arith.constant 1094795585 : index
+  %v = arith.constant dense<[1,2]> : vector<2xi32>
+  vector.store %v, %m[%c0, %c1] : memref<1x2xvector<2xi32>>, vector<2xi32>
+  return %m : memref<1x2xvector<2xi32>>
+}
+
+// CHECK-LABEL: @store_vector_memref
+// CHECK-NEXT: index out of bounds
diff --git a/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/vector/load.mlir b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/vector/load.mlir
new file mode 100644
index 00000000000000..3fabd7676d6413
--- /dev/null
+++ b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/vector/load.mlir
@@ -0,0 +1,27 @@
+// RUN: mlir-interpreter-runner %s -run-all | FileCheck %s
+
+func.func @load_vector_memref() -> vector<2x2xi32> {
+  %m = memref.alloc() : memref<4x4xvector<2x2xi32>>
+  %c3 = arith.constant 3 : index
+  %r = vector.load %m[%c3, %c3] : memref<4x4xvector<2x2xi32>>, vector<2x2xi32>
+  return %r : vector<2x2xi32>
+}
+
+// CHECK-LABEL: @load_vector_memref
+// CHECK-NEXT: Results
+// CHECK-NEXT{LITERAL}: vector<2x2xi32>: [[0, 0], [0, 0]]
+
+func.func @load_scalar_memref() -> vector<2x2xi32> {
+  %m = arith.constant dense<[[1, 2, 3, 4],
+                             [5, 6, 7, 8],
+                             [9, 10, 11, 12],
+                             [13, 14, 15, 16]]> : memref<4x4xi32>
+  %c1 = arith.constant 1 : index
+  %c2 = arith.constant 2 : index
+  %r = vector.load %m[%c1, %c2] : memref<4x4xi32>, vector<2x2xi32>
+  return %r : vector<2x2xi32>
+}
+
+// CHECK-LABEL: @load_scalar_memref
+// CHECK-NEXT: Results
+// CHECK-NEXT{LITERAL}: vector<2x2xi32>: [[7, 8], [11, 12]]
diff --git a/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/vector/maskedload.mlir b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/vector/maskedload.mlir
new file mode 100644
index 00000000000000..bb11f49d261b28
--- /dev/null
+++ b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/vector/maskedload.mlir
@@ -0,0 +1,19 @@
+// RUN: mlir-interpreter-runner %s -run-all | FileCheck %s
+
+func.func @maskedload() -> (vector<4xi32>, vector<4xi32>) {
+  %passthrough = arith.constant dense<[1,2,3,4]> : vector<4xi32>
+  %mask = arith.constant dense<[true, false, true, false]> : vector<4xi1>
+  %memref = arith.constant dense<[[10,11,12,13,14],
+                                  [15,16,17,18,19]]> : memref<2x5xi32>
+  %c1 = arith.constant 1 : index
+  %c2 = arith.constant 2 : index
+  %ret = vector.maskedload %memref[%c1, %c2], %mask, %passthrough
+    : memref<2x5xi32>, vector<4xi1>, vector<4xi32> into vector<4xi32>
+
+  return %passthrough, %ret : vector<4xi32>, vector<4xi32>
+}
+
+// CHECK-LABEL: @maskedload
+// CHECK-NEXT: Results
+// CHECK-NEXT: vector<4xi32>: [1, 2, 3, 4]
+// CHECK-NEXT: vector<4xi32>: [17, 2, 19, 4]
diff --git a/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/vector/maskedstore.mlir b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/vector/maskedstore.mlir
new file mode 100644
index 00000000000000..9ceaefc3a9c622
--- /dev/null
+++ b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/vector/maskedstore.mlir
@@ -0,0 +1,18 @@
+// RUN: mlir-interpreter-runner %s -run-all | FileCheck %s
+
+func.func @maskedstore() -> memref<2x5xi32> {
+  %value = arith.constant dense<[1,2,3,4]> : vector<4xi32>
+  %mask = arith.constant dense<[true, false, true, false]> : vector<4xi1>
+  %memref = arith.constant dense<[[10,11,12,13,14],
+                                  [15,16,17,18,19]]> : memref<2x5xi32>
+  %c1 = arith.constant 1 : index
+  %c2 = arith.constant 2 : index
+  vector.maskedstore %memref[%c1, %c2], %mask, %value
+    : memref<2x5xi32>, vector<4xi1>, vector<4xi32>
+
+  return %memref : memref<2x5xi32>
+}
+
+// CHECK-LABEL: @maskedstore
+// CHECK-NEXT: Results
+// CHECK-NEXT{LITERAL}: [[10, 11, 12, 13, 14], [15, 16, 1, 18, 3]]
diff --git a/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/vector/multi_reduction.mlir b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/vector/multi_reduction.mlir
new file mode 100644
index 00000000000000..94fc5481653583
--- /dev/null
+++ b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/vector/multi_reduction.mlir
@@ -0,0 +1,46 @@
+// RUN: mlir-interpreter-runner %s -run-all | FileCheck %s
+
+func.func @multi_reduction_1d_0d() -> i32 {
+  %a = arith.constant dense<[10, 4, 7]> : vector<3xi32>
+  %acc = arith.constant 1 : i32
+  %r = vector.multi_reduction <add>, %a, %acc [0] : vector<3xi32> to i32
+  return %r : i32
+}
+
+// CHECK-LABEL: @multi_reduction_1d_0d
+// CHECK-NEXT: Results
+// CHECK-NEXT: 22
+
+#dense234 = dense<[[[0,1,2,3],[4,5,6,7],[8,9,10,11]],
+                   [[12,13,14,15],[16,17,18,19],[20,21,22,23]]]>
+              : vector<2x3x4xi32>
+
+func.func @multi_reduction_3d_2d()
+    -> (vector<2x3xi32>, vector<2x4xi32>, vector<2x3xi32>) {
+  %a = arith.constant #dense234
+  %acc_23 = arith.constant dense<[[0, 1, 2], [3, 4, 5]]> : vector<2x3xi32>
+  %r1 = vector.multi_reduction <add>, %a, %acc_23 [2]
+    : vector<2x3x4xi32> to vector<2x3xi32>
+  %acc_24 = arith.constant dense<[[0, 1, 2, 3], [4, 5, 6, 7]]> : vector<2x4xi32>
+  %r2 = vector.multi_reduction <mul>, %a, %acc_24 [1]
+    : vector<2x3x4xi32> to vector<2x4xi32>
+  return %r1, %r2, %acc_23 : vector<2x3xi32>, vector<2x4xi32>, vector<2x3xi32>
+}
+
+// CHECK-LABEL: @multi_reduction_3d_2d
+// CHECK-NEXT: Results
+// CHECK-NEXT{LITERAL}: [[6, 23, 40], [57, 74, 91]]
+// CHECK-NEXT{LITERAL}: [[0, 45, 240, 693], [15360, 23205, 33264, 45885]]
+// CHECK-NEXT{LITERAL}: [[0, 1, 2], [3, 4, 5]]
+
+func.func @multi_reduction_3d_1d() -> vector<3xi32> {
+  %a = arith.constant #dense234
+  %acc_3 = arith.constant dense<[0, 1, 2]> : vector<3xi32>
+  %r1 = vector.multi_reduction <add>, %a, %acc_3 [2, 0]
+    : vector<2x3x4xi32> to vector<3xi32>
+  return %r1 : vector<3xi32>
+}
+
+// CHECK-LABEL: @multi_reduction_3d_1d
+// CHECK-NEXT: Results
+// CHECK-NEXT: [60, 93, 126]
diff --git a/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/vector/outerproduct.mlir b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/vector/outerproduct.mlir
new file mode 100644
index 00000000000000..2c47a16384c284
--- /dev/null
+++ b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/vector/outerproduct.mlir
@@ -0,0 +1,155 @@
+// RUN: mlir-interpreter-runner %s -run-all | FileCheck %s
+
+func.func @outerproduct_1d_1d() -> vector<2x3xi32> {
+  %a = arith.constant dense<[1, 2]> : vector<2xi32>
+  %b = arith.constant dense<[3, 4, 5]> : vector<3xi32>
+  %o = vector.outerproduct %a, %b : vector<2xi32>, vector<3xi32>
+  return %o : vector<2x3xi32>
+}
+
+// CHECK-LABEL: @outerproduct_1d_1d
+// CHECK-NEXT: Results
+// CHECK-NEXT{LITERAL}: [[3, 4, 5], [6, 8, 10]]
+
+func.func @outerproduct_1d_1d_add() -> vector<2x3xi32> {
+  %a = arith.constant dense<[1, 2]> : vector<2xi32>
+  %b = arith.constant dense<[3, 4, 5]> : vector<3xi32>
+  %init = arith.constant dense<[[100, 200, 300], [400, 500, 600]]>
+    : vector<2x3xi32>
+  %o = vector.outerproduct %a, %b, %init : vector<2xi32>, vector<3xi32>
+  return %o : vector<2x3xi32>
+}
+
+// CHECK-LABEL: @outerproduct_1d_1d_add
+// CHECK-NEXT: Results
+// CHECK-NEXT{LITERAL}: [[103, 204, 305], [406, 508, 610]]
+
+func.func @outerproduct_1d_0d_and() -> vector<1xi32> {
+  %a = arith.constant dense<[3]> : vector<1xi32>
+  %b = arith.constant 1 : i32
+  %init = arith.constant dense<[9]> : vector<1xi32>
+  %o = vector.outerproduct %a, %b, %init {kind = #vector.kind<and>}
+    : vector<1xi32>, i32
+  return %o : vector<1xi32>
+}
+
+// CHECK-LABEL: @outerproduct_1d_0d_and
+// CHECK-NEXT: Results
+// CHECK-NEXT{LITERAL}: [1]
+
+func.func @outerproduct_1d_0d_maxui() -> vector<3xi32> {
+  %a = arith.constant dense<[1, 2, 3]> : vector<3xi32>
+  %b = arith.constant 3 : i32
+  %init = arith.constant dense<[100, 0, -100]> : vector<3xi32>
+  %o = vector.outerproduct %a, %b, %init {kind = #vector.kind<maxui>}
+    : vector<3xi32>, i32
+  return %o : vector<3xi32>
+}
+
+// CHECK-LABEL: @outerproduct_1d_0d_maxui
+// CHECK-NEXT: Results
+// CHECK-NEXT{LITERAL}: [100, 6, -100]
+
+func.func @outerproduct_1d_0d_mul() -> vector<3xi32> {
+  %a = arith.constant dense<[1, 2, 3]> : vector<3xi32>
+  %b = arith.constant 3 : i32
+  %init = arith.constant dense<[1, 2, 3]> : vector<3xi32>
+  %o = vector.outerproduct %a, %b, %init {kind = #vector.kind<mul>}
+    : vector<3xi32>, i32
+  return %o : vector<3xi32>
+}
+
+// CHECK-LABEL: @outerproduct_1d_0d_mul
+// CHECK-NEXT: Results
+// CHECK-NEXT{LITERAL}: [3, 12, 27]
+
+func.func @outerproduct_1d_0d_minui() -> vector<3xi32> {
+  %a = arith.constant dense<[1, 2, 3]> : vector<3xi32>
+  %b = arith.constant 3 : i32
+  %init = arith.constant dense<[100, 0, -100]> : vector<3xi32>
+  %o = vector.outerproduct %a, %b, %init {kind = #vector.kind<minui>}
+    : vector<3xi32>, i32
+  return %o : vector<3xi32>
+}
+
+// CHECK-LABEL: @outerproduct_1d_0d_minui
+// CHECK-NEXT: Results
+// CHECK-NEXT{LITERAL}: [3, 0, 9]
+
+func.func @outerproduct_1d_0d_maxsi() -> vector<3xi32> {
+  %a = arith.constant dense<[1, 2, 3]> : vector<3xi32>
+  %b = arith.constant 3 : i32
+  %init = arith.constant dense<[100, 0, -100]> : vector<3xi32>
+  %o = vector.outerproduct %a, %b, %init {kind = #vector.kind<maxsi>}
+    : vector<3xi32>, i32
+  return %o : vector<3xi32>
+}
+
+// CHECK-LABEL: @outerproduct_1d_0d_maxsi
+// CHECK-NEXT: Results
+// CHECK-NEXT{LITERAL}: [100, 6, 9]
+
+func.func @outerproduct_1d_0d_minsi() -> vector<3xi32> {
+  %a = arith.constant dense<[1, 2, 3]> : vector<3xi32>
+  %b = arith.constant 3 : i32
+  %init = arith.constant dense<[100, 0, -100]> : vector<3xi32>
+  %o = vector.outerproduct %a, %b, %init {kind = #vector.kind<minsi>}
+    : vector<3xi32>, i32
+  return %o : vector<3xi32>
+}
+
+// CHECK-LABEL: @outerproduct_1d_0d_minsi
+// CHECK-NEXT: Results
+// CHECK-NEXT{LITERAL}: [3, 0, -100]
+
+func.func @outerproduct_1d_0d_maxf() -> vector<1xf32> {
+  %a = arith.constant dense<[1.0]> : vector<1xf32>
+  %b = arith.constant 3.0 : f32
+  %init = arith.constant dense<[10.0]> : vector<1xf32>
+  %o = vector.outerproduct %a, %b, %init {kind = #vector.kind<maximumf>}
+    : vector<1xf32>, f32
+  return %o : vector<1xf32>
+}
+
+// CHECK-LABEL: @outerproduct_1d_0d_maxf
+// CHECK-NEXT: Results
+// CHECK-NEXT{LITERAL}: [1.000000e+01]
+
+func.func @outerproduct_1d_0d_minf() -> vector<1xf32> {
+  %a = arith.constant dense<[1.0]> : vector<1xf32>
+  %b = arith.constant 3.0 : f32
+  %init = arith.constant dense<[10.0]> : vector<1xf32>
+  %o = vector.outerproduct %a, %b, %init {kind = #vector.kind<minimumf>}
+    : vector<1xf32>, f32
+  return %o : vector<1xf32>
+}
+
+// CHECK-LABEL: @outerproduct_1d_0d_minf
+// CHECK-NEXT: Results
+// CHECK-NEXT{LITERAL}: [3.000000e+00]
+
+func.func @outerproduct_1d_0d_or() -> vector<1xi32> {
+  %a = arith.constant dense<[3]> : vector<1xi32>
+  %b = arith.constant 1 : i32
+  %init = arith.constant dense<[9]> : vector<1xi32>
+  %o = vector.outerproduct %a, %b, %init {kind = #vector.kind<or>}
+    : vector<1xi32>, i32
+  return %o : vector<1xi32>
+}
+
+// CHECK-LABEL: @outerproduct_1d_0d_or
+// CHECK-NEXT: Results
+// CHECK-NEXT{LITERAL}: [11]
+
+func.func @outerproduct_1d_0d_xor() -> vector<1xi32> {
+  %a = arith.constant dense<[3]> : vector<1xi32>
+  %b = arith.constant 1 : i32
+  %init = arith.constant dense<[9]> : vector<1xi32>
+  %o = vector.outerproduct %a, %b, %init {kind = #vector.kind<xor>}
+    : vector<1xi32>, i32
+  return %o : vector<1xi32>
+}
+
+// CHECK-LABEL: @outerproduct_1d_0d_xor
+// CHECK-NEXT: Results
+// CHECK-NEXT{LITERAL}: [10]
diff --git a/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/vector/reduction.mlir b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/vector/reduction.mlir
new file mode 100644
index 00000000000000..c62fe869fc6459
--- /dev/null
+++ b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/vector/reduction.mlir
@@ -0,0 +1,235 @@
+// RUN: mlir-interpreter-runner %s -run-all | FileCheck %s
+
+func.func @reduction_add() -> i32 {
+  %a = arith.constant dense<[10, 4, 7]> : vector<3xi32>
+  %r = vector.reduction <add>, %a : vector<3xi32> into i32
+  return %r : i32
+}
+
+// CHECK-LABEL: @reduction_add
+// CHECK-NEXT: Results
+// CHECK-NEXT: 21
+
+func.func @reduction_minf_degenerate() -> f32 {
+  %a = arith.constant dense<[1.0]> : vector<1xf32>
+  %r = vector.reduction <minimumf>, %a : vector<1xf32> into f32
+  return %r : f32
+}
+
+// CHECK-LABEL: @reduction_minf_degenerate
+// CHECK-NEXT: Results
+// CHECK-NEXT: 1.000000e+00
+
+func.func @reduction_minf() -> f32 {
+  %a = arith.constant dense<[10.0, 4.0, 7.0]> : vector<3xf32>
+  %r = vector.reduction <minimumf>, %a : vector<3xf32> into f32
+  return %r : f32
+}
+
+// CHECK-LABEL: @reduction_minf
+// CHECK-NEXT: Results
+// CHECK-NEXT: 4.000000e+00
+
+func.func @reduction_acc() -> f32 {
+  %a = arith.constant dense<[10.0, 4.0, 7.0]> : vector<3xf32>
+  %acc = arith.constant 3.0 : f32
+  %r = vector.reduction <minimumf>, %a, %acc : vector<3xf32> into f32
+  return %r : f32
+}
+
+// CHECK-LABEL: @reduction_acc
+// CHECK-NEXT: Results
+// CHECK-NEXT: 3.000000e+00
+
+func.func @reduction_acc_first_is_minimum() -> f32 {
+  %a = arith.constant dense<[1.0, 4.0, 7.0]> : vector<3xf32>
+  %acc = arith.constant 3.0 : f32
+  %r = vector.reduction <minimumf>, %a, %acc : vector<3xf32> into f32
+  return %r : f32
+}
+
+// CHECK-LABEL: @reduction_acc_first_is_minimum
+// CHECK-NEXT: Results
+// CHECK-NEXT: 1.000000e+00
+
+func.func @masked_and() -> i32 {
+  %a = arith.constant dense<[255, 127, 6]> : vector<3xi32>
+  %m = arith.constant dense<[true, true, false]> : vector<3xi1>
+  %r = vector.mask %m {
+    vector.reduction <and>, %a : vector<3xi32> into i32
+  } : vector<3xi1> -> i32
+  return %r : i32
+}
+
+// CHECK-LABEL: @masked_and
+// CHECK-NEXT: Results
+// CHECK-NEXT: i32: 127
+
+func.func @masked_xor() -> i32 {
+  %a = arith.constant dense<[255, 1, 3]> : vector<3xi32>
+  %m = arith.constant dense<[false, true, true]> : vector<3xi1>
+  %r = vector.mask %m {
+    vector.reduction <xor>, %a : vector<3xi32> into i32
+  } : vector<3xi1> -> i32
+  return %r : i32
+}
+
+// CHECK-LABEL: @masked_xor
+// CHECK-NEXT: Results
+// CHECK-NEXT: i32: 2
+
+func.func @masked_or() -> i32 {
+  %a = arith.constant dense<[255, 1, 3]> : vector<3xi32>
+  %m = arith.constant dense<[false, true, true]> : vector<3xi1>
+  %r = vector.mask %m {
+    vector.reduction <or>, %a : vector<3xi32> into i32
+  } : vector<3xi1> -> i32
+  return %r : i32
+}
+
+// CHECK-LABEL: @masked_or
+// CHECK-NEXT: Results
+// CHECK-NEXT: i32: 3
+
+func.func @masked_add_i32() -> i32 {
+  %a = arith.constant dense<[255, 1, 3]> : vector<3xi32>
+  %m = arith.constant dense<[false, true, true]> : vector<3xi1>
+  %r = vector.mask %m {
+    vector.reduction <add>, %a : vector<3xi32> into i32
+  } : vector<3xi1> -> i32
+  return %r : i32
+}
+
+// CHECK-LABEL: @masked_add_i32
+// CHECK-NEXT: Results
+// CHECK-NEXT: i32: 4
+
+func.func @masked_add_f32() -> f32 {
+  %a = arith.constant dense<[255.0, 1.0, 3.0]> : vector<3xf32>
+  %m = arith.constant dense<[false, true, true]> : vector<3xi1>
+  %r = vector.mask %m {
+    vector.reduction <add>, %a : vector<3xf32> into f32
+  } : vector<3xi1> -> f32
+  return %r : f32
+}
+
+// CHECK-LABEL: @masked_add_f32
+// CHECK-NEXT: Results
+// CHECK-NEXT: f32: 4.0
+
+func.func @masked_mul_i32() -> i32 {
+  %a = arith.constant dense<[255, 2, 3]> : vector<3xi32>
+  %m = arith.constant dense<[false, true, true]> : vector<3xi1>
+  %r = vector.mask %m {
+    vector.reduction <mul>, %a : vector<3xi32> into i32
+  } : vector<3xi1> -> i32
+  return %r : i32
+}
+
+// CHECK-LABEL: @masked_mul_i32
+// CHECK-NEXT: Results
+// CHECK-NEXT: i32: 6
+
+func.func @masked_mul_f32() -> f32 {
+  %a = arith.constant dense<[255.0, 2.0, 3.0]> : vector<3xf32>
+  %m = arith.constant dense<[false, true, true]> : vector<3xi1>
+  %r = vector.mask %m {
+    vector.reduction <mul>, %a : vector<3xf32> into f32
+  } : vector<3xi1> -> f32
+  return %r : f32
+}
+
+// CHECK-LABEL: @masked_mul_f32
+// CHECK-NEXT: Results
+// CHECK-NEXT: f32: 6.0
+
+func.func @masked_minsi() -> i32 {
+  %a = arith.constant dense<[255, -2, -5]> : vector<3xi32>
+  %m = arith.constant dense<[true, true, false]> : vector<3xi1>
+  %r = vector.mask %m {
+    vector.reduction <minsi>, %a : vector<3xi32> into i32
+  } : vector<3xi1> -> i32
+  return %r : i32
+}
+
+// CHECK-LABEL: @masked_minsi
+// CHECK-NEXT: Results
+// CHECK-NEXT: i32: -2
+
+func.func @masked_minui() -> i32 {
+  %a = arith.constant dense<[255, -2, -5]> : vector<3xi32>
+  %m = arith.constant dense<[true, true, false]> : vector<3xi1>
+  %r = vector.mask %m {
+    vector.reduction <minui>, %a : vector<3xi32> into i32
+  } : vector<3xi1> -> i32
+  return %r : i32
+}
+
+// CHECK-LABEL: @masked_minui
+// CHECK-NEXT: Results
+// CHECK-NEXT: i32: 255
+
+func.func @masked_maxsi() -> i32 {
+  %a = arith.constant dense<[255, -2, 500]> : vector<3xi32>
+  %m = arith.constant dense<[true, true, false]> : vector<3xi1>
+  %r = vector.mask %m {
+    vector.reduction <maxsi>, %a : vector<3xi32> into i32
+  } : vector<3xi1> -> i32
+  return %r : i32
+}
+
+// CHECK-LABEL: @masked_maxsi
+// CHECK-NEXT: Results
+// CHECK-NEXT: i32: 255
+
+func.func @masked_maxui() -> i32 {
+  %a = arith.constant dense<[255, -5, -2]> : vector<3xi32>
+  %m = arith.constant dense<[true, true, false]> : vector<3xi1>
+  %r = vector.mask %m {
+    vector.reduction <maxui>, %a : vector<3xi32> into i32
+  } : vector<3xi1> -> i32
+  return %r : i32
+}
+
+// CHECK-LABEL: @masked_maxui
+// CHECK-NEXT: Results
+// CHECK-NEXT: i32: -5
+
+func.func @masked_minf() -> f32 {
+  %a = arith.constant dense<[255.0, 2.0, 3.0]> : vector<3xf32>
+  %m = arith.constant dense<[true, false, true]> : vector<3xi1>
+  %r = vector.mask %m {
+    vector.reduction <minimumf>, %a : vector<3xf32> into f32
+  } : vector<3xi1> -> f32
+  return %r : f32
+}
+
+// CHECK-LABEL: @masked_minf
+// CHECK-NEXT: Results
+// CHECK-NEXT: f32: 3.0
+
+func.func @masked_maxf() -> f32 {
+  %a = arith.constant dense<[255.0, 2.0, 3.0]> : vector<3xf32>
+  %m = arith.constant dense<[false, true, true]> : vector<3xi1>
+  %r = vector.mask %m {
+    vector.reduction <maximumf>, %a : vector<3xf32> into f32
+  } : vector<3xi1> -> f32
+  return %r : f32
+}
+
+// CHECK-LABEL: @masked_maxf
+// CHECK-NEXT: Results
+// CHECK-NEXT: f32: 3.0
+
+func.func @masked_maxf_empty() -> f32 {
+  %a = arith.constant dense<[255.0]> : vector<1xf32>
+  %m = arith.constant dense<[false]> : vector<1xi1>
+  %r = vector.mask %m {
+    vector.reduction <maximumf>, %a : vector<1xf32> into f32
+  } : vector<1xi1> -> f32
+  return %r : f32
+}
+
+// CHECK-LABEL: @masked_maxf_empty
+// CHECK-NEXT: Results
+// CHECK-NEXT: f32: -INF
diff --git a/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/vector/shape_cast.mlir b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/vector/shape_cast.mlir
new file mode 100644
index 00000000000000..333adcbc4b1dcb
--- /dev/null
+++ b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/vector/shape_cast.mlir
@@ -0,0 +1,23 @@
+// RUN: mlir-interpreter-runner %s -run-all | FileCheck %s
+
+func.func @shape_cast() -> vector<2x2xi32> {
+  %a = arith.constant dense<[[1, 2, 3, 4]]> : vector<1x4xi32>
+  %cast = vector.shape_cast %a : vector<1x4xi32> to vector<2x2xi32>
+  return %cast : vector<2x2xi32>
+}
+
+// CHECK-LABEL: @shape_cast
+// CHECK-NEXT: Results
+// CHECK-NEXT{LITERAL}: [[1, 2], [3, 4]]
+
+func.func @cast_of_transpose() -> (vector<3x2xi32>, vector<2x3xi32>) {
+  %a = arith.constant dense<[[1, 2, 3], [4, 5, 6]]> : vector<2x3xi32>
+  %b = vector.transpose %a, [1, 0] : vector<2x3xi32> to vector<3x2xi32>
+  %cast = vector.shape_cast %b : vector<3x2xi32> to vector<2x3xi32>
+  return %b, %cast : vector<3x2xi32>, vector<2x3xi32>
+}
+
+// CHECK-LABEL: @cast_of_transpose
+// CHECK-NEXT: Results
+// CHECK-NEXT{LITERAL}: [[1, 4], [2, 5], [3, 6]]
+// CHECK-NEXT{LITERAL}: [[1, 4, 2], [5, 3, 6]]
diff --git a/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/vector/shuffle.mlir b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/vector/shuffle.mlir
new file mode 100644
index 00000000000000..3fb3d881fbee51
--- /dev/null
+++ b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/vector/shuffle.mlir
@@ -0,0 +1,34 @@
+// RUN: mlir-interpreter-runner %s -run-all | FileCheck %s
+
+func.func @shuffle0d() -> vector<4xi32> {
+  %c1 = arith.constant dense<1> : vector<i32>
+  %c2 = arith.constant dense<2> : vector<i32>
+  %shuffle = vector.shuffle %c1, %c2[0, 1, 1, 0] : vector<i32>, vector<i32>
+  return %shuffle : vector<4xi32>
+}
+
+// CHECK-LABEL: @shuffle0d
+// CHECK-NEXT: Results
+// CHECK-NEXT: vector<4xi32>: [1, 2, 2, 1]
+
+func.func @shuffle1d() -> vector<4xi32> {
+  %c1 = arith.constant dense<[1, 2]> : vector<2xi32>
+  %c2 = arith.constant dense<[11, 12, 13, 14]> : vector<4xi32>
+  %shuffle = vector.shuffle %c1, %c2[0, 2, 5, 1] : vector<2xi32>, vector<4xi32>
+  return %shuffle : vector<4xi32>
+}
+
+// CHECK-LABEL: @shuffle1d
+// CHECK-NEXT: Results
+// CHECK-NEXT: vector<4xi32>: [1, 11, 14, 2]
+
+func.func @shuffle2d() -> vector<3x2xi32> {
+  %c1 = arith.constant dense<[[1, 2], [11, 12]]> : vector<2x2xi32>
+  %c2 = arith.constant dense<[[21, 22]]> : vector<1x2xi32>
+  %shuffle = vector.shuffle %c1, %c2[0, 2, 1] : vector<2x2xi32>, vector<1x2xi32>
+  return %shuffle : vector<3x2xi32>
+}
+
+// CHECK-LABEL: @shuffle2d
+// CHECK-NEXT: Results
+// CHECK-NEXT{LITERAL}: vector<3x2xi32>: [[1, 2], [21, 22], [11, 12]]
diff --git a/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/vector/splat.mlir b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/vector/splat.mlir
new file mode 100644
index 00000000000000..bfbb76b38af58f
--- /dev/null
+++ b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/vector/splat.mlir
@@ -0,0 +1,11 @@
+// RUN: mlir-interpreter-runner %s -run-all | FileCheck %s
+
+func.func @splat() -> vector<2x4xi32> {
+  %c42 = arith.constant 42 : i32
+  %splat = vector.splat %c42 : vector<2x4xi32>
+  return %splat : vector<2x4xi32>
+}
+
+// CHECK-LABEL: @splat
+// CHECK-NEXT: Results
+// CHECK-NEXT{LITERAL}: vector<2x4xi32>: [[42, 42, 42, 42], [42, 42, 42, 42]]
diff --git a/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/vector/store.mlir b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/vector/store.mlir
new file mode 100644
index 00000000000000..f97f6435d61230
--- /dev/null
+++ b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/vector/store.mlir
@@ -0,0 +1,39 @@
+// RUN: mlir-interpreter-runner %s -run-all | FileCheck %s
+
+func.func @store_vector_memref() -> memref<1x2xvector<2xi32>> {
+  %m = memref.alloc() : memref<1x2xvector<2xi32>>
+  %c0 = arith.constant 0 : index
+  %c1 = arith.constant 1 : index
+  %v = arith.constant dense<[1,2]> : vector<2xi32>
+  vector.store %v, %m[%c0, %c1] : memref<1x2xvector<2xi32>>, vector<2xi32>
+  return %m : memref<1x2xvector<2xi32>>
+}
+
+// CHECK-LABEL: @store_vector_memref
+// CHECK-NEXT: Results
+// CHECK-NEXT{LITERAL}: TensorOrMemref<1x2xvector<2xi32>>: [[[0, 0], [1, 2]]]
+
+func.func @store_scalar_memref() -> memref<2x2xi32> {
+  %m = memref.alloc() : memref<2x2xi32>
+  %v = arith.constant dense<[[1,2]]> : vector<1x2xi32>
+  %c0 = arith.constant 0 : index
+  %c1 = arith.constant 1 : index
+  vector.store %v, %m[%c1, %c0] : memref<2x2xi32>, vector<1x2xi32>
+  return %m : memref<2x2xi32>
+}
+
+// CHECK-LABEL: @store_scalar_memref
+// CHECK-NEXT: Results
+// CHECK-NEXT{LITERAL}: TensorOrMemref<2x2xi32>: [[0, 0], [1, 2]]
+
+func.func @store_oob() -> memref<2x2xi32> {
+  %m = memref.alloc() : memref<2x2xi32>
+  %v = arith.constant dense<[[1,2]]> : vector<1x2xi32>
+  %c1 = arith.constant 1 : index
+  vector.store %v, %m[%c1, %c1] : memref<2x2xi32>, vector<1x2xi32>
+  return %m : memref<2x2xi32>
+}
+
+// CHECK-LABEL: @store_oob
+// CHECK-NEXT: Results
+// CHECK-NEXT{LITERAL}: TensorOrMemref<2x2xi32>: [[0, 0], [0, 1]]
diff --git a/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/vector/transfer_read.mlir b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/vector/transfer_read.mlir
new file mode 100644
index 00000000000000..6cafe0b2f4cd06
--- /dev/null
+++ b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/vector/transfer_read.mlir
@@ -0,0 +1,118 @@
+// RUN: mlir-interpreter-runner %s -run-all | FileCheck %s
+
+// Adapted from mlir/test/Integration/Dialect/Vector/CPU/test-transfer-read.mlir.
+
+func.func @transfer_read_1d() -> vector<13xi32> {
+  %a = arith.constant dense<[0, 1, 2, 3, 4]> : memref<5xi32>
+  %c2 = arith.constant 2 : index
+  %c-42 = arith.constant -42 : i32
+  %f = vector.transfer_read %a[%c2], %c-42
+      {permutation_map = affine_map<(d0) -> (d0)>} :
+    memref<5xi32>, vector<13xi32>
+  return %f : vector<13xi32>
+}
+
+// CHECK-LABEL: @transfer_read_1d
+// CHECK-NEXT: Results
+// CHECK-NEXT: vector<13xi32>: [2, 3, 4, -42, -42, -42, -42, -42, -42, -42, -42, -42, -42]
+
+func.func @transfer_read_mask_1d() -> vector<13xi32> {
+  %a = arith.constant dense<[0, 1, 2, 3, 4]> : memref<5xi32>
+  %c2 = arith.constant 2 : index
+  %c-42 = arith.constant -42: i32
+  %m = arith.constant dense<[0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0]> : vector<13xi1>
+  %f = vector.transfer_read %a[%c2], %c-42, %m : memref<5xi32>, vector<13xi32>
+  return %f : vector<13xi32>
+}
+
+// CHECK-LABEL: @transfer_read_mask_1d
+// CHECK-NEXT: Results
+// CHECK-NEXT: vector<13xi32>: [-42, -42, 4, -42, -42, -42, -42, -42, -42, -42, -42, -42, -42]
+
+func.func @transfer_read_vector_mask() -> vector<6xi32> {
+  %a = arith.constant dense<[0, 1, 2, 3, 4]> : memref<5xi32>
+  %c2 = arith.constant 2 : index
+  %c-42 = arith.constant -42: i32
+  %m = arith.constant dense<[0, 0, 1, 1, 1, 1]> : vector<6xi1>
+  %passthrough = arith.constant dense<[-1, -2, -3, -4, -5, -6]> : vector<6xi32>
+  %f = vector.mask %m, %passthrough {
+    vector.transfer_read %a[%c2], %c-42 : memref<5xi32>, vector<6xi32>
+  } : vector<6xi1> -> vector<6xi32>
+  return %f : vector<6xi32>
+}
+
+// CHECK-LABEL: @transfer_read_vector_mask
+// CHECK-NEXT: Results
+// CHECK-NEXT: vector<6xi32>: [-1, -2, 4, -42, -42, -42]
+
+func.func @transfer_read_inbounds_4() -> vector<4xi32> {
+  %a = arith.constant dense<[0, 1, 2, 0, 0]> : memref<5xi32>
+  %c-42 = arith.constant -42: i32
+  %c1 = arith.constant 1 : index
+  %f = vector.transfer_read %a[%c1], %c-42
+      {permutation_map = affine_map<(d0) -> (d0)>, in_bounds = [true]} :
+    memref<5xi32>, vector<4xi32>
+  return %f : vector<4xi32>
+}
+
+// CHECK-LABEL: @transfer_read_inbounds_4
+// CHECK-NEXT: Results
+// CHECK-NEXT: vector<4xi32>: [1, 2, 0, 0]
+
+func.func @transfer_read_mask_inbounds_4() -> vector<4xi32> {
+  %a = arith.constant dense<[0, 1, 2, 0, 0]> : memref<5xi32>
+  %c-42 = arith.constant -42: i32
+  %c1 = arith.constant 1 : index
+  %m = arith.constant dense<[0, 1, 0, 1]> : vector<4xi1>
+  %f = vector.transfer_read %a[%c1], %c-42, %m {in_bounds = [true]}
+      : memref<5xi32>, vector<4xi32>
+  return %f : vector<4xi32>
+}
+
+// CHECK-LABEL: @transfer_read_mask_inbounds_4
+// CHECK-NEXT: Results
+// CHECK-NEXT: vector<4xi32>: [-42, 2, -42, 0]
+
+func.func @transfer_read_2d()-> vector<2x2xi32> {
+  %a = arith.constant dense<[[0, 1, 2, 3], [4, 5, 6, 7], [8, 9, 10, 11]]>
+    : memref<3x4xi32>
+  %c0 = arith.constant 0 : index
+  %c1 = arith.constant 1 : index
+  %c-42 = arith.constant -42: i32
+  %f = vector.transfer_read %a[%c1, %c0], %c-42
+      : memref<3x4xi32>, vector<2x2xi32>
+  return %f : vector<2x2xi32>
+}
+
+// CHECK-LABEL: @transfer_read_2d
+// CHECK-NEXT: Results
+// CHECK-NEXT{LITERAL}: vector<2x2xi32>: [[4, 5], [8, 9]]
+
+func.func @transfer_read_2d_1d() -> vector<2xi32> {
+  %a = arith.constant dense<[[0, 1, 2, 3], [4, 5, 6, 7]]> : memref<2x4xi32>
+  %c0 = arith.constant 0 : index
+  %c1 = arith.constant 1 : index
+  %c-42 = arith.constant -42: i32
+  %f = vector.transfer_read %a[%c1, %c0], %c-42
+      : memref<2x4xi32>, vector<2xi32>
+  return %f : vector<2xi32>
+}
+
+// CHECK-LABEL: @transfer_read_2d_1d
+// CHECK-NEXT: Results
+// CHECK-NEXT{LITERAL}: vector<2xi32>: [4, 5]
+
+func.func @transfer_read_i1() -> vector<2xi1> {
+  %a = arith.constant dense<[true, false]> : memref<2xi1>
+  %c0 = arith.constant 0 : index
+  %false = arith.constant false
+  // Note: this read is technically illegal, but it can occur in practice and
+  // persist for a while until it's lowered to something legal.
+  %ret = vector.transfer_read %a[%c0], %false {in_bounds = [true]}
+    : memref<2xi1>, vector<2xi1>
+  return %ret : vector<2xi1>
+}
+
+// CHECK-LABEL: @transfer_read_i1
+// CHECK-NEXT: Results
+// CHECK-NEXT{LITERAL}: [true, false]
diff --git a/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/vector/transfer_write.mlir b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/vector/transfer_write.mlir
new file mode 100644
index 00000000000000..6ca9eb8860395f
--- /dev/null
+++ b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/vector/transfer_write.mlir
@@ -0,0 +1,91 @@
+// RUN: mlir-interpreter-runner %s -run-all | FileCheck %s
+
+func.func @write_4_at_3_inbounds() -> memref<?xi32> {
+  %c8 = arith.constant 8 : index
+  %a = memref.alloc(%c8) : memref<?xi32>
+  %base = arith.constant 3 : index
+  %f = arith.constant dense<[1, 2, 3, 4]> : vector<4xi32>
+  vector.transfer_write %f, %a[%base]
+    {permutation_map = affine_map<(d0) -> (d0)>, in_bounds = [true]}
+    : vector<4xi32>, memref<?xi32>
+  return %a : memref<?xi32>
+}
+
+// CHECK-LABEL: @write_4_at_3_inbounds
+// CHECK-NEXT: Results
+// CHECK-NEXT: [0, 0, 0, 1, 2, 3, 4, 0]
+
+func.func @write_6_at_5() -> memref<8xi32> {
+  %a = memref.alloc() : memref<8xi32>
+  %base = arith.constant 5 : index
+  %f = arith.constant dense<[1, 2, 3, 4, 5, 6]> : vector<6xi32>
+  vector.transfer_write %f, %a[%base]
+    {permutation_map = affine_map<(d0) -> (d0)>}
+    : vector<6xi32>, memref<8xi32>
+  return %a : memref<8xi32>
+}
+
+// CHECK-LABEL: @write_6_at_5
+// CHECK-NEXT: Results
+// CHECK-NEXT: [0, 0, 0, 0, 0, 1, 2, 3]
+
+func.func @write_to_tensor() -> (tensor<3xi32>, tensor<3xi32>) {
+  %a = arith.constant dense<[1,2,3]> : tensor<3xi32>
+  %base = arith.constant 1 : index
+  %f = arith.constant dense<[4, 5]> : vector<2xi32>
+  %b = vector.transfer_write %f, %a[%base]
+    {permutation_map = affine_map<(d0) -> (d0)>}
+    : vector<2xi32>, tensor<3xi32>
+  return %a, %b : tensor<3xi32>, tensor<3xi32>
+}
+
+// CHECK-LABEL: @write_to_tensor
+// CHECK-NEXT: Results
+// CHECK-NEXT: [1, 2, 3]
+// CHECK-NEXT: [1, 4, 5]
+
+func.func @write_masked() -> memref<4xi32> {
+  %a = arith.constant dense<[10, 11, 12, 13]> : memref<4xi32>
+  %base = arith.constant 1 : index
+  %f = arith.constant dense<[1, 2, 3]> : vector<3xi32>
+  %mask = arith.constant dense<[true, false, true]> : vector<3xi1>
+  vector.transfer_write %f, %a[%base], %mask
+    {permutation_map = affine_map<(d0) -> (d0)>}
+    : vector<3xi32>, memref<4xi32>
+  return %a : memref<4xi32>
+}
+
+// CHECK-LABEL: @write_masked
+// CHECK-NEXT: Results
+// CHECK-NEXT: [10, 1, 12, 3]
+
+func.func @write_vector_mask() -> memref<4xi32> {
+  %a = arith.constant dense<[10, 11, 12, 13]> : memref<4xi32>
+  %base = arith.constant 1 : index
+  %f = arith.constant dense<[1, 2, 3]> : vector<3xi32>
+  %mask = arith.constant dense<[true, false, true]> : vector<3xi1>
+  vector.mask %mask {
+    vector.transfer_write %f, %a[%base]
+      {permutation_map = affine_map<(d0) -> (d0)>}
+      : vector<3xi32>, memref<4xi32>
+  } : vector<3xi1>
+  return %a : memref<4xi32>
+}
+
+// CHECK-LABEL: @write_vector_mask
+// CHECK-NEXT: Results
+// CHECK-NEXT: [10, 1, 12, 3]
+
+func.func @write_1d_to_2d() -> memref<2x4xi32> {
+  %a = memref.alloc() : memref<2x4xi32>
+  %c1 = arith.constant 1 : index
+  %c2 = arith.constant 2 : index
+  %f = arith.constant dense<[1, 2]> : vector<2xi32>
+  vector.transfer_write %f, %a[%c1, %c2]
+    : vector<2xi32>, memref<2x4xi32>
+  return %a : memref<2x4xi32>
+}
+
+// CHECK-LABEL: @write_1d_to_2d
+// CHECK-NEXT: Results
+// CHECK-NEXT{LITERAL}: [[0, 0, 0, 0], [0, 0, 1, 2]]
\ No newline at end of file
diff --git a/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/vector/transpose.mlir b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/vector/transpose.mlir
new file mode 100644
index 00000000000000..3c39327d67d92a
--- /dev/null
+++ b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/vector/transpose.mlir
@@ -0,0 +1,28 @@
+// RUN: mlir-interpreter-runner %s -run-all | FileCheck %s
+
+func.func @transpose() -> vector<2x1x4x3xi32> {
+  %0 = arith.constant dense<[[[
+      [000, 001, 002, 003],
+      [010, 011, 012, 013],
+      [020, 021, 022, 023]
+    ],
+    [
+      [100, 101, 102, 103],
+      [110, 111, 112, 113],
+      [120, 121, 122, 123]
+    ]]]> : vector<1x2x3x4xi32>
+  %1 = vector.transpose %0, [1, 0, 3, 2]
+    : vector<1x2x3x4xi32> to vector<2x1x4x3xi32>
+  return %1 : vector<2x1x4x3xi32>
+}
+
+// CHECK-LABEL: @transpose
+// CHECK-NEXT: Results
+// CHECK-NEXT{LITERAL}: [[[[0, 10, 20],
+// CHECK{LITERAL}:         [1, 11, 21],
+// CHECK{LITERAL}:         [2, 12, 22],
+// CHECK{LITERAL}:         [3, 13, 23]]],
+// CHECK{LITERAL}:       [[[100, 110, 120],
+// CHECK{LITERAL}:         [101, 111, 121],
+// CHECK{LITERAL}:         [102, 112, 122],
+// CHECK{LITERAL}:         [103, 113, 123]]]]
diff --git a/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/vector/type_cast.mlir b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/vector/type_cast.mlir
new file mode 100644
index 00000000000000..fecc8e21466152
--- /dev/null
+++ b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/vector/type_cast.mlir
@@ -0,0 +1,11 @@
+// RUN: mlir-interpreter-runner %s -run-all | FileCheck %s
+
+func.func @type_cast() -> memref<vector<2x2xi32>> {
+  %alloc = arith.constant dense<[[1, 2], [3, 4]]> : memref<2x2xi32>
+  %cast = vector.type_cast %alloc: memref<2x2xi32> to memref<vector<2x2xi32>>
+  return %cast : memref<vector<2x2xi32>>
+}
+
+// CHECK-LABEL: @type_cast
+// CHECK-NEXT: Results
+// CHECK-NEXT{LITERAL}: TensorOrMemref<vector<2x2xi32>>: [[1, 2], [3, 4]]
diff --git a/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/vector/vscale.mlir b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/vector/vscale.mlir
new file mode 100644
index 00000000000000..6d162357d57936
--- /dev/null
+++ b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/tests/vector/vscale.mlir
@@ -0,0 +1,12 @@
+// RUN: mlir-interpreter-runner %s -run-all | FileCheck %s
+
+func.func @vscale() -> (vector<[4]xi32>, index) {
+  %c = arith.constant dense<0> : vector<[4]xi32>
+  %vscale = vector.vscale
+  return %c, %vscale : vector<[4]xi32>, index
+}
+
+// CHECK-LABEL: @vscale
+// CHECK-NEXT: Results
+// CHECK-NEXT: [0, 0, 0, 0]
+// CHECK-NEXT: 1
\ No newline at end of file
diff --git a/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/util.cc b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/util.cc
new file mode 100644
index 00000000000000..209b14dff76a87
--- /dev/null
+++ b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/util.cc
@@ -0,0 +1,180 @@
+/* Copyright 2022 The OpenXLA Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+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.
+==============================================================================*/
+
+#include "xla/mlir/tools/mlir_interpreter/dialects/util.h"
+
+#include <cassert>
+#include <cstdint>
+#include <variant>
+
+#include "llvm/ADT/ArrayRef.h"
+#include "llvm/ADT/SmallVector.h"
+#include "mlir/IR/AffineExpr.h"  // from @llvm-project
+#include "mlir/IR/BuiltinTypeInterfaces.h"  // from @llvm-project
+#include "mlir/IR/BuiltinTypes.h"  // from @llvm-project
+#include "mlir/IR/Operation.h"  // from @llvm-project
+#include "mlir/IR/ValueRange.h"  // from @llvm-project
+#include "mlir/Interfaces/ViewLikeInterface.h"  // from @llvm-project
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+#include "mlir/Support/MathExtras.h"  // from @llvm-project
+#include "xla/mlir/tools/mlir_interpreter/framework/interpreter.h"
+#include "xla/mlir/tools/mlir_interpreter/framework/interpreter_value.h"
+
+namespace mlir {
+namespace interpreter {
+
+SmallVector<int64_t> ReplaceDynamicVals(llvm::ArrayRef<int64_t> static_vals,
+                                        ArrayRef<InterpreterValue>& args) {
+  llvm::SmallVector<int64_t> out;
+  for (int64_t val : static_vals) {
+    if (ShapedType::isDynamic(val)) {
+      out.push_back(std::get<int64_t>(args.front().storage));
+      args = args.drop_front(1);
+    } else {
+      out.push_back(val);
+    }
+  }
+  return out;
+}
+
+SmallVector<int64_t> ReplaceDynamicVals(ArrayRef<int64_t> static_vals,
+                                        ArrayRef<int64_t> dynamic_vals) {
+  llvm::SmallVector<int64_t> out;
+  for (int64_t val : static_vals) {
+    if (ShapedType::isDynamic(val)) {
+      out.push_back(dynamic_vals.front());
+      dynamic_vals = dynamic_vals.drop_front(1);
+    } else {
+      out.push_back(val);
+    }
+  }
+  assert(dynamic_vals.empty() && "expected no leftover dynamic values");
+  return out;
+}
+
+OffsetsSizesStrides ExtractOffsetsSizesStrides(
+    ArrayRef<InterpreterValue> args, OffsetSizeAndStrideOpInterface op) {
+  auto offsets = ReplaceDynamicVals(op.getStaticOffsets(), args);
+  auto sizes = ReplaceDynamicVals(op.getStaticSizes(), args);
+  auto strides = ReplaceDynamicVals(op.getStaticStrides(), args);
+  return {offsets, sizes, strides};
+}
+
+OffsetsSizesStrides ExtractOffsetsSizesStrides(
+    ArrayRef<int64_t> dynamic_offsets, ArrayRef<int64_t> dynamic_sizes,
+    ArrayRef<int64_t> dynamic_strides, OffsetSizeAndStrideOpInterface op) {
+  auto offsets = ReplaceDynamicVals(op.getStaticOffsets(), dynamic_offsets);
+  auto sizes = ReplaceDynamicVals(op.getStaticSizes(), dynamic_sizes);
+  auto strides = ReplaceDynamicVals(op.getStaticStrides(), dynamic_strides);
+  return {offsets, sizes, strides};
+}
+
+InterpreterValue ReshapeTensor(const InterpreterValue& in,
+                               ArrayRef<int64_t> shape) {
+  // This doesn't need a copy in many cases, but it's easier that way.
+  auto out = in.TypedAlike(shape);
+  for (const auto& [inIndex, outIndex] :
+       llvm::zip(in.View().Indices(), out.View().Indices())) {
+    out.InsertElement(outIndex, in.ExtractElement(inIndex));
+  }
+  return out;
+}
+
+InterpreterValue GetInitOperand(mlir::Operation* op, int64_t index,
+                                MutableArrayRef<InterpreterValue> args) {
+  return GetInitOperand(op->getOperands(), index, args);
+}
+
+InterpreterValue GetInitOperand(mlir::ValueRange values, int64_t index,
+                                ArrayRef<InterpreterValue> args) {
+  assert(args.size() == values.size() && "expected matching sizes");
+  return mlir::isa<TensorType>(values[index].getType()) ? args[index].Clone()
+                                                        : args[index];
+}
+
+InterpreterValue TransposeImpl(const InterpreterValue& in,
+                               ArrayRef<int64_t> permutation) {
+  auto out = in;
+  auto& view = out.View();
+
+  view.sizes.clear();
+  view.strides.clear();
+  for (int64_t p : permutation) {
+    view.sizes.push_back(in.View().sizes[p]);
+    view.strides.push_back(in.View().strides[p]);
+  }
+
+  return out;
+}
+
+int64_t DimImpl(const InterpreterValue& in, int64_t index,
+                InterpreterState& state) {
+  if (index < 0 || index >= in.View().Rank()) {
+    state.AddFailure("dimension index out of bounds");
+    return 0;
+  }
+  return in.View().sizes[index];
+}
+
+llvm::SmallVector<InterpreterValue> NoOpTerminator(
+    MutableArrayRef<InterpreterValue> args, mlir::Operation*,
+    InterpreterState&) {
+  return llvm::to_vector(args);
+}
+
+int64_t EvalAffineExpr(AffineExpr expr, ArrayRef<int64_t> dims,
+                       ArrayRef<int64_t> symbols) {
+  int64_t lhs = 0, rhs = 0;
+  if (auto bin = expr.dyn_cast<AffineBinaryOpExpr>()) {
+    lhs = EvalAffineExpr(bin.getLHS(), dims, symbols);
+    rhs = EvalAffineExpr(bin.getRHS(), dims, symbols);
+  }
+  switch (expr.getKind()) {
+    case AffineExprKind::Add:
+      return lhs + rhs;
+    case AffineExprKind::Mul:
+      return lhs * rhs;
+    case AffineExprKind::Mod:
+      return mod(lhs, rhs);
+    case AffineExprKind::FloorDiv:
+      return floorDiv(lhs, rhs);
+    case AffineExprKind::CeilDiv:
+      return ceilDiv(lhs, rhs);
+    case AffineExprKind::Constant:
+      return expr.cast<AffineConstantExpr>().getValue();
+    case AffineExprKind::DimId:
+      return dims[expr.cast<AffineDimExpr>().getPosition()];
+    case AffineExprKind::SymbolId:
+      return symbols[expr.cast<AffineSymbolExpr>().getPosition()];
+  }
+}
+
+SmallVector<int64_t> EvalAffineMap(AffineMap map, ArrayRef<int64_t> dims,
+                                   ArrayRef<int64_t> symbols) {
+  SmallVector<int64_t> result;
+  for (auto expr : map.getResults()) {
+    result.push_back(EvalAffineExpr(expr, dims, symbols));
+  }
+  return result;
+}
+
+llvm::SmallVector<int64_t> EvalAffineMap(AffineMap map,
+                                         ArrayRef<int64_t> operands) {
+  return EvalAffineMap(map, operands.take_front(map.getNumDims()),
+                       operands.drop_front(map.getNumDims()));
+}
+
+}  // namespace interpreter
+}  // namespace mlir
diff --git a/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/util.h b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/util.h
new file mode 100644
index 00000000000000..be769a2920d76e
--- /dev/null
+++ b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/util.h
@@ -0,0 +1,85 @@
+/* Copyright 2023 The OpenXLA Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+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.
+==============================================================================*/
+
+#ifndef XLA_MLIR_TOOLS_MLIR_INTERPRETER_DIALECTS_UTIL_H_
+#define XLA_MLIR_TOOLS_MLIR_INTERPRETER_DIALECTS_UTIL_H_
+
+#include <cstdint>
+
+#include "mlir/IR/AffineExpr.h"  // from @llvm-project
+#include "mlir/IR/AffineMap.h"  // from @llvm-project
+#include "mlir/IR/Operation.h"  // from @llvm-project
+#include "mlir/IR/ValueRange.h"  // from @llvm-project
+#include "mlir/Interfaces/ViewLikeInterface.h"  // from @llvm-project
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+#include "xla/mlir/tools/mlir_interpreter/framework/interpreter.h"
+#include "xla/mlir/tools/mlir_interpreter/framework/interpreter_value.h"
+
+namespace mlir {
+namespace interpreter {
+
+struct OffsetsSizesStrides {
+  llvm::SmallVector<int64_t> offsets;
+  llvm::SmallVector<int64_t> sizes;
+  llvm::SmallVector<int64_t> strides;
+};
+
+// Replaces dynamic placeholders in static_vals using elements from the front
+// of args, which are removed.
+SmallVector<int64_t> ReplaceDynamicVals(ArrayRef<int64_t> static_vals,
+                                        ArrayRef<InterpreterValue>& args);
+// Same as above, but the values are already unpacked. `dynamicVals.size` must
+// match the number of dynamic values in `staticVals`.
+SmallVector<int64_t> ReplaceDynamicVals(ArrayRef<int64_t> static_vals,
+                                        ArrayRef<int64_t> dynamic_vals);
+
+OffsetsSizesStrides ExtractOffsetsSizesStrides(
+    ArrayRef<InterpreterValue> args, OffsetSizeAndStrideOpInterface op);
+OffsetsSizesStrides ExtractOffsetsSizesStrides(
+    ArrayRef<int64_t> dynamic_offsets, ArrayRef<int64_t> dynamic_sizes,
+    ArrayRef<int64_t> dynamic_strides, OffsetSizeAndStrideOpInterface op);
+
+InterpreterValue ReshapeTensor(const InterpreterValue& in,
+                               ArrayRef<int64_t> shape);
+
+// gets the given operand, cloning its storage if it is a tensor.
+InterpreterValue GetInitOperand(mlir::Operation* op, int64_t index,
+                                MutableArrayRef<InterpreterValue> args);
+InterpreterValue GetInitOperand(mlir::ValueRange values, int64_t index,
+                                ArrayRef<InterpreterValue> args);
+
+// Common implementations for ops from different dialects but sharing the same
+// semantics.
+InterpreterValue TransposeImpl(const InterpreterValue& in,
+                               ArrayRef<int64_t> permutation);
+int64_t DimImpl(const InterpreterValue& in, int64_t index,
+                InterpreterState& state);
+
+// Terminator that just returns its args.
+llvm::SmallVector<InterpreterValue> NoOpTerminator(
+    MutableArrayRef<InterpreterValue> args, mlir::Operation*,
+    InterpreterState&);
+
+int64_t EvalAffineExpr(AffineExpr expr, ArrayRef<int64_t> dims,
+                       ArrayRef<int64_t> symbols);
+llvm::SmallVector<int64_t> EvalAffineMap(AffineMap map, ArrayRef<int64_t> dims,
+                                         ArrayRef<int64_t> symbols);
+llvm::SmallVector<int64_t> EvalAffineMap(AffineMap map,
+                                         ArrayRef<int64_t> operands);
+
+}  // namespace interpreter
+}  // namespace mlir
+
+#endif  // XLA_MLIR_TOOLS_MLIR_INTERPRETER_DIALECTS_UTIL_H_
diff --git a/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/vector.cc b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/vector.cc
new file mode 100644
index 00000000000000..0bfd5ddd17cd84
--- /dev/null
+++ b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/vector.cc
@@ -0,0 +1,875 @@
+/* Copyright 2023 The OpenXLA Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+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.
+==============================================================================*/
+
+#include <algorithm>
+#include <cassert>
+#include <cstdint>
+#include <cstring>
+#include <functional>
+#include <limits>
+#include <memory>
+#include <optional>
+#include <type_traits>
+#include <utility>
+
+#include "llvm/ADT/STLExtras.h"
+#include "llvm/ADT/SmallVector.h"
+#include "llvm/Support/ErrorHandling.h"
+#include "mlir/Dialect/Vector/IR/VectorOps.h"  // from @llvm-project
+#include "mlir/IR/BuiltinAttributes.h"  // from @llvm-project
+#include "mlir/IR/BuiltinTypeInterfaces.h"  // from @llvm-project
+#include "mlir/IR/BuiltinTypes.h"  // from @llvm-project
+#include "mlir/IR/TypeUtilities.h"  // from @llvm-project
+#include "mlir/IR/Types.h"  // from @llvm-project
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+#include "xla/mlir/tools/mlir_interpreter/dialects/comparators.h"
+#include "xla/mlir/tools/mlir_interpreter/dialects/util.h"
+#include "xla/mlir/tools/mlir_interpreter/framework/interpreter.h"
+#include "xla/mlir/tools/mlir_interpreter/framework/interpreter_value.h"
+#include "xla/mlir/tools/mlir_interpreter/framework/registration.h"
+#include "xla/mlir/tools/mlir_interpreter/framework/tensor_or_memref.h"
+
+namespace mlir {
+namespace interpreter {
+namespace {
+
+class MaskSideChannel : public InterpreterSideChannel {
+ public:
+  MaskSideChannel(TensorOrMemref<bool> mask,
+                  std::optional<InterpreterValue> passthrough)
+      : mask_(std::move(mask)), passthrough_(std::move(passthrough)) {}
+
+  const TensorOrMemref<bool>& GetMask() const { return mask_; }
+  const std::optional<InterpreterValue>& GetPassthrough() const {
+    return passthrough_;
+  }
+
+ private:
+  const TensorOrMemref<bool> mask_;
+  const std::optional<InterpreterValue> passthrough_;
+};
+
+template <typename T>
+using combiner_t = T (*)(T, T);
+
+template <typename T>
+combiner_t<T> GetCombiner(vector::CombiningKind kind) {
+  if constexpr (std::is_arithmetic_v<T> && !std::is_same_v<T, bool>) {
+    switch (kind) {
+      case vector::CombiningKind::ADD:
+        return +[](T a, T b) -> T { return a + b; };
+      case vector::CombiningKind::MUL:
+        return +[](T a, T b) -> T { return a * b; };
+      case vector::CombiningKind::MINSI:
+      case vector::CombiningKind::MINIMUMF:
+        return +[](T a, T b) -> T { return std::min(a, b); };
+      case vector::CombiningKind::MAXSI:
+      case vector::CombiningKind::MAXIMUMF:
+        return +[](T a, T b) -> T { return std::max(a, b); };
+      default: {
+      }
+    }
+  }
+
+  if constexpr (std::is_integral_v<T> && !std::is_same_v<T, bool>) {
+    switch (kind) {
+      case vector::CombiningKind::MINUI:
+        return &Iumin::apply<T>;
+      case vector::CombiningKind::MAXUI:
+        return &Iumax::apply<T>;
+      default: {
+      }
+    }
+  }
+
+  if constexpr (std::is_integral_v<T>) {
+    switch (kind) {
+      case vector::CombiningKind::AND:
+        return +[](T a, T b) -> T { return a & b; };
+      case vector::CombiningKind::OR:
+        return +[](T a, T b) -> T { return a | b; };
+      case vector::CombiningKind::XOR:
+        return +[](T a, T b) -> T { return a ^ b; };
+      default: {
+      }
+    }
+  }
+
+  llvm_unreachable("unknown combining kind");
+}
+
+InterpreterValue GetNeutralElement(vector::CombiningKind kind, mlir::Type ty) {
+  return DispatchScalarType(ty, [&](auto dummy) -> InterpreterValue {
+    using T = decltype(dummy);
+    switch (kind) {
+      case vector::CombiningKind::AND:
+      case vector::CombiningKind::MINUI:
+        if constexpr (std::is_same_v<T, bool>) {
+          return {true};
+        } else if constexpr (std::is_integral_v<T>) {
+          return {static_cast<T>(static_cast<std::make_unsigned_t<T>>(-1))};
+        }
+        break;
+      case vector::CombiningKind::ADD:
+      case vector::CombiningKind::MAXUI:
+      case vector::CombiningKind::OR:
+      case vector::CombiningKind::XOR:
+        return {T{0}};
+      case vector::CombiningKind::MUL:
+        return {T{1}};
+      case vector::CombiningKind::MINSI:
+        return {std::numeric_limits<T>::max()};
+      case vector::CombiningKind::MINIMUMF:
+        return {std::numeric_limits<T>::infinity()};
+      case vector::CombiningKind::MAXSI:
+        return {std::numeric_limits<T>::min()};  // NOLINT
+      case vector::CombiningKind::MAXIMUMF:
+        return {-std::numeric_limits<T>::infinity()};
+      default: {
+      }
+    }
+    llvm_unreachable("invalid combining kind");
+  });
+}
+
+template <typename IntType>
+SmallVector<IntType> ExtractVector(ArrayAttr array_attr) {
+  return llvm::to_vector<4>(llvm::map_range(
+      array_attr.getAsRange<IntegerAttr>(),
+      [](IntegerAttr attr) { return static_cast<IntType>(attr.getInt()); }));
+}
+
+InterpreterValue Bitcast(InterpreterState&, vector::BitCastOp op,
+                         const InterpreterValue& vector) {
+  ShapedType ty = cast<ShapedType>(op->getResultTypes()[0]);
+  auto flattened = vector.CoerceLayout({});
+  auto buffer = flattened.Buffer();
+  auto view = flattened.View();
+  view.sizes = llvm::to_vector(ty.getShape());
+  view.strides = BufferView::GetDefaultStrides(view.sizes);
+  return DispatchScalarType(ty, [&](auto dummy) -> InterpreterValue {
+    // TODO(jreiffers): i1 semantics are currently broken (both here and
+    // upstream).
+    using T = decltype(dummy);
+    return {TensorOrMemref<T>{buffer, view}};
+  });
+}
+
+InterpreterValue Broadcast(InterpreterState&, vector::BroadcastOp broadcast,
+                           const InterpreterValue& value) {
+  auto result =
+      value.IsTensor() ? value : value.AsUnitTensor(/*is_vector=*/true);
+  auto& result_view = result.View();
+
+  // Insert additional leading stride 0 dims.
+  SmallVector<int64_t> strides(broadcast.getResultVectorType().getRank());
+  llvm::copy(llvm::reverse(result_view.strides), strides.rbegin());
+  // Zero out broadcast dimension strides.
+  for (int64_t i : broadcast.computeBroadcastedUnitDims()) {
+    strides[i] = 0;
+  }
+
+  result_view.strides = std::move(strides);
+  result_view.sizes =
+      llvm::to_vector(broadcast.getResultVectorType().getShape());
+  return result;
+}
+
+void CompressStore(InterpreterState& state, vector::CompressStoreOp,
+                   InterpreterValue dst, ArrayRef<int64_t> indices,
+                   TensorOrMemref<bool> mask, InterpreterValue value) {
+  auto dst_buffer = dst.Buffer();
+  const auto& dst_view = dst.View();
+  if (dst_view.strides.back() != 1) {
+    state.AddFailure("trailing dimension must be continguous");
+    return;
+  }
+  auto offset = dst_view.GetPhysicalIndex(indices);
+  if (!offset) {
+    state.AddFailure("index out of bounds");
+    return;
+  }
+
+  auto src_buffer = value.Buffer();
+  const auto& src_view = value.View();
+
+  // TODO(jreiffers): Bounds checks.
+  int64_t n = src_view.sizes[0];
+  int64_t element_size = value.GetByteSizeOfElement();
+  for (int64_t i = 0; i < n; ++i) {
+    if (mask.at(i)) {
+      memcpy(dst_buffer->at(offset, element_size),
+             src_buffer->at(i, element_size), element_size);
+      ++*offset;
+    }
+  }
+}
+
+InterpreterValue MaskImpl(mlir::Operation* op, ArrayRef<int64_t> mask_sizes) {
+  auto out_sizes = cast<ShapedType>(op->getResultTypes()[0]).getShape();
+  auto result = TensorOrMemref<bool>::Empty(out_sizes);
+  result.view.is_vector = true;
+  BufferView iter;
+  iter.sizes = llvm::to_vector(mask_sizes);
+  for (const auto& indices : iter.Indices()) {
+    result.at(indices) = true;
+  }
+  return {result};
+}
+
+InterpreterValue ConstantMask(InterpreterState&, vector::ConstantMaskOp mask) {
+  return MaskImpl(mask, ExtractVector<int64_t>(mask.getMaskDimSizes()));
+}
+
+// TODO(jreiffers): Support masked contractions.
+InterpreterValue Contract(InterpreterState&, vector::ContractionOp contraction,
+                          const InterpreterValue& lhs,
+                          const InterpreterValue& rhs,
+                          const InterpreterValue& acc) {
+  BufferView iter;
+  contraction.getIterationBounds(iter.sizes);
+  auto maps = contraction.getIndexingMapsArray();
+  auto resultTy = contraction->getResultTypes()[0];
+  auto shapedTy = resultTy.dyn_cast<ShapedType>();
+  auto result =
+      DispatchScalarType(resultTy, [&](auto dummy) -> InterpreterValue {
+        using T = decltype(dummy);
+        using TT = TensorOrMemref<T>;
+        const auto& lhs_t = std::get<TT>(lhs.storage);
+        const auto& rhs_t = std::get<TT>(rhs.storage);
+        auto result_value = shapedTy ? acc.Clone() : acc.AsUnitTensor();
+        auto& result = std::get<TT>(result_value.storage);
+        auto combiner = *GetCombiner<T>(contraction.getKind());
+        for (const auto& indices : iter.Indices()) {
+          auto lhs_indices = EvalAffineMap(maps[0], indices);
+          auto rhs_indices = EvalAffineMap(maps[1], indices);
+          auto result_indices = EvalAffineMap(maps[2], indices);
+
+          auto& resultItem = result.at(result_indices);
+          resultItem = combiner(resultItem,
+                                lhs_t.at(lhs_indices) * rhs_t.at(rhs_indices));
+        }
+        return result_value;
+      });
+
+  return shapedTy ? result : result.ExtractElement({});
+}
+
+InterpreterValue CreateMask(InterpreterState&, vector::CreateMaskOp op,
+                            ArrayRef<int64_t> sizes) {
+  return MaskImpl(op, sizes);
+}
+
+InterpreterValue ExpandLoad(InterpreterState& state, vector::ExpandLoadOp,
+                            InterpreterValue memref,
+                            SmallVector<int64_t> offsets,
+                            TensorOrMemref<bool> mask,
+                            const InterpreterValue& passThrough) {
+  if (memref.View().strides.back() != 1) {
+    state.AddFailure("expected last dimension to be contiguous");
+    return {};
+  }
+
+  auto out = passThrough.Clone();
+  for (int64_t i = 0, e = out.View().sizes[0]; i < e; ++i) {
+    if (mask.at({i})) {
+      out.InsertElement({i}, memref.ExtractElement(offsets));
+      ++offsets.back();
+    }
+  }
+  return out;
+}
+
+InterpreterValue Extract(InterpreterState& state, vector::ExtractOp extract,
+                         const InterpreterValue& vector) {
+  auto result = vector;
+  auto& result_view = result.View();
+  for (int64_t offset : extract.getStaticPosition()) {
+    state.CheckSuccess(result_view.Slice(0, offset), "index out of bounds");
+  }
+  return result_view.Rank() == 0 ? result.ExtractElement({}) : result;
+}
+
+InterpreterValue ExtractElement(InterpreterState& state,
+                                vector::ExtractElementOp,
+                                const InterpreterValue& vector,
+                                std::optional<int64_t> index) {
+  if (!index) {
+    return vector.ExtractElement({});
+  }
+  if (!vector.View().InBounds(*index)) {
+    state.AddFailure("array index out of bounds");
+    return {};
+  }
+  return vector.ExtractElement(*index);
+}
+
+InterpreterValue ExtractSlice(InterpreterState& state,
+                              vector::ExtractStridedSliceOp extract,
+                              const InterpreterValue& vector) {
+  auto out = vector;
+  state.CheckSuccess(
+      out.View().Subview(ExtractVector<int64_t>(extract.getOffsets()),
+                         ExtractVector<int64_t>(extract.getSizes()),
+                         ExtractVector<int64_t>(extract.getStrides())),
+      "subview out of bounds");
+  return out;
+}
+
+InterpreterValue FlatTranspose(InterpreterState&,
+                               vector::FlatTransposeOp transpose,
+                               const InterpreterValue& vector) {
+  auto out = vector.Clone();
+  // We currently only implement -matrix-default-layout=column-major.
+  int64_t rows = transpose.getRows();
+  int64_t cols = transpose.getColumns();
+  for (int64_t i = 0; i < rows * cols; ++i) {
+    int64_t srcIndex = (i % cols) * rows + (i / cols);
+    out.InsertElement({i}, vector.ExtractElement({srcIndex}));
+  }
+  return out;
+}
+
+InterpreterValue FusedMultiplyAdd(InterpreterState&, vector::FMAOp op,
+                                  const InterpreterValue& lhs,
+                                  const InterpreterValue& rhs,
+                                  const InterpreterValue& acc) {
+  auto out = acc.Clone();
+  DispatchScalarType(op->getResultTypes()[0], [&](auto dummy) {
+    using TT = TensorOrMemref<decltype(dummy)>;
+    const auto& lhs_t = std::get<TT>(lhs.storage);
+    const auto& rhs_t = std::get<TT>(rhs.storage);
+    auto& result = std::get<TT>(out.storage);
+
+    for (const auto& indices : lhs_t.view.Indices()) {
+      result.at(indices) += lhs_t.at(indices) * rhs_t.at(indices);
+    }
+  });
+  return out;
+}
+
+InterpreterValue Gather(InterpreterState& state, vector::GatherOp op,
+                        const InterpreterValue& src, ArrayRef<int64_t> offsets,
+                        const InterpreterValue& indices,
+                        const TensorOrMemref<bool>& mask,
+                        const InterpreterValue& passThrough) {
+  if (isa<MemRefType>(op->getOperandTypes()[0]) &&
+      src.View().strides.back() != 1) {
+    state.AddFailure("expected trailing dimension to be contiguous");
+    return {};
+  }
+  auto out = passThrough.Clone();
+  for (const auto& outIndex : out.View().Indices()) {
+    if (!mask.at(outIndex)) {
+      continue;
+    }
+    auto inIndex = llvm::to_vector(offsets);
+    inIndex[0] += indices.ExtractElement(outIndex).AsInt();
+    out.InsertElement(outIndex, src.ExtractElement(inIndex));
+  }
+  return out;
+}
+
+InterpreterValue Insert(InterpreterState& state, vector::InsertOp insert,
+                        const InterpreterValue& src,
+                        const InterpreterValue& dst) {
+  auto result = dst.Clone();
+  auto resultSlice = result;
+  auto& resultSliceView = resultSlice.View();
+  for (int64_t offset : insert.getStaticPosition()) {
+    state.CheckSuccess(resultSliceView.Slice(0, offset), "index out of bounds");
+  }
+  resultSlice.Fill([&](auto indices) { return src.ExtractElement(indices); });
+  return result;
+}
+
+InterpreterValue InsertElement(InterpreterState& state, vector::InsertElementOp,
+                               const InterpreterValue& value,
+                               const InterpreterValue& vector,
+                               std::optional<int64_t> index) {
+  auto result = vector.Clone();
+  if (!index) {
+    result.InsertElement({}, value);
+    return result;
+  }
+  if (!result.View().InBounds(*index)) {
+    state.AddFailure("array index out of bounds");
+    return {};
+  }
+  result.InsertElement(*index, value);
+  return result;
+}
+
+InterpreterValue InsertSlice(InterpreterState& state,
+                             vector::InsertStridedSliceOp insert,
+                             InterpreterValue src,
+                             const InterpreterValue& dst) {
+  auto out = dst.Clone();
+  auto out_slice = out;
+  if (!out_slice.View()
+           .Subview(ExtractVector<int64_t>(insert.getOffsets()),
+                    insert.getSourceVectorType().getShape(),
+                    ExtractVector<int64_t>(insert.getStrides()))
+           .succeeded()) {
+    state.AddFailure("subview out of bounds");
+  }
+  for (const auto& index : src.View().Indices()) {
+    out_slice.InsertElement(index, src.ExtractElement(index));
+  }
+  return out;
+}
+
+InterpreterValue Load(InterpreterState& state, vector::LoadOp load,
+                      const InterpreterValue& memref,
+                      ArrayRef<int64_t> indices) {
+  if (memref.View().num_vector_dims > 0) {
+    return {memref.ExtractElement(indices)};
+  }
+  auto out = memref;
+  if (!out.View()
+           .Subview(indices, load.getVectorType().getShape(),
+                    SmallVector<int64_t>(load.getVectorType().getRank(), 1))
+           .succeeded()) {
+    // "Not all targets may support out-of-bounds vector loads"
+    state.AddFailure("out of bounds loads not supported");
+    return {};
+  }
+  out = out.Clone();
+  out.View().is_vector = true;
+  return out;
+}
+
+llvm::SmallVector<InterpreterValue> Mask(
+    InterpreterState& state, vector::MaskOp op, TensorOrMemref<bool> mask,
+    std::optional<InterpreterValue> passthrough) {
+  InterpreterScope scope(state);
+  scope.SetSideChannel(std::make_shared<MaskSideChannel>(mask, passthrough));
+  return Interpret(state, op.getMaskRegion(), {});
+}
+
+InterpreterValue MaskedLoad(InterpreterState& state, vector::MaskedLoadOp,
+                            const InterpreterValue& memref,
+                            SmallVector<int64_t> offsets,
+                            TensorOrMemref<bool> mask,
+                            const InterpreterValue& passThrough) {
+  if (memref.View().strides.back() != 1) {
+    state.AddFailure("expected last dimension to be contiguous");
+    return {};
+  }
+
+  auto out = passThrough.Clone();
+  for (int64_t i = 0, e = mask.view.sizes[0]; i < e; ++i) {
+    if (mask.at(i)) {
+      out.InsertElement(i, memref.ExtractElement(offsets));
+    }
+    ++offsets.back();
+  }
+  return out;
+}
+
+void MaskedStore(InterpreterState& state, vector::MaskedStoreOp,
+                 InterpreterValue memref, SmallVector<int64_t> offsets,
+                 TensorOrMemref<bool> mask, const InterpreterValue& vector) {
+  if (memref.View().strides.back() != 1) {
+    state.AddFailure("expected last dimension to be contiguous");
+    return;
+  }
+
+  for (int64_t i = 0, e = mask.view.sizes[0]; i < e; ++i) {
+    if (mask.at({i})) {
+      memref.InsertElement(offsets, vector.ExtractElement({i}));
+    }
+    ++offsets.back();
+  }
+}
+
+InterpreterValue ReductionImpl(InterpreterState& state,
+                               const InterpreterValue& v,
+                               const InterpreterValue* acc,
+                               vector::CombiningKind kind,
+                               SmallVector<int64_t> dims, Type elementType) {
+  llvm::sort(dims);
+  SmallVector<int64_t> keptDims;
+  SmallVector<int64_t> resultShape;
+  for (auto [dim, size] : llvm::enumerate(v.View().sizes)) {
+    if (!llvm::is_contained(dims, dim)) {
+      keptDims.push_back(dim);
+      resultShape.push_back(size);
+    }
+  }
+  return DispatchScalarType(elementType, [&](auto dummy) -> InterpreterValue {
+    using T = decltype(dummy);
+    using TT = TensorOrMemref<T>;
+
+    auto combiner = *GetCombiner<T>(kind);
+
+    TT result = acc ? std::get<TT>((keptDims.empty() ? acc->AsUnitTensor()
+                                                     : acc->Clone())
+                                       .storage)
+                    : TensorOrMemref<T>::Empty(resultShape);
+
+    for (const auto& result_index : result.view.Indices()) {
+      auto src = std::get<TT>(v.storage);
+      for (auto [dim, index] :
+           llvm::reverse(llvm::zip(keptDims, result_index))) {
+        state.CheckSuccess(src.view.Slice(dim, index), "index out of bounds");
+      }
+
+      T& item = result.at(result_index);
+      bool first = acc == nullptr;
+      for (const auto& srcIndex : src.view.Indices()) {
+        if (first) {
+          item = src.at(srcIndex);
+          first = false;
+        } else {
+          item = combiner(item, src.at(srcIndex));
+        }
+      }
+    }
+
+    if (keptDims.empty()) {
+      return {result.at({})};
+    }
+    return {result};
+  });
+}
+
+InterpreterValue MultiReduction(InterpreterState& state,
+                                vector::MultiDimReductionOp reduction,
+                                const InterpreterValue& source,
+                                const InterpreterValue& acc) {
+  auto elementTy = getElementTypeOrSelf(reduction->getResultTypes()[0]);
+  return {ReductionImpl(state, source, &acc, reduction.getKind(),
+                        ExtractVector<int64_t>(reduction.getReductionDims()),
+                        elementTy)};
+}
+
+InterpreterValue OuterProduct(InterpreterState&,
+                              vector::OuterProductOp outerproduct,
+                              const InterpreterValue& lhs,
+                              const InterpreterValue& rhs,
+                              std::optional<InterpreterValue> acc) {
+  ShapedType ty = cast<ShapedType>(outerproduct->getResultTypes()[0]);
+  return DispatchScalarType(ty, [&](auto dummy) -> InterpreterValue {
+    using T = decltype(dummy);
+    using TT = TensorOrMemref<T>;
+    const TT& lhs_t = std::get<TT>(lhs.storage);
+
+    auto combiner = GetCombiner<T>(outerproduct.getKind());
+    auto result =
+        acc ? std::get<TT>(acc->storage).Clone() : TT::Empty(ty.getShape());
+    if (std::holds_alternative<T>(rhs.storage)) {
+      T rhsS = std::get<T>(rhs.storage);
+      for (int64_t i : llvm::seq(int64_t{0}, lhs_t.view.sizes[0])) {
+        result.at(i) = combiner(result.at(i), lhs_t.at(i) * rhsS);
+      }
+    } else {
+      const TT& rhs_t = std::get<TT>(rhs.storage);
+      for (int64_t i : llvm::seq(int64_t{0}, lhs_t.view.sizes[0])) {
+        for (int64_t j : llvm::seq(int64_t{0}, rhs_t.view.sizes[0])) {
+          result.at({i, j}) =
+              combiner(result.at({i, j}), lhs_t.at(i) * rhs_t.at(j));
+        }
+      }
+    }
+    return {result};
+  });
+}
+
+InterpreterValue Reduction(InterpreterState& state,
+                           vector::ReductionOp reduction, InterpreterValue arg,
+                           std::optional<InterpreterValue> acc) {
+  auto* mask =
+      state.GetTopScope()->GetSideChannel<MaskSideChannel>(/*optional=*/true);
+  auto ty = reduction->getResultTypes()[0];
+  if (mask) {
+    if (mask->GetPassthrough()) {
+      state.AddFailure("passthrough should not be set with masked reduction");
+      return {};
+    }
+    arg = arg.Clone();
+    if (mask->GetMask().view.sizes != arg.View().sizes) {
+      state.AddFailure("mask shape should match argument shape");
+      return {};
+    }
+    auto neutral = GetNeutralElement(reduction.getKind(), ty);
+    for (const auto& idx : arg.View().Indices()) {
+      if (!mask->GetMask().at(idx)) {
+        arg.InsertElement(idx, neutral);
+      }
+    }
+  }
+
+  return ReductionImpl(state, arg, acc ? &acc.value() : nullptr,
+                       reduction.getKind(), {0}, ty);
+}
+
+InterpreterValue ShapeCast(InterpreterState&, vector::ShapeCastOp op,
+                           const InterpreterValue& in) {
+  auto out = in.CoerceLayout({});
+  auto& outView = out.View();
+  outView.sizes =
+      llvm::to_vector(op->getResultTypes()[0].cast<ShapedType>().getShape());
+  outView.strides = BufferView::GetDefaultStrides(outView.sizes);
+  return out;
+}
+
+InterpreterValue Shuffle(InterpreterState& state, vector::ShuffleOp shuffle,
+                         const InterpreterValue& v0,
+                         const InterpreterValue& v1) {
+  auto result = v0.TypedAlike(shuffle.getResultVectorType().getShape());
+  auto& result_view = result.View();
+  result_view.is_vector = true;
+
+  auto mask = ExtractVector<int64_t>(shuffle.getMask());
+  bool isZeroDim = v0.View().Rank() == 0;
+  int64_t size0 = isZeroDim ? 1 : v0.View().sizes[0];
+  for (auto [dstIndex, srcIndex] : llvm::enumerate(mask)) {
+    auto src = srcIndex < size0 ? v0 : v1;
+    if (!isZeroDim) {
+      state.CheckSuccess(
+          src.View().Slice(0, srcIndex < size0 ? srcIndex : srcIndex - size0),
+          "index out of bounds");
+    }
+    auto dst = result;
+    state.CheckSuccess(dst.View().Slice(0, dstIndex), "index out of bounds");
+    dst.Fill([&](auto indices) { return src.ExtractElement(indices); });
+  }
+  return result;
+}
+
+InterpreterValue Splat(InterpreterState&, vector::SplatOp op,
+                       const InterpreterValue& in) {
+  auto out = in.AsUnitTensor(/*is_vector=*/true);
+  auto& view = out.View();
+  view.sizes =
+      llvm::to_vector(op->getResultTypes()[0].cast<ShapedType>().getShape());
+  view.strides = SmallVector<int64_t>(view.sizes.size(), 0);
+  return out;
+}
+
+void Store(InterpreterState& state, vector::StoreOp,
+           const InterpreterValue& src, InterpreterValue dst,
+           ArrayRef<int64_t> offsets) {
+  if (!dst.View().InBounds(offsets)) {
+    state.AddFailure("array index out of bounds");
+    return;
+  }
+  const auto& outView = dst.View();
+  if (outView.num_vector_dims > 0) {
+    dst.InsertElement(offsets, src);
+  } else {
+    for (const auto& srcIndex : src.View().Indices()) {
+      auto dstIndex = srcIndex;
+      for (int64_t i = 0; i < dstIndex.size(); ++i) {
+        dstIndex[i] += offsets[i];
+      }
+      if (outView.InBounds(dstIndex)) {
+        dst.InsertElement(dstIndex, src.ExtractElement(srcIndex));
+      }
+    }
+  }
+}
+
+std::optional<InterpreterValue> ExtractMemorySlice(
+    InterpreterState& state, const AffineMap& map,
+    const InterpreterValue& memory, const InterpreterValue& vector,
+    ArrayRef<int64_t> offsets, std::optional<ArrayAttr> in_boundsAttr) {
+  llvm::SmallVector<bool> in_bounds(offsets.size());
+  if (in_boundsAttr) {
+    llvm::copy(in_boundsAttr->getAsValueRange<BoolAttr>(),
+               in_bounds.end() - in_boundsAttr->size());
+  }
+
+  auto memSlice = memory;
+  auto& memSliceView = memSlice.View();
+  auto& vectorView = vector.View();
+  for (int64_t i = 0; i < memSliceView.Rank(); ++i) {
+    bool found = false;
+    for (int64_t j = 0; !found && j < vectorView.Rank(); ++j) {
+      if (map.getResult(j).isFunctionOfDim(i)) {
+        int64_t size = memSliceView.sizes[i] - offsets[i];
+        bool isInBounds = size >= vectorView.sizes[j];
+        if (!isInBounds && in_bounds[i]) {
+          state.AddFailure("index out of bounds");
+          return std::nullopt;
+        }
+        (void)memSliceView.Slice(
+            i, offsets[i],
+            std::max(int64_t{0}, std::min(vectorView.sizes[j], size)));
+        found = true;
+      }
+    }
+    if (!found) {
+      bool isInBounds = memSliceView.sizes[i] > offsets[i];
+      if (!isInBounds) {
+        state.AddFailure("index out of bounds");
+        return std::nullopt;
+      }
+
+      (void)memSliceView.Slice(i, offsets[i], isInBounds ? 1 : 0);
+    }
+  }
+  return memSlice;
+}
+
+InterpreterValue TransferRead(InterpreterState& state,
+                              vector::TransferReadOp transfer,
+                              const InterpreterValue& src,
+                              ArrayRef<int64_t> offsets,
+                              const InterpreterValue& padding,
+                              std::optional<TensorOrMemref<bool>> mask) {
+  auto* maskChannel = state.GetTopScope()->GetSideChannel<MaskSideChannel>(
+      /*optional=*/true);
+  if (maskChannel) {
+    if (mask) {
+      state.AddFailure(
+          "vector.mask and transfer_read with mask should not be used "
+          "simultaneously");
+      return {};
+    }
+    mask = maskChannel->GetMask();
+  }
+
+  InterpreterValue dst = src.TypedAlike(transfer.getVectorType().getShape());
+  if (maskChannel && maskChannel->GetPassthrough()) {
+    dst.Fill([&](auto indices) {
+      if (mask->at(indices)) {
+        return padding;
+      }
+      return maskChannel->GetPassthrough()->ExtractElement(indices);
+    });
+  } else {
+    dst.Fill([&](auto) { return padding; });
+  }
+  dst.View().is_vector = true;
+
+  auto srcSlice = ExtractMemorySlice(state, transfer.getPermutationMap(), src,
+                                     dst, offsets, transfer.getInBounds());
+
+  if (!srcSlice) {
+    return {};
+  }
+  for (const auto& src_indices : srcSlice->View().Indices()) {
+    SmallVector<int64_t> dstIndices =
+        EvalAffineMap(transfer.getPermutationMap(), src_indices);
+
+    // Note: the handling of padding and passthrough values is somewhat
+    // arbitrary here. At the time of writing this, there seems to be little
+    // evidence of actual usage of this feature.
+    if (!mask || mask->at(dstIndices)) {
+      dst.InsertElement(dstIndices, srcSlice->ExtractElement(src_indices));
+    }
+  }
+
+  return dst;
+}
+
+llvm::SmallVector<InterpreterValue> TransferWrite(
+    InterpreterState& state, vector::TransferWriteOp transfer,
+    InterpreterValue src, InterpreterValue dst, ArrayRef<int64_t> offsets,
+    std::optional<TensorOrMemref<bool>> mask) {
+  if (auto* maskChannel = state.GetTopScope()->GetSideChannel<MaskSideChannel>(
+          /*optional=*/true)) {
+    if (mask) {
+      state.AddFailure(
+          "vector.mask and transfer_write with mask should not be used "
+          "simultaneously");
+      return {};
+    }
+    if (maskChannel->GetPassthrough()) {
+      state.AddFailure(
+          "vector.mask with passthrough should not be used with "
+          "transfer_write");
+      return {};
+    }
+    mask = maskChannel->GetMask();
+  }
+
+  const auto& src_view = src.View();
+  assert(transfer.getPermutationMap().getNumResults() == src_view.Rank() &&
+         "expected matching number of results");
+
+  dst = transfer.getSource().getType().isa<TensorType>() ? dst.Clone() : dst;
+  auto dst_slice = ExtractMemorySlice(state, transfer.getPermutationMap(), dst,
+                                      src, offsets, transfer.getInBounds());
+  if (!dst_slice) {
+    return {};
+  }
+
+  for (const auto& dstIndices : dst_slice->View().Indices()) {
+    SmallVector<int64_t> src_indices =
+        EvalAffineMap(transfer.getPermutationMap(), dstIndices);
+    if (src_view.InBounds(src_indices) && (!mask || mask->at(src_indices))) {
+      dst_slice->InsertElement(dstIndices, src.ExtractElement(src_indices));
+    }
+  }
+
+  if (transfer->getNumResults() == 0) {
+    return {};
+  }
+  return {dst};
+}
+
+InterpreterValue Transpose(InterpreterState&, vector::TransposeOp transpose,
+                           const InterpreterValue& vector) {
+  return TransposeImpl(vector, transpose.getPermutation());
+}
+
+InterpreterValue TypeCast(InterpreterState&, vector::TypeCastOp,
+                          InterpreterValue vector) {
+  vector.View().num_vector_dims = vector.View().Rank();
+  return vector;
+}
+
+uint64_t VScale(InterpreterState&, vector::VectorScaleOp) { return 1; }
+
+REGISTER_MLIR_INTERPRETER_OP("vector.yield", NoOpTerminator);
+REGISTER_MLIR_INTERPRETER_OP(Bitcast);
+REGISTER_MLIR_INTERPRETER_OP(Broadcast);
+REGISTER_MLIR_INTERPRETER_OP(CompressStore);
+REGISTER_MLIR_INTERPRETER_OP(ConstantMask);
+REGISTER_MLIR_INTERPRETER_OP(Contract);
+REGISTER_MLIR_INTERPRETER_OP(CreateMask);
+REGISTER_MLIR_INTERPRETER_OP(ExpandLoad);
+REGISTER_MLIR_INTERPRETER_OP(Extract);
+REGISTER_MLIR_INTERPRETER_OP(ExtractElement);
+REGISTER_MLIR_INTERPRETER_OP(ExtractSlice);
+REGISTER_MLIR_INTERPRETER_OP(FusedMultiplyAdd);
+REGISTER_MLIR_INTERPRETER_OP(FlatTranspose);
+REGISTER_MLIR_INTERPRETER_OP(Gather);
+REGISTER_MLIR_INTERPRETER_OP(Insert);
+REGISTER_MLIR_INTERPRETER_OP(InsertElement);
+REGISTER_MLIR_INTERPRETER_OP(InsertSlice);
+REGISTER_MLIR_INTERPRETER_OP(Load);
+REGISTER_MLIR_INTERPRETER_OP(Mask);
+REGISTER_MLIR_INTERPRETER_OP(MaskedLoad);
+REGISTER_MLIR_INTERPRETER_OP(MaskedStore);
+REGISTER_MLIR_INTERPRETER_OP(MultiReduction);
+REGISTER_MLIR_INTERPRETER_OP(OuterProduct);
+REGISTER_MLIR_INTERPRETER_OP(Reduction);
+REGISTER_MLIR_INTERPRETER_OP(ShapeCast);
+REGISTER_MLIR_INTERPRETER_OP(Shuffle);
+REGISTER_MLIR_INTERPRETER_OP(Splat);
+REGISTER_MLIR_INTERPRETER_OP(Store);
+REGISTER_MLIR_INTERPRETER_OP(TransferRead);
+REGISTER_MLIR_INTERPRETER_OP(TransferWrite);
+REGISTER_MLIR_INTERPRETER_OP(Transpose);
+REGISTER_MLIR_INTERPRETER_OP(TypeCast);
+REGISTER_MLIR_INTERPRETER_OP(VScale);
+
+}  // namespace
+}  // namespace interpreter
+}  // namespace mlir
diff --git a/third_party/xla/xla/mlir/tools/mlir_interpreter/framework/BUILD b/third_party/xla/xla/mlir/tools/mlir_interpreter/framework/BUILD
new file mode 100644
index 00000000000000..5ac54bd0703a7e
--- /dev/null
+++ b/third_party/xla/xla/mlir/tools/mlir_interpreter/framework/BUILD
@@ -0,0 +1,42 @@
+load("@local_tsl//tsl/platform:rules_cc.bzl", "cc_library")
+
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
+    default_visibility = [":friends"],
+    licenses = ["notice"],
+)
+
+package_group(
+    name = "friends",
+    includes = [
+        "//xla:friends",
+    ],
+)
+
+cc_library(
+    name = "framework",
+    srcs = [
+        "interpreter.cc",
+        "interpreter_value.cc",
+        "registration.cc",
+        "tensor_or_memref.cc",
+    ],
+    hdrs = [
+        "interpreter.h",
+        "interpreter_value.h",
+        "interpreter_value_util.h",
+        "registration.h",
+        "tensor_or_memref.h",
+    ],
+    deps = [
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/status:statusor",
+        "@com_google_absl//absl/strings",
+        "@llvm-project//llvm:Support",
+        "@llvm-project//mlir:BufferizationInterfaces",
+        "@llvm-project//mlir:DialectUtils",
+        "@llvm-project//mlir:FuncDialect",
+        "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:Support",
+    ],
+)
diff --git a/third_party/xla/xla/mlir/tools/mlir_interpreter/framework/interpreter.cc b/third_party/xla/xla/mlir/tools/mlir_interpreter/framework/interpreter.cc
new file mode 100644
index 00000000000000..a07322c17a5f86
--- /dev/null
+++ b/third_party/xla/xla/mlir/tools/mlir_interpreter/framework/interpreter.cc
@@ -0,0 +1,150 @@
+/* Copyright 2022 The OpenXLA Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+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.
+==============================================================================*/
+
+#include "xla/mlir/tools/mlir_interpreter/framework/interpreter.h"
+
+#include <cassert>
+#include <functional>
+#include <memory>
+#include <optional>
+#include <utility>
+
+#include "absl/status/status.h"
+#include "absl/status/statusor.h"
+#include "llvm/ADT/STLExtras.h"
+#include "llvm/ADT/SmallVector.h"
+#include "llvm/ADT/StringRef.h"
+#include "llvm/Support/ErrorHandling.h"
+#include "llvm/Support/raw_ostream.h"
+#include "mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h"  // from @llvm-project
+#include "mlir/IR/OpDefinition.h"  // from @llvm-project
+#include "mlir/IR/SymbolTable.h"  // from @llvm-project
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+#include "xla/mlir/tools/mlir_interpreter/framework/interpreter_value.h"
+#include "xla/mlir/tools/mlir_interpreter/framework/registration.h"
+
+namespace mlir {
+namespace interpreter {
+
+SmallVector<InterpreterValue> Interpret(InterpreterState& state,
+                                        Operation& op) {
+  auto fn = detail::GetFunction(op.getName().getStringRef());
+  if (!fn) {
+    llvm::errs() << "Unsupported op: " << op.getName().getStringRef() << "\n";
+    op.dump();
+    state.AddFailure("unsupported op");
+    return {};
+  }
+  SmallVector<InterpreterValue> operands;
+  for (auto operand : op.getOperands()) {
+    operands.push_back(state.GetTopScope()->Get(operand));
+  }
+  state.GetOptions().listener->BeforeOp(operands, &op);
+  auto results = fn(operands, &op, state);
+  for (auto* scope = state.GetTopScope(); scope != nullptr;
+       scope = scope->GetParentScope()) {
+    scope->Verify();
+  }
+  if (state.HasFailure()) {
+    llvm::errs() << "Encountered failure while executing " << op << "\n";
+  }
+  state.GetOptions().listener->AfterOp(results);
+  state.Step();
+  return results;
+}
+
+SmallVector<InterpreterValue> Interpret(InterpreterState& state, Region& region,
+                                        ArrayRef<InterpreterValue> bbargs) {
+  if (state.HasFailure()) return {};
+  assert(region.hasOneBlock() && "expected region to have one block");
+  state.GetOptions().listener->EnterRegion(bbargs, region);
+  InterpreterScope scope(state);
+
+  auto& block = region.getBlocks().front();
+  for (auto [value, interpreter_value] :
+       llvm::zip(block.getArguments(), bbargs)) {
+    scope.Set(value, interpreter_value);
+  }
+
+  std::optional<SmallVector<InterpreterValue>> blockResults;
+  for (mlir::Operation& op : block) {
+    auto results = Interpret(state, op);
+    if (state.HasFailure()) return {};
+    if (op.hasTrait<OpTrait::IsTerminator>()) {
+      assert(!blockResults.has_value() && "Expected at most one terminator");
+      blockResults = results;
+    } else {
+      if (results.size() != op.getNumResults()) {
+        llvm::errs() << "Unexpected number of results while interpreting "
+                     << op.getName().getStringRef() << ". Interpreter bug?\n";
+        llvm_unreachable("unexpected number of results");
+      }
+      for (auto [v, iv] : llvm::zip(op.getResults(), results)) {
+        scope.Set(v, iv);
+      }
+    }
+  }
+  if (!blockResults) {
+    blockResults = SmallVector<InterpreterValue>{};
+  }
+  state.GetOptions().listener->LeaveRegion(*blockResults);
+  return *std::move(blockResults);
+}
+
+InterpreterState::InterpreterState(const mlir::SymbolTable& symbols,
+                                   InterpreterOptions options)
+    : symbols_(symbols), options_(options) {
+  if (!options_.listener) {
+    static auto& noOpListener = *new InterpreterListener();
+    this->options_.listener = &noOpListener;
+  }
+  if (options_.max_steps) {
+    remaining_steps_ = *options_.max_steps;
+  }
+}
+
+void InterpreterState::AddFailure(llvm::StringRef failure) {
+  failed_ = true;
+  options_.error_handler(failure);
+}
+
+void InterpreterScope::Verify() const {
+  for (auto& [_, value] : values_) {
+    if (value.IsTensor() && value.Buffer() &&
+        !value.Buffer()->GetFailure().empty()) {
+      state_.AddFailure(value.Buffer()->GetFailure());
+      break;
+    }
+  }
+}
+
+InterpreterScope::~InterpreterScope() {
+  Verify();
+  state_.top_scope_ = parent_scope_;
+}
+
+absl::StatusOr<SmallVector<InterpreterValue>> RunInterpreter(
+    const mlir::SymbolTable& symbols, mlir::func::FuncOp function,
+    ArrayRef<InterpreterValue> args, InterpreterOptions options) {
+  InterpreterState state{symbols, std::move(options)};
+  auto results = Interpret(state, function.getBody(), args);
+  if (state.HasFailure()) {
+    return absl::InvalidArgumentError("Interpreter failed, check error logs");
+  }
+  return results;
+}
+
+}  // namespace interpreter
+}  // namespace mlir
diff --git a/third_party/xla/xla/mlir/tools/mlir_interpreter/framework/interpreter.h b/third_party/xla/xla/mlir/tools/mlir_interpreter/framework/interpreter.h
new file mode 100644
index 00000000000000..7f14f6df435434
--- /dev/null
+++ b/third_party/xla/xla/mlir/tools/mlir_interpreter/framework/interpreter.h
@@ -0,0 +1,186 @@
+/* Copyright 2022 The OpenXLA Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+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.
+==============================================================================*/
+
+#ifndef XLA_MLIR_TOOLS_MLIR_INTERPRETER_FRAMEWORK_INTERPRETER_H_
+#define XLA_MLIR_TOOLS_MLIR_INTERPRETER_FRAMEWORK_INTERPRETER_H_
+
+#include <cstdint>
+#include <functional>
+#include <limits>
+#include <memory>
+#include <optional>
+
+#include "absl/status/statusor.h"
+#include "llvm/ADT/StringRef.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+#include "xla/mlir/tools/mlir_interpreter/framework/interpreter_value.h"
+
+namespace mlir {
+namespace interpreter {
+
+class InterpreterScope;
+
+class InterpreterListener {
+ public:
+  virtual ~InterpreterListener() = default;
+  virtual void BeforeOp(ArrayRef<InterpreterValue> args, mlir::Operation*) {}
+  virtual void AfterOp(ArrayRef<InterpreterValue> results) {}
+  virtual void EnterRegion(ArrayRef<InterpreterValue> args,
+                           mlir::Region& region) {}
+  virtual void LeaveRegion(ArrayRef<InterpreterValue> terminatorArgs) {}
+};
+
+struct InterpreterStats {
+  // Memrefs only.
+  int64_t heap_size = 0;
+  int64_t peak_heap_size = 0;
+  int64_t num_allocations = 0;
+  int64_t num_deallocations = 0;
+};
+
+struct InterpreterOptions {
+  InterpreterListener* listener = nullptr;
+  std::optional<int64_t> max_steps = std::nullopt;
+  // If set, ignore deallocations. Normally, accessing a deallocated memref will
+  // trigger an assertion. This flag disables all allocations, which can be
+  // useful when debugging IR that includes a use-after-free bug.
+  bool disable_deallocations = false;
+  std::function<void(llvm::StringRef)> error_handler =
+      [](llvm::StringRef failure) {
+        llvm::errs() << "Interpreter failure: " << failure << "\n";
+      };
+  InterpreterStats* stats = nullptr;
+};
+
+class InterpreterState {
+ public:
+  InterpreterState(const mlir::SymbolTable& symbols,
+                   InterpreterOptions options);
+
+  void Step() {
+    if (remaining_steps_ == 0) {
+      AddFailure("maximum number of steps exceeded");
+      return;
+    }
+    --remaining_steps_;
+  }
+  void AddFailure(llvm::StringRef failure);
+  bool HasFailure() const { return failed_; }
+  void CheckSuccess(LogicalResult result, llvm::StringRef failure) {
+    if (!result.succeeded()) {
+      AddFailure(failure);
+    }
+  }
+
+  InterpreterScope* GetTopScope() { return top_scope_; }
+  const mlir::SymbolTable& GetSymbols() const { return symbols_; }
+  const InterpreterOptions& GetOptions() { return options_; }
+
+ private:
+  const mlir::SymbolTable& symbols_;
+  InterpreterScope* top_scope_ = nullptr;
+  bool failed_ = false;
+  InterpreterOptions options_;
+  int64_t remaining_steps_ = std::numeric_limits<int64_t>::max();
+
+  friend class InterpreterScope;
+  friend class InterpreterScopeStash;
+};
+
+// Used for passing arbitrary data to ops in sub-regions.
+class InterpreterSideChannel {
+ public:
+  virtual ~InterpreterSideChannel() = default;
+};
+
+// Holds a mapping from SSA values to InterpreterValues and registered side
+// channels. There's typically one scope per region, but ops can add additional
+// scopes if needed (for example, to register a side channel).
+class InterpreterScope {
+ public:
+  InterpreterScope(InterpreterScope&&) = delete;
+  explicit InterpreterScope(InterpreterState& state)
+      : state_(state), parent_scope_(state.top_scope_) {
+    state.top_scope_ = this;
+  }
+  ~InterpreterScope();
+
+  void Set(Value v, InterpreterValue iv) { values_[v] = std::move(iv); }
+
+  const InterpreterValue& Get(Value v) {
+    auto ret = values_.find(v);
+    if (ret == values_.end()) {
+      if (!parent_scope_) {
+        v.dump();
+      }
+
+      assert(parent_scope_ && "value not found");
+      return parent_scope_->Get(v);
+    }
+    return ret->second;
+  }
+
+  void Verify() const;
+
+  // Retrieves the side channel of the given type in this scope or one of its
+  // ancestor scopes. If `optional` is set, returns nullptr if not found,
+  // otherwise asserts.
+  template <typename T>
+  T* GetSideChannel(bool optional = false) {
+    for (auto& side_channel : side_channels_) {
+      if (auto it = dynamic_cast<T*>(side_channel.get())) {
+        return it;
+      }
+    }
+    if (!parent_scope_ && optional) return nullptr;
+    assert(parent_scope_ && "side channel not found");
+    return parent_scope_->GetSideChannel<T>(optional);
+  }
+
+  // Registers the given side channel. Will shadow a side channel of the same
+  // type if registered in an outer scope.
+  // The behavior of registering two side channels of the same type in the same
+  // scope is undefined.
+  void SetSideChannel(std::shared_ptr<InterpreterSideChannel> side_channel) {
+    side_channels_.push_back(std::move(side_channel));
+  }
+
+  InterpreterScope* GetParentScope() const { return parent_scope_; }
+
+ private:
+  DenseMap<Value, InterpreterValue> values_;
+  SmallVector<std::shared_ptr<InterpreterSideChannel>> side_channels_;
+
+  InterpreterState& state_;
+  InterpreterScope* parent_scope_;
+
+  friend class InterpreterScopeStash;
+};
+
+// Interprets the given region and returns the terminator's arguments. The
+// region must have a single block.
+SmallVector<InterpreterValue> Interpret(InterpreterState& state, Region& region,
+                                        ArrayRef<InterpreterValue> bbargs);
+
+// Interprets the given function.
+absl::StatusOr<SmallVector<InterpreterValue>> RunInterpreter(
+    const mlir::SymbolTable& symbols, mlir::func::FuncOp function,
+    ArrayRef<InterpreterValue> args, InterpreterOptions options = {});
+
+}  // namespace interpreter
+}  // namespace mlir
+
+#endif  // XLA_MLIR_TOOLS_MLIR_INTERPRETER_FRAMEWORK_INTERPRETER_H_
diff --git a/third_party/xla/xla/mlir/tools/mlir_interpreter/framework/interpreter_value.cc b/third_party/xla/xla/mlir/tools/mlir_interpreter/framework/interpreter_value.cc
new file mode 100644
index 00000000000000..55f9a84856f2bd
--- /dev/null
+++ b/third_party/xla/xla/mlir/tools/mlir_interpreter/framework/interpreter_value.cc
@@ -0,0 +1,382 @@
+/* Copyright 2022 The OpenXLA Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+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.
+==============================================================================*/
+
+#include "xla/mlir/tools/mlir_interpreter/framework/interpreter_value.h"
+
+#include <complex>
+#include <cstdint>
+#include <functional>
+#include <iterator>
+#include <memory>
+#include <string>
+#include <string_view>
+#include <variant>
+
+#include "llvm/Support/ErrorHandling.h"
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+#include "xla/mlir/tools/mlir_interpreter/framework/tensor_or_memref.h"
+
+namespace mlir {
+namespace interpreter {
+
+namespace {
+struct TypeStr {
+  static std::string_view Get(bool) { return "i1"; }
+  static std::string_view Get(int64_t) { return "i64"; }
+  static std::string_view Get(int32_t) { return "i32"; }
+  static std::string_view Get(int16_t) { return "i16"; }
+  static std::string_view Get(int8_t) { return "i8"; }
+  static std::string_view Get(uint64_t) { return "ui64"; }
+  static std::string_view Get(uint32_t) { return "ui32"; }
+  static std::string_view Get(uint16_t) { return "ui16"; }
+  static std::string_view Get(uint8_t) { return "ui8"; }
+  static std::string_view Get(float) { return "f32"; }
+  static std::string_view Get(double) { return "f64"; }
+  static std::string_view Get(std::complex<float>) { return "complex<f32>"; }
+  static std::string_view Get(std::complex<double>) { return "complex<f64>"; }
+};
+
+struct InterpreterValuePrinter {
+  llvm::raw_ostream& os;
+
+  template <typename T>
+  void operator()(const TensorOrMemref<T>& t) {
+    if (!t.buffer) {
+      os << "Memref: null";
+      return;
+    }
+
+    if (t.view.is_vector) {
+      os << "vector<";
+    } else {
+      os << "TensorOrMemref<";
+    }
+    ArrayRef<int64_t> sizes = t.view.sizes;
+    for (int64_t size : sizes.drop_back(t.view.num_vector_dims.value_or(0))) {
+      os << size << "x";
+    }
+    if (t.view.num_vector_dims) {
+      os << "vector<";
+      for (int64_t size : sizes.take_back(*t.view.num_vector_dims)) {
+        os << size << "x";
+      }
+      os << TypeStr::Get(T{}) << ">>: ";
+    } else {
+      os << TypeStr::Get(T{}) << ">: ";
+    }
+    SmallVector<int64_t> indices(t.view.Rank() +
+                                 t.view.num_vector_dims.value_or(0));
+    std::function<void(int64_t)> print;
+    print = [&](int64_t dim) {
+      if (dim == indices.size()) {
+        PrintScalar(t.at(indices));
+      } else {
+        os << "[";
+        for (int64_t i = 0; i < t.view.sizes[dim]; ++i) {
+          if (i > 0) os << ", ";
+          indices[dim] = i;
+          print(dim + 1);
+        }
+        os << "]";
+      }
+    };
+    if (t.buffer->Deallocated()) {
+      os << "<<deallocated>>";
+    } else {
+      print(0);
+    }
+  }
+
+  void operator()(const Tuple& t) {
+    os << "(";
+    bool first = true;
+    for (const auto& v : t.values) {
+      if (!first) os << ", ";
+      first = false;
+      v->Print(os);
+    }
+    os << ")";
+  }
+
+  template <typename T>
+  void operator()(const T& t) {
+    os << TypeStr::Get(t) << ": ";
+    PrintScalar(t);
+  }
+
+  template <typename T>
+  void PrintScalar(const T& v) {
+    os << v;
+  }
+
+  template <typename T>
+  void PrintScalar(const std::complex<T>& v) {
+    os << v.real() << (v.imag() >= 0 ? "+" : "") << v.imag() << "i";
+  }
+
+  void PrintScalar(bool v) { os << (v ? "true" : "false"); }
+
+  void PrintScalar(int8_t v) { os << (int)v; }
+  void PrintScalar(uint8_t v) { os << (int)v; }
+};
+}  // namespace
+
+void InterpreterValue::Print(llvm::raw_ostream& os) const {
+  std::visit(InterpreterValuePrinter{os}, storage);
+}
+
+std::string InterpreterValue::ToString() const {
+  std::string buf;
+  llvm::raw_string_ostream os(buf);
+  Print(os);
+  return buf;
+}
+
+InterpreterValue InterpreterValue::ExtractElement(
+    llvm::ArrayRef<int64_t> indices) const {
+  return std::visit(
+      [&](auto& it) -> InterpreterValue {
+        using T = std::decay_t<decltype(it)>;
+        if constexpr (is_tensor_or_memref_v<T>) {
+          if (it.view.num_vector_dims) {
+            return {it.VectorAt(indices)};
+          } else {
+            return {it.at(indices)};
+          }
+        } else if constexpr (std::is_same_v<T, Tuple>) {
+          llvm_unreachable("extracting from tuples is unsupported");
+        } else {
+          return {it};
+        }
+      },
+      storage);
+}
+
+void InterpreterValue::InsertElement(llvm::ArrayRef<int64_t> indices,
+                                     const InterpreterValue& value) {
+  std::visit(
+      [&](auto& it) {
+        using T = std::decay_t<decltype(it)>;
+        if constexpr (is_tensor_or_memref_v<T>) {
+          if (it.view.num_vector_dims) {
+            auto subview = it.VectorAt(indices);
+            const auto& values = std::get<T>(value.storage);
+            assert(values.view.sizes == subview.view.sizes &&
+                   "mismatched sizes");
+            for (const auto& index : subview.view.Indices()) {
+              subview.at(index) = values.at(index);
+            }
+          } else {
+            it.at(indices) = std::get<typename T::element_type>(value.storage);
+          }
+        } else if constexpr (std::is_same_v<T, Tuple>) {
+          llvm_unreachable("inserting into tuples is unsupported");
+        } else {
+          it = std::get<T>(value.storage);
+        }
+      },
+      storage);
+}
+
+void InterpreterValue::Fill(
+    const std::function<InterpreterValue(llvm::ArrayRef<int64_t> indices)>& f) {
+  std::visit(
+      [&](auto& it) {
+        using T = std::decay_t<decltype(it)>;
+        if constexpr (is_tensor_or_memref_v<T>) {
+          for (const auto& indices : it.view.Indices()) {
+            if (it.view.num_vector_dims) {
+              auto subview = it.VectorAt(indices);
+              auto value = std::get<T>(f(indices).storage);
+              for (const auto& index : subview.view.Indices()) {
+                subview.at(index) = value.at(index);
+              }
+            } else {
+              it.at(indices) =
+                  std::get<typename T::element_type>(f(indices).storage);
+            }
+          }
+        } else if constexpr (std::is_same_v<T, Tuple>) {
+          llvm_unreachable("Filling tuples is unsupported");
+        } else {
+          it = std::get<T>(f({}).storage);
+        }
+      },
+      storage);
+}
+
+InterpreterValue InterpreterValue::Clone(ArrayRef<int64_t> layout) const {
+  return std::visit(
+      [&](const auto& it) -> InterpreterValue {
+        using T = std::decay_t<decltype(it)>;
+        if constexpr (is_tensor_or_memref_v<T>) {
+          return {it.Clone(layout)};
+        } else if constexpr (std::is_same_v<T, Tuple>) {
+          llvm_unreachable("cloning tuples is unsupported");
+        } else {
+          return {it};
+        }
+      },
+      storage);
+}
+
+InterpreterValue InterpreterValue::CoerceLayout(
+    ArrayRef<int64_t> layout) const {
+  const auto& view = this->View();
+  if (view.strides == BufferView::GetStridesForLayout(view.sizes, layout)) {
+    return *this;
+  }
+  return Clone(layout);
+}
+
+InterpreterValue InterpreterValue::TypedAlike(
+    llvm::ArrayRef<int64_t> shape) const {
+  return std::visit(
+      [&](const auto& it) -> InterpreterValue {
+        using T = std::decay_t<decltype(it)>;
+        if constexpr (is_tensor_or_memref_v<T>) {
+          return {T::Empty(shape)};
+        } else if constexpr (std::is_same_v<T, Tuple>) {
+          llvm_unreachable("TypedAlike for tuples is unsupported");
+        } else {
+          return {TensorOrMemref<T>::Empty(shape)};
+        }
+      },
+      storage);
+}
+
+InterpreterValue InterpreterValue::MakeTensor(mlir::Type elementType,
+                                              SmallVector<int64_t> shape) {
+  auto vectorTy = llvm::dyn_cast<VectorType>(elementType);
+  if (vectorTy) {
+    llvm::copy(vectorTy.getShape(), std::back_inserter(shape));
+  }
+  return DispatchScalarType(elementType, [&](auto dummy) -> InterpreterValue {
+    auto tensor = TensorOrMemref<decltype(dummy)>::Empty(shape);
+    if (vectorTy) {
+      tensor.view.num_vector_dims = vectorTy.getRank();
+    }
+    return {tensor};
+  });
+}
+
+BufferView& InterpreterValue::View() {
+  return std::visit(
+      [](auto& it) -> BufferView& {
+        if constexpr (is_tensor_or_memref_v<decltype(it)>) {
+          return it.view;
+        }
+        llvm_unreachable("view is only supported for tensors");
+      },
+      storage);
+}
+
+const BufferView& InterpreterValue::View() const {
+  return std::visit(
+      [](const auto& it) -> const BufferView& {
+        if constexpr (is_tensor_or_memref_v<decltype(it)>) {
+          return it.view;
+        }
+        llvm_unreachable("view is only supported for tensors");
+      },
+      storage);
+}
+
+bool InterpreterValue::IsTensor() const {
+  return std::visit(
+      [](const auto& it) { return is_tensor_or_memref_v<decltype(it)>; },
+      storage);
+}
+
+InterpreterValue InterpreterValue::AsUnitTensor(bool is_vector) const {
+  auto result = TypedAlike({});
+  result.InsertElement({}, *this);
+  result.View().is_vector = is_vector;
+  return result;
+}
+
+bool Tuple::operator==(const Tuple& other) const {
+  if (other.values.size() != values.size()) return false;
+  for (const auto& [lhs, rhs] : llvm::zip(values, other.values)) {
+    if (!(*lhs == *rhs)) return false;
+  }
+  return true;
+}
+
+std::shared_ptr<Buffer> InterpreterValue::Buffer() const {
+  return std::visit(
+      [](const auto& it) -> std::shared_ptr<interpreter::Buffer> {
+        if constexpr (is_tensor_or_memref_v<decltype(it)>) {
+          return it.buffer;
+        } else {
+          llvm_unreachable("buffer() is only supported for tensors");
+        }
+      },
+      storage);
+}
+
+int64_t InterpreterValue::AsInt() const {
+  auto visit = [](auto value) -> int64_t {
+    if constexpr (std::is_integral_v<decltype(value)>) {
+      return static_cast<int64_t>(value);
+    } else {
+      llvm_unreachable("only integral types can be converted to ints");
+    }
+  };
+  return std::visit(visit, storage);
+}
+
+uint64_t InterpreterValue::AsUInt() const {
+  auto visit = [](auto value) -> uint64_t {
+    if constexpr (std::is_integral_v<decltype(value)>) {
+      if constexpr (std::is_signed_v<decltype(value)>) {
+        return static_cast<uint64_t>(
+            static_cast<std::make_unsigned_t<decltype(value)>>(value));
+      } else {
+        return static_cast<uint64_t>(value);
+      }
+    } else {
+      llvm_unreachable("only integral types can be converted to ints");
+    }
+  };
+  return std::visit(visit, storage);
+}
+
+double InterpreterValue::AsDouble() const {
+  auto visit = [](auto value) -> int64_t {
+    if constexpr (std::is_floating_point_v<decltype(value)>) {
+      return static_cast<double>(value);
+    } else {
+      llvm_unreachable("only float types can be converted to ints");
+    }
+  };
+  return std::visit(visit, storage);
+}
+
+int64_t InterpreterValue::GetByteSizeOfElement() const {
+  return std::visit(
+      [](const auto& it) -> int64_t {
+        using T = std::decay_t<decltype(it)>;
+        if constexpr (is_tensor_or_memref_v<T>) {
+          return sizeof(typename T::element_type);
+        } else {
+          llvm_unreachable("scalars have no element sizes");
+        }
+      },
+      storage);
+}
+
+}  // namespace interpreter
+}  // namespace mlir
diff --git a/third_party/xla/xla/mlir/tools/mlir_interpreter/framework/interpreter_value.h b/third_party/xla/xla/mlir/tools/mlir_interpreter/framework/interpreter_value.h
new file mode 100644
index 00000000000000..4c0e4788f499c4
--- /dev/null
+++ b/third_party/xla/xla/mlir/tools/mlir_interpreter/framework/interpreter_value.h
@@ -0,0 +1,227 @@
+/* Copyright 2022 The OpenXLA Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+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.
+==============================================================================*/
+
+#ifndef XLA_MLIR_TOOLS_MLIR_INTERPRETER_FRAMEWORK_INTERPRETER_VALUE_H_
+#define XLA_MLIR_TOOLS_MLIR_INTERPRETER_FRAMEWORK_INTERPRETER_VALUE_H_
+
+#include <cassert>
+#include <complex>
+#include <cstdint>
+#include <functional>
+#include <memory>
+#include <optional>
+#include <string>
+#include <variant>
+
+#include "llvm/ADT/ArrayRef.h"
+#include "llvm/Support/ErrorHandling.h"
+#include "llvm/Support/raw_ostream.h"
+#include "mlir/IR/BuiltinTypes.h"  // from @llvm-project
+#include "mlir/IR/TypeUtilities.h"  // from @llvm-project
+#include "mlir/IR/Types.h"  // from @llvm-project
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+#include "xla/mlir/tools/mlir_interpreter/framework/tensor_or_memref.h"
+
+namespace mlir {
+namespace interpreter {
+
+struct InterpreterValue;
+
+struct Tuple {
+  bool operator==(const Tuple& other) const;
+
+  SmallVector<std::shared_ptr<InterpreterValue>> values;
+};
+
+// Holds a scalar, a tensor/memref or a tuple. Tensors/memrefs can also
+// represent vectors.
+struct InterpreterValue {
+  void Print(llvm::raw_ostream& os) const;
+  std::string ToString() const;
+
+  // Returns the element at the given indices. If the value is a scalar, returns
+  // itself.
+  InterpreterValue ExtractElement(llvm::ArrayRef<int64_t> indices) const;
+  // Sets the element at the given index. If the value is a scalar, sets its
+  // value.
+  void InsertElement(llvm::ArrayRef<int64_t> indices,
+                     const InterpreterValue& value);
+  // Initializes all elements of the underlying tensor.
+  void Fill(
+      const std::function<InterpreterValue(llvm::ArrayRef<int64_t> indices)>&
+          f);
+
+  // Converts a scalar to a unit tensor or vector.
+  InterpreterValue AsUnitTensor(bool is_vector = false) const;
+  // For integral interpreter values, casts them to int64.
+  int64_t AsInt() const;
+  // For integral interpreter values, first casts them to the unsigned integer
+  // type of the same size, and then to uint64. For example, the result for
+  // int8_t{-1} is 255.
+  uint64_t AsUInt() const;
+  // For floating point scalars, casts them to double.
+  double AsDouble() const;
+  // Must be a tensor or memref.
+  int64_t GetByteSizeOfElement() const;
+
+  // Creates a new tensor InterpreterValue (backed a new buffer) with the same
+  // element type as this, but a different shape. If this is not a tensor, it is
+  // used as the element type.
+  // If `layout` is empty, the clone uses the default layout.
+  InterpreterValue Clone(ArrayRef<int64_t> layout = {}) const;
+  // Returns either this tensor InterpreterValue (if its layout matches the
+  // requested layout) or a clone.
+  InterpreterValue CoerceLayout(llvm::ArrayRef<int64_t> layout) const;
+  // Returns a tensor interpreter value with a newly allocated buffer of the
+  // given shape, with a default layout and the same element type as this
+  // interpreter value.
+  InterpreterValue TypedAlike(llvm::ArrayRef<int64_t> shape) const;
+
+  // Creates a tensor with the given element type and shape. `element_type` may
+  // be a vector type, in which case the shape only specifies the non-vector
+  // dimensions.
+  static InterpreterValue MakeTensor(mlir::Type elementType,
+                                     SmallVector<int64_t> shape);
+
+  // Returns the underlying tensor's view. Must be a tensor.
+  BufferView& View();
+  const BufferView& View() const;
+  // Returns the underlying tensor's buffer. Must be a tensor.
+  std::shared_ptr<Buffer> Buffer() const;
+
+  bool IsTensor() const;
+
+  bool operator==(const InterpreterValue& other) const {
+    if (storage.index() != other.storage.index()) return false;
+    if (IsTensor() || std::holds_alternative<Tuple>(storage))
+      return storage == other.storage;
+    // Tensors treat NaNs as equal, so just wrap the values.
+    return AsUnitTensor() == other.AsUnitTensor();
+  }
+
+  std::variant<
+      Tuple, bool, float, double, uint8_t, int8_t, uint16_t, int16_t, uint32_t,
+      int32_t, uint64_t, int64_t, std::complex<float>, std::complex<double>,
+      TensorOrMemref<bool>, TensorOrMemref<float>, TensorOrMemref<double>,
+      TensorOrMemref<uint8_t>, TensorOrMemref<int8_t>, TensorOrMemref<uint16_t>,
+      TensorOrMemref<int16_t>, TensorOrMemref<uint32_t>,
+      TensorOrMemref<int32_t>, TensorOrMemref<uint64_t>,
+      TensorOrMemref<int64_t>, TensorOrMemref<std::complex<float>>,
+      TensorOrMemref<std::complex<double>>>
+      storage;
+};
+
+template <typename T>
+constexpr static bool is_valid_interpreter_value_v =  // NOLINT
+    std::is_constructible_v<decltype(InterpreterValue::storage), T>;
+
+// Attempts to cast the given value to the requested type, returning nullopt if
+// no cast is possible. This allows casts to the concrete type of the value
+// (e.g. an `InterpreterValue` containing a `Tuple` can be cast to `Tuple`),
+// casts from a unit tensor to their contents, and casts of scalars to any
+// convertible type.
+// NOTE: When casting to an unsigned type, this behaves differently than
+// InterpreterValue::AsUint. That function preserves the content's bit width,
+// so InterpreterValueDynCast<uint64_t>({int8_t{-1}}) will return 2^64-1,
+// whereas asUInt will return 255.
+template <typename T>
+std::optional<T> InterpreterValueDynCast(InterpreterValue v) {
+  if constexpr (std::is_same_v<T, InterpreterValue>) {
+    return v;
+  }
+  if constexpr (is_valid_interpreter_value_v<T>) {
+    if (std::holds_alternative<T>(v.storage)) {
+      return std::get<T>(v.storage);
+    }
+  }
+  if (v.IsTensor() && !is_tensor_or_memref_v<T>) {
+    if (v.View().GetNumElements() != 1) {
+      return std::nullopt;
+    }
+    return InterpreterValueDynCast<T>(v.ExtractElement({}));
+  }
+  return std::visit(
+      [](auto v) -> std::optional<T> {
+        if constexpr (std::is_convertible_v<decltype(v), T>) {
+          return v;
+        } else {
+          return std::nullopt;
+        }
+      },
+      v.storage);
+}
+
+template <typename T>
+T InterpreterValueCast(InterpreterValue v) {
+  auto ret = InterpreterValueDynCast<T>(v);
+  assert(ret && "cast failed");
+  return *std::move(ret);
+}
+
+// Calls functor with a value of the C++ type corresponding to the given `Type`,
+// (or its element type).
+template <class Fn>
+auto DispatchScalarType(mlir::Type ty, Fn&& functor) {
+  ty = getElementTypeOrSelf(ty);
+  if (ty.isF32()) {
+    return functor(float{});
+  }
+  if (ty.isF64()) {
+    return functor(double{});
+  }
+  if (ty.isUnsignedInteger(64)) {
+    return functor(uint64_t{});
+  }
+  if (ty.isInteger(64) || ty.isIndex()) {
+    return functor(int64_t{});
+  }
+  if (ty.isUnsignedInteger(32)) {
+    return functor(uint32_t{});
+  }
+  if (ty.isInteger(32)) {
+    return functor(int32_t{});
+  }
+  if (ty.isUnsignedInteger(16)) {
+    return functor(uint16_t{});
+  }
+  if (ty.isInteger(16)) {
+    return functor(int16_t{});
+  }
+  if (ty.isUnsignedInteger(8)) {
+    return functor(uint8_t{});
+  }
+  if (ty.isInteger(8)) {
+    return functor(int8_t{});
+  }
+  if (ty.isInteger(1)) {
+    return functor(bool{});
+  }
+  if (auto complex = mlir::dyn_cast<ComplexType>(ty)) {
+    if (complex.getElementType().isF32()) {
+      return functor(std::complex<float>{});
+    }
+    if (complex.getElementType().isF64()) {
+      return functor(std::complex<double>{});
+    }
+  }
+
+  llvm::errs() << "DispatchScalarType unimplemented for " << ty << "\n";
+  llvm_unreachable("unimplemented");
+}
+
+}  // namespace interpreter
+}  // namespace mlir
+
+#endif  // XLA_MLIR_TOOLS_MLIR_INTERPRETER_FRAMEWORK_INTERPRETER_VALUE_H_
diff --git a/third_party/xla/xla/mlir/tools/mlir_interpreter/framework/interpreter_value_util.h b/third_party/xla/xla/mlir/tools/mlir_interpreter/framework/interpreter_value_util.h
new file mode 100644
index 00000000000000..5e9a6ba2903359
--- /dev/null
+++ b/third_party/xla/xla/mlir/tools/mlir_interpreter/framework/interpreter_value_util.h
@@ -0,0 +1,190 @@
+/* Copyright 2022 The OpenXLA Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+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.
+==============================================================================*/
+
+#ifndef XLA_MLIR_TOOLS_MLIR_INTERPRETER_FRAMEWORK_INTERPRETER_VALUE_UTIL_H_
+#define XLA_MLIR_TOOLS_MLIR_INTERPRETER_FRAMEWORK_INTERPRETER_VALUE_UTIL_H_
+
+#include <cassert>
+#include <complex>
+#include <type_traits>
+#include <utility>
+
+#include "llvm/ADT/STLExtras.h"
+#include "llvm/Support/ErrorHandling.h"
+#include "llvm/Support/raw_ostream.h"
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+#include "xla/mlir/tools/mlir_interpreter/framework/interpreter_value.h"
+#include "xla/mlir/tools/mlir_interpreter/framework/tensor_or_memref.h"
+
+namespace mlir {
+namespace interpreter {
+namespace detail {
+
+template <typename T>
+struct is_complex : std::false_type {};  // NOLINT
+
+template <typename T>
+struct is_complex<std::complex<T>> : std::true_type {};  // NOLINT
+
+template <typename Fn>
+struct InterpreterValueMapVisitor {
+  template <typename T>
+  InterpreterValue operator()(const TensorOrMemref<T>& t) {
+    if constexpr (Fn::template SupportedType<T>()) {
+      using out_elem_t = decltype(Fn::Apply(T()));
+      auto out = TensorOrMemref<out_elem_t>::EmptyLike(t.view);
+      for (const auto& index : out.view.Indices(true)) {
+        out.at(index) = Fn::Apply(t.at(index));
+      }
+      return {out};
+    } else {
+      llvm::errs() << llvm::getTypeName<Fn>()
+                   << " unsupported type: " << llvm::getTypeName<T>() << "\n";
+      llvm_unreachable("unsupported type");
+    }
+  }
+
+  InterpreterValue operator()(const Tuple& t) {
+    Tuple out;
+    for (const auto& value : t.values) {
+      out.values.push_back(std::make_unique<InterpreterValue>(
+          std::move(std::visit(*this, value->storage))));
+    }
+    return {out};
+  }
+
+  template <typename T>
+  InterpreterValue operator()(const T& t) {
+    if constexpr (Fn::template SupportedType<T>()) {
+      return {Fn::Apply(t)};
+    } else {
+      llvm::errs() << llvm::getTypeName<Fn>()
+                   << " unsupported type: " << llvm::getTypeName<T>() << "\n";
+      llvm_unreachable("unsupported type");
+    }
+  }
+};
+
+template <typename Fn>
+struct InterpreterValueBiMapVisitor {
+  const InterpreterValue& rhs;
+
+  template <typename T>
+  InterpreterValue operator()(const TensorOrMemref<T>& lhsT) {
+    if constexpr (Fn::template SupportedType<T>()) {
+      using OutElemT = decltype(Fn::Apply(T(), T()));
+      auto out = TensorOrMemref<OutElemT>::EmptyLike(lhsT.view);
+      const auto& rhsT = std::get<TensorOrMemref<T>>(rhs.storage);
+      for (const auto& index : out.view.Indices(true)) {
+        out.at(index) = Fn::Apply(lhsT.at(index), rhsT.at(index));
+      }
+      return {out};
+    } else {
+      llvm::errs() << llvm::getTypeName<Fn>()
+                   << " unsupported type: " << llvm::getTypeName<T>() << "\n";
+      llvm_unreachable("unsupported type");
+    }
+  }
+
+  InterpreterValue operator()(const Tuple& lhsT) {
+    const auto& rhsT = std::get<Tuple>(rhs.storage);
+    Tuple out;
+    for (const auto& [lhs_v, rhs_v] : llvm::zip(lhsT.values, rhsT.values)) {
+      out.values.push_back(std::make_unique<InterpreterValue>(std::move(
+          std::visit(InterpreterValueBiMapVisitor{*rhs_v}, lhs_v->storage))));
+    }
+    return {std::move(out)};
+  }
+
+  template <typename T>
+  InterpreterValue operator()(const T& t) {
+    if constexpr (Fn::template SupportedType<T>()) {
+      return {Fn::Apply(t, std::get<T>(rhs.storage))};
+    } else {
+      llvm::errs() << llvm::getTypeName<Fn>()
+                   << " unsupported type: " << llvm::getTypeName<T>() << "\n";
+      llvm_unreachable("unsupported type");
+    }
+  }
+};
+
+}  // namespace detail
+
+template <typename T>
+inline constexpr bool is_complex_v = detail::is_complex<T>::value;  // NOLINT
+
+template <bool allow_bools, bool allow_ints, bool allow_floats,
+          bool allow_complex, bool allow_unsigned = true>
+struct FilterMapTraits {
+  template <typename T>
+  static constexpr bool SupportedType() {
+    constexpr bool isBool = std::is_same_v<T, bool>;
+    constexpr bool isInt = std::is_integral_v<T> && !isBool;
+    constexpr bool isUnsigned = std::is_unsigned_v<T>;
+    return (allow_bools && isBool) ||
+           (allow_ints && isInt && (allow_unsigned || !isUnsigned)) ||
+           (allow_floats && std::is_floating_point_v<T>) ||
+           (allow_complex && is_complex_v<T>);
+  }
+};
+
+using CwiseAll = FilterMapTraits<true, true, true, true>;
+using CwiseArith = FilterMapTraits<false, true, true, true>;
+using CwiseComplex = FilterMapTraits<false, false, false, true>;
+using CwiseFloat = FilterMapTraits<false, false, true, false>;
+using CwiseInt = FilterMapTraits<false, true, false, false>;
+using CwiseIntegral = FilterMapTraits<true, true, false, false>;
+using CwiseNonIntegral = FilterMapTraits<false, false, true, true>;
+using CwiseReal = FilterMapTraits<false, true, true, false>;
+using CwiseSignedOrComplex = FilterMapTraits<false, true, true, true, false>;
+using CwiseSigned = FilterMapTraits<false, true, true, false, false>;
+
+template <typename Fn>
+InterpreterValue ApplyCwiseMap(const InterpreterValue& value) {
+  return std::visit(detail::InterpreterValueMapVisitor<Fn>{}, value.storage);
+}
+
+template <typename Fn>
+InterpreterValue ApplyCwiseBinaryMap(const InterpreterValue& lhs,
+                                     const InterpreterValue& rhs) {
+  assert(lhs.storage.index() == rhs.storage.index());
+  return std::visit(detail::InterpreterValueBiMapVisitor<Fn>{rhs}, lhs.storage);
+}
+
+// Unboxes (and casts if necessary) the given interpreter values. Asserts if the
+// types are incompatible.
+template <typename T>
+SmallVector<T> UnpackInterpreterValues(ArrayRef<InterpreterValue> values) {
+  SmallVector<T> result;
+  for (const auto& value : values) {
+    result.push_back(InterpreterValueCast<T>(value));
+  }
+  return result;
+}
+
+// Boxes the given values in InterpreterValues.
+template <typename T>
+SmallVector<InterpreterValue> PackInterpreterValues(ArrayRef<T> values) {
+  SmallVector<InterpreterValue> result;
+  for (const auto& value : values) {
+    result.push_back({value});
+  }
+  return result;
+}
+
+}  // namespace interpreter
+}  // namespace mlir
+
+#endif  // XLA_MLIR_TOOLS_MLIR_INTERPRETER_FRAMEWORK_INTERPRETER_VALUE_UTIL_H_
diff --git a/third_party/xla/xla/mlir/tools/mlir_interpreter/framework/registration.cc b/third_party/xla/xla/mlir/tools/mlir_interpreter/framework/registration.cc
new file mode 100644
index 00000000000000..454f08641e19e6
--- /dev/null
+++ b/third_party/xla/xla/mlir/tools/mlir_interpreter/framework/registration.cc
@@ -0,0 +1,122 @@
+/* Copyright 2023 The OpenXLA Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+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.
+==============================================================================*/
+
+#include "xla/mlir/tools/mlir_interpreter/framework/registration.h"
+
+#include <functional>
+#include <utility>
+
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+
+namespace mlir {
+namespace interpreter {
+namespace detail {
+namespace {
+
+// Aliases and function names are wrapped in functions because function
+// registrations are called from static initializers, whose execution order is
+// undefined.
+DenseMap<llvm::StringRef, llvm::StringRef>& GetOpAliases() {
+  static DenseMap<llvm::StringRef, llvm::StringRef>* aliases = nullptr;
+  if (!aliases) {
+    aliases = new DenseMap<llvm::StringRef, llvm::StringRef>();
+  }
+  return *aliases;
+}
+
+DenseMap<llvm::StringRef, InterpreterFunction>& GetFunctions() {
+  static DenseMap<llvm::StringRef, InterpreterFunction>* functions = nullptr;
+  if (!functions) {
+    functions = new DenseMap<llvm::StringRef, InterpreterFunction>();
+  }
+  return *functions;
+}
+
+}  // namespace
+
+InterpreterFunction GetFunction(llvm::StringRef name) {
+  const auto& fns = GetFunctions();
+  auto fn = fns.find(name);
+  if (fn != fns.end()) {
+    return fn->second;
+  }
+  const auto& aliases = GetOpAliases();
+  auto alias = aliases.find(name);
+  if (alias != aliases.end()) {
+    return fns.find(alias->second)->second;
+  }
+  return nullptr;
+}
+
+void RegisterInterpreterOp(llvm::StringRef name,
+                           InterpreterValue (*fn)(const InterpreterValue&)) {
+  RegisterInterpreterOp(
+      name,
+      [fn](MutableArrayRef<InterpreterValue> operands, mlir::Operation*,
+           InterpreterState&) -> SmallVector<InterpreterValue> {
+        assert(operands.size() == 1 && "unexpected number of operands");
+        return {fn(operands[0])};
+      });
+}
+
+void RegisterInterpreterOp(llvm::StringRef name,
+                           InterpreterValue (*fn)(const InterpreterValue&,
+                                                  const InterpreterValue&)) {
+  RegisterInterpreterOp(
+      name,
+      [fn](MutableArrayRef<InterpreterValue> operands, mlir::Operation*,
+           InterpreterState&) -> SmallVector<InterpreterValue> {
+        assert(operands.size() == 2 && "unexpected number of operands");
+        return {fn(operands[0], operands[1])};
+      });
+}
+
+void RegisterInterpreterOp(
+    llvm::StringRef name,
+    InterpreterValue (*fn)(MutableArrayRef<InterpreterValue>)) {
+  RegisterInterpreterOp(
+      name,
+      [fn](MutableArrayRef<InterpreterValue> operands, mlir::Operation*,
+           InterpreterState&) -> SmallVector<InterpreterValue> {
+        return {fn(operands)};
+      });
+}
+
+void RegisterInterpreterOp(llvm::StringRef name,
+                           void (*fn)(MutableArrayRef<InterpreterValue>)) {
+  RegisterInterpreterOp(
+      name,
+      [fn](MutableArrayRef<InterpreterValue> operands, mlir::Operation*,
+           InterpreterState&) -> SmallVector<InterpreterValue> {
+        fn(operands);
+        return {};
+      });
+}
+
+void RegisterInterpreterOp(
+    llvm::StringRef name,
+    std::function<llvm::SmallVector<InterpreterValue>(
+        MutableArrayRef<InterpreterValue>, mlir::Operation*, InterpreterState&)>
+        fn) {
+  GetFunctions()[name] = std::move(fn);
+}
+
+void RegisterInterpreterOp(llvm::StringRef name, llvm::StringRef original) {
+  GetOpAliases()[name] = original;
+}
+
+}  // namespace detail
+}  // namespace interpreter
+}  // namespace mlir
diff --git a/third_party/xla/xla/mlir/tools/mlir_interpreter/framework/registration.h b/third_party/xla/xla/mlir/tools/mlir_interpreter/framework/registration.h
new file mode 100644
index 00000000000000..75de217bddfb40
--- /dev/null
+++ b/third_party/xla/xla/mlir/tools/mlir_interpreter/framework/registration.h
@@ -0,0 +1,229 @@
+/* Copyright 2023 The OpenXLA Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+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.
+==============================================================================*/
+#ifndef XLA_MLIR_TOOLS_MLIR_INTERPRETER_FRAMEWORK_REGISTRATION_H_
+#define XLA_MLIR_TOOLS_MLIR_INTERPRETER_FRAMEWORK_REGISTRATION_H_
+
+#include <cstddef>
+#include <cstdint>
+#include <functional>
+#include <optional>
+#include <type_traits>
+#include <utility>
+
+#include "absl/strings/str_cat.h"
+#include "llvm/ADT/Sequence.h"
+#include "llvm/ADT/SmallVector.h"
+#include "llvm/ADT/StringRef.h"
+#include "llvm/Support/TypeName.h"
+#include "mlir/IR/Operation.h"  // from @llvm-project
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+#include "xla/mlir/tools/mlir_interpreter/framework/interpreter.h"
+#include "xla/mlir/tools/mlir_interpreter/framework/interpreter_value.h"
+#include "xla/mlir/tools/mlir_interpreter/framework/interpreter_value_util.h"
+
+#define MLIR_INTERPRETER_CONCAT_IMPL(x, y) x##y
+#define MLIR_INTERPRETER_CONCAT(x, y) MLIR_INTERPRETER_CONCAT_IMPL(x, y)
+#define REGISTER_MLIR_INTERPRETER_OP(args...)                     \
+  static int MLIR_INTERPRETER_CONCAT(init_, __COUNTER__) = []() { \
+    ::mlir::interpreter::detail::RegisterInterpreterOp(args);     \
+    return 1;                                                     \
+  }();
+
+namespace mlir {
+namespace interpreter {
+namespace detail {
+
+// The generic signature for interpreter functions. Typically the type-checked
+// form should be used instead.
+using InterpreterFunction = std::function<SmallVector<InterpreterValue>(
+    MutableArrayRef<InterpreterValue>, mlir::Operation*, InterpreterState&)>;
+
+// Returns the given registered function, or nullptr if not found.
+InterpreterFunction GetFunction(llvm::StringRef name);
+
+// Simple unary ops.
+void RegisterInterpreterOp(llvm::StringRef name,
+                           InterpreterValue (*fn)(const InterpreterValue&));
+
+// Simple binary ops.
+void RegisterInterpreterOp(llvm::StringRef name,
+                           InterpreterValue (*fn)(const InterpreterValue&,
+                                                  const InterpreterValue&));
+
+template <typename T>
+struct is_optional : std::false_type {};  // NOLINT
+
+template <typename T>
+struct is_optional<std::optional<T>> : std::true_type {};  // NOLINT
+
+// Converts the given arguments to the requested type. Supported target types:
+// - InterpreterValue storage types (e.g. uint8_t, TensorOrMemref<float>).
+//   Scalars will be cast if necessary.
+// - std::optional of InterpreterValue storage types
+// - ArrayRef of InterpreterValue storage types. Will cast if necessary (e.g.
+//   int32_t -> int64_t).
+// - no-op conversions: InterpreterValue, ArrayRef<InterpreterValue>
+template <typename ArgT>
+auto TypedInterpreterOpConvertArg(MutableArrayRef<InterpreterValue> args,
+                                  InterpreterState& state) {
+  constexpr bool optional = is_optional<ArgT>::value;
+  constexpr bool single = std::is_same_v<ArgT, InterpreterValue> ||
+                          is_valid_interpreter_value_v<ArgT>;
+  if constexpr (optional) {
+    if (args.empty()) {
+      return ArgT{};
+    }
+  }
+
+  if constexpr (single || optional) {
+    if (args.size() != 1) {
+      state.AddFailure("Expected a single argument for the operand");
+      return ArgT{};
+    }
+  }
+
+  auto fail = [&]() {
+    state.AddFailure(absl::StrCat("Unable to convert argument (variant index ",
+                                  args[0].storage.index(), ") to ",
+                                  llvm::getTypeName<ArgT>().str()));
+    return ArgT{};
+  };
+
+  if constexpr (single) {
+    if (auto arg = InterpreterValueDynCast<ArgT>(args[0])) {
+      return ArgT{*arg};
+    }
+    return fail();
+  } else if constexpr (optional) {
+    using T = std::decay_t<decltype(*std::declval<ArgT>())>;
+    if (auto arg = InterpreterValueDynCast<T>(args[0])) {
+      return arg;
+    }
+    return fail();
+  } else {
+    using E = std::decay_t<decltype(*std::declval<ArgT>().begin())>;
+    // Container argument types (e.g. MutableArrayRef<InterpreterValue>,
+    // ArrayRef<int64_t>).
+    if constexpr (std::is_same_v<E, InterpreterValue>) {
+      return ArgT{args};
+    } else {
+      // Note: we don't cast to ArgT here, because that's typically an ArrayRef,
+      // which would lead to returning a reference to a temporary.
+      return UnpackInterpreterValues<E>(args);
+    }
+  }
+}
+
+// Converts the given return value. Supported target types:
+// - InterpreterValue (no-op conversion)
+// - InterpreterValue storage types (the value will be boxed)
+// - SmallVector<InterpreterValue>
+template <typename RetT>
+SmallVector<InterpreterValue> TypedInterpreterOpConvertRet(RetT ret) {
+  if constexpr (std::is_same_v<RetT, InterpreterValue>) {
+    return {ret};
+  } else if constexpr (is_valid_interpreter_value_v<RetT>) {
+    return {InterpreterValue{ret}};
+  } else if constexpr (std::is_same_v<RetT, SmallVector<InterpreterValue>>) {
+    return ret;
+  } else {
+    using E = std::decay_t<decltype(*std::declval<RetT>().begin())>;
+    return PackInterpreterValues(ArrayRef<E>(ret));
+  }
+}
+
+// Adapts the given function to the generic handler signature
+// (SmallVector<InterpreterValue>(MutableArrayRef<InterpreterValue>, Operation*,
+// InterpreterState&)).
+// See the function below for usage.
+template <typename Op, typename Ret, typename... T, size_t... Indices>
+void RegisterTypedInterpreterOpImpl(Ret (*fn)(InterpreterState&, Op, T... args),
+                                    std::index_sequence<Indices...>) {
+  RegisterInterpreterOp(
+      Op::getOperationName(),
+      [fn](MutableArrayRef<InterpreterValue> args, mlir::Operation* op,
+           InterpreterState& state) -> SmallVector<InterpreterValue> {
+        auto cast = llvm::dyn_cast<Op>(op);
+        if (!cast) {
+          state.AddFailure(absl::StrCat(
+              "failed to cast op '", op->getName().getStringRef().str(),
+              "' to expected type (", llvm::getTypeName<Op>().str(), ")"));
+          return {};
+        }
+        int64_t used_args = 0;
+        for (auto i : llvm::seq(0ul, sizeof...(T))) {
+          used_args += cast.getODSOperandIndexAndLength(i).second;
+        }
+        if (args.size() != used_args) {
+          state.AddFailure("Op handler did not use all arguments");
+          return {};
+        }
+
+        auto extract_arg = [&](auto index, auto* dummy) {
+          auto [pos, length] = cast.getODSOperandIndexAndLength(index);
+          using ArgT = std::decay_t<decltype(*dummy)>;
+          return TypedInterpreterOpConvertArg<ArgT>(args.slice(pos, length),
+                                                    state);
+        };
+
+        if constexpr (std::is_same_v<Ret, void>) {
+          fn(state, cast, extract_arg(Indices, (std::decay_t<T>*){nullptr})...);
+          return {};
+        } else {
+          Ret ret = fn(state, cast,
+                       extract_arg(Indices, (std::decay_t<T>*){nullptr})...);
+          return TypedInterpreterOpConvertRet(ret);
+        }
+      });
+}
+
+// registers the given function. The function should take one argument per
+// Op operand, in the same order as the Op.
+// The argument types should match the operation's operand types:
+// - Variadic<...> becomes ArrayRef<...>
+// - Optional<...> becomes std::optional<...>
+// - Unboxing is optionally supported, e.g. an Optional<Index> operand can be
+//   passed to either a std::optional<int64_t> or a
+//   std::optional<InterpreterValue>.
+// Valid return types are InterpreterValue, SmallVector<InterpreterValue>, void,
+// and any type boxable in an InterpreterValue.
+template <typename Op, typename Ret, typename... T>
+void RegisterInterpreterOp(Ret (*fn)(InterpreterState&, Op, T...)) {
+  RegisterTypedInterpreterOpImpl(fn, std::index_sequence_for<T...>{});
+}
+
+// Simple variadic ops (single output).
+void RegisterInterpreterOp(
+    llvm::StringRef name,
+    InterpreterValue (*fn)(MutableArrayRef<InterpreterValue>));
+
+// Simple variadic ops(no output).
+void RegisterInterpreterOp(llvm::StringRef name,
+                           void (*fn)(MutableArrayRef<InterpreterValue>));
+
+// Generic ops.
+void RegisterInterpreterOp(
+    llvm::StringRef name,
+    std::function<llvm::SmallVector<InterpreterValue>(
+        MutableArrayRef<InterpreterValue>, mlir::Operation*, InterpreterState&)>
+        fn);
+
+void RegisterInterpreterOp(llvm::StringRef name, llvm::StringRef original);
+
+}  // namespace detail
+}  // namespace interpreter
+}  // namespace mlir
+
+#endif  // XLA_MLIR_TOOLS_MLIR_INTERPRETER_FRAMEWORK_REGISTRATION_H_
diff --git a/third_party/xla/xla/mlir/tools/mlir_interpreter/framework/tensor_or_memref.cc b/third_party/xla/xla/mlir/tools/mlir_interpreter/framework/tensor_or_memref.cc
new file mode 100644
index 00000000000000..7226ff73b9b14f
--- /dev/null
+++ b/third_party/xla/xla/mlir/tools/mlir_interpreter/framework/tensor_or_memref.cc
@@ -0,0 +1,162 @@
+/* Copyright 2022 The OpenXLA Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+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.
+==============================================================================*/
+
+#include "xla/mlir/tools/mlir_interpreter/framework/tensor_or_memref.h"
+
+#include <algorithm>
+#include <cstdint>
+#include <iostream>
+#include <numeric>
+#include <optional>
+#include <utility>
+
+#include "mlir/Dialect/Utils/IndexingUtils.h"
+#include "llvm/ADT/STLExtras.h"
+#include "llvm/ADT/SmallVector.h"
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+#include "mlir/Support/LogicalResult.h"  // from @llvm-project
+
+namespace mlir {
+namespace interpreter {
+
+std::optional<int64_t> BufferView::GetPhysicalIndex(
+    llvm::ArrayRef<int64_t> view_indices) const {
+  int64_t result = offset;
+  if (!InBounds(view_indices)) {
+    return std::nullopt;
+  }
+  for (int64_t i = 0; i < view_indices.size(); ++i) {
+    result += view_indices[i] * strides[i];
+  }
+  return result;
+}
+
+bool BufferView::InBounds(llvm::ArrayRef<int64_t> view_indices) const {
+  if (view_indices.size() > sizes.size()) return false;
+  for (auto [index, size] : llvm::zip(view_indices, sizes)) {
+    if (index < 0 || index >= size) return false;
+  }
+  return true;
+}
+
+SmallVector<int64_t> BufferView::GetDefaultStrides(ArrayRef<int64_t> sizes) {
+  SmallVector<int64_t> result(sizes.size());
+  int64_t stride = 1;
+  for (int64_t i = result.size() - 1; i >= 0; --i) {
+    result[i] = stride;
+    stride *= sizes[i];
+  }
+  return result;
+}
+
+SmallVector<int64_t> BufferView::GetStridesForLayout(ArrayRef<int64_t> sizes,
+                                                     ArrayRef<int64_t> layout) {
+  if (layout.empty()) return GetDefaultStrides(sizes);
+  auto inverseLayout = invertPermutationVector(layout);
+  SmallVector<int64_t> result(sizes.size());
+  int64_t stride = 1;
+  for (int64_t i = 0; i < layout.size(); ++i) {
+    result[inverseLayout[i]] = stride;
+    stride *= sizes[inverseLayout[i]];
+  }
+  return result;
+}
+
+LogicalResult BufferView::Slice(int64_t dim_index, int64_t dim_offset) {
+  llvm::SmallVector<int64_t> offsets(Rank(), 0);
+  offsets[dim_index] = dim_offset;
+  if (auto newOffset = GetPhysicalIndex(offsets)) {
+    offset = *newOffset;
+  } else {
+    return failure();
+  }
+  if (dim_index >= Rank()) --*num_vector_dims;
+  strides.erase(strides.begin() + dim_index);
+  sizes.erase(sizes.begin() + dim_index);
+  return success();
+}
+
+LogicalResult BufferView::Slice(int64_t dim_index, int64_t dim_offset,
+                                int64_t dim_size, int64_t dim_stride) {
+  llvm::SmallVector<int64_t> offsets(Rank(), 0);
+  offsets[dim_index] = dim_offset;
+  if (dim_size == 0) {
+    offset = 0;
+  } else if (auto newOffset = GetPhysicalIndex(offsets)) {
+    offset = *newOffset;
+  } else {
+    return failure();
+  }
+  sizes[dim_index] = dim_size;
+  strides[dim_index] *= dim_stride;
+  return success();
+}
+
+LogicalResult BufferView::Subview(ArrayRef<int64_t> subview_offsets,
+                                  ArrayRef<int64_t> subview_sizes,
+                                  ArrayRef<int64_t> subview_strides) {
+  if (auto new_offset = GetPhysicalIndex(subview_offsets)) {
+    offset = *new_offset;
+  } else {
+    return failure();
+  }
+
+  for (auto [in_size, subview_offset, subview_size, subview_stride] :
+       llvm::zip(sizes, subview_offsets, subview_sizes, subview_strides)) {
+    int64_t limitIndex = subview_offset + (subview_size - 1) * subview_stride;
+    if (subview_offset < 0 || subview_offset >= in_size || limitIndex < 0 ||
+        limitIndex >= in_size) {
+      return failure();
+    }
+  }
+
+  for (auto [in_stride, subview_stride] : llvm::zip(strides, subview_strides)) {
+    in_stride *= subview_stride;
+  }
+  sizes = llvm::to_vector(subview_sizes);
+  return success();
+}
+
+int64_t BufferView::GetNumElements(bool includeVectorDims) const {
+  size_t n = 1;
+  for (auto size : ArrayRef<int64_t>(sizes).drop_back(
+           includeVectorDims ? 0 : num_vector_dims.value_or(0)))
+    n *= size;
+  return n;
+}
+
+std::optional<int64_t> BufferView::GetCollapsedStride(
+    llvm::ArrayRef<int64_t> dims) const {
+  using StrideAndDim = std::pair<int64_t, int64_t>;
+  llvm::SmallVector<StrideAndDim> stridesAndDims;
+  for (auto dim : dims) {
+    if (sizes[dim] != 1) {
+      stridesAndDims.emplace_back(strides[dim], dim);
+    }
+  }
+
+  if (stridesAndDims.empty()) return 0;
+
+  llvm::sort(stridesAndDims);
+  int64_t nextStride = stridesAndDims.front().first;
+  for (auto [stride, dim] : stridesAndDims) {
+    if (stride != nextStride) return std::nullopt;
+    nextStride *= sizes[dim];
+  }
+  return stridesAndDims.front().first;
+}
+
+}  // namespace interpreter
+}  // namespace mlir
diff --git a/third_party/xla/xla/mlir/tools/mlir_interpreter/framework/tensor_or_memref.h b/third_party/xla/xla/mlir/tools/mlir_interpreter/framework/tensor_or_memref.h
new file mode 100644
index 00000000000000..6f801021a68101
--- /dev/null
+++ b/third_party/xla/xla/mlir/tools/mlir_interpreter/framework/tensor_or_memref.h
@@ -0,0 +1,364 @@
+/* Copyright 2022 The OpenXLA Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+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.
+==============================================================================*/
+
+#ifndef XLA_MLIR_TOOLS_MLIR_INTERPRETER_FRAMEWORK_TENSOR_OR_MEMREF_H_
+#define XLA_MLIR_TOOLS_MLIR_INTERPRETER_FRAMEWORK_TENSOR_OR_MEMREF_H_
+
+#include <math.h>
+
+#include <cmath>
+#include <complex>
+#include <cstddef>
+#include <cstdint>
+#include <iterator>
+#include <memory>
+#include <optional>
+#include <string>
+#include <type_traits>
+#include <utility>
+
+#include "llvm/ADT/ArrayRef.h"
+#include "llvm/ADT/StringRef.h"
+#include "llvm/Support/raw_ostream.h"
+#include "mlir/IR/Operation.h"  // from @llvm-project
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+#include "mlir/Support/LogicalResult.h"  // from @llvm-project
+
+namespace mlir {
+namespace interpreter {
+
+template <typename T>
+bool IsEqual(T a, T b) {
+  return a == b;
+}
+
+// TODO(jreiffers): Replace ifndef with a command line flag.
+#ifndef MLIR_INTERPRETER_COMPARE_DOUBLES_EXACT
+// Compare double precision float with a small tolerance, because complex
+// computations in the interpreter don't always produce the exact same result.
+template <>
+inline bool IsEqual(double a, double b) {
+  if (isinf(a) || isinf(b)) {
+    return a == b;
+  }
+
+  return fabs(a - b) < 1e-14;
+}
+
+template <>
+inline bool IsEqual(std::complex<double> a, std::complex<double> b) {
+  return IsEqual(a.real(), b.real()) && IsEqual(a.imag(), b.imag());
+}
+#endif
+
+// Represents a view into a physical buffer.
+struct BufferView {
+  int64_t offset;
+  llvm::SmallVector<int64_t> sizes;    // [10, 11, 12]
+  llvm::SmallVector<int64_t> strides;  // [132, 12, 1]
+  // Number of vector element dimensions in the tensor. nullopt if this is a
+  // vector itself (is_vector is set). {0} if this is a tensor of a unit vector.
+  std::optional<int64_t> num_vector_dims = std::nullopt;
+  bool is_vector = false;
+
+  int64_t Rank() const { return sizes.size() - num_vector_dims.value_or(0); }
+
+  // Removes the dimension from the view. If you need to keep it, use the
+  // overload below with dimSize = 1.
+  LogicalResult Slice(int64_t dim_index, int64_t dim_offset);
+  LogicalResult Slice(int64_t dim_index, int64_t dim_offset, int64_t dim_size,
+                      int64_t dim_stride = 1);
+  LogicalResult Subview(ArrayRef<int64_t> subview_offsets,
+                        ArrayRef<int64_t> subview_sizes,
+                        ArrayRef<int64_t> subview_strides);
+  int64_t GetNumElements(bool include_vector_dims = false) const;
+
+  class LogicalIndexView {
+   public:
+    class Iterator {
+     public:
+      using iterator_category = std::forward_iterator_tag;
+      using value_type = llvm::SmallVector<int64_t>;
+      using difference_type = std::ptrdiff_t;
+      using pointer = llvm::SmallVector<int64_t>*;
+      using reference = llvm::SmallVector<int64_t>&;
+
+      const llvm::SmallVector<int64_t>& operator*() const {
+        return view_indices_;
+      }
+      const llvm::SmallVector<int64_t>* operator->() const {
+        return &view_indices_;
+      }
+
+      Iterator& operator++() {
+        auto indexIt = view_indices_.rbegin();
+        auto sizeIt = view_->sizes.rbegin();
+        if (!include_vector_dims_) {
+          std::advance(sizeIt, view_->num_vector_dims.value_or(0));
+        }
+
+        for (auto e = view_indices_.rend(); indexIt != e; ++indexIt, ++sizeIt) {
+          ++*indexIt;
+          if (*indexIt < *sizeIt) {
+            return *this;
+          }
+          *indexIt = 0;
+        }
+
+        view_indices_.clear();
+        view_indices_.push_back(-1);
+        return *this;
+      }
+
+      Iterator operator++(int) {
+        auto tmp = *this;
+        ++(*this);
+        return tmp;
+      }
+
+      bool operator==(const Iterator& other) const {
+        return view_indices_ == other.view_indices_;
+      }
+
+      bool operator!=(const Iterator& other) const { return !(*this == other); }
+
+     private:
+      friend class LogicalIndexView;
+
+      Iterator(const BufferView* view, llvm::SmallVector<int64_t> indices,
+               bool include_vector_dims)
+          : view_(view),
+            view_indices_(std::move(indices)),
+            include_vector_dims_(include_vector_dims) {}
+
+      const BufferView* view_;
+      llvm::SmallVector<int64_t> view_indices_;
+      bool include_vector_dims_;
+    };
+
+    Iterator begin() const {
+      if (view_->GetNumElements() == 0) return end();
+      return {
+          view_,
+          llvm::SmallVector<int64_t>(
+              view_->Rank() +
+              (include_vector_dims_ ? view_->num_vector_dims.value_or(0) : 0)),
+          include_vector_dims_};
+    }
+    Iterator end() const { return {view_, {-1}, false}; }
+
+   private:
+    friend class BufferView;
+
+    LogicalIndexView(const BufferView* view, bool include_vector_dims)
+        : view_(view), include_vector_dims_(include_vector_dims) {}
+
+    const BufferView* view_;
+    bool include_vector_dims_;
+  };
+
+  // Returns nullopt if the index is out of bounds.
+  std::optional<int64_t> GetPhysicalIndex(
+      llvm::ArrayRef<int64_t> view_indices) const;
+  LogicalIndexView Indices(bool include_vector_dims = false) const {
+    return LogicalIndexView{this, include_vector_dims};
+  }
+  // Returns the stride resulting from collapsing the given dimensions, if
+  // possible.
+  std::optional<int64_t> GetCollapsedStride(llvm::ArrayRef<int64_t> dims) const;
+
+  bool InBounds(llvm::ArrayRef<int64_t> view_indices) const;
+  static SmallVector<int64_t> GetDefaultStrides(ArrayRef<int64_t> sizes);
+  static SmallVector<int64_t> GetStridesForLayout(ArrayRef<int64_t> sizes,
+                                                  ArrayRef<int64_t> layout);
+};
+
+// Backing for a TensorOrMemref.
+class Buffer {
+ private:
+  struct Dummy {};
+
+ public:
+  template <typename T>
+  static std::shared_ptr<Buffer> Allocate(size_t size) {
+    return std::make_shared<Buffer>(Dummy{}, size, sizeof(T));
+  }
+
+  char* at(std::optional<int64_t> idx, int64_t elementSize) {
+    auto byteOffset = GetByteOffset(idx, elementSize);
+    if (!byteOffset) {
+      return &storage_.data()[0];
+    }
+    return &storage_.data()[*byteOffset];
+  }
+
+  const char* at(std::optional<int64_t> idx, int64_t elementSize) const {
+    auto byteOffset = GetByteOffset(idx, elementSize);
+    if (!byteOffset) {
+      return &storage_.data()[0];
+    }
+    return &storage_.data()[*byteOffset];
+  }
+
+  Buffer(Dummy, size_t numElements, size_t elementSize)
+      : storage_(numElements * elementSize) {}
+
+  int64_t GetByteSize() const { return storage_.size(); }
+
+  void Deallocate(mlir::Operation* op) {
+    if (is_alloca_) {
+      SetFailure("deallocated stack buffer");
+    } else if (freed_by_ != nullptr) {
+      std::string failure;
+      llvm::raw_string_ostream os(failure);
+      os << "double-free\n";
+      os << "  Note: allocated by " << *allocated_by_ << "\n";
+      os << "  Note: previously freed by " << *freed_by_ << "\n";
+      SetFailure(failure);
+    } else {
+      freed_by_ = op;
+    }
+  }
+
+  bool Deallocated() const { return freed_by_ != nullptr; }
+  mlir::Operation* FreedByOp() const { return freed_by_; }
+  void SetAllocatedBy(mlir::Operation* allocated_by) {
+    this->allocated_by_ = allocated_by;
+  }
+
+  void SetFailure(llvm::StringRef failure) const {
+    this->failure_ = failure.str();
+  }
+  llvm::StringRef GetFailure() const { return failure_; }
+
+  void SetIsAlloca() { is_alloca_ = true; }
+
+ private:
+  std::optional<size_t> GetByteOffset(std::optional<int64_t> idx,
+                                      int64_t elementSize) const {
+    if (!idx) {
+      SetFailure("out of bounds access");
+      return std::nullopt;
+    }
+
+    if (freed_by_ != nullptr) {
+      std::string failure;
+      llvm::raw_string_ostream os(failure);
+      os << "use-after-free\n";
+      os << "  Note: allocated by " << *allocated_by_ << "\n";
+      os << "  Note: previously freed by " << *freed_by_ << "\n";
+      SetFailure(failure);
+      return std::nullopt;
+    }
+
+    return *idx * elementSize;
+  }
+
+  llvm::SmallVector<char> storage_;
+  mlir::Operation* freed_by_ = nullptr;
+  mlir::Operation* allocated_by_ = nullptr;
+  bool is_alloca_ = false;
+  mutable std::string failure_;
+};
+
+template <typename T>
+struct TensorOrMemref {
+  using element_type = T;
+
+  static TensorOrMemref<T> Empty(ArrayRef<int64_t> sizes,
+                                 ArrayRef<int64_t> layout = {}) {
+    BufferView dummy{0, SmallVector<int64_t>(sizes), {}};
+    return EmptyLike(dummy, layout);
+  }
+
+  static TensorOrMemref<T> EmptyLike(const BufferView& view,
+                                     ArrayRef<int64_t> layout = {}) {
+    BufferView newView = view;
+    newView.offset = 0;
+    newView.strides = BufferView::GetStridesForLayout(view.sizes, layout);
+    return {Buffer::Allocate<T>(view.GetNumElements(true)), newView};
+  }
+
+  TensorOrMemref<T> Clone(ArrayRef<int64_t> layout = {}) const {
+    auto out = EmptyLike(view, layout);
+    for (auto [src_index, dst_index] :
+         llvm::zip(view.Indices(true), out.view.Indices(true))) {
+      out.at(dst_index) = at(src_index);
+    }
+    return out;
+  }
+
+  const T& at(ArrayRef<int64_t> indices) const {
+    return *reinterpret_cast<const T*>(
+        buffer->at(view.GetPhysicalIndex(indices), sizeof(T)));
+  }
+
+  T& at(ArrayRef<int64_t> indices) {
+    return *reinterpret_cast<T*>(
+        buffer->at(view.GetPhysicalIndex(indices), sizeof(T)));
+  }
+
+  TensorOrMemref VectorAt(ArrayRef<int64_t> indices) const {
+    auto offset = view.GetPhysicalIndex(indices);
+    BufferView subview;
+    subview.strides = {view.strides.begin() + view.Rank(), view.strides.end()};
+    subview.sizes = {view.sizes.begin() + view.Rank(), view.sizes.end()};
+    if (offset) {
+      subview.offset = *offset;
+    } else {
+      buffer->SetFailure("out of bounds access");
+    }
+    subview.is_vector = true;
+    subview.num_vector_dims = std::nullopt;
+    return {buffer, subview};
+  }
+
+  bool operator==(const TensorOrMemref& other) const {
+    if (buffer->Deallocated() || other.buffer->Deallocated()) return false;
+    if (other.view.sizes != view.sizes) return false;
+    if (other.view.num_vector_dims != view.num_vector_dims) return false;
+    for (const auto& indices : view.Indices(true)) {
+      // Treat NaNs as equal.
+      if constexpr (std::is_floating_point_v<T>) {
+        bool thisnan = isnan(at(indices));
+        bool othernan = isnan(other.at(indices));
+        if (thisnan || othernan) {
+          if (thisnan && othernan) continue;
+          return false;
+        }
+      }
+      if (!IsEqual(at(indices), other.at(indices))) return false;
+    }
+    return true;
+  }
+
+  std::shared_ptr<Buffer> buffer;
+  BufferView view;
+};
+
+template <typename T>
+struct is_tensor_or_memref : std::false_type {};  // NOLINT
+
+template <typename T>
+struct is_tensor_or_memref<TensorOrMemref<T>> : std::true_type {};  // NOLINT
+
+template <typename T>
+inline constexpr bool is_tensor_or_memref_v =  // NOLINT
+    is_tensor_or_memref<std::decay_t<T>>::value;
+
+}  // namespace interpreter
+}  // namespace mlir
+
+#endif  // XLA_MLIR_TOOLS_MLIR_INTERPRETER_FRAMEWORK_TENSOR_OR_MEMREF_H_
diff --git a/third_party/xla/xla/mlir/tools/mlir_interpreter/framework/tests/BUILD b/third_party/xla/xla/mlir/tools/mlir_interpreter/framework/tests/BUILD
new file mode 100644
index 00000000000000..cd273637bc029b
--- /dev/null
+++ b/third_party/xla/xla/mlir/tools/mlir_interpreter/framework/tests/BUILD
@@ -0,0 +1,25 @@
+load("//xla:xla.bzl", "xla_cc_test")
+
+# copybara:uncomment package(default_applicable_licenses = ["//tensorflow:license"])
+
+xla_cc_test(
+    name = "tensor_or_memref_test",
+    srcs = ["tensor_or_memref_test.cc"],
+    deps = [
+        "//xla/mlir/tools/mlir_interpreter/framework",
+        "@com_google_absl//absl/strings",
+        "@com_google_googletest//:gtest_main",
+        "@llvm-project//llvm:Support",
+        "@llvm-project//mlir:Support",
+    ],
+)
+
+xla_cc_test(
+    name = "interpreter_value_test",
+    srcs = ["interpreter_value_test.cc"],
+    deps = [
+        "//xla/mlir/tools/mlir_interpreter/framework",
+        "@com_google_googletest//:gtest_main",
+        "@llvm-project//llvm:Support",
+    ],
+)
diff --git a/third_party/xla/xla/mlir/tools/mlir_interpreter/framework/tests/interpreter_value_test.cc b/third_party/xla/xla/mlir/tools/mlir_interpreter/framework/tests/interpreter_value_test.cc
new file mode 100644
index 00000000000000..93e72125440bc4
--- /dev/null
+++ b/third_party/xla/xla/mlir/tools/mlir_interpreter/framework/tests/interpreter_value_test.cc
@@ -0,0 +1,238 @@
+/* Copyright 2022 The OpenXLA Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+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.
+==============================================================================*/
+
+#include "xla/mlir/tools/mlir_interpreter/framework/interpreter_value.h"
+
+#include <complex>
+#include <cstdint>
+#include <variant>
+
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+#include "llvm/ADT/ArrayRef.h"
+#include "xla/mlir/tools/mlir_interpreter/framework/tensor_or_memref.h"
+
+namespace mlir {
+namespace interpreter {
+namespace {
+
+using ::testing::ElementsAre;
+using ::testing::IsEmpty;
+
+TEST(InterpreterValueTest, FillUnitTensor) {
+  auto t = TensorOrMemref<int64_t>::Empty({});
+  t.at({}) = 42;
+  InterpreterValue v{t};
+  v.Fill([](llvm::ArrayRef<int64_t>) { return InterpreterValue{int64_t{43}}; });
+  ASSERT_EQ(t.at({}), 43);
+}
+
+TEST(InterpreterValueTest, Fill1DTensor) {
+  auto t = TensorOrMemref<int64_t>::Empty({3});
+  InterpreterValue v{t};
+  v.Fill([](llvm::ArrayRef<int64_t> indices) {
+    return InterpreterValue{indices[0]};
+  });
+  ASSERT_EQ(t.at(0), 0);
+  ASSERT_EQ(t.at(1), 1);
+  ASSERT_EQ(t.at(2), 2);
+}
+
+TEST(InterpreterValueTest, FillTensorOfVector) {
+  auto t = TensorOrMemref<int64_t>::Empty({4, 2});
+  t.view.num_vector_dims = 1;
+
+  InterpreterValue v{t};
+  v.Fill([](llvm::ArrayRef<int64_t> indices) -> InterpreterValue {
+    EXPECT_EQ(indices.size(), 1);
+    auto r = TensorOrMemref<int64_t>::Empty({2});
+    r.view.is_vector = true;
+    r.at(0) = indices[0];
+    r.at(1) = indices[0] * 10;
+    return {r};
+  });
+  ASSERT_EQ(
+      v.ToString(),
+      "TensorOrMemref<4xvector<2xi64>>: [[0, 0], [1, 10], [2, 20], [3, 30]]");
+}
+
+TEST(InterpreterValueTest, FillZeroSizedTensor) {
+  auto t = TensorOrMemref<int64_t>::Empty({0, 1});
+  InterpreterValue v{t};
+  bool wasCalled = false;
+  v.Fill([&](llvm::ArrayRef<int64_t> indices) {
+    wasCalled = true;
+    return InterpreterValue{indices[0]};
+  });
+  EXPECT_FALSE(wasCalled);
+}
+
+TEST(InterpreterValueTest, TypedAlike) {
+  InterpreterValue v{TensorOrMemref<int32_t>::Empty({})};
+  auto TypedAlike = v.TypedAlike({1, 2, 3});
+  ASSERT_TRUE(
+      std::holds_alternative<TensorOrMemref<int32_t>>(TypedAlike.storage));
+  ASSERT_THAT(TypedAlike.View().sizes, ElementsAre(1, 2, 3));
+}
+
+TEST(InterpreterValueTest, AsUnitTensor) {
+  InterpreterValue v{42};
+  InterpreterValue wrapped = v.AsUnitTensor();
+  ASSERT_THAT(wrapped.View().sizes, IsEmpty());
+  ASSERT_EQ(std::get<TensorOrMemref<int32_t>>(wrapped.storage).at({}), 42);
+}
+
+TEST(InterpreterValueTest, IsTensor) {
+  ASSERT_FALSE(InterpreterValue{42}.IsTensor());
+  ASSERT_TRUE(InterpreterValue{TensorOrMemref<int32_t>::Empty({})}.IsTensor());
+}
+
+TEST(InterpreterValueTest, AsInt) {
+  ASSERT_EQ(InterpreterValue{int64_t{42}}.AsInt(), 42);
+  ASSERT_EQ(InterpreterValue{int32_t{42}}.AsInt(), 42);
+  ASSERT_EQ(InterpreterValue{int16_t{42}}.AsInt(), 42);
+  ASSERT_EQ(InterpreterValue{int8_t{42}}.AsInt(), 42);
+  ASSERT_EQ(InterpreterValue{int8_t{-1}}.AsInt(), -1);
+}
+
+TEST(InterpreterValueTest, AsUInt) {
+  ASSERT_EQ(InterpreterValue{int16_t{-1}}.AsUInt(), 65535);
+  ASSERT_EQ(InterpreterValue{int8_t{-1}}.AsUInt(), 255);
+}
+
+TEST(InterpreterValueTest, CloneTensor) {
+  auto tensor = TensorOrMemref<int64_t>::Empty({3});
+  tensor.at(0) = 1;
+  tensor.at(1) = 2;
+  tensor.at(2) = 3;
+
+  InterpreterValue wrapped{tensor};
+  auto clone = wrapped.Clone();
+  tensor.at(0) = 4;
+
+  auto& cloned_tensor = std::get<TensorOrMemref<int64_t>>(clone.storage);
+  ASSERT_EQ(cloned_tensor.at(0), 1);
+  ASSERT_EQ(cloned_tensor.at(1), 2);
+  ASSERT_EQ(cloned_tensor.at(2), 3);
+}
+
+TEST(InterpreterValueTest, CloneWithLayouts) {
+  auto tensor = TensorOrMemref<int64_t>::Empty({3, 5}, {0, 1});
+  tensor.at({2, 4}) = 42;
+
+  InterpreterValue wrapped{tensor};
+  auto clone = wrapped.Clone();
+  ASSERT_EQ(clone.View().strides,
+            BufferView::GetStridesForLayout({3, 5}, {1, 0}));
+  ASSERT_EQ(clone.ExtractElement({2, 4}).AsInt(), 42);
+}
+
+TEST(InterpreterValueTest, CoerceLayoutNoop) {
+  auto tensor = TensorOrMemref<int64_t>::Empty({3, 5}, {0, 1});
+  tensor.at({2, 4}) = 42;
+
+  InterpreterValue wrapped{tensor};
+  auto coerced = wrapped.CoerceLayout({0, 1});
+  ASSERT_EQ(tensor.buffer,
+            std::get<TensorOrMemref<int64_t>>(coerced.storage).buffer);
+}
+
+TEST(InterpreterValueTest, CoerceLayout) {
+  auto tensor = TensorOrMemref<int64_t>::Empty({3, 5});
+  tensor.at({2, 4}) = 42;
+
+  InterpreterValue wrapped{tensor};
+  auto clone = wrapped.CoerceLayout({0, 1});
+  ASSERT_EQ(clone.View().strides,
+            BufferView::GetStridesForLayout({3, 5}, {0, 1}));
+  ASSERT_EQ(clone.ExtractElement({2, 4}).AsInt(), 42);
+}
+
+TEST(InterpreterValueTest, CoerceLayoutSquare) {
+  auto tensor = TensorOrMemref<float>::Empty({2, 2});
+  tensor.at({0, 0}) = 1;
+  tensor.at({0, 1}) = 2;
+  tensor.at({1, 0}) = 3;
+  tensor.at({1, 1}) = 4;
+
+  InterpreterValue wrapped{tensor};
+  auto clone = wrapped.CoerceLayout({0, 1});
+  auto& cloned_tensor = std::get<TensorOrMemref<float>>(clone.storage);
+
+  EXPECT_EQ(
+      *reinterpret_cast<float*>(cloned_tensor.buffer->at(0, sizeof(float))), 1);
+  EXPECT_EQ(
+      *reinterpret_cast<float*>(cloned_tensor.buffer->at(1, sizeof(float))), 3);
+  EXPECT_EQ(
+      *reinterpret_cast<float*>(cloned_tensor.buffer->at(2, sizeof(float))), 2);
+  EXPECT_EQ(
+      *reinterpret_cast<float*>(cloned_tensor.buffer->at(3, sizeof(float))), 4);
+}
+
+TEST(InterpreterValueTest, CloneScalar) {
+  InterpreterValue value{42};
+  auto clone = value.Clone();
+  ASSERT_THAT(std::get<int32_t>(clone.storage), 42);
+}
+
+TEST(InterpreterValueTest, ToString) {
+  InterpreterValue value{TensorOrMemref<int64_t>::Empty({3})};
+  ASSERT_EQ(value.ToString(), "TensorOrMemref<3xi64>: [0, 0, 0]");
+}
+
+TEST(InterpreterValueTest, ToString2d) {
+  InterpreterValue value{TensorOrMemref<int64_t>::Empty({3, 2})};
+  ASSERT_EQ(value.ToString(),
+            "TensorOrMemref<3x2xi64>: [[0, 0], [0, 0], [0, 0]]");
+}
+
+TEST(InterpreterValueTest, ToString0d) {
+  InterpreterValue value{TensorOrMemref<int64_t>::Empty({})};
+  ASSERT_EQ(value.ToString(), "TensorOrMemref<i64>: 0");
+}
+
+TEST(InterpreterValueTest, ToStringComplex) {
+  InterpreterValue value{std::complex<float>{}};
+  ASSERT_EQ(value.ToString(), "complex<f32>: 0.000000e+00+0.000000e+00i");
+}
+
+TEST(CastTest, UnpackTensor) {
+  InterpreterValue value{TensorOrMemref<int8_t>::Empty({1, 1})};
+  value.InsertElement({0, 0}, {int8_t{1}});
+  ASSERT_EQ(InterpreterValueCast<int64_t>(value), 1);
+  ASSERT_EQ(InterpreterValueCast<uint8_t>(value), 1);
+  ASSERT_EQ(InterpreterValueCast<float>(value), 1.0f);
+  ASSERT_EQ(InterpreterValueCast<double>(value), 1.0);
+
+  InterpreterValue nonUnit{TensorOrMemref<int8_t>::Empty({2, 2})};
+  ASSERT_EQ(InterpreterValueDynCast<int64_t>(nonUnit), std::nullopt);
+}
+
+TEST(CastTest, IdentityCast) {
+  InterpreterValue value{TensorOrMemref<float>::Empty({1, 1})};
+  ASSERT_EQ(InterpreterValueCast<InterpreterValue>(value), value);
+}
+
+TEST(CastTest, CastToUnsigned) {
+  // Note: This is different from `AsUint`, which preserves the size of the
+  // original type (i.e. int8_t{-1} results in 255).
+  InterpreterValue value{int8_t{-1}};
+  ASSERT_EQ(InterpreterValueCast<uint8_t>(value), 255);
+  ASSERT_EQ(InterpreterValueCast<uint16_t>(value), 65535);
+}
+
+}  // namespace
+}  // namespace interpreter
+}  // namespace mlir
diff --git a/third_party/xla/xla/mlir/tools/mlir_interpreter/framework/tests/tensor_or_memref_test.cc b/third_party/xla/xla/mlir/tools/mlir_interpreter/framework/tests/tensor_or_memref_test.cc
new file mode 100644
index 00000000000000..6d431234a57ce1
--- /dev/null
+++ b/third_party/xla/xla/mlir/tools/mlir_interpreter/framework/tests/tensor_or_memref_test.cc
@@ -0,0 +1,104 @@
+/* Copyright 2022 The OpenXLA Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+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.
+==============================================================================*/
+
+#include "xla/mlir/tools/mlir_interpreter/framework/tensor_or_memref.h"
+
+#include <algorithm>
+#include <cstdint>
+#include <optional>
+
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+#include "absl/strings/str_join.h"
+#include "llvm/ADT/SmallBitVector.h"
+#include "llvm/ADT/SmallVector.h"
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+
+namespace mlir {
+namespace interpreter {
+namespace {
+
+using ::testing::ElementsAre;
+
+TEST(TensorOrMemrefTest, DefaultStrides) {
+  EXPECT_THAT(BufferView::GetDefaultStrides({1, 2, 3}), ElementsAre(6, 3, 1));
+}
+
+TEST(TensorOrMemrefTest, StridesForLayout) {
+  EXPECT_THAT(BufferView::GetStridesForLayout({1, 2, 3}, {2, 1, 0}),
+              ElementsAre(6, 3, 1));
+  EXPECT_THAT(BufferView::GetStridesForLayout({1, 2, 3}, {0, 1, 2}),
+              ElementsAre(1, 1, 2));
+  EXPECT_THAT(BufferView::GetStridesForLayout({3, 3, 3, 3}, {3, 0, 1, 2}),
+              ElementsAre(27, 1, 3, 9));
+}
+
+std::optional<int64_t> GetCollapsedStrideNaive(llvm::ArrayRef<int64_t> dims,
+                                               const BufferView& view) {
+  BufferView f;
+  for (int64_t dim : dims) {
+    f.sizes.push_back(view.sizes[dim]);
+  }
+
+  // Find all physical indices for the dimensions.
+  llvm::SmallBitVector v(view.GetNumElements());
+  for (const auto& indices : f.Indices()) {
+    SmallVector<int64_t> viewIndices(view.Rank());
+    for (auto [dim, index] : llvm::zip(dims, indices)) {
+      viewIndices[dim] = index;
+    }
+    v[*view.GetPhysicalIndex(viewIndices)] = true;
+  }
+
+  if (v.count() != f.GetNumElements()) return std::nullopt;
+  if (f.GetNumElements() <= 1) return 0;
+
+  // Check that they have a common stride.
+  int64_t min = v.find_first();
+  int64_t expectedStride = (v.find_last() - min) / (f.GetNumElements() - 1);
+  for (int64_t i = 0; i < f.GetNumElements(); ++i) {
+    if (!v[i * expectedStride + min]) return std::nullopt;
+  }
+
+  return expectedStride;
+}
+
+TEST(TensorOrMemrefTest, CollapsedStride) {
+  BufferView view{.sizes = {1, 2, 3, 1, 5},
+                  .strides = BufferView::GetDefaultStrides({1, 2, 3, 1, 5})};
+
+  auto check_all = [&]() {
+    for (int64_t i = 0; i < (1 << view.Rank()); ++i) {
+      SmallVector<int64_t> dims;
+      for (int64_t dim = 0; dim < view.Rank(); ++dim) {
+        if (i & (1 << dim)) dims.push_back(dim);
+      }
+
+      do {
+        auto v = view.GetCollapsedStride(dims);
+        auto n = GetCollapsedStrideNaive(dims, view);
+        EXPECT_EQ(n, v) << "checking " << absl::StrJoin(dims, ", ");
+      } while (std::next_permutation(dims.begin(), dims.end()));
+    }
+  };
+
+  check_all();
+  ASSERT_TRUE(view.Slice(3, 0).succeeded());
+  check_all();
+}
+
+}  // namespace
+}  // namespace interpreter
+}  // namespace mlir
diff --git a/third_party/xla/xla/mlir/tools/mlir_interpreter/mlir_interpreter_runner.cc b/third_party/xla/xla/mlir/tools/mlir_interpreter/mlir_interpreter_runner.cc
new file mode 100644
index 00000000000000..e75d0a603e810a
--- /dev/null
+++ b/third_party/xla/xla/mlir/tools/mlir_interpreter/mlir_interpreter_runner.cc
@@ -0,0 +1,138 @@
+/* Copyright 2022 The OpenXLA Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+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.
+==============================================================================*/
+
+#include <memory>
+#include <string>
+#include <utility>
+
+#include "mhlo/IR/register.h"
+#include "llvm/ADT/StringRef.h"
+#include "llvm/Support/Casting.h"
+#include "llvm/Support/CommandLine.h"
+#include "llvm/Support/SourceMgr.h"
+#include "llvm/Support/raw_ostream.h"
+#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
+#include "mlir/IR/DialectRegistry.h"  // from @llvm-project
+#include "mlir/IR/MLIRContext.h"  // from @llvm-project
+#include "mlir/IR/Operation.h"  // from @llvm-project
+#include "mlir/IR/OwningOpRef.h"  // from @llvm-project
+#include "mlir/IR/SymbolTable.h"  // from @llvm-project
+#include "mlir/InitAllDialects.h"  // from @llvm-project
+#include "mlir/InitAllPasses.h"  // from @llvm-project
+#include "mlir/Support/FileUtilities.h"  // from @llvm-project
+#include "mlir/Support/LogicalResult.h"  // from @llvm-project
+#include "mlir/Tools/ParseUtilities.h"  // from @llvm-project
+#include "xla/mlir/tools/mlir_interpreter/framework/interpreter.h"
+#include "xla/mlir_hlo/mhlo/IR/register.h"
+#include "xla/mlir_hlo/mhlo/transforms/passes.h"
+
+struct Options {
+  llvm::cl::opt<std::string> input_filename{llvm::cl::Positional,
+                                            llvm::cl::desc("<input file>"),
+                                            llvm::cl::init("-")};
+  llvm::cl::opt<bool> run_all_functions{
+      "run-all", llvm::cl::desc("Run all functions in the module"),
+      llvm::cl::init(false)};
+};
+
+static mlir::OwningOpRef<mlir::Operation *> ParseMlirInput(
+    llvm::StringRef inputFilename, mlir::MLIRContext *context) {
+  std::string errorMessage;
+  auto file = mlir::openInputFile(inputFilename, &errorMessage);
+  if (!file) {
+    llvm::errs() << errorMessage << "\n";
+    return {};
+  }
+
+  auto sourceMgr = std::make_shared<llvm::SourceMgr>();
+  sourceMgr->AddNewSourceBuffer(std::move(file), mlir::SMLoc());
+  return mlir::parseSourceFileForTool(sourceMgr, context,
+                                      /*insertImplicitModule=*/true);
+}
+
+static mlir::LogicalResult Run(mlir::ModuleOp module,
+                               mlir::func::FuncOp function) {
+  llvm::outs() << "@" << function.getName().str() << "()\n";
+  if (function.getBody().getBlocks().front().getNumArguments() > 0) {
+    llvm::errs() << "Function arguments are not supported.";
+    return mlir::failure();
+  }
+
+  mlir::SymbolTable symbolTable{module};
+  auto results =
+      mlir::interpreter::RunInterpreter(symbolTable, function, {}, {});
+  if (!results.ok()) {
+    llvm::errs() << "Interpreter failed\n";
+    return mlir::failure();
+  }
+
+  if (!results->empty()) {
+    llvm::outs() << "Results:\n";
+    for (const auto &result : *results) {
+      llvm::outs() << result.ToString() << "\n";
+    }
+  }
+
+  return mlir::success();
+}
+
+int main(int argc, char *argv[]) {
+  // Flush llvm::outs before writing errors.
+  llvm::errs().tie(&llvm::outs());
+
+  Options options;
+  llvm::cl::ParseCommandLineOptions(argc, argv, "MLIR interpreter driver\n");
+
+  mlir::DialectRegistry registry;
+  mlir::registerAllDialects(registry);
+  mlir::mhlo::registerAllMhloDialects(registry);
+  mlir::mhlo::registerAllMhloPasses();
+  mlir::registerAllPasses();
+
+  mlir::MLIRContext context(registry);
+  context.allowUnregisteredDialects();
+  auto parsedInput = ParseMlirInput(options.input_filename, &context);
+  if (!parsedInput) {
+    llvm::errs() << "Failed to parse module.\n";
+    return 1;
+  }
+  auto module = llvm::dyn_cast<mlir::ModuleOp>(**parsedInput);
+  if (!module) {
+    llvm::errs() << "Parsing returned something that's not a module.\n";
+    return 1;
+  }
+
+  if (options.run_all_functions) {
+    bool allSucceeded = true;
+    module.walk([&](mlir::func::FuncOp function) {
+      if (!function.isPrivate()) {
+        allSucceeded &= Run(module, function).succeeded();
+      }
+    });
+    if (!allSucceeded) {
+      return 1;
+    }
+  } else {
+    auto *main = module.lookupSymbol("main");
+    if (!main) {
+      llvm::errs() << "no main function found.\n";
+      return 1;
+    }
+    if (!Run(module, llvm::cast<mlir::func::FuncOp>(main)).succeeded()) {
+      return 1;
+    }
+  }
+  return 0;
+}
diff --git a/third_party/xla/xla/mlir/tools/mlir_replay/BUILD b/third_party/xla/xla/mlir/tools/mlir_replay/BUILD
new file mode 100644
index 00000000000000..c0bb7fe0e5c580
--- /dev/null
+++ b/third_party/xla/xla/mlir/tools/mlir_replay/BUILD
@@ -0,0 +1,63 @@
+load("@bazel_skylib//rules:build_test.bzl", "build_test")
+load("@local_tsl//tsl/platform:rules_cc.bzl", "cc_library")
+load("//xla:xla.bzl", "xla_cc_binary")
+
+# copybara:uncomment package(default_applicable_licenses = ["//tensorflow:license"])
+
+build_test(
+    name = "mlir_replay_build_test",
+    targets = [
+        ":mlir_replay",
+    ],
+)
+
+xla_cc_binary(
+    name = "mlir_replay",
+    testonly = True,
+    srcs = ["mlir_replay.cc"],
+    deps = [
+        ":mlir_replay_lib",
+        "//xla:debug_options_flags",
+        "//xla/mlir/tools/mlir_interpreter/dialects",
+        "//xla/mlir/tools/mlir_interpreter/framework",
+        "//xla/mlir/tools/mlir_replay/public:compiler_trace_proto_cc",
+        "//xla/mlir/tools/mlir_replay/public:compiler_trace_proto_cc_impl",
+        "//xla/mlir/tools/mlir_replay/public:execution_trace_proto_cc",
+        "//xla/mlir/tools/mlir_replay/public:execution_trace_utils",
+        "//xla/mlir_hlo:hlo_dialect_registration",
+        "//xla/service:hlo_proto_cc",
+        "//xla/tsl/util:command_line_flags",
+        "@com_google_absl//absl/strings",
+        "@com_google_absl//absl/strings:str_format",
+        "@llvm-project//mlir:AllPassesAndDialects",
+        "@llvm-project//mlir:IR",
+        "@local_tsl//tsl/platform:env",
+        "@local_tsl//tsl/platform:path",
+        "@local_tsl//tsl/platform:platform_port",
+        "@local_tsl//tsl/platform:status",
+    ],
+)
+
+cc_library(
+    name = "mlir_replay_lib",
+    srcs = ["mlir_replay_lib.cc"],
+    hdrs = ["mlir_replay_lib.h"],
+    deps = [
+        "//xla/mlir/tools/mlir_interpreter/framework",
+        "//xla/mlir/tools/mlir_replay/public:execution_trace_proto_cc",
+        "//xla/mlir/tools/mlir_replay/public:execution_trace_utils",
+        "//xla/service:hlo_proto_cc",
+        "@com_google_absl//absl/container:flat_hash_map",
+        "@com_google_absl//absl/random",
+        "@com_google_absl//absl/random:bit_gen_ref",
+        "@com_google_absl//absl/random:distributions",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/status:statusor",
+        "@llvm-project//llvm:Support",
+        "@llvm-project//mlir:FuncDialect",
+        "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:ParseUtilities",
+        "@llvm-project//mlir:Support",
+        "@local_tsl//tsl/platform:statusor",
+    ],
+)
diff --git a/third_party/xla/xla/mlir/tools/mlir_replay/README.md b/third_party/xla/xla/mlir/tools/mlir_replay/README.md
new file mode 100644
index 00000000000000..6c2091d526e92f
--- /dev/null
+++ b/third_party/xla/xla/mlir/tools/mlir_replay/README.md
@@ -0,0 +1,48 @@
+# MLIR Replay tool
+
+This tool is mainly intended for helping debug miscompiles. It takes as inputs
+an HLO snapshot proto with input tensors and a compiler trace proto with the
+state of the IR after each pass.
+
+This tool is built on top of
+[mlir-interpreter](https://github.com/tensorflow/mlir-hlo/tree/master/tools/mlir_interpreter/).
+
+Example usage:
+
+```
+# Run a JAX test with debug flags enabled:
+$ bazel test :some_jax_test --compilation_mode=opt \
+  --test_env=XLA_FLAGS="--xla_cpu_use_xla_runtime --xla_dump_to=/tmp/test-dump --xla_dump_hlo_snapshots" \
+  --test_filter=SomeSpecific.TestCase \
+  --test_sharding_strategy=disabled --test_strategy=local
+
+# JAX tends to compile many modules, so first check which one is broken:
+./mlir_replay \
+  --mlir-compilation-trace-dir=/tmp/test-dump
+
+Failures for /tmp/test-dump/module_1234.jit_something.mlir-trace.pb:
+  Result mismatch for /tmp/test-dump/module_1234.jit_something.snapshot.56.pb: TensorOrMemref<3xi32>: [1, 2, 3] != TensorOrMemref<3xi32>: [1, 1, 1]
+  run :mlir_replay -- --mlir-compilation-trace=/tmp/test-dump/module_1234.jit_something.mlir-trace.pb --hlo-snapshot=/tmp/test-dump/module_1234.jit_something.snapshot.56.pb --print-changes-only --stop-after-first-failure
+```
+
+There may be multiple failing modules. You can run the provided command to
+replay a particular one:
+
+```
+# Run the IR after each pass. Note that JAX typically compiles many modules, so
+# you may have check more than one.
+# There is one .mlir-trace.pb file per module (containing the intermediate IR)
+# and one .snapshot.pb file per execution (containing the inputs and outputs).
+$ ./mlir_replay \
+  --mlir-compilation-trace=/tmp/test-dump/module_1234.jit_something.mlir-trace.pb \
+  --hlo-snapshot=/tmp/test-dump/module_1234.jit_something.snapshot.56.pb \
+  --print-changes-only --stop-after-first-failure
+Running IR after APass
+Results: [1, 2, 3]
+
+Running IR after BPass
+Running IR after CPass
+Running IR after BrokenPass
+Results: [1, 1, 1]
+```
+
diff --git a/third_party/xla/xla/mlir/tools/mlir_replay/mlir_replay.cc b/third_party/xla/xla/mlir/tools/mlir_replay/mlir_replay.cc
new file mode 100644
index 00000000000000..e02e7fd3eb2641
--- /dev/null
+++ b/third_party/xla/xla/mlir/tools/mlir_replay/mlir_replay.cc
@@ -0,0 +1,236 @@
+/* Copyright 2023 The OpenXLA Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+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.
+==============================================================================*/
+
+#include <cstring>
+#include <string>
+#include <vector>
+
+#include "absl/strings/str_format.h"
+#include "absl/strings/str_split.h"
+#include "mlir/IR/DialectRegistry.h"  // from @llvm-project
+#include "mlir/IR/MLIRContext.h"  // from @llvm-project
+#include "mlir/InitAllDialects.h"  // from @llvm-project
+#include "xla/debug_options_flags.h"
+#include "xla/mlir/tools/mlir_interpreter/framework/interpreter_value.h"
+#include "xla/mlir/tools/mlir_replay/mlir_replay_lib.h"
+#include "xla/mlir/tools/mlir_replay/public/compiler_trace.pb.h"
+#include "xla/mlir/tools/mlir_replay/public/execution_trace.pb.h"
+#include "xla/mlir/tools/mlir_replay/public/execution_trace_utils.h"
+#include "xla/mlir_hlo/mhlo/IR/register.h"
+#include "xla/service/hlo.pb.h"
+#include "xla/tsl/util/command_line_flags.h"
+#include "tsl/platform/env.h"
+#include "tsl/platform/init_main.h"
+#include "tsl/platform/path.h"
+#include "tsl/platform/status.h"
+
+struct ReplayOptions {
+  std::string hlo_snapshot;
+  std::string mlir_compilation_trace;
+  std::string mlir_compilation_trace_dir = "";
+  std::string execution_trace_dir = "";
+  std::vector<std::string> entry_points = {"main", "main_xla_framework"};
+  bool print_changes_only = false;
+  bool stop_after_first_failure = false;
+  bool print_values = true;
+};
+
+bool ResultsMatch(const xla::HloSnapshot& snapshot,
+                  const llvm::SmallVector<mlir::interpreter::InterpreterValue>&
+                      first_pass_results,
+                  std::vector<std::string>& failures,
+                  const ReplayOptions& opts) {
+  auto actual = mlir::interpreter::LiteralToValue(snapshot.result());
+  TF_CHECK_OK(actual.status());
+
+  // We assume this is MHLO, so multiple results will be in a tuple.
+  if (first_pass_results.size() != 1) {
+    failures.push_back("expected one result");
+    return false;
+  }
+
+  if (!(*actual == first_pass_results[0])) {
+    if (opts.print_values) {
+      failures.push_back("result mismatch: " + actual->ToString() +
+                         " != " + first_pass_results[0].ToString());
+    } else {
+      failures.push_back("result mismatch");
+    }
+    return false;
+  }
+  return true;
+}
+
+void TestAll(mlir::MLIRContext& context, const ReplayOptions& opts) {
+  std::vector<std::string> traces;
+  TF_CHECK_OK(tsl::Env::Default()->GetMatchingPaths(
+      opts.mlir_compilation_trace_dir + "/*.mlir-trace.pb", &traces));
+
+  for (const auto& trace_path : traces) {
+    mlir::interpreter::MlirCompilationTrace trace;
+    TF_CHECK_OK(tsl::ReadBinaryProto(tsl::Env::Default(), trace_path, &trace))
+        << "Failed to load " << trace_path;
+
+    std::vector<std::string> snapshots;
+    std::string prefix =
+        trace_path.substr(0, trace_path.length() - strlen(".mlir-trace.pb"));
+    TF_CHECK_OK(tsl::Env::Default()->GetMatchingPaths(prefix + "*.snapshot.*",
+                                                      &snapshots));
+    CHECK_NE(snapshots.size(), 0)
+        << "No snapshots found for module " << trace_path << ".";
+
+    std::vector<std::string> failures;
+    for (const auto& snapshot_path : snapshots) {
+      xla::HloSnapshot snapshot;
+      TF_CHECK_OK(
+          tsl::ReadBinaryProto(tsl::Env::Default(), snapshot_path, &snapshot));
+
+      auto results =
+          mlir::interpreter::Run(context, trace.passes(0).mlir_module(),
+                                 snapshot, nullptr, opts.entry_points);
+      if (!results.status().ok()) {
+        failures.push_back("Failed to execute " + snapshot_path + ": " +
+                           results.status().ToString());
+      } else {
+        if (!ResultsMatch(snapshot, *results, failures, opts)) {
+          failures.push_back(
+              std::string("run :mlir_replay -- --mlir-compilation-trace=") +
+              trace_path + " --hlo-snapshot=" + snapshot_path +
+              " --print-changes-only --stop-after-first-failure");
+        }
+      }
+    }
+
+    if (!failures.empty()) {
+      llvm::errs() << "Failures for " << trace_path << ":\n  "
+                   << absl::StrJoin(failures, "\n  ") << "\n";
+    }
+  }
+}
+
+int main(int argc, char* argv[]) {
+  // Flush llvm::outs before writing errors.
+  llvm::errs().tie(&llvm::outs());
+
+  std::string entry_points;
+  ReplayOptions opts;
+  std::vector<tsl::Flag> flag_list = {
+      tsl::Flag("hlo-snapshot", &opts.hlo_snapshot,
+                "Filename of an HloSnapshot proto. Only used to read inputs."),
+      tsl::Flag("mlir-compilation-trace", &opts.mlir_compilation_trace,
+                "Filename of an MlirCompilerTrace proto."),
+      tsl::Flag("mlir-compilation-trace-dir", &opts.mlir_compilation_trace_dir,
+                "Directory from which to load MlirCompilerTrace and "
+                "HloSnapshot protos. The tool will run all snapshots and "
+                "report the ones with bugs."),
+      tsl::Flag("execution-trace-dir", &opts.execution_trace_dir,
+                "Directory where to store the execution traces (optional)."),
+      tsl::Flag("entry-point", &entry_points,
+                "Program entry function (optional, defaults to 'main')."),
+      tsl::Flag("print-changes-only", &opts.print_changes_only,
+                "If set, only print changed values"),
+      tsl::Flag("stop-after-first-failure", &opts.stop_after_first_failure,
+                "If set, stop after the first failed invocation."),
+      tsl::Flag("print-values", &opts.print_values, "If set, print values."),
+  };
+  xla::AppendDebugOptionsFlags(&flag_list);
+
+  // The usage string includes the message at the top of the file, the
+  // DebugOptions flags and the flags defined above.
+  std::string usage_string = tsl::Flags::Usage(argv[0], flag_list);
+  if (!tsl::Flags::Parse(&argc, argv, flag_list)) {
+    return 1;
+  }
+
+  if (!entry_points.empty()) {
+    opts.entry_points = absl::StrSplit(entry_points, ',');
+  }
+
+  tsl::port::InitMain(usage_string.c_str(), &argc, &argv);
+
+  CHECK(opts.mlir_compilation_trace.empty() !=
+        opts.mlir_compilation_trace_dir.empty())
+      << "Exactly one of --mlir-compilation-trace and "
+         "--mlir-compilation-trace-dir must be specified.";
+
+  CHECK(opts.mlir_compilation_trace_dir.empty() || opts.hlo_snapshot.empty())
+      << "If --mlir-compilation-trace-dir is set, --hlo-snapshot must not be.";
+
+  mlir::DialectRegistry registry;
+  mlir::registerAllDialects(registry);
+  mlir::mhlo::registerAllMhloDialects(registry);
+
+  mlir::MLIRContext context(registry);
+
+  if (!opts.mlir_compilation_trace_dir.empty()) {
+    TestAll(context, opts);
+    return 0;
+  }
+
+  xla::HloSnapshot snapshot;
+  if (!opts.hlo_snapshot.empty()) {
+    TF_CHECK_OK(tsl::ReadBinaryProto(tsl::Env::Default(), opts.hlo_snapshot,
+                                     &snapshot));
+  }
+  mlir::interpreter::MlirCompilationTrace trace;
+  TF_CHECK_OK(tsl::ReadBinaryProto(tsl::Env::Default(),
+                                   opts.mlir_compilation_trace, &trace));
+
+  llvm::SmallVector<mlir::interpreter::InterpreterValue> previous_results;
+  int pass_id = 0;
+  for (auto& state : trace.passes()) {
+    llvm::outs() << "Running IR after " << state.after_pass() << ".\n";
+    mlir::interpreter::ExecutionTrace execution_trace;
+    auto results = mlir::interpreter::Run(
+        context, state.mlir_module(), snapshot,
+        opts.execution_trace_dir.empty() ? nullptr : &execution_trace,
+        opts.entry_points);
+    if (results.status().ok()) {
+      if (opts.print_values &&
+          (!opts.print_changes_only || (*results != previous_results))) {
+        llvm::outs() << "Results:\n";
+        for (const auto& result : *results) {
+          llvm::outs() << result.ToString() << "\n";
+        }
+        previous_results = *results;
+        llvm::outs() << "\n";
+      }
+    } else {
+      llvm::errs() << results.status().ToString() << "\n";
+      if (opts.stop_after_first_failure) {
+        return 1;
+      }
+    }
+
+    if (!opts.execution_trace_dir.empty()) {
+      TF_CHECK_OK(
+          tsl::Env::Default()->RecursivelyCreateDir(opts.execution_trace_dir));
+      std::string filename = tsl::io::JoinPath(
+          opts.execution_trace_dir,
+          absl::StrFormat("%.4d.%s.mlir", pass_id, state.after_pass()));
+      TF_CHECK_OK(tsl::WriteStringToFile(tsl::Env::Default(), filename,
+                                         execution_trace.ir()));
+
+      filename = tsl::io::JoinPath(
+          opts.execution_trace_dir,
+          absl::StrFormat("%.4d.%s.trace.pb", pass_id, state.after_pass()));
+      TF_CHECK_OK(tsl::WriteBinaryProto(tsl::Env::Default(), filename,
+                                        execution_trace));
+    }
+    ++pass_id;
+  }
+
+  return 0;
+}
diff --git a/third_party/xla/xla/mlir/tools/mlir_replay/mlir_replay_lib.cc b/third_party/xla/xla/mlir/tools/mlir_replay/mlir_replay_lib.cc
new file mode 100644
index 00000000000000..c13f06a9b953e6
--- /dev/null
+++ b/third_party/xla/xla/mlir/tools/mlir_replay/mlir_replay_lib.cc
@@ -0,0 +1,254 @@
+/* Copyright 2022 The OpenXLA Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+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.
+==============================================================================*/
+
+#include "xla/mlir/tools/mlir_replay/mlir_replay_lib.h"
+
+#include <cstdint>
+#include <iterator>
+#include <memory>
+#include <numeric>
+#include <random>
+#include <string>
+#include <utility>
+#include <vector>
+
+#include "absl/container/flat_hash_map.h"
+#include "absl/random/bit_gen_ref.h"
+#include "absl/random/gaussian_distribution.h"
+#include "absl/random/random.h"
+#include "absl/status/status.h"
+#include "absl/status/statusor.h"
+#include "llvm/ADT/APInt.h"
+#include "llvm/ADT/Sequence.h"
+#include "llvm/Support/SourceMgr.h"
+#include "llvm/Support/raw_ostream.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
+#include "mlir/IR/BuiltinAttributes.h"  // from @llvm-project
+#include "mlir/IR/Operation.h"  // from @llvm-project
+#include "mlir/IR/OwningOpRef.h"  // from @llvm-project
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+#include "mlir/Support/LogicalResult.h"  // from @llvm-project
+#include "mlir/Tools/ParseUtilities.h"  // from @llvm-project
+#include "xla/mlir/tools/mlir_interpreter/framework/interpreter.h"
+#include "xla/mlir/tools/mlir_interpreter/framework/interpreter_value.h"
+#include "xla/mlir/tools/mlir_interpreter/framework/tensor_or_memref.h"
+#include "xla/mlir/tools/mlir_replay/public/execution_trace_utils.h"
+#include "xla/service/hlo.pb.h"
+#include "tsl/platform/statusor.h"
+
+namespace mlir {
+namespace interpreter {
+namespace {
+
+absl::StatusOr<SmallVector<InterpreterValue>> LoadArgs(
+    const xla::HloSnapshot& snapshot, TypeRange types) {
+  SmallVector<InterpreterValue> result;
+  for (const auto& [arg, type] : llvm::zip(snapshot.arguments(), types)) {
+    TF_ASSIGN_OR_RETURN(auto converted, LiteralToValue(arg, type));
+    result.push_back(std::move(converted));
+  }
+  return result;
+}
+
+namespace {
+template <typename T, template <typename _> class rng_t>
+mlir::interpreter::InterpreterValue RandomTensor(absl::BitGenRef bitgen,
+                                                 mlir::Type type) {
+  llvm::SmallVector<int64_t> shape;
+  auto shaped_ty = type.dyn_cast<mlir::ShapedType>();
+  if (shaped_ty) {
+    shape = llvm::to_vector(shaped_ty.getShape());
+  }
+
+  auto rng = rng_t<T>{};
+  auto result = mlir::interpreter::TensorOrMemref<T>::Empty(shape);
+  for (const auto& index : result.view.Indices()) {
+    auto& elem = result.at(index) = rng(bitgen);
+    // Ints are typically indices, so scale them down to a more reasonable
+    // range.
+    if constexpr (std::is_same_v<T, int64_t>) {
+      elem >>= 60;
+    }
+  }
+  if (shaped_ty) {
+    return {result};
+  }
+  return {result.at({})};
+}
+}  // namespace
+
+mlir::FailureOr<mlir::interpreter::InterpreterValue> MakeRandomInput(
+    absl::BitGenRef bitgen, mlir::Type type) {
+  auto elem_ty =
+      type.isa<ShapedType>() ? type.cast<ShapedType>().getElementType() : type;
+  if (elem_ty.isF32()) {
+    return RandomTensor<float, absl::gaussian_distribution>(bitgen, type);
+  }
+  if (elem_ty.isF64()) {
+    return RandomTensor<double, absl::gaussian_distribution>(bitgen, type);
+  }
+  if (elem_ty.isInteger(32)) {
+    return RandomTensor<int32_t, absl::uniform_int_distribution>(bitgen, type);
+  }
+  if (elem_ty.isInteger(16)) {
+    return RandomTensor<int16_t, absl::uniform_int_distribution>(bitgen, type);
+  }
+  if (elem_ty.isInteger(64)) {
+    return RandomTensor<int64_t, absl::uniform_int_distribution>(bitgen, type);
+  }
+  if (elem_ty.isInteger(1)) {
+    return {{TensorOrMemref<bool>::Empty(type.cast<ShapedType>().getShape())}};
+  }
+
+  llvm::errs() << "Unsupported type: ";
+  type.print(llvm::errs());
+  llvm::errs() << "\n";
+  return failure();
+}
+
+// TODO(jreiffers): Add a flag to intentionally alias as many buffers as
+// possible (in particular, all non-variable inputs).
+// Extracts a mapping from function arguments to allocated buffers.
+// The buffer assignment is only relevant once the program is bufferized and
+// memref results were converted to arguments.
+std::vector<int64_t> extractXlaBufferAssignment(func::FuncOp main) {
+  std::vector<int64_t> buffer_assignment(main.getNumArguments());
+  auto result_mapping =
+      main->getAttrOfType<IntegerAttr>("xla_framework.result_mapping");
+  if (!result_mapping) {
+    // No attribute, fall back to unique buffers for each argument.
+    std::iota(buffer_assignment.begin(), buffer_assignment.end(), 0);
+    return buffer_assignment;
+  }
+
+  std::vector<int64_t> result_to_buffer;
+  if (auto inner_mapping = main->getAttrOfType<ArrayAttr>(
+          "xla_framework.result_inner_mapping")) {
+    llvm::copy(llvm::map_range(inner_mapping.getAsValueRange<IntegerAttr>(),
+                               [](const llvm::APInt& value) {
+                                 return value.getSExtValue();
+                               }),
+               std::back_inserter(result_to_buffer));
+  } else {
+    result_to_buffer = {result_mapping.getInt()};
+  }
+
+  int64_t result_index = 0;
+  for (int64_t arg_index : llvm::seq<int64_t>(0, main.getNumArguments())) {
+    if (auto input_buffer_index = main.getArgAttrOfType<IntegerAttr>(
+            arg_index, "xla_framework.input_mapping")) {
+      buffer_assignment[arg_index] = input_buffer_index.getInt();
+    } else {
+      buffer_assignment[arg_index] = result_to_buffer[result_index++];
+    }
+  }
+
+  return buffer_assignment;
+}
+
+}  // namespace
+
+absl::StatusOr<SmallVector<InterpreterValue>> Run(
+    MLIRContext& context, const std::string& mlir_ir,
+    const xla::HloSnapshot& snapshot, ExecutionTrace* trace,
+    const std::vector<std::string>& entry) {
+  auto sourceMgr = std::make_shared<llvm::SourceMgr>();
+  sourceMgr->AddNewSourceBuffer(llvm::MemoryBuffer::getMemBuffer(mlir_ir),
+                                mlir::SMLoc());
+  mlir::OwningOpRef<mlir::Operation*> module =
+      mlir::parseSourceFileForTool(sourceMgr, &context, false);
+  if (!module) {
+    return absl::InvalidArgumentError("failed to parse MLIR");
+  }
+
+  SymbolTable symbols(*module);
+  func::FuncOp main;
+  for (const std::string& candidate : entry) {
+    main = llvm::dyn_cast_or_null<func::FuncOp>(symbols.lookup(candidate));
+    if (main && !main.getBody().empty()) {
+      break;
+    }
+  }
+
+  if (!main) {
+    return absl::InvalidArgumentError("failed to find entry point");
+  }
+
+  if (trace) {
+    llvm::raw_string_ostream os(*trace->mutable_ir());
+    (*module)->print(os, OpPrintingFlags().printGenericOpForm());
+  }
+
+  auto function_args = main.getBody().getBlocks().front().getArguments();
+  if (!llvm::all_of(function_args, [&](Value arg) {
+        return isa<ShapedType>(arg.getType());
+      })) {
+    return absl::InvalidArgumentError(
+        "expected all function arguments to be shaped types");
+  }
+
+  auto args_to_buffers = extractXlaBufferAssignment(main);
+  TF_ASSIGN_OR_RETURN(auto args,
+                      LoadArgs(snapshot, main.getBody().getArgumentTypes()));
+  auto out_args =
+      main.getBody().getBlocks().front().getArguments().drop_front(args.size());
+
+  absl::flat_hash_map<int64_t, InterpreterValue> buffer_to_value;
+  // None of the input arguments will be statically known to alias.
+  for (auto [index, value] : llvm::enumerate(args)) {
+    buffer_to_value[args_to_buffers[index]] = value;
+  }
+
+  std::seed_seq my_seed_seq({0});
+  absl::BitGen bitgen(my_seed_seq);
+  llvm::SmallVector<InterpreterValue> out_buffers;
+  // Add random inputs for output arguments and unspecified inputs.
+  for (auto arg : out_args) {
+    auto ty = arg.getType();
+
+    int64_t buffer_index = args_to_buffers[arg.getArgNumber()];
+    // If we already have a buffer for this argument, use it.
+    if (buffer_to_value.contains(buffer_index)) {
+      auto& value = buffer_to_value[buffer_index];
+      out_buffers.push_back(value);
+      args.push_back(value);
+      continue;
+    }
+
+    auto arg_or = MakeRandomInput(bitgen, ty);
+    if (!succeeded(arg_or)) {
+      return absl::InvalidArgumentError("failed to create input");
+    }
+    out_buffers.push_back(*arg_or);
+    args.push_back(*arg_or);
+    buffer_to_value[buffer_index] = *arg_or;
+  }
+
+  InterpreterOptions options;
+  ExecutionTraceListener tracer(trace);
+  if (trace) {
+    options.listener = &tracer;
+  }
+  TF_ASSIGN_OR_RETURN(auto results,
+                      RunInterpreter(symbols, main, args, options));
+
+  if (results.empty()) {
+    return out_buffers;
+  }
+  return results;
+}
+
+}  // namespace interpreter
+}  // namespace mlir
diff --git a/third_party/xla/xla/mlir/tools/mlir_replay/mlir_replay_lib.h b/third_party/xla/xla/mlir/tools/mlir_replay/mlir_replay_lib.h
new file mode 100644
index 00000000000000..c1ed5b228a53eb
--- /dev/null
+++ b/third_party/xla/xla/mlir/tools/mlir_replay/mlir_replay_lib.h
@@ -0,0 +1,40 @@
+/* Copyright 2022 The OpenXLA Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+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.
+==============================================================================*/
+
+#ifndef XLA_MLIR_TOOLS_MLIR_REPLAY_MLIR_REPLAY_LIB_H_
+#define XLA_MLIR_TOOLS_MLIR_REPLAY_MLIR_REPLAY_LIB_H_
+
+#include <string>
+#include <vector>
+
+#include "absl/status/statusor.h"
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+#include "xla/mlir/tools/mlir_interpreter/framework/interpreter_value.h"
+#include "xla/mlir/tools/mlir_replay/public/execution_trace.pb.h"
+#include "xla/service/hlo.pb.h"
+
+namespace mlir {
+namespace interpreter {
+
+// Runs the given IR on the inputs from `snapshot` and returns the result.
+absl::StatusOr<SmallVector<InterpreterValue>> Run(
+    MLIRContext& context, const std::string& mlir_ir,
+    const xla::HloSnapshot& snapshot, ExecutionTrace* trace,
+    const std::vector<std::string>& entry);
+
+}  // namespace interpreter
+}  // namespace mlir
+
+#endif  // XLA_MLIR_TOOLS_MLIR_REPLAY_MLIR_REPLAY_LIB_H_
diff --git a/third_party/xla/xla/mlir/tools/mlir_replay/public/BUILD b/third_party/xla/xla/mlir/tools/mlir_replay/public/BUILD
new file mode 100644
index 00000000000000..693c7e2f25f3e5
--- /dev/null
+++ b/third_party/xla/xla/mlir/tools/mlir_replay/public/BUILD
@@ -0,0 +1,56 @@
+load("@local_tsl//tsl/platform:build_config.bzl", "tf_proto_library")
+load("@local_tsl//tsl/platform:rules_cc.bzl", "cc_library")
+load("//xla:xla.bzl", "xla_cc_test")
+
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
+    default_visibility = ["//visibility:public"],
+)
+
+cc_library(
+    name = "execution_trace_utils",
+    srcs = ["execution_trace_utils.cc"],
+    hdrs = ["execution_trace_utils.h"],
+    deps = [
+        ":execution_trace_proto_cc",
+        ":execution_trace_proto_cc_impl",
+        "//xla:literal",
+        "//xla:xla_data_proto_cc",
+        "//xla/mlir/tools/mlir_interpreter/framework",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/status:statusor",
+        "@llvm-project//llvm:Support",
+        "@llvm-project//mlir:FuncDialect",
+        "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:Support",
+        "@local_tsl//tsl/platform:statusor",
+    ],
+)
+
+xla_cc_test(
+    name = "execution_trace_utils_test",
+    srcs = ["execution_trace_utils_test.cc"],
+    deps = [
+        ":execution_trace_utils",
+        "//xla:literal_util",
+        "//xla/mlir/tools/mlir_interpreter/framework",
+        "@com_google_googletest//:gtest_main",
+        "@llvm-project//llvm:Support",
+        "@llvm-project//mlir:Support",
+        "@local_tsl//tsl/platform:statusor",
+    ],
+)
+
+tf_proto_library(
+    name = "execution_trace_proto",
+    srcs = ["execution_trace.proto"],
+    cc_api_version = 2,
+    make_default_target_header_only = True,
+)
+
+tf_proto_library(
+    name = "compiler_trace_proto",
+    srcs = ["compiler_trace.proto"],
+    cc_api_version = 2,
+    make_default_target_header_only = True,
+)
diff --git a/third_party/xla/xla/mlir/tools/mlir_replay/public/compiler_trace.proto b/third_party/xla/xla/mlir/tools/mlir_replay/public/compiler_trace.proto
new file mode 100644
index 00000000000000..b8f75dd1ad355f
--- /dev/null
+++ b/third_party/xla/xla/mlir/tools/mlir_replay/public/compiler_trace.proto
@@ -0,0 +1,31 @@
+/* Copyright 2022 The OpenXLA Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+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.
+==============================================================================*/
+
+syntax = "proto2";
+
+package mlir.interpreter;
+
+message MlirCompilationTraceEntry {
+  // The name of the pass that was previously executed.
+  optional string after_pass = 1;
+
+  // MLIR module IR of the state after the pass.
+  optional string mlir_module = 2;
+}
+
+message MlirCompilationTrace {
+  // MLIR modules corresponding to each stage of the compilation pipeline.
+  repeated MlirCompilationTraceEntry passes = 1;
+}
diff --git a/third_party/xla/xla/mlir/tools/mlir_replay/public/execution_trace.proto b/third_party/xla/xla/mlir/tools/mlir_replay/public/execution_trace.proto
new file mode 100644
index 00000000000000..c4ff5ecd4fa939
--- /dev/null
+++ b/third_party/xla/xla/mlir/tools/mlir_replay/public/execution_trace.proto
@@ -0,0 +1,72 @@
+/* Copyright 2022 The OpenXLA Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+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.
+==============================================================================*/
+
+syntax = "proto2";
+
+package mlir.interpreter;
+
+message TracedValue {
+  // The shape - includes vector dimensions.
+  // TODO(jreiffers): Model vector dimensions separately.
+  repeated int64 shape = 1 [packed = true];
+  optional bool is_scalar = 2;
+
+  enum ElementType {
+    UNKNOWN = 0;
+    INTEGRAL = 1;
+    UNSIGNED = 2;
+    FLOAT = 3;
+    COMPLEX = 4;
+    TUPLE = 5;
+  }
+
+  optional int32 bit_width = 3;
+  optional ElementType element_type = 4;
+
+  repeated float floats = 5 [packed = true];
+  repeated double doubles = 6 [packed = true];
+  repeated int64 ints = 7 [packed = true];
+  repeated uint64 uints = 8 [packed = true];
+  repeated TracedValue tuple_elements = 9;
+}
+
+message InstructionTrace {
+  optional string name = 1;
+  repeated TracedValue args = 2;
+  repeated TracedValue results = 3;
+  // TODO(jreiffers): Model side effects (e.g. memref.store).
+
+  repeated RegionTrace regions = 4;
+}
+
+message RegionTrace {
+  // The number of the region that is being executed (within the parent op).
+  // For example: '1' for an scf.while's `after` region.
+  optional int32 region_number = 1;
+  // The arguments that were passed to the region.
+  repeated TracedValue bbargs = 2;
+  // One instruction per instruction in the region.
+  repeated InstructionTrace instructions = 3;
+  repeated TracedValue results = 4;
+}
+
+message ExecutionTrace {
+  // The IR that was executed. Note: this should always be filled in the generic
+  // format.
+  optional string ir = 1;
+
+  // The trace of the entry function execution.
+  optional RegionTrace trace = 2;
+}
diff --git a/third_party/xla/xla/mlir/tools/mlir_replay/public/execution_trace_utils.cc b/third_party/xla/xla/mlir/tools/mlir_replay/public/execution_trace_utils.cc
new file mode 100644
index 00000000000000..36482f40cadd01
--- /dev/null
+++ b/third_party/xla/xla/mlir/tools/mlir_replay/public/execution_trace_utils.cc
@@ -0,0 +1,449 @@
+/* Copyright 2022 The OpenXLA Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+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.
+==============================================================================*/
+
+#include "xla/mlir/tools/mlir_replay/public/execution_trace_utils.h"
+
+#include <complex>
+#include <cstdint>
+#include <functional>
+#include <iterator>
+#include <memory>
+#include <type_traits>
+#include <utility>
+#include <variant>
+
+#include "absl/status/status.h"
+#include "llvm/ADT/STLExtras.h"
+#include "llvm/ADT/SmallVector.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
+#include "mlir/IR/Attributes.h"  // from @llvm-project
+#include "mlir/IR/BuiltinAttributes.h"  // from @llvm-project
+#include "mlir/IR/BuiltinTypes.h"  // from @llvm-project
+#include "xla/mlir/tools/mlir_interpreter/framework/interpreter_value.h"
+#include "xla/mlir/tools/mlir_interpreter/framework/tensor_or_memref.h"
+#include "xla/mlir/tools/mlir_replay/public/execution_trace.pb.h"
+#include "tsl/platform/statusor.h"
+
+namespace mlir {
+namespace interpreter {
+namespace {
+
+// Visitor for converting an InterpreterValue to a TracedValue.
+struct TraceInterpreterValueVisitor {
+  TracedValue out;
+
+  void Add(float v) { out.add_floats(v); }
+  void Add(double v) { out.add_doubles(v); }
+  void Add(std::complex<float> v) {
+    out.add_floats(v.real());
+    out.add_floats(v.imag());
+  }
+  void Add(std::complex<double> v) {
+    out.add_doubles(v.real());
+    out.add_doubles(v.imag());
+  }
+  void Add(int64_t v) { out.add_ints(v); }
+  void Add(int32_t v) { out.add_ints(v); }
+  void Add(int16_t v) { out.add_ints(v); }
+  void Add(int8_t v) { out.add_ints(v); }
+  void Add(uint64_t v) { out.add_uints(v); }
+  void Add(uint32_t v) { out.add_uints(v); }
+  void Add(uint16_t v) { out.add_uints(v); }
+  void Add(uint8_t v) { out.add_uints(v); }
+  void Add(bool v) { out.add_ints(static_cast<int64_t>(v)); }
+
+  template <typename T>
+  void operator()(T v) {
+    SetElementType<T>();
+    out.set_is_scalar(true);
+    Add(v);
+  }
+
+  void operator()(const Tuple& t) {
+    out.set_element_type(TracedValue::TUPLE);
+    for (const auto& v : t.values) {
+      *out.add_tuple_elements() = ValueToTracedValue(*v);
+    }
+  }
+
+  template <typename T>
+  void operator()(const TensorOrMemref<T>& v) {
+    for (int64_t size : v.view.sizes) {
+      out.add_shape(size);
+    }
+    SetElementType<T>();
+    for (const auto& index : v.view.Indices()) {
+      Add(v.at(index));
+    }
+  }
+
+  template <typename T>
+  void SetElementType() {
+    out.set_element_type(GetElementType(T{}));
+    if constexpr (std::is_same_v<T, bool>) {
+      out.set_bit_width(1);
+    } else {
+      out.set_bit_width(sizeof(T) * 8);
+    }
+  }
+
+  template <typename T>
+  static TracedValue::ElementType GetElementType(const T&) {
+    if constexpr (std::is_floating_point_v<T>) {
+      return TracedValue::FLOAT;
+    } else if constexpr (std::is_integral_v<T>) {
+      if constexpr (std::is_unsigned_v<T>) {
+        return TracedValue::UNSIGNED;
+      } else {
+        return TracedValue::INTEGRAL;
+      }
+    } else {
+      T{"invalid type"} + 0;
+      return TracedValue::UNKNOWN;
+    }
+  }
+
+  template <typename T>
+  static TracedValue::ElementType GetElementType(const std::complex<T>&) {
+    return TracedValue::COMPLEX;
+  }
+
+  static TracedValue::ElementType GetElementType(const Tuple&) {
+    return TracedValue::UNKNOWN;
+  }
+};
+
+}  // namespace
+
+void ExecutionTraceListener::BeforeOp(ArrayRef<InterpreterValue> args,
+                                      Operation* op) {
+  auto* inst = regions_.back()->add_instructions();
+  inst->set_name(op->getName().getStringRef().str());
+  for (const auto& arg : args) {
+    *inst->add_args() = ValueToTracedValue(arg);
+  }
+}
+
+void ExecutionTraceListener::AfterOp(ArrayRef<InterpreterValue> results) {
+  auto* traced_results =
+      regions_.back()->mutable_instructions()->rbegin()->mutable_results();
+  for (const auto& result : results) {
+    *traced_results->Add() = ValueToTracedValue(result);
+  }
+}
+
+void ExecutionTraceListener::EnterRegion(ArrayRef<InterpreterValue> bbargs,
+                                         Region& region) {
+  if (regions_.empty()) {
+    regions_.push_back(trace_->mutable_trace());
+  } else {
+    regions_.push_back(
+        regions_.back()->mutable_instructions()->rbegin()->add_regions());
+  }
+
+  auto& traced_region = *regions_.back();
+  traced_region.set_region_number(region.getRegionNumber());
+  for (const auto& bbarg : bbargs) {
+    *traced_region.add_bbargs() = ValueToTracedValue(bbarg);
+  }
+}
+
+void ExecutionTraceListener::LeaveRegion(ArrayRef<InterpreterValue> yielded) {
+  for (const auto& result : yielded) {
+    *regions_.back()->add_results() = ValueToTracedValue(result);
+  }
+  regions_.pop_back();
+}
+
+llvm::SmallVector<mlir::Attribute> ValueToAttribute(
+    const InterpreterValue& value, mlir::Type type) {
+  if (std::holds_alternative<Tuple>(value.storage)) {
+    auto types = type.cast<TupleType>().getTypes();
+    const auto& t = std::get<Tuple>(value.storage);
+    llvm::SmallVector<mlir::Attribute> attrs;
+    for (const auto& [v, ty] : llvm::zip(t.values, types)) {
+      auto attr = ValueToAttribute(*v, ty);
+      assert(attr.size() == 1 && "nested tuples not supported");
+      attrs.push_back(attr.front());
+    }
+    return attrs;
+  }
+
+  if (!value.IsTensor()) {
+    return {cast<DenseElementsAttr>(
+                ValueToAttribute(value.AsUnitTensor(),
+                                 mlir::RankedTensorType::get({}, type))
+                    .front())
+                .getValues<mlir::Attribute>()[0]};
+  }
+
+  if (!type.isa<ShapedType>()) {
+    return {};
+  }
+
+  auto shaped_ty = type.cast<ShapedType>();
+  return {DispatchScalarType(shaped_ty, [&](auto dummy) -> mlir::Attribute {
+    using T = decltype(dummy);
+    auto& t = std::get<TensorOrMemref<T>>(value.storage);
+    SmallVector<T> vals;
+    for (const auto& index : t.view.Indices()) {
+      vals.push_back(t.at(index));
+    }
+    auto attr_ty =
+        shaped_ty.cloneWith(/*shape=*/t.view.sizes, shaped_ty.getElementType());
+    if constexpr (std::is_same_v<T, bool>) {
+      return mlir::DenseElementsAttr::get(attr_ty, vals);
+    } else {
+      return mlir::DenseElementsAttr::get<T>(attr_ty, vals);
+    }
+  })};
+}
+
+namespace {
+template <typename T>
+TensorOrMemref<T> ArrayLiteralToTensor(const xla::Literal& literal) {
+  SmallVector<int64_t> layout;
+  if (literal.shape().has_layout()) {
+    llvm::copy(literal.shape().layout().minor_to_major(),
+               std::back_inserter(layout));
+  }
+  SmallVector<int64_t> shape{literal.shape().dimensions().begin(),
+                             literal.shape().dimensions().end()};
+  auto result = TensorOrMemref<T>::Empty(shape, layout);
+  assert(literal.size_bytes() == result.buffer->GetByteSize() &&
+         "expected buffer sizes to match");
+  memcpy(result.buffer->at(0, 0), literal.untyped_data(),
+         result.buffer->GetByteSize());
+  return result;
+}
+}  // namespace
+
+absl::StatusOr<InterpreterValue> LiteralToValue(const xla::Literal& literal) {
+  if (literal.shape().IsTuple()) {
+    auto elements = literal.Clone().DecomposeTuple();
+    Tuple result;
+    for (auto& element : elements) {
+      TF_ASSIGN_OR_RETURN(auto converted, LiteralToValue(element));
+      result.values.push_back(
+          std::make_shared<InterpreterValue>(std::move(converted)));
+    }
+    return {{result}};
+  }
+
+  if (literal.shape().IsToken()) {
+    return absl::UnimplementedError("token arguments are not implemented");
+  }
+
+  if (literal.shape().IsArray()) {
+    switch (literal.shape().element_type()) {
+      case xla::PRED:
+        return {{ArrayLiteralToTensor<bool>(literal)}};
+      case xla::S8:
+        return {{ArrayLiteralToTensor<int8_t>(literal)}};
+      case xla::S16:
+        return {{ArrayLiteralToTensor<int16_t>(literal)}};
+      case xla::S32:
+        return {{ArrayLiteralToTensor<int32_t>(literal)}};
+      case xla::S64:
+        return {{ArrayLiteralToTensor<int64_t>(literal)}};
+      case xla::U8:
+        return {{ArrayLiteralToTensor<uint8_t>(literal)}};
+      case xla::U16:
+        return {{ArrayLiteralToTensor<uint16_t>(literal)}};
+      case xla::U32:
+        return {{ArrayLiteralToTensor<uint32_t>(literal)}};
+      case xla::U64:
+        return {{ArrayLiteralToTensor<uint64_t>(literal)}};
+      case xla::F16:
+        return absl::UnimplementedError("F16 not implemented");
+      case xla::F32:
+        return {{ArrayLiteralToTensor<float>(literal)}};
+      case xla::BF16:
+        return absl::UnimplementedError("BF16 not implemented");
+      case xla::F64:
+        return {{ArrayLiteralToTensor<double>(literal)}};
+      case xla::F8E5M2:
+        return absl::UnimplementedError("F8E5M2 not implemented");
+      case xla::F8E4M3FN:
+        return absl::UnimplementedError("F8E4M3FN not implemented");
+      case xla::F8E4M3B11FNUZ:
+        return absl::UnimplementedError("F8E4M3B11FNUZ not implemented");
+      case xla::F8E5M2FNUZ:
+        return absl::UnimplementedError("F8E5M2FNUZ not implemented");
+      case xla::F8E4M3FNUZ:
+        return absl::UnimplementedError("F8E4M3FNUZ not implemented");
+      case xla::C64:
+        return {{ArrayLiteralToTensor<std::complex<float>>(literal)}};
+      case xla::C128:
+        return {{ArrayLiteralToTensor<std::complex<double>>(literal)}};
+      default:
+        // Fallthrough intended.
+        break;
+    }
+  }
+
+  return absl::InvalidArgumentError("unexpected literal type");
+}
+
+absl::StatusOr<InterpreterValue> LiteralToValue(
+    const xla::LiteralProto& literal) {
+  TF_ASSIGN_OR_RETURN(auto deserialized,
+                      xla::Literal::CreateFromProto(literal));
+  return LiteralToValue(deserialized);
+}
+
+absl::StatusOr<InterpreterValue> LiteralToValue(
+    const xla::LiteralProto& literal, mlir::Type type) {
+  TF_ASSIGN_OR_RETURN(auto result, LiteralToValue(literal));
+  return {DispatchScalarType(type, [&](auto dummy) -> InterpreterValue {
+    TensorOrMemref<decltype(dummy)> cast;
+    cast.view = result.View();
+    cast.buffer = result.Buffer();
+    return {cast};
+  })};
+}
+
+TracedValue ValueToTracedValue(const InterpreterValue& value) {
+  TraceInterpreterValueVisitor visitor;
+  std::visit(visitor, value.storage);
+  return visitor.out;
+}
+
+absl::StatusOr<InterpreterValue> TracedValueToValue(
+    const TracedValue& traced_value) {
+  auto extract = [&](auto dummy, auto& elements) -> InterpreterValue {
+    using T = decltype(dummy);
+    if (traced_value.is_scalar()) {
+      return {static_cast<T>(elements[0])};
+    }
+
+    auto result =
+        TensorOrMemref<T>::Empty(llvm::to_vector(traced_value.shape()));
+    for (auto [index, element] : llvm::zip(result.view.Indices(), elements)) {
+      result.at(index) = element;
+    }
+    return {result};
+  };
+  auto extract_complex = [&](auto& elements) -> InterpreterValue {
+    using T = std::complex<std::decay_t<decltype(elements[0])>>;
+    if (traced_value.is_scalar()) {
+      return {T{elements[0], elements[1]}};
+    }
+
+    auto result =
+        TensorOrMemref<T>::Empty(llvm::to_vector(traced_value.shape()));
+    int64_t i = 0;
+    for (auto it = result.view.Indices().begin(),
+              end = result.view.Indices().end();
+         it != end; ++it, i += 2) {
+      result.at(*it) = {elements[i], elements[i + 1]};
+    }
+    return {result};
+  };
+  switch (traced_value.element_type()) {
+    case TracedValue::UNKNOWN:
+      break;
+    case TracedValue::FLOAT:
+      if (traced_value.bit_width() == 32) {
+        return extract(float{}, traced_value.floats());
+      }
+      return extract(double{}, traced_value.doubles());
+    case TracedValue::UNSIGNED:
+      switch (traced_value.bit_width()) {
+        case 1:
+          return extract(bool{}, traced_value.ints());
+        case 8:
+          return extract(uint8_t{}, traced_value.uints());
+        case 16:
+          return extract(uint16_t{}, traced_value.uints());
+        case 32:
+          return extract(uint32_t{}, traced_value.uints());
+        case 64:
+          return extract(uint64_t{}, traced_value.uints());
+      }
+      break;
+    case TracedValue::INTEGRAL:
+      switch (traced_value.bit_width()) {
+        case 8:
+          return extract(int8_t{}, traced_value.ints());
+        case 16:
+          return extract(int16_t{}, traced_value.ints());
+        case 32:
+          return extract(int32_t{}, traced_value.ints());
+        case 64:
+          return extract(int64_t{}, traced_value.ints());
+      }
+      break;
+    case TracedValue::COMPLEX:
+      switch (traced_value.bit_width()) {
+        case 64:
+          return extract_complex(traced_value.floats());
+        case 128:
+          return extract_complex(traced_value.doubles());
+      }
+      break;
+    case TracedValue::TUPLE:
+      Tuple result;
+      for (const auto& elem : traced_value.tuple_elements()) {
+        TF_ASSIGN_OR_RETURN(auto converted, TracedValueToValue(elem));
+        result.values.push_back(
+            std::make_shared<InterpreterValue>(std::move(converted)));
+      }
+      return {{std::move(result)}};
+  }
+  return absl::InvalidArgumentError("unexpected type: " +
+                                    traced_value.DebugString());
+}
+
+llvm::SmallVector<const InstructionTrace*> FindOpExecutionsInTrace(
+    const ExecutionTrace& trace, mlir::Operation* op) {
+  llvm::SmallVector<int64_t> region_indices;
+  llvm::SmallVector<int64_t> op_indices;
+
+  std::function<void(mlir::Operation*)> get_op_path;
+  get_op_path = [&](mlir::Operation* op) {
+    auto* parent = op->getParentOp();
+    if (!llvm::isa<func::FuncOp>(parent)) {
+      get_op_path(parent);
+      region_indices.push_back(op->getParentRegion()->getRegionNumber());
+    }
+
+    int64_t index = 0;
+    while ((op = op->getPrevNode()) != nullptr) ++index;
+    op_indices.push_back(index);
+  };
+  get_op_path(op);
+
+  llvm::SmallVector<const InstructionTrace*> result;
+  std::function<void(const RegionTrace& trace, int index)> step;
+  step = [&](const RegionTrace& trace, int index) {
+    auto& instruction_trace = trace.instructions(op_indices[index]);
+    if (region_indices.size() > index) {
+      for (const auto& region : instruction_trace.regions()) {
+        if (region.region_number() == region_indices[index]) {
+          step(region, index + 1);
+        }
+      }
+    } else {
+      result.push_back(&instruction_trace);
+    }
+  };
+  step(trace.trace(), 0);
+
+  return result;
+}
+
+}  // namespace interpreter
+}  // namespace mlir
diff --git a/third_party/xla/xla/mlir/tools/mlir_replay/public/execution_trace_utils.h b/third_party/xla/xla/mlir/tools/mlir_replay/public/execution_trace_utils.h
new file mode 100644
index 00000000000000..4c4c36e389795e
--- /dev/null
+++ b/third_party/xla/xla/mlir/tools/mlir_replay/public/execution_trace_utils.h
@@ -0,0 +1,76 @@
+/* Copyright 2022 The OpenXLA Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+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.
+==============================================================================*/
+
+#ifndef XLA_MLIR_TOOLS_MLIR_REPLAY_PUBLIC_EXECUTION_TRACE_UTILS_H_
+#define XLA_MLIR_TOOLS_MLIR_REPLAY_PUBLIC_EXECUTION_TRACE_UTILS_H_
+
+#include "absl/status/statusor.h"
+#include "mlir/IR/Attributes.h"  // from @llvm-project
+#include "mlir/IR/Operation.h"  // from @llvm-project
+#include "mlir/IR/Region.h"  // from @llvm-project
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+#include "xla/literal.h"
+#include "xla/mlir/tools/mlir_interpreter/framework/interpreter.h"
+#include "xla/mlir/tools/mlir_interpreter/framework/interpreter_value.h"
+#include "xla/mlir/tools/mlir_replay/public/execution_trace.pb.h"
+#include "xla/xla_data.pb.h"
+
+namespace mlir {
+namespace interpreter {
+
+// Interpreter listener that builds a trace of all executed ops and regions.
+class ExecutionTraceListener : public InterpreterListener {
+ public:
+  explicit ExecutionTraceListener(ExecutionTrace* trace) : trace_(trace) {}
+
+  void BeforeOp(ArrayRef<InterpreterValue> args, Operation* op) override;
+  void AfterOp(ArrayRef<InterpreterValue> results) override;
+  void EnterRegion(ArrayRef<InterpreterValue> bbargs, Region& region) override;
+  void LeaveRegion(ArrayRef<InterpreterValue> yielded) override;
+
+ private:
+  ExecutionTrace* trace_;
+  SmallVector<RegionTrace*> regions_;
+};
+
+// Returns an attribute with the given contents and type.
+llvm::SmallVector<mlir::Attribute> ValueToAttribute(
+    const InterpreterValue& value, mlir::Type type);
+
+// Deserializes the given literal.
+absl::StatusOr<InterpreterValue> LiteralToValue(
+    const xla::LiteralProto& literal);
+// Deserializes the given literal and then casts it to the given type.
+absl::StatusOr<InterpreterValue> LiteralToValue(
+    const xla::LiteralProto& literal, mlir::Type type);
+
+// Deserializes the given literal.
+absl::StatusOr<InterpreterValue> LiteralToValue(const xla::Literal& literal);
+
+// Serializes the given interpreter value.
+TracedValue ValueToTracedValue(const InterpreterValue& value);
+
+// Deserializes the given traced value.
+absl::StatusOr<InterpreterValue> TracedValueToValue(
+    const TracedValue& traced_value);
+
+// Returns all executions of the given op in the given trace.
+llvm::SmallVector<const InstructionTrace*> FindOpExecutionsInTrace(
+    const ExecutionTrace& trace, mlir::Operation* op);
+
+}  // namespace interpreter
+}  // namespace mlir
+
+#endif  // XLA_MLIR_TOOLS_MLIR_REPLAY_PUBLIC_EXECUTION_TRACE_UTILS_H_
diff --git a/third_party/xla/xla/mlir/tools/mlir_replay/public/execution_trace_utils_test.cc b/third_party/xla/xla/mlir/tools/mlir_replay/public/execution_trace_utils_test.cc
new file mode 100644
index 00000000000000..31a6ae074dea48
--- /dev/null
+++ b/third_party/xla/xla/mlir/tools/mlir_replay/public/execution_trace_utils_test.cc
@@ -0,0 +1,139 @@
+/* Copyright 2022 The OpenXLA Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+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.
+==============================================================================*/
+
+#include "xla/mlir/tools/mlir_replay/public/execution_trace_utils.h"
+
+#include <cmath>
+#include <complex>
+#include <cstdint>
+#include <memory>
+#include <utility>
+#include <vector>
+
+#include <gtest/gtest.h>
+#include "llvm/ADT/STLExtras.h"
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+#include "xla/literal_util.h"
+#include "xla/mlir/tools/mlir_interpreter/framework/interpreter_value.h"
+#include "xla/mlir/tools/mlir_interpreter/framework/tensor_or_memref.h"
+#include "tsl/platform/statusor.h"
+
+namespace mlir {
+namespace interpreter {
+namespace {
+
+class TracedValueRoundTripTest
+    : public ::testing::TestWithParam<InterpreterValue> {};
+
+TEST_P(TracedValueRoundTripTest, Run) {
+  auto traced_value = ValueToTracedValue(GetParam());
+  TF_ASSERT_OK_AND_ASSIGN(auto value, TracedValueToValue(traced_value));
+  EXPECT_EQ(GetParam(), value) << GetParam().ToString();
+}
+
+template <typename T>
+InterpreterValue MakeTensor(ArrayRef<int64_t> shape, ArrayRef<T> values) {
+  auto result = TensorOrMemref<T>::Empty(shape);
+  for (auto [indices, value] : llvm::zip(result.view.Indices(), values)) {
+    result.at(indices) = value;
+  }
+  return {result};
+}
+
+template <typename T>
+std::shared_ptr<T> WrapShared(T value) {
+  return std::make_shared<T>(std::move(value));
+}
+
+INSTANTIATE_TEST_SUITE_P(
+    RoundTrip, TracedValueRoundTripTest,
+    ::testing::ValuesIn(std::vector<InterpreterValue>{
+        {uint8_t{42}},
+        {uint16_t{43}},
+        {uint32_t{44}},
+        {uint64_t{45}},
+        {int8_t{-47}},
+        {int16_t{-48}},
+        {int32_t{-49}},
+        {int64_t{-50}},
+        {float{42.0}},
+        {double{42.0}},
+        {std::complex<float>{1.0, 2.0}},
+        {std::complex<double>{3.0, 4.0}},
+        {true},
+        {false},
+        {MakeTensor<int16_t>({1, 2}, {42, 43})},
+        {MakeTensor<double>({2, 2}, {1.0, -INFINITY, INFINITY, NAN})},
+        {MakeTensor<std::complex<double>>({}, {{1.0, 2.0}})},
+        {Tuple{SmallVector<std::shared_ptr<InterpreterValue>>{
+            WrapShared(InterpreterValue{42}),
+            WrapShared(InterpreterValue{43.0}),
+        }}}}));
+
+class FromLiteralTest
+    : public ::testing::TestWithParam<
+          std::pair<std::shared_ptr<xla::Literal>, InterpreterValue>> {};
+
+TEST_P(FromLiteralTest, Run) {
+  TF_ASSERT_OK_AND_ASSIGN(auto value, LiteralToValue(*GetParam().first));
+  EXPECT_EQ(value, GetParam().second)
+      << value.ToString() << " vs " << GetParam().second.ToString();
+}
+
+std::vector<std::pair<std::shared_ptr<xla::Literal>, InterpreterValue>>
+MakeInputs() {
+  using ::xla::LiteralUtil;
+  return {
+      {WrapShared(LiteralUtil::CreateR2<uint8_t>({{41, 42}})),
+       MakeTensor<uint8_t>({1, 2}, {41, 42})},
+      {WrapShared(LiteralUtil::CreateR0<uint16_t>(43)),
+       MakeTensor<uint16_t>({}, {43})},
+      {WrapShared(LiteralUtil::CreateR0<uint32_t>(44)),
+       MakeTensor<uint32_t>({}, {44})},
+      {WrapShared(LiteralUtil::CreateR0<uint64_t>(45)),
+       MakeTensor<uint64_t>({}, {45})},
+      {WrapShared(LiteralUtil::CreateR0<int8_t>(46)),
+       MakeTensor<int8_t>({}, {46})},
+      {WrapShared(LiteralUtil::CreateR0<int16_t>(47)),
+       MakeTensor<int16_t>({}, {47})},
+      {WrapShared(LiteralUtil::CreateR0<int32_t>(48)),
+       MakeTensor<int32_t>({}, {48})},
+      {WrapShared(LiteralUtil::CreateR0<int64_t>(49)),
+       MakeTensor<int64_t>({}, {49})},
+      {WrapShared(LiteralUtil::CreateR0<float>(50.0)),
+       MakeTensor<float>({}, {50.0})},
+      {WrapShared(LiteralUtil::CreateR0<double>(51.0)),
+       MakeTensor<double>({}, {51.0})},
+      {WrapShared(LiteralUtil::CreateR0<std::complex<float>>({52.0, 53.0})),
+       MakeTensor<std::complex<float>>({}, {{52.0, 53.0}})},
+      {WrapShared(LiteralUtil::CreateR0<std::complex<double>>({54.0, 55.0})),
+       MakeTensor<std::complex<double>>({}, {{54.0, 55.0}})},
+      {WrapShared(LiteralUtil::CreateR1<bool>({true, false})),
+       MakeTensor<bool>({2}, {true, false})},
+      {WrapShared(
+           LiteralUtil::MakeTupleOwned(LiteralUtil::CreateR0<bool>(true),
+                                       LiteralUtil::CreateR0<int8_t>(56))),
+       InterpreterValue{Tuple{SmallVector<std::shared_ptr<InterpreterValue>>{
+           std::make_shared<InterpreterValue>(MakeTensor<bool>({}, {true})),
+           std::make_shared<InterpreterValue>(
+               MakeTensor<int8_t>({}, {56}))}}}}};
+}
+
+INSTANTIATE_TEST_SUITE_P(Test, FromLiteralTest,
+                         ::testing::ValuesIn(MakeInputs()));
+
+}  // namespace
+}  // namespace interpreter
+}  // namespace mlir

From 3007c44f8911b61f8e20b7d64bd216fd0878545c Mon Sep 17 00:00:00 2001
From: ekuznetsov139 <nameless@fastmail.fm>
Date: Wed, 15 May 2024 05:58:51 -0700
Subject: [PATCH 177/478] PR #9666: [ROCm] Fused convolution+bias+activation

Imported from GitHub PR https://github.com/openxla/xla/pull/9666

This PR enables MIOpen convolution+bias+activation fusion for ROCm and updates fusion unit tests accordingly.
Copybara import of the project:

--
bd5be494abe6621dfd2c4ebcca4f8992077d8a89 by Eugene Kuznetsov <eugene.kuznetsov@amd.com>:

Switch to NHWC for ROCm and F16

--
8792f892770aba54e38a5352d9bec1d003e341d3 by Eugene Kuznetsov <eugene.kuznetsov@amd.com>:

Fused convolution+bias+activation

Merging this change closes #9666

PiperOrigin-RevId: 633923105
---
 third_party/xla/xla/service/gpu/BUILD         |    1 +
 .../xla/xla/service/gpu/amdgpu_compiler.cc    |    3 +
 .../service/gpu/cudnn_fused_conv_rewriter.cc  |   34 +-
 .../service/gpu/cudnn_fused_conv_rewriter.h   |    4 +-
 .../gpu/cudnn_fused_conv_rewriter_test.cc     |  184 +-
 .../xla/service/gpu/gpu_layout_assignment.cc  |    6 +
 .../service/gpu/runtime/convolution_thunk.cc  |    5 +-
 third_party/xla/xla/stream_executor/dnn.cc    |    3 +-
 third_party/xla/xla/stream_executor/dnn.h     |    9 +-
 .../xla/xla/stream_executor/rocm/rocm_dnn.cc  | 1640 +++++++++++------
 .../xla/xla/stream_executor/rocm/rocm_dnn.h   |   83 +-
 .../stream_executor/rocm/rocm_helpers.cu.cc   |   91 +-
 third_party/xla/xla/tests/convolution_test.cc |   45 +-
 13 files changed, 1378 insertions(+), 730 deletions(-)

diff --git a/third_party/xla/xla/service/gpu/BUILD b/third_party/xla/xla/service/gpu/BUILD
index 3883a34d7c4b6e..4dc41fcce52caf 100644
--- a/third_party/xla/xla/service/gpu/BUILD
+++ b/third_party/xla/xla/service/gpu/BUILD
@@ -3980,6 +3980,7 @@ cc_library(
         ":conv_algorithm_picker",
         ":cublas_pad_for_gemms",
         ":cublas_padding_requirements",
+        ":cudnn_fused_conv_rewriter",
         ":cusolver_rewriter",
         ":gemm_algorithm_picker",
         ":gpu_compiler",
diff --git a/third_party/xla/xla/service/gpu/amdgpu_compiler.cc b/third_party/xla/xla/service/gpu/amdgpu_compiler.cc
index f8621e76610000..d9eb60b3046658 100644
--- a/third_party/xla/xla/service/gpu/amdgpu_compiler.cc
+++ b/third_party/xla/xla/service/gpu/amdgpu_compiler.cc
@@ -34,6 +34,7 @@ limitations under the License.
 #include "xla/service/gpu/conv_algorithm_picker.h"
 #include "xla/service/gpu/cublas_pad_for_gemms.h"
 #include "xla/service/gpu/cublas_padding_requirements.h"
+#include "xla/service/gpu/cudnn_fused_conv_rewriter.h"
 #include "xla/service/gpu/cusolver_rewriter.h"
 #include "xla/service/gpu/gemm_algorithm_picker.h"
 #include "xla/service/gpu/gpu_compiler.h"
@@ -117,6 +118,8 @@ absl::Status AMDGPUCompiler::OptimizeHloConvolutionCanonicalization(
   pipeline.AddPass<GpusolverRewriter>();
   pipeline.AddPass<GpuConvRewriter>();
   pipeline.AddPass<GpuConvPaddingLegalization>();
+  auto rcc = std::get<se::RocmComputeCapability>(gpu_version);
+  pipeline.AddPass<CudnnFusedConvRewriter>(rcc);
 
   // The conv padding/vectorization passes which we need to get rid of.  They
   // also leave behind unnecessary tuple/get-tuple-element pairs that
diff --git a/third_party/xla/xla/service/gpu/cudnn_fused_conv_rewriter.cc b/third_party/xla/xla/service/gpu/cudnn_fused_conv_rewriter.cc
index cb8e0ec95e3b92..43226313cd7521 100644
--- a/third_party/xla/xla/service/gpu/cudnn_fused_conv_rewriter.cc
+++ b/third_party/xla/xla/service/gpu/cudnn_fused_conv_rewriter.cc
@@ -24,6 +24,7 @@ limitations under the License.
 #include <string>
 #include <tuple>
 #include <utility>
+#include <variant>
 #include <vector>
 
 #include "absl/algorithm/container.h"
@@ -96,6 +97,10 @@ bool IsNonDepthwiseConvCustomCall(const HloInstruction* instr) {
   return IsConvCustomCall(instr) && !IsConvDepthwise(instr);
 }
 
+bool IsROCm(se::GpuComputeCapability cc) {
+  return std::holds_alternative<se::RocmComputeCapability>(cc);
+}
+
 // elu, relu6, and leaky-relu activations are supported in cudnn via the
 // "runtime fusion" engine, which JIT compiles C++ code.  This can be slow to
 // compile, so we guard it with a debug option.
@@ -106,8 +111,12 @@ bool IsNonDepthwiseConvCustomCall(const HloInstruction* instr) {
 // Note that as of writing, xla_gpu_use_runtime_fusion is disabled by default
 // due to apparent bugs in cudnn 8.9.0.  See debug_options_flags.cc for details.
 bool ShouldUseCudnnRuntimeFusion(const DebugOptions& debug_opts,
-                                 se::CudaComputeCapability cc) {
-  return debug_opts.xla_gpu_use_runtime_fusion() && cc.IsAtLeast(7, 5);
+                                 se::GpuComputeCapability cc) {
+  const auto* cuda_cc = std::get_if<se::CudaComputeCapability>(&cc);
+  if (cuda_cc != nullptr)
+    return debug_opts.xla_gpu_use_runtime_fusion() && cuda_cc->IsAtLeast(7, 5);
+  else
+    return true;
 }
 
 bool IsSuitableForCudnnRuntimeFusion(HloInstruction* conv) {
@@ -984,7 +993,7 @@ absl::StatusOr<bool> FuseSideInputAlpha(HloComputation* comp) {
 }
 
 absl::StatusOr<bool> FuseElu(HloComputation* comp,
-                             se::CudaComputeCapability cc) {
+                             se::GpuComputeCapability cc) {
   if (!ShouldUseCudnnRuntimeFusion(comp->parent()->config().debug_options(),
                                    cc)) {
     return false;
@@ -1085,7 +1094,7 @@ absl::StatusOr<bool> FuseRelu(HloComputation* comp) {
 }
 
 absl::StatusOr<bool> FuseRelu6(HloComputation* comp,
-                               se::CudaComputeCapability cc) {
+                               se::GpuComputeCapability cc) {
   if (!ShouldUseCudnnRuntimeFusion(comp->parent()->config().debug_options(),
                                    cc)) {
     return false;
@@ -1134,7 +1143,7 @@ absl::StatusOr<bool> FuseRelu6(HloComputation* comp,
 }
 
 absl::StatusOr<bool> FuseLeakyRelu(HloComputation* comp,
-                                   se::CudaComputeCapability cc) {
+                                   se::GpuComputeCapability cc) {
   if (!ShouldUseCudnnRuntimeFusion(comp->parent()->config().debug_options(),
                                    cc)) {
     return false;
@@ -1254,7 +1263,9 @@ absl::StatusOr<bool> FuseConvertToF16(HloComputation* comp) {
   return changed;
 }
 
-absl::StatusOr<bool> FuseConvertToS8(HloComputation* comp) {
+absl::StatusOr<bool> FuseConvertToS8(HloComputation* comp,
+                                     se::GpuComputeCapability cc) {
+  if (IsROCm(cc)) return false;
   bool changed = false;
   for (HloInstruction* instr : comp->MakeInstructionPostOrder()) {
     HloInstruction* gte = nullptr;
@@ -1480,9 +1491,12 @@ absl::StatusOr<bool> CudnnFusedConvRewriter::Run(
     bool changed = false;
     // Rewrite FP8 convolutions and supported adjacent pointwise ops into a
     // ForwardGraph Custom Call.
-    TF_ASSIGN_OR_RETURN(changed, F8GraphConv(comp, compute_capability_));
-    if (changed) {
-      return changed;
+    if (!IsROCm(compute_capability_)) {
+      auto cc = std::get<se::CudaComputeCapability>(compute_capability_);
+      TF_ASSIGN_OR_RETURN(changed, F8GraphConv(comp, cc));
+      if (changed) {
+        return changed;
+      }
     }
     // Fuse "inside out" starting with the operations closest to the conv.
     TF_ASSIGN_OR_RETURN(changed, FuseRemoveConvertInConv(comp));
@@ -1516,7 +1530,7 @@ absl::StatusOr<bool> CudnnFusedConvRewriter::Run(
     TF_ASSIGN_OR_RETURN(changed, FuseConvertToF16(comp));
     any_changed |= changed;
 
-    TF_ASSIGN_OR_RETURN(changed, FuseConvertToS8(comp));
+    TF_ASSIGN_OR_RETURN(changed, FuseConvertToS8(comp, compute_capability_));
     any_changed |= changed;
 
     // f16 convs' bias+side-input can appear before or after conversion to f16.
diff --git a/third_party/xla/xla/service/gpu/cudnn_fused_conv_rewriter.h b/third_party/xla/xla/service/gpu/cudnn_fused_conv_rewriter.h
index bc7291d262a616..ff1d1565255399 100644
--- a/third_party/xla/xla/service/gpu/cudnn_fused_conv_rewriter.h
+++ b/third_party/xla/xla/service/gpu/cudnn_fused_conv_rewriter.h
@@ -100,6 +100,8 @@ class CudnnFusedConvRewriter : public HloModulePass {
  public:
   explicit CudnnFusedConvRewriter(se::CudaComputeCapability cc)
       : compute_capability_(cc) {}
+  explicit CudnnFusedConvRewriter(se::RocmComputeCapability cc)
+      : compute_capability_(cc) {}
 
   absl::string_view name() const override {
     return "cudnn-fused-convolution-rewriter";
@@ -111,7 +113,7 @@ class CudnnFusedConvRewriter : public HloModulePass {
       const absl::flat_hash_set<absl::string_view>& execution_threads) override;
 
  private:
-  const se::CudaComputeCapability compute_capability_;
+  const se::GpuComputeCapability compute_capability_;
 };
 
 }  // namespace gpu
diff --git a/third_party/xla/xla/service/gpu/cudnn_fused_conv_rewriter_test.cc b/third_party/xla/xla/service/gpu/cudnn_fused_conv_rewriter_test.cc
index 4a55b9dc9eb403..a1862bf2074d01 100644
--- a/third_party/xla/xla/service/gpu/cudnn_fused_conv_rewriter_test.cc
+++ b/third_party/xla/xla/service/gpu/cudnn_fused_conv_rewriter_test.cc
@@ -19,7 +19,10 @@ limitations under the License.
 #include <memory>
 #include <string>
 #include <string_view>
+#include <thread>  // NOLINT
 #include <utility>
+#include <variant>
+#include <vector>
 
 #include <gmock/gmock.h>
 #include <gtest/gtest.h>
@@ -68,8 +71,19 @@ namespace m = match;
 using ::testing::HasSubstr;
 using ::testing::Not;
 
+// TODO: Use constexpr vector once XLA is compiled with C++20.
+const auto* kf16f32f64 = new std::vector<std::string>({"f16", "f32", "f64"});
+const auto* kf16f32 = new std::vector<std::string>({"f16", "f32"});
+
 class CudnnFusedConvRewriterHloTest : public HloTestBase {
  public:
+  bool IsCuda() {
+    return std::holds_alternative<se::CudaComputeCapability>(
+        backend()
+            .default_stream_executor()
+            ->GetDeviceDescription()
+            .gpu_compute_capability());
+  }
   se::CudaComputeCapability GetCudaComputeCapability() {
     return backend()
         .default_stream_executor()
@@ -85,6 +99,13 @@ class CudnnFusedConvRewriterHloTest : public HloTestBase {
 
 class CudnnFusedConvRewriterTest : public GpuCodegenTest {
  public:
+  bool IsCuda() {
+    return std::holds_alternative<se::CudaComputeCapability>(
+        backend()
+            .default_stream_executor()
+            ->GetDeviceDescription()
+            .gpu_compute_capability());
+  }
   se::CudaComputeCapability GetCudaComputeCapability() {
     return backend()
         .default_stream_executor()
@@ -119,7 +140,7 @@ class CudnnFusedConvRewriterTest : public GpuCodegenTest {
   }
 
   void TestMatchWithAllTypes(absl::string_view hlo_string) {
-    for (absl::string_view type : {"f16", "f32", "f64"}) {
+    for (absl::string_view type : *(IsCuda() ? kf16f32f64 : kf16f32)) {
       const std::string hlo_with_new_type =
           absl::StrReplaceAll(hlo_string, {{"TYPE", type}});
       std::string optimized_hlo_string = GetOptimizedHlo(hlo_with_new_type);
@@ -157,7 +178,7 @@ class CudnnFusedConvRewriterTest : public GpuCodegenTest {
   }
 
   void TestNotMatchWithAllTypes(absl::string_view hlo_string) {
-    for (absl::string_view type : {"f16", "f32", "f64"}) {
+    for (absl::string_view type : *(IsCuda() ? kf16f32f64 : kf16f32)) {
       const std::string hlo_with_new_type =
           absl::StrReplaceAll(hlo_string, {{"TYPE", type}});
       std::string optimized_hlo_string = GetOptimizedHlo(hlo_with_new_type);
@@ -170,6 +191,7 @@ class CudnnFusedConvRewriterTest : public GpuCodegenTest {
 
   void TestF8(std::string pre_hlo_string, std::string custom_call_string,
               std::string serialized_graph_string) {
+    if (!IsCuda()) return;
     if (GetCudaComputeCapability().IsAtLeast(
             se::CudaComputeCapability::HOPPER)) {
       // On Hopper and newer architectures, test numerical correctness and
@@ -244,6 +266,23 @@ class CudnnFusedConvRewriterTest : public GpuCodegenTest {
   }
 };
 
+#if GOOGLE_CUDA
+#if (CUDA_VERSION < 12000 || CUDNN_VERSION < 8900)
+#define MAYBE_SKIP_TEST(CAUSE)                                           \
+  do {                                                                   \
+    if (absl::string_view(CAUSE) == "F8")                                \
+      GTEST_SKIP() << "FP8 convolutions require CUDA 12 and cuDNN 8.9."; \
+  } while (0)
+#else
+#define MAYBE_SKIP_TEST(CAUSE)
+#endif
+#else
+#define MAYBE_SKIP_TEST(CAUSE)                                \
+  do {                                                        \
+    GTEST_SKIP() << "ROCm does not support " CAUSE " fusion"; \
+  } while (0)
+#endif
+
 TEST_F(CudnnFusedConvRewriterTest, TestConvOnly) {
   // max(0, conv(x, w));
   TestMatchWithAllTypes(R"(
@@ -298,6 +337,81 @@ TEST_F(CudnnFusedConvRewriterTest, TestBias) {
     })");
 }
 
+TEST_F(CudnnFusedConvRewriterTest, Test3D) {
+  // max(0, conv(x, w) + bias);
+  std::string body = R"(
+    HloModule Test
+
+    ENTRY Test {
+      zero = TYPE[] constant(0)
+      zeros = TYPE[1,3,5,7,64] broadcast(zero), dimensions={}
+
+      input = TYPE[1,3,5,7,64] parameter(0)
+      filter = TYPE[3,3,3,64,64] parameter(1)
+      bias = TYPE[64] parameter(2)
+
+      conv = TYPE[1,3,5,7,64] convolution(input, filter), window={size=3x3x3 pad=1_1x1_1x1_1}, dim_labels=b012f_012io->b012f, feature_group_count=1
+      broadcasted_bias = TYPE[1,3,5,7,64] broadcast(bias), dimensions={4}
+      add1 = TYPE[1,3,5,7,64] add(conv, broadcasted_bias)
+    )";
+
+  std::string relu = R"(
+      ROOT relu = TYPE[1,3,5,7,64] maximum(zeros, add1)
+    })";
+
+  std::string elu = R"(
+      cmp = pred[1,3,5,7,64] compare(add1, zeros), direction=GT
+      expm1 = TYPE[1,3,5,7,64] exponential-minus-one(add1)
+      ROOT elu = TYPE[1,3,5,7,64] select(cmp, add1, expm1)
+    })";
+
+  TestMatchWithAllTypes(body + relu);
+  if (!IsCuda()) TestMatchWithAllTypes(body + elu);
+}
+
+TEST_F(CudnnFusedConvRewriterTest, TestBiasMultiCall) {
+  // max(0, conv(x, w) + bias);
+  std::string code = R"(
+    HloModule Test
+
+    ENTRY Test {
+      zero = TYPE[] constant(0)
+      zeros = TYPE[1,<<<format>>>,64] broadcast(zero), dimensions={}
+
+      input = TYPE[1,<<<format>>>,64] parameter(0)
+      filter = TYPE[3,3,64,64] parameter(1)
+      bias = TYPE[64] parameter(2)
+
+      conv = TYPE[1,<<<format>>>,64] convolution(input, filter), window={size=3x3 pad=1_1x1_1}, dim_labels=b01f_01io->b01f, feature_group_count=1
+      broadcasted_bias = TYPE[1,<<<format>>>,64] broadcast(bias), dimensions={3}
+      add1 = TYPE[1,<<<format>>>,64] add(conv, broadcasted_bias)
+      ROOT relu = TYPE[1,<<<format>>>,64] maximum(zeros, add1)
+    })";
+  absl::flat_hash_map<absl::string_view, absl::string_view> replacements;
+  replacements["<<<format>>>"] = "3,3";
+  TestMatchWithAllTypes(absl::StrReplaceAll(code, replacements));
+  replacements["<<<format>>>"] = "5,5";
+  TestMatchWithAllTypes(absl::StrReplaceAll(code, replacements));
+  replacements["<<<format>>>"] = "3,3";
+  TestMatchWithAllTypes(absl::StrReplaceAll(code, replacements));
+}
+
+TEST_F(CudnnFusedConvRewriterTest, TestBiasNoRelu) {
+  // conv(x, w) + bias;
+  TestMatchWithAllTypes(R"(
+    HloModule Test
+
+    ENTRY Test {
+      input = TYPE[1,3,3,64] parameter(0)
+      filter = TYPE[3,3,64,64] parameter(1)
+      bias = TYPE[64] parameter(2)
+
+      conv = TYPE[1,3,3,64] convolution(input, filter), window={size=3x3 pad=1_1x1_1}, dim_labels=b01f_01io->b01f, feature_group_count=1
+      broadcasted_bias = TYPE[1,3,3,64] broadcast(bias), dimensions={3}
+      ROOT add1 = TYPE[1,3,3,64] add(conv, broadcasted_bias)
+    })");
+}
+
 TEST_F(CudnnFusedConvRewriterTest, DontFuseBiasWithDepthwiseConv) {
   // conv(x, w) + bias;
   TestNotMatchWithAllTypes(R"(
@@ -365,8 +479,8 @@ TEST_F(CudnnFusedConvRewriterTest, DontFuseEluWithDepthwiseConv) {
 }
 
 TEST_F(CudnnFusedConvRewriterTest, TestRelu6) {
-  if (!GetCudaComputeCapability().IsAtLeast(
-          se::CudaComputeCapability::AMPERE)) {
+  if (IsCuda() && !GetCudaComputeCapability().IsAtLeast(
+                      se::CudaComputeCapability::AMPERE)) {
     GTEST_SKIP() << "Conv-Bias-Relu6 fusion is supported and recommended with "
                     "the Nvidia Ampere+ GPUs.";
   }
@@ -393,12 +507,11 @@ TEST_F(CudnnFusedConvRewriterTest, TestRelu6) {
 // number of input/output channels.  Check that we don't try to run this conv
 // with runtime fusion (or, if we do, that it works!).
 TEST_F(CudnnFusedConvRewriterTest, TestRelu6OddChannels) {
-  if (!GetCudaComputeCapability().IsAtLeast(
-          se::CudaComputeCapability::AMPERE)) {
+  if (IsCuda() && !GetCudaComputeCapability().IsAtLeast(
+                      se::CudaComputeCapability::AMPERE)) {
     GTEST_SKIP() << "Conv-Bias-Relu6 fusion is supported and recommended with "
                     "the Nvidia Ampere+ GPUs.";
   }
-
   TestMatchWithAllTypes(R"(
     HloModule Test
     ENTRY Test {
@@ -415,8 +528,8 @@ TEST_F(CudnnFusedConvRewriterTest, TestRelu6OddChannels) {
 }
 
 TEST_F(CudnnFusedConvRewriterTest, TestLeakyRelu) {
-  if (!GetCudaComputeCapability().IsAtLeast(
-          se::CudaComputeCapability::AMPERE)) {
+  if (IsCuda() && !GetCudaComputeCapability().IsAtLeast(
+                      se::CudaComputeCapability::AMPERE)) {
     GTEST_SKIP()
         << "Conv-Bias-LeakyRelu fusion is supported and recommended with "
            "the Nvidia Ampere+ GPUs.";
@@ -730,9 +843,7 @@ TEST_F(CudnnFusedConvRewriterTest, TestPreservesFeatureGroupCount) {
 }
 
 TEST_F(CudnnFusedConvRewriterTest, TestConvF8) {
-#if (CUDA_VERSION < 12000 || CUDNN_VERSION < 8900)
-  GTEST_SKIP() << "FP8 convolutions require CUDA 12 and cuDNN 8.9.";
-#endif
+  MAYBE_SKIP_TEST("F8");
   TestF8(
       // pre_hlo
       R"(
@@ -755,9 +866,7 @@ TEST_F(CudnnFusedConvRewriterTest, TestConvF8) {
 }
 
 TEST_F(CudnnFusedConvRewriterTest, TestConvScaledOutputF8) {
-#if (CUDA_VERSION < 12000 || CUDNN_VERSION < 8900)
-  GTEST_SKIP() << "FP8 convolutions require CUDA 12 and cuDNN 8.9.";
-#endif
+  MAYBE_SKIP_TEST("F8");
   TestF8(
       // pre_hlo
       R"(
@@ -791,9 +900,7 @@ TEST_F(CudnnFusedConvRewriterTest, TestConvScaledOutputF8) {
 }
 
 TEST_F(CudnnFusedConvRewriterTest, TestConvInvscaledOutputF8) {
-#if (CUDA_VERSION < 12000 || CUDNN_VERSION < 8900)
-  GTEST_SKIP() << "FP8 convolutions require CUDA 12 and cuDNN 8.9.";
-#endif
+  MAYBE_SKIP_TEST("F8");
   TestF8(
       // pre_hlo
       R"(
@@ -827,9 +934,7 @@ TEST_F(CudnnFusedConvRewriterTest, TestConvInvscaledOutputF8) {
 }
 
 TEST_F(CudnnFusedConvRewriterTest, TestConvScaledF8Parameterized) {
-#if (CUDA_VERSION < 12000 || CUDNN_VERSION < 8900)
-  GTEST_SKIP() << "FP8 convolutions require CUDA 12 and cuDNN 8.9.";
-#endif
+  MAYBE_SKIP_TEST("F8");
   TestF8Parameterized(
       // pre_hlo
       R"(
@@ -869,9 +974,7 @@ TEST_F(CudnnFusedConvRewriterTest, TestConvScaledF8Parameterized) {
 }
 
 TEST_F(CudnnFusedConvRewriterTest, TestConvScaledBiasF8) {
-#if (CUDA_VERSION < 12000 || CUDNN_VERSION < 8900)
-  GTEST_SKIP() << "FP8 convolutions require CUDA 12 and cuDNN 8.9.";
-#endif
+  MAYBE_SKIP_TEST("F8");
   TestF8(
       // pre_hlo
       R"(
@@ -913,9 +1016,7 @@ TEST_F(CudnnFusedConvRewriterTest, TestConvScaledBiasF8) {
 }
 
 TEST_F(CudnnFusedConvRewriterTest, TestConvScaledReluF8) {
-#if (CUDA_VERSION < 12000 || CUDNN_VERSION < 8900)
-  GTEST_SKIP() << "FP8 convolutions require CUDA 12 and cuDNN 8.9.";
-#endif
+  MAYBE_SKIP_TEST("F8");
   TestF8(
       // pre_hlo
       R"(
@@ -952,6 +1053,7 @@ TEST_F(CudnnFusedConvRewriterTest, TestConvScaledReluF8) {
 }
 
 TEST_F(CudnnFusedConvRewriterTest, TestConvAmaxF8) {
+  MAYBE_SKIP_TEST("F8");
   TestF8(
       // pre_hlo
       R"(
@@ -1001,6 +1103,7 @@ TEST_F(CudnnFusedConvRewriterTest, TestConvAmaxF8) {
 }
 
 TEST_F(CudnnFusedConvRewriterTest, TestConvReluAmaxF8) {
+  MAYBE_SKIP_TEST("F8");
   TestF8(
       // pre_hlo
       R"(
@@ -1053,9 +1156,7 @@ TEST_F(CudnnFusedConvRewriterTest, TestConvReluAmaxF8) {
 }
 
 TEST_F(CudnnFusedConvRewriterTest, TestConvScaledOutputMultipleUsersF8) {
-#if (CUDA_VERSION < 12000 || CUDNN_VERSION < 8900)
-  GTEST_SKIP() << "FP8 convolutions require CUDA 12 and cuDNN 8.9.";
-#endif
+  MAYBE_SKIP_TEST("F8");
   TestF8(
       // pre_hlo
       R"(
@@ -1095,9 +1196,7 @@ TEST_F(CudnnFusedConvRewriterTest, TestConvScaledOutputMultipleUsersF8) {
 }
 
 TEST_F(CudnnFusedConvRewriterTest, TestConvScaledOutputUnsupportedUserF8) {
-#if (CUDA_VERSION < 12000 || CUDNN_VERSION < 8900)
-  GTEST_SKIP() << "FP8 convolutions require CUDA 12 and cuDNN 8.9.";
-#endif
+  MAYBE_SKIP_TEST("F8");
   TestF8(
       // pre_hlo
       R"(
@@ -1133,6 +1232,7 @@ TEST_F(CudnnFusedConvRewriterTest, TestConvScaledOutputUnsupportedUserF8) {
 }
 
 TEST_F(CudnnFusedConvRewriterTest, TestConvInt8ToInt8) {
+  MAYBE_SKIP_TEST("I8");
   // max(0, clamp(conv(x, w)))); for int8_t
   TestClamp(
       // pre_hlo
@@ -1167,6 +1267,7 @@ TEST_F(CudnnFusedConvRewriterTest, TestConvInt8ToInt8) {
 }
 
 TEST_F(CudnnFusedConvRewriterHloTest, TestConvInt8ToFloat) {
+  MAYBE_SKIP_TEST("I8");
   const std::string module_str = R"(
     HloModule Test
 
@@ -1198,6 +1299,7 @@ TEST_F(CudnnFusedConvRewriterHloTest, TestConvInt8ToFloat) {
 }
 
 TEST_F(CudnnFusedConvRewriterHloTest, TestConvInt8ToInt8BiasSideInput) {
+  MAYBE_SKIP_TEST("I8");
   const std::string module_str = R"(
     HloModule Test
 
@@ -1238,6 +1340,7 @@ TEST_F(CudnnFusedConvRewriterHloTest, TestConvInt8ToInt8BiasSideInput) {
 }
 
 TEST_F(CudnnFusedConvRewriterHloTest, TestReluAfterConvert) {
+  MAYBE_SKIP_TEST("I8");
   const std::string module_str = R"(
     HloModule Test
 
@@ -1286,6 +1389,7 @@ TEST_F(CudnnFusedConvRewriterHloTest, TestReluAfterConvert) {
 }
 
 TEST_F(CudnnFusedConvRewriterHloTest, TestConvInt8ToFloatBiasSideInput) {
+  MAYBE_SKIP_TEST("I8");
   const std::string module_str = R"(
     HloModule Test
 
@@ -1333,6 +1437,7 @@ TEST_F(CudnnFusedConvRewriterHloTest, TestConvInt8ToFloatBiasSideInput) {
 //   reshape(side_input * alpha).
 // Make sure we can pattern-match this.
 TEST_F(CudnnFusedConvRewriterHloTest, Int8SideInputWithScaleAndReshape) {
+  MAYBE_SKIP_TEST("I8");
   const std::string module_str = R"(
     HloModule Test
 
@@ -1386,6 +1491,7 @@ TEST_F(CudnnFusedConvRewriterHloTest, Int8SideInputWithScaleAndReshape) {
 }
 
 TEST_F(CudnnFusedConvRewriterHloTest, FuseAlpha) {
+  MAYBE_SKIP_TEST("I8");
   const std::string module_str = R"(
     HloModule Test
 
@@ -2054,6 +2160,7 @@ TEST_F(CudnnFusedConvRewriterHloTest, DontFuseToS8IfMultipleUsers) {
 }
 
 TEST_F(CudnnFusedConvRewriterHloTest, RemoveConvertByFusingS32ToF32) {
+  MAYBE_SKIP_TEST("I8");
   const std::string_view module_str = R"(
     HloModule Test
 
@@ -2079,6 +2186,7 @@ TEST_F(CudnnFusedConvRewriterHloTest, RemoveConvertByFusingS32ToF32) {
 }
 
 TEST_F(CudnnFusedConvRewriterHloTest, RemoveConvertByFusingS8ToF32) {
+  MAYBE_SKIP_TEST("I8");
   const std::string_view module_str = R"(
     HloModule Test
 
@@ -2104,6 +2212,7 @@ TEST_F(CudnnFusedConvRewriterHloTest, RemoveConvertByFusingS8ToF32) {
 }
 
 TEST_F(CudnnFusedConvRewriterHloTest, RemoveConvertByFusingF32ToS8) {
+  MAYBE_SKIP_TEST("I8");
   const std::string_view module_str = R"(
     HloModule Test
 
@@ -2608,6 +2717,7 @@ TEST_F(CudnnFusedConvRewriterHloTest, F32ConstantsNotLosslesslyConvertible) {
 }
 
 TEST_F(CudnnFusedConvRewriterHloTest, FuseReluBeforeConvert) {
+  MAYBE_SKIP_TEST("I8");
   const std::string module_str = R"(
   HloModule Test
 
@@ -2663,6 +2773,7 @@ TEST_F(CudnnFusedConvRewriterHloTest, FuseReluBeforeConvert) {
 }
 
 TEST_F(CudnnFusedConvRewriterHloTest, BiasTypeMatchesConvTypeIfFp) {
+  MAYBE_SKIP_TEST("F64");
   const std::string module_str = R"(
   HloModule Test
 
@@ -2699,6 +2810,7 @@ TEST_F(CudnnFusedConvRewriterHloTest, BiasTypeMatchesConvTypeIfFp) {
 }
 
 TEST_F(CudnnFusedConvRewriterTest, TestFusedConvInt8ToInt8) {
+  MAYBE_SKIP_TEST("I8");
   // clamp(max(0, conv(x, w)+bias)); for int8_t
   TestClamp(
       // pre_hlo
@@ -2740,6 +2852,7 @@ TEST_F(CudnnFusedConvRewriterTest, TestFusedConvInt8ToInt8) {
 
 // Disabled per b/190854862 or nvbugs/3326122.
 TEST_F(CudnnFusedConvRewriterTest, DISABLED_TestFusedConvInt8ToFloat) {
+  MAYBE_SKIP_TEST("I8");
   // max(0, convert<float>(conv<int32_t>(int8_x),
   // conv<int32_t>(int8_w))+float_bias)); int8_t to float via bias.
   TestClamp(
@@ -2775,6 +2888,7 @@ TEST_F(CudnnFusedConvRewriterTest, DISABLED_TestFusedConvInt8ToFloat) {
 
 TEST_F(CudnnFusedConvRewriterTest,
        TestFusedConvWithScaledInt8SideInputBiasInt8ToInt8) {
+  MAYBE_SKIP_TEST("I8");
   // clamp(max(0, alpha_conv * conv(x, w) + alpha_side *
   // convert<int32_t>(int8_side_input) + bias)); for int8_t
   TestClamp(
@@ -2826,6 +2940,7 @@ TEST_F(CudnnFusedConvRewriterTest,
 
 TEST_F(CudnnFusedConvRewriterTest,
        TestFusedConvWithScaledFloatSideInputBiasInt8ToInt8) {
+  MAYBE_SKIP_TEST("I8");
   // From:
   // convert<int8_t>(clamp(max(0, alpha_conv * conv(x, w) + alpha_side *
   // float_side_input + bias))); To: convert<int8_t>(clamp(conv(int8_x, int8_w,
@@ -2878,6 +2993,7 @@ TEST_F(CudnnFusedConvRewriterTest,
 
 TEST_F(CudnnFusedConvRewriterTest,
        TestFusedConvWithScaledInt8SideInputBiasInt8ToFloat) {
+  MAYBE_SKIP_TEST("I8");
   // From:
   // clamp(max(0, alpha_conv * conv(x, w) + alpha_side *
   // convert<float>(int8_side_input) + bias)); To: clamp(conv(int8_x, int8_w,
@@ -2928,6 +3044,7 @@ TEST_F(CudnnFusedConvRewriterTest,
 }
 
 TEST_F(CudnnFusedConvRewriterTest, TestConvInt8ToInt8NoClamp) {
+  MAYBE_SKIP_TEST("I8");
   // Check that integer convolution without clamp to int8_t is not allowed.
   // convert<int8_t>(custom_call<int32_t>(int32_x, int32_w,
   // cudnnConvolutionForward))
@@ -2951,6 +3068,7 @@ TEST_F(CudnnFusedConvRewriterTest, TestConvInt8ToInt8NoClamp) {
 }
 
 TEST_F(CudnnFusedConvRewriterTest, TestFusedConvInt8ToInt8NoClamp) {
+  MAYBE_SKIP_TEST("I8");
   // Although bias and so on are fused with forward convolution,
   // it is still not allowed if the output is not clampped/converted to int8_t
   // max(0, alpha_conv * conv(x, w) + alpha_side * side_input + bias); for
diff --git a/third_party/xla/xla/service/gpu/gpu_layout_assignment.cc b/third_party/xla/xla/service/gpu/gpu_layout_assignment.cc
index ee5d7799a52814..227d965257f932 100644
--- a/third_party/xla/xla/service/gpu/gpu_layout_assignment.cc
+++ b/third_party/xla/xla/service/gpu/gpu_layout_assignment.cc
@@ -128,6 +128,12 @@ HeuristicLayoutAssignment(const HloInstruction* instr,
     return kAllNHWC;
   }
 
+  const auto* rocm_compute_capability =
+      std::get_if<se::RocmComputeCapability>(&gpu_version);
+  if (rocm_compute_capability && input_ty == F16) return kAllNHWC;
+
+  // If we're not Volta or not fp16/bfloat16, or not conv2D, the decision is
+  // easy: Use NCHW.
   const bool isFloat16 = (input_ty == F16) || (input_ty == BF16);
   if (std::holds_alternative<se::CudaComputeCapability>(gpu_version)) {
     // If we're not Volta or not fp16/bfloat16, or not conv2D, the decision is
diff --git a/third_party/xla/xla/service/gpu/runtime/convolution_thunk.cc b/third_party/xla/xla/service/gpu/runtime/convolution_thunk.cc
index 32e81855e7ffb8..8577fe9f55106e 100644
--- a/third_party/xla/xla/service/gpu/runtime/convolution_thunk.cc
+++ b/third_party/xla/xla/service/gpu/runtime/convolution_thunk.cc
@@ -99,6 +99,9 @@ absl::Status ConvolutionThunk::ExecuteOnStream(const ExecuteParams& params) {
     TF_ASSIGN_OR_RETURN(se::dnn::DataType input_type,
                         GetDNNDataTypeFromPrimitiveType(config_.input_type));
 
+    TF_ASSIGN_OR_RETURN(se::dnn::DataType output_type,
+                        GetDNNDataTypeFromPrimitiveType(config_.output_type));
+
     TF_ASSIGN_OR_RETURN(auto dnn,
                         se::dnn::internal::GetDnnFromStream(params.stream));
     se::OwningScratchAllocator<> scratch_allocator(
@@ -107,7 +110,7 @@ absl::Status ConvolutionThunk::ExecuteOnStream(const ExecuteParams& params) {
 
     std::vector<se::dnn::ProfileResult> profile_results;
     dnn->GetMIOpenConvolveAlgorithms(
-        kind, input_type, params.stream, config_.input_descriptor,
+        kind, input_type, output_type, params.stream, config_.input_descriptor,
         conv_params.input_buf, config_.filter_descriptor,
         conv_params.filter_buf, config_.output_descriptor,
         conv_params.output_buf, config_.conv_desc, &scratch_allocator,
diff --git a/third_party/xla/xla/stream_executor/dnn.cc b/third_party/xla/xla/stream_executor/dnn.cc
index ecd7555c94d7db..2b0d03516f4ca7 100644
--- a/third_party/xla/xla/stream_executor/dnn.cc
+++ b/third_party/xla/xla/stream_executor/dnn.cc
@@ -287,7 +287,8 @@ DnnSupport::FusedMHABackwardRunnerFromDesc(
 
 bool DnnSupport::GetMIOpenConvolveAlgorithms(
     dnn::ConvolutionKind /*kind*/, dnn::DataType /*element_type*/,
-    Stream* /*stream*/, const dnn::BatchDescriptor& /*input_descriptor*/,
+    dnn::DataType /*output_type*/, Stream* /*stream*/,
+    const dnn::BatchDescriptor& /*input_descriptor*/,
     DeviceMemoryBase input_data,
     const dnn::FilterDescriptor& /*filter_descriptor*/,
     DeviceMemoryBase filter_data,
diff --git a/third_party/xla/xla/stream_executor/dnn.h b/third_party/xla/xla/stream_executor/dnn.h
index 41fb3e34c8fc0c..299d8cd2049bfe 100644
--- a/third_party/xla/xla/stream_executor/dnn.h
+++ b/third_party/xla/xla/stream_executor/dnn.h
@@ -1767,10 +1767,11 @@ class DnnSupport {
       dnn::FMHAMaskKind mask_type);
 
   virtual bool GetMIOpenConvolveAlgorithms(
-      ConvolutionKind kind, DataType element_type, Stream* stream,
-      const BatchDescriptor& input_descriptor, DeviceMemoryBase input_data,
-      const FilterDescriptor& filter_descriptor, DeviceMemoryBase filter_data,
-      const BatchDescriptor& output_descriptor, DeviceMemoryBase output_data,
+      ConvolutionKind kind, DataType element_type, DataType output_type,
+      Stream* stream, const BatchDescriptor& input_descriptor,
+      DeviceMemoryBase input_data, const FilterDescriptor& filter_descriptor,
+      DeviceMemoryBase filter_data, const BatchDescriptor& output_descriptor,
+      DeviceMemoryBase output_data,
       const ConvolutionDescriptor& convolution_descriptor,
       ScratchAllocator* scratch_allocator,
       std::vector<ProfileResult>* out_algorithms);
diff --git a/third_party/xla/xla/stream_executor/rocm/rocm_dnn.cc b/third_party/xla/xla/stream_executor/rocm/rocm_dnn.cc
index fd3bcafa501390..cf5f9a9c80b72f 100644
--- a/third_party/xla/xla/stream_executor/rocm/rocm_dnn.cc
+++ b/third_party/xla/xla/stream_executor/rocm/rocm_dnn.cc
@@ -43,6 +43,7 @@ limitations under the License.
 #include "xla/stream_executor/scratch_allocator.h"
 #include "xla/stream_executor/stream.h"
 #include "xla/stream_executor/stream_executor.h"
+#include "xla/stream_executor/stream_interface.h"
 #include "xla/tsl/util/determinism.h"
 #include "xla/tsl/util/env_var.h"
 #include "tsl/platform/env.h"
@@ -112,6 +113,8 @@ string ToString(miopenStatus_t status) {
       return "miopenStatusNotImplemented";
     case miopenStatusUnknownError:
       return "miopenStatusUnknownError";
+    case miopenStatusUnsupportedOp:
+      return "miopenStatusUnsupportedOp";
     default:
       return absl::StrCat("<unknown miopen status: ", static_cast<int>(status),
                           ">");
@@ -377,6 +380,8 @@ namespace wrap {
   __macro(miopenConvolutionForwardGetSolutionWorkspaceSize)          \
   __macro(miopenConvolutionForwardCompileSolution)                   \
   __macro(miopenConvolutionForwardImmediate)                         \
+  __macro(miopenConvolutionForwardBias)                              \
+  __macro(miopenConvolutionBiasActivationForward)                    \
   __macro(miopenConvolutionBackwardDataGetSolutionCount)             \
   __macro(miopenConvolutionBackwardDataGetSolution)                  \
   __macro(miopenConvolutionBackwardDataGetSolutionWorkspaceSize)     \
@@ -495,11 +500,13 @@ namespace wrap {
   __macro(miopenExecuteFusionPlan)                                   \
   __macro(miopenDestroyOperatorArgs)                                 \
   __macro(miopenDestroyFusionPlan)                                   \
+  __macro(miopenConvolutionBiasActivationForward)                    \
   __macro(miopenConvolutionForwardGetSolutionCount)                  \
   __macro(miopenConvolutionForwardGetSolution)                       \
   __macro(miopenConvolutionForwardGetSolutionWorkspaceSize)          \
   __macro(miopenConvolutionForwardCompileSolution)                   \
   __macro(miopenConvolutionForwardImmediate)                         \
+  __macro(miopenConvolutionForwardBias)                              \
   __macro(miopenConvolutionBackwardDataGetSolutionCount)             \
   __macro(miopenConvolutionBackwardDataGetSolution)                  \
   __macro(miopenConvolutionBackwardDataGetSolutionWorkspaceSize)     \
@@ -565,7 +572,7 @@ uint64_t GetHashValue(miopenConvolutionDescriptor_t conv_desc) {
 
   uint64_t hash_value = tsl::hash<int>()(c_mode);
   auto hash64Combine = [&hash_value](int element) {
-    tsl::Hash64Combine(hash_value, tsl::hash<int>()(element));
+    hash_value = tsl::Hash64Combine(hash_value, tsl::hash<int>()(element));
   };
   std::for_each(pad.begin(), pad.end(), hash64Combine);
   std::for_each(stride.begin(), stride.end(), hash64Combine);
@@ -597,9 +604,11 @@ class CachedFusionPlans {
 
     auto it = cached_plans.find(hash);
     if (it != cached_plans.end()) {
+      VLOG(2) << "Found a cached plan for " << hash;
       *fusion_plan = it->second;
       found_cached_plan = true;
     } else {
+      VLOG(2) << "Creating a new plan for " << hash;
       auto status = wrap::miopenCreateFusionPlan(fusion_plan, fusion_direction,
                                                  input_descriptor);
       if (status != miopenStatusSuccess) {
@@ -675,6 +684,7 @@ dnn::ProfileResult GetProfileResultFromConvAlgoPerf(
   int64_t algo_id;
   switch (kind) {
     case dnn::ConvolutionKind::FORWARD:
+    case dnn::ConvolutionKind::FORWARD_BIAS_ACTIVATION:
       algo_id = algorithm.fwd_algo;
       break;
     case dnn::ConvolutionKind::BACKWARD_DATA:
@@ -798,425 +808,398 @@ absl::StatusOr<stream_executor::dnn::VersionInfo> MIOpenSupport::GetVersion() {
   return stream_executor::dnn::VersionInfo(1, 3, 0);
 }
 
-// Turns a BatchDescriptor structure into a miopen tensor handle within a scope.
-class ScopedTensorDescriptor {
- public:
-  ScopedTensorDescriptor(const BatchDescriptor& batch_descriptor,
-                         miopenDataType_t elem_type)
-      : handle_(nullptr) {
-    auto status = wrap::miopenCreateTensorDescriptor(&handle_);
-    if (status != miopenStatusSuccess) {
-      LOG(FATAL) << "could not create miopen tensor descriptor: "
-                 << ToString(status);
-    }
-
-    switch (batch_descriptor.layout()) {
-      case dnn::DataLayout::kBatchYXDepth:
-      case dnn::DataLayout::kBatchDepthYX: {
-        const int nd = batch_descriptor.ndims() + 2;
-
-        // MIOpen requires the strides and dims to be ordered as BDYX.
-        std::vector<int64_t> strides64 =
-            batch_descriptor.full_strides(dnn::DataLayout::kBatchDepthYX);
-        std::vector<int64_t> dims64 =
-            batch_descriptor.full_dims(dnn::DataLayout::kBatchDepthYX);
-
-        // MIOpen requires arrays of ints.
-        std::vector<int> strides(nd);
-        std::vector<int> dims(nd);
-        std::transform(strides64.cbegin(), strides64.cend(), strides.begin(),
-                       &CheckedNarrowing<int64_t, int>);
-        std::transform(dims64.cbegin(), dims64.cend(), dims.begin(),
-                       &CheckedNarrowing<int64_t, int>);
-        status = wrap::miopenSetTensorDescriptor(handle_, elem_type, nd,
-                                                 dims.data(), strides.data());
-
-        if (status != miopenStatusSuccess) {
-          LOG(FATAL) << "could not convert BatchDescriptor "
-                     << batch_descriptor.ToString()
-                     << " to miopen tensor descriptor: " << ToString(status);
-        }
-      } break;
-      default:
-        LOG(FATAL) << "Unsupported tensor format "
-                   << DataLayoutString(batch_descriptor.layout());
-        break;
-    }
-  }
-
-  ~ScopedTensorDescriptor() {
-    auto status = wrap::miopenDestroyTensorDescriptor(handle_);
-    if (status != miopenStatusSuccess) {
-      LOG(ERROR) << "could not destroy miopen tensor descriptor: "
-                 << ToString(status);
-    }
-  }
+template <typename T>
+miopenStatus_t miDestroyObject(T obj) {
+  return miopenStatusSuccess;
+}
 
-  miopenTensorDescriptor_t handle() const { return handle_; }
+template <>
+miopenStatus_t miDestroyObject(miopenTensorDescriptor_t obj) {
+  return wrap::miopenDestroyTensorDescriptor(obj);
+}
 
- private:
-  miopenTensorDescriptor_t handle_;  // Owned.
+template <>
+miopenStatus_t miDestroyObject(miopenConvolutionDescriptor_t obj) {
+  return wrap::miopenDestroyConvolutionDescriptor(obj);
+}
 
-  ScopedTensorDescriptor(const ScopedTensorDescriptor&) = delete;
-  void operator=(const ScopedTensorDescriptor&) = delete;
-};
+template <>
+miopenStatus_t miDestroyObject(miopenPoolingDescriptor_t obj) {
+  return wrap::miopenDestroyPoolingDescriptor(obj);
+}
 
-// Turns a FilterDescriptor structure into a miopen filter handle within a
-// scope.
-class ScopedFilterDescriptor {
- public:
-  ScopedFilterDescriptor(const FilterDescriptor& filter_descriptor,
-                         miopenDataType_t elem_type)
-      : handle_(nullptr) {
-    auto status = wrap::miopenCreateTensorDescriptor(&handle_);
-    if (status != miopenStatusSuccess) {
-      LOG(FATAL) << "could not create miopen filter descriptor: "
-                 << ToString(status);
-    }
+template <>
+miopenStatus_t miDestroyObject(miopenLRNDescriptor_t obj) {
+  return wrap::miopenDestroyLRNDescriptor(obj);
+}
 
-    // We need to pass two vectors to the miopenSetTensorDescriptor routine
-    // "dims" (length == number of dims, elem value == dimension size)
-    // "strides" (length == number of dims, elem value == stride size)
-    //
-    // Irrespective of the actual filter layout, the indexing of both those
-    // vectors must be the following (coz that is what MIOpen expects)
-    // dims[0] = strides[0] = N or output
-    // dims[1] = strides[1] = C or input
-    // dims[2] = strides[2] = H or spatial dim 0
-    // dims[3] = strides[3] = W or spatial dim 1
-    //
-    // assume you have a tensor with dimensions
-    // batch descriptor name    filter descriptor name    value
-    //   N (batch size)            O (output features)    256
-    //   C (channels)              I (input features)       3
-    //   H (height)                H (height)               7
-    //   W (width)                 W (width)                5
-    //
-    // The content of "dims" will be the same irrespective of layout
-    // layout (NCHW or NHWC), and MIOpen expects it should be
-    //                           NCHW layout   NHWC layout
-    // dims[0] = size of N dim =    256           256
-    // dims[1] = size of C dim =      3             3
-    // dims[2] = size of H dim =      7             7
-    // dims[3] = size of W dim =      5             5
-    //
-    // The content of "strides" will be different based on layout
-    //                                  NCHW layout   NHWC layout
-    //  strides[0] = stride of N dim =     7x5x3       7x5x3
-    //  strides[1] = stride of C dim =     7x5         1
-    //  strides[2] = stride of H dim =     5           5x3
-    //  strides[3] = stride of W dim =     1           3
-
-    switch (filter_descriptor.layout()) {
-      case dnn::FilterLayout::kOutputYXInput:
-      case dnn::FilterLayout::kOutputInputYX: {
-        const int nd = filter_descriptor.ndims() + 2;
-
-        // MIOpen requires the strides and dims to be ordered as BDYX.
-        std::vector<int64_t> strides64 =
-            filter_descriptor.full_strides(dnn::FilterLayout::kOutputInputYX);
-        std::vector<int64_t> dims64 =
-            filter_descriptor.full_dims(dnn::FilterLayout::kOutputInputYX);
-
-        // MIOpen requires arrays of ints.
-        std::vector<int> strides;
-        std::vector<int> dims;
-        absl::c_transform(strides64, std::back_inserter(strides),
-                          &CheckedNarrowing<int64_t, int>);
-        absl::c_transform(dims64, std::back_inserter(dims),
-                          &CheckedNarrowing<int64_t, int>);
-        status = wrap::miopenSetTensorDescriptor(handle_, elem_type, nd,
-                                                 dims.data(), strides.data());
+template <typename T>
+struct ScopedDescriptor {
+  ScopedDescriptor() : handle_(nullptr) {}
 
-        if (status != miopenStatusSuccess) {
-          LOG(FATAL) << "could not convert FilterDescriptor "
-                     << filter_descriptor.ToString()
-                     << " to miopen tensor descriptor: " << ToString(status);
-        }
-      } break;
-      default:
-        LOG(FATAL) << "Unsupported tensor format "
-                   << FilterLayoutString(filter_descriptor.layout());
-        break;
-    }
+  ScopedDescriptor(ScopedDescriptor<T>&& other) {
+    handle_ = other.handle_;
+    other.handle_ = nullptr;
   }
 
-  ~ScopedFilterDescriptor() {
-    auto status = wrap::miopenDestroyTensorDescriptor(handle_);
+  ~ScopedDescriptor() {
+    if (handle_ != nullptr) return;
+
+    auto status = miDestroyObject(
+        handle_);  // wrap::miopenDestroyTensorDescriptor(handle_);
     if (status != miopenStatusSuccess) {
-      LOG(ERROR) << "could not destroy miopen filter descriptor: "
+      LOG(ERROR) << "could not destroy miopen tensor descriptor: "
                  << ToString(status);
     }
   }
 
-  miopenTensorDescriptor_t handle() const { return handle_; }
+  T handle() const { return handle_; }
 
- private:
-  // miopen filter descriptor this object creates. Owned.
-  miopenTensorDescriptor_t handle_;
+  T handle_;  // Owned.
 
-  ScopedFilterDescriptor(const ScopedFilterDescriptor&) = delete;
-  void operator=(const ScopedFilterDescriptor&) = delete;
+  ScopedDescriptor(const ScopedDescriptor<T>&) = delete;
+  void operator=(const ScopedDescriptor<T>&) = delete;
 };
 
-// Turns a ConvolutionDescriptor structure into a miopen convolution handle
-// within a scope.
-class ScopedConvolutionDescriptor {
- public:
-  ScopedConvolutionDescriptor(
-      const ConvolutionDescriptor& convolution_descriptor,
-      miopenDataType_t data_type)
-      : handle_(nullptr) {
-    auto status = wrap::miopenCreateConvolutionDescriptor(&handle_);
-    if (status != miopenStatusSuccess) {
-      LOG(FATAL) << "could not create miopen convolution descriptor: "
-                 << ToString(status);
-    }
-    const auto& strides64 = convolution_descriptor.strides();
-    const auto& padding64 = convolution_descriptor.padding();
-    if (convolution_descriptor.pad_alignment() ==
-        dnn::PadAlignment::kTensorFlowPadding) {
-      LOG(ERROR) << "TensorFlow padding alignment is not supported.";
-    }
+using ScopedTensorDescriptor = ScopedDescriptor<miopenTensorDescriptor_t>;
+using ScopedFilterDescriptor = ScopedDescriptor<miopenTensorDescriptor_t>;
+using ScopedConvolutionDescriptor =
+    ScopedDescriptor<miopenConvolutionDescriptor_t>;
+using ScopedPoolingDescriptor = ScopedDescriptor<miopenPoolingDescriptor_t>;
+using ScopedNormalizeDescriptor = ScopedDescriptor<miopenLRNDescriptor_t>;
+
+absl::StatusOr<ScopedTensorDescriptor> scope(
+    const BatchDescriptor& batch_descriptor, miopenDataType_t data_type) {
+  ScopedTensorDescriptor obj;
+  auto status = wrap::miopenCreateTensorDescriptor(&obj.handle_);
+  if (status != miopenStatusSuccess) {
+    return absl::InternalError("could not create miopen tensor descriptor: " +
+                               ToString(status));
+  }
 
-    // MIOpen requires arrays of ints.
-    std::vector<int> strides(convolution_descriptor.ndims());
-    std::vector<int> padding(convolution_descriptor.ndims());
-    std::transform(strides64.cbegin(), strides64.cend(), strides.begin(),
-                   &CheckedNarrowing<int64_t, int>);
-    std::transform(padding64.cbegin(), padding64.cend(), padding.begin(),
-                   &CheckedNarrowing<int64_t, int>);
-
-    std::vector<int> upscale(convolution_descriptor.ndims());
-    const auto& dilations64 = convolution_descriptor.dilations();
-    std::transform(dilations64.cbegin(), dilations64.cend(), upscale.begin(),
-                   &CheckedNarrowing<int64_t, int>);
-
-    status = wrap::miopenInitConvolutionNdDescriptor(
-        handle_, convolution_descriptor.ndims(), padding.data(), strides.data(),
-        upscale.data(), miopenConvolution);
-    if (status != miopenStatusSuccess) {
-      LOG(FATAL) << "could not set miopen convolution descriptor: "
-                 << ToString(status);
-    }
+  switch (batch_descriptor.layout()) {
+    case dnn::DataLayout::kBatchYXDepth:
+    case dnn::DataLayout::kBatchDepthYX: {
+      const int nd = batch_descriptor.ndims() + 2;
 
-    VLOG(2) << "Requesting grouped convolution: "
-            << convolution_descriptor.group_count();
-    status = wrap::miopenSetConvolutionGroupCount(
-        handle_, convolution_descriptor.group_count());
-    if (status != miopenStatusSuccess) {
-      LOG(FATAL) << "could not set miopen convolution group count: "
-                 << ToString(status);
-    }
+      // MIOpen requires the strides and dims to be ordered as BDYX.
+      std::vector<int64_t> strides64 =
+          batch_descriptor.full_strides(dnn::DataLayout::kBatchDepthYX);
+      std::vector<int64_t> dims64 =
+          batch_descriptor.full_dims(dnn::DataLayout::kBatchDepthYX);
+
+      // MIOpen requires arrays of ints.
+      std::vector<int> strides(nd);
+      std::vector<int> dims(nd);
+      std::transform(strides64.cbegin(), strides64.cend(), strides.begin(),
+                     &CheckedNarrowing<int64_t, int>);
+      std::transform(dims64.cbegin(), dims64.cend(), dims.begin(),
+                     &CheckedNarrowing<int64_t, int>);
+      status = wrap::miopenSetTensorDescriptor(obj.handle_, data_type, nd,
+                                               dims.data(), strides.data());
 
-#if (TF_ROCM_VERSION >= 50300)
-    if (RequireMIOpenDeterminism()) {
-      status = wrap::miopenSetConvolutionAttribute(
-          handle_, MIOPEN_CONVOLUTION_ATTRIB_DETERMINISTIC, 1);
       if (status != miopenStatusSuccess) {
-        LOG(FATAL) << "could not set miopen convolution attribute: "
-                   << ToString(status);
+        return absl::InternalError(
+            "could not convert BatchDescriptor " + batch_descriptor.ToString() +
+            " to miopen tensor descriptor: " + ToString(status));
       }
-    }
-#endif
+    } break;
+    default:
+      return absl::InternalError("Unsupported tensor format " +
+                                 DataLayoutString(batch_descriptor.layout()));
+      break;
   }
-  ~ScopedConvolutionDescriptor() {
-    auto status = wrap::miopenDestroyConvolutionDescriptor(handle_);
-    if (status != miopenStatusSuccess) {
-      LOG(ERROR) << "could not destroy miopen convolution descriptor: "
-                 << ToString(status);
-    }
+  return obj;
+}
+
+absl::StatusOr<ScopedFilterDescriptor> scope(
+    const FilterDescriptor& filter_descriptor, miopenDataType_t data_type) {
+  ScopedFilterDescriptor obj;
+  auto status = wrap::miopenCreateTensorDescriptor(&obj.handle_);
+  if (status != miopenStatusSuccess) {
+    LOG(FATAL) << "could not create miopen filter descriptor: "
+               << ToString(status);
   }
 
-  miopenConvolutionDescriptor_t handle() const { return handle_; }
+  // We need to pass two vectors to the miopenSetTensorDescriptor routine
+  // "dims" (length == number of dims, elem value == dimension size)
+  // "strides" (length == number of dims, elem value == stride size)
+  //
+  // Irrespective of the actual filter layout, the indexing of both those
+  // vectors must be the following (coz that is what MIOpen expects)
+  // dims[0] = strides[0] = N or output
+  // dims[1] = strides[1] = C or input
+  // dims[2] = strides[2] = H or spatial dim 0
+  // dims[3] = strides[3] = W or spatial dim 1
+  //
+  // assume you have a tensor with dimensions
+  // batch descriptor name    filter descriptor name    value
+  //   N (batch size)            O (output features)    256
+  //   C (channels)              I (input features)       3
+  //   H (height)                H (height)               7
+  //   W (width)                 W (width)                5
+  //
+  // The content of "dims" will be the same irrespective of layout
+  // layout (NCHW or NHWC), and MIOpen expects it should be
+  //                           NCHW layout   NHWC layout
+  // dims[0] = size of N dim =    256           256
+  // dims[1] = size of C dim =      3             3
+  // dims[2] = size of H dim =      7             7
+  // dims[3] = size of W dim =      5             5
+  //
+  // The content of "strides" will be different based on layout
+  //                                  NCHW layout   NHWC layout
+  //  strides[0] = stride of N dim =     7x5x3       7x5x3
+  //  strides[1] = stride of C dim =     7x5         1
+  //  strides[2] = stride of H dim =     5           5x3
+  //  strides[3] = stride of W dim =     1           3
+
+  switch (filter_descriptor.layout()) {
+    case dnn::FilterLayout::kOutputYXInput:
+    case dnn::FilterLayout::kOutputInputYX: {
+      const int nd = filter_descriptor.ndims() + 2;
+
+      // MIOpen requires the strides and dims to be ordered as BDYX.
+      std::vector<int64_t> strides64 =
+          filter_descriptor.full_strides(dnn::FilterLayout::kOutputInputYX);
+      std::vector<int64_t> dims64 =
+          filter_descriptor.full_dims(dnn::FilterLayout::kOutputInputYX);
+
+      // MIOpen requires arrays of ints.
+      std::vector<int> strides;
+      std::vector<int> dims;
+      absl::c_transform(strides64, std::back_inserter(strides),
+                        &CheckedNarrowing<int64_t, int>);
+      absl::c_transform(dims64, std::back_inserter(dims),
+                        &CheckedNarrowing<int64_t, int>);
+      status = wrap::miopenSetTensorDescriptor(obj.handle_, data_type, nd,
+                                               dims.data(), strides.data());
 
- private:
-  miopenConvolutionDescriptor_t handle_;  // Owned.
+      if (status != miopenStatusSuccess) {
+        LOG(FATAL) << "could not convert FilterDescriptor "
+                   << filter_descriptor.ToString()
+                   << " to miopen tensor descriptor: " << ToString(status);
+      }
+    } break;
+    default:
+      LOG(FATAL) << "Unsupported tensor format "
+                 << FilterLayoutString(filter_descriptor.layout());
+      break;
+  }
+  return obj;
+}
 
-  ScopedConvolutionDescriptor(const ScopedConvolutionDescriptor&) = delete;
-  void operator=(const ScopedConvolutionDescriptor&) = delete;
-};
+absl::StatusOr<ScopedConvolutionDescriptor> scope(
+    const ConvolutionDescriptor& convolution_descriptor) {
+  ScopedConvolutionDescriptor obj;
+  auto status = wrap::miopenCreateConvolutionDescriptor(&obj.handle_);
+  if (status != miopenStatusSuccess) {
+    LOG(FATAL) << "could not create miopen convolution descriptor: "
+               << ToString(status);
+  }
+  const auto& strides64 = convolution_descriptor.strides();
+  const auto& padding64 = convolution_descriptor.padding();
+  if (convolution_descriptor.pad_alignment() ==
+      dnn::PadAlignment::kTensorFlowPadding) {
+    LOG(ERROR) << "TensorFlow padding alignment is not supported.";
+  }
 
-// Turns a PoolingDescriptor structure into a miopen pooling descriptor handle
-// within a scope.
-class ScopedPoolingDescriptor {
- public:
-  ScopedPoolingDescriptor(const PoolingDescriptor& pooling_descriptor)
-      : handle_(nullptr) {
-    auto status = wrap::miopenCreatePoolingDescriptor(&handle_);
-    if (status != miopenStatusSuccess) {
-      LOG(FATAL) << "could not create miopen pooling descriptor: "
-                 << ToString(status);
-    }
+  // MIOpen requires arrays of ints.
+  std::vector<int> strides(convolution_descriptor.ndims());
+  std::vector<int> padding(convolution_descriptor.ndims());
+  std::transform(strides64.cbegin(), strides64.cend(), strides.begin(),
+                 &CheckedNarrowing<int64_t, int>);
+  std::transform(padding64.cbegin(), padding64.cend(), padding.begin(),
+                 &CheckedNarrowing<int64_t, int>);
 
-    absl::Span<const int64_t> strides64 = pooling_descriptor.strides();
-    absl::Span<const int64_t> padding64 = pooling_descriptor.padding();
-    absl::Span<const int64_t> shape64 = pooling_descriptor.window();
-
-    const int nd = pooling_descriptor.ndims();
-    std::vector<int> shape(nd);
-    std::vector<int> padding(nd);
-    std::vector<int> strides(nd);
-    std::transform(strides64.cbegin(), strides64.cend(), strides.begin(),
-                   &CheckedNarrowing<int64_t, int>);
-    std::transform(padding64.cbegin(), padding64.cend(), padding.begin(),
-                   &CheckedNarrowing<int64_t, int>);
-    std::transform(shape64.cbegin(), shape64.cend(), shape.begin(),
-                   &CheckedNarrowing<int64_t, int>);
-
-    status = wrap::miopenSetNdPoolingDescriptor(
-        handle_,
-        (pooling_descriptor.mode() == dnn::PoolingMode::kMaximum
-             ? miopenPoolingMax
-             : miopenPoolingAverage),
-        nd, shape.data(), padding.data(), strides.data());
-
-    // Note: The index type has to be uint32 type for now because MIOpen
-    // API assumes all input indexes to be the same type. Since a tensor
-    // descriptor can only use int32 type, the index type here need to be
-    // aligned with the tensor index type of the (input) tensor descritptor
-    status = wrap::miopenSetPoolingIndexType(handle_, miopenIndexUint32);
+  std::vector<int> upscale(convolution_descriptor.ndims());
+  const auto& dilations64 = convolution_descriptor.dilations();
+  std::transform(dilations64.cbegin(), dilations64.cend(), upscale.begin(),
+                 &CheckedNarrowing<int64_t, int>);
 
-    if (status != miopenStatusSuccess) {
-      LOG(FATAL) << "could not set miopen pooling descriptor: "
-                 << ToString(status);
-    }
+  status = wrap::miopenInitConvolutionNdDescriptor(
+      obj.handle_, convolution_descriptor.ndims(), padding.data(),
+      strides.data(), upscale.data(), miopenConvolution);
+  if (status != miopenStatusSuccess) {
+    LOG(FATAL) << "could not set miopen convolution descriptor: "
+               << ToString(status);
   }
-  ~ScopedPoolingDescriptor() {
-    auto status = wrap::miopenDestroyPoolingDescriptor(handle_);
+
+  VLOG(2) << "Requesting grouped convolution: "
+          << convolution_descriptor.group_count();
+  status = wrap::miopenSetConvolutionGroupCount(
+      obj.handle_, convolution_descriptor.group_count());
+  if (status != miopenStatusSuccess) {
+    LOG(FATAL) << "could not set miopen convolution group count: "
+               << ToString(status);
+  }
+
+#if (TF_ROCM_VERSION >= 50300)
+  if (RequireMIOpenDeterminism()) {
+    status = wrap::miopenSetConvolutionAttribute(
+        obj.handle_, MIOPEN_CONVOLUTION_ATTRIB_DETERMINISTIC, 1);
     if (status != miopenStatusSuccess) {
-      LOG(ERROR) << "could not destroy miopen pooling descriptor: "
+      LOG(FATAL) << "could not set miopen convolution attribute: "
                  << ToString(status);
     }
   }
+#endif
+  return obj;
+}
 
-  miopenPoolingDescriptor_t handle() const { return handle_; }
+absl::StatusOr<ScopedPoolingDescriptor> scope(
+    const PoolingDescriptor& pooling_descriptor) {
+  ScopedPoolingDescriptor obj;
+  auto status = wrap::miopenCreatePoolingDescriptor(&obj.handle_);
+  if (status != miopenStatusSuccess) {
+    LOG(FATAL) << "could not create miopen pooling descriptor: "
+               << ToString(status);
+  }
 
- private:
-  miopenPoolingDescriptor_t handle_;  // Owned.
+  absl::Span<const int64_t> strides64 = pooling_descriptor.strides();
+  absl::Span<const int64_t> padding64 = pooling_descriptor.padding();
+  absl::Span<const int64_t> shape64 = pooling_descriptor.window();
 
-  ScopedPoolingDescriptor(const ScopedPoolingDescriptor&) = delete;
-  void operator=(const ScopedPoolingDescriptor&) = delete;
-};
+  const int nd = pooling_descriptor.ndims();
+  std::vector<int> shape(nd);
+  std::vector<int> padding(nd);
+  std::vector<int> strides(nd);
+  std::transform(strides64.cbegin(), strides64.cend(), strides.begin(),
+                 &CheckedNarrowing<int64_t, int>);
+  std::transform(padding64.cbegin(), padding64.cend(), padding.begin(),
+                 &CheckedNarrowing<int64_t, int>);
+  std::transform(shape64.cbegin(), shape64.cend(), shape.begin(),
+                 &CheckedNarrowing<int64_t, int>);
 
-// Turns a NormalizeDescriptor structure into a miopen LRN descriptor handle.
-class ScopedNormalizeDescriptor {
- public:
-  ScopedNormalizeDescriptor(const NormalizeDescriptor& normalize_descriptor)
-      : handle_(nullptr) {
-    auto status = wrap::miopenCreateLRNDescriptor(&handle_);
-    if (status != miopenStatusSuccess) {
-      LOG(FATAL) << "could not create miopen LRN descriptor: "
-                 << ToString(status);
-    }
+  status = wrap::miopenSetNdPoolingDescriptor(
+      obj.handle_,
+      (pooling_descriptor.mode() == dnn::PoolingMode::kMaximum
+           ? miopenPoolingMax
+           : miopenPoolingAverage),
+      nd, shape.data(), padding.data(), strides.data());
 
-    // The range specifies that the indices in the closed range
-    // [i - range, i + range] should be included in the normalization for index
-    // i. The lrnN value is the total number of elements in the range, so
-    // lrnN = 2*range + 1.
-    unsigned lrn_N = 2 * normalize_descriptor.range() + 1;
-
-    // Note that SE defines the normalization operation as
-    //
-    //  U_i = V_i / ((bias +  alpha      * (sum_j V_j^2)) ^ beta)
-    //
-    // but MIOpen defines it as
-    //
-    //  U_i = V_i / ((bias + (alpha / n) * (sum_j V_j^2)) ^ beta)
-    //
-    // i.e. there is a factor of n difference between the meaning of the alphas
-    // in the two contexts. The MIOpen alpha is n times the SE alpha.
-    double lrn_alpha = lrn_N * normalize_descriptor.alpha();
-
-    double lrn_beta = normalize_descriptor.beta();
-    double lrn_k = normalize_descriptor.bias();
-    status = wrap::miopenSetLRNDescriptor(handle_, miopenLRNCrossChannel, lrn_N,
-                                          lrn_alpha, lrn_beta, lrn_k);
-    if (status != miopenStatusSuccess) {
-      LOG(FATAL) << "could not set miopen LRN descriptor: " << ToString(status);
-    }
-  }
+  // Note: The index type has to be uint32 type for now because MIOpen
+  // API assumes all input indexes to be the same type. Since a tensor
+  // descriptor can only use int32 type, the index type here need to be
+  // aligned with the tensor index type of the (input) tensor descritptor
+  status = wrap::miopenSetPoolingIndexType(obj.handle_, miopenIndexUint32);
 
-  ~ScopedNormalizeDescriptor() {
-    auto status = wrap::miopenDestroyLRNDescriptor(handle_);
-    if (status != miopenStatusSuccess) {
-      LOG(ERROR) << "could not destroy miopen LRN descriptor: "
-                 << ToString(status);
-    }
+  if (status != miopenStatusSuccess) {
+    LOG(FATAL) << "could not set miopen pooling descriptor: "
+               << ToString(status);
   }
+  return obj;
+}
 
-  miopenLRNDescriptor_t handle() const { return handle_; }
+absl::StatusOr<ScopedNormalizeDescriptor> scope(
+    const NormalizeDescriptor& normalize_descriptor) {
+  ScopedNormalizeDescriptor obj;
+  auto status = wrap::miopenCreateLRNDescriptor(&obj.handle_);
+  if (status != miopenStatusSuccess) {
+    LOG(FATAL) << "could not create miopen LRN descriptor: "
+               << ToString(status);
+  }
 
- private:
-  miopenLRNDescriptor_t handle_;  // Owned.
+  // The range specifies that the indices in the closed range
+  // [i - range, i + range] should be included in the normalization for index
+  // i. The lrnN value is the total number of elements in the range, so
+  // lrnN = 2*range + 1.
+  unsigned lrn_N = 2 * normalize_descriptor.range() + 1;
 
-  ScopedNormalizeDescriptor(const ScopedNormalizeDescriptor&) = delete;
-  void operator=(const ScopedNormalizeDescriptor&) = delete;
-};
+  // Note that SE defines the normalization operation as
+  //
+  //  U_i = V_i / ((bias +  alpha      * (sum_j V_j^2)) ^ beta)
+  //
+  // but MIOpen defines it as
+  //
+  //  U_i = V_i / ((bias + (alpha / n) * (sum_j V_j^2)) ^ beta)
+  //
+  // i.e. there is a factor of n difference between the meaning of the alphas
+  // in the two contexts. The MIOpen alpha is n times the SE alpha.
+  double lrn_alpha = lrn_N * normalize_descriptor.alpha();
+
+  double lrn_beta = normalize_descriptor.beta();
+  double lrn_k = normalize_descriptor.bias();
+  status = wrap::miopenSetLRNDescriptor(obj.handle_, miopenLRNCrossChannel,
+                                        lrn_N, lrn_alpha, lrn_beta, lrn_k);
+  if (status != miopenStatusSuccess) {
+    LOG(FATAL) << "could not set miopen LRN descriptor: " << ToString(status);
+  }
+  return obj;
+}
 
 // Turns a activation mode into a miopen activation mode descriptor with a scope
 // around it
-class ScopedActivationDescriptor {
- public:
-  ScopedActivationDescriptor(dnn::ActivationMode activation_mode)
-      : handle_(nullptr),
-        miopen_activation_mode_(miopenActivationPASTHRU),
-        alpha_(0.0),
-        beta_(0.0),
-        gamma_(0.0) {
-    auto status = wrap::miopenCreateActivationDescriptor(&handle_);
+struct ScopedActivationDescriptor
+    : ScopedDescriptor<miopenActivationDescriptor_t> {
+  static absl::StatusOr<ScopedActivationDescriptor> Create(
+      dnn::ActivationMode activation_mode, double alpha = 0.0) {
+    ScopedActivationDescriptor obj;
+    obj.alpha_ = alpha;
+    auto status = wrap::miopenCreateActivationDescriptor(&obj.handle_);
     if (status != miopenStatusSuccess) {
-      LOG(FATAL) << "call to miopenCreateActivationDescriptor failed: "
-                 << ToString(status);
+      return absl::InternalError(
+          "call to miopenCreateActivationDescriptor failed: " +
+          ToString(status));
     } else {
       switch (activation_mode) {
         case dnn::ActivationMode::kNone:
-          miopen_activation_mode_ = miopenActivationPASTHRU;
+          obj.miopen_activation_mode_ = miopenActivationPASTHRU;
           break;
 
         case dnn::ActivationMode::kSigmoid:
-          miopen_activation_mode_ = miopenActivationLOGISTIC;
+          obj.miopen_activation_mode_ = miopenActivationLOGISTIC;
           break;
 
         case dnn::ActivationMode::kRelu:
-          miopen_activation_mode_ = miopenActivationRELU;
+        case dnn::ActivationMode::kReluX:
+          obj.miopen_activation_mode_ = miopenActivationRELU;
           break;
 
         case dnn::ActivationMode::kRelu6:
-          miopen_activation_mode_ = miopenActivationRELU;
-          alpha_ = 6.0;
+          obj.miopen_activation_mode_ = miopenActivationRELU;
+          obj.alpha_ = 6.0;
           break;
 
         case dnn::ActivationMode::kTanh:
-          miopen_activation_mode_ = miopenActivationTANH;
+          obj.miopen_activation_mode_ = miopenActivationTANH;
           break;
 
-        default:
-          LOG(FATAL) << "Activation mode ("
-                     << dnn::ActivationModeString(activation_mode)
-                     << ") not yet implemented";
+        case dnn::ActivationMode::kElu:
+          obj.miopen_activation_mode_ = miopenActivationELU;
+          break;
+
+        case dnn::ActivationMode::kLeakyRelu:
+          obj.miopen_activation_mode_ = miopenActivationLEAKYRELU;
           break;
+          // Check with MIOpen re: support: kBandPass, kGeluExact
+
+        default:
+          VLOG(1) << "Activation mode ("
+                  << dnn::ActivationModeString(activation_mode)
+                  << ") not yet implemented";
+          return absl::InternalError("Activation not implemented");
       }
 
       status = wrap::miopenSetActivationDescriptor(
-          handle_, miopen_activation_mode_, alpha_, beta_, gamma_);
+          obj.handle_, obj.miopen_activation_mode_, obj.alpha_, obj.beta_,
+          obj.gamma_);
       if (status != miopenStatusSuccess) {
-        LOG(FATAL) << "call to miopenSetActivationDescriptor failed: "
-                   << ToString(status);
+        return absl::InternalError(
+            "call to miopenSetActivationDescriptor failed: " +
+            ToString(status));
       }
     }
+    return obj;
   }
-
-  ~ScopedActivationDescriptor() {
-    auto status = wrap::miopenDestroyActivationDescriptor(handle_);
-    if (status != miopenStatusSuccess) {
-      LOG(FATAL) << "call to miopenDestroyActivationDescriptor failed: "
-                 << ToString(status);
-    }
+  ScopedActivationDescriptor(ScopedActivationDescriptor&& other)
+      : ScopedDescriptor<miopenActivationDescriptor_t>(std::move(other)) {
+    miopen_activation_mode_ = other.miopen_activation_mode_;
+    alpha_ = other.alpha_;
+    beta_ = other.beta_;
+    gamma_ = other.gamma_;
   }
 
-  miopenActivationDescriptor_t handle() const { return handle_; }
-
   uint64_t GetHashValue() {
     uint64_t hash_value = tsl::hash<int>()(miopen_activation_mode_);
     hash_value = tsl::Hash64Combine(hash_value, tsl::hash<double>()(alpha_));
@@ -1226,13 +1209,12 @@ class ScopedActivationDescriptor {
     return hash_value;
   }
 
- private:
-  miopenActivationDescriptor_t handle_;  // Owned.
-
-  ScopedActivationDescriptor(const ScopedActivationDescriptor&) = delete;
-  void operator=(const ScopedActivationDescriptor&) = delete;
+  ScopedActivationDescriptor()
+      : miopen_activation_mode_(miopenActivationPASTHRU),
+        alpha_(0.0),
+        beta_(0.0),
+        gamma_(0.0) {}
 
- public:
   // caching these values here to avoid calling miopenGetActivationDescriptor
   // to do the same. miopenGetActivationDescriptor gets called twice during each
   // call to execute a fusion plan (that involves the activation op)...once call
@@ -1262,6 +1244,7 @@ class ScopedFusionPlanBase {
   }
 
   virtual ~ScopedFusionPlanBase() {
+    if (fusion_args_ == nullptr) return;
     auto status = wrap::miopenDestroyOperatorArgs(fusion_args_);
     if (status != miopenStatusSuccess) {
       LOG(FATAL) << "call to miopenDestroyoperatorArgs failed: "
@@ -1286,7 +1269,6 @@ class ScopedFusionPlanBase {
 
   bool CompilationSucceeded() { return fusion_plan_compiled_; }
 
- protected:
   miopenStatus_t SetConvolutionArgs(const int op_idx, const float* alpha,
                                     const float* beta, const void* data) {
     miopenFusionOpDescriptor_t conv_op;
@@ -1437,11 +1419,20 @@ class ScopedFusionPlanBase {
     return status;
   }
 
+ public:
   miopenHandle_t miopen_handle_;
   miopenFusionPlanDescriptor_t fusion_plan_;
   miopenOperatorArgs_t fusion_args_;  // Owned.
   bool fusion_plan_compiled_;
 
+  ScopedFusionPlanBase(ScopedFusionPlanBase&& other) {
+    miopen_handle_ = other.miopen_handle_;
+    fusion_plan_ = other.fusion_plan_;
+    fusion_args_ = other.fusion_args_;
+    other.fusion_args_ = nullptr;
+    fusion_plan_compiled_ = other.fusion_plan_compiled_;
+  }
+
   ScopedFusionPlanBase(const ScopedFusionPlanBase&) = delete;
   void operator=(const ScopedFusionPlanBase&) = delete;
 };
@@ -1450,86 +1441,101 @@ class ScopedFusionPlanBase {
 class ScopedFusionPlanConvolutionBiasActivation : public ScopedFusionPlanBase {
  public:
   ScopedFusionPlanConvolutionBiasActivation(
+      miopenHandle_t miopen_handle, miopenTensorDescriptor_t input_descriptor)
+      : ScopedFusionPlanBase(miopen_handle, miopenVerticalFusion,
+                             input_descriptor) {}
+
+  ScopedFusionPlanConvolutionBiasActivation(
+      ScopedFusionPlanConvolutionBiasActivation&& other)
+      : ScopedFusionPlanBase(std::move(other)) {
+    conv_op = other.conv_op;
+    bias_op = other.bias_op;
+    actv_op = other.actv_op;
+  }
+
+  static absl::StatusOr<ScopedFusionPlanConvolutionBiasActivation> Create(
       miopenHandle_t miopen_handle, miopenTensorDescriptor_t input_descriptor,
       miopenTensorDescriptor_t filter_descriptor,
       miopenConvolutionDescriptor_t conv_descriptor,
       miopenTensorDescriptor_t bias_descriptor,
-      ScopedActivationDescriptor& activation_descriptor)
-      : ScopedFusionPlanBase(miopen_handle, miopenVerticalFusion,
-                             input_descriptor) {
-    uint64_t hash = GetFusionOpHashValue(
-        miopen_handle, input_descriptor, filter_descriptor, conv_descriptor,
-        bias_descriptor, activation_descriptor);
+      ScopedActivationDescriptor& act_descriptor) {
+    ScopedFusionPlanConvolutionBiasActivation obj(miopen_handle,
+                                                  input_descriptor);
+
+    VLOG(2) << "Fusion Plan compile begin";
+
+    uint64_t hash =
+        GetFusionOpHashValue(miopen_handle, input_descriptor, filter_descriptor,
+                             conv_descriptor, bias_descriptor, act_descriptor);
 
     bool is_compiled = CachedFusionPlans::FindOrCreate(
-        hash, &fusion_plan_, miopenVerticalFusion, input_descriptor);
+        hash, &obj.fusion_plan_, miopenVerticalFusion, input_descriptor);
+    if (is_compiled) VLOG(2) << "Cache hit";
     if (!is_compiled) {
-      miopenFusionOpDescriptor_t conv_op;
       auto status = wrap::miopenCreateOpConvForward(
-          fusion_plan_, &conv_op, conv_descriptor, filter_descriptor);
-      if (status != miopenStatusSuccess) {
-        LOG(FATAL) << "call to miopenCreateOpConvForward failed: "
-                   << ToString(status);
-      }
+          obj.fusion_plan_, &obj.conv_op, conv_descriptor, filter_descriptor);
+      if (status != miopenStatusSuccess)
+        return absl::InternalError("miopenCreateOpConvForward failed: " +
+                                   ToString(status));
 
-      miopenFusionOpDescriptor_t bias_op;
-      status = wrap::miopenCreateOpBiasForward(fusion_plan_, &bias_op,
+      status = wrap::miopenCreateOpBiasForward(obj.fusion_plan_, &obj.bias_op,
                                                bias_descriptor);
-      if (status != miopenStatusSuccess) {
-        LOG(FATAL) << "call to miopenCreateOpBiasForward failed: "
-                   << ToString(status);
-      }
-
-      miopenFusionOpDescriptor_t actv_op;
-      status = wrap::miopenCreateOpActivationForward(
-          fusion_plan_, &actv_op,
-          activation_descriptor.miopen_activation_mode_);
-      if (status != miopenStatusSuccess) {
-        LOG(FATAL) << "call to miopenCreateOpActivationForward failed: "
-                   << ToString(status);
+      if (status != miopenStatusSuccess)
+        return absl::InternalError("miopenCreateOpBiasForward failed: " +
+                                   ToString(status));
+
+      if (act_descriptor.miopen_activation_mode_ != miopenActivationPASTHRU) {
+        status = wrap::miopenCreateOpActivationForward(
+            obj.fusion_plan_, &obj.actv_op,
+            act_descriptor.miopen_activation_mode_);
+        if (status != miopenStatusSuccess)
+          return absl::InternalError(
+              "miopenCreateOpActivationForward failed: " + ToString(status));
       }
 
-      status = wrap::miopenCompileFusionPlan(miopen_handle_, fusion_plan_);
+      status = wrap::miopenCompileFusionPlan(miopen_handle, obj.fusion_plan_);
       if (status != miopenStatusSuccess) {
         VLOG(2) << "call to miopenCompileFusionPlan (CBA) failed: "
                 << ToString(status);
 
         CachedFusionPlans::MarkFusionPlanUnsupported(hash);
       } else {
-        VLOG(2) << "Fusion Plan compile succedded (CBA) ";
-        fusion_plan_compiled_ = true;
+        VLOG(2) << "Fusion Plan compile succeeded (CBA) ";
+        obj.fusion_plan_compiled_ = true;
       }
     } else {
       // fusion plan was already compiled...check whether it failed to compile
-      fusion_plan_compiled_ = !CachedFusionPlans::IsUnsupportedFusionPlan(hash);
+      obj.fusion_plan_compiled_ =
+          !CachedFusionPlans::IsUnsupportedFusionPlan(hash);
     }
+    return obj;
   }
 
   miopenStatus_t SetConvolutionArgs(const void* filter_data) {
-    float alpha = 1.0;
-    float beta = 0.0;
+    static const float alpha = 1.0;
+    static const float beta = 0.0;
     return ScopedFusionPlanBase::SetConvolutionArgs(k_conv_op_idx, &alpha,
                                                     &beta, filter_data);
   }
 
   miopenStatus_t SetBiasArgs(const void* bias_data) {
-    float alpha = 1.0;
-    float beta = 0.0;
+    static const float alpha = 1.0;
+    static const float beta = 0.0;
     return ScopedFusionPlanBase::SetBiasArgs(k_bias_op_idx, &alpha, &beta,
                                              bias_data);
   }
 
   miopenStatus_t SetActivationForwardArgs(
       ScopedActivationDescriptor& activation_descriptor) {
-    float alpha = 1.0;
-    float beta = 0.0;
+    static const float alpha = 1.0;
+    static const float beta = 0.0;
 
     return ScopedFusionPlanBase::SetActivationForwardArgs(
         k_actv_op_idx, &alpha, &beta, activation_descriptor.alpha_,
         activation_descriptor.beta_, activation_descriptor.gamma_);
   }
 
-  uint64_t GetFusionOpHashValue(
+  static uint64_t GetFusionOpHashValue(
       miopenHandle_t miopen_handle, miopenTensorDescriptor_t input_descriptor,
       miopenTensorDescriptor_t filter_descriptor,
       miopenConvolutionDescriptor_t conv_descriptor,
@@ -1550,6 +1556,11 @@ class ScopedFusionPlanConvolutionBiasActivation : public ScopedFusionPlanBase {
     return hash_value;
   }
 
+ public:
+  miopenFusionOpDescriptor_t conv_op;
+  miopenFusionOpDescriptor_t bias_op;
+  miopenFusionOpDescriptor_t actv_op;
+
  private:
   const int k_conv_op_idx = 0;
   const int k_bias_op_idx = 1;
@@ -1718,8 +1729,8 @@ class ScopedFusionPlanBatchNormActivationForward : public ScopedFusionPlanBase {
                                          void* batch_mean, void* batch_var,
                                          void* saved_mean, void* saved_var,
                                          double epsilon) {
-    float alpha = 1.0;
-    float beta = 0.0;
+    static const float alpha = 1.0;
+    static const float beta = 0.0;
     return ScopedFusionPlanBase::SetBatchNormForwardArgs(
         k_batchnorm_op_idx, &alpha, &beta, scale, offset, batch_mean, batch_var,
         saved_mean, saved_var, /*exponential_average_factor=*/1.0, epsilon);
@@ -1727,8 +1738,8 @@ class ScopedFusionPlanBatchNormActivationForward : public ScopedFusionPlanBase {
 
   miopenStatus_t SetActivationForwardArgs(
       ScopedActivationDescriptor& activation_descriptor) {
-    float alpha = 1.0;
-    float beta = 0.0;
+    static const float alpha = 1.0;
+    static const float beta = 0.0;
 
     return ScopedFusionPlanBase::SetActivationForwardArgs(
         k_actv_op_idx, &alpha, &beta, activation_descriptor.alpha_,
@@ -1867,6 +1878,24 @@ class ScopedFusionPlanBatchNormActivationBackward
 };
 
 namespace {
+
+const char* getTypeName(dnn::DataType data_type) {
+  switch (data_type) {
+    case dnn::DataType::kBF16:
+      return "BF16";
+    case dnn::DataType::kFloat:
+      return "F32";
+    case dnn::DataType::kHalf:
+      return "F16";
+    case dnn::DataType::kInt8:
+      return "I8";
+    case dnn::DataType::kDouble:
+      return "F64";
+    default:
+      return "Unknown";
+  }
+}
+
 miopenDataType_t ToMIOpenDataType(
     dnn::DataType data_type,
     dnn::DataLayout data_layout = dnn::DataLayout::kBatchDepthYX) {
@@ -3246,21 +3275,22 @@ class RocmConvRunner : public dnn::ConvRunner {
  public:
   RocmConvRunner(GpuExecutor* parent, MIOpenAccess* miopen, int64_t algo_id,
                  size_t workspace_size, dnn::ConvolutionKind kind,
-                 dnn::DataType input_type, bool use_immediate_mode,
-                 BatchDescriptor input_descriptor,
-                 BatchDescriptor output_descriptor,
-                 FilterDescriptor filter_descriptor,
-                 ConvolutionDescriptor conv_descriptor)
+                 dnn::DataType input_type, dnn::DataType output_type,
+                 bool use_immediate_mode,
+                 ScopedTensorDescriptor& scoped_input_desc,
+                 ScopedTensorDescriptor& scoped_output_desc,
+                 ScopedFilterDescriptor& scoped_filter_desc,
+                 ScopedConvolutionDescriptor& scoped_conv_desc)
       : parent_(parent),
         miopen_(miopen),
         algo_id_(algo_id),
         workspace_size_(workspace_size),
         kind_(kind),
         use_immediate_mode_(use_immediate_mode),
-        input_desc_{input_descriptor, ToMIOpenDataType(input_type)},
-        output_desc_{output_descriptor, ToMIOpenDataType(input_type)},
-        filter_desc_{filter_descriptor, ToMIOpenDataType(input_type)},
-        conv_desc_{conv_descriptor, ToMIOpenDataType(input_type)} {
+        input_desc_(std::move(scoped_input_desc)),
+        output_desc_(std::move(scoped_output_desc)),
+        filter_desc_(std::move(scoped_filter_desc)),
+        conv_desc_(std::move(scoped_conv_desc)) {
     bool is_backprop = ((kind == dnn::ConvolutionKind::BACKWARD_DATA) ||
                         (kind == dnn::ConvolutionKind::BACKWARD_FILTER));
     // #if TF_ROCM_VERSION >= 50000
@@ -3440,12 +3470,10 @@ absl::Status MIOpenSupport::GetConvolveRunners(
 
   std::vector<dnn::ProfileResult> profile_results;
   if (!GetMIOpenConvolveAlgorithms(
-          kind, input_type, stream, input_descriptor, input_data,
+          kind, input_type, output_type, stream, input_descriptor, input_data,
           filter_descriptor, filter_data, output_descriptor, output_data,
-          convolution_descriptor, scratch_allocator, &profile_results)) {
-    return absl::UnknownError(
-        "GetConvolveRunners: GetMIOpenConvolveAlgorithms failed");
-  }
+          convolution_descriptor, scratch_allocator, &profile_results))
+    return absl::InternalError("GetMIOpenConvolveAlgorithms failure");
 
   for (const auto& profile_result : profile_results) {
     TF_ASSIGN_OR_RETURN(
@@ -3467,27 +3495,24 @@ MIOpenSupport::ConvolveRunnerFromDesc(
     const dnn::FilterDescriptor& filter_descriptor,
     const dnn::BatchDescriptor& output_descriptor,
     const dnn::ConvolutionDescriptor& convolution_descriptor) {
-  if (input_type != output_type) {
-    return absl::UnimplementedError(
-        absl::StrFormat("MIOpen backend does not support different input and "
-                        "output types: %d != %d",
-                        input_type, output_type));
-  }
-
   auto workspace_size = algorithm_desc.workspace_size();
-  if (!workspace_size) {
-    return absl::InvalidArgumentError(
-        "MIOpenSupport::ConvolveRunnerFromDesc requires "
-        "AlgorithmProto.workspace_size, but it was missing.");
-  }
+  TF_ASSIGN_OR_RETURN(auto scoped_input_desc,
+                      scope(input_descriptor, ToMIOpenDataType(input_type)));
+  TF_ASSIGN_OR_RETURN(auto scoped_output_desc,
+                      scope(output_descriptor, ToMIOpenDataType(output_type)));
+  TF_ASSIGN_OR_RETURN(auto scoped_filter_desc,
+                      scope(filter_descriptor, ToMIOpenDataType(input_type)));
+  TF_ASSIGN_OR_RETURN(auto scoped_conv_desc, scope(convolution_descriptor));
+
   return {std::make_unique<RocmConvRunner>(
       parent_, miopen_.get(), algorithm_desc.algo_id(), *workspace_size, kind,
-      input_type, use_immediate_mode_, input_descriptor, output_descriptor,
-      filter_descriptor, convolution_descriptor)};
+      input_type, output_type, use_immediate_mode_, scoped_input_desc,
+      scoped_output_desc, scoped_filter_desc, scoped_conv_desc)};
 }
 
 bool MIOpenSupport::GetMIOpenConvolveAlgorithms(
-    dnn::ConvolutionKind kind, dnn::DataType element_type, Stream* stream,
+    dnn::ConvolutionKind kind, dnn::DataType input_type,
+    dnn::DataType output_type, Stream* stream,
     const dnn::BatchDescriptor& input_descriptor, DeviceMemoryBase input_data,
     const dnn::FilterDescriptor& filter_descriptor,
     DeviceMemoryBase filter_data, const dnn::BatchDescriptor& output_descriptor,
@@ -3497,19 +3522,22 @@ bool MIOpenSupport::GetMIOpenConvolveAlgorithms(
     std::vector<dnn::ProfileResult>* out_algorithms) {
   return use_immediate_mode_
              ? GetMIOpenConvolveAlgorithmsImmediateMode(
-                   kind, element_type, stream, input_descriptor, input_data,
-                   filter_descriptor, filter_data, output_descriptor,
-                   output_data, convolution_descriptor, scratch_allocator,
-                   out_algorithms)
+                   kind, input_type, output_type, stream, input_descriptor,
+                   input_data, filter_descriptor, filter_data,
+                   output_descriptor, output_data, convolution_descriptor,
+                   scratch_allocator, out_algorithms)
+                   .ok()
              : GetMIOpenConvolveAlgorithmsFindMode(
-                   kind, element_type, stream, input_descriptor, input_data,
-                   filter_descriptor, filter_data, output_descriptor,
-                   output_data, convolution_descriptor, scratch_allocator,
-                   out_algorithms);
+                   kind, input_type, output_type, stream, input_descriptor,
+                   input_data, filter_descriptor, filter_data,
+                   output_descriptor, output_data, convolution_descriptor,
+                   scratch_allocator, out_algorithms)
+                   .ok();
 }
 
-bool MIOpenSupport::GetMIOpenConvolveAlgorithmsImmediateMode(
-    dnn::ConvolutionKind kind, dnn::DataType element_type, Stream* stream,
+absl::Status MIOpenSupport::GetMIOpenConvolveAlgorithmsImmediateMode(
+    dnn::ConvolutionKind kind, dnn::DataType input_type,
+    dnn::DataType output_type, Stream* stream,
     const dnn::BatchDescriptor& input_descriptor, DeviceMemoryBase input_data,
     const dnn::FilterDescriptor& filter_descriptor,
     DeviceMemoryBase filter_data, const dnn::BatchDescriptor& output_descriptor,
@@ -3519,14 +3547,13 @@ bool MIOpenSupport::GetMIOpenConvolveAlgorithmsImmediateMode(
     std::vector<dnn::ProfileResult>* out_algorithms) {
   auto miopen = miopen_->GetHandle(parent_, stream);
 
-  ScopedTensorDescriptor input_nd{input_descriptor,
-                                  ToMIOpenDataType(element_type)};
-  ScopedTensorDescriptor output_nd{output_descriptor,
-                                   ToMIOpenDataType(element_type)};
-  ScopedFilterDescriptor filter{filter_descriptor,
-                                ToMIOpenDataType(element_type)};
-  ScopedConvolutionDescriptor conv{convolution_descriptor,
-                                   ToMIOpenDataType(element_type)};
+  TF_ASSIGN_OR_RETURN(auto input_nd,
+                      scope(input_descriptor, ToMIOpenDataType(input_type)));
+  TF_ASSIGN_OR_RETURN(auto output_nd,
+                      scope(output_descriptor, ToMIOpenDataType(output_type)));
+  TF_ASSIGN_OR_RETURN(auto filter,
+                      scope(filter_descriptor, ToMIOpenDataType(input_type)));
+  TF_ASSIGN_OR_RETURN(auto conv, scope(convolution_descriptor));
 
   bool is_backprop = ((kind == dnn::ConvolutionKind::BACKWARD_DATA) ||
                       (kind == dnn::ConvolutionKind::BACKWARD_FILTER));
@@ -3534,24 +3561,24 @@ bool MIOpenSupport::GetMIOpenConvolveAlgorithmsImmediateMode(
   //                   (call_context == dnn::CallContext::kBackpropFilter);
 
 #if TF_ROCM_VERSION >= 50000
-  if (is_backprop && (ToMIOpenDataType(element_type) == miopenHalf)) {
+  if (is_backprop && (ToMIOpenDataType(input_type) == miopenHalf)) {
     wrap::miopenSetConvolutionAttribute(
         conv.handle(), MIOPEN_CONVOLUTION_ATTRIB_FP16_ALT_IMPL, 1);
   }
 #endif
-  // First determine the number of algorityhms available
+  // First determine the number of algorithms available
   size_t maxSolutionCount = 0;
 
   switch (kind) {
-    case dnn::ConvolutionKind::FORWARD: {
+    case dnn::ConvolutionKind::FORWARD:
+    case dnn::ConvolutionKind::FORWARD_BIAS_ACTIVATION: {
       auto status = wrap::miopenConvolutionForwardGetSolutionCount(
           miopen.handle(), filter.handle(), input_nd.handle(), conv.handle(),
           output_nd.handle(), &maxSolutionCount);
       if (status != miopenStatusSuccess) {
-        LOG(FATAL)
-            << "call to miopenConvolutionForwardGetSolutionCount failed: "
-            << ToString(status);
-        return false;
+        return absl::InternalError(
+            "call to miopenConvolutionForwardGetSolutionCount failed: " +
+            ToString(status));
       }
       break;
     }
@@ -3560,10 +3587,10 @@ bool MIOpenSupport::GetMIOpenConvolveAlgorithmsImmediateMode(
           miopen.handle(), output_nd.handle(), filter.handle(), conv.handle(),
           input_nd.handle(), &maxSolutionCount);
       if (status != miopenStatusSuccess) {
-        LOG(FATAL) << "call to miopenConvolutionBackwardDataGetSolutionCount "
-                      "failed: "
-                   << ToString(status);
-        return false;
+        return absl::InternalError(
+            "call to miopenConvolutionBackwardDataGetSolutionCount "
+            "failed: " +
+            ToString(status));
       }
       break;
     }
@@ -3572,18 +3599,16 @@ bool MIOpenSupport::GetMIOpenConvolveAlgorithmsImmediateMode(
           miopen.handle(), output_nd.handle(), input_nd.handle(), conv.handle(),
           filter.handle(), &maxSolutionCount);
       if (status != miopenStatusSuccess) {
-        LOG(FATAL)
-            << "call to miopenConvolutionBackwardWeightsGetSolutionCount "
-               "failed: "
-            << ToString(status);
-        return false;
+        return absl::InternalError(
+            "call to miopenConvolutionBackwardWeightsGetSolutionCount "
+            "failed: " +
+            ToString(status));
       }
       break;
     }
     default: {
-      LOG(FATAL) << "Unexpected convolution kind " << static_cast<int>(kind);
-      return false;
-      break;
+      return absl::InternalError("Unexpected convolution kind " +
+                                 std::to_string(static_cast<int>(kind)));
     }
   }
 
@@ -3608,9 +3633,9 @@ bool MIOpenSupport::GetMIOpenConvolveAlgorithmsImmediateMode(
           solutions.get());
 
       if (status != miopenStatusSuccess) {
-        LOG(FATAL) << "call to miopenConvolutionForwardGetSolution failed: "
-                   << ToString(status);
-        return false;
+        return absl::InternalError(
+            "call to miopenConvolutionForwardGetSolution failed: " +
+            ToString(status));
       }
 
       VLOG(kConvDebugVlogLevel)
@@ -3629,10 +3654,9 @@ bool MIOpenSupport::GetMIOpenConvolveAlgorithmsImmediateMode(
             output_nd.handle(), solution.solution_id);
 
         if (status != miopenStatusSuccess) {
-          LOG(FATAL)
-              << "call to miopenConvolutionForwardCompileSolution failed: "
-              << ToString(status);
-          return false;
+          return absl::InternalError(
+              "call to miopenConvolutionForwardCompileSolution failed: " +
+              ToString(status));
         }
 
         out_algorithms->emplace_back(
@@ -3646,10 +3670,9 @@ bool MIOpenSupport::GetMIOpenConvolveAlgorithmsImmediateMode(
           miopen.handle(), output_nd.handle(), filter.handle(), conv.handle(),
           input_nd.handle(), maxSolutionCount, &solutionCount, solutions.get());
       if (status != miopenStatusSuccess) {
-        LOG(FATAL)
-            << "call to miopenConvolutionBackwardDataGetSolution failed: "
-            << ToString(status);
-        return false;
+        return absl::InternalError(
+            "call to miopenConvolutionBackwardDataGetSolution failed: " +
+            ToString(status));
       }
 
       VLOG(kConvDebugVlogLevel)
@@ -3668,10 +3691,10 @@ bool MIOpenSupport::GetMIOpenConvolveAlgorithmsImmediateMode(
             input_nd.handle(), solution.solution_id);
 
         if (status != miopenStatusSuccess) {
-          LOG(FATAL) << " call to miopenConvolutionBackwardDataCompileSolution "
-                        "failed: "
-                     << ToString(status);
-          return false;
+          return absl::InternalError(
+              " call to miopenConvolutionBackwardDataCompileSolution "
+              "failed: " +
+              ToString(status));
         }
 
         out_algorithms->emplace_back(
@@ -3684,10 +3707,9 @@ bool MIOpenSupport::GetMIOpenConvolveAlgorithmsImmediateMode(
           miopen.handle(), output_nd.handle(), input_nd.handle(), conv.handle(),
           filter.handle(), maxSolutionCount, &solutionCount, solutions.get());
       if (status != miopenStatusSuccess) {
-        LOG(FATAL)
-            << "call to miopenConvolutionBackwardWeightsGetSolution failed: "
-            << ToString(status);
-        return false;
+        return absl::InternalError(
+            "call to miopenConvolutionBackwardWeightsGetSolution failed: " +
+            ToString(status));
       }
 
       VLOG(kConvDebugVlogLevel)
@@ -3706,11 +3728,10 @@ bool MIOpenSupport::GetMIOpenConvolveAlgorithmsImmediateMode(
             conv.handle(), filter.handle(), solution.solution_id);
 
         if (status != miopenStatusSuccess) {
-          LOG(FATAL)
-              << "call to miopenConvolutionBackwardWeightsCompileSolution "
-                 "failed: "
-              << ToString(status);
-          return false;
+          return absl::InternalError(
+              "call to miopenConvolutionBackwardWeightsCompileSolution "
+              "failed: " +
+              ToString(status));
         }
 
         out_algorithms->emplace_back(
@@ -3719,17 +3740,17 @@ bool MIOpenSupport::GetMIOpenConvolveAlgorithmsImmediateMode(
       break;
     }
     default: {
-      LOG(FATAL) << "Unexpected convolution kind " << static_cast<int>(kind);
-      return false;
-      break;
+      return absl::InternalError("Unexpected convolution kind " +
+                                 std::to_string(static_cast<int>(kind)));
     }
   }
 
-  return true;
+  return absl::OkStatus();
 }
 
-bool MIOpenSupport::GetMIOpenConvolveAlgorithmsFindMode(
-    dnn::ConvolutionKind kind, dnn::DataType element_type, Stream* stream,
+absl::Status MIOpenSupport::GetMIOpenConvolveAlgorithmsFindMode(
+    dnn::ConvolutionKind kind, dnn::DataType input_type,
+    dnn::DataType output_type, Stream* stream,
     const dnn::BatchDescriptor& input_descriptor, DeviceMemoryBase input_data,
     const dnn::FilterDescriptor& filter_descriptor,
     DeviceMemoryBase filter_data, const dnn::BatchDescriptor& output_descriptor,
@@ -3739,14 +3760,13 @@ bool MIOpenSupport::GetMIOpenConvolveAlgorithmsFindMode(
     std::vector<dnn::ProfileResult>* out_algorithms) {
   auto miopen = miopen_->GetHandle(parent_, stream);
 
-  ScopedTensorDescriptor input_nd{input_descriptor,
-                                  ToMIOpenDataType(element_type)};
-  ScopedTensorDescriptor output_nd{output_descriptor,
-                                   ToMIOpenDataType(element_type)};
-  ScopedFilterDescriptor filter{filter_descriptor,
-                                ToMIOpenDataType(element_type)};
-  ScopedConvolutionDescriptor conv{convolution_descriptor,
-                                   ToMIOpenDataType(element_type)};
+  TF_ASSIGN_OR_RETURN(auto input_nd,
+                      scope(input_descriptor, ToMIOpenDataType(input_type)));
+  TF_ASSIGN_OR_RETURN(auto output_nd,
+                      scope(output_descriptor, ToMIOpenDataType(output_type)));
+  TF_ASSIGN_OR_RETURN(auto filter,
+                      scope(filter_descriptor, ToMIOpenDataType(input_type)));
+  TF_ASSIGN_OR_RETURN(auto conv, scope(convolution_descriptor));
 
   bool is_backprop = ((kind == dnn::ConvolutionKind::BACKWARD_DATA) ||
                       (kind == dnn::ConvolutionKind::BACKWARD_FILTER));
@@ -3754,7 +3774,7 @@ bool MIOpenSupport::GetMIOpenConvolveAlgorithmsFindMode(
   //                    (call_context == dnn::CallContext::kBackpropFilter);
 
 #if TF_ROCM_VERSION >= 50000
-  if (is_backprop && (ToMIOpenDataType(element_type) == miopenHalf)) {
+  if (is_backprop && (ToMIOpenDataType(input_type) == miopenHalf)) {
     wrap::miopenSetConvolutionAttribute(
         conv.handle(), MIOPEN_CONVOLUTION_ATTRIB_FP16_ALT_IMPL, 1);
   }
@@ -3763,15 +3783,15 @@ bool MIOpenSupport::GetMIOpenConvolveAlgorithmsFindMode(
   // Determine the workspace memory size that will need by the call to Find
   size_t scratch_memory_size = 0;
   switch (kind) {
-    case dnn::ConvolutionKind::FORWARD: {
+    case dnn::ConvolutionKind::FORWARD:
+    case dnn::ConvolutionKind::FORWARD_BIAS_ACTIVATION: {
       auto status = wrap::miopenConvolutionForwardGetWorkSpaceSize(
           miopen.handle(), filter.handle(), input_nd.handle(), conv.handle(),
           output_nd.handle(), &scratch_memory_size);
       if (status != miopenStatusSuccess) {
-        LOG(FATAL)
-            << "call to miopenConvolutionForwardGetWorkspaceSize failed: "
-            << ToString(status);
-        return false;
+        return absl::InternalError(
+            "call to miopenConvolutionForwardGetWorkspaceSize failed: " +
+            ToString(status));
       }
       break;
     }
@@ -3780,10 +3800,9 @@ bool MIOpenSupport::GetMIOpenConvolveAlgorithmsFindMode(
           miopen.handle(), output_nd.handle(), filter.handle(), conv.handle(),
           input_nd.handle(), &scratch_memory_size);
       if (status != miopenStatusSuccess) {
-        LOG(FATAL)
-            << "call to miopenConvolutionBackwardDataGetWorkspaceSize failed: "
-            << ToString(status);
-        return false;
+        return absl::InternalError(
+            "call to miopenConvolutionBackwardDataGetWorkspaceSize failed: " +
+            ToString(status));
       }
       break;
     }
@@ -3792,17 +3811,16 @@ bool MIOpenSupport::GetMIOpenConvolveAlgorithmsFindMode(
           miopen.handle(), output_nd.handle(), input_nd.handle(), conv.handle(),
           filter.handle(), &scratch_memory_size);
       if (status != miopenStatusSuccess) {
-        LOG(FATAL)
-            << "call to miopenConvolutionBackwardWeightsGetWorkspaceSize "
-               "failed: "
-            << ToString(status);
-        return false;
+        return absl::InternalError(
+            "call to miopenConvolutionBackwardWeightsGetWorkspaceSize "
+            "failed: " +
+            ToString(status));
       }
       break;
     }
     default: {
-      LOG(FATAL) << "Unexpected convolution kind " << static_cast<int>(kind);
-      return false;
+      return absl::InternalError("Unexpected convolution kind " +
+                                 static_cast<int>(kind));
       break;
     }
   }
@@ -3811,9 +3829,9 @@ bool MIOpenSupport::GetMIOpenConvolveAlgorithmsFindMode(
   DeviceMemory<uint8> scratch_memory;
   if (scratch_memory_size != 0) {
     if (scratch_allocator == nullptr) {
-      LOG(FATAL)
-          << "An allocator must be specified when scratch memory is needed";
-      return false;
+      return absl::InternalError(
+          "An allocator must be specified "
+          "when scratch memory is needed");
     }
     auto allocated = scratch_allocator->AllocateBytes(scratch_memory_size);
     if (allocated.ok()) {
@@ -3826,7 +3844,7 @@ bool MIOpenSupport::GetMIOpenConvolveAlgorithmsFindMode(
              "larger number (e.g. 8192) to increase the max memory limit.\n"
           << "\tIncreasing the max memory limit might help resolve this "
              "error";
-      return false;
+      return absl::InternalError("Out of memory");
     }
   }
 
@@ -3842,7 +3860,8 @@ bool MIOpenSupport::GetMIOpenConvolveAlgorithmsFindMode(
   bool exhaustiveSearch = false;
 
   switch (kind) {
-    case dnn::ConvolutionKind::FORWARD: {
+    case dnn::ConvolutionKind::FORWARD:
+    case dnn::ConvolutionKind::FORWARD_BIAS_ACTIVATION: {
       auto status = wrap::miopenFindConvolutionForwardAlgorithm(
           miopen.handle(), input_nd.handle(), input_data.opaque(),
           filter.handle(), filter_data.opaque(), conv.handle(),
@@ -3850,9 +3869,9 @@ bool MIOpenSupport::GetMIOpenConvolveAlgorithmsFindMode(
           &returnedAlgorithmCount, &returnedAlgorithm, scratch_memory.opaque(),
           scratch_memory_size, exhaustiveSearch);
       if (status != miopenStatusSuccess) {
-        LOG(FATAL) << "call to miopenFindConvolutionForwardAlgorithm failed: "
-                   << ToString(status);
-        return false;
+        return absl::InternalError(
+            "call to miopenFindConvolutionForwardAlgorithm failed: " +
+            ToString(status));
       }
       break;
     }
@@ -3864,10 +3883,9 @@ bool MIOpenSupport::GetMIOpenConvolveAlgorithmsFindMode(
           &returnedAlgorithmCount, &returnedAlgorithm, scratch_memory.opaque(),
           scratch_memory_size, exhaustiveSearch);
       if (status != miopenStatusSuccess) {
-        LOG(FATAL)
-            << "call to miopenFindConvolutionBackwardDataAlgorithm failed: "
-            << ToString(status);
-        return false;
+        return absl::InternalError(
+            "call to miopenFindConvolutionBackwardDataAlgorithm failed: " +
+            ToString(status));
       }
       break;
     }
@@ -3879,16 +3897,16 @@ bool MIOpenSupport::GetMIOpenConvolveAlgorithmsFindMode(
           &returnedAlgorithmCount, &returnedAlgorithm, scratch_memory.opaque(),
           scratch_memory_size, exhaustiveSearch);
       if (status != miopenStatusSuccess) {
-        LOG(FATAL) << "call to miopenConvolutionBackwardWeightsAlgorithm "
-                      "failed: "
-                   << ToString(status);
-        return false;
+        return absl::InternalError(
+            "call to miopenConvolutionBackwardWeightsAlgorithm "
+            "failed: " +
+            ToString(status));
       }
       break;
     }
     default: {
-      LOG(FATAL) << "Unexpected convolution kind " << static_cast<int>(kind);
-      return false;
+      return absl::InternalError("Unexpected convolution kind " +
+                                 std::to_string(static_cast<int>(kind)));
       break;
     }
   }
@@ -3896,7 +3914,7 @@ bool MIOpenSupport::GetMIOpenConvolveAlgorithmsFindMode(
   out_algorithms->emplace_back(
       GetProfileResultFromConvAlgoPerf(kind, returnedAlgorithm));
 
-  return true;
+  return absl::OkStatus();
 }
 
 bool MIOpenSupport::GetRnnAlgorithms(
@@ -3929,10 +3947,12 @@ bool MIOpenSupport::DoBatchNormalizationForward(
     bool is_training, ScratchAllocator* reserve_space_allocator,
     ScratchAllocator* workspace_allocator) {
   return DoBatchNormalizationForwardImpl<Eigen::bfloat16, float>(
-      stream, dnn::DataType::kBF16, dnn::DataType::kFloat, x, scale, offset,
-      estimated_mean, estimated_variance, side_input, x_desc, scale_offset_desc,
-      epsilon, exponential_average_factor, activation_mode, y, batch_mean,
-      batch_var, saved_mean, saved_inv_var, is_training);
+             stream, dnn::DataType::kBF16, dnn::DataType::kFloat, x, scale,
+             offset, estimated_mean, estimated_variance, side_input, x_desc,
+             scale_offset_desc, epsilon, exponential_average_factor,
+             activation_mode, y, batch_mean, batch_var, saved_mean,
+             saved_inv_var, is_training)
+      .ok();
 }
 
 bool MIOpenSupport::DoBatchNormalizationForward(
@@ -3950,10 +3970,12 @@ bool MIOpenSupport::DoBatchNormalizationForward(
     bool is_training, ScratchAllocator* reserve_space_allocator,
     ScratchAllocator* workspace_allocator) {
   return DoBatchNormalizationForwardImpl<Eigen::half, float>(
-      stream, dnn::DataType::kHalf, dnn::DataType::kFloat, x, scale, offset,
-      estimated_mean, estimated_variance, side_input, x_desc, scale_offset_desc,
-      epsilon, exponential_average_factor, activation_mode, y, batch_mean,
-      batch_var, saved_mean, saved_inv_var, is_training);
+             stream, dnn::DataType::kHalf, dnn::DataType::kFloat, x, scale,
+             offset, estimated_mean, estimated_variance, side_input, x_desc,
+             scale_offset_desc, epsilon, exponential_average_factor,
+             activation_mode, y, batch_mean, batch_var, saved_mean,
+             saved_inv_var, is_training)
+      .ok();
 }
 
 bool MIOpenSupport::DoBatchNormalizationForward(
@@ -3970,14 +3992,16 @@ bool MIOpenSupport::DoBatchNormalizationForward(
     bool is_training, ScratchAllocator* reserve_space_allocator,
     ScratchAllocator* workspace_allocator) {
   return DoBatchNormalizationForwardImpl<float, float>(
-      stream, dnn::DataType::kFloat, dnn::DataType::kFloat, x, scale, offset,
-      estimated_mean, estimated_variance, side_input, x_desc, scale_offset_desc,
-      epsilon, exponential_average_factor, activation_mode, y, batch_mean,
-      batch_var, saved_mean, saved_inv_var, is_training);
+             stream, dnn::DataType::kFloat, dnn::DataType::kFloat, x, scale,
+             offset, estimated_mean, estimated_variance, side_input, x_desc,
+             scale_offset_desc, epsilon, exponential_average_factor,
+             activation_mode, y, batch_mean, batch_var, saved_mean,
+             saved_inv_var, is_training)
+      .ok();
 }
 
 template <class T, class U>
-bool MIOpenSupport::DoBatchNormalizationForwardImpl(
+absl::Status MIOpenSupport::DoBatchNormalizationForwardImpl(
     Stream* stream, dnn::DataType input_data_type,
     dnn::DataType scale_data_type, const DeviceMemory<T>& x,
     const DeviceMemory<U>& scale, const DeviceMemory<U>& offset,
@@ -3992,10 +4016,11 @@ bool MIOpenSupport::DoBatchNormalizationForwardImpl(
     bool is_training) {
   auto miopen = miopen_->GetHandle(parent_, stream);
 
-  ScopedTensorDescriptor x_descriptor{x_desc,
-                                      ToMIOpenDataType(input_data_type)};
-  ScopedTensorDescriptor scale_offset_descriptor{
-      scale_offset_desc, ToMIOpenDataType(scale_data_type)};
+  TF_ASSIGN_OR_RETURN(auto x_descriptor,
+                      scope(x_desc, ToMIOpenDataType(input_data_type)));
+  TF_ASSIGN_OR_RETURN(
+      auto scale_offset_descriptor,
+      scope(scale_offset_desc, ToMIOpenDataType(scale_data_type)));
   miopenBatchNormMode_t mode = miopenBNSpatial;
   float one = 1.0;
   float zero = 0.0;
@@ -4018,11 +4043,11 @@ bool MIOpenSupport::DoBatchNormalizationForwardImpl(
         const_cast<void*>(maybe_inv_var), epsilon);
   }
   if (status != miopenStatusSuccess) {
-    LOG(ERROR) << "failed to enqueue forward batch normalization on stream: "
-               << ToString(status);
-    return false;
+    return absl::InternalError(
+        "failed to enqueue forward batch normalization on stream: " +
+        ToString(status));
   }
-  return true;
+  return absl::OkStatus();
 }
 
 bool MIOpenSupport::DoBatchNormalizationBackward(
@@ -4039,9 +4064,10 @@ bool MIOpenSupport::DoBatchNormalizationBackward(
     DeviceMemory<uint8_t>* reserve_space_data,
     ScratchAllocator* workspace_allocator) {
   return DoBatchNormalizationBackwardImpl<Eigen::bfloat16, float>(
-      stream, miopenBFloat16, miopenFloat, y_backprop, x, scale, mean, inv_var,
-      x_desc, scale_offset_desc, epsilon, x_backprop, scale_backprop,
-      offset_backprop);
+             stream, miopenBFloat16, miopenFloat, y_backprop, x, scale, mean,
+             inv_var, x_desc, scale_offset_desc, epsilon, x_backprop,
+             scale_backprop, offset_backprop)
+      .ok();
 }
 
 bool MIOpenSupport::DoBatchNormalizationBackward(
@@ -4057,9 +4083,10 @@ bool MIOpenSupport::DoBatchNormalizationBackward(
     DeviceMemory<uint8>* reserve_space_data,
     ScratchAllocator* workspace_allocator) {
   return DoBatchNormalizationBackwardImpl<Eigen::half, float>(
-      stream, miopenHalf, miopenFloat, y_backprop, x, scale, mean, inv_var,
-      x_desc, scale_offset_desc, epsilon, x_backprop, scale_backprop,
-      offset_backprop);
+             stream, miopenHalf, miopenFloat, y_backprop, x, scale, mean,
+             inv_var, x_desc, scale_offset_desc, epsilon, x_backprop,
+             scale_backprop, offset_backprop)
+      .ok();
 }
 
 bool MIOpenSupport::DoBatchNormalizationBackward(
@@ -4075,13 +4102,14 @@ bool MIOpenSupport::DoBatchNormalizationBackward(
     DeviceMemory<uint8>* reserve_space_data,
     ScratchAllocator* workspace_allocator) {
   return DoBatchNormalizationBackwardImpl<float, float>(
-      stream, miopenFloat, miopenFloat, y_backprop, x, scale, mean, variance,
-      x_desc, scale_offset_desc, epsilon, x_backprop, scale_backprop,
-      offset_backprop);
+             stream, miopenFloat, miopenFloat, y_backprop, x, scale, mean,
+             variance, x_desc, scale_offset_desc, epsilon, x_backprop,
+             scale_backprop, offset_backprop)
+      .ok();
 }
 
 template <class T, class U>
-bool MIOpenSupport::DoBatchNormalizationBackwardImpl(
+absl::Status MIOpenSupport::DoBatchNormalizationBackwardImpl(
     Stream* stream, int miopen_input_type, int miopen_scale_type,
     const DeviceMemory<T>& y_backprop, const DeviceMemory<T>& x,
     const DeviceMemory<U>& scale, const DeviceMemory<U>& mean,
@@ -4090,10 +4118,12 @@ bool MIOpenSupport::DoBatchNormalizationBackwardImpl(
     DeviceMemory<T>* x_backprop, DeviceMemory<U>* scale_backprop,
     DeviceMemory<U>* offset_backprop) {
   auto miopen = miopen_->GetHandle(parent_, stream);
-  ScopedTensorDescriptor x_descriptor{
-      x_desc, static_cast<miopenDataType_t>(miopen_input_type)};
-  ScopedTensorDescriptor scale_offset_descriptor{
-      scale_offset_desc, static_cast<miopenDataType_t>(miopen_scale_type)};
+  TF_ASSIGN_OR_RETURN(
+      auto x_descriptor,
+      scope(x_desc, static_cast<miopenDataType_t>(miopen_input_type)));
+  TF_ASSIGN_OR_RETURN(auto scale_offset_descriptor,
+                      scope(scale_offset_desc,
+                            static_cast<miopenDataType_t>(miopen_scale_type)));
   miopenBatchNormMode_t mode = miopenBNSpatial;
   float one = 1.0;
   float zero = 0.0;
@@ -4106,18 +4136,20 @@ bool MIOpenSupport::DoBatchNormalizationBackwardImpl(
       scale_backprop->opaque(), offset_backprop->opaque(), epsilon,
       mean.opaque(), variance.opaque());
   if (status != miopenStatusSuccess) {
-    LOG(ERROR) << "failed to enqueue backward batch normalization on stream: "
-               << ToString(status);
-    return false;
+    return absl::InternalError(
+        "failed to enqueue backward batch normalization on stream: " +
+        ToString(status));
   }
-  return true;
+  return absl::OkStatus();
 }
 
-template <typename T>
+template <typename T, typename Tbias>
 void launchInplaceBiasActivation(hipStream_t stream, void* c_data,
-                                 const void* bias_data, int activation_mode,
-                                 uint64_t m, uint64_t n, int64_t ldc,
-                                 float param);
+                                 const void* bias_data,
+                                 const void* side_input_data,
+                                 float side_input_scale, int activation_mode,
+                                 uint64_t batch, uint64_t m, uint64_t n,
+                                 int64_t ldc, float param);
 
 class ROCmFusedMatmulRunner : public dnn::FusedMatmulRunner {
   template <typename T>
@@ -4212,31 +4244,35 @@ absl::Status ROCmFusedMatmulRunner::gemm(Stream* stream,
                               NumericOptions{}, blas::CallContext::kNone);
 }
 
-template <typename T>
-absl::Status InplaceBiasActivation(Stream* stream, DeviceMemoryBase c_data,
-                                   DeviceMemoryBase bias_data,
-                                   dnn::ActivationMode activation_mode,
-                                   uint64_t m, uint64_t n, int64_t ldc,
-                                   float param) {
+template <typename T, typename Tbias = T>
+absl::Status InplaceBiasActivation(
+    Stream* stream, DeviceMemoryBase c_data, DeviceMemoryBase bias_data,
+    DeviceMemoryBase side_input_data, float side_input_scale,
+    dnn::ActivationMode activation_mode, uint64_t batch, uint64_t m, uint64_t n,
+    int64_t ldc, float param, bool transpose = false) {
   typedef typename std::conditional<
       std::is_same_v<T, Eigen::half>, __half,
       typename std::conditional<std::is_same_v<T, Eigen::bfloat16>,
                                 hip_bfloat16, T>::type>::type CT;
-
-  if (activation_mode == dnn::ActivationMode::kReluX ||
-      activation_mode == dnn::ActivationMode::kBandPass ||
-      activation_mode == dnn::ActivationMode::kLeakyRelu)
-
-    return absl::InvalidArgumentError(
-        "ROCm InplaceBiasActivation can't be used with "
-        "parametric activations yet");
-
-  launchInplaceBiasActivation<CT>(
+  typedef typename std::conditional<
+      std::is_same_v<Tbias, Eigen::half>, __half,
+      typename std::conditional<std::is_same_v<Tbias, Eigen::bfloat16>,
+                                hip_bfloat16, Tbias>::type>::type CTbias;
+  launchInplaceBiasActivation<CT, CTbias>(
       AsGpuStreamValue(stream), c_data.opaque(), bias_data.opaque(),
-      static_cast<int>(activation_mode), m, n, ldc, param);
+      side_input_data.opaque(), side_input_scale,
+      static_cast<int>(activation_mode) + (transpose ? 10 : 0), batch, m, n,
+      ldc, param);
   return absl::OkStatus();
 }
 
+template <typename Ta, typename Tb, typename... Args>
+absl::Status InplaceBiasActivation(Stream* stream, DeviceMemory<Ta> c_data,
+                                   DeviceMemory<Tb> bias_data, Args... args) {
+  return InplaceBiasActivation<Ta, Tb>(stream, DeviceMemoryBase(c_data),
+                                       DeviceMemoryBase(bias_data), args...);
+}
+
 // Launch the operation, with the signature determined by `Sig`.
 absl::Status ROCmFusedMatmulRunner::operator()(
     Stream* stream, dnn::ProfileResult* prof, DeviceMemoryBase scratch_memory,
@@ -4255,18 +4291,24 @@ absl::Status ROCmFusedMatmulRunner::operator()(
     return absl::InvalidArgumentError("Unsupported input type");
 
   if (!status.ok()) return status;
+
+  DeviceMemory<uint8_t> side_input;
   if (_input_type == dnn::DataType::kFloat)
-    return InplaceBiasActivation<float>(stream, c_data, bias_data,
-                                        _activation_mode, _m, _n, _ldc, 0.0f);
+    return InplaceBiasActivation<float>(stream, c_data, bias_data, side_input,
+                                        0.0f, _activation_mode, 1, _m, _n, _ldc,
+                                        0.0f);
   else if (_input_type == dnn::DataType::kHalf)
     return InplaceBiasActivation<Eigen::half>(
-        stream, c_data, bias_data, _activation_mode, _m, _n, _ldc, 0.0f);
+        stream, c_data, bias_data, side_input, 0.0f, _activation_mode, 1, _m,
+        _n, _ldc, 0.0f);
   else if (_input_type == dnn::DataType::kBF16)
     return InplaceBiasActivation<Eigen::bfloat16>(
-        stream, c_data, bias_data, _activation_mode, _m, _n, _ldc, 0.0f);
+        stream, c_data, bias_data, side_input, 0.0f, _activation_mode, 1, _m,
+        _n, _ldc, 0.0f);
   else if (_input_type == dnn::DataType::kDouble)
-    return InplaceBiasActivation<double>(stream, c_data, bias_data,
-                                         _activation_mode, _m, _n, _ldc, 0.0f);
+    return InplaceBiasActivation<double>(stream, c_data, bias_data, side_input,
+                                         0.0f, _activation_mode, 1, _m, _n,
+                                         _ldc, 0.0f);
   else
     return absl::InvalidArgumentError("Unsupported input type");
 }
@@ -4344,9 +4386,9 @@ absl::Status MIOpenSupport::DoPoolForward(
   auto miopen_dtype =
       element_type == dnn::DataType::kFloat ? miopenFloat : miopenHalf;
 
-  ScopedTensorDescriptor src_desc{input_dimensions, miopen_dtype};
-  ScopedTensorDescriptor dest_desc{output_dimensions, miopen_dtype};
-  ScopedPoolingDescriptor pooling_desc{pooling_dimensions};
+  TF_ASSIGN_OR_RETURN(auto src_desc, scope(input_dimensions, miopen_dtype));
+  TF_ASSIGN_OR_RETURN(auto dest_desc, scope(output_dimensions, miopen_dtype));
+  TF_ASSIGN_OR_RETURN(auto pooling_desc, scope(pooling_dimensions));
 
   bool do_backward = false;
   uint8* workspace = nullptr;
@@ -4498,9 +4540,9 @@ absl::Status MIOpenSupport::DoPoolBackward(
   auto miopen_dtype =
       element_type == dnn::DataType::kFloat ? miopenFloat : miopenHalf;
 
-  ScopedTensorDescriptor src_desc{input_dimensions, miopen_dtype};
-  ScopedTensorDescriptor dest_desc{output_dimensions, miopen_dtype};
-  ScopedPoolingDescriptor pooling_desc{pooling_dimensions};
+  TF_ASSIGN_OR_RETURN(auto src_desc, scope(input_dimensions, miopen_dtype));
+  TF_ASSIGN_OR_RETURN(auto dest_desc, scope(output_dimensions, miopen_dtype));
+  TF_ASSIGN_OR_RETURN(auto pooling_desc, scope(pooling_dimensions));
 
   uint8* workspace_ptr = 0;
   DeviceMemory<uint8> workspace;
@@ -4588,6 +4630,17 @@ absl::Status MIOpenSupport::DoPoolBackward(
   return absl::OkStatus();
 }
 
+#define ASSIGN_OR_RETURN_FALSE(lhs, rexpr) \
+  ASSIGN_OR_RETURN_FALSE_IMPL(             \
+      TF_STATUS_MACROS_CONCAT_NAME(_status_or_value, __COUNTER__), lhs, rexpr)
+
+#define ASSIGN_OR_RETURN_FALSE_IMPL(statusor, lhs, rexpr) \
+  auto statusor = (rexpr);                                \
+  if (TF_PREDICT_FALSE(!statusor.ok())) {                 \
+    return false;                                         \
+  }                                                       \
+  lhs = std::move(statusor).value()
+
 bool MIOpenSupport::DoNormalizeWithDimensions(
     Stream* stream, const dnn::NormalizeDescriptor& normalize_descriptor,
     const dnn::BatchDescriptor& dimensions,
@@ -4605,8 +4658,8 @@ bool MIOpenSupport::DoNormalizeWithDimensions(
   auto miopen = miopen_->GetHandle(parent_, stream);
 
   // Launch the normalization.
-  ScopedTensorDescriptor dims{dimensions, miopenFloat};
-  ScopedNormalizeDescriptor normalize{normalize_descriptor};
+  ASSIGN_OR_RETURN_FALSE(auto dims, scope(dimensions, miopenFloat));
+  ASSIGN_OR_RETURN_FALSE(auto normalize, scope(normalize_descriptor));
 
   // Alpha is the scaling factor for input.
   float alpha = 1.0f;
@@ -4643,8 +4696,8 @@ bool MIOpenSupport::DoNormalizeBackwardWithDimensions(
 
   auto miopen = miopen_->GetHandle(parent_, stream);
 
-  ScopedTensorDescriptor dims{dimensions, miopenFloat};
-  ScopedNormalizeDescriptor normalize{normalize_descriptor};
+  ASSIGN_OR_RETURN_FALSE(auto dims, scope(dimensions, miopenFloat));
+  ASSIGN_OR_RETURN_FALSE(auto normalize, scope(normalize_descriptor));
 
   float alpha = 1.0f;
   float beta = 0.0f;
@@ -4728,9 +4781,9 @@ bool MIOpenSupport::DeriveOutputBatchDescriptor(
     const FilterDescriptor& filter_descriptor,
     const dnn::ConvolutionDescriptor& convolution_descriptor,
     dnn::BatchDescriptor* output_batch_descriptor) {
-  ScopedTensorDescriptor input_nd{batch_descriptor, miopenFloat};
-  ScopedFilterDescriptor filter{filter_descriptor, miopenFloat};
-  ScopedConvolutionDescriptor conv{convolution_descriptor, miopenFloat};
+  ASSIGN_OR_RETURN_FALSE(auto input_nd, scope(batch_descriptor, miopenFloat));
+  ASSIGN_OR_RETURN_FALSE(auto filter, scope(filter_descriptor, miopenFloat));
+  ASSIGN_OR_RETURN_FALSE(auto conv, scope(convolution_descriptor));
 
   int dn = batch_descriptor.ndims() + 2;
   std::vector<int> dims(dn);  // in BDYX
@@ -4754,6 +4807,369 @@ bool MIOpenSupport::DeriveOutputBatchDescriptor(
   return true;
 }
 
+class RocmFusedConvRunner : public dnn::FusedConvRunner {
+ public:
+  std::string ToString() const override {
+    return MakeAlgorithmDesc().ToString();
+  }
+
+  uint64_t GetWorkspaceSize() const override { return workspace_size_; }
+
+  absl::StatusOr<dnn::AlgorithmDesc> ToAlgorithmDesc() const override {
+    return MakeAlgorithmDesc();
+  }
+
+  absl::Status operator()(Stream* stream, dnn::ProfileResult* profile_result,
+                          DeviceMemoryBase scratch_memory,
+                          DeviceMemoryBase input_data,
+                          DeviceMemoryBase filter_data,
+                          DeviceMemoryBase side_input_data,
+                          DeviceMemoryBase bias_data,
+                          DeviceMemoryBase output_data) const override {
+    VLOG(2) << "RocmFusedConvRunner()";
+    if (parent_ != stream->parent()) {
+      return absl::InternalError(
+          "RocmFusedConvRunner cached across multiple StreamExecutors.");
+    }
+
+    // We can't reliably detect whether this sequence can be fused until
+    // we come here and actually try to fuse it. So, we need a fallback.
+    bool do_unfused =
+        (side_input_scale_ != 0.0) || !fusion_plan_.CompilationSucceeded();
+
+    if (do_unfused)
+      return execute_unfused(stream, profile_result, scratch_memory, input_data,
+                             filter_data, side_input_data, bias_data,
+                             output_data);
+    auto algo = MakeAlgorithmDesc();
+    auto miopen = miopen_->GetHandle(parent_, stream);
+    fusion_plan_.SetConvolutionArgs(filter_data.opaque());
+    fusion_plan_.SetBiasArgs(bias_data.opaque());
+    if (activation_desc_.miopen_activation_mode_ != miopenActivationPASTHRU)
+      fusion_plan_.SetActivationForwardArgs(activation_desc_);
+
+    std::optional<GpuTimer> timer;
+    if (profile_result) {
+      auto timer_or_status = GpuTimer::Create(AsGpuStream(stream));
+      if (!timer_or_status.ok()) {
+        LOG(ERROR) << "Failed to create timer";
+        return absl::InternalError("Failed to start timer");
+      }
+      timer.emplace(std::move(*timer_or_status));
+    }
+
+    miopenStatus_t status;
+    status = wrap::miopenExecuteFusionPlan(
+        miopen.handle(), fusion_plan_.fusion_plan_, input_nd_.handle(),
+        input_data.opaque(), output_nd_.handle(), output_data.opaque(),
+        fusion_plan_.fusion_args_);
+
+    if (status != miopenStatusSuccess) {
+      LOG(ERROR) << "Failed to enqueue fused convolution on stream: "
+                 << stream_executor::gpu::ToString(status);
+      return absl::InternalError(
+          "Failed to enqueue fused convolution on stream: " +
+          stream_executor::gpu::ToString(status));
+    }
+
+    if (profile_result) {
+      absl::StatusOr<absl::Duration> elapsed = timer->GetElapsedDuration();
+      if (!elapsed.ok()) {
+        LOG(ERROR) << "Failed to get elapsed duration";
+        return absl::InternalError("Timer failure");
+      }
+      profile_result->set_elapsed_time_in_ms(
+          absl::ToDoubleMilliseconds(*elapsed));
+      profile_result->set_algorithm(algo);
+      profile_result->set_scratch_size(scratch_memory.size());
+    }
+
+    return absl::OkStatus();
+  }
+
+ public:
+  // Queries the workspace size and constructs a 'RocmFusedConvRunner'.
+  static absl::StatusOr<std::unique_ptr<const dnn::FusedConvRunner>> Create(
+      GpuExecutor* parent, Stream* stream, MIOpenAccess* miopen,
+      const dnn::AlgorithmDesc& algo, dnn::DataType input_type,
+      dnn::DataType bias_type, double conv_scale, double side_input_scale,
+      double leakyrelu_alpha, BatchDescriptor input_nd,
+      BatchDescriptor output_nd, FilterDescriptor filter,
+      BatchDescriptor bias_nd, ConvolutionDescriptor conv,
+      dnn::ActivationMode activation) {
+    TF_ASSIGN_OR_RETURN(
+        auto input_nd_,
+        scope(input_nd, ToMIOpenDataType(input_type, input_nd.layout())));
+    TF_ASSIGN_OR_RETURN(
+        auto output_nd_,
+        scope(output_nd, ToMIOpenDataType(input_type, input_nd.layout())));
+    TF_ASSIGN_OR_RETURN(auto filter_,
+                        scope(filter, ToMIOpenDataType(input_type)));
+    TF_ASSIGN_OR_RETURN(auto bias_nd_,
+                        scope(bias_nd, ToMIOpenDataType(bias_type)));
+    TF_ASSIGN_OR_RETURN(auto conv_, scope(conv));
+
+    TF_ASSIGN_OR_RETURN(
+        auto activation_desc,
+        ScopedActivationDescriptor::Create(activation, leakyrelu_alpha));
+
+    TF_ASSIGN_OR_RETURN(
+        auto fusion_plan,
+        ScopedFusionPlanConvolutionBiasActivation::Create(
+            miopen->GetHandle(parent, stream).handle(), input_nd_.handle(),
+            filter_.handle(), conv_.handle(), bias_nd_.handle(),
+            activation_desc));
+
+    VLOG(2) << "RocmFusedConvRunner";
+    auto mi = miopen->GetHandle(parent, stream);
+
+    size_t maxSolutionCount = 0;
+    auto status = wrap::miopenConvolutionForwardGetSolutionCount(
+        mi.handle(), filter_.handle(), input_nd_.handle(), conv_.handle(),
+        output_nd_.handle(), &maxSolutionCount);
+
+    size_t solutionCount = 0;
+    std::unique_ptr<miopenConvSolution_t[]> solutions(
+        new miopenConvSolution_t[maxSolutionCount]);
+
+    status = wrap::miopenConvolutionForwardGetSolution(
+        mi.handle(), filter_.handle(), input_nd_.handle(), conv_.handle(),
+        output_nd_.handle(), maxSolutionCount, &solutionCount, solutions.get());
+
+    VLOG(2) << solutionCount << " solutions";
+
+    if (solutionCount == 0) return absl::InternalError("No algorithms found");
+
+    size_t workspace_size_1 = solutions[0].workspace_size;
+    size_t true_workspace_size = 0;
+    status = wrap::miopenConvolutionForwardGetWorkSpaceSize(
+        mi.handle(), filter_.handle(), input_nd_.handle(), conv_.handle(),
+        output_nd_.handle(), &true_workspace_size);
+
+    VLOG(2) << "True workspace size " << workspace_size_1 << " "
+            << true_workspace_size;
+
+    auto obj = new RocmFusedConvRunner(
+        parent, stream, miopen, static_cast<int64_t>(solutions[0].solution_id),
+        true_workspace_size, input_type, bias_type, conv_scale,
+        side_input_scale, leakyrelu_alpha, input_nd, output_nd, filter, bias_nd,
+        conv, activation, input_nd_, output_nd_, filter_, bias_nd_, conv_,
+        activation_desc, fusion_plan);
+
+    return std::unique_ptr<const dnn::FusedConvRunner>(obj);
+  }
+
+ private:
+  // Private to prevent passing in the wrong workspace_size.
+  RocmFusedConvRunner(
+      GpuExecutor* parent, Stream* stream, MIOpenAccess* miopen,
+      int64_t algo_id, size_t workspace_size, dnn::DataType input_type,
+      dnn::DataType bias_type, double conv_scale, double side_input_scale,
+      double leakyrelu_alpha, BatchDescriptor dnn_input_nd,
+      BatchDescriptor dnn_output_nd, FilterDescriptor dnn_filter,
+      BatchDescriptor dnn_bias_nd, ConvolutionDescriptor dnn_conv,
+      dnn::ActivationMode activation, ScopedTensorDescriptor& input_nd,
+      ScopedTensorDescriptor& output_nd, ScopedFilterDescriptor& filter,
+      ScopedTensorDescriptor& bias_nd, ScopedConvolutionDescriptor& conv,
+      ScopedActivationDescriptor& activation_desc,
+      ScopedFusionPlanConvolutionBiasActivation& fusion_plan)
+      : parent_(parent),
+        miopen_(miopen),
+        algo_id_(algo_id),
+        workspace_size_(workspace_size),
+        input_type_(input_type),
+        bias_type_(bias_type),
+
+        conv_scale_(conv_scale),
+        side_input_scale_(side_input_scale),
+        leakyrelu_alpha_(leakyrelu_alpha),
+        side_input_scale_f32_(static_cast<float>(side_input_scale)),
+
+        dnn_input_nd_(dnn_input_nd),
+        dnn_output_nd_(dnn_output_nd),
+        dnn_filter_(dnn_filter),
+        dnn_bias_nd_(dnn_bias_nd),
+        dnn_conv_(dnn_conv),
+
+        activation_mode_(activation),
+
+        input_nd_(std::move(input_nd)),
+        output_nd_(std::move(output_nd)),
+        filter_(std::move(filter)),
+        bias_nd_(std::move(bias_nd)),
+        conv_(std::move(conv)),
+        activation_desc_(std::move(activation_desc)),
+        fusion_plan_(std::move(fusion_plan)) {}
+
+  absl::Status execute_unfused(
+      Stream* stream, dnn::ProfileResult* profile_result,
+      DeviceMemoryBase scratch_memory, DeviceMemoryBase input_data,
+      DeviceMemoryBase filter_data, DeviceMemoryBase side_input_data,
+      DeviceMemoryBase bias_data, DeviceMemoryBase output_data) const {
+    auto miopen = miopen_->GetHandle(parent_, stream);
+    auto status = wrap::miopenConvolutionForwardImmediate(
+        miopen.handle(), filter_.handle(), filter_data.opaque(),
+        input_nd_.handle(), input_data.opaque(), conv_.handle(),
+        output_nd_.handle(), output_data.opaque(), scratch_memory.opaque(),
+        scratch_memory.size(), static_cast<uint64_t>(algo_id_));
+    if (status != miopenStatusSuccess) {
+      VLOG(0) << "Failed to enqueue convolution: "
+              << stream_executor::gpu::ToString(status);
+      return absl::InternalError("Failed to enqueue convolution: " +
+                                 stream_executor::gpu::ToString(status));
+    }
+
+    int batch;
+    std::vector<int64_t> dims_output =
+        dnn_output_nd_.full_dims(dnn_output_nd_.layout());
+    int rank = dims_output.size();
+    if (rank != 4 && rank != 5)
+      return absl::InternalError(
+          "RocmFusedConvRunner expects 4d or 5d descriptors");
+    int d1 = 1, d2 = 1;
+    bool bNCHW = (dnn_output_nd_.layout() != dnn::DataLayout::kBatchYXDepth);
+    batch = dims_output[0];
+    int w, h = 1;
+    if (bNCHW) {
+      d1 = dims_output[1];
+      for (int i = 2; i < rank; i++) d2 *= dims_output[i];
+    } else {
+      d2 = dims_output[rank - 1];
+      for (int i = 1; i < rank - 1; i++) d1 *= dims_output[i];
+    }
+
+    float param = activation_desc_.alpha_;
+
+    auto inplace_call = [&](auto out, auto bias) {
+      return InplaceBiasActivation(stream, out, bias, side_input_data,
+                                   side_input_scale_f32_, activation_mode_,
+                                   batch, d1, d2, d2, param, bNCHW);
+    };
+
+    absl::Status biasActStatus;
+    if (input_type_ == dnn::DataType::kFloat &&
+        bias_type_ == dnn::DataType::kFloat)
+      biasActStatus = inplace_call(DeviceMemory<float>(output_data),
+                                   DeviceMemory<float>(bias_data));
+    else if (input_type_ == dnn::DataType::kHalf &&
+             bias_type_ == dnn::DataType::kFloat)
+      biasActStatus = inplace_call(DeviceMemory<Eigen::half>(output_data),
+                                   DeviceMemory<float>(bias_data));
+    else if (input_type_ == dnn::DataType::kHalf &&
+             bias_type_ == dnn::DataType::kHalf)
+      biasActStatus = inplace_call(DeviceMemory<Eigen::half>(output_data),
+                                   DeviceMemory<Eigen::half>(bias_data));
+    else if (input_type_ == dnn::DataType::kBF16 &&
+             bias_type_ == dnn::DataType::kFloat)
+      biasActStatus = inplace_call(DeviceMemory<Eigen::bfloat16>(output_data),
+                                   DeviceMemory<float>(bias_data));
+    else if (input_type_ == dnn::DataType::kBF16 &&
+             bias_type_ == dnn::DataType::kBF16)
+      biasActStatus = inplace_call(DeviceMemory<Eigen::bfloat16>(output_data),
+                                   DeviceMemory<Eigen::bfloat16>(bias_data));
+    else
+      return absl::InternalError("Unsupported data type");
+
+    return absl::OkStatus();
+  }
+
+  // Internal form of ToAlgorithmDesc without the StatusOr.
+  dnn::AlgorithmDesc MakeAlgorithmDesc() const {
+    return {algo_id_, /*tensor_ops_enabled_*/ true, workspace_size_};
+  }
+
+  std::string desc_;
+
+  GpuExecutor* parent_;
+  MIOpenAccess* miopen_;
+  int64_t algo_id_;
+  size_t workspace_size_;
+  dnn::DataType input_type_, bias_type_;
+  double conv_scale_, side_input_scale_, leakyrelu_alpha_;
+  float side_input_scale_f32_;
+  dnn::ActivationMode activation_mode_;
+
+  BatchDescriptor dnn_input_nd_;
+  BatchDescriptor dnn_output_nd_;
+  FilterDescriptor dnn_filter_;
+  BatchDescriptor dnn_bias_nd_;
+  ConvolutionDescriptor dnn_conv_;
+
+  ScopedTensorDescriptor input_nd_;
+  ScopedTensorDescriptor output_nd_;
+  ScopedFilterDescriptor filter_;
+  ScopedTensorDescriptor bias_nd_;
+  ScopedConvolutionDescriptor conv_;
+  mutable ScopedActivationDescriptor activation_desc_;
+  mutable ScopedFusionPlanConvolutionBiasActivation fusion_plan_;
+};
+
+absl::StatusOr<std::unique_ptr<const dnn::FusedConvRunner>>
+MIOpenSupport::FusedConvolveRunnerFromDesc(
+    Stream* stream, const dnn::AlgorithmDesc& algorithm_desc,
+    dnn::ConvolutionKind kind, dnn::DataType input_type,
+    dnn::DataType bias_type, dnn::DataType output_type, double conv_scale,
+    double side_input_scale, double leakyrelu_alpha,
+    const dnn::BatchDescriptor& input_descriptor,
+    const dnn::FilterDescriptor& filter_descriptor,
+    const dnn::BatchDescriptor& bias_descriptor,
+    const dnn::BatchDescriptor& output_descriptor,
+    const dnn::ConvolutionDescriptor& convolution_descriptor,
+    dnn::ActivationMode activation_mode) {
+  VLOG(2) << "MIOpenSupport::FusedConvolveRunnerFromDesc "
+          << filter_descriptor.ndims() << " " << side_input_scale << " "
+          << convolution_descriptor.ToString() << getTypeName(input_type) << " "
+          << getTypeName(bias_type) << " " << getTypeName(output_type);
+
+  // note: these checks need to be duplicated in XLA logic, because XLA calls
+  // this function directly and it terminates the process on error
+
+  return RocmFusedConvRunner::Create(
+      parent_, stream, miopen_.get(), algorithm_desc, input_type, bias_type,
+      conv_scale, side_input_scale, leakyrelu_alpha, input_descriptor,
+      output_descriptor, filter_descriptor, bias_descriptor,
+      convolution_descriptor, activation_mode);
+}
+
+absl::Status MIOpenSupport::GetFusedConvolveRunners(
+    bool use_cudnn_frontend, dnn::ConvolutionKind kind,
+    dnn::DataType input_type, dnn::DataType bias_type,
+    dnn::DataType output_type, double conv_scale, double side_input_scale,
+    double leakyrelu_alpha, Stream* stream,
+    const dnn::BatchDescriptor& input_descriptor,
+    const dnn::FilterDescriptor& filter_descriptor,
+    const dnn::BatchDescriptor& bias_descriptor,
+    const dnn::BatchDescriptor& output_descriptor,
+    const dnn::ConvolutionDescriptor& convolution_descriptor, bool use_fallback,
+    dnn::ActivationMode activation_mode, const NumericOptions& numeric_options,
+    std::vector<std::unique_ptr<const dnn::FusedConvRunner>>* out_exec_plans) {
+  VLOG(2) << "MIOpenSupport::GetFusedConvolveRunners";
+  VLOG(2) << "filter_descriptor " << filter_descriptor.ndims();
+
+  std::vector<dnn::AlgorithmDesc> algorithms{
+      // clang-format off
+      dnn::AlgorithmDesc(miopenConvolutionFwdAlgoGEMM, false, 0),
+      dnn::AlgorithmDesc(miopenConvolutionFwdAlgoDirect, false, 0),
+      dnn::AlgorithmDesc(miopenConvolutionFwdAlgoFFT, false, 0),
+      dnn::AlgorithmDesc(miopenConvolutionFwdAlgoWinograd, false, 0),
+      // clang-format on
+  };
+
+  for (const auto& algo : algorithms) {
+    auto runner_or = FusedConvolveRunnerFromDesc(
+        stream, algo, kind, input_type, bias_type, output_type, conv_scale,
+        side_input_scale, leakyrelu_alpha, input_descriptor, filter_descriptor,
+        bias_descriptor, output_descriptor, convolution_descriptor,
+        activation_mode);
+    if (!runner_or.ok()) continue;
+    out_exec_plans->push_back(std::move(runner_or).value());
+  }
+
+  VLOG(2) << "MIOpenSupport::GetFusedConvolveRunners returns "
+          << out_exec_plans->size() << " runners";
+  return absl::OkStatus();
+}
+
 bool UseNhwcLayoutForRocm() {
 #if TF_ROCM_VERSION >= 50100
   static bool is_enabled = [] {
diff --git a/third_party/xla/xla/stream_executor/rocm/rocm_dnn.h b/third_party/xla/xla/stream_executor/rocm/rocm_dnn.h
index ecaffd3cad392d..957c288f12c278 100644
--- a/third_party/xla/xla/stream_executor/rocm/rocm_dnn.h
+++ b/third_party/xla/xla/stream_executor/rocm/rocm_dnn.h
@@ -251,8 +251,34 @@ class MIOpenSupport : public dnn::DnnSupport {
       const dnn::BatchDescriptor& output_descriptor,
       const dnn::ConvolutionDescriptor& convolution_descriptor) override;
 
+  absl::StatusOr<std::unique_ptr<const dnn::FusedConvRunner>>
+  FusedConvolveRunnerFromDesc(
+      Stream* stream, const dnn::AlgorithmDesc& algorithm_desc,
+      dnn::ConvolutionKind kind, dnn::DataType input_type,
+      dnn::DataType bias_type, dnn::DataType output_type, double conv_scale,
+      double side_input_scale, double leakyrelu_alpha,
+      const dnn::BatchDescriptor& input_descriptor,
+      const dnn::FilterDescriptor& filter_descriptor,
+      const dnn::BatchDescriptor& bias_descriptor,
+      const dnn::BatchDescriptor& output_descriptor,
+      const dnn::ConvolutionDescriptor& convolution_descriptor,
+      dnn::ActivationMode activation_mode) override;
+
   bool GetMIOpenConvolveAlgorithms(
-      dnn::ConvolutionKind kind, dnn::DataType element_type, Stream* stream,
+      dnn::ConvolutionKind kind, dnn::DataType input_type,
+      dnn::DataType output_type, Stream* stream,
+      const dnn::BatchDescriptor& input_descriptor, DeviceMemoryBase input_data,
+      const dnn::FilterDescriptor& filter_descriptor,
+      DeviceMemoryBase filter_data,
+      const dnn::BatchDescriptor& output_descriptor,
+      DeviceMemoryBase output_data,
+      const dnn::ConvolutionDescriptor& convolution_descriptor,
+      ScratchAllocator* scratch_allocator,
+      std::vector<dnn::ProfileResult>* out_algorithms);
+
+  absl::Status GetMIOpenConvolveAlgorithmsImmediateMode(
+      dnn::ConvolutionKind kind, dnn::DataType input_type,
+      dnn::DataType output_type, Stream* stream,
       const dnn::BatchDescriptor& input_descriptor, DeviceMemoryBase input_data,
       const dnn::FilterDescriptor& filter_descriptor,
       DeviceMemoryBase filter_data,
@@ -260,7 +286,19 @@ class MIOpenSupport : public dnn::DnnSupport {
       DeviceMemoryBase output_data,
       const dnn::ConvolutionDescriptor& convolution_descriptor,
       ScratchAllocator* scratch_allocator,
-      std::vector<dnn::ProfileResult>* out_algorithms) override;
+      std::vector<dnn::ProfileResult>* out_algorithms);
+
+  absl::Status GetMIOpenConvolveAlgorithmsFindMode(
+      dnn::ConvolutionKind kind, dnn::DataType input_type,
+      dnn::DataType output_type, Stream* stream,
+      const dnn::BatchDescriptor& input_descriptor, DeviceMemoryBase input_data,
+      const dnn::FilterDescriptor& filter_descriptor,
+      DeviceMemoryBase filter_data,
+      const dnn::BatchDescriptor& output_descriptor,
+      DeviceMemoryBase output_data,
+      const dnn::ConvolutionDescriptor& convolution_descriptor,
+      ScratchAllocator* scratch_allocator,
+      std::vector<dnn::ProfileResult>* out_algorithms);
 
   bool GetRnnAlgorithms(
       std::vector<dnn::AlgorithmDesc>* out_algorithms) override;
@@ -389,6 +427,21 @@ class MIOpenSupport : public dnn::DnnSupport {
       std::vector<std::unique_ptr<const dnn::FusedMatmulRunner>>*
           out_exec_plans) override;
 
+  absl::Status GetFusedConvolveRunners(
+      bool use_cudnn_frontend, dnn::ConvolutionKind kind,
+      dnn::DataType input_type, dnn::DataType bias_type,
+      dnn::DataType output_type, double conv_scale, double side_input_scale,
+      double leakyrelu_alpha, Stream* stream,
+      const dnn::BatchDescriptor& input_descriptor,
+      const dnn::FilterDescriptor& filter_descriptor,
+      const dnn::BatchDescriptor& bias_descriptor,
+      const dnn::BatchDescriptor& output_descriptor,
+      const dnn::ConvolutionDescriptor& convolution_descriptor,
+      bool use_fallback, dnn::ActivationMode activation_mode,
+      const NumericOptions& numeric_options,
+      std::vector<std::unique_ptr<const dnn::FusedConvRunner>>* out_exec_plans)
+      override;
+
   absl::Status DoPoolForward(dnn::DataType element_type, Stream* stream,
                              const dnn::PoolingDescriptor& pooling_dimensions,
                              const dnn::BatchDescriptor& input_dimensions,
@@ -469,7 +522,7 @@ class MIOpenSupport : public dnn::DnnSupport {
   bool m_pooling_cache_enabled = false;
 
   template <class T, class U>
-  bool DoBatchNormalizationForwardImpl(
+  absl::Status DoBatchNormalizationForwardImpl(
       Stream* stream, dnn::DataType input_data_type,
       dnn::DataType scale_data_type, const DeviceMemory<T>& x,
       const DeviceMemory<U>& scale, const DeviceMemory<U>& offset,
@@ -484,7 +537,7 @@ class MIOpenSupport : public dnn::DnnSupport {
       bool is_training);
 
   template <class T, class U>
-  bool DoBatchNormalizationBackwardImpl(
+  absl::Status DoBatchNormalizationBackwardImpl(
       Stream* stream, int miopen_input_type, int miopen_scale_type,
       const DeviceMemory<T>& y_backprop, const DeviceMemory<T>& x,
       const DeviceMemory<U>& scale, const DeviceMemory<U>& mean,
@@ -569,28 +622,6 @@ class MIOpenSupport : public dnn::DnnSupport {
       ScratchAllocator* scratch_allocator, DeviceMemory<uint8>* scratch_memory,
       int* ctc_loss_algo_id) override;
 
-  bool GetMIOpenConvolveAlgorithmsImmediateMode(
-      dnn::ConvolutionKind kind, dnn::DataType element_type, Stream* stream,
-      const dnn::BatchDescriptor& input_descriptor, DeviceMemoryBase input_data,
-      const dnn::FilterDescriptor& filter_descriptor,
-      DeviceMemoryBase filter_data,
-      const dnn::BatchDescriptor& output_descriptor,
-      DeviceMemoryBase output_data,
-      const dnn::ConvolutionDescriptor& convolution_descriptor,
-      ScratchAllocator* scratch_allocator,
-      std::vector<dnn::ProfileResult>* out_algorithms);
-
-  bool GetMIOpenConvolveAlgorithmsFindMode(
-      dnn::ConvolutionKind kind, dnn::DataType element_type, Stream* stream,
-      const dnn::BatchDescriptor& input_descriptor, DeviceMemoryBase input_data,
-      const dnn::FilterDescriptor& filter_descriptor,
-      DeviceMemoryBase filter_data,
-      const dnn::BatchDescriptor& output_descriptor,
-      DeviceMemoryBase output_data,
-      const dnn::ConvolutionDescriptor& convolution_descriptor,
-      ScratchAllocator* scratch_allocator,
-      std::vector<dnn::ProfileResult>* out_algorithms);
-
   MIOpenSupport(const MIOpenSupport&) = delete;
   void operator=(const MIOpenSupport&) = delete;
 };
diff --git a/third_party/xla/xla/stream_executor/rocm/rocm_helpers.cu.cc b/third_party/xla/xla/stream_executor/rocm/rocm_helpers.cu.cc
index cf12ffdaeb47ca..ae287aa4025fb2 100644
--- a/third_party/xla/xla/stream_executor/rocm/rocm_helpers.cu.cc
+++ b/third_party/xla/xla/stream_executor/rocm/rocm_helpers.cu.cc
@@ -83,16 +83,23 @@ __device__ float sigmoid(float x) {
     return __expf(x) / (__expf(x) + 1.);
 }
 
-template <typename T, int act_mode>
-__global__ void launchInplaceBiasActivation_kernel(T* c_data,
-                                                   const T* bias_data,
-                                                   uint64_t m, uint64_t n,
-                                                   int64_t ldc, float param) {
+template <typename T, typename Tbias, int act_mode>
+__global__ void launchInplaceBiasActivation_kernel(
+    T* c_data, const Tbias* bias_data, const T* side_input_data,
+    float side_input_scale, uint64_t m, uint64_t n, int64_t ldc, float param,
+    int transpose) {
   uint64_t x = threadIdx.x + blockIdx.x * blockDim.x;
   uint64_t y = threadIdx.y + blockIdx.y * blockDim.y;
+  uint64_t z = blockIdx.z;
   if (x >= n || y >= m) return;
-  float v = static_cast<float>(c_data[x + y * ldc]) +
-            static_cast<float>(bias_data[x]);
+  float v;
+  uint64_t addr = x + y * ldc + z * m * n;
+  if (!transpose)
+    v = static_cast<float>(c_data[addr]) + static_cast<float>(bias_data[x]);
+  else
+    v = static_cast<float>(c_data[addr]) + static_cast<float>(bias_data[y]);
+  if (side_input_data != 0)
+    v += static_cast<float>(side_input_data[addr]) * side_input_scale;
   if (act_mode == 1)
     v = sigmoid(v);
   else if (act_mode == 2)
@@ -111,58 +118,62 @@ __global__ void launchInplaceBiasActivation_kernel(T* c_data,
     v = v > 0.0f ? v : param * v;
   else if (act_mode == 9)
     v = 0.5 * v * (1 + erf(v / sqrt(2.0f)));
-  c_data[x + y * ldc] = (T)v;
+  c_data[addr] = (T)v;
 }
 
-template <typename T>
+template <typename T, typename Tbias>
 void launchInplaceBiasActivation(hipStream_t stream, void* c_data,
-                                 const void* bias_data, int activation_mode,
-                                 uint64_t m, uint64_t n, int64_t ldc,
-                                 float param) {
+                                 const void* bias_data,
+                                 const void* side_input_data,
+                                 float side_input_scale, int activation_mode,
+                                 uint64_t batch, uint64_t m, uint64_t n,
+                                 int64_t ldc, float param) {
   uint64_t bx = min(n, static_cast<uint64_t>(256));
   uint64_t by = min(m, static_cast<uint64_t>(256) / bx);
   uint64_t gx = (n + bx - 1) / bx;
   uint64_t gy = (m + by - 1) / by;
-  auto kernel = launchInplaceBiasActivation_kernel<T, 0>;
+  int transpose = (activation_mode >= 10);
+  activation_mode %= 10;
+  auto kernel = launchInplaceBiasActivation_kernel<T, Tbias, 0>;
   if (activation_mode == 1)
-    kernel = launchInplaceBiasActivation_kernel<T, 1>;
+    kernel = launchInplaceBiasActivation_kernel<T, Tbias, 1>;
   else if (activation_mode == 2)
-    kernel = launchInplaceBiasActivation_kernel<T, 2>;
+    kernel = launchInplaceBiasActivation_kernel<T, Tbias, 2>;
   else if (activation_mode == 3)
-    kernel = launchInplaceBiasActivation_kernel<T, 3>;
+    kernel = launchInplaceBiasActivation_kernel<T, Tbias, 3>;
   else if (activation_mode == 4)
-    kernel = launchInplaceBiasActivation_kernel<T, 4>;
+    kernel = launchInplaceBiasActivation_kernel<T, Tbias, 4>;
   else if (activation_mode == 5)
-    kernel = launchInplaceBiasActivation_kernel<T, 5>;
+    kernel = launchInplaceBiasActivation_kernel<T, Tbias, 5>;
   else if (activation_mode == 6)
-    kernel = launchInplaceBiasActivation_kernel<T, 6>;
+    kernel = launchInplaceBiasActivation_kernel<T, Tbias, 6>;
   else if (activation_mode == 7)
-    kernel = launchInplaceBiasActivation_kernel<T, 7>;
+    kernel = launchInplaceBiasActivation_kernel<T, Tbias, 7>;
   else if (activation_mode == 8)
-    kernel = launchInplaceBiasActivation_kernel<T, 8>;
+    kernel = launchInplaceBiasActivation_kernel<T, Tbias, 8>;
   else if (activation_mode == 9)
-    kernel = launchInplaceBiasActivation_kernel<T, 9>;
+    kernel = launchInplaceBiasActivation_kernel<T, Tbias, 9>;
 
-  hipLaunchKernelGGL(kernel, dim3(gx, gy, 1), dim3(bx, by, 1), 0, stream,
-                     static_cast<T*>(c_data), static_cast<const T*>(bias_data),
-                     m, n, ldc, param);
+  hipLaunchKernelGGL(kernel, dim3(gx, gy, batch), dim3(bx, by, 1), 0, stream,
+                     static_cast<T*>(c_data),
+                     static_cast<const Tbias*>(bias_data),
+                     static_cast<const T*>(side_input_data), side_input_scale,
+                     m, n, ldc, param, transpose);
 }
 
-template void launchInplaceBiasActivation<__half>(
-    hipStream_t stream, void* c_data, const void* bias_data,
-    int activation_mode, uint64_t m, uint64_t n, int64_t ldc, float param);
-
-template void launchInplaceBiasActivation<hip_bfloat16>(
-    hipStream_t stream, void* c_data, const void* bias_data,
-    int activation_mode, uint64_t m, uint64_t n, int64_t ldc, float param);
-
-template void launchInplaceBiasActivation<float>(
-    hipStream_t stream, void* c_data, const void* bias_data,
-    int activation_mode, uint64_t m, uint64_t n, int64_t ldc, float param);
-
-template void launchInplaceBiasActivation<double>(
-    hipStream_t stream, void* c_data, const void* bias_data,
-    int activation_mode, uint64_t m, uint64_t n, int64_t ldc, float param);
+#define INSTANTIATE_BIAS_ACTIVATION(X, Y)                          \
+  template void launchInplaceBiasActivation<X, Y>(                 \
+      hipStream_t stream, void* c_data, const void* bias_data,     \
+      const void* side_input_data, float side_input_scale,         \
+      int activation_mode, uint64_t batch, uint64_t m, uint64_t n, \
+      int64_t ldc, float param);
+
+INSTANTIATE_BIAS_ACTIVATION(__half, __half)
+INSTANTIATE_BIAS_ACTIVATION(__half, float)
+INSTANTIATE_BIAS_ACTIVATION(hip_bfloat16, hip_bfloat16)
+INSTANTIATE_BIAS_ACTIVATION(hip_bfloat16, float)
+INSTANTIATE_BIAS_ACTIVATION(float, float)
+INSTANTIATE_BIAS_ACTIVATION(double, double)
 
 };  // namespace gpu
 };  // namespace stream_executor
diff --git a/third_party/xla/xla/tests/convolution_test.cc b/third_party/xla/xla/tests/convolution_test.cc
index c4f3ec42c1c648..3dd0f1db48f9c4 100644
--- a/third_party/xla/xla/tests/convolution_test.cc
+++ b/third_party/xla/xla/tests/convolution_test.cc
@@ -18,6 +18,7 @@ limitations under the License.
 
 #include <memory>
 
+#include "absl/strings/str_replace.h"
 #include "xla/array2d.h"
 #include "xla/array4d.h"
 #include "xla/client/global_data.h"
@@ -1755,7 +1756,7 @@ ENTRY TestComputation {
 }
 
 XLA_TEST_F(ConvolutionHloTest, TestFusedConv2D) {
-  constexpr char kHlo[] = R"(
+  std::string kHlo = R"(
 HloModule TestModule
 
 ENTRY TestComputation {
@@ -1765,11 +1766,51 @@ ENTRY TestComputation {
   %bias = f32[32] parameter(2)
   %broadcasted_bias = f32[8,5,5,32] broadcast(%bias), dimensions={3}
   %add = f32[8,5,5,32] add(%conv, %broadcasted_bias)
+)";
+
+  std::string kHloNoPad = R"(
+HloModule TestModule
+
+ENTRY TestComputation {
+  %p0 = f32[8,7,7,1] parameter(0)
+  %p1 = f32[3,3,1,32] parameter(1)
+  %conv = f32[8,5,5,32] convolution(p0, p1), window={size=3x3 pad=0_0x0_0}, dim_labels=b01f_01io->b01f
+  %bias = f32[32] parameter(2)
+  %broadcasted_bias = f32[8,5,5,32] broadcast(%bias), dimensions={3}
+  %add = f32[8,5,5,32] add(%conv, %broadcasted_bias)
+)";
+
+  std::string kHloRELU = R"(
+
   %zero = f32[] constant(0)
   %zeros = f32[8,5,5,32] broadcast(%zero), dimensions={}
   ROOT relu = f32[8,5,5,32] maximum(%zeros, %add)
 })";
-  EXPECT_TRUE(RunAndCompare(kHlo, ErrorSpec{0.01, 0.01}));
+
+  std::string kHloTANH = R"(
+  ROOT result = f32[8,5,5,32] tanh(%add)
+})";
+
+  std::string kHloELU = R"(
+  %zero = f32[] constant(0)
+  %zeros = f32[8,5,5,32] broadcast(%zero), dimensions={}
+  %one = f32[] constant(1)
+  %ones = f32[8,5,5,32] broadcast(%one), dimensions={}
+  %exp = f32[8,5,5,32] exponential(%add)
+  %expm1 = f32[8,5,5,32] subtract(%exp, %ones)
+  %sgn = pred[8,5,5,32] compare(%add, %zeros), direction=GT
+  ROOT elu = f32[8,5,5,32] select(%sgn, %add, %expm1)
+})";
+
+  EXPECT_TRUE(RunAndCompare(kHlo + kHloRELU, ErrorSpec{0.01, 0.01}));
+  EXPECT_TRUE(RunAndCompare(kHlo + kHloTANH, ErrorSpec{0.01, 0.01}));
+  EXPECT_TRUE(RunAndCompare(kHlo + kHloELU, ErrorSpec{0.01, 0.01}));
+  EXPECT_TRUE(
+      RunAndCompare(absl::StrReplaceAll(kHlo + kHloRELU, {{"f32", "f16"}}),
+                    ErrorSpec{0.03, 0.03}));
+  EXPECT_TRUE(
+      RunAndCompare(absl::StrReplaceAll(kHloNoPad + kHloRELU, {{"f32", "f16"}}),
+                    ErrorSpec{0.03, 0.03}));
 }
 
 XLA_TEST_F(ConvolutionHloTest, TestFusedConv3D) {

From 213a0ac1c84a5a479a91a640f21efff93ac3a966 Mon Sep 17 00:00:00 2001
From: Georg Stefan Schmid <gschmid@nvidia.com>
Date: Wed, 15 May 2024 06:51:16 -0700
Subject: [PATCH 178/478] PR #10803: Expose CompiledMemoryStats on CPU and GPU

Imported from GitHub PR https://github.com/openxla/xla/pull/10803

[CompiledMemoryStats](https://github.com/openxla/xla/blob/01073cf10b57c64615483545bfb34fb9ae0d3c06/xla/pjrt/pjrt_executable.h#L269-L269) is supposed to summarize the amount of memory required to run a given executable. Buffer-related fields in CompiledMemoryStats are currently missing for both the GPU and the CPU backend. This PR computes the missing fields using the executable's underlying [BufferAssignment](https://github.com/openxla/xla/blob/50c6489cb058881cc65622605c9c55029abebc5b/xla/service/buffer_assignment.h#L491) / buffer allocations.

# Example
```python
import jax
import jax.numpy as jnp

def f(a, b):
    ms = jnp.sum(a, axis=1, keepdims=True) / a.shape[1]
    return a * jax.lax.rsqrt(ms) + b

a0 = jax.random.uniform(jax.random.PRNGKey(0), (128, 128))
b0 = jax.random.uniform(jax.random.PRNGKey(0), (128, 128))
f_jit = jax.jit(f, donate_argnums=(0,))

print(f_jit.lower(a0, b0).compile().memory_analysis())
# CompiledMemoryStats(
#     generated_code_size_in_bytes=5792,
#     argument_size_in_bytes=131072,
#     output_size_in_bytes=65536,
#     alias_size_in_bytes=65536,
#     temp_size_in_bytes=512,
#     host_generated_code_size_in_bytes=0,
#     host_argument_size_in_bytes=0,
#     host_output_size_in_bytes=0,
#     host_alias_size_in_bytes=0,
#     host_temp_size_in_bytes=0)
```
Note the use of `donate_argnums=(0,)` which allows XLA to modify the first input buffer in-place and reuse it as the output. This is reflected in `alias_size_in_bytes` -- without buffer donation that number is zero.

# Open questions
- Does this match the semantics for the TPU backend?
- `dynamic_cast`-ing on `gpu::Executable` ain't great. Do we want `Executable`s to expose their buffer allocations publicly (i.e. GpuExecutable's [GetAllocations](https://github.com/openxla/xla/blob/f6ce6c11ceac904c9c4840fcec5f1ee715291425/xla/service/gpu/gpu_executable.h#L164) but for all Executables)?
Copybara import of the project:

--
f9ae2c4e367aed813af74ae50edfd27a6f9304eb by Georg Stefan Schmid <gschmid@nvidia.com>:

Expose CompiledMemoryStats on CPU and GPU

Merging this change closes #10803

PiperOrigin-RevId: 633937384
---
 third_party/xla/xla/pjrt/BUILD                |   2 +
 third_party/xla/xla/pjrt/cpu/cpu_client.h     |   2 +
 third_party/xla/xla/pjrt/pjrt_executable.cc   | 111 ++++++++++++++++++
 third_party/xla/xla/pjrt/pjrt_executable.h    |  37 +-----
 .../xla/pjrt/pjrt_stream_executor_client.h    |   2 +
 third_party/xla/xla/service/BUILD             |   1 +
 .../xla/xla/service/buffer_assignment.h       |   4 +
 .../xla/xla/service/cpu/cpu_executable.h      |   4 +
 third_party/xla/xla/service/executable.h      |   7 ++
 .../xla/xla/service/gpu/gpu_executable.h      |   2 +-
 10 files changed, 139 insertions(+), 33 deletions(-)

diff --git a/third_party/xla/xla/pjrt/BUILD b/third_party/xla/xla/pjrt/BUILD
index 662c77584df5e7..346fcb233ddd12 100644
--- a/third_party/xla/xla/pjrt/BUILD
+++ b/third_party/xla/xla/pjrt/BUILD
@@ -254,10 +254,12 @@ cc_library(
         "//xla/client:executable_build_options",
         "//xla/ffi:execution_context",
         "//xla/hlo/ir:hlo",
+        "//xla/service:buffer_assignment",
         "//xla/service:compiler",
         "//xla/service:computation_layout",
         "//xla/service:hlo_cost_analysis",
         "//xla/service:hlo_proto_cc",
+        "//xla/service:hlo_value",
         "@com_google_absl//absl/container:flat_hash_map",
         "@com_google_absl//absl/container:flat_hash_set",
         "@com_google_absl//absl/log",
diff --git a/third_party/xla/xla/pjrt/cpu/cpu_client.h b/third_party/xla/xla/pjrt/cpu/cpu_client.h
index 554d7bb9c3b0d6..c704f2b38a2f81 100644
--- a/third_party/xla/xla/pjrt/cpu/cpu_client.h
+++ b/third_party/xla/xla/pjrt/cpu/cpu_client.h
@@ -603,6 +603,8 @@ class TfrtCpuExecutable final : public PjRtLoadedExecutable {
           "cpu_executable_ has no hlo_proto.");
     }
     memory_stats.serialized_hlo_proto = proto->SerializeAsString();
+    memory_stats.PopulateBufferStatsFromAllocations(
+        cpu_executable_.get()->GetAllocations());
     return memory_stats;
   }
 
diff --git a/third_party/xla/xla/pjrt/pjrt_executable.cc b/third_party/xla/xla/pjrt/pjrt_executable.cc
index 4c4eae38f0cbc9..4caee29a4a862c 100644
--- a/third_party/xla/xla/pjrt/pjrt_executable.cc
+++ b/third_party/xla/xla/pjrt/pjrt_executable.cc
@@ -29,14 +29,17 @@ limitations under the License.
 #include "absl/strings/numbers.h"
 #include "absl/strings/str_cat.h"
 #include "absl/strings/string_view.h"
+#include "absl/types/span.h"
 #include "xla/client/executable_build_options.h"
 #include "xla/layout.h"
 #include "xla/pjrt/compile_options.pb.h"
 #include "xla/pjrt/execute_options.pb.h"
 #include "xla/pjrt/pjrt_common.h"
 #include "xla/pjrt/pjrt_layout.h"
+#include "xla/service/buffer_assignment.h"
 #include "xla/service/computation_layout.h"
 #include "xla/service/hlo_cost_analysis.h"
+#include "xla/service/hlo_value.h"
 #include "xla/shape.h"
 #include "xla/shape_layout.h"
 #include "xla/shape_util.h"
@@ -216,6 +219,114 @@ absl::StatusOr<ExecuteOptions> ExecuteOptions::FromProto(
   return options;
 }
 
+CompiledMemoryStatsProto CompiledMemoryStats::ToProto() {
+  CompiledMemoryStatsProto proto;
+  proto.set_generated_code_size_in_bytes(generated_code_size_in_bytes);
+  proto.set_argument_size_in_bytes(argument_size_in_bytes);
+  proto.set_output_size_in_bytes(output_size_in_bytes);
+  proto.set_alias_size_in_bytes(alias_size_in_bytes);
+  proto.set_temp_size_in_bytes(temp_size_in_bytes);
+  proto.mutable_hlo_proto()->ParseFromString(serialized_hlo_proto);
+  proto.set_host_generated_code_size_in_bytes(
+      host_generated_code_size_in_bytes);
+  proto.set_host_argument_size_in_bytes(host_argument_size_in_bytes);
+  proto.set_host_output_size_in_bytes(host_output_size_in_bytes);
+  proto.set_host_alias_size_in_bytes(host_alias_size_in_bytes);
+  proto.set_host_temp_size_in_bytes(host_temp_size_in_bytes);
+  return proto;
+}
+
+CompiledMemoryStats CompiledMemoryStats::FromProto(
+    const CompiledMemoryStatsProto& proto) {
+  CompiledMemoryStats stats;
+  stats.generated_code_size_in_bytes = proto.generated_code_size_in_bytes();
+  stats.argument_size_in_bytes = proto.argument_size_in_bytes();
+  stats.output_size_in_bytes = proto.output_size_in_bytes();
+  stats.alias_size_in_bytes = proto.alias_size_in_bytes();
+  stats.temp_size_in_bytes = proto.temp_size_in_bytes();
+  stats.serialized_hlo_proto = proto.hlo_proto().SerializeAsString();
+  stats.host_generated_code_size_in_bytes =
+      proto.host_generated_code_size_in_bytes();
+  stats.host_argument_size_in_bytes = proto.host_argument_size_in_bytes();
+  stats.host_output_size_in_bytes = proto.host_output_size_in_bytes();
+  stats.host_alias_size_in_bytes = proto.host_alias_size_in_bytes();
+  stats.host_temp_size_in_bytes = proto.host_temp_size_in_bytes();
+  return stats;
+}
+
+// Recomputes the memory stats from allocations. Why recompute?
+// Firstly, there are cases in which gpu::Executable inherits its allocations
+// from elsewhere, and no buffer assignment is available.
+// Secondly, exec->buffer_assignment()->GetStats() provides the statistics we
+// want, but does not distinguish between device and host memory, and does
+// not account for aliased memory.
+void CompiledMemoryStats::PopulateBufferStatsFromAllocations(
+    absl::Span<const BufferAllocation> allocs) {
+  argument_size_in_bytes = 0;
+  output_size_in_bytes = 0;
+  temp_size_in_bytes = 0;
+  alias_size_in_bytes = 0;
+  host_argument_size_in_bytes = 0;
+  host_output_size_in_bytes = 0;
+  host_temp_size_in_bytes = 0;
+  host_alias_size_in_bytes = 0;
+
+  for (auto& alloc : allocs) {
+    // All logical buffers assigned to a buffer allocation share a color.
+    // With buffer assigner's default colorer the color happens to be the
+    // memory space of the underlying HLO value. Callers may choose other
+    // colorers, however, e.g.:
+    // https://github.com/openxla/xla/blob/50c6489cb058881cc65622605c9c55029abebc5b/xla/service/gpu/compile_module_to_llvm_ir.cc#L152
+    // Until buffer allocations provide a stronger guarantee about colors,
+    // we sanity-check that the default coloring behavior was used.
+    int64_t alloc_memory_space = -1;
+    for (const auto& [value, _] : alloc.assigned_buffers()) {
+      const HloPosition& defining_position = value->defining_position();
+      int64_t memory_space = Layout::kDefaultMemorySpace;
+      if (defining_position.shape().has_layout()) {
+        memory_space = defining_position.shape().layout().memory_space();
+      }
+      if (alloc_memory_space == -1) {
+        alloc_memory_space = memory_space;
+      } else {
+        CHECK(alloc_memory_space == memory_space &&
+              "expected same memory space for all assignments in allocation");
+      }
+    }
+
+    bool is_host = alloc_memory_space == Layout::kHostMemorySpace;
+    int64_t size = alloc.size();
+    if (alloc.is_entry_computation_parameter()) {
+      if (is_host) {
+        host_argument_size_in_bytes += size;
+      } else {
+        argument_size_in_bytes += size;
+      }
+      if (alloc.is_parameter_aliased_with_output()) {
+        if (is_host) {
+          host_alias_size_in_bytes += size;
+        } else {
+          alias_size_in_bytes += size;
+        }
+      }
+    }
+    if (alloc.maybe_live_out()) {
+      if (is_host) {
+        host_output_size_in_bytes += size;
+      } else {
+        output_size_in_bytes += size;
+      }
+    }
+    if (alloc.IsPreallocatedTempBuffer()) {
+      if (is_host) {
+        host_temp_size_in_bytes += size;
+      } else {
+        temp_size_in_bytes += size;
+      }
+    }
+  }
+}
+
 void GetOpSharding(std::vector<OpSharding>& out, const OpSharding& sharding) {
   if (sharding.type() == OpSharding::TUPLE) {
     for (const OpSharding& s : sharding.tuple_shardings()) {
diff --git a/third_party/xla/xla/pjrt/pjrt_executable.h b/third_party/xla/xla/pjrt/pjrt_executable.h
index 176d1c11342007..e4005da2eee3b7 100644
--- a/third_party/xla/xla/pjrt/pjrt_executable.h
+++ b/third_party/xla/xla/pjrt/pjrt_executable.h
@@ -298,39 +298,12 @@ struct CompiledMemoryStats {
   std::string serialized_hlo_proto = "";
   std::string DebugString() const;
 
-  CompiledMemoryStatsProto ToProto() {
-    CompiledMemoryStatsProto proto;
-    proto.set_generated_code_size_in_bytes(generated_code_size_in_bytes);
-    proto.set_argument_size_in_bytes(argument_size_in_bytes);
-    proto.set_output_size_in_bytes(output_size_in_bytes);
-    proto.set_alias_size_in_bytes(alias_size_in_bytes);
-    proto.set_temp_size_in_bytes(temp_size_in_bytes);
-    proto.mutable_hlo_proto()->ParseFromString(serialized_hlo_proto);
-    proto.set_host_generated_code_size_in_bytes(
-        host_generated_code_size_in_bytes);
-    proto.set_host_argument_size_in_bytes(host_argument_size_in_bytes);
-    proto.set_host_output_size_in_bytes(host_output_size_in_bytes);
-    proto.set_host_alias_size_in_bytes(host_alias_size_in_bytes);
-    proto.set_host_temp_size_in_bytes(host_temp_size_in_bytes);
-    return proto;
-  }
+  CompiledMemoryStatsProto ToProto();
 
-  static CompiledMemoryStats FromProto(const CompiledMemoryStatsProto& proto) {
-    CompiledMemoryStats stats;
-    stats.generated_code_size_in_bytes = proto.generated_code_size_in_bytes();
-    stats.argument_size_in_bytes = proto.argument_size_in_bytes();
-    stats.output_size_in_bytes = proto.output_size_in_bytes();
-    stats.alias_size_in_bytes = proto.alias_size_in_bytes();
-    stats.temp_size_in_bytes = proto.temp_size_in_bytes();
-    stats.serialized_hlo_proto = proto.hlo_proto().SerializeAsString();
-    stats.host_generated_code_size_in_bytes =
-        proto.host_generated_code_size_in_bytes();
-    stats.host_argument_size_in_bytes = proto.host_argument_size_in_bytes();
-    stats.host_output_size_in_bytes = proto.host_output_size_in_bytes();
-    stats.host_alias_size_in_bytes = proto.host_alias_size_in_bytes();
-    stats.host_temp_size_in_bytes = proto.host_temp_size_in_bytes();
-    return stats;
-  }
+  static CompiledMemoryStats FromProto(const CompiledMemoryStatsProto& proto);
+
+  void PopulateBufferStatsFromAllocations(
+      absl::Span<const BufferAllocation> allocs);
 };
 
 class PjRtExecutable {
diff --git a/third_party/xla/xla/pjrt/pjrt_stream_executor_client.h b/third_party/xla/xla/pjrt/pjrt_stream_executor_client.h
index 913e38ef4c5bf4..9d95c5441b1f45 100644
--- a/third_party/xla/xla/pjrt/pjrt_stream_executor_client.h
+++ b/third_party/xla/xla/pjrt/pjrt_stream_executor_client.h
@@ -918,6 +918,8 @@ class PjRtStreamExecutorLoadedExecutable : public PjRtLoadedExecutable {
     if (proto != nullptr) {
       memory_stats.serialized_hlo_proto = proto->SerializeAsString();
     }
+    memory_stats.PopulateBufferStatsFromAllocations(
+        executables_[0]->executable()->GetAllocations());
     return memory_stats;
   }
 
diff --git a/third_party/xla/xla/service/BUILD b/third_party/xla/xla/service/BUILD
index 558257f0baf5cf..2cf56434539c13 100644
--- a/third_party/xla/xla/service/BUILD
+++ b/third_party/xla/xla/service/BUILD
@@ -1467,6 +1467,7 @@ cc_library(
         "service_executable_run_options.h",
     ],
     deps = [
+        ":buffer_assignment",
         ":computation_layout",
         ":dump",
         ":hlo_execution_profile",
diff --git a/third_party/xla/xla/service/buffer_assignment.h b/third_party/xla/xla/service/buffer_assignment.h
index e908b617641f05..e27ed3cb8f0bec 100644
--- a/third_party/xla/xla/service/buffer_assignment.h
+++ b/third_party/xla/xla/service/buffer_assignment.h
@@ -114,6 +114,10 @@ class BufferAllocation {
     return is_entry_computation_parameter_;
   }
 
+  bool is_parameter_aliased_with_output() const {
+    return is_parameter_aliased_with_output_;
+  }
+
   // Whether this allocation holds a constant.  On the CPU and GPU backends
   // constant allocations are not allocated dynamically, instead we resolve
   // references to these buffer allocations to a global in the readonly section
diff --git a/third_party/xla/xla/service/cpu/cpu_executable.h b/third_party/xla/xla/service/cpu/cpu_executable.h
index 34b4aa56ddf243..2dc7e2d97009e3 100644
--- a/third_party/xla/xla/service/cpu/cpu_executable.h
+++ b/third_party/xla/xla/service/cpu/cpu_executable.h
@@ -109,6 +109,10 @@ class CpuExecutable : public Executable {
 
   int64_t SizeOfGeneratedCodeInBytes() const override;
 
+  absl::Span<const BufferAllocation> GetAllocations() const override {
+    return assignment_->Allocations();
+  }
+
  private:
   // Creates an array suitable for passing as the "buffer_table" argument to the
   // JIT compiled function pointer.
diff --git a/third_party/xla/xla/service/executable.h b/third_party/xla/xla/service/executable.h
index 573ff79a564c3d..5c26581b9479e3 100644
--- a/third_party/xla/xla/service/executable.h
+++ b/third_party/xla/xla/service/executable.h
@@ -26,6 +26,7 @@ limitations under the License.
 #include "absl/types/variant.h"
 #include "xla/debug_options_flags.h"
 #include "xla/hlo/ir/hlo_module.h"
+#include "xla/service/buffer_assignment.h"
 #include "xla/service/computation_layout.h"
 #include "xla/service/hlo.pb.h"
 #include "xla/service/hlo_execution_profile.h"
@@ -406,6 +407,12 @@ class Executable {
   void MarkToBeReleasedArguments(absl::Span<ExecutionInput> arguments,
                                  ExecutionOutput& result);
 
+  // Returns the allocations resulting from buffer assignment, or an empty span
+  // if unimplemented.
+  virtual absl::Span<const BufferAllocation> GetAllocations() const {
+    return {};
+  }
+
  protected:
   // HloModule this was compiled from. BufferAssignment keeps pointers to
   // HloInstructions owned by the HloModule so we need to keep the HloModule
diff --git a/third_party/xla/xla/service/gpu/gpu_executable.h b/third_party/xla/xla/service/gpu/gpu_executable.h
index 7b5bfd7dcc2a8c..c76eacc809e4b7 100644
--- a/third_party/xla/xla/service/gpu/gpu_executable.h
+++ b/third_party/xla/xla/service/gpu/gpu_executable.h
@@ -159,7 +159,7 @@ class GpuExecutable : public Executable {
       const ServiceExecutableRunOptions* run_options,
       VariantArguments arguments);
 
-  absl::Span<const BufferAllocation> GetAllocations() const {
+  absl::Span<const BufferAllocation> GetAllocations() const override {
     // A GpuExecutable can get its allocations in three ways:
     // 1 - From a regular compilation that uses allocations from MLIR.
     // 2 - From a regular compilation that uses the original allocations from

From 453abe7da63cbd4ba673691a90a7a8d561f78672 Mon Sep 17 00:00:00 2001
From: Peter Hawkins <phawkins@google.com>
Date: Wed, 15 May 2024 07:10:41 -0700
Subject: [PATCH 179/478] [PJRT:CPU] Provide a compile-time thread pool.

The OneDNN matmul rewriter starts a thread pool if one wasn't provided as a compile option, spawning many threads per compilation. If we pass a preexisting thread pool, we can save the cost of creating and tearing down threads.

PiperOrigin-RevId: 633942672
---
 third_party/xla/xla/pjrt/cpu/cpu_client.cc | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/third_party/xla/xla/pjrt/cpu/cpu_client.cc b/third_party/xla/xla/pjrt/cpu/cpu_client.cc
index 3746dd4fd1c3d5..1b7b5fe12b7c88 100644
--- a/third_party/xla/xla/pjrt/cpu/cpu_client.cc
+++ b/third_party/xla/xla/pjrt/cpu/cpu_client.cc
@@ -826,6 +826,9 @@ absl::StatusOr<std::unique_ptr<PjRtLoadedExecutable>> TfrtCpuClient::Compile(
   xla::Compiler::CompileOptions compile_options{
       build_options.device_allocator(), build_options.compile_thread_pool(),
       build_options.layout_canonicalization_callback()};
+  if (!compile_options.thread_pool) {
+    compile_options.thread_pool = pjrt_client_thread_pool();
+  }
   TF_ASSIGN_OR_RETURN(
       std::unique_ptr<Executable> cpu_executable,
       JitCompile(computation, argument_layout_pointers, build_options,

From 2aef976b3e99eddc5dc0cd3c91026dc2b4832857 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Wed, 15 May 2024 07:20:39 -0700
Subject: [PATCH 180/478] Update ops-related pbtxt files.

PiperOrigin-RevId: 633945444
---
 tensorflow/core/ops/ops.pbtxt | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorflow/core/ops/ops.pbtxt b/tensorflow/core/ops/ops.pbtxt
index 0751aea36d13e7..cdc0191b8ad6c9 100644
--- a/tensorflow/core/ops/ops.pbtxt
+++ b/tensorflow/core/ops/ops.pbtxt
@@ -1,4 +1,4 @@
-go/debugonly   
+go/debugproto   
 op {
   name: "Abort"
   attr {

From 89b5367da720557e33e6da38703b7400c7223fa3 Mon Sep 17 00:00:00 2001
From: Goran Flegar <gflegar@google.com>
Date: Wed, 15 May 2024 07:59:49 -0700
Subject: [PATCH 181/478] Fix wrong register constraints produced for some
 types in WGMMAWaitGroupOpPattern

This is a minimal patch to unblock folks. I'll work on polishing it up and upstreaming a proper fix to OpenAI.

PiperOrigin-RevId: 633955310
---
 .../temporary/fix_register_constraints.patch  | 46 +++++++++++++++++++
 third_party/triton/temporary/series.bzl       |  1 +
 .../temporary/fix_register_constraints.patch  | 46 +++++++++++++++++++
 .../third_party/triton/temporary/series.bzl   |  1 +
 4 files changed, 94 insertions(+)
 create mode 100644 third_party/triton/temporary/fix_register_constraints.patch
 create mode 100644 third_party/xla/third_party/triton/temporary/fix_register_constraints.patch

diff --git a/third_party/triton/temporary/fix_register_constraints.patch b/third_party/triton/temporary/fix_register_constraints.patch
new file mode 100644
index 00000000000000..2caeac268e39db
--- /dev/null
+++ b/third_party/triton/temporary/fix_register_constraints.patch
@@ -0,0 +1,46 @@
+This should be fixed properly and upstreamed to OpenAI, tracked in b/336865530.
+
+diff --git a/third_party/nvidia/lib/NVGPUToLLVM/NVGPUToLLVMPass.cpp b/third_party/nvidia/lib/NVGPUToLLVM/NVGPUToLLVMPass.cpp
+--- a/third_party/nvidia/lib/NVGPUToLLVM/NVGPUToLLVMPass.cpp
++++ b/third_party/nvidia/lib/NVGPUToLLVM/NVGPUToLLVMPass.cpp
+@@ -480,10 +480,36 @@ public:
+   std::vector<std::string>
+   getOutputConstraints(ttn::WGMMAWaitGroupOp op) const {
+     auto outputStructType = cast<LLVM::LLVMStructType>(op.getType());
+-    uint32_t numOutputRegs = outputStructType.getBody().size();
+-    std::string output =
+-        outputStructType.getBody().front().isF32() ? "=f" : "=r";
+-    return std::vector<std::string>(numOutputRegs, output);
++    std::vector<std::string> outputConstraints;
++    outputConstraints.reserve(outputStructType.getBody().size());
++    for (mlir::Type type : outputStructType.getBody()) {
++      if (type.isF32()) {
++        outputConstraints.push_back("=f");
++        continue;
++      } else if (type.isF64()) {
++        outputConstraints.push_back("=d");
++        continue;
++      }
++      unsigned bitwidth = isa<LLVM::LLVMPointerType>(type) ?
++          64 : type.getIntOrFloatBitWidth();
++      switch (bitwidth) {
++        case 1:
++          outputConstraints.push_back("=b");
++          break;
++        case 16:
++          outputConstraints.push_back("=h");
++          break;
++        case 32:
++          outputConstraints.push_back("=r");
++          break;
++        case 64:
++          outputConstraints.push_back("=l");
++          break;
++        default:
++          assert(false && "unsupported bitwidth");
++      }
++    }
++    return outputConstraints;
+   }
+ 
+   OperandsAndConstraints
diff --git a/third_party/triton/temporary/series.bzl b/third_party/triton/temporary/series.bzl
index 214666767949de..382cb9e9de0154 100644
--- a/third_party/triton/temporary/series.bzl
+++ b/third_party/triton/temporary/series.bzl
@@ -7,4 +7,5 @@ internal patch during the next triton integration process.
 
 temporary_patch_list = [
     "//third_party/triton/temporary:mma_limit_pred.patch",
+    "//third_party/triton/temporary:fix_register_constraints.patch",
 ]
diff --git a/third_party/xla/third_party/triton/temporary/fix_register_constraints.patch b/third_party/xla/third_party/triton/temporary/fix_register_constraints.patch
new file mode 100644
index 00000000000000..2caeac268e39db
--- /dev/null
+++ b/third_party/xla/third_party/triton/temporary/fix_register_constraints.patch
@@ -0,0 +1,46 @@
+This should be fixed properly and upstreamed to OpenAI, tracked in b/336865530.
+
+diff --git a/third_party/nvidia/lib/NVGPUToLLVM/NVGPUToLLVMPass.cpp b/third_party/nvidia/lib/NVGPUToLLVM/NVGPUToLLVMPass.cpp
+--- a/third_party/nvidia/lib/NVGPUToLLVM/NVGPUToLLVMPass.cpp
++++ b/third_party/nvidia/lib/NVGPUToLLVM/NVGPUToLLVMPass.cpp
+@@ -480,10 +480,36 @@ public:
+   std::vector<std::string>
+   getOutputConstraints(ttn::WGMMAWaitGroupOp op) const {
+     auto outputStructType = cast<LLVM::LLVMStructType>(op.getType());
+-    uint32_t numOutputRegs = outputStructType.getBody().size();
+-    std::string output =
+-        outputStructType.getBody().front().isF32() ? "=f" : "=r";
+-    return std::vector<std::string>(numOutputRegs, output);
++    std::vector<std::string> outputConstraints;
++    outputConstraints.reserve(outputStructType.getBody().size());
++    for (mlir::Type type : outputStructType.getBody()) {
++      if (type.isF32()) {
++        outputConstraints.push_back("=f");
++        continue;
++      } else if (type.isF64()) {
++        outputConstraints.push_back("=d");
++        continue;
++      }
++      unsigned bitwidth = isa<LLVM::LLVMPointerType>(type) ?
++          64 : type.getIntOrFloatBitWidth();
++      switch (bitwidth) {
++        case 1:
++          outputConstraints.push_back("=b");
++          break;
++        case 16:
++          outputConstraints.push_back("=h");
++          break;
++        case 32:
++          outputConstraints.push_back("=r");
++          break;
++        case 64:
++          outputConstraints.push_back("=l");
++          break;
++        default:
++          assert(false && "unsupported bitwidth");
++      }
++    }
++    return outputConstraints;
+   }
+ 
+   OperandsAndConstraints
diff --git a/third_party/xla/third_party/triton/temporary/series.bzl b/third_party/xla/third_party/triton/temporary/series.bzl
index 214666767949de..382cb9e9de0154 100644
--- a/third_party/xla/third_party/triton/temporary/series.bzl
+++ b/third_party/xla/third_party/triton/temporary/series.bzl
@@ -7,4 +7,5 @@ internal patch during the next triton integration process.
 
 temporary_patch_list = [
     "//third_party/triton/temporary:mma_limit_pred.patch",
+    "//third_party/triton/temporary:fix_register_constraints.patch",
 ]

From 47bbe22a6b1e678c07c6b433d5ea848244e84a29 Mon Sep 17 00:00:00 2001
From: ekuznetsov139 <nameless@fastmail.fm>
Date: Wed, 15 May 2024 07:59:50 -0700
Subject: [PATCH 182/478] PR #9827: [ROCm] Stream executor API logging API

Imported from GitHub PR https://github.com/openxla/xla/pull/9827

Separated from https://github.com/openxla/xla/pull/9593.

This adds a low-overhead ability to log stream executor API calls, for testing purposes (unit tests can enable the functionality and use it to verify that high-level operations result in correct sequences of low-level calls). At present, this only covers BlasGemm / BlasLt calls on ROCm backend.
Copybara import of the project:

--
41732ac48c5383f33d428fe76fcebe02506d5b52 by Eugene Kuznetsov <eugene.kuznetsov@amd.com>:

Stream executor API logging API

Merging this change closes #9827

PiperOrigin-RevId: 633955314
---
 .../xla/stream_executor/gpu/gpu_executor.h    | 28 +++++++++++
 .../xla/stream_executor/rocm/hip_blas_lt.cc   |  5 ++
 .../xla/xla/stream_executor/rocm/rocm_blas.cc | 46 ++++++++++++++++++-
 .../xla/xla/stream_executor/rocm/rocm_blas.h  |  5 ++
 .../stream_executor_interface.h               | 32 +++++++++++++
 5 files changed, 115 insertions(+), 1 deletion(-)

diff --git a/third_party/xla/xla/stream_executor/gpu/gpu_executor.h b/third_party/xla/xla/stream_executor/gpu/gpu_executor.h
index 4fda798c54078f..c7bd541f2d0c8a 100644
--- a/third_party/xla/xla/stream_executor/gpu/gpu_executor.h
+++ b/third_party/xla/xla/stream_executor/gpu/gpu_executor.h
@@ -317,6 +317,28 @@ class GpuExecutor : public StreamExecutor {
   int cc_major() const { return cc_major_; }
   int cc_minor() const { return cc_minor_; }
 
+  absl::StatusOr<std::vector<ApiTrace>> ExtractApiTrace() override {
+    absl::MutexLock lock(&logger_mu_);
+    return std::move(argument_logs_);
+  }
+
+  absl::Status RecordApiTrace(ApiTrace call) override {
+    absl::MutexLock lock(&logger_mu_);
+    if (std::holds_alternative<GemmCallTrace>(call) &&
+        (argument_logging_mode_ & kLogGemm)) {
+      argument_logs_.push_back(call);
+    }
+    return absl::OkStatus();
+  }
+
+  bool SetArgumentLoggingMode(uint64_t mode) override {
+    absl::MutexLock lock(&logger_mu_);
+    argument_logging_mode_ = mode;
+    return true;
+  }
+
+  uint64_t GetArgumentLoggingMode() const { return argument_logging_mode_; }
+
  private:
   // Host callback landing routine invoked by CUDA.
   // data: User-provided callback provided to HostCallback() above, captured
@@ -437,6 +459,12 @@ class GpuExecutor : public StreamExecutor {
   // for a BLAS interface.
   std::unique_ptr<blas::BlasSupport> blas_ ABSL_GUARDED_BY(mu_);
 
+  absl::Mutex logger_mu_;
+
+  mutable std::vector<ApiTrace> argument_logs_ ABSL_GUARDED_BY(logger_mu_);
+
+  uint64_t argument_logging_mode_ = 0;
+
   GpuExecutor(const GpuExecutor&) = delete;
   void operator=(const GpuExecutor&) = delete;
 };
diff --git a/third_party/xla/xla/stream_executor/rocm/hip_blas_lt.cc b/third_party/xla/xla/stream_executor/rocm/hip_blas_lt.cc
index 6ce3ccad39b5da..2061708a736f6f 100644
--- a/third_party/xla/xla/stream_executor/rocm/hip_blas_lt.cc
+++ b/third_party/xla/xla/stream_executor/rocm/hip_blas_lt.cc
@@ -418,6 +418,11 @@ absl::Status BlasLt::MatmulPlan::DoMatmul(
     std::optional<DeviceMemoryBase> workspace,
     std::optional<ScratchAllocator*> scratch_allocator,
     blas::ProfileResult* profile_result) const {
+  absl::Status status = blas_lt_ref_.parent_->RecordApiTrace(
+      StreamExecutorInterface::GemmCallTrace{
+          StreamExecutorInterface::GemmCallTrace::GemmType::kBlasLt, 0,
+          a.size(), b.size()});
+
   TF_ASSIGN_OR_RETURN(
       std::optional<gpu::GpuTimer> timer,
       gpu::GpuTimer::CreateIfNeeded(
diff --git a/third_party/xla/xla/stream_executor/rocm/rocm_blas.cc b/third_party/xla/xla/stream_executor/rocm/rocm_blas.cc
index 5686a4791de9b4..8ca9f9ca9fc653 100644
--- a/third_party/xla/xla/stream_executor/rocm/rocm_blas.cc
+++ b/third_party/xla/xla/stream_executor/rocm/rocm_blas.cc
@@ -222,6 +222,26 @@ rocblas_side ROCMBlasSide(blas::Side side) {
   }
 }
 
+int DtypeSize(blas::DataType type) {
+  switch (type) {
+    case blas::DataType::kHalf:
+    case blas::DataType::kBF16:
+      return 2;
+    case blas::DataType::kFloat:
+      return 4;
+    case blas::DataType::kDouble:
+      return 8;
+    case blas::DataType::kInt8:
+      return 1;
+    case blas::DataType::kComplexFloat:
+      return 8;
+    case blas::DataType::kComplexDouble:
+      return 16;
+    default:
+      return 0;
+  }
+}
+
 absl::StatusOr<rocblas_datatype> AsRocBlasType(blas::DataType type) {
   switch (type) {
     case blas::DataType::kHalf:
@@ -432,6 +452,16 @@ Impl_DoBlasScal(wrap::rocblas_sscal, float,
  *    and ex functions expect the same type as the compute type (i.e. floats.)
  *
  **/
+using sei = StreamExecutorInterface;
+using GemmCallTrace = sei::GemmCallTrace;
+
+// Log the GEMM operation if the logging mode is enabled.
+void ROCMBlas::MaybeLogGemmOp(GemmCallTrace::GemmType op,
+                              blas::CallContext context, uint64_t size1,
+                              uint64_t size2) {
+  auto status =
+      parent_->RecordApiTrace(GemmCallTrace{op, (int)context, size1, size2});
+}
 
 absl::Status ROCMBlas::DoBlasGemm(
     Stream *stream, blas::Transpose transa, blas::Transpose transb, uint64_t m,
@@ -439,6 +469,9 @@ absl::Status ROCMBlas::DoBlasGemm(
     const DeviceMemoryBase &a, int lda, const DeviceMemoryBase &b, int ldb,
     const void *beta, DeviceMemoryBase *c, int ldc,
     const NumericOptions &numeric_options, blas::CallContext context) {
+  MaybeLogGemmOp(GemmCallTrace::GemmType::kPlain, context,
+                 m * k * DtypeSize(dtype), n * k * DtypeSize(dtype));
+
   VLOG(1) << absl::StreamFormat(
       "doing rocBLAS GEMM: at=%d bt=%d m=%u n=%u "
       "k=%llu alpha=%p a=%p lda=%d b=%p ldb=%d beta=%p "
@@ -542,6 +575,8 @@ absl::Status ROCMBlas::DoBlasGemmWithAlgorithm(
                                   numeric_options, context));
 
   } else {
+    MaybeLogGemmOp(GemmCallTrace::GemmType::kPlain, context,
+                   m * k * DtypeSize(type_a), n * k * DtypeSize(type_a));
     CheckPreconditions(transa, transb, m, n, k, type_a, lda, ldb);
     TF_ASSIGN_OR_RETURN(auto roc_type_a, AsRocBlasType(type_a));
     TF_ASSIGN_OR_RETURN(auto roc_type_c, AsRocBlasType(type_c));
@@ -601,6 +636,8 @@ absl::Status ROCMBlas::DoBlasGemmStridedBatchedWithAlgorithm(
         ldb, stride_b, beta, c, ldc, stride_c, batch_count, numeric_options,
         context));
   } else {
+    MaybeLogGemmOp(GemmCallTrace::GemmType::kStridedBatched, context, a.size(),
+                   b.size());
     VLOG(1) << absl::StreamFormat(
         "doing rocBLAS GEMM strided batched with Algorithm: at=%d bt=%d m=%u "
         "n=%u "
@@ -699,7 +736,6 @@ bool ROCMBlas::GetBlasGemmAlgorithms(
   ASSIGN_OR_FALSE(auto roc_comp_type, AsRocBlasComputeType(c->compute_type));
 
   if (c->batch_size == 1) {
-    // TODO: we should possibly use GemmFloat16Flags(type_a, context) here..
     return DoBlasInternalFailureOK(
         NameWrap{blas_lambda}, stream, true,
         wrap::rocblas_gemm_ex_get_solutions, ROCMBlasTranspose(a.transpose),
@@ -1032,6 +1068,8 @@ bool ROCMBlas::DoBlasGemmBatched(
     DeviceMemorySlice<Eigen::half> c, int ldc, int batch_count,
     const NumericOptions &numeric_options, ScratchAllocator *scratch_allocator,
     blas::CallContext context) {
+  MaybeLogGemmOp(GemmCallTrace::GemmType::kBatched, context, a.size(),
+                 b.size());
   const Eigen::half alpha_half(alpha);
   const Eigen::half beta_half(beta);
   absl::Status status;
@@ -1066,6 +1104,8 @@ bool ROCMBlas::DoBlasGemmBatched(
     DeviceMemorySlice<Eigen::bfloat16> c_array, int ldc, int batch_count,
     const NumericOptions &numeric_options, ScratchAllocator *scratch_allocator,
     blas::CallContext context) {
+  MaybeLogGemmOp(GemmCallTrace::GemmType::kBatched, context, a_array.size(),
+                 b_array.size());
   const Eigen::bfloat16 alpha_bf16(alpha);
   const Eigen::bfloat16 beta_bf16(beta);
 
@@ -1087,6 +1127,8 @@ bool ROCMBlas::DoBlasGemmBatched(
       DeviceMemorySlice<T> c_array, int ldc, int batch_count,                  \
       const NumericOptions &numeric_options,                                   \
       ScratchAllocator *scratch_allocator, blas::CallContext context) {        \
+    MaybeLogGemmOp(GemmCallTrace::GemmType::kBatched, context, a_array.size(), \
+                   b_array.size());                                            \
     absl::Status status = DoBlasGemmBatchedInternal(                           \
         Fun, stream, transa, transb, m, n, k, alpha, a_array, lda, b_array,    \
         ldb, beta, c_array, ldc, batch_count, scratch_allocator);              \
@@ -1153,6 +1195,8 @@ IMPL_DoBlasGemmBatched(float, wrap::rocblas_sgemm_strided_batched)
       static_cast<int>(transa), static_cast<int>(transb), m, n, k, alpha,
       a.opaque(), lda, b.opaque(), ldb, beta, c->opaque(), ldc, stride_a,
       stride_b, stride_c, batch_count);
+  MaybeLogGemmOp(GemmCallTrace::GemmType::kStridedBatched, context, a.size(),
+                 b.size());
 
   absl::Status status;
   auto call_gemm = [&](auto func, auto type) {
diff --git a/third_party/xla/xla/stream_executor/rocm/rocm_blas.h b/third_party/xla/xla/stream_executor/rocm/rocm_blas.h
index 537a3a7a46f07a..d8bb216ebeff8e 100644
--- a/third_party/xla/xla/stream_executor/rocm/rocm_blas.h
+++ b/third_party/xla/xla/stream_executor/rocm/rocm_blas.h
@@ -38,6 +38,7 @@ limitations under the License.
 #if TF_HIPBLASLT
 #include "xla/stream_executor/rocm/hip_blas_lt.h"
 #endif
+#include "xla/stream_executor/stream_executor_interface.h"
 
 namespace stream_executor {
 
@@ -199,6 +200,10 @@ class ROCMBlas : public blas::BlasSupport {
   // container holding solutions vector (to avoid reallocating it each time)
   std::vector<rocblas_int> solutions_;
 
+  void MaybeLogGemmOp(StreamExecutorInterface::GemmCallTrace::GemmType op,
+                      blas::CallContext context, uint64_t size1,
+                      uint64_t size2);
+
 #if TF_HIPBLASLT
   rocm::BlasLt blas_lt_;
 #endif
diff --git a/third_party/xla/xla/stream_executor/stream_executor_interface.h b/third_party/xla/xla/stream_executor/stream_executor_interface.h
index 65eff81b0a4d25..f310cce8b7eaca 100644
--- a/third_party/xla/xla/stream_executor/stream_executor_interface.h
+++ b/third_party/xla/xla/stream_executor/stream_executor_interface.h
@@ -390,6 +390,38 @@ class StreamExecutorInterface {
   // Returns the memory limit in bytes supported by this executor.
   virtual int64_t GetMemoryLimitBytes() const = 0;
 
+  // The following methods access an internal log of some subset
+  // of arguments passed to other class methods.
+  // Used for testing/debugging purposes.
+
+  struct GemmCallTrace {
+    enum class GemmType {
+      kPlain = 0,
+      kStridedBatched = 1,
+      kBatched = 2,
+      kBlasLt = 3
+    };
+    GemmType op;
+    int flags;
+    uint64_t size1, size2;
+  };
+  // This may be expanded as necessary to trace other calls
+  using ApiTrace = std::variant<GemmCallTrace>;
+
+  // Retrieves and clears internal argument logs.
+  virtual absl::StatusOr<std::vector<ApiTrace>> ExtractApiTrace() {
+    return absl::UnimplementedError("Not implemented");
+  }
+  virtual absl::Status RecordApiTrace(ApiTrace call) {
+    return absl::UnimplementedError("Not implemented");
+  }
+
+  static constexpr uint64_t kLogGemm = 1 << 0;
+
+  // Sets the argument logging mode. Returns true if 'mode' is valid.
+  // The mode is a bitmask of the kLog* constants.
+  virtual bool SetArgumentLoggingMode(uint64_t mode) { return false; }
+
  private:
   StreamExecutorInterface(const StreamExecutorInterface&) = delete;
   void operator=(const StreamExecutorInterface&) = delete;

From 132731343ad4b54dd86a2edfae2dac8c21a55adf Mon Sep 17 00:00:00 2001
From: Deqiang Chen <deqiangc@google.com>
Date: Wed, 15 May 2024 08:34:54 -0700
Subject: [PATCH 183/478] Add unit test for async_value_tensor.

PiperOrigin-RevId: 633965643
---
 tensorflow/core/tfrt/common/BUILD             | 13 ++++
 .../tfrt/common/async_value_tensor_test.cc    | 76 +++++++++++++++++++
 2 files changed, 89 insertions(+)
 create mode 100644 tensorflow/core/tfrt/common/async_value_tensor_test.cc

diff --git a/tensorflow/core/tfrt/common/BUILD b/tensorflow/core/tfrt/common/BUILD
index 2de9b7fcf904b3..bda3e3dde478e6 100644
--- a/tensorflow/core/tfrt/common/BUILD
+++ b/tensorflow/core/tfrt/common/BUILD
@@ -70,6 +70,19 @@ cc_library(
     ],
 )
 
+tf_cc_test(
+    name = "async_value_tensor_test",
+    srcs = ["async_value_tensor_test.cc"],
+    deps = [
+        ":async_value_tensor",
+        "//tensorflow/core:framework",
+        "//tensorflow/core:protos_all_cc",
+        "@com_google_googletest//:gtest_main",
+        "@local_xla//xla/pjrt:pjrt_client",
+        "@local_xla//xla/tsl/concurrency:async_value",
+    ],
+)
+
 cc_library(
     name = "pjrt_state",
     srcs = [
diff --git a/tensorflow/core/tfrt/common/async_value_tensor_test.cc b/tensorflow/core/tfrt/common/async_value_tensor_test.cc
new file mode 100644
index 00000000000000..a7c89af53c9da7
--- /dev/null
+++ b/tensorflow/core/tfrt/common/async_value_tensor_test.cc
@@ -0,0 +1,76 @@
+/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+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.
+==============================================================================*/
+
+#include "tensorflow/core/tfrt/common/async_value_tensor.h"
+
+#include <cstdint>
+#include <memory>
+
+#include <gtest/gtest.h>
+#include "xla/pjrt/pjrt_client.h"
+#include "xla/tsl/concurrency/async_value_ref.h"
+#include "tensorflow/core/framework/tensor.h"
+#include "tensorflow/core/framework/tensor_shape.h"
+#include "tensorflow/core/framework/types.pb.h"
+
+namespace tensorflow {
+namespace {
+
+TEST(AsyncValueTensorTest, InvalidTensor) {
+  tensorflow::Tensor tensor(tensorflow::DT_INT64, tensorflow::TensorShape({1}));
+
+  AsyncValueTensor* avt = AsyncValueTensor::FromTensor(&tensor);
+
+  ASSERT_EQ(avt, nullptr);
+}
+
+TEST(AsyncValueTensorTest, SetAndGetAsyncValue) {
+  AsyncValueAllocator allocator;
+  tensorflow::Tensor tensor(&allocator, tensorflow::DT_INT64,
+                            tensorflow::TensorShape({1}));
+
+  AsyncValueTensor* avt = AsyncValueTensor::FromTensor(&tensor);
+
+  ASSERT_NE(avt, nullptr);
+
+  tsl::AsyncValueRef<int32_t> value =
+      tsl::MakeConstructedAsyncValueRef<int32_t>(123);
+
+  avt->SetAsyncRef(value.CopyRCRef());
+
+  auto ret_value = avt->GetAsyncRef();
+  ASSERT_EQ(ret_value, value.CopyRCRef());
+}
+
+TEST(AsyncValueTensorTest, SetAndGetBuffer) {
+  AsyncValueAllocator allocator;
+  tensorflow::Tensor tensor(&allocator, tensorflow::DT_INT64,
+                            tensorflow::TensorShape({1}));
+
+  AsyncValueTensor* avt = AsyncValueTensor::FromTensor(&tensor);
+
+  ASSERT_NE(avt, nullptr);
+
+  std::shared_ptr<xla::PjRtBuffer> buffer;
+
+  avt->SetBuffer(buffer);
+
+  auto ret_buffer = avt->GetBuffer();
+
+  ASSERT_EQ(ret_buffer, buffer);
+}
+
+}  // namespace
+}  // namespace tensorflow

From b82d42294d2f747dc01dfc798a2c0c5b74d51f25 Mon Sep 17 00:00:00 2001
From: Derek Murray <mrry@google.com>
Date: Wed, 15 May 2024 09:04:03 -0700
Subject: [PATCH 184/478] Optimize gradient of `tf.reshape()` when one input
 dimension is dynamic.

Previously, if there was any unknown dimension in the shape of the input to `tf.reshape()` the gradient function would add a `tf.shape()` op that would introduce a data dependency on the input tensor, potentially keeping it live until the backward pass. However, in most cases* when only one of the input dimensions is unknown (e.g. a dynamic batch size) we can use the ability to pass a wildcard `-1` dimension to the `tf.reshape()` in the gradient function, and eliminate the data dependency.

(* Except when one or more of the input dimensions is 0, in which case the semantics of the wildcard are surprising, and resolve to 1.)

PiperOrigin-RevId: 633973978
---
 tensorflow/python/ops/array_grad.py      | 21 ++++++++++++++
 tensorflow/python/ops/array_grad_test.py | 35 ++++++++++++++++++++++++
 2 files changed, 56 insertions(+)

diff --git a/tensorflow/python/ops/array_grad.py b/tensorflow/python/ops/array_grad.py
index c5ba8157dc73e2..7f856d10b0c082 100644
--- a/tensorflow/python/ops/array_grad.py
+++ b/tensorflow/python/ops/array_grad.py
@@ -811,6 +811,27 @@ def _IdNGrad(_, *grad):
 
 @ops.RegisterGradient("Reshape")
 def _ReshapeGrad(op: ops.Operation, grad):
+  """Defines the gradient for `array_ops.reshape()`."""
+  input_shape = op.inputs[0].shape
+  if input_shape.rank is not None and not input_shape.is_fully_defined():
+    # If only one dimension is undefined, we can use a wildcard dimension in
+    # the argument to `reshape()` to avoid creating a data dependency via
+    # a dynamic `shape()` operation.
+    input_shape_as_list = input_shape.as_list()
+    undefined_dims = []
+    has_zero_dim = False
+    for i, dim in enumerate(input_shape_as_list):
+      if dim is None:
+        undefined_dims.append(i)
+      elif dim == 0:
+        # When the tensor has zero elements, the wildcard dimension
+        # is underspecified, and `reshape()` will arbitrarily set the unknown
+        # to `1`, triggering shape errors downstream.
+        has_zero_dim = True
+    if len(undefined_dims) == 1 and not has_zero_dim:
+      input_shape_as_list[undefined_dims[0]] = -1
+      return [array_ops.reshape(_IndexedSlicesToTensorNoWarning(grad),
+                                input_shape_as_list), None]
   return [
       array_ops.reshape(
           _IndexedSlicesToTensorNoWarning(grad), array_ops.shape(op.inputs[0])),
diff --git a/tensorflow/python/ops/array_grad_test.py b/tensorflow/python/ops/array_grad_test.py
index 3dfe394a2c15e3..1f7c65e5b6fe70 100644
--- a/tensorflow/python/ops/array_grad_test.py
+++ b/tensorflow/python/ops/array_grad_test.py
@@ -147,6 +147,41 @@ def f(x):
 
     self._testGrad(f, x)
 
+  def test_reshape_simple(self):
+    x = constant_op.constant([1., 2., 3.], dtype=dtypes.float64)
+    y = constant_op.constant([3, 1], dtype=dtypes.int64)
+
+    def f(x):
+      return array_ops.reshape(x, y)
+
+    self._testGrad(f, x)
+
+  def test_reshape_one_unknown_dim(self):
+    def f(x):
+      x_without_shape = array_ops.placeholder_with_default(x, shape=[None, 2])
+      return array_ops.reshape(x_without_shape, [3, 2])
+
+    x = constant_op.constant([[1., 2.], [3., 4.], [5., 6.]],
+                             dtype=dtypes.float64)
+    self._testGrad(f, x)
+
+  def test_reshape_two_unknown_dims(self):
+    def f(x):
+      x_without_shape = array_ops.placeholder_with_default(x,
+                                                           shape=[None, None])
+      return array_ops.reshape(x_without_shape, [6])
+
+    x = constant_op.constant([[1., 2., 3.], [4., 5., 6.]], dtype=dtypes.float64)
+    self._testGrad(f, x)
+
+  def test_reshape_one_unknown_dim_and_zero_elements(self):
+    def f(x):
+      x_without_shape = array_ops.placeholder_with_default(x, shape=[None, 0])
+      return array_ops.reshape(x_without_shape, [0])
+
+    x = constant_op.constant([], shape=[3, 0], dtype=dtypes.float64)
+    self._testGrad(f, x)
+
 
 if __name__ == "__main__":
   test.main()

From 5ec49b7e531b01320fe66403e7c498fd6ea487f5 Mon Sep 17 00:00:00 2001
From: Zixuan Jiang <zixuanjiang@google.com>
Date: Wed, 15 May 2024 09:18:31 -0700
Subject: [PATCH 185/478] In `xla::ShardingPropagation`, remove users and the
 parameters in the called computations in `already_inferred_from_operands`
 when clearing the cache.

Before this cl, shardings may not be completely propagated.

PiperOrigin-RevId: 633978152
---
 .../xla/xla/service/sharding_propagation.cc   |   8 +
 .../xla/service/sharding_propagation_test.cc  | 145 +++++++++++-------
 2 files changed, 97 insertions(+), 56 deletions(-)

diff --git a/third_party/xla/xla/service/sharding_propagation.cc b/third_party/xla/xla/service/sharding_propagation.cc
index 14bd73e50cbce5..dba8a717342596 100644
--- a/third_party/xla/xla/service/sharding_propagation.cc
+++ b/third_party/xla/xla/service/sharding_propagation.cc
@@ -3281,6 +3281,14 @@ absl::StatusOr<bool> ShardingPropagation::Run(
           }
           for (auto user : hlo_for_users->users()) {
             already_inferred_from_operands.erase(user);
+            // If the user has called computations, then the parameter
+            // instructions of these called computations are also removed from
+            // already_inferred_from_operands.
+            for (auto c : user->called_computations()) {
+              for (auto parameter : c->parameter_instructions()) {
+                already_inferred_from_operands.erase(parameter);
+              }
+            }
           }
           if (instruction_to_shard_group_id.contains(hlo)) {
             const int64_t shard_group_id =
diff --git a/third_party/xla/xla/service/sharding_propagation_test.cc b/third_party/xla/xla/service/sharding_propagation_test.cc
index 4861e4f0b75c5a..e06ecd3d3ca2a9 100644
--- a/third_party/xla/xla/service/sharding_propagation_test.cc
+++ b/third_party/xla/xla/service/sharding_propagation_test.cc
@@ -3843,26 +3843,43 @@ TEST_P(ParameterizedMetadataTest, Conditional) {
   const char* const hlo_string = R"(
 HloModule module
 
+%add-call {
+  %x = f32[4,4] parameter(0)
+  ROOT %add = f32[4,4] add(%x, %x)
+}
+
 %true_comp {
-  %tp = (f32[3,5]) parameter(0)
-  %tgte = f32[3,5] get-tuple-element(%tp), index=0
-  %ttr = f32[5,3] transpose(%tgte), dimensions={1,0}
-  ROOT %tr = (f32[5,3]) tuple(%ttr)
+  %tp = (f32[3,5], f32[4,4]) parameter(0)
+  %tgte.0 = f32[3,5] get-tuple-element(%tp), index=0
+  %ttr = f32[5,3] transpose(%tgte.0), dimensions={1,0}
+  %tgte.1 = f32[4,4] get-tuple-element(%tp), index=1
+  %tadd = f32[4,4] call(%tgte.1), to_apply=%add-call
+  ROOT %tr = (f32[5,3], f32[4,4]) tuple(%ttr, %tadd)
+}
+
+%mul-call {
+  %y = f32[4,4] parameter(0)
+  ROOT %mul = f32[4,4] multiply(%y, %y)
 }
 
 %false_comp {
-  %fp = (f32[5,3]) parameter(0)
-  %fgte = f32[5,3] get-tuple-element(%fp), index=0
-  ROOT %fr = (f32[5,3]) tuple(%fgte)
+  %fp = (f32[5,3], f32[4,4]) parameter(0)
+  %fgte.0 = f32[5,3] get-tuple-element(%fp), index=0
+  %fgte.1 = f32[4,4] get-tuple-element(%fp), index=1
+  %fmul = f32[4,4] call(%fgte.1), to_apply=%mul-call
+  ROOT %fr = (f32[5,3], f32[4,4]) tuple(%fgte.0, %fmul)
 }
 
 ENTRY entry {
   %cond = pred[] parameter(0)
-  %true_param = (f32[3,5]) parameter(1),
-    sharding={{devices=[1,2]0,1 metadata={op_name="a"}}}
-  %false_param = (f32[5,3]) parameter(2),
-    sharding={{devices=[1,3]0,1,2 metadata={op_name="b"}}}
-  %conditional = (f32[5,3]) conditional(
+  %tp.0 = f32[3,5] parameter(1), sharding={devices=[1,2]0,1 metadata={op_name="a"}}
+  %fp.0 = f32[5,3] parameter(2), sharding={devices=[1,3]0,1,2 metadata={op_name="b"}}
+  %constant = f32[4] constant({1,2,3,4}), sharding={devices=[4]0,1,2,3 metadata={op_name="c"}}
+  %broadcast = f32[4,4] broadcast(%constant), dimensions={1}
+  %add = f32[4,4] add(%broadcast, %broadcast)
+  %true_param = (f32[3,5], f32[4,4]) tuple(%tp.0, %add)
+  %false_param = (f32[5,3], f32[4,4]) tuple(%fp.0, %add)
+  %conditional = (f32[5,3], f32[4,4]) conditional(
       %cond, %true_param, %false_param),
     true_computation=%true_comp,
     false_computation=%false_comp
@@ -3879,60 +3896,76 @@ ENTRY entry {
           .Run(module.get()));
   XLA_VLOG_LINES(1, module->ToString());
   EXPECT_TRUE(changed);
+
   auto* tp = FindInstruction(module.get(), "tp");
-  ASSERT_NE(tp, nullptr);
-  EXPECT_THAT(tp, op::Sharding("{{devices=[1,2]0,1}}"));
-  auto* tgte = FindInstruction(module.get(), "tgte");
-  ASSERT_NE(tgte, nullptr);
-  EXPECT_THAT(tgte, op::Sharding("{devices=[1,2]0,1}"));
+  auto* tgte_0 = FindInstruction(module.get(), "tgte.0");
   auto* ttr = FindInstruction(module.get(), "ttr");
-  ASSERT_NE(ttr, nullptr);
-  EXPECT_THAT(ttr, op::Sharding("{devices=[2,1]0,1}"));
+  auto* tgte_1 = FindInstruction(module.get(), "tgte.1");
+  auto* tadd = FindInstruction(module.get(), "tadd");
   auto* tr = FindInstruction(module.get(), "tr");
-  ASSERT_NE(tr, nullptr);
-  EXPECT_THAT(tr, op::Sharding("{{devices=[1,3]0,1,2}}"));
+
   auto* fp = FindInstruction(module.get(), "fp");
-  ASSERT_NE(fp, nullptr);
-  EXPECT_THAT(fp, op::Sharding("{{devices=[1,3]0,1,2}}"));
-  auto* fgte = FindInstruction(module.get(), "fgte");
-  ASSERT_NE(fgte, nullptr);
-  EXPECT_THAT(fgte, op::Sharding("{devices=[1,3]0,1,2}"));
+  auto* fgte_0 = FindInstruction(module.get(), "fgte.0");
+  auto* fgte_1 = FindInstruction(module.get(), "fgte.1");
+  auto* fmul = FindInstruction(module.get(), "fmul");
   auto* fr = FindInstruction(module.get(), "fr");
-  ASSERT_NE(fr, nullptr);
-  EXPECT_THAT(fr, op::Sharding("{{devices=[1,3]0,1,2}}"));
+
+  auto* x = FindInstruction(module.get(), "x");
+  auto* add = FindInstruction(module.get(), "add");
+  auto* y = FindInstruction(module.get(), "y");
+  auto* mul = FindInstruction(module.get(), "mul");
+
   auto* conditional = FindInstruction(module.get(), "conditional");
-  ASSERT_NE(conditional, nullptr);
-  EXPECT_THAT(conditional, op::Sharding("{{devices=[1,3]0,1,2}}"));
 
-  auto check_metadata = [](const HloSharding& sharding,
-                           const OpMetadata& metadata) {
-    if (sharding.IsTuple()) {
-      EXPECT_THAT(sharding.tuple_elements()[0], ShardingMetadata({metadata}));
-    } else {
-      EXPECT_THAT(sharding, ShardingMetadata({metadata}));
-    }
-  };
+  const std::vector<HloInstruction*> instructions(
+      {tp, tgte_0, ttr, tgte_1, tadd, tr, fp, fgte_0, fgte_1, fmul, fr, x, add,
+       y, mul, conditional});
 
-  auto check_empty_metadata = [](const HloSharding& sharding) {
-    if (sharding.IsTuple()) {
-      EXPECT_THAT(sharding.tuple_elements()[0], ShardingMetadata({}));
-    } else {
-      EXPECT_THAT(sharding, ShardingMetadata({}));
-    }
-  };
+  for (HloInstruction* instruction : instructions) {
+    EXPECT_NE(instruction, nullptr);
+    EXPECT_TRUE(instruction->has_sharding());
+  }
 
-  for (HloInstruction* instruction : {tp, tgte, ttr}) {
-    if (GetParam().propagate_metadata && !GetParam().clear_metadata) {
-      check_metadata(instruction->sharding(), CreateMetadata("a"));
-    } else {
-      check_empty_metadata(instruction->sharding());
-    }
+  for (HloInstruction* instruction :
+       {tgte_1, tadd, fgte_1, fmul, x, add, y, mul}) {
+    EXPECT_THAT(instruction, op::Sharding("{devices=[1,4]0,1,2,3}"));
   }
-  for (HloInstruction* instruction : {tr, fp, fgte, fr, conditional}) {
-    if (GetParam().propagate_metadata && !GetParam().clear_metadata) {
-      check_metadata(instruction->sharding(), CreateMetadata("b"));
-    } else {
-      check_empty_metadata(instruction->sharding());
+  for (HloInstruction* instruction : {tr, fr, conditional, fp}) {
+    EXPECT_THAT(instruction,
+                op::Sharding("{{devices=[1,3]0,1,2}, {devices=[1,4]0,1,2,3}}"));
+  }
+  EXPECT_THAT(tp, op::Sharding("{{devices=[1,2]0,1}, {devices=[1,4]0,1,2,3}}"));
+  EXPECT_THAT(tgte_0, op::Sharding("{devices=[1,2]0,1}"));
+  EXPECT_THAT(ttr, op::Sharding("{devices=[2,1]0,1}"));
+  EXPECT_THAT(fgte_0, op::Sharding("{devices=[1,3]0,1,2}"));
+
+  if (GetParam().propagate_metadata && !GetParam().clear_metadata) {
+    for (HloInstruction* instruction :
+         {tgte_1, tadd, fgte_1, fmul, x, add, y, mul}) {
+      EXPECT_THAT(instruction->sharding(),
+                  ShardingMetadata({CreateMetadata("c")}));
+    }
+    for (HloInstruction* instruction : {tr, fr, conditional, fp}) {
+      const std::vector<HloSharding>& shardings =
+          instruction->sharding().tuple_elements();
+      EXPECT_THAT(shardings[0], ShardingMetadata({CreateMetadata("b")}));
+      EXPECT_THAT(shardings[1], ShardingMetadata({CreateMetadata("c")}));
+    }
+    for (HloInstruction* instruction : {tgte_0, ttr}) {
+      EXPECT_THAT(instruction->sharding(),
+                  ShardingMetadata({CreateMetadata("a")}));
+    }
+    EXPECT_THAT(fgte_0->sharding(), ShardingMetadata({CreateMetadata("b")}));
+  } else {
+    for (HloInstruction* instruction : instructions) {
+      if (instruction->sharding().IsTuple()) {
+        for (const HloSharding& tuple_element :
+             instruction->sharding().tuple_elements()) {
+          EXPECT_THAT(tuple_element, ShardingMetadata({}));
+        }
+      } else {
+        EXPECT_THAT(instruction->sharding(), ShardingMetadata({}));
+      }
     }
   }
 }

From 0fc519c9dbc86276b1b37d1d440aaddb6d8801f5 Mon Sep 17 00:00:00 2001
From: Johannes Reifferscheid <jreiffers@google.com>
Date: Wed, 15 May 2024 09:27:48 -0700
Subject: [PATCH 186/478] Fix a batch of ClangTidy issues.

Mostly naming, some includes, some miscellaneous other style guide
violations.

PiperOrigin-RevId: 633980451
---
 .../xla/xla/mlir/tools/mlir_bisect/BUILD      |   1 -
 .../xla/mlir/tools/mlir_bisect/bisect_lib.cc  |   4 +-
 .../xla/mlir/tools/mlir_bisect/mlir_bisect.cc |   5 +-
 .../xla/mlir/tools/mlir_bisect/test_passes.cc |   7 +
 .../tools/mlir_interpreter/dialects/BUILD     |   4 -
 .../tools/mlir_interpreter/dialects/arith.cc  |  26 +--
 .../dialects/bufferization.cc                 |   4 +
 .../tools/mlir_interpreter/dialects/linalg.cc |  20 +-
 .../tools/mlir_interpreter/dialects/memref.cc |  30 +--
 .../tools/mlir_interpreter/dialects/util.cc   |   6 +-
 .../tools/mlir_interpreter/dialects/vector.cc | 202 +++++++++---------
 .../mlir_interpreter/framework/interpreter.cc |  20 +-
 .../mlir_interpreter/framework/interpreter.h  |   2 +-
 .../framework/interpreter_value.cc            |  21 +-
 .../framework/interpreter_value.h             |   4 +-
 .../framework/interpreter_value_util.h        |  24 +--
 .../framework/registration.cc                 |   4 +
 .../framework/tensor_or_memref.cc             |  59 ++---
 .../framework/tensor_or_memref.h              |  51 ++---
 .../framework/tests/interpreter_value_test.cc |  11 +-
 .../framework/tests/tensor_or_memref_test.cc  |  15 +-
 .../mlir_interpreter_runner.cc                |  30 +--
 .../xla/xla/mlir/tools/mlir_replay/BUILD      |   2 +
 .../xla/mlir/tools/mlir_replay/mlir_replay.cc |   4 +
 .../mlir/tools/mlir_replay/mlir_replay_lib.cc |  21 +-
 .../xla/mlir/tools/mlir_replay/public/BUILD   |   1 +
 .../public/execution_trace_utils.cc           |   8 +
 .../public/execution_trace_utils_test.cc      |   1 +
 28 files changed, 322 insertions(+), 265 deletions(-)

diff --git a/third_party/xla/xla/mlir/tools/mlir_bisect/BUILD b/third_party/xla/xla/mlir/tools/mlir_bisect/BUILD
index c394e91ad24568..59fd7def032da0 100644
--- a/third_party/xla/xla/mlir/tools/mlir_bisect/BUILD
+++ b/third_party/xla/xla/mlir/tools/mlir_bisect/BUILD
@@ -32,7 +32,6 @@ xla_cc_binary(
         "@llvm-project//mlir:AllPassesAndDialects",
         "@llvm-project//mlir:ArithDialect",
         "@llvm-project//mlir:IR",
-        "@llvm-project//mlir:ParseUtilities",
         "@llvm-project//mlir:Parser",
         "@llvm-project//mlir:Pass",
         "@llvm-project//mlir:Support",
diff --git a/third_party/xla/xla/mlir/tools/mlir_bisect/bisect_lib.cc b/third_party/xla/xla/mlir/tools/mlir_bisect/bisect_lib.cc
index c76e226ff7a6fa..dbac610da22e30 100644
--- a/third_party/xla/xla/mlir/tools/mlir_bisect/bisect_lib.cc
+++ b/third_party/xla/xla/mlir/tools/mlir_bisect/bisect_lib.cc
@@ -76,8 +76,8 @@ CandidateVector GetCandidates(
     BisectState& state, ModuleOp op) {
   assert(strategy && "GetCandidates was passed a null strategy");
   CandidateVector result;
-  op.lookupSymbol("main")->walk([&](Operation* subOp) {
-    llvm::move(strategy(state, subOp), std::back_inserter(result));
+  op.lookupSymbol("main")->walk([&](Operation* sub_op) {
+    llvm::move(strategy(state, sub_op), std::back_inserter(result));
   });
   return result;
 }
diff --git a/third_party/xla/xla/mlir/tools/mlir_bisect/mlir_bisect.cc b/third_party/xla/xla/mlir/tools/mlir_bisect/mlir_bisect.cc
index ca1a699444b452..073513a7c56e9b 100644
--- a/third_party/xla/xla/mlir/tools/mlir_bisect/mlir_bisect.cc
+++ b/third_party/xla/xla/mlir/tools/mlir_bisect/mlir_bisect.cc
@@ -47,7 +47,6 @@ limitations under the License.
 #include "mlir/Support/FileUtilities.h"  // from @llvm-project
 #include "mlir/Support/LLVM.h"  // from @llvm-project
 #include "mlir/Support/LogicalResult.h"  // from @llvm-project
-#include "mlir/Tools/ParseUtilities.h"  // from @llvm-project
 #include "mlir/Transforms/Passes.h"  // from @llvm-project
 #include "xla/literal.h"
 #include "xla/mlir/tools/mlir_bisect/bisect_lib.h"
@@ -96,10 +95,10 @@ namespace mlir {
 namespace bisect {
 namespace {
 
-OwningOpRef<ModuleOp> ParseMlirInput(llvm::StringRef inputFilename,
+OwningOpRef<ModuleOp> ParseMlirInput(llvm::StringRef input_filename,
                                      MLIRContext* context) {
   std::string error_message;
-  auto file = mlir::openInputFile(inputFilename, &error_message);
+  auto file = mlir::openInputFile(input_filename, &error_message);
   if (!file) {
     llvm::errs() << error_message << "\n";
     return {};
diff --git a/third_party/xla/xla/mlir/tools/mlir_bisect/test_passes.cc b/third_party/xla/xla/mlir/tools/mlir_bisect/test_passes.cc
index a48b6a0b8cc2e9..507eb8a9f27d89 100644
--- a/third_party/xla/xla/mlir/tools/mlir_bisect/test_passes.cc
+++ b/third_party/xla/xla/mlir/tools/mlir_bisect/test_passes.cc
@@ -15,8 +15,15 @@ limitations under the License.
 
 #include "xla/mlir/tools/mlir_bisect/test_passes.h"
 
+#include <utility>
+
+#include "llvm/ADT/SmallVector.h"
 #include "mlir/Dialect/Linalg/IR/Linalg.h"  // from @llvm-project
+#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
 #include "mlir/Pass/Pass.h"  // from @llvm-project
+#include "mlir/Pass/PassRegistry.h"  // from @llvm-project
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+#include "mlir/Support/TypeID.h"  // from @llvm-project
 
 namespace mlir {
 namespace bisect {
diff --git a/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/BUILD b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/BUILD
index e6f5d71d0e19f6..53353594c97029 100644
--- a/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/BUILD
+++ b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/BUILD
@@ -55,13 +55,9 @@ cc_library(
     ],
     deps = [
         "//xla/mlir/tools/mlir_interpreter/framework",
-        "@com_google_absl//absl/strings",
         "@llvm-project//llvm:Support",
-        "@llvm-project//mlir:ArithDialect",
-        "@llvm-project//mlir:FuncDialect",
         "@llvm-project//mlir:IR",
         "@llvm-project//mlir:Support",
-        "@llvm-project//mlir:TensorDialect",
         "@llvm-project//mlir:ViewLikeInterface",
     ],
 )
diff --git a/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/arith.cc b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/arith.cc
index 17fd94c3019141..d48edff6f37634 100644
--- a/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/arith.cc
+++ b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/arith.cc
@@ -53,9 +53,9 @@ InterpreterValue Bitcast(InterpreterState&, arith::BitcastOp op,
   if (!shaped_ty) {
     return result.ExtractElement({});
   }
-  auto& outView = result.View();
-  outView.strides = BufferView::GetDefaultStrides(shaped_ty.getShape());
-  outView.sizes = llvm::to_vector(shaped_ty.getShape());
+  auto& out_view = result.View();
+  out_view.strides = BufferView::GetDefaultStrides(shaped_ty.getShape());
+  out_view.sizes = llvm::to_vector(shaped_ty.getShape());
   return result;
 }
 
@@ -68,11 +68,11 @@ InterpreterValue Constant(InterpreterState&, arith::ConstantOp constant) {
     if (shaped_ty) {
       auto values = cast<DenseElementsAttr>(constant.getValue()).getValues<T>();
       auto result = TensorOrMemref<T>::Empty(shaped_ty.getShape());
-      auto valueIt = values.begin();
+      auto value_it = values.begin();
       result.view.is_vector = isa<VectorType>(shaped_ty);
       for (const auto& index : result.view.Indices(true)) {
-        result.at(index) = *valueIt;
-        ++valueIt;
+        result.at(index) = *value_it;
+        ++value_it;
       }
       return {result};
     }
@@ -81,8 +81,8 @@ InterpreterValue Constant(InterpreterState&, arith::ConstantOp constant) {
     if (auto integer = value.dyn_cast<IntegerAttr>()) {
       return {static_cast<T>(integer.getInt())};
     }
-    if (auto floatValue = value.dyn_cast<FloatAttr>()) {
-      return {static_cast<T>(floatValue.getValueAsDouble())};
+    if (auto float_value = value.dyn_cast<FloatAttr>()) {
+      return {static_cast<T>(float_value.getValueAsDouble())};
     }
 
     llvm_unreachable("unsupported constant type");
@@ -230,10 +230,10 @@ InterpreterValue CmpF(InterpreterState&, arith::CmpFOp compare,
 
 InterpreterValue Select(InterpreterState& state, arith::SelectOp,
                         const InterpreterValue& cond,
-                        const InterpreterValue& trueValue,
-                        const InterpreterValue& falseValue) {
+                        const InterpreterValue& true_value,
+                        const InterpreterValue& false_value) {
   if (std::holds_alternative<bool>(cond.storage)) {
-    return std::get<bool>(cond.storage) ? trueValue : falseValue;
+    return std::get<bool>(cond.storage) ? true_value : false_value;
   }
 
   if (!cond.IsTensor() && !cond.View().is_vector) {
@@ -241,10 +241,10 @@ InterpreterValue Select(InterpreterState& state, arith::SelectOp,
     return {};
   }
 
-  auto ret = trueValue.Clone();
+  auto ret = true_value.Clone();
   for (const auto& index : cond.View().Indices(/*include_vector_dims=*/true)) {
     if (cond.ExtractElement(index).AsInt() == 0) {
-      ret.InsertElement(index, falseValue.ExtractElement(index));
+      ret.InsertElement(index, false_value.ExtractElement(index));
     }
   }
   return ret;
diff --git a/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/bufferization.cc b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/bufferization.cc
index 1561d54b4b31eb..6bd9239876e22c 100644
--- a/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/bufferization.cc
+++ b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/bufferization.cc
@@ -15,9 +15,13 @@ limitations under the License.
 
 #include "mlir/Dialect/Bufferization/IR/Bufferization.h"  // from @llvm-project
 
+// clang-format erroneously puts the Bufferization header above.
 #include <algorithm>  // NOLINT
+#include <cstdint>    // NOLINT
 #include <optional>   // NOLINT
 
+#include "mlir/IR/BuiltinTypeInterfaces.h"  // from @llvm-project
+#include "mlir/Support/LLVM.h"  // from @llvm-project
 #include "xla/mlir/tools/mlir_interpreter/dialects/util.h"
 #include "xla/mlir/tools/mlir_interpreter/framework/interpreter.h"
 #include "xla/mlir/tools/mlir_interpreter/framework/interpreter_value.h"
diff --git a/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/linalg.cc b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/linalg.cc
index f65cbaf4ef18de..2c55af769ac9ca 100644
--- a/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/linalg.cc
+++ b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/linalg.cc
@@ -100,16 +100,16 @@ llvm::SmallVector<InterpreterValue> Generic(
   auto output_maps =
       ArrayRef<AffineMap>(indexing_maps).drop_front(inputs.size());
   std::function<void(int64_t)> run;
-  run = [&](int64_t loopIndex) {
+  run = [&](int64_t loop_index) {
     // Abort recursion if we encountered some error previously.
     if (state.HasFailure()) {
       return;
     }
 
-    if (loopIndex < ranges.size()) {
-      for (int64_t index = 0; index < ranges[loopIndex]; ++index) {
-        ivs[loopIndex] = index;
-        run(loopIndex + 1);
+    if (loop_index < ranges.size()) {
+      for (int64_t index = 0; index < ranges[loop_index]; ++index) {
+        ivs[loop_index] = index;
+        run(loop_index + 1);
       }
     } else {
       llvm::SmallVector<InterpreterValue> bbargs;
@@ -118,9 +118,9 @@ llvm::SmallVector<InterpreterValue> Generic(
         auto indices = EvalAffineMap(map, ivs);
         bbargs.push_back(input.ExtractElement(indices));
       }
-      llvm::SmallVector<llvm::SmallVector<int64_t>> outputIndices;
+      llvm::SmallVector<llvm::SmallVector<int64_t>> output_indices;
       for (auto [output, map] : llvm::zip(outputs, output_maps)) {
-        auto& indices = outputIndices.emplace_back(EvalAffineMap(map, ivs));
+        auto& indices = output_indices.emplace_back(EvalAffineMap(map, ivs));
         bbargs.push_back(output.ExtractElement(indices));
       }
       // Evaluate region.
@@ -130,7 +130,7 @@ llvm::SmallVector<InterpreterValue> Generic(
       }
       // Insert yielded values in the outputs.
       for (auto [output, indices, yield] :
-           llvm::zip(outputs, outputIndices, yielded)) {
+           llvm::zip(outputs, output_indices, yielded)) {
         output.InsertElement(indices, yield);
       }
     }
@@ -192,12 +192,12 @@ llvm::SmallVector<InterpreterValue> Reduce(InterpreterState& state,
     for (auto& out : output) {
       args.push_back(out.ExtractElement(dst_index));
     }
-    auto newValues = Interpret(state, reduce.getRegion(), args);
+    auto new_values = Interpret(state, reduce.getRegion(), args);
     if (state.HasFailure()) {
       return {};
     }
 
-    for (auto [out, value] : llvm::zip(output, newValues)) {
+    for (auto [out, value] : llvm::zip(output, new_values)) {
       out.InsertElement(dst_index, value);
     }
   }
diff --git a/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/memref.cc b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/memref.cc
index 0fd938054c1b9b..64b28342f88707 100644
--- a/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/memref.cc
+++ b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/memref.cc
@@ -158,17 +158,17 @@ InterpreterValue Subview(InterpreterState& state, memref::SubViewOp subview,
 llvm::SmallVector<InterpreterValue> CollapseShape(
     InterpreterState& state, memref::CollapseShapeOp collapse,
     const InterpreterValue& memref) {
-  const BufferView& inputView = memref.View();
+  const BufferView& input_view = memref.View();
   InterpreterValue out = memref;
   auto& out_view = out.View();
   out_view.sizes.clear();
   out_view.strides.clear();
 
   for (const auto& group : collapse.getReassociationIndices()) {
-    if (auto stride = inputView.GetCollapsedStride(group)) {
+    if (auto stride = input_view.GetCollapsedStride(group)) {
       out_view.strides.push_back(*stride);
       int64_t& size = out_view.sizes.emplace_back(1);
-      for (int64_t dim : group) size *= inputView.sizes[dim];
+      for (int64_t dim : group) size *= input_view.sizes[dim];
     } else {
       state.AddFailure("cannot collapse dimensions without a common stride");
       return {};
@@ -186,9 +186,9 @@ InterpreterValue Cast(InterpreterState&, memref::CastOp,
 // TODO(jreiffers): Implement full expand_shape support.
 InterpreterValue ExpandShape(InterpreterState& state, memref::ExpandShapeOp op,
                              InterpreterValue memref) {
-  BufferView inputView = memref.View();
-  auto outTy = cast<MemRefType>(op->getResultTypes()[0]);
-  if (outTy.getNumDynamicDims() > 0) {
+  BufferView input_view = memref.View();
+  auto out_ty = cast<MemRefType>(op->getResultTypes()[0]);
+  if (out_ty.getNumDynamicDims() > 0) {
     state.AddFailure("dynamic dimensions unsupported.");
     return {};
   }
@@ -196,10 +196,10 @@ InterpreterValue ExpandShape(InterpreterState& state, memref::ExpandShapeOp op,
   InterpreterValue out = memref;
   auto& out_view = out.View();
   out_view.strides.clear();
-  out_view.sizes = llvm::to_vector(outTy.getShape());
+  out_view.sizes = llvm::to_vector(out_ty.getShape());
   int64_t dummy;
-  if (!getStridesAndOffset(outTy, out_view.strides, dummy).succeeded()) {
-    if (inputView.strides != BufferView::GetDefaultStrides(inputView.sizes)) {
+  if (!getStridesAndOffset(out_ty, out_view.strides, dummy).succeeded()) {
+    if (input_view.strides != BufferView::GetDefaultStrides(input_view.sizes)) {
       state.AddFailure("unsupported strides");
       return {};
     }
@@ -210,21 +210,21 @@ InterpreterValue ExpandShape(InterpreterState& state, memref::ExpandShapeOp op,
 }
 
 InterpreterValue GetGlobal(InterpreterState& state,
-                           memref::GetGlobalOp getGlobal) {
+                           memref::GetGlobalOp get_global) {
   auto global = llvm::cast<memref::GlobalOp>(
-      state.GetSymbols().lookup(getGlobal.getName()));
+      state.GetSymbols().lookup(get_global.getName()));
 
   auto value = global.getConstantInitValue();
   assert(value && "mutable globals are not implemented");
 
-  auto ty = getGlobal->getResultTypes()[0].cast<ShapedType>();
+  auto ty = cast<ShapedType>(get_global->getResultTypes()[0]);
   return DispatchScalarType(ty, [&](auto dummy) -> InterpreterValue {
     auto values = value.getValues<decltype(dummy)>();
     auto result = TensorOrMemref<decltype(dummy)>::Empty(ty.getShape());
-    auto valueIt = values.begin();
+    auto value_it = values.begin();
     for (const auto& index : result.view.Indices()) {
-      result.at(index) = *valueIt;
-      ++valueIt;
+      result.at(index) = *value_it;
+      ++value_it;
     }
     return {result};
   });
diff --git a/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/util.cc b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/util.cc
index 209b14dff76a87..e1f3dae982726e 100644
--- a/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/util.cc
+++ b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/util.cc
@@ -17,9 +17,9 @@ limitations under the License.
 
 #include <cassert>
 #include <cstdint>
-#include <variant>
 
 #include "llvm/ADT/ArrayRef.h"
+#include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/SmallVector.h"
 #include "mlir/IR/AffineExpr.h"  // from @llvm-project
 #include "mlir/IR/BuiltinTypeInterfaces.h"  // from @llvm-project
@@ -85,9 +85,9 @@ InterpreterValue ReshapeTensor(const InterpreterValue& in,
                                ArrayRef<int64_t> shape) {
   // This doesn't need a copy in many cases, but it's easier that way.
   auto out = in.TypedAlike(shape);
-  for (const auto& [inIndex, outIndex] :
+  for (const auto& [in_index, out_index] :
        llvm::zip(in.View().Indices(), out.View().Indices())) {
-    out.InsertElement(outIndex, in.ExtractElement(inIndex));
+    out.InsertElement(out_index, in.ExtractElement(in_index));
   }
   return out;
 }
diff --git a/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/vector.cc b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/vector.cc
index 0bfd5ddd17cd84..7305b7e1b4e649 100644
--- a/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/vector.cc
+++ b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/vector.cc
@@ -241,15 +241,15 @@ InterpreterValue Contract(InterpreterState&, vector::ContractionOp contraction,
   BufferView iter;
   contraction.getIterationBounds(iter.sizes);
   auto maps = contraction.getIndexingMapsArray();
-  auto resultTy = contraction->getResultTypes()[0];
-  auto shapedTy = resultTy.dyn_cast<ShapedType>();
+  auto result_ty = contraction->getResultTypes()[0];
+  auto shaped_ty = result_ty.dyn_cast<ShapedType>();
   auto result =
-      DispatchScalarType(resultTy, [&](auto dummy) -> InterpreterValue {
+      DispatchScalarType(result_ty, [&](auto dummy) -> InterpreterValue {
         using T = decltype(dummy);
         using TT = TensorOrMemref<T>;
         const auto& lhs_t = std::get<TT>(lhs.storage);
         const auto& rhs_t = std::get<TT>(rhs.storage);
-        auto result_value = shapedTy ? acc.Clone() : acc.AsUnitTensor();
+        auto result_value = shaped_ty ? acc.Clone() : acc.AsUnitTensor();
         auto& result = std::get<TT>(result_value.storage);
         auto combiner = *GetCombiner<T>(contraction.getKind());
         for (const auto& indices : iter.Indices()) {
@@ -257,14 +257,14 @@ InterpreterValue Contract(InterpreterState&, vector::ContractionOp contraction,
           auto rhs_indices = EvalAffineMap(maps[1], indices);
           auto result_indices = EvalAffineMap(maps[2], indices);
 
-          auto& resultItem = result.at(result_indices);
-          resultItem = combiner(resultItem,
-                                lhs_t.at(lhs_indices) * rhs_t.at(rhs_indices));
+          auto& result_item = result.at(result_indices);
+          result_item = combiner(result_item,
+                                 lhs_t.at(lhs_indices) * rhs_t.at(rhs_indices));
         }
         return result_value;
       });
 
-  return shapedTy ? result : result.ExtractElement({});
+  return shaped_ty ? result : result.ExtractElement({});
 }
 
 InterpreterValue CreateMask(InterpreterState&, vector::CreateMaskOp op,
@@ -276,13 +276,13 @@ InterpreterValue ExpandLoad(InterpreterState& state, vector::ExpandLoadOp,
                             InterpreterValue memref,
                             SmallVector<int64_t> offsets,
                             TensorOrMemref<bool> mask,
-                            const InterpreterValue& passThrough) {
+                            const InterpreterValue& passthrough) {
   if (memref.View().strides.back() != 1) {
     state.AddFailure("expected last dimension to be contiguous");
     return {};
   }
 
-  auto out = passThrough.Clone();
+  auto out = passthrough.Clone();
   for (int64_t i = 0, e = out.View().sizes[0]; i < e; ++i) {
     if (mask.at({i})) {
       out.InsertElement({i}, memref.ExtractElement(offsets));
@@ -336,8 +336,8 @@ InterpreterValue FlatTranspose(InterpreterState&,
   int64_t rows = transpose.getRows();
   int64_t cols = transpose.getColumns();
   for (int64_t i = 0; i < rows * cols; ++i) {
-    int64_t srcIndex = (i % cols) * rows + (i / cols);
-    out.InsertElement({i}, vector.ExtractElement({srcIndex}));
+    int64_t src_index = (i % cols) * rows + (i / cols);
+    out.InsertElement({i}, vector.ExtractElement({src_index}));
   }
   return out;
 }
@@ -364,20 +364,20 @@ InterpreterValue Gather(InterpreterState& state, vector::GatherOp op,
                         const InterpreterValue& src, ArrayRef<int64_t> offsets,
                         const InterpreterValue& indices,
                         const TensorOrMemref<bool>& mask,
-                        const InterpreterValue& passThrough) {
+                        const InterpreterValue& pass_through) {
   if (isa<MemRefType>(op->getOperandTypes()[0]) &&
       src.View().strides.back() != 1) {
     state.AddFailure("expected trailing dimension to be contiguous");
     return {};
   }
-  auto out = passThrough.Clone();
-  for (const auto& outIndex : out.View().Indices()) {
-    if (!mask.at(outIndex)) {
+  auto out = pass_through.Clone();
+  for (const auto& out_index : out.View().Indices()) {
+    if (!mask.at(out_index)) {
       continue;
     }
-    auto inIndex = llvm::to_vector(offsets);
-    inIndex[0] += indices.ExtractElement(outIndex).AsInt();
-    out.InsertElement(outIndex, src.ExtractElement(inIndex));
+    auto in_index = llvm::to_vector(offsets);
+    in_index[0] += indices.ExtractElement(out_index).AsInt();
+    out.InsertElement(out_index, src.ExtractElement(in_index));
   }
   return out;
 }
@@ -386,12 +386,13 @@ InterpreterValue Insert(InterpreterState& state, vector::InsertOp insert,
                         const InterpreterValue& src,
                         const InterpreterValue& dst) {
   auto result = dst.Clone();
-  auto resultSlice = result;
-  auto& resultSliceView = resultSlice.View();
+  auto result_slice = result;
+  auto& result_slice_view = result_slice.View();
   for (int64_t offset : insert.getStaticPosition()) {
-    state.CheckSuccess(resultSliceView.Slice(0, offset), "index out of bounds");
+    state.CheckSuccess(result_slice_view.Slice(0, offset),
+                       "index out of bounds");
   }
-  resultSlice.Fill([&](auto indices) { return src.ExtractElement(indices); });
+  result_slice.Fill([&](auto indices) { return src.ExtractElement(indices); });
   return result;
 }
 
@@ -463,13 +464,13 @@ InterpreterValue MaskedLoad(InterpreterState& state, vector::MaskedLoadOp,
                             const InterpreterValue& memref,
                             SmallVector<int64_t> offsets,
                             TensorOrMemref<bool> mask,
-                            const InterpreterValue& passThrough) {
+                            const InterpreterValue& passthrough) {
   if (memref.View().strides.back() != 1) {
     state.AddFailure("expected last dimension to be contiguous");
     return {};
   }
 
-  auto out = passThrough.Clone();
+  auto out = passthrough.Clone();
   for (int64_t i = 0, e = mask.view.sizes[0]; i < e; ++i) {
     if (mask.at(i)) {
       out.InsertElement(i, memref.ExtractElement(offsets));
@@ -499,47 +500,47 @@ InterpreterValue ReductionImpl(InterpreterState& state,
                                const InterpreterValue& v,
                                const InterpreterValue* acc,
                                vector::CombiningKind kind,
-                               SmallVector<int64_t> dims, Type elementType) {
+                               SmallVector<int64_t> dims, Type element_type) {
   llvm::sort(dims);
-  SmallVector<int64_t> keptDims;
-  SmallVector<int64_t> resultShape;
+  SmallVector<int64_t> kept_dims;
+  SmallVector<int64_t> result_shape;
   for (auto [dim, size] : llvm::enumerate(v.View().sizes)) {
     if (!llvm::is_contained(dims, dim)) {
-      keptDims.push_back(dim);
-      resultShape.push_back(size);
+      kept_dims.push_back(dim);
+      result_shape.push_back(size);
     }
   }
-  return DispatchScalarType(elementType, [&](auto dummy) -> InterpreterValue {
+  return DispatchScalarType(element_type, [&](auto dummy) -> InterpreterValue {
     using T = decltype(dummy);
     using TT = TensorOrMemref<T>;
 
     auto combiner = *GetCombiner<T>(kind);
 
-    TT result = acc ? std::get<TT>((keptDims.empty() ? acc->AsUnitTensor()
-                                                     : acc->Clone())
+    TT result = acc ? std::get<TT>((kept_dims.empty() ? acc->AsUnitTensor()
+                                                      : acc->Clone())
                                        .storage)
-                    : TensorOrMemref<T>::Empty(resultShape);
+                    : TensorOrMemref<T>::Empty(result_shape);
 
     for (const auto& result_index : result.view.Indices()) {
       auto src = std::get<TT>(v.storage);
       for (auto [dim, index] :
-           llvm::reverse(llvm::zip(keptDims, result_index))) {
+           llvm::reverse(llvm::zip(kept_dims, result_index))) {
         state.CheckSuccess(src.view.Slice(dim, index), "index out of bounds");
       }
 
       T& item = result.at(result_index);
       bool first = acc == nullptr;
-      for (const auto& srcIndex : src.view.Indices()) {
+      for (const auto& src_index : src.view.Indices()) {
         if (first) {
-          item = src.at(srcIndex);
+          item = src.at(src_index);
           first = false;
         } else {
-          item = combiner(item, src.at(srcIndex));
+          item = combiner(item, src.at(src_index));
         }
       }
     }
 
-    if (keptDims.empty()) {
+    if (kept_dims.empty()) {
       return {result.at({})};
     }
     return {result};
@@ -550,30 +551,30 @@ InterpreterValue MultiReduction(InterpreterState& state,
                                 vector::MultiDimReductionOp reduction,
                                 const InterpreterValue& source,
                                 const InterpreterValue& acc) {
-  auto elementTy = getElementTypeOrSelf(reduction->getResultTypes()[0]);
+  auto element_ty = getElementTypeOrSelf(reduction->getResultTypes()[0]);
   return {ReductionImpl(state, source, &acc, reduction.getKind(),
                         ExtractVector<int64_t>(reduction.getReductionDims()),
-                        elementTy)};
+                        element_ty)};
 }
 
 InterpreterValue OuterProduct(InterpreterState&,
-                              vector::OuterProductOp outerproduct,
+                              vector::OuterProductOp outer_product,
                               const InterpreterValue& lhs,
                               const InterpreterValue& rhs,
                               std::optional<InterpreterValue> acc) {
-  ShapedType ty = cast<ShapedType>(outerproduct->getResultTypes()[0]);
+  ShapedType ty = cast<ShapedType>(outer_product->getResultTypes()[0]);
   return DispatchScalarType(ty, [&](auto dummy) -> InterpreterValue {
     using T = decltype(dummy);
     using TT = TensorOrMemref<T>;
     const TT& lhs_t = std::get<TT>(lhs.storage);
 
-    auto combiner = GetCombiner<T>(outerproduct.getKind());
+    auto combiner = GetCombiner<T>(outer_product.getKind());
     auto result =
         acc ? std::get<TT>(acc->storage).Clone() : TT::Empty(ty.getShape());
     if (std::holds_alternative<T>(rhs.storage)) {
-      T rhsS = std::get<T>(rhs.storage);
+      T rhs_scalar = std::get<T>(rhs.storage);
       for (int64_t i : llvm::seq(int64_t{0}, lhs_t.view.sizes[0])) {
-        result.at(i) = combiner(result.at(i), lhs_t.at(i) * rhsS);
+        result.at(i) = combiner(result.at(i), lhs_t.at(i) * rhs_scalar);
       }
     } else {
       const TT& rhs_t = std::get<TT>(rhs.storage);
@@ -619,10 +620,10 @@ InterpreterValue Reduction(InterpreterState& state,
 InterpreterValue ShapeCast(InterpreterState&, vector::ShapeCastOp op,
                            const InterpreterValue& in) {
   auto out = in.CoerceLayout({});
-  auto& outView = out.View();
-  outView.sizes =
+  auto& out_view = out.View();
+  out_view.sizes =
       llvm::to_vector(op->getResultTypes()[0].cast<ShapedType>().getShape());
-  outView.strides = BufferView::GetDefaultStrides(outView.sizes);
+  out_view.strides = BufferView::GetDefaultStrides(out_view.sizes);
   return out;
 }
 
@@ -634,17 +635,18 @@ InterpreterValue Shuffle(InterpreterState& state, vector::ShuffleOp shuffle,
   result_view.is_vector = true;
 
   auto mask = ExtractVector<int64_t>(shuffle.getMask());
-  bool isZeroDim = v0.View().Rank() == 0;
-  int64_t size0 = isZeroDim ? 1 : v0.View().sizes[0];
-  for (auto [dstIndex, srcIndex] : llvm::enumerate(mask)) {
-    auto src = srcIndex < size0 ? v0 : v1;
-    if (!isZeroDim) {
+  bool is_zero_dim = v0.View().Rank() == 0;
+  int64_t size0 = is_zero_dim ? 1 : v0.View().sizes[0];
+  for (auto [dst_index, src_index] : llvm::enumerate(mask)) {
+    auto src = src_index < size0 ? v0 : v1;
+    if (!is_zero_dim) {
       state.CheckSuccess(
-          src.View().Slice(0, srcIndex < size0 ? srcIndex : srcIndex - size0),
+          src.View().Slice(0,
+                           src_index < size0 ? src_index : src_index - size0),
           "index out of bounds");
     }
     auto dst = result;
-    state.CheckSuccess(dst.View().Slice(0, dstIndex), "index out of bounds");
+    state.CheckSuccess(dst.View().Slice(0, dst_index), "index out of bounds");
     dst.Fill([&](auto indices) { return src.ExtractElement(indices); });
   }
   return result;
@@ -667,17 +669,17 @@ void Store(InterpreterState& state, vector::StoreOp,
     state.AddFailure("array index out of bounds");
     return;
   }
-  const auto& outView = dst.View();
-  if (outView.num_vector_dims > 0) {
+  const auto& out_view = dst.View();
+  if (out_view.num_vector_dims > 0) {
     dst.InsertElement(offsets, src);
   } else {
-    for (const auto& srcIndex : src.View().Indices()) {
-      auto dstIndex = srcIndex;
-      for (int64_t i = 0; i < dstIndex.size(); ++i) {
-        dstIndex[i] += offsets[i];
+    for (const auto& src_index : src.View().Indices()) {
+      auto dst_index = src_index;
+      for (int64_t i = 0; i < dst_index.size(); ++i) {
+        dst_index[i] += offsets[i];
       }
-      if (outView.InBounds(dstIndex)) {
-        dst.InsertElement(dstIndex, src.ExtractElement(srcIndex));
+      if (out_view.InBounds(dst_index)) {
+        dst.InsertElement(dst_index, src.ExtractElement(src_index));
       }
     }
   }
@@ -686,43 +688,43 @@ void Store(InterpreterState& state, vector::StoreOp,
 std::optional<InterpreterValue> ExtractMemorySlice(
     InterpreterState& state, const AffineMap& map,
     const InterpreterValue& memory, const InterpreterValue& vector,
-    ArrayRef<int64_t> offsets, std::optional<ArrayAttr> in_boundsAttr) {
+    ArrayRef<int64_t> offsets, std::optional<ArrayAttr> in_bounds_attr) {
   llvm::SmallVector<bool> in_bounds(offsets.size());
-  if (in_boundsAttr) {
-    llvm::copy(in_boundsAttr->getAsValueRange<BoolAttr>(),
-               in_bounds.end() - in_boundsAttr->size());
+  if (in_bounds_attr) {
+    llvm::copy(in_bounds_attr->getAsValueRange<BoolAttr>(),
+               in_bounds.end() - in_bounds_attr->size());
   }
 
-  auto memSlice = memory;
-  auto& memSliceView = memSlice.View();
-  auto& vectorView = vector.View();
-  for (int64_t i = 0; i < memSliceView.Rank(); ++i) {
+  auto mem_slice = memory;
+  auto& mem_slice_view = mem_slice.View();
+  auto& vector_view = vector.View();
+  for (int64_t i = 0; i < mem_slice_view.Rank(); ++i) {
     bool found = false;
-    for (int64_t j = 0; !found && j < vectorView.Rank(); ++j) {
+    for (int64_t j = 0; !found && j < vector_view.Rank(); ++j) {
       if (map.getResult(j).isFunctionOfDim(i)) {
-        int64_t size = memSliceView.sizes[i] - offsets[i];
-        bool isInBounds = size >= vectorView.sizes[j];
-        if (!isInBounds && in_bounds[i]) {
+        int64_t size = mem_slice_view.sizes[i] - offsets[i];
+        bool is_in_bounds = size >= vector_view.sizes[j];
+        if (!is_in_bounds && in_bounds[i]) {
           state.AddFailure("index out of bounds");
           return std::nullopt;
         }
-        (void)memSliceView.Slice(
+        (void)mem_slice_view.Slice(
             i, offsets[i],
-            std::max(int64_t{0}, std::min(vectorView.sizes[j], size)));
+            std::max(int64_t{0}, std::min(vector_view.sizes[j], size)));
         found = true;
       }
     }
     if (!found) {
-      bool isInBounds = memSliceView.sizes[i] > offsets[i];
-      if (!isInBounds) {
+      bool is_in_bounds = mem_slice_view.sizes[i] > offsets[i];
+      if (!is_in_bounds) {
         state.AddFailure("index out of bounds");
         return std::nullopt;
       }
 
-      (void)memSliceView.Slice(i, offsets[i], isInBounds ? 1 : 0);
+      (void)mem_slice_view.Slice(i, offsets[i], is_in_bounds ? 1 : 0);
     }
   }
-  return memSlice;
+  return mem_slice;
 }
 
 InterpreterValue TransferRead(InterpreterState& state,
@@ -731,46 +733,46 @@ InterpreterValue TransferRead(InterpreterState& state,
                               ArrayRef<int64_t> offsets,
                               const InterpreterValue& padding,
                               std::optional<TensorOrMemref<bool>> mask) {
-  auto* maskChannel = state.GetTopScope()->GetSideChannel<MaskSideChannel>(
+  auto* mask_channel = state.GetTopScope()->GetSideChannel<MaskSideChannel>(
       /*optional=*/true);
-  if (maskChannel) {
+  if (mask_channel) {
     if (mask) {
       state.AddFailure(
           "vector.mask and transfer_read with mask should not be used "
           "simultaneously");
       return {};
     }
-    mask = maskChannel->GetMask();
+    mask = mask_channel->GetMask();
   }
 
   InterpreterValue dst = src.TypedAlike(transfer.getVectorType().getShape());
-  if (maskChannel && maskChannel->GetPassthrough()) {
+  if (mask_channel && mask_channel->GetPassthrough()) {
     dst.Fill([&](auto indices) {
       if (mask->at(indices)) {
         return padding;
       }
-      return maskChannel->GetPassthrough()->ExtractElement(indices);
+      return mask_channel->GetPassthrough()->ExtractElement(indices);
     });
   } else {
     dst.Fill([&](auto) { return padding; });
   }
   dst.View().is_vector = true;
 
-  auto srcSlice = ExtractMemorySlice(state, transfer.getPermutationMap(), src,
-                                     dst, offsets, transfer.getInBounds());
+  auto src_slice = ExtractMemorySlice(state, transfer.getPermutationMap(), src,
+                                      dst, offsets, transfer.getInBounds());
 
-  if (!srcSlice) {
+  if (!src_slice) {
     return {};
   }
-  for (const auto& src_indices : srcSlice->View().Indices()) {
-    SmallVector<int64_t> dstIndices =
+  for (const auto& src_indices : src_slice->View().Indices()) {
+    SmallVector<int64_t> dst_indices =
         EvalAffineMap(transfer.getPermutationMap(), src_indices);
 
     // Note: the handling of padding and passthrough values is somewhat
     // arbitrary here. At the time of writing this, there seems to be little
     // evidence of actual usage of this feature.
-    if (!mask || mask->at(dstIndices)) {
-      dst.InsertElement(dstIndices, srcSlice->ExtractElement(src_indices));
+    if (!mask || mask->at(dst_indices)) {
+      dst.InsertElement(dst_indices, src_slice->ExtractElement(src_indices));
     }
   }
 
@@ -781,7 +783,7 @@ llvm::SmallVector<InterpreterValue> TransferWrite(
     InterpreterState& state, vector::TransferWriteOp transfer,
     InterpreterValue src, InterpreterValue dst, ArrayRef<int64_t> offsets,
     std::optional<TensorOrMemref<bool>> mask) {
-  if (auto* maskChannel = state.GetTopScope()->GetSideChannel<MaskSideChannel>(
+  if (auto* mask_channel = state.GetTopScope()->GetSideChannel<MaskSideChannel>(
           /*optional=*/true)) {
     if (mask) {
       state.AddFailure(
@@ -789,13 +791,13 @@ llvm::SmallVector<InterpreterValue> TransferWrite(
           "simultaneously");
       return {};
     }
-    if (maskChannel->GetPassthrough()) {
+    if (mask_channel->GetPassthrough()) {
       state.AddFailure(
           "vector.mask with passthrough should not be used with "
           "transfer_write");
       return {};
     }
-    mask = maskChannel->GetMask();
+    mask = mask_channel->GetMask();
   }
 
   const auto& src_view = src.View();
@@ -809,11 +811,11 @@ llvm::SmallVector<InterpreterValue> TransferWrite(
     return {};
   }
 
-  for (const auto& dstIndices : dst_slice->View().Indices()) {
+  for (const auto& dst_indices : dst_slice->View().Indices()) {
     SmallVector<int64_t> src_indices =
-        EvalAffineMap(transfer.getPermutationMap(), dstIndices);
+        EvalAffineMap(transfer.getPermutationMap(), dst_indices);
     if (src_view.InBounds(src_indices) && (!mask || mask->at(src_indices))) {
-      dst_slice->InsertElement(dstIndices, src.ExtractElement(src_indices));
+      dst_slice->InsertElement(dst_indices, src.ExtractElement(src_indices));
     }
   }
 
diff --git a/third_party/xla/xla/mlir/tools/mlir_interpreter/framework/interpreter.cc b/third_party/xla/xla/mlir/tools/mlir_interpreter/framework/interpreter.cc
index a07322c17a5f86..21de4ea6ffc7cf 100644
--- a/third_party/xla/xla/mlir/tools/mlir_interpreter/framework/interpreter.cc
+++ b/third_party/xla/xla/mlir/tools/mlir_interpreter/framework/interpreter.cc
@@ -30,6 +30,8 @@ limitations under the License.
 #include "llvm/Support/raw_ostream.h"
 #include "mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h"  // from @llvm-project
 #include "mlir/IR/OpDefinition.h"  // from @llvm-project
+#include "mlir/IR/Operation.h"  // from @llvm-project
+#include "mlir/IR/Region.h"  // from @llvm-project
 #include "mlir/IR/SymbolTable.h"  // from @llvm-project
 #include "mlir/Support/LLVM.h"  // from @llvm-project
 #include "xla/mlir/tools/mlir_interpreter/framework/interpreter_value.h"
@@ -78,13 +80,13 @@ SmallVector<InterpreterValue> Interpret(InterpreterState& state, Region& region,
     scope.Set(value, interpreter_value);
   }
 
-  std::optional<SmallVector<InterpreterValue>> blockResults;
+  std::optional<SmallVector<InterpreterValue>> block_results;
   for (mlir::Operation& op : block) {
     auto results = Interpret(state, op);
     if (state.HasFailure()) return {};
     if (op.hasTrait<OpTrait::IsTerminator>()) {
-      assert(!blockResults.has_value() && "Expected at most one terminator");
-      blockResults = results;
+      assert(!block_results.has_value() && "Expected at most one terminator");
+      block_results = results;
     } else {
       if (results.size() != op.getNumResults()) {
         llvm::errs() << "Unexpected number of results while interpreting "
@@ -96,19 +98,19 @@ SmallVector<InterpreterValue> Interpret(InterpreterState& state, Region& region,
       }
     }
   }
-  if (!blockResults) {
-    blockResults = SmallVector<InterpreterValue>{};
+  if (!block_results) {
+    block_results = SmallVector<InterpreterValue>{};
   }
-  state.GetOptions().listener->LeaveRegion(*blockResults);
-  return *std::move(blockResults);
+  state.GetOptions().listener->LeaveRegion(*block_results);
+  return *std::move(block_results);
 }
 
 InterpreterState::InterpreterState(const mlir::SymbolTable& symbols,
                                    InterpreterOptions options)
     : symbols_(symbols), options_(options) {
   if (!options_.listener) {
-    static auto& noOpListener = *new InterpreterListener();
-    this->options_.listener = &noOpListener;
+    static auto& no_op_listener = *new InterpreterListener();
+    this->options_.listener = &no_op_listener;
   }
   if (options_.max_steps) {
     remaining_steps_ = *options_.max_steps;
diff --git a/third_party/xla/xla/mlir/tools/mlir_interpreter/framework/interpreter.h b/third_party/xla/xla/mlir/tools/mlir_interpreter/framework/interpreter.h
index 7f14f6df435434..8251906f44b08f 100644
--- a/third_party/xla/xla/mlir/tools/mlir_interpreter/framework/interpreter.h
+++ b/third_party/xla/xla/mlir/tools/mlir_interpreter/framework/interpreter.h
@@ -40,7 +40,7 @@ class InterpreterListener {
   virtual void AfterOp(ArrayRef<InterpreterValue> results) {}
   virtual void EnterRegion(ArrayRef<InterpreterValue> args,
                            mlir::Region& region) {}
-  virtual void LeaveRegion(ArrayRef<InterpreterValue> terminatorArgs) {}
+  virtual void LeaveRegion(ArrayRef<InterpreterValue> terminator_args) {}
 };
 
 struct InterpreterStats {
diff --git a/third_party/xla/xla/mlir/tools/mlir_interpreter/framework/interpreter_value.cc b/third_party/xla/xla/mlir/tools/mlir_interpreter/framework/interpreter_value.cc
index 55f9a84856f2bd..cd5c2ac4190308 100644
--- a/third_party/xla/xla/mlir/tools/mlir_interpreter/framework/interpreter_value.cc
+++ b/third_party/xla/xla/mlir/tools/mlir_interpreter/framework/interpreter_value.cc
@@ -15,6 +15,7 @@ limitations under the License.
 
 #include "xla/mlir/tools/mlir_interpreter/framework/interpreter_value.h"
 
+#include <cassert>
 #include <complex>
 #include <cstdint>
 #include <functional>
@@ -22,9 +23,15 @@ limitations under the License.
 #include <memory>
 #include <string>
 #include <string_view>
+#include <type_traits>
 #include <variant>
 
+#include "llvm/ADT/STLExtras.h"
+#include "llvm/Support/Casting.h"
 #include "llvm/Support/ErrorHandling.h"
+#include "llvm/Support/raw_ostream.h"
+#include "mlir/IR/BuiltinTypes.h"  // from @llvm-project
+#include "mlir/IR/Types.h"  // from @llvm-project
 #include "mlir/Support/LLVM.h"  // from @llvm-project
 #include "xla/mlir/tools/mlir_interpreter/framework/tensor_or_memref.h"
 
@@ -257,16 +264,16 @@ InterpreterValue InterpreterValue::TypedAlike(
       storage);
 }
 
-InterpreterValue InterpreterValue::MakeTensor(mlir::Type elementType,
+InterpreterValue InterpreterValue::MakeTensor(mlir::Type element_type,
                                               SmallVector<int64_t> shape) {
-  auto vectorTy = llvm::dyn_cast<VectorType>(elementType);
-  if (vectorTy) {
-    llvm::copy(vectorTy.getShape(), std::back_inserter(shape));
+  auto vector_ty = llvm::dyn_cast<VectorType>(element_type);
+  if (vector_ty) {
+    llvm::copy(vector_ty.getShape(), std::back_inserter(shape));
   }
-  return DispatchScalarType(elementType, [&](auto dummy) -> InterpreterValue {
+  return DispatchScalarType(element_type, [&](auto dummy) -> InterpreterValue {
     auto tensor = TensorOrMemref<decltype(dummy)>::Empty(shape);
-    if (vectorTy) {
-      tensor.view.num_vector_dims = vectorTy.getRank();
+    if (vector_ty) {
+      tensor.view.num_vector_dims = vector_ty.getRank();
     }
     return {tensor};
   });
diff --git a/third_party/xla/xla/mlir/tools/mlir_interpreter/framework/interpreter_value.h b/third_party/xla/xla/mlir/tools/mlir_interpreter/framework/interpreter_value.h
index 4c0e4788f499c4..90fb8e74eceb33 100644
--- a/third_party/xla/xla/mlir/tools/mlir_interpreter/framework/interpreter_value.h
+++ b/third_party/xla/xla/mlir/tools/mlir_interpreter/framework/interpreter_value.h
@@ -92,7 +92,7 @@ struct InterpreterValue {
   // Creates a tensor with the given element type and shape. `element_type` may
   // be a vector type, in which case the shape only specifies the non-vector
   // dimensions.
-  static InterpreterValue MakeTensor(mlir::Type elementType,
+  static InterpreterValue MakeTensor(mlir::Type element_type,
                                      SmallVector<int64_t> shape);
 
   // Returns the underlying tensor's view. Must be a tensor.
@@ -135,7 +135,7 @@ constexpr static bool is_valid_interpreter_value_v =  // NOLINT
 // NOTE: When casting to an unsigned type, this behaves differently than
 // InterpreterValue::AsUint. That function preserves the content's bit width,
 // so InterpreterValueDynCast<uint64_t>({int8_t{-1}}) will return 2^64-1,
-// whereas asUInt will return 255.
+// whereas AsUInt will return 255.
 template <typename T>
 std::optional<T> InterpreterValueDynCast(InterpreterValue v) {
   if constexpr (std::is_same_v<T, InterpreterValue>) {
diff --git a/third_party/xla/xla/mlir/tools/mlir_interpreter/framework/interpreter_value_util.h b/third_party/xla/xla/mlir/tools/mlir_interpreter/framework/interpreter_value_util.h
index 5e9a6ba2903359..12f508dcd5fff3 100644
--- a/third_party/xla/xla/mlir/tools/mlir_interpreter/framework/interpreter_value_util.h
+++ b/third_party/xla/xla/mlir/tools/mlir_interpreter/framework/interpreter_value_util.h
@@ -82,13 +82,13 @@ struct InterpreterValueBiMapVisitor {
   const InterpreterValue& rhs;
 
   template <typename T>
-  InterpreterValue operator()(const TensorOrMemref<T>& lhsT) {
+  InterpreterValue operator()(const TensorOrMemref<T>& lhs_t) {
     if constexpr (Fn::template SupportedType<T>()) {
       using OutElemT = decltype(Fn::Apply(T(), T()));
-      auto out = TensorOrMemref<OutElemT>::EmptyLike(lhsT.view);
-      const auto& rhsT = std::get<TensorOrMemref<T>>(rhs.storage);
+      auto out = TensorOrMemref<OutElemT>::EmptyLike(lhs_t.view);
+      const auto& rhs_t = std::get<TensorOrMemref<T>>(rhs.storage);
       for (const auto& index : out.view.Indices(true)) {
-        out.at(index) = Fn::Apply(lhsT.at(index), rhsT.at(index));
+        out.at(index) = Fn::Apply(lhs_t.at(index), rhs_t.at(index));
       }
       return {out};
     } else {
@@ -98,10 +98,10 @@ struct InterpreterValueBiMapVisitor {
     }
   }
 
-  InterpreterValue operator()(const Tuple& lhsT) {
-    const auto& rhsT = std::get<Tuple>(rhs.storage);
+  InterpreterValue operator()(const Tuple& lhs_t) {
+    const auto& rhs_t = std::get<Tuple>(rhs.storage);
     Tuple out;
-    for (const auto& [lhs_v, rhs_v] : llvm::zip(lhsT.values, rhsT.values)) {
+    for (const auto& [lhs_v, rhs_v] : llvm::zip(lhs_t.values, rhs_t.values)) {
       out.values.push_back(std::make_unique<InterpreterValue>(std::move(
           std::visit(InterpreterValueBiMapVisitor{*rhs_v}, lhs_v->storage))));
     }
@@ -130,11 +130,11 @@ template <bool allow_bools, bool allow_ints, bool allow_floats,
 struct FilterMapTraits {
   template <typename T>
   static constexpr bool SupportedType() {
-    constexpr bool isBool = std::is_same_v<T, bool>;
-    constexpr bool isInt = std::is_integral_v<T> && !isBool;
-    constexpr bool isUnsigned = std::is_unsigned_v<T>;
-    return (allow_bools && isBool) ||
-           (allow_ints && isInt && (allow_unsigned || !isUnsigned)) ||
+    constexpr bool is_bool = std::is_same_v<T, bool>;
+    constexpr bool is_int = std::is_integral_v<T> && !is_bool;
+    constexpr bool is_unsigned = std::is_unsigned_v<T>;
+    return (allow_bools && is_bool) ||
+           (allow_ints && is_int && (allow_unsigned || !is_unsigned)) ||
            (allow_floats && std::is_floating_point_v<T>) ||
            (allow_complex && is_complex_v<T>);
   }
diff --git a/third_party/xla/xla/mlir/tools/mlir_interpreter/framework/registration.cc b/third_party/xla/xla/mlir/tools/mlir_interpreter/framework/registration.cc
index 454f08641e19e6..93ec63bbe06d53 100644
--- a/third_party/xla/xla/mlir/tools/mlir_interpreter/framework/registration.cc
+++ b/third_party/xla/xla/mlir/tools/mlir_interpreter/framework/registration.cc
@@ -15,10 +15,14 @@ limitations under the License.
 
 #include "xla/mlir/tools/mlir_interpreter/framework/registration.h"
 
+#include <cassert>
 #include <functional>
 #include <utility>
 
+#include "mlir/IR/Operation.h"  // from @llvm-project
 #include "mlir/Support/LLVM.h"  // from @llvm-project
+#include "xla/mlir/tools/mlir_interpreter/framework/interpreter.h"
+#include "xla/mlir/tools/mlir_interpreter/framework/interpreter_value.h"
 
 namespace mlir {
 namespace interpreter {
diff --git a/third_party/xla/xla/mlir/tools/mlir_interpreter/framework/tensor_or_memref.cc b/third_party/xla/xla/mlir/tools/mlir_interpreter/framework/tensor_or_memref.cc
index 7226ff73b9b14f..d66415e045572c 100644
--- a/third_party/xla/xla/mlir/tools/mlir_interpreter/framework/tensor_or_memref.cc
+++ b/third_party/xla/xla/mlir/tools/mlir_interpreter/framework/tensor_or_memref.cc
@@ -15,10 +15,8 @@ limitations under the License.
 
 #include "xla/mlir/tools/mlir_interpreter/framework/tensor_or_memref.h"
 
-#include <algorithm>
+#include <cstddef>
 #include <cstdint>
-#include <iostream>
-#include <numeric>
 #include <optional>
 #include <utility>
 
@@ -44,9 +42,13 @@ std::optional<int64_t> BufferView::GetPhysicalIndex(
 }
 
 bool BufferView::InBounds(llvm::ArrayRef<int64_t> view_indices) const {
-  if (view_indices.size() > sizes.size()) return false;
+  if (view_indices.size() > sizes.size()) {
+    return false;
+  }
   for (auto [index, size] : llvm::zip(view_indices, sizes)) {
-    if (index < 0 || index >= size) return false;
+    if (index < 0 || index >= size) {
+      return false;
+    }
   }
   return true;
 }
@@ -64,12 +66,12 @@ SmallVector<int64_t> BufferView::GetDefaultStrides(ArrayRef<int64_t> sizes) {
 SmallVector<int64_t> BufferView::GetStridesForLayout(ArrayRef<int64_t> sizes,
                                                      ArrayRef<int64_t> layout) {
   if (layout.empty()) return GetDefaultStrides(sizes);
-  auto inverseLayout = invertPermutationVector(layout);
+  auto inverse_layout = invertPermutationVector(layout);
   SmallVector<int64_t> result(sizes.size());
   int64_t stride = 1;
   for (int64_t i = 0; i < layout.size(); ++i) {
-    result[inverseLayout[i]] = stride;
-    stride *= sizes[inverseLayout[i]];
+    result[inverse_layout[i]] = stride;
+    stride *= sizes[inverse_layout[i]];
   }
   return result;
 }
@@ -77,8 +79,8 @@ SmallVector<int64_t> BufferView::GetStridesForLayout(ArrayRef<int64_t> sizes,
 LogicalResult BufferView::Slice(int64_t dim_index, int64_t dim_offset) {
   llvm::SmallVector<int64_t> offsets(Rank(), 0);
   offsets[dim_index] = dim_offset;
-  if (auto newOffset = GetPhysicalIndex(offsets)) {
-    offset = *newOffset;
+  if (auto new_offset = GetPhysicalIndex(offsets)) {
+    offset = *new_offset;
   } else {
     return failure();
   }
@@ -94,8 +96,8 @@ LogicalResult BufferView::Slice(int64_t dim_index, int64_t dim_offset,
   offsets[dim_index] = dim_offset;
   if (dim_size == 0) {
     offset = 0;
-  } else if (auto newOffset = GetPhysicalIndex(offsets)) {
-    offset = *newOffset;
+  } else if (auto new_offset = GetPhysicalIndex(offsets)) {
+    offset = *new_offset;
   } else {
     return failure();
   }
@@ -115,9 +117,9 @@ LogicalResult BufferView::Subview(ArrayRef<int64_t> subview_offsets,
 
   for (auto [in_size, subview_offset, subview_size, subview_stride] :
        llvm::zip(sizes, subview_offsets, subview_sizes, subview_strides)) {
-    int64_t limitIndex = subview_offset + (subview_size - 1) * subview_stride;
-    if (subview_offset < 0 || subview_offset >= in_size || limitIndex < 0 ||
-        limitIndex >= in_size) {
+    int64_t limit_index = subview_offset + (subview_size - 1) * subview_stride;
+    if (subview_offset < 0 || subview_offset >= in_size || limit_index < 0 ||
+        limit_index >= in_size) {
       return failure();
     }
   }
@@ -129,33 +131,38 @@ LogicalResult BufferView::Subview(ArrayRef<int64_t> subview_offsets,
   return success();
 }
 
-int64_t BufferView::GetNumElements(bool includeVectorDims) const {
+int64_t BufferView::GetNumElements(bool include_vector_dims) const {
   size_t n = 1;
   for (auto size : ArrayRef<int64_t>(sizes).drop_back(
-           includeVectorDims ? 0 : num_vector_dims.value_or(0)))
+           include_vector_dims ? 0 : num_vector_dims.value_or(0))) {
     n *= size;
+  }
   return n;
 }
 
 std::optional<int64_t> BufferView::GetCollapsedStride(
     llvm::ArrayRef<int64_t> dims) const {
   using StrideAndDim = std::pair<int64_t, int64_t>;
-  llvm::SmallVector<StrideAndDim> stridesAndDims;
+  llvm::SmallVector<StrideAndDim> strides_and_dims;
   for (auto dim : dims) {
     if (sizes[dim] != 1) {
-      stridesAndDims.emplace_back(strides[dim], dim);
+      strides_and_dims.emplace_back(strides[dim], dim);
     }
   }
 
-  if (stridesAndDims.empty()) return 0;
+  if (strides_and_dims.empty()) {
+    return 0;
+  }
 
-  llvm::sort(stridesAndDims);
-  int64_t nextStride = stridesAndDims.front().first;
-  for (auto [stride, dim] : stridesAndDims) {
-    if (stride != nextStride) return std::nullopt;
-    nextStride *= sizes[dim];
+  llvm::sort(strides_and_dims);
+  int64_t next_stride = strides_and_dims.front().first;
+  for (auto [stride, dim] : strides_and_dims) {
+    if (stride != next_stride) {
+      return std::nullopt;
+    }
+    next_stride *= sizes[dim];
   }
-  return stridesAndDims.front().first;
+  return strides_and_dims.front().first;
 }
 
 }  // namespace interpreter
diff --git a/third_party/xla/xla/mlir/tools/mlir_interpreter/framework/tensor_or_memref.h b/third_party/xla/xla/mlir/tools/mlir_interpreter/framework/tensor_or_memref.h
index 6f801021a68101..35afc96ce1b79c 100644
--- a/third_party/xla/xla/mlir/tools/mlir_interpreter/framework/tensor_or_memref.h
+++ b/third_party/xla/xla/mlir/tools/mlir_interpreter/framework/tensor_or_memref.h
@@ -76,14 +76,14 @@ struct BufferView {
   int64_t Rank() const { return sizes.size() - num_vector_dims.value_or(0); }
 
   // Removes the dimension from the view. If you need to keep it, use the
-  // overload below with dimSize = 1.
+  // overload below with dim_size = 1.
   LogicalResult Slice(int64_t dim_index, int64_t dim_offset);
   LogicalResult Slice(int64_t dim_index, int64_t dim_offset, int64_t dim_size,
                       int64_t dim_stride = 1);
   LogicalResult Subview(ArrayRef<int64_t> subview_offsets,
                         ArrayRef<int64_t> subview_sizes,
                         ArrayRef<int64_t> subview_strides);
-  int64_t GetNumElements(bool include_vector_dims = false) const;
+  int64_t GetNumElements(bool include_vector_dimsms = false) const;
 
   class LogicalIndexView {
    public:
@@ -103,18 +103,19 @@ struct BufferView {
       }
 
       Iterator& operator++() {
-        auto indexIt = view_indices_.rbegin();
-        auto sizeIt = view_->sizes.rbegin();
+        auto index_it = view_indices_.rbegin();
+        auto size_it = view_->sizes.rbegin();
         if (!include_vector_dims_) {
-          std::advance(sizeIt, view_->num_vector_dims.value_or(0));
+          std::advance(size_it, view_->num_vector_dims.value_or(0));
         }
 
-        for (auto e = view_indices_.rend(); indexIt != e; ++indexIt, ++sizeIt) {
-          ++*indexIt;
-          if (*indexIt < *sizeIt) {
+        for (auto e = view_indices_.rend(); index_it != e;
+             ++index_it, ++size_it) {
+          ++*index_it;
+          if (*index_it < *size_it) {
             return *this;
           }
-          *indexIt = 0;
+          *index_it = 0;
         }
 
         view_indices_.clear();
@@ -196,24 +197,24 @@ class Buffer {
     return std::make_shared<Buffer>(Dummy{}, size, sizeof(T));
   }
 
-  char* at(std::optional<int64_t> idx, int64_t elementSize) {
-    auto byteOffset = GetByteOffset(idx, elementSize);
-    if (!byteOffset) {
+  char* at(std::optional<int64_t> idx, int64_t element_size) {
+    auto byte_offset = GetByteOffset(idx, element_size);
+    if (!byte_offset) {
       return &storage_.data()[0];
     }
-    return &storage_.data()[*byteOffset];
+    return &storage_.data()[*byte_offset];
   }
 
-  const char* at(std::optional<int64_t> idx, int64_t elementSize) const {
-    auto byteOffset = GetByteOffset(idx, elementSize);
-    if (!byteOffset) {
+  const char* at(std::optional<int64_t> idx, int64_t element_size) const {
+    auto byte_offset = GetByteOffset(idx, element_size);
+    if (!byte_offset) {
       return &storage_.data()[0];
     }
-    return &storage_.data()[*byteOffset];
+    return &storage_.data()[*byte_offset];
   }
 
-  Buffer(Dummy, size_t numElements, size_t elementSize)
-      : storage_(numElements * elementSize) {}
+  Buffer(Dummy, size_t num_elements, size_t element_size)
+      : storage_(num_elements * element_size) {}
 
   int64_t GetByteSize() const { return storage_.size(); }
 
@@ -247,7 +248,7 @@ class Buffer {
 
  private:
   std::optional<size_t> GetByteOffset(std::optional<int64_t> idx,
-                                      int64_t elementSize) const {
+                                      int64_t element_size) const {
     if (!idx) {
       SetFailure("out of bounds access");
       return std::nullopt;
@@ -263,7 +264,7 @@ class Buffer {
       return std::nullopt;
     }
 
-    return *idx * elementSize;
+    return *idx * element_size;
   }
 
   llvm::SmallVector<char> storage_;
@@ -285,10 +286,10 @@ struct TensorOrMemref {
 
   static TensorOrMemref<T> EmptyLike(const BufferView& view,
                                      ArrayRef<int64_t> layout = {}) {
-    BufferView newView = view;
-    newView.offset = 0;
-    newView.strides = BufferView::GetStridesForLayout(view.sizes, layout);
-    return {Buffer::Allocate<T>(view.GetNumElements(true)), newView};
+    BufferView new_view = view;
+    new_view.offset = 0;
+    new_view.strides = BufferView::GetStridesForLayout(view.sizes, layout);
+    return {Buffer::Allocate<T>(view.GetNumElements(true)), new_view};
   }
 
   TensorOrMemref<T> Clone(ArrayRef<int64_t> layout = {}) const {
diff --git a/third_party/xla/xla/mlir/tools/mlir_interpreter/framework/tests/interpreter_value_test.cc b/third_party/xla/xla/mlir/tools/mlir_interpreter/framework/tests/interpreter_value_test.cc
index 93e72125440bc4..e049e672cbfa31 100644
--- a/third_party/xla/xla/mlir/tools/mlir_interpreter/framework/tests/interpreter_value_test.cc
+++ b/third_party/xla/xla/mlir/tools/mlir_interpreter/framework/tests/interpreter_value_test.cc
@@ -17,6 +17,7 @@ limitations under the License.
 
 #include <complex>
 #include <cstdint>
+#include <optional>
 #include <variant>
 
 #include <gmock/gmock.h>
@@ -71,12 +72,12 @@ TEST(InterpreterValueTest, FillTensorOfVector) {
 TEST(InterpreterValueTest, FillZeroSizedTensor) {
   auto t = TensorOrMemref<int64_t>::Empty({0, 1});
   InterpreterValue v{t};
-  bool wasCalled = false;
+  bool was_called = false;
   v.Fill([&](llvm::ArrayRef<int64_t> indices) {
-    wasCalled = true;
+    was_called = true;
     return InterpreterValue{indices[0]};
   });
-  EXPECT_FALSE(wasCalled);
+  EXPECT_FALSE(was_called);
 }
 
 TEST(InterpreterValueTest, TypedAlike) {
@@ -216,8 +217,8 @@ TEST(CastTest, UnpackTensor) {
   ASSERT_EQ(InterpreterValueCast<float>(value), 1.0f);
   ASSERT_EQ(InterpreterValueCast<double>(value), 1.0);
 
-  InterpreterValue nonUnit{TensorOrMemref<int8_t>::Empty({2, 2})};
-  ASSERT_EQ(InterpreterValueDynCast<int64_t>(nonUnit), std::nullopt);
+  InterpreterValue non_unit{TensorOrMemref<int8_t>::Empty({2, 2})};
+  ASSERT_EQ(InterpreterValueDynCast<int64_t>(non_unit), std::nullopt);
 }
 
 TEST(CastTest, IdentityCast) {
diff --git a/third_party/xla/xla/mlir/tools/mlir_interpreter/framework/tests/tensor_or_memref_test.cc b/third_party/xla/xla/mlir/tools/mlir_interpreter/framework/tests/tensor_or_memref_test.cc
index 6d431234a57ce1..ae19471933e742 100644
--- a/third_party/xla/xla/mlir/tools/mlir_interpreter/framework/tests/tensor_or_memref_test.cc
+++ b/third_party/xla/xla/mlir/tools/mlir_interpreter/framework/tests/tensor_or_memref_test.cc
@@ -22,6 +22,7 @@ limitations under the License.
 #include <gmock/gmock.h>
 #include <gtest/gtest.h>
 #include "absl/strings/str_join.h"
+#include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/SmallBitVector.h"
 #include "llvm/ADT/SmallVector.h"
 #include "mlir/Support/LLVM.h"  // from @llvm-project
@@ -55,11 +56,11 @@ std::optional<int64_t> GetCollapsedStrideNaive(llvm::ArrayRef<int64_t> dims,
   // Find all physical indices for the dimensions.
   llvm::SmallBitVector v(view.GetNumElements());
   for (const auto& indices : f.Indices()) {
-    SmallVector<int64_t> viewIndices(view.Rank());
+    SmallVector<int64_t> view_indices(view.Rank());
     for (auto [dim, index] : llvm::zip(dims, indices)) {
-      viewIndices[dim] = index;
+      view_indices[dim] = index;
     }
-    v[*view.GetPhysicalIndex(viewIndices)] = true;
+    v[*view.GetPhysicalIndex(view_indices)] = true;
   }
 
   if (v.count() != f.GetNumElements()) return std::nullopt;
@@ -67,12 +68,14 @@ std::optional<int64_t> GetCollapsedStrideNaive(llvm::ArrayRef<int64_t> dims,
 
   // Check that they have a common stride.
   int64_t min = v.find_first();
-  int64_t expectedStride = (v.find_last() - min) / (f.GetNumElements() - 1);
+  int64_t expected_stride = (v.find_last() - min) / (f.GetNumElements() - 1);
   for (int64_t i = 0; i < f.GetNumElements(); ++i) {
-    if (!v[i * expectedStride + min]) return std::nullopt;
+    if (!v[i * expected_stride + min]) {
+      return std::nullopt;
+    }
   }
 
-  return expectedStride;
+  return expected_stride;
 }
 
 TEST(TensorOrMemrefTest, CollapsedStride) {
diff --git a/third_party/xla/xla/mlir/tools/mlir_interpreter/mlir_interpreter_runner.cc b/third_party/xla/xla/mlir/tools/mlir_interpreter/mlir_interpreter_runner.cc
index e75d0a603e810a..e265c0c27f56bd 100644
--- a/third_party/xla/xla/mlir/tools/mlir_interpreter/mlir_interpreter_runner.cc
+++ b/third_party/xla/xla/mlir/tools/mlir_interpreter/mlir_interpreter_runner.cc
@@ -48,17 +48,17 @@ struct Options {
 };
 
 static mlir::OwningOpRef<mlir::Operation *> ParseMlirInput(
-    llvm::StringRef inputFilename, mlir::MLIRContext *context) {
-  std::string errorMessage;
-  auto file = mlir::openInputFile(inputFilename, &errorMessage);
+    llvm::StringRef input_filename, mlir::MLIRContext *context) {
+  std::string error_message;
+  auto file = mlir::openInputFile(input_filename, &error_message);
   if (!file) {
-    llvm::errs() << errorMessage << "\n";
+    llvm::errs() << error_message << "\n";
     return {};
   }
 
-  auto sourceMgr = std::make_shared<llvm::SourceMgr>();
-  sourceMgr->AddNewSourceBuffer(std::move(file), mlir::SMLoc());
-  return mlir::parseSourceFileForTool(sourceMgr, context,
+  auto source_mgr = std::make_shared<llvm::SourceMgr>();
+  source_mgr->AddNewSourceBuffer(std::move(file), mlir::SMLoc());
+  return mlir::parseSourceFileForTool(source_mgr, context,
                                       /*insertImplicitModule=*/true);
 }
 
@@ -70,9 +70,9 @@ static mlir::LogicalResult Run(mlir::ModuleOp module,
     return mlir::failure();
   }
 
-  mlir::SymbolTable symbolTable{module};
+  mlir::SymbolTable symbol_table{module};
   auto results =
-      mlir::interpreter::RunInterpreter(symbolTable, function, {}, {});
+      mlir::interpreter::RunInterpreter(symbol_table, function, {}, {});
   if (!results.ok()) {
     llvm::errs() << "Interpreter failed\n";
     return mlir::failure();
@@ -103,25 +103,25 @@ int main(int argc, char *argv[]) {
 
   mlir::MLIRContext context(registry);
   context.allowUnregisteredDialects();
-  auto parsedInput = ParseMlirInput(options.input_filename, &context);
-  if (!parsedInput) {
+  auto parsed_input = ParseMlirInput(options.input_filename, &context);
+  if (!parsed_input) {
     llvm::errs() << "Failed to parse module.\n";
     return 1;
   }
-  auto module = llvm::dyn_cast<mlir::ModuleOp>(**parsedInput);
+  auto module = llvm::dyn_cast<mlir::ModuleOp>(**parsed_input);
   if (!module) {
     llvm::errs() << "Parsing returned something that's not a module.\n";
     return 1;
   }
 
   if (options.run_all_functions) {
-    bool allSucceeded = true;
+    bool all_succeeded = true;
     module.walk([&](mlir::func::FuncOp function) {
       if (!function.isPrivate()) {
-        allSucceeded &= Run(module, function).succeeded();
+        all_succeeded &= Run(module, function).succeeded();
       }
     });
-    if (!allSucceeded) {
+    if (!all_succeeded) {
       return 1;
     }
   } else {
diff --git a/third_party/xla/xla/mlir/tools/mlir_replay/BUILD b/third_party/xla/xla/mlir/tools/mlir_replay/BUILD
index c0bb7fe0e5c580..0d58db72f2aad1 100644
--- a/third_party/xla/xla/mlir/tools/mlir_replay/BUILD
+++ b/third_party/xla/xla/mlir/tools/mlir_replay/BUILD
@@ -27,8 +27,10 @@ xla_cc_binary(
         "//xla/mlir_hlo:hlo_dialect_registration",
         "//xla/service:hlo_proto_cc",
         "//xla/tsl/util:command_line_flags",
+        "@com_google_absl//absl/log:check",
         "@com_google_absl//absl/strings",
         "@com_google_absl//absl/strings:str_format",
+        "@llvm-project//llvm:Support",
         "@llvm-project//mlir:AllPassesAndDialects",
         "@llvm-project//mlir:IR",
         "@local_tsl//tsl/platform:env",
diff --git a/third_party/xla/xla/mlir/tools/mlir_replay/mlir_replay.cc b/third_party/xla/xla/mlir/tools/mlir_replay/mlir_replay.cc
index e02e7fd3eb2641..b8a3f1e7e30694 100644
--- a/third_party/xla/xla/mlir/tools/mlir_replay/mlir_replay.cc
+++ b/third_party/xla/xla/mlir/tools/mlir_replay/mlir_replay.cc
@@ -17,8 +17,12 @@ limitations under the License.
 #include <string>
 #include <vector>
 
+#include "absl/log/check.h"
 #include "absl/strings/str_format.h"
+#include "absl/strings/str_join.h"
 #include "absl/strings/str_split.h"
+#include "llvm/ADT/SmallVector.h"
+#include "llvm/Support/raw_ostream.h"
 #include "mlir/IR/DialectRegistry.h"  // from @llvm-project
 #include "mlir/IR/MLIRContext.h"  // from @llvm-project
 #include "mlir/InitAllDialects.h"  // from @llvm-project
diff --git a/third_party/xla/xla/mlir/tools/mlir_replay/mlir_replay_lib.cc b/third_party/xla/xla/mlir/tools/mlir_replay/mlir_replay_lib.cc
index c13f06a9b953e6..8bf928f314d425 100644
--- a/third_party/xla/xla/mlir/tools/mlir_replay/mlir_replay_lib.cc
+++ b/third_party/xla/xla/mlir/tools/mlir_replay/mlir_replay_lib.cc
@@ -31,13 +31,22 @@ limitations under the License.
 #include "absl/status/status.h"
 #include "absl/status/statusor.h"
 #include "llvm/ADT/APInt.h"
+#include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/Sequence.h"
+#include "llvm/ADT/SmallVector.h"
+#include "llvm/Support/Casting.h"
+#include "llvm/Support/MemoryBuffer.h"
 #include "llvm/Support/SourceMgr.h"
 #include "llvm/Support/raw_ostream.h"
 #include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
 #include "mlir/IR/BuiltinAttributes.h"  // from @llvm-project
+#include "mlir/IR/BuiltinTypeInterfaces.h"  // from @llvm-project
+#include "mlir/IR/MLIRContext.h"  // from @llvm-project
 #include "mlir/IR/Operation.h"  // from @llvm-project
 #include "mlir/IR/OwningOpRef.h"  // from @llvm-project
+#include "mlir/IR/SymbolTable.h"  // from @llvm-project
+#include "mlir/IR/TypeRange.h"  // from @llvm-project
+#include "mlir/IR/Value.h"  // from @llvm-project
 #include "mlir/Support/LLVM.h"  // from @llvm-project
 #include "mlir/Support/LogicalResult.h"  // from @llvm-project
 #include "mlir/Tools/ParseUtilities.h"  // from @llvm-project
@@ -123,7 +132,7 @@ mlir::FailureOr<mlir::interpreter::InterpreterValue> MakeRandomInput(
 // Extracts a mapping from function arguments to allocated buffers.
 // The buffer assignment is only relevant once the program is bufferized and
 // memref results were converted to arguments.
-std::vector<int64_t> extractXlaBufferAssignment(func::FuncOp main) {
+std::vector<int64_t> ExtractXlaBufferAssignment(func::FuncOp main) {
   std::vector<int64_t> buffer_assignment(main.getNumArguments());
   auto result_mapping =
       main->getAttrOfType<IntegerAttr>("xla_framework.result_mapping");
@@ -164,11 +173,11 @@ absl::StatusOr<SmallVector<InterpreterValue>> Run(
     MLIRContext& context, const std::string& mlir_ir,
     const xla::HloSnapshot& snapshot, ExecutionTrace* trace,
     const std::vector<std::string>& entry) {
-  auto sourceMgr = std::make_shared<llvm::SourceMgr>();
-  sourceMgr->AddNewSourceBuffer(llvm::MemoryBuffer::getMemBuffer(mlir_ir),
-                                mlir::SMLoc());
+  auto source_mgr = std::make_shared<llvm::SourceMgr>();
+  source_mgr->AddNewSourceBuffer(llvm::MemoryBuffer::getMemBuffer(mlir_ir),
+                                 mlir::SMLoc());
   mlir::OwningOpRef<mlir::Operation*> module =
-      mlir::parseSourceFileForTool(sourceMgr, &context, false);
+      mlir::parseSourceFileForTool(source_mgr, &context, false);
   if (!module) {
     return absl::InvalidArgumentError("failed to parse MLIR");
   }
@@ -199,7 +208,7 @@ absl::StatusOr<SmallVector<InterpreterValue>> Run(
         "expected all function arguments to be shaped types");
   }
 
-  auto args_to_buffers = extractXlaBufferAssignment(main);
+  auto args_to_buffers = ExtractXlaBufferAssignment(main);
   TF_ASSIGN_OR_RETURN(auto args,
                       LoadArgs(snapshot, main.getBody().getArgumentTypes()));
   auto out_args =
diff --git a/third_party/xla/xla/mlir/tools/mlir_replay/public/BUILD b/third_party/xla/xla/mlir/tools/mlir_replay/public/BUILD
index 693c7e2f25f3e5..c9d831f9f61c31 100644
--- a/third_party/xla/xla/mlir/tools/mlir_replay/public/BUILD
+++ b/third_party/xla/xla/mlir/tools/mlir_replay/public/BUILD
@@ -32,6 +32,7 @@ xla_cc_test(
     srcs = ["execution_trace_utils_test.cc"],
     deps = [
         ":execution_trace_utils",
+        "//xla:literal",
         "//xla:literal_util",
         "//xla/mlir/tools/mlir_interpreter/framework",
         "@com_google_googletest//:gtest_main",
diff --git a/third_party/xla/xla/mlir/tools/mlir_replay/public/execution_trace_utils.cc b/third_party/xla/xla/mlir/tools/mlir_replay/public/execution_trace_utils.cc
index 36482f40cadd01..90357dde29d622 100644
--- a/third_party/xla/xla/mlir/tools/mlir_replay/public/execution_trace_utils.cc
+++ b/third_party/xla/xla/mlir/tools/mlir_replay/public/execution_trace_utils.cc
@@ -15,6 +15,7 @@ limitations under the License.
 
 #include "xla/mlir/tools/mlir_replay/public/execution_trace_utils.h"
 
+#include <cassert>
 #include <complex>
 #include <cstdint>
 #include <functional>
@@ -25,12 +26,19 @@ limitations under the License.
 #include <variant>
 
 #include "absl/status/status.h"
+#include "absl/status/statusor.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/SmallVector.h"
+#include "llvm/Support/Casting.h"
 #include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
 #include "mlir/IR/Attributes.h"  // from @llvm-project
 #include "mlir/IR/BuiltinAttributes.h"  // from @llvm-project
+#include "mlir/IR/BuiltinTypeInterfaces.h"  // from @llvm-project
 #include "mlir/IR/BuiltinTypes.h"  // from @llvm-project
+#include "mlir/IR/Region.h"  // from @llvm-project
+#include "mlir/IR/Types.h"  // from @llvm-project
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+#include "xla/literal.h"
 #include "xla/mlir/tools/mlir_interpreter/framework/interpreter_value.h"
 #include "xla/mlir/tools/mlir_interpreter/framework/tensor_or_memref.h"
 #include "xla/mlir/tools/mlir_replay/public/execution_trace.pb.h"
diff --git a/third_party/xla/xla/mlir/tools/mlir_replay/public/execution_trace_utils_test.cc b/third_party/xla/xla/mlir/tools/mlir_replay/public/execution_trace_utils_test.cc
index 31a6ae074dea48..094bad426493c3 100644
--- a/third_party/xla/xla/mlir/tools/mlir_replay/public/execution_trace_utils_test.cc
+++ b/third_party/xla/xla/mlir/tools/mlir_replay/public/execution_trace_utils_test.cc
@@ -25,6 +25,7 @@ limitations under the License.
 #include <gtest/gtest.h>
 #include "llvm/ADT/STLExtras.h"
 #include "mlir/Support/LLVM.h"  // from @llvm-project
+#include "xla/literal.h"
 #include "xla/literal_util.h"
 #include "xla/mlir/tools/mlir_interpreter/framework/interpreter_value.h"
 #include "xla/mlir/tools/mlir_interpreter/framework/tensor_or_memref.h"

From 85be21af7762f1c1b053cb595e690bcecb5aec9f Mon Sep 17 00:00:00 2001
From: Thomas Joerg <tjoerg@google.com>
Date: Wed, 15 May 2024 09:28:03 -0700
Subject: [PATCH 187/478] [XLA:GPU] Pass the CUDA / ROCm toolkit version
 explicitly for autotuning and GEMM rewriting.

This allows to remove more `#if GOOGLE_CUDA` preprocessor directives from HLO passes.

PiperOrigin-RevId: 633980533
---
 third_party/xla/xla/service/gpu/BUILD         |   4 +-
 .../service/gpu/gemm_algorithm_picker_test.cc |  36 ++-
 .../xla/service/gpu/gemm_fusion_autotuner.cc  |  27 +-
 .../xla/service/gpu/gemm_fusion_autotuner.h   |  14 +-
 .../service/gpu/gemm_fusion_autotuner_test.cc |  23 +-
 .../xla/xla/service/gpu/gemm_rewriter.cc      |  54 ++--
 .../xla/xla/service/gpu/gemm_rewriter.h       |   7 +-
 .../xla/xla/service/gpu/gpu_compiler.cc       |   8 +-
 .../xla/xla/service/gpu/nvptx_compiler.cc     |   3 +-
 .../gemm_broadcast_folding_rewrite_test.cc    |   8 +-
 .../service/gpu/tests/gemm_rewrite_test.cc    | 265 ++++++++++++------
 11 files changed, 301 insertions(+), 148 deletions(-)

diff --git a/third_party/xla/xla/service/gpu/BUILD b/third_party/xla/xla/service/gpu/BUILD
index 4dc41fcce52caf..475c1e680c00fe 100644
--- a/third_party/xla/xla/service/gpu/BUILD
+++ b/third_party/xla/xla/service/gpu/BUILD
@@ -814,7 +814,9 @@ xla_test(
         "@local_tsl//tsl/platform:platform_port",
         "@local_tsl//tsl/platform:status_matchers",
         "@local_tsl//tsl/platform:statusor",
-    ],
+    ] + if_cuda_is_configured([
+        "@local_config_cuda//cuda:cuda_headers",
+    ]),
 )
 
 cc_library(
diff --git a/third_party/xla/xla/service/gpu/gemm_algorithm_picker_test.cc b/third_party/xla/xla/service/gpu/gemm_algorithm_picker_test.cc
index 4da7b2d1792441..505681cbee9eeb 100644
--- a/third_party/xla/xla/service/gpu/gemm_algorithm_picker_test.cc
+++ b/third_party/xla/xla/service/gpu/gemm_algorithm_picker_test.cc
@@ -99,9 +99,12 @@ ENTRY main {
   se::StreamExecutor* stream_exec = executors[0];
   bool changed = false;
   TF_ASSERT_OK_AND_ASSIGN(
-      changed, RunHloPass(GemmRewriter(stream_exec->GetDeviceDescription()
-                                           .gpu_compute_capability()),
-                          m.get()));
+      changed,
+      RunHloPass(
+          GemmRewriter(
+              stream_exec->GetDeviceDescription().gpu_compute_capability(),
+              /*toolkit_version=*/12040),
+          m.get()));
   changed = false;
   DebugOptions opts;
   AutotuneConfig cfg{DeviceConfig{stream_exec, nullptr}, opts};
@@ -125,9 +128,12 @@ ENTRY main {
   TF_ASSERT_OK_AND_ASSIGN(m, ParseAndReturnVerifiedModule(kHlo, module_cfg));
   changed = false;
   TF_ASSERT_OK_AND_ASSIGN(
-      changed, RunHloPass(GemmRewriter(stream_exec->GetDeviceDescription()
-                                           .gpu_compute_capability()),
-                          m.get()));
+      changed,
+      RunHloPass(
+          GemmRewriter(
+              stream_exec->GetDeviceDescription().gpu_compute_capability(),
+              /*toolkit_version=*/12040),
+          m.get()));
   changed = false;
   TF_ASSERT_OK_AND_ASSIGN(changed,
                           RunHloPass(GemmAlgorithmPicker(cfg), m.get()));
@@ -174,9 +180,12 @@ ENTRY main {
 
   bool changed = false;
   TF_ASSERT_OK_AND_ASSIGN(
-      changed, RunHloPass(GemmRewriter(stream_exec->GetDeviceDescription()
-                                           .gpu_compute_capability()),
-                          m.get()));
+      changed,
+      RunHloPass(
+          GemmRewriter(
+              stream_exec->GetDeviceDescription().gpu_compute_capability(),
+              /*toolkit_version=*/12040),
+          m.get()));
   changed = false;
 
   DebugOptions opts;
@@ -208,9 +217,12 @@ ENTRY main {
       stream_exec->GetDeviceDescription().cuda_compute_capability()};
   AutotuneConfig deviceless_cfg{deviceless_config, opts};
   TF_ASSERT_OK_AND_ASSIGN(
-      changed, RunHloPass(GemmRewriter(stream_exec->GetDeviceDescription()
-                                           .gpu_compute_capability()),
-                          m.get()));
+      changed,
+      RunHloPass(
+          GemmRewriter(
+              stream_exec->GetDeviceDescription().gpu_compute_capability(),
+              /*toolkit_version=*/12040),
+          m.get()));
   changed = false;
   TF_ASSERT_OK_AND_ASSIGN(
       changed, RunHloPass(GemmAlgorithmPicker(deviceless_cfg), m.get()))
diff --git a/third_party/xla/xla/service/gpu/gemm_fusion_autotuner.cc b/third_party/xla/xla/service/gpu/gemm_fusion_autotuner.cc
index a2a1cc1c063493..a0d47feab41153 100644
--- a/third_party/xla/xla/service/gpu/gemm_fusion_autotuner.cc
+++ b/third_party/xla/xla/service/gpu/gemm_fusion_autotuner.cc
@@ -365,8 +365,8 @@ absl::StatusOr<std::unique_ptr<HloModule>> TritonGemmAutotuneExtractor(
 }
 
 absl::StatusOr<std::unique_ptr<HloModule>> CublasGemmAutotuneExtractor(
-    const AutotuneConfig& config, const HloFusionInstruction* fusion,
-    const DebugOptions& debug_opts) {
+    const AutotuneConfig& config, const int32_t toolkit_version,
+    const HloFusionInstruction* fusion, const DebugOptions& debug_opts) {
   const HloComputation* fusion_computation =
       fusion->called_computations().at(0);
   std::unique_ptr<HloModule> new_module =
@@ -386,7 +386,7 @@ absl::StatusOr<std::unique_ptr<HloModule>> CublasGemmAutotuneExtractor(
         PrecisionConfig::ALG_DOT_F32_F32_F32);
   }
 
-  GemmRewriter rewriter(config.GetGpuComputeCapability());
+  GemmRewriter rewriter(config.GetGpuComputeCapability(), toolkit_version);
   GpuInstructionFusion fusion_pass(
       /*may_duplicate=*/false, config.GetExecutor()->GetDeviceDescription());
   TF_RETURN_IF_ERROR(rewriter.Run(new_module.get()).status());
@@ -455,6 +455,7 @@ AutotuneResult FromConfig(const Config& config) {
 }
 
 absl::Status DumpAutotunedFusion(const AutotuneConfig& autotune_config,
+                                 const int32_t toolkit_version,
                                  AutotunerCompileUtil& util,
                                  const AutotuneResult result,
                                  const HloFusionInstruction* fusion,
@@ -482,8 +483,8 @@ absl::Status DumpAutotunedFusion(const AutotuneConfig& autotune_config,
               triton_gemm_config, device_desc, fusion, debug_opts,
               /*allow_filtering_kernels_spilling_registers=*/true);
         } else if (result.has_gemm()) {
-          return CublasGemmAutotuneExtractor(autotune_config, fusion,
-                                             debug_opts);
+          return CublasGemmAutotuneExtractor(autotune_config, toolkit_version,
+                                             fusion, debug_opts);
         } else {
           LOG(FATAL) << "Unknown result type: " << result.DebugString();
         }
@@ -736,10 +737,11 @@ GemmFusionAutotunerImpl::CompileAll(
                        })
                        .value_or(nullptr);
     } else if (std::holds_alternative<CuBlasConfig>(config)) {
-      TF_ASSIGN_OR_RETURN(
-          executable, compile_util.Compile([&](const DebugOptions& opts) {
-            return CublasGemmAutotuneExtractor(config_, fusion, opts);
-          }));
+      TF_ASSIGN_OR_RETURN(executable,
+                          compile_util.Compile([&](const DebugOptions& opts) {
+                            return CublasGemmAutotuneExtractor(
+                                config_, toolkit_version_, fusion, opts);
+                          }));
     } else {
       LOG(FATAL) << "Unsupported config type: " << config.index();
     }
@@ -1082,8 +1084,8 @@ absl::Status GemmFusionAutotunerImpl::Autotune(
             << tsl::proto_utils::FromDurationProto(best.run_time());
 
     if (debug_options_.xla_gpu_dump_autotuned_gemm_fusions()) {
-      TF_RETURN_IF_ERROR(DumpAutotunedFusion(config_, compile_util, best,
-                                             fusion, fusion_id++));
+      TF_RETURN_IF_ERROR(DumpAutotunedFusion(
+          config_, toolkit_version_, compile_util, best, fusion, fusion_id++));
     }
 
     const AutotuneCacheKey key = AutotunerUtil::GetKey(fusion, config_);
@@ -1123,7 +1125,8 @@ absl::StatusOr<bool> GemmFusionAutotuner::Run(
   XLA_SCOPED_LOGGING_TIMER("GEMM fusion autotuner");
 
   const DebugOptions& debug_options = module->config().debug_options();
-  GemmFusionAutotunerImpl autotuner(config_, debug_options, thread_pool_);
+  GemmFusionAutotunerImpl autotuner(config_, toolkit_version_, debug_options,
+                                    thread_pool_);
   GemmConfigSetCollector gemm_config_set_collector(&autotuner);
   absl::flat_hash_map<const HloFusionInstruction*, std::vector<Config>>
       gemm_config_sets;
diff --git a/third_party/xla/xla/service/gpu/gemm_fusion_autotuner.h b/third_party/xla/xla/service/gpu/gemm_fusion_autotuner.h
index 3f2a7a7d19a8e4..e7fbbbe6312408 100644
--- a/third_party/xla/xla/service/gpu/gemm_fusion_autotuner.h
+++ b/third_party/xla/xla/service/gpu/gemm_fusion_autotuner.h
@@ -48,8 +48,11 @@ namespace gpu {
 class GemmFusionAutotuner : public HloModulePass {
  public:
   explicit GemmFusionAutotuner(const AutotuneConfig& config,
+                               const int32_t toolkit_version,
                                tsl::thread::ThreadPool* thread_pool)
-      : config_(config), thread_pool_(thread_pool) {}
+      : config_(config),
+        toolkit_version_(toolkit_version),
+        thread_pool_(thread_pool) {}
 
   absl::string_view name() const override { return "triton-autotuner"; }
 
@@ -59,16 +62,20 @@ class GemmFusionAutotuner : public HloModulePass {
       const absl::flat_hash_set<absl::string_view>& execution_threads) override;
 
  private:
-  AutotuneConfig config_;
+  const AutotuneConfig config_;
+  const int32_t toolkit_version_;
   tsl::thread::ThreadPool* thread_pool_;
 };
 
 // Autotuner implementation.
 class GemmFusionAutotunerImpl {
  public:
-  GemmFusionAutotunerImpl(AutotuneConfig config, DebugOptions debug_options,
+  GemmFusionAutotunerImpl(const AutotuneConfig config,
+                          const int32_t toolkit_version,
+                          const DebugOptions debug_options,
                           tsl::thread::ThreadPool* thread_pool)
       : config_(std::move(config)),
+        toolkit_version_(toolkit_version),
         debug_options_(std::move(debug_options)),
         thread_pool_(thread_pool) {}
 
@@ -126,6 +133,7 @@ class GemmFusionAutotunerImpl {
   std::vector<TritonGemmConfig> GetExhaustiveTritonConfigs() const;
 
   const AutotuneConfig config_;
+  const int32_t toolkit_version_;
   const DebugOptions debug_options_;
   tsl::thread::ThreadPool* thread_pool_;
   std::vector<TritonGemmConfig> triton_configs_;
diff --git a/third_party/xla/xla/service/gpu/gemm_fusion_autotuner_test.cc b/third_party/xla/xla/service/gpu/gemm_fusion_autotuner_test.cc
index f18060b89b0507..4a5799aef8483f 100644
--- a/third_party/xla/xla/service/gpu/gemm_fusion_autotuner_test.cc
+++ b/third_party/xla/xla/service/gpu/gemm_fusion_autotuner_test.cc
@@ -15,6 +15,7 @@ limitations under the License.
 #include "xla/service/gpu/gemm_fusion_autotuner.h"
 
 #include <algorithm>
+#include <cstdint>
 #include <memory>
 #include <string>
 #include <utility>
@@ -26,6 +27,7 @@ limitations under the License.
 #include "absl/log/log.h"
 #include "absl/strings/str_format.h"
 #include "absl/strings/string_view.h"
+#include "third_party/gpus/cuda/include/cuda.h"
 #include "xla/autotuning.pb.h"
 #include "xla/error_spec.h"
 #include "xla/hlo/ir/hlo_casting_utils.h"
@@ -150,6 +152,8 @@ class StatelessAutotunerTest : public HloTestBase {
       : HloTestBase(/*verifier_layout_sensitive=*/true,
                     /*allow_mixed_precision_in_hlo_verifier=*/false) {}
 
+  int32_t GetToolkitVersion() const { return CUDA_VERSION; }
+
   void SetUp() override {
     AutotunerUtil::ClearAutotuneResults();
     HloTestBase::SetUp();
@@ -193,7 +197,7 @@ class GemmFusionAutotunerTest : public StatelessAutotunerTest {
         AutotuneConfig{DeviceConfig{backend().default_stream_executor(),
                                     backend().memory_allocator()},
                        opts},
-        &thread_pool);
+        GetToolkitVersion(), &thread_pool);
 
     RunAndFilecheckHloRewrite(
         hlo, std::move(pipeline), expected, [](const HloModule* m) {
@@ -233,10 +237,11 @@ class GemmFusionAutotunerTestWithMorePreciseReduction
 absl::StatusOr<std::vector<TritonGemmConfig>> GetPossibleMatmulAutotuneConfigs(
     const HloDotInstruction& dot,
     const se::CudaComputeCapability& compute_capability,
-    const DebugOptions& debug_options) {
+    const int32_t toolkit_version, const DebugOptions& debug_options) {
   DevicelessConfig test_config{/*model_str=*/"", compute_capability};
   AutotuneConfig autotune_config{test_config, debug_options};
-  GemmFusionAutotunerImpl autotuner(autotune_config, debug_options, nullptr);
+  GemmFusionAutotunerImpl autotuner(autotune_config, toolkit_version,
+                                    debug_options, nullptr);
   return autotuner.GenerateTritonConfigs(dot);
 }
 
@@ -256,7 +261,7 @@ ENTRY e {
       GetPossibleMatmulAutotuneConfigs(
           *Cast<HloDotInstruction>(
               module->entry_computation()->root_instruction()),
-          compute_capability, GetDebugOptionsForTest()));
+          compute_capability, GetToolkitVersion(), GetDebugOptionsForTest()));
   EXPECT_TRUE(std::any_of(
       configs.begin(), configs.end(),
       [](const TritonGemmConfig& config) { return config.num_stages > 2; }));
@@ -278,7 +283,7 @@ ENTRY e {
       GetPossibleMatmulAutotuneConfigs(
           *Cast<HloDotInstruction>(
               module->entry_computation()->root_instruction()),
-          compute_capability, GetDebugOptionsForTest()));
+          compute_capability, GetToolkitVersion(), GetDebugOptionsForTest()));
   EXPECT_TRUE(std::any_of(
       configs.begin(), configs.end(),
       [](const TritonGemmConfig& config) { return config.split_k >= 4; }));
@@ -300,7 +305,7 @@ ENTRY e {
       GetPossibleMatmulAutotuneConfigs(
           *Cast<HloDotInstruction>(
               module->entry_computation()->root_instruction()),
-          compute_capability, GetDebugOptionsForTest()));
+          compute_capability, GetToolkitVersion(), GetDebugOptionsForTest()));
   EXPECT_FALSE(std::any_of(
       configs.begin(), configs.end(),
       [](const TritonGemmConfig& config) { return config.split_k > 1; }));
@@ -655,7 +660,7 @@ ENTRY e {
                                           ->GetDeviceDescription()
                                           .cuda_compute_capability()},
                      opts},
-      &thread_pool);
+      GetToolkitVersion(), &thread_pool);
 
   TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr<VerifiedHloModule> module,
                           ParseAndReturnVerifiedModule(hlo));
@@ -741,7 +746,7 @@ ENTRY e {
       GetPossibleMatmulAutotuneConfigs(
           *Cast<HloDotInstruction>(
               module->entry_computation()->root_instruction()),
-          compute_capability, GetDebugOptionsForTest()));
+          compute_capability, GetToolkitVersion(), GetDebugOptionsForTest()));
   EXPECT_TRUE(std::all_of(
       configs.begin(), configs.end(),
       [](const TritonGemmConfig& config) { return config.split_k == 1; }));
@@ -772,7 +777,7 @@ ENTRY wais {
       GetPossibleMatmulAutotuneConfigs(
           *Cast<HloDotInstruction>(
               module->entry_computation()->root_instruction()),
-          compute_capability, debug_options));
+          compute_capability, GetToolkitVersion(), debug_options));
   for (const auto& config : configs) {
     int metadata_size = config.block_m * config.block_k / 16;
     EXPECT_LE(config.num_warps * WarpSize(), metadata_size);
diff --git a/third_party/xla/xla/service/gpu/gemm_rewriter.cc b/third_party/xla/xla/service/gpu/gemm_rewriter.cc
index 741075422b91ae..8d4064c14d5d5e 100644
--- a/third_party/xla/xla/service/gpu/gemm_rewriter.cc
+++ b/third_party/xla/xla/service/gpu/gemm_rewriter.cc
@@ -508,8 +508,11 @@ auto OptionalBitcast(HloInstruction **optional_bitcast, Pattern pattern) {
 class GemmRewriterVisitor : public DfsHloRewriteVisitor {
  public:
   explicit GemmRewriterVisitor(const se::GpuComputeCapability &gpu_version,
+                               const int32_t toolkit_version,
                                const bool f8_rewrite)
-      : gpu_version_(gpu_version), f8_rewrite_(f8_rewrite) {}
+      : gpu_version_(gpu_version),
+        toolkit_version_(toolkit_version),
+        f8_rewrite_(f8_rewrite) {}
 
   absl::Status HandleDot(HloInstruction *instr) override {
     if (!IsMatrixMultiplication(*instr) &&
@@ -933,25 +936,25 @@ class GemmRewriterVisitor : public DfsHloRewriteVisitor {
       std::vector<std::pair<HloInstruction *, int>> b_ops) {
     GemmBackendConfig &gemm_backend_config =
         *gpu_backend_config.mutable_gemm_backend_config();
-#if GOOGLE_CUDA
-    auto cuda_compute_capability_ =
-        std::get<se::CudaComputeCapability>(gpu_version_);
-    // FP8 GEMM kernels are only available on Ada, Hopper, and later
-    // architectures.
-    if (!cuda_compute_capability_.IsAtLeast(8, 9)) {
-      VLOG(1)
-          << "FP8 Custom Calls require Ada, Hopper, or later architectures.";
-      return false;
+    if (std::holds_alternative<se::CudaComputeCapability>(gpu_version_)) {
+      auto cuda_compute_capability =
+          std::get<se::CudaComputeCapability>(gpu_version_);
+      // FP8 GEMM kernels are only available on Ada, Hopper, and later
+      // architectures.
+      if (!cuda_compute_capability.IsAtLeast(8, 9)) {
+        VLOG(1) << "FP8 Custom Calls require Ada, Hopper, or later "
+                   "architectures. Got: "
+                << cuda_compute_capability.ToString()
+                << " and toolkit version: " << toolkit_version_;
+        return false;
+      }
+      // FP8 GEMM kernels are only available with CUDA 12.0 and above
+      if (toolkit_version_ < 12000) {
+        VLOG(1) << "FP8 Custom Calls require CUDA 12.0 or newer.";
+        return false;
+      }
     }
 
-#if CUDA_VERSION < 12000
-    // FP8 GEMM kernels are only available with CUDA 12.0 and above
-    VLOG(1) << "FP8 Custom Calls require CUDA 12.0 or newer.";
-    return false;
-#endif  // CUDA_VERSION < 12000
-
-#endif  // GOOGLE_CUDA
-
 #if TENSORFLOW_USE_ROCM
     auto isrocm = std::get_if<se::RocmComputeCapability>(&gpu_version_);
     if (!isrocm->has_fp8_support()) {
@@ -1748,6 +1751,7 @@ class GemmRewriterVisitor : public DfsHloRewriteVisitor {
 
  private:
   se::GpuComputeCapability gpu_version_;
+  int32_t toolkit_version_;
   bool f8_rewrite_;
 
   // Choose cublas or cublasLt for the target of the custom call that instr will
@@ -2316,8 +2320,9 @@ class GemmWorkspaceRewriteVisitor : public DfsHloRewriteVisitor {
 
 absl::StatusOr<bool> RunOnComputation(HloComputation *computation,
                                       se::GpuComputeCapability gpu_version,
+                                      int32_t toolkit_version,
                                       bool f8_rewrite) {
-  GemmRewriterVisitor visitor(gpu_version, f8_rewrite);
+  GemmRewriterVisitor visitor(gpu_version, toolkit_version, f8_rewrite);
   TF_RETURN_IF_ERROR(computation->Accept(&visitor));
   GemmWorkspaceRewriteVisitor workspace_visitor(gpu_version);
   TF_RETURN_IF_ERROR(computation->Accept(&workspace_visitor));
@@ -2327,8 +2332,10 @@ absl::StatusOr<bool> RunOnComputation(HloComputation *computation,
 }  // anonymous namespace
 
 GemmRewriter::GemmRewriter(se::GpuComputeCapability gpu_version,
-                           bool f8_rewrite)
-    : gpu_version_(gpu_version), f8_rewrite_(f8_rewrite) {}
+                           int32_t toolkit_version, bool f8_rewrite)
+    : gpu_version_(gpu_version),
+      toolkit_version_(toolkit_version),
+      f8_rewrite_(f8_rewrite) {}
 
 absl::StatusOr<bool> GemmRewriter::Run(
     HloModule *module,
@@ -2336,8 +2343,9 @@ absl::StatusOr<bool> GemmRewriter::Run(
   bool changed = false;
   for (HloComputation *computation :
        module->MakeNonfusionComputations(execution_threads)) {
-    TF_ASSIGN_OR_RETURN(
-        bool result, RunOnComputation(computation, gpu_version_, f8_rewrite_));
+    TF_ASSIGN_OR_RETURN(bool result,
+                        RunOnComputation(computation, gpu_version_,
+                                         toolkit_version_, f8_rewrite_));
     changed |= result;
   }
   return changed;
diff --git a/third_party/xla/xla/service/gpu/gemm_rewriter.h b/third_party/xla/xla/service/gpu/gemm_rewriter.h
index 576f75efaeb34a..161a29a3b26bbf 100644
--- a/third_party/xla/xla/service/gpu/gemm_rewriter.h
+++ b/third_party/xla/xla/service/gpu/gemm_rewriter.h
@@ -15,6 +15,8 @@ limitations under the License.
 #ifndef XLA_SERVICE_GPU_GEMM_REWRITER_H_
 #define XLA_SERVICE_GPU_GEMM_REWRITER_H_
 
+#include <cstdint>
+
 #include "absl/container/flat_hash_set.h"
 #include "absl/status/statusor.h"
 #include "absl/strings/string_view.h"
@@ -47,8 +49,8 @@ class GemmRewriter : public HloModulePass {
  public:
   // When f8_rewrite is true, only FP8 GEMMs are rewritten. Otherwise, non-FP8
   // GEMMs are rewritten.
-  explicit GemmRewriter(se::GpuComputeCapability gpu_version,
-                        bool f8_rewrite = false);
+  GemmRewriter(se::GpuComputeCapability gpu_version, int32_t toolkit_version,
+               bool f8_rewrite = false);
   absl::string_view name() const override { return "cublas-gemm-rewriter"; }
 
   using HloPassInterface::Run;
@@ -58,6 +60,7 @@ class GemmRewriter : public HloModulePass {
 
  private:
   se::GpuComputeCapability gpu_version_;
+  int32_t toolkit_version_;
   bool f8_rewrite_;
 };
 
diff --git a/third_party/xla/xla/service/gpu/gpu_compiler.cc b/third_party/xla/xla/service/gpu/gpu_compiler.cc
index 542c2ae0524d6f..d38cfdd2287e54 100644
--- a/third_party/xla/xla/service/gpu/gpu_compiler.cc
+++ b/third_party/xla/xla/service/gpu/gpu_compiler.cc
@@ -1368,14 +1368,16 @@ absl::Status GpuCompiler::OptimizeHloPostLayoutAssignment(
 
     // Rewrite FP8 GEMMs ahead of Triton which currently lacks support for FP8
     // and may rewrite quantized FP8 GEMMs as higher-precision GEMMs.
-    pipeline.AddPass<GemmRewriter>(gpu_version, /*f8_rewrite=*/true);
+    pipeline.AddPass<GemmRewriter>(gpu_version, GetToolkitVersion(),
+                                   /*f8_rewrite=*/true);
     if (debug_options.xla_gpu_enable_triton_gemm() && cuda_cc != nullptr &&
         cuda_cc->IsAtLeast(se::CudaComputeCapability::AMPERE)) {
       pipeline.AddPass<GemvRewriter>();
       pipeline.AddPass<GemmFusion>(gpu_version);
     }
     // Rewrite non-FP8 GEMMs.
-    pipeline.AddPass<GemmRewriter>(gpu_version, /*f8_rewrite=*/false);
+    pipeline.AddPass<GemmRewriter>(gpu_version, GetToolkitVersion(),
+                                   /*f8_rewrite=*/false);
 
     // Rewrite GEMMs with broadcasted inputs as strided GEMMs.
     pipeline.AddPass<GemmBroadcastFoldingRewriter>();
@@ -1440,7 +1442,7 @@ absl::Status GpuCompiler::OptimizeHloPostLayoutAssignment(
   pipeline.AddPass<CallInliner>();
   // TODO(tdanyluk): Apply CublasPadForGemms to the cuBLAS GEMMs generated
   // here for possibly better cuBLAS performance.
-  pipeline.AddPass<GemmRewriter>(gpu_version);
+  pipeline.AddPass<GemmRewriter>(gpu_version, GetToolkitVersion());
   // Rewrite GEMMs with broadcasted inputs as strided GEMMs.
   pipeline.AddPass<GemmBroadcastFoldingRewriter>();
 
diff --git a/third_party/xla/xla/service/gpu/nvptx_compiler.cc b/third_party/xla/xla/service/gpu/nvptx_compiler.cc
index 94b1314fd6cac5..d342a9621a5573 100644
--- a/third_party/xla/xla/service/gpu/nvptx_compiler.cc
+++ b/third_party/xla/xla/service/gpu/nvptx_compiler.cc
@@ -373,7 +373,8 @@ absl::Status NVPTXCompiler::AddConvAndGemmAutotuningPasses(
 absl::Status NVPTXCompiler::AddGemmFusionAutotuningPasses(
     HloPassPipeline* pipeline, HloModule* hlo_module,
     AutotuneConfig& autotune_config, tsl::thread::ThreadPool* thread_pool) {
-  pipeline->AddPass<GemmFusionAutotuner>(autotune_config, thread_pool);
+  pipeline->AddPass<GemmFusionAutotuner>(autotune_config, GetToolkitVersion(),
+                                         thread_pool);
   return absl::OkStatus();
 }
 
diff --git a/third_party/xla/xla/service/gpu/tests/gemm_broadcast_folding_rewrite_test.cc b/third_party/xla/xla/service/gpu/tests/gemm_broadcast_folding_rewrite_test.cc
index 7d6c68bbc3d244..7e98e974e12d5a 100644
--- a/third_party/xla/xla/service/gpu/tests/gemm_broadcast_folding_rewrite_test.cc
+++ b/third_party/xla/xla/service/gpu/tests/gemm_broadcast_folding_rewrite_test.cc
@@ -142,7 +142,7 @@ ENTRY AddDotsFunc {
   TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
                           ParseAndReturnVerifiedModule(hlo_text));
   // Use GemmRewriter to generate cublasGemm call.
-  GemmRewriter gemm_rewriter(GpuComputeComp());
+  GemmRewriter gemm_rewriter(GpuComputeComp(), /*toolkit_version=*/12040);
   TF_ASSERT_OK_AND_ASSIGN(bool changed,
                           this->RunHloPass(&gemm_rewriter, module.get()));
   EXPECT_TRUE(changed);
@@ -168,7 +168,7 @@ ENTRY AddDotsFunc {
   TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
                           ParseAndReturnVerifiedModule(hlo_text));
   // Use GemmRewriter to generate cublasGemm call.
-  GemmRewriter gemm_rewriter(GpuComputeComp());
+  GemmRewriter gemm_rewriter(GpuComputeComp(), /*toolkit_version=*/12040);
   TF_ASSERT_OK_AND_ASSIGN(bool changed,
                           this->RunHloPass(&gemm_rewriter, module.get()));
   EXPECT_TRUE(changed);
@@ -192,7 +192,7 @@ ENTRY %LHSBatchDimNonZero (Arg_1: f32[4,3], Arg_2: f32[4,7,3]) -> f32[4,7,7] {
   TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
                           ParseAndReturnVerifiedModule(hlo_text));
   // Use GemmRewriter to generate cublasGemm call.
-  GemmRewriter gemm_rewriter(GpuComputeComp());
+  GemmRewriter gemm_rewriter(GpuComputeComp(), /*toolkit_version=*/12040);
   TF_ASSERT_OK_AND_ASSIGN(bool changed,
                           this->RunHloPass(&gemm_rewriter, module.get()));
   EXPECT_TRUE(changed);
@@ -215,7 +215,7 @@ ENTRY %RHSBatchDimNonZero (Arg_1: f32[4,3], Arg_2: f32[4,7,3]) -> f32[4,7,7] {
   EXPECT_TRUE(RunAndCompare(hlo_text, ErrorSpec{1e-5, 1e-5}));
   TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
                           ParseAndReturnVerifiedModule(hlo_text));
-  GemmRewriter gemm_rewriter(GpuComputeComp());
+  GemmRewriter gemm_rewriter(GpuComputeComp(), /*toolkit_version=*/12040);
   TF_ASSERT_OK_AND_ASSIGN(bool changed,
                           this->RunHloPass(&gemm_rewriter, module.get()));
   EXPECT_TRUE(changed);
diff --git a/third_party/xla/xla/service/gpu/tests/gemm_rewrite_test.cc b/third_party/xla/xla/service/gpu/tests/gemm_rewrite_test.cc
index 2646ee1aa9353d..5c5fc750b101d6 100644
--- a/third_party/xla/xla/service/gpu/tests/gemm_rewrite_test.cc
+++ b/third_party/xla/xla/service/gpu/tests/gemm_rewrite_test.cc
@@ -72,6 +72,15 @@ class GemmRewriteTest : public GpuCodegenTest {
     return device_desc().gpu_compute_capability();
   }
 
+  int32_t GetToolkitVersion() const {
+#if GOOGLE_CUDA
+    return CUDA_VERSION;
+#elif TENSORFLOW_USE_ROCM
+    return TF_ROCM_VERSION;
+#endif
+    return 0;
+  }
+
   bool IsCuda() {
     return std::holds_alternative<se::CudaComputeCapability>(Capability());
   }
@@ -1210,7 +1219,7 @@ ENTRY test {
 
   TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
                           ParseAndReturnVerifiedModule(hlo_text));
-  GemmRewriter pass(Capability());
+  GemmRewriter pass(Capability(), GetToolkitVersion());
   TF_ASSERT_OK_AND_ASSIGN(bool changed, this->RunHloPass(&pass, module.get()));
   EXPECT_TRUE(changed);
 
@@ -1234,7 +1243,7 @@ ENTRY test {
 
   TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
                           ParseAndReturnVerifiedModule(hlo_text));
-  GemmRewriter pass(Capability());
+  GemmRewriter pass(Capability(), GetToolkitVersion());
   TF_ASSERT_OK_AND_ASSIGN(bool changed, this->RunHloPass(&pass, module.get()));
   EXPECT_TRUE(changed);
 
@@ -1258,7 +1267,7 @@ ENTRY test {
 
   TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
                           ParseAndReturnVerifiedModule(hlo_text));
-  GemmRewriter pass(Capability());
+  GemmRewriter pass(Capability(), GetToolkitVersion());
   TF_ASSERT_OK_AND_ASSIGN(bool changed, this->RunHloPass(&pass, module.get()));
   EXPECT_TRUE(changed);
   EXPECT_THAT(
@@ -1279,7 +1288,7 @@ ENTRY test {
 
   TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
                           ParseAndReturnVerifiedModule(hlo_text));
-  GemmRewriter pass(Capability());
+  GemmRewriter pass(Capability(), GetToolkitVersion());
   TF_ASSERT_OK_AND_ASSIGN(bool changed, this->RunHloPass(&pass, module.get()));
   EXPECT_TRUE(changed);
 
@@ -1303,7 +1312,7 @@ ENTRY test {
 
   TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
                           ParseAndReturnVerifiedModule(hlo_text));
-  GemmRewriter pass(Capability());
+  GemmRewriter pass(Capability(), GetToolkitVersion());
   TF_ASSERT_OK_AND_ASSIGN(bool changed, this->RunHloPass(&pass, module.get()));
   EXPECT_TRUE(changed);
 
@@ -1331,7 +1340,7 @@ ENTRY main {
 
   TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
                           ParseAndReturnVerifiedModule(hlo_text));
-  GemmRewriter pass(Capability());
+  GemmRewriter pass(Capability(), GetToolkitVersion());
   TF_ASSERT_OK_AND_ASSIGN(bool changed, this->RunHloPass(&pass, module.get()));
   EXPECT_TRUE(changed);
 
@@ -1359,7 +1368,7 @@ ENTRY main {
 
   TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
                           ParseAndReturnVerifiedModule(hlo_text));
-  GemmRewriter pass(Capability());
+  GemmRewriter pass(Capability(), GetToolkitVersion());
   TF_ASSERT_OK_AND_ASSIGN(bool changed, this->RunHloPass(&pass, module.get()));
   EXPECT_TRUE(changed);
 
@@ -1389,7 +1398,7 @@ ENTRY main {
 
   TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
                           ParseAndReturnVerifiedModule(hlo_text));
-  GemmRewriter pass(Capability());
+  GemmRewriter pass(Capability(), GetToolkitVersion());
   TF_ASSERT_OK_AND_ASSIGN(bool changed, this->RunHloPass(&pass, module.get()));
   EXPECT_TRUE(changed);
 
@@ -1427,10 +1436,12 @@ ENTRY e {
     lhs_contracting_dims={0}, rhs_contracting_dims={0}
 })";
 
-  RunAndFilecheckHloRewrite(hlo_text,
-                            GemmRewriter(se::CudaComputeCapability{
-                                se::CudaComputeCapability::AMPERE, 0}),
-                            R"(
+  RunAndFilecheckHloRewrite(
+      hlo_text,
+      GemmRewriter(
+          se::CudaComputeCapability{se::CudaComputeCapability::AMPERE, 0},
+          /*toolkit_version=*/12040),
+      R"(
 ; CHECK:  %[[P0:.+]] = f32[2048]{0} parameter(0)
 ; CHECK:  %[[P1:.+]] = f32[2048,16384]{1,0} parameter(1)
 ; CHECK:  %[[CUSTOM_CALL:.+]] = (f32[16384]{0}, s8[4194304]{0}) custom-call(%[[P0]], %[[P1]]), custom_call_target="__cublas$gemm"
@@ -1449,10 +1460,12 @@ ENTRY e {
    lhs_contracting_dims={2}, rhs_contracting_dims={2}
 })";
 
-  RunAndFilecheckHloRewrite(hlo_text,
-                            GemmRewriter(se::CudaComputeCapability{
-                                se::CudaComputeCapability::AMPERE, 0}),
-                            R"(
+  RunAndFilecheckHloRewrite(
+      hlo_text,
+      GemmRewriter(
+          se::CudaComputeCapability{se::CudaComputeCapability::AMPERE, 0},
+          /*toolkit_version=*/12040),
+      R"(
 ; CHECK:  %[[P0:.+]] = f32[10,10,2048]{2,1,0} parameter(0)
 ; CHECK:  %[[P1:.+]] = f32[10,10,2048,16384]{3,2,1,0} parameter(1)
 ; CHECK:  %[[CUSTOM_CALL:.+]] = (f32[10,10,16384]{2,1,0}, s8[4194304]{0}) custom-call(%[[P0]], %[[P1]]), custom_call_target="__cublas$gemm"
@@ -1471,7 +1484,8 @@ ENTRY main {
       lhs_contracting_dims={1}, rhs_contracting_dims={0}, sparsity=L.1@2:4
 })";
   auto hlo_pass = GemmRewriter(
-      se::CudaComputeCapability{se::CudaComputeCapability::AMPERE, 0});
+      se::CudaComputeCapability{se::CudaComputeCapability::AMPERE, 0},
+      /*toolkit_version=*/12040);
   TF_ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_text));
   TF_ASSERT_OK_AND_ASSIGN(bool changed, RunHloPass(&hlo_pass, module.get()));
   EXPECT_FALSE(changed);
@@ -2081,7 +2095,7 @@ ENTRY test {
 
   TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
                           ParseAndReturnVerifiedModule(hlo_text));
-  GemmRewriter pass(Capability());
+  GemmRewriter pass(Capability(), GetToolkitVersion());
   TF_ASSERT_OK_AND_ASSIGN(bool changed, this->RunHloPass(&pass, module.get()));
   EXPECT_TRUE(changed);
 
@@ -2131,7 +2145,7 @@ ENTRY test {
 
   TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
                           ParseAndReturnVerifiedModule(hlo_text));
-  GemmRewriter pass(Capability());
+  GemmRewriter pass(Capability(), GetToolkitVersion());
   TF_ASSERT_OK_AND_ASSIGN(bool changed, this->RunHloPass(&pass, module.get()));
   SCOPED_TRACE(module->ToString());
   EXPECT_TRUE(changed);
@@ -3785,7 +3799,7 @@ ENTRY test {
 
   TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
                           ParseAndReturnVerifiedModule(hlo_text));
-  GemmRewriter pass(Capability());
+  GemmRewriter pass(Capability(), GetToolkitVersion());
   TF_ASSERT_OK_AND_ASSIGN(bool changed, this->RunHloPass(&pass, module.get()));
   EXPECT_TRUE(changed);
 
@@ -3862,7 +3876,7 @@ ENTRY test {
 
   TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
                           ParseAndReturnVerifiedModule(hlo_text));
-  GemmRewriter pass(Capability());
+  GemmRewriter pass(Capability(), GetToolkitVersion());
   TF_ASSERT_OK_AND_ASSIGN(bool changed, this->RunHloPass(&pass, module.get()));
   EXPECT_TRUE(changed);
 
@@ -4145,7 +4159,7 @@ ENTRY test {
 
   TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
                           ParseAndReturnVerifiedModule(hlo_text));
-  GemmRewriter pass(Capability());
+  GemmRewriter pass(Capability(), GetToolkitVersion());
   TF_ASSERT_OK_AND_ASSIGN(bool changed, this->RunHloPass(&pass, module.get()));
   EXPECT_TRUE(changed);
 
@@ -4445,7 +4459,7 @@ ENTRY test {
 
   TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
                           ParseAndReturnVerifiedModule(hlo_text));
-  GemmRewriter pass(Capability());
+  GemmRewriter pass(Capability(), GetToolkitVersion());
   TF_ASSERT_OK_AND_ASSIGN(bool changed, this->RunHloPass(&pass, module.get()));
   SCOPED_TRACE(module->ToString());
   EXPECT_TRUE(changed);
@@ -4827,9 +4841,10 @@ TEST_P(ParameterizedFp8GemmRewriteTest, UnsupportedTypesF8) {
 )";
   EXPECT_TRUE(RunAndCompare(absl::StrReplaceAll(hlo_text, replacements_),
                             ErrorSpec{1e-2, 1e-2}));
-  RunAndFilecheckHloRewrite(hlo_text,
-                            GemmRewriter(Capability(), /*f8_rewrite=*/true),
-                            R"(
+  RunAndFilecheckHloRewrite(
+      hlo_text,
+      GemmRewriter(Capability(), GetToolkitVersion(), /*f8_rewrite=*/true),
+      R"(
 ; CHECK-LABEL: ENTRY %unsupported_types ({{.*}}: <<F8E5M2>>[16,16], {{.*}}: <<F8E5M2>>[16,16]) -> <<F8E5M2>>[16,16] {
 ; CHECK-NEXT:    [[P0:%[^ ]+]] = <<F8E5M2>>[16,16]{1,0} parameter(0)
 ; CHECK-NEXT:    [[P0_CONVERT:%[^ ]+]] = f16[16,16]{1,0} convert([[P0]])
@@ -4862,7 +4877,9 @@ TEST_P(ParameterizedFp8GemmRewriteTest, UnscaledABUnscaledDF8) {
 
   CheckFp8IfSupported(hlo_text);
   RunAndFilecheckHloRewrite(
-      hlo_text, GemmRewriter(CudaHopperOrRocmMI300(), /*f8_rewrite=*/true),
+      hlo_text,
+      GemmRewriter(CudaHopperOrRocmMI300(), GetToolkitVersion(),
+                   /*f8_rewrite=*/true),
       R"(
 ; CHECK-LABEL: ENTRY %test ({{.*}}: <<F8E4M3>>[16,32], {{.*}}: <<F8E4M3>>[32,16]) -> <<F8E4M3>>[16,16] {
 ; CHECK-NEXT:    [[P0:%[^ ]+]] = <<F8E4M3>>[16,32]{1,0} parameter(0)
@@ -4915,7 +4932,9 @@ TEST_P(ParameterizedFp8GemmRewriteTest, UnscaledABUnscaledDMatrixBiasF8) {
 
   CheckFp8IfSupported(hlo_text);
   RunAndFilecheckHloRewrite(
-      hlo_text, GemmRewriter(CudaHopperOrRocmMI300(), /*f8_rewrite=*/true),
+      hlo_text,
+      GemmRewriter(CudaHopperOrRocmMI300(), GetToolkitVersion(),
+                   /*f8_rewrite=*/true),
       R"(
 ; CHECK-LABEL: ENTRY %test ({{.*}}: <<F8E4M3>>[16,32], {{.*}}: <<F8E4M3>>[32,16], {{.*}}: <<F8E4M3>>[16,16]) -> <<F8E4M3>>[16,16] {
 ; CHECK-NEXT:    [[P0:%[^ ]+]] = <<F8E4M3>>[16,32]{1,0} parameter(0)
@@ -4976,7 +4995,9 @@ TEST_P(ParameterizedFp8GemmRewriteTest, ScaledABUnscaledDF8) {
 
   CheckFp8IfSupported(hlo_text);
   RunAndFilecheckHloRewrite(
-      hlo_text, GemmRewriter(CudaHopperOrRocmMI300(), /*f8_rewrite=*/true),
+      hlo_text,
+      GemmRewriter(CudaHopperOrRocmMI300(), GetToolkitVersion(),
+                   /*f8_rewrite=*/true),
       R"(
 ; CHECK-LABEL: ENTRY %test ({{.*}}: <<F8E4M3>>[16,32], {{.*}}: <<F8E4M3>>[32,16], {{.*}}: f32[], {{.*}}: f32[]) -> f32[16,16] {
 ; CHECK-NEXT:    [[P0:%[^ ]+]] = <<F8E4M3>>[16,32]{1,0} parameter(0)
@@ -5035,7 +5056,9 @@ TEST_P(ParameterizedFp8GemmRewriteTest, ScaledABUnscaledDPaddedF8) {
 
   CheckFp8IfSupported(hlo_text);
   RunAndFilecheckHloRewrite(
-      hlo_text, GemmRewriter(CudaHopperOrRocmMI300(), /*f8_rewrite=*/true),
+      hlo_text,
+      GemmRewriter(CudaHopperOrRocmMI300(), GetToolkitVersion(),
+                   /*f8_rewrite=*/true),
       R"(
 ; CHECK-LABEL: ENTRY %test ({{.*}}: <<F8E4M3>>[13,17], {{.*}}: <<F8E4M3>>[17,31], {{.*}}: f32[], {{.*}}: f32[]) -> f32[13,31] {
 ; CHECK-NEXT:    [[P0:%[^ ]+]] = <<F8E4M3>>[13,17]{1,0} parameter(0)
@@ -5101,7 +5124,8 @@ TEST_P(ParameterizedFp8GemmRewriteTest, ScaledABUnscaledDBitcastF8) {
 
   TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
                           ParseAndReturnVerifiedModule(hlo_text));
-  GemmRewriter pass(CudaHopperOrRocmMI300(), /*f8_rewrite=*/true);
+  GemmRewriter pass(CudaHopperOrRocmMI300(), GetToolkitVersion(),
+                    /*f8_rewrite=*/true);
   TF_ASSERT_OK_AND_ASSIGN(bool changed, this->RunHloPass(&pass, module.get()));
   EXPECT_TRUE(changed);
 
@@ -5146,7 +5170,9 @@ TEST_P(ParameterizedFp8GemmRewriteTest, ScaledABUnscaledDUnaryOpsF8) {
 
   CheckFp8IfSupported(hlo_text);
   RunAndFilecheckHloRewrite(
-      hlo_text, GemmRewriter(CudaHopperOrRocmMI300(), /*f8_rewrite=*/true),
+      hlo_text,
+      GemmRewriter(CudaHopperOrRocmMI300(), GetToolkitVersion(),
+                   /*f8_rewrite=*/true),
       R"(
 
 ; CHECK-LABEL: ENTRY %test ({{.*}}: <<F8E4M3>>[3], {{.*}}: <<F8E4M3>>[32,16], {{.*}}: f32[], {{.*}}: f32[]) -> f32[16,16] {
@@ -5213,13 +5239,16 @@ TEST_P(ParameterizedFp8GemmRewriteTest, ScaledABUnscaledDDynamicSliceF8) {
 
   TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
                           ParseAndReturnVerifiedModule(hlo_text));
-  GemmRewriter pass(CudaHopperOrRocmMI300(), /*f8_rewrite=*/true);
+  GemmRewriter pass(CudaHopperOrRocmMI300(), GetToolkitVersion(),
+                    /*f8_rewrite=*/true);
   TF_ASSERT_OK_AND_ASSIGN(bool changed, this->RunHloPass(&pass, module.get()));
   EXPECT_TRUE(changed);
 
   CheckFp8IfSupported(hlo_text);
   RunAndFilecheckHloRewrite(
-      hlo_text, GemmRewriter(CudaHopperOrRocmMI300(), /*f8_rewrite=*/true),
+      hlo_text,
+      GemmRewriter(CudaHopperOrRocmMI300(), GetToolkitVersion(),
+                   /*f8_rewrite=*/true),
       R"(
 ; CHECK-LABEL: ENTRY %test ({{.*}}: <<F8E4M3>>[32,32], {{.*}}: <<F8E4M3>>[16,32], {{.*}}: f32[], {{.*}}: f32[]) -> f32[16,16] {
 ; CHECK-NEXT:    [[P0:%[^ ]+]] = <<F8E4M3>>[32,32]{1,0} parameter(0)
@@ -5282,13 +5311,16 @@ TEST_P(ParameterizedFp8GemmRewriteTest, ScaledABUnscaledDSelectF8) {
 
   TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
                           ParseAndReturnVerifiedModule(hlo_text));
-  GemmRewriter pass(CudaHopperOrRocmMI300(), /*f8_rewrite=*/true);
+  GemmRewriter pass(CudaHopperOrRocmMI300(), GetToolkitVersion(),
+                    /*f8_rewrite=*/true);
   TF_ASSERT_OK_AND_ASSIGN(bool changed, this->RunHloPass(&pass, module.get()));
   EXPECT_TRUE(changed);
 
   CheckFp8IfSupported(hlo_text);
   RunAndFilecheckHloRewrite(
-      hlo_text, GemmRewriter(CudaHopperOrRocmMI300(), /*f8_rewrite=*/true),
+      hlo_text,
+      GemmRewriter(CudaHopperOrRocmMI300(), GetToolkitVersion(),
+                   /*f8_rewrite=*/true),
       R"(
 ; CHECK-LABEL: ENTRY %test ({{.*}}: <<F8E4M3>>[16,32], {{.*}}: <<F8E4M3>>[16,32], {{.*}}: f32[], {{.*}}: f32[], {{.*}}: pred[16,32]) -> f32[16,16] {
 ; CHECK-NEXT:    [[P0:%[^ ]+]] = <<F8E4M3>>[16,32]{1,0} parameter(0)
@@ -5355,7 +5387,8 @@ TEST_P(ParameterizedFp8GemmRewriteTest,
 
   TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
                           ParseAndReturnVerifiedModule(hlo_text));
-  GemmRewriter pass(CudaHopperOrRocmMI300(), /*f8_rewrite=*/true);
+  GemmRewriter pass(CudaHopperOrRocmMI300(), GetToolkitVersion(),
+                    /*f8_rewrite=*/true);
   TF_ASSERT_OK_AND_ASSIGN(bool changed, this->RunHloPass(&pass, module.get()));
   EXPECT_FALSE(changed);
 }
@@ -5390,7 +5423,9 @@ TEST_P(ParameterizedFp8GemmRewriteTest, BatchedScaledABUnscaledDF8) {
 
   CheckFp8IfSupported(hlo_text);
   RunAndFilecheckHloRewrite(
-      hlo_text, GemmRewriter(CudaHopperOrRocmMI300(), /*f8_rewrite=*/true),
+      hlo_text,
+      GemmRewriter(CudaHopperOrRocmMI300(), GetToolkitVersion(),
+                   /*f8_rewrite=*/true),
       R"(
 ; CHECK-LABEL: ENTRY %test ({{.*}}: <<F8E4M3>>[10,16,32], {{.*}}: <<F8E4M3>>[10,32,16], {{.*}}: f32[], {{.*}}: f32[]) -> f32[10,16,16] {
 ; CHECK-NEXT:    [[P0:%[^ ]+]] = <<F8E4M3>>[10,16,32]{2,1,0} parameter(0)
@@ -5452,7 +5487,9 @@ TEST_P(ParameterizedFp8GemmRewriteTest, ScaledABAlphaDF8) {
 
   CheckFp8IfSupported(hlo_text);
   RunAndFilecheckHloRewrite(
-      hlo_text, GemmRewriter(CudaHopperOrRocmMI300(), /*f8_rewrite=*/true),
+      hlo_text,
+      GemmRewriter(CudaHopperOrRocmMI300(), GetToolkitVersion(),
+                   /*f8_rewrite=*/true),
       R"(
 
 ; CHECK-LABEL: ENTRY %test ({{.*}}: <<F8E4M3>>[16,32], {{.*}}: <<F8E4M3>>[32,16], {{.*}}: f32[], {{.*}}: f32[]) -> f32[16,16] {
@@ -5515,7 +5552,9 @@ TEST_P(ParameterizedFp8GemmRewriteTest, ScaledABUnscaledDReluActivationF8) {
 
   CheckFp8IfSupported(hlo_text);
   RunAndFilecheckHloRewrite(
-      hlo_text, GemmRewriter(CudaHopperOrRocmMI300(), /*f8_rewrite=*/true),
+      hlo_text,
+      GemmRewriter(CudaHopperOrRocmMI300(), GetToolkitVersion(),
+                   /*f8_rewrite=*/true),
       R"(
 
 ; CHECK-LABEL: ENTRY %test ({{.*}}: <<F8E4M3>>[16,32], {{.*}}: <<F8E4M3>>[32,16], {{.*}}: f32[], {{.*}}: f32[]) -> f32[16,16] {
@@ -5599,7 +5638,9 @@ TEST_P(ParameterizedFp8GemmRewriteTest,
 // than 12.4.
 #if (GOOGLE_CUDA && CUDA_VERSION >= 12040) || TENSORFLOW_USE_ROCM
   RunAndFilecheckHloRewrite(
-      hlo_text, GemmRewriter(CudaHopperOrRocmMI300(), /*f8_rewrite=*/true),
+      hlo_text,
+      GemmRewriter(CudaHopperOrRocmMI300(), GetToolkitVersion(),
+                   /*f8_rewrite=*/true),
       R"(
 ; CHECK-LABEL: ENTRY %test ({{.*}}: <<F8E4M3>>[16,32], {{.*}}: <<F8E4M3>>[32,16], {{.*}}: bf16[], {{.*}}: bf16[], {{.*}}: bf16[16]) -> bf16[16,16] {
 ; CHECK-NEXT:    [[P0:%[^ ]+]] = <<F8E4M3>>[16,32]{1,0} parameter(0)
@@ -5686,7 +5727,9 @@ TEST_P(ParameterizedFp8GemmRewriteTest,
   // Currently, hipBlasLt does not support output datatype bf16 for fp8 matmul.
   // And no fusion was done for such cases.
   RunAndFilecheckHloRewrite(
-      hlo_text, GemmRewriter(CudaHopperOrRocmMI300(), /*f8_rewrite=*/true),
+      hlo_text,
+      GemmRewriter(CudaHopperOrRocmMI300(), GetToolkitVersion(),
+                   /*f8_rewrite=*/true),
       R"(
 ; CHECK-LABEL: ENTRY %test ({{.*}}: <<F8E4M3>>[16,32], {{.*}}: <<F8E4M3>>[32,16], {{.*}}: bf16[], {{.*}}: bf16[]) -> bf16[16,16] {
 ; CHECK-NEXT:    [[P0:%[^ ]+]] = <<F8E4M3>>[16,32]{1,0} parameter(0)
@@ -5750,7 +5793,9 @@ TEST_P(ParameterizedFp8GemmRewriteTest, InvScaledABUnscaledDF8) {
 
   CheckFp8IfSupported(hlo_text);
   RunAndFilecheckHloRewrite(
-      hlo_text, GemmRewriter(CudaHopperOrRocmMI300(), /*f8_rewrite=*/true),
+      hlo_text,
+      GemmRewriter(CudaHopperOrRocmMI300(), GetToolkitVersion(),
+                   /*f8_rewrite=*/true),
       R"(
 ; CHECK:           custom_call_target="__cublas$lt$matmul$f8",
       )");
@@ -5788,7 +5833,9 @@ TEST_P(ParameterizedFp8GemmRewriteTest, ScaledABUnscaledDMatrixBiasF8) {
 
   CheckFp8IfSupported(hlo_text);
   RunAndFilecheckHloRewrite(
-      hlo_text, GemmRewriter(CudaHopperOrRocmMI300(), /*f8_rewrite=*/true),
+      hlo_text,
+      GemmRewriter(CudaHopperOrRocmMI300(), GetToolkitVersion(),
+                   /*f8_rewrite=*/true),
       R"(
 
 ; CHECK-LABEL: ENTRY %test ({{.*}}: <<F8E4M3>>[16,32], {{.*}}: <<F8E4M3>>[32,16], {{.*}}: f32[16,16], {{.*}}: f32[], {{.*}}: f32[]) -> f32[16,16] {
@@ -5851,7 +5898,9 @@ TEST_P(ParameterizedFp8GemmRewriteTest, ScaledABUnscaledDMatrixBiasPaddedF8) {
 
   CheckFp8IfSupported(hlo_text);
   RunAndFilecheckHloRewrite(
-      hlo_text, GemmRewriter(CudaHopperOrRocmMI300(), /*f8_rewrite=*/true),
+      hlo_text,
+      GemmRewriter(CudaHopperOrRocmMI300(), GetToolkitVersion(),
+                   /*f8_rewrite=*/true),
       R"(
 
 ; CHECK-LABEL: ENTRY %test ({{.*}}: <<F8E4M3>>[14,31], {{.*}}: <<F8E4M3>>[31,14], {{.*}}: f32[14,14], {{.*}}: f32[], {{.*}}: f32[]) -> f32[14,14] {
@@ -5921,7 +5970,9 @@ TEST_P(ParameterizedFp8GemmRewriteTest, UnscaledABScaledDF8) {
 
   CheckFp8IfSupported(hlo_text, ErrorSpec{1e-2, 1e-1});
   RunAndFilecheckHloRewrite(
-      hlo_text, GemmRewriter(CudaHopperOrRocmMI300(), /*f8_rewrite=*/true),
+      hlo_text,
+      GemmRewriter(CudaHopperOrRocmMI300(), GetToolkitVersion(),
+                   /*f8_rewrite=*/true),
       R"(
 ; CHECK-LABEL: ENTRY %test ({{.*}}: <<F8E4M3>>[16,32], {{.*}}: <<F8E4M3>>[32,16], {{.*}}: f32[]) -> <<F8E4M3>>[16,16] {
 ; CHECK:         [[P0:%[^ ]+]] = <<F8E4M3>>[16,32]{1,0} parameter(0)
@@ -5977,7 +6028,9 @@ TEST_P(ParameterizedFp8GemmRewriteTest, UnscaledABScaledF32DF8) {
 
   CheckFp8IfSupported(hlo_text, ErrorSpec{1e-2, 1e-1});
   RunAndFilecheckHloRewrite(
-      hlo_text, GemmRewriter(CudaHopperOrRocmMI300(), /*f8_rewrite=*/true),
+      hlo_text,
+      GemmRewriter(CudaHopperOrRocmMI300(), GetToolkitVersion(),
+                   /*f8_rewrite=*/true),
       R"(
 ; CHECK-LABEL: ENTRY %test ({{.*}}: <<F8E4M3>>[16,32], {{.*}}: <<F8E4M3>>[32,16], {{.*}}: f32[]) -> f32[16,16] {
 ; CHECK:         [[P0:%[^ ]+]] = <<F8E4M3>>[16,32]{1,0} parameter(0)
@@ -6033,7 +6086,9 @@ TEST_P(ParameterizedFp8GemmRewriteTest, UnscaledABInvScaledF32DF8) {
 
   CheckFp8IfSupported(hlo_text, ErrorSpec{1e-2, 1e-1});
   RunAndFilecheckHloRewrite(
-      hlo_text, GemmRewriter(CudaHopperOrRocmMI300(), /*f8_rewrite=*/true),
+      hlo_text,
+      GemmRewriter(CudaHopperOrRocmMI300(), GetToolkitVersion(),
+                   /*f8_rewrite=*/true),
       R"(
 ; CHECK-LABEL: ENTRY %test ({{.*}}: <<F8E4M3>>[16,32], {{.*}}: <<F8E4M3>>[32,16], {{.*}}: f32[]) -> f32[16,16] {
 ; CHECK:         [[P0:%[^ ]+]] = <<F8E4M3>>[16,32]{1,0} parameter(0)
@@ -6091,7 +6146,9 @@ TEST_P(ParameterizedFp8GemmRewriteTest, UnscaledABScaledF32DMatrixBiasF8) {
 
   CheckFp8IfSupported(hlo_text, ErrorSpec{1e-2, 1e-1});
   RunAndFilecheckHloRewrite(
-      hlo_text, GemmRewriter(CudaHopperOrRocmMI300(), /*f8_rewrite=*/true),
+      hlo_text,
+      GemmRewriter(CudaHopperOrRocmMI300(), GetToolkitVersion(),
+                   /*f8_rewrite=*/true),
       R"(
 ; CHECK-LABEL: ENTRY %test ({{.*}}: <<F8E4M3>>[16,32], {{.*}}: <<F8E4M3>>[32,16], {{.*}}: f32[]) -> f32[16,16] {
 ; CHECK:         [[P0:%[^ ]+]] = <<F8E4M3>>[16,32]{1,0} parameter(0)
@@ -6163,7 +6220,9 @@ TEST_P(ParameterizedFp8GemmRewriteTest, ScaledABScaledDF8) {
 
   CheckFp8IfSupported(hlo_text);
   RunAndFilecheckHloRewrite(
-      hlo_text, GemmRewriter(CudaHopperOrRocmMI300(), /*f8_rewrite=*/true),
+      hlo_text,
+      GemmRewriter(CudaHopperOrRocmMI300(), GetToolkitVersion(),
+                   /*f8_rewrite=*/true),
       R"(
 ; CHECK-LABEL: ENTRY %test ({{.*}}: <<F8E4M3>>[16,32], {{.*}}: <<F8E4M3>>[32,16], {{.*}}: f32[], {{.*}}: f32[], {{.*}}: f32[]) -> <<F8E4M3>>[16,16] {
 ; CHECK:         [[P0:%[^ ]+]] = <<F8E4M3>>[16,32]{1,0} parameter(0)
@@ -6235,7 +6294,9 @@ TEST_P(ParameterizedFp8GemmRewriteTest, ScaledABInvScaledDF8) {
 
   CheckFp8IfSupported(hlo_text);
   RunAndFilecheckHloRewrite(
-      hlo_text, GemmRewriter(CudaHopperOrRocmMI300(), /*f8_rewrite=*/true),
+      hlo_text,
+      GemmRewriter(CudaHopperOrRocmMI300(), GetToolkitVersion(),
+                   /*f8_rewrite=*/true),
       R"(
 
 ; CHECK-NOT:     divide
@@ -6285,7 +6346,9 @@ TEST_P(ParameterizedFp8GemmRewriteTest, ScaledABScaledDReluActivationF8) {
 
   CheckFp8IfSupported(hlo_text);
   RunAndFilecheckHloRewrite(
-      hlo_text, GemmRewriter(CudaHopperOrRocmMI300(), /*f8_rewrite=*/true),
+      hlo_text,
+      GemmRewriter(CudaHopperOrRocmMI300(), GetToolkitVersion(),
+                   /*f8_rewrite=*/true),
       R"(
 ; CHECK-LABEL: ENTRY %test ({{.*}}: <<F8E4M3>>[16,32], {{.*}}: <<F8E4M3>>[32,16], {{.*}}: f32[], {{.*}}: f32[], {{.*}}: f32[]) -> <<F8E4M3>>[16,16] {
 ; CHECK:         [[P0:%[^ ]+]] = <<F8E4M3>>[16,32]{1,0} parameter(0)
@@ -6359,7 +6422,9 @@ TEST_P(ParameterizedFp8GemmRewriteTest, ScaledABScaledDMatrixBiasF8) {
 
   CheckFp8IfSupported(hlo_text, ErrorSpec{0.1, 0.1});
   RunAndFilecheckHloRewrite(
-      hlo_text, GemmRewriter(CudaHopperOrRocmMI300(), /*f8_rewrite=*/true),
+      hlo_text,
+      GemmRewriter(CudaHopperOrRocmMI300(), GetToolkitVersion(),
+                   /*f8_rewrite=*/true),
       R"(
 
 ; CHECK-LABEL: ENTRY %test ({{.*}}: <<F8E4M3>>[16,32], {{.*}}: <<F8E4M3>>[32,16], {{.*}}: f16[16,16], {{.*}}: f16[], {{.*}}: f16[], {{.*}}: f16[]) -> <<F8E4M3>>[16,16] {
@@ -6434,7 +6499,9 @@ TEST_P(ParameterizedFp8GemmRewriteTest, ScaledABScaledDVectorBiasF8) {
 
   CheckFp8IfSupported(hlo_text, ErrorSpec{0.1, 0.1});
   RunAndFilecheckHloRewrite(
-      hlo_text, GemmRewriter(CudaHopperOrRocmMI300(), /*f8_rewrite=*/true),
+      hlo_text,
+      GemmRewriter(CudaHopperOrRocmMI300(), GetToolkitVersion(),
+                   /*f8_rewrite=*/true),
       R"(
 
 ; CHECK-LABEL: ENTRY %test ({{.*}}: <<F8E4M3>>[16,32], {{.*}}: <<F8E4M3>>[32,16], {{.*}}: f16[16], {{.*}}: f16[], {{.*}}: f16[], {{.*}}: f16[]) -> <<F8E4M3>>[16,16] {
@@ -6507,7 +6574,9 @@ TEST_P(ParameterizedFp8GemmRewriteTest, ScaledABUnscaledDF32VectorBiasF8) {
 
   CheckFp8IfSupported(hlo_text);
   RunAndFilecheckHloRewrite(
-      hlo_text, GemmRewriter(CudaHopperOrRocmMI300(), /*f8_rewrite=*/true),
+      hlo_text,
+      GemmRewriter(CudaHopperOrRocmMI300(), GetToolkitVersion(),
+                   /*f8_rewrite=*/true),
       R"(
 ; CHECK-LABEL: ENTRY %test ({{.*}}: <<F8E4M3>>[16,32], {{.*}}: <<F8E4M3>>[32,16], {{.*}}: f32[16], {{.*}}: f32[], {{.*}}: f32[]) -> f32[16,16] {
 ; CHECK:         [[P0:%[^ ]+]] = <<F8E4M3>>[16,32]{1,0} parameter(0)
@@ -6574,7 +6643,9 @@ TEST_P(ParameterizedFp8GemmRewriteTest,
 
   CheckFp8IfSupported(hlo_text, ErrorSpec{2e-3, 0.});
   RunAndFilecheckHloRewrite(
-      hlo_text, GemmRewriter(CudaHopperOrRocmMI300(), /*f8_rewrite=*/true),
+      hlo_text,
+      GemmRewriter(CudaHopperOrRocmMI300(), GetToolkitVersion(),
+                   /*f8_rewrite=*/true),
       R"(
 ; CHECK-LABEL: ENTRY %test ({{.*}}: <<F8E4M3>>[16,32], {{.*}}: <<F8E4M3>>[32,16], {{.*}}: f16[16], {{.*}}: f16[], {{.*}}: f16[]) -> f16[16,16] {
 ; CHECK-NEXT:    [[P0:%[^ ]+]] = <<F8E4M3>>[16,32]{1,0} parameter(0)
@@ -6640,7 +6711,8 @@ TEST_P(ParameterizedFp8GemmRewriteTest, Rank3ScaledABUnscaledDVectorBiasF8) {
 
   TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
                           ParseAndReturnVerifiedModule(hlo_text));
-  GemmRewriter pass(CudaHopperOrRocmMI300(), /*f8_rewrite=*/true);
+  GemmRewriter pass(CudaHopperOrRocmMI300(), GetToolkitVersion(),
+                    /*f8_rewrite=*/true);
   TF_ASSERT_OK_AND_ASSIGN(bool changed, this->RunHloPass(&pass, module.get()));
   EXPECT_TRUE(changed);
 
@@ -6652,7 +6724,9 @@ TEST_P(ParameterizedFp8GemmRewriteTest, Rank3ScaledABUnscaledDVectorBiasF8) {
                      .WithShape(F16, {4, 16, 32})));
 
   RunAndFilecheckHloRewrite(
-      hlo_text, GemmRewriter(CudaHopperOrRocmMI300(), /*f8_rewrite=*/true),
+      hlo_text,
+      GemmRewriter(CudaHopperOrRocmMI300(), GetToolkitVersion(),
+                   /*f8_rewrite=*/true),
       R"(
 ; CHECK-LABEL: ENTRY %test ({{.*}}: <<F8E4M3>>[4,16,16], {{.*}}: <<F8E4M3>>[16,32], {{.*}}: f32[32], {{.*}}: f16[], {{.*}}: f16[]) -> f16[4,16,32] {
 ; CHECK-NEXT:    [[P0:%[^ ]+]] = <<F8E4M3>>[4,16,16]{2,1,0} parameter(0)
@@ -6723,7 +6797,8 @@ TEST_P(ParameterizedFp8GemmRewriteTest,
 
   TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
                           ParseAndReturnVerifiedModule(hlo_text));
-  GemmRewriter pass(CudaHopperOrRocmMI300(), /*f8_rewrite=*/true);
+  GemmRewriter pass(CudaHopperOrRocmMI300(), GetToolkitVersion(),
+                    /*f8_rewrite=*/true);
   TF_ASSERT_OK_AND_ASSIGN(bool changed, this->RunHloPass(&pass, module.get()));
   EXPECT_TRUE(changed);
 
@@ -6737,7 +6812,9 @@ TEST_P(ParameterizedFp8GemmRewriteTest,
               .WithShape(F16, {4, 15, 31})));
 
   RunAndFilecheckHloRewrite(
-      hlo_text, GemmRewriter(CudaHopperOrRocmMI300(), /*f8_rewrite=*/true),
+      hlo_text,
+      GemmRewriter(CudaHopperOrRocmMI300(), GetToolkitVersion(),
+                   /*f8_rewrite=*/true),
       R"(
 ; CHECK-LABEL: ENTRY %test ({{.*}}: <<F8E4M3>>[4,15,15], {{.*}}: <<F8E4M3>>[15,31], {{.*}}: f32[31], {{.*}}: f16[], {{.*}}: f16[]) -> f16[4,15,31] {
 ; CHECK-NEXT:    [[P0:%[^ ]+]] = <<F8E4M3>>[4,15,15]{2,1,0} parameter(0)
@@ -6812,7 +6889,8 @@ TEST_P(ParameterizedFp8GemmRewriteTest, Rank3ScaledABUnscaledDMatrixBiasF8) {
 
   TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
                           ParseAndReturnVerifiedModule(hlo_text));
-  GemmRewriter pass(CudaHopperOrRocmMI300(), /*f8_rewrite=*/true);
+  GemmRewriter pass(CudaHopperOrRocmMI300(), GetToolkitVersion(),
+                    /*f8_rewrite=*/true);
   TF_ASSERT_OK_AND_ASSIGN(bool changed, this->RunHloPass(&pass, module.get()));
   EXPECT_TRUE(changed);
 
@@ -6824,7 +6902,9 @@ TEST_P(ParameterizedFp8GemmRewriteTest, Rank3ScaledABUnscaledDMatrixBiasF8) {
                      .WithShape(F32, {4, 16, 32})));
 
   RunAndFilecheckHloRewrite(
-      hlo_text, GemmRewriter(CudaHopperOrRocmMI300(), /*f8_rewrite=*/true),
+      hlo_text,
+      GemmRewriter(CudaHopperOrRocmMI300(), GetToolkitVersion(),
+                   /*f8_rewrite=*/true),
       R"(
 ; CHECK-LABEL: ENTRY %test ({{.*}}: <<F8E4M3>>[4,16,16], {{.*}}: <<F8E4M3>>[16,32], {{.*}}: f32[4,16,32], {{.*}}: f32[], {{.*}}: f32[]) -> f32[4,16,32] {
 ; CHECK-NEXT:    [[P0:%[^ ]+]] = <<F8E4M3>>[4,16,16]{2,1,0} parameter(0)
@@ -6891,7 +6971,8 @@ TEST_P(ParameterizedFp8GemmRewriteTest,
 
   TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
                           ParseAndReturnVerifiedModule(hlo_text));
-  GemmRewriter pass(CudaHopperOrRocmMI300(), /*f8_rewrite=*/true);
+  GemmRewriter pass(CudaHopperOrRocmMI300(), GetToolkitVersion(),
+                    /*f8_rewrite=*/true);
   TF_ASSERT_OK_AND_ASSIGN(bool changed, this->RunHloPass(&pass, module.get()));
   EXPECT_TRUE(changed);
 
@@ -6905,7 +6986,9 @@ TEST_P(ParameterizedFp8GemmRewriteTest,
               .WithShape(F32, {3, 15, 31})));
 
   RunAndFilecheckHloRewrite(
-      hlo_text, GemmRewriter(CudaHopperOrRocmMI300(), /*f8_rewrite=*/true),
+      hlo_text,
+      GemmRewriter(CudaHopperOrRocmMI300(), GetToolkitVersion(),
+                   /*f8_rewrite=*/true),
       R"(
 ; CHECK-LABEL: ENTRY %test ({{.*}}: <<F8E4M3>>[3,15,15], {{.*}}: <<F8E4M3>>[15,31], {{.*}}: f32[3,15,31], {{.*}}: f32[], {{.*}}: f32[]) -> f32[3,15,31] {
 ; CHECK-NEXT:    [[P0:%[^ ]+]] = <<F8E4M3>>[3,15,15]{2,1,0} parameter(0)
@@ -6980,12 +7063,15 @@ TEST_P(ParameterizedFp8GemmRewriteTest,
 
   TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
                           ParseAndReturnVerifiedModule(hlo_text));
-  GemmRewriter pass(CudaHopperOrRocmMI300(), /*f8_rewrite=*/true);
+  GemmRewriter pass(CudaHopperOrRocmMI300(), GetToolkitVersion(),
+                    /*f8_rewrite=*/true);
   TF_ASSERT_OK_AND_ASSIGN(bool changed, this->RunHloPass(&pass, module.get()));
   EXPECT_TRUE(changed);
 
   RunAndFilecheckHloRewrite(
-      hlo_text, GemmRewriter(CudaHopperOrRocmMI300(), /*f8_rewrite=*/true),
+      hlo_text,
+      GemmRewriter(CudaHopperOrRocmMI300(), GetToolkitVersion(),
+                   /*f8_rewrite=*/true),
       R"(
 ; CHECK-LABEL: ENTRY %test ({{.*}}: <<F8E4M3>>[48,16], {{.*}}: <<F8E4M3>>[16,32], {{.*}}: f32[32,16], {{.*}}: f32[], {{.*}}: f32[]) -> f32[32,16] {
 ; CHECK-NEXT:    [[P0:%[^ ]+]] = <<F8E4M3>>[48,16]{1,0} parameter(0)
@@ -7052,7 +7138,9 @@ TEST_P(ParameterizedFp8GemmRewriteTest, ScaledABUnscaledDWithAllGatherF8) {
   config.set_num_partitions(8);
 
   RunAndFilecheckHloRewrite(
-      hlo_text, GemmRewriter(CudaHopperOrRocmMI300(), /*f8_rewrite=*/true),
+      hlo_text,
+      GemmRewriter(CudaHopperOrRocmMI300(), GetToolkitVersion(),
+                   /*f8_rewrite=*/true),
       R"(
 ; CHECK-LABEL: ENTRY %test ({{.*}}: <<F8E4M3>>[16,32], {{.*}}: <<F8E4M3>>[16,32], {{.*}}: f32[], {{.*}}: f32[]) -> f32[16,32] {
 ; CHECK:         [[P0:%[^ ]+]] = <<F8E4M3>>[16,32]{1,0} parameter(0)
@@ -7118,7 +7206,9 @@ TEST_P(ParameterizedFp8GemmRewriteTest, ScaledABUnscaledDWithAllToAllF8) {
   config.set_num_partitions(8);
 
   RunAndFilecheckHloRewrite(
-      hlo_text, GemmRewriter(CudaHopperOrRocmMI300(), /*f8_rewrite=*/true),
+      hlo_text,
+      GemmRewriter(CudaHopperOrRocmMI300(), GetToolkitVersion(),
+                   /*f8_rewrite=*/true),
       R"(
 ; CHECK-LABEL: ENTRY %test ({{.*}}: <<F8E4M3>>[16,32], {{.*}}: <<F8E4M3>>[16,32], {{.*}}: f32[], {{.*}}: f32[]) -> f32[16,16] {
 ; CHECK:         [[P0:%[^ ]+]] = <<F8E4M3>>[16,32]{1,0} parameter(0)
@@ -7182,7 +7272,9 @@ TEST_P(ParameterizedFp8GemmRewriteTest,
   config.set_num_partitions(8);
 
   RunAndFilecheckHloRewrite(
-      hlo_text, GemmRewriter(CudaHopperOrRocmMI300(), /*f8_rewrite=*/true),
+      hlo_text,
+      GemmRewriter(CudaHopperOrRocmMI300(), GetToolkitVersion(),
+                   /*f8_rewrite=*/true),
       R"(
 ; CHECK-LABEL: ENTRY %test ({{.*}}: <<F8E4M3>>[16,32], {{.*}}: <<F8E4M3>>[16,32], {{.*}}: f32[], {{.*}}: f32[]) -> f32[16,16] {
 ; CHECK:         [[P0:%[^ ]+]] = <<F8E4M3>>[16,32]{1,0} parameter(0)
@@ -7248,7 +7340,9 @@ TEST_P(ParameterizedFp8GemmRewriteTest,
 
   CheckFp8IfSupported(hlo_text, ErrorSpec{2e-3, 0.});
   RunAndFilecheckHloRewrite(
-      hlo_text, GemmRewriter(CudaHopperOrRocmMI300(), /*f8_rewrite=*/true),
+      hlo_text,
+      GemmRewriter(CudaHopperOrRocmMI300(), GetToolkitVersion(),
+                   /*f8_rewrite=*/true),
       R"(
 ; CHECK-LABEL:   ENTRY %test ({{.*}}: <<F8E4M3>>[16,32], {{.*}}: <<F8E4M3>>[32,16], {{.*}}: f16[16], {{.*}}: f16[16,16], {{.*}}: f16[], {{.*}}: f16[]) -> f16[16,16] {
 ; CHECK-DAG:     [[P0:%[^ ]+]] = <<F8E4M3>>[16,32]{1,0} parameter(0)
@@ -7333,7 +7427,9 @@ TEST_P(ParameterizedFp8GemmRewriteTest, ScaledABScaledDWithDAmaxF8) {
 
   CheckFp8IfSupported(hlo_text);
   RunAndFilecheckHloRewrite(
-      hlo_text, GemmRewriter(CudaHopperOrRocmMI300(), /*f8_rewrite=*/true),
+      hlo_text,
+      GemmRewriter(CudaHopperOrRocmMI300(), GetToolkitVersion(),
+                   /*f8_rewrite=*/true),
       R"(
 ; CHECK-LABEL: ENTRY %test ({{.*}}: <<F8E4M3>>[16,32], {{.*}}: <<F8E4M3>>[32,16], {{.*}}: f32[], {{.*}}: f32[], {{.*}}: f32[]) -> (<<F8E4M3>>[16,16], f32[]) {
 ; CHECK:    [[P0:%[^ ]+]] = <<F8E4M3>>[16,32]{1,0} parameter(0)
@@ -7418,7 +7514,9 @@ TEST_P(ParameterizedFp8GemmRewriteTest,
 
   CheckFp8IfSupported(hlo_text);
   RunAndFilecheckHloRewrite(
-      hlo_text, GemmRewriter(CudaHopperOrRocmMI300(), /*f8_rewrite=*/true),
+      hlo_text,
+      GemmRewriter(CudaHopperOrRocmMI300(), GetToolkitVersion(),
+                   /*f8_rewrite=*/true),
       R"(
 ; CHECK-LABEL: ENTRY %test ({{.*}}: <<F8E4M3>>[16,32], {{.*}}: <<F8E4M3>>[32,16], {{.*}}: f16[], {{.*}}: f16[], {{.*}}: f16[]) -> (<<F8E4M3>>[16,16], f16[]) {
 ; CHECK:    [[P0:%[^ ]+]] = <<F8E4M3>>[16,32]{1,0} parameter(0)
@@ -7506,7 +7604,9 @@ TEST_P(ParameterizedFp8GemmRewriteTest,
 
   CheckFp8IfSupported(hlo_text);
   RunAndFilecheckHloRewrite(
-      hlo_text, GemmRewriter(CudaHopperOrRocmMI300(), /*f8_rewrite=*/true),
+      hlo_text,
+      GemmRewriter(CudaHopperOrRocmMI300(), GetToolkitVersion(),
+                   /*f8_rewrite=*/true),
       R"(
 ; CHECK-LABEL: ENTRY %test ({{.*}}: <<F8E4M3>>[16,32], {{.*}}: <<F8E4M3>>[32,16], {{.*}}: f32[], {{.*}}: f32[], {{.*}}: f32[]) -> (<<F8E4M3>>[16,16], f32[]) {
 ; CHECK:    [[P0:%[^ ]+]] = <<F8E4M3>>[16,32]{1,0} parameter(0)
@@ -7645,7 +7745,9 @@ TEST_P(ParameterizedFp8GemmRewriteTest, ScaledABUnscaledDF8Parameterized) {
     CheckFp8IfSupported(hlo_text);
 
     RunAndFilecheckHloRewrite(
-        hlo_text, GemmRewriter(CudaHopperOrRocmMI300(), /*f8_rewrite=*/true),
+        hlo_text,
+        GemmRewriter(CudaHopperOrRocmMI300(), GetToolkitVersion(),
+                     /*f8_rewrite=*/true),
         R"(
     ; CHECK:           custom_call_target="__cublas$lt$matmul$f8",
           )");
@@ -7719,7 +7821,9 @@ ENTRY f {
     CheckFp8IfSupported(hlo_text);
 
     RunAndFilecheckHloRewrite(
-        hlo_text, GemmRewriter(CudaHopperOrRocmMI300(), /*f8_rewrite=*/true),
+        hlo_text,
+        GemmRewriter(CudaHopperOrRocmMI300(), GetToolkitVersion(),
+                     /*f8_rewrite=*/true),
         R"(
     ; CHECK:           custom_call_target="__cublas$lt$matmul$f8",
           )");
@@ -7756,7 +7860,9 @@ TEST_P(ParameterizedFp8GemmRewriteTest, ScaledABUnscaledDF8TF32E5M2) {
 
   CheckFp8IfSupported(hlo_text);
   RunAndFilecheckHloRewrite(
-      hlo_text, GemmRewriter(CudaHopperOrRocmMI300(), /*f8_rewrite=*/true),
+      hlo_text,
+      GemmRewriter(CudaHopperOrRocmMI300(), GetToolkitVersion(),
+                   /*f8_rewrite=*/true),
       R"(
     ; CHECK:           custom_call_target="__cublas$lt$matmul$f8",
           )");
@@ -7792,14 +7898,17 @@ TEST_P(ParameterizedFp8GemmRewriteTest, FnuzTypeF8) {
 #if GOOGLE_CUDA
   TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
                           ParseAndReturnVerifiedModule(hlo_text));
-  GemmRewriter pass(CudaHopperOrRocmMI300(), /*f8_rewrite=*/true);
+  GemmRewriter pass(CudaHopperOrRocmMI300(), GetToolkitVersion(),
+                    /*f8_rewrite=*/true);
   TF_ASSERT_OK_AND_ASSIGN(bool changed, this->RunHloPass(&pass, module.get()));
   EXPECT_FALSE(changed);
 #endif
 #if TENSORFLOW_USE_ROCM
   EXPECT_TRUE(RunAndCompare(hlo_text, ErrorSpec{1e-2, 1e-2}));
   RunAndFilecheckHloRewrite(
-      hlo_text, GemmRewriter(CudaHopperOrRocmMI300(), /*f8_rewrite=*/true),
+      hlo_text,
+      GemmRewriter(CudaHopperOrRocmMI300(), GetToolkitVersion(),
+                   /*f8_rewrite=*/true),
       R"(
 ; CHECK-LABEL: ENTRY %test ({{.*}}: f8e4m3fnuz[16,32], {{.*}}: f8e4m3fnuz[32,16], {{.*}}: f32[], {{.*}}: f32[]) -> f32[16,16] {
 ; CHECK-NEXT:    [[P0:%[^ ]+]] = f8e4m3fnuz[16,32]{1,0} parameter(0)

From c01a4be55297bcbe9bf6801a30eacdbb416ce3d5 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Wed, 15 May 2024 09:33:22 -0700
Subject: [PATCH 188/478] Update ops-related pbtxt files.

PiperOrigin-RevId: 633982202
---
 tensorflow/core/ops/ops.pbtxt | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorflow/core/ops/ops.pbtxt b/tensorflow/core/ops/ops.pbtxt
index cdc0191b8ad6c9..975f4a22eca161 100644
--- a/tensorflow/core/ops/ops.pbtxt
+++ b/tensorflow/core/ops/ops.pbtxt
@@ -1,4 +1,4 @@
-go/debugproto   
+go/debugstr  
 op {
   name: "Abort"
   attr {

From 462d6fd50e2405383828bb419e3d0dcc80362371 Mon Sep 17 00:00:00 2001
From: Kyle Lucke <klucke@google.com>
Date: Wed, 15 May 2024 09:39:57 -0700
Subject: [PATCH 189/478] Use absl::Status instead of xla::Status now that
 they're identical.

PiperOrigin-RevId: 633984152
---
 .../python/ifrt_proxy/common/ifrt_service.proto  | 16 ++++++++--------
 1 file changed, 8 insertions(+), 8 deletions(-)

diff --git a/third_party/xla/xla/python/ifrt_proxy/common/ifrt_service.proto b/third_party/xla/xla/python/ifrt_proxy/common/ifrt_service.proto
index 0b3e1e65e9b304..6c37451979dbda 100644
--- a/third_party/xla/xla/python/ifrt_proxy/common/ifrt_service.proto
+++ b/third_party/xla/xla/python/ifrt_proxy/common/ifrt_service.proto
@@ -163,21 +163,21 @@ message ResponseMetadata {
   // ID of the operation this response belongs to.
   fixed64 op_id = 1;
 
-  // Status of the operation.
+  // absl::Status of the operation.
   //
   //    In case of "chunked" responses (i.e., the full logical response is
   //    spread across a sequence of IfrtResponse protos), the actual sequence of
-  //    IfrtResponse messages will follow only if this Status is OK in the very
-  //    first message. That is, in case of errors, server sends a single
+  //    IfrtResponse messages will follow only if this absl::Status is OK in the
+  //    very first message. That is, in case of errors, server sends a single
   //    IfrtResponse with the appropriate error included.
   //
   //    In case of "batched" operations (i.e., where the response is carrying
   //    the outcomes of multiple requests that were "batched" in the same
-  //    IfrtRequest proto - such as deleting a bunch of Arrays) this Status
-  //    field provides a way to quickly check if none of the individual
-  //    operations encountered errors. Clients should not rely on specific error
-  //    type or string when this is not OK, they should check the response
-  //    message for individual Statuses.
+  //    IfrtRequest proto - such as deleting a bunch of Arrays) this
+  //    absl::Status field provides a way to quickly check if none of the
+  //    individual operations encountered errors. Clients should not rely on
+  //    specific error type or string when this is not OK, they should check the
+  //    response message for individual absl::Statuses.
   tensorflow.StatusProto status = 2;
 }
 

From 537256146ebb061f00c711d322e2885408afa3c3 Mon Sep 17 00:00:00 2001
From: David Majnemer <majnemer@google.com>
Date: Wed, 15 May 2024 09:47:15 -0700
Subject: [PATCH 190/478] Use absl/numeric/bits.h for bit manipulation.

This reduces code duplication and allows us to reuse optimized code.

PiperOrigin-RevId: 633986138
---
 tensorflow/compiler/jit/BUILD                 |  1 +
 tensorflow/compiler/jit/device_util.h         | 17 +---
 .../gpu/gpu_bfc_allocator_test.cc             | 16 ----
 tensorflow/core/framework/BUILD               |  2 +-
 tensorflow/core/framework/types.h             | 50 ++---------
 tensorflow/core/lib/strings/BUILD             |  1 +
 tensorflow/core/lib/strings/ordered_code.cc   | 41 +--------
 .../tsl/tsl/framework/bfc_allocator.h         | 25 +-----
 .../xla/third_party/tsl/tsl/lib/core/BUILD    |  2 +-
 .../third_party/tsl/tsl/lib/core/bitmap.cc    | 41 ++-------
 .../xla/third_party/tsl/tsl/lib/core/bits.h   | 84 +++----------------
 11 files changed, 35 insertions(+), 245 deletions(-)

diff --git a/tensorflow/compiler/jit/BUILD b/tensorflow/compiler/jit/BUILD
index 76f3c147903748..36552f4ae3027e 100644
--- a/tensorflow/compiler/jit/BUILD
+++ b/tensorflow/compiler/jit/BUILD
@@ -1149,6 +1149,7 @@ cc_library(
         "@com_google_absl//absl/algorithm:container",
         "@com_google_absl//absl/container:flat_hash_map",
         "@com_google_absl//absl/container:flat_hash_set",
+        "@com_google_absl//absl/numeric:bits",
         "@com_google_absl//absl/strings",
         "@com_google_absl//absl/types:span",
         "@local_xla//xla:status_macros",
diff --git a/tensorflow/compiler/jit/device_util.h b/tensorflow/compiler/jit/device_util.h
index df3b7d04fbfe7b..ec4d9484ae8854 100644
--- a/tensorflow/compiler/jit/device_util.h
+++ b/tensorflow/compiler/jit/device_util.h
@@ -20,6 +20,7 @@ limitations under the License.
 #include <memory>
 
 #include "absl/container/flat_hash_map.h"
+#include "absl/numeric/bits.h"
 #include "absl/strings/string_view.h"
 #include "absl/types/span.h"
 #include "tensorflow/compiler/tf2xla/xla_op_registry.h"
@@ -79,7 +80,7 @@ class DeviceSet {
         uint64 only_lowest_bit_set = word & -word;
         // The number of trailing zeros in a non-zero word is the index of the
         // least significant 1.
-        int bit_index = ctz_uint64(word);
+        int bit_index = absl::countr_zero(word);
         if (!func(DeviceId(word_index * kWordSize + bit_index))) {
           return;
         }
@@ -89,20 +90,6 @@ class DeviceSet {
   }
 
  private:
-  static int ctz_uint64(uint64 x) {
-    DCHECK_NE(x, 0);
-#ifdef __GNUC__
-    return __builtin_ctzl(x);
-#else
-    int result = 0u;
-    while ((x & 1u) == 0u) {
-      x >>= 1;
-      ++result;
-    }
-    return result;
-#endif
-  }
-
   absl::InlinedVector<uint64, 1> storage_;
 
   const int kWordSize = 64;
diff --git a/tensorflow/core/common_runtime/gpu/gpu_bfc_allocator_test.cc b/tensorflow/core/common_runtime/gpu/gpu_bfc_allocator_test.cc
index 48aba783f1f526..b291186f0fba96 100644
--- a/tensorflow/core/common_runtime/gpu/gpu_bfc_allocator_test.cc
+++ b/tensorflow/core/common_runtime/gpu/gpu_bfc_allocator_test.cc
@@ -512,18 +512,6 @@ class GPUBFCAllocatorPrivateMethodsTest
     }
   }
 
-  void TestLog2FloorNonZeroSlow() {
-    GPUBFCAllocator a(GetParam()(1ull << 32), 1 /* total_memory */, "GPU_0_bfc",
-                      {});
-    EXPECT_EQ(-1, a.Log2FloorNonZeroSlow(0));
-    EXPECT_EQ(0, a.Log2FloorNonZeroSlow(1));
-    EXPECT_EQ(1, a.Log2FloorNonZeroSlow(2));
-    EXPECT_EQ(1, a.Log2FloorNonZeroSlow(3));
-    EXPECT_EQ(9, a.Log2FloorNonZeroSlow(1023));
-    EXPECT_EQ(10, a.Log2FloorNonZeroSlow(1024));
-    EXPECT_EQ(10, a.Log2FloorNonZeroSlow(1025));
-  }
-
   void TestForceAllowGrowth() {
     // Unset flag value uses provided option.
     unsetenv("TF_FORCE_GPU_ALLOW_GROWTH");
@@ -563,10 +551,6 @@ class GPUBFCAllocatorPrivateMethodsTest
 
 TEST_P(GPUBFCAllocatorPrivateMethodsTest, BinDebugInfo) { TestBinDebugInfo(); }
 
-TEST_P(GPUBFCAllocatorPrivateMethodsTest, Log2FloorNonZeroSlow) {
-  TestLog2FloorNonZeroSlow();
-}
-
 TEST_P(GPUBFCAllocatorPrivateMethodsTest, ForceAllowGrowth) {
   TestForceAllowGrowth();
 }
diff --git a/tensorflow/core/framework/BUILD b/tensorflow/core/framework/BUILD
index 617404f2722585..2051a3d14ea024 100644
--- a/tensorflow/core/framework/BUILD
+++ b/tensorflow/core/framework/BUILD
@@ -873,8 +873,8 @@ tf_cuda_library(
         "//tensorflow/core/platform:tensor_coding",
         "//tensorflow/core/platform:types",
         "//tensorflow/core/public:version",
-        "//tensorflow/core/util:managed_stack_trace",
         "@com_google_absl//absl/memory",
+        "@com_google_absl//absl/numeric:bits",
         "@com_google_absl//absl/strings",
         "@eigen_archive//:eigen3",
         "@local_tsl//tsl/framework:device_type",
diff --git a/tensorflow/core/framework/types.h b/tensorflow/core/framework/types.h
index 90e148a2e0784b..142ac9913ddc93 100644
--- a/tensorflow/core/framework/types.h
+++ b/tensorflow/core/framework/types.h
@@ -16,10 +16,12 @@ limitations under the License.
 #ifndef TENSORFLOW_CORE_FRAMEWORK_TYPES_H_
 #define TENSORFLOW_CORE_FRAMEWORK_TYPES_H_
 
+#include <cstddef>
 #include <map>
 #include <set>
 #include <string>
 
+#include "absl/numeric/bits.h"
 #include "unsupported/Eigen/CXX11/Tensor"  // from @eigen_archive
 #include "tensorflow/core/framework/bfloat16.h"
 #include "tensorflow/core/framework/full_type.pb.h"
@@ -124,7 +126,7 @@ class DataTypeSet {
       if (pos_ < kNumBits) {
         uint32 remaining_mask = set_.mask_ >> pos_;
         if (remaining_mask != 0u) {
-          pos_ += ctz_uint32(remaining_mask);
+          pos_ += absl::countr_zero(remaining_mask);
         }
       }
       DCHECK_LE(pos_, kNumBits);
@@ -137,39 +139,11 @@ class DataTypeSet {
     }
   };
 
-  static uint32 ctz_uint32(uint32 x) {
-    DCHECK_NE(x, 0u);
-#ifdef __GNUC__
-    return __builtin_ctz(x);
-#else
-    uint32 n = 0u;
-    while ((x & 1u) == 0u) {
-      x >>= 1;
-      ++n;
-    }
-    return n;
-#endif
-  }
-
-  static uint32 clz_uint32(uint32 x) {
-    DCHECK_NE(x, 0u);
-#ifdef __GNUC__
-    return __builtin_clz(x);
-#else
-    uint32 n = 0u;
-    while ((x >> (kNumBits - 1u)) == 0u) {
-      x <<= 1;
-      ++n;
-    }
-    return n;
-#endif
-  }
-
   Iterator begin() const {
     // The begin position is the index of the first bit set to 1 in the entire
     // bit mask. If there are no bits set to 1, then the index is 0.
     if (mask_ != 0) {
-      return Iterator(*this, ctz_uint32(mask_));
+      return Iterator(*this, absl::countr_zero(mask_));
     }
     // The set is empty.
     return Iterator(*this, 0);
@@ -179,25 +153,13 @@ class DataTypeSet {
     // The end position is the index of the highest bit that is set, plus 1.
     // If there are no bits set to 1, then the index is 0.
     if (mask_ != 0) {
-      return Iterator(*this, kNumBits - clz_uint32(mask_));
+      return Iterator(*this, kNumBits - absl::countl_zero(mask_));
     }
     // The set is empty.
     return Iterator(*this, 0);
   }
 
-  size_t size() const {
-#if defined(__GNUC__)
-    return __builtin_popcount(mask_);
-#else
-    size_t n = 0;
-    uint32 x = mask_;
-    while (x > 0) {
-      n += x & 1u;
-      x >>= 1;
-    }
-    return n;
-#endif
-  }
+  size_t size() const { return absl::popcount(mask_); }
 
   constexpr DataTypeSet operator|(const DataTypeSet& other) const {
     return DataTypeSet(mask_ | other.mask_);
diff --git a/tensorflow/core/lib/strings/BUILD b/tensorflow/core/lib/strings/BUILD
index 2f04ae43b027a3..72eb0a6dac308c 100644
--- a/tensorflow/core/lib/strings/BUILD
+++ b/tensorflow/core/lib/strings/BUILD
@@ -43,6 +43,7 @@ cc_library(
         "//tensorflow/core/platform:macros",
         "//tensorflow/core/platform:stringpiece",
         "//tensorflow/core/platform:types",
+        "@local_tsl//tsl/lib/core:bits",
     ],
 )
 
diff --git a/tensorflow/core/lib/strings/ordered_code.cc b/tensorflow/core/lib/strings/ordered_code.cc
index 8d3f0a66d7977d..65210dc28d1467 100644
--- a/tensorflow/core/lib/strings/ordered_code.cc
+++ b/tensorflow/core/lib/strings/ordered_code.cc
@@ -20,6 +20,7 @@ limitations under the License.
 
 #include "tensorflow/core/platform/logging.h"
 #include "tensorflow/core/platform/stringpiece.h"
+#include "tsl/lib/core/bits.h"
 
 namespace tensorflow {
 namespace strings {
@@ -403,49 +404,15 @@ static const uint64 kLengthToMask[1 + kMaxSigned64Length] = {
 // For positive numbers, the number of bits is 1 plus the most significant
 // bit position (the highest bit position in a positive int64 is 63).
 // For a negative number n, we count the bits in ~n.
-// That is, length = kBitsToLength[Bits::Log2Floor64(n < 0 ? ~n : n) + 1].
+// That is, length = kBitsToLength[tsl::Log2Floor64(n < 0 ? ~n : n) + 1].
 static const int8 kBitsToLength[1 + 63] = {
     1, 1, 1, 1, 1, 1, 1, 2, 2, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3, 3, 3, 4,
     4, 4, 4, 4, 4, 4, 5, 5, 5, 5, 5, 5, 5, 6, 6, 6, 6, 6, 6, 6, 7, 7,
     7, 7, 7, 7, 7, 8, 8, 8, 8, 8, 8, 8, 9, 9, 9, 9, 9, 9, 9, 10};
 
-#if defined(__GNUC__)
-// Returns floor(lg(n)).  Returns -1 if n == 0.
-static int Log2Floor64(uint64 n) {
-  return n == 0 ? -1 : 63 ^ __builtin_clzll(n);
-}
-#else
-// Portable slow version
-static int Log2Floor32_Portable(uint32 n) {
-  if (n == 0) return -1;
-  int log = 0;
-  uint32 value = n;
-  for (int i = 4; i >= 0; --i) {
-    int shift = (1 << i);
-    uint32 x = value >> shift;
-    if (x != 0) {
-      value = x;
-      log += shift;
-    }
-  }
-  assert(value == 1);
-  return log;
-}
-// Returns floor(lg(n)).  Returns -1 if n == 0.
-static int Log2Floor64(uint64 n) {
-  const uint32 topbits = static_cast<uint32>(n >> 32);
-  if (topbits == 0) {
-    // Top bits are zero, so scan in bottom bits
-    return Log2Floor32_Portable(static_cast<uint32>(n));
-  } else {
-    return 32 + Log2Floor32_Portable(topbits);
-  }
-}
-#endif
-
 // Calculates the encoding length in bytes of the signed number n.
 static inline int SignedEncodingLength(int64_t n) {
-  return kBitsToLength[Log2Floor64(n < 0 ? ~n : n) + 1];
+  return kBitsToLength[tsl::Log2Floor64(n < 0 ? ~n : n) + 1];
 }
 
 static void StoreBigEndian64(char* dst, uint64 v) {
@@ -494,7 +461,7 @@ bool OrderedCode::ReadSignedNumIncreasing(StringPiece* src, int64_t* result) {
   int len;
   uint64 x;
   if (first_byte != 0xff) {
-    len = 7 - Log2Floor64(first_byte ^ 0xff);
+    len = 7 - tsl::Log2Floor64(first_byte ^ 0xff);
     if (src->size() < static_cast<size_t>(len)) return false;
     x = xor_mask;  // sign extend using xor_mask
     for (int i = 0; i < len; ++i)
diff --git a/third_party/xla/third_party/tsl/tsl/framework/bfc_allocator.h b/third_party/xla/third_party/tsl/tsl/framework/bfc_allocator.h
index 76921c5f04a79d..6ea913059c472a 100644
--- a/third_party/xla/third_party/tsl/tsl/framework/bfc_allocator.h
+++ b/third_party/xla/third_party/tsl/tsl/framework/bfc_allocator.h
@@ -27,6 +27,7 @@ limitations under the License.
 #include "tsl/framework/allocator.h"
 #include "tsl/framework/allocator_retry.h"
 #include "tsl/framework/shared_counter.h"
+#include "tsl/lib/core/bits.h"
 #include "tsl/platform/macros.h"
 #include "tsl/platform/mutex.h"
 #include "tsl/platform/numbers.h"
@@ -536,28 +537,6 @@ class BFCAllocator : public Allocator {
   // Structures immutable after construction
   size_t memory_limit_ = 0;
 
-  inline int Log2FloorNonZeroSlow(uint64 n) {
-    int r = 0;
-    while (n > 0) {
-      r++;
-      n >>= 1;
-    }
-    return r - 1;
-  }
-
-  // Returns floor(log2(n)).
-  inline int Log2FloorNonZero(uint64 n) {
-#if defined(__GNUC__)
-    return 63 ^ __builtin_clzll(n);
-#elif defined(PLATFORM_WINDOWS) && (_WIN64)
-    unsigned long index;
-    _BitScanReverse64(&index, n);
-    return index;
-#else
-    return Log2FloorNonZeroSlow(n);
-#endif
-  }
-
   // Map from bin size to Bin
   Bin* BinFromIndex(BinNum index) {
     return reinterpret_cast<Bin*>(&(bins_space_[index * sizeof(Bin)]));
@@ -567,7 +546,7 @@ class BFCAllocator : public Allocator {
   }
   BinNum BinNumForSize(size_t bytes) {
     uint64 v = std::max<size_t>(bytes, 256) >> kMinAllocationBits;
-    int b = std::min(kNumBins - 1, Log2FloorNonZero(v));
+    int b = std::min(kNumBins - 1, tsl::Log2Floor64(v));
     return b;
   }
   Bin* BinForSize(size_t bytes) { return BinFromIndex(BinNumForSize(bytes)); }
diff --git a/third_party/xla/third_party/tsl/tsl/lib/core/BUILD b/third_party/xla/third_party/tsl/tsl/lib/core/BUILD
index c8bf89d1aa9bd4..0fee1ed113152a 100644
--- a/third_party/xla/third_party/tsl/tsl/lib/core/BUILD
+++ b/third_party/xla/third_party/tsl/tsl/lib/core/BUILD
@@ -83,6 +83,7 @@ cc_library(
     compatible_with = get_compatible_with_portable(),
     deps = [
         "//tsl/platform:logging",
+        "@com_google_absl//absl/numeric:bits",
     ],
     alwayslink = 1,
 )
@@ -102,7 +103,6 @@ cc_library(
     hdrs = ["bits.h"],
     deps = [
         "//tsl/platform:logging",
-        "//tsl/platform:types",
         "@com_google_absl//absl/numeric:bits",
     ],
 )
diff --git a/third_party/xla/third_party/tsl/tsl/lib/core/bitmap.cc b/third_party/xla/third_party/tsl/tsl/lib/core/bitmap.cc
index f8648d95e3ec5e..8681428cbd879d 100644
--- a/third_party/xla/third_party/tsl/tsl/lib/core/bitmap.cc
+++ b/third_party/xla/third_party/tsl/tsl/lib/core/bitmap.cc
@@ -15,8 +15,13 @@ limitations under the License.
 
 #include "tsl/lib/core/bitmap.h"
 
+#include <cstddef>
+#include <cstdint>
+#include <cstring>
 #include <string>
 
+#include "absl/numeric/bits.h"
+
 namespace tsl {
 namespace core {
 
@@ -34,41 +39,7 @@ void Bitmap::Reset(size_t n) {
 
 // Return 1+index of the first set bit in w; return 0 if w == 0.
 static size_t FindFirstSet(uint32_t w) {
-  // TODO(jeff,sanjay): If this becomes a performance issue, we could
-  // use the __builtin_ffs(w) routine on GCC, or the ffs(w) routine on
-  // some other platforms.
-
-  // clang-format off
-  static uint8_t kLowestBitSet[256] = {
-    /*  0*/ 0,  1,  2,  1,  3,  1,  2,  1,  4,  1,  2,  1,  3,  1,  2,  1,
-    /* 16*/ 5,  1,  2,  1,  3,  1,  2,  1,  4,  1,  2,  1,  3,  1,  2,  1,
-    /* 32*/ 6,  1,  2,  1,  3,  1,  2,  1,  4,  1,  2,  1,  3,  1,  2,  1,
-    /* 48*/ 5,  1,  2,  1,  3,  1,  2,  1,  4,  1,  2,  1,  3,  1,  2,  1,
-    /* 64*/ 7,  1,  2,  1,  3,  1,  2,  1,  4,  1,  2,  1,  3,  1,  2,  1,
-    /* 80*/ 5,  1,  2,  1,  3,  1,  2,  1,  4,  1,  2,  1,  3,  1,  2,  1,
-    /* 96*/ 6,  1,  2,  1,  3,  1,  2,  1,  4,  1,  2,  1,  3,  1,  2,  1,
-    /*112*/ 5,  1,  2,  1,  3,  1,  2,  1,  4,  1,  2,  1,  3,  1,  2,  1,
-    /*128*/ 8,  1,  2,  1,  3,  1,  2,  1,  4,  1,  2,  1,  3,  1,  2,  1,
-    /*144*/ 5,  1,  2,  1,  3,  1,  2,  1,  4,  1,  2,  1,  3,  1,  2,  1,
-    /*160*/ 6,  1,  2,  1,  3,  1,  2,  1,  4,  1,  2,  1,  3,  1,  2,  1,
-    /*176*/ 5,  1,  2,  1,  3,  1,  2,  1,  4,  1,  2,  1,  3,  1,  2,  1,
-    /*192*/ 7,  1,  2,  1,  3,  1,  2,  1,  4,  1,  2,  1,  3,  1,  2,  1,
-    /*208*/ 5,  1,  2,  1,  3,  1,  2,  1,  4,  1,  2,  1,  3,  1,  2,  1,
-    /*224*/ 6,  1,  2,  1,  3,  1,  2,  1,  4,  1,  2,  1,  3,  1,  2,  1,
-    /*240*/ 5,  1,  2,  1,  3,  1,  2,  1,  4,  1,  2,  1,  3,  1,  2,  1,
-  };
-  // clang-format on
-  if (w & 0xff) {
-    return kLowestBitSet[w & 0xff];
-  } else if ((w >> 8) & 0xff) {
-    return kLowestBitSet[(w >> 8) & 0xff] + 8;
-  } else if ((w >> 16) & 0xff) {
-    return kLowestBitSet[(w >> 16) & 0xff] + 16;
-  } else if ((w >> 24) & 0xff) {
-    return kLowestBitSet[(w >> 24) & 0xff] + 24;
-  } else {
-    return 0;
-  }
+  return w == 0 ? 0 : absl::countr_zero(w) + 1;
 }
 
 size_t Bitmap::FirstUnset(size_t start) const {
diff --git a/third_party/xla/third_party/tsl/tsl/lib/core/bits.h b/third_party/xla/third_party/tsl/tsl/lib/core/bits.h
index 9a31ae5815febe..f179117201abad 100644
--- a/third_party/xla/third_party/tsl/tsl/lib/core/bits.h
+++ b/third_party/xla/third_party/tsl/tsl/lib/core/bits.h
@@ -20,91 +20,29 @@ limitations under the License.
 
 #include "absl/numeric/bits.h"
 #include "tsl/platform/logging.h"
-#include "tsl/platform/types.h"
 
 namespace tsl {
 
 // Return floor(log2(n)) for positive integer n.  Returns -1 iff n == 0.
-int Log2Floor(uint32 n);
-int Log2Floor64(uint64 n);
-
-// Return ceiling(log2(n)) for positive integer n.  Returns -1 iff n == 0.
-int Log2Ceiling(uint32 n);
-int Log2Ceiling64(uint64 n);
-
-// ------------------------------------------------------------------------
-// Implementation details follow
-// ------------------------------------------------------------------------
-
-#if defined(__GNUC__)
-
-// Return floor(log2(n)) for positive integer n.  Returns -1 iff n == 0.
-inline int Log2Floor(uint32 n) { return n == 0 ? -1 : 31 ^ __builtin_clz(n); }
-
-// Return floor(log2(n)) for positive integer n.  Returns -1 iff n == 0.
-inline int Log2Floor64(uint64 n) {
-  return n == 0 ? -1 : 63 ^ __builtin_clzll(n);
-}
-
-#else
+inline int Log2Floor(uint32_t n) { return absl::bit_width(n) - 1; }
 
 // Return floor(log2(n)) for positive integer n.  Returns -1 iff n == 0.
-inline int Log2Floor(uint32 n) {
-  if (n == 0) return -1;
-  int log = 0;
-  uint32 value = n;
-  for (int i = 4; i >= 0; --i) {
-    int shift = (1 << i);
-    uint32 x = value >> shift;
-    if (x != 0) {
-      value = x;
-      log += shift;
-    }
-  }
-  assert(value == 1);
-  return log;
-}
-
-// Return floor(log2(n)) for positive integer n.  Returns -1 iff n == 0.
-// Log2Floor64() is defined in terms of Log2Floor32()
-inline int Log2Floor64(uint64 n) {
-  const uint32 topbits = static_cast<uint32>(n >> 32);
-  if (topbits == 0) {
-    // Top bits are zero, so scan in bottom bits
-    return Log2Floor(static_cast<uint32>(n));
-  } else {
-    return 32 + Log2Floor(topbits);
-  }
-}
+inline int Log2Floor64(uint64_t n) { return absl::bit_width(n) - 1; }
 
-#endif
-
-inline int Log2Ceiling(uint32 n) {
-  int floor = Log2Floor(n);
-  if (n == (n & ~(n - 1)))  // zero or a power of two
-    return floor;
-  else
-    return floor + 1;
+// Return ceiling(log2(n)) for positive integer n.  Returns -1 iff n == 0.
+inline int Log2Ceiling(uint32_t n) {
+  return n == 0 ? -1 : absl::bit_width(n - 1);
 }
 
-inline int Log2Ceiling64(uint64 n) {
-  int floor = Log2Floor64(n);
-  if (n == (n & ~(n - 1)))  // zero or a power of two
-    return floor;
-  else
-    return floor + 1;
+// Return ceiling(log2(n)) for positive integer n.  Returns -1 iff n == 0.
+inline int Log2Ceiling64(uint64_t n) {
+  return n == 0 ? -1 : absl::bit_width(n - 1);
 }
 
-inline uint32 NextPowerOfTwo(uint32 value) {
-  int exponent = Log2Ceiling(value);
-  DCHECK_LT(exponent, std::numeric_limits<uint32>::digits);
-  return 1 << exponent;
-}
+inline uint32_t NextPowerOfTwo(uint32_t value) { return absl::bit_ceil(value); }
 
-inline uint64 NextPowerOfTwo64(uint64 value) {
-  int exponent = Log2Ceiling(value);
-  DCHECK_LT(exponent, std::numeric_limits<uint64>::digits);
-  return 1LL << exponent;
+inline uint64_t NextPowerOfTwo64(uint64_t value) {
+  return absl::bit_ceil(value);
 }
 
 inline int64_t NextPowerOfTwoS64(int64_t value) {

From 98976d397b052b39c6203dae707d1b68f43c2e73 Mon Sep 17 00:00:00 2001
From: Kyle Lucke <klucke@google.com>
Date: Wed, 15 May 2024 10:15:15 -0700
Subject: [PATCH 191/478] Remove unused Stream methods and includes in
 stream.h.

PiperOrigin-RevId: 633995201
---
 third_party/xla/xla/stream_executor/stream.h | 39 --------------------
 1 file changed, 39 deletions(-)

diff --git a/third_party/xla/xla/stream_executor/stream.h b/third_party/xla/xla/stream_executor/stream.h
index f6bf0a49335d82..317846ea3bdfef 100644
--- a/third_party/xla/xla/stream_executor/stream.h
+++ b/third_party/xla/xla/stream_executor/stream.h
@@ -23,13 +23,10 @@ limitations under the License.
 
 #include <cstdint>
 #include <memory>
-#include <optional>
-#include <string>
 #include <utility>
 #include <variant>
 #include <vector>
 
-#include "absl/base/attributes.h"
 #include "absl/base/thread_annotations.h"
 #include "absl/functional/any_invocable.h"
 #include "absl/log/check.h"
@@ -145,22 +142,12 @@ class Stream {
   absl::Status ThenLaunch(ThreadDim thread_dims, BlockDim block_dims,
                           const TypedKernel<Params...> &kernel, Args... args);
 
-  template <typename... Params, typename... Args>
-  absl::Status ThenLaunch(ThreadDim thread_dims, BlockDim block_dims,
-                          ClusterDim cluster_dims,
-                          const TypedKernel<Params...> &kernel, Args... args);
-
   // Same as above, with an explicit argument for shared memory size in bytes.
   template <typename... Params, typename... Args>
   absl::Status ThenLaunch(ThreadDim thread_dims, BlockDim block_dims,
                           int32_t shmem_bytes,
                           const TypedKernel<Params...> &kernel, Args... args);
 
-  template <typename... Params, typename... Args>
-  absl::Status ThenLaunch(ThreadDim thread_dims, BlockDim block_dims,
-                          ClusterDim cluster_dims, int32_t shmem_bytes,
-                          const TypedKernel<Params...> &kernel, Args... args);
-
   // Create a dependency for this stream's next work on the other stream
   // completing. Does not take ownership of other, and other must not be
   // null.
@@ -279,7 +266,6 @@ class Stream {
     return parent_;
   }
 
-  //
   CudaComputeCapability GetCudaComputeCapability() const {
     return parent()->GetDeviceDescription().cuda_compute_capability();
   }
@@ -329,9 +315,6 @@ class Stream {
   void operator=(const Stream &) = delete;
 };
 
-////////////
-// Inlines
-
 template <typename... Params, typename... Args>
 inline absl::Status Stream::ThenLaunch(ThreadDim thread_dims,
                                        BlockDim block_dims,
@@ -354,28 +337,6 @@ inline absl::Status Stream::ThenLaunch(ThreadDim thread_dims,
   return absl::OkStatus();
 }
 
-template <typename... Params, typename... Args>
-inline absl::Status Stream::ThenLaunch(ThreadDim thread_dims,
-                                       BlockDim block_dims,
-                                       ClusterDim cluster_dims,
-                                       const TypedKernel<Params...> &kernel,
-                                       Args... args) {
-  auto kernel_args = PackKernelArgs(kernel, args...);
-  TF_RETURN_IF_ERROR(parent_->Launch(this, thread_dims, block_dims,
-                                     cluster_dims, *kernel, *kernel_args));
-  return absl::OkStatus();
-}
-
-template <typename... Params, typename... Args>
-inline absl::Status Stream::ThenLaunch(
-    ThreadDim thread_dims, BlockDim block_dims, ClusterDim cluster_dims,
-    int32_t shmem_bytes, const TypedKernel<Params...> &kernel, Args... args) {
-  auto kernel_args = PackKernelArgs(shmem_bytes, args...);
-  TF_RETURN_IF_ERROR(parent_->Launch(this, thread_dims, block_dims,
-                                     cluster_dims, *kernel, *kernel_args));
-  return absl::OkStatus();
-}
-
 }  // namespace stream_executor
 
 #endif  // XLA_STREAM_EXECUTOR_STREAM_H_

From 7993feaac8feea268cd713cb3503b1a3cdeeb0d0 Mon Sep 17 00:00:00 2001
From: Kyle Lucke <klucke@google.com>
Date: Wed, 15 May 2024 10:50:36 -0700
Subject: [PATCH 192/478] Use absl::Status instead of xla::Status now that
 they're identical.

PiperOrigin-RevId: 634006752
---
 third_party/xla/xla/python/ifrt/tuple.h | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/third_party/xla/xla/python/ifrt/tuple.h b/third_party/xla/xla/python/ifrt/tuple.h
index f86800e6cb3497..54cb1a1fd0284b 100644
--- a/third_party/xla/xla/python/ifrt/tuple.h
+++ b/third_party/xla/xla/python/ifrt/tuple.h
@@ -22,10 +22,10 @@ limitations under the License.
 #include <utility>
 #include <vector>
 
+#include "absl/status/status.h"
 #include "absl/types/span.h"
 #include "llvm/Support/ExtensibleRTTI.h"
 #include "xla/python/ifrt/value.h"
-#include "xla/status.h"
 #include "xla/tsl/concurrency/ref_count.h"
 
 namespace xla {
@@ -53,7 +53,7 @@ class Tuple : public llvm::RTTIExtends<Tuple, Value> {
   virtual int Arity() = 0;
 
   // Unpacks the tuple into its constituent pieces.
-  virtual Status Unpack(absl::Span<tsl::RCReference<Value>> values) = 0;
+  virtual absl::Status Unpack(absl::Span<tsl::RCReference<Value>> values) = 0;
 
   static char ID;  // NOLINT
 };

From dbed71af7fcf7ddf95c5e42ff9f0137c7948f60e Mon Sep 17 00:00:00 2001
From: Kyle Lucke <klucke@google.com>
Date: Wed, 15 May 2024 11:01:43 -0700
Subject: [PATCH 193/478] Use absl::Status instead of xla::Status now that
 they're identical.

PiperOrigin-RevId: 634010613
---
 .../address_computation_fusion_rewriter.cc    |  9 +++--
 .../collective_permute_cycle_decomposer.cc    | 16 ++++----
 .../xla/service/gpu/conv_algorithm_picker.cc  | 11 +++---
 .../service/gpu/cudnn_fused_mha_rewriter.cc   |  2 +-
 .../xla/xla/service/gpu/cusolver_context.cc   |  2 +-
 .../xla/service/gpu/dot_sparsity_rewriter.cc  |  2 +-
 .../xla/xla/service/gpu/fusion_wrapper.cc     |  2 +-
 .../xla/xla/service/gpu/gpu_p2p_pipeliner.cc  |  6 +--
 .../xla/xla/service/gpu/ir_emitter_triton.h   |  2 +-
 .../xla/service/gpu/ir_emitter_unnested.cc    |  4 +-
 .../memory_space_assignment/algorithm.cc      |  9 +++--
 .../memory_space_assignment/algorithm.h       |  6 +--
 .../memory_space_assignment/allocation.cc     | 20 +++++-----
 .../memory_space_assignment/allocation.h      | 37 ++++++++++---------
 .../cost_analysis_test.cc                     |  4 +-
 .../memory_bound_loop_optimizer.cc            |  2 +-
 .../memory_bound_loop_optimizer.h             |  2 +-
 .../memory_bound_loop_optimizer_test.cc       |  8 ++--
 .../memory_space_assignment.cc                | 23 ++++++------
 .../memory_space_assignment.h                 | 15 ++++----
 .../memory_space_assignment_test.cc           | 22 +++++------
 21 files changed, 104 insertions(+), 100 deletions(-)

diff --git a/third_party/xla/xla/service/gpu/address_computation_fusion_rewriter.cc b/third_party/xla/xla/service/gpu/address_computation_fusion_rewriter.cc
index b065166555e7fe..2f2a9e835e102e 100644
--- a/third_party/xla/xla/service/gpu/address_computation_fusion_rewriter.cc
+++ b/third_party/xla/xla/service/gpu/address_computation_fusion_rewriter.cc
@@ -275,10 +275,11 @@ absl::InlinedVector<HloInstruction*, 4> GetPatternCaptures(
   return captures;
 }
 
-Status CreateRootTuple(HloInstruction* hero, HloComputation::Builder& builder,
-                       DataflowPathsView sliced_user_paths,
-                       absl::flat_hash_map<const HloInstruction*,
-                                           HloInstruction*>& instr_mapping) {
+absl::Status CreateRootTuple(
+    HloInstruction* hero, HloComputation::Builder& builder,
+    DataflowPathsView sliced_user_paths,
+    absl::flat_hash_map<const HloInstruction*, HloInstruction*>&
+        instr_mapping) {
   unsigned tuple_size = hero->shape().tuple_shapes_size();
 
   std::vector<HloInstruction*> sliced_elems(tuple_size, nullptr);
diff --git a/third_party/xla/xla/service/gpu/collective_permute_cycle_decomposer.cc b/third_party/xla/xla/service/gpu/collective_permute_cycle_decomposer.cc
index 8edf682481baf9..d595352c7522a5 100644
--- a/third_party/xla/xla/service/gpu/collective_permute_cycle_decomposer.cc
+++ b/third_party/xla/xla/service/gpu/collective_permute_cycle_decomposer.cc
@@ -126,10 +126,10 @@ CycleType ShouldDecomposeWithCycleType(
 
 // Constructs the frontend attributes for the two decomposed CollectivePermute
 // instructions.
-Status GetFrontendAttributes(HloCollectivePermuteInstruction* cp,
-                             CycleType cycle_type,
-                             xla::FrontendAttributes& cp1_attr,
-                             xla::FrontendAttributes& cp2_attr) {
+absl::Status GetFrontendAttributes(HloCollectivePermuteInstruction* cp,
+                                   CycleType cycle_type,
+                                   xla::FrontendAttributes& cp1_attr,
+                                   xla::FrontendAttributes& cp2_attr) {
   cp1_attr = cp->frontend_attributes();
   cp2_attr = cp->frontend_attributes();
   auto validation_it =
@@ -177,11 +177,9 @@ Status GetFrontendAttributes(HloCollectivePermuteInstruction* cp,
 
 // Decomposes a CollectivePermute instruction with a cycle in its source-target
 // pairs into two CollectivePermute instructions.
-Status DecomposeCollectivePermuteCycle(HloCollectivePermuteInstruction* cp,
-                                       HloComputation* computation,
-                                       HloModule* module,
-                                       int64_t next_channel_id,
-                                       CycleType cycle_type) {
+absl::Status DecomposeCollectivePermuteCycle(
+    HloCollectivePermuteInstruction* cp, HloComputation* computation,
+    HloModule* module, int64_t next_channel_id, CycleType cycle_type) {
   const SourceTargetPairs& pairs = cp->source_target_pairs();
   int64_t num_pairs = pairs.size();
   // A forward cycle has its backedge at the end as in
diff --git a/third_party/xla/xla/service/gpu/conv_algorithm_picker.cc b/third_party/xla/xla/service/gpu/conv_algorithm_picker.cc
index ea4a961e753b0c..2a6dcae8417b9b 100644
--- a/third_party/xla/xla/service/gpu/conv_algorithm_picker.cc
+++ b/third_party/xla/xla/service/gpu/conv_algorithm_picker.cc
@@ -466,11 +466,12 @@ GpuConvAlgorithmPicker::AutotuneRuntimeArguments::FromInstruction(
   return runtime_arguments;
 }
 
-// There are three tiers of errors possible here: returning a failed StatusOr
-// means autotuning fails immediately; returning an AutotuneResult with a
-// failure code other than DISQUALIFIED means autotuning fails if
-// crash_on_checking_failure is set; and returning a DISQUALIFIED AutotuneResult
-// simply skips the engine/algorithm while recording a reason for skipping it.
+// There are three tiers of errors possible here: returning a failed
+// absl::StatusOr means autotuning fails immediately; returning an
+// AutotuneResult with a failure code other than DISQUALIFIED means autotuning
+// fails if crash_on_checking_failure is set; and returning a DISQUALIFIED
+// AutotuneResult simply skips the engine/algorithm while recording a reason for
+// skipping it.
 absl::StatusOr<AutotuneResult> GpuConvAlgorithmPicker::AutotuneOneConvRunner(
     GenericConvRunner* const runner,
     std::optional<ReferenceResult>* reference_result,
diff --git a/third_party/xla/xla/service/gpu/cudnn_fused_mha_rewriter.cc b/third_party/xla/xla/service/gpu/cudnn_fused_mha_rewriter.cc
index 0dea3331bc11cd..ebbe483725bdc4 100644
--- a/third_party/xla/xla/service/gpu/cudnn_fused_mha_rewriter.cc
+++ b/third_party/xla/xla/service/gpu/cudnn_fused_mha_rewriter.cc
@@ -1574,7 +1574,7 @@ absl::StatusOr<bool> FuseBwdMultiHeadedAttentionBlock(
   return true;
 }
 
-Status RestoreFwdGraph(
+absl::Status RestoreFwdGraph(
     HloComputation* comp, HloInstruction* fwd_fmha_call, HloInstruction* bmm2,
     HloInstruction* activation, HloInstruction* original_bmm2_producer0,
     HloInstruction* original_bmm2_producer1,
diff --git a/third_party/xla/xla/service/gpu/cusolver_context.cc b/third_party/xla/xla/service/gpu/cusolver_context.cc
index 8f9642ba4ea051..5814e33ec122cf 100644
--- a/third_party/xla/xla/service/gpu/cusolver_context.cc
+++ b/third_party/xla/xla/service/gpu/cusolver_context.cc
@@ -143,7 +143,7 @@ cublasFillMode_t GpuBlasUpperLower(se::blas::UpperLower uplo) {
   }
 }
 
-// Converts a cuSolver absl::Status to a Status.
+// Converts a cuSolver absl::Status to a absl::Status.
 absl::Status ConvertStatus(cusolverStatus_t status) {
   switch (status) {
     case CUSOLVER_STATUS_SUCCESS:
diff --git a/third_party/xla/xla/service/gpu/dot_sparsity_rewriter.cc b/third_party/xla/xla/service/gpu/dot_sparsity_rewriter.cc
index 7c1e7546482ea9..00f3de94fab99c 100644
--- a/third_party/xla/xla/service/gpu/dot_sparsity_rewriter.cc
+++ b/third_party/xla/xla/service/gpu/dot_sparsity_rewriter.cc
@@ -37,7 +37,7 @@ namespace {
 
 class SparseDotRewriterImpl : public DfsHloRewriteVisitor {
  public:
-  Status HandleDot(HloInstruction* instr) override {
+  absl::Status HandleDot(HloInstruction* instr) override {
     // Only handle sparse dots with a single RHS sparse descriptor.
     HloDotInstruction* dot = Cast<HloDotInstruction>(instr);
     if (dot->sparse_operands() != 1 || dot->sparsity().front().index() != 1) {
diff --git a/third_party/xla/xla/service/gpu/fusion_wrapper.cc b/third_party/xla/xla/service/gpu/fusion_wrapper.cc
index 5b5120397a2665..940597be29376e 100644
--- a/third_party/xla/xla/service/gpu/fusion_wrapper.cc
+++ b/third_party/xla/xla/service/gpu/fusion_wrapper.cc
@@ -36,7 +36,7 @@ absl::StatusOr<bool> FusionWrapper::Run(
   auto instructions = module->entry_computation()->MakeInstructionPostOrder();
   bool changed = false;
 
-  std::function<Status(HloInstruction*)> handle_instruction;
+  std::function<absl::Status(HloInstruction*)> handle_instruction;
   handle_instruction = [&](HloInstruction* instruction) -> absl::Status {
     switch (instruction->opcode()) {
       case HloOpcode::kConditional:
diff --git a/third_party/xla/xla/service/gpu/gpu_p2p_pipeliner.cc b/third_party/xla/xla/service/gpu/gpu_p2p_pipeliner.cc
index 843b020ab29c8f..5c48574307faa0 100644
--- a/third_party/xla/xla/service/gpu/gpu_p2p_pipeliner.cc
+++ b/third_party/xla/xla/service/gpu/gpu_p2p_pipeliner.cc
@@ -76,7 +76,7 @@ bool ShouldAllowLoopVariantParameterInChain(const HloInstruction* instr) {
   return true;
 }
 
-Status PostprocessP2PImpl(
+absl::Status PostprocessP2PImpl(
     HloInstruction* instr,
     std::function<std::string(std::vector<ReplicaGroup>&)> transformer) {
   // The input instruction is a Done instruction.
@@ -106,7 +106,7 @@ Status PostprocessP2PImpl(
 
 // Modifies the loop iteration frontend attribute for the peeled off Send and
 // Recv for the first iteration of a loop.
-Status PostprocessPeeledP2P(HloInstruction* instr) {
+absl::Status PostprocessPeeledP2P(HloInstruction* instr) {
   auto transform_bounds = [&](std::vector<ReplicaGroup>& replica_groups) {
     std::vector<std::pair<int64_t, int64_t>> bounds;
     bounds.reserve(replica_groups.size());
@@ -149,7 +149,7 @@ Status PostprocessPeeledP2P(HloInstruction* instr) {
 
 // Modifies the loop iteration frontend attribute for the rotated Send and Recv
 // for the remaining iterations in a loop.
-Status PostprocessRotatedP2P(HloInstruction* instr) {
+absl::Status PostprocessRotatedP2P(HloInstruction* instr) {
   auto transform_bounds = [&](std::vector<ReplicaGroup>& replica_groups) {
     std::vector<std::pair<int64_t, int64_t>> bounds;
     bounds.reserve(replica_groups.size());
diff --git a/third_party/xla/xla/service/gpu/ir_emitter_triton.h b/third_party/xla/xla/service/gpu/ir_emitter_triton.h
index 91cd73e70da7c3..a384db6a27965b 100644
--- a/third_party/xla/xla/service/gpu/ir_emitter_triton.h
+++ b/third_party/xla/xla/service/gpu/ir_emitter_triton.h
@@ -76,7 +76,7 @@ absl::Status EmitSoftMax(mlir::OpBuilder b, absl::string_view libdevice_path,
                          mlir::triton::FuncOp fn,
                          const TritonGemmConfig& config);
 
-using TritonIrEmitter = std::function<Status(
+using TritonIrEmitter = std::function<absl::Status(
     mlir::OpBuilder, absl::string_view, const se::DeviceDescription&,
     const TritonFusionAnalysis& analysis, const HloComputation*,
     mlir::triton::FuncOp, const TritonGemmConfig&)>;
diff --git a/third_party/xla/xla/service/gpu/ir_emitter_unnested.cc b/third_party/xla/xla/service/gpu/ir_emitter_unnested.cc
index 85e39381bac1d9..5961fe8d68bce1 100644
--- a/third_party/xla/xla/service/gpu/ir_emitter_unnested.cc
+++ b/third_party/xla/xla/service/gpu/ir_emitter_unnested.cc
@@ -238,7 +238,7 @@ static ConditionalThunkConfig GetConditionalThunkConfig(
   return config;
 }
 
-Status IrEmitterUnnested::EmitConditional(const HloInstruction* instr) {
+absl::Status IrEmitterUnnested::EmitConditional(const HloInstruction* instr) {
   std::vector<ThunkSequence> branch_thunks;
   branch_thunks.reserve(instr->branch_count());
 
@@ -2144,7 +2144,7 @@ absl::Status IrEmitterUnnested::EmitReplicaOrPartitionId(
   return absl::OkStatus();
 }
 
-Status IrEmitterUnnested::EmitCollectivePermute(
+absl::Status IrEmitterUnnested::EmitCollectivePermute(
     const HloCollectivePermuteInstruction* instr) {
   TF_RET_CHECK(instr->operand_count() == 1);
   auto* operand = instr->operand(0);
diff --git a/third_party/xla/xla/service/memory_space_assignment/algorithm.cc b/third_party/xla/xla/service/memory_space_assignment/algorithm.cc
index b170014f568f6b..c8cafe0bf8410e 100644
--- a/third_party/xla/xla/service/memory_space_assignment/algorithm.cc
+++ b/third_party/xla/xla/service/memory_space_assignment/algorithm.cc
@@ -942,8 +942,9 @@ void MsaAlgorithm::DumpDebugStringsIfEnabled() const {
   options_.dump_fn("scheduleinfo", instruction_schedule_str_);
 }
 
-Status MsaAlgorithm::OptimizeMemoryBoundLoop(int loop_start_idx,
-                                             int loop_end_idx, int loop_size) {
+absl::Status MsaAlgorithm::OptimizeMemoryBoundLoop(int loop_start_idx,
+                                                   int loop_end_idx,
+                                                   int loop_size) {
   // The MemoryBoundLoopOptimizer works with a minimum of three unrolled loop
   // iterations: previous, current, and next. So, we pick the second iteration
   // out of the loop as the current iteration.
@@ -1912,7 +1913,7 @@ absl::StatusOr<MsaAlgorithm::Result> MsaAlgorithm::AllocateAllocationValues(
         result_mark(AllocateSegment(request), result);
         if (request.require_no_copy_alternate_mem_allocation &&
             result != Result::kSuccess) {
-          Status failed_precondition = FailedPrecondition(
+          absl::Status failed_precondition = FailedPrecondition(
               "The value defined at %s requires allocation in the alternate "
               "memory, which could not be satisfied. This typically happens "
               "because more pinned buffers are live than the alternate memory "
@@ -3373,7 +3374,7 @@ void MsaAlgorithm::ImportRepackedSlicedAllocation(
           << "; Allocation: " << allocation->ToString();
 }
 
-Status MsaAlgorithm::AreRepackedSlicesValid(
+absl::Status MsaAlgorithm::AreRepackedSlicesValid(
     const RepackAllocationBlock& block) {
   if (!block.repacked_slice_data.has_value()) {
     return OkStatus();
diff --git a/third_party/xla/xla/service/memory_space_assignment/algorithm.h b/third_party/xla/xla/service/memory_space_assignment/algorithm.h
index b9e577225fffd6..ea775b8a9a7369 100644
--- a/third_party/xla/xla/service/memory_space_assignment/algorithm.h
+++ b/third_party/xla/xla/service/memory_space_assignment/algorithm.h
@@ -684,8 +684,8 @@ class MsaAlgorithm : public GlobalDecreasingSizeBestFitHeap<HloValue> {
   // For the given loop with the start and end index and loop size, run the
   // MemoryBoundLoopOptimizer and record its outputs into
   // optimized_allocations_map_.
-  Status OptimizeMemoryBoundLoop(int loop_start_idx, int loop_end_idx,
-                                 int loop_size);
+  absl::Status OptimizeMemoryBoundLoop(int loop_start_idx, int loop_end_idx,
+                                       int loop_size);
 
   // Identify memory-bound loops in the graph and call OptimizeMemoryBoundLoop
   // for the found loops.
@@ -945,7 +945,7 @@ class MsaAlgorithm : public GlobalDecreasingSizeBestFitHeap<HloValue> {
   // Helper functions to implement ImportRepackedAllocations.
   void ImportRepackedNonSlicedAllocation(RepackAllocationBlock& block);
   void ImportRepackedSlicedAllocation(RepackAllocationBlock& block);
-  Status AreRepackedSlicesValid(const RepackAllocationBlock& block);
+  absl::Status AreRepackedSlicesValid(const RepackAllocationBlock& block);
 
   // Adds an asynchronous copy to allocations.
   void AddAsyncCopy(
diff --git a/third_party/xla/xla/service/memory_space_assignment/allocation.cc b/third_party/xla/xla/service/memory_space_assignment/allocation.cc
index 1ca05bfa8c5d14..9f01ddbd9ab791 100644
--- a/third_party/xla/xla/service/memory_space_assignment/allocation.cc
+++ b/third_party/xla/xla/service/memory_space_assignment/allocation.cc
@@ -152,8 +152,8 @@ void Allocation::AddUse(HloUse use) {
   uses_.push_back(use);
 }
 
-Status Allocation::UpdateUses(HloComputation* computation,
-                              HloInstruction* producing_instruction) {
+absl::Status Allocation::UpdateUses(HloComputation* computation,
+                                    HloInstruction* producing_instruction) {
   for (const HloUse& use : uses()) {
     HloInstruction* replacement_instruction = producing_instruction;
     Shape operand_shape = use.instruction->operand(use.operand_number)->shape();
@@ -259,7 +259,7 @@ bool PinnedAllocation::operator==(const Allocation& other) const {
   return casted_other != nullptr && (*this) == (*casted_other);
 }
 
-Status PinnedAllocation::Process() {
+absl::Status PinnedAllocation::Process() {
   if (is_scoped_allocation()) {
     // Nothing to do here for scoped allocations.
     return OkStatus();
@@ -313,7 +313,7 @@ int64_t CopyAllocation::earliest_available_time() const {
   return copy_done_schedule_before_;
 }
 
-Status CopyAllocation::Process() {
+absl::Status CopyAllocation::Process() {
   // Copy allocations need to insert asynchronous copy nodes.
   Shape shape = defining_position().shape();
   HloInstruction* producing_instruction = AddGetTupleElements();
@@ -439,7 +439,7 @@ SlicedCopyAllocation::SlicedCopyAllocation(
   }
 }
 
-Status SlicedCopyAllocation::Process() {
+absl::Status SlicedCopyAllocation::Process() {
   Shape shape = defining_position().shape();
   HloInstruction* producing_instruction = AddGetTupleElements();
 
@@ -619,7 +619,7 @@ std::string SlicedCopyAllocation::ToString() const {
       ", uses: ", UsesToString(uses()), ", from ", prev_allocation_.ToString());
 }
 
-Status SlicedCopyAllocation::CreateBitcastConcat(
+absl::Status SlicedCopyAllocation::CreateBitcastConcat(
     const Shape& shape, absl::Span<HloInstruction* const> slices) {
   CHECK(!slices.empty());
   concat_ =
@@ -649,7 +649,7 @@ bool SlicedCopyAllocation::SliceDetail::operator==(
   return SliceDetailToTuple(*this) == SliceDetailToTuple(other);
 }
 
-Status SlicedCopyAllocation::SliceDetail::CreateAsyncSlice(
+absl::Status SlicedCopyAllocation::SliceDetail::CreateAsyncSlice(
     const Shape& original_shape, HloInstruction& producer,
     HloComputation& parent) {
   if (original_shape.rank() != slice_decision.sizing.slice_params.size()) {
@@ -728,7 +728,7 @@ MirroredAllocation::MirroredAllocation(const Allocation& original_allocation,
                  /*cross_program_prefetch_index=*/std::nullopt),
       original_allocation_(original_allocation) {}
 
-Status MirroredAllocation::Process() {
+absl::Status MirroredAllocation::Process() {
   set_original_defining_position(original_allocation_.defining_position());
   if (is_scoped_allocation()) {
     // Nothing to do here for scoped allocations.
@@ -754,7 +754,7 @@ HloPosition ParentAllocation::defining_position() const {
   return original_defining_position();
 }
 
-Status ParentAllocation::Process() {
+absl::Status ParentAllocation::Process() {
   // Add an additional parameter to the while HLO with a reference to the buffer
   // in the default memory space.
   HloInstruction* producing_instruction =
@@ -796,7 +796,7 @@ Status ParentAllocation::Process() {
   return UpdateUses(computation, final_instruction);
 }
 
-Status ParentAllocation::PostProcess() {
+absl::Status ParentAllocation::PostProcess() {
   // Update the root of the while body with the new parameter. The reason why we
   // need a separate post-process for this is because other allocations may have
   // while body root as a use, so they would update the old root instead of the
diff --git a/third_party/xla/xla/service/memory_space_assignment/allocation.h b/third_party/xla/xla/service/memory_space_assignment/allocation.h
index accbe1bcff581f..3d01439aee7623 100644
--- a/third_party/xla/xla/service/memory_space_assignment/allocation.h
+++ b/third_party/xla/xla/service/memory_space_assignment/allocation.h
@@ -122,8 +122,8 @@ class Allocation {
   // Adds a use to this allocation.
   void AddUse(HloUse use);
   // Replaces all uses of the allocation with the copy_complete instruction.
-  Status UpdateUses(HloComputation* computation,
-                    HloInstruction* producing_instruction);
+  absl::Status UpdateUses(HloComputation* computation,
+                          HloInstruction* producing_instruction);
 
   // Allocation type methods
   // --------------------------------------------------------------------------
@@ -140,10 +140,10 @@ class Allocation {
   // After all of the time ranges for the allocations have been assigned,
   // Process morphs the instructions affected to assign the memory spaces and
   // insert asynchronous copy instructions if necessary.
-  virtual Status Process() = 0;
+  virtual absl::Status Process() = 0;
   // An optional post-process step that will be called after all allocations
   // have been processed.
-  virtual Status PostProcess() = 0;
+  virtual absl::Status PostProcess() = 0;
   // Marks (adds this allocation to needed_allocations) if this allocation is
   // needed. PinnedAllocation and CopyAllocations are always needed and
   // ParentAllocations are needed if they have any uses or if other
@@ -209,8 +209,8 @@ class PinnedAllocation final : public Allocation {
   int64_t earliest_available_time() const override { return start_time(); }
   bool is_copy_allocation() const override { return false; }
   bool is_sliced_copy_allocation() const override { return false; }
-  Status Process() override;
-  Status PostProcess() override { return OkStatus(); }
+  absl::Status Process() override;
+  absl::Status PostProcess() override { return OkStatus(); }
   void MarkIfNeeded(absl::flat_hash_set<const Allocation*>& needed_allocations)
       const override;
   void MarkNeeded(absl::flat_hash_set<const Allocation*>& needed_allocations)
@@ -246,8 +246,8 @@ class CopyAllocation final : public Allocation {
   int64_t earliest_available_time() const override;
   bool is_copy_allocation() const override { return true; }
   bool is_sliced_copy_allocation() const override { return false; }
-  Status Process() override;
-  Status PostProcess() override { return OkStatus(); }
+  absl::Status Process() override;
+  absl::Status PostProcess() override { return OkStatus(); }
   void MarkIfNeeded(absl::flat_hash_set<const Allocation*>& needed_allocations)
       const override;
   void MarkNeeded(absl::flat_hash_set<const Allocation*>& needed_allocations)
@@ -317,8 +317,9 @@ class SlicedCopyAllocation final : public Allocation {
 
     // Create the instructions to copy the slice. This method updates
     // copy_start and copy_done.
-    Status CreateAsyncSlice(const Shape& original_shape,
-                            HloInstruction& producer, HloComputation& parent);
+    absl::Status CreateAsyncSlice(const Shape& original_shape,
+                                  HloInstruction& producer,
+                                  HloComputation& parent);
 
     SliceDecision slice_decision;
     int64_t copy_start_after_time = -1;
@@ -347,8 +348,8 @@ class SlicedCopyAllocation final : public Allocation {
   bool is_sliced_copy_allocation() const override { return true; }
   // MemorySpaceAssignment::Process() calls Process() to create asynchronous
   // slice copies, and a bitcast-concat call to glue the slices back together.
-  Status Process() override;
-  Status PostProcess() override { return OkStatus(); }
+  absl::Status Process() override;
+  absl::Status PostProcess() override { return OkStatus(); }
   // Marks the allocation as needed.
   void MarkIfNeeded(absl::flat_hash_set<const Allocation*>& needed_allocations)
       const override;
@@ -372,8 +373,8 @@ class SlicedCopyAllocation final : public Allocation {
   SlicedCopyAllocation() = delete;
 
   // Create an instruction to concatenate the slices. Populates concat_.
-  Status CreateBitcastConcat(const Shape& shape,
-                             absl::Span<HloInstruction* const> slices);
+  absl::Status CreateBitcastConcat(const Shape& shape,
+                                   absl::Span<HloInstruction* const> slices);
 
   Shape original_shape_to_slice_;
   const Allocation& prev_allocation_;
@@ -403,8 +404,8 @@ class MirroredAllocation final : public Allocation {
   int64_t earliest_available_time() const override { return start_time(); }
   bool is_copy_allocation() const override { return false; }
   bool is_sliced_copy_allocation() const override { return false; }
-  Status Process() override;
-  Status PostProcess() override { return OkStatus(); }
+  absl::Status Process() override;
+  absl::Status PostProcess() override { return OkStatus(); }
   void MarkIfNeeded(absl::flat_hash_set<const Allocation*>& needed_allocations)
       const override;
   void MarkNeeded(absl::flat_hash_set<const Allocation*>& needed_allocations)
@@ -435,8 +436,8 @@ class ParentAllocation final : public Allocation {
   int64_t earliest_available_time() const override { return start_time(); }
   bool is_copy_allocation() const override { return false; }
   bool is_sliced_copy_allocation() const override { return false; }
-  Status Process() override;
-  Status PostProcess() override;
+  absl::Status Process() override;
+  absl::Status PostProcess() override;
   void MarkIfNeeded(absl::flat_hash_set<const Allocation*>& needed_allocations)
       const override;
   void MarkNeeded(absl::flat_hash_set<const Allocation*>& needed_allocations)
diff --git a/third_party/xla/xla/service/memory_space_assignment/cost_analysis_test.cc b/third_party/xla/xla/service/memory_space_assignment/cost_analysis_test.cc
index f4dcd98cdf0993..5340551b0c1374 100644
--- a/third_party/xla/xla/service/memory_space_assignment/cost_analysis_test.cc
+++ b/third_party/xla/xla/service/memory_space_assignment/cost_analysis_test.cc
@@ -45,8 +45,8 @@ int64_t ShapeSize(const Shape& shape) {
 
 class MemorySpaceAssignmentCostAnalysisTest : public HloTestBase {
  protected:
-  Status Initialize(const HloModule* module,
-                    float pipeline_overhead_window_size_mib = 0.0) {
+  absl::Status Initialize(const HloModule* module,
+                          float pipeline_overhead_window_size_mib = 0.0) {
     HloCostAnalysis::Options options;
     options_.alternate_mem_bandwidth_bytes_per_second = 128;
     options_.async_copy_bandwidth_bytes_per_second = 32;
diff --git a/third_party/xla/xla/service/memory_space_assignment/memory_bound_loop_optimizer.cc b/third_party/xla/xla/service/memory_space_assignment/memory_bound_loop_optimizer.cc
index c2615a7d0d0841..55035f996def94 100644
--- a/third_party/xla/xla/service/memory_space_assignment/memory_bound_loop_optimizer.cc
+++ b/third_party/xla/xla/service/memory_space_assignment/memory_bound_loop_optimizer.cc
@@ -92,7 +92,7 @@ MemoryBoundLoopOptimizer::MemoryBoundLoopOptimizer(
       size_function_(size_function),
       reserved_scoped_memory_fn_(reserved_scoped_memory_fn) {}
 
-Status MemoryBoundLoopOptimizer::Initialize() {
+absl::Status MemoryBoundLoopOptimizer::Initialize() {
   const auto& instruction_sequence =
       hlo_live_range_.flattened_instruction_sequence().instructions();
   VLOG(3) << "MemoryBoundLoopOptimizer::Initialize, loop start: " << loop_start_
diff --git a/third_party/xla/xla/service/memory_space_assignment/memory_bound_loop_optimizer.h b/third_party/xla/xla/service/memory_space_assignment/memory_bound_loop_optimizer.h
index 1f75ae693d34b3..cef1dffa52c693 100644
--- a/third_party/xla/xla/service/memory_space_assignment/memory_bound_loop_optimizer.h
+++ b/third_party/xla/xla/service/memory_space_assignment/memory_bound_loop_optimizer.h
@@ -207,7 +207,7 @@ class MemoryBoundLoopOptimizer {
       const ReservedScopedMemoryFunction& reserved_scoped_memory_fn);
 
   // Initializes the data structures used by the optimizer.
-  Status Initialize();
+  absl::Status Initialize();
 
   // Given an HloBuffer object, determines if this buffer represents a LoopValue
   // that can be optimized by the optimizer, and if so it adds a LoopValue to
diff --git a/third_party/xla/xla/service/memory_space_assignment/memory_bound_loop_optimizer_test.cc b/third_party/xla/xla/service/memory_space_assignment/memory_bound_loop_optimizer_test.cc
index f3a2ca55ab6fb9..15b50eac85f4d0 100644
--- a/third_party/xla/xla/service/memory_space_assignment/memory_bound_loop_optimizer_test.cc
+++ b/third_party/xla/xla/service/memory_space_assignment/memory_bound_loop_optimizer_test.cc
@@ -91,8 +91,8 @@ class MemoryBoundLoopOptimizerTest : public HloTestBase {
   const int64_t kAlternateMemorySpace = 1;
   const int64_t kDefaultMemorySpace = 0;
 
-  Status Initialize(const HloModule* module,
-                    uint64_t alternate_memory_size = 256) {
+  absl::Status Initialize(const HloModule* module,
+                          uint64_t alternate_memory_size = 256) {
     HloCostAnalysis::Options options;
     MemoryBoundLoopOptimizerOptions optimizer_options;
     optimizer_options.set_enabled(true);
@@ -370,8 +370,8 @@ ENTRY Entry {
     return preset_assignments;
   }
 
-  Status VerifyMsaEquivalence(HloModule* module,
-                              bool expect_unsupported_allocations = false) {
+  absl::Status VerifyMsaEquivalence(
+      HloModule* module, bool expect_unsupported_allocations = false) {
     // Create a map indexed by instruction number and operand number.
     absl::flat_hash_map<std::pair<int, int>, const Allocation*> allocation_map;
     for (const MemoryBoundLoopOptimizer::LoopValue& value :
diff --git a/third_party/xla/xla/service/memory_space_assignment/memory_space_assignment.cc b/third_party/xla/xla/service/memory_space_assignment/memory_space_assignment.cc
index 4f4fdc29e50d42..b2ede602ceb643 100644
--- a/third_party/xla/xla/service/memory_space_assignment/memory_space_assignment.cc
+++ b/third_party/xla/xla/service/memory_space_assignment/memory_space_assignment.cc
@@ -70,14 +70,14 @@ namespace xla {
 namespace memory_space_assignment {
 namespace {
 
-Status InsertInstructionAndEnsureOperandsInserted(
+absl::Status InsertInstructionAndEnsureOperandsInserted(
     HloInstruction* new_instruction, HloInstructionSequence* new_sequence,
     absl::flat_hash_set<HloInstruction*>* inserted_instructions);
 
 // Insert an instruction to the schedule, and make sure its dependencies
 // (operands) are already in the schedule. If not, insert these operands
 // before the instruction.
-Status EnsureInstructionAndOperandsInserted(
+absl::Status EnsureInstructionAndOperandsInserted(
     HloInstruction* new_instruction, HloInstructionSequence* new_sequence,
     absl::flat_hash_set<HloInstruction*>* inserted_instructions) {
   if (inserted_instructions->contains(new_instruction)) {
@@ -90,7 +90,7 @@ Status EnsureInstructionAndOperandsInserted(
 // Same as above, but does not check if instruction is already inserted. This is
 // used when the caller already knows the instruction isn't inserted yet, to
 // speed up compilation.
-Status InsertInstructionAndEnsureOperandsInserted(
+absl::Status InsertInstructionAndEnsureOperandsInserted(
     HloInstruction* new_instruction, HloInstructionSequence* new_sequence,
     absl::flat_hash_set<HloInstruction*>* inserted_instructions) {
   for (HloInstruction* operand : new_instruction->operands()) {
@@ -373,7 +373,7 @@ MemorySpaceAssignment::RunMemorySpaceAssignment(
   return std::move(preset_assignments_);
 }
 
-Status MemorySpaceAssignment::FindAllocationSequence(
+absl::Status MemorySpaceAssignment::FindAllocationSequence(
     const HloLiveRange& hlo_live_range,
     const HloAliasAnalysis& alias_analysis) {
   auto algorithm = std::make_unique<MsaAlgorithm>(
@@ -442,7 +442,8 @@ float MemorySpaceAssignment::ComputeEstimatedElapsedTime(
   return total_elapsed;
 }
 
-Status MemorySpaceAssignment::Process(const HloLiveRange& hlo_live_range) {
+absl::Status MemorySpaceAssignment::Process(
+    const HloLiveRange& hlo_live_range) {
   VLOG(1) << "Processing assigned buffers...";
   // Since some parent allocations may not be needed (e.g. when they don't have
   // any uses and if there is no other (non-parent) allocation that depends on
@@ -512,7 +513,7 @@ Status MemorySpaceAssignment::Process(const HloLiveRange& hlo_live_range) {
   return OkStatus();
 }
 
-Status MemorySpaceAssignment::ExportAndColorBuffers() {
+absl::Status MemorySpaceAssignment::ExportAndColorBuffers() {
   VLOG(1) << "Exporting buffers...";
   TF_ASSIGN_OR_RETURN(auto alias_analysis, HloAliasAnalysis::Run(module_));
   absl::flat_hash_map<int64_t, int64_t> seen_buffer_offsets;
@@ -603,7 +604,7 @@ void MemorySpaceAssignment::RemoveAssignmentForInstruction(
   }
 }
 
-Status MemorySpaceAssignment::SimplifyGraph() {
+absl::Status MemorySpaceAssignment::SimplifyGraph() {
   VLOG(1) << "Simplifying graph...";
   for (HloComputation* computation : module_->MakeNonfusionComputations()) {
     // Parallel computations aren't in the schedule and don't need to be
@@ -927,7 +928,7 @@ void MemorySpaceAssignment::ScheduleAsynchronousCopies() {
   }
 }
 
-Status MemorySpaceAssignment::FixSchedule() {
+absl::Status MemorySpaceAssignment::FixSchedule() {
   VLOG(1) << "Fixing schedule...";
   TF_RET_CHECK(module_->has_schedule());
   HloSchedule& schedule = module_->schedule();
@@ -1021,7 +1022,7 @@ Status MemorySpaceAssignment::FixSchedule() {
   return OkStatus();
 }
 
-Status MemorySpaceAssignment::VerifyAndExportHeapSimulatorTrace() {
+absl::Status MemorySpaceAssignment::VerifyAndExportHeapSimulatorTrace() {
   VLOG(1) << "Verifying...";
   TF_ASSIGN_OR_RETURN(std::unique_ptr<HloAliasAnalysis> alias_analysis,
                       HloAliasAnalysis::Run(module_));
@@ -1113,8 +1114,8 @@ Status MemorySpaceAssignment::VerifyAndExportHeapSimulatorTrace() {
         }
       }
 
-      std::function<Status(const HloInstruction*, int64_t, int64_t,
-                           absl::string_view)>
+      std::function<absl::Status(const HloInstruction*, int64_t, int64_t,
+                                 absl::string_view)>
           split_conditional_buffer;
       split_conditional_buffer = [&](const HloInstruction* use_instruction,
                                      int64_t start_time, int64_t end_time,
diff --git a/third_party/xla/xla/service/memory_space_assignment/memory_space_assignment.h b/third_party/xla/xla/service/memory_space_assignment/memory_space_assignment.h
index 9e0e6d1a9a45b2..e37f5ad53aeae0 100644
--- a/third_party/xla/xla/service/memory_space_assignment/memory_space_assignment.h
+++ b/third_party/xla/xla/service/memory_space_assignment/memory_space_assignment.h
@@ -297,7 +297,7 @@ class MemorySpaceAssignment {
 
   // Verify that the memory space assignment is free of overlapping buffers and
   // export heap simulator trace to be used by buffer_assignment.
-  Status VerifyAndExportHeapSimulatorTrace();
+  absl::Status VerifyAndExportHeapSimulatorTrace();
 
  protected:
   // Main driver of the memory space assignment pass.
@@ -307,8 +307,9 @@ class MemorySpaceAssignment {
 
   // Finds an AllocationSequence for placing buffers in alternate memory using
   // the MsaAlgorithm algorithm. Must be set before Process() is called.
-  virtual Status FindAllocationSequence(const HloLiveRange& hlo_live_range,
-                                        const HloAliasAnalysis& alias_analysis);
+  virtual absl::Status FindAllocationSequence(
+      const HloLiveRange& hlo_live_range,
+      const HloAliasAnalysis& alias_analysis);
 
   const Options& options() const { return options_; }
 
@@ -337,20 +338,20 @@ class MemorySpaceAssignment {
  private:
   // Process calls Process methods of the allocations after the allocations have
   // been finalized.
-  Status Process(const HloLiveRange& hlo_live_range);
+  absl::Status Process(const HloLiveRange& hlo_live_range);
 
   // Process() might have altered the computation graph by inserting kTuple and
   // kGetTupleElement instructions. SimplifyGraph performs a simple DCE and
   // tuple simplification operation (e.g., given GetTupleElement(Tuple(a, b),
   // 1), simply forwards b). Runs to fixed point.
-  Status SimplifyGraph();
+  absl::Status SimplifyGraph();
 
   // FixSchedule inserts asynchronous copies in the schedule.
-  Status FixSchedule();
+  absl::Status FixSchedule();
 
   // Export the alternate memory assignments to the PresetAssignments and color
   // the HLO graph with the determined memory spaces.
-  Status ExportAndColorBuffers();
+  absl::Status ExportAndColorBuffers();
 
   // Schedules asynchronous copies and ensures that the CopyStarts and their
   // corresponding CopyDones follow the same order.
diff --git a/third_party/xla/xla/service/memory_space_assignment/memory_space_assignment_test.cc b/third_party/xla/xla/service/memory_space_assignment/memory_space_assignment_test.cc
index 155ee2b597260b..b694fb852c5d50 100644
--- a/third_party/xla/xla/service/memory_space_assignment/memory_space_assignment_test.cc
+++ b/third_party/xla/xla/service/memory_space_assignment/memory_space_assignment_test.cc
@@ -113,7 +113,7 @@ int64_t ReservedScopedMemoryFn(
 }
 
 template <typename MessageType>
-StatusOr<MessageType> ParseTextProto(const std::string& text_proto) {
+absl::StatusOr<MessageType> ParseTextProto(const std::string& text_proto) {
   tsl::protobuf::TextFormat::Parser parser;
   MessageType parsed_proto;
   tsl::protobuf::io::ArrayInputStream input_stream(text_proto.data(),
@@ -10344,7 +10344,7 @@ class SlicedPrefetchTest : public MemorySpaceAssignmentTestBase {
   // REQUIRES:
   // - Concat-bitcast and all slices were found in the schedule used to
   //   construct schedule_to_class.
-  static Status ConcatBitcastAndSlicesAfterInstruction(
+  static absl::Status ConcatBitcastAndSlicesAfterInstruction(
       const std::vector<HloInstruction*>& schedule,
       const std::vector<InstructionClass>& schedule_to_class,
       int slices_start_after_index) {
@@ -10370,7 +10370,7 @@ class SlicedPrefetchTest : public MemorySpaceAssignmentTestBase {
   // REQUIRES:
   // - Concat-bitcast and all slices were found in the schedule used to
   //   construct schedule_to_class.
-  static Status AtLeastOneNonCopyLikeInstructionBetweenSliceStarts(
+  static absl::Status AtLeastOneNonCopyLikeInstructionBetweenSliceStarts(
       const std::vector<HloInstruction*>& schedule,
       const std::vector<InstructionClass>& schedule_to_class) {
     bool found_non_copy_since_last_slice_start = true;
@@ -10400,7 +10400,7 @@ class SlicedPrefetchTest : public MemorySpaceAssignmentTestBase {
   // REQUIRES:
   // - Concat-bitcast and all slices were found in the schedule used to
   //   construct schedule_to_class.
-  static Status OneSliceStartAfterInstructionWithNoCopyLikeBetween(
+  static absl::Status OneSliceStartAfterInstructionWithNoCopyLikeBetween(
       const std::vector<HloInstruction*>& schedule,
       const std::vector<InstructionClass>& schedule_to_class,
       int slices_start_after_index) {
@@ -10453,7 +10453,7 @@ class SlicedPrefetchTest : public MemorySpaceAssignmentTestBase {
   // REQUIRES:
   // - Concat-bitcast and all slices were found in the schedule used to
   //   construct schedule_to_class.
-  static Status ConcatBitcastAndSlicesBeforeInstruction(
+  static absl::Status ConcatBitcastAndSlicesBeforeInstruction(
       const std::vector<HloInstruction*>& schedule,
       const std::vector<InstructionClass>& schedule_to_class,
       int slices_done_before_index) {
@@ -10480,7 +10480,7 @@ class SlicedPrefetchTest : public MemorySpaceAssignmentTestBase {
   // REQUIRES:
   // - Concat-bitcast and all slices were found in the schedule used to
   //   construct schedule_to_class.
-  static Status
+  static absl::Status
   ConcatBitcastAndSliceDonesBeforeInstructionWithNoCopyLikeBetween(
       const std::vector<HloInstruction*>& schedule,
       const std::vector<InstructionClass>& schedule_to_class,
@@ -10513,7 +10513,7 @@ class SlicedPrefetchTest : public MemorySpaceAssignmentTestBase {
   // REQUIRES:
   // - Concat-bitcast and all slices were found in the schedule used to
   //   construct schedule_to_class.
-  static Status ConcatBitcastAfterSliceDones(
+  static absl::Status ConcatBitcastAfterSliceDones(
       const std::vector<HloInstruction*>& schedule,
       const std::vector<InstructionClass>& schedule_to_class) {
     int concat_bitcast_index = -1;
@@ -10553,7 +10553,7 @@ class SlicedPrefetchTest : public MemorySpaceAssignmentTestBase {
   //   slices_done_before_instruction_name in the schedule, with no
   //   non-copy-like instruction between AND
   // - concat_bitcast comes after all slice dones AND
-  static Status CheckSchedule(
+  static absl::Status CheckSchedule(
       const HloModule& module, const HloInstruction* concat_bitcast,
       std::string_view slices_start_after_instruction_name,
       std::string_view slices_done_before_instruction_name,
@@ -10643,9 +10643,9 @@ class SlicedPrefetchTest : public MemorySpaceAssignmentTestBase {
   //   the slice chunks AND
   // - The size of the chunk assigned to the sliced_copy_result has the same
   //   size as the instruction's shape
-  static Status CheckSliceChunks(const PresetAssignments& assignments,
-                                 const HloInstruction* sliced_copy_result,
-                                 bool expect_bitcasted_io = false) {
+  static absl::Status CheckSliceChunks(const PresetAssignments& assignments,
+                                       const HloInstruction* sliced_copy_result,
+                                       bool expect_bitcasted_io = false) {
     const HloInstruction* concat_bitcast =
         (expect_bitcasted_io ? sliced_copy_result->operand(0)
                              : sliced_copy_result);

From 40e4d40cad4a92de321831039e8cf9791bca5288 Mon Sep 17 00:00:00 2001
From: Kyle Lucke <klucke@google.com>
Date: Wed, 15 May 2024 11:06:22 -0700
Subject: [PATCH 194/478] Move Event creation into StreamExecutorInterface.

This is the first step in eliminating the pimpl pattern from Events.  Backward references made me had to temporarily eliminate the event target.  Eventually having Event inherit from EventInterface will enable it to be separated out again.

PiperOrigin-RevId: 634012404
---
 .../xla/service/gpu/gpu_transfer_manager.cc   |  5 +--
 third_party/xla/xla/service/gpu/runtime/BUILD |  4 +--
 .../xla/xla/service/gpu/runtime/copy_thunk.cc | 32 +++++++------------
 .../xla/xla/service/gpu/runtime/copy_thunk.h  | 10 +++---
 .../runtime/nccl_collective_permute_thunk.h   |  1 +
 .../gpu/runtime/nccl_collective_thunk.cc      |  8 ++---
 .../gpu/runtime/nccl_collective_thunk.h       |  4 +--
 third_party/xla/xla/stream_executor/BUILD     | 18 ++---------
 .../stream_executor/mock_stream_executor.h    |  2 ++
 .../stream_executor_interface.h               |  4 +++
 .../stream_executor/stream_executor_pimpl.cc  |  9 ++++++
 .../stream_executor/stream_executor_pimpl.h   |  1 +
 12 files changed, 44 insertions(+), 54 deletions(-)

diff --git a/third_party/xla/xla/service/gpu/gpu_transfer_manager.cc b/third_party/xla/xla/service/gpu/gpu_transfer_manager.cc
index 89dd6343810bb2..f628e17aa5e6c8 100644
--- a/third_party/xla/xla/service/gpu/gpu_transfer_manager.cc
+++ b/third_party/xla/xla/service/gpu/gpu_transfer_manager.cc
@@ -318,10 +318,7 @@ GpuTransferManager::GetOrCreateStagingBuffer(se::StreamExecutor* executor) {
   TF_ASSIGN_OR_RETURN(auto staging_buffer,
                       executor->HostMemoryAllocate(kStagingBufferSize));
 
-  auto transfer_completed = std::make_unique<se::Event>(executor);
-  if (!transfer_completed->Init()) {
-    return absl::InternalError("Failed to initialize transfer completed event");
-  }
+  TF_ASSIGN_OR_RETURN(auto transfer_completed, executor->CreateEvent());
 
   auto emplaced = staging_buffers_.try_emplace(
       executor, std::move(staging_buffer), std::move(transfer_completed));
diff --git a/third_party/xla/xla/service/gpu/runtime/BUILD b/third_party/xla/xla/service/gpu/runtime/BUILD
index 4327ce5583e62b..3bc13cac2f9665 100644
--- a/third_party/xla/xla/service/gpu/runtime/BUILD
+++ b/third_party/xla/xla/service/gpu/runtime/BUILD
@@ -844,6 +844,7 @@ cc_library(
         ":nccl_api",
         ":nccl_collective_thunk",
         ":nccl_p2p_thunk_common",
+        ":thunk",
         "//xla:status_macros",
         "//xla:xla_data_proto_cc",
         "//xla/hlo/ir:hlo",
@@ -851,11 +852,11 @@ cc_library(
         "//xla/service:global_device_id",
         "//xla/service/gpu:backend_configs_cc",
         "//xla/service/gpu:ir_emission_utils",
-        "//xla/service/gpu/runtime:thunk",
         "//xla/stream_executor",
         "//xla/translate/mhlo_to_hlo:attribute_exporter",
         "//xla/tsl/concurrency:async_value",
         "@com_google_absl//absl/algorithm:container",
+        "@com_google_absl//absl/container:node_hash_map",
         "@com_google_absl//absl/log",
         "@com_google_absl//absl/status",
         "@com_google_absl//absl/strings:str_format",
@@ -902,7 +903,6 @@ cc_library(
         "@com_google_absl//absl/base:core_headers",
         "@com_google_absl//absl/container:flat_hash_map",
         "@com_google_absl//absl/container:flat_hash_set",
-        "@com_google_absl//absl/container:node_hash_map",
         "@com_google_absl//absl/log",
         "@com_google_absl//absl/status",
         "@com_google_absl//absl/status:statusor",
diff --git a/third_party/xla/xla/service/gpu/runtime/copy_thunk.cc b/third_party/xla/xla/service/gpu/runtime/copy_thunk.cc
index ee90d4befc60d2..89b29d75736f82 100644
--- a/third_party/xla/xla/service/gpu/runtime/copy_thunk.cc
+++ b/third_party/xla/xla/service/gpu/runtime/copy_thunk.cc
@@ -79,10 +79,10 @@ absl::Status CopyThunk::ExecuteOnStream(const ExecuteParams& params) {
 // and return the event in order to do RecordEvent() for async memcpy.
 absl::Status CopyThunk::AsyncEvents::Emplace(se::StreamExecutor* executor,
                                              const HloInstruction* instr,
-                                             se::Event&& event) {
+                                             std::unique_ptr<se::Event> event) {
   Key key = {executor, instr};
   absl::MutexLock lock(&mutex_);
-  VLOG(3) << "Emplace event " << event.implementation();
+  VLOG(3) << "Emplace event " << event->implementation();
   if (auto [it, inserted] = events_.try_emplace(key, std::move(event));
       inserted) {
     return absl::OkStatus();
@@ -92,12 +92,12 @@ absl::Status CopyThunk::AsyncEvents::Emplace(se::StreamExecutor* executor,
 
 // Retrieve a completion event started by copy-start instruction
 // `instr`, and remove the event from the collection.
-absl::StatusOr<se::Event> CopyThunk::AsyncEvents::Extract(
+absl::StatusOr<std::unique_ptr<se::Event>> CopyThunk::AsyncEvents::Extract(
     se::StreamExecutor* executor, const HloInstruction* instr) {
   Key key = {executor, instr};
   absl::MutexLock lock(&mutex_);
   if (auto event = events_.extract(key)) {
-    VLOG(3) << "Extract event " << event.mapped().implementation();
+    VLOG(3) << "Extract event " << event.mapped()->implementation();
     return std::move(event.mapped());
   }
   return absl::InternalError("Async copy event was not found!");
@@ -133,14 +133,10 @@ absl::Status DeviceToHostCopyThunk::ExecuteOnStream(
   }
   VLOG(2) << "Memcpy D2H from the other stream";
   se::StreamExecutor* executor = params.stream->parent();
-  se::Event event(executor);
-  if (!event.Init()) {
-    return absl::InternalError(
-        "Failed to initialize copy operation async completion event!");
-  }
+  TF_ASSIGN_OR_RETURN(auto event, executor->CreateEvent());
   // Record memcpy operation completion.
-  TF_RETURN_IF_ERROR(stream->RecordEvent(&event));
-  VLOG(3) << "Emplace events: " << event.implementation()
+  TF_RETURN_IF_ERROR(stream->RecordEvent(event.get()));
+  VLOG(3) << "Emplace events: " << event->implementation()
           << " for instr: " << instr_->ToString();
   return async_events_->Emplace(executor, instr_, std::move(event));
 }
@@ -175,14 +171,10 @@ absl::Status HostToDeviceCopyThunk::ExecuteOnStream(
   }
   VLOG(2) << "Memcpy H2D from the other stream";
   se::StreamExecutor* executor = params.stream->parent();
-  se::Event event(executor);
-  if (!event.Init()) {
-    return absl::InternalError(
-        "Failed to initialize copy operation async completion event!");
-  }
+  TF_ASSIGN_OR_RETURN(auto event, executor->CreateEvent());
   // Record memcpy operation completion.
-  TF_RETURN_IF_ERROR(stream->RecordEvent(&event));
-  VLOG(3) << "Emplace events: " << event.implementation()
+  TF_RETURN_IF_ERROR(stream->RecordEvent(event.get()));
+  VLOG(3) << "Emplace events: " << event->implementation()
           << " for instr: " << instr_->ToString();
   return async_events_->Emplace(executor, instr_, std::move(event));
 }
@@ -202,9 +194,9 @@ absl::Status CopyDoneThunk::ExecuteOnStream(const ExecuteParams& params) {
   VLOG(3) << "CopyDone thunk between a host and a device for: "
           << copy_start_instr_->ToString();
   se::StreamExecutor* executor = params.stream->parent();
-  TF_ASSIGN_OR_RETURN(se::Event event,
+  TF_ASSIGN_OR_RETURN(std::unique_ptr<se::Event> event,
                       async_events_->Extract(executor, copy_start_instr_));
-  return params.stream->WaitFor(&event);
+  return params.stream->WaitFor(event.get());
 }
 
 }  // namespace gpu
diff --git a/third_party/xla/xla/service/gpu/runtime/copy_thunk.h b/third_party/xla/xla/service/gpu/runtime/copy_thunk.h
index d8490f6280a8c7..4bbd962ed6e91a 100644
--- a/third_party/xla/xla/service/gpu/runtime/copy_thunk.h
+++ b/third_party/xla/xla/service/gpu/runtime/copy_thunk.h
@@ -68,17 +68,19 @@ class CopyThunk : public Thunk {
    public:
     // Add a new copy-start completion event.
     absl::Status Emplace(se::StreamExecutor* executor,
-                         const HloInstruction* instr, se::Event&& event);
+                         const HloInstruction* instr,
+                         std::unique_ptr<se::Event> event);
 
     // Retrieve a completion event started by copy-start instruction
     // `instr`, and remove the event from the collection.
-    absl::StatusOr<se::Event> Extract(se::StreamExecutor* executor,
-                                      const HloInstruction* instr);
+    absl::StatusOr<std::unique_ptr<se::Event>> Extract(
+        se::StreamExecutor* executor, const HloInstruction* instr);
 
    private:
     using Key = std::pair<se::StreamExecutor*, const HloInstruction*>;
     absl::Mutex mutex_;
-    absl::flat_hash_map<Key, se::Event> events_ ABSL_GUARDED_BY(mutex_);
+    absl::flat_hash_map<Key, std::unique_ptr<se::Event>> events_
+        ABSL_GUARDED_BY(mutex_);
   };
   CopyThunk(ThunkInfo thunk_info, const BufferAllocation::Slice& source_buffer,
             const BufferAllocation::Slice& destination_buffer,
diff --git a/third_party/xla/xla/service/gpu/runtime/nccl_collective_permute_thunk.h b/third_party/xla/xla/service/gpu/runtime/nccl_collective_permute_thunk.h
index 425777f68df187..98719401a57aff 100644
--- a/third_party/xla/xla/service/gpu/runtime/nccl_collective_permute_thunk.h
+++ b/third_party/xla/xla/service/gpu/runtime/nccl_collective_permute_thunk.h
@@ -18,6 +18,7 @@ limitations under the License.
 
 #include <cstdint>
 
+#include "absl/container/node_hash_map.h"
 #include "absl/status/status.h"
 #include "absl/strings/string_view.h"
 #include "xla/hlo/ir/hlo_instructions.h"
diff --git a/third_party/xla/xla/service/gpu/runtime/nccl_collective_thunk.cc b/third_party/xla/xla/service/gpu/runtime/nccl_collective_thunk.cc
index 453fc21d305d1f..f4c1ebf3df99dd 100644
--- a/third_party/xla/xla/service/gpu/runtime/nccl_collective_thunk.cc
+++ b/third_party/xla/xla/service/gpu/runtime/nccl_collective_thunk.cc
@@ -353,11 +353,7 @@ absl::Status NcclCollectiveThunk::AsyncEvents::Initialize(
   absl::MutexLock lock(&mu_);
   if (events_.contains(executor)) return absl::OkStatus();
 
-  se::Event event(executor);
-  if (!event.Init()) {
-    return absl::InternalError(
-        "Failed to initialize collective operation async completion event");
-  }
+  TF_ASSIGN_OR_RETURN(auto event, executor->CreateEvent());
 
   events_.try_emplace(executor, std::move(event));
   return absl::OkStatus();
@@ -373,7 +369,7 @@ absl::StatusOr<se::Event*> NcclCollectiveThunk::AsyncEvents::GetEvent(
         "Collective operation async completion event not initialized");
   }
 
-  return &event->second;
+  return event->second.get();
 }
 
 absl::Status NcclCollectiveThunk::Prepare(const PrepareParams& params,
diff --git a/third_party/xla/xla/service/gpu/runtime/nccl_collective_thunk.h b/third_party/xla/xla/service/gpu/runtime/nccl_collective_thunk.h
index c53fbd8a5db9d0..977f791b87328f 100644
--- a/third_party/xla/xla/service/gpu/runtime/nccl_collective_thunk.h
+++ b/third_party/xla/xla/service/gpu/runtime/nccl_collective_thunk.h
@@ -25,7 +25,7 @@ limitations under the License.
 #include <vector>
 
 #include "absl/base/thread_annotations.h"
-#include "absl/container/node_hash_map.h"
+#include "absl/container/flat_hash_map.h"
 #include "absl/status/status.h"
 #include "absl/status/statusor.h"
 #include "absl/synchronization/mutex.h"
@@ -155,7 +155,7 @@ class NcclCollectiveThunk : public Thunk {
 
    private:
     absl::Mutex mu_;
-    absl::node_hash_map<se::StreamExecutor*, se::Event> events_
+    absl::flat_hash_map<se::StreamExecutor*, std::unique_ptr<se::Event>> events_
         ABSL_GUARDED_BY(mu_);
   };
 
diff --git a/third_party/xla/xla/stream_executor/BUILD b/third_party/xla/xla/stream_executor/BUILD
index 7137b54991ad52..248876a1104577 100644
--- a/third_party/xla/xla/stream_executor/BUILD
+++ b/third_party/xla/xla/stream_executor/BUILD
@@ -553,20 +553,6 @@ cc_library(
     ],
 )
 
-cc_library(
-    name = "event",
-    srcs = ["event.cc"],
-    hdrs = ["event.h"],
-    visibility = ["//visibility:private"],
-    deps = [
-        ":event_interface",
-        ":stream_executor_headers",
-        ":stream_executor_interface",
-        "@com_google_absl//absl/log",
-        "@com_google_absl//absl/status",
-    ],
-)
-
 cc_library(
     name = "executor_cache",
     srcs = ["executor_cache.cc"],
@@ -680,6 +666,7 @@ transitive_hdrs(
 cc_library(
     name = "stream_executor_pimpl",
     srcs = [
+        "event.cc",
         "stream.cc",
         "stream_executor_pimpl.cc",
     ],
@@ -690,6 +677,7 @@ cc_library(
         ":blas",  # build_cleaner: keep
         ":command_buffer",  # build_cleaner: keep
         ":dnn",  # build_cleaner: keep
+        ":event_interface",
         ":fft",
         ":host_memory_allocation",  # build_cleaner: keep
         ":kernel_spec",
@@ -728,7 +716,6 @@ cc_library(
         ":device_description",
         ":device_memory",
         ":dnn",
-        ":event",
         ":executor_cache",
         ":kernel",
         ":kernel_spec",
@@ -843,7 +830,6 @@ cc_library(
     visibility = [":internal"],
     deps = [
         ":dnn",
-        ":event",
         ":platform_manager",
         ":scratch_allocator",
         ":stream_executor",
diff --git a/third_party/xla/xla/stream_executor/mock_stream_executor.h b/third_party/xla/xla/stream_executor/mock_stream_executor.h
index 89464874408753..21f9147b9f6440 100644
--- a/third_party/xla/xla/stream_executor/mock_stream_executor.h
+++ b/third_party/xla/xla/stream_executor/mock_stream_executor.h
@@ -177,6 +177,8 @@ class MockStreamExecutor : public StreamExecutorInterface {
   MOCK_METHOD(int64_t, GetMemoryLimitBytes, (), (const.override));
   MOCK_METHOD(const DeviceDescription&, GetDeviceDescription, (),
               (const, override));
+  MOCK_METHOD(absl::StatusOr<std::unique_ptr<Event>>, CreateEvent, (),
+              (override));
 };
 
 }  // namespace stream_executor
diff --git a/third_party/xla/xla/stream_executor/stream_executor_interface.h b/third_party/xla/xla/stream_executor/stream_executor_interface.h
index f310cce8b7eaca..628577725014d5 100644
--- a/third_party/xla/xla/stream_executor/stream_executor_interface.h
+++ b/third_party/xla/xla/stream_executor/stream_executor_interface.h
@@ -42,6 +42,7 @@ limitations under the License.
 
 namespace stream_executor {
 
+class Event;
 class Stream;
 
 // Interface which defines the method for interacting with an accelerator device
@@ -65,6 +66,9 @@ class StreamExecutorInterface {
       std::optional<std::variant<StreamPriority, int>> priority =
           std::nullopt) = 0;
 
+  // Creates and initializes an Event.
+  virtual absl::StatusOr<std::unique_ptr<Event>> CreateEvent() = 0;
+
   // Obtains metadata about the underlying device.
   // The value is cached on first use.
   virtual const DeviceDescription& GetDeviceDescription() const = 0;
diff --git a/third_party/xla/xla/stream_executor/stream_executor_pimpl.cc b/third_party/xla/xla/stream_executor/stream_executor_pimpl.cc
index efd27f4ec49f48..2866a7d7d50aeb 100644
--- a/third_party/xla/xla/stream_executor/stream_executor_pimpl.cc
+++ b/third_party/xla/xla/stream_executor/stream_executor_pimpl.cc
@@ -78,4 +78,13 @@ const DeviceDescription& StreamExecutor::GetDeviceDescription() const {
   return *device_description_;
 }
 
+absl::StatusOr<std::unique_ptr<Event>> StreamExecutor::CreateEvent() {
+  auto event = std::make_unique<Event>(this);
+  if (event->Init()) {
+    return std::move(event);
+  }
+
+  return absl::InternalError("Failed to create event.");
+}
+
 }  // namespace stream_executor
diff --git a/third_party/xla/xla/stream_executor/stream_executor_pimpl.h b/third_party/xla/xla/stream_executor/stream_executor_pimpl.h
index f57ace700cc9a0..41546f59ea5321 100644
--- a/third_party/xla/xla/stream_executor/stream_executor_pimpl.h
+++ b/third_party/xla/xla/stream_executor/stream_executor_pimpl.h
@@ -71,6 +71,7 @@ class StreamExecutor : public StreamExecutorInterface {
   const Platform* GetPlatform() const override { return platform_; }
   const DeviceDescription& GetDeviceDescription() const override;
   int64_t GetMemoryLimitBytes() const override { return memory_limit_bytes_; }
+  absl::StatusOr<std::unique_ptr<Event>> CreateEvent() override;
 
  private:
   // Reader/writer lock for mutable data structures on this StreamExecutor.

From 87f576ff3f951899286a8a92637b0363b5a6684c Mon Sep 17 00:00:00 2001
From: Kyle Lucke <klucke@google.com>
Date: Wed, 15 May 2024 11:07:53 -0700
Subject: [PATCH 195/478] Use absl::Status instead of xla::Status now that
 they're identical.

PiperOrigin-RevId: 634012954
---
 .../xla/xla/tests/client_library_test_base.cc       | 13 +++++++------
 .../xla/xla/tests/client_library_test_base.h        |  8 ++++----
 third_party/xla/xla/tests/compute_constant_test.cc  |  5 +++--
 third_party/xla/xla/tests/conv_depthwise_test.cc    |  2 +-
 .../xla/xla/tests/grouped_convolution_test.cc       |  2 +-
 third_party/xla/xla/tests/hlo_test_base.cc          |  6 +++---
 third_party/xla/xla/tests/hlo_test_base.h           |  4 ++--
 third_party/xla/xla/tests/literal_test_util.cc      |  2 +-
 third_party/xla/xla/tests/llvm_compiler_test.cc     |  4 ++--
 third_party/xla/xla/tests/llvm_irgen_test_base.cc   |  6 +++---
 third_party/xla/xla/tests/llvm_irgen_test_base.h    |  2 +-
 third_party/xla/xla/tests/local_client_test_base.cc |  3 ++-
 third_party/xla/xla/tests/local_client_test_base.h  |  3 ++-
 third_party/xla/xla/tests/test_utils.cc             |  6 +++---
 third_party/xla/xla/tests/test_utils.h              |  4 ++--
 third_party/xla/xla/tests/verified_hlo_module.cc    |  7 ++++---
 third_party/xla/xla/tests/verified_hlo_module.h     |  4 ++--
 third_party/xla/xla/tests/xla_hlo_profile_test.cc   |  2 +-
 18 files changed, 44 insertions(+), 39 deletions(-)

diff --git a/third_party/xla/xla/tests/client_library_test_base.cc b/third_party/xla/xla/tests/client_library_test_base.cc
index 86dc33b33d1958..7df9f2d1d9d530 100644
--- a/third_party/xla/xla/tests/client_library_test_base.cc
+++ b/third_party/xla/xla/tests/client_library_test_base.cc
@@ -181,7 +181,8 @@ void ClientLibraryTestBase::ComputeAndCompareLiteral(
                                                   error, shape_with_layout));
 }
 
-Status ClientLibraryTestBase::ComputeAndCompareLiteralWithAllOutputLayouts(
+absl::Status
+ClientLibraryTestBase::ComputeAndCompareLiteralWithAllOutputLayouts(
     const xla::XlaComputation& computation, const Literal& expected,
     absl::Span<GlobalData* const> arguments,
     const std::function<void(const Literal& actual,
@@ -207,7 +208,7 @@ Status ClientLibraryTestBase::ComputeAndCompareLiteralWithAllOutputLayouts(
   return OkStatus();
 }
 
-Status ClientLibraryTestBase::ComputeAndCompareLiteralWithAllInputLayouts(
+absl::Status ClientLibraryTestBase::ComputeAndCompareLiteralWithAllInputLayouts(
     const xla::XlaComputation& computation, const Literal& /*expected*/,
     absl::Span<GlobalData* const> arguments,
     const std::function<void(const Literal& actual,
@@ -218,8 +219,8 @@ Status ClientLibraryTestBase::ComputeAndCompareLiteralWithAllInputLayouts(
   // This is a recursive function. It's an std::function instead of a lambda
   // because it needs to capture itself. The index is the index of the argument
   // to try all layouts for.
-  std::function<Status(int64_t)> choose;
-  choose = [&, this](int64_t index) -> Status {
+  std::function<absl::Status(int64_t)> choose;
+  choose = [&, this](int64_t index) -> absl::Status {
     if (index < arguments.size()) {
       // Try out all layouts for the operand.
       TF_ASSIGN_OR_RETURN(auto literal,
@@ -293,7 +294,7 @@ absl::StatusOr<Literal> ClientLibraryTestBase::ComputeAndTransfer(
   return ExecuteAndTransfer(computation, arguments, shape_with_layout);
 }
 
-Status ClientLibraryTestBase::ComputeAndCompareLiteralWithStatus(
+absl::Status ClientLibraryTestBase::ComputeAndCompareLiteralWithStatus(
     XlaBuilder* builder, const Literal& expected,
     absl::Span<GlobalData* const> arguments_passed_in,
     const Shape* shape_with_layout) {
@@ -355,7 +356,7 @@ Status ClientLibraryTestBase::ComputeAndCompareLiteralWithStatus(
   return OkStatus();
 }
 
-Status ClientLibraryTestBase::ComputeAndCompareLiteralWithStatus(
+absl::Status ClientLibraryTestBase::ComputeAndCompareLiteralWithStatus(
     XlaBuilder* builder, const Literal& expected,
     absl::Span<GlobalData* const> arguments_passed_in, ErrorSpec error,
     const Shape* shape_with_layout) {
diff --git a/third_party/xla/xla/tests/client_library_test_base.h b/third_party/xla/xla/tests/client_library_test_base.h
index b7be83b5c9c40f..e60ea01b493b44 100644
--- a/third_party/xla/xla/tests/client_library_test_base.h
+++ b/third_party/xla/xla/tests/client_library_test_base.h
@@ -198,11 +198,11 @@ class ClientLibraryTestBase : public ManifestCheckingTest {
       const Shape* shape_with_layout = nullptr);
 
   // ComputeAndCompare variant which returns an error status.
-  Status ComputeAndCompareLiteralWithStatus(
+  absl::Status ComputeAndCompareLiteralWithStatus(
       XlaBuilder* builder, const Literal& expected,
       absl::Span<GlobalData* const> arguments,
       const Shape* shape_with_layout = nullptr);
-  Status ComputeAndCompareLiteralWithStatus(
+  absl::Status ComputeAndCompareLiteralWithStatus(
       XlaBuilder* builder, const Literal& expected,
       absl::Span<GlobalData* const> arguments, ErrorSpec error,
       const Shape* shape_with_layout = nullptr);
@@ -422,13 +422,13 @@ class ClientLibraryTestBase : public ManifestCheckingTest {
   ExecutionOptions execution_options_;
 
  private:
-  Status ComputeAndCompareLiteralWithAllOutputLayouts(
+  absl::Status ComputeAndCompareLiteralWithAllOutputLayouts(
       const xla::XlaComputation& computation, const Literal& expected,
       absl::Span<GlobalData* const> arguments,
       const std::function<void(const Literal& actual,
                                const std::string& error_message)>&
           verify_output);
-  Status ComputeAndCompareLiteralWithAllInputLayouts(
+  absl::Status ComputeAndCompareLiteralWithAllInputLayouts(
       const xla::XlaComputation& computation, const Literal& expected,
       absl::Span<GlobalData* const> arguments,
       const std::function<void(const Literal& actual,
diff --git a/third_party/xla/xla/tests/compute_constant_test.cc b/third_party/xla/xla/tests/compute_constant_test.cc
index b37cf2d9b83950..53946543cdcccc 100644
--- a/third_party/xla/xla/tests/compute_constant_test.cc
+++ b/third_party/xla/xla/tests/compute_constant_test.cc
@@ -78,8 +78,9 @@ class ComputeConstantTest : public ::testing::Test {
   }
 
   template <class Scalar>
-  StatusOr<Scalar> ComputeConstantScalar(Client* client, const XlaOp operand,
-                                         XlaBuilder* builder) {
+  absl::StatusOr<Scalar> ComputeConstantScalar(Client* client,
+                                               const XlaOp operand,
+                                               XlaBuilder* builder) {
     TF_ASSIGN_OR_RETURN(auto literal, ComputeConstantLiteral(client, operand,
                                                              builder, nullptr));
     return literal.Get<Scalar>({});
diff --git a/third_party/xla/xla/tests/conv_depthwise_test.cc b/third_party/xla/xla/tests/conv_depthwise_test.cc
index b44e089b596ce9..e1826bd902505d 100644
--- a/third_party/xla/xla/tests/conv_depthwise_test.cc
+++ b/third_party/xla/xla/tests/conv_depthwise_test.cc
@@ -97,7 +97,7 @@ XLA_TEST_P(DepthwiseConvolution2DTest, DoIt) {
       BuildHloTextDepthwiseConvolution2D(spec, use_bfloat16);
 
   EXPECT_TRUE(RunAndCompare(hlo_text, ErrorSpec{0.01, 0.01},
-                            [](HloModule* module) -> Status {
+                            [](HloModule* module) -> absl::Status {
                               BFloat16MixedPrecisionRemoval remover;
                               TF_RETURN_IF_ERROR(remover.Run(module).status());
                               Despecializer despecializer;
diff --git a/third_party/xla/xla/tests/grouped_convolution_test.cc b/third_party/xla/xla/tests/grouped_convolution_test.cc
index 45cc3fe4112eb5..24f41c3c96cfeb 100644
--- a/third_party/xla/xla/tests/grouped_convolution_test.cc
+++ b/third_party/xla/xla/tests/grouped_convolution_test.cc
@@ -241,7 +241,7 @@ XLA_TEST_P(GroupedConvolution2DTest, DoIt) {
       BuildHloTextGroupedConvolution2D(spec, use_bfloat16);
 
   EXPECT_TRUE(RunAndCompare(hlo_text, ErrorSpec{0.01, 0.01},
-                            [](HloModule* module) -> Status {
+                            [](HloModule* module) -> absl::Status {
                               BFloat16MixedPrecisionRemoval remover;
                               TF_RETURN_IF_ERROR(remover.Run(module).status());
                               Despecializer despecializer;
diff --git a/third_party/xla/xla/tests/hlo_test_base.cc b/third_party/xla/xla/tests/hlo_test_base.cc
index ab9888f23fd419..12f8b6fde817b1 100644
--- a/third_party/xla/xla/tests/hlo_test_base.cc
+++ b/third_party/xla/xla/tests/hlo_test_base.cc
@@ -687,7 +687,7 @@ ::testing::AssertionResult HloTestBase::Run(
     // Set backend configuration if it is given.
     HloInstruction* instruction =
         module->entry_computation()->root_instruction();
-    Status s = instruction->set_backend_config(*backend_config);
+    absl::Status s = instruction->set_backend_config(*backend_config);
     return s.ok() ? ::testing::AssertionSuccess()
                   : ::testing::AssertionFailure() << s.message();
   }
@@ -723,7 +723,7 @@ ::testing::AssertionResult HloTestBase::RunReplicated(
     // Set backend configuration if it is given.
     HloInstruction* instruction =
         module->entry_computation()->root_instruction();
-    Status s = instruction->set_backend_config(*backend_config);
+    absl::Status s = instruction->set_backend_config(*backend_config);
     return s.ok() ? ::testing::AssertionSuccess()
                   : ::testing::AssertionFailure() << s.message();
   }
@@ -779,7 +779,7 @@ ::testing::AssertionResult HloTestBase::RunMultipleTimes(
       // Set backend configuration if it is given.
       HloInstruction* instruction =
           module->entry_computation()->root_instruction();
-      Status s = instruction->set_backend_config(*backend_config);
+      absl::Status s = instruction->set_backend_config(*backend_config);
       return s.ok() ? ::testing::AssertionSuccess()
                     : ::testing::AssertionFailure() << s.message();
     }
diff --git a/third_party/xla/xla/tests/hlo_test_base.h b/third_party/xla/xla/tests/hlo_test_base.h
index d5207a0d638a38..68b44a208c3470 100644
--- a/third_party/xla/xla/tests/hlo_test_base.h
+++ b/third_party/xla/xla/tests/hlo_test_base.h
@@ -101,7 +101,7 @@ class HloTestBase : public ManifestCheckingTest {
 
   // Runs the hlo_pass with the provided module and returns the result. This
   // function also verifies that the module remains unchanged when hlo_pass
-  // returns false as the StatusOr value.
+  // returns false as the absl::StatusOr value.
   //
   // These three overloads all do the same thing.  The && overload lets you do
   // `RunHloPass(MyPass(), module)` all in one line.  The reason for the
@@ -121,7 +121,7 @@ class HloTestBase : public ManifestCheckingTest {
   // Runs the hlo_pass with the provided module group and returns the result.
   // This method runs the input HLO module group pass for a `HloModuleGroup` and
   // it also verifies the module group remains unchanged when hlo_pass returns
-  // false as the StatusOr value.
+  // false as the absl::StatusOr value.
   static absl::StatusOr<bool> RunHloPass(HloPassInterface&& hlo_pass,
                                          HloModuleGroup* module_group);
 
diff --git a/third_party/xla/xla/tests/literal_test_util.cc b/third_party/xla/xla/tests/literal_test_util.cc
index df898b1cedea84..1d5c99fd712e88 100644
--- a/third_party/xla/xla/tests/literal_test_util.cc
+++ b/third_party/xla/xla/tests/literal_test_util.cc
@@ -63,7 +63,7 @@ void OnMiscompare(const LiteralSlice& expected, const LiteralSlice& actual,
   WriteLiteralToTempFile(mismatches, "mismatches");
 }
 
-::testing::AssertionResult StatusToAssertion(const Status& s) {
+::testing::AssertionResult StatusToAssertion(const absl::Status& s) {
   if (s.ok()) {
     return ::testing::AssertionSuccess();
   }
diff --git a/third_party/xla/xla/tests/llvm_compiler_test.cc b/third_party/xla/xla/tests/llvm_compiler_test.cc
index 32b2a5be85e3d1..70274abf1c6ac3 100644
--- a/third_party/xla/xla/tests/llvm_compiler_test.cc
+++ b/third_party/xla/xla/tests/llvm_compiler_test.cc
@@ -59,14 +59,14 @@ class GpuDummyCompiler : public GpuCompiler {
 
   int32_t GetToolkitVersion() const override { return 0; }
 
-  Status OptimizeHloConvolutionCanonicalization(
+  absl::Status OptimizeHloConvolutionCanonicalization(
       HloModule* hlo_module, se::GpuComputeCapability gpu_version,
       se::dnn::VersionInfo dnn_version,
       se::DeviceMemoryAllocator* device_allocator) {
     return OkStatus();
   }
 
-  Status OptimizeHloPostLayoutAssignment(
+  absl::Status OptimizeHloPostLayoutAssignment(
       HloModule* hlo_module, se::StreamExecutor* stream_executor,
       const CompileOptions& options, const TargetConfig& gpu_target_config,
       tsl::thread::ThreadPool* thread_pool) override {
diff --git a/third_party/xla/xla/tests/llvm_irgen_test_base.cc b/third_party/xla/xla/tests/llvm_irgen_test_base.cc
index 59b16529b4df1f..b6c7694d353237 100644
--- a/third_party/xla/xla/tests/llvm_irgen_test_base.cc
+++ b/third_party/xla/xla/tests/llvm_irgen_test_base.cc
@@ -50,7 +50,7 @@ void LlvmIrGenTestBase::CompileAndVerifyIr(
     std::unique_ptr<HloModule> hlo_module, const std::string& pattern,
     bool match_optimized_ir, bool run_optimization_passes) {
   SetIrHook(match_optimized_ir);
-  Status status =
+  absl::Status status =
       CompileToExecutable(std::move(hlo_module), run_optimization_passes)
           .status();
   ResetIrHook();
@@ -77,7 +77,7 @@ void LlvmIrGenTestBase::CompileAheadOfTimeAndVerifyIr(
     std::unique_ptr<HloModule> hlo_module, const AotCompilationOptions& options,
     const std::string& pattern, bool match_optimized_ir) {
   SetIrHook(match_optimized_ir);
-  Status status =
+  absl::Status status =
       CompileToAotCompilationResult(std::move(hlo_module), options).status();
   ResetIrHook();
   TF_ASSERT_OK(status);
@@ -91,7 +91,7 @@ LLVMCompiler* LlvmIrGenTestBase::GetLLVMCompiler() {
   return static_cast<LLVMCompiler*>(backend().compiler());
 }
 
-Status LlvmIrGenTestBase::IrHook(const llvm::Module& module) {
+absl::Status LlvmIrGenTestBase::IrHook(const llvm::Module& module) {
   ir_ = llvm_ir::DumpToString(&module);
   return OkStatus();
 }
diff --git a/third_party/xla/xla/tests/llvm_irgen_test_base.h b/third_party/xla/xla/tests/llvm_irgen_test_base.h
index 890e5bacb65ed5..718eb360490ef4 100644
--- a/third_party/xla/xla/tests/llvm_irgen_test_base.h
+++ b/third_party/xla/xla/tests/llvm_irgen_test_base.h
@@ -67,7 +67,7 @@ class LlvmIrGenTestBase : public CodegenTestBase {
   void ResetIrHook();
 
   std::string ir_;
-  Status IrHook(const llvm::Module& module);
+  absl::Status IrHook(const llvm::Module& module);
 };
 
 }  // namespace xla
diff --git a/third_party/xla/xla/tests/local_client_test_base.cc b/third_party/xla/xla/tests/local_client_test_base.cc
index 8df3e5eb11aca0..60d0310fa7d0b5 100644
--- a/third_party/xla/xla/tests/local_client_test_base.cc
+++ b/third_party/xla/xla/tests/local_client_test_base.cc
@@ -52,7 +52,8 @@ absl::StatusOr<se::OwningDeviceMemory> TestAllocator::Allocate(
       device_ordinal, size, retry_on_failure, memory_space);
 }
 
-Status TestAllocator::Deallocate(int device_ordinal, se::DeviceMemoryBase mem) {
+absl::Status TestAllocator::Deallocate(int device_ordinal,
+                                       se::DeviceMemoryBase mem) {
   VLOG(2) << "Deallocate(" << device_ordinal << ")";
   {
     absl::MutexLock lock(&count_mutex_);
diff --git a/third_party/xla/xla/tests/local_client_test_base.h b/third_party/xla/xla/tests/local_client_test_base.h
index d2e58ff445eb94..02504d0ab0b7bd 100644
--- a/third_party/xla/xla/tests/local_client_test_base.h
+++ b/third_party/xla/xla/tests/local_client_test_base.h
@@ -52,7 +52,8 @@ class TestAllocator : public se::StreamExecutorMemoryAllocator {
   absl::StatusOr<se::OwningDeviceMemory> Allocate(
       int device_ordinal, uint64_t size, bool retry_on_failure,
       int64_t memory_space) override;
-  Status Deallocate(int device_ordinal, se::DeviceMemoryBase mem) override;
+  absl::Status Deallocate(int device_ordinal,
+                          se::DeviceMemoryBase mem) override;
 
   // Return the number of allocations that have been performed.
   int64_t allocation_count() const;
diff --git a/third_party/xla/xla/tests/test_utils.cc b/third_party/xla/xla/tests/test_utils.cc
index d2934f0395c6e3..b078841d26182c 100644
--- a/third_party/xla/xla/tests/test_utils.cc
+++ b/third_party/xla/xla/tests/test_utils.cc
@@ -261,7 +261,7 @@ absl::StatusOr<Literal> MakeFakeLiteralInternal(
   Literal literal(new_shape);
 
   TF_RETURN_IF_ERROR(primitive_util::PrimitiveTypeSwitch<Status>(
-      [&](auto primitive_type_constant) -> Status {
+      [&](auto primitive_type_constant) -> absl::Status {
         if constexpr (primitive_util::IsArrayType(primitive_type_constant)) {
           using NativeT = primitive_util::NativeTypeOf<primitive_type_constant>;
           if constexpr (primitive_util::IsFloatingPointType(
@@ -639,8 +639,8 @@ absl::StatusOr<std::vector<Literal>> MakeFakeArguments(
   return std::move(arguments);
 }
 
-Status VerifyHloModule(HloModule* const module, bool layout_sensitive,
-                       bool allow_mixed_precision) {
+absl::Status VerifyHloModule(HloModule* const module, bool layout_sensitive,
+                             bool allow_mixed_precision) {
   return HloVerifier(/*layout_sensitive=*/layout_sensitive,
                      /*allow_mixed_precision=*/allow_mixed_precision)
       .Run(module)
diff --git a/third_party/xla/xla/tests/test_utils.h b/third_party/xla/xla/tests/test_utils.h
index fb101102d6381d..8ea594511cf3fa 100644
--- a/third_party/xla/xla/tests/test_utils.h
+++ b/third_party/xla/xla/tests/test_utils.h
@@ -109,8 +109,8 @@ absl::StatusOr<std::vector<Literal>> MakeFakeArguments(
 
 // Check that a given module satisfies various constraints before trying to
 // execute it.
-Status VerifyHloModule(HloModule* const module, bool layout_sensitive,
-                       bool allow_mixed_precision);
+absl::Status VerifyHloModule(HloModule* const module, bool layout_sensitive,
+                             bool allow_mixed_precision);
 
 // Creates a dot op with operands 'lhs' and 'rhs' that contracts dimension 1 of
 // the LHS with dimension 0 of the RHS with no batch dimensions.
diff --git a/third_party/xla/xla/tests/verified_hlo_module.cc b/third_party/xla/xla/tests/verified_hlo_module.cc
index 3eaacbc0ac7a1f..86635697a9e70d 100644
--- a/third_party/xla/xla/tests/verified_hlo_module.cc
+++ b/third_party/xla/xla/tests/verified_hlo_module.cc
@@ -26,7 +26,8 @@ limitations under the License.
 
 namespace xla {
 
-Status VerifiedHloModule::ParseHloStringAndVerifyModule(absl::string_view str) {
+absl::Status VerifiedHloModule::ParseHloStringAndVerifyModule(
+    absl::string_view str) {
   TF_RET_CHECK(computation_count() == 0);
   auto parser = HloParser::CreateHloParserForTests(str);
   TF_RETURN_IF_ERROR(parser->Run(this));
@@ -34,7 +35,7 @@ Status VerifiedHloModule::ParseHloStringAndVerifyModule(absl::string_view str) {
 }
 
 void VerifiedHloModule::VerifyOrAddFailure(absl::string_view message) {
-  Status status = Verify();
+  absl::Status status = Verify();
   if (!status.ok()) {
     ADD_FAILURE() << "HloVerifier failed on module " << name()
                   << (message.empty() ? "" : absl::StrCat(" (", message, ")"))
@@ -44,7 +45,7 @@ void VerifiedHloModule::VerifyOrAddFailure(absl::string_view message) {
   }
 }
 
-Status VerifiedHloModule::Verify() {
+absl::Status VerifiedHloModule::Verify() {
   if (computation_count() == 0) {
     // The computation was never built. Nothing to verify.
     return OkStatus();
diff --git a/third_party/xla/xla/tests/verified_hlo_module.h b/third_party/xla/xla/tests/verified_hlo_module.h
index 52fce003b9ccf1..837166a36d3e1d 100644
--- a/third_party/xla/xla/tests/verified_hlo_module.h
+++ b/third_party/xla/xla/tests/verified_hlo_module.h
@@ -48,14 +48,14 @@ class VerifiedHloModule : public HloModule {
   // builds the VerifiedHloModule in place. Before calling this method, the
   // module must be empty (no computations). Finally verifies the module using
   // HloVerifier and returns the status.
-  Status ParseHloStringAndVerifyModule(absl::string_view str);
+  absl::Status ParseHloStringAndVerifyModule(absl::string_view str);
 
   // Verifies the module and flags any error with ADD_FAILURE. 'message' is
   // included in the failure message.
   void VerifyOrAddFailure(absl::string_view message);
 
   // Verifies the module using HloVerifier and returns the status.
-  Status Verify();
+  absl::Status Verify();
 
  private:
   HloVerifier verifier_;
diff --git a/third_party/xla/xla/tests/xla_hlo_profile_test.cc b/third_party/xla/xla/tests/xla_hlo_profile_test.cc
index 2806e63f94c074..5533b95f3c265c 100644
--- a/third_party/xla/xla/tests/xla_hlo_profile_test.cc
+++ b/third_party/xla/xla/tests/xla_hlo_profile_test.cc
@@ -78,7 +78,7 @@ ::testing::AssertionResult HasTrops(
          << parsed_line.trops << "'";
 }
 
-Status ParseOneProfileOutputLine(
+absl::Status ParseOneProfileOutputLine(
     const std::string& line, bool expect_hlo,
     absl::flat_hash_map<std::string, ParsedProfileOutputLine>* parsed_results,
     absl::Span<const absl::string_view> opcodes_to_ignore = {}) {

From 14bb818be7745c5df9d9046b0353f75520899a59 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Wed, 15 May 2024 11:12:00 -0700
Subject: [PATCH 196/478] Disable pjrt runner for scatter test as it fails on
 some cases

PiperOrigin-RevId: 634014544
---
 third_party/xla/xla/tests/BUILD | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/third_party/xla/xla/tests/BUILD b/third_party/xla/xla/tests/BUILD
index 58600c91d25852..563a8ea6b203c7 100644
--- a/third_party/xla/xla/tests/BUILD
+++ b/third_party/xla/xla/tests/BUILD
@@ -994,6 +994,7 @@ xla_test(
 xla_test(
     name = "gather_operation_test",
     srcs = ["gather_operation_test.cc"],
+    shard_count = 20,
     tags = ["test_hlo_pjrt_runner"],
     deps = [
         ":client_library_test_base",
@@ -1012,7 +1013,7 @@ xla_test(
 xla_test(
     name = "scatter_test",
     srcs = ["scatter_test.cc"],
-    tags = ["test_hlo_pjrt_runner"],
+    # TODO(b/245550554): enable Pjrt runner for scatter test once it's fixed.
     deps = [
         ":client_library_test_base",
         ":hlo_test_base",

From 683ff7d2f35d37c96a1bc9265f5c4349c5ddedb0 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Wed, 15 May 2024 11:15:21 -0700
Subject: [PATCH 197/478] Reduces the # of memory constraints to a subset of
 times (instead of all possible times).

PiperOrigin-RevId: 634015658
---
 .../auto_sharding/auto_sharding_memory.cc     |  44 ++++++++
 .../auto_sharding/auto_sharding_memory.h      |  12 ++
 .../auto_sharding_memory_test.cc              | 106 ++++++++++++++----
 .../auto_sharding/auto_sharding_solver.cc     |  19 +++-
 4 files changed, 155 insertions(+), 26 deletions(-)

diff --git a/third_party/xla/xla/hlo/experimental/auto_sharding/auto_sharding_memory.cc b/third_party/xla/xla/hlo/experimental/auto_sharding/auto_sharding_memory.cc
index 5c12fdfbf3c791..411e54d4a9d44e 100644
--- a/third_party/xla/xla/hlo/experimental/auto_sharding/auto_sharding_memory.cc
+++ b/third_party/xla/xla/hlo/experimental/auto_sharding/auto_sharding_memory.cc
@@ -25,6 +25,7 @@ limitations under the License.
 
 #include "absl/container/btree_map.h"
 #include "absl/container/btree_set.h"
+#include "absl/container/flat_hash_set.h"
 #include "absl/log/log.h"
 #include "tsl/platform/protobuf.h"
 
@@ -280,5 +281,48 @@ MemoryTermReducer::GetReducedGroups() const {
   return reduced_groups_;
 }
 
+// Retrieves a reduced subset of time points (i.e., along the liveness
+// dimension) that are sufficient to establish memory constraints.
+absl::flat_hash_set<int64_t> MemoryTermReducer::GetReducedTimes(
+    int64_t num_primitives) {
+  return GetReducedTimes(num_primitives, reduced_intervals_, reduced_groups_);
+}
+
+absl::flat_hash_set<int64_t> MemoryTermReducer::GetReducedTimes(
+    int64_t num_primitives,
+    const std::vector<std::pair<int64_t, int64_t>>& reduced_intervals,
+    const std::vector<absl::btree_set<int64_t>>& reduced_groups) {
+  // First, reconstruct the original intervals of the primitives alone (since
+  // the start of a group does not necessarily correspond to the start of a
+  // primitive).
+  std::vector<std::pair<int64_t, int64_t>> intervals;
+  for (int64_t reduced_interval_idx = 0;
+       reduced_interval_idx < reduced_intervals.size();
+       ++reduced_interval_idx) {
+    const Interval& reduced_interval = reduced_intervals[reduced_interval_idx];
+    if (reduced_interval_idx < num_primitives) {
+      intervals.push_back(reduced_interval);
+      continue;
+    }
+    const GroupIdx group_idx = reduced_interval_idx - num_primitives;
+    for (const PrimIdx prim_idx : reduced_groups[group_idx]) {
+      Interval& interval = intervals[prim_idx];
+      if (!IsValid(interval)) {
+        interval.first = reduced_interval.first;
+        interval.second = reduced_interval.second;
+        continue;
+      }
+      interval.first = std::min(interval.first, reduced_interval.first);
+      interval.second = std::max(interval.second, reduced_interval.second);
+    }
+  }
+  // Then, collect the entering times across all primitive intervals.
+  absl::flat_hash_set<int64_t> reduced_times;
+  for (const Interval& interval : intervals) {
+    if (IsValid(interval)) reduced_times.insert(interval.first);
+  }
+  return reduced_times;
+}
+
 }  // namespace spmd
 }  // namespace xla
diff --git a/third_party/xla/xla/hlo/experimental/auto_sharding/auto_sharding_memory.h b/third_party/xla/xla/hlo/experimental/auto_sharding/auto_sharding_memory.h
index 563f18bb190fc6..a9e0938cf7a0d8 100644
--- a/third_party/xla/xla/hlo/experimental/auto_sharding/auto_sharding_memory.h
+++ b/third_party/xla/xla/hlo/experimental/auto_sharding/auto_sharding_memory.h
@@ -23,6 +23,7 @@ limitations under the License.
 #include <vector>
 
 #include "absl/container/btree_set.h"
+#include "absl/container/flat_hash_set.h"
 #include "tsl/platform/protobuf.h"
 
 namespace xla {
@@ -64,6 +65,17 @@ class MemoryTermReducer {
   const std::vector<std::pair<int64_t, int64_t>>& GetReducedIntervals() const;
   const std::vector<absl::btree_set<int64_t>>& GetReducedGroups() const;
 
+  // Retrieves a reduced subset of time points along the liveness profile that
+  // are sufficient to establish memory constraints.
+  absl::flat_hash_set<int64_t> GetReducedTimes(int64_t num_primitives);
+
+  // A static version of the above method (in case we're using a precomputed
+  // memory term reduction).
+  static absl::flat_hash_set<int64_t> GetReducedTimes(
+      int64_t num_primitives,
+      const std::vector<std::pair<int64_t, int64_t>>& reduced_intervals,
+      const std::vector<absl::btree_set<int64_t>>& reduced_groups);
+
  private:
   // The internal implementation, agnostic to whether the client uses a liveness
   // matrix or primitive intervals.
diff --git a/third_party/xla/xla/hlo/experimental/auto_sharding/auto_sharding_memory_test.cc b/third_party/xla/xla/hlo/experimental/auto_sharding/auto_sharding_memory_test.cc
index cfa7b4d2af23ee..a1c9bc27694226 100644
--- a/third_party/xla/xla/hlo/experimental/auto_sharding/auto_sharding_memory_test.cc
+++ b/third_party/xla/xla/hlo/experimental/auto_sharding/auto_sharding_memory_test.cc
@@ -22,6 +22,7 @@ limitations under the License.
 
 #include <gtest/gtest.h>
 #include "absl/container/btree_set.h"
+#include "absl/container/flat_hash_set.h"
 
 namespace xla {
 namespace spmd {
@@ -49,6 +50,7 @@ std::function<std::pair<int64_t, int64_t>(int64_t)> Convert(
 //  +------->  ==>  +------->
 //    (time)          (time)
 TEST(AutoShardingMemoryTest, WithoutOverlap) {
+  const int num_primitives = 2;
   const std::vector<std::vector<int64_t>> live =
       {{0   },
        {0   },
@@ -59,7 +61,7 @@ TEST(AutoShardingMemoryTest, WithoutOverlap) {
 
   MemoryTermReducer reducer;
   const auto num_terms =
-          reducer.Reduce(live.size(), /*num_primitives=*/2, Convert(live));
+          reducer.Reduce(live.size(), num_primitives, Convert(live));
 
   const std::vector<std::vector<int64_t>> expected_reduced_live =
       {{0   },
@@ -72,10 +74,12 @@ TEST(AutoShardingMemoryTest, WithoutOverlap) {
       {{0, 2}, {3, 5}};
   const std::vector<absl::btree_set<int64_t>> expected_reduced_groups = {};
   const std::pair<int64_t, int64_t> expected_num_terms = {6, 6};
+  const absl::flat_hash_set<int64_t> expected_reduced_times = {0, 3};
   EXPECT_EQ(num_terms, expected_num_terms);
   EXPECT_EQ(reducer.GetReducedLive(), expected_reduced_live);
   EXPECT_EQ(reducer.GetReducedIntervals(), expected_reduced_intervals);
   EXPECT_EQ(reducer.GetReducedGroups(), expected_reduced_groups);
+  EXPECT_EQ(reducer.GetReducedTimes(num_primitives), expected_reduced_times);
 }
 
 //  |                  2222
@@ -84,6 +88,7 @@ TEST(AutoShardingMemoryTest, WithoutOverlap) {
 //  +------->  ==>  +------->
 //    (time)          (time)
 TEST(AutoShardingMemoryTest, PartialOverlap) {
+  const int num_primitives = 2;
   const std::vector<std::vector<int64_t>> live =
       {{0   },
        {0, 1},
@@ -94,7 +99,7 @@ TEST(AutoShardingMemoryTest, PartialOverlap) {
 
   MemoryTermReducer reducer;
   const auto num_terms =
-          reducer.Reduce(live.size(), /*num_primitives=*/2, Convert(live));
+          reducer.Reduce(live.size(), num_primitives, Convert(live));
 
   const std::vector<std::vector<int64_t>> expected_reduced_live =
       {{0      },
@@ -108,10 +113,12 @@ TEST(AutoShardingMemoryTest, PartialOverlap) {
   const std::vector<absl::btree_set<int64_t>> expected_reduced_groups =
       {{0, 1}};
   const std::pair<int64_t, int64_t> expected_num_terms = {10, 8};
+  const absl::flat_hash_set<int64_t> expected_reduced_times = {0, 1};
   EXPECT_EQ(num_terms, expected_num_terms);
   EXPECT_EQ(reducer.GetReducedLive(), expected_reduced_live);
   EXPECT_EQ(reducer.GetReducedIntervals(), expected_reduced_intervals);
   EXPECT_EQ(reducer.GetReducedGroups(), expected_reduced_groups);
+  EXPECT_EQ(reducer.GetReducedTimes(num_primitives), expected_reduced_times);
 }
 
 //  |                  2222
@@ -120,6 +127,7 @@ TEST(AutoShardingMemoryTest, PartialOverlap) {
 //  +------->  ==>  +------->
 //    (time)          (time)
 TEST(AutoShardingMemoryTest, PartialOverlapReversed) {
+  const int num_primitives = 2;
   const std::vector<std::vector<int64_t>> live =
       {{   1},
        {0, 1},
@@ -130,7 +138,7 @@ TEST(AutoShardingMemoryTest, PartialOverlapReversed) {
 
   MemoryTermReducer reducer;
   const auto num_terms =
-          reducer.Reduce(live.size(), /*num_primitives=*/2, Convert(live));
+          reducer.Reduce(live.size(), num_primitives, Convert(live));
 
   const std::vector<std::vector<int64_t>> expected_reduced_live =
       {{   1   },
@@ -144,10 +152,12 @@ TEST(AutoShardingMemoryTest, PartialOverlapReversed) {
   const std::vector<absl::btree_set<int64_t>> expected_reduced_groups =
       {{0, 1}};
   const std::pair<int64_t, int64_t> expected_num_terms = {10, 8};
+  const absl::flat_hash_set<int64_t> expected_reduced_times = {0, 1};
   EXPECT_EQ(num_terms, expected_num_terms);
   EXPECT_EQ(reducer.GetReducedLive(), expected_reduced_live);
   EXPECT_EQ(reducer.GetReducedIntervals(), expected_reduced_intervals);
   EXPECT_EQ(reducer.GetReducedGroups(), expected_reduced_groups);
+  EXPECT_EQ(reducer.GetReducedTimes(num_primitives), expected_reduced_times);
 }
 
 //  |  1111    ==>  |  1111
@@ -155,6 +165,7 @@ TEST(AutoShardingMemoryTest, PartialOverlapReversed) {
 //  +------->  ==>  +------->
 //    (time)          (time)
 TEST(AutoShardingMemoryTest, DoesNotSplitPrimitive) {
+  const int num_primitives = 2;
   const std::vector<std::vector<int64_t>> live =
       {{0   },
        {0, 1},
@@ -165,7 +176,7 @@ TEST(AutoShardingMemoryTest, DoesNotSplitPrimitive) {
 
   MemoryTermReducer reducer;
   const auto num_terms =
-          reducer.Reduce(live.size(), /*num_primitives=*/2, Convert(live));
+          reducer.Reduce(live.size(), num_primitives, Convert(live));
 
   const std::vector<std::vector<int64_t>> expected_reduced_live =
       {{0   },
@@ -178,10 +189,12 @@ TEST(AutoShardingMemoryTest, DoesNotSplitPrimitive) {
       {{0, 5}, {1, 4}};
   const std::vector<absl::btree_set<int64_t>> expected_reduced_groups = {};
   const std::pair<int64_t, int64_t> expected_num_terms = {10, 10};
+  const absl::flat_hash_set<int64_t> expected_reduced_times = {0, 1};
   EXPECT_EQ(num_terms, expected_num_terms);
   EXPECT_EQ(reducer.GetReducedLive(), expected_reduced_live);
   EXPECT_EQ(reducer.GetReducedIntervals(), expected_reduced_intervals);
   EXPECT_EQ(reducer.GetReducedGroups(), expected_reduced_groups);
+  EXPECT_EQ(reducer.GetReducedTimes(num_primitives), expected_reduced_times);
 }
 
 //  |          ==>  |  22222
@@ -190,6 +203,7 @@ TEST(AutoShardingMemoryTest, DoesNotSplitPrimitive) {
 //  +------->  ==>  +------->
 //    (time)          (time)
 TEST(AutoShardingMemoryTest, OnePrimitiveVanishes) {
+  const int num_primitives = 2;
   const std::vector<std::vector<int64_t>> live =
       {{0   },
        {0, 1},
@@ -200,7 +214,7 @@ TEST(AutoShardingMemoryTest, OnePrimitiveVanishes) {
 
   MemoryTermReducer reducer;
   const auto num_terms =
-          reducer.Reduce(live.size(), /*num_primitives=*/2, Convert(live));
+          reducer.Reduce(live.size(), num_primitives, Convert(live));
 
   const std::vector<std::vector<int64_t>> expected_reduced_live =
       {{0   },
@@ -214,10 +228,12 @@ TEST(AutoShardingMemoryTest, OnePrimitiveVanishes) {
   const std::vector<absl::btree_set<int64_t>> expected_reduced_groups =
       {{0, 1}};
   const std::pair<int64_t, int64_t> expected_num_terms = {11, 8};
+  const absl::flat_hash_set<int64_t> expected_reduced_times = {0, 1};
   EXPECT_EQ(num_terms, expected_num_terms);
   EXPECT_EQ(reducer.GetReducedLive(), expected_reduced_live);
   EXPECT_EQ(reducer.GetReducedIntervals(), expected_reduced_intervals);
   EXPECT_EQ(reducer.GetReducedGroups(), expected_reduced_groups);
+  EXPECT_EQ(reducer.GetReducedTimes(num_primitives), expected_reduced_times);
 }
 
 //  |          ==>  | 222222
@@ -226,6 +242,7 @@ TEST(AutoShardingMemoryTest, OnePrimitiveVanishes) {
 //  +------->  ==>  +------->
 //    (time)          (time)
 TEST(AutoShardingMemoryTest, BothPrimitivesVanish) {
+  const int num_primitives = 2;
   const std::vector<std::vector<int64_t>> live =
       {{0, 1},
        {0, 1},
@@ -236,7 +253,7 @@ TEST(AutoShardingMemoryTest, BothPrimitivesVanish) {
 
   MemoryTermReducer reducer;
   const auto num_terms =
-          reducer.Reduce(live.size(), /*num_primitives=*/2, Convert(live));
+          reducer.Reduce(live.size(), num_primitives, Convert(live));
 
   const std::vector<std::vector<int64_t>> expected_reduced_live =
       {{2},
@@ -250,10 +267,12 @@ TEST(AutoShardingMemoryTest, BothPrimitivesVanish) {
   const std::vector<absl::btree_set<int64_t>> expected_reduced_groups =
       {{0, 1}};
   const std::pair<int64_t, int64_t> expected_num_terms = {12, 8};
+  const absl::flat_hash_set<int64_t> expected_reduced_times = {0};
   EXPECT_EQ(num_terms, expected_num_terms);
   EXPECT_EQ(reducer.GetReducedLive(), expected_reduced_live);
   EXPECT_EQ(reducer.GetReducedIntervals(), expected_reduced_intervals);
   EXPECT_EQ(reducer.GetReducedGroups(), expected_reduced_groups);
+  EXPECT_EQ(reducer.GetReducedTimes(num_primitives), expected_reduced_times);
 }
 
 //  |            ==>  | 33333
@@ -263,6 +282,7 @@ TEST(AutoShardingMemoryTest, BothPrimitivesVanish) {
 //  +--------->  ==>  +--------->
 //    (time)            (time)
 TEST(AutoShardingMemoryTest, OneGroupingPreventsAnother) {
+  const int num_primitives = 3;
   const std::vector<std::vector<int64_t>> live =
       {{0, 1   },
        {0, 1   },
@@ -276,7 +296,7 @@ TEST(AutoShardingMemoryTest, OneGroupingPreventsAnother) {
 
   MemoryTermReducer reducer;
   const auto num_terms =
-          reducer.Reduce(live.size(), /*num_primitives=*/3, Convert(live));
+          reducer.Reduce(live.size(), num_primitives, Convert(live));
 
   const std::vector<std::vector<int64_t>> expected_reduced_live =
       {{      3},
@@ -293,10 +313,12 @@ TEST(AutoShardingMemoryTest, OneGroupingPreventsAnother) {
   const std::vector<absl::btree_set<int64_t>> expected_reduced_groups =
       {{0, 1}};
   const std::pair<int64_t, int64_t> expected_num_terms = {18, 15};
+  const absl::flat_hash_set<int64_t> expected_reduced_times = {0, 4};
   EXPECT_EQ(num_terms, expected_num_terms);
   EXPECT_EQ(reducer.GetReducedLive(), expected_reduced_live);
   EXPECT_EQ(reducer.GetReducedIntervals(), expected_reduced_intervals);
   EXPECT_EQ(reducer.GetReducedGroups(), expected_reduced_groups);
+  EXPECT_EQ(reducer.GetReducedTimes(num_primitives), expected_reduced_times);
 }
 
 //  |          ==>  | 333444
@@ -306,6 +328,7 @@ TEST(AutoShardingMemoryTest, OneGroupingPreventsAnother) {
 //  +------->  ==>  +------->
 //    (time)          (time)
 TEST(AutoShardingMemoryTest, TwoGroups) {
+  const int num_primitives = 3;
   const std::vector<std::vector<int64_t>> live =
       {{0, 1   },
        {0, 1   },
@@ -316,7 +339,7 @@ TEST(AutoShardingMemoryTest, TwoGroups) {
 
   MemoryTermReducer reducer;
   const auto num_terms =
-          reducer.Reduce(live.size(), /*num_primitives=*/3, Convert(live));
+          reducer.Reduce(live.size(), num_primitives, Convert(live));
 
   const std::vector<std::vector<int64_t>> expected_reduced_live =
       {{3},
@@ -330,10 +353,12 @@ TEST(AutoShardingMemoryTest, TwoGroups) {
   const std::vector<absl::btree_set<int64_t>> expected_reduced_groups =
       {{0, 1}, {0, 2}};
   const std::pair<int64_t, int64_t> expected_num_terms = {12, 10};
+  const absl::flat_hash_set<int64_t> expected_reduced_times = {0, 3};
   EXPECT_EQ(num_terms, expected_num_terms);
   EXPECT_EQ(reducer.GetReducedLive(), expected_reduced_live);
   EXPECT_EQ(reducer.GetReducedIntervals(), expected_reduced_intervals);
   EXPECT_EQ(reducer.GetReducedGroups(), expected_reduced_groups);
+  EXPECT_EQ(reducer.GetReducedTimes(num_primitives), expected_reduced_times);
 }
 
 //  |           ==>  |  444555
@@ -344,6 +369,7 @@ TEST(AutoShardingMemoryTest, TwoGroups) {
 //  +------->   ==>  +------->
 //    (time)           (time)
 TEST(AutoShardingMemoryTest, TwoGroupsMutuallyExclusive) {
+  const int num_primitives = 4;
   const std::vector<std::vector<int64_t>> live =
       {{0         },
        {0, 1      },
@@ -356,7 +382,7 @@ TEST(AutoShardingMemoryTest, TwoGroupsMutuallyExclusive) {
 
   MemoryTermReducer reducer;
   const auto num_terms =
-          reducer.Reduce(live.size(), /*num_primitives=*/4, Convert(live));
+          reducer.Reduce(live.size(), num_primitives, Convert(live));
 
   const std::vector<std::vector<int64_t>> expected_reduced_live =
       {{0      },
@@ -372,10 +398,12 @@ TEST(AutoShardingMemoryTest, TwoGroupsMutuallyExclusive) {
   const std::vector<absl::btree_set<int64_t>> expected_reduced_groups =
       {{0, 1}, {2, 3}};
   const std::pair<int64_t, int64_t> expected_num_terms = {14, 12};
+  const absl::flat_hash_set<int64_t> expected_reduced_times = {0, 1, 4};
   EXPECT_EQ(num_terms, expected_num_terms);
   EXPECT_EQ(reducer.GetReducedLive(), expected_reduced_live);
   EXPECT_EQ(reducer.GetReducedIntervals(), expected_reduced_intervals);
   EXPECT_EQ(reducer.GetReducedGroups(), expected_reduced_groups);
+  EXPECT_EQ(reducer.GetReducedTimes(num_primitives), expected_reduced_times);
 }
 
 //  | 11      ==>  | 11
@@ -383,13 +411,14 @@ TEST(AutoShardingMemoryTest, TwoGroupsMutuallyExclusive) {
 //  +------>  ==>  +------>
 //   (time)         (time)
 TEST(AutoShardingMemoryTest, MergingPrimitivesWouldNotReduceTerms) {
+  const int num_primitives = 2;
   const std::vector<std::vector<int64_t>> live =
       {{0, 1},
        {0, 1}};
 
   MemoryTermReducer reducer;
   const auto num_terms =
-          reducer.Reduce(live.size(), /*num_primitives=*/2, Convert(live));
+          reducer.Reduce(live.size(), num_primitives, Convert(live));
 
   const std::vector<std::vector<int64_t>> expected_reduced_live =
       {{0, 1},
@@ -398,10 +427,12 @@ TEST(AutoShardingMemoryTest, MergingPrimitivesWouldNotReduceTerms) {
       {{0, 1}, {0, 1}};
   const std::vector<absl::btree_set<int64_t>> expected_reduced_groups = {};
   const std::pair<int64_t, int64_t> expected_num_terms = {4, 4};
+  const absl::flat_hash_set<int64_t> expected_reduced_times = {0};
   EXPECT_EQ(num_terms, expected_num_terms);
   EXPECT_EQ(reducer.GetReducedLive(), expected_reduced_live);
   EXPECT_EQ(reducer.GetReducedIntervals(), expected_reduced_intervals);
   EXPECT_EQ(reducer.GetReducedGroups(), expected_reduced_groups);
+  EXPECT_EQ(reducer.GetReducedTimes(num_primitives), expected_reduced_times);
 }
 
 //  |          ==>  | 333333
@@ -411,6 +442,7 @@ TEST(AutoShardingMemoryTest, MergingPrimitivesWouldNotReduceTerms) {
 //  +------->  ==>  +------->
 //    (time)          (time)
 TEST(AutoShardingMemoryTest, AllPrimitivesVanish) {
+  const int num_primitives = 3;
   const std::vector<std::vector<int64_t>> live =
       {{0, 1, 2},
        {0, 1, 2},
@@ -421,7 +453,7 @@ TEST(AutoShardingMemoryTest, AllPrimitivesVanish) {
 
   MemoryTermReducer reducer;
   const auto num_terms =
-          reducer.Reduce(live.size(), /*num_primitives=*/3, Convert(live));
+          reducer.Reduce(live.size(), num_primitives, Convert(live));
 
   const std::vector<std::vector<int64_t>> expected_reduced_live =
       {{3},
@@ -435,10 +467,12 @@ TEST(AutoShardingMemoryTest, AllPrimitivesVanish) {
   const std::vector<absl::btree_set<int64_t>> expected_reduced_groups =
       {{0, 1, 2}};
   const std::pair<int64_t, int64_t> expected_num_terms = {18, 9};
+  const absl::flat_hash_set<int64_t> expected_reduced_times = {0};
   EXPECT_EQ(num_terms, expected_num_terms);
   EXPECT_EQ(reducer.GetReducedLive(), expected_reduced_live);
   EXPECT_EQ(reducer.GetReducedIntervals(), expected_reduced_intervals);
   EXPECT_EQ(reducer.GetReducedGroups(), expected_reduced_groups);
+  EXPECT_EQ(reducer.GetReducedTimes(num_primitives), expected_reduced_times);
 }
 
 //  |            ==>  |    555555
@@ -450,6 +484,7 @@ TEST(AutoShardingMemoryTest, AllPrimitivesVanish) {
 //  +--------->  ==>  +--------->
 //     (time)            (time)
 TEST(AutoShardingMemoryTest, MergingGroupsWouldNotReduceTerms) {
+  const int num_primitives = 4;
   const std::vector<std::vector<int64_t>> live =
       {{0, 1      },
        {0, 1      },
@@ -463,7 +498,7 @@ TEST(AutoShardingMemoryTest, MergingGroupsWouldNotReduceTerms) {
 
   MemoryTermReducer reducer;
   const auto num_terms =
-          reducer.Reduce(live.size(), /*num_primitives=*/4, Convert(live));
+          reducer.Reduce(live.size(), num_primitives, Convert(live));
 
   const std::vector<std::vector<int64_t>> expected_reduced_live =
       {{4   },
@@ -480,10 +515,12 @@ TEST(AutoShardingMemoryTest, MergingGroupsWouldNotReduceTerms) {
   const std::vector<absl::btree_set<int64_t>> expected_reduced_groups =
       {{0, 1}, {2, 3}};
   const std::pair<int64_t, int64_t> expected_num_terms = {26, 17};
+  const absl::flat_hash_set<int64_t> expected_reduced_times = {0, 3};
   EXPECT_EQ(num_terms, expected_num_terms);
   EXPECT_EQ(reducer.GetReducedLive(), expected_reduced_live);
   EXPECT_EQ(reducer.GetReducedIntervals(), expected_reduced_intervals);
   EXPECT_EQ(reducer.GetReducedGroups(), expected_reduced_groups);
+  EXPECT_EQ(reducer.GetReducedTimes(num_primitives), expected_reduced_times);
 }
 
 //  |                      |  444466666555
@@ -494,6 +531,7 @@ TEST(AutoShardingMemoryTest, MergingGroupsWouldNotReduceTerms) {
 //  +-------------->  ==>  +-------------->
 //       (time)                 (time)
 TEST(AutoShardingMemoryTest, ExampleFromDocumentation) {
+  const int num_primitives = 4;
   const std::vector<std::vector<int64_t>> live =
       {{0         },
        {0, 1      },
@@ -512,7 +550,7 @@ TEST(AutoShardingMemoryTest, ExampleFromDocumentation) {
 
   MemoryTermReducer reducer;
   const auto num_terms =
-          reducer.Reduce(live.size(), /*num_primitives=*/4, Convert(live));
+          reducer.Reduce(live.size(), num_primitives, Convert(live));
 
   const std::vector<std::vector<int64_t>> expected_reduced_live =
       {{0      },
@@ -534,10 +572,12 @@ TEST(AutoShardingMemoryTest, ExampleFromDocumentation) {
   const std::vector<absl::btree_set<int64_t>> expected_reduced_groups =
       {{0, 1}, {2, 3}, {0, 1, 2, 3}};
   const std::pair<int64_t, int64_t> expected_num_terms = {36, 22};
+  const absl::flat_hash_set<int64_t> expected_reduced_times = {0, 1, 5};
   EXPECT_EQ(num_terms, expected_num_terms);
   EXPECT_EQ(reducer.GetReducedLive(), expected_reduced_live);
   EXPECT_EQ(reducer.GetReducedIntervals(), expected_reduced_intervals);
   EXPECT_EQ(reducer.GetReducedGroups(), expected_reduced_groups);
+  EXPECT_EQ(reducer.GetReducedTimes(num_primitives), expected_reduced_times);
 }
 
 //  |         ==>  | 333
@@ -547,6 +587,7 @@ TEST(AutoShardingMemoryTest, ExampleFromDocumentation) {
 //  +------>  ==>  +------>
 //   (time)         (time)
 TEST(AutoShardingMemoryTest, MergesWithRightmost) {
+  const int num_primitives = 3;
   const std::vector<std::vector<int64_t>> live =
       {{0,    2},
        {0,    2},
@@ -555,7 +596,7 @@ TEST(AutoShardingMemoryTest, MergesWithRightmost) {
 
   MemoryTermReducer reducer;
   const auto num_terms =
-          reducer.Reduce(live.size(), /*num_primitives=*/3, Convert(live));
+          reducer.Reduce(live.size(), num_primitives, Convert(live));
 
   const std::vector<std::vector<int64_t>> expected_reduced_live =
       {{      3},
@@ -567,10 +608,12 @@ TEST(AutoShardingMemoryTest, MergesWithRightmost) {
   const std::vector<absl::btree_set<int64_t>> expected_reduced_groups =
       {{0, 2}};
   const std::pair<int64_t, int64_t> expected_num_terms = {8, 7};
+  const absl::flat_hash_set<int64_t> expected_reduced_times = {0, 3};
   EXPECT_EQ(num_terms, expected_num_terms);
   EXPECT_EQ(reducer.GetReducedLive(), expected_reduced_live);
   EXPECT_EQ(reducer.GetReducedIntervals(), expected_reduced_intervals);
   EXPECT_EQ(reducer.GetReducedGroups(), expected_reduced_groups);
+  EXPECT_EQ(reducer.GetReducedTimes(num_primitives), expected_reduced_times);
 }
 
 //  |                      |  444466666555
@@ -581,11 +624,12 @@ TEST(AutoShardingMemoryTest, MergesWithRightmost) {
 //  +-------------->  ==>  +-------------->
 //       (time)                 (time)
 TEST(AutoShardingMemoryTest, ExampleFromDocumentationUsingIntervals) {
+  const int num_primitives = 4;
   const std::vector<std::pair<int64_t, int64_t>> intervals =
       {{0, 9}, {1, 9}, {5, 12}, {5, 13}};
 
   MemoryTermReducer reducer;
-  const auto num_terms = reducer.Reduce(/*num_lives=*/14, /*num_primitives=*/4,
+  const auto num_terms = reducer.Reduce(/*num_lives=*/14, num_primitives,
                                         Convert(intervals));
 
   const std::vector<std::vector<int64_t>> expected_reduced_live = {};
@@ -594,18 +638,21 @@ TEST(AutoShardingMemoryTest, ExampleFromDocumentationUsingIntervals) {
   const std::vector<absl::btree_set<int64_t>> expected_reduced_groups =
       {{0, 1}, {2, 3}, {0, 1, 2, 3}};
   const std::pair<int64_t, int64_t> expected_num_terms = {36, 22};
+  const absl::flat_hash_set<int64_t> expected_reduced_times = {0, 1, 5};
   EXPECT_EQ(num_terms, expected_num_terms);
   EXPECT_EQ(reducer.GetReducedLive(), expected_reduced_live);
   EXPECT_EQ(reducer.GetReducedIntervals(), expected_reduced_intervals);
   EXPECT_EQ(reducer.GetReducedGroups(), expected_reduced_groups);
+  EXPECT_EQ(reducer.GetReducedTimes(num_primitives), expected_reduced_times);
 }
 
 TEST(AutoShardingMemoryTest, InvalidIntervals) {
+  const int num_primitives = 3;
   const std::vector<std::pair<int64_t, int64_t>> intervals =
       {{0, 4}, {9223372036854775807, 0}, {9223372036854775807, 0}};
 
   MemoryTermReducer reducer;
-  const auto num_terms = reducer.Reduce(/*num_lives=*/5, /*num_primitives=*/3,
+  const auto num_terms = reducer.Reduce(/*num_lives=*/5, num_primitives,
                                         Convert(intervals));
 
   const std::vector<std::vector<int64_t>> expected_reduced_live = {};
@@ -613,10 +660,12 @@ TEST(AutoShardingMemoryTest, InvalidIntervals) {
       {{0, 4}, {9223372036854775807, 0}, {9223372036854775807, 0}};
   const std::vector<absl::btree_set<int64_t>> expected_reduced_groups = {};
   const std::pair<int64_t, int64_t> expected_num_terms = {5, 5};
+  const absl::flat_hash_set<int64_t> expected_reduced_times = {0};
   EXPECT_EQ(num_terms, expected_num_terms);
   EXPECT_EQ(reducer.GetReducedLive(), expected_reduced_live);
   EXPECT_EQ(reducer.GetReducedIntervals(), expected_reduced_intervals);
   EXPECT_EQ(reducer.GetReducedGroups(), expected_reduced_groups);
+  EXPECT_EQ(reducer.GetReducedTimes(num_primitives), expected_reduced_times);
 }
 
 //  |                 ==>  |      55555555
@@ -628,11 +677,12 @@ TEST(AutoShardingMemoryTest, InvalidIntervals) {
 //  +-------------->  ==>  +-------------->
 //       (time)                 (time)
 TEST(AutoShardingMemoryTest, OneIterationOnly) {
+  const int num_primitives = 4;
   const std::vector<std::pair<int64_t, int64_t>> intervals =
       {{0, 9}, {1, 9}, {5, 12}, {5, 13}};
 
   MemoryTermReducer reducer;
-  const auto num_terms = reducer.Reduce(/*num_lives=*/14, /*num_primitives=*/4,
+  const auto num_terms = reducer.Reduce(/*num_lives=*/14, num_primitives,
                                         Convert(intervals),
                                         /*max_iterations=*/1);
 
@@ -642,10 +692,12 @@ TEST(AutoShardingMemoryTest, OneIterationOnly) {
   const std::vector<absl::btree_set<int64_t>> expected_reduced_groups =
       {{0, 1}, {2, 3}};
   const std::pair<int64_t, int64_t> expected_num_terms = {36, 23};
+  const absl::flat_hash_set<int64_t> expected_reduced_times = {0, 1, 5};
   EXPECT_EQ(num_terms, expected_num_terms);
   EXPECT_EQ(reducer.GetReducedLive(), expected_reduced_live);
   EXPECT_EQ(reducer.GetReducedIntervals(), expected_reduced_intervals);
   EXPECT_EQ(reducer.GetReducedGroups(), expected_reduced_groups);
+  EXPECT_EQ(reducer.GetReducedTimes(num_primitives), expected_reduced_times);
 }
 
 //  |                 ==>  | 55555
@@ -657,11 +709,12 @@ TEST(AutoShardingMemoryTest, OneIterationOnly) {
 //  +-------------->  ==>  +-------------->
 //       (time)                 (time)
 TEST(AutoShardingMemoryTest, StairsBottomLeft) {
+  const int num_primitives = 4;
   const std::vector<std::pair<int64_t, int64_t>> intervals =
       {{0, 13}, {0, 10}, {0, 7}, {0, 4}};
 
   MemoryTermReducer reducer;
-  const auto num_terms = reducer.Reduce(/*num_lives=*/14, /*num_primitives=*/4,
+  const auto num_terms = reducer.Reduce(/*num_lives=*/14, num_primitives,
                                         Convert(intervals),
                                         /*max_iterations=*/1);
 
@@ -671,10 +724,12 @@ TEST(AutoShardingMemoryTest, StairsBottomLeft) {
   const std::vector<absl::btree_set<int64_t>> expected_reduced_groups =
       {{0, 1}, {2, 3}};
   const std::pair<int64_t, int64_t> expected_num_terms = {38, 26};
+  const absl::flat_hash_set<int64_t> expected_reduced_times = {0};
   EXPECT_EQ(num_terms, expected_num_terms);
   EXPECT_EQ(reducer.GetReducedLive(), expected_reduced_live);
   EXPECT_EQ(reducer.GetReducedIntervals(), expected_reduced_intervals);
   EXPECT_EQ(reducer.GetReducedGroups(), expected_reduced_groups);
+  EXPECT_EQ(reducer.GetReducedTimes(num_primitives), expected_reduced_times);
 }
 
 //  |                 ==>  | 55555
@@ -686,11 +741,12 @@ TEST(AutoShardingMemoryTest, StairsBottomLeft) {
 //  +-------------->  ==>  +-------------->
 //       (time)                 (time)
 TEST(AutoShardingMemoryTest, StairsTopLeft) {
+  const int num_primitives = 4;
   const std::vector<std::pair<int64_t, int64_t>> intervals =
       {{0, 4}, {0, 7}, {0, 10}, {0, 13}};
 
   MemoryTermReducer reducer;
-  const auto num_terms = reducer.Reduce(/*num_lives=*/14, /*num_primitives=*/4,
+  const auto num_terms = reducer.Reduce(/*num_lives=*/14, num_primitives,
                                         Convert(intervals),
                                         /*max_iterations=*/1);
 
@@ -700,10 +756,12 @@ TEST(AutoShardingMemoryTest, StairsTopLeft) {
   const std::vector<absl::btree_set<int64_t>> expected_reduced_groups =
       {{2, 3}, {0, 1}};
   const std::pair<int64_t, int64_t> expected_num_terms = {38, 26};
+  const absl::flat_hash_set<int64_t> expected_reduced_times = {0};
   EXPECT_EQ(num_terms, expected_num_terms);
   EXPECT_EQ(reducer.GetReducedLive(), expected_reduced_live);
   EXPECT_EQ(reducer.GetReducedIntervals(), expected_reduced_intervals);
   EXPECT_EQ(reducer.GetReducedGroups(), expected_reduced_groups);
+  EXPECT_EQ(reducer.GetReducedTimes(num_primitives), expected_reduced_times);
 }
 
 //  |                 ==>  |          55555
@@ -715,11 +773,12 @@ TEST(AutoShardingMemoryTest, StairsTopLeft) {
 //  +-------------->  ==>  +-------------->
 //       (time)                 (time)
 TEST(AutoShardingMemoryTest, StairsTopRight) {
+  const int num_primitives = 4;
   const std::vector<std::pair<int64_t, int64_t>> intervals =
       {{9, 13}, {6, 13}, {3, 13}, {0, 13}};
 
   MemoryTermReducer reducer;
-  const auto num_terms = reducer.Reduce(/*num_lives=*/14, /*num_primitives=*/4,
+  const auto num_terms = reducer.Reduce(/*num_lives=*/14, num_primitives,
                                         Convert(intervals),
                                         /*max_iterations=*/1);
 
@@ -729,10 +788,12 @@ TEST(AutoShardingMemoryTest, StairsTopRight) {
   const std::vector<absl::btree_set<int64_t>> expected_reduced_groups =
       {{2, 3}, {0, 1}};
   const std::pair<int64_t, int64_t> expected_num_terms = {38, 26};
+  const absl::flat_hash_set<int64_t> expected_reduced_times = {0, 3, 6, 9};
   EXPECT_EQ(num_terms, expected_num_terms);
   EXPECT_EQ(reducer.GetReducedLive(), expected_reduced_live);
   EXPECT_EQ(reducer.GetReducedIntervals(), expected_reduced_intervals);
   EXPECT_EQ(reducer.GetReducedGroups(), expected_reduced_groups);
+  EXPECT_EQ(reducer.GetReducedTimes(num_primitives), expected_reduced_times);
 }
 
 //  |                 ==>  |          55555
@@ -744,11 +805,12 @@ TEST(AutoShardingMemoryTest, StairsTopRight) {
 //  +-------------->  ==>  +-------------->
 //       (time)                 (time)
 TEST(AutoShardingMemoryTest, StairsBottomRight) {
+  const int num_primitives = 4;
   const std::vector<std::pair<int64_t, int64_t>> intervals =
       {{0, 13}, {3, 13}, {6, 13}, {9, 13}};
 
   MemoryTermReducer reducer;
-  const auto num_terms = reducer.Reduce(/*num_lives=*/14, /*num_primitives=*/4,
+  const auto num_terms = reducer.Reduce(/*num_lives=*/14, num_primitives,
                                         Convert(intervals),
                                         /*max_iterations=*/1);
 
@@ -758,10 +820,12 @@ TEST(AutoShardingMemoryTest, StairsBottomRight) {
   const std::vector<absl::btree_set<int64_t>> expected_reduced_groups =
       {{0, 1}, {2, 3}};
   const std::pair<int64_t, int64_t> expected_num_terms = {38, 26};
+  const absl::flat_hash_set<int64_t> expected_reduced_times = {0, 3, 6, 9};
   EXPECT_EQ(num_terms, expected_num_terms);
   EXPECT_EQ(reducer.GetReducedLive(), expected_reduced_live);
   EXPECT_EQ(reducer.GetReducedIntervals(), expected_reduced_intervals);
   EXPECT_EQ(reducer.GetReducedGroups(), expected_reduced_groups);
+  EXPECT_EQ(reducer.GetReducedTimes(num_primitives), expected_reduced_times);
 }
 
 // clang-format on
diff --git a/third_party/xla/xla/hlo/experimental/auto_sharding/auto_sharding_solver.cc b/third_party/xla/xla/hlo/experimental/auto_sharding/auto_sharding_solver.cc
index 181687da2f545d..16816e000e3ffd 100644
--- a/third_party/xla/xla/hlo/experimental/auto_sharding/auto_sharding_solver.cc
+++ b/third_party/xla/xla/hlo/experimental/auto_sharding/auto_sharding_solver.cc
@@ -231,7 +231,8 @@ std::optional<std::pair<int64_t, int64_t>> ReduceMemoryTerms(
     std::string_view prim_type,
     std::vector<std::vector<MPVariable*>>& prim_vars,
     std::vector<std::pair<int64_t, int64_t>>& reduced_intervals,
-    std::vector<MPVariable*>& group_vars) {
+    std::vector<MPVariable*>& group_vars,
+    absl::flat_hash_set<int64_t>& reduced_times) {
   std::optional<std::pair<int64_t, int64_t>> num_terms = std::nullopt;
   std::vector<absl::btree_set<int64_t>> reduced_groups;
   if (groups.empty()) {
@@ -278,6 +279,9 @@ std::optional<std::pair<int64_t, int64_t>> ReduceMemoryTerms(
       }
     }
   }
+  const absl::flat_hash_set<int64_t> times = MemoryTermReducer::GetReducedTimes(
+      num_primitives, reduced_intervals, reduced_groups);
+  reduced_times.insert(times.begin(), times.end());
   return num_terms;
 }
 
@@ -289,12 +293,14 @@ void AddMemoryTerms(
     const tsl::protobuf::RepeatedPtrField<  // NOLINT
         AutoShardingSolverRequest_Costs>& memory_costs,
     const MPVariable* overbudget_var,
+    const absl::flat_hash_set<int64_t>& reduced_times,
     std::vector<std::vector<MPVariable*>>& prim_vars,
     std::vector<MPVariable*>& group_vars,
     absl::flat_hash_map<LivenessIdx, MPConstraint*>& constraints) {
   for (int64_t prim_idx = 0; prim_idx < intervals.size(); ++prim_idx) {
     for (int64_t time_idx = intervals[prim_idx].first;
          time_idx <= intervals[prim_idx].second; ++time_idx) {
+      if (!reduced_times.contains(time_idx)) continue;
       if (!constraints.contains(time_idx)) {
         MPConstraint* constraint = solver.MakeRowConstraint(
             -MPSolver::infinity(), 100.0, absl::StrCat("mem[", time_idx, "]"));
@@ -581,18 +587,19 @@ AutoShardingSolverResult CallORToolsSolver(
     };
     std::vector<std::pair<int64_t, int64_t>> reduced_intervals_nodes,
         reduced_intervals_edges;
+    absl::flat_hash_set<int64_t> reduced_times;
     std::vector<MPVariable*> group_node_vars, group_edge_vars;
     const absl::Time term_reduction_start_time = absl::Now();
     auto num_node_terms = ReduceMemoryTerms(
         request, *solver, request.live_size(), request.num_nodes(),
         std::move(LiveNodes), request.node_intervals(), request.node_groups(),
         request.memory_costs(), "node", s, reduced_intervals_nodes,
-        group_node_vars);
+        group_node_vars, reduced_times);
     auto num_edge_terms = ReduceMemoryTerms(
         request, *solver, request.live_edges_size(), request.edges_size(),
         std::move(LiveEdges), request.edge_intervals(), request.edge_groups(),
         request.memory_edge_costs(), "edge", e, reduced_intervals_edges,
-        group_edge_vars);
+        group_edge_vars, reduced_times);
     const absl::Time term_reduction_end_time = absl::Now();
     if (num_node_terms && num_edge_terms) {
       const auto term_reduction_duration =
@@ -606,11 +613,13 @@ AutoShardingSolverResult CallORToolsSolver(
     absl::flat_hash_map<LivenessIdx, MPConstraint*> constraints;
     AddMemoryTerms(request, *solver, request.num_nodes(),
                    reduced_intervals_nodes, request.memory_costs(),
-                   overbudget_var, s, group_node_vars, constraints);
+                   overbudget_var, reduced_times, s, group_node_vars,
+                   constraints);
     if (request.enable_memory_edge_costs()) {
       AddMemoryTerms(request, *solver, request.edges_size(),
                      reduced_intervals_edges, request.memory_edge_costs(),
-                     overbudget_var, e, group_edge_vars, constraints);
+                     overbudget_var, reduced_times, e, group_edge_vars,
+                     constraints);
     }
     if (overbudget_var) {
       solver->MutableObjective()->SetCoefficient(

From 94d270e9138cd0de2f87ebe7c0379bbc6a33f01f Mon Sep 17 00:00:00 2001
From: Peter Hawkins <phawkins@google.com>
Date: Wed, 15 May 2024 12:02:34 -0700
Subject: [PATCH 198/478] [PJRT:CPU] Bump the stack size for the CPU client
 thread pool to 8MB.

2MB was found to be causing crashes in JAX linear algebra tests that call into lapack/blas.
OpenBLAS appears to have fairly large stack size requirements.

PiperOrigin-RevId: 634031821
---
 third_party/xla/xla/pjrt/cpu/cpu_client.cc | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/third_party/xla/xla/pjrt/cpu/cpu_client.cc b/third_party/xla/xla/pjrt/cpu/cpu_client.cc
index 1b7b5fe12b7c88..192f92a752edb3 100644
--- a/third_party/xla/xla/pjrt/cpu/cpu_client.cc
+++ b/third_party/xla/xla/pjrt/cpu/cpu_client.cc
@@ -410,7 +410,9 @@ static tsl::ThreadOptions GetThreadOptions() {
   tsl::ThreadOptions thread_options;
   // On Mac OS the default stack size is 512KiB, which is too small for some
   // BLAS and LAPACK functions (https://github.com/google/jax/issues/20428).
-  thread_options.stack_size = 2 * 1024 * 1024;
+  // On Linux we also observed that 2MB wasn't enough to run some OpenBLAS
+  // functions.
+  thread_options.stack_size = 8 * 1024 * 1024;
   return thread_options;
 }
 

From ca2f9ae1f16185a6be1444b971df4467526b7c76 Mon Sep 17 00:00:00 2001
From: Kyle Lucke <klucke@google.com>
Date: Wed, 15 May 2024 12:05:58 -0700
Subject: [PATCH 199/478] Use absl::Status instead of xla::Status now that
 they're identical.

PiperOrigin-RevId: 634033147
---
 third_party/xla/xla/python/pjrt_ifrt/pjrt_array.cc      | 7 ++++---
 third_party/xla/xla/python/pjrt_ifrt/pjrt_executable.cc | 4 ++--
 third_party/xla/xla/python/pjrt_ifrt/pjrt_tuple.cc      | 2 +-
 third_party/xla/xla/python/pjrt_ifrt/pjrt_tuple.h       | 2 +-
 4 files changed, 8 insertions(+), 7 deletions(-)

diff --git a/third_party/xla/xla/python/pjrt_ifrt/pjrt_array.cc b/third_party/xla/xla/python/pjrt_ifrt/pjrt_array.cc
index 12cdf80e7705b7..2d595ab45e3d22 100644
--- a/third_party/xla/xla/python/pjrt_ifrt/pjrt_array.cc
+++ b/third_party/xla/xla/python/pjrt_ifrt/pjrt_array.cc
@@ -60,8 +60,9 @@ namespace {
 
 // Validates the sharding and PjRtBuffers have consistent device and memory
 // kind.
-Status ValidateArrayCreationInput(std::shared_ptr<const Sharding> sharding,
-                                  const PjRtArray::PjRtBuffers& pjrt_buffers) {
+absl::Status ValidateArrayCreationInput(
+    std::shared_ptr<const Sharding> sharding,
+    const PjRtArray::PjRtBuffers& pjrt_buffers) {
   if (pjrt_buffers.empty()) {
     return InvalidArgument("pjrt_buffers must be non-empty");
   }
@@ -416,7 +417,7 @@ Future<> PjRtArray::CopyToHostBuffer(
   // TODO(hyeontaek): Handle semantics == kDonateInput.
   pjrt_buffer->ToLiteral(literal_ptr)
       .OnReady([literal = std::move(literal),
-                promise = std::move(promise)](Status s) mutable {
+                promise = std::move(promise)](absl::Status s) mutable {
         promise.Set(std::move(s));
         literal = nullptr;
       });
diff --git a/third_party/xla/xla/python/pjrt_ifrt/pjrt_executable.cc b/third_party/xla/xla/python/pjrt_ifrt/pjrt_executable.cc
index bf2930516a61bb..02b4cea85f504f 100644
--- a/third_party/xla/xla/python/pjrt_ifrt/pjrt_executable.cc
+++ b/third_party/xla/xla/python/pjrt_ifrt/pjrt_executable.cc
@@ -339,7 +339,7 @@ PjRtLoadedExecutable::CreateInternal(
   auto append_arg = [&](const xla::PrimitiveType& element_type,
                         const xla::DimensionVector& dimensions,
                         const xla::HloSharding* sharding,
-                        MemoryKind memory_kind) -> Status {
+                        MemoryKind memory_kind) -> absl::Status {
     TF_ASSIGN_OR_RETURN(auto dtype, ToDType(element_type));
     output_dtypes.push_back(dtype);
     output_shapes.push_back(Shape(dimensions));
@@ -611,7 +611,7 @@ PjRtLoadedExecutable::Execute(absl::Span<tsl::RCReference<Array>> args,
     result.status.OnReady(
         [all_loaded_host_callbacks = all_loaded_host_callbacks_,
          host_callback_states = std::move(host_callback_states)](
-            Status) mutable { all_loaded_host_callbacks.reset(); });
+            absl::Status) mutable { all_loaded_host_callbacks.reset(); });
   }
 
   // Convert 2-level PjRtBuffer vectors into an Array vector.
diff --git a/third_party/xla/xla/python/pjrt_ifrt/pjrt_tuple.cc b/third_party/xla/xla/python/pjrt_ifrt/pjrt_tuple.cc
index 9909dcd665fcad..ab5d935147e2ed 100644
--- a/third_party/xla/xla/python/pjrt_ifrt/pjrt_tuple.cc
+++ b/third_party/xla/xla/python/pjrt_ifrt/pjrt_tuple.cc
@@ -81,7 +81,7 @@ std::string PjRtTuple::DebugString() const {
 }
 int PjRtTuple::Arity() { return values_.size(); }
 
-Status PjRtTuple::Unpack(absl::Span<tsl::RCReference<Value>> values_out) {
+absl::Status PjRtTuple::Unpack(absl::Span<tsl::RCReference<Value>> values_out) {
   if (values_out.size() != values_.size()) {
     return InvalidArgument(
         "Wrong number of output values for "
diff --git a/third_party/xla/xla/python/pjrt_ifrt/pjrt_tuple.h b/third_party/xla/xla/python/pjrt_ifrt/pjrt_tuple.h
index 6ad3deac70415e..2a6359be9507b2 100644
--- a/third_party/xla/xla/python/pjrt_ifrt/pjrt_tuple.h
+++ b/third_party/xla/xla/python/pjrt_ifrt/pjrt_tuple.h
@@ -56,7 +56,7 @@ class PjRtTuple final : public llvm::RTTIExtends<PjRtTuple, Tuple> {
 
   int Arity() override;
 
-  Status Unpack(absl::Span<tsl::RCReference<Value>> values) override;
+  absl::Status Unpack(absl::Span<tsl::RCReference<Value>> values) override;
 
   static char ID;  // NOLINT
 

From 70b6a4811eb0f964f55ed9ccec9fda7dc3c3b2a2 Mon Sep 17 00:00:00 2001
From: CJ Carey <cjcarey@google.com>
Date: Wed, 15 May 2024 12:08:32 -0700
Subject: [PATCH 200/478] Fix a bug in `validate_colocate` when a variable is
 not created in a strategy scope.

In some cases `v._distribute_strategy` can be None, which was causing an AttributeError when trying to check if the `extended` field matched the expected value.

PiperOrigin-RevId: 634034018
---
 tensorflow/python/distribute/distribute_utils.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorflow/python/distribute/distribute_utils.py b/tensorflow/python/distribute/distribute_utils.py
index 73ef2fcf0f5454..fd577236210a90 100644
--- a/tensorflow/python/distribute/distribute_utils.py
+++ b/tensorflow/python/distribute/distribute_utils.py
@@ -260,7 +260,7 @@ def is_distributed_table(v):
 
 def _validate_colocate_extended(v, extended):
   variable_strategy = v._distribute_strategy  # pylint: disable=protected-access
-  if variable_strategy.extended is not extended:
+  if not variable_strategy or variable_strategy.extended is not extended:
     raise ValueError(
         "`colocate_vars_with` must only be passed a variable created in this "
         "tf.distribute.Strategy.scope(), not %s created in scope: %s" %

From 59ad508b15bf881e99006380421fdb71490c2756 Mon Sep 17 00:00:00 2001
From: Kevin Gleason <gleasonk@google.com>
Date: Wed, 15 May 2024 12:25:04 -0700
Subject: [PATCH 201/478] Bump ODML to target StableHLO v1.0.0

PiperOrigin-RevId: 634038627
---
 .../lite/stablehlo/transforms/legalize_stablehlo_to_vhlo.cc     | 2 +-
 tensorflow/lite/core/macros.h                                   | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/tensorflow/compiler/mlir/lite/stablehlo/transforms/legalize_stablehlo_to_vhlo.cc b/tensorflow/compiler/mlir/lite/stablehlo/transforms/legalize_stablehlo_to_vhlo.cc
index a03d1a1d958f65..b3a85259cc482a 100644
--- a/tensorflow/compiler/mlir/lite/stablehlo/transforms/legalize_stablehlo_to_vhlo.cc
+++ b/tensorflow/compiler/mlir/lite/stablehlo/transforms/legalize_stablehlo_to_vhlo.cc
@@ -280,7 +280,7 @@ struct LegalizeStablehloToVhloPass
     VhloToStablehloTypeConverter to_builtin_converter;
 
     // StableHLO --> VHLO (allow funcs)
-    //   VHLO -> Downgrade to 0.19.0 / tflite_supported_stablehlo_version
+    //   VHLO -> Downgrade to tflite_supported_stablehlo_version
     //     VHLO Tensor --> Builtin Tensor
     //       Remove cast(tensor->vhlo) -> cast(vhlo->tensor) pattern
     if (failed(ApplyStablehloToVhloPatterns(module,
diff --git a/tensorflow/lite/core/macros.h b/tensorflow/lite/core/macros.h
index f1d10aec93229e..9eab6be877d2e5 100644
--- a/tensorflow/lite/core/macros.h
+++ b/tensorflow/lite/core/macros.h
@@ -74,7 +74,7 @@ inline constexpr char tflite_metadata_buffer_location[] = "buffer_location";
 inline constexpr char tflite_metadata_min_runtime_version[] =
     "min_runtime_version";
 // the stablehlo op version is supported by the tflite runtime
-inline constexpr char tflite_supported_stablehlo_version[] = "0.19.0";
+inline constexpr char tflite_supported_stablehlo_version[] = "1.0.0";
 #endif
 
 #endif  // TENSORFLOW_LITE_CORE_MACROS_H_

From a1b4c95e237b2a40ad3e653b05604fd782e88454 Mon Sep 17 00:00:00 2001
From: George Karpenkov <cheshire@google.com>
Date: Wed, 15 May 2024 12:27:54 -0700
Subject: [PATCH 202/478] [XLA:GPU] Error out early on missing autotuning cache
 when required

This is technically a performance optimization: we would still error out later
on otherwise, but after recompiling & running a bunch of modules, hence no test.

PiperOrigin-RevId: 634039337
---
 .../xla/service/gpu/gemm_fusion_autotuner.cc  | 27 +++++++++++++------
 1 file changed, 19 insertions(+), 8 deletions(-)

diff --git a/third_party/xla/xla/service/gpu/gemm_fusion_autotuner.cc b/third_party/xla/xla/service/gpu/gemm_fusion_autotuner.cc
index a0d47feab41153..dada3783df6b2e 100644
--- a/third_party/xla/xla/service/gpu/gemm_fusion_autotuner.cc
+++ b/third_party/xla/xla/service/gpu/gemm_fusion_autotuner.cc
@@ -203,16 +203,22 @@ class GemmFusionAutotunerVisitor : public DfsHloRewriteVisitor {
   AutotuneConfig config_;
 };
 
+using TilingConfigsMap =
+    absl::flat_hash_map<const HloFusionInstruction*, std::vector<Config>>;
+
 class GemmConfigSetCollector : public ConstDfsHloVisitorWithDefault {
  public:
   explicit GemmConfigSetCollector(GemmFusionAutotunerImpl* impl)
       : impl_(impl) {}
 
-  absl::StatusOr<
-      absl::flat_hash_map<const HloFusionInstruction*, std::vector<Config>>>
-  CollectGemmConfigSets(
+  // Find configurations to tune.
+  absl::StatusOr<TilingConfigsMap> CollectGemmConfigSets(
       const HloModule* module,
       const absl::flat_hash_set<absl::string_view>& execution_threads = {}) {
+    error_out_on_cache_miss_ =
+        module->config()
+            .debug_options()
+            .xla_gpu_require_complete_aot_autotune_results();
     gemm_config_sets_.clear();
     for (HloComputation* computation :
          module->MakeNonfusionComputations(execution_threads)) {
@@ -249,6 +255,13 @@ class GemmConfigSetCollector : public ConstDfsHloVisitorWithDefault {
                           (backend_config.kind() == kCuDnnFusionKind &&
                            !backend_config.has_cudnn_fusion_config());
     if (missing_config) {
+      if (error_out_on_cache_miss_) {
+        return absl::NotFoundError(absl::StrCat(
+            "Complete autotuning results are required, but no cache result "
+            "found for key: ",
+            key.ToString()));
+      }
+
       TF_ASSIGN_OR_RETURN(std::vector<Config> configs,
                           impl_->GenerateConfigs(*fusion));
       TF_RET_CHECK(
@@ -264,9 +277,9 @@ class GemmConfigSetCollector : public ConstDfsHloVisitorWithDefault {
   }
 
  private:
+  bool error_out_on_cache_miss_;
   GemmFusionAutotunerImpl* impl_;
-  absl::flat_hash_map<const HloFusionInstruction*, std::vector<Config>>
-      gemm_config_sets_;
+  TilingConfigsMap gemm_config_sets_;
   AutoTuneCacheKeyCount fusion_count_map_;
   absl::flat_hash_set<AutotuneCacheKey> handled_fusions_;
 };
@@ -1128,9 +1141,7 @@ absl::StatusOr<bool> GemmFusionAutotuner::Run(
   GemmFusionAutotunerImpl autotuner(config_, toolkit_version_, debug_options,
                                     thread_pool_);
   GemmConfigSetCollector gemm_config_set_collector(&autotuner);
-  absl::flat_hash_map<const HloFusionInstruction*, std::vector<Config>>
-      gemm_config_sets;
-  TF_ASSIGN_OR_RETURN(gemm_config_sets,
+  TF_ASSIGN_OR_RETURN(TilingConfigsMap gemm_config_sets,
                       gemm_config_set_collector.CollectGemmConfigSets(
                           module, execution_threads));
 

From 16cfd19eba2d2bfd50487dc9bb7b82734a1413dd Mon Sep 17 00:00:00 2001
From: Trevor Morris <tmorris@nvidia.com>
Date: Wed, 15 May 2024 12:33:15 -0700
Subject: [PATCH 203/478] PR #12439: [GPU] Multinode support for
 multihost_hlo_runner

Imported from GitHub PR https://github.com/openxla/xla/pull/12439

This PR adds multinode support to the multihost_hlo_runner. To use, there is a new command line argument `address` for the address of the coordinator/root node.

Example usage with SLURM:
```
bazel run //xla/tools/multihost_hlo_runner:hlo_runner_main -- \
  --task_id=${SLURM_PROCID} \
  --num_nodes=${SLURM_NTASKS} \
  --address="${SLURM_LAUNCH_NODE_IPADDR}:12345" \
  ...
  ```
Copybara import of the project:

--
ad1705053185c32abbf4d917f2b1fd669d5b6893 by Trevor Morris <tmorris@nvidia.com>:

Multinode changes for hlo_runner_main

Merging this change closes #12439

PiperOrigin-RevId: 634040710
---
 .../xla/xla/tools/multihost_hlo_runner/BUILD  |  1 +
 .../multihost_hlo_runner/hlo_runner_main.cc   | 37 +++++++++++++++++--
 2 files changed, 35 insertions(+), 3 deletions(-)

diff --git a/third_party/xla/xla/tools/multihost_hlo_runner/BUILD b/third_party/xla/xla/tools/multihost_hlo_runner/BUILD
index 8014fbe8ed3f0d..3ec5ee2b60ffbf 100644
--- a/third_party/xla/xla/tools/multihost_hlo_runner/BUILD
+++ b/third_party/xla/xla/tools/multihost_hlo_runner/BUILD
@@ -39,6 +39,7 @@ xla_cc_binary(
         "//xla:status",
         "//xla:statusor",
         "//xla/pjrt:pjrt_client",
+        "//xla/pjrt/distributed",
         "//xla/service:cpu_plugin",
         "//xla/tsl/util:command_line_flags",
         "@com_google_absl//absl/log:check",
diff --git a/third_party/xla/xla/tools/multihost_hlo_runner/hlo_runner_main.cc b/third_party/xla/xla/tools/multihost_hlo_runner/hlo_runner_main.cc
index b0666ff491fcba..54394d6dcfa1a2 100644
--- a/third_party/xla/xla/tools/multihost_hlo_runner/hlo_runner_main.cc
+++ b/third_party/xla/xla/tools/multihost_hlo_runner/hlo_runner_main.cc
@@ -23,6 +23,7 @@ limitations under the License.
 #include "absl/log/check.h"
 #include "absl/strings/str_cat.h"
 #include "xla/debug_options_flags.h"
+#include "xla/pjrt/distributed/distributed.h"
 #include "xla/pjrt/pjrt_client.h"
 #include "xla/status.h"
 #include "xla/statusor.h"
@@ -81,6 +82,7 @@ int main(int argc, char** argv) {
   int task_id = 0;
   int num_nodes = 1;
   std::string device_type_str = "gpu";
+  std::string address_str = "";
   xla::FunctionalHloRunner::PreprocessingOptions preproc_options;
   xla::FunctionalHloRunner::RawCompileOptions raw_compile_options;
   xla::FunctionalHloRunner::RunningOptions running_options;
@@ -97,10 +99,15 @@ int main(int argc, char** argv) {
                 "Example: /a/b/literal.txt."),
       tsl::Flag("task_id", &task_id, "Borg task id."),
       tsl::Flag("device_type", &device_type_str, "Device type: gpu, host"),
-      tsl::Flag("num_nodes", &num_nodes, "Number of nodes (hosts)"),
+      tsl::Flag("num_nodes", &num_nodes,
+                "Number of nodes (hosts). If greater than 1, a distributed "
+                "service will be created for task_id 0"),
       tsl::Flag(
           "enable_mock_nccl", &enable_mock_nccl,
           "Should we simulate multi-hosts run with mock nccl collectives?"),
+      tsl::Flag("address", &address_str,
+                "Coordinator address with port for when num_nodes > 1. "
+                "Example: 127.0.0.1:12345"),
   };
 
   xla::MultiHostHloRunnerFlags hlo_runner_flags;
@@ -131,6 +138,8 @@ int main(int argc, char** argv) {
     }
   }
 
+  std::unique_ptr<xla::DistributedRuntimeService> service;
+
   // The main logic:
   absl::StatusOr<std::unique_ptr<xla::PjRtClient>> client = [&] {
     if (device_type_str == "host") {
@@ -144,8 +153,30 @@ int main(int argc, char** argv) {
       CHECK_GT(num_nodes, 1);
       return xla::FunctionalHloRunner::CreateMockGpuClient(num_nodes);
     } else {
-      CHECK_EQ(num_nodes, 1);
-      return xla::FunctionalHloRunner::CreateGpuClient();
+      if (num_nodes == 1) {
+        return xla::FunctionalHloRunner::CreateGpuClient();
+      } else {
+        CHECK_GT(address_str.length(), 0);
+        // Multinode. Start service on task 0.
+        if (task_id == 0) {
+          std::string coordinator_bind_address =
+              "[::]:" + address_str.substr(address_str.rfind(":") + 1);
+          xla::CoordinationServiceImpl::Options options;
+          options.num_nodes = num_nodes;
+          auto status_or = xla::GetDistributedRuntimeService(
+              coordinator_bind_address, options);
+          TF_QCHECK_OK(status_or.status());
+          service = std::move(status_or.value());
+        }
+        xla::DistributedRuntimeClient::Options options;
+        options.node_id = task_id;
+        options.init_timeout = absl::Seconds(300);
+        auto distributed_client =
+            xla::GetDistributedRuntimeClient(address_str, options);
+        TF_QCHECK_OK(distributed_client->Connect());
+        return xla::FunctionalHloRunner::CreateGpuClient(distributed_client,
+                                                         task_id, num_nodes);
+      }
     }
   }();
 

From f49186fa978f469207137c9f658b010304a07e53 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Wed, 15 May 2024 12:57:19 -0700
Subject: [PATCH 204/478] When computing the set of reduced times, keep just
 the latest one if multiple primitives enter the liveness profile
 consecutively.

PiperOrigin-RevId: 634047168
---
 .../auto_sharding/auto_sharding_memory.cc     | 19 ++++++++++++---
 .../auto_sharding_memory_test.cc              | 24 +++++++++----------
 2 files changed, 28 insertions(+), 15 deletions(-)

diff --git a/third_party/xla/xla/hlo/experimental/auto_sharding/auto_sharding_memory.cc b/third_party/xla/xla/hlo/experimental/auto_sharding/auto_sharding_memory.cc
index 411e54d4a9d44e..4d5fc89ad7ab21 100644
--- a/third_party/xla/xla/hlo/experimental/auto_sharding/auto_sharding_memory.cc
+++ b/third_party/xla/xla/hlo/experimental/auto_sharding/auto_sharding_memory.cc
@@ -316,11 +316,24 @@ absl::flat_hash_set<int64_t> MemoryTermReducer::GetReducedTimes(
       interval.second = std::max(interval.second, reduced_interval.second);
     }
   }
-  // Then, collect the entering times across all primitive intervals.
-  absl::flat_hash_set<int64_t> reduced_times;
+  // Next, created a sorted set of all entering & eviction times.
+  absl::btree_set<std::pair<int64_t, /*is_eviction=*/bool>> times;
   for (const Interval& interval : intervals) {
-    if (IsValid(interval)) reduced_times.insert(interval.first);
+    if (!IsValid(interval)) continue;
+    times.insert({interval.first, false});
+    times.insert({interval.second, true});
+  }
+  // Finally, collect the latest entering times across all primitive intervals.
+  int64_t last_entering_time = -1;
+  absl::flat_hash_set<int64_t> reduced_times;
+  for (const auto& time : times) {
+    if (/*is_eviction*/ time.second) {
+      reduced_times.insert(last_entering_time);
+    } else {
+      last_entering_time = time.first;
+    }
   }
+  reduced_times.insert(last_entering_time);
   return reduced_times;
 }
 
diff --git a/third_party/xla/xla/hlo/experimental/auto_sharding/auto_sharding_memory_test.cc b/third_party/xla/xla/hlo/experimental/auto_sharding/auto_sharding_memory_test.cc
index a1c9bc27694226..a99945ec9bbc69 100644
--- a/third_party/xla/xla/hlo/experimental/auto_sharding/auto_sharding_memory_test.cc
+++ b/third_party/xla/xla/hlo/experimental/auto_sharding/auto_sharding_memory_test.cc
@@ -113,7 +113,7 @@ TEST(AutoShardingMemoryTest, PartialOverlap) {
   const std::vector<absl::btree_set<int64_t>> expected_reduced_groups =
       {{0, 1}};
   const std::pair<int64_t, int64_t> expected_num_terms = {10, 8};
-  const absl::flat_hash_set<int64_t> expected_reduced_times = {0, 1};
+  const absl::flat_hash_set<int64_t> expected_reduced_times = {1};
   EXPECT_EQ(num_terms, expected_num_terms);
   EXPECT_EQ(reducer.GetReducedLive(), expected_reduced_live);
   EXPECT_EQ(reducer.GetReducedIntervals(), expected_reduced_intervals);
@@ -152,7 +152,7 @@ TEST(AutoShardingMemoryTest, PartialOverlapReversed) {
   const std::vector<absl::btree_set<int64_t>> expected_reduced_groups =
       {{0, 1}};
   const std::pair<int64_t, int64_t> expected_num_terms = {10, 8};
-  const absl::flat_hash_set<int64_t> expected_reduced_times = {0, 1};
+  const absl::flat_hash_set<int64_t> expected_reduced_times = {1};
   EXPECT_EQ(num_terms, expected_num_terms);
   EXPECT_EQ(reducer.GetReducedLive(), expected_reduced_live);
   EXPECT_EQ(reducer.GetReducedIntervals(), expected_reduced_intervals);
@@ -189,7 +189,7 @@ TEST(AutoShardingMemoryTest, DoesNotSplitPrimitive) {
       {{0, 5}, {1, 4}};
   const std::vector<absl::btree_set<int64_t>> expected_reduced_groups = {};
   const std::pair<int64_t, int64_t> expected_num_terms = {10, 10};
-  const absl::flat_hash_set<int64_t> expected_reduced_times = {0, 1};
+  const absl::flat_hash_set<int64_t> expected_reduced_times = {1};
   EXPECT_EQ(num_terms, expected_num_terms);
   EXPECT_EQ(reducer.GetReducedLive(), expected_reduced_live);
   EXPECT_EQ(reducer.GetReducedIntervals(), expected_reduced_intervals);
@@ -228,7 +228,7 @@ TEST(AutoShardingMemoryTest, OnePrimitiveVanishes) {
   const std::vector<absl::btree_set<int64_t>> expected_reduced_groups =
       {{0, 1}};
   const std::pair<int64_t, int64_t> expected_num_terms = {11, 8};
-  const absl::flat_hash_set<int64_t> expected_reduced_times = {0, 1};
+  const absl::flat_hash_set<int64_t> expected_reduced_times = {1};
   EXPECT_EQ(num_terms, expected_num_terms);
   EXPECT_EQ(reducer.GetReducedLive(), expected_reduced_live);
   EXPECT_EQ(reducer.GetReducedIntervals(), expected_reduced_intervals);
@@ -313,7 +313,7 @@ TEST(AutoShardingMemoryTest, OneGroupingPreventsAnother) {
   const std::vector<absl::btree_set<int64_t>> expected_reduced_groups =
       {{0, 1}};
   const std::pair<int64_t, int64_t> expected_num_terms = {18, 15};
-  const absl::flat_hash_set<int64_t> expected_reduced_times = {0, 4};
+  const absl::flat_hash_set<int64_t> expected_reduced_times = {4};
   EXPECT_EQ(num_terms, expected_num_terms);
   EXPECT_EQ(reducer.GetReducedLive(), expected_reduced_live);
   EXPECT_EQ(reducer.GetReducedIntervals(), expected_reduced_intervals);
@@ -398,7 +398,7 @@ TEST(AutoShardingMemoryTest, TwoGroupsMutuallyExclusive) {
   const std::vector<absl::btree_set<int64_t>> expected_reduced_groups =
       {{0, 1}, {2, 3}};
   const std::pair<int64_t, int64_t> expected_num_terms = {14, 12};
-  const absl::flat_hash_set<int64_t> expected_reduced_times = {0, 1, 4};
+  const absl::flat_hash_set<int64_t> expected_reduced_times = {1, 4};
   EXPECT_EQ(num_terms, expected_num_terms);
   EXPECT_EQ(reducer.GetReducedLive(), expected_reduced_live);
   EXPECT_EQ(reducer.GetReducedIntervals(), expected_reduced_intervals);
@@ -515,7 +515,7 @@ TEST(AutoShardingMemoryTest, MergingGroupsWouldNotReduceTerms) {
   const std::vector<absl::btree_set<int64_t>> expected_reduced_groups =
       {{0, 1}, {2, 3}};
   const std::pair<int64_t, int64_t> expected_num_terms = {26, 17};
-  const absl::flat_hash_set<int64_t> expected_reduced_times = {0, 3};
+  const absl::flat_hash_set<int64_t> expected_reduced_times = {3};
   EXPECT_EQ(num_terms, expected_num_terms);
   EXPECT_EQ(reducer.GetReducedLive(), expected_reduced_live);
   EXPECT_EQ(reducer.GetReducedIntervals(), expected_reduced_intervals);
@@ -572,7 +572,7 @@ TEST(AutoShardingMemoryTest, ExampleFromDocumentation) {
   const std::vector<absl::btree_set<int64_t>> expected_reduced_groups =
       {{0, 1}, {2, 3}, {0, 1, 2, 3}};
   const std::pair<int64_t, int64_t> expected_num_terms = {36, 22};
-  const absl::flat_hash_set<int64_t> expected_reduced_times = {0, 1, 5};
+  const absl::flat_hash_set<int64_t> expected_reduced_times = {5};
   EXPECT_EQ(num_terms, expected_num_terms);
   EXPECT_EQ(reducer.GetReducedLive(), expected_reduced_live);
   EXPECT_EQ(reducer.GetReducedIntervals(), expected_reduced_intervals);
@@ -638,7 +638,7 @@ TEST(AutoShardingMemoryTest, ExampleFromDocumentationUsingIntervals) {
   const std::vector<absl::btree_set<int64_t>> expected_reduced_groups =
       {{0, 1}, {2, 3}, {0, 1, 2, 3}};
   const std::pair<int64_t, int64_t> expected_num_terms = {36, 22};
-  const absl::flat_hash_set<int64_t> expected_reduced_times = {0, 1, 5};
+  const absl::flat_hash_set<int64_t> expected_reduced_times = {5};
   EXPECT_EQ(num_terms, expected_num_terms);
   EXPECT_EQ(reducer.GetReducedLive(), expected_reduced_live);
   EXPECT_EQ(reducer.GetReducedIntervals(), expected_reduced_intervals);
@@ -692,7 +692,7 @@ TEST(AutoShardingMemoryTest, OneIterationOnly) {
   const std::vector<absl::btree_set<int64_t>> expected_reduced_groups =
       {{0, 1}, {2, 3}};
   const std::pair<int64_t, int64_t> expected_num_terms = {36, 23};
-  const absl::flat_hash_set<int64_t> expected_reduced_times = {0, 1, 5};
+  const absl::flat_hash_set<int64_t> expected_reduced_times = {5};
   EXPECT_EQ(num_terms, expected_num_terms);
   EXPECT_EQ(reducer.GetReducedLive(), expected_reduced_live);
   EXPECT_EQ(reducer.GetReducedIntervals(), expected_reduced_intervals);
@@ -788,7 +788,7 @@ TEST(AutoShardingMemoryTest, StairsTopRight) {
   const std::vector<absl::btree_set<int64_t>> expected_reduced_groups =
       {{2, 3}, {0, 1}};
   const std::pair<int64_t, int64_t> expected_num_terms = {38, 26};
-  const absl::flat_hash_set<int64_t> expected_reduced_times = {0, 3, 6, 9};
+  const absl::flat_hash_set<int64_t> expected_reduced_times = {9};
   EXPECT_EQ(num_terms, expected_num_terms);
   EXPECT_EQ(reducer.GetReducedLive(), expected_reduced_live);
   EXPECT_EQ(reducer.GetReducedIntervals(), expected_reduced_intervals);
@@ -820,7 +820,7 @@ TEST(AutoShardingMemoryTest, StairsBottomRight) {
   const std::vector<absl::btree_set<int64_t>> expected_reduced_groups =
       {{0, 1}, {2, 3}};
   const std::pair<int64_t, int64_t> expected_num_terms = {38, 26};
-  const absl::flat_hash_set<int64_t> expected_reduced_times = {0, 3, 6, 9};
+  const absl::flat_hash_set<int64_t> expected_reduced_times = {9};
   EXPECT_EQ(num_terms, expected_num_terms);
   EXPECT_EQ(reducer.GetReducedLive(), expected_reduced_live);
   EXPECT_EQ(reducer.GetReducedIntervals(), expected_reduced_intervals);

From bb53564b30bee0f206960829347f41fe61c41b2d Mon Sep 17 00:00:00 2001
From: Kyle Lucke <klucke@google.com>
Date: Wed, 15 May 2024 13:00:24 -0700
Subject: [PATCH 205/478] Make TpuExecutor use StreamExecutorInterface to
 create Events.

PiperOrigin-RevId: 634048195
---
 third_party/xla/xla/stream_executor/tpu/c_api_defn.h    | 4 ++--
 third_party/xla/xla/stream_executor/tpu/tpu_executor.cc | 8 +++++++-
 2 files changed, 9 insertions(+), 3 deletions(-)

diff --git a/third_party/xla/xla/stream_executor/tpu/c_api_defn.h b/third_party/xla/xla/stream_executor/tpu/c_api_defn.h
index 407d6c326707ac..59ecd662196daf 100644
--- a/third_party/xla/xla/stream_executor/tpu/c_api_defn.h
+++ b/third_party/xla/xla/stream_executor/tpu/c_api_defn.h
@@ -45,8 +45,8 @@ struct SE_Stream {
 };
 
 struct SE_Event {
-  explicit SE_Event(stream_executor::StreamExecutor* parent) : event(parent) {}
-  stream_executor::Event event;
+  explicit SE_Event(stream_executor::StreamExecutor* parent) {}
+  std::unique_ptr<stream_executor::Event> event;
 };
 
 struct SE_StreamExecutorConfig {
diff --git a/third_party/xla/xla/stream_executor/tpu/tpu_executor.cc b/third_party/xla/xla/stream_executor/tpu/tpu_executor.cc
index c59c0c45db1b82..ffab1983d879a0 100644
--- a/third_party/xla/xla/stream_executor/tpu/tpu_executor.cc
+++ b/third_party/xla/xla/stream_executor/tpu/tpu_executor.cc
@@ -95,7 +95,13 @@ bool TpuExecutor::CreateStreamDependency(Stream* dependent, Stream* other) {
       get_stream(other->implementation()));
 }
 
-Status TpuExecutor::AllocateEvent(Event* event) { return absl::OkStatus(); }
+Status TpuExecutor::AllocateEvent(Event* event) {
+  StatusHelper status;
+  auto se_event = tpu_platform().LookupEvent(event->implementation());
+  ExecutorApiFn()->TpuExecutor_AllocateEventFn(executor_, se_event,
+                                               status.c_status);
+  return status.status();
+}
 
 Status TpuExecutor::DeallocateEvent(Event* event) {
   tpu_platform().EraseEvent(event->implementation());

From cc54b2ccac9443736d7b6c97c4deacd13bb6c5ef Mon Sep 17 00:00:00 2001
From: Eugene Zhulenev <ezhulenev@google.com>
Date: Wed, 15 May 2024 13:15:58 -0700
Subject: [PATCH 206/478] [xla:ffi] Add support for FFI execution context to
 PjRtCpuClient

Add a test for passing arbitrary user data via execution context side channel to XLA FFI handlers registered for a HOST platform

PiperOrigin-RevId: 634053080
---
 third_party/xla/xla/pjrt/cpu/BUILD            |  5 ++
 third_party/xla/xla/pjrt/cpu/cpu_client.cc    |  3 +
 .../xla/xla/pjrt/cpu/cpu_client_test.cc       | 76 ++++++++++++++++++-
 third_party/xla/xla/service/cpu/BUILD         |  2 +
 third_party/xla/xla/service/cpu/ir_emitter.cc |  7 +-
 .../service/cpu/runtime_handle_ffi_call.cc    | 40 ++++++----
 .../xla/service/cpu/runtime_handle_ffi_call.h | 10 +--
 7 files changed, 116 insertions(+), 27 deletions(-)

diff --git a/third_party/xla/xla/pjrt/cpu/BUILD b/third_party/xla/xla/pjrt/cpu/BUILD
index a06396197fd617..7cc103260d4bd3 100644
--- a/third_party/xla/xla/pjrt/cpu/BUILD
+++ b/third_party/xla/xla/pjrt/cpu/BUILD
@@ -228,11 +228,16 @@ xla_cc_test(
         "//xla:literal_util",
         "//xla:shape_util",
         "//xla:status",
+        "//xla:status_macros",
         "//xla:util",
+        "//xla/client:xla_computation",
         "//xla/ffi",
         "//xla/ffi:ffi_api",
         "//xla/pjrt:host_memory_spaces",
+        "//xla/pjrt:pjrt_client",
+        "//xla/pjrt:pjrt_executable",
         "//xla/service:hlo_parser",
+        "//xla/tests:literal_test_util",
         "//xla/tests:test_utils",
         "@com_google_absl//absl/status",
         "@com_google_absl//absl/synchronization",
diff --git a/third_party/xla/xla/pjrt/cpu/cpu_client.cc b/third_party/xla/xla/pjrt/cpu/cpu_client.cc
index 192f92a752edb3..7416654899de7f 100644
--- a/third_party/xla/xla/pjrt/cpu/cpu_client.cc
+++ b/third_party/xla/xla/pjrt/cpu/cpu_client.cc
@@ -1531,6 +1531,9 @@ absl::StatusOr<PjRtLoadedExecutable::Result> TfrtCpuExecutable::ExecuteHelper(
       input_deps.push_back(std::move(last_enqueue_event));
     }
   }
+  if (options.context != nullptr) {
+    run_options.set_ffi_execution_context(&options.context->ffi_context());
+  }
 
   bool execute_inline = cheap_computation_ || !client_->asynchronous_;
 
diff --git a/third_party/xla/xla/pjrt/cpu/cpu_client_test.cc b/third_party/xla/xla/pjrt/cpu/cpu_client_test.cc
index 1b8075b11013d5..e70c5b4f370919 100644
--- a/third_party/xla/xla/pjrt/cpu/cpu_client_test.cc
+++ b/third_party/xla/xla/pjrt/cpu/cpu_client_test.cc
@@ -15,14 +15,15 @@ limitations under the License.
 
 #include "xla/pjrt/cpu/cpu_client.h"
 
-#include "xla/pjrt/host_memory_spaces.h"
-
 #ifndef _WIN32
 #include <unistd.h>
 #endif
 
 #include <algorithm>
+#include <cstdint>
 #include <cstring>
+#include <memory>
+#include <optional>
 #include <string>
 #include <vector>
 
@@ -30,14 +31,19 @@ limitations under the License.
 #include <gtest/gtest.h>
 #include "absl/status/status.h"
 #include "absl/synchronization/notification.h"
+#include "xla/client/xla_computation.h"
 #include "xla/ffi/ffi.h"
 #include "xla/ffi/ffi_api.h"
 #include "xla/literal.h"
 #include "xla/literal_util.h"
+#include "xla/pjrt/host_memory_spaces.h"
+#include "xla/pjrt/pjrt_client.h"
+#include "xla/pjrt/pjrt_executable.h"
 #include "xla/service/hlo_parser.h"
 #include "xla/shape.h"
 #include "xla/shape_util.h"
 #include "xla/status.h"
+#include "xla/tests/literal_test_util.h"
 #include "xla/tests/test_utils.h"
 #include "xla/util.h"
 #include "tsl/lib/core/status_test_util.h"
@@ -91,10 +97,16 @@ TEST(TfrtCpuClientTest, MemorySpace) {
 
 TEST(TfrtCpuClientTest, DonationWithExecutionError) {
   constexpr char kProgram[] =
-      R"(HloModule DonationWithExecutionError, input_output_alias={ {}: (0, {}, must-alias) }
+      R"(
+HloModule DonationWithExecutionError,
+          input_output_alias={ {}: (0, {}, must-alias) }
+
 ENTRY DonationWithExecutionError() -> f32[2, 2] {
     %input = f32[2, 2] parameter(0)
-    %custom-call = (f32[2, 2], u8[0]) custom-call(%input), custom_call_target="__xla_test$$TestError", api_version=API_VERSION_TYPED_FFI, output_to_operand_aliasing={{0}: (0, {})}
+    %custom-call = (f32[2, 2], u8[0]) custom-call(%input),
+                      custom_call_target="__xla_test$$TestError",
+                      api_version=API_VERSION_TYPED_FFI,
+                      output_to_operand_aliasing={{0}: (0, {})}
     ROOT %result = f32[2, 2] get-tuple-element(%custom-call), index=0
 })";
 
@@ -345,5 +357,61 @@ TEST(TfrtCpuClientTest, AsyncTransferRawDataToSubBuffer) {
   EXPECT_THAT(literal->data<uint32_t>(), Each(0x42424242));
 }
 
+// User-defined data type to be passed to FFI handler via the execute context
+// side channel.
+struct MemsetValue {
+  explicit MemsetValue(float value) : value(value) {}
+  float value;
+};
+
+static absl::Status MemsetFromValue(
+    ffi::Result<ffi::BufferR1<PrimitiveType::F32>> result,
+    MemsetValue* memset_value) {
+  for (size_t i = 0; i < result->dimensions.at(0); ++i) {
+    result->data.base()[i] = memset_value->value;
+  }
+  return absl::OkStatus();
+}
+
+XLA_FFI_DEFINE_HANDLER(kMemsetFromValue, MemsetFromValue,
+                       ffi::Ffi::Bind()
+                           .Ret<ffi::BufferR1<PrimitiveType::F32>>()
+                           .Ctx<ffi::UserData<MemsetValue>>());
+
+XLA_FFI_REGISTER_HANDLER(ffi::GetXlaFfiApi(), "MemsetFromValue", "HOST",
+                         kMemsetFromValue);
+
+TEST(TfrtCpuClientTest, ForwardUserDataToFfiHandler) {
+  static constexpr char const* kProgram = R"(
+    HloModule ffi_handler
+    ENTRY main {
+      ROOT %custom-call = f32[4] custom-call(),
+                          custom_call_target="MemsetFromValue",
+                          api_version=API_VERSION_TYPED_FFI
+    })";
+
+  TF_ASSERT_OK_AND_ASSIGN(auto client, GetTfrtCpuClient(CpuClientOptions()));
+
+  TF_ASSERT_OK_AND_ASSIGN(auto hlo_module,
+                          ParseAndReturnUnverifiedModule(kProgram, {}));
+  XlaComputation xla_computation(hlo_module->ToProto());
+  TF_ASSERT_OK_AND_ASSIGN(auto executable,
+                          client->Compile(xla_computation, {}));
+
+  ExecuteContext context;
+  TF_ASSERT_OK(context.ffi_context().Emplace<MemsetValue>(42.0f));
+
+  ExecuteOptions opts;
+  opts.context = &context;
+
+  auto result = executable->Execute(/*argument_handles=*/{{}}, opts);
+
+  TF_ASSERT_OK_AND_ASSIGN(std::shared_ptr<xla::Literal> result_literal,
+                          result->at(0).at(0)->ToLiteralSync());
+  EXPECT_TRUE(LiteralTestUtil::Equal(
+      LiteralUtil::CreateR1<float>({42.0f, 42.0f, 42.0f, 42.0f}),
+      *result_literal));
+}
+
 }  // namespace
 }  // namespace xla
diff --git a/third_party/xla/xla/service/cpu/BUILD b/third_party/xla/xla/service/cpu/BUILD
index 3f2bdc0f211410..24dc3d2d833696 100644
--- a/third_party/xla/xla/service/cpu/BUILD
+++ b/third_party/xla/xla/service/cpu/BUILD
@@ -1182,12 +1182,14 @@ cc_library(
     copts = runtime_copts(),
     visibility = ["//visibility:public"],
     deps = [
+        "//xla:executable_run_options",
         "//xla:shape_util",
         "//xla:xla_data_proto_cc",
         "//xla/ffi:attribute_map",
         "//xla/ffi:call_frame",
         "//xla/ffi:ffi_api",
         "//xla/service:custom_call_status_public_headers",
+        "//xla/service:executable",
         "@com_google_absl//absl/algorithm:container",
         "@com_google_absl//absl/base:core_headers",
         "@com_google_absl//absl/base:dynamic_annotations",
diff --git a/third_party/xla/xla/service/cpu/ir_emitter.cc b/third_party/xla/xla/service/cpu/ir_emitter.cc
index a2d5a84d2c3468..dff2ee5e7cde59 100644
--- a/third_party/xla/xla/service/cpu/ir_emitter.cc
+++ b/third_party/xla/xla/service/cpu/ir_emitter.cc
@@ -3178,6 +3178,7 @@ llvm::Value* IrEmitter::EmitCallToFfi(HloCustomCallInstruction* custom_call,
   const auto opaque_ref = llvm_ir::AsStringRef(opaque);
 
   std::vector<llvm::Value*> arguments = {
+      GetExecutableRunOptionsArgument(),     // run_options_ptr
       b_.CreateGlobalStringPtr(target_ref),  // target_name_ptr
       b_.getInt64(target.size()),            // target_name_len
       output_address,                        // output
@@ -3193,10 +3194,10 @@ llvm::Value* IrEmitter::EmitCallToFfi(HloCustomCallInstruction* custom_call,
       result_shapes_alloca,                  // result_dims
   };
 
-  return EmitCallToFunc("__xla_cpu_runtime_HandleFfiCall", arguments,
+  return EmitCallToFunc(runtime::kHandleFfiCallSymbolName, arguments,
                         b_.getVoidTy(),
-                        /* does_not_throw = */ false,
-                        /* only_accesses_arg_memory = */ true);
+                        /*does_not_throw=*/false,
+                        /*only_accesses_arg_memory=*/true);
 }
 
 void IrEmitter::EmitTransferElements(llvm::Value* target, llvm::Value* source,
diff --git a/third_party/xla/xla/service/cpu/runtime_handle_ffi_call.cc b/third_party/xla/xla/service/cpu/runtime_handle_ffi_call.cc
index 7d3271db600634..b8dd5c628bec77 100644
--- a/third_party/xla/xla/service/cpu/runtime_handle_ffi_call.cc
+++ b/third_party/xla/xla/service/cpu/runtime_handle_ffi_call.cc
@@ -27,19 +27,18 @@ limitations under the License.
 #include "absl/strings/str_cat.h"
 #include "absl/strings/string_view.h"
 #include "absl/types/span.h"
-#include "llvm/ADT/TypeSwitch.h"
 #include "mlir/AsmParser/AsmParser.h"  // from @llvm-project
 #include "mlir/IR/Attributes.h"  // from @llvm-project
-#include "mlir/IR/Builders.h"  // from @llvm-project
 #include "mlir/IR/BuiltinAttributes.h"  // from @llvm-project
 #include "mlir/IR/MLIRContext.h"  // from @llvm-project
+#include "xla/executable_run_options.h"
 #include "xla/ffi/attribute_map.h"
 #include "xla/ffi/call_frame.h"
 #include "xla/ffi/ffi_api.h"
 #include "xla/primitive_util.h"
 #include "xla/service/custom_call_status.h"
+#include "xla/service/service_executable_run_options.h"
 #include "xla/xla_data.pb.h"
-#include "tsl/platform/errors.h"
 #include "tsl/platform/logging.h"
 #include "tsl/platform/statusor.h"
 
@@ -108,11 +107,12 @@ void BuildBuffers(absl::Span<const int32_t> types, int64_t* encoded_dims,
   }
 }
 
-inline absl::Status BuildAndCallFfi(
-    std::string_view target_name, std::string_view backend_config,
-    absl::Span<void* const> outputs, absl::Span<void* const> inputs,
-    absl::Span<const int32_t> result_types, int64_t* result_dims,
-    absl::Span<const int32_t> operand_types, int64_t* operand_dims) {
+static absl::Status BuildAndCallFfi(
+    const xla::ExecutableRunOptions* run_options, std::string_view target_name,
+    std::string_view backend_config, absl::Span<void* const> outputs,
+    absl::Span<void* const> inputs, absl::Span<const int32_t> result_types,
+    int64_t* result_dims, absl::Span<const int32_t> operand_types,
+    int64_t* operand_dims) {
   CHECK_EQ(outputs.size(), result_types.size());
   CHECK_EQ(inputs.size(), operand_types.size());
 
@@ -161,18 +161,24 @@ inline absl::Status BuildAndCallFfi(
   BuildBuffers(operand_types, operand_dims, inputs, ArgInserter(builder));
   BuildBuffers(result_types, result_dims, outputs, RetInserter(builder));
 
+  // Forward executable run options to the FFI handlers via the call options.
+  xla::ServiceExecutableRunOptions service_run_options(*run_options);
+
+  ffi::CallOptions call_options;
+  call_options.run_options = &service_run_options;
+
   ffi::CallFrame call_frame = builder.Build();
-  return ffi::Call(registration->handler, call_frame);  // Status
+  return ffi::Call(registration->handler, call_frame, call_options);
 }
 
 }  // namespace
 
 ABSL_ATTRIBUTE_NO_SANITIZE_MEMORY void __xla_cpu_runtime_HandleFfiCall(
-    const char* target_name_ptr, int64_t target_name_len, void* output,
-    void** inputs, const char* opaque_str_ptr, int64_t opaque_str_len,
-    void* status_opaque, int32_t* operand_types, int64_t operand_count,
-    int64_t* operand_dims, int32_t* result_types, int64_t result_count,
-    int64_t* result_dims) {
+    const void* run_options_ptr, const char* target_name_ptr,
+    int64_t target_name_len, void* output, void** inputs,
+    const char* opaque_str_ptr, int64_t opaque_str_len, void* status_opaque,
+    int32_t* operand_types, int64_t operand_count, int64_t* operand_dims,
+    int32_t* result_types, int64_t result_count, int64_t* result_dims) {
   auto target_name = absl::string_view(target_name_ptr, target_name_len);
   auto backend_config = absl::string_view(opaque_str_ptr, opaque_str_len);
   auto xla_status = reinterpret_cast<XlaCustomCallStatus*>(status_opaque);
@@ -189,8 +195,12 @@ ABSL_ATTRIBUTE_NO_SANITIZE_MEMORY void __xla_cpu_runtime_HandleFfiCall(
   ABSL_ANNOTATE_MEMORY_IS_INITIALIZED(operand_types,
                                       operand_count * sizeof(int32_t));
 
+  const xla::ExecutableRunOptions* run_options =
+      reinterpret_cast<const xla::ExecutableRunOptions*>(run_options_ptr);
+
   absl::Status status = BuildAndCallFfi(
-      target_name, backend_config, absl::MakeSpan(outputs, result_count),
+      run_options, target_name, backend_config,
+      absl::MakeSpan(outputs, result_count),
       absl::MakeSpan(inputs, operand_count),
       absl::MakeSpan(result_types, result_count), result_dims,
       absl::MakeSpan(operand_types, operand_count), operand_dims);
diff --git a/third_party/xla/xla/service/cpu/runtime_handle_ffi_call.h b/third_party/xla/xla/service/cpu/runtime_handle_ffi_call.h
index e8afb236b2bb26..22511ebfa5c8f6 100644
--- a/third_party/xla/xla/service/cpu/runtime_handle_ffi_call.h
+++ b/third_party/xla/xla/service/cpu/runtime_handle_ffi_call.h
@@ -21,11 +21,11 @@ limitations under the License.
 extern "C" {
 
 extern void __xla_cpu_runtime_HandleFfiCall(
-    const char* target_name_ptr, int64_t target_name_len, void* output,
-    void** inputs, const char* opaque_str_ptr, int64_t opaque_str_len,
-    void* status_opaque, int32_t* operand_types, int64_t operand_count,
-    int64_t* operand_dims, int32_t* result_types, int64_t result_count,
-    int64_t* result_dims);
+    const void* run_options_ptr, const char* target_name_ptr,
+    int64_t target_name_len, void* output, void** inputs,
+    const char* opaque_str_ptr, int64_t opaque_str_len, void* status_opaque,
+    int32_t* operand_types, int64_t operand_count, int64_t* operand_dims,
+    int32_t* result_types, int64_t result_count, int64_t* result_dims);
 
 }  // extern "C"
 

From 8cf028f777baaf2df87a662503b71bcaf6960ab6 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Wed, 15 May 2024 13:18:59 -0700
Subject: [PATCH 207/478] Update ops-related pbtxt files.

PiperOrigin-RevId: 634053943
---
 tensorflow/core/ops/ops.pbtxt | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorflow/core/ops/ops.pbtxt b/tensorflow/core/ops/ops.pbtxt
index 975f4a22eca161..2073f2e00ecd81 100644
--- a/tensorflow/core/ops/ops.pbtxt
+++ b/tensorflow/core/ops/ops.pbtxt
@@ -1,4 +1,4 @@
-go/debugstr  
+go/debugonly    
 op {
   name: "Abort"
   attr {

From f8c0326cc7ad3fb082de0628f835e74ec3dce0be Mon Sep 17 00:00:00 2001
From: Ashiq Imran <ashiq.imran@intel.com>
Date: Wed, 15 May 2024 14:04:28 -0700
Subject: [PATCH 208/478] upgrading oneDNN version 3.4.1

---
 tensorflow/workspace2.bzl           | 6 +++---
 third_party/mkl_dnn/mkldnn_v1.BUILD | 4 ++--
 2 files changed, 5 insertions(+), 5 deletions(-)

diff --git a/tensorflow/workspace2.bzl b/tensorflow/workspace2.bzl
index 77eea2ac869167..18e5e4a582e33b 100644
--- a/tensorflow/workspace2.bzl
+++ b/tensorflow/workspace2.bzl
@@ -205,9 +205,9 @@ def _tf_repositories():
     tf_http_archive(
         name = "onednn",
         build_file = "//third_party/mkl_dnn:mkldnn_v1.BUILD",
-        sha256 = "e291fa4702f4bcfa6c8c23cb5b6599f0fefa8f23bc08edb9e15ddc5254ab7843",
-        strip_prefix = "oneDNN-3.3.4",
-        urls = tf_mirror_urls("https://github.com/oneapi-src/oneDNN/archive/refs/tags/v3.3.4.tar.gz"),
+        sha256 = "906559a25581b292352420721112e1656d21029b66e8597816f9e741fbcdeadb",
+        strip_prefix = "oneDNN-3.4.1",
+        urls = tf_mirror_urls("https://github.com/oneapi-src/oneDNN/archive/refs/tags/v3.4.1.tar.gz"),
     )
 
     tf_http_archive(
diff --git a/third_party/mkl_dnn/mkldnn_v1.BUILD b/third_party/mkl_dnn/mkldnn_v1.BUILD
index cc9e66d77e3a77..d36831ea4b9f53 100644
--- a/third_party/mkl_dnn/mkldnn_v1.BUILD
+++ b/third_party/mkl_dnn/mkldnn_v1.BUILD
@@ -94,8 +94,8 @@ expand_template(
     out = "include/oneapi/dnnl/dnnl_version.h",
     substitutions = {
         "@DNNL_VERSION_MAJOR@": "3",
-        "@DNNL_VERSION_MINOR@": "3",
-        "@DNNL_VERSION_PATCH@": "4",
+        "@DNNL_VERSION_MINOR@": "4",
+        "@DNNL_VERSION_PATCH@": "1",
         "@DNNL_VERSION_HASH@": "N/A",
     },
     template = "include/oneapi/dnnl/dnnl_version.h.in",

From 9d7e195200d8bd32516445a653575473eae7e7da Mon Sep 17 00:00:00 2001
From: Sania Nagpal <sanianagpal@google.com>
Date: Wed, 15 May 2024 13:48:13 -0700
Subject: [PATCH 209/478] [Testing] Code Coverage: Increase Code Coverage of
 //third_party/tensorflow/core/runtime_fallback/util to at least 75%

PiperOrigin-RevId: 634063098
---
 tensorflow/core/runtime_fallback/util/attr_util_test.cc | 8 +++++++-
 1 file changed, 7 insertions(+), 1 deletion(-)

diff --git a/tensorflow/core/runtime_fallback/util/attr_util_test.cc b/tensorflow/core/runtime_fallback/util/attr_util_test.cc
index b3febbb9373774..91811eb63dd01c 100644
--- a/tensorflow/core/runtime_fallback/util/attr_util_test.cc
+++ b/tensorflow/core/runtime_fallback/util/attr_util_test.cc
@@ -283,6 +283,7 @@ TEST(UtilsTest, FillAttrValueMapOk) {
   attrs.SetArray("shape", tfrt::ArrayRef<int64_t>{2, 2});
   attrs.SetArray("values", tfrt::ArrayRef<float>{2});
   attrs.SetArray("flags", tfrt::ArrayRef<bool>{false, true});
+  attrs.SetArray("baz", tfrt::ArrayRef<char>{'a'});
 
   attrs.Set<bool>("transpose_a", false);
   attrs.Set<bool>("transpose_b", true);
@@ -290,6 +291,9 @@ TEST(UtilsTest, FillAttrValueMapOk) {
   attrs.Set<float>("foo", 2);
   attrs.Set<int64_t>("bar", 2);
 
+  tfrt::AggregateAttr aggAttr;
+  attrs.Set<tfrt::AggregateAttr>("aggAttr", aggAttr);
+
   AttrValueMap map;
   auto host_context = CreateTestHostContext();
 
@@ -303,10 +307,12 @@ TEST(UtilsTest, FillAttrValueMapOk) {
           Pair(Eq("shape"), EqualsProto(R"pb(list { i: 2 i: 2 })pb")),
           Pair(Eq("values"), EqualsProto(R"pb(list { f: 2 })pb")),
           Pair(Eq("flags"), EqualsProto(R"pb(list { b: false b: true })pb")),
+          Pair(Eq("baz"), EqualsProto(R"pb(s: "a")pb")),
           Pair(Eq("transpose_a"), EqualsProto(R"pb(b: false)pb")),
           Pair(Eq("transpose_b"), EqualsProto(R"pb(b: true)pb")),
           Pair(Eq("foo"), EqualsProto(R"pb(f: 2)pb")),
-          Pair(Eq("bar"), EqualsProto(R"pb(i: 2)pb"))));
+          Pair(Eq("bar"), EqualsProto(R"pb(i: 2)pb")),
+          Pair(Eq("aggAttr"), EqualsProto(R"pb(list {})pb"))));
 }
 
 }  // namespace

From 3c4348dfbfb62ba919f9654c35d0c7d96ff99a8c Mon Sep 17 00:00:00 2001
From: Fangrui Song <maskray@google.com>
Date: Wed, 15 May 2024 13:48:34 -0700
Subject: [PATCH 210/478] Integrate LLVM at llvm/llvm-project@45726c1a3a3d

Updates LLVM usage to match
[45726c1a3a3d](https://github.com/llvm/llvm-project/commit/45726c1a3a3d)

PiperOrigin-RevId: 634063221
---
 third_party/llvm/workspace.bzl | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/third_party/llvm/workspace.bzl b/third_party/llvm/workspace.bzl
index a7d6c7b299393d..573719561ffed1 100644
--- a/third_party/llvm/workspace.bzl
+++ b/third_party/llvm/workspace.bzl
@@ -4,8 +4,8 @@ load("//third_party:repo.bzl", "tf_http_archive")
 
 def repo(name):
     """Imports LLVM."""
-    LLVM_COMMIT = "f89b1b8a68065c4b880417abb0563bce21399b52"
-    LLVM_SHA256 = "eba609768a9fbce79c7096d156fc96e999296cc61aca9ca736df7a9c5236215c"
+    LLVM_COMMIT = "45726c1a3a3d89ff9f6ebe657c3cb7bcd59b88db"
+    LLVM_SHA256 = "61b637626b1d0c7f67b5387a9e6a49097dbdfff85f578b53aa3a4fb63d8cb362"
 
     tf_http_archive(
         name = name,

From 53d603baeb378a66c592ed02fb110140324a2558 Mon Sep 17 00:00:00 2001
From: Alexander Lyashuk <crem@google.com>
Date: Wed, 15 May 2024 13:54:54 -0700
Subject: [PATCH 211/478] Make it possible to lower fp8 `tt.splat`.

Before the fix, `tt.splat` was lowered to e.g.
```
%14 = "llvm.mlir.constant"() <{value = 0.000000e+00 : f8E4M3FNUZ}> : () -> f8E4M3FNUZ
```
which LLVM rejected.

Translating the result type through typeConverter is what is done in other similar places. It results in
```
%14 = "llvm.mlir.constant"() <{value = 0.000000e+00 : f8E4M3FNUZ}> : () -> i8
```
and it's accepted by LLVM. During the MLIR to LLVM lowering, the fp8 value is converted to i8 with the correct binary representation.

The `isFloat()` function that is updated happened to have just one caller (in ArithConstantSplatOpConversion)

PiperOrigin-RevId: 634065311
---
 third_party/triton/temporary/fp8_splat.patch  | 41 +++++++++++++++++++
 third_party/triton/temporary/series.bzl       |  1 +
 .../triton/temporary/fp8_splat.patch          | 41 +++++++++++++++++++
 .../third_party/triton/temporary/series.bzl   |  1 +
 4 files changed, 84 insertions(+)
 create mode 100644 third_party/triton/temporary/fp8_splat.patch
 create mode 100644 third_party/xla/third_party/triton/temporary/fp8_splat.patch

diff --git a/third_party/triton/temporary/fp8_splat.patch b/third_party/triton/temporary/fp8_splat.patch
new file mode 100644
index 00000000000000..7a709e85419821
--- /dev/null
+++ b/third_party/triton/temporary/fp8_splat.patch
@@ -0,0 +1,41 @@
+diff --git a/include/triton/Conversion/MLIRTypes.h b/include/triton/Conversion/MLIRTypes.h
+--- a/include/triton/Conversion/MLIRTypes.h
++++ b/include/triton/Conversion/MLIRTypes.h
+@@ -26,12 +26,6 @@ inline Type f32Ty(MLIRContext *ctx) { re
+ inline Type f64Ty(MLIRContext *ctx) { return FloatType::getF64(ctx); }
+ inline Type bf16Ty(MLIRContext *ctx) { return FloatType::getBF16(ctx); }
+ 
+-inline bool isFloat(Type type) {
+-  return type.isF32() || type.isF64() || type.isF16() || type.isF128();
+-}
+-
+-inline bool isInt(Type type) { return type.isIntOrFloat() && !isFloat(type); }
+-
+ } // namespace type
+ } // namespace triton
+ } // namespace mlir
+diff --git a/lib/Conversion/TritonGPUToLLVM/ViewOpToLLVM.cpp b/lib/Conversion/TritonGPUToLLVM/ViewOpToLLVM.cpp
+--- a/lib/Conversion/TritonGPUToLLVM/ViewOpToLLVM.cpp
++++ b/lib/Conversion/TritonGPUToLLVM/ViewOpToLLVM.cpp
+@@ -74,17 +74,18 @@ struct ArithConstantSplatOpConversion
+     auto values = mlir::dyn_cast<SplatElementsAttr>(op.getValue());
+     auto elemType = values.getElementType();
+     Attribute val;
+-    if (elemType.isBF16() || type::isFloat(elemType)) {
++    if (isa<FloatType>(elemType)) {
+       val = values.getValues<FloatAttr>()[0];
+-    } else if (type::isInt(elemType)) {
++    } else if (isa<IntegerType>(elemType)) {
+       val = values.getValues<IntegerAttr>()[0];
+     } else {
+       llvm::errs() << "ArithConstantSplatOpConversion get unsupported type: "
+                    << value.getType() << "\n";
+       return failure();
+     }
+-    auto constOp = rewriter.create<LLVM::ConstantOp>(loc, elemType, val);
+     auto typeConverter = getTypeConverter();
++    auto constOp = rewriter.create<LLVM::ConstantOp>(
++        loc, typeConverter->convertType(elemType), val);
+     auto llStruct = SplatOpConversion::convertSplatLikeOp(
+         elemType, op.getType(), constOp, typeConverter, rewriter, loc);
+     rewriter.replaceOp(op, llStruct);
diff --git a/third_party/triton/temporary/series.bzl b/third_party/triton/temporary/series.bzl
index 382cb9e9de0154..d68c23d5babbd6 100644
--- a/third_party/triton/temporary/series.bzl
+++ b/third_party/triton/temporary/series.bzl
@@ -6,6 +6,7 @@ internal patch during the next triton integration process.
 """
 
 temporary_patch_list = [
+    "//third_party/triton/temporary:fp8_splat.patch",
     "//third_party/triton/temporary:mma_limit_pred.patch",
     "//third_party/triton/temporary:fix_register_constraints.patch",
 ]
diff --git a/third_party/xla/third_party/triton/temporary/fp8_splat.patch b/third_party/xla/third_party/triton/temporary/fp8_splat.patch
new file mode 100644
index 00000000000000..7a709e85419821
--- /dev/null
+++ b/third_party/xla/third_party/triton/temporary/fp8_splat.patch
@@ -0,0 +1,41 @@
+diff --git a/include/triton/Conversion/MLIRTypes.h b/include/triton/Conversion/MLIRTypes.h
+--- a/include/triton/Conversion/MLIRTypes.h
++++ b/include/triton/Conversion/MLIRTypes.h
+@@ -26,12 +26,6 @@ inline Type f32Ty(MLIRContext *ctx) { re
+ inline Type f64Ty(MLIRContext *ctx) { return FloatType::getF64(ctx); }
+ inline Type bf16Ty(MLIRContext *ctx) { return FloatType::getBF16(ctx); }
+ 
+-inline bool isFloat(Type type) {
+-  return type.isF32() || type.isF64() || type.isF16() || type.isF128();
+-}
+-
+-inline bool isInt(Type type) { return type.isIntOrFloat() && !isFloat(type); }
+-
+ } // namespace type
+ } // namespace triton
+ } // namespace mlir
+diff --git a/lib/Conversion/TritonGPUToLLVM/ViewOpToLLVM.cpp b/lib/Conversion/TritonGPUToLLVM/ViewOpToLLVM.cpp
+--- a/lib/Conversion/TritonGPUToLLVM/ViewOpToLLVM.cpp
++++ b/lib/Conversion/TritonGPUToLLVM/ViewOpToLLVM.cpp
+@@ -74,17 +74,18 @@ struct ArithConstantSplatOpConversion
+     auto values = mlir::dyn_cast<SplatElementsAttr>(op.getValue());
+     auto elemType = values.getElementType();
+     Attribute val;
+-    if (elemType.isBF16() || type::isFloat(elemType)) {
++    if (isa<FloatType>(elemType)) {
+       val = values.getValues<FloatAttr>()[0];
+-    } else if (type::isInt(elemType)) {
++    } else if (isa<IntegerType>(elemType)) {
+       val = values.getValues<IntegerAttr>()[0];
+     } else {
+       llvm::errs() << "ArithConstantSplatOpConversion get unsupported type: "
+                    << value.getType() << "\n";
+       return failure();
+     }
+-    auto constOp = rewriter.create<LLVM::ConstantOp>(loc, elemType, val);
+     auto typeConverter = getTypeConverter();
++    auto constOp = rewriter.create<LLVM::ConstantOp>(
++        loc, typeConverter->convertType(elemType), val);
+     auto llStruct = SplatOpConversion::convertSplatLikeOp(
+         elemType, op.getType(), constOp, typeConverter, rewriter, loc);
+     rewriter.replaceOp(op, llStruct);
diff --git a/third_party/xla/third_party/triton/temporary/series.bzl b/third_party/xla/third_party/triton/temporary/series.bzl
index 382cb9e9de0154..d68c23d5babbd6 100644
--- a/third_party/xla/third_party/triton/temporary/series.bzl
+++ b/third_party/xla/third_party/triton/temporary/series.bzl
@@ -6,6 +6,7 @@ internal patch during the next triton integration process.
 """
 
 temporary_patch_list = [
+    "//third_party/triton/temporary:fp8_splat.patch",
     "//third_party/triton/temporary:mma_limit_pred.patch",
     "//third_party/triton/temporary:fix_register_constraints.patch",
 ]

From 6a519962a4b7b6fa31bb54914463d5726a9230c3 Mon Sep 17 00:00:00 2001
From: Bixia Zheng <bixia@google.com>
Date: Wed, 15 May 2024 13:57:38 -0700
Subject: [PATCH 212/478] [xla:spmd] Fix a problem in generating code for
 concat that requires padding.

When an internal concat operand requires padding, previously, we don't zero out
the padded elements and the generated code produces incorrect results.

Fix two existing tests.

PiperOrigin-RevId: 634066344
---
 third_party/xla/xla/service/spmd/spmd_partitioner.cc |  3 ++-
 .../xla/xla/service/spmd/spmd_partitioner_test.cc    | 12 ++++++++++--
 2 files changed, 12 insertions(+), 3 deletions(-)

diff --git a/third_party/xla/xla/service/spmd/spmd_partitioner.cc b/third_party/xla/xla/service/spmd/spmd_partitioner.cc
index d0cc58af4e341d..2d99310671f39b 100644
--- a/third_party/xla/xla/service/spmd/spmd_partitioner.cc
+++ b/third_party/xla/xla/service/spmd/spmd_partitioner.cc
@@ -2581,7 +2581,8 @@ Status SpmdPartitioningVisitor::HandleConcatenate(HloInstruction* hlo) {
   int64_t offset = 0;
   auto state = MakePartitioningState();
   for (HloInstruction* operand : hlo->operands()) {
-    auto spmd_operand = GetPartitionedHlo(operand).Reshard(sharding).hlo();
+    auto spmd_operand =
+        GetPartitionedHlo(operand).Reshard(sharding).PadWithZero().hlo();
     std::vector<HloInstruction*> start_indices(
         hlo->shape().rank(), b_.AddInstruction(HloInstruction::CreateConstant(
                                  LiteralUtil::Zero(S32))));
diff --git a/third_party/xla/xla/service/spmd/spmd_partitioner_test.cc b/third_party/xla/xla/service/spmd/spmd_partitioner_test.cc
index 406580a9ad302e..e542984ae13152 100644
--- a/third_party/xla/xla/service/spmd/spmd_partitioner_test.cc
+++ b/third_party/xla/xla/service/spmd/spmd_partitioner_test.cc
@@ -2207,13 +2207,17 @@ ENTRY entry {
       AllOf(op::Copy(op::DynamicSlice(op::Pad(op::Parameter(), op::Constant()),
                                       op::Constant(), op::Reshape())),
             op::Shape("f32[14,129]"));
+  auto param0_adjusted =
+      AllOf(op::Select(op::Compare(op::Add(), op::Broadcast(op::Constant())),
+                       param0, op::Broadcast(op::Constant())),
+            op::Shape("f32[14,129]"));
   auto param1 = AllOf(op::Copy(op::DynamicSlice(op::Parameter(), op::Constant(),
                                                 op::Reshape())),
                       op::Shape("f32[14,58]"));
   EXPECT_THAT(root, AllOf(op::DynamicSlice(
                               AllOf(op::AllReduce(op::DynamicUpdateSlice(
                                         op::DynamicUpdateSlice(
-                                            op::Broadcast(), param0,
+                                            op::Broadcast(), param0_adjusted,
                                             op::Constant(), op::Multiply()),
                                         param1, op::Constant(), op::Add())),
                                     op::Shape("f32[14,374]")),
@@ -2238,11 +2242,15 @@ ENTRY entry {
 
   const auto root = module->entry_computation()->root_instruction();
   auto param0 = AllOf(op::Parameter(0), op::Shape("f32[7,129]"));
+  auto param0_adjusted =
+      AllOf(op::Select(op::Compare(op::Add(), op::Broadcast(op::Constant())),
+                       param0, op::Broadcast(op::Constant())),
+            op::Shape("f32[7,129]"));
   auto param1 = AllOf(op::Parameter(1), op::Shape("f32[7,58]"));
   EXPECT_THAT(root, AllOf(op::DynamicSlice(
                               AllOf(op::AllReduce(op::DynamicUpdateSlice(
                                         op::DynamicUpdateSlice(
-                                            op::Broadcast(), param0,
+                                            op::Broadcast(), param0_adjusted,
                                             op::Constant(), op::Multiply()),
                                         param1, op::Constant(), op::Add())),
                                     op::Shape("f32[7,374]")),

From b95622d541412d93916135fe4194393887717d1a Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Wed, 15 May 2024 13:59:36 -0700
Subject: [PATCH 213/478] Update TFRT dependency to use revision
 http://github.com/tensorflow/runtime/commit/ac3f51614db030270303c593a41f46bea68dab3d.

PiperOrigin-RevId: 634067000
---
 third_party/tf_runtime/workspace.bzl                          | 4 ++--
 .../xla/third_party/tsl/third_party/tf_runtime/workspace.bzl  | 4 ++--
 2 files changed, 4 insertions(+), 4 deletions(-)

diff --git a/third_party/tf_runtime/workspace.bzl b/third_party/tf_runtime/workspace.bzl
index 8cd9762125eee0..ee7e50c4cb4178 100644
--- a/third_party/tf_runtime/workspace.bzl
+++ b/third_party/tf_runtime/workspace.bzl
@@ -6,8 +6,8 @@ def repo():
     """Imports TFRT."""
 
     # Attention: tools parse and update these lines.
-    TFRT_COMMIT = "7bdf48f1aac0b48ff85a4e0fb5ff7f98a703f8d6"
-    TFRT_SHA256 = "dc02eeae8f6c3e72bf51cad0f04676eeefabe241389b6f08b8310c999eeb64b3"
+    TFRT_COMMIT = "ac3f51614db030270303c593a41f46bea68dab3d"
+    TFRT_SHA256 = "a904f3ee171c84c2b5e1bdbaf3f7a0faa43f036383b17f1f69c8abbda3f0936f"
 
     tf_http_archive(
         name = "tf_runtime",
diff --git a/third_party/xla/third_party/tsl/third_party/tf_runtime/workspace.bzl b/third_party/xla/third_party/tsl/third_party/tf_runtime/workspace.bzl
index 8cd9762125eee0..ee7e50c4cb4178 100644
--- a/third_party/xla/third_party/tsl/third_party/tf_runtime/workspace.bzl
+++ b/third_party/xla/third_party/tsl/third_party/tf_runtime/workspace.bzl
@@ -6,8 +6,8 @@ def repo():
     """Imports TFRT."""
 
     # Attention: tools parse and update these lines.
-    TFRT_COMMIT = "7bdf48f1aac0b48ff85a4e0fb5ff7f98a703f8d6"
-    TFRT_SHA256 = "dc02eeae8f6c3e72bf51cad0f04676eeefabe241389b6f08b8310c999eeb64b3"
+    TFRT_COMMIT = "ac3f51614db030270303c593a41f46bea68dab3d"
+    TFRT_SHA256 = "a904f3ee171c84c2b5e1bdbaf3f7a0faa43f036383b17f1f69c8abbda3f0936f"
 
     tf_http_archive(
         name = "tf_runtime",

From 46b882a708974e276796193d7e916cce0ad137c5 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Wed, 15 May 2024 14:27:18 -0700
Subject: [PATCH 214/478] Disable some tests that time out on certain
 platforms.

PiperOrigin-RevId: 634076395
---
 third_party/xla/xla/tests/BUILD               |   4 +-
 .../xla/xla/tests/dot_operation_test.cc       |   2 +-
 .../xla/xla/tests/reduce_window_test.cc       | 111 +++++++++++-------
 .../xla/xla/tests/transfer_manager_test.cc    |   2 +-
 4 files changed, 71 insertions(+), 48 deletions(-)

diff --git a/third_party/xla/xla/tests/BUILD b/third_party/xla/xla/tests/BUILD
index 563a8ea6b203c7..ea9aa2f35d7ef4 100644
--- a/third_party/xla/xla/tests/BUILD
+++ b/third_party/xla/xla/tests/BUILD
@@ -966,6 +966,7 @@ xla_test(
         "gpu",
         "interpreter",
     ],
+    shard_count = 50,
     tags = [
         "nozapfhahn",
         "optonly",
@@ -1647,7 +1648,7 @@ xla_test(
     name = "reduce_window_test",
     timeout = "long",
     srcs = [],
-    shard_count = 20,
+    shard_count = 40,
     tags = [
         "optonly",
     ],
@@ -2653,6 +2654,7 @@ xla_cc_test(
 xla_test(
     name = "transfer_manager_test",
     srcs = ["transfer_manager_test.cc"],
+    shard_count = 50,
     deps = [
         ":literal_test_util",
         ":local_client_test_base",
diff --git a/third_party/xla/xla/tests/dot_operation_test.cc b/third_party/xla/xla/tests/dot_operation_test.cc
index fc6ccb8f8f5b01..f272d8da4e00d3 100644
--- a/third_party/xla/xla/tests/dot_operation_test.cc
+++ b/third_party/xla/xla/tests/dot_operation_test.cc
@@ -472,7 +472,7 @@ std::vector<DotTestParam> CreateDotTestParameters() {
 XLA_TEST_P(ParametricDotTest, TestF16) { TestImpl<Eigen::half>(); }
 #endif
 XLA_TEST_P(ParametricDotTest, TestF32) { TestImpl<float>(); }
-XLA_TEST_P(ParametricDotTest, TestF64) { TestImpl<double>(); }
+XLA_TEST_P(ParametricDotTest, OVERSIZE_ON_GRM(TestF64)) { TestImpl<double>(); }
 XLA_TEST_P(ParametricDotTest, TestC64) { TestImpl<std::complex<float>>(); }
 #ifndef XLA_BACKEND_DOES_NOT_SUPPORT_COMPLEX128
 XLA_TEST_P(ParametricDotTest, TestC128) { TestImpl<std::complex<double>>(); }
diff --git a/third_party/xla/xla/tests/reduce_window_test.cc b/third_party/xla/xla/tests/reduce_window_test.cc
index b76065f55649f7..ccbb8f4f3cb8ba 100644
--- a/third_party/xla/xla/tests/reduce_window_test.cc
+++ b/third_party/xla/xla/tests/reduce_window_test.cc
@@ -889,7 +889,10 @@ INSTANTIATE_TEST_CASE_P(
 
 class R4ReduceWindowLargeTest : public R4ReduceWindowTest {};
 
-XLA_TEST_P(R4ReduceWindowLargeTest, DISABLED_ON_INTERPRETER(DoIt)) { DoIt(); }
+XLA_TEST_P(R4ReduceWindowLargeTest,
+           OVERSIZE_ON_GRM(DISABLED_ON_INTERPRETER(DoIt))) {
+  DoIt();
+}
 
 // Test cases that are large/slow/failed.
 const R4ReduceWindowTestData kR4ReduceWindowLargeTestValues[] = {
@@ -980,7 +983,9 @@ struct R3ReduceWindowTestData {
   int64_t layout[3];
   Padding padding;
   Reducer reducer;
-} kR3TestCases[] = {
+};
+
+R3ReduceWindowTestData kR3TestCases[] = {
     {/*base_bounds=*/{2, 1, 2}, /*window_bounds=*/{1, 1, 2},
      /*strides=*/{1, 1, 1}, /*layout=*/{2, 1, 0},
      /*padding=*/Padding::kValid, /*reducer=*/Reducer::kAdd},
@@ -1014,15 +1019,6 @@ struct R3ReduceWindowTestData {
     {/*base_bounds=*/{63, 261, 257}, /*window_bounds=*/{63, 261, 257},
      /*strides=*/{1, 1, 1}, /*layout=*/{2, 1, 0},
      /*padding=*/Padding::kValid, /*reducer=*/Reducer::kMax},
-    {/*base_bounds=*/{10003, 10, 5}, /*window_bounds=*/{9999, 7, 3},
-     /*strides=*/{1, 1, 1}, /*layout=*/{2, 1, 0},
-     /*padding=*/Padding::kValid, /*reducer=*/Reducer::kAdd},
-    {/*base_bounds=*/{9999, 1, 1}, /*window_bounds=*/{9999, 1, 1},
-     /*strides=*/{1, 1, 1}, /*layout=*/{2, 1, 0},
-     /*padding=*/Padding::kValid, /*reducer=*/Reducer::kAdd},
-    {/*base_bounds=*/{10003, 10, 5}, /*window_bounds=*/{9999, 7, 3},
-     /*strides=*/{2, 2, 2}, /*layout=*/{2, 1, 0},
-     /*padding=*/Padding::kValid, /*reducer=*/Reducer::kAdd},
 };
 
 std::string R3ReduceWindowTestDataToString(
@@ -1047,46 +1043,48 @@ class R3ReduceWindowTest : public ReduceWindowTestBase,
                                ::testing::tuple<R3ReduceWindowTestData, bool>> {
  protected:
   R3ReduceWindowTest() { set_use_bfloat16(::testing::get<1>(GetParam())); }
-};
 
-XLA_TEST_P(R3ReduceWindowTest, DoIt) {
-  XlaBuilder b(TestName());
-  const auto& param = ::testing::get<0>(GetParam());
+  void DoIt() {
+    XlaBuilder b(TestName());
+    const auto& param = ::testing::get<0>(GetParam());
 
-  const float kInitValue = 0.0f;
-  Array3D<float> input(param.base_bounds[0], param.base_bounds[1],
-                       param.base_bounds[2]);
-  // Choose a prime iota length so that each window sees a unique set of values.
-  // (Technically, the requirement is that the iota length is relatively prime
-  // to all of the dimensions involved in the reduce-window.)
-  input.FillRepeatedIota(0, 137);
-  Literal input_literal = LiteralUtil::CreateR3FromArray3DWithLayout(
-      input, LayoutUtil::MakeLayout(param.layout));
-  auto reducer = param.reducer;
-  if (use_bfloat16()) {
-    input_literal = LiteralUtil::ConvertF32ToBF16(input_literal);
-
-    // To avoid numerical issues, force the reducer to be kMax for bf16
-    // inputs.
-    reducer = kMax;
-  }
+    const float kInitValue = 0.0f;
+    Array3D<float> input(param.base_bounds[0], param.base_bounds[1],
+                         param.base_bounds[2]);
+    // Choose a prime iota length so that each window sees a unique set of
+    // values. (Technically, the requirement is that the iota length is
+    // relatively prime to all of the dimensions involved in the reduce-window.)
+    input.FillRepeatedIota(0, 137);
+    Literal input_literal = LiteralUtil::CreateR3FromArray3DWithLayout(
+        input, LayoutUtil::MakeLayout(param.layout));
+    auto reducer = param.reducer;
+    if (use_bfloat16()) {
+      input_literal = LiteralUtil::ConvertF32ToBF16(input_literal);
 
-  XlaOp parameter = Parameter(&b, 0, input_literal.shape(), "input");
-  auto init_value =
-      CreateConstantFromLiteral(LiteralUtil::CreateR0(kInitValue), &b);
+      // To avoid numerical issues, force the reducer to be kMax for bf16
+      // inputs.
+      reducer = kMax;
+    }
 
-  auto computation = reducer == kAdd
-                         ? CreateScalarAddComputation(FloatType(), &b)
-                         : CreateScalarMaxComputation(FloatType(), &b);
+    XlaOp parameter = Parameter(&b, 0, input_literal.shape(), "input");
+    auto init_value =
+        CreateConstantFromLiteral(LiteralUtil::CreateR0(kInitValue), &b);
 
-  ReduceWindow(/*operand=*/parameter,
-               /*init_value=*/init_value,
-               /*computation=*/computation,
-               /*window_dimensions=*/param.window_bounds,
-               /*window_strides=*/param.strides, /*padding=*/param.padding);
+    auto computation = reducer == kAdd
+                           ? CreateScalarAddComputation(FloatType(), &b)
+                           : CreateScalarMaxComputation(FloatType(), &b);
 
-  ComputeAndCompare(&b, {std::move(input_literal)}, DefaultErrorSpec());
-}
+    ReduceWindow(/*operand=*/parameter,
+                 /*init_value=*/init_value,
+                 /*computation=*/computation,
+                 /*window_dimensions=*/param.window_bounds,
+                 /*window_strides=*/param.strides, /*padding=*/param.padding);
+
+    ComputeAndCompare(&b, {std::move(input_literal)}, DefaultErrorSpec());
+  }
+};
+
+XLA_TEST_P(R3ReduceWindowTest, DoIt) { DoIt(); }
 
 INSTANTIATE_TEST_CASE_P(
     R3ReduceWindowTestInstantiation, R3ReduceWindowTest,
@@ -1094,6 +1092,29 @@ INSTANTIATE_TEST_CASE_P(
                        ::testing::ValuesIn(use_bfloat16_params)),
     R3ReduceWindowTestDataToString);
 
+class R3ReduceWindowLargeTest : public R3ReduceWindowTest {};
+
+XLA_TEST_P(R3ReduceWindowLargeTest, OVERSIZE_ON_GRM(DoIt)) { DoIt(); }
+
+// Test cases that are large/slow/failed.
+const R3ReduceWindowTestData kR3ReduceWindowLargeTestValues[] = {
+    {/*base_bounds=*/{10003, 10, 5}, /*window_bounds=*/{9999, 7, 3},
+     /*strides=*/{1, 1, 1}, /*layout=*/{2, 1, 0},
+     /*padding=*/Padding::kValid, /*reducer=*/Reducer::kAdd},
+    {/*base_bounds=*/{9999, 1, 1}, /*window_bounds=*/{9999, 1, 1},
+     /*strides=*/{1, 1, 1}, /*layout=*/{2, 1, 0},
+     /*padding=*/Padding::kValid, /*reducer=*/Reducer::kAdd},
+    {/*base_bounds=*/{10003, 10, 5}, /*window_bounds=*/{9999, 7, 3},
+     /*strides=*/{2, 2, 2}, /*layout=*/{2, 1, 0},
+     /*padding=*/Padding::kValid, /*reducer=*/Reducer::kAdd},
+};
+
+INSTANTIATE_TEST_CASE_P(
+    R3ReduceWindowLargeTestInstantiation, R3ReduceWindowLargeTest,
+    ::testing::Combine(::testing::ValuesIn(kR3ReduceWindowLargeTestValues),
+                       ::testing::ValuesIn(use_bfloat16_params)),
+    R3ReduceWindowTestDataToString);
+
 struct R2ReduceWindowTestData {
   int64_t base_bounds[2];
   int64_t window_bounds[2];
diff --git a/third_party/xla/xla/tests/transfer_manager_test.cc b/third_party/xla/xla/tests/transfer_manager_test.cc
index 0e4c7793b72ce2..e67e391be026d1 100644
--- a/third_party/xla/xla/tests/transfer_manager_test.cc
+++ b/third_party/xla/xla/tests/transfer_manager_test.cc
@@ -309,7 +309,7 @@ XLA_TEST_F(TransferManagerTest, TransferTokenFromDevice) {
   EXPECT_TRUE(LiteralTestUtil::Equal(LiteralUtil::CreateToken(), result));
 }
 
-XLA_TEST_F(TransferManagerTest, MultiStreamRoundTripSoak) {
+XLA_TEST_F(TransferManagerTest, OVERSIZE_ON_GRM(MultiStreamRoundTripSoak)) {
   const int64_t kIterationCount = 5000;
   Literal literal1 = LiteralUtil::MakeTupleFromSlices(
       {LiteralUtil::CreateR0<float>(123.0f),

From b74de8d7dbb0e5365075a4863dd245a45cde7d7d Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Wed, 15 May 2024 14:32:56 -0700
Subject: [PATCH 215/478] Updates the benchmarking tool documentation with
 discovery text about flex ops, more information about flex ops, and an
 additional section regarding creating a custom version of the app with
 user-supplied custom ops.

PiperOrigin-RevId: 634078145
---
 tensorflow/lite/tools/benchmark/README.md | 67 ++++++++++++++++++++---
 1 file changed, 60 insertions(+), 7 deletions(-)

diff --git a/tensorflow/lite/tools/benchmark/README.md b/tensorflow/lite/tools/benchmark/README.md
index aa607ea952e682..f25da51705d6b8 100644
--- a/tensorflow/lite/tools/benchmark/README.md
+++ b/tensorflow/lite/tools/benchmark/README.md
@@ -250,6 +250,11 @@ delegate first, and then the XNNPACK delegate secondly.
 
 ## To build/install/run
 
+Note: The benchmarking tool must be compiled with a TFLite runtime that
+supports the ops found in the model to be tested.<br/>
+If Tensorflow Ops ("flex ops")
+or other custom ops are used in the model, please see the section [below](#build-the-benchmark-tool-with-tensorflow-ops-support).
+
 ### On Android:
 
 (0) Refer to https://www.tensorflow.org/lite/guide/build_android to edit the
@@ -308,8 +313,8 @@ adb shell /data/local/tmp/benchmark_model \
 bazel build -c opt tensorflow/lite/tools/benchmark:benchmark_model
 ```
 
-(2) Run on your compute graph, similar to the Android case but without the need of adb shell.
-For example:
+(2) Run on your compute graph, similar to the Android case but without the need
+of adb shell. For example:
 
 ```
 bazel-bin/tensorflow/lite/tools/benchmark/benchmark_model \
@@ -442,12 +447,19 @@ some additional parameters as detailed below.
 
 ## Build the benchmark tool with Tensorflow ops support
 
-You can build the benchmark tool with [Tensorflow operators support](https://www.tensorflow.org/lite/guide/ops_select).
+If you see an error that says: `ERROR: Select TensorFlow op(s), included in the
+given model, is(are) not supported by this interpreter.` you will need to
+build with [Tensorflow operators support](https://www.tensorflow.org/lite/guide/ops_select).
+
+Having Tensorflow ops in the TFLite file works when the benchmark tool is built
+with Tensorflow ops support. It doesn't require any additional option to use it.
 
 ### How to build
 
-To build the tool, you need to use 'benchmark_model_plus_flex' target with
-'--config=monolithic' option.
+To build the tool, you need to use the `benchmark_model_plus_flex` target with
+the `--config=monolithic` flag.
+
+**Desktop**
 
 ```
 bazel build -c opt \
@@ -455,12 +467,53 @@ bazel build -c opt \
   tensorflow/lite/tools/benchmark:benchmark_model_plus_flex
 ```
 
+**Android**
+
+```
+bazel build -c opt \
+  --config=monolithic --config=android_arm64 \
+  tensorflow/lite/tools/benchmark:benchmark_model_plus_flex
+```
+
 ### How to benchmark tflite model with Tensorflow ops
 
-Tensorflow ops support just works the benchmark tool is built with Tensorflow
-ops support. It doesn't require any additional option to use it.
+Follow the further instructions [above](#to-buildinstallrun) replacing
+`benchmark_model` with the `benchmark_model_plus_flex` file created here.
+
+For example, on desktop it's very easy:
 
 ```
 bazel-bin/tensorflow/lite/tools/benchmark/benchmark_model_plus_flex \
   --graph=model_converted_with_TF_ops.tflite \
 ```
+
+## Build the benchmark tool with Custom ops support
+
+If you see an error that says `ERROR: Op type not registered 'XXXXXXXX'
+in binary running on localhost.` for custom ops running in your TFLite model,
+you will need to manually build the tool to include your libraries providing
+the custom ops.
+
+### How to build
+
+While possible, this is not necessarily supported.
+
+However, you should be able to create a new `cc_binary` rule that depends on
+`tensorflow/lite/tools/benchmark:benchmark_model_main` along with your custom op
+rules.
+
+```
+cc_binary(
+    name = "benchmark_model_plus_custom_ops",
+    deps = [
+        ":my_custom_ops_provider",
+        "//tensorflow/lite/tools/benchmark:benchmark_model_main",
+    ],
+)
+```
+
+### How to benchmark tflite model with Custom ops
+
+Use the `benchmark_model_plus_custom_ops` (or whatever) file created by your
+custom rule instead of the `benchmark_model` file in the instructions,
+[above](#to-buildinstallrun).

From 907bfb84509d56d4f7209881b7e7aca1e7165d49 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Wed, 15 May 2024 14:42:40 -0700
Subject: [PATCH 216/478] [Multi-host GPU]Add a utility function to build
 GpuTopologyProto from GlobalTopologyProto.

PiperOrigin-RevId: 634081336
---
 third_party/xla/xla/pjrt/distributed/BUILD    |  4 +-
 .../xla/xla/pjrt/distributed/topology_util.cc | 46 +++++++++++
 .../xla/xla/pjrt/distributed/topology_util.h  |  5 ++
 .../pjrt/distributed/topology_util_test.cc    | 76 +++++++++++++++++++
 4 files changed, 130 insertions(+), 1 deletion(-)

diff --git a/third_party/xla/xla/pjrt/distributed/BUILD b/third_party/xla/xla/pjrt/distributed/BUILD
index f75e76df583df9..c0ae37c0b3e3b7 100644
--- a/third_party/xla/xla/pjrt/distributed/BUILD
+++ b/third_party/xla/xla/pjrt/distributed/BUILD
@@ -50,10 +50,11 @@ xla_cc_test(
         ":in_memory_key_value_store",
         ":protocol_proto_cc",
         ":topology_util",
-        "@com_google_absl//absl/status",
+        "//xla:test_helpers",
         "@com_google_absl//absl/time",
         "@local_tsl//tsl/lib/core:status_test_util",
         "@local_tsl//tsl/platform:env",
+        "@local_tsl//tsl/platform:statusor",
         "@local_tsl//tsl/platform:test",
         "@local_tsl//tsl/platform:test_main",
     ],
@@ -124,6 +125,7 @@ cc_library(
         "//xla:util",
         "//xla/pjrt:pjrt_client",
         "//xla/pjrt:utils",
+        "//xla/pjrt/gpu:gpu_topology_proto_cc",
         "@com_google_absl//absl/container:flat_hash_map",
         "@com_google_absl//absl/status",
         "@com_google_absl//absl/strings",
diff --git a/third_party/xla/xla/pjrt/distributed/topology_util.cc b/third_party/xla/xla/pjrt/distributed/topology_util.cc
index bec8c617010ce6..e8649c68c5bb86 100644
--- a/third_party/xla/xla/pjrt/distributed/topology_util.cc
+++ b/third_party/xla/xla/pjrt/distributed/topology_util.cc
@@ -16,6 +16,7 @@ limitations under the License.
 #include "xla/pjrt/distributed/topology_util.h"
 
 #include <fstream>
+#include <map>
 #include <string>
 #include <string_view>
 #include <vector>
@@ -197,4 +198,49 @@ Status ExchangeTopologies(std::string_view platform, int node_id, int num_nodes,
   return absl::OkStatus();
 }
 
+absl::StatusOr<GpuTopologyProto> BuildGpuTopology(
+    const GlobalTopologyProto& global_topology) {
+  GpuTopologyProto gpu_topology;
+  std::map<int, std::vector<int>> slice_id_to_node_ids;
+  std::vector<int> device_ids;
+  for (int i = 0; i < global_topology.nodes_size(); ++i) {
+    const LocalTopologyProto& local_topology = global_topology.nodes(i);
+
+    slice_id_to_node_ids[local_topology.devices(0).slice_index()].push_back(
+        local_topology.node_id());
+
+    // Initializes GPU topology with the first local topology.
+    if (i == 0) {
+      gpu_topology.set_platform_version((local_topology.devices(0).name()));
+      gpu_topology.set_num_devices_per_host(local_topology.devices_size());
+    } else {
+      // Check for consistent number of devices per host.
+      if (gpu_topology.num_devices_per_host() !=
+          local_topology.devices_size()) {
+        return absl::InternalError(
+            "GpuTopology doesn't support multi-host with different number of "
+            "devices per host.");
+      }
+    }
+
+    for (const DeviceProto& device : local_topology.devices()) {
+      device_ids.push_back(device.global_device_id());
+    }
+  }
+
+  gpu_topology.set_num_slices(slice_id_to_node_ids.size());
+  gpu_topology.set_num_hosts_per_slice(
+      slice_id_to_node_ids.begin()->second.size());
+  // Check for consistent number of hosts per slice.
+  for (const auto& [boot_id, node_ids] : slice_id_to_node_ids) {
+    if (node_ids.size() != gpu_topology.num_hosts_per_slice()) {
+      return absl::InternalError(
+          "GpuTopology doesn't support multi-host with different number of "
+          "hosts per slice.");
+    }
+  }
+  gpu_topology.mutable_device_ids()->Add(device_ids.begin(), device_ids.end());
+  return gpu_topology;
+}
+
 }  // namespace xla
diff --git a/third_party/xla/xla/pjrt/distributed/topology_util.h b/third_party/xla/xla/pjrt/distributed/topology_util.h
index a74478cdb0b1ca..c25788b9015adc 100644
--- a/third_party/xla/xla/pjrt/distributed/topology_util.h
+++ b/third_party/xla/xla/pjrt/distributed/topology_util.h
@@ -23,6 +23,7 @@ limitations under the License.
 #include "absl/types/span.h"
 #include "xla/pjrt/distributed/key_value_store_interface.h"
 #include "xla/pjrt/distributed/protocol.pb.h"
+#include "xla/pjrt/gpu/gpu_topology.pb.h"
 #include "xla/status.h"
 #include "xla/statusor.h"
 
@@ -50,6 +51,10 @@ Status ExchangeTopologies(std::string_view platform, int node_id, int num_nodes,
 GlobalTopologyProto BuildGlobalTopology(
     absl::Span<LocalTopologyProto> local_topologies);
 
+// Builds a GpuTopologyProto representing the GPU configuration described in the
+// given GlobalTopologyProto.
+absl::StatusOr<GpuTopologyProto> BuildGpuTopology(
+    const GlobalTopologyProto& global_topology);
 }  // namespace xla
 
 #endif  // XLA_PJRT_DISTRIBUTED_TOPOLOGY_UTIL_H_
diff --git a/third_party/xla/xla/pjrt/distributed/topology_util_test.cc b/third_party/xla/xla/pjrt/distributed/topology_util_test.cc
index f311b398d5309e..e28df2c5e93e2f 100644
--- a/third_party/xla/xla/pjrt/distributed/topology_util_test.cc
+++ b/third_party/xla/xla/pjrt/distributed/topology_util_test.cc
@@ -22,8 +22,10 @@ limitations under the License.
 #include "absl/time/time.h"
 #include "xla/pjrt/distributed/in_memory_key_value_store.h"
 #include "xla/pjrt/distributed/protocol.pb.h"
+#include "xla/test_helpers.h"
 #include "tsl/lib/core/status_test_util.h"
 #include "tsl/platform/env.h"
+#include "tsl/platform/statusor.h"
 #include "tsl/platform/test.h"
 #include "tsl/platform/threadpool.h"
 
@@ -82,5 +84,79 @@ TEST(TopologyTest, ExchangeTopology) {
   }
 }
 
+TEST(TopologyTest, BuildGpuTopology) {
+  std::string slice_0_boot_id = "foo";
+  std::string slice_1_boot_id = "bar";
+  std::vector<LocalTopologyProto> locals(2);
+  // Adds 1 host to slice 0 and 1 host to slice 1.
+  locals[0].set_boot_id(slice_0_boot_id);
+  locals[1].set_boot_id(slice_1_boot_id);
+  locals[0].set_node_id(0);
+  locals[1].set_node_id(1);
+  // Adds 2 devices to host 0 and 2 devices to host 1.
+  DeviceProto* d0 = locals[0].add_devices();
+  d0->set_local_device_ordinal(0);
+  DeviceProto* d1 = locals[0].add_devices();
+  d1->set_local_device_ordinal(1);
+  DeviceProto* d2 = locals[1].add_devices();
+  d2->set_local_device_ordinal(0);
+  DeviceProto* d3 = locals[1].add_devices();
+  d3->set_local_device_ordinal(1);
+
+  GlobalTopologyProto global =
+      BuildGlobalTopology(absl::Span<LocalTopologyProto>(locals));
+
+  TF_ASSERT_OK_AND_ASSIGN(auto gpu_topology, BuildGpuTopology(global));
+  EXPECT_EQ(gpu_topology.device_ids_size(), 4);
+  EXPECT_EQ(gpu_topology.num_slices(), 2);
+  EXPECT_EQ(gpu_topology.num_hosts_per_slice(), 1);
+  EXPECT_EQ(gpu_topology.num_devices_per_host(), 2);
+}
+
+TEST(TopologyTest, BuildGpuTopologyWithDifferentNumHostsPerSlice) {
+  std::string slice_0_boot_id = "foo";
+  std::string slice_1_boot_id = "bar";
+  std::vector<LocalTopologyProto> locals(3);
+  // Adds 2 hosts to slice 0 and 1 host to slice 1.
+  locals[0].set_boot_id(slice_0_boot_id);
+  locals[1].set_boot_id(slice_0_boot_id);
+  locals[2].set_boot_id(slice_1_boot_id);
+  locals[0].set_node_id(0);
+  locals[1].set_node_id(1);
+  locals[2].set_node_id(2);
+  DeviceProto* d0 = locals[0].add_devices();
+  d0->set_local_device_ordinal(0);
+  DeviceProto* d1 = locals[1].add_devices();
+  d1->set_local_device_ordinal(0);
+  DeviceProto* d2 = locals[2].add_devices();
+  d2->set_local_device_ordinal(0);
+
+  GlobalTopologyProto global =
+      BuildGlobalTopology(absl::Span<LocalTopologyProto>(locals));
+
+  EXPECT_IS_NOT_OK(BuildGpuTopology(global).status());
+}
+
+TEST(TopologyTest, BuildGpuTopologyWithDifferentNumDevicesPerHost) {
+  std::string slice_0_boot_id = "foo";
+  std::string slice_1_boot_id = "bar";
+  std::vector<LocalTopologyProto> locals(2);
+  locals[0].set_boot_id(slice_0_boot_id);
+  locals[1].set_boot_id(slice_1_boot_id);
+  locals[0].set_node_id(0);
+  locals[1].set_node_id(1);
+  // Adds 2 devices to host 0 and 1 device to host 1.
+  DeviceProto* d0 = locals[0].add_devices();
+  d0->set_local_device_ordinal(0);
+  DeviceProto* d1 = locals[0].add_devices();
+  d1->set_local_device_ordinal(1);
+  DeviceProto* d2 = locals[1].add_devices();
+  d2->set_local_device_ordinal(0);
+
+  GlobalTopologyProto global =
+      BuildGlobalTopology(absl::Span<LocalTopologyProto>(locals));
+
+  EXPECT_IS_NOT_OK(BuildGpuTopology(global).status());
+}
 }  // namespace
 }  // namespace xla

From 406ea63fd038892c01528674b97eeb2854c0f9cf Mon Sep 17 00:00:00 2001
From: Terry Heo <terryheo@google.com>
Date: Wed, 15 May 2024 14:49:45 -0700
Subject: [PATCH 217/478] gpu_delegate: Allow broadcast of [1, 1] to [1, x, y,
 z]

PiperOrigin-RevId: 634083568
---
 .../tools/versioning/gpu_compatibility.cc     |  7 ++++-
 .../versioning/gpu_compatibility_test.cc      | 30 +++++++++++++++++++
 2 files changed, 36 insertions(+), 1 deletion(-)

diff --git a/tensorflow/lite/tools/versioning/gpu_compatibility.cc b/tensorflow/lite/tools/versioning/gpu_compatibility.cc
index d8e1b12879a85b..c5b50ac728cdad 100644
--- a/tensorflow/lite/tools/versioning/gpu_compatibility.cc
+++ b/tensorflow/lite/tools/versioning/gpu_compatibility.cc
@@ -449,9 +449,14 @@ absl::Status CheckAddMulBroadcastCompatibility(
     }
     bool is_broadcastable = false;
 
-    // Broadcasting 3D to 4D with batch 1 works.
     if (longer_dims->size() == 4 && shorter_dims->size() == 3 &&
         longer_dims->at(0) == 1) {
+      // Broadcasting 3D to 4D with batch 1 works.
+      is_broadcastable = true;
+    } else if (longer_dims->size() == 4 && shorter_dims->size() == 2 &&
+               longer_dims->at(0) == 1 && shorter_dims->at(0) == 1 &&
+               shorter_dims->at(1) == 1) {
+      // Broadcasting 2D [1, 1] to 4D [1, x, y, z] works.
       is_broadcastable = true;
     }
 
diff --git a/tensorflow/lite/tools/versioning/gpu_compatibility_test.cc b/tensorflow/lite/tools/versioning/gpu_compatibility_test.cc
index 8493137c1e0f90..c5283e8e4adfab 100644
--- a/tensorflow/lite/tools/versioning/gpu_compatibility_test.cc
+++ b/tensorflow/lite/tools/versioning/gpu_compatibility_test.cc
@@ -156,4 +156,34 @@ TEST(CheckGpuDelegateCompatibility, Add3Dto4DBroadcastSuccess) {
   EXPECT_TRUE(CheckGpuDelegateCompatibility(op_sig).message().empty());
 }
 
+TEST(CheckGpuDelegateCompatibility, Add2Dto4DBroadcastSuccess) {
+  OpSignature op_sig = OpSignature();
+  op_sig.op = BuiltinOperator_ADD;
+  auto params = std::make_unique<TfLiteAddParams>();
+  op_sig.builtin_data = static_cast<void*>(params.get());
+  op_sig.inputs = std::vector<OpSignatureTensorSpec>(2);
+  op_sig.inputs[0] = OpSignatureTensorSpec();
+  op_sig.inputs[0].dims = {1, 512, 512, 1};
+  op_sig.inputs[1] = OpSignatureTensorSpec();
+  // Can be broadcasted to {1, 1, 1, 1}
+  op_sig.inputs[1].dims = {1, 1};
+
+  EXPECT_TRUE(CheckGpuDelegateCompatibility(op_sig).message().empty());
+}
+
+TEST(CheckGpuDelegateCompatibility, Add2Dto4DBroadcastSuccess2) {
+  OpSignature op_sig = OpSignature();
+  op_sig.op = BuiltinOperator_ADD;
+  auto params = std::make_unique<TfLiteAddParams>();
+  op_sig.builtin_data = static_cast<void*>(params.get());
+  op_sig.inputs = std::vector<OpSignatureTensorSpec>(2);
+  op_sig.inputs[0] = OpSignatureTensorSpec();
+  op_sig.inputs[0].dims = {1, 384, 384, 3};
+  op_sig.inputs[1] = OpSignatureTensorSpec();
+  // Can be broadcasted to {1, 1, 1, 1}
+  op_sig.inputs[1].dims = {1, 1};
+
+  EXPECT_TRUE(CheckGpuDelegateCompatibility(op_sig).message().empty());
+}
+
 }  // namespace tflite

From 61bc508ddc9979def2cadc4c633831c910732e18 Mon Sep 17 00:00:00 2001
From: Kuy Mainwaring <kuym@google.com>
Date: Wed, 15 May 2024 14:49:50 -0700
Subject: [PATCH 218/478] [XLA:GPU] Clang-tidy cleanup for shadowed class
 member in xla/service/gpu/tests/gpu_fused_mha_test.cc

PiperOrigin-RevId: 634083594
---
 third_party/xla/xla/service/gpu/tests/gpu_fused_mha_test.cc | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/third_party/xla/xla/service/gpu/tests/gpu_fused_mha_test.cc b/third_party/xla/xla/service/gpu/tests/gpu_fused_mha_test.cc
index 7c1dc38c20a453..a0cb225ff911ff 100644
--- a/third_party/xla/xla/service/gpu/tests/gpu_fused_mha_test.cc
+++ b/third_party/xla/xla/service/gpu/tests/gpu_fused_mha_test.cc
@@ -83,7 +83,7 @@ class MultiHeadedAttentionTest : public GpuCodegenTest {
         .cuda_compute_capability();
   }
 
-  ErrorSpec error_spec_{2.5E-3, 1e-5};
+  ErrorSpec mha_error_spec_{2.5E-3, 1e-5};
 
  protected:
   DebugOptions GetDebugOptionsForTest() override {
@@ -131,7 +131,7 @@ class MultiHeadedAttentionTest : public GpuCodegenTest {
         ExecuteAndTransfer(std::move(test_module), literals);
 
     EXPECT_TRUE(
-        LiteralTestUtil::Near(expected_result, actual_result, error_spec_));
+        LiteralTestUtil::Near(expected_result, actual_result, mha_error_spec_));
   }
 
   template <typename T>

From 0f752edea8f3d84a3253bc4cfa0a0ad431f80a29 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Wed, 15 May 2024 14:51:31 -0700
Subject: [PATCH 219/478] Updating schema location.

PiperOrigin-RevId: 634084054
---
 .../experimental/acceleration/mini_benchmark/build_defs.bzl   | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/tensorflow/lite/experimental/acceleration/mini_benchmark/build_defs.bzl b/tensorflow/lite/experimental/acceleration/mini_benchmark/build_defs.bzl
index 1c2a1a3f8cb561..6bc40c3c19571b 100644
--- a/tensorflow/lite/experimental/acceleration/mini_benchmark/build_defs.bzl
+++ b/tensorflow/lite/experimental/acceleration/mini_benchmark/build_defs.bzl
@@ -104,7 +104,7 @@ def validation_model(
         srcs = [
             main_model,
             jpegs,
-            "//tensorflow/lite/schema:schema.fbs",
+            "//tensorflow/compiler/mlir/lite/schema:schema.fbs",
             metrics_model,
         ],
         outs = [name + ".tflite"],
@@ -112,7 +112,7 @@ def validation_model(
           JPEGS='$(locations %s)'
           JPEGS=$${JPEGS// /,}
           $(location //tensorflow/lite/experimental/acceleration/mini_benchmark/model_modifier:embedder_cmdline) \
-              --schema=$(location //tensorflow/lite/schema:schema.fbs) \
+              --schema=$(location //tensorflow/compiler/mlir/lite/schema:schema.fbs) \
               --main_model=$(location %s) \
               --metrics_model=$(location %s) \
               %s %s \

From a64922e50d17fa76ddc3db1f62fdbf5134c62dc7 Mon Sep 17 00:00:00 2001
From: Jacques Pienaar <jpienaar@google.com>
Date: Wed, 15 May 2024 14:59:41 -0700
Subject: [PATCH 220/478] Change to add attribute before building

Avoids issues when converting to properties early (which happens along most
generated build paths today and should be all in future).

PiperOrigin-RevId: 634086476
---
 .../transforms/rewrite_tpu_embedding_ops.cc     | 17 +++++++++--------
 1 file changed, 9 insertions(+), 8 deletions(-)

diff --git a/tensorflow/compiler/mlir/tensorflow/transforms/rewrite_tpu_embedding_ops.cc b/tensorflow/compiler/mlir/tensorflow/transforms/rewrite_tpu_embedding_ops.cc
index 2ff6c78896fff2..faedd25114807e 100644
--- a/tensorflow/compiler/mlir/tensorflow/transforms/rewrite_tpu_embedding_ops.cc
+++ b/tensorflow/compiler/mlir/tensorflow/transforms/rewrite_tpu_embedding_ops.cc
@@ -37,12 +37,13 @@ struct RewriteTPUEmbeddingOps
 
 // Rewrites the given op to `OpT` op after adding the given operand at the end.
 template <typename OpT>
-OpT AddOperandAndRewriteAs(Operation* op, Value operand, OpBuilder* builder) {
+OpT AddOperandAndRewriteAs(Operation* op, Value operand, NamedAttrList attr,
+                           OpBuilder* builder) {
   builder->setInsertionPoint(op);
   auto operands = llvm::to_vector<4>(op->getOperands());
   operands.push_back(operand);
   auto new_op = builder->create<OpT>(op->getLoc(), op->getResultTypes(),
-                                     operands, op->getAttrs());
+                                     operands, attr.getAttrs());
   op->replaceAllUsesWith(new_op.getOperation()->getResults());
   op->erase();
   return new_op;
@@ -83,8 +84,8 @@ LogicalResult RunOnRegion(Region* region) {
 
   // Rewrite RecvTPUEmbeddingActivations op to the corresponding internal op.
   if (recv_op)
-    AddOperandAndRewriteAs<XlaRecvTPUEmbeddingActivationsOp>(recv_op, dedup_op,
-                                                             &builder);
+    AddOperandAndRewriteAs<XlaRecvTPUEmbeddingActivationsOp>(
+        recv_op, dedup_op, recv_op->getAttrs(), &builder);
 
   // Rewrite SendTPUEmbeddingGradients op to the corresponding internal op and
   // then update the OperandSegmentSize attribute.
@@ -92,11 +93,11 @@ LogicalResult RunOnRegion(Region* region) {
     int32_t operand_sizes[] = {static_cast<int32_t>(send_op.getN()),
                                static_cast<int32_t>(send_op.getNN()), 1};
     auto operand_size_attr = builder.getDenseI32ArrayAttr(operand_sizes);
+    NamedAttrList attrs(send_op->getAttrs());
+    attrs.set(send_op.getOperandSegmentSizeAttr(), operand_size_attr);
 
-    auto new_send_op = AddOperandAndRewriteAs<XlaSendTPUEmbeddingGradientsOp>(
-        send_op, dedup_op, &builder);
-    new_send_op->setAttr(new_send_op.getOperandSegmentSizeAttr(),
-                         operand_size_attr);
+    AddOperandAndRewriteAs<XlaSendTPUEmbeddingGradientsOp>(send_op, dedup_op,
+                                                           attrs, &builder);
   }
   return success();
 }

From bf1b31c12e65913a0f7717e074161fed843a87cd Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Wed, 15 May 2024 15:17:36 -0700
Subject: [PATCH 221/478] Update ops-related pbtxt files.

PiperOrigin-RevId: 634092262
---
 tensorflow/core/ops/ops.pbtxt | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorflow/core/ops/ops.pbtxt b/tensorflow/core/ops/ops.pbtxt
index 2073f2e00ecd81..a15dada136af38 100644
--- a/tensorflow/core/ops/ops.pbtxt
+++ b/tensorflow/core/ops/ops.pbtxt
@@ -1,4 +1,4 @@
-go/debugonly    
+go/nodeserialize  
 op {
   name: "Abort"
   attr {

From 8799c03211fbeb420e8f07ed957f2ddf8fe07c2d Mon Sep 17 00:00:00 2001
From: Vamsi Manchala <vamsimanchala@google.com>
Date: Wed, 15 May 2024 15:17:55 -0700
Subject: [PATCH 222/478] Add a pattern to lower a stablehlo.composite with
 `jax.image.resize` in `nearest` mode to a tflite.resize_nearest_neighbor op.

PiperOrigin-RevId: 634092342
---
 .../stablehlo/tests/composite-lowering.mlir   | 53 +++++++++++++++++++
 .../transforms/composite_lowering_patterns.td | 15 ++++++
 2 files changed, 68 insertions(+)

diff --git a/tensorflow/compiler/mlir/lite/stablehlo/tests/composite-lowering.mlir b/tensorflow/compiler/mlir/lite/stablehlo/tests/composite-lowering.mlir
index 292e6f5f8ee612..4121caa60a8e0d 100644
--- a/tensorflow/compiler/mlir/lite/stablehlo/tests/composite-lowering.mlir
+++ b/tensorflow/compiler/mlir/lite/stablehlo/tests/composite-lowering.mlir
@@ -187,3 +187,56 @@ func.func @gelu(%arg0: tensor<2xf32>) -> tensor<2xf32> {
 // CHECK-LABEL: gelu
 // CHECK: %0 = "tfl.gelu"(%arg0) <{approximate = false}> : (tensor<2xf32>) -> tensor<2xf32>
 
+// CHECK-LABEL  func.func @jax_image_resize_nearest
+func.func @jax_image_resize_nearest(%arg0: tensor<1x2x2x10xf32>) -> (tensor<1x4x4x10xf32>) {
+  %1 = mhlo.composite "odml.jax_resize_nearest_neighbor2d" %arg0 {composite_attributes = {output_size = dense<4> : tensor<2xi64>}, decomposition = @XlaCallModule_odml.jax_resize_nearest_neighbor2d.impl_0} : (tensor<1x2x2x10xf32>) -> tensor<1x4x4x10xf32>
+  return %1 : tensor<1x4x4x10xf32>
+}
+func.func private @XlaCallModule_odml.jax_resize_nearest_neighbor2d.impl_0(%arg0: tensor<1x2x2x10xf32>) -> tensor<1x4x4x10xf32> {
+  %0 = call @XlaCallModule__resize_0(%arg0) : (tensor<1x2x2x10xf32>) -> tensor<1x4x4x10xf32>
+  return %0 : tensor<1x4x4x10xf32>
+}
+func.func private @XlaCallModule__resize_0(%arg0: tensor<1x2x2x10xf32>) -> (tensor<1x4x4x10xf32>) {
+  %0 = mhlo.constant dense<2> : tensor<i32>
+  %1 = mhlo.constant dense<0> : tensor<i32>
+  %2 = mhlo.constant dense<4.000000e+00> : tensor<f32>
+  %3 = mhlo.constant dense<2.000000e+00> : tensor<f32>
+  %4 = mhlo.constant dense<5.000000e-01> : tensor<f32>
+  %5 = "mhlo.iota"() <{iota_dimension = 0 : i64}> : () -> tensor<4xf32>
+  %6 = "mhlo.broadcast_in_dim"(%4) <{broadcast_dimensions = dense<> : tensor<0xi64>}> : (tensor<f32>) -> tensor<4xf32>
+  %7 = mhlo.add %5, %6 : tensor<4xf32>
+  %8 = "mhlo.broadcast_in_dim"(%3) <{broadcast_dimensions = dense<> : tensor<0xi64>}> : (tensor<f32>) -> tensor<4xf32>
+  %9 = mhlo.multiply %7, %8 : tensor<4xf32>
+  %10 = "mhlo.broadcast_in_dim"(%2) <{broadcast_dimensions = dense<> : tensor<0xi64>}> : (tensor<f32>) -> tensor<4xf32>
+  %11 = mhlo.divide %9, %10 : tensor<4xf32>
+  %12 = mhlo.floor %11 : tensor<4xf32>
+  %13 = mhlo.convert %12 : (tensor<4xf32>) -> tensor<4xi32>
+  %14 = "mhlo.broadcast_in_dim"(%1) <{broadcast_dimensions = dense<> : tensor<0xi64>}> : (tensor<i32>) -> tensor<4xi32>
+  %15 = mhlo.compare  LT, %13, %14,  SIGNED : (tensor<4xi32>, tensor<4xi32>) -> tensor<4xi1>
+  %16 = "mhlo.broadcast_in_dim"(%0) <{broadcast_dimensions = dense<> : tensor<0xi64>}> : (tensor<i32>) -> tensor<4xi32>
+  %17 = mhlo.add %13, %16 : tensor<4xi32>
+  %18 = mhlo.select %15, %17, %13 : tensor<4xi1>, tensor<4xi32>
+  %19 = "mhlo.broadcast_in_dim"(%18) <{broadcast_dimensions = dense<0> : tensor<1xi64>}> : (tensor<4xi32>) -> tensor<4x1xi32>
+  %20 = "mhlo.gather"(%arg0, %19) <{dimension_numbers = #mhlo.gather<offset_dims = [0, 2, 3], collapsed_slice_dims = [1], start_index_map = [1], index_vector_dim = 1>, slice_sizes = dense<[1, 1, 2, 10]> : tensor<4xi64>}> : (tensor<1x2x2x10xf32>, tensor<4x1xi32>) -> tensor<1x4x2x10xf32>
+  %21 = "mhlo.iota"() <{iota_dimension = 0 : i64}> : () -> tensor<4xf32>
+  %22 = "mhlo.broadcast_in_dim"(%4) <{broadcast_dimensions = dense<> : tensor<0xi64>}> : (tensor<f32>) -> tensor<4xf32>
+  %23 = mhlo.add %21, %22 : tensor<4xf32>
+  %24 = "mhlo.broadcast_in_dim"(%3) <{broadcast_dimensions = dense<> : tensor<0xi64>}> : (tensor<f32>) -> tensor<4xf32>
+  %25 = mhlo.multiply %23, %24 : tensor<4xf32>
+  %26 = "mhlo.broadcast_in_dim"(%2) <{broadcast_dimensions = dense<> : tensor<0xi64>}> : (tensor<f32>) -> tensor<4xf32>
+  %27 = mhlo.divide %25, %26 : tensor<4xf32>
+  %28 = mhlo.floor %27 : tensor<4xf32>
+  %29 = mhlo.convert %28 : (tensor<4xf32>) -> tensor<4xi32>
+  %30 = "mhlo.broadcast_in_dim"(%1) <{broadcast_dimensions = dense<> : tensor<0xi64>}> : (tensor<i32>) -> tensor<4xi32>
+  %31 = mhlo.compare  LT, %29, %30,  SIGNED : (tensor<4xi32>, tensor<4xi32>) -> tensor<4xi1>
+  %32 = "mhlo.broadcast_in_dim"(%0) <{broadcast_dimensions = dense<> : tensor<0xi64>}> : (tensor<i32>) -> tensor<4xi32>
+  %33 = mhlo.add %29, %32 : tensor<4xi32>
+  %34 = mhlo.select %31, %33, %29 : tensor<4xi1>, tensor<4xi32>
+  %35 = "mhlo.broadcast_in_dim"(%34) <{broadcast_dimensions = dense<0> : tensor<1xi64>}> : (tensor<4xi32>) -> tensor<4x1xi32>
+  %36 = "mhlo.gather"(%20, %35) <{dimension_numbers = #mhlo.gather<offset_dims = [0, 1, 3], collapsed_slice_dims = [2], start_index_map = [2], index_vector_dim = 1>, slice_sizes = dense<[1, 4, 1, 10]> : tensor<4xi64>}> : (tensor<1x4x2x10xf32>, tensor<4x1xi32>) -> tensor<1x4x4x10xf32>
+  return %36 : tensor<1x4x4x10xf32>
+}
+
+// CHECK:  %cst = arith.constant dense<4> : tensor<2xi32>
+// CHECK:  %0 = "tfl.resize_nearest_neighbor"(%arg0, %cst) <{align_corners = false, half_pixel_centers = true}> : (tensor<1x2x2x10xf32>, tensor<2xi32>) -> tensor<1x4x4x10xf32>
+// CHECK:  return %0 : tensor<1x4x4x10xf32>
diff --git a/tensorflow/compiler/mlir/lite/stablehlo/transforms/composite_lowering_patterns.td b/tensorflow/compiler/mlir/lite/stablehlo/transforms/composite_lowering_patterns.td
index 89327e492d72aa..5b9324c2a1782b 100644
--- a/tensorflow/compiler/mlir/lite/stablehlo/transforms/composite_lowering_patterns.td
+++ b/tensorflow/compiler/mlir/lite/stablehlo/transforms/composite_lowering_patterns.td
@@ -50,6 +50,21 @@ def LegalizeTorchUpsampleBlinear2dComposite: Pat<
                             ConstantAttr<RankedI32ElementsAttr<[4]>,"{0, 3, 1, 2}">)),
                     [(IsSupportedNchwUpsampleBlinear $input, $old_val, $attrs)]>;
 
+// TODO(b/333961789): Add support for NCHW layout for PyTorch resize, plus jax
+// supports NCHW inputs as well, so we need to add reliable way of checking the
+// layout.
+// pattern to lower a stablehlo.composite with `jax.image.resize` in `nearest`
+// mode to a tflite.resize_nearest_neighbor op.
+def LegalizeJaxResizeNearestNeighbor2dComposite: Pat<
+    (MHLO_CompositeOp:$old_val
+    (variadic $input),
+    ConstantStrAttr<StrAttr, "odml.jax_resize_nearest_neighbor2d">, $attrs, $_, $_),
+    (TFL_ResizeNearestNeighborOp
+      $input,
+      (Arith_ConstantOp:$output_size (GetI32DenseAttr (GetAsVectorAttr<"output_size"> $attrs))),
+      ConstBoolAttrFalse,
+      ConstBoolAttrTrue)>;
+
 
 def LegalizeCompositeGELU : Pat<(MHLO_CompositeOp:$composite 
                                   (variadic $inputs), 

From 3bf2ac3b776c05471f1b4f14f50dddfbcc79b38b Mon Sep 17 00:00:00 2001
From: Vadym Matsishevskyi <vam@google.com>
Date: Wed, 15 May 2024 15:25:35 -0700
Subject: [PATCH 223/478] Use hermetic Python in TSL and XLA

PiperOrigin-RevId: 634094641
---
 tensorflow/opensource_only.files              |   9 +-
 third_party/llvm/toolchains.patch             |   8 +-
 third_party/py/BUILD                          |  40 +++
 third_party/py/non_hermetic/BUILD             |   0
 third_party/py/non_hermetic/BUILD.tpl         |  80 -----
 third_party/py/non_hermetic/README            |   3 -
 .../py/non_hermetic/python_configure.bzl      | 312 ------------------
 third_party/py/python_init_pip.bzl            |  34 ++
 third_party/py/python_init_repositories.bzl   |  12 +
 third_party/py/python_init_rules.bzl          |  11 +
 third_party/py/python_init_toolchains.bzl     |  13 +
 third_party/py/python_repo.bzl                | 206 ++++++++++++
 third_party/xla/WORKSPACE                     |  26 ++
 third_party/xla/opensource_only.files         |  18 +-
 third_party/xla/requirements_lock_3_11.txt    |  49 +++
 third_party/xla/third_party/py/BUILD          |  40 +++
 .../xla/third_party/py/non_hermetic/BUILD     |   0
 .../xla/third_party/py/non_hermetic/BUILD.tpl |  80 -----
 .../xla/third_party/py/non_hermetic/README    |   3 -
 .../py/non_hermetic/python_configure.bzl      | 312 ------------------
 .../xla/third_party/py/python_init_pip.bzl    |  34 ++
 .../py/python_init_repositories.bzl           |  12 +
 .../xla/third_party/py/python_init_rules.bzl  |  11 +
 .../third_party/py/python_init_toolchains.bzl |  13 +
 .../xla/third_party/py/python_repo.bzl        | 206 ++++++++++++
 third_party/xla/third_party/tsl/WORKSPACE     |  25 ++
 .../xla/third_party/tsl/opensource_only.files |  18 +-
 .../tsl/requirements_lock_3_11.txt            |  49 +++
 .../xla/third_party/tsl/third_party/py/BUILD  |  40 +++
 .../tsl/third_party/py/non_hermetic/BUILD     |   0
 .../tsl/third_party/py/non_hermetic/BUILD.tpl |  80 -----
 .../tsl/third_party/py/non_hermetic/README    |   3 -
 .../py/non_hermetic/python_configure.bzl      | 312 ------------------
 .../tsl/third_party/py/python_init_pip.bzl    |  34 ++
 .../py/python_init_repositories.bzl           |  12 +
 .../tsl/third_party/py/python_init_rules.bzl  |  11 +
 .../third_party/py/python_init_toolchains.bzl |  13 +
 .../tsl/third_party/py/python_repo.bzl        | 206 ++++++++++++
 .../xla/third_party/tsl/workspace2.bzl        |   7 +-
 third_party/xla/xla/lit.bzl                   |  21 +-
 third_party/xla/xla/mlir_hlo/tests/BUILD      |   1 +
 41 files changed, 1151 insertions(+), 1213 deletions(-)
 delete mode 100644 third_party/py/non_hermetic/BUILD
 delete mode 100644 third_party/py/non_hermetic/BUILD.tpl
 delete mode 100644 third_party/py/non_hermetic/README
 delete mode 100644 third_party/py/non_hermetic/python_configure.bzl
 create mode 100644 third_party/py/python_init_pip.bzl
 create mode 100644 third_party/py/python_init_repositories.bzl
 create mode 100644 third_party/py/python_init_rules.bzl
 create mode 100644 third_party/py/python_init_toolchains.bzl
 create mode 100644 third_party/py/python_repo.bzl
 create mode 100644 third_party/xla/requirements_lock_3_11.txt
 delete mode 100644 third_party/xla/third_party/py/non_hermetic/BUILD
 delete mode 100644 third_party/xla/third_party/py/non_hermetic/BUILD.tpl
 delete mode 100644 third_party/xla/third_party/py/non_hermetic/README
 delete mode 100644 third_party/xla/third_party/py/non_hermetic/python_configure.bzl
 create mode 100644 third_party/xla/third_party/py/python_init_pip.bzl
 create mode 100644 third_party/xla/third_party/py/python_init_repositories.bzl
 create mode 100644 third_party/xla/third_party/py/python_init_rules.bzl
 create mode 100644 third_party/xla/third_party/py/python_init_toolchains.bzl
 create mode 100644 third_party/xla/third_party/py/python_repo.bzl
 create mode 100644 third_party/xla/third_party/tsl/requirements_lock_3_11.txt
 delete mode 100644 third_party/xla/third_party/tsl/third_party/py/non_hermetic/BUILD
 delete mode 100644 third_party/xla/third_party/tsl/third_party/py/non_hermetic/BUILD.tpl
 delete mode 100644 third_party/xla/third_party/tsl/third_party/py/non_hermetic/README
 delete mode 100644 third_party/xla/third_party/tsl/third_party/py/non_hermetic/python_configure.bzl
 create mode 100644 third_party/xla/third_party/tsl/third_party/py/python_init_pip.bzl
 create mode 100644 third_party/xla/third_party/tsl/third_party/py/python_init_repositories.bzl
 create mode 100644 third_party/xla/third_party/tsl/third_party/py/python_init_rules.bzl
 create mode 100644 third_party/xla/third_party/tsl/third_party/py/python_init_toolchains.bzl
 create mode 100644 third_party/xla/third_party/tsl/third_party/py/python_repo.bzl

diff --git a/tensorflow/opensource_only.files b/tensorflow/opensource_only.files
index b83a48134d02e1..49adcabcd51766 100644
--- a/tensorflow/opensource_only.files
+++ b/tensorflow/opensource_only.files
@@ -300,15 +300,16 @@ tf_staging/third_party/py/BUILD.tpl:
 tf_staging/third_party/py/BUILD:
 tf_staging/third_party/py/ml_dtypes/BUILD:
 tf_staging/third_party/py/ml_dtypes/LICENSE:
-tf_staging/third_party/py/non_hermetic/BUILD.tpl:
-tf_staging/third_party/py/non_hermetic/BUILD:
-tf_staging/third_party/py/non_hermetic/README:
 tf_staging/third_party/py/non_hermetic/ml_dtypes/BUILD:
 tf_staging/third_party/py/non_hermetic/ml_dtypes/LICENSE:
 tf_staging/third_party/py/non_hermetic/numpy/BUILD:
-tf_staging/third_party/py/non_hermetic/python_configure.bzl:
 tf_staging/third_party/py/numpy/BUILD:
 tf_staging/third_party/py/python_configure.bzl:
+tf_staging/third_party/py/python_init_pip.bzl:
+tf_staging/third_party/py/python_init_repositories.bzl:
+tf_staging/third_party/py/python_init_rules.bzl:
+tf_staging/third_party/py/python_init_toolchains.bzl:
+tf_staging/third_party/py/python_repo.bzl:
 tf_staging/third_party/pybind11.BUILD:
 tf_staging/third_party/pybind11_bazel/BUILD:
 tf_staging/third_party/pybind11_protobuf/BUILD:
diff --git a/third_party/llvm/toolchains.patch b/third_party/llvm/toolchains.patch
index a4de4eaaff343a..1cbcfe31d3d072 100644
--- a/third_party/llvm/toolchains.patch
+++ b/third_party/llvm/toolchains.patch
@@ -2,7 +2,7 @@ diff --git a/utils/bazel/llvm-project-overlay/llvm/BUILD.bazel b/utils/bazel/llv
 index c43ab727e285..7d848d2dffae 100644
 --- a/utils/bazel/llvm-project-overlay/llvm/BUILD.bazel
 +++ b/utils/bazel/llvm-project-overlay/llvm/BUILD.bazel
-@@ -30,6 +30,30 @@ exports_files([
+@@ -30,6 +30,36 @@ exports_files([
      "utils/lit/lit.py",
  ])
  
@@ -22,6 +22,12 @@ index c43ab727e285..7d848d2dffae 100644
 +    },
 +)
 +
++py_library(
++    name = "lit_lib",
++    testonly = True,
++    srcs = ["utils/lit/lit.py"] + glob(["utils/lit/lit/**/*.py"]),
++)
++
 +config_setting(
 +    name = "macos_x86_64",
 +    values = {
diff --git a/third_party/py/BUILD b/third_party/py/BUILD
index e69de29bb2d1d6..84eba77ce1a7af 100644
--- a/third_party/py/BUILD
+++ b/third_party/py/BUILD
@@ -0,0 +1,40 @@
+load("@python//:defs.bzl", "compile_pip_requirements")
+load("@python_version_repo//:py_version.bzl", "REQUIREMENTS")
+
+compile_pip_requirements(
+    name = "requirements",
+    extra_args = [
+        "--allow-unsafe",
+        "--build-isolation",
+    ],
+    generate_hashes = True,
+    requirements_in = "requirements.in",
+    requirements_txt = REQUIREMENTS,
+)
+
+compile_pip_requirements(
+    name = "requirements_nightly",
+    data = ["test-requirements.txt"],
+    extra_args = [
+        "--allow-unsafe",
+        "--build-isolation",
+        "--extra-index-url=https://pypi.anaconda.org/scientific-python-nightly-wheels/simple",
+        "--pre",
+        "--upgrade",
+    ],
+    generate_hashes = False,
+    requirements_in = "requirements.in",
+    requirements_txt = REQUIREMENTS,
+)
+
+compile_pip_requirements(
+    name = "requirements_dev",
+    extra_args = [
+        "--allow-unsafe",
+        "--build-isolation",
+        "--upgrade",
+    ],
+    generate_hashes = False,
+    requirements_in = "requirements.in",
+    requirements_txt = REQUIREMENTS,
+)
diff --git a/third_party/py/non_hermetic/BUILD b/third_party/py/non_hermetic/BUILD
deleted file mode 100644
index e69de29bb2d1d6..00000000000000
diff --git a/third_party/py/non_hermetic/BUILD.tpl b/third_party/py/non_hermetic/BUILD.tpl
deleted file mode 100644
index 45480bd4a31cf8..00000000000000
--- a/third_party/py/non_hermetic/BUILD.tpl
+++ /dev/null
@@ -1,80 +0,0 @@
-licenses(["restricted"])
-
-package(default_visibility = ["//visibility:public"])
-
-# Point both runtimes to the same python binary to ensure we always
-# use the python binary specified by ./configure.py script.
-load("@bazel_tools//tools/python:toolchain.bzl", "py_runtime_pair")
-
-py_runtime(
-    name = "py2_runtime",
-    interpreter_path = "%{PYTHON_BIN_PATH}",
-    python_version = "PY2",
-)
-
-py_runtime(
-    name = "py3_runtime",
-    interpreter_path = "%{PYTHON_BIN_PATH}",
-    python_version = "PY3",
-)
-
-py_runtime_pair(
-    name = "py_runtime_pair",
-    py2_runtime = ":py2_runtime",
-    py3_runtime = ":py3_runtime",
-)
-
-toolchain(
-    name = "py_toolchain",
-    toolchain = ":py_runtime_pair",
-    toolchain_type = "@bazel_tools//tools/python:toolchain_type",
-    target_compatible_with = [%{PLATFORM_CONSTRAINT}],
-    exec_compatible_with = [%{PLATFORM_CONSTRAINT}],
-)
-
-# To build Python C/C++ extension on Windows, we need to link to python import library pythonXY.lib
-# See https://docs.python.org/3/extending/windows.html
-cc_import(
-    name = "python_lib",
-    interface_library = select({
-        ":windows": ":python_import_lib",
-        # A placeholder for Unix platforms which makes --no_build happy.
-        "//conditions:default": "not-existing.lib",
-    }),
-    system_provided = 1,
-)
-
-cc_library(
-    name = "python_headers",
-    hdrs = [":python_include"],
-    deps = select({
-        ":windows": [":python_lib"],
-        "//conditions:default": [],
-    }),
-    includes = ["python_include"],
-)
-
-# This alias is exists for the use of targets in the @llvm-project dependency,
-# which expect a python_headers target called @python_runtime//:headers. We use
-# a repo_mapping to alias python_runtime to this package, and an alias to create
-# the correct target.
-alias(
-    name = "headers",
-    actual = ":python_headers",
-)
-
-cc_library(
-    name = "numpy_headers",
-    hdrs = [":numpy_include"],
-    includes = ["numpy_include"],
-)
-
-config_setting(
-    name = "windows",
-    values = {"cpu": "x64_windows"},
-    visibility = ["//visibility:public"],
-)
-
-%{PYTHON_INCLUDE_GENRULE}
-%{NUMPY_INCLUDE_GENRULE}
-%{PYTHON_IMPORT_LIB_GENRULE}
\ No newline at end of file
diff --git a/third_party/py/non_hermetic/README b/third_party/py/non_hermetic/README
deleted file mode 100644
index 62188a5817a09e..00000000000000
--- a/third_party/py/non_hermetic/README
+++ /dev/null
@@ -1,3 +0,0 @@
-This is a temporary copy of python_configure repository rule. It allows
-XLA and TSL to keep non-hermetic Python while TF uses hermetic Python.
-DO NOT DEPEND ON THIS COPY as it will be deleted soon.
\ No newline at end of file
diff --git a/third_party/py/non_hermetic/python_configure.bzl b/third_party/py/non_hermetic/python_configure.bzl
deleted file mode 100644
index 89732c3e33d8ee..00000000000000
--- a/third_party/py/non_hermetic/python_configure.bzl
+++ /dev/null
@@ -1,312 +0,0 @@
-"""Repository rule for Python autoconfiguration.
-
-`python_configure` depends on the following environment variables:
-
-  * `PYTHON_BIN_PATH`: location of python binary.
-  * `PYTHON_LIB_PATH`: Location of python libraries.
-"""
-
-load(
-    "//third_party/remote_config:common.bzl",
-    "BAZEL_SH",
-    "PYTHON_BIN_PATH",
-    "PYTHON_LIB_PATH",
-    "TF_PYTHON_CONFIG_REPO",
-    "auto_config_fail",
-    "config_repo_label",
-    "execute",
-    "get_bash_bin",
-    "get_host_environ",
-    "get_python_bin",
-    "is_windows",
-    "raw_exec",
-    "read_dir",
-)
-
-def _genrule(src_dir, genrule_name, command, outs):
-    """Returns a string with a genrule.
-
-    Genrule executes the given command and produces the given outputs.
-    """
-    return (
-        "genrule(\n" +
-        '    name = "' +
-        genrule_name + '",\n' +
-        "    outs = [\n" +
-        outs +
-        "\n    ],\n" +
-        '    cmd = """\n' +
-        command +
-        '\n   """,\n' +
-        ")\n"
-    )
-
-def _norm_path(path):
-    """Returns a path with '/' and remove the trailing slash."""
-    path = path.replace("\\", "/")
-    if path[-1] == "/":
-        path = path[:-1]
-    return path
-
-def _symlink_genrule_for_dir(
-        repository_ctx,
-        src_dir,
-        dest_dir,
-        genrule_name,
-        src_files = [],
-        dest_files = []):
-    """Returns a genrule to symlink(or copy if on Windows) a set of files.
-
-    If src_dir is passed, files will be read from the given directory; otherwise
-    we assume files are in src_files and dest_files
-    """
-    if src_dir != None:
-        src_dir = _norm_path(src_dir)
-        dest_dir = _norm_path(dest_dir)
-        files = "\n".join(read_dir(repository_ctx, src_dir))
-
-        # Create a list with the src_dir stripped to use for outputs.
-        dest_files = files.replace(src_dir, "").splitlines()
-        src_files = files.splitlines()
-    command = []
-    outs = []
-    for i in range(len(dest_files)):
-        if dest_files[i] != "":
-            # If we have only one file to link we do not want to use the dest_dir, as
-            # $(@D) will include the full path to the file.
-            dest = "$(@D)/" + dest_dir + dest_files[i] if len(dest_files) != 1 else "$(@D)/" + dest_files[i]
-
-            # Copy the headers to create a sandboxable setup.
-            cmd = "cp -f"
-            command.append(cmd + ' "%s" "%s"' % (src_files[i], dest))
-            outs.append('        "' + dest_dir + dest_files[i] + '",')
-    genrule = _genrule(
-        src_dir,
-        genrule_name,
-        " && ".join(command),
-        "\n".join(outs),
-    )
-    return genrule
-
-def _get_python_lib(repository_ctx, python_bin):
-    """Gets the python lib path."""
-    python_lib = get_host_environ(repository_ctx, PYTHON_LIB_PATH)
-    if python_lib != None:
-        return python_lib
-
-    # The interesting program to execute.
-    print_lib = [
-        "from __future__ import print_function",
-        "import site",
-        "import os",
-        "python_paths = []",
-        "if os.getenv('PYTHONPATH') is not None:",
-        "  python_paths = os.getenv('PYTHONPATH').split(':')",
-        "try:",
-        "  library_paths = site.getsitepackages()",
-        "except AttributeError:",
-        "  from distutils.sysconfig import get_python_lib",
-        "  library_paths = [get_python_lib()]",
-        "all_paths = set(python_paths + library_paths)",
-        "paths = []",
-        "for path in all_paths:",
-        "  if os.path.isdir(path):",
-        "    paths.append(path)",
-        "if len(paths) >=1:",
-        "  print(paths[0])",
-    ]
-
-    # The below script writes the above program to a file
-    # and executes it. This is to work around the limitation
-    # of not being able to upload files as part of execute.
-    cmd = "from os import linesep;"
-    cmd += "f = open('script.py', 'w');"
-    for line in print_lib:
-        cmd += "f.write(\"%s\" + linesep);" % line
-    cmd += "f.close();"
-    cmd += "from subprocess import call;"
-    cmd += "call([\"%s\", \"script.py\"]);" % python_bin
-
-    result = execute(repository_ctx, [python_bin, "-c", cmd])
-    return result.stdout.strip()
-
-def _check_python_lib(repository_ctx, python_lib):
-    """Checks the python lib path."""
-    cmd = 'test -d "%s" -a -x "%s"' % (python_lib, python_lib)
-    result = raw_exec(repository_ctx, [get_bash_bin(repository_ctx), "-c", cmd])
-    if result.return_code == 1:
-        auto_config_fail("Invalid python library path: %s" % python_lib)
-
-def _check_python_bin(repository_ctx, python_bin):
-    """Checks the python bin path."""
-    cmd = '[[ -x "%s" ]] && [[ ! -d "%s" ]]' % (python_bin, python_bin)
-    result = raw_exec(repository_ctx, [get_bash_bin(repository_ctx), "-c", cmd])
-    if result.return_code == 1:
-        auto_config_fail("--define %s='%s' is not executable. Is it the python binary?" % (
-            PYTHON_BIN_PATH,
-            python_bin,
-        ))
-
-def _get_python_include(repository_ctx, python_bin):
-    """Gets the python include path."""
-    result = execute(
-        repository_ctx,
-        [
-            python_bin,
-            "-Wignore",
-            "-c",
-            "import sysconfig; " +
-            "print(sysconfig.get_path('include'))",
-        ],
-        error_msg = "Problem getting python include path.",
-        error_details = ("Is the Python binary path set up right? " +
-                         "(See ./configure or " + PYTHON_BIN_PATH + ".) " +
-                         "Is distutils installed?"),
-    )
-    return result.stdout.splitlines()[0]
-
-def _get_python_import_lib_name(repository_ctx, python_bin):
-    """Get Python import library name (pythonXY.lib) on Windows."""
-    result = execute(
-        repository_ctx,
-        [
-            python_bin,
-            "-c",
-            "import sys;" +
-            'print("python" + str(sys.version_info[0]) + ' +
-            '      str(sys.version_info[1]) + ".lib")',
-        ],
-        error_msg = "Problem getting python import library.",
-        error_details = ("Is the Python binary path set up right? " +
-                         "(See ./configure or " + PYTHON_BIN_PATH + ".) "),
-    )
-    return result.stdout.splitlines()[0]
-
-def _get_numpy_include(repository_ctx, python_bin):
-    """Gets the numpy include path."""
-    return execute(
-        repository_ctx,
-        [
-            python_bin,
-            "-c",
-            "from __future__ import print_function;" +
-            "import numpy;" +
-            " print(numpy.get_include());",
-        ],
-        error_msg = "Problem getting numpy include path.",
-        error_details = "Is numpy installed?",
-    ).stdout.splitlines()[0]
-
-def _create_local_python_repository(repository_ctx):
-    """Creates the repository containing files set up to build with Python."""
-
-    # Resolve all labels before doing any real work. Resolving causes the
-    # function to be restarted with all previous state being lost. This
-    # can easily lead to a O(n^2) runtime in the number of labels.
-    build_tpl = repository_ctx.path(Label("//third_party/py:BUILD.tpl"))
-
-    python_bin = get_python_bin(repository_ctx)
-    _check_python_bin(repository_ctx, python_bin)
-    python_lib = _get_python_lib(repository_ctx, python_bin)
-    _check_python_lib(repository_ctx, python_lib)
-    python_include = _get_python_include(repository_ctx, python_bin)
-    numpy_include = _get_numpy_include(repository_ctx, python_bin) + "/numpy"
-    python_include_rule = _symlink_genrule_for_dir(
-        repository_ctx,
-        python_include,
-        "python_include",
-        "python_include",
-    )
-    python_import_lib_genrule = ""
-
-    # To build Python C/C++ extension on Windows, we need to link to python import library pythonXY.lib
-    # See https://docs.python.org/3/extending/windows.html
-    if is_windows(repository_ctx):
-        python_bin = python_bin.replace("\\", "/")
-        python_include = _norm_path(python_include)
-        python_import_lib_name = _get_python_import_lib_name(repository_ctx, python_bin)
-        python_import_lib_src = python_include.rsplit("/", 1)[0] + "/libs/" + python_import_lib_name
-        python_import_lib_genrule = _symlink_genrule_for_dir(
-            repository_ctx,
-            None,
-            "",
-            "python_import_lib",
-            [python_import_lib_src],
-            [python_import_lib_name],
-        )
-    numpy_include_rule = _symlink_genrule_for_dir(
-        repository_ctx,
-        numpy_include,
-        "numpy_include/numpy",
-        "numpy_include",
-    )
-
-    platform_constraint = ""
-    if repository_ctx.attr.platform_constraint:
-        platform_constraint = "\"%s\"" % repository_ctx.attr.platform_constraint
-    repository_ctx.template("BUILD", build_tpl, {
-        "%{PYTHON_BIN_PATH}": python_bin,
-        "%{PYTHON_INCLUDE_GENRULE}": python_include_rule,
-        "%{PYTHON_IMPORT_LIB_GENRULE}": python_import_lib_genrule,
-        "%{NUMPY_INCLUDE_GENRULE}": numpy_include_rule,
-        "%{PLATFORM_CONSTRAINT}": platform_constraint,
-    })
-
-def _create_remote_python_repository(repository_ctx, remote_config_repo):
-    """Creates pointers to a remotely configured repo set up to build with Python.
-    """
-    repository_ctx.template("BUILD", config_repo_label(remote_config_repo, ":BUILD"), {})
-
-def _python_autoconf_impl(repository_ctx):
-    """Implementation of the python_autoconf repository rule."""
-    if get_host_environ(repository_ctx, TF_PYTHON_CONFIG_REPO) != None:
-        _create_remote_python_repository(
-            repository_ctx,
-            get_host_environ(repository_ctx, TF_PYTHON_CONFIG_REPO),
-        )
-    else:
-        _create_local_python_repository(repository_ctx)
-
-_ENVIRONS = [
-    BAZEL_SH,
-    PYTHON_BIN_PATH,
-    PYTHON_LIB_PATH,
-]
-
-local_python_configure = repository_rule(
-    implementation = _create_local_python_repository,
-    environ = _ENVIRONS,
-    attrs = {
-        "environ": attr.string_dict(),
-        "platform_constraint": attr.string(),
-    },
-)
-
-remote_python_configure = repository_rule(
-    implementation = _create_local_python_repository,
-    environ = _ENVIRONS,
-    remotable = True,
-    attrs = {
-        "environ": attr.string_dict(),
-        "platform_constraint": attr.string(),
-    },
-)
-
-python_configure = repository_rule(
-    implementation = _python_autoconf_impl,
-    environ = _ENVIRONS + [TF_PYTHON_CONFIG_REPO],
-    attrs = {
-        "platform_constraint": attr.string(),
-    },
-)
-"""Detects and configures the local Python.
-
-Add the following to your WORKSPACE FILE:
-
-```python
-python_configure(name = "local_config_python")
-```
-
-Args:
-  name: A unique name for this workspace rule.
-"""
diff --git a/third_party/py/python_init_pip.bzl b/third_party/py/python_init_pip.bzl
new file mode 100644
index 00000000000000..efc2bf8233cf61
--- /dev/null
+++ b/third_party/py/python_init_pip.bzl
@@ -0,0 +1,34 @@
+"""Hermetic Python initialization. Consult the WORKSPACE on how to use it."""
+
+load("@python//:defs.bzl", "interpreter")
+load("@python_version_repo//:py_version.bzl", "REQUIREMENTS")
+load("@rules_python//python:pip.bzl", "package_annotation", "pip_parse")
+
+def python_init_pip():
+    numpy_annotations = {
+        "numpy": package_annotation(
+            additive_build_content = """\
+cc_library(
+    name = "numpy_headers_2",
+    hdrs = glob(["site-packages/numpy/_core/include/**/*.h"]),
+    strip_include_prefix="site-packages/numpy/_core/include/",
+)
+cc_library(
+    name = "numpy_headers_1",
+    hdrs = glob(["site-packages/numpy/core/include/**/*.h"]),
+    strip_include_prefix="site-packages/numpy/core/include/",
+)
+cc_library(
+    name = "numpy_headers",
+    deps = [":numpy_headers_2", ":numpy_headers_1"],
+)
+""",
+        ),
+    }
+
+    pip_parse(
+        name = "pypi",
+        annotations = numpy_annotations,
+        python_interpreter_target = interpreter,
+        requirements_lock = REQUIREMENTS,
+    )
diff --git a/third_party/py/python_init_repositories.bzl b/third_party/py/python_init_repositories.bzl
new file mode 100644
index 00000000000000..5a405f2c2aba4c
--- /dev/null
+++ b/third_party/py/python_init_repositories.bzl
@@ -0,0 +1,12 @@
+"""Hermetic Python initialization. Consult the WORKSPACE on how to use it."""
+
+load("@rules_python//python:repositories.bzl", "py_repositories")
+load("//third_party/py:python_repo.bzl", "python_repository")
+
+def python_init_repositories(requirements = {}):
+    python_repository(
+        name = "python_version_repo",
+        requirements_versions = requirements.keys(),
+        requirements_locks = requirements.values(),
+    )
+    py_repositories()
diff --git a/third_party/py/python_init_rules.bzl b/third_party/py/python_init_rules.bzl
new file mode 100644
index 00000000000000..98a7b8bc3c315a
--- /dev/null
+++ b/third_party/py/python_init_rules.bzl
@@ -0,0 +1,11 @@
+"""Hermetic Python initialization. Consult the WORKSPACE on how to use it."""
+
+load("@bazel_tools//tools/build_defs/repo:http.bzl", "http_archive")
+
+def python_init_rules():
+    http_archive(
+        name = "rules_python",
+        sha256 = "9d04041ac92a0985e344235f5d946f71ac543f1b1565f2cdbc9a2aaee8adf55b",
+        strip_prefix = "rules_python-0.26.0",
+        url = "https://github.com/bazelbuild/rules_python/releases/download/0.26.0/rules_python-0.26.0.tar.gz",
+    )
diff --git a/third_party/py/python_init_toolchains.bzl b/third_party/py/python_init_toolchains.bzl
new file mode 100644
index 00000000000000..c1f800db4c01e7
--- /dev/null
+++ b/third_party/py/python_init_toolchains.bzl
@@ -0,0 +1,13 @@
+"""Hermetic Python initialization. Consult the WORKSPACE on how to use it."""
+
+load("@python_version_repo//:py_version.bzl", "HERMETIC_PYTHON_VERSION")
+load("@rules_python//python:repositories.bzl", "python_register_toolchains")
+load("@rules_python//python:versions.bzl", "MINOR_MAPPING")
+
+def python_init_toolchains():
+    if HERMETIC_PYTHON_VERSION in MINOR_MAPPING:
+        python_register_toolchains(
+            name = "python",
+            ignore_root_user_error = True,
+            python_version = HERMETIC_PYTHON_VERSION,
+        )
diff --git a/third_party/py/python_repo.bzl b/third_party/py/python_repo.bzl
new file mode 100644
index 00000000000000..77a6ce9ce50b60
--- /dev/null
+++ b/third_party/py/python_repo.bzl
@@ -0,0 +1,206 @@
+"""
+Repository rule to manage hermetic Python interpreter under Bazel.
+
+Version can be set via build parameter "--repo_env=HERMETIC_PYTHON_VERSION=3.11"
+Defaults to 3.11.
+
+To set wheel name, add "--repo_env=WHEEL_NAME=tensorflow_cpu"
+"""
+
+VERSIONS = ["3.9", "3.10", "3.11", "3.12", "3.13"]
+DEFAULT_VERSION = "3.11"
+WARNING = """
+HERMETIC_PYTHON_VERSION variable was not set correctly, using default version. 
+Python {} will be used.
+To select Python version, either set HERMETIC_PYTHON_VERSION env variable in 
+your shell:
+  export HERMETIC_PYTHON_VERSION=3.12
+OR pass it as an argument to bazel command directly or inside your .bazelrc 
+file:
+  --repo_env=HERMETIC_PYTHON_VERSION=3.12
+""".format(DEFAULT_VERSION)
+
+content = """TF_PYTHON_VERSION = "{version}"
+HERMETIC_PYTHON_VERSION = "{version}"
+WHEEL_NAME = "{wheel_name}"
+WHEEL_COLLAB = "{wheel_collab}"
+REQUIREMENTS = "{requirements}"
+"""
+
+def _python_repository_impl(ctx):
+    ctx.file("BUILD", "")
+    version_legacy = ctx.os.environ.get("TF_PYTHON_VERSION", "")
+    version = ctx.os.environ.get("HERMETIC_PYTHON_VERSION", "")
+    if not version:
+        version = version_legacy
+    else:
+        version_legacy = version
+
+    wheel_name = ctx.os.environ.get("WHEEL_NAME", "tensorflow")
+    wheel_collab = ctx.os.environ.get("WHEEL_COLLAB", False)
+    if version not in VERSIONS:
+        print(WARNING)  # buildifier: disable=print
+        version = DEFAULT_VERSION
+    else:
+        print("Using hermetic Python %s" % version)  # buildifier: disable=print
+
+    requirements = ""
+    for i in range(0, len(ctx.attr.requirements_locks)):
+        if ctx.attr.requirements_versions[i] == version:
+            requirements = ctx.attr.requirements_locks[i]
+            break
+
+    ctx.file(
+        "py_version.bzl",
+        content.format(
+            version = version,
+            wheel_name = wheel_name,
+            wheel_collab = wheel_collab,
+            requirements = str(requirements),
+        ),
+    )
+
+python_repository = repository_rule(
+    implementation = _python_repository_impl,
+    attrs = {
+        "requirements_versions": attr.string_list(
+            mandatory = False,
+            default = [],
+        ),
+        "requirements_locks": attr.label_list(
+            mandatory = False,
+            default = [],
+        ),
+    },
+    environ = [
+        "TF_PYTHON_VERSION",
+        "HERMETIC_PYTHON_VERSION",
+        "WHEEL_NAME",
+        "WHEEL_COLLAB",
+    ],
+)
+
+def _custom_python_interpreter_impl(ctx):
+    version = ctx.attr.version
+    strip_prefix = ctx.attr.strip_prefix.format(version = version)
+    urls = [url.format(version = version) for url in ctx.attr.urls]
+    binary_name = ctx.attr.binary_name
+    if not binary_name:
+        ver_chunks = version.split(".")
+        binary_name = "python%s.%s" % (ver_chunks[0], ver_chunks[1])
+
+    install_dir = "{name}-{version}".format(name = ctx.attr.name, version = version)
+    _exec_and_check(ctx, ["mkdir", install_dir])
+    install_path = ctx.path(install_dir)
+    srcs_dir = "srcs"
+    ctx.download_and_extract(
+        url = urls,
+        stripPrefix = strip_prefix,
+        output = srcs_dir,
+    )
+
+    configure_params = []
+    if "CC" in ctx.os.environ:
+        configure_params.append("CC={}".format(ctx.os.environ["CC"]))
+    if "CXX" in ctx.os.environ:
+        configure_params.append("CXX={}".format(ctx.os.environ["CXX"]))
+
+    configure_params.append("--enable-optimizations")
+    configure_params.append("--prefix=%s" % install_path.realpath)
+    _exec_and_check(
+        ctx,
+        ["./configure"] + configure_params,
+        working_directory = srcs_dir,
+        quiet = False,
+    )
+    res = _exec_and_check(ctx, ["nproc"])
+    cores = 12 if res.return_code != 0 else max(1, int(res.stdout.strip()) - 1)
+    _exec_and_check(ctx, ["make", "-j%s" % cores], working_directory = srcs_dir)
+    _exec_and_check(ctx, ["make", "altinstall"], working_directory = srcs_dir)
+    _exec_and_check(ctx, ["ln", "-s", binary_name, "python3"], working_directory = install_dir + "/bin")
+    tar = "{install_dir}.tgz".format(install_dir = install_dir)
+    _exec_and_check(ctx, ["tar", "czpf", tar, install_dir])
+    _exec_and_check(ctx, ["rm", "-rf", srcs_dir])
+    res = _exec_and_check(ctx, ["sha256sum", tar])
+
+    sha256 = res.stdout.split(" ")[0].strip()
+    tar_path = ctx.path(tar)
+
+    example = """\n\n
+To use newly built Python interpreter add the following code snippet RIGHT AFTER
+python_init_toolchains() in your WORKSPACE file. The code sample should work as
+is but it may need some tuning, if you have special requirements.
+
+```
+load("@rules_python//python:repositories.bzl", "python_register_toolchains")
+python_register_toolchains(
+    name = "python",
+    # By default assume the interpreter is on the local file system, replace
+    # with proper URL if it is not the case.
+    base_url = "file://",
+    ignore_root_user_error = True,
+    python_version = "{version}",
+    tool_versions = {{
+        "{version}": {{
+            # Path to .tar.gz with Python binary. By default it points to .tgz
+            # file in cache where it was built originally; replace with proper
+            # file location, if you moved it somewhere else.
+            "url": "{tar_path}",
+            "sha256": {{
+                # By default we assume Linux x86_64 architecture, eplace with
+                # proper architecture if you were building on a different platform.
+                "x86_64-unknown-linux-gnu": "{sha256}",
+            }},
+            "strip_prefix": "{install_dir}",
+        }},
+    }},
+)
+```
+\n\n""".format(version = version, tar_path = tar_path, sha256 = sha256, install_dir = install_dir)
+
+    instructions = "INSTRUCTIONS-{version}.md".format(version = version)
+    ctx.file(instructions + ".tmpl", example, executable = False)
+    ctx.file(
+        "BUILD.bazel",
+        """
+genrule(
+    name = "{name}",
+    srcs = ["{tar}", "{instructions}.tmpl"],
+    outs = ["{install_dir}.tar.gz", "{instructions}"],
+    cmd = "cp $(location {tar}) $(location {install_dir}.tar.gz); cp $(location {instructions}.tmpl) $(location {instructions})",
+    visibility = ["//visibility:public"],
+)
+     """.format(
+            name = ctx.attr.name,
+            tar = tar,
+            install_dir = install_dir,
+            instructions = instructions,
+        ),
+        executable = False,
+    )
+
+    print(example)  # buildifier: disable=print
+
+custom_python_interpreter = repository_rule(
+    implementation = _custom_python_interpreter_impl,
+    attrs = {
+        "urls": attr.string_list(),
+        "strip_prefix": attr.string(),
+        "binary_name": attr.string(mandatory = False),
+        "version": attr.string(),
+    },
+)
+
+def _exec_and_check(ctx, command, fail_on_error = True, quiet = False, **kwargs):
+    res = ctx.execute(command, quiet = quiet, **kwargs)
+    if fail_on_error and res.return_code != 0:
+        fail("""
+Failed to execute command: `{command}`
+Exit Code: {code}
+STDERR: {stderr}
+        """.format(
+            command = command,
+            code = res.return_code,
+            stderr = res.stderr,
+        ))
+    return res
diff --git a/third_party/xla/WORKSPACE b/third_party/xla/WORKSPACE
index 7ba74d6276c2e4..9d046e22949091 100644
--- a/third_party/xla/WORKSPACE
+++ b/third_party/xla/WORKSPACE
@@ -7,6 +7,32 @@ workspace(name = "xla")
 # restriction that load() statements need to be at the top of .bzl files.
 # E.g. we can not retrieve a new repository with http_archive and then load()
 # a macro from that repository in the same file.
+
+# Initialize hermetic Python
+load("//third_party/py:python_init_rules.bzl", "python_init_rules")
+
+python_init_rules()
+
+load("//third_party/py:python_init_repositories.bzl", "python_init_repositories")
+
+python_init_repositories(
+    requirements = {
+        "3.11": "//:requirements_lock_3_11.txt",
+    },
+)
+
+load("//third_party/py:python_init_toolchains.bzl", "python_init_toolchains")
+
+python_init_toolchains()
+
+load("//third_party/py:python_init_pip.bzl", "python_init_pip")
+
+python_init_pip()
+
+load("@pypi//:requirements.bzl", "install_deps")
+
+install_deps()
+
 load(":workspace4.bzl", "xla_workspace4")
 
 xla_workspace4()
diff --git a/third_party/xla/opensource_only.files b/third_party/xla/opensource_only.files
index 7655cabdafeb6b..666a858f3e4bd9 100644
--- a/third_party/xla/opensource_only.files
+++ b/third_party/xla/opensource_only.files
@@ -19,13 +19,17 @@ third_party/llvm_openmp/openmp.bzl:
 third_party/ortools/BUILD:
 third_party/ortools/glpk.BUILD:
 third_party/ortools/ortools.patch:
-third_party/py/non_hermetic/BUILD.tpl:
-third_party/py/non_hermetic/BUILD:
-third_party/py/non_hermetic/README:
-third_party/py/non_hermetic/ml_dtypes/BUILD:
-third_party/py/non_hermetic/ml_dtypes/LICENSE:
-third_party/py/non_hermetic/numpy/BUILD:
-third_party/py/non_hermetic/python_configure.bzl:
+third_party/py/BUILD.tpl:
+third_party/py/BUILD:
+third_party/py/ml_dtypes/BUILD:
+third_party/py/ml_dtypes/LICENSE:
+third_party/py/numpy/BUILD:
+third_party/py/python_configure.bzl:
+third_party/py/python_init_pip.bzl:
+third_party/py/python_init_repositories.bzl:
+third_party/py/python_init_rules.bzl:
+third_party/py/python_init_toolchains.bzl:
+third_party/py/python_repo.bzl:
 third_party/python_runtime/BUILD:
 third_party/repo.bzl:
 third_party/stablehlo/BUILD:
diff --git a/third_party/xla/requirements_lock_3_11.txt b/third_party/xla/requirements_lock_3_11.txt
new file mode 100644
index 00000000000000..5c4bb687dfecae
--- /dev/null
+++ b/third_party/xla/requirements_lock_3_11.txt
@@ -0,0 +1,49 @@
+numpy==1.24.3 \
+    --hash=sha256:0ec87a7084caa559c36e0a2309e4ecb1baa03b687201d0a847c8b0ed476a7187 \
+    --hash=sha256:1a7d6acc2e7524c9955e5c903160aa4ea083736fde7e91276b0e5d98e6332812 \
+    --hash=sha256:202de8f38fc4a45a3eea4b63e2f376e5f2dc64ef0fa692838e31a808520efaf7 \
+    --hash=sha256:210461d87fb02a84ef243cac5e814aad2b7f4be953b32cb53327bb49fd77fbb4 \
+    --hash=sha256:2d926b52ba1367f9acb76b0df6ed21f0b16a1ad87c6720a1121674e5cf63e2b6 \
+    --hash=sha256:352ee00c7f8387b44d19f4cada524586f07379c0d49270f87233983bc5087ca0 \
+    --hash=sha256:35400e6a8d102fd07c71ed7dcadd9eb62ee9a6e84ec159bd48c28235bbb0f8e4 \
+    --hash=sha256:3c1104d3c036fb81ab923f507536daedc718d0ad5a8707c6061cdfd6d184e570 \
+    --hash=sha256:4719d5aefb5189f50887773699eaf94e7d1e02bf36c1a9d353d9f46703758ca4 \
+    --hash=sha256:4749e053a29364d3452c034827102ee100986903263e89884922ef01a0a6fd2f \
+    --hash=sha256:5342cf6aad47943286afa6f1609cad9b4266a05e7f2ec408e2cf7aea7ff69d80 \
+    --hash=sha256:56e48aec79ae238f6e4395886b5eaed058abb7231fb3361ddd7bfdf4eed54289 \
+    --hash=sha256:76e3f4e85fc5d4fd311f6e9b794d0c00e7002ec122be271f2019d63376f1d385 \
+    --hash=sha256:7776ea65423ca6a15255ba1872d82d207bd1e09f6d0894ee4a64678dd2204078 \
+    --hash=sha256:784c6da1a07818491b0ffd63c6bbe5a33deaa0e25a20e1b3ea20cf0e43f8046c \
+    --hash=sha256:8535303847b89aa6b0f00aa1dc62867b5a32923e4d1681a35b5eef2d9591a463 \
+    --hash=sha256:9a7721ec204d3a237225db3e194c25268faf92e19338a35f3a224469cb6039a3 \
+    --hash=sha256:a1d3c026f57ceaad42f8231305d4653d5f05dc6332a730ae5c0bea3513de0950 \
+    --hash=sha256:ab344f1bf21f140adab8e47fdbc7c35a477dc01408791f8ba00d018dd0bc5155 \
+    --hash=sha256:ab5f23af8c16022663a652d3b25dcdc272ac3f83c3af4c02eb8b824e6b3ab9d7 \
+    --hash=sha256:ae8d0be48d1b6ed82588934aaaa179875e7dc4f3d84da18d7eae6eb3f06c242c \
+    --hash=sha256:c91c4afd8abc3908e00a44b2672718905b8611503f7ff87390cc0ac3423fb096 \
+    --hash=sha256:d5036197ecae68d7f491fcdb4df90082b0d4960ca6599ba2659957aafced7c17 \
+    --hash=sha256:d6cc757de514c00b24ae8cf5c876af2a7c3df189028d68c0cb4eaa9cd5afc2bf \
+    --hash=sha256:d933fabd8f6a319e8530d0de4fcc2e6a61917e0b0c271fded460032db42a0fe4 \
+    --hash=sha256:ea8282b9bcfe2b5e7d491d0bf7f3e2da29700cec05b49e64d6246923329f2b02 \
+    --hash=sha256:ecde0f8adef7dfdec993fd54b0f78183051b6580f606111a6d789cd14c61ea0c \
+    --hash=sha256:f21c442fdd2805e91799fbe044a7b999b8571bb0ab0f7850d0cb9641a687092b
+lit==17.0.6 \
+    --hash=sha256:dfa9af9b55fc4509a56be7bf2346f079d7f4a242d583b9f2e0b078fd0abae31b
+ml-dtypes==0.3.2 \
+    --hash=sha256:2c34f2ba9660b21fe1034b608308a01be82bbef2a92fb8199f24dc6bad0d5226 \
+    --hash=sha256:3a17ef2322e60858d93584e9c52a5be7dd6236b056b7fa1ec57f1bb6ba043e33 \
+    --hash=sha256:533059bc5f1764fac071ef54598db358c167c51a718f68f5bb55e3dee79d2967 \
+    --hash=sha256:6604877d567a29bfe7cc02969ae0f2425260e5335505cf5e7fefc3e5465f5655 \
+    --hash=sha256:6b35c4e8ca957c877ac35c79ffa77724ecc3702a1e4b18b08306c03feae597bb \
+    --hash=sha256:763697ab8a88d47443997a7cdf3aac7340049aed45f7521f6b0ec8a0594821fe \
+    --hash=sha256:7a4c3fcbf86fa52d0204f07cfd23947ef05b4ad743a1a988e163caa34a201e5e \
+    --hash=sha256:7afde548890a92b41c0fed3a6c525f1200a5727205f73dc21181a2726571bb53 \
+    --hash=sha256:7ba8e1fafc7fff3e643f453bffa7d082df1678a73286ce8187d3e825e776eb94 \
+    --hash=sha256:91f8783fd1f2c23fd3b9ee5ad66b785dafa58ba3cdb050c4458021fa4d1eb226 \
+    --hash=sha256:93b78f53431c93953f7850bb1b925a17f0ab5d97527e38a7e865b5b4bc5cfc18 \
+    --hash=sha256:961134ea44c7b8ca63eda902a44b58cd8bd670e21d62e255c81fba0a8e70d9b7 \
+    --hash=sha256:b89b194e9501a92d289c1ffd411380baf5daafb9818109a4f49b0a1b6dce4462 \
+    --hash=sha256:c7b3fb3d4f6b39bcd4f6c4b98f406291f0d681a895490ee29a0f95bab850d53c \
+    --hash=sha256:d1a746fe5fb9cd974a91070174258f0be129c592b93f9ce7df6cc336416c3fbd \
+    --hash=sha256:e8505946df1665db01332d885c2020b4cb9e84a8b1241eb4ba69d59591f65855 \
+    --hash=sha256:f47619d978ab1ae7dfdc4052ea97c636c6263e1f19bd1be0e42c346b98d15ff4
\ No newline at end of file
diff --git a/third_party/xla/third_party/py/BUILD b/third_party/xla/third_party/py/BUILD
index e69de29bb2d1d6..84eba77ce1a7af 100644
--- a/third_party/xla/third_party/py/BUILD
+++ b/third_party/xla/third_party/py/BUILD
@@ -0,0 +1,40 @@
+load("@python//:defs.bzl", "compile_pip_requirements")
+load("@python_version_repo//:py_version.bzl", "REQUIREMENTS")
+
+compile_pip_requirements(
+    name = "requirements",
+    extra_args = [
+        "--allow-unsafe",
+        "--build-isolation",
+    ],
+    generate_hashes = True,
+    requirements_in = "requirements.in",
+    requirements_txt = REQUIREMENTS,
+)
+
+compile_pip_requirements(
+    name = "requirements_nightly",
+    data = ["test-requirements.txt"],
+    extra_args = [
+        "--allow-unsafe",
+        "--build-isolation",
+        "--extra-index-url=https://pypi.anaconda.org/scientific-python-nightly-wheels/simple",
+        "--pre",
+        "--upgrade",
+    ],
+    generate_hashes = False,
+    requirements_in = "requirements.in",
+    requirements_txt = REQUIREMENTS,
+)
+
+compile_pip_requirements(
+    name = "requirements_dev",
+    extra_args = [
+        "--allow-unsafe",
+        "--build-isolation",
+        "--upgrade",
+    ],
+    generate_hashes = False,
+    requirements_in = "requirements.in",
+    requirements_txt = REQUIREMENTS,
+)
diff --git a/third_party/xla/third_party/py/non_hermetic/BUILD b/third_party/xla/third_party/py/non_hermetic/BUILD
deleted file mode 100644
index e69de29bb2d1d6..00000000000000
diff --git a/third_party/xla/third_party/py/non_hermetic/BUILD.tpl b/third_party/xla/third_party/py/non_hermetic/BUILD.tpl
deleted file mode 100644
index 45480bd4a31cf8..00000000000000
--- a/third_party/xla/third_party/py/non_hermetic/BUILD.tpl
+++ /dev/null
@@ -1,80 +0,0 @@
-licenses(["restricted"])
-
-package(default_visibility = ["//visibility:public"])
-
-# Point both runtimes to the same python binary to ensure we always
-# use the python binary specified by ./configure.py script.
-load("@bazel_tools//tools/python:toolchain.bzl", "py_runtime_pair")
-
-py_runtime(
-    name = "py2_runtime",
-    interpreter_path = "%{PYTHON_BIN_PATH}",
-    python_version = "PY2",
-)
-
-py_runtime(
-    name = "py3_runtime",
-    interpreter_path = "%{PYTHON_BIN_PATH}",
-    python_version = "PY3",
-)
-
-py_runtime_pair(
-    name = "py_runtime_pair",
-    py2_runtime = ":py2_runtime",
-    py3_runtime = ":py3_runtime",
-)
-
-toolchain(
-    name = "py_toolchain",
-    toolchain = ":py_runtime_pair",
-    toolchain_type = "@bazel_tools//tools/python:toolchain_type",
-    target_compatible_with = [%{PLATFORM_CONSTRAINT}],
-    exec_compatible_with = [%{PLATFORM_CONSTRAINT}],
-)
-
-# To build Python C/C++ extension on Windows, we need to link to python import library pythonXY.lib
-# See https://docs.python.org/3/extending/windows.html
-cc_import(
-    name = "python_lib",
-    interface_library = select({
-        ":windows": ":python_import_lib",
-        # A placeholder for Unix platforms which makes --no_build happy.
-        "//conditions:default": "not-existing.lib",
-    }),
-    system_provided = 1,
-)
-
-cc_library(
-    name = "python_headers",
-    hdrs = [":python_include"],
-    deps = select({
-        ":windows": [":python_lib"],
-        "//conditions:default": [],
-    }),
-    includes = ["python_include"],
-)
-
-# This alias is exists for the use of targets in the @llvm-project dependency,
-# which expect a python_headers target called @python_runtime//:headers. We use
-# a repo_mapping to alias python_runtime to this package, and an alias to create
-# the correct target.
-alias(
-    name = "headers",
-    actual = ":python_headers",
-)
-
-cc_library(
-    name = "numpy_headers",
-    hdrs = [":numpy_include"],
-    includes = ["numpy_include"],
-)
-
-config_setting(
-    name = "windows",
-    values = {"cpu": "x64_windows"},
-    visibility = ["//visibility:public"],
-)
-
-%{PYTHON_INCLUDE_GENRULE}
-%{NUMPY_INCLUDE_GENRULE}
-%{PYTHON_IMPORT_LIB_GENRULE}
\ No newline at end of file
diff --git a/third_party/xla/third_party/py/non_hermetic/README b/third_party/xla/third_party/py/non_hermetic/README
deleted file mode 100644
index 62188a5817a09e..00000000000000
--- a/third_party/xla/third_party/py/non_hermetic/README
+++ /dev/null
@@ -1,3 +0,0 @@
-This is a temporary copy of python_configure repository rule. It allows
-XLA and TSL to keep non-hermetic Python while TF uses hermetic Python.
-DO NOT DEPEND ON THIS COPY as it will be deleted soon.
\ No newline at end of file
diff --git a/third_party/xla/third_party/py/non_hermetic/python_configure.bzl b/third_party/xla/third_party/py/non_hermetic/python_configure.bzl
deleted file mode 100644
index 89732c3e33d8ee..00000000000000
--- a/third_party/xla/third_party/py/non_hermetic/python_configure.bzl
+++ /dev/null
@@ -1,312 +0,0 @@
-"""Repository rule for Python autoconfiguration.
-
-`python_configure` depends on the following environment variables:
-
-  * `PYTHON_BIN_PATH`: location of python binary.
-  * `PYTHON_LIB_PATH`: Location of python libraries.
-"""
-
-load(
-    "//third_party/remote_config:common.bzl",
-    "BAZEL_SH",
-    "PYTHON_BIN_PATH",
-    "PYTHON_LIB_PATH",
-    "TF_PYTHON_CONFIG_REPO",
-    "auto_config_fail",
-    "config_repo_label",
-    "execute",
-    "get_bash_bin",
-    "get_host_environ",
-    "get_python_bin",
-    "is_windows",
-    "raw_exec",
-    "read_dir",
-)
-
-def _genrule(src_dir, genrule_name, command, outs):
-    """Returns a string with a genrule.
-
-    Genrule executes the given command and produces the given outputs.
-    """
-    return (
-        "genrule(\n" +
-        '    name = "' +
-        genrule_name + '",\n' +
-        "    outs = [\n" +
-        outs +
-        "\n    ],\n" +
-        '    cmd = """\n' +
-        command +
-        '\n   """,\n' +
-        ")\n"
-    )
-
-def _norm_path(path):
-    """Returns a path with '/' and remove the trailing slash."""
-    path = path.replace("\\", "/")
-    if path[-1] == "/":
-        path = path[:-1]
-    return path
-
-def _symlink_genrule_for_dir(
-        repository_ctx,
-        src_dir,
-        dest_dir,
-        genrule_name,
-        src_files = [],
-        dest_files = []):
-    """Returns a genrule to symlink(or copy if on Windows) a set of files.
-
-    If src_dir is passed, files will be read from the given directory; otherwise
-    we assume files are in src_files and dest_files
-    """
-    if src_dir != None:
-        src_dir = _norm_path(src_dir)
-        dest_dir = _norm_path(dest_dir)
-        files = "\n".join(read_dir(repository_ctx, src_dir))
-
-        # Create a list with the src_dir stripped to use for outputs.
-        dest_files = files.replace(src_dir, "").splitlines()
-        src_files = files.splitlines()
-    command = []
-    outs = []
-    for i in range(len(dest_files)):
-        if dest_files[i] != "":
-            # If we have only one file to link we do not want to use the dest_dir, as
-            # $(@D) will include the full path to the file.
-            dest = "$(@D)/" + dest_dir + dest_files[i] if len(dest_files) != 1 else "$(@D)/" + dest_files[i]
-
-            # Copy the headers to create a sandboxable setup.
-            cmd = "cp -f"
-            command.append(cmd + ' "%s" "%s"' % (src_files[i], dest))
-            outs.append('        "' + dest_dir + dest_files[i] + '",')
-    genrule = _genrule(
-        src_dir,
-        genrule_name,
-        " && ".join(command),
-        "\n".join(outs),
-    )
-    return genrule
-
-def _get_python_lib(repository_ctx, python_bin):
-    """Gets the python lib path."""
-    python_lib = get_host_environ(repository_ctx, PYTHON_LIB_PATH)
-    if python_lib != None:
-        return python_lib
-
-    # The interesting program to execute.
-    print_lib = [
-        "from __future__ import print_function",
-        "import site",
-        "import os",
-        "python_paths = []",
-        "if os.getenv('PYTHONPATH') is not None:",
-        "  python_paths = os.getenv('PYTHONPATH').split(':')",
-        "try:",
-        "  library_paths = site.getsitepackages()",
-        "except AttributeError:",
-        "  from distutils.sysconfig import get_python_lib",
-        "  library_paths = [get_python_lib()]",
-        "all_paths = set(python_paths + library_paths)",
-        "paths = []",
-        "for path in all_paths:",
-        "  if os.path.isdir(path):",
-        "    paths.append(path)",
-        "if len(paths) >=1:",
-        "  print(paths[0])",
-    ]
-
-    # The below script writes the above program to a file
-    # and executes it. This is to work around the limitation
-    # of not being able to upload files as part of execute.
-    cmd = "from os import linesep;"
-    cmd += "f = open('script.py', 'w');"
-    for line in print_lib:
-        cmd += "f.write(\"%s\" + linesep);" % line
-    cmd += "f.close();"
-    cmd += "from subprocess import call;"
-    cmd += "call([\"%s\", \"script.py\"]);" % python_bin
-
-    result = execute(repository_ctx, [python_bin, "-c", cmd])
-    return result.stdout.strip()
-
-def _check_python_lib(repository_ctx, python_lib):
-    """Checks the python lib path."""
-    cmd = 'test -d "%s" -a -x "%s"' % (python_lib, python_lib)
-    result = raw_exec(repository_ctx, [get_bash_bin(repository_ctx), "-c", cmd])
-    if result.return_code == 1:
-        auto_config_fail("Invalid python library path: %s" % python_lib)
-
-def _check_python_bin(repository_ctx, python_bin):
-    """Checks the python bin path."""
-    cmd = '[[ -x "%s" ]] && [[ ! -d "%s" ]]' % (python_bin, python_bin)
-    result = raw_exec(repository_ctx, [get_bash_bin(repository_ctx), "-c", cmd])
-    if result.return_code == 1:
-        auto_config_fail("--define %s='%s' is not executable. Is it the python binary?" % (
-            PYTHON_BIN_PATH,
-            python_bin,
-        ))
-
-def _get_python_include(repository_ctx, python_bin):
-    """Gets the python include path."""
-    result = execute(
-        repository_ctx,
-        [
-            python_bin,
-            "-Wignore",
-            "-c",
-            "import sysconfig; " +
-            "print(sysconfig.get_path('include'))",
-        ],
-        error_msg = "Problem getting python include path.",
-        error_details = ("Is the Python binary path set up right? " +
-                         "(See ./configure or " + PYTHON_BIN_PATH + ".) " +
-                         "Is distutils installed?"),
-    )
-    return result.stdout.splitlines()[0]
-
-def _get_python_import_lib_name(repository_ctx, python_bin):
-    """Get Python import library name (pythonXY.lib) on Windows."""
-    result = execute(
-        repository_ctx,
-        [
-            python_bin,
-            "-c",
-            "import sys;" +
-            'print("python" + str(sys.version_info[0]) + ' +
-            '      str(sys.version_info[1]) + ".lib")',
-        ],
-        error_msg = "Problem getting python import library.",
-        error_details = ("Is the Python binary path set up right? " +
-                         "(See ./configure or " + PYTHON_BIN_PATH + ".) "),
-    )
-    return result.stdout.splitlines()[0]
-
-def _get_numpy_include(repository_ctx, python_bin):
-    """Gets the numpy include path."""
-    return execute(
-        repository_ctx,
-        [
-            python_bin,
-            "-c",
-            "from __future__ import print_function;" +
-            "import numpy;" +
-            " print(numpy.get_include());",
-        ],
-        error_msg = "Problem getting numpy include path.",
-        error_details = "Is numpy installed?",
-    ).stdout.splitlines()[0]
-
-def _create_local_python_repository(repository_ctx):
-    """Creates the repository containing files set up to build with Python."""
-
-    # Resolve all labels before doing any real work. Resolving causes the
-    # function to be restarted with all previous state being lost. This
-    # can easily lead to a O(n^2) runtime in the number of labels.
-    build_tpl = repository_ctx.path(Label("//third_party/py:BUILD.tpl"))
-
-    python_bin = get_python_bin(repository_ctx)
-    _check_python_bin(repository_ctx, python_bin)
-    python_lib = _get_python_lib(repository_ctx, python_bin)
-    _check_python_lib(repository_ctx, python_lib)
-    python_include = _get_python_include(repository_ctx, python_bin)
-    numpy_include = _get_numpy_include(repository_ctx, python_bin) + "/numpy"
-    python_include_rule = _symlink_genrule_for_dir(
-        repository_ctx,
-        python_include,
-        "python_include",
-        "python_include",
-    )
-    python_import_lib_genrule = ""
-
-    # To build Python C/C++ extension on Windows, we need to link to python import library pythonXY.lib
-    # See https://docs.python.org/3/extending/windows.html
-    if is_windows(repository_ctx):
-        python_bin = python_bin.replace("\\", "/")
-        python_include = _norm_path(python_include)
-        python_import_lib_name = _get_python_import_lib_name(repository_ctx, python_bin)
-        python_import_lib_src = python_include.rsplit("/", 1)[0] + "/libs/" + python_import_lib_name
-        python_import_lib_genrule = _symlink_genrule_for_dir(
-            repository_ctx,
-            None,
-            "",
-            "python_import_lib",
-            [python_import_lib_src],
-            [python_import_lib_name],
-        )
-    numpy_include_rule = _symlink_genrule_for_dir(
-        repository_ctx,
-        numpy_include,
-        "numpy_include/numpy",
-        "numpy_include",
-    )
-
-    platform_constraint = ""
-    if repository_ctx.attr.platform_constraint:
-        platform_constraint = "\"%s\"" % repository_ctx.attr.platform_constraint
-    repository_ctx.template("BUILD", build_tpl, {
-        "%{PYTHON_BIN_PATH}": python_bin,
-        "%{PYTHON_INCLUDE_GENRULE}": python_include_rule,
-        "%{PYTHON_IMPORT_LIB_GENRULE}": python_import_lib_genrule,
-        "%{NUMPY_INCLUDE_GENRULE}": numpy_include_rule,
-        "%{PLATFORM_CONSTRAINT}": platform_constraint,
-    })
-
-def _create_remote_python_repository(repository_ctx, remote_config_repo):
-    """Creates pointers to a remotely configured repo set up to build with Python.
-    """
-    repository_ctx.template("BUILD", config_repo_label(remote_config_repo, ":BUILD"), {})
-
-def _python_autoconf_impl(repository_ctx):
-    """Implementation of the python_autoconf repository rule."""
-    if get_host_environ(repository_ctx, TF_PYTHON_CONFIG_REPO) != None:
-        _create_remote_python_repository(
-            repository_ctx,
-            get_host_environ(repository_ctx, TF_PYTHON_CONFIG_REPO),
-        )
-    else:
-        _create_local_python_repository(repository_ctx)
-
-_ENVIRONS = [
-    BAZEL_SH,
-    PYTHON_BIN_PATH,
-    PYTHON_LIB_PATH,
-]
-
-local_python_configure = repository_rule(
-    implementation = _create_local_python_repository,
-    environ = _ENVIRONS,
-    attrs = {
-        "environ": attr.string_dict(),
-        "platform_constraint": attr.string(),
-    },
-)
-
-remote_python_configure = repository_rule(
-    implementation = _create_local_python_repository,
-    environ = _ENVIRONS,
-    remotable = True,
-    attrs = {
-        "environ": attr.string_dict(),
-        "platform_constraint": attr.string(),
-    },
-)
-
-python_configure = repository_rule(
-    implementation = _python_autoconf_impl,
-    environ = _ENVIRONS + [TF_PYTHON_CONFIG_REPO],
-    attrs = {
-        "platform_constraint": attr.string(),
-    },
-)
-"""Detects and configures the local Python.
-
-Add the following to your WORKSPACE FILE:
-
-```python
-python_configure(name = "local_config_python")
-```
-
-Args:
-  name: A unique name for this workspace rule.
-"""
diff --git a/third_party/xla/third_party/py/python_init_pip.bzl b/third_party/xla/third_party/py/python_init_pip.bzl
new file mode 100644
index 00000000000000..efc2bf8233cf61
--- /dev/null
+++ b/third_party/xla/third_party/py/python_init_pip.bzl
@@ -0,0 +1,34 @@
+"""Hermetic Python initialization. Consult the WORKSPACE on how to use it."""
+
+load("@python//:defs.bzl", "interpreter")
+load("@python_version_repo//:py_version.bzl", "REQUIREMENTS")
+load("@rules_python//python:pip.bzl", "package_annotation", "pip_parse")
+
+def python_init_pip():
+    numpy_annotations = {
+        "numpy": package_annotation(
+            additive_build_content = """\
+cc_library(
+    name = "numpy_headers_2",
+    hdrs = glob(["site-packages/numpy/_core/include/**/*.h"]),
+    strip_include_prefix="site-packages/numpy/_core/include/",
+)
+cc_library(
+    name = "numpy_headers_1",
+    hdrs = glob(["site-packages/numpy/core/include/**/*.h"]),
+    strip_include_prefix="site-packages/numpy/core/include/",
+)
+cc_library(
+    name = "numpy_headers",
+    deps = [":numpy_headers_2", ":numpy_headers_1"],
+)
+""",
+        ),
+    }
+
+    pip_parse(
+        name = "pypi",
+        annotations = numpy_annotations,
+        python_interpreter_target = interpreter,
+        requirements_lock = REQUIREMENTS,
+    )
diff --git a/third_party/xla/third_party/py/python_init_repositories.bzl b/third_party/xla/third_party/py/python_init_repositories.bzl
new file mode 100644
index 00000000000000..5a405f2c2aba4c
--- /dev/null
+++ b/third_party/xla/third_party/py/python_init_repositories.bzl
@@ -0,0 +1,12 @@
+"""Hermetic Python initialization. Consult the WORKSPACE on how to use it."""
+
+load("@rules_python//python:repositories.bzl", "py_repositories")
+load("//third_party/py:python_repo.bzl", "python_repository")
+
+def python_init_repositories(requirements = {}):
+    python_repository(
+        name = "python_version_repo",
+        requirements_versions = requirements.keys(),
+        requirements_locks = requirements.values(),
+    )
+    py_repositories()
diff --git a/third_party/xla/third_party/py/python_init_rules.bzl b/third_party/xla/third_party/py/python_init_rules.bzl
new file mode 100644
index 00000000000000..98a7b8bc3c315a
--- /dev/null
+++ b/third_party/xla/third_party/py/python_init_rules.bzl
@@ -0,0 +1,11 @@
+"""Hermetic Python initialization. Consult the WORKSPACE on how to use it."""
+
+load("@bazel_tools//tools/build_defs/repo:http.bzl", "http_archive")
+
+def python_init_rules():
+    http_archive(
+        name = "rules_python",
+        sha256 = "9d04041ac92a0985e344235f5d946f71ac543f1b1565f2cdbc9a2aaee8adf55b",
+        strip_prefix = "rules_python-0.26.0",
+        url = "https://github.com/bazelbuild/rules_python/releases/download/0.26.0/rules_python-0.26.0.tar.gz",
+    )
diff --git a/third_party/xla/third_party/py/python_init_toolchains.bzl b/third_party/xla/third_party/py/python_init_toolchains.bzl
new file mode 100644
index 00000000000000..c1f800db4c01e7
--- /dev/null
+++ b/third_party/xla/third_party/py/python_init_toolchains.bzl
@@ -0,0 +1,13 @@
+"""Hermetic Python initialization. Consult the WORKSPACE on how to use it."""
+
+load("@python_version_repo//:py_version.bzl", "HERMETIC_PYTHON_VERSION")
+load("@rules_python//python:repositories.bzl", "python_register_toolchains")
+load("@rules_python//python:versions.bzl", "MINOR_MAPPING")
+
+def python_init_toolchains():
+    if HERMETIC_PYTHON_VERSION in MINOR_MAPPING:
+        python_register_toolchains(
+            name = "python",
+            ignore_root_user_error = True,
+            python_version = HERMETIC_PYTHON_VERSION,
+        )
diff --git a/third_party/xla/third_party/py/python_repo.bzl b/third_party/xla/third_party/py/python_repo.bzl
new file mode 100644
index 00000000000000..77a6ce9ce50b60
--- /dev/null
+++ b/third_party/xla/third_party/py/python_repo.bzl
@@ -0,0 +1,206 @@
+"""
+Repository rule to manage hermetic Python interpreter under Bazel.
+
+Version can be set via build parameter "--repo_env=HERMETIC_PYTHON_VERSION=3.11"
+Defaults to 3.11.
+
+To set wheel name, add "--repo_env=WHEEL_NAME=tensorflow_cpu"
+"""
+
+VERSIONS = ["3.9", "3.10", "3.11", "3.12", "3.13"]
+DEFAULT_VERSION = "3.11"
+WARNING = """
+HERMETIC_PYTHON_VERSION variable was not set correctly, using default version. 
+Python {} will be used.
+To select Python version, either set HERMETIC_PYTHON_VERSION env variable in 
+your shell:
+  export HERMETIC_PYTHON_VERSION=3.12
+OR pass it as an argument to bazel command directly or inside your .bazelrc 
+file:
+  --repo_env=HERMETIC_PYTHON_VERSION=3.12
+""".format(DEFAULT_VERSION)
+
+content = """TF_PYTHON_VERSION = "{version}"
+HERMETIC_PYTHON_VERSION = "{version}"
+WHEEL_NAME = "{wheel_name}"
+WHEEL_COLLAB = "{wheel_collab}"
+REQUIREMENTS = "{requirements}"
+"""
+
+def _python_repository_impl(ctx):
+    ctx.file("BUILD", "")
+    version_legacy = ctx.os.environ.get("TF_PYTHON_VERSION", "")
+    version = ctx.os.environ.get("HERMETIC_PYTHON_VERSION", "")
+    if not version:
+        version = version_legacy
+    else:
+        version_legacy = version
+
+    wheel_name = ctx.os.environ.get("WHEEL_NAME", "tensorflow")
+    wheel_collab = ctx.os.environ.get("WHEEL_COLLAB", False)
+    if version not in VERSIONS:
+        print(WARNING)  # buildifier: disable=print
+        version = DEFAULT_VERSION
+    else:
+        print("Using hermetic Python %s" % version)  # buildifier: disable=print
+
+    requirements = ""
+    for i in range(0, len(ctx.attr.requirements_locks)):
+        if ctx.attr.requirements_versions[i] == version:
+            requirements = ctx.attr.requirements_locks[i]
+            break
+
+    ctx.file(
+        "py_version.bzl",
+        content.format(
+            version = version,
+            wheel_name = wheel_name,
+            wheel_collab = wheel_collab,
+            requirements = str(requirements),
+        ),
+    )
+
+python_repository = repository_rule(
+    implementation = _python_repository_impl,
+    attrs = {
+        "requirements_versions": attr.string_list(
+            mandatory = False,
+            default = [],
+        ),
+        "requirements_locks": attr.label_list(
+            mandatory = False,
+            default = [],
+        ),
+    },
+    environ = [
+        "TF_PYTHON_VERSION",
+        "HERMETIC_PYTHON_VERSION",
+        "WHEEL_NAME",
+        "WHEEL_COLLAB",
+    ],
+)
+
+def _custom_python_interpreter_impl(ctx):
+    version = ctx.attr.version
+    strip_prefix = ctx.attr.strip_prefix.format(version = version)
+    urls = [url.format(version = version) for url in ctx.attr.urls]
+    binary_name = ctx.attr.binary_name
+    if not binary_name:
+        ver_chunks = version.split(".")
+        binary_name = "python%s.%s" % (ver_chunks[0], ver_chunks[1])
+
+    install_dir = "{name}-{version}".format(name = ctx.attr.name, version = version)
+    _exec_and_check(ctx, ["mkdir", install_dir])
+    install_path = ctx.path(install_dir)
+    srcs_dir = "srcs"
+    ctx.download_and_extract(
+        url = urls,
+        stripPrefix = strip_prefix,
+        output = srcs_dir,
+    )
+
+    configure_params = []
+    if "CC" in ctx.os.environ:
+        configure_params.append("CC={}".format(ctx.os.environ["CC"]))
+    if "CXX" in ctx.os.environ:
+        configure_params.append("CXX={}".format(ctx.os.environ["CXX"]))
+
+    configure_params.append("--enable-optimizations")
+    configure_params.append("--prefix=%s" % install_path.realpath)
+    _exec_and_check(
+        ctx,
+        ["./configure"] + configure_params,
+        working_directory = srcs_dir,
+        quiet = False,
+    )
+    res = _exec_and_check(ctx, ["nproc"])
+    cores = 12 if res.return_code != 0 else max(1, int(res.stdout.strip()) - 1)
+    _exec_and_check(ctx, ["make", "-j%s" % cores], working_directory = srcs_dir)
+    _exec_and_check(ctx, ["make", "altinstall"], working_directory = srcs_dir)
+    _exec_and_check(ctx, ["ln", "-s", binary_name, "python3"], working_directory = install_dir + "/bin")
+    tar = "{install_dir}.tgz".format(install_dir = install_dir)
+    _exec_and_check(ctx, ["tar", "czpf", tar, install_dir])
+    _exec_and_check(ctx, ["rm", "-rf", srcs_dir])
+    res = _exec_and_check(ctx, ["sha256sum", tar])
+
+    sha256 = res.stdout.split(" ")[0].strip()
+    tar_path = ctx.path(tar)
+
+    example = """\n\n
+To use newly built Python interpreter add the following code snippet RIGHT AFTER
+python_init_toolchains() in your WORKSPACE file. The code sample should work as
+is but it may need some tuning, if you have special requirements.
+
+```
+load("@rules_python//python:repositories.bzl", "python_register_toolchains")
+python_register_toolchains(
+    name = "python",
+    # By default assume the interpreter is on the local file system, replace
+    # with proper URL if it is not the case.
+    base_url = "file://",
+    ignore_root_user_error = True,
+    python_version = "{version}",
+    tool_versions = {{
+        "{version}": {{
+            # Path to .tar.gz with Python binary. By default it points to .tgz
+            # file in cache where it was built originally; replace with proper
+            # file location, if you moved it somewhere else.
+            "url": "{tar_path}",
+            "sha256": {{
+                # By default we assume Linux x86_64 architecture, eplace with
+                # proper architecture if you were building on a different platform.
+                "x86_64-unknown-linux-gnu": "{sha256}",
+            }},
+            "strip_prefix": "{install_dir}",
+        }},
+    }},
+)
+```
+\n\n""".format(version = version, tar_path = tar_path, sha256 = sha256, install_dir = install_dir)
+
+    instructions = "INSTRUCTIONS-{version}.md".format(version = version)
+    ctx.file(instructions + ".tmpl", example, executable = False)
+    ctx.file(
+        "BUILD.bazel",
+        """
+genrule(
+    name = "{name}",
+    srcs = ["{tar}", "{instructions}.tmpl"],
+    outs = ["{install_dir}.tar.gz", "{instructions}"],
+    cmd = "cp $(location {tar}) $(location {install_dir}.tar.gz); cp $(location {instructions}.tmpl) $(location {instructions})",
+    visibility = ["//visibility:public"],
+)
+     """.format(
+            name = ctx.attr.name,
+            tar = tar,
+            install_dir = install_dir,
+            instructions = instructions,
+        ),
+        executable = False,
+    )
+
+    print(example)  # buildifier: disable=print
+
+custom_python_interpreter = repository_rule(
+    implementation = _custom_python_interpreter_impl,
+    attrs = {
+        "urls": attr.string_list(),
+        "strip_prefix": attr.string(),
+        "binary_name": attr.string(mandatory = False),
+        "version": attr.string(),
+    },
+)
+
+def _exec_and_check(ctx, command, fail_on_error = True, quiet = False, **kwargs):
+    res = ctx.execute(command, quiet = quiet, **kwargs)
+    if fail_on_error and res.return_code != 0:
+        fail("""
+Failed to execute command: `{command}`
+Exit Code: {code}
+STDERR: {stderr}
+        """.format(
+            command = command,
+            code = res.return_code,
+            stderr = res.stderr,
+        ))
+    return res
diff --git a/third_party/xla/third_party/tsl/WORKSPACE b/third_party/xla/third_party/tsl/WORKSPACE
index 6ad0d6e0e2b7be..19350e3dbba762 100644
--- a/third_party/xla/third_party/tsl/WORKSPACE
+++ b/third_party/xla/third_party/tsl/WORKSPACE
@@ -10,6 +10,31 @@ workspace(name = "tsl")
 
 # buildifier: disable=load-on-top
 
+# Initialize hermetic Python
+load("//third_party/py:python_init_rules.bzl", "python_init_rules")
+
+python_init_rules()
+
+load("//third_party/py:python_init_repositories.bzl", "python_init_repositories")
+
+python_init_repositories(
+    requirements = {
+        "3.11": "//:requirements_lock_3_11.txt",
+    },
+)
+
+load("//third_party/py:python_init_toolchains.bzl", "python_init_toolchains")
+
+python_init_toolchains()
+
+load("//third_party/py:python_init_pip.bzl", "python_init_pip")
+
+python_init_pip()
+
+load("@pypi//:requirements.bzl", "install_deps")
+
+install_deps()
+
 load(":workspace3.bzl", "tsl_workspace3")
 
 tsl_workspace3()
diff --git a/third_party/xla/third_party/tsl/opensource_only.files b/third_party/xla/third_party/tsl/opensource_only.files
index 3f2bcf8431edc0..709cacf5c1756e 100644
--- a/third_party/xla/third_party/tsl/opensource_only.files
+++ b/third_party/xla/third_party/tsl/opensource_only.files
@@ -73,13 +73,17 @@ third_party/nccl/system.BUILD.tpl:
 third_party/nvtx/BUILD:
 third_party/nvtx/LICENSE:
 third_party/protobuf/BUILD:
-third_party/py/non_hermetic/BUILD.tpl:
-third_party/py/non_hermetic/BUILD:
-third_party/py/non_hermetic/README:
-third_party/py/non_hermetic/ml_dtypes/BUILD:
-third_party/py/non_hermetic/ml_dtypes/LICENSE:
-third_party/py/non_hermetic/numpy/BUILD:
-third_party/py/non_hermetic/python_configure.bzl:
+third_party/py/BUILD.tpl:
+third_party/py/BUILD:
+third_party/py/ml_dtypes/BUILD:
+third_party/py/ml_dtypes/LICENSE:
+third_party/py/numpy/BUILD:
+third_party/py/python_configure.bzl:
+third_party/py/python_init_pip.bzl:
+third_party/py/python_init_repositories.bzl:
+third_party/py/python_init_rules.bzl:
+third_party/py/python_init_toolchains.bzl:
+third_party/py/python_repo.bzl:
 third_party/pybind11.BUILD:
 third_party/pybind11_bazel/BUILD:
 third_party/python_runtime/BUILD:
diff --git a/third_party/xla/third_party/tsl/requirements_lock_3_11.txt b/third_party/xla/third_party/tsl/requirements_lock_3_11.txt
new file mode 100644
index 00000000000000..5c4bb687dfecae
--- /dev/null
+++ b/third_party/xla/third_party/tsl/requirements_lock_3_11.txt
@@ -0,0 +1,49 @@
+numpy==1.24.3 \
+    --hash=sha256:0ec87a7084caa559c36e0a2309e4ecb1baa03b687201d0a847c8b0ed476a7187 \
+    --hash=sha256:1a7d6acc2e7524c9955e5c903160aa4ea083736fde7e91276b0e5d98e6332812 \
+    --hash=sha256:202de8f38fc4a45a3eea4b63e2f376e5f2dc64ef0fa692838e31a808520efaf7 \
+    --hash=sha256:210461d87fb02a84ef243cac5e814aad2b7f4be953b32cb53327bb49fd77fbb4 \
+    --hash=sha256:2d926b52ba1367f9acb76b0df6ed21f0b16a1ad87c6720a1121674e5cf63e2b6 \
+    --hash=sha256:352ee00c7f8387b44d19f4cada524586f07379c0d49270f87233983bc5087ca0 \
+    --hash=sha256:35400e6a8d102fd07c71ed7dcadd9eb62ee9a6e84ec159bd48c28235bbb0f8e4 \
+    --hash=sha256:3c1104d3c036fb81ab923f507536daedc718d0ad5a8707c6061cdfd6d184e570 \
+    --hash=sha256:4719d5aefb5189f50887773699eaf94e7d1e02bf36c1a9d353d9f46703758ca4 \
+    --hash=sha256:4749e053a29364d3452c034827102ee100986903263e89884922ef01a0a6fd2f \
+    --hash=sha256:5342cf6aad47943286afa6f1609cad9b4266a05e7f2ec408e2cf7aea7ff69d80 \
+    --hash=sha256:56e48aec79ae238f6e4395886b5eaed058abb7231fb3361ddd7bfdf4eed54289 \
+    --hash=sha256:76e3f4e85fc5d4fd311f6e9b794d0c00e7002ec122be271f2019d63376f1d385 \
+    --hash=sha256:7776ea65423ca6a15255ba1872d82d207bd1e09f6d0894ee4a64678dd2204078 \
+    --hash=sha256:784c6da1a07818491b0ffd63c6bbe5a33deaa0e25a20e1b3ea20cf0e43f8046c \
+    --hash=sha256:8535303847b89aa6b0f00aa1dc62867b5a32923e4d1681a35b5eef2d9591a463 \
+    --hash=sha256:9a7721ec204d3a237225db3e194c25268faf92e19338a35f3a224469cb6039a3 \
+    --hash=sha256:a1d3c026f57ceaad42f8231305d4653d5f05dc6332a730ae5c0bea3513de0950 \
+    --hash=sha256:ab344f1bf21f140adab8e47fdbc7c35a477dc01408791f8ba00d018dd0bc5155 \
+    --hash=sha256:ab5f23af8c16022663a652d3b25dcdc272ac3f83c3af4c02eb8b824e6b3ab9d7 \
+    --hash=sha256:ae8d0be48d1b6ed82588934aaaa179875e7dc4f3d84da18d7eae6eb3f06c242c \
+    --hash=sha256:c91c4afd8abc3908e00a44b2672718905b8611503f7ff87390cc0ac3423fb096 \
+    --hash=sha256:d5036197ecae68d7f491fcdb4df90082b0d4960ca6599ba2659957aafced7c17 \
+    --hash=sha256:d6cc757de514c00b24ae8cf5c876af2a7c3df189028d68c0cb4eaa9cd5afc2bf \
+    --hash=sha256:d933fabd8f6a319e8530d0de4fcc2e6a61917e0b0c271fded460032db42a0fe4 \
+    --hash=sha256:ea8282b9bcfe2b5e7d491d0bf7f3e2da29700cec05b49e64d6246923329f2b02 \
+    --hash=sha256:ecde0f8adef7dfdec993fd54b0f78183051b6580f606111a6d789cd14c61ea0c \
+    --hash=sha256:f21c442fdd2805e91799fbe044a7b999b8571bb0ab0f7850d0cb9641a687092b
+lit==17.0.6 \
+    --hash=sha256:dfa9af9b55fc4509a56be7bf2346f079d7f4a242d583b9f2e0b078fd0abae31b
+ml-dtypes==0.3.2 \
+    --hash=sha256:2c34f2ba9660b21fe1034b608308a01be82bbef2a92fb8199f24dc6bad0d5226 \
+    --hash=sha256:3a17ef2322e60858d93584e9c52a5be7dd6236b056b7fa1ec57f1bb6ba043e33 \
+    --hash=sha256:533059bc5f1764fac071ef54598db358c167c51a718f68f5bb55e3dee79d2967 \
+    --hash=sha256:6604877d567a29bfe7cc02969ae0f2425260e5335505cf5e7fefc3e5465f5655 \
+    --hash=sha256:6b35c4e8ca957c877ac35c79ffa77724ecc3702a1e4b18b08306c03feae597bb \
+    --hash=sha256:763697ab8a88d47443997a7cdf3aac7340049aed45f7521f6b0ec8a0594821fe \
+    --hash=sha256:7a4c3fcbf86fa52d0204f07cfd23947ef05b4ad743a1a988e163caa34a201e5e \
+    --hash=sha256:7afde548890a92b41c0fed3a6c525f1200a5727205f73dc21181a2726571bb53 \
+    --hash=sha256:7ba8e1fafc7fff3e643f453bffa7d082df1678a73286ce8187d3e825e776eb94 \
+    --hash=sha256:91f8783fd1f2c23fd3b9ee5ad66b785dafa58ba3cdb050c4458021fa4d1eb226 \
+    --hash=sha256:93b78f53431c93953f7850bb1b925a17f0ab5d97527e38a7e865b5b4bc5cfc18 \
+    --hash=sha256:961134ea44c7b8ca63eda902a44b58cd8bd670e21d62e255c81fba0a8e70d9b7 \
+    --hash=sha256:b89b194e9501a92d289c1ffd411380baf5daafb9818109a4f49b0a1b6dce4462 \
+    --hash=sha256:c7b3fb3d4f6b39bcd4f6c4b98f406291f0d681a895490ee29a0f95bab850d53c \
+    --hash=sha256:d1a746fe5fb9cd974a91070174258f0be129c592b93f9ce7df6cc336416c3fbd \
+    --hash=sha256:e8505946df1665db01332d885c2020b4cb9e84a8b1241eb4ba69d59591f65855 \
+    --hash=sha256:f47619d978ab1ae7dfdc4052ea97c636c6263e1f19bd1be0e42c346b98d15ff4
\ No newline at end of file
diff --git a/third_party/xla/third_party/tsl/third_party/py/BUILD b/third_party/xla/third_party/tsl/third_party/py/BUILD
index e69de29bb2d1d6..84eba77ce1a7af 100644
--- a/third_party/xla/third_party/tsl/third_party/py/BUILD
+++ b/third_party/xla/third_party/tsl/third_party/py/BUILD
@@ -0,0 +1,40 @@
+load("@python//:defs.bzl", "compile_pip_requirements")
+load("@python_version_repo//:py_version.bzl", "REQUIREMENTS")
+
+compile_pip_requirements(
+    name = "requirements",
+    extra_args = [
+        "--allow-unsafe",
+        "--build-isolation",
+    ],
+    generate_hashes = True,
+    requirements_in = "requirements.in",
+    requirements_txt = REQUIREMENTS,
+)
+
+compile_pip_requirements(
+    name = "requirements_nightly",
+    data = ["test-requirements.txt"],
+    extra_args = [
+        "--allow-unsafe",
+        "--build-isolation",
+        "--extra-index-url=https://pypi.anaconda.org/scientific-python-nightly-wheels/simple",
+        "--pre",
+        "--upgrade",
+    ],
+    generate_hashes = False,
+    requirements_in = "requirements.in",
+    requirements_txt = REQUIREMENTS,
+)
+
+compile_pip_requirements(
+    name = "requirements_dev",
+    extra_args = [
+        "--allow-unsafe",
+        "--build-isolation",
+        "--upgrade",
+    ],
+    generate_hashes = False,
+    requirements_in = "requirements.in",
+    requirements_txt = REQUIREMENTS,
+)
diff --git a/third_party/xla/third_party/tsl/third_party/py/non_hermetic/BUILD b/third_party/xla/third_party/tsl/third_party/py/non_hermetic/BUILD
deleted file mode 100644
index e69de29bb2d1d6..00000000000000
diff --git a/third_party/xla/third_party/tsl/third_party/py/non_hermetic/BUILD.tpl b/third_party/xla/third_party/tsl/third_party/py/non_hermetic/BUILD.tpl
deleted file mode 100644
index 45480bd4a31cf8..00000000000000
--- a/third_party/xla/third_party/tsl/third_party/py/non_hermetic/BUILD.tpl
+++ /dev/null
@@ -1,80 +0,0 @@
-licenses(["restricted"])
-
-package(default_visibility = ["//visibility:public"])
-
-# Point both runtimes to the same python binary to ensure we always
-# use the python binary specified by ./configure.py script.
-load("@bazel_tools//tools/python:toolchain.bzl", "py_runtime_pair")
-
-py_runtime(
-    name = "py2_runtime",
-    interpreter_path = "%{PYTHON_BIN_PATH}",
-    python_version = "PY2",
-)
-
-py_runtime(
-    name = "py3_runtime",
-    interpreter_path = "%{PYTHON_BIN_PATH}",
-    python_version = "PY3",
-)
-
-py_runtime_pair(
-    name = "py_runtime_pair",
-    py2_runtime = ":py2_runtime",
-    py3_runtime = ":py3_runtime",
-)
-
-toolchain(
-    name = "py_toolchain",
-    toolchain = ":py_runtime_pair",
-    toolchain_type = "@bazel_tools//tools/python:toolchain_type",
-    target_compatible_with = [%{PLATFORM_CONSTRAINT}],
-    exec_compatible_with = [%{PLATFORM_CONSTRAINT}],
-)
-
-# To build Python C/C++ extension on Windows, we need to link to python import library pythonXY.lib
-# See https://docs.python.org/3/extending/windows.html
-cc_import(
-    name = "python_lib",
-    interface_library = select({
-        ":windows": ":python_import_lib",
-        # A placeholder for Unix platforms which makes --no_build happy.
-        "//conditions:default": "not-existing.lib",
-    }),
-    system_provided = 1,
-)
-
-cc_library(
-    name = "python_headers",
-    hdrs = [":python_include"],
-    deps = select({
-        ":windows": [":python_lib"],
-        "//conditions:default": [],
-    }),
-    includes = ["python_include"],
-)
-
-# This alias is exists for the use of targets in the @llvm-project dependency,
-# which expect a python_headers target called @python_runtime//:headers. We use
-# a repo_mapping to alias python_runtime to this package, and an alias to create
-# the correct target.
-alias(
-    name = "headers",
-    actual = ":python_headers",
-)
-
-cc_library(
-    name = "numpy_headers",
-    hdrs = [":numpy_include"],
-    includes = ["numpy_include"],
-)
-
-config_setting(
-    name = "windows",
-    values = {"cpu": "x64_windows"},
-    visibility = ["//visibility:public"],
-)
-
-%{PYTHON_INCLUDE_GENRULE}
-%{NUMPY_INCLUDE_GENRULE}
-%{PYTHON_IMPORT_LIB_GENRULE}
\ No newline at end of file
diff --git a/third_party/xla/third_party/tsl/third_party/py/non_hermetic/README b/third_party/xla/third_party/tsl/third_party/py/non_hermetic/README
deleted file mode 100644
index 62188a5817a09e..00000000000000
--- a/third_party/xla/third_party/tsl/third_party/py/non_hermetic/README
+++ /dev/null
@@ -1,3 +0,0 @@
-This is a temporary copy of python_configure repository rule. It allows
-XLA and TSL to keep non-hermetic Python while TF uses hermetic Python.
-DO NOT DEPEND ON THIS COPY as it will be deleted soon.
\ No newline at end of file
diff --git a/third_party/xla/third_party/tsl/third_party/py/non_hermetic/python_configure.bzl b/third_party/xla/third_party/tsl/third_party/py/non_hermetic/python_configure.bzl
deleted file mode 100644
index 89732c3e33d8ee..00000000000000
--- a/third_party/xla/third_party/tsl/third_party/py/non_hermetic/python_configure.bzl
+++ /dev/null
@@ -1,312 +0,0 @@
-"""Repository rule for Python autoconfiguration.
-
-`python_configure` depends on the following environment variables:
-
-  * `PYTHON_BIN_PATH`: location of python binary.
-  * `PYTHON_LIB_PATH`: Location of python libraries.
-"""
-
-load(
-    "//third_party/remote_config:common.bzl",
-    "BAZEL_SH",
-    "PYTHON_BIN_PATH",
-    "PYTHON_LIB_PATH",
-    "TF_PYTHON_CONFIG_REPO",
-    "auto_config_fail",
-    "config_repo_label",
-    "execute",
-    "get_bash_bin",
-    "get_host_environ",
-    "get_python_bin",
-    "is_windows",
-    "raw_exec",
-    "read_dir",
-)
-
-def _genrule(src_dir, genrule_name, command, outs):
-    """Returns a string with a genrule.
-
-    Genrule executes the given command and produces the given outputs.
-    """
-    return (
-        "genrule(\n" +
-        '    name = "' +
-        genrule_name + '",\n' +
-        "    outs = [\n" +
-        outs +
-        "\n    ],\n" +
-        '    cmd = """\n' +
-        command +
-        '\n   """,\n' +
-        ")\n"
-    )
-
-def _norm_path(path):
-    """Returns a path with '/' and remove the trailing slash."""
-    path = path.replace("\\", "/")
-    if path[-1] == "/":
-        path = path[:-1]
-    return path
-
-def _symlink_genrule_for_dir(
-        repository_ctx,
-        src_dir,
-        dest_dir,
-        genrule_name,
-        src_files = [],
-        dest_files = []):
-    """Returns a genrule to symlink(or copy if on Windows) a set of files.
-
-    If src_dir is passed, files will be read from the given directory; otherwise
-    we assume files are in src_files and dest_files
-    """
-    if src_dir != None:
-        src_dir = _norm_path(src_dir)
-        dest_dir = _norm_path(dest_dir)
-        files = "\n".join(read_dir(repository_ctx, src_dir))
-
-        # Create a list with the src_dir stripped to use for outputs.
-        dest_files = files.replace(src_dir, "").splitlines()
-        src_files = files.splitlines()
-    command = []
-    outs = []
-    for i in range(len(dest_files)):
-        if dest_files[i] != "":
-            # If we have only one file to link we do not want to use the dest_dir, as
-            # $(@D) will include the full path to the file.
-            dest = "$(@D)/" + dest_dir + dest_files[i] if len(dest_files) != 1 else "$(@D)/" + dest_files[i]
-
-            # Copy the headers to create a sandboxable setup.
-            cmd = "cp -f"
-            command.append(cmd + ' "%s" "%s"' % (src_files[i], dest))
-            outs.append('        "' + dest_dir + dest_files[i] + '",')
-    genrule = _genrule(
-        src_dir,
-        genrule_name,
-        " && ".join(command),
-        "\n".join(outs),
-    )
-    return genrule
-
-def _get_python_lib(repository_ctx, python_bin):
-    """Gets the python lib path."""
-    python_lib = get_host_environ(repository_ctx, PYTHON_LIB_PATH)
-    if python_lib != None:
-        return python_lib
-
-    # The interesting program to execute.
-    print_lib = [
-        "from __future__ import print_function",
-        "import site",
-        "import os",
-        "python_paths = []",
-        "if os.getenv('PYTHONPATH') is not None:",
-        "  python_paths = os.getenv('PYTHONPATH').split(':')",
-        "try:",
-        "  library_paths = site.getsitepackages()",
-        "except AttributeError:",
-        "  from distutils.sysconfig import get_python_lib",
-        "  library_paths = [get_python_lib()]",
-        "all_paths = set(python_paths + library_paths)",
-        "paths = []",
-        "for path in all_paths:",
-        "  if os.path.isdir(path):",
-        "    paths.append(path)",
-        "if len(paths) >=1:",
-        "  print(paths[0])",
-    ]
-
-    # The below script writes the above program to a file
-    # and executes it. This is to work around the limitation
-    # of not being able to upload files as part of execute.
-    cmd = "from os import linesep;"
-    cmd += "f = open('script.py', 'w');"
-    for line in print_lib:
-        cmd += "f.write(\"%s\" + linesep);" % line
-    cmd += "f.close();"
-    cmd += "from subprocess import call;"
-    cmd += "call([\"%s\", \"script.py\"]);" % python_bin
-
-    result = execute(repository_ctx, [python_bin, "-c", cmd])
-    return result.stdout.strip()
-
-def _check_python_lib(repository_ctx, python_lib):
-    """Checks the python lib path."""
-    cmd = 'test -d "%s" -a -x "%s"' % (python_lib, python_lib)
-    result = raw_exec(repository_ctx, [get_bash_bin(repository_ctx), "-c", cmd])
-    if result.return_code == 1:
-        auto_config_fail("Invalid python library path: %s" % python_lib)
-
-def _check_python_bin(repository_ctx, python_bin):
-    """Checks the python bin path."""
-    cmd = '[[ -x "%s" ]] && [[ ! -d "%s" ]]' % (python_bin, python_bin)
-    result = raw_exec(repository_ctx, [get_bash_bin(repository_ctx), "-c", cmd])
-    if result.return_code == 1:
-        auto_config_fail("--define %s='%s' is not executable. Is it the python binary?" % (
-            PYTHON_BIN_PATH,
-            python_bin,
-        ))
-
-def _get_python_include(repository_ctx, python_bin):
-    """Gets the python include path."""
-    result = execute(
-        repository_ctx,
-        [
-            python_bin,
-            "-Wignore",
-            "-c",
-            "import sysconfig; " +
-            "print(sysconfig.get_path('include'))",
-        ],
-        error_msg = "Problem getting python include path.",
-        error_details = ("Is the Python binary path set up right? " +
-                         "(See ./configure or " + PYTHON_BIN_PATH + ".) " +
-                         "Is distutils installed?"),
-    )
-    return result.stdout.splitlines()[0]
-
-def _get_python_import_lib_name(repository_ctx, python_bin):
-    """Get Python import library name (pythonXY.lib) on Windows."""
-    result = execute(
-        repository_ctx,
-        [
-            python_bin,
-            "-c",
-            "import sys;" +
-            'print("python" + str(sys.version_info[0]) + ' +
-            '      str(sys.version_info[1]) + ".lib")',
-        ],
-        error_msg = "Problem getting python import library.",
-        error_details = ("Is the Python binary path set up right? " +
-                         "(See ./configure or " + PYTHON_BIN_PATH + ".) "),
-    )
-    return result.stdout.splitlines()[0]
-
-def _get_numpy_include(repository_ctx, python_bin):
-    """Gets the numpy include path."""
-    return execute(
-        repository_ctx,
-        [
-            python_bin,
-            "-c",
-            "from __future__ import print_function;" +
-            "import numpy;" +
-            " print(numpy.get_include());",
-        ],
-        error_msg = "Problem getting numpy include path.",
-        error_details = "Is numpy installed?",
-    ).stdout.splitlines()[0]
-
-def _create_local_python_repository(repository_ctx):
-    """Creates the repository containing files set up to build with Python."""
-
-    # Resolve all labels before doing any real work. Resolving causes the
-    # function to be restarted with all previous state being lost. This
-    # can easily lead to a O(n^2) runtime in the number of labels.
-    build_tpl = repository_ctx.path(Label("//third_party/py:BUILD.tpl"))
-
-    python_bin = get_python_bin(repository_ctx)
-    _check_python_bin(repository_ctx, python_bin)
-    python_lib = _get_python_lib(repository_ctx, python_bin)
-    _check_python_lib(repository_ctx, python_lib)
-    python_include = _get_python_include(repository_ctx, python_bin)
-    numpy_include = _get_numpy_include(repository_ctx, python_bin) + "/numpy"
-    python_include_rule = _symlink_genrule_for_dir(
-        repository_ctx,
-        python_include,
-        "python_include",
-        "python_include",
-    )
-    python_import_lib_genrule = ""
-
-    # To build Python C/C++ extension on Windows, we need to link to python import library pythonXY.lib
-    # See https://docs.python.org/3/extending/windows.html
-    if is_windows(repository_ctx):
-        python_bin = python_bin.replace("\\", "/")
-        python_include = _norm_path(python_include)
-        python_import_lib_name = _get_python_import_lib_name(repository_ctx, python_bin)
-        python_import_lib_src = python_include.rsplit("/", 1)[0] + "/libs/" + python_import_lib_name
-        python_import_lib_genrule = _symlink_genrule_for_dir(
-            repository_ctx,
-            None,
-            "",
-            "python_import_lib",
-            [python_import_lib_src],
-            [python_import_lib_name],
-        )
-    numpy_include_rule = _symlink_genrule_for_dir(
-        repository_ctx,
-        numpy_include,
-        "numpy_include/numpy",
-        "numpy_include",
-    )
-
-    platform_constraint = ""
-    if repository_ctx.attr.platform_constraint:
-        platform_constraint = "\"%s\"" % repository_ctx.attr.platform_constraint
-    repository_ctx.template("BUILD", build_tpl, {
-        "%{PYTHON_BIN_PATH}": python_bin,
-        "%{PYTHON_INCLUDE_GENRULE}": python_include_rule,
-        "%{PYTHON_IMPORT_LIB_GENRULE}": python_import_lib_genrule,
-        "%{NUMPY_INCLUDE_GENRULE}": numpy_include_rule,
-        "%{PLATFORM_CONSTRAINT}": platform_constraint,
-    })
-
-def _create_remote_python_repository(repository_ctx, remote_config_repo):
-    """Creates pointers to a remotely configured repo set up to build with Python.
-    """
-    repository_ctx.template("BUILD", config_repo_label(remote_config_repo, ":BUILD"), {})
-
-def _python_autoconf_impl(repository_ctx):
-    """Implementation of the python_autoconf repository rule."""
-    if get_host_environ(repository_ctx, TF_PYTHON_CONFIG_REPO) != None:
-        _create_remote_python_repository(
-            repository_ctx,
-            get_host_environ(repository_ctx, TF_PYTHON_CONFIG_REPO),
-        )
-    else:
-        _create_local_python_repository(repository_ctx)
-
-_ENVIRONS = [
-    BAZEL_SH,
-    PYTHON_BIN_PATH,
-    PYTHON_LIB_PATH,
-]
-
-local_python_configure = repository_rule(
-    implementation = _create_local_python_repository,
-    environ = _ENVIRONS,
-    attrs = {
-        "environ": attr.string_dict(),
-        "platform_constraint": attr.string(),
-    },
-)
-
-remote_python_configure = repository_rule(
-    implementation = _create_local_python_repository,
-    environ = _ENVIRONS,
-    remotable = True,
-    attrs = {
-        "environ": attr.string_dict(),
-        "platform_constraint": attr.string(),
-    },
-)
-
-python_configure = repository_rule(
-    implementation = _python_autoconf_impl,
-    environ = _ENVIRONS + [TF_PYTHON_CONFIG_REPO],
-    attrs = {
-        "platform_constraint": attr.string(),
-    },
-)
-"""Detects and configures the local Python.
-
-Add the following to your WORKSPACE FILE:
-
-```python
-python_configure(name = "local_config_python")
-```
-
-Args:
-  name: A unique name for this workspace rule.
-"""
diff --git a/third_party/xla/third_party/tsl/third_party/py/python_init_pip.bzl b/third_party/xla/third_party/tsl/third_party/py/python_init_pip.bzl
new file mode 100644
index 00000000000000..efc2bf8233cf61
--- /dev/null
+++ b/third_party/xla/third_party/tsl/third_party/py/python_init_pip.bzl
@@ -0,0 +1,34 @@
+"""Hermetic Python initialization. Consult the WORKSPACE on how to use it."""
+
+load("@python//:defs.bzl", "interpreter")
+load("@python_version_repo//:py_version.bzl", "REQUIREMENTS")
+load("@rules_python//python:pip.bzl", "package_annotation", "pip_parse")
+
+def python_init_pip():
+    numpy_annotations = {
+        "numpy": package_annotation(
+            additive_build_content = """\
+cc_library(
+    name = "numpy_headers_2",
+    hdrs = glob(["site-packages/numpy/_core/include/**/*.h"]),
+    strip_include_prefix="site-packages/numpy/_core/include/",
+)
+cc_library(
+    name = "numpy_headers_1",
+    hdrs = glob(["site-packages/numpy/core/include/**/*.h"]),
+    strip_include_prefix="site-packages/numpy/core/include/",
+)
+cc_library(
+    name = "numpy_headers",
+    deps = [":numpy_headers_2", ":numpy_headers_1"],
+)
+""",
+        ),
+    }
+
+    pip_parse(
+        name = "pypi",
+        annotations = numpy_annotations,
+        python_interpreter_target = interpreter,
+        requirements_lock = REQUIREMENTS,
+    )
diff --git a/third_party/xla/third_party/tsl/third_party/py/python_init_repositories.bzl b/third_party/xla/third_party/tsl/third_party/py/python_init_repositories.bzl
new file mode 100644
index 00000000000000..5a405f2c2aba4c
--- /dev/null
+++ b/third_party/xla/third_party/tsl/third_party/py/python_init_repositories.bzl
@@ -0,0 +1,12 @@
+"""Hermetic Python initialization. Consult the WORKSPACE on how to use it."""
+
+load("@rules_python//python:repositories.bzl", "py_repositories")
+load("//third_party/py:python_repo.bzl", "python_repository")
+
+def python_init_repositories(requirements = {}):
+    python_repository(
+        name = "python_version_repo",
+        requirements_versions = requirements.keys(),
+        requirements_locks = requirements.values(),
+    )
+    py_repositories()
diff --git a/third_party/xla/third_party/tsl/third_party/py/python_init_rules.bzl b/third_party/xla/third_party/tsl/third_party/py/python_init_rules.bzl
new file mode 100644
index 00000000000000..98a7b8bc3c315a
--- /dev/null
+++ b/third_party/xla/third_party/tsl/third_party/py/python_init_rules.bzl
@@ -0,0 +1,11 @@
+"""Hermetic Python initialization. Consult the WORKSPACE on how to use it."""
+
+load("@bazel_tools//tools/build_defs/repo:http.bzl", "http_archive")
+
+def python_init_rules():
+    http_archive(
+        name = "rules_python",
+        sha256 = "9d04041ac92a0985e344235f5d946f71ac543f1b1565f2cdbc9a2aaee8adf55b",
+        strip_prefix = "rules_python-0.26.0",
+        url = "https://github.com/bazelbuild/rules_python/releases/download/0.26.0/rules_python-0.26.0.tar.gz",
+    )
diff --git a/third_party/xla/third_party/tsl/third_party/py/python_init_toolchains.bzl b/third_party/xla/third_party/tsl/third_party/py/python_init_toolchains.bzl
new file mode 100644
index 00000000000000..c1f800db4c01e7
--- /dev/null
+++ b/third_party/xla/third_party/tsl/third_party/py/python_init_toolchains.bzl
@@ -0,0 +1,13 @@
+"""Hermetic Python initialization. Consult the WORKSPACE on how to use it."""
+
+load("@python_version_repo//:py_version.bzl", "HERMETIC_PYTHON_VERSION")
+load("@rules_python//python:repositories.bzl", "python_register_toolchains")
+load("@rules_python//python:versions.bzl", "MINOR_MAPPING")
+
+def python_init_toolchains():
+    if HERMETIC_PYTHON_VERSION in MINOR_MAPPING:
+        python_register_toolchains(
+            name = "python",
+            ignore_root_user_error = True,
+            python_version = HERMETIC_PYTHON_VERSION,
+        )
diff --git a/third_party/xla/third_party/tsl/third_party/py/python_repo.bzl b/third_party/xla/third_party/tsl/third_party/py/python_repo.bzl
new file mode 100644
index 00000000000000..77a6ce9ce50b60
--- /dev/null
+++ b/third_party/xla/third_party/tsl/third_party/py/python_repo.bzl
@@ -0,0 +1,206 @@
+"""
+Repository rule to manage hermetic Python interpreter under Bazel.
+
+Version can be set via build parameter "--repo_env=HERMETIC_PYTHON_VERSION=3.11"
+Defaults to 3.11.
+
+To set wheel name, add "--repo_env=WHEEL_NAME=tensorflow_cpu"
+"""
+
+VERSIONS = ["3.9", "3.10", "3.11", "3.12", "3.13"]
+DEFAULT_VERSION = "3.11"
+WARNING = """
+HERMETIC_PYTHON_VERSION variable was not set correctly, using default version. 
+Python {} will be used.
+To select Python version, either set HERMETIC_PYTHON_VERSION env variable in 
+your shell:
+  export HERMETIC_PYTHON_VERSION=3.12
+OR pass it as an argument to bazel command directly or inside your .bazelrc 
+file:
+  --repo_env=HERMETIC_PYTHON_VERSION=3.12
+""".format(DEFAULT_VERSION)
+
+content = """TF_PYTHON_VERSION = "{version}"
+HERMETIC_PYTHON_VERSION = "{version}"
+WHEEL_NAME = "{wheel_name}"
+WHEEL_COLLAB = "{wheel_collab}"
+REQUIREMENTS = "{requirements}"
+"""
+
+def _python_repository_impl(ctx):
+    ctx.file("BUILD", "")
+    version_legacy = ctx.os.environ.get("TF_PYTHON_VERSION", "")
+    version = ctx.os.environ.get("HERMETIC_PYTHON_VERSION", "")
+    if not version:
+        version = version_legacy
+    else:
+        version_legacy = version
+
+    wheel_name = ctx.os.environ.get("WHEEL_NAME", "tensorflow")
+    wheel_collab = ctx.os.environ.get("WHEEL_COLLAB", False)
+    if version not in VERSIONS:
+        print(WARNING)  # buildifier: disable=print
+        version = DEFAULT_VERSION
+    else:
+        print("Using hermetic Python %s" % version)  # buildifier: disable=print
+
+    requirements = ""
+    for i in range(0, len(ctx.attr.requirements_locks)):
+        if ctx.attr.requirements_versions[i] == version:
+            requirements = ctx.attr.requirements_locks[i]
+            break
+
+    ctx.file(
+        "py_version.bzl",
+        content.format(
+            version = version,
+            wheel_name = wheel_name,
+            wheel_collab = wheel_collab,
+            requirements = str(requirements),
+        ),
+    )
+
+python_repository = repository_rule(
+    implementation = _python_repository_impl,
+    attrs = {
+        "requirements_versions": attr.string_list(
+            mandatory = False,
+            default = [],
+        ),
+        "requirements_locks": attr.label_list(
+            mandatory = False,
+            default = [],
+        ),
+    },
+    environ = [
+        "TF_PYTHON_VERSION",
+        "HERMETIC_PYTHON_VERSION",
+        "WHEEL_NAME",
+        "WHEEL_COLLAB",
+    ],
+)
+
+def _custom_python_interpreter_impl(ctx):
+    version = ctx.attr.version
+    strip_prefix = ctx.attr.strip_prefix.format(version = version)
+    urls = [url.format(version = version) for url in ctx.attr.urls]
+    binary_name = ctx.attr.binary_name
+    if not binary_name:
+        ver_chunks = version.split(".")
+        binary_name = "python%s.%s" % (ver_chunks[0], ver_chunks[1])
+
+    install_dir = "{name}-{version}".format(name = ctx.attr.name, version = version)
+    _exec_and_check(ctx, ["mkdir", install_dir])
+    install_path = ctx.path(install_dir)
+    srcs_dir = "srcs"
+    ctx.download_and_extract(
+        url = urls,
+        stripPrefix = strip_prefix,
+        output = srcs_dir,
+    )
+
+    configure_params = []
+    if "CC" in ctx.os.environ:
+        configure_params.append("CC={}".format(ctx.os.environ["CC"]))
+    if "CXX" in ctx.os.environ:
+        configure_params.append("CXX={}".format(ctx.os.environ["CXX"]))
+
+    configure_params.append("--enable-optimizations")
+    configure_params.append("--prefix=%s" % install_path.realpath)
+    _exec_and_check(
+        ctx,
+        ["./configure"] + configure_params,
+        working_directory = srcs_dir,
+        quiet = False,
+    )
+    res = _exec_and_check(ctx, ["nproc"])
+    cores = 12 if res.return_code != 0 else max(1, int(res.stdout.strip()) - 1)
+    _exec_and_check(ctx, ["make", "-j%s" % cores], working_directory = srcs_dir)
+    _exec_and_check(ctx, ["make", "altinstall"], working_directory = srcs_dir)
+    _exec_and_check(ctx, ["ln", "-s", binary_name, "python3"], working_directory = install_dir + "/bin")
+    tar = "{install_dir}.tgz".format(install_dir = install_dir)
+    _exec_and_check(ctx, ["tar", "czpf", tar, install_dir])
+    _exec_and_check(ctx, ["rm", "-rf", srcs_dir])
+    res = _exec_and_check(ctx, ["sha256sum", tar])
+
+    sha256 = res.stdout.split(" ")[0].strip()
+    tar_path = ctx.path(tar)
+
+    example = """\n\n
+To use newly built Python interpreter add the following code snippet RIGHT AFTER
+python_init_toolchains() in your WORKSPACE file. The code sample should work as
+is but it may need some tuning, if you have special requirements.
+
+```
+load("@rules_python//python:repositories.bzl", "python_register_toolchains")
+python_register_toolchains(
+    name = "python",
+    # By default assume the interpreter is on the local file system, replace
+    # with proper URL if it is not the case.
+    base_url = "file://",
+    ignore_root_user_error = True,
+    python_version = "{version}",
+    tool_versions = {{
+        "{version}": {{
+            # Path to .tar.gz with Python binary. By default it points to .tgz
+            # file in cache where it was built originally; replace with proper
+            # file location, if you moved it somewhere else.
+            "url": "{tar_path}",
+            "sha256": {{
+                # By default we assume Linux x86_64 architecture, eplace with
+                # proper architecture if you were building on a different platform.
+                "x86_64-unknown-linux-gnu": "{sha256}",
+            }},
+            "strip_prefix": "{install_dir}",
+        }},
+    }},
+)
+```
+\n\n""".format(version = version, tar_path = tar_path, sha256 = sha256, install_dir = install_dir)
+
+    instructions = "INSTRUCTIONS-{version}.md".format(version = version)
+    ctx.file(instructions + ".tmpl", example, executable = False)
+    ctx.file(
+        "BUILD.bazel",
+        """
+genrule(
+    name = "{name}",
+    srcs = ["{tar}", "{instructions}.tmpl"],
+    outs = ["{install_dir}.tar.gz", "{instructions}"],
+    cmd = "cp $(location {tar}) $(location {install_dir}.tar.gz); cp $(location {instructions}.tmpl) $(location {instructions})",
+    visibility = ["//visibility:public"],
+)
+     """.format(
+            name = ctx.attr.name,
+            tar = tar,
+            install_dir = install_dir,
+            instructions = instructions,
+        ),
+        executable = False,
+    )
+
+    print(example)  # buildifier: disable=print
+
+custom_python_interpreter = repository_rule(
+    implementation = _custom_python_interpreter_impl,
+    attrs = {
+        "urls": attr.string_list(),
+        "strip_prefix": attr.string(),
+        "binary_name": attr.string(mandatory = False),
+        "version": attr.string(),
+    },
+)
+
+def _exec_and_check(ctx, command, fail_on_error = True, quiet = False, **kwargs):
+    res = ctx.execute(command, quiet = quiet, **kwargs)
+    if fail_on_error and res.return_code != 0:
+        fail("""
+Failed to execute command: `{command}`
+Exit Code: {code}
+STDERR: {stderr}
+        """.format(
+            command = command,
+            code = res.return_code,
+            stderr = res.stderr,
+        ))
+    return res
diff --git a/third_party/xla/third_party/tsl/workspace2.bzl b/third_party/xla/third_party/tsl/workspace2.bzl
index 001ce018d87066..67011d3d2bb3ea 100644
--- a/third_party/xla/third_party/tsl/workspace2.bzl
+++ b/third_party/xla/third_party/tsl/workspace2.bzl
@@ -505,12 +505,6 @@ def _tf_repositories():
         urls = tf_mirror_urls("https://github.com/google/double-conversion/archive/v3.2.0.tar.gz"),
     )
 
-    tf_http_archive(
-        name = "rules_python",
-        sha256 = "aa96a691d3a8177f3215b14b0edc9641787abaaa30363a080165d06ab65e1161",
-        urls = tf_mirror_urls("https://github.com/bazelbuild/rules_python/releases/download/0.0.1/rules_python-0.0.1.tar.gz"),
-    )
-
     tf_http_archive(
         name = "build_bazel_rules_android",
         sha256 = "cd06d15dd8bb59926e4d65f9003bfc20f9da4b2519985c27e190cddc8b7a7806",
@@ -597,6 +591,7 @@ def _tf_repositories():
         urls = tf_mirror_urls("https://github.com/google/glog/archive/refs/tags/v0.4.0.tar.gz"),
     )
 
+# buildifier: disable=unnamed-macro
 def workspace():
     # Check the bazel version before executing any repository rules, in case
     # those rules rely on the version we require here.
diff --git a/third_party/xla/xla/lit.bzl b/third_party/xla/xla/lit.bzl
index 57d25fff261242..962baf4181de6f 100644
--- a/third_party/xla/xla/lit.bzl
+++ b/third_party/xla/xla/lit.bzl
@@ -178,10 +178,25 @@ def lit_test(
         visibility = ["//visibility:private"],
         **kwargs
     )
+    lit_name = "//third_party/py/lit:lit"
+
+    # copybara:comment_begin(oss-only)
+    lit_name = "lit_custom_" + name
+    native.py_binary(
+        name = lit_name,
+        main = "@llvm-project//llvm:utils/lit/lit.py",
+        srcs = ["@llvm-project//llvm:utils/lit/lit.py"],
+        testonly = True,
+        deps = [
+            "@llvm-project//llvm:lit_lib",
+            "@pypi_lit//:pkg",
+        ],
+    )
 
+    # copybara:comment_end
     native_test(
         name = name,
-        src = "@llvm-project//llvm:lit",
+        src = lit_name,
         args = [
             "-a",
             "--path",
@@ -189,14 +204,14 @@ def lit_test(
             "$(location {})".format(test_file),
         ] + args,
         data = [
-            "@llvm-project//llvm:lit",
+            lit_name,
             test_file,
 
             # TODO(cheshire): Config is not passed properly when it's not
             # called lit.cfg.py
             cfg,
             tools_on_path_target_name,
-        ] + data,
+        ] + data + ["@pypi_lit//:pkg"],
         visibility = visibility,
         env = env,
         timeout = timeout,
diff --git a/third_party/xla/xla/mlir_hlo/tests/BUILD b/third_party/xla/xla/mlir_hlo/tests/BUILD
index 4db81efeb2dd52..08df267d05a427 100644
--- a/third_party/xla/xla/mlir_hlo/tests/BUILD
+++ b/third_party/xla/xla/mlir_hlo/tests/BUILD
@@ -27,6 +27,7 @@ package(
             "nomsan",  # The execution engine doesn't work with msan, see b/248097619.
         ],
         deps = ["@pypi_lit//:pkg"],
+        deps = ["@pypi_lit//:pkg"],
     )
     for src in glob(["**/*.mlir"])
 ]

From abca74fa7e9f826bc9aae5c5858df724e539d3ff Mon Sep 17 00:00:00 2001
From: Kyle Lucke <klucke@google.com>
Date: Wed, 15 May 2024 15:29:59 -0700
Subject: [PATCH 224/478] Use StreamExecutorInterface::CreateEvent in
 event_pool.cc.

PiperOrigin-RevId: 634095995
---
 third_party/xla/xla/pjrt/event_pool.cc | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/third_party/xla/xla/pjrt/event_pool.cc b/third_party/xla/xla/pjrt/event_pool.cc
index a5e2f83c3e9d87..fadfe8bd56ee79 100644
--- a/third_party/xla/xla/pjrt/event_pool.cc
+++ b/third_party/xla/xla/pjrt/event_pool.cc
@@ -46,8 +46,7 @@ absl::StatusOr<EventPool::Handle> EventPool::AllocateEvent(
     }
   }
   if (!event.event_) {
-    event.event_ = std::make_unique<se::Event>(executor);
-    TF_RET_CHECK(event.event_->Init()) << "Event initialization failed";
+    TF_ASSIGN_OR_RETURN(event.event_, executor->CreateEvent());
   }
   return event;
 }

From c031655ca9514e7946bb87eda17cf5932c381a8b Mon Sep 17 00:00:00 2001
From: Luke Boyer <lukeboyer@google.com>
Date: Wed, 15 May 2024 15:52:18 -0700
Subject: [PATCH 225/478] Add pattern to fuse quant weights into transpose
 conv.

PiperOrigin-RevId: 634102361
---
 tensorflow/compiler/mlir/lite/tests/optimize.mlir   | 13 +++++++++++++
 .../mlir/lite/transforms/optimize_patterns.td       | 10 ++++++++++
 2 files changed, 23 insertions(+)

diff --git a/tensorflow/compiler/mlir/lite/tests/optimize.mlir b/tensorflow/compiler/mlir/lite/tests/optimize.mlir
index 3c2c24baba8972..96cffcd1371473 100644
--- a/tensorflow/compiler/mlir/lite/tests/optimize.mlir
+++ b/tensorflow/compiler/mlir/lite/tests/optimize.mlir
@@ -103,6 +103,19 @@ func.func @fuseSubIntoConv2d(%arg0: tensor<256x32x32x3xf32>, %arg1: tensor<16x3x
   // CHECK: %0 = "tfl.conv_2d"(%arg0, %arg1, %cst)
 }
 
+// CHECK-LABEL: foldQuantWeightsIntoTposeConv
+func.func @foldQuantWeightsIntoTposeConv(%arg0: tensor<2x2x3x2048xf32>) -> tensor<2x3x2x2048xf32> {
+  %output_shape = arith.constant dense<[2, 3, 2, 2048]> : tensor<4xi32>
+  %q_weighs = "tfl.pseudo_qconst"() {qtype = tensor<4x2x2x2048x!quant.uniform<u8<1:255>:f32, 0.15:151>>, value = dense<-76> : tensor<4x2x2x2048xi8>} : () -> tensor<4x2x2x2048x!quant.uniform<u8<1:255>:f32, 0.15:151>>
+  %dq_weights = "tfl.dequantize"(%q_weighs) : (tensor<4x2x2x2048x!quant.uniform<u8<1:255>:f32, 0.15:151>>) -> tensor<4x2x2x2048xf32>
+  %bias = "tfl.no_value"() {value} : () -> none
+  %out = "tfl.transpose_conv"(%output_shape, %dq_weights, %arg0, %bias) {fused_activation_function = "NONE", padding = "SAME", stride_h = 1 : i32, stride_w = 1 : i32} : (tensor<4xi32>, tensor<4x2x2x2048xf32>, tensor<2x2x3x2048xf32>, none) -> tensor<2x3x2x2048xf32>
+  func.return %out : tensor<2x3x2x2048xf32>
+
+  // CHECK-NOT: "tfl.dequantize"
+  // CHECK: "tfl.transpose_conv"(%cst, %1, %arg0, %0) <{fused_activation_function = "NONE", padding = "SAME", stride_h = 1 : i32, stride_w = 1 : i32}> : (tensor<4xi32>, tensor<4x2x2x2048x!quant.uniform<u8<1:255>:f32
+}
+
 // CHECK-LABEL: fuseAddIntoTransposeConv
 func.func @fuseAddIntoTransposeConv(%arg0: tensor<1x32x42x128xf32>) -> tensor<1x64x84x32xf32> {
   %cst = arith.constant dense<1.5> : tensor<32xf32>
diff --git a/tensorflow/compiler/mlir/lite/transforms/optimize_patterns.td b/tensorflow/compiler/mlir/lite/transforms/optimize_patterns.td
index 4353b82e2fb901..d3da883d0d8a02 100644
--- a/tensorflow/compiler/mlir/lite/transforms/optimize_patterns.td
+++ b/tensorflow/compiler/mlir/lite/transforms/optimize_patterns.td
@@ -813,6 +813,16 @@ def MatchReluPattern : Pat<
   (TFL_ReluOp $input),
   [(FloatValueEquals<"0"> $Zero)]>;
 
+// Transpose conv supports hybrid computation with quantized weights.
+def FoldQuantWeightsIntoTposeConv : Pat<
+  (TFL_TransposeConvOp
+    $output_shape,
+    (TFL_DequantizeOp $quant_weights),
+    $quant_input,
+    $bias, $padding, $stride_h, $stride_w, $faf),
+  (TFL_TransposeConvOp $output_shape, $quant_weights,
+    $quant_input, $bias, $padding, $stride_h, $stride_w, $faf)>;
+
 // Optimize Minimum of tf.Relu and constant six to tf.Relu6
 def MinimumOfReluAnd6ToRelu6 :
   Pat<(TFL_MinimumOp (TFL_ReluOp $x), (Arith_ConstantOp $y)),

From c62f485a70c2906af3de5845fe48f4b20bb24e94 Mon Sep 17 00:00:00 2001
From: Kyle Lucke <klucke@google.com>
Date: Wed, 15 May 2024 16:00:29 -0700
Subject: [PATCH 226/478] Use absl::Status instead of xla::Status now that
 they're identical.

PiperOrigin-RevId: 634104645
---
 .../llvm_ir/dynamic_update_slice_util.cc      | 23 ++++---
 .../llvm_ir/dynamic_update_slice_util.h       | 16 ++---
 .../service/llvm_ir/kernel_support_library.cc | 16 ++---
 .../service/llvm_ir/kernel_support_library.h  | 65 ++++++++++---------
 .../xla/xla/service/llvm_ir/llvm_util.cc      |  6 +-
 .../xla/xla/service/llvm_ir/loop_emitter.cc   |  6 +-
 .../xla/xla/service/llvm_ir/loop_emitter.h    |  6 +-
 .../xla/xla/service/llvm_ir/sort_util.cc      |  8 +--
 .../xla/xla/service/llvm_ir/sort_util.h       |  4 +-
 .../dumped_computation_to_operation_list.cc   |  2 +-
 .../tools/extract_collective_operations.cc    |  4 +-
 .../xla/tools/hlo_control_flow_flattening.cc  |  8 +--
 .../xla/tools/hlo_control_flow_flattening.h   | 10 +--
 third_party/xla/xla/tools/hlo_extractor.cc    | 14 ++--
 .../xla/xla/tools/prepare_reference_module.cc |  4 +-
 .../xla/xla/tools/prepare_reference_module.h  |  4 +-
 third_party/xla/xla/tools/run_hlo_module.cc   | 31 +++++----
 third_party/xla/xla/tools/run_hlo_module.h    | 13 ++--
 .../hlo_to_mhlo/hlo_function_importer.cc      | 14 ++--
 .../hlo_to_mhlo/hlo_function_importer.h       | 25 +++----
 .../hlo_to_mhlo/hlo_module_importer.cc        |  4 +-
 .../hlo_to_mhlo/hlo_module_importer.h         |  4 +-
 .../translate/hlo_to_mhlo/hlo_to_mlir_hlo.cc  | 15 +++--
 .../translate/hlo_to_mhlo/hlo_to_mlir_hlo.h   | 15 +++--
 .../xla/xla/translate/hlo_to_mhlo/hlo_utils.h |  4 +-
 25 files changed, 166 insertions(+), 155 deletions(-)

diff --git a/third_party/xla/xla/service/llvm_ir/dynamic_update_slice_util.cc b/third_party/xla/xla/service/llvm_ir/dynamic_update_slice_util.cc
index 6b17b942364799..b7b802fb0a6d0e 100644
--- a/third_party/xla/xla/service/llvm_ir/dynamic_update_slice_util.cc
+++ b/third_party/xla/xla/service/llvm_ir/dynamic_update_slice_util.cc
@@ -100,7 +100,7 @@ bool CanEmitFusedDynamicUpdateSliceInPlace(HloInstruction* fusion,
 // Emits a sequential loop if launch_dimensions is null.
 using IndexGenerator = std::function<absl::StatusOr<llvm::Value*>(int64_t)>;
 
-static Status EmitDynamicUpdateSliceInPlaceImpl(
+static absl::Status EmitDynamicUpdateSliceInPlaceImpl(
     const Shape& update_shape, const IndexGenerator& start_indices_generator,
     bool is_signed, ElementGenerator update_array_generator,
     const IrArray& output_array, const gpu::LaunchDimensions* launch_dimensions,
@@ -135,7 +135,8 @@ static Status EmitDynamicUpdateSliceInPlaceImpl(
                         max_bound, start_multi_index[i]);
   }
 
-  auto loop_body_emitter = [&](const IrArray::Index& update_index) -> Status {
+  auto loop_body_emitter =
+      [&](const IrArray::Index& update_index) -> absl::Status {
     // Calculate output_index, where we'll write the value from update.  For
     // each dimension,
     //
@@ -165,10 +166,9 @@ static Status EmitDynamicUpdateSliceInPlaceImpl(
   return LoopEmitter(loop_body_emitter, update_shape, b).EmitLoop(name);
 }
 
-Status EmitDynamicUpdateSliceInPlace(absl::Span<const IrArray> operand_arrays,
-                                     const IrArray& output_array,
-                                     absl::string_view name,
-                                     llvm::IRBuilder<>* b) {
+absl::Status EmitDynamicUpdateSliceInPlace(
+    absl::Span<const IrArray> operand_arrays, const IrArray& output_array,
+    absl::string_view name, llvm::IRBuilder<>* b) {
   VLOG(2) << "EmitDynamicUpdateSliceInPlace for " << name;
 
   // No need to use operand_arrays[0], the input array of the
@@ -196,7 +196,7 @@ Status EmitDynamicUpdateSliceInPlace(absl::Span<const IrArray> operand_arrays,
 // EmitParallelFusedDynamicUpdateSliceInPlace.
 //
 // Emits a sequential loop if launch_dimensions is null.
-static Status EmitFusedDynamicUpdateSliceInPlaceImpl(
+static absl::Status EmitFusedDynamicUpdateSliceInPlaceImpl(
     const HloComputation* fusion,
     const std::vector<std::pair<const HloInstruction*, const IrArray>>&
         dus_and_output_array,
@@ -252,10 +252,9 @@ static Status EmitFusedDynamicUpdateSliceInPlaceImpl(
   return OkStatus();
 }
 
-Status EmitFusedDynamicUpdateSliceInPlace(HloInstruction* fusion,
-                                          const IrArray& fusion_output_array,
-                                          FusedIrEmitter* fused_emitter,
-                                          llvm::IRBuilder<>* b) {
+absl::Status EmitFusedDynamicUpdateSliceInPlace(
+    HloInstruction* fusion, const IrArray& fusion_output_array,
+    FusedIrEmitter* fused_emitter, llvm::IRBuilder<>* b) {
   HloInstruction* dus = fusion->called_computations()[0]->root_instruction();
   CHECK_EQ(dus->opcode(), HloOpcode::kDynamicUpdateSlice);
   std::vector<std::pair<const HloInstruction*, const IrArray>>
@@ -266,7 +265,7 @@ Status EmitFusedDynamicUpdateSliceInPlace(HloInstruction* fusion,
       /*launch_dimensions=*/nullptr, b);
 }
 
-Status EmitParallelFusedDynamicUpdateSliceInPlace(
+absl::Status EmitParallelFusedDynamicUpdateSliceInPlace(
     const HloComputation* fusion,
     const std::vector<std::pair<const HloInstruction*, const IrArray>>&
         dus_and_output_array,
diff --git a/third_party/xla/xla/service/llvm_ir/dynamic_update_slice_util.h b/third_party/xla/xla/service/llvm_ir/dynamic_update_slice_util.h
index cf015c8692226e..ea70465d0b5207 100644
--- a/third_party/xla/xla/service/llvm_ir/dynamic_update_slice_util.h
+++ b/third_party/xla/xla/service/llvm_ir/dynamic_update_slice_util.h
@@ -64,24 +64,22 @@ bool CanEmitFusedDynamicUpdateSliceInPlace(HloInstruction* fusion,
 // Emits IR for running the given dynamic-update-slice op in-place -- that is,
 // where the input and output buffers share the same slice, so we can simply
 // modify the input/output buffer without touching any of the other elements.
-Status EmitDynamicUpdateSliceInPlace(absl::Span<const IrArray> operand_arrays,
-                                     const IrArray& output_array,
-                                     absl::string_view name,
-                                     llvm::IRBuilder<>* b);
+absl::Status EmitDynamicUpdateSliceInPlace(
+    absl::Span<const IrArray> operand_arrays, const IrArray& output_array,
+    absl::string_view name, llvm::IRBuilder<>* b);
 
 // Given a loop-fusion node whose root is a dynamic-update-slice op whose
 // array-to-be-updated and output share the same buffer slice, emits
 // (sequential) code for a fusion node that does the dynamic-update-slice in
 // place.
-Status EmitFusedDynamicUpdateSliceInPlace(HloInstruction* fusion,
-                                          const IrArray& fusion_output_array,
-                                          FusedIrEmitter* fused_emitter,
-                                          llvm::IRBuilder<>* b);
+absl::Status EmitFusedDynamicUpdateSliceInPlace(
+    HloInstruction* fusion, const IrArray& fusion_output_array,
+    FusedIrEmitter* fused_emitter, llvm::IRBuilder<>* b);
 
 // Same as EmitFusedDynamicUpdateSliceInPlace, except emits a parallel loop with
 // the given launch dimensions for arbitrarily many independent dynamic slice
 // updates.
-Status EmitParallelFusedDynamicUpdateSliceInPlace(
+absl::Status EmitParallelFusedDynamicUpdateSliceInPlace(
     const HloComputation* fusion,
     const std::vector<std::pair<const HloInstruction*, const IrArray>>&
         dus_and_output_array,
diff --git a/third_party/xla/xla/service/llvm_ir/kernel_support_library.cc b/third_party/xla/xla/service/llvm_ir/kernel_support_library.cc
index 71ce2d7a4495e6..782d308417a765 100644
--- a/third_party/xla/xla/service/llvm_ir/kernel_support_library.cc
+++ b/third_party/xla/xla/service/llvm_ir/kernel_support_library.cc
@@ -19,11 +19,11 @@ limitations under the License.
 #include "xla/service/llvm_ir/llvm_util.h"
 
 namespace xla {
-Status KernelSupportLibrary::ForWithStatus(
+absl::Status KernelSupportLibrary::ForWithStatus(
     absl::string_view name, llvm::Value* start, llvm::Value* end,
     llvm::Value* step,
-    const std::function<Status(llvm::Value*, bool)>& for_body_generator) {
-  return IfWithStatus(b_->CreateICmpSLT(start, end), [&]() -> Status {
+    const std::function<absl::Status(llvm::Value*, bool)>& for_body_generator) {
+  return IfWithStatus(b_->CreateICmpSLT(start, end), [&]() -> absl::Status {
     TF_RETURN_IF_ERROR(for_body_generator(start, /*is_first_iteration=*/true));
     return ForWithStatus(
         name, b_->CreateAdd(start, step), end, step,
@@ -31,10 +31,10 @@ Status KernelSupportLibrary::ForWithStatus(
   });
 }
 
-Status KernelSupportLibrary::ForWithStatus(
+absl::Status KernelSupportLibrary::ForWithStatus(
     absl::string_view name, llvm::Value* start, llvm::Value* end,
     llvm::Value* step,
-    const std::function<Status(llvm::Value*)>& for_body_generator) {
+    const std::function<absl::Status(llvm::Value*)>& for_body_generator) {
   std::unique_ptr<llvm_ir::ForLoop> loop = llvm_ir::ForLoop::EmitForLoop(
       name, start, end, step, b_,
       /*unroll_mode=*/unroll_mode_,
@@ -45,10 +45,10 @@ Status KernelSupportLibrary::ForWithStatus(
   return OkStatus();
 }
 
-Status KernelSupportLibrary::IfWithStatus(
+absl::Status KernelSupportLibrary::IfWithStatus(
     absl::string_view name, llvm::Value* condition,
-    const std::function<Status()>& true_block_generator,
-    const std::function<Status()>& false_block_generator) {
+    const std::function<absl::Status()>& true_block_generator,
+    const std::function<absl::Status()>& false_block_generator) {
   llvm_ir::LlvmIfData if_data =
       llvm_ir::EmitIfThenElse(condition, name, b_,
                               /*emit_else=*/false_block_generator != nullptr);
diff --git a/third_party/xla/xla/service/llvm_ir/kernel_support_library.h b/third_party/xla/xla/service/llvm_ir/kernel_support_library.h
index 2864101aec0955..3580e76f2512a4 100644
--- a/third_party/xla/xla/service/llvm_ir/kernel_support_library.h
+++ b/third_party/xla/xla/service/llvm_ir/kernel_support_library.h
@@ -48,11 +48,11 @@ class KernelSupportLibrary {
   //     for (i64 i = `start` + `step`; i < `end`; i += `step`)
   //       `for_body_generator(/*ind_var=*/i, /*is_first_iteration=*/false)`;
   //   }
-  Status ForWithStatus(
+  absl::Status ForWithStatus(
       absl::string_view name, llvm::Value* start, llvm::Value* end,
       llvm::Value* step,
-      const std::function<Status(llvm::Value* ind_var,
-                                 bool is_first_iteration)>& for_body_generator);
+      const std::function<absl::Status(
+          llvm::Value* ind_var, bool is_first_iteration)>& for_body_generator);
 
   void For(
       absl::string_view name, llvm::Value* start, llvm::Value* end,
@@ -60,17 +60,17 @@ class KernelSupportLibrary {
       const std::function<void(llvm::Value* ind_var, bool is_first_iteration)>&
           for_body_generator) {
     CHECK_EQ(OkStatus(),
-             ForWithStatus(
-                 name, start, end, step,
-                 [&](llvm::Value* ind_var, bool is_first_iteration) -> Status {
-                   for_body_generator(ind_var, is_first_iteration);
-                   return OkStatus();
-                 }));
+             ForWithStatus(name, start, end, step,
+                           [&](llvm::Value* ind_var,
+                               bool is_first_iteration) -> absl::Status {
+                             for_body_generator(ind_var, is_first_iteration);
+                             return OkStatus();
+                           }));
   }
 
-  Status ForWithStatus(
+  absl::Status ForWithStatus(
       absl::string_view name, int64_t start, int64_t end, int64_t step,
-      const std::function<Status(
+      const std::function<absl::Status(
           llvm::Value* ind_var, bool is_first_iteration)>& for_body_generator) {
     return ForWithStatus(name, /*start=*/b_->getInt64(start),
                          /*end=*/b_->getInt64(end),
@@ -90,29 +90,32 @@ class KernelSupportLibrary {
   //
   //   for (i64 i = `start`; i < `end`; i += `step`)
   //     `for_body_generator(/*ind_var=*/i)`;
-  Status ForWithStatus(
+  absl::Status ForWithStatus(
       absl::string_view name, llvm::Value* start, llvm::Value* end,
       llvm::Value* step,
-      const std::function<Status(llvm::Value* ind_var)>& for_body_generator);
+      const std::function<absl::Status(llvm::Value* ind_var)>&
+          for_body_generator);
 
   void For(
       absl::string_view name, llvm::Value* start, llvm::Value* end,
       llvm::Value* step,
       const std::function<void(llvm::Value* ind_var)>& for_body_generator) {
-    CHECK_EQ(OkStatus(), ForWithStatus(name, start, end, step,
-                                       [&](llvm::Value* ind_var) -> Status {
-                                         for_body_generator(ind_var);
-                                         return OkStatus();
-                                       }));
+    CHECK_EQ(OkStatus(),
+             ForWithStatus(name, start, end, step,
+                           [&](llvm::Value* ind_var) -> absl::Status {
+                             for_body_generator(ind_var);
+                             return OkStatus();
+                           }));
   }
 
-  Status ForWithStatus(
+  absl::Status ForWithStatus(
       absl::string_view name, llvm::Value* start, llvm::Value* end,
       int64_t step,
-      const std::function<Status(llvm::Value* ind_var)>& for_body_generator) {
+      const std::function<absl::Status(llvm::Value* ind_var)>&
+          for_body_generator) {
     return ForWithStatus(name, start, end,
                          llvm::ConstantInt::get(start->getType(), step),
-                         [&](llvm::Value* indvar) -> Status {
+                         [&](llvm::Value* indvar) -> absl::Status {
                            return for_body_generator(indvar);
                          });
   }
@@ -125,9 +128,10 @@ class KernelSupportLibrary {
         for_body_generator);
   }
 
-  Status ForWithStatus(
+  absl::Status ForWithStatus(
       absl::string_view name, int64_t start, int64_t end, int64_t step,
-      const std::function<Status(llvm::Value* ind_var)>& for_body_generator) {
+      const std::function<absl::Status(llvm::Value* ind_var)>&
+          for_body_generator) {
     return ForWithStatus(name, /*start=*/b_->getInt64(start),
                          /*end=*/b_->getInt64(end),
                          /*step=*/b_->getInt64(step), for_body_generator);
@@ -148,17 +152,16 @@ class KernelSupportLibrary {
   //   else
   //      `false_block_generator()`;
   // The else is skipped if false_block_generator is null.
-  Status IfWithStatus(
+  absl::Status IfWithStatus(
       absl::string_view name, llvm::Value* condition,
-      const std::function<Status()>& true_block_generator,
-      const std::function<Status()>& false_block_generator = nullptr);
+      const std::function<absl::Status()>& true_block_generator,
+      const std::function<absl::Status()>& false_block_generator = nullptr);
 
-  Status IfWithStatus(
+  absl::Status IfWithStatus(
       llvm::Value* condition,
-      const std::function<Status()>& true_block_generator,
-      const std::function<Status()>& false_block_generator = []() -> Status {
-        return OkStatus();
-      }) {
+      const std::function<absl::Status()>& true_block_generator,
+      const std::function<absl::Status()>& false_block_generator =
+          []() -> absl::Status { return OkStatus(); }) {
     return IfWithStatus("", condition, true_block_generator,
                         false_block_generator);
   }
diff --git a/third_party/xla/xla/service/llvm_ir/llvm_util.cc b/third_party/xla/xla/service/llvm_ir/llvm_util.cc
index 88df89021d3db4..498f224a13758e 100644
--- a/third_party/xla/xla/service/llvm_ir/llvm_util.cc
+++ b/third_party/xla/xla/service/llvm_ir/llvm_util.cc
@@ -709,9 +709,9 @@ std::map<int, llvm::MDNode*> MergeMetadata(
   return result;
 }
 
-static Status CreateAndWriteStringToFile(const std::string& directory_name,
-                                         const std::string& file_name,
-                                         const std::string& text) {
+static absl::Status CreateAndWriteStringToFile(
+    const std::string& directory_name, const std::string& file_name,
+    const std::string& text) {
   std::unique_ptr<tsl::WritableFile> f;
   TF_RETURN_IF_ERROR(tsl::Env::Default()->RecursivelyCreateDir(directory_name));
   TF_RETURN_IF_ERROR(tsl::Env::Default()->NewWritableFile(file_name, &f));
diff --git a/third_party/xla/xla/service/llvm_ir/loop_emitter.cc b/third_party/xla/xla/service/llvm_ir/loop_emitter.cc
index 8e07365a88c034..c1b8eed3ac1c94 100644
--- a/third_party/xla/xla/service/llvm_ir/loop_emitter.cc
+++ b/third_party/xla/xla/service/llvm_ir/loop_emitter.cc
@@ -72,7 +72,7 @@ BodyEmitter MakeBodyEmitter(const ElementGenerator& target_element_generator,
                                          target_arrays.end());
   if (!is_tuple) {
     CHECK_EQ(target_arrays.size(), 1);
-    return [=](const llvm_ir::IrArray::Index array_index) -> Status {
+    return [=](const llvm_ir::IrArray::Index array_index) -> absl::Status {
       // Convert target_element_generator to a BodyEmitter.
       TF_ASSIGN_OR_RETURN(llvm::Value * target_element,
                           target_element_generator(array_index));
@@ -178,8 +178,8 @@ std::vector<IrArray::Index> LoopEmitter::EmitIndexAndSetExitBasicBlock(
   return {array_index};
 }
 
-Status LoopEmitter::EmitLoop(absl::string_view loop_name,
-                             llvm::Type* index_type) {
+absl::Status LoopEmitter::EmitLoop(absl::string_view loop_name,
+                                   llvm::Type* index_type) {
   if (index_type == nullptr) {
     index_type = b_->getInt64Ty();
   }
diff --git a/third_party/xla/xla/service/llvm_ir/loop_emitter.h b/third_party/xla/xla/service/llvm_ir/loop_emitter.h
index da23eeb5ce37c2..3182de779f9987 100644
--- a/third_party/xla/xla/service/llvm_ir/loop_emitter.h
+++ b/third_party/xla/xla/service/llvm_ir/loop_emitter.h
@@ -35,7 +35,7 @@ namespace llvm_ir {
 // llvm::Value*.
 using ElementGenerator =
     std::function<absl::StatusOr<llvm::Value*>(const IrArray::Index& index)>;
-using BodyEmitter = std::function<Status(const IrArray::Index& index)>;
+using BodyEmitter = std::function<absl::Status(const IrArray::Index& index)>;
 
 // Creates the body emitter from target arrays.
 BodyEmitter MakeBodyEmitter(const ElementGenerator& target_element_generator,
@@ -79,8 +79,8 @@ class LoopEmitter {
       llvm::Value* base_index);
 
   // Emits a complete loop nest for every element in the given shape.
-  Status EmitLoop(absl::string_view loop_name = "",
-                  llvm::Type* index_type = nullptr);
+  absl::Status EmitLoop(absl::string_view loop_name = "",
+                        llvm::Type* index_type = nullptr);
 
  protected:
   // An IR emitter that generates the loop body.
diff --git a/third_party/xla/xla/service/llvm_ir/sort_util.cc b/third_party/xla/xla/service/llvm_ir/sort_util.cc
index 22ed179ae40958..b9773a2646ddf4 100644
--- a/third_party/xla/xla/service/llvm_ir/sort_util.cc
+++ b/third_party/xla/xla/service/llvm_ir/sort_util.cc
@@ -48,7 +48,7 @@ namespace llvm_ir {
 namespace {
 
 // Adds the inner comparison loop body where we compare elements.
-Status EmitCompareLoopBody(
+absl::Status EmitCompareLoopBody(
     int64_t iteration_bound, int64_t num_values,
     llvm::Value* element_pair_index, int64_t xor_mask, llvm::Type* index_type,
     std::function<llvm::Value*(int64_t operand, llvm::Value* index)>
@@ -153,7 +153,7 @@ Status EmitCompareLoopBody(
   });
 }
 
-Status EmitTiledCompareLoop(
+absl::Status EmitTiledCompareLoop(
     const IrArray::Index& tiled_keys_index, int64_t dimension_to_sort,
     int64_t dimension_to_sort_bound, absl::Span<const int64_t> xor_masks,
     const std::vector<IrArray>& params,
@@ -319,7 +319,7 @@ Status EmitTiledCompareLoop(
 }
 }  // namespace
 
-Status EmitSortInPlace(
+absl::Status EmitSortInPlace(
     int64_t dimension_to_sort, const std::vector<IrArray>& values_arrays,
     absl::string_view name, absl::Span<const int64_t> xor_masks,
     llvm::IRBuilder<>* b, const gpu::LaunchDimensions& launch_dimensions,
@@ -369,7 +369,7 @@ Status EmitSortInPlace(
   }
 
   auto compare_loop_body_emitter =
-      [&](const IrArray::Index& tiles_index) -> Status {
+      [&](const IrArray::Index& tiles_index) -> absl::Status {
     // Naive C++ code for the inner compare loop:
     //
     // for (int64_t i = 0; i < dimension_to_sort_bound; ++i) {
diff --git a/third_party/xla/xla/service/llvm_ir/sort_util.h b/third_party/xla/xla/service/llvm_ir/sort_util.h
index e84c2c9194144b..7cf78f9e7bc425 100644
--- a/third_party/xla/xla/service/llvm_ir/sort_util.h
+++ b/third_party/xla/xla/service/llvm_ir/sort_util.h
@@ -28,12 +28,12 @@ limitations under the License.
 namespace xla {
 namespace llvm_ir {
 using EmitCallToNestedComputationCallback =
-    std::function<Status(absl::Span<llvm::Value* const>, llvm::Value*)>;
+    std::function<absl::Status(absl::Span<llvm::Value* const>, llvm::Value*)>;
 // Emits llvm IR to do pairwise comparisons/swaps in the 'dimension_to_sort'
 // dimension of each array in 'values_arrays'. All other dimensions are kept
 // as-is. This implements the inner loop of BitonicSort. It is assumed that
 // 'xor_masks' contains only powers of 2, or values 2^k - 1 (k > 0).
-Status EmitSortInPlace(
+absl::Status EmitSortInPlace(
     int64_t dimension_to_sort, const std::vector<IrArray>& values_arrays,
     absl::string_view name, absl::Span<const int64_t> xor_masks,
     llvm::IRBuilder<>* b, const gpu::LaunchDimensions& launch_dimensions,
diff --git a/third_party/xla/xla/tools/dumped_computation_to_operation_list.cc b/third_party/xla/xla/tools/dumped_computation_to_operation_list.cc
index a916d796b0559d..48b31b993d6fe3 100644
--- a/third_party/xla/xla/tools/dumped_computation_to_operation_list.cc
+++ b/third_party/xla/xla/tools/dumped_computation_to_operation_list.cc
@@ -43,7 +43,7 @@ class OperationDumper : public DfsHloVisitorWithDefault {
  public:
   explicit OperationDumper(const std::string& path) : path_(path) {}
 
-  Status DefaultAction(HloInstruction* hlo) override {
+  absl::Status DefaultAction(HloInstruction* hlo) override {
     std::string params = absl::StrJoin(
         hlo->operands(), ", ",
         [](std::string* out, const HloInstruction* operand) {
diff --git a/third_party/xla/xla/tools/extract_collective_operations.cc b/third_party/xla/xla/tools/extract_collective_operations.cc
index 1a15fa8fbc1fe6..3c46c3983d9291 100644
--- a/third_party/xla/xla/tools/extract_collective_operations.cc
+++ b/third_party/xla/xla/tools/extract_collective_operations.cc
@@ -48,8 +48,8 @@ bazel run extract_collective_operations -- --input=path/to/hlo_module
 }  // namespace
 
 namespace xla {
-Status ExtractCollectiveOperations(const std::string& input,
-                                   const std::string& output) {
+absl::Status ExtractCollectiveOperations(const std::string& input,
+                                         const std::string& output) {
   TF_ASSIGN_OR_RETURN(
       std::unique_ptr<HloModule> test_module,
       LoadModuleFromFile(input, std::string(tsl::io::Extension(input)),
diff --git a/third_party/xla/xla/tools/hlo_control_flow_flattening.cc b/third_party/xla/xla/tools/hlo_control_flow_flattening.cc
index 09c9f6e7355ca5..50eb5bed832397 100644
--- a/third_party/xla/xla/tools/hlo_control_flow_flattening.cc
+++ b/third_party/xla/xla/tools/hlo_control_flow_flattening.cc
@@ -145,7 +145,7 @@ int GetLoopBoundWithOuterLoopMax(const HloInstruction& while_hlo,
   return loop_bound;
 }
 
-Status HloControlFlowFlattening::FlattenWhileLoop(
+absl::Status HloControlFlowFlattening::FlattenWhileLoop(
     HloInstruction* while_hlo, const CallGraph& call_graph) const {
   CHECK_EQ(while_hlo->opcode(), HloOpcode::kWhile);
   HloComputation* computation = while_hlo->parent();
@@ -264,7 +264,7 @@ Status HloControlFlowFlattening::FlattenWhileLoop(
   return OkStatus();
 }
 
-Status HloControlFlowFlattening::RemoveInfeed(
+absl::Status HloControlFlowFlattening::RemoveInfeed(
     HloInstruction* infeed_hlo) const {
   CHECK_EQ(infeed_hlo->opcode(), HloOpcode::kInfeed);
   HloComputation* computation = infeed_hlo->parent();
@@ -326,7 +326,7 @@ HloControlFlowFlattening::RemoveRecvAndRecvDone(
   return std::make_pair(custom_call_recv, custom_call_recv_done);
 }
 
-Status HloControlFlowFlattening::RemoveOutfeed(
+absl::Status HloControlFlowFlattening::RemoveOutfeed(
     HloInstruction* outfeed_hlo) const {
   CHECK_EQ(outfeed_hlo->opcode(), HloOpcode::kOutfeed);
   HloComputation* computation = outfeed_hlo->parent();
@@ -409,7 +409,7 @@ absl::StatusOr<HloInstruction*> HloControlFlowFlattening::RemoveCollective(
   return custom_call;
 }
 
-Status HloControlFlowFlattening::RemoveId(HloInstruction* hlo) const {
+absl::Status HloControlFlowFlattening::RemoveId(HloInstruction* hlo) const {
   HloComputation* computation = hlo->parent();
   HloInstruction* zero = CreateConstant(hlo->shape(), computation);
   std::string original_op_name(hlo->name());
diff --git a/third_party/xla/xla/tools/hlo_control_flow_flattening.h b/third_party/xla/xla/tools/hlo_control_flow_flattening.h
index fee2aa6c5685ea..e48cce0b865a8a 100644
--- a/third_party/xla/xla/tools/hlo_control_flow_flattening.h
+++ b/third_party/xla/xla/tools/hlo_control_flow_flattening.h
@@ -67,16 +67,16 @@ class HloControlFlowFlattening : public HloModulePass {
 
  private:
   // Replaces an infeed with a custom call.
-  Status RemoveInfeed(HloInstruction* infeed_hlo) const;
+  absl::Status RemoveInfeed(HloInstruction* infeed_hlo) const;
   // Removes outfeeds and replaces the outfeed HLO with a side-effecting custom
   // call that ensures that XLA doesn't dead-code-eliminate the outfeeded values
   // but lowers to a no-op.
-  Status RemoveOutfeed(HloInstruction* outfeed_hlo) const;
+  absl::Status RemoveOutfeed(HloInstruction* outfeed_hlo) const;
   // Flattens the while loop. Precondition: while_hlo is a while instruction.
-  Status FlattenWhileLoop(HloInstruction* while_hlo,
-                          const CallGraph& call_graph) const;
+  absl::Status FlattenWhileLoop(HloInstruction* while_hlo,
+                                const CallGraph& call_graph) const;
   // Replaces an id with a zero constant.
-  Status RemoveId(HloInstruction* hlo) const;
+  absl::Status RemoveId(HloInstruction* hlo) const;
 
   int while_execution_count_;
   int max_outer_loop_count_;
diff --git a/third_party/xla/xla/tools/hlo_extractor.cc b/third_party/xla/xla/tools/hlo_extractor.cc
index f9f07153cd9c9e..48935a72b67c26 100644
--- a/third_party/xla/xla/tools/hlo_extractor.cc
+++ b/third_party/xla/xla/tools/hlo_extractor.cc
@@ -94,12 +94,12 @@ class ExtractionVisitor : public ConstDfsHloVisitorWithDefault {
     }
   }
 
-  Status HandleParameter(const HloInstruction* parameter) override {
+  absl::Status HandleParameter(const HloInstruction* parameter) override {
     // Entry parameters need renumbering.
     return ReplaceWithParameter(parameter);
   }
 
-  Status DefaultAction(const HloInstruction* hlo) override {
+  absl::Status DefaultAction(const HloInstruction* hlo) override {
     // Replace the following two types of instructions with parameters/constants
     // (1) the instructions at the boundary with (2) the instructions that are
     // not selected by the hlo_selector.
@@ -154,7 +154,7 @@ class ExtractionVisitor : public ConstDfsHloVisitorWithDefault {
     return OkStatus();
   }
 
-  Status FinishVisit(const HloInstruction* /*root*/) override {
+  absl::Status FinishVisit(const HloInstruction* /*root*/) override {
     // Create the entry computation for the extracted module.
     auto new_entry_computation = module_->AddEntryComputation(
         old_computations_to_builders_.at(root_instruction_->parent())->Build());
@@ -188,7 +188,7 @@ class ExtractionVisitor : public ConstDfsHloVisitorWithDefault {
 
  private:
   // Replace the `hlo` with Constant of the same shape.
-  Status ReplaceWithConstant(const HloInstruction* hlo) {
+  absl::Status ReplaceWithConstant(const HloInstruction* hlo) {
     absl::StatusOr<Literal> literal_status = MakeFakeLiteral(hlo->shape());
     TF_CHECK_OK(literal_status.status());
     auto new_const =
@@ -202,7 +202,7 @@ class ExtractionVisitor : public ConstDfsHloVisitorWithDefault {
   }
 
   // Replace the `hlo` with Parameter of the same shape.
-  Status ReplaceWithParameter(const HloInstruction* hlo) {
+  absl::Status ReplaceWithParameter(const HloInstruction* hlo) {
     CHECK(parameter_numbers_.contains(hlo->parent()));
     auto new_parameter = HloInstruction::CreateParameter(
         parameter_numbers_.at(hlo->parent())++, hlo->shape(), hlo->name());
@@ -268,8 +268,8 @@ class ExtractionVisitor : public ConstDfsHloVisitorWithDefault {
   // Replace with `hlo` with a broadcasted constant of the same shape. The
   // constant could be either a zero or a random number, depending on
   // `replace_type`.
-  Status ReplaceWithConstantBroadcast(const HloInstruction* hlo,
-                                      ReplaceType replace_type) {
+  absl::Status ReplaceWithConstantBroadcast(const HloInstruction* hlo,
+                                            ReplaceType replace_type) {
     CHECK(replace_type == ReplaceType::kReplaceZeroBroadcast ||
           replace_type == ReplaceType::kReplaceRandomBroadcast);
     CHECK(old_computations_to_builders_.contains(hlo->parent()));
diff --git a/third_party/xla/xla/tools/prepare_reference_module.cc b/third_party/xla/xla/tools/prepare_reference_module.cc
index fc904d9fea1c60..baf17e72c84488 100644
--- a/third_party/xla/xla/tools/prepare_reference_module.cc
+++ b/third_party/xla/xla/tools/prepare_reference_module.cc
@@ -33,8 +33,8 @@ namespace xla {
 absl::StatusOr<std::unique_ptr<HloModule>> PrepareReferenceModule(
     const HloModule& test_module, HloRunnerInterface* test_runner,
     const std::function<void(HloModuleConfig*)>& config_modifier_hook,
-    const std::function<Status(const HloModule&, HloRunnerInterface*,
-                               HloModule*)>& module_modifier_hook) {
+    const std::function<absl::Status(const HloModule&, HloRunnerInterface*,
+                                     HloModule*)>& module_modifier_hook) {
   DebugOptions debug_options = GetDebugOptionsFromFlags();
   // The combination of fast math and optimizations leads to unsound code
   // transformations (see third_party/tensorflow/compiler/xla/xla.proto for
diff --git a/third_party/xla/xla/tools/prepare_reference_module.h b/third_party/xla/xla/tools/prepare_reference_module.h
index 5fac3559a196e0..4b67d6eb947763 100644
--- a/third_party/xla/xla/tools/prepare_reference_module.h
+++ b/third_party/xla/xla/tools/prepare_reference_module.h
@@ -36,8 +36,8 @@ namespace xla {
 absl::StatusOr<std::unique_ptr<HloModule>> PrepareReferenceModule(
     const HloModule& test_module, HloRunnerInterface* test_runner,
     const std::function<void(HloModuleConfig*)>& config_modifier_hook = {},
-    const std::function<Status(const HloModule&, HloRunnerInterface*,
-                               HloModule*)>& module_modifier_hook = {});
+    const std::function<absl::Status(const HloModule&, HloRunnerInterface*,
+                                     HloModule*)>& module_modifier_hook = {});
 
 }  // namespace xla
 
diff --git a/third_party/xla/xla/tools/run_hlo_module.cc b/third_party/xla/xla/tools/run_hlo_module.cc
index 2c0a8c61668b9c..10b2b70afbfc24 100644
--- a/third_party/xla/xla/tools/run_hlo_module.cc
+++ b/third_party/xla/xla/tools/run_hlo_module.cc
@@ -167,13 +167,14 @@ absl::StatusOr<Literal> ExecuteWithRunner(
   return std::move(result_status).value();
 }
 
-Status RunAndCompareInternal(
+absl::Status RunAndCompareInternal(
     std::unique_ptr<HloModule> test_module,
     const BufferAssignmentProto* buffer_assignment_proto,
     HloRunnerInterface* test_runner, HloRunnerInterface* reference_runner,
     std::minstd_rand0* engine, const RunHloModuleOptions& options,
     xla::RunHloModuleIterationLiterals* iteration_literals_proto,
-    std::function<Status(const HloModule&, HloRunnerInterface*, HloModule*)>
+    std::function<absl::Status(const HloModule&, HloRunnerInterface*,
+                               HloModule*)>
         reference_module_modifier_hook,
     std::function<void(HloModuleConfig*)> config_modifier_hook,
     ModuleResult* test_run_result, ModuleResult* reference_run_result) {
@@ -325,7 +326,7 @@ Status RunAndCompareInternal(
   ErrorSpec error_spec(static_cast<float>(options.abs_error_bound),
                        static_cast<float>(options.rel_error_bound));
 
-  Status comparison_status =
+  absl::Status comparison_status =
       literal_comparison::Near(/*expected=*/reference_result,
                                /*actual=*/test_result,
                                /*error=*/error_spec,
@@ -345,7 +346,7 @@ struct ChunkResult {
   std::string module_name;
   ModuleResult test_result = ModuleResult::kDidntRun;
   ModuleResult reference_result = ModuleResult::kDidntRun;
-  Status status;
+  absl::Status status;
 
   bool operator<(const ChunkResult& other) const {
     if (test_result != other.test_result) {
@@ -412,13 +413,14 @@ std::string BuildResultsTable(absl::Span<const ChunkResult> chunk_results,
   return strstr.str();
 }
 
-Status RunIsolatedAndCompare(
+absl::Status RunIsolatedAndCompare(
     std::unique_ptr<HloModule> test_module,
     const BufferAssignmentProto* buffer_assignment_proto,
     HloRunnerInterface* test_runner, HloRunnerInterface* reference_runner,
     std::minstd_rand0* engine, const RunHloModuleOptions& options,
     xla::RunHloModuleIterationLiterals* iteration_literals_proto,
-    std::function<Status(const HloModule&, HloRunnerInterface*, HloModule*)>
+    std::function<absl::Status(const HloModule&, HloRunnerInterface*,
+                               HloModule*)>
         reference_module_modifier_hook,
     std::function<void(HloModuleConfig*)> config_modifier_hook) {
   CHECK(test_module);
@@ -439,12 +441,12 @@ Status RunIsolatedAndCompare(
       std::vector<std::unique_ptr<HloModule>> modules,
       DecomposeHloModule(*test_module, /*deduplicate_modules=*/true));
 
-  Status status = OkStatus();
+  absl::Status status = OkStatus();
   for (std::unique_ptr<HloModule>& module : modules) {
     const std::string module_name = module->name();
     ModuleResult test_module_result = ModuleResult::kDidntRun;
     ModuleResult reference_module_result = ModuleResult::kDidntRun;
-    Status chunk_status = RunAndCompareInternal(
+    absl::Status chunk_status = RunAndCompareInternal(
         std::move(module), buffer_assignment_proto, test_runner,
         reference_runner, engine, options, iteration_literals_proto,
         reference_module_modifier_hook, config_modifier_hook,
@@ -460,13 +462,14 @@ Status RunIsolatedAndCompare(
 
 }  // namespace
 
-Status RunAndCompare(
+absl::Status RunAndCompare(
     std::unique_ptr<HloModule> test_module,
     const BufferAssignmentProto* buffer_assignment_proto,
     HloRunnerInterface* test_runner, HloRunnerInterface* reference_runner,
     std::minstd_rand0* engine, const RunHloModuleOptions& options,
     xla::RunHloModuleIterationLiterals* iteration_literals_proto,
-    std::function<Status(const HloModule&, HloRunnerInterface*, HloModule*)>
+    std::function<absl::Status(const HloModule&, HloRunnerInterface*,
+                               HloModule*)>
         reference_module_modifier_hook,
     std::function<void(HloModuleConfig*)> config_modifier_hook) {
   if (options.isolate_instructions) {
@@ -481,15 +484,17 @@ Status RunAndCompare(
       reference_module_modifier_hook, config_modifier_hook, nullptr, nullptr);
 }
 
-Status RunAndCompare(
+absl::Status RunAndCompare(
     const std::string& hlo_filename, HloRunnerInterface* test_runner,
     HloRunnerInterface* reference_runner, std::minstd_rand0* engine,
     const RunHloModuleOptions& options,
     xla::RunHloModuleIterationLiterals* iteration_literals_proto,
-    std::function<Status(const HloModule&, HloRunnerInterface*, HloModule*)>
+    std::function<absl::Status(const HloModule&, HloRunnerInterface*,
+                               HloModule*)>
         reference_module_modifier_hook,
     std::function<void(HloModuleConfig*)> config_modifier_hook,
-    std::function<Status(const RunHloModuleOptions& options, HloModule& module)>
+    std::function<absl::Status(const RunHloModuleOptions& options,
+                               HloModule& module)>
         compilation_env_modifier_hook) {
   BufferAssignmentProto buffer_assignment_proto;
   TF_ASSIGN_OR_RETURN(
diff --git a/third_party/xla/xla/tools/run_hlo_module.h b/third_party/xla/xla/tools/run_hlo_module.h
index 80b16b277c7f08..898b489effb6f6 100644
--- a/third_party/xla/xla/tools/run_hlo_module.h
+++ b/third_party/xla/xla/tools/run_hlo_module.h
@@ -65,13 +65,14 @@ struct RunHloModuleOptions {
 // the results. 'reference_module_modifier_hook' can be used to transform the
 // HloModule before it is run on the reference platform. This may be necessary
 // to match the numerics of the test platform.
-Status RunAndCompare(
+absl::Status RunAndCompare(
     std::unique_ptr<HloModule> test_module,
     const BufferAssignmentProto* buffer_assignment_proto,
     HloRunnerInterface* test_runner, HloRunnerInterface* reference_runner,
     std::minstd_rand0* engine, const RunHloModuleOptions& options,
     xla::RunHloModuleIterationLiterals* iteration_literals_proto = nullptr,
-    std::function<Status(const HloModule&, HloRunnerInterface*, HloModule*)>
+    std::function<absl::Status(const HloModule&, HloRunnerInterface*,
+                               HloModule*)>
         reference_module_modifier_hook = {},
     std::function<void(HloModuleConfig*)> config_modifier_hook = {});
 
@@ -79,15 +80,17 @@ Status RunAndCompare(
 // an argument, a function 'compilation_env_modifier_hook' that potentially sets
 // various fields in compilation environments, for an HLO module being loaded
 // from the file.
-Status RunAndCompare(
+absl::Status RunAndCompare(
     const std::string& hlo_filename, HloRunnerInterface* test_runner,
     HloRunnerInterface* reference_runner, std::minstd_rand0* engine,
     const RunHloModuleOptions& options,
     xla::RunHloModuleIterationLiterals* iteration_literals_proto = nullptr,
-    std::function<Status(const HloModule&, HloRunnerInterface*, HloModule*)>
+    std::function<absl::Status(const HloModule&, HloRunnerInterface*,
+                               HloModule*)>
         reference_module_modifier_hook = {},
     std::function<void(HloModuleConfig*)> config_modifier_hook = {},
-    std::function<Status(const RunHloModuleOptions& options, HloModule& module)>
+    std::function<absl::Status(const RunHloModuleOptions& options,
+                               HloModule& module)>
         compilation_env_modifier_hook = {});
 }  // namespace xla
 
diff --git a/third_party/xla/xla/translate/hlo_to_mhlo/hlo_function_importer.cc b/third_party/xla/xla/translate/hlo_to_mhlo/hlo_function_importer.cc
index 3b1263ad079d31..8407deef789292 100644
--- a/third_party/xla/xla/translate/hlo_to_mhlo/hlo_function_importer.cc
+++ b/third_party/xla/xla/translate/hlo_to_mhlo/hlo_function_importer.cc
@@ -182,7 +182,7 @@ absl::StatusOr<mlir::Operation*> HloFunctionImporter::ImportOldStyleAsyncStart(
     llvm::SmallVectorImpl<mlir::NamedAttribute>& attributes,
     const llvm::SmallVectorImpl<mlir::Value>& operands, mlir::Location loc,
     mlir::Type result_type, mlir::OpBuilder* func_builder,
-    std::string func_name, std::function<Status(sync_op)> mutate_op) {
+    std::string func_name, std::function<absl::Status(sync_op)> mutate_op) {
   auto result_types = result_type.cast<mlir::TupleType>().getTypes();
   if (result_types.size() < 2) {
     return tsl::errors::InvalidArgument(
@@ -442,7 +442,7 @@ absl::StatusOr<mlir::func::FuncOp> HloFunctionImporter::ImportAsFunc(
   return importer.ImportAsFunc(computation, is_main);
 }
 
-Status HloFunctionImporter::ImportAsRegion(
+absl::Status HloFunctionImporter::ImportAsRegion(
     const HloComputation& computation, mlir::SymbolTable& symbol_table,
     mlir::Region* region, mlir::Builder* builder,
     bool flatten_computation_args_result) {
@@ -623,8 +623,8 @@ absl::StatusOr<FuncOp> HloFunctionImporter::ImportAsFunc(
   return function;
 }
 
-Status HloFunctionImporter::ImportAsRegion(const HloComputation& computation,
-                                           mlir::Region* region) {
+absl::Status HloFunctionImporter::ImportAsRegion(
+    const HloComputation& computation, mlir::Region* region) {
   auto loc = region->getLoc();
   // TODO(hinsu): Store computation name as an attribute for round-trip.
   auto* block = new mlir::Block;
@@ -681,7 +681,7 @@ absl::StatusOr<Value> HloFunctionImporter::ImportInstructionsImpl(
   return GetMlirValue(computation.root_instruction());
 }
 
-Status HloFunctionImporter::ImportInstructions(
+absl::Status HloFunctionImporter::ImportInstructions(
     const HloComputation& computation, mlir::Block* block) {
   llvm::SmallVector<Value, 4> arguments(block->args_begin(), block->args_end());
   mlir::OpBuilder builder = mlir::OpBuilder::atBlockEnd(block);
@@ -2271,7 +2271,7 @@ HloFunctionImporter::GetOperands(const HloInstruction* instruction) {
   return operands;
 }
 
-Status HloFunctionImporter::GetMlirTypes(
+absl::Status HloFunctionImporter::GetMlirTypes(
     absl::Span<const HloInstruction* const> instructions,
     llvm::SmallVectorImpl<mlir::Type>* types) {
   for (auto instruction : instructions) {
@@ -2430,7 +2430,7 @@ void HloFunctionImporter::SetLayoutForMlir(mlir::Operation* op,
   op->setAttr(attr_name, GetLayoutAttribute(b, shape));
 }
 
-Status HloFunctionImporter::ConvertShapeToMlirLayout(
+absl::Status HloFunctionImporter::ConvertShapeToMlirLayout(
     const Shape& shape,
     llvm::SmallVectorImpl<mlir::Attribute>& flattened_attr) {
   if (shape.IsToken()) {
diff --git a/third_party/xla/xla/translate/hlo_to_mhlo/hlo_function_importer.h b/third_party/xla/xla/translate/hlo_to_mhlo/hlo_function_importer.h
index 0a051aafb1a70a..df8a7780c84621 100644
--- a/third_party/xla/xla/translate/hlo_to_mhlo/hlo_function_importer.h
+++ b/third_party/xla/xla/translate/hlo_to_mhlo/hlo_function_importer.h
@@ -66,10 +66,10 @@ class HloFunctionImporter {
   // Imports the given hlo computation to the specified region.
   //
   // Flattens the tuple-typed region argument(s) and return value(s).
-  static Status ImportAsRegion(const HloComputation& computation,
-                               mlir::SymbolTable& symbol_table,
-                               mlir::Region* region, mlir::Builder* builder,
-                               bool flatten_computation_args_result = false);
+  static absl::Status ImportAsRegion(
+      const HloComputation& computation, mlir::SymbolTable& symbol_table,
+      mlir::Region* region, mlir::Builder* builder,
+      bool flatten_computation_args_result = false);
 
   // Imports the given computation to the given place specified by `builder`.
   // `arguments` contains values for all parameters.
@@ -170,13 +170,13 @@ class HloFunctionImporter {
       const HloComputation& computation, bool is_main);
 
   // Imports the given computation in the specified region.
-  Status ImportAsRegion(const HloComputation& computation,
-                        mlir::Region* region);
+  absl::Status ImportAsRegion(const HloComputation& computation,
+                              mlir::Region* region);
 
   // Imports instructions from the given computation in the specified block.
   // Assumes that the block already has correct arguments populated.
-  Status ImportInstructions(const HloComputation& computation,
-                            mlir::Block* block);
+  absl::Status ImportInstructions(const HloComputation& computation,
+                                  mlir::Block* block);
   absl::StatusOr<mlir::Value> ImportInstructionsImpl(
       const HloComputation& computation,
       const llvm::SmallVectorImpl<mlir::Value>& arguments,
@@ -204,7 +204,7 @@ class HloFunctionImporter {
 
   // Converts an XLA shape/layout to the corresponding MLIR layout, in
   // flattened_attr, while flattening the tuple layout.
-  Status ConvertShapeToMlirLayout(
+  absl::Status ConvertShapeToMlirLayout(
       const Shape& shape,
       llvm::SmallVectorImpl<mlir::Attribute>& flattened_attr);
 
@@ -213,8 +213,9 @@ class HloFunctionImporter {
 
   // Takes a list of HloInstructions and generates the list of types used for
   // input, bypassing tuples to subsets.
-  Status GetMlirTypes(absl::Span<const HloInstruction* const> instructions,
-                      llvm::SmallVectorImpl<mlir::Type>* types);
+  absl::Status GetMlirTypes(
+      absl::Span<const HloInstruction* const> instructions,
+      llvm::SmallVectorImpl<mlir::Type>* types);
 
   // Returns the Mlir Value for the corresponding HloInstruction.
   absl::StatusOr<mlir::Value> GetMlirValue(const HloInstruction* instruction);
@@ -272,7 +273,7 @@ class HloFunctionImporter {
       llvm::SmallVectorImpl<mlir::NamedAttribute>& attributes,
       const llvm::SmallVectorImpl<mlir::Value>& operands, mlir::Location loc,
       mlir::Type result_type, mlir::OpBuilder* func_builder,
-      std::string func_name, std::function<Status(SyncOp)> mutate_op);
+      std::string func_name, std::function<absl::Status(SyncOp)> mutate_op);
 
   // Imports an old-style async done op
   absl::StatusOr<mlir::Operation*> ImportOldStyleAsyncDone(
diff --git a/third_party/xla/xla/translate/hlo_to_mhlo/hlo_module_importer.cc b/third_party/xla/xla/translate/hlo_to_mhlo/hlo_module_importer.cc
index 1e80399cb7253a..c0fb1ffbcf5c9d 100644
--- a/third_party/xla/xla/translate/hlo_to_mhlo/hlo_module_importer.cc
+++ b/third_party/xla/xla/translate/hlo_to_mhlo/hlo_module_importer.cc
@@ -99,7 +99,7 @@ mlir::ArrayAttr ConvertCrossProgramPrefetches(
 }
 }  // namespace
 
-Status HloModuleImporter::Import(const HloModule& hlo_module) {
+absl::Status HloModuleImporter::Import(const HloModule& hlo_module) {
   auto module = llvm::cast<mlir::ModuleOp>(symbol_table_.getOp());
   module.setName(hlo_module.name());
   module->setAttr(
@@ -157,7 +157,7 @@ Status HloModuleImporter::Import(const HloModule& hlo_module) {
   return OkStatus();
 }
 
-Status HloModuleImporter::Import(const HloModuleProto& module_proto) {
+absl::Status HloModuleImporter::Import(const HloModuleProto& module_proto) {
   DebugOptions debug_options;
   TF_ASSIGN_OR_RETURN(
       auto module_config,
diff --git a/third_party/xla/xla/translate/hlo_to_mhlo/hlo_module_importer.h b/third_party/xla/xla/translate/hlo_to_mhlo/hlo_module_importer.h
index ec6e6bc79b85f5..bc608c5f42153d 100644
--- a/third_party/xla/xla/translate/hlo_to_mhlo/hlo_module_importer.h
+++ b/third_party/xla/xla/translate/hlo_to_mhlo/hlo_module_importer.h
@@ -43,10 +43,10 @@ class HloModuleImporter {
                              bool flatten_computation_args_result = false);
 
   // Import the HloModule into the MLIR Module.
-  Status Import(const xla::HloModule& module);
+  absl::Status Import(const xla::HloModule& module);
 
   // Import the HloModuleProto into the MLIR Module.
-  Status Import(const xla::HloModuleProto& module);
+  absl::Status Import(const xla::HloModuleProto& module);
 
  private:
   void ImportFrontendAttributes(const xla::HloModule& hlo_module,
diff --git a/third_party/xla/xla/translate/hlo_to_mhlo/hlo_to_mlir_hlo.cc b/third_party/xla/xla/translate/hlo_to_mhlo/hlo_to_mlir_hlo.cc
index 6706f2f3695188..e8d81dcc4a92d2 100644
--- a/third_party/xla/xla/translate/hlo_to_mhlo/hlo_to_mlir_hlo.cc
+++ b/third_party/xla/xla/translate/hlo_to_mhlo/hlo_to_mlir_hlo.cc
@@ -21,19 +21,20 @@ limitations under the License.
 
 namespace xla {
 
-Status ConvertHloToMlirHlo(mlir::ModuleOp module,
-                           xla::HloModuleProto const* hlo_module_proto,
-                           bool import_all_computation,
-                           bool flatten_computation_args_result) {
+absl::Status ConvertHloToMlirHlo(mlir::ModuleOp module,
+                                 xla::HloModuleProto const* hlo_module_proto,
+                                 bool import_all_computation,
+                                 bool flatten_computation_args_result) {
   mlir::BaseScopedDiagnosticHandler diag_handler(module.getContext());
   return HloModuleImporter(module, import_all_computation,
                            flatten_computation_args_result)
       .Import(*hlo_module_proto);
 }
 
-Status ConvertHloToMlirHlo(mlir::ModuleOp module, xla::HloModule* hlo_module,
-                           bool import_all_computation,
-                           bool flatten_computation_args_result) {
+absl::Status ConvertHloToMlirHlo(mlir::ModuleOp module,
+                                 xla::HloModule* hlo_module,
+                                 bool import_all_computation,
+                                 bool flatten_computation_args_result) {
   mlir::BaseScopedDiagnosticHandler diag_handler(module.getContext());
   return HloModuleImporter(module, import_all_computation,
                            flatten_computation_args_result)
diff --git a/third_party/xla/xla/translate/hlo_to_mhlo/hlo_to_mlir_hlo.h b/third_party/xla/xla/translate/hlo_to_mhlo/hlo_to_mlir_hlo.h
index caec1a8e5143df..277f522331ffda 100644
--- a/third_party/xla/xla/translate/hlo_to_mhlo/hlo_to_mlir_hlo.h
+++ b/third_party/xla/xla/translate/hlo_to_mhlo/hlo_to_mlir_hlo.h
@@ -35,10 +35,10 @@ class HloModuleProto;
 //
 // If `flatten_computation_args_result` is set to true, flattens all tuple
 // arguments and result of every computation when importing them as func ops.
-Status ConvertHloToMlirHlo(mlir::ModuleOp module,
-                           xla::HloModuleProto const* hlo_module,
-                           bool import_all_computations = false,
-                           bool flatten_computation_args_result = false);
+absl::Status ConvertHloToMlirHlo(mlir::ModuleOp module,
+                                 xla::HloModuleProto const* hlo_module,
+                                 bool import_all_computations = false,
+                                 bool flatten_computation_args_result = false);
 
 // Converts an HLO module to a MLIR module in HLO dialect.
 //
@@ -47,9 +47,10 @@ Status ConvertHloToMlirHlo(mlir::ModuleOp module,
 //
 // If `flatten_computation_args_result` is set to true, flattens all tuple
 // arguments and result of every computation when importing them as func ops.
-Status ConvertHloToMlirHlo(mlir::ModuleOp module, xla::HloModule* hlo_module,
-                           bool import_all_computations = false,
-                           bool flatten_computation_args_result = false);
+absl::Status ConvertHloToMlirHlo(mlir::ModuleOp module,
+                                 xla::HloModule* hlo_module,
+                                 bool import_all_computations = false,
+                                 bool flatten_computation_args_result = false);
 
 }  // namespace xla
 
diff --git a/third_party/xla/xla/translate/hlo_to_mhlo/hlo_utils.h b/third_party/xla/xla/translate/hlo_to_mhlo/hlo_utils.h
index b570125c654da9..7e0f21c97578b9 100644
--- a/third_party/xla/xla/translate/hlo_to_mhlo/hlo_utils.h
+++ b/third_party/xla/xla/translate/hlo_to_mhlo/hlo_utils.h
@@ -43,8 +43,8 @@ mlir::DenseIntElementsAttr CreateDenseIntElementsAttrFromVector(
 
 // Converts the given XLA shape for tensors to the template MLIR type.
 template <typename TypeT>
-static StatusOr<TypeT> ConvertTensorShapeToType(const Shape& xla_ty,
-                                                mlir::Builder builder) {
+static absl::StatusOr<TypeT> ConvertTensorShapeToType(const Shape& xla_ty,
+                                                      mlir::Builder builder) {
   auto element_type_or =
       ConvertPrimitiveTypeToMlirType(xla_ty.element_type(), builder);
   if (!element_type_or.ok()) return element_type_or.status();

From d18c491313468e90fd0ebeb2ee28926f913a7b72 Mon Sep 17 00:00:00 2001
From: Kyle Lucke <klucke@google.com>
Date: Wed, 15 May 2024 16:18:44 -0700
Subject: [PATCH 227/478] Make compiler/jit targets use
 StreamExecutorInterface::CreateEvent to create stream_executor::Events.

PiperOrigin-RevId: 634109825
---
 tensorflow/compiler/jit/BUILD                 | 2 ++
 tensorflow/compiler/jit/xla_device_context.cc | 5 +++--
 tensorflow/compiler/jit/xla_launch_util.cc    | 5 +----
 tensorflow/compiler/jit/xla_tpu_device.cc     | 6 +++---
 4 files changed, 9 insertions(+), 9 deletions(-)

diff --git a/tensorflow/compiler/jit/BUILD b/tensorflow/compiler/jit/BUILD
index 36552f4ae3027e..623334534567de 100644
--- a/tensorflow/compiler/jit/BUILD
+++ b/tensorflow/compiler/jit/BUILD
@@ -199,6 +199,7 @@ cc_library(
         "//tensorflow/core/tpu:tpu_node_device_util",
         "//tensorflow/core/tpu:virtual_device",
         "@com_google_absl//absl/types:optional",
+        "@local_tsl//tsl/platform:statusor",
         "@local_xla//xla/stream_executor/tpu:c_api_conversions",
         "@local_xla//xla/stream_executor/tpu:status_helper",
         "@local_xla//xla/stream_executor/tpu:tpu_api",
@@ -314,6 +315,7 @@ cc_library(
         "//tensorflow/core/common_runtime:dma_helper",
         "//tensorflow/core/framework:allocator",
         "@com_google_absl//absl/synchronization",
+        "@local_tsl//tsl/platform:statusor",
         "@local_xla//xla:util",
         "@local_xla//xla/client:global_data",
         "@local_xla//xla/client:local_client",
diff --git a/tensorflow/compiler/jit/xla_device_context.cc b/tensorflow/compiler/jit/xla_device_context.cc
index 821d294af90f66..faf3b65d407a7e 100644
--- a/tensorflow/compiler/jit/xla_device_context.cc
+++ b/tensorflow/compiler/jit/xla_device_context.cc
@@ -30,6 +30,7 @@ limitations under the License.
 #include "tensorflow/core/common_runtime/device.h"
 #include "tensorflow/core/common_runtime/dma_helper.h"
 #include "tensorflow/core/framework/tensor_reference.h"
+#include "tsl/platform/statusor.h"
 
 namespace tensorflow {
 
@@ -171,8 +172,8 @@ void XlaDeviceContext::CopyCPUTensorToDevice(const Tensor* cpu_tensor,
         host_to_device_stream_.get(), literal, xla_tensor->shaped_buffer()));
 
     if (UseMultipleStreams()) {
-      auto event = std::make_shared<se::Event>(stream_->parent());
-      TF_RET_CHECK(event->Init()) << "Event failed to initialize!";
+      TF_ASSIGN_OR_RETURN(std::shared_ptr<se::Event> event,
+                          stream_->parent()->CreateEvent());
       TF_RETURN_IF_ERROR(host_to_device_stream_->RecordEvent(event.get()));
       xla_tensor->ResetDefinitionEvent(std::move(event),
                                        host_to_device_stream_.get());
diff --git a/tensorflow/compiler/jit/xla_launch_util.cc b/tensorflow/compiler/jit/xla_launch_util.cc
index fcfad5c0a7eb67..cfeaa937024b32 100644
--- a/tensorflow/compiler/jit/xla_launch_util.cc
+++ b/tensorflow/compiler/jit/xla_launch_util.cc
@@ -390,10 +390,7 @@ Status XlaComputationLaunchContext::PopulateOutputs(
 
   std::shared_ptr<se::Event> definition_event;
   if (use_multiple_streams_ && stream) {
-    definition_event = std::make_shared<se::Event>(stream->parent());
-    if (!definition_event->Init()) {
-      return errors::Internal("Failed to initialize tensor definition event.");
-    }
+    TF_ASSIGN_OR_RETURN(definition_event, stream->parent()->CreateEvent());
     TF_RETURN_IF_ERROR(stream->RecordEvent(definition_event.get()));
   }
 
diff --git a/tensorflow/compiler/jit/xla_tpu_device.cc b/tensorflow/compiler/jit/xla_tpu_device.cc
index dfedd586df69aa..403e6b17e6fc00 100644
--- a/tensorflow/compiler/jit/xla_tpu_device.cc
+++ b/tensorflow/compiler/jit/xla_tpu_device.cc
@@ -43,6 +43,7 @@ limitations under the License.
 #include "tensorflow/core/tpu/tpu_defs.h"
 #include "tensorflow/core/tpu/tpu_node_device_util.h"
 #include "tensorflow/core/tpu/virtual_device.h"
+#include "tsl/platform/statusor.h"
 
 namespace tensorflow {
 namespace {
@@ -271,9 +272,8 @@ void TpuDeviceToDeviceCopy(DeviceContext* src_dev_context,
           dst_xla_context->host_to_device_stream()));
     }
 
-    auto definition_event =
-        std::make_shared<se::Event>(dst_xla_context->stream()->parent());
-    TF_RET_CHECK(definition_event->Init()) << "Event failed to initialize!";
+    TF_ASSIGN_OR_RETURN(std::shared_ptr<se::Event> definition_event,
+                        dst_xla_context->stream()->parent()->CreateEvent());
     TF_RETURN_IF_ERROR(
         dst_device_to_device_stream->RecordEvent(definition_event.get()));
     xla_output->ResetDefinitionEvent(std::move(definition_event),

From b2830923b03f73f1e1a955d92f0c77a822f1d21a Mon Sep 17 00:00:00 2001
From: Kyle Lucke <klucke@google.com>
Date: Wed, 15 May 2024 16:57:44 -0700
Subject: [PATCH 228/478] Use absl::Status instead of xla::Status now that
 they're identical.

PiperOrigin-RevId: 634123462
---
 .../auto_sharding/auto_sharding.cc            | 17 +++---
 .../auto_sharding/auto_sharding.h             | 61 ++++++++++---------
 .../auto_sharding_dot_handler.cc              | 46 +++++++-------
 .../auto_sharding/auto_sharding_runner.cc     |  2 +-
 .../auto_sharding/auto_sharding_solver.cc     |  4 +-
 .../auto_sharding/auto_sharding_solver.h      |  2 +-
 .../auto_sharding/auto_sharding_strategy.cc   |  7 ++-
 .../auto_sharding/auto_sharding_util.cc       |  8 +--
 8 files changed, 77 insertions(+), 70 deletions(-)

diff --git a/third_party/xla/xla/hlo/experimental/auto_sharding/auto_sharding.cc b/third_party/xla/xla/hlo/experimental/auto_sharding/auto_sharding.cc
index e041ff800ded20..5b1d5f0229ae6e 100644
--- a/third_party/xla/xla/hlo/experimental/auto_sharding/auto_sharding.cc
+++ b/third_party/xla/xla/hlo/experimental/auto_sharding/auto_sharding.cc
@@ -2320,7 +2320,7 @@ void SetHloSharding(
   }
 }
 
-Status SetHloShardingPostProcessing(
+absl::Status SetHloShardingPostProcessing(
     const HloInstructionSequence& sequence,
     const absl::flat_hash_set<const HloInstruction*>& instructions_to_shard,
     const StrategyMap& strategy_map, const CostGraph& cost_graph,
@@ -3400,11 +3400,11 @@ void AnnotateShardingWithSimpleHeuristic(
 
 // Filter strategies according to the option.force_batch_dim_to_mesh_dim.
 // This can be used to forcibly generate data-parallel strategies.
-Status FilterStrategy(const HloInstruction* ins, const Shape& shape,
-                      std::unique_ptr<StrategyGroup>& strategy_group,
-                      const ClusterEnvironment& cluster_env,
-                      const InstructionBatchDimMap& batch_map,
-                      const AutoShardingOption& option) {
+absl::Status FilterStrategy(const HloInstruction* ins, const Shape& shape,
+                            std::unique_ptr<StrategyGroup>& strategy_group,
+                            const ClusterEnvironment& cluster_env,
+                            const InstructionBatchDimMap& batch_map,
+                            const AutoShardingOption& option) {
   int mesh_dim = option.force_batch_dim_to_mesh_dim;
   int batch_dim = batch_map.at(GetBatchDimMapKey(ins));
   const Array<int64_t>& device_mesh = cluster_env.device_mesh_;
@@ -3731,7 +3731,8 @@ AutoShardingImplementation::SaveAndRemoveShardingAnnotation(
   return std::make_pair(preserve_shardings, module_is_changed);
 }
 
-Status AutoShardingImplementation::CanonicalizeLayouts(HloModule* module) {
+absl::Status AutoShardingImplementation::CanonicalizeLayouts(
+    HloModule* module) {
   if (!module->layout_canonicalization_callback()) {
     LOG(INFO) << "There is no registered layout_canonicalization_callback.";
     return OkStatus();
@@ -4070,7 +4071,7 @@ absl::StatusOr<AutoShardingResult> AutoShardingImplementation::RunAutoSharding(
                              option_.try_multiple_mesh_shapes));
     spmd::AliasSet alias_set =
         spmd::BuildAliasSet(module, input_output_alias_config, strategy_map);
-    if (Status alias_set_status = CheckAliasSetCompatibility(
+    if (absl::Status alias_set_status = CheckAliasSetCompatibility(
             alias_set, strategy_groups, sequence,
             /* crash_at_error */ !option_.try_multiple_mesh_shapes);
         !alias_set_status.ok()) {
diff --git a/third_party/xla/xla/hlo/experimental/auto_sharding/auto_sharding.h b/third_party/xla/xla/hlo/experimental/auto_sharding/auto_sharding.h
index 1dbe24348f0454..b7266fbbdad321 100644
--- a/third_party/xla/xla/hlo/experimental/auto_sharding/auto_sharding.h
+++ b/third_party/xla/xla/hlo/experimental/auto_sharding/auto_sharding.h
@@ -95,7 +95,7 @@ class AutoShardingImplementation {
   // PJRT which assigns layouts based on runtime shapes: see
   // DetermineArgumentLayoutsFromCompileOptions() in
   //     tensorflow/compiler/xla/pjrt/utils.cc
-  Status CanonicalizeLayouts(HloModule* module);
+  absl::Status CanonicalizeLayouts(HloModule* module);
 
   // Returns the optimal objective value that the ILP solver computes.
   double GetSolverOptimalObjectiveValue() {
@@ -164,30 +164,33 @@ void SetInNodesWithInstruction(std::unique_ptr<StrategyGroup>& strategy_group,
 
 void RemoveDuplicatedStrategy(std::unique_ptr<StrategyGroup>& strategy_group);
 
-Status FilterStrategy(const HloInstruction* ins, const Shape& shape,
-                      std::unique_ptr<StrategyGroup>& strategy_group,
-                      const ClusterEnvironment& cluster_env,
-                      const InstructionBatchDimMap& batch_map,
-                      const AutoShardingOption& option);
-
-Status HandleDot(std::unique_ptr<StrategyGroup>& strategy_group,
-                 StrategyGroups& strategy_groups, StrategyMap& strategy_map,
-                 const HloInstruction* ins, size_t instruction_id,
-                 const HloInstructionSequence& instruction_sequence,
-                 const HloCostAnalysis& hlo_cost_analysis,
-                 const ClusterEnvironment& cluster_env,
-                 const InstructionBatchDimMap& batch_map,
-                 const AutoShardingOption& option, const CallGraph& call_graph);
-
-Status HandleConv(std::unique_ptr<StrategyGroup>& strategy_group,
-                  StrategyGroups& strategy_groups, StrategyMap& strategy_map,
-                  const HloInstruction* ins, size_t instruction_id,
-                  const HloInstructionSequence& instruction_sequence,
-                  const HloCostAnalysis& hlo_cost_analysis,
-                  const ClusterEnvironment& cluster_env,
-                  const InstructionBatchDimMap& batch_map,
-                  const AutoShardingOption& option,
-                  const CallGraph& call_graph);
+absl::Status FilterStrategy(const HloInstruction* ins, const Shape& shape,
+                            std::unique_ptr<StrategyGroup>& strategy_group,
+                            const ClusterEnvironment& cluster_env,
+                            const InstructionBatchDimMap& batch_map,
+                            const AutoShardingOption& option);
+
+absl::Status HandleDot(std::unique_ptr<StrategyGroup>& strategy_group,
+                       StrategyGroups& strategy_groups,
+                       StrategyMap& strategy_map, const HloInstruction* ins,
+                       size_t instruction_id,
+                       const HloInstructionSequence& instruction_sequence,
+                       const HloCostAnalysis& hlo_cost_analysis,
+                       const ClusterEnvironment& cluster_env,
+                       const InstructionBatchDimMap& batch_map,
+                       const AutoShardingOption& option,
+                       const CallGraph& call_graph);
+
+absl::Status HandleConv(std::unique_ptr<StrategyGroup>& strategy_group,
+                        StrategyGroups& strategy_groups,
+                        StrategyMap& strategy_map, const HloInstruction* ins,
+                        size_t instruction_id,
+                        const HloInstructionSequence& instruction_sequence,
+                        const HloCostAnalysis& hlo_cost_analysis,
+                        const ClusterEnvironment& cluster_env,
+                        const InstructionBatchDimMap& batch_map,
+                        const AutoShardingOption& option,
+                        const CallGraph& call_graph);
 
 void AnnotateShardingWithSimpleHeuristic(HloModule* module,
                                          const std::string& heuristic,
@@ -207,10 +210,10 @@ AliasMap BuildAliasMap(const HloModule* module);
 AliasSet BuildAliasSet(const HloModule* module,
                        const StrategyMap& strategy_map);
 
-Status CheckAliasSetCompatibility(const AliasSet& alias_set,
-                                  const StrategyGroups& strategy_groups,
-                                  const HloInstructionSequence& sequence,
-                                  bool crash_on_error);
+absl::Status CheckAliasSetCompatibility(const AliasSet& alias_set,
+                                        const StrategyGroups& strategy_groups,
+                                        const HloInstructionSequence& sequence,
+                                        bool crash_on_error);
 
 absl::Status GenerateReduceScatter(
     const HloInstructionSequence& sequence, const AliasMap& alias_map,
diff --git a/third_party/xla/xla/hlo/experimental/auto_sharding/auto_sharding_dot_handler.cc b/third_party/xla/xla/hlo/experimental/auto_sharding/auto_sharding_dot_handler.cc
index 8a8bf39312fd5a..1405ceffe14a88 100644
--- a/third_party/xla/xla/hlo/experimental/auto_sharding/auto_sharding_dot_handler.cc
+++ b/third_party/xla/xla/hlo/experimental/auto_sharding/auto_sharding_dot_handler.cc
@@ -250,7 +250,7 @@ class DotHandler : public HandlerBase {
       const DimMap& rhs_dim_map, const DimMap& output_dim_map,
       const Array<int64_t>& device_mesh, double compute_cost) override;
 
-  Status RegisterStrategies();
+  absl::Status RegisterStrategies();
 
   // Dimension information
   bool is_dot_;
@@ -285,7 +285,7 @@ class ConvHandler : public HandlerBase {
 
   void SplitDepthwise(bool forward);
 
-  Status RegisterStrategies();
+  absl::Status RegisterStrategies();
 
   // Dimension information
   const ConvolutionDimensionNumbers& conv_dnums_;
@@ -949,7 +949,7 @@ void DotHandler::AppendReduceScatterWindowedEinsumStrategy(
   }
 }
 
-Status DotHandler::RegisterStrategies() {
+absl::Status DotHandler::RegisterStrategies() {
   // SS = SR x RS
   // Split lhs space dim and rhs space dim.
   SplitLhsSpaceRhsSpace();
@@ -1061,7 +1061,7 @@ ConvHandler::ConvHandler(std::unique_ptr<StrategyGroup>& strategy_group,
   out_out_channel_dim_ = conv_dnums_.output_feature_dimension();
 }
 
-Status ConvHandler::RegisterStrategies() {
+absl::Status ConvHandler::RegisterStrategies() {
   // For 1D sharding
   if ((ins_->feature_group_count() ==
            lhs_->shape().dimensions(lhs_in_channel_dim_) &&
@@ -1216,15 +1216,16 @@ void ConvHandler::SplitDepthwise(bool forward) {
 }  // namespace
 
 // Register strategies for dot instructions.
-Status HandleDot(std::unique_ptr<StrategyGroup>& strategy_group,
-                 StrategyGroups& strategy_groups, StrategyMap& strategy_map,
-                 const HloInstruction* ins, size_t instruction_id,
-                 const HloInstructionSequence& instruction_sequence,
-                 const HloCostAnalysis& hlo_cost_analysis,
-                 const ClusterEnvironment& cluster_env,
-                 const InstructionBatchDimMap& batch_map,
-                 const AutoShardingOption& option,
-                 const CallGraph& call_graph) {
+absl::Status HandleDot(std::unique_ptr<StrategyGroup>& strategy_group,
+                       StrategyGroups& strategy_groups,
+                       StrategyMap& strategy_map, const HloInstruction* ins,
+                       size_t instruction_id,
+                       const HloInstructionSequence& instruction_sequence,
+                       const HloCostAnalysis& hlo_cost_analysis,
+                       const ClusterEnvironment& cluster_env,
+                       const InstructionBatchDimMap& batch_map,
+                       const AutoShardingOption& option,
+                       const CallGraph& call_graph) {
   strategy_group = CreateLeafStrategyGroup(instruction_id, ins, strategy_map,
                                            strategy_groups);
 
@@ -1236,15 +1237,16 @@ Status HandleDot(std::unique_ptr<StrategyGroup>& strategy_group,
 }
 
 // Register strategies for convolution instructions.
-Status HandleConv(std::unique_ptr<StrategyGroup>& strategy_group,
-                  StrategyGroups& strategy_groups, StrategyMap& strategy_map,
-                  const HloInstruction* ins, size_t instruction_id,
-                  const HloInstructionSequence& instruction_sequence,
-                  const HloCostAnalysis& hlo_cost_analysis,
-                  const ClusterEnvironment& cluster_env,
-                  const InstructionBatchDimMap& batch_map,
-                  const AutoShardingOption& option,
-                  const CallGraph& call_graph) {
+absl::Status HandleConv(std::unique_ptr<StrategyGroup>& strategy_group,
+                        StrategyGroups& strategy_groups,
+                        StrategyMap& strategy_map, const HloInstruction* ins,
+                        size_t instruction_id,
+                        const HloInstructionSequence& instruction_sequence,
+                        const HloCostAnalysis& hlo_cost_analysis,
+                        const ClusterEnvironment& cluster_env,
+                        const InstructionBatchDimMap& batch_map,
+                        const AutoShardingOption& option,
+                        const CallGraph& call_graph) {
   strategy_group = CreateLeafStrategyGroup(instruction_id, ins, strategy_map,
                                            strategy_groups);
 
diff --git a/third_party/xla/xla/hlo/experimental/auto_sharding/auto_sharding_runner.cc b/third_party/xla/xla/hlo/experimental/auto_sharding/auto_sharding_runner.cc
index b7a721780b998c..58494978c14816 100644
--- a/third_party/xla/xla/hlo/experimental/auto_sharding/auto_sharding_runner.cc
+++ b/third_party/xla/xla/hlo/experimental/auto_sharding/auto_sharding_runner.cc
@@ -28,7 +28,7 @@ namespace xla {
 namespace spmd {
 namespace {
 
-Status RunAutoShardingPassFromFile(const std::string& file_name) {
+absl::Status RunAutoShardingPassFromFile(const std::string& file_name) {
   std::string hlo_text;
   TF_RETURN_IF_ERROR(
       tsl::ReadFileToString(tsl::Env::Default(), file_name, &hlo_text));
diff --git a/third_party/xla/xla/hlo/experimental/auto_sharding/auto_sharding_solver.cc b/third_party/xla/xla/hlo/experimental/auto_sharding/auto_sharding_solver.cc
index 16816e000e3ffd..df4ebe5b9ac40e 100644
--- a/third_party/xla/xla/hlo/experimental/auto_sharding/auto_sharding_solver.cc
+++ b/third_party/xla/xla/hlo/experimental/auto_sharding/auto_sharding_solver.cc
@@ -865,7 +865,7 @@ AutoShardingSolverResult SolveAndExtractSolution(
     const MPVariable* overbudget_var, const MPVariable* makespan_var,
     MPSolver& solver) {
   auto status = solver.Solve();
-  LOG(INFO) << "Solver Status: " << status;
+  LOG(INFO) << "Solver absl::Status: " << status;
 
   if (status == operations_research::MPSolver::INFEASIBLE) {
     LOG(ERROR) << "MPSolver could not find any feasible solution.";
@@ -1231,7 +1231,7 @@ std::vector<std::string> Rationalize(const AutoShardingSolverRequest& request,
   return rationales;
 }
 
-Status ValidateRequest(const AutoShardingSolverRequest& request) {
+absl::Status ValidateRequest(const AutoShardingSolverRequest& request) {
   const int num_nodes = request.num_nodes();
   const int num_edges = request.edges_size();
   TF_RET_CHECK(num_nodes == request.computation_costs_size());
diff --git a/third_party/xla/xla/hlo/experimental/auto_sharding/auto_sharding_solver.h b/third_party/xla/xla/hlo/experimental/auto_sharding/auto_sharding_solver.h
index 32809503603c75..9b3c099cd92521 100644
--- a/third_party/xla/xla/hlo/experimental/auto_sharding/auto_sharding_solver.h
+++ b/third_party/xla/xla/hlo/experimental/auto_sharding/auto_sharding_solver.h
@@ -142,7 +142,7 @@ class StrategyShaver {
 
 // Check fail if `request` is invalid (e.g., because of negative node costs).
 // Note: This does not include checks for valid variable aliasing yet.
-Status ValidateRequest(const AutoShardingSolverRequest& request);
+absl::Status ValidateRequest(const AutoShardingSolverRequest& request);
 
 }  // namespace spmd
 }  // namespace xla
diff --git a/third_party/xla/xla/hlo/experimental/auto_sharding/auto_sharding_strategy.cc b/third_party/xla/xla/hlo/experimental/auto_sharding/auto_sharding_strategy.cc
index e0c8c51c8bc2ed..b5c41f7600e78d 100644
--- a/third_party/xla/xla/hlo/experimental/auto_sharding/auto_sharding_strategy.cc
+++ b/third_party/xla/xla/hlo/experimental/auto_sharding/auto_sharding_strategy.cc
@@ -897,9 +897,10 @@ BuildStrategyAndCost(
           << ins->ToString() << " does not have any valid strategies.";
     } else if (!(strategy_group->is_tuple ||
                  !strategy_group->strategies.empty())) {
-      return Status(absl::StatusCode::kFailedPrecondition,
-                    "Could not generate any shardings for an instruction due "
-                    "to mismatched mesh shapes.");
+      return absl::Status(
+          absl::StatusCode::kFailedPrecondition,
+          "Could not generate any shardings for an instruction due "
+          "to mismatched mesh shapes.");
     }
     // Checks the shape of resharding_costs is valid. It will check fail if the
     // shape is not as expected.
diff --git a/third_party/xla/xla/hlo/experimental/auto_sharding/auto_sharding_util.cc b/third_party/xla/xla/hlo/experimental/auto_sharding/auto_sharding_util.cc
index 71775ed9fd67de..e43fc36aed518c 100644
--- a/third_party/xla/xla/hlo/experimental/auto_sharding/auto_sharding_util.cc
+++ b/third_party/xla/xla/hlo/experimental/auto_sharding/auto_sharding_util.cc
@@ -1803,10 +1803,10 @@ AliasSet BuildAliasSet(const HloModule* module,
   return alias_set;
 }
 
-Status CheckAliasSetCompatibility(const AliasSet& alias_set,
-                                  const StrategyGroups& strategy_groups,
-                                  const HloInstructionSequence& sequence,
-                                  bool crash_on_error) {
+absl::Status CheckAliasSetCompatibility(const AliasSet& alias_set,
+                                        const StrategyGroups& strategy_groups,
+                                        const HloInstructionSequence& sequence,
+                                        bool crash_on_error) {
   const std::vector<HloInstruction*>& instructions = sequence.instructions();
   // Checks the compatibility
   for (const auto& pair : alias_set) {

From 0f51ba77ec5c231d3b83e0f7e8215fa02b9e75fc Mon Sep 17 00:00:00 2001
From: Fangrui Song <maskray@google.com>
Date: Wed, 15 May 2024 17:11:34 -0700
Subject: [PATCH 229/478] Integrate LLVM at llvm/llvm-project@e1ed138a67a9

Updates LLVM usage to match
[e1ed138a67a9](https://github.com/llvm/llvm-project/commit/e1ed138a67a9)

PiperOrigin-RevId: 634127563
---
 third_party/llvm/workspace.bzl | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/third_party/llvm/workspace.bzl b/third_party/llvm/workspace.bzl
index 573719561ffed1..92ab0a674b5fc1 100644
--- a/third_party/llvm/workspace.bzl
+++ b/third_party/llvm/workspace.bzl
@@ -4,8 +4,8 @@ load("//third_party:repo.bzl", "tf_http_archive")
 
 def repo(name):
     """Imports LLVM."""
-    LLVM_COMMIT = "45726c1a3a3d89ff9f6ebe657c3cb7bcd59b88db"
-    LLVM_SHA256 = "61b637626b1d0c7f67b5387a9e6a49097dbdfff85f578b53aa3a4fb63d8cb362"
+    LLVM_COMMIT = "e1ed138a67a92ef1ff0214ca198094be13045090"
+    LLVM_SHA256 = "3a55ccdd3e90334b1181bb4950972cf005a2b99641cd2f0bc153d91d3a62a727"
 
     tf_http_archive(
         name = name,

From 1ad1a6b67d82885ed2d6e4ed85778708a0a40379 Mon Sep 17 00:00:00 2001
From: Samuel Agyakwa <sagyakwa@google.com>
Date: Wed, 15 May 2024 17:14:33 -0700
Subject: [PATCH 230/478] [Cleanup] Update all instances of tfrt::t:: to use
 tfrt::tfrt_tensor

PiperOrigin-RevId: 634128347
---
 tensorflow/compiler/mlir/tfrt/ir/tfrt_fallback_sync.cc | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/tensorflow/compiler/mlir/tfrt/ir/tfrt_fallback_sync.cc b/tensorflow/compiler/mlir/tfrt/ir/tfrt_fallback_sync.cc
index 8083fcac076745..6a429ef275e869 100644
--- a/tensorflow/compiler/mlir/tfrt/ir/tfrt_fallback_sync.cc
+++ b/tensorflow/compiler/mlir/tfrt/ir/tfrt_fallback_sync.cc
@@ -30,6 +30,7 @@ limitations under the License.
 #include "tfrt/core_runtime/opdefs/attributes.h"  // from @tf_runtime
 #include "tfrt/core_runtime/opdefs/core_runtime.h"  // from @tf_runtime
 #include "tfrt/core_runtime/opdefs/types.h"  // from @tf_runtime
+#include "tfrt/tensor/opdefs/tensor.h"  // from @tf_runtime
 
 namespace tfrt {
 namespace fallback_sync {
@@ -50,7 +51,7 @@ FallbackSyncDialect::FallbackSyncDialect(MLIRContext *context)
 }
 
 static Type GetTensorType(Builder *builder) {
-  return tfrt::t::TensorType::get(builder->getContext());
+  return tfrt::tfrt_tensor::TensorType::get(builder->getContext());
 }
 
 }  // namespace fallback_sync

From 46769b13686f4620fa7950e8f97f1f4a1de423d9 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Wed, 15 May 2024 17:17:38 -0700
Subject: [PATCH 231/478] Update ops-related pbtxt files.

PiperOrigin-RevId: 634129215
---
 tensorflow/core/ops/ops.pbtxt | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorflow/core/ops/ops.pbtxt b/tensorflow/core/ops/ops.pbtxt
index a15dada136af38..846207488b92f3 100644
--- a/tensorflow/core/ops/ops.pbtxt
+++ b/tensorflow/core/ops/ops.pbtxt
@@ -1,4 +1,4 @@
-go/nodeserialize  
+go/debugstr   
 op {
   name: "Abort"
   attr {

From 232e49fbc9cbfe586c8619a3531ea03ce5326741 Mon Sep 17 00:00:00 2001
From: Yang Chen <yangchen@google.com>
Date: Wed, 15 May 2024 18:16:05 -0700
Subject: [PATCH 232/478] [P0] Clean up unused dependency on xla_data.proto.h.

This is attempting to fix this error:
```
File already exists in database: xla/xla_data.proto.
```

PiperOrigin-RevId: 634144823
---
 third_party/xla/xla/ffi/BUILD            | 1 -
 third_party/xla/xla/ffi/attribute_map.cc | 1 -
 2 files changed, 2 deletions(-)

diff --git a/third_party/xla/xla/ffi/BUILD b/third_party/xla/xla/ffi/BUILD
index 4da1564a334ad1..2f597ceb904bee 100644
--- a/third_party/xla/xla/ffi/BUILD
+++ b/third_party/xla/xla/ffi/BUILD
@@ -108,7 +108,6 @@ cc_library(
     hdrs = ["attribute_map.h"],
     deps = [
         ":call_frame",
-        "//xla:xla_data_proto_cc_impl",
         "@com_google_absl//absl/status",
         "@com_google_absl//absl/status:statusor",
         "@com_google_absl//absl/strings",
diff --git a/third_party/xla/xla/ffi/attribute_map.cc b/third_party/xla/xla/ffi/attribute_map.cc
index 9443a8c010d542..8ef42c30085b19 100644
--- a/third_party/xla/xla/ffi/attribute_map.cc
+++ b/third_party/xla/xla/ffi/attribute_map.cc
@@ -26,7 +26,6 @@ limitations under the License.
 #include "mlir/IR/BuiltinAttributes.h"  // from @llvm-project
 #include "mlir/Support/LLVM.h"  // from @llvm-project
 #include "xla/ffi/call_frame.h"
-#include "xla/xla_data.pb.h"
 #include "tsl/platform/errors.h"
 
 using FlatAttribute = xla::ffi::CallFrameBuilder::FlatAttribute;

From 09675d1e4f27a889eee7b596d7565033bab02b29 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Wed, 15 May 2024 19:17:36 -0700
Subject: [PATCH 233/478] Update ops-related pbtxt files.

PiperOrigin-RevId: 634161496
---
 tensorflow/core/ops/ops.pbtxt | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorflow/core/ops/ops.pbtxt b/tensorflow/core/ops/ops.pbtxt
index 846207488b92f3..cdc0191b8ad6c9 100644
--- a/tensorflow/core/ops/ops.pbtxt
+++ b/tensorflow/core/ops/ops.pbtxt
@@ -1,4 +1,4 @@
-go/debugstr   
+go/debugproto   
 op {
   name: "Abort"
   attr {

From 7201a35ce827ce4a40be9de71753958a3310c87b Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Wed, 15 May 2024 19:48:21 -0700
Subject: [PATCH 234/478] Update TFRT dependency to use revision
 http://github.com/tensorflow/runtime/commit/76dfd9f8757c12739429c8178d1b0c5167862ae7.

PiperOrigin-RevId: 634168734
---
 third_party/tf_runtime/workspace.bzl                          | 4 ++--
 .../xla/third_party/tsl/third_party/tf_runtime/workspace.bzl  | 4 ++--
 2 files changed, 4 insertions(+), 4 deletions(-)

diff --git a/third_party/tf_runtime/workspace.bzl b/third_party/tf_runtime/workspace.bzl
index ee7e50c4cb4178..2b6fd2f3c8a0d5 100644
--- a/third_party/tf_runtime/workspace.bzl
+++ b/third_party/tf_runtime/workspace.bzl
@@ -6,8 +6,8 @@ def repo():
     """Imports TFRT."""
 
     # Attention: tools parse and update these lines.
-    TFRT_COMMIT = "ac3f51614db030270303c593a41f46bea68dab3d"
-    TFRT_SHA256 = "a904f3ee171c84c2b5e1bdbaf3f7a0faa43f036383b17f1f69c8abbda3f0936f"
+    TFRT_COMMIT = "76dfd9f8757c12739429c8178d1b0c5167862ae7"
+    TFRT_SHA256 = "e607dd6d42ed911021615a7ac5e41ad26190724888e2742544e33bc28a97ff71"
 
     tf_http_archive(
         name = "tf_runtime",
diff --git a/third_party/xla/third_party/tsl/third_party/tf_runtime/workspace.bzl b/third_party/xla/third_party/tsl/third_party/tf_runtime/workspace.bzl
index ee7e50c4cb4178..2b6fd2f3c8a0d5 100644
--- a/third_party/xla/third_party/tsl/third_party/tf_runtime/workspace.bzl
+++ b/third_party/xla/third_party/tsl/third_party/tf_runtime/workspace.bzl
@@ -6,8 +6,8 @@ def repo():
     """Imports TFRT."""
 
     # Attention: tools parse and update these lines.
-    TFRT_COMMIT = "ac3f51614db030270303c593a41f46bea68dab3d"
-    TFRT_SHA256 = "a904f3ee171c84c2b5e1bdbaf3f7a0faa43f036383b17f1f69c8abbda3f0936f"
+    TFRT_COMMIT = "76dfd9f8757c12739429c8178d1b0c5167862ae7"
+    TFRT_SHA256 = "e607dd6d42ed911021615a7ac5e41ad26190724888e2742544e33bc28a97ff71"
 
     tf_http_archive(
         name = "tf_runtime",

From b55d4f19569e850cb11f27de8f7596a7bb6046d5 Mon Sep 17 00:00:00 2001
From: Anlun Xu <anlunx@google.com>
Date: Wed, 15 May 2024 19:55:33 -0700
Subject: [PATCH 235/478] Reverts adf106883e61d1c0d5583b08491fec91cd74420e

PiperOrigin-RevId: 634170429
---
 third_party/xla/xla/service/gpu/BUILD         |  4 +-
 .../xla/xla/service/gpu/gemm_rewriter.cc      |  6 +--
 .../xla/xla/service/gpu/gpu_compiler.cc       | 54 +++++++------------
 .../xla/service/gpu/ir_emitter_triton_test.cc | 35 +++---------
 .../xla/xla/service/gpu/nvptx_compiler.cc     | 18 +++----
 5 files changed, 36 insertions(+), 81 deletions(-)

diff --git a/third_party/xla/xla/service/gpu/BUILD b/third_party/xla/xla/service/gpu/BUILD
index 475c1e680c00fe..9e99084bf14e92 100644
--- a/third_party/xla/xla/service/gpu/BUILD
+++ b/third_party/xla/xla/service/gpu/BUILD
@@ -3388,7 +3388,6 @@ cc_library(
         ":pipelined_p2p_rewriter",
         ":collective_permute_cycle_decomposer",
         ":address_computation_fusion_rewriter",
-        ":gemv_rewriter",
         ":algorithm_checker",
         ":alias_passthrough_params",
         ":all_reduce_blueconnect",
@@ -3407,7 +3406,6 @@ cc_library(
         ":gemm_broadcast_folding_rewriter",
         ":gemm_fusion",
         ":gemm_rewriter",
-        ":gpu_algebraic_simplifier",
         ":gpu_all_gather_optimizer",
         ":gpu_async_collective_annotator",
         ":gpu_constants",
@@ -3758,7 +3756,6 @@ cc_library(
         ":dot_sparsity_rewriter",
         ":gemm_algorithm_picker",
         ":gemm_fusion_autotuner",
-        ":gpu_algebraic_simplifier",
         ":gpu_asm_opts_util",
         ":gpu_compiler",
         ":gpu_conv_padding_legalization",
@@ -3797,6 +3794,7 @@ cc_library(
         "//xla:util",
         "//xla:xla_proto_cc",
         "//xla/hlo/ir:hlo",
+        "//xla/service:algebraic_simplifier",
         "//xla/service:call_inliner",
         "//xla/service:convert_mover",
         "//xla/service:dot_dimension_merger",
diff --git a/third_party/xla/xla/service/gpu/gemm_rewriter.cc b/third_party/xla/xla/service/gpu/gemm_rewriter.cc
index 8d4064c14d5d5e..2f9c0ac6b35839 100644
--- a/third_party/xla/xla/service/gpu/gemm_rewriter.cc
+++ b/third_party/xla/xla/service/gpu/gemm_rewriter.cc
@@ -537,10 +537,8 @@ class GemmRewriterVisitor : public DfsHloRewriteVisitor {
     }
 
     CHECK(!instr->IsRank2Transpose());
-    if (instr->operand(0)->IsRank2Transpose() ||
-        instr->operand(1)->IsRank2Transpose()) {
-      return absl::OkStatus();
-    }
+    CHECK(!instr->mutable_operand(0)->IsRank2Transpose());
+    CHECK(!instr->mutable_operand(1)->IsRank2Transpose());
     // Create a GemmBackendConfig based on the instruction.
     TF_ASSIGN_OR_RETURN(GpuBackendConfig gpu_backend_config,
                         instr->backend_config<GpuBackendConfig>());
diff --git a/third_party/xla/xla/service/gpu/gpu_compiler.cc b/third_party/xla/xla/service/gpu/gpu_compiler.cc
index d38cfdd2287e54..769e6de904613a 100644
--- a/third_party/xla/xla/service/gpu/gpu_compiler.cc
+++ b/third_party/xla/xla/service/gpu/gpu_compiler.cc
@@ -66,6 +66,7 @@ limitations under the License.
 #include "xla/hlo/ir/hlo_schedule.h"
 #include "xla/hlo/transforms/hlo_constant_splitter.h"
 #include "xla/maybe_owning.h"
+#include "xla/service/algebraic_simplifier.h"
 #include "xla/service/all_gather_broadcast_reorder.h"
 #include "xla/service/all_gather_combiner.h"
 #include "xla/service/all_reduce_combiner.h"
@@ -124,8 +125,6 @@ limitations under the License.
 #include "xla/service/gpu/gemm_broadcast_folding_rewriter.h"
 #include "xla/service/gpu/gemm_fusion.h"
 #include "xla/service/gpu/gemm_rewriter.h"
-#include "xla/service/gpu/gemv_rewriter.h"
-#include "xla/service/gpu/gpu_algebraic_simplifier.h"
 #include "xla/service/gpu/gpu_all_gather_optimizer.h"
 #include "xla/service/gpu/gpu_async_collective_annotator.h"
 #include "xla/service/gpu/gpu_conv_rewriter.h"
@@ -563,10 +562,8 @@ absl::Status RunSPMDPasses(
     HloPassPipeline& spmd_simplify =
         spmd_pipeline.AddPass<HloPassFix<HloPassPipeline>>("spmd-simplify");
 
-    se::GpuComputeCapability gpu_version =
-        gpu_target_config.device_description.gpu_compute_capability();
-    spmd_simplify.AddPass<GpuAlgebraicSimplifier>(
-        layout_insensitive_algsimp_opts, gpu_version);
+    spmd_simplify.AddPass<AlgebraicSimplifier>(layout_insensitive_algsimp_opts);
+
     spmd_simplify.AddPass<SortSimplifier>();
     spmd_simplify.AddPass<TupleSimplifier>();
     spmd_simplify.AddPass<ScatterExpander>(
@@ -766,8 +763,6 @@ absl::Status RunOptimizationPasses(
   }
 
   pipeline.AddPass<DynamicPadder>(dynamic_padder_options);
-  se::GpuComputeCapability gpu_version =
-      gpu_target_config.device_description.gpu_compute_capability();
 
   // Build simplification pipeline.  The passes in here are run to a fixed
   // point.
@@ -785,8 +780,7 @@ absl::Status RunOptimizationPasses(
     pipeline.AddPass<ScatterExpander>(
         ScatterExpander::kEliminateSimpleScatters);
     pipeline.AddPass<ScatterSliceSimplifier>();
-    pipeline.AddPass<GpuAlgebraicSimplifier>(layout_insensitive_algsimp_opts,
-                                             gpu_version);
+    pipeline.AddPass<AlgebraicSimplifier>(layout_insensitive_algsimp_opts);
     pipeline.AddPass<BitcastDtypesExpander>();
     // AlgebraicSimplifier may add contracting dimensions to a dot.
     pipeline.AddPass<DotDimensionSorter>();
@@ -818,8 +812,7 @@ absl::Status RunOptimizationPasses(
   [&, &pipeline =
           pipeline.AddPass<HloPassFix<HloPassPipeline>>("simplification-2")] {
     pipeline.AddPass<ConvertMover>();
-    pipeline.AddPass<GpuAlgebraicSimplifier>(layout_insensitive_algsimp_opts,
-                                             gpu_version);
+    pipeline.AddPass<AlgebraicSimplifier>(layout_insensitive_algsimp_opts);
   }();
 
   pipeline.AddPass<HloComputationDeduplicator>(
@@ -829,8 +822,7 @@ absl::Status RunOptimizationPasses(
 
 absl::Status RunCollectiveOptimizationPasses(
     HloModule* hlo_module,
-    const AlgebraicSimplifierOptions& layout_insensitive_algsimp_opts,
-    se::GpuComputeCapability gpu_version) {
+    const AlgebraicSimplifierOptions& layout_insensitive_algsimp_opts) {
   // Optimize collectives generated by SPMD partitioning. Enable these passes
   // otherwise as well so that all collectives can get these optimizations.
   const DebugOptions& debug_options = hlo_module->config().debug_options();
@@ -913,8 +905,8 @@ absl::Status RunCollectiveOptimizationPasses(
   // Run algebraic simplifier to reshape(broadcast) into a broadcast when
   // the reshape is just adding a unit dimension. This will help with the
   // AllGatherBroadcastReorder pass.
-  collectives_pipeline.AddPass<GpuAlgebraicSimplifier>(
-      layout_insensitive_algsimp_opts, gpu_version);
+  collectives_pipeline.AddPass<AlgebraicSimplifier>(
+      layout_insensitive_algsimp_opts);
 
   collectives_pipeline.AddPass<AllGatherBroadcastReorder>();
 
@@ -1115,12 +1107,11 @@ absl::Status RunPostFusionCollectiveOptimizationPasses(HloModule* hlo_module) {
 
 absl::Status RunPostFusionSimplificationPasses(
     HloModule* hlo_module,
-    const AlgebraicSimplifierOptions& layout_insensitive_algsimp_opts,
-    se::GpuComputeCapability gpu_version) {
+    const AlgebraicSimplifierOptions& layout_insensitive_algsimp_opts) {
   HloPassPipeline pipeline("post-fusion-simplification-pipeline optimization");
   AlgebraicSimplifierOptions options = layout_insensitive_algsimp_opts;
   options.set_is_layout_sensitive(true);
-  pipeline.AddPass<GpuAlgebraicSimplifier>(options, gpu_version);
+  pipeline.AddPass<AlgebraicSimplifier>(options);
 
   // This invocation is used to populate deduplicated_name for fusions that
   // are considered duplicates according to the comparator in this pass.
@@ -1184,13 +1175,13 @@ absl::Status GpuCompiler::OptimizeHloModule(
                                    layout_insensitive_algsimp_opts));
   TF_RETURN_IF_ERROR(RunOptimizationPasses(hlo_module, gpu_target_config,
                                            layout_insensitive_algsimp_opts));
-  se::GpuComputeCapability gpu_version =
-      gpu_target_config.device_description.gpu_compute_capability();
   TF_RETURN_IF_ERROR(RunCollectiveOptimizationPasses(
-      hlo_module, layout_insensitive_algsimp_opts, gpu_version));
+      hlo_module, layout_insensitive_algsimp_opts));
 
   // Run target-specific HLO optimization passes for convolution
   // canonicalization.
+  se::GpuComputeCapability gpu_version =
+      gpu_target_config.device_description.gpu_compute_capability();
   se::dnn::VersionInfo dnn_version = gpu_target_config.dnn_version_info;
   if (stream_exec != nullptr) {
     gpu_version = GetGpuVersion(stream_exec);
@@ -1231,8 +1222,8 @@ absl::Status GpuCompiler::OptimizeHloModule(
         &NormalizeLayoutForGpuCustomCalls);
     // The LayoutAssignment pass may leave behind kCopy instructions which are
     // duplicate or NOPs, so remove them with algebraic simplification and CSE.
-    layout_normalization_pipeline.AddPass<HloPassFix<GpuAlgebraicSimplifier>>(
-        simplifier_options, gpu_version);
+    layout_normalization_pipeline.AddPass<HloPassFix<AlgebraicSimplifier>>(
+        simplifier_options);
   }
   TF_RETURN_IF_ERROR(layout_normalization_pipeline.Run(hlo_module).status());
   // Run target-specific HLO optimization passes after layout assignment.
@@ -1261,7 +1252,7 @@ absl::Status GpuCompiler::OptimizeHloModule(
       }));
   TF_RETURN_IF_ERROR(RunPostFusionCollectiveOptimizationPasses(hlo_module));
   TF_RETURN_IF_ERROR(RunPostFusionSimplificationPasses(
-      hlo_module, layout_insensitive_algsimp_opts, gpu_version));
+      hlo_module, layout_insensitive_algsimp_opts));
 
   TF_RETURN_IF_ERROR(RunPostFusionVerificationPasses(
       hlo_module, stream_exec, options, gpu_target_config));
@@ -1334,8 +1325,7 @@ absl::Status GpuCompiler::OptimizeHloPostLayoutAssignment(
 
     // The LayoutAssignment pass may leave behind kCopy instructions which are
     // duplicate or NOPs, so remove them with algebraic simplification and CSE.
-    pipeline.AddPass<HloPassFix<GpuAlgebraicSimplifier>>(simplifier_options,
-                                                         gpu_version);
+    pipeline.AddPass<HloPassFix<AlgebraicSimplifier>>(simplifier_options);
 
     // GemmRewriter assumes that all transposes are folded into gemms, but,
     // since commit 7d529df, this is not always true at this point.
@@ -1372,7 +1362,6 @@ absl::Status GpuCompiler::OptimizeHloPostLayoutAssignment(
                                    /*f8_rewrite=*/true);
     if (debug_options.xla_gpu_enable_triton_gemm() && cuda_cc != nullptr &&
         cuda_cc->IsAtLeast(se::CudaComputeCapability::AMPERE)) {
-      pipeline.AddPass<GemvRewriter>();
       pipeline.AddPass<GemmFusion>(gpu_version);
     }
     // Rewrite non-FP8 GEMMs.
@@ -1386,8 +1375,7 @@ absl::Status GpuCompiler::OptimizeHloPostLayoutAssignment(
       pipeline.AddPass<LayoutNormalization>(&NormalizeLayoutForGpuCustomCalls);
       // Remove any redundant operations (such as bitcasts) introduced by layout
       // normalization.
-      pipeline.AddPass<HloPassFix<GpuAlgebraicSimplifier>>(simplifier_options,
-                                                           gpu_version);
+      pipeline.AddPass<HloPassFix<AlgebraicSimplifier>>(simplifier_options);
     }
     pipeline.AddPass<BroadcastCanonicalizer>();
 
@@ -1400,8 +1388,7 @@ absl::Status GpuCompiler::OptimizeHloPostLayoutAssignment(
     if (debug_options.xla_gpu_enable_triton_softmax_fusion() &&
         cuda_cc != nullptr &&
         cuda_cc->IsAtLeast(se::CudaComputeCapability::AMPERE)) {
-      pipeline.AddPass<HloPassFix<GpuAlgebraicSimplifier>>(simplifier_options,
-                                                           gpu_version);
+      pipeline.AddPass<HloPassFix<AlgebraicSimplifier>>(simplifier_options);
       pipeline.AddPass<SoftmaxRewriterTriton>(gpu_version);
     }
 
@@ -1464,8 +1451,7 @@ absl::Status GpuCompiler::OptimizeHloPostLayoutAssignment(
 
   // The LayoutAssignment pass may leave behind kCopy instructions which are
   // duplicate or NOPs, so remove them with algebraic simplification and CSE.
-  pipeline.AddPass<HloPassFix<GpuAlgebraicSimplifier>>(simplifier_options,
-                                                       gpu_version);
+  pipeline.AddPass<HloPassFix<AlgebraicSimplifier>>(simplifier_options);
 
   if (debug_options.xla_allow_excess_precision()) {
     // This pass cleans up chains of compiler-generated converts
diff --git a/third_party/xla/xla/service/gpu/ir_emitter_triton_test.cc b/third_party/xla/xla/service/gpu/ir_emitter_triton_test.cc
index a9a0ed4ba69e9d..04f0cda2d3e3bd 100644
--- a/third_party/xla/xla/service/gpu/ir_emitter_triton_test.cc
+++ b/third_party/xla/xla/service/gpu/ir_emitter_triton_test.cc
@@ -2861,19 +2861,6 @@ ENTRY e {
   EXPECT_TRUE(RunAndCompare(kHloText, ErrorSpec{/*aabs=*/1e-2, /*arel=*/1e-2}));
 }
 
-TEST_F(TritonGemmTestWithSplitK, SplitKWithTrivialDimension) {
-  const std::string kHloText = R"(
-ENTRY entry_computation {
-  p0 = f16[1001,1]{1,0} parameter(0)
-  convert = f32[1001,1]{1,0} convert(p0)
-  p1 = f32[1001,2048]{1,0} parameter(1)
-  ROOT dot = f32[1,2048]{1,0} dot(convert, p1),
-    lhs_contracting_dims={0}, rhs_contracting_dims={0}
-})";
-
-  EXPECT_TRUE(RunAndCompare(kHloText, ErrorSpec{/*aabs=*/1e-2, /*arel=*/1e-2}));
-}
-
 TEST_F(TritonGemmLevel2Test, NarrowingConvertOutputIsFused) {
   const std::string kHloText = R"(
 HloModule m
@@ -3066,7 +3053,7 @@ ENTRY e {
 }
 
 TEST_F(TritonGemmTestAny,
-       LowerDotWithLhsWithoutNonContractingDimThroughTriton) {
+       ShouldNotLowerDotWithLhsWithoutNonContractingDimThroughTriton) {
   const std::string hlo_text = R"(
 HloModule t
 
@@ -3076,18 +3063,13 @@ ENTRY e {
   ROOT dot = f32[32,64] dot(f32[32,40] parameter_0, f32[32,40,64] parameter_1), lhs_batch_dims={0}, lhs_contracting_dims={1}, rhs_batch_dims={0}, rhs_contracting_dims={1}
 })";
 
-  TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
-                          GetOptimizedModule(hlo_text));
+  MatchOptimizedHlo(hlo_text, "CHECK-NOT: triton");
 
-  EXPECT_THAT(
-      module->entry_computation()->root_instruction(),
-      GmockMatch(m::Fusion(m::Parameter(), m::Parameter())
-                     .WithFusionKind(HloInstruction::FusionKind::kCustom)));
-  EXPECT_TRUE(RunAndCompare(hlo_text, ErrorSpec{/*aabs=*/1e-3, /*arel=*/1e-3}));
+  EXPECT_TRUE(RunAndCompare(hlo_text, ErrorSpec{/*aabs=*/1e-6, /*arel=*/1e-6}));
 }
 
 TEST_F(TritonGemmTestAny,
-       LowerDotWithRhsWithoutNonContractingDimThroughTriton) {
+       ShouldNotLowerDotWithRhsWithoutNonContractingDimThroughTriton) {
   const std::string hlo_text = R"(
 HloModule t
 
@@ -3097,14 +3079,9 @@ ENTRY e {
   ROOT dot = f32[32,64] dot(f32[32,40,64] parameter_0, f32[32,40] parameter_1), lhs_batch_dims={0}, lhs_contracting_dims={1}, rhs_batch_dims={0}, rhs_contracting_dims={1}
 })";
 
-  TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
-                          GetOptimizedModule(hlo_text));
+  MatchOptimizedHlo(hlo_text, "CHECK-NOT: triton");
 
-  EXPECT_THAT(
-      module->entry_computation()->root_instruction(),
-      GmockMatch(m::Fusion(m::Parameter(), m::Parameter())
-                     .WithFusionKind(HloInstruction::FusionKind::kCustom)));
-  EXPECT_TRUE(RunAndCompare(hlo_text, ErrorSpec{/*aabs=*/1e-3, /*arel=*/1e-3}));
+  EXPECT_TRUE(RunAndCompare(hlo_text, ErrorSpec{/*aabs=*/1e-6, /*arel=*/1e-6}));
 }
 
 // This group of tests compares GPU results of dots already rewritten
diff --git a/third_party/xla/xla/service/gpu/nvptx_compiler.cc b/third_party/xla/xla/service/gpu/nvptx_compiler.cc
index d342a9621a5573..8b02969fca6131 100644
--- a/third_party/xla/xla/service/gpu/nvptx_compiler.cc
+++ b/third_party/xla/xla/service/gpu/nvptx_compiler.cc
@@ -42,6 +42,7 @@ limitations under the License.
 #include "llvm/Support/raw_ostream.h"
 #include "xla/hlo/ir/hlo_computation.h"
 #include "xla/hlo/ir/hlo_opcode.h"
+#include "xla/service/algebraic_simplifier.h"
 #include "xla/service/call_inliner.h"
 #include "xla/service/convert_mover.h"
 #include "xla/service/dot_dimension_merger.h"
@@ -66,7 +67,6 @@ limitations under the License.
 #include "xla/service/gpu/dot_sparsity_rewriter.h"
 #include "xla/service/gpu/gemm_algorithm_picker.h"
 #include "xla/service/gpu/gemm_fusion_autotuner.h"
-#include "xla/service/gpu/gpu_algebraic_simplifier.h"
 #include "xla/service/gpu/gpu_asm_opts_util.h"
 #include "xla/service/gpu/gpu_compiler.h"
 #include "xla/service/gpu/gpu_conv_padding_legalization.h"
@@ -216,8 +216,7 @@ absl::Status NVPTXCompiler::OptimizeHloConvolutionCanonicalization(
       GetAlgebraicSimplifierOptions(hlo_module->config());
   algsimp_options.set_enable_conv_operand_swap(false);
   algsimp_options.set_enable_unconditional_reduce_of_concat_replacement(false);
-  pipeline.AddPass<HloPassFix<GpuAlgebraicSimplifier>>(algsimp_options,
-                                                       gpu_version);
+  pipeline.AddPass<HloPassFix<AlgebraicSimplifier>>(algsimp_options);
 
   // CudnnSimplifyPadding gets rid of some padding introduced by
   // CudnnPadForConvolutions and used by CudnnVectorizeConvolutions.  The
@@ -234,7 +233,7 @@ absl::Status NVPTXCompiler::OptimizeHloConvolutionCanonicalization(
     ReshapeMoverOptions reshape_mover_options;
     reshape_mover_options.reshape_of_1d_broadcast_is_cheap = true;
     pipeline.AddPass<ReshapeMover>(reshape_mover_options);
-    pipeline.AddPass<GpuAlgebraicSimplifier>(algsimp_options, gpu_version);
+    pipeline.AddPass<AlgebraicSimplifier>(algsimp_options);
   }();
 
   // The reshapes and transposes can possibly be eliminated using
@@ -245,7 +244,7 @@ absl::Status NVPTXCompiler::OptimizeHloConvolutionCanonicalization(
   [&, &pipeline = pipeline.AddPass<HloPassFix<HloPassPipeline>>(
           "simplify_after_conv_canonicalization")] {
     pipeline.AddPass<ConvertMover>();
-    pipeline.AddPass<GpuAlgebraicSimplifier>(algsimp_options, gpu_version);
+    pipeline.AddPass<AlgebraicSimplifier>(algsimp_options);
   }();
 
   // GpuConvRewriter, GpuConvPaddingLegalization and
@@ -283,10 +282,8 @@ absl::Status NVPTXCompiler::OptimizeHloPostLayoutAssignment(
         false);
 
     mha_fusion_pipeline.AddPass<HloCSE>(/*is_layout_sensitive=*/true);
-    se::GpuComputeCapability gpu_version =
-        gpu_target_config.device_description.gpu_compute_capability();
-    mha_fusion_pipeline.AddPass<HloPassFix<GpuAlgebraicSimplifier>>(
-        alg_sim_options, gpu_version);
+    mha_fusion_pipeline.AddPass<HloPassFix<AlgebraicSimplifier>>(
+        alg_sim_options);
     mha_fusion_pipeline.AddPass<HloCSE>(/*is_layout_sensitive=*/true);
     // Rewrite Multi-Headed Attention modules to Fused MHA custom-calls.
     if (stream_exec) {
@@ -296,8 +293,7 @@ absl::Status NVPTXCompiler::OptimizeHloPostLayoutAssignment(
       mha_fusion_pipeline.AddPass<CudnnFusedMHARewriter>(
           cuda_compute_capability, gpu_target_config.dnn_version_info);
     }
-    mha_fusion_pipeline.AddPass<GpuAlgebraicSimplifier>(alg_sim_options,
-                                                        gpu_version);
+    mha_fusion_pipeline.AddPass<AlgebraicSimplifier>(alg_sim_options);
     mha_fusion_pipeline.AddPass<CudnnFusedMHATransposeFusion>();
     mha_fusion_pipeline.AddPass<HloDCE>();
     mha_fusion_pipeline.AddPass<HloCSE>(/*is_layout_sensitive=*/true);

From 4e800be1bdc2bc4fb07123527c6b5b3361c9cca7 Mon Sep 17 00:00:00 2001
From: Junwhan Ahn <junwhan@google.com>
Date: Wed, 15 May 2024 21:08:19 -0700
Subject: [PATCH 236/478] Check for `GetEnableMemories()` outside
 `PyDeviceList::PopulateMemoryKindInfo()`

Checking `GetEnableMemories()` inside `PyDeviceList::PopulateMemoryKindInfo()` causes the memory kind to be built/cached based on the value of `enable_memories` at the time the method was called for the first time. This is problematic when users later change the value of `enable_memories` but reuse the same device list.

This CL fixes the problem by making `PyDeviceList::MemoryKinds()` and `PyDeviceList::DefaultMemoryKind()` always check for `GetEnableMemories()` outside the cached method to make it behave correctly even if `enable_memories` changes.

PiperOrigin-RevId: 634191502
---
 third_party/xla/xla/python/py_device_list.cc | 17 ++++++-----------
 1 file changed, 6 insertions(+), 11 deletions(-)

diff --git a/third_party/xla/xla/python/py_device_list.cc b/third_party/xla/xla/python/py_device_list.cc
index f89ff54adb8c04..27110f439b9d98 100644
--- a/third_party/xla/xla/python/py_device_list.cc
+++ b/third_party/xla/xla/python/py_device_list.cc
@@ -326,11 +326,6 @@ void PyDeviceList::PopulateMemoryKindInfo() {
     throw nb::value_error("Unrecognized DeviceList type");
   }
   MemoryKindInfo info;
-  if (!GetEnableMemories()) {
-    info.default_memory_kind = nb::none();
-    memory_kind_info_ = std::move(info);
-    return;
-  }
   xla::ifrt::Device* addressable_device = nullptr;
   const int process_index = py_client_ ? py_client_->process_index() : 0;
   for (xla::ifrt::Device* device : std::get<0>(device_list_).devices()) {
@@ -366,12 +361,6 @@ void PyDeviceList::PopulateMemoryKindInfo() {
 
 void PyDeviceList::PopulateMemoryKindInfoForDuckTypedDevices() {
   MemoryKindInfo info;
-  if (!GetEnableMemories()) {
-    info.default_memory_kind = nb::none();
-    // info.memory_kinds is default-initialized to an empty tuple.
-    memory_kind_info_ = std::move(info);
-    return;
-  }
   try {
     nb::handle addressable_device;
     for (nb::handle device : std::get<1>(device_list_)) {
@@ -399,6 +388,9 @@ void PyDeviceList::PopulateMemoryKindInfoForDuckTypedDevices() {
 }
 
 absl::StatusOr<nb::tuple> PyDeviceList::MemoryKinds() {
+  if (!GetEnableMemories()) {
+    return nb::tuple();
+  }
   if (!memory_kind_info_.has_value()) {
     PopulateMemoryKindInfo();
   }
@@ -409,6 +401,9 @@ absl::StatusOr<nb::tuple> PyDeviceList::MemoryKinds() {
 }
 
 absl::StatusOr<nb::object> PyDeviceList::DefaultMemoryKind() {
+  if (!GetEnableMemories()) {
+    return nb::none();
+  }
   if (!memory_kind_info_.has_value()) {
     PopulateMemoryKindInfo();
   }

From 099c4653a5347e59e6791508f9cad39bd6bea903 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Wed, 15 May 2024 21:17:31 -0700
Subject: [PATCH 237/478] Update ops-related pbtxt files.

PiperOrigin-RevId: 634193569
---
 tensorflow/core/ops/ops.pbtxt | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorflow/core/ops/ops.pbtxt b/tensorflow/core/ops/ops.pbtxt
index cdc0191b8ad6c9..846207488b92f3 100644
--- a/tensorflow/core/ops/ops.pbtxt
+++ b/tensorflow/core/ops/ops.pbtxt
@@ -1,4 +1,4 @@
-go/debugproto   
+go/debugstr   
 op {
   name: "Abort"
   attr {

From c13cbeba32d39fcedb298c4b74119302272b75ff Mon Sep 17 00:00:00 2001
From: Jieying Luo <jieying@google.com>
Date: Wed, 15 May 2024 21:25:01 -0700
Subject: [PATCH 238/478] Propagate error to output if the input buffer has
 error.

PiperOrigin-RevId: 634195071
---
 .../xla/pjrt/gpu/se_gpu_pjrt_client_test.cc   | 35 +++++++++++++++++++
 .../xla/pjrt/pjrt_stream_executor_client.cc   | 32 +++++++++++++++++
 2 files changed, 67 insertions(+)

diff --git a/third_party/xla/xla/pjrt/gpu/se_gpu_pjrt_client_test.cc b/third_party/xla/xla/pjrt/gpu/se_gpu_pjrt_client_test.cc
index d3bb175754e1fb..616ab665e34604 100644
--- a/third_party/xla/xla/pjrt/gpu/se_gpu_pjrt_client_test.cc
+++ b/third_party/xla/xla/pjrt/gpu/se_gpu_pjrt_client_test.cc
@@ -140,6 +140,41 @@ TEST(StreamExecutorGpuClientTest, MemorySpace) {
   }
 }
 
+TEST(StreamExecutorGpuClientTest, PropagateError) {
+  TF_ASSERT_OK_AND_ASSIGN(auto client,
+                          GetStreamExecutorGpuClient(GpuClientOptions()));
+  auto shape = xla::ShapeUtil::MakeScalarShape(xla::F32);
+  absl::Status input_error = absl::InvalidArgumentError("input error");
+  TF_ASSERT_OK_AND_ASSIGN(
+      auto buffer,
+      client->CreateErrorBuffer(
+          input_error, shape,
+          *client->addressable_devices()[0]->default_memory_space()));
+
+  static constexpr char const* kAddProgram =
+      R"(
+HloModule Add.6, entry_computation_layout={(f32[], f32[])->(f32[], f32[])}
+
+ENTRY %Add.6 (a.1: f32[], b.2: f32[]) -> (f32[], f32[]) {
+  %a.1 = f32[] parameter(0)
+  %b.2 = f32[] parameter(1)
+  %add.3 = f32[] add(f32[] %a.1, f32[] %b.2)
+  %add.4 = f32[] add(f32[] %add.3, f32[] %add.3)
+  ROOT %tuple.5 = (f32[], f32[]) tuple(f32[] %add.3, f32[] %add.4)
+}
+)";
+  TF_ASSERT_OK_AND_ASSIGN(auto executable,
+                          CompileExecutable(kAddProgram, *client));
+
+  TF_ASSERT_OK_AND_ASSIGN(
+      auto result,
+      executable->Execute({{buffer.get(), buffer.get()}}, /*options=*/{}));
+
+  ASSERT_EQ(result.size(), 1);
+  ASSERT_EQ(result[0].size(), 1);
+  EXPECT_EQ(result[0][0]->GetReadyFuture().Await(), input_error);
+}
+
 TEST(StreamExecutorGpuClientTest, SendRecvChunked) {
   TF_ASSERT_OK_AND_ASSIGN(auto client,
                           GetStreamExecutorGpuClient(GpuClientOptions()));
diff --git a/third_party/xla/xla/pjrt/pjrt_stream_executor_client.cc b/third_party/xla/xla/pjrt/pjrt_stream_executor_client.cc
index b59fd9f145de05..fcf28b0773a118 100644
--- a/third_party/xla/xla/pjrt/pjrt_stream_executor_client.cc
+++ b/third_party/xla/xla/pjrt/pjrt_stream_executor_client.cc
@@ -2875,6 +2875,21 @@ PjRtStreamExecutorLoadedExecutable::MakeOutputBuffers(
   return outputs;
 }
 
+static Status GetFirstInputError(
+    absl::Span<PjRtBuffer* const> argument_handles) {
+  for (auto* handle : argument_handles) {
+    auto* buffer = tensorflow::down_cast<PjRtStreamExecutorBuffer*>(handle);
+    PjRtStreamExecutorBuffer::ScopedHold hold =
+        buffer->GetBufferWithUsageHold();
+    for (const auto& event : hold->definition_events()) {
+      if (event->IsPredeterminedError()) {
+        return event->GetDefinedStatus();
+      }
+    }
+  }
+  return OkStatus();
+}
+
 StatusOr<PjRtLoadedExecutable::Result>
 PjRtStreamExecutorLoadedExecutable::ExecuteHelper(
     absl::Span<PjRtBuffer* const> argument_handles, int replica, int partition,
@@ -2897,6 +2912,23 @@ PjRtStreamExecutorLoadedExecutable::ExecuteHelper(
     (*device_assignment)(0, 0) = device->id();
   }
 
+  Status input_error = GetFirstInputError(argument_handles);
+  if (!input_error.ok()) {
+    TF_ASSIGN_OR_RETURN(PjRtMemorySpace * memory_space,
+                        device->default_memory_space());
+    std::vector<std::unique_ptr<PjRtBuffer>> outputs;
+    TF_ASSIGN_OR_RETURN(auto hlo_modules, GetHloModules());
+    for (const auto& hlo_module : hlo_modules) {
+      TF_ASSIGN_OR_RETURN(
+          auto error_buffer,
+          client_->CreateErrorBuffer(input_error, hlo_module->result_shape(),
+                                     memory_space));
+      outputs.push_back(std::move(error_buffer));
+    }
+    auto future = std::make_optional(PjRtFuture<>(input_error));
+    return Result({std::move(future), /*buffers=*/std::move(outputs)});
+  }
+
   CHECK_EQ(device->process_index(), client_->process_index());
   int device_ordinal = tensorflow::down_cast<PjRtStreamExecutorDevice*>(device)
                            ->local_device_state()

From 04ab4ee799529cf1826b41d18a3e38800295fda9 Mon Sep 17 00:00:00 2001
From: Hyeontaek Lim <hyeontaek@google.com>
Date: Wed, 15 May 2024 22:38:14 -0700
Subject: [PATCH 239/478] [IFRT] Wrap RemapPlan::mappings with shared_ptr

`xla::ifrt::RemapPlan`'s `mappings` is frequently reused across different
`xla::ifrt::RemapPlan`. It can be costly to keep a copy of `mappings` for each
plan if the mappings is not trivially small.

This change wraps `mappings` with `std::shared_ptr` so that it can be shared
more cheaply across multiple plans.

PiperOrigin-RevId: 634214931
---
 .../xla/python/ifrt/remap_impl_test_lib.cc    | 15 +++++---
 third_party/xla/xla/python/ifrt/remap_plan.cc | 15 ++++----
 third_party/xla/xla/python/ifrt/remap_plan.h  |  3 +-
 .../xla/xla/python/ifrt/remap_plan_test.cc    | 37 ++++++++++++-------
 .../xla/xla/python/pjrt_ifrt/pjrt_remap.cc    |  4 +-
 5 files changed, 44 insertions(+), 30 deletions(-)

diff --git a/third_party/xla/xla/python/ifrt/remap_impl_test_lib.cc b/third_party/xla/xla/python/ifrt/remap_impl_test_lib.cc
index 7ea29058621f79..9486e0409412b7 100644
--- a/third_party/xla/xla/python/ifrt/remap_impl_test_lib.cc
+++ b/third_party/xla/xla/python/ifrt/remap_impl_test_lib.cc
@@ -160,7 +160,8 @@ TEST(RemapImplTest, ExtractSingleShard) {
                     MemoryKind(), /*shape=*/Shape({2, 3}),
                     /*shard_shape=*/Shape({2, 3}))});
   // arrays[0].shards[1:2:1] is mapped into out_arrays[0].shards[0:1:1].
-  plan.mappings.push_back(
+  plan.mappings = std::make_shared<std::vector<RemapPlan::Mapping>>();
+  plan.mappings->push_back(
       RemapPlan::Mapping{/*in_array=*/0, /*out_array=*/0,
                          /*from=*/{RemapPlan::Interval{1, 2, 1}},
                          /*to=*/{RemapPlan::Interval{0, 1, 1}}});
@@ -221,12 +222,14 @@ TEST(RemapImplTest, InterleaveArrays) {
                     MemoryKind(), /*shape=*/Shape({8, 3}),
                     /*shard_shape=*/Shape({2, 3}))});
   // arrays[0].shards[0:2:1] is mapped into out_arrays[0].shards[0:4:2].
-  plan.mappings.push_back(
+  plan.mappings = std::make_shared<std::vector<RemapPlan::Mapping>>();
+  plan.mappings->reserve(2);
+  plan.mappings->push_back(
       RemapPlan::Mapping{/*in_array=*/0, /*out_array=*/0,
                          /*from=*/{RemapPlan::Interval{0, 2, 1}},
                          /*to=*/{RemapPlan::Interval{0, 4, 2}}});
   // arrays[1].shards[0:2:1] is mapped into out_arrays[0].shards[1:4:2].
-  plan.mappings.push_back(
+  plan.mappings->push_back(
       RemapPlan::Mapping{/*in_array=*/1, /*out_array=*/0,
                          /*from=*/{RemapPlan::Interval{0, 2, 1}},
                          /*to=*/{RemapPlan::Interval{1, 4, 2}}});
@@ -289,12 +292,14 @@ TEST(RemapImplTest, DeinterleaveArrays) {
                     MemoryKind(), /*shape=*/Shape({4, 3}),
                     /*shard_shape=*/Shape({2, 3}))});
   // arrays[0].shards[0:4:2] is mapped into out_arrays[0].shards[0:2:1].
-  plan.mappings.push_back(
+  plan.mappings = std::make_shared<std::vector<RemapPlan::Mapping>>();
+  plan.mappings->reserve(2);
+  plan.mappings->push_back(
       RemapPlan::Mapping{/*in_array=*/0, /*out_array=*/0,
                          /*from=*/{RemapPlan::Interval{0, 4, 2}},
                          /*to=*/{RemapPlan::Interval{0, 2, 1}}});
   // arrays[0].shards[1:4:2] is mapped into out_arrays[1].shards[0:2:1].
-  plan.mappings.push_back(
+  plan.mappings->push_back(
       RemapPlan::Mapping{/*in_array=*/0, /*out_array=*/1,
                          /*from=*/{RemapPlan::Interval{1, 4, 2}},
                          /*to=*/{RemapPlan::Interval{0, 2, 1}}});
diff --git a/third_party/xla/xla/python/ifrt/remap_plan.cc b/third_party/xla/xla/python/ifrt/remap_plan.cc
index 2f019ea88a5f11..af4ac039d770b0 100644
--- a/third_party/xla/xla/python/ifrt/remap_plan.cc
+++ b/third_party/xla/xla/python/ifrt/remap_plan.cc
@@ -180,8 +180,8 @@ absl::Status RemapPlan::Validate() const {
         /*v=*/nullptr);
   }
 
-  for (int64_t i = 0; i < mappings.size(); ++i) {
-    const RemapPlan::Mapping& mapping = mappings[i];
+  for (int64_t i = 0; i < mappings->size(); ++i) {
+    const RemapPlan::Mapping& mapping = (*mappings)[i];
     if (mapping.in_array < 0 || mapping.in_array >= num_inputs) {
       return InvalidArgument(
           "mappings[%d].in_array must be in [0, %d], but is %d", i,
@@ -280,10 +280,11 @@ absl::StatusOr<RemapPlan> RemapPlan::FromProto(
     plan.output_specs.push_back(std::move(output_spec));
   }
 
-  plan.mappings.reserve(proto.mappings_size());
+  plan.mappings = std::make_shared<std::vector<Mapping>>();
+  plan.mappings->reserve(proto.mappings_size());
   for (const auto& mapping_proto : proto.mappings()) {
     TF_ASSIGN_OR_RETURN(auto mapping, MappingFromProto(mapping_proto));
-    plan.mappings.push_back(std::move(mapping));
+    plan.mappings->push_back(std::move(mapping));
   }
 
   return plan;
@@ -301,8 +302,8 @@ absl::StatusOr<RemapPlanProto> RemapPlan::ToProto() const {
     TF_ASSIGN_OR_RETURN(*proto.add_output_specs(), output_spec.ToProto());
   }
 
-  proto.mutable_mappings()->Reserve(mappings.size());
-  for (const auto& mapping : mappings) {
+  proto.mutable_mappings()->Reserve(mappings->size());
+  for (const auto& mapping : *mappings) {
     TF_ASSIGN_OR_RETURN(*proto.add_mappings(), MappingToProto(mapping));
   }
 
@@ -330,7 +331,7 @@ std::string RemapPlan::DebugString() const {
   };
   return absl::StrCat(
       "RemapPlan(output_specs=", format_array_specs(output_specs), ",",
-      "mappings=", format_mappings(mappings), ")");
+      "mappings=", format_mappings(*mappings), ")");
 }
 
 }  // namespace ifrt
diff --git a/third_party/xla/xla/python/ifrt/remap_plan.h b/third_party/xla/xla/python/ifrt/remap_plan.h
index 7d862e151124ab..22679c6f27f14f 100644
--- a/third_party/xla/xla/python/ifrt/remap_plan.h
+++ b/third_party/xla/xla/python/ifrt/remap_plan.h
@@ -17,6 +17,7 @@ limitations under the License.
 #define XLA_PYTHON_IFRT_REMAP_PLAN_H_
 
 #include <cstdint>
+#include <memory>
 #include <string>
 #include <vector>
 
@@ -82,7 +83,7 @@ struct RemapPlan {
   std::vector<ArraySpec> output_specs;
 
   // Mappings.
-  std::vector<Mapping> mappings;
+  std::shared_ptr<std::vector<Mapping>> mappings;
 
   // Validates this plan against the requirements (see `RemapPlan` comment).
   // This is a slow operation. It should not be performed repeatedly.
diff --git a/third_party/xla/xla/python/ifrt/remap_plan_test.cc b/third_party/xla/xla/python/ifrt/remap_plan_test.cc
index b94b11e0e31063..69865ad659419e 100644
--- a/third_party/xla/xla/python/ifrt/remap_plan_test.cc
+++ b/third_party/xla/xla/python/ifrt/remap_plan_test.cc
@@ -67,12 +67,13 @@ TEST_P(RemapPlanTest, ToFromProto) {
   plan.output_specs.push_back(ArraySpec{
       /*dtype=*/DType(DType::kF32), /*shape=*/shape, /*sharding=*/sharding});
 
-  plan.mappings.reserve(2);
-  plan.mappings.push_back(RemapPlan::Mapping{
+  plan.mappings = std::make_shared<std::vector<RemapPlan::Mapping>>();
+  plan.mappings->reserve(2);
+  plan.mappings->push_back(RemapPlan::Mapping{
       /*in_array=*/0, /*out_array=*/1,
       /*from=*/{RemapPlan::Interval{0, 2, 1}, RemapPlan::Interval{2, 4, 1}},
       /*to=*/{RemapPlan::Interval{1, 4, 2}, RemapPlan::Interval{0, 4, 2}}});
-  plan.mappings.push_back(RemapPlan::Mapping{
+  plan.mappings->push_back(RemapPlan::Mapping{
       /*in_array=*/1, /*out_array=*/0,
       /*from=*/{RemapPlan::Interval{0, 4, 2}, RemapPlan::Interval{1, 4, 2}},
       /*to=*/{RemapPlan::Interval{0, 2, 1}, RemapPlan::Interval{2, 4, 1}}});
@@ -83,7 +84,7 @@ TEST_P(RemapPlanTest, ToFromProto) {
       RemapPlan::FromProto(absl::bind_front(&Client::LookupDevice, client()),
                            plan_proto));
 
-  EXPECT_THAT(plan_copy.mappings, ElementsAreArray(plan.mappings));
+  EXPECT_THAT(*plan_copy.mappings, ElementsAreArray(*plan.mappings));
 
   EXPECT_THAT(plan_copy.output_specs, SizeIs(2));
   for (const auto& spec : plan_copy.input_specs) {
@@ -166,7 +167,8 @@ TEST_P(RemapPlanTest, InvalidInputArrayIndex) {
                 ConcreteEvenSharding::Create(GetDevices({0}), MemoryKind(),
                                              /*shape=*/Shape({2, 3}),
                                              /*shard_shape=*/Shape({2, 3}))});
-  plan.mappings.push_back(
+  plan.mappings = std::make_shared<std::vector<RemapPlan::Mapping>>();
+  plan.mappings->push_back(
       RemapPlan::Mapping{/*in_array=*/1,  // Invalid in_array
                          /*out_array=*/0,
                          /*from=*/{RemapPlan::Interval{0, 1, 1}},
@@ -193,7 +195,8 @@ TEST_P(RemapPlanTest, InvalidOutputArrayIndex) {
                 ConcreteEvenSharding::Create(GetDevices({0}), MemoryKind(),
                                              /*shape=*/Shape({2, 3}),
                                              /*shard_shape=*/Shape({2, 3}))});
-  plan.mappings.push_back(
+  plan.mappings = std::make_shared<std::vector<RemapPlan::Mapping>>();
+  plan.mappings->push_back(
       RemapPlan::Mapping{/*in_array=*/0,
                          /*out_array=*/1,  // Invalid out_array
                          /*from=*/{RemapPlan::Interval{0, 1, 1}},
@@ -220,7 +223,8 @@ TEST_P(RemapPlanTest, InvalidIntervalCount) {
                 ConcreteEvenSharding::Create(GetDevices({0}), MemoryKind(),
                                              /*shape=*/Shape({2, 3}),
                                              /*shard_shape=*/Shape({2, 3}))});
-  plan.mappings.push_back(RemapPlan::Mapping{
+  plan.mappings = std::make_shared<std::vector<RemapPlan::Mapping>>();
+  plan.mappings->push_back(RemapPlan::Mapping{
       /*in_array=*/0,
       /*out_array=*/0,
       /*from=*/{RemapPlan::Interval{0, 1, 1}, RemapPlan::Interval{0, 1, 1}},
@@ -250,9 +254,10 @@ TEST_P(RemapPlanTest, InvalidShardIndex) {
                   ConcreteEvenSharding::Create(GetDevices({0}), MemoryKind(),
                                                /*shape=*/Shape({2, 3}),
                                                /*shard_shape=*/Shape({2, 3}))});
-    plan.mappings.push_back(RemapPlan::Mapping{/*in_array=*/0, /*out_array=*/0,
-                                               /*from=*/{from},
-                                               /*to=*/{to}});
+    plan.mappings = std::make_shared<std::vector<RemapPlan::Mapping>>();
+    plan.mappings->push_back(RemapPlan::Mapping{/*in_array=*/0, /*out_array=*/0,
+                                                /*from=*/{from},
+                                                /*to=*/{to}});
     return plan.Validate();
   };
 
@@ -306,7 +311,8 @@ TEST_P(RemapPlanTest, AlreadyUsedInputShard) {
                 ConcreteEvenSharding::Create(GetDevices({0, 1}), MemoryKind(),
                                              /*shape=*/Shape({4, 3}),
                                              /*shard_shape=*/Shape({2, 3}))});
-  plan.mappings.push_back(RemapPlan::Mapping{
+  plan.mappings = std::make_shared<std::vector<RemapPlan::Mapping>>();
+  plan.mappings->push_back(RemapPlan::Mapping{
       /*in_array=*/0,
       /*out_array=*/0,
       /*from=*/{RemapPlan::Interval{0, 1, 1}, RemapPlan::Interval{0, 1, 1}},
@@ -332,7 +338,8 @@ TEST_P(RemapPlanTest, UnassignedOutputShard) {
                 ConcreteEvenSharding::Create(GetDevices({0, 1}), MemoryKind(),
                                              /*shape=*/Shape({4, 3}),
                                              /*shard_shape=*/Shape({2, 3}))});
-  plan.mappings.push_back(
+  plan.mappings = std::make_shared<std::vector<RemapPlan::Mapping>>();
+  plan.mappings->push_back(
       RemapPlan::Mapping{/*in_array=*/0,
                          /*out_array=*/0,
                          /*from=*/{RemapPlan::Interval{0, 1, 1}},
@@ -358,7 +365,8 @@ TEST_P(RemapPlanTest, AlreadyAssignedOutputShard) {
                 ConcreteEvenSharding::Create(GetDevices({0}), MemoryKind(),
                                              /*shape=*/Shape({2, 3}),
                                              /*shard_shape=*/Shape({2, 3}))});
-  plan.mappings.push_back(RemapPlan::Mapping{
+  plan.mappings = std::make_shared<std::vector<RemapPlan::Mapping>>();
+  plan.mappings->push_back(RemapPlan::Mapping{
       /*in_array=*/0,
       /*out_array=*/0,
       /*from=*/{RemapPlan::Interval{0, 1, 1}, RemapPlan::Interval{1, 2, 1}},
@@ -385,7 +393,8 @@ TEST_P(RemapPlanTest, InvalidOutputDevices) {
                 ConcreteEvenSharding::Create(GetDevices({1, 0}), MemoryKind(),
                                              /*shape=*/Shape({4, 3}),
                                              /*shard_shape=*/Shape({2, 3}))});
-  plan.mappings.push_back(
+  plan.mappings = std::make_shared<std::vector<RemapPlan::Mapping>>();
+  plan.mappings->push_back(
       RemapPlan::Mapping{/*in_array=*/0,
                          /*out_array=*/0,
                          /*from=*/{RemapPlan::Interval{0, 2, 1}},
diff --git a/third_party/xla/xla/python/pjrt_ifrt/pjrt_remap.cc b/third_party/xla/xla/python/pjrt_ifrt/pjrt_remap.cc
index 394eacaf855c06..f72a1761b6e75b 100644
--- a/third_party/xla/xla/python/pjrt_ifrt/pjrt_remap.cc
+++ b/third_party/xla/xla/python/pjrt_ifrt/pjrt_remap.cc
@@ -52,8 +52,6 @@ PjRtCompatibleClientRemapArrays(
           arrays[i]->DebugString());
     }
   }
-  VLOG(0) << "plan.input_specs.size(): " << plan.input_specs.size();
-  VLOG(0) << "plan.output_specs.size(): " << plan.output_specs.size();
   if (plan.input_specs.size() > 1 || plan.output_specs.size() > 1) {
     if (semantics != ArrayCopySemantics::kDonateInput) {
       return InvalidArgument(
@@ -67,7 +65,7 @@ PjRtCompatibleClientRemapArrays(
     out_buffers_list[i].resize(plan.output_specs[i].sharding->devices().size());
   }
 
-  for (const RemapPlan::Mapping& mapping : plan.mappings) {
+  for (const RemapPlan::Mapping& mapping : *plan.mappings) {
     TF_ASSIGN_OR_RETURN(
         absl::Span<std::shared_ptr<xla::PjRtBuffer>> in_buffers,
         static_cast<PjRtCompatibleArray*>(arrays[mapping.in_array].get())

From c0cbc2fc9ab638b46c9c01fc0c1891db4ac2b3f4 Mon Sep 17 00:00:00 2001
From: Christian Sigg <csigg@google.com>
Date: Wed, 15 May 2024 22:45:26 -0700
Subject: [PATCH 240/478] Update to Triton version that contains previously
 patched fixes.

PiperOrigin-RevId: 634216731
---
 .../triton/temporary/mma_limit_pred.patch     | 23 -------------------
 third_party/triton/temporary/series.bzl       |  1 -
 third_party/triton/workspace.bzl              |  4 ++--
 .../triton/temporary/mma_limit_pred.patch     | 23 -------------------
 .../third_party/triton/temporary/series.bzl   |  1 -
 .../xla/third_party/triton/workspace.bzl      |  4 ++--
 6 files changed, 4 insertions(+), 52 deletions(-)
 delete mode 100644 third_party/triton/temporary/mma_limit_pred.patch
 delete mode 100644 third_party/xla/third_party/triton/temporary/mma_limit_pred.patch

diff --git a/third_party/triton/temporary/mma_limit_pred.patch b/third_party/triton/temporary/mma_limit_pred.patch
deleted file mode 100644
index cde501d0a9b6fd..00000000000000
--- a/third_party/triton/temporary/mma_limit_pred.patch
+++ /dev/null
@@ -1,23 +0,0 @@
-// This patch fixes issues that broke internal benchmarks as well as individual
-// fusions reported. It should be moved to our internal patches in the next
-// integration.
-
-diff --git a/lib/Dialect/TritonGPU/Transforms/AccelerateMatmul.cpp b/lib/Dialect/TritonGPU/Transforms/AccelerateMatmul.cpp
---- a/lib/Dialect/TritonGPU/Transforms/AccelerateMatmul.cpp
-+++ b/lib/Dialect/TritonGPU/Transforms/AccelerateMatmul.cpp
-@@ -133,6 +133,15 @@ class BlockedToMMA : public mlir::Rewrit
-   mutable llvm::DenseMap<Operation *, unsigned> dotOpInstNs;
- 
-   static bool bwdFilter(Operation *op) {
-+    // Dot operand layout assignment to Predicates are not currently supported
-+    // during lowering from TritonGPU to LLVM in Triton for MMA cases. This
-+    // condition limits visibility of the original bit-width so that predicate
-+    // are not considered, hence, kwidth can never be = 32.
-+    if (isa<arith::UIToFPOp>(op)) {
-+      Type srcType = getElementTypeOrSelf(op->getOperand(0));
-+      if (srcType.isInteger(1))
-+        return false;
-+    }
-     return op->getNumOperands() == 1 &&
-            (isa<tt::FpToFpOp, tt::BitcastOp, ttg::ConvertLayoutOp>(op) ||
-             isPureUnaryInlineAsm(op) ||
diff --git a/third_party/triton/temporary/series.bzl b/third_party/triton/temporary/series.bzl
index d68c23d5babbd6..a5dfa7d5eae3f6 100644
--- a/third_party/triton/temporary/series.bzl
+++ b/third_party/triton/temporary/series.bzl
@@ -7,6 +7,5 @@ internal patch during the next triton integration process.
 
 temporary_patch_list = [
     "//third_party/triton/temporary:fp8_splat.patch",
-    "//third_party/triton/temporary:mma_limit_pred.patch",
     "//third_party/triton/temporary:fix_register_constraints.patch",
 ]
diff --git a/third_party/triton/workspace.bzl b/third_party/triton/workspace.bzl
index 933b815253ffcd..4c21ab7d7c9099 100644
--- a/third_party/triton/workspace.bzl
+++ b/third_party/triton/workspace.bzl
@@ -8,8 +8,8 @@ load("//third_party/triton/xla_extensions:series.bzl", "extensions_files_patch_l
 def repo():
     """Imports Triton."""
 
-    TRITON_COMMIT = "cl632952902"
-    TRITON_SHA256 = "f83c4f197cc2ae3b034070ec8189007451093edc445d1e383eb17a1e3808de9a"
+    TRITON_COMMIT = "cl633653394"
+    TRITON_SHA256 = "4bf1b25ce6f7ce6ab6fb4e2403b24c8c0bf71f6a9f5c08f0125fdb44bcb371a8"
     tf_http_archive(
         name = "triton",
         sha256 = TRITON_SHA256,
diff --git a/third_party/xla/third_party/triton/temporary/mma_limit_pred.patch b/third_party/xla/third_party/triton/temporary/mma_limit_pred.patch
deleted file mode 100644
index cde501d0a9b6fd..00000000000000
--- a/third_party/xla/third_party/triton/temporary/mma_limit_pred.patch
+++ /dev/null
@@ -1,23 +0,0 @@
-// This patch fixes issues that broke internal benchmarks as well as individual
-// fusions reported. It should be moved to our internal patches in the next
-// integration.
-
-diff --git a/lib/Dialect/TritonGPU/Transforms/AccelerateMatmul.cpp b/lib/Dialect/TritonGPU/Transforms/AccelerateMatmul.cpp
---- a/lib/Dialect/TritonGPU/Transforms/AccelerateMatmul.cpp
-+++ b/lib/Dialect/TritonGPU/Transforms/AccelerateMatmul.cpp
-@@ -133,6 +133,15 @@ class BlockedToMMA : public mlir::Rewrit
-   mutable llvm::DenseMap<Operation *, unsigned> dotOpInstNs;
- 
-   static bool bwdFilter(Operation *op) {
-+    // Dot operand layout assignment to Predicates are not currently supported
-+    // during lowering from TritonGPU to LLVM in Triton for MMA cases. This
-+    // condition limits visibility of the original bit-width so that predicate
-+    // are not considered, hence, kwidth can never be = 32.
-+    if (isa<arith::UIToFPOp>(op)) {
-+      Type srcType = getElementTypeOrSelf(op->getOperand(0));
-+      if (srcType.isInteger(1))
-+        return false;
-+    }
-     return op->getNumOperands() == 1 &&
-            (isa<tt::FpToFpOp, tt::BitcastOp, ttg::ConvertLayoutOp>(op) ||
-             isPureUnaryInlineAsm(op) ||
diff --git a/third_party/xla/third_party/triton/temporary/series.bzl b/third_party/xla/third_party/triton/temporary/series.bzl
index d68c23d5babbd6..a5dfa7d5eae3f6 100644
--- a/third_party/xla/third_party/triton/temporary/series.bzl
+++ b/third_party/xla/third_party/triton/temporary/series.bzl
@@ -7,6 +7,5 @@ internal patch during the next triton integration process.
 
 temporary_patch_list = [
     "//third_party/triton/temporary:fp8_splat.patch",
-    "//third_party/triton/temporary:mma_limit_pred.patch",
     "//third_party/triton/temporary:fix_register_constraints.patch",
 ]
diff --git a/third_party/xla/third_party/triton/workspace.bzl b/third_party/xla/third_party/triton/workspace.bzl
index 933b815253ffcd..4c21ab7d7c9099 100644
--- a/third_party/xla/third_party/triton/workspace.bzl
+++ b/third_party/xla/third_party/triton/workspace.bzl
@@ -8,8 +8,8 @@ load("//third_party/triton/xla_extensions:series.bzl", "extensions_files_patch_l
 def repo():
     """Imports Triton."""
 
-    TRITON_COMMIT = "cl632952902"
-    TRITON_SHA256 = "f83c4f197cc2ae3b034070ec8189007451093edc445d1e383eb17a1e3808de9a"
+    TRITON_COMMIT = "cl633653394"
+    TRITON_SHA256 = "4bf1b25ce6f7ce6ab6fb4e2403b24c8c0bf71f6a9f5c08f0125fdb44bcb371a8"
     tf_http_archive(
         name = "triton",
         sha256 = TRITON_SHA256,

From b12b5e99f0805167ee07411deabf8f842d098eb1 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Wed, 15 May 2024 23:19:11 -0700
Subject: [PATCH 241/478] Update ops-related pbtxt files.

PiperOrigin-RevId: 634224863
---
 tensorflow/core/ops/ops.pbtxt | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorflow/core/ops/ops.pbtxt b/tensorflow/core/ops/ops.pbtxt
index 846207488b92f3..2073f2e00ecd81 100644
--- a/tensorflow/core/ops/ops.pbtxt
+++ b/tensorflow/core/ops/ops.pbtxt
@@ -1,4 +1,4 @@
-go/debugstr   
+go/debugonly    
 op {
   name: "Abort"
   attr {

From ecaea0e443944d0a84e43ec6a1d53ca0e11f3411 Mon Sep 17 00:00:00 2001
From: Ilia Sergachev <isergachev@nvidia.com>
Date: Thu, 16 May 2024 00:41:51 -0700
Subject: [PATCH 242/478] PR #12538: [GPU] Fix OSS compilation of previously
 disabled tests.

Imported from GitHub PR https://github.com/openxla/xla/pull/12538

Copybara import of the project:

--
78d74154a1693426014eef2940e759d1e36902ed by Ilia Sergachev <isergachev@nvidia.com>:

[GPU] Fix OSS compilation of previously disabled tests.

Merging this change closes #12538

PiperOrigin-RevId: 634248275
---
 third_party/xla/xla/service/gpu/BUILD             |  1 +
 .../xla/xla/service/gpu/ir_emitter_triton_test.cc |  2 +-
 .../xla/xla/service/gpu/triton_support_test.cc    | 15 ++++++++-------
 3 files changed, 10 insertions(+), 8 deletions(-)

diff --git a/third_party/xla/xla/service/gpu/BUILD b/third_party/xla/xla/service/gpu/BUILD
index 9e99084bf14e92..c36ed77518c480 100644
--- a/third_party/xla/xla/service/gpu/BUILD
+++ b/third_party/xla/xla/service/gpu/BUILD
@@ -1220,6 +1220,7 @@ xla_test(
         "@com_google_googletest//:gtest_main",
         "@llvm-project//llvm:ir_headers",
         "@llvm-project//mlir:IR",
+        "@local_tsl//tsl/lib/core:status_test_util",
         "@local_tsl//tsl/platform:status_matchers",
         "@local_tsl//tsl/platform:statusor",
     ],
diff --git a/third_party/xla/xla/service/gpu/ir_emitter_triton_test.cc b/third_party/xla/xla/service/gpu/ir_emitter_triton_test.cc
index 04f0cda2d3e3bd..d98bda3e05fb70 100644
--- a/third_party/xla/xla/service/gpu/ir_emitter_triton_test.cc
+++ b/third_party/xla/xla/service/gpu/ir_emitter_triton_test.cc
@@ -5212,7 +5212,7 @@ ENTRY e {
 })";
 
   TritonGemmConfig config(32, 16, 128, 1, 1, 4);
-  ASSERT_OK(
+  TF_ASSERT_OK(
       CreateTritonIrAndFileCheck(hlo_text, config, EmitMatMul, "triton_dot", R"(
 CHECK:      tt.dot
 CHECK-NOT:  inputPrecision = tf32
diff --git a/third_party/xla/xla/service/gpu/triton_support_test.cc b/third_party/xla/xla/service/gpu/triton_support_test.cc
index e0c7f43831e320..8a8b2242040a97 100644
--- a/third_party/xla/xla/service/gpu/triton_support_test.cc
+++ b/third_party/xla/xla/service/gpu/triton_support_test.cc
@@ -48,6 +48,7 @@ limitations under the License.
 #include "xla/stream_executor/device_description.h"
 #include "xla/xla.pb.h"
 #include "xla/xla_data.pb.h"
+#include "tsl/lib/core/status_test_util.h"
 #include "tsl/platform/status_matchers.h"
 #include "tsl/platform/statusor.h"
 
@@ -165,7 +166,7 @@ ENTRY e {
       hlo_query::GetFirstInstructionWithOpcode(*computation, opcode);
   if (IsTritonSupportedInstruction(*instr, GetCudaComputeCapability())) {
     float tolerance = getTolerance(data_type);
-    EXPECT_OK(ApplyFloatNormalization(module.get()));
+    TF_EXPECT_OK(ApplyFloatNormalization(module.get()));
     EXPECT_TRUE(RunAndCompareNoHloPasses(
         std::move(module), ErrorSpec{/*aabs=*/tolerance, /*arel=*/tolerance}));
   } else {
@@ -246,7 +247,7 @@ ENTRY e {
       hlo_query::GetFirstInstructionWithOpcode(*computation, opcode);
   if (IsTritonSupportedInstruction(*instr, GetCudaComputeCapability())) {
     float tolerance = getTolerance(data_type);
-    EXPECT_OK(ApplyFloatNormalization(module.get()));
+    TF_EXPECT_OK(ApplyFloatNormalization(module.get()));
     EXPECT_TRUE(RunAndCompareNoHloPasses(
         std::move(module), ErrorSpec{/*aabs=*/tolerance, /*arel=*/tolerance}));
   } else {
@@ -329,7 +330,7 @@ ENTRY e {
       hlo_query::GetFirstInstructionWithOpcode(*computation, opcode);
   if (IsTritonSupportedInstruction(*instr, GetCudaComputeCapability())) {
     float tolerance = getTolerance(data_type);
-    EXPECT_OK(ApplyFloatNormalization(module.get()));
+    TF_EXPECT_OK(ApplyFloatNormalization(module.get()));
     EXPECT_TRUE(RunAndCompareNoHloPasses(
         std::move(module), ErrorSpec{/*aabs=*/tolerance, /*arel=*/tolerance}));
   } else {
@@ -393,7 +394,7 @@ ENTRY e {
       hlo_query::GetFirstInstructionWithOpcode(*computation, opcode);
   if (IsTritonSupportedInstruction(*instr, GetCudaComputeCapability())) {
     float tolerance = getTolerance(data_type);
-    EXPECT_OK(ApplyFloatNormalization(module.get()));
+    TF_EXPECT_OK(ApplyFloatNormalization(module.get()));
     EXPECT_TRUE(RunAndCompareNoHloPasses(
         std::move(module), ErrorSpec{/*aabs=*/tolerance, /*arel=*/tolerance}));
   } else {
@@ -449,7 +450,7 @@ ENTRY e {
   const HloInstruction* instr =
       hlo_query::GetFirstInstructionWithOpcode(*computation, opcode);
   if (IsTritonSupportedInstruction(*instr, GetCudaComputeCapability())) {
-    EXPECT_OK(ApplyFloatNormalization(module.get()));
+    TF_EXPECT_OK(ApplyFloatNormalization(module.get()));
     EXPECT_TRUE(RunAndCompareNoHloPasses(
         std::move(module), ErrorSpec{/*aabs=*/2e-4, /*arel=*/2e-4}));
   } else {
@@ -630,7 +631,7 @@ ENTRY main {
       hlo_query::GetFirstInstructionWithOpcode(*computation, opcode);
   if (IsTritonSupportedInstruction(*instr, GetCudaComputeCapability())) {
     float tolerance = getTolerance(data_type);
-    EXPECT_OK(ApplyFloatNormalization(module.get()));
+    TF_EXPECT_OK(ApplyFloatNormalization(module.get()));
     EXPECT_TRUE(RunAndCompareNoHloPasses(
         std::move(module), ErrorSpec{/*aabs=*/tolerance, /*arel=*/tolerance}));
   } else {
@@ -693,7 +694,7 @@ ENTRY main {
       *computation, HloOpcode::kReduce);
   EXPECT_TRUE(IsTritonSupportedInstruction(*instr, GetCudaComputeCapability())
                   .CanFuse());
-  EXPECT_OK(ApplyFloatNormalization(hlo_module.get()));
+  TF_EXPECT_OK(ApplyFloatNormalization(hlo_module.get()));
   EXPECT_TRUE(RunAndCompareNoHloPasses(
       std::move(hlo_module), ErrorSpec{/*aabs=*/2e-4, /*arel=*/2e-4}));
 }

From 6ffda33d9e4561486d8c8da706c7bced81573caa Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Thu, 16 May 2024 01:18:56 -0700
Subject: [PATCH 243/478] Update ops-related pbtxt files.

PiperOrigin-RevId: 634257201
---
 tensorflow/core/ops/ops.pbtxt | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorflow/core/ops/ops.pbtxt b/tensorflow/core/ops/ops.pbtxt
index 2073f2e00ecd81..975f4a22eca161 100644
--- a/tensorflow/core/ops/ops.pbtxt
+++ b/tensorflow/core/ops/ops.pbtxt
@@ -1,4 +1,4 @@
-go/debugonly    
+go/debugstr  
 op {
   name: "Abort"
   attr {

From 6c538a6ecf2668ad81e9874cc14ab8e97038f87e Mon Sep 17 00:00:00 2001
From: George Karpenkov <cheshire@google.com>
Date: Thu, 16 May 2024 01:28:46 -0700
Subject: [PATCH 244/478] [XLA] [NFC] Serialize all autotuning results

The previous logic for filtering by-module wasn't correct, as instructions
could be modified after autotuning, resuling in not all relevant information
serialized.

This could result in excessive data serialized, but autotuning results are very
small compared to the module size and the compiled artifact.

PiperOrigin-RevId: 634259538
---
 .../xla/xla/service/gpu/autotuner_util.cc     | 28 +------------------
 .../xla/xla/service/gpu/autotuner_util.h      | 12 --------
 .../xla/xla/service/gpu/gpu_compiler.cc       | 28 +++++++++----------
 3 files changed, 14 insertions(+), 54 deletions(-)

diff --git a/third_party/xla/xla/service/gpu/autotuner_util.cc b/third_party/xla/xla/service/gpu/autotuner_util.cc
index fe5904ee55a6fe..f3737a60bb996a 100644
--- a/third_party/xla/xla/service/gpu/autotuner_util.cc
+++ b/third_party/xla/xla/service/gpu/autotuner_util.cc
@@ -120,7 +120,7 @@ static void SerializeAutotuneEntry(AutotuneResults* results,
 /*static*/ absl::Status AutotunerUtil::LoadAutotuneResults(
     const AutotuneResults& results) {
   absl::MutexLock lock(&autotune_cache_mu);
-  for (const auto& result : results.results()) {
+  for (const AutotuneResults::Entry& result : results.results()) {
     autotune_cache[AutotuneCacheKey(result.device(), result.hlo())] =
         result.result();
   }
@@ -332,31 +332,5 @@ AutotunerUtil::CreateRedzoneAllocator(const AutotuneConfig& config,
           : 0);
 }
 
-/*static*/ absl::StatusOr<std::string>
-AutotunerUtil::SerializeAutotuneResultsForModule(
-    const HloModule& module, const AutotuneConfig& autotune_config,
-    bool as_textproto) {
-  AutotuneResults results =
-      SerializeAutotuneResultsForModule(module, autotune_config);
-  return AutotuneResultsToString(results, as_textproto);
-}
-
-/*static*/ AutotuneResults AutotunerUtil::SerializeAutotuneResultsForModule(
-    const HloModule& module, const AutotuneConfig& autotune_config) {
-  AutotuneResults results;
-  results.set_version(kVersion);
-
-  for (const HloInstruction* instr :
-       module.entry_computation()->instructions()) {
-    AutotuneCacheKey k(autotune_config.GetModelStr(), *instr);
-    if (const AutotuneResult* res = TryFindInCache(k)) {
-      SerializeAutotuneEntry(&results, k, res);
-    }
-  }
-
-  SortAutotuneResults(&results);
-  return results;
-}
-
 }  // namespace gpu
 }  // namespace xla
diff --git a/third_party/xla/xla/service/gpu/autotuner_util.h b/third_party/xla/xla/service/gpu/autotuner_util.h
index 5570bbe5086435..d9ab1989ddaefc 100644
--- a/third_party/xla/xla/service/gpu/autotuner_util.h
+++ b/third_party/xla/xla/service/gpu/autotuner_util.h
@@ -259,18 +259,6 @@ struct AutotunerUtil {
   static absl::StatusOr<std::string> SerializeAutotuneResults(
       bool as_textproto = false);
 
-  // As above, but only performs serialization for instructions found in the
-  // module.
-  //
-  // Only serializes autotuning results for instructions found in the module:
-  // while this is more expensive than serializing all cache, this avoids
-  // quadratic blow-up when serializing cache for a large number of modules.
-  static absl::StatusOr<std::string> SerializeAutotuneResultsForModule(
-      const HloModule& module, const AutotuneConfig& autotune_config,
-      bool as_textproto);
-  static AutotuneResults SerializeAutotuneResultsForModule(
-      const HloModule& module, const AutotuneConfig& autotune_config);
-
   static absl::Status SerializeAutotuneResults(AutotuneResults* results);
   static absl::Status LoadAutotuneResults(absl::string_view data,
                                           bool as_textproto = false);
diff --git a/third_party/xla/xla/service/gpu/gpu_compiler.cc b/third_party/xla/xla/service/gpu/gpu_compiler.cc
index 769e6de904613a..92c04a51754ce9 100644
--- a/third_party/xla/xla/service/gpu/gpu_compiler.cc
+++ b/third_party/xla/xla/service/gpu/gpu_compiler.cc
@@ -1486,16 +1486,6 @@ absl::Status GpuCompiler::OptimizeHloPostLayoutAssignment(
 #endif  // NDEBUG
 
   TF_RETURN_IF_ERROR(pipeline.Run(hlo_module).status());
-
-  if (DumpingEnabledForHloModule(*hlo_module)) {
-    TF_ASSIGN_OR_RETURN(
-        std::string autotune_results,
-        AutotunerUtil::SerializeAutotuneResultsForModule(
-            *hlo_module, autotune_config, /*as_textproto=*/true));
-    DumpToFileInDirOrStdout(*hlo_module, "", "autotune_results.pbtxt",
-                            autotune_results);
-  }
-
   return absl::OkStatus();
 }
 
@@ -1580,12 +1570,12 @@ absl::StatusOr<std::unique_ptr<HloModule>> GpuCompiler::RunHloPasses(
   DumpHloModuleMetadataIfEnabled({module.get()});
 
   AutotuneResults autotune_results;
+  TF_ASSIGN_OR_RETURN(
+      AutotuneConfig autotune_config,
+      GetAutotuneConfig(stream_exec, debug_opts, options, gpu_target_config));
   if (!is_deviceless) {
-    TF_ASSIGN_OR_RETURN(
-        AutotuneConfig autotune_config,
-        GetAutotuneConfig(stream_exec, debug_opts, options, gpu_target_config));
-    autotune_results = AutotunerUtil::SerializeAutotuneResultsForModule(
-        *module, autotune_config);
+    TF_RETURN_IF_ERROR(
+        AutotunerUtil::SerializeAutotuneResults(&autotune_results));
     TF_RETURN_IF_ERROR(SerializeAutotuneResultsToFile(debug_opts));
   }
   const std::optional<std::string> optimized_fingerprint =
@@ -1596,6 +1586,14 @@ absl::StatusOr<std::unique_ptr<HloModule>> GpuCompiler::RunHloPasses(
                                 *optimized_fingerprint);
   }
 
+  if (DumpingEnabledForHloModule(*module)) {
+    TF_ASSIGN_OR_RETURN(
+        std::string autotune_results,
+        AutotunerUtil::SerializeAutotuneResults(/*as_textproto=*/true));
+    DumpToFileInDirOrStdout(*module, "", "autotune_results.pbtxt",
+                            autotune_results);
+  }
+
   return std::move(module);
 }
 

From 31c77e92e155ec8df24aa44557e0f05d23c1b63c Mon Sep 17 00:00:00 2001
From: Thomas Joerg <tjoerg@google.com>
Date: Thu, 16 May 2024 01:42:41 -0700
Subject: [PATCH 245/478] [XLA:GPU] Remove preprocessor `#if GOOGLE_CUDA`s and
 similar from GemmRewriter pass.

PiperOrigin-RevId: 634263367
---
 .../xla/xla/service/gpu/gemm_rewriter.cc      | 390 +++++++++---------
 1 file changed, 205 insertions(+), 185 deletions(-)

diff --git a/third_party/xla/xla/service/gpu/gemm_rewriter.cc b/third_party/xla/xla/service/gpu/gemm_rewriter.cc
index 2f9c0ac6b35839..dfb72f2dc0cba9 100644
--- a/third_party/xla/xla/service/gpu/gemm_rewriter.cc
+++ b/third_party/xla/xla/service/gpu/gemm_rewriter.cc
@@ -21,6 +21,7 @@ limitations under the License.
 #include <cmath>
 #include <cstddef>
 #include <cstdint>
+#include <initializer_list>
 #include <limits>
 #include <memory>
 #include <optional>
@@ -68,9 +69,6 @@ limitations under the License.
 #include "tsl/platform/statusor.h"
 #include "tsl/protobuf/dnn.pb.h"
 
-#if GOOGLE_CUDA
-#include "third_party/gpus/cuda/include/cuda.h"
-#endif
 
 namespace xla {
 namespace gpu {
@@ -599,13 +597,11 @@ class GemmRewriterVisitor : public DfsHloRewriteVisitor {
                              const_cast<HloInstruction *>(instr), b, b_scale,
                              b_mult_scale, b_ops);
                        })))) {
-#if TENSORFLOW_USE_ROCM
-        if (instr->shape().element_type() != F16 &&
+        if (IsRocm(gpu_version_) && instr->shape().element_type() != F16 &&
             instr->shape().element_type() != F32) {
           TF_ASSIGN_OR_RETURN(instr,
                               TurnF8DotWithUnsupportedOutputTypeIntoF32(instr));
         }
-#endif  // TENSORFLOW_USE_ROCM
         TF_ASSIGN_OR_RETURN(
             bool created_call,
             CreateF8CustomCall(instr, gpu_backend_config, a, b, a_scale,
@@ -926,6 +922,32 @@ class GemmRewriterVisitor : public DfsHloRewriteVisitor {
     return absl::OkStatus();
   }
 
+  static bool IsCuda(const se::GpuComputeCapability &gpu_version) {
+    return std::holds_alternative<se::CudaComputeCapability>(gpu_version);
+  }
+
+  static absl::StatusOr<se::CudaComputeCapability> GetCudaComputeCapability(
+      const se::GpuComputeCapability &gpu_version) {
+    auto *cuda_cc = std::get_if<se::CudaComputeCapability>(&gpu_version);
+    if (cuda_cc == nullptr) {
+      return absl::InvalidArgumentError("Compute Capability is not CUDA.");
+    }
+    return *cuda_cc;
+  }
+
+  static bool IsRocm(const se::GpuComputeCapability &gpu_version) {
+    return std::holds_alternative<se::RocmComputeCapability>(gpu_version);
+  }
+
+  static absl::StatusOr<se::RocmComputeCapability> GetRocmComputeCapability(
+      const se::GpuComputeCapability &gpu_version) {
+    auto rocm_cc = std::get_if<se::RocmComputeCapability>(&gpu_version);
+    if (rocm_cc == nullptr) {
+      return absl::InvalidArgumentError("Compute Capability is not ROCm.");
+    }
+    return *rocm_cc;
+  }
+
   absl::StatusOr<bool> CreateF8CustomCall(
       HloInstruction *instr, GpuBackendConfig &gpu_backend_config,
       HloInstruction *a, HloInstruction *b, HloInstruction *a_scale,
@@ -934,9 +956,9 @@ class GemmRewriterVisitor : public DfsHloRewriteVisitor {
       std::vector<std::pair<HloInstruction *, int>> b_ops) {
     GemmBackendConfig &gemm_backend_config =
         *gpu_backend_config.mutable_gemm_backend_config();
-    if (std::holds_alternative<se::CudaComputeCapability>(gpu_version_)) {
-      auto cuda_compute_capability =
-          std::get<se::CudaComputeCapability>(gpu_version_);
+    if (IsCuda(gpu_version_)) {
+      TF_ASSIGN_OR_RETURN(auto cuda_compute_capability,
+                          GetCudaComputeCapability(gpu_version_));
       // FP8 GEMM kernels are only available on Ada, Hopper, and later
       // architectures.
       if (!cuda_compute_capability.IsAtLeast(8, 9)) {
@@ -953,65 +975,63 @@ class GemmRewriterVisitor : public DfsHloRewriteVisitor {
       }
     }
 
-#if TENSORFLOW_USE_ROCM
-    auto isrocm = std::get_if<se::RocmComputeCapability>(&gpu_version_);
-    if (!isrocm->has_fp8_support()) {
-      VLOG(1) << "FP8 Custom Calls require MI300, or later architectures.";
-      return false;
+    if (IsRocm(gpu_version_)) {
+      TF_ASSIGN_OR_RETURN(auto rocm_compute_capability,
+                          GetRocmComputeCapability(gpu_version_));
+      if (!rocm_compute_capability.has_fp8_support()) {
+        VLOG(1) << "FP8 Custom Calls require MI300, or later architectures.";
+        return false;
+      }
+      if (toolkit_version_ < 60000) {
+        // FP8 GEMM kernels are only available with ROCm 6.0 and above
+        VLOG(1) << "FP8 Custom Calls require ROCm 6.0 or newer.";
+        return false;
+      }
     }
 
-#if TF_ROCM_VERSION < 60000
-    // FP8 GEMM kernels are only available with ROCm 6.0 and above
-    VLOG(1) << "FP8 Custom Calls require ROCm 6.0 or newer.";
-    return false;
-#endif  // TF_ROCM_VERSION < 60000
-
-#endif  // TENSORFLOW_USE_ROCM
-
-#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
     PrimitiveType a_type = a->shape().element_type();
     PrimitiveType b_type = b->shape().element_type();
-#endif  // GOOGLE_CUDA || TENSORFLOW_USE_ROCM
 
     // cuBLASLt FP8 GEMM kernels require one of the two operands to be in
     // F8E4M3FN format.
-#if GOOGLE_CUDA
-    if (a_type == F8E5M2 && b_type == F8E5M2) {
-      VLOG(1)
-          << "Failed to rewrite " << instr->ToShortString()
-          << " into FP8 Custom Call. The element type of one of the operands "
-             "must be F8E4M3FN.";
-      return false;
-    }
-    if ((a_type != F8E5M2 && a_type != F8E4M3FN) ||
-        (b_type != F8E5M2 && b_type != F8E4M3FN)) {
-      VLOG(1) << "Failed to rewrite " << instr->ToShortString()
-              << " into FP8 Custom Call. The input types must be F8E5M2 or "
-                 "F8E4M3FN, but got "
-              << PrimitiveType_Name(a_type) << " and "
-              << PrimitiveType_Name(b_type);
-      return false;
+    if (IsCuda(gpu_version_)) {
+      if (a_type == F8E5M2 && b_type == F8E5M2) {
+        VLOG(1)
+            << "Failed to rewrite " << instr->ToShortString()
+            << " into FP8 Custom Call. The element type of one of the operands "
+               "must be F8E4M3FN.";
+        return false;
+      }
+      if ((a_type != F8E5M2 && a_type != F8E4M3FN) ||
+          (b_type != F8E5M2 && b_type != F8E4M3FN)) {
+        VLOG(1) << "Failed to rewrite " << instr->ToShortString()
+                << " into FP8 Custom Call. The input types must be F8E5M2 or "
+                   "F8E4M3FN, but got "
+                << PrimitiveType_Name(a_type) << " and "
+                << PrimitiveType_Name(b_type);
+        return false;
+      }
     }
-#endif  // GOOGLE_CUDA
 
-#if TENSORFLOW_USE_ROCM
-    if (a_type == F8E5M2FNUZ && b_type == F8E5M2FNUZ) {
-      VLOG(1)
-          << "Failed to rewrite " << instr->ToShortString()
-          << " into FP8 Custom Call. The element type of one of the operands "
-             "must be F8E4M3FNUZ.";
-      return false;
-    }
-    if ((a_type != F8E5M2FNUZ && a_type != F8E4M3FNUZ) ||
-        (b_type != F8E5M2FNUZ && b_type != F8E4M3FNUZ)) {
-      VLOG(1) << "Failed to rewrite " << instr->ToShortString()
-              << " into FP8 Custom Call. The input types must be F8E5M2FNUZ or "
-                 "F8E4M3FNUZ, but got "
-              << PrimitiveType_Name(a_type) << " and "
-              << PrimitiveType_Name(b_type);
-      return false;
+    if (IsRocm(gpu_version_)) {
+      if (a_type == F8E5M2FNUZ && b_type == F8E5M2FNUZ) {
+        VLOG(1)
+            << "Failed to rewrite " << instr->ToShortString()
+            << " into FP8 Custom Call. The element type of one of the operands "
+               "must be F8E4M3FNUZ.";
+        return false;
+      }
+      if ((a_type != F8E5M2FNUZ && a_type != F8E4M3FNUZ) ||
+          (b_type != F8E5M2FNUZ && b_type != F8E4M3FNUZ)) {
+        VLOG(1)
+            << "Failed to rewrite " << instr->ToShortString()
+            << " into FP8 Custom Call. The input types must be F8E5M2FNUZ or "
+               "F8E4M3FNUZ, but got "
+            << PrimitiveType_Name(a_type) << " and "
+            << PrimitiveType_Name(b_type);
+        return false;
+      }
     }
-#endif  // TENSORFLOW_USE_ROCM
 
     absl::Span<const int64_t> batch_dims =
         gemm_backend_config.dot_dimension_numbers().rhs_batch_dimensions();
@@ -1696,16 +1716,13 @@ class GemmRewriterVisitor : public DfsHloRewriteVisitor {
     if (!SupportsEpilogueFusion(gemm->shape().element_type())) {
       return absl::OkStatus();
     }
-
-#if CUDA_VERSION < 12040
     // For CUDA versions less than 12.3.2, cuBLAS LT returns
     // CUBLAS_STATUS_NOT_SUPPORTED in some cases when fusing gelu into an FP8
     // matmul. We cannot check the patch version, so disable this fusion with
     // CUDA versions less than 12.4.
-    if (IsCublasLtMatmulF8(*gemm)) {
+    if (toolkit_version_ < 12040 && IsCublasLtMatmulF8(*gemm)) {
       return absl::OkStatus();
     }
-#endif
 
     // There are four users of the gemm output within the GELU calculation.
     bool has_aux = gemm->user_count() > 4;
@@ -1911,131 +1928,136 @@ class GemmRewriterVisitor : public DfsHloRewriteVisitor {
 
     using se::blas::ComputationType;
     using se::blas::DataType;
+    using TypeCombinations = std::initializer_list<std::tuple<
+        ComputationType, DataType /*scale_type*/, PrimitiveType /*a_dtype*/,
+        PrimitiveType /*b_dtype*/, DataType /*output_dtype*/>>;
     // This matrix of supported types is taken directly from cublasLt
     // documentation.
     // https://docs.nvidia.com/cuda/cublas/index.html#cublasLtMatmul
-    const std::array<
-        std::tuple<ComputationType, DataType /*scale_type*/,
-                   PrimitiveType /*a_dtype*/, PrimitiveType /*b_dtype*/,
-                   DataType /*output_dtype*/>,
-        32>
-        supported_type_combinations = {{
-#if GOOGLE_CUDA
-            // FP8 types:
-            {ComputationType::kF32, DataType::kFloat, PrimitiveType::F8E4M3FN,
-             PrimitiveType::F8E4M3FN, DataType::kBF16},
-            {ComputationType::kF32, DataType::kFloat, PrimitiveType::F8E4M3FN,
-             PrimitiveType::F8E4M3FN, DataType::kF8E4M3FN},
-            {ComputationType::kF32, DataType::kFloat, PrimitiveType::F8E4M3FN,
-             PrimitiveType::F8E4M3FN, DataType::kHalf},
-            {ComputationType::kF32, DataType::kFloat, PrimitiveType::F8E4M3FN,
-             PrimitiveType::F8E4M3FN, DataType::kFloat},
-
-            {ComputationType::kF32, DataType::kFloat, PrimitiveType::F8E4M3FN,
-             PrimitiveType::F8E5M2, DataType::kBF16},
-            {ComputationType::kF32, DataType::kFloat, PrimitiveType::F8E4M3FN,
-             PrimitiveType::F8E5M2, DataType::kF8E4M3FN},
-            {ComputationType::kF32, DataType::kFloat, PrimitiveType::F8E4M3FN,
-             PrimitiveType::F8E5M2, DataType::kF8E5M2},
-            {ComputationType::kF32, DataType::kFloat, PrimitiveType::F8E4M3FN,
-             PrimitiveType::F8E5M2, DataType::kHalf},
-            {ComputationType::kF32, DataType::kFloat, PrimitiveType::F8E4M3FN,
-             PrimitiveType::F8E5M2, DataType::kFloat},
-
-            {ComputationType::kF32, DataType::kFloat, PrimitiveType::F8E5M2,
-             PrimitiveType::F8E4M3FN, DataType::kBF16},
-            {ComputationType::kF32, DataType::kFloat, PrimitiveType::F8E5M2,
-             PrimitiveType::F8E4M3FN, DataType::kF8E4M3FN},
-            {ComputationType::kF32, DataType::kFloat, PrimitiveType::F8E5M2,
-             PrimitiveType::F8E4M3FN, DataType::kF8E5M2},
-            {ComputationType::kF32, DataType::kFloat, PrimitiveType::F8E5M2,
-             PrimitiveType::F8E4M3FN, DataType::kHalf},
-            {ComputationType::kF32, DataType::kFloat, PrimitiveType::F8E5M2,
-             PrimitiveType::F8E4M3FN, DataType::kFloat},
-#endif  // GOOGLE_CUDA
-#if TENSORFLOW_USE_ROCM
-            // FP8 types:
-            {ComputationType::kF32, DataType::kFloat, PrimitiveType::F8E4M3FNUZ,
-             PrimitiveType::F8E4M3FNUZ, DataType::kBF16},
-            {ComputationType::kF32, DataType::kFloat, PrimitiveType::F8E4M3FNUZ,
-             PrimitiveType::F8E4M3FNUZ, DataType::kF8E4M3FNUZ},
-            {ComputationType::kF32, DataType::kFloat, PrimitiveType::F8E4M3FNUZ,
-             PrimitiveType::F8E4M3FNUZ, DataType::kHalf},
-            {ComputationType::kF32, DataType::kFloat, PrimitiveType::F8E4M3FNUZ,
-             PrimitiveType::F8E4M3FNUZ, DataType::kFloat},
-
-            {ComputationType::kF32, DataType::kFloat, PrimitiveType::F8E4M3FNUZ,
-             PrimitiveType::F8E5M2FNUZ, DataType::kBF16},
-            {ComputationType::kF32, DataType::kFloat, PrimitiveType::F8E4M3FNUZ,
-             PrimitiveType::F8E5M2FNUZ, DataType::kF8E4M3FNUZ},
-            {ComputationType::kF32, DataType::kFloat, PrimitiveType::F8E4M3FNUZ,
-             PrimitiveType::F8E5M2FNUZ, DataType::kF8E5M2FNUZ},
-            {ComputationType::kF32, DataType::kFloat, PrimitiveType::F8E4M3FNUZ,
-             PrimitiveType::F8E5M2FNUZ, DataType::kHalf},
-            {ComputationType::kF32, DataType::kFloat, PrimitiveType::F8E4M3FNUZ,
-             PrimitiveType::F8E5M2FNUZ, DataType::kFloat},
-
-            {ComputationType::kF32, DataType::kFloat, PrimitiveType::F8E5M2FNUZ,
-             PrimitiveType::F8E4M3FNUZ, DataType::kBF16},
-            {ComputationType::kF32, DataType::kFloat, PrimitiveType::F8E5M2FNUZ,
-             PrimitiveType::F8E4M3FNUZ, DataType::kF8E4M3FNUZ},
-            {ComputationType::kF32, DataType::kFloat, PrimitiveType::F8E5M2FNUZ,
-             PrimitiveType::F8E4M3FNUZ, DataType::kF8E5M2FNUZ},
-            {ComputationType::kF32, DataType::kFloat, PrimitiveType::F8E5M2FNUZ,
-             PrimitiveType::F8E4M3FNUZ, DataType::kHalf},
-            {ComputationType::kF32, DataType::kFloat, PrimitiveType::F8E5M2FNUZ,
-             PrimitiveType::F8E4M3FNUZ, DataType::kFloat},
-#endif  // TENSORFLOW_USE_ROCM
+    const TypeCombinations supported_cublas_type_combinations = {
+        // FP8 types:
+        {ComputationType::kF32, DataType::kFloat, PrimitiveType::F8E4M3FN,
+         PrimitiveType::F8E4M3FN, DataType::kBF16},
+        {ComputationType::kF32, DataType::kFloat, PrimitiveType::F8E4M3FN,
+         PrimitiveType::F8E4M3FN, DataType::kF8E4M3FN},
+        {ComputationType::kF32, DataType::kFloat, PrimitiveType::F8E4M3FN,
+         PrimitiveType::F8E4M3FN, DataType::kHalf},
+        {ComputationType::kF32, DataType::kFloat, PrimitiveType::F8E4M3FN,
+         PrimitiveType::F8E4M3FN, DataType::kFloat},
+
+        {ComputationType::kF32, DataType::kFloat, PrimitiveType::F8E4M3FN,
+         PrimitiveType::F8E5M2, DataType::kBF16},
+        {ComputationType::kF32, DataType::kFloat, PrimitiveType::F8E4M3FN,
+         PrimitiveType::F8E5M2, DataType::kF8E4M3FN},
+        {ComputationType::kF32, DataType::kFloat, PrimitiveType::F8E4M3FN,
+         PrimitiveType::F8E5M2, DataType::kF8E5M2},
+        {ComputationType::kF32, DataType::kFloat, PrimitiveType::F8E4M3FN,
+         PrimitiveType::F8E5M2, DataType::kHalf},
+        {ComputationType::kF32, DataType::kFloat, PrimitiveType::F8E4M3FN,
+         PrimitiveType::F8E5M2, DataType::kFloat},
+
+        {ComputationType::kF32, DataType::kFloat, PrimitiveType::F8E5M2,
+         PrimitiveType::F8E4M3FN, DataType::kBF16},
+        {ComputationType::kF32, DataType::kFloat, PrimitiveType::F8E5M2,
+         PrimitiveType::F8E4M3FN, DataType::kF8E4M3FN},
+        {ComputationType::kF32, DataType::kFloat, PrimitiveType::F8E5M2,
+         PrimitiveType::F8E4M3FN, DataType::kF8E5M2},
+        {ComputationType::kF32, DataType::kFloat, PrimitiveType::F8E5M2,
+         PrimitiveType::F8E4M3FN, DataType::kHalf},
+        {ComputationType::kF32, DataType::kFloat, PrimitiveType::F8E5M2,
+         PrimitiveType::F8E4M3FN, DataType::kFloat},
+        // There would be an entry here for A/BType complex int8, but we do
+        // not support that type.
+        {ComputationType::kF32, DataType::kComplexFloat, PrimitiveType::C64,
+         PrimitiveType::C64, DataType::kComplexFloat},
+
+        {ComputationType::kF16AsF32, DataType::kFloat, PrimitiveType::F32,
+         PrimitiveType::F32, DataType::kFloat},
+        {ComputationType::kF16AsF32, DataType::kComplexFloat,
+         PrimitiveType::C64, PrimitiveType::C64, DataType::kComplexFloat},
+        // The next 4 may be supported by hipblaslt, but they are not
+        // covered by any unit tests
+        {ComputationType::kBF16AsF32, DataType::kFloat, PrimitiveType::F32,
+         PrimitiveType::F32, DataType::kFloat},
+        {ComputationType::kBF16AsF32, DataType::kComplexFloat,
+         PrimitiveType::C64, PrimitiveType::C64, DataType::kComplexFloat},
+
+        {ComputationType::kTF32AsF32, DataType::kFloat, PrimitiveType::F32,
+         PrimitiveType::F32, DataType::kFloat},
+        {ComputationType::kTF32AsF32, DataType::kComplexFloat,
+         PrimitiveType::C64, PrimitiveType::C64, DataType::kComplexFloat},
+
+        {ComputationType::kF64, DataType::kDouble, PrimitiveType::F64,
+         PrimitiveType::F64, DataType::kDouble},
+        {ComputationType::kF64, DataType::kComplexDouble, PrimitiveType::C128,
+         PrimitiveType::C128, DataType::kComplexDouble},
+    };
+    if (absl::c_linear_search(supported_cublas_type_combinations,
+                              std::make_tuple(compute_type, scale_type, a_dtype,
+                                              b_dtype, output_dtype))) {
+      return true;
+    }
+    const TypeCombinations supported_hipblas_type_combinations = {
+        // FP8 types:
+        {ComputationType::kF32, DataType::kFloat, PrimitiveType::F8E4M3FNUZ,
+         PrimitiveType::F8E4M3FNUZ, DataType::kBF16},
+        {ComputationType::kF32, DataType::kFloat, PrimitiveType::F8E4M3FNUZ,
+         PrimitiveType::F8E4M3FNUZ, DataType::kF8E4M3FNUZ},
+        {ComputationType::kF32, DataType::kFloat, PrimitiveType::F8E4M3FNUZ,
+         PrimitiveType::F8E4M3FNUZ, DataType::kHalf},
+        {ComputationType::kF32, DataType::kFloat, PrimitiveType::F8E4M3FNUZ,
+         PrimitiveType::F8E4M3FNUZ, DataType::kFloat},
+
+        {ComputationType::kF32, DataType::kFloat, PrimitiveType::F8E4M3FNUZ,
+         PrimitiveType::F8E5M2FNUZ, DataType::kBF16},
+        {ComputationType::kF32, DataType::kFloat, PrimitiveType::F8E4M3FNUZ,
+         PrimitiveType::F8E5M2FNUZ, DataType::kF8E4M3FNUZ},
+        {ComputationType::kF32, DataType::kFloat, PrimitiveType::F8E4M3FNUZ,
+         PrimitiveType::F8E5M2FNUZ, DataType::kF8E5M2FNUZ},
+        {ComputationType::kF32, DataType::kFloat, PrimitiveType::F8E4M3FNUZ,
+         PrimitiveType::F8E5M2FNUZ, DataType::kHalf},
+        {ComputationType::kF32, DataType::kFloat, PrimitiveType::F8E4M3FNUZ,
+         PrimitiveType::F8E5M2FNUZ, DataType::kFloat},
+
+        {ComputationType::kF32, DataType::kFloat, PrimitiveType::F8E5M2FNUZ,
+         PrimitiveType::F8E4M3FNUZ, DataType::kBF16},
+        {ComputationType::kF32, DataType::kFloat, PrimitiveType::F8E5M2FNUZ,
+         PrimitiveType::F8E4M3FNUZ, DataType::kF8E4M3FNUZ},
+        {ComputationType::kF32, DataType::kFloat, PrimitiveType::F8E5M2FNUZ,
+         PrimitiveType::F8E4M3FNUZ, DataType::kF8E5M2FNUZ},
+        {ComputationType::kF32, DataType::kFloat, PrimitiveType::F8E5M2FNUZ,
+         PrimitiveType::F8E4M3FNUZ, DataType::kHalf},
+        {ComputationType::kF32, DataType::kFloat, PrimitiveType::F8E5M2FNUZ,
+         PrimitiveType::F8E4M3FNUZ, DataType::kFloat},
+    };
+    if (absl::c_linear_search(supported_hipblas_type_combinations,
+                              std::make_tuple(compute_type, scale_type, a_dtype,
+                                              b_dtype, output_dtype))) {
+      return true;
+    }
+    const TypeCombinations supported_type_combinations = {
         // Other data types:
-            {ComputationType::kF16, DataType::kHalf, PrimitiveType::F16,
-             PrimitiveType::F16, DataType::kHalf},
-
-            {ComputationType::kI32, DataType::kInt32, PrimitiveType::S8,
-             PrimitiveType::S8, DataType::kInt32},
-            {ComputationType::kI32, DataType::kFloat, PrimitiveType::S8,
-             PrimitiveType::S8, DataType::kInt8},
-
-            {ComputationType::kF32, DataType::kFloat, PrimitiveType::BF16,
-             PrimitiveType::BF16, DataType::kBF16},
-            {ComputationType::kF32, DataType::kFloat, PrimitiveType::F16,
-             PrimitiveType::F16, DataType::kHalf},
-            {ComputationType::kF32, DataType::kFloat, PrimitiveType::S8,
-             PrimitiveType::S8, DataType::kFloat},
-            {ComputationType::kF32, DataType::kFloat, PrimitiveType::BF16,
-             PrimitiveType::BF16, DataType::kFloat},
-            {ComputationType::kF32, DataType::kFloat, PrimitiveType::F16,
-             PrimitiveType::F16, DataType::kFloat},
-            {ComputationType::kF32, DataType::kFloat, PrimitiveType::F32,
-             PrimitiveType::F32, DataType::kFloat},
-#if GOOGLE_CUDA
-            // There would be an entry here for A/BType complex int8, but we do
-            // not support that type.
-            {ComputationType::kF32, DataType::kComplexFloat, PrimitiveType::C64,
-             PrimitiveType::C64, DataType::kComplexFloat},
-
-            {ComputationType::kF16AsF32, DataType::kFloat, PrimitiveType::F32,
-             PrimitiveType::F32, DataType::kFloat},
-            {ComputationType::kF16AsF32, DataType::kComplexFloat,
-             PrimitiveType::C64, PrimitiveType::C64, DataType::kComplexFloat},
-            // The next 4 may be supported by hipblaslt, but they are not
-            // covered by any unit tests
-            {ComputationType::kBF16AsF32, DataType::kFloat, PrimitiveType::F32,
-             PrimitiveType::F32, DataType::kFloat},
-            {ComputationType::kBF16AsF32, DataType::kComplexFloat,
-             PrimitiveType::C64, PrimitiveType::C64, DataType::kComplexFloat},
-
-            {ComputationType::kTF32AsF32, DataType::kFloat, PrimitiveType::F32,
-             PrimitiveType::F32, DataType::kFloat},
-            {ComputationType::kTF32AsF32, DataType::kComplexFloat,
-             PrimitiveType::C64, PrimitiveType::C64, DataType::kComplexFloat},
-
-            {ComputationType::kF64, DataType::kDouble, PrimitiveType::F64,
-             PrimitiveType::F64, DataType::kDouble},
-            {ComputationType::kF64, DataType::kComplexDouble,
-             PrimitiveType::C128, PrimitiveType::C128,
-             DataType::kComplexDouble},
-#endif  // GOOGLE_CUDA
-        }};
+        {ComputationType::kF16, DataType::kHalf, PrimitiveType::F16,
+         PrimitiveType::F16, DataType::kHalf},
+
+        {ComputationType::kI32, DataType::kInt32, PrimitiveType::S8,
+         PrimitiveType::S8, DataType::kInt32},
+        {ComputationType::kI32, DataType::kFloat, PrimitiveType::S8,
+         PrimitiveType::S8, DataType::kInt8},
+
+        {ComputationType::kF32, DataType::kFloat, PrimitiveType::BF16,
+         PrimitiveType::BF16, DataType::kBF16},
+        {ComputationType::kF32, DataType::kFloat, PrimitiveType::F16,
+         PrimitiveType::F16, DataType::kHalf},
+        {ComputationType::kF32, DataType::kFloat, PrimitiveType::S8,
+         PrimitiveType::S8, DataType::kFloat},
+        {ComputationType::kF32, DataType::kFloat, PrimitiveType::BF16,
+         PrimitiveType::BF16, DataType::kFloat},
+        {ComputationType::kF32, DataType::kFloat, PrimitiveType::F16,
+         PrimitiveType::F16, DataType::kFloat},
+        {ComputationType::kF32, DataType::kFloat, PrimitiveType::F32,
+         PrimitiveType::F32, DataType::kFloat},
+    };
 
     return absl::c_linear_search(
         supported_type_combinations,
@@ -2172,7 +2194,6 @@ class GemmRewriterVisitor : public DfsHloRewriteVisitor {
     return lhs_non_contracting_dimension_size <= kMaxDimensionSize;
   }
 
-#if TENSORFLOW_USE_ROCM
   // Turns an F8 dot with unsupported output type into an F8 dot with F32
   // output, and converting the F32 output to unsupported output types.
   absl::StatusOr<HloInstruction *> TurnF8DotWithUnsupportedOutputTypeIntoF32(
@@ -2186,7 +2207,6 @@ class GemmRewriterVisitor : public DfsHloRewriteVisitor {
     TF_RETURN_IF_ERROR(ReplaceInstruction(instr, convert));
     return f32_dot;
   }
-#endif  // TENSORFLOW_USE_ROCM
 
   // Turns an F8 dot into an F16 dot, converting operands to F16 and
   // converting the output back to F8.

From 1f944897c735decf807217139a643eaf75d8e352 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Tam=C3=A1s=20Danyluk?= <tdanyluk@google.com>
Date: Thu, 16 May 2024 01:47:06 -0700
Subject: [PATCH 246/478] [XLA:GPU] Use IndexingMap::RemoveUnusedSymbols
 instead of WithoutRangeVars

PiperOrigin-RevId: 634264428
---
 .../xla/service/gpu/model/symbolic_tile.cc    | 53 +++----------------
 1 file changed, 8 insertions(+), 45 deletions(-)

diff --git a/third_party/xla/xla/service/gpu/model/symbolic_tile.cc b/third_party/xla/xla/service/gpu/model/symbolic_tile.cc
index ef0216d77fd778..bedf16c1389d3d 100644
--- a/third_party/xla/xla/service/gpu/model/symbolic_tile.cc
+++ b/third_party/xla/xla/service/gpu/model/symbolic_tile.cc
@@ -49,43 +49,6 @@ using ::mlir::getAffineConstantExpr;
 using ::mlir::getAffineDimExpr;
 using ::mlir::MLIRContext;
 
-// Gets a modified version of `expressions` where the indices of the RTVars are
-// decreased, because no RangeVars appear anymore.
-//
-// (dimensions)[RangeVars, RTVars] -> (dimensions)[RTVars]
-//
-// Precondition: `expressions` must not contain any RangeVar symbols.
-//
-// This handles a vector of expressions at once, because we don't want to
-// regenerate `symbol_map` every time.
-std::vector<AffineExpr> WithoutRangeVars(std::vector<AffineExpr> expressions,
-                                         const IndexingMap& indexing_map) {
-  // Precondition check:
-  for (AffineExpr expression : expressions) {
-    expression.walk([&indexing_map](AffineExpr expr) {
-      CHECK(!(expr.getKind() == AffineExprKind::SymbolId &&
-              indexing_map.IsRangeVarSymbol(
-                  llvm::cast<mlir::AffineSymbolExpr>(expr))));
-    });
-  }
-
-  MLIRContext* mlir_context = indexing_map.GetMLIRContext();
-
-  // Cannot use AffineExpr::shiftSymbols, because it doesn't support negative
-  // shifts.
-  llvm::DenseMap<AffineExpr, AffineExpr> symbol_map;
-  for (int i = 0; i < indexing_map.GetRTVarsCount(); i++) {
-    symbol_map[getAffineSymbolExpr(indexing_map.GetRangeVarsCount() + i,
-                                   mlir_context)] =
-        getAffineSymbolExpr(i, mlir_context);
-  }
-  for (AffineExpr& expression : expressions) {
-    expression = expression.replace(symbol_map);
-  }
-
-  return expressions;
-}
-
 // Gets a modified version of `expressions` where both the original dimensions
 // and symbols are replaced with symbols.
 //
@@ -418,16 +381,13 @@ std::optional<SizeAndStrideExpression> ExtractSizeAndStride(
   }
 
   std::vector<AffineExpr> results;
-  absl::c_move(WithoutRangeVars(std::move(offset_expressions), indexing_map),
-               std::back_inserter(results));
-  absl::c_move(WithoutRangeVars(std::move(size_expressions), indexing_map),
-               std::back_inserter(results));
-  absl::c_move(WithoutRangeVars(std::move(stride_expressions), indexing_map),
-               std::back_inserter(results));
+  absl::c_move(std::move(offset_expressions), std::back_inserter(results));
+  absl::c_move(std::move(size_expressions), std::back_inserter(results));
+  absl::c_move(std::move(stride_expressions), std::back_inserter(results));
 
   AffineMap tile_affine_map =
       AffineMap::get(/*dimCount=*/tile_sizes.size(),
-                     /*symbolCount=*/indexing_map.GetRTVarsCount(),
+                     /*symbolCount=*/indexing_map.GetSymbolCount(),
                      /*results=*/results,
                      /*context=*/indexing_map.GetMLIRContext());
 
@@ -437,8 +397,11 @@ std::optional<SizeAndStrideExpression> ExtractSizeAndStride(
   IndexingMap tile_map(
       /*affine_map=*/std::move(tile_affine_map),
       /*dimensions=*/std::move(tile_sizes),
-      /*range_vars=*/{},
+      /*range_vars=*/indexing_map.GetRangeVars(),
       /*rt_vars=*/indexing_map.GetRTVars());
+  tile_map.RemoveUnusedSymbols();
+  CHECK_EQ(tile_map.GetRangeVarsCount(), 0);
+
   VLOG(1) << "tile_map: " << tile_map.ToString();
   return SymbolicTile(std::move(tile_map));
 }

From 6c2e476491ecf24ab69abf88edf2f97dbbfd0ca8 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Thu, 16 May 2024 02:02:12 -0700
Subject: [PATCH 247/478] Update GraphDef version to 1864.

PiperOrigin-RevId: 634268437
---
 tensorflow/core/public/version.h | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorflow/core/public/version.h b/tensorflow/core/public/version.h
index a60f2dd95fe679..ea887161e0298c 100644
--- a/tensorflow/core/public/version.h
+++ b/tensorflow/core/public/version.h
@@ -108,7 +108,7 @@ limitations under the License.
 
 #define TF_GRAPH_DEF_VERSION_MIN_PRODUCER 0
 #define TF_GRAPH_DEF_VERSION_MIN_CONSUMER 0
-#define TF_GRAPH_DEF_VERSION 1863  // Updated: 2024/5/15
+#define TF_GRAPH_DEF_VERSION 1864  // Updated: 2024/5/16
 
 // Checkpoint compatibility versions (the versions field in SavedSliceMeta).
 //

From 175131a63a8b55dc14179e6ed80b56a94751a1f2 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Thu, 16 May 2024 02:02:17 -0700
Subject: [PATCH 248/478] compat: Update forward compatibility horizon to
 2024-05-16

PiperOrigin-RevId: 634268459
---
 tensorflow/python/compat/compat.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorflow/python/compat/compat.py b/tensorflow/python/compat/compat.py
index 8416badea05697..a236aaf853f105 100644
--- a/tensorflow/python/compat/compat.py
+++ b/tensorflow/python/compat/compat.py
@@ -29,7 +29,7 @@
 # This value changes every day with an automatic CL. It can be modified in code
 # via `forward_compatibility_horizon()` or with the environment variable
 # TF_FORWARD_COMPATIBILITY_DELTA_DAYS, which is added to the compatibility date.
-_FORWARD_COMPATIBILITY_HORIZON = datetime.date(2024, 5, 15)
+_FORWARD_COMPATIBILITY_HORIZON = datetime.date(2024, 5, 16)
 _FORWARD_COMPATIBILITY_DELTA_DAYS_VAR_NAME = "TF_FORWARD_COMPATIBILITY_DELTA_DAYS"
 _FORWARD_COMPATIBILITY_DATE_NUMBER = None
 

From 5720ac7dccdaa81fc5c56257698766cedba6e688 Mon Sep 17 00:00:00 2001
From: Aliia Khasanova <aliia@google.com>
Date: Thu, 16 May 2024 02:41:44 -0700
Subject: [PATCH 249/478] [XLA:GPU] Remove GpuStatus.

XLA should not include any platform specific headers in XLA header files. This change moves us one step closer to this goal by getting rid of type aliases in `gpu_types.h`.

PiperOrigin-RevId: 634278649
---
 .../xla/xla/stream_executor/cuda/BUILD        |  2 +-
 .../xla/stream_executor/cuda/cuda_driver.cc   | 23 +++++++++---------
 .../xla/stream_executor/cuda/cuda_driver.h    |  5 ++++
 .../xla/stream_executor/cuda/cuda_event.cc    |  4 ++--
 .../xla/xla/stream_executor/gpu/gpu_driver.h  |  6 -----
 .../xla/xla/stream_executor/gpu/gpu_types.h   |  2 --
 .../xla/stream_executor/rocm/rocm_driver.cc   | 24 +++++++++----------
 .../xla/stream_executor/rocm/rocm_driver.h    |  2 ++
 .../xla/stream_executor/rocm/rocm_event.cc    |  3 ++-
 9 files changed, 34 insertions(+), 37 deletions(-)

diff --git a/third_party/xla/xla/stream_executor/cuda/BUILD b/third_party/xla/xla/stream_executor/cuda/BUILD
index 789ce0c444327b..a756c861863cdf 100644
--- a/third_party/xla/xla/stream_executor/cuda/BUILD
+++ b/third_party/xla/xla/stream_executor/cuda/BUILD
@@ -458,8 +458,8 @@ cuda_only_cc_library(
     srcs = ["cuda_event.cc"],
     hdrs = ["cuda_event.h"],
     deps = [
+        ":cuda_driver",
         "//xla/stream_executor",
-        "//xla/stream_executor/gpu:gpu_driver_header",
         "//xla/stream_executor/gpu:gpu_event",
         "//xla/stream_executor/gpu:gpu_executor_header",
         "@com_google_absl//absl/log",
diff --git a/third_party/xla/xla/stream_executor/cuda/cuda_driver.cc b/third_party/xla/xla/stream_executor/cuda/cuda_driver.cc
index f665f39155041a..ac6a7fb4e0e389 100644
--- a/third_party/xla/xla/stream_executor/cuda/cuda_driver.cc
+++ b/third_party/xla/xla/stream_executor/cuda/cuda_driver.cc
@@ -306,6 +306,17 @@ int GetFlagsFromEnv() {
 
 }  // namespace
 
+absl::StatusOr<CUresult> QueryEvent(GpuContext* context, CUevent event) {
+  ScopedActivateContext activated{context};
+  CUresult res = cuEventQuery(event);
+  if (res != CUDA_SUCCESS && res != CUDA_ERROR_NOT_READY) {
+    return absl::InternalError(
+        absl::StrFormat("failed to query event: %s", ToString(res)));
+  }
+
+  return res;
+}
+
 /* static */ absl::Status GpuDriver::Init() {
   // Cached return value from calling InternalInit(), as cuInit need only be
   // called once, but GpuDriver::Init may be called many times.
@@ -1840,18 +1851,6 @@ GpuDriver::CreateMemoryHandle(GpuContext* context, uint64_t bytes) {
   return absl::OkStatus();
 }
 
-/* static */ absl::StatusOr<CUresult> GpuDriver::QueryEvent(GpuContext* context,
-                                                            CUevent event) {
-  ScopedActivateContext activated{context};
-  CUresult res = cuEventQuery(event);
-  if (res != CUDA_SUCCESS && res != CUDA_ERROR_NOT_READY) {
-    return absl::InternalError(
-        absl::StrFormat("failed to query event: %s", ToString(res)));
-  }
-
-  return res;
-}
-
 /* static */ bool GpuDriver::GetEventElapsedTime(GpuContext* context,
                                                  float* elapsed_milliseconds,
                                                  CUevent start, CUevent stop) {
diff --git a/third_party/xla/xla/stream_executor/cuda/cuda_driver.h b/third_party/xla/xla/stream_executor/cuda/cuda_driver.h
index a72740ef4ead86..110cbe238fcc90 100644
--- a/third_party/xla/xla/stream_executor/cuda/cuda_driver.h
+++ b/third_party/xla/xla/stream_executor/cuda/cuda_driver.h
@@ -48,6 +48,11 @@ static std::string ToString(CUresult result) {
   return absl::StrCat(error_name, ": ", error_string);
 }
 
+// Polls (without blocking) to determine the status of an event - pending or
+// complete (or an error status).
+// http://docs.nvidia.com/cuda/cuda-driver-api/group__CUDA__EVENT.html#group__CUDA__EVENT_1g6f0704d755066b0ee705749ae911deef
+absl::StatusOr<CUresult> QueryEvent(GpuContext* context, CUevent event);
+
 // CUDAContext wraps a cuda CUcontext handle, and includes a unique id. The
 // unique id is positive, and ids are not repeated within the process.
 class GpuContext {
diff --git a/third_party/xla/xla/stream_executor/cuda/cuda_event.cc b/third_party/xla/xla/stream_executor/cuda/cuda_event.cc
index f42cf47e86e1ab..8f80e387387672 100644
--- a/third_party/xla/xla/stream_executor/cuda/cuda_event.cc
+++ b/third_party/xla/xla/stream_executor/cuda/cuda_event.cc
@@ -16,8 +16,8 @@ limitations under the License.
 #include "absl/log/log.h"
 #include "absl/status/statusor.h"
 #include "third_party/gpus/cuda/include/cuda.h"
+#include "xla/stream_executor/cuda/cuda_driver.h"
 #include "xla/stream_executor/event.h"
-#include "xla/stream_executor/gpu/gpu_driver.h"
 #include "xla/stream_executor/gpu/gpu_event.h"
 #include "xla/stream_executor/gpu/gpu_executor.h"
 
@@ -26,7 +26,7 @@ namespace gpu {
 
 Event::Status GpuEvent::PollForStatus() {
   absl::StatusOr<CUresult> status =
-      GpuDriver::QueryEvent(parent_->gpu_context(), gpu_event_);
+      QueryEvent(parent_->gpu_context(), gpu_event_);
   if (!status.ok()) {
     LOG(ERROR) << "Error polling for event status: "
                << status.status().message();
diff --git a/third_party/xla/xla/stream_executor/gpu/gpu_driver.h b/third_party/xla/xla/stream_executor/gpu/gpu_driver.h
index cb07b4bcc7ae80..c44f47b47999a7 100644
--- a/third_party/xla/xla/stream_executor/gpu/gpu_driver.h
+++ b/third_party/xla/xla/stream_executor/gpu/gpu_driver.h
@@ -777,12 +777,6 @@ class GpuDriver {
   static absl::Status RecordEvent(GpuContext* context, GpuEventHandle event,
                                   GpuStreamHandle stream);
 
-  // Polls (without blocking) to determine the status of an event - pending or
-  // complete (or an error status).
-  // http://docs.nvidia.com/cuda/cuda-driver-api/group__CUDA__EVENT.html#group__CUDA__EVENT_1g6f0704d755066b0ee705749ae911deef
-  static absl::StatusOr<GpuStatus> QueryEvent(GpuContext* context,
-                                              GpuEventHandle event);
-
   // -- Pointer-specific calls.
 
   // Returns the context in which pointer was allocated or registered.
diff --git a/third_party/xla/xla/stream_executor/gpu/gpu_types.h b/third_party/xla/xla/stream_executor/gpu/gpu_types.h
index 68e08bb41035fa..3964b0777eaeff 100644
--- a/third_party/xla/xla/stream_executor/gpu/gpu_types.h
+++ b/third_party/xla/xla/stream_executor/gpu/gpu_types.h
@@ -51,7 +51,6 @@ using GpuDevicePtr = hipDeviceptr_t;
 using GpuDeviceAttribute = hipDeviceAttribute_t;
 using GpuDeviceProperty = hipDeviceProp_t;
 using GpuModuleHandle = hipModule_t;
-using GpuStatus = hipError_t;
 using GpuFuncCachePreference = hipFuncCache_t;
 using GpuSharedMemConfig = hipSharedMemConfig;
 using GpuComplexType = hipComplex;
@@ -72,7 +71,6 @@ using GpuDevicePtr = CUdeviceptr;
 using GpuDeviceAttribute = CUdevice_attribute;
 using GpuDeviceProperty = CUdevprop;
 using GpuModuleHandle = CUmodule;
-using GpuStatus = CUresult;
 using GpuFuncCachePreference = CUfunc_cache;
 using GpuSharedMemConfig = CUsharedconfig;
 using GpuComplexType = cuComplex;
diff --git a/third_party/xla/xla/stream_executor/rocm/rocm_driver.cc b/third_party/xla/xla/stream_executor/rocm/rocm_driver.cc
index ec4404a215a60d..9a5e13cab06e1c 100644
--- a/third_party/xla/xla/stream_executor/rocm/rocm_driver.cc
+++ b/third_party/xla/xla/stream_executor/rocm/rocm_driver.cc
@@ -115,6 +115,17 @@ string ToString(hipError_t result) {
   }
 }
 
+absl::StatusOr<hipError_t> QueryEvent(GpuContext* context, hipEvent_t event) {
+  ScopedActivateContext activated{context};
+  hipError_t res = wrap::hipEventQuery(event);
+  if (res != hipSuccess && res != hipErrorNotReady) {
+    return absl::Status{
+        absl::StatusCode::kInternal,
+        absl::StrFormat("failed to query event: %s", ToString(res).c_str())};
+  }
+  return res;
+}
+
 namespace {
 
 // Returns the current context and checks that it is in the set of HIP contexts
@@ -1466,19 +1477,6 @@ struct BitPatternToValue {
   }
 }
 
-/* static */ absl::StatusOr<hipError_t> GpuDriver::QueryEvent(
-    GpuContext* context, GpuEventHandle event) {
-  ScopedActivateContext activated{context};
-  hipError_t res = wrap::hipEventQuery(event);
-  if (res != hipSuccess && res != hipErrorNotReady) {
-    return absl::Status{
-        absl::StatusCode::kInternal,
-        absl::StrFormat("failed to query event: %s", ToString(res).c_str())};
-  }
-
-  return res;
-}
-
 /* static */ bool GpuDriver::GetEventElapsedTime(GpuContext* context,
                                                  float* elapsed_milliseconds,
                                                  GpuEventHandle start,
diff --git a/third_party/xla/xla/stream_executor/rocm/rocm_driver.h b/third_party/xla/xla/stream_executor/rocm/rocm_driver.h
index 80822446a58d80..c5277ac798c097 100644
--- a/third_party/xla/xla/stream_executor/rocm/rocm_driver.h
+++ b/third_party/xla/xla/stream_executor/rocm/rocm_driver.h
@@ -32,6 +32,8 @@ namespace gpu {
 // Error summaries taken from:
 std::string ToString(hipError_t result);
 
+absl::StatusOr<hipError_t> QueryEvent(GpuContext* context, hipEvent_t event);
+
 // GpuContext wraps the device_ordinal and hipCtx_t handle.
 class GpuContext {
  public:
diff --git a/third_party/xla/xla/stream_executor/rocm/rocm_event.cc b/third_party/xla/xla/stream_executor/rocm/rocm_event.cc
index 497dee272cc1c0..730f4066ec7f62 100644
--- a/third_party/xla/xla/stream_executor/rocm/rocm_event.cc
+++ b/third_party/xla/xla/stream_executor/rocm/rocm_event.cc
@@ -16,13 +16,14 @@ limitations under the License.
 #include "xla/stream_executor/gpu/gpu_event.h"
 #include "xla/stream_executor/gpu/gpu_executor.h"
 #include "xla/stream_executor/gpu/gpu_stream.h"
+#include "xla/stream_executor/rocm/rocm_driver.h"
 
 namespace stream_executor {
 namespace gpu {
 
 Event::Status GpuEvent::PollForStatus() {
   absl::StatusOr<hipError_t> status =
-      GpuDriver::QueryEvent(parent_->gpu_context(), gpu_event_);
+      QueryEvent(parent_->gpu_context(), gpu_event_);
   if (!status.ok()) {
     LOG(ERROR) << "Error polling for event status: "
                << status.status().message();

From bf59de73ccd779bff8740113012403427eb15b48 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Thu, 16 May 2024 02:45:20 -0700
Subject: [PATCH 250/478] Automated Code Change

PiperOrigin-RevId: 634279450
---
 tensorflow/compiler/mlir/init_mlir.cc | 1 +
 1 file changed, 1 insertion(+)

diff --git a/tensorflow/compiler/mlir/init_mlir.cc b/tensorflow/compiler/mlir/init_mlir.cc
index 938cd52359b9d6..ce7cefabcdcf73 100644
--- a/tensorflow/compiler/mlir/init_mlir.cc
+++ b/tensorflow/compiler/mlir/init_mlir.cc
@@ -15,6 +15,7 @@ limitations under the License.
 
 #include "tensorflow/compiler/mlir/init_mlir.h"
 
+#include "llvm/ADT/StringRef.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/PrettyStackTrace.h"
 #include "tensorflow/core/platform/init_main.h"

From b3388cd3cfb4c79c25360286f7b1870a5afbc600 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Thu, 16 May 2024 03:08:01 -0700
Subject: [PATCH 251/478] Automated Code Change

PiperOrigin-RevId: 634285627
---
 tensorflow/core/lib/monitoring/BUILD | 15 +--------------
 1 file changed, 1 insertion(+), 14 deletions(-)

diff --git a/tensorflow/core/lib/monitoring/BUILD b/tensorflow/core/lib/monitoring/BUILD
index b832e86aedaf46..62e26fa28d30c9 100644
--- a/tensorflow/core/lib/monitoring/BUILD
+++ b/tensorflow/core/lib/monitoring/BUILD
@@ -9,14 +9,6 @@ package(
     default_visibility = [
         # tensorflow/core:lib effectively exposes all targets under tensorflow/core/lib/**
         "//tensorflow/core:__pkg__",
-        # tensorflow/core/platform:monitoring depends on this package
-        "//tensorflow/core/platform:__subpackages__",
-        # tensorflow/compiler/xla/pjrt:metrics depends on this package
-        "@local_xla//xla/pjrt:__subpackages__",
-        "@local_xla//xla/service/gpu:__subpackages__",
-        # tensorflow/compiler/mlir/tfrt:tf_jitrt depends on this package
-        "//tensorflow/compiler/mlir/tfrt:__subpackages__",
-        "@local_tsl//tsl/lib/monitoring:__subpackages__",
     ],
     licenses = ["notice"],
 )
@@ -98,7 +90,6 @@ cc_library(
     visibility = [
         "//learning/brain/google/monitoring:__subpackages__",
         "//tensorflow/core:__subpackages__",
-        "@local_tsl//tsl:__subpackages__",
     ],
     deps = [
         ":collected_metrics",
@@ -148,11 +139,7 @@ cc_library(
 cc_library(
     name = "metric_def",
     hdrs = ["metric_def.h"],
-    visibility = [
-        "//learning/brain/google/monitoring:__subpackages__",
-        "//tensorflow/core:__subpackages__",
-        "@local_tsl//tsl/lib/monitoring:__subpackages__",
-    ],
+    visibility = ["//tensorflow/core:__subpackages__"],
     deps = [
         ":types",
         "//tensorflow/core/framework:summary_proto_cc",

From 31f477032e1a7229b970f200410943b0d4034846 Mon Sep 17 00:00:00 2001
From: Ilia Sergachev <isergachev@nvidia.com>
Date: Thu, 16 May 2024 03:09:03 -0700
Subject: [PATCH 252/478] PR #12501: [NFC] Fix mistypes in scatter expander.

Imported from GitHub PR https://github.com/openxla/xla/pull/12501

Copybara import of the project:

--
97b50d47877dd3dc46535c6fa34f0e449ee57bfe by Ilia Sergachev <isergachev@nvidia.com>:

[NFC] Fix mistypes in scatter expander.

Merging this change closes #12501

PiperOrigin-RevId: 634285954
---
 third_party/xla/xla/service/gpu/gpu_compiler.cc      | 2 +-
 third_party/xla/xla/service/scatter_expander.cc      | 2 +-
 third_party/xla/xla/service/scatter_expander.h       | 8 ++++----
 third_party/xla/xla/service/scatter_expander_test.cc | 6 +++---
 4 files changed, 9 insertions(+), 9 deletions(-)

diff --git a/third_party/xla/xla/service/gpu/gpu_compiler.cc b/third_party/xla/xla/service/gpu/gpu_compiler.cc
index 92c04a51754ce9..b26d9a6cd50524 100644
--- a/third_party/xla/xla/service/gpu/gpu_compiler.cc
+++ b/third_party/xla/xla/service/gpu/gpu_compiler.cc
@@ -695,7 +695,7 @@ absl::Status RunOptimizationPasses(
     // Scatter can be indeterministic if indices are not unique or a non
     // associative combiner function is used. Eliminate these Scatter ops.
     pipeline.AddPass<ScatterExpander>(
-        ScatterExpander::kEliminateIndeterminisitcScatters);
+        ScatterExpander::kEliminateIndeterministicScatters);
   }
   // Scatters unsupported on XLA:GPU are eliminated.
   pipeline.AddPass<GpuScatterExpander>();
diff --git a/third_party/xla/xla/service/scatter_expander.cc b/third_party/xla/xla/service/scatter_expander.cc
index f76c88cabda01f..ec5d3a7726c2bc 100644
--- a/third_party/xla/xla/service/scatter_expander.cc
+++ b/third_party/xla/xla/service/scatter_expander.cc
@@ -525,7 +525,7 @@ bool ScatterExpander::InstructionMatchesPattern(HloInstruction* inst) {
   return (scatter != nullptr) && (mode_ == kEliminateAllScatters ||
                                   (mode_ == kEliminateSimpleScatters &&
                                    ScatterTripCount(scatter) == 1) ||
-                                  (mode_ == kEliminateIndeterminisitcScatters &&
+                                  (mode_ == kEliminateIndeterministicScatters &&
                                    !IsDeterministic(scatter)));
 }
 
diff --git a/third_party/xla/xla/service/scatter_expander.h b/third_party/xla/xla/service/scatter_expander.h
index d87bb52fc07380..0658e962017454 100644
--- a/third_party/xla/xla/service/scatter_expander.h
+++ b/third_party/xla/xla/service/scatter_expander.h
@@ -23,7 +23,7 @@ namespace xla {
 // This pass rewrites scatter operations into (roughly) while loops of
 // dynamic-update-slices.
 //
-// This pass can be used in two ways:
+// This pass can be used in three ways:
 //
 //   - kEliminateAllScatters: For backends that don't support scatter, this pass
 //     can convert every scatter into a loop.
@@ -32,8 +32,8 @@ namespace xla {
 //     pass can strength-reduce "simple" scatters -- specifically, scatters that
 //     can be represented without a loop -- to dynamic-update-slices.
 //
-//   - kEliminateIndeterminisitcScatters: For backends that *do* support
-//     scatter, this pass converts scatters with potentially indeterminisitc
+//   - kEliminateIndeterministicScatters: For backends that *do* support
+//     scatter, this pass converts scatters with potentially indeterministic
 //     behavior, because of non-unique indices or non-associative combiner
 //     functions. There may be false positives, but no false negatives, i.e.
 //     some scatters are converted even when deterministic in practice.
@@ -46,7 +46,7 @@ class ScatterExpander : public OpExpanderPass {
   enum Mode {
     kEliminateAllScatters,
     kEliminateSimpleScatters,
-    kEliminateIndeterminisitcScatters,
+    kEliminateIndeterministicScatters,
   };
 
   explicit ScatterExpander(Mode m) : mode_(m) {}
diff --git a/third_party/xla/xla/service/scatter_expander_test.cc b/third_party/xla/xla/service/scatter_expander_test.cc
index 71c76374f0d085..a74eabf4080ef7 100644
--- a/third_party/xla/xla/service/scatter_expander_test.cc
+++ b/third_party/xla/xla/service/scatter_expander_test.cc
@@ -270,7 +270,7 @@ TEST_F(ScatterExpanderTest, DoNotEliminateScatterWithAssociativeCombiner) {
                           ParseAndReturnVerifiedModule(kModuleStr));
 
   ScatterExpander scatter_expander(
-      ScatterExpander::kEliminateIndeterminisitcScatters);
+      ScatterExpander::kEliminateIndeterministicScatters);
   TF_ASSERT_OK_AND_ASSIGN(bool result,
                           RunHloPass(&scatter_expander, module.get()));
   EXPECT_FALSE(result);
@@ -300,7 +300,7 @@ TEST_F(ScatterExpanderTest, EliminateScatterWithNonAssociativeCombiner) {
                           ParseAndReturnVerifiedModule(kModuleStr));
 
   ScatterExpander scatter_expander(
-      ScatterExpander::kEliminateIndeterminisitcScatters);
+      ScatterExpander::kEliminateIndeterministicScatters);
   TF_ASSERT_OK_AND_ASSIGN(bool result,
                           RunHloPass(&scatter_expander, module.get()));
   EXPECT_TRUE(result);
@@ -330,7 +330,7 @@ TEST_F(ScatterExpanderTest, DoNotEliminateScatterWithAssociativeFp32Combiner) {
                           ParseAndReturnVerifiedModule(kModuleStr));
 
   ScatterExpander scatter_expander(
-      ScatterExpander::kEliminateIndeterminisitcScatters);
+      ScatterExpander::kEliminateIndeterministicScatters);
   TF_ASSERT_OK_AND_ASSIGN(bool result,
                           RunHloPass(&scatter_expander, module.get()));
   EXPECT_FALSE(result);

From 9d9aa1623bacaab2e94a2e9719e78e87b0ecf2bc Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Thu, 16 May 2024 03:17:53 -0700
Subject: [PATCH 253/478] Update ops-related pbtxt files.

PiperOrigin-RevId: 634288224
---
 tensorflow/core/ops/ops.pbtxt | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorflow/core/ops/ops.pbtxt b/tensorflow/core/ops/ops.pbtxt
index 975f4a22eca161..0751aea36d13e7 100644
--- a/tensorflow/core/ops/ops.pbtxt
+++ b/tensorflow/core/ops/ops.pbtxt
@@ -1,4 +1,4 @@
-go/debugstr  
+go/debugonly   
 op {
   name: "Abort"
   attr {

From e641445101e7cabc4f50303b3e3e9caf917337ca Mon Sep 17 00:00:00 2001
From: Quentin Khan <qkhan@google.com>
Date: Thu, 16 May 2024 03:31:26 -0700
Subject: [PATCH 254/478] XNNPack weight cache provider: refactor code.

PiperOrigin-RevId: 634291787
---
 .../lite/delegates/xnnpack/weight_cache.cc    | 65 ++++++++++++-------
 .../lite/delegates/xnnpack/weight_cache.h     | 11 ++--
 2 files changed, 47 insertions(+), 29 deletions(-)

diff --git a/tensorflow/lite/delegates/xnnpack/weight_cache.cc b/tensorflow/lite/delegates/xnnpack/weight_cache.cc
index cb178662831839..51dd6441d4afa2 100644
--- a/tensorflow/lite/delegates/xnnpack/weight_cache.cc
+++ b/tensorflow/lite/delegates/xnnpack/weight_cache.cc
@@ -19,6 +19,7 @@ limitations under the License.
 
 #if defined(_MSC_VER)
 #include <io.h>
+#define F_OK 0
 #else
 #include <sys/mman.h>
 #include <unistd.h>
@@ -76,6 +77,12 @@ class ScopeGuard {
 template <class F>
 ScopeGuard(F&&) -> ScopeGuard<F>;
 
+// Returns true if the given path exists.
+[[nodiscard]]
+bool FileExists(const char* path) {
+  return access(path, F_OK) != -1;
+}
+
 }  // namespace
 
 void swap(MMapHandle& a, MMapHandle& b) {
@@ -302,19 +309,35 @@ void MMapWeightCacheProvider::SetFilePath(const char* path) {
   XNNPACK_ABORT_CHECK(
       !IsFinalized(),
       "Cannot change the path of a cache that has already been loaded.");
-  file_path_ = path;
+  // We try to keep file_path_'s data as stable as possible. Don't overwrite
+  // if the path hasn't changed.
+  if (file_path_ != path) {
+    file_path_ = path;
+  }
 }
 
 bool MMapWeightCacheProvider::Load(const std::string& path) {
-  file_path_ = path;
-  if (mmap_handle_.Map(path.c_str())) {
-    return Load(std::move(mmap_handle_));
-  }
-  return false;
+  SetFilePath(path.c_str());
+  return Load();
 }
 
-bool MMapWeightCacheProvider::Load(MMapHandle&& handle) {
-  swap(mmap_handle_, handle);
+bool MMapWeightCacheProvider::Load() {
+  XNNPACK_ABORT_CHECK(!file_path_.empty(),
+                      "Path wasn't provided to weight cache provider.");
+  mmap_buffer_base_offset_ = 0;
+  cache_key_to_offset_.clear();
+
+  if (!FileExists(file_path_.c_str())) {
+    TFLITE_LOG_PROD(tflite::TFLITE_LOG_WARNING,
+                    "Could not load weight cache (%s): %s.", file_path_.c_str(),
+                    strerror(errno));
+    return false;
+  }
+
+  if (!mmap_handle_.Map(file_path_.c_str())) {
+    return false;
+  }
+
   // Verifiy the flabuffer part of the file.
   const size_t verifier_size =
       std::min(mmap_handle_.size(),
@@ -409,9 +432,12 @@ void* MMapWeightCacheProvider::OffsetToAddr(const size_t offset) {
   return mmap_handle_.data() + mmap_buffer_base_offset_ + offset;
 }
 
-void MMapWeightCacheProvider::Reset() {
-  MMapWeightCacheProvider empty;
-  std::swap(*this, empty);
+void MMapWeightCacheProvider::Release() {
+  buffer_address_to_identifier_.clear();
+  cache_key_to_offset_.clear();
+  mmap_handle_ = MMapHandle();
+  mmap_buffer_base_offset_ = 0;
+  builder_ = WeightCacheBuilder();
 }
 
 bool MMapWeightCacheProvider::Finalize() {
@@ -423,19 +449,12 @@ bool MMapWeightCacheProvider::Finalize() {
                     "File path wasn't set. Cannot finalize the cache.");
     return false;
   }
-  std::string file_path = file_path_;
-  if (!builder_.Write(file_path.c_str())) {
+  if (!builder_.Write(file_path_.c_str())) {
     return false;
   }
-  // The buffer mapping needs to be kept. We save it and restore it after the
-  // Reset.
-  std::unordered_map<const void*, uint64_t>
-      buffer_address_to_identifier_backup =
-          std::move(buffer_address_to_identifier_);
-  Reset();
-  buffer_address_to_identifier_ =
-      std::move(buffer_address_to_identifier_backup);
-  return Load(file_path);
+  builder_ = WeightCacheBuilder();
+
+  return Load();
 }
 
 bool MMapWeightCacheProvider::IsFinalized() const {
@@ -468,7 +487,7 @@ void* MMapWeightCacheProvider::offset_to_addr(void* context, size_t offset) {
 }
 
 enum xnn_status MMapWeightCacheProvider::delete_cache(void* context) {
-  reinterpret_cast<MMapWeightCacheProvider*>(context)->Reset();
+  reinterpret_cast<MMapWeightCacheProvider*>(context)->Release();
   return xnn_status_success;
 }
 
diff --git a/tensorflow/lite/delegates/xnnpack/weight_cache.h b/tensorflow/lite/delegates/xnnpack/weight_cache.h
index 8942ea3beb3e9e..e8c885c85cac13 100644
--- a/tensorflow/lite/delegates/xnnpack/weight_cache.h
+++ b/tensorflow/lite/delegates/xnnpack/weight_cache.h
@@ -179,14 +179,13 @@ class MMapWeightCacheProvider {
   // WARNING: Can only be called if the cache isn't finalized.
   void SetFilePath(const char* file_path);
 
-  // Loads a flatbuffer following the layout in weight_cache_schema.fbs and set
-  // the file path.
+  // Set the weight file path and loads it.
   [[nodiscard /*Loading a cache file may fail.*/]]
   bool Load(const std::string& path);
 
-  // Loads an MMap allocation following the layout in weight_cache_schema.fbs.
+  // Loads the weight cache previouslt set with `SetFilePath`.
   [[nodiscard /*Loading cache data may fail.*/]]
-  bool Load(MMapHandle&& mmap_handle);
+  bool Load();
 
   // Creates the tensor map.
   void MapTensorIdentifiers(
@@ -223,8 +222,8 @@ class MMapWeightCacheProvider {
   // WARNING: This does not check the validity of the passed offset.
   void* OffsetToAddr(size_t offset);
 
-  // Resets the weight cache provider as if it had been default constructed.
-  void Reset();
+  // Releases the weight cache's memory.
+  void Release();
 
   // Ensures that the cache is ready.
   //

From 62fe318da438ef01831e3ed2d260c63d6077638a Mon Sep 17 00:00:00 2001
From: Thomas Joerg <tjoerg@google.com>
Date: Thu, 16 May 2024 04:06:44 -0700
Subject: [PATCH 255/478] [XLA:GPU] Remove GOOGLE_CUDA and TENSORFLOW_USE_ROCM
 guards from hlo_fusion_analysis.

This is a pure HLO pass that does not require a GPU.

PiperOrigin-RevId: 634303515
---
 third_party/xla/xla/service/gpu/BUILD                  | 1 -
 third_party/xla/xla/service/gpu/hlo_fusion_analysis.cc | 2 --
 2 files changed, 3 deletions(-)

diff --git a/third_party/xla/xla/service/gpu/BUILD b/third_party/xla/xla/service/gpu/BUILD
index c36ed77518c480..a87fec8982bd4b 100644
--- a/third_party/xla/xla/service/gpu/BUILD
+++ b/third_party/xla/xla/service/gpu/BUILD
@@ -4414,7 +4414,6 @@ cc_library(
     name = "hlo_fusion_analysis",
     srcs = ["hlo_fusion_analysis.cc"],
     hdrs = ["hlo_fusion_analysis.h"],
-    local_defines = if_cuda_is_configured(["GOOGLE_CUDA=1"]),
     deps = [
         ":backend_configs_cc",
         ":hlo_traversal",
diff --git a/third_party/xla/xla/service/gpu/hlo_fusion_analysis.cc b/third_party/xla/xla/service/gpu/hlo_fusion_analysis.cc
index 86285cf2006fdb..d6460c00861fa6 100644
--- a/third_party/xla/xla/service/gpu/hlo_fusion_analysis.cc
+++ b/third_party/xla/xla/service/gpu/hlo_fusion_analysis.cc
@@ -199,7 +199,6 @@ HloFusionAnalysis::EmitterFusionKind HloFusionAnalysis::GetEmitterFusionKind()
     return EmitterFusionKind::kCustomFusion;
   }
 
-#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
   if (fusion_backend_config_.kind() == kTritonGemmFusionKind ||
       fusion_backend_config_.kind() == kTritonSoftmaxFusionKind) {
     return EmitterFusionKind::kTriton;
@@ -208,7 +207,6 @@ HloFusionAnalysis::EmitterFusionKind HloFusionAnalysis::GetEmitterFusionKind()
   if (fusion_backend_config_.kind() == kCuDnnFusionKind) {
     return EmitterFusionKind::kCuDnn;
   }
-#endif
 
   if (input_output_info_.smallest_input_dtype_bits < 8 ||
       input_output_info_.smallest_output_dtype_bits < 8) {

From b43e06101949e272ec947c0271282684128891a9 Mon Sep 17 00:00:00 2001
From: Johannes Reifferscheid <jreiffers@google.com>
Date: Thu, 16 May 2024 04:09:20 -0700
Subject: [PATCH 256/478] Optionally dump the MLIR compilation pipeline steps.

This is used as an input to mlir_replay. Note that GPU-specific ops are
not implemented in the interpreter yet (at HEAD).

PiperOrigin-RevId: 634304231
---
 .../xla/mlir/tools/mlir_replay/public/BUILD   | 19 +++++
 .../public/compiler_trace_instrumentation.cc  | 46 ++++++++++++
 .../public/compiler_trace_instrumentation.h   | 42 +++++++++++
 .../xla/xla/service/gpu/fusions/mlir/BUILD    |  4 +-
 .../gpu/fusions/mlir/mlir_fusion_emitter.cc   | 70 ++++++++++++-------
 .../gpu/fusions/mlir/mlir_fusion_emitter.h    | 14 ++--
 6 files changed, 162 insertions(+), 33 deletions(-)
 create mode 100644 third_party/xla/xla/mlir/tools/mlir_replay/public/compiler_trace_instrumentation.cc
 create mode 100644 third_party/xla/xla/mlir/tools/mlir_replay/public/compiler_trace_instrumentation.h

diff --git a/third_party/xla/xla/mlir/tools/mlir_replay/public/BUILD b/third_party/xla/xla/mlir/tools/mlir_replay/public/BUILD
index c9d831f9f61c31..39b902613bfeb5 100644
--- a/third_party/xla/xla/mlir/tools/mlir_replay/public/BUILD
+++ b/third_party/xla/xla/mlir/tools/mlir_replay/public/BUILD
@@ -7,6 +7,25 @@ package(
     default_visibility = ["//visibility:public"],
 )
 
+cc_library(
+    name = "compiler_trace_instrumentation",
+    srcs = ["compiler_trace_instrumentation.cc"],
+    hdrs = ["compiler_trace_instrumentation.h"],
+    deps = [
+        ":compiler_trace_proto_cc",
+        ":compiler_trace_proto_cc_impl",
+        "//xla/service/llvm_ir:llvm_util",
+        "@com_google_absl//absl/strings:str_format",
+        "@llvm-project//llvm:Support",
+        "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:Pass",
+        "@local_tsl//tsl/platform:env",
+        "@local_tsl//tsl/platform:logging",
+        "@local_tsl//tsl/platform:path",
+        "@local_tsl//tsl/platform:protobuf",
+    ],
+)
+
 cc_library(
     name = "execution_trace_utils",
     srcs = ["execution_trace_utils.cc"],
diff --git a/third_party/xla/xla/mlir/tools/mlir_replay/public/compiler_trace_instrumentation.cc b/third_party/xla/xla/mlir/tools/mlir_replay/public/compiler_trace_instrumentation.cc
new file mode 100644
index 00000000000000..cc3c3460c4d97b
--- /dev/null
+++ b/third_party/xla/xla/mlir/tools/mlir_replay/public/compiler_trace_instrumentation.cc
@@ -0,0 +1,46 @@
+/* Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+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.
+==============================================================================*/
+
+#include "xla/mlir/tools/mlir_replay/public/compiler_trace_instrumentation.h"
+
+#include <string>
+
+#include "absl/strings/str_format.h"
+#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
+#include "xla/service/llvm_ir/llvm_util.h"
+#include "tsl/platform/env.h"
+#include "tsl/platform/logging.h"
+#include "tsl/platform/path.h"
+
+namespace mlir {
+namespace interpreter {
+
+void MlirCompilerTraceInstrumentation::runAfterPass(Pass* pass, Operation* op) {
+  ModuleOp module = llvm::dyn_cast<ModuleOp>(op);
+  if (!module) {
+    module = op->getParentOfType<mlir::ModuleOp>();
+  }
+  if (!module) {
+    LOG(ERROR) << "Failed to find a ModuleOp: " << pass->getName().str() << ".";
+    return;
+  }
+
+  auto* item = trace_.mutable_passes()->Add();
+  item->set_after_pass(pass->getName().str());
+  *item->mutable_mlir_module() = xla::llvm_ir::DumpToString(module);
+}
+
+}  // namespace interpreter
+}  // namespace mlir
diff --git a/third_party/xla/xla/mlir/tools/mlir_replay/public/compiler_trace_instrumentation.h b/third_party/xla/xla/mlir/tools/mlir_replay/public/compiler_trace_instrumentation.h
new file mode 100644
index 00000000000000..8260c53b39965e
--- /dev/null
+++ b/third_party/xla/xla/mlir/tools/mlir_replay/public/compiler_trace_instrumentation.h
@@ -0,0 +1,42 @@
+/* Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+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.
+==============================================================================*/
+
+#ifndef XLA_MLIR_TOOLS_MLIR_REPLAY_PUBLIC_COMPILER_TRACE_INSTRUMENTATION_H_
+#define XLA_MLIR_TOOLS_MLIR_REPLAY_PUBLIC_COMPILER_TRACE_INSTRUMENTATION_H_
+
+#include <string>
+
+#include "mlir/Pass/Pass.h"  // from @llvm-project
+#include "mlir/Pass/PassInstrumentation.h"  // from @llvm-project
+#include "xla/mlir/tools/mlir_replay/public/compiler_trace.pb.h"
+
+namespace mlir {
+namespace interpreter {
+
+// Instrumentation that logs the state of the IR after each pass.
+class MlirCompilerTraceInstrumentation : public PassInstrumentation {
+ public:
+  explicit MlirCompilerTraceInstrumentation(MlirCompilationTrace& trace)
+      : trace_(trace) {}
+  void runAfterPass(Pass* pass, Operation* op) override;
+
+ private:
+  MlirCompilationTrace& trace_;
+};
+
+}  // namespace interpreter
+}  // namespace mlir
+
+#endif  // XLA_MLIR_TOOLS_MLIR_REPLAY_PUBLIC_COMPILER_TRACE_INSTRUMENTATION_H_
diff --git a/third_party/xla/xla/service/gpu/fusions/mlir/BUILD b/third_party/xla/xla/service/gpu/fusions/mlir/BUILD
index c53dd2089747de..645bafc935eac7 100644
--- a/third_party/xla/xla/service/gpu/fusions/mlir/BUILD
+++ b/third_party/xla/xla/service/gpu/fusions/mlir/BUILD
@@ -158,9 +158,12 @@ cc_library(
         "//xla:status_macros",
         "//xla:util",
         "//xla/hlo/ir:hlo",
+        "//xla/mlir/tools/mlir_replay/public:compiler_trace_instrumentation",
+        "//xla/mlir/tools/mlir_replay/public:compiler_trace_proto_cc",
         "//xla/mlir_hlo",
         "//xla/mlir_hlo:mhlo_passes",
         "//xla/service:buffer_assignment",
+        "//xla/service:dump",
         "//xla/service/gpu:ir_emitter_context",
         "//xla/service/gpu:kernel_arguments",
         "//xla/service/gpu:kernel_reuse_cache",
@@ -174,7 +177,6 @@ cc_library(
         "//xla/stream_executor:device_description",
         "@com_google_absl//absl/algorithm:container",
         "@com_google_absl//absl/container:flat_hash_map",
-        "@com_google_absl//absl/container:flat_hash_set",
         "@com_google_absl//absl/log",
         "@com_google_absl//absl/log:check",
         "@com_google_absl//absl/status",
diff --git a/third_party/xla/xla/service/gpu/fusions/mlir/mlir_fusion_emitter.cc b/third_party/xla/xla/service/gpu/fusions/mlir/mlir_fusion_emitter.cc
index 7daca77a7bd37a..9d78aa64418868 100644
--- a/third_party/xla/xla/service/gpu/fusions/mlir/mlir_fusion_emitter.cc
+++ b/third_party/xla/xla/service/gpu/fusions/mlir/mlir_fusion_emitter.cc
@@ -75,9 +75,12 @@ limitations under the License.
 #include "mlir/Transforms/Passes.h"  // from @llvm-project
 #include "xla/hlo/ir/hlo_instruction.h"
 #include "xla/hlo/ir/hlo_instructions.h"
+#include "xla/mlir/tools/mlir_replay/public/compiler_trace.pb.h"
+#include "xla/mlir/tools/mlir_replay/public/compiler_trace_instrumentation.h"
 #include "xla/mlir_hlo/mhlo/IR/hlo_ops.h"
 #include "xla/mlir_hlo/mhlo/transforms/passes.h"
 #include "xla/service/buffer_assignment.h"
+#include "xla/service/dump.h"
 #include "xla/service/gpu/fusions/fusion_emitter.h"
 #include "xla/service/gpu/fusions/mlir/computation_partitioner.h"
 #include "xla/service/gpu/fusions/mlir/elemental_hlo_to_mlir.h"
@@ -282,15 +285,19 @@ MlirFusionEmitterBase::CreateLLVMModule(
     const BufferAssignment* buffer_assignment) const {
   bool is_amd = std::holds_alternative<se::RocmComputeCapability>(
       device.gpu_compute_capability());
+  auto* hlo_module = fusion.GetModule();
+  std::unique_ptr<mlir::interpreter::MlirCompilationTrace> trace = nullptr;
+  if (DumpingEnabledForHloModule(*hlo_module) &&
+      DumpingEnabledForHloPass("mlir-fusion-emitter",
+                               hlo_module->config().debug_options())) {
+    trace = std::make_unique<mlir::interpreter::MlirCompilationTrace>();
+  }
   TF_RET_CHECK(!is_amd) << "Unsupported device type: " << device.name();
   TF_ASSIGN_OR_RETURN(
       auto module, CreateMLIRModule(mlir_context, fusion, entry_function_name,
                                     buffer_assignment));
 
   mlir::PassManager pm(&mlir_context);
-  if (VLOG_IS_ON(5)) {
-    pm.enableIRPrinting();
-  }
   pm.addPass(mlir::createCanonicalizerPass());
   pm.addPass(mlir::createInlinerPass());
   pm.addPass(mlir::createCanonicalizerPass());
@@ -328,15 +335,13 @@ MlirFusionEmitterBase::CreateLLVMModule(
   pm.addPass(CreateLowerToLLVMPass());
   pm.addPass(mlir::createReconcileUnrealizedCastsPass());
 
-  if (pm.run(module.get()).failed()) {
-    std::string module_dump;
-    llvm::raw_string_ostream os(module_dump);
-    module->print(os);
-    return absl::InternalError(absl::StrFormat(
-        "Failed create LLVM module.\nHloFusionInstruction "
-        "computation:\n%s\nMLIR module:\n%s",
-        fusion.fused_instructions_computation()->ToString(), module_dump));
+  auto pipeline_status = RunPassPipeline(module.get(), pm, trace.get());
+  if (trace) {
+    DumpPerModuleProtobufToFile(
+        *hlo_module, *trace, hlo_module->config().debug_options(),
+        absl::StrCat(entry_function_name, ".mlir-trace"));
   }
+  TF_RETURN_IF_ERROR(pipeline_status);
 
   auto llvm_module = mlir::translateModuleToLLVMIR(module.get(), llvm_context);
   TF_RET_CHECK(llvm_module != nullptr)
@@ -349,7 +354,8 @@ absl::StatusOr<mlir::OwningOpRef<mlir::ModuleOp>>
 MlirFusionEmitterBase::CreateMLIRModule(
     mlir::MLIRContext& context, const HloFusionInstruction& fusion,
     const std::string& entry_function_name,
-    const BufferAssignment* buffer_assignment) const {
+    const BufferAssignment* buffer_assignment,
+    mlir::interpreter::MlirCompilationTrace* trace) const {
   context.loadDialect<mlir::DLTIDialect, mlir::tensor::TensorDialect,
                       mlir::func::FuncDialect, mlir::affine::AffineDialect,
                       mlir::arith::ArithDialect, mlir::cf::ControlFlowDialect,
@@ -434,22 +440,13 @@ MlirFusionEmitterBase::CreateMLIRModule(
 
   // Run a minimal simplification pipeline.
   mlir::PassManager pm(&context);
-  if (VLOG_IS_ON(5)) {
-    context.disableMultithreading();
-    pm.enableIRPrinting();
-  }
   pm.addPass(CreateSimplifyArithPass());
   pm.addPass(mlir::createCanonicalizerPass());
   pm.addPass(mlir::createCSEPass());
-  if (pm.run(module.get()).failed()) {
-    std::string module_dump;
-    llvm::raw_string_ostream os(module_dump);
-    module->print(os);
-    return absl::InternalError(absl::StrFormat(
-        "Failed to simplify module.\nHloFusionInstruction "
-        "computation:\n%s\nMLIR module:\n%s",
-        fusion.fused_instructions_computation()->ToString(), module_dump));
-  }
+  // We won't dump the trace here if the pipeline fails. This is acceptable,
+  // since failures this early are usually easy to debug from the single MLIR
+  // snapshot that is dumped in RunPassPipeline.
+  TF_RETURN_IF_ERROR(RunPassPipeline(module.get(), pm, trace));
 
   return module;
 }
@@ -561,5 +558,28 @@ MlirFusionEmitterBase::EmitEpilogue(
   return results_per_root;
 }
 
+absl::Status MlirFusionEmitterBase::RunPassPipeline(
+    mlir::ModuleOp module, mlir::PassManager& pm,
+    mlir::interpreter::MlirCompilationTrace* trace) const {
+  if (VLOG_IS_ON(5)) {
+    module.getContext()->disableMultithreading();
+    pm.enableIRPrinting();
+  }
+  if (trace) {
+    module.getContext()->disableMultithreading();
+    pm.addInstrumentation(
+        std::make_unique<mlir::interpreter::MlirCompilerTraceInstrumentation>(
+            *trace));
+  }
+  if (pm.run(module).failed()) {
+    std::string module_dump;
+    llvm::raw_string_ostream os(module_dump);
+    module->print(os);
+    return absl::InternalError(absl::StrFormat(
+        "Failed to run pass pipeline.\nMLIR module:\n%s", module_dump));
+  }
+  return absl::OkStatus();
+}
+
 }  // namespace gpu
 }  // namespace xla
diff --git a/third_party/xla/xla/service/gpu/fusions/mlir/mlir_fusion_emitter.h b/third_party/xla/xla/service/gpu/fusions/mlir/mlir_fusion_emitter.h
index 1ea02ac98cd120..a555cd0b1b8cbe 100644
--- a/third_party/xla/xla/service/gpu/fusions/mlir/mlir_fusion_emitter.h
+++ b/third_party/xla/xla/service/gpu/fusions/mlir/mlir_fusion_emitter.h
@@ -15,15 +15,10 @@ limitations under the License.
 #ifndef XLA_SERVICE_GPU_FUSIONS_MLIR_MLIR_FUSION_EMITTER_H_
 #define XLA_SERVICE_GPU_FUSIONS_MLIR_MLIR_FUSION_EMITTER_H_
 
-#include <functional>
-#include <memory>
 #include <string>
-#include <vector>
 
-#include "absl/container/flat_hash_set.h"
 #include "absl/status/status.h"
 #include "absl/status/statusor.h"
-#include "llvm/ADT/SmallVector.h"
 #include "llvm/IR/LLVMContext.h"
 #include "llvm/IR/Module.h"
 #include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
@@ -34,8 +29,9 @@ limitations under the License.
 #include "mlir/IR/OwningOpRef.h"  // from @llvm-project
 #include "mlir/IR/Value.h"  // from @llvm-project
 #include "mlir/IR/ValueRange.h"  // from @llvm-project
-#include "xla/hlo/ir/hlo_instruction.h"
+#include "mlir/Pass/PassManager.h"  // from @llvm-project
 #include "xla/hlo/ir/hlo_instructions.h"
+#include "xla/mlir/tools/mlir_replay/public/compiler_trace.pb.h"
 #include "xla/service/buffer_assignment.h"
 #include "xla/service/gpu/fusions/fusion_emitter.h"
 #include "xla/service/gpu/fusions/mlir/computation_partitioner.h"
@@ -65,7 +61,8 @@ class MlirFusionEmitterBase : public KernelFusionInterface {
   absl::StatusOr<mlir::OwningOpRef<mlir::ModuleOp>> CreateMLIRModule(
       mlir::MLIRContext& context, const HloFusionInstruction& fusion,
       const std::string& entry_function_name,
-      const BufferAssignment* buffer_assignment) const;
+      const BufferAssignment* buffer_assignment,
+      mlir::interpreter::MlirCompilationTrace* trace = nullptr) const;
 
  protected:
   // Returns the set of instructions that will be isolated in the partitioned,
@@ -118,6 +115,9 @@ class MlirFusionEmitterBase : public KernelFusionInterface {
   absl::Status EmitMlir(mlir::ModuleOp module,
                         mlir::func::FuncOp entry_function,
                         const HloFusionInstruction& fusion) const;
+  absl::Status RunPassPipeline(
+      mlir::ModuleOp module, mlir::PassManager& pm,
+      mlir::interpreter::MlirCompilationTrace* trace) const;
 };
 
 }  // namespace gpu

From f9d867e34fb00b1860959be8f794e36b1a8e0eb6 Mon Sep 17 00:00:00 2001
From: Quentin Khan <qkhan@google.com>
Date: Thu, 16 May 2024 04:28:56 -0700
Subject: [PATCH 257/478] XNNPack weight cache provider: update the internal
 options object to keep track of the cache path.

PiperOrigin-RevId: 634308967
---
 tensorflow/lite/delegates/xnnpack/BUILD               | 1 +
 tensorflow/lite/delegates/xnnpack/weight_cache.h      | 2 ++
 tensorflow/lite/delegates/xnnpack/xnnpack_delegate.cc | 7 +++++--
 3 files changed, 8 insertions(+), 2 deletions(-)

diff --git a/tensorflow/lite/delegates/xnnpack/BUILD b/tensorflow/lite/delegates/xnnpack/BUILD
index 911a9358129a29..9067998f1d57d3 100644
--- a/tensorflow/lite/delegates/xnnpack/BUILD
+++ b/tensorflow/lite/delegates/xnnpack/BUILD
@@ -254,6 +254,7 @@ cc_library(
         ":tflite_with_xnnpack_transient_indirection_buffer",
         ":weight_cache",
         "//tensorflow/lite:kernel_api",
+        "//tensorflow/lite:logger",
         "//tensorflow/lite:minimal_logging",
         "//tensorflow/lite/c:c_api_types",
         "//tensorflow/lite/core:subgraph",
diff --git a/tensorflow/lite/delegates/xnnpack/weight_cache.h b/tensorflow/lite/delegates/xnnpack/weight_cache.h
index e8c885c85cac13..0eb66308a37d85 100644
--- a/tensorflow/lite/delegates/xnnpack/weight_cache.h
+++ b/tensorflow/lite/delegates/xnnpack/weight_cache.h
@@ -179,6 +179,8 @@ class MMapWeightCacheProvider {
   // WARNING: Can only be called if the cache isn't finalized.
   void SetFilePath(const char* file_path);
 
+  const std::string& GetFilePath() const { return file_path_; }
+
   // Set the weight file path and loads it.
   [[nodiscard /*Loading a cache file may fail.*/]]
   bool Load(const std::string& path);
diff --git a/tensorflow/lite/delegates/xnnpack/xnnpack_delegate.cc b/tensorflow/lite/delegates/xnnpack/xnnpack_delegate.cc
index 6ed0696d8d9e5b..30d4c6dec9b9d6 100644
--- a/tensorflow/lite/delegates/xnnpack/xnnpack_delegate.cc
+++ b/tensorflow/lite/delegates/xnnpack/xnnpack_delegate.cc
@@ -47,6 +47,7 @@ limitations under the License.
 #include "tensorflow/lite/kernels/internal/utils/sparsity_format_converter.h"
 #include "tensorflow/lite/kernels/kernel_util.h"
 #include "tensorflow/lite/kernels/padding.h"
+#include "tensorflow/lite/logger.h"
 #include "tensorflow/lite/minimal_logging.h"
 #include "tensorflow/lite/schema/schema_generated.h"
 #include "tensorflow/lite/tools/optimize/reduced_precision_support.h"
@@ -522,8 +523,8 @@ class Delegate {
     }
 
 #endif
-    TFLITE_LOG_PROD_ONCE(tflite::TFLITE_LOG_INFO,
-                         "Created TensorFlow Lite XNNPACK delegate for CPU.");
+    TFLITE_LOG_PROD(tflite::TFLITE_LOG_INFO,
+                    "Created TensorFlow Lite XNNPACK delegate for CPU.");
 
     options_ =
         options != nullptr ? *options : TfLiteXNNPackDelegateOptionsDefault();
@@ -547,6 +548,8 @@ class Delegate {
         options_.weights_cache =
             reinterpret_cast<TfLiteXNNPackDelegateWeightsCache*>(
                 weight_cache_provider_.GetCacheProvider().context);
+        options_.experimental_weight_cache_file_path =
+            weight_cache_provider_.GetFilePath().data();
       } else {
         TFLITE_LOG_PROD(tflite::TFLITE_LOG_INFO,
                         "XNNPack weight cache not enabled.");

From feae9a42520d3fdf36fca7505c0672d057ce8706 Mon Sep 17 00:00:00 2001
From: Johannes Reifferscheid <jreiffers@google.com>
Date: Thu, 16 May 2024 05:09:48 -0700
Subject: [PATCH 258/478] Use apply_indexing for input index computations.

For this to work, we need a folder, so add that too. We kind of
rely on the folder to work most of the time (for caching of
instructions). There may be some dragons here, I'll probably rewrite
that stuff at some point.

This works around the remaining known miscompiles. The issue
must be somewhere in LLVM and be related to GEPs with subtractions,
but I lack the skills and resources to actually track it down.

PiperOrigin-RevId: 634326550
---
 third_party/xla/xla/service/gpu/fusions/BUILD  |  1 +
 .../xla/xla/service/gpu/fusions/loop_mlir.cc   | 14 ++++++++------
 .../gpu/fusions/mlir/elemental_hlo_to_mlir.cc  |  4 ++--
 .../fusions/mlir/elemental_hlo_to_mlir_test.cc | 18 ++++++++----------
 .../service/gpu/fusions/mlir/ir/xla_gpu_ops.cc | 16 ++++++++++++++++
 .../service/gpu/fusions/mlir/ir/xla_gpu_ops.td |  1 +
 6 files changed, 36 insertions(+), 18 deletions(-)

diff --git a/third_party/xla/xla/service/gpu/fusions/BUILD b/third_party/xla/xla/service/gpu/fusions/BUILD
index 17b48188bed3ea..6ba5e41a4e9f2e 100644
--- a/third_party/xla/xla/service/gpu/fusions/BUILD
+++ b/third_party/xla/xla/service/gpu/fusions/BUILD
@@ -359,6 +359,7 @@ cc_library(
     hdrs = ["loop_mlir.h"],
     deps = [
         ":loop",
+        "//xla:shape_util",
         "//xla:status",
         "//xla:status_macros",
         "//xla:xla_data_proto_cc",
diff --git a/third_party/xla/xla/service/gpu/fusions/loop_mlir.cc b/third_party/xla/xla/service/gpu/fusions/loop_mlir.cc
index 0861ab46057d45..d52043f117efec 100644
--- a/third_party/xla/xla/service/gpu/fusions/loop_mlir.cc
+++ b/third_party/xla/xla/service/gpu/fusions/loop_mlir.cc
@@ -42,6 +42,7 @@ limitations under the License.
 #include "xla/service/gpu/launch_dimensions.h"
 #include "xla/service/gpu/model/indexing_analysis.h"
 #include "xla/service/gpu/model/indexing_map.h"
+#include "xla/shape.h"
 #include "xla/status_macros.h"
 #include "xla/xla_data.pb.h"
 
@@ -124,8 +125,9 @@ absl::Status MlirLoopFusion::EmitEntryFunction(
 
   auto body_builder = [&](ValueRange output_tensors, ValueRange dim_values,
                           ValueRange symbol_values) -> SmallVector<Value> {
-    auto first_output_indices = mlir_converter::ApplyAffineMap(
-        indexing->GetAffineMap(), dim_values, symbol_values, builder);
+    llvm::SmallVector<Value> first_output_indices;
+    builder.createOrFold<ApplyIndexingOp>(first_output_indices, dim_values,
+                                          symbol_values, *indexing);
     auto root_fn = call_targets(
         fusion.fused_instructions_computation()->root_instruction());
 
@@ -141,11 +143,11 @@ absl::Status MlirLoopFusion::EmitEntryFunction(
     result_tensors.reserve(output_tensor_args.size());
     for (auto [root_shape, tensor, value] :
          llvm::zip(result_shapes, output_tensors, result_scalars)) {
-      auto output_indices = mlir_converter::ApplyAffineMap(
+      llvm::SmallVector<Value> output_indices;
+      builder.createOrFold<ApplyIndexingOp>(
+          output_indices, first_output_indices, ValueRange{},
           GetBitcastMap(*result_shapes.front(), *root_shape,
-                        builder.getContext())
-              .GetAffineMap(),
-          first_output_indices, {}, builder);
+                        builder.getContext()));
       result_tensors.push_back(builder.create<mlir::tensor::InsertOp>(
           value, tensor, output_indices));
     }
diff --git a/third_party/xla/xla/service/gpu/fusions/mlir/elemental_hlo_to_mlir.cc b/third_party/xla/xla/service/gpu/fusions/mlir/elemental_hlo_to_mlir.cc
index ae7071d7dadb46..5433fa6d31f9ef 100644
--- a/third_party/xla/xla/service/gpu/fusions/mlir/elemental_hlo_to_mlir.cc
+++ b/third_party/xla/xla/service/gpu/fusions/mlir/elemental_hlo_to_mlir.cc
@@ -483,9 +483,9 @@ SmallVector<SmallVector<Value>> GetInputIndices(
   SmallVector<SmallVector<Value>> indices;
   for (auto& maps : indexing.indexing_maps) {
     CHECK_EQ(maps.size(), 1);
-    auto map = maps.begin()->GetAffineMap();
     CHECK(!maps.begin()->IsUndefined());
-    indices.emplace_back() = ApplyAffineMap(map, output_indices, {}, b);
+    b.createOrFold<ApplyIndexingOp>(indices.emplace_back(), output_indices,
+                                    ValueRange{}, *maps.begin());
   }
   return indices;
 }
diff --git a/third_party/xla/xla/service/gpu/fusions/mlir/elemental_hlo_to_mlir_test.cc b/third_party/xla/xla/service/gpu/fusions/mlir/elemental_hlo_to_mlir_test.cc
index 5a865cf6847cd4..4ffe6e5613a436 100644
--- a/third_party/xla/xla/service/gpu/fusions/mlir/elemental_hlo_to_mlir_test.cc
+++ b/third_party/xla/xla/service/gpu/fusions/mlir/elemental_hlo_to_mlir_test.cc
@@ -418,11 +418,10 @@ TEST_F(ElementalHloToMlirTest, Pad) {
     // CHECK:        %[[Y_BOUNDS:.*]] = arith.andi %[[Y_L]], %[[Y_H]]
     // CHECK:        %[[FROM_INPUT:.*]] = arith.andi %[[X_AND_CONSTRAINT]], %[[Y_BOUNDS]]
     // CHECK:        %[[RET:.*]] = scf.if %[[FROM_INPUT]]
-    // CHECK:          %[[X_IN:.*]] = affine.apply
-    // CHECK-SAME:         <()[s0] -> ((s0 - 1) floordiv 2)>()[%[[X]]]
-    // CHECK:          %[[Y_IN:.*]] = affine.apply
-    // CHECK-SAME:         <()[s0] -> (s0 - 4)>()[%[[Y]]]
-    // CHECK:          %[[VAL:.*]] = tensor.extract %[[ARG0]][%[[X_IN]], %[[Y_IN]]]
+    // CHECK:          %[[IN:.*]]:2 = xla_gpu.apply_indexing
+    // CHECK-SAME:         <(d0, d1) -> ((d0 - 1) floordiv 2, d1 - 4)>
+    // CHECK-SAME:         (%[[X]] in [1, 7], %[[Y]] in [4, 7])
+    // CHECK:          %[[VAL:.*]] = tensor.extract %[[ARG0]][%[[IN]]#0, %[[IN]]#1]
     // CHECK:          scf.yield %[[VAL]]
     // CHECK:        } else {
     // CHECK:          %[[PAD_VAL:.*]] = tensor.extract %[[ARG1]][]
@@ -461,11 +460,10 @@ TEST_F(ElementalHloToMlirTest, PadUnsigned) {
     // CHECK:        %[[Y_BOUNDS:.*]] = arith.andi %[[Y_L]], %[[Y_H]]
     // CHECK:        %[[FROM_INPUT:.*]] = arith.andi %[[X_AND_CONSTRAINT]], %[[Y_BOUNDS]]
     // CHECK:        %[[RET:.*]] = scf.if %[[FROM_INPUT]]
-    // CHECK:          %[[X_IN:.*]] = affine.apply
-    // CHECK-SAME:         <()[s0] -> ((s0 - 1) floordiv 2)>()[%[[X]]]
-    // CHECK:          %[[Y_IN:.*]] = affine.apply
-    // CHECK-SAME:         <()[s0] -> (s0 - 4)>()[%[[Y]]]
-    // CHECK:          %[[VAL:.*]] = tensor.extract %[[ARG0]][%[[X_IN]], %[[Y_IN]]]
+    // CHECK:          %[[IN:.*]]:2 = xla_gpu.apply_indexing
+    // CHECK-SAME:         <(d0, d1) -> ((d0 - 1) floordiv 2, d1 - 4)>
+    // CHECK-SAME:         (%[[X]] in [1, 7], %[[Y]] in [4, 7])
+    // CHECK:          %[[VAL:.*]] = tensor.extract %[[ARG0]][%[[IN]]#0, %[[IN]]#1]
     // CHECK:          %[[CAST0:.*]] = builtin.unrealized_conversion_cast %[[VAL]]
     // CHECK:          scf.yield %[[CAST0]]
     // CHECK:        } else {
diff --git a/third_party/xla/xla/service/gpu/fusions/mlir/ir/xla_gpu_ops.cc b/third_party/xla/xla/service/gpu/fusions/mlir/ir/xla_gpu_ops.cc
index b419f6e3b95605..4cc02c974cb420 100644
--- a/third_party/xla/xla/service/gpu/fusions/mlir/ir/xla_gpu_ops.cc
+++ b/third_party/xla/xla/service/gpu/fusions/mlir/ir/xla_gpu_ops.cc
@@ -531,6 +531,22 @@ void ApplyIndexingOp::getCanonicalizationPatterns(
               SimplifyIndexingMap>(context);
 }
 
+mlir::LogicalResult ApplyIndexingOp::fold(
+    FoldAdaptor adaptor, llvm::SmallVectorImpl<mlir::OpFoldResult> &results) {
+  auto map = getAffineMap();
+  for (auto expr : map.getResults()) {
+    if (auto dim = mlir::dyn_cast<mlir::AffineDimExpr>(expr)) {
+      results.push_back(getOperand(dim.getPosition()));
+    } else if (auto sym = mlir::dyn_cast<mlir::AffineSymbolExpr>(expr)) {
+      results.push_back(getOperand(map.getNumDims() + sym.getPosition()));
+    } else {
+      results.clear();
+      return failure();
+    }
+  }
+  return success();
+}
+
 //===----------------------------------------------------------------------===//
 // AtomicRMWOp
 //===----------------------------------------------------------------------===//
diff --git a/third_party/xla/xla/service/gpu/fusions/mlir/ir/xla_gpu_ops.td b/third_party/xla/xla/service/gpu/fusions/mlir/ir/xla_gpu_ops.td
index 317f978b748d88..fb89e2bed2164d 100644
--- a/third_party/xla/xla/service/gpu/fusions/mlir/ir/xla_gpu_ops.td
+++ b/third_party/xla/xla/service/gpu/fusions/mlir/ir/xla_gpu_ops.td
@@ -278,6 +278,7 @@ def ApplyIndexingOp : XLAGPU_Op<"apply_indexing", [Pure]> {
   let hasCustomAssemblyFormat = 1;
   let hasVerifier = 1;
   let hasCanonicalizer = 1;
+  let hasFolder = 1;
 }
 
 #endif // MLIR_HLO_DIALECT_MHLO_IR_HLO_OPS

From 8621ef58f120b7b5f3eec734457366e01187d89b Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Thu, 16 May 2024 05:18:00 -0700
Subject: [PATCH 259/478] Update ops-related pbtxt files.

PiperOrigin-RevId: 634330969
---
 tensorflow/core/ops/ops.pbtxt | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorflow/core/ops/ops.pbtxt b/tensorflow/core/ops/ops.pbtxt
index 0751aea36d13e7..2a7a7c62c9d010 100644
--- a/tensorflow/core/ops/ops.pbtxt
+++ b/tensorflow/core/ops/ops.pbtxt
@@ -1,4 +1,4 @@
-go/debugonly   
+go/debugstr    
 op {
   name: "Abort"
   attr {

From 3fa9094d9ddcbdcc579ea1ef674fe6fcfd2ca09b Mon Sep 17 00:00:00 2001
From: Quentin Khan <qkhan@google.com>
Date: Thu, 16 May 2024 05:58:45 -0700
Subject: [PATCH 260/478] XNNPack weight cache: make logging message format
 consistent across the module.

PiperOrigin-RevId: 634349913
---
 .../lite/delegates/xnnpack/weight_cache.cc    | 63 +++++++++++--------
 1 file changed, 38 insertions(+), 25 deletions(-)

diff --git a/tensorflow/lite/delegates/xnnpack/weight_cache.cc b/tensorflow/lite/delegates/xnnpack/weight_cache.cc
index 51dd6441d4afa2..aa7307386ff2cd 100644
--- a/tensorflow/lite/delegates/xnnpack/weight_cache.cc
+++ b/tensorflow/lite/delegates/xnnpack/weight_cache.cc
@@ -105,9 +105,10 @@ bool MMapHandle::Map(const char* path) {
 
   const int fd = open(path, O_RDONLY);
   if (fd == -1) {
-    TFLITE_LOG_PROD(tflite::TFLITE_LOG_ERROR,
-                    "Could not open file to mmap: %s (%s).", strerror(errno),
-                    path)
+    TFLITE_LOG_PROD(
+        tflite::TFLITE_LOG_ERROR,
+        "XNNPack weight cache: could not open file to mmap ('%s'): %s.", path,
+        strerror(errno))
     return false;
   }
 
@@ -120,8 +121,9 @@ bool MMapHandle::Map(const char* path) {
   struct stat file_stats;
   if (fstat(fd, &file_stats)) {
     TFLITE_LOG_PROD(tflite::TFLITE_LOG_ERROR,
-                    "Could not access file stats to get size: %s (%s).",
-                    strerror(errno), path)
+                    "XNNPack weight cache: could not access file stats to get "
+                    "size ('%s'): %s.",
+                    path, strerror(errno))
     return false;
   }
 
@@ -135,8 +137,8 @@ bool MMapHandle::Map(const char* path) {
       const auto bytes = read(fd, data_reader, remaining_bytes);
       if (bytes == -1) {
         TFLITE_LOG_PROD(tflite::TFLITE_LOG_ERROR,
-                        "Could not read file ('%s'): %s.", path,
-                        strerror(errno))
+                        "XNNPack weight cache: could not read file ('%s'): %s.",
+                        path, strerror(errno))
         UnMap();
         return false;
       }
@@ -148,8 +150,9 @@ bool MMapHandle::Map(const char* path) {
   data_ = static_cast<uint8_t*>(
       mmap(/*addr=*/nullptr, size_, PROT_READ, MAP_SHARED, fd, /*offset=*/0));
   if (data_ == MAP_FAILED) {
-    TFLITE_LOG_PROD(tflite::TFLITE_LOG_ERROR, "Could not mmap file: %s (%s).",
-                    strerror(errno), path)
+    TFLITE_LOG_PROD(tflite::TFLITE_LOG_ERROR,
+                    "XNNPack weight cache: could not mmap file (%s): %s.", path,
+                    strerror(errno))
     data_ = nullptr;
     size_ = 0;
     return false;
@@ -221,9 +224,10 @@ bool WriteData(const int fd, const uint8_t* data, size_t size,
   for (size_t bytes = 0; bytes < size;) {
     const auto written_bytes = write(fd, data + bytes, size - bytes);
     if (written_bytes == -1) {
-      TFLITE_LOG_PROD(tflite::TFLITE_LOG_ERROR,
-                      "Cache file write incomplete (%s). %s: %s", file_path,
-                      step_description, strerror(errno))
+      TFLITE_LOG_PROD(
+          tflite::TFLITE_LOG_ERROR,
+          "XNNPack weight cache: file write incomplete (%s). %s: %s.",
+          file_path, step_description, strerror(errno))
     }
     bytes += written_bytes;
   }
@@ -236,9 +240,11 @@ bool WriteData(const int fd, const uint8_t* data, size_t size,
 bool WeightCacheBuilder::Write(const char* path) {
   const int fd = open(path, O_WRONLY | O_CREAT | O_TRUNC, 0644);
   if (fd == -1) {
-    TFLITE_LOG_PROD(tflite::TFLITE_LOG_ERROR,
-                    "Could not open cache file ('%s') for writing: %s", path,
-                    strerror(errno))
+    TFLITE_LOG_PROD(
+        tflite::TFLITE_LOG_ERROR,
+        "XNNPack weight cache: could not open cache file ('%s') for "
+        "writing: %s.",
+        path, strerror(errno))
     return false;
   }
 
@@ -265,7 +271,8 @@ bool WeightCacheBuilder::Write(const char* path) {
   mutable_packed_weights->mutate_base_offset(builder.GetSize() +
                                              alignment_offset);
 
-  // Write the flatbuffer which serves as a header to index the following data.
+  // Write the flatbuffer which serves as a header to index the following
+  // data.
   if (!WriteData(fd, builder.GetBufferPointer(), builder.GetSize(), path,
                  "Header")) {
     return false;
@@ -281,6 +288,8 @@ bool WeightCacheBuilder::Write(const char* path) {
                  "Buffer data")) {
     return false;
   }
+  TFLITE_LOG_PROD(tflite::TFLITE_LOG_INFO,
+                  "XNNPack weight cache: written to '%s'.", path);
   return true;
 }
 
@@ -329,8 +338,8 @@ bool MMapWeightCacheProvider::Load() {
 
   if (!FileExists(file_path_.c_str())) {
     TFLITE_LOG_PROD(tflite::TFLITE_LOG_WARNING,
-                    "Could not load weight cache (%s): %s.", file_path_.c_str(),
-                    strerror(errno));
+                    "XNNPack weight cache: could not load '%s': %s.",
+                    file_path_.c_str(), strerror(errno));
     return false;
   }
 
@@ -345,7 +354,7 @@ bool MMapWeightCacheProvider::Load() {
   flatbuffers::Verifier verifier(mmap_handle_.data(), verifier_size);
   if (!cache::schema::VerifyPackedWeightsBuffer(verifier)) {
     TFLITE_LOG_PROD(tflite::TFLITE_LOG_ERROR,
-                    "Packed weights buffer validation failed.");
+                    "XNNPack weight cache: header validation failed.");
     return false;
   }
 
@@ -353,16 +362,18 @@ bool MMapWeightCacheProvider::Load() {
   const cache::schema::PackedWeights* packed_weights =
       cache::schema::GetPackedWeights(mmap_handle_.data());
   if (!packed_weights) {
-    TFLITE_LOG_PROD(tflite::TFLITE_LOG_ERROR,
-                    "Could not get packed weights from flatbuffer.");
+    TFLITE_LOG_PROD(
+        tflite::TFLITE_LOG_ERROR,
+        "XNNPack weight cache: could not get packed weights from flatbuffer.");
     return false;
   }
   mmap_buffer_base_offset_ = packed_weights->base_offset();
   if (const auto buffers = packed_weights->buffers(); buffers) {
     for (auto* buffer : *buffers) {
       if (!buffer) {
-        TFLITE_LOG_PROD(tflite::TFLITE_LOG_ERROR,
-                        "Invalid buffer address in buffer list.");
+        TFLITE_LOG_PROD(
+            tflite::TFLITE_LOG_ERROR,
+            "XNNPack weight cache: Invalid buffer address in buffer list.");
         return false;
       }
       cache_key_to_offset_.emplace(
@@ -388,7 +399,8 @@ void MMapWeightCacheProvider::MapTensorIdentifiers(
 size_t MMapWeightCacheProvider::LookUp(
     const xnn_weights_cache_look_up_key* cache_key) {
   if (!cache_key) {
-    TFLITE_LOG_PROD(tflite::TFLITE_LOG_ERROR, "A null cache key was provided.");
+    TFLITE_LOG_PROD(tflite::TFLITE_LOG_ERROR,
+                    "XNNPack weight cache: a null cache key was provided.");
     return SIZE_MAX;
   }
   const PackIdentifier pack_id = BuildPackIdentifier(*cache_key);
@@ -446,7 +458,8 @@ bool MMapWeightCacheProvider::Finalize() {
   }
   if (file_path_.empty()) {
     TFLITE_LOG_PROD(tflite::TFLITE_LOG_ERROR,
-                    "File path wasn't set. Cannot finalize the cache.");
+                    "XNNPack weight cache: file path wasn't set. Cannot "
+                    "finalize the cache.");
     return false;
   }
   if (!builder_.Write(file_path_.c_str())) {

From 15da74aa06b1979128bead6d405d537941987698 Mon Sep 17 00:00:00 2001
From: Thomas Joerg <tjoerg@google.com>
Date: Thu, 16 May 2024 06:06:41 -0700
Subject: [PATCH 261/478] [XLA:GPU] Remove GOOGLE_CUDA and TENSORFLOW_USE_ROCM
 defines from gemm_algorithm_picker.

These defines are not needed since the BUILD target is anyway guarded by if_gpu_is_configured.

PiperOrigin-RevId: 634353818
---
 third_party/xla/xla/service/gpu/BUILD                  | 10 +---------
 .../xla/xla/service/gpu/gemm_algorithm_picker.cc       |  9 +--------
 .../xla/xla/service/gpu/gemm_algorithm_picker.h        |  5 +----
 3 files changed, 3 insertions(+), 21 deletions(-)

diff --git a/third_party/xla/xla/service/gpu/BUILD b/third_party/xla/xla/service/gpu/BUILD
index a87fec8982bd4b..a28626e06d3725 100644
--- a/third_party/xla/xla/service/gpu/BUILD
+++ b/third_party/xla/xla/service/gpu/BUILD
@@ -1129,9 +1129,6 @@ cc_library(
     name = "gemm_rewriter",
     srcs = ["gemm_rewriter.cc"],
     hdrs = ["gemm_rewriter.h"],
-    local_defines = if_cuda_is_configured(["GOOGLE_CUDA=1"]) + if_rocm_is_configured([
-        "TENSORFLOW_USE_ROCM=1",
-    ]),
     deps = [
         ":backend_configs_cc",
         ":cublas_cudnn",
@@ -1166,9 +1163,7 @@ cc_library(
         "@local_tsl//tsl/platform:ml_dtypes",
         "@local_tsl//tsl/platform:statusor",
         "@local_tsl//tsl/protobuf:dnn_proto_cc",
-    ] + if_cuda_is_configured([
-        "@local_config_cuda//cuda:cuda_headers",
-    ]),
+    ],
 )
 
 cc_library(
@@ -1534,9 +1529,6 @@ cc_library(
     name = "gemm_algorithm_picker",
     srcs = if_gpu_is_configured(["gemm_algorithm_picker.cc"]),
     hdrs = if_gpu_is_configured(["gemm_algorithm_picker.h"]),
-    local_defines = if_cuda_is_configured(["GOOGLE_CUDA=1"]) + if_rocm_is_configured([
-        "TENSORFLOW_USE_ROCM=1",
-    ]),
     deps = if_gpu_is_configured([
         ":backend_configs_cc",
         ":buffer_comparator",
diff --git a/third_party/xla/xla/service/gpu/gemm_algorithm_picker.cc b/third_party/xla/xla/service/gpu/gemm_algorithm_picker.cc
index c70d2a35d91522..e43703f362b525 100644
--- a/third_party/xla/xla/service/gpu/gemm_algorithm_picker.cc
+++ b/third_party/xla/xla/service/gpu/gemm_algorithm_picker.cc
@@ -36,6 +36,7 @@ limitations under the License.
 #include "xla/service/gpu/autotuner_compile_util.h"
 #include "xla/service/gpu/autotuner_util.h"
 #include "xla/service/gpu/backend_configs.pb.h"
+#include "xla/service/gpu/buffer_comparator.h"
 #include "xla/service/gpu/cublas_cudnn.h"
 #include "xla/service/gpu/matmul_utils.h"
 #include "xla/service/gpu/stream_executor_util.h"
@@ -55,10 +56,6 @@ limitations under the License.
 #include "tsl/platform/statusor.h"
 #include "tsl/profiler/lib/scoped_annotation.h"
 
-#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
-#include "xla/service/gpu/buffer_comparator.h"
-#endif
-
 namespace xla {
 namespace gpu {
 namespace {
@@ -89,8 +86,6 @@ absl::StatusOr<BlasLt::Epilogue> AsBlasLtEpilogue(
   }
 }
 
-#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
-
 class GemmAutotuner {
   const AutotuneConfig& autotune_config_;
   RedzoneBuffers rz_buffers_;
@@ -371,8 +366,6 @@ class GemmAutotuner {
   }  // GetBestAlgorithm
 };  // GemmAutotuner
 
-#endif  // GOOGLE_CUDA || TENSORFLOW_USE_ROCM
-
 // Do Gemm Autotune without stream executor. Use results from autotune cache
 // only.
 absl::StatusOr<bool> RunOnInstruction(HloInstruction* gemm,
diff --git a/third_party/xla/xla/service/gpu/gemm_algorithm_picker.h b/third_party/xla/xla/service/gpu/gemm_algorithm_picker.h
index 6610a95279b6ba..cc70e0d16a21af 100644
--- a/third_party/xla/xla/service/gpu/gemm_algorithm_picker.h
+++ b/third_party/xla/xla/service/gpu/gemm_algorithm_picker.h
@@ -33,11 +33,8 @@ limitations under the License.
 #include "xla/stream_executor/blas.h"
 #include "xla/stream_executor/device_memory.h"
 #include "xla/stream_executor/device_memory_allocator.h"
-#include "xla/stream_executor/stream_executor.h"
-
-#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
 #include "xla/stream_executor/gpu/redzone_allocator.h"
-#endif
+#include "xla/stream_executor/stream_executor.h"
 
 namespace xla {
 namespace gpu {

From ed5ebc6c3a47d33b01ba3e5f993020588a9b7ae3 Mon Sep 17 00:00:00 2001
From: Quentin Khan <qkhan@google.com>
Date: Thu, 16 May 2024 06:09:31 -0700
Subject: [PATCH 262/478] Add a function to get delegate options using the C
 API.

- Remove redundant helper function.
- Store the weight cache path inside the delegate and update the internal
  options object to use it.

PiperOrigin-RevId: 634354940
---
 tensorflow/lite/core/kernels/BUILD            |  1 -
 tensorflow/lite/core/kernels/register_test.cc | 25 +++++++++---------
 tensorflow/lite/delegates/xnnpack/BUILD       |  9 -------
 .../delegates/xnnpack/xnnpack_delegate.cc     | 19 +++++++-------
 .../lite/delegates/xnnpack/xnnpack_delegate.h |  7 +++++
 .../delegates/xnnpack/xnnpack_delegate_test.h | 26 -------------------
 6 files changed, 29 insertions(+), 58 deletions(-)
 delete mode 100644 tensorflow/lite/delegates/xnnpack/xnnpack_delegate_test.h

diff --git a/tensorflow/lite/core/kernels/BUILD b/tensorflow/lite/core/kernels/BUILD
index fcb5a458edcbce..71f9980fe1b957 100644
--- a/tensorflow/lite/core/kernels/BUILD
+++ b/tensorflow/lite/core/kernels/BUILD
@@ -25,7 +25,6 @@ cc_test(
         ":builtin_ops",
         "//tensorflow/lite:mutable_op_resolver",
         "//tensorflow/lite/delegates/xnnpack:xnnpack_delegate",
-        "//tensorflow/lite/delegates/xnnpack:xnnpack_delegate_test",
         "//tensorflow/lite/schema:schema_fbs",
         "@com_google_googletest//:gtest_main",
     ],
diff --git a/tensorflow/lite/core/kernels/register_test.cc b/tensorflow/lite/core/kernels/register_test.cc
index f47a7f84397c35..6bbc1bd0db07ae 100644
--- a/tensorflow/lite/core/kernels/register_test.cc
+++ b/tensorflow/lite/core/kernels/register_test.cc
@@ -20,7 +20,6 @@ limitations under the License.
 #include <gmock/gmock.h>
 #include <gtest/gtest.h>
 #include "tensorflow/lite/delegates/xnnpack/xnnpack_delegate.h"
-#include "tensorflow/lite/delegates/xnnpack/xnnpack_delegate_test.h"
 #include "tensorflow/lite/mutable_op_resolver.h"
 #include "tensorflow/lite/schema/schema_generated.h"
 
@@ -57,13 +56,13 @@ TEST(BuiltinOpResolverTest, HasXNNPACKDelegate_QS8) {
       builtin_op_resolver.GetDelegateCreators()[0];
   std::unique_ptr<TfLiteDelegate, void (*)(TfLiteDelegate *)> delegate =
       delegate_creator(nullptr);
-  const void *delegate_data = TfLiteOpaqueDelegateGetData(delegate.get());
-  TfLiteXNNPackDelegateOptions options = GetOptions(delegate_data);
+  const TfLiteXNNPackDelegateOptions *options =
+      TfLiteXNNPackDelegateGetOptions(delegate.get());
 
-  ASSERT_EQ(options.flags & TFLITE_XNNPACK_DELEGATE_FLAG_QU8,
+  ASSERT_EQ(options->flags & TFLITE_XNNPACK_DELEGATE_FLAG_QU8,
             TFLITE_XNNPACK_DELEGATE_FLAG_QU8);
 
-  ASSERT_EQ(options.flags & TFLITE_XNNPACK_DELEGATE_FLAG_QS8,
+  ASSERT_EQ(options->flags & TFLITE_XNNPACK_DELEGATE_FLAG_QS8,
             TFLITE_XNNPACK_DELEGATE_FLAG_QS8);
 }
 
@@ -74,13 +73,13 @@ TEST(BuiltinOpResolverTest, HasXNNPACKDelegate_QS8_QU8) {
       builtin_op_resolver.GetDelegateCreators()[0];
   std::unique_ptr<TfLiteDelegate, void (*)(TfLiteDelegate *)> delegate =
       delegate_creator(nullptr);
-  const void *delegate_data = TfLiteOpaqueDelegateGetData(delegate.get());
-  TfLiteXNNPackDelegateOptions options = GetOptions(delegate_data);
+  const TfLiteXNNPackDelegateOptions *options =
+      TfLiteXNNPackDelegateGetOptions(delegate.get());
 
-  ASSERT_EQ(options.flags & TFLITE_XNNPACK_DELEGATE_FLAG_QU8,
+  ASSERT_EQ(options->flags & TFLITE_XNNPACK_DELEGATE_FLAG_QU8,
             TFLITE_XNNPACK_DELEGATE_FLAG_QU8);
 
-  ASSERT_EQ(options.flags & TFLITE_XNNPACK_DELEGATE_FLAG_QS8,
+  ASSERT_EQ(options->flags & TFLITE_XNNPACK_DELEGATE_FLAG_QS8,
             TFLITE_XNNPACK_DELEGATE_FLAG_QS8);
 }
 
@@ -91,12 +90,12 @@ TEST(BuiltinOpResolverTest, Disable_QU8) {
       builtin_op_resolver.GetDelegateCreators()[0];
   std::unique_ptr<TfLiteDelegate, void (*)(TfLiteDelegate *)> delegate =
       delegate_creator(nullptr);
-  const void *delegate_data = TfLiteOpaqueDelegateGetData(delegate.get());
-  TfLiteXNNPackDelegateOptions options = GetOptions(delegate_data);
+  const TfLiteXNNPackDelegateOptions *options =
+      TfLiteXNNPackDelegateGetOptions(delegate.get());
 
-  ASSERT_EQ(options.flags & TFLITE_XNNPACK_DELEGATE_FLAG_QU8, 0);
+  ASSERT_EQ(options->flags & TFLITE_XNNPACK_DELEGATE_FLAG_QU8, 0);
 
-  ASSERT_EQ(options.flags & TFLITE_XNNPACK_DELEGATE_FLAG_QS8,
+  ASSERT_EQ(options->flags & TFLITE_XNNPACK_DELEGATE_FLAG_QS8,
             TFLITE_XNNPACK_DELEGATE_FLAG_QS8);
 }
 #endif  // TFLITE_WITHOUT_XNNPACK
diff --git a/tensorflow/lite/delegates/xnnpack/BUILD b/tensorflow/lite/delegates/xnnpack/BUILD
index 9067998f1d57d3..a2841aecca92cd 100644
--- a/tensorflow/lite/delegates/xnnpack/BUILD
+++ b/tensorflow/lite/delegates/xnnpack/BUILD
@@ -215,15 +215,6 @@ cc_library(
     }),
 )
 
-cc_library(
-    name = "xnnpack_delegate_test",
-    testonly = True,
-    hdrs = ["xnnpack_delegate_test.h"],
-    deps = [
-        ":xnnpack_delegate",
-    ],
-)
-
 cc_library(
     name = "xnnpack_delegate",
     srcs = ["xnnpack_delegate.cc"],
diff --git a/tensorflow/lite/delegates/xnnpack/xnnpack_delegate.cc b/tensorflow/lite/delegates/xnnpack/xnnpack_delegate.cc
index 30d4c6dec9b9d6..26d24bca18032d 100644
--- a/tensorflow/lite/delegates/xnnpack/xnnpack_delegate.cc
+++ b/tensorflow/lite/delegates/xnnpack/xnnpack_delegate.cc
@@ -694,7 +694,7 @@ class Delegate {
 #endif
   }
 
-  TfLiteXNNPackDelegateOptions options() const { return options_; }
+  const TfLiteXNNPackDelegateOptions& options() const { return options_; }
 
   int64_t GetXNNPackDelegateFlags() {
     if (enable_subgraph_reshaping()) {
@@ -7880,14 +7880,6 @@ TfLiteXNNPackDelegateOptions TfLiteXNNPackDelegateOptionsDefault() {
   return options;
 }
 
-TfLiteXNNPackDelegateOptions GetOptions(const void* delegate_data) {
-  if (delegate_data == nullptr) {
-    return TfLiteXNNPackDelegateOptionsDefault();
-  }
-  return static_cast<const tflite::xnnpack::Delegate*>(delegate_data)
-      ->options();
-}
-
 TfLiteDelegate* TfLiteXNNPackDelegateCreate(
     const TfLiteXNNPackDelegateOptions* options) {
   return TfLiteXNNPackDelegateCreateWithThreadpool(options, nullptr);
@@ -7919,6 +7911,15 @@ void* TfLiteXNNPackDelegateGetThreadPool(TfLiteDelegate* delegate) {
       static_cast<::tflite::xnnpack::Delegate*>(delegate->data_)->threadpool());
 }
 
+const TfLiteXNNPackDelegateOptions* TfLiteXNNPackDelegateGetOptions(
+    TfLiteDelegate* delegate) {
+  if (delegate == nullptr) {
+    return nullptr;
+  }
+  return &(static_cast<const tflite::xnnpack::Delegate*>(delegate->data_)
+               ->options());
+}
+
 int TfLiteXNNPackDelegateGetFlags(TfLiteDelegate* delegate) {
   if (delegate == nullptr) {
     return 0;
diff --git a/tensorflow/lite/delegates/xnnpack/xnnpack_delegate.h b/tensorflow/lite/delegates/xnnpack/xnnpack_delegate.h
index 1d40c04db940ac..55eddcf1a54d67 100644
--- a/tensorflow/lite/delegates/xnnpack/xnnpack_delegate.h
+++ b/tensorflow/lite/delegates/xnnpack/xnnpack_delegate.h
@@ -101,6 +101,13 @@ TfLiteDelegate* TfLiteXNNPackDelegateCreateWithThreadpool(
 TFL_CAPI_EXPORT void* TfLiteXNNPackDelegateGetThreadPool(
     TfLiteDelegate* delegate);
 
+// Returns the options in the delegate.
+// Returns NULL if the delegate is NULL.
+//
+// WARNING: This API is experimental and subject to change.
+TFL_CAPI_EXPORT const TfLiteXNNPackDelegateOptions*
+TfLiteXNNPackDelegateGetOptions(TfLiteDelegate* delegate);
+
 // Returns the flags used for an XNNPack delegate.
 // See documentation for TfLiteXNNPackDelegateOptions.flags.
 //
diff --git a/tensorflow/lite/delegates/xnnpack/xnnpack_delegate_test.h b/tensorflow/lite/delegates/xnnpack/xnnpack_delegate_test.h
deleted file mode 100644
index b20d024f06debe..00000000000000
--- a/tensorflow/lite/delegates/xnnpack/xnnpack_delegate_test.h
+++ /dev/null
@@ -1,26 +0,0 @@
-/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
-
-Licensed under the Apache License, Version 2.0 (the "License");
-you may not use this file except in compliance with the License.
-You may obtain a copy of the License at
-
-    http://www.apache.org/licenses/LICENSE-2.0
-
-Unless required by applicable law or agreed to in writing, software
-distributed under the License is distributed on an "AS IS" BASIS,
-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.
-==============================================================================*/
-
-#ifndef TENSORFLOW_LITE_DELEGATES_XNNPACK_XNNPACK_DELEGATE_TEST_H_
-#define TENSORFLOW_LITE_DELEGATES_XNNPACK_XNNPACK_DELEGATE_TEST_H_
-
-#include "tensorflow/lite/delegates/xnnpack/xnnpack_delegate.h"
-
-// This test only function has undefined behaviour unless delegate was obtained
-// by calling TfLiteOpaqueDelegateGetData on a TfLiteOpaqueDelegate* that refers
-// to an XNNPACK delegate, or from the 'data' field of a TfLiteDelegate* that
-// refers to an XNNPACK delegate.
-TfLiteXNNPackDelegateOptions GetOptions(const void* delegate_data);
-#endif  // TENSORFLOW_LITE_DELEGATES_XNNPACK_XNNPACK_DELEGATE_TEST_H_

From 228e8a43debffee6ddda30bbe4a404acd92dd6c2 Mon Sep 17 00:00:00 2001
From: Jane Liu <janeliu@nvidia.com>
Date: Thu, 16 May 2024 06:26:52 -0700
Subject: [PATCH 263/478] PR #11895: Offloading 3/3: enable rematerialization
 using xla flags

Imported from GitHub PR https://github.com/openxla/xla/pull/11895

This enables activation offloading automatically on GPU through HLO Rematerialization.
There is one xla flag needed to trigger rematerialization offloading:
1. A bool to enable host memory offloading. It's false by default to disable the offloading.

The other information is achieved by querying the GPU device:
2. Memory bandwidth in bytes per second.
3. The number of floating point operations on a device.

Copybara import of the project:

--
46b6232dd0c9dd34d7b170f049e6074ccee002bd by Jane Liu <janeliu@nvidia.com>:

Offloading 3/3: enable rematerialization using a bool xla flag.
Query the memory bandwidth and flops from GPUs.

Merging this change closes #11895

PiperOrigin-RevId: 634360725
---
 third_party/xla/xla/debug_options_flags.cc    |  7 ++++
 .../xla/xla/service/gpu/gpu_compiler.cc       | 34 +++++++++++++++++--
 third_party/xla/xla/xla.proto                 |  5 ++-
 3 files changed, 42 insertions(+), 4 deletions(-)

diff --git a/third_party/xla/xla/debug_options_flags.cc b/third_party/xla/xla/debug_options_flags.cc
index 8b782dbbe5dc39..7232a0a8b157bc 100644
--- a/third_party/xla/xla/debug_options_flags.cc
+++ b/third_party/xla/xla/debug_options_flags.cc
@@ -251,6 +251,8 @@ DebugOptions DefaultDebugOptionsIgnoringFlags() {
 
   opts.set_xla_gpu_require_complete_aot_autotune_results(false);
 
+  opts.set_xla_gpu_enable_host_memory_offloading(false);
+
   return opts;
 }
 
@@ -1670,6 +1672,11 @@ void MakeDebugOptionsFlags(std::vector<tsl::Flag>* flag_list,
           &DebugOptions::set_xla_reduce_window_rewrite_base_length),
       debug_options->xla_reduce_window_rewrite_base_length(),
       "Base length to rewrite the reduce window to, no rewrite if set to 0."));
+  flag_list->push_back(tsl::Flag(
+      "xla_gpu_enable_host_memory_offloading",
+      bool_setter_for(&DebugOptions::set_xla_gpu_enable_host_memory_offloading),
+      debug_options->xla_gpu_enable_host_memory_offloading(),
+      "Whether to trigger host memory offloading on a device."));
 }  // NOLINT(readability/fn_size)
 
 // Allocates flag_values and flag_objects; this function must not be called more
diff --git a/third_party/xla/xla/service/gpu/gpu_compiler.cc b/third_party/xla/xla/service/gpu/gpu_compiler.cc
index b26d9a6cd50524..609c92872f05e1 100644
--- a/third_party/xla/xla/service/gpu/gpu_compiler.cc
+++ b/third_party/xla/xla/service/gpu/gpu_compiler.cc
@@ -2161,16 +2161,44 @@ absl::Status GpuCompiler::RunPostSchedulingPipelines(
   {
     HloPassPipeline pipeline("remat-pipeline");
 
-    HloCostAnalysis hlo_cost_analysis(ShapeSizeBytesFunction());
+    const bool enable_offloading = module->config()
+                                       .debug_options()
+                                       .xla_gpu_enable_host_memory_offloading();
     HloRematerialization::RematerializationModeConfig
         rematerialization_mode_config(/*recompute=*/true, /*compress=*/true,
-                                      /*host_offload=*/false);
+                                      /*host_offload=*/enable_offloading);
+    HloCostAnalysis::Options hlo_cost_analysis_options;
+    hlo_cost_analysis_options.shape_size = ShapeSizeBytesFunction();
+    std::optional<HloRematerialization::HostMemoryOffloadConfig>
+        offloading_config = std::nullopt;
+    if (enable_offloading) {
+      constexpr float kGiga = 1e+9;
+      // Fused multiply-add means that these two instructions are computed as
+      // one, so for this case the maximum flops is doubled.
+      constexpr float kFma = 2;
+      float flops_per_sec = gpu_device_info.core_count() *
+                            gpu_device_info.fpus_per_core() *
+                            gpu_device_info.clock_rate_ghz() * kGiga * kFma;
+      int64_t host_memory_space_color =
+          static_cast<int64_t>(se::MemoryType::kHost);
+      hlo_cost_analysis_options.set_flops_per_second(flops_per_sec);
+      hlo_cost_analysis_options.set_transcendentals_per_second(flops_per_sec);
+      offloading_config =
+          std::make_optional<HloRematerialization::HostMemoryOffloadConfig>(
+              /*host_memory_space=*/host_memory_space_color,
+              /*bandwidth_to_host_bytes_per_second=*/
+              gpu_device_info.memory_bandwidth(),
+              /*bandwidth_from_host_bytes_per_second=*/
+              gpu_device_info.memory_bandwidth());
+    }
+    HloCostAnalysis hlo_cost_analysis(hlo_cost_analysis_options);
     HloRematerialization::Options options(
         hlo_cost_analysis, rematerialization_mode_config,
         // Assume 75% of the total device memory is available for XLA.
         /*memory_limit_bytes=*/scheduler_mem_limit,
         /*block_size_limit=*/1, /*block_rematerialization_factor=*/1,
-        /*min_remat_size=*/0, /*compact_shape_function=*/nullptr);
+        /*min_remat_size=*/0, /*compact_shape_function=*/nullptr,
+        /*host_memory_offload_config=*/offloading_config);
     HloRematerialization::RematerializationSizes sizes;
     pipeline.AddPass<HloRematerialization>(options, sizes);
     pipeline.AddPass<StreamAttributeAnnotator>();
diff --git a/third_party/xla/xla/xla.proto b/third_party/xla/xla/xla.proto
index 391ddd053ffcad..ea742ebc06cd61 100644
--- a/third_party/xla/xla/xla.proto
+++ b/third_party/xla/xla/xla.proto
@@ -768,7 +768,10 @@ message DebugOptions {
   // Base length to rewrite the reduce window to, no rewrite if set to 0.
   int64 xla_reduce_window_rewrite_base_length = 293;
 
-  // Next id: 296
+  // If true, will enable host memory offloading on a device.
+  bool xla_gpu_enable_host_memory_offloading = 296;
+
+  // Next id: 297
 
   // Extra options to pass to the compilation backend (e.g. LLVM); specific
   // interpretation of these values is left to the backend.

From c9e8a7f37a0ab305316ef8173b509055a51e3690 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Tam=C3=A1s=20Danyluk?= <tdanyluk@google.com>
Date: Thu, 16 May 2024 06:27:06 -0700
Subject: [PATCH 264/478] [XLA:GPU] Use IndexingMap's simplifier in tiling

PiperOrigin-RevId: 634360780
---
 third_party/xla/xla/service/gpu/model/BUILD   |  1 +
 .../xla/service/gpu/model/symbolic_tile.cc    | 42 +++++++++++++------
 .../service/gpu/model/symbolic_tile_test.cc   | 16 ++++---
 3 files changed, 38 insertions(+), 21 deletions(-)

diff --git a/third_party/xla/xla/service/gpu/model/BUILD b/third_party/xla/xla/service/gpu/model/BUILD
index 2a889bddd38b24..82c40cf5b23c5f 100644
--- a/third_party/xla/xla/service/gpu/model/BUILD
+++ b/third_party/xla/xla/service/gpu/model/BUILD
@@ -526,6 +526,7 @@ cc_library(
     hdrs = ["symbolic_tile.h"],
     deps = [
         ":affine_map_printer",
+        ":indexing_analysis",
         ":indexing_map",
         "@com_google_absl//absl/algorithm:container",
         "@com_google_absl//absl/log",
diff --git a/third_party/xla/xla/service/gpu/model/symbolic_tile.cc b/third_party/xla/xla/service/gpu/model/symbolic_tile.cc
index bedf16c1389d3d..18ddd718c107f8 100644
--- a/third_party/xla/xla/service/gpu/model/symbolic_tile.cc
+++ b/third_party/xla/xla/service/gpu/model/symbolic_tile.cc
@@ -36,6 +36,7 @@ limitations under the License.
 #include "mlir/IR/MLIRContext.h"  // from @llvm-project
 #include "mlir/Support/LLVM.h"  // from @llvm-project
 #include "xla/service/gpu/model/affine_map_printer.h"
+#include "xla/service/gpu/model/indexing_analysis.h"
 #include "xla/service/gpu/model/indexing_map.h"
 
 namespace xla {
@@ -278,6 +279,30 @@ std::optional<SizeAndStrideExpression> ExtractSizeAndStride(
   LOG(FATAL) << "unreachable";
 }
 
+// Simplifies the given affine expression using the constraints / bounds of
+// the reference indexing map.
+//
+// The dimensions and symbols of the expression should correspond to the
+// dimensions and symbols of the reference indexing map.
+AffineExpr SimplifyAffineExpr(const AffineExpr& expr,
+                              const IndexingMap& reference) {
+  AffineMap tmp_affine_map =
+      AffineMap::get(/*dimCount=*/reference.GetDimVars().size(),
+                     /*symbolCount=*/reference.GetSymbolCount(),
+                     /*results=*/{expr},
+                     /*context=*/reference.GetMLIRContext());
+  IndexingMap tmp_indexing_map(
+      /*affine_map=*/std::move(tmp_affine_map),
+      /*dimensions=*/reference.GetDimVars(),
+      /*range_vars=*/reference.GetRangeVars(),
+      /*rt_vars=*/reference.GetRTVars(),
+      /*constraints=*/reference.GetConstraints());
+  tmp_indexing_map.Simplify(GetIndexingMapForInstruction);
+
+  CHECK_EQ(tmp_indexing_map.GetAffineMap().getResults().size(), 1);
+  return tmp_indexing_map.GetAffineMap().getResults().back();
+}
+
 }  // anonymous namespace
 
 /*static*/ std::optional<SymbolicTile> SymbolicTile::FromIndexingMap(
@@ -322,6 +347,9 @@ std::optional<SizeAndStrideExpression> ExtractSizeAndStride(
           input_affine_map, getAffineConstantExpr(0, mlir_context),
           indexing_map.GetRangeVarsCount())
           .getResults();
+  for (AffineExpr& expr : offset_expressions) {
+    expr = SimplifyAffineExpr(expr, indexing_map);
+  }
 
   std::vector<AffineExpr> size_expressions;
   std::vector<AffineExpr> stride_expressions;
@@ -333,19 +361,9 @@ std::optional<SizeAndStrideExpression> ExtractSizeAndStride(
   for (auto [composite_indexing, offset] :
        llvm::zip(input_affine_map.getResults(), offset_expressions)) {
     std::optional<SizeAndStrideExpression> maybe_size_and_stride =
-        ExtractSizeAndStride(composite_indexing - offset,
+        ExtractSizeAndStride(SimplifyAffineExpr(composite_indexing - offset,
+                                                /*reference=*/indexing_map),
                              indexing_map.GetSymbolBounds());
-    if (!maybe_size_and_stride.has_value()) {
-      // Retry with a simplified expression.
-      // For example `(d0 + s0 - s0)` will be simplified to `d0`.
-      // But the simplification doesn't help when it rewrites `mod` to
-      // `floordiv` & `add`, so at first we try without simplification.
-      maybe_size_and_stride = ExtractSizeAndStride(
-          simplifyAffineExpr(composite_indexing - offset,
-                             input_affine_map.getNumDims(),
-                             input_affine_map.getNumSymbols()),
-          indexing_map.GetSymbolBounds());
-    }
     if (!maybe_size_and_stride.has_value()) {
       VLOG(1) << "No size and stride extracted";
       return std::nullopt;
diff --git a/third_party/xla/xla/service/gpu/model/symbolic_tile_test.cc b/third_party/xla/xla/service/gpu/model/symbolic_tile_test.cc
index f5703ecb8d662a..d3611b33d676db 100644
--- a/third_party/xla/xla/service/gpu/model/symbolic_tile_test.cc
+++ b/third_party/xla/xla/service/gpu/model/symbolic_tile_test.cc
@@ -314,20 +314,18 @@ TEST_F(SymbolicTileTest, CanPropagateTileThroughDynamicSlice) {
   EXPECT_THAT(
       SymbolicTile::FromIndexingMap(*input_indexing.indexing_maps[0].begin()),
       // s0, s1, s2: tile sizes
-      // s3, s4, s5: runtime parameters
-      // TODO(tdanyluk): If the RTVar can only have 1 value, maybe we should
-      // optimize it out?
+      // s3, s4: runtime parameters
+      // Note: We don't have s0 in the size map's rhs, because the first dim
+      // of the tile size can only be 1. The second offset is optimized to 0,
+      // because that is the only possible value.
       Optional(MatchSymbolicTileWithRtVars(
-          "()[s0, s1, s2, s3, s4, s5] -> (s3, s4, s5)",
-          "()[s0, s1, s2] -> (s0, s1, s2)", "()[s0, s1, s2] -> (1, 1, 1)",
+          "()[s0, s1, s2, s3, s4] -> (s3, 0, s4)",
+          "()[s0, s1, s2] -> (1, s1, s2)", "()[s0, s1, s2] -> (0, 1, 1)",
           R"(
 s3 in [0, 1]
   hlo: %of1 = s32[] parameter(1)
   (d0, d1, d2) -> ()
-s4 in [0, 0]
-  hlo: %of2 = s32[] parameter(2)
-  (d0, d1, d2) -> ()
-s5 in [0, 226]
+s4 in [0, 226]
   hlo: %of3 = s32[] parameter(3)
   (d0, d1, d2) -> ()
 )")));

From dbf4e387f0e3ba64157e3b4951601d8c38841776 Mon Sep 17 00:00:00 2001
From: Ruturaj Vaidya <ruturaj.vaidya@amd.com>
Date: Thu, 16 May 2024 07:03:07 -0700
Subject: [PATCH 265/478] PR #12550: [ROCM] Avoid hard coded specifications for
 mi300 transpose fusion

Imported from GitHub PR https://github.com/openxla/xla/pull/12550

Copybara import of the project:

--
426c2bee41800651ab191d830865c210d9f2cbea by Ruturaj4 <ruturaj.vaidya@amd.com>:

[ROCm] Avoid hard coded code for mi300 tranpose fusion

Merging this change closes #12550

PiperOrigin-RevId: 634374108
---
 .../xla/xla/service/gpu/fusions/fusions.cc    |  3 ++-
 .../xla/xla/service/gpu/fusions/transpose.cc  | 23 ++++++++++++++++---
 .../xla/xla/service/gpu/fusions/transpose.h   |  3 ++-
 3 files changed, 24 insertions(+), 5 deletions(-)

diff --git a/third_party/xla/xla/service/gpu/fusions/fusions.cc b/third_party/xla/xla/service/gpu/fusions/fusions.cc
index 3567463575539f..3f524403bd70ce 100644
--- a/third_party/xla/xla/service/gpu/fusions/fusions.cc
+++ b/third_party/xla/xla/service/gpu/fusions/fusions.cc
@@ -221,7 +221,8 @@ absl::StatusOr<std::unique_ptr<FusionInterface>> GetFusionEmitter(
       if (check_mlir_emitters()) {
         return std::make_unique<MlirTransposeFusion>(analysis);
       }
-      return std::make_unique<TransposeFusion>(analysis);
+      return std::make_unique<TransposeFusion>(analysis.device_info(),
+                                               analysis);
     }
     case HloFusionAnalysis::EmitterFusionKind::kConcatenate: {
       if (check_mlir_emitters()) {
diff --git a/third_party/xla/xla/service/gpu/fusions/transpose.cc b/third_party/xla/xla/service/gpu/fusions/transpose.cc
index 326f7b0ae30c5b..319941d8e5d59f 100644
--- a/third_party/xla/xla/service/gpu/fusions/transpose.cc
+++ b/third_party/xla/xla/service/gpu/fusions/transpose.cc
@@ -55,7 +55,8 @@ namespace xla {
 namespace gpu {
 namespace {
 
-Tiling ComputeTransposeTiling(const TransposeDescription& tiled_transpose) {
+Tiling ComputeTransposeTiling(const se::DeviceDescription& gpu_device_info,
+                              const TransposeDescription& tiled_transpose) {
   constexpr int kNumRows = 4;
   static_assert(WarpSize() % kNumRows == 0);
 
@@ -77,6 +78,20 @@ Tiling ComputeTransposeTiling(const TransposeDescription& tiled_transpose) {
   absl::InlinedVector<int64_t, 4> num_threads{1, 1, WarpSize()};
   num_threads[permutation[2]] = kNumRows;
 
+  auto capability = gpu_device_info.gpu_compute_capability();
+  std::visit(
+      [&](const auto& capability) {
+        if constexpr (std::is_same_v<std::decay_t<decltype(capability)>,
+                                     stream_executor::RocmComputeCapability>) {
+          // kNumRows = 8 works well on MI300 with wavefront size 64.
+          if (capability.gfx9_mi300()) {
+            tile_sizes[permutation[2]] = gpu_device_info.threads_per_warp() / 8;
+            num_threads[permutation[2]] = 8;
+          }
+        }
+      },
+      capability);
+
   return Tiling(input_dims, tile_sizes, num_threads);
 }
 
@@ -106,9 +121,11 @@ llvm_ir::IrArray::Index PermuteIndex(const llvm_ir::IrArray::Index& index,
 
 }  // namespace
 
-TransposeFusion::TransposeFusion(const HloFusionAnalysis& analysis)
+TransposeFusion::TransposeFusion(const se::DeviceDescription& gpu_device_info,
+                                 const HloFusionAnalysis& analysis)
     : analysis_(analysis),
-      tiling_(ComputeTransposeTiling(analysis.tiled_transpose())) {
+      tiling_(
+          ComputeTransposeTiling(gpu_device_info, analysis.tiled_transpose())) {
   for (auto [root, hero] :
        llvm::zip(analysis_.fusion_roots(), analysis_.fusion_heroes())) {
     if (auto transpose = GetDescriptionForTiledTransposeEmitter(*root, *hero)) {
diff --git a/third_party/xla/xla/service/gpu/fusions/transpose.h b/third_party/xla/xla/service/gpu/fusions/transpose.h
index 899b1cb94390ae..f9e7fca7a1aa6f 100644
--- a/third_party/xla/xla/service/gpu/fusions/transpose.h
+++ b/third_party/xla/xla/service/gpu/fusions/transpose.h
@@ -60,7 +60,8 @@ namespace gpu {
 // efficient to launch fewer blocks so each transposes many tiles.
 class TransposeFusion : public KernelFusionEmitterBase {
  public:
-  explicit TransposeFusion(const HloFusionAnalysis& analysis);
+  explicit TransposeFusion(const se::DeviceDescription& gpu_device_info,
+                           const HloFusionAnalysis& analysis);
   LaunchDimensions launch_dimensions() const override;
 
   std::optional<IndexingMap> ComputeThreadIdToOutputIndexing(

From f6a70491868200fa1c1c27f8e0748e40689d396d Mon Sep 17 00:00:00 2001
From: Ruturaj Vaidya <ruturaj.vaidya@amd.com>
Date: Thu, 16 May 2024 07:07:34 -0700
Subject: [PATCH 266/478] PR #12494: [ROCM] GPU performance cost model specs
 for MI300

Imported from GitHub PR https://github.com/openxla/xla/pull/12494

@ddunl  @xla-rotation  : would you take a look please?
Copybara import of the project:

--
265d19a155907357fe89ff4e0f5dbbe52f329c4c by Ruturaj4 <ruturaj.vaidya@amd.com>:

[ROCm] set performance cost model specifications for mi300

--
0f0a6fa282bf4ce729d8f4d570fa9a1d6f8fad50 by Ruturaj4 <ruturaj.vaidya@amd.com>:

[ROCm] set performance cost model specifications for mi100 and mi200

Merging this change closes #12494

PiperOrigin-RevId: 634375698
---
 .../gpu/model/gpu_performance_model_base.cc   | 24 +++---
 .../gpu/model/gpu_performance_model_base.h    |  6 --
 .../xla/stream_executor/device_description.h  | 75 +++++++++++++++++--
 3 files changed, 82 insertions(+), 23 deletions(-)

diff --git a/third_party/xla/xla/service/gpu/model/gpu_performance_model_base.cc b/third_party/xla/xla/service/gpu/model/gpu_performance_model_base.cc
index 66bbea0b8402b5..f1b7a9025cef6a 100644
--- a/third_party/xla/xla/service/gpu/model/gpu_performance_model_base.cc
+++ b/third_party/xla/xla/service/gpu/model/gpu_performance_model_base.cc
@@ -57,22 +57,18 @@ bool FusionUsesParameterElementwiseFromRoot(
              fusion->fused_expression_root()) == 1.f;
 }
 
-int GetCoalescingWasteFactor(PrimitiveType element_type) {
+int GetCoalescingWasteFactor(PrimitiveType element_type,
+                             const se::DeviceDescription& gpu_device_info) {
   int64_t element_size_bytes =
       element_type == PrimitiveType::TUPLE ||
               element_type == PrimitiveType::TOKEN
           ? 4 /* Dummy value. TODO(jreiffers): Model this case. */
           : ShapeUtil::ByteSizeOfPrimitiveType(element_type);
-  // Cache line is 128B that is split into 4 sectors of 32B. Default transaction
-  // size from DRAM -> L2 = 64 Bytes = 2 sectors, since V100, but it can be also
-  // configured.
-  // https://developer.download.nvidia.com/video/gputechconf/gtc/2020/presentations/s21819-optimizing-applications-for-nvidia-ampere-gpu-architecture.pdf
-  // (page 10).
-  constexpr int kDRAMToL2TransactionSizeBytes = 64;
   // Assume we use one element from the cache line and waste the remaining
   // bandwidth. For example, if we're reading f32s, we use 1/16nd of the cache
   // line.
-  return kDRAMToL2TransactionSizeBytes / element_size_bytes;
+  return gpu_device_info.dram_to_l2_transaction_size_bytes() /
+         element_size_bytes;
 }
 
 // Limit the bandwidth for low occupancy cases. Each SM can issue at most
@@ -80,7 +76,8 @@ int GetCoalescingWasteFactor(PrimitiveType element_type) {
 // (1830 MHz) to saturate the memory bandwidth (3.35 TB/s).
 float AdjustBandwidth(const se::DeviceDescription& gpu_device_info,
                       float bandwidth, int64_t num_blocks) {
-  float per_block_bandwidth = gpu_device_info.clock_rate_ghz() * 1.0e9f * 32;
+  float per_block_bandwidth = gpu_device_info.clock_rate_ghz() * 1.0e9f *
+                              gpu_device_info.memory_transactions_per_clock();
   float max_bandwidth = num_blocks * per_block_bandwidth;
 
   return std::min(bandwidth, max_bandwidth);
@@ -286,7 +283,8 @@ absl::Duration GpuPerformanceModelBase::ReadTime(
   float bandwidth = gpu_device_info.memory_bandwidth();
   if (n_bytes_net < gpu_device_info.l2_cache_size()) {
     bandwidth *= kL2CacheSpeedup;
-    if (n_bytes_net < kL1CacheSizePerSM * gpu_device_info.core_count()) {
+    if (n_bytes_net <
+        gpu_device_info.l1_cache_size_per_SM() * gpu_device_info.core_count()) {
       bandwidth *= kL1CacheSpeedup;
     }
   }
@@ -300,7 +298,8 @@ absl::Duration GpuPerformanceModelBase::ReadTimeWithDRAMHeuristic(
     const se::DeviceDescription& gpu_device_info, int64_t num_blocks,
     int64_t n_bytes_net, int64_t n_bytes_total, PrimitiveType element_type,
     bool coalesced) {
-  int waste_factor = coalesced ? 1 : GetCoalescingWasteFactor(element_type);
+  int waste_factor =
+      coalesced ? 1 : GetCoalescingWasteFactor(element_type, gpu_device_info);
 
   // The first read of the input buffer always happens from DRAM. If reads are
   // no coaleced, bandwidth is reduced by the waste factor.
@@ -313,7 +312,8 @@ absl::Duration GpuPerformanceModelBase::ReadTimeWithDRAMHeuristic(
   float rest_bandwidth = gpu_device_info.memory_bandwidth();
   if (n_bytes_net < gpu_device_info.l2_cache_size()) {
     rest_bandwidth *= kL2CacheSpeedup;
-    if (n_bytes_net < kL1CacheSizePerSM * gpu_device_info.core_count()) {
+    if (n_bytes_net <
+        gpu_device_info.l1_cache_size_per_SM() * gpu_device_info.core_count()) {
       rest_bandwidth *= kL1CacheSpeedup;
     }
   } else {
diff --git a/third_party/xla/xla/service/gpu/model/gpu_performance_model_base.h b/third_party/xla/xla/service/gpu/model/gpu_performance_model_base.h
index c7185e1fee04a5..835746b9c25316 100644
--- a/third_party/xla/xla/service/gpu/model/gpu_performance_model_base.h
+++ b/third_party/xla/xla/service/gpu/model/gpu_performance_model_base.h
@@ -144,12 +144,6 @@ class GpuPerformanceModelBase {
       absl::Microseconds(5);
   static constexpr float kL2CacheSpeedup = 2.5;
   static constexpr float kL1CacheSpeedup = 8;
-  // A very conservative estimate. L1 size varies because it can be dynamically
-  // configured as shared memory; there is no easy way to query its actual size;
-  // also we do not count what occupies cache, but rather claim that what is
-  // much smaller than the cache size will likely stay in it.
-  // For reference, it can be up to 256 kB per SM on RTX A6000.
-  static constexpr float kL1CacheSizePerSM = 2 * 1024;
 
   // Uses HloFusionAnalysis for computing the actual number of threads and
   // blocks that the IR emitter will use.
diff --git a/third_party/xla/xla/stream_executor/device_description.h b/third_party/xla/xla/stream_executor/device_description.h
index 7a7c35bc55659b..74227faf3dae76 100644
--- a/third_party/xla/xla/stream_executor/device_description.h
+++ b/third_party/xla/xla/stream_executor/device_description.h
@@ -183,6 +183,15 @@ class RocmComputeCapability {
     return absl::StrJoin(kSupportedGfxVersions, ", ");
   }
 
+  bool gfx9_mi100() const { return gfx_version() == "gfx908"; }
+
+  bool gfx9_mi200() const { return gfx_version() == "gfx90a"; }
+
+  bool gfx9_mi300() const {
+    static constexpr absl::string_view kList[] = {"gfx940", "gfx941", "gfx942"};
+    return absl::c_count(kList, gfx_version()) != 0;
+  }
+
   bool gfx9_mi100_or_later() const {
     static constexpr absl::string_view kList[] = {"gfx908", "gfx90a", "gfx940",
                                                   "gfx941", "gfx942"};
@@ -195,11 +204,6 @@ class RocmComputeCapability {
     return absl::c_count(kList, gfx_version()) != 0;
   }
 
-  bool gfx9_mi300() const {
-    static constexpr absl::string_view kList[] = {"gfx940", "gfx941", "gfx942"};
-    return absl::c_count(kList, gfx_version()) != 0;
-  }
-
   bool navi21() const { return gfx_version() == "gfx1030"; }
 
   bool navi31() const { return gfx_version() == "gfx1100"; }
@@ -396,6 +400,67 @@ class DeviceDescription {
     return shared_memory_per_block_optin_;
   }
 
+  // L1 size varies because it can be dynamically
+  // configured as shared memory; there is no easy way to query its actual size;
+  // also we do not count what occupies cache, but rather claim that what is
+  // much smaller than the cache size will likely stay in it.
+  constexpr int64_t l1_cache_size_per_SM() const {
+    return std::visit(
+        [](const auto &capability) -> int64_t {
+          if constexpr (std::is_same_v<std::decay_t<decltype(capability)>,
+                                       RocmComputeCapability>) {
+            // MI100 and MI200 has 16KB L1 cache per CU.
+            if (capability.gfx9_mi100() || capability.gfx9_mi200()) {
+              return 16 * 1024;
+            }
+            // MI300 has 32KB L1 cache per CU.
+            if (capability.gfx9_mi300()) {
+              return 32 * 1024;
+            }
+          }
+          // Default return for other GPUs (e.g., RTX A6000).
+          return 2 * 1024;
+        },
+        gpu_compute_capability_);
+  }
+
+  constexpr int64_t dram_to_l2_transaction_size_bytes() const {
+    return std::visit(
+        [](const auto &capability) -> int {
+          if constexpr (std::is_same_v<std::decay_t<decltype(capability)>,
+                                       RocmComputeCapability>) {
+            // DRAM->L2 bus is 128 Byte width for MI300.
+            if (capability.gfx9_mi300()) {
+              return 128;
+            }
+          }
+          // Cache line is 128B that is split into 4 sectors of 32B. Default
+          // transaction size from DRAM -> L2 = 64 Bytes = 2 sectors, since
+          // V100, but it can be also configured.
+          // https://developer.download.nvidia.com/video/gputechconf/gtc/2020/presentations/s21819-optimizing-applications-for-nvidia-ampere-gpu-architecture.pdf
+          // (page 10).
+          // return 64 Bytes by default.
+          return 64;
+        },
+        gpu_compute_capability_);
+  }
+
+  constexpr int64_t memory_transactions_per_clock() const {
+    return std::visit(
+        [](const auto &capability) -> int {
+          if constexpr (std::is_same_v<std::decay_t<decltype(capability)>,
+                                       RocmComputeCapability>) {
+            // 16 works well on MI300.
+            if (capability.gfx9_mi300()) {
+              return 16;
+            }
+          }
+          // Default return for other GPUs.
+          return 32;
+        },
+        gpu_compute_capability_);
+  }
+
   GpuDeviceInfoProto ToGpuProto() const;
 
   std::string ToString() const;

From 7ba79f064fe22db6b7014a77148bb0a0cfe4cb75 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Thu, 16 May 2024 07:20:54 -0700
Subject: [PATCH 267/478] Update ops-related pbtxt files.

PiperOrigin-RevId: 634380668
---
 tensorflow/core/ops/ops.pbtxt | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorflow/core/ops/ops.pbtxt b/tensorflow/core/ops/ops.pbtxt
index 2a7a7c62c9d010..cdc0191b8ad6c9 100644
--- a/tensorflow/core/ops/ops.pbtxt
+++ b/tensorflow/core/ops/ops.pbtxt
@@ -1,4 +1,4 @@
-go/debugstr    
+go/debugproto   
 op {
   name: "Abort"
   attr {

From 3c92d1e861d76ce5a8ed10d0d36d03de5c260021 Mon Sep 17 00:00:00 2001
From: Deqiang Chen <deqiangc@google.com>
Date: Thu, 16 May 2024 07:22:23 -0700
Subject: [PATCH 268/478] Add a new test case in pjrt_util_test

PiperOrigin-RevId: 634381150
---
 tensorflow/core/tfrt/common/BUILD             |  2 ++
 tensorflow/core/tfrt/common/pjrt_util_test.cc | 18 ++++++++++++++++++
 2 files changed, 20 insertions(+)

diff --git a/tensorflow/core/tfrt/common/BUILD b/tensorflow/core/tfrt/common/BUILD
index bda3e3dde478e6..8129cdb0ea0f65 100644
--- a/tensorflow/core/tfrt/common/BUILD
+++ b/tensorflow/core/tfrt/common/BUILD
@@ -176,6 +176,8 @@ tf_cc_test(
         "//tensorflow/core:framework",
         "@local_tsl//tsl/lib/core:status_test_util",
         "@local_tsl//tsl/platform:status_matchers",
+        "@local_tsl//tsl/platform:statusor",
+        "@local_tsl//tsl/platform:test",
         "@local_tsl//tsl/platform:test_main",
         "@local_tsl//tsl/protobuf:error_codes_proto_impl_cc",
         "@local_xla//xla/pjrt:tfrt_cpu_pjrt_client",
diff --git a/tensorflow/core/tfrt/common/pjrt_util_test.cc b/tensorflow/core/tfrt/common/pjrt_util_test.cc
index 68cc2e98dd1a5d..f8de14dd034812 100644
--- a/tensorflow/core/tfrt/common/pjrt_util_test.cc
+++ b/tensorflow/core/tfrt/common/pjrt_util_test.cc
@@ -14,6 +14,9 @@ limitations under the License.
 ==============================================================================*/
 #include "tensorflow/core/tfrt/common/pjrt_util.h"
 
+#include <memory>
+#include <utility>
+
 #include "xla/pjrt/tfrt_cpu_pjrt_client.h"
 #include "tensorflow/core/framework/resource_mgr.h"
 #include "tensorflow/core/framework/types.h"
@@ -21,11 +24,14 @@ limitations under the License.
 #include "tensorflow/core/tfrt/common/pjrt_state.h"
 #include "tsl/lib/core/status_test_util.h"
 #include "tsl/platform/status_matchers.h"
+#include "tsl/platform/statusor.h"
+#include "tsl/platform/test.h"
 #include "tsl/protobuf/error_codes.pb.h"
 
 namespace tensorflow {
 namespace {
 
+using ::testing::ElementsAre;
 using ::testing::HasSubstr;
 using ::tsl::testing::StatusIs;
 
@@ -44,5 +50,17 @@ TEST(PjRtStateResourceManagerTest, SetNullPjRtClient) {
       StatusIs(error::INVALID_ARGUMENT, HasSubstr("PJRT client is nullptr")));
 }
 
+TEST(PjRtGpuClientCreationInfoTest, SetAndGet) {
+  auto info = std::make_unique<PjRtGpuClientCreationInfo>();
+  info->allowed_devices.insert(123);
+  TF_ASSERT_OK(
+      SetPjRtGpuClientCreationInfoInTFGlobalResourceManager(std::move(info)));
+
+  TF_ASSERT_OK_AND_ASSIGN(PjRtGpuClientCreationInfo * retrieved_info,
+                          GetPjRtGpuClientCreationInfo());
+
+  EXPECT_THAT(retrieved_info->allowed_devices, ElementsAre(123));
+}
+
 }  // namespace
 }  // namespace tensorflow

From 92b7d2cf2ebfafd7416966afaf64ba5c99e5ea79 Mon Sep 17 00:00:00 2001
From: Greg Olechwierowicz <olechwierowicz@google.com>
Date: Thu, 16 May 2024 07:39:23 -0700
Subject: [PATCH 269/478] [XLA:GPU][NFC] Print non-trivial reduction for
 all-reduce key.

PiperOrigin-RevId: 634387161
---
 third_party/xla/xla/service/BUILD             |  1 +
 third_party/xla/xla/service/all_reduce_key.cc | 11 ++++++++++-
 2 files changed, 11 insertions(+), 1 deletion(-)

diff --git a/third_party/xla/xla/service/BUILD b/third_party/xla/xla/service/BUILD
index 2cf56434539c13..ad05eb8777cd74 100644
--- a/third_party/xla/xla/service/BUILD
+++ b/third_party/xla/xla/service/BUILD
@@ -145,6 +145,7 @@ cc_library(
     deps = [
         ":hlo_domain_map",
         "//xla/hlo/ir:hlo",
+        "@com_google_absl//absl/log",
     ],
 )
 
diff --git a/third_party/xla/xla/service/all_reduce_key.cc b/third_party/xla/xla/service/all_reduce_key.cc
index ddc0ac79a41aa2..bd2fd49dc6be51 100644
--- a/third_party/xla/xla/service/all_reduce_key.cc
+++ b/third_party/xla/xla/service/all_reduce_key.cc
@@ -15,8 +15,16 @@ limitations under the License.
 
 #include "xla/service/all_reduce_key.h"
 
+#include <cstdint>
+#include <optional>
+#include <vector>
+
+#include "absl/log/log.h"
 #include "xla/hlo/ir/hlo_casting_utils.h"
+#include "xla/hlo/ir/hlo_instruction.h"
 #include "xla/hlo/ir/hlo_instructions.h"
+#include "xla/hlo/ir/hlo_opcode.h"
+#include "xla/service/hlo_domain_map.h"
 
 namespace xla {
 
@@ -32,7 +40,8 @@ std::optional<AllReduceKey> GetAllReduceKey(const HloInstruction* instruction,
 
   if (instruction->to_apply()->instruction_count() != 3 ||
       instruction->to_apply()->num_parameters() != 2) {
-    VLOG(1) << "Skipping due to non-trivial reduction function.";
+    VLOG(1) << "Skipping due to non-trivial reduction function: "
+            << instruction->to_apply()->ToString();
     return std::nullopt;
   }
 

From 73cf4c11b2f89b3a425860a60e989e6a3e4dc36a Mon Sep 17 00:00:00 2001
From: Oleg Shyshkov <shyshkov@google.com>
Date: Thu, 16 May 2024 07:39:25 -0700
Subject: [PATCH 270/478] [XLA:GPU] Introduce TiledHloComputation

TiledHloComputation is a wrapper around a vector of TiledHloInstructions. It provides a convenient way to access the instructions and the root instruction.

PiperOrigin-RevId: 634387167
---
 third_party/xla/xla/service/gpu/BUILD         |  1 +
 .../xla/xla/service/gpu/ir_emitter_triton.cc  | 20 +++---
 third_party/xla/xla/service/gpu/model/BUILD   | 12 ++++
 .../gpu/model/symbolic_tile_analysis.cc       |  6 +-
 .../gpu/model/symbolic_tile_analysis.h        |  6 +-
 .../gpu/model/symbolic_tile_analysis_test.cc  | 17 ++---
 .../service/gpu/model/tiled_hlo_computation.h | 71 +++++++++++++++++++
 7 files changed, 109 insertions(+), 24 deletions(-)
 create mode 100644 third_party/xla/xla/service/gpu/model/tiled_hlo_computation.h

diff --git a/third_party/xla/xla/service/gpu/BUILD b/third_party/xla/xla/service/gpu/BUILD
index a28626e06d3725..1c2712b6fd3538 100644
--- a/third_party/xla/xla/service/gpu/BUILD
+++ b/third_party/xla/xla/service/gpu/BUILD
@@ -537,6 +537,7 @@ cc_library(
         "//xla/service/gpu/model:indexing_map",
         "//xla/service/gpu/model:symbolic_tile_analysis",
         "//xla/service/gpu/model:symbolic_tiled_hlo_instruction",
+        "//xla/service/gpu/model:tiled_hlo_computation",
         "//xla/service/gpu/model:tiled_hlo_instruction",
         "//xla/service/llvm_ir:llvm_util",
         "//xla/stream_executor:device_description",
diff --git a/third_party/xla/xla/service/gpu/ir_emitter_triton.cc b/third_party/xla/xla/service/gpu/ir_emitter_triton.cc
index 20f290e27ff18a..3692c5e51d4207 100644
--- a/third_party/xla/xla/service/gpu/ir_emitter_triton.cc
+++ b/third_party/xla/xla/service/gpu/ir_emitter_triton.cc
@@ -116,6 +116,7 @@ limitations under the License.
 #include "xla/service/gpu/matmul_utils.h"
 #include "xla/service/gpu/model/indexing_map.h"
 #include "xla/service/gpu/model/symbolic_tile_analysis.h"
+#include "xla/service/gpu/model/tiled_hlo_computation.h"
 #include "xla/service/gpu/model/tiled_hlo_instruction.h"
 #include "xla/service/gpu/target_util.h"
 #include "xla/service/gpu/triton_fusion_analysis.h"
@@ -856,22 +857,22 @@ absl::StatusOr<Value> EmitTiledHloInstruction(
 absl::StatusOr<Value> EmitTiledScope(
     ImplicitLocOpBuilder& b, absl::string_view libdevice_path,
     const se::DeviceDescription& device_info,
-    const std::vector<std::unique_ptr<TiledHloInstruction>>&
-        tiled_hlo_instructions,
+    const TiledHloComputation& tiled_computation,
     std::function<absl::StatusOr<Value>(const TiledHloInstruction&)>
         emit_param_load_fn,
     absl::flat_hash_map<const TiledHloInstruction*, Value>& values) {
-  for (const auto& tiled_hlo : tiled_hlo_instructions) {
+  for (const TiledHloInstruction* tiled_hlo :
+       tiled_computation.instructions()) {
     TF_ASSIGN_OR_RETURN(
         Value result,
         EmitTiledHloInstruction(b, libdevice_path, device_info, *tiled_hlo,
                                 emit_param_load_fn, values));
-    TF_RET_CHECK(values.insert({tiled_hlo.get(), result}).second)
+    TF_RET_CHECK(values.insert({tiled_hlo, result}).second)
         << tiled_hlo->hlo()->ToString();
     VLOG(8) << "Emitted "
             << tiled_hlo->hlo()->ToString(HloPrintOptions::ShortParsable());
   }
-  return values[tiled_hlo_instructions.back().get()];
+  return values[tiled_computation.GetRoot()];
 }
 
 // Emit sequence of instructions using compatible tiling ordered producers
@@ -2323,9 +2324,8 @@ absl::Status EmitTiledSoftMax(mlir::OpBuilder builder,
       computation->root_instruction()->shape().rank(), 1);
   output_tile_sizes.back() = row_len;
 
-  TF_ASSIGN_OR_RETURN(
-      std::vector<std::unique_ptr<TiledHloInstruction>> tiled_hlo_instructions,
-      analysis->ComputeTiledHloInstructions(output_tile_sizes));
+  TF_ASSIGN_OR_RETURN(TiledHloComputation tiled_hlo_computation,
+                      analysis->ComputeTiledHloInstructions(output_tile_sizes));
 
   // block_size must be a power of two.
   int result_block_size = llvm::PowerOf2Ceil(row_len);
@@ -2373,11 +2373,11 @@ absl::Status EmitTiledSoftMax(mlir::OpBuilder builder,
   absl::flat_hash_map<const TiledHloInstruction*, Value> values_out;
   TF_ASSIGN_OR_RETURN(
       Value result,
-      EmitTiledScope(b, libdevice_path, device_info, tiled_hlo_instructions,
+      EmitTiledScope(b, libdevice_path, device_info, tiled_hlo_computation,
                      emit_param_load, values_out));
 
   Value ptr_offset =
-      ComputeBasePtrOffset(b, pid, *tiled_hlo_instructions.back());
+      ComputeBasePtrOffset(b, pid, *tiled_hlo_computation.GetRoot());
 
   Value store_tensor = b.create<mt::MakeTensorPtrOp>(
       /*base=*/AddPtr(b, fn.getArgument(computation->num_parameters()),
diff --git a/third_party/xla/xla/service/gpu/model/BUILD b/third_party/xla/xla/service/gpu/model/BUILD
index 82c40cf5b23c5f..ae876e97aa6832 100644
--- a/third_party/xla/xla/service/gpu/model/BUILD
+++ b/third_party/xla/xla/service/gpu/model/BUILD
@@ -624,6 +624,16 @@ xla_cc_test(
     ],
 )
 
+cc_library(
+    name = "tiled_hlo_computation",
+    hdrs = ["tiled_hlo_computation.h"],
+    deps = [
+        ":tiled_hlo_instruction",
+        "//xla:util",
+        "@local_tsl//tsl/lib/gtl:iterator_range",
+    ],
+)
+
 cc_library(
     name = "symbolic_tile_analysis",
     srcs = ["symbolic_tile_analysis.cc"],
@@ -633,6 +643,7 @@ cc_library(
         ":indexing_map",
         ":symbolic_tile",
         ":symbolic_tiled_hlo_instruction",
+        ":tiled_hlo_computation",
         ":tiled_hlo_instruction",
         "//xla:status",
         "//xla/hlo/ir:hlo",
@@ -658,6 +669,7 @@ xla_cc_test(
     deps = [
         ":indexing_test_utils",
         ":symbolic_tile_analysis",
+        ":tiled_hlo_computation",
         ":tiled_hlo_instruction",
         "//xla/hlo/ir:hlo",
         "//xla/tests:hlo_test_base",
diff --git a/third_party/xla/xla/service/gpu/model/symbolic_tile_analysis.cc b/third_party/xla/xla/service/gpu/model/symbolic_tile_analysis.cc
index 3b9e376c8b32e4..71d4c285ae1103 100644
--- a/third_party/xla/xla/service/gpu/model/symbolic_tile_analysis.cc
+++ b/third_party/xla/xla/service/gpu/model/symbolic_tile_analysis.cc
@@ -41,6 +41,7 @@ limitations under the License.
 #include "xla/service/gpu/model/indexing_map.h"
 #include "xla/service/gpu/model/symbolic_tile.h"
 #include "xla/service/gpu/model/symbolic_tiled_hlo_instruction.h"
+#include "xla/service/gpu/model/tiled_hlo_computation.h"
 #include "xla/service/gpu/model/tiled_hlo_instruction.h"
 #include "xla/service/instruction_fusion.h"
 #include "xla/status.h"
@@ -240,7 +241,7 @@ absl::StatusOr<IndexingMap> ComputeBlockIdToTileOffsetIndexing(
   return SymbolicTileAnalysis(std::move(tiled_hlo_instructions), ctx);
 }
 
-absl::StatusOr<std::vector<std::unique_ptr<TiledHloInstruction>>>
+absl::StatusOr<TiledHloComputation>
 SymbolicTileAnalysis::ComputeTiledHloInstructions(
     const std::vector<int64_t>& tile_parameters) const {
   IndexingMap block_id_to_root_tile_offset = ComputeBlockIdToOutputTileIndexing(
@@ -311,7 +312,8 @@ SymbolicTileAnalysis::ComputeTiledHloInstructions(
     return topological_order.at(i1.get()) < topological_order.at(i2.get());
   });
 
-  return tiled_hlo_instructions;
+  return TiledHloComputation::FromSortedTiledHloInstructions(
+      std::move(tiled_hlo_instructions));
 }
 
 }  // namespace gpu
diff --git a/third_party/xla/xla/service/gpu/model/symbolic_tile_analysis.h b/third_party/xla/xla/service/gpu/model/symbolic_tile_analysis.h
index 40b184d795f382..a880437356c4f0 100644
--- a/third_party/xla/xla/service/gpu/model/symbolic_tile_analysis.h
+++ b/third_party/xla/xla/service/gpu/model/symbolic_tile_analysis.h
@@ -26,7 +26,7 @@ limitations under the License.
 #include "mlir/IR/MLIRContext.h"  // from @llvm-project
 #include "xla/hlo/ir/hlo_instruction.h"
 #include "xla/service/gpu/model/symbolic_tiled_hlo_instruction.h"
-#include "xla/service/gpu/model/tiled_hlo_instruction.h"
+#include "xla/service/gpu/model/tiled_hlo_computation.h"
 #include "xla/service/instruction_fusion.h"
 
 namespace xla {
@@ -50,9 +50,7 @@ class SymbolicTileAnalysis {
       const HloComputation& computation, mlir::MLIRContext* ctx);
 
   // Returns a graph of HLO instructions tiled with the given tile parameters.
-  // Result vector has instructions in def-before-use order.
-  absl::StatusOr<std::vector<std::unique_ptr<TiledHloInstruction>>>
-  ComputeTiledHloInstructions(
+  absl::StatusOr<TiledHloComputation> ComputeTiledHloInstructions(
       const std::vector<int64_t>& tile_parameters) const;
 
   // Returns the tiled root instruction.
diff --git a/third_party/xla/xla/service/gpu/model/symbolic_tile_analysis_test.cc b/third_party/xla/xla/service/gpu/model/symbolic_tile_analysis_test.cc
index b6340a83636e33..4480481ae935af 100644
--- a/third_party/xla/xla/service/gpu/model/symbolic_tile_analysis_test.cc
+++ b/third_party/xla/xla/service/gpu/model/symbolic_tile_analysis_test.cc
@@ -27,6 +27,7 @@ limitations under the License.
 #include "mlir/IR/MLIRContext.h"  // from @llvm-project
 #include "xla/hlo/ir/hlo_module.h"
 #include "xla/service/gpu/model/indexing_test_utils.h"
+#include "xla/service/gpu/model/tiled_hlo_computation.h"
 #include "xla/service/gpu/model/tiled_hlo_instruction.h"
 #include "xla/tests/hlo_test_base.h"
 #include "xla/tests/verified_hlo_module.h"
@@ -76,10 +77,10 @@ ENTRY main {
   EXPECT_TRUE(SetAnalysis(module.get()));
 
   TF_ASSERT_OK_AND_ASSIGN(
-      std::vector<std::unique_ptr<TiledHloInstruction>> tiled_hlo_instructions,
+      TiledHloComputation tiled_hlo_computation,
       analysis_->ComputeTiledHloInstructions(/*tile_parameters=*/{1, 10}));
 
-  TiledHloInstruction* root = tiled_hlo_instructions.back().get();
+  const TiledHloInstruction* root = tiled_hlo_computation.GetRoot();
 
   EXPECT_THAT(root->block_id_to_tile_offsets_indexing(), MatchIndexingMap(R"(
     (d0) -> (d0 floordiv 10, (d0 mod 10) * 10)
@@ -124,10 +125,10 @@ ENTRY main {
   EXPECT_TRUE(SetAnalysis(module.get()));
 
   TF_ASSERT_OK_AND_ASSIGN(
-      std::vector<std::unique_ptr<TiledHloInstruction>> tiled_hlo_instructions,
+      TiledHloComputation tiled_hlo_computation,
       analysis_->ComputeTiledHloInstructions(/*tile_parameters=*/{1, 10}));
 
-  TiledHloInstruction* root = tiled_hlo_instructions.back().get();
+  const TiledHloInstruction* root = tiled_hlo_computation.GetRoot();
 
   auto p0_from_subtract0 = root->operand(0)->operand(0);
   auto p0_from_subtract1 = root->operand(1)->operand(0);
@@ -146,10 +147,10 @@ ENTRY main {
   EXPECT_TRUE(SetAnalysis(module.get()));
 
   TF_ASSERT_OK_AND_ASSIGN(
-      std::vector<std::unique_ptr<TiledHloInstruction>> tiled_hlo_instructions,
+      TiledHloComputation tiled_hlo_computation,
       analysis_->ComputeTiledHloInstructions(/*tile_parameters=*/{2, 4, 2}));
 
-  TiledHloInstruction* root = tiled_hlo_instructions.back().get();
+  const TiledHloInstruction* root = tiled_hlo_computation.GetRoot();
 
   EXPECT_THAT(root->block_id_to_tile_offsets_indexing(), MatchIndexingMap(R"(
     (d0) -> ((d0 floordiv 16) * 2, ((d0 floordiv 8) mod 2) * 4, (d0 mod 8) * 2)
@@ -178,10 +179,10 @@ ENTRY main {
   EXPECT_TRUE(SetAnalysis(module.get()));
 
   TF_ASSERT_OK_AND_ASSIGN(
-      std::vector<std::unique_ptr<TiledHloInstruction>> tiled_hlo_instructions,
+      TiledHloComputation tiled_hlo_computation,
       analysis_->ComputeTiledHloInstructions(/*tile_parameters=*/{2, 2}));
 
-  TiledHloInstruction* root = tiled_hlo_instructions.back().get();
+  const TiledHloInstruction* root = tiled_hlo_computation.GetRoot();
   const TiledHloInstruction* p0_from_slice0 = root->operand(0)->operand(0);
   const TiledHloInstruction* p0_from_slice1 = root->operand(1)->operand(0);
 
diff --git a/third_party/xla/xla/service/gpu/model/tiled_hlo_computation.h b/third_party/xla/xla/service/gpu/model/tiled_hlo_computation.h
new file mode 100644
index 00000000000000..d7fe5e71789311
--- /dev/null
+++ b/third_party/xla/xla/service/gpu/model/tiled_hlo_computation.h
@@ -0,0 +1,71 @@
+/* Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+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.
+==============================================================================*/
+
+#ifndef XLA_SERVICE_GPU_MODEL_TILED_HLO_COMPUTATION_H_
+#define XLA_SERVICE_GPU_MODEL_TILED_HLO_COMPUTATION_H_
+
+#include <memory>
+#include <utility>
+#include <vector>
+
+#include "xla/iterator_util.h"
+#include "xla/service/gpu/model/tiled_hlo_instruction.h"
+#include "tsl/lib/gtl/iterator_range.h"
+
+namespace xla {
+namespace gpu {
+
+// Stores TiledHloInstructions in the computation.
+//  * Instructions reference each other with non-owning pointers.
+//  * Instructions with the same tiling parameters are CSE-ed during
+//  construction.
+//  * Instructions are stored in def-before-use order.
+//  * The last element in the vector in the root instruction.
+class TiledHloComputation {
+ public:
+  // Creates a computation from a list of instructions. The instructions are
+  // expected to be sorted in def-before-use order.
+  static TiledHloComputation FromSortedTiledHloInstructions(
+      std::vector<std::unique_ptr<TiledHloInstruction>> instructions) {
+    return TiledHloComputation(std::move(instructions));
+  }
+
+  // Returns an iterator range over the instructions in the computation in
+  // def-before-use order.
+  tsl::gtl::iterator_range<UnwrappingIterator<
+      std::vector<std::unique_ptr<TiledHloInstruction>>::const_iterator>>
+  instructions() const {
+    return {MakeUnwrappingIterator(instructions_.begin()),
+            MakeUnwrappingIterator(instructions_.end())};
+  }
+
+  // Returns the root instruction of the computation.
+  const TiledHloInstruction* GetRoot() const {
+    return instructions_.back().get();
+  }
+
+ private:
+  explicit TiledHloComputation(
+      std::vector<std::unique_ptr<TiledHloInstruction>> instructions)
+      : instructions_(std::move(instructions)) {}
+
+  // Stores instructions in the computation in def-before-use order.
+  std::vector<std::unique_ptr<TiledHloInstruction>> instructions_;
+};
+
+}  // namespace gpu
+}  // namespace xla
+
+#endif  // XLA_SERVICE_GPU_MODEL_TILED_HLO_COMPUTATION_H_

From 23efbfea447b0147c4395e9d377992d977677a40 Mon Sep 17 00:00:00 2001
From: Benjamin Chetioui <bchetioui@google.com>
Date: Thu, 16 May 2024 07:39:25 -0700
Subject: [PATCH 271/478] [XLA:GPU][NFC] Remove unused includes in
 `symbolic_tile.h`.

PiperOrigin-RevId: 634387169
---
 third_party/xla/xla/service/gpu/model/symbolic_tile.h | 2 --
 1 file changed, 2 deletions(-)

diff --git a/third_party/xla/xla/service/gpu/model/symbolic_tile.h b/third_party/xla/xla/service/gpu/model/symbolic_tile.h
index 8a7225598de104..85059586ef5dd6 100644
--- a/third_party/xla/xla/service/gpu/model/symbolic_tile.h
+++ b/third_party/xla/xla/service/gpu/model/symbolic_tile.h
@@ -20,9 +20,7 @@ limitations under the License.
 #include <ostream>
 #include <string>
 #include <utility>
-#include <vector>
 
-#include "llvm/ADT/DenseMap.h"
 #include "mlir/IR/AffineExpr.h"  // from @llvm-project
 #include "mlir/IR/AffineMap.h"  // from @llvm-project
 #include "xla/service/gpu/model/affine_map_printer.h"

From 6bf19f6d7a660ca35a225b84a5c4212464f6fa18 Mon Sep 17 00:00:00 2001
From: Kyle Lucke <klucke@google.com>
Date: Thu, 16 May 2024 07:47:56 -0700
Subject: [PATCH 272/478] Use StreamExecutorInterface::CreateEvent to create
 stream_executor::Events.

PiperOrigin-RevId: 634390248
---
 tensorflow/core/common_runtime/device/device_event_mgr.cc | 3 +--
 tensorflow/core/tpu/kernels/tpu_execute_op.cc             | 5 ++---
 tensorflow/core/tpu/kernels/tpu_reshard_variables_op.cc   | 5 ++---
 3 files changed, 5 insertions(+), 8 deletions(-)

diff --git a/tensorflow/core/common_runtime/device/device_event_mgr.cc b/tensorflow/core/common_runtime/device/device_event_mgr.cc
index 1941cea9bc74f4..e786e5c9915bc9 100644
--- a/tensorflow/core/common_runtime/device/device_event_mgr.cc
+++ b/tensorflow/core/common_runtime/device/device_event_mgr.cc
@@ -172,8 +172,7 @@ void EventMgr::EnqueueCallback(se::Stream* stream, std::function<void()> func) {
   // Events are created on demand, and repeatedly reused.  There is no
   // limit placed here on the number of allocated Events.
   if (free_events_.empty()) {
-    free_events_.push_back(std::make_unique<se::Event>(exec_));
-    free_events_.back()->Init();
+    free_events_.emplace_back(exec_->CreateEvent().value());
   }
 
   std::unique_ptr<se::Event> e = std::move(free_events_.back());
diff --git a/tensorflow/core/tpu/kernels/tpu_execute_op.cc b/tensorflow/core/tpu/kernels/tpu_execute_op.cc
index 02654c0d6a2cc6..7c23a7558eea05 100644
--- a/tensorflow/core/tpu/kernels/tpu_execute_op.cc
+++ b/tensorflow/core/tpu/kernels/tpu_execute_op.cc
@@ -739,9 +739,8 @@ Status TPUExecuteOp::DoWork(OpKernelContext* context) {
         });
   }
 
-  auto definition_event = std::make_shared<se::Event>(stream->parent());
-  TF_RET_CHECK(definition_event->Init())
-      << "TPU definition event initialization failed";
+  TF_ASSIGN_OR_RETURN(std::shared_ptr<se::Event> definition_event,
+                      stream->parent()->CreateEvent());
 
   trace_me_init.Stop();
 
diff --git a/tensorflow/core/tpu/kernels/tpu_reshard_variables_op.cc b/tensorflow/core/tpu/kernels/tpu_reshard_variables_op.cc
index 96ae248347474f..e311a802a158f3 100644
--- a/tensorflow/core/tpu/kernels/tpu_reshard_variables_op.cc
+++ b/tensorflow/core/tpu/kernels/tpu_reshard_variables_op.cc
@@ -216,9 +216,8 @@ Status TPUReshardVariablesOpKernel::DoTpuExecute(
   TF_RET_CHECK(!executable->has_session_module())
       << "session module not supported in sharding/unsharding program.";
 
-  auto definition_event = std::make_shared<se::Event>(stream->parent());
-  TF_RET_CHECK(definition_event->Init())
-      << "TPU definition event initialization failed";
+  TF_ASSIGN_OR_RETURN(std::shared_ptr<se::Event> definition_event,
+                      stream->parent()->CreateEvent());
 
   trace_me_init.Stop();
 

From e8dcadf49bbb0cb8a77688bb778884cf07ca44a0 Mon Sep 17 00:00:00 2001
From: Dmitri Gribenko <dmitrig@google.com>
Date: Thu, 16 May 2024 08:11:03 -0700
Subject: [PATCH 273/478] Integrate LLVM at llvm/llvm-project@e27f9bb31984

Updates LLVM usage to match
[e27f9bb31984](https://github.com/llvm/llvm-project/commit/e27f9bb31984)

PiperOrigin-RevId: 634398955
---
 third_party/llvm/generated.patch | 836 +++++++++++++++++++++++++++++++
 third_party/llvm/workspace.bzl   |   4 +-
 2 files changed, 838 insertions(+), 2 deletions(-)

diff --git a/third_party/llvm/generated.patch b/third_party/llvm/generated.patch
index 509398da979e83..9ed9c58a9c4eb8 100644
--- a/third_party/llvm/generated.patch
+++ b/third_party/llvm/generated.patch
@@ -1 +1,837 @@
 Auto generated patch. Do not edit or delete it, even if empty.
+diff -ruN --strip-trailing-cr a/llvm/include/llvm/CodeGen/BasicTTIImpl.h b/llvm/include/llvm/CodeGen/BasicTTIImpl.h
+--- a/llvm/include/llvm/CodeGen/BasicTTIImpl.h
++++ b/llvm/include/llvm/CodeGen/BasicTTIImpl.h
+@@ -612,13 +612,7 @@
+     if (PartialUnrollingThreshold.getNumOccurrences() > 0)
+       MaxOps = PartialUnrollingThreshold;
+     else if (ST->getSchedModel().LoopMicroOpBufferSize > 0)
+-      // Upper bound by the default PartialThreshold, which is the same as
+-      // the default full-unroll Threshold. Even if the loop micro-op buffer
+-      // is very large, this does not mean that we want to unroll all loops
+-      // to that length, as it would increase code size beyond the limits of
+-      // what unrolling normally allows.
+-      MaxOps = std::min(ST->getSchedModel().LoopMicroOpBufferSize,
+-                        UP.PartialThreshold);
++      MaxOps = ST->getSchedModel().LoopMicroOpBufferSize;
+     else
+       return;
+ 
+diff -ruN --strip-trailing-cr a/llvm/test/Transforms/LoopUnroll/X86/znver3.ll b/llvm/test/Transforms/LoopUnroll/X86/znver3.ll
+--- a/llvm/test/Transforms/LoopUnroll/X86/znver3.ll
++++ b/llvm/test/Transforms/LoopUnroll/X86/znver3.ll
+@@ -9,8 +9,8 @@
+ ; CHECK-NEXT:  entry:
+ ; CHECK-NEXT:    br label [[FOR_BODY:%.*]]
+ ; CHECK:       for.body:
+-; CHECK-NEXT:    [[INDVARS_IV:%.*]] = phi i64 [ 0, [[ENTRY:%.*]] ], [ [[INDVARS_IV_NEXT_31:%.*]], [[FOR_BODY]] ]
+-; CHECK-NEXT:    [[SUM:%.*]] = phi i32 [ 0, [[ENTRY]] ], [ [[SUM_NEXT_31:%.*]], [[FOR_BODY]] ]
++; CHECK-NEXT:    [[INDVARS_IV:%.*]] = phi i64 [ 0, [[ENTRY:%.*]] ], [ [[INDVARS_IV_NEXT_127:%.*]], [[FOR_BODY]] ]
++; CHECK-NEXT:    [[SUM:%.*]] = phi i32 [ 0, [[ENTRY]] ], [ [[SUM_NEXT_127:%.*]], [[FOR_BODY]] ]
+ ; CHECK-NEXT:    [[ARRAYIDX:%.*]] = getelementptr inbounds i32, ptr [[ARY]], i64 [[INDVARS_IV]]
+ ; CHECK-NEXT:    [[VAL:%.*]] = load i32, ptr [[ARRAYIDX]], align 4
+ ; CHECK-NEXT:    [[SUM_NEXT:%.*]] = add nsw i32 [[VAL]], [[SUM]]
+@@ -137,12 +137,396 @@
+ ; CHECK-NEXT:    [[INDVARS_IV_NEXT_30:%.*]] = add nuw nsw i64 [[INDVARS_IV]], 31
+ ; CHECK-NEXT:    [[ARRAYIDX_31:%.*]] = getelementptr inbounds i32, ptr [[ARY]], i64 [[INDVARS_IV_NEXT_30]]
+ ; CHECK-NEXT:    [[VAL_31:%.*]] = load i32, ptr [[ARRAYIDX_31]], align 4
+-; CHECK-NEXT:    [[SUM_NEXT_31]] = add nsw i32 [[VAL_31]], [[SUM_NEXT_30]]
+-; CHECK-NEXT:    [[INDVARS_IV_NEXT_31]] = add nuw nsw i64 [[INDVARS_IV]], 32
+-; CHECK-NEXT:    [[EXITCOND_NOT_31:%.*]] = icmp eq i64 [[INDVARS_IV_NEXT_31]], 8192
+-; CHECK-NEXT:    br i1 [[EXITCOND_NOT_31]], label [[FOR_COND_CLEANUP:%.*]], label [[FOR_BODY]]
++; CHECK-NEXT:    [[SUM_NEXT_31:%.*]] = add nsw i32 [[VAL_31]], [[SUM_NEXT_30]]
++; CHECK-NEXT:    [[INDVARS_IV_NEXT_31:%.*]] = add nuw nsw i64 [[INDVARS_IV]], 32
++; CHECK-NEXT:    [[ARRAYIDX_32:%.*]] = getelementptr inbounds i32, ptr [[ARY]], i64 [[INDVARS_IV_NEXT_31]]
++; CHECK-NEXT:    [[VAL_32:%.*]] = load i32, ptr [[ARRAYIDX_32]], align 4
++; CHECK-NEXT:    [[SUM_NEXT_32:%.*]] = add nsw i32 [[VAL_32]], [[SUM_NEXT_31]]
++; CHECK-NEXT:    [[INDVARS_IV_NEXT_32:%.*]] = add nuw nsw i64 [[INDVARS_IV]], 33
++; CHECK-NEXT:    [[ARRAYIDX_33:%.*]] = getelementptr inbounds i32, ptr [[ARY]], i64 [[INDVARS_IV_NEXT_32]]
++; CHECK-NEXT:    [[VAL_33:%.*]] = load i32, ptr [[ARRAYIDX_33]], align 4
++; CHECK-NEXT:    [[SUM_NEXT_33:%.*]] = add nsw i32 [[VAL_33]], [[SUM_NEXT_32]]
++; CHECK-NEXT:    [[INDVARS_IV_NEXT_33:%.*]] = add nuw nsw i64 [[INDVARS_IV]], 34
++; CHECK-NEXT:    [[ARRAYIDX_34:%.*]] = getelementptr inbounds i32, ptr [[ARY]], i64 [[INDVARS_IV_NEXT_33]]
++; CHECK-NEXT:    [[VAL_34:%.*]] = load i32, ptr [[ARRAYIDX_34]], align 4
++; CHECK-NEXT:    [[SUM_NEXT_34:%.*]] = add nsw i32 [[VAL_34]], [[SUM_NEXT_33]]
++; CHECK-NEXT:    [[INDVARS_IV_NEXT_34:%.*]] = add nuw nsw i64 [[INDVARS_IV]], 35
++; CHECK-NEXT:    [[ARRAYIDX_35:%.*]] = getelementptr inbounds i32, ptr [[ARY]], i64 [[INDVARS_IV_NEXT_34]]
++; CHECK-NEXT:    [[VAL_35:%.*]] = load i32, ptr [[ARRAYIDX_35]], align 4
++; CHECK-NEXT:    [[SUM_NEXT_35:%.*]] = add nsw i32 [[VAL_35]], [[SUM_NEXT_34]]
++; CHECK-NEXT:    [[INDVARS_IV_NEXT_35:%.*]] = add nuw nsw i64 [[INDVARS_IV]], 36
++; CHECK-NEXT:    [[ARRAYIDX_36:%.*]] = getelementptr inbounds i32, ptr [[ARY]], i64 [[INDVARS_IV_NEXT_35]]
++; CHECK-NEXT:    [[VAL_36:%.*]] = load i32, ptr [[ARRAYIDX_36]], align 4
++; CHECK-NEXT:    [[SUM_NEXT_36:%.*]] = add nsw i32 [[VAL_36]], [[SUM_NEXT_35]]
++; CHECK-NEXT:    [[INDVARS_IV_NEXT_36:%.*]] = add nuw nsw i64 [[INDVARS_IV]], 37
++; CHECK-NEXT:    [[ARRAYIDX_37:%.*]] = getelementptr inbounds i32, ptr [[ARY]], i64 [[INDVARS_IV_NEXT_36]]
++; CHECK-NEXT:    [[VAL_37:%.*]] = load i32, ptr [[ARRAYIDX_37]], align 4
++; CHECK-NEXT:    [[SUM_NEXT_37:%.*]] = add nsw i32 [[VAL_37]], [[SUM_NEXT_36]]
++; CHECK-NEXT:    [[INDVARS_IV_NEXT_37:%.*]] = add nuw nsw i64 [[INDVARS_IV]], 38
++; CHECK-NEXT:    [[ARRAYIDX_38:%.*]] = getelementptr inbounds i32, ptr [[ARY]], i64 [[INDVARS_IV_NEXT_37]]
++; CHECK-NEXT:    [[VAL_38:%.*]] = load i32, ptr [[ARRAYIDX_38]], align 4
++; CHECK-NEXT:    [[SUM_NEXT_38:%.*]] = add nsw i32 [[VAL_38]], [[SUM_NEXT_37]]
++; CHECK-NEXT:    [[INDVARS_IV_NEXT_38:%.*]] = add nuw nsw i64 [[INDVARS_IV]], 39
++; CHECK-NEXT:    [[ARRAYIDX_39:%.*]] = getelementptr inbounds i32, ptr [[ARY]], i64 [[INDVARS_IV_NEXT_38]]
++; CHECK-NEXT:    [[VAL_39:%.*]] = load i32, ptr [[ARRAYIDX_39]], align 4
++; CHECK-NEXT:    [[SUM_NEXT_39:%.*]] = add nsw i32 [[VAL_39]], [[SUM_NEXT_38]]
++; CHECK-NEXT:    [[INDVARS_IV_NEXT_39:%.*]] = add nuw nsw i64 [[INDVARS_IV]], 40
++; CHECK-NEXT:    [[ARRAYIDX_40:%.*]] = getelementptr inbounds i32, ptr [[ARY]], i64 [[INDVARS_IV_NEXT_39]]
++; CHECK-NEXT:    [[VAL_40:%.*]] = load i32, ptr [[ARRAYIDX_40]], align 4
++; CHECK-NEXT:    [[SUM_NEXT_40:%.*]] = add nsw i32 [[VAL_40]], [[SUM_NEXT_39]]
++; CHECK-NEXT:    [[INDVARS_IV_NEXT_40:%.*]] = add nuw nsw i64 [[INDVARS_IV]], 41
++; CHECK-NEXT:    [[ARRAYIDX_41:%.*]] = getelementptr inbounds i32, ptr [[ARY]], i64 [[INDVARS_IV_NEXT_40]]
++; CHECK-NEXT:    [[VAL_41:%.*]] = load i32, ptr [[ARRAYIDX_41]], align 4
++; CHECK-NEXT:    [[SUM_NEXT_41:%.*]] = add nsw i32 [[VAL_41]], [[SUM_NEXT_40]]
++; CHECK-NEXT:    [[INDVARS_IV_NEXT_41:%.*]] = add nuw nsw i64 [[INDVARS_IV]], 42
++; CHECK-NEXT:    [[ARRAYIDX_42:%.*]] = getelementptr inbounds i32, ptr [[ARY]], i64 [[INDVARS_IV_NEXT_41]]
++; CHECK-NEXT:    [[VAL_42:%.*]] = load i32, ptr [[ARRAYIDX_42]], align 4
++; CHECK-NEXT:    [[SUM_NEXT_42:%.*]] = add nsw i32 [[VAL_42]], [[SUM_NEXT_41]]
++; CHECK-NEXT:    [[INDVARS_IV_NEXT_42:%.*]] = add nuw nsw i64 [[INDVARS_IV]], 43
++; CHECK-NEXT:    [[ARRAYIDX_43:%.*]] = getelementptr inbounds i32, ptr [[ARY]], i64 [[INDVARS_IV_NEXT_42]]
++; CHECK-NEXT:    [[VAL_43:%.*]] = load i32, ptr [[ARRAYIDX_43]], align 4
++; CHECK-NEXT:    [[SUM_NEXT_43:%.*]] = add nsw i32 [[VAL_43]], [[SUM_NEXT_42]]
++; CHECK-NEXT:    [[INDVARS_IV_NEXT_43:%.*]] = add nuw nsw i64 [[INDVARS_IV]], 44
++; CHECK-NEXT:    [[ARRAYIDX_44:%.*]] = getelementptr inbounds i32, ptr [[ARY]], i64 [[INDVARS_IV_NEXT_43]]
++; CHECK-NEXT:    [[VAL_44:%.*]] = load i32, ptr [[ARRAYIDX_44]], align 4
++; CHECK-NEXT:    [[SUM_NEXT_44:%.*]] = add nsw i32 [[VAL_44]], [[SUM_NEXT_43]]
++; CHECK-NEXT:    [[INDVARS_IV_NEXT_44:%.*]] = add nuw nsw i64 [[INDVARS_IV]], 45
++; CHECK-NEXT:    [[ARRAYIDX_45:%.*]] = getelementptr inbounds i32, ptr [[ARY]], i64 [[INDVARS_IV_NEXT_44]]
++; CHECK-NEXT:    [[VAL_45:%.*]] = load i32, ptr [[ARRAYIDX_45]], align 4
++; CHECK-NEXT:    [[SUM_NEXT_45:%.*]] = add nsw i32 [[VAL_45]], [[SUM_NEXT_44]]
++; CHECK-NEXT:    [[INDVARS_IV_NEXT_45:%.*]] = add nuw nsw i64 [[INDVARS_IV]], 46
++; CHECK-NEXT:    [[ARRAYIDX_46:%.*]] = getelementptr inbounds i32, ptr [[ARY]], i64 [[INDVARS_IV_NEXT_45]]
++; CHECK-NEXT:    [[VAL_46:%.*]] = load i32, ptr [[ARRAYIDX_46]], align 4
++; CHECK-NEXT:    [[SUM_NEXT_46:%.*]] = add nsw i32 [[VAL_46]], [[SUM_NEXT_45]]
++; CHECK-NEXT:    [[INDVARS_IV_NEXT_46:%.*]] = add nuw nsw i64 [[INDVARS_IV]], 47
++; CHECK-NEXT:    [[ARRAYIDX_47:%.*]] = getelementptr inbounds i32, ptr [[ARY]], i64 [[INDVARS_IV_NEXT_46]]
++; CHECK-NEXT:    [[VAL_47:%.*]] = load i32, ptr [[ARRAYIDX_47]], align 4
++; CHECK-NEXT:    [[SUM_NEXT_47:%.*]] = add nsw i32 [[VAL_47]], [[SUM_NEXT_46]]
++; CHECK-NEXT:    [[INDVARS_IV_NEXT_47:%.*]] = add nuw nsw i64 [[INDVARS_IV]], 48
++; CHECK-NEXT:    [[ARRAYIDX_48:%.*]] = getelementptr inbounds i32, ptr [[ARY]], i64 [[INDVARS_IV_NEXT_47]]
++; CHECK-NEXT:    [[VAL_48:%.*]] = load i32, ptr [[ARRAYIDX_48]], align 4
++; CHECK-NEXT:    [[SUM_NEXT_48:%.*]] = add nsw i32 [[VAL_48]], [[SUM_NEXT_47]]
++; CHECK-NEXT:    [[INDVARS_IV_NEXT_48:%.*]] = add nuw nsw i64 [[INDVARS_IV]], 49
++; CHECK-NEXT:    [[ARRAYIDX_49:%.*]] = getelementptr inbounds i32, ptr [[ARY]], i64 [[INDVARS_IV_NEXT_48]]
++; CHECK-NEXT:    [[VAL_49:%.*]] = load i32, ptr [[ARRAYIDX_49]], align 4
++; CHECK-NEXT:    [[SUM_NEXT_49:%.*]] = add nsw i32 [[VAL_49]], [[SUM_NEXT_48]]
++; CHECK-NEXT:    [[INDVARS_IV_NEXT_49:%.*]] = add nuw nsw i64 [[INDVARS_IV]], 50
++; CHECK-NEXT:    [[ARRAYIDX_50:%.*]] = getelementptr inbounds i32, ptr [[ARY]], i64 [[INDVARS_IV_NEXT_49]]
++; CHECK-NEXT:    [[VAL_50:%.*]] = load i32, ptr [[ARRAYIDX_50]], align 4
++; CHECK-NEXT:    [[SUM_NEXT_50:%.*]] = add nsw i32 [[VAL_50]], [[SUM_NEXT_49]]
++; CHECK-NEXT:    [[INDVARS_IV_NEXT_50:%.*]] = add nuw nsw i64 [[INDVARS_IV]], 51
++; CHECK-NEXT:    [[ARRAYIDX_51:%.*]] = getelementptr inbounds i32, ptr [[ARY]], i64 [[INDVARS_IV_NEXT_50]]
++; CHECK-NEXT:    [[VAL_51:%.*]] = load i32, ptr [[ARRAYIDX_51]], align 4
++; CHECK-NEXT:    [[SUM_NEXT_51:%.*]] = add nsw i32 [[VAL_51]], [[SUM_NEXT_50]]
++; CHECK-NEXT:    [[INDVARS_IV_NEXT_51:%.*]] = add nuw nsw i64 [[INDVARS_IV]], 52
++; CHECK-NEXT:    [[ARRAYIDX_52:%.*]] = getelementptr inbounds i32, ptr [[ARY]], i64 [[INDVARS_IV_NEXT_51]]
++; CHECK-NEXT:    [[VAL_52:%.*]] = load i32, ptr [[ARRAYIDX_52]], align 4
++; CHECK-NEXT:    [[SUM_NEXT_52:%.*]] = add nsw i32 [[VAL_52]], [[SUM_NEXT_51]]
++; CHECK-NEXT:    [[INDVARS_IV_NEXT_52:%.*]] = add nuw nsw i64 [[INDVARS_IV]], 53
++; CHECK-NEXT:    [[ARRAYIDX_53:%.*]] = getelementptr inbounds i32, ptr [[ARY]], i64 [[INDVARS_IV_NEXT_52]]
++; CHECK-NEXT:    [[VAL_53:%.*]] = load i32, ptr [[ARRAYIDX_53]], align 4
++; CHECK-NEXT:    [[SUM_NEXT_53:%.*]] = add nsw i32 [[VAL_53]], [[SUM_NEXT_52]]
++; CHECK-NEXT:    [[INDVARS_IV_NEXT_53:%.*]] = add nuw nsw i64 [[INDVARS_IV]], 54
++; CHECK-NEXT:    [[ARRAYIDX_54:%.*]] = getelementptr inbounds i32, ptr [[ARY]], i64 [[INDVARS_IV_NEXT_53]]
++; CHECK-NEXT:    [[VAL_54:%.*]] = load i32, ptr [[ARRAYIDX_54]], align 4
++; CHECK-NEXT:    [[SUM_NEXT_54:%.*]] = add nsw i32 [[VAL_54]], [[SUM_NEXT_53]]
++; CHECK-NEXT:    [[INDVARS_IV_NEXT_54:%.*]] = add nuw nsw i64 [[INDVARS_IV]], 55
++; CHECK-NEXT:    [[ARRAYIDX_55:%.*]] = getelementptr inbounds i32, ptr [[ARY]], i64 [[INDVARS_IV_NEXT_54]]
++; CHECK-NEXT:    [[VAL_55:%.*]] = load i32, ptr [[ARRAYIDX_55]], align 4
++; CHECK-NEXT:    [[SUM_NEXT_55:%.*]] = add nsw i32 [[VAL_55]], [[SUM_NEXT_54]]
++; CHECK-NEXT:    [[INDVARS_IV_NEXT_55:%.*]] = add nuw nsw i64 [[INDVARS_IV]], 56
++; CHECK-NEXT:    [[ARRAYIDX_56:%.*]] = getelementptr inbounds i32, ptr [[ARY]], i64 [[INDVARS_IV_NEXT_55]]
++; CHECK-NEXT:    [[VAL_56:%.*]] = load i32, ptr [[ARRAYIDX_56]], align 4
++; CHECK-NEXT:    [[SUM_NEXT_56:%.*]] = add nsw i32 [[VAL_56]], [[SUM_NEXT_55]]
++; CHECK-NEXT:    [[INDVARS_IV_NEXT_56:%.*]] = add nuw nsw i64 [[INDVARS_IV]], 57
++; CHECK-NEXT:    [[ARRAYIDX_57:%.*]] = getelementptr inbounds i32, ptr [[ARY]], i64 [[INDVARS_IV_NEXT_56]]
++; CHECK-NEXT:    [[VAL_57:%.*]] = load i32, ptr [[ARRAYIDX_57]], align 4
++; CHECK-NEXT:    [[SUM_NEXT_57:%.*]] = add nsw i32 [[VAL_57]], [[SUM_NEXT_56]]
++; CHECK-NEXT:    [[INDVARS_IV_NEXT_57:%.*]] = add nuw nsw i64 [[INDVARS_IV]], 58
++; CHECK-NEXT:    [[ARRAYIDX_58:%.*]] = getelementptr inbounds i32, ptr [[ARY]], i64 [[INDVARS_IV_NEXT_57]]
++; CHECK-NEXT:    [[VAL_58:%.*]] = load i32, ptr [[ARRAYIDX_58]], align 4
++; CHECK-NEXT:    [[SUM_NEXT_58:%.*]] = add nsw i32 [[VAL_58]], [[SUM_NEXT_57]]
++; CHECK-NEXT:    [[INDVARS_IV_NEXT_58:%.*]] = add nuw nsw i64 [[INDVARS_IV]], 59
++; CHECK-NEXT:    [[ARRAYIDX_59:%.*]] = getelementptr inbounds i32, ptr [[ARY]], i64 [[INDVARS_IV_NEXT_58]]
++; CHECK-NEXT:    [[VAL_59:%.*]] = load i32, ptr [[ARRAYIDX_59]], align 4
++; CHECK-NEXT:    [[SUM_NEXT_59:%.*]] = add nsw i32 [[VAL_59]], [[SUM_NEXT_58]]
++; CHECK-NEXT:    [[INDVARS_IV_NEXT_59:%.*]] = add nuw nsw i64 [[INDVARS_IV]], 60
++; CHECK-NEXT:    [[ARRAYIDX_60:%.*]] = getelementptr inbounds i32, ptr [[ARY]], i64 [[INDVARS_IV_NEXT_59]]
++; CHECK-NEXT:    [[VAL_60:%.*]] = load i32, ptr [[ARRAYIDX_60]], align 4
++; CHECK-NEXT:    [[SUM_NEXT_60:%.*]] = add nsw i32 [[VAL_60]], [[SUM_NEXT_59]]
++; CHECK-NEXT:    [[INDVARS_IV_NEXT_60:%.*]] = add nuw nsw i64 [[INDVARS_IV]], 61
++; CHECK-NEXT:    [[ARRAYIDX_61:%.*]] = getelementptr inbounds i32, ptr [[ARY]], i64 [[INDVARS_IV_NEXT_60]]
++; CHECK-NEXT:    [[VAL_61:%.*]] = load i32, ptr [[ARRAYIDX_61]], align 4
++; CHECK-NEXT:    [[SUM_NEXT_61:%.*]] = add nsw i32 [[VAL_61]], [[SUM_NEXT_60]]
++; CHECK-NEXT:    [[INDVARS_IV_NEXT_61:%.*]] = add nuw nsw i64 [[INDVARS_IV]], 62
++; CHECK-NEXT:    [[ARRAYIDX_62:%.*]] = getelementptr inbounds i32, ptr [[ARY]], i64 [[INDVARS_IV_NEXT_61]]
++; CHECK-NEXT:    [[VAL_62:%.*]] = load i32, ptr [[ARRAYIDX_62]], align 4
++; CHECK-NEXT:    [[SUM_NEXT_62:%.*]] = add nsw i32 [[VAL_62]], [[SUM_NEXT_61]]
++; CHECK-NEXT:    [[INDVARS_IV_NEXT_62:%.*]] = add nuw nsw i64 [[INDVARS_IV]], 63
++; CHECK-NEXT:    [[ARRAYIDX_63:%.*]] = getelementptr inbounds i32, ptr [[ARY]], i64 [[INDVARS_IV_NEXT_62]]
++; CHECK-NEXT:    [[VAL_63:%.*]] = load i32, ptr [[ARRAYIDX_63]], align 4
++; CHECK-NEXT:    [[SUM_NEXT_63:%.*]] = add nsw i32 [[VAL_63]], [[SUM_NEXT_62]]
++; CHECK-NEXT:    [[INDVARS_IV_NEXT_63:%.*]] = add nuw nsw i64 [[INDVARS_IV]], 64
++; CHECK-NEXT:    [[ARRAYIDX_64:%.*]] = getelementptr inbounds i32, ptr [[ARY]], i64 [[INDVARS_IV_NEXT_63]]
++; CHECK-NEXT:    [[VAL_64:%.*]] = load i32, ptr [[ARRAYIDX_64]], align 4
++; CHECK-NEXT:    [[SUM_NEXT_64:%.*]] = add nsw i32 [[VAL_64]], [[SUM_NEXT_63]]
++; CHECK-NEXT:    [[INDVARS_IV_NEXT_64:%.*]] = add nuw nsw i64 [[INDVARS_IV]], 65
++; CHECK-NEXT:    [[ARRAYIDX_65:%.*]] = getelementptr inbounds i32, ptr [[ARY]], i64 [[INDVARS_IV_NEXT_64]]
++; CHECK-NEXT:    [[VAL_65:%.*]] = load i32, ptr [[ARRAYIDX_65]], align 4
++; CHECK-NEXT:    [[SUM_NEXT_65:%.*]] = add nsw i32 [[VAL_65]], [[SUM_NEXT_64]]
++; CHECK-NEXT:    [[INDVARS_IV_NEXT_65:%.*]] = add nuw nsw i64 [[INDVARS_IV]], 66
++; CHECK-NEXT:    [[ARRAYIDX_66:%.*]] = getelementptr inbounds i32, ptr [[ARY]], i64 [[INDVARS_IV_NEXT_65]]
++; CHECK-NEXT:    [[VAL_66:%.*]] = load i32, ptr [[ARRAYIDX_66]], align 4
++; CHECK-NEXT:    [[SUM_NEXT_66:%.*]] = add nsw i32 [[VAL_66]], [[SUM_NEXT_65]]
++; CHECK-NEXT:    [[INDVARS_IV_NEXT_66:%.*]] = add nuw nsw i64 [[INDVARS_IV]], 67
++; CHECK-NEXT:    [[ARRAYIDX_67:%.*]] = getelementptr inbounds i32, ptr [[ARY]], i64 [[INDVARS_IV_NEXT_66]]
++; CHECK-NEXT:    [[VAL_67:%.*]] = load i32, ptr [[ARRAYIDX_67]], align 4
++; CHECK-NEXT:    [[SUM_NEXT_67:%.*]] = add nsw i32 [[VAL_67]], [[SUM_NEXT_66]]
++; CHECK-NEXT:    [[INDVARS_IV_NEXT_67:%.*]] = add nuw nsw i64 [[INDVARS_IV]], 68
++; CHECK-NEXT:    [[ARRAYIDX_68:%.*]] = getelementptr inbounds i32, ptr [[ARY]], i64 [[INDVARS_IV_NEXT_67]]
++; CHECK-NEXT:    [[VAL_68:%.*]] = load i32, ptr [[ARRAYIDX_68]], align 4
++; CHECK-NEXT:    [[SUM_NEXT_68:%.*]] = add nsw i32 [[VAL_68]], [[SUM_NEXT_67]]
++; CHECK-NEXT:    [[INDVARS_IV_NEXT_68:%.*]] = add nuw nsw i64 [[INDVARS_IV]], 69
++; CHECK-NEXT:    [[ARRAYIDX_69:%.*]] = getelementptr inbounds i32, ptr [[ARY]], i64 [[INDVARS_IV_NEXT_68]]
++; CHECK-NEXT:    [[VAL_69:%.*]] = load i32, ptr [[ARRAYIDX_69]], align 4
++; CHECK-NEXT:    [[SUM_NEXT_69:%.*]] = add nsw i32 [[VAL_69]], [[SUM_NEXT_68]]
++; CHECK-NEXT:    [[INDVARS_IV_NEXT_69:%.*]] = add nuw nsw i64 [[INDVARS_IV]], 70
++; CHECK-NEXT:    [[ARRAYIDX_70:%.*]] = getelementptr inbounds i32, ptr [[ARY]], i64 [[INDVARS_IV_NEXT_69]]
++; CHECK-NEXT:    [[VAL_70:%.*]] = load i32, ptr [[ARRAYIDX_70]], align 4
++; CHECK-NEXT:    [[SUM_NEXT_70:%.*]] = add nsw i32 [[VAL_70]], [[SUM_NEXT_69]]
++; CHECK-NEXT:    [[INDVARS_IV_NEXT_70:%.*]] = add nuw nsw i64 [[INDVARS_IV]], 71
++; CHECK-NEXT:    [[ARRAYIDX_71:%.*]] = getelementptr inbounds i32, ptr [[ARY]], i64 [[INDVARS_IV_NEXT_70]]
++; CHECK-NEXT:    [[VAL_71:%.*]] = load i32, ptr [[ARRAYIDX_71]], align 4
++; CHECK-NEXT:    [[SUM_NEXT_71:%.*]] = add nsw i32 [[VAL_71]], [[SUM_NEXT_70]]
++; CHECK-NEXT:    [[INDVARS_IV_NEXT_71:%.*]] = add nuw nsw i64 [[INDVARS_IV]], 72
++; CHECK-NEXT:    [[ARRAYIDX_72:%.*]] = getelementptr inbounds i32, ptr [[ARY]], i64 [[INDVARS_IV_NEXT_71]]
++; CHECK-NEXT:    [[VAL_72:%.*]] = load i32, ptr [[ARRAYIDX_72]], align 4
++; CHECK-NEXT:    [[SUM_NEXT_72:%.*]] = add nsw i32 [[VAL_72]], [[SUM_NEXT_71]]
++; CHECK-NEXT:    [[INDVARS_IV_NEXT_72:%.*]] = add nuw nsw i64 [[INDVARS_IV]], 73
++; CHECK-NEXT:    [[ARRAYIDX_73:%.*]] = getelementptr inbounds i32, ptr [[ARY]], i64 [[INDVARS_IV_NEXT_72]]
++; CHECK-NEXT:    [[VAL_73:%.*]] = load i32, ptr [[ARRAYIDX_73]], align 4
++; CHECK-NEXT:    [[SUM_NEXT_73:%.*]] = add nsw i32 [[VAL_73]], [[SUM_NEXT_72]]
++; CHECK-NEXT:    [[INDVARS_IV_NEXT_73:%.*]] = add nuw nsw i64 [[INDVARS_IV]], 74
++; CHECK-NEXT:    [[ARRAYIDX_74:%.*]] = getelementptr inbounds i32, ptr [[ARY]], i64 [[INDVARS_IV_NEXT_73]]
++; CHECK-NEXT:    [[VAL_74:%.*]] = load i32, ptr [[ARRAYIDX_74]], align 4
++; CHECK-NEXT:    [[SUM_NEXT_74:%.*]] = add nsw i32 [[VAL_74]], [[SUM_NEXT_73]]
++; CHECK-NEXT:    [[INDVARS_IV_NEXT_74:%.*]] = add nuw nsw i64 [[INDVARS_IV]], 75
++; CHECK-NEXT:    [[ARRAYIDX_75:%.*]] = getelementptr inbounds i32, ptr [[ARY]], i64 [[INDVARS_IV_NEXT_74]]
++; CHECK-NEXT:    [[VAL_75:%.*]] = load i32, ptr [[ARRAYIDX_75]], align 4
++; CHECK-NEXT:    [[SUM_NEXT_75:%.*]] = add nsw i32 [[VAL_75]], [[SUM_NEXT_74]]
++; CHECK-NEXT:    [[INDVARS_IV_NEXT_75:%.*]] = add nuw nsw i64 [[INDVARS_IV]], 76
++; CHECK-NEXT:    [[ARRAYIDX_76:%.*]] = getelementptr inbounds i32, ptr [[ARY]], i64 [[INDVARS_IV_NEXT_75]]
++; CHECK-NEXT:    [[VAL_76:%.*]] = load i32, ptr [[ARRAYIDX_76]], align 4
++; CHECK-NEXT:    [[SUM_NEXT_76:%.*]] = add nsw i32 [[VAL_76]], [[SUM_NEXT_75]]
++; CHECK-NEXT:    [[INDVARS_IV_NEXT_76:%.*]] = add nuw nsw i64 [[INDVARS_IV]], 77
++; CHECK-NEXT:    [[ARRAYIDX_77:%.*]] = getelementptr inbounds i32, ptr [[ARY]], i64 [[INDVARS_IV_NEXT_76]]
++; CHECK-NEXT:    [[VAL_77:%.*]] = load i32, ptr [[ARRAYIDX_77]], align 4
++; CHECK-NEXT:    [[SUM_NEXT_77:%.*]] = add nsw i32 [[VAL_77]], [[SUM_NEXT_76]]
++; CHECK-NEXT:    [[INDVARS_IV_NEXT_77:%.*]] = add nuw nsw i64 [[INDVARS_IV]], 78
++; CHECK-NEXT:    [[ARRAYIDX_78:%.*]] = getelementptr inbounds i32, ptr [[ARY]], i64 [[INDVARS_IV_NEXT_77]]
++; CHECK-NEXT:    [[VAL_78:%.*]] = load i32, ptr [[ARRAYIDX_78]], align 4
++; CHECK-NEXT:    [[SUM_NEXT_78:%.*]] = add nsw i32 [[VAL_78]], [[SUM_NEXT_77]]
++; CHECK-NEXT:    [[INDVARS_IV_NEXT_78:%.*]] = add nuw nsw i64 [[INDVARS_IV]], 79
++; CHECK-NEXT:    [[ARRAYIDX_79:%.*]] = getelementptr inbounds i32, ptr [[ARY]], i64 [[INDVARS_IV_NEXT_78]]
++; CHECK-NEXT:    [[VAL_79:%.*]] = load i32, ptr [[ARRAYIDX_79]], align 4
++; CHECK-NEXT:    [[SUM_NEXT_79:%.*]] = add nsw i32 [[VAL_79]], [[SUM_NEXT_78]]
++; CHECK-NEXT:    [[INDVARS_IV_NEXT_79:%.*]] = add nuw nsw i64 [[INDVARS_IV]], 80
++; CHECK-NEXT:    [[ARRAYIDX_80:%.*]] = getelementptr inbounds i32, ptr [[ARY]], i64 [[INDVARS_IV_NEXT_79]]
++; CHECK-NEXT:    [[VAL_80:%.*]] = load i32, ptr [[ARRAYIDX_80]], align 4
++; CHECK-NEXT:    [[SUM_NEXT_80:%.*]] = add nsw i32 [[VAL_80]], [[SUM_NEXT_79]]
++; CHECK-NEXT:    [[INDVARS_IV_NEXT_80:%.*]] = add nuw nsw i64 [[INDVARS_IV]], 81
++; CHECK-NEXT:    [[ARRAYIDX_81:%.*]] = getelementptr inbounds i32, ptr [[ARY]], i64 [[INDVARS_IV_NEXT_80]]
++; CHECK-NEXT:    [[VAL_81:%.*]] = load i32, ptr [[ARRAYIDX_81]], align 4
++; CHECK-NEXT:    [[SUM_NEXT_81:%.*]] = add nsw i32 [[VAL_81]], [[SUM_NEXT_80]]
++; CHECK-NEXT:    [[INDVARS_IV_NEXT_81:%.*]] = add nuw nsw i64 [[INDVARS_IV]], 82
++; CHECK-NEXT:    [[ARRAYIDX_82:%.*]] = getelementptr inbounds i32, ptr [[ARY]], i64 [[INDVARS_IV_NEXT_81]]
++; CHECK-NEXT:    [[VAL_82:%.*]] = load i32, ptr [[ARRAYIDX_82]], align 4
++; CHECK-NEXT:    [[SUM_NEXT_82:%.*]] = add nsw i32 [[VAL_82]], [[SUM_NEXT_81]]
++; CHECK-NEXT:    [[INDVARS_IV_NEXT_82:%.*]] = add nuw nsw i64 [[INDVARS_IV]], 83
++; CHECK-NEXT:    [[ARRAYIDX_83:%.*]] = getelementptr inbounds i32, ptr [[ARY]], i64 [[INDVARS_IV_NEXT_82]]
++; CHECK-NEXT:    [[VAL_83:%.*]] = load i32, ptr [[ARRAYIDX_83]], align 4
++; CHECK-NEXT:    [[SUM_NEXT_83:%.*]] = add nsw i32 [[VAL_83]], [[SUM_NEXT_82]]
++; CHECK-NEXT:    [[INDVARS_IV_NEXT_83:%.*]] = add nuw nsw i64 [[INDVARS_IV]], 84
++; CHECK-NEXT:    [[ARRAYIDX_84:%.*]] = getelementptr inbounds i32, ptr [[ARY]], i64 [[INDVARS_IV_NEXT_83]]
++; CHECK-NEXT:    [[VAL_84:%.*]] = load i32, ptr [[ARRAYIDX_84]], align 4
++; CHECK-NEXT:    [[SUM_NEXT_84:%.*]] = add nsw i32 [[VAL_84]], [[SUM_NEXT_83]]
++; CHECK-NEXT:    [[INDVARS_IV_NEXT_84:%.*]] = add nuw nsw i64 [[INDVARS_IV]], 85
++; CHECK-NEXT:    [[ARRAYIDX_85:%.*]] = getelementptr inbounds i32, ptr [[ARY]], i64 [[INDVARS_IV_NEXT_84]]
++; CHECK-NEXT:    [[VAL_85:%.*]] = load i32, ptr [[ARRAYIDX_85]], align 4
++; CHECK-NEXT:    [[SUM_NEXT_85:%.*]] = add nsw i32 [[VAL_85]], [[SUM_NEXT_84]]
++; CHECK-NEXT:    [[INDVARS_IV_NEXT_85:%.*]] = add nuw nsw i64 [[INDVARS_IV]], 86
++; CHECK-NEXT:    [[ARRAYIDX_86:%.*]] = getelementptr inbounds i32, ptr [[ARY]], i64 [[INDVARS_IV_NEXT_85]]
++; CHECK-NEXT:    [[VAL_86:%.*]] = load i32, ptr [[ARRAYIDX_86]], align 4
++; CHECK-NEXT:    [[SUM_NEXT_86:%.*]] = add nsw i32 [[VAL_86]], [[SUM_NEXT_85]]
++; CHECK-NEXT:    [[INDVARS_IV_NEXT_86:%.*]] = add nuw nsw i64 [[INDVARS_IV]], 87
++; CHECK-NEXT:    [[ARRAYIDX_87:%.*]] = getelementptr inbounds i32, ptr [[ARY]], i64 [[INDVARS_IV_NEXT_86]]
++; CHECK-NEXT:    [[VAL_87:%.*]] = load i32, ptr [[ARRAYIDX_87]], align 4
++; CHECK-NEXT:    [[SUM_NEXT_87:%.*]] = add nsw i32 [[VAL_87]], [[SUM_NEXT_86]]
++; CHECK-NEXT:    [[INDVARS_IV_NEXT_87:%.*]] = add nuw nsw i64 [[INDVARS_IV]], 88
++; CHECK-NEXT:    [[ARRAYIDX_88:%.*]] = getelementptr inbounds i32, ptr [[ARY]], i64 [[INDVARS_IV_NEXT_87]]
++; CHECK-NEXT:    [[VAL_88:%.*]] = load i32, ptr [[ARRAYIDX_88]], align 4
++; CHECK-NEXT:    [[SUM_NEXT_88:%.*]] = add nsw i32 [[VAL_88]], [[SUM_NEXT_87]]
++; CHECK-NEXT:    [[INDVARS_IV_NEXT_88:%.*]] = add nuw nsw i64 [[INDVARS_IV]], 89
++; CHECK-NEXT:    [[ARRAYIDX_89:%.*]] = getelementptr inbounds i32, ptr [[ARY]], i64 [[INDVARS_IV_NEXT_88]]
++; CHECK-NEXT:    [[VAL_89:%.*]] = load i32, ptr [[ARRAYIDX_89]], align 4
++; CHECK-NEXT:    [[SUM_NEXT_89:%.*]] = add nsw i32 [[VAL_89]], [[SUM_NEXT_88]]
++; CHECK-NEXT:    [[INDVARS_IV_NEXT_89:%.*]] = add nuw nsw i64 [[INDVARS_IV]], 90
++; CHECK-NEXT:    [[ARRAYIDX_90:%.*]] = getelementptr inbounds i32, ptr [[ARY]], i64 [[INDVARS_IV_NEXT_89]]
++; CHECK-NEXT:    [[VAL_90:%.*]] = load i32, ptr [[ARRAYIDX_90]], align 4
++; CHECK-NEXT:    [[SUM_NEXT_90:%.*]] = add nsw i32 [[VAL_90]], [[SUM_NEXT_89]]
++; CHECK-NEXT:    [[INDVARS_IV_NEXT_90:%.*]] = add nuw nsw i64 [[INDVARS_IV]], 91
++; CHECK-NEXT:    [[ARRAYIDX_91:%.*]] = getelementptr inbounds i32, ptr [[ARY]], i64 [[INDVARS_IV_NEXT_90]]
++; CHECK-NEXT:    [[VAL_91:%.*]] = load i32, ptr [[ARRAYIDX_91]], align 4
++; CHECK-NEXT:    [[SUM_NEXT_91:%.*]] = add nsw i32 [[VAL_91]], [[SUM_NEXT_90]]
++; CHECK-NEXT:    [[INDVARS_IV_NEXT_91:%.*]] = add nuw nsw i64 [[INDVARS_IV]], 92
++; CHECK-NEXT:    [[ARRAYIDX_92:%.*]] = getelementptr inbounds i32, ptr [[ARY]], i64 [[INDVARS_IV_NEXT_91]]
++; CHECK-NEXT:    [[VAL_92:%.*]] = load i32, ptr [[ARRAYIDX_92]], align 4
++; CHECK-NEXT:    [[SUM_NEXT_92:%.*]] = add nsw i32 [[VAL_92]], [[SUM_NEXT_91]]
++; CHECK-NEXT:    [[INDVARS_IV_NEXT_92:%.*]] = add nuw nsw i64 [[INDVARS_IV]], 93
++; CHECK-NEXT:    [[ARRAYIDX_93:%.*]] = getelementptr inbounds i32, ptr [[ARY]], i64 [[INDVARS_IV_NEXT_92]]
++; CHECK-NEXT:    [[VAL_93:%.*]] = load i32, ptr [[ARRAYIDX_93]], align 4
++; CHECK-NEXT:    [[SUM_NEXT_93:%.*]] = add nsw i32 [[VAL_93]], [[SUM_NEXT_92]]
++; CHECK-NEXT:    [[INDVARS_IV_NEXT_93:%.*]] = add nuw nsw i64 [[INDVARS_IV]], 94
++; CHECK-NEXT:    [[ARRAYIDX_94:%.*]] = getelementptr inbounds i32, ptr [[ARY]], i64 [[INDVARS_IV_NEXT_93]]
++; CHECK-NEXT:    [[VAL_94:%.*]] = load i32, ptr [[ARRAYIDX_94]], align 4
++; CHECK-NEXT:    [[SUM_NEXT_94:%.*]] = add nsw i32 [[VAL_94]], [[SUM_NEXT_93]]
++; CHECK-NEXT:    [[INDVARS_IV_NEXT_94:%.*]] = add nuw nsw i64 [[INDVARS_IV]], 95
++; CHECK-NEXT:    [[ARRAYIDX_95:%.*]] = getelementptr inbounds i32, ptr [[ARY]], i64 [[INDVARS_IV_NEXT_94]]
++; CHECK-NEXT:    [[VAL_95:%.*]] = load i32, ptr [[ARRAYIDX_95]], align 4
++; CHECK-NEXT:    [[SUM_NEXT_95:%.*]] = add nsw i32 [[VAL_95]], [[SUM_NEXT_94]]
++; CHECK-NEXT:    [[INDVARS_IV_NEXT_95:%.*]] = add nuw nsw i64 [[INDVARS_IV]], 96
++; CHECK-NEXT:    [[ARRAYIDX_96:%.*]] = getelementptr inbounds i32, ptr [[ARY]], i64 [[INDVARS_IV_NEXT_95]]
++; CHECK-NEXT:    [[VAL_96:%.*]] = load i32, ptr [[ARRAYIDX_96]], align 4
++; CHECK-NEXT:    [[SUM_NEXT_96:%.*]] = add nsw i32 [[VAL_96]], [[SUM_NEXT_95]]
++; CHECK-NEXT:    [[INDVARS_IV_NEXT_96:%.*]] = add nuw nsw i64 [[INDVARS_IV]], 97
++; CHECK-NEXT:    [[ARRAYIDX_97:%.*]] = getelementptr inbounds i32, ptr [[ARY]], i64 [[INDVARS_IV_NEXT_96]]
++; CHECK-NEXT:    [[VAL_97:%.*]] = load i32, ptr [[ARRAYIDX_97]], align 4
++; CHECK-NEXT:    [[SUM_NEXT_97:%.*]] = add nsw i32 [[VAL_97]], [[SUM_NEXT_96]]
++; CHECK-NEXT:    [[INDVARS_IV_NEXT_97:%.*]] = add nuw nsw i64 [[INDVARS_IV]], 98
++; CHECK-NEXT:    [[ARRAYIDX_98:%.*]] = getelementptr inbounds i32, ptr [[ARY]], i64 [[INDVARS_IV_NEXT_97]]
++; CHECK-NEXT:    [[VAL_98:%.*]] = load i32, ptr [[ARRAYIDX_98]], align 4
++; CHECK-NEXT:    [[SUM_NEXT_98:%.*]] = add nsw i32 [[VAL_98]], [[SUM_NEXT_97]]
++; CHECK-NEXT:    [[INDVARS_IV_NEXT_98:%.*]] = add nuw nsw i64 [[INDVARS_IV]], 99
++; CHECK-NEXT:    [[ARRAYIDX_99:%.*]] = getelementptr inbounds i32, ptr [[ARY]], i64 [[INDVARS_IV_NEXT_98]]
++; CHECK-NEXT:    [[VAL_99:%.*]] = load i32, ptr [[ARRAYIDX_99]], align 4
++; CHECK-NEXT:    [[SUM_NEXT_99:%.*]] = add nsw i32 [[VAL_99]], [[SUM_NEXT_98]]
++; CHECK-NEXT:    [[INDVARS_IV_NEXT_99:%.*]] = add nuw nsw i64 [[INDVARS_IV]], 100
++; CHECK-NEXT:    [[ARRAYIDX_100:%.*]] = getelementptr inbounds i32, ptr [[ARY]], i64 [[INDVARS_IV_NEXT_99]]
++; CHECK-NEXT:    [[VAL_100:%.*]] = load i32, ptr [[ARRAYIDX_100]], align 4
++; CHECK-NEXT:    [[SUM_NEXT_100:%.*]] = add nsw i32 [[VAL_100]], [[SUM_NEXT_99]]
++; CHECK-NEXT:    [[INDVARS_IV_NEXT_100:%.*]] = add nuw nsw i64 [[INDVARS_IV]], 101
++; CHECK-NEXT:    [[ARRAYIDX_101:%.*]] = getelementptr inbounds i32, ptr [[ARY]], i64 [[INDVARS_IV_NEXT_100]]
++; CHECK-NEXT:    [[VAL_101:%.*]] = load i32, ptr [[ARRAYIDX_101]], align 4
++; CHECK-NEXT:    [[SUM_NEXT_101:%.*]] = add nsw i32 [[VAL_101]], [[SUM_NEXT_100]]
++; CHECK-NEXT:    [[INDVARS_IV_NEXT_101:%.*]] = add nuw nsw i64 [[INDVARS_IV]], 102
++; CHECK-NEXT:    [[ARRAYIDX_102:%.*]] = getelementptr inbounds i32, ptr [[ARY]], i64 [[INDVARS_IV_NEXT_101]]
++; CHECK-NEXT:    [[VAL_102:%.*]] = load i32, ptr [[ARRAYIDX_102]], align 4
++; CHECK-NEXT:    [[SUM_NEXT_102:%.*]] = add nsw i32 [[VAL_102]], [[SUM_NEXT_101]]
++; CHECK-NEXT:    [[INDVARS_IV_NEXT_102:%.*]] = add nuw nsw i64 [[INDVARS_IV]], 103
++; CHECK-NEXT:    [[ARRAYIDX_103:%.*]] = getelementptr inbounds i32, ptr [[ARY]], i64 [[INDVARS_IV_NEXT_102]]
++; CHECK-NEXT:    [[VAL_103:%.*]] = load i32, ptr [[ARRAYIDX_103]], align 4
++; CHECK-NEXT:    [[SUM_NEXT_103:%.*]] = add nsw i32 [[VAL_103]], [[SUM_NEXT_102]]
++; CHECK-NEXT:    [[INDVARS_IV_NEXT_103:%.*]] = add nuw nsw i64 [[INDVARS_IV]], 104
++; CHECK-NEXT:    [[ARRAYIDX_104:%.*]] = getelementptr inbounds i32, ptr [[ARY]], i64 [[INDVARS_IV_NEXT_103]]
++; CHECK-NEXT:    [[VAL_104:%.*]] = load i32, ptr [[ARRAYIDX_104]], align 4
++; CHECK-NEXT:    [[SUM_NEXT_104:%.*]] = add nsw i32 [[VAL_104]], [[SUM_NEXT_103]]
++; CHECK-NEXT:    [[INDVARS_IV_NEXT_104:%.*]] = add nuw nsw i64 [[INDVARS_IV]], 105
++; CHECK-NEXT:    [[ARRAYIDX_105:%.*]] = getelementptr inbounds i32, ptr [[ARY]], i64 [[INDVARS_IV_NEXT_104]]
++; CHECK-NEXT:    [[VAL_105:%.*]] = load i32, ptr [[ARRAYIDX_105]], align 4
++; CHECK-NEXT:    [[SUM_NEXT_105:%.*]] = add nsw i32 [[VAL_105]], [[SUM_NEXT_104]]
++; CHECK-NEXT:    [[INDVARS_IV_NEXT_105:%.*]] = add nuw nsw i64 [[INDVARS_IV]], 106
++; CHECK-NEXT:    [[ARRAYIDX_106:%.*]] = getelementptr inbounds i32, ptr [[ARY]], i64 [[INDVARS_IV_NEXT_105]]
++; CHECK-NEXT:    [[VAL_106:%.*]] = load i32, ptr [[ARRAYIDX_106]], align 4
++; CHECK-NEXT:    [[SUM_NEXT_106:%.*]] = add nsw i32 [[VAL_106]], [[SUM_NEXT_105]]
++; CHECK-NEXT:    [[INDVARS_IV_NEXT_106:%.*]] = add nuw nsw i64 [[INDVARS_IV]], 107
++; CHECK-NEXT:    [[ARRAYIDX_107:%.*]] = getelementptr inbounds i32, ptr [[ARY]], i64 [[INDVARS_IV_NEXT_106]]
++; CHECK-NEXT:    [[VAL_107:%.*]] = load i32, ptr [[ARRAYIDX_107]], align 4
++; CHECK-NEXT:    [[SUM_NEXT_107:%.*]] = add nsw i32 [[VAL_107]], [[SUM_NEXT_106]]
++; CHECK-NEXT:    [[INDVARS_IV_NEXT_107:%.*]] = add nuw nsw i64 [[INDVARS_IV]], 108
++; CHECK-NEXT:    [[ARRAYIDX_108:%.*]] = getelementptr inbounds i32, ptr [[ARY]], i64 [[INDVARS_IV_NEXT_107]]
++; CHECK-NEXT:    [[VAL_108:%.*]] = load i32, ptr [[ARRAYIDX_108]], align 4
++; CHECK-NEXT:    [[SUM_NEXT_108:%.*]] = add nsw i32 [[VAL_108]], [[SUM_NEXT_107]]
++; CHECK-NEXT:    [[INDVARS_IV_NEXT_108:%.*]] = add nuw nsw i64 [[INDVARS_IV]], 109
++; CHECK-NEXT:    [[ARRAYIDX_109:%.*]] = getelementptr inbounds i32, ptr [[ARY]], i64 [[INDVARS_IV_NEXT_108]]
++; CHECK-NEXT:    [[VAL_109:%.*]] = load i32, ptr [[ARRAYIDX_109]], align 4
++; CHECK-NEXT:    [[SUM_NEXT_109:%.*]] = add nsw i32 [[VAL_109]], [[SUM_NEXT_108]]
++; CHECK-NEXT:    [[INDVARS_IV_NEXT_109:%.*]] = add nuw nsw i64 [[INDVARS_IV]], 110
++; CHECK-NEXT:    [[ARRAYIDX_110:%.*]] = getelementptr inbounds i32, ptr [[ARY]], i64 [[INDVARS_IV_NEXT_109]]
++; CHECK-NEXT:    [[VAL_110:%.*]] = load i32, ptr [[ARRAYIDX_110]], align 4
++; CHECK-NEXT:    [[SUM_NEXT_110:%.*]] = add nsw i32 [[VAL_110]], [[SUM_NEXT_109]]
++; CHECK-NEXT:    [[INDVARS_IV_NEXT_110:%.*]] = add nuw nsw i64 [[INDVARS_IV]], 111
++; CHECK-NEXT:    [[ARRAYIDX_111:%.*]] = getelementptr inbounds i32, ptr [[ARY]], i64 [[INDVARS_IV_NEXT_110]]
++; CHECK-NEXT:    [[VAL_111:%.*]] = load i32, ptr [[ARRAYIDX_111]], align 4
++; CHECK-NEXT:    [[SUM_NEXT_111:%.*]] = add nsw i32 [[VAL_111]], [[SUM_NEXT_110]]
++; CHECK-NEXT:    [[INDVARS_IV_NEXT_111:%.*]] = add nuw nsw i64 [[INDVARS_IV]], 112
++; CHECK-NEXT:    [[ARRAYIDX_112:%.*]] = getelementptr inbounds i32, ptr [[ARY]], i64 [[INDVARS_IV_NEXT_111]]
++; CHECK-NEXT:    [[VAL_112:%.*]] = load i32, ptr [[ARRAYIDX_112]], align 4
++; CHECK-NEXT:    [[SUM_NEXT_112:%.*]] = add nsw i32 [[VAL_112]], [[SUM_NEXT_111]]
++; CHECK-NEXT:    [[INDVARS_IV_NEXT_112:%.*]] = add nuw nsw i64 [[INDVARS_IV]], 113
++; CHECK-NEXT:    [[ARRAYIDX_113:%.*]] = getelementptr inbounds i32, ptr [[ARY]], i64 [[INDVARS_IV_NEXT_112]]
++; CHECK-NEXT:    [[VAL_113:%.*]] = load i32, ptr [[ARRAYIDX_113]], align 4
++; CHECK-NEXT:    [[SUM_NEXT_113:%.*]] = add nsw i32 [[VAL_113]], [[SUM_NEXT_112]]
++; CHECK-NEXT:    [[INDVARS_IV_NEXT_113:%.*]] = add nuw nsw i64 [[INDVARS_IV]], 114
++; CHECK-NEXT:    [[ARRAYIDX_114:%.*]] = getelementptr inbounds i32, ptr [[ARY]], i64 [[INDVARS_IV_NEXT_113]]
++; CHECK-NEXT:    [[VAL_114:%.*]] = load i32, ptr [[ARRAYIDX_114]], align 4
++; CHECK-NEXT:    [[SUM_NEXT_114:%.*]] = add nsw i32 [[VAL_114]], [[SUM_NEXT_113]]
++; CHECK-NEXT:    [[INDVARS_IV_NEXT_114:%.*]] = add nuw nsw i64 [[INDVARS_IV]], 115
++; CHECK-NEXT:    [[ARRAYIDX_115:%.*]] = getelementptr inbounds i32, ptr [[ARY]], i64 [[INDVARS_IV_NEXT_114]]
++; CHECK-NEXT:    [[VAL_115:%.*]] = load i32, ptr [[ARRAYIDX_115]], align 4
++; CHECK-NEXT:    [[SUM_NEXT_115:%.*]] = add nsw i32 [[VAL_115]], [[SUM_NEXT_114]]
++; CHECK-NEXT:    [[INDVARS_IV_NEXT_115:%.*]] = add nuw nsw i64 [[INDVARS_IV]], 116
++; CHECK-NEXT:    [[ARRAYIDX_116:%.*]] = getelementptr inbounds i32, ptr [[ARY]], i64 [[INDVARS_IV_NEXT_115]]
++; CHECK-NEXT:    [[VAL_116:%.*]] = load i32, ptr [[ARRAYIDX_116]], align 4
++; CHECK-NEXT:    [[SUM_NEXT_116:%.*]] = add nsw i32 [[VAL_116]], [[SUM_NEXT_115]]
++; CHECK-NEXT:    [[INDVARS_IV_NEXT_116:%.*]] = add nuw nsw i64 [[INDVARS_IV]], 117
++; CHECK-NEXT:    [[ARRAYIDX_117:%.*]] = getelementptr inbounds i32, ptr [[ARY]], i64 [[INDVARS_IV_NEXT_116]]
++; CHECK-NEXT:    [[VAL_117:%.*]] = load i32, ptr [[ARRAYIDX_117]], align 4
++; CHECK-NEXT:    [[SUM_NEXT_117:%.*]] = add nsw i32 [[VAL_117]], [[SUM_NEXT_116]]
++; CHECK-NEXT:    [[INDVARS_IV_NEXT_117:%.*]] = add nuw nsw i64 [[INDVARS_IV]], 118
++; CHECK-NEXT:    [[ARRAYIDX_118:%.*]] = getelementptr inbounds i32, ptr [[ARY]], i64 [[INDVARS_IV_NEXT_117]]
++; CHECK-NEXT:    [[VAL_118:%.*]] = load i32, ptr [[ARRAYIDX_118]], align 4
++; CHECK-NEXT:    [[SUM_NEXT_118:%.*]] = add nsw i32 [[VAL_118]], [[SUM_NEXT_117]]
++; CHECK-NEXT:    [[INDVARS_IV_NEXT_118:%.*]] = add nuw nsw i64 [[INDVARS_IV]], 119
++; CHECK-NEXT:    [[ARRAYIDX_119:%.*]] = getelementptr inbounds i32, ptr [[ARY]], i64 [[INDVARS_IV_NEXT_118]]
++; CHECK-NEXT:    [[VAL_119:%.*]] = load i32, ptr [[ARRAYIDX_119]], align 4
++; CHECK-NEXT:    [[SUM_NEXT_119:%.*]] = add nsw i32 [[VAL_119]], [[SUM_NEXT_118]]
++; CHECK-NEXT:    [[INDVARS_IV_NEXT_119:%.*]] = add nuw nsw i64 [[INDVARS_IV]], 120
++; CHECK-NEXT:    [[ARRAYIDX_120:%.*]] = getelementptr inbounds i32, ptr [[ARY]], i64 [[INDVARS_IV_NEXT_119]]
++; CHECK-NEXT:    [[VAL_120:%.*]] = load i32, ptr [[ARRAYIDX_120]], align 4
++; CHECK-NEXT:    [[SUM_NEXT_120:%.*]] = add nsw i32 [[VAL_120]], [[SUM_NEXT_119]]
++; CHECK-NEXT:    [[INDVARS_IV_NEXT_120:%.*]] = add nuw nsw i64 [[INDVARS_IV]], 121
++; CHECK-NEXT:    [[ARRAYIDX_121:%.*]] = getelementptr inbounds i32, ptr [[ARY]], i64 [[INDVARS_IV_NEXT_120]]
++; CHECK-NEXT:    [[VAL_121:%.*]] = load i32, ptr [[ARRAYIDX_121]], align 4
++; CHECK-NEXT:    [[SUM_NEXT_121:%.*]] = add nsw i32 [[VAL_121]], [[SUM_NEXT_120]]
++; CHECK-NEXT:    [[INDVARS_IV_NEXT_121:%.*]] = add nuw nsw i64 [[INDVARS_IV]], 122
++; CHECK-NEXT:    [[ARRAYIDX_122:%.*]] = getelementptr inbounds i32, ptr [[ARY]], i64 [[INDVARS_IV_NEXT_121]]
++; CHECK-NEXT:    [[VAL_122:%.*]] = load i32, ptr [[ARRAYIDX_122]], align 4
++; CHECK-NEXT:    [[SUM_NEXT_122:%.*]] = add nsw i32 [[VAL_122]], [[SUM_NEXT_121]]
++; CHECK-NEXT:    [[INDVARS_IV_NEXT_122:%.*]] = add nuw nsw i64 [[INDVARS_IV]], 123
++; CHECK-NEXT:    [[ARRAYIDX_123:%.*]] = getelementptr inbounds i32, ptr [[ARY]], i64 [[INDVARS_IV_NEXT_122]]
++; CHECK-NEXT:    [[VAL_123:%.*]] = load i32, ptr [[ARRAYIDX_123]], align 4
++; CHECK-NEXT:    [[SUM_NEXT_123:%.*]] = add nsw i32 [[VAL_123]], [[SUM_NEXT_122]]
++; CHECK-NEXT:    [[INDVARS_IV_NEXT_123:%.*]] = add nuw nsw i64 [[INDVARS_IV]], 124
++; CHECK-NEXT:    [[ARRAYIDX_124:%.*]] = getelementptr inbounds i32, ptr [[ARY]], i64 [[INDVARS_IV_NEXT_123]]
++; CHECK-NEXT:    [[VAL_124:%.*]] = load i32, ptr [[ARRAYIDX_124]], align 4
++; CHECK-NEXT:    [[SUM_NEXT_124:%.*]] = add nsw i32 [[VAL_124]], [[SUM_NEXT_123]]
++; CHECK-NEXT:    [[INDVARS_IV_NEXT_124:%.*]] = add nuw nsw i64 [[INDVARS_IV]], 125
++; CHECK-NEXT:    [[ARRAYIDX_125:%.*]] = getelementptr inbounds i32, ptr [[ARY]], i64 [[INDVARS_IV_NEXT_124]]
++; CHECK-NEXT:    [[VAL_125:%.*]] = load i32, ptr [[ARRAYIDX_125]], align 4
++; CHECK-NEXT:    [[SUM_NEXT_125:%.*]] = add nsw i32 [[VAL_125]], [[SUM_NEXT_124]]
++; CHECK-NEXT:    [[INDVARS_IV_NEXT_125:%.*]] = add nuw nsw i64 [[INDVARS_IV]], 126
++; CHECK-NEXT:    [[ARRAYIDX_126:%.*]] = getelementptr inbounds i32, ptr [[ARY]], i64 [[INDVARS_IV_NEXT_125]]
++; CHECK-NEXT:    [[VAL_126:%.*]] = load i32, ptr [[ARRAYIDX_126]], align 4
++; CHECK-NEXT:    [[SUM_NEXT_126:%.*]] = add nsw i32 [[VAL_126]], [[SUM_NEXT_125]]
++; CHECK-NEXT:    [[INDVARS_IV_NEXT_126:%.*]] = add nuw nsw i64 [[INDVARS_IV]], 127
++; CHECK-NEXT:    [[ARRAYIDX_127:%.*]] = getelementptr inbounds i32, ptr [[ARY]], i64 [[INDVARS_IV_NEXT_126]]
++; CHECK-NEXT:    [[VAL_127:%.*]] = load i32, ptr [[ARRAYIDX_127]], align 4
++; CHECK-NEXT:    [[SUM_NEXT_127]] = add nsw i32 [[VAL_127]], [[SUM_NEXT_126]]
++; CHECK-NEXT:    [[INDVARS_IV_NEXT_127]] = add nuw nsw i64 [[INDVARS_IV]], 128
++; CHECK-NEXT:    [[EXITCOND_NOT_127:%.*]] = icmp eq i64 [[INDVARS_IV_NEXT_127]], 8192
++; CHECK-NEXT:    br i1 [[EXITCOND_NOT_127]], label [[FOR_COND_CLEANUP:%.*]], label [[FOR_BODY]]
+ ; CHECK:       for.cond.cleanup:
+-; CHECK-NEXT:    [[SUM_NEXT_LCSSA:%.*]] = phi i32 [ [[SUM_NEXT_31]], [[FOR_BODY]] ]
++; CHECK-NEXT:    [[SUM_NEXT_LCSSA:%.*]] = phi i32 [ [[SUM_NEXT_127]], [[FOR_BODY]] ]
+ ; CHECK-NEXT:    ret i32 [[SUM_NEXT_LCSSA]]
+ ;
+ entry:
+@@ -167,16 +551,16 @@
+ ; CHECK-SAME: ptr [[ARY:%.*]], i64 [[N:%.*]]) #[[ATTR0]] {
+ ; CHECK-NEXT:  entry:
+ ; CHECK-NEXT:    [[TMP0:%.*]] = add i64 [[N]], -1
+-; CHECK-NEXT:    [[XTRAITER:%.*]] = and i64 [[N]], 1
+-; CHECK-NEXT:    [[TMP1:%.*]] = icmp ult i64 [[TMP0]], 1
++; CHECK-NEXT:    [[XTRAITER:%.*]] = and i64 [[N]], 7
++; CHECK-NEXT:    [[TMP1:%.*]] = icmp ult i64 [[TMP0]], 7
+ ; CHECK-NEXT:    br i1 [[TMP1]], label [[FOR_COND_CLEANUP_UNR_LCSSA:%.*]], label [[ENTRY_NEW:%.*]]
+ ; CHECK:       entry.new:
+ ; CHECK-NEXT:    [[UNROLL_ITER:%.*]] = sub i64 [[N]], [[XTRAITER]]
+ ; CHECK-NEXT:    br label [[FOR_BODY:%.*]]
+ ; CHECK:       for.body:
+-; CHECK-NEXT:    [[INDVARS_IV:%.*]] = phi i64 [ 0, [[ENTRY_NEW]] ], [ [[INDVARS_IV_NEXT_1:%.*]], [[FOR_BODY]] ]
+-; CHECK-NEXT:    [[SUM:%.*]] = phi i32 [ 0, [[ENTRY_NEW]] ], [ [[SUM_NEXT_1:%.*]], [[FOR_BODY]] ]
+-; CHECK-NEXT:    [[NITER:%.*]] = phi i64 [ 0, [[ENTRY_NEW]] ], [ [[NITER_NEXT_1:%.*]], [[FOR_BODY]] ]
++; CHECK-NEXT:    [[INDVARS_IV:%.*]] = phi i64 [ 0, [[ENTRY_NEW]] ], [ [[INDVARS_IV_NEXT_7:%.*]], [[FOR_BODY]] ]
++; CHECK-NEXT:    [[SUM:%.*]] = phi i32 [ 0, [[ENTRY_NEW]] ], [ [[SUM_NEXT_7:%.*]], [[FOR_BODY]] ]
++; CHECK-NEXT:    [[NITER:%.*]] = phi i64 [ 0, [[ENTRY_NEW]] ], [ [[NITER_NEXT_7:%.*]], [[FOR_BODY]] ]
+ ; CHECK-NEXT:    [[ARRAYIDX:%.*]] = getelementptr inbounds i32, ptr [[ARY]], i64 [[INDVARS_IV]]
+ ; CHECK-NEXT:    [[VAL:%.*]] = load i32, ptr [[ARRAYIDX]], align 4
+ ; CHECK-NEXT:    [[DUMMY1:%.*]] = mul i32 [[VAL]], [[VAL]]
+@@ -283,15 +667,339 @@
+ ; CHECK-NEXT:    [[DUMMY48_1:%.*]] = mul i32 [[DUMMY47_1]], [[DUMMY47_1]]
+ ; CHECK-NEXT:    [[DUMMY49_1:%.*]] = mul i32 [[DUMMY48_1]], [[DUMMY48_1]]
+ ; CHECK-NEXT:    [[DUMMY50_1:%.*]] = mul i32 [[DUMMY49_1]], [[DUMMY49_1]]
+-; CHECK-NEXT:    [[SUM_NEXT_1]] = add nsw i32 [[DUMMY50_1]], [[SUM_NEXT]]
+-; CHECK-NEXT:    [[INDVARS_IV_NEXT_1]] = add nuw nsw i64 [[INDVARS_IV]], 2
+-; CHECK-NEXT:    [[NITER_NEXT_1]] = add i64 [[NITER]], 2
+-; CHECK-NEXT:    [[NITER_NCMP_1:%.*]] = icmp eq i64 [[NITER_NEXT_1]], [[UNROLL_ITER]]
+-; CHECK-NEXT:    br i1 [[NITER_NCMP_1]], label [[FOR_COND_CLEANUP_UNR_LCSSA_LOOPEXIT:%.*]], label [[FOR_BODY]]
++; CHECK-NEXT:    [[SUM_NEXT_1:%.*]] = add nsw i32 [[DUMMY50_1]], [[SUM_NEXT]]
++; CHECK-NEXT:    [[INDVARS_IV_NEXT_1:%.*]] = add nuw nsw i64 [[INDVARS_IV]], 2
++; CHECK-NEXT:    [[ARRAYIDX_2:%.*]] = getelementptr inbounds i32, ptr [[ARY]], i64 [[INDVARS_IV_NEXT_1]]
++; CHECK-NEXT:    [[VAL_2:%.*]] = load i32, ptr [[ARRAYIDX_2]], align 4
++; CHECK-NEXT:    [[DUMMY1_2:%.*]] = mul i32 [[VAL_2]], [[VAL_2]]
++; CHECK-NEXT:    [[DUMMY2_2:%.*]] = mul i32 [[DUMMY1_2]], [[DUMMY1_2]]
++; CHECK-NEXT:    [[DUMMY3_2:%.*]] = mul i32 [[DUMMY2_2]], [[DUMMY2_2]]
++; CHECK-NEXT:    [[DUMMY4_2:%.*]] = mul i32 [[DUMMY3_2]], [[DUMMY3_2]]
++; CHECK-NEXT:    [[DUMMY5_2:%.*]] = mul i32 [[DUMMY4_2]], [[DUMMY4_2]]
++; CHECK-NEXT:    [[DUMMY6_2:%.*]] = mul i32 [[DUMMY5_2]], [[DUMMY5_2]]
++; CHECK-NEXT:    [[DUMMY7_2:%.*]] = mul i32 [[DUMMY6_2]], [[DUMMY6_2]]
++; CHECK-NEXT:    [[DUMMY8_2:%.*]] = mul i32 [[DUMMY7_2]], [[DUMMY7_2]]
++; CHECK-NEXT:    [[DUMMY9_2:%.*]] = mul i32 [[DUMMY8_2]], [[DUMMY8_2]]
++; CHECK-NEXT:    [[DUMMY10_2:%.*]] = mul i32 [[DUMMY9_2]], [[DUMMY9_2]]
++; CHECK-NEXT:    [[DUMMY11_2:%.*]] = mul i32 [[DUMMY10_2]], [[DUMMY10_2]]
++; CHECK-NEXT:    [[DUMMY12_2:%.*]] = mul i32 [[DUMMY11_2]], [[DUMMY11_2]]
++; CHECK-NEXT:    [[DUMMY13_2:%.*]] = mul i32 [[DUMMY12_2]], [[DUMMY12_2]]
++; CHECK-NEXT:    [[DUMMY14_2:%.*]] = mul i32 [[DUMMY13_2]], [[DUMMY13_2]]
++; CHECK-NEXT:    [[DUMMY15_2:%.*]] = mul i32 [[DUMMY14_2]], [[DUMMY14_2]]
++; CHECK-NEXT:    [[DUMMY16_2:%.*]] = mul i32 [[DUMMY15_2]], [[DUMMY15_2]]
++; CHECK-NEXT:    [[DUMMY17_2:%.*]] = mul i32 [[DUMMY16_2]], [[DUMMY16_2]]
++; CHECK-NEXT:    [[DUMMY18_2:%.*]] = mul i32 [[DUMMY17_2]], [[DUMMY17_2]]
++; CHECK-NEXT:    [[DUMMY19_2:%.*]] = mul i32 [[DUMMY18_2]], [[DUMMY18_2]]
++; CHECK-NEXT:    [[DUMMY20_2:%.*]] = mul i32 [[DUMMY19_2]], [[DUMMY19_2]]
++; CHECK-NEXT:    [[DUMMY21_2:%.*]] = mul i32 [[DUMMY20_2]], [[DUMMY20_2]]
++; CHECK-NEXT:    [[DUMMY22_2:%.*]] = mul i32 [[DUMMY21_2]], [[DUMMY21_2]]
++; CHECK-NEXT:    [[DUMMY23_2:%.*]] = mul i32 [[DUMMY22_2]], [[DUMMY22_2]]
++; CHECK-NEXT:    [[DUMMY24_2:%.*]] = mul i32 [[DUMMY23_2]], [[DUMMY23_2]]
++; CHECK-NEXT:    [[DUMMY25_2:%.*]] = mul i32 [[DUMMY24_2]], [[DUMMY24_2]]
++; CHECK-NEXT:    [[DUMMY26_2:%.*]] = mul i32 [[DUMMY25_2]], [[DUMMY25_2]]
++; CHECK-NEXT:    [[DUMMY27_2:%.*]] = mul i32 [[DUMMY26_2]], [[DUMMY26_2]]
++; CHECK-NEXT:    [[DUMMY28_2:%.*]] = mul i32 [[DUMMY27_2]], [[DUMMY27_2]]
++; CHECK-NEXT:    [[DUMMY29_2:%.*]] = mul i32 [[DUMMY28_2]], [[DUMMY28_2]]
++; CHECK-NEXT:    [[DUMMY30_2:%.*]] = mul i32 [[DUMMY29_2]], [[DUMMY29_2]]
++; CHECK-NEXT:    [[DUMMY31_2:%.*]] = mul i32 [[DUMMY30_2]], [[DUMMY30_2]]
++; CHECK-NEXT:    [[DUMMY32_2:%.*]] = mul i32 [[DUMMY31_2]], [[DUMMY31_2]]
++; CHECK-NEXT:    [[DUMMY33_2:%.*]] = mul i32 [[DUMMY32_2]], [[DUMMY32_2]]
++; CHECK-NEXT:    [[DUMMY34_2:%.*]] = mul i32 [[DUMMY33_2]], [[DUMMY33_2]]
++; CHECK-NEXT:    [[DUMMY35_2:%.*]] = mul i32 [[DUMMY34_2]], [[DUMMY34_2]]
++; CHECK-NEXT:    [[DUMMY36_2:%.*]] = mul i32 [[DUMMY35_2]], [[DUMMY35_2]]
++; CHECK-NEXT:    [[DUMMY37_2:%.*]] = mul i32 [[DUMMY36_2]], [[DUMMY36_2]]
++; CHECK-NEXT:    [[DUMMY38_2:%.*]] = mul i32 [[DUMMY37_2]], [[DUMMY37_2]]
++; CHECK-NEXT:    [[DUMMY39_2:%.*]] = mul i32 [[DUMMY38_2]], [[DUMMY38_2]]
++; CHECK-NEXT:    [[DUMMY40_2:%.*]] = mul i32 [[DUMMY39_2]], [[DUMMY39_2]]
++; CHECK-NEXT:    [[DUMMY41_2:%.*]] = mul i32 [[DUMMY40_2]], [[DUMMY40_2]]
++; CHECK-NEXT:    [[DUMMY42_2:%.*]] = mul i32 [[DUMMY41_2]], [[DUMMY41_2]]
++; CHECK-NEXT:    [[DUMMY43_2:%.*]] = mul i32 [[DUMMY42_2]], [[DUMMY42_2]]
++; CHECK-NEXT:    [[DUMMY44_2:%.*]] = mul i32 [[DUMMY43_2]], [[DUMMY43_2]]
++; CHECK-NEXT:    [[DUMMY45_2:%.*]] = mul i32 [[DUMMY44_2]], [[DUMMY44_2]]
++; CHECK-NEXT:    [[DUMMY46_2:%.*]] = mul i32 [[DUMMY45_2]], [[DUMMY45_2]]
++; CHECK-NEXT:    [[DUMMY47_2:%.*]] = mul i32 [[DUMMY46_2]], [[DUMMY46_2]]
++; CHECK-NEXT:    [[DUMMY48_2:%.*]] = mul i32 [[DUMMY47_2]], [[DUMMY47_2]]
++; CHECK-NEXT:    [[DUMMY49_2:%.*]] = mul i32 [[DUMMY48_2]], [[DUMMY48_2]]
++; CHECK-NEXT:    [[DUMMY50_2:%.*]] = mul i32 [[DUMMY49_2]], [[DUMMY49_2]]
++; CHECK-NEXT:    [[SUM_NEXT_2:%.*]] = add nsw i32 [[DUMMY50_2]], [[SUM_NEXT_1]]
++; CHECK-NEXT:    [[INDVARS_IV_NEXT_2:%.*]] = add nuw nsw i64 [[INDVARS_IV]], 3
++; CHECK-NEXT:    [[ARRAYIDX_3:%.*]] = getelementptr inbounds i32, ptr [[ARY]], i64 [[INDVARS_IV_NEXT_2]]
++; CHECK-NEXT:    [[VAL_3:%.*]] = load i32, ptr [[ARRAYIDX_3]], align 4
++; CHECK-NEXT:    [[DUMMY1_3:%.*]] = mul i32 [[VAL_3]], [[VAL_3]]
++; CHECK-NEXT:    [[DUMMY2_3:%.*]] = mul i32 [[DUMMY1_3]], [[DUMMY1_3]]
++; CHECK-NEXT:    [[DUMMY3_3:%.*]] = mul i32 [[DUMMY2_3]], [[DUMMY2_3]]
++; CHECK-NEXT:    [[DUMMY4_3:%.*]] = mul i32 [[DUMMY3_3]], [[DUMMY3_3]]
++; CHECK-NEXT:    [[DUMMY5_3:%.*]] = mul i32 [[DUMMY4_3]], [[DUMMY4_3]]
++; CHECK-NEXT:    [[DUMMY6_3:%.*]] = mul i32 [[DUMMY5_3]], [[DUMMY5_3]]
++; CHECK-NEXT:    [[DUMMY7_3:%.*]] = mul i32 [[DUMMY6_3]], [[DUMMY6_3]]
++; CHECK-NEXT:    [[DUMMY8_3:%.*]] = mul i32 [[DUMMY7_3]], [[DUMMY7_3]]
++; CHECK-NEXT:    [[DUMMY9_3:%.*]] = mul i32 [[DUMMY8_3]], [[DUMMY8_3]]
++; CHECK-NEXT:    [[DUMMY10_3:%.*]] = mul i32 [[DUMMY9_3]], [[DUMMY9_3]]
++; CHECK-NEXT:    [[DUMMY11_3:%.*]] = mul i32 [[DUMMY10_3]], [[DUMMY10_3]]
++; CHECK-NEXT:    [[DUMMY12_3:%.*]] = mul i32 [[DUMMY11_3]], [[DUMMY11_3]]
++; CHECK-NEXT:    [[DUMMY13_3:%.*]] = mul i32 [[DUMMY12_3]], [[DUMMY12_3]]
++; CHECK-NEXT:    [[DUMMY14_3:%.*]] = mul i32 [[DUMMY13_3]], [[DUMMY13_3]]
++; CHECK-NEXT:    [[DUMMY15_3:%.*]] = mul i32 [[DUMMY14_3]], [[DUMMY14_3]]
++; CHECK-NEXT:    [[DUMMY16_3:%.*]] = mul i32 [[DUMMY15_3]], [[DUMMY15_3]]
++; CHECK-NEXT:    [[DUMMY17_3:%.*]] = mul i32 [[DUMMY16_3]], [[DUMMY16_3]]
++; CHECK-NEXT:    [[DUMMY18_3:%.*]] = mul i32 [[DUMMY17_3]], [[DUMMY17_3]]
++; CHECK-NEXT:    [[DUMMY19_3:%.*]] = mul i32 [[DUMMY18_3]], [[DUMMY18_3]]
++; CHECK-NEXT:    [[DUMMY20_3:%.*]] = mul i32 [[DUMMY19_3]], [[DUMMY19_3]]
++; CHECK-NEXT:    [[DUMMY21_3:%.*]] = mul i32 [[DUMMY20_3]], [[DUMMY20_3]]
++; CHECK-NEXT:    [[DUMMY22_3:%.*]] = mul i32 [[DUMMY21_3]], [[DUMMY21_3]]
++; CHECK-NEXT:    [[DUMMY23_3:%.*]] = mul i32 [[DUMMY22_3]], [[DUMMY22_3]]
++; CHECK-NEXT:    [[DUMMY24_3:%.*]] = mul i32 [[DUMMY23_3]], [[DUMMY23_3]]
++; CHECK-NEXT:    [[DUMMY25_3:%.*]] = mul i32 [[DUMMY24_3]], [[DUMMY24_3]]
++; CHECK-NEXT:    [[DUMMY26_3:%.*]] = mul i32 [[DUMMY25_3]], [[DUMMY25_3]]
++; CHECK-NEXT:    [[DUMMY27_3:%.*]] = mul i32 [[DUMMY26_3]], [[DUMMY26_3]]
++; CHECK-NEXT:    [[DUMMY28_3:%.*]] = mul i32 [[DUMMY27_3]], [[DUMMY27_3]]
++; CHECK-NEXT:    [[DUMMY29_3:%.*]] = mul i32 [[DUMMY28_3]], [[DUMMY28_3]]
++; CHECK-NEXT:    [[DUMMY30_3:%.*]] = mul i32 [[DUMMY29_3]], [[DUMMY29_3]]
++; CHECK-NEXT:    [[DUMMY31_3:%.*]] = mul i32 [[DUMMY30_3]], [[DUMMY30_3]]
++; CHECK-NEXT:    [[DUMMY32_3:%.*]] = mul i32 [[DUMMY31_3]], [[DUMMY31_3]]
++; CHECK-NEXT:    [[DUMMY33_3:%.*]] = mul i32 [[DUMMY32_3]], [[DUMMY32_3]]
++; CHECK-NEXT:    [[DUMMY34_3:%.*]] = mul i32 [[DUMMY33_3]], [[DUMMY33_3]]
++; CHECK-NEXT:    [[DUMMY35_3:%.*]] = mul i32 [[DUMMY34_3]], [[DUMMY34_3]]
++; CHECK-NEXT:    [[DUMMY36_3:%.*]] = mul i32 [[DUMMY35_3]], [[DUMMY35_3]]
++; CHECK-NEXT:    [[DUMMY37_3:%.*]] = mul i32 [[DUMMY36_3]], [[DUMMY36_3]]
++; CHECK-NEXT:    [[DUMMY38_3:%.*]] = mul i32 [[DUMMY37_3]], [[DUMMY37_3]]
++; CHECK-NEXT:    [[DUMMY39_3:%.*]] = mul i32 [[DUMMY38_3]], [[DUMMY38_3]]
++; CHECK-NEXT:    [[DUMMY40_3:%.*]] = mul i32 [[DUMMY39_3]], [[DUMMY39_3]]
++; CHECK-NEXT:    [[DUMMY41_3:%.*]] = mul i32 [[DUMMY40_3]], [[DUMMY40_3]]
++; CHECK-NEXT:    [[DUMMY42_3:%.*]] = mul i32 [[DUMMY41_3]], [[DUMMY41_3]]
++; CHECK-NEXT:    [[DUMMY43_3:%.*]] = mul i32 [[DUMMY42_3]], [[DUMMY42_3]]
++; CHECK-NEXT:    [[DUMMY44_3:%.*]] = mul i32 [[DUMMY43_3]], [[DUMMY43_3]]
++; CHECK-NEXT:    [[DUMMY45_3:%.*]] = mul i32 [[DUMMY44_3]], [[DUMMY44_3]]
++; CHECK-NEXT:    [[DUMMY46_3:%.*]] = mul i32 [[DUMMY45_3]], [[DUMMY45_3]]
++; CHECK-NEXT:    [[DUMMY47_3:%.*]] = mul i32 [[DUMMY46_3]], [[DUMMY46_3]]
++; CHECK-NEXT:    [[DUMMY48_3:%.*]] = mul i32 [[DUMMY47_3]], [[DUMMY47_3]]
++; CHECK-NEXT:    [[DUMMY49_3:%.*]] = mul i32 [[DUMMY48_3]], [[DUMMY48_3]]
++; CHECK-NEXT:    [[DUMMY50_3:%.*]] = mul i32 [[DUMMY49_3]], [[DUMMY49_3]]
++; CHECK-NEXT:    [[SUM_NEXT_3:%.*]] = add nsw i32 [[DUMMY50_3]], [[SUM_NEXT_2]]
++; CHECK-NEXT:    [[INDVARS_IV_NEXT_3:%.*]] = add nuw nsw i64 [[INDVARS_IV]], 4
++; CHECK-NEXT:    [[ARRAYIDX_4:%.*]] = getelementptr inbounds i32, ptr [[ARY]], i64 [[INDVARS_IV_NEXT_3]]
++; CHECK-NEXT:    [[VAL_4:%.*]] = load i32, ptr [[ARRAYIDX_4]], align 4
++; CHECK-NEXT:    [[DUMMY1_4:%.*]] = mul i32 [[VAL_4]], [[VAL_4]]
++; CHECK-NEXT:    [[DUMMY2_4:%.*]] = mul i32 [[DUMMY1_4]], [[DUMMY1_4]]
++; CHECK-NEXT:    [[DUMMY3_4:%.*]] = mul i32 [[DUMMY2_4]], [[DUMMY2_4]]
++; CHECK-NEXT:    [[DUMMY4_4:%.*]] = mul i32 [[DUMMY3_4]], [[DUMMY3_4]]
++; CHECK-NEXT:    [[DUMMY5_4:%.*]] = mul i32 [[DUMMY4_4]], [[DUMMY4_4]]
++; CHECK-NEXT:    [[DUMMY6_4:%.*]] = mul i32 [[DUMMY5_4]], [[DUMMY5_4]]
++; CHECK-NEXT:    [[DUMMY7_4:%.*]] = mul i32 [[DUMMY6_4]], [[DUMMY6_4]]
++; CHECK-NEXT:    [[DUMMY8_4:%.*]] = mul i32 [[DUMMY7_4]], [[DUMMY7_4]]
++; CHECK-NEXT:    [[DUMMY9_4:%.*]] = mul i32 [[DUMMY8_4]], [[DUMMY8_4]]
++; CHECK-NEXT:    [[DUMMY10_4:%.*]] = mul i32 [[DUMMY9_4]], [[DUMMY9_4]]
++; CHECK-NEXT:    [[DUMMY11_4:%.*]] = mul i32 [[DUMMY10_4]], [[DUMMY10_4]]
++; CHECK-NEXT:    [[DUMMY12_4:%.*]] = mul i32 [[DUMMY11_4]], [[DUMMY11_4]]
++; CHECK-NEXT:    [[DUMMY13_4:%.*]] = mul i32 [[DUMMY12_4]], [[DUMMY12_4]]
++; CHECK-NEXT:    [[DUMMY14_4:%.*]] = mul i32 [[DUMMY13_4]], [[DUMMY13_4]]
++; CHECK-NEXT:    [[DUMMY15_4:%.*]] = mul i32 [[DUMMY14_4]], [[DUMMY14_4]]
++; CHECK-NEXT:    [[DUMMY16_4:%.*]] = mul i32 [[DUMMY15_4]], [[DUMMY15_4]]
++; CHECK-NEXT:    [[DUMMY17_4:%.*]] = mul i32 [[DUMMY16_4]], [[DUMMY16_4]]
++; CHECK-NEXT:    [[DUMMY18_4:%.*]] = mul i32 [[DUMMY17_4]], [[DUMMY17_4]]
++; CHECK-NEXT:    [[DUMMY19_4:%.*]] = mul i32 [[DUMMY18_4]], [[DUMMY18_4]]
++; CHECK-NEXT:    [[DUMMY20_4:%.*]] = mul i32 [[DUMMY19_4]], [[DUMMY19_4]]
++; CHECK-NEXT:    [[DUMMY21_4:%.*]] = mul i32 [[DUMMY20_4]], [[DUMMY20_4]]
++; CHECK-NEXT:    [[DUMMY22_4:%.*]] = mul i32 [[DUMMY21_4]], [[DUMMY21_4]]
++; CHECK-NEXT:    [[DUMMY23_4:%.*]] = mul i32 [[DUMMY22_4]], [[DUMMY22_4]]
++; CHECK-NEXT:    [[DUMMY24_4:%.*]] = mul i32 [[DUMMY23_4]], [[DUMMY23_4]]
++; CHECK-NEXT:    [[DUMMY25_4:%.*]] = mul i32 [[DUMMY24_4]], [[DUMMY24_4]]
++; CHECK-NEXT:    [[DUMMY26_4:%.*]] = mul i32 [[DUMMY25_4]], [[DUMMY25_4]]
++; CHECK-NEXT:    [[DUMMY27_4:%.*]] = mul i32 [[DUMMY26_4]], [[DUMMY26_4]]
++; CHECK-NEXT:    [[DUMMY28_4:%.*]] = mul i32 [[DUMMY27_4]], [[DUMMY27_4]]
++; CHECK-NEXT:    [[DUMMY29_4:%.*]] = mul i32 [[DUMMY28_4]], [[DUMMY28_4]]
++; CHECK-NEXT:    [[DUMMY30_4:%.*]] = mul i32 [[DUMMY29_4]], [[DUMMY29_4]]
++; CHECK-NEXT:    [[DUMMY31_4:%.*]] = mul i32 [[DUMMY30_4]], [[DUMMY30_4]]
++; CHECK-NEXT:    [[DUMMY32_4:%.*]] = mul i32 [[DUMMY31_4]], [[DUMMY31_4]]
++; CHECK-NEXT:    [[DUMMY33_4:%.*]] = mul i32 [[DUMMY32_4]], [[DUMMY32_4]]
++; CHECK-NEXT:    [[DUMMY34_4:%.*]] = mul i32 [[DUMMY33_4]], [[DUMMY33_4]]
++; CHECK-NEXT:    [[DUMMY35_4:%.*]] = mul i32 [[DUMMY34_4]], [[DUMMY34_4]]
++; CHECK-NEXT:    [[DUMMY36_4:%.*]] = mul i32 [[DUMMY35_4]], [[DUMMY35_4]]
++; CHECK-NEXT:    [[DUMMY37_4:%.*]] = mul i32 [[DUMMY36_4]], [[DUMMY36_4]]
++; CHECK-NEXT:    [[DUMMY38_4:%.*]] = mul i32 [[DUMMY37_4]], [[DUMMY37_4]]
++; CHECK-NEXT:    [[DUMMY39_4:%.*]] = mul i32 [[DUMMY38_4]], [[DUMMY38_4]]
++; CHECK-NEXT:    [[DUMMY40_4:%.*]] = mul i32 [[DUMMY39_4]], [[DUMMY39_4]]
++; CHECK-NEXT:    [[DUMMY41_4:%.*]] = mul i32 [[DUMMY40_4]], [[DUMMY40_4]]
++; CHECK-NEXT:    [[DUMMY42_4:%.*]] = mul i32 [[DUMMY41_4]], [[DUMMY41_4]]
++; CHECK-NEXT:    [[DUMMY43_4:%.*]] = mul i32 [[DUMMY42_4]], [[DUMMY42_4]]
++; CHECK-NEXT:    [[DUMMY44_4:%.*]] = mul i32 [[DUMMY43_4]], [[DUMMY43_4]]
++; CHECK-NEXT:    [[DUMMY45_4:%.*]] = mul i32 [[DUMMY44_4]], [[DUMMY44_4]]
++; CHECK-NEXT:    [[DUMMY46_4:%.*]] = mul i32 [[DUMMY45_4]], [[DUMMY45_4]]
++; CHECK-NEXT:    [[DUMMY47_4:%.*]] = mul i32 [[DUMMY46_4]], [[DUMMY46_4]]
++; CHECK-NEXT:    [[DUMMY48_4:%.*]] = mul i32 [[DUMMY47_4]], [[DUMMY47_4]]
++; CHECK-NEXT:    [[DUMMY49_4:%.*]] = mul i32 [[DUMMY48_4]], [[DUMMY48_4]]
++; CHECK-NEXT:    [[DUMMY50_4:%.*]] = mul i32 [[DUMMY49_4]], [[DUMMY49_4]]
++; CHECK-NEXT:    [[SUM_NEXT_4:%.*]] = add nsw i32 [[DUMMY50_4]], [[SUM_NEXT_3]]
++; CHECK-NEXT:    [[INDVARS_IV_NEXT_4:%.*]] = add nuw nsw i64 [[INDVARS_IV]], 5
++; CHECK-NEXT:    [[ARRAYIDX_5:%.*]] = getelementptr inbounds i32, ptr [[ARY]], i64 [[INDVARS_IV_NEXT_4]]
++; CHECK-NEXT:    [[VAL_5:%.*]] = load i32, ptr [[ARRAYIDX_5]], align 4
++; CHECK-NEXT:    [[DUMMY1_5:%.*]] = mul i32 [[VAL_5]], [[VAL_5]]
++; CHECK-NEXT:    [[DUMMY2_5:%.*]] = mul i32 [[DUMMY1_5]], [[DUMMY1_5]]
++; CHECK-NEXT:    [[DUMMY3_5:%.*]] = mul i32 [[DUMMY2_5]], [[DUMMY2_5]]
++; CHECK-NEXT:    [[DUMMY4_5:%.*]] = mul i32 [[DUMMY3_5]], [[DUMMY3_5]]
++; CHECK-NEXT:    [[DUMMY5_5:%.*]] = mul i32 [[DUMMY4_5]], [[DUMMY4_5]]
++; CHECK-NEXT:    [[DUMMY6_5:%.*]] = mul i32 [[DUMMY5_5]], [[DUMMY5_5]]
++; CHECK-NEXT:    [[DUMMY7_5:%.*]] = mul i32 [[DUMMY6_5]], [[DUMMY6_5]]
++; CHECK-NEXT:    [[DUMMY8_5:%.*]] = mul i32 [[DUMMY7_5]], [[DUMMY7_5]]
++; CHECK-NEXT:    [[DUMMY9_5:%.*]] = mul i32 [[DUMMY8_5]], [[DUMMY8_5]]
++; CHECK-NEXT:    [[DUMMY10_5:%.*]] = mul i32 [[DUMMY9_5]], [[DUMMY9_5]]
++; CHECK-NEXT:    [[DUMMY11_5:%.*]] = mul i32 [[DUMMY10_5]], [[DUMMY10_5]]
++; CHECK-NEXT:    [[DUMMY12_5:%.*]] = mul i32 [[DUMMY11_5]], [[DUMMY11_5]]
++; CHECK-NEXT:    [[DUMMY13_5:%.*]] = mul i32 [[DUMMY12_5]], [[DUMMY12_5]]
++; CHECK-NEXT:    [[DUMMY14_5:%.*]] = mul i32 [[DUMMY13_5]], [[DUMMY13_5]]
++; CHECK-NEXT:    [[DUMMY15_5:%.*]] = mul i32 [[DUMMY14_5]], [[DUMMY14_5]]
++; CHECK-NEXT:    [[DUMMY16_5:%.*]] = mul i32 [[DUMMY15_5]], [[DUMMY15_5]]
++; CHECK-NEXT:    [[DUMMY17_5:%.*]] = mul i32 [[DUMMY16_5]], [[DUMMY16_5]]
++; CHECK-NEXT:    [[DUMMY18_5:%.*]] = mul i32 [[DUMMY17_5]], [[DUMMY17_5]]
++; CHECK-NEXT:    [[DUMMY19_5:%.*]] = mul i32 [[DUMMY18_5]], [[DUMMY18_5]]
++; CHECK-NEXT:    [[DUMMY20_5:%.*]] = mul i32 [[DUMMY19_5]], [[DUMMY19_5]]
++; CHECK-NEXT:    [[DUMMY21_5:%.*]] = mul i32 [[DUMMY20_5]], [[DUMMY20_5]]
++; CHECK-NEXT:    [[DUMMY22_5:%.*]] = mul i32 [[DUMMY21_5]], [[DUMMY21_5]]
++; CHECK-NEXT:    [[DUMMY23_5:%.*]] = mul i32 [[DUMMY22_5]], [[DUMMY22_5]]
++; CHECK-NEXT:    [[DUMMY24_5:%.*]] = mul i32 [[DUMMY23_5]], [[DUMMY23_5]]
++; CHECK-NEXT:    [[DUMMY25_5:%.*]] = mul i32 [[DUMMY24_5]], [[DUMMY24_5]]
++; CHECK-NEXT:    [[DUMMY26_5:%.*]] = mul i32 [[DUMMY25_5]], [[DUMMY25_5]]
++; CHECK-NEXT:    [[DUMMY27_5:%.*]] = mul i32 [[DUMMY26_5]], [[DUMMY26_5]]
++; CHECK-NEXT:    [[DUMMY28_5:%.*]] = mul i32 [[DUMMY27_5]], [[DUMMY27_5]]
++; CHECK-NEXT:    [[DUMMY29_5:%.*]] = mul i32 [[DUMMY28_5]], [[DUMMY28_5]]
++; CHECK-NEXT:    [[DUMMY30_5:%.*]] = mul i32 [[DUMMY29_5]], [[DUMMY29_5]]
++; CHECK-NEXT:    [[DUMMY31_5:%.*]] = mul i32 [[DUMMY30_5]], [[DUMMY30_5]]
++; CHECK-NEXT:    [[DUMMY32_5:%.*]] = mul i32 [[DUMMY31_5]], [[DUMMY31_5]]
++; CHECK-NEXT:    [[DUMMY33_5:%.*]] = mul i32 [[DUMMY32_5]], [[DUMMY32_5]]
++; CHECK-NEXT:    [[DUMMY34_5:%.*]] = mul i32 [[DUMMY33_5]], [[DUMMY33_5]]
++; CHECK-NEXT:    [[DUMMY35_5:%.*]] = mul i32 [[DUMMY34_5]], [[DUMMY34_5]]
++; CHECK-NEXT:    [[DUMMY36_5:%.*]] = mul i32 [[DUMMY35_5]], [[DUMMY35_5]]
++; CHECK-NEXT:    [[DUMMY37_5:%.*]] = mul i32 [[DUMMY36_5]], [[DUMMY36_5]]
++; CHECK-NEXT:    [[DUMMY38_5:%.*]] = mul i32 [[DUMMY37_5]], [[DUMMY37_5]]
++; CHECK-NEXT:    [[DUMMY39_5:%.*]] = mul i32 [[DUMMY38_5]], [[DUMMY38_5]]
++; CHECK-NEXT:    [[DUMMY40_5:%.*]] = mul i32 [[DUMMY39_5]], [[DUMMY39_5]]
++; CHECK-NEXT:    [[DUMMY41_5:%.*]] = mul i32 [[DUMMY40_5]], [[DUMMY40_5]]
++; CHECK-NEXT:    [[DUMMY42_5:%.*]] = mul i32 [[DUMMY41_5]], [[DUMMY41_5]]
++; CHECK-NEXT:    [[DUMMY43_5:%.*]] = mul i32 [[DUMMY42_5]], [[DUMMY42_5]]
++; CHECK-NEXT:    [[DUMMY44_5:%.*]] = mul i32 [[DUMMY43_5]], [[DUMMY43_5]]
++; CHECK-NEXT:    [[DUMMY45_5:%.*]] = mul i32 [[DUMMY44_5]], [[DUMMY44_5]]
++; CHECK-NEXT:    [[DUMMY46_5:%.*]] = mul i32 [[DUMMY45_5]], [[DUMMY45_5]]
++; CHECK-NEXT:    [[DUMMY47_5:%.*]] = mul i32 [[DUMMY46_5]], [[DUMMY46_5]]
++; CHECK-NEXT:    [[DUMMY48_5:%.*]] = mul i32 [[DUMMY47_5]], [[DUMMY47_5]]
++; CHECK-NEXT:    [[DUMMY49_5:%.*]] = mul i32 [[DUMMY48_5]], [[DUMMY48_5]]
++; CHECK-NEXT:    [[DUMMY50_5:%.*]] = mul i32 [[DUMMY49_5]], [[DUMMY49_5]]
++; CHECK-NEXT:    [[SUM_NEXT_5:%.*]] = add nsw i32 [[DUMMY50_5]], [[SUM_NEXT_4]]
++; CHECK-NEXT:    [[INDVARS_IV_NEXT_5:%.*]] = add nuw nsw i64 [[INDVARS_IV]], 6
++; CHECK-NEXT:    [[ARRAYIDX_6:%.*]] = getelementptr inbounds i32, ptr [[ARY]], i64 [[INDVARS_IV_NEXT_5]]
++; CHECK-NEXT:    [[VAL_6:%.*]] = load i32, ptr [[ARRAYIDX_6]], align 4
++; CHECK-NEXT:    [[DUMMY1_6:%.*]] = mul i32 [[VAL_6]], [[VAL_6]]
++; CHECK-NEXT:    [[DUMMY2_6:%.*]] = mul i32 [[DUMMY1_6]], [[DUMMY1_6]]
++; CHECK-NEXT:    [[DUMMY3_6:%.*]] = mul i32 [[DUMMY2_6]], [[DUMMY2_6]]
++; CHECK-NEXT:    [[DUMMY4_6:%.*]] = mul i32 [[DUMMY3_6]], [[DUMMY3_6]]
++; CHECK-NEXT:    [[DUMMY5_6:%.*]] = mul i32 [[DUMMY4_6]], [[DUMMY4_6]]
++; CHECK-NEXT:    [[DUMMY6_6:%.*]] = mul i32 [[DUMMY5_6]], [[DUMMY5_6]]
++; CHECK-NEXT:    [[DUMMY7_6:%.*]] = mul i32 [[DUMMY6_6]], [[DUMMY6_6]]
++; CHECK-NEXT:    [[DUMMY8_6:%.*]] = mul i32 [[DUMMY7_6]], [[DUMMY7_6]]
++; CHECK-NEXT:    [[DUMMY9_6:%.*]] = mul i32 [[DUMMY8_6]], [[DUMMY8_6]]
++; CHECK-NEXT:    [[DUMMY10_6:%.*]] = mul i32 [[DUMMY9_6]], [[DUMMY9_6]]
++; CHECK-NEXT:    [[DUMMY11_6:%.*]] = mul i32 [[DUMMY10_6]], [[DUMMY10_6]]
++; CHECK-NEXT:    [[DUMMY12_6:%.*]] = mul i32 [[DUMMY11_6]], [[DUMMY11_6]]
++; CHECK-NEXT:    [[DUMMY13_6:%.*]] = mul i32 [[DUMMY12_6]], [[DUMMY12_6]]
++; CHECK-NEXT:    [[DUMMY14_6:%.*]] = mul i32 [[DUMMY13_6]], [[DUMMY13_6]]
++; CHECK-NEXT:    [[DUMMY15_6:%.*]] = mul i32 [[DUMMY14_6]], [[DUMMY14_6]]
++; CHECK-NEXT:    [[DUMMY16_6:%.*]] = mul i32 [[DUMMY15_6]], [[DUMMY15_6]]
++; CHECK-NEXT:    [[DUMMY17_6:%.*]] = mul i32 [[DUMMY16_6]], [[DUMMY16_6]]
++; CHECK-NEXT:    [[DUMMY18_6:%.*]] = mul i32 [[DUMMY17_6]], [[DUMMY17_6]]
++; CHECK-NEXT:    [[DUMMY19_6:%.*]] = mul i32 [[DUMMY18_6]], [[DUMMY18_6]]
++; CHECK-NEXT:    [[DUMMY20_6:%.*]] = mul i32 [[DUMMY19_6]], [[DUMMY19_6]]
++; CHECK-NEXT:    [[DUMMY21_6:%.*]] = mul i32 [[DUMMY20_6]], [[DUMMY20_6]]
++; CHECK-NEXT:    [[DUMMY22_6:%.*]] = mul i32 [[DUMMY21_6]], [[DUMMY21_6]]
++; CHECK-NEXT:    [[DUMMY23_6:%.*]] = mul i32 [[DUMMY22_6]], [[DUMMY22_6]]
++; CHECK-NEXT:    [[DUMMY24_6:%.*]] = mul i32 [[DUMMY23_6]], [[DUMMY23_6]]
++; CHECK-NEXT:    [[DUMMY25_6:%.*]] = mul i32 [[DUMMY24_6]], [[DUMMY24_6]]
++; CHECK-NEXT:    [[DUMMY26_6:%.*]] = mul i32 [[DUMMY25_6]], [[DUMMY25_6]]
++; CHECK-NEXT:    [[DUMMY27_6:%.*]] = mul i32 [[DUMMY26_6]], [[DUMMY26_6]]
++; CHECK-NEXT:    [[DUMMY28_6:%.*]] = mul i32 [[DUMMY27_6]], [[DUMMY27_6]]
++; CHECK-NEXT:    [[DUMMY29_6:%.*]] = mul i32 [[DUMMY28_6]], [[DUMMY28_6]]
++; CHECK-NEXT:    [[DUMMY30_6:%.*]] = mul i32 [[DUMMY29_6]], [[DUMMY29_6]]
++; CHECK-NEXT:    [[DUMMY31_6:%.*]] = mul i32 [[DUMMY30_6]], [[DUMMY30_6]]
++; CHECK-NEXT:    [[DUMMY32_6:%.*]] = mul i32 [[DUMMY31_6]], [[DUMMY31_6]]
++; CHECK-NEXT:    [[DUMMY33_6:%.*]] = mul i32 [[DUMMY32_6]], [[DUMMY32_6]]
++; CHECK-NEXT:    [[DUMMY34_6:%.*]] = mul i32 [[DUMMY33_6]], [[DUMMY33_6]]
++; CHECK-NEXT:    [[DUMMY35_6:%.*]] = mul i32 [[DUMMY34_6]], [[DUMMY34_6]]
++; CHECK-NEXT:    [[DUMMY36_6:%.*]] = mul i32 [[DUMMY35_6]], [[DUMMY35_6]]
++; CHECK-NEXT:    [[DUMMY37_6:%.*]] = mul i32 [[DUMMY36_6]], [[DUMMY36_6]]
++; CHECK-NEXT:    [[DUMMY38_6:%.*]] = mul i32 [[DUMMY37_6]], [[DUMMY37_6]]
++; CHECK-NEXT:    [[DUMMY39_6:%.*]] = mul i32 [[DUMMY38_6]], [[DUMMY38_6]]
++; CHECK-NEXT:    [[DUMMY40_6:%.*]] = mul i32 [[DUMMY39_6]], [[DUMMY39_6]]
++; CHECK-NEXT:    [[DUMMY41_6:%.*]] = mul i32 [[DUMMY40_6]], [[DUMMY40_6]]
++; CHECK-NEXT:    [[DUMMY42_6:%.*]] = mul i32 [[DUMMY41_6]], [[DUMMY41_6]]
++; CHECK-NEXT:    [[DUMMY43_6:%.*]] = mul i32 [[DUMMY42_6]], [[DUMMY42_6]]
++; CHECK-NEXT:    [[DUMMY44_6:%.*]] = mul i32 [[DUMMY43_6]], [[DUMMY43_6]]
++; CHECK-NEXT:    [[DUMMY45_6:%.*]] = mul i32 [[DUMMY44_6]], [[DUMMY44_6]]
++; CHECK-NEXT:    [[DUMMY46_6:%.*]] = mul i32 [[DUMMY45_6]], [[DUMMY45_6]]
++; CHECK-NEXT:    [[DUMMY47_6:%.*]] = mul i32 [[DUMMY46_6]], [[DUMMY46_6]]
++; CHECK-NEXT:    [[DUMMY48_6:%.*]] = mul i32 [[DUMMY47_6]], [[DUMMY47_6]]
++; CHECK-NEXT:    [[DUMMY49_6:%.*]] = mul i32 [[DUMMY48_6]], [[DUMMY48_6]]
++; CHECK-NEXT:    [[DUMMY50_6:%.*]] = mul i32 [[DUMMY49_6]], [[DUMMY49_6]]
++; CHECK-NEXT:    [[SUM_NEXT_6:%.*]] = add nsw i32 [[DUMMY50_6]], [[SUM_NEXT_5]]
++; CHECK-NEXT:    [[INDVARS_IV_NEXT_6:%.*]] = add nuw nsw i64 [[INDVARS_IV]], 7
++; CHECK-NEXT:    [[ARRAYIDX_7:%.*]] = getelementptr inbounds i32, ptr [[ARY]], i64 [[INDVARS_IV_NEXT_6]]
++; CHECK-NEXT:    [[VAL_7:%.*]] = load i32, ptr [[ARRAYIDX_7]], align 4
++; CHECK-NEXT:    [[DUMMY1_7:%.*]] = mul i32 [[VAL_7]], [[VAL_7]]
++; CHECK-NEXT:    [[DUMMY2_7:%.*]] = mul i32 [[DUMMY1_7]], [[DUMMY1_7]]
++; CHECK-NEXT:    [[DUMMY3_7:%.*]] = mul i32 [[DUMMY2_7]], [[DUMMY2_7]]
++; CHECK-NEXT:    [[DUMMY4_7:%.*]] = mul i32 [[DUMMY3_7]], [[DUMMY3_7]]
++; CHECK-NEXT:    [[DUMMY5_7:%.*]] = mul i32 [[DUMMY4_7]], [[DUMMY4_7]]
++; CHECK-NEXT:    [[DUMMY6_7:%.*]] = mul i32 [[DUMMY5_7]], [[DUMMY5_7]]
++; CHECK-NEXT:    [[DUMMY7_7:%.*]] = mul i32 [[DUMMY6_7]], [[DUMMY6_7]]
++; CHECK-NEXT:    [[DUMMY8_7:%.*]] = mul i32 [[DUMMY7_7]], [[DUMMY7_7]]
++; CHECK-NEXT:    [[DUMMY9_7:%.*]] = mul i32 [[DUMMY8_7]], [[DUMMY8_7]]
++; CHECK-NEXT:    [[DUMMY10_7:%.*]] = mul i32 [[DUMMY9_7]], [[DUMMY9_7]]
++; CHECK-NEXT:    [[DUMMY11_7:%.*]] = mul i32 [[DUMMY10_7]], [[DUMMY10_7]]
++; CHECK-NEXT:    [[DUMMY12_7:%.*]] = mul i32 [[DUMMY11_7]], [[DUMMY11_7]]
++; CHECK-NEXT:    [[DUMMY13_7:%.*]] = mul i32 [[DUMMY12_7]], [[DUMMY12_7]]
++; CHECK-NEXT:    [[DUMMY14_7:%.*]] = mul i32 [[DUMMY13_7]], [[DUMMY13_7]]
++; CHECK-NEXT:    [[DUMMY15_7:%.*]] = mul i32 [[DUMMY14_7]], [[DUMMY14_7]]
++; CHECK-NEXT:    [[DUMMY16_7:%.*]] = mul i32 [[DUMMY15_7]], [[DUMMY15_7]]
++; CHECK-NEXT:    [[DUMMY17_7:%.*]] = mul i32 [[DUMMY16_7]], [[DUMMY16_7]]
++; CHECK-NEXT:    [[DUMMY18_7:%.*]] = mul i32 [[DUMMY17_7]], [[DUMMY17_7]]
++; CHECK-NEXT:    [[DUMMY19_7:%.*]] = mul i32 [[DUMMY18_7]], [[DUMMY18_7]]
++; CHECK-NEXT:    [[DUMMY20_7:%.*]] = mul i32 [[DUMMY19_7]], [[DUMMY19_7]]
++; CHECK-NEXT:    [[DUMMY21_7:%.*]] = mul i32 [[DUMMY20_7]], [[DUMMY20_7]]
++; CHECK-NEXT:    [[DUMMY22_7:%.*]] = mul i32 [[DUMMY21_7]], [[DUMMY21_7]]
++; CHECK-NEXT:    [[DUMMY23_7:%.*]] = mul i32 [[DUMMY22_7]], [[DUMMY22_7]]
++; CHECK-NEXT:    [[DUMMY24_7:%.*]] = mul i32 [[DUMMY23_7]], [[DUMMY23_7]]
++; CHECK-NEXT:    [[DUMMY25_7:%.*]] = mul i32 [[DUMMY24_7]], [[DUMMY24_7]]
++; CHECK-NEXT:    [[DUMMY26_7:%.*]] = mul i32 [[DUMMY25_7]], [[DUMMY25_7]]
++; CHECK-NEXT:    [[DUMMY27_7:%.*]] = mul i32 [[DUMMY26_7]], [[DUMMY26_7]]
++; CHECK-NEXT:    [[DUMMY28_7:%.*]] = mul i32 [[DUMMY27_7]], [[DUMMY27_7]]
++; CHECK-NEXT:    [[DUMMY29_7:%.*]] = mul i32 [[DUMMY28_7]], [[DUMMY28_7]]
++; CHECK-NEXT:    [[DUMMY30_7:%.*]] = mul i32 [[DUMMY29_7]], [[DUMMY29_7]]
++; CHECK-NEXT:    [[DUMMY31_7:%.*]] = mul i32 [[DUMMY30_7]], [[DUMMY30_7]]
++; CHECK-NEXT:    [[DUMMY32_7:%.*]] = mul i32 [[DUMMY31_7]], [[DUMMY31_7]]
++; CHECK-NEXT:    [[DUMMY33_7:%.*]] = mul i32 [[DUMMY32_7]], [[DUMMY32_7]]
++; CHECK-NEXT:    [[DUMMY34_7:%.*]] = mul i32 [[DUMMY33_7]], [[DUMMY33_7]]
++; CHECK-NEXT:    [[DUMMY35_7:%.*]] = mul i32 [[DUMMY34_7]], [[DUMMY34_7]]
++; CHECK-NEXT:    [[DUMMY36_7:%.*]] = mul i32 [[DUMMY35_7]], [[DUMMY35_7]]
++; CHECK-NEXT:    [[DUMMY37_7:%.*]] = mul i32 [[DUMMY36_7]], [[DUMMY36_7]]
++; CHECK-NEXT:    [[DUMMY38_7:%.*]] = mul i32 [[DUMMY37_7]], [[DUMMY37_7]]
++; CHECK-NEXT:    [[DUMMY39_7:%.*]] = mul i32 [[DUMMY38_7]], [[DUMMY38_7]]
++; CHECK-NEXT:    [[DUMMY40_7:%.*]] = mul i32 [[DUMMY39_7]], [[DUMMY39_7]]
++; CHECK-NEXT:    [[DUMMY41_7:%.*]] = mul i32 [[DUMMY40_7]], [[DUMMY40_7]]
++; CHECK-NEXT:    [[DUMMY42_7:%.*]] = mul i32 [[DUMMY41_7]], [[DUMMY41_7]]
++; CHECK-NEXT:    [[DUMMY43_7:%.*]] = mul i32 [[DUMMY42_7]], [[DUMMY42_7]]
++; CHECK-NEXT:    [[DUMMY44_7:%.*]] = mul i32 [[DUMMY43_7]], [[DUMMY43_7]]
++; CHECK-NEXT:    [[DUMMY45_7:%.*]] = mul i32 [[DUMMY44_7]], [[DUMMY44_7]]
++; CHECK-NEXT:    [[DUMMY46_7:%.*]] = mul i32 [[DUMMY45_7]], [[DUMMY45_7]]
++; CHECK-NEXT:    [[DUMMY47_7:%.*]] = mul i32 [[DUMMY46_7]], [[DUMMY46_7]]
++; CHECK-NEXT:    [[DUMMY48_7:%.*]] = mul i32 [[DUMMY47_7]], [[DUMMY47_7]]
++; CHECK-NEXT:    [[DUMMY49_7:%.*]] = mul i32 [[DUMMY48_7]], [[DUMMY48_7]]
++; CHECK-NEXT:    [[DUMMY50_7:%.*]] = mul i32 [[DUMMY49_7]], [[DUMMY49_7]]
++; CHECK-NEXT:    [[SUM_NEXT_7]] = add nsw i32 [[DUMMY50_7]], [[SUM_NEXT_6]]
++; CHECK-NEXT:    [[INDVARS_IV_NEXT_7]] = add nuw nsw i64 [[INDVARS_IV]], 8
++; CHECK-NEXT:    [[NITER_NEXT_7]] = add i64 [[NITER]], 8
++; CHECK-NEXT:    [[NITER_NCMP_7:%.*]] = icmp eq i64 [[NITER_NEXT_7]], [[UNROLL_ITER]]
++; CHECK-NEXT:    br i1 [[NITER_NCMP_7]], label [[FOR_COND_CLEANUP_UNR_LCSSA_LOOPEXIT:%.*]], label [[FOR_BODY]]
+ ; CHECK:       for.cond.cleanup.unr-lcssa.loopexit:
+-; CHECK-NEXT:    [[SUM_NEXT_LCSSA_PH_PH:%.*]] = phi i32 [ [[SUM_NEXT_1]], [[FOR_BODY]] ]
+-; CHECK-NEXT:    [[INDVARS_IV_UNR_PH:%.*]] = phi i64 [ [[INDVARS_IV_NEXT_1]], [[FOR_BODY]] ]
+-; CHECK-NEXT:    [[SUM_UNR_PH:%.*]] = phi i32 [ [[SUM_NEXT_1]], [[FOR_BODY]] ]
++; CHECK-NEXT:    [[SUM_NEXT_LCSSA_PH_PH:%.*]] = phi i32 [ [[SUM_NEXT_7]], [[FOR_BODY]] ]
++; CHECK-NEXT:    [[INDVARS_IV_UNR_PH:%.*]] = phi i64 [ [[INDVARS_IV_NEXT_7]], [[FOR_BODY]] ]
++; CHECK-NEXT:    [[SUM_UNR_PH:%.*]] = phi i32 [ [[SUM_NEXT_7]], [[FOR_BODY]] ]
+ ; CHECK-NEXT:    br label [[FOR_COND_CLEANUP_UNR_LCSSA]]
+ ; CHECK:       for.cond.cleanup.unr-lcssa:
+ ; CHECK-NEXT:    [[SUM_NEXT_LCSSA_PH:%.*]] = phi i32 [ undef, [[ENTRY:%.*]] ], [ [[SUM_NEXT_LCSSA_PH_PH]], [[FOR_COND_CLEANUP_UNR_LCSSA_LOOPEXIT]] ]
+@@ -302,7 +1010,10 @@
+ ; CHECK:       for.body.epil.preheader:
+ ; CHECK-NEXT:    br label [[FOR_BODY_EPIL:%.*]]
+ ; CHECK:       for.body.epil:
+-; CHECK-NEXT:    [[ARRAYIDX_EPIL:%.*]] = getelementptr inbounds i32, ptr [[ARY]], i64 [[INDVARS_IV_UNR]]
++; CHECK-NEXT:    [[INDVARS_IV_EPIL:%.*]] = phi i64 [ [[INDVARS_IV_UNR]], [[FOR_BODY_EPIL_PREHEADER]] ], [ [[INDVARS_IV_NEXT_EPIL:%.*]], [[FOR_BODY_EPIL]] ]
++; CHECK-NEXT:    [[SUM_EPIL:%.*]] = phi i32 [ [[SUM_UNR]], [[FOR_BODY_EPIL_PREHEADER]] ], [ [[SUM_NEXT_EPIL:%.*]], [[FOR_BODY_EPIL]] ]
++; CHECK-NEXT:    [[EPIL_ITER:%.*]] = phi i64 [ 0, [[FOR_BODY_EPIL_PREHEADER]] ], [ [[EPIL_ITER_NEXT:%.*]], [[FOR_BODY_EPIL]] ]
++; CHECK-NEXT:    [[ARRAYIDX_EPIL:%.*]] = getelementptr inbounds i32, ptr [[ARY]], i64 [[INDVARS_IV_EPIL]]
+ ; CHECK-NEXT:    [[VAL_EPIL:%.*]] = load i32, ptr [[ARRAYIDX_EPIL]], align 4
+ ; CHECK-NEXT:    [[DUMMY1_EPIL:%.*]] = mul i32 [[VAL_EPIL]], [[VAL_EPIL]]
+ ; CHECK-NEXT:    [[DUMMY2_EPIL:%.*]] = mul i32 [[DUMMY1_EPIL]], [[DUMMY1_EPIL]]
+@@ -354,10 +1065,17 @@
+ ; CHECK-NEXT:    [[DUMMY48_EPIL:%.*]] = mul i32 [[DUMMY47_EPIL]], [[DUMMY47_EPIL]]
+ ; CHECK-NEXT:    [[DUMMY49_EPIL:%.*]] = mul i32 [[DUMMY48_EPIL]], [[DUMMY48_EPIL]]
+ ; CHECK-NEXT:    [[DUMMY50_EPIL:%.*]] = mul i32 [[DUMMY49_EPIL]], [[DUMMY49_EPIL]]
+-; CHECK-NEXT:    [[SUM_NEXT_EPIL:%.*]] = add nsw i32 [[DUMMY50_EPIL]], [[SUM_UNR]]
++; CHECK-NEXT:    [[SUM_NEXT_EPIL]] = add nsw i32 [[DUMMY50_EPIL]], [[SUM_EPIL]]
++; CHECK-NEXT:    [[INDVARS_IV_NEXT_EPIL]] = add nuw nsw i64 [[INDVARS_IV_EPIL]], 1
++; CHECK-NEXT:    [[EXITCOND_NOT_EPIL:%.*]] = icmp eq i64 [[INDVARS_IV_NEXT_EPIL]], [[N]]
++; CHECK-NEXT:    [[EPIL_ITER_NEXT]] = add i64 [[EPIL_ITER]], 1
++; CHECK-NEXT:    [[EPIL_ITER_CMP:%.*]] = icmp ne i64 [[EPIL_ITER_NEXT]], [[XTRAITER]]
++; CHECK-NEXT:    br i1 [[EPIL_ITER_CMP]], label [[FOR_BODY_EPIL]], label [[FOR_COND_CLEANUP_EPILOG_LCSSA:%.*]], !llvm.loop [[LOOP0:![0-9]+]]
++; CHECK:       for.cond.cleanup.epilog-lcssa:
++; CHECK-NEXT:    [[SUM_NEXT_LCSSA_PH1:%.*]] = phi i32 [ [[SUM_NEXT_EPIL]], [[FOR_BODY_EPIL]] ]
+ ; CHECK-NEXT:    br label [[FOR_COND_CLEANUP]]
+ ; CHECK:       for.cond.cleanup:
+-; CHECK-NEXT:    [[SUM_NEXT_LCSSA:%.*]] = phi i32 [ [[SUM_NEXT_LCSSA_PH]], [[FOR_COND_CLEANUP_UNR_LCSSA]] ], [ [[SUM_NEXT_EPIL]], [[FOR_BODY_EPIL]] ]
++; CHECK-NEXT:    [[SUM_NEXT_LCSSA:%.*]] = phi i32 [ [[SUM_NEXT_LCSSA_PH]], [[FOR_COND_CLEANUP_UNR_LCSSA]] ], [ [[SUM_NEXT_LCSSA_PH1]], [[FOR_COND_CLEANUP_EPILOG_LCSSA]] ]
+ ; CHECK-NEXT:    ret i32 [[SUM_NEXT_LCSSA]]
+ ;
+ entry:
diff --git a/third_party/llvm/workspace.bzl b/third_party/llvm/workspace.bzl
index 92ab0a674b5fc1..632e27d09333f5 100644
--- a/third_party/llvm/workspace.bzl
+++ b/third_party/llvm/workspace.bzl
@@ -4,8 +4,8 @@ load("//third_party:repo.bzl", "tf_http_archive")
 
 def repo(name):
     """Imports LLVM."""
-    LLVM_COMMIT = "e1ed138a67a92ef1ff0214ca198094be13045090"
-    LLVM_SHA256 = "3a55ccdd3e90334b1181bb4950972cf005a2b99641cd2f0bc153d91d3a62a727"
+    LLVM_COMMIT = "e27f9bb31984927d10995e3d7e5a8fcde60cf5d4"
+    LLVM_SHA256 = "bba7c0ab0932cdb31b4ebf93f46bf5509fb173aa691d8a822e4717b001cb9791"
 
     tf_http_archive(
         name = name,

From 41b7a8b2d83292c6e216a5faf096c45a7bdd4efe Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Thu, 16 May 2024 09:18:42 -0700
Subject: [PATCH 274/478] Update ops-related pbtxt files.

PiperOrigin-RevId: 634424923
---
 tensorflow/core/ops/ops.pbtxt | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorflow/core/ops/ops.pbtxt b/tensorflow/core/ops/ops.pbtxt
index cdc0191b8ad6c9..2a7a7c62c9d010 100644
--- a/tensorflow/core/ops/ops.pbtxt
+++ b/tensorflow/core/ops/ops.pbtxt
@@ -1,4 +1,4 @@
-go/debugproto   
+go/debugstr    
 op {
   name: "Abort"
   attr {

From 0cdc6a5b538f221f8bc4fe87405cea61a7604102 Mon Sep 17 00:00:00 2001
From: George Karpenkov <cheshire@google.com>
Date: Thu, 16 May 2024 09:33:04 -0700
Subject: [PATCH 275/478] [XLA:GPU] [NFC] Better error message for slow
 autotuned kernels

PiperOrigin-RevId: 634430923
---
 third_party/xla/xla/service/gpu/gemm_fusion_autotuner.cc | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/third_party/xla/xla/service/gpu/gemm_fusion_autotuner.cc b/third_party/xla/xla/service/gpu/gemm_fusion_autotuner.cc
index dada3783df6b2e..1644433f48e779 100644
--- a/third_party/xla/xla/service/gpu/gemm_fusion_autotuner.cc
+++ b/third_party/xla/xla/service/gpu/gemm_fusion_autotuner.cc
@@ -904,7 +904,8 @@ absl::StatusOr<std::vector<AutotuneResult>> GemmFusionAutotunerImpl::Profile(
 
       VLOG(5) << "Running the kernel took: " << profiling_output->duration;
       if (profiling_output->duration >= absl::Seconds(1)) {
-        LOG(WARNING) << "Slow kernel for " << fusion.name()
+        LOG(WARNING) << "Slow kernel for "
+                     << fusion.called_computations()[0]->ToString()
                      << " took: " << profiling_output->duration << ". "
                      << ToString(candidate.config);
       }

From bdec6a7cba18167f5e91ce3fc7d1c246cfc67e84 Mon Sep 17 00:00:00 2001
From: Dateng Lin <datenglin@google.com>
Date: Thu, 16 May 2024 09:36:42 -0700
Subject: [PATCH 276/478] Removed all the visibilities to
 `//learning/infra/mira/distributed` in MLRT.

PiperOrigin-RevId: 634432048
---
 tensorflow/compiler/mlir/tfrt/ir/mlrt/BUILD         | 6 ------
 tensorflow/compiler/mlir/tfrt/transforms/mlrt/BUILD | 1 -
 2 files changed, 7 deletions(-)

diff --git a/tensorflow/compiler/mlir/tfrt/ir/mlrt/BUILD b/tensorflow/compiler/mlir/tfrt/ir/mlrt/BUILD
index 5bc69ce7a70758..374aad2a242d9b 100644
--- a/tensorflow/compiler/mlir/tfrt/ir/mlrt/BUILD
+++ b/tensorflow/compiler/mlir/tfrt/ir/mlrt/BUILD
@@ -14,7 +14,6 @@ td_library(
     includes = ["."],
     visibility = [
         # copybara:uncomment "//learning/brain/tfrt/mlir:__subpackages__",
-        "//learning/infra/mira/distributed:__subpackages__",
     ],
     deps = [
         "@llvm-project//mlir:OpBaseTdFiles",
@@ -51,7 +50,6 @@ cc_library(
     ],
     visibility = [
         # copybara:uncomment "//learning/brain/tfrt/mlir:__subpackages__",
-        "//learning/infra/mira/distributed:__subpackages__",
         "//tensorflow/compiler/mlir/tfrt:__subpackages__",
     ],
     deps = [
@@ -71,9 +69,6 @@ td_library(
         "tf_ops.td",
     ],
     includes = ["."],
-    visibility = [
-        # copybara:uncomment "//learning/infra/mira/distributed:__subpackages__",
-    ],
     deps = [
         ":mlrt_td_files",
         "//tensorflow/compiler/mlir/tensorflow:tensorflow_ops_td_files",
@@ -156,7 +151,6 @@ cc_library(
     hdrs = ["tf_mlrt_ops.h"],
     visibility = [
         # copybara:uncomment "//learning/brain/experimental/tfrt/mlrt/application/tensorflow/tests:__subpackages__",
-        # copybara:uncomment "//learning/infra/mira/distributed:__subpackages__",
         "//tensorflow/compiler/mlir/tfrt:__subpackages__",
     ],
     deps = [
diff --git a/tensorflow/compiler/mlir/tfrt/transforms/mlrt/BUILD b/tensorflow/compiler/mlir/tfrt/transforms/mlrt/BUILD
index d7fafb49ee6cdd..ed518285828d1a 100644
--- a/tensorflow/compiler/mlir/tfrt/transforms/mlrt/BUILD
+++ b/tensorflow/compiler/mlir/tfrt/transforms/mlrt/BUILD
@@ -1,7 +1,6 @@
 package(
     # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
     default_visibility = [
-        # copybara:uncomment "//learning/infra/mira/distributed:__subpackages__",
         "//tensorflow/compiler/mlir/tfrt:__subpackages__",
         "//tensorflow/core/tfrt:__subpackages__",
     ],

From 0a7be3e5715e5cfaf232d081fbf1bbd232be37be Mon Sep 17 00:00:00 2001
From: Kuy Mainwaring <kuym@google.com>
Date: Thu, 16 May 2024 10:10:17 -0700
Subject: [PATCH 277/478] [XLA:GPU] #includes cleanup for
 xla/service/gpu/tests/gemm_rewrite_test.cc, nop_custom_call_test.cc

PiperOrigin-RevId: 634446511
---
 third_party/xla/xla/service/gpu/tests/BUILD                 | 4 ++--
 third_party/xla/xla/service/gpu/tests/gemm_rewrite_test.cc  | 6 ++++--
 .../xla/xla/service/gpu/tests/nop_custom_call_test.cc       | 2 ++
 3 files changed, 8 insertions(+), 4 deletions(-)

diff --git a/third_party/xla/xla/service/gpu/tests/BUILD b/third_party/xla/xla/service/gpu/tests/BUILD
index 257e3ddf87973f..11819ccfeba7b5 100644
--- a/third_party/xla/xla/service/gpu/tests/BUILD
+++ b/third_party/xla/xla/service/gpu/tests/BUILD
@@ -177,16 +177,16 @@ xla_test(
         "//xla/service:buffer_assignment",
         "//xla/service:executable",
         "//xla/service:hlo_module_config",
+        "//xla/service:hlo_pass",
         "//xla/service:pattern_matcher",
         "//xla/service:pattern_matcher_gmock",
         "//xla/service/gpu:gemm_rewriter",
         "//xla/service/gpu:gpu_executable",
-        "//xla/service/gpu:variant_visitor",
         "//xla/stream_executor:device_description",
         "//xla/stream_executor:device_memory_allocator",
         "//xla/tests:filecheck",
+        "//xla/tests:verified_hlo_module",
         "@com_google_absl//absl/container:flat_hash_map",
-        "@com_google_absl//absl/functional:any_invocable",
         "@com_google_absl//absl/status:statusor",
         "@com_google_absl//absl/strings",
         "@com_google_absl//absl/types:span",
diff --git a/third_party/xla/xla/service/gpu/tests/gemm_rewrite_test.cc b/third_party/xla/xla/service/gpu/tests/gemm_rewrite_test.cc
index 5c5fc750b101d6..cca4fc6aa957c3 100644
--- a/third_party/xla/xla/service/gpu/tests/gemm_rewrite_test.cc
+++ b/third_party/xla/xla/service/gpu/tests/gemm_rewrite_test.cc
@@ -15,7 +15,9 @@ limitations under the License.
 
 #include <array>
 #include <cstdint>
+#include <functional>
 #include <memory>
+#include <optional>
 #include <string>
 #include <tuple>
 #include <utility>
@@ -23,7 +25,6 @@ limitations under the License.
 #include <vector>
 
 #include "absl/container/flat_hash_map.h"
-#include "absl/functional/any_invocable.h"
 #include "absl/status/statusor.h"
 #include "absl/strings/str_replace.h"
 #include "absl/strings/string_view.h"
@@ -37,14 +38,15 @@ limitations under the License.
 #include "xla/service/gpu/gemm_rewriter.h"
 #include "xla/service/gpu/gpu_executable.h"
 #include "xla/service/gpu/tests/gpu_codegen_test.h"
-#include "xla/service/gpu/variant_visitor.h"
 #include "xla/service/hlo_module_config.h"
+#include "xla/service/hlo_pass_interface.h"
 #include "xla/service/pattern_matcher.h"
 #include "xla/service/pattern_matcher_gmock.h"
 #include "xla/stream_executor/device_description.h"
 #include "xla/stream_executor/device_memory_allocator.h"
 #include "xla/test.h"
 #include "xla/tests/filecheck.h"
+#include "xla/tests/verified_hlo_module.h"
 #include "xla/xla.pb.h"
 #include "tsl/lib/core/status_test_util.h"
 #include "tsl/platform/statusor.h"
diff --git a/third_party/xla/xla/service/gpu/tests/nop_custom_call_test.cc b/third_party/xla/xla/service/gpu/tests/nop_custom_call_test.cc
index 0da317ef3d1398..d979d18aa8ac9d 100644
--- a/third_party/xla/xla/service/gpu/tests/nop_custom_call_test.cc
+++ b/third_party/xla/xla/service/gpu/tests/nop_custom_call_test.cc
@@ -13,6 +13,8 @@ See the License for the specific language governing permissions and
 limitations under the License.
 ==============================================================================*/
 
+#include <utility>
+
 #include "xla/tests/hlo_test_base.h"
 
 namespace xla {

From 7f174f1ca5e693163d94b8aeacb51514b67ac57a Mon Sep 17 00:00:00 2001
From: Kuy Mainwaring <kuym@google.com>
Date: Thu, 16 May 2024 10:10:37 -0700
Subject: [PATCH 278/478] [XLA:GPU] Clang-tidy cleanup for
 xla/service/gpu/model/indexing_test_utils.h

PiperOrigin-RevId: 634446637
---
 .../xla/xla/service/gpu/model/indexing_test_utils.cc        | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/third_party/xla/xla/service/gpu/model/indexing_test_utils.cc b/third_party/xla/xla/service/gpu/model/indexing_test_utils.cc
index 72becac80c075e..76073e48089755 100644
--- a/third_party/xla/xla/service/gpu/model/indexing_test_utils.cc
+++ b/third_party/xla/xla/service/gpu/model/indexing_test_utils.cc
@@ -18,8 +18,10 @@ limitations under the License.
 #include <cctype>
 #include <cstddef>
 #include <string>
+#include <string_view>
+#include <utility>
+#include <vector>
 
-#include <gtest/gtest.h>
 #include "absl/container/flat_hash_set.h"
 #include "absl/log/check.h"
 #include "absl/strings/str_cat.h"
@@ -31,10 +33,8 @@ limitations under the License.
 #include "mlir/IR/MLIRContext.h"  // from @llvm-project
 #include "mlir/Support/LLVM.h"  // from @llvm-project
 #include "xla/hlo/ir/hlo_instruction.h"
-#include "xla/hlo/ir/hlo_opcode.h"
 #include "xla/service/gpu/model/indexing_analysis.h"
 #include "xla/service/gpu/model/indexing_map.h"
-#include "xla/tests/hlo_test_base.h"
 
 namespace xla {
 namespace gpu {

From 8fb62972655dd2a95d08eab99ea7bd879dd8abd7 Mon Sep 17 00:00:00 2001
From: Kuy Mainwaring <kuym@google.com>
Date: Thu, 16 May 2024 10:12:03 -0700
Subject: [PATCH 279/478] [XLA:GPU] Clang-tidy cleanup for
 xla/service/gpu/model/indexing_analysis.cc

PiperOrigin-RevId: 634447146
---
 third_party/xla/xla/service/gpu/model/indexing_analysis.cc | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/third_party/xla/xla/service/gpu/model/indexing_analysis.cc b/third_party/xla/xla/service/gpu/model/indexing_analysis.cc
index 2f9318c447ccaa..fd971af23bb3c6 100644
--- a/third_party/xla/xla/service/gpu/model/indexing_analysis.cc
+++ b/third_party/xla/xla/service/gpu/model/indexing_analysis.cc
@@ -1171,7 +1171,7 @@ IndexingMap GetIndexingMapFromPhysicalLayoutToLogical(
     const Shape& shape, MLIRContext* mlir_context) {
   if (shape.rank() == 0) {
     return IndexingMap(AffineMap::get(mlir_context),
-                       /*dim_vars=*/{}, /*range vars=*/{}, /*rt_vars=*/{});
+                       /*dimensions=*/{}, /*range vars=*/{}, /*rt_vars=*/{});
   }
   return IndexingMap::FromTensorSizes(
       ComputeTransposeIndexingMap(
@@ -1186,7 +1186,7 @@ IndexingMap GetIndexingMapFromLogicalToPhysicalLayout(
     const Shape& shape, MLIRContext* mlir_context) {
   if (shape.rank() == 0) {
     return IndexingMap(AffineMap::get(mlir_context),
-                       /*dim_vars=*/{}, /*range vars=*/{}, /*rt_vars=*/{});
+                       /*dimensions=*/{}, /*range vars=*/{}, /*rt_vars=*/{});
   }
   return IndexingMap::FromTensorSizes(
       ComputeTransposeIndexingMap(ToTransposeDimensions(shape.layout()),

From a6f127809b9864a335c273617516883613edad5a Mon Sep 17 00:00:00 2001
From: Henning Becker <hebecker@google.com>
Date: Thu, 16 May 2024 10:15:21 -0700
Subject: [PATCH 280/478] Remove duplicate dependency from PJRT test

A dependency from `se_gpu_pjrt_compiler_aot_test` to `amdgpu_compiler_impl`
had been recently added - I believe by accident.

The test shouldn't generally depend on the AMDGPU compiler - only if it's
a ROCm build in which case the dependency is already there.

PiperOrigin-RevId: 634448333
---
 third_party/xla/xla/pjrt/gpu/BUILD | 1 -
 1 file changed, 1 deletion(-)

diff --git a/third_party/xla/xla/pjrt/gpu/BUILD b/third_party/xla/xla/pjrt/gpu/BUILD
index f8991a93a189e5..69eb9f9a635d64 100644
--- a/third_party/xla/xla/pjrt/gpu/BUILD
+++ b/third_party/xla/xla/pjrt/gpu/BUILD
@@ -338,7 +338,6 @@ xla_cc_test(
         "//xla/service:compiler",
         "//xla/service:gpu_plugin",
         "//xla/service:hlo_parser",
-        "//xla/service/gpu:amdgpu_compiler_impl",
         "//xla/tests:literal_test_util",
         "@com_google_absl//absl/memory",
         "@com_google_absl//absl/status",

From 5b91995308c2b92152e0fbaf30fc2c259bb58062 Mon Sep 17 00:00:00 2001
From: Kuy Mainwaring <kuym@google.com>
Date: Thu, 16 May 2024 10:15:46 -0700
Subject: [PATCH 281/478] [XLA:GPU] Clang-tidy cleanup for
 xla/service/gpu/model/gpu_performance_model_base.h

PiperOrigin-RevId: 634448529
---
 .../xla/xla/service/gpu/model/gpu_performance_model_base.h  | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/third_party/xla/xla/service/gpu/model/gpu_performance_model_base.h b/third_party/xla/xla/service/gpu/model/gpu_performance_model_base.h
index 835746b9c25316..8a1cbcd26abe5e 100644
--- a/third_party/xla/xla/service/gpu/model/gpu_performance_model_base.h
+++ b/third_party/xla/xla/service/gpu/model/gpu_performance_model_base.h
@@ -54,9 +54,9 @@ struct EstimateRunTimeData {
         " compute_time: %s\n"
         " exec_time: %s\n"
         "}",
-        flops, bytes_written, num_threads, FormatDuration(read_time),
-        FormatDuration(write_time), FormatDuration(compute_time),
-        FormatDuration(exec_time));
+        flops, bytes_written, num_threads, absl::FormatDuration(read_time),
+        absl::FormatDuration(write_time), absl::FormatDuration(compute_time),
+        absl::FormatDuration(exec_time));
   }
 };
 

From be9bf7d8c21a92bf9142c5ef828fa20bbeab77e1 Mon Sep 17 00:00:00 2001
From: Kuy Mainwaring <kuym@google.com>
Date: Thu, 16 May 2024 10:16:10 -0700
Subject: [PATCH 282/478] [XLA:GPU] Clang-tidy cleanup for
 xla/service/gpu/llvm_gpu_backend/gpu_backend_lib.cc

PiperOrigin-RevId: 634448673
---
 .../xla/xla/service/gpu/llvm_gpu_backend/gpu_backend_lib.cc      | 1 +
 1 file changed, 1 insertion(+)

diff --git a/third_party/xla/xla/service/gpu/llvm_gpu_backend/gpu_backend_lib.cc b/third_party/xla/xla/service/gpu/llvm_gpu_backend/gpu_backend_lib.cc
index 6f60f4b3a5ac5d..1c5929166c8a39 100644
--- a/third_party/xla/xla/service/gpu/llvm_gpu_backend/gpu_backend_lib.cc
+++ b/third_party/xla/xla/service/gpu/llvm_gpu_backend/gpu_backend_lib.cc
@@ -24,6 +24,7 @@ limitations under the License.
 #include <mutex>  // NOLINT
 #include <optional>
 #include <string>
+#include <string_view>
 #include <system_error>  // NOLINT
 #include <utility>
 #include <variant>

From 5aade4cda97856b12d1c6819b9f747e7977b1550 Mon Sep 17 00:00:00 2001
From: Michael Hudgins <michaelhudgins@google.com>
Date: Thu, 16 May 2024 10:25:56 -0700
Subject: [PATCH 283/478] Modify docker pull backoff in CI

PiperOrigin-RevId: 634452150
---
 ci/official/utilities/setup_docker.sh | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/ci/official/utilities/setup_docker.sh b/ci/official/utilities/setup_docker.sh
index 36afa2545eb244..91618c75f3ba51 100755
--- a/ci/official/utilities/setup_docker.sh
+++ b/ci/official/utilities/setup_docker.sh
@@ -14,11 +14,12 @@
 # limitations under the License.
 # ==============================================================================
 if [[ "$TFCI_DOCKER_PULL_ENABLE" == 1 ]]; then
-  # Simple retry logic for docker-pull errors. Sleeps for 15s if a pull fails.
+  # Simple retry logic for docker-pull errors. Sleeps if a pull fails.
   # Pulling an already-pulled container image will finish instantly, so
   # repeating the command costs nothing.
   docker pull "$TFCI_DOCKER_IMAGE" || sleep 15
-  docker pull "$TFCI_DOCKER_IMAGE" || sleep 15
+  docker pull "$TFCI_DOCKER_IMAGE" || sleep 30
+  docker pull "$TFCI_DOCKER_IMAGE" || sleep 60
   docker pull "$TFCI_DOCKER_IMAGE"
 fi 
 

From c05f5ad3d7da282726597343b19e8c6af6e71b1c Mon Sep 17 00:00:00 2001
From: Matthias Kramm <kramm@google.com>
Date: Thu, 16 May 2024 10:33:57 -0700
Subject: [PATCH 284/478] Change DebugString call to
 proto2::TextFormat::PrintToString.

PiperOrigin-RevId: 634454849
---
 tensorflow/core/ops/compat/op_compatibility_lib.h | 6 +++++-
 1 file changed, 5 insertions(+), 1 deletion(-)

diff --git a/tensorflow/core/ops/compat/op_compatibility_lib.h b/tensorflow/core/ops/compat/op_compatibility_lib.h
index 2f26fd63bdd7a5..1693f2bd5cfe27 100644
--- a/tensorflow/core/ops/compat/op_compatibility_lib.h
+++ b/tensorflow/core/ops/compat/op_compatibility_lib.h
@@ -52,7 +52,11 @@ class OpCompatibilityLib {
 
   // Should match the contents of ops_file().  Run before calling
   // ValidateCompatible().
-  string OpsString() const { return op_list_.DebugString(); }
+  string OpsString() const {
+    string result;
+    google::protobuf::TextFormat::PrintToString(op_list_, &result);
+    return result;
+  }
 
   // Returns the number of ops in OpsString(), includes all ops, not
   // just stable ops.

From 4c5b5197f8ff2d79d251f3755c9c0d826c3e55a3 Mon Sep 17 00:00:00 2001
From: Henning Becker <hebecker@google.com>
Date: Thu, 16 May 2024 10:44:06 -0700
Subject: [PATCH 285/478] Make buffer_comparator_kernel target a
 gpu_kernel_library

PiperOrigin-RevId: 634458480
---
 third_party/xla/xla/service/gpu/BUILD | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/third_party/xla/xla/service/gpu/BUILD b/third_party/xla/xla/service/gpu/BUILD
index 1c2712b6fd3538..5f10710157ddf1 100644
--- a/third_party/xla/xla/service/gpu/BUILD
+++ b/third_party/xla/xla/service/gpu/BUILD
@@ -4473,7 +4473,7 @@ cc_library(
     ]),
 )
 
-cuda_library(
+gpu_kernel_library(
     name = "buffer_comparator_kernel",
     srcs = if_gpu_is_configured(["buffer_comparator.cu.cc"]),
     copts = rocm_copts(),

From e1cd322f067a222c7cb91a4213196d1ecbc057b1 Mon Sep 17 00:00:00 2001
From: Benjamin Chetioui <bchetioui@google.com>
Date: Thu, 16 May 2024 10:51:11 -0700
Subject: [PATCH 286/478] [XLA:GPU] Remove filter checking if an indexing map
 is tileable in `SymbolicTile`.

Symbolic tiles are now robust and well-reasoned enough that the filter should
no longer be necessary.

PiperOrigin-RevId: 634460882
---
 .../xla/service/gpu/model/symbolic_tile.cc    | 40 -------------------
 1 file changed, 40 deletions(-)

diff --git a/third_party/xla/xla/service/gpu/model/symbolic_tile.cc b/third_party/xla/xla/service/gpu/model/symbolic_tile.cc
index 18ddd718c107f8..07c6881f13a5b1 100644
--- a/third_party/xla/xla/service/gpu/model/symbolic_tile.cc
+++ b/third_party/xla/xla/service/gpu/model/symbolic_tile.cc
@@ -78,39 +78,6 @@ std::vector<AffineExpr> DimsToSymbols(std::vector<AffineExpr> expressions,
   return expressions;
 }
 
-// Internal helper that checks whether an affine map describes a tileable space.
-// In simple terms, this currently returns true if "dimensions don't mix", i.e.,
-// every result expression only refers to a single dimension (or symbol).
-//
-// TODO(b/328427138): this is too restrictive for expressions involving e.g.
-// (output-to-input) split reshapes, where several symbols may appear within the
-// same expression but still yield a tileable space. This will be handled in a
-// forthcoming change.
-bool IndexingMapDescribesTileableSpace(const IndexingMap& indexing_map) {
-  for (AffineExpr result_expr : indexing_map.GetAffineMap().getResults()) {
-    // Using a simple integer here might be overly restrictive, since there may
-    // be cases where the same symbol appears in several places within the
-    // expression. It is a bit unclear whether this is a case that would happen
-    // in practice and whether we would be able to handle it well in all cases
-    // if it did. For that reason, we err on the side of conservatism and
-    // explicitly do not support such cases.
-    int64_t num_hits = 0;
-    result_expr.walk([&num_hits, &indexing_map](AffineExpr expr) {
-      if ((expr.getKind() == AffineExprKind::DimId) ||
-          (expr.getKind() == AffineExprKind::SymbolId &&
-           indexing_map.IsRangeVarSymbol(
-               llvm::cast<mlir::AffineSymbolExpr>(expr)))) {
-        ++num_hits;
-      }
-    });
-
-    if (num_hits > 1) {
-      return false;
-    }
-  }
-  return true;
-}
-
 // Helper to perform function application to using the same parameter for every
 // dimension and symbol parameter.
 AffineMap SubstituteAllIndicesAndRangeVarSymbolsWithSameValue(
@@ -309,13 +276,6 @@ AffineExpr SimplifyAffineExpr(const AffineExpr& expr,
     const IndexingMap& indexing_map) {
   VLOG(1) << "SymbolicTile::FromIndexingMap: " << indexing_map.ToString();
 
-  // TODO(b/328427138): handle multiple symbols in a single tile to support
-  // merging dimensions.
-  if (!IndexingMapDescribesTileableSpace(indexing_map)) {
-    VLOG(1) << "Not a tileable indexing map";
-    return std::nullopt;
-  }
-
   AffineMap input_affine_map = indexing_map.GetAffineMap();
   MLIRContext* mlir_context = input_affine_map.getContext();
 

From da32b633de3bcf37a8e1f0b1caf778ffcaff5025 Mon Sep 17 00:00:00 2001
From: Kyle Lucke <klucke@google.com>
Date: Thu, 16 May 2024 11:09:30 -0700
Subject: [PATCH 287/478] Use absl::Status instead of xla::Status now that
 they're identical.

PiperOrigin-RevId: 634467593
---
 third_party/xla/xla/hlo/ir/dfs_hlo_visitor.cc |  10 +-
 third_party/xla/xla/hlo/ir/dfs_hlo_visitor.h  | 253 +++++++++---------
 .../xla/hlo/ir/dfs_hlo_visitor_with_default.h | 182 +++++++------
 .../xla/hlo/ir/dynamic_parameter_binding.cc   |   8 +-
 .../xla/hlo/ir/dynamic_parameter_binding.h    |  12 +-
 third_party/xla/xla/hlo/ir/hlo_computation.cc |  35 +--
 third_party/xla/xla/hlo/ir/hlo_computation.h  |  58 ++--
 .../hlo/ir/hlo_input_output_alias_config.cc   |  18 +-
 third_party/xla/xla/hlo/ir/hlo_instruction.cc |  82 +++---
 third_party/xla/xla/hlo/ir/hlo_instruction.h  |  92 ++++---
 .../xla/xla/hlo/ir/hlo_instructions.cc        |   2 +-
 third_party/xla/xla/hlo/ir/hlo_instructions.h |   2 +-
 third_party/xla/xla/hlo/ir/hlo_module.cc      |   9 +-
 third_party/xla/xla/hlo/ir/hlo_module.h       |  11 +-
 .../xla/xla/hlo/ir/hlo_module_metadata.cc     |   6 +-
 .../xla/xla/hlo/ir/hlo_module_metadata.h      |  20 +-
 third_party/xla/xla/hlo/ir/hlo_schedule.cc    |   6 +-
 third_party/xla/xla/hlo/ir/hlo_schedule.h     |   6 +-
 third_party/xla/xla/hlo/ir/hlo_sharding.cc    |  22 +-
 third_party/xla/xla/hlo/ir/hlo_sharding.h     |  14 +-
 .../xla/xla/hlo/ir/hlo_sharding_metadata.cc   |  16 +-
 .../xla/xla/hlo/ir/hlo_sharding_metadata.h    |   4 +-
 third_party/xla/xla/hlo/ir/tile_assignment.cc |   5 +-
 third_party/xla/xla/hlo/ir/tile_assignment.h  |   5 +-
 24 files changed, 452 insertions(+), 426 deletions(-)

diff --git a/third_party/xla/xla/hlo/ir/dfs_hlo_visitor.cc b/third_party/xla/xla/hlo/ir/dfs_hlo_visitor.cc
index 22f031eb7f77fb..a41ec006618271 100644
--- a/third_party/xla/xla/hlo/ir/dfs_hlo_visitor.cc
+++ b/third_party/xla/xla/hlo/ir/dfs_hlo_visitor.cc
@@ -25,14 +25,14 @@ limitations under the License.
 namespace xla {
 
 template <typename HloInstructionPtr>
-Status DfsHloVisitorBase<HloInstructionPtr>::HandleElementwiseUnary(
+absl::Status DfsHloVisitorBase<HloInstructionPtr>::HandleElementwiseUnary(
     HloInstructionPtr hlo) {
   return Unimplemented("DfsHloVisitor::HandleElementwiseUnary: %s",
                        HloOpcodeString(hlo->opcode()));
 }
 
 template <typename HloInstructionPtr>
-Status DfsHloVisitorBase<HloInstructionPtr>::HandleElementwiseBinary(
+absl::Status DfsHloVisitorBase<HloInstructionPtr>::HandleElementwiseBinary(
     HloInstructionPtr hlo) {
   return Unimplemented("DfsHloVisitor::HandleElementwiseBinary: %s",
                        HloOpcodeString(hlo->opcode()));
@@ -62,12 +62,14 @@ void DfsHloVisitorBase<HloInstructionPtr>::SetVisited(
 }
 
 template <typename HloInstructionPtr>
-Status DfsHloVisitorBase<HloInstructionPtr>::Preprocess(HloInstructionPtr) {
+absl::Status DfsHloVisitorBase<HloInstructionPtr>::Preprocess(
+    HloInstructionPtr) {
   return OkStatus();
 }
 
 template <typename HloInstructionPtr>
-Status DfsHloVisitorBase<HloInstructionPtr>::Postprocess(HloInstructionPtr) {
+absl::Status DfsHloVisitorBase<HloInstructionPtr>::Postprocess(
+    HloInstructionPtr) {
   return OkStatus();
 }
 
diff --git a/third_party/xla/xla/hlo/ir/dfs_hlo_visitor.h b/third_party/xla/xla/hlo/ir/dfs_hlo_visitor.h
index 29d5715716bd5f..3382aeef22a163 100644
--- a/third_party/xla/xla/hlo/ir/dfs_hlo_visitor.h
+++ b/third_party/xla/xla/hlo/ir/dfs_hlo_visitor.h
@@ -77,256 +77,257 @@ class DfsHloVisitorBase {
   // These routines are self-descriptive, see class comment for usage
   // information.
 
-  virtual Status HandleElementwiseUnary(HloInstructionPtr hlo);
-  virtual Status HandleElementwiseBinary(HloInstructionPtr hlo);
+  virtual absl::Status HandleElementwiseUnary(HloInstructionPtr hlo);
+  virtual absl::Status HandleElementwiseBinary(HloInstructionPtr hlo);
 
-  virtual Status HandleClamp(HloInstructionPtr hlo) = 0;
-  virtual Status HandleSelect(HloInstructionPtr hlo) = 0;
-  virtual Status HandleMaximum(HloInstructionPtr hlo) {
+  virtual absl::Status HandleClamp(HloInstructionPtr hlo) = 0;
+  virtual absl::Status HandleSelect(HloInstructionPtr hlo) = 0;
+  virtual absl::Status HandleMaximum(HloInstructionPtr hlo) {
     return HandleElementwiseBinary(hlo);
   }
-  virtual Status HandleMinimum(HloInstructionPtr hlo) {
+  virtual absl::Status HandleMinimum(HloInstructionPtr hlo) {
     return HandleElementwiseBinary(hlo);
   }
-  virtual Status HandleConcatenate(HloInstructionPtr hlo) = 0;
-  virtual Status HandleConvert(HloInstructionPtr hlo) {
+  virtual absl::Status HandleConcatenate(HloInstructionPtr hlo) = 0;
+  virtual absl::Status HandleConvert(HloInstructionPtr hlo) {
     return HandleElementwiseUnary(hlo);
   }
-  virtual Status HandleBitcastConvert(HloInstructionPtr hlo) {
+  virtual absl::Status HandleBitcastConvert(HloInstructionPtr hlo) {
     return HandleElementwiseUnary(hlo);
   }
-  virtual Status HandleStochasticConvert(HloInstructionPtr hlo) {
+  virtual absl::Status HandleStochasticConvert(HloInstructionPtr hlo) {
     return HandleElementwiseBinary(hlo);
   }
-  virtual Status HandleCopy(HloInstructionPtr hlo) {
+  virtual absl::Status HandleCopy(HloInstructionPtr hlo) {
     return HandleElementwiseUnary(hlo);
   }
-  virtual Status HandleComplex(HloInstructionPtr hlo) {
+  virtual absl::Status HandleComplex(HloInstructionPtr hlo) {
     return HandleElementwiseBinary(hlo);
   }
-  virtual Status HandleMultiply(HloInstructionPtr hlo) {
+  virtual absl::Status HandleMultiply(HloInstructionPtr hlo) {
     return HandleElementwiseBinary(hlo);
   }
-  virtual Status HandleDot(HloInstructionPtr hlo) = 0;
-  virtual Status HandlePower(HloInstructionPtr hlo) {
+  virtual absl::Status HandleDot(HloInstructionPtr hlo) = 0;
+  virtual absl::Status HandlePower(HloInstructionPtr hlo) {
     return HandleElementwiseBinary(hlo);
   }
-  virtual Status HandleSqrt(HloInstructionPtr hlo) {
+  virtual absl::Status HandleSqrt(HloInstructionPtr hlo) {
     return HandleElementwiseUnary(hlo);
   }
-  virtual Status HandleRsqrt(HloInstructionPtr hlo) {
+  virtual absl::Status HandleRsqrt(HloInstructionPtr hlo) {
     return HandleElementwiseUnary(hlo);
   }
-  virtual Status HandleCbrt(HloInstructionPtr hlo) {
+  virtual absl::Status HandleCbrt(HloInstructionPtr hlo) {
     return HandleElementwiseUnary(hlo);
   }
   /* go/keep-sorted start */
-  virtual Status HandleAllGather(HloInstructionPtr hlo) = 0;
-  virtual Status HandleAllGatherDone(HloInstructionPtr hlo) = 0;
-  virtual Status HandleAllGatherStart(HloInstructionPtr hlo) = 0;
-  virtual Status HandleAllReduce(HloInstructionPtr hlo) = 0;
-  virtual Status HandleAllReduceDone(HloInstructionPtr hlo) = 0;
-  virtual Status HandleAllReduceStart(HloInstructionPtr hlo) = 0;
-  virtual Status HandleAllToAll(HloInstructionPtr hlo) = 0;
-  virtual Status HandleCollectiveBroadcast(HloInstructionPtr hlo) = 0;
-  virtual Status HandleCollectivePermute(HloInstructionPtr hlo) = 0;
-  virtual Status HandleCollectivePermuteDone(HloInstructionPtr hlo) = 0;
-  virtual Status HandleCollectivePermuteStart(HloInstructionPtr hlo) = 0;
-  virtual Status HandleConvolution(HloInstructionPtr hlo) = 0;
-  virtual Status HandleOptimizationBarrier(HloInstructionPtr hlo) = 0;
-  virtual Status HandlePartitionId(HloInstructionPtr hlo) = 0;
-  virtual Status HandleReduceScatter(HloInstructionPtr hlo) = 0;
-  virtual Status HandleReplicaId(HloInstructionPtr hlo) = 0;
+  virtual absl::Status HandleAllGather(HloInstructionPtr hlo) = 0;
+  virtual absl::Status HandleAllGatherDone(HloInstructionPtr hlo) = 0;
+  virtual absl::Status HandleAllGatherStart(HloInstructionPtr hlo) = 0;
+  virtual absl::Status HandleAllReduce(HloInstructionPtr hlo) = 0;
+  virtual absl::Status HandleAllReduceDone(HloInstructionPtr hlo) = 0;
+  virtual absl::Status HandleAllReduceStart(HloInstructionPtr hlo) = 0;
+  virtual absl::Status HandleAllToAll(HloInstructionPtr hlo) = 0;
+  virtual absl::Status HandleCollectiveBroadcast(HloInstructionPtr hlo) = 0;
+  virtual absl::Status HandleCollectivePermute(HloInstructionPtr hlo) = 0;
+  virtual absl::Status HandleCollectivePermuteDone(HloInstructionPtr hlo) = 0;
+  virtual absl::Status HandleCollectivePermuteStart(HloInstructionPtr hlo) = 0;
+  virtual absl::Status HandleConvolution(HloInstructionPtr hlo) = 0;
+  virtual absl::Status HandleOptimizationBarrier(HloInstructionPtr hlo) = 0;
+  virtual absl::Status HandlePartitionId(HloInstructionPtr hlo) = 0;
+  virtual absl::Status HandleReduceScatter(HloInstructionPtr hlo) = 0;
+  virtual absl::Status HandleReplicaId(HloInstructionPtr hlo) = 0;
   /* go/keep-sorted end */
 
   /* go/keep-sorted start */
-  virtual Status HandleCholesky(HloInstructionPtr hlo) = 0;
-  virtual Status HandleFft(HloInstructionPtr fft) = 0;
-  virtual Status HandleTopK(HloInstructionPtr hlo) = 0;
-  virtual Status HandleTriangularSolve(HloInstructionPtr hlo) = 0;
+  virtual absl::Status HandleCholesky(HloInstructionPtr hlo) = 0;
+  virtual absl::Status HandleFft(HloInstructionPtr fft) = 0;
+  virtual absl::Status HandleTopK(HloInstructionPtr hlo) = 0;
+  virtual absl::Status HandleTriangularSolve(HloInstructionPtr hlo) = 0;
   /* go/keep-sorted end */
 
-  virtual Status HandleGetDimensionSize(HloInstructionPtr hlo) = 0;
-  virtual Status HandleSetDimensionSize(HloInstructionPtr hlo) = 0;
+  virtual absl::Status HandleGetDimensionSize(HloInstructionPtr hlo) = 0;
+  virtual absl::Status HandleSetDimensionSize(HloInstructionPtr hlo) = 0;
 
-  virtual Status HandleCompare(HloInstructionPtr hlo) {
+  virtual absl::Status HandleCompare(HloInstructionPtr hlo) {
     return HandleElementwiseBinary(hlo);
   }
-  virtual Status HandleAdd(HloInstructionPtr hlo) {
+  virtual absl::Status HandleAdd(HloInstructionPtr hlo) {
     return HandleElementwiseBinary(hlo);
   }
-  virtual Status HandleDivide(HloInstructionPtr hlo) {
+  virtual absl::Status HandleDivide(HloInstructionPtr hlo) {
     return HandleElementwiseBinary(hlo);
   }
-  virtual Status HandleRemainder(HloInstructionPtr hlo) {
+  virtual absl::Status HandleRemainder(HloInstructionPtr hlo) {
     return HandleElementwiseBinary(hlo);
   }
-  virtual Status HandleSubtract(HloInstructionPtr hlo) {
+  virtual absl::Status HandleSubtract(HloInstructionPtr hlo) {
     return HandleElementwiseBinary(hlo);
   }
-  virtual Status HandleAbs(HloInstructionPtr hlo) {
+  virtual absl::Status HandleAbs(HloInstructionPtr hlo) {
     return HandleElementwiseUnary(hlo);
   }
-  virtual Status HandleAtan2(HloInstructionPtr hlo) {
+  virtual absl::Status HandleAtan2(HloInstructionPtr hlo) {
     return HandleElementwiseBinary(hlo);
   }
-  virtual Status HandleRound(HloInstructionPtr hlo) {
+  virtual absl::Status HandleRound(HloInstructionPtr hlo) {
     return HandleElementwiseUnary(hlo);
   }
-  virtual Status HandleRoundNearestEven(HloInstructionPtr hlo) {
+  virtual absl::Status HandleRoundNearestEven(HloInstructionPtr hlo) {
     return HandleElementwiseUnary(hlo);
   }
-  virtual Status HandleErf(HloInstructionPtr hlo) {
+  virtual absl::Status HandleErf(HloInstructionPtr hlo) {
     return HandleElementwiseUnary(hlo);
   }
-  virtual Status HandleLogistic(HloInstructionPtr hlo) {
+  virtual absl::Status HandleLogistic(HloInstructionPtr hlo) {
     return HandleElementwiseUnary(hlo);
   }
-  virtual Status HandleSign(HloInstructionPtr hlo) {
+  virtual absl::Status HandleSign(HloInstructionPtr hlo) {
     return HandleElementwiseUnary(hlo);
   }
-  virtual Status HandleNegate(HloInstructionPtr hlo) {
+  virtual absl::Status HandleNegate(HloInstructionPtr hlo) {
     return HandleElementwiseUnary(hlo);
   }
-  virtual Status HandleExp(HloInstructionPtr hlo) {
+  virtual absl::Status HandleExp(HloInstructionPtr hlo) {
     return HandleElementwiseUnary(hlo);
   }
-  virtual Status HandleExpm1(HloInstructionPtr hlo) {
+  virtual absl::Status HandleExpm1(HloInstructionPtr hlo) {
     return HandleElementwiseUnary(hlo);
   }
-  virtual Status HandleFloor(HloInstructionPtr hlo) {
+  virtual absl::Status HandleFloor(HloInstructionPtr hlo) {
     return HandleElementwiseUnary(hlo);
   }
-  virtual Status HandleCeil(HloInstructionPtr hlo) {
+  virtual absl::Status HandleCeil(HloInstructionPtr hlo) {
     return HandleElementwiseUnary(hlo);
   }
-  virtual Status HandleLog(HloInstructionPtr hlo) {
+  virtual absl::Status HandleLog(HloInstructionPtr hlo) {
     return HandleElementwiseUnary(hlo);
   }
-  virtual Status HandleClz(HloInstructionPtr hlo) {
+  virtual absl::Status HandleClz(HloInstructionPtr hlo) {
     return HandleElementwiseUnary(hlo);
   }
-  virtual Status HandleLog1p(HloInstructionPtr hlo) {
+  virtual absl::Status HandleLog1p(HloInstructionPtr hlo) {
     return HandleElementwiseUnary(hlo);
   }
-  virtual Status HandleCos(HloInstructionPtr hlo) {
+  virtual absl::Status HandleCos(HloInstructionPtr hlo) {
     return HandleElementwiseUnary(hlo);
   }
-  virtual Status HandleSin(HloInstructionPtr hlo) {
+  virtual absl::Status HandleSin(HloInstructionPtr hlo) {
     return HandleElementwiseUnary(hlo);
   }
-  virtual Status HandleTan(HloInstructionPtr hlo) {
+  virtual absl::Status HandleTan(HloInstructionPtr hlo) {
     return HandleElementwiseUnary(hlo);
   }
-  virtual Status HandleTanh(HloInstructionPtr hlo) {
+  virtual absl::Status HandleTanh(HloInstructionPtr hlo) {
     return HandleElementwiseUnary(hlo);
   }
-  virtual Status HandleReal(HloInstructionPtr hlo) {
+  virtual absl::Status HandleReal(HloInstructionPtr hlo) {
     return HandleElementwiseUnary(hlo);
   }
-  virtual Status HandleImag(HloInstructionPtr hlo) {
+  virtual absl::Status HandleImag(HloInstructionPtr hlo) {
     return HandleElementwiseUnary(hlo);
   }
-  virtual Status HandleIsFinite(HloInstructionPtr hlo) {
+  virtual absl::Status HandleIsFinite(HloInstructionPtr hlo) {
     return HandleElementwiseUnary(hlo);
   }
-  virtual Status HandleAnd(HloInstructionPtr hlo) {
+  virtual absl::Status HandleAnd(HloInstructionPtr hlo) {
     return HandleElementwiseBinary(hlo);
   }
-  virtual Status HandleNot(HloInstructionPtr hlo) {
+  virtual absl::Status HandleNot(HloInstructionPtr hlo) {
     return HandleElementwiseUnary(hlo);
   }
-  virtual Status HandleOr(HloInstructionPtr hlo) {
+  virtual absl::Status HandleOr(HloInstructionPtr hlo) {
     return HandleElementwiseBinary(hlo);
   }
-  virtual Status HandleXor(HloInstructionPtr hlo) {
+  virtual absl::Status HandleXor(HloInstructionPtr hlo) {
     return HandleElementwiseBinary(hlo);
   }
-  virtual Status HandlePopulationCount(HloInstructionPtr hlo) {
+  virtual absl::Status HandlePopulationCount(HloInstructionPtr hlo) {
     return HandleElementwiseUnary(hlo);
   }
-  virtual Status HandleShiftLeft(HloInstructionPtr hlo) {
+  virtual absl::Status HandleShiftLeft(HloInstructionPtr hlo) {
     return HandleElementwiseBinary(hlo);
   }
-  virtual Status HandleShiftRightArithmetic(HloInstructionPtr hlo) {
+  virtual absl::Status HandleShiftRightArithmetic(HloInstructionPtr hlo) {
     return HandleElementwiseBinary(hlo);
   }
-  virtual Status HandleShiftRightLogical(HloInstructionPtr hlo) {
+  virtual absl::Status HandleShiftRightLogical(HloInstructionPtr hlo) {
     return HandleElementwiseBinary(hlo);
   }
 
-  virtual Status HandleReducePrecision(HloInstructionPtr hlo) {
+  virtual absl::Status HandleReducePrecision(HloInstructionPtr hlo) {
     return HandleElementwiseUnary(hlo);
   }
 
-  virtual Status HandleDomain(HloInstructionPtr hlo) {
+  virtual absl::Status HandleDomain(HloInstructionPtr hlo) {
     return HandleElementwiseUnary(hlo);
   }
 
   /* go/keep-sorted start */
-  virtual Status HandleInfeed(HloInstructionPtr hlo) = 0;
-  virtual Status HandleOutfeed(HloInstructionPtr hlo) = 0;
+  virtual absl::Status HandleInfeed(HloInstructionPtr hlo) = 0;
+  virtual absl::Status HandleOutfeed(HloInstructionPtr hlo) = 0;
   /* go/keep-sorted end */
 
   /* go/keep-sorted start */
-  virtual Status HandleBitcast(HloInstructionPtr hlo) = 0;
-  virtual Status HandleBroadcast(HloInstructionPtr hlo) = 0;
-  virtual Status HandleCall(HloInstructionPtr hlo) = 0;
-  virtual Status HandleConditional(HloInstructionPtr hlo) = 0;
-  virtual Status HandleConstant(HloInstructionPtr hlo) = 0;
-  virtual Status HandleCustomCall(HloInstructionPtr hlo) = 0;
-  virtual Status HandleDynamicReshape(HloInstructionPtr hlo) = 0;
-  virtual Status HandleDynamicSlice(HloInstructionPtr hlo) = 0;
-  virtual Status HandleDynamicUpdateSlice(HloInstructionPtr hlo) = 0;
-  virtual Status HandleFusion(HloInstructionPtr hlo) = 0;
-  virtual Status HandleGather(HloInstructionPtr hlo) = 0;
-  virtual Status HandleGetTupleElement(HloInstructionPtr hlo) = 0;
-  virtual Status HandleIota(HloInstructionPtr hlo) = 0;
-  virtual Status HandleMap(HloInstructionPtr hlo) = 0;
-  virtual Status HandleParameter(HloInstructionPtr hlo) = 0;
-  virtual Status HandleReduce(HloInstructionPtr hlo) = 0;
-  virtual Status HandleReduceWindow(HloInstructionPtr hlo) = 0;
-  virtual Status HandleReshape(HloInstructionPtr hlo) = 0;
-  virtual Status HandleReverse(HloInstructionPtr hlo) = 0;
-  virtual Status HandleRng(HloInstructionPtr hlo) = 0;
-  virtual Status HandleRngBitGenerator(HloInstructionPtr hlo) = 0;
-  virtual Status HandleRngGetAndUpdateState(HloInstructionPtr hlo) = 0;
-  virtual Status HandleScatter(HloInstructionPtr hlo) = 0;
-  virtual Status HandleSelectAndScatter(HloInstructionPtr hlo) = 0;
-  virtual Status HandleSlice(HloInstructionPtr hlo) = 0;
-  virtual Status HandleSort(HloInstructionPtr hlo) = 0;
-  virtual Status HandleTranspose(HloInstructionPtr hlo) = 0;
-  virtual Status HandleTuple(HloInstructionPtr hlo) = 0;
-  virtual Status HandleWhile(HloInstructionPtr hlo) = 0;
+  virtual absl::Status HandleBitcast(HloInstructionPtr hlo) = 0;
+  virtual absl::Status HandleBroadcast(HloInstructionPtr hlo) = 0;
+  virtual absl::Status HandleCall(HloInstructionPtr hlo) = 0;
+  virtual absl::Status HandleConditional(HloInstructionPtr hlo) = 0;
+  virtual absl::Status HandleConstant(HloInstructionPtr hlo) = 0;
+  virtual absl::Status HandleCustomCall(HloInstructionPtr hlo) = 0;
+  virtual absl::Status HandleDynamicReshape(HloInstructionPtr hlo) = 0;
+  virtual absl::Status HandleDynamicSlice(HloInstructionPtr hlo) = 0;
+  virtual absl::Status HandleDynamicUpdateSlice(HloInstructionPtr hlo) = 0;
+  virtual absl::Status HandleFusion(HloInstructionPtr hlo) = 0;
+  virtual absl::Status HandleGather(HloInstructionPtr hlo) = 0;
+  virtual absl::Status HandleGetTupleElement(HloInstructionPtr hlo) = 0;
+  virtual absl::Status HandleIota(HloInstructionPtr hlo) = 0;
+  virtual absl::Status HandleMap(HloInstructionPtr hlo) = 0;
+  virtual absl::Status HandleParameter(HloInstructionPtr hlo) = 0;
+  virtual absl::Status HandleReduce(HloInstructionPtr hlo) = 0;
+  virtual absl::Status HandleReduceWindow(HloInstructionPtr hlo) = 0;
+  virtual absl::Status HandleReshape(HloInstructionPtr hlo) = 0;
+  virtual absl::Status HandleReverse(HloInstructionPtr hlo) = 0;
+  virtual absl::Status HandleRng(HloInstructionPtr hlo) = 0;
+  virtual absl::Status HandleRngBitGenerator(HloInstructionPtr hlo) = 0;
+  virtual absl::Status HandleRngGetAndUpdateState(HloInstructionPtr hlo) = 0;
+  virtual absl::Status HandleScatter(HloInstructionPtr hlo) = 0;
+  virtual absl::Status HandleSelectAndScatter(HloInstructionPtr hlo) = 0;
+  virtual absl::Status HandleSlice(HloInstructionPtr hlo) = 0;
+  virtual absl::Status HandleSort(HloInstructionPtr hlo) = 0;
+  virtual absl::Status HandleTranspose(HloInstructionPtr hlo) = 0;
+  virtual absl::Status HandleTuple(HloInstructionPtr hlo) = 0;
+  virtual absl::Status HandleWhile(HloInstructionPtr hlo) = 0;
   /* go/keep-sorted end */
 
-  virtual Status HandlePad(HloInstructionPtr hlo) = 0;
+  virtual absl::Status HandlePad(HloInstructionPtr hlo) = 0;
 
-  virtual Status HandleAsyncStart(HloInstructionPtr hlo) = 0;
-  virtual Status HandleAsyncUpdate(HloInstructionPtr hlo) = 0;
-  virtual Status HandleAsyncDone(HloInstructionPtr hlo) = 0;
+  virtual absl::Status HandleAsyncStart(HloInstructionPtr hlo) = 0;
+  virtual absl::Status HandleAsyncUpdate(HloInstructionPtr hlo) = 0;
+  virtual absl::Status HandleAsyncDone(HloInstructionPtr hlo) = 0;
 
-  virtual Status HandleCopyStart(HloInstructionPtr copy_start) = 0;
-  virtual Status HandleCopyDone(HloInstructionPtr copy_done) = 0;
+  virtual absl::Status HandleCopyStart(HloInstructionPtr copy_start) = 0;
+  virtual absl::Status HandleCopyDone(HloInstructionPtr copy_done) = 0;
 
-  virtual Status HandleSend(HloInstructionPtr send) = 0;
-  virtual Status HandleSendDone(HloInstructionPtr send_done) = 0;
+  virtual absl::Status HandleSend(HloInstructionPtr send) = 0;
+  virtual absl::Status HandleSendDone(HloInstructionPtr send_done) = 0;
 
-  virtual Status HandleRecv(HloInstructionPtr recv) = 0;
-  virtual Status HandleRecvDone(HloInstructionPtr recv_done) = 0;
+  virtual absl::Status HandleRecv(HloInstructionPtr recv) = 0;
+  virtual absl::Status HandleRecvDone(HloInstructionPtr recv_done) = 0;
 
-  virtual Status HandleBatchNormTraining(HloInstructionPtr hlo) = 0;
+  virtual absl::Status HandleBatchNormTraining(HloInstructionPtr hlo) = 0;
 
-  virtual Status HandleBatchNormInference(HloInstructionPtr hlo) = 0;
+  virtual absl::Status HandleBatchNormInference(HloInstructionPtr hlo) = 0;
 
-  virtual Status HandleBatchNormGrad(HloInstructionPtr hlo) = 0;
+  virtual absl::Status HandleBatchNormGrad(HloInstructionPtr hlo) = 0;
 
-  virtual Status HandleAddDependency(HloInstructionPtr add_dependency) = 0;
-  virtual Status HandleAfterAll(HloInstructionPtr token) = 0;
+  virtual absl::Status HandleAddDependency(
+      HloInstructionPtr add_dependency) = 0;
+  virtual absl::Status HandleAfterAll(HloInstructionPtr token) = 0;
 
   // Invoked to inform the visitor that the traversal has completed, and that
   // the root was "root".
-  virtual Status FinishVisit(HloInstructionPtr root) = 0;
+  virtual absl::Status FinishVisit(HloInstructionPtr root) = 0;
 
   // 3 possible visitation states of HLO instructions. Each instruction's
   // state only flows one way: kNotVisited -> kVisiting -> kVisited.
@@ -404,7 +405,7 @@ class DfsHloVisitorBase {
   //
   // Overriding methods should call DfsHloVisitor::Preprocess before doing their
   // own preprocessing.
-  virtual Status Preprocess(HloInstructionPtr hlo);
+  virtual absl::Status Preprocess(HloInstructionPtr hlo);
 
   // This method should be overridden by subclasses that wish to run some
   // operation on an op after its Handle* visitor method is called. See
@@ -412,7 +413,7 @@ class DfsHloVisitorBase {
   //
   // Overriding methods should call DfsHloVisitor::Postprocess after doing their
   // own postprocessing.
-  virtual Status Postprocess(HloInstructionPtr hlo);
+  virtual absl::Status Postprocess(HloInstructionPtr hlo);
 
  private:
   absl::flat_hash_map<int, VisitState> visit_state_;
diff --git a/third_party/xla/xla/hlo/ir/dfs_hlo_visitor_with_default.h b/third_party/xla/xla/hlo/ir/dfs_hlo_visitor_with_default.h
index 93b085772cd624..5a6b46fbf331c5 100644
--- a/third_party/xla/xla/hlo/ir/dfs_hlo_visitor_with_default.h
+++ b/third_party/xla/xla/hlo/ir/dfs_hlo_visitor_with_default.h
@@ -48,239 +48,244 @@ class DfsHloVisitorWithDefaultBase
   ~DfsHloVisitorWithDefaultBase() override = default;
 
   // Default action performed on HloInstruction.
-  virtual Status DefaultAction(HloInstructionPtr hlo_instruction) = 0;
+  virtual absl::Status DefaultAction(HloInstructionPtr hlo_instruction) = 0;
 
-  Status HandleElementwiseUnary(HloInstructionPtr hlo) override {
+  absl::Status HandleElementwiseUnary(HloInstructionPtr hlo) override {
     return DefaultAction(hlo);
   }
-  Status HandleElementwiseBinary(HloInstructionPtr hlo) override {
+  absl::Status HandleElementwiseBinary(HloInstructionPtr hlo) override {
     return DefaultAction(hlo);
   }
 
-  Status HandleBatchNormTraining(HloInstructionPtr hlo) override {
+  absl::Status HandleBatchNormTraining(HloInstructionPtr hlo) override {
     return DefaultAction(hlo);
   }
 
-  Status HandleBatchNormInference(HloInstructionPtr hlo) override {
+  absl::Status HandleBatchNormInference(HloInstructionPtr hlo) override {
     return DefaultAction(hlo);
   }
 
-  Status HandleBatchNormGrad(HloInstructionPtr hlo) override {
+  absl::Status HandleBatchNormGrad(HloInstructionPtr hlo) override {
     return DefaultAction(hlo);
   }
 
-  Status HandleClamp(HloInstructionPtr clamp) override {
+  absl::Status HandleClamp(HloInstructionPtr clamp) override {
     return DefaultAction(clamp);
   }
-  Status HandleConcatenate(HloInstructionPtr concatenate) override {
+  absl::Status HandleConcatenate(HloInstructionPtr concatenate) override {
     return DefaultAction(concatenate);
   }
-  Status HandleSelect(HloInstructionPtr select) override {
+  absl::Status HandleSelect(HloInstructionPtr select) override {
     return DefaultAction(select);
   }
-  Status HandleDot(HloInstructionPtr dot) override {
+  absl::Status HandleDot(HloInstructionPtr dot) override {
     return DefaultAction(dot);
   }
-  Status HandleConvolution(HloInstructionPtr convolution) override {
+  absl::Status HandleConvolution(HloInstructionPtr convolution) override {
     return DefaultAction(convolution);
   }
-  Status HandleFft(HloInstructionPtr fft) override {
+  absl::Status HandleFft(HloInstructionPtr fft) override {
     return DefaultAction(fft);
   }
-  Status HandleTriangularSolve(HloInstructionPtr hlo) override {
+  absl::Status HandleTriangularSolve(HloInstructionPtr hlo) override {
     return DefaultAction(hlo);
   }
-  Status HandleCholesky(HloInstructionPtr hlo) override {
+  absl::Status HandleCholesky(HloInstructionPtr hlo) override {
     return DefaultAction(hlo);
   }
-  Status HandleOptimizationBarrier(HloInstructionPtr hlo) override {
+  absl::Status HandleOptimizationBarrier(HloInstructionPtr hlo) override {
     return DefaultAction(hlo);
   }
-  Status HandleAllGather(HloInstructionPtr crs) override {
+  absl::Status HandleAllGather(HloInstructionPtr crs) override {
     return DefaultAction(crs);
   }
-  Status HandleAllGatherStart(HloInstructionPtr crs) override {
+  absl::Status HandleAllGatherStart(HloInstructionPtr crs) override {
     return DefaultAction(crs);
   }
-  Status HandleAllGatherDone(HloInstructionPtr crs) override {
+  absl::Status HandleAllGatherDone(HloInstructionPtr crs) override {
     return DefaultAction(crs);
   }
-  Status HandleAllReduce(HloInstructionPtr crs) override {
+  absl::Status HandleAllReduce(HloInstructionPtr crs) override {
     return DefaultAction(crs);
   }
-  Status HandleReduceScatter(HloInstructionPtr hlo) override {
+  absl::Status HandleReduceScatter(HloInstructionPtr hlo) override {
     return DefaultAction(hlo);
   }
-  Status HandleAllReduceStart(HloInstructionPtr hlo) override {
+  absl::Status HandleAllReduceStart(HloInstructionPtr hlo) override {
     return DefaultAction(hlo);
   }
-  Status HandleAllReduceDone(HloInstructionPtr hlo) override {
+  absl::Status HandleAllReduceDone(HloInstructionPtr hlo) override {
     return DefaultAction(hlo);
   }
-  Status HandleAllToAll(HloInstructionPtr hlo) override {
+  absl::Status HandleAllToAll(HloInstructionPtr hlo) override {
     return DefaultAction(hlo);
   }
-  Status HandleCollectiveBroadcast(HloInstructionPtr hlo) override {
+  absl::Status HandleCollectiveBroadcast(HloInstructionPtr hlo) override {
     return DefaultAction(hlo);
   }
-  Status HandleCollectivePermute(HloInstructionPtr hlo) override {
+  absl::Status HandleCollectivePermute(HloInstructionPtr hlo) override {
     return DefaultAction(hlo);
   }
-  Status HandleCollectivePermuteStart(HloInstructionPtr hlo) override {
+  absl::Status HandleCollectivePermuteStart(HloInstructionPtr hlo) override {
     return DefaultAction(hlo);
   }
-  Status HandleCollectivePermuteDone(HloInstructionPtr hlo) override {
+  absl::Status HandleCollectivePermuteDone(HloInstructionPtr hlo) override {
     return DefaultAction(hlo);
   }
-  Status HandleReplicaId(HloInstructionPtr hlo) override {
+  absl::Status HandleReplicaId(HloInstructionPtr hlo) override {
     return DefaultAction(hlo);
   }
-  Status HandlePartitionId(HloInstructionPtr hlo) override {
+  absl::Status HandlePartitionId(HloInstructionPtr hlo) override {
     return DefaultAction(hlo);
   }
-  Status HandleRng(HloInstructionPtr random) override {
+  absl::Status HandleRng(HloInstructionPtr random) override {
     return DefaultAction(random);
   }
-  Status HandleRngBitGenerator(HloInstructionPtr random) override {
+  absl::Status HandleRngBitGenerator(HloInstructionPtr random) override {
     return DefaultAction(random);
   }
-  Status HandleRngGetAndUpdateState(HloInstructionPtr random) override {
+  absl::Status HandleRngGetAndUpdateState(HloInstructionPtr random) override {
     return DefaultAction(random);
   }
-  Status HandleInfeed(HloInstructionPtr infeed) override {
+  absl::Status HandleInfeed(HloInstructionPtr infeed) override {
     return DefaultAction(infeed);
   }
-  Status HandleOutfeed(HloInstructionPtr outfeed) override {
+  absl::Status HandleOutfeed(HloInstructionPtr outfeed) override {
     return DefaultAction(outfeed);
   }
-  Status HandleReverse(HloInstructionPtr reverse) override {
+  absl::Status HandleReverse(HloInstructionPtr reverse) override {
     return DefaultAction(reverse);
   }
-  Status HandleSort(HloInstructionPtr sort) override {
+  absl::Status HandleSort(HloInstructionPtr sort) override {
     return DefaultAction(sort);
   }
-  Status HandleConstant(HloInstructionPtr constant) override {
+  absl::Status HandleConstant(HloInstructionPtr constant) override {
     return DefaultAction(constant);
   }
-  Status HandleIota(HloInstructionPtr iota) override {
+  absl::Status HandleIota(HloInstructionPtr iota) override {
     return DefaultAction(iota);
   }
-  Status HandleGetTupleElement(HloInstructionPtr get_tuple_element) override {
+  absl::Status HandleGetTupleElement(
+      HloInstructionPtr get_tuple_element) override {
     return DefaultAction(get_tuple_element);
   }
-  Status HandleParameter(HloInstructionPtr parameter) override {
+  absl::Status HandleParameter(HloInstructionPtr parameter) override {
     return DefaultAction(parameter);
   }
-  Status HandleFusion(HloInstructionPtr fusion) override {
+  absl::Status HandleFusion(HloInstructionPtr fusion) override {
     return DefaultAction(fusion);
   }
-  Status HandleCall(HloInstructionPtr call) override {
+  absl::Status HandleCall(HloInstructionPtr call) override {
     return DefaultAction(call);
   }
-  Status HandleCustomCall(HloInstructionPtr custom_call) override {
+  absl::Status HandleCustomCall(HloInstructionPtr custom_call) override {
     return DefaultAction(custom_call);
   }
-  Status HandleSlice(HloInstructionPtr slice) override {
+  absl::Status HandleSlice(HloInstructionPtr slice) override {
     return DefaultAction(slice);
   }
-  Status HandleDynamicSlice(HloInstructionPtr dynamic_slice) override {
+  absl::Status HandleDynamicSlice(HloInstructionPtr dynamic_slice) override {
     return DefaultAction(dynamic_slice);
   }
-  Status HandleDynamicUpdateSlice(
+  absl::Status HandleDynamicUpdateSlice(
       HloInstructionPtr dynamic_update_slice) override {
     return DefaultAction(dynamic_update_slice);
   }
-  Status HandleTuple(HloInstructionPtr tuple) override {
+  absl::Status HandleTuple(HloInstructionPtr tuple) override {
     return DefaultAction(tuple);
   }
-  Status HandleMap(HloInstructionPtr map) override {
+  absl::Status HandleMap(HloInstructionPtr map) override {
     return DefaultAction(map);
   }
-  Status HandleReduce(HloInstructionPtr reduce) override {
+  absl::Status HandleReduce(HloInstructionPtr reduce) override {
     return DefaultAction(reduce);
   }
-  Status HandleReduceWindow(HloInstructionPtr reduce_window) override {
+  absl::Status HandleReduceWindow(HloInstructionPtr reduce_window) override {
     return DefaultAction(reduce_window);
   }
-  Status HandleSelectAndScatter(HloInstructionPtr select_and_scatter) override {
+  absl::Status HandleSelectAndScatter(
+      HloInstructionPtr select_and_scatter) override {
     return DefaultAction(select_and_scatter);
   }
-  Status HandleBitcast(HloInstructionPtr bitcast) override {
+  absl::Status HandleBitcast(HloInstructionPtr bitcast) override {
     return DefaultAction(bitcast);
   }
-  Status HandleBroadcast(HloInstructionPtr broadcast) override {
+  absl::Status HandleBroadcast(HloInstructionPtr broadcast) override {
     return DefaultAction(broadcast);
   }
-  Status HandlePad(HloInstructionPtr pad) override {
+  absl::Status HandlePad(HloInstructionPtr pad) override {
     return DefaultAction(pad);
   }
-  Status HandleDynamicReshape(HloInstructionPtr dynamic_reshape) override {
+  absl::Status HandleDynamicReshape(
+      HloInstructionPtr dynamic_reshape) override {
     return DefaultAction(dynamic_reshape);
   }
-  Status HandleReshape(HloInstructionPtr reshape) override {
+  absl::Status HandleReshape(HloInstructionPtr reshape) override {
     return DefaultAction(reshape);
   }
-  Status HandleTranspose(HloInstructionPtr transpose) override {
+  absl::Status HandleTranspose(HloInstructionPtr transpose) override {
     return DefaultAction(transpose);
   }
-  Status HandleWhile(HloInstructionPtr xla_while) override {
+  absl::Status HandleWhile(HloInstructionPtr xla_while) override {
     return DefaultAction(xla_while);
   }
-  Status HandleConditional(HloInstructionPtr conditional) override {
+  absl::Status HandleConditional(HloInstructionPtr conditional) override {
     return DefaultAction(conditional);
   }
-  Status HandleAsyncStart(HloInstructionPtr async_start) override {
+  absl::Status HandleAsyncStart(HloInstructionPtr async_start) override {
     return DefaultAction(async_start);
   }
-  Status HandleAsyncUpdate(HloInstructionPtr async_update) override {
+  absl::Status HandleAsyncUpdate(HloInstructionPtr async_update) override {
     return DefaultAction(async_update);
   }
-  Status HandleAsyncDone(HloInstructionPtr async_done) override {
+  absl::Status HandleAsyncDone(HloInstructionPtr async_done) override {
     return DefaultAction(async_done);
   }
-  Status HandleCopyStart(HloInstructionPtr copy_start) override {
+  absl::Status HandleCopyStart(HloInstructionPtr copy_start) override {
     return DefaultAction(copy_start);
   }
-  Status HandleCopyDone(HloInstructionPtr copy_done) override {
+  absl::Status HandleCopyDone(HloInstructionPtr copy_done) override {
     return DefaultAction(copy_done);
   }
-  Status HandleRecv(HloInstructionPtr recv) override {
+  absl::Status HandleRecv(HloInstructionPtr recv) override {
     return DefaultAction(recv);
   }
-  Status HandleRecvDone(HloInstructionPtr recv_done) override {
+  absl::Status HandleRecvDone(HloInstructionPtr recv_done) override {
     return DefaultAction(recv_done);
   }
-  Status HandleSend(HloInstructionPtr send) override {
+  absl::Status HandleSend(HloInstructionPtr send) override {
     return DefaultAction(send);
   }
-  Status HandleTopK(HloInstructionPtr topk) override {
+  absl::Status HandleTopK(HloInstructionPtr topk) override {
     return DefaultAction(topk);
   }
-  Status HandleSendDone(HloInstructionPtr send_done) override {
+  absl::Status HandleSendDone(HloInstructionPtr send_done) override {
     return DefaultAction(send_done);
   }
-  Status HandleGather(HloInstructionPtr gather) override {
+  absl::Status HandleGather(HloInstructionPtr gather) override {
     return DefaultAction(gather);
   }
-  Status HandleScatter(HloInstructionPtr scatter) override {
+  absl::Status HandleScatter(HloInstructionPtr scatter) override {
     return DefaultAction(scatter);
   }
-  Status HandleAfterAll(HloInstructionPtr token) override {
+  absl::Status HandleAfterAll(HloInstructionPtr token) override {
     return DefaultAction(token);
   }
-  Status HandleGetDimensionSize(HloInstructionPtr get_size) override {
+  absl::Status HandleGetDimensionSize(HloInstructionPtr get_size) override {
     return DefaultAction(get_size);
   }
-  Status HandleSetDimensionSize(HloInstructionPtr get_size) override {
+  absl::Status HandleSetDimensionSize(HloInstructionPtr get_size) override {
     return DefaultAction(get_size);
   }
-  Status HandleAddDependency(HloInstructionPtr add_dependency) override {
+  absl::Status HandleAddDependency(HloInstructionPtr add_dependency) override {
     return DefaultAction(add_dependency);
   }
 
   // Invoked to inform the visitor that the traversal has completed, and that
   // the root was "root".
-  Status FinishVisit(HloInstructionPtr /*root*/) override { return OkStatus(); }
+  absl::Status FinishVisit(HloInstructionPtr /*root*/) override {
+    return OkStatus();
+  }
 
  private:
   DfsHloVisitorWithDefaultBase(const DfsHloVisitorWithDefaultBase&) = delete;
@@ -304,7 +309,7 @@ class DfsHloRewriteVisitor : public DfsHloVisitorWithDefault {
   absl::StatusOr<bool> RunOnModule(
       HloModule* module,
       const absl::flat_hash_set<absl::string_view>& execution_threads = {}) {
-    Status status;
+    absl::Status status;
     for (HloComputation* computation :
          module->MakeNonfusionComputations(execution_threads)) {
       status = computation->Accept(this);
@@ -314,7 +319,7 @@ class DfsHloRewriteVisitor : public DfsHloVisitorWithDefault {
   }
 
   // Default visitor action is to do nothing and return OK.
-  Status DefaultAction(HloInstruction* /*hlo_instruction*/) override {
+  absl::Status DefaultAction(HloInstruction* /*hlo_instruction*/) override {
     return OkStatus();
   }
 
@@ -323,14 +328,14 @@ class DfsHloRewriteVisitor : public DfsHloVisitorWithDefault {
  protected:
   // Replaces the existing HLO instruction old_instruction, with
   // new_instruction, and marks the optimizer status as changed.
-  // Returns the Status representing the result of the replace operation.
-  Status ReplaceWithNewInstruction(
+  // Returns the absl::Status representing the result of the replace operation.
+  absl::Status ReplaceWithNewInstruction(
       HloInstruction* old_instruction,
       std::unique_ptr<HloInstruction> new_instruction) {
     VLOG(3) << "Replacing instruction:" << "\n  old: "
             << old_instruction->ToString()
             << "\n  new: " << new_instruction->ToString();
-    Status status = old_instruction->parent()->ReplaceWithNewInstruction(
+    absl::Status status = old_instruction->parent()->ReplaceWithNewInstruction(
         old_instruction, std::move(new_instruction));
     if (ABSL_PREDICT_TRUE(status.ok())) {
       changed_ = true;
@@ -340,7 +345,7 @@ class DfsHloRewriteVisitor : public DfsHloVisitorWithDefault {
 
   // Replaces the existing HLO instruction old_instruction, with
   // new_instruction, and marks the optimizer status as changed.
-  // Returns the Status representing the result of the replace operation.
+  // Returns the absl::Status representing the result of the replace operation.
   absl::StatusOr<bool> ReplaceInstruction(HloInstruction* old_instruction,
                                           HloInstruction* new_instruction,
                                           bool preserve_sharding) {
@@ -356,8 +361,8 @@ class DfsHloRewriteVisitor : public DfsHloVisitorWithDefault {
     return changed_or;
   }
 
-  Status ReplaceInstruction(HloInstruction* old_instruction,
-                            HloInstruction* new_instruction) {
+  absl::Status ReplaceInstruction(HloInstruction* old_instruction,
+                                  HloInstruction* new_instruction) {
     absl::StatusOr<bool> changed_or =
         ReplaceInstruction(old_instruction, new_instruction,
                            /*preserve_sharding=*/false);
@@ -375,7 +380,8 @@ class DfsHloRewriteVisitor : public DfsHloVisitorWithDefault {
 };
 
 // (Const)FunctionVisitor lets you transform an
-// std::function<Status((const) HloInstruction*)> into a (Const)DfsHloVisitor.
+// std::function<absl::Status((const) HloInstruction*)> into a
+// (Const)DfsHloVisitor.
 //
 // This is useful if you have code that needs to handle visitors in the form of
 // both std::function and DfsHloVisitor.  You can wrap the function in a
@@ -385,10 +391,10 @@ class FunctionVisitorBase
     : public DfsHloVisitorWithDefaultBase<HloInstructionPtr> {
  public:
   explicit FunctionVisitorBase(
-      std::function<Status(HloInstructionPtr)> visitor_func)
+      std::function<absl::Status(HloInstructionPtr)> visitor_func)
       : visitor_func_(std::move(visitor_func)) {}
 
-  Status DefaultAction(HloInstructionPtr hlo_instruction) override {
+  absl::Status DefaultAction(HloInstructionPtr hlo_instruction) override {
     return visitor_func_(hlo_instruction);
   }
 
@@ -396,7 +402,7 @@ class FunctionVisitorBase
   FunctionVisitorBase(const FunctionVisitorBase&) = delete;
   FunctionVisitorBase& operator=(const FunctionVisitorBase&) = delete;
 
-  std::function<Status(HloInstructionPtr)> visitor_func_;
+  std::function<absl::Status(HloInstructionPtr)> visitor_func_;
 };
 
 using FunctionVisitor = FunctionVisitorBase<HloInstruction*>;
diff --git a/third_party/xla/xla/hlo/ir/dynamic_parameter_binding.cc b/third_party/xla/xla/hlo/ir/dynamic_parameter_binding.cc
index ca96fbe1351136..b1fc5fcab3aa5b 100644
--- a/third_party/xla/xla/hlo/ir/dynamic_parameter_binding.cc
+++ b/third_party/xla/xla/hlo/ir/dynamic_parameter_binding.cc
@@ -31,7 +31,7 @@ limitations under the License.
 
 namespace xla {
 
-Status DynamicParameterBinding::Bind(
+absl::Status DynamicParameterBinding::Bind(
     const DynamicSizeParameter& dynamic_parameter,
     const DynamicDimension& dynamic_dimension) {
   auto result = bindings_.emplace(dynamic_dimension, dynamic_parameter);
@@ -67,18 +67,18 @@ std::string DynamicParameterBinding::ToString() const {
   return absl::StrJoin(pieces, "\n");
 }
 
-Status DynamicParameterBinding::ForEachBinding(BindingFn fn) const {
+absl::Status DynamicParameterBinding::ForEachBinding(BindingFn fn) const {
   for (const auto& binding : bindings_) {
     TF_RETURN_IF_ERROR(fn(binding.second, binding.first));
   }
   return OkStatus();
 }
 
-Status DynamicParameterBinding::Verify(
+absl::Status DynamicParameterBinding::Verify(
     const HloComputation& computation) const {
   return ForEachBinding([&](const DynamicSizeParameter& dynamic_parameter,
                             const DynamicDimension& dynamic_dimension)
-                            -> Status {
+                            -> absl::Status {
     TF_RET_CHECK(dynamic_parameter.parameter_num >= 0 &&
                  dynamic_parameter.parameter_num <
                      computation.num_parameters());
diff --git a/third_party/xla/xla/hlo/ir/dynamic_parameter_binding.h b/third_party/xla/xla/hlo/ir/dynamic_parameter_binding.h
index 70a199db635321..7c78ef3c8d659d 100644
--- a/third_party/xla/xla/hlo/ir/dynamic_parameter_binding.h
+++ b/third_party/xla/xla/hlo/ir/dynamic_parameter_binding.h
@@ -81,8 +81,8 @@ class DynamicParameterBinding {
   // Adds binding which indicates that the dimension indicated by
   // `dynamic_dimension` is dynamic, and its runtime size is represented by
   // `dynamic_parameter`.
-  Status Bind(const DynamicSizeParameter& dynamic_parameter,
-              const DynamicDimension& dynamic_dimension);
+  absl::Status Bind(const DynamicSizeParameter& dynamic_parameter,
+                    const DynamicDimension& dynamic_dimension);
 
   // Returns the parameter and the index representing the runtime size of
   // dimension `dim_num` of parameter `param_num` at `param_index`.
@@ -92,17 +92,17 @@ class DynamicParameterBinding {
       const DynamicDimension& dynamic_dimension) const;
 
   using BindingFn =
-      std::function<Status(const DynamicSizeParameter& dynamic_parameter,
-                           const DynamicDimension& dynamic_dimension)>;
+      std::function<absl::Status(const DynamicSizeParameter& dynamic_parameter,
+                                 const DynamicDimension& dynamic_dimension)>;
 
   // Iterate through each binding.
-  Status ForEachBinding(BindingFn fn) const;
+  absl::Status ForEachBinding(BindingFn fn) const;
 
   std::string ToString() const;
 
   // Verifies that the given binding is valid for the given computation.
   // Specifically, the binding's parameter and parameter size should be valid.
-  Status Verify(const HloComputation& computation) const;
+  absl::Status Verify(const HloComputation& computation) const;
 
   // Returns true iff there are no bindings.
   bool empty() const { return bindings_.empty(); }
diff --git a/third_party/xla/xla/hlo/ir/hlo_computation.cc b/third_party/xla/xla/hlo/ir/hlo_computation.cc
index f04850e93ee7bc..b424a4b374dc76 100644
--- a/third_party/xla/xla/hlo/ir/hlo_computation.cc
+++ b/third_party/xla/xla/hlo/ir/hlo_computation.cc
@@ -277,7 +277,7 @@ HloInstruction* HloComputation::AddEntryComputationParameter(
   return instructions_.back().get();
 }
 
-Status HloComputation::ReplaceEntryComputationParameter(
+absl::Status HloComputation::ReplaceEntryComputationParameter(
     int64_t param_no, HloInstruction* old_instruction,
     std::unique_ptr<HloInstruction> instruction) {
   CHECK_GE(param_no, 0);
@@ -297,7 +297,7 @@ Status HloComputation::ReplaceEntryComputationParameter(
   return ForceRemoveInstruction(old_instruction);
 }
 
-Status HloComputation::RemoveParameter(int64_t param_no) {
+absl::Status HloComputation::RemoveParameter(int64_t param_no) {
   CHECK_GE(param_no, 0);
   CHECK_LT(param_no, param_instructions_.size());
   HloInstruction* param_instruction = param_instructions_[param_no];
@@ -339,15 +339,15 @@ HloInstruction* HloComputation::ReplaceParameter(
   return new_instruction;
 }
 
-Status HloComputation::RemoveUnusedParametersFromFusedComputation() {
+absl::Status HloComputation::RemoveUnusedParametersFromFusedComputation() {
   return RemoveUnusedParametersImpl(/*allow_non_fusion=*/false);
 }
 
-Status HloComputation::RemoveUnusedParametersFromAnyComputation() {
+absl::Status HloComputation::RemoveUnusedParametersFromAnyComputation() {
   return RemoveUnusedParametersImpl(/*allow_non_fusion=*/true);
 }
 
-Status HloComputation::RemoveUnusedParametersImpl(bool allow_non_fusion) {
+absl::Status HloComputation::RemoveUnusedParametersImpl(bool allow_non_fusion) {
   CHECK(allow_non_fusion || IsFusionComputation());
   int64_t removed = 0;
   for (int64_t i = 0; i < param_instructions_.size(); ++i) {
@@ -404,7 +404,7 @@ bool HloComputation::IsMarkedAsDead(const HloInstruction* inst) {
   return inst->IsMarkedAsDead();
 }
 
-Status HloComputation::RemoveInstructionAndUnusedOperands(
+absl::Status HloComputation::RemoveInstructionAndUnusedOperands(
     HloInstruction* instruction,
     std::optional<absl::FunctionRef<void(HloInstruction*)>> cleanup,
     bool ignore_control_dependencies) {
@@ -474,16 +474,17 @@ Status HloComputation::RemoveInstructionAndUnusedOperands(
   return OkStatus();
 }
 
-Status HloComputation::RemoveInstruction(HloInstruction* instruction) {
+absl::Status HloComputation::RemoveInstruction(HloInstruction* instruction) {
   return RemoveInstructionImpl(instruction, /*ignore_safety_check=*/false);
 }
 
-Status HloComputation::ForceRemoveInstruction(HloInstruction* instruction) {
+absl::Status HloComputation::ForceRemoveInstruction(
+    HloInstruction* instruction) {
   return RemoveInstructionImpl(instruction, /*ignore_safety_check=*/true);
 }
 
-Status HloComputation::RemoveInstructionImpl(HloInstruction* instruction,
-                                             bool ignore_safety_check) {
+absl::Status HloComputation::RemoveInstructionImpl(HloInstruction* instruction,
+                                                   bool ignore_safety_check) {
   VLOG(2) << "Removing instruction " << instruction << " "
           << instruction->name() << " from computation " << name();
   TF_RET_CHECK(ignore_safety_check || IsSafelyRemovable(instruction))
@@ -1029,7 +1030,7 @@ HloComputation::CreateFromProto(
     return to_proto_id[a.get()] < to_proto_id[b.get()];
   });
 
-  TF_RETURN_IF_ERROR([&]() -> Status {
+  TF_RETURN_IF_ERROR([&]() -> absl::Status {
     std::vector<bool> parameters_seen(parameter_count);
     int parameters_seen_count = 0;
     for (auto& instruction : instructions) {
@@ -1342,14 +1343,14 @@ bool HloComputation::EqualInternal(
   return true;
 }
 
-Status HloComputation::ReplaceWithNewInstruction(
+absl::Status HloComputation::ReplaceWithNewInstruction(
     HloInstruction* old_instruction,
     std::unique_ptr<HloInstruction> new_instruction) {
   return ReplaceInstruction(old_instruction,
                             AddInstruction(std::move(new_instruction)));
 }
 
-Status HloComputation::ReplaceWithNewEntryComputationParameter(
+absl::Status HloComputation::ReplaceWithNewEntryComputationParameter(
     HloInstruction* old_instruction,
     std::unique_ptr<HloInstruction> new_instruction) {
   return ReplaceInstruction(old_instruction, AddEntryComputationParameter(
@@ -1369,8 +1370,8 @@ absl::StatusOr<bool> HloComputation::ReplaceInstruction(
       relay_control_dependency, remove_unused_operands);
 }
 
-Status HloComputation::ReplaceInstruction(HloInstruction* old_instruction,
-                                          HloInstruction* new_instruction) {
+absl::Status HloComputation::ReplaceInstruction(
+    HloInstruction* old_instruction, HloInstruction* new_instruction) {
   TF_ASSIGN_OR_RETURN(bool changed,
                       ReplaceInstruction(old_instruction, new_instruction,
                                          /*preserve_sharding=*/false));
@@ -1444,7 +1445,7 @@ absl::StatusOr<bool> HloComputation::ReplaceInstructionWithDifferentShape(
   return true;
 }
 
-Status HloComputation::ReplaceInstructionWithDifferentShape(
+absl::Status HloComputation::ReplaceInstructionWithDifferentShape(
     HloInstruction* old_instruction, HloInstruction* new_instruction) {
   TF_ASSIGN_OR_RETURN(bool changed, ReplaceInstructionWithDifferentShape(
                                         old_instruction, new_instruction,
@@ -1468,7 +1469,7 @@ std::vector<HloInstruction*> HloComputation::CollectUnreachableRoots() const {
   return unreachable_roots;
 }
 
-Status HloComputation::AcceptWithOperandOrder(
+absl::Status HloComputation::AcceptWithOperandOrder(
     DfsHloVisitor* visitor,
     const HloInstruction::CompareFunction& operand_order) const {
   // Visit unreachable roots. Beware that the visitor might delete the currently
diff --git a/third_party/xla/xla/hlo/ir/hlo_computation.h b/third_party/xla/xla/hlo/ir/hlo_computation.h
index 9c63f104d58ab5..8c5b2aa9b29e72 100644
--- a/third_party/xla/xla/hlo/ir/hlo_computation.h
+++ b/third_party/xla/xla/hlo/ir/hlo_computation.h
@@ -111,8 +111,8 @@ class HloComputation {
       return AddInstruction(std::move(parameter));
     }
 
-    Status ForEachInstruction(
-        absl::FunctionRef<Status(const HloInstruction*)> func) const {
+    absl::Status ForEachInstruction(
+        absl::FunctionRef<absl::Status(const HloInstruction*)> func) const {
       for (const auto& instruction : instructions_) {
         TF_RETURN_IF_ERROR(func(instruction.get()));
       }
@@ -216,17 +216,17 @@ class HloComputation {
   // Remove the param_no'th parameter from the computation.
   // Note this is only applicatable to the computation for the fusion
   // instruction.
-  Status RemoveParameter(int64_t param_no);
+  absl::Status RemoveParameter(int64_t param_no);
 
   // Remove unused parameters from the computation.
   // Note this is only applicatable to the computation for the fusion
   // instruction.
-  Status RemoveUnusedParametersFromFusedComputation();
+  absl::Status RemoveUnusedParametersFromFusedComputation();
 
   // Remove unused parameters from the computation. Unlike
   // RemoveUnusedParametersFromFusedComputation, this function can be used
   // to remove parameters from non-fusion computations.
-  Status RemoveUnusedParametersFromAnyComputation();
+  absl::Status RemoveUnusedParametersFromAnyComputation();
 
   // Adds a new parameter instruction to a fusion computation.
   //
@@ -244,18 +244,18 @@ class HloComputation {
 
   // Replaces an old parameter with a new parameter. Adds the new parameter
   // instruction to the entry computation.  Updates users instruction.
-  Status ReplaceEntryComputationParameter(
+  absl::Status ReplaceEntryComputationParameter(
       int64_t param_no, HloInstruction* old_instruction,
       std::unique_ptr<HloInstruction> instruction);
 
   // Remove an instruction from the computation. The instruction must have no
   // users. Instruction is deallocated with this call.
-  Status RemoveInstruction(HloInstruction* instruction);
+  absl::Status RemoveInstruction(HloInstruction* instruction);
 
   // Removes an instruction from the computation. The instruction must have no
   // users. Instruction is deallocated with this call. The instruction will be
   // removed even if it is marked as not removable.
-  Status ForceRemoveInstruction(HloInstruction* instruction);
+  absl::Status ForceRemoveInstruction(HloInstruction* instruction);
 
   // Remove an instruction (including side effecting ones) from the computation
   // and also transitively any operand that has no side effect and no users post
@@ -267,7 +267,7 @@ class HloComputation {
   // the predecessors to the succesors of the removed instructions, so that the
   // logical exeuction order of the remaining unremoved instructions are
   // preserved.
-  Status RemoveInstructionAndUnusedOperands(
+  absl::Status RemoveInstructionAndUnusedOperands(
       HloInstruction* instruction,
       std::optional<absl::FunctionRef<void(HloInstruction*)>> cleanup =
           std::nullopt,
@@ -560,14 +560,14 @@ class HloComputation {
 
   // Replaces old instruction with newly created instruction. Removes old
   // instruction from computation. Updates uses and root instruction.
-  Status ReplaceWithNewInstruction(
+  absl::Status ReplaceWithNewInstruction(
       HloInstruction* old_instruction,
       std::unique_ptr<HloInstruction> new_instruction);
 
   // Replaces an old instruction with a newly created instruction, and adds the
   // new instruction as an entry computation's parameter. Removes old
   // instruction from computation. Updates uses and root instruction.
-  Status ReplaceWithNewEntryComputationParameter(
+  absl::Status ReplaceWithNewEntryComputationParameter(
       HloInstruction* old_instruction,
       std::unique_ptr<HloInstruction> new_instruction);
 
@@ -588,9 +588,9 @@ class HloComputation {
                                           bool remove_unused_operands = true);
 
   // Same as above, with preserve_sharding=false. Since this replacement always
-  // happens, it returns just a Status as opposed to StatusOr<bool>
-  Status ReplaceInstruction(HloInstruction* old_instruction,
-                            HloInstruction* new_instruction);
+  // happens, it returns just a absl::Status as opposed to absl::StatusOr<bool>
+  absl::Status ReplaceInstruction(HloInstruction* old_instruction,
+                                  HloInstruction* new_instruction);
 
   // Same as ReplaceInstruction, but the new instruction can have a different
   // shape.
@@ -598,8 +598,8 @@ class HloComputation {
       HloInstruction* old_instruction, HloInstruction* new_instruction,
       bool preserve_sharding, bool relay_control_dependency = false,
       bool remove_unused_operands = true);
-  Status ReplaceInstructionWithDifferentShape(HloInstruction* old_instruction,
-                                              HloInstruction* new_instruction);
+  absl::Status ReplaceInstructionWithDifferentShape(
+      HloInstruction* old_instruction, HloInstruction* new_instruction);
 
   // Set/get the module containing this computation.
   void set_parent(HloModule* module) { parent_ = module; }
@@ -613,20 +613,20 @@ class HloComputation {
   // the visitor's FinishVisit method is called once upon completion (with the
   // root instruction as the argument).
   template <typename HloInstructionPtr>
-  Status Accept(DfsHloVisitorBase<HloInstructionPtr>* visitor) const;
+  absl::Status Accept(DfsHloVisitorBase<HloInstructionPtr>* visitor) const;
 
   // Same as Accept() above, but the order of operand and control predecessor
   // visitation is determined by the given operand order; if compare(A, B) ==
   // true, A is visited before B.
-  Status AcceptWithOperandOrder(
+  absl::Status AcceptWithOperandOrder(
       DfsHloVisitor* visitor,
       const HloInstruction::CompareFunction& operand_order) const;
 
   // Visit every node in the computation in the given order. 'order' must
   // be a topological sort of all instructions in the computation.
   template <typename HloInstructionPtr>
-  Status AcceptOrdered(DfsHloVisitorBase<HloInstructionPtr>* visitor,
-                       absl::Span<HloInstruction* const> order) const;
+  absl::Status AcceptOrdered(DfsHloVisitorBase<HloInstructionPtr>* visitor,
+                             absl::Span<HloInstruction* const> order) const;
 
   // Returns a deep copy of this computation including all instructions.
   // If the clone context is specified, it will be populated with the cloned
@@ -908,10 +908,10 @@ class HloComputation {
       const ChannelDependencies& channel_dependencies, VisitMap& visited,
       std::vector<HloInstruction*>* dfs_stack_scratch) const;
 
-  Status RemoveUnusedParametersImpl(bool allow_non_fusion);
+  absl::Status RemoveUnusedParametersImpl(bool allow_non_fusion);
 
-  Status RemoveInstructionImpl(HloInstruction* instruction,
-                               bool ignore_safety_check);
+  absl::Status RemoveInstructionImpl(HloInstruction* instruction,
+                                     bool ignore_safety_check);
 
   enum class InstructionType : uint8_t {
     kUnset,
@@ -989,7 +989,7 @@ class HloComputation {
 };
 
 template <typename HloInstructionPtr>
-Status HloComputation::Accept(
+absl::Status HloComputation::Accept(
     DfsHloVisitorBase<HloInstructionPtr>* visitor) const {
   // Visit unreachable roots. Beware that the visitor might delete the currently
   // visited root, which would invalidate iterators if the unreachable roots
@@ -1004,11 +1004,11 @@ Status HloComputation::Accept(
 }
 
 // Explicit instantiations.
-template Status HloComputation::Accept(DfsHloVisitor* visitor) const;
-template Status HloComputation::Accept(ConstDfsHloVisitor* visitor) const;
+template absl::Status HloComputation::Accept(DfsHloVisitor* visitor) const;
+template absl::Status HloComputation::Accept(ConstDfsHloVisitor* visitor) const;
 
 template <typename HloInstructionPtr>
-Status HloComputation::AcceptOrdered(
+absl::Status HloComputation::AcceptOrdered(
     DfsHloVisitorBase<HloInstructionPtr>* visitor,
     absl::Span<HloInstruction* const> order) const {
   VLOG(3) << "Accepting visitor with order.";
@@ -1034,9 +1034,9 @@ Status HloComputation::AcceptOrdered(
 }
 
 // Explicit instantiations.
-template Status HloComputation::AcceptOrdered(
+template absl::Status HloComputation::AcceptOrdered(
     DfsHloVisitor*, absl::Span<HloInstruction* const>) const;
-template Status HloComputation::AcceptOrdered(
+template absl::Status HloComputation::AcceptOrdered(
     ConstDfsHloVisitor*, absl::Span<HloInstruction* const>) const;
 
 }  // namespace xla
diff --git a/third_party/xla/xla/hlo/ir/hlo_input_output_alias_config.cc b/third_party/xla/xla/hlo/ir/hlo_input_output_alias_config.cc
index 3dca857c43a84e..23a886814ae08e 100644
--- a/third_party/xla/xla/hlo/ir/hlo_input_output_alias_config.cc
+++ b/third_party/xla/xla/hlo/ir/hlo_input_output_alias_config.cc
@@ -43,7 +43,7 @@ bool HloInputOutputAliasConfig::OutputHasAlias(
   return alias_.element(output_index).has_value();
 }
 
-Status HloInputOutputAliasConfig::SetUpAlias(
+absl::Status HloInputOutputAliasConfig::SetUpAlias(
     const ShapeIndex& output_index, int64_t param_number,
     const ShapeIndex& param_index,
     HloInputOutputAliasConfig::AliasKind must_alias) {
@@ -182,7 +182,7 @@ void HloInputOutputAliasConfig::ForEachAlias(AliasFn fn) const {
       });
 }
 
-Status HloInputOutputAliasConfig::ForEachAliasWithStatus(
+absl::Status HloInputOutputAliasConfig::ForEachAliasWithStatus(
     AliasFnWithStatus fn) const {
   return alias_.ForEachElementWithStatus(
       [&](const ShapeIndex& output_index, std::optional<Alias> aliased) {
@@ -193,7 +193,7 @@ Status HloInputOutputAliasConfig::ForEachAliasWithStatus(
       });
 }
 
-Status HloInputOutputAliasConfig::Verify(
+absl::Status HloInputOutputAliasConfig::Verify(
     const HloModule& module,
     absl::FunctionRef<int64_t(const Shape&)> size_func) const {
   std::vector<ShapeTree<bool>> param_has_seen;
@@ -203,7 +203,7 @@ Status HloInputOutputAliasConfig::Verify(
     param_has_seen.emplace_back(param->shape());
   }
   return ForEachAliasWithStatus([&](const ShapeIndex& output_index,
-                                    const Alias& alias) -> Status {
+                                    const Alias& alias) -> absl::Status {
     TF_RET_CHECK(0 <= alias.parameter_number);
     TF_RET_CHECK(entry->num_parameters() > alias.parameter_number);
     const Shape& param_shape =
@@ -250,8 +250,8 @@ std::ostream& operator<<(std::ostream& out,
   return out;
 }
 
-Status HloBufferDonorConfig::AddBufferDonor(int64_t param_number,
-                                            const ShapeIndex& param_index) {
+absl::Status HloBufferDonorConfig::AddBufferDonor(
+    int64_t param_number, const ShapeIndex& param_index) {
   TF_RET_CHECK(param_number >= 0) << param_number;
   VLOG(4) << "Register the parameter " << param_number << " at index "
           << param_index.ToString() << " as a buffer donor.";
@@ -259,8 +259,8 @@ Status HloBufferDonorConfig::AddBufferDonor(int64_t param_number,
   return OkStatus();
 }
 
-Status HloBufferDonorConfig::RemoveBufferDonor(int64_t param_number,
-                                               const ShapeIndex& param_index) {
+absl::Status HloBufferDonorConfig::RemoveBufferDonor(
+    int64_t param_number, const ShapeIndex& param_index) {
   TF_RET_CHECK(param_number >= 0) << param_number;
   buffer_donor_.erase(BufferDonor(param_number, param_index));
   return OkStatus();
@@ -319,7 +319,7 @@ bool HloBufferDonorConfig::ParameterIsBufferDonor(
   return it != buffer_donor_.end();
 }
 
-Status HloBufferDonorConfig::Verify(const HloModule& module) const {
+absl::Status HloBufferDonorConfig::Verify(const HloModule& module) const {
   const HloComputation* entry = module.entry_computation();
   const auto& alias_config = module.input_output_alias_config();
   for (const auto& donor : buffer_donor_) {
diff --git a/third_party/xla/xla/hlo/ir/hlo_instruction.cc b/third_party/xla/xla/hlo/ir/hlo_instruction.cc
index b5124abd24c7c6..9b9bdcbaa7d9d4 100644
--- a/third_party/xla/xla/hlo/ir/hlo_instruction.cc
+++ b/third_party/xla/xla/hlo/ir/hlo_instruction.cc
@@ -97,7 +97,7 @@ const HloInstruction::Rare* const HloInstruction::kEmptyRare =
 namespace {
 // Specialization for erasing from PtrVec<T>.
 template <typename T>
-Status EraseElementFromVector(PtrVec<T>* container, T value) {
+absl::Status EraseElementFromVector(PtrVec<T>* container, T value) {
   // absl::c_find returns a const_iterator which does not seem to work on
   // gcc 4.8.4, and this breaks the ubuntu/xla_gpu build bot.
   auto it = std::find(container->begin(), container->end(), value);
@@ -2698,7 +2698,8 @@ HloInstruction::InstructionVector HloInstruction::unique_operands() const {
   return unique;
 }
 
-Status HloInstruction::AddControlDependencyTo(HloInstruction* instruction) {
+absl::Status HloInstruction::AddControlDependencyTo(
+    HloInstruction* instruction) {
   TF_RET_CHECK(instruction->parent() == parent());
   if (!absl::c_linear_search(control_successors(), instruction)) {
     mutable_rare()->control_successors.push_back(instruction);
@@ -2709,7 +2710,8 @@ Status HloInstruction::AddControlDependencyTo(HloInstruction* instruction) {
   return OkStatus();
 }
 
-Status HloInstruction::RemoveControlDependencyTo(HloInstruction* instruction) {
+absl::Status HloInstruction::RemoveControlDependencyTo(
+    HloInstruction* instruction) {
   TF_RET_CHECK(instruction->parent() == parent());
   if (has_rare()) {
     TF_RETURN_IF_ERROR(EraseElementFromVector(
@@ -2722,7 +2724,7 @@ Status HloInstruction::RemoveControlDependencyTo(HloInstruction* instruction) {
   return OkStatus();
 }
 
-Status HloInstruction::DropAllControlDeps() {
+absl::Status HloInstruction::DropAllControlDeps() {
   if (has_rare()) {
     for (auto* ctrl_succ : rare()->control_successors) {
       TF_RETURN_IF_ERROR(EraseElementFromVector(
@@ -2739,7 +2741,7 @@ Status HloInstruction::DropAllControlDeps() {
   return OkStatus();
 }
 
-Status HloInstruction::SafelyDropAllControlDependencies() {
+absl::Status HloInstruction::SafelyDropAllControlDependencies() {
   // Add all pairs of transitive dependencies from predecessors to successors.
   if (has_rare()) {
     for (HloInstruction* predecessor : rare()->control_predecessors) {
@@ -2757,8 +2759,8 @@ bool HloInstruction::HasControlDependencies() const {
   return (!r->control_predecessors.empty() || !r->control_successors.empty());
 }
 
-Status HloInstruction::CopyAllControlDepsTo(HloInstruction* start,
-                                            HloInstruction* end) const {
+absl::Status HloInstruction::CopyAllControlDepsTo(HloInstruction* start,
+                                                  HloInstruction* end) const {
   for (auto* ctrl_pred : control_predecessors()) {
     TF_RETURN_IF_ERROR(ctrl_pred->AddControlDependencyTo(start));
   }
@@ -3029,8 +3031,8 @@ bool HloInstruction::IdenticalSlowPath(
   return false;
 }
 
-Status HloInstruction::ReplaceUseWith(HloInstruction* user,
-                                      HloInstruction* new_producer) {
+absl::Status HloInstruction::ReplaceUseWith(HloInstruction* user,
+                                            HloInstruction* new_producer) {
   TF_RET_CHECK(
       ShapeUtil::CompatibleIgnoringFpPrecision(shape(), new_producer->shape()))
       << "this shape: " << ShapeUtil::HumanString(shape())
@@ -3039,7 +3041,7 @@ Status HloInstruction::ReplaceUseWith(HloInstruction* user,
   return ReplaceUseWithDifferentShape(user, new_producer);
 }
 
-Status HloInstruction::ReplaceUseWithDifferentShape(
+absl::Status HloInstruction::ReplaceUseWithDifferentShape(
     HloInstruction* user, HloInstruction* new_producer) {
   VLOG(3) << "Replacing uses of " << name() << " in " << user->name()
           << " with " << new_producer->name();
@@ -3058,8 +3060,9 @@ Status HloInstruction::ReplaceUseWithDifferentShape(
   return OkStatus();
 }
 
-Status HloInstruction::ReplaceUseWith(HloInstruction* user, int operand_number,
-                                      HloInstruction* new_producer) {
+absl::Status HloInstruction::ReplaceUseWith(HloInstruction* user,
+                                            int operand_number,
+                                            HloInstruction* new_producer) {
   TF_RET_CHECK(
       ShapeUtil::CompatibleIgnoringFpPrecision(shape(), new_producer->shape()))
       << "this shape: " << ShapeUtil::HumanString(shape())
@@ -3068,7 +3071,7 @@ Status HloInstruction::ReplaceUseWith(HloInstruction* user, int operand_number,
   return ReplaceUseWithDifferentShape(user, operand_number, new_producer);
 }
 
-Status HloInstruction::ReplaceUseWithDifferentShape(
+absl::Status HloInstruction::ReplaceUseWithDifferentShape(
     HloInstruction* user, int operand_number, HloInstruction* new_producer) {
   VLOG(3) << "Replacing operand " << operand_number << " of " << name()
           << " in " << user->name() << " with " << new_producer->name();
@@ -3085,8 +3088,8 @@ Status HloInstruction::ReplaceUseWithDifferentShape(
   return OkStatus();
 }
 
-Status HloInstruction::ReplaceOperandWith(int64_t operand_num,
-                                          HloInstruction* new_operand) {
+absl::Status HloInstruction::ReplaceOperandWith(int64_t operand_num,
+                                                HloInstruction* new_operand) {
   auto old_operand = operand(operand_num);
   TF_RET_CHECK(ShapeUtil::CompatibleIgnoringFpPrecision(old_operand->shape(),
                                                         new_operand->shape()))
@@ -3095,7 +3098,7 @@ Status HloInstruction::ReplaceOperandWith(int64_t operand_num,
   return ReplaceOperandWithDifferentShape(operand_num, new_operand);
 }
 
-Status HloInstruction::ReplaceOperandWithDifferentShape(
+absl::Status HloInstruction::ReplaceOperandWithDifferentShape(
     int64_t operand_num, HloInstruction* new_operand) {
   TF_RET_CHECK(operand_num >= 0);
   TF_RET_CHECK(operand_num < operand_count());
@@ -3119,7 +3122,7 @@ Status HloInstruction::ReplaceOperandWithDifferentShape(
 // Copy all the instructions in the given fusion instruction into the fusion
 // instruction's parent computation and replace the use of the fusion
 // instruction with the copy of the fusion expression root.
-Status HloInstruction::Defuse() {
+absl::Status HloInstruction::Defuse() {
   if (opcode() != HloOpcode::kFusion) {
     return OkStatus();
   }
@@ -3164,15 +3167,15 @@ Status HloInstruction::Defuse() {
   return module->RemoveEmbeddedComputation(fused_computation);
 }
 
-Status HloInstruction::ReplaceUsesWith(absl::Span<HloInstruction* const> users,
-                                       HloInstruction* new_producer) {
+absl::Status HloInstruction::ReplaceUsesWith(
+    absl::Span<HloInstruction* const> users, HloInstruction* new_producer) {
   TF_RET_CHECK(
       ShapeUtil::CompatibleIgnoringFpPrecision(shape(), new_producer->shape()))
       << shape() << " is not compatible with " << new_producer->shape();
   return ReplaceAllUsesWithDifferentShape(users, new_producer);
 }
 
-Status HloInstruction::ReplaceAllUsesWithDifferentShape(
+absl::Status HloInstruction::ReplaceAllUsesWithDifferentShape(
     absl::Span<HloInstruction* const> users, HloInstruction* new_producer) {
   // Make a copy since users span might get mutated during the loop
   std::vector<HloInstruction*> users_vector(users.begin(), users.end());
@@ -3187,8 +3190,8 @@ Status HloInstruction::ReplaceAllUsesWithDifferentShape(
   return OkStatus();
 }
 
-Status HloInstruction::ReplaceAllUsesWith(HloInstruction* new_producer,
-                                          absl::string_view trigger) {
+absl::Status HloInstruction::ReplaceAllUsesWith(HloInstruction* new_producer,
+                                                absl::string_view trigger) {
   auto print_options = HloPrintOptions::ShortParsable()
                            .set_print_operand_shape(true)
                            .set_print_extra_attributes(false);
@@ -3201,7 +3204,7 @@ Status HloInstruction::ReplaceAllUsesWith(HloInstruction* new_producer,
   return ReplaceAllUsesWithDifferentShape(new_producer);
 }
 
-Status HloInstruction::ReplaceAllUsesWithDifferentShape(
+absl::Status HloInstruction::ReplaceAllUsesWithDifferentShape(
     HloInstruction* new_producer) {
   bool new_producer_is_user = false;
   // Make a copy since users span might get mutated during the loop
@@ -4058,7 +4061,8 @@ HloInstruction::HloInstruction(HloOpcode opcode, const Shape& shape)
 }
 
 template <typename HloInstructionPtr>
-Status HloInstruction::Visit(DfsHloVisitorBase<HloInstructionPtr>* visitor) {
+absl::Status HloInstruction::Visit(
+    DfsHloVisitorBase<HloInstructionPtr>* visitor) {
   switch (opcode_) {
     case HloOpcode::kAbs:
       return visitor->HandleAbs(this);
@@ -4311,8 +4315,8 @@ Status HloInstruction::Visit(DfsHloVisitorBase<HloInstructionPtr>* visitor) {
 }
 
 // Explicit instantiations.
-template Status HloInstruction::Visit(DfsHloVisitor* visitor);
-template Status HloInstruction::Visit(ConstDfsHloVisitor* visitor);
+template absl::Status HloInstruction::Visit(DfsHloVisitor* visitor);
+template absl::Status HloInstruction::Visit(ConstDfsHloVisitor* visitor);
 
 // Push "child" onto the dfs_stack if not already visited.  Returns false if a
 // cycle was detected, and true otherwise.
@@ -4340,10 +4344,10 @@ using InternalCompareFunction =
     absl::FunctionRef<bool(std::pair<int, const HloInstruction*>,
                            std::pair<int, const HloInstruction*>)>;
 template <typename Visitor>
-static Status PostOrderDFS(HloInstruction* root, Visitor* visitor,
-                           std::optional<InternalCompareFunction> operand_order,
-                           bool ignore_control_predecessors,
-                           bool cross_computation) {
+static absl::Status PostOrderDFS(
+    HloInstruction* root, Visitor* visitor,
+    std::optional<InternalCompareFunction> operand_order,
+    bool ignore_control_predecessors, bool cross_computation) {
   visitor->ReserveVisitStates(root->parent()->instruction_count());
 
   // dfs_stack holds pairs of <HloInstruction*->unique_id(), HloInstruction*>.
@@ -4438,10 +4442,9 @@ static Status PostOrderDFS(HloInstruction* root, Visitor* visitor,
 }
 
 template <typename HloInstructionPtr>
-Status HloInstruction::Accept(DfsHloVisitorBase<HloInstructionPtr>* visitor,
-                              bool call_finish_visit,
-                              bool ignore_control_predecessors,
-                              bool cross_computation) {
+absl::Status HloInstruction::Accept(
+    DfsHloVisitorBase<HloInstructionPtr>* visitor, bool call_finish_visit,
+    bool ignore_control_predecessors, bool cross_computation) {
   VLOG(3) << "HloInstruction::Accept(%" << name() << ")";
   TF_RETURN_IF_ERROR(PostOrderDFS(this, visitor, std::nullopt,
                                   ignore_control_predecessors,
@@ -4453,12 +4456,13 @@ Status HloInstruction::Accept(DfsHloVisitorBase<HloInstructionPtr>* visitor,
 }
 
 // Explicit instantiations.
-template Status HloInstruction::Accept(DfsHloVisitor*, bool, bool, bool);
-template Status HloInstruction::Accept(ConstDfsHloVisitor*, bool, bool, bool);
+template absl::Status HloInstruction::Accept(DfsHloVisitor*, bool, bool, bool);
+template absl::Status HloInstruction::Accept(ConstDfsHloVisitor*, bool, bool,
+                                             bool);
 
-Status HloInstruction::AcceptWithOperandOrder(DfsHloVisitor* visitor,
-                                              CompareFunction operand_order,
-                                              bool call_finish_visit) {
+absl::Status HloInstruction::AcceptWithOperandOrder(
+    DfsHloVisitor* visitor, CompareFunction operand_order,
+    bool call_finish_visit) {
   VLOG(2) << "HloInstruction::AcceptWithOperandOrder(%" << name() << ")";
   auto func = [operand_order](std::pair<int, const HloInstruction*> a,
                               std::pair<int, const HloInstruction*> b) {
diff --git a/third_party/xla/xla/hlo/ir/hlo_instruction.h b/third_party/xla/xla/hlo/ir/hlo_instruction.h
index c452e31ad9288d..a1f7e2e98312fc 100644
--- a/third_party/xla/xla/hlo/ir/hlo_instruction.h
+++ b/third_party/xla/xla/hlo/ir/hlo_instruction.h
@@ -1498,19 +1498,19 @@ class HloInstruction {
   // This is used to enforce an additional ordering requirement that is not
   // captured by normal data dependencies, such as ordering among Send or Recv
   // operations to avoid deadlock.
-  Status AddControlDependencyTo(HloInstruction* instruction);
+  absl::Status AddControlDependencyTo(HloInstruction* instruction);
 
   // Removes a previously added control dependency from this instruction to
   // 'instruction'.
-  Status RemoveControlDependencyTo(HloInstruction* instruction);
+  absl::Status RemoveControlDependencyTo(HloInstruction* instruction);
 
   // Drops all control predecessors and successors from this HLO instruction.
-  Status DropAllControlDeps();
+  absl::Status DropAllControlDeps();
 
   // Drops all control predecessors and successors from this HLO instruction,
   // and the maintain the transitivie control dependencies between
   // control predecessors and control successors.
-  Status SafelyDropAllControlDependencies();
+  absl::Status SafelyDropAllControlDependencies();
 
   // Returns if instruction has any control dependencies.
   bool HasControlDependencies() const;
@@ -1522,7 +1522,7 @@ class HloInstruction {
   // Depending on the use cases we see in practice, in the future we may
   // consider folding the logic here into Clone, CloneWithNewOperands and
   // ReplaceAllUsesWith by treating control dependencies like data dependencies.
-  Status CopyAllControlDepsFrom(const HloInstruction* inst) {
+  absl::Status CopyAllControlDepsFrom(const HloInstruction* inst) {
     return inst->CopyAllControlDepsTo(this, this);
   }
 
@@ -1530,7 +1530,8 @@ class HloInstruction {
   // all control predecessors of this instruction to control predecessors of
   // `start` and copies all control successors of this instruction to control
   // successors of `end`.
-  Status CopyAllControlDepsTo(HloInstruction* start, HloInstruction* end) const;
+  absl::Status CopyAllControlDepsTo(HloInstruction* start,
+                                    HloInstruction* end) const;
 
   // Returns the set of control predecessors (successors) of this
   // instruction. Control predecessors (successors) must execute before (after)
@@ -1626,32 +1627,35 @@ class HloInstruction {
   //
   // If user is a fusion instruction, this function will remove any duplicated
   // operands of it which could be created due to this replacement.
-  Status ReplaceUseWith(HloInstruction* user, HloInstruction* new_producer);
+  absl::Status ReplaceUseWith(HloInstruction* user,
+                              HloInstruction* new_producer);
 
   // Same as ReplaceUseWith(), but new_producer can have a different shape.
-  Status ReplaceUseWithDifferentShape(HloInstruction* user,
-                                      HloInstruction* new_producer);
+  absl::Status ReplaceUseWithDifferentShape(HloInstruction* user,
+                                            HloInstruction* new_producer);
 
   // Same as ReplaceUseWith but only replaces the use at the given operand
   // number.
-  Status ReplaceUseWith(HloInstruction* user, int operand_number,
-                        HloInstruction* new_producer);
-  Status ReplaceUseWithDifferentShape(HloInstruction* user, int operand_number,
-                                      HloInstruction* new_producer);
+  absl::Status ReplaceUseWith(HloInstruction* user, int operand_number,
+                              HloInstruction* new_producer);
+  absl::Status ReplaceUseWithDifferentShape(HloInstruction* user,
+                                            int operand_number,
+                                            HloInstruction* new_producer);
 
   // Replaces the specified operand with new_operand. The old and new operands
   // must have compatible shapes ignoring floating-point precision.
   //
   // This function does NOT remove duplicated operands even if this instruction
   // is a fusion, so that the existing operand numbers do not change.
-  Status ReplaceOperandWith(int64_t operand_num, HloInstruction* new_operand);
+  absl::Status ReplaceOperandWith(int64_t operand_num,
+                                  HloInstruction* new_operand);
 
   // Same as ReplaceOperandWith(), but new_operand can have a different shape.
-  Status ReplaceOperandWithDifferentShape(int64_t operand_num,
-                                          HloInstruction* new_operand);
+  absl::Status ReplaceOperandWithDifferentShape(int64_t operand_num,
+                                                HloInstruction* new_operand);
 
   // Decomposes fusion back to individual parts.
-  Status Defuse();
+  absl::Status Defuse();
 
   // Replaces all uses of this instruction with the new producer. If
   // new_producer is a user of this instruction then new_producer remains a use
@@ -1668,16 +1672,16 @@ class HloInstruction {
   //
   // trigger is a string used in the error message if the new and the
   // current instruction don't have a compatible shape.
-  Status ReplaceAllUsesWith(HloInstruction* new_producer,
-                            absl::string_view trigger = "");
+  absl::Status ReplaceAllUsesWith(HloInstruction* new_producer,
+                                  absl::string_view trigger = "");
 
   // Same as ReplaceAllUsesWith, but new_producer can have a different shape.
-  Status ReplaceAllUsesWithDifferentShape(HloInstruction* new_producer);
+  absl::Status ReplaceAllUsesWithDifferentShape(HloInstruction* new_producer);
 
   // Same as ReplaceAllUsesWith, but only replace given set of users.
-  Status ReplaceUsesWith(absl::Span<HloInstruction* const> users,
-                         HloInstruction* new_producer);
-  Status ReplaceAllUsesWithDifferentShape(
+  absl::Status ReplaceUsesWith(absl::Span<HloInstruction* const> users,
+                               HloInstruction* new_producer);
+  absl::Status ReplaceAllUsesWithDifferentShape(
       absl::Span<HloInstruction* const> users, HloInstruction* new_producer);
 
   // Performs a postorder DFS visit using this node as the root. If
@@ -1689,13 +1693,14 @@ class HloInstruction {
   // true, DFS will go across the computation boundary (i.e., from an
   // instruction to the root instruction of a computation it calls).
   template <typename HloInstructionPtr>
-  Status Accept(DfsHloVisitorBase<HloInstructionPtr>* visitor,
-                bool call_finish_visit = true,
-                bool ignore_control_predecessors = false,
-                bool cross_computation = false);
-  Status Accept(ConstDfsHloVisitor* visitor, bool call_finish_visit = true,
-                bool ignore_control_predecessors = false,
-                bool cross_computation = false) const {
+  absl::Status Accept(DfsHloVisitorBase<HloInstructionPtr>* visitor,
+                      bool call_finish_visit = true,
+                      bool ignore_control_predecessors = false,
+                      bool cross_computation = false);
+  absl::Status Accept(ConstDfsHloVisitor* visitor,
+                      bool call_finish_visit = true,
+                      bool ignore_control_predecessors = false,
+                      bool cross_computation = false) const {
     return const_cast<HloInstruction*>(this)->Accept(
         visitor, call_finish_visit, ignore_control_predecessors,
         cross_computation);
@@ -1706,14 +1711,14 @@ class HloInstruction {
   // true, A is visited before B.
   using CompareFunction =
       absl::FunctionRef<bool(const HloInstruction*, const HloInstruction*)>;
-  Status AcceptWithOperandOrder(DfsHloVisitor* visitor,
-                                CompareFunction operand_order,
-                                bool call_finish_visit = true);
+  absl::Status AcceptWithOperandOrder(DfsHloVisitor* visitor,
+                                      CompareFunction operand_order,
+                                      bool call_finish_visit = true);
 
   // Visit this instruction and only this instruction with the given visitor.
   template <typename HloInstructionPtr>
-  Status Visit(DfsHloVisitorBase<HloInstructionPtr>* visitor);
-  Status Visit(ConstDfsHloVisitor* visitor) const {
+  absl::Status Visit(DfsHloVisitorBase<HloInstructionPtr>* visitor);
+  absl::Status Visit(ConstDfsHloVisitor* visitor) const {
     return const_cast<HloInstruction*>(this)->Visit(visitor);
   }
 
@@ -2130,13 +2135,13 @@ class HloInstruction {
   //
   // ConfigProto should be a protobuf Message type.
   template <typename ConfigProto, EnableIfProto<ConfigProto>* = nullptr>
-  StatusOr<ConfigProto> backend_config() const {
+  absl::StatusOr<ConfigProto> backend_config() const {
     ConfigProto proto;
     TF_RETURN_IF_ERROR(backend_config_.GetProto(&proto));
     return std::move(proto);
   }
 
-  Status set_backend_config(const tsl::protobuf::Message& proto) {
+  absl::Status set_backend_config(const tsl::protobuf::Message& proto) {
     backend_config_ = BackendConfigWrapper(proto);
     return OkStatus();
   }
@@ -2631,7 +2636,7 @@ class HloInstruction {
 
   // Helper for implementing backend_config().  Parses backend_config_ into the
   // given proto.
-  Status GetBackendConfigInternal(tsl::protobuf::Message* proto) const;
+  absl::Status GetBackendConfigInternal(tsl::protobuf::Message* proto) const;
 
   // Mark this instruction as dead. Accessed by friend class HloInstruction.
   void MarkAsDead() { marked_as_dead_ = true; }
@@ -2789,11 +2794,12 @@ class HloInstruction {
 };
 
 // Explicit instantiations in hlo_instruction.cc.
-extern template Status HloInstruction::Accept(DfsHloVisitor*, bool, bool, bool);
-extern template Status HloInstruction::Accept(ConstDfsHloVisitor*, bool, bool,
-                                              bool);
-extern template Status HloInstruction::Visit(DfsHloVisitor* visitor);
-extern template Status HloInstruction::Visit(ConstDfsHloVisitor* visitor);
+extern template absl::Status HloInstruction::Accept(DfsHloVisitor*, bool, bool,
+                                                    bool);
+extern template absl::Status HloInstruction::Accept(ConstDfsHloVisitor*, bool,
+                                                    bool, bool);
+extern template absl::Status HloInstruction::Visit(DfsHloVisitor* visitor);
+extern template absl::Status HloInstruction::Visit(ConstDfsHloVisitor* visitor);
 
 absl::string_view ToString(HloInstruction::FusionKind kind);
 absl::StatusOr<HloInstruction::FusionKind> StringToFusionKind(
diff --git a/third_party/xla/xla/hlo/ir/hlo_instructions.cc b/third_party/xla/xla/hlo/ir/hlo_instructions.cc
index 9837719f825d14..b1a44f092a4c4c 100644
--- a/third_party/xla/xla/hlo/ir/hlo_instructions.cc
+++ b/third_party/xla/xla/hlo/ir/hlo_instructions.cc
@@ -2516,7 +2516,7 @@ std::unique_ptr<HloInstruction> HloFusionInstruction::CloneWithNewOperandsImpl(
   return new_fusion_instruction;
 }
 
-Status HloFusionInstruction::DeduplicateFusionOperands() {
+absl::Status HloFusionInstruction::DeduplicateFusionOperands() {
   if (IsCustomFusion()) {
     return OkStatus();
   }
diff --git a/third_party/xla/xla/hlo/ir/hlo_instructions.h b/third_party/xla/xla/hlo/ir/hlo_instructions.h
index f11e536f45e675..c84754fb87a114 100644
--- a/third_party/xla/xla/hlo/ir/hlo_instructions.h
+++ b/third_party/xla/xla/hlo/ir/hlo_instructions.h
@@ -1516,7 +1516,7 @@ class HloFusionInstruction : public HloCallableInstruction {
   void set_fusion_kind(FusionKind kind) { fusion_kind_ = kind; }
 
   // If multiple operands are the same instruction, keeps only one of them.
-  Status DeduplicateFusionOperands();
+  absl::Status DeduplicateFusionOperands();
 
   static bool ClassOf(const HloInstruction* hlo) {
     return hlo->opcode() == HloOpcode::kFusion;
diff --git a/third_party/xla/xla/hlo/ir/hlo_module.cc b/third_party/xla/xla/hlo/ir/hlo_module.cc
index 2deb35aa19cd48..b5fdd4d074f9f1 100644
--- a/third_party/xla/xla/hlo/ir/hlo_module.cc
+++ b/third_party/xla/xla/hlo/ir/hlo_module.cc
@@ -79,7 +79,7 @@ HloModule::HloModule(const std::string& name,
   metadata_.set_canonical_module_id(unique_id_);
 }
 
-Status HloModule::set_schedule(HloSchedule schedule) {
+absl::Status HloModule::set_schedule(HloSchedule schedule) {
   TF_RET_CHECK(schedule.module() == this);
   TF_RETURN_IF_ERROR(schedule.Verify());
   schedule_ = std::move(schedule);
@@ -187,7 +187,7 @@ HloComputation* HloModule::AddEntryComputationWithLayouts(
                                 /*preserve_entry_layouts=*/true);
 }
 
-Status HloModule::RemoveEmbeddedComputation(HloComputation* to_remove) {
+absl::Status HloModule::RemoveEmbeddedComputation(HloComputation* to_remove) {
   if (has_schedule()) {
     schedule_->remove_computation(to_remove);
   }
@@ -511,7 +511,8 @@ absl::StatusOr<HloModuleProtoWithConfig> HloModule::ToProtoWithConfig() const {
   return result;
 }
 
-Status HloModule::CheckUniqueNamesAndIdsForComputationsAndInstructions() const {
+absl::Status HloModule::CheckUniqueNamesAndIdsForComputationsAndInstructions()
+    const {
   absl::flat_hash_set<absl::string_view> computation_names;
   absl::flat_hash_set<int> computation_ids;
   absl::flat_hash_set<absl::string_view> instruction_names;
@@ -1118,7 +1119,7 @@ std::unique_ptr<HloModule> HloModule::Clone(
   return module;
 }
 
-Status HloModule::RemoveUnusedComputations() {
+absl::Status HloModule::RemoveUnusedComputations() {
   std::string suffix = "tmp";
   auto module = std::make_unique<HloModule>(
       absl::StrCat(name_, "-", suffix), config(),
diff --git a/third_party/xla/xla/hlo/ir/hlo_module.h b/third_party/xla/xla/hlo/ir/hlo_module.h
index d232fb1462ce00..cc884d25a3474a 100644
--- a/third_party/xla/xla/hlo/ir/hlo_module.h
+++ b/third_party/xla/xla/hlo/ir/hlo_module.h
@@ -163,10 +163,10 @@ class HloModule {
       std::unique_ptr<HloComputation> computation);
 
   // Removes an embedded computation.
-  Status RemoveEmbeddedComputation(HloComputation* to_remove);
+  absl::Status RemoveEmbeddedComputation(HloComputation* to_remove);
 
   // Removes unused computations.
-  Status RemoveUnusedComputations();
+  absl::Status RemoveUnusedComputations();
 
   // Marks duplicate fusions with the same name to be able to group them for
   // analysis purposes (e.g. through Xprof).
@@ -524,7 +524,7 @@ class HloModule {
   int unique_id() const { return unique_id_; }
 
   // Sets the schedule of the module to the given schedule.
-  Status set_schedule(HloSchedule schedule);
+  absl::Status set_schedule(HloSchedule schedule);
 
   // Clears the schedule of the module.
   void clear_schedule() { schedule_.reset(); }
@@ -558,7 +558,7 @@ class HloModule {
     computation->UniquifyName(&computation_name_uniquer_);
   }
 
-  Status CheckUniqueNamesAndIdsForComputationsAndInstructions() const;
+  absl::Status CheckUniqueNamesAndIdsForComputationsAndInstructions() const;
 
   // Checks if this config has a list of entry parameters' HLO shardings for
   // SPMD.
@@ -611,7 +611,8 @@ class HloModule {
         CrossProgramPrefetchInfo{parameter, index, alt_memory_offset});
   }
 
-  Status SetCrossProgramPrefetchOffset(int64_t prefetch_index, int64_t offset) {
+  absl::Status SetCrossProgramPrefetchOffset(int64_t prefetch_index,
+                                             int64_t offset) {
     TF_RET_CHECK(prefetch_index < cross_program_prefetches_.size());
     auto& [parameter, index, optional_offset] =
         cross_program_prefetches_[prefetch_index];
diff --git a/third_party/xla/xla/hlo/ir/hlo_module_metadata.cc b/third_party/xla/xla/hlo/ir/hlo_module_metadata.cc
index aa910c10919249..9765dee1fcf096 100644
--- a/third_party/xla/xla/hlo/ir/hlo_module_metadata.cc
+++ b/third_party/xla/xla/hlo/ir/hlo_module_metadata.cc
@@ -37,7 +37,7 @@ HloModuleMetadata::GetCurrentHloPassMetadata() {
   return running_passes_.back();
 }
 
-Status HloModuleMetadata::MutateCurrentHloPassMetadata(
+absl::Status HloModuleMetadata::MutateCurrentHloPassMetadata(
     absl::FunctionRef<void(HloPassMetadata*)> mutator) {
   TF_ASSIGN_OR_RETURN(HloPassMetadata * pass_metadata,
                       GetCurrentHloPassMetadata());
@@ -52,7 +52,7 @@ void HloModuleMetadata::RecordPassStart() {
   running_passes_.push_back(pass_metadata);
 }
 
-Status HloModuleMetadata::RecordPassEnd() {
+absl::Status HloModuleMetadata::RecordPassEnd() {
   TF_ASSIGN_OR_RETURN(HloPassMetadata * pass_metadata,
                       GetCurrentHloPassMetadata());
   pass_metadata->set_end_timestamp_usec(env_->NowMicros());
@@ -88,7 +88,7 @@ void HloModuleMetadata::set_prepartitioning_metadata(
   }
 }
 
-Status HloModuleMetadata::set_custom_metadata(
+absl::Status HloModuleMetadata::set_custom_metadata(
     const ::tsl::protobuf::Message& message) {
   TF_ASSIGN_OR_RETURN(HloPassMetadata * pass_metadata,
                       GetCurrentHloPassMetadata());
diff --git a/third_party/xla/xla/hlo/ir/hlo_module_metadata.h b/third_party/xla/xla/hlo/ir/hlo_module_metadata.h
index a18e9090db4d01..14231f29897e59 100644
--- a/third_party/xla/xla/hlo/ir/hlo_module_metadata.h
+++ b/third_party/xla/xla/hlo/ir/hlo_module_metadata.h
@@ -44,7 +44,7 @@ class HloModuleMetadata {
 
   // Marks the currently running pass as finished. Returns NotFound if metadata
   // for the currently running pass cannot be found.
-  Status RecordPassEnd();
+  absl::Status RecordPassEnd();
 
   const std::optional<HloModuleMetadataProto>& prepartitioning_metadata()
       const {
@@ -63,7 +63,7 @@ class HloModuleMetadata {
   void add_partitioned_module_id(int64_t id) {
     module_metadata_.add_partitioned_module_ids(id);
   }
-  Status set_custom_metadata(const ::tsl::protobuf::Message& message);
+  absl::Status set_custom_metadata(const ::tsl::protobuf::Message& message);
 
   absl::StatusOr<int64_t> current_pass_id() {
     TF_ASSIGN_OR_RETURN(HloPassMetadata * pass_metadata,
@@ -72,37 +72,39 @@ class HloModuleMetadata {
   }
 
   // Setters for the current HloPassMetadata.
-  Status set_current_pass_name(const std::string& pass_name) {
+  absl::Status set_current_pass_name(const std::string& pass_name) {
     return MutateCurrentHloPassMetadata(
         [&pass_name](HloPassMetadata* pass_metadata) {
           pass_metadata->set_pass_name(pass_name);
         });
   }
-  Status set_current_pass_pipeline_name(const std::string& pipeline_name) {
+  absl::Status set_current_pass_pipeline_name(
+      const std::string& pipeline_name) {
     return MutateCurrentHloPassMetadata(
         [&pipeline_name](HloPassMetadata* pass_metadata) {
           pass_metadata->set_pipeline_name(pipeline_name);
         });
   }
-  Status add_current_pass_dump_filename(const std::string& dump_filename) {
+  absl::Status add_current_pass_dump_filename(
+      const std::string& dump_filename) {
     return MutateCurrentHloPassMetadata(
         [&dump_filename](HloPassMetadata* pass_metadata) {
           pass_metadata->add_dump_filenames(dump_filename);
         });
   }
-  Status set_current_pass_module_changed(bool module_changed) {
+  absl::Status set_current_pass_module_changed(bool module_changed) {
     return MutateCurrentHloPassMetadata(
         [&module_changed](HloPassMetadata* pass_metadata) {
           pass_metadata->set_module_changed(module_changed);
         });
   }
-  Status set_current_pass_module_id(int64_t module_id) {
+  absl::Status set_current_pass_module_id(int64_t module_id) {
     return MutateCurrentHloPassMetadata(
         [&module_id](HloPassMetadata* pass_metadata) {
           pass_metadata->set_module_id(module_id);
         });
   }
-  Status add_current_pass_module_group_module_id(int64_t module_id) {
+  absl::Status add_current_pass_module_group_module_id(int64_t module_id) {
     return MutateCurrentHloPassMetadata(
         [&module_id](HloPassMetadata* pass_metadata) {
           pass_metadata->add_module_group_module_ids(module_id);
@@ -115,7 +117,7 @@ class HloModuleMetadata {
   // currently running pass cannot be found.
   absl::StatusOr<HloPassMetadata*> GetCurrentHloPassMetadata();
 
-  Status MutateCurrentHloPassMetadata(
+  absl::Status MutateCurrentHloPassMetadata(
       absl::FunctionRef<void(HloPassMetadata*)> mutator);
 
   HloModuleMetadataProto module_metadata_;
diff --git a/third_party/xla/xla/hlo/ir/hlo_schedule.cc b/third_party/xla/xla/hlo/ir/hlo_schedule.cc
index 2922a1b1286f37..1e0dce4f73a0b1 100644
--- a/third_party/xla/xla/hlo/ir/hlo_schedule.cc
+++ b/third_party/xla/xla/hlo/ir/hlo_schedule.cc
@@ -124,7 +124,7 @@ const HloInstructionSequence& HloSchedule::sequence(
   return sequences_.at(computation->unique_id());
 }
 
-Status HloSchedule::UpdateComputationSchedule(
+absl::Status HloSchedule::UpdateComputationSchedule(
     const HloComputation* computation) {
   // Map from unique ID to HloInstruction pointer for instructions in the
   // computation.
@@ -212,7 +212,7 @@ Status HloSchedule::UpdateComputationSchedule(
   return OkStatus();
 }
 
-Status HloSchedule::Update(
+absl::Status HloSchedule::Update(
     const absl::flat_hash_set<absl::string_view>& execution_threads) {
   // The schedule must contain a sequence for every non-fusion computation in
   // the module for the specified threads, but can have sequences for
@@ -282,7 +282,7 @@ HloSchedule::num_sequences_by_execution_thread() const {
   return sequence_num_by_execution_threads;
 }
 
-Status HloSchedule::Verify() const {
+absl::Status HloSchedule::Verify() const {
   VLOG(2) << "VerifySchedule()";
   XLA_VLOG_LINES(2, ToString());
 
diff --git a/third_party/xla/xla/hlo/ir/hlo_schedule.h b/third_party/xla/xla/hlo/ir/hlo_schedule.h
index 3e04c843d577de..bfc2ce5342c5b4 100644
--- a/third_party/xla/xla/hlo/ir/hlo_schedule.h
+++ b/third_party/xla/xla/hlo/ir/hlo_schedule.h
@@ -190,14 +190,14 @@ class HloSchedule {
   // remain in the same order in the updated schedule. Instructions which exist
   // in the module but not in the given schedule will be placed as early as
   // possible in the updated schedule.
-  Status Update(
+  absl::Status Update(
       const absl::flat_hash_set<absl::string_view>& execution_threads = {});
 
   // Verifies that the given schedule is valid for the given module.
   // Specifically, the schedule contains exactly the instructions in the
   // non-fusion computations in the module and every dependency in the module is
   // satisfied in the schedule.
-  Status Verify() const;
+  absl::Status Verify() const;
 
   std::string ToString() const;
 
@@ -207,7 +207,7 @@ class HloSchedule {
 
  private:
   // Updates the instruction sequence for the given computation.
-  Status UpdateComputationSchedule(const HloComputation* computation);
+  absl::Status UpdateComputationSchedule(const HloComputation* computation);
 
   const HloModule* module_;
 
diff --git a/third_party/xla/xla/hlo/ir/hlo_sharding.cc b/third_party/xla/xla/hlo/ir/hlo_sharding.cc
index 0577e5de93d0d1..26b2702885b042 100644
--- a/third_party/xla/xla/hlo/ir/hlo_sharding.cc
+++ b/third_party/xla/xla/hlo/ir/hlo_sharding.cc
@@ -613,7 +613,7 @@ int64_t HloSharding::RequiredLeaves(const Shape& shape) {
   return (leaf_count == 0) ? 1 : leaf_count;
 }
 
-Status HloSharding::CheckLeafCount(const Shape& shape) const {
+absl::Status HloSharding::CheckLeafCount(const Shape& shape) const {
   int64_t leaf_count = ShapeUtil::GetLeafCount(shape);
   if (leaf_count == 0 && tuple_elements_.size() == 1) {
     // Allow (but don't require) empty tuples to have a single sharding
@@ -683,8 +683,8 @@ int64_t HloSharding::GetUniqueDevice() const {
   return *device;
 }
 
-Status HloSharding::ValidateTuple(const Shape& shape,
-                                  std::optional<int64_t> num_devices) const {
+absl::Status HloSharding::ValidateTuple(
+    const Shape& shape, std::optional<int64_t> num_devices) const {
   if (!shape.IsTuple()) {
     return tsl::errors::InvalidArgument(
         "Sharding is tuple-shaped but validation shape is not.");
@@ -699,7 +699,7 @@ Status HloSharding::ValidateTuple(const Shape& shape,
   // shape tree.
   ShapeTree<HloSharding> shape_tree = GetAsShapeTree(shape);
   for (const auto& index_to_sharding : shape_tree.leaves()) {
-    Status status = index_to_sharding.second.ValidateNonTuple(
+    absl::Status status = index_to_sharding.second.ValidateNonTuple(
         ShapeUtil::GetSubshape(shape, index_to_sharding.first), num_devices);
     if (!status.ok()) {
       tsl::errors::AppendToMessage(
@@ -712,13 +712,13 @@ Status HloSharding::ValidateTuple(const Shape& shape,
   return OkStatus();
 }
 
-Status HloSharding::Validate(const Shape& shape,
-                             std::optional<int64_t> num_devices) const {
+absl::Status HloSharding::Validate(const Shape& shape,
+                                   std::optional<int64_t> num_devices) const {
   if (shape.IsToken()) {
     return OkStatus();
   }
-  Status status = IsTuple() ? ValidateTuple(shape, num_devices)
-                            : ValidateNonTuple(shape, num_devices);
+  absl::Status status = IsTuple() ? ValidateTuple(shape, num_devices)
+                                  : ValidateNonTuple(shape, num_devices);
   if (!status.ok()) {
     tsl::errors::AppendToMessage(
         &status, StrCat("Note: While validating sharding ", ToString(),
@@ -727,8 +727,8 @@ Status HloSharding::Validate(const Shape& shape,
   return status;
 }
 
-Status HloSharding::ValidateNonTuple(const Shape& shape,
-                                     std::optional<int64_t> num_devices) const {
+absl::Status HloSharding::ValidateNonTuple(
+    const Shape& shape, std::optional<int64_t> num_devices) const {
   if (shape.IsTuple()) {
     return absl::InvalidArgumentError(
         "Validation shape is a tuple but sharding is not.");
@@ -742,7 +742,7 @@ Status HloSharding::ValidateNonTuple(const Shape& shape,
   bool all_devices_seen;
   if (!tile_assignment_.iota_) {
     absl::flat_hash_set<int64_t> seen_devices;
-    Status status = tile_assignment_.array().EachStatus(
+    absl::Status status = tile_assignment_.array().EachStatus(
         [&num_devices, &seen_devices](absl::Span<const int64_t> indices,
                                       int32_t device) {
           if (num_devices.has_value() && device >= *num_devices) {
diff --git a/third_party/xla/xla/hlo/ir/hlo_sharding.h b/third_party/xla/xla/hlo/ir/hlo_sharding.h
index 4237b1e182ba9d..a15d3b33e4c44f 100644
--- a/third_party/xla/xla/hlo/ir/hlo_sharding.h
+++ b/third_party/xla/xla/hlo/ir/hlo_sharding.h
@@ -164,8 +164,8 @@ class HloSharding {
   std::string ToString(bool include_metadata = false) const;
 
   // Validate that this sharding can be applied to a tensor with shape `shape`.
-  Status Validate(const Shape& shape,
-                  std::optional<int64_t> num_devices = {}) const;
+  absl::Status Validate(const Shape& shape,
+                        std::optional<int64_t> num_devices = {}) const;
 
   // Returns true if the sharding has tuple type.
   bool IsTuple() const { return tuple_; }
@@ -629,15 +629,15 @@ class HloSharding {
 
   // Checks that the number of elements in tuple_elements_ is consistent with
   // the tuple shape passes as argument.
-  Status CheckLeafCount(const Shape& shape) const;
+  absl::Status CheckLeafCount(const Shape& shape) const;
 
   // Internal helper to validate a tuple sharding.
-  Status ValidateTuple(const Shape& shape,
-                       std::optional<int64_t> num_devices) const;
+  absl::Status ValidateTuple(const Shape& shape,
+                             std::optional<int64_t> num_devices) const;
 
   // Internal helper to validate a non-tuple (leaf) sharding.
-  Status ValidateNonTuple(const Shape& shape,
-                          std::optional<int64_t> num_devices) const;
+  absl::Status ValidateNonTuple(const Shape& shape,
+                                std::optional<int64_t> num_devices) const;
 
   // This field is only used if replicated_ is false. If maximal_ is true, then
   // the field contains a rank 1 array with a single element, which is the
diff --git a/third_party/xla/xla/hlo/ir/hlo_sharding_metadata.cc b/third_party/xla/xla/hlo/ir/hlo_sharding_metadata.cc
index 85f5fc04e579dc..3a9eebaadaac26 100644
--- a/third_party/xla/xla/hlo/ir/hlo_sharding_metadata.cc
+++ b/third_party/xla/xla/hlo/ir/hlo_sharding_metadata.cc
@@ -121,8 +121,8 @@ std::vector<PassThrough> LocatePassThroughDomainLinks(
   return pass_through;
 }
 
-Status FixupPassThroughDomainLinks(const DomainMetadata::Domain& domain,
-                                   const HloSharding& sharding) {
+absl::Status FixupPassThroughDomainLinks(const DomainMetadata::Domain& domain,
+                                         const HloSharding& sharding) {
   for (auto& pass_through : LocatePassThroughDomainLinks(domain)) {
     HloInstruction* tuple = pass_through.operand->parent()->AddInstruction(
         HloInstruction::CreateTuple({pass_through.operand}));
@@ -153,8 +153,8 @@ std::shared_ptr<const HloSharding> CloneShardingForDomain(
   return std::make_shared<const HloSharding>(*single_sharding);
 }
 
-Status ApplyDomainSingleSharding(const DomainMetadata::Domain& domain,
-                                 const HloSharding& sharding) {
+absl::Status ApplyDomainSingleSharding(const DomainMetadata::Domain& domain,
+                                       const HloSharding& sharding) {
   VLOG(4) << "Applying " << sharding << " sharding";
   for (HloInstruction* instruction : domain.instructions) {
     // We only change instructions without sharding, since otherwise we might
@@ -342,8 +342,8 @@ absl::StatusOr<int64_t> ApplyDomainShardingPass(
   return assigned;
 }
 
-Status ApplyDomainSharding(const DomainMetadata::Domain& domain,
-                           const HloSharding& sharding) {
+absl::Status ApplyDomainSharding(const DomainMetadata::Domain& domain,
+                                 const HloSharding& sharding) {
   // None of the external normalizers handled the domain sharding, try to see
   // whether this is a single sharding first.
   auto single_sharding = sharding.ExtractSingleSharding();
@@ -438,14 +438,14 @@ std::string ShardingMetadata::ToString() const {
 /*static*/ absl::StatusOr<const ShardingMetadata*>
 ShardingMetadata::ToShardingMetadata(const DomainMetadata* metadata) {
   if (metadata->Kind() != ShardingMetadata::KindName()) {
-    return Status(
+    return absl::Status(
         absl::StatusCode::kInvalidArgument,
         "ShardingMetadata normalizer called with incorrect domain metadata");
   }
   return static_cast<const ShardingMetadata*>(metadata);
 }
 
-Status ShardingMetadata::NormalizeShardingDomain(
+absl::Status ShardingMetadata::NormalizeShardingDomain(
     const DomainMetadata::Domain& domain, const DomainMetadata* metadata) {
   if (metadata != nullptr) {
     TF_ASSIGN_OR_RETURN(const auto& sharding_metadata,
diff --git a/third_party/xla/xla/hlo/ir/hlo_sharding_metadata.h b/third_party/xla/xla/hlo/ir/hlo_sharding_metadata.h
index f8d1b9368eb5d1..689a1d28694479 100644
--- a/third_party/xla/xla/hlo/ir/hlo_sharding_metadata.h
+++ b/third_party/xla/xla/hlo/ir/hlo_sharding_metadata.h
@@ -67,8 +67,8 @@ class ShardingMetadata : public DomainMetadata {
   // be the original one. Policy wise, HLO passes are allowed to create new
   // unassigned instructions, but if they do create assigned ones, they have to
   // conform to the ones around.
-  static Status NormalizeShardingDomain(const DomainMetadata::Domain& domain,
-                                        const DomainMetadata* metadata);
+  static absl::Status NormalizeShardingDomain(
+      const DomainMetadata::Domain& domain, const DomainMetadata* metadata);
 
  private:
   std::shared_ptr<const HloSharding> sharding_;
diff --git a/third_party/xla/xla/hlo/ir/tile_assignment.cc b/third_party/xla/xla/hlo/ir/tile_assignment.cc
index 0647736abeb9f7..2db424833bc85f 100644
--- a/third_party/xla/xla/hlo/ir/tile_assignment.cc
+++ b/third_party/xla/xla/hlo/ir/tile_assignment.cc
@@ -327,8 +327,9 @@ void TileAssignment::Each(
   array_->Each(f);
 }
 
-Status TileAssignment::EachStatus(
-    absl::FunctionRef<Status(absl::Span<const int64_t>, int64_t)> f) const {
+absl::Status TileAssignment::EachStatus(
+    absl::FunctionRef<absl::Status(absl::Span<const int64_t>, int64_t)> f)
+    const {
   MaybeMaterializeFullArray();
   return array_->EachStatus(f);
 }
diff --git a/third_party/xla/xla/hlo/ir/tile_assignment.h b/third_party/xla/xla/hlo/ir/tile_assignment.h
index 339f408f248e95..ca7f23bda2dc14 100644
--- a/third_party/xla/xla/hlo/ir/tile_assignment.h
+++ b/third_party/xla/xla/hlo/ir/tile_assignment.h
@@ -205,8 +205,9 @@ class TileAssignment {
   void Each(
       absl::FunctionRef<void(absl::Span<const int64_t>, int64_t)> f) const;
 
-  Status EachStatus(
-      absl::FunctionRef<Status(absl::Span<const int64_t>, int64_t)> f) const;
+  absl::Status EachStatus(
+      absl::FunctionRef<absl::Status(absl::Span<const int64_t>, int64_t)> f)
+      const;
 
   // Returns a tile assignment reshaped to the given dimensions.
   // REQUIRES: new shape has the same number of elements.

From 8d97018e3fca33a15e105b50c2040f5cb1d1a3dd Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Thu, 16 May 2024 11:16:07 -0700
Subject: [PATCH 288/478] Enables additional solver output if we're removing
 user shardings, which usually implies internal testing / debugging.

PiperOrigin-RevId: 634469824
---
 .../xla/xla/hlo/experimental/auto_sharding/auto_sharding.cc  | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/third_party/xla/xla/hlo/experimental/auto_sharding/auto_sharding.cc b/third_party/xla/xla/hlo/experimental/auto_sharding/auto_sharding.cc
index 5b1d5f0229ae6e..1564d1834f3cc1 100644
--- a/third_party/xla/xla/hlo/experimental/auto_sharding/auto_sharding.cc
+++ b/third_party/xla/xla/hlo/experimental/auto_sharding/auto_sharding.cc
@@ -1979,6 +1979,11 @@ AutoShardingSolverResult CallSolver(
   request.set_deterministic_mode(deterministic_mode);
   request.set_request_name(std::string(request_name));
   request.set_enable_memory_edge_costs(option.model_resharding_memory_costs);
+  // If we're removing user shardings, we are probably doing internal testing /
+  // debugging where additional output from the solver might be helpful.
+  request.set_enable_output(
+      option.preserve_shardings ==
+      AutoShardingOption::PreserveShardingsType::kRemoveAllShardings);
   if (max_cost) {
     request.mutable_max_cost()->set_coeff(*max_cost);
   }

From 62abb367e83c3a73bfe6704af65e87eae048232f Mon Sep 17 00:00:00 2001
From: Quentin Khan <qkhan@google.com>
Date: Thu, 16 May 2024 11:17:07 -0700
Subject: [PATCH 289/478] Add a command line parameter in benchmark_tool to set
 the XNNPack cache file.

PiperOrigin-RevId: 634470123
---
 .../configuration/configuration.proto         |   4 +
 .../configuration/configuration_generated.h   |  37 +++++-
 .../testdata/configuration.proto_prev         | 110 ++++++++++++++++++
 .../configuration/c/xnnpack_plugin.cc         |   4 +
 tensorflow/lite/tools/delegates/BUILD         |  14 +++
 .../delegates/xnnpack_delegate_provider.cc    |  11 +-
 .../xnnpack_delegate_provider_test.cc         |  82 +++++++++++++
 tensorflow/lite/tools/evaluation/BUILD        |   4 +-
 tensorflow/lite/tools/evaluation/utils.cc     |  36 +++++-
 tensorflow/lite/tools/evaluation/utils.h      |   4 +-
 tensorflow/lite/tools/tool_params.h           |   2 +-
 11 files changed, 294 insertions(+), 14 deletions(-)
 create mode 100644 tensorflow/lite/tools/delegates/xnnpack_delegate_provider_test.cc

diff --git a/tensorflow/lite/acceleration/configuration/configuration.proto b/tensorflow/lite/acceleration/configuration/configuration.proto
index a2bbd88004f8f6..cf82771c225328 100644
--- a/tensorflow/lite/acceleration/configuration/configuration.proto
+++ b/tensorflow/lite/acceleration/configuration/configuration.proto
@@ -317,6 +317,10 @@ enum XNNPackFlags {
 message XNNPackSettings {
   optional int32 num_threads = 1;
   optional XNNPackFlags flags = 2 [default = TFLITE_XNNPACK_DELEGATE_NO_FLAGS];
+  // Path to the experimental XNNPack cache file. XNNPack packed buffers are
+  // saved to and reloaded from this cache which can reduce initialization time
+  // and the packing memory footprint.
+  optional string experimental_weight_cache_file_path = 3;
 }
 
 // CoreML Delegate settings.
diff --git a/tensorflow/lite/acceleration/configuration/configuration_generated.h b/tensorflow/lite/acceleration/configuration/configuration_generated.h
index 6d65817edf2f99..5fd634acad5837 100644
--- a/tensorflow/lite/acceleration/configuration/configuration_generated.h
+++ b/tensorflow/lite/acceleration/configuration/configuration_generated.h
@@ -1692,6 +1692,7 @@ struct XNNPackSettingsT : public ::flatbuffers::NativeTable {
   typedef XNNPackSettings TableType;
   int32_t num_threads = 0;
   tflite::XNNPackFlags flags = tflite::XNNPackFlags_TFLITE_XNNPACK_DELEGATE_NO_FLAGS;
+  std::string experimental_weight_cache_file_path{};
 };
 
 struct XNNPackSettings FLATBUFFERS_FINAL_CLASS : private ::flatbuffers::Table {
@@ -1699,7 +1700,8 @@ struct XNNPackSettings FLATBUFFERS_FINAL_CLASS : private ::flatbuffers::Table {
   typedef XNNPackSettingsBuilder Builder;
   enum FlatBuffersVTableOffset FLATBUFFERS_VTABLE_UNDERLYING_TYPE {
     VT_NUM_THREADS = 4,
-    VT_FLAGS = 6
+    VT_FLAGS = 6,
+    VT_EXPERIMENTAL_WEIGHT_CACHE_FILE_PATH = 8
   };
   int32_t num_threads() const {
     return GetField<int32_t>(VT_NUM_THREADS, 0);
@@ -1707,10 +1709,15 @@ struct XNNPackSettings FLATBUFFERS_FINAL_CLASS : private ::flatbuffers::Table {
   tflite::XNNPackFlags flags() const {
     return static_cast<tflite::XNNPackFlags>(GetField<int32_t>(VT_FLAGS, 0));
   }
+  const ::flatbuffers::String *experimental_weight_cache_file_path() const {
+    return GetPointer<const ::flatbuffers::String *>(VT_EXPERIMENTAL_WEIGHT_CACHE_FILE_PATH);
+  }
   bool Verify(::flatbuffers::Verifier &verifier) const {
     return VerifyTableStart(verifier) &&
            VerifyField<int32_t>(verifier, VT_NUM_THREADS, 4) &&
            VerifyField<int32_t>(verifier, VT_FLAGS, 4) &&
+           VerifyOffset(verifier, VT_EXPERIMENTAL_WEIGHT_CACHE_FILE_PATH) &&
+           verifier.VerifyString(experimental_weight_cache_file_path()) &&
            verifier.EndTable();
   }
   XNNPackSettingsT *UnPack(const ::flatbuffers::resolver_function_t *_resolver = nullptr) const;
@@ -1728,6 +1735,9 @@ struct XNNPackSettingsBuilder {
   void add_flags(tflite::XNNPackFlags flags) {
     fbb_.AddElement<int32_t>(XNNPackSettings::VT_FLAGS, static_cast<int32_t>(flags), 0);
   }
+  void add_experimental_weight_cache_file_path(::flatbuffers::Offset<::flatbuffers::String> experimental_weight_cache_file_path) {
+    fbb_.AddOffset(XNNPackSettings::VT_EXPERIMENTAL_WEIGHT_CACHE_FILE_PATH, experimental_weight_cache_file_path);
+  }
   explicit XNNPackSettingsBuilder(::flatbuffers::FlatBufferBuilder &_fbb)
         : fbb_(_fbb) {
     start_ = fbb_.StartTable();
@@ -1742,13 +1752,28 @@ struct XNNPackSettingsBuilder {
 inline ::flatbuffers::Offset<XNNPackSettings> CreateXNNPackSettings(
     ::flatbuffers::FlatBufferBuilder &_fbb,
     int32_t num_threads = 0,
-    tflite::XNNPackFlags flags = tflite::XNNPackFlags_TFLITE_XNNPACK_DELEGATE_NO_FLAGS) {
+    tflite::XNNPackFlags flags = tflite::XNNPackFlags_TFLITE_XNNPACK_DELEGATE_NO_FLAGS,
+    ::flatbuffers::Offset<::flatbuffers::String> experimental_weight_cache_file_path = 0) {
   XNNPackSettingsBuilder builder_(_fbb);
+  builder_.add_experimental_weight_cache_file_path(experimental_weight_cache_file_path);
   builder_.add_flags(flags);
   builder_.add_num_threads(num_threads);
   return builder_.Finish();
 }
 
+inline ::flatbuffers::Offset<XNNPackSettings> CreateXNNPackSettingsDirect(
+    ::flatbuffers::FlatBufferBuilder &_fbb,
+    int32_t num_threads = 0,
+    tflite::XNNPackFlags flags = tflite::XNNPackFlags_TFLITE_XNNPACK_DELEGATE_NO_FLAGS,
+    const char *experimental_weight_cache_file_path = nullptr) {
+  auto experimental_weight_cache_file_path__ = experimental_weight_cache_file_path ? _fbb.CreateString(experimental_weight_cache_file_path) : 0;
+  return tflite::CreateXNNPackSettings(
+      _fbb,
+      num_threads,
+      flags,
+      experimental_weight_cache_file_path__);
+}
+
 ::flatbuffers::Offset<XNNPackSettings> CreateXNNPackSettings(::flatbuffers::FlatBufferBuilder &_fbb, const XNNPackSettingsT *_o, const ::flatbuffers::rehasher_function_t *_rehasher = nullptr);
 
 struct CoreMLSettingsT : public ::flatbuffers::NativeTable {
@@ -4911,7 +4936,8 @@ inline ::flatbuffers::Offset<HexagonSettings> CreateHexagonSettings(::flatbuffer
 inline bool operator==(const XNNPackSettingsT &lhs, const XNNPackSettingsT &rhs) {
   return
       (lhs.num_threads == rhs.num_threads) &&
-      (lhs.flags == rhs.flags);
+      (lhs.flags == rhs.flags) &&
+      (lhs.experimental_weight_cache_file_path == rhs.experimental_weight_cache_file_path);
 }
 
 inline bool operator!=(const XNNPackSettingsT &lhs, const XNNPackSettingsT &rhs) {
@@ -4930,6 +4956,7 @@ inline void XNNPackSettings::UnPackTo(XNNPackSettingsT *_o, const ::flatbuffers:
   (void)_resolver;
   { auto _e = num_threads(); _o->num_threads = _e; }
   { auto _e = flags(); _o->flags = _e; }
+  { auto _e = experimental_weight_cache_file_path(); if (_e) _o->experimental_weight_cache_file_path = _e->str(); }
 }
 
 inline ::flatbuffers::Offset<XNNPackSettings> XNNPackSettings::Pack(::flatbuffers::FlatBufferBuilder &_fbb, const XNNPackSettingsT* _o, const ::flatbuffers::rehasher_function_t *_rehasher) {
@@ -4942,10 +4969,12 @@ inline ::flatbuffers::Offset<XNNPackSettings> CreateXNNPackSettings(::flatbuffer
   struct _VectorArgs { ::flatbuffers::FlatBufferBuilder *__fbb; const XNNPackSettingsT* __o; const ::flatbuffers::rehasher_function_t *__rehasher; } _va = { &_fbb, _o, _rehasher}; (void)_va;
   auto _num_threads = _o->num_threads;
   auto _flags = _o->flags;
+  auto _experimental_weight_cache_file_path = _o->experimental_weight_cache_file_path.empty() ? 0 : _fbb.CreateString(_o->experimental_weight_cache_file_path);
   return tflite::CreateXNNPackSettings(
       _fbb,
       _num_threads,
-      _flags);
+      _flags,
+      _experimental_weight_cache_file_path);
 }
 
 
diff --git a/tensorflow/lite/acceleration/configuration/testdata/configuration.proto_prev b/tensorflow/lite/acceleration/configuration/testdata/configuration.proto_prev
index 655d95f5990120..e9616ecb4afb65 100644
--- a/tensorflow/lite/acceleration/configuration/testdata/configuration.proto_prev
+++ b/tensorflow/lite/acceleration/configuration/testdata/configuration.proto_prev
@@ -97,6 +97,8 @@ enum Delegate {
   CORE_ML = 7;
   // Arm NN Delegate.
   ARMNN = 8;
+  // MediaTek Neuron Delegate.
+  MTK_NEURON = 9;
 }
 
 enum NNAPIExecutionPreference {
@@ -662,6 +664,111 @@ message ArmNNSettings {
   optional string additional_parameters = 3;
 }
 
+// MediaTek Neuron Delegate Settings.
+// See https://neuropilot.mediatek.com/ for more information.
+message MtkNeuronSettings {
+  enum ExecutionPreference {
+    PREFERENCE_UNDEFINED = 0;
+
+    // Prefer execution in a power-efficient mode, optimizing for low power
+    // consumption.
+    PREFERENCE_LOW_POWER = 1;
+
+    // Prefer execution that provides shorter single-shot latency, optimizing
+    // for fast response times.
+    PREFERENCE_FAST_SINGLE_ANSWER = 2;
+
+    // Prefer execution that provides sustained speed for continuous operation
+    // and higher throughput, optimizing for overall performance in ongoing or
+    // repetitive tasks.
+    PREFERENCE_SUSTAINED_SPEED = 3;
+
+    // Prefer execution in the turbo boost mode, which may boost the frequencies
+    // of APU and other system components such as CPU and DRAM, to achieve
+    // maximum performance. If boosting is not supported in the underlying
+    // system, it falls back to the behavior of PREFERENCE_FAST_SINGLE_ANSWER.
+    PREFERENCE_TURBO_BOOST = 4;
+  }
+
+  enum ExecutionPriority {
+    PRIORITY_UNDEFINED = 0;
+    PRIORITY_LOW = 90;
+    PRIORITY_MEDIUM = 100;
+    PRIORITY_HIGH = 110;
+  }
+
+  enum OptimizationHint {
+    OPTIMIZATION_NONE = 0;
+
+    // Optimization hint for reducing latency. This hint may distribute the
+    // workload across multiple APU cores in the compiled model to achieve
+    // faster execution.
+    OPTIMIZATION_LOW_LATENCY = 1;
+
+    // Optimization hint for reducing DRAM access and minimizing memory
+    // bandwidth usage through kernel fusion and data fusion techniques.
+    OPTIMIZATION_DEEP_FUSION = 2;
+
+    // Optimization hint for processing multiple input samples in parallel
+    // across available APU cores in the batch dimension. This optimization is
+    // effective for models with a batch size greater than 1.
+    OPTIMIZATION_BATCH_PROCESSING = 3;
+  }
+
+  // How to check the operator compatibility with the underlying accelerator.
+  enum OperationCheckMode {
+    NO_OPERATION_CHECK = 0;
+
+    // Checks each node separately with multiple queries to the backend.
+    PER_NODE_OPERATION_CHECK = 1;
+
+    // Checks all nodes in the graph at once with a batched query to the
+    // backend.
+    PRE_OPERATION_CHECK = 2;
+  }
+
+  // The preferred execution mode. The system-wide default will be used when
+  // PREFERENCE_UNDEFINED is passed to the delegate.
+  optional ExecutionPreference execution_preference = 1;
+
+  // The execution priority of the inference request. The system-wide default
+  // will be used when PRIORITY_UNDEFINED is passed to the delegate.
+  optional ExecutionPriority execution_priority = 2;
+
+  // The optimization hints that will instruct the model compiler.
+  repeated OptimizationHint optimization_hints = 3 [packed = true];
+
+  // Whether and how to check the operator compatibility with the underlying
+  // accelerator.
+  optional OperationCheckMode operation_check_mode = 4;
+
+  // Whether to allow the accelerator to optionally use lower-precision FP16
+  // arithmetic when performing calculations on FP32 data.
+  optional bool allow_fp16_precision_for_fp32 = 5;
+
+  // Whether to use AHardwareBuffer_* API to manage buffers. Requires Android
+  // API level >= 26, or a dedicated AHardwareBuffer API shim on non-Android
+  // platforms.
+  optional bool use_ahwb = 6;
+
+  // Whether to use cachable (consistent / coherent) memory. This will affect
+  // both buffer allocation and buffer importing behaviors.
+  optional bool use_cacheable_buffer = 7 [default = true];
+
+  // Extra options for the Neuron compiler, such as "--opt-bw".
+  // See docs at https://neuropilot.mediatek.com/ for available options.
+  repeated string compile_options = 8;
+
+  // Optional list of target accelerator device names.
+  // If empty, the delegate will automatically select the accelerator.
+  // See docs at https://neuropilot.mediatek.com/ for available accelerators.
+  repeated string accelerator_names = 9;
+
+  // Optional path to the platform-dependent Neuron configuration file.
+  // See docs at https://neuropilot.mediatek.com/ for more details.
+  optional string neuron_config_path = 10;
+}
+
 // How to configure TFLite.
 message TFLiteSettings {
   // Which delegate to use.
@@ -719,6 +826,9 @@ message TFLiteSettings {
 
   // For configuring the Arm NN delegate.
   optional ArmNNSettings armnn_settings = 16;
+
+  // For configuring MediaTek Neuron delegate.
+  optional MtkNeuronSettings mtk_neuron_settings = 17;
 }
 
 // Whether to automatically fallback to TFLite CPU path on delegation errors.
diff --git a/tensorflow/lite/core/acceleration/configuration/c/xnnpack_plugin.cc b/tensorflow/lite/core/acceleration/configuration/c/xnnpack_plugin.cc
index c82b0a7964af04..c2e9ec540fcbb0 100644
--- a/tensorflow/lite/core/acceleration/configuration/c/xnnpack_plugin.cc
+++ b/tensorflow/lite/core/acceleration/configuration/c/xnnpack_plugin.cc
@@ -39,6 +39,10 @@ static TfLiteDelegate* CreateDelegate(const void* settings) {
     if (xnnpack_settings->flags()) {
       options.flags = xnnpack_settings->flags();
     }
+    if (xnnpack_settings->experimental_weight_cache_file_path()) {
+      options.experimental_weight_cache_file_path =
+          xnnpack_settings->experimental_weight_cache_file_path()->c_str();
+    }
   }
   return TfLiteXNNPackDelegateCreate(&options);
 }
diff --git a/tensorflow/lite/tools/delegates/BUILD b/tensorflow/lite/tools/delegates/BUILD
index a6ba60436ebcee..97ab1d7b0787b1 100644
--- a/tensorflow/lite/tools/delegates/BUILD
+++ b/tensorflow/lite/tools/delegates/BUILD
@@ -173,6 +173,20 @@ cc_library_with_tflite(
     alwayslink = 1,
 )
 
+cc_test(
+    name = "xnnpack_delegate_provider_test",
+    srcs = ["xnnpack_delegate_provider_test.cc"],
+    copts = tflite_copts(),
+    visibility = ["//visibility:public"],
+    deps = [
+        ":delegate_provider_hdr",
+        ":xnnpack_delegate_provider",
+        "//tensorflow/lite/delegates/xnnpack:xnnpack_delegate",
+        "//tensorflow/lite/tools:tool_params",
+        "@com_google_googletest//:gtest_main",
+    ],
+)
+
 cc_library(
     name = "external_delegate_provider",
     srcs = ["external_delegate_provider.cc"],
diff --git a/tensorflow/lite/tools/delegates/xnnpack_delegate_provider.cc b/tensorflow/lite/tools/delegates/xnnpack_delegate_provider.cc
index 7d0ef212966397..9a9f7236991438 100644
--- a/tensorflow/lite/tools/delegates/xnnpack_delegate_provider.cc
+++ b/tensorflow/lite/tools/delegates/xnnpack_delegate_provider.cc
@@ -27,6 +27,8 @@ class XnnpackDelegateProvider : public DelegateProvider {
     default_params_.AddParam("use_xnnpack", ToolParam::Create<bool>(false));
     default_params_.AddParam("xnnpack_force_fp16",
                              ToolParam::Create<bool>(false));
+    default_params_.AddParam("xnnpack_experimental_weight_cache_file_path",
+                             ToolParam::Create<std::string>(""));
   }
 
   std::vector<Flag> CreateFlags(ToolParams* params) const final;
@@ -54,6 +56,8 @@ std::vector<Flag> XnnpackDelegateProvider::CreateFlags(
                        "false explicitly."),
       CreateFlag<bool>("xnnpack_force_fp16", params,
                        "enforce float16 inference."),
+      CreateFlag<std::string>("xnnpack_experimental_weight_cache_file_path",
+                              params, "enable file-backed weight caching."),
   };
   return flags;
 }
@@ -63,6 +67,9 @@ void XnnpackDelegateProvider::LogParams(const ToolParams& params,
   LOG_TOOL_PARAM(params, bool, "use_xnnpack", "Use xnnpack", verbose);
   LOG_TOOL_PARAM(params, bool, "xnnpack_force_fp16", "xnnpack_force_fp16",
                  verbose);
+  LOG_TOOL_PARAM(params, std::string,
+                 "xnnpack_experimental_weight_cache_file_path",
+                 "xnnpack_experimental_weight_cache_file_path", verbose);
 }
 
 TfLiteDelegatePtr XnnpackDelegateProvider::CreateTfLiteDelegate(
@@ -70,7 +77,9 @@ TfLiteDelegatePtr XnnpackDelegateProvider::CreateTfLiteDelegate(
   if (params.Get<bool>("use_xnnpack")) {
     return evaluation::CreateXNNPACKDelegate(
         params.Get<int32_t>("num_threads"),
-        params.Get<bool>("xnnpack_force_fp16"));
+        params.Get<bool>("xnnpack_force_fp16"),
+        params.Get<std::string>("xnnpack_experimental_weight_cache_file_path")
+            .c_str());
   }
   return CreateNullDelegate();
 }
diff --git a/tensorflow/lite/tools/delegates/xnnpack_delegate_provider_test.cc b/tensorflow/lite/tools/delegates/xnnpack_delegate_provider_test.cc
new file mode 100644
index 00000000000000..19fd593fa51136
--- /dev/null
+++ b/tensorflow/lite/tools/delegates/xnnpack_delegate_provider_test.cc
@@ -0,0 +1,82 @@
+/* Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+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.
+==============================================================================*/
+
+#include <cstdint>
+#include <string>
+
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+#include "tensorflow/lite/delegates/xnnpack/xnnpack_delegate.h"
+#include "tensorflow/lite/tools/delegates/delegate_provider.h"
+#include "tensorflow/lite/tools/tool_params.h"
+
+namespace tflite {
+namespace tools {
+namespace {
+
+TEST(XNNPackDelegateProviderTest, Test) {
+  const char kFakeCacheParam[] = "/tmp/random/path";
+
+  const auto& providers = GetRegisteredDelegateProviders();
+  ASSERT_EQ(providers.size(), 1);
+  ToolParams params;
+
+  const auto& xnnpack_provider = providers[0];
+  ASSERT_NE(xnnpack_provider, nullptr);
+
+  params.Merge(xnnpack_provider->DefaultParams());
+  params.AddParam("num_threads", ToolParam::Create<int32_t>(-1));
+
+  EXPECT_TRUE(params.HasParam("use_xnnpack"));
+  EXPECT_FALSE(params.HasValueSet<bool>("use_xnnpack"));
+  ASSERT_NE(params.GetParam("use_xnnpack"), nullptr);
+
+  EXPECT_TRUE(params.HasParam("xnnpack_force_fp16"));
+  EXPECT_FALSE(params.HasValueSet<bool>("xnnpack_force_fp16"));
+  ASSERT_NE(params.GetParam("xnnpack_force_fp16"), nullptr);
+
+  EXPECT_TRUE(params.HasParam("xnnpack_experimental_weight_cache_file_path"));
+  EXPECT_FALSE(params.HasValueSet<std::string>(
+      "xnnpack_experimental_weight_cache_file_path"));
+  ASSERT_NE(params.GetParam("xnnpack_experimental_weight_cache_file_path"),
+            nullptr);
+
+  params.Set<bool>("use_xnnpack", true, /*position=*/0);
+
+  {
+    TfLiteDelegatePtr delegate = xnnpack_provider->CreateTfLiteDelegate(params);
+    const TfLiteXNNPackDelegateOptions* options =
+        TfLiteXNNPackDelegateGetOptions(delegate.get());
+    ASSERT_NE(options, nullptr);
+    EXPECT_EQ(options->experimental_weight_cache_file_path, nullptr);
+  }
+
+  params.Set<bool>("xnnpack_force_fp16", true, /*position=*/1);
+  params.Set<std::string>("xnnpack_experimental_weight_cache_file_path",
+                          kFakeCacheParam, /*position=*/2);
+  {
+    TfLiteDelegatePtr delegate = xnnpack_provider->CreateTfLiteDelegate(params);
+    const TfLiteXNNPackDelegateOptions* options =
+        TfLiteXNNPackDelegateGetOptions(delegate.get());
+    ASSERT_NE(options, nullptr);
+    EXPECT_THAT(options->experimental_weight_cache_file_path,
+                testing::StrEq(kFakeCacheParam));
+    EXPECT_TRUE(options->flags & TFLITE_XNNPACK_DELEGATE_FLAG_FORCE_FP16);
+  }
+}
+
+}  // namespace
+}  // namespace tools
+}  // namespace tflite
diff --git a/tensorflow/lite/tools/evaluation/BUILD b/tensorflow/lite/tools/evaluation/BUILD
index 914c33471616e8..a095d0934fa4bf 100644
--- a/tensorflow/lite/tools/evaluation/BUILD
+++ b/tensorflow/lite/tools/evaluation/BUILD
@@ -75,7 +75,9 @@ cc_library_with_stable_tflite_abi(
             "//tensorflow/lite/acceleration/configuration/c:xnnpack_plugin",
         ],
     }],
-    deps = select({
+    deps = [
+        "@flatbuffers",
+    ] + select({
         "//tensorflow/lite:tflite_with_xnnpack_explicit_false": [],
         "//conditions:default": [
             "//tensorflow/lite/acceleration/configuration:configuration_fbs",
diff --git a/tensorflow/lite/tools/evaluation/utils.cc b/tensorflow/lite/tools/evaluation/utils.cc
index 20221df8d37c91..98f46247ab1c5b 100644
--- a/tensorflow/lite/tools/evaluation/utils.cc
+++ b/tensorflow/lite/tools/evaluation/utils.cc
@@ -15,7 +15,19 @@ limitations under the License.
 
 #include "tensorflow/lite/tools/evaluation/utils.h"
 
+#include <algorithm>
+#include <cctype>
+#include <cstddef>
+#include <cstdint>
+#include <fstream>
+#include <string>
+#include <unordered_set>
+#include <vector>
+
+#include "flatbuffers/buffer.h"  // from @flatbuffers
+#include "flatbuffers/string.h"  // from @flatbuffers
 #include "tensorflow/lite/tools/delegates/delegate_provider.h"
+
 #if defined(__APPLE__)
 #include "TargetConditionals.h"
 #if (TARGET_OS_IPHONE && !TARGET_IPHONE_SIMULATOR) || \
@@ -31,6 +43,7 @@ limitations under the License.
 #include "tensorflow/lite/acceleration/configuration/c/xnnpack_plugin.h"
 #include "tensorflow/lite/acceleration/configuration/configuration_generated.h"
 #include "tensorflow/lite/c/common.h"
+#include "tensorflow/lite/delegates/xnnpack/xnnpack_delegate.h"
 #endif  // !defined(TFLITE_WITHOUT_XNNPACK)
 
 #if !defined(_WIN32)
@@ -38,11 +51,6 @@ limitations under the License.
 #endif
 #include <sys/stat.h>
 
-#include <algorithm>
-#include <fstream>
-#include <memory>
-#include <string>
-
 namespace tflite {
 namespace evaluation {
 
@@ -212,6 +220,12 @@ TfLiteDelegatePtr CreateXNNPACKDelegate() {
 TfLiteDelegatePtr CreateXNNPACKDelegate(
     const TfLiteXNNPackDelegateOptions* xnnpack_options) {
   flatbuffers::FlatBufferBuilder flatbuffer_builder;
+  flatbuffers::Offset<flatbuffers::String> experimental_weight_cache_file_path;
+  if (xnnpack_options->experimental_weight_cache_file_path) {
+    experimental_weight_cache_file_path = flatbuffer_builder.CreateString(
+        xnnpack_options->experimental_weight_cache_file_path);
+  }
+
   tflite::XNNPackSettingsBuilder xnnpack_settings_builder(flatbuffer_builder);
   int num_threads = xnnpack_options->num_threads;
   if (num_threads >= 0) {
@@ -220,6 +234,8 @@ TfLiteDelegatePtr CreateXNNPACKDelegate(
   xnnpack_settings_builder.fbb_.AddElement<int32_t>(
       XNNPackSettings::VT_FLAGS, static_cast<int32_t>(xnnpack_options->flags),
       0);
+  xnnpack_settings_builder.add_experimental_weight_cache_file_path(
+      experimental_weight_cache_file_path);
   flatbuffers::Offset<tflite::XNNPackSettings> xnnpack_settings =
       xnnpack_settings_builder.Finish();
   tflite::TFLiteSettingsBuilder tflite_settings_builder(flatbuffer_builder);
@@ -240,7 +256,9 @@ TfLiteDelegatePtr CreateXNNPACKDelegate(
   return TfLiteDelegatePtr(delegate, delegate_deleter);
 }
 
-TfLiteDelegatePtr CreateXNNPACKDelegate(int num_threads, bool force_fp16) {
+TfLiteDelegatePtr CreateXNNPACKDelegate(
+    int num_threads, bool force_fp16,
+    const char* experimental_weight_cache_file_path) {
   auto opts = XNNPackDelegateOptionsDefault();
   // Note that we don't want to use the thread pool for num_threads == 1.
   opts.num_threads = num_threads > 1 ? num_threads : 0;
@@ -248,6 +266,12 @@ TfLiteDelegatePtr CreateXNNPACKDelegate(int num_threads, bool force_fp16) {
     TFLITE_LOG(INFO) << "XNNPack FP16 inference enabled.";
     opts.flags |= TFLITE_XNNPACK_DELEGATE_FLAG_FORCE_FP16;
   }
+  if (experimental_weight_cache_file_path &&
+      experimental_weight_cache_file_path[0] != '\0') {
+    TFLITE_LOG(INFO) << "XNNPack file-backed weight cache enabled.";
+    opts.experimental_weight_cache_file_path =
+        experimental_weight_cache_file_path;
+  }
   return CreateXNNPACKDelegate(&opts);
 }
 #endif
diff --git a/tensorflow/lite/tools/evaluation/utils.h b/tensorflow/lite/tools/evaluation/utils.h
index 1735b9ee08a0fc..ed89f609770e4e 100644
--- a/tensorflow/lite/tools/evaluation/utils.h
+++ b/tensorflow/lite/tools/evaluation/utils.h
@@ -105,7 +105,9 @@ TfLiteDelegatePtr CreateXNNPACKDelegate();
 TfLiteDelegatePtr CreateXNNPACKDelegate(
     const TfLiteXNNPackDelegateOptions* options);
 #endif  // !defined(TFLITE_WITHOUT_XNNPACK)
-TfLiteDelegatePtr CreateXNNPACKDelegate(int num_threads, bool force_fp16);
+TfLiteDelegatePtr CreateXNNPACKDelegate(
+    int num_threads, bool force_fp16,
+    const char* experimental_weight_cache_file_path = nullptr);
 
 TfLiteDelegatePtr CreateCoreMlDelegate();
 }  // namespace evaluation
diff --git a/tensorflow/lite/tools/tool_params.h b/tensorflow/lite/tools/tool_params.h
index 70a16ccd2a73f6..09f64797f3eebf 100644
--- a/tensorflow/lite/tools/tool_params.h
+++ b/tensorflow/lite/tools/tool_params.h
@@ -94,7 +94,7 @@ class TypedToolParam : public ToolParam {
     has_value_set_ = true;
   }
 
-  T Get() const { return value_; }
+  const T& Get() const { return value_; }
 
   void Set(const ToolParam& other) override {
     Set(other.AsConstTyped<T>()->Get());

From 74fd07ce4919084fa703725a588db4b8a941ea56 Mon Sep 17 00:00:00 2001
From: Kyle Lucke <klucke@google.com>
Date: Thu, 16 May 2024 11:26:47 -0700
Subject: [PATCH 290/478] Use absl::Status instead of xla::Status now that
 they're identical.

PiperOrigin-RevId: 634473357
---
 third_party/xla/xla/pjrt/exceptions.h         |   4 +-
 .../xla/xla/pjrt/host_callback_test.cc        |  13 +-
 .../xla/xla/pjrt/interpreter_device.cc        |   2 +-
 third_party/xla/xla/pjrt/interpreter_device.h |   2 +-
 third_party/xla/xla/pjrt/layout_mode.cc       |  11 +-
 third_party/xla/xla/pjrt/layout_mode.h        |   2 +-
 .../xla/xla/pjrt/local_device_state.cc        |  10 +-
 third_party/xla/xla/pjrt/local_device_state.h |  11 +-
 third_party/xla/xla/pjrt/mlir_to_hlo.cc       |  12 +-
 third_party/xla/xla/pjrt/mlir_to_hlo.h        |  12 +-
 third_party/xla/xla/pjrt/pjrt_c_api_client.cc | 114 ++++++-----
 third_party/xla/xla/pjrt/pjrt_c_api_client.h  | 158 +++++++-------
 third_party/xla/xla/pjrt/pjrt_client.cc       |   5 +-
 third_party/xla/xla/pjrt/pjrt_client.h        | 192 +++++++++---------
 third_party/xla/xla/pjrt/pjrt_client_test.cc  |   4 +-
 third_party/xla/xla/pjrt/pjrt_compiler.cc     |   4 +-
 third_party/xla/xla/pjrt/pjrt_compiler.h      |  10 +-
 .../xla/xla/pjrt/pjrt_compiler_test.cc        |   8 +-
 third_party/xla/xla/pjrt/pjrt_executable.cc   |  25 +--
 third_party/xla/xla/pjrt/pjrt_executable.h    |  46 +++--
 third_party/xla/xla/pjrt/pjrt_future.h        |   5 +-
 third_party/xla/xla/pjrt/pjrt_layout.h        |   3 +-
 .../xla/pjrt/pjrt_stream_executor_client.cc   | 170 ++++++++--------
 .../xla/pjrt/pjrt_stream_executor_client.h    | 152 +++++++-------
 .../pjrt/pjrt_stream_executor_client_test.cc  |   4 +-
 third_party/xla/xla/pjrt/status_casters.h     |   2 +-
 .../xla/pjrt/stream_executor_executable.cc    |   3 +-
 third_party/xla/xla/pjrt/tf_pjrt_client.cc    |  17 +-
 third_party/xla/xla/pjrt/tf_pjrt_client.h     |  91 +++++----
 .../xla/xla/pjrt/tracked_device_buffer.cc     |   2 +-
 .../xla/xla/pjrt/tracked_device_buffer.h      |   6 +-
 .../xla/pjrt/tracked_device_buffer_test.cc    |   6 +-
 third_party/xla/xla/pjrt/transpose.cc         |   8 +-
 third_party/xla/xla/pjrt/transpose.h          |   9 +-
 third_party/xla/xla/pjrt/utils.cc             |  66 +++---
 third_party/xla/xla/pjrt/utils.h              |  40 ++--
 36 files changed, 636 insertions(+), 593 deletions(-)

diff --git a/third_party/xla/xla/pjrt/exceptions.h b/third_party/xla/xla/pjrt/exceptions.h
index bf72c4f4fe5ec7..4ef8d544958a82 100644
--- a/third_party/xla/xla/pjrt/exceptions.h
+++ b/third_party/xla/xla/pjrt/exceptions.h
@@ -33,7 +33,7 @@ namespace xla {
 // Python code instead of RuntimeError.
 class XlaRuntimeError : public std::runtime_error {
  public:
-  explicit XlaRuntimeError(Status status)
+  explicit XlaRuntimeError(absl::Status status)
       : std::runtime_error(StatusToString(status)), status_(status) {
     CHECK(!status_->ok());
   }
@@ -43,7 +43,7 @@ class XlaRuntimeError : public std::runtime_error {
   std::optional<Status> status() const { return status_; }
 
  private:
-  static std::string StatusToString(const Status& st) {
+  static std::string StatusToString(const absl::Status& st) {
     if (!ShowStackTraces()) {
       return st.ToString(absl::StatusToStringMode::kWithNoExtraData);
     }
diff --git a/third_party/xla/xla/pjrt/host_callback_test.cc b/third_party/xla/xla/pjrt/host_callback_test.cc
index d142b878471539..d134b216e400df 100644
--- a/third_party/xla/xla/pjrt/host_callback_test.cc
+++ b/third_party/xla/xla/pjrt/host_callback_test.cc
@@ -32,17 +32,18 @@ class TestPjRtHostMemoryForDeviceManager
  public:
   ~TestPjRtHostMemoryForDeviceManager() override = default;
 
-  StatusOr<PjRtChunk> ToDeviceLayout(const void* src_data, size_t src_size,
-                                     const Shape& host_shape,
-                                     const Shape& device_shape) override {
+  absl::StatusOr<PjRtChunk> ToDeviceLayout(const void* src_data,
+                                           size_t src_size,
+                                           const Shape& host_shape,
+                                           const Shape& device_shape) override {
     auto chunk = PjRtChunk::AllocateDefault(src_size);
     std::memcpy(chunk.data(), src_data, src_size);
     return chunk;
   }
 
-  Status ToHostLayout(const void* src_data, size_t src_size,
-                      const Shape& src_shape, void* dst_data, size_t dst_size,
-                      const Shape& dst_shape) override {
+  absl::Status ToHostLayout(const void* src_data, size_t src_size,
+                            const Shape& src_shape, void* dst_data,
+                            size_t dst_size, const Shape& dst_shape) override {
     CHECK_EQ(src_size, dst_size);
     std::memcpy(dst_data, src_data, src_size);
     return OkStatus();
diff --git a/third_party/xla/xla/pjrt/interpreter_device.cc b/third_party/xla/xla/pjrt/interpreter_device.cc
index 48a233da0948ca..cfa12b96392645 100644
--- a/third_party/xla/xla/pjrt/interpreter_device.cc
+++ b/third_party/xla/xla/pjrt/interpreter_device.cc
@@ -33,7 +33,7 @@ InterpreterDevice::InterpreterDevice(
     : PjRtStreamExecutorDevice(id, std::move(local_device_state),
                                /*device_kind=*/kInterpreterPlatformName) {}
 
-StatusOr<std::unique_ptr<PjRtClient>> GetInterpreterClient() {
+absl::StatusOr<std::unique_ptr<PjRtClient>> GetInterpreterClient() {
   TF_ASSIGN_OR_RETURN(se::Platform * platform,
                       PlatformUtil::GetPlatform("Interpreter"));
   if (platform->VisibleDeviceCount() != 1) {
diff --git a/third_party/xla/xla/pjrt/interpreter_device.h b/third_party/xla/xla/pjrt/interpreter_device.h
index 1842a816e9eb28..7952f362cf9a6b 100644
--- a/third_party/xla/xla/pjrt/interpreter_device.h
+++ b/third_party/xla/xla/pjrt/interpreter_device.h
@@ -29,7 +29,7 @@ class InterpreterDevice : public PjRtStreamExecutorDevice {
                     std::unique_ptr<LocalDeviceState> local_device_state);
 };
 
-StatusOr<std::unique_ptr<PjRtClient>> GetInterpreterClient();
+absl::StatusOr<std::unique_ptr<PjRtClient>> GetInterpreterClient();
 
 }  // namespace xla
 
diff --git a/third_party/xla/xla/pjrt/layout_mode.cc b/third_party/xla/xla/pjrt/layout_mode.cc
index 864a03799364db..b236f066598995 100644
--- a/third_party/xla/xla/pjrt/layout_mode.cc
+++ b/third_party/xla/xla/pjrt/layout_mode.cc
@@ -50,7 +50,7 @@ std::string LayoutMode::ToString() const {
   }
 }
 
-StatusOr<LayoutMode> LayoutMode::FromString(std::string s) {
+absl::StatusOr<LayoutMode> LayoutMode::FromString(std::string s) {
   if (s == "default") {
     return LayoutMode(Mode::kDefault);
   }
@@ -58,11 +58,12 @@ StatusOr<LayoutMode> LayoutMode::FromString(std::string s) {
     return LayoutMode(Mode::kAuto);
   }
   // LayoutMode is user-specified; parse Layout string
-  StatusOr<Layout> layout = ParseLayout(s);
+  absl::StatusOr<Layout> layout = ParseLayout(s);
   if (!layout.ok()) {
-    Status new_status(layout.status().code(),
-                      absl::StrCat("Error parsing user-specified layout mode '",
-                                   s, "': ", layout.status().message()));
+    absl::Status new_status(
+        layout.status().code(),
+        absl::StrCat("Error parsing user-specified layout mode '", s,
+                     "': ", layout.status().message()));
     return new_status;
   }
   return LayoutMode(*layout);
diff --git a/third_party/xla/xla/pjrt/layout_mode.h b/third_party/xla/xla/pjrt/layout_mode.h
index 156932a2f71c40..3e8d208fa89859 100644
--- a/third_party/xla/xla/pjrt/layout_mode.h
+++ b/third_party/xla/xla/pjrt/layout_mode.h
@@ -59,7 +59,7 @@ struct LayoutMode {
   // the correct format for the "mhlo.layout_mode" attribute.
   std::string ToString() const;
   // Parses a string produced by LayoutMode::ToString() or Layout::ToString().
-  static StatusOr<LayoutMode> FromString(std::string s);
+  static absl::StatusOr<LayoutMode> FromString(std::string s);
 };
 
 }  // namespace xla
diff --git a/third_party/xla/xla/pjrt/local_device_state.cc b/third_party/xla/xla/pjrt/local_device_state.cc
index 5c0d2b30e70325..fca99bcd391f58 100644
--- a/third_party/xla/xla/pjrt/local_device_state.cc
+++ b/third_party/xla/xla/pjrt/local_device_state.cc
@@ -91,14 +91,14 @@ LocalDeviceState::LocalDeviceState(se::StreamExecutor* executor,
 }
 
 LocalDeviceState::~LocalDeviceState() {
-  Status status = SynchronizeAllActivity();
+  absl::Status status = SynchronizeAllActivity();
   if (!status.ok()) {
     LOG(ERROR) << "Error when closing device: " << status;
   }
 }
 
-Status LocalDeviceState::SynchronizeAllActivity() {
-  Status status;
+absl::Status LocalDeviceState::SynchronizeAllActivity() {
+  absl::Status status;
   // TODO(phawkins): in theory the call to SynchronizeAllActivity below should
   // suffice. However on the Host platform SynchronizeAllActivity is a dummy
   // implementation that doesn't actually block. To make sure activity has
@@ -121,7 +121,7 @@ Status LocalDeviceState::SynchronizeAllActivity() {
   return status;
 }
 
-Status LocalDeviceState::ThenMemcpyDeviceToDevice(
+absl::Status LocalDeviceState::ThenMemcpyDeviceToDevice(
     se::Stream* transfer_stream, se::Stream* dst_stream,
     se::DeviceMemoryBase src_buffer, se::DeviceMemoryBase dst_buffer) {
   // The default implementation simply calls MemcpyD2D, and assumes that
@@ -175,7 +175,7 @@ se::Stream* LocalDeviceState::GetExternalReadyEventStream() {
   return external_ready_event_streams_.at(i).get();
 }
 
-StatusOr<se::Stream*> LocalDeviceState::GetStreamFromExternalStream(
+absl::StatusOr<se::Stream*> LocalDeviceState::GetStreamFromExternalStream(
     std::intptr_t stream) {
   // TODO(skyewm): replace with map lookup if performance is an issue (currently
   // it just iterates over 4 streams).
diff --git a/third_party/xla/xla/pjrt/local_device_state.h b/third_party/xla/xla/pjrt/local_device_state.h
index 73206d4fbd3abf..deece8131acecb 100644
--- a/third_party/xla/xla/pjrt/local_device_state.h
+++ b/third_party/xla/xla/pjrt/local_device_state.h
@@ -146,7 +146,7 @@ class LocalDeviceState {
   // returned by GetExternalReadyEventStream.
   // TODO(skyewm): this function could map other raw streams if needed. It's
   // currently only used with external ready event streams.
-  StatusOr<se::Stream*> GetStreamFromExternalStream(std::intptr_t stream);
+  absl::StatusOr<se::Stream*> GetStreamFromExternalStream(std::intptr_t stream);
 
   // Returns a vector of device to device streams.
   std::vector<se::Stream*> GetDeviceToDeviceStreams();
@@ -159,10 +159,9 @@ class LocalDeviceState {
   void ReturnStreamToPool(std::unique_ptr<se::Stream> stream);
 
   // Enqueues a copy of `src_buffer` to `dst_buffer` onto `transfer_stream`.
-  virtual Status ThenMemcpyDeviceToDevice(se::Stream* transfer_stream,
-                                          se::Stream* dst_stream,
-                                          se::DeviceMemoryBase src_buffer,
-                                          se::DeviceMemoryBase dst_buffer);
+  virtual absl::Status ThenMemcpyDeviceToDevice(
+      se::Stream* transfer_stream, se::Stream* dst_stream,
+      se::DeviceMemoryBase src_buffer, se::DeviceMemoryBase dst_buffer);
 
   WorkerThread* execute_thread() const { return execute_thread_.get(); }
 
@@ -196,7 +195,7 @@ class LocalDeviceState {
   int GetNewPrngSeed();
 
  private:
-  Status SynchronizeAllActivity();
+  absl::Status SynchronizeAllActivity();
 
   AllocationModel allocation_model_;
 
diff --git a/third_party/xla/xla/pjrt/mlir_to_hlo.cc b/third_party/xla/xla/pjrt/mlir_to_hlo.cc
index 2c27c5daa65d5d..d2d74da471083a 100644
--- a/third_party/xla/xla/pjrt/mlir_to_hlo.cc
+++ b/third_party/xla/xla/pjrt/mlir_to_hlo.cc
@@ -193,9 +193,9 @@ void UpgradeStablehlo(mlir::ModuleOp module) {
 
 }  // namespace
 
-Status MlirToXlaComputation(mlir::ModuleOp module,
-                            XlaComputation& xla_computation,
-                            bool use_tuple_args, bool return_tuple) {
+absl::Status MlirToXlaComputation(mlir::ModuleOp module,
+                                  XlaComputation& xla_computation,
+                                  bool use_tuple_args, bool return_tuple) {
   mlir::BaseScopedDiagnosticHandler diagnostic_handler(module->getContext());
   {
     mlir::PassManager pm(module->getContext());
@@ -227,7 +227,7 @@ Status MlirToXlaComputation(mlir::ModuleOp module,
   return OkStatus();
 }
 
-StatusOr<mlir::OwningOpRef<mlir::ModuleOp>> ParseMlirModuleString(
+absl::StatusOr<mlir::OwningOpRef<mlir::ModuleOp>> ParseMlirModuleString(
     absl::string_view mlir_module_str, mlir::MLIRContext& context) {
   mlir::DialectRegistry registry;
   registry.insert<mlir::arith::ArithDialect>();
@@ -267,7 +267,7 @@ StatusOr<mlir::OwningOpRef<mlir::ModuleOp>> ParseMlirModuleString(
   return std::move(module);
 }
 
-Status ParseMlirModuleStringAndConvertToXlaComputation(
+absl::Status ParseMlirModuleStringAndConvertToXlaComputation(
     absl::string_view mlir_module_str, XlaComputation& xla_computation,
     bool use_tuple_args, bool return_tuple) {
   mlir::MLIRContext context;
@@ -334,7 +334,7 @@ absl::StatusOr<std::string> SerializeUsingVersionedStablehlo(
   return buffer;
 }
 
-Status UpgradeVersionedStablehlo(mlir::ModuleOp mlir_module) {
+absl::Status UpgradeVersionedStablehlo(mlir::ModuleOp mlir_module) {
   // Apply StableHLO bytecode patch
   UpgradeStablehlo(mlir_module);
 
diff --git a/third_party/xla/xla/pjrt/mlir_to_hlo.h b/third_party/xla/xla/pjrt/mlir_to_hlo.h
index 9bad5399717422..efce0dc48f28d9 100644
--- a/third_party/xla/xla/pjrt/mlir_to_hlo.h
+++ b/third_party/xla/xla/pjrt/mlir_to_hlo.h
@@ -24,16 +24,16 @@ limitations under the License.
 namespace xla {
 
 // Converts an MHLO/CHLO module string to an mlir::Module.
-StatusOr<mlir::OwningOpRef<mlir::ModuleOp>> ParseMlirModuleString(
+absl::StatusOr<mlir::OwningOpRef<mlir::ModuleOp>> ParseMlirModuleString(
     absl::string_view mlir_module_str, mlir::MLIRContext& context);
 
 // Converts an CHLO/MHLO module to XLA HLO.
-Status MlirToXlaComputation(mlir::ModuleOp module,
-                            XlaComputation& xla_computation,
-                            bool use_tuple_args, bool return_tuple);
+absl::Status MlirToXlaComputation(mlir::ModuleOp module,
+                                  XlaComputation& xla_computation,
+                                  bool use_tuple_args, bool return_tuple);
 
 // Converts an MHLO/CHLO module string to an XLA computation.
-Status ParseMlirModuleStringAndConvertToXlaComputation(
+absl::Status ParseMlirModuleStringAndConvertToXlaComputation(
     absl::string_view mlir_module_str, XlaComputation& xla_computation,
     bool use_tuple_args, bool return_tuple);
 
@@ -62,7 +62,7 @@ absl::StatusOr<std::string> SerializeUsingVersionedStablehlo(
 // on portable artifacts that are outside of the compatibility window.
 // `ParseMlirModuleString` uses this method, and should be preferred to directly
 // calling `UpgradeVersionedStablehlo` where possible.
-Status UpgradeVersionedStablehlo(mlir::ModuleOp mlir_module);
+absl::Status UpgradeVersionedStablehlo(mlir::ModuleOp mlir_module);
 
 }  // namespace xla
 
diff --git a/third_party/xla/xla/pjrt/pjrt_c_api_client.cc b/third_party/xla/xla/pjrt/pjrt_c_api_client.cc
index 7248e9fe153b8b..e8f5d3926c0465 100644
--- a/third_party/xla/xla/pjrt/pjrt_c_api_client.cc
+++ b/third_party/xla/xla/pjrt/pjrt_c_api_client.cc
@@ -99,9 +99,9 @@ namespace xla {
 
 // ---------------------------------- Client -----------------------------------
 
-static StatusOr<const PjRtCApiTopologyDescription> InitClientTopoDesc(
+static absl::StatusOr<const PjRtCApiTopologyDescription> InitClientTopoDesc(
     const PJRT_Api* c_api, PJRT_Client* c_client) {
-  StatusOr<PJRT_TopologyDescription*> c_topo =
+  absl::StatusOr<PJRT_TopologyDescription*> c_topo =
       pjrt::GetTopologyDescription(c_client, c_api);
   TF_RETURN_IF_ERROR(c_topo.status());
   return PjRtCApiTopologyDescription(c_api, *c_topo, /*owned=*/false);
@@ -303,7 +303,7 @@ static DeviceAssignment CalculateDefaultAssignment(
   return cpp_device_assignment;
 }
 
-StatusOr<DeviceAssignment> PjRtCApiClient::GetDefaultDeviceAssignment(
+absl::StatusOr<DeviceAssignment> PjRtCApiClient::GetDefaultDeviceAssignment(
     int num_replicas, int num_partitions) const {
   PJRT_Client_DefaultDeviceAssignment_Args args;
   args.struct_size = PJRT_Client_DefaultDeviceAssignment_Args_STRUCT_SIZE;
@@ -322,7 +322,7 @@ StatusOr<DeviceAssignment> PjRtCApiClient::GetDefaultDeviceAssignment(
                                     param);
 }
 
-StatusOr<PjRtDevice*> PjRtCApiClient::LookupDevice(
+absl::StatusOr<PjRtDevice*> PjRtCApiClient::LookupDevice(
     PjRtGlobalDeviceId global_device_id) const {
   PJRT_Client_LookupDevice_Args args;
   args.struct_size = PJRT_Client_LookupDevice_Args_STRUCT_SIZE;
@@ -333,12 +333,12 @@ StatusOr<PjRtDevice*> PjRtCApiClient::LookupDevice(
   return GetCppDevice(args.device);
 }
 
-StatusOr<PjRtDevice*> PjRtCApiClient::LookupAddressableDevice(
+absl::StatusOr<PjRtDevice*> PjRtCApiClient::LookupAddressableDevice(
     int local_hardware_id) const {
   return LookupAddressableDevice(PjRtLocalDeviceId(local_hardware_id));
 }
 
-StatusOr<PjRtDevice*> PjRtCApiClient::LookupAddressableDevice(
+absl::StatusOr<PjRtDevice*> PjRtCApiClient::LookupAddressableDevice(
     PjRtLocalDeviceId local_device_id) const {
   PJRT_Client_LookupAddressableDevice_Args args;
   args.struct_size = PJRT_Client_LookupAddressableDevice_Args_STRUCT_SIZE;
@@ -356,10 +356,10 @@ absl::Span<PjRtMemorySpace* const> PjRtCApiClient::memory_spaces() const {
 
 // Initializes `PJRT_Client_Compile_Args`, which will be used to call
 // API PJRT_Client_Compile().
-static StatusOr<std::unique_ptr<PjRtLoadedExecutable>> InitializeArgsAndCompile(
-    PjRtCApiClient* api_client, const PJRT_Api* c_api, PJRT_Client* client,
-    const CompileOptions& options, const std::string& code,
-    const std::string& format) {
+static absl::StatusOr<std::unique_ptr<PjRtLoadedExecutable>>
+InitializeArgsAndCompile(PjRtCApiClient* api_client, const PJRT_Api* c_api,
+                         PJRT_Client* client, const CompileOptions& options,
+                         const std::string& code, const std::string& format) {
   PJRT_Client_Compile_Args args;
   args.struct_size = PJRT_Client_Compile_Args_STRUCT_SIZE;
   PJRT_Profiler_Extension profiler_extension =
@@ -388,7 +388,7 @@ static StatusOr<std::unique_ptr<PjRtLoadedExecutable>> InitializeArgsAndCompile(
   return ret;
 }
 
-StatusOr<std::unique_ptr<PjRtLoadedExecutable>> PjRtCApiClient::Compile(
+absl::StatusOr<std::unique_ptr<PjRtLoadedExecutable>> PjRtCApiClient::Compile(
     const XlaComputation& computation, CompileOptions options) {
   std::string module_str = computation.proto().SerializeAsString();
   std::string format(pjrt::kHloFormat);
@@ -396,7 +396,7 @@ StatusOr<std::unique_ptr<PjRtLoadedExecutable>> PjRtCApiClient::Compile(
                                   module_str, format);
 }
 
-StatusOr<std::unique_ptr<PjRtLoadedExecutable>> PjRtCApiClient::Compile(
+absl::StatusOr<std::unique_ptr<PjRtLoadedExecutable>> PjRtCApiClient::Compile(
     mlir::ModuleOp module, CompileOptions options) {
   // TODO: Once plugins are ready, use SerializeUsingVersionedStablehlo.
   if (!pjrt_c_api()) llvm::report_fatal_error("pjrt_c_api is null");
@@ -409,7 +409,7 @@ StatusOr<std::unique_ptr<PjRtLoadedExecutable>> PjRtCApiClient::Compile(
                                   serialized, format);
 }
 
-StatusOr<std::unique_ptr<PjRtLoadedExecutable>>
+absl::StatusOr<std::unique_ptr<PjRtLoadedExecutable>>
 PjRtCApiClient::DeserializeExecutable(absl::string_view serialized,
                                       std::optional<CompileOptions> options) {
   PJRT_Executable_DeserializeAndLoad_Args des_args;
@@ -430,7 +430,7 @@ PjRtCApiClient::DeserializeExecutable(absl::string_view serialized,
       std::make_unique<PjRtCApiLoadedExecutable>(this, c_exec));
 }
 
-StatusOr<const PjRtTopologyDescription*>
+absl::StatusOr<const PjRtTopologyDescription*>
 PjRtCApiClient::GetTopologyDescription() const {
   if (!topo_desc_.ok()) {
     return topo_desc_.status();
@@ -438,7 +438,7 @@ PjRtCApiClient::GetTopologyDescription() const {
   return &(*topo_desc_);
 }
 
-StatusOr<std::uintptr_t> PjRtCApiClient::UnsafeBufferPointer(
+absl::StatusOr<std::uintptr_t> PjRtCApiClient::UnsafeBufferPointer(
     PjRtBuffer* buffer) {
   // Validate that the buffer's client matches the function call's client, since
   // that could be a common error.
@@ -464,7 +464,7 @@ StatusOr<std::uintptr_t> PjRtCApiClient::UnsafeBufferPointer(
   return args.buffer_pointer;
 }
 
-StatusOr<std::unique_ptr<PjRtBuffer>>
+absl::StatusOr<std::unique_ptr<PjRtBuffer>>
 PjRtCApiClient::BufferFromHostBufferInternalImpl(
     const void* data, PrimitiveType type, absl::Span<int64_t const> dims,
     std::optional<absl::Span<int64_t const>> byte_strides,
@@ -559,7 +559,8 @@ PjRtCApiClient::BufferFromHostBufferInternalImpl(
   return buffer;
 }
 
-StatusOr<std::unique_ptr<PjRtBuffer>> PjRtCApiClient::BufferFromHostBuffer(
+absl::StatusOr<std::unique_ptr<PjRtBuffer>>
+PjRtCApiClient::BufferFromHostBuffer(
     const void* data, PrimitiveType type, absl::Span<int64_t const> dims,
     std::optional<absl::Span<int64_t const>> byte_strides,
     HostBufferSemantics host_buffer_semantics,
@@ -570,7 +571,8 @@ StatusOr<std::unique_ptr<PjRtBuffer>> PjRtCApiClient::BufferFromHostBuffer(
       std::move(on_done_with_host_buffer), memory_space, device_layout);
 }
 
-StatusOr<std::unique_ptr<PjRtBuffer>> PjRtCApiClient::BufferFromHostBuffer(
+absl::StatusOr<std::unique_ptr<PjRtBuffer>>
+PjRtCApiClient::BufferFromHostBuffer(
     const void* data, PrimitiveType type, absl::Span<int64_t const> dims,
     std::optional<absl::Span<int64_t const>> byte_strides,
     HostBufferSemantics host_buffer_semantics,
@@ -581,7 +583,8 @@ StatusOr<std::unique_ptr<PjRtBuffer>> PjRtCApiClient::BufferFromHostBuffer(
       std::move(on_done_with_host_buffer), device, device_layout);
 }
 
-StatusOr<std::unique_ptr<PjRtBuffer>> PjRtCApiClient::BufferFromHostBuffer(
+absl::StatusOr<std::unique_ptr<PjRtBuffer>>
+PjRtCApiClient::BufferFromHostBuffer(
     const void* data, PrimitiveType type, absl::Span<int64_t const> dims,
     std::optional<absl::Span<int64_t const>> byte_strides,
     HostBufferSemantics host_buffer_semantics,
@@ -592,7 +595,8 @@ StatusOr<std::unique_ptr<PjRtBuffer>> PjRtCApiClient::BufferFromHostBuffer(
       std::move(on_done_with_host_buffer), device, /*device_layout=*/nullptr);
 }
 
-StatusOr<std::unique_ptr<PjRtBuffer>> PjRtCApiClient::CreateViewOfDeviceBuffer(
+absl::StatusOr<std::unique_ptr<PjRtBuffer>>
+PjRtCApiClient::CreateViewOfDeviceBuffer(
     void* device_ptr, const Shape& shape, PjRtDevice* device,
     std::function<void()> on_delete_callback,
     std::optional<std::intptr_t> stream) {
@@ -835,7 +839,7 @@ PjRtLocalHardwareId PjRtCApiDevice::local_hardware_id_typed() const {
   return PjRtLocalHardwareId(args.local_hardware_id);
 }
 
-StatusOr<PjRtMemorySpace*> PjRtCApiDevice::default_memory_space() const {
+absl::StatusOr<PjRtMemorySpace*> PjRtCApiDevice::default_memory_space() const {
   PJRT_Device_DefaultMemory_Args args;
   args.struct_size = PJRT_Device_DefaultMemory_Args_STRUCT_SIZE;
   args.extension_start = nullptr;
@@ -845,7 +849,7 @@ StatusOr<PjRtMemorySpace*> PjRtCApiDevice::default_memory_space() const {
   return client_->GetCppMemory(args.memory);
 }
 
-StatusOr<tsl::AllocatorStats> PjRtCApiDevice::GetAllocatorStats() const {
+absl::StatusOr<tsl::AllocatorStats> PjRtCApiDevice::GetAllocatorStats() const {
   PJRT_Device_MemoryStats_Args args;
   args.struct_size = PJRT_Device_MemoryStats_Args_STRUCT_SIZE;
   args.extension_start = nullptr;
@@ -1042,7 +1046,7 @@ int64_t PjRtCApiExecutable::SizeOfGeneratedCodeInBytes() const {
   return args.size_in_bytes;
 }
 
-StatusOr<absl::flat_hash_map<std::string, PjRtValueType>>
+absl::StatusOr<absl::flat_hash_map<std::string, PjRtValueType>>
 PjRtCApiExecutable::GetCostAnalysis() const {
   // Initialize function call args
   PJRT_Executable_GetCostAnalysis_Args args;
@@ -1060,7 +1064,7 @@ PjRtCApiExecutable::GetCostAnalysis() const {
                                              args.num_properties);
 }
 
-StatusOr<std::vector<std::vector<PrimitiveType>>>
+absl::StatusOr<std::vector<std::vector<PrimitiveType>>>
 PjRtCApiExecutable::GetOutputElementTypes() const {
   PJRT_Executable_OutputElementTypes_Args args;
   args.struct_size = PJRT_Executable_OutputElementTypes_Args_STRUCT_SIZE;
@@ -1080,7 +1084,7 @@ PjRtCApiExecutable::GetOutputElementTypes() const {
   return std::vector<std::vector<PrimitiveType>>{std::move(out)};
 }
 
-StatusOr<std::vector<std::vector<DimensionVector>>>
+absl::StatusOr<std::vector<std::vector<DimensionVector>>>
 PjRtCApiExecutable::GetOutputDimensions() const {
   PJRT_Executable_OutputDimensions_Args args;
   args.struct_size = PJRT_Executable_OutputDimensions_Args_STRUCT_SIZE;
@@ -1106,7 +1110,7 @@ PjRtCApiExecutable::GetOutputDimensions() const {
   return std::vector<std::vector<DimensionVector>>{std::move(out)};
 }
 
-StatusOr<std::vector<std::vector<absl::string_view>>>
+absl::StatusOr<std::vector<std::vector<absl::string_view>>>
 PjRtCApiExecutable::GetOutputMemoryKinds() const {
   PJRT_Executable_OutputMemoryKinds_Args args;
   args.struct_size = PJRT_Executable_OutputMemoryKinds_Args_STRUCT_SIZE;
@@ -1126,7 +1130,7 @@ PjRtCApiExecutable::GetOutputMemoryKinds() const {
   return std::vector<std::vector<absl::string_view>>{std::move(out)};
 }
 
-StatusOr<std::vector<std::shared_ptr<HloModule>>>
+absl::StatusOr<std::vector<std::shared_ptr<HloModule>>>
 PjRtCApiExecutable::GetHloModules() const {
   auto* c_api = pjrt_c_api();
   auto* executable = c_executable();
@@ -1195,7 +1199,7 @@ PjRtCApiExecutable::GetHloModules() const {
   return out;
 }
 
-StatusOr<std::string> PjRtCApiExecutable::SerializeExecutable() const {
+absl::StatusOr<std::string> PjRtCApiExecutable::SerializeExecutable() const {
   auto* c_api = pjrt_c_api();
   auto* executable = c_executable();
   PJRT_Executable_Serialize_Args ser_args;
@@ -1212,7 +1216,7 @@ StatusOr<std::string> PjRtCApiExecutable::SerializeExecutable() const {
   return std::string(ser_args.serialized_bytes, ser_args.serialized_bytes_size);
 }
 
-StatusOr<std::string> PjRtCApiExecutable::FingerprintExecutable() const {
+absl::StatusOr<std::string> PjRtCApiExecutable::FingerprintExecutable() const {
   const PJRT_Api* c_api_ = pjrt_c_api();
   PJRT_Executable_Fingerprint_Args args;
   args.struct_size = PJRT_Executable_Fingerprint_Args_STRUCT_SIZE;
@@ -1585,7 +1589,7 @@ PjRtCApiLoadedExecutable::GetCommonExecuteArgs(
   return args;
 }
 
-StatusOr<std::vector<std::vector<std::unique_ptr<PjRtBuffer>>>>
+absl::StatusOr<std::vector<std::vector<std::unique_ptr<PjRtBuffer>>>>
 PjRtCApiLoadedExecutable::Execute(
     absl::Span<const std::vector<PjRtBuffer*>> argument_handles,
     const ExecuteOptions& options,
@@ -1648,15 +1652,15 @@ PjRtCApiLoadedExecutable::Execute(
                                        client_);
 }
 
-StatusOr<std::vector<std::unique_ptr<PjRtBuffer>>>
+absl::StatusOr<std::vector<std::unique_ptr<PjRtBuffer>>>
 PjRtCApiLoadedExecutable::ExecuteWithSingleDevice(
     absl::Span<PjRtBuffer* const> argument_handles, PjRtDevice* device,
     const ExecuteOptions& options, std::optional<PjRtFuture<>>& returned_future,
     bool fill_future) {
   if (!options.send_callbacks.empty() || !options.recv_callbacks.empty()) {
-    return Status(absl::StatusCode::kUnimplemented,
-                  "Send/recv callbacks not implemented for "
-                  "PjRtCApiLoadedExecutable::ExecuteWithSingleDevice.");
+    return absl::Status(absl::StatusCode::kUnimplemented,
+                        "Send/recv callbacks not implemented for "
+                        "PjRtCApiLoadedExecutable::ExecuteWithSingleDevice.");
   }
 
   std::vector<std::vector<PjRtBuffer*>> argument_handles_vec = {
@@ -1704,7 +1708,7 @@ PjRtCApiLoadedExecutable::ExecuteWithSingleDevice(
       client_)[0]);
 }
 
-StatusOr<std::vector<std::unique_ptr<PjRtBuffer>>>
+absl::StatusOr<std::vector<std::unique_ptr<PjRtBuffer>>>
 PjRtCApiLoadedExecutable::ExecuteSharded(
     absl::Span<PjRtBuffer* const> argument_handles, PjRtDevice* device,
     const ExecuteOptions& options, std::optional<PjRtFuture<>>& returned_future,
@@ -1713,7 +1717,7 @@ PjRtCApiLoadedExecutable::ExecuteSharded(
                                  returned_future, fill_future);
 }
 
-StatusOr<std::vector<std::unique_ptr<PjRtBuffer>>>
+absl::StatusOr<std::vector<std::unique_ptr<PjRtBuffer>>>
 PjRtCApiLoadedExecutable::ExecutePortable(
     absl::Span<PjRtBuffer* const> argument_handles, PjRtDevice* device,
     const ExecuteOptions& options, std::optional<PjRtFuture<>>& returned_future,
@@ -1743,8 +1747,10 @@ bool PjRtCApiLoadedExecutable::IsDeleted() {
   return args.is_deleted;
 }
 
-StatusOr<std::string> PjRtCApiLoadedExecutable::FingerprintExecutable() const {
-  StatusOr<std::string> fingerprint = executable_->FingerprintExecutable();
+absl::StatusOr<std::string> PjRtCApiLoadedExecutable::FingerprintExecutable()
+    const {
+  absl::StatusOr<std::string> fingerprint =
+      executable_->FingerprintExecutable();
   if (fingerprint.ok()) {
     return *fingerprint;
   }
@@ -1866,7 +1872,7 @@ absl::Span<const bool> PjRtCApiBuffer::is_dynamic_dimension() const {
   return *is_dynamic_dimension_;
 }
 
-StatusOr<std::vector<int64_t>> PjRtCApiBuffer::logical_dimensions() {
+absl::StatusOr<std::vector<int64_t>> PjRtCApiBuffer::logical_dimensions() {
   PJRT_Buffer_UnpaddedDimensions_Args args;
   args.struct_size = PJRT_Buffer_UnpaddedDimensions_Args_STRUCT_SIZE;
   args.extension_start = nullptr;
@@ -1927,7 +1933,7 @@ PjRtFuture<> PjRtCApiBuffer::ToLiteral(MutableLiteralBase* literal) {
   return pjrt::ConvertCEventToCppFuture(args.event, api);
 }
 
-StatusOr<size_t> PjRtCApiBuffer::GetOnDeviceSizeInBytes() const {
+absl::StatusOr<size_t> PjRtCApiBuffer::GetOnDeviceSizeInBytes() const {
   PJRT_Buffer_OnDeviceSizeInBytes_Args args;
   args.struct_size = PJRT_Buffer_OnDeviceSizeInBytes_Args_STRUCT_SIZE;
   args.extension_start = nullptr;
@@ -1985,7 +1991,7 @@ bool PjRtCApiBuffer::IsDeleted() {
   return args.is_deleted;
 }
 
-StatusOr<std::unique_ptr<PjRtBuffer>> PjRtCApiBuffer::CopyToDevice(
+absl::StatusOr<std::unique_ptr<PjRtBuffer>> PjRtCApiBuffer::CopyToDevice(
     PjRtDevice* dst_device) {
   if (dst_device->client() == client_) {
     PJRT_Buffer_CopyToDevice_Args args;
@@ -2016,7 +2022,7 @@ StatusOr<std::unique_ptr<PjRtBuffer>> PjRtCApiBuffer::CopyToDevice(
   }
 }
 
-StatusOr<std::unique_ptr<PjRtBuffer>> PjRtCApiBuffer::CopyToMemorySpace(
+absl::StatusOr<std::unique_ptr<PjRtBuffer>> PjRtCApiBuffer::CopyToMemorySpace(
     PjRtMemorySpace* dst_memory) {
   const PJRT_Api* api = pjrt_c_api();
 
@@ -2108,7 +2114,7 @@ PjRtFuture<> PjRtCApiBuffer::GetReadyFuture() {
   return PjRtFuture<>{*readiness_promise_};
 }
 
-StatusOr<std::unique_ptr<PjRtBuffer::ExternalReference>>
+absl::StatusOr<std::unique_ptr<PjRtBuffer::ExternalReference>>
 PjRtCApiBuffer::AcquireExternalReference() {
   PJRT_Buffer_IncreaseExternalReferenceCount_Args increase_reference_count_args;
   increase_reference_count_args.buffer = c_buffer();
@@ -2148,7 +2154,7 @@ PjRtCApiExternalReference::~PjRtCApiExternalReference() {
       client_->pjrt_c_api());
 }
 
-Status PjRtCApiExternalReference::WaitUntilBufferReadyOnStream(
+absl::Status PjRtCApiExternalReference::WaitUntilBufferReadyOnStream(
     std::intptr_t stream) {
   const PJRT_Api* c_api = buffer_->pjrt_c_api();
   PJRT_Stream_Extension* extension = pjrt::FindExtension<PJRT_Stream_Extension>(
@@ -2220,7 +2226,7 @@ PjRtCApiTopologyDescription::DeviceDescriptions() const {
   return out;
 }
 
-StatusOr<std::string> PjRtCApiTopologyDescription::Serialize() const {
+absl::StatusOr<std::string> PjRtCApiTopologyDescription::Serialize() const {
   PJRT_TopologyDescription_Serialize_Args args;
   args.struct_size = PJRT_TopologyDescription_Serialize_Args_STRUCT_SIZE;
   args.extension_start = nullptr;
@@ -2245,10 +2251,12 @@ void PjRtCApiTopologyDescription::InitAttributes() {
 
 // Initializes `PJRT_Compile_Args`, which will be used to call
 // API PJRT_Compile().
-static StatusOr<std::unique_ptr<PjRtExecutable>> InitializeArgsAndCompileAot(
-    const PJRT_Api* c_api, PjRtClient* client, const CompileOptions& options,
-    const PjRtTopologyDescription& topology, const std::string& code,
-    const std::string& format) {
+static absl::StatusOr<std::unique_ptr<PjRtExecutable>>
+InitializeArgsAndCompileAot(const PJRT_Api* c_api, PjRtClient* client,
+                            const CompileOptions& options,
+                            const PjRtTopologyDescription& topology,
+                            const std::string& code,
+                            const std::string& format) {
   PJRT_Compile_Args args;
   args.struct_size = PJRT_Compile_Args_STRUCT_SIZE;
   args.extension_start = nullptr;
@@ -2282,7 +2290,7 @@ static StatusOr<std::unique_ptr<PjRtExecutable>> InitializeArgsAndCompileAot(
   return ret;
 }
 
-StatusOr<std::unique_ptr<PjRtExecutable>> PjRtCApiCompiler::Compile(
+absl::StatusOr<std::unique_ptr<PjRtExecutable>> PjRtCApiCompiler::Compile(
     CompileOptions options, const XlaComputation& computation,
     const PjRtTopologyDescription& topology, PjRtClient* client) {
   std::string module_str = computation.proto().SerializeAsString();
@@ -2291,7 +2299,7 @@ StatusOr<std::unique_ptr<PjRtExecutable>> PjRtCApiCompiler::Compile(
                                      module_str, format);
 }
 
-StatusOr<std::unique_ptr<PjRtExecutable>> PjRtCApiCompiler::Compile(
+absl::StatusOr<std::unique_ptr<PjRtExecutable>> PjRtCApiCompiler::Compile(
     CompileOptions options, mlir::ModuleOp module,
     const PjRtTopologyDescription& topology, PjRtClient* client) {
   // TODO: Once plugins are ready, use SerializeUsingVersionedStablehlo.
@@ -2308,7 +2316,7 @@ StatusOr<std::unique_ptr<PjRtExecutable>> PjRtCApiCompiler::Compile(
 
 // -------------------------------- API access ---------------------------------
 
-StatusOr<std::unique_ptr<PjRtClient>> GetCApiClient(
+absl::StatusOr<std::unique_ptr<PjRtClient>> GetCApiClient(
     absl::string_view device_type,
     const absl::flat_hash_map<std::string, PjRtValueType>& create_options,
     std::shared_ptr<KeyValueStoreInterface> kv_store) {
@@ -2341,7 +2349,7 @@ StatusOr<std::unique_ptr<PjRtClient>> GetCApiClient(
       c_api, c_client, std::move(kv_callback_data)));
 }
 
-StatusOr<std::unique_ptr<PjRtTopologyDescription>> GetCApiTopology(
+absl::StatusOr<std::unique_ptr<PjRtTopologyDescription>> GetCApiTopology(
     absl::string_view device_type, absl::string_view topology_name,
     const absl::flat_hash_map<std::string, PjRtValueType>& create_options) {
   TF_ASSIGN_OR_RETURN(const PJRT_Api* c_api, pjrt::PjrtApi(device_type));
diff --git a/third_party/xla/xla/pjrt/pjrt_c_api_client.h b/third_party/xla/xla/pjrt/pjrt_c_api_client.h
index 1210713080a8c6..0a2c0dc0350ece 100644
--- a/third_party/xla/xla/pjrt/pjrt_c_api_client.h
+++ b/third_party/xla/xla/pjrt/pjrt_c_api_client.h
@@ -132,13 +132,13 @@ class PjRtCApiDevice : public PjRtDevice {
   int local_hardware_id() const override;
   PjRtLocalHardwareId local_hardware_id_typed() const override;
 
-  Status TransferToInfeed(const LiteralSlice& literal) override {
+  absl::Status TransferToInfeed(const LiteralSlice& literal) override {
     return Unimplemented(
         "PJRT C API does not support TransferToInfeed. Please report an issue "
         "at https://github.com/google/jax/issues if you need this feature.");
   }
 
-  Status TransferFromOutfeed(MutableBorrowingLiteral literal) override {
+  absl::Status TransferFromOutfeed(MutableBorrowingLiteral literal) override {
     return Unimplemented(
         "PJRT C API does not support TransferFromOutfeed. Please report an "
         "issue at https://github.com/google/jax/issues if you need this "
@@ -149,7 +149,7 @@ class PjRtCApiDevice : public PjRtDevice {
     return memory_spaces_;
   }
 
-  StatusOr<PjRtMemorySpace*> default_memory_space() const override;
+  absl::StatusOr<PjRtMemorySpace*> default_memory_space() const override;
 
   std::unique_ptr<ScopedAsyncTrackingEvent> CreateAsyncTrackingEvent(
       absl::string_view description) const override {
@@ -166,7 +166,7 @@ class PjRtCApiDevice : public PjRtDevice {
     return description_;
   }
 
-  StatusOr<tsl::AllocatorStats> GetAllocatorStats() const override;
+  absl::StatusOr<tsl::AllocatorStats> GetAllocatorStats() const override;
 
   absl::StatusOr<std::intptr_t> GetStreamForExternalReadyEvents()
       const override;
@@ -185,11 +185,11 @@ class PjRtCApiCompiler : public PjRtCompiler {
  public:
   explicit PjRtCApiCompiler(const PJRT_Api* c_api) : c_api_(c_api) {}
 
-  StatusOr<std::unique_ptr<PjRtExecutable>> Compile(
+  absl::StatusOr<std::unique_ptr<PjRtExecutable>> Compile(
       CompileOptions options, const XlaComputation& computation,
       const PjRtTopologyDescription& topology, PjRtClient* client) override;
 
-  StatusOr<std::unique_ptr<PjRtExecutable>> Compile(
+  absl::StatusOr<std::unique_ptr<PjRtExecutable>> Compile(
       CompileOptions options, mlir::ModuleOp module,
       const PjRtTopologyDescription& topology, PjRtClient* client) override;
 
@@ -231,7 +231,7 @@ class PjRtCApiTopologyDescription : public PjRtTopologyDescription {
     return attributes_;
   }
 
-  StatusOr<Layout> GetDefaultLayout(
+  absl::StatusOr<Layout> GetDefaultLayout(
       PrimitiveType element_type,
       absl::Span<const int64_t> dims) const override {
     return Unimplemented("PJRT C API does not support GetDefaultLayout");
@@ -267,12 +267,12 @@ class PjRtCApiClient : public PjRtClient {
   absl::Span<PjRtDevice* const> devices() const override;
   absl::Span<PjRtDevice* const> addressable_devices() const override;
 
-  StatusOr<PjRtDevice*> LookupDevice(
+  absl::StatusOr<PjRtDevice*> LookupDevice(
       PjRtGlobalDeviceId global_device_id) const override;
 
-  StatusOr<PjRtDevice*> LookupAddressableDevice(
+  absl::StatusOr<PjRtDevice*> LookupAddressableDevice(
       int local_hardware_id) const override;
-  StatusOr<PjRtDevice*> LookupAddressableDevice(
+  absl::StatusOr<PjRtDevice*> LookupAddressableDevice(
       PjRtLocalDeviceId local_device_id) const override;
 
   absl::Span<PjRtMemorySpace* const> memory_spaces() const override;
@@ -290,10 +290,10 @@ class PjRtCApiClient : public PjRtClient {
     return PjRtRuntimeType::kTfrt;
   }
 
-  StatusOr<DeviceAssignment> GetDefaultDeviceAssignment(
+  absl::StatusOr<DeviceAssignment> GetDefaultDeviceAssignment(
       int num_replicas, int num_partitions) const override;
 
-  StatusOr<std::unique_ptr<HloCostAnalysis>> GetHloCostAnalysis()
+  absl::StatusOr<std::unique_ptr<HloCostAnalysis>> GetHloCostAnalysis()
       const override {
     return Unimplemented(
         "PJRT C API does not support GetHloCostAnalysis. Please report an "
@@ -301,21 +301,21 @@ class PjRtCApiClient : public PjRtClient {
         "feature.");
   }
 
-  StatusOr<Layout> GetDefaultLayout(PrimitiveType element_type,
-                                    absl::Span<const int64_t> dims) override;
+  absl::StatusOr<Layout> GetDefaultLayout(
+      PrimitiveType element_type, absl::Span<const int64_t> dims) override;
 
-  StatusOr<std::unique_ptr<PjRtLoadedExecutable>> Compile(
+  absl::StatusOr<std::unique_ptr<PjRtLoadedExecutable>> Compile(
       const XlaComputation& computation, CompileOptions options) override;
 
-  StatusOr<std::unique_ptr<PjRtLoadedExecutable>> Compile(
+  absl::StatusOr<std::unique_ptr<PjRtLoadedExecutable>> Compile(
       mlir::ModuleOp module, CompileOptions options) override;
 
   // `PjRtCApiClient::DeserializeExecutable()` ignores `CompileOptions` arg
-  StatusOr<std::unique_ptr<PjRtLoadedExecutable>> DeserializeExecutable(
+  absl::StatusOr<std::unique_ptr<PjRtLoadedExecutable>> DeserializeExecutable(
       absl::string_view serialized,
       std::optional<CompileOptions> options) override;
 
-  StatusOr<std::unique_ptr<PjRtBuffer>> CreateUninitializedBuffer(
+  absl::StatusOr<std::unique_ptr<PjRtBuffer>> CreateUninitializedBuffer(
       const Shape& shape, PjRtDevice* device) override {
     return Unimplemented(
         "PJRT C API does not support CreateUninitializedBuffer. Please report "
@@ -323,10 +323,10 @@ class PjRtCApiClient : public PjRtClient {
         "feature.");
   }
 
-  StatusOr<const PjRtTopologyDescription*> GetTopologyDescription()
+  absl::StatusOr<const PjRtTopologyDescription*> GetTopologyDescription()
       const override;
 
-  StatusOr<std::unique_ptr<AsyncHostToDeviceTransferManager>>
+  absl::StatusOr<std::unique_ptr<AsyncHostToDeviceTransferManager>>
   CreateBuffersForAsyncHostToDevice(absl::Span<const Shape> shapes,
                                     PjRtDevice* device) override {
     return Unimplemented(
@@ -344,28 +344,28 @@ class PjRtCApiClient : public PjRtClient {
         "this feature.");
   }
 
-  StatusOr<std::unique_ptr<PjRtBuffer>> BufferFromHostBuffer(
+  absl::StatusOr<std::unique_ptr<PjRtBuffer>> BufferFromHostBuffer(
       const void* data, PrimitiveType type, absl::Span<int64_t const> dims,
       std::optional<absl::Span<int64_t const>> byte_strides,
       HostBufferSemantics host_buffer_semantics,
       absl::AnyInvocable<void() &&> on_done_with_host_buffer,
       PjRtDevice* device) override;
 
-  StatusOr<std::unique_ptr<PjRtBuffer>> BufferFromHostBuffer(
+  absl::StatusOr<std::unique_ptr<PjRtBuffer>> BufferFromHostBuffer(
       const void* data, PrimitiveType type, absl::Span<int64_t const> dims,
       std::optional<absl::Span<int64_t const>> byte_strides,
       HostBufferSemantics host_buffer_semantics,
       absl::AnyInvocable<void() &&> on_done_with_host_buffer,
       PjRtDevice* device, const Layout* device_layout) override;
 
-  StatusOr<std::unique_ptr<PjRtBuffer>> BufferFromHostBuffer(
+  absl::StatusOr<std::unique_ptr<PjRtBuffer>> BufferFromHostBuffer(
       const void* data, PrimitiveType type, absl::Span<int64_t const> dims,
       std::optional<absl::Span<int64_t const>> byte_strides,
       HostBufferSemantics host_buffer_semantics,
       absl::AnyInvocable<void() &&> on_done_with_host_buffer,
       PjRtMemorySpace* memory_space, const Layout* device_layout) override;
 
-  StatusOr<std::unique_ptr<PjRtBuffer>> BufferFromHostLiteral(
+  absl::StatusOr<std::unique_ptr<PjRtBuffer>> BufferFromHostLiteral(
       const LiteralSlice& literal, PjRtDevice* device) override {
     return Unimplemented(
         "PJRT C API does not support BufferFromHostLiteral. Please report an "
@@ -373,14 +373,15 @@ class PjRtCApiClient : public PjRtClient {
         "feature.");
   }
 
-  StatusOr<std::unique_ptr<PjRtBuffer>> CreateViewOfDeviceBuffer(
+  absl::StatusOr<std::unique_ptr<PjRtBuffer>> CreateViewOfDeviceBuffer(
       void* device_ptr, const Shape& shape, PjRtDevice* device,
       std::function<void()> on_delete_callback,
       std::optional<std::intptr_t> stream) override;
 
-  StatusOr<std::uintptr_t> UnsafeBufferPointer(PjRtBuffer* buffer) override;
+  absl::StatusOr<std::uintptr_t> UnsafeBufferPointer(
+      PjRtBuffer* buffer) override;
 
-  StatusOr<std::vector<std::unique_ptr<PjRtBuffer>>>
+  absl::StatusOr<std::vector<std::unique_ptr<PjRtBuffer>>>
   MakeCrossHostReceiveBuffers(absl::Span<const Shape> shapes,
                               PjRtDevice* device,
                               PjRtCrossHostRecvNotifier notifier) override {
@@ -390,7 +391,7 @@ class PjRtCApiClient : public PjRtClient {
         "this feature.");
   }
 
-  StatusOr<std::vector<std::unique_ptr<PjRtBuffer>>>
+  absl::StatusOr<std::vector<std::unique_ptr<PjRtBuffer>>>
   MakeCrossHostReceiveBuffersForGather(
       absl::Span<const Shape> shapes, std::vector<GatherDetails> gather_details,
       PjRtDevice* device, PjRtCrossHostRecvNotifier notifier) override {
@@ -400,28 +401,28 @@ class PjRtCApiClient : public PjRtClient {
         "this feature.");
   }
 
-  StatusOr<ChannelHandle> CreateChannelHandle() override {
+  absl::StatusOr<ChannelHandle> CreateChannelHandle() override {
     return Unimplemented(
         "PJRT C API does not support CreateChannelHandle. Please report an "
         "issue at https://github.com/google/jax/issues if you need this "
         "feature.");
   }
 
-  StatusOr<ChannelHandle> CreateDeviceToHostChannelHandle() override {
+  absl::StatusOr<ChannelHandle> CreateDeviceToHostChannelHandle() override {
     return Unimplemented(
         "PJRT C API does not support CreateDeviceToHostChannelHandle. Please "
         "report an issue at https://github.com/google/jax/issues if you need "
         "this feature.");
   }
 
-  StatusOr<ChannelHandle> CreateHostToDeviceChannelHandle() override {
+  absl::StatusOr<ChannelHandle> CreateHostToDeviceChannelHandle() override {
     return Unimplemented(
         "PJRT C API does not support CreateHostToDeviceChannelHandle. Please "
         "report an issue at https://github.com/google/jax/issues if you need "
         "this feature.");
   }
 
-  Status Defragment() override {
+  absl::Status Defragment() override {
     return Unimplemented(
         "PJRT C API does not support Defragment. Please report an issue at "
         "https://github.com/google/jax/issues if you need this feature.");
@@ -454,7 +455,7 @@ class PjRtCApiClient : public PjRtClient {
   void InitDevicesAndMemorySpaces();
   void InitAttributes();
 
-  StatusOr<std::unique_ptr<PjRtBuffer>> BufferFromHostBufferInternalImpl(
+  absl::StatusOr<std::unique_ptr<PjRtBuffer>> BufferFromHostBufferInternalImpl(
       const void* data, PrimitiveType type, absl::Span<int64_t const> dims,
       std::optional<absl::Span<int64_t const>> byte_strides,
       HostBufferSemantics host_buffer_semantics,
@@ -477,7 +478,7 @@ class PjRtCApiClient : public PjRtClient {
   // There may be an error fetching the topology desc via the C API
   // (e.g. unimplemented). Save the error during client init so we can return it
   // from GetTopologyDescription().
-  StatusOr<const PjRtCApiTopologyDescription> topo_desc_;
+  absl::StatusOr<const PjRtCApiTopologyDescription> topo_desc_;
 
   const std::string platform_version_;
   const std::string platform_name_;
@@ -506,9 +507,9 @@ class PjRtCApiBuffer : public PjRtBuffer {
 
   absl::Span<const bool> is_dynamic_dimension() const override;
 
-  StatusOr<std::vector<int64_t>> logical_dimensions() override;
+  absl::StatusOr<std::vector<int64_t>> logical_dimensions() override;
 
-  StatusOr<Shape> logical_on_device_shape() override {
+  absl::StatusOr<Shape> logical_on_device_shape() override {
     LOG(FATAL) << "PjRtBuffer::on_logical_device_shape() not implemented in "
                   "PJRT C API";
   }
@@ -519,7 +520,7 @@ class PjRtCApiBuffer : public PjRtBuffer {
 
   PjRtClient* client() const override { return client_; }
 
-  StatusOr<std::unique_ptr<ExternalReference>> AcquireExternalReference()
+  absl::StatusOr<std::unique_ptr<ExternalReference>> AcquireExternalReference()
       override;
 
   PjRtFuture<> ToLiteral(MutableLiteralBase* literal) override;
@@ -527,7 +528,7 @@ class PjRtCApiBuffer : public PjRtBuffer {
       absl::AnyInvocable<absl::StatusOr<MutableLiteralBase*>() &&> generator)
       override;
 
-  StatusOr<size_t> GetOnDeviceSizeInBytes() const override;
+  absl::StatusOr<size_t> GetOnDeviceSizeInBytes() const override;
 
   PjRtFuture<> CopyRawToHost(void* dst, int64_t offset,
                              int64_t transfer_size) override {
@@ -538,18 +539,18 @@ class PjRtCApiBuffer : public PjRtBuffer {
 
   void Delete() override;
 
-  StatusOr<std::unique_ptr<ExternalReference>> ReleaseDeviceMemoryOwnership(
-      bool wait_for_operations_to_complete) override {
+  absl::StatusOr<std::unique_ptr<ExternalReference>>
+  ReleaseDeviceMemoryOwnership(bool wait_for_operations_to_complete) override {
     return Unimplemented(
         "PJRT C API does not support ReleaseDeviceMemoryOwnership");
   }
 
   bool IsDeleted() override;
 
-  StatusOr<std::unique_ptr<PjRtBuffer>> CopyToDevice(
+  absl::StatusOr<std::unique_ptr<PjRtBuffer>> CopyToDevice(
       PjRtDevice* dst_device) override;
 
-  StatusOr<std::unique_ptr<PjRtBuffer>> CopyToMemorySpace(
+  absl::StatusOr<std::unique_ptr<PjRtBuffer>> CopyToMemorySpace(
       PjRtMemorySpace* dst_memory_space) override;
 
   void CopyToRemoteDevice(PjRtFuture<std::string> serialized_descriptor,
@@ -612,7 +613,7 @@ class PjRtCApiExternalReference : public PjRtBuffer::ExternalReference {
   }
   ~PjRtCApiExternalReference() override;
 
-  Status WaitUntilBufferReadyOnStream(std::intptr_t stream) override;
+  absl::Status WaitUntilBufferReadyOnStream(std::intptr_t stream) override;
 
  private:
   PjRtCApiClient* client_;
@@ -629,37 +630,37 @@ class PjRtCApiExecutable : public PjRtExecutable {
 
   int64_t SizeOfGeneratedCodeInBytes() const override;
 
-  StatusOr<absl::flat_hash_map<std::string, PjRtValueType>> GetCostAnalysis()
-      const override;
+  absl::StatusOr<absl::flat_hash_map<std::string, PjRtValueType>>
+  GetCostAnalysis() const override;
 
-  StatusOr<std::vector<std::shared_ptr<HloModule>>> GetHloModules()
+  absl::StatusOr<std::vector<std::shared_ptr<HloModule>>> GetHloModules()
       const override;
 
-  StatusOr<CompiledMemoryStats> GetCompiledMemoryStats() const override {
+  absl::StatusOr<CompiledMemoryStats> GetCompiledMemoryStats() const override {
     return pjrt::GetCompiledMemoryStats(c_api_, executable_.get());
   }
 
-  StatusOr<std::vector<Shape>> GetOutputShapes() const override {
+  absl::StatusOr<std::vector<Shape>> GetOutputShapes() const override {
     LOG(FATAL) << "PjRtExecutable::GetOutputShapes() not implemented in PJRT C "
                   "API. Please use PjRtExecutable::GetOutputElementTypes() or "
                   "PjRtExecutable::GetOutputDimensions().";
   }
 
-  StatusOr<std::vector<std::vector<PrimitiveType>>> GetOutputElementTypes()
-      const override;
+  absl::StatusOr<std::vector<std::vector<PrimitiveType>>>
+  GetOutputElementTypes() const override;
 
-  StatusOr<std::vector<std::vector<DimensionVector>>> GetOutputDimensions()
-      const override;
+  absl::StatusOr<std::vector<std::vector<DimensionVector>>>
+  GetOutputDimensions() const override;
 
-  StatusOr<std::vector<std::vector<absl::string_view>>> GetOutputMemoryKinds()
-      const override;
+  absl::StatusOr<std::vector<std::vector<absl::string_view>>>
+  GetOutputMemoryKinds() const override;
 
   const PJRT_Api* pjrt_c_api() const { return c_api_; }
   PJRT_Executable* c_executable() const { return executable_.get(); }
 
-  StatusOr<std::string> SerializeExecutable() const override;
+  absl::StatusOr<std::string> SerializeExecutable() const override;
 
-  StatusOr<std::string> FingerprintExecutable() const override;
+  absl::StatusOr<std::string> FingerprintExecutable() const override;
 
  private:
   const PJRT_Api* c_api_;
@@ -680,8 +681,8 @@ class PjRtCApiLoadedExecutable : public PjRtLoadedExecutable {
     return executable_->SizeOfGeneratedCodeInBytes();
   }
 
-  StatusOr<absl::flat_hash_map<std::string, PjRtValueType>> GetCostAnalysis()
-      const override {
+  absl::StatusOr<absl::flat_hash_map<std::string, PjRtValueType>>
+  GetCostAnalysis() const override {
     return executable_->GetCostAnalysis();
   }
 
@@ -699,48 +700,48 @@ class PjRtCApiLoadedExecutable : public PjRtLoadedExecutable {
     return addressable_devices_;
   }
 
-  StatusOr<std::vector<std::shared_ptr<HloModule>>> GetHloModules()
+  absl::StatusOr<std::vector<std::shared_ptr<HloModule>>> GetHloModules()
       const override {
     return executable_->GetHloModules();
   }
 
-  StatusOr<CompiledMemoryStats> GetCompiledMemoryStats() const override {
+  absl::StatusOr<CompiledMemoryStats> GetCompiledMemoryStats() const override {
     return executable_->GetCompiledMemoryStats();
   }
 
-  StatusOr<std::vector<Shape>> GetOutputShapes() const override {
+  absl::StatusOr<std::vector<Shape>> GetOutputShapes() const override {
     LOG(FATAL)
         << "PjRtLoadedExecutable::GetOutputShapes() not implemented in PJRT C "
            "API. Please use PjRtLoadedExecutable::GetOutputElementTypes() or "
            "PjRtLoadedExecutable::GetOutputDimensions().";
   }
 
-  StatusOr<std::vector<std::vector<PrimitiveType>>> GetOutputElementTypes()
-      const override {
+  absl::StatusOr<std::vector<std::vector<PrimitiveType>>>
+  GetOutputElementTypes() const override {
     return executable_->GetOutputElementTypes();
   }
 
-  StatusOr<std::vector<std::vector<DimensionVector>>> GetOutputDimensions()
-      const override {
+  absl::StatusOr<std::vector<std::vector<DimensionVector>>>
+  GetOutputDimensions() const override {
     return executable_->GetOutputDimensions();
   }
 
-  StatusOr<std::vector<std::vector<absl::string_view>>> GetOutputMemoryKinds()
-      const override {
+  absl::StatusOr<std::vector<std::vector<absl::string_view>>>
+  GetOutputMemoryKinds() const override {
     return executable_->GetOutputMemoryKinds();
   }
 
-  StatusOr<std::vector<std::vector<std::unique_ptr<PjRtBuffer>>>> Execute(
+  absl::StatusOr<std::vector<std::vector<std::unique_ptr<PjRtBuffer>>>> Execute(
       absl::Span<const std::vector<PjRtBuffer*>> argument_handles,
       const ExecuteOptions& options,
       std::optional<std::vector<PjRtFuture<>>>& returned_futures) override;
 
-  StatusOr<std::vector<std::unique_ptr<PjRtBuffer>>> ExecuteSharded(
+  absl::StatusOr<std::vector<std::unique_ptr<PjRtBuffer>>> ExecuteSharded(
       absl::Span<PjRtBuffer* const> argument_handles, PjRtDevice* device,
       const ExecuteOptions& options,
       std::optional<PjRtFuture<>>& returned_future, bool fill_future) override;
 
-  StatusOr<std::vector<std::unique_ptr<PjRtBuffer>>> ExecutePortable(
+  absl::StatusOr<std::vector<std::unique_ptr<PjRtBuffer>>> ExecutePortable(
       absl::Span<PjRtBuffer* const> argument_handles, PjRtDevice* device,
       const ExecuteOptions& options,
       std::optional<PjRtFuture<>>& returned_future, bool fill_future) override;
@@ -748,7 +749,7 @@ class PjRtCApiLoadedExecutable : public PjRtLoadedExecutable {
   void Delete() override;
   bool IsDeleted() override;
 
-  StatusOr<std::string> SerializeExecutable() const override {
+  absl::StatusOr<std::string> SerializeExecutable() const override {
     return executable_->SerializeExecutable();
   }
 
@@ -771,7 +772,7 @@ class PjRtCApiLoadedExecutable : public PjRtLoadedExecutable {
 
   // Override to call FingerprintExecutable through the wrapped
   // PjRtCApiExecutable.
-  StatusOr<std::string> FingerprintExecutable() const override;
+  absl::StatusOr<std::string> FingerprintExecutable() const override;
 
  private:
   // Groups data needed to support send/recv execution callbacks.
@@ -798,10 +799,11 @@ class PjRtCApiLoadedExecutable : public PjRtLoadedExecutable {
       SendRecvCallbackData& send_recv_callback_data,
       std::vector<int64_t>& non_donatable_input_indices_storage);
 
-  StatusOr<std::vector<std::unique_ptr<PjRtBuffer>>> ExecuteWithSingleDevice(
-      absl::Span<PjRtBuffer* const> argument_handles, PjRtDevice* device,
-      const ExecuteOptions& options,
-      std::optional<PjRtFuture<>>& returned_future, bool fill_future);
+  absl::StatusOr<std::vector<std::unique_ptr<PjRtBuffer>>>
+  ExecuteWithSingleDevice(absl::Span<PjRtBuffer* const> argument_handles,
+                          PjRtDevice* device, const ExecuteOptions& options,
+                          std::optional<PjRtFuture<>>& returned_future,
+                          bool fill_future);
 
   PjRtCApiClient* client_;
   std::unique_ptr<PJRT_LoadedExecutable, ::pjrt::PJRT_LoadedExecutableDeleter>
@@ -825,7 +827,7 @@ class CApiCopyToDeviceStream : public CopyToDeviceStream {
   const PJRT_Api* c_api_;
 };
 
-StatusOr<std::unique_ptr<PjRtClient>> GetCApiClient(
+absl::StatusOr<std::unique_ptr<PjRtClient>> GetCApiClient(
     absl::string_view device_type,
     const absl::flat_hash_map<std::string, PjRtValueType>& create_options = {},
     std::shared_ptr<KeyValueStoreInterface> kv_store = nullptr);
@@ -838,7 +840,7 @@ absl::StatusOr<std::unique_ptr<PjRtTopologyDescription>> GetCApiTopology(
 // registered with standard way (xla_bridge.register_plugin).
 // TODO(b/322357665): Delete this method after TPU plugin changes to use the
 // standard registration.
-StatusOr<std::unique_ptr<PjRtTopologyDescription>> GetCApiTopology(
+absl::StatusOr<std::unique_ptr<PjRtTopologyDescription>> GetCApiTopology(
     absl::string_view device_type, absl::string_view topology_name,
     const absl::flat_hash_map<std::string, PjRtValueType>& create_options = {});
 
diff --git a/third_party/xla/xla/pjrt/pjrt_client.cc b/third_party/xla/xla/pjrt/pjrt_client.cc
index 32bf9e91cd0e5c..ea3534ff8f241b 100644
--- a/third_party/xla/xla/pjrt/pjrt_client.cc
+++ b/third_party/xla/xla/pjrt/pjrt_client.cc
@@ -31,7 +31,8 @@ namespace xla {
 
 PjRtBuffer::ExternalReference::~ExternalReference() = default;
 
-StatusOr<std::uintptr_t> PjRtClient::UnsafeBufferPointer(PjRtBuffer* buffer) {
+absl::StatusOr<std::uintptr_t> PjRtClient::UnsafeBufferPointer(
+    PjRtBuffer* buffer) {
   if (buffer->on_device_shape().IsTuple()) {
     return Unimplemented(
         "unsafe_buffer_pointer is not implemented for tuple buffers.");
@@ -78,7 +79,7 @@ PjRtHostMemoryForDeviceManager::~PjRtHostMemoryForDeviceManager() = default;
 
 CopyToDeviceStream::~CopyToDeviceStream() = default;
 
-StatusOr<absl::flat_hash_map<std::string, PjRtValueType>>
+absl::StatusOr<absl::flat_hash_map<std::string, PjRtValueType>>
 PjRtLoadedExecutable::GetCostAnalysis() const {
   TF_ASSIGN_OR_RETURN(std::unique_ptr<HloCostAnalysis> hlo_cost_analysis,
                       client()->GetHloCostAnalysis());
diff --git a/third_party/xla/xla/pjrt/pjrt_client.h b/third_party/xla/xla/pjrt/pjrt_client.h
index b287ad25755923..c5a2a0b1e2ff76 100644
--- a/third_party/xla/xla/pjrt/pjrt_client.h
+++ b/third_party/xla/xla/pjrt/pjrt_client.h
@@ -209,15 +209,15 @@ class PjRtDevice {
       absl::string_view description) const = 0;
 
   // Transfer the given literal to the infeed queue.
-  virtual Status TransferToInfeed(const LiteralSlice& literal) = 0;
+  virtual absl::Status TransferToInfeed(const LiteralSlice& literal) = 0;
 
   // Transfer and return a value of the given shape from the outfeed queue.
-  virtual Status TransferFromOutfeed(MutableBorrowingLiteral literal) = 0;
+  virtual absl::Status TransferFromOutfeed(MutableBorrowingLiteral literal) = 0;
 
   // Returns allocator stats for the device. Only some PjRtDevice
   // implementations support allocator_stats, and those that do not will return
   // an Unimplemented error.
-  virtual StatusOr<tsl::AllocatorStats> GetAllocatorStats() const {
+  virtual absl::StatusOr<tsl::AllocatorStats> GetAllocatorStats() const {
     return Unimplemented("GetAllocatorStats is not supported");
   }
 
@@ -226,7 +226,7 @@ class PjRtDevice {
   virtual absl::Span<PjRtMemorySpace* const> memory_spaces() const = 0;
 
   // Returns the default memory space attached to this device.
-  virtual StatusOr<PjRtMemorySpace*> default_memory_space() const = 0;
+  virtual absl::StatusOr<PjRtMemorySpace*> default_memory_space() const = 0;
 
   virtual absl::StatusOr<PjRtMemorySpace*> memory_space_by_kind(
       absl::string_view memory_space_kind) const {
@@ -237,7 +237,8 @@ class PjRtDevice {
   // an externally-managed buffer is ready to use on this device. This is
   // intended to support dlpack on GPU and is not expected to be implemented for
   // all hardware platforms.
-  virtual StatusOr<std::intptr_t> GetStreamForExternalReadyEvents() const {
+  virtual absl::StatusOr<std::intptr_t> GetStreamForExternalReadyEvents()
+      const {
     return Unimplemented(
         "PjRtDevice::GetStreamForExternalReadyEvents only implemented for "
         "GPU");
@@ -255,7 +256,8 @@ class PjRtDevice {
   //
   // Returns error otherwise, including in the case that poisoning is not
   // implemented by this client.
-  virtual StatusOr<bool> PoisonExecution(int32_t launch_id, Status error) {
+  virtual absl::StatusOr<bool> PoisonExecution(int32_t launch_id,
+                                               absl::Status error) {
     return Unimplemented("PoisonExecution is not supported");
   }
 };
@@ -286,9 +288,9 @@ struct PjRtCrossHostRecvDescriptors {
 // duplicate (e.g., both send and cancel) then the system will be left in an
 // undefined state. If there is no send or cancellation then the system will
 // hang indefinitely.
-using PjRtCrossHostSendCancelNotifier =
-    std::function<void(absl::string_view serialized_descriptor, Status reason,
-                       std::function<void(Status)> on_canceled)>;
+using PjRtCrossHostSendCancelNotifier = std::function<void(
+    absl::string_view serialized_descriptor, absl::Status reason,
+    std::function<void(Status)> on_canceled)>;
 // State asynchronously returned by MakeCrossHostReceiveBuffers. "descriptors"
 // will match the returned PjRtBuffer objects 1:1. Specifically, each PjRtBuffer
 // returned by MakeCrossHostReceiveBuffers will have one
@@ -426,17 +428,17 @@ class PjRtHostMemoryForDeviceManager {
   // Transforms the host memory representations of a shape with the host layout
   // to the host memory representation of the same shape with the device layout.
   // `src_shape` and `dst_shape` may only differ in their layouts.
-  virtual StatusOr<PjRtChunk> ToDeviceLayout(const void* src_data,
-                                             size_t src_size,
-                                             const Shape& host_shape,
-                                             const Shape& device_shape) = 0;
+  virtual absl::StatusOr<PjRtChunk> ToDeviceLayout(
+      const void* src_data, size_t src_size, const Shape& host_shape,
+      const Shape& device_shape) = 0;
 
   // Transforms the host memory representations of a shape with the device
   // layout to the host memory representation of the same shape with the host
   // layout. `src_shape` and `dst_shape` may only differ in their layouts.
-  virtual Status ToHostLayout(const void* src_data, size_t src_size,
-                              const Shape& src_shape, void* dst_data,
-                              size_t dst_size, const Shape& dst_shape) = 0;
+  virtual absl::Status ToHostLayout(const void* src_data, size_t src_size,
+                                    const Shape& src_shape, void* dst_data,
+                                    size_t dst_size,
+                                    const Shape& dst_shape) = 0;
 };
 
 class PjRtLoadedExecutable;
@@ -525,17 +527,17 @@ class PjRtClient {
   virtual absl::Span<PjRtDevice* const> addressable_devices() const = 0;
 
   // Lookup any PjRtDevice for a given PjRtDevice::id().
-  virtual StatusOr<PjRtDevice*> LookupDevice(
+  virtual absl::StatusOr<PjRtDevice*> LookupDevice(
       PjRtGlobalDeviceId global_device_id) const = 0;
 
   // Return an addressable PjRtDevice for a given
   // PjRtDevice::local_hardware_id().
   ABSL_DEPRECATED("Use LookupAddressableDevice(PjRtLocalDeviceId) instead")
-  virtual StatusOr<PjRtDevice*> LookupAddressableDevice(
+  virtual absl::StatusOr<PjRtDevice*> LookupAddressableDevice(
       int local_hardware_id) const {
     return LookupAddressableDevice(PjRtLocalDeviceId(local_hardware_id));
   }
-  virtual StatusOr<PjRtDevice*> LookupAddressableDevice(
+  virtual absl::StatusOr<PjRtDevice*> LookupAddressableDevice(
       PjRtLocalDeviceId local_device_id) const = 0;
 
   // Return all memory spaces owned by the client.
@@ -565,7 +567,7 @@ class PjRtClient {
 
   // Return a device-specific default device assignment, e.g., GPU and TPU may
   // be different.
-  virtual StatusOr<DeviceAssignment> GetDefaultDeviceAssignment(
+  virtual absl::StatusOr<DeviceAssignment> GetDefaultDeviceAssignment(
       int num_replicas, int num_partitions) const = 0;
 
   // Returns a device-specific default device assignment for multi-slice system.
@@ -574,7 +576,7 @@ class PjRtClient {
   // communication happens across replicas assuming then that
   // num_replicas_per_slice is going to be "num_replicas / num_slices".
   // TODO(zhangqiaorjc): Convert this to pure virtual and push down.
-  virtual StatusOr<DeviceAssignment> GetDefaultDeviceAssignment(
+  virtual absl::StatusOr<DeviceAssignment> GetDefaultDeviceAssignment(
       int num_replicas, std::optional<int> num_replicas_per_slice,
       int num_partitions, const MultiSliceConfig* multi_slice_config) const {
     return Unimplemented("Multi slice device assignment is not supported.");
@@ -586,19 +588,19 @@ class PjRtClient {
   // default layout is used for program arguments and outputs unless
   // user-specified or compiler-chosen layouts are requested via the
   // "mhlo.layout_mode" attribute.
-  virtual StatusOr<Layout> GetDefaultLayout(PrimitiveType element_type,
-                                            absl::Span<const int64_t> dims) = 0;
+  virtual absl::StatusOr<Layout> GetDefaultLayout(
+      PrimitiveType element_type, absl::Span<const int64_t> dims) = 0;
 
   // Returns a backend-specific HLO cost analysis visitor.
-  virtual StatusOr<std::unique_ptr<HloCostAnalysis>> GetHloCostAnalysis()
+  virtual absl::StatusOr<std::unique_ptr<HloCostAnalysis>> GetHloCostAnalysis()
       const = 0;
 
   // Compile `computation` with given `options`.
-  virtual StatusOr<std::unique_ptr<PjRtLoadedExecutable>> Compile(
+  virtual absl::StatusOr<std::unique_ptr<PjRtLoadedExecutable>> Compile(
       const XlaComputation& computation, CompileOptions options) = 0;
 
   // Variant of `Compile` that accepts an MLIR module.
-  virtual StatusOr<std::unique_ptr<PjRtLoadedExecutable>> Compile(
+  virtual absl::StatusOr<std::unique_ptr<PjRtLoadedExecutable>> Compile(
       mlir::ModuleOp module, CompileOptions options) = 0;
 
   // Deserializes a serialized executable as produced by
@@ -608,13 +610,14 @@ class PjRtClient {
   // Pending completion of b/237720161, `options` is a mandatory argument in
   // most implementations of this interface. They _are_ optional for
   // implementations related to the PJRT C API.
-  virtual StatusOr<std::unique_ptr<PjRtLoadedExecutable>> DeserializeExecutable(
-      absl::string_view serialized, std::optional<CompileOptions> options) = 0;
+  virtual absl::StatusOr<std::unique_ptr<PjRtLoadedExecutable>>
+  DeserializeExecutable(absl::string_view serialized,
+                        std::optional<CompileOptions> options) = 0;
 
   // LoadSerializedExecutable takes the serialized output of PjRtExecutable. The
   // returned executable is loaded by this client. The same checks are made as
   // in Load that the serialized executable is compatible with the client.
-  virtual StatusOr<std::unique_ptr<PjRtLoadedExecutable>>
+  virtual absl::StatusOr<std::unique_ptr<PjRtLoadedExecutable>>
   LoadSerializedExecutable(absl::string_view serialized,
                            std::optional<CompileOptions> options,
                            const LoadOptions& load_options) {
@@ -627,28 +630,28 @@ class PjRtClient {
   // PjRtExecutable contains a copy of the CompileOptions that was used to
   // generate the executable. Load will use the CompileOptions from within the
   // executable.
-  virtual StatusOr<std::unique_ptr<PjRtLoadedExecutable>> Load(
+  virtual absl::StatusOr<std::unique_ptr<PjRtLoadedExecutable>> Load(
       std::unique_ptr<PjRtExecutable> executable,
       const LoadOptions& load_options) {
     return Unimplemented("Loading executable not supported.");
   }
 
   // Creates a buffer on the device without initializing or copying any data.
-  virtual StatusOr<std::unique_ptr<PjRtBuffer>> CreateUninitializedBuffer(
+  virtual absl::StatusOr<std::unique_ptr<PjRtBuffer>> CreateUninitializedBuffer(
       const Shape& shape, PjRtDevice* device) = 0;
 
   // Creates buffer in the given memory space that carries an error future
   // without allocating memory.
-  virtual StatusOr<std::unique_ptr<PjRtBuffer>> CreateErrorBuffer(
-      Status error, const Shape& shape, PjRtMemorySpace* memory) {
+  virtual absl::StatusOr<std::unique_ptr<PjRtBuffer>> CreateErrorBuffer(
+      absl::Status error, const Shape& shape, PjRtMemorySpace* memory) {
     return Unimplemented("CreateErrorBuffer not supported.");
   }
 
   // Creates buffer in the given device that carries an error future without
   // allocating memory.
   ABSL_DEPRECATED("Use CreateErrorBuffer(Status, Shape, PjRtMemorySpace*)")
-  virtual StatusOr<std::unique_ptr<PjRtBuffer>> CreateErrorBuffer(
-      Status error, const Shape& shape, PjRtDevice* device) {
+  virtual absl::StatusOr<std::unique_ptr<PjRtBuffer>> CreateErrorBuffer(
+      absl::Status error, const Shape& shape, PjRtDevice* device) {
     auto default_memory_space = device->default_memory_space();
     if (!default_memory_space.ok()) {
       return default_memory_space.status();
@@ -657,8 +660,8 @@ class PjRtClient {
   }
 
   // Gets the pointer to the topology description held by the client.
-  virtual StatusOr<const PjRtTopologyDescription*> GetTopologyDescription()
-      const {
+  virtual absl::StatusOr<const PjRtTopologyDescription*>
+  GetTopologyDescription() const {
     return Unimplemented("GetTopologyDescription not supported on platform %s",
                          platform_name());
   }
@@ -708,7 +711,7 @@ class PjRtClient {
     // transfer is complete but before the buffers are made available to
     // their consumers. 'literal' must remain in scope until on_done is
     // called.
-    virtual Status TransferLiteralToBuffer(
+    virtual absl::Status TransferLiteralToBuffer(
         int buffer_index, const LiteralSlice& literal,
         absl::AnyInvocable<void() &&> on_done) = 0;
 
@@ -721,7 +724,7 @@ class PjRtClient {
     // buffer_index can be made after this call. on_done is called when the
     // transfer is complete but before the buffers are made available to their
     // consumers. 'data' must remain in scope until on_done is called.
-    virtual Status TransferRawDataToBuffer(
+    virtual absl::Status TransferRawDataToBuffer(
         int buffer_index, absl::string_view data,
         absl::AnyInvocable<void() &&> on_done) = 0;
 
@@ -735,7 +738,7 @@ class PjRtClient {
     // calls) into buffer_index can be made after this call. on_done is called
     // when the transfer is complete but before the buffers are made available
     // to their consumers. 'data' must remain in scope until on_done is called.
-    virtual Status TransferRawDataToSubBuffer(
+    virtual absl::Status TransferRawDataToSubBuffer(
         int buffer_index, const void* data, int64_t offset,
         int64_t transfer_size, bool is_last_transfer,
         absl::AnyInvocable<void() &&> on_done) = 0;
@@ -743,7 +746,7 @@ class PjRtClient {
     // Indicates that a specific buffer should result in an error status. No
     // transfer calls (or further SetBufferError calls) into buffer_index can
     // be made after this call.
-    virtual void SetBufferError(int buffer_index, Status error) = 0;
+    virtual void SetBufferError(int buffer_index, absl::Status error) = 0;
 
     // Adds the specified key/value metadata for the transfer operation.
     // This is typically used for debugging purposes, such as adding a handle
@@ -754,12 +757,12 @@ class PjRtClient {
 
   // Returns a manager for async transfers into a set of buffers with on-host
   // shapes 'shapes'.
-  virtual StatusOr<std::unique_ptr<AsyncHostToDeviceTransferManager>>
+  virtual absl::StatusOr<std::unique_ptr<AsyncHostToDeviceTransferManager>>
   CreateBuffersForAsyncHostToDevice(absl::Span<const Shape> shapes,
                                     PjRtDevice* device) = 0;
 
   // Variant of CreateBuffersForAsyncHostToDevice with PjRtMemorySpace.
-  virtual StatusOr<std::unique_ptr<AsyncHostToDeviceTransferManager>>
+  virtual absl::StatusOr<std::unique_ptr<AsyncHostToDeviceTransferManager>>
   CreateBuffersForAsyncHostToDevice(absl::Span<const Shape> shapes,
                                     PjRtMemorySpace* memory_space) = 0;
 
@@ -778,23 +781,23 @@ class PjRtClient {
 
     // Transfers data to the device buffer. Data should already be in the
     // device layout.
-    virtual Status TransferRawDataToSubBuffer(
+    virtual absl::Status TransferRawDataToSubBuffer(
         const void* data, int64_t offset, int64_t transfer_size,
         bool is_last_transfer, absl::AnyInvocable<void() &&> on_done) = 0;
 
     // The resulting buffer becomes ready when all transfers complete.
-    virtual StatusOr<std::unique_ptr<PjRtBuffer>> Slice(
+    virtual absl::StatusOr<std::unique_ptr<PjRtBuffer>> Slice(
         int64_t offset, PrimitiveType type, absl::Span<int64_t const> dims,
         const Layout& layout) = 0;
   };
   // Creates a raw device buffer of a given size in bytes.
-  virtual StatusOr<std::unique_ptr<PjRtRawDeviceBuffer>> CreateRawDeviceBuffer(
-      int64_t size, PjRtDevice* device) {
+  virtual absl::StatusOr<std::unique_ptr<PjRtRawDeviceBuffer>>
+  CreateRawDeviceBuffer(int64_t size, PjRtDevice* device) {
     return Unimplemented("CreateRawDeviceBuffer is not implemented.");
   }
 
   // On-device bytes required for a PjRt buffer with these `Shape` attributes.
-  virtual StatusOr<int64_t> GetOnDeviceSizeInBytes(
+  virtual absl::StatusOr<int64_t> GetOnDeviceSizeInBytes(
       PrimitiveType type, absl::Span<int64_t const> dims,
       const Layout& layout) {
     return Unimplemented("GetOnDeviceSizeInBytes is not implemented.");
@@ -849,7 +852,7 @@ class PjRtClient {
   //
   // If byte_strides is omitted, the array is assumed to have a dense layout
   // with dimensions in major-to-minor order.
-  virtual StatusOr<std::unique_ptr<PjRtBuffer>> BufferFromHostBuffer(
+  virtual absl::StatusOr<std::unique_ptr<PjRtBuffer>> BufferFromHostBuffer(
       const void* data, PrimitiveType type, absl::Span<int64_t const> dims,
       std::optional<absl::Span<int64_t const>> byte_strides,
       HostBufferSemantics host_buffer_semantics,
@@ -860,7 +863,7 @@ class PjRtClient {
   // used when non-compact layout is preferred.
   // TODO(b/275645543): remove BufferFromHostBuffer without optional device
   // layout after all the inherited classes and call sites are updated.
-  virtual StatusOr<std::unique_ptr<PjRtBuffer>> BufferFromHostBuffer(
+  virtual absl::StatusOr<std::unique_ptr<PjRtBuffer>> BufferFromHostBuffer(
       const void* data, PrimitiveType type, absl::Span<int64_t const> dims,
       std::optional<absl::Span<int64_t const>> byte_strides,
       HostBufferSemantics host_buffer_semantics,
@@ -874,7 +877,7 @@ class PjRtClient {
 
   // TODO(b/277820585): remove BufferFromHostBuffer with PjRtDevice after the
   // migration is done.
-  virtual StatusOr<std::unique_ptr<PjRtBuffer>> BufferFromHostBuffer(
+  virtual absl::StatusOr<std::unique_ptr<PjRtBuffer>> BufferFromHostBuffer(
       const void* data, PrimitiveType type, absl::Span<int64_t const> dims,
       std::optional<absl::Span<int64_t const>> byte_strides,
       HostBufferSemantics host_buffer_semantics,
@@ -889,12 +892,12 @@ class PjRtClient {
   // Note that literal must remain in scope until the transfer has completed, so
   // the caller should, for example, wait for GetReadyFuture().Await()
   // completes on the return value before letting literal go out of scope.
-  virtual StatusOr<std::unique_ptr<PjRtBuffer>> BufferFromHostLiteral(
+  virtual absl::StatusOr<std::unique_ptr<PjRtBuffer>> BufferFromHostLiteral(
       const LiteralSlice& literal, PjRtDevice* device) = 0;
 
   // TODO(b/277820585): remove BufferFromHostLiteral with PjRtDevice after the
   // migration is done.
-  virtual StatusOr<std::unique_ptr<PjRtBuffer>> BufferFromHostLiteral(
+  virtual absl::StatusOr<std::unique_ptr<PjRtBuffer>> BufferFromHostLiteral(
       const LiteralSlice& literal, PjRtMemorySpace* memory_space) {
     return tsl::errors::Unimplemented(
         "BufferFromHostLiteral with PjRtMemorySpace is not implemented on "
@@ -914,14 +917,15 @@ class PjRtClient {
   // indicates when the returned buffer is ready to use. This is intended to
   // support dlpack on GPU and is not expected to be supported on all hardware
   // platforms.
-  virtual StatusOr<std::unique_ptr<PjRtBuffer>> CreateViewOfDeviceBuffer(
+  virtual absl::StatusOr<std::unique_ptr<PjRtBuffer>> CreateViewOfDeviceBuffer(
       void* device_ptr, const Shape& shape, PjRtDevice* device,
       std::function<void()> on_delete_callback,
       std::optional<std::intptr_t> stream = std::nullopt) = 0;
 
   // Returns platform-dependent address for the given buffer that is often but
   // not guaranteed to be the physical/device address.
-  virtual StatusOr<std::uintptr_t> UnsafeBufferPointer(PjRtBuffer* buffer);
+  virtual absl::StatusOr<std::uintptr_t> UnsafeBufferPointer(
+      PjRtBuffer* buffer);
 
   // Returns a vector of PjRtBuffers that can be used to receive
   // cross host transfers using `client` on `device'. Asynchronously calls
@@ -941,7 +945,7 @@ class PjRtClient {
   //
   // See note on semantics of cross-device copies in the class definition
   // comment for PjRtClient.
-  virtual StatusOr<std::vector<std::unique_ptr<PjRtBuffer>>>
+  virtual absl::StatusOr<std::vector<std::unique_ptr<PjRtBuffer>>>
   MakeCrossHostReceiveBuffers(absl::Span<const Shape> shapes,
                               PjRtDevice* device,
                               PjRtCrossHostRecvNotifier notifier) = 0;
@@ -975,19 +979,19 @@ class PjRtClient {
     // completes.
     std::vector<int64_t> slice_boundaries;
   };
-  virtual StatusOr<std::vector<std::unique_ptr<PjRtBuffer>>>
+  virtual absl::StatusOr<std::vector<std::unique_ptr<PjRtBuffer>>>
   MakeCrossHostReceiveBuffersForGather(
       absl::Span<const Shape> shapes, std::vector<GatherDetails> gather_details,
       PjRtDevice* device, PjRtCrossHostRecvNotifier notifier) = 0;
 
   // Create ChannelHandles for XLA send/recv.
-  virtual StatusOr<ChannelHandle> CreateChannelHandle() = 0;
-  virtual StatusOr<ChannelHandle> CreateDeviceToHostChannelHandle() = 0;
-  virtual StatusOr<ChannelHandle> CreateHostToDeviceChannelHandle() = 0;
+  virtual absl::StatusOr<ChannelHandle> CreateChannelHandle() = 0;
+  virtual absl::StatusOr<ChannelHandle> CreateDeviceToHostChannelHandle() = 0;
+  virtual absl::StatusOr<ChannelHandle> CreateHostToDeviceChannelHandle() = 0;
 
   // TODO(zhangqiaorjc): Experimental API to be removed.
   // Defragment device memory.
-  virtual Status Defragment() = 0;
+  virtual absl::Status Defragment() = 0;
 
   // If false, this client does not support send/recv host callbacks, and
   // callers should not set the `send_callbacks` and `recv_callbacks` arguments
@@ -1057,7 +1061,7 @@ class PjRtBuffer {
   //
   // Since this method actually acquires locks and communicate with the device,
   // it does not have the const qualifier, similar to what ToLiteral does.
-  virtual StatusOr<std::vector<int64_t>> logical_dimensions() {
+  virtual absl::StatusOr<std::vector<int64_t>> logical_dimensions() {
     TF_ASSIGN_OR_RETURN(Shape logical_shape, logical_on_device_shape());
     absl::Span<const int64_t> dims = logical_shape.dimensions();
     return std::vector<int64_t>(dims.begin(), dims.end());
@@ -1070,7 +1074,7 @@ class PjRtBuffer {
   //
   // Since this method actually acquires locks and communicate with the device,
   // it does not have the const qualifier, similar to what ToLiteral does.
-  virtual StatusOr<Shape> logical_on_device_shape() {
+  virtual absl::StatusOr<Shape> logical_on_device_shape() {
     const Shape& shape = on_device_shape();
     CHECK(shape.is_static())
         << "logical_on_device_shape needs to be overridden for platform '"
@@ -1098,7 +1102,7 @@ class PjRtBuffer {
 
     // Stream is platform-specific. This is intended to support dlpack on GPU
     // and is not expected to be implemented for all hardware platforms.
-    virtual Status WaitUntilBufferReadyOnStream(std::intptr_t stream) {
+    virtual absl::Status WaitUntilBufferReadyOnStream(std::intptr_t stream) {
       return Unimplemented(
           "WaitUntilBufferReadyOnStream is only implemented for GPU.");
     }
@@ -1106,7 +1110,7 @@ class PjRtBuffer {
    protected:
     void* data_ptr_;
   };
-  virtual StatusOr<std::unique_ptr<ExternalReference>>
+  virtual absl::StatusOr<std::unique_ptr<ExternalReference>>
   AcquireExternalReference() = 0;
 
   // Asynchronously copies the buffer's value into `literal`.
@@ -1125,10 +1129,10 @@ class PjRtBuffer {
           generator) = 0;
 
   // Synchronous overload of ToLiteral, as a convenience.
-  Status ToLiteralSync(MutableLiteralBase* literal) {
+  absl::Status ToLiteralSync(MutableLiteralBase* literal) {
     absl::Notification done;
-    Status status;
-    ToLiteral(literal).OnReady([&](Status s) {
+    absl::Status status;
+    ToLiteral(literal).OnReady([&](absl::Status s) {
       status = std::move(s);
       done.Notify();
     });
@@ -1174,7 +1178,7 @@ class PjRtBuffer {
   }
 
   // Returns the number of bytes of the buffer storage on the device.
-  virtual StatusOr<size_t> GetOnDeviceSizeInBytes() const = 0;
+  virtual absl::StatusOr<size_t> GetOnDeviceSizeInBytes() const = 0;
 
   // Transfers a sub-range of the on-device representation of the buffer.
   // offset+transfer_size must be less than GetOnDeviceSizeInBytes. The
@@ -1229,7 +1233,7 @@ class PjRtBuffer {
   // If the buffer was shared via an external reference it is the client's
   // responsibility that accesses via that reference do not interfere with
   // accesses via the buffer returned from ReleaseDeviceMemoryOwnership.
-  virtual StatusOr<std::unique_ptr<ExternalReference>>
+  virtual absl::StatusOr<std::unique_ptr<ExternalReference>>
   ReleaseDeviceMemoryOwnership(bool wait_for_operations_to_complete) = 0;
 
   // True if and only if Delete or Release has previously been called.
@@ -1242,7 +1246,7 @@ class PjRtBuffer {
   //
   // See note on semantics of cross-device copies in the class definition
   // comment for PjRtClient.
-  virtual StatusOr<std::unique_ptr<PjRtBuffer>> CopyToDevice(
+  virtual absl::StatusOr<std::unique_ptr<PjRtBuffer>> CopyToDevice(
       PjRtDevice* dst_device) = 0;
 
   // Copies the buffer to memory space `dst_memory_space`.
@@ -1254,7 +1258,7 @@ class PjRtBuffer {
   //
   // See note on semantics of cross-device copies in the class definition
   // comment for PjRtClient.
-  virtual StatusOr<std::unique_ptr<PjRtBuffer>> CopyToMemorySpace(
+  virtual absl::StatusOr<std::unique_ptr<PjRtBuffer>> CopyToMemorySpace(
       PjRtMemorySpace* dst_memory_space) = 0;
 
   // Prepares to send a copy of the buffer to a remote device. The destination
@@ -1287,7 +1291,7 @@ class PjRtBuffer {
   // See note on semantics of cross-device copies in the class definition
   // comment for PjRtClient.
   using RemoteSendCallback =
-      std::function<void(Status status, bool sends_were_enqueued)>;
+      std::function<void(absl::Status status, bool sends_were_enqueued)>;
   virtual void CopyToRemoteDevice(PjRtFuture<std::string> serialized_descriptor,
                                   RemoteSendCallback on_done) = 0;
   struct ScatterDetails {
@@ -1325,8 +1329,8 @@ class PjRtBuffer {
   //
   // If either 'this' or 'dependency' transitions to error, then the returned
   // buffer will transition to error.
-  virtual StatusOr<std::unique_ptr<PjRtBuffer>> DonateWithControlDependency(
-      PjRtFuture<> dependency) {
+  virtual absl::StatusOr<std::unique_ptr<PjRtBuffer>>
+  DonateWithControlDependency(PjRtFuture<> dependency) {
     return Unimplemented("DonateWithControlDependency is not supported.");
   }
 
@@ -1370,7 +1374,7 @@ class PjRtBuffer {
         std::vector<PjRtBuffer::RemoteSendCallback> callbacks,
         const ScatterDetails& scatter_details) = 0;
   };
-  virtual StatusOr<std::unique_ptr<AsyncSendPlaceholder>>
+  virtual absl::StatusOr<std::unique_ptr<AsyncSendPlaceholder>>
   CreateAsyncSendPlaceholder() {
     return Unimplemented("AsyncSendPlaceholder is not supported.");
   }
@@ -1389,7 +1393,7 @@ class PjRtBuffer {
   // Blocks the host until the buffer's value has been computed and is ready for
   // immediate use on the device. Useful in particular for timing benchmarks.
   ABSL_DEPRECATED("Use GetReadyFuture().Await() instead")
-  Status BlockHostUntilReady() {
+  absl::Status BlockHostUntilReady() {
     auto s = GetReadyFuture().Await();
     // Fix up error string because some clients rely on it.
     if (!s.ok() &&
@@ -1419,8 +1423,8 @@ class PjRtLoadedExecutable : public PjRtExecutable {
   // Returns named values for cost properties of this executable (such as
   // operations, size of input/outputs, and run time estimate). Properties may
   // differ for different platforms.
-  StatusOr<absl::flat_hash_map<std::string, PjRtValueType>> GetCostAnalysis()
-      const override;
+  absl::StatusOr<absl::flat_hash_map<std::string, PjRtValueType>>
+  GetCostAnalysis() const override;
 
   // The replica and partition indices of device_assignment to be run by this
   // client. On single-host platforms without partitioning, this is all replicas
@@ -1467,12 +1471,12 @@ class PjRtLoadedExecutable : public PjRtExecutable {
   //
   // The caller is *NOT* required to ensure that PjRtLoadedExecutable stays
   // alive until futures are ready.
-  virtual StatusOr<std::vector<std::vector<std::unique_ptr<PjRtBuffer>>>>
+  virtual absl::StatusOr<std::vector<std::vector<std::unique_ptr<PjRtBuffer>>>>
   Execute(absl::Span<const std::vector<PjRtBuffer*>> argument_handles,
           const ExecuteOptions& options,
           std::optional<std::vector<PjRtFuture<>>>& returned_futures) = 0;
   // Convenience wrapper for Execute that never returns futures.
-  StatusOr<std::vector<std::vector<std::unique_ptr<PjRtBuffer>>>> Execute(
+  absl::StatusOr<std::vector<std::vector<std::unique_ptr<PjRtBuffer>>>> Execute(
       absl::Span<const std::vector<PjRtBuffer*>> argument_handles,
       const ExecuteOptions& options) {
     std::optional<std::vector<PjRtFuture<>>> returned_futures;
@@ -1489,12 +1493,13 @@ class PjRtLoadedExecutable : public PjRtExecutable {
   //     once the execution has completed.
   //    else:
   //     returned_future will not be modified.
-  virtual StatusOr<std::vector<std::unique_ptr<PjRtBuffer>>> ExecuteSharded(
-      absl::Span<PjRtBuffer* const> argument_handles, PjRtDevice* device,
-      const ExecuteOptions& options,
-      std::optional<PjRtFuture<>>& returned_future, bool fill_future) = 0;
+  virtual absl::StatusOr<std::vector<std::unique_ptr<PjRtBuffer>>>
+  ExecuteSharded(absl::Span<PjRtBuffer* const> argument_handles,
+                 PjRtDevice* device, const ExecuteOptions& options,
+                 std::optional<PjRtFuture<>>& returned_future,
+                 bool fill_future) = 0;
   // Convenience wrapper for ExecuteSharded that always returns a future.
-  StatusOr<std::vector<std::unique_ptr<PjRtBuffer>>> ExecuteSharded(
+  absl::StatusOr<std::vector<std::unique_ptr<PjRtBuffer>>> ExecuteSharded(
       absl::Span<PjRtBuffer* const> argument_handles, PjRtDevice* device,
       const ExecuteOptions& options,
       std::optional<PjRtFuture<>>& returned_future) {
@@ -1502,7 +1507,7 @@ class PjRtLoadedExecutable : public PjRtExecutable {
                           returned_future, /*fill_future=*/true);
   }
   // Convenience wrapper for ExecuteSharded that never returns a future.
-  StatusOr<std::vector<std::unique_ptr<PjRtBuffer>>> ExecuteSharded(
+  absl::StatusOr<std::vector<std::unique_ptr<PjRtBuffer>>> ExecuteSharded(
       absl::Span<PjRtBuffer* const> argument_handles, PjRtDevice* device,
       const ExecuteOptions& options) {
     std::optional<PjRtFuture<>> returned_future;
@@ -1520,12 +1525,13 @@ class PjRtLoadedExecutable : public PjRtExecutable {
   //     once the execution has completed.
   //    else:
   //     returned_future will not be modified.
-  virtual StatusOr<std::vector<std::unique_ptr<PjRtBuffer>>> ExecutePortable(
-      absl::Span<PjRtBuffer* const> argument_handles, PjRtDevice* device,
-      const ExecuteOptions& options,
-      std::optional<PjRtFuture<>>& returned_future, bool fill_future) = 0;
+  virtual absl::StatusOr<std::vector<std::unique_ptr<PjRtBuffer>>>
+  ExecutePortable(absl::Span<PjRtBuffer* const> argument_handles,
+                  PjRtDevice* device, const ExecuteOptions& options,
+                  std::optional<PjRtFuture<>>& returned_future,
+                  bool fill_future) = 0;
   // Convenience wrapper for ExecutePortable that always returns a future.
-  StatusOr<std::vector<std::unique_ptr<PjRtBuffer>>> ExecutePortable(
+  absl::StatusOr<std::vector<std::unique_ptr<PjRtBuffer>>> ExecutePortable(
       absl::Span<PjRtBuffer* const> argument_handles, PjRtDevice* device,
       const ExecuteOptions& options,
       std::optional<PjRtFuture<>>& returned_future) {
@@ -1533,7 +1539,7 @@ class PjRtLoadedExecutable : public PjRtExecutable {
                            returned_future, /*fill_future=*/true);
   }
   // Convenience wrapper for ExecutePortable that never returns a future.
-  StatusOr<std::vector<std::unique_ptr<PjRtBuffer>>> ExecutePortable(
+  absl::StatusOr<std::vector<std::unique_ptr<PjRtBuffer>>> ExecutePortable(
       absl::Span<PjRtBuffer* const> argument_handles, PjRtDevice* device,
       const ExecuteOptions& options) {
     std::optional<PjRtFuture<>> returned_future;
diff --git a/third_party/xla/xla/pjrt/pjrt_client_test.cc b/third_party/xla/xla/pjrt/pjrt_client_test.cc
index 1bc5d6704abbf3..2f0ccf833f369d 100644
--- a/third_party/xla/xla/pjrt/pjrt_client_test.cc
+++ b/third_party/xla/xla/pjrt/pjrt_client_test.cc
@@ -64,7 +64,7 @@ TestClientFactory& GetGlobalTestClientFactory() {
   return *factory;
 }
 
-StatusOr<std::unique_ptr<PjRtClient>> GetClient() {
+absl::StatusOr<std::unique_ptr<PjRtClient>> GetClient() {
   return GetGlobalTestClientFactory().Get()();
 }
 
@@ -486,7 +486,7 @@ TEST(PjRtClientTest, CopyToDeviceAsyncExternalCpuOnly) {
   }
 }
 
-StatusOr<std::unique_ptr<PjRtBuffer>> MakeFloatBuffer(
+absl::StatusOr<std::unique_ptr<PjRtBuffer>> MakeFloatBuffer(
     PjRtClient* client, const std::vector<float>& data,
     absl::Span<const int64_t> dimensions) {
   Shape shape = ShapeUtil::MakeShape(F32, {2, 2});
diff --git a/third_party/xla/xla/pjrt/pjrt_compiler.cc b/third_party/xla/xla/pjrt/pjrt_compiler.cc
index 9501e6f5908c01..7896b9ab902795 100644
--- a/third_party/xla/xla/pjrt/pjrt_compiler.cc
+++ b/third_party/xla/xla/pjrt/pjrt_compiler.cc
@@ -72,7 +72,7 @@ void PjRtRegisterCompiler(absl::string_view platform_name,
   (*compiler_registry)[platform_name] = std::move(compiler);
 }
 
-StatusOr<std::unique_ptr<PjRtExecutable>> PjRtCompile(
+absl::StatusOr<std::unique_ptr<PjRtExecutable>> PjRtCompile(
     CompileOptions options, const XlaComputation& computation,
     const PjRtTopologyDescription& topology, PjRtClient* client) {
   auto topology_compiler = topology.compiler();
@@ -91,7 +91,7 @@ StatusOr<std::unique_ptr<PjRtExecutable>> PjRtCompile(
   return it->second->Compile(std::move(options), computation, topology, client);
 }
 
-StatusOr<std::unique_ptr<PjRtExecutable>> PjRtCompile(
+absl::StatusOr<std::unique_ptr<PjRtExecutable>> PjRtCompile(
     CompileOptions options, mlir::ModuleOp module,
     const PjRtTopologyDescription& topology, PjRtClient* client) {
   auto topology_compiler = topology.compiler();
diff --git a/third_party/xla/xla/pjrt/pjrt_compiler.h b/third_party/xla/xla/pjrt/pjrt_compiler.h
index 46c363ace13619..0c8a21e7d0976f 100644
--- a/third_party/xla/xla/pjrt/pjrt_compiler.h
+++ b/third_party/xla/xla/pjrt/pjrt_compiler.h
@@ -146,7 +146,7 @@ class PjRtTopologyDescription {
   // default layout is used for program arguments and outputs unless
   // user-specified or compiler-chosen layouts are requested via the
   // "mhlo.layout_mode" attribute.
-  virtual StatusOr<Layout> GetDefaultLayout(
+  virtual absl::StatusOr<Layout> GetDefaultLayout(
       PrimitiveType element_type, absl::Span<const int64_t> dims) const = 0;
 };
 
@@ -157,12 +157,12 @@ class PjRtCompiler {
 
   // Compiles the 'computation' and returns a 'PjRtExecutable'. The returned
   // PjRtExecutable must be loaded by a compatible client before execution.
-  virtual StatusOr<std::unique_ptr<PjRtExecutable>> Compile(
+  virtual absl::StatusOr<std::unique_ptr<PjRtExecutable>> Compile(
       CompileOptions options, const XlaComputation& computation,
       const PjRtTopologyDescription& topology, PjRtClient* client) = 0;
 
   // Variant of `Compile` that accepts an MLIR module.
-  virtual StatusOr<std::unique_ptr<PjRtExecutable>> Compile(
+  virtual absl::StatusOr<std::unique_ptr<PjRtExecutable>> Compile(
       CompileOptions options, mlir::ModuleOp module,
       const PjRtTopologyDescription& topology, PjRtClient* client) = 0;
 };
@@ -183,12 +183,12 @@ void PjRtRegisterCompiler(absl::string_view platform_name,
 // Returns error::NotFound if a compiler has not been registered for the
 // platform. Forwards errors returned from the registered compiler in case of a
 // compilation failure.
-StatusOr<std::unique_ptr<PjRtExecutable>> PjRtCompile(
+absl::StatusOr<std::unique_ptr<PjRtExecutable>> PjRtCompile(
     CompileOptions options, const XlaComputation& computation,
     const PjRtTopologyDescription& topology, PjRtClient* client = nullptr);
 
 // Variant of `PjRtCompile` that accepts an MLIR module.
-StatusOr<std::unique_ptr<PjRtExecutable>> PjRtCompile(
+absl::StatusOr<std::unique_ptr<PjRtExecutable>> PjRtCompile(
     CompileOptions options, mlir::ModuleOp module,
     const PjRtTopologyDescription& topology, PjRtClient* client = nullptr);
 
diff --git a/third_party/xla/xla/pjrt/pjrt_compiler_test.cc b/third_party/xla/xla/pjrt/pjrt_compiler_test.cc
index 98a2b8e8d5e16b..14de671708a528 100644
--- a/third_party/xla/xla/pjrt/pjrt_compiler_test.cc
+++ b/third_party/xla/xla/pjrt/pjrt_compiler_test.cc
@@ -56,7 +56,7 @@ class PjRtTestTopology : public PjRtTopologyDescription {
       const override {
     LOG(FATAL) << "Unused";
   }
-  StatusOr<Layout> GetDefaultLayout(
+  absl::StatusOr<Layout> GetDefaultLayout(
       PrimitiveType element_type,
       absl::Span<const int64_t> dims) const override {
     return Unimplemented("TestTopology does not support GetDefaultLayout");
@@ -90,7 +90,7 @@ TEST(PjRtCompilerTest, CompilerRegistered) {
         const override {
       LOG(FATAL) << "Unused";
     }
-    StatusOr<Layout> GetDefaultLayout(
+    absl::StatusOr<Layout> GetDefaultLayout(
         PrimitiveType element_type,
         absl::Span<const int64_t> dims) const override {
       return Unimplemented("TestTopology does not support GetDefaultLayout");
@@ -100,12 +100,12 @@ TEST(PjRtCompilerTest, CompilerRegistered) {
 
   class PjRtTestCompiler : public PjRtCompiler {
    public:
-    StatusOr<std::unique_ptr<PjRtExecutable>> Compile(
+    absl::StatusOr<std::unique_ptr<PjRtExecutable>> Compile(
         CompileOptions options, const XlaComputation& computation,
         const PjRtTopologyDescription& topology, PjRtClient* client) override {
       return tsl::errors::Unimplemented("test compiler!");
     }
-    StatusOr<std::unique_ptr<PjRtExecutable>> Compile(
+    absl::StatusOr<std::unique_ptr<PjRtExecutable>> Compile(
         CompileOptions options, mlir::ModuleOp module,
         const PjRtTopologyDescription& topology, PjRtClient* client) override {
       return tsl::errors::Unimplemented("test compiler!");
diff --git a/third_party/xla/xla/pjrt/pjrt_executable.cc b/third_party/xla/xla/pjrt/pjrt_executable.cc
index 4caee29a4a862c..be77360fe17196 100644
--- a/third_party/xla/xla/pjrt/pjrt_executable.cc
+++ b/third_party/xla/xla/pjrt/pjrt_executable.cc
@@ -72,7 +72,7 @@ void SetOptionOverride(OptionOverrideProto& option, double value) {
 
 }  // namespace
 
-StatusOr<CompileOptionsProto> CompileOptions::ToProto() const {
+absl::StatusOr<CompileOptionsProto> CompileOptions::ToProto() const {
   CompileOptionsProto output;
   if (argument_layouts.has_value()) {
     for (const auto& layout : *argument_layouts) {
@@ -110,7 +110,7 @@ void CompileOptions::SerializeEnvOptionOverrides(
   }
 }
 
-StatusOr<CompileOptions> CompileOptions::FromProto(
+absl::StatusOr<CompileOptions> CompileOptions::FromProto(
     const CompileOptionsProto& proto) {
   if (!proto.serialized_multi_slice_config().empty()) {
     return Unimplemented(
@@ -365,7 +365,7 @@ std::optional<std::vector<OpSharding>> PjRtExecutable::GetParameterShardings()
   return out;
 }
 
-StatusOr<std::vector<Shape>> PjRtExecutable::GetOutputShapes() const {
+absl::StatusOr<std::vector<Shape>> PjRtExecutable::GetOutputShapes() const {
   TF_ASSIGN_OR_RETURN(auto modules, GetHloModules());
   std::vector<Shape> output_shapes;
   output_shapes.reserve(modules.size());
@@ -375,7 +375,7 @@ StatusOr<std::vector<Shape>> PjRtExecutable::GetOutputShapes() const {
   return output_shapes;
 }
 
-StatusOr<std::vector<std::vector<PrimitiveType>>>
+absl::StatusOr<std::vector<std::vector<PrimitiveType>>>
 PjRtExecutable::GetOutputElementTypes() const {
   TF_ASSIGN_OR_RETURN(auto output_shapes, GetOutputShapes());
   std::vector<std::vector<PrimitiveType>> output_element_types;
@@ -403,7 +403,7 @@ PjRtExecutable::GetOutputElementTypes() const {
   return output_element_types;
 }
 
-StatusOr<std::vector<std::vector<DimensionVector>>>
+absl::StatusOr<std::vector<std::vector<DimensionVector>>>
 PjRtExecutable::GetOutputDimensions() const {
   TF_ASSIGN_OR_RETURN(auto output_shapes, GetOutputShapes());
   std::vector<std::vector<DimensionVector>> output_dimensions;
@@ -483,7 +483,7 @@ PjRtExecutable::GetOutputLayouts() const {
   return result;
 }
 
-StatusOr<absl::flat_hash_map<std::string, PjRtValueType>>
+absl::StatusOr<absl::flat_hash_map<std::string, PjRtValueType>>
 PjRtExecutableUtil::RunHloCostAnalysis(const PjRtExecutable& executable,
                                        HloCostAnalysis* hlo_cost_analysis) {
   TF_ASSIGN_OR_RETURN(std::vector<std::shared_ptr<HloModule>> modules,
@@ -502,7 +502,7 @@ PjRtExecutableUtil::RunHloCostAnalysis(const PjRtExecutable& executable,
   return RunHloCostAnalysis(modules, hlo_cost_analysis);
 }
 
-StatusOr<absl::flat_hash_map<std::string, PjRtValueType>>
+absl::StatusOr<absl::flat_hash_map<std::string, PjRtValueType>>
 PjRtExecutableUtil::RunHloCostAnalysis(
     const std::vector<std::shared_ptr<xla::HloModule>>& hlo_modules,
     HloCostAnalysis* hlo_cost_analysis) {
@@ -524,7 +524,8 @@ PjRtExecutableUtil::RunHloCostAnalysis(
   return ret;
 }
 
-StatusOr<std::vector<std::pair<std::string, CompileOptions::OptionOverride>>>
+absl::StatusOr<
+    std::vector<std::pair<std::string, CompileOptions::OptionOverride>>>
 CompileOptions::LoadEnvOptionOverrides(
     const google::protobuf::Map<std::string, xla::OptionOverrideProto>&
         env_option_overrides) {
@@ -558,8 +559,8 @@ CompileOptions::LoadEnvOptionOverrides(
   return result;
 }
 
-Status CompileOptions::ApplyOption(const std::string& key,
-                                   const OptionOverride& value) {
+absl::Status CompileOptions::ApplyOption(const std::string& key,
+                                         const OptionOverride& value) {
   if (auto* xla_field = xla::DebugOptions::descriptor()->FindFieldByName(key)) {
     xla::DebugOptions& debug_options =
         *executable_build_options.mutable_debug_options();
@@ -607,14 +608,14 @@ Status CompileOptions::ApplyOption(const std::string& key,
   }
 }
 
-Status CompileOptions::ApplyAllOptionOverrides() {
+absl::Status CompileOptions::ApplyAllOptionOverrides() {
   for (auto& option : env_option_overrides) {
     TF_RETURN_IF_ERROR(ApplyOption(option.first, option.second));
   }
   return OkStatus();
 }
 
-Status CompileOptions::ApplyOptionFromString(
+absl::Status CompileOptions::ApplyOptionFromString(
     const tsl::protobuf::FieldDescriptor* field, const std::string& value) {
   xla::DebugOptions& debug_options =
       *executable_build_options.mutable_debug_options();
diff --git a/third_party/xla/xla/pjrt/pjrt_executable.h b/third_party/xla/xla/pjrt/pjrt_executable.h
index e4005da2eee3b7..d69b27033512dd 100644
--- a/third_party/xla/xla/pjrt/pjrt_executable.h
+++ b/third_party/xla/xla/pjrt/pjrt_executable.h
@@ -111,15 +111,15 @@ struct CompileOptions {
       PrecisionConfig::DEFAULT;
 
   // Applies env_option_overrides to executable_build_options.debug_options().
-  Status ApplyAllOptionOverrides();
+  absl::Status ApplyAllOptionOverrides();
 
   // Applies a single option to executable_build_options.debug_options().
-  Status ApplyOption(const std::string& key, const OptionOverride& value);
+  absl::Status ApplyOption(const std::string& key, const OptionOverride& value);
 
-  Status ApplyOptionFromString(const tsl::protobuf::FieldDescriptor* field,
-                               const std::string& value);
+  absl::Status ApplyOptionFromString(
+      const tsl::protobuf::FieldDescriptor* field, const std::string& value);
 
-  static StatusOr<
+  static absl::StatusOr<
       std::vector<std::pair<std::string, CompileOptions::OptionOverride>>>
   LoadEnvOptionOverrides(
       const google::protobuf::Map<std::string, xla::OptionOverrideProto>&
@@ -130,10 +130,11 @@ struct CompileOptions {
           output_env_option_overrides) const;
 
   // Serialize the CompileOptions into a CompileOptionsProto.
-  StatusOr<CompileOptionsProto> ToProto() const;
+  absl::StatusOr<CompileOptionsProto> ToProto() const;
 
   // Deserialize the CompileOptionsProto into a CompileOptions.
-  static StatusOr<CompileOptions> FromProto(const CompileOptionsProto& proto);
+  static absl::StatusOr<CompileOptions> FromProto(
+      const CompileOptionsProto& proto);
 };
 
 struct LoadOptions {
@@ -191,8 +192,9 @@ struct SendCallback {
   // TODO(chky): Currently the callback invocation order may not be consistent
   // with the HLO send op invocation order, due to limitations in some PjRt
   // implementation. Consider making it strictly the same order as HLO program.
-  std::function<Status(const PjRtTransferMetadata& metadata, PjRtChunk chunk,
-                       size_t total_size_in_bytes, bool done)>
+  std::function<absl::Status(const PjRtTransferMetadata& metadata,
+                             PjRtChunk chunk, size_t total_size_in_bytes,
+                             bool done)>
       callback;
 };
 
@@ -320,21 +322,21 @@ class PjRtExecutable {
   virtual absl::string_view name() const = 0;
 
   // Return an HloModule (optimized) per partition.
-  virtual StatusOr<std::vector<std::shared_ptr<HloModule>>> GetHloModules()
-      const = 0;
+  virtual absl::StatusOr<std::vector<std::shared_ptr<HloModule>>>
+  GetHloModules() const = 0;
 
   // Returns an output Shape per program, the size should be equal to
   // `GetHloModules()`.
-  virtual StatusOr<std::vector<Shape>> GetOutputShapes() const;
+  virtual absl::StatusOr<std::vector<Shape>> GetOutputShapes() const;
 
   // Returns a list of element types for each output, the size of the outer list
   // should be equal to `GetHloModules()`.
-  virtual StatusOr<std::vector<std::vector<PrimitiveType>>>
+  virtual absl::StatusOr<std::vector<std::vector<PrimitiveType>>>
   GetOutputElementTypes() const;
 
   // Returns a list of dimensions for each output, the size of the outer list
   // should be equal to `GetHloModules()`.
-  virtual StatusOr<std::vector<std::vector<DimensionVector>>>
+  virtual absl::StatusOr<std::vector<std::vector<DimensionVector>>>
   GetOutputDimensions() const;
 
   // Returns the layout of each input parameter.
@@ -349,7 +351,7 @@ class PjRtExecutable {
   // value is `[num_programs, num_output]`. The size of the outer list should be
   // equal to `GetHloModules()`. Under SPMD, one can use
   // `GetOutputMemoryKinds().front()`.
-  virtual StatusOr<std::vector<std::vector<absl::string_view>>>
+  virtual absl::StatusOr<std::vector<std::vector<absl::string_view>>>
   GetOutputMemoryKinds() const = 0;
 
   // Returns a list of parameter OpSharding protos.
@@ -360,38 +362,38 @@ class PjRtExecutable {
 
   // Return memory stats that allow callers to estimate device memory usage
   // when running this executable.
-  virtual StatusOr<CompiledMemoryStats> GetCompiledMemoryStats() const {
+  virtual absl::StatusOr<CompiledMemoryStats> GetCompiledMemoryStats() const {
     return Unimplemented("Retrieving CompiledMemoryStats is not supported.");
   }
 
   // Returns named values for cost properties of this executable (such as
   // operations, size of input/outputs, and run time estimate). Properties may
   // differ for different platforms.
-  virtual StatusOr<absl::flat_hash_map<std::string, PjRtValueType>>
+  virtual absl::StatusOr<absl::flat_hash_map<std::string, PjRtValueType>>
   GetCostAnalysis() const = 0;
 
   // Serialize this executable into a string and return the value.
-  virtual StatusOr<std::string> SerializeExecutable() const {
+  virtual absl::StatusOr<std::string> SerializeExecutable() const {
     return Unimplemented("Serializing executable is not supported.");
   }
 
   // Return a fingerprint of this executable.
-  virtual StatusOr<std::string> FingerprintExecutable() const {
+  virtual absl::StatusOr<std::string> FingerprintExecutable() const {
     return Unimplemented("Fingerprinting executable is not supported.");
   }
 
-  virtual StatusOr<struct CompileOptions> GetCompileOptions() const {
+  virtual absl::StatusOr<struct CompileOptions> GetCompileOptions() const {
     return Unimplemented("CompileOptions not available.");
   }
 };
 
 class PjRtExecutableUtil {
  public:
-  static StatusOr<absl::flat_hash_map<std::string, PjRtValueType>>
+  static absl::StatusOr<absl::flat_hash_map<std::string, PjRtValueType>>
   RunHloCostAnalysis(const PjRtExecutable& executable,
                      HloCostAnalysis* hlo_cost_analysis);
 
-  static StatusOr<absl::flat_hash_map<std::string, PjRtValueType>>
+  static absl::StatusOr<absl::flat_hash_map<std::string, PjRtValueType>>
   RunHloCostAnalysis(
       const std::vector<std::shared_ptr<xla::HloModule>>& hlo_modules,
       HloCostAnalysis* hlo_cost_analysis);
diff --git a/third_party/xla/xla/pjrt/pjrt_future.h b/third_party/xla/xla/pjrt/pjrt_future.h
index 9d112c049125bc..ac76be51859a1d 100644
--- a/third_party/xla/xla/pjrt/pjrt_future.h
+++ b/third_party/xla/xla/pjrt/pjrt_future.h
@@ -382,8 +382,9 @@ class PjRtFuture : public internal::PjRtFutureBase<absl::StatusOr<T>> {
   static_assert(!std::is_same_v<T, absl::Status>,
                 "Use PjRtFuture<> specialization for stateless futures");
 
-  static_assert(!internal::IsStatusOr<T>::value,
-                "PjRtFuture<T> already has an implicit StatusOr<T> semantics");
+  static_assert(
+      !internal::IsStatusOr<T>::value,
+      "PjRtFuture<T> already has an implicit absl::StatusOr<T> semantics");
 
  public:
   class Promise : public Base::Promise {
diff --git a/third_party/xla/xla/pjrt/pjrt_layout.h b/third_party/xla/xla/pjrt/pjrt_layout.h
index 0fbf205de8b6e7..60d25a3430b42b 100644
--- a/third_party/xla/xla/pjrt/pjrt_layout.h
+++ b/third_party/xla/xla/pjrt/pjrt_layout.h
@@ -70,7 +70,8 @@ class PjRtXlaLayout : public PjRtLayout {
 
   std::string Serialize() const override { return xla_layout_.ToString(); }
 
-  static StatusOr<PjRtXlaLayout> Deserialize(absl::string_view serialized) {
+  static absl::StatusOr<PjRtXlaLayout> Deserialize(
+      absl::string_view serialized) {
     TF_ASSIGN_OR_RETURN(Layout xla_layout, ParseLayout(serialized));
     return PjRtXlaLayout(std::move(xla_layout));
   }
diff --git a/third_party/xla/xla/pjrt/pjrt_stream_executor_client.cc b/third_party/xla/xla/pjrt/pjrt_stream_executor_client.cc
index fcf28b0773a118..b0e11d18f260a7 100644
--- a/third_party/xla/xla/pjrt/pjrt_stream_executor_client.cc
+++ b/third_party/xla/xla/pjrt/pjrt_stream_executor_client.cc
@@ -172,15 +172,15 @@ absl::string_view PjRtStreamExecutorDevice::platform_name() const {
   return client_->platform_name();
 }
 
-StatusOr<LocalDeviceState*> PjRtStreamExecutorDevice::GetLocalDeviceState()
-    const {
+absl::StatusOr<LocalDeviceState*>
+PjRtStreamExecutorDevice::GetLocalDeviceState() const {
   if (local_device_state_) {
     return local_device_state_.get();
   }
   return InvalidArgument("Device %s is not a local device.", DebugString());
 }
 
-StatusOr<DeviceAssignment> DevicesToDeviceAssignment(
+absl::StatusOr<DeviceAssignment> DevicesToDeviceAssignment(
     absl::Span<const std::vector<PjRtDevice*>> devices) {
   if (devices.empty()) {
     return InvalidArgument(
@@ -295,13 +295,14 @@ PjRtStreamExecutorClient::PjRtStreamExecutorClient(
                });
 }
 
-StatusOr<DeviceAssignment> PjRtStreamExecutorClient::GetDefaultDeviceAssignment(
-    int num_replicas, int num_partitions) const {
+absl::StatusOr<DeviceAssignment>
+PjRtStreamExecutorClient::GetDefaultDeviceAssignment(int num_replicas,
+                                                     int num_partitions) const {
   return client_->backend().computation_placer()->AssignDevices(num_replicas,
                                                                 num_partitions);
 }
 
-StatusOr<Layout> PjRtStreamExecutorClient::GetDefaultLayout(
+absl::StatusOr<Layout> PjRtStreamExecutorClient::GetDefaultLayout(
     PrimitiveType element_type, absl::Span<const int64_t> dims) {
   Shape shape = ShapeUtil::MakeShape(element_type, dims);
   TF_ASSIGN_OR_RETURN(
@@ -311,7 +312,7 @@ StatusOr<Layout> PjRtStreamExecutorClient::GetDefaultLayout(
   return shape.layout();
 }
 
-StatusOr<std::unique_ptr<HloCostAnalysis>>
+absl::StatusOr<std::unique_ptr<HloCostAnalysis>>
 PjRtStreamExecutorClient::GetHloCostAnalysis() const {
   return std::make_unique<HloCostAnalysis>(
       client_->backend().compiler()->ShapeSizeBytesFunction());
@@ -425,7 +426,8 @@ void RecordUsage(PjRtStreamExecutorBuffer::ScopedHold device_buffer,
 // The caller may optionally provide a definition event to be recorded in
 // the buffer.
 // TODO(phawkins): replace on_host_shape here with on_device_shape.
-StatusOr<std::unique_ptr<PjRtStreamExecutorBuffer>> AllocateDestinationBuffer(
+absl::StatusOr<std::unique_ptr<PjRtStreamExecutorBuffer>>
+AllocateDestinationBuffer(
     const Shape& on_host_shape, PjRtDevice* device,
     LocalDeviceState* local_device, se::Stream* copy_stream,
     bool is_uninitialized_create, PjRtStreamExecutorClient* client,
@@ -510,7 +512,7 @@ StatusOr<std::unique_ptr<PjRtStreamExecutorBuffer>> AllocateDestinationBuffer(
 
     definition_events.emplace_back(
         std::make_shared<BufferSequencingEvent>(client->thread_pool()));
-    StatusOr<EventPool::Handle> event_or =
+    absl::StatusOr<EventPool::Handle> event_or =
         local_device->event_pool().ThenAllocateAndRecordEvent(
             tuple_table_stream);
     if (!event_or.ok()) {
@@ -548,12 +550,12 @@ StatusOr<std::unique_ptr<PjRtStreamExecutorBuffer>> AllocateDestinationBuffer(
 // Adds necessary synchronization after a copy has been enqueued to a buffer.
 // definition_event was added when the buffer was allocated, but has not yet
 // had an event recorded.
-Status AddDestinationBufferSynchronization(
+absl::Status AddDestinationBufferSynchronization(
     LocalDeviceState* local_device,
     PjRtStreamExecutorBuffer::ScopedHold device_buffer,
     std::shared_ptr<BufferSequencingEvent> definition_event,
     se::Stream* copy_stream) {
-  StatusOr<EventPool::Handle> event_or =
+  absl::StatusOr<EventPool::Handle> event_or =
       local_device->event_pool().ThenAllocateAndRecordEvent(copy_stream);
   if (!event_or.ok()) {
     StallStreamOnError(local_device, copy_stream);
@@ -592,7 +594,7 @@ PjRtStreamExecutorBuffer::ScopedHold::ScopedHold(ScopedHold&& other)
 }
 
 void PjRtStreamExecutorBuffer::ScopedHold::Acquire(
-    StatusOr<std::shared_ptr<TrackedDeviceBuffer>>&& buffer_or) {
+    absl::StatusOr<std::shared_ptr<TrackedDeviceBuffer>>&& buffer_or) {
   CHECK(!ok());
   if (buffer_or.ok()) {
     buffer_ = buffer_or.value();
@@ -648,7 +650,7 @@ void PjRtStreamExecutorBuffer::ScopedHold::AddToInput(
 
 bool PjRtStreamExecutorBuffer::IsOnCpu() const { return false; }
 
-StatusOr<Shape> PjRtStreamExecutorBuffer::logical_on_device_shape() {
+absl::StatusOr<Shape> PjRtStreamExecutorBuffer::logical_on_device_shape() {
   if (on_device_shape_.is_static()) {
     return on_device_shape_;
   }
@@ -668,7 +670,7 @@ StatusOr<Shape> PjRtStreamExecutorBuffer::logical_on_device_shape() {
 
   WaitForBufferDefinitionEventsOnStream(*device_buffer, stream);
   ShapedBuffer shaped_buffer = device_buffer->AsShapedBuffer(on_device_shape_);
-  StatusOr<EventPool::Handle> event_or =
+  absl::StatusOr<EventPool::Handle> event_or =
       local_device->event_pool().AllocateEvent(stream->parent());
   if (!event_or.ok()) {
     return event_or.status();
@@ -697,7 +699,7 @@ class ScopedHoldAsExternalReference : public PjRtBuffer::ExternalReference {
 
   ~ScopedHoldAsExternalReference() override = default;
 
-  Status WaitUntilBufferReadyOnStream(std::intptr_t stream) override {
+  absl::Status WaitUntilBufferReadyOnStream(std::intptr_t stream) override {
     for (const std::shared_ptr<BufferSequencingEvent>& event :
          external_reference_->definition_events()) {
       TF_RETURN_IF_ERROR(event->WaitForEventOnExternalStream(stream));
@@ -711,10 +713,10 @@ class ScopedHoldAsExternalReference : public PjRtBuffer::ExternalReference {
 
 }  // namespace
 
-StatusOr<std::unique_ptr<PjRtBuffer::ExternalReference>>
+absl::StatusOr<std::unique_ptr<PjRtBuffer::ExternalReference>>
 PjRtStreamExecutorBuffer::AcquireExternalReference() {
   ScopedHold hold = GetBufferWithExternalReference();
-  Status hold_status = hold.status();
+  absl::Status hold_status = hold.status();
   if (!hold_status.ok()) return hold_status;
   return std::unique_ptr<ExternalReference>(
       std::make_unique<ScopedHoldAsExternalReference>(std::move(hold)));
@@ -735,7 +737,7 @@ class TrackedDeviceBufferExternalReference
   std::shared_ptr<TrackedDeviceBuffer> tracked_device_buffer_;
 };
 
-StatusOr<std::unique_ptr<PjRtBuffer::ExternalReference>>
+absl::StatusOr<std::unique_ptr<PjRtBuffer::ExternalReference>>
 PjRtStreamExecutorBuffer::ReleaseDeviceMemoryOwnership(
     bool wait_for_operations_to_complete) {
   if (on_device_shape_.IsTuple()) {
@@ -817,7 +819,7 @@ PjRtStreamExecutorBuffer::DonateWithControlDependency(PjRtFuture<> dependency) {
   return new_buffer;
 }
 
-StatusOr<std::unique_ptr<PjRtBuffer>>
+absl::StatusOr<std::unique_ptr<PjRtBuffer>>
 PjRtStreamExecutorClient::BufferFromHostBuffer(
     const void* data, PrimitiveType type, absl::Span<int64_t const> dims,
     std::optional<absl::Span<int64_t const>> byte_strides,
@@ -1025,7 +1027,7 @@ PjRtStreamExecutorClient::BufferFromHostBuffer(
   return std::unique_ptr<PjRtBuffer>(std::move(py_buffer));
 }
 
-StatusOr<std::unique_ptr<PjRtBuffer>>
+absl::StatusOr<std::unique_ptr<PjRtBuffer>>
 PjRtStreamExecutorClient::BufferFromHostBuffer(
     const void* data, PrimitiveType type, absl::Span<int64_t const> dims,
     std::optional<absl::Span<int64_t const>> byte_strides,
@@ -1056,13 +1058,13 @@ PjRtStreamExecutorClient::BufferFromHostBuffer(
       platform_name()));
 }
 
-StatusOr<std::unique_ptr<PjRtBuffer>>
+absl::StatusOr<std::unique_ptr<PjRtBuffer>>
 PjRtStreamExecutorClient::CreateUninitializedBuffer(const Shape& shape,
                                                     PjRtDevice* device) {
   return CreateUninitializedBuffer(shape, device, nullptr);
 }
 
-StatusOr<std::unique_ptr<PjRtBuffer>>
+absl::StatusOr<std::unique_ptr<PjRtBuffer>>
 PjRtStreamExecutorClient::CreateUninitializedBuffer(
     const Shape& shape, PjRtDevice* device,
     std::shared_ptr<BufferSequencingEvent> definition_event) {
@@ -1087,8 +1089,9 @@ PjRtStreamExecutorClient::CreateUninitializedBuffer(
   return std::unique_ptr<PjRtBuffer>(std::move(py_buffer));
 }
 
-StatusOr<std::unique_ptr<PjRtBuffer>>
-PjRtStreamExecutorClient::CreateErrorBuffer(Status error, const Shape& shape,
+absl::StatusOr<std::unique_ptr<PjRtBuffer>>
+PjRtStreamExecutorClient::CreateErrorBuffer(absl::Status error,
+                                            const Shape& shape,
                                             PjRtMemorySpace* memory) {
   if (memory->client() != this) {
     return absl::InvalidArgumentError(
@@ -1120,7 +1123,7 @@ PjRtStreamExecutorClient::CreateErrorBuffer(Status error, const Shape& shape,
   return py_buffer;
 }
 
-StatusOr<std::unique_ptr<PjRtBuffer>>
+absl::StatusOr<std::unique_ptr<PjRtBuffer>>
 PjRtStreamExecutorClient::BufferFromHostLiteral(const LiteralSlice& literal,
                                                 PjRtDevice* device) {
   tsl::profiler::TraceMe traceme(
@@ -1195,7 +1198,7 @@ PjRtStreamExecutorClient::BufferFromHostLiteral(const LiteralSlice& literal,
       platform_name()));
 }
 
-StatusOr<std::vector<std::unique_ptr<PjRtBuffer>>>
+absl::StatusOr<std::vector<std::unique_ptr<PjRtBuffer>>>
 PjRtStreamExecutorClient::MakeCrossHostReceiveBuffers(
     absl::Span<const Shape> shapes, PjRtDevice* device,
     PjRtCrossHostRecvNotifier notifier) {
@@ -1226,7 +1229,7 @@ PjRtStreamExecutorClient::MakeCrossHostReceiveBuffers(
   return buffers;
 }
 
-StatusOr<std::vector<std::unique_ptr<PjRtBuffer>>>
+absl::StatusOr<std::vector<std::unique_ptr<PjRtBuffer>>>
 PjRtStreamExecutorClient::MakeCrossHostReceiveBuffersForGather(
     absl::Span<const Shape> shapes, std::vector<GatherDetails> gather_details,
     PjRtDevice* device, PjRtCrossHostRecvNotifier notifier) {
@@ -1269,7 +1272,7 @@ PjRtStreamExecutorClient::MakeCrossHostReceiveBuffersForGather(
   return buffers;
 }
 
-StatusOr<std::unique_ptr<PjRtBuffer>>
+absl::StatusOr<std::unique_ptr<PjRtBuffer>>
 PjRtStreamExecutorClient::CreateViewOfDeviceBuffer(
     void* device_ptr, const Shape& shape, PjRtDevice* device,
     std::function<void()> on_delete_callback,
@@ -1308,14 +1311,15 @@ PjRtStreamExecutorClient::CreateViewOfDeviceBuffer(
 }
 
 // Transfer the given literal to the infeed queue of the given local device.
-Status PjRtStreamExecutorDevice::TransferToInfeed(const LiteralSlice& literal) {
+absl::Status PjRtStreamExecutorDevice::TransferToInfeed(
+    const LiteralSlice& literal) {
   // Only support infeed to local device.
   TF_ASSIGN_OR_RETURN(LocalDeviceState * local_device, GetLocalDeviceState());
   return local_device->client()->TransferToInfeedLocal(
       literal, local_device->local_hardware_id().value());
 }
 
-Status PjRtStreamExecutorDevice::TransferFromOutfeed(
+absl::Status PjRtStreamExecutorDevice::TransferFromOutfeed(
     MutableBorrowingLiteral literal) {
   VLOG(1) << "PjRtStreamExecutorDevice::TransferFromOutfeed";
   TF_ASSIGN_OR_RETURN(LocalDeviceState * local_device, GetLocalDeviceState());
@@ -1341,8 +1345,8 @@ absl::Span<PjRtMemorySpace* const> PjRtStreamExecutorDevice::memory_spaces()
   return memory_spaces_;
 }
 
-StatusOr<PjRtMemorySpace*> PjRtStreamExecutorDevice::default_memory_space()
-    const {
+absl::StatusOr<PjRtMemorySpace*>
+PjRtStreamExecutorDevice::default_memory_space() const {
   return Unimplemented("default_memory_space is not supported.");
 }
 
@@ -1369,7 +1373,7 @@ PjRtStreamExecutorDevice::memory_space_by_kind_id(int id) const {
   return it->second;
 }
 
-StatusOr<std::intptr_t>
+absl::StatusOr<std::intptr_t>
 PjRtStreamExecutorDevice::GetStreamForExternalReadyEvents() const {
   TF_ASSIGN_OR_RETURN(LocalDeviceState * local_device, GetLocalDeviceState());
   se::Stream* stream = local_device->GetExternalReadyEventStream();
@@ -1382,12 +1386,12 @@ PjRtStreamExecutorDevice::GetStreamForExternalReadyEvents() const {
   return absl::bit_cast<std::intptr_t>(raw_stream);
 }
 
-StatusOr<PjRtDevice*> PjRtStreamExecutorClient::LookupAddressableDevice(
+absl::StatusOr<PjRtDevice*> PjRtStreamExecutorClient::LookupAddressableDevice(
     int local_hardware_id) const {
   return LookupAddressableDevice(PjRtLocalDeviceId(local_hardware_id));
 }
 
-StatusOr<PjRtDevice*> PjRtStreamExecutorClient::LookupAddressableDevice(
+absl::StatusOr<PjRtDevice*> PjRtStreamExecutorClient::LookupAddressableDevice(
     xla::PjRtLocalDeviceId local_device_id) const {
   for (auto* device : addressable_devices_) {
     if (local_device_id == device->local_device_id()) {
@@ -1438,7 +1442,7 @@ void PjRtStreamExecutorBuffer::WaitForOutstandingDonationHold() {
   mu_.Await(absl::Condition(&not_in_donation_hold));
 }
 
-StatusOr<std::shared_ptr<TrackedDeviceBuffer>>
+absl::StatusOr<std::shared_ptr<TrackedDeviceBuffer>>
 PjRtStreamExecutorBuffer::Release(bool wait_for_operations_to_complete) {
   tsl::profiler::TraceMe trace_me("PjRtStreamExecutorBuffer::Release");
   std::shared_ptr<TrackedDeviceBuffer> device_buffer;
@@ -1526,7 +1530,7 @@ bool PjRtStreamExecutorBuffer::IsDeleted() {
   return device_buffer_ == nullptr;
 }
 
-StatusOr<std::shared_ptr<TrackedDeviceBuffer>>
+absl::StatusOr<std::shared_ptr<TrackedDeviceBuffer>>
 PjRtStreamExecutorBuffer::GetBufferForHoldLocked(ScopedHold::Type type) {
   // All callers should have called WaitForOutstandingDonationHold().
   CHECK_EQ(holds_[ScopedHold::kDonation], 0);
@@ -1658,14 +1662,14 @@ PjRtFuture<> PjRtStreamExecutorBuffer::ToLiteral(MutableLiteralBase* literal) {
                            transfer_manager = std::move(transfer_manager),
                            on_device_shape{on_device_shape_}, literal, promise,
                            local_device]() mutable {
-    StatusOr<EventPool::Handle> event_or =
+    absl::StatusOr<EventPool::Handle> event_or =
         local_device->event_pool().AllocateEvent(stream->parent());
     if (!event_or.ok()) {
       promise.Set(event_or.status());
       return;
     }
 
-    Status defined_status =
+    absl::Status defined_status =
         tracked_device_buffer->definition_events()[0]->GetDefinedStatus();
     if (!defined_status.ok()) {
       promise.Set(defined_status);
@@ -1687,7 +1691,9 @@ PjRtFuture<> PjRtStreamExecutorBuffer::ToLiteral(MutableLiteralBase* literal) {
 
     transfer_manager->TransferLiteralFromDevice(
         stream, shaped_buffer, literal,
-        [promise](Status status) mutable { promise.Set(std::move(status)); },
+        [promise](absl::Status status) mutable {
+          promise.Set(std::move(status));
+        },
         transfer_metadata_ptr);
 
     local_device->event_pool().ThenRecordEvent(stream, event_or.value());
@@ -1720,7 +1726,8 @@ PjRtFuture<> PjRtStreamExecutorBuffer::ToLiteral(MutableLiteralBase* literal) {
       });
 }
 
-StatusOr<size_t> PjRtStreamExecutorBuffer::GetOnDeviceSizeInBytes() const {
+absl::StatusOr<size_t> PjRtStreamExecutorBuffer::GetOnDeviceSizeInBytes()
+    const {
   absl::MutexLock lock(&mu_);
   if (device_buffer_ == nullptr) {
     return InvalidArgument(
@@ -1744,7 +1751,7 @@ PjRtFuture<> PjRtStreamExecutorBuffer::CopyRawToHostFuture(
   return client_->CopyRawSubBufferToHost(this, dst, offset, transfer_size);
 }
 
-StatusOr<ShapedBuffer> PjRtStreamExecutorBuffer::AsShapedBuffer() const {
+absl::StatusOr<ShapedBuffer> PjRtStreamExecutorBuffer::AsShapedBuffer() const {
   absl::MutexLock lock(&mu_);
   if (device_buffer_ == nullptr) {
     return InvalidArgument(
@@ -1763,8 +1770,8 @@ PjRtStreamExecutorBuffer::GetBufferWithHold(ScopedHold::Type type) {
   return hold;
 }
 
-StatusOr<std::pair<std::unique_ptr<PjRtBuffer>,
-                   std::shared_ptr<BufferSequencingEvent>>>
+absl::StatusOr<std::pair<std::unique_ptr<PjRtBuffer>,
+                         std::shared_ptr<BufferSequencingEvent>>>
 PjRtStreamExecutorBuffer::CopyToDeviceHelper(
     PjRtDevice* dst_device, LocalDeviceState* dst_local_device,
     LocalDeviceState* transfer_local_device, LocalDeviceState* src_local_device,
@@ -1800,7 +1807,7 @@ PjRtStreamExecutorBuffer::CopyToDeviceHelper(
     VLOG(1)
         << "PjRtStreamExecutorBuffer::CopyToDeviceHelper::async_copy_to_device";
 
-    Status defined_status =
+    absl::Status defined_status =
         src_device_buffer->definition_events()[0]->GetDefinedStatus();
     // Only proceeds to transfer when the buffer doesn't hold an error.
     if (defined_status.ok()) {
@@ -1840,7 +1847,7 @@ PjRtStreamExecutorBuffer::CopyToDeviceHelper(
         }
       }
 
-      StatusOr<EventPool::Handle> event_or =
+      absl::StatusOr<EventPool::Handle> event_or =
           transfer_local_device->event_pool().ThenAllocateAndRecordEvent(
               transfer_stream);
       if (!event_or.ok()) {
@@ -1876,8 +1883,8 @@ PjRtStreamExecutorBuffer::CopyToDeviceHelper(
       std::move(copy_event));
 }
 
-StatusOr<std::unique_ptr<PjRtBuffer>> PjRtStreamExecutorBuffer::CopyToDevice(
-    PjRtDevice* dst_device) {
+absl::StatusOr<std::unique_ptr<PjRtBuffer>>
+PjRtStreamExecutorBuffer::CopyToDevice(PjRtDevice* dst_device) {
   tsl::profiler::TraceMe traceme("PjRtStreamExecutorBuffer::CopyToDevice");
   VLOG(1) << "PjRtStreamExecutorBuffer::CopyToDevice";
   if (dst_device == device_) {
@@ -1927,8 +1934,8 @@ StatusOr<std::unique_ptr<PjRtBuffer>> PjRtStreamExecutorBuffer::CopyToDevice(
     AcquireHoldLocked(&src_device_buffer);
   }
 
-  StatusOr<std::pair<std::unique_ptr<PjRtBuffer>,
-                     std::shared_ptr<BufferSequencingEvent>>>
+  absl::StatusOr<std::pair<std::unique_ptr<PjRtBuffer>,
+                           std::shared_ptr<BufferSequencingEvent>>>
       buffer_and_event_or = CopyToDeviceHelper(
           dst_device, dst_local_device, transfer_local_device,
           device_->local_device_state(), transfer_stream,
@@ -2010,7 +2017,7 @@ PjRtFuture<> PjRtStreamExecutorBuffer::GetReadyFuture() {
         [device_buffer, local_device_state = std::move(local_device_state),
          definition_promise]() mutable {
           std::unique_ptr<se::Stream> stream;
-          Status defined_status =
+          absl::Status defined_status =
               device_buffer->definition_events()[0]->GetDefinedStatus();
           if (!defined_status.ok()) {
             definition_promise.Set(defined_status);
@@ -2085,11 +2092,11 @@ struct TupleHandle {
   std::shared_ptr<BufferSequencingEvent> event;
 };
 
-Status CheckCompatibleShapes(bool strict_shape_checking,
-                             const Shape& buffer_on_device_shape,
-                             const Shape& execution_shape,
-                             const TransferManager& transfer_manager,
-                             int parameter_index) {
+absl::Status CheckCompatibleShapes(bool strict_shape_checking,
+                                   const Shape& buffer_on_device_shape,
+                                   const Shape& execution_shape,
+                                   const TransferManager& transfer_manager,
+                                   int parameter_index) {
   // Handle the special case: the underlying pjrt buffer of a JAX token may have
   // shape `pred[0]`.
   if (execution_shape.IsToken() &&
@@ -2134,7 +2141,7 @@ Status CheckCompatibleShapes(bool strict_shape_checking,
 }
 
 // Makes a tuple from the arguments to an execution.
-StatusOr<TupleHandle> MakeTupleHelper(
+absl::StatusOr<TupleHandle> MakeTupleHelper(
     PjRtStreamExecutorClient* client, LocalDeviceState* local_device,
     bool strict_shape_checking, const Shape& tupled_parameter_shape,
     absl::Span<PjRtBuffer* const> py_buffers,
@@ -2190,7 +2197,7 @@ StatusOr<TupleHandle> MakeTupleHelper(
 
   TF_RETURN_IF_ERROR(transfer_manager->WriteRootTupleIndexTable(
       stream, execution_input.Buffers()));
-  StatusOr<EventPool::Handle> event_or =
+  absl::StatusOr<EventPool::Handle> event_or =
       local_device->event_pool().ThenAllocateAndRecordEvent(stream);
   if (!event_or.ok()) {
     StallStreamOnError(local_device, stream);
@@ -2301,7 +2308,8 @@ PjRtStreamExecutorLoadedExecutable::PjRtStreamExecutorLoadedExecutable(
   }
 }
 
-Status PjRtStreamExecutorLoadedExecutable::SetUpDonation(bool tuple_inputs) {
+absl::Status PjRtStreamExecutorLoadedExecutable::SetUpDonation(
+    bool tuple_inputs) {
   parameters_that_must_be_donated_.reserve(executables_.size());
   for (auto& executable : executables_) {
     TF_ASSIGN_OR_RETURN(std::vector<int> parameters_to_donate,
@@ -2328,7 +2336,7 @@ PjRtStreamExecutorLoadedExecutable::ParametersThatMustBeDonated(
   return parameters_that_must_be_donated_[executable_idx];
 }
 
-StatusOr<std::vector<ExecutionInput>>
+absl::StatusOr<std::vector<ExecutionInput>>
 PjRtStreamExecutorLoadedExecutable::MakeExecutionInputsAndWaitForEvents(
     int device_ordinal, const ExecuteOptions& options,
     absl::Span<const Shape> executable_parameter_shapes,
@@ -2419,7 +2427,7 @@ static SendDeviceMemoryFunction ConvertSendCallbacksToSendFunction(
              int64_t channel_id, se::Stream* stream, const Shape& shape,
              const se::DeviceMemoryBase& src,
              const absl::flat_hash_map<std::string, std::string>&)
-             -> StatusOr<AsyncValueRef<se::Event>> {
+             -> absl::StatusOr<AsyncValueRef<se::Event>> {
     VLOG(3) << "Send " << src.size() << " bytes to channel #" << channel_id
             << " (shape=" << shape.ToString() << ")";
 
@@ -2592,7 +2600,7 @@ static RecvDeviceMemoryFunction ConvertRecvCallbacksToRecvFunction(
   return [callbacks](int64_t channel_id, se::Stream* stream, const Shape& shape,
                      se::DeviceMemoryBase* dst,
                      const absl::flat_hash_map<std::string, std::string>&)
-             -> StatusOr<AsyncValueRef<se::Event>> {
+             -> absl::StatusOr<AsyncValueRef<se::Event>> {
     VLOG(3) << "Recv from channel #" << channel_id
             << " (shape=" << shape.ToString() << ")";
 
@@ -2629,7 +2637,7 @@ static RecvDeviceMemoryFunction ConvertRecvCallbacksToRecvFunction(
 // converted on success.
 // When `options` has non-zero `launch_id`, use `launch_id` instead of `run_id`
 // to initialize `run_options`.
-StatusOr<ScopedShapedBuffer>
+absl::StatusOr<ScopedShapedBuffer>
 PjRtStreamExecutorLoadedExecutable::EnqueueExecution(
     absl::Span<PjRtBuffer* const> argument_handles, int replica, int partition,
     int executable_idx, const RunId& run_id, const ExecuteOptions& options,
@@ -2787,7 +2795,7 @@ PjRtStreamExecutorLoadedExecutable::EnqueueExecution(
         device_state->compute_semaphore().ScopedAcquire(1));
   }
 
-  StatusOr<ExecutionOutput> result_buffer_or_status =
+  absl::StatusOr<ExecutionOutput> result_buffer_or_status =
       executables_[executable_idx]->RunAsync(std::move(execution_inputs),
                                              run_options);
 
@@ -2890,7 +2898,7 @@ static Status GetFirstInputError(
   return OkStatus();
 }
 
-StatusOr<PjRtLoadedExecutable::Result>
+absl::StatusOr<PjRtLoadedExecutable::Result>
 PjRtStreamExecutorLoadedExecutable::ExecuteHelper(
     absl::Span<PjRtBuffer* const> argument_handles, int replica, int partition,
     const RunId& run_id, const ExecuteOptions& options, bool fill_future,
@@ -2945,7 +2953,7 @@ PjRtStreamExecutorLoadedExecutable::ExecuteHelper(
   std::vector<std::function<void()>> compute_callbacks;
   std::vector<PjRtStreamExecutorBuffer::ScopedHold> device_buffers;
   device_buffers.reserve(argument_handles.size());
-  StatusOr<ScopedShapedBuffer> result_buffer_or_status = EnqueueExecution(
+  absl::StatusOr<ScopedShapedBuffer> result_buffer_or_status = EnqueueExecution(
       argument_handles, replica, partition, executable_idx, run_id, options,
       device, &device_buffers, std::move(device_assignment), compute_callbacks);
 
@@ -2958,7 +2966,7 @@ PjRtStreamExecutorLoadedExecutable::ExecuteHelper(
 
   LocalDeviceState* device_state = &(client_->device_state(device_ordinal));
   se::Stream* stream = device_state->compute_stream();
-  StatusOr<EventPool::Handle> event_or =
+  absl::StatusOr<EventPool::Handle> event_or =
       device_state->event_pool().ThenAllocateAndRecordEvent(stream);
   if (!event_or.ok()) {
     StallStreamOnError(device_state, stream);
@@ -3014,7 +3022,7 @@ PjRtStreamExecutorLoadedExecutable::ExecuteHelper(
   return Result({/*future=*/std::move(future), /*buffers=*/std::move(outputs)});
 }
 
-StatusOr<std::vector<std::vector<std::unique_ptr<PjRtBuffer>>>>
+absl::StatusOr<std::vector<std::vector<std::unique_ptr<PjRtBuffer>>>>
 PjRtStreamExecutorLoadedExecutable::Execute(
     absl::Span<const std::vector<PjRtBuffer*>> argument_handles,
     const ExecuteOptions& options,
@@ -3054,7 +3062,7 @@ PjRtStreamExecutorLoadedExecutable::Execute(
     absl::Mutex mu;
     int running = num_addressable_devices;
     int failed = 0;
-    Status first_failure_status;
+    absl::Status first_failure_status;
 
     for (int i = 0; i < num_addressable_devices; ++i) {
       const int replica = addressable_device_logical_ids_[i].replica;
@@ -3139,7 +3147,7 @@ PjRtStreamExecutorLoadedExecutable::Execute(
   return wrapped_results;
 }
 
-StatusOr<std::vector<std::unique_ptr<PjRtBuffer>>>
+absl::StatusOr<std::vector<std::unique_ptr<PjRtBuffer>>>
 PjRtStreamExecutorLoadedExecutable::ExecuteSharded(
     absl::Span<PjRtBuffer* const> argument_handles, PjRtDevice* device,
     const ExecuteOptions& options, std::optional<PjRtFuture<>>& returned_future,
@@ -3168,7 +3176,7 @@ PjRtStreamExecutorLoadedExecutable::ExecuteSharded(
       device->id());
 }
 
-StatusOr<std::vector<std::unique_ptr<PjRtBuffer>>>
+absl::StatusOr<std::vector<std::unique_ptr<PjRtBuffer>>>
 PjRtStreamExecutorLoadedExecutable::ExecutePortable(
     absl::Span<PjRtBuffer* const> argument_handles, PjRtDevice* device,
     const ExecuteOptions& options, std::optional<PjRtFuture<>>& returned_future,
@@ -3195,7 +3203,7 @@ PjRtStreamExecutorLoadedExecutable::ExecutePortable(
   return std::move(result.buffers);
 }
 
-StatusOr<std::vector<std::shared_ptr<HloModule>>>
+absl::StatusOr<std::vector<std::shared_ptr<HloModule>>>
 PjRtStreamExecutorLoadedExecutable::GetHloModules() const {
   std::vector<std::shared_ptr<HloModule>> modules;
   modules.reserve(executables().size());
@@ -3208,12 +3216,12 @@ PjRtStreamExecutorLoadedExecutable::GetHloModules() const {
   return std::move(modules);
 }
 
-StatusOr<std::vector<std::vector<absl::string_view>>>
+absl::StatusOr<std::vector<std::vector<absl::string_view>>>
 PjRtStreamExecutorLoadedExecutable::GetOutputMemoryKinds() const {
   return Unimplemented("GetOutputMemoryKinds is not supported.");
 }
 
-StatusOr<std::string>
+absl::StatusOr<std::string>
 PjRtStreamExecutorLoadedExecutable::FingerprintExecutable() const {
   if (executables_.size() != 1) {
     return absl::InternalError(
@@ -3229,7 +3237,7 @@ PjRtStreamExecutorLoadedExecutable::FingerprintExecutable() const {
   }
 }
 
-StatusOr<PjRtStreamExecutorClient::ExecutableExtras>
+absl::StatusOr<PjRtStreamExecutorClient::ExecutableExtras>
 PjRtStreamExecutorClient::GetExecutableExtras(CompileOptions* options) {
   ExecutableExtras extras;
   std::shared_ptr<DeviceAssignment>& device_assignment =
@@ -3247,7 +3255,7 @@ PjRtStreamExecutorClient::GetExecutableExtras(CompileOptions* options) {
   }
 
   auto layout_callback = [local_client = client()](const HloModule& module)
-      -> StatusOr<std::pair<std::vector<Shape>, Shape>> {
+      -> absl::StatusOr<std::pair<std::vector<Shape>, Shape>> {
     ExecutableBuildOptions build_options;
     std::vector<const Shape*> argument_layout_pointers;
     std::optional<std::vector<Shape>> argument_layouts;
@@ -3311,7 +3319,7 @@ PjRtStreamExecutorClient::GetExecutableExtras(CompileOptions* options) {
   return extras;
 }
 
-StatusOr<std::unique_ptr<PjRtLoadedExecutable>>
+absl::StatusOr<std::unique_ptr<PjRtLoadedExecutable>>
 PjRtStreamExecutorClient::Compile(const XlaComputation& computation,
                                   CompileOptions options) {
   tsl::profiler::TraceMe traceme("PjRtStreamExecutorClient::Compile");
@@ -3354,7 +3362,7 @@ PjRtStreamExecutorClient::Compile(const XlaComputation& computation,
   return std::unique_ptr<PjRtLoadedExecutable>(std::move(executable));
 }
 
-StatusOr<std::unique_ptr<PjRtLoadedExecutable>>
+absl::StatusOr<std::unique_ptr<PjRtLoadedExecutable>>
 PjRtStreamExecutorClient::Compile(mlir::ModuleOp module,
                                   CompileOptions options) {
   XlaComputation xla_computation;
@@ -3365,7 +3373,7 @@ PjRtStreamExecutorClient::Compile(mlir::ModuleOp module,
   return Compile(xla_computation, options);
 }
 
-StatusOr<std::string> PjRtStreamExecutorClient::SerializeExecutable(
+absl::StatusOr<std::string> PjRtStreamExecutorClient::SerializeExecutable(
     const PjRtLoadedExecutable& executable) const {
   const PjRtStreamExecutorLoadedExecutable* se_executable =
       tensorflow::down_cast<const PjRtStreamExecutorLoadedExecutable*>(
@@ -3399,7 +3407,7 @@ StatusOr<std::string> PjRtStreamExecutorClient::SerializeExecutable(
   return proto.SerializeAsString();
 }
 
-StatusOr<std::unique_ptr<PjRtLoadedExecutable>>
+absl::StatusOr<std::unique_ptr<PjRtLoadedExecutable>>
 PjRtStreamExecutorClient::DeserializeExecutable(
     absl::string_view serialized, std::optional<CompileOptions> options) {
   ExecutableAndOptionsProto proto;
@@ -3456,7 +3464,7 @@ PjRtStreamExecutorClient::DeserializeExecutable(
   return std::unique_ptr<PjRtLoadedExecutable>(std::move(executable));
 }
 
-StatusOr<std::unique_ptr<PjRtLoadedExecutable>>
+absl::StatusOr<std::unique_ptr<PjRtLoadedExecutable>>
 PjRtStreamExecutorClient::LoadSerializedExecutable(
     absl::string_view serialized, std::optional<CompileOptions> options,
     const LoadOptions& load_options) {
diff --git a/third_party/xla/xla/pjrt/pjrt_stream_executor_client.h b/third_party/xla/xla/pjrt/pjrt_stream_executor_client.h
index 9d95c5441b1f45..cc879e4f773b1e 100644
--- a/third_party/xla/xla/pjrt/pjrt_stream_executor_client.h
+++ b/third_party/xla/xla/pjrt/pjrt_stream_executor_client.h
@@ -188,24 +188,25 @@ class PjRtStreamExecutorDevice : public PjRtDevice {
   // If this is a device local to this host, returns a LocalDeviceState object
   // that can be used to manipulate the device. Returns an error if the device
   // is not local to this host.
-  StatusOr<LocalDeviceState*> GetLocalDeviceState() const;
+  absl::StatusOr<LocalDeviceState*> GetLocalDeviceState() const;
 
-  Status TransferToInfeed(const LiteralSlice& literal) override;
+  absl::Status TransferToInfeed(const LiteralSlice& literal) override;
 
-  Status TransferFromOutfeed(MutableBorrowingLiteral literal) override;
+  absl::Status TransferFromOutfeed(MutableBorrowingLiteral literal) override;
 
   void AttachMemorySpace(PjRtMemorySpace* memory_space);
 
   absl::Span<PjRtMemorySpace* const> memory_spaces() const override;
 
-  StatusOr<PjRtMemorySpace*> default_memory_space() const override;
+  absl::StatusOr<PjRtMemorySpace*> default_memory_space() const override;
 
   absl::StatusOr<PjRtMemorySpace*> memory_space_by_kind(
       absl::string_view memory_space_kind) const override;
 
   absl::StatusOr<PjRtMemorySpace*> memory_space_by_kind_id(int id) const;
 
-  StatusOr<std::intptr_t> GetStreamForExternalReadyEvents() const override;
+  absl::StatusOr<std::intptr_t> GetStreamForExternalReadyEvents()
+      const override;
 
   std::unique_ptr<ScopedAsyncTrackingEvent> CreateAsyncTrackingEvent(
       absl::string_view description) const override {
@@ -279,7 +280,7 @@ class PjRtStreamExecutorClient : public PjRtClient {
     return addressable_devices_;
   }
 
-  StatusOr<PjRtDevice*> LookupDevice(
+  absl::StatusOr<PjRtDevice*> LookupDevice(
       PjRtGlobalDeviceId global_device_id) const override {
     auto it = id_to_device_.find(global_device_id.value());
     if (it != id_to_device_.end()) {
@@ -289,9 +290,9 @@ class PjRtStreamExecutorClient : public PjRtClient {
                            global_device_id.value());
   }
 
-  StatusOr<PjRtDevice*> LookupAddressableDevice(
+  absl::StatusOr<PjRtDevice*> LookupAddressableDevice(
       int local_hardware_id) const override;
-  StatusOr<PjRtDevice*> LookupAddressableDevice(
+  absl::StatusOr<PjRtDevice*> LookupAddressableDevice(
       PjRtLocalDeviceId local_device_id) const override;
 
   absl::Span<PjRtMemorySpace* const> memory_spaces() const override;
@@ -306,67 +307,68 @@ class PjRtStreamExecutorClient : public PjRtClient {
   // function specifies which one the platform expects.
   virtual bool EnqueueD2DTransfersOnSrcStream() const { return true; }
 
-  StatusOr<DeviceAssignment> GetDefaultDeviceAssignment(
+  absl::StatusOr<DeviceAssignment> GetDefaultDeviceAssignment(
       int num_replicas, int num_partitions) const override;
 
-  StatusOr<Layout> GetDefaultLayout(PrimitiveType element_type,
-                                    absl::Span<const int64_t> dims) override;
+  absl::StatusOr<Layout> GetDefaultLayout(
+      PrimitiveType element_type, absl::Span<const int64_t> dims) override;
 
-  StatusOr<std::unique_ptr<PjRtLoadedExecutable>> Compile(
+  absl::StatusOr<std::unique_ptr<PjRtLoadedExecutable>> Compile(
       const XlaComputation& computation, CompileOptions options) override;
-  StatusOr<std::unique_ptr<PjRtLoadedExecutable>> Compile(
+  absl::StatusOr<std::unique_ptr<PjRtLoadedExecutable>> Compile(
       mlir::ModuleOp mlir_module, CompileOptions options) override;
 
-  virtual StatusOr<std::string> SerializeExecutable(
+  virtual absl::StatusOr<std::string> SerializeExecutable(
       const PjRtLoadedExecutable& executable) const;
 
   // For PjRtStreamExecutorClient, `options` is mandatory.
   // This function returns an InvalidArgument error if `std::nullopt` is passed.
   // TODO(b/237720161): make it actually optional
-  StatusOr<std::unique_ptr<PjRtLoadedExecutable>> DeserializeExecutable(
+  absl::StatusOr<std::unique_ptr<PjRtLoadedExecutable>> DeserializeExecutable(
       absl::string_view serialized,
       std::optional<CompileOptions> options) override;
 
-  StatusOr<std::unique_ptr<PjRtLoadedExecutable>> LoadSerializedExecutable(
-      absl::string_view serialized, std::optional<CompileOptions> options,
-      const LoadOptions& load_options) override;
+  absl::StatusOr<std::unique_ptr<PjRtLoadedExecutable>>
+  LoadSerializedExecutable(absl::string_view serialized,
+                           std::optional<CompileOptions> options,
+                           const LoadOptions& load_options) override;
 
-  StatusOr<std::unique_ptr<HloCostAnalysis>> GetHloCostAnalysis()
+  absl::StatusOr<std::unique_ptr<HloCostAnalysis>> GetHloCostAnalysis()
       const override;
 
   // Creates a buffer on the device without initializing or copying any data.
   // An optional `definition_event` may be speficied that can be used to
   // ensure the buffer isn't referenced until some external mechanism has
   // initialized the data.
-  StatusOr<std::unique_ptr<PjRtBuffer>> CreateUninitializedBuffer(
+  absl::StatusOr<std::unique_ptr<PjRtBuffer>> CreateUninitializedBuffer(
       const Shape& shape, PjRtDevice* device) override;
-  StatusOr<std::unique_ptr<PjRtBuffer>> CreateUninitializedBuffer(
+  absl::StatusOr<std::unique_ptr<PjRtBuffer>> CreateUninitializedBuffer(
       const Shape& shape, PjRtDevice* device,
       std::shared_ptr<BufferSequencingEvent> definition_event);
 
-  StatusOr<std::unique_ptr<PjRtBuffer>> CreateErrorBuffer(
-      Status error, const Shape& shape, PjRtMemorySpace* memory) override;
+  absl::StatusOr<std::unique_ptr<PjRtBuffer>> CreateErrorBuffer(
+      absl::Status error, const Shape& shape, PjRtMemorySpace* memory) override;
 
-  StatusOr<std::unique_ptr<PjRtClient::AsyncHostToDeviceTransferManager>>
+  absl::StatusOr<std::unique_ptr<PjRtClient::AsyncHostToDeviceTransferManager>>
   CreateBuffersForAsyncHostToDevice(absl::Span<const Shape> shapes,
                                     PjRtDevice* device) override {
     return Unimplemented("Async transfer to buffers not implemented");
   };
 
-  StatusOr<std::unique_ptr<PjRtClient::AsyncHostToDeviceTransferManager>>
+  absl::StatusOr<std::unique_ptr<PjRtClient::AsyncHostToDeviceTransferManager>>
   CreateBuffersForAsyncHostToDevice(absl::Span<const Shape> shapes,
                                     PjRtMemorySpace* memory_space) override {
     return Unimplemented("Async transfer to buffers not implemented");
   };
 
-  StatusOr<std::unique_ptr<PjRtBuffer>> BufferFromHostBuffer(
+  absl::StatusOr<std::unique_ptr<PjRtBuffer>> BufferFromHostBuffer(
       const void* data, PrimitiveType type, absl::Span<int64_t const> dims,
       std::optional<absl::Span<int64_t const>> byte_strides,
       HostBufferSemantics host_buffer_semantics,
       absl::AnyInvocable<void() &&> on_done_with_host_buffer,
       PjRtDevice* device, const Layout* device_layout) override;
 
-  StatusOr<std::unique_ptr<PjRtBuffer>> BufferFromHostBuffer(
+  absl::StatusOr<std::unique_ptr<PjRtBuffer>> BufferFromHostBuffer(
       const void* data, PrimitiveType type, absl::Span<int64_t const> dims,
       std::optional<absl::Span<int64_t const>> byte_strides,
       HostBufferSemantics host_buffer_semantics,
@@ -380,39 +382,39 @@ class PjRtStreamExecutorClient : public PjRtClient {
       absl::AnyInvocable<void() &&> on_done_with_host_buffer,
       PjRtMemorySpace* memory_space, const Layout* device_layout) override;
 
-  StatusOr<std::unique_ptr<PjRtBuffer>> BufferFromHostLiteral(
+  absl::StatusOr<std::unique_ptr<PjRtBuffer>> BufferFromHostLiteral(
       const LiteralSlice& literal, PjRtDevice* device) override;
 
   absl::StatusOr<std::unique_ptr<PjRtBuffer>> BufferFromHostLiteral(
       const LiteralSlice& literal, PjRtMemorySpace* memory_space) override;
 
-  StatusOr<std::vector<std::unique_ptr<PjRtBuffer>>>
+  absl::StatusOr<std::vector<std::unique_ptr<PjRtBuffer>>>
   MakeCrossHostReceiveBuffers(absl::Span<const Shape> shapes,
                               PjRtDevice* device,
                               PjRtCrossHostRecvNotifier notifier) override;
 
-  StatusOr<std::vector<std::unique_ptr<PjRtBuffer>>>
+  absl::StatusOr<std::vector<std::unique_ptr<PjRtBuffer>>>
   MakeCrossHostReceiveBuffersForGather(
       absl::Span<const Shape> shapes, std::vector<GatherDetails> gather_details,
       PjRtDevice* device, PjRtCrossHostRecvNotifier notifier) override;
 
-  StatusOr<std::unique_ptr<PjRtBuffer>> CreateViewOfDeviceBuffer(
+  absl::StatusOr<std::unique_ptr<PjRtBuffer>> CreateViewOfDeviceBuffer(
       void* device_ptr, const Shape& shape, PjRtDevice* device,
       std::function<void()> on_delete_callback,
       std::optional<std::intptr_t> stream) override;
 
-  StatusOr<ChannelHandle> CreateChannelHandle() override {
+  absl::StatusOr<ChannelHandle> CreateChannelHandle() override {
     return client()->CreateChannelHandle();
   }
-  StatusOr<ChannelHandle> CreateDeviceToHostChannelHandle() override {
+  absl::StatusOr<ChannelHandle> CreateDeviceToHostChannelHandle() override {
     return client()->CreateDeviceToHostChannelHandle();
   }
-  StatusOr<ChannelHandle> CreateHostToDeviceChannelHandle() override {
+  absl::StatusOr<ChannelHandle> CreateHostToDeviceChannelHandle() override {
     return client()->CreateHostToDeviceChannelHandle();
   }
 
   // TODO(zhangqiaorjc): Experimental. Will be removed.
-  Status Defragment() override {
+  absl::Status Defragment() override {
     return Unimplemented("Defragment not implemented");
   }
 
@@ -439,7 +441,7 @@ class PjRtStreamExecutorClient : public PjRtClient {
  protected:
   friend class PjRtStreamExecutorBuffer;
 
-  virtual Status EnqueueCrossHostReceive(
+  virtual absl::Status EnqueueCrossHostReceive(
       absl::Span<const std::unique_ptr<PjRtBuffer>> buffers,
       std::shared_ptr<BufferSequencingEvent> definition_event,
       PjRtCrossHostRecvNotifier notifier,
@@ -479,7 +481,7 @@ class PjRtStreamExecutorClient : public PjRtClient {
         addressable_device_logical_ids;
     std::vector<PjRtDevice*> addressable_devices;
   };
-  StatusOr<ExecutableExtras> GetExecutableExtras(CompileOptions* options);
+  absl::StatusOr<ExecutableExtras> GetExecutableExtras(CompileOptions* options);
 
   const PjRtPlatformId platform_id_;
   const std::string platform_name_;
@@ -524,7 +526,7 @@ class PjRtStreamExecutorClient : public PjRtClient {
 
 // Converts a 2D set of Device objects indexed by [replica][partition] into an
 // xla::DeviceAssignment.
-StatusOr<DeviceAssignment> DevicesToDeviceAssignment(
+absl::StatusOr<DeviceAssignment> DevicesToDeviceAssignment(
     absl::Span<const std::vector<PjRtDevice*>> devices);
 
 class PjRtStreamExecutorBuffer : public PjRtBuffer {
@@ -578,9 +580,10 @@ class PjRtStreamExecutorBuffer : public PjRtBuffer {
   class ScopedHold {
    public:
     enum Type { kUsage = 0, kExternalReference, kDonation, kMaxValue };
-    // Use a State enum instead of encoding the state in an error Status to
-    // avoid creating Status values in non-error cases. Creating a Status
-    // entails several allocations and can add O(us) to every use of a hold.
+    // Use a State enum instead of encoding the state in an error absl::Status
+    // to avoid creating absl::Status values in non-error cases. Creating a
+    // absl::Status entails several allocations and can add O(us) to every use
+    // of a hold.
     enum State {
       kUninitialized = 0,
       kValid,
@@ -598,8 +601,8 @@ class PjRtStreamExecutorBuffer : public PjRtBuffer {
 
     Type type() const { return type_; }
 
-    Status status() const {
-      // Lazily create Status values only when they are requested.
+    absl::Status status() const {
+      // Lazily create absl::Status values only when they are requested.
       switch (state_) {
         case kUninitialized:
           return InvalidArgument("Buffer has not been initialized");
@@ -667,8 +670,9 @@ class PjRtStreamExecutorBuffer : public PjRtBuffer {
 
     // Helper struct that makes it possible to move a ScopedHold through a
     // closure.
-    using ForClosure = std::tuple<PjRtStreamExecutorBuffer*, Type, State,
-                                  Status, std::shared_ptr<TrackedDeviceBuffer>>;
+    using ForClosure =
+        std::tuple<PjRtStreamExecutorBuffer*, Type, State, absl::Status,
+                   std::shared_ptr<TrackedDeviceBuffer>>;
 
     ScopedHold(PjRtStreamExecutorBuffer* parent, Type type)
         : parent_(parent), type_(type), state_(kUninitialized) {}
@@ -699,7 +703,7 @@ class PjRtStreamExecutorBuffer : public PjRtBuffer {
     // There is an invariant that if ok() then
     // buffer_.value() != nullptr.
     State state_;
-    Status status_;
+    absl::Status status_;
     std::shared_ptr<TrackedDeviceBuffer> buffer_;
   };
 
@@ -715,7 +719,7 @@ class PjRtStreamExecutorBuffer : public PjRtBuffer {
   PjRtStreamExecutorBuffer& operator=(PjRtStreamExecutorBuffer&&) = delete;
 
   const Shape& on_device_shape() const override { return on_device_shape_; }
-  StatusOr<Shape> logical_on_device_shape() override;
+  absl::StatusOr<Shape> logical_on_device_shape() override;
   PjRtMemorySpace* memory_space() const override { return memory_space_; }
   PjRtStreamExecutorDevice* device() const override { return device_; }
   PjRtPlatformId platform_id() const { return client_->platform_id(); }
@@ -726,11 +730,11 @@ class PjRtStreamExecutorBuffer : public PjRtBuffer {
            on_device_shape_.tuple_shapes_size() == 0;
   }
 
-  StatusOr<std::unique_ptr<ExternalReference>> AcquireExternalReference()
+  absl::StatusOr<std::unique_ptr<ExternalReference>> AcquireExternalReference()
       override;
 
-  StatusOr<std::unique_ptr<ExternalReference>> ReleaseDeviceMemoryOwnership(
-      bool wait_for_operations_to_complete) override;
+  absl::StatusOr<std::unique_ptr<ExternalReference>>
+  ReleaseDeviceMemoryOwnership(bool wait_for_operations_to_complete) override;
 
   using PjRtBuffer::ToLiteralSync;
   PjRtFuture<> ToLiteral(MutableLiteralBase* literal) override;
@@ -738,7 +742,7 @@ class PjRtStreamExecutorBuffer : public PjRtBuffer {
       absl::AnyInvocable<absl::StatusOr<MutableLiteralBase*>() &&> generator)
       override;
 
-  StatusOr<size_t> GetOnDeviceSizeInBytes() const override;
+  absl::StatusOr<size_t> GetOnDeviceSizeInBytes() const override;
 
   PjRtFuture<> CopyRawToHost(void* dst, int64_t offset,
                              int64_t transfer_size) override;
@@ -761,7 +765,7 @@ class PjRtStreamExecutorBuffer : public PjRtBuffer {
 
   // Returns a view of the PjRtBuffer device memory as a ShapedBuffer. The
   // PjRtBuffer retains ownership of the device buffers.
-  StatusOr<ShapedBuffer> AsShapedBuffer() const;
+  absl::StatusOr<ShapedBuffer> AsShapedBuffer() const;
 
   // Returns a hold on the TrackedDeviceBuffer holding the device
   // buffers. See comment on ScopedHold.
@@ -773,10 +777,10 @@ class PjRtStreamExecutorBuffer : public PjRtBuffer {
     return GetBufferWithHold(ScopedHold::kExternalReference);
   }
 
-  StatusOr<std::unique_ptr<PjRtBuffer>> CopyToDevice(
+  absl::StatusOr<std::unique_ptr<PjRtBuffer>> CopyToDevice(
       PjRtDevice* dst_device) override;
 
-  StatusOr<std::unique_ptr<PjRtBuffer>> CopyToMemorySpace(
+  absl::StatusOr<std::unique_ptr<PjRtBuffer>> CopyToMemorySpace(
       PjRtMemorySpace* dst_memory_space) override;
 
   void CopyToRemoteDevice(PjRtFuture<std::string> serialized_descriptor,
@@ -805,7 +809,7 @@ class PjRtStreamExecutorBuffer : public PjRtBuffer {
   // If the buffer was shared via an external reference it is the client's
   // responsibility that accesses via that reference do not interfere with
   // accesses via the buffer returned from Release.
-  StatusOr<std::shared_ptr<TrackedDeviceBuffer>> Release(
+  absl::StatusOr<std::shared_ptr<TrackedDeviceBuffer>> Release(
       bool wait_for_operations_to_complete);
 
   absl::StatusOr<std::unique_ptr<PjRtBuffer>> DonateWithControlDependency(
@@ -825,7 +829,7 @@ class PjRtStreamExecutorBuffer : public PjRtBuffer {
   // an outstanding external hold.
   // Requires holds_[kDonation] == 0 (i.e., WaitForOutstandingDonationHolds()
   // must be called first.)
-  StatusOr<std::shared_ptr<TrackedDeviceBuffer>> GetBufferForHoldLocked(
+  absl::StatusOr<std::shared_ptr<TrackedDeviceBuffer>> GetBufferForHoldLocked(
       ScopedHold::Type type) ABSL_EXCLUSIVE_LOCKS_REQUIRED(mu_);
 
   // Adds a hold of hold->type() and initializes `hold` with device_buffer_.
@@ -851,8 +855,8 @@ class PjRtStreamExecutorBuffer : public PjRtBuffer {
   // buffer==device_buffer_ or device_buffer_==nullptr.
   void DropHold(ScopedHold::Type type, TrackedDeviceBuffer* buffer);
 
-  StatusOr<std::pair<std::unique_ptr<PjRtBuffer>,
-                     std::shared_ptr<BufferSequencingEvent>>>
+  absl::StatusOr<std::pair<std::unique_ptr<PjRtBuffer>,
+                           std::shared_ptr<BufferSequencingEvent>>>
   CopyToDeviceHelper(PjRtDevice* dst_device, LocalDeviceState* dst_local_device,
                      LocalDeviceState* transfer_local_device,
                      LocalDeviceState* src_local_device,
@@ -906,7 +910,7 @@ class PjRtStreamExecutorLoadedExecutable : public PjRtLoadedExecutable {
     return size;
   }
 
-  StatusOr<CompiledMemoryStats> GetCompiledMemoryStats() const override {
+  absl::StatusOr<CompiledMemoryStats> GetCompiledMemoryStats() const override {
     if (executables_.size() != 1) {
       return Unimplemented(
           "Retrieving CompiledMemoryStats is not supported for multiple "
@@ -937,26 +941,26 @@ class PjRtStreamExecutorLoadedExecutable : public PjRtLoadedExecutable {
   }
 
   // Return an HloModule per partition.
-  StatusOr<std::vector<std::shared_ptr<HloModule>>> GetHloModules()
+  absl::StatusOr<std::vector<std::shared_ptr<HloModule>>> GetHloModules()
       const override;
 
-  StatusOr<std::vector<std::vector<absl::string_view>>> GetOutputMemoryKinds()
-      const override;
+  absl::StatusOr<std::vector<std::vector<absl::string_view>>>
+  GetOutputMemoryKinds() const override;
 
   using PjRtLoadedExecutable::Execute;
-  StatusOr<std::vector<std::vector<std::unique_ptr<PjRtBuffer>>>> Execute(
+  absl::StatusOr<std::vector<std::vector<std::unique_ptr<PjRtBuffer>>>> Execute(
       absl::Span<const std::vector<PjRtBuffer*>> argument_handles,
       const ExecuteOptions& options,
       std::optional<std::vector<PjRtFuture<>>>& returned_futures) override;
 
   using PjRtLoadedExecutable::ExecuteSharded;
-  StatusOr<std::vector<std::unique_ptr<PjRtBuffer>>> ExecuteSharded(
+  absl::StatusOr<std::vector<std::unique_ptr<PjRtBuffer>>> ExecuteSharded(
       absl::Span<PjRtBuffer* const> argument_handles, PjRtDevice* device,
       const ExecuteOptions& options,
       std::optional<PjRtFuture<>>& returned_future, bool fill_future) override;
 
   using PjRtLoadedExecutable::ExecutePortable;
-  StatusOr<std::vector<std::unique_ptr<PjRtBuffer>>> ExecutePortable(
+  absl::StatusOr<std::vector<std::unique_ptr<PjRtBuffer>>> ExecutePortable(
       absl::Span<PjRtBuffer* const> argument_handles, PjRtDevice* device,
       const ExecuteOptions& options,
       std::optional<PjRtFuture<>>& returned_future, bool fill_future) override;
@@ -965,7 +969,7 @@ class PjRtStreamExecutorLoadedExecutable : public PjRtLoadedExecutable {
 
   bool IsDeleted() override { return executables_.empty(); }
 
-  StatusOr<std::string> SerializeExecutable() const override {
+  absl::StatusOr<std::string> SerializeExecutable() const override {
     return client_->SerializeExecutable(*this);
   }
 
@@ -993,14 +997,14 @@ class PjRtStreamExecutorLoadedExecutable : public PjRtLoadedExecutable {
   friend class StreamExecutorGpuClient;
   // Initializes information about which arguments to which executables must be
   // donated due to aliases that were specified by the computation.
-  Status SetUpDonation(bool tuple_inputs);
+  absl::Status SetUpDonation(bool tuple_inputs);
 
   // Returns a sorted list of the parameters that must be donated. Derived
   // classes may use custom logic.
   virtual absl::Span<int const> ParametersThatMustBeDonated(
       int executable_idx) const;
 
-  virtual StatusOr<std::vector<ExecutionInput>>
+  virtual absl::StatusOr<std::vector<ExecutionInput>>
   MakeExecutionInputsAndWaitForEvents(
       int device_ordinal, const ExecuteOptions& options,
       absl::Span<const Shape> executable_parameter_shapes,
@@ -1008,7 +1012,7 @@ class PjRtStreamExecutorLoadedExecutable : public PjRtLoadedExecutable {
       absl::Span<const PjRtStreamExecutorBuffer::ScopedHold> device_buffers,
       absl::flat_hash_set<BufferSequencingEvent*>& events) const;
 
-  StatusOr<ScopedShapedBuffer> EnqueueExecution(
+  absl::StatusOr<ScopedShapedBuffer> EnqueueExecution(
       absl::Span<PjRtBuffer* const> argument_handles, int replica,
       int partition, int executable_idx, const RunId& run_id,
       const ExecuteOptions& options, PjRtDevice* device,
@@ -1024,12 +1028,10 @@ class PjRtStreamExecutorLoadedExecutable : public PjRtLoadedExecutable {
       std::vector<std::shared_ptr<TrackedDeviceBuffer>>& buffers_to_release)
       const;
 
-  StatusOr<Result> ExecuteHelper(absl::Span<PjRtBuffer* const> argument_handles,
-                                 int replica, int partition,
-                                 const RunId& run_id,
-                                 const ExecuteOptions& options,
-                                 bool fill_future,
-                                 PjRtDevice* device = nullptr) const;
+  absl::StatusOr<Result> ExecuteHelper(
+      absl::Span<PjRtBuffer* const> argument_handles, int replica,
+      int partition, const RunId& run_id, const ExecuteOptions& options,
+      bool fill_future, PjRtDevice* device = nullptr) const;
 
   // Create shared pointers so we can free them after the execution: with
   // asynchronous execution, the process being executed can outlive the
diff --git a/third_party/xla/xla/pjrt/pjrt_stream_executor_client_test.cc b/third_party/xla/xla/pjrt/pjrt_stream_executor_client_test.cc
index 7d494bcbcf7464..26df8698348323 100644
--- a/third_party/xla/xla/pjrt/pjrt_stream_executor_client_test.cc
+++ b/third_party/xla/xla/pjrt/pjrt_stream_executor_client_test.cc
@@ -70,7 +70,7 @@ absl::StatusOr<std::unique_ptr<PjRtStreamExecutorClient>> GetClient() {
       /*gpu_run_options=*/nullptr);
 }
 
-StatusOr<std::unique_ptr<PjRtLoadedExecutable>> ToyExecutable(
+absl::StatusOr<std::unique_ptr<PjRtLoadedExecutable>> ToyExecutable(
     PjRtStreamExecutorClient& client, Shape shape,
     absl::AnyInvocable<void(XlaBuilder&)> set_up_aliases) {
   CompileOptions compile_options;
@@ -88,7 +88,7 @@ StatusOr<std::unique_ptr<PjRtLoadedExecutable>> ToyExecutable(
   return executable;
 }
 
-Status ExecuteWithSameInputBuffer(
+absl::Status ExecuteWithSameInputBuffer(
     absl::AnyInvocable<void(XlaBuilder&)> set_up_aliases) {
   auto shape = xla::ShapeUtil::MakeScalarShape(xla::F32);
   TF_ASSIGN_OR_RETURN(auto client, GetClient());
diff --git a/third_party/xla/xla/pjrt/status_casters.h b/third_party/xla/xla/pjrt/status_casters.h
index bd4a044f6c3042..7e45a310dfa5ba 100644
--- a/third_party/xla/xla/pjrt/status_casters.h
+++ b/third_party/xla/xla/pjrt/status_casters.h
@@ -25,7 +25,7 @@ namespace xla {
 
 // C++ -> Python caster helpers.
 //
-// Failing statuses become Python exceptions; OK Status() becomes None.
+// Failing statuses become Python exceptions; OK absl::Status() becomes None.
 //
 // Given there can be only a single global pybind11 type_caster for the
 // `absl::Status` type, and given XLA wants a custom exception being raised,
diff --git a/third_party/xla/xla/pjrt/stream_executor_executable.cc b/third_party/xla/xla/pjrt/stream_executor_executable.cc
index d706de1b9be6c8..766cd0d2147a07 100644
--- a/third_party/xla/xla/pjrt/stream_executor_executable.cc
+++ b/third_party/xla/xla/pjrt/stream_executor_executable.cc
@@ -26,7 +26,8 @@ limitations under the License.
 #include "tsl/platform/statusor.h"
 
 namespace xla {
-StatusOr<std::string> StreamExecutorExecutable::SerializeExecutable() const {
+absl::StatusOr<std::string> StreamExecutorExecutable::SerializeExecutable()
+    const {
   if (aot_executables_.empty()) {
     return absl::InternalError("No local executable");
   }
diff --git a/third_party/xla/xla/pjrt/tf_pjrt_client.cc b/third_party/xla/xla/pjrt/tf_pjrt_client.cc
index 816a6b8b3dbedd..c35bee43e91bcc 100644
--- a/third_party/xla/xla/pjrt/tf_pjrt_client.cc
+++ b/third_party/xla/xla/pjrt/tf_pjrt_client.cc
@@ -38,7 +38,7 @@ TfPjRtBuffer::~TfPjRtBuffer() { client_->UntrackBuffer(this); }
 PjRtClient* TfPjRtBuffer::client() const { return client_; }
 PjRtClient* TfPjRtExecutable::client() const { return client_; }
 
-StatusOr<std::unique_ptr<PjRtBuffer>> TfPjRtBuffer::CopyToDevice(
+absl::StatusOr<std::unique_ptr<PjRtBuffer>> TfPjRtBuffer::CopyToDevice(
     PjRtDevice* dst_device) {
   TF_ASSIGN_OR_RETURN(std::unique_ptr<PjRtBuffer> result,
                       wrapped_->CopyToDevice(dst_device));
@@ -50,7 +50,7 @@ TfPjRtExecutable::TfPjRtExecutable(
     TfPjRtClient* client, std::unique_ptr<PjRtLoadedExecutable> wrapped)
     : client_(client), wrapped_(std::move(wrapped)) {}
 
-StatusOr<std::vector<std::vector<std::unique_ptr<PjRtBuffer>>>>
+absl::StatusOr<std::vector<std::vector<std::unique_ptr<PjRtBuffer>>>>
 TfPjRtExecutable::Execute(
     absl::Span<const std::vector<PjRtBuffer*>> argument_handles,
     const ExecuteOptions& options,
@@ -76,7 +76,7 @@ TfPjRtExecutable::Execute(
   return out;
 }
 
-StatusOr<std::vector<std::unique_ptr<PjRtBuffer>>>
+absl::StatusOr<std::vector<std::unique_ptr<PjRtBuffer>>>
 TfPjRtExecutable::ExecuteSharded(absl::Span<PjRtBuffer* const> argument_handles,
                                  PjRtDevice* device,
                                  const ExecuteOptions& options,
@@ -96,7 +96,7 @@ TfPjRtExecutable::ExecuteSharded(absl::Span<PjRtBuffer* const> argument_handles,
   }
   return out;
 }
-StatusOr<std::vector<std::unique_ptr<PjRtBuffer>>>
+absl::StatusOr<std::vector<std::unique_ptr<PjRtBuffer>>>
 TfPjRtExecutable::ExecutePortable(
     absl::Span<PjRtBuffer* const> argument_handles, PjRtDevice* device,
     const ExecuteOptions& options, std::optional<PjRtFuture<>>& returned_future,
@@ -129,15 +129,16 @@ TfPjRtClient::TfPjRtClient(std::unique_ptr<PjRtClient> wrapped)
 
 TfPjRtClient::~TfPjRtClient() { LOG(INFO) << "TfPjRtClient destroyed."; }
 
-StatusOr<std::unique_ptr<PjRtBuffer>> TfPjRtClient::WrapBuffer(
-    StatusOr<std::unique_ptr<PjRtBuffer>> to_wrap) {
+absl::StatusOr<std::unique_ptr<PjRtBuffer>> TfPjRtClient::WrapBuffer(
+    absl::StatusOr<std::unique_ptr<PjRtBuffer>> to_wrap) {
   TF_ASSIGN_OR_RETURN(std::unique_ptr<PjRtBuffer> buffer, std::move(to_wrap));
   return std::unique_ptr<PjRtBuffer>(
       std::make_unique<TfPjRtBuffer>(this, std::move(buffer)));
 }
 
-StatusOr<std::unique_ptr<PjRtLoadedExecutable>> TfPjRtClient::WrapExecutable(
-    StatusOr<std::unique_ptr<PjRtLoadedExecutable>> to_wrap) {
+absl::StatusOr<std::unique_ptr<PjRtLoadedExecutable>>
+TfPjRtClient::WrapExecutable(
+    absl::StatusOr<std::unique_ptr<PjRtLoadedExecutable>> to_wrap) {
   TF_ASSIGN_OR_RETURN(std::unique_ptr<PjRtLoadedExecutable> executable,
                       std::move(to_wrap));
   return std::unique_ptr<PjRtLoadedExecutable>(
diff --git a/third_party/xla/xla/pjrt/tf_pjrt_client.h b/third_party/xla/xla/pjrt/tf_pjrt_client.h
index 510a7141c98e36..d1dfd3eace20e7 100644
--- a/third_party/xla/xla/pjrt/tf_pjrt_client.h
+++ b/third_party/xla/xla/pjrt/tf_pjrt_client.h
@@ -45,7 +45,7 @@ class TfPjRtBuffer : public PjRtBuffer {
   const Shape& on_device_shape() const override {
     return wrapped_->on_device_shape();
   }
-  StatusOr<Shape> logical_on_device_shape() override {
+  absl::StatusOr<Shape> logical_on_device_shape() override {
     return wrapped_->logical_on_device_shape();
   }
   PjRtMemorySpace* memory_space() const override {
@@ -53,7 +53,7 @@ class TfPjRtBuffer : public PjRtBuffer {
   }
   PjRtDevice* device() const override { return wrapped_->device(); }
   PjRtClient* client() const override;
-  StatusOr<std::unique_ptr<ExternalReference>> AcquireExternalReference()
+  absl::StatusOr<std::unique_ptr<ExternalReference>> AcquireExternalReference()
       override {
     return wrapped_->AcquireExternalReference();
   }
@@ -65,7 +65,7 @@ class TfPjRtBuffer : public PjRtBuffer {
       override {
     return wrapped_->LazyToLiteral(std::move(generator));
   }
-  StatusOr<size_t> GetOnDeviceSizeInBytes() const override {
+  absl::StatusOr<size_t> GetOnDeviceSizeInBytes() const override {
     return wrapped_->GetOnDeviceSizeInBytes();
   }
   PjRtFuture<> CopyRawToHost(void* dst, int64_t offset,
@@ -73,15 +73,15 @@ class TfPjRtBuffer : public PjRtBuffer {
     return wrapped_->CopyRawToHost(dst, offset, transfer_size);
   }
   void Delete() override { wrapped_->Delete(); }
-  StatusOr<std::unique_ptr<ExternalReference>> ReleaseDeviceMemoryOwnership(
-      bool wait_for_operations_to_complete) override {
+  absl::StatusOr<std::unique_ptr<ExternalReference>>
+  ReleaseDeviceMemoryOwnership(bool wait_for_operations_to_complete) override {
     return wrapped_->ReleaseDeviceMemoryOwnership(
         wait_for_operations_to_complete);
   }
   bool IsDeleted() override { return wrapped_->IsDeleted(); }
-  StatusOr<std::unique_ptr<PjRtBuffer>> CopyToDevice(
+  absl::StatusOr<std::unique_ptr<PjRtBuffer>> CopyToDevice(
       PjRtDevice* dst_device) override;
-  StatusOr<std::unique_ptr<PjRtBuffer>> CopyToMemorySpace(
+  absl::StatusOr<std::unique_ptr<PjRtBuffer>> CopyToMemorySpace(
       PjRtMemorySpace* dst_memory_space) override {
     return Unimplemented("CopyToMemorySpace not implemented");
   }
@@ -138,26 +138,26 @@ class TfPjRtExecutable : public PjRtLoadedExecutable {
   absl::Span<PjRtDevice* const> addressable_devices() const override {
     return wrapped_->addressable_devices();
   }
-  StatusOr<std::vector<std::shared_ptr<HloModule>>> GetHloModules()
+  absl::StatusOr<std::vector<std::shared_ptr<HloModule>>> GetHloModules()
       const override {
     return wrapped_->GetHloModules();
   }
-  StatusOr<std::vector<std::vector<absl::string_view>>> GetOutputMemoryKinds()
-      const override {
+  absl::StatusOr<std::vector<std::vector<absl::string_view>>>
+  GetOutputMemoryKinds() const override {
     return wrapped_->GetOutputMemoryKinds();
   }
   using PjRtLoadedExecutable::Execute;
-  StatusOr<std::vector<std::vector<std::unique_ptr<PjRtBuffer>>>> Execute(
+  absl::StatusOr<std::vector<std::vector<std::unique_ptr<PjRtBuffer>>>> Execute(
       absl::Span<const std::vector<PjRtBuffer*>> argument_handles,
       const ExecuteOptions& options,
       std::optional<std::vector<PjRtFuture<>>>& returned_futures) override;
   using PjRtLoadedExecutable::ExecuteSharded;
-  StatusOr<std::vector<std::unique_ptr<PjRtBuffer>>> ExecuteSharded(
+  absl::StatusOr<std::vector<std::unique_ptr<PjRtBuffer>>> ExecuteSharded(
       absl::Span<PjRtBuffer* const> argument_handles, PjRtDevice* device,
       const ExecuteOptions& options,
       std::optional<PjRtFuture<>>& returned_future, bool fill_future) override;
   using PjRtLoadedExecutable::ExecutePortable;
-  StatusOr<std::vector<std::unique_ptr<PjRtBuffer>>> ExecutePortable(
+  absl::StatusOr<std::vector<std::unique_ptr<PjRtBuffer>>> ExecutePortable(
       absl::Span<PjRtBuffer* const> argument_handles, PjRtDevice* device,
       const ExecuteOptions& options,
       std::optional<PjRtFuture<>>& returned_future, bool fill_future) override;
@@ -168,15 +168,15 @@ class TfPjRtExecutable : public PjRtLoadedExecutable {
     return wrapped_->IsReturnedFutureSupported();
   }
 
-  StatusOr<std::string> SerializeExecutable() const override {
+  absl::StatusOr<std::string> SerializeExecutable() const override {
     return wrapped_->SerializeExecutable();
   }
 
-  StatusOr<struct CompileOptions> GetCompileOptions() const override {
+  absl::StatusOr<struct CompileOptions> GetCompileOptions() const override {
     return wrapped_->GetCompileOptions();
   }
 
-  StatusOr<std::string> FingerprintExecutable() const override {
+  absl::StatusOr<std::string> FingerprintExecutable() const override {
     return wrapped_->FingerprintExecutable();
   }
 
@@ -204,15 +204,15 @@ class TfPjRtClient : public PjRtClient {
   absl::Span<PjRtDevice* const> addressable_devices() const override {
     return wrapped_->addressable_devices();
   }
-  StatusOr<PjRtDevice*> LookupDevice(
+  absl::StatusOr<PjRtDevice*> LookupDevice(
       PjRtGlobalDeviceId global_device_id) const override {
     return wrapped_->LookupDevice(global_device_id);
   }
-  StatusOr<PjRtDevice*> LookupAddressableDevice(
+  absl::StatusOr<PjRtDevice*> LookupAddressableDevice(
       int local_hardware_id) const override {
     return LookupAddressableDevice(PjRtLocalDeviceId(local_hardware_id));
   }
-  StatusOr<PjRtDevice*> LookupAddressableDevice(
+  absl::StatusOr<PjRtDevice*> LookupAddressableDevice(
       PjRtLocalDeviceId local_device_id) const override {
     if (wrapped_ == nullptr) {
       return tsl::errors::Internal(
@@ -235,51 +235,51 @@ class TfPjRtClient : public PjRtClient {
   PjRtRuntimeType runtime_type() const override {
     return wrapped_->runtime_type();
   }
-  StatusOr<DeviceAssignment> GetDefaultDeviceAssignment(
+  absl::StatusOr<DeviceAssignment> GetDefaultDeviceAssignment(
       int num_replicas, int num_partitions) const override {
     return wrapped_->GetDefaultDeviceAssignment(num_replicas, num_partitions);
   }
-  StatusOr<Layout> GetDefaultLayout(PrimitiveType element_type,
-                                    absl::Span<const int64_t> dims) override {
+  absl::StatusOr<Layout> GetDefaultLayout(
+      PrimitiveType element_type, absl::Span<const int64_t> dims) override {
     return wrapped_->GetDefaultLayout(element_type, dims);
   }
-  StatusOr<std::unique_ptr<HloCostAnalysis>> GetHloCostAnalysis()
+  absl::StatusOr<std::unique_ptr<HloCostAnalysis>> GetHloCostAnalysis()
       const override {
     return wrapped_->GetHloCostAnalysis();
   }
-  StatusOr<std::unique_ptr<PjRtLoadedExecutable>> Compile(
+  absl::StatusOr<std::unique_ptr<PjRtLoadedExecutable>> Compile(
       const XlaComputation& computation, CompileOptions options) override {
     return WrapExecutable(wrapped_->Compile(computation, options));
   }
-  StatusOr<std::unique_ptr<PjRtLoadedExecutable>> Compile(
+  absl::StatusOr<std::unique_ptr<PjRtLoadedExecutable>> Compile(
       mlir::ModuleOp module, CompileOptions options) override {
     return WrapExecutable(wrapped_->Compile(std::move(module), options));
   }
 
-  StatusOr<std::unique_ptr<PjRtLoadedExecutable>> DeserializeExecutable(
+  absl::StatusOr<std::unique_ptr<PjRtLoadedExecutable>> DeserializeExecutable(
       absl::string_view serialized,
       std::optional<CompileOptions> options) override {
     return WrapExecutable(wrapped_->DeserializeExecutable(serialized, options));
   }
 
-  StatusOr<std::unique_ptr<PjRtBuffer>> CreateUninitializedBuffer(
+  absl::StatusOr<std::unique_ptr<PjRtBuffer>> CreateUninitializedBuffer(
       const Shape& shape, PjRtDevice* device) override {
     return Unimplemented(
         "CreateUninitializedBuffer not supported for TfPjRtClient.");
   }
-  StatusOr<std::unique_ptr<AsyncHostToDeviceTransferManager>>
+  absl::StatusOr<std::unique_ptr<AsyncHostToDeviceTransferManager>>
   CreateBuffersForAsyncHostToDevice(absl::Span<const Shape> shapes,
                                     PjRtDevice* device) override {
     return Unimplemented(
         "AsyncHostToDeviceTransferManager not supported for Tf.");
   }
-  StatusOr<std::unique_ptr<AsyncHostToDeviceTransferManager>>
+  absl::StatusOr<std::unique_ptr<AsyncHostToDeviceTransferManager>>
   CreateBuffersForAsyncHostToDevice(absl::Span<const Shape> shapes,
                                     PjRtMemorySpace* memory_space) override {
     return Unimplemented(
         "AsyncHostToDeviceTransferManager not supported for Tf.");
   }
-  StatusOr<std::unique_ptr<PjRtBuffer>> BufferFromHostBuffer(
+  absl::StatusOr<std::unique_ptr<PjRtBuffer>> BufferFromHostBuffer(
       const void* data, PrimitiveType type, absl::Span<int64_t const> dims,
       std::optional<absl::Span<int64_t const>> byte_strides,
       HostBufferSemantics host_buffer_semantics,
@@ -289,7 +289,7 @@ class TfPjRtClient : public PjRtClient {
         data, type, dims, byte_strides, host_buffer_semantics,
         std::move(on_done_with_host_buffer), device));
   }
-  StatusOr<std::unique_ptr<PjRtBuffer>> BufferFromHostBuffer(
+  absl::StatusOr<std::unique_ptr<PjRtBuffer>> BufferFromHostBuffer(
       const void* data, PrimitiveType type, absl::Span<int64_t const> dims,
       std::optional<absl::Span<int64_t const>> byte_strides,
       HostBufferSemantics host_buffer_semantics,
@@ -299,53 +299,54 @@ class TfPjRtClient : public PjRtClient {
         data, type, dims, byte_strides, host_buffer_semantics,
         std::move(on_done_with_host_buffer), device, device_layout));
   }
-  StatusOr<std::unique_ptr<PjRtBuffer>> BufferFromHostLiteral(
+  absl::StatusOr<std::unique_ptr<PjRtBuffer>> BufferFromHostLiteral(
       const LiteralSlice& literal, PjRtDevice* device) override {
     return WrapBuffer(wrapped_->BufferFromHostLiteral(literal, device));
   }
-  StatusOr<std::unique_ptr<PjRtBuffer>> CreateViewOfDeviceBuffer(
+  absl::StatusOr<std::unique_ptr<PjRtBuffer>> CreateViewOfDeviceBuffer(
       void* device_ptr, const Shape& shape, PjRtDevice* device,
       std::function<void()> on_delete_callback,
       std::optional<std::intptr_t> stream) override {
     return WrapBuffer(wrapped_->CreateViewOfDeviceBuffer(
         device_ptr, shape, device, on_delete_callback, stream));
   }
-  StatusOr<std::uintptr_t> UnsafeBufferPointer(PjRtBuffer* buffer) override {
+  absl::StatusOr<std::uintptr_t> UnsafeBufferPointer(
+      PjRtBuffer* buffer) override {
     return wrapped_->UnsafeBufferPointer(UnwrapBuffer(buffer));
   }
-  StatusOr<std::vector<std::unique_ptr<PjRtBuffer>>>
+  absl::StatusOr<std::vector<std::unique_ptr<PjRtBuffer>>>
   MakeCrossHostReceiveBuffers(absl::Span<const Shape> shapes,
                               PjRtDevice* device,
                               PjRtCrossHostRecvNotifier notifier) override {
     return wrapped_->MakeCrossHostReceiveBuffers(shapes, device,
                                                  std::move(notifier));
   }
-  StatusOr<std::vector<std::unique_ptr<PjRtBuffer>>>
+  absl::StatusOr<std::vector<std::unique_ptr<PjRtBuffer>>>
   MakeCrossHostReceiveBuffersForGather(
       absl::Span<const Shape> shapes, std::vector<GatherDetails> gather_details,
       PjRtDevice* device, PjRtCrossHostRecvNotifier notifier) override {
     return wrapped_->MakeCrossHostReceiveBuffersForGather(
         shapes, std::move(gather_details), device, std::move(notifier));
   }
-  StatusOr<ChannelHandle> CreateChannelHandle() override {
+  absl::StatusOr<ChannelHandle> CreateChannelHandle() override {
     return wrapped_->CreateChannelHandle();
   }
-  StatusOr<ChannelHandle> CreateDeviceToHostChannelHandle() override {
+  absl::StatusOr<ChannelHandle> CreateDeviceToHostChannelHandle() override {
     return wrapped_->CreateDeviceToHostChannelHandle();
   }
-  StatusOr<ChannelHandle> CreateHostToDeviceChannelHandle() override {
+  absl::StatusOr<ChannelHandle> CreateHostToDeviceChannelHandle() override {
     return wrapped_->CreateHostToDeviceChannelHandle();
   }
-  StatusOr<const PjRtTopologyDescription*> GetTopologyDescription()
+  absl::StatusOr<const PjRtTopologyDescription*> GetTopologyDescription()
       const override {
     return wrapped_->GetTopologyDescription();
   }
-  Status Defragment() override { return wrapped_->Defragment(); }
+  absl::Status Defragment() override { return wrapped_->Defragment(); }
 
   PjRtClient* wrapped() const { return wrapped_.get(); }
 
-  StatusOr<std::unique_ptr<PjRtBuffer>> WrapBuffer(
-      StatusOr<std::unique_ptr<PjRtBuffer>> to_wrap);
+  absl::StatusOr<std::unique_ptr<PjRtBuffer>> WrapBuffer(
+      absl::StatusOr<std::unique_ptr<PjRtBuffer>> to_wrap);
 
   // Tracks a non-owning pointer of TfPjRtBuffer in TfPjRtClient.
   void TrackBuffer(TfPjRtBuffer* buffer);
@@ -370,8 +371,8 @@ class TfPjRtClient : public PjRtClient {
                 ->wrapped();
   }
 
-  StatusOr<std::unique_ptr<PjRtLoadedExecutable>> WrapExecutable(
-      StatusOr<std::unique_ptr<PjRtLoadedExecutable>> to_wrap);
+  absl::StatusOr<std::unique_ptr<PjRtLoadedExecutable>> WrapExecutable(
+      absl::StatusOr<std::unique_ptr<PjRtLoadedExecutable>> to_wrap);
 
   std::unique_ptr<PjRtClient> wrapped_;
 
diff --git a/third_party/xla/xla/pjrt/tracked_device_buffer.cc b/third_party/xla/xla/pjrt/tracked_device_buffer.cc
index 2f7a20392da018..11db8c45b7ae27 100644
--- a/third_party/xla/xla/pjrt/tracked_device_buffer.cc
+++ b/third_party/xla/xla/pjrt/tracked_device_buffer.cc
@@ -88,7 +88,7 @@ void BufferSequencingEvent::WaitForEventOnStream(se::Stream* stream) {
   streams_defined_on_.push_back(stream);
 }
 
-Status BufferSequencingEvent::WaitForEventOnExternalStream(
+absl::Status BufferSequencingEvent::WaitForEventOnExternalStream(
     std::intptr_t stream) {
   absl::MutexLock lock(&mu_);
 
diff --git a/third_party/xla/xla/pjrt/tracked_device_buffer.h b/third_party/xla/xla/pjrt/tracked_device_buffer.h
index 379e4840fea6ea..3b04f44ecdfe98 100644
--- a/third_party/xla/xla/pjrt/tracked_device_buffer.h
+++ b/third_party/xla/xla/pjrt/tracked_device_buffer.h
@@ -84,7 +84,7 @@ class BufferSequencingEvent {
   // Same as WaitForEventOnStream, but takes a raw platform-specific
   // stream. Currently on implemented for CUDA and ROCM GPU, where stream is a
   // GpuStreamHandle (e.g. a cudaStream_t).
-  Status WaitForEventOnExternalStream(std::intptr_t stream);
+  absl::Status WaitForEventOnExternalStream(std::intptr_t stream);
 
   // Returns true if the event is known to have occurred by the tail of
   // 'stream'. If RecordOnStream has not yet been called, blocks the calling
@@ -130,7 +130,7 @@ class BufferSequencingEvent {
     return defined_status_.IsConcrete();
   }
 
-  void SetDefinedStatus(Status status) {
+  void SetDefinedStatus(absl::Status status) {
     {
       absl::MutexLock lock(&mu_);
       defined_status_.emplace(status);
@@ -139,7 +139,7 @@ class BufferSequencingEvent {
     this->ExecuteFutureTasks();
   }
 
-  Status GetDefinedStatus() {
+  absl::Status GetDefinedStatus() {
     absl::MutexLock lock(&mu_);
     CHECK(defined_status_.IsConcrete());
     return defined_status_.get();
diff --git a/third_party/xla/xla/pjrt/tracked_device_buffer_test.cc b/third_party/xla/xla/pjrt/tracked_device_buffer_test.cc
index 9b4b238775443b..814307b1ea90d5 100644
--- a/third_party/xla/xla/pjrt/tracked_device_buffer_test.cc
+++ b/third_party/xla/xla/pjrt/tracked_device_buffer_test.cc
@@ -28,12 +28,12 @@ limitations under the License.
 namespace xla {
 namespace {
 
-StatusOr<std::shared_ptr<TrackedDeviceBuffer>> MakeArray(const Shape& shape,
-                                                         LocalClient* client) {
+absl::StatusOr<std::shared_ptr<TrackedDeviceBuffer>> MakeArray(
+    const Shape& shape, LocalClient* client) {
   std::vector<stream_executor::DeviceMemoryBase> device_buffers;
   TF_RETURN_IF_ERROR(ShapeUtil::ForEachSubshapeWithStatus(
       client->backend().transfer_manager()->HostShapeToDeviceShape(shape),
-      [&](const Shape& subshape, const ShapeIndex&) -> Status {
+      [&](const Shape& subshape, const ShapeIndex&) -> absl::Status {
         TF_ASSIGN_OR_RETURN(
             se::OwningDeviceMemory device_memory,
             client->backend().memory_allocator()->Allocate(
diff --git a/third_party/xla/xla/pjrt/transpose.cc b/third_party/xla/xla/pjrt/transpose.cc
index 67dae968451d27..d45fab4fb1051a 100644
--- a/third_party/xla/xla/pjrt/transpose.cc
+++ b/third_party/xla/xla/pjrt/transpose.cc
@@ -677,7 +677,7 @@ int64_t TransposePlan::OutputNumElems() const {
 }
 
 // Parses and validates a tiling specification, and populates `tiling`.
-static Status ParseTilingSpecification(
+static absl::Status ParseTilingSpecification(
     int ndim, absl::Span<int64_t const> tiling_spec,
     absl::InlinedVector<int64_t, 4>& tiling) {
   tiling.resize(ndim, 1);
@@ -894,7 +894,7 @@ void TransposePlan::BuildPlanNodes(
   }
 }
 
-StatusOr<std::unique_ptr<TransposePlan>> TransposePlan::Create(
+absl::StatusOr<std::unique_ptr<TransposePlan>> TransposePlan::Create(
     const Options& o) {
   auto is_negative = [](int d) { return d < 0; };
   if (absl::c_find_if(o.dims, is_negative) != o.dims.end()) {
@@ -1360,7 +1360,7 @@ TransposePlanCache::TransposePlanCache(int capacity)
 
 TransposePlanCache::~TransposePlanCache() = default;
 
-StatusOr<std::shared_ptr<TransposePlan>> TransposePlanCache::GetOrCreate(
+absl::StatusOr<std::shared_ptr<TransposePlan>> TransposePlanCache::GetOrCreate(
     const TransposePlan::Options& o) {
   TransposePlanCacheKey key;
   key.elem_size_in_bytes = o.elem_size_in_bytes;
@@ -1388,7 +1388,7 @@ StatusOr<std::shared_ptr<TransposePlan>> TransposePlanCache::GetOrCreate(
   return cache_.GetOrCreateIfAbsent(
       key,
       [&](const TransposePlanCacheKey& key)
-          -> StatusOr<std::shared_ptr<TransposePlan>> {
+          -> absl::StatusOr<std::shared_ptr<TransposePlan>> {
         TF_ASSIGN_OR_RETURN(std::unique_ptr<TransposePlan> plan,
                             TransposePlan::Create(o));
         return std::shared_ptr<TransposePlan>(std::move(plan));
diff --git a/third_party/xla/xla/pjrt/transpose.h b/third_party/xla/xla/pjrt/transpose.h
index 767d707fdb88e5..0df1f8a7402289 100644
--- a/third_party/xla/xla/pjrt/transpose.h
+++ b/third_party/xla/xla/pjrt/transpose.h
@@ -98,7 +98,7 @@ class TransposePlan {
     int num_threads = 1;
   };
 
-  static StatusOr<std::unique_ptr<TransposePlan>> Create(
+  static absl::StatusOr<std::unique_ptr<TransposePlan>> Create(
       const Options& options);
 
   TransposePlan();
@@ -282,13 +282,14 @@ class TransposePlanCache {
   TransposePlanCache& operator=(TransposePlanCache&&) = delete;
 
   // Creates or returns a cached copy of a transpose plan.
-  StatusOr<std::shared_ptr<TransposePlan>> GetOrCreate(
+  absl::StatusOr<std::shared_ptr<TransposePlan>> GetOrCreate(
       const TransposePlan::Options& options);
 
  private:
   LRUCache<TransposePlanCacheKey,
-           StatusOr<std::shared_ptr<TransposePlan>>>::LRUList lru_list_;
-  LRUCache<TransposePlanCacheKey, StatusOr<std::shared_ptr<TransposePlan>>>
+           absl::StatusOr<std::shared_ptr<TransposePlan>>>::LRUList lru_list_;
+  LRUCache<TransposePlanCacheKey,
+           absl::StatusOr<std::shared_ptr<TransposePlan>>>
       cache_;
 };
 
diff --git a/third_party/xla/xla/pjrt/utils.cc b/third_party/xla/xla/pjrt/utils.cc
index 9bcbcfc89be21d..647b19b112cadf 100644
--- a/third_party/xla/xla/pjrt/utils.cc
+++ b/third_party/xla/xla/pjrt/utils.cc
@@ -63,8 +63,8 @@ limitations under the License.
 namespace xla {
 
 namespace {
-StatusOr<Shape> GetShardedShape(const Shape& shape,
-                                const OpSharding& sharding) {
+absl::StatusOr<Shape> GetShardedShape(const Shape& shape,
+                                      const OpSharding& sharding) {
   TF_ASSIGN_OR_RETURN(HloSharding hlo_sharding,
                       HloSharding::FromProto(sharding));
   if (shape.IsTuple()) {
@@ -82,7 +82,7 @@ StatusOr<Shape> GetShardedShape(const Shape& shape,
   }
 }
 
-StatusOr<Shape> GetShardedShape(const HloInstructionProto& instr) {
+absl::StatusOr<Shape> GetShardedShape(const HloInstructionProto& instr) {
   const Shape unsharded_shape(instr.shape());
   Shape sharded_shape;
   if (instr.has_sharding()) {
@@ -96,7 +96,7 @@ StatusOr<Shape> GetShardedShape(const HloInstructionProto& instr) {
 }
 
 // Returns sharded (argument shapes, result shape) without layouts.
-StatusOr<std::pair<std::vector<Shape>, Shape>> GetShardedProgramShapes(
+absl::StatusOr<std::pair<std::vector<Shape>, Shape>> GetShardedProgramShapes(
     const XlaComputation& computation, const ProgramShape& program_shape) {
   std::vector<Shape> arg_shapes;
   arg_shapes.resize(program_shape.parameters_size());
@@ -135,7 +135,7 @@ StatusOr<std::pair<std::vector<Shape>, Shape>> GetShardedProgramShapes(
 }
 }  // namespace
 
-Status ParseDeviceAssignmentCompileOptions(
+absl::Status ParseDeviceAssignmentCompileOptions(
     bool compile_portable_executable, ExecutableBuildOptions* build_options,
     std::function<StatusOr<DeviceAssignment>(int, int)>
         GetDefaultDeviceAssignmentFunction,
@@ -179,7 +179,7 @@ Status ParseDeviceAssignmentCompileOptions(
 // Helper method that takes an ArrayAttr of DictionaryAttrs for each arg or
 // result of a function, and looks for "mhlo.layout_mode". `all_attrs` can be
 // nullptr. `num_values` is the number of arguments or results.
-static StatusOr<std::vector<LayoutMode>> MlirAttrsToLayoutModes(
+static absl::StatusOr<std::vector<LayoutMode>> MlirAttrsToLayoutModes(
     mlir::ArrayAttr all_attrs, size_t num_values) {
   if (all_attrs == nullptr) {
     return std::vector<LayoutMode>(num_values);
@@ -210,7 +210,8 @@ static StatusOr<std::vector<LayoutMode>> MlirAttrsToLayoutModes(
 
 // TODO(b/329428415): Make this generic enough to be used by the GPU and TPU
 // compilers.
-StatusOr<MemorySpaceColor> GetMemorySpaceColor(const std::string& memory_kind) {
+absl::StatusOr<MemorySpaceColor> GetMemorySpaceColor(
+    const std::string& memory_kind) {
   // TODO(yashkatariya,zce): Unpinned_host is not valid for compiler. Only
   // pinned_host matters. So should there be a different lowering for
   // unpinned_host?
@@ -259,8 +260,9 @@ static absl::StatusOr<std::vector<MemorySpaceColor>> MlirAttrsToMemoryKinds(
 // Helper function for getting default LayoutModes for tupled arguments or
 // outputs. Returns nullopt if the arguments/outputs are not tupled. Raises an
 // error if layout modes are requested on tupled values.
-static StatusOr<std::optional<std::vector<LayoutMode>>> GetTupleLayoutModes(
-    mlir::ArrayRef<mlir::Type> types, mlir::ArrayAttr all_attrs) {
+static absl::StatusOr<std::optional<std::vector<LayoutMode>>>
+GetTupleLayoutModes(mlir::ArrayRef<mlir::Type> types,
+                    mlir::ArrayAttr all_attrs) {
   if (types.size() != 1 || !llvm::isa<mlir::TupleType>(types[0])) {
     return std::nullopt;
   }
@@ -307,7 +309,8 @@ GetTupleMemoryKinds(mlir::ArrayRef<mlir::Type> types,
       xla::Layout::kDefaultMemorySpace);
 }
 
-StatusOr<std::vector<LayoutMode>> GetArgLayoutModes(mlir::ModuleOp module) {
+absl::StatusOr<std::vector<LayoutMode>> GetArgLayoutModes(
+    mlir::ModuleOp module) {
   mlir::func::FuncOp main = module.lookupSymbol<mlir::func::FuncOp>("main");
   if (main == nullptr) {
     return InvalidArgument(
@@ -323,7 +326,8 @@ StatusOr<std::vector<LayoutMode>> GetArgLayoutModes(mlir::ModuleOp module) {
   return MlirAttrsToLayoutModes(main.getAllArgAttrs(), main.getNumArguments());
 }
 
-StatusOr<std::vector<LayoutMode>> GetOutputLayoutModes(mlir::ModuleOp module) {
+absl::StatusOr<std::vector<LayoutMode>> GetOutputLayoutModes(
+    mlir::ModuleOp module) {
   mlir::func::FuncOp main = module.lookupSymbol<mlir::func::FuncOp>("main");
   if (main == nullptr) {
     return InvalidArgument(
@@ -385,8 +389,8 @@ static std::string GetFrontendAttr(absl::Span<const LayoutMode> layout_modes) {
                        });
 }
 
-Status AddLayoutModesToFrontendAttrs(mlir::ModuleOp module,
-                                     XlaComputation& xla_computation) {
+absl::Status AddLayoutModesToFrontendAttrs(mlir::ModuleOp module,
+                                           XlaComputation& xla_computation) {
   TF_ASSIGN_OR_RETURN(std::vector<LayoutMode> arg_layout_modes,
                       GetArgLayoutModes(module));
   TF_ASSIGN_OR_RETURN(std::vector<LayoutMode> out_layout_modes,
@@ -411,8 +415,8 @@ static std::string GetFrontendAttrForMemorySpace(
       });
 }
 
-Status AddMemoryKindsToFrontendAttrs(mlir::ModuleOp module,
-                                     XlaComputation& xla_computation) {
+absl::Status AddMemoryKindsToFrontendAttrs(mlir::ModuleOp module,
+                                           XlaComputation& xla_computation) {
   TF_ASSIGN_OR_RETURN(std::vector<MemorySpaceColor> arg_memory_spaces,
                       GetArgMemoryKinds(module));
   TF_ASSIGN_OR_RETURN(std::vector<MemorySpaceColor> out_memory_spaces,
@@ -430,7 +434,7 @@ Status AddMemoryKindsToFrontendAttrs(mlir::ModuleOp module,
   return OkStatus();
 }
 
-static StatusOr<std::vector<LayoutMode>> GetLayoutModesFromFrontendAttr(
+static absl::StatusOr<std::vector<LayoutMode>> GetLayoutModesFromFrontendAttr(
     absl::string_view attr) {
   // SkipEmpty() needed to avoid returning the empty string when attr is empty.
   std::vector<std::string> str_modes =
@@ -443,7 +447,7 @@ static StatusOr<std::vector<LayoutMode>> GetLayoutModesFromFrontendAttr(
   return result;
 }
 
-static StatusOr<std::vector<LayoutMode>> GetLayoutModes(
+static absl::StatusOr<std::vector<LayoutMode>> GetLayoutModes(
     const XlaComputation& computation, absl::string_view frontend_attr_name,
     size_t num_values) {
   const auto& frontend_attrs = computation.proto().frontend_attributes().map();
@@ -455,8 +459,8 @@ static StatusOr<std::vector<LayoutMode>> GetLayoutModes(
   return GetLayoutModesFromFrontendAttr(iter->second);
 }
 
-static StatusOr<std::vector<MemorySpaceColor>> GetMemoryKindsFromFrontendAttr(
-    absl::string_view attr) {
+static absl::StatusOr<std::vector<MemorySpaceColor>>
+GetMemoryKindsFromFrontendAttr(absl::string_view attr) {
   // SkipEmpty() needed to avoid returning the empty string when attr is empty.
   std::vector<std::string> str_memory_spaces =
       absl::StrSplit(attr, kDelimiter, absl::SkipEmpty());
@@ -471,7 +475,7 @@ static StatusOr<std::vector<MemorySpaceColor>> GetMemoryKindsFromFrontendAttr(
   return result;
 }
 
-static StatusOr<std::vector<MemorySpaceColor>> GetMemoryKinds(
+static absl::StatusOr<std::vector<MemorySpaceColor>> GetMemoryKinds(
     const XlaComputation& computation, absl::string_view frontend_attr_name,
     size_t num_values) {
   const auto& frontend_attrs = computation.proto().frontend_attributes().map();
@@ -484,7 +488,7 @@ static StatusOr<std::vector<MemorySpaceColor>> GetMemoryKinds(
   return GetMemoryKindsFromFrontendAttr(iter->second);
 }
 
-StatusOr<std::vector<LayoutMode>> GetArgLayoutModes(
+absl::StatusOr<std::vector<LayoutMode>> GetArgLayoutModes(
     const XlaComputation& computation) {
   TF_ASSIGN_OR_RETURN(ProgramShape program_shape,
                       computation.GetProgramShape());
@@ -495,7 +499,7 @@ StatusOr<std::vector<LayoutMode>> GetArgLayoutModes(
   return GetLayoutModes(computation, "arg_layout_modes", num_args);
 }
 
-StatusOr<std::vector<MemorySpaceColor>> GetArgMemoryKinds(
+absl::StatusOr<std::vector<MemorySpaceColor>> GetArgMemoryKinds(
     const XlaComputation& computation) {
   TF_ASSIGN_OR_RETURN(ProgramShape program_shape,
                       computation.GetProgramShape());
@@ -506,7 +510,7 @@ StatusOr<std::vector<MemorySpaceColor>> GetArgMemoryKinds(
   return GetMemoryKinds(computation, "arg_memory_spaces", num_args);
 }
 
-StatusOr<std::vector<LayoutMode>> GetOutputLayoutModes(
+absl::StatusOr<std::vector<LayoutMode>> GetOutputLayoutModes(
     const XlaComputation& computation) {
   TF_ASSIGN_OR_RETURN(ProgramShape program_shape,
                       computation.GetProgramShape());
@@ -516,7 +520,7 @@ StatusOr<std::vector<LayoutMode>> GetOutputLayoutModes(
   return GetLayoutModes(computation, "out_layout_modes", num_outputs);
 }
 
-StatusOr<std::vector<MemorySpaceColor>> GetOutputMemoryKinds(
+absl::StatusOr<std::vector<MemorySpaceColor>> GetOutputMemoryKinds(
     const XlaComputation& computation) {
   TF_ASSIGN_OR_RETURN(ProgramShape program_shape,
                       computation.GetProgramShape());
@@ -526,7 +530,7 @@ StatusOr<std::vector<MemorySpaceColor>> GetOutputMemoryKinds(
   return GetMemoryKinds(computation, "out_memory_spaces", num_outputs);
 }
 
-static StatusOr<Shape> LayoutModeToXlaShape(
+static absl::StatusOr<Shape> LayoutModeToXlaShape(
     const LayoutMode& layout_mode, const Shape& unsharded_shape,
     const Shape& sharded_shape,
     std::function<StatusOr<Shape>(Shape)>
@@ -565,7 +569,7 @@ static StatusOr<Shape> LayoutModeToXlaShape(
   return result;
 }
 
-StatusOr<std::pair<std::vector<Shape>, Shape>> LayoutModesToXlaShapes(
+absl::StatusOr<std::pair<std::vector<Shape>, Shape>> LayoutModesToXlaShapes(
     const XlaComputation& computation, std::vector<LayoutMode> arg_layout_modes,
     std::vector<LayoutMode> out_layout_modes,
     const std::vector<MemorySpaceColor>& arg_memory_spaces,
@@ -673,7 +677,7 @@ StatusOr<std::pair<std::vector<Shape>, Shape>> LayoutModesToXlaShapes(
                                               std::move(out_layout)};
 }
 
-StatusOr<std::pair<std::vector<Shape>, std::vector<const Shape*>>>
+absl::StatusOr<std::pair<std::vector<Shape>, std::vector<const Shape*>>>
 LayoutModesToXla(const XlaComputation& computation,
                  std::vector<LayoutMode> arg_layout_modes,
                  std::vector<LayoutMode> out_layout_modes,
@@ -704,7 +708,7 @@ LayoutModesToXla(const XlaComputation& computation,
       std::move(arg_layouts), std::move(arg_layout_pointers)};
 }
 
-Status DetermineArgumentLayoutsFromCompileOptions(
+absl::Status DetermineArgumentLayoutsFromCompileOptions(
     const XlaComputation& computation,
     std::function<StatusOr<Shape>(Shape)>
         choose_compact_layout_for_shape_function,
@@ -767,7 +771,7 @@ Status DetermineArgumentLayoutsFromCompileOptions(
   return OkStatus();
 }
 
-StatusOr<std::vector<int>> ComputeParametersThatMustBeDonated(
+absl::StatusOr<std::vector<int>> ComputeParametersThatMustBeDonated(
     const HloModule& module, bool tuple_inputs) {
   HloComputation* computation = module.entry_computation();
   int number_of_parameters = [&]() -> int {
@@ -858,7 +862,7 @@ bool HasMajorToMinorLayout(PrimitiveType type, absl::Span<int64_t const> dims,
   return true;
 }
 
-StatusOr<Shape> MakeShapeWithTrivialByteStrides(
+absl::StatusOr<Shape> MakeShapeWithTrivialByteStrides(
     PrimitiveType element_type, absl::Span<const int64_t> dimensions,
     absl::Span<const int64_t> byte_strides) {
   TF_RET_CHECK(dimensions.size() == byte_strides.size());
@@ -894,7 +898,7 @@ StatusOr<Shape> MakeShapeWithTrivialByteStrides(
                                              minor_to_major);
 }
 
-Status TestBufferDonationClashes(
+absl::Status TestBufferDonationClashes(
     void* opaque_key,
     absl::flat_hash_map<const void*, std::pair<bool, int>>& donation_clashes,
     bool is_donated, int arg_idx, int replica, int partition) {
diff --git a/third_party/xla/xla/pjrt/utils.h b/third_party/xla/xla/pjrt/utils.h
index 3134b49702eee5..0271924420913b 100644
--- a/third_party/xla/xla/pjrt/utils.h
+++ b/third_party/xla/xla/pjrt/utils.h
@@ -41,7 +41,7 @@ using MemorySpaceColor = int;
 
 // Returns the num_replicas, num_partitions and device assignment given a
 // ExecutableBuildOptions and whether we want a portable executable.
-Status ParseDeviceAssignmentCompileOptions(
+absl::Status ParseDeviceAssignmentCompileOptions(
     bool compile_portable_executable, ExecutableBuildOptions* build_options,
     std::function<StatusOr<DeviceAssignment>(int, int)>
         GetDefaultDeviceAssignmentFunction,
@@ -51,56 +51,58 @@ Status ParseDeviceAssignmentCompileOptions(
 // Returns the LayoutMode for each argument of the main function in the
 // module. Checks for the "mhlo.layout_mode" attr, and if not present, assumes
 // LayoutMode::Mode::kDefault.
-StatusOr<std::vector<LayoutMode>> GetArgLayoutModes(mlir::ModuleOp module);
+absl::StatusOr<std::vector<LayoutMode>> GetArgLayoutModes(
+    mlir::ModuleOp module);
 // Returns the LayoutMode for each output of the main function in the
 // module. Checks for the "mhlo.layout_mode" attr, and if not present, assumes
 // LayoutMode::Mode::kDefault.
-StatusOr<std::vector<LayoutMode>> GetOutputLayoutModes(mlir::ModuleOp module);
+absl::StatusOr<std::vector<LayoutMode>> GetOutputLayoutModes(
+    mlir::ModuleOp module);
 
 // Returns the memory space for each argument of the computations. Checks
 // for the "mhlo.memory_kind" frontend attribute, and if not present, assumes 0.
-StatusOr<std::vector<MemorySpaceColor>> GetArgMemoryKinds(
+absl::StatusOr<std::vector<MemorySpaceColor>> GetArgMemoryKinds(
     mlir::ModuleOp module);
 // Returns the memory space for each output of the computations. Checks for
 // the "mhlo.memory_kind" frontend attribute, and if not present, assumes 0.
-StatusOr<std::vector<MemorySpaceColor>> GetOutputMemoryKinds(
+absl::StatusOr<std::vector<MemorySpaceColor>> GetOutputMemoryKinds(
     mlir::ModuleOp module);
 
 // Populates the frontend attributes "arg_layout_mode" and "out_layout_mode" in
 // xla_computation based on `module`. This function must be called before the
 // LayoutMode getters below work correctly on `computation`.
-Status AddLayoutModesToFrontendAttrs(mlir::ModuleOp module,
-                                     XlaComputation& xla_computation);
+absl::Status AddLayoutModesToFrontendAttrs(mlir::ModuleOp module,
+                                           XlaComputation& xla_computation);
 
 // Populates the frontend attributes "arg_memory_kinds" and "out_memory_kinds"
 // in xla_computation based on `module`. This function must be called before the
 // LayoutMode getters below work correctly on `computation`.
-Status AddMemoryKindsToFrontendAttrs(mlir::ModuleOp module,
-                                     XlaComputation& xla_computation);
+absl::Status AddMemoryKindsToFrontendAttrs(mlir::ModuleOp module,
+                                           XlaComputation& xla_computation);
 
 // Returns the LayoutMode for each argument of the computations. Checks for the
 // "arg_layout_mode" frontend attribute, and if not present, assumes
 // LayoutMode::Mode::kDefault.
-StatusOr<std::vector<LayoutMode>> GetArgLayoutModes(
+absl::StatusOr<std::vector<LayoutMode>> GetArgLayoutModes(
     const XlaComputation& computation);
 // Returns the LayoutMode for each argument of the computations. Checks for the
 // "out_layout_mode" frontend attribute, and if not present, assumes
 // LayoutMode::Mode::kDefault.
-StatusOr<std::vector<LayoutMode>> GetOutputLayoutModes(
+absl::StatusOr<std::vector<LayoutMode>> GetOutputLayoutModes(
     const XlaComputation& computation);
 
 // Returns the memory space for each argument of the computations. Checks for
 // the "arg_memory_kind" frontend attribute, and if not present, assumes 0.
-StatusOr<std::vector<MemorySpaceColor>> GetArgMemoryKinds(
+absl::StatusOr<std::vector<MemorySpaceColor>> GetArgMemoryKinds(
     const XlaComputation& computation);
 // Returns the memory space for each argument of the computations. Checks for
 // the "out_memory_kind" frontend attribute, and if not present, assumes 0.
-StatusOr<std::vector<MemorySpaceColor>> GetOutputMemoryKinds(
+absl::StatusOr<std::vector<MemorySpaceColor>> GetOutputMemoryKinds(
     const XlaComputation& computation);
 
 // Returns (arg shapes, output shape) with properly-set Layouts that can
 // be passed to XLA to reflect arg_layout_modes and out_layout_modes.
-StatusOr<std::pair<std::vector<Shape>, Shape>> LayoutModesToXlaShapes(
+absl::StatusOr<std::pair<std::vector<Shape>, Shape>> LayoutModesToXlaShapes(
     const XlaComputation& computation, std::vector<LayoutMode> arg_layout_modes,
     std::vector<LayoutMode> out_layout_modes,
     const std::vector<MemorySpaceColor>& arg_memory_spaces,
@@ -112,7 +114,7 @@ StatusOr<std::pair<std::vector<Shape>, Shape>> LayoutModesToXlaShapes(
 // * Returns a vector of argument xla::Shapes with properly-set Layouts
 // * Returns vector of pointers to those Shapes to create HloModuleConfig
 // * Modifies `build_options` to have the correct result_layout set or unset
-StatusOr<std::pair<std::vector<Shape>, std::vector<const Shape*>>>
+absl::StatusOr<std::pair<std::vector<Shape>, std::vector<const Shape*>>>
 LayoutModesToXla(const XlaComputation& computation,
                  std::vector<LayoutMode> arg_layout_modes,
                  std::vector<LayoutMode> out_layout_modes,
@@ -124,7 +126,7 @@ LayoutModesToXla(const XlaComputation& computation,
 
 // Returns pointers to the argument layouts given an XlaComputation and
 // ExecutableBuildOptions.
-Status DetermineArgumentLayoutsFromCompileOptions(
+absl::Status DetermineArgumentLayoutsFromCompileOptions(
     const XlaComputation& computation,
     std::function<StatusOr<Shape>(Shape)>
         choose_compact_layout_for_shape_function,
@@ -136,7 +138,7 @@ Status DetermineArgumentLayoutsFromCompileOptions(
 // to outputs. This function returns a sorted vector of parameters that must be
 // donated when executable is run. tuple_inputs reflects the option that
 // executable was compiled with.
-StatusOr<std::vector<int>> ComputeParametersThatMustBeDonated(
+absl::StatusOr<std::vector<int>> ComputeParametersThatMustBeDonated(
     const HloModule& hlo_module, bool tuple_inputs);
 
 // Return max parallelism level.
@@ -150,7 +152,7 @@ bool HasMajorToMinorLayout(PrimitiveType type, absl::Span<int64_t const> dims,
 // Constructs a new dense array shape with the given byte strides. Supports only
 // trivial (compact) byte_strides that represents a transposition of a dense
 // buffer.
-StatusOr<Shape> MakeShapeWithTrivialByteStrides(
+absl::StatusOr<Shape> MakeShapeWithTrivialByteStrides(
     PrimitiveType element_type, absl::Span<const int64_t> dimensions,
     absl::Span<const int64_t> byte_strides);
 
@@ -159,7 +161,7 @@ StatusOr<Shape> MakeShapeWithTrivialByteStrides(
 // Multiple uses are valid iff they are all not donations.  The provided map
 // stores the opaque buffer identity, a bool to denote if the previous use is a
 // donation, and the index of the previous use for better error messages.
-Status TestBufferDonationClashes(
+absl::Status TestBufferDonationClashes(
     void* opaque_key,
     absl::flat_hash_map<const void*, std::pair<bool, int>>& donation_clashes,
     bool is_donated, int arg_idx, int replica, int partition);

From 3321aa430175492e018972720e4f96eb7e4eef68 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Thu, 16 May 2024 11:28:50 -0700
Subject: [PATCH 291/478] Update ops-related pbtxt files.

PiperOrigin-RevId: 634474219
---
 tensorflow/core/ops/ops.pbtxt | 1 -
 1 file changed, 1 deletion(-)

diff --git a/tensorflow/core/ops/ops.pbtxt b/tensorflow/core/ops/ops.pbtxt
index 2a7a7c62c9d010..042577243a4e90 100644
--- a/tensorflow/core/ops/ops.pbtxt
+++ b/tensorflow/core/ops/ops.pbtxt
@@ -1,4 +1,3 @@
-go/debugstr    
 op {
   name: "Abort"
   attr {

From 2c05f7889c71d863f2679dec13d8e05dab4e1b1e Mon Sep 17 00:00:00 2001
From: Henning Becker <hebecker@google.com>
Date: Thu, 16 May 2024 11:35:55 -0700
Subject: [PATCH 292/478] Add missing rocm_config.h include to
 amdgpu_compiler.cc

rocm_config.h defines `TF_ROCM_VERSION` which is used in `amdgpu_compiler.cc`.

PiperOrigin-RevId: 634476794
---
 third_party/xla/xla/service/gpu/amdgpu_compiler.cc | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/third_party/xla/xla/service/gpu/amdgpu_compiler.cc b/third_party/xla/xla/service/gpu/amdgpu_compiler.cc
index d9eb60b3046658..4224fe11d372d3 100644
--- a/third_party/xla/xla/service/gpu/amdgpu_compiler.cc
+++ b/third_party/xla/xla/service/gpu/amdgpu_compiler.cc
@@ -61,6 +61,10 @@ limitations under the License.
 #include "tsl/platform/statusor.h"
 #include "tsl/platform/threadpool.h"
 
+#if TENSORFLOW_USE_ROCM
+#include "rocm/rocm_config.h"
+#endif
+
 namespace xla {
 namespace gpu {
 

From 4ce1a064f2e4b0c141aa64f590716db1cc9d74df Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Thu, 16 May 2024 11:42:20 -0700
Subject: [PATCH 293/478] Reverts c031655ca9514e7946bb87eda17cf5932c381a8b

PiperOrigin-RevId: 634478686
---
 tensorflow/compiler/mlir/lite/tests/optimize.mlir   | 13 -------------
 .../mlir/lite/transforms/optimize_patterns.td       | 10 ----------
 2 files changed, 23 deletions(-)

diff --git a/tensorflow/compiler/mlir/lite/tests/optimize.mlir b/tensorflow/compiler/mlir/lite/tests/optimize.mlir
index 96cffcd1371473..3c2c24baba8972 100644
--- a/tensorflow/compiler/mlir/lite/tests/optimize.mlir
+++ b/tensorflow/compiler/mlir/lite/tests/optimize.mlir
@@ -103,19 +103,6 @@ func.func @fuseSubIntoConv2d(%arg0: tensor<256x32x32x3xf32>, %arg1: tensor<16x3x
   // CHECK: %0 = "tfl.conv_2d"(%arg0, %arg1, %cst)
 }
 
-// CHECK-LABEL: foldQuantWeightsIntoTposeConv
-func.func @foldQuantWeightsIntoTposeConv(%arg0: tensor<2x2x3x2048xf32>) -> tensor<2x3x2x2048xf32> {
-  %output_shape = arith.constant dense<[2, 3, 2, 2048]> : tensor<4xi32>
-  %q_weighs = "tfl.pseudo_qconst"() {qtype = tensor<4x2x2x2048x!quant.uniform<u8<1:255>:f32, 0.15:151>>, value = dense<-76> : tensor<4x2x2x2048xi8>} : () -> tensor<4x2x2x2048x!quant.uniform<u8<1:255>:f32, 0.15:151>>
-  %dq_weights = "tfl.dequantize"(%q_weighs) : (tensor<4x2x2x2048x!quant.uniform<u8<1:255>:f32, 0.15:151>>) -> tensor<4x2x2x2048xf32>
-  %bias = "tfl.no_value"() {value} : () -> none
-  %out = "tfl.transpose_conv"(%output_shape, %dq_weights, %arg0, %bias) {fused_activation_function = "NONE", padding = "SAME", stride_h = 1 : i32, stride_w = 1 : i32} : (tensor<4xi32>, tensor<4x2x2x2048xf32>, tensor<2x2x3x2048xf32>, none) -> tensor<2x3x2x2048xf32>
-  func.return %out : tensor<2x3x2x2048xf32>
-
-  // CHECK-NOT: "tfl.dequantize"
-  // CHECK: "tfl.transpose_conv"(%cst, %1, %arg0, %0) <{fused_activation_function = "NONE", padding = "SAME", stride_h = 1 : i32, stride_w = 1 : i32}> : (tensor<4xi32>, tensor<4x2x2x2048x!quant.uniform<u8<1:255>:f32
-}
-
 // CHECK-LABEL: fuseAddIntoTransposeConv
 func.func @fuseAddIntoTransposeConv(%arg0: tensor<1x32x42x128xf32>) -> tensor<1x64x84x32xf32> {
   %cst = arith.constant dense<1.5> : tensor<32xf32>
diff --git a/tensorflow/compiler/mlir/lite/transforms/optimize_patterns.td b/tensorflow/compiler/mlir/lite/transforms/optimize_patterns.td
index d3da883d0d8a02..4353b82e2fb901 100644
--- a/tensorflow/compiler/mlir/lite/transforms/optimize_patterns.td
+++ b/tensorflow/compiler/mlir/lite/transforms/optimize_patterns.td
@@ -813,16 +813,6 @@ def MatchReluPattern : Pat<
   (TFL_ReluOp $input),
   [(FloatValueEquals<"0"> $Zero)]>;
 
-// Transpose conv supports hybrid computation with quantized weights.
-def FoldQuantWeightsIntoTposeConv : Pat<
-  (TFL_TransposeConvOp
-    $output_shape,
-    (TFL_DequantizeOp $quant_weights),
-    $quant_input,
-    $bias, $padding, $stride_h, $stride_w, $faf),
-  (TFL_TransposeConvOp $output_shape, $quant_weights,
-    $quant_input, $bias, $padding, $stride_h, $stride_w, $faf)>;
-
 // Optimize Minimum of tf.Relu and constant six to tf.Relu6
 def MinimumOfReluAnd6ToRelu6 :
   Pat<(TFL_MinimumOp (TFL_ReluOp $x), (Arith_ConstantOp $y)),

From ebff15dc59ef5759cda3794425395071058242b6 Mon Sep 17 00:00:00 2001
From: Tongfei Guo <tongfei@google.com>
Date: Thu, 16 May 2024 11:58:06 -0700
Subject: [PATCH 294/478] Allow overriding some methods in
 all_gather_decomposer.

PiperOrigin-RevId: 634483343
---
 third_party/xla/xla/service/BUILD             |  8 +++--
 .../xla/xla/service/all_gather_decomposer.cc  | 34 ++++++++++++-------
 .../xla/xla/service/all_gather_decomposer.h   | 13 +++++++
 3 files changed, 40 insertions(+), 15 deletions(-)

diff --git a/third_party/xla/xla/service/BUILD b/third_party/xla/xla/service/BUILD
index ad05eb8777cd74..e9c6d033688b6d 100644
--- a/third_party/xla/xla/service/BUILD
+++ b/third_party/xla/xla/service/BUILD
@@ -3625,12 +3625,14 @@ cc_library(
         ":hlo_pass",
         "//xla:literal_util",
         "//xla:shape_util",
-        "//xla:status_macros",
-        "//xla:types",
+        "//xla:status",
+        "//xla:util",
         "//xla/hlo/ir:hlo",
-        "@com_google_absl//absl/algorithm:container",
+        "@com_google_absl//absl/container:flat_hash_set",
         "@com_google_absl//absl/strings",
+        "@local_tsl//tsl/platform:errors",
         "@local_tsl//tsl/platform:logging",
+        "@local_tsl//tsl/platform:statusor",
     ],
 )
 
diff --git a/third_party/xla/xla/service/all_gather_decomposer.cc b/third_party/xla/xla/service/all_gather_decomposer.cc
index f0a9c625ad8a4e..56c9115ef6f821 100644
--- a/third_party/xla/xla/service/all_gather_decomposer.cc
+++ b/third_party/xla/xla/service/all_gather_decomposer.cc
@@ -15,10 +15,12 @@ limitations under the License.
 
 #include "xla/service/all_gather_decomposer.h"
 
+#include <cstdint>
+#include <optional>
 #include <vector>
 
-#include "absl/algorithm/container.h"
-#include "absl/strings/str_join.h"
+#include "absl/container/flat_hash_set.h"
+#include "absl/strings/string_view.h"
 #include "xla/hlo/ir/hlo_casting_utils.h"
 #include "xla/hlo/ir/hlo_computation.h"
 #include "xla/hlo/ir/hlo_instruction.h"
@@ -27,13 +29,17 @@ limitations under the License.
 #include "xla/literal_util.h"
 #include "xla/service/collective_decomposer_utils.h"
 #include "xla/service/collective_ops_utils.h"
+#include "xla/shape.h"
 #include "xla/shape_util.h"
-#include "xla/status_macros.h"
-#include "xla/types.h"
+#include "xla/status.h"
+#include "xla/util.h"
+#include "tsl/platform/errors.h"
 #include "tsl/platform/logging.h"
+#include "tsl/platform/statusor.h"
 
 namespace xla {
 
+namespace {
 // Creates a computation of x + y.
 HloComputation* MakeBinaryAdd(PrimitiveType type, HloModule* module) {
   HloComputation::Builder sum_b("add");
@@ -51,14 +57,15 @@ HloComputation* MakeBinaryAdd(PrimitiveType type, HloModule* module) {
   HloComputation* reduction = module->AddEmbeddedComputation(sum_b.Build());
   return reduction;
 }
+}  // namespace
 
-HloInstruction* TranslateAllGatherToAllReducePerOperand(
+HloInstruction* AllGatherDecomposer::TranslateAllGatherToAllReducePerOperand(
     CollectiveOpGroupMode group_mode, const HloAllGatherInstruction& ag,
-    const Shape& output_shape, HloInstruction* operand, HloComputation* comp) {
+    const Shape& output_shape, HloInstruction* operand, HloComputation* comp,
+    int64_t ag_dim) {
   std::vector<HloInstruction*> start_indices =
       CreateStartIndicesForCollectiveDecomposition(
-          group_mode, ag.replica_groups(), operand->shape(),
-          ag.all_gather_dimension(), comp)
+          group_mode, ag.replica_groups(), operand->shape(), ag_dim, comp)
           .value();
 
   auto zero = comp->AddInstruction(HloInstruction::CreateConstant(
@@ -77,7 +84,8 @@ HloInstruction* TranslateAllGatherToAllReducePerOperand(
   return ar;
 }
 
-Status DecomposeAllGather(HloAllGatherInstruction* ag, HloComputation* comp) {
+Status AllGatherDecomposer::DecomposeAllGather(HloAllGatherInstruction* ag,
+                                               HloComputation* comp) {
   TF_ASSIGN_OR_RETURN(CollectiveOpGroupMode group_mode,
                       GetCollectiveOpGroupMode(ag->channel_id().has_value(),
                                                ag->use_global_device_ids()));
@@ -87,14 +95,16 @@ Status DecomposeAllGather(HloAllGatherInstruction* ag, HloComputation* comp) {
       auto* input_operand = ag->mutable_operand(i);
       const auto& output_shape = ag->shape().tuple_shapes(i);
       auto* ar = TranslateAllGatherToAllReducePerOperand(
-          group_mode, *ag, output_shape, input_operand, comp);
+          group_mode, *ag, output_shape, input_operand, comp,
+          ag->all_gather_dimension());
       tuple_inputs.push_back(ar);
     }
     auto tup = comp->AddInstruction(HloInstruction::CreateTuple(tuple_inputs));
     TF_RETURN_IF_ERROR(ag->ReplaceAllUsesWith(tup));
   } else {
     auto* ar = TranslateAllGatherToAllReducePerOperand(
-        group_mode, *ag, ag->shape(), ag->mutable_operand(0), comp);
+        group_mode, *ag, ag->shape(), ag->mutable_operand(0), comp,
+        ag->all_gather_dimension());
     TF_RETURN_IF_ERROR(ag->ReplaceAllUsesWith(ar));
   }
   TF_RETURN_IF_ERROR(comp->RemoveInstructionAndUnusedOperands(ag));
@@ -111,7 +121,7 @@ absl::StatusOr<bool> AllGatherDecomposer::Run(
         continue;
       }
       auto ag = Cast<HloAllGatherInstruction>(hlo);
-      if (should_decompose_(*ag)) {
+      if (ShouldDecompose(*ag)) {
         TF_RETURN_IF_ERROR(DecomposeAllGather(ag, comp));
         changed = true;
       }
diff --git a/third_party/xla/xla/service/all_gather_decomposer.h b/third_party/xla/xla/service/all_gather_decomposer.h
index da56d0c4023037..924113fc23d2de 100644
--- a/third_party/xla/xla/service/all_gather_decomposer.h
+++ b/third_party/xla/xla/service/all_gather_decomposer.h
@@ -18,6 +18,7 @@ limitations under the License.
 
 #include "xla/hlo/ir/hlo_instructions.h"
 #include "xla/hlo/ir/hlo_module.h"
+#include "xla/service/collective_ops_utils.h"
 #include "xla/service/hlo_pass_interface.h"
 
 namespace xla {
@@ -41,6 +42,18 @@ class AllGatherDecomposer : public HloModulePass {
       HloModule* module,
       const absl::flat_hash_set<absl::string_view>& execution_threads) override;
 
+ protected:
+  virtual HloInstruction* TranslateAllGatherToAllReducePerOperand(
+      CollectiveOpGroupMode group_mode, const HloAllGatherInstruction& ag,
+      const Shape& output_shape, HloInstruction* operand, HloComputation* comp,
+      int64_t ag_dim);
+
+  virtual bool ShouldDecompose(const HloAllGatherInstruction& ag) const {
+    return should_decompose_(ag);
+  }
+
+  Status DecomposeAllGather(HloAllGatherInstruction* ag, HloComputation* comp);
+
  private:
   std::function<bool(const HloAllGatherInstruction&)> should_decompose_;
 };

From 43085540b1af0926ed392e3dbedb75b7cdf5af81 Mon Sep 17 00:00:00 2001
From: Eugene Zhulenev <ezhulenev@google.com>
Date: Thu, 16 May 2024 12:03:46 -0700
Subject: [PATCH 295/478] [xla_client] Port all custom calls used in test to
 FFI

- Remove legacy custom calls and always use FFI for xla_client tests.
- Add a little bit more documentation to cover custom call registration process

PiperOrigin-RevId: 634485582
---
 third_party/xla/xla/python/BUILD              | 18 ++--
 .../xla/xla/python/custom_call_for_test.pyx   | 33 -------
 .../xla/xla/python/custom_calls_testlib.cc    | 85 +++++++++++++++++++
 .../python/typed_ffi_custom_call_for_test.cc  | 67 ---------------
 third_party/xla/xla/python/xla_client_test.py | 61 +++----------
 5 files changed, 101 insertions(+), 163 deletions(-)
 delete mode 100644 third_party/xla/xla/python/custom_call_for_test.pyx
 create mode 100644 third_party/xla/xla/python/custom_calls_testlib.cc
 delete mode 100644 third_party/xla/xla/python/typed_ffi_custom_call_for_test.cc

diff --git a/third_party/xla/xla/python/BUILD b/third_party/xla/xla/python/BUILD
index ec8631655c1243..48501d4600b401 100644
--- a/third_party/xla/xla/python/BUILD
+++ b/third_party/xla/xla/python/BUILD
@@ -1,7 +1,7 @@
 load("@bazel_skylib//rules:common_settings.bzl", "bool_flag")
 load("@local_config_cuda//cuda:build_defs.bzl", "if_cuda")
 load("@local_config_rocm//rocm:build_defs.bzl", "if_rocm")
-load("@local_tsl//tsl/platform:build_config.bzl", "pyx_library", "tf_proto_library")
+load("@local_tsl//tsl/platform:build_config.bzl", "tf_proto_library")
 load("@local_tsl//tsl/platform:rules_cc.bzl", "cc_library")
 load(
     "@local_tsl//tsl/platform/default:cuda_build_defs.bzl",
@@ -55,19 +55,14 @@ exports_files([
     "xla_client.pyi",
 ])
 
-pyx_library(
-    name = "custom_call_for_test",
-    testonly = True,
-    srcs = ["custom_call_for_test.pyx"],
-)
-
 tsl_pybind_extension(
-    name = "typed_ffi_custom_calls",
-    srcs = ["typed_ffi_custom_call_for_test.cc"],
+    name = "custom_calls_testlib",
+    srcs = ["custom_calls_testlib.cc"],
     visibility = ["//visibility:private"],
     deps = [
         "//third_party/nanobind",
-        "//xla/ffi",
+        "//xla/ffi/api:c_api",
+        "//xla/ffi/api:ffi",
         "@com_google_absl//absl/status",
     ],
 )
@@ -111,8 +106,7 @@ py_strict_test(
         "not_run:arm",
     ],  # TODO(phawkins): This test passes, but requires --config=monolithic.
     deps = [
-        ":custom_call_for_test",
-        ":typed_ffi_custom_calls",
+        ":custom_calls_testlib",
         ":xla_client",
         ":xla_extension",
         "@absl_py//absl/flags",
diff --git a/third_party/xla/xla/python/custom_call_for_test.pyx b/third_party/xla/xla/python/custom_call_for_test.pyx
deleted file mode 100644
index 306b424067d115..00000000000000
--- a/third_party/xla/xla/python/custom_call_for_test.pyx
+++ /dev/null
@@ -1,33 +0,0 @@
-# cython: language_level=2
-# distutils: language = c++
-
-# Test case for defining a XLA custom call target in Cython, and registering
-# it via the xla_client SWIG API.
-
-from cpython.pycapsule cimport PyCapsule_New
-
-cdef void test_subtract_f32(void* out_ptr, void** data_ptr,
-                            void* xla_custom_call_status) nogil:
-  cdef float a = (<float*>(data_ptr[0]))[0]
-  cdef float b = (<float*>(data_ptr[1]))[0]
-  cdef float* out = <float*>(out_ptr)
-  out[0] = a - b
-
-cdef void test_add_input_and_opaque_len(void* out_buffer, void** ins,
-                                        const char* opaque, size_t opaque_len,
-                                        void* xla_custom_call_status):
-  cdef float a = (<float*>(ins[0]))[0]
-  cdef float b = <float>opaque_len
-  cdef float* out = <float*>(out_buffer)
-  out[0] = a + b
-
-
-cpu_custom_call_targets = {}
-
-cdef register_custom_call_target(fn_name, void* fn):
-  cdef const char* name = "xla._CUSTOM_CALL_TARGET"
-  cpu_custom_call_targets[fn_name] = PyCapsule_New(fn, name, NULL)
-
-register_custom_call_target(b"test_subtract_f32", <void*>(test_subtract_f32))
-register_custom_call_target(b"test_add_input_and_opaque_len",
-                            <void*>(test_add_input_and_opaque_len))
diff --git a/third_party/xla/xla/python/custom_calls_testlib.cc b/third_party/xla/xla/python/custom_calls_testlib.cc
new file mode 100644
index 00000000000000..ee32a9864681fe
--- /dev/null
+++ b/third_party/xla/xla/python/custom_calls_testlib.cc
@@ -0,0 +1,85 @@
+/* Copyright 2024 The OpenXLA Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+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.
+==============================================================================*/
+
+#include "third_party/nanobind/include/nanobind/nanobind.h"
+#include "xla/ffi/api/c_api.h"
+#include "xla/ffi/api/ffi.h"
+
+namespace xla::ffi {
+namespace nb = ::nanobind;
+
+// Implement custom calls as static functions with XLA FFI types in the function
+// signature that gives access to the arguments and results buffers together
+// with their types and dimensions. See `ffi/api/ffi_test.cc` for more XLA FFI
+// examples and features (e.g. binding attributes, custom user-defined structs
+// and arbitrary execution context).
+
+static Error AlwaysFail(Result<BufferBase>) {
+  return Error(XLA_FFI_Error_Code_INTERNAL, "Failed intentionally");
+}
+
+static Error AlwaysSucceed(Result<BufferBase>) { return Error::Success(); }
+
+static Error Subtract(BufferR0<DataType::F32> a, BufferR0<DataType::F32> b,
+                      Result<BufferR0<DataType::F32>> out) {
+  *out->data = *a.data - *b.data;
+  return Error::Success();
+}
+
+static Error SubtractCst(BufferR0<DataType::F32> a,
+                         Result<BufferR0<DataType::F32>> out, float cst) {
+  *out->data = *a.data - cst;
+  return Error::Success();
+}
+
+// Define XLA FFI handlers from the implementations defined above using explicit
+// XLA FFI binding API to describe type signatures of custom calls.
+
+XLA_FFI_DEFINE_HANDLER(kAlwaysFail, AlwaysFail, Ffi::Bind().Ret<BufferBase>());
+
+XLA_FFI_DEFINE_HANDLER(kAlwaysSucceed, AlwaysSucceed,
+                       Ffi::Bind().Ret<BufferBase>());
+
+XLA_FFI_DEFINE_HANDLER(kSubtract, Subtract,
+                       Ffi::Bind()
+                           .Arg<BufferR0<DataType::F32>>()
+                           .Arg<BufferR0<DataType::F32>>()
+                           .Ret<BufferR0<DataType::F32>>());
+
+XLA_FFI_DEFINE_HANDLER(kSubtractCst, SubtractCst,
+                       Ffi::Bind()
+                           .Arg<BufferR0<DataType::F32>>()
+                           .Ret<BufferR0<DataType::F32>>()
+                           .Attr<float>("cst"));
+
+template <typename T>
+static auto BindFunction(T* fn) {
+  return nb::capsule(reinterpret_cast<void*>(fn));
+}
+
+// Custom calls registration library that exports function pointers to XLA FFI
+// handlers to the python users.
+NB_MODULE(custom_calls_testlib, m) {
+  m.def("registrations", []() {
+    nb::dict dict;
+    dict["always_fail"] = BindFunction(kAlwaysFail);
+    dict["always_succeed"] = BindFunction(kAlwaysSucceed);
+    dict["subtract_f32"] = BindFunction(kSubtract);
+    dict["subtract_f32_cst"] = BindFunction(kSubtractCst);
+    return dict;
+  });
+}
+
+}  // namespace xla::ffi
diff --git a/third_party/xla/xla/python/typed_ffi_custom_call_for_test.cc b/third_party/xla/xla/python/typed_ffi_custom_call_for_test.cc
deleted file mode 100644
index f399b4e2253f77..00000000000000
--- a/third_party/xla/xla/python/typed_ffi_custom_call_for_test.cc
+++ /dev/null
@@ -1,67 +0,0 @@
-/* Copyright 2024 The OpenXLA Authors.
-
-Licensed under the Apache License, Version 2.0 (the "License");
-you may not use this file except in compliance with the License.
-You may obtain a copy of the License at
-
-    http://www.apache.org/licenses/LICENSE-2.0
-
-Unless required by applicable law or agreed to in writing, software
-distributed under the License is distributed on an "AS IS" BASIS,
-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.
-==============================================================================*/
-#include "absl/status/status.h"
-#include "third_party/nanobind/include/nanobind/nanobind.h"
-#include "xla/ffi/ffi.h"
-
-namespace xla {
-
-namespace {
-
-using F32Scalar = ffi::BufferR0<PrimitiveType::F32>;
-
-auto AlwaysFail(ffi::Result<ffi::BufferBase>) {
-  return absl::InternalError("Failed intentionally");
-}
-
-XLA_FFI_DEFINE_HANDLER(kAlwaysFail, AlwaysFail,
-                       ffi::Ffi::Bind().Ret<ffi::BufferBase>());
-
-auto AlwaysSucceed(ffi::Result<ffi::BufferBase>) { return absl::OkStatus(); }
-
-XLA_FFI_DEFINE_HANDLER(kAlwaysSucceed, AlwaysSucceed,
-                       ffi::Ffi::Bind().Ret<ffi::BufferBase>());
-
-auto SubtractN(F32Scalar in, ffi::Result<F32Scalar> out, float n) {
-  *out->data.base() = *in.data.base() - n;
-  return absl::OkStatus();
-}
-
-XLA_FFI_DEFINE_HANDLER(kSubtractN, SubtractN,
-                       ffi::Ffi::Bind()
-                           .Arg<F32Scalar>()  // in
-                           .Ret<F32Scalar>()  // out
-                           .Attr<float>("n"));
-
-namespace nb = ::nanobind;
-
-template <typename T>
-auto BindFunction(T* fn) {
-  return nb::capsule(absl::bit_cast<void*>(fn));
-}
-
-NB_MODULE(typed_ffi_custom_calls, m) {
-  m.def("registrations", []() {
-    nb::dict dict;
-    dict["xla_client_test$$AlwaysFail"] = BindFunction(kAlwaysFail);
-    dict["xla_client_test$$AlwaysSucceed"] = BindFunction(kAlwaysSucceed);
-    dict["xla_client_test$$subtract_n_f32"] = BindFunction(kSubtractN);
-    return dict;
-  });
-}
-
-}  // namespace
-
-}  // namespace xla
diff --git a/third_party/xla/xla/python/xla_client_test.py b/third_party/xla/xla/python/xla_client_test.py
index 71549ecb3e2979..c6e8e586198cc6 100644
--- a/third_party/xla/xla/python/xla_client_test.py
+++ b/third_party/xla/xla/python/xla_client_test.py
@@ -33,14 +33,9 @@
 
 # pylint: disable=g-import-not-at-top
 try:
-  from xla.python import custom_call_for_test
+  from xla.python import custom_calls_testlib
 except ImportError:
-  custom_call_for_test = None
-
-try:
-  from xla.python import typed_ffi_custom_calls
-except ImportError:
-  typed_ffi_custom_calls = None
+  custom_calls_testlib = None
 
 try:
   from xla.python import xla_extension
@@ -129,11 +124,7 @@ def setUp(self):
 
       global _CUSTOM_CALLS_REGISTERED
       if self.backend.platform == "cpu" and not _CUSTOM_CALLS_REGISTERED:
-        for name, fn in custom_call_for_test.cpu_custom_call_targets.items():
-          xla_client.register_custom_call_target(
-              name, fn, platform="cpu", api_version=0
-          )
-        for name, fn in typed_ffi_custom_calls.registrations().items():
+        for name, fn in custom_calls_testlib.registrations().items():
           xla_client.register_custom_call_target(
               name, fn, platform="cpu", api_version=1
           )
@@ -435,7 +426,7 @@ def testCustomCall(self):
       c = self._NewComputation()
       ops.CustomCallWithLayout(
           c,
-          b"test_subtract_f32",
+          b"subtract_f32",
           operands=[
               ops.Constant(c, np.float32(1.25)),
               ops.Constant(c, np.float32(0.5))
@@ -447,7 +438,7 @@ def testCustomCall(self):
               xla_client.Shape.array_shape(np.dtype(np.float32), (), ()),
           ],
           api_version=xla_client.ops.CustomCallApiVersion
-          .API_VERSION_STATUS_RETURNING)
+          .API_VERSION_TYPED_FFI)
       self._ExecuteAndCompareClose(c, expected=[0.75])
 
     def testCustomCallWithUnifiedApiUnknownTarget(self):
@@ -471,31 +462,6 @@ def testCustomCallWithUnifiedApiUnknownTarget(self):
       ):
         self._Execute(c, arguments=())
 
-    def testCustomCallWithUnifiedApi(self):
-      if self.backend.platform != "cpu":
-        self.skipTest("Test requires cpu platform")
-      c = self._NewComputation()
-
-      opaque_str = b"foo"
-      ops.CustomCallWithLayout(
-          c,
-          b"test_add_input_and_opaque_len",
-          operands=[
-              ops.Constant(c, np.float32(1.25)),
-              ops.Constant(c, np.float32(0.5))
-          ],
-          shape_with_layout=xla_client.Shape.array_shape(
-              np.dtype(np.float32), (), ()),
-          operand_shapes_with_layout=[
-              xla_client.Shape.array_shape(np.dtype(np.float32), (), ()),
-              xla_client.Shape.array_shape(np.dtype(np.float32), (), ()),
-          ],
-          # With opaque length = 3.0
-          opaque=opaque_str,
-          api_version=xla_client.ops.CustomCallApiVersion
-          .API_VERSION_STATUS_RETURNING_UNIFIED)
-      self._ExecuteAndCompareClose(c, expected=[1.25 + len(opaque_str)])
-
     def testCustomCallTypedFfiUnknownTarget(self):
       if self.backend.platform != "cpu":
         self.skipTest("Test requires cpu platform")
@@ -521,7 +487,7 @@ def testCustomCallTypedFfiAlwaysFail(self):
 
       ops.CustomCallWithLayout(
           c,
-          b"xla_client_test$$AlwaysFail",
+          b"always_fail",
           operands=[],
           shape_with_layout=xla_client.Shape.array_shape(
               np.dtype(np.float32), (), ()
@@ -542,7 +508,7 @@ def testCustomCallTypedFfiAlwaysSucceed(self):
 
       ops.CustomCallWithLayout(
           c,
-          b"xla_client_test$$AlwaysSucceed",
+          b"always_succeed",
           operands=[],
           shape_with_layout=xla_client.Shape.array_shape(
               np.dtype(np.float32), (), ()
@@ -560,7 +526,7 @@ def testCustomCallTypedFfiSubtract(self):
 
       ops.CustomCallWithLayout(
           c,
-          b"xla_client_test$$subtract_n_f32",
+          b"subtract_f32_cst",
           operands=[ops.Constant(c, np.float32(1.25))],
           shape_with_layout=xla_client.Shape.array_shape(
               np.dtype(np.float32), (), ()
@@ -568,7 +534,7 @@ def testCustomCallTypedFfiSubtract(self):
           operand_shapes_with_layout=[
               xla_client.Shape.array_shape(np.dtype(np.float32), (), ()),
           ],
-          opaque=b"{n = 3.0 : f32}",
+          opaque=b"{cst = 3.0 : f32}",
           api_version=xla_client.ops.CustomCallApiVersion.API_VERSION_TYPED_FFI,
       )
       self._ExecuteAndCompareClose(c, expected=[-1.75])
@@ -582,14 +548,7 @@ def testCustomCallLookup(self):
       self.assertTrue(_CUSTOM_CALLS_REGISTERED)
       xla_client.make_cpu_client()
       self.assertContainsSubset(
-          [
-              call.decode()
-              for call in custom_call_for_test.cpu_custom_call_targets.keys()
-          ],
-          xla_client.custom_call_targets("Host").keys(),
-      )
-      self.assertContainsSubset(
-          list(typed_ffi_custom_calls.registrations().keys()),
+          list(custom_calls_testlib.registrations().keys()),
           xla_client.custom_call_targets("Host").keys(),
       )
 

From 3b1d506d01c0fc50699a1f29e7504b0b66810dfe Mon Sep 17 00:00:00 2001
From: Anshuman Goswami <anshumang@google.com>
Date: Thu, 16 May 2024 12:38:38 -0700
Subject: [PATCH 296/478] Fix typo in comment.

PiperOrigin-RevId: 634495833
---
 .../core/common_runtime/eager/custom_device_op_handler.cc       | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorflow/core/common_runtime/eager/custom_device_op_handler.cc b/tensorflow/core/common_runtime/eager/custom_device_op_handler.cc
index e3e7f455b01123..34b89b4da61eb0 100644
--- a/tensorflow/core/common_runtime/eager/custom_device_op_handler.cc
+++ b/tensorflow/core/common_runtime/eager/custom_device_op_handler.cc
@@ -67,7 +67,7 @@ Status CustomDeviceOpHandler::Execute(ImmediateExecutionOperation* op,
   }
 
   // The op will be placed on physical device. However, it contains custom
-  // device tensor handles. The tensor handles will be copy to physical device
+  // device tensor handles. The tensor handles will be copied to physical device
   // first.
   if (op->HasCustomDeviceInput()) {
     auto inputs = op->GetInputs();

From 17ea55bbb32ab17ac82d408d0b54a8ec7d0801fa Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Thu, 16 May 2024 12:54:12 -0700
Subject: [PATCH 297/478] Move schema_fbs_with_mutable to mlir/lite/schema.

PiperOrigin-RevId: 634500044
---
 tensorflow/compiler/mlir/lite/BUILD                  |  6 +++---
 tensorflow/compiler/mlir/lite/flatbuffer_export.cc   |  2 +-
 tensorflow/compiler/mlir/lite/flatbuffer_import.cc   |  2 +-
 tensorflow/compiler/mlir/lite/flatbuffer_operator.cc |  2 +-
 tensorflow/compiler/mlir/lite/flatbuffer_operator.h  |  2 +-
 tensorflow/compiler/mlir/lite/schema/BUILD           | 12 ++++++++++++
 tensorflow/lite/schema/BUILD                         |  2 +-
 7 files changed, 20 insertions(+), 8 deletions(-)

diff --git a/tensorflow/compiler/mlir/lite/BUILD b/tensorflow/compiler/mlir/lite/BUILD
index 2cfbecb6cbbc89..0114f28b476694 100644
--- a/tensorflow/compiler/mlir/lite/BUILD
+++ b/tensorflow/compiler/mlir/lite/BUILD
@@ -1026,6 +1026,7 @@ cc_library(
         ":converter_inc",
         ":tensorflow_lite",
         "//tensorflow/compiler/mlir/lite/schema:schema_fbs",
+        "//tensorflow/compiler/mlir/lite/schema:schema_fbs_with_mutable",
         "//tensorflow/compiler/mlir/tensorflow:dynamic_shape_utils",
         "//tensorflow/compiler/mlir/tensorflow:tensorflow_types",
         "//tensorflow/core/platform:errors",
@@ -1033,7 +1034,6 @@ cc_library(
         "//tensorflow/core/platform:statusor",
         "//tensorflow/lite/core/c:private_common",
         "//tensorflow/lite/kernels/internal:kernel_utils",
-        "//tensorflow/lite/schema:schema_fbs_with_mutable",
         "//tensorflow/lite/schema:schema_utils",
         "@com_google_absl//absl/container:flat_hash_map",
         "@com_google_absl//absl/strings",
@@ -1088,6 +1088,7 @@ cc_library(
         "//tensorflow/compiler/mlir:op_or_arg_name_mapper",
         "//tensorflow/compiler/mlir/lite/metrics:error_collector_inst",
         "//tensorflow/compiler/mlir/lite/quantization/ir:QuantOps",
+        "//tensorflow/compiler/mlir/lite/schema:schema_fbs_with_mutable",
         "//tensorflow/compiler/mlir/tensorflow",
         "//tensorflow/compiler/mlir/tensorflow:convert_tensor",
         "//tensorflow/compiler/mlir/tensorflow:dynamic_shape_utils",
@@ -1105,7 +1106,6 @@ cc_library(
         "//tensorflow/lite/experimental/remat:metadata_util",
         "//tensorflow/lite/python/metrics:converter_error_data_proto_cc",
         "//tensorflow/lite/schema:schema_conversion_utils",
-        "//tensorflow/lite/schema:schema_fbs_with_mutable",
         "//tensorflow/lite/toco:toco_flags_proto_cc",
         "//tensorflow/lite/tools/versioning",
         "//tensorflow/lite/tools/versioning:gpu_compatibility",
@@ -1148,6 +1148,7 @@ cc_library(
         ":size_utils",
         ":tensorflow_lite",
         "//tensorflow/compiler/mlir/lite/quantization/ir:QuantOps",
+        "//tensorflow/compiler/mlir/lite/schema:schema_fbs_with_mutable",
         "//tensorflow/compiler/mlir/lite/stablehlo:legalize_stablehlo_composite_to_tfl_custom",
         "//tensorflow/compiler/mlir/lite/stablehlo:legalize_stablehlo_to_vhlo_pass",
         "//tensorflow/compiler/mlir/quantization/common/quantization_lib",
@@ -1162,7 +1163,6 @@ cc_library(
         "//tensorflow/core/platform:status",
         "//tensorflow/lite:framework",
         "//tensorflow/lite/experimental/remat:metadata_util",
-        "//tensorflow/lite/schema:schema_fbs_with_mutable",
         "//tensorflow/lite/schema:schema_utils",
         "@com_google_absl//absl/container:flat_hash_map",
         "@com_google_absl//absl/container:flat_hash_set",
diff --git a/tensorflow/compiler/mlir/lite/flatbuffer_export.cc b/tensorflow/compiler/mlir/lite/flatbuffer_export.cc
index 84bc7ab2daae19..0b7bd8cc7177e6 100644
--- a/tensorflow/compiler/mlir/lite/flatbuffer_export.cc
+++ b/tensorflow/compiler/mlir/lite/flatbuffer_export.cc
@@ -85,6 +85,7 @@ limitations under the License.
 #include "tensorflow/compiler/mlir/lite/ir/tfl_ops.h"
 #include "tensorflow/compiler/mlir/lite/metrics/error_collector_inst.h"
 #include "tensorflow/compiler/mlir/lite/quantization/ir/QuantOps.h"
+#include "tensorflow/compiler/mlir/lite/schema/mutable/schema_generated.h"
 #include "tensorflow/compiler/mlir/lite/utils/convert_type.h"
 #include "tensorflow/compiler/mlir/lite/utils/low_bit_utils.h"
 #include "tensorflow/compiler/mlir/lite/utils/stateful_ops_utils.h"
@@ -110,7 +111,6 @@ limitations under the License.
 #include "tensorflow/lite/experimental/remat/metadata_util.h"
 #include "tensorflow/lite/graph_info.h"
 #include "tensorflow/lite/python/metrics/converter_error_data.pb.h"
-#include "tensorflow/lite/schema/mutable/schema_generated.h"
 #include "tensorflow/lite/schema/schema_conversion_utils.h"
 #include "tensorflow/lite/string_util.h"
 #include "tensorflow/lite/toco/toco_flags.pb.h"
diff --git a/tensorflow/compiler/mlir/lite/flatbuffer_import.cc b/tensorflow/compiler/mlir/lite/flatbuffer_import.cc
index 7858c42d0673d3..bcc0244194dccd 100644
--- a/tensorflow/compiler/mlir/lite/flatbuffer_import.cc
+++ b/tensorflow/compiler/mlir/lite/flatbuffer_import.cc
@@ -77,6 +77,7 @@ limitations under the License.
 #include "tensorflow/compiler/mlir/lite/ir/tfl_ops.h"
 #include "tensorflow/compiler/mlir/lite/offset_buffer.h"
 #include "tensorflow/compiler/mlir/lite/quantization/ir/QuantOps.h"
+#include "tensorflow/compiler/mlir/lite/schema/mutable/schema_generated.h"
 #include "tensorflow/compiler/mlir/lite/stablehlo/transforms/passes.h"
 #include "tensorflow/compiler/mlir/lite/utils/const_tensor_utils.h"
 #include "tensorflow/compiler/mlir/lite/utils/convert_type.h"
@@ -96,7 +97,6 @@ limitations under the License.
 #include "tensorflow/lite/experimental/remat/metadata_util.h"
 #include "tensorflow/lite/graph_info.h"
 #include "tensorflow/lite/model_builder.h"
-#include "tensorflow/lite/schema/mutable/schema_generated.h"
 #include "tensorflow/lite/schema/schema_utils.h"
 #include "tsl/platform/status.h"
 #include "tsl/platform/statusor.h"
diff --git a/tensorflow/compiler/mlir/lite/flatbuffer_operator.cc b/tensorflow/compiler/mlir/lite/flatbuffer_operator.cc
index 0734ebca106228..3dfc21f5c2a07c 100644
--- a/tensorflow/compiler/mlir/lite/flatbuffer_operator.cc
+++ b/tensorflow/compiler/mlir/lite/flatbuffer_operator.cc
@@ -46,6 +46,7 @@ limitations under the License.
 #include "stablehlo/dialect/StablehloOps.h"  // from @stablehlo
 #include "stablehlo/dialect/VhloOps.h"  // from @stablehlo
 #include "tensorflow/compiler/mlir/lite/ir/tfl_ops.h"
+#include "tensorflow/compiler/mlir/lite/schema/mutable/schema_generated.h"
 #include "tensorflow/compiler/mlir/lite/utils/convert_type.h"
 #include "tensorflow/compiler/mlir/tensorflow/ir/tf_types.h"
 #include "tensorflow/compiler/mlir/tensorflow/utils/dynamic_shape_utils.h"
@@ -54,7 +55,6 @@ limitations under the License.
 #include "tensorflow/core/platform/status.h"
 #include "tensorflow/lite/core/c/builtin_op_data.h"
 #include "tensorflow/lite/kernels/internal/kernel_utils.h"
-#include "tensorflow/lite/schema/mutable/schema_generated.h"
 #include "tensorflow/lite/schema/schema_utils.h"
 #include "tsl/platform/status.h"
 
diff --git a/tensorflow/compiler/mlir/lite/flatbuffer_operator.h b/tensorflow/compiler/mlir/lite/flatbuffer_operator.h
index 4a54264b6244f2..64865eb77b5c43 100644
--- a/tensorflow/compiler/mlir/lite/flatbuffer_operator.h
+++ b/tensorflow/compiler/mlir/lite/flatbuffer_operator.h
@@ -34,9 +34,9 @@ limitations under the License.
 #include "stablehlo/dialect/StablehloOps.h"  // from @stablehlo
 #include "stablehlo/dialect/VhloOps.h"  // from @stablehlo
 #include "stablehlo/dialect/VhloTypes.h"  // from @stablehlo
+#include "tensorflow/compiler/mlir/lite/schema/mutable/schema_generated.h"
 #include "tensorflow/core/platform/status.h"
 #include "tensorflow/core/platform/statusor.h"
-#include "tensorflow/lite/schema/mutable/schema_generated.h"
 
 namespace mlir {
 
diff --git a/tensorflow/compiler/mlir/lite/schema/BUILD b/tensorflow/compiler/mlir/lite/schema/BUILD
index d9aed023196ad9..bb5b4434f6f5ad 100644
--- a/tensorflow/compiler/mlir/lite/schema/BUILD
+++ b/tensorflow/compiler/mlir/lite/schema/BUILD
@@ -18,3 +18,15 @@ flatbuffer_cc_library(
     srcs = ["schema.fbs"],
     compatible_with = get_compatible_with_portable(),
 )
+
+# Generic schema for flatbuffer converter (but with mutable makes bigger).
+flatbuffer_cc_library(
+    name = "schema_fbs_with_mutable",
+    srcs = ["schema.fbs"],
+    compatible_with = get_compatible_with_portable(),
+    flatc_args = [
+        "--gen-mutable",
+        "--gen-object-api",
+    ],
+    out_prefix = "mutable/",
+)
diff --git a/tensorflow/lite/schema/BUILD b/tensorflow/lite/schema/BUILD
index fafd7a24e0edb7..6a005359ccba7a 100644
--- a/tensorflow/lite/schema/BUILD
+++ b/tensorflow/lite/schema/BUILD
@@ -122,7 +122,7 @@ flatbuffer_cc_library(
 # Generic schema for flatbuffer converter (but with mutable makes bigger).
 flatbuffer_cc_library(
     name = "schema_fbs_with_mutable",
-    srcs = ["schema.fbs"],
+    srcs = ["//tensorflow/compiler/mlir/lite/schema:schema.fbs"],
     compatible_with = get_compatible_with_portable(),
     flatc_args = [
         "--gen-mutable",

From 163edc28a9216712886706e389384671ec266510 Mon Sep 17 00:00:00 2001
From: Kyle Lucke <klucke@google.com>
Date: Thu, 16 May 2024 13:23:57 -0700
Subject: [PATCH 298/478] Use absl::Status instead of xla::Status now that
 they're identical.

PiperOrigin-RevId: 634508957
---
 .../xla/xla/backends/interpreter/compiler.cc  |  2 +-
 .../xla/xla/backends/interpreter/compiler.h   |  2 +-
 .../profiler/cpu/metadata_collector.cc        |  6 +--
 .../backends/profiler/plugin/plugin_tracer.cc |  6 +--
 .../backends/profiler/plugin/plugin_tracer.h  |  6 +--
 .../xla/backends/profiler/tpu/tpu_tracer.cc   | 12 +++---
 .../xla/xla/hlo/utils/hlo_sharding_util.cc    | 12 +++---
 .../xla/xla/hlo/utils/hlo_sharding_util.h     |  2 +-
 third_party/xla/xla/mlir/utils/error_util.h   | 19 ++++-----
 third_party/xla/xla/mlir_hlo/README.md        |  4 +-
 .../xla/xla/pjrt/c/pjrt_c_api_helpers.cc      |  2 +-
 .../xla/xla/service/gpu/fusions/reduction.cc  |  8 ++--
 .../gpu/llvm_gpu_backend/gpu_backend_lib.cc   |  4 +-
 .../service/heap_simulator/heap_simulator.cc  | 24 ++++++-----
 .../service/heap_simulator/heap_simulator.h   | 16 ++++----
 .../stream_executor/device_memory_allocator.h |  5 ++-
 .../host/host_execution_engine.cc             |  3 +-
 .../xla/xla/stream_executor/lazy_op_runner.h  |  2 +-
 third_party/xla/xla/tools/hlo_opt/opt_main.cc |  2 +-
 .../xla/translate/mhlo_to_hlo/translate.cc    | 10 ++---
 .../preemption/preemption_sync_manager.cc     |  2 +-
 .../grpc_coordination_service_impl.h          |  2 +-
 .../tsl/distributed_runtime/rpc/grpc_util.h   |  2 +-
 third_party/xla/xla/tsl/util/use_cudnn.cc     | 40 ++++++++++---------
 24 files changed, 99 insertions(+), 94 deletions(-)

diff --git a/third_party/xla/xla/backends/interpreter/compiler.cc b/third_party/xla/xla/backends/interpreter/compiler.cc
index e48f63f77929ae..2409d9602671e0 100644
--- a/third_party/xla/xla/backends/interpreter/compiler.cc
+++ b/third_party/xla/xla/backends/interpreter/compiler.cc
@@ -89,7 +89,7 @@ absl::StatusOr<Literal> HandleEvaluatorCustomCall(
 
 }  // namespace
 
-Status InterpreterCompiler::RunHloOptimization(HloModule* hlo_module) {
+absl::Status InterpreterCompiler::RunHloOptimization(HloModule* hlo_module) {
   HloPassPipeline pipeline("Interpreter");
 
   // The TopkDecomposer generates a compare op with type=TOTALORDER and must
diff --git a/third_party/xla/xla/backends/interpreter/compiler.h b/third_party/xla/xla/backends/interpreter/compiler.h
index cfdbaa4fd23928..1f234f161befb1 100644
--- a/third_party/xla/xla/backends/interpreter/compiler.h
+++ b/third_party/xla/xla/backends/interpreter/compiler.h
@@ -62,7 +62,7 @@ class InterpreterCompiler : public Compiler {
   se::Platform::Id PlatformId() const override;
 
  private:
-  Status RunHloOptimization(HloModule* hlo_module);
+  absl::Status RunHloOptimization(HloModule* hlo_module);
 
   InterpreterCompiler(const InterpreterCompiler&) = delete;
   InterpreterCompiler& operator=(const InterpreterCompiler&) = delete;
diff --git a/third_party/xla/xla/backends/profiler/cpu/metadata_collector.cc b/third_party/xla/xla/backends/profiler/cpu/metadata_collector.cc
index a19d4527ec930c..0fc767d5d69e7f 100644
--- a/third_party/xla/xla/backends/profiler/cpu/metadata_collector.cc
+++ b/third_party/xla/xla/backends/profiler/cpu/metadata_collector.cc
@@ -42,7 +42,7 @@ class MetadataCollector : public tsl::profiler::ProfilerInterface {
  public:
   MetadataCollector() = default;
 
-  Status Start() override {
+  absl::Status Start() override {
     if (!trace_active_) {
       xla::XlaDebugInfoManager::Get()->StartTracing();
       trace_active_ = true;
@@ -50,7 +50,7 @@ class MetadataCollector : public tsl::profiler::ProfilerInterface {
     return OkStatus();
   }
 
-  Status Stop() override {
+  absl::Status Stop() override {
     if (trace_active_) {
       xla::XlaDebugInfoManager::Get()->StopTracing(&debug_info_);
       trace_active_ = false;
@@ -58,7 +58,7 @@ class MetadataCollector : public tsl::profiler::ProfilerInterface {
     return OkStatus();
   }
 
-  Status CollectData(tsl::profiler::XSpace* space) override {
+  absl::Status CollectData(tsl::profiler::XSpace* space) override {
     if (!debug_info_.empty()) {
       tsl::profiler::XPlane* plane =
           tsl::profiler::FindOrAddMutablePlaneWithName(
diff --git a/third_party/xla/xla/backends/profiler/plugin/plugin_tracer.cc b/third_party/xla/xla/backends/profiler/plugin/plugin_tracer.cc
index 8176ac55f553e3..5ce73a669a785f 100644
--- a/third_party/xla/xla/backends/profiler/plugin/plugin_tracer.cc
+++ b/third_party/xla/xla/backends/profiler/plugin/plugin_tracer.cc
@@ -141,7 +141,7 @@ PluginTracer::~PluginTracer() {
   }
 }
 
-Status PluginTracer::Start() {
+absl::Status PluginTracer::Start() {
   PLUGIN_Profiler_Start_Args args;
   args.profiler = profiler_;
   RETURN_STATUS_IF_PLUGIN_PROFILER_ERROR(profiler_api_->start(&args),
@@ -149,7 +149,7 @@ Status PluginTracer::Start() {
   return OkStatus();
 }
 
-Status PluginTracer::Stop() {
+absl::Status PluginTracer::Stop() {
   PLUGIN_Profiler_Stop_Args args;
   args.profiler = profiler_;
   RETURN_STATUS_IF_PLUGIN_PROFILER_ERROR(profiler_api_->stop(&args),
@@ -157,7 +157,7 @@ Status PluginTracer::Stop() {
   return OkStatus();
 }
 
-Status PluginTracer::CollectData(XSpace* space) {
+absl::Status PluginTracer::CollectData(XSpace* space) {
   PLUGIN_Profiler_CollectData_Args args;
   args.profiler = profiler_;
   args.buffer = nullptr;
diff --git a/third_party/xla/xla/backends/profiler/plugin/plugin_tracer.h b/third_party/xla/xla/backends/profiler/plugin/plugin_tracer.h
index acd88644ac090f..832a72e6546ea0 100644
--- a/third_party/xla/xla/backends/profiler/plugin/plugin_tracer.h
+++ b/third_party/xla/xla/backends/profiler/plugin/plugin_tracer.h
@@ -33,11 +33,11 @@ class PluginTracer : public tsl::profiler::ProfilerInterface {
                         const tensorflow::ProfileOptions& options);
   ~PluginTracer() override;
 
-  Status Start() override;
+  absl::Status Start() override;
 
-  Status Stop() override;
+  absl::Status Stop() override;
 
-  Status CollectData(tensorflow::profiler::XSpace* space) override;
+  absl::Status CollectData(tensorflow::profiler::XSpace* space) override;
 
  private:
   const PLUGIN_Profiler_Api* profiler_api_;
diff --git a/third_party/xla/xla/backends/profiler/tpu/tpu_tracer.cc b/third_party/xla/xla/backends/profiler/tpu/tpu_tracer.cc
index 19c7e57e49a13e..031cf651bf408d 100644
--- a/third_party/xla/xla/backends/profiler/tpu/tpu_tracer.cc
+++ b/third_party/xla/xla/backends/profiler/tpu/tpu_tracer.cc
@@ -95,11 +95,11 @@ class TpuTracer : public ProfilerInterface {
   explicit TpuTracer();
   ~TpuTracer() override;
 
-  Status Start() override;
+  absl::Status Start() override;
 
-  Status Stop() override;
+  absl::Status Stop() override;
 
-  Status CollectData(XSpace* space) override;
+  absl::Status CollectData(XSpace* space) override;
 
  private:
   TpuProfiler* tpu_profiler_;
@@ -118,7 +118,7 @@ TpuTracer::~TpuTracer() {
   stream_executor::tpu::ProfilerApiFn()->TpuProfiler_DestroyFn(tpu_profiler_);
 }
 
-Status TpuTracer::Start() {
+absl::Status TpuTracer::Start() {
   ProfilerStatusHelper status;
   stream_executor::tpu::ProfilerApiFn()->TpuProfiler_StartFn(tpu_profiler_,
                                                              status.c_status);
@@ -129,7 +129,7 @@ Status TpuTracer::Start() {
   return OkStatus();
 }
 
-Status TpuTracer::Stop() {
+absl::Status TpuTracer::Stop() {
   ProfilerStatusHelper status;
   stream_executor::tpu::ProfilerApiFn()->TpuProfiler_StopFn(tpu_profiler_,
                                                             status.c_status);
@@ -140,7 +140,7 @@ Status TpuTracer::Stop() {
   return OkStatus();
 }
 
-Status TpuTracer::CollectData(XSpace* space) {
+absl::Status TpuTracer::CollectData(XSpace* space) {
   ProfilerStatusHelper status;
   // Get size of buffer required for TPU driver to serialize XSpace into.
   size_t size_in_bytes;
diff --git a/third_party/xla/xla/hlo/utils/hlo_sharding_util.cc b/third_party/xla/xla/hlo/utils/hlo_sharding_util.cc
index 0133703e0abd68..40cd1ccb77e603 100644
--- a/third_party/xla/xla/hlo/utils/hlo_sharding_util.cc
+++ b/third_party/xla/xla/hlo/utils/hlo_sharding_util.cc
@@ -496,7 +496,7 @@ bool MergeShardingIfCompatible(const HloSharding& to_merge,
     };
     // Try to find the intersection of to_merge and dst replication groups, in
     // order to determine the merged tile assignment.
-    Status compatible =
+    absl::Status compatible =
         new_tile_array.EachStatus([&](absl::Span<const int64_t> indices,
                                       int64_t* device) -> absl::Status {
           DimensionVector to_merge_index(
@@ -1866,7 +1866,7 @@ IdentityValueAndHloOpcodeForScatterReduceComputation(
   auto computation = scatter.to_apply();
   // We only handle computations with 2 parameters and only 1 calculation.
   if (computation->instruction_count() != 3) {
-    return Status(
+    return absl::Status(
         absl::StatusCode::kInvalidArgument,
         "Expected scatter reduce computation with 2 parameters and only 1 "
         "calculation");
@@ -1893,9 +1893,9 @@ IdentityValueAndHloOpcodeForScatterReduceComputation(
                           root_instruction->opcode());
   }
 
-  return Status(absl::StatusCode::kInvalidArgument,
-                "Expected scatter reduce computation which is "
-                "add/or/multiply/add/min/max");
+  return absl::Status(absl::StatusCode::kInvalidArgument,
+                      "Expected scatter reduce computation which is "
+                      "add/or/multiply/add/min/max");
 }
 
 namespace {
@@ -3156,7 +3156,7 @@ Shape TileLeafShape(const HloSharding& sharding, const Shape& shape) {
   return result_shape;
 }
 
-Status CanonicalizeLayoutAfterShardingPropagation(
+absl::Status CanonicalizeLayoutAfterShardingPropagation(
     HloModule* module, bool update_output_layout,
     bool update_parameters_layout) {
   if (!update_output_layout && !update_parameters_layout) {
diff --git a/third_party/xla/xla/hlo/utils/hlo_sharding_util.h b/third_party/xla/xla/hlo/utils/hlo_sharding_util.h
index f4ad671e50fba8..b1269f7562cfa5 100644
--- a/third_party/xla/xla/hlo/utils/hlo_sharding_util.h
+++ b/third_party/xla/xla/hlo/utils/hlo_sharding_util.h
@@ -494,7 +494,7 @@ Shape TileLeafShape(const HloSharding& sharding, const Shape& shape);
 // that assigns layouts based on runtime shapes. Refer to
 // DetermineArgumentLayoutsFromCompileOptions() in
 // tensorflow/compiler/xla/pjrt/utils.h.
-Status CanonicalizeLayoutAfterShardingPropagation(
+absl::Status CanonicalizeLayoutAfterShardingPropagation(
     HloModule* module, bool update_output_layout,
     bool update_parameters_layout);
 
diff --git a/third_party/xla/xla/mlir/utils/error_util.h b/third_party/xla/xla/mlir/utils/error_util.h
index b37f478e173c48..d48c80ed59d0e6 100644
--- a/third_party/xla/xla/mlir/utils/error_util.h
+++ b/third_party/xla/xla/mlir/utils/error_util.h
@@ -25,13 +25,14 @@ limitations under the License.
 #include "mlir/IR/Location.h"  // from @llvm-project
 #include "mlir/IR/MLIRContext.h"  // from @llvm-project
 
-// Error utilities for MLIR when interacting with code using Status returns.
+// Error utilities for MLIR when interacting with code using absl::Status
+// returns.
 namespace mlir {
 // Diagnostic handler that collects all the diagnostics reported and can
-// produce a Status to return to callers. This is for the case where MLIR
-// functions are called from a function that will return a Status: MLIR code
-// still uses the default error reporting, and the final return function can
-// return the Status constructed from the diagnostics collected.
+// produce a absl::Status to return to callers. This is for the case where MLIR
+// functions are called from a function that will return a absl::Status: MLIR
+// code still uses the default error reporting, and the final return function
+// can return the absl::Status constructed from the diagnostics collected.
 class BaseScopedDiagnosticHandler : public SourceMgrDiagnosticHandler {
  public:
   explicit BaseScopedDiagnosticHandler(MLIRContext* context,
@@ -42,10 +43,10 @@ class BaseScopedDiagnosticHandler : public SourceMgrDiagnosticHandler {
   // Returns whether any errors were reported.
   bool ok() const;
 
-  // Returns Status corresponding to the diagnostics reported. This consumes
-  // the diagnostics reported and returns a Status of type Unknown. It is
-  // required to consume the error status, if there is one, before destroying
-  // the object.
+  // Returns absl::Status corresponding to the diagnostics reported. This
+  // consumes the diagnostics reported and returns a absl::Status of type
+  // Unknown. It is required to consume the error status, if there is one,
+  // before destroying the object.
   absl::Status ConsumeStatus();
 
   // Returns the combination of the passed in status and consumed diagnostics.
diff --git a/third_party/xla/xla/mlir_hlo/README.md b/third_party/xla/xla/mlir_hlo/README.md
index 3164f0895d41c4..0a8770aae99330 100644
--- a/third_party/xla/xla/mlir_hlo/README.md
+++ b/third_party/xla/xla/mlir_hlo/README.md
@@ -167,7 +167,7 @@ The latter could be implemented by adding the op and a legalization of the op to
 other known ops, from which a helper function can get generated that could be
 used as regular.
 
-Status: Exists but need to be cleaned up.
+absl::Status: Exists but need to be cleaned up.
 
 ### Meta HLO Dialect `mhlo`
 
@@ -210,7 +210,7 @@ intermediate optimization dialect/format. It is also where we can experiment
 cheaply with new ops. This format will be where the representation would differ
 from existing endpoints.
 
-Status: Exists but need to be cleaned up and evolved, in particular with respect
+absl::Status: Exists but need to be cleaned up and evolved, in particular with respect
 to supporting dynamic shapes.
 
 MHLO differs from XLA HLO op set in multiple ways, including:
diff --git a/third_party/xla/xla/pjrt/c/pjrt_c_api_helpers.cc b/third_party/xla/xla/pjrt/c/pjrt_c_api_helpers.cc
index bfdcda79bf427d..ac45d750adf394 100644
--- a/third_party/xla/xla/pjrt/c/pjrt_c_api_helpers.cc
+++ b/third_party/xla/xla/pjrt/c/pjrt_c_api_helpers.cc
@@ -420,7 +420,7 @@ xla::PjRtClient::HostBufferSemantics ConvertFromPjRtHostBufferSemantics(
 
 xla::PjRtFuture<> ConvertCEventToCppFuture(PJRT_Event* c_event,
                                            const PJRT_Api* c_api) {
-  using absl::Status, xla::PjRtFuture;
+  using xla::PjRtFuture;
   PJRT_Event_OnReady_Args event_onready_args;
   event_onready_args.struct_size = PJRT_Event_OnReady_Args_STRUCT_SIZE;
   event_onready_args.extension_start = nullptr;
diff --git a/third_party/xla/xla/service/gpu/fusions/reduction.cc b/third_party/xla/xla/service/gpu/fusions/reduction.cc
index c799d184483198..9856b45a516d8a 100644
--- a/third_party/xla/xla/service/gpu/fusions/reduction.cc
+++ b/third_party/xla/xla/service/gpu/fusions/reduction.cc
@@ -168,8 +168,8 @@ class ReductionEmitter {
   absl::StatusOr<std::unique_ptr<Thunk>> BuildKernelThunkForFusion(
       const LaunchDimensions& launch_dimensions,
       absl::string_view discriminator,
-      std::function<Status(std::vector<llvm_ir::IrArray>,
-                           std::vector<llvm_ir::IrArray>)>
+      std::function<absl::Status(std::vector<llvm_ir::IrArray>,
+                                 std::vector<llvm_ir::IrArray>)>
           kernel_builder_fn);
 
   absl::StatusOr<std::unique_ptr<Thunk>> BuildFusedInitializerThunk(
@@ -387,8 +387,8 @@ void ReductionEmitter::EmitSyncThreads() {
 absl::StatusOr<std::unique_ptr<Thunk>>
 ReductionEmitter::BuildKernelThunkForFusion(
     const LaunchDimensions& launch_dimensions, absl::string_view discriminator,
-    std::function<Status(std::vector<llvm_ir::IrArray>,
-                         std::vector<llvm_ir::IrArray>)>
+    std::function<absl::Status(std::vector<llvm_ir::IrArray>,
+                               std::vector<llvm_ir::IrArray>)>
         kernel_builder_fn) {
   const HloComputation* fused_computation =
       fusion_.fused_instructions_computation();
diff --git a/third_party/xla/xla/service/gpu/llvm_gpu_backend/gpu_backend_lib.cc b/third_party/xla/xla/service/gpu/llvm_gpu_backend/gpu_backend_lib.cc
index 1c5929166c8a39..7f4deb747d7c42 100644
--- a/third_party/xla/xla/service/gpu/llvm_gpu_backend/gpu_backend_lib.cc
+++ b/third_party/xla/xla/service/gpu/llvm_gpu_backend/gpu_backend_lib.cc
@@ -383,8 +383,8 @@ std::unique_ptr<llvm::TargetMachine> NVPTXGetTargetMachine(
 }
 
 using TargetModuleLinker =
-    std::function<Status(llvm::Module*, se::GpuComputeCapability,
-                         const DebugOptions&, const std::string&)>;
+    std::function<absl::Status(llvm::Module*, se::GpuComputeCapability,
+                               const DebugOptions&, const std::string&)>;
 
 void DumpModule(const std::string output_filename, const llvm::Module* module) {
   std::error_code ec;
diff --git a/third_party/xla/xla/service/heap_simulator/heap_simulator.cc b/third_party/xla/xla/service/heap_simulator/heap_simulator.cc
index 1745082331f0b4..737b9e9351b894 100644
--- a/third_party/xla/xla/service/heap_simulator/heap_simulator.cc
+++ b/third_party/xla/xla/service/heap_simulator/heap_simulator.cc
@@ -196,7 +196,7 @@ absl::StatusOr<HeapSimulator::Result<HloValue>> HeapSimulator::Run(
 
 // Runs a heap simulation for the given 'computation', assuming the given
 // 'instruction_sequence'.
-Status HeapSimulator::RunComputation(
+absl::Status HeapSimulator::RunComputation(
     const HloComputation& computation,
     const HloInstructionSequence& instruction_sequence,
     const HloAliasAnalysis& alias_analysis, HloLiveRange* hlo_live_range) {
@@ -557,7 +557,7 @@ void NoFragmentationStatsHeap<BufferType>::Free(const BufferType* buffer,
 }
 
 template <typename BufferType>
-StatusOr<HeapSimulator::Result<BufferType>>
+absl::StatusOr<HeapSimulator::Result<BufferType>>
 NoFragmentationStatsHeap<BufferType>::Finish() {
   // The result.chunk_map is empty, since we only collect stats, and don't
   // actually compute chunk assignments.
@@ -1938,10 +1938,11 @@ GlobalDecreasingSizeBestFitHeap<
 }
 
 template <typename BufferType>
-Status GlobalDecreasingSizeBestFitHeap<BufferType>::SlicedAllocationFinder::
-    DoesPermutationFit(absl::Span<const int64_t> permutation_of_slice_times,
-                       const FreeChunkRoot& root, int64_t offset) const {
-  Status result =
+absl::Status GlobalDecreasingSizeBestFitHeap<BufferType>::
+    SlicedAllocationFinder::DoesPermutationFit(
+        absl::Span<const int64_t> permutation_of_slice_times,
+        const FreeChunkRoot& root, int64_t offset) const {
+  absl::Status result =
       DoesPermutationFitImpl(permutation_of_slice_times, root, offset);
   VLOG(3) << "SlicedAllocationFinder::DoesPermutationFit\n"
           << "  permutation of slice times: [ "
@@ -1953,9 +1954,10 @@ Status GlobalDecreasingSizeBestFitHeap<BufferType>::SlicedAllocationFinder::
 }
 
 template <typename BufferType>
-Status GlobalDecreasingSizeBestFitHeap<BufferType>::SlicedAllocationFinder::
-    DoesPermutationFitImpl(absl::Span<const int64_t> permutation_of_slice_times,
-                           const FreeChunkRoot& root, int64_t offset) const {
+absl::Status GlobalDecreasingSizeBestFitHeap<BufferType>::
+    SlicedAllocationFinder::DoesPermutationFitImpl(
+        absl::Span<const int64_t> permutation_of_slice_times,
+        const FreeChunkRoot& root, int64_t offset) const {
   if (permutation_of_slice_times.size() != sorted_slice_sizes_.size()) {
     return InvalidArgumentStrCat(
         sorted_slice_sizes_.size(), " slices times expected in permutation. ",
@@ -2119,7 +2121,7 @@ GlobalDecreasingSizeBestFitHeap<BufferType>::SlicedAllocationFinder::
 }
 
 template <typename BufferType>
-StatusOr<HeapSimulator::Result<BufferType>>
+absl::StatusOr<HeapSimulator::Result<BufferType>>
 GlobalDecreasingSizeBestFitHeap<BufferType>::Finish() {
   std::vector<BufferInterval> sorted_buffer_intervals =
       GetSortedBufferIntervals();
@@ -2396,7 +2398,7 @@ ConstrainedGlobalDecreasingSizeBestFitHeap::Finish() {
 }
 
 template <typename BufferType>
-StatusOr<HeapSimulator::Result<BufferType>>
+absl::StatusOr<HeapSimulator::Result<BufferType>>
 ChooseBestHeapAlgorithm<BufferType>::Finish() {
   DCHECK(!algorithms_.empty());
   std::vector<Result> results(algorithms_.size());
diff --git a/third_party/xla/xla/service/heap_simulator/heap_simulator.h b/third_party/xla/xla/service/heap_simulator/heap_simulator.h
index 849be24cb4b086..d6eae82ce8aa6f 100644
--- a/third_party/xla/xla/service/heap_simulator/heap_simulator.h
+++ b/third_party/xla/xla/service/heap_simulator/heap_simulator.h
@@ -215,10 +215,10 @@ class HeapSimulator {
                     memory_by_computation = nullptr);
   ~HeapSimulator();
 
-  Status RunComputation(const HloComputation& computation,
-                        const HloInstructionSequence& instruction_sequence,
-                        const HloAliasAnalysis& alias_analysis,
-                        HloLiveRange* live_range);
+  absl::Status RunComputation(
+      const HloComputation& computation,
+      const HloInstructionSequence& instruction_sequence,
+      const HloAliasAnalysis& alias_analysis, HloLiveRange* live_range);
 
   bool IgnoreBuffer(const HloValue* buffer) const;
   void Alloc(const HloValue* buffer, const HloInstruction* instruction);
@@ -756,14 +756,14 @@ class GlobalDecreasingSizeBestFitHeap : public HeapAlgorithm<BufferType> {
     // (spatially) should be allocated. Such a slice has size
     // sorted_slice_sizes_[i] and would be allocated at offset +
     // sum(sorted_slice_sizes[j], for j in [0, i-1]).
-    Status DoesPermutationFit(
+    absl::Status DoesPermutationFit(
         absl::Span<const int64_t> permutation_of_slice_times,
         const FreeChunkRoot& root, int64_t offset) const;
 
     // Only DoesSlicedPermutationFit() should call this method directly. Other
     // callers should call DoesSlicedPermutationFit(), which contains some
     // wrapper VLOGGING.
-    Status DoesPermutationFitImpl(
+    absl::Status DoesPermutationFitImpl(
         absl::Span<const int64_t> permutation_of_slice_times,
         const FreeChunkRoot& root, int64_t offset) const;
 
@@ -811,7 +811,7 @@ class GlobalDecreasingSizeBestFitHeap : public HeapAlgorithm<BufferType> {
   void ShareWith(const BufferType* buffer, const BufferType* share_with,
                  int64_t size) override;
 
-  StatusOr<Result> Finish() override;
+  absl::StatusOr<Result> Finish() override;
 
   // Return a BufferIntervalCompare function that sort by spatial size. We don't
   // look at co-locates as they should have the same size.
@@ -992,7 +992,7 @@ class ChooseBestHeapAlgorithm : public HeapAlgorithm<BufferType> {
     }
   }
 
-  StatusOr<Result> Finish() override;
+  absl::StatusOr<Result> Finish() override;
 
  private:
   std::vector<std::unique_ptr<HeapAlgorithm<BufferType>>> algorithms_;
diff --git a/third_party/xla/xla/stream_executor/device_memory_allocator.h b/third_party/xla/xla/stream_executor/device_memory_allocator.h
index bc3b20d566d4d5..bab6812458036b 100644
--- a/third_party/xla/xla/stream_executor/device_memory_allocator.h
+++ b/third_party/xla/xla/stream_executor/device_memory_allocator.h
@@ -161,8 +161,9 @@ class DeviceMemoryAllocator {
 
   // Allocates memory on the device.
   //
-  // If size > 0 and the returned StatusOr is OK, the wrapped OwningDeviceMemory
-  // must not be null.  If size == 0, must return a null OwningDeviceMemory.
+  // If size > 0 and the returned absl::StatusOr is OK, the wrapped
+  // OwningDeviceMemory must not be null.  If size == 0, must return a null
+  // OwningDeviceMemory.
   //
   // 'retry_on_failure': If false, and the first attempt to allocate the memory
   // fails, the allocation should return immediately without retrying.  An
diff --git a/third_party/xla/xla/stream_executor/host/host_execution_engine.cc b/third_party/xla/xla/stream_executor/host/host_execution_engine.cc
index f48203bfb3166c..0ed748e633bd05 100644
--- a/third_party/xla/xla/stream_executor/host/host_execution_engine.cc
+++ b/third_party/xla/xla/stream_executor/host/host_execution_engine.cc
@@ -67,7 +67,6 @@ namespace stream_executor {
 namespace host {
 
 using absl::InternalError;
-using absl::StatusOr;
 using absl::StrFormat;
 
 using llvm::Expected;
@@ -198,7 +197,7 @@ static std::string ToString(const llvm::Error &err) {
   return str;
 }
 
-/*static*/ StatusOr<std::unique_ptr<RuntimeExecutionEngine>>
+/*static*/ absl::StatusOr<std::unique_ptr<RuntimeExecutionEngine>>
 RuntimeExecutionEngine::CreateFromModule(
     std::unique_ptr<llvm::LLVMContext> ctx,
     std::unique_ptr<llvm::Module> module, JitOptions options,
diff --git a/third_party/xla/xla/stream_executor/lazy_op_runner.h b/third_party/xla/xla/stream_executor/lazy_op_runner.h
index d4265f6852f705..c211244b426542 100644
--- a/third_party/xla/xla/stream_executor/lazy_op_runner.h
+++ b/third_party/xla/xla/stream_executor/lazy_op_runner.h
@@ -66,7 +66,7 @@ inline absl::StatusOr<DnnSupport*> GetDnnFromStream(Stream* stream) {
 //   struct Config;
 //
 //   // Use a StreamExecutor to create an OpRunner.
-//   static StatusOr<OpRunner<Config>> OpRunnerFromDesc(
+//   static absl::StatusOr<OpRunner<Config>> OpRunnerFromDesc(
 //       const AlgorithmDesc& desc, Config config, StreamExecutor* stream);
 // };
 template <typename Op>
diff --git a/third_party/xla/xla/tools/hlo_opt/opt_main.cc b/third_party/xla/xla/tools/hlo_opt/opt_main.cc
index fca335e60647d9..57ff4480f3a085 100644
--- a/third_party/xla/xla/tools/hlo_opt/opt_main.cc
+++ b/third_party/xla/xla/tools/hlo_opt/opt_main.cc
@@ -168,7 +168,7 @@ absl::StatusOr<std::string> TranslateToStage(int argc, char** argv,
   return out_combined;
 }
 
-Status RunOpt(int argc, char** argv, const HloOptConfig& opts) {
+absl::Status RunOpt(int argc, char** argv, const HloOptConfig& opts) {
   TF_ASSIGN_OR_RETURN(std::string output, TranslateToStage(argc, argv, opts));
   if (opts.output_file == "-") {
     std::cout << output << std::endl;
diff --git a/third_party/xla/xla/translate/mhlo_to_hlo/translate.cc b/third_party/xla/xla/translate/mhlo_to_hlo/translate.cc
index 3d56686a97e2ac..31a0e1b789b4cc 100644
--- a/third_party/xla/xla/translate/mhlo_to_hlo/translate.cc
+++ b/third_party/xla/xla/translate/mhlo_to_hlo/translate.cc
@@ -32,7 +32,7 @@ mlir::LogicalResult MlirHloToHloTranslateFunction(mlir::ModuleOp module,
   if (!module) return mlir::failure();
 
   HloProto hloProto;
-  Status status = mlir::ConvertMlirHloToHlo(
+  absl::Status status = mlir::ConvertMlirHloToHlo(
       module, &hloProto, emit_use_tuple_arg, emit_return_tuple);
   if (!status.ok()) {
     module.emitOpError() << status.message();
@@ -55,9 +55,9 @@ absl::StatusOr<std::unique_ptr<HloModule>> HloModuleFromProto(
 
 // Wraps BuildHloFromMlirHlo to output an HloProto that's the same as
 // ConvertMlirHloToHlo.
-Status ConvertMlirHloToHloViaBuilder(mlir::ModuleOp module,
-                                     ::xla::HloProto* hlo_proto,
-                                     mlir::MlirToHloConversionOptions options) {
+absl::Status ConvertMlirHloToHloViaBuilder(
+    mlir::ModuleOp module, ::xla::HloProto* hlo_proto,
+    mlir::MlirToHloConversionOptions options) {
   mlir::func::FuncOp main = module.lookupSymbol<mlir::func::FuncOp>("main");
   mlir::Block& block = main.getRegion().front();
   xla::XlaBuilder builder("main");
@@ -117,7 +117,7 @@ mlir::LogicalResult MlirHloToHloTextTranslateFunction(
   HloProto hloProto;
   mlir::MlirToHloConversionOptions options;
   options.propagate_layouts = with_layouts;
-  Status status =
+  absl::Status status =
       via_builder
           ? ConvertMlirHloToHloViaBuilder(module, &hloProto, options)
           : mlir::ConvertMlirHloToHlo(module, &hloProto, emit_use_tuple_arg,
diff --git a/third_party/xla/xla/tsl/distributed_runtime/preemption/preemption_sync_manager.cc b/third_party/xla/xla/tsl/distributed_runtime/preemption/preemption_sync_manager.cc
index 923c7236f83c84..dbe23826d355e8 100644
--- a/third_party/xla/xla/tsl/distributed_runtime/preemption/preemption_sync_manager.cc
+++ b/third_party/xla/xla/tsl/distributed_runtime/preemption/preemption_sync_manager.cc
@@ -153,7 +153,7 @@ absl::Status PreemptionSyncManagerImpl::Initialize(
             kPreemptionNoticeKey, absl::FormatTime(*death_time));
         LOG(INFO) << "Notified coordination service that this task will "
                      "be preempted at "
-                  << *death_time << ". Status: " << s;
+                  << *death_time << ". absl::Status: " << s;
       });
 
   /* Listen for preemption notice (death time) from coordination service, which
diff --git a/third_party/xla/xla/tsl/distributed_runtime/rpc/coordination/grpc_coordination_service_impl.h b/third_party/xla/xla/tsl/distributed_runtime/rpc/coordination/grpc_coordination_service_impl.h
index ac27e6ba28d76b..21441ad290a26c 100644
--- a/third_party/xla/xla/tsl/distributed_runtime/rpc/coordination/grpc_coordination_service_impl.h
+++ b/third_party/xla/xla/tsl/distributed_runtime/rpc/coordination/grpc_coordination_service_impl.h
@@ -67,7 +67,7 @@ class GrpcCoordinationServiceImpl : public AsyncServiceInterface {
     }                                                                         \
     compute_pool_.Schedule([this, call]() {                                   \
       rpc_handler_.method##Async(&call->request, &call->response,             \
-                                 [call](const Status& s) {                    \
+                                 [call](const absl::Status& s) {              \
                                    call->ClearCancelCallback();               \
                                    call->SendResponse(ToGrpcStatus(s));       \
                                  });                                          \
diff --git a/third_party/xla/xla/tsl/distributed_runtime/rpc/grpc_util.h b/third_party/xla/xla/tsl/distributed_runtime/rpc/grpc_util.h
index 4c5173c79a9a80..1b83ce6190408d 100644
--- a/third_party/xla/xla/tsl/distributed_runtime/rpc/grpc_util.h
+++ b/third_party/xla/xla/tsl/distributed_runtime/rpc/grpc_util.h
@@ -33,7 +33,7 @@ namespace tsl {
 
 // Proto: tensorflow::distributed_runtime::GrpcPayloadsLost
 // Location: tsl/protobuf/distributed_runtime_payloads.proto
-// Usage: Flags the Status to have lost payloads during GRPC conversion.
+// Usage: Flags the absl::Status to have lost payloads during GRPC conversion.
 constexpr char kGrpcPayloadsLost[] =
     "type.googleapis.com/tensorflow.distributed_runtime.GrpcPayloadsLost";
 
diff --git a/third_party/xla/xla/tsl/util/use_cudnn.cc b/third_party/xla/xla/tsl/util/use_cudnn.cc
index a3e1b4d25d2667..da15cadca6315f 100644
--- a/third_party/xla/xla/tsl/util/use_cudnn.cc
+++ b/third_party/xla/xla/tsl/util/use_cudnn.cc
@@ -27,14 +27,15 @@ limitations under the License.
 
 namespace tsl {
 
-#define ADD_BOOL_CUDNN_FLAG(func_name, flag_name, default_value)           \
-  bool func_name() {                                                       \
-    bool value = default_value;                                            \
-    Status status = ReadBoolFromEnvVar(#flag_name, default_value, &value); \
-    if (!status.ok()) {                                                    \
-      LOG(ERROR) << status;                                                \
-    }                                                                      \
-    return value;                                                          \
+#define ADD_BOOL_CUDNN_FLAG(func_name, flag_name, default_value) \
+  bool func_name() {                                             \
+    bool value = default_value;                                  \
+    absl::Status status =                                        \
+        ReadBoolFromEnvVar(#flag_name, default_value, &value);   \
+    if (!status.ok()) {                                          \
+      LOG(ERROR) << status;                                      \
+    }                                                            \
+    return value;                                                \
   }
 
 bool CudnnUseFrontend() {
@@ -43,8 +44,8 @@ bool CudnnUseFrontend() {
 #if GOOGLE_CUDA
     if (CUDNN_VERSION >= 8100) {
       // cuDNN 8.1.0 + the frontend has issues regarding fused convolution.
-      Status status = ReadBoolFromEnvVar("TF_CUDNN_USE_FRONTEND",
-                                         CUDNN_VERSION >= 8200, &value);
+      absl::Status status = ReadBoolFromEnvVar("TF_CUDNN_USE_FRONTEND",
+                                               CUDNN_VERSION >= 8200, &value);
       if (!status.ok()) {
         LOG(ERROR) << status;
       }
@@ -64,7 +65,7 @@ bool CudnnUseRuntimeFusion() {
     bool value = false;
 #if GOOGLE_CUDA
     if (CUDNN_VERSION >= 8400) {
-      Status status =
+      absl::Status status =
           ReadBoolFromEnvVar("TF_CUDNN_USE_RUNTIME_FUSION", false, &value);
       if (!status.ok()) {
         LOG(ERROR) << status;
@@ -99,14 +100,15 @@ ADD_BOOL_CUDNN_FLAG(DebugCudnnRnnUseTensorOps,
                     TF_DEBUG_CUDNN_RNN_USE_TENSOR_OPS, false);
 #undef ADD_BOOL_CUDNN_FLAG
 
-#define ADD_INT64_CUDNN_FLAG(func_name, flag_name, default_value)           \
-  int64_t func_name() {                                                     \
-    int64_t value = default_value;                                          \
-    Status status = ReadInt64FromEnvVar(#flag_name, default_value, &value); \
-    if (!status.ok()) {                                                     \
-      LOG(ERROR) << status;                                                 \
-    }                                                                       \
-    return value;                                                           \
+#define ADD_INT64_CUDNN_FLAG(func_name, flag_name, default_value) \
+  int64_t func_name() {                                           \
+    int64_t value = default_value;                                \
+    absl::Status status =                                         \
+        ReadInt64FromEnvVar(#flag_name, default_value, &value);   \
+    if (!status.ok()) {                                           \
+      LOG(ERROR) << status;                                       \
+    }                                                             \
+    return value;                                                 \
   }
 // Cudnn RNN algorithm to use for both forward and backward pass. Only effective
 // when TF_DEBUG_CUDNN_RNN is true. See Nvidia Cudnn manual for allowed

From fbe2a689eacfb6d7aede1e220f6292659d954ace Mon Sep 17 00:00:00 2001
From: Kyle Lucke <klucke@google.com>
Date: Thu, 16 May 2024 13:30:16 -0700
Subject: [PATCH 299/478] Use absl::Status instead of xla::Status now that
 they're identical.

PiperOrigin-RevId: 634511028
---
 .../xla/pjrt/cpu/abstract_tfrt_cpu_buffer.cc   | 15 +++++++++------
 .../xla/pjrt/cpu/abstract_tfrt_cpu_buffer.h    | 10 +++++-----
 third_party/xla/xla/pjrt/cpu/cpu_client.cc     | 17 +++++++++--------
 third_party/xla/xla/pjrt/cpu/cpu_client.h      | 14 +++++++-------
 third_party/xla/xla/pjrt/distributed/client.cc | 16 ++++++++--------
 third_party/xla/xla/pjrt/distributed/client.h  |  5 +++--
 .../xla/pjrt/distributed/client_server_test.cc |  2 +-
 .../xla/xla/pjrt/distributed/topology_util.cc  | 13 +++++++------
 .../xla/xla/pjrt/distributed/topology_util.h   | 13 +++++++------
 third_party/xla/xla/pjrt/distributed/util.h    |  8 ++++----
 third_party/xla/xla/pjrt/gpu/gpu_helpers.cc    |  8 ++++----
 .../xla/xla/pjrt/gpu/se_gpu_pjrt_client.cc     | 10 +++++-----
 .../xla/xla/pjrt/gpu/se_gpu_pjrt_client.h      |  2 +-
 .../xla/pjrt/gpu/se_gpu_pjrt_client_test.cc    |  8 ++++----
 .../xla/xla/service/gpu/runtime/annotation.cc  |  2 +-
 .../runtime/nccl_collective_broadcast_thunk.cc | 11 ++++++-----
 .../runtime/nccl_collective_broadcast_thunk.h  | 18 ++++++++++--------
 .../gpu/runtime/nccl_collective_thunk.cc       | 14 +++++++-------
 .../gpu/runtime/nccl_collective_thunk.h        |  8 ++++----
 .../gpu/runtime/nccl_p2p_thunk_common.cc       |  6 +++---
 20 files changed, 105 insertions(+), 95 deletions(-)

diff --git a/third_party/xla/xla/pjrt/cpu/abstract_tfrt_cpu_buffer.cc b/third_party/xla/xla/pjrt/cpu/abstract_tfrt_cpu_buffer.cc
index a52622a28e1fb3..d6f851f8c6febd 100644
--- a/third_party/xla/xla/pjrt/cpu/abstract_tfrt_cpu_buffer.cc
+++ b/third_party/xla/xla/pjrt/cpu/abstract_tfrt_cpu_buffer.cc
@@ -349,7 +349,7 @@ AbstractTfrtCpuBuffer::Release(bool wait_for_operations_to_complete) {
     // Block the host until all usage events have completed. Usage events
     // dominate definition events, so this also waits for the buffer to be
     // defined. Return the first error encountered.
-    Status first_error;
+    absl::Status first_error;
     for (const auto& av : events) {
       BlockUntilReady(av.GetAsyncValue());
       if (auto* error = av.GetErrorIfPresent()) {
@@ -879,7 +879,8 @@ AbstractAsyncHostToHostMemoryTransferManager::RetrieveBuffer(int buffer_index) {
   return std::move(buffers_[buffer_index]);
 }
 
-Status AbstractAsyncHostToHostMemoryTransferManager::TransferLiteralToBuffer(
+absl::Status
+AbstractAsyncHostToHostMemoryTransferManager::TransferLiteralToBuffer(
     int buffer_index, const LiteralSlice& literal,
     absl::AnyInvocable<void() &&> on_done) {
   return TransferRawDataToSubBuffer(buffer_index, literal.untyped_data(),
@@ -888,7 +889,8 @@ Status AbstractAsyncHostToHostMemoryTransferManager::TransferLiteralToBuffer(
                                     std::move(on_done));
 }
 
-Status AbstractAsyncHostToHostMemoryTransferManager::TransferRawDataToBuffer(
+absl::Status
+AbstractAsyncHostToHostMemoryTransferManager::TransferRawDataToBuffer(
     int buffer_index, absl::string_view data,
     absl::AnyInvocable<void() &&> on_done) {
   return TransferRawDataToSubBuffer(
@@ -898,7 +900,8 @@ Status AbstractAsyncHostToHostMemoryTransferManager::TransferRawDataToBuffer(
 
 // The definition events of `device_buffers_` must be ready before calling this
 // function.
-Status AbstractAsyncHostToHostMemoryTransferManager::TransferRawDataToSubBuffer(
+absl::Status
+AbstractAsyncHostToHostMemoryTransferManager::TransferRawDataToSubBuffer(
     int buffer_index, const void* data, int64_t offset, int64_t transfer_size,
     bool is_last_transfer, absl::AnyInvocable<void() &&> on_done) {
   {
@@ -947,12 +950,12 @@ Status AbstractAsyncHostToHostMemoryTransferManager::TransferRawDataToSubBuffer(
 }
 
 void AbstractAsyncHostToHostMemoryTransferManager::SetBufferError(
-    int buffer_index, Status error) {
+    int buffer_index, absl::Status error) {
   absl::MutexLock l(&mu_);
   avs_[buffer_index]->SetError(error);
 }
 
-/*static*/ Status
+/*static*/ absl::Status
 AbstractAsyncHostToHostMemoryTransferManager::PopulateAsyncTransferManagerData(
     absl::Span<const std::unique_ptr<AbstractTfrtCpuBuffer>> buffers,
     absl::InlinedVector<TrackedTfrtCpuDeviceBuffer*, 4>& device_buffers,
diff --git a/third_party/xla/xla/pjrt/cpu/abstract_tfrt_cpu_buffer.h b/third_party/xla/xla/pjrt/cpu/abstract_tfrt_cpu_buffer.h
index 5c5c72eff731a2..f21e4546d47de2 100644
--- a/third_party/xla/xla/pjrt/cpu/abstract_tfrt_cpu_buffer.h
+++ b/third_party/xla/xla/pjrt/cpu/abstract_tfrt_cpu_buffer.h
@@ -320,19 +320,19 @@ class AbstractAsyncHostToHostMemoryTransferManager
 
   std::unique_ptr<PjRtBuffer> RetrieveBuffer(int buffer_index) override;
 
-  Status TransferLiteralToBuffer(
+  absl::Status TransferLiteralToBuffer(
       int buffer_index, const LiteralSlice& literal,
       absl::AnyInvocable<void() &&> on_done) override;
 
-  Status TransferRawDataToBuffer(
+  absl::Status TransferRawDataToBuffer(
       int buffer_index, absl::string_view data,
       absl::AnyInvocable<void() &&> on_done) override;
 
-  Status TransferRawDataToSubBuffer(
+  absl::Status TransferRawDataToSubBuffer(
       int buffer_index, const void* data, int64_t offset, int64_t transfer_size,
       bool is_last_transfer, absl::AnyInvocable<void() &&> on_done) override;
 
-  void SetBufferError(int buffer_index, Status error) override;
+  void SetBufferError(int buffer_index, absl::Status error) override;
 
   void AddTransferMetadata(const TransferMetadata& meta) override {
     LOG(WARNING) << "AddTransferMetadata not implemented for "
@@ -351,7 +351,7 @@ class AbstractAsyncHostToHostMemoryTransferManager
 
   // Initialize `device_buffers`, `buffer_sizes`, `buffer_transfers_in_flight`,
   // and `last_transfer_finished` from `buffers`.
-  static Status PopulateAsyncTransferManagerData(
+  static absl::Status PopulateAsyncTransferManagerData(
       absl::Span<const std::unique_ptr<AbstractTfrtCpuBuffer>> buffers,
       absl::InlinedVector<TrackedTfrtCpuDeviceBuffer*, 4>& device_buffers,
       absl::InlinedVector<size_t, 4>& buffer_sizes,
diff --git a/third_party/xla/xla/pjrt/cpu/cpu_client.cc b/third_party/xla/xla/pjrt/cpu/cpu_client.cc
index 7416654899de7f..a0972f795f3e03 100644
--- a/third_party/xla/xla/pjrt/cpu/cpu_client.cc
+++ b/third_party/xla/xla/pjrt/cpu/cpu_client.cc
@@ -304,11 +304,12 @@ TfrtCpuDevice::TfrtCpuDevice(int id, int process_index, int local_hardware_id,
       max_inflight_computations_semaphore_(
           /*capacity=*/max_inflight_computations) {}
 
-Status TfrtCpuDevice::TransferToInfeed(const LiteralSlice& literal) {
+absl::Status TfrtCpuDevice::TransferToInfeed(const LiteralSlice& literal) {
   return TransferLiteralToInfeedOnCpu(local_hardware_id(), literal);
 }
 
-Status TfrtCpuDevice::TransferFromOutfeed(MutableBorrowingLiteral literal) {
+absl::Status TfrtCpuDevice::TransferFromOutfeed(
+    MutableBorrowingLiteral literal) {
   return TransferLiteralFromOutfeedOnCpu(local_hardware_id(), literal);
 }
 
@@ -903,7 +904,7 @@ TfrtCpuClient::CreateViewOfDeviceBuffer(
 }
 
 absl::StatusOr<std::unique_ptr<PjRtBuffer>> TfrtCpuClient::CreateErrorBuffer(
-    Status error, const Shape& shape, PjRtDevice* device) {
+    absl::Status error, const Shape& shape, PjRtDevice* device) {
   if (device->client() != this) {
     return absl::InvalidArgumentError("Device is not attached to this client");
   }
@@ -920,7 +921,7 @@ absl::StatusOr<std::unique_ptr<PjRtBuffer>> TfrtCpuClient::CreateErrorBuffer(
 }
 
 absl::StatusOr<std::unique_ptr<PjRtBuffer>> TfrtCpuClient::CreateErrorBuffer(
-    Status error, const Shape& shape, PjRtMemorySpace* memory) {
+    absl::Status error, const Shape& shape, PjRtMemorySpace* memory) {
   return CreateErrorBuffer(std::move(error), shape, memory->devices()[0]);
 }
 
@@ -1171,7 +1172,7 @@ absl::StatusOr<std::optional<std::string>> TfrtCpuExecutable::Fingerprint()
   return std::optional<std::string>();
 }
 
-Status TfrtCpuExecutable::SetUpDonation(bool tuple_inputs) {
+absl::Status TfrtCpuExecutable::SetUpDonation(bool tuple_inputs) {
   TF_ASSIGN_OR_RETURN(parameters_that_must_be_donated_,
                       ComputeParametersThatMustBeDonated(
                           *cpu_executable_->shared_module(), tuple_inputs));
@@ -1311,7 +1312,7 @@ static absl::InlinedVector<BufferInfo, 4> CreateResultBufferInfo(
   return output_buffer_info;
 }
 
-Status TfrtCpuExecutable::CheckBufferCompatibilities(
+absl::Status TfrtCpuExecutable::CheckBufferCompatibilities(
     absl::Span<std::pair<bool, TrackedTfrtCpuDeviceBuffer*> const>
         input_buffers) const {
   if (input_buffers.size() != input_buffer_sizes_in_bytes_.size()) {
@@ -1409,7 +1410,7 @@ absl::StatusOr<PjRtLoadedExecutable::Result> TfrtCpuExecutable::ExecuteHelper(
     }
 
     TrackedTfrtCpuDeviceBuffer* tracked_buffer;
-    auto get_buffer = [&](int i) -> Status {
+    auto get_buffer = [&](int i) -> absl::Status {
       bool must_donate = donate_it != parameters_that_must_be_donated_.end() &&
                          *donate_it == i;
       TF_RETURN_IF_ERROR(TestBufferDonationClashes(
@@ -1838,7 +1839,7 @@ TfrtCpuExecutable::Execute(
     absl::Mutex mu;
     int running = num_addressable_devices;
     int failed = 0;
-    Status first_failure_status;
+    absl::Status first_failure_status;
 
     for (int i = 0; i < num_addressable_devices; ++i) {
       const int replica = addressable_device_logical_ids_[i].replica;
diff --git a/third_party/xla/xla/pjrt/cpu/cpu_client.h b/third_party/xla/xla/pjrt/cpu/cpu_client.h
index c704f2b38a2f81..310c6114467877 100644
--- a/third_party/xla/xla/pjrt/cpu/cpu_client.h
+++ b/third_party/xla/xla/pjrt/cpu/cpu_client.h
@@ -229,9 +229,9 @@ class TfrtCpuDevice final : public PjRtDevice {
     return PjRtLocalHardwareId(description_.local_hardware_id());
   }
 
-  Status TransferToInfeed(const LiteralSlice& literal) override;
+  absl::Status TransferToInfeed(const LiteralSlice& literal) override;
 
-  Status TransferFromOutfeed(MutableBorrowingLiteral literal) override;
+  absl::Status TransferFromOutfeed(MutableBorrowingLiteral literal) override;
 
   void AttachMemorySpace(PjRtMemorySpace* memory_space);
 
@@ -330,10 +330,10 @@ class TfrtCpuClient final : public PjRtClient {
       std::optional<CompileOptions> options) override;
 
   absl::StatusOr<std::unique_ptr<PjRtBuffer>> CreateErrorBuffer(
-      Status error, const Shape& shape, PjRtDevice* device) override;
+      absl::Status error, const Shape& shape, PjRtDevice* device) override;
 
   absl::StatusOr<std::unique_ptr<PjRtBuffer>> CreateErrorBuffer(
-      Status error, const Shape& shape, PjRtMemorySpace* memory) override;
+      absl::Status error, const Shape& shape, PjRtMemorySpace* memory) override;
 
   absl::StatusOr<std::unique_ptr<PjRtBuffer>> CreateUninitializedBuffer(
       const Shape& shape, PjRtDevice* device) override;
@@ -403,7 +403,7 @@ class TfrtCpuClient final : public PjRtClient {
     return Unimplemented("CreateHostToDeviceChannelHandle not implemented.");
   }
 
-  Status Defragment() override {
+  absl::Status Defragment() override {
     return Unimplemented("Defragment not implemented.");
   }
 
@@ -649,11 +649,11 @@ class TfrtCpuExecutable final : public PjRtLoadedExecutable {
  private:
   friend class TfrtCpuClient;
 
-  Status SetUpDonation(bool tuple_inputs);
+  absl::Status SetUpDonation(bool tuple_inputs);
 
   // Checks that the input buffers passed in by the user have the correct size
   // on device for the compiled program.
-  Status CheckBufferCompatibilities(
+  absl::Status CheckBufferCompatibilities(
       absl::Span<std::pair<bool, TrackedTfrtCpuDeviceBuffer*> const>
           input_buffers) const;
 
diff --git a/third_party/xla/xla/pjrt/distributed/client.cc b/third_party/xla/xla/pjrt/distributed/client.cc
index bb55b4d02ffdb4..74f30ff06ccc4e 100644
--- a/third_party/xla/xla/pjrt/distributed/client.cc
+++ b/third_party/xla/xla/pjrt/distributed/client.cc
@@ -89,16 +89,16 @@ DistributedRuntimeCoordinationServiceClient::
       absl::ToInt64Milliseconds(options.shutdown_timeout));
   config.set_agent_destruction_without_shutdown(
       !options.shutdown_on_destruction);
-  auto error_fn =
-      [timeout_fn = options.missed_heartbeat_callback](const Status& status) {
-        LOG(ERROR) << "Coordination service agent in error status: " << status;
-        timeout_fn(status, /*coordinator_reported_failure=*/true);
-      };
+  auto error_fn = [timeout_fn = options.missed_heartbeat_callback](
+                      const absl::Status& status) {
+    LOG(ERROR) << "Coordination service agent in error status: " << status;
+    timeout_fn(status, /*coordinator_reported_failure=*/true);
+  };
 
   std::unique_ptr<tsl::CoordinationClient> leader_client;
   leader_client.reset(tsl::NewGrpcCoordinationClient(channel));
   coord_agent_ = tsl::CreateCoordinationServiceAgent();
-  const Status status =
+  const absl::Status status =
       coord_agent_->Initialize(options.env, "jax_worker", options.node_id,
                                config, std::move(leader_client), error_fn);
   if (!status.ok()) {
@@ -116,7 +116,7 @@ absl::Status DistributedRuntimeCoordinationServiceClient::Connect() {
       absl::Now() +
       absl::Milliseconds(config_.cluster_register_timeout_in_ms());
 
-  Status s = coord_agent_->Connect();
+  absl::Status s = coord_agent_->Connect();
   if (s.ok()) {
     absl::Duration barrier_timeout = deadline - absl::Now();
     // Note: `init_timeout` in client options may be set to 0 so that the
@@ -136,7 +136,7 @@ absl::Status DistributedRuntimeCoordinationServiceClient::Connect() {
 
 absl::Status DistributedRuntimeCoordinationServiceClient::Shutdown() {
   LOG(INFO) << "Distributed task shutdown initiated.";
-  Status s = coord_agent_->Shutdown();
+  absl::Status s = coord_agent_->Shutdown();
   LOG(INFO) << "Distributed task shutdown result: " << s;
   return s;
 }
diff --git a/third_party/xla/xla/pjrt/distributed/client.h b/third_party/xla/xla/pjrt/distributed/client.h
index 4dc341c311014d..03560cd1ca59f7 100644
--- a/third_party/xla/xla/pjrt/distributed/client.h
+++ b/third_party/xla/xla/pjrt/distributed/client.h
@@ -83,7 +83,8 @@ class DistributedRuntimeClient {
                        "missing heartbeats. This most likely indicates that "
                        "another task died; see the other task logs for more "
                        "details. Disable Python buffering, i.e. `python -u`, "
-                       "to be sure to see all the previous output. Status: "
+                       "to be sure to see all the previous output. "
+                       "absl::Status: "
                     << status;
               } else {
                 LOG(QFATAL)
@@ -92,7 +93,7 @@ class DistributedRuntimeClient {
                        "indicates that the coordinator task died; see the "
                        "coordinator's task logs for more details. "
                        "Disable Python buffering, i.e. `python -u`, to be "
-                       "sure to see all the previous output. Status: "
+                       "sure to see all the previous output. absl::Status: "
                     << status;
               }
             };
diff --git a/third_party/xla/xla/pjrt/distributed/client_server_test.cc b/third_party/xla/xla/pjrt/distributed/client_server_test.cc
index 344b2a2239f1a8..f6a4b1fae9828c 100644
--- a/third_party/xla/xla/pjrt/distributed/client_server_test.cc
+++ b/third_party/xla/xla/pjrt/distributed/client_server_test.cc
@@ -612,7 +612,7 @@ TEST_F(ClientServerTest, WaitAtBarrier_Timeout) {
     if (node_id == 1) {
       n.WaitForNotification();
     }
-    Status barrier_status =
+    absl::Status barrier_status =
         client->WaitAtBarrier("barrier_1", kBarrierTimeout, std::nullopt);
     // Node 0 notifies that barrier has already timed out.
     if (node_id == 0) {
diff --git a/third_party/xla/xla/pjrt/distributed/topology_util.cc b/third_party/xla/xla/pjrt/distributed/topology_util.cc
index e8649c68c5bb86..e2fbd874756223 100644
--- a/third_party/xla/xla/pjrt/distributed/topology_util.cc
+++ b/third_party/xla/xla/pjrt/distributed/topology_util.cc
@@ -156,12 +156,13 @@ GlobalTopologyProto BuildGlobalTopology(
   return global_topology;
 }
 
-Status ExchangeTopologies(std::string_view platform, int node_id, int num_nodes,
-                          absl::Duration get_local_topology_timeout,
-                          absl::Duration get_global_topology_timeout,
-                          KeyValueStoreInterface* kv_store,
-                          const LocalTopologyProto& local_topology,
-                          GlobalTopologyProto* global_topology) {
+absl::Status ExchangeTopologies(std::string_view platform, int node_id,
+                                int num_nodes,
+                                absl::Duration get_local_topology_timeout,
+                                absl::Duration get_global_topology_timeout,
+                                KeyValueStoreInterface* kv_store,
+                                const LocalTopologyProto& local_topology,
+                                GlobalTopologyProto* global_topology) {
   VLOG(3) << "Local Topology for platform" << platform << ":\n"
           << local_topology.DebugString();
   if (num_nodes == 1) {
diff --git a/third_party/xla/xla/pjrt/distributed/topology_util.h b/third_party/xla/xla/pjrt/distributed/topology_util.h
index c25788b9015adc..99cc8eab8012e1 100644
--- a/third_party/xla/xla/pjrt/distributed/topology_util.h
+++ b/third_party/xla/xla/pjrt/distributed/topology_util.h
@@ -36,12 +36,13 @@ absl::StatusOr<std::string> GetBootIdString();
 // Performs a distributed exchange of topologies using a KV store. Each process
 // provides its local topology, and the local topologies are exchanged to
 // form a global topology.
-Status ExchangeTopologies(std::string_view platform, int node_id, int num_nodes,
-                          absl::Duration get_local_topology_timeout,
-                          absl::Duration get_global_topology_timeout,
-                          KeyValueStoreInterface* kv_store,
-                          const LocalTopologyProto& local_topology,
-                          GlobalTopologyProto* global_topology);
+absl::Status ExchangeTopologies(std::string_view platform, int node_id,
+                                int num_nodes,
+                                absl::Duration get_local_topology_timeout,
+                                absl::Duration get_global_topology_timeout,
+                                KeyValueStoreInterface* kv_store,
+                                const LocalTopologyProto& local_topology,
+                                GlobalTopologyProto* global_topology);
 
 // Functions below this point are public only for testing.
 
diff --git a/third_party/xla/xla/pjrt/distributed/util.h b/third_party/xla/xla/pjrt/distributed/util.h
index 8e0f290795207b..be07f318ab7246 100644
--- a/third_party/xla/xla/pjrt/distributed/util.h
+++ b/third_party/xla/xla/pjrt/distributed/util.h
@@ -21,16 +21,16 @@ limitations under the License.
 
 namespace xla {
 
-inline Status FromGrpcStatus(const ::grpc::Status& s) {
+inline absl::Status FromGrpcStatus(const ::grpc::Status& s) {
   if (s.ok()) {
     return OkStatus();
   } else {
-    return Status(static_cast<absl::StatusCode>(s.error_code()),
-                  s.error_message());
+    return absl::Status(static_cast<absl::StatusCode>(s.error_code()),
+                        s.error_message());
   }
 }
 
-inline ::grpc::Status ToGrpcStatus(const Status& s) {
+inline ::grpc::Status ToGrpcStatus(const absl::Status& s) {
   if (s.ok()) {
     return ::grpc::Status::OK;
   } else {
diff --git a/third_party/xla/xla/pjrt/gpu/gpu_helpers.cc b/third_party/xla/xla/pjrt/gpu/gpu_helpers.cc
index b9477360de1190..bbbb4be4107e22 100644
--- a/third_party/xla/xla/pjrt/gpu/gpu_helpers.cc
+++ b/third_party/xla/xla/pjrt/gpu/gpu_helpers.cc
@@ -59,7 +59,7 @@ void EnablePeerAccess(absl::Span<se::StreamExecutor* const> executors) {
       se::StreamExecutor* from = executors[i];
       se::StreamExecutor* to = executors[j];
       if (from->CanEnablePeerAccessTo(to)) {
-        Status status = from->EnablePeerAccessTo(to);
+        absl::Status status = from->EnablePeerAccessTo(to);
         if (!status.ok()) {
           LOG(WARNING) << "Unable to enable peer access between GPUs " << i
                        << " and " << j << "; status: " << status;
@@ -76,8 +76,8 @@ absl::StatusOr<std::unique_ptr<tsl::BFCAllocator>> CreateBFCAllocator(
     se::StreamExecutor* executor, double memory_fraction, bool preallocate,
     std::optional<int64_t> gpu_system_memory_size) {
   bool enable_unified_memory;
-  Status status = tsl::ReadBoolFromEnvVar("TF_FORCE_UNIFIED_MEMORY", false,
-                                          &enable_unified_memory);
+  absl::Status status = tsl::ReadBoolFromEnvVar("TF_FORCE_UNIFIED_MEMORY",
+                                                false, &enable_unified_memory);
   if (!status.ok()) {
     LOG(ERROR) << "Unable to read TF_FORCE_UNIFIED_MEMORY: "
                << status.message();
@@ -173,7 +173,7 @@ std::unique_ptr<tsl::BFCAllocator> GetGpuHostAllocator(
                                   /*free_visitors=*/{}));
 
   int64_t xla_pjrt_gpu_host_memory_limit_gb;
-  Status status =
+  absl::Status status =
       tsl::ReadInt64FromEnvVar("XLA_PJRT_GPU_HOST_MEMORY_LIMIT_GB", 64,
                                &xla_pjrt_gpu_host_memory_limit_gb);
   if (!status.ok()) {
diff --git a/third_party/xla/xla/pjrt/gpu/se_gpu_pjrt_client.cc b/third_party/xla/xla/pjrt/gpu/se_gpu_pjrt_client.cc
index bfa11cb315c8c5..a1be7fe67058db 100644
--- a/third_party/xla/xla/pjrt/gpu/se_gpu_pjrt_client.cc
+++ b/third_party/xla/xla/pjrt/gpu/se_gpu_pjrt_client.cc
@@ -203,7 +203,7 @@ class AsyncHostToDeviceTransferManager
     return std::move(buffers_[buffer_index]);
   };
 
-  Status TransferLiteralToBuffer(
+  absl::Status TransferLiteralToBuffer(
       int buffer_index, const LiteralSlice& literal,
       absl::AnyInvocable<void() &&> on_done) override {
     tsl::profiler::TraceMe traceme(
@@ -295,7 +295,7 @@ class AsyncHostToDeviceTransferManager
     return OkStatus();
   }
 
-  Status TransferRawDataToBuffer(
+  absl::Status TransferRawDataToBuffer(
       int buffer_index, absl::string_view data,
       absl::AnyInvocable<void() &&> on_done) override {
     return TransferRawDataToSubBuffer(buffer_index, data.data(),
@@ -304,7 +304,7 @@ class AsyncHostToDeviceTransferManager
                                       std::move(on_done));
   }
 
-  Status TransferRawDataToSubBuffer(
+  absl::Status TransferRawDataToSubBuffer(
       int buffer_index, const void* data, int64_t offset, int64_t transfer_size,
       bool is_last_transfer, absl::AnyInvocable<void() &&> on_done) override {
     auto* stream = device_->local_device_state()->host_to_device_stream();
@@ -401,7 +401,7 @@ class AsyncHostToDeviceTransferManager
     return stream->DoHostCallback(std::move(cleanup));
   }
 
-  void SetBufferError(int buffer_index, Status error) override {
+  void SetBufferError(int buffer_index, absl::Status error) override {
     {
       absl::MutexLock l(&mu_);
       // For a given buffer_index, SetBufferError can't be called twice, or
@@ -894,7 +894,7 @@ GetStreamExecutorGpuDeviceAllocator(
 
 }  // namespace
 
-Status BuildDistributedDevices(
+absl::Status BuildDistributedDevices(
     std::string_view platform_name,
     std::map<int, std::unique_ptr<LocalDeviceState>> local_device_states,
     int node_id, int num_nodes,
diff --git a/third_party/xla/xla/pjrt/gpu/se_gpu_pjrt_client.h b/third_party/xla/xla/pjrt/gpu/se_gpu_pjrt_client.h
index 2101501050a670..79af364f3bea33 100644
--- a/third_party/xla/xla/pjrt/gpu/se_gpu_pjrt_client.h
+++ b/third_party/xla/xla/pjrt/gpu/se_gpu_pjrt_client.h
@@ -152,7 +152,7 @@ class StreamExecutorGpuTopologyDescription : public PjRtTopologyDescription {
     return attributes_;
   }
 
-  StatusOr<Layout> GetDefaultLayout(
+  absl::StatusOr<Layout> GetDefaultLayout(
       PrimitiveType element_type,
       absl::Span<const int64_t> dims) const override;
 
diff --git a/third_party/xla/xla/pjrt/gpu/se_gpu_pjrt_client_test.cc b/third_party/xla/xla/pjrt/gpu/se_gpu_pjrt_client_test.cc
index 616ab665e34604..42a5170dcc2541 100644
--- a/third_party/xla/xla/pjrt/gpu/se_gpu_pjrt_client_test.cc
+++ b/third_party/xla/xla/pjrt/gpu/se_gpu_pjrt_client_test.cc
@@ -374,7 +374,7 @@ TEST(StreamExecutorGpuClientTest, ToLiteralAsync) {
   TF_ASSERT_OK(
       transfer_manager->TransferLiteralToBuffer(0, src_literal, [&]() {}));
 
-  buffer->ToLiteral(literal.get()).OnReady([&](Status s) {
+  buffer->ToLiteral(literal.get()).OnReady([&](absl::Status s) {
     absl::MutexLock l(&mu);
     TF_ASSERT_OK(s);
     got_literal = true;
@@ -408,7 +408,7 @@ TEST(StreamExecutorGpuClientTest, ToLiteralAsyncBeforeBufferReady) {
       ShapeUtil::DeviceShapeToHostShape(buffer->on_device_shape()));
   bool got_literal = false;
 
-  buffer->ToLiteral(literal.get()).OnReady([&](Status s) {
+  buffer->ToLiteral(literal.get()).OnReady([&](absl::Status s) {
     absl::MutexLock l(&mu);
     TF_ASSERT_OK(s);
     got_literal = true;
@@ -468,12 +468,12 @@ TEST(StreamExecutorGpuClientTest, FromHostAsync) {
   for (auto& buffer : buffers) {
     literals.push_back(std::make_shared<Literal>(
         ShapeUtil::DeviceShapeToHostShape(buffer->on_device_shape())));
-    buffer->ToLiteral(literals.back().get()).OnReady([&](Status s) {
+    buffer->ToLiteral(literals.back().get()).OnReady([&](absl::Status s) {
       absl::MutexLock l(&mu);
       TF_ASSERT_OK(s);
       ++got_literal_count;
     });
-    buffer->GetReadyFuture().OnReady([&](Status s) {
+    buffer->GetReadyFuture().OnReady([&](absl::Status s) {
       absl::MutexLock l(&mu);
       TF_ASSERT_OK(s);
       ++got_callback_count;
diff --git a/third_party/xla/xla/service/gpu/runtime/annotation.cc b/third_party/xla/xla/service/gpu/runtime/annotation.cc
index 809f02d30d10b3..f0a9600ccd882d 100644
--- a/third_party/xla/xla/service/gpu/runtime/annotation.cc
+++ b/third_party/xla/xla/service/gpu/runtime/annotation.cc
@@ -154,7 +154,7 @@ class SourceLocationVisitor : public ConstDfsHloVisitorWithDefault {
     return std::move(oss).str();
   }
 
-  Status DefaultAction(HloInstruction const* inst) final {
+  absl::Status DefaultAction(HloInstruction const* inst) final {
     OpMetadata const& meta = inst->metadata();
     // The full op_name is split across three places: the module-level
     // annotation shows the prefix that is common to the whole module, the
diff --git a/third_party/xla/xla/service/gpu/runtime/nccl_collective_broadcast_thunk.cc b/third_party/xla/xla/service/gpu/runtime/nccl_collective_broadcast_thunk.cc
index f11b8c850417d6..76d316feededf3 100644
--- a/third_party/xla/xla/service/gpu/runtime/nccl_collective_broadcast_thunk.cc
+++ b/third_party/xla/xla/service/gpu/runtime/nccl_collective_broadcast_thunk.cc
@@ -43,7 +43,7 @@ NcclCollectiveBroadcastStartThunk::NcclCollectiveBroadcastStartThunk(
       config_(GetNcclCollectiveConfig(instr, std::nullopt)),
       buffers_(std::move(buffers)) {}
 
-/*static*/ Status NcclCollectiveBroadcastStartThunk::CheckImplementable(
+/*static*/ absl::Status NcclCollectiveBroadcastStartThunk::CheckImplementable(
     const HloInstruction* instr, int64_t replica_count,
     int64_t partition_count) {
   return OkStatus();
@@ -55,7 +55,7 @@ NcclCollectiveBroadcastStartThunk::GetGroupMode(
   return GetNcclCollectiveConfig(inst, std::nullopt).group_mode;
 }
 
-Status NcclCollectiveBroadcastStartThunk::RunNcclCollective(
+absl::Status NcclCollectiveBroadcastStartThunk::RunNcclCollective(
     const ExecuteParams& params, se::Stream& stream,
     NcclCommHandleWrapper comm_wrapper) {
   TF_ASSIGN_OR_RETURN(
@@ -65,9 +65,10 @@ Status NcclCollectiveBroadcastStartThunk::RunNcclCollective(
       device_buffers, stream, comm_wrapper.comm_handle, nccl_api());
 }
 
-Status RunCollectiveBroadcast(std::vector<DeviceBufferPair>& buffers,
-                              se::Stream& stream, NcclApi::NcclCommHandle comm,
-                              NcclApi* nccl_api) {
+absl::Status RunCollectiveBroadcast(std::vector<DeviceBufferPair>& buffers,
+                                    se::Stream& stream,
+                                    NcclApi::NcclCommHandle comm,
+                                    NcclApi* nccl_api) {
   TF_RETURN_IF_ERROR(nccl_api->GroupStart());
   for (auto buffer : buffers) {
     se::DeviceMemoryBase src_addr = buffer.source_buffer;
diff --git a/third_party/xla/xla/service/gpu/runtime/nccl_collective_broadcast_thunk.h b/third_party/xla/xla/service/gpu/runtime/nccl_collective_broadcast_thunk.h
index 1bf7a021bd46b7..8c0cc19215557d 100644
--- a/third_party/xla/xla/service/gpu/runtime/nccl_collective_broadcast_thunk.h
+++ b/third_party/xla/xla/service/gpu/runtime/nccl_collective_broadcast_thunk.h
@@ -32,9 +32,9 @@ namespace xla::gpu {
 // Thunk that performs a NCCL-based collective broadcast.
 class NcclCollectiveBroadcastStartThunk : public NcclCollectiveThunk {
  public:
-  static Status CheckImplementable(const HloInstruction* instr,
-                                   int64_t replica_count,
-                                   int64_t partition_count);
+  static absl::Status CheckImplementable(const HloInstruction* instr,
+                                         int64_t replica_count,
+                                         int64_t partition_count);
 
   static CollectiveOpGroupMode GetGroupMode(
       const HloCollectiveBroadcastInstruction* inst);
@@ -50,17 +50,19 @@ class NcclCollectiveBroadcastStartThunk : public NcclCollectiveThunk {
       std::vector<Buffer> buffers);
 
  protected:
-  Status RunNcclCollective(const ExecuteParams& params, se::Stream& stream,
-                           NcclCommHandleWrapper comm_wrapper) override;
+  absl::Status RunNcclCollective(const ExecuteParams& params,
+                                 se::Stream& stream,
+                                 NcclCommHandleWrapper comm_wrapper) override;
 
  private:
   const NcclCollectiveConfig config_;
   const std::vector<Buffer> buffers_;
 };
 
-Status RunCollectiveBroadcast(std::vector<DeviceBufferPair>& buffers,
-                              se::Stream& stream, NcclApi::NcclCommHandle comm,
-                              NcclApi* nccl_api);
+absl::Status RunCollectiveBroadcast(std::vector<DeviceBufferPair>& buffers,
+                                    se::Stream& stream,
+                                    NcclApi::NcclCommHandle comm,
+                                    NcclApi* nccl_api);
 
 }  // namespace xla::gpu
 
diff --git a/third_party/xla/xla/service/gpu/runtime/nccl_collective_thunk.cc b/third_party/xla/xla/service/gpu/runtime/nccl_collective_thunk.cc
index f4c1ebf3df99dd..10a7ef41fd3fcf 100644
--- a/third_party/xla/xla/service/gpu/runtime/nccl_collective_thunk.cc
+++ b/third_party/xla/xla/service/gpu/runtime/nccl_collective_thunk.cc
@@ -290,9 +290,9 @@ absl::StatusOr<std::vector<DeviceBufferPair>> ConvertToDeviceBuffers(
   return device_buffers;
 }
 
-Status RegisterBufferOnce(NcclApi* nccl_api, int device_ordinal,
-                          NcclApi::NcclCommHandle comm,
-                          se::DeviceMemoryBase buffer) {
+absl::Status RegisterBufferOnce(NcclApi* nccl_api, int device_ordinal,
+                                NcclApi::NcclCommHandle comm,
+                                se::DeviceMemoryBase buffer) {
   // Keep track of which communicators we have registered for already.
   // Each ncclMemAlloc'd buffer needs to be registered once per comm.
   struct RegisteredBuffers {
@@ -332,9 +332,9 @@ Status RegisterBufferOnce(NcclApi* nccl_api, int device_ordinal,
   return OkStatus();
 }
 
-Status MaybeRegisterBuffers(NcclApi* nccl_api, int device_ordinal,
-                            const std::vector<DeviceBufferPair>& buffers,
-                            NcclApi::NcclCommHandle comm) {
+absl::Status MaybeRegisterBuffers(NcclApi* nccl_api, int device_ordinal,
+                                  const std::vector<DeviceBufferPair>& buffers,
+                                  NcclApi::NcclCommHandle comm) {
   for (int i = 0; i < buffers.size(); ++i) {
     if (buffers[i].source_memory_space == kCollectiveMemorySpaceColor) {
       TF_RETURN_IF_ERROR(RegisterBufferOnce(nccl_api, device_ordinal, comm,
@@ -428,7 +428,7 @@ bool operator==(const FirstCallRendezvousKey& a,
 }
 }  // namespace
 
-Status NcclCollectiveThunk::ExecuteOnStream(const ExecuteParams& params) {
+absl::Status NcclCollectiveThunk::ExecuteOnStream(const ExecuteParams& params) {
   VLOG(1) << absl::StreamFormat("Starting %s %s.", IsAsync() ? "async" : "sync",
                                 Thunk::KindToString(kind()));
   const NcclStreamId stream_id = nccl_stream_id();
diff --git a/third_party/xla/xla/service/gpu/runtime/nccl_collective_thunk.h b/third_party/xla/xla/service/gpu/runtime/nccl_collective_thunk.h
index 977f791b87328f..d428966e784d14 100644
--- a/third_party/xla/xla/service/gpu/runtime/nccl_collective_thunk.h
+++ b/third_party/xla/xla/service/gpu/runtime/nccl_collective_thunk.h
@@ -271,7 +271,7 @@ absl::Status AddOpDescription(absl::Status status, OpT op,
     str = llvm_ir::DumpToString(op.getOperation());
   }
 
-  return Status(
+  return absl::Status(
       status.code(),
       absl::StrFormat(
           "%s\n"
@@ -320,9 +320,9 @@ absl::StatusOr<std::vector<DeviceBufferPair>> ConvertToDeviceBuffers(
 // communicator to enable zero-copy collectives.
 //
 // https://docs.nvidia.com/deeplearning/nccl/user-guide/docs/usage/bufferreg.html
-Status MaybeRegisterBuffers(NcclApi* nccl_api, int device_ordinal,
-                            const std::vector<DeviceBufferPair>& buffers,
-                            NcclApi::NcclCommHandle comm);
+absl::Status MaybeRegisterBuffers(NcclApi* nccl_api, int device_ordinal,
+                                  const std::vector<DeviceBufferPair>& buffers,
+                                  NcclApi::NcclCommHandle comm);
 
 }  // namespace gpu
 }  // namespace xla
diff --git a/third_party/xla/xla/service/gpu/runtime/nccl_p2p_thunk_common.cc b/third_party/xla/xla/service/gpu/runtime/nccl_p2p_thunk_common.cc
index f4ef9a77205255..6d0ae9584fc51a 100644
--- a/third_party/xla/xla/service/gpu/runtime/nccl_p2p_thunk_common.cc
+++ b/third_party/xla/xla/service/gpu/runtime/nccl_p2p_thunk_common.cc
@@ -111,8 +111,8 @@ NcclP2PConfig GetNcclP2PConfigForSendRecv(const HloSendRecvInstruction* instr,
       instr->frontend_attributes().map().at(kSendRecvSourceTargetPairsAttr);
 
   // We currently ignore problems related to the source-target-pair string to
-  // avoid using StatusOr for the return type. This should be ok as Send/Recv
-  // are generated by the compiler.
+  // avoid using absl::StatusOr for the return type. This should be ok as
+  // Send/Recv are generated by the compiler.
   if (!source_target_pairs_string.has_value()) {
     return p2p_config;
   }
@@ -135,7 +135,7 @@ NcclP2PConfig GetNcclP2PConfigForSendRecv(const HloSendRecvInstruction* instr,
       if (!statusor_bounds.ok() ||
           statusor_bounds.value().size() != replica_groups.size()) {
         // Ignore problems related to the source-target-pair string to avoid
-        // using StatusOr for the return type.
+        // using absl::StatusOr for the return type.
         return p2p_config;
       }
       validation_kind = NcclP2PConfig::ValidationKind::kConditional;

From caf11ae374460ba7e4477cacaba03d7f35af7881 Mon Sep 17 00:00:00 2001
From: Eugene Zhulenev <ezhulenev@google.com>
Date: Thu, 16 May 2024 13:56:06 -0700
Subject: [PATCH 300/478] [xla] Update XLA custom call documentation to use XLA
 FFI

Now when XLA FFI is supported on both XLA:GPU and XLA:CPU backends update custom call documentation to new XLA FFI APIs instead of a mix of legacy ones.

PiperOrigin-RevId: 634519378
---
 third_party/xla/docs/custom_call.md | 258 +++++++++-------------------
 1 file changed, 79 insertions(+), 179 deletions(-)

diff --git a/third_party/xla/docs/custom_call.md b/third_party/xla/docs/custom_call.md
index 84633d697daa20..7f4bf765b88b67 100644
--- a/third_party/xla/docs/custom_call.md
+++ b/third_party/xla/docs/custom_call.md
@@ -1,18 +1,18 @@
 # XLA custom calls
 
-This document describes how to write and use XLA custom calls. Custom calls let
-you invoke code written in a programming language like C++ or CUDA from an XLA
-program.
-
-> **Caution:** Custom calls are a low-level power-user feature. It is easy to
-> break your program in difficult-to-debug (and even difficult-to-notice) ways
-> using custom calls. You shouldn't use custom calls unless you're prepared to
-> debug XLA yourself when something goes wrong, and you should expect relatively
-> less assistance from XLA developers if you run into trouble.
-
-> **Caution:** The custom-call API/ABI is not currently stable. We don't intend
-> to change it capriciously, but it may change. Some possible future changes are
-> described below.
+This document describes how to write and use XLA custom calls using XLA FFI
+library. Custom call is a mechanism to describe an external "operation" in the
+HLO module to the XLA compiler (at compile time), and XLA FFI is a mechanism to
+register implementation of such operations with XLA (at run time). FFI stands
+for "foreign function interface" and it is a set of C APIs that define a binary
+interface (ABI) for XLA to call into external code written in other programming
+languages. XLA provides header-only bindings for XLA FFI written in C++, which
+hides all the low level details of underlying C APIs from the end user.
+
+> **Caution:** The custom-call API/ABI uses PJRT-style versioning (major, minor),
+> however at this point it is still experimental and can be broken at any time.
+> Once API/ABI is finalized we intend to provide stability guarantees
+> similar to PJRT.
 
 > **Caution** The HLO-visible names of functions registered with the custom-call
 > macros API do not respect C++ namespaces. As a result, accidental collisions
@@ -42,30 +42,54 @@ void do_it() {
       xla::Parameter(&b, 1, xla::ShapeUtil::MakeShape(xla::F32, {2048}), "p1");
   xla::XlaOp custom_call =
       xla::CustomCall(&b, "do_custom_call", /*operands=*/{param0, param1},
-                      /*shape=*/xla::ShapeUtil::MakeShape(xla::F32, {2048}));
+        /*shape=*/xla::ShapeUtil::MakeShape(xla::F32, {2048}),
+        /*opaque=*/"", /*has_side_effect=*/false,
+        /*output_operand_aliasing=*/{}, /*literal=*/nullptr,
+        /*schedule=*/CustomCallSchedule::SCHEDULE_NONE,
+        /*api_version=*/CustomCallApiVersion::API_VERSION_TYPED_FFI);
 }
 
-void do_custom_call(void* out, const void** in) {
-  float* out_buf = reinterpret_cast<float*>(out);
-  const float* in0 = reinterpret_cast<const float*>(in[0]);
-  const float* in1 = reinterpret_cast<const float*>(in[1]);
-  for (int i = 0; i < 2048; ++i) {
-    out_buf[i] = in0[i % 128] + in1[i];
+// Constrain custom call arguments to rank-1 buffers of F32 data type.
+using BufferF32 = xla::ffi::BufferR1<xla::ffi::DataType::F32>;
+
+// Implement a custom call as a C+ function. Note that we can use `Buffer` type
+// defined by XLA FFI that gives us access to buffer data type and shape.
+xla::ffi::Error do_custom_call(BufferF32 in0, BufferF32 in1,
+                               xla::ffi::Result<BufferF32> out) {
+  size_t d0 = in0.dimensions[0];
+  size_t d1 = in1.dimensions[0];
+
+  // Check that dimensions are compatible.
+  assert(out->dimensions[0] == d1 && "unexpected dimensions");
+
+  for (size_t i = 0; i < d1; ++i) {
+    out->data[i] = in0.data[i % d0] + in1.data[i];
   }
 }
-XLA_REGISTER_CUSTOM_CALL_TARGET(do_custom_call, "Host");
-```
 
-Notice that the function `do_custom_call` needs to know the dimensions of the
-buffers it operates over. In this example we hardcode the sizes `128` and
-`2048`. If you don't want to do this, you can pass the dimensions in as
-parameters to the call.
+// Explicitly define an XLA FFI handler signature and bind it to the
+// `do_custom_call` implementation. XLA FFI handler can automatically infer
+// type signature from the custom call function, but it relies on magical
+// template metaprogramming an explicit binding provides and extra level of
+// type checking and clearly states custom call author intentions.
+XLA_FFI_DEFINE_HANDLER(handler, do_custom_call,
+                       ffi::Ffi::Bind()
+                           .Arg<Buffer>()
+                           .Arg<Buffer>()
+                           .Ret<Buffer>();
+
+// Registers `handler` with and XLA FFI on a "Host" platform.
+XLA_FFI_REGISTER_HANDLER(xla::ffi::GetXlaFfiApi(), "do_custom_call",
+                         "Host", handler);
+```
 
 ## Create a custom call on GPU
 
-The GPU custom call framework is somewhat different than that on the CPU. Here
-is a CUDA example that does the same computation (`A[i] = B[i % 128] + C[i]`) as
-the CPU code above.
+The GPU custom call registration with XLA FFI is almost identical, the only
+difference is that for GPU you need to ask for an underlying platform stream
+(CUDA or ROCM stream) to be able to launch kernel on device. Here is a CUDA
+example that does the same computation (`A[i] = B[i % 128] + C[i]`) as the CPU
+code above.
 
 ```c++
 void do_it() { /* same implementation as above */ }
@@ -75,18 +99,29 @@ __global__ custom_call_kernel(const float* in0, const float* in1, float* out) {
   out[idx] = in0[idx % 128] + in1[idx];
 }
 
-void do_custom_call(CUstream stream, void** buffers,
-                    const char* opaque, size_t opaque_len) {
-  const float* in0 = reinterpret_cast<const float*>(buffers[0]);
-  const float* in1 = reinterpret_cast<const float*>(buffers[1]);
-  float* out = reinterpret_cast<float*>(buffers[2]);
+void do_custom_call(CUstream stream, BufferF32 in0, BufferF32 in1,
+                    xla::ffi::Result<BufferF32> out) {
+  size_t d0 = in0.dimensions[0];
+  size_t d1 = in1.dimensions[0];
+  size_t d2 = out->dimensions[0];
+
+  assert(d0 == 128 && d1 == 2048 && d2 == 2048 && "unexpected dimensions");
 
   const int64_t block_dim = 64;
   const int64_t grid_dim = 2048 / block_dim;
-  custom_call_kernel<<<grid_dim, block_dim,
-                       /*dynamic_shared_mem_bytes=*/0, stream>>>(in0, in1, out);
+  custom_call_kernel<<<grid_dim, block_dim, 0, stream>>>(
+    in0.data, in1.data, out.data);
 }
-XLA_REGISTER_CUSTOM_CALL_TARGET(do_custom_call, "CUDA");
+
+XLA_FFI_DEFINE_HANDLER(handler, do_custom_call,
+                       ffi::Ffi::Bind()
+                           .Ctx<xla::ffi::PlatformStream<CUstream>>()
+                           .Arg<Buffer>()
+                           .Arg<Buffer>()
+                           .Ret<Buffer>();
+
+XLA_FFI_REGISTER_HANDLER(xla::ffi::GetXlaFfiApi(), "do_custom_call",
+                         "CUDA", handler);
 ```
 
 Notice first that the GPU custom call function *is still a function executed on
@@ -94,99 +129,10 @@ the CPU*. The `do_custom_call` CPU function is responsible for enqueueing work
 on the GPU. Here it launches a CUDA kernel, but it could also do something else,
 like call cuBLAS.
 
-`buffers` is an array of pointers that lives on the host, and each element it
-contains points to device (i.e. GPU) memory. The parameters come first, followed
-by the output value. This is notably different from the CPU calling convention,
-which has two params, `ins` and `out`. The GPU calling convention makes it
-possible to handle tuple-shaped inputs/outputs efficiently.
-
-As in the CPU example, we've hardcoded the input and output buffer sizes into
-our custom call. However unlike in the CPU case, passing the buffer sizes in as
-operands to the custom call would not work well. Usually we need the buffer
-sizes available to us on the CPU (e.g. when launching a kernel, we need to know
-the block/grid dimensions to use). But if we were to pass the buffer sizes as
-operands to our custom call, their values would live in GPU memory. We'd then
-have to do an expensive synchronous device-to-host `memcpy` at the start of our
-operation just to read the sizes.
-
-To let you work around this, we provide the `opaque` parameter. You can set this
-to an arbitrary string of bytes when you create the custom call:
-
-```c++
-std::string opaque = "...";
-xla::CustomCall(&b, "do_custom_call", /*operands=*/{param0, param1},
-                /*output_shape=*/xla::ShapeUtil::MakeShape(xla::F32, {2048}),
-                opaque);
-```
-
-Because `xla::Shape` has a protocol buffer representation, you could store this
-serialized proto inside of `opaque` and deserialize it within your GPU custom
-call. Note however that although `xla::ShapeProto` does not change frequently,
-it *does* change. Check the Git log to see how it has changed in the past.
-
-## Signalling an error
-
-If your custom call encounters an error, you can signal the error to the XLA
-runtime (instead of e.g. crashing or returning nonsense in the output buffers)
-by using the following signature for your function:
-
-**On CPU:**
-
-```c++
-#include "xla/service/custom_call_status.h"
-
-void do_custom_call(void* out, const void** in, XlaCustomCallStatus* status);
-```
-
-**on GPU:**
-
-```c++
-#include "xla/service/custom_call_status.h"
-
-void do_custom_call(CUstream stream, void** buffers, const char* opaque,
-                    size_t opaque_len, xla::XlaCustomCallStatus* status);
-```
-
-You can signal failure by using `XlaCustomCallStatusSetFailure`, e.g.:
-
-```c++
-void do_custom_call(void* out, const void** in, XlaCustomCallStatus* status) {
-  // ... do some work.
-
-  if (bad_condition) {
-    char* error_message = "An error occurred";
-    XlaCustomCallStatusSetFailure(status, error_message, strlen(error_message));
-    return;
-  }
-
-  // ... continue.
-}
-```
-
-You can also use `XlaCustomCallStatusSetSuccess` to indicate success, but the
-`XlaCustomCallStatus` is in a success state by default, so ignoring it
-completely will also indicate success.
-
-When using custom call functions with this signature, you must create the
-corresponding `custom-call` op with the appropriate API version set, e.g.:
-
-```c++
-xla::CustomCall(&b, "do_custom_call", /*operands=*/{param0, param1},
-                /*output_shape=*/xla::ShapeUtil::MakeShape(F32, {2048}),
-                opaque, /*has_side_effect=*/false,
-                /*output_operand_aliasing=*/{}, /*literal=*/nullptr,
-                /*schedule=*/xla::CustomCallSchedule::SCHEDULE_NONE,
-                /*api_version=*/API_VERSION_STATUS_RETURNING);
-```
-
-> **Note:** In the future all clients will be required to migrate their custom
-> call functions to the new API version and the old one will be deprecated. For
-> custom calls that can't fail, you can simply add the new
-> `XlaCustomCallStatus*` parameter and then ignore it.
-
-On failure, none of the custom call outputs will be used; the XLA runtime will
-terminate the computation. It is not possible for an HLO computation to recover
-from the error (e.g. by catching and handling it).
+Arguments and results also live on the host, and data member contains a pointer
+to device (i.e. GPU) memory. Buffers passed to custom call handler have the
+shape of the underlying device buffers, so the custom call can compute kernel
+launch parameters from them.
 
 ## Passing tuples to custom calls
 
@@ -209,32 +155,12 @@ Shape out_shape = ShapeUtil::MakeTuple({
   ShapeUtil::MakeShape(F32, {512}),
   ShapeUtil::MakeShape(F32, {1024}),
 });
-xla::CustomCall(&b, "do_custom_call", /*operands=*/{p0}, out_shape);
+xla::CustomCall(&b, "do_custom_call", /*operands=*/{p0}, out_shape, ...);
 ```
 
 On both CPU and GPU, a tuple is represented in memory as an array of pointers.
-In C++ pseudocode, parameter 0 above is laid out as follows.
-
-```c++
-// In-memory layout of parameter 0 from custom call above. True on both CPU
-// and GPU.
-float* subbuf0 = new float[32];
-float* subbuf1 = new float[64];
-float* subbuf2 = new float[128]
-float* subbuf3 = new float[256];
-
-void* subtuple = new void*[2];
-(*subtuple)[0] = subbuf1;
-(*subtuple)[1] = subbuf2;
-
-void* p0 = new void*[3];
-(*p0)[0] = subbuf0;
-(*p0)[1] = subtuple;
-(*p0)[2] = subbuf3;
-```
-
-Although the in-memory representation of tuples is the same in CPU and GPU, they
-are handled differently in the CPU and GPU custom-call calling conventions.
+When XLA calls custom calls with tuple arguments or results it flattens them and
+passes as regular buffer arguments or results.
 
 ### Tuple outputs as temp buffers
 
@@ -256,29 +182,3 @@ been written to. That's exactly what you want from a temp buffer.
 In the example above, suppose we wanted to use the `F32[1024]` as a temp buffer.
 Then we'd write the HLO just as above, and we'd simply never read tuple index 1
 of the custom call's output.
-
-### Tuples in CPU custom calls
-
-In CPU code, we have a function `do_custom_call(const void** ins, void* out)`.
-`ins` is an array with just one element, which points to `param0`. The
-subbuffers of `param0` are accessible by dereferencing that pointer, and the
-subbuffers of `output_tuple` are accessible by dereferencing `out`.
-
-### Tuples in GPU custom calls
-
-In GPU code, we have a function `do_custom_call(..., void** buffers, ...)`. In
-this case `buffers` is a host array of *six* device pointers, one for each leaf
-buffer in the input/output. To generate the flat list, we iterate over the
-parameters and output, and for each we do a preorder traversal of its shape.
-Concretely:
-
-```c++
-// Layout of `buffers` parameter to GPU custom call function for custom-call
-// above.
-buffers[0] == subbuf0
-buffers[1] == subbuf1
-buffers[2] == subbuf2
-buffers[3] == subbuf3
-buffers[4] == output_subbuf0
-buffers[5] == output_subbuf1
-```

From 3eef59e29b6385902a8a2eaaaed97ce57d277c98 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Thu, 16 May 2024 14:00:48 -0700
Subject: [PATCH 301/478] Introduce ExecutionStreamAssignment.

`ExecutionStreamAssignments` represent a mapping from `HloInstructions` to `ExecutionStreamIds`. Asynchronous calls (`async-start`, `async-update`, and `async-done`) result in the target computations being assigned new `ExecutionStreamIds` to support concurrent execution.

PiperOrigin-RevId: 634520855
---
 third_party/xla/xla/service/gpu/BUILD         |  33 ++++
 .../gpu/execution_stream_assignment.cc        | 147 +++++++++++++++
 .../service/gpu/execution_stream_assignment.h |  80 +++++++++
 .../gpu/execution_stream_assignment_test.cc   | 169 ++++++++++++++++++
 4 files changed, 429 insertions(+)
 create mode 100644 third_party/xla/xla/service/gpu/execution_stream_assignment.cc
 create mode 100644 third_party/xla/xla/service/gpu/execution_stream_assignment.h
 create mode 100644 third_party/xla/xla/service/gpu/execution_stream_assignment_test.cc

diff --git a/third_party/xla/xla/service/gpu/BUILD b/third_party/xla/xla/service/gpu/BUILD
index 5f10710157ddf1..2af7bf63b6dfe7 100644
--- a/third_party/xla/xla/service/gpu/BUILD
+++ b/third_party/xla/xla/service/gpu/BUILD
@@ -6170,3 +6170,36 @@ xla_cc_test(
         "@local_tsl//tsl/platform:test_main",
     ],
 )
+
+cc_library(
+    name = "execution_stream_assignment",
+    srcs = ["execution_stream_assignment.cc"],
+    hdrs = ["execution_stream_assignment.h"],
+    deps = [
+        "//xla/hlo/ir:hlo",
+        "//xla/service:call_graph",
+        "//xla/service/gpu/runtime:thunk",
+        "@com_google_absl//absl/container:flat_hash_map",
+        "@com_google_absl//absl/log:check",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/status:statusor",
+        "@com_google_absl//absl/strings",
+    ],
+)
+
+xla_cc_test(
+    name = "execution_stream_assignment_test",
+    srcs = ["execution_stream_assignment_test.cc"],
+    deps = [
+        ":execution_stream_assignment",
+        "//xla:shape_util",
+        "//xla/hlo/ir:hlo",
+        "//xla/service/gpu/runtime:thunk",
+        "//xla/tests:hlo_test_base",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/strings:string_view",
+        "@com_google_googletest//:gtest_main",
+        "@local_tsl//tsl/platform:status_matchers",
+        "@local_tsl//tsl/platform:statusor",
+    ],
+)
diff --git a/third_party/xla/xla/service/gpu/execution_stream_assignment.cc b/third_party/xla/xla/service/gpu/execution_stream_assignment.cc
new file mode 100644
index 00000000000000..a41a680563ae6c
--- /dev/null
+++ b/third_party/xla/xla/service/gpu/execution_stream_assignment.cc
@@ -0,0 +1,147 @@
+/* Copyright 2024 The OpenXLA Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+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.
+==============================================================================*/
+
+#include "xla/service/gpu/execution_stream_assignment.h"
+
+#include <deque>
+#include <memory>
+#include <utility>
+
+#include "absl/container/flat_hash_map.h"
+#include "absl/log/check.h"
+#include "absl/status/status.h"
+#include "absl/strings/str_cat.h"
+#include "xla/hlo/ir/hlo_casting_utils.h"
+#include "xla/hlo/ir/hlo_computation.h"
+#include "xla/hlo/ir/hlo_instruction.h"
+#include "xla/hlo/ir/hlo_instructions.h"
+#include "xla/hlo/ir/hlo_module.h"
+#include "xla/hlo/ir/hlo_opcode.h"
+#include "xla/service/call_graph.h"
+#include "xla/service/gpu/runtime/thunk.h"
+
+namespace xla::gpu {
+
+ExecutionStreamAssignment::ExecutionStreamAssignment(const HloModule* module) {
+  std::unique_ptr<CallGraph> call_graph = CallGraph::Build(module);
+
+  // We'll walk the `CallGraph` starting from the entrypoint. The instructions
+  // on the entrypoint computation will be assigned `ExecutionStreamId(0)`, and
+  // each invocation of `async-start` will result in the target computation
+  // being assigned a new `ExecutionStreamId`.
+  ExecutionStreamId next_stream_id = ExecutionStreamId(1);
+
+  // Each `Pending` item represents an `HloComputation` that needs to be
+  // processed. We start with the entrypoint and add callees as we discover
+  // them.
+  struct Pending {
+    Pending(HloComputation* node, ExecutionStreamId stream_id)
+        : node(node), stream_id(stream_id) {}
+    HloComputation* node;
+    ExecutionStreamId stream_id;
+  };
+  std::deque<Pending> queue;
+  queue.emplace_back(module->entry_computation(), ExecutionStreamId(0));
+
+  while (!queue.empty()) {
+    Pending pending = queue.front();
+    queue.pop_front();
+
+    // First, we'll assign the current `ExecutionStreamId` to all synchronous
+    // instructions. Asynchronous instructions will be handled afterwards.
+    for (HloInstruction* instruction : pending.node->instructions()) {
+      if (instruction->IsAsynchronous()) continue;
+      CHECK(sync_instructions_.try_emplace(instruction, pending.stream_id)
+                .second);
+    }
+
+    // Next, we'll process all callsites in the current computation.
+    for (const CallSite& callsite :
+         call_graph->GetNode(pending.node).callsites()) {
+      if (callsite.instruction()->IsAsynchronous()) {
+        // Asynchronous calls will result in a new `ExecutionStreamId` being
+        // dispensed for the target computation.
+        CHECK_EQ(callsite.instruction()->opcode(), HloOpcode::kAsyncStart);
+        const ExecutionStreamId async_stream_id = next_stream_id++;
+
+        // Because the `HloModule` is assumed to be flat, all computations must
+        // be invoked from a single call-like instruction.
+        CHECK_EQ(callsite.called_computations().size(), 1);
+        queue.emplace_back(callsite.called_computations()[0], async_stream_id);
+
+        AsyncExecutionStreamIds streams;
+        streams.source_stream_id = pending.stream_id;
+        streams.destination_stream_id = async_stream_id;
+        CHECK(async_instructions_.try_emplace(callsite.instruction(), streams)
+                  .second);
+      } else {
+        // Synchronous calls will result in the target computation being invoked
+        // using the same `ExecutionStreamId`.
+        queue.emplace_back(callsite.called_computations()[0],
+                           pending.stream_id);
+      }
+    }
+
+    // And finally, we need to assign `ExecutionStreamIds` to all asynchronous
+    // instructions that are were not handled by the iteration over callsites
+    // above. These are the `async-updates` and `async-dones`. Both of these
+    // should share the `ExecutionStreamId` as the originating `async-starts`.
+    for (HloInstruction* instruction : pending.node->instructions()) {
+      if (!instruction->IsAsynchronous()) continue;
+      if (instruction->opcode() == HloOpcode::kAsyncStart) {
+        CHECK(async_instructions_.find(instruction) !=
+              async_instructions_.end());
+      } else {
+        HloInstruction* async_start =
+            Cast<HloAsyncInstruction>(instruction)->async_chain_start();
+        AsyncExecutionStreamIds async_start_streams =
+            async_instructions_.at(async_start);
+        CHECK(async_instructions_.try_emplace(instruction, async_start_streams)
+                  .second);
+      }
+    }
+  }
+}
+
+namespace {
+absl::Status StreamNotFoundError(const HloInstruction* instruction) {
+  return absl::NotFoundError(
+      absl::StrCat("No stream found for ", instruction->ToString(),
+                   "; it may not be reachable from the module's entrypoint."));
+}
+}  // namespace
+
+absl::StatusOr<ExecutionStreamId>
+ExecutionStreamAssignment::GetSyncExecutionStreamId(
+    const HloInstruction* instruction) const {
+  CHECK(!instruction->IsAsynchronous());
+  auto stream = sync_instructions_.find(instruction);
+  if (stream == sync_instructions_.end()) {
+    return StreamNotFoundError(instruction);
+  }
+  return stream->second;
+}
+
+absl::StatusOr<ExecutionStreamAssignment::AsyncExecutionStreamIds>
+ExecutionStreamAssignment::GetAsyncExecutionStreamIds(
+    const HloAsyncInstruction* instruction) const {
+  auto streams = async_instructions_.find(instruction);
+  if (streams == async_instructions_.end()) {
+    return StreamNotFoundError(instruction);
+  }
+  return streams->second;
+}
+
+}  // namespace xla::gpu
diff --git a/third_party/xla/xla/service/gpu/execution_stream_assignment.h b/third_party/xla/xla/service/gpu/execution_stream_assignment.h
new file mode 100644
index 00000000000000..8200b7730d9eb8
--- /dev/null
+++ b/third_party/xla/xla/service/gpu/execution_stream_assignment.h
@@ -0,0 +1,80 @@
+/* Copyright 2024 The OpenXLA Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+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.
+==============================================================================*/
+
+#ifndef XLA_SERVICE_GPU_EXECUTION_STREAM_ASSIGNMENT_H_
+#define XLA_SERVICE_GPU_EXECUTION_STREAM_ASSIGNMENT_H_
+
+#include "absl/container/flat_hash_map.h"
+#include "absl/log/check.h"
+#include "absl/status/statusor.h"
+#include "xla/hlo/ir/hlo_instruction.h"
+#include "xla/hlo/ir/hlo_instructions.h"
+#include "xla/hlo/ir/hlo_module.h"
+#include "xla/service/gpu/runtime/thunk.h"
+
+namespace xla::gpu {
+
+// `ExecutionStreamAssignments` represent a mapping from `HloInstructions` to
+// `ExecutionStreamIds`. Asynchronous calls (`async-start`, `async-update`, and
+// `async-done`) result in the target computations being assigned new
+// `ExecutionStreamIds` to support concurrent execution.
+class ExecutionStreamAssignment {
+ public:
+  // The `HloModule` must be flat. In other words, there must be a one-to-one
+  // mapping between callsites and computations. One way to guarantee this is to
+  // pass the module through the `FlattenCallGraph` pass.
+  //
+  // The ExecutionStreamAssignment does not take ownership of the `HloModule`.
+  explicit ExecutionStreamAssignment(const HloModule* module);
+
+  // Returns the `ExecutionStreamId` for the given instruction, which *must* be
+  // synchronous. Returns an error if the instruction is not reachable
+  // from the module's entrypoint.
+  absl::StatusOr<ExecutionStreamId> GetSyncExecutionStreamId(
+      const HloInstruction* instruction) const;
+
+  // Returns the source and destination `ExecutionStreamIds` for the given
+  // instruction, which *must* be asynchronous. Returns an error if the
+  // instruction is not reachable from the module's entrypoint.
+  struct AsyncExecutionStreamIds {
+    // The `ExecutionStreamId` for the calling instruction (e.g. the computation
+    // that invokes `async-start`).
+    ExecutionStreamId source_stream_id;
+    // The `ExecutionStreamId` for the callee computation (e.g. the callee of an
+    // `async-start` instruction).
+    ExecutionStreamId destination_stream_id;
+  };
+  absl::StatusOr<AsyncExecutionStreamIds> GetAsyncExecutionStreamIds(
+      const HloAsyncInstruction* instruction) const;
+
+ private:
+  // Maps from `HloInstructions` to `ExecutionStreamIds` for synchronous and
+  // asynchronous instructions, respectively. Instructions that are not
+  // reachable from the module's entrypoint will not be present.
+  absl::flat_hash_map<HloInstruction*, ExecutionStreamId> sync_instructions_;
+  absl::flat_hash_map<HloInstruction*, AsyncExecutionStreamIds>
+      async_instructions_;
+};
+
+inline bool operator==(
+    const ExecutionStreamAssignment::AsyncExecutionStreamIds& first,
+    const ExecutionStreamAssignment::AsyncExecutionStreamIds& second) {
+  return first.source_stream_id == second.source_stream_id &&
+         first.destination_stream_id == second.destination_stream_id;
+}
+
+}  // namespace xla::gpu
+
+#endif  // XLA_SERVICE_GPU_EXECUTION_STREAM_ASSIGNMENT_H_
diff --git a/third_party/xla/xla/service/gpu/execution_stream_assignment_test.cc b/third_party/xla/xla/service/gpu/execution_stream_assignment_test.cc
new file mode 100644
index 00000000000000..507d6e9fcef6b4
--- /dev/null
+++ b/third_party/xla/xla/service/gpu/execution_stream_assignment_test.cc
@@ -0,0 +1,169 @@
+/* Copyright 2024 The OpenXLA Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+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.
+==============================================================================*/
+
+#include "xla/service/gpu/execution_stream_assignment.h"
+
+#include <memory>
+
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+#include "absl/status/status.h"
+#include "absl/strings/string_view.h"
+#include "xla/hlo/ir/hlo_casting_utils.h"
+#include "xla/hlo/ir/hlo_computation.h"
+#include "xla/hlo/ir/hlo_instruction.h"
+#include "xla/hlo/ir/hlo_instructions.h"
+#include "xla/hlo/ir/hlo_module.h"
+#include "xla/service/gpu/runtime/thunk.h"
+#include "xla/shape.h"
+#include "xla/shape_util.h"
+#include "xla/tests/hlo_test_base.h"
+#include "tsl/platform/status_matchers.h"
+#include "tsl/platform/statusor.h"
+
+using ::tsl::testing::IsOkAndHolds;
+using ::tsl::testing::StatusIs;
+
+using AsyncExecutionStreamIds =
+    ::xla::gpu::ExecutionStreamAssignment::AsyncExecutionStreamIds;
+
+namespace xla::gpu {
+namespace {
+
+class ExecutionStreamAssignmentTest : public HloTestBase {
+ protected:
+  // Adds expectations for the `ExecutionStreamId` for all synchronous
+  // `HloInstructions` in the given `HloComputation`.
+  void ExpectExecutionStreamForSyncInstructions(
+      const ExecutionStreamAssignment& assignment, HloComputation* computation,
+      ExecutionStreamId stream) const {
+    for (const HloInstruction* instruction : computation->instructions()) {
+      if (instruction->IsAsynchronous()) continue;
+      EXPECT_THAT(assignment.GetSyncExecutionStreamId(instruction),
+                  IsOkAndHolds(stream));
+    }
+  }
+
+  // Adds expectations for the `ExecutionStreamId` for all asynchronous
+  // `HloInstructions` in the given `HloComputation`.
+  void ExpectExecutionStreamForAsyncInstructions(
+      const ExecutionStreamAssignment& assignment, HloComputation* computation,
+      ExecutionStreamId source_stream,
+      ExecutionStreamId destination_stream) const {
+    for (const HloInstruction* instruction : computation->instructions()) {
+      if (!instruction->IsAsynchronous()) continue;
+      AsyncExecutionStreamIds expected_stream_ids;
+      expected_stream_ids.source_stream_id = source_stream;
+      expected_stream_ids.destination_stream_id = destination_stream;
+      EXPECT_THAT(assignment.GetAsyncExecutionStreamIds(
+                      Cast<HloAsyncInstruction>(instruction)),
+                  IsOkAndHolds(expected_stream_ids));
+    }
+  }
+
+  const Shape kTensorShape = ShapeUtil::MakeShape(F32, {2, 2});
+};
+
+TEST_F(ExecutionStreamAssignmentTest, AsyncFusion) {
+  // We'll create an `HloModule` with two nested `async-fusions`.
+  // ENTRY -> ASYNC-FUSION -> ASYNC-FUSION -> BINARY_OP
+  const char* kModuleStr = R"(
+    HloModule m
+
+    // Leaf computation.
+    leaf {
+      p0 = f32[2,2] parameter(0)
+      ROOT add = f32[2,2] add(p0, p0)
+    }
+
+    // Innermost `async-fusion`.
+    fusion {
+      p0 = f32[2,2] parameter(0)
+      start = ((f32[2,2]), f32[2,2], s32[]) fusion-start(p0),
+          kind=kLoop, calls=leaf
+      update = ((f32[2,2]), f32[2,2], s32[]) fusion-update(start)
+      ROOT done = f32[2,2] fusion-done(update)
+    }
+
+    // Outermost `async-fusion` and entrypoint for the module.
+    ENTRY entry {
+      p0 = f32[2,2] parameter(0)
+      start = ((f32[2,2]), f32[2,2], s32[]) fusion-start(p0),
+          kind=kLoop, calls=fusion
+      update = ((f32[2,2]), f32[2,2], s32[]) fusion-update(start)
+      ROOT done = f32[2,2] fusion-done(update)
+    }
+  )";
+  TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
+                          ParseAndReturnVerifiedModule(kModuleStr));
+
+  ExecutionStreamAssignment assignment(module.get());
+
+  // The outermost computation should run on `ExecutionStreamId(0)` and launch
+  // asynchronous work on `ExecutionStreamId(1)`.
+  ExpectExecutionStreamForSyncInstructions(
+      assignment, FindComputation(module.get(), "entry"), ExecutionStreamId(0));
+  ExpectExecutionStreamForAsyncInstructions(
+      assignment, FindComputation(module.get(), "entry"), ExecutionStreamId(0),
+      ExecutionStreamId(1));
+
+  // The nested computation should run on `ExecutionStreamId(1)` and launch
+  // asynchronous work on `ExecutionStreamId(2)`.
+  ExpectExecutionStreamForSyncInstructions(
+      assignment, FindComputation(module.get(), "fusion"),
+      ExecutionStreamId(1));
+  ExpectExecutionStreamForAsyncInstructions(
+      assignment, FindComputation(module.get(), "fusion"), ExecutionStreamId(1),
+      ExecutionStreamId(2));
+
+  // The innermost computation should run on `ExecutionStreamId(2)`
+  ExpectExecutionStreamForSyncInstructions(
+      assignment, FindComputation(module.get(), "leaf"), ExecutionStreamId(2));
+}
+
+TEST_F(ExecutionStreamAssignmentTest, UnreachableComputation) {
+  // We'll create an `HloModule` with two computations: the ENTRY computation,
+  // and another unreachable embedded computation.
+  const char* kModuleStr = R"(
+    HloModule m
+
+    // Unreachable computation.
+    unreachable {
+      p0 = f32[2,2] parameter(0)
+      ROOT add = f32[2,2] add(p0, p0)
+    }
+
+    ENTRY entry {
+      p0 = f32[2,2] parameter(0)
+      ROOT add = f32[2,2] add(p0, p0)
+    }
+  )";
+  TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
+                          ParseAndReturnVerifiedModule(kModuleStr));
+
+  ExecutionStreamAssignment assignment(module.get());
+  ExpectExecutionStreamForSyncInstructions(
+      assignment, FindComputation(module.get(), "entry"), ExecutionStreamId(0));
+
+  // Unreachable instructions should have no assigned `ExecutionStreamId`.
+  for (const HloInstruction* instruction :
+       FindComputation(module.get(), "unreachable")->instructions()) {
+    EXPECT_THAT(assignment.GetSyncExecutionStreamId(instruction),
+                StatusIs(absl::StatusCode::kNotFound));
+  }
+}
+
+}  // namespace
+}  // namespace xla::gpu

From d1459f2a6e6b5e3c72cfbcd3e634291947de3221 Mon Sep 17 00:00:00 2001
From: David Dunleavy <ddunleavy@google.com>
Date: Thu, 16 May 2024 14:08:37 -0700
Subject: [PATCH 302/478] Remove broken TSL tests from `TARGET_FILTERS`, tag
 them as `no_oss` instead

They will be fixed post TSL move

PiperOrigin-RevId: 634523661
---
 third_party/xla/.kokoro/linux/build.sh                   | 2 +-
 third_party/xla/third_party/tsl/tsl/platform/BUILD       | 1 +
 third_party/xla/third_party/tsl/tsl/platform/cloud/BUILD | 2 ++
 3 files changed, 4 insertions(+), 1 deletion(-)

diff --git a/third_party/xla/.kokoro/linux/build.sh b/third_party/xla/.kokoro/linux/build.sh
index f7e572c5046281..dc495910bfc4e5 100644
--- a/third_party/xla/.kokoro/linux/build.sh
+++ b/third_party/xla/.kokoro/linux/build.sh
@@ -51,7 +51,7 @@ docker run --name xla -w /tf/xla -itd --rm \
 TAGS_FILTER="-no_oss"
 ADDITIONAL_FLAGS=""
 RBE_FLAGS=""
-TARGET_FILTERS="-@local_tsl//tsl/platform:subprocess_test -@local_tsl//tsl/platform/cloud:google_auth_provider_test -@local_tsl//tsl/platform/cloud:oauth_client_test"
+TARGET_FILTERS=""
 
 if is_linux_gpu_job ; then
     TAGS_FILTER="$TAGS_FILTER,requires-gpu-nvidia"
diff --git a/third_party/xla/third_party/tsl/tsl/platform/BUILD b/third_party/xla/third_party/tsl/tsl/platform/BUILD
index 5f4dfaecbc1013..273fd530636675 100644
--- a/third_party/xla/third_party/tsl/tsl/platform/BUILD
+++ b/third_party/xla/third_party/tsl/tsl/platform/BUILD
@@ -1455,6 +1455,7 @@ tsl_cc_test(
         "//tsl/platform/testdata:test_noop",
         "//tsl/platform/testdata:test_stderr",
     ],
+    tags = ["no_oss"],  # TODO(b/327036247): revisit after this moves to XLA
     deps = [
         ":path",
         ":strcat",
diff --git a/third_party/xla/third_party/tsl/tsl/platform/cloud/BUILD b/third_party/xla/third_party/tsl/tsl/platform/cloud/BUILD
index 2ce4463f65b7ee..bff2db41d43626 100644
--- a/third_party/xla/third_party/tsl/tsl/platform/cloud/BUILD
+++ b/third_party/xla/third_party/tsl/tsl/platform/cloud/BUILD
@@ -455,6 +455,7 @@ tsl_cc_test(
         "//tsl/platform/cloud/testdata:service_account_credentials",
         "//tsl/platform/cloud/testdata:service_account_public_key",
     ],
+    tags = ["no_oss"],  # TODO(b/327036247): revisit after this moves to XLA
     deps = [
         ":http_request_fake",
         ":oauth_client",
@@ -478,6 +479,7 @@ tsl_cc_test(
         "//tsl/platform/cloud/testdata:application_default_credentials",
         "//tsl/platform/cloud/testdata:service_account_credentials",
     ],
+    tags = ["no_oss"],  # TODO(b/327036247): revisit after this moves to XLA
     deps = [
         ":google_auth_provider",
         ":http_request_fake",

From d835f88ca73cd23dda977e9effeb2a7ba19d0acd Mon Sep 17 00:00:00 2001
From: Kyle Lucke <klucke@google.com>
Date: Thu, 16 May 2024 14:11:24 -0700
Subject: [PATCH 303/478] Use absl::Status instead of xla::Status now that
 they're identical.

PiperOrigin-RevId: 634524626
---
 .../profiler/gpu/cupti_buffer_events.cc       |  14 +-
 .../xla/backends/profiler/gpu/cupti_tracer.cc |  55 +++---
 .../profiler/gpu/device_tracer_cuda.cc        |  25 ++-
 .../profiler/gpu/device_tracer_rocm.cc        |  24 +--
 third_party/xla/xla/client/lib/math.cc        |   3 +-
 third_party/xla/xla/client/lib/svd.cc         |   2 +-
 third_party/xla/xla/client/lib/tridiagonal.cc |  12 +-
 third_party/xla/xla/ffi/ffi.h                 |   2 +-
 third_party/xla/xla/ffi/ffi_api.cc            |  28 +--
 third_party/xla/xla/ffi/ffi_api.h             |  10 +-
 .../xla/xla/hlo/evaluator/hlo_evaluator.cc    | 167 +++++++++---------
 .../xla/xla/hlo/evaluator/hlo_evaluator.h     | 114 ++++++------
 .../evaluator/hlo_evaluator_typed_visitor.h   | 130 +++++++-------
 .../xla/xla/stream_executor/cuda/cuda_dnn.cc  |   8 +-
 .../xla/stream_executor/cuda/cuda_driver.cc   |   2 +-
 .../xla/stream_executor/cuda/cuda_executor.cc |   6 +-
 .../xla/xla/stream_executor/rocm/rocm_dnn.cc  |   2 +-
 .../xla/stream_executor/rocm/rocm_driver.cc   |   2 +-
 .../xla/stream_executor/rocm/rocm_executor.cc |   6 +-
 .../xla/tools/hlo_bisect/hlo_bisect_state.cc  |  15 +-
 .../xla/tools/hlo_bisect/hlo_bisect_state.h   |   2 +-
 .../xla/tools/hlo_bisect/hlo_bisect_utils.cc  |  19 +-
 .../functional_hlo_runner.cc                  |  16 +-
 .../functional_hlo_runner.h                   |  17 +-
 24 files changed, 349 insertions(+), 332 deletions(-)

diff --git a/third_party/xla/xla/backends/profiler/gpu/cupti_buffer_events.cc b/third_party/xla/xla/backends/profiler/gpu/cupti_buffer_events.cc
index 8b0d6977eba605..4c3ce94e1651a4 100644
--- a/third_party/xla/xla/backends/profiler/gpu/cupti_buffer_events.cc
+++ b/third_party/xla/xla/backends/profiler/gpu/cupti_buffer_events.cc
@@ -26,7 +26,6 @@ namespace profiler {
 namespace {
 
 using absl::OkStatus;
-using absl::Status;
 using absl::StatusCode;
 
 // CUPTI from CUDA 11.6 adds information about the hardware channel that ops
@@ -365,11 +364,10 @@ void AddSynchronizationActivityEvent(
   collector.receive(std::move(event));
 }
 
-static Status ConvertActivityBuffer(CuptiEventCollectorDelegate &collector,
-                                    uint8_t *buffer, const size_t size,
-                                    const size_t max_activity_event_count,
-                                    size_t &total_activity_event_count,
-                                    size_t &dropped_activity_event_count) {
+static absl::Status ConvertActivityBuffer(
+    CuptiEventCollectorDelegate &collector, uint8_t *buffer, const size_t size,
+    const size_t max_activity_event_count, size_t &total_activity_event_count,
+    size_t &dropped_activity_event_count) {
   CuptiInterface *cupti_interface = GetCuptiInterface();
   CUpti_Activity *record = nullptr;
   while (true) {
@@ -430,8 +428,8 @@ static Status ConvertActivityBuffer(CuptiEventCollectorDelegate &collector,
       break;
     } else {
       LOG(WARNING) << "CUPTI parse ACTIVITY buffer error: " << status;
-      return Status(StatusCode::kInternal,
-                    "Parse cupti activity buffer error.");
+      return absl::Status(StatusCode::kInternal,
+                          "Parse cupti activity buffer error.");
     }
   }
   VLOG(3) << "CUPTI tracer post-process one ACTIVITY buffer of size: " << size
diff --git a/third_party/xla/xla/backends/profiler/gpu/cupti_tracer.cc b/third_party/xla/xla/backends/profiler/gpu/cupti_tracer.cc
index 67e854ed9d3049..913171fc960d76 100644
--- a/third_party/xla/xla/backends/profiler/gpu/cupti_tracer.cc
+++ b/third_party/xla/xla/backends/profiler/gpu/cupti_tracer.cc
@@ -35,7 +35,6 @@ namespace profiler {
 namespace {
 
 using absl::OkStatus;
-using absl::Status;
 using tsl::Env;
 using tsl::profiler::AnnotationStack;
 
@@ -49,7 +48,7 @@ class CuptiApiTracingDisabler {
   ~CuptiApiTracingDisabler() { internalCuCall--; }
 };
 
-Status ToStatus(CUptiResult result) {
+absl::Status ToStatus(CUptiResult result) {
   if (result == CUPTI_SUCCESS) {
     return OkStatus();
   }
@@ -58,7 +57,7 @@ Status ToStatus(CUptiResult result) {
   return tsl::errors::Unavailable("CUPTI error: ", str ? str : "<unknown>");
 }
 
-Status ToStatus(CUresult result) {
+absl::Status ToStatus(CUresult result) {
   if (result == CUDA_SUCCESS) {
     return OkStatus();
   }
@@ -67,7 +66,7 @@ Status ToStatus(CUresult result) {
   return tsl::errors::Unavailable("CUDA error: ", str ? str : "<unknown>");
 }
 
-inline void LogIfError(const Status &status) {
+inline void LogIfError(const absl::Status &status) {
   if (status.ok()) return;
   LOG(ERROR) << status.message();
 }
@@ -302,8 +301,9 @@ void CUPTIAPI ProcessCuptiActivityBuffer(CUcontext context, uint32_t stream_id,
           << " size: " << size << " valid_size: " << valid_size;
   VLOG(3) << "Activity profile for stream " << stream_id;
 
-  Status status = CuptiTracer::GetCuptiTracerSingleton()->ProcessActivityBuffer(
-      context, stream_id, buffer, valid_size);
+  absl::Status status =
+      CuptiTracer::GetCuptiTracerSingleton()->ProcessActivityBuffer(
+          context, stream_id, buffer, valid_size);
   if (!status.ok()) {
     LOG(ERROR) << status;
   }
@@ -657,17 +657,17 @@ class CuptiDriverApiHookWithActivityApi : public CuptiDriverApiHook {
         cupti_interface_(cupti_interface),
         collector_(collector) {}
 
-  Status OnDriverApiEnter(int device_id, CUpti_CallbackDomain domain,
-                          CUpti_CallbackId cbid,
-                          const CUpti_CallbackData *cbdata) override {
+  absl::Status OnDriverApiEnter(int device_id, CUpti_CallbackDomain domain,
+                                CUpti_CallbackId cbid,
+                                const CUpti_CallbackData *cbdata) override {
     // Stash away the current Cupti timestamp into cbdata.
     *cbdata->correlationData =
         option_.required_callback_api_events ? CuptiTracer::GetTimestamp() : 0;
     return OkStatus();
   }
-  Status OnDriverApiExit(int device_id, CUpti_CallbackDomain domain,
-                         CUpti_CallbackId cbid,
-                         const CUpti_CallbackData *cbdata) override {
+  absl::Status OnDriverApiExit(int device_id, CUpti_CallbackDomain domain,
+                               CUpti_CallbackId cbid,
+                               const CUpti_CallbackData *cbdata) override {
     // If we are not collecting CPU events from Callback API, we can return now.
     if (!option_.required_callback_api_events) {
       return OkStatus();
@@ -680,7 +680,7 @@ class CuptiDriverApiHookWithActivityApi : public CuptiDriverApiHook {
     return AddDriverApiCallbackEvent(collector_, cupti_interface_, device_id,
                                      start_tsc, end_tsc, domain, cbid, cbdata);
   }
-  Status SyncAndFlush() override {
+  absl::Status SyncAndFlush() override {
     if (option_.sync_devices_before_stop) {
       CuptiApiTracingDisabler disabler;
       absl::MutexLock lock(&mutex_);
@@ -724,7 +724,7 @@ class CuptiDriverApiHookWithActivityApi : public CuptiDriverApiHook {
 
 }  // namespace
 
-/*static*/ Status CuptiDriverApiHook::AddDriverApiCallbackEvent(
+/*static*/ absl::Status CuptiDriverApiHook::AddDriverApiCallbackEvent(
     CuptiTraceCollector *collector, CuptiInterface *cupti_interface,
     int device_id, uint64_t start_tsc, uint64_t end_tsc,
     CUpti_CallbackDomain domain, CUpti_CallbackId cbid,
@@ -904,7 +904,7 @@ void CuptiTracer::Enable(const CuptiTracerOptions &option,
   cupti_driver_api_hook_.reset(new CuptiDriverApiHookWithActivityApi(
       option, cupti_interface_, collector));
 
-  Status status = EnableApiTracing();
+  absl::Status status = EnableApiTracing();
   need_root_access_ |= status.code() == tsl::error::PERMISSION_DENIED;
   if (!status.ok()) return;
 
@@ -936,7 +936,7 @@ void CuptiTracer::Disable() {
   tsl::profiler::AnnotationStack::Enable(false);
 }
 
-Status CuptiTracer::EnableApiTracing() {
+absl::Status CuptiTracer::EnableApiTracing() {
   if (api_tracing_enabled_) return OkStatus();
 
   VLOG(1) << "Enable subscriber";
@@ -964,7 +964,7 @@ Status CuptiTracer::EnableApiTracing() {
   return OkStatus();
 }
 
-Status CuptiTracer::DisableApiTracing() {
+absl::Status CuptiTracer::DisableApiTracing() {
   if (!api_tracing_enabled_) return OkStatus();
 
   api_tracing_enabled_ = false;
@@ -989,7 +989,7 @@ Status CuptiTracer::DisableApiTracing() {
   return OkStatus();
 }
 
-Status CuptiTracer::EnableActivityTracing() {
+absl::Status CuptiTracer::EnableActivityTracing() {
   if (activity_tracing_enabled_) return OkStatus();
   PrepareActivityStart();
   if (!option_->activities_selected.empty()) {
@@ -1018,7 +1018,7 @@ Status CuptiTracer::EnableActivityTracing() {
   return OkStatus();
 }
 
-Status CuptiTracer::DisableActivityTracing() {
+absl::Status CuptiTracer::DisableActivityTracing() {
   if (activity_tracing_enabled_) {
     VLOG(1) << "Disabling activity tracing for "
             << option_->activities_selected.size() << " activities";
@@ -1040,7 +1040,7 @@ Status CuptiTracer::DisableActivityTracing() {
   return OkStatus();
 }
 
-Status CuptiTracer::Finalize() {
+absl::Status CuptiTracer::Finalize() {
   if (option_->cupti_finalize) {
     VLOG(1) << "CuptiFinalize";
     RETURN_IF_CUPTI_ERROR(cupti_interface_->Finalize());
@@ -1059,8 +1059,8 @@ Status CuptiTracer::Finalize() {
   return 0;
 }
 
-Status CuptiTracer::HandleNVTXCallback(CUpti_CallbackId cbid,
-                                       const CUpti_CallbackData *cbdata) {
+absl::Status CuptiTracer::HandleNVTXCallback(CUpti_CallbackId cbid,
+                                             const CUpti_CallbackData *cbdata) {
   const CUpti_NvtxData *pdata =
       reinterpret_cast<const CUpti_NvtxData *>(cbdata);
   if (cbid == CUPTI_CBID_NVTX_nvtxDomainRangePushEx) {
@@ -1078,9 +1078,9 @@ Status CuptiTracer::HandleNVTXCallback(CUpti_CallbackId cbid,
   return OkStatus();
 }
 
-Status CuptiTracer::HandleCallback(CUpti_CallbackDomain domain,
-                                   CUpti_CallbackId cbid,
-                                   const CUpti_CallbackData *cbdata) {
+absl::Status CuptiTracer::HandleCallback(CUpti_CallbackDomain domain,
+                                         CUpti_CallbackId cbid,
+                                         const CUpti_CallbackData *cbdata) {
   if (!api_tracing_enabled_) return OkStatus();    // already unsubscribed.
   if (!cupti_driver_api_hook_) return OkStatus();  // already unsubscribed.
   if (domain == CUPTI_CB_DOMAIN_NVTX) return HandleNVTXCallback(cbid, cbdata);
@@ -1190,8 +1190,9 @@ static size_t CountCuptiActivityEvent(uint8_t *buffer, size_t size) {
   return total_event_count;
 }
 
-Status CuptiTracer::ProcessActivityBuffer(CUcontext context, uint32_t stream_id,
-                                          uint8_t *buffer, size_t size) {
+absl::Status CuptiTracer::ProcessActivityBuffer(CUcontext context,
+                                                uint32_t stream_id,
+                                                uint8_t *buffer, size_t size) {
   absl::Cleanup buffer_cleanup = [&]() {
     if (buffer) activity_buffers_->ReclaimBuffer(buffer);
   };
diff --git a/third_party/xla/xla/backends/profiler/gpu/device_tracer_cuda.cc b/third_party/xla/xla/backends/profiler/gpu/device_tracer_cuda.cc
index d7bb2524b66762..741a4129044e74 100644
--- a/third_party/xla/xla/backends/profiler/gpu/device_tracer_cuda.cc
+++ b/third_party/xla/xla/backends/profiler/gpu/device_tracer_cuda.cc
@@ -40,7 +40,6 @@ namespace xla {
 namespace profiler {
 
 using absl::OkStatus;
-using absl::Status;
 using tensorflow::ProfileOptions;
 using tensorflow::profiler::XSpace;
 using tsl::ReadBoolFromEnvVar;
@@ -55,13 +54,13 @@ class GpuTracer : public tsl::profiler::ProfilerInterface {
   ~GpuTracer() override {}
 
   // GpuTracer interface:
-  Status Start() override;
-  Status Stop() override;
-  Status CollectData(XSpace* space) override;
+  absl::Status Start() override;
+  absl::Status Stop() override;
+  absl::Status CollectData(XSpace* space) override;
 
  private:
-  Status DoStart();
-  Status DoStop();
+  absl::Status DoStart();
+  absl::Status DoStop();
 
   enum State {
     kNotStarted,
@@ -77,7 +76,7 @@ class GpuTracer : public tsl::profiler::ProfilerInterface {
   std::unique_ptr<CuptiTraceCollector> cupti_collector_;
 };
 
-Status GpuTracer::DoStart() {
+absl::Status GpuTracer::DoStart() {
   if (!cupti_tracer_->IsAvailable()) {
     return tsl::errors::Unavailable("Another profile session running.");
   }
@@ -159,8 +158,8 @@ Status GpuTracer::DoStart() {
   return OkStatus();
 }
 
-Status GpuTracer::Start() {
-  Status status = DoStart();
+absl::Status GpuTracer::Start() {
+  absl::Status status = DoStart();
   if (status.ok()) {
     profiling_state_ = State::kStartedOk;
     return OkStatus();
@@ -170,20 +169,20 @@ Status GpuTracer::Start() {
   }
 }
 
-Status GpuTracer::DoStop() {
+absl::Status GpuTracer::DoStop() {
   cupti_tracer_->Disable();
   return OkStatus();
 }
 
-Status GpuTracer::Stop() {
+absl::Status GpuTracer::Stop() {
   if (profiling_state_ == State::kStartedOk) {
-    Status status = DoStop();
+    absl::Status status = DoStop();
     profiling_state_ = status.ok() ? State::kStoppedOk : State::kStoppedError;
   }
   return OkStatus();
 }
 
-Status GpuTracer::CollectData(XSpace* space) {
+absl::Status GpuTracer::CollectData(XSpace* space) {
   VLOG(2) << "Collecting data to XSpace from GpuTracer.";
   switch (profiling_state_) {
     case State::kNotStarted:
diff --git a/third_party/xla/xla/backends/profiler/gpu/device_tracer_rocm.cc b/third_party/xla/xla/backends/profiler/gpu/device_tracer_rocm.cc
index 81eb2d192ea09a..560ae8033af66e 100644
--- a/third_party/xla/xla/backends/profiler/gpu/device_tracer_rocm.cc
+++ b/third_party/xla/xla/backends/profiler/gpu/device_tracer_rocm.cc
@@ -75,13 +75,13 @@ class GpuTracer : public profiler::ProfilerInterface {
   ~GpuTracer() override {}
 
   // GpuTracer interface:
-  Status Start() override;
-  Status Stop() override;
-  Status CollectData(XSpace* space) override;
+  absl::Status Start() override;
+  absl::Status Stop() override;
+  absl::Status CollectData(XSpace* space) override;
 
  private:
-  Status DoStart();
-  Status DoStop();
+  absl::Status DoStart();
+  absl::Status DoStop();
 
   RocmTracerOptions GetRocmTracerOptions();
 
@@ -184,7 +184,7 @@ RocmTraceCollectorOptions GpuTracer::GetRocmTraceCollectorOptions(
   return options;
 }
 
-Status GpuTracer::DoStart() {
+absl::Status GpuTracer::DoStart() {
   if (!rocm_tracer_->IsAvailable()) {
     return tsl::errors::Unavailable("Another profile session running.");
   }
@@ -208,8 +208,8 @@ Status GpuTracer::DoStart() {
   return OkStatus();
 }
 
-Status GpuTracer::Start() {
-  Status status = DoStart();
+absl::Status GpuTracer::Start() {
+  absl::Status status = DoStart();
   if (status.ok()) {
     profiling_state_ = State::kStartedOk;
     return OkStatus();
@@ -219,21 +219,21 @@ Status GpuTracer::Start() {
   }
 }
 
-Status GpuTracer::DoStop() {
+absl::Status GpuTracer::DoStop() {
   rocm_tracer_->Disable();
   AnnotationStack::Enable(false);
   return OkStatus();
 }
 
-Status GpuTracer::Stop() {
+absl::Status GpuTracer::Stop() {
   if (profiling_state_ == State::kStartedOk) {
-    Status status = DoStop();
+    absl::Status status = DoStop();
     profiling_state_ = status.ok() ? State::kStoppedOk : State::kStoppedError;
   }
   return OkStatus();
 }
 
-Status GpuTracer::CollectData(XSpace* space) {
+absl::Status GpuTracer::CollectData(XSpace* space) {
   switch (profiling_state_) {
     case State::kNotStarted:
       VLOG(3) << "No trace data collected, session wasn't started";
diff --git a/third_party/xla/xla/client/lib/math.cc b/third_party/xla/xla/client/lib/math.cc
index 60592de4ddc80f..772faba50d78c6 100644
--- a/third_party/xla/xla/client/lib/math.cc
+++ b/third_party/xla/xla/client/lib/math.cc
@@ -93,7 +93,8 @@ static XlaOp DoWithUpcastToF32(XlaOp operand,
 
 // TODO(jlebar): Use this function in more places in this file to restrict the
 // domain of other functions.
-static Status EnsureOperandIsRealFp(absl::string_view op_name, XlaOp operand) {
+static absl::Status EnsureOperandIsRealFp(absl::string_view op_name,
+                                          XlaOp operand) {
   auto& b = *operand.builder();
   TF_ASSIGN_OR_RETURN(auto shape, b.GetShape(operand));
   auto elem_ty = shape.element_type();
diff --git a/third_party/xla/xla/client/lib/svd.cc b/third_party/xla/xla/client/lib/svd.cc
index 9e7a33209d0b75..2304ec3b14c253 100644
--- a/third_party/xla/xla/client/lib/svd.cc
+++ b/third_party/xla/xla/client/lib/svd.cc
@@ -829,7 +829,7 @@ absl::StatusOr<SVDResult> SortBySingularValuesAndPostProcessing(
 SVDResult SVD(XlaOp a, int64_t max_iter, float epsilon,
               PrecisionConfig::Precision precision) {
   XlaBuilder* builder = a.builder();
-  auto return_error = [&](const Status& status) {
+  auto return_error = [&](const absl::Status& status) {
     SVDResult result;
     result.u = builder->ReportError(status);
     result.v = builder->ReportError(status);
diff --git a/third_party/xla/xla/client/lib/tridiagonal.cc b/third_party/xla/xla/client/lib/tridiagonal.cc
index 4daab7b4e8408b..459becdcfa4734 100644
--- a/third_party/xla/xla/client/lib/tridiagonal.cc
+++ b/third_party/xla/xla/client/lib/tridiagonal.cc
@@ -36,9 +36,9 @@ namespace tridiagonal {
 
 namespace {
 
-Status CheckSecondToLastDimension(const Shape& op_shape, int64_t rank,
-                                  int64_t expected,
-                                  const std::string& op_name) {
+absl::Status CheckSecondToLastDimension(const Shape& op_shape, int64_t rank,
+                                        int64_t expected,
+                                        const std::string& op_name) {
   const auto actual_num_dims = ShapeUtil::GetDimension(op_shape, rank - 2);
 
   if (actual_num_dims != expected) {
@@ -119,9 +119,9 @@ struct TridiagonalMatMulShapeParams {
   PrimitiveType element_type;
 };
 
-Status ValidateTridiagonalMatMulDiagonal(const Shape& diagonal_shape,
-                                         const std::string_view diagonal_name,
-                                         const Shape& rhs_shape) {
+absl::Status ValidateTridiagonalMatMulDiagonal(
+    const Shape& diagonal_shape, const std::string_view diagonal_name,
+    const Shape& rhs_shape) {
   const int64_t diagonal_rank = diagonal_shape.rank();
   const int64_t rhs_rank = rhs_shape.rank();
   if (diagonal_rank != rhs_rank) {
diff --git a/third_party/xla/xla/ffi/ffi.h b/third_party/xla/xla/ffi/ffi.h
index 629e0f02071d48..3100bc3a4fa3b7 100644
--- a/third_party/xla/xla/ffi/ffi.h
+++ b/third_party/xla/xla/ffi/ffi.h
@@ -360,7 +360,7 @@ struct CtxDecoding<UserData<T>> {
 
 template <>
 struct ResultEncoding<Status> {
-  static XLA_FFI_Error* Encode(const XLA_FFI_Api* api, Status status) {
+  static XLA_FFI_Error* Encode(const XLA_FFI_Api* api, absl::Status status) {
     return api->internal_api->XLA_FFI_INTERNAL_Error_Forward(&status);
   }
 };
diff --git a/third_party/xla/xla/ffi/ffi_api.cc b/third_party/xla/xla/ffi/ffi_api.cc
index 73c6c7aba412af..0bc59e99f51248 100644
--- a/third_party/xla/xla/ffi/ffi_api.cc
+++ b/third_party/xla/xla/ffi/ffi_api.cc
@@ -63,14 +63,15 @@ bool IsCommandBufferCompatible(XLA_FFI_Handler_Traits traits) {
 // WARNING: These functions defined in `call_frame.h` as we need to make them
 // available without having to depend on `ffi.h` header.
 
-Status TakeStatus(XLA_FFI_Error* error) {
+absl::Status TakeStatus(XLA_FFI_Error* error) {
   if (error == nullptr) return absl::OkStatus();
-  Status status = std::move(error->status);
+  absl::Status status = std::move(error->status);
   delete error;
   return status;
 }
 
-Status Call(Ffi& handler, CallFrame& call_frame, const CallOptions& options) {
+absl::Status Call(Ffi& handler, CallFrame& call_frame,
+                  const CallOptions& options) {
   XLA_FFI_ExecutionContext ctx = {
       options.run_options, options.called_computation,
       internal::ScopedExecutionContext::GetCallExecutionContext(options)};
@@ -78,8 +79,8 @@ Status Call(Ffi& handler, CallFrame& call_frame, const CallOptions& options) {
   return TakeStatus(handler.Call(&ffi_call_frame));
 }
 
-Status Call(XLA_FFI_Handler* handler, CallFrame& call_frame,
-            const CallOptions& options) {
+absl::Status Call(XLA_FFI_Handler* handler, CallFrame& call_frame,
+                  const CallOptions& options) {
   XLA_FFI_ExecutionContext ctx = {
       options.run_options, options.called_computation,
       internal::ScopedExecutionContext::GetCallExecutionContext(options)};
@@ -127,9 +128,10 @@ static HandlerRegistry& GetHandlerRegistry() {
   return *registry;
 }
 
-static Status RegisterHandler(std::string_view name, std::string_view platform,
-                              XLA_FFI_Handler* handler,
-                              XLA_FFI_Handler_Traits traits) {
+static absl::Status RegisterHandler(std::string_view name,
+                                    std::string_view platform,
+                                    XLA_FFI_Handler* handler,
+                                    XLA_FFI_Handler_Traits traits) {
   auto emplaced = GetHandlerRegistry().try_emplace(
       MakeHandlerKey(name, platform), HandlerRegistration{handler, traits});
   if (!emplaced.second)
@@ -172,8 +174,8 @@ static std::string StructSizeErrorMsg(std::string_view struct_name,
                       XLA_FFI_API_MAJOR, ".", XLA_FFI_API_MINOR, ".");
 }
 
-static Status ActualStructSizeIsGreaterOrEqual(std::string_view struct_name,
-                                               size_t expected, size_t actual) {
+static absl::Status ActualStructSizeIsGreaterOrEqual(
+    std::string_view struct_name, size_t expected, size_t actual) {
   if (actual < expected) {
     return absl::InvalidArgumentError(
         StructSizeErrorMsg(struct_name, expected, actual));
@@ -225,7 +227,7 @@ static absl::StatusCode ToStatusCode(XLA_FFI_Error_Code errc) {
 
 #define XLA_FFI_RETURN_IF_ERROR(expr)                                   \
   do {                                                                  \
-    Status _status = (expr);                                            \
+    absl::Status _status = (expr);                                      \
     if (!_status.ok()) {                                                \
       XLA_FFI_Error* _c_status = new XLA_FFI_Error{std::move(_status)}; \
       return _c_status;                                                 \
@@ -241,7 +243,7 @@ static XLA_FFI_Error* XLA_FFI_Error_Create(XLA_FFI_Error_Create_Args* args) {
 }
 
 static void XLA_FFI_Error_GetMessage(XLA_FFI_Error_GetMessage_Args* args) {
-  Status struct_size_check = ActualStructSizeIsGreaterOrEqual(
+  absl::Status struct_size_check = ActualStructSizeIsGreaterOrEqual(
       "XLA_FFI_Error_GetMessage", XLA_FFI_Error_GetMessage_Args_STRUCT_SIZE,
       args->struct_size);
   if (!struct_size_check.ok()) {
@@ -253,7 +255,7 @@ static void XLA_FFI_Error_GetMessage(XLA_FFI_Error_GetMessage_Args* args) {
 }
 
 static void XLA_FFI_Error_Destroy(XLA_FFI_Error_Destroy_Args* args) {
-  Status struct_size_check = ActualStructSizeIsGreaterOrEqual(
+  absl::Status struct_size_check = ActualStructSizeIsGreaterOrEqual(
       "XLA_FFI_Error_Destroy", XLA_FFI_Error_Destroy_Args_STRUCT_SIZE,
       args->struct_size);
   if (!struct_size_check.ok()) {
diff --git a/third_party/xla/xla/ffi/ffi_api.h b/third_party/xla/xla/ffi/ffi_api.h
index dc50916f8ab5b1..766fb83c83026c 100644
--- a/third_party/xla/xla/ffi/ffi_api.h
+++ b/third_party/xla/xla/ffi/ffi_api.h
@@ -50,13 +50,13 @@ struct CallOptions {
 
 // Takes ownership of the XLA FFI error and returns underlying status. Frees
 // `error` if it's not nullptr; returns OK status otherwise.
-Status TakeStatus(XLA_FFI_Error* error);
+absl::Status TakeStatus(XLA_FFI_Error* error);
 
-Status Call(Ffi& handler, CallFrame& call_frame,
-            const CallOptions& options = {});
+absl::Status Call(Ffi& handler, CallFrame& call_frame,
+                  const CallOptions& options = {});
 
-Status Call(XLA_FFI_Handler* handler, CallFrame& call_frame,
-            const CallOptions& options = {});
+absl::Status Call(XLA_FFI_Handler* handler, CallFrame& call_frame,
+                  const CallOptions& options = {});
 
 namespace internal {
 // This is an internal workaround to override FFI execution context for FFI
diff --git a/third_party/xla/xla/hlo/evaluator/hlo_evaluator.cc b/third_party/xla/xla/hlo/evaluator/hlo_evaluator.cc
index 4ec41dc7d81b4a..7e0a9b1627e5a4 100644
--- a/third_party/xla/xla/hlo/evaluator/hlo_evaluator.cc
+++ b/third_party/xla/xla/hlo/evaluator/hlo_evaluator.cc
@@ -159,9 +159,10 @@ std::optional<bool> GetInstructionStaticValueAsBool(
 template <PrimitiveType kType>
 struct PopulateParallelImpl {
   using NativeT = NativeTypeOf<kType>;
-  static Status Run(Literal& literal,
-                    absl::FunctionRef<Literal(absl::Span<const int64_t>, int)>
-                        literal_generator) {
+  static absl::Status Run(
+      Literal& literal,
+      absl::FunctionRef<Literal(absl::Span<const int64_t>, int)>
+          literal_generator) {
     return literal.PopulateParallel<NativeT>(
         [&literal_generator](absl::Span<const int64_t> output_index,
                              int thread_id) {
@@ -174,7 +175,7 @@ struct PopulateParallelImpl {
 template <PrimitiveType kType>
 struct PopulateImpl {
   using NativeT = NativeTypeOf<kType>;
-  static Status Run(
+  static absl::Status Run(
       Literal& literal,
       absl::FunctionRef<Literal(absl::Span<const int64_t>)> literal_generator) {
     return literal.Populate<NativeT>(
@@ -192,10 +193,10 @@ struct PopulateImpl {
 // to small templated helpers just for the parts that require manipulating the
 // native types to avoid templating the whole implementations.
 template <template <PrimitiveType> typename Trait, typename F>
-Status Apply(Literal& literal, F&& literal_generator) {
+absl::Status Apply(Literal& literal, F&& literal_generator) {
   return primitive_util::PrimitiveTypeSwitch<Status>(
       [&, literal_generator = std::forward<F>(literal_generator)](
-          auto primitive_type_constant) -> Status {
+          auto primitive_type_constant) -> absl::Status {
         if constexpr (primitive_util::IsArrayType(primitive_type_constant)) {
           return Trait<primitive_type_constant>::Run(
               literal, std::move(literal_generator));
@@ -216,7 +217,7 @@ enum class EvalErrorDetail : uint32_t {
   kDynamicValueDependence = 0,
 };
 
-std::optional<EvalErrorDetail> ParseEvalErrorDetail(const Status& error) {
+std::optional<EvalErrorDetail> ParseEvalErrorDetail(const absl::Status& error) {
   auto error_detail = error.GetPayload(kEvalErrorDetailUrl);
   if (!error_detail.has_value() && error_detail->empty()) {
     return std::nullopt;
@@ -225,11 +226,12 @@ std::optional<EvalErrorDetail> ParseEvalErrorDetail(const Status& error) {
       absl::little_endian::Load32(error_detail->Flatten().data()));
 }
 
-Status MakeEvalErrorDueToParamOrInfeed(const HloInstruction& eval_instruction) {
-  Status error = tsl::errors::FailedPrecondition(
+absl::Status MakeEvalErrorDueToParamOrInfeed(
+    const HloInstruction& eval_instruction) {
+  absl::Status error = absl::FailedPreconditionError(absl::StrCat(
       "Failed to evaluate instruction (", eval_instruction.name(),
       ") since it depends on infeed or parameters to its parent computation (",
-      eval_instruction.parent()->name(), ").");
+      eval_instruction.parent()->name(), ")."));
   std::string error_payload;
   error_payload.resize(sizeof(EvalErrorDetail));
   absl::little_endian::Store32(
@@ -1067,7 +1069,7 @@ absl::StatusOr<Literal> HloEvaluator::EvaluateDotOp(
   return Evaluate(cloned_instruction.get());
 }
 
-Status HloEvaluator::EvaluateParameterFromCallerArgument(
+absl::Status HloEvaluator::EvaluateParameterFromCallerArgument(
     const HloInstruction* parameter, const ShapeIndex& shape_index) {
   CHECK(!evaluated_.contains(parameter));
   const HloComputation* parent_computation = parameter->parent();
@@ -1155,7 +1157,7 @@ DimensionVector HloEvaluator::MakeDimMultipliers(const Shape& shape) {
   return v;
 }
 
-Status HloEvaluator::EvaluateInternal(
+absl::Status HloEvaluator::EvaluateInternal(
     const HloInstruction* instruction, const ShapeIndex& shape_index,
     bool recursively_evaluate_nonconstant_operands) {
   // Don't need to evaluate this instruction again if it has already been
@@ -1198,7 +1200,7 @@ Status HloEvaluator::EvaluateInternal(
         }
       }
       if (call_graph_cache_ && tuple_points_to_analysis_cache_) {
-        Status argument_eval_status =
+        absl::Status argument_eval_status =
             EvaluateParameterFromCallerArgument(instruction, shape_index);
         if (!argument_eval_status.ok()) {
           VLOG(4) << "Failed to evaluate parameter " << instruction->name()
@@ -1238,7 +1240,7 @@ Status HloEvaluator::EvaluateInternal(
   return OkStatus();
 }
 
-Status HloEvaluator::HandleBitcast(const HloInstruction* bitcast) {
+absl::Status HloEvaluator::HandleBitcast(const HloInstruction* bitcast) {
   const Literal& operand_literal = GetEvaluatedLiteralFor(bitcast->operand(0));
   Literal result(bitcast->shape());
   // Bitcast output is allowed to be smaller than the input if the backend-
@@ -1252,7 +1254,7 @@ Status HloEvaluator::HandleBitcast(const HloInstruction* bitcast) {
   return OkStatus();
 }
 
-Status HloEvaluator::HandleBitcastConvert(const HloInstruction* convert) {
+absl::Status HloEvaluator::HandleBitcastConvert(const HloInstruction* convert) {
   const HloInstruction* operand = convert->operand(0);
   TF_ASSIGN_OR_RETURN(
       Literal result,
@@ -1262,7 +1264,7 @@ Status HloEvaluator::HandleBitcastConvert(const HloInstruction* convert) {
   return OkStatus();
 }
 
-Status HloEvaluator::HandleGetDimensionSize(
+absl::Status HloEvaluator::HandleGetDimensionSize(
     const HloInstruction* get_dimension_size) {
   const HloInstruction* operand = get_dimension_size->operand(0);
   int64_t dim = get_dimension_size->dimension();
@@ -1287,7 +1289,7 @@ Status HloEvaluator::HandleGetDimensionSize(
   return OkStatus();
 }
 
-Status HloEvaluator::HandleSetDimensionSize(
+absl::Status HloEvaluator::HandleSetDimensionSize(
     const HloInstruction* set_dimension_size) {
   const Literal& operand_literal =
       GetEvaluatedLiteralFor(set_dimension_size->operand(0));
@@ -1302,7 +1304,7 @@ Status HloEvaluator::HandleSetDimensionSize(
   return OkStatus();
 }
 
-Status HloEvaluator::HandleParameter(const HloInstruction* parameter) {
+absl::Status HloEvaluator::HandleParameter(const HloInstruction* parameter) {
   if (!IsAlreadyEvaluated(parameter, visitor_shape_index_)) {
     if (!enable_partial_evaluation_) {
       return tsl::errors::FailedPrecondition(
@@ -1333,7 +1335,7 @@ Status HloEvaluator::HandleParameter(const HloInstruction* parameter) {
   return OkStatus();
 }
 
-Status HloEvaluator::HandleInfeed(const HloInstruction* infeed) {
+absl::Status HloEvaluator::HandleInfeed(const HloInstruction* infeed) {
   if (!enable_partial_evaluation_) {
     return tsl::errors::FailedPrecondition(
         "Failed to evaluate instruction since its operands are unknown "
@@ -1344,24 +1346,25 @@ Status HloEvaluator::HandleInfeed(const HloInstruction* infeed) {
   return OkStatus();
 }
 
-Status HloEvaluator::HandleConstant(const HloInstruction*) {
+absl::Status HloEvaluator::HandleConstant(const HloInstruction*) {
   return OkStatus();
 }
 
-Status HloEvaluator::HandleReshape(const HloInstruction* reshape) {
+absl::Status HloEvaluator::HandleReshape(const HloInstruction* reshape) {
   TF_ASSIGN_OR_RETURN(evaluated_[reshape],
                       GetEvaluatedLiteralFor(reshape->operand(0))
                           .Reshape(reshape->shape().dimensions()));
   return OkStatus();
 }
 
-Status HloEvaluator::HandleTranspose(const HloInstruction* transpose) {
+absl::Status HloEvaluator::HandleTranspose(const HloInstruction* transpose) {
   evaluated_[transpose] = GetEvaluatedLiteralFor(transpose->operand(0))
                               .Transpose(transpose->dimensions());
   return OkStatus();
 }
 
-Status HloEvaluator::HandleConcatenate(const HloInstruction* concatenate) {
+absl::Status HloEvaluator::HandleConcatenate(
+    const HloInstruction* concatenate) {
   absl::Span<HloInstruction* const> operands(concatenate->operands());
   // The result concatenate dimension is going to be the sum of all
   // concatenate dimensions of the operands taking part of the operation.
@@ -1402,11 +1405,11 @@ Status HloEvaluator::HandleConcatenate(const HloInstruction* concatenate) {
   return OkStatus();
 }
 
-Status HloEvaluator::HandleIsFinite(const HloInstruction* is_finite) {
+absl::Status HloEvaluator::HandleIsFinite(const HloInstruction* is_finite) {
   auto operand = is_finite->operand(0);
   auto elem_ty = operand->shape().element_type();
   return primitive_util::PrimitiveTypeSwitch<Status>(
-      [&](auto primitive_type_constant) -> Status {
+      [&](auto primitive_type_constant) -> absl::Status {
         if constexpr (primitive_util::IsFloatingPointType(
                           primitive_type_constant)) {
           using NativeT = primitive_util::NativeTypeOf<primitive_type_constant>;
@@ -1426,10 +1429,10 @@ Status HloEvaluator::HandleIsFinite(const HloInstruction* is_finite) {
       elem_ty);
 }
 
-Status HloEvaluator::HandleReal(const HloInstruction* real) {
+absl::Status HloEvaluator::HandleReal(const HloInstruction* real) {
   auto operand = real->operand(0);
   return primitive_util::PrimitiveTypeSwitch<Status>(
-      [&](auto primitive_type_constant) -> Status {
+      [&](auto primitive_type_constant) -> absl::Status {
         if constexpr (primitive_util::IsFloatingPointType(
                           primitive_type_constant)) {
           using NativeT = primitive_util::NativeTypeOf<primitive_type_constant>;
@@ -1455,10 +1458,10 @@ Status HloEvaluator::HandleReal(const HloInstruction* real) {
       operand->shape().element_type());
 }
 
-Status HloEvaluator::HandleImag(const HloInstruction* imag) {
+absl::Status HloEvaluator::HandleImag(const HloInstruction* imag) {
   auto operand = imag->operand(0);
   return primitive_util::PrimitiveTypeSwitch<Status>(
-      [&](auto primitive_type_constant) -> Status {
+      [&](auto primitive_type_constant) -> absl::Status {
         if constexpr (primitive_util::IsFloatingPointType(
                           primitive_type_constant)) {
           using NativeT = primitive_util::NativeTypeOf<primitive_type_constant>;
@@ -1484,14 +1487,14 @@ Status HloEvaluator::HandleImag(const HloInstruction* imag) {
       operand->shape().element_type());
 }
 
-Status HloEvaluator::HandleComplex(const HloInstruction* complex) {
+absl::Status HloEvaluator::HandleComplex(const HloInstruction* complex) {
   const Literal& real = GetEvaluatedLiteralFor(complex->operand(0));
   const Literal& imag = GetEvaluatedLiteralFor(complex->operand(1));
   TF_RET_CHECK(ShapeUtil::Compatible(real.shape(), imag.shape()));
 
   Literal result(complex->shape());
   return primitive_util::PrimitiveTypeSwitch<Status>(
-      [&](auto primitive_type_constant) -> Status {
+      [&](auto primitive_type_constant) -> absl::Status {
         if constexpr (primitive_util::IsComplexType(primitive_type_constant)) {
           using NativeT = primitive_util::NativeTypeOf<primitive_type_constant>;
           TF_RETURN_IF_ERROR(result.Populate<NativeT>(
@@ -1509,7 +1512,7 @@ Status HloEvaluator::HandleComplex(const HloInstruction* complex) {
       complex->shape().element_type());
 }
 
-Status HloEvaluator::HandleCompare(const HloInstruction* compare) {
+absl::Status HloEvaluator::HandleCompare(const HloInstruction* compare) {
   ComparisonDirection direction = compare->comparison_direction();
   ComparisonOrder order = compare->comparison_order();
   auto lhs = compare->operand(0);
@@ -1525,7 +1528,7 @@ Status HloEvaluator::HandleCompare(const HloInstruction* compare) {
   const Literal& rhs_literal = GetEvaluatedLiteralFor(rhs);
   // Note here we switch on the operand's type.
   return primitive_util::PrimitiveTypeSwitch<Status>(
-      [&](auto primitive_type_constant) -> Status {
+      [&](auto primitive_type_constant) -> absl::Status {
         if constexpr (primitive_util::IsArrayType(primitive_type_constant)) {
           using NativeT = primitive_util::NativeTypeOf<primitive_type_constant>;
           TF_ASSIGN_OR_RETURN(evaluated_[compare],
@@ -1539,7 +1542,7 @@ Status HloEvaluator::HandleCompare(const HloInstruction* compare) {
       element_type);
 }
 
-Status HloEvaluator::HandleTuple(const HloInstruction* tuple) {
+absl::Status HloEvaluator::HandleTuple(const HloInstruction* tuple) {
   std::vector<const Literal*> operand_literals;
   std::vector<Literal> operand_literal_values;
   if (!visitor_shape_index_.empty()) {
@@ -1671,8 +1674,9 @@ class FftTransform {
     absl::c_reverse(fft_lengths_);
   }
 
-  Status ComputeFft(const HloInstruction* fft, const Literal& input_literal,
-                    Literal* output_literal) {
+  absl::Status ComputeFft(const HloInstruction* fft,
+                          const Literal& input_literal,
+                          Literal* output_literal) {
     const Shape& input_shape = input_literal.shape();
     const Shape& output_shape = fft->shape();
 
@@ -2171,7 +2175,8 @@ class FftTransform {
     }
   }
 
-  Status CheckParameters(const Shape& input_shape, const Shape& output_shape) {
+  absl::Status CheckParameters(const Shape& input_shape,
+                               const Shape& output_shape) {
     // Check FFT parameters.
     if (fft_rank_ <= 0) {
       return InvalidArgument("Zero or negative FFT rank.");
@@ -2243,7 +2248,7 @@ class FftTransform {
 
 }  // namespace
 
-Status HloEvaluator::HandleFft(const HloInstruction* fft) {
+absl::Status HloEvaluator::HandleFft(const HloInstruction* fft) {
   const Literal& input_literal = GetEvaluatedLiteralFor(fft->operand(0));
   Literal output_literal = Literal::CreateFromShape(fft->shape());
 
@@ -2379,7 +2384,7 @@ class OutputBatchIndexToInputIndex {
 
   // Populates index_vector_ by iterating over start_indices_ according to
   // index_vector_index_.
-  Status FetchIndexVector() {
+  absl::Status FetchIndexVector() {
     int64_t index_vector_dim = dim_numbers_.index_vector_dim();
     for (int64_t i = 0, e = index_vector_.size(); i < e; i++) {
       index_vector_index_[index_vector_dim] = i;
@@ -2530,7 +2535,7 @@ ReshapedGatherIndices(int64_t index_vector_dim, const Literal& start_indices,
   return std::cref(*reshaped_start_indices);
 }
 
-Status HloEvaluator::HandleGather(const HloInstruction* gather) {
+absl::Status HloEvaluator::HandleGather(const HloInstruction* gather) {
   Literal result = Literal::CreateFromShape(gather->shape());
   const Shape& shape = gather->shape();
   const GatherDimensionNumbers& dim_numbers =
@@ -2784,7 +2789,7 @@ class UpdateScatterIndexToInputIndex {
 
   // Populates index_vector_ by iterating over scatter_indices_ according to
   // index_vector_index_.
-  Status FetchIndexVector() {
+  absl::Status FetchIndexVector() {
     int64_t index_vector_dim = dim_numbers_.index_vector_dim();
     for (int64_t i = 0, e = index_vector_.size(); i < e; i++) {
       index_vector_index_[index_vector_dim] = i;
@@ -2914,7 +2919,7 @@ class UpdateWindowIndexToInputIndex {
 };
 }  // namespace
 
-Status HloEvaluator::HandleScatter(const HloInstruction* hlo) {
+absl::Status HloEvaluator::HandleScatter(const HloInstruction* hlo) {
   auto* scatter = DynCast<HloScatterInstruction>(hlo);
   const ScatterDimensionNumbers& dim_numbers =
       scatter->scatter_dimension_numbers();
@@ -3043,7 +3048,7 @@ Status HloEvaluator::HandleScatter(const HloInstruction* hlo) {
   return OkStatus();
 }
 
-Status HloEvaluator::HandleBroadcast(const HloInstruction* broadcast) {
+absl::Status HloEvaluator::HandleBroadcast(const HloInstruction* broadcast) {
   const Literal& operand = GetEvaluatedLiteralFor(broadcast->operand(0));
   TF_RET_CHECK(broadcast->shape().element_type() ==
                operand.shape().element_type())
@@ -3071,19 +3076,20 @@ Status HloEvaluator::HandleBroadcast(const HloInstruction* broadcast) {
   return OkStatus();
 }
 
-Status HloEvaluator::HandleAfterAll(const HloInstruction* after_all) {
+absl::Status HloEvaluator::HandleAfterAll(const HloInstruction* after_all) {
   evaluated_[after_all] = LiteralUtil::CreateToken();
   return OkStatus();
 }
 
-Status HloEvaluator::HandleAddDependency(const HloInstruction* add_dependency) {
+absl::Status HloEvaluator::HandleAddDependency(
+    const HloInstruction* add_dependency) {
   // AddDedendency just forwards its zero-th operand.
   evaluated_[add_dependency] =
       GetEvaluatedLiteralFor(add_dependency->operand(0)).Clone();
   return OkStatus();
 }
 
-Status HloEvaluator::HandleGetTupleElement(
+absl::Status HloEvaluator::HandleGetTupleElement(
     const HloInstruction* get_tuple_element) {
   const auto result_shape = get_tuple_element->shape();
   const int64_t index = get_tuple_element->tuple_index();
@@ -3106,13 +3112,13 @@ Status HloEvaluator::HandleGetTupleElement(
                                                 /*src_shape_index=*/{index});
 }
 
-Status HloEvaluator::HandleCopy(const HloInstruction* copy) {
+absl::Status HloEvaluator::HandleCopy(const HloInstruction* copy) {
   TF_RET_CHECK(ShapeUtil::Compatible(copy->shape(), copy->operand(0)->shape()));
   evaluated_[copy] = GetEvaluatedLiteralFor(copy->operand(0)).Clone();
   return OkStatus();
 }
 
-Status HloEvaluator::HandleAsyncStart(const HloInstruction* async_start) {
+absl::Status HloEvaluator::HandleAsyncStart(const HloInstruction* async_start) {
   std::vector<const Literal*> arg_literals;
   arg_literals.reserve(async_start->operands().size());
   for (auto operand : async_start->operands()) {
@@ -3143,7 +3149,8 @@ Status HloEvaluator::HandleAsyncStart(const HloInstruction* async_start) {
   return OkStatus();
 }
 
-Status HloEvaluator::HandleAsyncUpdate(const HloInstruction* async_update) {
+absl::Status HloEvaluator::HandleAsyncUpdate(
+    const HloInstruction* async_update) {
   const Literal& operand_tuple_literal =
       GetEvaluatedLiteralFor(async_update->operand(0));
   evaluated_[async_update] = Literal(async_update->shape());
@@ -3153,7 +3160,7 @@ Status HloEvaluator::HandleAsyncUpdate(const HloInstruction* async_update) {
   return OkStatus();
 }
 
-Status HloEvaluator::HandleAsyncDone(const HloInstruction* async_done) {
+absl::Status HloEvaluator::HandleAsyncDone(const HloInstruction* async_done) {
   const Literal& operand_tuple_literal =
       GetEvaluatedLiteralFor(async_done->operand(0));
   evaluated_[async_done] = Literal(async_done->shape());
@@ -3163,7 +3170,7 @@ Status HloEvaluator::HandleAsyncDone(const HloInstruction* async_done) {
   return OkStatus();
 }
 
-Status HloEvaluator::HandleCopyStart(const HloInstruction* copy_start) {
+absl::Status HloEvaluator::HandleCopyStart(const HloInstruction* copy_start) {
   if (copy_start->user_count() != 1 ||
       copy_start->users().at(0)->opcode() != HloOpcode::kCopyDone) {
     return tsl::errors::FailedPrecondition(
@@ -3183,7 +3190,7 @@ Status HloEvaluator::HandleCopyStart(const HloInstruction* copy_start) {
   return OkStatus();
 }
 
-Status HloEvaluator::HandleCopyDone(const HloInstruction* copy_done) {
+absl::Status HloEvaluator::HandleCopyDone(const HloInstruction* copy_done) {
   const HloInstruction* operand = copy_done->operand(0);
   if (operand->opcode() != HloOpcode::kCopyStart) {
     return tsl::errors::FailedPrecondition(
@@ -3200,7 +3207,7 @@ Status HloEvaluator::HandleCopyDone(const HloInstruction* copy_done) {
   return OkStatus();
 }
 
-Status HloEvaluator::HandleCall(const HloInstruction* call) {
+absl::Status HloEvaluator::HandleCall(const HloInstruction* call) {
   auto* computation = call->to_apply();
   auto operands = call->operands();
 
@@ -3222,7 +3229,7 @@ Status HloEvaluator::HandleCall(const HloInstruction* call) {
   return OkStatus();
 }
 
-Status HloEvaluator::HandleFusion(const HloInstruction* fusion) {
+absl::Status HloEvaluator::HandleFusion(const HloInstruction* fusion) {
   HloModuleConfig config;
   // Attach cloned computation to an empty HLO module so the existing ones are
   // not modified.
@@ -3260,7 +3267,8 @@ Status HloEvaluator::HandleFusion(const HloInstruction* fusion) {
   return OkStatus();
 }
 
-Status HloEvaluator::HandleConditional(const HloInstruction* conditional) {
+absl::Status HloEvaluator::HandleConditional(
+    const HloInstruction* conditional) {
   const auto& branch_index_literal =
       GetEvaluatedLiteralFor(conditional->operand(0));
   int branch_index;
@@ -3288,7 +3296,7 @@ Status HloEvaluator::HandleConditional(const HloInstruction* conditional) {
   return OkStatus();
 }
 
-Status HloEvaluator::HandleConvert(const HloInstruction* convert) {
+absl::Status HloEvaluator::HandleConvert(const HloInstruction* convert) {
   const HloInstruction* operand = convert->operand(0);
   TF_RET_CHECK(ShapeUtil::SameDimensions(operand->shape(), convert->shape()));
   TF_ASSIGN_OR_RETURN(Literal result, GetEvaluatedLiteralFor(operand).Convert(
@@ -3297,7 +3305,8 @@ Status HloEvaluator::HandleConvert(const HloInstruction* convert) {
   return OkStatus();
 }
 
-Status HloEvaluator::HandleDynamicSlice(const HloInstruction* dynamic_slice) {
+absl::Status HloEvaluator::HandleDynamicSlice(
+    const HloInstruction* dynamic_slice) {
   auto operand = dynamic_slice->operand(0);
   auto start_indices = dynamic_slice->operand(1);
   auto result_shape = dynamic_slice->shape();
@@ -3349,7 +3358,7 @@ Status HloEvaluator::HandleDynamicSlice(const HloInstruction* dynamic_slice) {
   return OkStatus();
 }
 
-Status HloEvaluator::HandleDynamicUpdateSlice(const HloInstruction* dus) {
+absl::Status HloEvaluator::HandleDynamicUpdateSlice(const HloInstruction* dus) {
   auto operand = dus->operand(0);
   auto update = dus->operand(1);
   auto start_indices = dus->operand(2);
@@ -3401,7 +3410,7 @@ Status HloEvaluator::HandleDynamicUpdateSlice(const HloInstruction* dus) {
   return OkStatus();
 }
 
-Status HloEvaluator::HandleSelect(const HloInstruction* select) {
+absl::Status HloEvaluator::HandleSelect(const HloInstruction* select) {
   const auto& pred = GetEvaluatedLiteralFor(select->operand(0));
   const auto& on_true = GetEvaluatedLiteralFor(select->operand(1));
   const auto& on_false = GetEvaluatedLiteralFor(select->operand(2));
@@ -3474,7 +3483,7 @@ absl::StatusOr<Literal> TryParseAndEvaluateWhileInductionVar(
 
 }  // namespace
 
-Status HloEvaluator::HandleWhile(const HloInstruction* while_hlo) {
+absl::Status HloEvaluator::HandleWhile(const HloInstruction* while_hlo) {
   const HloComputation* cond_comp = while_hlo->while_condition();
   const HloComputation* body_comp = while_hlo->while_body();
   // Initialize the loop carried valued with the input to the While instruction.
@@ -3734,7 +3743,7 @@ absl::StatusOr<Literal> StochasticConvertOp(const Literal& operand_literal,
 }
 }  // namespace
 
-Status HloEvaluator::HandleReverse(const HloInstruction* reverse) {
+absl::Status HloEvaluator::HandleReverse(const HloInstruction* reverse) {
   const Shape& result_shape = reverse->shape();
   const auto reverse_dimensions = reverse->dimensions();
 
@@ -3770,7 +3779,7 @@ Status HloEvaluator::HandleReverse(const HloInstruction* reverse) {
   return OkStatus();
 }
 
-Status HloEvaluator::HandleSelectAndScatter(
+absl::Status HloEvaluator::HandleSelectAndScatter(
     const HloInstruction* select_and_scatter) {
   auto operand = select_and_scatter->operand(0);
   auto source = select_and_scatter->operand(1);
@@ -3864,7 +3873,7 @@ Status HloEvaluator::HandleSelectAndScatter(
   return OkStatus();
 }
 
-Status HloEvaluator::HandleSlice(const HloInstruction* slice) {
+absl::Status HloEvaluator::HandleSlice(const HloInstruction* slice) {
   auto operand = slice->operand(0);
   const Shape& shape = slice->shape();
   TF_ASSIGN_OR_RETURN(auto inferred_return_shape,
@@ -3899,7 +3908,7 @@ Status HloEvaluator::HandleSlice(const HloInstruction* slice) {
   return OkStatus();
 }
 
-Status HloEvaluator::HandleSort(const HloInstruction* sort) {
+absl::Status HloEvaluator::HandleSort(const HloInstruction* sort) {
   TF_RET_CHECK(sort->operand_count() >= 1)
       << "Expected at least 1 operand for sort";
   for (int64_t i = 1; i < sort->operand_count(); ++i) {
@@ -3974,13 +3983,13 @@ Status HloEvaluator::HandleSort(const HloInstruction* sort) {
 #endif
     return a_is_smaller;
   };
-  std::function<Status(absl::Span<const Literal>, absl::Span<int64_t>,
-                       absl::Span<int64_t>, absl::Span<int64_t>,
-                       std::vector<int64_t>&, HloEvaluator*)>
+  std::function<absl::Status(absl::Span<const Literal>, absl::Span<int64_t>,
+                             absl::Span<int64_t>, absl::Span<int64_t>,
+                             std::vector<int64_t>&, HloEvaluator*)>
       merge = [&](absl::Span<const Literal> literals_to_sort,
                   absl::Span<int64_t> lhs, absl::Span<int64_t> rhs,
                   absl::Span<int64_t> output, std::vector<int64_t>& tmp,
-                  HloEvaluator* embedded_evaluator) -> Status {
+                  HloEvaluator* embedded_evaluator) -> absl::Status {
     tmp.clear();
     tmp.reserve(output.size());
     // Keep picking between elements.
@@ -4013,13 +4022,13 @@ Status HloEvaluator::HandleSort(const HloInstruction* sort) {
   const size_t work_per_thread = useful_parallelism > 1
                                      ? sort_dim_elements / useful_parallelism
                                      : std::numeric_limits<size_t>::max();
-  std::function<Status(absl::Span<const Literal>, absl::Span<int64_t>,
-                       std::vector<int64_t>*, HloEvaluator*)>
+  std::function<absl::Status(absl::Span<const Literal>, absl::Span<int64_t>,
+                             std::vector<int64_t>*, HloEvaluator*)>
       mergesort = [&merge, &mergesort, &less_than, this, env, work_per_thread](
                       absl::Span<const Literal> literals_to_sort,
                       absl::Span<int64_t> to_sort,
                       std::vector<int64_t>* scratch,
-                      HloEvaluator* embedded_evaluator) -> Status {
+                      HloEvaluator* embedded_evaluator) -> absl::Status {
     // Base case: inputs with 0 or 1 elements are already sorted.
     if (to_sort.size() < 2) {
       return OkStatus();
@@ -4046,7 +4055,7 @@ Status HloEvaluator::HandleSort(const HloInstruction* sort) {
       // Overlap sorting the LHS with the RHS if we have enough work to
       // do. The recursive call for to `mergesort(rhs)` will potentially
       // create more threads.
-      Status lhs_status;
+      absl::Status lhs_status;
       if (to_sort.size() >= work_per_thread) {
         std::unique_ptr<tsl::Thread> thread = absl::WrapUnique(env->StartThread(
             tsl::ThreadOptions(), "XLA_mergesort",
@@ -4157,7 +4166,7 @@ Status HloEvaluator::HandleSort(const HloInstruction* sort) {
   return OkStatus();
 }
 
-Status HloEvaluator::HandleStochasticConvert(
+absl::Status HloEvaluator::HandleStochasticConvert(
     const HloInstruction* stochastic_convert) {
   const HloInstruction* operand = stochastic_convert->operand(0);
   const HloInstruction* random = stochastic_convert->operand(1);
@@ -4312,7 +4321,7 @@ static absl::StatusOr<bool> GenerateReduceOutputElement(
   return true;
 }
 
-Status HloEvaluator::HandleReduce(const HloInstruction* hlo) {
+absl::Status HloEvaluator::HandleReduce(const HloInstruction* hlo) {
   const HloReduceInstruction* reduce = Cast<HloReduceInstruction>(hlo);
   int64_t num_args = reduce->inputs().size();
   absl::Span<const int64_t> dimensions_to_reduce(reduce->dimensions());
@@ -4413,7 +4422,7 @@ Status HloEvaluator::HandleReduce(const HloInstruction* hlo) {
   return OkStatus();
 }
 
-Status HloEvaluator::HandleReduceWindow(const HloInstruction* hlo) {
+absl::Status HloEvaluator::HandleReduceWindow(const HloInstruction* hlo) {
   auto* reduce_window = Cast<HloReduceWindowInstruction>(hlo);
   const Window& window = reduce_window->window();
   HloComputation* function = reduce_window->to_apply();
@@ -4554,7 +4563,7 @@ Status HloEvaluator::HandleReduceWindow(const HloInstruction* hlo) {
   return OkStatus();
 }
 
-Status HloEvaluator::HandleMap(const HloInstruction* map) {
+absl::Status HloEvaluator::HandleMap(const HloInstruction* map) {
   auto operands = map->operands();
   const HloComputation* computation = map->to_apply();
 
@@ -4586,7 +4595,7 @@ Status HloEvaluator::HandleMap(const HloInstruction* map) {
   return OkStatus();
 }
 
-Status HloEvaluator::HandleCustomCall(const HloInstruction* custom_call) {
+absl::Status HloEvaluator::HandleCustomCall(const HloInstruction* custom_call) {
   if (!custom_call_handler_) {
     // No handler is registered; this means custom-calls are not allowed.
     return DefaultAction(custom_call);
@@ -4607,7 +4616,7 @@ Status HloEvaluator::HandleCustomCall(const HloInstruction* custom_call) {
   return OkStatus();
 }
 
-Status HloEvaluator::Preprocess(const HloInstruction* hlo) {
+absl::Status HloEvaluator::Preprocess(const HloInstruction* hlo) {
   VLOG(3) << "About to visit HLO: " << hlo->ToString();
   if (!enable_partial_evaluation_) {
     for (const HloInstruction* operand : hlo->operands()) {
@@ -4622,7 +4631,7 @@ Status HloEvaluator::Preprocess(const HloInstruction* hlo) {
   return ShapeUtil::ValidateShape(hlo->shape());
 }
 
-Status HloEvaluator::Postprocess(const HloInstruction* hlo) {
+absl::Status HloEvaluator::Postprocess(const HloInstruction* hlo) {
   VLOG(3) << "Finished visiting " << hlo->ToString()
           << "; evaluated value is: " << GetEvaluatedLiteralFor(hlo).ToString();
   // Out of convenience the literal may have been produced with a different
diff --git a/third_party/xla/xla/hlo/evaluator/hlo_evaluator.h b/third_party/xla/xla/hlo/evaluator/hlo_evaluator.h
index ed6accfacf96e9..5d527c092aa8be 100644
--- a/third_party/xla/xla/hlo/evaluator/hlo_evaluator.h
+++ b/third_party/xla/xla/hlo/evaluator/hlo_evaluator.h
@@ -244,12 +244,12 @@ class HloEvaluator : public ConstDfsHloVisitorWithDefault {
   // EvaluateInternal. Such partial evaluation reduces the computation and
   // memory overhead in cases where we need only one tuple element by avoiding
   // the evaluation of a full tuple.
-  Status EvaluateInternal(
+  absl::Status EvaluateInternal(
       const HloInstruction* instruction, const ShapeIndex& shape_index = {},
       bool recursively_evaluate_nonconstant_operands = false);
 
-  Status EvaluateParameterFromCallerArgument(const HloInstruction* parameter,
-                                             const ShapeIndex& shape_index);
+  absl::Status EvaluateParameterFromCallerArgument(
+      const HloInstruction* parameter, const ShapeIndex& shape_index);
 
   // Helper method to extract a list of int64_t from evaluated instruction for
   // start_indices for DynamicSlice and DynamicUpdateSlice.
@@ -280,86 +280,88 @@ class HloEvaluator : public ConstDfsHloVisitorWithDefault {
 
   // Wraps around instruction handling to infer types before dispatching to
   // the corresponding typed Visitor.
-  Status DefaultAction(const HloInstruction* hlo) override {
+  absl::Status DefaultAction(const HloInstruction* hlo) override {
     return hlo->Visit(typed_visitors_[hlo->shape().element_type()].get());
   }
 
-  Status Preprocess(const HloInstruction* hlo) override;
-  Status Postprocess(const HloInstruction* hlo) override;
+  absl::Status Preprocess(const HloInstruction* hlo) override;
+  absl::Status Postprocess(const HloInstruction* hlo) override;
 
   // Operations that are type-agnostic or always return a specific type, such as
   // HandleIsFinite where boolean is always returned.
   //
-  Status HandleBitcast(const HloInstruction* bitcast) override;
-  Status HandleBitcastConvert(const HloInstruction* convert) override;
-  Status HandleGetDimensionSize(
+  absl::Status HandleBitcast(const HloInstruction* bitcast) override;
+  absl::Status HandleBitcastConvert(const HloInstruction* convert) override;
+  absl::Status HandleGetDimensionSize(
       const HloInstruction* get_dimension_size) override;
-  Status HandleSetDimensionSize(
+  absl::Status HandleSetDimensionSize(
       const HloInstruction* set_dimension_size) override;
-  Status HandleParameter(const HloInstruction* parameter) override;
-  Status HandleInfeed(const HloInstruction* infeed) override;
-  Status HandleConstant(const HloInstruction* constant) override;
-  Status HandleConcatenate(const HloInstruction* concatenate) override;
-  Status HandleReshape(const HloInstruction* reshape) override;
-  Status HandleTranspose(const HloInstruction* transpose) override;
-  Status HandleIsFinite(const HloInstruction* is_finite) override;
-  Status HandleCompare(const HloInstruction* compare) override;
-  Status HandleTuple(const HloInstruction* tuple) override;
-  Status HandleFft(const HloInstruction* fft) override;
-  Status HandleGather(const HloInstruction* gather) override;
-  Status HandleScatter(const HloInstruction* hlo) override;
-  Status HandleGetTupleElement(
+  absl::Status HandleParameter(const HloInstruction* parameter) override;
+  absl::Status HandleInfeed(const HloInstruction* infeed) override;
+  absl::Status HandleConstant(const HloInstruction* constant) override;
+  absl::Status HandleConcatenate(const HloInstruction* concatenate) override;
+  absl::Status HandleReshape(const HloInstruction* reshape) override;
+  absl::Status HandleTranspose(const HloInstruction* transpose) override;
+  absl::Status HandleIsFinite(const HloInstruction* is_finite) override;
+  absl::Status HandleCompare(const HloInstruction* compare) override;
+  absl::Status HandleTuple(const HloInstruction* tuple) override;
+  absl::Status HandleFft(const HloInstruction* fft) override;
+  absl::Status HandleGather(const HloInstruction* gather) override;
+  absl::Status HandleScatter(const HloInstruction* hlo) override;
+  absl::Status HandleGetTupleElement(
       const HloInstruction* get_tuple_element) override;
-  Status HandleAsyncStart(const HloInstruction* async_start) override;
-  Status HandleAsyncUpdate(const HloInstruction* async_update) override;
-  Status HandleAsyncDone(const HloInstruction* async_done) override;
-  Status HandleCopy(const HloInstruction* copy) override;
-  Status HandleCopyStart(const HloInstruction* copy_start) override;
-  Status HandleCopyDone(const HloInstruction* copy_done) override;
-  Status HandleConditional(const HloInstruction* conditional) override;
-  Status HandleConvert(const HloInstruction* convert) override;
-  Status HandleCall(const HloInstruction* call) override;
-  Status HandleDynamicSlice(const HloInstruction* dynamic_slice) override;
-  Status HandleDynamicUpdateSlice(const HloInstruction* dus) override;
-  Status HandleFusion(const HloInstruction* fusion) override;
-  Status HandleWhile(const HloInstruction* while_hlo) override;
-  Status HandleSelect(const HloInstruction* select) override;
-  Status HandleBroadcast(const HloInstruction* broadcast) override;
-  Status HandleAfterAll(const HloInstruction* after_all) override;
-  Status HandleAddDependency(const HloInstruction* add_dependency) override;
-  Status HandleReverse(const HloInstruction* reverse) override;
-  Status HandleSelectAndScatter(
+  absl::Status HandleAsyncStart(const HloInstruction* async_start) override;
+  absl::Status HandleAsyncUpdate(const HloInstruction* async_update) override;
+  absl::Status HandleAsyncDone(const HloInstruction* async_done) override;
+  absl::Status HandleCopy(const HloInstruction* copy) override;
+  absl::Status HandleCopyStart(const HloInstruction* copy_start) override;
+  absl::Status HandleCopyDone(const HloInstruction* copy_done) override;
+  absl::Status HandleConditional(const HloInstruction* conditional) override;
+  absl::Status HandleConvert(const HloInstruction* convert) override;
+  absl::Status HandleCall(const HloInstruction* call) override;
+  absl::Status HandleDynamicSlice(const HloInstruction* dynamic_slice) override;
+  absl::Status HandleDynamicUpdateSlice(const HloInstruction* dus) override;
+  absl::Status HandleFusion(const HloInstruction* fusion) override;
+  absl::Status HandleWhile(const HloInstruction* while_hlo) override;
+  absl::Status HandleSelect(const HloInstruction* select) override;
+  absl::Status HandleBroadcast(const HloInstruction* broadcast) override;
+  absl::Status HandleAfterAll(const HloInstruction* after_all) override;
+  absl::Status HandleAddDependency(
+      const HloInstruction* add_dependency) override;
+  absl::Status HandleReverse(const HloInstruction* reverse) override;
+  absl::Status HandleSelectAndScatter(
       const HloInstruction* select_and_scatter) override;
-  Status HandleSlice(const HloInstruction* slice) override;
-  Status HandleSort(const HloInstruction* sort) override;
-  Status HandleStochasticConvert(
+  absl::Status HandleSlice(const HloInstruction* slice) override;
+  absl::Status HandleSort(const HloInstruction* sort) override;
+  absl::Status HandleStochasticConvert(
       const HloInstruction* stochastic_convert) override;
-  Status HandleReal(const HloInstruction* real) override;
-  Status HandleImag(const HloInstruction* imag) override;
-  Status HandleComplex(const HloInstruction* complex) override;
-  Status HandleReduce(const HloInstruction* hlo) override;
-  Status HandleReduceWindow(const HloInstruction* hlo) override;
-  Status HandleMap(const HloInstruction* map) override;
-  Status HandleCustomCall(const HloInstruction* custom_call) override;
+  absl::Status HandleReal(const HloInstruction* real) override;
+  absl::Status HandleImag(const HloInstruction* imag) override;
+  absl::Status HandleComplex(const HloInstruction* complex) override;
+  absl::Status HandleReduce(const HloInstruction* hlo) override;
+  absl::Status HandleReduceWindow(const HloInstruction* hlo) override;
+  absl::Status HandleMap(const HloInstruction* map) override;
+  absl::Status HandleCustomCall(const HloInstruction* custom_call) override;
 
   // Unsupported HLOs, note some of them (such as BatchNorm*) are typically
   // expanded in a semantic-preserving way into other HLOs by adding expansion
   // HLO pass to the HLO optimization pass during compilation, which can then be
   // handled by the evaluator.
-  Status HandleBatchNormGrad(const HloInstruction* batch_norm_grad) override {
+  absl::Status HandleBatchNormGrad(
+      const HloInstruction* batch_norm_grad) override {
     return Unimplemented("BatchNormGrad HLO is unsupported by the evaluator.");
   }
-  Status HandleBatchNormInference(
+  absl::Status HandleBatchNormInference(
       const HloInstruction* batch_norm_inference) override {
     return Unimplemented(
         "BatchNormInference HLO is unsupported by the evaluator.");
   }
-  Status HandleBatchNormTraining(
+  absl::Status HandleBatchNormTraining(
       const HloInstruction* batch_norm_training) override {
     return Unimplemented(
         "BatchNormTraining HLO is unsupported by the evaluator.");
   }
-  Status HandleOutfeed(const HloInstruction* outfeed) override {
+  absl::Status HandleOutfeed(const HloInstruction* outfeed) override {
     return Unimplemented("Outfeed HLO is unsupported by the evaluator.");
   }
 
diff --git a/third_party/xla/xla/hlo/evaluator/hlo_evaluator_typed_visitor.h b/third_party/xla/xla/hlo/evaluator/hlo_evaluator_typed_visitor.h
index 1ebf786af7bcb9..9d11bbc5d0f906 100644
--- a/third_party/xla/xla/hlo/evaluator/hlo_evaluator_typed_visitor.h
+++ b/third_party/xla/xla/hlo/evaluator/hlo_evaluator_typed_visitor.h
@@ -119,8 +119,8 @@ auto ToArithmeticSafeType(T t) {
 template <typename ReturnT, typename ElementwiseT = ReturnT>
 class HloEvaluatorTypedVisitor : public ConstDfsHloVisitorWithDefault {
  private:
-  ABSL_ATTRIBUTE_NOINLINE Status
-  UnsupportedTypeError(const HloInstruction* instruction) {
+  ABSL_ATTRIBUTE_NOINLINE absl::Status UnsupportedTypeError(
+      const HloInstruction* instruction) {
     return InvalidArgument(
         "Unsupported type for %s: %s", HloOpcodeString(instruction->opcode()),
         PrimitiveType_Name(instruction->shape().element_type()));
@@ -155,14 +155,14 @@ class HloEvaluatorTypedVisitor : public ConstDfsHloVisitorWithDefault {
     };
   }
 
-  Status DefaultAction(const HloInstruction* hlo_instruction) override {
+  absl::Status DefaultAction(const HloInstruction* hlo_instruction) override {
     return Unimplemented("unhandled HLO ops for HloEvaluator: %s.",
                          HloOpcodeString(hlo_instruction->opcode()));
   }
 
   template <typename NativeT,
             typename std::enable_if_t<std::is_unsigned_v<NativeT>>* = nullptr>
-  Status HandleAbs(const HloInstruction* abs) {
+  absl::Status HandleAbs(const HloInstruction* abs) {
     TF_ASSIGN_OR_RETURN(parent_->evaluated_[abs],
                         ElementWiseUnaryOp(abs, [](NativeT elem_operand) {
                           return elem_operand;
@@ -172,7 +172,7 @@ class HloEvaluatorTypedVisitor : public ConstDfsHloVisitorWithDefault {
 
   template <typename NativeT,
             typename std::enable_if_t<std::is_signed_v<NativeT>>* = nullptr>
-  Status HandleAbs(const HloInstruction* abs) {
+  absl::Status HandleAbs(const HloInstruction* abs) {
     TF_ASSIGN_OR_RETURN(parent_->evaluated_[abs],
                         ElementWiseUnaryOp(abs, [](NativeT elem_operand) {
                           return std::abs(elem_operand);
@@ -182,7 +182,7 @@ class HloEvaluatorTypedVisitor : public ConstDfsHloVisitorWithDefault {
 
   template <typename NativeT,
             typename std::enable_if_t<is_complex_v<NativeT>>* = nullptr>
-  Status HandleAbs(const HloInstruction* abs) {
+  absl::Status HandleAbs(const HloInstruction* abs) {
     const Literal& operand_literal =
         parent_->GetEvaluatedLiteralFor(abs->operand(0));
     TF_ASSIGN_OR_RETURN(
@@ -195,7 +195,7 @@ class HloEvaluatorTypedVisitor : public ConstDfsHloVisitorWithDefault {
     return OkStatus();
   }
 
-  Status HandleAbs(const HloInstruction* abs) override {
+  absl::Status HandleAbs(const HloInstruction* abs) override {
     // If the operand is of C64 type, the return type of abs will be F32.
     // However, ElementwiseT would still be the return type, F32, and thus
     // specifying the ElementwiseT explicitly as C64 is needed below.
@@ -207,7 +207,7 @@ class HloEvaluatorTypedVisitor : public ConstDfsHloVisitorWithDefault {
     return HandleAbs<ElementwiseT>(abs);
   }
 
-  Status HandleRound(const HloInstruction* round) override {
+  absl::Status HandleRound(const HloInstruction* round) override {
     if constexpr (!is_complex_v<ReturnT>) {
       TF_ASSIGN_OR_RETURN(
           parent_->evaluated_[round],
@@ -219,7 +219,7 @@ class HloEvaluatorTypedVisitor : public ConstDfsHloVisitorWithDefault {
     return UnsupportedTypeError(round);
   }
 
-  Status HandleRoundNearestEven(const HloInstruction* round) override {
+  absl::Status HandleRoundNearestEven(const HloInstruction* round) override {
     if constexpr (!is_complex_v<ReturnT>) {
       // Verify the current rounding direction.
       TF_RET_CHECK(fegetround() == FE_TONEAREST);
@@ -233,7 +233,7 @@ class HloEvaluatorTypedVisitor : public ConstDfsHloVisitorWithDefault {
     return UnsupportedTypeError(round);
   }
 
-  Status HandleCeil(const HloInstruction* ceil) override {
+  absl::Status HandleCeil(const HloInstruction* ceil) override {
     if constexpr (!is_complex_v<ReturnT>) {
       TF_ASSIGN_OR_RETURN(
           parent_->evaluated_[ceil],
@@ -245,7 +245,7 @@ class HloEvaluatorTypedVisitor : public ConstDfsHloVisitorWithDefault {
     return UnsupportedTypeError(ceil);
   }
 
-  Status HandleErf(const HloInstruction* erf) override {
+  absl::Status HandleErf(const HloInstruction* erf) override {
     if constexpr (!is_complex_v<ReturnT>) {
       TF_ASSIGN_OR_RETURN(
           parent_->evaluated_[erf],
@@ -257,7 +257,7 @@ class HloEvaluatorTypedVisitor : public ConstDfsHloVisitorWithDefault {
     return UnsupportedTypeError(erf);
   }
 
-  Status HandleExp(const HloInstruction* exp) override {
+  absl::Status HandleExp(const HloInstruction* exp) override {
     TF_ASSIGN_OR_RETURN(parent_->evaluated_[exp],
                         ElementWiseUnaryOp(exp, [](ElementwiseT elem_operand) {
                           return std::exp(elem_operand);
@@ -265,7 +265,7 @@ class HloEvaluatorTypedVisitor : public ConstDfsHloVisitorWithDefault {
     return OkStatus();
   }
 
-  Status HandleExpm1(const HloInstruction* expm1) override {
+  absl::Status HandleExpm1(const HloInstruction* expm1) override {
     if constexpr (!is_complex_v<ReturnT>) {
       TF_ASSIGN_OR_RETURN(
           parent_->evaluated_[expm1],
@@ -277,7 +277,7 @@ class HloEvaluatorTypedVisitor : public ConstDfsHloVisitorWithDefault {
     return UnsupportedTypeError(expm1);
   }
 
-  Status HandleFloor(const HloInstruction* floor) override {
+  absl::Status HandleFloor(const HloInstruction* floor) override {
     if constexpr (!is_complex_v<ReturnT>) {
       TF_ASSIGN_OR_RETURN(
           parent_->evaluated_[floor],
@@ -289,7 +289,7 @@ class HloEvaluatorTypedVisitor : public ConstDfsHloVisitorWithDefault {
     return UnsupportedTypeError(floor);
   }
 
-  Status HandleLog(const HloInstruction* log) override {
+  absl::Status HandleLog(const HloInstruction* log) override {
     TF_ASSIGN_OR_RETURN(parent_->evaluated_[log],
                         ElementWiseUnaryOp(log, [](ElementwiseT elem_operand) {
                           return std::log(elem_operand);
@@ -297,7 +297,7 @@ class HloEvaluatorTypedVisitor : public ConstDfsHloVisitorWithDefault {
     return OkStatus();
   }
 
-  Status HandleLog1p(const HloInstruction* log1p) override {
+  absl::Status HandleLog1p(const HloInstruction* log1p) override {
     if constexpr (!is_complex_v<ReturnT>) {
       TF_ASSIGN_OR_RETURN(
           parent_->evaluated_[log1p],
@@ -309,7 +309,7 @@ class HloEvaluatorTypedVisitor : public ConstDfsHloVisitorWithDefault {
     return UnsupportedTypeError(log1p);
   }
 
-  Status HandleNot(const HloInstruction* not_) override {
+  absl::Status HandleNot(const HloInstruction* not_) override {
     if constexpr (std::is_arithmetic_v<ElementwiseT>) {
       TF_ASSIGN_OR_RETURN(
           parent_->evaluated_[not_],
@@ -331,7 +331,7 @@ class HloEvaluatorTypedVisitor : public ConstDfsHloVisitorWithDefault {
       typename NativeT,
       typename std::enable_if_t<std::is_signed_v<NativeT> &&
                                 !std::is_floating_point_v<NativeT>>* = nullptr>
-  Status HandleNegate(const HloInstruction* negate) {
+  absl::Status HandleNegate(const HloInstruction* negate) {
     using type = std::make_unsigned_t<NativeT>;
     TF_ASSIGN_OR_RETURN(
         parent_->evaluated_[negate],
@@ -344,7 +344,7 @@ class HloEvaluatorTypedVisitor : public ConstDfsHloVisitorWithDefault {
   template <typename NativeT, typename std::enable_if_t<
                                   !std::is_signed_v<NativeT> ||
                                   std::is_floating_point_v<NativeT>>* = nullptr>
-  Status HandleNegate(const HloInstruction* negate) {
+  absl::Status HandleNegate(const HloInstruction* negate) {
     TF_ASSIGN_OR_RETURN(
         parent_->evaluated_[negate],
         ElementWiseUnaryOp(
@@ -352,11 +352,11 @@ class HloEvaluatorTypedVisitor : public ConstDfsHloVisitorWithDefault {
     return OkStatus();
   }
 
-  Status HandleNegate(const HloInstruction* negate) override {
+  absl::Status HandleNegate(const HloInstruction* negate) override {
     return HandleNegate<ReturnT>(negate);
   }
 
-  Status HandleLogistic(const HloInstruction* logistic) override {
+  absl::Status HandleLogistic(const HloInstruction* logistic) override {
     TF_ASSIGN_OR_RETURN(
         parent_->evaluated_[logistic],
         ElementWiseUnaryOp(logistic, [](ElementwiseT elem_operand) {
@@ -366,7 +366,7 @@ class HloEvaluatorTypedVisitor : public ConstDfsHloVisitorWithDefault {
     return OkStatus();
   }
 
-  Status HandleSign(const HloInstruction* sign) override {
+  absl::Status HandleSign(const HloInstruction* sign) override {
     using NativeT = ElementwiseT;
     TF_ASSIGN_OR_RETURN(
         parent_->evaluated_[sign],
@@ -389,7 +389,7 @@ class HloEvaluatorTypedVisitor : public ConstDfsHloVisitorWithDefault {
     return OkStatus();
   }
 
-  Status HandleAtan2(const HloInstruction* atan2) override {
+  absl::Status HandleAtan2(const HloInstruction* atan2) override {
     if constexpr (std::is_floating_point_v<ElementwiseT>) {
       TF_ASSIGN_OR_RETURN(parent_->evaluated_[atan2],
                           ElementWiseBinaryOp(atan2, [](ElementwiseT lhs_elem,
@@ -411,7 +411,7 @@ class HloEvaluatorTypedVisitor : public ConstDfsHloVisitorWithDefault {
     return UnsupportedTypeError(atan2);
   }
 
-  Status HandleTanh(const HloInstruction* tanh) override {
+  absl::Status HandleTanh(const HloInstruction* tanh) override {
     TF_ASSIGN_OR_RETURN(parent_->evaluated_[tanh],
                         ElementWiseUnaryOp(tanh, [](ElementwiseT elem_operand) {
                           return std::tanh(elem_operand);
@@ -419,7 +419,7 @@ class HloEvaluatorTypedVisitor : public ConstDfsHloVisitorWithDefault {
     return OkStatus();
   }
 
-  Status HandleMultiply(const HloInstruction* multiply) override {
+  absl::Status HandleMultiply(const HloInstruction* multiply) override {
     TF_ASSIGN_OR_RETURN(
         parent_->evaluated_[multiply],
         ElementWiseBinaryOp(
@@ -430,7 +430,7 @@ class HloEvaluatorTypedVisitor : public ConstDfsHloVisitorWithDefault {
     return OkStatus();
   }
 
-  Status HandleSubtract(const HloInstruction* subtract) override {
+  absl::Status HandleSubtract(const HloInstruction* subtract) override {
     TF_ASSIGN_OR_RETURN(
         parent_->evaluated_[subtract],
         ElementWiseBinaryOp(
@@ -441,7 +441,7 @@ class HloEvaluatorTypedVisitor : public ConstDfsHloVisitorWithDefault {
     return OkStatus();
   }
 
-  Status HandleAdd(const HloInstruction* add) override {
+  absl::Status HandleAdd(const HloInstruction* add) override {
     TF_ASSIGN_OR_RETURN(parent_->evaluated_[add],
                         ElementWiseBinaryOp(add, [](ElementwiseT lhs_elem,
                                                     ElementwiseT rhs_elem) {
@@ -451,7 +451,7 @@ class HloEvaluatorTypedVisitor : public ConstDfsHloVisitorWithDefault {
     return OkStatus();
   }
 
-  Status HandleDivide(const HloInstruction* divide) override {
+  absl::Status HandleDivide(const HloInstruction* divide) override {
     TF_ASSIGN_OR_RETURN(
         parent_->evaluated_[divide],
         ElementWiseBinaryOp(
@@ -478,7 +478,7 @@ class HloEvaluatorTypedVisitor : public ConstDfsHloVisitorWithDefault {
     return OkStatus();
   }
 
-  Status HandleMaximum(const HloInstruction* maximum) override {
+  absl::Status HandleMaximum(const HloInstruction* maximum) override {
     if constexpr (!is_complex_v<ElementwiseT>) {
       TF_ASSIGN_OR_RETURN(
           parent_->evaluated_[maximum],
@@ -498,7 +498,7 @@ class HloEvaluatorTypedVisitor : public ConstDfsHloVisitorWithDefault {
     return UnsupportedTypeError(maximum);
   }
 
-  Status HandleMinimum(const HloInstruction* minimum) override {
+  absl::Status HandleMinimum(const HloInstruction* minimum) override {
     if constexpr (!is_complex_v<ElementwiseT>) {
       TF_ASSIGN_OR_RETURN(
           parent_->evaluated_[minimum],
@@ -518,7 +518,7 @@ class HloEvaluatorTypedVisitor : public ConstDfsHloVisitorWithDefault {
     return UnsupportedTypeError(minimum);
   }
 
-  Status HandlePower(const HloInstruction* power) override {
+  absl::Status HandlePower(const HloInstruction* power) override {
     TF_ASSIGN_OR_RETURN(
         parent_->evaluated_[power],
         ElementWiseBinaryOp(power, [](ElementwiseT lhs_el,
@@ -569,7 +569,7 @@ class HloEvaluatorTypedVisitor : public ConstDfsHloVisitorWithDefault {
     return OkStatus();
   }
 
-  Status HandleSqrt(const HloInstruction* sqrt) override {
+  absl::Status HandleSqrt(const HloInstruction* sqrt) override {
     TF_ASSIGN_OR_RETURN(parent_->evaluated_[sqrt],
                         ElementWiseUnaryOp(sqrt, [](ElementwiseT elem_operand) {
                           return std::sqrt(elem_operand);
@@ -577,7 +577,7 @@ class HloEvaluatorTypedVisitor : public ConstDfsHloVisitorWithDefault {
     return OkStatus();
   }
 
-  Status HandleCbrt(const HloInstruction* cbrt) override {
+  absl::Status HandleCbrt(const HloInstruction* cbrt) override {
     if constexpr (!is_complex_v<ElementwiseT>) {
       TF_ASSIGN_OR_RETURN(
           parent_->evaluated_[cbrt],
@@ -589,7 +589,7 @@ class HloEvaluatorTypedVisitor : public ConstDfsHloVisitorWithDefault {
     return UnsupportedTypeError(cbrt);
   }
 
-  Status HandleRsqrt(const HloInstruction* rsqrt) override {
+  absl::Status HandleRsqrt(const HloInstruction* rsqrt) override {
     TF_ASSIGN_OR_RETURN(
         parent_->evaluated_[rsqrt],
         ElementWiseUnaryOp(rsqrt, [](ElementwiseT elem_operand) {
@@ -598,7 +598,7 @@ class HloEvaluatorTypedVisitor : public ConstDfsHloVisitorWithDefault {
     return OkStatus();
   }
 
-  Status HandleRemainder(const HloInstruction* remainder) override {
+  absl::Status HandleRemainder(const HloInstruction* remainder) override {
     if constexpr (!is_complex_v<ElementwiseT>) {
       TF_ASSIGN_OR_RETURN(
           parent_->evaluated_[remainder],
@@ -626,7 +626,7 @@ class HloEvaluatorTypedVisitor : public ConstDfsHloVisitorWithDefault {
     return UnsupportedTypeError(remainder);
   }
 
-  Status HandleAnd(const HloInstruction* and_inst) override {
+  absl::Status HandleAnd(const HloInstruction* and_inst) override {
     if constexpr (std::is_integral_v<ElementwiseT>) {
       TF_ASSIGN_OR_RETURN(
           parent_->evaluated_[and_inst],
@@ -639,7 +639,7 @@ class HloEvaluatorTypedVisitor : public ConstDfsHloVisitorWithDefault {
     return UnsupportedTypeError(and_inst);
   }
 
-  Status HandleOr(const HloInstruction* or_inst) override {
+  absl::Status HandleOr(const HloInstruction* or_inst) override {
     if constexpr (std::is_integral_v<ElementwiseT>) {
       TF_ASSIGN_OR_RETURN(parent_->evaluated_[or_inst],
                           ElementWiseBinaryOp(or_inst, [](ElementwiseT lhs_el,
@@ -651,7 +651,7 @@ class HloEvaluatorTypedVisitor : public ConstDfsHloVisitorWithDefault {
     return UnsupportedTypeError(or_inst);
   }
 
-  Status HandleXor(const HloInstruction* xor_inst) override {
+  absl::Status HandleXor(const HloInstruction* xor_inst) override {
     if constexpr (std::is_integral_v<ElementwiseT>) {
       TF_ASSIGN_OR_RETURN(
           parent_->evaluated_[xor_inst],
@@ -664,7 +664,7 @@ class HloEvaluatorTypedVisitor : public ConstDfsHloVisitorWithDefault {
     return UnsupportedTypeError(xor_inst);
   }
 
-  Status HandleShiftLeft(const HloInstruction* shl) override {
+  absl::Status HandleShiftLeft(const HloInstruction* shl) override {
     if constexpr (std::is_integral_v<ElementwiseT> &&
                   !std::is_same_v<ElementwiseT, bool>) {
       TF_ASSIGN_OR_RETURN(parent_->evaluated_[shl],
@@ -679,7 +679,7 @@ class HloEvaluatorTypedVisitor : public ConstDfsHloVisitorWithDefault {
     return UnsupportedTypeError(shl);
   }
 
-  Status HandleShiftRightArithmetic(const HloInstruction* shr) override {
+  absl::Status HandleShiftRightArithmetic(const HloInstruction* shr) override {
     if constexpr (std::is_integral_v<ElementwiseT> &&
                   !std::is_same_v<ElementwiseT, bool>) {
       using SignedT = make_specialized_signed_t<ReturnT>;
@@ -699,7 +699,7 @@ class HloEvaluatorTypedVisitor : public ConstDfsHloVisitorWithDefault {
     return UnsupportedTypeError(shr);
   }
 
-  Status HandleShiftRightLogical(const HloInstruction* shr) override {
+  absl::Status HandleShiftRightLogical(const HloInstruction* shr) override {
     if constexpr (std::is_integral_v<ElementwiseT> &&
                   !std::is_same_v<ElementwiseT, bool>) {
       using UnsignedT = make_specialized_unsigned_t<ReturnT>;
@@ -719,7 +719,7 @@ class HloEvaluatorTypedVisitor : public ConstDfsHloVisitorWithDefault {
     return UnsupportedTypeError(shr);
   }
 
-  Status HandleClamp(const HloInstruction* clamp) override {
+  absl::Status HandleClamp(const HloInstruction* clamp) override {
     if constexpr (!is_complex_v<ElementwiseT>) {
       auto clamp_op = [](ElementwiseT low, ElementwiseT value,
                          ElementwiseT high) {
@@ -745,7 +745,7 @@ class HloEvaluatorTypedVisitor : public ConstDfsHloVisitorWithDefault {
     return UnsupportedTypeError(clamp);
   }
 
-  Status HandleSelect(const HloInstruction* select) override {
+  absl::Status HandleSelect(const HloInstruction* select) override {
     CHECK(!ShapeUtil::IsScalar(select->operand(0)->shape()));
     CHECK(select->shape().IsArray());
     std::function<ReturnT(bool, ReturnT, ReturnT)> select_op =
@@ -760,9 +760,9 @@ class HloEvaluatorTypedVisitor : public ConstDfsHloVisitorWithDefault {
     return OkStatus();
   }
 
-  Status HandleConvolutionWithLiterals(const HloInstruction* conv,
-                                       const Literal& lhs_literal,
-                                       const Literal& rhs_literal) {
+  absl::Status HandleConvolutionWithLiterals(const HloInstruction* conv,
+                                             const Literal& lhs_literal,
+                                             const Literal& rhs_literal) {
     const auto& window = conv->window();
     const Shape& result_shape = conv->shape();
     const Shape& lhs_shape = lhs_literal.shape();
@@ -961,7 +961,7 @@ class HloEvaluatorTypedVisitor : public ConstDfsHloVisitorWithDefault {
     return OkStatus();
   }
 
-  Status HandleConvolution(const HloInstruction* conv) override {
+  absl::Status HandleConvolution(const HloInstruction* conv) override {
     auto lhs = conv->operand(0);
     auto rhs = conv->operand(1);
     const auto& window = conv->window();
@@ -1020,7 +1020,7 @@ class HloEvaluatorTypedVisitor : public ConstDfsHloVisitorWithDefault {
         rhs_literal.Convert(result_shape.element_type()).value());
   }
 
-  Status HandleDot(const HloInstruction* dot) override {
+  absl::Status HandleDot(const HloInstruction* dot) override {
     if (dot->dot_dimension_numbers().rhs_contracting_dimensions_size() == 1 &&
         parent_->use_fast_path_ &&
         ShapeUtil::SameElementType(dot->operand(0)->shape(), dot->shape()) &&
@@ -1032,7 +1032,7 @@ class HloEvaluatorTypedVisitor : public ConstDfsHloVisitorWithDefault {
 
   template <typename NativeT, typename std::enable_if_t<
                                   std::is_same_v<NativeT, float>>* = nullptr>
-  Status HandleDot(const HloInstruction* dot) {
+  absl::Status HandleDot(const HloInstruction* dot) {
     const HloInstruction* lhs = dot->operand(0);
     const HloInstruction* rhs = dot->operand(1);
     CHECK(dot->shape().IsArray());
@@ -1097,13 +1097,13 @@ class HloEvaluatorTypedVisitor : public ConstDfsHloVisitorWithDefault {
 
   template <typename NativeT, typename std::enable_if_t<
                                   !std::is_same_v<NativeT, float>>* = nullptr>
-  Status HandleDot(const HloInstruction* dot) {
+  absl::Status HandleDot(const HloInstruction* dot) {
     return HandleDotSlowPath(dot);
   }
 
-  Status HandleDotSlowPathWithLiterals(const HloInstruction* dot,
-                                       const Literal& lhs_literal,
-                                       const Literal& rhs_literal) {
+  absl::Status HandleDotSlowPathWithLiterals(const HloInstruction* dot,
+                                             const Literal& lhs_literal,
+                                             const Literal& rhs_literal) {
     const auto& dnums = dot->dot_dimension_numbers();
 
     const auto lhs_rank = lhs_literal.shape().rank();
@@ -1212,7 +1212,7 @@ class HloEvaluatorTypedVisitor : public ConstDfsHloVisitorWithDefault {
     return OkStatus();
   }
 
-  Status HandleDotSlowPath(const HloInstruction* dot) {
+  absl::Status HandleDotSlowPath(const HloInstruction* dot) {
     auto lhs = dot->operand(0);
     auto rhs = dot->operand(1);
     CHECK(dot->shape().IsArray());
@@ -1242,7 +1242,7 @@ class HloEvaluatorTypedVisitor : public ConstDfsHloVisitorWithDefault {
         rhs_literal.Convert(dot->shape().element_type()).value());
   }
 
-  Status HandlePad(const HloInstruction* pad) override {
+  absl::Status HandlePad(const HloInstruction* pad) override {
     CHECK(pad->operand(0)->shape().IsArray());
     // Padding value must be scalar.
     CHECK(ShapeUtil::IsScalar(pad->operand(1)->shape()));
@@ -1310,7 +1310,7 @@ class HloEvaluatorTypedVisitor : public ConstDfsHloVisitorWithDefault {
     return OkStatus();
   }
 
-  Status HandleClz(const HloInstruction* clz) override {
+  absl::Status HandleClz(const HloInstruction* clz) override {
     // Enable CLZ only for integer types.
     if constexpr (std::is_integral_v<ElementwiseT> &&
                   !std::is_same_v<ElementwiseT, bool>) {
@@ -1326,7 +1326,7 @@ class HloEvaluatorTypedVisitor : public ConstDfsHloVisitorWithDefault {
     return UnsupportedTypeError(clz);
   }
 
-  Status HandlePopulationCount(const HloInstruction* popcnt) override {
+  absl::Status HandlePopulationCount(const HloInstruction* popcnt) override {
     if constexpr (std::is_integral_v<ElementwiseT> &&
                   !std::is_same_v<ElementwiseT, bool>) {
       TF_ASSIGN_OR_RETURN(
@@ -1340,7 +1340,7 @@ class HloEvaluatorTypedVisitor : public ConstDfsHloVisitorWithDefault {
     return UnsupportedTypeError(popcnt);
   }
 
-  Status HandleSin(const HloInstruction* sin) override {
+  absl::Status HandleSin(const HloInstruction* sin) override {
     if constexpr (std::is_floating_point_v<ElementwiseT> ||
                   is_complex_v<ElementwiseT>) {
       TF_ASSIGN_OR_RETURN(
@@ -1353,7 +1353,7 @@ class HloEvaluatorTypedVisitor : public ConstDfsHloVisitorWithDefault {
     return UnsupportedTypeError(sin);
   }
 
-  Status HandleCos(const HloInstruction* cos) override {
+  absl::Status HandleCos(const HloInstruction* cos) override {
     if constexpr (std::is_floating_point_v<ElementwiseT> ||
                   is_complex_v<ElementwiseT>) {
       TF_ASSIGN_OR_RETURN(
@@ -1366,7 +1366,7 @@ class HloEvaluatorTypedVisitor : public ConstDfsHloVisitorWithDefault {
     return UnsupportedTypeError(cos);
   }
 
-  Status HandleTan(const HloInstruction* tan) override {
+  absl::Status HandleTan(const HloInstruction* tan) override {
     if constexpr (std::is_floating_point_v<ElementwiseT>) {
       TF_ASSIGN_OR_RETURN(
           parent_->evaluated_[tan],
@@ -1380,7 +1380,7 @@ class HloEvaluatorTypedVisitor : public ConstDfsHloVisitorWithDefault {
 
   template <typename NativeT, typename std::enable_if_t<
                                   std::is_floating_point_v<NativeT>>* = nullptr>
-  Status HandleReducePrecision(const HloInstruction* reduce_precision) {
+  absl::Status HandleReducePrecision(const HloInstruction* reduce_precision) {
     TF_ASSIGN_OR_RETURN(
         parent_->evaluated_[reduce_precision],
         ElementWiseUnaryOp(reduce_precision, [&](ElementwiseT elem) {
@@ -1480,16 +1480,16 @@ class HloEvaluatorTypedVisitor : public ConstDfsHloVisitorWithDefault {
   template <typename NativeT,
             typename std::enable_if_t<std::is_integral_v<NativeT> ||
                                       is_complex_v<NativeT>>* = nullptr>
-  Status HandleReducePrecision(const HloInstruction* reduce_precision) {
+  absl::Status HandleReducePrecision(const HloInstruction* reduce_precision) {
     return UnsupportedTypeError(reduce_precision);
   }
 
-  Status HandleReducePrecision(
+  absl::Status HandleReducePrecision(
       const HloInstruction* reduce_precision) override {
     return HandleReducePrecision<ElementwiseT>(reduce_precision);
   }
 
-  Status HandleIota(const HloInstruction* instruction) override {
+  absl::Status HandleIota(const HloInstruction* instruction) override {
     auto* iota = Cast<HloIotaInstruction>(instruction);
     if constexpr (std::is_integral_v<ElementwiseT> ||
                   is_complex_v<ElementwiseT> ||
@@ -1506,7 +1506,7 @@ class HloEvaluatorTypedVisitor : public ConstDfsHloVisitorWithDefault {
     return UnsupportedTypeError(iota);
   }
 
-  Status HandleRng(const HloInstruction* random) override {
+  absl::Status HandleRng(const HloInstruction* random) override {
     RandomDistribution distribution = random->random_distribution();
     const Shape& result_shape = random->shape();
     Literal result(result_shape);
diff --git a/third_party/xla/xla/stream_executor/cuda/cuda_dnn.cc b/third_party/xla/xla/stream_executor/cuda/cuda_dnn.cc
index 07c53b1abbefd1..9c33472189148d 100644
--- a/third_party/xla/xla/stream_executor/cuda/cuda_dnn.cc
+++ b/third_party/xla/xla/stream_executor/cuda/cuda_dnn.cc
@@ -4984,13 +4984,13 @@ static absl::StatusOr<cudnn_frontend::ExecutionPlan> GetExecPlanFromHeuristics(
     } else {
       VLOG(4) << "Failed to build cuDNN execution plan for opGraph "
               << opGraph.getTag()
-              << ". Status: " << CudnnStatusToString(status);
+              << ". absl::Status: " << CudnnStatusToString(status);
     }
   }
 
   LOG(FATAL) << "Failed to generate cuDNN execution plan for opGraph "
              << opGraph.getTag()
-             << ". Status of final plan: " << CudnnStatusToString(status);
+             << ". absl::Status of final plan: " << CudnnStatusToString(status);
 #else
   return absl::UnimplementedError("Supported only for cuDNN >= 8.8.0");
 #endif
@@ -5699,7 +5699,7 @@ class CudnnLegacyConvRunner : public dnn::ConvRunner {
         filter_(std::move(filter)),
         conv_(std::move(conv)) {}
 
-  // Internal form of ToAlgorithmDesc without the StatusOr.
+  // Internal form of ToAlgorithmDesc without the absl::StatusOr.
   dnn::AlgorithmDesc MakeAlgorithmDesc() const {
     return {algo_id_, tensor_ops_enabled_, workspace_size_};
   }
@@ -6712,7 +6712,7 @@ class CudnnLegacyFusedConvRunner : public dnn::FusedConvRunner {
         conv_(std::move(conv)),
         activation_desc_(std::move(activation_desc)) {}
 
-  // Internal form of ToAlgorithmDesc without the StatusOr.
+  // Internal form of ToAlgorithmDesc without the absl::StatusOr.
   dnn::AlgorithmDesc MakeAlgorithmDesc() const {
     return {algo_id_, tensor_ops_enabled_, workspace_size_};
   }
diff --git a/third_party/xla/xla/stream_executor/cuda/cuda_driver.cc b/third_party/xla/xla/stream_executor/cuda/cuda_driver.cc
index ac6a7fb4e0e389..336e821eba2bbb 100644
--- a/third_party/xla/xla/stream_executor/cuda/cuda_driver.cc
+++ b/third_party/xla/xla/stream_executor/cuda/cuda_driver.cc
@@ -2259,7 +2259,7 @@ GpuDriver::CreateMemoryHandle(GpuContext* context, uint64_t bytes) {
 }
 
 // Helper function that turns the integer output of cuDeviceGetAttribute to type
-// T and wraps it in a StatusOr.
+// T and wraps it in a absl::StatusOr.
 template <typename T>
 static absl::StatusOr<T> GetSimpleAttribute(CUdevice device,
                                             CUdevice_attribute attribute) {
diff --git a/third_party/xla/xla/stream_executor/cuda/cuda_executor.cc b/third_party/xla/xla/stream_executor/cuda/cuda_executor.cc
index 596a98fa028e14..cdf5e85bb2e2df 100644
--- a/third_party/xla/xla/stream_executor/cuda/cuda_executor.cc
+++ b/third_party/xla/xla/stream_executor/cuda/cuda_executor.cc
@@ -709,7 +709,8 @@ absl::Status GpuExecutor::Memcpy(Stream* stream, void* host_dst,
   bool ok = GpuDriver::AsynchronousMemcpyD2H(context_, host_dst,
                                              AsCudaDevicePtr(gpu_src), size,
                                              AsGpuStreamValue(stream));
-  // TODO(b/326130105): Change AsynchronousMemcpyD2H calls to return Status.
+  // TODO(b/326130105): Change AsynchronousMemcpyD2H calls to return
+  // absl::Status.
   if (!ok) {
     return absl::InternalError("Failed to memcpy from device to host.");
   }
@@ -721,7 +722,8 @@ absl::Status GpuExecutor::Memcpy(Stream* stream, DeviceMemoryBase* gpu_dst,
   bool ok = GpuDriver::AsynchronousMemcpyH2D(context_, AsCudaDevicePtr(gpu_dst),
                                              host_src, size,
                                              AsGpuStreamValue(stream));
-  // TODO(b/326130105): Change AsynchronousMemcpyD2H calls to return Status.
+  // TODO(b/326130105): Change AsynchronousMemcpyD2H calls to return
+  // absl::Status.
   if (!ok) {
     return absl::InternalError("Failed to memcpy from device to host.");
   }
diff --git a/third_party/xla/xla/stream_executor/rocm/rocm_dnn.cc b/third_party/xla/xla/stream_executor/rocm/rocm_dnn.cc
index cf5f9a9c80b72f..046c86df225408 100644
--- a/third_party/xla/xla/stream_executor/rocm/rocm_dnn.cc
+++ b/third_party/xla/xla/stream_executor/rocm/rocm_dnn.cc
@@ -1975,7 +1975,7 @@ template <typename Base>
 class MIOpenDescriptorCommon : public MixinBase<Base> {
  public:
   bool ok() const { return status_.ok(); }
-  absl::Status Status() const { return status_; }
+  absl::Status absl::Status() const { return status_; }
 
  protected:
   void SetFailure(const absl::Status& status) { status_.Update(status); }
diff --git a/third_party/xla/xla/stream_executor/rocm/rocm_driver.cc b/third_party/xla/xla/stream_executor/rocm/rocm_driver.cc
index 9a5e13cab06e1c..e2e1db3355d56d 100644
--- a/third_party/xla/xla/stream_executor/rocm/rocm_driver.cc
+++ b/third_party/xla/xla/stream_executor/rocm/rocm_driver.cc
@@ -1860,7 +1860,7 @@ struct BitPatternToValue {
 }
 
 // Helper function that turns the integer output of hipDeviceGetAttribute to
-// type T and wraps it in a StatusOr.
+// type T and wraps it in a absl::StatusOr.
 template <typename T>
 static absl::StatusOr<T> GetSimpleAttribute(hipDevice_t device,
                                             hipDeviceAttribute_t attribute) {
diff --git a/third_party/xla/xla/stream_executor/rocm/rocm_executor.cc b/third_party/xla/xla/stream_executor/rocm/rocm_executor.cc
index 226cf3afe231b7..6096c2b1e85183 100644
--- a/third_party/xla/xla/stream_executor/rocm/rocm_executor.cc
+++ b/third_party/xla/xla/stream_executor/rocm/rocm_executor.cc
@@ -611,7 +611,8 @@ absl::Status GpuExecutor::Memcpy(Stream* stream, void* host_dst,
                                              AsROCmDevicePtr(gpu_src), size,
                                              AsGpuStreamValue(stream));
 
-  // TODO(b/326130105): Change AsynchronousMemcpyD2H calls to return Status.
+  // TODO(b/326130105): Change AsynchronousMemcpyD2H calls to return
+  // absl::Status.
   if (!ok) {
     return absl::InternalError("Failed to memcpy from device to host.");
   }
@@ -623,7 +624,8 @@ absl::Status GpuExecutor::Memcpy(Stream* stream, DeviceMemoryBase* gpu_dst,
   bool ok = GpuDriver::AsynchronousMemcpyH2D(context_, AsROCmDevicePtr(gpu_dst),
                                              host_src, size,
                                              AsGpuStreamValue(stream));
-  // TODO(b/326130105): Change AsynchronousMemcpyD2H calls to return Status.
+  // TODO(b/326130105): Change AsynchronousMemcpyD2H calls to return
+  // absl::Status.
   if (!ok) {
     return absl::InternalError("Failed to memcpy from device to host.");
   }
diff --git a/third_party/xla/xla/tools/hlo_bisect/hlo_bisect_state.cc b/third_party/xla/xla/tools/hlo_bisect/hlo_bisect_state.cc
index 19f1b29fd1b77f..a8dfde80578563 100644
--- a/third_party/xla/xla/tools/hlo_bisect/hlo_bisect_state.cc
+++ b/third_party/xla/xla/tools/hlo_bisect/hlo_bisect_state.cc
@@ -55,8 +55,8 @@ std::vector<HloInstruction*> GetModifiedInstructionPostOrder(
 // Changes the module by replacing the original root instruction of the entry
 // computation with a new root instruction that is a tuple containing the values
 // in `outputs`.
-Status MorphModuleWithOutputs(HloModule* module,
-                              absl::Span<HloInstruction* const> outputs) {
+absl::Status MorphModuleWithOutputs(HloModule* module,
+                                    absl::Span<HloInstruction* const> outputs) {
   HloComputation* entry_computation = module->entry_computation();
   HloInstruction* new_root = outputs.size() == 1
                                  ? outputs[0]
@@ -75,7 +75,7 @@ Status MorphModuleWithOutputs(HloModule* module,
 // Changes the module by keeping only the provided instructions of the entry
 // computation (should be sorted in the modified instruction post order),
 // inserting a new root instruction to keep all values live.
-Status MorphModuleWithInstructions(
+absl::Status MorphModuleWithInstructions(
     HloModule* module, absl::Span<HloInstruction* const> instructions) {
   ConstHloInstructionSet in_range_instructions(instructions.begin(),
                                                instructions.end());
@@ -94,7 +94,8 @@ Status MorphModuleWithInstructions(
   return MorphModuleWithOutputs(module, outputs);
 }
 
-Status MorphModuleWithInstructions(HloModule* module, size_t num_instructions) {
+absl::Status MorphModuleWithInstructions(HloModule* module,
+                                         size_t num_instructions) {
   std::vector<HloInstruction*> ordered_instructions =
       GetModifiedInstructionPostOrder(module->entry_computation());
   HloInstruction* const* instructions_begin = &ordered_instructions.front();
@@ -104,7 +105,7 @@ Status MorphModuleWithInstructions(HloModule* module, size_t num_instructions) {
 
 // Changes the module by replacing some instructions in the entry computation
 // with literals.
-Status MorphModuleWithLiterals(
+absl::Status MorphModuleWithLiterals(
     HloModule* module, absl::flat_hash_map<std::string, Literal> literal_map) {
   HloComputation* entry_computation = module->entry_computation();
 
@@ -120,7 +121,7 @@ Status MorphModuleWithLiterals(
     if (!instruction->IsDead()) {
       HloInstruction* new_instruction = entry_computation->AddInstruction(
           HloInstruction::CreateConstant(std::move(literal)));
-      Status replace_status =
+      absl::Status replace_status =
           entry_computation->ReplaceInstruction(instruction, new_instruction);
       TF_RETURN_IF_ERROR(replace_status);
     }
@@ -320,7 +321,7 @@ absl::StatusOr<bool> HloBisectState::TrimByUsingConstants() {
   return has_bug;
 }
 
-Status HloBisectState::ExpectModuleIsBuggy() {
+absl::Status HloBisectState::ExpectModuleIsBuggy() {
   // Verify that the current module has a bug.
   TF_ASSIGN_OR_RETURN(bool has_bug, RunModule(*module_));
   if (has_bug) {
diff --git a/third_party/xla/xla/tools/hlo_bisect/hlo_bisect_state.h b/third_party/xla/xla/tools/hlo_bisect/hlo_bisect_state.h
index 6cc19dd7bcbee9..8b489f795a30bc 100644
--- a/third_party/xla/xla/tools/hlo_bisect/hlo_bisect_state.h
+++ b/third_party/xla/xla/tools/hlo_bisect/hlo_bisect_state.h
@@ -82,7 +82,7 @@ class HloBisectState {
 
   // Asserts that the module still has the bug. If negative, runs the bug
   // checker repeatedly to verify that it's deterministic.
-  Status ExpectModuleIsBuggy();
+  absl::Status ExpectModuleIsBuggy();
 
   std::unique_ptr<xla::HloModule> module_;
   BugCheckerInterface* bug_checker_;
diff --git a/third_party/xla/xla/tools/hlo_bisect/hlo_bisect_utils.cc b/third_party/xla/xla/tools/hlo_bisect/hlo_bisect_utils.cc
index 338d4755ce47c9..c86f4cc0435286 100644
--- a/third_party/xla/xla/tools/hlo_bisect/hlo_bisect_utils.cc
+++ b/third_party/xla/xla/tools/hlo_bisect/hlo_bisect_utils.cc
@@ -109,7 +109,8 @@ absl::StatusOr<ModuleWithInputs> GetModuleAndInputData(
   }
   LOG(INFO) << input_file << " is not HloProto. Trying HLO text.\n";
   std::string hlo_string;
-  Status to_string_status = tsl::ReadFileToString(env, input_file, &hlo_string);
+  absl::Status to_string_status =
+      tsl::ReadFileToString(env, input_file, &hlo_string);
   if (!to_string_status.ok()) {
     LOG(ERROR) << input_file << " problem in reading file to string: "
                << to_string_status.message();
@@ -130,9 +131,9 @@ absl::StatusOr<ModuleWithInputs> GetModuleAndInputData(
 }
 
 // Outputs the given HloModule as HloProto to the given file.
-Status DumpHloModule(HloModule* module, const std::string& file_name,
-                     absl::string_view dir_path,
-                     absl::string_view output_format) {
+absl::Status DumpHloModule(HloModule* module, const std::string& file_name,
+                           absl::string_view dir_path,
+                           absl::string_view output_format) {
   HloProto proto = MakeHloProto(*module);
   if (output_format == "hlo") {
     tsl::Env* env = tsl::Env::Default();
@@ -347,7 +348,7 @@ void RunBisect(std::unique_ptr<BisectRunner> runner, bool all_computations,
   CHECK(bisect_status.ok()) << bisect_status.status().message();
 
   std::unique_ptr<HloModule> new_module = std::move(bisect_status.value());
-  Status dump_status =
+  absl::Status dump_status =
       DumpHloModule(new_module.get(), new_module->name() + "_trimmed",
                     dump_path, output_format);
   CHECK(dump_status.ok()) << dump_status.message();
@@ -368,10 +369,10 @@ absl::StatusOr<ModuleWithInputs> GetVerifiedModuleAndInputData(
       }
     }
   }
-  Status verified_status = HloVerifier(/*layout_sensitive=*/false,
-                                       /*allow_mixed_precision=*/false)
-                               .Run(module.get())
-                               .status();
+  absl::Status verified_status = HloVerifier(/*layout_sensitive=*/false,
+                                             /*allow_mixed_precision=*/false)
+                                     .Run(module.get())
+                                     .status();
   if (!verified_status.ok()) {
     LOG(ERROR) << "Failed to verify hlo module " << verified_status.message();
     return verified_status;
diff --git a/third_party/xla/xla/tools/multihost_hlo_runner/functional_hlo_runner.cc b/third_party/xla/xla/tools/multihost_hlo_runner/functional_hlo_runner.cc
index 989f0476543930..a44fc62f300217 100644
--- a/third_party/xla/xla/tools/multihost_hlo_runner/functional_hlo_runner.cc
+++ b/third_party/xla/xla/tools/multihost_hlo_runner/functional_hlo_runner.cc
@@ -424,7 +424,7 @@ FunctionalHloRunner::CreateExecutableBuildOptionsFromExecutionOptions(
   return build_options;
 }
 
-Status FunctionalHloRunner::DumpOutput(
+absl::Status FunctionalHloRunner::DumpOutput(
     const FunctionalHloRunner::PerDeviceLiteralVecType& output,
     absl::string_view dump_output_to, int task_id) {
   std::vector<std::string> output_path_vec =
@@ -443,7 +443,7 @@ Status FunctionalHloRunner::DumpOutput(
       output_path_vec[literal_id_index] = absl::StrCat("literal_", literal_id);
       std::string literal_path = absl::StrJoin(output_path_vec, ".");
       CHECK_EQ(suffix, std::string("txt"));
-      Status write_status =
+      absl::Status write_status =
           tsl::WriteStringToFile(tsl::Env::Default(), literal_path,
                                  literal_vec[literal_id].ToString());
       if (!write_status.ok()) {
@@ -488,7 +488,7 @@ FunctionalHloRunner::LoadHloModuleAndArguments(absl::string_view hlo_file,
   return hlo_module_and_arguments;
 }
 
-Status FunctionalHloRunner::LoadAndRunAndDump(
+absl::Status FunctionalHloRunner::LoadAndRunAndDump(
     PjRtClient& client, const DebugOptions& debug_options,
     const xla::FunctionalHloRunner::PreprocessingOptions& preproc_options,
     const xla::FunctionalHloRunner::RawCompileOptions& raw_compile_options,
@@ -543,7 +543,7 @@ FunctionalHloRunner::LoadAndRun(PjRtClient& client,
       hlo_module_and_arguments.hlo_module.get(), loaded_arguments);
 }
 
-Status FunctionalHloRunner::LoadAndCompile(
+absl::Status FunctionalHloRunner::LoadAndCompile(
     PjRtClient& client, const DebugOptions& debug_options,
     const PreprocessingOptions& preproc_options,
     const RawCompileOptions& raw_compile_options, std::string_view hlo_file,
@@ -719,7 +719,7 @@ ParameterType GetParameterType(const HloModule& module) {
 
 }  // namespace
 
-Status FunctionalHloRunner::PrepareHloModuleForCompilation(
+absl::Status FunctionalHloRunner::PrepareHloModuleForCompilation(
     HloModule* hlo_module, const DebugOptions& debug_options,
     const PreprocessingOptions& preproc_options) {
   hlo_module->mutable_config().set_debug_options(debug_options);
@@ -894,7 +894,7 @@ std::vector<Shape> GetArgumentShapes(const HloModule& module) {
   return argument_shapes;
 }
 
-Status EnsureSingleTupleForFlattening(const HloModule& module) {
+absl::Status EnsureSingleTupleForFlattening(const HloModule& module) {
   if (module.entry_computation()->num_parameters() != 1) {
     return InvalidArgument(
         "Flattening arguments requires the number of parameters to be 1. "
@@ -1354,7 +1354,7 @@ FunctionalHloRunner::FetchAndLogOutput(
     ModuleOutputMode module_output_mode, bool log_output) {
   CHECK(!output_buffers.empty());
   absl::Mutex mu;
-  Status status;
+  absl::Status status;
   size_t num_pending_transfers = 0;
   bool device_0_is_local = false;
   for (PjRtDevice* device : GetLocalDevices(client)) {
@@ -1389,7 +1389,7 @@ FunctionalHloRunner::FetchAndLogOutput(
                "same device";
         output_slice.emplace_back(
             ShapeUtil::DeviceShapeToHostShape(buffer->on_device_shape()));
-        buffer->ToLiteral(&output_slice.back()).OnReady([&](Status s) {
+        buffer->ToLiteral(&output_slice.back()).OnReady([&](absl::Status s) {
           absl::MutexLock lock(&mu);
           --num_pending_transfers;
           status.Update(s);
diff --git a/third_party/xla/xla/tools/multihost_hlo_runner/functional_hlo_runner.h b/third_party/xla/xla/tools/multihost_hlo_runner/functional_hlo_runner.h
index 04560dcb981eb2..367897a6e2f48f 100644
--- a/third_party/xla/xla/tools/multihost_hlo_runner/functional_hlo_runner.h
+++ b/third_party/xla/xla/tools/multihost_hlo_runner/functional_hlo_runner.h
@@ -241,7 +241,7 @@ class FunctionalHloRunner {
   // Runs on HLO module and dumps the output if needed.
   //
   // This is the highest level API in this file.
-  static Status LoadAndRunAndDump(
+  static absl::Status LoadAndRunAndDump(
       PjRtClient& client, const DebugOptions& debug_options,
       const xla::FunctionalHloRunner::PreprocessingOptions& preproc_options,
       const xla::FunctionalHloRunner::RawCompileOptions& raw_compile_options,
@@ -266,12 +266,11 @@ class FunctionalHloRunner {
   //
   // This function allows compiling multi-device HLOs on machines with fewer
   // devices.
-  static Status LoadAndCompile(PjRtClient& client,
-                               const DebugOptions& debug_options,
-                               const PreprocessingOptions& preproc_options,
-                               const RawCompileOptions& raw_compile_options,
-                               std::string_view hlo_file,
-                               InputFormat input_format, int task_id = 0);
+  static absl::Status LoadAndCompile(
+      PjRtClient& client, const DebugOptions& debug_options,
+      const PreprocessingOptions& preproc_options,
+      const RawCompileOptions& raw_compile_options, std::string_view hlo_file,
+      InputFormat input_format, int task_id = 0);
 
   // Compiles and runs the given HLO module with the given arguments for each
   // device. The given arguments is a map from device ID to a list of arguments.
@@ -326,7 +325,7 @@ class FunctionalHloRunner {
 
   // This would ideally be private, but we need it for the implementation of
   // MultihostHloRunner.
-  static Status PrepareHloModuleForCompilation(
+  static absl::Status PrepareHloModuleForCompilation(
       HloModule* hlo_module, const DebugOptions& debug_options,
       const PreprocessingOptions& preproc_options);
   // This would ideally be private, but we need it for the implementation of
@@ -334,7 +333,7 @@ class FunctionalHloRunner {
   static CompileOptions CompleteCompileOptions(const HloModule& hlo_module,
                                                CompileOptions compile_options);
 
-  static Status DumpOutput(
+  static absl::Status DumpOutput(
       const FunctionalHloRunner::PerDeviceLiteralVecType& output,
       absl::string_view dump_output_to, int task_id);
 

From 01e821d0db18f3f081603fe41ac8ede78cd5fe45 Mon Sep 17 00:00:00 2001
From: Fangrui Song <maskray@google.com>
Date: Thu, 16 May 2024 14:19:27 -0700
Subject: [PATCH 304/478] Integrate LLVM at llvm/llvm-project@a383b3cca338

Updates LLVM usage to match
[a383b3cca338](https://github.com/llvm/llvm-project/commit/a383b3cca338)

PiperOrigin-RevId: 634527229
---
 third_party/llvm/generated.patch | 836 -------------------------------
 third_party/llvm/workspace.bzl   |   4 +-
 2 files changed, 2 insertions(+), 838 deletions(-)

diff --git a/third_party/llvm/generated.patch b/third_party/llvm/generated.patch
index 9ed9c58a9c4eb8..509398da979e83 100644
--- a/third_party/llvm/generated.patch
+++ b/third_party/llvm/generated.patch
@@ -1,837 +1 @@
 Auto generated patch. Do not edit or delete it, even if empty.
-diff -ruN --strip-trailing-cr a/llvm/include/llvm/CodeGen/BasicTTIImpl.h b/llvm/include/llvm/CodeGen/BasicTTIImpl.h
---- a/llvm/include/llvm/CodeGen/BasicTTIImpl.h
-+++ b/llvm/include/llvm/CodeGen/BasicTTIImpl.h
-@@ -612,13 +612,7 @@
-     if (PartialUnrollingThreshold.getNumOccurrences() > 0)
-       MaxOps = PartialUnrollingThreshold;
-     else if (ST->getSchedModel().LoopMicroOpBufferSize > 0)
--      // Upper bound by the default PartialThreshold, which is the same as
--      // the default full-unroll Threshold. Even if the loop micro-op buffer
--      // is very large, this does not mean that we want to unroll all loops
--      // to that length, as it would increase code size beyond the limits of
--      // what unrolling normally allows.
--      MaxOps = std::min(ST->getSchedModel().LoopMicroOpBufferSize,
--                        UP.PartialThreshold);
-+      MaxOps = ST->getSchedModel().LoopMicroOpBufferSize;
-     else
-       return;
- 
-diff -ruN --strip-trailing-cr a/llvm/test/Transforms/LoopUnroll/X86/znver3.ll b/llvm/test/Transforms/LoopUnroll/X86/znver3.ll
---- a/llvm/test/Transforms/LoopUnroll/X86/znver3.ll
-+++ b/llvm/test/Transforms/LoopUnroll/X86/znver3.ll
-@@ -9,8 +9,8 @@
- ; CHECK-NEXT:  entry:
- ; CHECK-NEXT:    br label [[FOR_BODY:%.*]]
- ; CHECK:       for.body:
--; CHECK-NEXT:    [[INDVARS_IV:%.*]] = phi i64 [ 0, [[ENTRY:%.*]] ], [ [[INDVARS_IV_NEXT_31:%.*]], [[FOR_BODY]] ]
--; CHECK-NEXT:    [[SUM:%.*]] = phi i32 [ 0, [[ENTRY]] ], [ [[SUM_NEXT_31:%.*]], [[FOR_BODY]] ]
-+; CHECK-NEXT:    [[INDVARS_IV:%.*]] = phi i64 [ 0, [[ENTRY:%.*]] ], [ [[INDVARS_IV_NEXT_127:%.*]], [[FOR_BODY]] ]
-+; CHECK-NEXT:    [[SUM:%.*]] = phi i32 [ 0, [[ENTRY]] ], [ [[SUM_NEXT_127:%.*]], [[FOR_BODY]] ]
- ; CHECK-NEXT:    [[ARRAYIDX:%.*]] = getelementptr inbounds i32, ptr [[ARY]], i64 [[INDVARS_IV]]
- ; CHECK-NEXT:    [[VAL:%.*]] = load i32, ptr [[ARRAYIDX]], align 4
- ; CHECK-NEXT:    [[SUM_NEXT:%.*]] = add nsw i32 [[VAL]], [[SUM]]
-@@ -137,12 +137,396 @@
- ; CHECK-NEXT:    [[INDVARS_IV_NEXT_30:%.*]] = add nuw nsw i64 [[INDVARS_IV]], 31
- ; CHECK-NEXT:    [[ARRAYIDX_31:%.*]] = getelementptr inbounds i32, ptr [[ARY]], i64 [[INDVARS_IV_NEXT_30]]
- ; CHECK-NEXT:    [[VAL_31:%.*]] = load i32, ptr [[ARRAYIDX_31]], align 4
--; CHECK-NEXT:    [[SUM_NEXT_31]] = add nsw i32 [[VAL_31]], [[SUM_NEXT_30]]
--; CHECK-NEXT:    [[INDVARS_IV_NEXT_31]] = add nuw nsw i64 [[INDVARS_IV]], 32
--; CHECK-NEXT:    [[EXITCOND_NOT_31:%.*]] = icmp eq i64 [[INDVARS_IV_NEXT_31]], 8192
--; CHECK-NEXT:    br i1 [[EXITCOND_NOT_31]], label [[FOR_COND_CLEANUP:%.*]], label [[FOR_BODY]]
-+; CHECK-NEXT:    [[SUM_NEXT_31:%.*]] = add nsw i32 [[VAL_31]], [[SUM_NEXT_30]]
-+; CHECK-NEXT:    [[INDVARS_IV_NEXT_31:%.*]] = add nuw nsw i64 [[INDVARS_IV]], 32
-+; CHECK-NEXT:    [[ARRAYIDX_32:%.*]] = getelementptr inbounds i32, ptr [[ARY]], i64 [[INDVARS_IV_NEXT_31]]
-+; CHECK-NEXT:    [[VAL_32:%.*]] = load i32, ptr [[ARRAYIDX_32]], align 4
-+; CHECK-NEXT:    [[SUM_NEXT_32:%.*]] = add nsw i32 [[VAL_32]], [[SUM_NEXT_31]]
-+; CHECK-NEXT:    [[INDVARS_IV_NEXT_32:%.*]] = add nuw nsw i64 [[INDVARS_IV]], 33
-+; CHECK-NEXT:    [[ARRAYIDX_33:%.*]] = getelementptr inbounds i32, ptr [[ARY]], i64 [[INDVARS_IV_NEXT_32]]
-+; CHECK-NEXT:    [[VAL_33:%.*]] = load i32, ptr [[ARRAYIDX_33]], align 4
-+; CHECK-NEXT:    [[SUM_NEXT_33:%.*]] = add nsw i32 [[VAL_33]], [[SUM_NEXT_32]]
-+; CHECK-NEXT:    [[INDVARS_IV_NEXT_33:%.*]] = add nuw nsw i64 [[INDVARS_IV]], 34
-+; CHECK-NEXT:    [[ARRAYIDX_34:%.*]] = getelementptr inbounds i32, ptr [[ARY]], i64 [[INDVARS_IV_NEXT_33]]
-+; CHECK-NEXT:    [[VAL_34:%.*]] = load i32, ptr [[ARRAYIDX_34]], align 4
-+; CHECK-NEXT:    [[SUM_NEXT_34:%.*]] = add nsw i32 [[VAL_34]], [[SUM_NEXT_33]]
-+; CHECK-NEXT:    [[INDVARS_IV_NEXT_34:%.*]] = add nuw nsw i64 [[INDVARS_IV]], 35
-+; CHECK-NEXT:    [[ARRAYIDX_35:%.*]] = getelementptr inbounds i32, ptr [[ARY]], i64 [[INDVARS_IV_NEXT_34]]
-+; CHECK-NEXT:    [[VAL_35:%.*]] = load i32, ptr [[ARRAYIDX_35]], align 4
-+; CHECK-NEXT:    [[SUM_NEXT_35:%.*]] = add nsw i32 [[VAL_35]], [[SUM_NEXT_34]]
-+; CHECK-NEXT:    [[INDVARS_IV_NEXT_35:%.*]] = add nuw nsw i64 [[INDVARS_IV]], 36
-+; CHECK-NEXT:    [[ARRAYIDX_36:%.*]] = getelementptr inbounds i32, ptr [[ARY]], i64 [[INDVARS_IV_NEXT_35]]
-+; CHECK-NEXT:    [[VAL_36:%.*]] = load i32, ptr [[ARRAYIDX_36]], align 4
-+; CHECK-NEXT:    [[SUM_NEXT_36:%.*]] = add nsw i32 [[VAL_36]], [[SUM_NEXT_35]]
-+; CHECK-NEXT:    [[INDVARS_IV_NEXT_36:%.*]] = add nuw nsw i64 [[INDVARS_IV]], 37
-+; CHECK-NEXT:    [[ARRAYIDX_37:%.*]] = getelementptr inbounds i32, ptr [[ARY]], i64 [[INDVARS_IV_NEXT_36]]
-+; CHECK-NEXT:    [[VAL_37:%.*]] = load i32, ptr [[ARRAYIDX_37]], align 4
-+; CHECK-NEXT:    [[SUM_NEXT_37:%.*]] = add nsw i32 [[VAL_37]], [[SUM_NEXT_36]]
-+; CHECK-NEXT:    [[INDVARS_IV_NEXT_37:%.*]] = add nuw nsw i64 [[INDVARS_IV]], 38
-+; CHECK-NEXT:    [[ARRAYIDX_38:%.*]] = getelementptr inbounds i32, ptr [[ARY]], i64 [[INDVARS_IV_NEXT_37]]
-+; CHECK-NEXT:    [[VAL_38:%.*]] = load i32, ptr [[ARRAYIDX_38]], align 4
-+; CHECK-NEXT:    [[SUM_NEXT_38:%.*]] = add nsw i32 [[VAL_38]], [[SUM_NEXT_37]]
-+; CHECK-NEXT:    [[INDVARS_IV_NEXT_38:%.*]] = add nuw nsw i64 [[INDVARS_IV]], 39
-+; CHECK-NEXT:    [[ARRAYIDX_39:%.*]] = getelementptr inbounds i32, ptr [[ARY]], i64 [[INDVARS_IV_NEXT_38]]
-+; CHECK-NEXT:    [[VAL_39:%.*]] = load i32, ptr [[ARRAYIDX_39]], align 4
-+; CHECK-NEXT:    [[SUM_NEXT_39:%.*]] = add nsw i32 [[VAL_39]], [[SUM_NEXT_38]]
-+; CHECK-NEXT:    [[INDVARS_IV_NEXT_39:%.*]] = add nuw nsw i64 [[INDVARS_IV]], 40
-+; CHECK-NEXT:    [[ARRAYIDX_40:%.*]] = getelementptr inbounds i32, ptr [[ARY]], i64 [[INDVARS_IV_NEXT_39]]
-+; CHECK-NEXT:    [[VAL_40:%.*]] = load i32, ptr [[ARRAYIDX_40]], align 4
-+; CHECK-NEXT:    [[SUM_NEXT_40:%.*]] = add nsw i32 [[VAL_40]], [[SUM_NEXT_39]]
-+; CHECK-NEXT:    [[INDVARS_IV_NEXT_40:%.*]] = add nuw nsw i64 [[INDVARS_IV]], 41
-+; CHECK-NEXT:    [[ARRAYIDX_41:%.*]] = getelementptr inbounds i32, ptr [[ARY]], i64 [[INDVARS_IV_NEXT_40]]
-+; CHECK-NEXT:    [[VAL_41:%.*]] = load i32, ptr [[ARRAYIDX_41]], align 4
-+; CHECK-NEXT:    [[SUM_NEXT_41:%.*]] = add nsw i32 [[VAL_41]], [[SUM_NEXT_40]]
-+; CHECK-NEXT:    [[INDVARS_IV_NEXT_41:%.*]] = add nuw nsw i64 [[INDVARS_IV]], 42
-+; CHECK-NEXT:    [[ARRAYIDX_42:%.*]] = getelementptr inbounds i32, ptr [[ARY]], i64 [[INDVARS_IV_NEXT_41]]
-+; CHECK-NEXT:    [[VAL_42:%.*]] = load i32, ptr [[ARRAYIDX_42]], align 4
-+; CHECK-NEXT:    [[SUM_NEXT_42:%.*]] = add nsw i32 [[VAL_42]], [[SUM_NEXT_41]]
-+; CHECK-NEXT:    [[INDVARS_IV_NEXT_42:%.*]] = add nuw nsw i64 [[INDVARS_IV]], 43
-+; CHECK-NEXT:    [[ARRAYIDX_43:%.*]] = getelementptr inbounds i32, ptr [[ARY]], i64 [[INDVARS_IV_NEXT_42]]
-+; CHECK-NEXT:    [[VAL_43:%.*]] = load i32, ptr [[ARRAYIDX_43]], align 4
-+; CHECK-NEXT:    [[SUM_NEXT_43:%.*]] = add nsw i32 [[VAL_43]], [[SUM_NEXT_42]]
-+; CHECK-NEXT:    [[INDVARS_IV_NEXT_43:%.*]] = add nuw nsw i64 [[INDVARS_IV]], 44
-+; CHECK-NEXT:    [[ARRAYIDX_44:%.*]] = getelementptr inbounds i32, ptr [[ARY]], i64 [[INDVARS_IV_NEXT_43]]
-+; CHECK-NEXT:    [[VAL_44:%.*]] = load i32, ptr [[ARRAYIDX_44]], align 4
-+; CHECK-NEXT:    [[SUM_NEXT_44:%.*]] = add nsw i32 [[VAL_44]], [[SUM_NEXT_43]]
-+; CHECK-NEXT:    [[INDVARS_IV_NEXT_44:%.*]] = add nuw nsw i64 [[INDVARS_IV]], 45
-+; CHECK-NEXT:    [[ARRAYIDX_45:%.*]] = getelementptr inbounds i32, ptr [[ARY]], i64 [[INDVARS_IV_NEXT_44]]
-+; CHECK-NEXT:    [[VAL_45:%.*]] = load i32, ptr [[ARRAYIDX_45]], align 4
-+; CHECK-NEXT:    [[SUM_NEXT_45:%.*]] = add nsw i32 [[VAL_45]], [[SUM_NEXT_44]]
-+; CHECK-NEXT:    [[INDVARS_IV_NEXT_45:%.*]] = add nuw nsw i64 [[INDVARS_IV]], 46
-+; CHECK-NEXT:    [[ARRAYIDX_46:%.*]] = getelementptr inbounds i32, ptr [[ARY]], i64 [[INDVARS_IV_NEXT_45]]
-+; CHECK-NEXT:    [[VAL_46:%.*]] = load i32, ptr [[ARRAYIDX_46]], align 4
-+; CHECK-NEXT:    [[SUM_NEXT_46:%.*]] = add nsw i32 [[VAL_46]], [[SUM_NEXT_45]]
-+; CHECK-NEXT:    [[INDVARS_IV_NEXT_46:%.*]] = add nuw nsw i64 [[INDVARS_IV]], 47
-+; CHECK-NEXT:    [[ARRAYIDX_47:%.*]] = getelementptr inbounds i32, ptr [[ARY]], i64 [[INDVARS_IV_NEXT_46]]
-+; CHECK-NEXT:    [[VAL_47:%.*]] = load i32, ptr [[ARRAYIDX_47]], align 4
-+; CHECK-NEXT:    [[SUM_NEXT_47:%.*]] = add nsw i32 [[VAL_47]], [[SUM_NEXT_46]]
-+; CHECK-NEXT:    [[INDVARS_IV_NEXT_47:%.*]] = add nuw nsw i64 [[INDVARS_IV]], 48
-+; CHECK-NEXT:    [[ARRAYIDX_48:%.*]] = getelementptr inbounds i32, ptr [[ARY]], i64 [[INDVARS_IV_NEXT_47]]
-+; CHECK-NEXT:    [[VAL_48:%.*]] = load i32, ptr [[ARRAYIDX_48]], align 4
-+; CHECK-NEXT:    [[SUM_NEXT_48:%.*]] = add nsw i32 [[VAL_48]], [[SUM_NEXT_47]]
-+; CHECK-NEXT:    [[INDVARS_IV_NEXT_48:%.*]] = add nuw nsw i64 [[INDVARS_IV]], 49
-+; CHECK-NEXT:    [[ARRAYIDX_49:%.*]] = getelementptr inbounds i32, ptr [[ARY]], i64 [[INDVARS_IV_NEXT_48]]
-+; CHECK-NEXT:    [[VAL_49:%.*]] = load i32, ptr [[ARRAYIDX_49]], align 4
-+; CHECK-NEXT:    [[SUM_NEXT_49:%.*]] = add nsw i32 [[VAL_49]], [[SUM_NEXT_48]]
-+; CHECK-NEXT:    [[INDVARS_IV_NEXT_49:%.*]] = add nuw nsw i64 [[INDVARS_IV]], 50
-+; CHECK-NEXT:    [[ARRAYIDX_50:%.*]] = getelementptr inbounds i32, ptr [[ARY]], i64 [[INDVARS_IV_NEXT_49]]
-+; CHECK-NEXT:    [[VAL_50:%.*]] = load i32, ptr [[ARRAYIDX_50]], align 4
-+; CHECK-NEXT:    [[SUM_NEXT_50:%.*]] = add nsw i32 [[VAL_50]], [[SUM_NEXT_49]]
-+; CHECK-NEXT:    [[INDVARS_IV_NEXT_50:%.*]] = add nuw nsw i64 [[INDVARS_IV]], 51
-+; CHECK-NEXT:    [[ARRAYIDX_51:%.*]] = getelementptr inbounds i32, ptr [[ARY]], i64 [[INDVARS_IV_NEXT_50]]
-+; CHECK-NEXT:    [[VAL_51:%.*]] = load i32, ptr [[ARRAYIDX_51]], align 4
-+; CHECK-NEXT:    [[SUM_NEXT_51:%.*]] = add nsw i32 [[VAL_51]], [[SUM_NEXT_50]]
-+; CHECK-NEXT:    [[INDVARS_IV_NEXT_51:%.*]] = add nuw nsw i64 [[INDVARS_IV]], 52
-+; CHECK-NEXT:    [[ARRAYIDX_52:%.*]] = getelementptr inbounds i32, ptr [[ARY]], i64 [[INDVARS_IV_NEXT_51]]
-+; CHECK-NEXT:    [[VAL_52:%.*]] = load i32, ptr [[ARRAYIDX_52]], align 4
-+; CHECK-NEXT:    [[SUM_NEXT_52:%.*]] = add nsw i32 [[VAL_52]], [[SUM_NEXT_51]]
-+; CHECK-NEXT:    [[INDVARS_IV_NEXT_52:%.*]] = add nuw nsw i64 [[INDVARS_IV]], 53
-+; CHECK-NEXT:    [[ARRAYIDX_53:%.*]] = getelementptr inbounds i32, ptr [[ARY]], i64 [[INDVARS_IV_NEXT_52]]
-+; CHECK-NEXT:    [[VAL_53:%.*]] = load i32, ptr [[ARRAYIDX_53]], align 4
-+; CHECK-NEXT:    [[SUM_NEXT_53:%.*]] = add nsw i32 [[VAL_53]], [[SUM_NEXT_52]]
-+; CHECK-NEXT:    [[INDVARS_IV_NEXT_53:%.*]] = add nuw nsw i64 [[INDVARS_IV]], 54
-+; CHECK-NEXT:    [[ARRAYIDX_54:%.*]] = getelementptr inbounds i32, ptr [[ARY]], i64 [[INDVARS_IV_NEXT_53]]
-+; CHECK-NEXT:    [[VAL_54:%.*]] = load i32, ptr [[ARRAYIDX_54]], align 4
-+; CHECK-NEXT:    [[SUM_NEXT_54:%.*]] = add nsw i32 [[VAL_54]], [[SUM_NEXT_53]]
-+; CHECK-NEXT:    [[INDVARS_IV_NEXT_54:%.*]] = add nuw nsw i64 [[INDVARS_IV]], 55
-+; CHECK-NEXT:    [[ARRAYIDX_55:%.*]] = getelementptr inbounds i32, ptr [[ARY]], i64 [[INDVARS_IV_NEXT_54]]
-+; CHECK-NEXT:    [[VAL_55:%.*]] = load i32, ptr [[ARRAYIDX_55]], align 4
-+; CHECK-NEXT:    [[SUM_NEXT_55:%.*]] = add nsw i32 [[VAL_55]], [[SUM_NEXT_54]]
-+; CHECK-NEXT:    [[INDVARS_IV_NEXT_55:%.*]] = add nuw nsw i64 [[INDVARS_IV]], 56
-+; CHECK-NEXT:    [[ARRAYIDX_56:%.*]] = getelementptr inbounds i32, ptr [[ARY]], i64 [[INDVARS_IV_NEXT_55]]
-+; CHECK-NEXT:    [[VAL_56:%.*]] = load i32, ptr [[ARRAYIDX_56]], align 4
-+; CHECK-NEXT:    [[SUM_NEXT_56:%.*]] = add nsw i32 [[VAL_56]], [[SUM_NEXT_55]]
-+; CHECK-NEXT:    [[INDVARS_IV_NEXT_56:%.*]] = add nuw nsw i64 [[INDVARS_IV]], 57
-+; CHECK-NEXT:    [[ARRAYIDX_57:%.*]] = getelementptr inbounds i32, ptr [[ARY]], i64 [[INDVARS_IV_NEXT_56]]
-+; CHECK-NEXT:    [[VAL_57:%.*]] = load i32, ptr [[ARRAYIDX_57]], align 4
-+; CHECK-NEXT:    [[SUM_NEXT_57:%.*]] = add nsw i32 [[VAL_57]], [[SUM_NEXT_56]]
-+; CHECK-NEXT:    [[INDVARS_IV_NEXT_57:%.*]] = add nuw nsw i64 [[INDVARS_IV]], 58
-+; CHECK-NEXT:    [[ARRAYIDX_58:%.*]] = getelementptr inbounds i32, ptr [[ARY]], i64 [[INDVARS_IV_NEXT_57]]
-+; CHECK-NEXT:    [[VAL_58:%.*]] = load i32, ptr [[ARRAYIDX_58]], align 4
-+; CHECK-NEXT:    [[SUM_NEXT_58:%.*]] = add nsw i32 [[VAL_58]], [[SUM_NEXT_57]]
-+; CHECK-NEXT:    [[INDVARS_IV_NEXT_58:%.*]] = add nuw nsw i64 [[INDVARS_IV]], 59
-+; CHECK-NEXT:    [[ARRAYIDX_59:%.*]] = getelementptr inbounds i32, ptr [[ARY]], i64 [[INDVARS_IV_NEXT_58]]
-+; CHECK-NEXT:    [[VAL_59:%.*]] = load i32, ptr [[ARRAYIDX_59]], align 4
-+; CHECK-NEXT:    [[SUM_NEXT_59:%.*]] = add nsw i32 [[VAL_59]], [[SUM_NEXT_58]]
-+; CHECK-NEXT:    [[INDVARS_IV_NEXT_59:%.*]] = add nuw nsw i64 [[INDVARS_IV]], 60
-+; CHECK-NEXT:    [[ARRAYIDX_60:%.*]] = getelementptr inbounds i32, ptr [[ARY]], i64 [[INDVARS_IV_NEXT_59]]
-+; CHECK-NEXT:    [[VAL_60:%.*]] = load i32, ptr [[ARRAYIDX_60]], align 4
-+; CHECK-NEXT:    [[SUM_NEXT_60:%.*]] = add nsw i32 [[VAL_60]], [[SUM_NEXT_59]]
-+; CHECK-NEXT:    [[INDVARS_IV_NEXT_60:%.*]] = add nuw nsw i64 [[INDVARS_IV]], 61
-+; CHECK-NEXT:    [[ARRAYIDX_61:%.*]] = getelementptr inbounds i32, ptr [[ARY]], i64 [[INDVARS_IV_NEXT_60]]
-+; CHECK-NEXT:    [[VAL_61:%.*]] = load i32, ptr [[ARRAYIDX_61]], align 4
-+; CHECK-NEXT:    [[SUM_NEXT_61:%.*]] = add nsw i32 [[VAL_61]], [[SUM_NEXT_60]]
-+; CHECK-NEXT:    [[INDVARS_IV_NEXT_61:%.*]] = add nuw nsw i64 [[INDVARS_IV]], 62
-+; CHECK-NEXT:    [[ARRAYIDX_62:%.*]] = getelementptr inbounds i32, ptr [[ARY]], i64 [[INDVARS_IV_NEXT_61]]
-+; CHECK-NEXT:    [[VAL_62:%.*]] = load i32, ptr [[ARRAYIDX_62]], align 4
-+; CHECK-NEXT:    [[SUM_NEXT_62:%.*]] = add nsw i32 [[VAL_62]], [[SUM_NEXT_61]]
-+; CHECK-NEXT:    [[INDVARS_IV_NEXT_62:%.*]] = add nuw nsw i64 [[INDVARS_IV]], 63
-+; CHECK-NEXT:    [[ARRAYIDX_63:%.*]] = getelementptr inbounds i32, ptr [[ARY]], i64 [[INDVARS_IV_NEXT_62]]
-+; CHECK-NEXT:    [[VAL_63:%.*]] = load i32, ptr [[ARRAYIDX_63]], align 4
-+; CHECK-NEXT:    [[SUM_NEXT_63:%.*]] = add nsw i32 [[VAL_63]], [[SUM_NEXT_62]]
-+; CHECK-NEXT:    [[INDVARS_IV_NEXT_63:%.*]] = add nuw nsw i64 [[INDVARS_IV]], 64
-+; CHECK-NEXT:    [[ARRAYIDX_64:%.*]] = getelementptr inbounds i32, ptr [[ARY]], i64 [[INDVARS_IV_NEXT_63]]
-+; CHECK-NEXT:    [[VAL_64:%.*]] = load i32, ptr [[ARRAYIDX_64]], align 4
-+; CHECK-NEXT:    [[SUM_NEXT_64:%.*]] = add nsw i32 [[VAL_64]], [[SUM_NEXT_63]]
-+; CHECK-NEXT:    [[INDVARS_IV_NEXT_64:%.*]] = add nuw nsw i64 [[INDVARS_IV]], 65
-+; CHECK-NEXT:    [[ARRAYIDX_65:%.*]] = getelementptr inbounds i32, ptr [[ARY]], i64 [[INDVARS_IV_NEXT_64]]
-+; CHECK-NEXT:    [[VAL_65:%.*]] = load i32, ptr [[ARRAYIDX_65]], align 4
-+; CHECK-NEXT:    [[SUM_NEXT_65:%.*]] = add nsw i32 [[VAL_65]], [[SUM_NEXT_64]]
-+; CHECK-NEXT:    [[INDVARS_IV_NEXT_65:%.*]] = add nuw nsw i64 [[INDVARS_IV]], 66
-+; CHECK-NEXT:    [[ARRAYIDX_66:%.*]] = getelementptr inbounds i32, ptr [[ARY]], i64 [[INDVARS_IV_NEXT_65]]
-+; CHECK-NEXT:    [[VAL_66:%.*]] = load i32, ptr [[ARRAYIDX_66]], align 4
-+; CHECK-NEXT:    [[SUM_NEXT_66:%.*]] = add nsw i32 [[VAL_66]], [[SUM_NEXT_65]]
-+; CHECK-NEXT:    [[INDVARS_IV_NEXT_66:%.*]] = add nuw nsw i64 [[INDVARS_IV]], 67
-+; CHECK-NEXT:    [[ARRAYIDX_67:%.*]] = getelementptr inbounds i32, ptr [[ARY]], i64 [[INDVARS_IV_NEXT_66]]
-+; CHECK-NEXT:    [[VAL_67:%.*]] = load i32, ptr [[ARRAYIDX_67]], align 4
-+; CHECK-NEXT:    [[SUM_NEXT_67:%.*]] = add nsw i32 [[VAL_67]], [[SUM_NEXT_66]]
-+; CHECK-NEXT:    [[INDVARS_IV_NEXT_67:%.*]] = add nuw nsw i64 [[INDVARS_IV]], 68
-+; CHECK-NEXT:    [[ARRAYIDX_68:%.*]] = getelementptr inbounds i32, ptr [[ARY]], i64 [[INDVARS_IV_NEXT_67]]
-+; CHECK-NEXT:    [[VAL_68:%.*]] = load i32, ptr [[ARRAYIDX_68]], align 4
-+; CHECK-NEXT:    [[SUM_NEXT_68:%.*]] = add nsw i32 [[VAL_68]], [[SUM_NEXT_67]]
-+; CHECK-NEXT:    [[INDVARS_IV_NEXT_68:%.*]] = add nuw nsw i64 [[INDVARS_IV]], 69
-+; CHECK-NEXT:    [[ARRAYIDX_69:%.*]] = getelementptr inbounds i32, ptr [[ARY]], i64 [[INDVARS_IV_NEXT_68]]
-+; CHECK-NEXT:    [[VAL_69:%.*]] = load i32, ptr [[ARRAYIDX_69]], align 4
-+; CHECK-NEXT:    [[SUM_NEXT_69:%.*]] = add nsw i32 [[VAL_69]], [[SUM_NEXT_68]]
-+; CHECK-NEXT:    [[INDVARS_IV_NEXT_69:%.*]] = add nuw nsw i64 [[INDVARS_IV]], 70
-+; CHECK-NEXT:    [[ARRAYIDX_70:%.*]] = getelementptr inbounds i32, ptr [[ARY]], i64 [[INDVARS_IV_NEXT_69]]
-+; CHECK-NEXT:    [[VAL_70:%.*]] = load i32, ptr [[ARRAYIDX_70]], align 4
-+; CHECK-NEXT:    [[SUM_NEXT_70:%.*]] = add nsw i32 [[VAL_70]], [[SUM_NEXT_69]]
-+; CHECK-NEXT:    [[INDVARS_IV_NEXT_70:%.*]] = add nuw nsw i64 [[INDVARS_IV]], 71
-+; CHECK-NEXT:    [[ARRAYIDX_71:%.*]] = getelementptr inbounds i32, ptr [[ARY]], i64 [[INDVARS_IV_NEXT_70]]
-+; CHECK-NEXT:    [[VAL_71:%.*]] = load i32, ptr [[ARRAYIDX_71]], align 4
-+; CHECK-NEXT:    [[SUM_NEXT_71:%.*]] = add nsw i32 [[VAL_71]], [[SUM_NEXT_70]]
-+; CHECK-NEXT:    [[INDVARS_IV_NEXT_71:%.*]] = add nuw nsw i64 [[INDVARS_IV]], 72
-+; CHECK-NEXT:    [[ARRAYIDX_72:%.*]] = getelementptr inbounds i32, ptr [[ARY]], i64 [[INDVARS_IV_NEXT_71]]
-+; CHECK-NEXT:    [[VAL_72:%.*]] = load i32, ptr [[ARRAYIDX_72]], align 4
-+; CHECK-NEXT:    [[SUM_NEXT_72:%.*]] = add nsw i32 [[VAL_72]], [[SUM_NEXT_71]]
-+; CHECK-NEXT:    [[INDVARS_IV_NEXT_72:%.*]] = add nuw nsw i64 [[INDVARS_IV]], 73
-+; CHECK-NEXT:    [[ARRAYIDX_73:%.*]] = getelementptr inbounds i32, ptr [[ARY]], i64 [[INDVARS_IV_NEXT_72]]
-+; CHECK-NEXT:    [[VAL_73:%.*]] = load i32, ptr [[ARRAYIDX_73]], align 4
-+; CHECK-NEXT:    [[SUM_NEXT_73:%.*]] = add nsw i32 [[VAL_73]], [[SUM_NEXT_72]]
-+; CHECK-NEXT:    [[INDVARS_IV_NEXT_73:%.*]] = add nuw nsw i64 [[INDVARS_IV]], 74
-+; CHECK-NEXT:    [[ARRAYIDX_74:%.*]] = getelementptr inbounds i32, ptr [[ARY]], i64 [[INDVARS_IV_NEXT_73]]
-+; CHECK-NEXT:    [[VAL_74:%.*]] = load i32, ptr [[ARRAYIDX_74]], align 4
-+; CHECK-NEXT:    [[SUM_NEXT_74:%.*]] = add nsw i32 [[VAL_74]], [[SUM_NEXT_73]]
-+; CHECK-NEXT:    [[INDVARS_IV_NEXT_74:%.*]] = add nuw nsw i64 [[INDVARS_IV]], 75
-+; CHECK-NEXT:    [[ARRAYIDX_75:%.*]] = getelementptr inbounds i32, ptr [[ARY]], i64 [[INDVARS_IV_NEXT_74]]
-+; CHECK-NEXT:    [[VAL_75:%.*]] = load i32, ptr [[ARRAYIDX_75]], align 4
-+; CHECK-NEXT:    [[SUM_NEXT_75:%.*]] = add nsw i32 [[VAL_75]], [[SUM_NEXT_74]]
-+; CHECK-NEXT:    [[INDVARS_IV_NEXT_75:%.*]] = add nuw nsw i64 [[INDVARS_IV]], 76
-+; CHECK-NEXT:    [[ARRAYIDX_76:%.*]] = getelementptr inbounds i32, ptr [[ARY]], i64 [[INDVARS_IV_NEXT_75]]
-+; CHECK-NEXT:    [[VAL_76:%.*]] = load i32, ptr [[ARRAYIDX_76]], align 4
-+; CHECK-NEXT:    [[SUM_NEXT_76:%.*]] = add nsw i32 [[VAL_76]], [[SUM_NEXT_75]]
-+; CHECK-NEXT:    [[INDVARS_IV_NEXT_76:%.*]] = add nuw nsw i64 [[INDVARS_IV]], 77
-+; CHECK-NEXT:    [[ARRAYIDX_77:%.*]] = getelementptr inbounds i32, ptr [[ARY]], i64 [[INDVARS_IV_NEXT_76]]
-+; CHECK-NEXT:    [[VAL_77:%.*]] = load i32, ptr [[ARRAYIDX_77]], align 4
-+; CHECK-NEXT:    [[SUM_NEXT_77:%.*]] = add nsw i32 [[VAL_77]], [[SUM_NEXT_76]]
-+; CHECK-NEXT:    [[INDVARS_IV_NEXT_77:%.*]] = add nuw nsw i64 [[INDVARS_IV]], 78
-+; CHECK-NEXT:    [[ARRAYIDX_78:%.*]] = getelementptr inbounds i32, ptr [[ARY]], i64 [[INDVARS_IV_NEXT_77]]
-+; CHECK-NEXT:    [[VAL_78:%.*]] = load i32, ptr [[ARRAYIDX_78]], align 4
-+; CHECK-NEXT:    [[SUM_NEXT_78:%.*]] = add nsw i32 [[VAL_78]], [[SUM_NEXT_77]]
-+; CHECK-NEXT:    [[INDVARS_IV_NEXT_78:%.*]] = add nuw nsw i64 [[INDVARS_IV]], 79
-+; CHECK-NEXT:    [[ARRAYIDX_79:%.*]] = getelementptr inbounds i32, ptr [[ARY]], i64 [[INDVARS_IV_NEXT_78]]
-+; CHECK-NEXT:    [[VAL_79:%.*]] = load i32, ptr [[ARRAYIDX_79]], align 4
-+; CHECK-NEXT:    [[SUM_NEXT_79:%.*]] = add nsw i32 [[VAL_79]], [[SUM_NEXT_78]]
-+; CHECK-NEXT:    [[INDVARS_IV_NEXT_79:%.*]] = add nuw nsw i64 [[INDVARS_IV]], 80
-+; CHECK-NEXT:    [[ARRAYIDX_80:%.*]] = getelementptr inbounds i32, ptr [[ARY]], i64 [[INDVARS_IV_NEXT_79]]
-+; CHECK-NEXT:    [[VAL_80:%.*]] = load i32, ptr [[ARRAYIDX_80]], align 4
-+; CHECK-NEXT:    [[SUM_NEXT_80:%.*]] = add nsw i32 [[VAL_80]], [[SUM_NEXT_79]]
-+; CHECK-NEXT:    [[INDVARS_IV_NEXT_80:%.*]] = add nuw nsw i64 [[INDVARS_IV]], 81
-+; CHECK-NEXT:    [[ARRAYIDX_81:%.*]] = getelementptr inbounds i32, ptr [[ARY]], i64 [[INDVARS_IV_NEXT_80]]
-+; CHECK-NEXT:    [[VAL_81:%.*]] = load i32, ptr [[ARRAYIDX_81]], align 4
-+; CHECK-NEXT:    [[SUM_NEXT_81:%.*]] = add nsw i32 [[VAL_81]], [[SUM_NEXT_80]]
-+; CHECK-NEXT:    [[INDVARS_IV_NEXT_81:%.*]] = add nuw nsw i64 [[INDVARS_IV]], 82
-+; CHECK-NEXT:    [[ARRAYIDX_82:%.*]] = getelementptr inbounds i32, ptr [[ARY]], i64 [[INDVARS_IV_NEXT_81]]
-+; CHECK-NEXT:    [[VAL_82:%.*]] = load i32, ptr [[ARRAYIDX_82]], align 4
-+; CHECK-NEXT:    [[SUM_NEXT_82:%.*]] = add nsw i32 [[VAL_82]], [[SUM_NEXT_81]]
-+; CHECK-NEXT:    [[INDVARS_IV_NEXT_82:%.*]] = add nuw nsw i64 [[INDVARS_IV]], 83
-+; CHECK-NEXT:    [[ARRAYIDX_83:%.*]] = getelementptr inbounds i32, ptr [[ARY]], i64 [[INDVARS_IV_NEXT_82]]
-+; CHECK-NEXT:    [[VAL_83:%.*]] = load i32, ptr [[ARRAYIDX_83]], align 4
-+; CHECK-NEXT:    [[SUM_NEXT_83:%.*]] = add nsw i32 [[VAL_83]], [[SUM_NEXT_82]]
-+; CHECK-NEXT:    [[INDVARS_IV_NEXT_83:%.*]] = add nuw nsw i64 [[INDVARS_IV]], 84
-+; CHECK-NEXT:    [[ARRAYIDX_84:%.*]] = getelementptr inbounds i32, ptr [[ARY]], i64 [[INDVARS_IV_NEXT_83]]
-+; CHECK-NEXT:    [[VAL_84:%.*]] = load i32, ptr [[ARRAYIDX_84]], align 4
-+; CHECK-NEXT:    [[SUM_NEXT_84:%.*]] = add nsw i32 [[VAL_84]], [[SUM_NEXT_83]]
-+; CHECK-NEXT:    [[INDVARS_IV_NEXT_84:%.*]] = add nuw nsw i64 [[INDVARS_IV]], 85
-+; CHECK-NEXT:    [[ARRAYIDX_85:%.*]] = getelementptr inbounds i32, ptr [[ARY]], i64 [[INDVARS_IV_NEXT_84]]
-+; CHECK-NEXT:    [[VAL_85:%.*]] = load i32, ptr [[ARRAYIDX_85]], align 4
-+; CHECK-NEXT:    [[SUM_NEXT_85:%.*]] = add nsw i32 [[VAL_85]], [[SUM_NEXT_84]]
-+; CHECK-NEXT:    [[INDVARS_IV_NEXT_85:%.*]] = add nuw nsw i64 [[INDVARS_IV]], 86
-+; CHECK-NEXT:    [[ARRAYIDX_86:%.*]] = getelementptr inbounds i32, ptr [[ARY]], i64 [[INDVARS_IV_NEXT_85]]
-+; CHECK-NEXT:    [[VAL_86:%.*]] = load i32, ptr [[ARRAYIDX_86]], align 4
-+; CHECK-NEXT:    [[SUM_NEXT_86:%.*]] = add nsw i32 [[VAL_86]], [[SUM_NEXT_85]]
-+; CHECK-NEXT:    [[INDVARS_IV_NEXT_86:%.*]] = add nuw nsw i64 [[INDVARS_IV]], 87
-+; CHECK-NEXT:    [[ARRAYIDX_87:%.*]] = getelementptr inbounds i32, ptr [[ARY]], i64 [[INDVARS_IV_NEXT_86]]
-+; CHECK-NEXT:    [[VAL_87:%.*]] = load i32, ptr [[ARRAYIDX_87]], align 4
-+; CHECK-NEXT:    [[SUM_NEXT_87:%.*]] = add nsw i32 [[VAL_87]], [[SUM_NEXT_86]]
-+; CHECK-NEXT:    [[INDVARS_IV_NEXT_87:%.*]] = add nuw nsw i64 [[INDVARS_IV]], 88
-+; CHECK-NEXT:    [[ARRAYIDX_88:%.*]] = getelementptr inbounds i32, ptr [[ARY]], i64 [[INDVARS_IV_NEXT_87]]
-+; CHECK-NEXT:    [[VAL_88:%.*]] = load i32, ptr [[ARRAYIDX_88]], align 4
-+; CHECK-NEXT:    [[SUM_NEXT_88:%.*]] = add nsw i32 [[VAL_88]], [[SUM_NEXT_87]]
-+; CHECK-NEXT:    [[INDVARS_IV_NEXT_88:%.*]] = add nuw nsw i64 [[INDVARS_IV]], 89
-+; CHECK-NEXT:    [[ARRAYIDX_89:%.*]] = getelementptr inbounds i32, ptr [[ARY]], i64 [[INDVARS_IV_NEXT_88]]
-+; CHECK-NEXT:    [[VAL_89:%.*]] = load i32, ptr [[ARRAYIDX_89]], align 4
-+; CHECK-NEXT:    [[SUM_NEXT_89:%.*]] = add nsw i32 [[VAL_89]], [[SUM_NEXT_88]]
-+; CHECK-NEXT:    [[INDVARS_IV_NEXT_89:%.*]] = add nuw nsw i64 [[INDVARS_IV]], 90
-+; CHECK-NEXT:    [[ARRAYIDX_90:%.*]] = getelementptr inbounds i32, ptr [[ARY]], i64 [[INDVARS_IV_NEXT_89]]
-+; CHECK-NEXT:    [[VAL_90:%.*]] = load i32, ptr [[ARRAYIDX_90]], align 4
-+; CHECK-NEXT:    [[SUM_NEXT_90:%.*]] = add nsw i32 [[VAL_90]], [[SUM_NEXT_89]]
-+; CHECK-NEXT:    [[INDVARS_IV_NEXT_90:%.*]] = add nuw nsw i64 [[INDVARS_IV]], 91
-+; CHECK-NEXT:    [[ARRAYIDX_91:%.*]] = getelementptr inbounds i32, ptr [[ARY]], i64 [[INDVARS_IV_NEXT_90]]
-+; CHECK-NEXT:    [[VAL_91:%.*]] = load i32, ptr [[ARRAYIDX_91]], align 4
-+; CHECK-NEXT:    [[SUM_NEXT_91:%.*]] = add nsw i32 [[VAL_91]], [[SUM_NEXT_90]]
-+; CHECK-NEXT:    [[INDVARS_IV_NEXT_91:%.*]] = add nuw nsw i64 [[INDVARS_IV]], 92
-+; CHECK-NEXT:    [[ARRAYIDX_92:%.*]] = getelementptr inbounds i32, ptr [[ARY]], i64 [[INDVARS_IV_NEXT_91]]
-+; CHECK-NEXT:    [[VAL_92:%.*]] = load i32, ptr [[ARRAYIDX_92]], align 4
-+; CHECK-NEXT:    [[SUM_NEXT_92:%.*]] = add nsw i32 [[VAL_92]], [[SUM_NEXT_91]]
-+; CHECK-NEXT:    [[INDVARS_IV_NEXT_92:%.*]] = add nuw nsw i64 [[INDVARS_IV]], 93
-+; CHECK-NEXT:    [[ARRAYIDX_93:%.*]] = getelementptr inbounds i32, ptr [[ARY]], i64 [[INDVARS_IV_NEXT_92]]
-+; CHECK-NEXT:    [[VAL_93:%.*]] = load i32, ptr [[ARRAYIDX_93]], align 4
-+; CHECK-NEXT:    [[SUM_NEXT_93:%.*]] = add nsw i32 [[VAL_93]], [[SUM_NEXT_92]]
-+; CHECK-NEXT:    [[INDVARS_IV_NEXT_93:%.*]] = add nuw nsw i64 [[INDVARS_IV]], 94
-+; CHECK-NEXT:    [[ARRAYIDX_94:%.*]] = getelementptr inbounds i32, ptr [[ARY]], i64 [[INDVARS_IV_NEXT_93]]
-+; CHECK-NEXT:    [[VAL_94:%.*]] = load i32, ptr [[ARRAYIDX_94]], align 4
-+; CHECK-NEXT:    [[SUM_NEXT_94:%.*]] = add nsw i32 [[VAL_94]], [[SUM_NEXT_93]]
-+; CHECK-NEXT:    [[INDVARS_IV_NEXT_94:%.*]] = add nuw nsw i64 [[INDVARS_IV]], 95
-+; CHECK-NEXT:    [[ARRAYIDX_95:%.*]] = getelementptr inbounds i32, ptr [[ARY]], i64 [[INDVARS_IV_NEXT_94]]
-+; CHECK-NEXT:    [[VAL_95:%.*]] = load i32, ptr [[ARRAYIDX_95]], align 4
-+; CHECK-NEXT:    [[SUM_NEXT_95:%.*]] = add nsw i32 [[VAL_95]], [[SUM_NEXT_94]]
-+; CHECK-NEXT:    [[INDVARS_IV_NEXT_95:%.*]] = add nuw nsw i64 [[INDVARS_IV]], 96
-+; CHECK-NEXT:    [[ARRAYIDX_96:%.*]] = getelementptr inbounds i32, ptr [[ARY]], i64 [[INDVARS_IV_NEXT_95]]
-+; CHECK-NEXT:    [[VAL_96:%.*]] = load i32, ptr [[ARRAYIDX_96]], align 4
-+; CHECK-NEXT:    [[SUM_NEXT_96:%.*]] = add nsw i32 [[VAL_96]], [[SUM_NEXT_95]]
-+; CHECK-NEXT:    [[INDVARS_IV_NEXT_96:%.*]] = add nuw nsw i64 [[INDVARS_IV]], 97
-+; CHECK-NEXT:    [[ARRAYIDX_97:%.*]] = getelementptr inbounds i32, ptr [[ARY]], i64 [[INDVARS_IV_NEXT_96]]
-+; CHECK-NEXT:    [[VAL_97:%.*]] = load i32, ptr [[ARRAYIDX_97]], align 4
-+; CHECK-NEXT:    [[SUM_NEXT_97:%.*]] = add nsw i32 [[VAL_97]], [[SUM_NEXT_96]]
-+; CHECK-NEXT:    [[INDVARS_IV_NEXT_97:%.*]] = add nuw nsw i64 [[INDVARS_IV]], 98
-+; CHECK-NEXT:    [[ARRAYIDX_98:%.*]] = getelementptr inbounds i32, ptr [[ARY]], i64 [[INDVARS_IV_NEXT_97]]
-+; CHECK-NEXT:    [[VAL_98:%.*]] = load i32, ptr [[ARRAYIDX_98]], align 4
-+; CHECK-NEXT:    [[SUM_NEXT_98:%.*]] = add nsw i32 [[VAL_98]], [[SUM_NEXT_97]]
-+; CHECK-NEXT:    [[INDVARS_IV_NEXT_98:%.*]] = add nuw nsw i64 [[INDVARS_IV]], 99
-+; CHECK-NEXT:    [[ARRAYIDX_99:%.*]] = getelementptr inbounds i32, ptr [[ARY]], i64 [[INDVARS_IV_NEXT_98]]
-+; CHECK-NEXT:    [[VAL_99:%.*]] = load i32, ptr [[ARRAYIDX_99]], align 4
-+; CHECK-NEXT:    [[SUM_NEXT_99:%.*]] = add nsw i32 [[VAL_99]], [[SUM_NEXT_98]]
-+; CHECK-NEXT:    [[INDVARS_IV_NEXT_99:%.*]] = add nuw nsw i64 [[INDVARS_IV]], 100
-+; CHECK-NEXT:    [[ARRAYIDX_100:%.*]] = getelementptr inbounds i32, ptr [[ARY]], i64 [[INDVARS_IV_NEXT_99]]
-+; CHECK-NEXT:    [[VAL_100:%.*]] = load i32, ptr [[ARRAYIDX_100]], align 4
-+; CHECK-NEXT:    [[SUM_NEXT_100:%.*]] = add nsw i32 [[VAL_100]], [[SUM_NEXT_99]]
-+; CHECK-NEXT:    [[INDVARS_IV_NEXT_100:%.*]] = add nuw nsw i64 [[INDVARS_IV]], 101
-+; CHECK-NEXT:    [[ARRAYIDX_101:%.*]] = getelementptr inbounds i32, ptr [[ARY]], i64 [[INDVARS_IV_NEXT_100]]
-+; CHECK-NEXT:    [[VAL_101:%.*]] = load i32, ptr [[ARRAYIDX_101]], align 4
-+; CHECK-NEXT:    [[SUM_NEXT_101:%.*]] = add nsw i32 [[VAL_101]], [[SUM_NEXT_100]]
-+; CHECK-NEXT:    [[INDVARS_IV_NEXT_101:%.*]] = add nuw nsw i64 [[INDVARS_IV]], 102
-+; CHECK-NEXT:    [[ARRAYIDX_102:%.*]] = getelementptr inbounds i32, ptr [[ARY]], i64 [[INDVARS_IV_NEXT_101]]
-+; CHECK-NEXT:    [[VAL_102:%.*]] = load i32, ptr [[ARRAYIDX_102]], align 4
-+; CHECK-NEXT:    [[SUM_NEXT_102:%.*]] = add nsw i32 [[VAL_102]], [[SUM_NEXT_101]]
-+; CHECK-NEXT:    [[INDVARS_IV_NEXT_102:%.*]] = add nuw nsw i64 [[INDVARS_IV]], 103
-+; CHECK-NEXT:    [[ARRAYIDX_103:%.*]] = getelementptr inbounds i32, ptr [[ARY]], i64 [[INDVARS_IV_NEXT_102]]
-+; CHECK-NEXT:    [[VAL_103:%.*]] = load i32, ptr [[ARRAYIDX_103]], align 4
-+; CHECK-NEXT:    [[SUM_NEXT_103:%.*]] = add nsw i32 [[VAL_103]], [[SUM_NEXT_102]]
-+; CHECK-NEXT:    [[INDVARS_IV_NEXT_103:%.*]] = add nuw nsw i64 [[INDVARS_IV]], 104
-+; CHECK-NEXT:    [[ARRAYIDX_104:%.*]] = getelementptr inbounds i32, ptr [[ARY]], i64 [[INDVARS_IV_NEXT_103]]
-+; CHECK-NEXT:    [[VAL_104:%.*]] = load i32, ptr [[ARRAYIDX_104]], align 4
-+; CHECK-NEXT:    [[SUM_NEXT_104:%.*]] = add nsw i32 [[VAL_104]], [[SUM_NEXT_103]]
-+; CHECK-NEXT:    [[INDVARS_IV_NEXT_104:%.*]] = add nuw nsw i64 [[INDVARS_IV]], 105
-+; CHECK-NEXT:    [[ARRAYIDX_105:%.*]] = getelementptr inbounds i32, ptr [[ARY]], i64 [[INDVARS_IV_NEXT_104]]
-+; CHECK-NEXT:    [[VAL_105:%.*]] = load i32, ptr [[ARRAYIDX_105]], align 4
-+; CHECK-NEXT:    [[SUM_NEXT_105:%.*]] = add nsw i32 [[VAL_105]], [[SUM_NEXT_104]]
-+; CHECK-NEXT:    [[INDVARS_IV_NEXT_105:%.*]] = add nuw nsw i64 [[INDVARS_IV]], 106
-+; CHECK-NEXT:    [[ARRAYIDX_106:%.*]] = getelementptr inbounds i32, ptr [[ARY]], i64 [[INDVARS_IV_NEXT_105]]
-+; CHECK-NEXT:    [[VAL_106:%.*]] = load i32, ptr [[ARRAYIDX_106]], align 4
-+; CHECK-NEXT:    [[SUM_NEXT_106:%.*]] = add nsw i32 [[VAL_106]], [[SUM_NEXT_105]]
-+; CHECK-NEXT:    [[INDVARS_IV_NEXT_106:%.*]] = add nuw nsw i64 [[INDVARS_IV]], 107
-+; CHECK-NEXT:    [[ARRAYIDX_107:%.*]] = getelementptr inbounds i32, ptr [[ARY]], i64 [[INDVARS_IV_NEXT_106]]
-+; CHECK-NEXT:    [[VAL_107:%.*]] = load i32, ptr [[ARRAYIDX_107]], align 4
-+; CHECK-NEXT:    [[SUM_NEXT_107:%.*]] = add nsw i32 [[VAL_107]], [[SUM_NEXT_106]]
-+; CHECK-NEXT:    [[INDVARS_IV_NEXT_107:%.*]] = add nuw nsw i64 [[INDVARS_IV]], 108
-+; CHECK-NEXT:    [[ARRAYIDX_108:%.*]] = getelementptr inbounds i32, ptr [[ARY]], i64 [[INDVARS_IV_NEXT_107]]
-+; CHECK-NEXT:    [[VAL_108:%.*]] = load i32, ptr [[ARRAYIDX_108]], align 4
-+; CHECK-NEXT:    [[SUM_NEXT_108:%.*]] = add nsw i32 [[VAL_108]], [[SUM_NEXT_107]]
-+; CHECK-NEXT:    [[INDVARS_IV_NEXT_108:%.*]] = add nuw nsw i64 [[INDVARS_IV]], 109
-+; CHECK-NEXT:    [[ARRAYIDX_109:%.*]] = getelementptr inbounds i32, ptr [[ARY]], i64 [[INDVARS_IV_NEXT_108]]
-+; CHECK-NEXT:    [[VAL_109:%.*]] = load i32, ptr [[ARRAYIDX_109]], align 4
-+; CHECK-NEXT:    [[SUM_NEXT_109:%.*]] = add nsw i32 [[VAL_109]], [[SUM_NEXT_108]]
-+; CHECK-NEXT:    [[INDVARS_IV_NEXT_109:%.*]] = add nuw nsw i64 [[INDVARS_IV]], 110
-+; CHECK-NEXT:    [[ARRAYIDX_110:%.*]] = getelementptr inbounds i32, ptr [[ARY]], i64 [[INDVARS_IV_NEXT_109]]
-+; CHECK-NEXT:    [[VAL_110:%.*]] = load i32, ptr [[ARRAYIDX_110]], align 4
-+; CHECK-NEXT:    [[SUM_NEXT_110:%.*]] = add nsw i32 [[VAL_110]], [[SUM_NEXT_109]]
-+; CHECK-NEXT:    [[INDVARS_IV_NEXT_110:%.*]] = add nuw nsw i64 [[INDVARS_IV]], 111
-+; CHECK-NEXT:    [[ARRAYIDX_111:%.*]] = getelementptr inbounds i32, ptr [[ARY]], i64 [[INDVARS_IV_NEXT_110]]
-+; CHECK-NEXT:    [[VAL_111:%.*]] = load i32, ptr [[ARRAYIDX_111]], align 4
-+; CHECK-NEXT:    [[SUM_NEXT_111:%.*]] = add nsw i32 [[VAL_111]], [[SUM_NEXT_110]]
-+; CHECK-NEXT:    [[INDVARS_IV_NEXT_111:%.*]] = add nuw nsw i64 [[INDVARS_IV]], 112
-+; CHECK-NEXT:    [[ARRAYIDX_112:%.*]] = getelementptr inbounds i32, ptr [[ARY]], i64 [[INDVARS_IV_NEXT_111]]
-+; CHECK-NEXT:    [[VAL_112:%.*]] = load i32, ptr [[ARRAYIDX_112]], align 4
-+; CHECK-NEXT:    [[SUM_NEXT_112:%.*]] = add nsw i32 [[VAL_112]], [[SUM_NEXT_111]]
-+; CHECK-NEXT:    [[INDVARS_IV_NEXT_112:%.*]] = add nuw nsw i64 [[INDVARS_IV]], 113
-+; CHECK-NEXT:    [[ARRAYIDX_113:%.*]] = getelementptr inbounds i32, ptr [[ARY]], i64 [[INDVARS_IV_NEXT_112]]
-+; CHECK-NEXT:    [[VAL_113:%.*]] = load i32, ptr [[ARRAYIDX_113]], align 4
-+; CHECK-NEXT:    [[SUM_NEXT_113:%.*]] = add nsw i32 [[VAL_113]], [[SUM_NEXT_112]]
-+; CHECK-NEXT:    [[INDVARS_IV_NEXT_113:%.*]] = add nuw nsw i64 [[INDVARS_IV]], 114
-+; CHECK-NEXT:    [[ARRAYIDX_114:%.*]] = getelementptr inbounds i32, ptr [[ARY]], i64 [[INDVARS_IV_NEXT_113]]
-+; CHECK-NEXT:    [[VAL_114:%.*]] = load i32, ptr [[ARRAYIDX_114]], align 4
-+; CHECK-NEXT:    [[SUM_NEXT_114:%.*]] = add nsw i32 [[VAL_114]], [[SUM_NEXT_113]]
-+; CHECK-NEXT:    [[INDVARS_IV_NEXT_114:%.*]] = add nuw nsw i64 [[INDVARS_IV]], 115
-+; CHECK-NEXT:    [[ARRAYIDX_115:%.*]] = getelementptr inbounds i32, ptr [[ARY]], i64 [[INDVARS_IV_NEXT_114]]
-+; CHECK-NEXT:    [[VAL_115:%.*]] = load i32, ptr [[ARRAYIDX_115]], align 4
-+; CHECK-NEXT:    [[SUM_NEXT_115:%.*]] = add nsw i32 [[VAL_115]], [[SUM_NEXT_114]]
-+; CHECK-NEXT:    [[INDVARS_IV_NEXT_115:%.*]] = add nuw nsw i64 [[INDVARS_IV]], 116
-+; CHECK-NEXT:    [[ARRAYIDX_116:%.*]] = getelementptr inbounds i32, ptr [[ARY]], i64 [[INDVARS_IV_NEXT_115]]
-+; CHECK-NEXT:    [[VAL_116:%.*]] = load i32, ptr [[ARRAYIDX_116]], align 4
-+; CHECK-NEXT:    [[SUM_NEXT_116:%.*]] = add nsw i32 [[VAL_116]], [[SUM_NEXT_115]]
-+; CHECK-NEXT:    [[INDVARS_IV_NEXT_116:%.*]] = add nuw nsw i64 [[INDVARS_IV]], 117
-+; CHECK-NEXT:    [[ARRAYIDX_117:%.*]] = getelementptr inbounds i32, ptr [[ARY]], i64 [[INDVARS_IV_NEXT_116]]
-+; CHECK-NEXT:    [[VAL_117:%.*]] = load i32, ptr [[ARRAYIDX_117]], align 4
-+; CHECK-NEXT:    [[SUM_NEXT_117:%.*]] = add nsw i32 [[VAL_117]], [[SUM_NEXT_116]]
-+; CHECK-NEXT:    [[INDVARS_IV_NEXT_117:%.*]] = add nuw nsw i64 [[INDVARS_IV]], 118
-+; CHECK-NEXT:    [[ARRAYIDX_118:%.*]] = getelementptr inbounds i32, ptr [[ARY]], i64 [[INDVARS_IV_NEXT_117]]
-+; CHECK-NEXT:    [[VAL_118:%.*]] = load i32, ptr [[ARRAYIDX_118]], align 4
-+; CHECK-NEXT:    [[SUM_NEXT_118:%.*]] = add nsw i32 [[VAL_118]], [[SUM_NEXT_117]]
-+; CHECK-NEXT:    [[INDVARS_IV_NEXT_118:%.*]] = add nuw nsw i64 [[INDVARS_IV]], 119
-+; CHECK-NEXT:    [[ARRAYIDX_119:%.*]] = getelementptr inbounds i32, ptr [[ARY]], i64 [[INDVARS_IV_NEXT_118]]
-+; CHECK-NEXT:    [[VAL_119:%.*]] = load i32, ptr [[ARRAYIDX_119]], align 4
-+; CHECK-NEXT:    [[SUM_NEXT_119:%.*]] = add nsw i32 [[VAL_119]], [[SUM_NEXT_118]]
-+; CHECK-NEXT:    [[INDVARS_IV_NEXT_119:%.*]] = add nuw nsw i64 [[INDVARS_IV]], 120
-+; CHECK-NEXT:    [[ARRAYIDX_120:%.*]] = getelementptr inbounds i32, ptr [[ARY]], i64 [[INDVARS_IV_NEXT_119]]
-+; CHECK-NEXT:    [[VAL_120:%.*]] = load i32, ptr [[ARRAYIDX_120]], align 4
-+; CHECK-NEXT:    [[SUM_NEXT_120:%.*]] = add nsw i32 [[VAL_120]], [[SUM_NEXT_119]]
-+; CHECK-NEXT:    [[INDVARS_IV_NEXT_120:%.*]] = add nuw nsw i64 [[INDVARS_IV]], 121
-+; CHECK-NEXT:    [[ARRAYIDX_121:%.*]] = getelementptr inbounds i32, ptr [[ARY]], i64 [[INDVARS_IV_NEXT_120]]
-+; CHECK-NEXT:    [[VAL_121:%.*]] = load i32, ptr [[ARRAYIDX_121]], align 4
-+; CHECK-NEXT:    [[SUM_NEXT_121:%.*]] = add nsw i32 [[VAL_121]], [[SUM_NEXT_120]]
-+; CHECK-NEXT:    [[INDVARS_IV_NEXT_121:%.*]] = add nuw nsw i64 [[INDVARS_IV]], 122
-+; CHECK-NEXT:    [[ARRAYIDX_122:%.*]] = getelementptr inbounds i32, ptr [[ARY]], i64 [[INDVARS_IV_NEXT_121]]
-+; CHECK-NEXT:    [[VAL_122:%.*]] = load i32, ptr [[ARRAYIDX_122]], align 4
-+; CHECK-NEXT:    [[SUM_NEXT_122:%.*]] = add nsw i32 [[VAL_122]], [[SUM_NEXT_121]]
-+; CHECK-NEXT:    [[INDVARS_IV_NEXT_122:%.*]] = add nuw nsw i64 [[INDVARS_IV]], 123
-+; CHECK-NEXT:    [[ARRAYIDX_123:%.*]] = getelementptr inbounds i32, ptr [[ARY]], i64 [[INDVARS_IV_NEXT_122]]
-+; CHECK-NEXT:    [[VAL_123:%.*]] = load i32, ptr [[ARRAYIDX_123]], align 4
-+; CHECK-NEXT:    [[SUM_NEXT_123:%.*]] = add nsw i32 [[VAL_123]], [[SUM_NEXT_122]]
-+; CHECK-NEXT:    [[INDVARS_IV_NEXT_123:%.*]] = add nuw nsw i64 [[INDVARS_IV]], 124
-+; CHECK-NEXT:    [[ARRAYIDX_124:%.*]] = getelementptr inbounds i32, ptr [[ARY]], i64 [[INDVARS_IV_NEXT_123]]
-+; CHECK-NEXT:    [[VAL_124:%.*]] = load i32, ptr [[ARRAYIDX_124]], align 4
-+; CHECK-NEXT:    [[SUM_NEXT_124:%.*]] = add nsw i32 [[VAL_124]], [[SUM_NEXT_123]]
-+; CHECK-NEXT:    [[INDVARS_IV_NEXT_124:%.*]] = add nuw nsw i64 [[INDVARS_IV]], 125
-+; CHECK-NEXT:    [[ARRAYIDX_125:%.*]] = getelementptr inbounds i32, ptr [[ARY]], i64 [[INDVARS_IV_NEXT_124]]
-+; CHECK-NEXT:    [[VAL_125:%.*]] = load i32, ptr [[ARRAYIDX_125]], align 4
-+; CHECK-NEXT:    [[SUM_NEXT_125:%.*]] = add nsw i32 [[VAL_125]], [[SUM_NEXT_124]]
-+; CHECK-NEXT:    [[INDVARS_IV_NEXT_125:%.*]] = add nuw nsw i64 [[INDVARS_IV]], 126
-+; CHECK-NEXT:    [[ARRAYIDX_126:%.*]] = getelementptr inbounds i32, ptr [[ARY]], i64 [[INDVARS_IV_NEXT_125]]
-+; CHECK-NEXT:    [[VAL_126:%.*]] = load i32, ptr [[ARRAYIDX_126]], align 4
-+; CHECK-NEXT:    [[SUM_NEXT_126:%.*]] = add nsw i32 [[VAL_126]], [[SUM_NEXT_125]]
-+; CHECK-NEXT:    [[INDVARS_IV_NEXT_126:%.*]] = add nuw nsw i64 [[INDVARS_IV]], 127
-+; CHECK-NEXT:    [[ARRAYIDX_127:%.*]] = getelementptr inbounds i32, ptr [[ARY]], i64 [[INDVARS_IV_NEXT_126]]
-+; CHECK-NEXT:    [[VAL_127:%.*]] = load i32, ptr [[ARRAYIDX_127]], align 4
-+; CHECK-NEXT:    [[SUM_NEXT_127]] = add nsw i32 [[VAL_127]], [[SUM_NEXT_126]]
-+; CHECK-NEXT:    [[INDVARS_IV_NEXT_127]] = add nuw nsw i64 [[INDVARS_IV]], 128
-+; CHECK-NEXT:    [[EXITCOND_NOT_127:%.*]] = icmp eq i64 [[INDVARS_IV_NEXT_127]], 8192
-+; CHECK-NEXT:    br i1 [[EXITCOND_NOT_127]], label [[FOR_COND_CLEANUP:%.*]], label [[FOR_BODY]]
- ; CHECK:       for.cond.cleanup:
--; CHECK-NEXT:    [[SUM_NEXT_LCSSA:%.*]] = phi i32 [ [[SUM_NEXT_31]], [[FOR_BODY]] ]
-+; CHECK-NEXT:    [[SUM_NEXT_LCSSA:%.*]] = phi i32 [ [[SUM_NEXT_127]], [[FOR_BODY]] ]
- ; CHECK-NEXT:    ret i32 [[SUM_NEXT_LCSSA]]
- ;
- entry:
-@@ -167,16 +551,16 @@
- ; CHECK-SAME: ptr [[ARY:%.*]], i64 [[N:%.*]]) #[[ATTR0]] {
- ; CHECK-NEXT:  entry:
- ; CHECK-NEXT:    [[TMP0:%.*]] = add i64 [[N]], -1
--; CHECK-NEXT:    [[XTRAITER:%.*]] = and i64 [[N]], 1
--; CHECK-NEXT:    [[TMP1:%.*]] = icmp ult i64 [[TMP0]], 1
-+; CHECK-NEXT:    [[XTRAITER:%.*]] = and i64 [[N]], 7
-+; CHECK-NEXT:    [[TMP1:%.*]] = icmp ult i64 [[TMP0]], 7
- ; CHECK-NEXT:    br i1 [[TMP1]], label [[FOR_COND_CLEANUP_UNR_LCSSA:%.*]], label [[ENTRY_NEW:%.*]]
- ; CHECK:       entry.new:
- ; CHECK-NEXT:    [[UNROLL_ITER:%.*]] = sub i64 [[N]], [[XTRAITER]]
- ; CHECK-NEXT:    br label [[FOR_BODY:%.*]]
- ; CHECK:       for.body:
--; CHECK-NEXT:    [[INDVARS_IV:%.*]] = phi i64 [ 0, [[ENTRY_NEW]] ], [ [[INDVARS_IV_NEXT_1:%.*]], [[FOR_BODY]] ]
--; CHECK-NEXT:    [[SUM:%.*]] = phi i32 [ 0, [[ENTRY_NEW]] ], [ [[SUM_NEXT_1:%.*]], [[FOR_BODY]] ]
--; CHECK-NEXT:    [[NITER:%.*]] = phi i64 [ 0, [[ENTRY_NEW]] ], [ [[NITER_NEXT_1:%.*]], [[FOR_BODY]] ]
-+; CHECK-NEXT:    [[INDVARS_IV:%.*]] = phi i64 [ 0, [[ENTRY_NEW]] ], [ [[INDVARS_IV_NEXT_7:%.*]], [[FOR_BODY]] ]
-+; CHECK-NEXT:    [[SUM:%.*]] = phi i32 [ 0, [[ENTRY_NEW]] ], [ [[SUM_NEXT_7:%.*]], [[FOR_BODY]] ]
-+; CHECK-NEXT:    [[NITER:%.*]] = phi i64 [ 0, [[ENTRY_NEW]] ], [ [[NITER_NEXT_7:%.*]], [[FOR_BODY]] ]
- ; CHECK-NEXT:    [[ARRAYIDX:%.*]] = getelementptr inbounds i32, ptr [[ARY]], i64 [[INDVARS_IV]]
- ; CHECK-NEXT:    [[VAL:%.*]] = load i32, ptr [[ARRAYIDX]], align 4
- ; CHECK-NEXT:    [[DUMMY1:%.*]] = mul i32 [[VAL]], [[VAL]]
-@@ -283,15 +667,339 @@
- ; CHECK-NEXT:    [[DUMMY48_1:%.*]] = mul i32 [[DUMMY47_1]], [[DUMMY47_1]]
- ; CHECK-NEXT:    [[DUMMY49_1:%.*]] = mul i32 [[DUMMY48_1]], [[DUMMY48_1]]
- ; CHECK-NEXT:    [[DUMMY50_1:%.*]] = mul i32 [[DUMMY49_1]], [[DUMMY49_1]]
--; CHECK-NEXT:    [[SUM_NEXT_1]] = add nsw i32 [[DUMMY50_1]], [[SUM_NEXT]]
--; CHECK-NEXT:    [[INDVARS_IV_NEXT_1]] = add nuw nsw i64 [[INDVARS_IV]], 2
--; CHECK-NEXT:    [[NITER_NEXT_1]] = add i64 [[NITER]], 2
--; CHECK-NEXT:    [[NITER_NCMP_1:%.*]] = icmp eq i64 [[NITER_NEXT_1]], [[UNROLL_ITER]]
--; CHECK-NEXT:    br i1 [[NITER_NCMP_1]], label [[FOR_COND_CLEANUP_UNR_LCSSA_LOOPEXIT:%.*]], label [[FOR_BODY]]
-+; CHECK-NEXT:    [[SUM_NEXT_1:%.*]] = add nsw i32 [[DUMMY50_1]], [[SUM_NEXT]]
-+; CHECK-NEXT:    [[INDVARS_IV_NEXT_1:%.*]] = add nuw nsw i64 [[INDVARS_IV]], 2
-+; CHECK-NEXT:    [[ARRAYIDX_2:%.*]] = getelementptr inbounds i32, ptr [[ARY]], i64 [[INDVARS_IV_NEXT_1]]
-+; CHECK-NEXT:    [[VAL_2:%.*]] = load i32, ptr [[ARRAYIDX_2]], align 4
-+; CHECK-NEXT:    [[DUMMY1_2:%.*]] = mul i32 [[VAL_2]], [[VAL_2]]
-+; CHECK-NEXT:    [[DUMMY2_2:%.*]] = mul i32 [[DUMMY1_2]], [[DUMMY1_2]]
-+; CHECK-NEXT:    [[DUMMY3_2:%.*]] = mul i32 [[DUMMY2_2]], [[DUMMY2_2]]
-+; CHECK-NEXT:    [[DUMMY4_2:%.*]] = mul i32 [[DUMMY3_2]], [[DUMMY3_2]]
-+; CHECK-NEXT:    [[DUMMY5_2:%.*]] = mul i32 [[DUMMY4_2]], [[DUMMY4_2]]
-+; CHECK-NEXT:    [[DUMMY6_2:%.*]] = mul i32 [[DUMMY5_2]], [[DUMMY5_2]]
-+; CHECK-NEXT:    [[DUMMY7_2:%.*]] = mul i32 [[DUMMY6_2]], [[DUMMY6_2]]
-+; CHECK-NEXT:    [[DUMMY8_2:%.*]] = mul i32 [[DUMMY7_2]], [[DUMMY7_2]]
-+; CHECK-NEXT:    [[DUMMY9_2:%.*]] = mul i32 [[DUMMY8_2]], [[DUMMY8_2]]
-+; CHECK-NEXT:    [[DUMMY10_2:%.*]] = mul i32 [[DUMMY9_2]], [[DUMMY9_2]]
-+; CHECK-NEXT:    [[DUMMY11_2:%.*]] = mul i32 [[DUMMY10_2]], [[DUMMY10_2]]
-+; CHECK-NEXT:    [[DUMMY12_2:%.*]] = mul i32 [[DUMMY11_2]], [[DUMMY11_2]]
-+; CHECK-NEXT:    [[DUMMY13_2:%.*]] = mul i32 [[DUMMY12_2]], [[DUMMY12_2]]
-+; CHECK-NEXT:    [[DUMMY14_2:%.*]] = mul i32 [[DUMMY13_2]], [[DUMMY13_2]]
-+; CHECK-NEXT:    [[DUMMY15_2:%.*]] = mul i32 [[DUMMY14_2]], [[DUMMY14_2]]
-+; CHECK-NEXT:    [[DUMMY16_2:%.*]] = mul i32 [[DUMMY15_2]], [[DUMMY15_2]]
-+; CHECK-NEXT:    [[DUMMY17_2:%.*]] = mul i32 [[DUMMY16_2]], [[DUMMY16_2]]
-+; CHECK-NEXT:    [[DUMMY18_2:%.*]] = mul i32 [[DUMMY17_2]], [[DUMMY17_2]]
-+; CHECK-NEXT:    [[DUMMY19_2:%.*]] = mul i32 [[DUMMY18_2]], [[DUMMY18_2]]
-+; CHECK-NEXT:    [[DUMMY20_2:%.*]] = mul i32 [[DUMMY19_2]], [[DUMMY19_2]]
-+; CHECK-NEXT:    [[DUMMY21_2:%.*]] = mul i32 [[DUMMY20_2]], [[DUMMY20_2]]
-+; CHECK-NEXT:    [[DUMMY22_2:%.*]] = mul i32 [[DUMMY21_2]], [[DUMMY21_2]]
-+; CHECK-NEXT:    [[DUMMY23_2:%.*]] = mul i32 [[DUMMY22_2]], [[DUMMY22_2]]
-+; CHECK-NEXT:    [[DUMMY24_2:%.*]] = mul i32 [[DUMMY23_2]], [[DUMMY23_2]]
-+; CHECK-NEXT:    [[DUMMY25_2:%.*]] = mul i32 [[DUMMY24_2]], [[DUMMY24_2]]
-+; CHECK-NEXT:    [[DUMMY26_2:%.*]] = mul i32 [[DUMMY25_2]], [[DUMMY25_2]]
-+; CHECK-NEXT:    [[DUMMY27_2:%.*]] = mul i32 [[DUMMY26_2]], [[DUMMY26_2]]
-+; CHECK-NEXT:    [[DUMMY28_2:%.*]] = mul i32 [[DUMMY27_2]], [[DUMMY27_2]]
-+; CHECK-NEXT:    [[DUMMY29_2:%.*]] = mul i32 [[DUMMY28_2]], [[DUMMY28_2]]
-+; CHECK-NEXT:    [[DUMMY30_2:%.*]] = mul i32 [[DUMMY29_2]], [[DUMMY29_2]]
-+; CHECK-NEXT:    [[DUMMY31_2:%.*]] = mul i32 [[DUMMY30_2]], [[DUMMY30_2]]
-+; CHECK-NEXT:    [[DUMMY32_2:%.*]] = mul i32 [[DUMMY31_2]], [[DUMMY31_2]]
-+; CHECK-NEXT:    [[DUMMY33_2:%.*]] = mul i32 [[DUMMY32_2]], [[DUMMY32_2]]
-+; CHECK-NEXT:    [[DUMMY34_2:%.*]] = mul i32 [[DUMMY33_2]], [[DUMMY33_2]]
-+; CHECK-NEXT:    [[DUMMY35_2:%.*]] = mul i32 [[DUMMY34_2]], [[DUMMY34_2]]
-+; CHECK-NEXT:    [[DUMMY36_2:%.*]] = mul i32 [[DUMMY35_2]], [[DUMMY35_2]]
-+; CHECK-NEXT:    [[DUMMY37_2:%.*]] = mul i32 [[DUMMY36_2]], [[DUMMY36_2]]
-+; CHECK-NEXT:    [[DUMMY38_2:%.*]] = mul i32 [[DUMMY37_2]], [[DUMMY37_2]]
-+; CHECK-NEXT:    [[DUMMY39_2:%.*]] = mul i32 [[DUMMY38_2]], [[DUMMY38_2]]
-+; CHECK-NEXT:    [[DUMMY40_2:%.*]] = mul i32 [[DUMMY39_2]], [[DUMMY39_2]]
-+; CHECK-NEXT:    [[DUMMY41_2:%.*]] = mul i32 [[DUMMY40_2]], [[DUMMY40_2]]
-+; CHECK-NEXT:    [[DUMMY42_2:%.*]] = mul i32 [[DUMMY41_2]], [[DUMMY41_2]]
-+; CHECK-NEXT:    [[DUMMY43_2:%.*]] = mul i32 [[DUMMY42_2]], [[DUMMY42_2]]
-+; CHECK-NEXT:    [[DUMMY44_2:%.*]] = mul i32 [[DUMMY43_2]], [[DUMMY43_2]]
-+; CHECK-NEXT:    [[DUMMY45_2:%.*]] = mul i32 [[DUMMY44_2]], [[DUMMY44_2]]
-+; CHECK-NEXT:    [[DUMMY46_2:%.*]] = mul i32 [[DUMMY45_2]], [[DUMMY45_2]]
-+; CHECK-NEXT:    [[DUMMY47_2:%.*]] = mul i32 [[DUMMY46_2]], [[DUMMY46_2]]
-+; CHECK-NEXT:    [[DUMMY48_2:%.*]] = mul i32 [[DUMMY47_2]], [[DUMMY47_2]]
-+; CHECK-NEXT:    [[DUMMY49_2:%.*]] = mul i32 [[DUMMY48_2]], [[DUMMY48_2]]
-+; CHECK-NEXT:    [[DUMMY50_2:%.*]] = mul i32 [[DUMMY49_2]], [[DUMMY49_2]]
-+; CHECK-NEXT:    [[SUM_NEXT_2:%.*]] = add nsw i32 [[DUMMY50_2]], [[SUM_NEXT_1]]
-+; CHECK-NEXT:    [[INDVARS_IV_NEXT_2:%.*]] = add nuw nsw i64 [[INDVARS_IV]], 3
-+; CHECK-NEXT:    [[ARRAYIDX_3:%.*]] = getelementptr inbounds i32, ptr [[ARY]], i64 [[INDVARS_IV_NEXT_2]]
-+; CHECK-NEXT:    [[VAL_3:%.*]] = load i32, ptr [[ARRAYIDX_3]], align 4
-+; CHECK-NEXT:    [[DUMMY1_3:%.*]] = mul i32 [[VAL_3]], [[VAL_3]]
-+; CHECK-NEXT:    [[DUMMY2_3:%.*]] = mul i32 [[DUMMY1_3]], [[DUMMY1_3]]
-+; CHECK-NEXT:    [[DUMMY3_3:%.*]] = mul i32 [[DUMMY2_3]], [[DUMMY2_3]]
-+; CHECK-NEXT:    [[DUMMY4_3:%.*]] = mul i32 [[DUMMY3_3]], [[DUMMY3_3]]
-+; CHECK-NEXT:    [[DUMMY5_3:%.*]] = mul i32 [[DUMMY4_3]], [[DUMMY4_3]]
-+; CHECK-NEXT:    [[DUMMY6_3:%.*]] = mul i32 [[DUMMY5_3]], [[DUMMY5_3]]
-+; CHECK-NEXT:    [[DUMMY7_3:%.*]] = mul i32 [[DUMMY6_3]], [[DUMMY6_3]]
-+; CHECK-NEXT:    [[DUMMY8_3:%.*]] = mul i32 [[DUMMY7_3]], [[DUMMY7_3]]
-+; CHECK-NEXT:    [[DUMMY9_3:%.*]] = mul i32 [[DUMMY8_3]], [[DUMMY8_3]]
-+; CHECK-NEXT:    [[DUMMY10_3:%.*]] = mul i32 [[DUMMY9_3]], [[DUMMY9_3]]
-+; CHECK-NEXT:    [[DUMMY11_3:%.*]] = mul i32 [[DUMMY10_3]], [[DUMMY10_3]]
-+; CHECK-NEXT:    [[DUMMY12_3:%.*]] = mul i32 [[DUMMY11_3]], [[DUMMY11_3]]
-+; CHECK-NEXT:    [[DUMMY13_3:%.*]] = mul i32 [[DUMMY12_3]], [[DUMMY12_3]]
-+; CHECK-NEXT:    [[DUMMY14_3:%.*]] = mul i32 [[DUMMY13_3]], [[DUMMY13_3]]
-+; CHECK-NEXT:    [[DUMMY15_3:%.*]] = mul i32 [[DUMMY14_3]], [[DUMMY14_3]]
-+; CHECK-NEXT:    [[DUMMY16_3:%.*]] = mul i32 [[DUMMY15_3]], [[DUMMY15_3]]
-+; CHECK-NEXT:    [[DUMMY17_3:%.*]] = mul i32 [[DUMMY16_3]], [[DUMMY16_3]]
-+; CHECK-NEXT:    [[DUMMY18_3:%.*]] = mul i32 [[DUMMY17_3]], [[DUMMY17_3]]
-+; CHECK-NEXT:    [[DUMMY19_3:%.*]] = mul i32 [[DUMMY18_3]], [[DUMMY18_3]]
-+; CHECK-NEXT:    [[DUMMY20_3:%.*]] = mul i32 [[DUMMY19_3]], [[DUMMY19_3]]
-+; CHECK-NEXT:    [[DUMMY21_3:%.*]] = mul i32 [[DUMMY20_3]], [[DUMMY20_3]]
-+; CHECK-NEXT:    [[DUMMY22_3:%.*]] = mul i32 [[DUMMY21_3]], [[DUMMY21_3]]
-+; CHECK-NEXT:    [[DUMMY23_3:%.*]] = mul i32 [[DUMMY22_3]], [[DUMMY22_3]]
-+; CHECK-NEXT:    [[DUMMY24_3:%.*]] = mul i32 [[DUMMY23_3]], [[DUMMY23_3]]
-+; CHECK-NEXT:    [[DUMMY25_3:%.*]] = mul i32 [[DUMMY24_3]], [[DUMMY24_3]]
-+; CHECK-NEXT:    [[DUMMY26_3:%.*]] = mul i32 [[DUMMY25_3]], [[DUMMY25_3]]
-+; CHECK-NEXT:    [[DUMMY27_3:%.*]] = mul i32 [[DUMMY26_3]], [[DUMMY26_3]]
-+; CHECK-NEXT:    [[DUMMY28_3:%.*]] = mul i32 [[DUMMY27_3]], [[DUMMY27_3]]
-+; CHECK-NEXT:    [[DUMMY29_3:%.*]] = mul i32 [[DUMMY28_3]], [[DUMMY28_3]]
-+; CHECK-NEXT:    [[DUMMY30_3:%.*]] = mul i32 [[DUMMY29_3]], [[DUMMY29_3]]
-+; CHECK-NEXT:    [[DUMMY31_3:%.*]] = mul i32 [[DUMMY30_3]], [[DUMMY30_3]]
-+; CHECK-NEXT:    [[DUMMY32_3:%.*]] = mul i32 [[DUMMY31_3]], [[DUMMY31_3]]
-+; CHECK-NEXT:    [[DUMMY33_3:%.*]] = mul i32 [[DUMMY32_3]], [[DUMMY32_3]]
-+; CHECK-NEXT:    [[DUMMY34_3:%.*]] = mul i32 [[DUMMY33_3]], [[DUMMY33_3]]
-+; CHECK-NEXT:    [[DUMMY35_3:%.*]] = mul i32 [[DUMMY34_3]], [[DUMMY34_3]]
-+; CHECK-NEXT:    [[DUMMY36_3:%.*]] = mul i32 [[DUMMY35_3]], [[DUMMY35_3]]
-+; CHECK-NEXT:    [[DUMMY37_3:%.*]] = mul i32 [[DUMMY36_3]], [[DUMMY36_3]]
-+; CHECK-NEXT:    [[DUMMY38_3:%.*]] = mul i32 [[DUMMY37_3]], [[DUMMY37_3]]
-+; CHECK-NEXT:    [[DUMMY39_3:%.*]] = mul i32 [[DUMMY38_3]], [[DUMMY38_3]]
-+; CHECK-NEXT:    [[DUMMY40_3:%.*]] = mul i32 [[DUMMY39_3]], [[DUMMY39_3]]
-+; CHECK-NEXT:    [[DUMMY41_3:%.*]] = mul i32 [[DUMMY40_3]], [[DUMMY40_3]]
-+; CHECK-NEXT:    [[DUMMY42_3:%.*]] = mul i32 [[DUMMY41_3]], [[DUMMY41_3]]
-+; CHECK-NEXT:    [[DUMMY43_3:%.*]] = mul i32 [[DUMMY42_3]], [[DUMMY42_3]]
-+; CHECK-NEXT:    [[DUMMY44_3:%.*]] = mul i32 [[DUMMY43_3]], [[DUMMY43_3]]
-+; CHECK-NEXT:    [[DUMMY45_3:%.*]] = mul i32 [[DUMMY44_3]], [[DUMMY44_3]]
-+; CHECK-NEXT:    [[DUMMY46_3:%.*]] = mul i32 [[DUMMY45_3]], [[DUMMY45_3]]
-+; CHECK-NEXT:    [[DUMMY47_3:%.*]] = mul i32 [[DUMMY46_3]], [[DUMMY46_3]]
-+; CHECK-NEXT:    [[DUMMY48_3:%.*]] = mul i32 [[DUMMY47_3]], [[DUMMY47_3]]
-+; CHECK-NEXT:    [[DUMMY49_3:%.*]] = mul i32 [[DUMMY48_3]], [[DUMMY48_3]]
-+; CHECK-NEXT:    [[DUMMY50_3:%.*]] = mul i32 [[DUMMY49_3]], [[DUMMY49_3]]
-+; CHECK-NEXT:    [[SUM_NEXT_3:%.*]] = add nsw i32 [[DUMMY50_3]], [[SUM_NEXT_2]]
-+; CHECK-NEXT:    [[INDVARS_IV_NEXT_3:%.*]] = add nuw nsw i64 [[INDVARS_IV]], 4
-+; CHECK-NEXT:    [[ARRAYIDX_4:%.*]] = getelementptr inbounds i32, ptr [[ARY]], i64 [[INDVARS_IV_NEXT_3]]
-+; CHECK-NEXT:    [[VAL_4:%.*]] = load i32, ptr [[ARRAYIDX_4]], align 4
-+; CHECK-NEXT:    [[DUMMY1_4:%.*]] = mul i32 [[VAL_4]], [[VAL_4]]
-+; CHECK-NEXT:    [[DUMMY2_4:%.*]] = mul i32 [[DUMMY1_4]], [[DUMMY1_4]]
-+; CHECK-NEXT:    [[DUMMY3_4:%.*]] = mul i32 [[DUMMY2_4]], [[DUMMY2_4]]
-+; CHECK-NEXT:    [[DUMMY4_4:%.*]] = mul i32 [[DUMMY3_4]], [[DUMMY3_4]]
-+; CHECK-NEXT:    [[DUMMY5_4:%.*]] = mul i32 [[DUMMY4_4]], [[DUMMY4_4]]
-+; CHECK-NEXT:    [[DUMMY6_4:%.*]] = mul i32 [[DUMMY5_4]], [[DUMMY5_4]]
-+; CHECK-NEXT:    [[DUMMY7_4:%.*]] = mul i32 [[DUMMY6_4]], [[DUMMY6_4]]
-+; CHECK-NEXT:    [[DUMMY8_4:%.*]] = mul i32 [[DUMMY7_4]], [[DUMMY7_4]]
-+; CHECK-NEXT:    [[DUMMY9_4:%.*]] = mul i32 [[DUMMY8_4]], [[DUMMY8_4]]
-+; CHECK-NEXT:    [[DUMMY10_4:%.*]] = mul i32 [[DUMMY9_4]], [[DUMMY9_4]]
-+; CHECK-NEXT:    [[DUMMY11_4:%.*]] = mul i32 [[DUMMY10_4]], [[DUMMY10_4]]
-+; CHECK-NEXT:    [[DUMMY12_4:%.*]] = mul i32 [[DUMMY11_4]], [[DUMMY11_4]]
-+; CHECK-NEXT:    [[DUMMY13_4:%.*]] = mul i32 [[DUMMY12_4]], [[DUMMY12_4]]
-+; CHECK-NEXT:    [[DUMMY14_4:%.*]] = mul i32 [[DUMMY13_4]], [[DUMMY13_4]]
-+; CHECK-NEXT:    [[DUMMY15_4:%.*]] = mul i32 [[DUMMY14_4]], [[DUMMY14_4]]
-+; CHECK-NEXT:    [[DUMMY16_4:%.*]] = mul i32 [[DUMMY15_4]], [[DUMMY15_4]]
-+; CHECK-NEXT:    [[DUMMY17_4:%.*]] = mul i32 [[DUMMY16_4]], [[DUMMY16_4]]
-+; CHECK-NEXT:    [[DUMMY18_4:%.*]] = mul i32 [[DUMMY17_4]], [[DUMMY17_4]]
-+; CHECK-NEXT:    [[DUMMY19_4:%.*]] = mul i32 [[DUMMY18_4]], [[DUMMY18_4]]
-+; CHECK-NEXT:    [[DUMMY20_4:%.*]] = mul i32 [[DUMMY19_4]], [[DUMMY19_4]]
-+; CHECK-NEXT:    [[DUMMY21_4:%.*]] = mul i32 [[DUMMY20_4]], [[DUMMY20_4]]
-+; CHECK-NEXT:    [[DUMMY22_4:%.*]] = mul i32 [[DUMMY21_4]], [[DUMMY21_4]]
-+; CHECK-NEXT:    [[DUMMY23_4:%.*]] = mul i32 [[DUMMY22_4]], [[DUMMY22_4]]
-+; CHECK-NEXT:    [[DUMMY24_4:%.*]] = mul i32 [[DUMMY23_4]], [[DUMMY23_4]]
-+; CHECK-NEXT:    [[DUMMY25_4:%.*]] = mul i32 [[DUMMY24_4]], [[DUMMY24_4]]
-+; CHECK-NEXT:    [[DUMMY26_4:%.*]] = mul i32 [[DUMMY25_4]], [[DUMMY25_4]]
-+; CHECK-NEXT:    [[DUMMY27_4:%.*]] = mul i32 [[DUMMY26_4]], [[DUMMY26_4]]
-+; CHECK-NEXT:    [[DUMMY28_4:%.*]] = mul i32 [[DUMMY27_4]], [[DUMMY27_4]]
-+; CHECK-NEXT:    [[DUMMY29_4:%.*]] = mul i32 [[DUMMY28_4]], [[DUMMY28_4]]
-+; CHECK-NEXT:    [[DUMMY30_4:%.*]] = mul i32 [[DUMMY29_4]], [[DUMMY29_4]]
-+; CHECK-NEXT:    [[DUMMY31_4:%.*]] = mul i32 [[DUMMY30_4]], [[DUMMY30_4]]
-+; CHECK-NEXT:    [[DUMMY32_4:%.*]] = mul i32 [[DUMMY31_4]], [[DUMMY31_4]]
-+; CHECK-NEXT:    [[DUMMY33_4:%.*]] = mul i32 [[DUMMY32_4]], [[DUMMY32_4]]
-+; CHECK-NEXT:    [[DUMMY34_4:%.*]] = mul i32 [[DUMMY33_4]], [[DUMMY33_4]]
-+; CHECK-NEXT:    [[DUMMY35_4:%.*]] = mul i32 [[DUMMY34_4]], [[DUMMY34_4]]
-+; CHECK-NEXT:    [[DUMMY36_4:%.*]] = mul i32 [[DUMMY35_4]], [[DUMMY35_4]]
-+; CHECK-NEXT:    [[DUMMY37_4:%.*]] = mul i32 [[DUMMY36_4]], [[DUMMY36_4]]
-+; CHECK-NEXT:    [[DUMMY38_4:%.*]] = mul i32 [[DUMMY37_4]], [[DUMMY37_4]]
-+; CHECK-NEXT:    [[DUMMY39_4:%.*]] = mul i32 [[DUMMY38_4]], [[DUMMY38_4]]
-+; CHECK-NEXT:    [[DUMMY40_4:%.*]] = mul i32 [[DUMMY39_4]], [[DUMMY39_4]]
-+; CHECK-NEXT:    [[DUMMY41_4:%.*]] = mul i32 [[DUMMY40_4]], [[DUMMY40_4]]
-+; CHECK-NEXT:    [[DUMMY42_4:%.*]] = mul i32 [[DUMMY41_4]], [[DUMMY41_4]]
-+; CHECK-NEXT:    [[DUMMY43_4:%.*]] = mul i32 [[DUMMY42_4]], [[DUMMY42_4]]
-+; CHECK-NEXT:    [[DUMMY44_4:%.*]] = mul i32 [[DUMMY43_4]], [[DUMMY43_4]]
-+; CHECK-NEXT:    [[DUMMY45_4:%.*]] = mul i32 [[DUMMY44_4]], [[DUMMY44_4]]
-+; CHECK-NEXT:    [[DUMMY46_4:%.*]] = mul i32 [[DUMMY45_4]], [[DUMMY45_4]]
-+; CHECK-NEXT:    [[DUMMY47_4:%.*]] = mul i32 [[DUMMY46_4]], [[DUMMY46_4]]
-+; CHECK-NEXT:    [[DUMMY48_4:%.*]] = mul i32 [[DUMMY47_4]], [[DUMMY47_4]]
-+; CHECK-NEXT:    [[DUMMY49_4:%.*]] = mul i32 [[DUMMY48_4]], [[DUMMY48_4]]
-+; CHECK-NEXT:    [[DUMMY50_4:%.*]] = mul i32 [[DUMMY49_4]], [[DUMMY49_4]]
-+; CHECK-NEXT:    [[SUM_NEXT_4:%.*]] = add nsw i32 [[DUMMY50_4]], [[SUM_NEXT_3]]
-+; CHECK-NEXT:    [[INDVARS_IV_NEXT_4:%.*]] = add nuw nsw i64 [[INDVARS_IV]], 5
-+; CHECK-NEXT:    [[ARRAYIDX_5:%.*]] = getelementptr inbounds i32, ptr [[ARY]], i64 [[INDVARS_IV_NEXT_4]]
-+; CHECK-NEXT:    [[VAL_5:%.*]] = load i32, ptr [[ARRAYIDX_5]], align 4
-+; CHECK-NEXT:    [[DUMMY1_5:%.*]] = mul i32 [[VAL_5]], [[VAL_5]]
-+; CHECK-NEXT:    [[DUMMY2_5:%.*]] = mul i32 [[DUMMY1_5]], [[DUMMY1_5]]
-+; CHECK-NEXT:    [[DUMMY3_5:%.*]] = mul i32 [[DUMMY2_5]], [[DUMMY2_5]]
-+; CHECK-NEXT:    [[DUMMY4_5:%.*]] = mul i32 [[DUMMY3_5]], [[DUMMY3_5]]
-+; CHECK-NEXT:    [[DUMMY5_5:%.*]] = mul i32 [[DUMMY4_5]], [[DUMMY4_5]]
-+; CHECK-NEXT:    [[DUMMY6_5:%.*]] = mul i32 [[DUMMY5_5]], [[DUMMY5_5]]
-+; CHECK-NEXT:    [[DUMMY7_5:%.*]] = mul i32 [[DUMMY6_5]], [[DUMMY6_5]]
-+; CHECK-NEXT:    [[DUMMY8_5:%.*]] = mul i32 [[DUMMY7_5]], [[DUMMY7_5]]
-+; CHECK-NEXT:    [[DUMMY9_5:%.*]] = mul i32 [[DUMMY8_5]], [[DUMMY8_5]]
-+; CHECK-NEXT:    [[DUMMY10_5:%.*]] = mul i32 [[DUMMY9_5]], [[DUMMY9_5]]
-+; CHECK-NEXT:    [[DUMMY11_5:%.*]] = mul i32 [[DUMMY10_5]], [[DUMMY10_5]]
-+; CHECK-NEXT:    [[DUMMY12_5:%.*]] = mul i32 [[DUMMY11_5]], [[DUMMY11_5]]
-+; CHECK-NEXT:    [[DUMMY13_5:%.*]] = mul i32 [[DUMMY12_5]], [[DUMMY12_5]]
-+; CHECK-NEXT:    [[DUMMY14_5:%.*]] = mul i32 [[DUMMY13_5]], [[DUMMY13_5]]
-+; CHECK-NEXT:    [[DUMMY15_5:%.*]] = mul i32 [[DUMMY14_5]], [[DUMMY14_5]]
-+; CHECK-NEXT:    [[DUMMY16_5:%.*]] = mul i32 [[DUMMY15_5]], [[DUMMY15_5]]
-+; CHECK-NEXT:    [[DUMMY17_5:%.*]] = mul i32 [[DUMMY16_5]], [[DUMMY16_5]]
-+; CHECK-NEXT:    [[DUMMY18_5:%.*]] = mul i32 [[DUMMY17_5]], [[DUMMY17_5]]
-+; CHECK-NEXT:    [[DUMMY19_5:%.*]] = mul i32 [[DUMMY18_5]], [[DUMMY18_5]]
-+; CHECK-NEXT:    [[DUMMY20_5:%.*]] = mul i32 [[DUMMY19_5]], [[DUMMY19_5]]
-+; CHECK-NEXT:    [[DUMMY21_5:%.*]] = mul i32 [[DUMMY20_5]], [[DUMMY20_5]]
-+; CHECK-NEXT:    [[DUMMY22_5:%.*]] = mul i32 [[DUMMY21_5]], [[DUMMY21_5]]
-+; CHECK-NEXT:    [[DUMMY23_5:%.*]] = mul i32 [[DUMMY22_5]], [[DUMMY22_5]]
-+; CHECK-NEXT:    [[DUMMY24_5:%.*]] = mul i32 [[DUMMY23_5]], [[DUMMY23_5]]
-+; CHECK-NEXT:    [[DUMMY25_5:%.*]] = mul i32 [[DUMMY24_5]], [[DUMMY24_5]]
-+; CHECK-NEXT:    [[DUMMY26_5:%.*]] = mul i32 [[DUMMY25_5]], [[DUMMY25_5]]
-+; CHECK-NEXT:    [[DUMMY27_5:%.*]] = mul i32 [[DUMMY26_5]], [[DUMMY26_5]]
-+; CHECK-NEXT:    [[DUMMY28_5:%.*]] = mul i32 [[DUMMY27_5]], [[DUMMY27_5]]
-+; CHECK-NEXT:    [[DUMMY29_5:%.*]] = mul i32 [[DUMMY28_5]], [[DUMMY28_5]]
-+; CHECK-NEXT:    [[DUMMY30_5:%.*]] = mul i32 [[DUMMY29_5]], [[DUMMY29_5]]
-+; CHECK-NEXT:    [[DUMMY31_5:%.*]] = mul i32 [[DUMMY30_5]], [[DUMMY30_5]]
-+; CHECK-NEXT:    [[DUMMY32_5:%.*]] = mul i32 [[DUMMY31_5]], [[DUMMY31_5]]
-+; CHECK-NEXT:    [[DUMMY33_5:%.*]] = mul i32 [[DUMMY32_5]], [[DUMMY32_5]]
-+; CHECK-NEXT:    [[DUMMY34_5:%.*]] = mul i32 [[DUMMY33_5]], [[DUMMY33_5]]
-+; CHECK-NEXT:    [[DUMMY35_5:%.*]] = mul i32 [[DUMMY34_5]], [[DUMMY34_5]]
-+; CHECK-NEXT:    [[DUMMY36_5:%.*]] = mul i32 [[DUMMY35_5]], [[DUMMY35_5]]
-+; CHECK-NEXT:    [[DUMMY37_5:%.*]] = mul i32 [[DUMMY36_5]], [[DUMMY36_5]]
-+; CHECK-NEXT:    [[DUMMY38_5:%.*]] = mul i32 [[DUMMY37_5]], [[DUMMY37_5]]
-+; CHECK-NEXT:    [[DUMMY39_5:%.*]] = mul i32 [[DUMMY38_5]], [[DUMMY38_5]]
-+; CHECK-NEXT:    [[DUMMY40_5:%.*]] = mul i32 [[DUMMY39_5]], [[DUMMY39_5]]
-+; CHECK-NEXT:    [[DUMMY41_5:%.*]] = mul i32 [[DUMMY40_5]], [[DUMMY40_5]]
-+; CHECK-NEXT:    [[DUMMY42_5:%.*]] = mul i32 [[DUMMY41_5]], [[DUMMY41_5]]
-+; CHECK-NEXT:    [[DUMMY43_5:%.*]] = mul i32 [[DUMMY42_5]], [[DUMMY42_5]]
-+; CHECK-NEXT:    [[DUMMY44_5:%.*]] = mul i32 [[DUMMY43_5]], [[DUMMY43_5]]
-+; CHECK-NEXT:    [[DUMMY45_5:%.*]] = mul i32 [[DUMMY44_5]], [[DUMMY44_5]]
-+; CHECK-NEXT:    [[DUMMY46_5:%.*]] = mul i32 [[DUMMY45_5]], [[DUMMY45_5]]
-+; CHECK-NEXT:    [[DUMMY47_5:%.*]] = mul i32 [[DUMMY46_5]], [[DUMMY46_5]]
-+; CHECK-NEXT:    [[DUMMY48_5:%.*]] = mul i32 [[DUMMY47_5]], [[DUMMY47_5]]
-+; CHECK-NEXT:    [[DUMMY49_5:%.*]] = mul i32 [[DUMMY48_5]], [[DUMMY48_5]]
-+; CHECK-NEXT:    [[DUMMY50_5:%.*]] = mul i32 [[DUMMY49_5]], [[DUMMY49_5]]
-+; CHECK-NEXT:    [[SUM_NEXT_5:%.*]] = add nsw i32 [[DUMMY50_5]], [[SUM_NEXT_4]]
-+; CHECK-NEXT:    [[INDVARS_IV_NEXT_5:%.*]] = add nuw nsw i64 [[INDVARS_IV]], 6
-+; CHECK-NEXT:    [[ARRAYIDX_6:%.*]] = getelementptr inbounds i32, ptr [[ARY]], i64 [[INDVARS_IV_NEXT_5]]
-+; CHECK-NEXT:    [[VAL_6:%.*]] = load i32, ptr [[ARRAYIDX_6]], align 4
-+; CHECK-NEXT:    [[DUMMY1_6:%.*]] = mul i32 [[VAL_6]], [[VAL_6]]
-+; CHECK-NEXT:    [[DUMMY2_6:%.*]] = mul i32 [[DUMMY1_6]], [[DUMMY1_6]]
-+; CHECK-NEXT:    [[DUMMY3_6:%.*]] = mul i32 [[DUMMY2_6]], [[DUMMY2_6]]
-+; CHECK-NEXT:    [[DUMMY4_6:%.*]] = mul i32 [[DUMMY3_6]], [[DUMMY3_6]]
-+; CHECK-NEXT:    [[DUMMY5_6:%.*]] = mul i32 [[DUMMY4_6]], [[DUMMY4_6]]
-+; CHECK-NEXT:    [[DUMMY6_6:%.*]] = mul i32 [[DUMMY5_6]], [[DUMMY5_6]]
-+; CHECK-NEXT:    [[DUMMY7_6:%.*]] = mul i32 [[DUMMY6_6]], [[DUMMY6_6]]
-+; CHECK-NEXT:    [[DUMMY8_6:%.*]] = mul i32 [[DUMMY7_6]], [[DUMMY7_6]]
-+; CHECK-NEXT:    [[DUMMY9_6:%.*]] = mul i32 [[DUMMY8_6]], [[DUMMY8_6]]
-+; CHECK-NEXT:    [[DUMMY10_6:%.*]] = mul i32 [[DUMMY9_6]], [[DUMMY9_6]]
-+; CHECK-NEXT:    [[DUMMY11_6:%.*]] = mul i32 [[DUMMY10_6]], [[DUMMY10_6]]
-+; CHECK-NEXT:    [[DUMMY12_6:%.*]] = mul i32 [[DUMMY11_6]], [[DUMMY11_6]]
-+; CHECK-NEXT:    [[DUMMY13_6:%.*]] = mul i32 [[DUMMY12_6]], [[DUMMY12_6]]
-+; CHECK-NEXT:    [[DUMMY14_6:%.*]] = mul i32 [[DUMMY13_6]], [[DUMMY13_6]]
-+; CHECK-NEXT:    [[DUMMY15_6:%.*]] = mul i32 [[DUMMY14_6]], [[DUMMY14_6]]
-+; CHECK-NEXT:    [[DUMMY16_6:%.*]] = mul i32 [[DUMMY15_6]], [[DUMMY15_6]]
-+; CHECK-NEXT:    [[DUMMY17_6:%.*]] = mul i32 [[DUMMY16_6]], [[DUMMY16_6]]
-+; CHECK-NEXT:    [[DUMMY18_6:%.*]] = mul i32 [[DUMMY17_6]], [[DUMMY17_6]]
-+; CHECK-NEXT:    [[DUMMY19_6:%.*]] = mul i32 [[DUMMY18_6]], [[DUMMY18_6]]
-+; CHECK-NEXT:    [[DUMMY20_6:%.*]] = mul i32 [[DUMMY19_6]], [[DUMMY19_6]]
-+; CHECK-NEXT:    [[DUMMY21_6:%.*]] = mul i32 [[DUMMY20_6]], [[DUMMY20_6]]
-+; CHECK-NEXT:    [[DUMMY22_6:%.*]] = mul i32 [[DUMMY21_6]], [[DUMMY21_6]]
-+; CHECK-NEXT:    [[DUMMY23_6:%.*]] = mul i32 [[DUMMY22_6]], [[DUMMY22_6]]
-+; CHECK-NEXT:    [[DUMMY24_6:%.*]] = mul i32 [[DUMMY23_6]], [[DUMMY23_6]]
-+; CHECK-NEXT:    [[DUMMY25_6:%.*]] = mul i32 [[DUMMY24_6]], [[DUMMY24_6]]
-+; CHECK-NEXT:    [[DUMMY26_6:%.*]] = mul i32 [[DUMMY25_6]], [[DUMMY25_6]]
-+; CHECK-NEXT:    [[DUMMY27_6:%.*]] = mul i32 [[DUMMY26_6]], [[DUMMY26_6]]
-+; CHECK-NEXT:    [[DUMMY28_6:%.*]] = mul i32 [[DUMMY27_6]], [[DUMMY27_6]]
-+; CHECK-NEXT:    [[DUMMY29_6:%.*]] = mul i32 [[DUMMY28_6]], [[DUMMY28_6]]
-+; CHECK-NEXT:    [[DUMMY30_6:%.*]] = mul i32 [[DUMMY29_6]], [[DUMMY29_6]]
-+; CHECK-NEXT:    [[DUMMY31_6:%.*]] = mul i32 [[DUMMY30_6]], [[DUMMY30_6]]
-+; CHECK-NEXT:    [[DUMMY32_6:%.*]] = mul i32 [[DUMMY31_6]], [[DUMMY31_6]]
-+; CHECK-NEXT:    [[DUMMY33_6:%.*]] = mul i32 [[DUMMY32_6]], [[DUMMY32_6]]
-+; CHECK-NEXT:    [[DUMMY34_6:%.*]] = mul i32 [[DUMMY33_6]], [[DUMMY33_6]]
-+; CHECK-NEXT:    [[DUMMY35_6:%.*]] = mul i32 [[DUMMY34_6]], [[DUMMY34_6]]
-+; CHECK-NEXT:    [[DUMMY36_6:%.*]] = mul i32 [[DUMMY35_6]], [[DUMMY35_6]]
-+; CHECK-NEXT:    [[DUMMY37_6:%.*]] = mul i32 [[DUMMY36_6]], [[DUMMY36_6]]
-+; CHECK-NEXT:    [[DUMMY38_6:%.*]] = mul i32 [[DUMMY37_6]], [[DUMMY37_6]]
-+; CHECK-NEXT:    [[DUMMY39_6:%.*]] = mul i32 [[DUMMY38_6]], [[DUMMY38_6]]
-+; CHECK-NEXT:    [[DUMMY40_6:%.*]] = mul i32 [[DUMMY39_6]], [[DUMMY39_6]]
-+; CHECK-NEXT:    [[DUMMY41_6:%.*]] = mul i32 [[DUMMY40_6]], [[DUMMY40_6]]
-+; CHECK-NEXT:    [[DUMMY42_6:%.*]] = mul i32 [[DUMMY41_6]], [[DUMMY41_6]]
-+; CHECK-NEXT:    [[DUMMY43_6:%.*]] = mul i32 [[DUMMY42_6]], [[DUMMY42_6]]
-+; CHECK-NEXT:    [[DUMMY44_6:%.*]] = mul i32 [[DUMMY43_6]], [[DUMMY43_6]]
-+; CHECK-NEXT:    [[DUMMY45_6:%.*]] = mul i32 [[DUMMY44_6]], [[DUMMY44_6]]
-+; CHECK-NEXT:    [[DUMMY46_6:%.*]] = mul i32 [[DUMMY45_6]], [[DUMMY45_6]]
-+; CHECK-NEXT:    [[DUMMY47_6:%.*]] = mul i32 [[DUMMY46_6]], [[DUMMY46_6]]
-+; CHECK-NEXT:    [[DUMMY48_6:%.*]] = mul i32 [[DUMMY47_6]], [[DUMMY47_6]]
-+; CHECK-NEXT:    [[DUMMY49_6:%.*]] = mul i32 [[DUMMY48_6]], [[DUMMY48_6]]
-+; CHECK-NEXT:    [[DUMMY50_6:%.*]] = mul i32 [[DUMMY49_6]], [[DUMMY49_6]]
-+; CHECK-NEXT:    [[SUM_NEXT_6:%.*]] = add nsw i32 [[DUMMY50_6]], [[SUM_NEXT_5]]
-+; CHECK-NEXT:    [[INDVARS_IV_NEXT_6:%.*]] = add nuw nsw i64 [[INDVARS_IV]], 7
-+; CHECK-NEXT:    [[ARRAYIDX_7:%.*]] = getelementptr inbounds i32, ptr [[ARY]], i64 [[INDVARS_IV_NEXT_6]]
-+; CHECK-NEXT:    [[VAL_7:%.*]] = load i32, ptr [[ARRAYIDX_7]], align 4
-+; CHECK-NEXT:    [[DUMMY1_7:%.*]] = mul i32 [[VAL_7]], [[VAL_7]]
-+; CHECK-NEXT:    [[DUMMY2_7:%.*]] = mul i32 [[DUMMY1_7]], [[DUMMY1_7]]
-+; CHECK-NEXT:    [[DUMMY3_7:%.*]] = mul i32 [[DUMMY2_7]], [[DUMMY2_7]]
-+; CHECK-NEXT:    [[DUMMY4_7:%.*]] = mul i32 [[DUMMY3_7]], [[DUMMY3_7]]
-+; CHECK-NEXT:    [[DUMMY5_7:%.*]] = mul i32 [[DUMMY4_7]], [[DUMMY4_7]]
-+; CHECK-NEXT:    [[DUMMY6_7:%.*]] = mul i32 [[DUMMY5_7]], [[DUMMY5_7]]
-+; CHECK-NEXT:    [[DUMMY7_7:%.*]] = mul i32 [[DUMMY6_7]], [[DUMMY6_7]]
-+; CHECK-NEXT:    [[DUMMY8_7:%.*]] = mul i32 [[DUMMY7_7]], [[DUMMY7_7]]
-+; CHECK-NEXT:    [[DUMMY9_7:%.*]] = mul i32 [[DUMMY8_7]], [[DUMMY8_7]]
-+; CHECK-NEXT:    [[DUMMY10_7:%.*]] = mul i32 [[DUMMY9_7]], [[DUMMY9_7]]
-+; CHECK-NEXT:    [[DUMMY11_7:%.*]] = mul i32 [[DUMMY10_7]], [[DUMMY10_7]]
-+; CHECK-NEXT:    [[DUMMY12_7:%.*]] = mul i32 [[DUMMY11_7]], [[DUMMY11_7]]
-+; CHECK-NEXT:    [[DUMMY13_7:%.*]] = mul i32 [[DUMMY12_7]], [[DUMMY12_7]]
-+; CHECK-NEXT:    [[DUMMY14_7:%.*]] = mul i32 [[DUMMY13_7]], [[DUMMY13_7]]
-+; CHECK-NEXT:    [[DUMMY15_7:%.*]] = mul i32 [[DUMMY14_7]], [[DUMMY14_7]]
-+; CHECK-NEXT:    [[DUMMY16_7:%.*]] = mul i32 [[DUMMY15_7]], [[DUMMY15_7]]
-+; CHECK-NEXT:    [[DUMMY17_7:%.*]] = mul i32 [[DUMMY16_7]], [[DUMMY16_7]]
-+; CHECK-NEXT:    [[DUMMY18_7:%.*]] = mul i32 [[DUMMY17_7]], [[DUMMY17_7]]
-+; CHECK-NEXT:    [[DUMMY19_7:%.*]] = mul i32 [[DUMMY18_7]], [[DUMMY18_7]]
-+; CHECK-NEXT:    [[DUMMY20_7:%.*]] = mul i32 [[DUMMY19_7]], [[DUMMY19_7]]
-+; CHECK-NEXT:    [[DUMMY21_7:%.*]] = mul i32 [[DUMMY20_7]], [[DUMMY20_7]]
-+; CHECK-NEXT:    [[DUMMY22_7:%.*]] = mul i32 [[DUMMY21_7]], [[DUMMY21_7]]
-+; CHECK-NEXT:    [[DUMMY23_7:%.*]] = mul i32 [[DUMMY22_7]], [[DUMMY22_7]]
-+; CHECK-NEXT:    [[DUMMY24_7:%.*]] = mul i32 [[DUMMY23_7]], [[DUMMY23_7]]
-+; CHECK-NEXT:    [[DUMMY25_7:%.*]] = mul i32 [[DUMMY24_7]], [[DUMMY24_7]]
-+; CHECK-NEXT:    [[DUMMY26_7:%.*]] = mul i32 [[DUMMY25_7]], [[DUMMY25_7]]
-+; CHECK-NEXT:    [[DUMMY27_7:%.*]] = mul i32 [[DUMMY26_7]], [[DUMMY26_7]]
-+; CHECK-NEXT:    [[DUMMY28_7:%.*]] = mul i32 [[DUMMY27_7]], [[DUMMY27_7]]
-+; CHECK-NEXT:    [[DUMMY29_7:%.*]] = mul i32 [[DUMMY28_7]], [[DUMMY28_7]]
-+; CHECK-NEXT:    [[DUMMY30_7:%.*]] = mul i32 [[DUMMY29_7]], [[DUMMY29_7]]
-+; CHECK-NEXT:    [[DUMMY31_7:%.*]] = mul i32 [[DUMMY30_7]], [[DUMMY30_7]]
-+; CHECK-NEXT:    [[DUMMY32_7:%.*]] = mul i32 [[DUMMY31_7]], [[DUMMY31_7]]
-+; CHECK-NEXT:    [[DUMMY33_7:%.*]] = mul i32 [[DUMMY32_7]], [[DUMMY32_7]]
-+; CHECK-NEXT:    [[DUMMY34_7:%.*]] = mul i32 [[DUMMY33_7]], [[DUMMY33_7]]
-+; CHECK-NEXT:    [[DUMMY35_7:%.*]] = mul i32 [[DUMMY34_7]], [[DUMMY34_7]]
-+; CHECK-NEXT:    [[DUMMY36_7:%.*]] = mul i32 [[DUMMY35_7]], [[DUMMY35_7]]
-+; CHECK-NEXT:    [[DUMMY37_7:%.*]] = mul i32 [[DUMMY36_7]], [[DUMMY36_7]]
-+; CHECK-NEXT:    [[DUMMY38_7:%.*]] = mul i32 [[DUMMY37_7]], [[DUMMY37_7]]
-+; CHECK-NEXT:    [[DUMMY39_7:%.*]] = mul i32 [[DUMMY38_7]], [[DUMMY38_7]]
-+; CHECK-NEXT:    [[DUMMY40_7:%.*]] = mul i32 [[DUMMY39_7]], [[DUMMY39_7]]
-+; CHECK-NEXT:    [[DUMMY41_7:%.*]] = mul i32 [[DUMMY40_7]], [[DUMMY40_7]]
-+; CHECK-NEXT:    [[DUMMY42_7:%.*]] = mul i32 [[DUMMY41_7]], [[DUMMY41_7]]
-+; CHECK-NEXT:    [[DUMMY43_7:%.*]] = mul i32 [[DUMMY42_7]], [[DUMMY42_7]]
-+; CHECK-NEXT:    [[DUMMY44_7:%.*]] = mul i32 [[DUMMY43_7]], [[DUMMY43_7]]
-+; CHECK-NEXT:    [[DUMMY45_7:%.*]] = mul i32 [[DUMMY44_7]], [[DUMMY44_7]]
-+; CHECK-NEXT:    [[DUMMY46_7:%.*]] = mul i32 [[DUMMY45_7]], [[DUMMY45_7]]
-+; CHECK-NEXT:    [[DUMMY47_7:%.*]] = mul i32 [[DUMMY46_7]], [[DUMMY46_7]]
-+; CHECK-NEXT:    [[DUMMY48_7:%.*]] = mul i32 [[DUMMY47_7]], [[DUMMY47_7]]
-+; CHECK-NEXT:    [[DUMMY49_7:%.*]] = mul i32 [[DUMMY48_7]], [[DUMMY48_7]]
-+; CHECK-NEXT:    [[DUMMY50_7:%.*]] = mul i32 [[DUMMY49_7]], [[DUMMY49_7]]
-+; CHECK-NEXT:    [[SUM_NEXT_7]] = add nsw i32 [[DUMMY50_7]], [[SUM_NEXT_6]]
-+; CHECK-NEXT:    [[INDVARS_IV_NEXT_7]] = add nuw nsw i64 [[INDVARS_IV]], 8
-+; CHECK-NEXT:    [[NITER_NEXT_7]] = add i64 [[NITER]], 8
-+; CHECK-NEXT:    [[NITER_NCMP_7:%.*]] = icmp eq i64 [[NITER_NEXT_7]], [[UNROLL_ITER]]
-+; CHECK-NEXT:    br i1 [[NITER_NCMP_7]], label [[FOR_COND_CLEANUP_UNR_LCSSA_LOOPEXIT:%.*]], label [[FOR_BODY]]
- ; CHECK:       for.cond.cleanup.unr-lcssa.loopexit:
--; CHECK-NEXT:    [[SUM_NEXT_LCSSA_PH_PH:%.*]] = phi i32 [ [[SUM_NEXT_1]], [[FOR_BODY]] ]
--; CHECK-NEXT:    [[INDVARS_IV_UNR_PH:%.*]] = phi i64 [ [[INDVARS_IV_NEXT_1]], [[FOR_BODY]] ]
--; CHECK-NEXT:    [[SUM_UNR_PH:%.*]] = phi i32 [ [[SUM_NEXT_1]], [[FOR_BODY]] ]
-+; CHECK-NEXT:    [[SUM_NEXT_LCSSA_PH_PH:%.*]] = phi i32 [ [[SUM_NEXT_7]], [[FOR_BODY]] ]
-+; CHECK-NEXT:    [[INDVARS_IV_UNR_PH:%.*]] = phi i64 [ [[INDVARS_IV_NEXT_7]], [[FOR_BODY]] ]
-+; CHECK-NEXT:    [[SUM_UNR_PH:%.*]] = phi i32 [ [[SUM_NEXT_7]], [[FOR_BODY]] ]
- ; CHECK-NEXT:    br label [[FOR_COND_CLEANUP_UNR_LCSSA]]
- ; CHECK:       for.cond.cleanup.unr-lcssa:
- ; CHECK-NEXT:    [[SUM_NEXT_LCSSA_PH:%.*]] = phi i32 [ undef, [[ENTRY:%.*]] ], [ [[SUM_NEXT_LCSSA_PH_PH]], [[FOR_COND_CLEANUP_UNR_LCSSA_LOOPEXIT]] ]
-@@ -302,7 +1010,10 @@
- ; CHECK:       for.body.epil.preheader:
- ; CHECK-NEXT:    br label [[FOR_BODY_EPIL:%.*]]
- ; CHECK:       for.body.epil:
--; CHECK-NEXT:    [[ARRAYIDX_EPIL:%.*]] = getelementptr inbounds i32, ptr [[ARY]], i64 [[INDVARS_IV_UNR]]
-+; CHECK-NEXT:    [[INDVARS_IV_EPIL:%.*]] = phi i64 [ [[INDVARS_IV_UNR]], [[FOR_BODY_EPIL_PREHEADER]] ], [ [[INDVARS_IV_NEXT_EPIL:%.*]], [[FOR_BODY_EPIL]] ]
-+; CHECK-NEXT:    [[SUM_EPIL:%.*]] = phi i32 [ [[SUM_UNR]], [[FOR_BODY_EPIL_PREHEADER]] ], [ [[SUM_NEXT_EPIL:%.*]], [[FOR_BODY_EPIL]] ]
-+; CHECK-NEXT:    [[EPIL_ITER:%.*]] = phi i64 [ 0, [[FOR_BODY_EPIL_PREHEADER]] ], [ [[EPIL_ITER_NEXT:%.*]], [[FOR_BODY_EPIL]] ]
-+; CHECK-NEXT:    [[ARRAYIDX_EPIL:%.*]] = getelementptr inbounds i32, ptr [[ARY]], i64 [[INDVARS_IV_EPIL]]
- ; CHECK-NEXT:    [[VAL_EPIL:%.*]] = load i32, ptr [[ARRAYIDX_EPIL]], align 4
- ; CHECK-NEXT:    [[DUMMY1_EPIL:%.*]] = mul i32 [[VAL_EPIL]], [[VAL_EPIL]]
- ; CHECK-NEXT:    [[DUMMY2_EPIL:%.*]] = mul i32 [[DUMMY1_EPIL]], [[DUMMY1_EPIL]]
-@@ -354,10 +1065,17 @@
- ; CHECK-NEXT:    [[DUMMY48_EPIL:%.*]] = mul i32 [[DUMMY47_EPIL]], [[DUMMY47_EPIL]]
- ; CHECK-NEXT:    [[DUMMY49_EPIL:%.*]] = mul i32 [[DUMMY48_EPIL]], [[DUMMY48_EPIL]]
- ; CHECK-NEXT:    [[DUMMY50_EPIL:%.*]] = mul i32 [[DUMMY49_EPIL]], [[DUMMY49_EPIL]]
--; CHECK-NEXT:    [[SUM_NEXT_EPIL:%.*]] = add nsw i32 [[DUMMY50_EPIL]], [[SUM_UNR]]
-+; CHECK-NEXT:    [[SUM_NEXT_EPIL]] = add nsw i32 [[DUMMY50_EPIL]], [[SUM_EPIL]]
-+; CHECK-NEXT:    [[INDVARS_IV_NEXT_EPIL]] = add nuw nsw i64 [[INDVARS_IV_EPIL]], 1
-+; CHECK-NEXT:    [[EXITCOND_NOT_EPIL:%.*]] = icmp eq i64 [[INDVARS_IV_NEXT_EPIL]], [[N]]
-+; CHECK-NEXT:    [[EPIL_ITER_NEXT]] = add i64 [[EPIL_ITER]], 1
-+; CHECK-NEXT:    [[EPIL_ITER_CMP:%.*]] = icmp ne i64 [[EPIL_ITER_NEXT]], [[XTRAITER]]
-+; CHECK-NEXT:    br i1 [[EPIL_ITER_CMP]], label [[FOR_BODY_EPIL]], label [[FOR_COND_CLEANUP_EPILOG_LCSSA:%.*]], !llvm.loop [[LOOP0:![0-9]+]]
-+; CHECK:       for.cond.cleanup.epilog-lcssa:
-+; CHECK-NEXT:    [[SUM_NEXT_LCSSA_PH1:%.*]] = phi i32 [ [[SUM_NEXT_EPIL]], [[FOR_BODY_EPIL]] ]
- ; CHECK-NEXT:    br label [[FOR_COND_CLEANUP]]
- ; CHECK:       for.cond.cleanup:
--; CHECK-NEXT:    [[SUM_NEXT_LCSSA:%.*]] = phi i32 [ [[SUM_NEXT_LCSSA_PH]], [[FOR_COND_CLEANUP_UNR_LCSSA]] ], [ [[SUM_NEXT_EPIL]], [[FOR_BODY_EPIL]] ]
-+; CHECK-NEXT:    [[SUM_NEXT_LCSSA:%.*]] = phi i32 [ [[SUM_NEXT_LCSSA_PH]], [[FOR_COND_CLEANUP_UNR_LCSSA]] ], [ [[SUM_NEXT_LCSSA_PH1]], [[FOR_COND_CLEANUP_EPILOG_LCSSA]] ]
- ; CHECK-NEXT:    ret i32 [[SUM_NEXT_LCSSA]]
- ;
- entry:
diff --git a/third_party/llvm/workspace.bzl b/third_party/llvm/workspace.bzl
index 632e27d09333f5..4aca2a545e9fbf 100644
--- a/third_party/llvm/workspace.bzl
+++ b/third_party/llvm/workspace.bzl
@@ -4,8 +4,8 @@ load("//third_party:repo.bzl", "tf_http_archive")
 
 def repo(name):
     """Imports LLVM."""
-    LLVM_COMMIT = "e27f9bb31984927d10995e3d7e5a8fcde60cf5d4"
-    LLVM_SHA256 = "bba7c0ab0932cdb31b4ebf93f46bf5509fb173aa691d8a822e4717b001cb9791"
+    LLVM_COMMIT = "a383b3cca3383128123b33bc73ffc0199ec573e4"
+    LLVM_SHA256 = "f70bda6dd60f40263a377850a2c43556e4e36ed2a30b16c6ccb211b29eab2230"
 
     tf_http_archive(
         name = name,

From 632a143a0caf8c2c54a5050b5e9c0e3973ad5484 Mon Sep 17 00:00:00 2001
From: Luke Boyer <lukeboyer@google.com>
Date: Thu, 16 May 2024 14:29:49 -0700
Subject: [PATCH 305/478] Remove the test only tag from odml converter target.

PiperOrigin-RevId: 634530469
---
 tensorflow/compiler/mlir/lite/stablehlo/odml_converter/BUILD | 1 -
 1 file changed, 1 deletion(-)

diff --git a/tensorflow/compiler/mlir/lite/stablehlo/odml_converter/BUILD b/tensorflow/compiler/mlir/lite/stablehlo/odml_converter/BUILD
index eb0d86946e56eb..5add6c730cac5e 100644
--- a/tensorflow/compiler/mlir/lite/stablehlo/odml_converter/BUILD
+++ b/tensorflow/compiler/mlir/lite/stablehlo/odml_converter/BUILD
@@ -18,7 +18,6 @@ package_group(
 
 tf_cc_binary(
     name = "odml-converter",
-    testonly = True,
     srcs = ["odml_converter_main.cc"],
     compatible_with = get_compatible_with_portable(),
     visibility = [

From 1ab908fc086b31cab508b76c058aee985c2d3d84 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Thu, 16 May 2024 14:36:41 -0700
Subject: [PATCH 306/478] Express tests tolerances for exhaustive tests
 relative to the accuracy and subnormal boundary of the data type.

The following table shows the default absolute and relative error tolerances for each numeric datatype.

| Type             | atol before | rtol before | atol after | rtol after |
| ---------------- | ----------- | ---------- | ----------- | --------- |
| double           | 1e-4        | 1e-4       | 4.45e-308   | 1.11e-14 |
| float            | 1e-4        | 1e-4       | 2.35e-38    | 5.96e-6  |
| float16          | 1e-3        | 1e-3       | 1.22e-4     | 5e-4     |
| bfloat16         | 2e-3        | 2e-2       | 2.35e-38    | 1.6e-2   |
| complex<double>  | 1e-4        | 1e-4       | 4.45e-308   | 1.11e-14 |
| complex<float>   | 1e-4        | 1e-4       | 2.35e-38    | 5.96e-6  |

PiperOrigin-RevId: 634532580
---
 third_party/xla/xla/tests/exhaustive/BUILD    |  10 +-
 .../exhaustive/exhaustive_op_test_utils.h     |  89 +++-
 .../exhaustive_unary_test_complex.cc          | 142 ++++--
 .../exhaustive_unary_test_f32_or_smaller.cc   | 448 +++++++-----------
 .../exhaustive/exhaustive_unary_test_f64.cc   |   1 +
 5 files changed, 354 insertions(+), 336 deletions(-)

diff --git a/third_party/xla/xla/tests/exhaustive/BUILD b/third_party/xla/xla/tests/exhaustive/BUILD
index e995732bee51f5..e83f089cf18899 100644
--- a/third_party/xla/xla/tests/exhaustive/BUILD
+++ b/third_party/xla/xla/tests/exhaustive/BUILD
@@ -46,7 +46,7 @@ xla_test(
     name = "exhaustive_unary_test_f32_or_smaller",
     srcs = ["exhaustive_unary_test_f32_or_smaller.cc"],
     real_hardware_only = True,  # Very slow on the interpreter.
-    shard_count = 48,
+    shard_count = 50,
     tags = [
         "optonly",
         # This is a big test that we skip for capacity reasons in OSS testing.
@@ -55,6 +55,7 @@ xla_test(
     deps = [
         ":exhaustive_op_test_utils",
         "//xla:util",
+        "//xla/client:xla_builder",
         "//xla/tests:client_library_test_base",
     ],
 )
@@ -77,6 +78,7 @@ xla_test(
     deps = [
         ":exhaustive_op_test_utils",
         "//xla:util",
+        "//xla/client:xla_builder",
         "//xla/tests:client_library_test_base",
     ],
 )
@@ -96,8 +98,12 @@ xla_test(
     ],
     deps = [
         ":exhaustive_op_test_utils",
-        "//xla:util",
+        "//xla:types",
+        "//xla/client:xla_builder",
         "//xla/tests:client_library_test_base",
+        "//xla/tests:literal_test_util",
+        "//xla/tests:test_macros_header",
+        "@com_google_absl//absl/types:span",
     ],
 )
 
diff --git a/third_party/xla/xla/tests/exhaustive/exhaustive_op_test_utils.h b/third_party/xla/xla/tests/exhaustive/exhaustive_op_test_utils.h
index 363ddb357e1458..52fbeb2379827d 100644
--- a/third_party/xla/xla/tests/exhaustive/exhaustive_op_test_utils.h
+++ b/third_party/xla/xla/tests/exhaustive/exhaustive_op_test_utils.h
@@ -22,6 +22,7 @@ limitations under the License.
 #include <cstdint>
 #include <functional>
 #include <iterator>
+#include <limits>
 #include <string>
 #include <utility>
 #include <vector>
@@ -41,16 +42,13 @@ namespace xla {
 namespace exhaustive_op_test {
 
 struct ErrorSpec {
-  float abs_err;
-  float rel_err;
-
+  double abs_err = 0;
+  double rel_err = 0;
   // If true, will consider -0 not near to +0 and vice versa.  Note that
   // +epsilon may still be considered close to -0, depending on the error
   // spec; this only covers the case when both `expected` and `actual` are
   // equal to 0.
   bool strict_signed_zeros = false;
-
-  ErrorSpec(float a, float r) : abs_err(a), rel_err(r) {}
 };
 
 // Representations of the reference function passed in by the user.
@@ -940,66 +938,109 @@ inline std::vector<std::pair<int64_t, int64_t>> CreateExhaustiveF32Ranges() {
 }
 
 template <PrimitiveType T, size_t N>
-inline ErrorSpec DefaultSpecGenerator(
-    typename ExhaustiveOpTestBase<T, N>::NativeT) {
-  LOG(FATAL) << "Unhandled Type";
-}
+ErrorSpec DefaultSpecGenerator(typename ExhaustiveOpTestBase<T, N>::NativeT);
 
 template <PrimitiveType T, size_t N>
-inline ErrorSpec DefaultSpecGenerator(
-    typename ExhaustiveOpTestBase<T, N>::NativeT,
-    typename ExhaustiveOpTestBase<T, N>::NativeT) {
-  LOG(FATAL) << "Unhandled Type";
-}
+ErrorSpec DefaultSpecGenerator(typename ExhaustiveOpTestBase<T, N>::NativeT,
+                               typename ExhaustiveOpTestBase<T, N>::NativeT);
+
+// The following two constants set the default absolute and relative error
+// tolerance in units of the smallest normalized value and the relative accuracy
+// of the format, respectively. Notice that setting an absolute tolerance above
+// the value of the smallest normalized float means that we effectively ignore
+// relative errors in values at or below the subnormal boundary (e.g. for values
+// less than ~1e-38 for FP32).
+static constexpr float kDefaultAbsoluteToleranceSlackFactor = 2;
+static constexpr float kDefaultRelativeToleranceSlackFactor = 50;
 
 template <>
 inline ErrorSpec DefaultSpecGenerator<C128, 1>(complex128) {
-  return ErrorSpec{0.0001, 0.0001};
+  double atol =
+      kDefaultAbsoluteToleranceSlackFactor * std::numeric_limits<double>::min();
+  double rtol = kDefaultRelativeToleranceSlackFactor *
+                std::numeric_limits<double>::epsilon();
+  return ErrorSpec{atol, rtol};
 }
 
 template <>
 inline ErrorSpec DefaultSpecGenerator<C64, 1>(complex64) {
-  return ErrorSpec{0.0001, 0.0001};
+  double atol =
+      kDefaultAbsoluteToleranceSlackFactor * std::numeric_limits<float>::min();
+  double rtol = kDefaultRelativeToleranceSlackFactor *
+                std::numeric_limits<float>::epsilon();
+  return ErrorSpec{atol, rtol};
 }
 
 template <>
 inline ErrorSpec DefaultSpecGenerator<F64, 1>(double) {
-  return ErrorSpec{0.0001, 0.0001};
+  double atol =
+      kDefaultAbsoluteToleranceSlackFactor * std::numeric_limits<double>::min();
+  double rtol = kDefaultRelativeToleranceSlackFactor *
+                std::numeric_limits<double>::epsilon();
+  return ErrorSpec{atol, rtol};
 }
 
 template <>
 inline ErrorSpec DefaultSpecGenerator<F32, 1>(float) {
-  return ErrorSpec{0.0001, 0.0001};
+  double atol =
+      kDefaultAbsoluteToleranceSlackFactor * std::numeric_limits<float>::min();
+  double rtol = kDefaultRelativeToleranceSlackFactor *
+                std::numeric_limits<float>::epsilon();
+  return ErrorSpec{atol, rtol};
 }
 
 template <>
 inline ErrorSpec DefaultSpecGenerator<F16, 1>(Eigen::half) {
-  return ErrorSpec{0.001, 0.001};
+  double atol = kDefaultAbsoluteToleranceSlackFactor *
+                std::numeric_limits<Eigen::half>::min();
+  // epsilon for FP16 is quite large, so a slack factor of 5 suffices.
+  double rtol = 5 * std::numeric_limits<Eigen::half>::epsilon();
+  return ErrorSpec{atol, rtol};
 }
 
 template <>
 inline ErrorSpec DefaultSpecGenerator<BF16, 1>(bfloat16) {
-  return ErrorSpec{0.002, 0.02};
+  double atol = kDefaultAbsoluteToleranceSlackFactor *
+                std::numeric_limits<bfloat16>::min();
+  // epsilon for BF16 is quite large, so a slack factor of 2 suffices.
+  double rtol = 2 * std::numeric_limits<bfloat16>::epsilon();
+  return ErrorSpec{atol, rtol};
 }
 
 template <>
 inline ErrorSpec DefaultSpecGenerator<F64, 2>(double, double) {
-  return ErrorSpec{0.001, 0.001};
+  double atol =
+      kDefaultAbsoluteToleranceSlackFactor * std::numeric_limits<double>::min();
+  double rtol = kDefaultRelativeToleranceSlackFactor *
+                std::numeric_limits<double>::epsilon();
+  return ErrorSpec{atol, rtol};
 }
 
 template <>
 inline ErrorSpec DefaultSpecGenerator<F32, 2>(float, float) {
-  return ErrorSpec{0.001, 0.001};
+  double atol =
+      kDefaultAbsoluteToleranceSlackFactor * std::numeric_limits<float>::min();
+  double rtol = kDefaultRelativeToleranceSlackFactor *
+                std::numeric_limits<float>::epsilon();
+  return ErrorSpec{atol, rtol};
 }
 
 template <>
 inline ErrorSpec DefaultSpecGenerator<F16, 2>(Eigen::half, Eigen::half) {
-  return ErrorSpec{0.001, 0.001};
+  double atol = kDefaultAbsoluteToleranceSlackFactor *
+                std::numeric_limits<Eigen::half>::min();
+  // epsilon for FP16 is quite large, so a slack factor of 5 suffices.
+  double rtol = 5 * std::numeric_limits<Eigen::half>::epsilon();
+  return ErrorSpec{atol, rtol};
 }
 
 template <>
 inline ErrorSpec DefaultSpecGenerator<BF16, 2>(bfloat16, bfloat16) {
-  return ErrorSpec{0.002, 0.02};
+  double atol = kDefaultAbsoluteToleranceSlackFactor *
+                std::numeric_limits<bfloat16>::min();
+  // epsilon for BF16 is quite large, so a slack factor of 5 suffices.
+  double rtol = 2 * std::numeric_limits<bfloat16>::epsilon();
+  return ErrorSpec{atol, rtol};
 }
 
 template <PrimitiveType T, size_t N>
diff --git a/third_party/xla/xla/tests/exhaustive/exhaustive_unary_test_complex.cc b/third_party/xla/xla/tests/exhaustive/exhaustive_unary_test_complex.cc
index 0c078d667da787..6672747e30aa4b 100644
--- a/third_party/xla/xla/tests/exhaustive/exhaustive_unary_test_complex.cc
+++ b/third_party/xla/xla/tests/exhaustive/exhaustive_unary_test_complex.cc
@@ -13,9 +13,18 @@ See the License for the specific language governing permissions and
 limitations under the License.
 ==============================================================================*/
 
+#include <array>
+#include <cmath>
+#include <cstdint>
+#include <limits>
+#include <tuple>
+
+#include "absl/types/span.h"
+#include "xla/client/xla_builder.h"
 #include "xla/tests/client_library_test_base.h"
 #include "xla/tests/exhaustive/exhaustive_op_test_utils.h"
-#include "xla/util.h"
+#include "xla/tests/test_macros.h"
+#include "xla/types.h"
 
 #ifdef __FAST_MATH__
 #error "Can't be compiled with fast math on"
@@ -95,41 +104,68 @@ using ExhaustiveC128UnaryTest = ExhaustiveComplexUnaryTestBase<C128>;
   XLA_TEST_P(ExhaustiveC64UnaryTest, test_name) \
   __VA_ARGS__
 
-// TODO(b/138578594): Enable the test for the CPU backend after fixing the bug.
-UNARY_TEST_COMPLEX_64(DISABLED_ON_CPU(Log), {
-  // TODO(timshen): see b/162664705.
+UNARY_TEST_COMPLEX_64(Log, {
+  // TODO(rmlarsen): see b/162664705 and b/138578594
   known_incorrect_fn_ = [this](int64_t val) {
-    return std::isnan(this->ConvertValue(val));
+    complex64 x = this->ConvertValue(val);
+    return std::isnan(x.real()) || std::isnan(x.imag()) ||
+           (platform_ == "Host" &&
+            std::abs(x) < std::numeric_limits<float>::min());
   };
   ErrorSpecGen error_spec_gen = +[](complex64 x) {
-    // std::log uses std::abs, which uses hypotf. The latter upcasts its
-    // argument to double, so there are cases where there's excess precision, so
-    // abs returns numeric_limits<float>::max, but it's not possible to arrive
-    // at this number without using doubles internally.
-    // Obviously, the default relative error spec makes no sense for this case.
-    // `log(max) ~ 88.72`, which is very far from `inf`, but `max` is considered
-    // very close to `inf`.
-    if (std::abs(x) == std::numeric_limits<float>::max()) {
+    // The reference implementation overflows to infinity for arguments near
+    // FLT_MAX.
+    if (std::abs(x) >= std::numeric_limits<float>::max()) {
       float inf = std::numeric_limits<float>::infinity();
-      return ErrorSpec{inf, inf};
+      return ErrorSpec{.abs_err = inf, .rel_err = inf};
     }
-    return GetDefaultSpecGenerator()(x);
+    return ErrorSpec{.abs_err = std::numeric_limits<float>::epsilon(),
+                     .rel_err = 50 * std::numeric_limits<float>::epsilon()};
   };
-  Run(Log, [](complex64 x) { return std::log<float>(x); }, error_spec_gen);
+  Run(Log, [](complex64 x) { return std::log(x); }, error_spec_gen);
 })
 
 UNARY_TEST_COMPLEX_64(Sqrt, {
-  Run(Sqrt, [](complex64 x) {
-    return static_cast<complex64>(
-        std::sqrt<double>(static_cast<complex128>(x)));
-  });
+  // The reference implementation overflows to infinity for arguments near
+  // FLT_MAX.
+  ErrorSpecGen error_spec_gen = +[](complex64 x) {
+    if (std::abs(x) >= std::numeric_limits<float>::max()) {
+      float inf = std::numeric_limits<float>::infinity();
+      return ErrorSpec{.abs_err = inf, .rel_err = inf};
+    }
+    // The reference implementation appears to be very poor on inputs with
+    // subnormal entries. Allowing an absolute error of ~sqrt(FLT_DENORM_MIN)
+    // allows such cases to pass, effectively letting the relative error
+    // increase gradually until it reaches 100% at abs(x) == FLT_DENORM_MIN.
+    return ErrorSpec{
+        .abs_err = std::sqrt(std::numeric_limits<float>::denorm_min()),
+        .rel_err = 50 * std::numeric_limits<float>::epsilon()};
+  };
+  Run(Sqrt, [](complex64 x) { return std::sqrt(x); }, error_spec_gen);
 })
 
 UNARY_TEST_COMPLEX_64(Rsqrt, {
-  Run(Rsqrt, [](complex64 x) {
-    return static_cast<complex64>(
-        complex128(1, 0) / std::sqrt<double>(static_cast<complex128>(x)));
-  });
+  known_incorrect_fn_ = [this](int64_t val) {
+    complex64 x = this->ConvertValue(val);
+    return (platform_ == "Host" && (x.imag() == 0.0f || x.real() == 0.0f));
+  };
+  ErrorSpecGen error_spec_gen = +[](complex64 x) {
+    // As noted above for Sqrt, the accuracy of sqrt degrades severely for
+    // inputs with inputs with subnormals entries.
+    constexpr double norm_min = std::numeric_limits<float>::min();
+    constexpr double denorm_min = std::numeric_limits<float>::denorm_min();
+    if (std::abs(x) < norm_min) {
+      // Gradually loosen the relative tolerance as abs(x) becomes smaller
+      // than norm_min, letting it reach 100% when abs(x) = 10 * denorm_min.
+      return ErrorSpec{.abs_err = std::sqrt(std::numeric_limits<float>::min()),
+                       .rel_err = 10 * denorm_min / std::abs(x)};
+    }
+    return ErrorSpec{.abs_err = std::sqrt(std::numeric_limits<float>::min()),
+                     .rel_err = 50 * std::numeric_limits<float>::epsilon()};
+  };
+  Run(
+      Rsqrt, [](complex64 x) { return complex64(1, 0) / std::sqrt(x); },
+      error_spec_gen);
 })
 
 // The current libc++ implementation of the complex tanh function provides
@@ -142,7 +178,7 @@ UNARY_TEST_COMPLEX_64(Tanh, {
     // This implementation of Tanh becomes less accurate when the denominator
     // is small.
     if (std::cosh(2 * x.real()) + std::cos(2 * x.imag()) < 1e-4) {
-      return ErrorSpec{5e-2, 5e-2};
+      return ErrorSpec{.abs_err = 5e-2, .rel_err = 5e-2};
     }
 
     return GetDefaultSpecGenerator()(x);
@@ -193,47 +229,65 @@ INSTANTIATE_TEST_SUITE_P(
   __VA_ARGS__
 
 UNARY_TEST_COMPLEX_128(Log, {
-  // TODO(b/138578313): Enable the test for all values after fixing the bug.
+  // TODO(rmlarsen): see b/162664705 and b/138578594
   known_incorrect_fn_ = [&](int64_t v) {
     double f = this->ConvertValue(v);
     return std::fpclassify(f) == FP_NAN || std::abs(f) > 1.0e+300 ||
            std::abs(f) < 1.0e-300;
   };
-  Run(Log, [](complex128 x) { return std::log<double>(x); });
+  Run(Log, [](complex128 x) { return std::log(x); });
 })
 
 UNARY_TEST_COMPLEX_128(Sqrt, {
+  ErrorSpecGen error_spec_gen = +[](complex128 x) {
+    // TODO(rmlarsen): Investigate why we only get ~float accuracy here.
+
+    // The reference implementation appears to be very poor on inputs with
+    // subnormal entries. Allowing an absolute error of ~sqrt(DBL_DENORM_MIN)
+    // allows such cases to pass, effectively letting the relative error
+    // increase gradually until it reaches 100% at abs(x) == DBL_DENORM_MIN.
+    return ErrorSpec{
+        .abs_err = std::sqrt(std::numeric_limits<double>::denorm_min()),
+        .rel_err = 50 * std::numeric_limits<double>::epsilon()};
+  };
   // Similar to the Tanh bug.
   known_incorrect_fn_ = [&](int64_t v) {
     double f = this->ConvertValue(v);
     return std::abs(f) > std::numeric_limits<double>::max() / 2;
   };
-  Run(Sqrt, [](complex128 x) { return std::sqrt<double>(x); });
+  Run(Sqrt, [](complex128 x) { return std::sqrt(x); }, error_spec_gen);
 })
 
 UNARY_TEST_COMPLEX_128(Rsqrt, {
-  ErrorSpecGen error_spec_gen = GetDefaultSpecGenerator();
-  if (platform_ == "CUDA") {
-    // Edge case on CUDA backend where the Log of a complex number made up of
-    // the smallest denormals is more accurate than the interpreter backend.
-    error_spec_gen = [](complex128 x) {
-      constexpr double denorm_min = std::numeric_limits<double>::denorm_min();
-      if (std::abs(x.real()) == denorm_min &&
-          std::abs(x.imag()) == denorm_min) {
-        return ErrorSpec(0.5, 0.5);
-      }
-      return GetDefaultSpecGenerator()(x);
-    };
-  }
+  ErrorSpecGen error_spec_gen = +[](complex128 x) {
+    // As noted above for Sqrt, the accuracy of sqrt degrades severely for
+    // inputs with inputs with subnormals entries.
+    constexpr double norm_min = std::numeric_limits<double>::min();
+    constexpr double denorm_min = std::numeric_limits<double>::denorm_min();
+    if (std::abs(x) < norm_min) {
+      // Gradually loosen the relative tolerance as abs(x) becomes smaller
+      // than norm_min, letting it reach 100% when abs(x) = 10 * denorm_min.
+      return ErrorSpec{.abs_err = std::sqrt(std::numeric_limits<double>::min()),
+                       .rel_err = 10 * denorm_min / std::abs(x)};
+    }
+    return ErrorSpec{.abs_err = std::sqrt(std::numeric_limits<double>::min()),
+                     .rel_err = 50 * std::numeric_limits<double>::epsilon()};
+  };
   Run(
-      Rsqrt,
-      [](complex128 x) { return complex128(1, 0) / std::sqrt<double>(x); },
+      Rsqrt, [](complex128 x) { return complex128(1, 0) / std::sqrt(x); },
       error_spec_gen);
 })
 
 UNARY_TEST_COMPLEX_128(Tanh, {
+  ErrorSpecGen error_spec_gen = [](complex128 x) {
+    // TODO(rmlarsen): Investigate why we only get slightly better than
+    // float accuracy here.
+    return ErrorSpec{.abs_err = 2 * std::numeric_limits<double>::min(),
+                     .rel_err = 2e-8};
+  };
+
   SetParamsForTanh();
-  Run(Tanh, +[](complex128 x) { return std::tanh(x); });
+  Run(Tanh, +[](complex128 x) { return std::tanh(x); }, error_spec_gen);
 })
 
 INSTANTIATE_TEST_SUITE_P(
diff --git a/third_party/xla/xla/tests/exhaustive/exhaustive_unary_test_f32_or_smaller.cc b/third_party/xla/xla/tests/exhaustive/exhaustive_unary_test_f32_or_smaller.cc
index d908c4e167aeec..db3fc8d4d8aff4 100644
--- a/third_party/xla/xla/tests/exhaustive/exhaustive_unary_test_f32_or_smaller.cc
+++ b/third_party/xla/xla/tests/exhaustive/exhaustive_unary_test_f32_or_smaller.cc
@@ -15,6 +15,7 @@ limitations under the License.
 
 #include <fenv.h>  // NOLINT
 
+#include <algorithm>
 #include <array>
 #include <cmath>
 #include <limits>
@@ -22,6 +23,7 @@ limitations under the License.
 #include <tuple>
 #include <utility>
 
+#include "xla/client/xla_builder.h"
 #include "xla/tests/client_library_test_base.h"
 #include "xla/tests/exhaustive/exhaustive_op_test_utils.h"
 #include "xla/util.h"
@@ -305,96 +307,60 @@ using ExhaustiveBF16UnaryTest = Exhaustive32BitOrLessUnaryTest<BF16>;
   UNARY_TEST_BF16(test_name, __VA_ARGS__)
 
 UNARY_TEST_FLOAT_32_BITS_OR_LESS(Log, {
-  ErrorSpecGen error_spec_gen = GetDefaultSpecGenerator();
-  if (platform_ != "Host" && platform_ != "CUDA" && ty_ == F32) {
-    error_spec_gen = +[](NativeT x) { return ErrorSpec{0.001, 0.001}; };
+  auto error_spec_gen = GetDefaultSpecGenerator();
+  if (platform_ != "Host" && platform_ != "CUDA") {
+    error_spec_gen = +[](NativeT x) {
+      NativeT eps = std::numeric_limits<NativeT>::epsilon();
+      return ErrorSpec{.abs_err = 2e-4, .rel_err = eps};
+    };
   }
   Run(Log, std::log, error_spec_gen);
 })
-
 UNARY_TEST_FLOAT_32_BITS_OR_LESS(Log1p, {
-  ErrorSpecGen error_spec_gen = GetDefaultSpecGenerator();
-  if (platform_ != "Host" && platform_ != "CUDA" && ty_ == F32) {
-    error_spec_gen = +[](NativeT x) { return ErrorSpec{0.001, 0.001}; };
+  auto error_spec_gen = GetDefaultSpecGenerator();
+  if (platform_ != "Host" && platform_ != "CUDA") {
+    error_spec_gen = +[](NativeT x) {
+      NativeT eps = std::numeric_limits<NativeT>::epsilon();
+      return ErrorSpec{.abs_err = 2e-4, .rel_err = eps};
+    };
   }
   Run(Log1p, std::log1p, error_spec_gen);
 })
-
 UNARY_TEST_FLOAT_32_BITS_OR_LESS(Exp, {
-  // When x < -105, the true value of exp(x) is smaller than the smallest F32,
-  // so exp(x) should return exactly 0. We want our implementation of exp to
-  // return exactly 0 as well, as not doing so implies either that our
-  // implementation of exp is not following the asymptotic behavior that exp(x)
-  // approaches 0 as x approaches -inf, or that our implementation is not
-  // approaching 0 fast enough.
-  ErrorSpecGen error_spec_gen = +[](NativeT x) {
-    if (x < static_cast<NativeT>(-105)) {
-      return ErrorSpec{0, 0};
-    }
-    return GetDefaultSpecGenerator()(x);
-  };
-
-  // Our CPU implementation of exp returns one incorrect value: says
-  // exp(88.7228394) = max-float, but the correct answer is inf.  We deem this
-  // acceptable and check for it explicitly so that we can be aware if anything
-  // changes.
-  if (platform_ == "Host") {
-    auto host_exp_with_overflow = +[](float f) {
-      if (f == 88.7228394f) {
-        return 3.40282347e+38f;
-      }
-      return std::exp(f);
+  auto error_spec_gen = GetDefaultSpecGenerator();
+  if (platform_ != "Host" && platform_ != "CUDA") {
+    error_spec_gen = +[](NativeT x) {
+      NativeT min = std::numeric_limits<NativeT>::min();
+      NativeT eps = std::numeric_limits<NativeT>::epsilon();
+      return ErrorSpec{.abs_err = min, .rel_err = 75 * eps};
     };
-    Run(Exp, host_exp_with_overflow, error_spec_gen);
-  } else {
-    Run(Exp, std::exp, error_spec_gen);
   }
+  Run(Exp, std::exp, error_spec_gen);
 })
 
 UNARY_TEST_FLOAT_32_BITS_OR_LESS(Expm1, {
-  ErrorSpecGen error_spec_gen = GetDefaultSpecGenerator();
-  if (ty_ == F32) {
-    if (platform_ == "Host") {
-      error_spec_gen = +[](NativeT x) {
-        // We expect no worse than an error of 8 ULPs.
-        return ErrorSpec{
-            0.0, std::scalbn(8.0f, -std::numeric_limits<float>::digits)};
-      };
-    } else {
-      error_spec_gen = +[](NativeT x) { return ErrorSpec{0, 0.00015}; };
-    }
-  }
-
-  // Our CPU implementation of expm1 returns one incorrect value: says
-  // exp(88.7228394) = max-float, but the correct answer is inf.  We deem this
-  // acceptable and check for it explicitly so that we can be aware if anything
-  // changes.
-  if (platform_ == "Host") {
-    auto host_expm1_with_overflow = +[](float f) {
-      if (f == 88.7228394f) {
-        return 3.40282347e+38f;
-      }
-      return std::expm1(f);
+  auto error_spec_gen = GetDefaultSpecGenerator();
+  if (platform_ != "Host" && platform_ != "CUDA") {
+    error_spec_gen = +[](NativeT x) {
+      // FIXME(rmlarsen): Break into region around zero and everything else.
+      NativeT eps = std::numeric_limits<NativeT>::epsilon();
+      return ErrorSpec{.abs_err = 5e-4, .rel_err = 200 * eps};
     };
-    Run(Expm1, host_expm1_with_overflow, error_spec_gen);
-  } else {
-    Run(Expm1, std::expm1, error_spec_gen);
   }
+  Run(Expm1, std::expm1, error_spec_gen);
 })
-
 UNARY_TEST_FLOAT_32_BITS_OR_LESS(Logistic, {
+  // FIXME(rmlarsen): Break into region around zero and everything else.
+  auto error_spec_gen = GetDefaultSpecGenerator();
+  if (platform_ != "Host" && platform_ != "CUDA") {
+    error_spec_gen = +[](NativeT x) {
+      float eps = std::numeric_limits<NativeT>::epsilon();
+      float atol = std::min(0.004f, 200 * eps);
+      return ErrorSpec{.abs_err = atol, .rel_err = 0};
+    };
+  }
   EvaluateOp fn = +[](float x) { return 1.0f / (1.0f + std::exp(-x)); };
-
-  Run(
-      Logistic, fn, +[](NativeT) {
-        // Notice that we use the same absolute and relative tolerance.
-        // Since Logistic(x) -> 0 for x -> -Inf, the relative error
-        // is actually enormous for x < -5, say.
-        // For example, Logistic(-13.8183346) = 1.9967556e-06 while the expected
-        // value is 9.97178972e-07.
-        float tol = 200.0f * std::numeric_limits<NativeT>::epsilon();
-        return ErrorSpec(tol, tol);
-      });
+  Run(Logistic, fn, error_spec_gen);
 })
 
 // It feels a little overkill to exhaustively test sqrt and pow(x, 0.5), but
@@ -402,127 +368,127 @@ UNARY_TEST_FLOAT_32_BITS_OR_LESS(Logistic, {
 // pow(x, 0.5), but this is not true for x == -inf.
 UNARY_TEST_FLOAT_32_BITS_OR_LESS(PowOneHalf, {
   EvaluateOp fn = +[](float x) { return std::pow(x, 0.5f); };
-
   Run([](XlaOp x) { return Pow(x, ScalarLike(x, 0.5)); }, fn);
 })
 
 UNARY_TEST_FLOAT_32_BITS_OR_LESS(Rsqrt, {
-  Run(
-      Rsqrt, +[](float x) { return 1 / std::sqrt(x); });
+  Run(Rsqrt, +[](float x) { return 1 / std::sqrt(x); });
 })
 
 UNARY_TEST_FLOAT_32_BITS_OR_LESS(Sqrt, {
-  ErrorSpecGen error_spec_gen = GetDefaultSpecGenerator();
-  if (platform_ == "Host" || platform_ == "CUDA") {
-    error_spec_gen = +[](NativeT x) {
-      auto spec = GetDefaultSpecGenerator()(x);
-      spec.strict_signed_zeros = true;
-      return spec;
-    };
-  }
-
+  auto error_spec_gen = +[](NativeT x) {
+    auto spec = GetDefaultSpecGenerator()(x);
+    spec.strict_signed_zeros = true;
+    return spec;
+  };
   Run(Sqrt, std::sqrt, error_spec_gen);
 })
 
 UNARY_TEST_FLOAT_32_BITS_OR_LESS(Cbrt, {
-  if (platform_ == "Host" || platform_ == "CUDA") {
-    ErrorSpecGen error_spec_gen = +[](NativeT x) {
-      return ErrorSpec{0.01, 0.01};
-    };
-    Run(Cbrt, std::cbrt, error_spec_gen);
-  } else {
-    Run(Cbrt, std::cbrt);
-  }
+  auto error_spec_gen = +[](NativeT x) {
+    NativeT eps = std::numeric_limits<NativeT>::epsilon();
+    NativeT min = std::numeric_limits<NativeT>::min();
+    // Allow a small absollute error (e.g. 9e-16 for F32).
+    // This corresponds to a 0.5% relative error for the
+    // smallest normalized floating point values.
+    return ErrorSpec{.abs_err = std::cbrt(min) / 200, .rel_err = 10 * eps};
+  };
+  Run(Cbrt, std::cbrt, error_spec_gen);
 })
 
-// TODO(jlebar): Test trig functions over complex inputs.
-XLA_TEST_P(ExhaustiveF32UnaryTest, Acosh) {
-  // Error inherited from Log, which our implementation of Acosh uses.
-  ErrorSpecGen error_spec_gen = GetDefaultSpecGenerator();
+// Tests for inverse hyperbolic functions.
+UNARY_TEST_FLOAT_32_BITS_OR_LESS(Acosh, {
+  auto error_spec_gen = +[](NativeT x) {
+    NativeT eps = std::numeric_limits<NativeT>::epsilon();
+    return ErrorSpec{.abs_err = 1e-7, .rel_err = 50 * eps};
+  };
   if (platform_ != "Host" && platform_ != "CUDA") {
-    error_spec_gen = +[](float x) { return ErrorSpec{0.001, 0.001}; };
+    error_spec_gen = +[](NativeT x) {
+      NativeT eps = std::numeric_limits<NativeT>::epsilon();
+      return ErrorSpec{2e-4, eps};
+    };
   }
-
   Run(Acosh, std::acosh, error_spec_gen);
-}
-
-#if !defined(XLA_BACKEND_DOES_NOT_SUPPORT_FLOAT16)
-XLA_TEST_P(ExhaustiveF16UnaryTest, Acosh) { Run(Acosh, std::acosh); }
-#endif
-
-#if defined(XLA_BACKEND_SUPPORTS_BFLOAT16)
-XLA_TEST_P(ExhaustiveBF16UnaryTest, Acosh) { Run(Acosh, std::acosh); }
-#endif
-
-// Tests for Asinh
-XLA_TEST_P(ExhaustiveF32UnaryTest, Asinh) {
-  ErrorSpecGen error_spec_gen = GetDefaultSpecGenerator();
+})
+UNARY_TEST_FLOAT_32_BITS_OR_LESS(Asinh, {
+  auto error_spec_gen = GetDefaultSpecGenerator();
   if (platform_ != "Host" && platform_ != "CUDA") {
-    error_spec_gen = +[](float x) { return ErrorSpec{0.001, 0.001}; };
+    error_spec_gen = +[](NativeT x) {
+      NativeT eps = std::numeric_limits<NativeT>::epsilon();
+      return ErrorSpec{.abs_err = 2e-4, .rel_err = eps};
+    };
   }
-
   Run(Asinh, std::asinh, error_spec_gen);
-}
-
-#if !defined(XLA_BACKEND_DOES_NOT_SUPPORT_FLOAT16)
-XLA_TEST_P(ExhaustiveF16UnaryTest, Asinh) { Run(Asinh, std::asinh); }
-#endif
-
-#if defined(XLA_BACKEND_SUPPORTS_BFLOAT16)
-XLA_TEST_P(ExhaustiveBF16UnaryTest, Asinh) { Run(Asinh, std::asinh); }
-#endif
+})
+UNARY_TEST_FLOAT_32_BITS_OR_LESS(Atanh, {
+  auto error_spec_gen = GetDefaultSpecGenerator();
+  if (platform_ != "Host" && platform_ != "CUDA") {
+    error_spec_gen = +[](NativeT x) {
+      NativeT eps = std::numeric_limits<NativeT>::epsilon();
+      return ErrorSpec{.abs_err = 1e-4, .rel_err = eps};
+    };
+  }
+  Run(Atanh, std::atanh, error_spec_gen);
+})
 
-UNARY_TEST_FLOAT_32_BITS_OR_LESS(Atanh, { Run(Atanh, std::atanh); })
-UNARY_TEST_FLOAT_32_BITS_OR_LESS(Acos, { Run(Acos, std::acos); })
-UNARY_TEST_FLOAT_32_BITS_OR_LESS(Asin, { Run(Asin, std::asin); })
+// Tests for inverse trogonometric functions.
+UNARY_TEST_FLOAT_32_BITS_OR_LESS(Acos, {
+  ErrorSpecGen error_spec_gen = GetDefaultSpecGenerator();
+  if (platform_ != "Host") {
+    error_spec_gen = +[](NativeT x) {
+      NativeT eps = std::numeric_limits<NativeT>::epsilon();
+      return ErrorSpec{.abs_err = 1e-6, .rel_err = 10 * eps};
+    };
+  }
+  Run(Acos, std::acos, error_spec_gen);
+})
+UNARY_TEST_FLOAT_32_BITS_OR_LESS(Asin, {
+  auto error_spec_gen = +[](NativeT x) {
+    NativeT min = std::numeric_limits<NativeT>::min();
+    NativeT eps = std::numeric_limits<NativeT>::epsilon();
+    return ErrorSpec{.abs_err = 2.0f * min, .rel_err = 10 * eps};
+  };
+  Run(Asin, std::asin, error_spec_gen);
+})
+UNARY_TEST_FLOAT_32_BITS_OR_LESS(Atan, {
+  auto error_spec_gen = +[](NativeT x) {
+    NativeT min = std::numeric_limits<NativeT>::min();
+    NativeT eps = std::numeric_limits<NativeT>::epsilon();
+    return ErrorSpec{.abs_err = 2.0f * min, .rel_err = 20 * eps};
+  };
+  Run(Atan, std::atan, error_spec_gen);
+})
 
 UNARY_TEST_FLOAT_32_BITS_OR_LESS(Cosh, {
-  // Our cosh implementation incorrectly overflows to inf for +/-89.4159851.
-  // The correct answer of 3.40281961e+38 (0x7f7fffec) is very close to
-  // max-float, so we deem this acceptable.
-  //
-  // This does not occur on CPU because we have an offsetting error in our
-  // implementation of exp.
-  float (*host_cosh)(float);
-  if (platform_ == "Host") {
-    host_cosh = &std::cosh;
-  } else {
-    host_cosh = +[](float x) {
-      if (std::abs(x) == 89.4159851f) {
-        return std::numeric_limits<float>::infinity();
-      }
-      return std::cosh(x);
+  auto error_spec_gen = GetDefaultSpecGenerator();
+  if (platform_ != "Host" && platform_ != "CUDA") {
+    error_spec_gen = +[](NativeT x) {
+      NativeT eps = std::numeric_limits<NativeT>::epsilon();
+      // Cosh is always greater than or equal to 1, so an absolute
+      // tolerance does not make sense.
+      return ErrorSpec{.abs_err = 0, .rel_err = 100 * eps};
     };
   }
-  Run(Cosh, host_cosh);
+  Run(Cosh, std::cosh, error_spec_gen);
 })
 
 UNARY_TEST_FLOAT_32_BITS_OR_LESS(Sinh, {
-  // Our sinh implementation incorrectly overflows to +/-inf for +/-89.4159851.
-  // The correct answer of 3.40281961e+38 (0x7f7fffec) is very close to
-  // max-float, so we deem this acceptable.
-  //
-  // This does not occur on CPU because we have an offsetting error in our
-  // implementation of exp.
-  float (*host_sinh)(float);
-  if (platform_ == "Host") {
-    host_sinh = &std::sinh;
-  } else {
-    host_sinh = +[](float x) {
-      if (std::abs(x) == 89.4159851f) {
-        return std::copysign(std::numeric_limits<float>::infinity(), x);
-      }
-      return std::sinh(x);
+  auto error_spec_gen = GetDefaultSpecGenerator();
+  if (platform_ != "Host" && platform_ != "CUDA") {
+    error_spec_gen = +[](NativeT x) {
+      NativeT eps = std::numeric_limits<NativeT>::epsilon();
+      return ErrorSpec{.abs_err = 1e-5, .rel_err = 100 * eps};
     };
   }
-  Run(Sinh, host_sinh);
+  Run(Sinh, std::sinh, error_spec_gen);
 })
 
 UNARY_TEST_FLOAT_32_BITS_OR_LESS(TanhBounderTestUpperBound, {
   SetBounder(8, 9);
   ErrorSpecGen error_spec_gen = GetDefaultSpecGenerator();
-  if (platform_ == "CUDA" || platform_ == "CPU") {
-    error_spec_gen = +[](NativeT x) { return ErrorSpec{0, 0}; };
+  if (platform_ == "CUDA" || platform_ == "Host") {
+    error_spec_gen =
+        +[](NativeT x) { return ErrorSpec{.abs_err = 0, .rel_err = 0}; };
   }
   Run(
       Tanh, +[](float) { return 1.0f; }, error_spec_gen,
@@ -532,7 +498,7 @@ UNARY_TEST_FLOAT_32_BITS_OR_LESS(TanhBounderTestUpperBound, {
 UNARY_TEST_FLOAT_32_BITS_OR_LESS(TanhBounderTestLowerBound, {
   SetBounder(-9, -8);
   ErrorSpecGen error_spec_gen = GetDefaultSpecGenerator();
-  if (platform_ == "CUDA" || platform_ == "CPU") {
+  if (platform_ == "CUDA" || platform_ == "Host") {
     error_spec_gen = +[](NativeT x) { return ErrorSpec{0, 0}; };
   }
   Run(
@@ -542,154 +508,104 @@ UNARY_TEST_FLOAT_32_BITS_OR_LESS(TanhBounderTestLowerBound, {
 
 UNARY_TEST_FLOAT_32_BITS_OR_LESS(TanhNormalTest, {
   ErrorSpecGen error_spec_gen = GetDefaultSpecGenerator();
-  if (platform_ == "CUDA") {
+  if (platform_ != "CUDA" && platform_ != "Host") {
     error_spec_gen = +[](NativeT x) {
-      return x <= static_cast<NativeT>(-20.0) || x >= static_cast<NativeT>(20.0)
-                 ? ErrorSpec{0, 0}
-                 : GetDefaultSpecGenerator()(x);
+      // The range of tanh is [-1:1], so no point in giving a relative
+      // tolerance when we have an absolute one.
+      return ErrorSpec{.abs_err = 5e-4, .rel_err = 0};
     };
   }
   Run(Tanh, std::tanh, error_spec_gen);
 })
 
-UNARY_TEST_F32(Cos, {
+UNARY_TEST_FLOAT_32_BITS_OR_LESS(Cos, {
   Run(
       Cos, std::cos, +[](NativeT) {
         // This error spec corresponds to a maximum relative error of 2 ULP.
-        return ErrorSpec(0, 2 * std::numeric_limits<float>::epsilon());
+        NativeT eps = std::numeric_limits<NativeT>::epsilon();
+        return ErrorSpec{.abs_err = 0, .rel_err = 2 * eps};
       });
 })
 
-UNARY_TEST_F16(Cos, { Run(Cos, std::cos); })
-
-UNARY_TEST_BF16(Cos, { Run(Cos, std::cos); })
-
-UNARY_TEST_F32(Sin, {
+UNARY_TEST_FLOAT_32_BITS_OR_LESS(Sin, {
   Run(
       Sin, std::sin, +[](NativeT) {
         // This error spec corresponds to a maximum relative error of 2 ULP.
-        return ErrorSpec(0, 2 * std::numeric_limits<float>::epsilon());
+        NativeT eps = std::numeric_limits<NativeT>::epsilon();
+        return ErrorSpec{.abs_err = 0, .rel_err = 2 * eps};
       });
 })
 
-UNARY_TEST_F16(Sin, { Run(Sin, std::sin); })
-
-UNARY_TEST_BF16(Sin, { Run(Sin, std::sin); })
-
-UNARY_TEST_F32(Tan, {
+UNARY_TEST_FLOAT_32_BITS_OR_LESS(Tan, {
   Run(
       Tan, std::tan, +[](NativeT) {
         // This error spec corresponds to a maximum relative error of 4 ULP.
-        return ErrorSpec(0, 4 * std::numeric_limits<float>::epsilon());
+        NativeT eps = std::numeric_limits<NativeT>::epsilon();
+        return ErrorSpec{.abs_err = 0, .rel_err = 4 * eps};
       });
 })
 
-UNARY_TEST_F16(Tan, {
-  Run(
-      Tan, std::tan, +[](NativeT) {
-        // This error spec corresponds to a maximum relative error of 2 ULP.
-        return ErrorSpec(0, 2 * std::numeric_limits<Eigen::half>::epsilon());
-      });
-})
-
-UNARY_TEST_BF16(Tan, {
-  Run(
-      Tan, std::tan, +[](NativeT) {
-        // This error spec corresponds to a maximum relative error of 1 ULP.
-        return ErrorSpec(0, std::numeric_limits<Eigen::bfloat16>::epsilon());
-      });
-})
-
-// TODO(jlebar): Enable these.
-// UNARY_TEST_FLOAT_32_BITS_OR_LESS(Atan) { Run(Atan, std::atan); }
-// UNARY_TEST_FLOAT_32_BITS_OR_LESS(Atan2) { Run(Atan2, std::atan2); }
-
-UNARY_TEST_FLOAT_32_BITS_OR_LESS(Erf, {
-  Run(
-      Erf, std::erf, +[](NativeT x) {
-        NativeT tol =
-            std::max(std::numeric_limits<NativeT>::epsilon(),
-                     NativeT(4 * std::numeric_limits<float>::epsilon()));
-        NativeT abs_tol = std::min(tol, NativeT(1 - std::abs(std::erf(x))));
-        return ErrorSpec(abs_tol, tol);
-      });
+UNARY_TEST_FLOAT_32_BITS_OR_LESS(Erf, { Run(Erf, std::erf); })
+UNARY_TEST_FLOAT_32_BITS_OR_LESS(Erfc, {
+  auto error_spec_gen = GetDefaultSpecGenerator();
+  if (platform_ != "Host" && platform_ != "CUDA") {
+    error_spec_gen = +[](NativeT x) {
+      NativeT min = std::numeric_limits<NativeT>::min();
+      NativeT eps = std::numeric_limits<NativeT>::epsilon();
+      return ErrorSpec{.abs_err = 2 * min, .rel_err = 100 * eps};
+    };
+  }
+  Run(Erfc, std::erfc, error_spec_gen);
 })
-
-UNARY_TEST_FLOAT_32_BITS_OR_LESS(Erfc, { Run(Erfc, std::erfc); })
-
-UNARY_TEST_F32(ErfInv, { Run(ErfInv, HostErfInv); })
-
-UNARY_TEST_F16(ErfInv, {
-  Run(ErfInv, HostErfInv, [](Eigen::half) { return ErrorSpec{0.002, 0.002}; });
+UNARY_TEST_FLOAT_32_BITS_OR_LESS(ErfInv, {
+  auto error_spec_gen = GetDefaultSpecGenerator();
+  if (platform_ != "Host" && platform_ != "CUDA") {
+    error_spec_gen = +[](NativeT x) {
+      NativeT eps = std::numeric_limits<NativeT>::epsilon();
+      return ErrorSpec{.abs_err = 5e-5, .rel_err = 2 * eps};
+    };
+  }
+  Run(ErfInv, HostErfInv, error_spec_gen);
 })
 
-UNARY_TEST_BF16(ErfInv, { Run(ErfInv, HostErfInv); })
-
 UNARY_TEST_FLOAT_32_BITS_OR_LESS(Digamma, {
-  ErrorSpecGen error_spec_gen = GetDefaultSpecGenerator();
+  auto error_spec_gen = +[](NativeT x) {
+    NativeT eps = std::numeric_limits<NativeT>::epsilon();
+    return ErrorSpec{2e-5, 10 * eps};
+  };
   if (platform_ != "Host" && platform_ != "CUDA") {
-    // TODO(b/123956399): This is a fairly high error, significantly higher than
-    // we see on CPU/GPU.
-    error_spec_gen = +[](NativeT) { return ErrorSpec{0.01, 0.01}; };
-  }
-
-  if (platform_ == "CUDA") {
-    // On GPU we get a wrong answer for the denormal inputs +/-2.93873588e-39
-    // (0x00200000 and 0x80200000).  These should return -/+inf (at least
-    // according to our reference implementation!) but XLA:GPU returns
-    // -/+3.40282326e+38 (0xff7ffffe and 0x7f7ffffe).
-    //
-    // I deem this an acceptable result, as XLA:GPU flushes denormals, and as
-    // the results we get here are very close to MAX_FLOAT.  We just hardcode
-    // these results, as this is better than ignoring these inputs altogether.
-    auto host_digamma_with_gpu_ftz_errors = +[](float x) {
-      if (BitCast<uint32_t>(x) == 0x00200000 ||
-          BitCast<uint32_t>(x) == 0x80200000) {
-        return std::copysign(std::numeric_limits<float>::max(), -x);
-      }
-      return HostDigamma(x);
+    error_spec_gen = +[](NativeT x) {
+      NativeT eps = std::numeric_limits<NativeT>::epsilon();
+      return ErrorSpec{.abs_err = 2e-4, .rel_err = 10 * eps};
     };
-    Run(Digamma, host_digamma_with_gpu_ftz_errors, error_spec_gen);
-  } else {
-    Run(Digamma, HostDigamma, error_spec_gen);
   }
+  Run(Digamma, HostDigamma, error_spec_gen);
 })
 
 UNARY_TEST_FLOAT_32_BITS_OR_LESS(Lgamma, {
-  // Our implementation gets within 0.0001 rel error except for ~20 denormal
-  // inputs on GPU.  Anyway 0.001 rel error should be good enough for lgamma.
-  ErrorSpecGen error_spec_gen = GetDefaultSpecGenerator();
-  if (platform_ == "CUDA" && (ty_ == F32 || ty_ == F16)) {
+  auto error_spec_gen = +[](NativeT x) {
+    NativeT eps = std::numeric_limits<NativeT>::epsilon();
+    return ErrorSpec{.abs_err = 1e-5, .rel_err = 150 * eps};
+  };
+  if (platform_ == "CUDA") {
     error_spec_gen = +[](NativeT x) {
-      auto spec = GetDefaultSpecGenerator()(x);
-      spec.rel_err = 0.001;
-      return spec;
+      NativeT eps = std::numeric_limits<NativeT>::epsilon();
+      return ErrorSpec{.abs_err = 1e-5, .rel_err = 5000 * eps};
+    };
+  } else if (platform_ != "Host") {
+    error_spec_gen = +[](NativeT x) {
+      NativeT eps = std::numeric_limits<NativeT>::epsilon();
+      return ErrorSpec{.abs_err = 5e-4, .rel_err = 5000 * eps};
     };
   }
-
-  float (*host_lgamma)(float) = std::lgamma;
-  if (platform_ != "Host" && platform_ != "CUDA") {
-    // TODO(b/123956399): This is a fairly high error, significantly higher than
-    // we see on CPU/GPU.
-    error_spec_gen = +[](NativeT) { return ErrorSpec{0.01, 0.01}; };
-
-    // Overflows to inf for input 4.08500343e+36 (0x7c44af8e).
-    if (ty_ == F32) {
-      host_lgamma = +[](float v) {
-        if (BitCast<uint32_t>(v) == 0x7c44af8e) {
-          return std::numeric_limits<float>::infinity();
-        }
-        return std::lgamma(v);
-      };
-    }
-  }
-  Run(Lgamma, host_lgamma, error_spec_gen);
+  Run(Lgamma, std::lgamma, error_spec_gen);
 })
 
 UNARY_TEST_FLOAT_32_BITS_OR_LESS(Round, { Run(Round, std::round); })
 
 UNARY_TEST_FLOAT_32_BITS_OR_LESS(RoundNearestEven, {
-  ErrorSpecGen error_spec_gen = +[](NativeT) { return ErrorSpec{0.0, 0.0}; };
+  ErrorSpecGen error_spec_gen =
+      +[](NativeT) { return ErrorSpec{.abs_err = 0.0, .rel_err = 0.0}; };
   int curr_direction = fegetround();
   fesetround(FE_TONEAREST);
   Run(RoundNearestEven, std::nearbyint, error_spec_gen);
diff --git a/third_party/xla/xla/tests/exhaustive/exhaustive_unary_test_f64.cc b/third_party/xla/xla/tests/exhaustive/exhaustive_unary_test_f64.cc
index f95a03c209e876..7d06c3121b4c3d 100644
--- a/third_party/xla/xla/tests/exhaustive/exhaustive_unary_test_f64.cc
+++ b/third_party/xla/xla/tests/exhaustive/exhaustive_unary_test_f64.cc
@@ -13,6 +13,7 @@ See the License for the specific language governing permissions and
 limitations under the License.
 ==============================================================================*/
 
+#include "xla/client/xla_builder.h"
 #include "xla/tests/client_library_test_base.h"
 #include "xla/tests/exhaustive/exhaustive_op_test_utils.h"
 #include "xla/util.h"

From 1c969ddae00c86e16299a08e253e100e1d314ac7 Mon Sep 17 00:00:00 2001
From: Kyle Lucke <klucke@google.com>
Date: Thu, 16 May 2024 14:42:21 -0700
Subject: [PATCH 307/478] Use absl::Status instead of xla::Status now that
 they're identical.

PiperOrigin-RevId: 634534302
---
 .../xla/xla/service/algebraic_simplifier.cc   | 126 ++++---
 .../xla/xla/service/algebraic_simplifier.h    | 107 +++---
 .../xla/xla/service/all_gather_combiner.cc    |   6 +-
 .../xla/xla/service/all_gather_decomposer.cc  |   4 +-
 .../xla/xla/service/all_reduce_combiner.cc    |   2 +-
 .../xla/xla/service/all_reduce_contiguous.cc  |   3 +-
 .../xla/xla/service/allocation_tracker.cc     |   6 +-
 .../xla/xla/service/allocation_tracker.h      |   6 +-
 .../xla/xla/service/ar_crs_combiner.cc        |   4 +-
 third_party/xla/xla/service/ar_crs_combiner.h |   4 +-
 third_party/xla/xla/service/backend.cc        |   2 +-
 third_party/xla/xla/service/backend.h         |   2 +-
 .../xla/xla/service/batchnorm_expander.cc     |  12 +-
 .../service/bfloat16_conversion_folding.cc    |  23 +-
 .../xla/xla/service/bfloat16_propagation.cc   |   6 +-
 .../xla/xla/service/bfloat16_propagation.h    |   6 +-
 .../xla/xla/service/buffer_assignment.cc      |  14 +-
 .../xla/xla/service/buffer_assignment.h       |  15 +-
 .../xla/xla/service/buffer_assignment_test.cc |   2 +-
 third_party/xla/xla/service/call_graph.cc     |   6 +-
 third_party/xla/xla/service/call_graph.h      |   8 +-
 .../xla/xla/service/call_graph_test.cc        |   2 +-
 third_party/xla/xla/service/call_inliner.cc   |  12 +-
 .../xla/service/collective_combiner_utils.h   |   3 +-
 .../xla/xla/service/collective_ops_utils.h    |   6 +-
 .../service/collective_permute_decomposer.cc  |   2 +-
 .../xla/xla/service/collective_pipeliner.cc   |  36 +-
 .../xla/xla/service/collective_pipeliner.h    |   2 +-
 .../xla/service/compilation_environments.cc   |   4 +-
 .../xla/service/compilation_environments.h    |   6 +-
 .../xla/xla/service/compile_only_service.h    |  21 +-
 .../xla/xla/service/computation_placer.cc     |   2 +-
 .../xla/xla/service/computation_placer.h      |   4 +-
 .../xla/service/conditional_canonicalizer.cc  |   2 +-
 .../xla/service/conditional_code_motion.cc    |  10 +-
 .../xla/xla/service/conditional_to_select.cc  |   2 +-
 .../convert_async_collectives_to_sync.cc      |   2 +-
 .../convert_async_collectives_to_sync.h       |   4 +-
 .../convert_async_collectives_to_sync_test.cc |   4 +-
 .../service/convolution_group_converter.cc    |  12 +-
 third_party/xla/xla/service/copy_insertion.cc |  28 +-
 third_party/xla/xla/service/copy_insertion.h  |  12 +-
 .../xla/xla/service/cpu_gpu_shape_verifier.cc |   4 +-
 .../xla/xla/service/cpu_gpu_shape_verifier.h  |   2 +-
 third_party/xla/xla/service/defuser.cc        |   2 +-
 .../dfs_hlo_visitor_with_default_test.cc      |   6 +-
 third_party/xla/xla/service/dot_decomposer.cc |   2 +-
 .../xla/xla/service/dot_dimension_merger.cc   |   2 +-
 third_party/xla/xla/service/dump.cc           |  11 +-
 .../service/dynamic_dimension_inference.cc    | 305 ++++++++--------
 .../xla/service/dynamic_dimension_inference.h |   9 +-
 .../dynamic_dimension_inference_test.cc       |   4 +-
 third_party/xla/xla/service/dynamic_padder.cc |  61 ++--
 .../xla/xla/service/dynamic_padder_test.cc    |   2 +-
 third_party/xla/xla/service/eigh_expander.cc  |   8 +-
 third_party/xla/xla/service/eigh_expander.h   |   2 +-
 third_party/xla/xla/service/executable.cc     |  16 +-
 third_party/xla/xla/service/executable.h      |   2 +-
 .../xla/xla/service/execution_tracker.cc      |   4 +-
 .../xla/xla/service/execution_tracker.h       |   4 +-
 third_party/xla/xla/service/export_hlo.h      |   4 +-
 .../xla/xla/service/flatten_call_graph.cc     |   2 +-
 .../xla/xla/service/float_normalization.cc    |  47 +--
 .../xla/xla/service/gather_expander_test.cc   |   6 +-
 .../xla/service/generic_transfer_manager.cc   |  25 +-
 .../xla/service/generic_transfer_manager.h    |  43 ++-
 .../service/gpu_compilation_environment.cc    |   3 +-
 .../xla/service/gpu_compilation_environment.h |   3 +-
 .../xla/xla/service/hlo_alias_analysis.cc     |   2 +-
 .../xla/xla/service/hlo_alias_analysis.h      |   2 +-
 .../xla/xla/service/hlo_computation_test.cc   |   4 +-
 .../xla/xla/service/hlo_cost_analysis.cc      | 198 ++++++-----
 .../xla/xla/service/hlo_cost_analysis.h       | 188 +++++-----
 .../xla/xla/service/hlo_dataflow_analysis.cc  |   4 +-
 .../xla/xla/service/hlo_dataflow_analysis.h   |   4 +-
 .../xla/service/hlo_dataflow_analysis_test.cc |   2 +-
 third_party/xla/xla/service/hlo_dce.cc        |   2 +-
 third_party/xla/xla/service/hlo_dce.h         |   2 +-
 third_party/xla/xla/service/hlo_domain_map.cc |  12 +-
 third_party/xla/xla/service/hlo_domain_map.h  |  12 +-
 .../xla/xla/service/hlo_domain_remover.cc     |   4 +-
 .../xla/xla/service/hlo_domain_remover.h      |   8 +-
 .../xla/xla/service/hlo_domain_test.cc        |   4 +-
 .../xla/xla/service/hlo_domain_verifier.cc    |   9 +-
 .../xla/xla/service/hlo_graph_dumper.cc       |   8 +-
 .../xla/xla/service/hlo_instruction_test.cc   |  18 +-
 .../xla/xla/service/hlo_module_config_test.cc |   2 +-
 .../xla/service/hlo_module_group_metadata.cc  |  18 +-
 .../xla/service/hlo_module_group_metadata.h   |  14 +-
 .../xla/xla/service/hlo_module_group_util.cc  |   9 +-
 .../xla/xla/service/hlo_module_group_util.h   |  13 +-
 .../xla/xla/service/hlo_module_util.cc        |   4 +-
 third_party/xla/xla/service/hlo_parser.cc     |  10 +-
 third_party/xla/xla/service/hlo_parser.h      |   2 +-
 third_party/xla/xla/service/hlo_pass_fix.h    |  11 +-
 .../xla/xla/service/hlo_pass_interface.h      |   6 +-
 .../xla/xla/service/hlo_pass_pipeline.cc      |  23 +-
 .../xla/xla/service/hlo_pass_pipeline.h       |   5 +-
 .../xla/xla/service/hlo_pass_pipeline_test.cc |   6 +-
 .../xla/xla/service/hlo_rematerialization.cc  |  44 +--
 .../xla/service/hlo_replication_analysis.cc   |   4 +-
 .../xla/service/hlo_replication_analysis.h    |   2 +-
 .../service/hlo_value_semantics_analysis.cc   |   2 +-
 .../service/hlo_value_semantics_analysis.h    | 215 ++++++------
 third_party/xla/xla/service/hlo_verifier.cc   | 329 +++++++++---------
 third_party/xla/xla/service/hlo_verifier.h    | 214 ++++++------
 .../xla/xla/service/hlo_verifier_test.cc      |  14 +-
 .../host_memory_transfer_asyncifier.cc        |   8 +-
 .../xla/xla/service/host_offload_legalize.cc  |  17 +-
 third_party/xla/xla/service/host_offloader.cc |  21 +-
 third_party/xla/xla/service/host_offloader.h  |  18 +-
 .../xla/xla/service/indexed_array_analysis.cc |   2 +-
 .../xla/xla/service/indexed_array_analysis.h  |   2 +-
 .../xla/service/latency_hiding_scheduler.cc   |   6 +-
 .../xla/service/latency_hiding_scheduler.h    |   6 +-
 .../xla/xla/service/layout_assignment.cc      | 134 +++----
 .../xla/xla/service/layout_assignment.h       | 143 ++++----
 .../xla/xla/service/layout_assignment_test.cc |   8 +-
 .../xla/xla/service/layout_normalization.cc   |  42 +--
 third_party/xla/xla/service/llvm_compiler.h   |   4 +-
 .../xla/service/logical_buffer_analysis.cc    |  31 +-
 .../xla/xla/service/logical_buffer_analysis.h |  29 +-
 third_party/xla/xla/service/map_inliner.cc    |   6 +-
 .../xla/service/mapped_ptr_container_sorter.h |  15 +-
 .../service/memory_space_propagation_test.cc  |   4 +-
 .../xla/service/p2p_schedule_preparation.cc   |  50 +--
 third_party/xla/xla/service/qr_expander.cc    |   4 +-
 .../xla/xla/service/reduce_decomposer.cc      |   4 +-
 .../xla/service/reduce_scatter_combiner.cc    |   3 +-
 .../xla/xla/service/reduce_window_rewriter.cc |   5 +-
 .../xla/xla/service/reduce_window_rewriter.h  |   4 +-
 .../xla/xla/service/reshape_decomposer.cc     |   2 +-
 .../xla/xla/service/reshape_mover_test.cc     |   4 +-
 third_party/xla/xla/service/service.cc        |  78 +++--
 third_party/xla/xla/service/service.h         |  65 ++--
 .../service/service_executable_run_options.h  |   6 +-
 .../xla/xla/service/shape_inference.cc        |  18 +-
 .../xla/xla/service/shaped_buffer_test.cc     |   3 +-
 .../xla/xla/service/sharding_propagation.cc   |   4 +-
 .../xla/xla/service/sharding_propagation.h    |   4 +-
 third_party/xla/xla/service/slice_sinker.cc   |   5 +-
 third_party/xla/xla/service/source_map_util.h |  13 +-
 .../xla/service/space_to_batch_converter.cc   |  32 +-
 .../service/stochastic_convert_decomposer.cc  |   4 +-
 .../xla/xla/service/sub_byte_normalization.cc |   2 +-
 third_party/xla/xla/service/topk_rewriter.cc  |   8 +-
 .../xla/xla/service/tpu_computation_placer.cc |  13 +-
 .../xla/xla/service/tpu_computation_placer.h  |  13 +-
 .../xla/xla/service/transfer_manager.cc       |  29 +-
 .../xla/xla/service/transfer_manager.h        |  46 +--
 .../xla/xla/service/transpose_folding.cc      |   2 +-
 .../xla/service/tree_reduction_rewriter.cc    |   2 +-
 .../xla/service/tuple_points_to_analysis.cc   |  49 +--
 .../xla/service/tuple_points_to_analysis.h    |  49 +--
 .../while_loop_all_reduce_code_motion.cc      |   4 +-
 .../service/while_loop_concat_code_motion.cc  |  14 +-
 .../service/while_loop_constant_sinking.cc    |   7 +-
 157 files changed, 1874 insertions(+), 1717 deletions(-)

diff --git a/third_party/xla/xla/service/algebraic_simplifier.cc b/third_party/xla/xla/service/algebraic_simplifier.cc
index 0e22ed1cbd4a38..90c9a6812717a3 100644
--- a/third_party/xla/xla/service/algebraic_simplifier.cc
+++ b/third_party/xla/xla/service/algebraic_simplifier.cc
@@ -556,7 +556,7 @@ std::unique_ptr<HloInstruction> MakeScalarInstruction(HloInstruction* target,
 
 }  // namespace
 
-Status AlgebraicSimplifierVisitor::ScalarMultiplyReduction(
+absl::Status AlgebraicSimplifierVisitor::ScalarMultiplyReduction(
     HloInstruction* dot) {
   // We only process bfloat16 and float32 for now.
   if (dot->shape().element_type() != BF16 &&
@@ -778,7 +778,7 @@ bool AlgebraicSimplifierVisitor::ReplaceInstructionIfCompatible(
       .value();
 }
 
-Status AlgebraicSimplifierVisitor::HandleAbs(HloInstruction* abs) {
+absl::Status AlgebraicSimplifierVisitor::HandleAbs(HloInstruction* abs) {
   HloInstruction* abs_operand = abs->mutable_operand(0);
   VLOG(10) << "trying transform [Abs(A) => A] " << abs->ToString()
            << " Abs operand is: " << abs_operand->ToString();
@@ -788,7 +788,7 @@ Status AlgebraicSimplifierVisitor::HandleAbs(HloInstruction* abs) {
   return OkStatus();
 }
 
-Status AlgebraicSimplifierVisitor::HandleAdd(HloInstruction* add) {
+absl::Status AlgebraicSimplifierVisitor::HandleAdd(HloInstruction* add) {
   HloInstruction *lhs, *rhs;
   CHECK(Match(add, m::Add(m::Op(&lhs), m::Op(&rhs))));
 
@@ -1125,7 +1125,8 @@ absl::StatusOr<bool> AlgebraicSimplifierVisitor::TrySimplifyTautologicalCompare(
   return false;
 }
 
-Status AlgebraicSimplifierVisitor::HandleAllToAll(HloInstruction* all_to_all) {
+absl::Status AlgebraicSimplifierVisitor::HandleAllToAll(
+    HloInstruction* all_to_all) {
   if (all_to_all->shape().IsArray() &&
       Match(all_to_all->mutable_operand(0),
             m::Broadcast(m::ConstantScalar()))) {
@@ -1134,7 +1135,8 @@ Status AlgebraicSimplifierVisitor::HandleAllToAll(HloInstruction* all_to_all) {
   return OkStatus();
 }
 
-Status AlgebraicSimplifierVisitor::HandleAnd(HloInstruction* logical_and) {
+absl::Status AlgebraicSimplifierVisitor::HandleAnd(
+    HloInstruction* logical_and) {
   HloInstruction *lhs, *rhs;
   CHECK(Match(logical_and, m::And(m::Op(&lhs), m::Op(&rhs))));
   // Simplify logical and
@@ -1178,7 +1180,8 @@ Status AlgebraicSimplifierVisitor::HandleAnd(HloInstruction* logical_and) {
   return OkStatus();
 }
 
-Status AlgebraicSimplifierVisitor::HandleBitcast(HloInstruction* bitcast) {
+absl::Status AlgebraicSimplifierVisitor::HandleBitcast(
+    HloInstruction* bitcast) {
   // It's tricky for the simplifier to determine whether
   // it should remove the op when control deps are present. I.e.
   // control deps might be added to preserve a certain order.
@@ -1479,7 +1482,7 @@ bool AlgebraicSimplifierVisitor::SwapCopyBitcastCopy(
   return false;
 }
 
-Status AlgebraicSimplifierVisitor::HandleBitcastConvert(
+absl::Status AlgebraicSimplifierVisitor::HandleBitcastConvert(
     HloInstruction* bitcast) {
   TF_ASSIGN_OR_RETURN(bool replaced,
                       TrySimplifyTautologicalBitcastConvert(bitcast));
@@ -1491,7 +1494,7 @@ Status AlgebraicSimplifierVisitor::HandleBitcastConvert(
   return OkStatus();
 }
 
-Status AlgebraicSimplifierVisitor::HandleCopy(HloInstruction* copy) {
+absl::Status AlgebraicSimplifierVisitor::HandleCopy(HloInstruction* copy) {
   if (SwapCopyBitcastCopy(copy)) {
     return OkStatus();
   }
@@ -1633,7 +1636,7 @@ Status AlgebraicSimplifierVisitor::HandleCopy(HloInstruction* copy) {
   return OkStatus();
 }
 
-Status AlgebraicSimplifierVisitor::HandleConcatenate(
+absl::Status AlgebraicSimplifierVisitor::HandleConcatenate(
     HloInstruction* concatenate) {
   absl::Span<HloInstruction* const> operands(concatenate->operands());
   if (operands.size() == 1) {
@@ -1843,7 +1846,7 @@ AlgebraicSimplifierVisitor::TrySimplifyTautologicalBitcastConvert(
   return true;
 }
 
-Status
+absl::Status
 AlgebraicSimplifierVisitor::TryRemoveUpcastAndDowncastSurroundingBinaryOp(
     HloInstruction* convert_instruction) {
   HloInstruction* arg_1 = nullptr;
@@ -1939,7 +1942,8 @@ static HloInstruction* BuildTupleConstant(HloComputation* computation,
   }
 }
 
-Status AlgebraicSimplifierVisitor::HandleConstant(HloInstruction* constant) {
+absl::Status AlgebraicSimplifierVisitor::HandleConstant(
+    HloInstruction* constant) {
   // Tuple constants aren't directly supported by any backend. Expand them into
   // explicit Tuple instructions.
   if (constant->shape().IsTuple()) {
@@ -1985,7 +1989,7 @@ Status AlgebraicSimplifierVisitor::HandleConstant(HloInstruction* constant) {
   return OkStatus();
 }
 
-Status AlgebraicSimplifierVisitor::HandleSubtract(HloInstruction* sub) {
+absl::Status AlgebraicSimplifierVisitor::HandleSubtract(HloInstruction* sub) {
   HloInstruction *lhs, *rhs;
   CHECK(Match(sub, m::Subtract(m::Op(&lhs), m::Op(&rhs))));
   // A - 0 => A
@@ -2021,7 +2025,7 @@ Status AlgebraicSimplifierVisitor::HandleSubtract(HloInstruction* sub) {
 }
 namespace {
 template <typename T>
-Status InvertConstant(const HloInstruction& constant, Literal* result) {
+absl::Status InvertConstant(const HloInstruction& constant, Literal* result) {
   return result->Populate<T>([&](absl::Span<const int64_t> indices) {
     return T{1.0} / constant.literal().Get<T>(indices);
   });
@@ -2084,7 +2088,7 @@ std::unique_ptr<HloInstruction> TryDivideToShift(
 }
 }  // namespace
 
-Status AlgebraicSimplifierVisitor::HandleDivide(HloInstruction* divide) {
+absl::Status AlgebraicSimplifierVisitor::HandleDivide(HloInstruction* divide) {
   HloInstruction *a, *b, *c, *d;
   CHECK(Match(divide, m::Divide(m::Op(&a), m::Op(&b))));
   // A/1 => A
@@ -2181,7 +2185,7 @@ Status AlgebraicSimplifierVisitor::HandleDivide(HloInstruction* divide) {
     Shape result_shape = c->literal().shape();
     Literal new_literal(result_shape);
     return primitive_util::PrimitiveTypeSwitch<Status>(
-        [&](auto primitive_type_constant) -> Status {
+        [&](auto primitive_type_constant) -> absl::Status {
           if constexpr (primitive_util::IsFloatingPointType(
                             primitive_type_constant) ||
                         primitive_util::IsComplexType(
@@ -2373,7 +2377,7 @@ AlgebraicSimplifierVisitor::RemoveDegenerateDimensionFromDot(
 // We want to transform this into
 //
 //  bcast' = f32[2,3,1,4] broadcast(f32[2,4] x), dimensions={0,3}
-Status AlgebraicSimplifierVisitor::SimplifyTransposeOfBroadcast(
+absl::Status AlgebraicSimplifierVisitor::SimplifyTransposeOfBroadcast(
     HloInstruction* transpose, absl::Span<const int64_t> dimensions) {
   HloInstruction* broadcast = transpose->mutable_operand(0);
   if (broadcast->opcode() != HloOpcode::kBroadcast ||
@@ -3295,7 +3299,7 @@ AlgebraicSimplifierVisitor::AssociativeReorderDotOperator(
   return nullptr;
 }
 
-Status AlgebraicSimplifierVisitor::HandleDot(HloInstruction* dot) {
+absl::Status AlgebraicSimplifierVisitor::HandleDot(HloInstruction* dot) {
   CHECK(computation_ == dot->parent());
   HloDotInstruction* dot_cast = Cast<HloDotInstruction>(dot);
   const auto& dnums = dot->dot_dimension_numbers();
@@ -3858,7 +3862,7 @@ std::vector<int64_t> GetPaddedDims(const HloInstruction* pad) {
 }
 }  // namespace
 
-Status AlgebraicSimplifierVisitor::HandleGather(HloInstruction* gather) {
+absl::Status AlgebraicSimplifierVisitor::HandleGather(HloInstruction* gather) {
   const Shape& operand_shape = gather->operand(0)->shape();
   if (ShapeUtil::IsZeroElementArray(operand_shape)) {
     return ReplaceInstruction(gather, MakeScalarLike(gather, 0));
@@ -4135,7 +4139,8 @@ absl::StatusOr<std::unique_ptr<HloInstruction>> MinMaxToClamp(
 }
 }  // namespace
 
-Status AlgebraicSimplifierVisitor::HandleMaximum(HloInstruction* maximum) {
+absl::Status AlgebraicSimplifierVisitor::HandleMaximum(
+    HloInstruction* maximum) {
   HloInstruction *lhs, *rhs;
   CHECK(Match(maximum, m::Maximum(m::Op(&lhs), m::Op(&rhs))));
 
@@ -4207,7 +4212,8 @@ Status AlgebraicSimplifierVisitor::HandleMaximum(HloInstruction* maximum) {
   return OkStatus();
 }
 
-Status AlgebraicSimplifierVisitor::HandleMinimum(HloInstruction* minimum) {
+absl::Status AlgebraicSimplifierVisitor::HandleMinimum(
+    HloInstruction* minimum) {
   HloInstruction *lhs, *rhs;
   CHECK(Match(minimum, m::Minimum(m::Op(&lhs), m::Op(&rhs))));
 
@@ -4264,7 +4270,7 @@ Status AlgebraicSimplifierVisitor::HandleMinimum(HloInstruction* minimum) {
   return OkStatus();
 }
 
-Status AlgebraicSimplifierVisitor::HandleClamp(HloInstruction* clamp) {
+absl::Status AlgebraicSimplifierVisitor::HandleClamp(HloInstruction* clamp) {
   HloInstruction* clamp_lower_bound;
   HloInstruction* clamp_upper_bound;
   HloInstruction* to_clamp;
@@ -4301,7 +4307,8 @@ Status AlgebraicSimplifierVisitor::HandleClamp(HloInstruction* clamp) {
   return OkStatus();
 }
 
-Status AlgebraicSimplifierVisitor::HandleMultiply(HloInstruction* multiply) {
+absl::Status AlgebraicSimplifierVisitor::HandleMultiply(
+    HloInstruction* multiply) {
   HloInstruction *lhs, *rhs;
   CHECK(Match(multiply, m::Multiply(m::Op(&lhs), m::Op(&rhs))));
   // LHS*1 => LHS
@@ -4501,7 +4508,7 @@ Status AlgebraicSimplifierVisitor::HandleMultiply(HloInstruction* multiply) {
   return OkStatus();
 }
 
-Status AlgebraicSimplifierVisitor::HandleNegate(HloInstruction* negate) {
+absl::Status AlgebraicSimplifierVisitor::HandleNegate(HloInstruction* negate) {
   // negate(negate(x)) => x
   HloInstruction* x;
   if (Match(negate, m::Negate(m::Negate(m::Op(&x)))) &&
@@ -4511,7 +4518,8 @@ Status AlgebraicSimplifierVisitor::HandleNegate(HloInstruction* negate) {
   return OkStatus();
 }
 
-Status AlgebraicSimplifierVisitor::HandleNot(HloInstruction* logical_not) {
+absl::Status AlgebraicSimplifierVisitor::HandleNot(
+    HloInstruction* logical_not) {
   // not(not(x)) => x
   HloInstruction* x;
   if (Match(logical_not, m::Not(m::Not(m::Op(&x)))) &&
@@ -4521,7 +4529,7 @@ Status AlgebraicSimplifierVisitor::HandleNot(HloInstruction* logical_not) {
   return OkStatus();
 }
 
-Status AlgebraicSimplifierVisitor::HandleOr(HloInstruction* logical_or) {
+absl::Status AlgebraicSimplifierVisitor::HandleOr(HloInstruction* logical_or) {
   HloInstruction *lhs, *rhs;
   CHECK(Match(logical_or, m::Or(m::Op(&lhs), m::Op(&rhs))));
 
@@ -4557,7 +4565,7 @@ Status AlgebraicSimplifierVisitor::HandleOr(HloInstruction* logical_or) {
   return OkStatus();
 }
 
-Status AlgebraicSimplifierVisitor::HandleLog(HloInstruction* log) {
+absl::Status AlgebraicSimplifierVisitor::HandleLog(HloInstruction* log) {
   // ln(exp(A)) => A
   VLOG(10) << "trying transform [ln(exp(A)) => A]: " << log->ToString();
   HloInstruction *a, *b;
@@ -4607,7 +4615,7 @@ Status AlgebraicSimplifierVisitor::HandleLog(HloInstruction* log) {
   return OkStatus();
 }
 
-Status AlgebraicSimplifierVisitor::HandleGetTupleElement(
+absl::Status AlgebraicSimplifierVisitor::HandleGetTupleElement(
     HloInstruction* get_tuple_element) {
   auto operand = get_tuple_element->mutable_operand(0);
   if (operand->opcode() == HloOpcode::kTuple) {
@@ -4624,7 +4632,7 @@ Status AlgebraicSimplifierVisitor::HandleGetTupleElement(
   return OkStatus();
 }
 
-Status AlgebraicSimplifierVisitor::HandleOptimizationBarrier(
+absl::Status AlgebraicSimplifierVisitor::HandleOptimizationBarrier(
     HloInstruction* barrier) {
   if (!barrier->shape().IsTuple() ||
       barrier == computation_->root_instruction()) {
@@ -4743,7 +4751,8 @@ bool OutputIsSubsetOfOperandElements(HloInstruction* instruction,
 
 }  // namespace
 
-Status AlgebraicSimplifierVisitor::HandleBroadcast(HloInstruction* broadcast) {
+absl::Status AlgebraicSimplifierVisitor::HandleBroadcast(
+    HloInstruction* broadcast) {
   HloInstruction* operand;
   CHECK(Match(broadcast, m::Broadcast(m::Op(&operand))));
   auto dims = *broadcast->mutable_dimensions();
@@ -4867,7 +4876,8 @@ Status AlgebraicSimplifierVisitor::HandleBroadcast(HloInstruction* broadcast) {
   return OkStatus();
 }
 
-Status AlgebraicSimplifierVisitor::HandleCompare(HloInstruction* compare) {
+absl::Status AlgebraicSimplifierVisitor::HandleCompare(
+    HloInstruction* compare) {
   HloInstruction* lhs;
   HloInstruction* rhs;
   CHECK(Match(compare, m::Compare(m::Op(&lhs), m::Op(&rhs))));
@@ -5013,7 +5023,8 @@ Status AlgebraicSimplifierVisitor::HandleCompare(HloInstruction* compare) {
   return OkStatus();
 }
 
-Status AlgebraicSimplifierVisitor::HandleConvert(HloInstruction* convert) {
+absl::Status AlgebraicSimplifierVisitor::HandleConvert(
+    HloInstruction* convert) {
   PrimitiveType src_type = convert->operand(0)->shape().element_type();
   PrimitiveType dest_type = convert->shape().element_type();
   // A conversion to the same element type as the operand is a nop and can be
@@ -5042,7 +5053,7 @@ Status AlgebraicSimplifierVisitor::HandleConvert(HloInstruction* convert) {
   return TryRemoveUpcastAndDowncastSurroundingBinaryOp(convert);
 }
 
-Status AlgebraicSimplifierVisitor::HandleCustomCall(
+absl::Status AlgebraicSimplifierVisitor::HandleCustomCall(
     HloInstruction* custom_call) {
   // Remove redundant slice to dynamic of pad to static
   HloInstruction *pad_to_static0, *pad_to_static1, *pad_to_static_operand;
@@ -5064,7 +5075,8 @@ Status AlgebraicSimplifierVisitor::HandleCustomCall(
 }
 
 // Complex(Real(c), Imag(c)) -> c
-Status AlgebraicSimplifierVisitor::HandleComplex(HloInstruction* complex) {
+absl::Status AlgebraicSimplifierVisitor::HandleComplex(
+    HloInstruction* complex) {
   HloInstruction *c0, *c1;
   if (Match(complex, m::Complex(m::Real(m::Op(&c0)), m::Imag(m::Op(&c1)))) &&
       c0 == c1) {
@@ -5074,7 +5086,7 @@ Status AlgebraicSimplifierVisitor::HandleComplex(HloInstruction* complex) {
 }
 
 // Real(Complex(r, i)) -> r
-Status AlgebraicSimplifierVisitor::HandleReal(HloInstruction* real) {
+absl::Status AlgebraicSimplifierVisitor::HandleReal(HloInstruction* real) {
   HloInstruction* op;
   if (Match(real, m::Real(m::Complex(m::Op(&op), m::Op())))) {
     return ReplaceInstruction(real, op);
@@ -5083,7 +5095,7 @@ Status AlgebraicSimplifierVisitor::HandleReal(HloInstruction* real) {
 }
 
 // Imag(Complex(r, i)) -> i
-Status AlgebraicSimplifierVisitor::HandleImag(HloInstruction* imag) {
+absl::Status AlgebraicSimplifierVisitor::HandleImag(HloInstruction* imag) {
   HloInstruction* op;
   if (Match(imag, m::Imag(m::Complex(m::Op(), m::Op(&op))))) {
     return ReplaceInstruction(imag, op);
@@ -5091,7 +5103,8 @@ Status AlgebraicSimplifierVisitor::HandleImag(HloInstruction* imag) {
   return OkStatus();
 }
 
-Status AlgebraicSimplifierVisitor::HandleIota(HloInstruction* instruction) {
+absl::Status AlgebraicSimplifierVisitor::HandleIota(
+    HloInstruction* instruction) {
   // iota -> zero if the iota dimension never produces an element other than
   // zero.
   auto* iota = Cast<HloIotaInstruction>(instruction);
@@ -5101,7 +5114,7 @@ Status AlgebraicSimplifierVisitor::HandleIota(HloInstruction* instruction) {
   return OkStatus();
 }
 
-Status AlgebraicSimplifierVisitor::HandlePad(HloInstruction* pad) {
+absl::Status AlgebraicSimplifierVisitor::HandlePad(HloInstruction* pad) {
   if (ShapeUtil::IsZeroElementArray(pad->operand(0)->shape())) {
     return ReplaceWithNewInstruction(
         pad, HloInstruction::CreateBroadcast(pad->shape(),
@@ -5304,7 +5317,7 @@ Status AlgebraicSimplifierVisitor::HandlePad(HloInstruction* pad) {
   return OkStatus();
 }
 
-Status AlgebraicSimplifierVisitor::HandlePower(HloInstruction* power) {
+absl::Status AlgebraicSimplifierVisitor::HandlePower(HloInstruction* power) {
   VLOG(10) << "trying transform [pow(A, 0) => 1]: " << power->ToString();
   HloInstruction *lhs, *rhs;
   CHECK(Match(power, m::Power(m::Op(&lhs), m::Op(&rhs))));
@@ -5521,7 +5534,8 @@ std::unique_ptr<HloInstruction> TryRemainderToAnd(
 }
 }  // namespace
 
-Status AlgebraicSimplifierVisitor::HandleRemainder(HloInstruction* remainder) {
+absl::Status AlgebraicSimplifierVisitor::HandleRemainder(
+    HloInstruction* remainder) {
   HloInstruction *a, *b;
   CHECK(Match(remainder, m::Remainder(m::Op(&a), m::Op(&b))));
 
@@ -5607,7 +5621,8 @@ Status AlgebraicSimplifierVisitor::HandleRemainder(HloInstruction* remainder) {
   return OkStatus();
 }
 
-Status AlgebraicSimplifierVisitor::HandleReshape(HloInstruction* reshape) {
+absl::Status AlgebraicSimplifierVisitor::HandleReshape(
+    HloInstruction* reshape) {
   auto operand = reshape->mutable_operand(0);
 
   // Reshape directly to empty constant if the shape contains zero-element
@@ -5874,7 +5889,8 @@ int64_t CountElementsLessThan(absl::Span<const int64_t> elements,
   return count;
 }
 
-Status AlgebraicSimplifierVisitor::HandleReverse(HloInstruction* reverse) {
+absl::Status AlgebraicSimplifierVisitor::HandleReverse(
+    HloInstruction* reverse) {
   // When all the dimensions to reverse are trivial (i.e. the bound is 1),
   // there is nothing to be done.
   auto dim_is_one = [&](int64_t i) -> bool {
@@ -6132,7 +6148,7 @@ absl::StatusOr<bool> AlgebraicSimplifierVisitor::TryToReorderSliceAndReverse(
   return false;
 }
 
-Status AlgebraicSimplifierVisitor::HandleSlice(HloInstruction* slice) {
+absl::Status AlgebraicSimplifierVisitor::HandleSlice(HloInstruction* slice) {
   // Delete no-op slices, i.e. where shape = operand shape.
   if (ReplaceInstructionIfCompatible(slice, slice->mutable_operand(0))) {
     return OkStatus();
@@ -6486,7 +6502,7 @@ Status AlgebraicSimplifierVisitor::HandleSlice(HloInstruction* slice) {
   return OkStatus();
 }
 
-Status AlgebraicSimplifierVisitor::HandleRsqrt(HloInstruction* rsqrt) {
+absl::Status AlgebraicSimplifierVisitor::HandleRsqrt(HloInstruction* rsqrt) {
   VLOG(10) << "trying transform [rsqrt(Pow(A, -2)) => |A|] "
            << rsqrt->ToString();
   HloInstruction* rsqrt_operand = rsqrt->mutable_operand(0);
@@ -6511,7 +6527,7 @@ Status AlgebraicSimplifierVisitor::HandleRsqrt(HloInstruction* rsqrt) {
   return OkStatus();
 }
 
-Status AlgebraicSimplifierVisitor::HandleDynamicSlice(
+absl::Status AlgebraicSimplifierVisitor::HandleDynamicSlice(
     HloInstruction* dynamic_slice) {
   // Skip optimizations for async dynamic-slices.
   if (dynamic_slice->parent()->IsAsyncComputation()) {
@@ -6777,7 +6793,7 @@ Status AlgebraicSimplifierVisitor::HandleDynamicSlice(
   return OkStatus();
 }
 
-Status AlgebraicSimplifierVisitor::HandleDynamicUpdateSlice(
+absl::Status AlgebraicSimplifierVisitor::HandleDynamicUpdateSlice(
     HloInstruction* dynamic_update_slice) {
   // Skip optimizations for async dynamic update slices
   if (dynamic_update_slice->parent()->IsAsyncComputation()) {
@@ -7061,7 +7077,7 @@ static bool ReductionComputationsEquivalent(const HloComputation& a,
          category_a == category_b;
 }
 
-Status AlgebraicSimplifierVisitor::HandleReduce(HloInstruction* hlo) {
+absl::Status AlgebraicSimplifierVisitor::HandleReduce(HloInstruction* hlo) {
   HloReduceInstruction* reduce = Cast<HloReduceInstruction>(hlo);
   bool multi_output_reduce = reduce->shape().IsTuple();
   // For tuple reduce, we require all reduce shapes to be the same, up to the
@@ -7603,7 +7619,8 @@ Status AlgebraicSimplifierVisitor::HandleReduce(HloInstruction* hlo) {
   return OkStatus();
 }
 
-Status AlgebraicSimplifierVisitor::HandleReduceWindow(HloInstruction* hlo) {
+absl::Status AlgebraicSimplifierVisitor::HandleReduceWindow(
+    HloInstruction* hlo) {
   auto* reduce_window = Cast<HloReduceWindowInstruction>(hlo);
   const bool multi_output_reduce_window = reduce_window->shape().IsTuple();
   auto inputs = reduce_window->inputs();
@@ -7937,7 +7954,7 @@ Status AlgebraicSimplifierVisitor::HandleReduceWindow(HloInstruction* hlo) {
                          /*reduce_computation=*/function));
 }
 
-Status AlgebraicSimplifierVisitor::HandleSelect(HloInstruction* select) {
+absl::Status AlgebraicSimplifierVisitor::HandleSelect(HloInstruction* select) {
   // select(x, y, y) -> y.
   if (select->operand(1) == select->operand(2) &&
       ReplaceInstructionIfCompatible(select, select->mutable_operand(1))) {
@@ -8024,7 +8041,7 @@ Status AlgebraicSimplifierVisitor::HandleSelect(HloInstruction* select) {
   return OkStatus();
 }
 
-Status AlgebraicSimplifierVisitor::HandleScatter(HloInstruction* hlo) {
+absl::Status AlgebraicSimplifierVisitor::HandleScatter(HloInstruction* hlo) {
   auto* scatter = Cast<HloScatterInstruction>(hlo);
 
   if (absl::c_all_of(scatter->scatter_updates(),
@@ -8047,7 +8064,7 @@ Status AlgebraicSimplifierVisitor::HandleScatter(HloInstruction* hlo) {
   return OkStatus();
 }
 
-Status AlgebraicSimplifierVisitor::HandleSort(HloInstruction* sort) {
+absl::Status AlgebraicSimplifierVisitor::HandleSort(HloInstruction* sort) {
   auto operand = sort->mutable_operand(0);
   int64_t dimension_to_sort = sort->dimensions(0);
   if (ShapeUtil::IsZeroElementArray(operand->shape()) ||
@@ -8062,7 +8079,7 @@ Status AlgebraicSimplifierVisitor::HandleSort(HloInstruction* sort) {
   return OkStatus();
 }
 
-Status AlgebraicSimplifierVisitor::HandleSqrt(HloInstruction* sqrt) {
+absl::Status AlgebraicSimplifierVisitor::HandleSqrt(HloInstruction* sqrt) {
   VLOG(10) << "trying transform [sqrt(A*A) => |A|] " << sqrt->ToString();
   HloInstruction* sqrt_operand = sqrt->mutable_operand(0);
   if (sqrt_operand->opcode() == HloOpcode::kMultiply &&
@@ -8117,7 +8134,8 @@ bool IsPermutationOfIota(absl::Span<const int64_t> elems) {
 
 }  // namespace
 
-Status AlgebraicSimplifierVisitor::HandleTranspose(HloInstruction* transpose) {
+absl::Status AlgebraicSimplifierVisitor::HandleTranspose(
+    HloInstruction* transpose) {
   auto operand = transpose->mutable_operand(0);
   if (std::is_sorted(transpose->dimensions().begin(),
                      transpose->dimensions().end())) {
@@ -8914,7 +8932,7 @@ absl::StatusOr<bool> AlgebraicSimplifierVisitor::SimplifyConvToMultiply(
   return true;
 }
 
-Status AlgebraicSimplifierVisitor::HandleConvolution(
+absl::Status AlgebraicSimplifierVisitor::HandleConvolution(
     HloInstruction* convolution) {
   if (options_.enable_scalar_multiply_reduction()) {
     TF_RETURN_IF_ERROR(ScalarMultiplyReduction(convolution));
@@ -8957,7 +8975,7 @@ Status AlgebraicSimplifierVisitor::HandleConvolution(
   return OkStatus();
 }
 
-Status AlgebraicSimplifierVisitor::HandleMap(HloInstruction* map) {
+absl::Status AlgebraicSimplifierVisitor::HandleMap(HloInstruction* map) {
   auto* map_computation = map->to_apply();
   auto* map_root = map_computation->root_instruction();
   if (map_root->opcode() == HloOpcode::kParameter) {
diff --git a/third_party/xla/xla/service/algebraic_simplifier.h b/third_party/xla/xla/service/algebraic_simplifier.h
index 6ec43b9928dceb..b23250c644f60d 100644
--- a/third_party/xla/xla/service/algebraic_simplifier.h
+++ b/third_party/xla/xla/service/algebraic_simplifier.h
@@ -315,103 +315,104 @@ class AlgebraicSimplifierVisitor : public DfsHloRewriteVisitor {
                                       AlgebraicSimplifier* simplifier)
       : options_(options), simplifier_(simplifier) {}
 
-  Status HandleAbs(HloInstruction* abs) override;
+  absl::Status HandleAbs(HloInstruction* abs) override;
 
-  Status HandleAdd(HloInstruction* add) override;
+  absl::Status HandleAdd(HloInstruction* add) override;
 
-  Status HandleAllToAll(HloInstruction* all_to_all) override;
+  absl::Status HandleAllToAll(HloInstruction* all_to_all) override;
 
-  Status HandleAnd(HloInstruction* logical_and) override;
+  absl::Status HandleAnd(HloInstruction* logical_and) override;
 
-  Status HandleBitcast(HloInstruction* bitcast) override;
+  absl::Status HandleBitcast(HloInstruction* bitcast) override;
 
-  Status HandleBitcastConvert(HloInstruction* bitcast) override;
+  absl::Status HandleBitcastConvert(HloInstruction* bitcast) override;
 
-  Status HandleBroadcast(HloInstruction* broadcast) override;
+  absl::Status HandleBroadcast(HloInstruction* broadcast) override;
 
-  Status HandleCompare(HloInstruction* compare) override;
+  absl::Status HandleCompare(HloInstruction* compare) override;
 
-  Status HandleConcatenate(HloInstruction* concatenate) override;
+  absl::Status HandleConcatenate(HloInstruction* concatenate) override;
 
-  Status HandleConstant(HloInstruction* constant) override;
+  absl::Status HandleConstant(HloInstruction* constant) override;
 
-  Status HandleCopy(HloInstruction* copy) override;
+  absl::Status HandleCopy(HloInstruction* copy) override;
 
-  Status HandleConvert(HloInstruction* convert) override;
+  absl::Status HandleConvert(HloInstruction* convert) override;
 
-  Status HandleComplex(HloInstruction* complex) override;
+  absl::Status HandleComplex(HloInstruction* complex) override;
 
-  Status HandleCustomCall(HloInstruction* custom_call) override;
+  absl::Status HandleCustomCall(HloInstruction* custom_call) override;
 
-  Status HandleReal(HloInstruction* real) override;
+  absl::Status HandleReal(HloInstruction* real) override;
 
-  Status HandleImag(HloInstruction* imag) override;
+  absl::Status HandleImag(HloInstruction* imag) override;
 
-  Status HandleIota(HloInstruction* instruction) override;
+  absl::Status HandleIota(HloInstruction* instruction) override;
 
-  Status HandleConvolution(HloInstruction* convolution) override;
+  absl::Status HandleConvolution(HloInstruction* convolution) override;
 
-  Status HandleDivide(HloInstruction* divide) override;
+  absl::Status HandleDivide(HloInstruction* divide) override;
 
-  Status HandleDot(HloInstruction* dot) override;
+  absl::Status HandleDot(HloInstruction* dot) override;
 
-  Status HandleGather(HloInstruction* gather) override;
+  absl::Status HandleGather(HloInstruction* gather) override;
 
-  Status HandleGetTupleElement(HloInstruction* get_tuple_element) override;
+  absl::Status HandleGetTupleElement(
+      HloInstruction* get_tuple_element) override;
 
-  Status HandleLog(HloInstruction* log) override;
+  absl::Status HandleLog(HloInstruction* log) override;
 
-  Status HandleMaximum(HloInstruction* maximum) override;
+  absl::Status HandleMaximum(HloInstruction* maximum) override;
 
-  Status HandleMinimum(HloInstruction* minimum) override;
+  absl::Status HandleMinimum(HloInstruction* minimum) override;
 
-  Status HandleClamp(HloInstruction* clamp) override;
+  absl::Status HandleClamp(HloInstruction* clamp) override;
 
-  Status HandleMultiply(HloInstruction* multiply) override;
+  absl::Status HandleMultiply(HloInstruction* multiply) override;
 
-  Status HandleNegate(HloInstruction* negate) override;
+  absl::Status HandleNegate(HloInstruction* negate) override;
 
-  Status HandleNot(HloInstruction* logical_not) override;
+  absl::Status HandleNot(HloInstruction* logical_not) override;
 
-  Status HandleOptimizationBarrier(HloInstruction* barrier) override;
+  absl::Status HandleOptimizationBarrier(HloInstruction* barrier) override;
 
-  Status HandleOr(HloInstruction* logical_or) override;
+  absl::Status HandleOr(HloInstruction* logical_or) override;
 
-  Status HandlePad(HloInstruction* pad) override;
+  absl::Status HandlePad(HloInstruction* pad) override;
 
-  Status HandlePower(HloInstruction* power) override;
+  absl::Status HandlePower(HloInstruction* power) override;
 
-  Status HandleRemainder(HloInstruction* remainder) override;
+  absl::Status HandleRemainder(HloInstruction* remainder) override;
 
-  Status HandleReshape(HloInstruction* reshape) override;
+  absl::Status HandleReshape(HloInstruction* reshape) override;
 
-  Status HandleReduce(HloInstruction* hlo) override;
+  absl::Status HandleReduce(HloInstruction* hlo) override;
 
-  Status HandleReduceWindow(HloInstruction* hlo) override;
+  absl::Status HandleReduceWindow(HloInstruction* hlo) override;
 
-  Status HandleReverse(HloInstruction* reverse) override;
+  absl::Status HandleReverse(HloInstruction* reverse) override;
 
-  Status HandleRsqrt(HloInstruction* rsqrt) override;
+  absl::Status HandleRsqrt(HloInstruction* rsqrt) override;
 
-  Status HandleSlice(HloInstruction* slice) override;
+  absl::Status HandleSlice(HloInstruction* slice) override;
 
-  Status HandleSqrt(HloInstruction* sqrt) override;
+  absl::Status HandleSqrt(HloInstruction* sqrt) override;
 
-  Status HandleDynamicSlice(HloInstruction* dynamic_slice) override;
+  absl::Status HandleDynamicSlice(HloInstruction* dynamic_slice) override;
 
-  Status HandleDynamicUpdateSlice(
+  absl::Status HandleDynamicUpdateSlice(
       HloInstruction* dynamic_update_slice) override;
-  Status HandleScatter(HloInstruction* hlo) override;
+  absl::Status HandleScatter(HloInstruction* hlo) override;
 
-  Status HandleSelect(HloInstruction* select) override;
+  absl::Status HandleSelect(HloInstruction* select) override;
 
-  Status HandleSort(HloInstruction* sort) override;
+  absl::Status HandleSort(HloInstruction* sort) override;
 
-  Status HandleTranspose(HloInstruction* transpose) override;
+  absl::Status HandleTranspose(HloInstruction* transpose) override;
 
-  Status HandleSubtract(HloInstruction* sub) override;
+  absl::Status HandleSubtract(HloInstruction* sub) override;
 
-  Status HandleMap(HloInstruction* map) override;
+  absl::Status HandleMap(HloInstruction* map) override;
 
   // Runs the visitor on a computation.
   bool Run(HloComputation* computation,
@@ -455,8 +456,8 @@ class AlgebraicSimplifierVisitor : public DfsHloRewriteVisitor {
 
   // Moves the transpose to the broadcast if possible. Can also be called with a
   // bitcast transpose.
-  Status SimplifyTransposeOfBroadcast(HloInstruction* transpose,
-                                      absl::Span<const int64_t> dimensions);
+  absl::Status SimplifyTransposeOfBroadcast(
+      HloInstruction* transpose, absl::Span<const int64_t> dimensions);
 
   // Converts to primitive type if the input hlo is not that type, otherwise
   // returns the original hlo.
@@ -494,7 +495,7 @@ class AlgebraicSimplifierVisitor : public DfsHloRewriteVisitor {
 
   // Move scalar multiply to the smallest side of convolution to
   // reduce multiply computations.
-  Status ScalarMultiplyReduction(HloInstruction* dot);
+  absl::Status ScalarMultiplyReduction(HloInstruction* dot);
 
   // Convenience method for replacing an instruction with a bitcast. If operand
   // is not null, then the bitcast will use the specified operand instead of the
@@ -610,7 +611,7 @@ class AlgebraicSimplifierVisitor : public DfsHloRewriteVisitor {
   //  where TS is a smaller point type than TL (ex, TS=fp16, TL=fp32)
   // ->
   // bin_op<TS>(data1<TS>, data2<TS>)
-  Status TryRemoveUpcastAndDowncastSurroundingBinaryOp(
+  absl::Status TryRemoveUpcastAndDowncastSurroundingBinaryOp(
       HloInstruction* convert_instruction);
 
   // Useful when we want to use the same visitor over multiple computations.
diff --git a/third_party/xla/xla/service/all_gather_combiner.cc b/third_party/xla/xla/service/all_gather_combiner.cc
index e13de8c521e223..6d2e86cac466ad 100644
--- a/third_party/xla/xla/service/all_gather_combiner.cc
+++ b/third_party/xla/xla/service/all_gather_combiner.cc
@@ -73,8 +73,8 @@ int64_t FindMostFrequentGatherDim(
 // Combines the elements of to_combine into a single AllGather op. All entries
 // in to_combine must be AllGather ops with exactly one operand and the same
 // preferred all_gather_dimension.
-Status CombineAllGathers(absl::Span<HloInstruction* const> to_combine,
-                         bool combine_by_dim) {
+absl::Status CombineAllGathers(absl::Span<HloInstruction* const> to_combine,
+                               bool combine_by_dim) {
   if (to_combine.size() < 2) {
     return OkStatus();
   }
@@ -225,7 +225,7 @@ absl::StatusOr<bool> AllGatherCombiner::Run(
       return CombineKey(instruction, *domain_map, combine_by_dim_);
     };
     auto combine_fn =
-        [&](absl::Span<HloInstruction* const> to_combine) -> Status {
+        [&](absl::Span<HloInstruction* const> to_combine) -> absl::Status {
       return CombineAllGathers(to_combine, combine_by_dim_);
     };
 
diff --git a/third_party/xla/xla/service/all_gather_decomposer.cc b/third_party/xla/xla/service/all_gather_decomposer.cc
index 56c9115ef6f821..3f3183fe3920af 100644
--- a/third_party/xla/xla/service/all_gather_decomposer.cc
+++ b/third_party/xla/xla/service/all_gather_decomposer.cc
@@ -84,8 +84,8 @@ HloInstruction* AllGatherDecomposer::TranslateAllGatherToAllReducePerOperand(
   return ar;
 }
 
-Status AllGatherDecomposer::DecomposeAllGather(HloAllGatherInstruction* ag,
-                                               HloComputation* comp) {
+absl::Status AllGatherDecomposer::DecomposeAllGather(
+    HloAllGatherInstruction* ag, HloComputation* comp) {
   TF_ASSIGN_OR_RETURN(CollectiveOpGroupMode group_mode,
                       GetCollectiveOpGroupMode(ag->channel_id().has_value(),
                                                ag->use_global_device_ids()));
diff --git a/third_party/xla/xla/service/all_reduce_combiner.cc b/third_party/xla/xla/service/all_reduce_combiner.cc
index 333ab5d515b29f..6417717b13fef5 100644
--- a/third_party/xla/xla/service/all_reduce_combiner.cc
+++ b/third_party/xla/xla/service/all_reduce_combiner.cc
@@ -45,7 +45,7 @@ namespace {
 // Combines the elements of to_combine into a single AllReduce op. All
 // entries in to_combine must be AllReduce ops with exactly one operand
 // and the same reduction operation.
-Status CombineAllReduces(absl::Span<HloInstruction* const> to_combine) {
+absl::Status CombineAllReduces(absl::Span<HloInstruction* const> to_combine) {
   if (to_combine.size() < 2) {
     return OkStatus();
   }
diff --git a/third_party/xla/xla/service/all_reduce_contiguous.cc b/third_party/xla/xla/service/all_reduce_contiguous.cc
index 91613a9626088e..b40a4d04371316 100644
--- a/third_party/xla/xla/service/all_reduce_contiguous.cc
+++ b/third_party/xla/xla/service/all_reduce_contiguous.cc
@@ -28,7 +28,8 @@ limitations under the License.
 namespace xla {
 namespace {
 
-Status ReplaceWithContiguousAllReduce(HloAllReduceInstruction* all_reduce) {
+absl::Status ReplaceWithContiguousAllReduce(
+    HloAllReduceInstruction* all_reduce) {
   TF_RET_CHECK(all_reduce);
   TF_RET_CHECK(!all_reduce->has_sharding());
 
diff --git a/third_party/xla/xla/service/allocation_tracker.cc b/third_party/xla/xla/service/allocation_tracker.cc
index a2cbad64b75964..2bc816b0415968 100644
--- a/third_party/xla/xla/service/allocation_tracker.cc
+++ b/third_party/xla/xla/service/allocation_tracker.cc
@@ -94,7 +94,7 @@ absl::StatusOr<GlobalDataHandle> AllocationTracker::RegisterInternal(
   return result;
 }
 
-Status AllocationTracker::Unregister(const GlobalDataHandle& data) {
+absl::Status AllocationTracker::Unregister(const GlobalDataHandle& data) {
   absl::MutexLock lock(&mutex_);
   VLOG(2) << "Unregister("
           << "handle: " << data.handle() << ")";
@@ -218,8 +218,8 @@ void AllocationTracker::AddAllocationOrIncrementRefCount(
   }
 }
 
-Status AllocationTracker::DecrementRefCount(se::DeviceMemoryBase device_memory,
-                                            int device_ordinal) {
+absl::Status AllocationTracker::DecrementRefCount(
+    se::DeviceMemoryBase device_memory, int device_ordinal) {
   AllocationMap& allocation_map = opaque_to_allocation_map_[device_ordinal];
   auto it = allocation_map.find(device_memory.opaque());
   TF_RET_CHECK(it != allocation_map.end());
diff --git a/third_party/xla/xla/service/allocation_tracker.h b/third_party/xla/xla/service/allocation_tracker.h
index c8359c0fc30971..035641e8c35eb5 100644
--- a/third_party/xla/xla/service/allocation_tracker.h
+++ b/third_party/xla/xla/service/allocation_tracker.h
@@ -52,7 +52,7 @@ class AllocationTracker {
       const std::string& tag);
 
   // Unregister the allocation for the given data handle.
-  Status Unregister(const GlobalDataHandle& data);
+  absl::Status Unregister(const GlobalDataHandle& data);
 
   // Returns a vector of global data handles that point to the tuple elements.
   absl::StatusOr<std::vector<GlobalDataHandle>> DeconstructTuple(
@@ -103,8 +103,8 @@ class AllocationTracker {
 
   // Decrements the reference count of the given device memory. Then, if it is
   // zero, deallocate the memory.
-  Status DecrementRefCount(se::DeviceMemoryBase device_memory,
-                           int device_ordinal)
+  absl::Status DecrementRefCount(se::DeviceMemoryBase device_memory,
+                                 int device_ordinal)
       ABSL_EXCLUSIVE_LOCKS_REQUIRED(mutex_);
 
   // A map from device memory opaque value to allocation. One such map is
diff --git a/third_party/xla/xla/service/ar_crs_combiner.cc b/third_party/xla/xla/service/ar_crs_combiner.cc
index 1ad71ca8d7e56c..e64d746861ab7e 100644
--- a/third_party/xla/xla/service/ar_crs_combiner.cc
+++ b/third_party/xla/xla/service/ar_crs_combiner.cc
@@ -462,7 +462,7 @@ void ArCrsCombiner::GroupAllReducesById(HloModule* module) {
   }
 }
 
-Status ArCrsCombiner::KeepProvablyEqualInstructionGroupsMPMD() {
+absl::Status ArCrsCombiner::KeepProvablyEqualInstructionGroupsMPMD() {
   for (auto it = all_reduce_map_.begin(); it != all_reduce_map_.end();) {
     auto copy_it = it++;  // Advance `it` before invalidation from erase.
     auto channel_id = copy_it->first;
@@ -496,7 +496,7 @@ Status ArCrsCombiner::KeepProvablyEqualInstructionGroupsMPMD() {
   return OkStatus();
 }
 
-Status ArCrsCombiner::KeepProvablyEqualInstructionGroupsSPMD(
+absl::Status ArCrsCombiner::KeepProvablyEqualInstructionGroupsSPMD(
     HloModule* module) {
   // For SPMD mode, use HloReplicationAnalysis to figure out HLO value
   // equivalence across partitions.
diff --git a/third_party/xla/xla/service/ar_crs_combiner.h b/third_party/xla/xla/service/ar_crs_combiner.h
index 7b537b7dd87429..925dcb9d8b2784 100644
--- a/third_party/xla/xla/service/ar_crs_combiner.h
+++ b/third_party/xla/xla/service/ar_crs_combiner.h
@@ -154,10 +154,10 @@ class ArCrsCombiner : public HloModulePass {
 
   // Looks at each AllReduce group in all_reduce_map_, and keeps only the
   // groups for which it's safe to move the AllReduce later in the HLO graph.
-  Status KeepProvablyEqualInstructionGroupsMPMD();
+  absl::Status KeepProvablyEqualInstructionGroupsMPMD();
 
   // Same as above, but runs on SPMD partitioned module instead of MPMD.
-  Status KeepProvablyEqualInstructionGroupsSPMD(HloModule* module);
+  absl::Status KeepProvablyEqualInstructionGroupsSPMD(HloModule* module);
 
   // Performs the graph rewrite that eliminates the early AllReduce and turns
   // the later CRS into an AllReduce.
diff --git a/third_party/xla/xla/service/backend.cc b/third_party/xla/xla/service/backend.cc
index efccee7f73e738..b3999ae2364001 100644
--- a/third_party/xla/xla/service/backend.cc
+++ b/third_party/xla/xla/service/backend.cc
@@ -213,7 +213,7 @@ absl::StatusOr<bool> Backend::devices_equivalent(int device_ordinal_a,
           executor_b->GetDeviceDescription().name());
 }
 
-Status Backend::ResetDevices() {
+absl::Status Backend::ResetDevices() {
   return transfer_manager_->ResetDevices(stream_executors_);
 }
 
diff --git a/third_party/xla/xla/service/backend.h b/third_party/xla/xla/service/backend.h
index fb8a324a4d320d..0c919517d98bb7 100644
--- a/third_party/xla/xla/service/backend.h
+++ b/third_party/xla/xla/service/backend.h
@@ -168,7 +168,7 @@ class Backend {
   tsl::thread::ThreadPool* eigen_intra_op_thread_pool() const;
 
   // Resets the devices associated with this backend.
-  Status ResetDevices();
+  absl::Status ResetDevices();
 
  private:
   Backend(se::Platform* platform, Compiler* compiler,
diff --git a/third_party/xla/xla/service/batchnorm_expander.cc b/third_party/xla/xla/service/batchnorm_expander.cc
index 592cb4a210cfe0..4fcd3b22b6d2ef 100644
--- a/third_party/xla/xla/service/batchnorm_expander.cc
+++ b/third_party/xla/xla/service/batchnorm_expander.cc
@@ -52,11 +52,11 @@ using std::optional;
 // operations into smaller operations.
 class BatchNormExpanderVisitor : public DfsHloRewriteVisitor {
  public:
-  Status HandleBatchNormTraining(HloInstruction* batch_norm) override;
+  absl::Status HandleBatchNormTraining(HloInstruction* batch_norm) override;
 
-  Status HandleBatchNormInference(HloInstruction* batch_norm) override;
+  absl::Status HandleBatchNormInference(HloInstruction* batch_norm) override;
 
-  Status HandleBatchNormGrad(HloInstruction* batch_norm) override;
+  absl::Status HandleBatchNormGrad(HloInstruction* batch_norm) override;
 
   // Runs the visitor on a computation.
   static bool Run(HloComputation* computation, bool rewrite_training_op,
@@ -150,7 +150,7 @@ bool BatchNormExpanderVisitor::Run(HloComputation* computation,
   return visitor.changed();
 }
 
-Status BatchNormExpanderVisitor::HandleBatchNormTraining(
+absl::Status BatchNormExpanderVisitor::HandleBatchNormTraining(
     HloInstruction* batch_norm) {
   if (!rewrite_training_op_) {
     return OkStatus();
@@ -298,7 +298,7 @@ Status BatchNormExpanderVisitor::HandleBatchNormTraining(
   return OkStatus();
 }
 
-Status BatchNormExpanderVisitor::HandleBatchNormInference(
+absl::Status BatchNormExpanderVisitor::HandleBatchNormInference(
     HloInstruction* batch_norm) {
   if (!rewrite_inference_op_) {
     return OkStatus();
@@ -390,7 +390,7 @@ Status BatchNormExpanderVisitor::HandleBatchNormInference(
   return OkStatus();
 }
 
-Status BatchNormExpanderVisitor::HandleBatchNormGrad(
+absl::Status BatchNormExpanderVisitor::HandleBatchNormGrad(
     HloInstruction* batch_norm) {
   // Use the following formulas to calculate gradients:
   // scale_grad =
diff --git a/third_party/xla/xla/service/bfloat16_conversion_folding.cc b/third_party/xla/xla/service/bfloat16_conversion_folding.cc
index c8e94d576c5c0b..9e39947c3341ae 100644
--- a/third_party/xla/xla/service/bfloat16_conversion_folding.cc
+++ b/third_party/xla/xla/service/bfloat16_conversion_folding.cc
@@ -36,10 +36,10 @@ class BFloat16ConversionFoldingVisitor : public DfsHloVisitorWithDefault {
         bfloat16_support_(bfloat16_support),
         bfloat16_conversion_folding_(bfloat16_conversion_folding) {}
 
-  Status DefaultAction(HloInstruction* hlo) override;
+  absl::Status DefaultAction(HloInstruction* hlo) override;
 
   // Special handling for all-reduce which can have a tuple output.
-  Status HandleAllReduce(HloInstruction* crs) override;
+  absl::Status HandleAllReduce(HloInstruction* crs) override;
 
   static bool Run(HloComputation* computation,
                   const FloatSupport* bfloat16_support,
@@ -53,17 +53,18 @@ class BFloat16ConversionFoldingVisitor : public DfsHloVisitorWithDefault {
  private:
   // Checks if the HLO has a BF16 -> F32 conversion as input, or a F32 -> BF16
   // conversion as output, and folds them to the HLO itself if feasible.
-  Status TryFoldBF16Conversions(HloInstruction* hlo);
+  absl::Status TryFoldBF16Conversions(HloInstruction* hlo);
 
   // Folds the F32 -> BF16 conversions from the HLO's output.
   //
   // Precondition: all of the HLO's users are F32 -> BF16 conversions.
-  Status FoldOutputConversions(HloInstruction* hlo);
+  absl::Status FoldOutputConversions(HloInstruction* hlo);
 
   // Folds the BF16 -> F32 conversion operand to the HLO.
   //
   // Precondition: the operand is a BF16 -> F32 conversion.
-  Status FoldOperandConversion(HloInstruction* hlo, int64_t operand_index);
+  absl::Status FoldOperandConversion(HloInstruction* hlo,
+                                     int64_t operand_index);
 
   HloComputation* computation_;
   const FloatSupport* bfloat16_support_;
@@ -71,7 +72,7 @@ class BFloat16ConversionFoldingVisitor : public DfsHloVisitorWithDefault {
   bool changed_ = false;
 };
 
-Status BFloat16ConversionFoldingVisitor::FoldOutputConversions(
+absl::Status BFloat16ConversionFoldingVisitor::FoldOutputConversions(
     HloInstruction* hlo) {
   std::vector<HloInstruction*> materialized_users = hlo->users();
   hlo->mutable_shape()->set_element_type(BF16);
@@ -84,7 +85,7 @@ Status BFloat16ConversionFoldingVisitor::FoldOutputConversions(
   return OkStatus();
 }
 
-Status BFloat16ConversionFoldingVisitor::FoldOperandConversion(
+absl::Status BFloat16ConversionFoldingVisitor::FoldOperandConversion(
     HloInstruction* hlo, int64_t operand_index) {
   // The operand is a convert from BF16 to F32.
   auto operand = hlo->mutable_operand(operand_index);
@@ -114,7 +115,7 @@ bool AllUsersAreF32ToBF16Converts(const HloInstruction* hlo) {
 
 }  // namespace
 
-Status BFloat16ConversionFoldingVisitor::TryFoldBF16Conversions(
+absl::Status BFloat16ConversionFoldingVisitor::TryFoldBF16Conversions(
     HloInstruction* hlo) {
   std::vector<int64_t> bf16_to_f32_operands;
   bool has_other_f32_operands = false;
@@ -157,7 +158,8 @@ Status BFloat16ConversionFoldingVisitor::TryFoldBF16Conversions(
   return OkStatus();
 }
 
-Status BFloat16ConversionFoldingVisitor::DefaultAction(HloInstruction* hlo) {
+absl::Status BFloat16ConversionFoldingVisitor::DefaultAction(
+    HloInstruction* hlo) {
   // Do not fold BF16 conversions for instructions related to tuples, entry and
   // exit of a computation, fusion, convert, side-effecting instructions,
   // in-place operations and control flow.
@@ -186,7 +188,8 @@ Status BFloat16ConversionFoldingVisitor::DefaultAction(HloInstruction* hlo) {
   return TryFoldBF16Conversions(hlo);
 }
 
-Status BFloat16ConversionFoldingVisitor::HandleAllReduce(HloInstruction* crs) {
+absl::Status BFloat16ConversionFoldingVisitor::HandleAllReduce(
+    HloInstruction* crs) {
   if (crs->HasSideEffectNoRecurse()) {
     // Do not perform optimization on side-effected AllReduce.
     return OkStatus();
diff --git a/third_party/xla/xla/service/bfloat16_propagation.cc b/third_party/xla/xla/service/bfloat16_propagation.cc
index 38eb493081f302..8bcf979a7c24a9 100644
--- a/third_party/xla/xla/service/bfloat16_propagation.cc
+++ b/third_party/xla/xla/service/bfloat16_propagation.cc
@@ -708,7 +708,7 @@ void BFloat16Propagation::ResolveInconsistencyOfAliasingBuffers(
   }
 }
 
-Status BFloat16Propagation::ResolveInconsistentFusions(
+absl::Status BFloat16Propagation::ResolveInconsistentFusions(
     HloModule* module,
     const absl::flat_hash_set<absl::string_view>& execution_threads) {
   // We could have changed a fusion computation's root shape to have a different
@@ -771,7 +771,7 @@ Status BFloat16Propagation::ResolveInconsistentFusions(
   return OkStatus();
 }
 
-Status BFloat16Propagation::ResolveConvertedConstants(
+absl::Status BFloat16Propagation::ResolveConvertedConstants(
     HloModule* module,
     const absl::flat_hash_set<absl::string_view>& execution_threads) {
   // We may have converted some constants from F32 to BF16, so adjust the
@@ -803,7 +803,7 @@ Status BFloat16Propagation::ResolveConvertedConstants(
   return OkStatus();
 }
 
-Status BFloat16Propagation::SkipNoopConversions(
+absl::Status BFloat16Propagation::SkipNoopConversions(
     HloModule* module,
     const absl::flat_hash_set<absl::string_view>& execution_threads) {
   for (auto computation : module->computations(execution_threads)) {
diff --git a/third_party/xla/xla/service/bfloat16_propagation.h b/third_party/xla/xla/service/bfloat16_propagation.h
index 21625e7337573d..b7d1fdabfac706 100644
--- a/third_party/xla/xla/service/bfloat16_propagation.h
+++ b/third_party/xla/xla/service/bfloat16_propagation.h
@@ -159,19 +159,19 @@ class BFloat16Propagation : public HloModulePass {
 
   // Resolves inconsistencies introduced by this pass for fusions with
   // tuple-type output.
-  Status ResolveInconsistentFusions(
+  absl::Status ResolveInconsistentFusions(
       HloModule* module,
       const absl::flat_hash_set<absl::string_view>& execution_threads);
 
   // Converts the literals in kConstant HLOs which have their types changed to
   // BF16 by this pass.
-  Status ResolveConvertedConstants(
+  absl::Status ResolveConvertedConstants(
       HloModule* module,
       const absl::flat_hash_set<absl::string_view>& execution_threads);
 
   // Skips no-op conversions (same source and target shapes) that can be
   // produced this pass, i.e., replaces them in their uses with their operands.
-  Status SkipNoopConversions(
+  absl::Status SkipNoopConversions(
       HloModule* module,
       const absl::flat_hash_set<absl::string_view>& execution_threads);
 
diff --git a/third_party/xla/xla/service/buffer_assignment.cc b/third_party/xla/xla/service/buffer_assignment.cc
index 1f09360a6b07ce..5b326950691898 100644
--- a/third_party/xla/xla/service/buffer_assignment.cc
+++ b/third_party/xla/xla/service/buffer_assignment.cc
@@ -117,7 +117,7 @@ BuildIdToLogicalBufferMap(
 
 }  // namespace
 
-Status GatherComputationsByAllocationType(
+absl::Status GatherComputationsByAllocationType(
     const HloModule* module,
     std::vector<const HloComputation*>* thread_local_computations,
     std::vector<const HloComputation*>* global_computations) {
@@ -522,7 +522,7 @@ bool BufferAssignment::HaveDisjointSlices(const HloInstruction* hlo_a,
   // assigned slice, returns the empty set.
   auto collect_slices = [&](const HloInstruction* instr) -> SliceSet {
     SliceSet slices;
-    Status status = ShapeUtil::ForEachSubshapeWithStatus(
+    absl::Status status = ShapeUtil::ForEachSubshapeWithStatus(
         instr->shape(),
         [&](const Shape& /*subshape*/,
             const ShapeIndex& index) -> absl::Status {
@@ -723,7 +723,7 @@ void BufferAssignment::CombineTempAllocations(
   }
 }
 
-Status BufferAssignment::ComputeSummaryStats() {
+absl::Status BufferAssignment::ComputeSummaryStats() {
   for (auto& allocation : Allocations()) {
     if (allocation.is_entry_computation_parameter()) {
       stats_.parameter_allocation_count++;
@@ -1277,7 +1277,7 @@ bool BufferAssigner::MaybeAssignBuffer(BufferAllocation* allocation,
   return true;
 }  // namespace xla
 
-Status BufferAssigner::AssignSingleHloBuffer(
+absl::Status BufferAssigner::AssignSingleHloBuffer(
     const HloBuffer* hlo_buffer, bool is_thread_local,
     absl::flat_hash_map<const HloComputation*,
                         absl::flat_hash_set<const HloValue*>>*
@@ -1415,7 +1415,7 @@ Status BufferAssigner::AssignSingleHloBuffer(
   return OkStatus();
 }
 
-Status BufferAssigner::AssignBuffersForComputations(
+absl::Status BufferAssigner::AssignBuffersForComputations(
     const std::vector<const HloComputation*>& computations,
     bool is_thread_local,
     absl::flat_hash_map<const HloComputation*,
@@ -1565,7 +1565,7 @@ BufferAssigner::SplitBuffersByPrivateStackComputation(
   return computation_map;
 }
 
-Status BufferAssigner::AssignPresetBuffers(
+absl::Status BufferAssigner::AssignPresetBuffers(
     absl::flat_hash_set<const HloBuffer*>* assigned_buffers,
     BufferAssignment* assignment) {
   if (!preset_assignments_) {
@@ -1616,7 +1616,7 @@ Status BufferAssigner::AssignPresetBuffers(
   return OkStatus();
 }
 
-Status BufferAssigner::AssignBuffersWithSequentialOrdering(
+absl::Status BufferAssigner::AssignBuffersWithSequentialOrdering(
     const flat_hash_map<const HloComputation*, flat_hash_set<const HloValue*>>&
         buffers_to_assign_sequentially,
     bool run_whole_module_heap_simulation, BufferAssignment* assignment,
diff --git a/third_party/xla/xla/service/buffer_assignment.h b/third_party/xla/xla/service/buffer_assignment.h
index e27ed3cb8f0bec..bf60d9676abf2e 100644
--- a/third_party/xla/xla/service/buffer_assignment.h
+++ b/third_party/xla/xla/service/buffer_assignment.h
@@ -49,7 +49,7 @@ namespace xla {
 // elements in thread_local_computations and global_computations are in post
 // order (if computation A has an instruction which calls computation B, then A
 // will appear after B in the vector).
-Status GatherComputationsByAllocationType(
+absl::Status GatherComputationsByAllocationType(
     const HloModule* module,
     std::vector<const HloComputation*>* thread_local_computations,
     std::vector<const HloComputation*>* global_computations);
@@ -572,7 +572,7 @@ class BufferAssignment {
       const absl::flat_hash_set<BufferValue::Color>& private_stack_colors);
 
   // Computes stats for the assignment, to be retrieved by GetStats.
-  Status ComputeSummaryStats();
+  absl::Status ComputeSummaryStats();
 
   // The vector of buffer allocations. Indexed by BufferAllocation::Index.
   std::vector<BufferAllocation> allocations_;
@@ -611,7 +611,8 @@ class BufferAssignment {
 // A class which constructs a buffer assignment.
 class BufferAssigner {
  public:
-  using Colorer = std::function<Status(HloAliasAnalysis*, const HloOrdering&)>;
+  using Colorer =
+      std::function<absl::Status(HloAliasAnalysis*, const HloOrdering&)>;
   using MustNotLiveOut =
       std::function<bool(const HloInstruction*, const ShapeIndex&)>;
   using PrivateStacks = absl::flat_hash_map<BufferValue::Color,
@@ -683,7 +684,7 @@ class BufferAssigner {
   // is modified to reflect the new buffer assignments. If is_thread_local is
   // true, then all assigned buffers have the is_thread_local flag set to
   // true.
-  Status AssignBuffersForComputations(
+  absl::Status AssignBuffersForComputations(
       const std::vector<const HloComputation*>& computations,
       bool is_thread_local,
       absl::flat_hash_map<const HloComputation*,
@@ -697,12 +698,12 @@ class BufferAssigner {
 
   // Assigns pre-set assignments, if provided. These assignments will be added
   // to assigned_buffers and skip buffer allocation.
-  Status AssignPresetBuffers(
+  absl::Status AssignPresetBuffers(
       absl::flat_hash_set<const HloBuffer*>* assigned_buffers,
       BufferAssignment* assignment);
 
   // Assigns a single hlo buffer to an HLO allocation.
-  Status AssignSingleHloBuffer(
+  absl::Status AssignSingleHloBuffer(
       const HloBuffer* hlo_buffer, bool is_thread_local,
       absl::flat_hash_map<const HloComputation*,
                           absl::flat_hash_set<const HloValue*>>*
@@ -715,7 +716,7 @@ class BufferAssigner {
   // assignment->liveness().hlo_ordering().SequentialOrder. If
   // 'run_whole_module_heap_simulation' is true, the heap simulation will be run
   // assuming all global computations are sequentially ordered.
-  Status AssignBuffersWithSequentialOrdering(
+  absl::Status AssignBuffersWithSequentialOrdering(
       const absl::flat_hash_map<const HloComputation*,
                                 absl::flat_hash_set<const HloValue*>>&
           buffers_to_assign_sequentially,
diff --git a/third_party/xla/xla/service/buffer_assignment_test.cc b/third_party/xla/xla/service/buffer_assignment_test.cc
index 33e6297eefdaf7..02b7707dbf7fdf 100644
--- a/third_party/xla/xla/service/buffer_assignment_test.cc
+++ b/third_party/xla/xla/service/buffer_assignment_test.cc
@@ -62,7 +62,7 @@ class InstructionListVisitor : public DfsHloVisitorWithDefault {
  public:
   explicit InstructionListVisitor(const HloInstruction* root) : root_(root) {}
 
-  Status DefaultAction(HloInstruction* hlo) override {
+  absl::Status DefaultAction(HloInstruction* hlo) override {
     // For each instruction, just push it on the list after walking the
     // operands.
     instructions_.push_back(hlo);
diff --git a/third_party/xla/xla/service/call_graph.cc b/third_party/xla/xla/service/call_graph.cc
index fcf0d239a2dabb..805de50bff9838 100644
--- a/third_party/xla/xla/service/call_graph.cc
+++ b/third_party/xla/xla/service/call_graph.cc
@@ -363,7 +363,7 @@ std::unique_ptr<CallGraph> CallGraph::Build(
   return call_graph;
 }
 
-Status CallGraph::VisitNodesInternal(
+absl::Status CallGraph::VisitNodesInternal(
     VisitorFunction visitor_func, const CallGraphNode& node,
     absl::flat_hash_set<const CallGraphNode*>* visited) const {
   auto pair = visited->insert(&node);
@@ -380,8 +380,8 @@ Status CallGraph::VisitNodesInternal(
   return visitor_func(node);
 }
 
-Status CallGraph::VisitNodes(VisitorFunction visitor_func,
-                             bool visit_unreachable_nodes) const {
+absl::Status CallGraph::VisitNodes(VisitorFunction visitor_func,
+                                   bool visit_unreachable_nodes) const {
   absl::flat_hash_set<const CallGraphNode*> visited;
   if (visit_unreachable_nodes) {
     // Traverse from all roots in the call graph.
diff --git a/third_party/xla/xla/service/call_graph.h b/third_party/xla/xla/service/call_graph.h
index d0fa22157f363f..a64346e5d0dc3d 100644
--- a/third_party/xla/xla/service/call_graph.h
+++ b/third_party/xla/xla/service/call_graph.h
@@ -188,7 +188,7 @@ class CallGraphNode {
 // computation in the module.
 class CallGraph {
  public:
-  using VisitorFunction = absl::FunctionRef<Status(const CallGraphNode&)>;
+  using VisitorFunction = absl::FunctionRef<absl::Status(const CallGraphNode&)>;
 
   // Builds and returns a call graph for the given HLO module. If a non-empty
   // execution_threads is provided, only computations that are in
@@ -208,8 +208,8 @@ class CallGraph {
   // in post order (callees before callers). If visit_unreachable_nodes is true
   // then all nodes in the call graph are visited. Otherwise only those nodes
   // reachable from the entry computation are visited.
-  Status VisitNodes(VisitorFunction visitor_func,
-                    bool visit_unreachable_nodes = true) const;
+  absl::Status VisitNodes(VisitorFunction visitor_func,
+                          bool visit_unreachable_nodes = true) const;
 
   // Returns true if 'a' dominates 'b' in the call graph. Computation 'a'
   // dominates computation 'b' iff all callgraph paths in the caller-to-callee
@@ -360,7 +360,7 @@ class CallGraph {
   // post order (callee before caller) calling visitor_func on each node. Adds
   // nodes to 'visited' as each node is visited. Skips nodes already in
   // 'visited'.
-  Status VisitNodesInternal(
+  absl::Status VisitNodesInternal(
       VisitorFunction visitor_func, const CallGraphNode& node,
       absl::flat_hash_set<const CallGraphNode*>* visited) const;
 
diff --git a/third_party/xla/xla/service/call_graph_test.cc b/third_party/xla/xla/service/call_graph_test.cc
index 9337bec456fa5b..f97dc23510b093 100644
--- a/third_party/xla/xla/service/call_graph_test.cc
+++ b/third_party/xla/xla/service/call_graph_test.cc
@@ -706,7 +706,7 @@ TEST_F(CallGraphTest, VisitWithError) {
   module->AddEntryComputation(MakeScalarComputation());
   std::unique_ptr<CallGraph> call_graph = CallGraph::Build(module.get());
 
-  Status status = call_graph->VisitNodes(
+  absl::Status status = call_graph->VisitNodes(
       [](const CallGraphNode&) { return Internal("Visitation failed"); });
 
   ASSERT_FALSE(status.ok());
diff --git a/third_party/xla/xla/service/call_inliner.cc b/third_party/xla/xla/service/call_inliner.cc
index f92369d8605e91..63edc5c1f8f952 100644
--- a/third_party/xla/xla/service/call_inliner.cc
+++ b/third_party/xla/xla/service/call_inliner.cc
@@ -55,7 +55,7 @@ class SubcomputationInsertionVisitor : public DfsHloVisitorWithDefault {
 
   // Resolves the operands to the HLO instruction in the inlined (caller) graph,
   // and clones the HLO instruction into that graph with the new operands.
-  Status DefaultAction(HloInstruction* hlo) override {
+  absl::Status DefaultAction(HloInstruction* hlo) override {
     std::vector<HloInstruction*> new_operands;
     for (HloInstruction* operand : hlo->operands()) {
       TF_ASSIGN_OR_RETURN(HloInstruction * new_operand, Resolve(operand));
@@ -81,7 +81,7 @@ class SubcomputationInsertionVisitor : public DfsHloVisitorWithDefault {
   // Does not create new nodes for the parameter; rather, notes the mapping from
   // the subcomputation parameter node to the call operands in the caller
   // computation.
-  Status HandleParameter(HloInstruction* parameter) override {
+  absl::Status HandleParameter(HloInstruction* parameter) override {
     TF_RETURN_IF_ERROR(NoteMapping(
         parameter, call_->mutable_operand(parameter->parameter_number())));
     return OkStatus();
@@ -89,7 +89,7 @@ class SubcomputationInsertionVisitor : public DfsHloVisitorWithDefault {
 
   // Wires the consumers of the call to instead point at the newly created root,
   // replacing the call operation in the caller computation.
-  Status FinishVisit(HloInstruction* root) override {
+  absl::Status FinishVisit(HloInstruction* root) override {
     TF_ASSIGN_OR_RETURN(HloInstruction * new_root, Resolve(root));
     VLOG(1) << "Replacing all uses of " << call_->ToString()
             << " with new root " << new_root->ToString();
@@ -119,8 +119,8 @@ class SubcomputationInsertionVisitor : public DfsHloVisitorWithDefault {
   //
   // Returns an error status if the subcomputation_hlo is mapped more than
   // once.
-  Status NoteMapping(HloInstruction* subcomputation_hlo,
-                     HloInstruction* new_hlo) {
+  absl::Status NoteMapping(HloInstruction* subcomputation_hlo,
+                           HloInstruction* new_hlo) {
     auto result = subcomputation_hlo_to_new_hlo_.insert(
         std::make_pair(subcomputation_hlo, new_hlo));
     TF_RET_CHECK(result.second)
@@ -161,7 +161,7 @@ absl::StatusOr<bool> CallInliner::Run(
   // we'll always inline kCalls into their callers in the appropriate order.
   bool did_mutate = false;
   TF_RETURN_IF_ERROR(call_graph->VisitNodes([&](const CallGraphNode& node)
-                                                -> Status {
+                                                -> absl::Status {
     if (!HloInstruction::IsThreadIncluded(
             node.computation()->execution_thread(), execution_threads)) {
       return OkStatus();
diff --git a/third_party/xla/xla/service/collective_combiner_utils.h b/third_party/xla/xla/service/collective_combiner_utils.h
index b54296d764a87a..e7b2ecc573670a 100644
--- a/third_party/xla/xla/service/collective_combiner_utils.h
+++ b/third_party/xla/xla/service/collective_combiner_utils.h
@@ -46,7 +46,8 @@ template <typename K>
 absl::StatusOr<bool> CombineInstructionsByKey(
     HloComputation* computation,
     absl::FunctionRef<std::optional<K>(const HloInstruction*)> key_fn,
-    absl::FunctionRef<Status(absl::Span<HloInstruction* const>)> combine_fn,
+    absl::FunctionRef<absl::Status(absl::Span<HloInstruction* const>)>
+        combine_fn,
     int64_t combine_threshold_bytes, int64_t combine_threshold_count) {
   // Cache keys for each instruction and build sets of instructions with the
   // same key that might be combined together.
diff --git a/third_party/xla/xla/service/collective_ops_utils.h b/third_party/xla/xla/service/collective_ops_utils.h
index d8e046bee4df90..315d0f3330b836 100644
--- a/third_party/xla/xla/service/collective_ops_utils.h
+++ b/third_party/xla/xla/service/collective_ops_utils.h
@@ -308,7 +308,7 @@ class Rendezvous {
 
   // Submit a participant to the rendezvous. We get the rendezvous from
   // `rendezvous_getter`, which we can then use to drop the existing reference.
-  static StatusOr<O> SubmitParticipant(
+  static absl::StatusOr<O> SubmitParticipant(
       absl::FunctionRef<std::shared_ptr<Rendezvous<I, O>>()> rendezvous_getter,
       I participant) {
     std::shared_ptr<Rendezvous<I, O>> rendezvous = rendezvous_getter();
@@ -336,7 +336,7 @@ class Rendezvous {
 
  protected:
   // Returns domain-specific output O and whether this replica is primary.
-  virtual StatusOr<O> RunCollectiveOp(const I& participant) = 0;
+  virtual absl::StatusOr<O> RunCollectiveOp(const I& participant) = 0;
 
   // Adding participants_ requires holding mu_.
   // Not annotated with ABSL_GUARDED_BY(mu_) because we do not require the lock
@@ -353,7 +353,7 @@ class Rendezvous {
   //  - a BlockingCounter initialized to the number of participants, so that
   //    the caller can coordinate with the participants one last time if it
   //    chooses.  This is useful for coordinating destruction of the Rendezvous.
-  StatusOr<std::pair<O, std::shared_ptr<tsl::BlockingCounter>>>
+  absl::StatusOr<std::pair<O, std::shared_ptr<tsl::BlockingCounter>>>
   SubmitParticipant(const I& participant) {
     {
       absl::MutexLock lock(&mu_);
diff --git a/third_party/xla/xla/service/collective_permute_decomposer.cc b/third_party/xla/xla/service/collective_permute_decomposer.cc
index 1660163b948b11..39f4b241c7f361 100644
--- a/third_party/xla/xla/service/collective_permute_decomposer.cc
+++ b/third_party/xla/xla/service/collective_permute_decomposer.cc
@@ -110,7 +110,7 @@ bool MayPipeline(const HloCollectivePermuteInstruction& collective_permute) {
 // collective-permute will be pipelined and the value of the attribute
 // represents the runtime stream to execute the instruction. Without the
 // frontend attribute, the collective-permute will not be pipelined.
-Status DecomposeCollectivePermute(
+absl::Status DecomposeCollectivePermute(
     HloCollectivePermuteInstruction* collective_permute,
     HloComputation* computation, const std::string& pipeline_decision) {
   // We currently only decompose collective-permute with a channel_id.
diff --git a/third_party/xla/xla/service/collective_pipeliner.cc b/third_party/xla/xla/service/collective_pipeliner.cc
index 68e6348b44d4fa..ffd988752d7edb 100644
--- a/third_party/xla/xla/service/collective_pipeliner.cc
+++ b/third_party/xla/xla/service/collective_pipeliner.cc
@@ -76,9 +76,9 @@ using LoopVariantParameterInfo =
 
 // Update all control dependencies for a cloned instruction to connect other
 // cloned instructions rather than originals.
-Status UpdateControlDependencies(HloInstruction* original,
-                                 HloInstruction* new_instr,
-                                 const InstructionMap& cloned_map) {
+absl::Status UpdateControlDependencies(HloInstruction* original,
+                                       HloInstruction* new_instr,
+                                       const InstructionMap& cloned_map) {
   for (auto* pred : original->control_predecessors()) {
     auto it = cloned_map.find(pred);
     if (it == cloned_map.end()) {
@@ -1176,14 +1176,12 @@ HloInstruction* CreateZero(HloComputation* comp, const Shape& shape,
 // }
 // xg_last = all-reduce(x)
 // yg_last = all-reduce(y)
-Status TransformLoopForward(const WhileLoopAnalysis& loop_analysis,
-                            bool insert_non_alias_custom_call,
-                            int64_t level_to_operate_on, bool pipeline_use_tree,
-                            bool process_different_sized_ops,
-                            HloPredicate should_process,
-                            HloPredicate acceptable_formatting,
-                            HloPredicate reuse_output_buffer,
-                            int64_t& next_channel_id) {
+absl::Status TransformLoopForward(
+    const WhileLoopAnalysis& loop_analysis, bool insert_non_alias_custom_call,
+    int64_t level_to_operate_on, bool pipeline_use_tree,
+    bool process_different_sized_ops, HloPredicate should_process,
+    HloPredicate acceptable_formatting, HloPredicate reuse_output_buffer,
+    int64_t& next_channel_id) {
   // Defining some maps/sets to keep track of instructions duplicated.
   InstructionMap while_body_to_peeled;
   absl::flat_hash_set<HloInstruction*> to_skip_set;
@@ -1584,13 +1582,13 @@ Status TransformLoopForward(const WhileLoopAnalysis& loop_analysis,
 // }
 // xg_all = all-reduce(x_all)
 // yg_all = all-reduce(y_all)
-Status TransformLoopForwardSink(const WhileLoopAnalysis& loop_analysis,
-                                bool insert_non_alias_custom_call,
-                                int64_t level_to_operate_on,
-                                bool pipeline_use_tree,
-                                bool process_different_sized_ops,
-                                HloPredicate should_process,
-                                int64_t& next_channel_id) {
+absl::Status TransformLoopForwardSink(const WhileLoopAnalysis& loop_analysis,
+                                      bool insert_non_alias_custom_call,
+                                      int64_t level_to_operate_on,
+                                      bool pipeline_use_tree,
+                                      bool process_different_sized_ops,
+                                      HloPredicate should_process,
+                                      int64_t& next_channel_id) {
   // Defining some maps/sets to keep track of instructions duplicated.
   absl::flat_hash_map<HloInstruction*, int64_t> is_output_instruction;
   absl::flat_hash_map<const HloInstruction*, bool> invariant_cache;
@@ -2082,7 +2080,7 @@ Status TransformLoopForwardSink(const WhileLoopAnalysis& loop_analysis,
 //   x_ag = p0_ag_next
 // }
 // x_last = computation(p0_ag_next)
-static Status TransformLoopBackward(
+static absl::Status TransformLoopBackward(
     const WhileLoopAnalysis& loop_analysis, bool insert_non_alias_custom_call,
     int64_t level_to_operate_on, bool process_different_sized_ops,
     HloPredicate should_process, HloPredicate acceptable_formatting,
diff --git a/third_party/xla/xla/service/collective_pipeliner.h b/third_party/xla/xla/service/collective_pipeliner.h
index 0e87137f5ef59a..433686fdb5a1b9 100644
--- a/third_party/xla/xla/service/collective_pipeliner.h
+++ b/third_party/xla/xla/service/collective_pipeliner.h
@@ -65,7 +65,7 @@ class CollectivePipeliner : public HloModulePass {
   // Postprocessing cloned collective instructions, such as for modifying loop
   // iteration related frontend attributes to reflect loop pipelining.
   using HloPostprocessor =
-      std::optional<std::function<Status(HloInstruction* instr)>>;
+      std::optional<std::function<absl::Status(HloInstruction* instr)>>;
 
   struct Config {
     int64_t level_to_operate_on = 0;
diff --git a/third_party/xla/xla/service/compilation_environments.cc b/third_party/xla/xla/service/compilation_environments.cc
index 0c2569b92dfcfb..a86209a01a7a69 100644
--- a/third_party/xla/xla/service/compilation_environments.cc
+++ b/third_party/xla/xla/service/compilation_environments.cc
@@ -185,7 +185,7 @@ void CompilationEnvironments::RegisterProcessNewEnvFn(
                   << descriptor->full_name() << "' has already been registered";
 }
 
-Status CompilationEnvironments::AddEnv(
+absl::Status CompilationEnvironments::AddEnv(
     std::unique_ptr<tsl::protobuf::Message> env) {
   if (!env) {
     return tsl::errors::InvalidArgument(
@@ -239,7 +239,7 @@ void CompilationEnvironments::EnvAdded(std::string_view env_type) {
   GlobalCompEnvStats::GetSingleton().EnvAdded(env_type);
 }
 
-Status CompilationEnvironments::AddEnvImpl(
+absl::Status CompilationEnvironments::AddEnvImpl(
     const tsl::protobuf::Descriptor& descriptor,
     std::unique_ptr<tsl::protobuf::Message> env) {
   // Check if we already have an environment of env's type
diff --git a/third_party/xla/xla/service/compilation_environments.h b/third_party/xla/xla/service/compilation_environments.h
index 3ffea24bb53b11..9dbb73df94292e 100644
--- a/third_party/xla/xla/service/compilation_environments.h
+++ b/third_party/xla/xla/service/compilation_environments.h
@@ -84,7 +84,7 @@ class CompilationEnvironments {
   // All added environments are processed via registered ProcessNewEnvFns. If
   // such a function was not regitered for env's proto descriptor or env's
   // proto type is unknown, an error will be returned.
-  Status AddEnv(std::unique_ptr<tsl::protobuf::Message> env);
+  absl::Status AddEnv(std::unique_ptr<tsl::protobuf::Message> env);
 
   // Returns the CompilationEnvironment corresponding to T. If such an
   // environment has not been added, ProcessNewEnvFn(nullptr) will be added and
@@ -122,8 +122,8 @@ class CompilationEnvironments {
   // are added to CompilationEnvironments.
   static void EnvAdded(std::string_view env_type);
 
-  Status AddEnvImpl(const tsl::protobuf::Descriptor& descriptor,
-                    std::unique_ptr<tsl::protobuf::Message> env);
+  absl::Status AddEnvImpl(const tsl::protobuf::Descriptor& descriptor,
+                          std::unique_ptr<tsl::protobuf::Message> env);
 
   absl::flat_hash_map<const tsl::protobuf::Descriptor*,
                       std::unique_ptr<tsl::protobuf::Message>>
diff --git a/third_party/xla/xla/service/compile_only_service.h b/third_party/xla/xla/service/compile_only_service.h
index 09ca0534454b46..72a63ec2e23b58 100644
--- a/third_party/xla/xla/service/compile_only_service.h
+++ b/third_party/xla/xla/service/compile_only_service.h
@@ -53,27 +53,28 @@ class CompileOnlyService : public Service {
                      const AotCompilationOptions& options,
                      std::unique_ptr<AotCompilationMetadata>* metadata);
 
-  Status GetDeviceHandles(const GetDeviceHandlesRequest* arg,
-                          GetDeviceHandlesResponse* result) override {
+  absl::Status GetDeviceHandles(const GetDeviceHandlesRequest* arg,
+                                GetDeviceHandlesResponse* result) override {
     return Unimplemented("CompileOnlyService does not support devices.");
   }
-  Status TransferToServer(const TransferToServerRequest* arg,
-                          TransferToServerResponse* result) override {
+  absl::Status TransferToServer(const TransferToServerRequest* arg,
+                                TransferToServerResponse* result) override {
     return Unimplemented(
         "CompileOnlyService does not support device data transfers.");
   }
-  Status TransferToInfeed(const TransferToInfeedRequest* arg,
-                          TransferToInfeedResponse* result) override {
+  absl::Status TransferToInfeed(const TransferToInfeedRequest* arg,
+                                TransferToInfeedResponse* result) override {
     return Unimplemented(
         "CompileOnlyService does not support device data transfers.");
   }
-  Status TransferFromOutfeed(const TransferFromOutfeedRequest* arg,
-                             TransferFromOutfeedResponse* result) override {
+  absl::Status TransferFromOutfeed(
+      const TransferFromOutfeedRequest* arg,
+      TransferFromOutfeedResponse* result) override {
     return Unimplemented(
         "CompileOnlyService does not support device data transfers.");
   }
-  Status ResetDevice(const ResetDeviceRequest* arg,
-                     ResetDeviceResponse* result) override {
+  absl::Status ResetDevice(const ResetDeviceRequest* arg,
+                           ResetDeviceResponse* result) override {
     return Unimplemented("CompileOnlyService does not support devices.");
   }
 
diff --git a/third_party/xla/xla/service/computation_placer.cc b/third_party/xla/xla/service/computation_placer.cc
index d0f2e93fdb7ebd..40243118504763 100644
--- a/third_party/xla/xla/service/computation_placer.cc
+++ b/third_party/xla/xla/service/computation_placer.cc
@@ -84,7 +84,7 @@ DeviceAssignment::GetDeviceToLogicalIdMap() const {
   return device_to_logical_id;
 }
 
-Status DeviceAssignment::Serialize(DeviceAssignmentProto* proto) const {
+absl::Status DeviceAssignment::Serialize(DeviceAssignmentProto* proto) const {
   proto->set_replica_count(replica_count());
   proto->set_computation_count(computation_count());
   for (int computation = 0; computation < computation_count(); ++computation) {
diff --git a/third_party/xla/xla/service/computation_placer.h b/third_party/xla/xla/service/computation_placer.h
index 178f581650a5a7..3ee82afd157d3e 100644
--- a/third_party/xla/xla/service/computation_placer.h
+++ b/third_party/xla/xla/service/computation_placer.h
@@ -64,11 +64,11 @@ class DeviceAssignment : public Array2D<int64_t> {
       const;
 
   // Protocol buffer serialization and deserialization.
-  Status Serialize(DeviceAssignmentProto* proto) const;
+  absl::Status Serialize(DeviceAssignmentProto* proto) const;
 
   // Return a std::unique_ptr<DeviceAssignment> instead of a DeviceAssignment
   // directly because one of the supported TF platforms (mac) does not compile
-  // due to a StatusOr of an incomplete type (DeviceAssignment).
+  // due to a absl::StatusOr of an incomplete type (DeviceAssignment).
   static absl::StatusOr<std::unique_ptr<DeviceAssignment>> Deserialize(
       const DeviceAssignmentProto& proto);
 
diff --git a/third_party/xla/xla/service/conditional_canonicalizer.cc b/third_party/xla/xla/service/conditional_canonicalizer.cc
index 22eb2cb6500fea..9b87eceb763dda 100644
--- a/third_party/xla/xla/service/conditional_canonicalizer.cc
+++ b/third_party/xla/xla/service/conditional_canonicalizer.cc
@@ -21,7 +21,7 @@ limitations under the License.
 
 namespace xla {
 namespace {
-Status CanonicalizeNonTupleConditional(HloInstruction* conditional) {
+absl::Status CanonicalizeNonTupleConditional(HloInstruction* conditional) {
   TF_RET_CHECK(conditional->opcode() == HloOpcode::kConditional);
   for (auto* branch : conditional->called_computations()) {
     HloInstruction* root = branch->root_instruction();
diff --git a/third_party/xla/xla/service/conditional_code_motion.cc b/third_party/xla/xla/service/conditional_code_motion.cc
index 1fed4183420b64..a1429c38b19bbe 100644
--- a/third_party/xla/xla/service/conditional_code_motion.cc
+++ b/third_party/xla/xla/service/conditional_code_motion.cc
@@ -498,8 +498,8 @@ absl::flat_hash_set<int64_t> FindSpecialConverts(HloInstruction* old_root,
 // z = tuple(y.0, y.1, ...y.n)
 // Doing so ensures that we can accommodate the possible shape-change of the
 // conditional when the instructions are hoisted.
-Status RestructureConditionalInstruction(HloComputation* computation,
-                                         HloInstruction* conditional) {
+absl::Status RestructureConditionalInstruction(HloComputation* computation,
+                                               HloInstruction* conditional) {
   HloInstruction* old_root = computation->root_instruction();
   std::vector<HloInstruction*> new_operands;
   int cur_index = 0;
@@ -937,7 +937,7 @@ absl::StatusOr<bool> ConditionalCodeMotion::MoveUserInstructionsIn(
 class MoveOperandIntoBranch {
  public:
   MoveOperandIntoBranch() = default;
-  Status operator()(HloInstruction* inst, HloInstruction*& user) {
+  absl::Status operator()(HloInstruction* inst, HloInstruction*& user) {
     VLOG(1) << "operand to move into branch: " << inst->ToString();
     VLOG(2) << "MoveIntoBranches user =" << user->ToString() << "\n";
     CHECK(inst->user_count() == 1 || inst->opcode() == HloOpcode::kBroadcast);
@@ -1061,7 +1061,7 @@ class MoveOperandIntoBranch {
   // Replace input with its operands inside user. Use matching_tuple_indices to
   // remember which operands of input is matched to which entries in the new
   // user.
-  Status ReplaceInputInUser(
+  absl::Status ReplaceInputInUser(
       HloInstruction* input, HloInstruction*& user,
       absl::InlinedVector<HloInstruction*, 4>& new_operands,
       std::vector<std::vector<int64_t>>& matching_tuple_indices) {
@@ -1137,7 +1137,7 @@ class MoveOperandIntoBranch {
     }
     return OkStatus();
   }
-  Status MoveInputIntoBranch(
+  absl::Status MoveInputIntoBranch(
       HloInstruction* input, HloInstruction*& user,
       absl::InlinedVector<HloInstruction*, 4>& new_operands,
       std::vector<std::vector<int64_t>>& matching_tuple_indices) {
diff --git a/third_party/xla/xla/service/conditional_to_select.cc b/third_party/xla/xla/service/conditional_to_select.cc
index 3b2fd710038300..b32b96efe76c49 100644
--- a/third_party/xla/xla/service/conditional_to_select.cc
+++ b/third_party/xla/xla/service/conditional_to_select.cc
@@ -73,7 +73,7 @@ absl::StatusOr<bool> ConditionalToSelect::Run(
   bool did_mutate = false;
   VLOG(1) << "Running conditional-to-select pass";
   TF_RETURN_IF_ERROR(
-      call_graph->VisitNodes([&](const CallGraphNode& node) -> Status {
+      call_graph->VisitNodes([&](const CallGraphNode& node) -> absl::Status {
         std::vector<HloInstruction*> ToInline;
         if (node.context() != CallContext::kEmbedded) {
           return OkStatus();
diff --git a/third_party/xla/xla/service/convert_async_collectives_to_sync.cc b/third_party/xla/xla/service/convert_async_collectives_to_sync.cc
index 6d633211c06dd3..6da602e1239918 100644
--- a/third_party/xla/xla/service/convert_async_collectives_to_sync.cc
+++ b/third_party/xla/xla/service/convert_async_collectives_to_sync.cc
@@ -106,7 +106,7 @@ absl::StatusOr<HloInstruction*> CreateSyncVariant(HloInstruction* async_start,
   return sync_instruction;
 }
 
-/*static*/ Status
+/*static*/ absl::Status
 ConvertAsyncCollectivesToSync::ReplaceAsyncInstructionsWithSync(
     HloComputation* computation,
     absl::Span<const std::pair<HloInstruction*, HloInstruction*>> async_pairs) {
diff --git a/third_party/xla/xla/service/convert_async_collectives_to_sync.h b/third_party/xla/xla/service/convert_async_collectives_to_sync.h
index 2b37c6ee7fa469..c359b2bda38db4 100644
--- a/third_party/xla/xla/service/convert_async_collectives_to_sync.h
+++ b/third_party/xla/xla/service/convert_async_collectives_to_sync.h
@@ -40,7 +40,7 @@ class ConvertAsyncCollectivesToSync : public HloModulePass {
       HloModule* module,
       const absl::flat_hash_set<absl::string_view>& execution_threads) override;
 
-  virtual Status ConvertAsyncInstructionsToSync(
+  virtual absl::Status ConvertAsyncInstructionsToSync(
       HloComputation* computation,
       absl::Span<const std::pair<HloInstruction*, HloInstruction*>> async_pairs)
       const {
@@ -49,7 +49,7 @@ class ConvertAsyncCollectivesToSync : public HloModulePass {
 
   // Helper utility to replace a list of pairs of async-start/done ops in a
   // computation with their synchronous variants and update the schedule.
-  static Status ReplaceAsyncInstructionsWithSync(
+  static absl::Status ReplaceAsyncInstructionsWithSync(
       HloComputation* computation,
       absl::Span<const std::pair<HloInstruction*, HloInstruction*>>
           async_pairs);
diff --git a/third_party/xla/xla/service/convert_async_collectives_to_sync_test.cc b/third_party/xla/xla/service/convert_async_collectives_to_sync_test.cc
index 355ba78b7e2fc9..2ba38a727559e2 100644
--- a/third_party/xla/xla/service/convert_async_collectives_to_sync_test.cc
+++ b/third_party/xla/xla/service/convert_async_collectives_to_sync_test.cc
@@ -35,8 +35,8 @@ namespace m = xla::testing::opcode_matchers;
 // to the HLO module string.
 class ConvertAsyncCollectivesToSyncTest : public HloTestBase {
  public:
-  Status RunPass(HloModule *module, bool expect_change,
-                 HloPredicate is_nop = {}) {
+  absl::Status RunPass(HloModule *module, bool expect_change,
+                       HloPredicate is_nop = {}) {
     TF_ASSIGN_OR_RETURN(bool changed,
                         ConvertAsyncCollectivesToSync{is_nop}.Run(module));
     EXPECT_EQ(changed, expect_change);
diff --git a/third_party/xla/xla/service/convolution_group_converter.cc b/third_party/xla/xla/service/convolution_group_converter.cc
index 4dfba405d7a7f2..c4e192101c8800 100644
--- a/third_party/xla/xla/service/convolution_group_converter.cc
+++ b/third_party/xla/xla/service/convolution_group_converter.cc
@@ -45,13 +45,13 @@ namespace {
 class ConvolutionVisitor : public DfsHloVisitorWithDefault {
  public:
   // Default visitor action is to do nothing and return OK.
-  Status DefaultAction(HloInstruction* /*hlo_instruction*/) override {
+  absl::Status DefaultAction(HloInstruction* /*hlo_instruction*/) override {
     return OkStatus();
   }
 
-  Status HandleConvolution(HloInstruction* convolution) override;
+  absl::Status HandleConvolution(HloInstruction* convolution) override;
 
-  Status HandleBatchGroupCount(HloInstruction* convolution);
+  absl::Status HandleBatchGroupCount(HloInstruction* convolution);
 
   // Runs the visitor on a computation.
   static bool Run(HloComputation* computation,
@@ -201,7 +201,8 @@ HloInstruction* GetExpandedFilterMask(
 
 // This function handles batch_group_counts which are relevant only for
 // depthwise backprop filter convolutions.
-Status ConvolutionVisitor::HandleBatchGroupCount(HloInstruction* convolution) {
+absl::Status ConvolutionVisitor::HandleBatchGroupCount(
+    HloInstruction* convolution) {
   auto dim_numbers = convolution->convolution_dimension_numbers();
   auto activation = convolution->mutable_operand(0);
   auto filter = convolution->mutable_operand(1);
@@ -418,7 +419,8 @@ Status ConvolutionVisitor::HandleBatchGroupCount(HloInstruction* convolution) {
   return OkStatus();
 }
 
-Status ConvolutionVisitor::HandleConvolution(HloInstruction* convolution) {
+absl::Status ConvolutionVisitor::HandleConvolution(
+    HloInstruction* convolution) {
   if (convert_batch_groups_only_) {
     return HandleBatchGroupCount(convolution);
   }
diff --git a/third_party/xla/xla/service/copy_insertion.cc b/third_party/xla/xla/service/copy_insertion.cc
index 5e99bcb5fa4a2a..5c8e9bad01d5b5 100644
--- a/third_party/xla/xla/service/copy_insertion.cc
+++ b/third_party/xla/xla/service/copy_insertion.cc
@@ -284,8 +284,8 @@ bool IndicesToCopyForConditional(const HloDataflowAnalysis& dataflow,
 // If the loop state is a tuple then the above kCopy instructions are a deep
 // copy constructed of kCopy, kGetTupleElement, and kTuple instruction as
 // constructed by HloInstruction::DeepCopyInstruction.
-Status AddCopiesForWhile(const HloAliasAnalysis& alias_analysis,
-                         HloInstruction* xla_while) {
+absl::Status AddCopiesForWhile(const HloAliasAnalysis& alias_analysis,
+                               HloInstruction* xla_while) {
   VLOG(2) << "Adding copies for kWhile instruction " << xla_while->name();
   TF_RET_CHECK(xla_while->opcode() == HloOpcode::kWhile);
 
@@ -342,9 +342,9 @@ Status AddCopiesForWhile(const HloAliasAnalysis& alias_analysis,
 
 // Add copies for the operands of in-place operations. RemoveUnnecessaryCopies
 // will remove the unnecessary copies.
-Status AddCopiesForInPlaceOperation(const HloAliasAnalysis& alias_analysis,
-                                    HloInstruction* in_place_op,
-                                    int64_t operand_number) {
+absl::Status AddCopiesForInPlaceOperation(
+    const HloAliasAnalysis& alias_analysis, HloInstruction* in_place_op,
+    int64_t operand_number) {
   VLOG(2) << "Adding copies for in-place operation " << in_place_op->name();
   HloInstruction* operand = in_place_op->mutable_operand(operand_number);
   TF_ASSIGN_OR_RETURN(HloInstruction * deep_copy,
@@ -358,7 +358,7 @@ Status AddCopiesForInPlaceOperation(const HloAliasAnalysis& alias_analysis,
 // each aliased parameter to resolve interference of aliased input and output
 // buffer. We later rely on RemoveUnnecessaryCopies to drop the unnecessary
 // ones.
-Status AddCopiesForAliasedInputOutputs(
+absl::Status AddCopiesForAliasedInputOutputs(
     HloModule* module,
     const absl::flat_hash_set<absl::string_view>& execution_threads) {
   HloComputation* entry = module->entry_computation();
@@ -429,7 +429,7 @@ Status AddCopiesForAliasedInputOutputs(
   // Add control dependencies between the input/output copies.
   TF_RETURN_IF_ERROR(module->input_output_alias_config().ForEachAliasWithStatus(
       [&](const ShapeIndex& output_index,
-          const HloInputOutputAliasConfig::Alias& alias) -> Status {
+          const HloInputOutputAliasConfig::Alias& alias) -> absl::Status {
         if (!copied_parameters[alias.parameter_number]) {
           return OkStatus();
         }
@@ -450,7 +450,7 @@ Status AddCopiesForAliasedInputOutputs(
 }
 
 // Removes any control dependencies to or from the given instruction.
-Status StripControlDependenciesFrom(HloInstruction* instruction) {
+absl::Status StripControlDependenciesFrom(HloInstruction* instruction) {
   while (!instruction->control_successors().empty()) {
     TF_RETURN_IF_ERROR(instruction->RemoveControlDependencyTo(
         instruction->control_successors().front()));
@@ -1454,7 +1454,7 @@ class CopyRemover {
   }
 
   // Verify invariants within the linked lists.
-  Status Verify() const {
+  absl::Status Verify() const {
     for (const ValueNode* head : value_lists_) {
       const ValueNode* p = head;
       do {
@@ -1967,7 +1967,7 @@ class CopyRemover {
 // We add copies for all phi indices of the true and false computation
 // roots, in order to resolve interference. We later rely on
 // RemoveUnnecessaryCopies to drop the unnecessary ones.
-Status CopyInsertion::AddCopiesForConditional(
+absl::Status CopyInsertion::AddCopiesForConditional(
     const HloAliasAnalysis& alias_analysis, HloInstruction* conditional) {
   VLOG(2) << "Adding copies for kConditional instruction "
           << conditional->name();
@@ -1997,7 +1997,7 @@ Status CopyInsertion::AddCopiesForConditional(
 // Add kCopy instructions to the given module to guarantee there is no
 // live-range interference. Generally interference can only occur around kWhile
 // instructions which have update-in-place semantics.
-Status CopyInsertion::AddCopiesToResolveInterference(
+absl::Status CopyInsertion::AddCopiesToResolveInterference(
     HloModule* module,
     const absl::flat_hash_set<absl::string_view>& execution_threads) {
   TF_ASSIGN_OR_RETURN(std::unique_ptr<HloAliasAnalysis> alias_analysis,
@@ -2071,14 +2071,14 @@ Status CopyInsertion::AddCopiesToResolveInterference(
   return OkStatus();
 }
 
-Status CopyInsertion::AddSpecialCaseCopies(
+absl::Status CopyInsertion::AddSpecialCaseCopies(
     HloModule* module,
     const absl::flat_hash_set<absl::string_view>& execution_threads) {
   std::unique_ptr<CallGraph> call_graph = CallGraph::Build(module);
   return AddSpecialCaseCopies(*call_graph, execution_threads, module);
 }
 
-Status CopyInsertion::AddSpecialCaseCopies(
+absl::Status CopyInsertion::AddSpecialCaseCopies(
     const CallGraph& call_graph,
     const absl::flat_hash_set<absl::string_view>& execution_threads,
     HloModule* module) {
@@ -2245,7 +2245,7 @@ static int64_t GetNumExistingCopies(
   return num_existing_copies;
 }
 
-Status CopyInsertion::RemoveUnnecessaryCopies(
+absl::Status CopyInsertion::RemoveUnnecessaryCopies(
     HloModule* module, bool check_live_range_ordering,
     const absl::flat_hash_set<absl::string_view>& execution_threads) {
   XLA_VLOG_LINES(
diff --git a/third_party/xla/xla/service/copy_insertion.h b/third_party/xla/xla/service/copy_insertion.h
index a062e6d133d5de..cba44576d9283d 100644
--- a/third_party/xla/xla/service/copy_insertion.h
+++ b/third_party/xla/xla/service/copy_insertion.h
@@ -72,7 +72,7 @@ class CopyInsertion : public HloModulePass {
   // eligible for copy elision are considered for removal.
   // If check_live_range_ordering is true, check that live ranges are ordered
   // in all the existing aliased buffers.
-  Status RemoveUnnecessaryCopies(
+  absl::Status RemoveUnnecessaryCopies(
       HloModule* module, bool check_live_range_ordering = false,
       const absl::flat_hash_set<absl::string_view>& execution_threads = {});
 
@@ -86,27 +86,27 @@ class CopyInsertion : public HloModulePass {
   //
   //    (3) Constants and parameters cannot be live out of the entry computation
   //
-  Status AddSpecialCaseCopies(
+  absl::Status AddSpecialCaseCopies(
       HloModule* module,
       const absl::flat_hash_set<absl::string_view>& execution_threads = {});
 
  protected:
   // Override which requires the caller to pass in a call graph.
-  virtual Status AddSpecialCaseCopies(
+  virtual absl::Status AddSpecialCaseCopies(
       const CallGraph& call_graph,
       const absl::flat_hash_set<absl::string_view>& execution_threads,
       HloModule* module);
 
   // Add copies for conditional instructions.
-  virtual Status AddCopiesForConditional(const HloAliasAnalysis& alias_analysis,
-                                         HloInstruction* conditional);
+  virtual absl::Status AddCopiesForConditional(
+      const HloAliasAnalysis& alias_analysis, HloInstruction* conditional);
 
   // Backend specific function that decides whether an instruction can share
   // buffer with its operand.
   HloDataflowAnalysis::CanShareBuffer can_share_buffer_;
 
  private:
-  Status AddCopiesToResolveInterference(
+  absl::Status AddCopiesToResolveInterference(
       HloModule* module,
       const absl::flat_hash_set<absl::string_view>& execution_threads);
   int64_t use_region_based_live_range_analysis_;
diff --git a/third_party/xla/xla/service/cpu_gpu_shape_verifier.cc b/third_party/xla/xla/service/cpu_gpu_shape_verifier.cc
index 4ccfb47a2eb089..d3966c1cf3d000 100644
--- a/third_party/xla/xla/service/cpu_gpu_shape_verifier.cc
+++ b/third_party/xla/xla/service/cpu_gpu_shape_verifier.cc
@@ -28,7 +28,7 @@ limitations under the License.
 namespace xla {
 
 namespace {
-Status VerifyS4U4Usage(HloInstruction* instruction) {
+absl::Status VerifyS4U4Usage(HloInstruction* instruction) {
   switch (instruction->opcode()) {
     case HloOpcode::kBitcast:
     case HloOpcode::kConstant:
@@ -62,7 +62,7 @@ Status VerifyS4U4Usage(HloInstruction* instruction) {
 }
 }  // namespace
 
-Status CpuGpuShapeVerifier::Preprocess(HloInstruction* hlo) {
+absl::Status CpuGpuShapeVerifier::Preprocess(HloInstruction* hlo) {
   TF_RETURN_IF_ERROR(ShapeUtil::ForEachSubshapeWithStatus(
       hlo->shape(), [&](const Shape& shape, const ShapeIndex&) {
         if (shape.has_layout()) {
diff --git a/third_party/xla/xla/service/cpu_gpu_shape_verifier.h b/third_party/xla/xla/service/cpu_gpu_shape_verifier.h
index 78fd43649b37d9..7adee3b49495f5 100644
--- a/third_party/xla/xla/service/cpu_gpu_shape_verifier.h
+++ b/third_party/xla/xla/service/cpu_gpu_shape_verifier.h
@@ -29,7 +29,7 @@ class CpuGpuShapeVerifier : public ShapeVerifier {
   explicit CpuGpuShapeVerifier(const HloVerifierOpts& opts)
       : ShapeVerifier(opts) {}
 
-  Status Preprocess(HloInstruction* hlo) override;
+  absl::Status Preprocess(HloInstruction* hlo) override;
 };
 
 // A verifier metadata class that uses the CpuGpuShapeVerifier.
diff --git a/third_party/xla/xla/service/defuser.cc b/third_party/xla/xla/service/defuser.cc
index 2e8f9f03a82374..f539c7b67d1d66 100644
--- a/third_party/xla/xla/service/defuser.cc
+++ b/third_party/xla/xla/service/defuser.cc
@@ -45,7 +45,7 @@ absl::StatusOr<bool> Defuser::Run(
   bool changed = false;
   std::unique_ptr<CallGraph> call_graph = CallGraph::Build(module);
   TF_RETURN_IF_ERROR(call_graph->VisitNodes(
-      [&](const CallGraphNode& call_graph_node) -> Status {
+      [&](const CallGraphNode& call_graph_node) -> absl::Status {
         if (call_graph_node.computation()->IsFusionComputation()) {
           TF_RET_CHECK(call_graph_node.caller_callsites().size() == 1);
           HloInstruction* fusion_instruction =
diff --git a/third_party/xla/xla/service/dfs_hlo_visitor_with_default_test.cc b/third_party/xla/xla/service/dfs_hlo_visitor_with_default_test.cc
index a3a2329992ffd5..21821ea1f22e18 100644
--- a/third_party/xla/xla/service/dfs_hlo_visitor_with_default_test.cc
+++ b/third_party/xla/xla/service/dfs_hlo_visitor_with_default_test.cc
@@ -37,7 +37,7 @@ TEST_F(DfsHloVisitorWithDefaultTest, DefaultElementwiseTest) {
 
   class ElementwiseTestVisitor : public DfsHloVisitorWithDefault {
    public:
-    Status DefaultAction(HloInstruction* hlo) override {
+    absl::Status DefaultAction(HloInstruction* hlo) override {
       // The HLO should be neither an elementwise unary nor binary op. These
       // cases are handled in HandleElementwiseBinary/Unary.
       TF_RET_CHECK(!(hlo->IsElementwise() && hlo->operand_count() == 2))
@@ -47,13 +47,13 @@ TEST_F(DfsHloVisitorWithDefaultTest, DefaultElementwiseTest) {
       return OkStatus();
     }
 
-    Status HandleElementwiseBinary(HloInstruction* hlo) override {
+    absl::Status HandleElementwiseBinary(HloInstruction* hlo) override {
       // HLO should be elementwise binary.
       TF_RET_CHECK(hlo->IsElementwise() && hlo->operand_count() == 2)
           << hlo->ToString();
       return OkStatus();
     }
-    Status HandleElementwiseUnary(HloInstruction* hlo) override {
+    absl::Status HandleElementwiseUnary(HloInstruction* hlo) override {
       // HLO should be elementwise unary.
       TF_RET_CHECK(hlo->IsElementwise() && hlo->operand_count() == 1)
           << hlo->ToString();
diff --git a/third_party/xla/xla/service/dot_decomposer.cc b/third_party/xla/xla/service/dot_decomposer.cc
index 33e77654ff60d8..2ddf54ddce3ac3 100644
--- a/third_party/xla/xla/service/dot_decomposer.cc
+++ b/third_party/xla/xla/service/dot_decomposer.cc
@@ -48,7 +48,7 @@ namespace {
 // * Batch dimensions are the most major dimensions.
 // This requires transposing and reshaping of the lhs and rhs, and reshaping the
 // output batch to the original shape.
-Status CanonicalizeDot(HloDotInstruction* original_dot) {
+absl::Status CanonicalizeDot(HloDotInstruction* original_dot) {
   auto computation = original_dot->parent();
   const auto& original_dnums = original_dot->dot_dimension_numbers();
   const int64_t num_batch_dims = original_dnums.lhs_batch_dimensions_size();
diff --git a/third_party/xla/xla/service/dot_dimension_merger.cc b/third_party/xla/xla/service/dot_dimension_merger.cc
index ba2f26bcca4a22..5dc254f440576e 100644
--- a/third_party/xla/xla/service/dot_dimension_merger.cc
+++ b/third_party/xla/xla/service/dot_dimension_merger.cc
@@ -60,7 +60,7 @@ std::vector<int64_t> ShiftDimensions(absl::Span<const int64_t> dimensions,
 // Merge all batch dimensions into logically first one for both operands.
 class BatchDimensionMerger : public DfsHloRewriteVisitor {
  public:
-  Status HandleDot(HloInstruction* dot) override {
+  absl::Status HandleDot(HloInstruction* dot) override {
     const DotDimensionNumbers& dnums = dot->dot_dimension_numbers();
     const Shape& lhs_shape = dot->operand(0)->shape();
     const Shape& rhs_shape = dot->operand(1)->shape();
diff --git a/third_party/xla/xla/service/dump.cc b/third_party/xla/xla/service/dump.cc
index bce64ef2ecbb67..dd2843966d0792 100644
--- a/third_party/xla/xla/service/dump.cc
+++ b/third_party/xla/xla/service/dump.cc
@@ -227,8 +227,9 @@ class DataProducer {
   std::queue<std::function<std::string()>> produce_funcs_;
 };
 
-static Status WriteStringToFile(tsl::Env* env, const std::string& fname,
-                                DataProducer& data_producer, bool compressed) {
+static absl::Status WriteStringToFile(tsl::Env* env, const std::string& fname,
+                                      DataProducer& data_producer,
+                                      bool compressed) {
   std::unique_ptr<tsl::WritableFile> file;
   TF_RETURN_IF_ERROR(env->NewWritableFile(fname, &file));
   if (compressed) {
@@ -248,8 +249,8 @@ static Status WriteStringToFile(tsl::Env* env, const std::string& fname,
   }
 }
 
-static Status WriteStringToFile(tsl::Env* env, const std::string& fname,
-                                absl::string_view data, bool compressed) {
+static absl::Status WriteStringToFile(tsl::Env* env, const std::string& fname,
+                                      absl::string_view data, bool compressed) {
   if (!compressed) {
     return tsl::WriteStringToFile(env, fname, data);
   }
@@ -665,7 +666,7 @@ void DumpProtobufToFile(const tsl::protobuf::Message& proto,
     return;
   }
   const std::string path = tsl::io::JoinPath(dir, filename);
-  Status status;
+  absl::Status status;
   if (opts.dump_as_text) {
     if (text_formatter) {
       auto written_proto = text_formatter(env, proto);
diff --git a/third_party/xla/xla/service/dynamic_dimension_inference.cc b/third_party/xla/xla/service/dynamic_dimension_inference.cc
index ee74ae367d93f8..ce80a0ecd4a10a 100644
--- a/third_party/xla/xla/service/dynamic_dimension_inference.cc
+++ b/third_party/xla/xla/service/dynamic_dimension_inference.cc
@@ -117,7 +117,7 @@ class DynamicDimensionInferenceVisitor : public DfsHloRewriteVisitor {
         shape_check_mode_(shape_check_mode),
         assertion_generator_(assertion_generator) {}
 
-  Status DefaultAction(HloInstruction* hlo) override;
+  absl::Status DefaultAction(HloInstruction* hlo) override;
 
   static absl::StatusOr<bool> Run(
       HloComputation* computation, HloDataflowAnalysis& dataflow_analysis,
@@ -145,86 +145,86 @@ class DynamicDimensionInferenceVisitor : public DfsHloRewriteVisitor {
     return visitor.changed();
   }
 
-  Status HandleParameter(HloInstruction* hlo) override;
+  absl::Status HandleParameter(HloInstruction* hlo) override;
 
-  Status HandleInfeed(HloInstruction* hlo) override;
+  absl::Status HandleInfeed(HloInstruction* hlo) override;
 
-  Status HandleConstant(HloInstruction* hlo) override;
+  absl::Status HandleConstant(HloInstruction* hlo) override;
 
-  Status HandleReduce(HloInstruction* hlo) override;
+  absl::Status HandleReduce(HloInstruction* hlo) override;
 
-  Status HandleDot(HloInstruction* hlo) override;
+  absl::Status HandleDot(HloInstruction* hlo) override;
 
-  Status HandleTuple(HloInstruction* hlo) override;
+  absl::Status HandleTuple(HloInstruction* hlo) override;
 
-  Status HandleTranspose(HloInstruction* hlo) override;
+  absl::Status HandleTranspose(HloInstruction* hlo) override;
 
-  Status HandleDynamicReshape(HloInstruction* hlo) override;
+  absl::Status HandleDynamicReshape(HloInstruction* hlo) override;
 
-  Status HandleReshape(HloInstruction* hlo) override;
+  absl::Status HandleReshape(HloInstruction* hlo) override;
 
-  Status HandleSort(HloInstruction* hlo) override;
+  absl::Status HandleSort(HloInstruction* hlo) override;
 
-  Status HandlePad(HloInstruction* hlo) override;
+  absl::Status HandlePad(HloInstruction* hlo) override;
 
-  Status HandleCustomCall(HloInstruction* hlo) override;
+  absl::Status HandleCustomCall(HloInstruction* hlo) override;
 
-  Status HandleBroadcast(HloInstruction* hlo) override;
+  absl::Status HandleBroadcast(HloInstruction* hlo) override;
 
-  Status HandleGetDimensionSize(HloInstruction* hlo) override;
+  absl::Status HandleGetDimensionSize(HloInstruction* hlo) override;
 
-  Status HandleSetDimensionSize(HloInstruction* hlo) override;
+  absl::Status HandleSetDimensionSize(HloInstruction* hlo) override;
 
-  Status HandleSelect(HloInstruction* hlo) override;
+  absl::Status HandleSelect(HloInstruction* hlo) override;
 
-  Status HandleConvolution(HloInstruction* hlo) override;
+  absl::Status HandleConvolution(HloInstruction* hlo) override;
 
-  Status HandleConcatenate(HloInstruction* hlo) override;
+  absl::Status HandleConcatenate(HloInstruction* hlo) override;
 
-  Status HandleReduceWindow(HloInstruction* hlo) override;
+  absl::Status HandleReduceWindow(HloInstruction* hlo) override;
 
-  Status HandleReverse(HloInstruction* hlo) override;
+  absl::Status HandleReverse(HloInstruction* hlo) override;
 
-  Status HandleSelectAndScatter(HloInstruction* hlo) override;
+  absl::Status HandleSelectAndScatter(HloInstruction* hlo) override;
 
-  Status HandleGetTupleElement(HloInstruction* hlo) override;
+  absl::Status HandleGetTupleElement(HloInstruction* hlo) override;
 
-  Status HandleElementwiseUnary(HloInstruction* hlo) override;
+  absl::Status HandleElementwiseUnary(HloInstruction* hlo) override;
 
-  Status HandleElementwiseNary(HloInstruction* hlo);
+  absl::Status HandleElementwiseNary(HloInstruction* hlo);
 
-  Status HandleElementwiseBinary(HloInstruction* hlo) override;
+  absl::Status HandleElementwiseBinary(HloInstruction* hlo) override;
 
-  Status HandleClamp(HloInstruction* hlo) override;
+  absl::Status HandleClamp(HloInstruction* hlo) override;
 
-  Status HandleConditional(HloInstruction* hlo) override;
+  absl::Status HandleConditional(HloInstruction* hlo) override;
 
-  Status HandleWhile(HloInstruction* hlo) override;
+  absl::Status HandleWhile(HloInstruction* hlo) override;
 
-  Status HandleSlice(HloInstruction* hlo) override;
+  absl::Status HandleSlice(HloInstruction* hlo) override;
 
-  Status HandleDynamicSlice(HloInstruction* hlo) override;
+  absl::Status HandleDynamicSlice(HloInstruction* hlo) override;
 
-  Status HandleDynamicUpdateSlice(HloInstruction* hlo) override;
+  absl::Status HandleDynamicUpdateSlice(HloInstruction* hlo) override;
 
-  Status HandleGather(HloInstruction* hlo) override;
+  absl::Status HandleGather(HloInstruction* hlo) override;
 
-  Status HandleScatter(HloInstruction* hlo) override;
+  absl::Status HandleScatter(HloInstruction* hlo) override;
 
-  Status HandleMap(HloInstruction* hlo) override;
+  absl::Status HandleMap(HloInstruction* hlo) override;
 
-  Status HandleDomain(HloInstruction* hlo) override;
+  absl::Status HandleDomain(HloInstruction* hlo) override;
 
-  Status HandleAsyncStart(HloInstruction* hlo) override;
+  absl::Status HandleAsyncStart(HloInstruction* hlo) override;
 
-  Status HandleAsyncDone(HloInstruction* hlo) override;
+  absl::Status HandleAsyncDone(HloInstruction* hlo) override;
 
  private:
-  using OperandDynamicDimensionFn = absl::FunctionRef<Status(
+  using OperandDynamicDimensionFn = absl::FunctionRef<absl::Status(
       HloInstruction* operand, ShapeIndex index, int64_t dimension,
       int64_t operand_index, HloInstruction* dynamic_size)>;
 
-  using DynamicDimensionFn = std::function<Status(
+  using DynamicDimensionFn = std::function<absl::Status(
       ShapeIndex index, int64_t dimension, HloInstruction* dynamic_size)>;
 
   void SetDynamicSize(HloInstruction* inst, const ShapeIndex& index,
@@ -234,31 +234,31 @@ class DynamicDimensionInferenceVisitor : public DfsHloRewriteVisitor {
   void SetDynamicSizes(HloInstruction* inst, const ShapeIndex& index,
                        absl::Span<HloInstruction* const> sizes);
 
-  Status HandleDynamicConvolutionForward(HloInstruction* hlo,
-                                         int64_t operand_index,
-                                         int64_t dimension,
-                                         HloInstruction* dynamic_size);
+  absl::Status HandleDynamicConvolutionForward(HloInstruction* hlo,
+                                               int64_t operand_index,
+                                               int64_t dimension,
+                                               HloInstruction* dynamic_size);
 
-  Status HandleDynamicConvolutionKernelGrad(HloInstruction* hlo,
-                                            int64_t operand_index,
-                                            int64_t dimension);
+  absl::Status HandleDynamicConvolutionKernelGrad(HloInstruction* hlo,
+                                                  int64_t operand_index,
+                                                  int64_t dimension);
 
-  Status HandleDynamicConvolutionInputGrad(HloInstruction* hlo,
-                                           int64_t operand_index,
-                                           int64_t dimension);
+  absl::Status HandleDynamicConvolutionInputGrad(HloInstruction* hlo,
+                                                 int64_t operand_index,
+                                                 int64_t dimension);
 
-  Status HandleDynamicWindowSamePadding(HloInstruction* hlo,
-                                        HloInstruction* dynamic_size,
-                                        int64_t operand_index,
-                                        int64_t dimension);
+  absl::Status HandleDynamicWindowSamePadding(HloInstruction* hlo,
+                                              HloInstruction* dynamic_size,
+                                              int64_t operand_index,
+                                              int64_t dimension);
 
-  Status ForEachOperandDynamicDimension(HloInstruction* inst,
-                                        OperandDynamicDimensionFn);
-  Status ForEachDynamicDimensionInOperand(HloInstruction* inst,
-                                          int64_t operand_index,
-                                          OperandDynamicDimensionFn);
-  Status ForEachDynamicDimension(HloInstruction* inst,
-                                 const DynamicDimensionFn& fn);
+  absl::Status ForEachOperandDynamicDimension(HloInstruction* inst,
+                                              OperandDynamicDimensionFn);
+  absl::Status ForEachDynamicDimensionInOperand(HloInstruction* inst,
+                                                int64_t operand_index,
+                                                OperandDynamicDimensionFn);
+  absl::Status ForEachDynamicDimension(HloInstruction* inst,
+                                       const DynamicDimensionFn& fn);
 
   bool CanInfer(HloInstruction* hlo) { return parent_->CanInfer(hlo); }
 
@@ -274,18 +274,18 @@ class DynamicDimensionInferenceVisitor : public DfsHloRewriteVisitor {
   // produces a tuple, each tuple component will be considered independently.
   // Returns the original instruction, with all arrays converted to static
   // shapes.
-  Status InsertPadToStaticOnInstruction(HloInstruction* inst);
+  absl::Status InsertPadToStaticOnInstruction(HloInstruction* inst);
 
   // Insert shape check to make sure `dim1` is equal to `dim2`. If
   // support_implicit_broadcast is true, the check will pass if either of them
   // is 1, even if they are different.
-  Status InsertShapeCheck(HloInstruction* dim1, HloInstruction* dim2,
-                          bool support_implicit_broadcast);
+  absl::Status InsertShapeCheck(HloInstruction* dim1, HloInstruction* dim2,
+                                bool support_implicit_broadcast);
 
   // Pass through a dynamic dimension from the input to the output with the
   // same value and index in the shape. This is a helper function to handle
   // trivial instructions like elementwise operations.
-  Status PassThroughDynamicDimension(HloInstruction*);
+  absl::Status PassThroughDynamicDimension(HloInstruction*);
 
   // The dynamic parameter bindings of this computation.
   const DynamicParameterBinding& param_bindings_;
@@ -333,7 +333,8 @@ void DynamicDimensionInferenceVisitor::SetDynamicSizes(
   }
 }
 
-Status DynamicDimensionInferenceVisitor::DefaultAction(HloInstruction* hlo) {
+absl::Status DynamicDimensionInferenceVisitor::DefaultAction(
+    HloInstruction* hlo) {
   return ForEachOperandDynamicDimension(
       hlo, [&](HloInstruction* operand, ShapeIndex index, int64_t dimension,
                int64_t operand_index, HloInstruction* dynamic_size) {
@@ -344,7 +345,7 @@ Status DynamicDimensionInferenceVisitor::DefaultAction(HloInstruction* hlo) {
       });
 }
 
-Status DynamicDimensionInferenceVisitor::HandleGetTupleElement(
+absl::Status DynamicDimensionInferenceVisitor::HandleGetTupleElement(
     HloInstruction* hlo) {
   if (!CanInfer(hlo)) {
     return OkStatus();
@@ -352,7 +353,7 @@ Status DynamicDimensionInferenceVisitor::HandleGetTupleElement(
   return ForEachOperandDynamicDimension(
       hlo,
       [&](HloInstruction* operand, ShapeIndex index, int64_t dimension,
-          int64_t operand_index, HloInstruction* dynamic_size) -> Status {
+          int64_t operand_index, HloInstruction* dynamic_size) -> absl::Status {
         if (hlo->tuple_index() != index[0]) {
           return OkStatus();
         }
@@ -362,7 +363,8 @@ Status DynamicDimensionInferenceVisitor::HandleGetTupleElement(
       });
 }
 
-Status DynamicDimensionInferenceVisitor::HandleTuple(HloInstruction* hlo) {
+absl::Status DynamicDimensionInferenceVisitor::HandleTuple(
+    HloInstruction* hlo) {
   if (!CanInfer(hlo)) {
     return OkStatus();
   }
@@ -376,7 +378,8 @@ Status DynamicDimensionInferenceVisitor::HandleTuple(HloInstruction* hlo) {
   return OkStatus();
 }
 
-Status DynamicDimensionInferenceVisitor::HandleBroadcast(HloInstruction* hlo) {
+absl::Status DynamicDimensionInferenceVisitor::HandleBroadcast(
+    HloInstruction* hlo) {
   if (!CanInfer(hlo)) {
     return OkStatus();
   }
@@ -389,7 +392,8 @@ Status DynamicDimensionInferenceVisitor::HandleBroadcast(HloInstruction* hlo) {
       });
 }
 
-Status DynamicDimensionInferenceVisitor::HandleConstant(HloInstruction* hlo) {
+absl::Status DynamicDimensionInferenceVisitor::HandleConstant(
+    HloInstruction* hlo) {
   if (!hlo->shape().is_dynamic()) {
     return OkStatus();
   }
@@ -398,7 +402,8 @@ Status DynamicDimensionInferenceVisitor::HandleConstant(HloInstruction* hlo) {
   Shape padded_shape = constant->shape();
   bool pad_any = false;
   TF_RETURN_IF_ERROR(ShapeUtil::ForEachMutableSubshapeWithStatus(
-      &padded_shape, [&](Shape* subshape, const ShapeIndex& index) -> Status {
+      &padded_shape,
+      [&](Shape* subshape, const ShapeIndex& index) -> absl::Status {
         if (!subshape->IsArray()) {
           return OkStatus();
         }
@@ -437,7 +442,7 @@ Status DynamicDimensionInferenceVisitor::HandleConstant(HloInstruction* hlo) {
   TF_RETURN_IF_ERROR(constant->ReplaceAllUsesWith(padded_constant));
   SetVisited(*padded_constant);
   TF_RETURN_IF_ERROR(do_pad.ForEachElementWithStatus(
-      [&](const ShapeIndex& index, bool requires_pad) -> Status {
+      [&](const ShapeIndex& index, bool requires_pad) -> absl::Status {
         if (!requires_pad) {
           return OkStatus();
         }
@@ -460,7 +465,8 @@ Status DynamicDimensionInferenceVisitor::HandleConstant(HloInstruction* hlo) {
   return OkStatus();
 }
 
-Status DynamicDimensionInferenceVisitor::HandleCustomCall(HloInstruction* hlo) {
+absl::Status DynamicDimensionInferenceVisitor::HandleCustomCall(
+    HloInstruction* hlo) {
   if (hlo->custom_call_target() == "PadToStatic") {
     for (int64_t i = 0; i < hlo->operand(0)->shape().rank(); ++i) {
       if (hlo->operand(0)->shape().is_dynamic_dimension(i)) {
@@ -543,7 +549,7 @@ Status DynamicDimensionInferenceVisitor::HandleCustomCall(HloInstruction* hlo) {
   return InsertPadToStaticOnInstruction(hlo);
 }
 
-Status DynamicDimensionInferenceVisitor::HandleSort(HloInstruction* hlo) {
+absl::Status DynamicDimensionInferenceVisitor::HandleSort(HloInstruction* hlo) {
   if (!CanInfer(hlo)) {
     return OkStatus();
   }
@@ -562,7 +568,7 @@ Status DynamicDimensionInferenceVisitor::HandleSort(HloInstruction* hlo) {
       });
 }
 
-Status DynamicDimensionInferenceVisitor::HandlePad(HloInstruction* hlo) {
+absl::Status DynamicDimensionInferenceVisitor::HandlePad(HloInstruction* hlo) {
   if (!CanInfer(hlo)) {
     return OkStatus();
   }
@@ -620,7 +626,8 @@ Status DynamicDimensionInferenceVisitor::HandlePad(HloInstruction* hlo) {
       });
 }
 
-Status DynamicDimensionInferenceVisitor::HandleReduce(HloInstruction* hlo) {
+absl::Status DynamicDimensionInferenceVisitor::HandleReduce(
+    HloInstruction* hlo) {
   if (!CanInfer(hlo)) {
     return OkStatus();
   }
@@ -628,7 +635,7 @@ Status DynamicDimensionInferenceVisitor::HandleReduce(HloInstruction* hlo) {
   int64_t rank = -1;
   TF_RETURN_IF_ERROR(ShapeUtil::ForEachSubshapeWithStatus(
       reduce->shape(),
-      [&](const Shape& subshape, const ShapeIndex& index) -> Status {
+      [&](const Shape& subshape, const ShapeIndex& index) -> absl::Status {
         if (!subshape.IsArray()) {
           return OkStatus();
         }
@@ -686,7 +693,7 @@ Status DynamicDimensionInferenceVisitor::HandleReduce(HloInstruction* hlo) {
   return OkStatus();
 }
 
-Status DynamicDimensionInferenceVisitor::HandleDot(HloInstruction* hlo) {
+absl::Status DynamicDimensionInferenceVisitor::HandleDot(HloInstruction* hlo) {
   if (!CanInfer(hlo)) {
     return OkStatus();
   }
@@ -696,7 +703,7 @@ Status DynamicDimensionInferenceVisitor::HandleDot(HloInstruction* hlo) {
       hlo,
       [&](HloInstruction* operand, ShapeIndex operand_shape_index,
           int64_t operand_dimension, int64_t operand_index,
-          HloInstruction* dynamic_size) -> Status {
+          HloInstruction* dynamic_size) -> absl::Status {
         // There are three types of dimensions in a dot:
         // A. batch dims
         // B. contracting dims
@@ -780,14 +787,15 @@ Status DynamicDimensionInferenceVisitor::HandleDot(HloInstruction* hlo) {
   return OkStatus();
 }
 
-Status DynamicDimensionInferenceVisitor::HandleTranspose(HloInstruction* hlo) {
+absl::Status DynamicDimensionInferenceVisitor::HandleTranspose(
+    HloInstruction* hlo) {
   if (!CanInfer(hlo)) {
     return OkStatus();
   }
   return ForEachOperandDynamicDimension(
       hlo,
       [&](HloInstruction* operand, ShapeIndex index, int64_t dimension,
-          int64_t operand_index, HloInstruction* dynamic_size) -> Status {
+          int64_t operand_index, HloInstruction* dynamic_size) -> absl::Status {
         int64_t permuted_dim = -1;
         for (int64_t i = 0; i < hlo->dimensions().size(); ++i) {
           if (hlo->dimensions()[i] == dimension) {
@@ -800,7 +808,7 @@ Status DynamicDimensionInferenceVisitor::HandleTranspose(HloInstruction* hlo) {
       });
 }
 
-Status DynamicDimensionInferenceVisitor::HandleConvolution(
+absl::Status DynamicDimensionInferenceVisitor::HandleConvolution(
     HloInstruction* hlo) {
   if (!CanInfer(hlo)) {
     return OkStatus();
@@ -833,7 +841,7 @@ Status DynamicDimensionInferenceVisitor::HandleConvolution(
       });
 }
 
-Status DynamicDimensionInferenceVisitor::HandleConcatenate(
+absl::Status DynamicDimensionInferenceVisitor::HandleConcatenate(
     HloInstruction* hlo) {
   if (!CanInfer(hlo)) {
     return OkStatus();
@@ -882,7 +890,7 @@ Status DynamicDimensionInferenceVisitor::HandleConcatenate(
   TF_RETURN_IF_ERROR(ForEachOperandDynamicDimension(
       hlo,
       [&](HloInstruction* operand, ShapeIndex index, int64_t dimension,
-          int64_t operand_index, HloInstruction* dynamic_size) -> Status {
+          int64_t operand_index, HloInstruction* dynamic_size) -> absl::Status {
         TF_RET_CHECK(index.empty());
         int64_t concatenate_dimension = hlo->concatenate_dimension();
         if (concatenate_dimension == dimension) {
@@ -897,7 +905,7 @@ Status DynamicDimensionInferenceVisitor::HandleConcatenate(
   return OkStatus();
 }
 
-Status DynamicDimensionInferenceVisitor::HandleGetDimensionSize(
+absl::Status DynamicDimensionInferenceVisitor::HandleGetDimensionSize(
     HloInstruction* gds) {
   // Dynamic dimension doesn't propagate through GetDimensionSize:
   //
@@ -934,7 +942,7 @@ Status DynamicDimensionInferenceVisitor::HandleGetDimensionSize(
   return OkStatus();
 }
 
-Status DynamicDimensionInferenceVisitor::HandleSetDimensionSize(
+absl::Status DynamicDimensionInferenceVisitor::HandleSetDimensionSize(
     HloInstruction* hlo) {
   bool dimension_is_static = false;
   const HloInstruction* size = hlo->operand(1);
@@ -965,7 +973,7 @@ Status DynamicDimensionInferenceVisitor::HandleSetDimensionSize(
   TF_RETURN_IF_ERROR(ForEachOperandDynamicDimension(
       hlo,
       [&](HloInstruction* operand, ShapeIndex index, int64_t dimension,
-          int64_t operand_index, HloInstruction* dynamic_size) -> Status {
+          int64_t operand_index, HloInstruction* dynamic_size) -> absl::Status {
         TF_RET_CHECK(operand_index == 0);
         if (dimension != hlo->dimension()) {
           SetDynamicSize(hlo, index, dimension, dynamic_size,
@@ -977,7 +985,7 @@ Status DynamicDimensionInferenceVisitor::HandleSetDimensionSize(
   return OkStatus();
 }
 
-Status DynamicDimensionInferenceVisitor::HandleDynamicConvolutionForward(
+absl::Status DynamicDimensionInferenceVisitor::HandleDynamicConvolutionForward(
     HloInstruction* hlo, int64_t operand_index, int64_t dimension,
     HloInstruction* dynamic_size) {
   if (!CanInfer(hlo)) {
@@ -1017,7 +1025,7 @@ Status DynamicDimensionInferenceVisitor::HandleDynamicConvolutionForward(
   return OkStatus();
 }
 
-Status DynamicDimensionInferenceVisitor::HandleDynamicWindowSamePadding(
+absl::Status DynamicDimensionInferenceVisitor::HandleDynamicWindowSamePadding(
     HloInstruction* hlo, HloInstruction* dynamic_size, int64_t operand_index,
     int64_t dimension) {
   if (!CanInfer(hlo)) {
@@ -1037,7 +1045,8 @@ Status DynamicDimensionInferenceVisitor::HandleDynamicWindowSamePadding(
   return OkStatus();
 }
 
-Status DynamicDimensionInferenceVisitor::HandleDynamicConvolutionInputGrad(
+absl::Status
+DynamicDimensionInferenceVisitor::HandleDynamicConvolutionInputGrad(
     HloInstruction* hlo, int64_t operand_index, int64_t dimension) {
   if (!CanInfer(hlo)) {
     return OkStatus();
@@ -1060,13 +1069,14 @@ Status DynamicDimensionInferenceVisitor::HandleDynamicConvolutionInputGrad(
   return OkStatus();
 }
 
-Status DynamicDimensionInferenceVisitor::HandleDynamicConvolutionKernelGrad(
+absl::Status
+DynamicDimensionInferenceVisitor::HandleDynamicConvolutionKernelGrad(
     HloInstruction* hlo, int64_t operand_index, int64_t dimension) {
   // Dynamic convolution kernel grad produces static shape outputs.
   return OkStatus();
 }
 
-Status DynamicDimensionInferenceVisitor::PassThroughDynamicDimension(
+absl::Status DynamicDimensionInferenceVisitor::PassThroughDynamicDimension(
     HloInstruction* hlo) {
   if (!CanInfer(hlo)) {
     return OkStatus();
@@ -1093,11 +1103,13 @@ Status DynamicDimensionInferenceVisitor::PassThroughDynamicDimension(
   return OkStatus();
 }
 
-Status DynamicDimensionInferenceVisitor::HandleDomain(HloInstruction* hlo) {
+absl::Status DynamicDimensionInferenceVisitor::HandleDomain(
+    HloInstruction* hlo) {
   return PassThroughDynamicDimension(hlo);
 }
 
-Status DynamicDimensionInferenceVisitor::HandleAsyncStart(HloInstruction* hlo) {
+absl::Status DynamicDimensionInferenceVisitor::HandleAsyncStart(
+    HloInstruction* hlo) {
   if (!HloInstruction::IsThreadIncluded(hlo->async_execution_thread(),
                                         parent_->execution_threads_)) {
     // Async-start not included in specified execution thread set will use
@@ -1108,7 +1120,8 @@ Status DynamicDimensionInferenceVisitor::HandleAsyncStart(HloInstruction* hlo) {
   return DefaultAction(hlo);
 }
 
-Status DynamicDimensionInferenceVisitor::HandleAsyncDone(HloInstruction* hlo) {
+absl::Status DynamicDimensionInferenceVisitor::HandleAsyncDone(
+    HloInstruction* hlo) {
   if (!HloInstruction::IsThreadIncluded(hlo->async_execution_thread(),
                                         parent_->execution_threads_)) {
     // Other threads can return a dynamic shape directly, so we may need to
@@ -1118,16 +1131,17 @@ Status DynamicDimensionInferenceVisitor::HandleAsyncDone(HloInstruction* hlo) {
   return DefaultAction(hlo);
 }
 
-Status DynamicDimensionInferenceVisitor::HandleElementwiseUnary(
+absl::Status DynamicDimensionInferenceVisitor::HandleElementwiseUnary(
     HloInstruction* hlo) {
   return PassThroughDynamicDimension(hlo);
 }
 
-Status DynamicDimensionInferenceVisitor::HandleSelect(HloInstruction* hlo) {
+absl::Status DynamicDimensionInferenceVisitor::HandleSelect(
+    HloInstruction* hlo) {
   return HandleElementwiseNary(hlo);
 }
 
-Status DynamicDimensionInferenceVisitor::HandleElementwiseNary(
+absl::Status DynamicDimensionInferenceVisitor::HandleElementwiseNary(
     HloInstruction* hlo) {
   if (!CanInfer(hlo)) {
     return OkStatus();
@@ -1141,7 +1155,7 @@ Status DynamicDimensionInferenceVisitor::HandleElementwiseNary(
   TF_RETURN_IF_ERROR(ForEachOperandDynamicDimension(
       hlo,
       [&](HloInstruction* operand, ShapeIndex index, int64_t dimension,
-          int64_t operand_index, HloInstruction* dynamic_size) -> Status {
+          int64_t operand_index, HloInstruction* dynamic_size) -> absl::Status {
         TF_RET_CHECK(index.empty());
         operand_sizes[dimension][operand_index] = dynamic_size;
         return OkStatus();
@@ -1214,16 +1228,17 @@ Status DynamicDimensionInferenceVisitor::HandleElementwiseNary(
   return OkStatus();
 }
 
-Status DynamicDimensionInferenceVisitor::HandleElementwiseBinary(
+absl::Status DynamicDimensionInferenceVisitor::HandleElementwiseBinary(
     HloInstruction* hlo) {
   return HandleElementwiseNary(hlo);
 }
 
-Status DynamicDimensionInferenceVisitor::HandleClamp(HloInstruction* hlo) {
+absl::Status DynamicDimensionInferenceVisitor::HandleClamp(
+    HloInstruction* hlo) {
   return PassThroughDynamicDimension(hlo);
 }
 
-Status DynamicDimensionInferenceVisitor::HandleDynamicReshape(
+absl::Status DynamicDimensionInferenceVisitor::HandleDynamicReshape(
     HloInstruction* hlo) {
   if (!CanInfer(hlo)) {
     return OkStatus();
@@ -1238,7 +1253,7 @@ Status DynamicDimensionInferenceVisitor::HandleDynamicReshape(
   return OkStatus();
 }
 
-Status DynamicDimensionInferenceVisitor::HandleReshape(
+absl::Status DynamicDimensionInferenceVisitor::HandleReshape(
     HloInstruction* const hlo) {
   if (!CanInfer(hlo)) {
     return OkStatus();
@@ -1326,7 +1341,7 @@ Status DynamicDimensionInferenceVisitor::HandleReshape(
                   op,
                   [&](HloInstruction* operand, ShapeIndex index,
                       int64_t op_dynamic_dimension, int64_t operand_index,
-                      HloInstruction* operand_dynamic_size) -> Status {
+                      HloInstruction* operand_dynamic_size) -> absl::Status {
                     if (op_dynamic_dimension >= orig_reshape_pair.first.first &&
                         op_dynamic_dimension < orig_reshape_pair.first.second) {
                       auto dynamic_size =
@@ -1368,7 +1383,7 @@ Status DynamicDimensionInferenceVisitor::HandleReshape(
       hlo,
       [&](HloInstruction* operand, ShapeIndex index,
           int64_t input_dynamic_dimension, int64_t operand_index,
-          HloInstruction* operand_dynamic_size) -> Status {
+          HloInstruction* operand_dynamic_size) -> absl::Status {
         auto reshape_pair =
             find_reshape_group_pair(hlo, input_dynamic_dimension);
         reshape_group_pairs[input_dynamic_dimension] = reshape_pair;
@@ -1429,7 +1444,7 @@ Status DynamicDimensionInferenceVisitor::HandleReshape(
       hlo,
       [&](HloInstruction* operand, ShapeIndex index,
           int64_t input_dynamic_dimension, int64_t operand_index,
-          HloInstruction* operand_dynamic_size) -> Status {
+          HloInstruction* operand_dynamic_size) -> absl::Status {
         HloInstruction* const reshape = hlo;
         if (reshape->shape().rank() == 0) {
           VLOG(0) << "Reshaping a dynamic dimension into a scalar, which has "
@@ -1619,7 +1634,7 @@ Status DynamicDimensionInferenceVisitor::HandleReshape(
   return OkStatus();
 }
 
-Status DynamicDimensionInferenceVisitor::HandleReduceWindow(
+absl::Status DynamicDimensionInferenceVisitor::HandleReduceWindow(
     HloInstruction* hlo) {
   if (!CanInfer(hlo)) {
     return OkStatus();
@@ -1675,7 +1690,7 @@ Status DynamicDimensionInferenceVisitor::HandleReduceWindow(
   return OkStatus();
 }
 
-Status DynamicDimensionInferenceVisitor::HandleSelectAndScatter(
+absl::Status DynamicDimensionInferenceVisitor::HandleSelectAndScatter(
     HloInstruction* hlo) {
   if (!CanInfer(hlo)) {
     return OkStatus();
@@ -1694,14 +1709,16 @@ Status DynamicDimensionInferenceVisitor::HandleSelectAndScatter(
       });
 }
 
-Status DynamicDimensionInferenceVisitor::HandleSlice(HloInstruction* hlo) {
+absl::Status DynamicDimensionInferenceVisitor::HandleSlice(
+    HloInstruction* hlo) {
   if (!CanInfer(hlo)) {
     return OkStatus();
   }
   return ForEachOperandDynamicDimension(
       hlo,
       [&](HloInstruction* operand, ShapeIndex /*index*/, int64_t dimension,
-          int64_t /*operand_index*/, HloInstruction* dynamic_size) -> Status {
+          int64_t /*operand_index*/,
+          HloInstruction* dynamic_size) -> absl::Status {
         int64_t start = hlo->slice_starts(dimension);
         int64_t limit = hlo->slice_limits(dimension);
         int64_t stride = hlo->slice_strides(dimension);
@@ -1735,7 +1752,7 @@ Status DynamicDimensionInferenceVisitor::HandleSlice(HloInstruction* hlo) {
       });
 }
 
-Status DynamicDimensionInferenceVisitor::HandleDynamicSlice(
+absl::Status DynamicDimensionInferenceVisitor::HandleDynamicSlice(
     HloInstruction* hlo) {
   if (!CanInfer(hlo)) {
     return OkStatus();
@@ -1743,7 +1760,7 @@ Status DynamicDimensionInferenceVisitor::HandleDynamicSlice(
   return ForEachOperandDynamicDimension(
       hlo,
       [&](HloInstruction* operand, ShapeIndex index, int64_t dimension,
-          int64_t operand_index, HloInstruction* dynamic_size) -> Status {
+          int64_t operand_index, HloInstruction* dynamic_size) -> absl::Status {
         // Slicing a single element out kills the dynamic dimension.
         if (hlo->shape().dimensions(dimension) == 1) {
           return OkStatus();
@@ -1766,7 +1783,7 @@ Status DynamicDimensionInferenceVisitor::HandleDynamicSlice(
       });
 }
 
-Status DynamicDimensionInferenceVisitor::HandleDynamicUpdateSlice(
+absl::Status DynamicDimensionInferenceVisitor::HandleDynamicUpdateSlice(
     HloInstruction* hlo) {
   if (!CanInfer(hlo)) {
     return OkStatus();
@@ -1776,7 +1793,7 @@ Status DynamicDimensionInferenceVisitor::HandleDynamicUpdateSlice(
   TF_RETURN_IF_ERROR(ForEachOperandDynamicDimension(
       hlo,
       [&](HloInstruction* operand, ShapeIndex index, int64_t dimension,
-          int64_t operand_index, HloInstruction* dynamic_size) -> Status {
+          int64_t operand_index, HloInstruction* dynamic_size) -> absl::Status {
         TF_RET_CHECK(index.empty());
 
         if (hlo->shape().dimensions(dimension) !=
@@ -1808,11 +1825,13 @@ Status DynamicDimensionInferenceVisitor::HandleDynamicUpdateSlice(
   return OkStatus();
 }
 
-Status DynamicDimensionInferenceVisitor::HandleReverse(HloInstruction* hlo) {
+absl::Status DynamicDimensionInferenceVisitor::HandleReverse(
+    HloInstruction* hlo) {
   return PassThroughDynamicDimension(hlo);
 }
 
-Status DynamicDimensionInferenceVisitor::HandleGather(HloInstruction* hlo) {
+absl::Status DynamicDimensionInferenceVisitor::HandleGather(
+    HloInstruction* hlo) {
   if (!CanInfer(hlo)) {
     return OkStatus();
   }
@@ -1822,7 +1841,7 @@ Status DynamicDimensionInferenceVisitor::HandleGather(HloInstruction* hlo) {
       hlo,
       [&](HloInstruction* operand, ShapeIndex /*index*/,
           int64_t input_dynamic_dimension, int64_t operand_index,
-          HloInstruction* dynamic_size) -> Status {
+          HloInstruction* dynamic_size) -> absl::Status {
         const GatherDimensionNumbers& gather_dims =
             hlo->gather_dimension_numbers();
         if (operand_index == 0) {
@@ -1887,7 +1906,7 @@ Status DynamicDimensionInferenceVisitor::HandleGather(HloInstruction* hlo) {
   return OkStatus();
 }
 
-Status DynamicDimensionInferenceVisitor::HandleConditional(
+absl::Status DynamicDimensionInferenceVisitor::HandleConditional(
     HloInstruction* hlo) {
   if (!CanInfer(hlo)) {
     return OkStatus();
@@ -1923,7 +1942,7 @@ Status DynamicDimensionInferenceVisitor::HandleConditional(
     TF_RETURN_IF_ERROR(ForEachDynamicDimensionInOperand(
         hlo, operand_index,
         [&](HloInstruction*, ShapeIndex, int64_t, int64_t,
-            HloInstruction* dynamic_size) -> Status {
+            HloInstruction* dynamic_size) -> absl::Status {
           TF_RET_CHECK(hlo->operand(operand_index)->shape().IsTuple())
               << "Only tuple typed inputs can have dynamic dimension. Please "
                  "file a bug against XLA team.";
@@ -2110,14 +2129,15 @@ Status DynamicDimensionInferenceVisitor::HandleConditional(
   return OkStatus();
 }
 
-Status DynamicDimensionInferenceVisitor::HandleMap(HloInstruction* hlo) {
+absl::Status DynamicDimensionInferenceVisitor::HandleMap(HloInstruction* hlo) {
   if (!CanInfer(hlo)) {
     return OkStatus();
   }
   return HandleElementwiseNary(hlo);
 }
 
-Status DynamicDimensionInferenceVisitor::HandleScatter(HloInstruction* hlo) {
+absl::Status DynamicDimensionInferenceVisitor::HandleScatter(
+    HloInstruction* hlo) {
   if (!CanInfer(hlo)) {
     return OkStatus();
   }
@@ -2186,7 +2206,8 @@ Status DynamicDimensionInferenceVisitor::HandleScatter(HloInstruction* hlo) {
       });
 }
 
-Status DynamicDimensionInferenceVisitor::HandleWhile(HloInstruction* hlo) {
+absl::Status DynamicDimensionInferenceVisitor::HandleWhile(
+    HloInstruction* hlo) {
   if (!CanInfer(hlo)) {
     return OkStatus();
   }
@@ -2207,7 +2228,7 @@ Status DynamicDimensionInferenceVisitor::HandleWhile(HloInstruction* hlo) {
   TF_RETURN_IF_ERROR(ForEachOperandDynamicDimension(
       hlo,
       [&](HloInstruction* operand, ShapeIndex index, int64_t dim,
-          int64_t operand_num, HloInstruction* dynamic_size) -> Status {
+          int64_t operand_num, HloInstruction* dynamic_size) -> absl::Status {
         TF_RET_CHECK(operand_num == 0);
         operands_to_add.push_back(dynamic_size);
         dynamic_output_mapping.mutable_element(index)->emplace(dim,
@@ -2298,7 +2319,8 @@ Status DynamicDimensionInferenceVisitor::HandleWhile(HloInstruction* hlo) {
   // Add dynamic dimension size as new outputs of the while loop body.
   TF_RETURN_IF_ERROR(dynamic_output_mapping.ForEachElementWithStatus(
       [&](const ShapeIndex& index,
-          const absl::flat_hash_map<int64_t, int64_t>& dim_to_size) -> Status {
+          const absl::flat_hash_map<int64_t, int64_t>& dim_to_size)
+          -> absl::Status {
         for (auto [dimension, output_index] : dim_to_size) {
           TF_RET_CHECK(new_root_operands[output_index] == nullptr);
           HloInstruction* dynamic_size =
@@ -2317,7 +2339,7 @@ Status DynamicDimensionInferenceVisitor::HandleWhile(HloInstruction* hlo) {
     TF_RETURN_IF_ERROR(ForEachDynamicDimension(
         body_root,
         [&](ShapeIndex index, int64_t dimension,
-            HloInstruction* dynamic_size) -> Status {
+            HloInstruction* dynamic_size) -> absl::Status {
           SetDynamicSize(new_body_root, index, dimension, dynamic_size);
           if (index.empty() || index.front() != i) {
             return OkStatus();
@@ -2333,7 +2355,8 @@ Status DynamicDimensionInferenceVisitor::HandleWhile(HloInstruction* hlo) {
   // Record the dynamic sizes of while loop output.
   return dynamic_output_mapping.ForEachElementWithStatus(
       [&](const ShapeIndex& index,
-          const absl::flat_hash_map<int64_t, int64_t>& dim_to_size) -> Status {
+          const absl::flat_hash_map<int64_t, int64_t>& dim_to_size)
+          -> absl::Status {
         for (auto [dimension, output_index] : dim_to_size) {
           HloInstruction* dynamic_size = hlo->AddInstruction(
               HloInstruction::CreateGetTupleElement(hlo, output_index));
@@ -2354,7 +2377,8 @@ Status DynamicDimensionInferenceVisitor::HandleWhile(HloInstruction* hlo) {
       });
 }
 
-Status DynamicDimensionInferenceVisitor::HandleParameter(HloInstruction* hlo) {
+absl::Status DynamicDimensionInferenceVisitor::HandleParameter(
+    HloInstruction* hlo) {
   if (hlo->parent()->IsEntryComputation()) {
     TF_RET_CHECK(param_bindings_.empty());
     return InsertPadToStaticOnInstruction(hlo);
@@ -2363,7 +2387,7 @@ Status DynamicDimensionInferenceVisitor::HandleParameter(HloInstruction* hlo) {
   return param_bindings_.ForEachBinding(
       [&](const DynamicParameterBinding::DynamicSizeParameter& dynamic_size,
           const DynamicParameterBinding::DynamicDimension& dynamic_dimension)
-          -> Status {
+          -> absl::Status {
         if (dynamic_dimension.parameter_num == hlo->parameter_number()) {
           SetDynamicSize(
               hlo, dynamic_dimension.parameter_index,
@@ -2378,11 +2402,12 @@ Status DynamicDimensionInferenceVisitor::HandleParameter(HloInstruction* hlo) {
       });
 }
 
-Status DynamicDimensionInferenceVisitor::HandleInfeed(HloInstruction* hlo) {
+absl::Status DynamicDimensionInferenceVisitor::HandleInfeed(
+    HloInstruction* hlo) {
   return InsertPadToStaticOnInstruction(hlo);
 }
 
-Status DynamicDimensionInferenceVisitor::ForEachDynamicDimension(
+absl::Status DynamicDimensionInferenceVisitor::ForEachDynamicDimension(
     HloInstruction* inst, const DynamicDimensionFn& fn) {
   auto iter = parent_->per_hlo_dynamic_dimensions_.find(inst);
   if (iter != parent_->per_hlo_dynamic_dimensions_.end()) {
@@ -2454,7 +2479,7 @@ absl::StatusOr<bool> DynamicDimensionInferenceVisitor::RequiresPadToStatic(
 // Insert pad-to-static after `inst` if `inst` has dynamic dimensions in it.
 // If the instruction produces a tuple, each tuple component will be considered
 // independently.
-Status DynamicDimensionInferenceVisitor::InsertPadToStaticOnInstruction(
+absl::Status DynamicDimensionInferenceVisitor::InsertPadToStaticOnInstruction(
     HloInstruction* inst) {
   if (inst->shape().is_static()) {
     return OkStatus();
@@ -2490,7 +2515,8 @@ Status DynamicDimensionInferenceVisitor::InsertPadToStaticOnInstruction(
   ShapeTree<HloInstruction*> padded(inst->shape(), nullptr);
   TF_RETURN_IF_ERROR(ShapeUtil::ForEachSubshapePostOrderWithStatus(
       inst->shape(),
-      [&](const Shape& subshape, const ShapeIndex& shape_index) -> Status {
+      [&](const Shape& subshape,
+          const ShapeIndex& shape_index) -> absl::Status {
         HloInstruction* element = gtes.element(shape_index);
         SetVisited(*gtes.element(shape_index));
         if (subshape.IsTuple()) {
@@ -2574,7 +2600,7 @@ Status DynamicDimensionInferenceVisitor::InsertPadToStaticOnInstruction(
   return OkStatus();
 }
 
-Status DynamicDimensionInferenceVisitor::InsertShapeCheck(
+absl::Status DynamicDimensionInferenceVisitor::InsertShapeCheck(
     HloInstruction* dim1, HloInstruction* dim2,
     bool support_implicit_broadcast) {
   switch (shape_check_mode_) {
@@ -2603,7 +2629,7 @@ Status DynamicDimensionInferenceVisitor::InsertShapeCheck(
   }
 }
 
-Status DynamicDimensionInferenceVisitor::ForEachDynamicDimensionInOperand(
+absl::Status DynamicDimensionInferenceVisitor::ForEachDynamicDimensionInOperand(
     HloInstruction* inst, int64_t operand_index, OperandDynamicDimensionFn fn) {
   auto iter =
       parent_->per_hlo_dynamic_dimensions_.find(inst->operand(operand_index));
@@ -2620,7 +2646,7 @@ Status DynamicDimensionInferenceVisitor::ForEachDynamicDimensionInOperand(
   return OkStatus();
 }
 
-Status DynamicDimensionInferenceVisitor::ForEachOperandDynamicDimension(
+absl::Status DynamicDimensionInferenceVisitor::ForEachOperandDynamicDimension(
     HloInstruction* inst, OperandDynamicDimensionFn fn) {
   for (int64_t operand_index = 0; operand_index < inst->operand_count();
        ++operand_index) {
@@ -2714,7 +2740,7 @@ DynamicDimensionInference::DynamicDimensionInference(
       assertion_generator_(assertion_generator),
       execution_threads_(execution_threads) {}
 
-Status DynamicDimensionInference::AnalyzeDynamicDimensions() {
+absl::Status DynamicDimensionInference::AnalyzeDynamicDimensions() {
   TF_ASSIGN_OR_RETURN(
       std::unique_ptr<HloDataflowAnalysis> dataflow_analysis,
       HloDataflowAnalysis::Run(*module_, /*ssa_form=*/false,
@@ -2749,9 +2775,8 @@ void DynamicDimensionInference::ReplaceAllDynamicDimensionUsesWith(
   }
 }
 
-Status DynamicDimensionInference::ForwardDynamicSize(HloInstruction* inst,
-                                                     HloInstruction* new_inst,
-                                                     const ShapeIndex& index) {
+absl::Status DynamicDimensionInference::ForwardDynamicSize(
+    HloInstruction* inst, HloInstruction* new_inst, const ShapeIndex& index) {
   TF_RET_CHECK(ShapeUtil::Compatible(inst->shape(), new_inst->shape()));
 
   for (int64_t dim = 0; dim < inst->shape().rank(); ++dim) {
diff --git a/third_party/xla/xla/service/dynamic_dimension_inference.h b/third_party/xla/xla/service/dynamic_dimension_inference.h
index 4ce0b0c3087c16..cab9502ccccbf9 100644
--- a/third_party/xla/xla/service/dynamic_dimension_inference.h
+++ b/third_party/xla/xla/service/dynamic_dimension_inference.h
@@ -71,7 +71,7 @@ class DynamicDimensionInference {
     kIgnore,
   };
   using CustomCallInferenceHandler =
-      std::function<Status(HloInstruction*, DynamicDimensionInference*)>;
+      std::function<absl::Status(HloInstruction*, DynamicDimensionInference*)>;
 
   // Generate an assertion which fails the execution if the instruction value is
   // false.
@@ -108,8 +108,9 @@ class DynamicDimensionInference {
                            ShapeIndexView index = {}) const;
 
   // Forward dynamic dimension size at `dim` from `inst` to `new_inst`.
-  Status ForwardDynamicSize(HloInstruction* inst, HloInstruction* new_inst,
-                            const ShapeIndex& index);
+  absl::Status ForwardDynamicSize(HloInstruction* inst,
+                                  HloInstruction* new_inst,
+                                  const ShapeIndex& index);
 
   // Update the dynamic mapping so that we know dimension `dim` of instruction
   // `inst` at `index` has a dynamic size, and its runtime size is represented
@@ -193,7 +194,7 @@ class DynamicDimensionInference {
 
   // AnalyzeDynamicDimensions starts the analysis of the dynamic dimensions in
   // module_.
-  Status AnalyzeDynamicDimensions();
+  absl::Status AnalyzeDynamicDimensions();
 
   // HloModule being analyzed.
   HloModule* module_;
diff --git a/third_party/xla/xla/service/dynamic_dimension_inference_test.cc b/third_party/xla/xla/service/dynamic_dimension_inference_test.cc
index 24b65c063c95e8..3c6b9d73d06a2b 100644
--- a/third_party/xla/xla/service/dynamic_dimension_inference_test.cc
+++ b/third_party/xla/xla/service/dynamic_dimension_inference_test.cc
@@ -46,7 +46,7 @@ class DynamicDimensionInferenceTest : public HloTestBase {
     module_ = CreateNewVerifiedModule();
   }
 
-  Status RunInference(
+  absl::Status RunInference(
       OpSupportsDynamismHandler op_supports_dynamism_handler = nullptr,
       DynamicDimensionInference::CustomCallInferenceHandler handler = nullptr,
       DynamicDimensionInference::ShapeCheckMode shape_check_mode =
@@ -603,7 +603,7 @@ TEST_F(DynamicDimensionInferenceTest, ReshapeTestMajorDimension) {
   module_->AddEntryComputation(builder.Build());
 
   SCOPED_TRACE(module_->ToString());
-  Status status = RunInference();
+  absl::Status status = RunInference();
   EXPECT_NE(inference_->GetDynamicSize(reshape, {}, 0), nullptr);
 }
 
diff --git a/third_party/xla/xla/service/dynamic_padder.cc b/third_party/xla/xla/service/dynamic_padder.cc
index f5e7ae0fbba8a0..7b4021c2690873 100644
--- a/third_party/xla/xla/service/dynamic_padder.cc
+++ b/third_party/xla/xla/service/dynamic_padder.cc
@@ -1850,25 +1850,25 @@ class DynamicShapeRemovingVisitor : public DfsHloRewriteVisitor {
         dynamic_dimension_inference_(dynamic_dimension_inference),
         execution_threads_(execution_threads) {}
 
-  Status DefaultAction(HloInstruction* hlo) override;
+  absl::Status DefaultAction(HloInstruction* hlo) override;
 
-  Status HandleCustomCall(HloInstruction* hlo) override;
+  absl::Status HandleCustomCall(HloInstruction* hlo) override;
 
-  Status HandleTuple(HloInstruction* hlo) override;
-  Status HandleGetTupleElement(HloInstruction* hlo) override;
+  absl::Status HandleTuple(HloInstruction* hlo) override;
+  absl::Status HandleGetTupleElement(HloInstruction* hlo) override;
 
-  Status HandleParameter(HloInstruction* hlo) override;
-  Status HandleInfeed(HloInstruction* hlo) override;
+  absl::Status HandleParameter(HloInstruction* hlo) override;
+  absl::Status HandleInfeed(HloInstruction* hlo) override;
 
-  Status HandleAsyncStart(HloInstruction* hlo) override;
-  Status HandleAsyncUpdate(HloInstruction* hlo) override;
-  Status HandleAsyncDone(HloInstruction* hlo) override;
+  absl::Status HandleAsyncStart(HloInstruction* hlo) override;
+  absl::Status HandleAsyncUpdate(HloInstruction* hlo) override;
+  absl::Status HandleAsyncDone(HloInstruction* hlo) override;
 
-  Status HandleWhile(HloInstruction* hlo) override;
-  Status HandleConditional(HloInstruction* hlo) override;
+  absl::Status HandleWhile(HloInstruction* hlo) override;
+  absl::Status HandleConditional(HloInstruction* hlo) override;
 
-  Status HandleGetDimensionSize(HloInstruction* hlo) override;
-  Status HandleSetDimensionSize(HloInstruction* hlo) override;
+  absl::Status HandleGetDimensionSize(HloInstruction* hlo) override;
+  absl::Status HandleSetDimensionSize(HloInstruction* hlo) override;
 
   static absl::StatusOr<bool> Run(
       HloComputation* computation,
@@ -1900,7 +1900,7 @@ class DynamicShapeRemovingVisitor : public DfsHloRewriteVisitor {
 
   // Same as above, but for all of the instructions operands.  The operands will
   // be replaced by dynamic operands as needed.
-  Status ConvertOperandsToDynamic(HloInstruction* inst);
+  absl::Status ConvertOperandsToDynamic(HloInstruction* inst);
 
   const OpSupportsDynamismHandler& op_supports_dynamism_handler_;
 
@@ -1947,7 +1947,7 @@ absl::StatusOr<HloInstruction*> DynamicShapeRemovingVisitor::ConvertToDynamic(
   return TupleUtil::AssembleTupleInstruction(inst->parent(), std::move(gtes));
 }
 
-Status DynamicShapeRemovingVisitor::ConvertOperandsToDynamic(
+absl::Status DynamicShapeRemovingVisitor::ConvertOperandsToDynamic(
     HloInstruction* inst) {
   for (int64_t i = 0; i < inst->operand_count(); ++i) {
     auto operand = inst->mutable_operand(i);
@@ -1961,7 +1961,7 @@ Status DynamicShapeRemovingVisitor::ConvertOperandsToDynamic(
   return OkStatus();
 }
 
-Status DynamicShapeRemovingVisitor::DefaultAction(HloInstruction* hlo) {
+absl::Status DynamicShapeRemovingVisitor::DefaultAction(HloInstruction* hlo) {
   // By default, ops don't support dynamic lowering.
   OpDynamismSupport op_support = OpDynamismSupport::kNoSupport;
   if (op_supports_dynamism_handler_) {
@@ -1993,23 +1993,25 @@ Status DynamicShapeRemovingVisitor::DefaultAction(HloInstruction* hlo) {
   return OkStatus();
 }
 
-Status DynamicShapeRemovingVisitor::HandleGetTupleElement(HloInstruction* hlo) {
+absl::Status DynamicShapeRemovingVisitor::HandleGetTupleElement(
+    HloInstruction* hlo) {
   return OkStatus();
 }
 
-Status DynamicShapeRemovingVisitor::HandleTuple(HloInstruction* hlo) {
+absl::Status DynamicShapeRemovingVisitor::HandleTuple(HloInstruction* hlo) {
   return OkStatus();
 }
 
-Status DynamicShapeRemovingVisitor::HandleInfeed(HloInstruction* hlo) {
+absl::Status DynamicShapeRemovingVisitor::HandleInfeed(HloInstruction* hlo) {
   return OkStatus();
 }
 
-Status DynamicShapeRemovingVisitor::HandleParameter(HloInstruction* hlo) {
+absl::Status DynamicShapeRemovingVisitor::HandleParameter(HloInstruction* hlo) {
   return OkStatus();
 }
 
-Status DynamicShapeRemovingVisitor::HandleCustomCall(HloInstruction* hlo) {
+absl::Status DynamicShapeRemovingVisitor::HandleCustomCall(
+    HloInstruction* hlo) {
   if (hlo->custom_call_target() == "SliceToDynamic" ||
       hlo->custom_call_target() == "PadToStatic") {
     // Those ops support are created to handle dynamic tensors so by their
@@ -2020,7 +2022,8 @@ Status DynamicShapeRemovingVisitor::HandleCustomCall(HloInstruction* hlo) {
   return DefaultAction(hlo);
 }
 
-Status DynamicShapeRemovingVisitor::HandleAsyncStart(HloInstruction* hlo) {
+absl::Status DynamicShapeRemovingVisitor::HandleAsyncStart(
+    HloInstruction* hlo) {
   if (HloInstruction::IsThreadIncluded(hlo->async_execution_thread(),
                                        execution_threads_)) {
     return OkStatus();
@@ -2028,28 +2031,30 @@ Status DynamicShapeRemovingVisitor::HandleAsyncStart(HloInstruction* hlo) {
   return ConvertOperandsToDynamic(hlo);
 }
 
-Status DynamicShapeRemovingVisitor::HandleAsyncUpdate(HloInstruction* hlo) {
+absl::Status DynamicShapeRemovingVisitor::HandleAsyncUpdate(
+    HloInstruction* hlo) {
   return OkStatus();
 }
 
-Status DynamicShapeRemovingVisitor::HandleAsyncDone(HloInstruction* hlo) {
+absl::Status DynamicShapeRemovingVisitor::HandleAsyncDone(HloInstruction* hlo) {
   return OkStatus();
 }
 
-Status DynamicShapeRemovingVisitor::HandleWhile(HloInstruction* hlo) {
+absl::Status DynamicShapeRemovingVisitor::HandleWhile(HloInstruction* hlo) {
   return OkStatus();
 }
 
-Status DynamicShapeRemovingVisitor::HandleConditional(HloInstruction* hlo) {
+absl::Status DynamicShapeRemovingVisitor::HandleConditional(
+    HloInstruction* hlo) {
   return OkStatus();
 }
 
-Status DynamicShapeRemovingVisitor::HandleGetDimensionSize(
+absl::Status DynamicShapeRemovingVisitor::HandleGetDimensionSize(
     HloInstruction* hlo) {
   return OkStatus();
 }
 
-Status DynamicShapeRemovingVisitor::HandleSetDimensionSize(
+absl::Status DynamicShapeRemovingVisitor::HandleSetDimensionSize(
     HloInstruction* hlo) {
   *hlo->mutable_shape() = hlo->operand(0)->shape();
   hlo->mutable_shape()->set_dynamic_dimension(hlo->dimension(), false);
diff --git a/third_party/xla/xla/service/dynamic_padder_test.cc b/third_party/xla/xla/service/dynamic_padder_test.cc
index 20421c57f0116e..e0a70f8bdc4fe6 100644
--- a/third_party/xla/xla/service/dynamic_padder_test.cc
+++ b/third_party/xla/xla/service/dynamic_padder_test.cc
@@ -78,7 +78,7 @@ OpDynamismSupport OpHasDynamismSupport(HloInstruction* hlo) {
   return OpDynamismSupport::kNoSupport;
 }
 
-Status CustomCallDynamicDimensionInference(
+absl::Status CustomCallDynamicDimensionInference(
     HloInstruction* hlo, DynamicDimensionInference* inferencer) {
   if (hlo->custom_call_target() == "OpWithDynamicLowering") {
     if (hlo->shape().IsTuple()) {
diff --git a/third_party/xla/xla/service/eigh_expander.cc b/third_party/xla/xla/service/eigh_expander.cc
index 93d08aa184f40f..84a75bd6b4ffda 100644
--- a/third_party/xla/xla/service/eigh_expander.cc
+++ b/third_party/xla/xla/service/eigh_expander.cc
@@ -236,9 +236,9 @@ void PermuteColumnsInRow(XlaOp& left, XlaOp& right) {
 // implicit way of computing a tournament for n players such that each player
 // plays every other player exactly once in n - 1 rounds. See the Brent/Luk
 // paper for more details.
-Status ApplyRotations(int64_t n, XlaOp& w_tl, XlaOp& w_tr, XlaOp& w_bl,
-                      XlaOp& w_br, XlaOp& v_tl, XlaOp& v_tr, XlaOp& v_bl,
-                      XlaOp& v_br) {
+absl::Status ApplyRotations(int64_t n, XlaOp& w_tl, XlaOp& w_tr, XlaOp& w_bl,
+                            XlaOp& w_br, XlaOp& v_tl, XlaOp& v_tr, XlaOp& v_bl,
+                            XlaOp& v_br) {
   TF_ASSIGN_OR_RETURN(Eigh2x2 rotation,
                       HermitianEigenDecomposition2x2(w_tl, w_tr, w_br));
 
@@ -351,7 +351,7 @@ absl::StatusOr<std::vector<XlaOp>> Sweeps(
 
 }  // namespace
 
-Status EighExpander::SortByEigenvalues(XlaOp& v, XlaOp& w) {
+absl::Status EighExpander::SortByEigenvalues(XlaOp& v, XlaOp& w) {
   XlaBuilder* builder = v.builder();
   TF_ASSIGN_OR_RETURN(Shape v_shape, builder->GetShape(v));
   TF_ASSIGN_OR_RETURN(Shape w_shape, builder->GetShape(w));
diff --git a/third_party/xla/xla/service/eigh_expander.h b/third_party/xla/xla/service/eigh_expander.h
index ed53fac030afd8..074e09305354d9 100644
--- a/third_party/xla/xla/service/eigh_expander.h
+++ b/third_party/xla/xla/service/eigh_expander.h
@@ -35,7 +35,7 @@ class EighExpander : public OpExpanderPass {
   virtual XlaOp BuildEigh(XlaOp a, bool lower, int64_t max_iter, float tol,
                           bool sort_eigenvalues);
 
-  Status SortByEigenvalues(XlaOp& v, XlaOp& w);
+  absl::Status SortByEigenvalues(XlaOp& v, XlaOp& w);
 
  private:
   // Mapping from op signatures to existing computations.
diff --git a/third_party/xla/xla/service/executable.cc b/third_party/xla/xla/service/executable.cc
index 7349471ea24396..d0dc2bf5b2e24f 100644
--- a/third_party/xla/xla/service/executable.cc
+++ b/third_party/xla/xla/service/executable.cc
@@ -41,7 +41,7 @@ ExecutionInput::~ExecutionInput() {
   }
 }
 
-Status ExecutionInput::SetDynamicShape(Shape dynamic_shape) {
+absl::Status ExecutionInput::SetDynamicShape(Shape dynamic_shape) {
   const Shape& input_shape = shape();
   if (!ShapeUtil::DynamicShapeIsCompatible(input_shape, dynamic_shape)) {
     return tsl::errors::InvalidArgument(
@@ -83,7 +83,7 @@ absl::StatusOr<ScopedShapedBuffer> Executable::ExecuteOnStream(
     HloExecutionProfile* hlo_execution_profile) {
   absl::StatusOr<ScopedShapedBuffer> result =
       ExecuteAsyncOnStream(run_options, arguments, hlo_execution_profile);
-  Status blocking_status = run_options->stream()->BlockHostUntilDone();
+  absl::Status blocking_status = run_options->stream()->BlockHostUntilDone();
   TF_RETURN_IF_ERROR(result.status());
   TF_RETURN_IF_ERROR(blocking_status);
   return result;
@@ -120,7 +120,7 @@ absl::StatusOr<ExecutionOutput> Executable::ExecuteOnStream(
     HloExecutionProfile* hlo_execution_profile) {
   absl::StatusOr<ExecutionOutput> result = ExecuteAsyncOnStream(
       run_options, std::move(arguments), hlo_execution_profile);
-  Status blocking_status = run_options->stream()->BlockHostUntilDone();
+  absl::Status blocking_status = run_options->stream()->BlockHostUntilDone();
   TF_RETURN_IF_ERROR(result.status());
   TF_RETURN_IF_ERROR(blocking_status);
   return result;
@@ -163,7 +163,7 @@ absl::StatusOr<ScopedShapedBuffer> Executable::ExecuteOnStreamWrapper(
     absl::Span<const ShapedBuffer* const> arguments) {
   absl::StatusOr<ScopedShapedBuffer> result =
       ExecuteAsyncOnStreamWrapper(run_options, arguments);
-  Status block_status = run_options->stream()->BlockHostUntilDone();
+  absl::Status block_status = run_options->stream()->BlockHostUntilDone();
   TF_RETURN_IF_ERROR(result.status());
   TF_RETURN_IF_ERROR(block_status);
   return result;
@@ -174,7 +174,7 @@ absl::StatusOr<ExecutionOutput> Executable::ExecuteOnStreamWrapper(
     std::vector<ExecutionInput> arguments) {
   absl::StatusOr<ExecutionOutput> result =
       ExecuteAsyncOnStreamWrapper(run_options, std::move(arguments));
-  Status block_status = run_options->stream()->BlockHostUntilDone();
+  absl::Status block_status = run_options->stream()->BlockHostUntilDone();
   TF_RETURN_IF_ERROR(result.status());
   TF_RETURN_IF_ERROR(block_status);
   return result;
@@ -194,12 +194,12 @@ static ExecuteAsyncOnStreamWrapperState ExecuteWrapperBeforeExecution(
   return state;
 }
 
-Status ExecuteWrapperAfterExecution(
+absl::Status ExecuteWrapperAfterExecution(
     Executable* executable, const ExecuteAsyncOnStreamWrapperState& state,
-    Status return_status, se::Stream* stream) {
+    absl::Status return_status, se::Stream* stream) {
   if (!return_status.ok()) {
     if (state.profile != nullptr) {
-      Status status = stream->BlockHostUntilDone();
+      absl::Status status = stream->BlockHostUntilDone();
       if (!status.ok()) {
         LOG(ERROR) << "Failed to BlockHostUntilDone: " << status;
       }
diff --git a/third_party/xla/xla/service/executable.h b/third_party/xla/xla/service/executable.h
index 5c26581b9479e3..e7d9cd29244055 100644
--- a/third_party/xla/xla/service/executable.h
+++ b/third_party/xla/xla/service/executable.h
@@ -96,7 +96,7 @@ class ExecutionInput {
     return host_shape_ != nullptr ? *host_shape_ : shape();
   }
 
-  Status SetDynamicShape(Shape dynamic_shape);
+  absl::Status SetDynamicShape(Shape dynamic_shape);
 
   absl::StatusOr<xla::ShapedBuffer> ToShapedBuffer(
       se::DeviceMemoryAllocator* allocator, int device_ordinal) const;
diff --git a/third_party/xla/xla/service/execution_tracker.cc b/third_party/xla/xla/service/execution_tracker.cc
index 655aca424f1906..9bef365644d886 100644
--- a/third_party/xla/xla/service/execution_tracker.cc
+++ b/third_party/xla/xla/service/execution_tracker.cc
@@ -37,7 +37,7 @@ AsyncExecution::AsyncExecution(Backend* backend,
   }
 }
 
-Status AsyncExecution::BlockUntilDone() const {
+absl::Status AsyncExecution::BlockUntilDone() const {
   for (auto& stream : streams_) {
     TF_RETURN_IF_ERROR(stream->BlockHostUntilDone());
   }
@@ -62,7 +62,7 @@ ExecutionHandle ExecutionTracker::Register(Backend* backend,
   return execution_handle;
 }
 
-Status ExecutionTracker::Unregister(const ExecutionHandle& handle) {
+absl::Status ExecutionTracker::Unregister(const ExecutionHandle& handle) {
   absl::MutexLock lock(&execution_mutex_);
   auto it = handle_to_execution_.find(handle.handle());
   if (it == handle_to_execution_.end()) {
diff --git a/third_party/xla/xla/service/execution_tracker.h b/third_party/xla/xla/service/execution_tracker.h
index 3bd5f32edced51..06c9279299ef7d 100644
--- a/third_party/xla/xla/service/execution_tracker.h
+++ b/third_party/xla/xla/service/execution_tracker.h
@@ -39,7 +39,7 @@ class AsyncExecution {
   AsyncExecution(Backend* backend, std::vector<StreamPool::Ptr> streams,
                  const ExecutionProfile& profile, GlobalDataHandle result);
 
-  Status BlockUntilDone() const;
+  absl::Status BlockUntilDone() const;
 
   const GlobalDataHandle& result() const { return result_; }
 
@@ -73,7 +73,7 @@ class ExecutionTracker {
                            GlobalDataHandle data);
 
   // Unregisters the execution for the given handle.
-  Status Unregister(const ExecutionHandle& handle);
+  absl::Status Unregister(const ExecutionHandle& handle);
 
   // Resolves the given ExecutionHandle to an AsyncExecution. Returns an
   // error status if the given handle is not found, which means that the
diff --git a/third_party/xla/xla/service/export_hlo.h b/third_party/xla/xla/service/export_hlo.h
index da5cfe89c927f0..e0e0456cd506e6 100644
--- a/third_party/xla/xla/service/export_hlo.h
+++ b/third_party/xla/xla/service/export_hlo.h
@@ -39,8 +39,8 @@ class SymbolUploader {
   virtual ~SymbolUploader() = default;
 
   // Returns a string identifying the uploaded HLO, or empty if the upload did
-  // not complete. We use optional rather than StatusOr because an upload error
-  // is not a compiler error.
+  // not complete. We use optional rather than absl::StatusOr because an upload
+  // error is not a compiler error.
   virtual std::optional<std::string> MaybeUploadUnoptimizedHloModule(
       HloModule* module,
       const stream_executor::GpuTargetConfigProto& gpu_target_config) = 0;
diff --git a/third_party/xla/xla/service/flatten_call_graph.cc b/third_party/xla/xla/service/flatten_call_graph.cc
index 191aca4b46e13c..39a3ca7054137d 100644
--- a/third_party/xla/xla/service/flatten_call_graph.cc
+++ b/third_party/xla/xla/service/flatten_call_graph.cc
@@ -68,7 +68,7 @@ void ReplaceCalledComputation(HloInstruction* instruction,
 }
 
 // Flatten a single call graph node. Expects to visit nodes in postorder.
-Status FlattenNode(const CallGraphNode& node) {
+absl::Status FlattenNode(const CallGraphNode& node) {
   HloComputation* computation = node.computation();
   HloModule* module = computation->parent();
   // Clone callee for all call-sites except the first one.
diff --git a/third_party/xla/xla/service/float_normalization.cc b/third_party/xla/xla/service/float_normalization.cc
index 97fbc5bbc26888..7110c7b866decd 100644
--- a/third_party/xla/xla/service/float_normalization.cc
+++ b/third_party/xla/xla/service/float_normalization.cc
@@ -43,18 +43,18 @@ class FloatNormalizationVisitor : public DfsHloVisitorWithDefault {
         float_normalization_(float_normalization) {}
 
   bool changed() const { return changed_; }
-  Status DefaultAction(HloInstruction* hlo) override;
-  Status Preprocess(HloInstruction* hlo) override;
+  absl::Status DefaultAction(HloInstruction* hlo) override;
+  absl::Status Preprocess(HloInstruction* hlo) override;
 
  private:
   // Checks if the HLO uses low-precision in an unsupported way, and if so,
   // inserts conversions between the low- and high-precision types to make it
   // supported.
-  Status HandleInstruction(HloInstruction* hlo);
+  absl::Status HandleInstruction(HloInstruction* hlo);
 
   // Handle instructions with tuple outputs by examining each output
   // independently.
-  Status HandleMultipleOutputs(HloInstruction* hlo);
+  absl::Status HandleMultipleOutputs(HloInstruction* hlo);
 
   // Creates a copy of `hlo` with subshapes matching `from` type converted to
   // `to` type. If no matching subshapes are found, returns the original `hlo`.
@@ -65,27 +65,27 @@ class FloatNormalizationVisitor : public DfsHloVisitorWithDefault {
 
   // Inserts a conversion HLO that changes the given HLO's output type. If the
   // output is a tuple, change all elements that match the from type.
-  Status InsertConvertAfterOutput(HloInstruction* hlo, PrimitiveType from,
-                                  PrimitiveType to,
-                                  HloComputation* computation);
+  absl::Status InsertConvertAfterOutput(HloInstruction* hlo, PrimitiveType from,
+                                        PrimitiveType to,
+                                        HloComputation* computation);
 
   // Changes the output type to the specified type, then inserts a conversion
   // to the original type. If the output is a tuple, change all elements that
   // match the from type.
-  Status ChangeOutputTypeThenInsertConvertBack(HloInstruction* hlo,
-                                               PrimitiveType from,
-                                               PrimitiveType to,
-                                               HloComputation* computation);
+  absl::Status ChangeOutputTypeThenInsertConvertBack(
+      HloInstruction* hlo, PrimitiveType from, PrimitiveType to,
+      HloComputation* computation);
 
   // Inserts a conversion HLO that changes the given HLO's operand type. If the
   // operand is a tuple, change all elements that match the from type.
-  Status InsertConvertBeforeOperand(HloInstruction* hlo, int64_t operand_idx,
-                                    PrimitiveType from, PrimitiveType to,
-                                    HloComputation* computation);
+  absl::Status InsertConvertBeforeOperand(HloInstruction* hlo,
+                                          int64_t operand_idx,
+                                          PrimitiveType from, PrimitiveType to,
+                                          HloComputation* computation);
 
   // Inserts conversion HLOs to replace the called computations' low-precision
   // operands/outputs to high-precision.
-  Status ConvertCalledComputations(
+  absl::Status ConvertCalledComputations(
       HloInstruction* hlo,
       absl::Span<HloComputation* const> low_precision_called_comps);
 
@@ -156,7 +156,7 @@ absl::StatusOr<HloInstruction*> FloatNormalizationVisitor::ConvertType(
   return new_hlo;
 }
 
-Status FloatNormalizationVisitor::InsertConvertAfterOutput(
+absl::Status FloatNormalizationVisitor::InsertConvertAfterOutput(
     HloInstruction* hlo, PrimitiveType from, PrimitiveType to,
     HloComputation* computation) {
   bool is_root = computation->root_instruction() == hlo;
@@ -177,7 +177,7 @@ Status FloatNormalizationVisitor::InsertConvertAfterOutput(
   return OkStatus();
 }
 
-Status FloatNormalizationVisitor::ChangeOutputTypeThenInsertConvertBack(
+absl::Status FloatNormalizationVisitor::ChangeOutputTypeThenInsertConvertBack(
     HloInstruction* hlo, PrimitiveType from, PrimitiveType to,
     HloComputation* computation) {
   auto original_shape = hlo->shape();
@@ -246,7 +246,7 @@ Status FloatNormalizationVisitor::ChangeOutputTypeThenInsertConvertBack(
   return OkStatus();
 }
 
-Status FloatNormalizationVisitor::InsertConvertBeforeOperand(
+absl::Status FloatNormalizationVisitor::InsertConvertBeforeOperand(
     HloInstruction* hlo, int64_t operand_idx, PrimitiveType from,
     PrimitiveType to, HloComputation* computation) {
   auto operand = hlo->mutable_operand(operand_idx);
@@ -261,7 +261,7 @@ Status FloatNormalizationVisitor::InsertConvertBeforeOperand(
   return OkStatus();
 }
 
-Status FloatNormalizationVisitor::ConvertCalledComputations(
+absl::Status FloatNormalizationVisitor::ConvertCalledComputations(
     HloInstruction* hlo,
     absl::Span<HloComputation* const> low_precision_called_comps) {
   absl::flat_hash_map<HloComputation*, HloComputation*> cloned_computations;
@@ -300,7 +300,8 @@ bool ShouldAvoidNormalizingComputationsForInstruction(HloInstruction* hlo) {
          hlo->opcode() == HloOpcode::kReduceScatter;
 }
 
-Status FloatNormalizationVisitor::HandleMultipleOutputs(HloInstruction* hlo) {
+absl::Status FloatNormalizationVisitor::HandleMultipleOutputs(
+    HloInstruction* hlo) {
   std::vector<PrimitiveType> operand_types(hlo->operand_count());
   std::vector<PrimitiveType> output_types(hlo->operand_count());
   int64_t high_prec_count = 0;
@@ -426,7 +427,7 @@ Status FloatNormalizationVisitor::HandleMultipleOutputs(HloInstruction* hlo) {
   return ConvertCalledComputations(hlo, low_precision_called_comps);
 }
 
-Status FloatNormalizationVisitor::HandleInstruction(HloInstruction* hlo) {
+absl::Status FloatNormalizationVisitor::HandleInstruction(HloInstruction* hlo) {
   int high_prec_count = 0;
   int low_prec_count = 0;
 
@@ -540,7 +541,7 @@ Status FloatNormalizationVisitor::HandleInstruction(HloInstruction* hlo) {
   return ConvertCalledComputations(hlo, low_precision_called_comps);
 }
 
-Status FloatNormalizationVisitor::DefaultAction(HloInstruction* hlo) {
+absl::Status FloatNormalizationVisitor::DefaultAction(HloInstruction* hlo) {
   // Do not change instructions related to entry and exit of a computation,
   // tuples, fusion, convert, side-effecting instructions, control flow, and
   // bitcast-convert.
@@ -569,7 +570,7 @@ Status FloatNormalizationVisitor::DefaultAction(HloInstruction* hlo) {
   return HandleInstruction(hlo);
 }
 
-Status FloatNormalizationVisitor::Preprocess(HloInstruction* hlo) {
+absl::Status FloatNormalizationVisitor::Preprocess(HloInstruction* hlo) {
   computation_ = hlo->parent();
   return OkStatus();
 }
diff --git a/third_party/xla/xla/service/gather_expander_test.cc b/third_party/xla/xla/service/gather_expander_test.cc
index 2a215559e281d3..271fabdbda2c62 100644
--- a/third_party/xla/xla/service/gather_expander_test.cc
+++ b/third_party/xla/xla/service/gather_expander_test.cc
@@ -43,9 +43,9 @@ ENTRY main {
   TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
                           ParseAndReturnVerifiedModule(hlo_text));
 
-  Status status = GatherExpander{GatherExpander::kEliminateAllGathers}
-                      .Run(module.get())
-                      .status();
+  absl::Status status = GatherExpander{GatherExpander::kEliminateAllGathers}
+                            .Run(module.get())
+                            .status();
   EXPECT_EQ(status.code(), tsl::error::UNIMPLEMENTED);
 
   ASSERT_THAT(
diff --git a/third_party/xla/xla/service/generic_transfer_manager.cc b/third_party/xla/xla/service/generic_transfer_manager.cc
index f285c1be04f045..28a0e4e7f47dc7 100644
--- a/third_party/xla/xla/service/generic_transfer_manager.cc
+++ b/third_party/xla/xla/service/generic_transfer_manager.cc
@@ -53,7 +53,7 @@ se::Platform::Id GenericTransferManager::PlatformId() const {
   return platform_id_;
 }
 
-Status GenericTransferManager::WriteSingleTupleIndexTable(
+absl::Status GenericTransferManager::WriteSingleTupleIndexTable(
     se::Stream* stream, absl::Span<const se::DeviceMemoryBase> elements,
     const Shape& shape, se::DeviceMemoryBase* region) {
   TF_RET_CHECK(elements.size() == ShapeUtil::TupleElementCount(shape));
@@ -81,13 +81,13 @@ void GenericTransferManager::TransferLiteralFromDevice(
           << stream->parent()->device_ordinal()
           << "; device buffer: " << device_buffer;
 
-  Status status = [&]() -> Status {
+  absl::Status status = [&]() -> absl::Status {
     TF_RET_CHECK(stream->parent()->device_ordinal() ==
                  device_buffer.device_ordinal());
 
     TF_RETURN_IF_ERROR(ShapeUtil::ForEachSubshapeWithStatus(
         device_buffer.on_device_shape(),
-        [&](const Shape& subshape, const ShapeIndex& index) -> Status {
+        [&](const Shape& subshape, const ShapeIndex& index) -> absl::Status {
           if (subshape.IsArray()) {
             if (PackSubbyteTypes() &&
                 primitive_util::IsSubByteNonPredType(subshape.element_type())) {
@@ -143,7 +143,7 @@ void GenericTransferManager::TransferLiteralFromDevice(
   }
 }
 
-Status GenericTransferManager::TransferLiteralToDeviceAsync(
+absl::Status GenericTransferManager::TransferLiteralToDeviceAsync(
     se::Stream* stream, const LiteralSlice& literal,
     const ShapedBuffer& device_buffer,
     const TransferMetadata* /*transfer_metadata*/) {
@@ -161,7 +161,8 @@ Status GenericTransferManager::TransferLiteralToDeviceAsync(
 
   return ShapeUtil::ForEachSubshapeWithStatus(
       device_buffer.on_device_shape(),
-      [&](const Shape& device_subshape, const ShapeIndex& index) -> Status {
+      [&](const Shape& device_subshape,
+          const ShapeIndex& index) -> absl::Status {
         if (device_subshape.IsArray()) {
           int64_t size = GetByteSizeRequirement(device_subshape);
           se::DeviceMemoryBase device_memory = device_buffer.buffer(index);
@@ -205,24 +206,24 @@ Status GenericTransferManager::TransferLiteralToDeviceAsync(
       });
 }
 
-Status GenericTransferManager::TransferLiteralToInfeed(
+absl::Status GenericTransferManager::TransferLiteralToInfeed(
     se::StreamExecutor* executor, const LiteralSlice& literal) {
   return Unimplemented("Generic transfer to Infeed");
 }
 
-Status GenericTransferManager::TransferLiteralFromOutfeed(
+absl::Status GenericTransferManager::TransferLiteralFromOutfeed(
     se::StreamExecutor* executor, MutableBorrowingLiteral literal) {
   return Unimplemented("Generic transfer from Outfeed");
 }
 
-Status GenericTransferManager::ResetDevices(
+absl::Status GenericTransferManager::ResetDevices(
     absl::Span<se::StreamExecutor* const>
     /*executors*/) {
   return Unimplemented(
       "Device reset is not yet supported on this platform (b/30481585)");
 }
 
-Status GenericTransferManager::TransferBufferFromDevice(
+absl::Status GenericTransferManager::TransferBufferFromDevice(
     se::Stream* stream, const se::DeviceMemoryBase& source, int64_t size,
     void* destination) {
   if (source.size() < size) {
@@ -234,7 +235,7 @@ Status GenericTransferManager::TransferBufferFromDevice(
   return stream->Memcpy(destination, source, size);
 }
 
-Status GenericTransferManager::TransferBufferToDevice(
+absl::Status GenericTransferManager::TransferBufferToDevice(
     se::Stream* stream, int64_t size, const void* source,
     se::DeviceMemoryBase* destination) {
   if (destination->size() < size) {
@@ -246,7 +247,7 @@ Status GenericTransferManager::TransferBufferToDevice(
   return stream->Memcpy(destination, source, size);
 }
 
-Status GenericTransferManager::TransferIntNArrayFromDevice(
+absl::Status GenericTransferManager::TransferIntNArrayFromDevice(
     se::Stream* stream, const se::DeviceMemoryBase& source,
     PrimitiveType element_type, int64_t num_elements, void* destination) {
   int bit_width = primitive_util::BitWidth(element_type);
@@ -265,7 +266,7 @@ Status GenericTransferManager::TransferIntNArrayFromDevice(
   return OkStatus();
 }
 
-Status GenericTransferManager::TransferIntNArrayToDevice(
+absl::Status GenericTransferManager::TransferIntNArrayToDevice(
     se::Stream* stream, PrimitiveType element_type, int64_t num_elements,
     const void* source, se::DeviceMemoryBase* destination) {
   int bit_width = primitive_util::BitWidth(element_type);
diff --git a/third_party/xla/xla/service/generic_transfer_manager.h b/third_party/xla/xla/service/generic_transfer_manager.h
index 7b212aa63aee98..a00f46d4f6fc45 100644
--- a/third_party/xla/xla/service/generic_transfer_manager.h
+++ b/third_party/xla/xla/service/generic_transfer_manager.h
@@ -61,22 +61,23 @@ class GenericTransferManager : public TransferManager {
       MutableBorrowingLiteral literal, std::function<void(Status)> done,
       const TransferMetadata* transfer_metadata) override;
 
-  Status TransferLiteralToDeviceAsync(
+  absl::Status TransferLiteralToDeviceAsync(
       se::Stream* stream, const LiteralSlice& literal,
       const ShapedBuffer& device_buffer,
       const TransferMetadata* transfer_metadata) override;
 
-  Status TransferLiteralToInfeed(se::StreamExecutor* executor,
-                                 const LiteralSlice& literal) override;
+  absl::Status TransferLiteralToInfeed(se::StreamExecutor* executor,
+                                       const LiteralSlice& literal) override;
 
-  Status TransferLiteralFromOutfeed(se::StreamExecutor* executor,
-                                    MutableBorrowingLiteral literal) override;
+  absl::Status TransferLiteralFromOutfeed(
+      se::StreamExecutor* executor, MutableBorrowingLiteral literal) override;
 
-  Status ResetDevices(absl::Span<se::StreamExecutor* const> executors) override;
+  absl::Status ResetDevices(
+      absl::Span<se::StreamExecutor* const> executors) override;
 
   int64_t GetByteSizeRequirement(const Shape& shape) const override;
 
-  Status WriteSingleTupleIndexTable(
+  absl::Status WriteSingleTupleIndexTable(
       se::Stream* stream, absl::Span<const se::DeviceMemoryBase> elements,
       const Shape& shape, se::DeviceMemoryBase* region) override;
 
@@ -87,39 +88,35 @@ class GenericTransferManager : public TransferManager {
   // 'destination' buffer.
   //
   // size is the size to transfer to destination in bytes.
-  virtual Status TransferBufferFromDevice(se::Stream* stream,
-                                          const se::DeviceMemoryBase& source,
-                                          int64_t size, void* destination);
+  virtual absl::Status TransferBufferFromDevice(
+      se::Stream* stream, const se::DeviceMemoryBase& source, int64_t size,
+      void* destination);
 
   // Transfer a memory block of the given size from 'source' buffer to the given
   // destination of the device.
   //
   // size is the size to transfer from source in bytes.
-  virtual Status TransferBufferToDevice(se::Stream* stream, int64_t size,
-                                        const void* source,
-                                        se::DeviceMemoryBase* destination);
+  virtual absl::Status TransferBufferToDevice(
+      se::Stream* stream, int64_t size, const void* source,
+      se::DeviceMemoryBase* destination);
 
   // Transfers a buffer of packed int4 values from the device to the host, then
   // unpacks them on the host. 'source' is a buffer with (num_elements+1)/2
   // bytes where each byte stores two int4 values. 'destination' is a buffer
   // with num_elements bytes, where a single int4 value will be written to each
   // byte in the lower 4 bits.
-  virtual Status TransferIntNArrayFromDevice(se::Stream* stream,
-                                             const se::DeviceMemoryBase& source,
-                                             PrimitiveType element_type,
-                                             int64_t num_elements,
-                                             void* destination);
+  virtual absl::Status TransferIntNArrayFromDevice(
+      se::Stream* stream, const se::DeviceMemoryBase& source,
+      PrimitiveType element_type, int64_t num_elements, void* destination);
 
   // Packs an array of int4 values then transfers the packed buffer from the
   // host to the device. 'source' is a buffer with num_elements bytes, where the
   // lower 4 bits of each byte stores an int4 value. 'destination' is a buffer
   // with (num_elements+1)/2 bytes, where two int4 values will be written into
   // each byte.
-  virtual Status TransferIntNArrayToDevice(se::Stream* stream,
-                                           PrimitiveType element_type,
-                                           int64_t num_elements,
-                                           const void* source,
-                                           se::DeviceMemoryBase* destination);
+  virtual absl::Status TransferIntNArrayToDevice(
+      se::Stream* stream, PrimitiveType element_type, int64_t num_elements,
+      const void* source, se::DeviceMemoryBase* destination);
 
   // The platform this transfer manager targets.
   const se::Platform::Id platform_id_;
diff --git a/third_party/xla/xla/service/gpu_compilation_environment.cc b/third_party/xla/xla/service/gpu_compilation_environment.cc
index d598c02df3d5f6..4376c975544c0c 100644
--- a/third_party/xla/xla/service/gpu_compilation_environment.cc
+++ b/third_party/xla/xla/service/gpu_compilation_environment.cc
@@ -79,7 +79,8 @@ GpuCompilationEnvironment CreateGpuCompEnvWithDefaultValues() {
   return env;
 }
 
-Status InitializeMissingFieldsFromXLAFlags(GpuCompilationEnvironment& env) {
+absl::Status InitializeMissingFieldsFromXLAFlags(
+    GpuCompilationEnvironment& env) {
   TF_ASSIGN_OR_RETURN(GpuCompilationEnvironment from_env,
                       CreateGpuCompEnvFromEnvVar());
 
diff --git a/third_party/xla/xla/service/gpu_compilation_environment.h b/third_party/xla/xla/service/gpu_compilation_environment.h
index 23f2a30273c8a9..a5ca8f67c2a87e 100644
--- a/third_party/xla/xla/service/gpu_compilation_environment.h
+++ b/third_party/xla/xla/service/gpu_compilation_environment.h
@@ -32,7 +32,8 @@ GpuCompilationEnvironment CreateGpuCompEnvWithDefaultValues();
 
 // Returns non-OK status if XLA_FLAGS env var has malformed values or
 // if it has conflict with the GpuCompilationEnvironment proto
-Status InitializeMissingFieldsFromXLAFlags(GpuCompilationEnvironment& env);
+absl::Status InitializeMissingFieldsFromXLAFlags(
+    GpuCompilationEnvironment& env);
 
 }  // namespace xla
 #endif  // XLA_SERVICE_GPU_COMPILATION_ENVIRONMENT_H_
diff --git a/third_party/xla/xla/service/hlo_alias_analysis.cc b/third_party/xla/xla/service/hlo_alias_analysis.cc
index 444255d9c344b2..3d6f5a87e1a760 100644
--- a/third_party/xla/xla/service/hlo_alias_analysis.cc
+++ b/third_party/xla/xla/service/hlo_alias_analysis.cc
@@ -316,7 +316,7 @@ std::vector<const HloBuffer*> HloAliasAnalysis::ComputeBuffersAt(
   return buffers;
 }
 
-Status HloAliasAnalysis::Verify() const {
+absl::Status HloAliasAnalysis::Verify() const {
   // Verify consistency between the value_to_buffer_ map and
   // HloBuffer::values().
   for (const auto& pair : value_to_buffer_) {
diff --git a/third_party/xla/xla/service/hlo_alias_analysis.h b/third_party/xla/xla/service/hlo_alias_analysis.h
index b0896a93c6883e..197a4d60109ec8 100644
--- a/third_party/xla/xla/service/hlo_alias_analysis.h
+++ b/third_party/xla/xla/service/hlo_alias_analysis.h
@@ -103,7 +103,7 @@ class HloAliasAnalysis {
   explicit HloAliasAnalysis(const HloModule* module);
 
   // Verify various invariants of the alias analysis.
-  Status Verify() const;
+  absl::Status Verify() const;
 
   const HloModule* module_;
 
diff --git a/third_party/xla/xla/service/hlo_computation_test.cc b/third_party/xla/xla/service/hlo_computation_test.cc
index 0cc83289383a4e..37d2c2c2e5c345 100644
--- a/third_party/xla/xla/service/hlo_computation_test.cc
+++ b/third_party/xla/xla/service/hlo_computation_test.cc
@@ -253,14 +253,14 @@ TEST_F(HloComputationTest, VisitWithMultipleRoots) {
     explicit TestVisitor(HloComputation* computation)
         : computation_(computation) {}
 
-    Status DefaultAction(HloInstruction* hlo_instruction) override {
+    absl::Status DefaultAction(HloInstruction* hlo_instruction) override {
       EXPECT_FALSE(visited_set_.contains(hlo_instruction));
       visited_set_.insert(hlo_instruction);
       last_visited_ = hlo_instruction;
       return OkStatus();
     }
 
-    Status FinishVisit(HloInstruction* root) override {
+    absl::Status FinishVisit(HloInstruction* root) override {
       EXPECT_EQ(computation_->root_instruction(), root);
       ++finish_visit_calls_;
       return OkStatus();
diff --git a/third_party/xla/xla/service/hlo_cost_analysis.cc b/third_party/xla/xla/service/hlo_cost_analysis.cc
index c60a4193aaee20..74b584851d815b 100644
--- a/third_party/xla/xla/service/hlo_cost_analysis.cc
+++ b/third_party/xla/xla/service/hlo_cost_analysis.cc
@@ -47,7 +47,7 @@ HloCostAnalysis::HloCostAnalysis(ShapeSizeFunction shape_size,
                                  const Properties& per_second_rates)
     : HloCostAnalysis(Options{shape_size, per_second_rates}) {}
 
-Status HloCostAnalysis::Preprocess(const HloInstruction* hlo) {
+absl::Status HloCostAnalysis::Preprocess(const HloInstruction* hlo) {
   // Set current instruction cost values to reasonable default values. Each
   // handler can overwrite these values. In Postprocess, these values are
   // accumulated and written to the per-instruction maps.
@@ -71,7 +71,7 @@ Status HloCostAnalysis::Preprocess(const HloInstruction* hlo) {
   return OkStatus();
 }
 
-Status HloCostAnalysis::Postprocess(const HloInstruction* hlo) {
+absl::Status HloCostAnalysis::Postprocess(const HloInstruction* hlo) {
   if (current_should_compute_bottleneck_time_) {
     // Compute the time as the time of the bottleneck, i.e. the slowest property
     // given the per-second rate of each property.
@@ -101,7 +101,7 @@ Status HloCostAnalysis::Postprocess(const HloInstruction* hlo) {
   return OkStatus();
 }
 
-Status HloCostAnalysis::RemoveInstruction(HloInstruction* instruction) {
+absl::Status HloCostAnalysis::RemoveInstruction(HloInstruction* instruction) {
   // Subtract the previously calculated properties of the instruction
   // from HLO graph's total properties_sum_ if instruction was analyzed before.
   auto it = hlo_properties_.find(instruction);
@@ -114,7 +114,7 @@ Status HloCostAnalysis::RemoveInstruction(HloInstruction* instruction) {
   return OkStatus();
 }
 
-Status HloCostAnalysis::RevisitInstruction(HloInstruction* instruction) {
+absl::Status HloCostAnalysis::RevisitInstruction(HloInstruction* instruction) {
   TF_RETURN_IF_ERROR(RemoveInstruction(instruction));
   // Now do Preprocess() -> Visit() -> Postprocess() for the instruction same
   // way it is done during the complete analysis.
@@ -124,7 +124,7 @@ Status HloCostAnalysis::RevisitInstruction(HloInstruction* instruction) {
   return OkStatus();
 }
 
-Status HloCostAnalysis::HandleElementwiseOp(
+absl::Status HloCostAnalysis::HandleElementwiseOp(
     const HloInstruction* hlo_instruction) {
   const auto& shape = hlo_instruction->shape();
   // For element-wise operations, the number of computations is the same as the
@@ -241,7 +241,7 @@ int64_t HloCostAnalysis::FusionParameterReadBytes(
   return size;
 }
 
-Status HloCostAnalysis::FusionCalculateUtilizations(
+absl::Status HloCostAnalysis::FusionCalculateUtilizations(
     const HloInstruction* fusion) {
   // Default trivial implementation: assume 100% utilization of every fusion
   // instruction.
@@ -256,27 +256,29 @@ Status HloCostAnalysis::FusionCalculateUtilizations(
   return OkStatus();
 }
 
-Status HloCostAnalysis::HandleElementwiseUnary(const HloInstruction* hlo) {
+absl::Status HloCostAnalysis::HandleElementwiseUnary(
+    const HloInstruction* hlo) {
   return HandleElementwiseOp(hlo);
 }
 
-Status HloCostAnalysis::HandleElementwiseBinary(const HloInstruction* hlo) {
+absl::Status HloCostAnalysis::HandleElementwiseBinary(
+    const HloInstruction* hlo) {
   return HandleElementwiseOp(hlo);
 }
 
-Status HloCostAnalysis::HandleCompare(const HloInstruction* compare) {
+absl::Status HloCostAnalysis::HandleCompare(const HloInstruction* compare) {
   return HandleElementwiseOp(compare);
 }
 
-Status HloCostAnalysis::HandleClamp(const HloInstruction* clamp) {
+absl::Status HloCostAnalysis::HandleClamp(const HloInstruction* clamp) {
   return HandleElementwiseOp(clamp);
 }
 
-Status HloCostAnalysis::HandleReducePrecision(const HloInstruction* hlo) {
+absl::Status HloCostAnalysis::HandleReducePrecision(const HloInstruction* hlo) {
   return HandleElementwiseOp(hlo);
 }
 
-Status HloCostAnalysis::HandleParameter(const HloInstruction*) {
+absl::Status HloCostAnalysis::HandleParameter(const HloInstruction*) {
   current_should_compute_bottleneck_time_ = false;
   current_properties_[kBytesAccessedKey] = 0;
   current_properties_.set_output_bytes_accessed(0);
@@ -284,7 +286,7 @@ Status HloCostAnalysis::HandleParameter(const HloInstruction*) {
   return OkStatus();
 }
 
-Status HloCostAnalysis::HandleConstant(const HloInstruction*) {
+absl::Status HloCostAnalysis::HandleConstant(const HloInstruction*) {
   current_should_compute_bottleneck_time_ = false;
   current_properties_[kBytesAccessedKey] = 0;
   current_properties_.set_output_bytes_accessed(0);
@@ -292,9 +294,11 @@ Status HloCostAnalysis::HandleConstant(const HloInstruction*) {
   return OkStatus();
 }
 
-Status HloCostAnalysis::HandleIota(const HloInstruction*) { return OkStatus(); }
+absl::Status HloCostAnalysis::HandleIota(const HloInstruction*) {
+  return OkStatus();
+}
 
-Status HloCostAnalysis::HandleGetTupleElement(
+absl::Status HloCostAnalysis::HandleGetTupleElement(
     const HloInstruction* get_tuple_element) {
   // GetTupleElement forwards a pointer and does not touch each element in the
   // output.
@@ -306,15 +310,15 @@ Status HloCostAnalysis::HandleGetTupleElement(
   return OkStatus();
 }
 
-Status HloCostAnalysis::HandleSelect(const HloInstruction* hlo) {
+absl::Status HloCostAnalysis::HandleSelect(const HloInstruction* hlo) {
   return HandleElementwiseOp(hlo);
 }
 
-Status HloCostAnalysis::HandleReverse(const HloInstruction*) {
+absl::Status HloCostAnalysis::HandleReverse(const HloInstruction*) {
   return OkStatus();
 }
 
-Status HloCostAnalysis::HandleSlice(const HloInstruction* slice) {
+absl::Status HloCostAnalysis::HandleSlice(const HloInstruction* slice) {
   const int64_t output_shape_size = GetShapeSize(slice->shape());
 
   const int64_t num_input_elements =
@@ -329,7 +333,7 @@ Status HloCostAnalysis::HandleSlice(const HloInstruction* slice) {
   return OkStatus();
 }
 
-Status HloCostAnalysis::HandleDynamicSlice(
+absl::Status HloCostAnalysis::HandleDynamicSlice(
     const HloInstruction* dynamic_slice) {
   const int64_t output_shape_size = GetShapeSize(dynamic_slice->shape());
   const int64_t start_indices_shape_size =
@@ -350,7 +354,7 @@ Status HloCostAnalysis::HandleDynamicSlice(
   return OkStatus();
 }
 
-Status HloCostAnalysis::HandleDynamicUpdateSlice(
+absl::Status HloCostAnalysis::HandleDynamicUpdateSlice(
     const HloInstruction* dynamic_update_slice) {
   const int64_t update_shape_size =
       GetShapeSize(dynamic_update_slice->operand(1)->shape());
@@ -377,7 +381,7 @@ Status HloCostAnalysis::HandleDynamicUpdateSlice(
   return OkStatus();
 }
 
-Status HloCostAnalysis::HandleTuple(const HloInstruction* tuple) {
+absl::Status HloCostAnalysis::HandleTuple(const HloInstruction* tuple) {
   // The tuple instruction only gathers pointers from inputs (it doesn't iterate
   // through them). The memory touched is then only the size of the output
   // index table of the tuple.
@@ -390,17 +394,19 @@ Status HloCostAnalysis::HandleTuple(const HloInstruction* tuple) {
   return OkStatus();
 }
 
-Status HloCostAnalysis::HandleConcatenate(const HloInstruction*) {
+absl::Status HloCostAnalysis::HandleConcatenate(const HloInstruction*) {
   return OkStatus();
 }
 
-Status HloCostAnalysis::HandleConvert(const HloInstruction* convert) {
+absl::Status HloCostAnalysis::HandleConvert(const HloInstruction* convert) {
   return HandleElementwiseOp(convert);
 }
 
-Status HloCostAnalysis::HandleCopy(const HloInstruction*) { return OkStatus(); }
+absl::Status HloCostAnalysis::HandleCopy(const HloInstruction*) {
+  return OkStatus();
+}
 
-Status HloCostAnalysis::HandleDomain(const HloInstruction* domain) {
+absl::Status HloCostAnalysis::HandleDomain(const HloInstruction* domain) {
   // Domain does not have any computation or data transfer.
   current_should_compute_bottleneck_time_ = false;
   current_properties_[kBytesAccessedKey] = 0;
@@ -425,13 +431,13 @@ int64_t HloCostAnalysis::GetDotFlops(const Shape& lhs_shape,
   return kFmaFlops * ShapeUtil::ElementsIn(result_shape) * reduction_width;
 }
 
-Status HloCostAnalysis::HandleDot(const HloInstruction* dot) {
+absl::Status HloCostAnalysis::HandleDot(const HloInstruction* dot) {
   current_properties_[kFlopsKey] = GetDotFlops(
       dot->operand(0)->shape(), dot->shape(), dot->dot_dimension_numbers());
   return OkStatus();
 }
 
-Status HloCostAnalysis::HandleInfeed(const HloInstruction* infeed) {
+absl::Status HloCostAnalysis::HandleInfeed(const HloInstruction* infeed) {
   // Count nested infeed output tuples.
   int64_t size = 0;
   ShapeUtil::ForEachLeafShape(
@@ -445,7 +451,7 @@ Status HloCostAnalysis::HandleInfeed(const HloInstruction* infeed) {
   return OkStatus();
 }
 
-Status HloCostAnalysis::HandleOutfeed(const HloInstruction* outfeed) {
+absl::Status HloCostAnalysis::HandleOutfeed(const HloInstruction* outfeed) {
   // Count nested outfeed operand tuples.
   current_properties_[kBytesAccessedKey] = 0;
   for (int64_t i = 0; i < outfeed->operand_count(); ++i) {
@@ -465,7 +471,7 @@ Status HloCostAnalysis::HandleOutfeed(const HloInstruction* outfeed) {
   return OkStatus();
 }
 
-Status HloCostAnalysis::HandleMap(const HloInstruction* map) {
+absl::Status HloCostAnalysis::HandleMap(const HloInstruction* map) {
   // Compute properties of the mapped function.
   TF_ASSIGN_OR_RETURN(const Properties sub_properties,
                       ProcessSubcomputation(map->to_apply()));
@@ -480,7 +486,7 @@ Status HloCostAnalysis::HandleMap(const HloInstruction* map) {
   return OkStatus();
 }
 
-Status HloCostAnalysis::HandleReduce(const HloInstruction* reduce) {
+absl::Status HloCostAnalysis::HandleReduce(const HloInstruction* reduce) {
   HloComputation* function = reduce->to_apply();
   // Compute the cost of the user function.
   TF_ASSIGN_OR_RETURN(const Properties sub_properties,
@@ -504,7 +510,7 @@ Status HloCostAnalysis::HandleReduce(const HloInstruction* reduce) {
   return OkStatus();
 }
 
-Status HloCostAnalysis::HandleReduceWindow(
+absl::Status HloCostAnalysis::HandleReduceWindow(
     const HloInstruction* reduce_window) {
   const Window& window = reduce_window->window();
   auto function = reduce_window->to_apply();
@@ -587,7 +593,7 @@ Status HloCostAnalysis::HandleReduceWindow(
   return OkStatus();
 }
 
-Status HloCostAnalysis::HandleSelectAndScatter(
+absl::Status HloCostAnalysis::HandleSelectAndScatter(
     const HloInstruction* instruction) {
   // Compute the properties of the select and scatter function.
   // Compute the properties of the reduction function.
@@ -620,7 +626,7 @@ Status HloCostAnalysis::HandleSelectAndScatter(
   return OkStatus();
 }
 
-Status HloCostAnalysis::HandleBitcast(const HloInstruction*) {
+absl::Status HloCostAnalysis::HandleBitcast(const HloInstruction*) {
   // A bitcast does no computation and touches no memory.
   current_properties_[kBytesAccessedKey] = 0;
   current_properties_.set_output_bytes_accessed(0);
@@ -629,7 +635,7 @@ Status HloCostAnalysis::HandleBitcast(const HloInstruction*) {
   return OkStatus();
 }
 
-Status HloCostAnalysis::HandleBroadcast(const HloInstruction* broadcast) {
+absl::Status HloCostAnalysis::HandleBroadcast(const HloInstruction* broadcast) {
   if (options_.count_multiple_input_accesses) {
     current_properties_.set_operand_bytes_accessed(
         0, GetShapeSize(broadcast->shape()));
@@ -640,74 +646,81 @@ Status HloCostAnalysis::HandleBroadcast(const HloInstruction* broadcast) {
   return OkStatus();
 }
 
-Status HloCostAnalysis::HandlePad(const HloInstruction*) { return OkStatus(); }
+absl::Status HloCostAnalysis::HandlePad(const HloInstruction*) {
+  return OkStatus();
+}
 
-Status HloCostAnalysis::HandleAsyncStart(const HloInstruction* async_start) {
+absl::Status HloCostAnalysis::HandleAsyncStart(
+    const HloInstruction* async_start) {
   TF_ASSIGN_OR_RETURN(
       current_properties_,
       ProcessSubcomputation(async_start->called_computations()[0]));
   return OkStatus();
 }
 
-Status HloCostAnalysis::HandleAsyncUpdate(const HloInstruction*) {
+absl::Status HloCostAnalysis::HandleAsyncUpdate(const HloInstruction*) {
   return OkStatus();
 }
 
-Status HloCostAnalysis::HandleAsyncDone(const HloInstruction*) {
+absl::Status HloCostAnalysis::HandleAsyncDone(const HloInstruction*) {
   return OkStatus();
 }
 
-Status HloCostAnalysis::HandleCopyStart(const HloInstruction*) {
+absl::Status HloCostAnalysis::HandleCopyStart(const HloInstruction*) {
   return OkStatus();
 }
 
-Status HloCostAnalysis::HandleCopyDone(const HloInstruction*) {
+absl::Status HloCostAnalysis::HandleCopyDone(const HloInstruction*) {
   return OkStatus();
 }
 
-Status HloCostAnalysis::HandleSend(const HloInstruction*) { return OkStatus(); }
+absl::Status HloCostAnalysis::HandleSend(const HloInstruction*) {
+  return OkStatus();
+}
 
-Status HloCostAnalysis::HandleSendDone(const HloInstruction*) {
+absl::Status HloCostAnalysis::HandleSendDone(const HloInstruction*) {
   return OkStatus();
 }
 
-Status HloCostAnalysis::HandleRecv(const HloInstruction*) { return OkStatus(); }
+absl::Status HloCostAnalysis::HandleRecv(const HloInstruction*) {
+  return OkStatus();
+}
 
-Status HloCostAnalysis::HandleRecvDone(const HloInstruction*) {
+absl::Status HloCostAnalysis::HandleRecvDone(const HloInstruction*) {
   return OkStatus();
 }
 
-Status HloCostAnalysis::HandleReshape(const HloInstruction*) {
+absl::Status HloCostAnalysis::HandleReshape(const HloInstruction*) {
   return OkStatus();
 }
 
-Status HloCostAnalysis::HandleDynamicReshape(const HloInstruction*) {
+absl::Status HloCostAnalysis::HandleDynamicReshape(const HloInstruction*) {
   return OkStatus();
 }
 
-Status HloCostAnalysis::HandleBatchNormTraining(const HloInstruction*) {
+absl::Status HloCostAnalysis::HandleBatchNormTraining(const HloInstruction*) {
   // TODO(b/62294698): Implement cost analysis for batch-norm-training.
   return OkStatus();
 }
 
-Status HloCostAnalysis::HandleBatchNormInference(const HloInstruction*) {
+absl::Status HloCostAnalysis::HandleBatchNormInference(const HloInstruction*) {
   // TODO(b/62294698): Implement cost analysis for batch-norm-inference.
   return OkStatus();
 }
 
-Status HloCostAnalysis::HandleBatchNormGrad(const HloInstruction*) {
+absl::Status HloCostAnalysis::HandleBatchNormGrad(const HloInstruction*) {
   // TODO(b/62294698): Implement cost analysis for batch-norm-grad.
   return OkStatus();
 }
 
-Status HloCostAnalysis::HandleTranspose(const HloInstruction* transpose) {
+absl::Status HloCostAnalysis::HandleTranspose(const HloInstruction* transpose) {
   if (transpose->IsEffectiveBitcast()) {
     return HandleBitcast(transpose);
   }
   return OkStatus();
 }
 
-Status HloCostAnalysis::HandleAfterAll(const HloInstruction* token) {
+absl::Status HloCostAnalysis::HandleAfterAll(const HloInstruction* token) {
   // This instruction is used to enforce ordering at compile time. No code is
   // emitted.
   current_should_compute_bottleneck_time_ = false;
@@ -720,7 +733,7 @@ Status HloCostAnalysis::HandleAfterAll(const HloInstruction* token) {
   return OkStatus();
 }
 
-Status HloCostAnalysis::HandleAddDependency(
+absl::Status HloCostAnalysis::HandleAddDependency(
     const HloInstruction* add_dependency) {
   // This instruction is used to enforce ordering at compile time. No code is
   // emitted.
@@ -879,12 +892,13 @@ int64_t HloCostAnalysis::GetConvolutionFlops(const HloInstruction* convolution,
   return fma_count * kFmaFlops;
 }
 
-Status HloCostAnalysis::HandleConvolution(const HloInstruction* convolution) {
+absl::Status HloCostAnalysis::HandleConvolution(
+    const HloInstruction* convolution) {
   current_properties_[kFlopsKey] = GetConvolutionFlops(convolution);
   return OkStatus();
 }
 
-Status HloCostAnalysis::HandleFft(const HloInstruction* fft) {
+absl::Status HloCostAnalysis::HandleFft(const HloInstruction* fft) {
   auto real_shape =
       fft->operand(0)->shape().IsTuple()
           ? ShapeUtil::GetTupleElementShape(fft->operand(0)->shape(), 0)
@@ -899,7 +913,7 @@ Status HloCostAnalysis::HandleFft(const HloInstruction* fft) {
   return OkStatus();
 }
 
-Status HloCostAnalysis::HandleTriangularSolve(const HloInstruction* hlo) {
+absl::Status HloCostAnalysis::HandleTriangularSolve(const HloInstruction* hlo) {
   // Half of operand 0 is read.
   float bytes_accessed = GetShapeSize(hlo->shape());
   current_properties_.set_output_bytes_accessed(GetShapeSize(hlo->shape()));
@@ -920,7 +934,7 @@ Status HloCostAnalysis::HandleTriangularSolve(const HloInstruction* hlo) {
   return OkStatus();
 }
 
-Status HloCostAnalysis::HandleCholesky(const HloInstruction* hlo) {
+absl::Status HloCostAnalysis::HandleCholesky(const HloInstruction* hlo) {
   // Half of operand 0 is read and half of the output will be written.
   float bytes_accessed = GetShapeSize(hlo->operand(0)->shape()) / 2.0f;
   current_properties_.set_output_bytes_accessed(
@@ -938,24 +952,25 @@ Status HloCostAnalysis::HandleCholesky(const HloInstruction* hlo) {
   return OkStatus();
 }
 
-Status HloCostAnalysis::HandleOptimizationBarrier(
+absl::Status HloCostAnalysis::HandleOptimizationBarrier(
     const HloInstruction* /*hlo*/) {
   return OkStatus();
 }
 
-Status HloCostAnalysis::HandleAllGather(const HloInstruction* /*hlo*/) {
+absl::Status HloCostAnalysis::HandleAllGather(const HloInstruction* /*hlo*/) {
   return OkStatus();
 }
 
-Status HloCostAnalysis::HandleAllGatherStart(const HloInstruction* hlo) {
+absl::Status HloCostAnalysis::HandleAllGatherStart(const HloInstruction* hlo) {
   return HandleAllGather(hlo);
 }
 
-Status HloCostAnalysis::HandleAllGatherDone(const HloInstruction* /*hlo*/) {
+absl::Status HloCostAnalysis::HandleAllGatherDone(
+    const HloInstruction* /*hlo*/) {
   return OkStatus();
 }
 
-Status HloCostAnalysis::HandleAllReduce(const HloInstruction* crs) {
+absl::Status HloCostAnalysis::HandleAllReduce(const HloInstruction* crs) {
   // We assume 2 replicas, so that each output element is the sum of two input
   // elements.
   //
@@ -980,50 +995,52 @@ Status HloCostAnalysis::HandleAllReduce(const HloInstruction* crs) {
   return OkStatus();
 }
 
-Status HloCostAnalysis::HandleReduceScatter(const HloInstruction* hlo) {
+absl::Status HloCostAnalysis::HandleReduceScatter(const HloInstruction* hlo) {
   return OkStatus();
 }
 
-Status HloCostAnalysis::HandleAllReduceStart(const HloInstruction* hlo) {
+absl::Status HloCostAnalysis::HandleAllReduceStart(const HloInstruction* hlo) {
   return HandleAllReduce(hlo);
 }
 
-Status HloCostAnalysis::HandleAllReduceDone(const HloInstruction* /*hlo*/) {
+absl::Status HloCostAnalysis::HandleAllReduceDone(
+    const HloInstruction* /*hlo*/) {
   return OkStatus();
 }
 
-Status HloCostAnalysis::HandleAllToAll(const HloInstruction* hlo) {
+absl::Status HloCostAnalysis::HandleAllToAll(const HloInstruction* hlo) {
   return OkStatus();
 }
 
-Status HloCostAnalysis::HandleCollectiveBroadcast(
+absl::Status HloCostAnalysis::HandleCollectiveBroadcast(
     const HloInstruction* /*hlo*/) {
   return OkStatus();
 }
 
-Status HloCostAnalysis::HandleCollectivePermute(const HloInstruction* /*hlo*/) {
+absl::Status HloCostAnalysis::HandleCollectivePermute(
+    const HloInstruction* /*hlo*/) {
   return OkStatus();
 }
 
-Status HloCostAnalysis::HandleCollectivePermuteStart(
+absl::Status HloCostAnalysis::HandleCollectivePermuteStart(
     const HloInstruction* /*hlo*/) {
   return OkStatus();
 }
 
-Status HloCostAnalysis::HandleCollectivePermuteDone(
+absl::Status HloCostAnalysis::HandleCollectivePermuteDone(
     const HloInstruction* /*hlo*/) {
   return OkStatus();
 }
 
-Status HloCostAnalysis::HandlePartitionId(const HloInstruction* /*hlo*/) {
+absl::Status HloCostAnalysis::HandlePartitionId(const HloInstruction* /*hlo*/) {
   return OkStatus();
 }
 
-Status HloCostAnalysis::HandleReplicaId(const HloInstruction* /*hlo*/) {
+absl::Status HloCostAnalysis::HandleReplicaId(const HloInstruction* /*hlo*/) {
   return OkStatus();
 }
 
-Status HloCostAnalysis::HandleRng(const HloInstruction* random) {
+absl::Status HloCostAnalysis::HandleRng(const HloInstruction* random) {
   // TODO(b/26346211): Implement better estimates for the RNG cost, since the
   // cost changes with the implementation and the distribution. For now, assume
   // the cost of each RNG is same as a transcendental operation.
@@ -1032,7 +1049,8 @@ Status HloCostAnalysis::HandleRng(const HloInstruction* random) {
   return OkStatus();
 }
 
-Status HloCostAnalysis::HandleRngBitGenerator(const HloInstruction* random) {
+absl::Status HloCostAnalysis::HandleRngBitGenerator(
+    const HloInstruction* random) {
   // TODO(b/26346211): Implement better estimates for the RNG cost, since the
   // cost changes with the implementation and the distribution. For now, assume
   // the cost of each RNG is same as a transcendental operation.
@@ -1041,12 +1059,12 @@ Status HloCostAnalysis::HandleRngBitGenerator(const HloInstruction* random) {
   return OkStatus();
 }
 
-Status HloCostAnalysis::HandleRngGetAndUpdateState(
+absl::Status HloCostAnalysis::HandleRngGetAndUpdateState(
     const HloInstruction* random) {
   return OkStatus();
 }
 
-Status HloCostAnalysis::FusionProcessOutputBytesAccessed(
+absl::Status HloCostAnalysis::FusionProcessOutputBytesAccessed(
     const HloInstruction* fusion) {
   // Fusion nodes that produce a tuple also produce the entries in the tuple.
   // Ignore the memory accessed inside fused ops, since fusion is supposed to
@@ -1118,7 +1136,7 @@ Status HloCostAnalysis::FusionProcessOutputBytesAccessed(
   return OkStatus();
 }
 
-Status HloCostAnalysis::FusionProcessOperandBytesRead(
+absl::Status HloCostAnalysis::FusionProcessOperandBytesRead(
     const HloInstruction* fusion) {
   for (int64_t i = 0; i < fusion->fused_parameters().size(); ++i) {
     const HloInstruction* operand = fusion->fused_parameter(i);
@@ -1160,7 +1178,7 @@ Status HloCostAnalysis::FusionProcessOperandBytesRead(
   return OkStatus();
 }
 
-Status HloCostAnalysis::FusionCountConstantsMemoryAccess(
+absl::Status HloCostAnalysis::FusionCountConstantsMemoryAccess(
     const HloInstruction* fusion) {
   // Count memory access to all large constants.
   for (const HloInstruction* instr :
@@ -1179,7 +1197,7 @@ Status HloCostAnalysis::FusionCountConstantsMemoryAccess(
   return OkStatus();
 }
 
-Status HloCostAnalysis::HandleFusion(const HloInstruction* fusion) {
+absl::Status HloCostAnalysis::HandleFusion(const HloInstruction* fusion) {
   VLOG(8) << "Processing fusion " << fusion->ToString();
 
   if (fusion->IsCustomFusion()) {
@@ -1206,14 +1224,15 @@ Status HloCostAnalysis::HandleFusion(const HloInstruction* fusion) {
   return OkStatus();
 }
 
-Status HloCostAnalysis::HandleCall(const HloInstruction* call) {
+absl::Status HloCostAnalysis::HandleCall(const HloInstruction* call) {
   TF_ASSIGN_OR_RETURN(current_properties_,
                       ProcessSubcomputation(call->to_apply()));
   current_should_compute_bottleneck_time_ = false;
   return OkStatus();
 }
 
-Status HloCostAnalysis::HandleCustomCall(const HloInstruction* custom_call) {
+absl::Status HloCostAnalysis::HandleCustomCall(
+    const HloInstruction* custom_call) {
   // Mark applicable fields as "unknown", since we don't know what this
   // CustomCall does.  This is better than returning an error, which would stop
   // iteration, and therefore would prevent us from getting *any* stats for a
@@ -1229,7 +1248,7 @@ Status HloCostAnalysis::HandleCustomCall(const HloInstruction* custom_call) {
   return OkStatus();
 }
 
-Status HloCostAnalysis::HandleSort(const HloInstruction* sort) {
+absl::Status HloCostAnalysis::HandleSort(const HloInstruction* sort) {
   // This assumes a comparison based N*log(N) algorithm. As for all ops, the
   // actual properties of the op depend on the backend implementation.
   int64_t elements = ShapeUtil::ElementsIn(sort->operand(0)->shape());
@@ -1237,12 +1256,12 @@ Status HloCostAnalysis::HandleSort(const HloInstruction* sort) {
   return OkStatus();
 }
 
-Status HloCostAnalysis::HandleTopK(const HloInstruction* topk) {
+absl::Status HloCostAnalysis::HandleTopK(const HloInstruction* topk) {
   // TODO(cheshire): Cost analysis for TopK.
   return OkStatus();
 }
 
-Status HloCostAnalysis::HandleWhile(const HloInstruction* xla_while) {
+absl::Status HloCostAnalysis::HandleWhile(const HloInstruction* xla_while) {
   // Since the number of iterations of the while node will not always be
   // something that we can statically analyze, we cannot precisely compute the
   // cost of a while node. For now compute the cost of a single iteration.
@@ -1264,7 +1283,8 @@ Status HloCostAnalysis::HandleWhile(const HloInstruction* xla_while) {
   return OkStatus();
 }
 
-Status HloCostAnalysis::HandleConditional(const HloInstruction* conditional) {
+absl::Status HloCostAnalysis::HandleConditional(
+    const HloInstruction* conditional) {
   // Compute the cost of the branch computations and take the maximum from those
   // for each property.
   TF_ASSIGN_OR_RETURN(
@@ -1286,7 +1306,7 @@ Status HloCostAnalysis::HandleConditional(const HloInstruction* conditional) {
   return OkStatus();
 }
 
-Status HloCostAnalysis::HandleGather(const HloInstruction* gather) {
+absl::Status HloCostAnalysis::HandleGather(const HloInstruction* gather) {
   // Gather doesn't read the whole input buffer, it's equivalent to a copy the
   // size of the output shape and a read of the gather indices.
   int64_t output_size = GetShapeSize(gather->shape());
@@ -1303,7 +1323,7 @@ Status HloCostAnalysis::HandleGather(const HloInstruction* gather) {
   return OkStatus();
 }
 
-Status HloCostAnalysis::HandleScatter(const HloInstruction* hlo) {
+absl::Status HloCostAnalysis::HandleScatter(const HloInstruction* hlo) {
   auto* scatter = Cast<HloScatterInstruction>(hlo);
   // Scatter accesses the equivalent of 3N update shapes (input, output, and
   // updates), and the scatter indices.
@@ -1333,17 +1353,17 @@ Status HloCostAnalysis::HandleScatter(const HloInstruction* hlo) {
   return OkStatus();
 }
 
-Status HloCostAnalysis::HandleGetDimensionSize(
+absl::Status HloCostAnalysis::HandleGetDimensionSize(
     const HloInstruction* /*get_size*/) {
   return OkStatus();
 }
 
-Status HloCostAnalysis::HandleSetDimensionSize(
+absl::Status HloCostAnalysis::HandleSetDimensionSize(
     const HloInstruction* /*set_size*/) {
   return OkStatus();
 }
 
-Status HloCostAnalysis::FinishVisit(const HloInstruction*) {
+absl::Status HloCostAnalysis::FinishVisit(const HloInstruction*) {
   return OkStatus();
 }
 
diff --git a/third_party/xla/xla/service/hlo_cost_analysis.h b/third_party/xla/xla/service/hlo_cost_analysis.h
index c0d5a78e1c32d1..5c79b3ddcb1371 100644
--- a/third_party/xla/xla/service/hlo_cost_analysis.h
+++ b/third_party/xla/xla/service/hlo_cost_analysis.h
@@ -440,100 +440,104 @@ class HloCostAnalysis : public ConstDfsHloVisitor {
   explicit HloCostAnalysis(ShapeSizeFunction shape_size,
                            const Properties& per_second_rates = {});
 
-  Status HandleElementwiseUnary(const HloInstruction* hlo) override;
-  Status HandleElementwiseBinary(const HloInstruction* hlo) override;
-  Status HandleConstant(const HloInstruction* constant) override;
-  Status HandleIota(const HloInstruction* iota) override;
-  Status HandleGetTupleElement(
+  absl::Status HandleElementwiseUnary(const HloInstruction* hlo) override;
+  absl::Status HandleElementwiseBinary(const HloInstruction* hlo) override;
+  absl::Status HandleConstant(const HloInstruction* constant) override;
+  absl::Status HandleIota(const HloInstruction* iota) override;
+  absl::Status HandleGetTupleElement(
       const HloInstruction* get_tuple_element) override;
-  Status HandleSelect(const HloInstruction* hlo) override;
-  Status HandleCompare(const HloInstruction* compare) override;
-  Status HandleClamp(const HloInstruction* clamp) override;
-  Status HandleReducePrecision(const HloInstruction* hlo) override;
-  Status HandleConcatenate(const HloInstruction* concatenate) override;
-  Status HandleAsyncStart(const HloInstruction* async_start) override;
-  Status HandleAsyncUpdate(const HloInstruction* async_update) override;
-  Status HandleAsyncDone(const HloInstruction* async_done) override;
-  Status HandleCopyStart(const HloInstruction* send) override;
-  Status HandleCopyDone(const HloInstruction* send_done) override;
-  Status HandleSend(const HloInstruction* send) override;
-  Status HandleSendDone(const HloInstruction* send_done) override;
-  Status HandleRecv(const HloInstruction* recv) override;
-  Status HandleRecvDone(const HloInstruction* recv_done) override;
-  Status HandleConvert(const HloInstruction* convert) override;
-  Status HandleCopy(const HloInstruction* copy) override;
-  Status HandleDomain(const HloInstruction* domain) override;
-  Status HandleDot(const HloInstruction* dot) override;
-  Status HandleConvolution(const HloInstruction* convolution) override;
-  Status HandleFft(const HloInstruction* fft) override;
-  Status HandleTriangularSolve(const HloInstruction* hlo) override;
-  Status HandleCholesky(const HloInstruction* hlo) override;
-  Status HandleOptimizationBarrier(const HloInstruction* hlo) override;
-  Status HandleAllGather(const HloInstruction* hlo) override;
-  Status HandleAllGatherStart(const HloInstruction* hlo) override;
-  Status HandleAllGatherDone(const HloInstruction* hlo) override;
-  Status HandleAllReduce(const HloInstruction* crs) override;
-  Status HandleReduceScatter(const HloInstruction* hlo) override;
-  Status HandleAllReduceStart(const HloInstruction* hlo) override;
-  Status HandleAllReduceDone(const HloInstruction* hlo) override;
-  Status HandleAllToAll(const HloInstruction* hlo) override;
-  Status HandleCollectiveBroadcast(const HloInstruction* hlo) override;
-  Status HandleCollectivePermute(const HloInstruction* hlo) override;
-  Status HandleCollectivePermuteStart(const HloInstruction* hlo) override;
-  Status HandleCollectivePermuteDone(const HloInstruction* hlo) override;
-  Status HandleReplicaId(const HloInstruction* hlo) override;
-  Status HandlePartitionId(const HloInstruction* hlo) override;
-  Status HandleInfeed(const HloInstruction* infeed) override;
-  Status HandleOutfeed(const HloInstruction* outfeed) override;
-  Status HandleRng(const HloInstruction* random) override;
-  Status HandleRngBitGenerator(const HloInstruction* random) override;
-  Status HandleRngGetAndUpdateState(const HloInstruction* random) override;
-  Status HandleReverse(const HloInstruction* reverse) override;
-  Status HandleSort(const HloInstruction* sort) override;
-  Status HandleParameter(const HloInstruction* parameter) override;
-  Status HandleReduce(const HloInstruction* reduce) override;
-  Status HandleBatchNormTraining(
+  absl::Status HandleSelect(const HloInstruction* hlo) override;
+  absl::Status HandleCompare(const HloInstruction* compare) override;
+  absl::Status HandleClamp(const HloInstruction* clamp) override;
+  absl::Status HandleReducePrecision(const HloInstruction* hlo) override;
+  absl::Status HandleConcatenate(const HloInstruction* concatenate) override;
+  absl::Status HandleAsyncStart(const HloInstruction* async_start) override;
+  absl::Status HandleAsyncUpdate(const HloInstruction* async_update) override;
+  absl::Status HandleAsyncDone(const HloInstruction* async_done) override;
+  absl::Status HandleCopyStart(const HloInstruction* send) override;
+  absl::Status HandleCopyDone(const HloInstruction* send_done) override;
+  absl::Status HandleSend(const HloInstruction* send) override;
+  absl::Status HandleSendDone(const HloInstruction* send_done) override;
+  absl::Status HandleRecv(const HloInstruction* recv) override;
+  absl::Status HandleRecvDone(const HloInstruction* recv_done) override;
+  absl::Status HandleConvert(const HloInstruction* convert) override;
+  absl::Status HandleCopy(const HloInstruction* copy) override;
+  absl::Status HandleDomain(const HloInstruction* domain) override;
+  absl::Status HandleDot(const HloInstruction* dot) override;
+  absl::Status HandleConvolution(const HloInstruction* convolution) override;
+  absl::Status HandleFft(const HloInstruction* fft) override;
+  absl::Status HandleTriangularSolve(const HloInstruction* hlo) override;
+  absl::Status HandleCholesky(const HloInstruction* hlo) override;
+  absl::Status HandleOptimizationBarrier(const HloInstruction* hlo) override;
+  absl::Status HandleAllGather(const HloInstruction* hlo) override;
+  absl::Status HandleAllGatherStart(const HloInstruction* hlo) override;
+  absl::Status HandleAllGatherDone(const HloInstruction* hlo) override;
+  absl::Status HandleAllReduce(const HloInstruction* crs) override;
+  absl::Status HandleReduceScatter(const HloInstruction* hlo) override;
+  absl::Status HandleAllReduceStart(const HloInstruction* hlo) override;
+  absl::Status HandleAllReduceDone(const HloInstruction* hlo) override;
+  absl::Status HandleAllToAll(const HloInstruction* hlo) override;
+  absl::Status HandleCollectiveBroadcast(const HloInstruction* hlo) override;
+  absl::Status HandleCollectivePermute(const HloInstruction* hlo) override;
+  absl::Status HandleCollectivePermuteStart(const HloInstruction* hlo) override;
+  absl::Status HandleCollectivePermuteDone(const HloInstruction* hlo) override;
+  absl::Status HandleReplicaId(const HloInstruction* hlo) override;
+  absl::Status HandlePartitionId(const HloInstruction* hlo) override;
+  absl::Status HandleInfeed(const HloInstruction* infeed) override;
+  absl::Status HandleOutfeed(const HloInstruction* outfeed) override;
+  absl::Status HandleRng(const HloInstruction* random) override;
+  absl::Status HandleRngBitGenerator(const HloInstruction* random) override;
+  absl::Status HandleRngGetAndUpdateState(
+      const HloInstruction* random) override;
+  absl::Status HandleReverse(const HloInstruction* reverse) override;
+  absl::Status HandleSort(const HloInstruction* sort) override;
+  absl::Status HandleParameter(const HloInstruction* parameter) override;
+  absl::Status HandleReduce(const HloInstruction* reduce) override;
+  absl::Status HandleBatchNormTraining(
       const HloInstruction* batch_norm_training) override;
-  Status HandleBatchNormInference(
+  absl::Status HandleBatchNormInference(
       const HloInstruction* batch_norm_inference) override;
-  Status HandleBatchNormGrad(const HloInstruction* batch_norm_grad) override;
-  Status HandleFusion(const HloInstruction* fusion) override;
-  Status HandleCall(const HloInstruction* call) override;
-  Status HandleCustomCall(const HloInstruction* custom_call) override;
-  Status HandleSlice(const HloInstruction* slice) override;
-  Status HandleDynamicSlice(const HloInstruction* dynamic_slice) override;
-  Status HandleDynamicUpdateSlice(
+  absl::Status HandleBatchNormGrad(
+      const HloInstruction* batch_norm_grad) override;
+  absl::Status HandleFusion(const HloInstruction* fusion) override;
+  absl::Status HandleCall(const HloInstruction* call) override;
+  absl::Status HandleCustomCall(const HloInstruction* custom_call) override;
+  absl::Status HandleSlice(const HloInstruction* slice) override;
+  absl::Status HandleDynamicSlice(const HloInstruction* dynamic_slice) override;
+  absl::Status HandleDynamicUpdateSlice(
       const HloInstruction* dynamic_update_slice) override;
-  Status HandleTuple(const HloInstruction* tuple) override;
-  Status HandleMap(const HloInstruction* map) override;
-  Status HandleReduceWindow(const HloInstruction* reduce_window) override;
-  Status HandleSelectAndScatter(const HloInstruction* instruction) override;
-  Status HandleBitcast(const HloInstruction* bitcast) override;
-  Status HandleBroadcast(const HloInstruction* broadcast) override;
-  Status HandlePad(const HloInstruction* pad) override;
-  Status HandleReshape(const HloInstruction* reshape) override;
-  Status HandleDynamicReshape(const HloInstruction* reshape) override;
-  Status HandleAddDependency(const HloInstruction* add_dependency) override;
-  Status HandleAfterAll(const HloInstruction* token) override;
-  Status HandleTranspose(const HloInstruction* transpose) override;
-  Status HandleWhile(const HloInstruction* xla_while) override;
-  Status HandleConditional(const HloInstruction* conditional) override;
-  Status HandleGather(const HloInstruction* gather) override;
-  Status HandleScatter(const HloInstruction* hlo) override;
-  Status HandleGetDimensionSize(const HloInstruction* get_size) override;
-  Status HandleSetDimensionSize(const HloInstruction* set_size) override;
-  Status HandleTopK(const HloInstruction* topk) override;
-  Status FinishVisit(const HloInstruction* root) override;
-
-  Status Preprocess(const HloInstruction* hlo) override;
-  Status Postprocess(const HloInstruction* hlo) override;
+  absl::Status HandleTuple(const HloInstruction* tuple) override;
+  absl::Status HandleMap(const HloInstruction* map) override;
+  absl::Status HandleReduceWindow(const HloInstruction* reduce_window) override;
+  absl::Status HandleSelectAndScatter(
+      const HloInstruction* instruction) override;
+  absl::Status HandleBitcast(const HloInstruction* bitcast) override;
+  absl::Status HandleBroadcast(const HloInstruction* broadcast) override;
+  absl::Status HandlePad(const HloInstruction* pad) override;
+  absl::Status HandleReshape(const HloInstruction* reshape) override;
+  absl::Status HandleDynamicReshape(const HloInstruction* reshape) override;
+  absl::Status HandleAddDependency(
+      const HloInstruction* add_dependency) override;
+  absl::Status HandleAfterAll(const HloInstruction* token) override;
+  absl::Status HandleTranspose(const HloInstruction* transpose) override;
+  absl::Status HandleWhile(const HloInstruction* xla_while) override;
+  absl::Status HandleConditional(const HloInstruction* conditional) override;
+  absl::Status HandleGather(const HloInstruction* gather) override;
+  absl::Status HandleScatter(const HloInstruction* hlo) override;
+  absl::Status HandleGetDimensionSize(const HloInstruction* get_size) override;
+  absl::Status HandleSetDimensionSize(const HloInstruction* set_size) override;
+  absl::Status HandleTopK(const HloInstruction* topk) override;
+  absl::Status FinishVisit(const HloInstruction* root) override;
+
+  absl::Status Preprocess(const HloInstruction* hlo) override;
+  absl::Status Postprocess(const HloInstruction* hlo) override;
 
   // Enable efficient updates if a known small set of instructions within an
   // HLO graph was modified.
   // Updates the cost analysis by removing one instruction.
-  Status RemoveInstruction(HloInstruction* instruction);
+  absl::Status RemoveInstruction(HloInstruction* instruction);
   // Updates the cost analysis by re-doing the analysis of one instruction.
-  Status RevisitInstruction(HloInstruction* instruction);
+  absl::Status RevisitInstruction(HloInstruction* instruction);
 
   // Decorates shape_size_ by returning 0 immediately if the shape does not have
   // a layout.
@@ -607,14 +611,17 @@ class HloCostAnalysis : public ConstDfsHloVisitor {
  protected:
   // Computes the bytes accessed based on the outputs produced by the fusion
   // instruction.
-  virtual Status FusionProcessOutputBytesAccessed(const HloInstruction* fusion);
+  virtual absl::Status FusionProcessOutputBytesAccessed(
+      const HloInstruction* fusion);
 
   // Computes the bytes accessed (read) based on the inputs consumed by the
   // fusion instruction.
-  virtual Status FusionProcessOperandBytesRead(const HloInstruction* fusion);
+  virtual absl::Status FusionProcessOperandBytesRead(
+      const HloInstruction* fusion);
 
   // Computes memory access to all larger constants in the fusion instruction.
-  virtual Status FusionCountConstantsMemoryAccess(const HloInstruction* fusion);
+  virtual absl::Status FusionCountConstantsMemoryAccess(
+      const HloInstruction* fusion);
 
   // Allows exclusion of certain types of inputs from bytes accessed during
   // FusionProcessOperandBytesRead.
@@ -657,7 +664,7 @@ class HloCostAnalysis : public ConstDfsHloVisitor {
       HloComputation* computation);
 
   // Utility function to handle all element-wise operations.
-  Status HandleElementwiseOp(const HloInstruction* hlo_instruction);
+  absl::Status HandleElementwiseOp(const HloInstruction* hlo_instruction);
 
   // Returns 0.0f if the hlo is not present in hlo_to_properties or if the key
   // is not present in hlo_to_properties[hlo]. Otherwise, returns the value that
@@ -672,7 +679,8 @@ class HloCostAnalysis : public ConstDfsHloVisitor {
 
   // Traverses a fusion counting total utilization of every instruction inside.
   // Currently implemented non-trivially only in the GPU cost analysis.
-  virtual Status FusionCalculateUtilizations(const HloInstruction* fusion);
+  virtual absl::Status FusionCalculateUtilizations(
+      const HloInstruction* fusion);
 
   HloToProperties hlo_properties_;
 
diff --git a/third_party/xla/xla/service/hlo_dataflow_analysis.cc b/third_party/xla/xla/service/hlo_dataflow_analysis.cc
index a60e0211ca5648..7d23d9d60427d3 100644
--- a/third_party/xla/xla/service/hlo_dataflow_analysis.cc
+++ b/third_party/xla/xla/service/hlo_dataflow_analysis.cc
@@ -1418,7 +1418,7 @@ InstructionValueSet& HloDataflowAnalysis::GetInstructionValueSet(
   return *value_sets_.find(instruction)->second;
 }
 
-Status HloDataflowAnalysis::InitializeInstructionValueSets() {
+absl::Status HloDataflowAnalysis::InitializeInstructionValueSets() {
   for (const HloComputation* computation : module_.MakeComputationSorted()) {
     if (!HloInstruction::IsThreadIncluded(computation->execution_thread(),
                                           execution_threads_)) {
@@ -1731,7 +1731,7 @@ absl::StatusOr<std::unique_ptr<HloDataflowAnalysis>> HloDataflowAnalysis::Run(
   return std::move(dataflow_analysis);
 }
 
-Status HloDataflowAnalysis::Verify() const {
+absl::Status HloDataflowAnalysis::Verify() const {
   // Verify each HloValue appears in the value sets that the value's positions()
   // indicate.
   for (const HloValue* value : values()) {
diff --git a/third_party/xla/xla/service/hlo_dataflow_analysis.h b/third_party/xla/xla/service/hlo_dataflow_analysis.h
index a5cf84409aae4d..ddc97b0c7df12c 100644
--- a/third_party/xla/xla/service/hlo_dataflow_analysis.h
+++ b/third_party/xla/xla/service/hlo_dataflow_analysis.h
@@ -226,7 +226,7 @@ class HloDataflowAnalysis {
   GetInPlaceInputOutputPairs(const HloInstruction* instruction);
 
   // Verifies various invariants of the dataflow analysis.
-  Status Verify() const;
+  absl::Status Verify() const;
 
  private:
   static bool AreTransitiveUsesElementwiseOrTuple(const HloInstruction* inst);
@@ -264,7 +264,7 @@ class HloDataflowAnalysis {
   // Constructs and initializes the InstructionValueSets of all instructions to
   // contain exactly the HloValues defined by each instruction. These values can
   // then propagated throughout the HLO graph by calling Propagate.
-  Status InitializeInstructionValueSets();
+  absl::Status InitializeInstructionValueSets();
 
   // Updates the value set of the given instruction based on the values flowing
   // into the instruction (operands and cross-computation dataflow).
diff --git a/third_party/xla/xla/service/hlo_dataflow_analysis_test.cc b/third_party/xla/xla/service/hlo_dataflow_analysis_test.cc
index 2244f5e37fb923..89ce2c4bea22f8 100644
--- a/third_party/xla/xla/service/hlo_dataflow_analysis_test.cc
+++ b/third_party/xla/xla/service/hlo_dataflow_analysis_test.cc
@@ -2105,7 +2105,7 @@ TEST_F(HloDataflowAnalysisTest, AllGatherStartAndDoneWithTuple) {
   )";
   TF_ASSERT_OK_AND_ASSIGN(module_, ParseAndReturnVerifiedModule(hlo_text));
   const HloDataflowAnalysis& analysis = RunAnalysis(/*ssa_form=*/false);
-  Status status = analysis.Verify();
+  absl::Status status = analysis.Verify();
   EXPECT_TRUE(status.ok()) << status.ToString();
 
   HloInstruction* done = module_->entry_computation()->root_instruction();
diff --git a/third_party/xla/xla/service/hlo_dce.cc b/third_party/xla/xla/service/hlo_dce.cc
index 4a7e9af403bf37..bd71c1d484452d 100644
--- a/third_party/xla/xla/service/hlo_dce.cc
+++ b/third_party/xla/xla/service/hlo_dce.cc
@@ -183,7 +183,7 @@ bool IsRemovableWhile(HloInstruction* instruction,
   return changed;
 }
 
-Status HloDCE::RecursivelyRemoveDeadComputation(
+absl::Status HloDCE::RecursivelyRemoveDeadComputation(
     HloModule* module, HloComputation* computation,
     absl::flat_hash_map<HloComputation*, int>& live_call_counts) {
   std::vector<HloComputation*> to_be_deleted;
diff --git a/third_party/xla/xla/service/hlo_dce.h b/third_party/xla/xla/service/hlo_dce.h
index 8e8e3fab4d58ec..c6330d30c38ef8 100644
--- a/third_party/xla/xla/service/hlo_dce.h
+++ b/third_party/xla/xla/service/hlo_dce.h
@@ -66,7 +66,7 @@ class HloDCE : public HloModulePass {
   // computations and checks if any of the subcomputations become dead after the
   // removal. Returns whether all dead computations were successfully removed
   // from the module.
-  Status RecursivelyRemoveDeadComputation(
+  absl::Status RecursivelyRemoveDeadComputation(
       HloModule* module, HloComputation* computation,
       absl::flat_hash_map<HloComputation*, int>& live_call_counts);
 
diff --git a/third_party/xla/xla/service/hlo_domain_map.cc b/third_party/xla/xla/service/hlo_domain_map.cc
index 6aa3933244de8f..79ae9f04476ad1 100644
--- a/third_party/xla/xla/service/hlo_domain_map.cc
+++ b/third_party/xla/xla/service/hlo_domain_map.cc
@@ -60,7 +60,7 @@ int64_t HloDomainMap::GetDomainMetadataId(
   return FindOrDie(domain_metadata_id_, instruction);
 }
 
-Status HloDomainMap::TryProcessEmptyDomain(HloInstruction* instruction) {
+absl::Status HloDomainMap::TryProcessEmptyDomain(HloInstruction* instruction) {
   TF_RET_CHECK(instruction->opcode() == HloOpcode::kDomain);
   // We only check operands, so we are sure to not process the empty domain from
   // both sides.
@@ -80,7 +80,7 @@ Status HloDomainMap::TryProcessEmptyDomain(HloInstruction* instruction) {
   return OkStatus();
 }
 
-Status HloDomainMap::Populate(HloComputation* computation) {
+absl::Status HloDomainMap::Populate(HloComputation* computation) {
   InstructionOrderMap instructions_post_order;
   int64_t count = 0;
   for (HloInstruction* instruction : computation->MakeInstructionPostOrder()) {
@@ -106,7 +106,7 @@ Status HloDomainMap::Populate(HloComputation* computation) {
   return OkStatus();
 }
 
-Status HloDomainMap::PopulateDomainMetadataMap() {
+absl::Status HloDomainMap::PopulateDomainMetadataMap() {
   auto hash = [](const DomainMetadata* m) { return m->Hash(); };
   auto equal = [](const DomainMetadata* a, const DomainMetadata* b) {
     return a->Matches(*b);
@@ -142,7 +142,7 @@ Status HloDomainMap::PopulateDomainMetadataMap() {
   return OkStatus();
 }
 
-Status HloDomainMap::InsertDomain(
+absl::Status HloDomainMap::InsertDomain(
     std::unique_ptr<DomainMetadata::Domain> domain) {
   int64_t domain_id = instruction_domains_.size();
   instruction_domains_.push_back(std::move(domain));
@@ -152,8 +152,8 @@ Status HloDomainMap::InsertDomain(
   return OkStatus();
 }
 
-Status HloDomainMap::ExpandDomain(HloInstruction* instruction,
-                                  DomainMetadata::Domain* domain) const {
+absl::Status HloDomainMap::ExpandDomain(HloInstruction* instruction,
+                                        DomainMetadata::Domain* domain) const {
   std::vector<HloInstruction*> in_queue;
   in_queue.push_back(instruction);
   while (!in_queue.empty()) {
diff --git a/third_party/xla/xla/service/hlo_domain_map.h b/third_party/xla/xla/service/hlo_domain_map.h
index 97814571274337..5dcae5ff4cf01a 100644
--- a/third_party/xla/xla/service/hlo_domain_map.h
+++ b/third_party/xla/xla/service/hlo_domain_map.h
@@ -87,13 +87,13 @@ class HloDomainMap {
   // kDomain instruction of the same kind, hence defining an empty domain.
   // If that is the case, create the empty domain and call the proper
   // normalizer.
-  Status TryProcessEmptyDomain(HloInstruction* instruction);
+  absl::Status TryProcessEmptyDomain(HloInstruction* instruction);
 
-  Status Populate(HloComputation* computation);
+  absl::Status Populate(HloComputation* computation);
 
   // Inserts the provided domain into the ones tracked by this object,
   // creating a new domain ID.
-  Status InsertDomain(std::unique_ptr<DomainMetadata::Domain> domain);
+  absl::Status InsertDomain(std::unique_ptr<DomainMetadata::Domain> domain);
 
   // From the given instruction, expands operand and user wise, the set of
   // instructions which can be reached without crossing a kDomain instruction
@@ -101,8 +101,8 @@ class HloDomainMap {
   // The domain data structure will be populated with all the reached
   // instructions, and the boundaries of the domain, with the kDomain
   // instructions encountered while expanding the reach.
-  Status ExpandDomain(HloInstruction* instruction,
-                      DomainMetadata::Domain* domain) const;
+  absl::Status ExpandDomain(HloInstruction* instruction,
+                            DomainMetadata::Domain* domain) const;
 
   // Creates a domain data structure using the ExpandDomain() API.
   absl::StatusOr<std::unique_ptr<DomainMetadata::Domain>> CreateDomain(
@@ -117,7 +117,7 @@ class HloDomainMap {
 
   // Populates domain_metadata_id_ that maps each HloInstruction to the unique
   // ID of its associated domain metatadata.
-  Status PopulateDomainMetadataMap();
+  absl::Status PopulateDomainMetadataMap();
 
   std::string domain_kind_;
   std::vector<std::unique_ptr<DomainMetadata::Domain>> instruction_domains_;
diff --git a/third_party/xla/xla/service/hlo_domain_remover.cc b/third_party/xla/xla/service/hlo_domain_remover.cc
index 3d35a25595e1ba..dad5f41ea71885 100644
--- a/third_party/xla/xla/service/hlo_domain_remover.cc
+++ b/third_party/xla/xla/service/hlo_domain_remover.cc
@@ -36,13 +36,13 @@ class HloDomainRemover::RunContext {
  private:
   // Verifies the consistency of the domain, and normalizes the instructions
   // within it.
-  Status VerifyAndNormalizeDomain(const DomainMetadata::Domain& domain);
+  absl::Status VerifyAndNormalizeDomain(const DomainMetadata::Domain& domain);
 
   HloModule* module_;
   HloDomainRemover* remover_;
 };
 
-Status HloDomainRemover::RunContext::VerifyAndNormalizeDomain(
+absl::Status HloDomainRemover::RunContext::VerifyAndNormalizeDomain(
     const DomainMetadata::Domain& domain) {
   TF_ASSIGN_OR_RETURN(const DomainMetadata* ref_metadata,
                       HloDomainVerifier::VerifyDomain(domain));
diff --git a/third_party/xla/xla/service/hlo_domain_remover.h b/third_party/xla/xla/service/hlo_domain_remover.h
index 9cda66d36e8451..6510ab684f9c37 100644
--- a/third_party/xla/xla/service/hlo_domain_remover.h
+++ b/third_party/xla/xla/service/hlo_domain_remover.h
@@ -37,8 +37,8 @@ class HloDomainRemover : public HloModulePass {
   // function is tasked at applying attribute normalization on the instructions
   // within such domain.
   HloDomainRemover(absl::string_view kind,
-                   std::function<Status(const DomainMetadata::Domain&,
-                                        const DomainMetadata* metadata)>
+                   std::function<absl::Status(const DomainMetadata::Domain&,
+                                              const DomainMetadata* metadata)>
                        normalizer)
       : kind_(kind), normalizer_(std::move(normalizer)) {}
 
@@ -58,8 +58,8 @@ class HloDomainRemover : public HloModulePass {
   class RunContext;
 
   std::string kind_;
-  std::function<Status(const DomainMetadata::Domain&,
-                       const DomainMetadata* metadata)>
+  std::function<absl::Status(const DomainMetadata::Domain&,
+                             const DomainMetadata* metadata)>
       normalizer_;
 };
 
diff --git a/third_party/xla/xla/service/hlo_domain_test.cc b/third_party/xla/xla/service/hlo_domain_test.cc
index 8e82878602f1a6..d0c7312868d826 100644
--- a/third_party/xla/xla/service/hlo_domain_test.cc
+++ b/third_party/xla/xla/service/hlo_domain_test.cc
@@ -121,8 +121,8 @@ class OpNameDomainCreator {
   }
 };
 
-Status OpNameDomainNormalizer(const DomainMetadata::Domain& domain,
-                              const DomainMetadata* metadata) {
+absl::Status OpNameDomainNormalizer(const DomainMetadata::Domain& domain,
+                                    const DomainMetadata* metadata) {
   // Nothing to do for the particular use this test make of the OpName domains.
   return OkStatus();
 }
diff --git a/third_party/xla/xla/service/hlo_domain_verifier.cc b/third_party/xla/xla/service/hlo_domain_verifier.cc
index 0d83a318dbf9ea..50f874945bf382 100644
--- a/third_party/xla/xla/service/hlo_domain_verifier.cc
+++ b/third_party/xla/xla/service/hlo_domain_verifier.cc
@@ -31,19 +31,20 @@ class HloDomainVerifier::RunContext {
   RunContext(HloModule* module, HloDomainVerifier* verifier)
       : module_(module), verifier_(verifier) {}
 
-  Status Run(const absl::flat_hash_set<absl::string_view>& execution_threads);
+  absl::Status Run(
+      const absl::flat_hash_set<absl::string_view>& execution_threads);
 
  private:
   // If the verifier caller passed an empty vector for kinds, we collect all the
   // available domain types.
-  Status PopulateDomainKinds(
+  absl::Status PopulateDomainKinds(
       const absl::flat_hash_set<absl::string_view>& execution_threads);
 
   HloModule* module_;
   HloDomainVerifier* verifier_;
 };
 
-Status HloDomainVerifier::RunContext::PopulateDomainKinds(
+absl::Status HloDomainVerifier::RunContext::PopulateDomainKinds(
     const absl::flat_hash_set<absl::string_view>& execution_threads) {
   if (verifier_->kinds_.empty()) {
     // The caller specified no domain kinds, collect all the ones available.
@@ -65,7 +66,7 @@ Status HloDomainVerifier::RunContext::PopulateDomainKinds(
   return OkStatus();
 }
 
-Status HloDomainVerifier::RunContext::Run(
+absl::Status HloDomainVerifier::RunContext::Run(
     const absl::flat_hash_set<absl::string_view>& execution_threads) {
   VLOG(4) << "Running HLO Domain Verifier";
   TF_RETURN_IF_ERROR(PopulateDomainKinds(execution_threads));
diff --git a/third_party/xla/xla/service/hlo_graph_dumper.cc b/third_party/xla/xla/service/hlo_graph_dumper.cc
index c16dfb4328b66f..16fd21fd506821 100644
--- a/third_party/xla/xla/service/hlo_graph_dumper.cc
+++ b/third_party/xla/xla/service/hlo_graph_dumper.cc
@@ -1843,14 +1843,14 @@ static absl::StatusOr<std::string> CompressAndEncode(absl::string_view input) {
     explicit WritableStringFile(std::string* data) : data_(data){};
     ~WritableStringFile() override = default;
 
-    Status Append(absl::string_view data) override {
+    absl::Status Append(absl::string_view data) override {
       absl::StrAppend(data_, data);
       return OkStatus();
     }
 
-    Status Close() override { return OkStatus(); }
-    Status Flush() override { return OkStatus(); }
-    Status Sync() override { return OkStatus(); }
+    absl::Status Close() override { return OkStatus(); }
+    absl::Status Flush() override { return OkStatus(); }
+    absl::Status Sync() override { return OkStatus(); }
 
    private:
     std::string* data_;
diff --git a/third_party/xla/xla/service/hlo_instruction_test.cc b/third_party/xla/xla/service/hlo_instruction_test.cc
index 9d555475776704..45f19b0145ef19 100644
--- a/third_party/xla/xla/service/hlo_instruction_test.cc
+++ b/third_party/xla/xla/service/hlo_instruction_test.cc
@@ -59,24 +59,24 @@ class HloInstructionTest : public HloTestBase {
 // before their users, nodes not visited twice, etc.)
 class OpAndUserCollectingVisitor : public DfsHloVisitorWithDefault {
  public:
-  Status DefaultAction(HloInstruction* hlo_instruction) override {
+  absl::Status DefaultAction(HloInstruction* hlo_instruction) override {
     return Unimplemented("not implemented %s",
                          HloOpcodeString(hlo_instruction->opcode()));
   }
 
-  Status HandleParameter(HloInstruction* parameter) override {
+  absl::Status HandleParameter(HloInstruction* parameter) override {
     EXPECT_FALSE(count_.contains(parameter));
     count_[parameter] = GetCountsForNode(parameter);
     return OkStatus();
   }
 
-  Status HandleConstant(HloInstruction* constant) override {
+  absl::Status HandleConstant(HloInstruction* constant) override {
     EXPECT_FALSE(count_.contains(constant));
     count_[constant] = GetCountsForNode(constant);
     return OkStatus();
   }
 
-  Status HandleAdd(HloInstruction* add) override {
+  absl::Status HandleAdd(HloInstruction* add) override {
     auto lhs = add->operand(0);
     auto rhs = add->operand(1);
     EXPECT_FALSE(count_.contains(add));
@@ -86,7 +86,7 @@ class OpAndUserCollectingVisitor : public DfsHloVisitorWithDefault {
     return OkStatus();
   }
 
-  Status HandleNegate(HloInstruction* negate) override {
+  absl::Status HandleNegate(HloInstruction* negate) override {
     auto operand = negate->operand(0);
     EXPECT_FALSE(count_.contains(negate));
     EXPECT_TRUE(count_.contains(operand));
@@ -94,7 +94,7 @@ class OpAndUserCollectingVisitor : public DfsHloVisitorWithDefault {
     return OkStatus();
   }
 
-  Status HandleMap(HloInstruction* map) override {
+  absl::Status HandleMap(HloInstruction* map) override {
     EXPECT_FALSE(count_.contains(map));
     for (HloInstruction* arg : map->operands()) {
       EXPECT_TRUE(count_.contains(arg));
@@ -103,7 +103,7 @@ class OpAndUserCollectingVisitor : public DfsHloVisitorWithDefault {
     return OkStatus();
   }
 
-  Status HandleReduce(HloInstruction* reduce) override {
+  absl::Status HandleReduce(HloInstruction* reduce) override {
     auto arg = reduce->operand(0);
     auto init_value = reduce->operand(1);
     EXPECT_FALSE(count_.contains(reduce));
@@ -576,12 +576,12 @@ class NodeCollectorAndPostProcessor : public DfsHloVisitorWithDefault {
  public:
   NodeCollectorAndPostProcessor() {}
 
-  Status Postprocess(HloInstruction* hlo) override {
+  absl::Status Postprocess(HloInstruction* hlo) override {
     post_processed_nodes_.push_back(hlo);
     return OkStatus();
   }
 
-  Status DefaultAction(HloInstruction* hlo_instruction) override {
+  absl::Status DefaultAction(HloInstruction* hlo_instruction) override {
     visited_nodes_.push_back(hlo_instruction);
     return OkStatus();
   }
diff --git a/third_party/xla/xla/service/hlo_module_config_test.cc b/third_party/xla/xla/service/hlo_module_config_test.cc
index 952aec1244feca..6481dc2b511a1a 100644
--- a/third_party/xla/xla/service/hlo_module_config_test.cc
+++ b/third_party/xla/xla/service/hlo_module_config_test.cc
@@ -27,7 +27,7 @@ namespace {
 using ::testing::EqualsProto;
 
 template <typename MessageType>
-StatusOr<MessageType> ParseTextProto(const std::string& text_proto) {
+absl::StatusOr<MessageType> ParseTextProto(const std::string& text_proto) {
   tsl::protobuf::TextFormat::Parser parser;
   MessageType parsed_proto;
   tsl::protobuf::io::ArrayInputStream input_stream(text_proto.data(),
diff --git a/third_party/xla/xla/service/hlo_module_group_metadata.cc b/third_party/xla/xla/service/hlo_module_group_metadata.cc
index 2ab8c03c7ec030..1a21714791de56 100644
--- a/third_party/xla/xla/service/hlo_module_group_metadata.cc
+++ b/third_party/xla/xla/service/hlo_module_group_metadata.cc
@@ -64,12 +64,12 @@ HloModuleGroupMetadata::Build(absl::Span<HloModule* const> modules) {
   return std::move(metadata);
 }
 
-Status HloModuleGroupMetadata::Build() {
+absl::Status HloModuleGroupMetadata::Build() {
   TF_RETURN_IF_ERROR(RecordInstructions());
   TF_RETURN_IF_ERROR(VerifyChannelInstructions());
 
   // Record all companion while instructions.
-  const auto visitor = [this](HloInstruction* hlo) -> Status {
+  const auto visitor = [this](HloInstruction* hlo) -> absl::Status {
     // We only need to process if the instruction is within the computation
     // of a companion instruction, like in the condition or body computation
     // of a While.
@@ -159,7 +159,7 @@ Status HloModuleGroupMetadata::Build() {
   return OkStatus();
 }
 
-Status HloModuleGroupMetadata::VerifyCompanionSets() const {
+absl::Status HloModuleGroupMetadata::VerifyCompanionSets() const {
   for (const auto& companions : companion_sets_) {
     // A companion set must be composed at most of an instruction per
     // device/module.
@@ -320,8 +320,8 @@ int64_t HloModuleGroupMetadata::GetDeviceModulesCount() const {
   return modules_.size();
 }
 
-Status HloModuleGroupMetadata::RecordInstructions() {
-  const auto visitor = [this](HloInstruction* hlo) -> Status {
+absl::Status HloModuleGroupMetadata::RecordInstructions() {
+  const auto visitor = [this](HloInstruction* hlo) -> absl::Status {
     if (hlo->opcode() == HloOpcode::kWhile) {
       tracked_instructions_[hlo->while_condition()] =
           TrackedInstruction(hlo, ComputationKind::kWhileCondition);
@@ -404,8 +404,8 @@ Status HloModuleGroupMetadata::RecordInstructions() {
   return OkStatus();
 }
 
-Status HloModuleGroupMetadata::AddCompanion(HloInstruction* instruction1,
-                                            HloInstruction* instruction2) {
+absl::Status HloModuleGroupMetadata::AddCompanion(
+    HloInstruction* instruction1, HloInstruction* instruction2) {
   TF_RET_CHECK(instruction1->opcode() == HloOpcode::kWhile ||
                instruction1->opcode() == HloOpcode::kConditional ||
                instruction1->opcode() == HloOpcode::kCall);
@@ -449,7 +449,7 @@ Status HloModuleGroupMetadata::AddCompanion(HloInstruction* instruction1,
   return OkStatus();
 }
 
-Status HloModuleGroupMetadata::VerifyChannelInstructions() {
+absl::Status HloModuleGroupMetadata::VerifyChannelInstructions() {
   for (const Channel& channel : channels_) {
     if (channel.send == nullptr) {
       return FailedPrecondition("missing send for id : %d", channel.id);
@@ -530,7 +530,7 @@ Status HloModuleGroupMetadata::VerifyChannelInstructions() {
   return OkStatus();
 }
 
-Status HloModuleGroupMetadata::CheckCommunicatingInstruction(
+absl::Status HloModuleGroupMetadata::CheckCommunicatingInstruction(
     HloInstruction* instruction) const {
   HloComputation* computation = instruction->parent();
   const HloModule* module = computation->parent();
diff --git a/third_party/xla/xla/service/hlo_module_group_metadata.h b/third_party/xla/xla/service/hlo_module_group_metadata.h
index f977b19ab2ed77..1b19d0b069c359 100644
--- a/third_party/xla/xla/service/hlo_module_group_metadata.h
+++ b/third_party/xla/xla/service/hlo_module_group_metadata.h
@@ -215,11 +215,11 @@ class HloModuleGroupMetadata {
   }
 
  private:
-  Status Build();
+  absl::Status Build();
 
   // Record all channel instructions, cross-module AllReduce instructions, and
   // While/Conditional/Call instructions.
-  Status RecordInstructions();
+  absl::Status RecordInstructions();
 
   // Verifies the given HloModules are well-formed and follow the specification,
   // in particular with respect to using channel instructions.
@@ -229,20 +229,20 @@ class HloModuleGroupMetadata {
   // * The nest level of channel instructions match.
   // * Channel instructions are used in allowed computations, i.e., in the
   //   entry computation of the module or condition/body of While computations.
-  Status VerifyChannelInstructions();
+  absl::Status VerifyChannelInstructions();
 
   // Adds metadata that the given two instructions are companions.
-  Status AddCompanion(HloInstruction* instruction1,
-                      HloInstruction* instruction2);
+  absl::Status AddCompanion(HloInstruction* instruction1,
+                            HloInstruction* instruction2);
 
   // Checks whether a communicating instruction is placed in a valid position
   // within the graph.
-  Status CheckCommunicatingInstruction(HloInstruction* instruction) const;
+  absl::Status CheckCommunicatingInstruction(HloInstruction* instruction) const;
 
   // Performs a consistency check on the companion sets built for the input
   // modules. Checks that each instruction in a companion set is in a different
   // module/device.
-  Status VerifyCompanionSets() const;
+  absl::Status VerifyCompanionSets() const;
 
   // Retrieves a pointer to the stored TrackedInstruction associated with a
   // tracked computation, or nullptr in case such computation is not tracked.
diff --git a/third_party/xla/xla/service/hlo_module_group_util.cc b/third_party/xla/xla/service/hlo_module_group_util.cc
index 362c3ce46c3e31..3c199c44fd1cba 100644
--- a/third_party/xla/xla/service/hlo_module_group_util.cc
+++ b/third_party/xla/xla/service/hlo_module_group_util.cc
@@ -235,10 +235,9 @@ std::string HloModuleGroupUtil::CycleToString(
   return absl::StrJoin(names, " --> ");
 }
 
-Status HloModuleGroupUtil::VisitTopologicalOrder(VisitStates* visit_state,
-                                                 VisitFunction visit_function,
-                                                 HloInstruction* root,
-                                                 bool send_recv_as_one_group) {
+absl::Status HloModuleGroupUtil::VisitTopologicalOrder(
+    VisitStates* visit_state, VisitFunction visit_function,
+    HloInstruction* root, bool send_recv_as_one_group) {
   // Stack of HLO instructions visited in DFS order.
   std::stack<HloInstruction*> stack;
   stack.push(root);
@@ -332,7 +331,7 @@ Status HloModuleGroupUtil::VisitTopologicalOrder(VisitStates* visit_state,
   return OkStatus();
 }
 
-Status HloModuleGroupUtil::VerifyComputations(
+absl::Status HloModuleGroupUtil::VerifyComputations(
     absl::Span<HloComputation* const> computations) {
   auto visit_function =
       [&](HloInstruction* instruction,
diff --git a/third_party/xla/xla/service/hlo_module_group_util.h b/third_party/xla/xla/service/hlo_module_group_util.h
index 9d5e9af41e478b..dc084615791950 100644
--- a/third_party/xla/xla/service/hlo_module_group_util.h
+++ b/third_party/xla/xla/service/hlo_module_group_util.h
@@ -68,7 +68,7 @@ class HloModuleGroupUtil {
 
   // Function called on each instruction group during the DFS traversal. See the
   // comment for VisitTopologicalOrder()).
-  using VisitFunction = absl::FunctionRef<Status(
+  using VisitFunction = absl::FunctionRef<absl::Status(
       HloInstruction* hlo,
       const std::vector<HloInstruction*>& instruction_group)>;
 
@@ -91,13 +91,14 @@ class HloModuleGroupUtil {
   // * send_recv_as_one_group: if true, treat (Recv, Send, RecvDone, SendDone)
   // as one group.
   using VisitStates = absl::flat_hash_map<HloInstruction*, VisitState>;
-  Status VisitTopologicalOrder(VisitStates* visit_state,
-                               VisitFunction visit_function,
-                               HloInstruction* root,
-                               bool send_recv_as_one_group = false);
+  absl::Status VisitTopologicalOrder(VisitStates* visit_state,
+                                     VisitFunction visit_function,
+                                     HloInstruction* root,
+                                     bool send_recv_as_one_group = false);
 
   // Verifies that the computations are well-formed (e.g., no cycles).
-  Status VerifyComputations(absl::Span<HloComputation* const> computations);
+  absl::Status VerifyComputations(
+      absl::Span<HloComputation* const> computations);
 
   // Below Reachability utils resemble those in HloComputation, except that
   // they can handle instructions across multiple computations.
diff --git a/third_party/xla/xla/service/hlo_module_util.cc b/third_party/xla/xla/service/hlo_module_util.cc
index eca1d682061e6e..ef9fdae98bb48a 100644
--- a/third_party/xla/xla/service/hlo_module_util.cc
+++ b/third_party/xla/xla/service/hlo_module_util.cc
@@ -32,8 +32,8 @@ namespace xla {
 
 namespace {
 
-Status ValidateResultShape(const Shape& client_shape,
-                           const Shape& result_shape) {
+absl::Status ValidateResultShape(const Shape& client_shape,
+                                 const Shape& result_shape) {
   TF_RETURN_IF_ERROR(ShapeUtil::ValidateShapeWithOptionalLayout(client_shape));
   if (!ShapeUtil::Compatible(client_shape, result_shape)) {
     return InvalidArgument(
diff --git a/third_party/xla/xla/service/hlo_parser.cc b/third_party/xla/xla/service/hlo_parser.cc
index 53c5aac2c8adcb..9102e59601bf4d 100644
--- a/third_party/xla/xla/service/hlo_parser.cc
+++ b/third_party/xla/xla/service/hlo_parser.cc
@@ -244,7 +244,7 @@ class HloParserImpl : public HloParser {
 
   // Runs the parser and constructs the resulting HLO in the given (empty)
   // HloModule. Returns the error status in case an error occurred.
-  Status Run(HloModule* module) override;
+  absl::Status Run(HloModule* module) override;
 
   // Returns the error information.
   std::string GetError() const { return StrJoin(error_, "\n"); }
@@ -695,7 +695,7 @@ bool HloParserImpl::TokenError(absl::string_view msg) {
   return Error(lexer_.GetLoc(), msg);
 }
 
-Status HloParserImpl::Run(HloModule* module) {
+absl::Status HloParserImpl::Run(HloModule* module) {
   lexer_.Lex();
   if ((lexer_.GetKind() == TokKind::kw_HloModule) ||
       (lexer_.GetKind() == TokKind::kw_ENTRY) ||
@@ -1150,7 +1150,7 @@ bool HloParserImpl::ParseHloModule(HloModule* module,
   if (aliasing_data) {
     HloInputOutputAliasConfig alias_config(module->result_shape());
     for (auto& p : *aliasing_data) {
-      Status st =
+      absl::Status st =
           alias_config.SetUpAlias(p.first, p.second.parameter_number,
                                   p.second.parameter_index, p.second.kind);
       if (!st.ok()) {
@@ -1162,7 +1162,7 @@ bool HloParserImpl::ParseHloModule(HloModule* module,
   if (buffer_donor_data) {
     HloBufferDonorConfig buffer_donor_config;
     for (auto& p : *buffer_donor_data) {
-      Status st =
+      absl::Status st =
           buffer_donor_config.AddBufferDonor(p.param_number, p.param_index);
       if (!st.ok()) {
         return TokenError(st.message());
@@ -1424,7 +1424,7 @@ bool HloParserImpl::ParseInstructionRhs(HloComputation::Builder* builder,
   }
   if (predecessors) {
     for (auto* pre : *predecessors) {
-      Status status = pre->AddControlDependencyTo(instruction);
+      absl::Status status = pre->AddControlDependencyTo(instruction);
       if (!status.ok()) {
         return Error(name_loc, StrCat("error adding control dependency for: ",
                                       name, " status: ", status.ToString()));
diff --git a/third_party/xla/xla/service/hlo_parser.h b/third_party/xla/xla/service/hlo_parser.h
index 1a29800d6fedcd..bf95bfaa3f8129 100644
--- a/third_party/xla/xla/service/hlo_parser.h
+++ b/third_party/xla/xla/service/hlo_parser.h
@@ -93,7 +93,7 @@ class HloParser {
  public:
   // Runs the parser and constructs the resulting HLO in the given (empty)
   // HloModule. Returns the error status in case an error occurred.
-  virtual Status Run(HloModule* module) = 0;
+  virtual absl::Status Run(HloModule* module) = 0;
   virtual ~HloParser() {}
 
  private:
diff --git a/third_party/xla/xla/service/hlo_pass_fix.h b/third_party/xla/xla/service/hlo_pass_fix.h
index 3835a013b9031b..f360da9fff2c97 100644
--- a/third_party/xla/xla/service/hlo_pass_fix.h
+++ b/third_party/xla/xla/service/hlo_pass_fix.h
@@ -38,9 +38,10 @@ class HloPassFix : public Pass {
   template <typename... Args>
   explicit HloPassFix(Args&&... args) : Pass(args...) {}
 
-  Status RunOnChangedComputations(HloModule* module, RunState* outer_run_state,
-                                  const absl::flat_hash_set<absl::string_view>&
-                                      execution_threads) override {
+  absl::Status RunOnChangedComputations(
+      HloModule* module, RunState* outer_run_state,
+      const absl::flat_hash_set<absl::string_view>& execution_threads)
+      override {
     RunState run_state;
     run_state.changed_last_iteration = outer_run_state->changed_last_iteration;
     TF_RETURN_IF_ERROR(RunToFixPoint(module, &run_state, execution_threads));
@@ -85,7 +86,7 @@ class HloPassFix : public Pass {
   }
 
  private:
-  Status RunToFixPoint(
+  absl::Status RunToFixPoint(
       HloModule* module, RunState* run_state,
       const absl::flat_hash_set<absl::string_view>& execution_threads) {
     VLOG(3) << "Running HloPassFix on " << Pass::name();
@@ -108,7 +109,7 @@ class HloPassFix : public Pass {
     return OkStatus();
   }
 
-  Status RunOnChangedComputationsOnce(
+  absl::Status RunOnChangedComputationsOnce(
       HloModule* module, RunState* run_state,
       const absl::flat_hash_set<absl::string_view>& execution_threads) {
     // If Pass overrides RunOnChangedComputations, just forward to it.
diff --git a/third_party/xla/xla/service/hlo_pass_interface.h b/third_party/xla/xla/service/hlo_pass_interface.h
index 64fee1155c7f14..7e6aaa3804fd23 100644
--- a/third_party/xla/xla/service/hlo_pass_interface.h
+++ b/third_party/xla/xla/service/hlo_pass_interface.h
@@ -77,7 +77,7 @@ class HloPassInterface {
   //      absl::string_view name() const override { return "my-new-pass"; }
   //
   //      using HloPassInterface::Run;
-  //      StatusOr<bool> Run(
+  //      absl::StatusOr<bool> Run(
   //        HloModule* module,
   //        const absl::flat_hash_set<absl::string_view>& execution_threads)
   //        override;
@@ -97,7 +97,7 @@ class HloPassInterface {
   // NOTE: This is a temporary default implementation that conservatively treats
   // all computations as changed. Eventually all passes should override this
   // method instead of Run() and Run() will call into this method instead.
-  virtual Status RunOnChangedComputations(
+  virtual absl::Status RunOnChangedComputations(
       HloModule* module, RunState* run_state,
       const absl::flat_hash_set<absl::string_view>& execution_threads) {
     TF_ASSIGN_OR_RETURN(bool changed, Run(module, execution_threads));
@@ -126,7 +126,7 @@ class HloPassInterface {
   //      absl::string_view name() const override { return "my-new-pass"; }
   //
   //      using HloPassInterface::RunOnModuleGroup;
-  //      StatusOr<bool> RunOnModuleGroup(
+  //      absl::StatusOr<bool> RunOnModuleGroup(
   //        HloModuleGroup* module_group,
   //        const absl::flat_hash_set<absl::string_view>& execution_threads)
   //        override;
diff --git a/third_party/xla/xla/service/hlo_pass_pipeline.cc b/third_party/xla/xla/service/hlo_pass_pipeline.cc
index 021e3e8337a43e..f45f81cd186a42 100644
--- a/third_party/xla/xla/service/hlo_pass_pipeline.cc
+++ b/third_party/xla/xla/service/hlo_pass_pipeline.cc
@@ -40,7 +40,7 @@ namespace {
 void RecordPassStartMetadata(HloModule& module, const std::string& pass_name,
                              const std::string& pipeline_name) {
   module.metadata()->RecordPassStart();
-  // An HloPassMetadata was just created so Status should always be OK.
+  // An HloPassMetadata was just created so absl::Status should always be OK.
   TF_CHECK_OK(module.metadata()->set_current_pass_name(pass_name));
   TF_CHECK_OK(module.metadata()->set_current_pass_pipeline_name(pipeline_name));
 }
@@ -53,9 +53,9 @@ void RecordPassStartMetadata(HloModuleGroup& module_group,
   }
 }
 
-Status AttemptRecordPassEndMetadata(HloModule& module,
-                                    const std::string& pass_name,
-                                    bool module_changed) {
+absl::Status AttemptRecordPassEndMetadata(HloModule& module,
+                                          const std::string& pass_name,
+                                          bool module_changed) {
   // Module id is set here instead of RecordPassStartMetadata because it may
   // change in the middle of the pass, and we want the final id.
   TF_RETURN_IF_ERROR(
@@ -68,16 +68,16 @@ Status AttemptRecordPassEndMetadata(HloModule& module,
 
 void RecordPassEndMetadata(HloModule& module, const std::string& pass_name,
                            bool module_changed) {
-  Status status =
+  absl::Status status =
       AttemptRecordPassEndMetadata(module, pass_name, module_changed);
   if (!status.ok()) {
     LOG(FATAL) << status;
   }
 }
 
-Status AttemptRecordPassEndMetadata(HloModuleGroup& module_group,
-                                    const std::string& pass_name,
-                                    bool module_changed) {
+absl::Status AttemptRecordPassEndMetadata(HloModuleGroup& module_group,
+                                          const std::string& pass_name,
+                                          bool module_changed) {
   for (HloModule* module : module_group.modules()) {
     for (HloModule* other_module : module_group.modules()) {
       TF_RETURN_IF_ERROR(
@@ -92,7 +92,7 @@ Status AttemptRecordPassEndMetadata(HloModuleGroup& module_group,
 
 void RecordPassEndMetadata(HloModuleGroup& module_group,
                            const std::string& pass_name, bool module_changed) {
-  Status status =
+  absl::Status status =
       AttemptRecordPassEndMetadata(module_group, pass_name, module_changed);
   if (!status.ok()) {
     LOG(FATAL) << status;
@@ -102,7 +102,7 @@ void RecordPassEndMetadata(HloModuleGroup& module_group,
 }  // namespace
 
 template <typename HloT>
-Status HloPassPipeline::RunInvariantCheckers(
+absl::Status HloPassPipeline::RunInvariantCheckers(
     HloT* hlo, absl::string_view after_pass_name,
     const absl::flat_hash_set<absl::string_view>& execution_threads) {
   for (auto& invariant_checker : invariant_checkers_) {
@@ -271,7 +271,8 @@ void HloPassPipeline::MaybeDumpHloAndSaveFilenames(
     absl::string_view before_pass_name) {
   for (const std::string& filename : DumpHloModuleBetweenPassesIfEnabled(
            name(), before_pass_name, after_pass_name, module)) {
-    Status status = module.metadata()->add_current_pass_dump_filename(filename);
+    absl::Status status =
+        module.metadata()->add_current_pass_dump_filename(filename);
     if (!status.ok()) {
       LOG(FATAL) << status;
     }
diff --git a/third_party/xla/xla/service/hlo_pass_pipeline.h b/third_party/xla/xla/service/hlo_pass_pipeline.h
index 6ee1fc5e200566..e9fdec03dd21ea 100644
--- a/third_party/xla/xla/service/hlo_pass_pipeline.h
+++ b/third_party/xla/xla/service/hlo_pass_pipeline.h
@@ -114,11 +114,12 @@ class HloPassPipeline : public HloPassInterface {
   // `execution_threads`. Empty `execution_threads` means all execution threads
   // are included. HloT can be either HloModule or HloModuleGroup.
   template <typename HloT>
-  Status RunInvariantCheckers(HloT* hlo, absl::string_view after_pass_name) {
+  absl::Status RunInvariantCheckers(HloT* hlo,
+                                    absl::string_view after_pass_name) {
     return RunInvariantCheckers(hlo, after_pass_name, /*execution_threads=*/{});
   }
   template <typename HloT>
-  Status RunInvariantCheckers(
+  absl::Status RunInvariantCheckers(
       HloT* hlo, absl::string_view after_pass_name,
       const absl::flat_hash_set<absl::string_view>& execution_threads);
 
diff --git a/third_party/xla/xla/service/hlo_pass_pipeline_test.cc b/third_party/xla/xla/service/hlo_pass_pipeline_test.cc
index 252f9ccb7124e8..d5ad880f72c7a7 100644
--- a/third_party/xla/xla/service/hlo_pass_pipeline_test.cc
+++ b/third_party/xla/xla/service/hlo_pass_pipeline_test.cc
@@ -319,7 +319,7 @@ ENTRY main {
     pipeline.AddInvariantChecker<BarBlowerUpper>();
     pipeline.AddPass<FooToBarModulePass>();
 
-    Status status = pipeline.Run(module.get()).status();
+    absl::Status status = pipeline.Run(module.get()).status();
     ASSERT_IS_NOT_OK(status);
     EXPECT_THAT(status.message(),
                 ::testing::HasSubstr("Module has instruction named bar"));
@@ -331,7 +331,7 @@ ENTRY main {
     HloPassPipeline pipeline(TestName());
     pipeline.AddInvariantChecker<BarBlowerUpper>();
 
-    Status status = pipeline.Run(module.get()).status();
+    absl::Status status = pipeline.Run(module.get()).status();
     ASSERT_IS_NOT_OK(status);
     EXPECT_THAT(status.message(),
                 ::testing::HasSubstr("Module has instruction named bar"));
@@ -356,7 +356,7 @@ ENTRY main {
   HloPassPipeline pipeline(TestName());
   pipeline.AddPass<BazToQuxModuleGroupPass>();
 
-  Status status = pipeline.Run(module.get()).status();
+  absl::Status status = pipeline.Run(module.get()).status();
   ASSERT_IS_NOT_OK(status);
   EXPECT_THAT(
       status.message(),
diff --git a/third_party/xla/xla/service/hlo_rematerialization.cc b/third_party/xla/xla/service/hlo_rematerialization.cc
index a948cde80aa7d2..69cfeed1b05852 100644
--- a/third_party/xla/xla/service/hlo_rematerialization.cc
+++ b/third_party/xla/xla/service/hlo_rematerialization.cc
@@ -521,7 +521,7 @@ class MemoryUsageTracker {
   // EndInstruction) to accurately model memory usage. At BeginInstruction the
   // memory for the output value(s) of the current instruction is allocated. At
   // EndInstruction memory for dead operand(s) is freed.
-  Status BeginInstruction(Item* item);
+  absl::Status BeginInstruction(Item* item);
 
   int64_t RematerializationCost(const std::vector<Item*>& items,
                                 int64_t memory_reduced,
@@ -551,7 +551,7 @@ class MemoryUsageTracker {
   // Finishes the placement of the current instruction. This frees any dead
   // operands or dead result of the instruction. This must be called after
   // each call to BeginInstruction.
-  Status EndInstruction();
+  absl::Status EndInstruction();
 
   // Returns the number of bytes that the current memory usage will be reduced
   // if the given instruction is compact.
@@ -563,16 +563,18 @@ class MemoryUsageTracker {
   int64_t MemoryReducedIfRematerialized(
       absl::Span<const Item* const> items) const;
 
-  Status AddCompressInstructions(Item* original_item, Item* compressed_item,
-                                 Item* uncompressed_item);
+  absl::Status AddCompressInstructions(Item* original_item,
+                                       Item* compressed_item,
+                                       Item* uncompressed_item);
 
   // Adjusts memory usage to account for the rematerialization of
   // original_item for all remaining unplaced uses. The rematerialization
   // is remat_item. This method should be called after the HLO graph has
   // been transformed (rematerialization instruction created and connected
   // to uses).
-  Status AddRematerializedInstruction(Item* original_item, Item* remat_item,
-                                      absl::Span<Item*> indirect_users);
+  absl::Status AddRematerializedInstruction(Item* original_item,
+                                            Item* remat_item,
+                                            absl::Span<Item*> indirect_users);
 
   // Given a list of uses return two lists where one is the ones which are
   // placed and the other is ones which are not yet placed.
@@ -583,11 +585,11 @@ class MemoryUsageTracker {
   // Given the newly created instructions for host memory offload, create new
   // buffers, link their uses to their users, and update the current memory
   // usage.
-  Status AddHostOffloadCopyInstructions(Item* original_item,
-                                        Item* copy_start_to_host_item,
-                                        Item* copy_done_to_host_item,
-                                        Item* copy_start_to_device_item,
-                                        Item* copy_done_to_device_item);
+  absl::Status AddHostOffloadCopyInstructions(Item* original_item,
+                                              Item* copy_start_to_host_item,
+                                              Item* copy_done_to_host_item,
+                                              Item* copy_start_to_device_item,
+                                              Item* copy_done_to_device_item);
 
   // Counts the bytes that this item occupies by summing up the buffers defined
   // by this item. If only_count_unplaced_users is true, only count users of
@@ -714,7 +716,7 @@ class MemoryUsageTracker {
 
   // Adjust our tracked memory usage as a result of this item going out of
   // scope.
-  Status CountFreedMemory(Item* item);
+  absl::Status CountFreedMemory(Item* item);
 
   // Buffers have users and users have buffers used, this function resolves
   // outstanding issues in that bidirectional dependency.
@@ -940,7 +942,7 @@ void MemoryUsageTracker::CountAllocatedMemory(Item* item) {
   }
 }
 
-Status MemoryUsageTracker::CountFreedMemory(Item* item) {
+absl::Status MemoryUsageTracker::CountFreedMemory(Item* item) {
   for (BufferId buffer_id : item->buffers_used) {
     Buffer& buffer = buffers_.at(buffer_id);
     buffer.unfinished_user_count--;
@@ -969,7 +971,7 @@ Status MemoryUsageTracker::CountFreedMemory(Item* item) {
   return OkStatus();
 }
 
-Status MemoryUsageTracker::BeginInstruction(Item* item) {
+absl::Status MemoryUsageTracker::BeginInstruction(Item* item) {
   const HloInstruction* instruction = item->instruction;
   VLOG(3) << "BeginInstruction " << instruction->name();
   TF_RET_CHECK(in_progress_item_ == nullptr);
@@ -991,7 +993,7 @@ Status MemoryUsageTracker::BeginInstruction(Item* item) {
   return OkStatus();
 }
 
-Status MemoryUsageTracker::EndInstruction() {
+absl::Status MemoryUsageTracker::EndInstruction() {
   TF_RET_CHECK(in_progress_item_ != nullptr);
   VLOG(3) << "EndInstruction " << in_progress_item_->instruction->name();
 
@@ -1116,9 +1118,8 @@ void MemoryUsageTracker::ReplaceUsesInUsersOfBuffer(Buffer& buffer,
   }
 }
 
-Status MemoryUsageTracker::AddCompressInstructions(Item* original_item,
-                                                   Item* compressed_item,
-                                                   Item* uncompressed_item) {
+absl::Status MemoryUsageTracker::AddCompressInstructions(
+    Item* original_item, Item* compressed_item, Item* uncompressed_item) {
   CHECK(original_item->placed)
       << "Compressing instruction, but the original is not yet placed.";
   CHECK_EQ(original_item->buffers_output.size(), 1)
@@ -1185,7 +1186,7 @@ Status MemoryUsageTracker::AddCompressInstructions(Item* original_item,
   return OkStatus();
 }
 
-Status MemoryUsageTracker::AddRematerializedInstruction(
+absl::Status MemoryUsageTracker::AddRematerializedInstruction(
     Item* original_item, Item* remat_item, absl::Span<Item*> indirect_users) {
   VLOG(3) << "AddRematerializedInstruction: original_instruction = "
           << original_item->instruction->name()
@@ -1330,7 +1331,7 @@ Status MemoryUsageTracker::AddRematerializedInstruction(
   return OkStatus();
 }
 
-Status MemoryUsageTracker::AddHostOffloadCopyInstructions(
+absl::Status MemoryUsageTracker::AddHostOffloadCopyInstructions(
     Item* original_item, Item* copy_start_to_host_item,
     Item* copy_done_to_host_item, Item* copy_start_to_device_item,
     Item* copy_done_to_device_item) {
@@ -2925,7 +2926,8 @@ absl::StatusOr<bool> HloRematerialization::Run(
   // Compute peak memory usage of all computations in the module called in a
   // sequential context.
   TF_RETURN_IF_ERROR(call_graph_->VisitNodes(
-      [this, module, &execution_threads](const CallGraphNode& node) -> Status {
+      [this, module,
+       &execution_threads](const CallGraphNode& node) -> absl::Status {
         if (node.context() == CallContext::kControlFlow &&
             HloInstruction::IsThreadIncluded(
                 node.computation()->execution_thread(), execution_threads)) {
diff --git a/third_party/xla/xla/service/hlo_replication_analysis.cc b/third_party/xla/xla/service/hlo_replication_analysis.cc
index e8dd9c7393da47..e8b19284f4ca10 100644
--- a/third_party/xla/xla/service/hlo_replication_analysis.cc
+++ b/third_party/xla/xla/service/hlo_replication_analysis.cc
@@ -399,7 +399,7 @@ bool HloReplicationAnalysis::ComputeHloReplicationOnComputation(
   return changed;
 }
 
-Status HloReplicationAnalysis::ComputeHloReplication() {
+absl::Status HloReplicationAnalysis::ComputeHloReplication() {
   // Add entry parameters to the above sets according to user annotation.
   // Replicated modules read from `parameter_replicated_at_leaf_buffers` whereas
   // SPMD partitioned modules read from HloSharding attributes.
@@ -410,7 +410,7 @@ Status HloReplicationAnalysis::ComputeHloReplication() {
                                          HloReplication::UniqueOnAllDevices());
     const auto& replication = param->parameter_replicated_at_leaf_buffers();
     int leaf_index = 0;
-    Status status = ShapeUtil::ForEachSubshapeWithStatus(
+    absl::Status status = ShapeUtil::ForEachSubshapeWithStatus(
         param->shape(), [&](const Shape& subshape, const ShapeIndex& index) {
           if (!ShapeUtil::IsLeafIndex(param->shape(), index)) {
             return OkStatus();
diff --git a/third_party/xla/xla/service/hlo_replication_analysis.h b/third_party/xla/xla/service/hlo_replication_analysis.h
index e6f680214d7992..c16a1093baf9dd 100644
--- a/third_party/xla/xla/service/hlo_replication_analysis.h
+++ b/third_party/xla/xla/service/hlo_replication_analysis.h
@@ -114,7 +114,7 @@ class HloReplicationAnalysis {
         support_partial_replication_(support_partial_replication) {}
 
   // Computes hlo_replication_.
-  Status ComputeHloReplication();
+  absl::Status ComputeHloReplication();
 
   // A helper function to recursively compute hlo_replication on a computation.
   // Returns whether hlo_replication_ is changed.
diff --git a/third_party/xla/xla/service/hlo_value_semantics_analysis.cc b/third_party/xla/xla/service/hlo_value_semantics_analysis.cc
index 093fd302b42adb..26317ebd471ae5 100644
--- a/third_party/xla/xla/service/hlo_value_semantics_analysis.cc
+++ b/third_party/xla/xla/service/hlo_value_semantics_analysis.cc
@@ -1746,7 +1746,7 @@ absl::Status HloValueSemanticsPropagation::HandleConditional(
   std::vector<HloValueSemantics> merged_semantics_leaves;
   TF_RETURN_IF_ERROR(semantics_tree_vec[0].ForEachElementWithStatus(
       [&](const ShapeIndex& index,
-          const HloValueSemantics* semantics) -> Status {
+          const HloValueSemantics* semantics) -> absl::Status {
         std::vector<HloValueSemantics> semantics_vector;
         for (size_t i = 0; i < semantics_tree_vec.size(); ++i) {
           semantics_vector.push_back(
diff --git a/third_party/xla/xla/service/hlo_value_semantics_analysis.h b/third_party/xla/xla/service/hlo_value_semantics_analysis.h
index 48b0bbd68e1809..b005be32603728 100644
--- a/third_party/xla/xla/service/hlo_value_semantics_analysis.h
+++ b/third_party/xla/xla/service/hlo_value_semantics_analysis.h
@@ -61,8 +61,8 @@ class HloPreOrderDFS {
  public:
   HloPreOrderDFS() = default;
   ~HloPreOrderDFS() = default;
-  Status Run(const HloComputation& computation,
-             DfsHloVisitorBase<HloInstruction*>* visitor);
+  absl::Status Run(const HloComputation& computation,
+                   DfsHloVisitorBase<HloInstruction*>* visitor);
 
  private:
   bool IsReady(const HloInstruction* instruction) const;
@@ -93,43 +93,46 @@ class EinsumDepthAnalysis : public DfsHloVisitorWithDefault {
       const HloComputation& computation,
       const SendRecvGroupMap& send_recv_group_map);
   ~EinsumDepthAnalysis() override = default;
-  Status DefaultAction(HloInstruction* instruction) override;
-  Status HandleTuple(HloInstruction* tuple) override;
-  Status HandleGetTupleElement(HloInstruction* get_tuple_element) override;
-  Status HandleDot(HloInstruction* dot) override;
-  Status HandleConvolution(HloInstruction* convolution) override;
-  Status HandleCall(HloInstruction* call) override;
-  Status HandleFusion(HloInstruction* fusion) override;
-  Status HandleWhile(HloInstruction* xla_while) override;
-  Status HandleConditional(HloInstruction* conditional) override;
-  Status HandleAfterAll(HloInstruction* after_all) override;
-  Status HandleSend(HloInstruction* send) override;
-  Status HandleRecv(HloInstruction* recv) override;
-  Status HandleSendDone(HloInstruction* send_done) override;
-  Status HandleRecvDone(HloInstruction* recv_done) override;
-  Status HandleAllReduce(HloInstruction* all_reduce) override;
-  Status HandleAsyncStart(HloInstruction* async_start) override;
-  Status HandleAsyncDone(HloInstruction* async_done) override;
+  absl::Status DefaultAction(HloInstruction* instruction) override;
+  absl::Status HandleTuple(HloInstruction* tuple) override;
+  absl::Status HandleGetTupleElement(
+      HloInstruction* get_tuple_element) override;
+  absl::Status HandleDot(HloInstruction* dot) override;
+  absl::Status HandleConvolution(HloInstruction* convolution) override;
+  absl::Status HandleCall(HloInstruction* call) override;
+  absl::Status HandleFusion(HloInstruction* fusion) override;
+  absl::Status HandleWhile(HloInstruction* xla_while) override;
+  absl::Status HandleConditional(HloInstruction* conditional) override;
+  absl::Status HandleAfterAll(HloInstruction* after_all) override;
+  absl::Status HandleSend(HloInstruction* send) override;
+  absl::Status HandleRecv(HloInstruction* recv) override;
+  absl::Status HandleSendDone(HloInstruction* send_done) override;
+  absl::Status HandleRecvDone(HloInstruction* recv_done) override;
+  absl::Status HandleAllReduce(HloInstruction* all_reduce) override;
+  absl::Status HandleAsyncStart(HloInstruction* async_start) override;
+  absl::Status HandleAsyncDone(HloInstruction* async_done) override;
   const EinsumDepthMap& GetEinsumDepthMap() const { return einsum_depth_map_; }
 
  private:
   explicit EinsumDepthAnalysis(const SendRecvGroupMap& send_recv_group_map)
       : send_recv_group_map_(&send_recv_group_map) {}
-  Status RunInternal(const HloComputation& computation,
-                     const std::optional<ShapeTree<int>>& root_depth);
+  absl::Status RunInternal(const HloComputation& computation,
+                           const std::optional<ShapeTree<int>>& root_depth);
   ShapeTree<int>& GetOrCreateDepthTree(const HloInstruction* instruction);
   ShapeTree<int>& GetDepthTreeOrDie(const HloInstruction* instruction);
-  Status SetInstructionDepth(const HloInstruction* instruction, int depth);
-  Status SetInstructionDepth(const HloInstruction* instruction,
-                             const ShapeTree<int>& depth);
-  Status SetInstructionDepthFromTupleDepth(
+  absl::Status SetInstructionDepth(const HloInstruction* instruction,
+                                   int depth);
+  absl::Status SetInstructionDepth(const HloInstruction* instruction,
+                                   const ShapeTree<int>& depth);
+  absl::Status SetInstructionDepthFromTupleDepth(
       const HloInstruction* instruction, const ShapeTree<int>& tuple_depth_tree,
       int tuple_index);
-  Status HandleDepthIncrementInstruction(HloInstruction* instruction);
-  Status HandleCalledComputation(const HloComputation& called_computation,
-                                 const ShapeTree<int>& root_depth,
-                                 absl::Span<HloInstruction* const> operands);
-  Status HandleTupleLike(HloInstruction* tuple_like);
+  absl::Status HandleDepthIncrementInstruction(HloInstruction* instruction);
+  absl::Status HandleCalledComputation(
+      const HloComputation& called_computation,
+      const ShapeTree<int>& root_depth,
+      absl::Span<HloInstruction* const> operands);
+  absl::Status HandleTupleLike(HloInstruction* tuple_like);
   EinsumDepthMap einsum_depth_map_;
   const SendRecvGroupMap* const send_recv_group_map_;
 };
@@ -146,22 +149,23 @@ class EinsumHeightAnalysis : public DfsHloVisitorWithDefault {
       const HloComputation& computation,
       const SendRecvGroupMap& send_recv_group_map);
   ~EinsumHeightAnalysis() override = default;
-  Status DefaultAction(HloInstruction* instruction) override;
-  Status HandleTuple(HloInstruction* tuple) override;
-  Status HandleGetTupleElement(HloInstruction* get_tuple_element) override;
-  Status HandleDot(HloInstruction* dot) override;
-  Status HandleConvolution(HloInstruction* convolution) override;
-  Status HandleCall(HloInstruction* call) override;
-  Status HandleFusion(HloInstruction* fusion) override;
-  Status HandleWhile(HloInstruction* xla_while) override;
-  Status HandleConditional(HloInstruction* conditional) override;
-  Status HandleSend(HloInstruction* send) override;
-  Status HandleRecv(HloInstruction* recv) override;
-  Status HandleSendDone(HloInstruction* send_done) override;
-  Status HandleRecvDone(HloInstruction* recv_done) override;
-  Status HandleAllReduce(HloInstruction* all_reduce) override;
-  Status HandleAsyncStart(HloInstruction* async_start) override;
-  Status HandleAsyncDone(HloInstruction* async_done) override;
+  absl::Status DefaultAction(HloInstruction* instruction) override;
+  absl::Status HandleTuple(HloInstruction* tuple) override;
+  absl::Status HandleGetTupleElement(
+      HloInstruction* get_tuple_element) override;
+  absl::Status HandleDot(HloInstruction* dot) override;
+  absl::Status HandleConvolution(HloInstruction* convolution) override;
+  absl::Status HandleCall(HloInstruction* call) override;
+  absl::Status HandleFusion(HloInstruction* fusion) override;
+  absl::Status HandleWhile(HloInstruction* xla_while) override;
+  absl::Status HandleConditional(HloInstruction* conditional) override;
+  absl::Status HandleSend(HloInstruction* send) override;
+  absl::Status HandleRecv(HloInstruction* recv) override;
+  absl::Status HandleSendDone(HloInstruction* send_done) override;
+  absl::Status HandleRecvDone(HloInstruction* recv_done) override;
+  absl::Status HandleAllReduce(HloInstruction* all_reduce) override;
+  absl::Status HandleAsyncStart(HloInstruction* async_start) override;
+  absl::Status HandleAsyncDone(HloInstruction* async_done) override;
   const EinsumHeightMap& GetEinsumHeightMap() const {
     return einsum_height_map_;
   }
@@ -169,18 +173,20 @@ class EinsumHeightAnalysis : public DfsHloVisitorWithDefault {
  private:
   explicit EinsumHeightAnalysis(const SendRecvGroupMap& send_recv_group_map)
       : send_recv_group_map_(&send_recv_group_map) {}
-  Status RunInternal(const HloComputation& computation,
-                     absl::Span<HloInstruction* const> operands);
+  absl::Status RunInternal(const HloComputation& computation,
+                           absl::Span<HloInstruction* const> operands);
   ShapeTree<int>& GetOrCreateHeightTree(const HloInstruction* instruction);
   ShapeTree<int>& GetHeightTreeOrDie(const HloInstruction* instruction);
   bool HasHeightFor(const HloInstruction* instruction) const;
-  Status SetInstructionHeight(const HloInstruction* instruction, int height);
-  Status SetInstructionHeight(const HloInstruction* instruction,
-                              const ShapeTree<int>& height);
-  Status HandleHeightIncrementInstruction(HloInstruction* instruction);
-  Status HandleCalledComputation(const HloComputation& computation,
-                                 absl::Span<HloInstruction* const> operands);
-  Status HandleTupleLike(HloInstruction* tuple_like);
+  absl::Status SetInstructionHeight(const HloInstruction* instruction,
+                                    int height);
+  absl::Status SetInstructionHeight(const HloInstruction* instruction,
+                                    const ShapeTree<int>& height);
+  absl::Status HandleHeightIncrementInstruction(HloInstruction* instruction);
+  absl::Status HandleCalledComputation(
+      const HloComputation& computation,
+      absl::Span<HloInstruction* const> operands);
+  absl::Status HandleTupleLike(HloInstruction* tuple_like);
 
   EinsumHeightMap einsum_height_map_;
   const SendRecvGroupMap* const send_recv_group_map_;
@@ -269,19 +275,20 @@ class HloValueSemanticsAnalysis {
  protected:
   friend class HloValueSemanticsPropagation;
   explicit HloValueSemanticsAnalysis(const HloModule& module);
-  virtual Status InitializeEinsumDepth();
-  virtual Status InitializeEinsumHeight();
+  virtual absl::Status InitializeEinsumDepth();
+  virtual absl::Status InitializeEinsumHeight();
   // We match send and recv HLOs to propagate semantics from send to recv.
   virtual void InitializeSendRecvGroups();
   void AnnotateWeights();
 
   // Infer semantics for all instructions in the computation. Computation
   // parameters are assigned the semantics of the corresponding operand.
-  Status RunOnComputation(const HloComputation& computation,
-                          absl::Span<const HloInstruction* const> operands);
+  absl::Status RunOnComputation(
+      const HloComputation& computation,
+      absl::Span<const HloInstruction* const> operands);
   // Same as the above RunOnComputation, but computation parameters have
   // already been assigned with semantics.
-  virtual Status RunOnComputation(const HloComputation& computation);
+  virtual absl::Status RunOnComputation(const HloComputation& computation);
   HloValueSemantics::Id NextId();
   const HloValueSemantics* NewHloValueSemantics(HloValueSemanticLabel label,
                                                 const HloPosition& origin);
@@ -314,50 +321,52 @@ class HloValueSemanticsAnalysis {
 class HloValueSemanticsPropagation : public DfsHloVisitorWithDefault {
  public:
   explicit HloValueSemanticsPropagation(HloValueSemanticsAnalysis* analysis);
-  Status Run(const HloComputation& computation);
+  absl::Status Run(const HloComputation& computation);
   // Infer the output semantics from all operands of the instruction.
-  Status DefaultAction(HloInstruction* instruction) override;
-  Status HandleParameter(HloInstruction* parameter) override;
-  Status HandleConstant(HloInstruction* constant) override;
-  Status HandleIota(HloInstruction* iota) override;
-  Status HandlePartitionId(HloInstruction* partition_id) override;
-  Status HandleReplicaId(HloInstruction* replica_id) override;
-  Status HandleClamp(HloInstruction* clamp) override;
-  Status HandleTuple(HloInstruction* tuple) override;
-  Status HandleGetTupleElement(HloInstruction* get_tuple_element) override;
-  Status HandleCall(HloInstruction* call) override;
-  Status HandleFusion(HloInstruction* fusion) override;
-  Status HandleCustomCall(HloInstruction* custom_call) override;
-  Status HandleWhile(HloInstruction* xla_while) override;
-  Status HandleConditional(HloInstruction* conditional) override;
-  Status HandleSelect(HloInstruction* select) override;
-  Status HandleConcatenate(HloInstruction* concatenate) override;
-  Status HandleDynamicSlice(HloInstruction* dynamic_slice) override;
-  Status HandleDynamicUpdateSlice(
+  absl::Status DefaultAction(HloInstruction* instruction) override;
+  absl::Status HandleParameter(HloInstruction* parameter) override;
+  absl::Status HandleConstant(HloInstruction* constant) override;
+  absl::Status HandleIota(HloInstruction* iota) override;
+  absl::Status HandlePartitionId(HloInstruction* partition_id) override;
+  absl::Status HandleReplicaId(HloInstruction* replica_id) override;
+  absl::Status HandleClamp(HloInstruction* clamp) override;
+  absl::Status HandleTuple(HloInstruction* tuple) override;
+  absl::Status HandleGetTupleElement(
+      HloInstruction* get_tuple_element) override;
+  absl::Status HandleCall(HloInstruction* call) override;
+  absl::Status HandleFusion(HloInstruction* fusion) override;
+  absl::Status HandleCustomCall(HloInstruction* custom_call) override;
+  absl::Status HandleWhile(HloInstruction* xla_while) override;
+  absl::Status HandleConditional(HloInstruction* conditional) override;
+  absl::Status HandleSelect(HloInstruction* select) override;
+  absl::Status HandleConcatenate(HloInstruction* concatenate) override;
+  absl::Status HandleDynamicSlice(HloInstruction* dynamic_slice) override;
+  absl::Status HandleDynamicUpdateSlice(
       HloInstruction* dynamic_update_slice) override;
-  Status HandleCopyStart(HloInstruction* copy_start) override;
-  Status HandleCopyDone(HloInstruction* copy_done) override;
-  Status HandleAllGatherStart(HloInstruction* all_gather_start) override;
-  Status HandleAllGatherDone(HloInstruction* all_gather_done) override;
-  Status HandleCollectivePermuteStart(
+  absl::Status HandleCopyStart(HloInstruction* copy_start) override;
+  absl::Status HandleCopyDone(HloInstruction* copy_done) override;
+  absl::Status HandleAllGatherStart(HloInstruction* all_gather_start) override;
+  absl::Status HandleAllGatherDone(HloInstruction* all_gather_done) override;
+  absl::Status HandleCollectivePermuteStart(
       HloInstruction* collective_permute_start) override;
-  Status HandleCollectivePermuteDone(
+  absl::Status HandleCollectivePermuteDone(
       HloInstruction* collective_permute_done) override;
-  Status HandleGather(HloInstruction* gather) override;
-  Status HandleScatter(HloInstruction* scatter) override;
-  Status HandleAfterAll(HloInstruction* after_all) override;
-  Status HandleAllReduce(HloInstruction* all_reduce) override;
-  Status HandleAsyncStart(HloInstruction* async_start) override;
-  Status HandleAsyncDone(HloInstruction* async_done) override;
-  Status HandleInfeed(HloInstruction* infeed) override;
-  Status HandleOutfeed(HloInstruction* outfeed) override;
-  Status HandleDomain(HloInstruction* domain) override;
-  Status HandleOptimizationBarrier(HloInstruction* opt_barrier) override;
-  Status HandleRngBitGenerator(HloInstruction* rng_bit_generator) override;
-  Status HandleSend(HloInstruction* send) override;
-  Status HandleRecv(HloInstruction* recv) override;
-  Status HandleSendDone(HloInstruction* send_done) override;
-  Status HandleRecvDone(HloInstruction* recv_done) override;
+  absl::Status HandleGather(HloInstruction* gather) override;
+  absl::Status HandleScatter(HloInstruction* scatter) override;
+  absl::Status HandleAfterAll(HloInstruction* after_all) override;
+  absl::Status HandleAllReduce(HloInstruction* all_reduce) override;
+  absl::Status HandleAsyncStart(HloInstruction* async_start) override;
+  absl::Status HandleAsyncDone(HloInstruction* async_done) override;
+  absl::Status HandleInfeed(HloInstruction* infeed) override;
+  absl::Status HandleOutfeed(HloInstruction* outfeed) override;
+  absl::Status HandleDomain(HloInstruction* domain) override;
+  absl::Status HandleOptimizationBarrier(HloInstruction* opt_barrier) override;
+  absl::Status HandleRngBitGenerator(
+      HloInstruction* rng_bit_generator) override;
+  absl::Status HandleSend(HloInstruction* send) override;
+  absl::Status HandleRecv(HloInstruction* recv) override;
+  absl::Status HandleSendDone(HloInstruction* send_done) override;
+  absl::Status HandleRecvDone(HloInstruction* recv_done) override;
 
  protected:
   HloValueSemantics CopySemantics(const HloValueSemantics& semantics) const;
@@ -419,9 +428,9 @@ class HloValueSemanticsPropagation : public DfsHloVisitorWithDefault {
   absl::StatusOr<HloValueSemantics> ComputeSemanticsFromOperands(
       HloInstruction* instruction, absl::Span<const int64_t> operand_indices,
       absl::Span<const ShapeIndex> operand_shape_indices = {}) const;
-  Status HandleTupleLike(HloInstruction* tuple_like);
-  Status HandleCollectiveOrCopyStart(HloInstruction* op_start);
-  Status HandleCollectiveOrCopyDone(HloInstruction* op_done);
+  absl::Status HandleTupleLike(HloInstruction* tuple_like);
+  absl::Status HandleCollectiveOrCopyStart(HloInstruction* op_start);
+  absl::Status HandleCollectiveOrCopyDone(HloInstruction* op_done);
   HloValueSemanticsAnalysis* analysis_;
 };
 
diff --git a/third_party/xla/xla/service/hlo_verifier.cc b/third_party/xla/xla/service/hlo_verifier.cc
index 523d49ed19a83b..841f5ee0a93548 100644
--- a/third_party/xla/xla/service/hlo_verifier.cc
+++ b/third_party/xla/xla/service/hlo_verifier.cc
@@ -72,7 +72,7 @@ bool IsCallerInstruction(HloInstruction* hlo) {
   return HloInstruction::MightHaveCalledComputations(hlo->opcode());
 }
 
-Status CheckOperandCount(const HloInstruction* hlo, int expected) {
+absl::Status CheckOperandCount(const HloInstruction* hlo, int expected) {
   if (hlo->operand_count() != expected) {
     return Internal("Expected %d operands for %s instruction: %s", expected,
                     HloOpcodeString(hlo->opcode()), hlo->ToString());
@@ -106,8 +106,8 @@ int64_t GetSubgroupSize(HloCollectiveInstruction* hlo,
   }
 }
 
-Status CheckNestedComputationThreadNameEqual(const HloComputation* comp,
-                                             bool skip_nested_async_op_check) {
+absl::Status CheckNestedComputationThreadNameEqual(
+    const HloComputation* comp, bool skip_nested_async_op_check) {
   for (const HloInstruction* instr : comp->instructions()) {
     if (skip_nested_async_op_check && instr->IsAsynchronous()) {
       continue;
@@ -127,7 +127,7 @@ Status CheckNestedComputationThreadNameEqual(const HloComputation* comp,
 }
 }  // namespace
 
-/*static*/ Status ShapeVerifier::CheckParameterCount(
+/*static*/ absl::Status ShapeVerifier::CheckParameterCount(
     const HloInstruction* calling_instruction,
     const HloComputation* computation, int expected) {
   if (computation->num_parameters() != expected) {
@@ -139,7 +139,7 @@ Status CheckNestedComputationThreadNameEqual(const HloComputation* comp,
   return OkStatus();
 }
 
-Status ShapeVerifier::Preprocess(HloInstruction* hlo) {
+absl::Status ShapeVerifier::Preprocess(HloInstruction* hlo) {
   if (!hlo->called_computations().empty() && !IsCallerInstruction(hlo)) {
     return Internal(
         "Called computations specified for non-caller instruction %s",
@@ -156,23 +156,23 @@ Status ShapeVerifier::Preprocess(HloInstruction* hlo) {
   return OkStatus();
 }
 
-Status ShapeVerifier::HandleElementwiseUnary(HloInstruction* hlo) {
+absl::Status ShapeVerifier::HandleElementwiseUnary(HloInstruction* hlo) {
   return CheckUnaryShape(hlo);
 }
 
-Status ShapeVerifier::HandleElementwiseBinary(HloInstruction* hlo) {
+absl::Status ShapeVerifier::HandleElementwiseBinary(HloInstruction* hlo) {
   return CheckBinaryShape(hlo);
 }
 
-Status ShapeVerifier::HandleClamp(HloInstruction* clamp) {
+absl::Status ShapeVerifier::HandleClamp(HloInstruction* clamp) {
   return CheckTernaryShape(clamp);
 }
 
-Status ShapeVerifier::HandleSelect(HloInstruction* select) {
+absl::Status ShapeVerifier::HandleSelect(HloInstruction* select) {
   return CheckTernaryShape(select);
 }
 
-Status ShapeVerifier::HandleConcatenate(HloInstruction* concatenate) {
+absl::Status ShapeVerifier::HandleConcatenate(HloInstruction* concatenate) {
   std::vector<const Shape*> operand_shapes;
   for (const HloInstruction* operand : concatenate->operands()) {
     operand_shapes.push_back(&operand->shape());
@@ -182,30 +182,30 @@ Status ShapeVerifier::HandleConcatenate(HloInstruction* concatenate) {
                         operand_shapes, concatenate->concatenate_dimension()));
 }
 
-Status ShapeVerifier::HandleConvert(HloInstruction* convert) {
+absl::Status ShapeVerifier::HandleConvert(HloInstruction* convert) {
   return CheckShape(convert, ShapeInference::InferConvertShape(
                                  convert->operand(0)->shape(),
                                  convert->shape().element_type()));
 }
 
-Status ShapeVerifier::HandleBitcastConvert(HloInstruction* convert) {
+absl::Status ShapeVerifier::HandleBitcastConvert(HloInstruction* convert) {
   return CheckShape(convert, ShapeInference::InferBitcastConvertShape(
                                  convert->operand(0)->shape(),
                                  convert->shape().element_type()));
 }
 
-Status ShapeVerifier::HandleStochasticConvert(HloInstruction* convert) {
+absl::Status ShapeVerifier::HandleStochasticConvert(HloInstruction* convert) {
   return CheckShape(
       convert, ShapeInference::InferStochasticConvertShape(
                    convert->operand(0)->shape(), convert->operand(1)->shape(),
                    convert->shape().element_type()));
 }
 
-Status ShapeVerifier::HandleCopy(HloInstruction* copy) {
+absl::Status ShapeVerifier::HandleCopy(HloInstruction* copy) {
   return CheckUnaryShape(copy);
 }
 
-Status ShapeVerifier::HandleDot(HloInstruction* dot) {
+absl::Status ShapeVerifier::HandleDot(HloInstruction* dot) {
   auto sparsity = Cast<HloDotInstruction>(dot)->sparsity();
   TF_RETURN_IF_ERROR(
       CheckOperandCount(dot, HloDotInstruction::kOperands + sparsity.size()));
@@ -257,7 +257,7 @@ Status ShapeVerifier::HandleDot(HloInstruction* dot) {
   return CheckShape(dot, expected);
 }
 
-Status ShapeVerifier::HandleConvolution(HloInstruction* convolution) {
+absl::Status ShapeVerifier::HandleConvolution(HloInstruction* convolution) {
   TF_ASSIGN_OR_RETURN(
       Shape expected,
       ShapeInference::InferConvolveShape(
@@ -304,7 +304,7 @@ Status ShapeVerifier::HandleConvolution(HloInstruction* convolution) {
   return CheckShape(convolution, expected);
 }
 
-Status ShapeVerifier::HandleFft(HloInstruction* fft) {
+absl::Status ShapeVerifier::HandleFft(HloInstruction* fft) {
   TF_ASSIGN_OR_RETURN(
       const Shape expected,
       ShapeInference::InferFftShape(fft->operand(0)->shape(), fft->fft_type(),
@@ -312,7 +312,7 @@ Status ShapeVerifier::HandleFft(HloInstruction* fft) {
   return CheckShape(fft, expected);
 }
 
-Status ShapeVerifier::HandleTriangularSolve(HloInstruction* hlo) {
+absl::Status ShapeVerifier::HandleTriangularSolve(HloInstruction* hlo) {
   TF_ASSIGN_OR_RETURN(const Shape expected,
                       ShapeInference::InferTriangularSolveShape(
                           hlo->operand(0)->shape(), hlo->operand(1)->shape(),
@@ -320,14 +320,14 @@ Status ShapeVerifier::HandleTriangularSolve(HloInstruction* hlo) {
   return CheckShape(hlo, expected);
 }
 
-Status ShapeVerifier::HandleCholesky(HloInstruction* hlo) {
+absl::Status ShapeVerifier::HandleCholesky(HloInstruction* hlo) {
   TF_RETURN_IF_ERROR(CheckOperandCount(hlo, 1));
   TF_ASSIGN_OR_RETURN(const Shape expected, ShapeInference::InferCholeskyShape(
                                                 hlo->operand(0)->shape()));
   return CheckShape(hlo, expected);
 }
 
-Status ShapeVerifier::HandleOptimizationBarrier(HloInstruction* hlo) {
+absl::Status ShapeVerifier::HandleOptimizationBarrier(HloInstruction* hlo) {
   TF_RETURN_IF_ERROR(CheckOperandCount(hlo, 1));
   return CheckShape(hlo, hlo->operand(0)->shape());
 }
@@ -353,9 +353,9 @@ bool ShapeVerifier::ShapesSame(const Shape& a, const Shape& b,
 // In general, if replica groups is not empty, all replica groups should be of
 // the same size. The exception is all-reduce, where non-uniform replica groups
 // are allowed. This is controlled by `uniform_replica_group_size`.
-static Status CheckReplicaGroups(HloInstruction* hlo,
-                                 CollectiveOpGroupMode group_mode,
-                                 bool uniform_replica_group_size = true) {
+static absl::Status CheckReplicaGroups(HloInstruction* hlo,
+                                       CollectiveOpGroupMode group_mode,
+                                       bool uniform_replica_group_size = true) {
   if (!hlo->replica_groups().empty()) {
     absl::flat_hash_set<int64_t> replicas_seen;
     for (const ReplicaGroup& g : hlo->replica_groups()) {
@@ -427,8 +427,8 @@ static Status CheckReplicaGroups(HloInstruction* hlo,
   return OkStatus();
 }
 
-static Status CheckCommonAllGatherInvariants(HloInstruction* hlo,
-                                             int64_t* computed_shard_count) {
+static absl::Status CheckCommonAllGatherInvariants(
+    HloInstruction* hlo, int64_t* computed_shard_count) {
   auto ag = Cast<HloAllGatherInstruction>(hlo);
   CHECK_NE(computed_shard_count, nullptr) << "Expected a shard count as input";
   TF_ASSIGN_OR_RETURN(CollectiveOpGroupMode group_mode,
@@ -471,7 +471,7 @@ static Status CheckCommonAllGatherInvariants(HloInstruction* hlo,
   return OkStatus();
 }
 
-Status ShapeVerifier::HandleAllGather(HloInstruction* hlo) {
+absl::Status ShapeVerifier::HandleAllGather(HloInstruction* hlo) {
   auto ag = Cast<HloAllGatherInstruction>(hlo);
   int64_t shard_count;
   TF_RETURN_IF_ERROR(CheckCommonAllGatherInvariants(hlo, &shard_count));
@@ -484,7 +484,7 @@ Status ShapeVerifier::HandleAllGather(HloInstruction* hlo) {
               operand_shapes, ag->all_gather_dimension(), shard_count));
 }
 
-Status ShapeVerifier::HandleAllGatherStart(HloInstruction* hlo) {
+absl::Status ShapeVerifier::HandleAllGatherStart(HloInstruction* hlo) {
   auto ag = Cast<HloAllGatherInstruction>(hlo);
   int64_t shard_count;
   TF_RETURN_IF_ERROR(CheckCommonAllGatherInvariants(hlo, &shard_count));
@@ -497,12 +497,12 @@ Status ShapeVerifier::HandleAllGatherStart(HloInstruction* hlo) {
               operand_shapes, ag->all_gather_dimension(), shard_count));
 }
 
-Status ShapeVerifier::HandleAllGatherDone(HloInstruction* hlo) {
+absl::Status ShapeVerifier::HandleAllGatherDone(HloInstruction* hlo) {
   return CheckShape(
       hlo, ShapeInference::InferAllGatherDoneShape(hlo->operand(0)->shape()));
 }
 
-Status ShapeVerifier::HandleAllReduce(HloInstruction* hlo) {
+absl::Status ShapeVerifier::HandleAllReduce(HloInstruction* hlo) {
   auto ar = Cast<HloAllReduceInstruction>(hlo);
   TF_ASSIGN_OR_RETURN(CollectiveOpGroupMode group_mode,
                       GetCollectiveOpGroupMode(ar->channel_id().has_value(),
@@ -517,7 +517,7 @@ Status ShapeVerifier::HandleAllReduce(HloInstruction* hlo) {
   return CheckShape(hlo, ShapeInference::InferAllReduceShape(operand_shapes));
 }
 
-Status ShapeVerifier::HandleReduceScatter(HloInstruction* hlo) {
+absl::Status ShapeVerifier::HandleReduceScatter(HloInstruction* hlo) {
   auto ars = Cast<HloReduceScatterInstruction>(hlo);
   TF_ASSIGN_OR_RETURN(CollectiveOpGroupMode group_mode,
                       GetCollectiveOpGroupMode(ars->channel_id().has_value(),
@@ -557,7 +557,7 @@ Status ShapeVerifier::HandleReduceScatter(HloInstruction* hlo) {
                         operand_shapes, ars->scatter_dimension(), shard_count));
 }
 
-Status ShapeVerifier::HandleAllReduceStart(HloInstruction* hlo) {
+absl::Status ShapeVerifier::HandleAllReduceStart(HloInstruction* hlo) {
   auto ar = Cast<HloAllReduceInstruction>(hlo);
   TF_ASSIGN_OR_RETURN(CollectiveOpGroupMode group_mode,
                       GetCollectiveOpGroupMode(ar->channel_id().has_value(),
@@ -573,12 +573,12 @@ Status ShapeVerifier::HandleAllReduceStart(HloInstruction* hlo) {
                     ShapeInference::InferAllReduceStartShape(operand_shapes));
 }
 
-Status ShapeVerifier::HandleAllReduceDone(HloInstruction* hlo) {
+absl::Status ShapeVerifier::HandleAllReduceDone(HloInstruction* hlo) {
   return CheckShape(
       hlo, ShapeInference::InferAllReduceDoneShape(hlo->operand(0)->shape()));
 }
 
-Status ShapeVerifier::HandleAllToAll(HloInstruction* hlo) {
+absl::Status ShapeVerifier::HandleAllToAll(HloInstruction* hlo) {
   auto* all_to_all = Cast<HloAllToAllInstruction>(hlo);
   TF_ASSIGN_OR_RETURN(CollectiveOpGroupMode group_mode,
                       GetCollectiveOpGroupMode(
@@ -605,18 +605,18 @@ Status ShapeVerifier::HandleAllToAll(HloInstruction* hlo) {
   }
 }
 
-Status ShapeVerifier::HandlePartitionId(HloInstruction* hlo) {
+absl::Status ShapeVerifier::HandlePartitionId(HloInstruction* hlo) {
   return CheckShape(hlo, ShapeUtil::MakeShape(U32, {}));
 }
 
-Status ShapeVerifier::HandleReplicaId(HloInstruction* hlo) {
+absl::Status ShapeVerifier::HandleReplicaId(HloInstruction* hlo) {
   return CheckShape(hlo, ShapeUtil::MakeShape(U32, {}));
 }
 
 namespace {
 
-Status CheckBufferOffset(const Shape& buffer_shape,
-                         const Shape& buffer_offset_shape) {
+absl::Status CheckBufferOffset(const Shape& buffer_shape,
+                               const Shape& buffer_offset_shape) {
   if (!buffer_offset_shape.IsTuple()) {
     return Internal("Buffer offset is not tuple.");
   }
@@ -652,7 +652,7 @@ Status CheckBufferOffset(const Shape& buffer_shape,
   return OkStatus();
 }
 
-Status CheckInplaceCollectivePermute(HloInstruction* collective_permute) {
+absl::Status CheckInplaceCollectivePermute(HloInstruction* collective_permute) {
   if (collective_permute->operand_count() == 1) {
     return OkStatus();
   }
@@ -667,12 +667,12 @@ Status CheckInplaceCollectivePermute(HloInstruction* collective_permute) {
   const Shape& output_offset_shape = collective_permute->operand(3)->shape();
 
   if (input_buffer_shape.IsArray() && output_buffer_shape.IsArray()) {
-    Status check_input_buffer_offset =
+    absl::Status check_input_buffer_offset =
         CheckBufferOffset(input_buffer_shape, input_offset_shape);
     if (!check_input_buffer_offset.ok()) {
       return check_input_buffer_offset;
     }
-    Status check_output_buffer_offset =
+    absl::Status check_output_buffer_offset =
         CheckBufferOffset(output_buffer_shape, output_offset_shape);
     if (!check_output_buffer_offset.ok()) {
       return check_output_buffer_offset;
@@ -688,7 +688,7 @@ Status CheckInplaceCollectivePermute(HloInstruction* collective_permute) {
       return Internal("Unmatching input buffers and input offset.");
     }
     for (int i = 0; i < input_buffer_shape.tuple_shapes_size(); ++i) {
-      Status check_input_buffer_offset =
+      absl::Status check_input_buffer_offset =
           CheckBufferOffset(input_buffer_shape.tuple_shapes(i),
                             input_offset_shape.tuple_shapes(i));
       if (!check_input_buffer_offset.ok()) {
@@ -701,7 +701,7 @@ Status CheckInplaceCollectivePermute(HloInstruction* collective_permute) {
       return Internal("Unmatching output buffers and output offset.");
     }
     for (int i = 0; i < output_buffer_shape.tuple_shapes_size(); ++i) {
-      Status check_output_buffer_offset =
+      absl::Status check_output_buffer_offset =
           CheckBufferOffset(output_buffer_shape.tuple_shapes(i),
                             output_offset_shape.tuple_shapes(i));
       if (!check_output_buffer_offset.ok()) {
@@ -714,8 +714,8 @@ Status CheckInplaceCollectivePermute(HloInstruction* collective_permute) {
   return OkStatus();
 }
 
-Status CheckDuplicatedSourceOrTarget(HloInstruction* hlo,
-                                     CollectiveOpGroupMode group_mode) {
+absl::Status CheckDuplicatedSourceOrTarget(HloInstruction* hlo,
+                                           CollectiveOpGroupMode group_mode) {
   // A source or target cannot appear twice in the collective-permute's
   // source-target pairs. Also, based on the group formation mode, check if the
   // source and target IDs are within expected range.
@@ -795,7 +795,7 @@ Status CheckDuplicatedSourceOrTarget(HloInstruction* hlo,
 
 }  // namespace
 
-Status ShapeVerifier::HandleCollectiveBroadcast(HloInstruction* hlo) {
+absl::Status ShapeVerifier::HandleCollectiveBroadcast(HloInstruction* hlo) {
   std::vector<const Shape*> operand_shapes;
   for (const HloInstruction* operand : hlo->operands()) {
     operand_shapes.push_back(&operand->shape());
@@ -804,7 +804,7 @@ Status ShapeVerifier::HandleCollectiveBroadcast(HloInstruction* hlo) {
       hlo, ShapeInference::InferCollectiveBroadcastShape(operand_shapes));
 }
 
-Status ShapeVerifier::HandleCollectivePermute(HloInstruction* hlo) {
+absl::Status ShapeVerifier::HandleCollectivePermute(HloInstruction* hlo) {
   TF_ASSIGN_OR_RETURN(
       CollectiveOpGroupMode group_mode,
       GetCollectiveOpGroupMode(hlo->channel_id().has_value(),
@@ -819,7 +819,7 @@ Status ShapeVerifier::HandleCollectivePermute(HloInstruction* hlo) {
       hlo, ShapeInference::InferCollectivePermuteShape(operand_shapes));
 }
 
-Status ShapeVerifier::HandleCollectivePermuteStart(HloInstruction* hlo) {
+absl::Status ShapeVerifier::HandleCollectivePermuteStart(HloInstruction* hlo) {
   TF_ASSIGN_OR_RETURN(
       CollectiveOpGroupMode group_mode,
       GetCollectiveOpGroupMode(hlo->channel_id().has_value(),
@@ -839,20 +839,21 @@ Status ShapeVerifier::HandleCollectivePermuteStart(HloInstruction* hlo) {
                              operand_shapes, context_shapes));
 }
 
-Status ShapeVerifier::HandleCollectivePermuteDone(HloInstruction* hlo) {
+absl::Status ShapeVerifier::HandleCollectivePermuteDone(HloInstruction* hlo) {
   return CheckShape(hlo, ShapeInference::InferCollectivePermuteDoneShape(
                              hlo->operand(0)->shape()));
 }
 
-Status ShapeVerifier::HandleReducePrecision(HloInstruction* reduce_precision) {
+absl::Status ShapeVerifier::HandleReducePrecision(
+    HloInstruction* reduce_precision) {
   return CheckShape(reduce_precision, ShapeInference::InferReducePrecisionShape(
                                           reduce_precision->operand(0)->shape(),
                                           reduce_precision->exponent_bits(),
                                           reduce_precision->mantissa_bits()));
 }
 
-Status ShapeVerifier::CheckIsTokenOperand(const HloInstruction* instruction,
-                                          int64_t operand_no) {
+absl::Status ShapeVerifier::CheckIsTokenOperand(
+    const HloInstruction* instruction, int64_t operand_no) {
   const HloInstruction* token = instruction->operand(operand_no);
   if (!ShapeUtil::Equal(token->shape(), ShapeUtil::MakeTokenShape())) {
     return Internal(
@@ -863,7 +864,7 @@ Status ShapeVerifier::CheckIsTokenOperand(const HloInstruction* instruction,
   return OkStatus();
 }
 
-Status ShapeVerifier::CheckOperandAndParameter(
+absl::Status ShapeVerifier::CheckOperandAndParameter(
     const HloInstruction* instruction, int64_t operand_number,
     const HloComputation* computation, int64_t parameter_number) {
   const HloInstruction* operand = instruction->operand(operand_number);
@@ -877,7 +878,7 @@ Status ShapeVerifier::CheckOperandAndParameter(
   return OkStatus();
 }
 
-Status ShapeVerifier::HandleInfeed(HloInstruction* instruction) {
+absl::Status ShapeVerifier::HandleInfeed(HloInstruction* instruction) {
   HloInfeedInstruction* infeed = Cast<HloInfeedInstruction>(instruction);
   TF_RETURN_IF_ERROR(CheckIsTokenOperand(instruction, 0));
 
@@ -888,7 +889,7 @@ Status ShapeVerifier::HandleInfeed(HloInstruction* instruction) {
                     /*only_compare_minor_to_major_in_layout=*/true);
 }
 
-Status ShapeVerifier::HandleOutfeed(HloInstruction* instruction) {
+absl::Status ShapeVerifier::HandleOutfeed(HloInstruction* instruction) {
   HloOutfeedInstruction* outfeed = Cast<HloOutfeedInstruction>(instruction);
   TF_RETURN_IF_ERROR(CheckIsTokenOperand(instruction, 1));
 
@@ -914,7 +915,7 @@ bool ShapeVerifier::HasCompatibleElementTypes(const Shape& shape_0,
                                                          result_shape)));
 }
 
-Status ShapeVerifier::HandleRng(HloInstruction* instruction) {
+absl::Status ShapeVerifier::HandleRng(HloInstruction* instruction) {
   TF_RETURN_IF_ERROR(CheckOperandCount(instruction, 2));
 
   const Shape& shape_0 = instruction->operand(0)->shape();
@@ -963,7 +964,7 @@ Status ShapeVerifier::HandleRng(HloInstruction* instruction) {
   return OkStatus();
 }
 
-Status ShapeVerifier::HandleRngBitGenerator(HloInstruction* hlo) {
+absl::Status ShapeVerifier::HandleRngBitGenerator(HloInstruction* hlo) {
   if (!hlo->shape().IsTuple()) {
     return OkStatus();
   }
@@ -983,7 +984,8 @@ Status ShapeVerifier::HandleRngBitGenerator(HloInstruction* hlo) {
   return OkStatus();
 }
 
-Status ShapeVerifier::HandleRngGetAndUpdateState(HloInstruction* instruction) {
+absl::Status ShapeVerifier::HandleRngGetAndUpdateState(
+    HloInstruction* instruction) {
   TF_RETURN_IF_ERROR(CheckOperandCount(instruction, 0));
   const Shape& result_shape = instruction->shape();
   const Shape expected_shape = ShapeUtil::MakeShape(U64, {2});
@@ -996,19 +998,19 @@ Status ShapeVerifier::HandleRngGetAndUpdateState(HloInstruction* instruction) {
   return OkStatus();
 }
 
-Status ShapeVerifier::HandleReverse(HloInstruction* reverse) {
+absl::Status ShapeVerifier::HandleReverse(HloInstruction* reverse) {
   return CheckShape(
       reverse, ShapeInference::InferReverseShape(reverse->operand(0)->shape(),
                                                  reverse->dimensions()));
 }
 
-Status ShapeVerifier::HandleTopK(HloInstruction* hlo) {
+absl::Status ShapeVerifier::HandleTopK(HloInstruction* hlo) {
   return CheckShape(
       hlo, ShapeInference::InferTopKShape(hlo->operand(0)->shape(),
                                           Cast<HloTopKInstruction>(hlo)->k()));
 }
 
-Status ShapeVerifier::HandleSort(HloInstruction* hlo) {
+absl::Status ShapeVerifier::HandleSort(HloInstruction* hlo) {
   HloSortInstruction* sort = Cast<HloSortInstruction>(hlo);
   if (sort->operand_count() < 1) {
     return Internal("Expected at least 1 operand for %s instruction: %s",
@@ -1072,7 +1074,7 @@ Status ShapeVerifier::HandleSort(HloInstruction* hlo) {
   return CheckVariadicShape(sort);
 }
 
-Status ShapeVerifier::HandleConstant(HloInstruction* constant) {
+absl::Status ShapeVerifier::HandleConstant(HloInstruction* constant) {
   if (!Cast<HloConstantInstruction>(constant)->HasLiteral()) {
     return Internal("Constant is required to have a valid literal: %s",
                     constant->ToString());
@@ -1081,7 +1083,7 @@ Status ShapeVerifier::HandleConstant(HloInstruction* constant) {
                     /*only_compare_minor_to_major_in_layout=*/true);
 }
 
-Status ShapeVerifier::HandleIota(HloInstruction* hlo) {
+absl::Status ShapeVerifier::HandleIota(HloInstruction* hlo) {
   auto* iota = Cast<HloIotaInstruction>(hlo);
   if (!iota->shape().IsArray()) {
     return Internal("Iota does not support non-array result.");
@@ -1110,7 +1112,8 @@ Status ShapeVerifier::HandleIota(HloInstruction* hlo) {
   return OkStatus();
 }
 
-Status ShapeVerifier::HandleGetTupleElement(HloInstruction* get_tuple_element) {
+absl::Status ShapeVerifier::HandleGetTupleElement(
+    HloInstruction* get_tuple_element) {
   return CheckShape(get_tuple_element,
                     ShapeInference::InferGetTupleElementShape(
                         get_tuple_element->operand(0)->shape(),
@@ -1118,7 +1121,7 @@ Status ShapeVerifier::HandleGetTupleElement(HloInstruction* get_tuple_element) {
 }
 
 namespace {
-Status SameElementTypesForOperandsAndToApplyParameters(
+absl::Status SameElementTypesForOperandsAndToApplyParameters(
     const HloInstruction& instruction, int64_t num_operands_to_check) {
   const ProgramShape& to_apply = instruction.to_apply()->ComputeProgramShape();
   for (int i = 0; i < num_operands_to_check; ++i) {
@@ -1135,7 +1138,7 @@ Status SameElementTypesForOperandsAndToApplyParameters(
 }
 }  // namespace
 
-Status ShapeVerifier::HandleReduce(HloInstruction* reduce) {
+absl::Status ShapeVerifier::HandleReduce(HloInstruction* reduce) {
   if (reduce->operand_count() % 2 != 0) {
     return Internal(
         "Expected an even number of operands for %s instruction: %s",
@@ -1157,7 +1160,7 @@ Status ShapeVerifier::HandleReduce(HloInstruction* reduce) {
                    *reduce, reduce->operand_count());
 }
 
-Status ShapeVerifier::HandleBitcast(HloInstruction* bitcast) {
+absl::Status ShapeVerifier::HandleBitcast(HloInstruction* bitcast) {
   const Shape& output_shape = bitcast->shape();
   const Shape& operand_shape = bitcast->operand(0)->shape();
   if (opts_.layout_sensitive &&
@@ -1180,7 +1183,7 @@ Status ShapeVerifier::HandleBitcast(HloInstruction* bitcast) {
   return OkStatus();
 }
 
-Status ShapeVerifier::HandleBroadcast(HloInstruction* broadcast) {
+absl::Status ShapeVerifier::HandleBroadcast(HloInstruction* broadcast) {
   // HLO broadcast has no exact analog at the client level so there is no
   // ShapeInference method. Check the output shape explicitly.
   const Shape& operand_shape = broadcast->operand(0)->shape();
@@ -1201,7 +1204,8 @@ Status ShapeVerifier::HandleBroadcast(HloInstruction* broadcast) {
   return OkStatus();
 }
 
-Status ShapeVerifier::HandleDynamicReshape(HloInstruction* dynamic_reshape) {
+absl::Status ShapeVerifier::HandleDynamicReshape(
+    HloInstruction* dynamic_reshape) {
   // Check for mixed precision.
   const Shape& operand_shape = dynamic_reshape->operand(0)->shape();
   TF_RET_CHECK(SameElementType(dynamic_reshape->shape(), operand_shape));
@@ -1215,7 +1219,7 @@ Status ShapeVerifier::HandleDynamicReshape(HloInstruction* dynamic_reshape) {
   return OkStatus();
 }
 
-Status ShapeVerifier::HandleReshape(HloInstruction* reshape) {
+absl::Status ShapeVerifier::HandleReshape(HloInstruction* reshape) {
   // Check for mixed precision.
   const Shape& operand_shape = reshape->operand(0)->shape();
   TF_RET_CHECK(SameElementType(reshape->shape(), operand_shape));
@@ -1224,17 +1228,17 @@ Status ShapeVerifier::HandleReshape(HloInstruction* reshape) {
   return OkStatus();
 }
 
-Status ShapeVerifier::HandleTranspose(HloInstruction* transpose) {
+absl::Status ShapeVerifier::HandleTranspose(HloInstruction* transpose) {
   return CheckShape(
       transpose, ShapeInference::InferTransposeShape(
                      transpose->operand(0)->shape(), transpose->dimensions()));
 }
 
-Status ShapeVerifier::HandleParameter(HloInstruction* hlo) {
+absl::Status ShapeVerifier::HandleParameter(HloInstruction* hlo) {
   return OkStatus();
 }
 
-Status ShapeVerifier::HandleFusion(HloInstruction* fusion) {
+absl::Status ShapeVerifier::HandleFusion(HloInstruction* fusion) {
   if (fusion->called_computations().size() != 1) {
     return Internal("Fusion has a non-unary number of called computations (%s)",
                     fusion->ToString().c_str());
@@ -1291,7 +1295,7 @@ Status ShapeVerifier::HandleFusion(HloInstruction* fusion) {
   return OkStatus();
 }
 
-Status ShapeVerifier::HandleCall(HloInstruction* call) {
+absl::Status ShapeVerifier::HandleCall(HloInstruction* call) {
   TF_RETURN_IF_ERROR(
       CheckParameterCount(call, call->to_apply(), call->operand_count()));
   for (int64_t i = 0; i < call->to_apply()->num_parameters(); ++i) {
@@ -1301,7 +1305,7 @@ Status ShapeVerifier::HandleCall(HloInstruction* call) {
   return CheckShape(call, call->to_apply()->root_instruction()->shape());
 }
 
-Status ShapeVerifier::HandleCustomCall(HloInstruction* instruction) {
+absl::Status ShapeVerifier::HandleCustomCall(HloInstruction* instruction) {
   const HloCustomCallInstruction* custom_call =
       DynCast<const HloCustomCallInstruction>(instruction);
   TF_RET_CHECK(custom_call != nullptr);
@@ -1347,14 +1351,14 @@ Status ShapeVerifier::HandleCustomCall(HloInstruction* instruction) {
   return OkStatus();
 }
 
-Status ShapeVerifier::HandleSlice(HloInstruction* slice) {
+absl::Status ShapeVerifier::HandleSlice(HloInstruction* slice) {
   return CheckShape(slice,
                     ShapeInference::InferSliceShape(
                         slice->operand(0)->shape(), slice->slice_starts(),
                         slice->slice_limits(), slice->slice_strides()));
 }
 
-Status ShapeVerifier::HandleDynamicSlice(HloInstruction* dynamic_slice) {
+absl::Status ShapeVerifier::HandleDynamicSlice(HloInstruction* dynamic_slice) {
   return CheckShape(
       dynamic_slice,
       ShapeInference::InferDynamicSliceShape(
@@ -1363,7 +1367,7 @@ Status ShapeVerifier::HandleDynamicSlice(HloInstruction* dynamic_slice) {
           dynamic_slice->dynamic_slice_sizes()));
 }
 
-Status ShapeVerifier::HandleDynamicUpdateSlice(
+absl::Status ShapeVerifier::HandleDynamicUpdateSlice(
     HloInstruction* dynamic_update_slice) {
   return CheckShape(
       dynamic_update_slice,
@@ -1374,11 +1378,11 @@ Status ShapeVerifier::HandleDynamicUpdateSlice(
               ->index_shapes()));
 }
 
-Status ShapeVerifier::HandleTuple(HloInstruction* tuple) {
+absl::Status ShapeVerifier::HandleTuple(HloInstruction* tuple) {
   return CheckVariadicShape(tuple);
 }
 
-Status ShapeVerifier::HandleMap(HloInstruction* map) {
+absl::Status ShapeVerifier::HandleMap(HloInstruction* map) {
   std::vector<const Shape*> operand_shapes;
   int64_t max_operand_rank = 0;
   for (const HloInstruction* operand : map->operands()) {
@@ -1401,7 +1405,7 @@ Status ShapeVerifier::HandleMap(HloInstruction* map) {
                    *map, map->operand_count());
 }
 
-Status ShapeVerifier::HandleReduceWindow(HloInstruction* reduce_window) {
+absl::Status ShapeVerifier::HandleReduceWindow(HloInstruction* reduce_window) {
   auto reduce_window_instr = Cast<HloReduceWindowInstruction>(reduce_window);
   auto input_shapes = reduce_window_instr->input_shapes();
   auto init_shapes = reduce_window_instr->init_value_shapes();
@@ -1416,7 +1420,8 @@ Status ShapeVerifier::HandleReduceWindow(HloInstruction* reduce_window) {
                    *reduce_window, reduce_window->operand_count());
 }
 
-Status ShapeVerifier::HandleSelectAndScatter(HloInstruction* instruction) {
+absl::Status ShapeVerifier::HandleSelectAndScatter(
+    HloInstruction* instruction) {
   return CheckShape(
       instruction,
       ShapeInference::InferSelectAndScatterShape(
@@ -1426,7 +1431,7 @@ Status ShapeVerifier::HandleSelectAndScatter(HloInstruction* instruction) {
           instruction->scatter()->ComputeProgramShape()));
 }
 
-Status ShapeVerifier::HandleWhile(HloInstruction* xla_while) {
+absl::Status ShapeVerifier::HandleWhile(HloInstruction* xla_while) {
   TF_RETURN_IF_ERROR(
       CheckParameterCount(xla_while, xla_while->while_body(), 1));
   TF_RETURN_IF_ERROR(
@@ -1450,7 +1455,7 @@ Status ShapeVerifier::HandleWhile(HloInstruction* xla_while) {
                     xla_while->while_body()->root_instruction()->shape());
 }
 
-Status ShapeVerifier::HandleConditional(HloInstruction* conditional) {
+absl::Status ShapeVerifier::HandleConditional(HloInstruction* conditional) {
   if (!ShapeUtil::IsScalar(conditional->operand(0)->shape())) {
     return InvalidArgument(
         "The first operand of conditional must be a scalar. Got %s",
@@ -1482,14 +1487,14 @@ Status ShapeVerifier::HandleConditional(HloInstruction* conditional) {
   return OkStatus();
 }
 
-Status ShapeVerifier::HandlePad(HloInstruction* pad) {
+absl::Status ShapeVerifier::HandlePad(HloInstruction* pad) {
   return CheckShape(pad, ShapeInference::InferPadShape(pad->operand(0)->shape(),
                                                        pad->operand(1)->shape(),
                                                        pad->padding_config()));
 }
 
 namespace {
-Status CheckAsyncOpOperand(const HloInstruction* async_op) {
+absl::Status CheckAsyncOpOperand(const HloInstruction* async_op) {
   const HloInstruction* operand = async_op->operand(0);
   if (operand->opcode() != HloOpcode::kAsyncStart &&
       operand->opcode() != HloOpcode::kAsyncUpdate) {
@@ -1514,7 +1519,7 @@ Status CheckAsyncOpOperand(const HloInstruction* async_op) {
   return OkStatus();
 }
 
-Status CheckAsyncOpComputationThreadName(const HloInstruction* async_op) {
+absl::Status CheckAsyncOpComputationThreadName(const HloInstruction* async_op) {
   absl::string_view async_execution_thread = async_op->async_execution_thread();
   if (async_execution_thread !=
       async_op->async_wrapped_computation()->execution_thread()) {
@@ -1529,8 +1534,8 @@ Status CheckAsyncOpComputationThreadName(const HloInstruction* async_op) {
       /*skip_nested_async_op_check=*/false);
 }
 
-Status CheckCallableInstructionThreadName(const HloInstruction* instruction,
-                                          bool skip_nested_async_op_check) {
+absl::Status CheckCallableInstructionThreadName(
+    const HloInstruction* instruction, bool skip_nested_async_op_check) {
   for (const HloComputation* computation : instruction->called_computations()) {
     if (instruction->parent() != nullptr) {
       if (instruction->parent()->execution_thread() !=
@@ -1549,7 +1554,7 @@ Status CheckCallableInstructionThreadName(const HloInstruction* instruction,
 }
 }  // namespace
 
-Status ShapeVerifier::CheckAsyncOpComputationShapes(
+absl::Status ShapeVerifier::CheckAsyncOpComputationShapes(
     const HloInstruction* async_op, const Shape& async_shape) {
   if (!async_shape.IsTuple() || async_shape.tuple_shapes_size() < 2) {
     return Internal(
@@ -1586,7 +1591,7 @@ Status ShapeVerifier::CheckAsyncOpComputationShapes(
   return OkStatus();
 }
 
-Status ShapeVerifier::HandleAsyncStart(HloInstruction* async_start) {
+absl::Status ShapeVerifier::HandleAsyncStart(HloInstruction* async_start) {
   TF_RETURN_IF_ERROR(
       CheckAsyncOpComputationShapes(async_start, async_start->shape()));
   TF_RETURN_IF_ERROR(CheckAsyncOpComputationThreadName(async_start));
@@ -1605,7 +1610,7 @@ Status ShapeVerifier::HandleAsyncStart(HloInstruction* async_start) {
   return OkStatus();
 }
 
-Status ShapeVerifier::HandleAsyncUpdate(HloInstruction* async_update) {
+absl::Status ShapeVerifier::HandleAsyncUpdate(HloInstruction* async_update) {
   TF_RETURN_IF_ERROR(CheckAsyncOpComputationThreadName(async_update));
   if (!ShapesSame(async_update->operand(0)->shape(), async_update->shape())) {
     return Internal(
@@ -1619,7 +1624,7 @@ Status ShapeVerifier::HandleAsyncUpdate(HloInstruction* async_update) {
   return CheckAsyncOpOperand(async_update);
 }
 
-Status ShapeVerifier::HandleAsyncDone(HloInstruction* async_done) {
+absl::Status ShapeVerifier::HandleAsyncDone(HloInstruction* async_done) {
   TF_RETURN_IF_ERROR(CheckAsyncOpComputationThreadName(async_done));
   TF_RETURN_IF_ERROR(CheckAsyncOpComputationShapes(
       async_done, async_done->operand(0)->shape()));
@@ -1634,7 +1639,7 @@ Status ShapeVerifier::HandleAsyncDone(HloInstruction* async_done) {
   return CheckAsyncOpOperand(async_done);
 }
 
-Status ShapeVerifier::HandleCopyStart(HloInstruction* copy_start) {
+absl::Status ShapeVerifier::HandleCopyStart(HloInstruction* copy_start) {
   return CheckShape(copy_start,
                     ShapeUtil::MakeTupleShape({copy_start->operand(0)->shape(),
                                                copy_start->operand(0)->shape(),
@@ -1642,7 +1647,7 @@ Status ShapeVerifier::HandleCopyStart(HloInstruction* copy_start) {
                     /*only_compare_minor_to_major_in_layout=*/true);
 }
 
-Status ShapeVerifier::HandleCopyDone(HloInstruction* copy_done) {
+absl::Status ShapeVerifier::HandleCopyDone(HloInstruction* copy_done) {
   const Shape& operand_shape = copy_done->operand(0)->shape();
   const Shape& dest_shape = ShapeUtil::GetTupleElementShape(operand_shape, 0);
   const Shape& src_shape = ShapeUtil::GetTupleElementShape(operand_shape, 1);
@@ -1658,7 +1663,7 @@ Status ShapeVerifier::HandleCopyDone(HloInstruction* copy_done) {
                                    copy_done->operand(0)->shape(), 0));
 }
 
-Status ShapeVerifier::HandleSend(HloInstruction* send) {
+absl::Status ShapeVerifier::HandleSend(HloInstruction* send) {
   return CheckShape(send,
                     ShapeUtil::MakeTupleShape({send->operand(0)->shape(),
                                                ShapeUtil::MakeShape(U32, {}),
@@ -1666,11 +1671,11 @@ Status ShapeVerifier::HandleSend(HloInstruction* send) {
                     /*only_compare_minor_to_major_in_layout=*/true);
 }
 
-Status ShapeVerifier::HandleSendDone(HloInstruction* send_done) {
+absl::Status ShapeVerifier::HandleSendDone(HloInstruction* send_done) {
   return CheckShape(send_done, ShapeUtil::MakeTokenShape());
 }
 
-Status ShapeVerifier::HandleRecv(HloInstruction* recv) {
+absl::Status ShapeVerifier::HandleRecv(HloInstruction* recv) {
   return CheckShape(
       recv,
       ShapeUtil::MakeTupleShape(
@@ -1679,7 +1684,7 @@ Status ShapeVerifier::HandleRecv(HloInstruction* recv) {
       /*only_compare_minor_to_major_in_layout=*/true);
 }
 
-Status ShapeVerifier::HandleRecvDone(HloInstruction* recv_done) {
+absl::Status ShapeVerifier::HandleRecvDone(HloInstruction* recv_done) {
   return CheckShape(
       recv_done,
       ShapeUtil::MakeTupleShape(
@@ -1687,7 +1692,7 @@ Status ShapeVerifier::HandleRecvDone(HloInstruction* recv_done) {
            ShapeUtil::MakeTokenShape()}));
 }
 
-Status ShapeVerifier::HandleBatchNormTraining(
+absl::Status ShapeVerifier::HandleBatchNormTraining(
     HloInstruction* batch_norm_training) {
   return CheckShape(batch_norm_training,
                     ShapeInference::InferBatchNormTrainingShape(
@@ -1697,7 +1702,7 @@ Status ShapeVerifier::HandleBatchNormTraining(
                         batch_norm_training->feature_index()));
 }
 
-Status ShapeVerifier::HandleBatchNormInference(
+absl::Status ShapeVerifier::HandleBatchNormInference(
     HloInstruction* batch_norm_inference) {
   return CheckShape(batch_norm_inference,
                     ShapeInference::InferBatchNormInferenceShape(
@@ -1709,7 +1714,8 @@ Status ShapeVerifier::HandleBatchNormInference(
                         batch_norm_inference->feature_index()));
 }
 
-Status ShapeVerifier::HandleBatchNormGrad(HloInstruction* batch_norm_grad) {
+absl::Status ShapeVerifier::HandleBatchNormGrad(
+    HloInstruction* batch_norm_grad) {
   return CheckShape(batch_norm_grad, ShapeInference::InferBatchNormGradShape(
                                          batch_norm_grad->operand(0)->shape(),
                                          batch_norm_grad->operand(1)->shape(),
@@ -1723,7 +1729,7 @@ namespace {
 
 // Checks that the instruction does not have mixed precision floating point
 // inputs.
-Status CheckMixedPrecisionOperands(const HloInstruction* instruction) {
+absl::Status CheckMixedPrecisionOperands(const HloInstruction* instruction) {
   switch (instruction->opcode()) {
     // Allow-list the following opcodes for mixed-precision check, because
     // they involve data pass through or grouping via tuples, where the
@@ -1790,7 +1796,7 @@ Status CheckMixedPrecisionOperands(const HloInstruction* instruction) {
 
 }  // namespace
 
-Status ShapeVerifier::HandleGather(HloInstruction* gather) {
+absl::Status ShapeVerifier::HandleGather(HloInstruction* gather) {
   return CheckShape(
       gather,
       ShapeInference::InferGatherShape(
@@ -1798,7 +1804,7 @@ Status ShapeVerifier::HandleGather(HloInstruction* gather) {
           gather->gather_dimension_numbers(), gather->gather_slice_sizes()));
 }
 
-Status ShapeVerifier::HandleScatter(HloInstruction* scatter) {
+absl::Status ShapeVerifier::HandleScatter(HloInstruction* scatter) {
   absl::InlinedVector<const Shape*, 3> arg_shapes;
   arg_shapes.reserve(scatter->operand_count());
   for (const HloInstruction* operand : scatter->operands()) {
@@ -1810,7 +1816,7 @@ Status ShapeVerifier::HandleScatter(HloInstruction* scatter) {
                         scatter->scatter_dimension_numbers()));
 }
 
-Status ShapeVerifier::HandleAfterAll(HloInstruction* token) {
+absl::Status ShapeVerifier::HandleAfterAll(HloInstruction* token) {
   std::vector<const Shape*> operand_shapes;
   for (const HloInstruction* operand : token->operands()) {
     operand_shapes.push_back(&operand->shape());
@@ -1818,27 +1824,28 @@ Status ShapeVerifier::HandleAfterAll(HloInstruction* token) {
   return CheckShape(token, ShapeUtil::MakeTokenShape());
 }
 
-Status ShapeVerifier::HandleAddDependency(HloInstruction* add_dependency) {
+absl::Status ShapeVerifier::HandleAddDependency(
+    HloInstruction* add_dependency) {
   TF_RETURN_IF_ERROR(CheckIsTokenOperand(add_dependency, 1));
   return CheckShape(add_dependency, add_dependency->operand(0)->shape());
 }
 
-Status ShapeVerifier::HandleGetDimensionSize(HloInstruction* get_size) {
+absl::Status ShapeVerifier::HandleGetDimensionSize(HloInstruction* get_size) {
   return CheckShape(get_size,
                     ShapeInference::InferGetDimensionSizeShape(
                         get_size->operand(0)->shape(), get_size->dimension()));
 }
 
-Status ShapeVerifier::HandleSetDimensionSize(HloInstruction* set_size) {
+absl::Status ShapeVerifier::HandleSetDimensionSize(HloInstruction* set_size) {
   return CheckShape(set_size,
                     ShapeInference::InferSetDimensionSizeShape(
                         set_size->operand(0)->shape(),
                         set_size->operand(1)->shape(), set_size->dimension()));
 }
 
-Status ShapeVerifier::CheckShape(const HloInstruction* instruction,
-                                 const Shape& inferred_shape,
-                                 bool only_compare_minor_to_major_in_layout) {
+absl::Status ShapeVerifier::CheckShape(
+    const HloInstruction* instruction, const Shape& inferred_shape,
+    bool only_compare_minor_to_major_in_layout) {
   // If allow_mixed_precision_ is false, check if there are operands with
   // different precisions. We need this check because ShapeInference allows
   // mixed precision inputs.
@@ -1917,11 +1924,11 @@ Status ShapeVerifier::CheckShape(const HloInstruction* instruction,
   return OkStatus();
 }
 
-Status ShapeVerifier::CheckShape(
+absl::Status ShapeVerifier::CheckShape(
     const HloInstruction* instruction,
     const absl::StatusOr<Shape>& inferred_shape_status) {
   if (!inferred_shape_status.ok()) {
-    Status s = inferred_shape_status.status();
+    absl::Status s = inferred_shape_status.status();
     tsl::errors::AppendToMessage(&s, ", for instruction ",
                                  instruction->ToString());
     return s;
@@ -1929,33 +1936,37 @@ Status ShapeVerifier::CheckShape(
   return CheckShape(instruction, inferred_shape_status.value());
 }
 
-Status ShapeVerifier::CheckUnaryShape(const HloInstruction* instruction) {
+absl::Status ShapeVerifier::CheckUnaryShape(const HloInstruction* instruction) {
   return CheckShape(instruction,
                     ShapeInference::InferUnaryOpShape(instruction->opcode(),
                                                       instruction->operand(0)));
 }
 
-Status ShapeVerifier::CheckBinaryShape(const HloInstruction* instruction) {
+absl::Status ShapeVerifier::CheckBinaryShape(
+    const HloInstruction* instruction) {
   return CheckShape(
       instruction, ShapeInference::InferBinaryOpShape(instruction->opcode(),
                                                       instruction->operand(0),
                                                       instruction->operand(1)));
 }
 
-Status ShapeVerifier::CheckTernaryShape(const HloInstruction* instruction) {
+absl::Status ShapeVerifier::CheckTernaryShape(
+    const HloInstruction* instruction) {
   return CheckShape(instruction,
                     ShapeInference::InferTernaryOpShape(
                         instruction->opcode(), instruction->operand(0),
                         instruction->operand(1), instruction->operand(2)));
 }
 
-Status ShapeVerifier::CheckVariadicShape(const HloInstruction* instruction) {
+absl::Status ShapeVerifier::CheckVariadicShape(
+    const HloInstruction* instruction) {
   return CheckShape(instruction,
                     ShapeInference::InferVariadicOpShape(
                         instruction->opcode(), instruction->operands()));
 }
 
-Status ShapeVerifier::VerifyEntryComputationLayout(const HloModule& module) {
+absl::Status ShapeVerifier::VerifyEntryComputationLayout(
+    const HloModule& module) {
   const HloComputation* computation = module.entry_computation();
   const auto& layout = module.entry_computation_layout();
   const ShapeLayout& result_layout = layout.result_layout();
@@ -2057,7 +2068,7 @@ std::string ComputationsToString(
 //
 // (3) the operands of each instruction are non-null and are in the same
 // computation as the instruction.
-Status VerifyHloStructure(HloModule* module) {
+absl::Status VerifyHloStructure(HloModule* module) {
   for (const HloComputation* computation : module->computations()) {
     if (computation == nullptr) {
       return Internal("Computation in module %s is a null pointer",
@@ -2138,8 +2149,8 @@ bool ShapeContainsToken(const Shape& shape) {
 }
 
 // Checks if the given two instructions share the same channel id.
-Status CheckSameChannel(const HloInstruction* instr1,
-                        const HloInstruction* instr2) {
+absl::Status CheckSameChannel(const HloInstruction* instr1,
+                              const HloInstruction* instr2) {
   if (instr1->channel_id() != instr2->channel_id()) {
     return Internal(
         "Expected to have the same channel id, actual channel ids are: %s "
@@ -2153,8 +2164,8 @@ Status CheckSameChannel(const HloInstruction* instr1,
 // Checks if the given two instructions have the same is_host_transfer
 // attribute value. Instructions must be send/recv instructions or their
 // 'done' variant.
-Status CheckSameIsHostTransfer(const HloInstruction* instr1,
-                               const HloInstruction* instr2) {
+absl::Status CheckSameIsHostTransfer(const HloInstruction* instr1,
+                                     const HloInstruction* instr2) {
   const HloSendRecvInstruction* send_recv1 =
       DynCast<const HloSendRecvInstruction>(instr1);
   const HloSendRecvInstruction* send_recv2 =
@@ -2171,8 +2182,9 @@ Status CheckSameIsHostTransfer(const HloInstruction* instr1,
   return OkStatus();
 }
 
-Status VerifySingleUser(const HloInstruction* instruction,
-                        const absl::flat_hash_set<HloOpcode>& expected_users) {
+absl::Status VerifySingleUser(
+    const HloInstruction* instruction,
+    const absl::flat_hash_set<HloOpcode>& expected_users) {
   TF_RET_CHECK(instruction->users().size() == 1)
       << "The " << instruction->opcode()
       << " instruction requires one consumer, found "
@@ -2190,8 +2202,9 @@ Status VerifySingleUser(const HloInstruction* instruction,
   return OkStatus();
 }
 
-Status VerifySingleOperand(const HloInstruction* instruction,
-                           const std::vector<HloOpcode>& expected_operands) {
+absl::Status VerifySingleOperand(
+    const HloInstruction* instruction,
+    const std::vector<HloOpcode>& expected_operands) {
   TF_RET_CHECK(instruction->operands().size() == 1)
       << "The " << instruction->opcode()
       << " instruction requires one consumer, found "
@@ -2211,7 +2224,7 @@ Status VerifySingleOperand(const HloInstruction* instruction,
 }
 
 // Checks asynchronous instruction pairs.
-Status VerifyAsynchronousInstructionPairs(const HloModule& module) {
+absl::Status VerifyAsynchronousInstructionPairs(const HloModule& module) {
   // CopyStart must have a single CopyDone user.
   for (const HloComputation* computation : module.computations()) {
     for (const HloInstruction* instruction : computation->instructions()) {
@@ -2273,7 +2286,7 @@ Status VerifyAsynchronousInstructionPairs(const HloModule& module) {
 
 // Checks that the asynchronous computation only has a root and parameter
 // instructions.
-Status VerifyAsyncComputation(const HloComputation* async_computation) {
+absl::Status VerifyAsyncComputation(const HloComputation* async_computation) {
   if (!async_computation->CanExpandIntoSingleInstruction()) {
     return FailedPrecondition(
         "Asynchronous computation %s expected to contain only the root and "
@@ -2285,7 +2298,7 @@ Status VerifyAsyncComputation(const HloComputation* async_computation) {
 
 // Checks that AllReduce instructions in the module are either all layout
 // constrained or all unconstrained.
-Status VerifyLayoutConstrainedAllReduce(const HloModule& module) {
+absl::Status VerifyLayoutConstrainedAllReduce(const HloModule& module) {
   const HloAllReduceInstruction* reference = nullptr;
   for (const HloComputation* computation : module.computations()) {
     for (const HloInstruction* instruction : computation->instructions()) {
@@ -2309,7 +2322,7 @@ Status VerifyLayoutConstrainedAllReduce(const HloModule& module) {
 
 // Checks various invariants of channel instructions (send/recv and
 // collectives).
-Status VerifyChannels(const HloModule& module) {
+absl::Status VerifyChannels(const HloModule& module) {
   absl::flat_hash_map<int64_t, std::vector<const HloInstruction*>>
       channel_instructions;
 
@@ -2477,7 +2490,7 @@ Status VerifyChannels(const HloModule& module) {
 }
 
 // CHECKs various invariants of a fusion instruction.
-Status CheckFusionInstruction(HloInstruction* fusion) {
+absl::Status CheckFusionInstruction(HloInstruction* fusion) {
   // The parent fusion instruction of the fusion computation must be 'fusion'.
   HloComputation* fused_computation = fusion->fused_instructions_computation();
   if (fusion != fused_computation->FusionInstruction()) {
@@ -2603,7 +2616,7 @@ Status CheckFusionInstruction(HloInstruction* fusion) {
 
 // Checks that the operand shapes are compatible to the output shape, i.e.,
 // that there are no implicit broadcasts.
-Status CheckElementwiseInstruction(HloInstruction* instruction) {
+absl::Status CheckElementwiseInstruction(HloInstruction* instruction) {
   const Shape& out_shape = instruction->shape();
   for (HloInstruction* operand : instruction->operands()) {
     const Shape& operand_shape = operand->shape();
@@ -2660,15 +2673,15 @@ class InstructionVerifier : public DfsHloVisitorWithDefault {
     }
   }
 
-  Status DefaultAction(HloInstruction*) override { return OkStatus(); }
+  absl::Status DefaultAction(HloInstruction*) override { return OkStatus(); }
 
-  Status HandleFusion(HloInstruction* fusion) override {
+  absl::Status HandleFusion(HloInstruction* fusion) override {
     TF_RETURN_IF_ERROR(CheckCallableInstructionThreadName(
         fusion, /*skip_nested_async_op_check*/ false));
     return CheckFusionInstruction(fusion);
   }
 
-  Status HandleBroadcast(HloInstruction* broadcast) override {
+  absl::Status HandleBroadcast(HloInstruction* broadcast) override {
     // If you see this failure then someone has confused the difference
     // between the HLO broadcast op, and the UserComputation broadcast
     // op. See https://groups.google.com/forum/#!topic/xla-dev/9LqijHmTt_I
@@ -2687,12 +2700,12 @@ class InstructionVerifier : public DfsHloVisitorWithDefault {
     return OkStatus();
   }
 
-  Status HandleBitcastConvert(HloInstruction* c) override {
+  absl::Status HandleBitcastConvert(HloInstruction* c) override {
     // Shape verifier will check all we need.
     return OkStatus();
   }
 
-  Status HandleWhile(HloInstruction* xla_while) override {
+  absl::Status HandleWhile(HloInstruction* xla_while) override {
     auto* while_cond = xla_while->while_condition();
     auto* while_body = xla_while->while_body();
     if (while_cond->num_parameters() != 1) {
@@ -2724,13 +2737,13 @@ class InstructionVerifier : public DfsHloVisitorWithDefault {
     return OkStatus();
   }
 
-  Status HandleCall(HloInstruction* call) override {
+  absl::Status HandleCall(HloInstruction* call) override {
     // Allow kCall to contain computations on separate thread.
     return CheckCallableInstructionThreadName(
         call, /*skip_nested_async_op_check=*/true);
   }
 
-  Status HandleConditional(HloInstruction* conditional) override {
+  absl::Status HandleConditional(HloInstruction* conditional) override {
     const std::vector<HloComputation*> branch_computations =
         conditional->branch_computations();
     std::vector<const HloInstruction*> sharding_check_instructions;
@@ -2759,20 +2772,20 @@ class InstructionVerifier : public DfsHloVisitorWithDefault {
     return OkStatus();
   }
 
-  Status HandleElementwiseUnary(HloInstruction* instruction) override {
+  absl::Status HandleElementwiseUnary(HloInstruction* instruction) override {
     return CheckElementwiseInstruction(instruction);
   }
 
-  Status HandleElementwiseBinary(HloInstruction* instruction) override {
+  absl::Status HandleElementwiseBinary(HloInstruction* instruction) override {
     return CheckElementwiseInstruction(instruction);
   }
 
-  Status HandleGetTupleElement(HloInstruction* gte) override {
+  absl::Status HandleGetTupleElement(HloInstruction* gte) override {
     TF_RET_CHECK(gte->operand(0)->shape().IsTuple());
     return OkStatus();
   }
 
-  Status HandleTranspose(HloInstruction* transpose) override {
+  absl::Status HandleTranspose(HloInstruction* transpose) override {
     const Shape& shape = transpose->shape();
     const HloInstruction* operand = transpose->operand(0);
     TF_RET_CHECK(shape.dimensions().size() == transpose->dimensions().size());
@@ -2788,7 +2801,7 @@ class InstructionVerifier : public DfsHloVisitorWithDefault {
     return OkStatus();
   }
 
-  Status HandleAllReduce(HloInstruction* crs) override {
+  absl::Status HandleAllReduce(HloInstruction* crs) override {
     if (crs->channel_id().has_value()) {
       TF_RET_CHECK(crs->channel_id().value() > 0)
           << "All reduce channel id must be greater than 0 for "
@@ -2797,7 +2810,7 @@ class InstructionVerifier : public DfsHloVisitorWithDefault {
     return OkStatus();
   }
 
-  Status HandleReshape(HloInstruction* hlo) override {
+  absl::Status HandleReshape(HloInstruction* hlo) override {
     if (opts_.verify_reshape_is_bitcast && !hlo->IsFused()) {
       TF_RET_CHECK(
           ShapeUtil::ReshapeIsBitcast(hlo->operand(0)->shape(), hlo->shape()))
@@ -2806,7 +2819,7 @@ class InstructionVerifier : public DfsHloVisitorWithDefault {
     return OkStatus();
   }
 
-  Status HandleCustomCall(HloInstruction* hlo) override {
+  absl::Status HandleCustomCall(HloInstruction* hlo) override {
     if (opts_.verify_custom_call_nested_computation_thread_name) {
       // Allow kCustomCall to contain computations on separate thread.
       return CheckCallableInstructionThreadName(
@@ -2815,7 +2828,7 @@ class InstructionVerifier : public DfsHloVisitorWithDefault {
     return OkStatus();
   }
 
-  Status HandleScatter(HloInstruction* scatter) override {
+  absl::Status HandleScatter(HloInstruction* scatter) override {
     int64_t rank = scatter->operand(0)->shape().rank();
     for (int64_t operand_dim :
          scatter->scatter_dimension_numbers().scatter_dims_to_operand_dims()) {
@@ -2828,7 +2841,7 @@ class InstructionVerifier : public DfsHloVisitorWithDefault {
     return OkStatus();
   }
 
-  Status Preprocess(HloInstruction* instruction) override {
+  absl::Status Preprocess(HloInstruction* instruction) override {
     auto [it, inserted] =
         instructions_by_name_.emplace(instruction->name(), instruction);
     TF_RET_CHECK(inserted) << "HLO has name that is not unique within module:\n"
@@ -2839,10 +2852,10 @@ class InstructionVerifier : public DfsHloVisitorWithDefault {
                            << it->second->parent()->name();
 
     if (instruction->has_sharding()) {
-      Status status =
+      absl::Status status =
           instruction->sharding().Validate(instruction->shape(), num_devices_);
       if (!status.ok()) {
-        return Status(
+        return absl::Status(
             status.code(),
             absl::StrCat("Invalid sharding for instruction: ",
                          instruction->ToString(), ": ", status.message()));
@@ -2852,7 +2865,7 @@ class InstructionVerifier : public DfsHloVisitorWithDefault {
     return OkStatus();
   }
 
-  Status Postprocess(HloInstruction* instruction) override {
+  absl::Status Postprocess(HloInstruction* instruction) override {
     if (!opts_.InstructionCanChangeLayout(instruction) &&
         LayoutUtil::IsDenseArray(instruction->shape()) &&
         instruction->shape().has_layout()) {
@@ -2888,7 +2901,7 @@ class InstructionVerifier : public DfsHloVisitorWithDefault {
   }
 
  private:
-  static Status VerifyConsistentSharding(
+  static absl::Status VerifyConsistentSharding(
       const HloInstruction* parent,
       absl::Span<const HloInstruction* const> instructions) {
     const HloInstruction* common_sharding_inst = nullptr;
@@ -2913,7 +2926,7 @@ class InstructionVerifier : public DfsHloVisitorWithDefault {
   // memory space from `operand_memory_space` to `result_memory_space`.
   // Returns OkStatus() if the instruction's layout changes are valid;
   // otherwise, returns an appropriate error status.
-  static Status HostOffloadInstructionCanChangeMemorySpace(
+  static absl::Status HostOffloadInstructionCanChangeMemorySpace(
       const HloInstruction* instruction, const int64_t operand_memory_space,
       const int64_t result_memory_space) {
     TF_RET_CHECK(!(operand_memory_space == Layout::kGenericFastMemorySpace &&
@@ -3007,9 +3020,9 @@ absl::StatusOr<bool> HloVerifier::Run(
   if (status_or_changed.ok()) {
     return status_or_changed.value();
   }
-  return Status(status_or_changed.status().code(),
-                absl::StrCat("during context [", context_,
-                             "]: ", status_or_changed.status().message()));
+  return absl::Status(status_or_changed.status().code(),
+                      absl::StrCat("during context [", context_, "]: ",
+                                   status_or_changed.status().message()));
 }
 
 MetadataTracker::MetadataTracker(absl::string_view prefix) : prefix_(prefix) {}
@@ -3071,7 +3084,7 @@ void MetadataTracker::HandleMetadata(const OpMetadata& metadata) {
   }
 }
 
-Status MetadataTracker::DefaultAction(HloInstruction* instruction) {
+absl::Status MetadataTracker::DefaultAction(HloInstruction* instruction) {
   HandleMetadata(instruction->metadata());
   return OkStatus();
 }
diff --git a/third_party/xla/xla/service/hlo_verifier.h b/third_party/xla/xla/service/hlo_verifier.h
index 91806e35976a06..2604645ec7ad54 100644
--- a/third_party/xla/xla/service/hlo_verifier.h
+++ b/third_party/xla/xla/service/hlo_verifier.h
@@ -155,93 +155,95 @@ class ShapeVerifier : public DfsHloVisitor {
 
   // Verifies that entry computation layout matches parameters and root shape of
   // the module's entry computation.
-  virtual Status VerifyEntryComputationLayout(const HloModule& module);
-
-  Status Preprocess(HloInstruction* hlo) override;
-
-  Status HandleElementwiseUnary(HloInstruction* hlo) override;
-  Status HandleElementwiseBinary(HloInstruction* hlo) override;
-  Status HandleClamp(HloInstruction* clamp) override;
-  Status HandleSelect(HloInstruction* select) override;
-  Status HandleConcatenate(HloInstruction* concatenate) override;
-  Status HandleIota(HloInstruction* hlo) override;
-  Status HandleConvert(HloInstruction* convert) override;
-  Status HandleBitcastConvert(HloInstruction* convert) override;
-  Status HandleStochasticConvert(HloInstruction* convert) override;
-  Status HandleCopy(HloInstruction* copy) override;
-  Status HandleDot(HloInstruction* dot) override;
-  Status HandleConvolution(HloInstruction* convolution) override;
-  Status HandleFft(HloInstruction* fft) override;
-  Status HandleCholesky(HloInstruction* hlo) override;
-  Status HandleTriangularSolve(HloInstruction* hlo) override;
-  Status HandleAllGather(HloInstruction* hlo) override;
-  Status HandleAllGatherStart(HloInstruction* hlo) override;
-  Status HandleAllGatherDone(HloInstruction* hlo) override;
-  Status HandleAllReduce(HloInstruction* hlo) override;
-  Status HandleAllReduceStart(HloInstruction* hlo) override;
-  Status HandleAllReduceDone(HloInstruction* hlo) override;
-  Status HandleAllToAll(HloInstruction* hlo) override;
-  Status HandleCollectiveBroadcast(HloInstruction* hlo) override;
-  Status HandleCollectivePermute(HloInstruction* hlo) override;
-  Status HandleCollectivePermuteStart(HloInstruction* hlo) override;
-  Status HandleCollectivePermuteDone(HloInstruction* hlo) override;
-  Status HandlePartitionId(HloInstruction* hlo) override;
-  Status HandleReplicaId(HloInstruction* hlo) override;
-  Status HandleReducePrecision(HloInstruction* reduce_precision) override;
-  Status HandleInfeed(HloInstruction*) override;
-  Status HandleOptimizationBarrier(HloInstruction* hlo) override;
-  Status HandleOutfeed(HloInstruction*) override;
-  Status HandleRng(HloInstruction*) override;
-  Status HandleRngBitGenerator(HloInstruction*) override;
-  Status HandleRngGetAndUpdateState(HloInstruction*) override;
-  Status HandleReverse(HloInstruction* reverse) override;
-  Status HandleSort(HloInstruction* hlo) override;
-  Status HandleTopK(HloInstruction* hlo) override;
-  Status HandleConstant(HloInstruction* constant) override;
-  Status HandleGetTupleElement(HloInstruction* get_tuple_element) override;
-  Status HandleReduce(HloInstruction* reduce) override;
-  Status HandleBitcast(HloInstruction* bitcast) override;
-  Status HandleBroadcast(HloInstruction* broadcast) override;
-  Status HandleReshape(HloInstruction* reshape) override;
-  Status HandleDynamicReshape(HloInstruction* dynamic_reshape) override;
-  Status HandleTranspose(HloInstruction* transpose) override;
-  Status HandleParameter(HloInstruction*) override;
-  Status HandleFusion(HloInstruction*) override;
-  Status HandleCall(HloInstruction* call) override;
-  Status HandleCustomCall(HloInstruction*) override;
-  Status HandleSlice(HloInstruction* slice) override;
-  Status HandleDynamicSlice(HloInstruction* dynamic_slice) override;
-  Status HandleDynamicUpdateSlice(
+  virtual absl::Status VerifyEntryComputationLayout(const HloModule& module);
+
+  absl::Status Preprocess(HloInstruction* hlo) override;
+
+  absl::Status HandleElementwiseUnary(HloInstruction* hlo) override;
+  absl::Status HandleElementwiseBinary(HloInstruction* hlo) override;
+  absl::Status HandleClamp(HloInstruction* clamp) override;
+  absl::Status HandleSelect(HloInstruction* select) override;
+  absl::Status HandleConcatenate(HloInstruction* concatenate) override;
+  absl::Status HandleIota(HloInstruction* hlo) override;
+  absl::Status HandleConvert(HloInstruction* convert) override;
+  absl::Status HandleBitcastConvert(HloInstruction* convert) override;
+  absl::Status HandleStochasticConvert(HloInstruction* convert) override;
+  absl::Status HandleCopy(HloInstruction* copy) override;
+  absl::Status HandleDot(HloInstruction* dot) override;
+  absl::Status HandleConvolution(HloInstruction* convolution) override;
+  absl::Status HandleFft(HloInstruction* fft) override;
+  absl::Status HandleCholesky(HloInstruction* hlo) override;
+  absl::Status HandleTriangularSolve(HloInstruction* hlo) override;
+  absl::Status HandleAllGather(HloInstruction* hlo) override;
+  absl::Status HandleAllGatherStart(HloInstruction* hlo) override;
+  absl::Status HandleAllGatherDone(HloInstruction* hlo) override;
+  absl::Status HandleAllReduce(HloInstruction* hlo) override;
+  absl::Status HandleAllReduceStart(HloInstruction* hlo) override;
+  absl::Status HandleAllReduceDone(HloInstruction* hlo) override;
+  absl::Status HandleAllToAll(HloInstruction* hlo) override;
+  absl::Status HandleCollectiveBroadcast(HloInstruction* hlo) override;
+  absl::Status HandleCollectivePermute(HloInstruction* hlo) override;
+  absl::Status HandleCollectivePermuteStart(HloInstruction* hlo) override;
+  absl::Status HandleCollectivePermuteDone(HloInstruction* hlo) override;
+  absl::Status HandlePartitionId(HloInstruction* hlo) override;
+  absl::Status HandleReplicaId(HloInstruction* hlo) override;
+  absl::Status HandleReducePrecision(HloInstruction* reduce_precision) override;
+  absl::Status HandleInfeed(HloInstruction*) override;
+  absl::Status HandleOptimizationBarrier(HloInstruction* hlo) override;
+  absl::Status HandleOutfeed(HloInstruction*) override;
+  absl::Status HandleRng(HloInstruction*) override;
+  absl::Status HandleRngBitGenerator(HloInstruction*) override;
+  absl::Status HandleRngGetAndUpdateState(HloInstruction*) override;
+  absl::Status HandleReverse(HloInstruction* reverse) override;
+  absl::Status HandleSort(HloInstruction* hlo) override;
+  absl::Status HandleTopK(HloInstruction* hlo) override;
+  absl::Status HandleConstant(HloInstruction* constant) override;
+  absl::Status HandleGetTupleElement(
+      HloInstruction* get_tuple_element) override;
+  absl::Status HandleReduce(HloInstruction* reduce) override;
+  absl::Status HandleBitcast(HloInstruction* bitcast) override;
+  absl::Status HandleBroadcast(HloInstruction* broadcast) override;
+  absl::Status HandleReshape(HloInstruction* reshape) override;
+  absl::Status HandleDynamicReshape(HloInstruction* dynamic_reshape) override;
+  absl::Status HandleTranspose(HloInstruction* transpose) override;
+  absl::Status HandleParameter(HloInstruction*) override;
+  absl::Status HandleFusion(HloInstruction*) override;
+  absl::Status HandleCall(HloInstruction* call) override;
+  absl::Status HandleCustomCall(HloInstruction*) override;
+  absl::Status HandleSlice(HloInstruction* slice) override;
+  absl::Status HandleDynamicSlice(HloInstruction* dynamic_slice) override;
+  absl::Status HandleDynamicUpdateSlice(
       HloInstruction* dynamic_update_slice) override;
-  Status HandleTuple(HloInstruction* tuple) override;
-  Status HandleMap(HloInstruction* map) override;
-  Status HandleReduceScatter(HloInstruction* hlo) override;
-  Status HandleReduceWindow(HloInstruction* reduce_window) override;
-  Status HandleSelectAndScatter(HloInstruction* instruction) override;
-  Status HandleWhile(HloInstruction* xla_while) override;
-  Status HandleConditional(HloInstruction* conditional) override;
-  Status HandlePad(HloInstruction* pad) override;
-  Status HandleAsyncStart(HloInstruction* async_start) override;
-  Status HandleAsyncUpdate(HloInstruction* async_update) override;
-  Status HandleAsyncDone(HloInstruction* async_done) override;
-  Status HandleCopyStart(HloInstruction* copy_start) override;
-  Status HandleCopyDone(HloInstruction* copy_done) override;
-  Status HandleSend(HloInstruction* send) override;
-  Status HandleSendDone(HloInstruction* send_done) override;
-  Status HandleRecv(HloInstruction* recv) override;
-  Status HandleRecvDone(HloInstruction* recv_done) override;
-  Status HandleBatchNormTraining(HloInstruction* batch_norm_training) override;
-  Status HandleBatchNormInference(
+  absl::Status HandleTuple(HloInstruction* tuple) override;
+  absl::Status HandleMap(HloInstruction* map) override;
+  absl::Status HandleReduceScatter(HloInstruction* hlo) override;
+  absl::Status HandleReduceWindow(HloInstruction* reduce_window) override;
+  absl::Status HandleSelectAndScatter(HloInstruction* instruction) override;
+  absl::Status HandleWhile(HloInstruction* xla_while) override;
+  absl::Status HandleConditional(HloInstruction* conditional) override;
+  absl::Status HandlePad(HloInstruction* pad) override;
+  absl::Status HandleAsyncStart(HloInstruction* async_start) override;
+  absl::Status HandleAsyncUpdate(HloInstruction* async_update) override;
+  absl::Status HandleAsyncDone(HloInstruction* async_done) override;
+  absl::Status HandleCopyStart(HloInstruction* copy_start) override;
+  absl::Status HandleCopyDone(HloInstruction* copy_done) override;
+  absl::Status HandleSend(HloInstruction* send) override;
+  absl::Status HandleSendDone(HloInstruction* send_done) override;
+  absl::Status HandleRecv(HloInstruction* recv) override;
+  absl::Status HandleRecvDone(HloInstruction* recv_done) override;
+  absl::Status HandleBatchNormTraining(
+      HloInstruction* batch_norm_training) override;
+  absl::Status HandleBatchNormInference(
       HloInstruction* batch_norm_inference) override;
-  Status HandleBatchNormGrad(HloInstruction* batch_norm_grad) override;
-  Status HandleGather(HloInstruction* gather) override;
-  Status HandleScatter(HloInstruction* scatter) override;
-  Status HandleAfterAll(HloInstruction* token) override;
-  Status HandleGetDimensionSize(HloInstruction* get_size) override;
-  Status HandleSetDimensionSize(HloInstruction* set_size) override;
-  Status HandleAddDependency(HloInstruction* add_dependency) override;
+  absl::Status HandleBatchNormGrad(HloInstruction* batch_norm_grad) override;
+  absl::Status HandleGather(HloInstruction* gather) override;
+  absl::Status HandleScatter(HloInstruction* scatter) override;
+  absl::Status HandleAfterAll(HloInstruction* token) override;
+  absl::Status HandleGetDimensionSize(HloInstruction* get_size) override;
+  absl::Status HandleSetDimensionSize(HloInstruction* set_size) override;
+  absl::Status HandleAddDependency(HloInstruction* add_dependency) override;
 
-  Status FinishVisit(HloInstruction*) override { return OkStatus(); }
+  absl::Status FinishVisit(HloInstruction*) override { return OkStatus(); }
 
  protected:
   // Helpers that switch on layout_sensitive_.
@@ -249,23 +251,23 @@ class ShapeVerifier : public DfsHloVisitor {
 
   // Check the instruction's shape against the shape given by ShapeInference
   // and return an appropriate error if there is a mismatch.
-  Status CheckShape(const HloInstruction* instruction,
-                    const Shape& inferred_shape,
-                    bool only_compare_minor_to_major_in_layout = false);
+  absl::Status CheckShape(const HloInstruction* instruction,
+                          const Shape& inferred_shape,
+                          bool only_compare_minor_to_major_in_layout = false);
 
-  // Overload which takes a StatusOr to reduce boilerplate in the caller.
-  Status CheckShape(const HloInstruction* instruction,
-                    const absl::StatusOr<Shape>& inferred_shape_status);
+  // Overload which takes a absl::StatusOr to reduce boilerplate in the caller.
+  absl::Status CheckShape(const HloInstruction* instruction,
+                          const absl::StatusOr<Shape>& inferred_shape_status);
 
-  static Status CheckParameterCount(const HloInstruction* calling_instruction,
-                                    const HloComputation* computation,
-                                    int expected);
+  static absl::Status CheckParameterCount(
+      const HloInstruction* calling_instruction,
+      const HloComputation* computation, int expected);
 
   // Check a unary (binary, etc) instruction's shape against the inferred shape.
-  Status CheckUnaryShape(const HloInstruction* instruction);
-  Status CheckBinaryShape(const HloInstruction* instruction);
-  Status CheckTernaryShape(const HloInstruction* instruction);
-  Status CheckVariadicShape(const HloInstruction* instruction);
+  absl::Status CheckUnaryShape(const HloInstruction* instruction);
+  absl::Status CheckBinaryShape(const HloInstruction* instruction);
+  absl::Status CheckTernaryShape(const HloInstruction* instruction);
+  absl::Status CheckVariadicShape(const HloInstruction* instruction);
 
  private:
   std::string StringifyShape(const Shape& s) {
@@ -281,20 +283,20 @@ class ShapeVerifier : public DfsHloVisitor {
   }
 
   // Checks that the given operand of the given instruction is of type TOKEN.
-  Status CheckIsTokenOperand(const HloInstruction* instruction,
-                             int64_t operand_no);
+  absl::Status CheckIsTokenOperand(const HloInstruction* instruction,
+                                   int64_t operand_no);
 
   // Checks that the shape of the given operand of the given instruction matches
   // the given parameter of the given computation.
-  Status CheckOperandAndParameter(const HloInstruction* instruction,
-                                  int64_t operand_number,
-                                  const HloComputation* computation,
-                                  int64_t parameter_number);
+  absl::Status CheckOperandAndParameter(const HloInstruction* instruction,
+                                        int64_t operand_number,
+                                        const HloComputation* computation,
+                                        int64_t parameter_number);
 
   // Checks that the shape of async op operands and results match the called
   // computation parameters and root.
-  Status CheckAsyncOpComputationShapes(const HloInstruction* async_op,
-                                       const Shape& async_shape);
+  absl::Status CheckAsyncOpComputationShapes(const HloInstruction* async_op,
+                                             const Shape& async_shape);
 
   // Returns true if the shapes of the two operands have the same element type,
   // and the result shape either has the same element type as the operand shapes
@@ -395,7 +397,7 @@ class MetadataTracker : public DfsHloVisitorWithDefault {
  public:
   explicit MetadataTracker(absl::string_view prefix);
   ~MetadataTracker() override;
-  Status DefaultAction(HloInstruction* instruction) override;
+  absl::Status DefaultAction(HloInstruction* instruction) override;
   void HandleMetadata(const OpMetadata& metadata);
 
  private:
diff --git a/third_party/xla/xla/service/hlo_verifier_test.cc b/third_party/xla/xla/service/hlo_verifier_test.cc
index 3f7bfc23243dfd..c7e125a75d36c9 100644
--- a/third_party/xla/xla/service/hlo_verifier_test.cc
+++ b/third_party/xla/xla/service/hlo_verifier_test.cc
@@ -2832,7 +2832,7 @@ ENTRY entry {
 
   TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
                           ParseAndReturnUnverifiedModule(kHlo));
-  Status status = verifier().Run(module.get()).status();
+  absl::Status status = verifier().Run(module.get()).status();
 
   TF_ASSERT_OK(status);
 }
@@ -2850,7 +2850,7 @@ ENTRY entry {
 
   TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
                           ParseAndReturnUnverifiedModule(kHlo));
-  Status status = verifier().Run(module.get()).status();
+  absl::Status status = verifier().Run(module.get()).status();
 
   ASSERT_FALSE(status.ok());
   EXPECT_THAT(status.message(), HasSubstr("should be compatible"));
@@ -2869,7 +2869,7 @@ ENTRY entry {
 
   TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
                           ParseAndReturnUnverifiedModule(kHlo));
-  Status status = verifier().Run(module.get()).status();
+  absl::Status status = verifier().Run(module.get()).status();
 
   TF_ASSERT_OK(status);
 }
@@ -2887,7 +2887,7 @@ ENTRY entry {
 
   TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
                           ParseAndReturnUnverifiedModule(kHlo));
-  Status status = verifier().Run(module.get()).status();
+  absl::Status status = verifier().Run(module.get()).status();
 
   ASSERT_FALSE(status.ok());
   EXPECT_THAT(status.message(), HasSubstr("should be compatible"));
@@ -2905,7 +2905,7 @@ ENTRY entry {
 
   TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
                           ParseAndReturnUnverifiedModule(kHlo));
-  Status status = verifier().Run(module.get()).status();
+  absl::Status status = verifier().Run(module.get()).status();
 
   ASSERT_FALSE(status.ok());
   EXPECT_THAT(status.message(),
@@ -2925,7 +2925,7 @@ ENTRY entry {
 
   TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
                           ParseAndReturnUnverifiedModule(kHlo));
-  Status status = verifier().Run(module.get()).status();
+  absl::Status status = verifier().Run(module.get()).status();
 
   TF_ASSERT_OK(status);
 }
@@ -2942,7 +2942,7 @@ ENTRY entry {
 
   TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
                           ParseAndReturnUnverifiedModule(kHlo));
-  Status status = verifier().Run(module.get()).status();
+  absl::Status status = verifier().Run(module.get()).status();
 
   TF_ASSERT_OK(status);
 }
diff --git a/third_party/xla/xla/service/host_memory_transfer_asyncifier.cc b/third_party/xla/xla/service/host_memory_transfer_asyncifier.cc
index c704183e325d61..6fe52707d88ede 100644
--- a/third_party/xla/xla/service/host_memory_transfer_asyncifier.cc
+++ b/third_party/xla/xla/service/host_memory_transfer_asyncifier.cc
@@ -42,13 +42,13 @@ class HostMemoryTransferAsyncifierVisitor : public DfsHloVisitorWithDefault {
       : kHostMemorySpaceColor(host_memory_space_color) {}
   bool Changed() const { return changed_; }
 
-  Status DefaultAction(HloInstruction* hlo_instruction) override {
+  absl::Status DefaultAction(HloInstruction* hlo_instruction) override {
     return OkStatus();
   }
 
   // Replace all dynamic-slice ops which slice from host memory to device memory
   // with an asynchronous dynamic-slice.
-  Status HandleDynamicSlice(HloInstruction* dynamic_slice) override {
+  absl::Status HandleDynamicSlice(HloInstruction* dynamic_slice) override {
     // Check that the dynamic_slice and its first operand have layouts. This
     // pass must only be run after LayoutAssignment.
     HloInstruction* dynamic_slice_operand = dynamic_slice->mutable_operand(0);
@@ -95,7 +95,7 @@ class HostMemoryTransferAsyncifierVisitor : public DfsHloVisitorWithDefault {
 
   // Replace all dynamic-update-slice ops which update into host memory from
   // device memory with an asynchronous dynamic-update-slice.
-  Status HandleDynamicUpdateSlice(
+  absl::Status HandleDynamicUpdateSlice(
       HloInstruction* dynamic_update_slice) override {
     // Check that the dynamic-update-slice and its first two operands have
     // layouts. This pass must only be run after LayoutAssignment.
@@ -152,7 +152,7 @@ class HostMemoryTransferAsyncifierVisitor : public DfsHloVisitorWithDefault {
 
   // Replace all copy ops which copy from host memory to device memory or from
   // device memory to host memory with an asynchronous copy.
-  Status HandleCopy(HloInstruction* copy) override {
+  absl::Status HandleCopy(HloInstruction* copy) override {
     HloInstruction* operand = copy->mutable_operand(0);
     if (!operand->shape().has_layout()) {
       return InternalStrCat(operand->name(), " does not have a layout.");
diff --git a/third_party/xla/xla/service/host_offload_legalize.cc b/third_party/xla/xla/service/host_offload_legalize.cc
index fa881a1b271450..b460f93e41d015 100644
--- a/third_party/xla/xla/service/host_offload_legalize.cc
+++ b/third_party/xla/xla/service/host_offload_legalize.cc
@@ -144,9 +144,9 @@ absl::StatusOr<bool> DuplicateBroadcastForEachUse(HloModule* module) {
   return split_at_least_one;
 }
 
-// Walk up in the chain of memory offloaded instructions. Status not-ok when
-// an instructions not supported or end of chain reached.
-// Walks one instruction at a time.
+// Walk up in the chain of memory offloaded instructions. absl::Status not-ok
+// when an instructions not supported or end of chain reached. Walks one
+// instruction at a time.
 absl::StatusOr<std::pair<HloInstruction*, int>> WalkUpMemoryOffload(
     std::pair<HloInstruction*, int> current_value,
     const CallGraph& call_graph) {
@@ -220,10 +220,10 @@ absl::StatusOr<std::pair<HloInstruction*, int>> WalkUpMemoryOffload(
   }
 }
 
-// Walk down in the chain of memory offloaded instructions. Status not-ok when
-// an instructions not supported or end of chain reached.
-// Walks one instruction at a time, but returns multiple instructions for each
-// conforming user.
+// Walk down in the chain of memory offloaded instructions. absl::Status not-ok
+// when an instructions not supported or end of chain reached. Walks one
+// instruction at a time, but returns multiple instructions for each conforming
+// user.
 absl::StatusOr<std::vector<std::pair<HloInstruction*, int>>>
 WalkDownMemoryOffload(const std::pair<HloInstruction*, int64_t>& current_value,
                       const CallGraph& call_graph) {
@@ -232,7 +232,8 @@ WalkDownMemoryOffload(const std::pair<HloInstruction*, int64_t>& current_value,
   VLOG(5) << "Current value in progress: " << current_value.first->ToString()
           << " idx: " << current_value.second;
   std::vector<std::pair<HloInstruction*, int>> results;
-  auto add_gte_for_idx = [&results](HloInstruction* instr, int idx) -> Status {
+  auto add_gte_for_idx = [&results](HloInstruction* instr,
+                                    int idx) -> absl::Status {
     HloInstruction* gte = nullptr;
     for (HloInstruction* user : instr->users()) {
       if (user->opcode() != HloOpcode::kGetTupleElement) {
diff --git a/third_party/xla/xla/service/host_offloader.cc b/third_party/xla/xla/service/host_offloader.cc
index 77f5629915c52d..d1c12b8888d1f9 100644
--- a/third_party/xla/xla/service/host_offloader.cc
+++ b/third_party/xla/xla/service/host_offloader.cc
@@ -239,7 +239,8 @@ std::optional<HloInstruction*> FindAnnotationFromDUS(HloInstruction* hlo) {
 
 }  // namespace
 
-Status HostOffloader::HandleMoveToHostCustomCall(HloInstruction* custom_call) {
+absl::Status HostOffloader::HandleMoveToHostCustomCall(
+    HloInstruction* custom_call) {
   VLOG(2) << "Found a custom call annotating start-of-host-offload: "
           << custom_call->ToString();
   // Save a pointer to this custom call for when we want to remove it later.
@@ -279,7 +280,7 @@ Status HostOffloader::HandleMoveToHostCustomCall(HloInstruction* custom_call) {
   return OkStatus();
 }
 
-Status HostOffloader::MemoryOnlyOffloadStartingWithDus(
+absl::Status HostOffloader::MemoryOnlyOffloadStartingWithDus(
     const HloInstruction* dynamic_update_slice) {
   // The user wants to offload the data defined by this dynamic-update-slice.
   VLOG(2) << "Host memory offload starts with a dynamic-update-slice: "
@@ -435,7 +436,7 @@ void HostOffloader::AddAllPositionsToBeMovedToHostMemory(
   }
 }
 
-Status HostOffloader::MemoryOnlyOffloadStartingWithCopy(
+absl::Status HostOffloader::MemoryOnlyOffloadStartingWithCopy(
     const HloInstruction* copy) {
   // The user wants to offload the data defined by this copy.
   VLOG(2) << "Host memory offload starts with a copy: " << copy->name();
@@ -518,7 +519,7 @@ Status HostOffloader::MemoryOnlyOffloadStartingWithCopy(
   return OkStatus();
 }
 
-Status HostOffloader::MemoryOnlyOffloadInsertCopies(
+absl::Status HostOffloader::MemoryOnlyOffloadInsertCopies(
     HloInstruction* custom_call) {
   VLOG(3) << "Found an offload annotation (" << custom_call->name()
           << "). Expecting that we'll need to insert copies";
@@ -556,7 +557,7 @@ Status HostOffloader::MemoryOnlyOffloadInsertCopies(
   return OkStatus();
 }
 
-Status HostOffloader::DynamifySlice(HloInstruction* slice) {
+absl::Status HostOffloader::DynamifySlice(HloInstruction* slice) {
   VLOG(3) << "Dynamifying slice " << slice->ToString();
   std::vector<HloInstruction*> start_constants;
   for (int64_t start : slice->slice_starts()) {
@@ -582,7 +583,8 @@ Status HostOffloader::DynamifySlice(HloInstruction* slice) {
 // Taking an instruction representing a move-to-device custom call, creates a
 // copy to device for that operand and replaces all uses of the operand of the
 // load annotation with the copy.
-Status HostOffloader::CreateCopyForInputStreaming(HloInstruction* custom_call) {
+absl::Status HostOffloader::CreateCopyForInputStreaming(
+    HloInstruction* custom_call) {
   HloInstruction* operand_of_load_annotation = custom_call->mutable_operand(0);
   Shape copy_shape = operand_of_load_annotation->shape();
   SetMemorySpace(&copy_shape, Layout::kDefaultMemorySpace);
@@ -631,7 +633,8 @@ Status HostOffloader::CreateCopyForInputStreaming(HloInstruction* custom_call) {
 
 // From a unique buffer on host memory, finds move-to-device custom calls
 // for this buffer and inserts the appropriate copies.
-Status HostOffloader::HandleStreamedBuffer(const HloBuffer& unique_buffer) {
+absl::Status HostOffloader::HandleStreamedBuffer(
+    const HloBuffer& unique_buffer) {
   // Find all move-to-device custom calls that are using this buffer.
   for (const HloValue* value : unique_buffer.values()) {
     // First, handle the defining instruction of this buffer, as a potential
@@ -684,7 +687,7 @@ Status HostOffloader::HandleStreamedBuffer(const HloBuffer& unique_buffer) {
 // corresponding move-to-device custom calls for these parameters. Once found,
 // adds these move-to-device custom calls to the expected host-to-device
 // annotations, and creates the necessary copies for input streaming.
-Status HostOffloader::HandleInputStreaming(HloComputation* computation) {
+absl::Status HostOffloader::HandleInputStreaming(HloComputation* computation) {
   const ComputationLayout& entry_computation_layout =
       computation->parent()->entry_computation_layout();
 
@@ -712,7 +715,7 @@ Status HostOffloader::HandleInputStreaming(HloComputation* computation) {
 // Starts from the result of the entry computation and looks for a case of
 // output streaming. This function will not change any hlo, it will only mark
 // instructions to be converted to host memory space.
-Status HostOffloader::HandleOutputStreaming(HloComputation* computation) {
+absl::Status HostOffloader::HandleOutputStreaming(HloComputation* computation) {
   const ComputationLayout& entry_computation_layout =
       computation->parent()->entry_computation_layout();
 
diff --git a/third_party/xla/xla/service/host_offloader.h b/third_party/xla/xla/service/host_offloader.h
index bc8ac288427025..f4fba494c2bfc9 100644
--- a/third_party/xla/xla/service/host_offloader.h
+++ b/third_party/xla/xla/service/host_offloader.h
@@ -72,33 +72,33 @@ class HostOffloader : public HloModulePass {
 
   // Process streamed inputs for the given computation, finding the relevant
   // move-to-device custom calls and inserting the appropriate copies.
-  Status HandleInputStreaming(HloComputation* computation);
+  absl::Status HandleInputStreaming(HloComputation* computation);
   // Process streamed outputs for the given computation, finding the relevant
   // move-to-host custom calls and inserting the appropriate copies.
-  Status HandleOutputStreaming(HloComputation* computation);
+  absl::Status HandleOutputStreaming(HloComputation* computation);
   // From a unique buffer on host memory, finds move-to-device custom calls
   // for this buffer and inserts the appropriate copies.
-  Status HandleStreamedBuffer(const HloBuffer& unique_buffer);
+  absl::Status HandleStreamedBuffer(const HloBuffer& unique_buffer);
   // Creates a copy to device for the input streaming custom call.
-  Status CreateCopyForInputStreaming(HloInstruction* custom_call);
+  absl::Status CreateCopyForInputStreaming(HloInstruction* custom_call);
   absl::StatusOr<bool> TryParameterStreaming(HloInstruction* custom_call);
   absl::StatusOr<bool> TryOutputStreaming(HloInstruction* custom_call);
-  Status HandleMoveToHostCustomCall(HloInstruction* custom_call);
+  absl::Status HandleMoveToHostCustomCall(HloInstruction* custom_call);
 
   // Handle memory-only offloading where the data is written to the host via a
   // dynamic-update-slice and is read back via a dynamic-slice.
-  Status MemoryOnlyOffloadStartingWithDus(
+  absl::Status MemoryOnlyOffloadStartingWithDus(
       const HloInstruction* dynamic_update_slice);
 
   // Handle memory-only offloading where the data is written to the host via a
   // copy and is read back via a copy.
-  Status MemoryOnlyOffloadStartingWithCopy(const HloInstruction* copy);
+  absl::Status MemoryOnlyOffloadStartingWithCopy(const HloInstruction* copy);
 
   // Handle memory-only offloading where there are no ops yet for data movement.
   // We will insert copies at the points where the annotations are.
-  Status MemoryOnlyOffloadInsertCopies(HloInstruction* custom_call);
+  absl::Status MemoryOnlyOffloadInsertCopies(HloInstruction* custom_call);
 
-  Status DynamifySlice(HloInstruction* slice);
+  absl::Status DynamifySlice(HloInstruction* slice);
 
   static constexpr std::array kAllowedPositionOpcodes = {
       HloOpcode::kBitcast,
diff --git a/third_party/xla/xla/service/indexed_array_analysis.cc b/third_party/xla/xla/service/indexed_array_analysis.cc
index 786060daa8b834..1b711e59082a9c 100644
--- a/third_party/xla/xla/service/indexed_array_analysis.cc
+++ b/third_party/xla/xla/service/indexed_array_analysis.cc
@@ -92,7 +92,7 @@ absl::StatusOr<Analysis::Array*> IndexedArrayAnalysis::GetArrayFor(
   return FindOrDie(cache_, instr);
 }
 
-Status IndexedArrayAnalysis::TraverseAndPopulateCache(
+absl::Status IndexedArrayAnalysis::TraverseAndPopulateCache(
     const HloInstruction* root) {
   // Depth first search over the DAG, invoking ComputeArrayFor in post order.
   // The HLO instructions already in the cache are considered leaves.
diff --git a/third_party/xla/xla/service/indexed_array_analysis.h b/third_party/xla/xla/service/indexed_array_analysis.h
index a65e8f185b1e52..ce56c6b287246c 100644
--- a/third_party/xla/xla/service/indexed_array_analysis.h
+++ b/third_party/xla/xla/service/indexed_array_analysis.h
@@ -253,7 +253,7 @@ class IndexedArrayAnalysis {
  private:
   // Helper function that ensures that every HLO instruction that is
   // transitively used by `root` has an entry in `cache_`.
-  Status TraverseAndPopulateCache(const HloInstruction* root);
+  absl::Status TraverseAndPopulateCache(const HloInstruction* root);
 
   // Creates an Array instance for `instr` under the assumption that all
   // operations of `instr` are present in `cache_`.
diff --git a/third_party/xla/xla/service/latency_hiding_scheduler.cc b/third_party/xla/xla/service/latency_hiding_scheduler.cc
index ffe96757859f99..83eb5c5ebffc38 100644
--- a/third_party/xla/xla/service/latency_hiding_scheduler.cc
+++ b/third_party/xla/xla/service/latency_hiding_scheduler.cc
@@ -1810,7 +1810,8 @@ void HloScheduleGraph::InitializeGraphAnalysis(
   }
 }
 
-Status DefaultSchedulerCore::InitializeScheduler(const HloModule* module) {
+absl::Status DefaultSchedulerCore::InitializeScheduler(
+    const HloModule* module) {
   TF_ASSIGN_OR_RETURN(alias_analysis_, HloAliasAnalysis::Run(module));
   module_pressure_state_ = std::make_unique<ModulePressureState>(
       module, alias_analysis_.get(), shape_size_bytes_);
@@ -1819,7 +1820,8 @@ Status DefaultSchedulerCore::InitializeScheduler(const HloModule* module) {
   return OkStatus();
 }
 
-Status DefaultSchedulerCore::SchedulingStep(SchedulingState* sched_state) {
+absl::Status DefaultSchedulerCore::SchedulingStep(
+    SchedulingState* sched_state) {
   // Get the first available node for scheduling that is the node that
   // satisfies our ready heuristic the best.
   TF_ASSIGN_OR_RETURN(HloGraphNode * node,
diff --git a/third_party/xla/xla/service/latency_hiding_scheduler.h b/third_party/xla/xla/service/latency_hiding_scheduler.h
index dd05f7ff521c2b..99e148cfb64018 100644
--- a/third_party/xla/xla/service/latency_hiding_scheduler.h
+++ b/third_party/xla/xla/service/latency_hiding_scheduler.h
@@ -277,7 +277,7 @@ class AsyncTracker {
 // Base class for the core scheduling algorithm.
 class SchedulerCore {
  public:
-  virtual Status InitializeScheduler(const HloModule* module) = 0;
+  virtual absl::Status InitializeScheduler(const HloModule* module) = 0;
   virtual absl::StatusOr<std::vector<HloInstruction*>> ScheduleComputation(
       const HloComputation* computation) = 0;
   virtual ~SchedulerCore() = default;
@@ -841,7 +841,7 @@ class DefaultSchedulerCore : public SchedulerCore {
         target_scheduling_rule_(target_scheduling_rule),
         early_target_scheduling_rule_(early_target_scheduling_rule),
         post_processing_fn_(post_processing_fn) {}
-  Status InitializeScheduler(const HloModule* module) override;
+  absl::Status InitializeScheduler(const HloModule* module) override;
   absl::StatusOr<std::vector<HloInstruction*>> ScheduleComputation(
       const HloComputation* computation) override;
   static bool AddOccupierToResource(
@@ -866,7 +866,7 @@ class DefaultSchedulerCore : public SchedulerCore {
       HloGraphNode* n, SchedulingState* sched_state) const;
   // Perform the scheduling of one or more instructions. Called every time the
   // ready set is not empty.
-  virtual Status SchedulingStep(SchedulingState* sched_state);
+  virtual absl::Status SchedulingStep(SchedulingState* sched_state);
   // Pick a node to schedule according to cost model.
   virtual absl::StatusOr<HloGraphNode*> FindAndExtractBestNodeAvailable(
       SchedulingState& sched_state,
diff --git a/third_party/xla/xla/service/layout_assignment.cc b/third_party/xla/xla/service/layout_assignment.cc
index bd03b70faa7cec..4335df8aa8642d 100644
--- a/third_party/xla/xla/service/layout_assignment.cc
+++ b/third_party/xla/xla/service/layout_assignment.cc
@@ -266,11 +266,11 @@ bool LayoutAssignment::AllOperandBuffersForwarded(
   });
 }
 
-Status LayoutAssignment::SetBufferLayout(const Layout& layout,
-                                         const LogicalBuffer& buffer,
-                                         bool mandatory, bool dfs,
-                                         int64_t priority,
-                                         const HloInstruction* user) {
+absl::Status LayoutAssignment::SetBufferLayout(const Layout& layout,
+                                               const LogicalBuffer& buffer,
+                                               bool mandatory, bool dfs,
+                                               int64_t priority,
+                                               const HloInstruction* user) {
   VLOG(3) << "SetBufferLayout : " << buffer << " : "
           << LayoutUtil::HumanString(layout) << " with priority " << priority
           << "; mandatory = " << mandatory << "; dfs = " << dfs << "\n";
@@ -341,10 +341,9 @@ Status LayoutAssignment::SetBufferLayout(const Layout& layout,
   return OkStatus();
 }
 
-Status LayoutAssignment::SetOperandLayout(const Shape& shape_with_layout,
-                                          const HloInstruction* instruction,
-                                          int64_t operand_no, bool mandatory,
-                                          bool dfs, int64_t priority) {
+absl::Status LayoutAssignment::SetOperandLayout(
+    const Shape& shape_with_layout, const HloInstruction* instruction,
+    int64_t operand_no, bool mandatory, bool dfs, int64_t priority) {
   if (shape_with_layout.IsArray() && shape_with_layout.rank() == 0) {
     return OkStatus();
   }
@@ -419,7 +418,7 @@ void LayoutAssignment::PushAddedConstraints(
   }
 }
 
-Status LayoutAssignment::SetArrayOperandLayout(
+absl::Status LayoutAssignment::SetArrayOperandLayout(
     const Layout& layout, const HloInstruction* instruction, int64_t operand_no,
     bool mandatory, bool dfs, int64_t priority) {
   const HloInstruction* operand = instruction->operand(operand_no);
@@ -431,7 +430,7 @@ Status LayoutAssignment::SetArrayOperandLayout(
                           priority);
 }
 
-Status LayoutAssignment::LayoutConstraints::SetResultLayout(
+absl::Status LayoutAssignment::LayoutConstraints::SetResultLayout(
     LayoutAssignment* assignment, const Shape& shape_with_layout,
     int64_t priority) {
   VLOG(3) << "  : " << ShapeUtil::HumanStringWithLayout(shape_with_layout)
@@ -443,11 +442,9 @@ Status LayoutAssignment::LayoutConstraints::SetResultLayout(
   return OkStatus();
 }
 
-Status LayoutAssignment::SetInstructionLayout(const Layout& layout,
-                                              const HloInstruction* instruction,
-                                              bool mandatory, bool dfs,
-                                              bool allow_alias,
-                                              int64_t priority) {
+absl::Status LayoutAssignment::SetInstructionLayout(
+    const Layout& layout, const HloInstruction* instruction, bool mandatory,
+    bool dfs, bool allow_alias, int64_t priority) {
   if (priority < 0) {
     priority = current_priority_;
   }
@@ -467,7 +464,7 @@ Status LayoutAssignment::SetInstructionLayout(const Layout& layout,
   return ShapeUtil::ForEachSubshapeWithStatus(
       instruction->shape(),
       [this, layout, instruction, mandatory, allow_alias, priority](
-          const Shape& subshape, const ShapeIndex& index) -> Status {
+          const Shape& subshape, const ShapeIndex& index) -> absl::Status {
         auto buffers =
             points_to_analysis_->GetPointsToSet(instruction).element(index);
         CHECK_EQ(1, buffers.size());
@@ -483,11 +480,9 @@ Status LayoutAssignment::SetInstructionLayout(const Layout& layout,
       });
 }
 
-Status LayoutAssignment::SetInstructionLayout(const Shape& shape_with_layout,
-                                              const HloInstruction* instruction,
-                                              bool mandatory, bool dfs,
-                                              bool allow_alias,
-                                              int64_t priority) {
+absl::Status LayoutAssignment::SetInstructionLayout(
+    const Shape& shape_with_layout, const HloInstruction* instruction,
+    bool mandatory, bool dfs, bool allow_alias, int64_t priority) {
   VLOG(3) << "SetInstructionLayout : " << instruction->name() << ", "
           << ShapeUtil::HumanStringWithLayout(shape_with_layout)
           << ": priority = " << priority << " : mandatory = " << mandatory
@@ -505,7 +500,7 @@ Status LayoutAssignment::SetInstructionLayout(const Shape& shape_with_layout,
   TF_RETURN_IF_ERROR(ShapeUtil::ForEachSubshapeWithStatus(
       shape_with_layout,
       [this, dfs, instruction, mandatory, allow_alias, priority](
-          const Shape& subshape, const ShapeIndex& index) -> Status {
+          const Shape& subshape, const ShapeIndex& index) -> absl::Status {
         auto buffers =
             points_to_analysis_->GetPointsToSet(instruction).element(index);
         CHECK_EQ(1, buffers.size());
@@ -614,7 +609,7 @@ bool IsHostSendRecv(const HloInstruction* instruction) {
 
 }  // namespace
 
-Status LayoutAssignment::BuildHostChannelConstraints(
+absl::Status LayoutAssignment::BuildHostChannelConstraints(
     HloComputation* computation) {
   for (auto* instruction : computation->instructions()) {
     const HloSendRecvInstruction* send_recv_instr =
@@ -649,9 +644,9 @@ bool IsLayoutConstrainedCustomCall(HloInstruction* instruction) {
   return custom_call != nullptr && custom_call->layout_constrained();
 }
 
-Status PropagateParameterLayoutToUsers(const HloInstruction* instruction,
-                                       const Shape& shape,
-                                       LayoutAssignment* constraints) {
+absl::Status PropagateParameterLayoutToUsers(const HloInstruction* instruction,
+                                             const Shape& shape,
+                                             LayoutAssignment* constraints) {
   for (auto* user : instruction->users()) {
     // Excluding tuple operations as they do not participate in layout
     // propagations (they do not create or aliase buffers).
@@ -679,7 +674,7 @@ Status PropagateParameterLayoutToUsers(const HloInstruction* instruction,
 
 }  // namespace
 
-Status LayoutAssignment::AddMandatoryConstraints(
+absl::Status LayoutAssignment::AddMandatoryConstraints(
     ChannelLayoutConstraints* channel_constraints,
     LayoutConstraints* constraints) {
   VLOG(2) << "Adding mandatory layout constraints to computation "
@@ -968,7 +963,7 @@ bool LayoutsInShapesEqual(const Shape& lhs, const Shape& rhs) {
 
 // Operands of layout-constrained custom calls must match the expected
 // constrained layouts.
-Status CheckCustomCallLayout(HloInstruction* instruction) {
+absl::Status CheckCustomCallLayout(HloInstruction* instruction) {
   if (IsLayoutConstrainedCustomCall(instruction)) {
     const HloCustomCallInstruction* custom_call =
         DynCast<HloCustomCallInstruction>(instruction);
@@ -987,9 +982,10 @@ Status CheckCustomCallLayout(HloInstruction* instruction) {
 //   (3) body computation parameter
 //   (4) body computation result
 //   (5) while instruction result
-Status CheckWhileLayout(HloInstruction* while_inst,
-                        const ComputationLayout& condition_computation_layout,
-                        const ComputationLayout& body_computation_layout) {
+absl::Status CheckWhileLayout(
+    HloInstruction* while_inst,
+    const ComputationLayout& condition_computation_layout,
+    const ComputationLayout& body_computation_layout) {
   auto init_shape = while_inst->operand(0)->shape();
   TF_RET_CHECK(
       condition_computation_layout.parameter_layout(0).MatchesLayoutInShape(
@@ -1002,12 +998,12 @@ Status CheckWhileLayout(HloInstruction* while_inst,
   return OkStatus();
 }
 
-Status CheckOptimizationBarrierLayout(HloInstruction* inst) {
+absl::Status CheckOptimizationBarrierLayout(HloInstruction* inst) {
   TF_RET_CHECK(LayoutsInShapesEqual(inst->operand(0)->shape(), inst->shape()));
   return OkStatus();
 }
 
-Status CheckConditionalLayout(
+absl::Status CheckConditionalLayout(
     HloInstruction* instruction,
     absl::Span<const ComputationLayout> branch_computation_layouts) {
   for (int j = 0; j < instruction->branch_count(); ++j) {
@@ -1035,7 +1031,7 @@ Status CheckConditionalLayout(
 // Fusion parameters must match the layout of the fusion instructions operands,
 // and the root of the fusion expression must match the layout of the fusion
 // instruction.
-Status CheckFusionLayout(HloInstruction* fusion) {
+absl::Status CheckFusionLayout(HloInstruction* fusion) {
   TF_RET_CHECK(HloOpcode::kFusion == fusion->opcode());
 
   TF_RET_CHECK(LayoutsInShapesEqual(fusion->shape(),
@@ -1049,8 +1045,8 @@ Status CheckFusionLayout(HloInstruction* fusion) {
 
 // The layout of a parameter must match the respective layout in the
 // computation's ComputationLayout.
-Status CheckParameterLayout(HloInstruction* parameter,
-                            const ComputationLayout& computation_layout) {
+absl::Status CheckParameterLayout(HloInstruction* parameter,
+                                  const ComputationLayout& computation_layout) {
   const ShapeLayout& parameter_layout =
       computation_layout.parameter_layout(parameter->parameter_number());
   return ShapeUtil::ForEachSubshapeWithStatus(
@@ -1074,7 +1070,7 @@ Status CheckParameterLayout(HloInstruction* parameter,
 }
 
 // The layout of a constant instruction must match the layout of its literal.
-Status CheckConstantLayout(HloInstruction* constant) {
+absl::Status CheckConstantLayout(HloInstruction* constant) {
   if (!LayoutsInShapesEqual(constant->literal().shape(), constant->shape())) {
     return Internal(
         "constant instruction %s does not match the layout of its literal %s",
@@ -1097,7 +1093,7 @@ Layout GetBroadcastLayoutFromOutput(const Layout& layout,
   return shape.layout();
 }
 
-Status CheckBroadcastLayout(HloInstruction* broadcast) {
+absl::Status CheckBroadcastLayout(HloInstruction* broadcast) {
   CHECK_EQ(broadcast->opcode(), HloOpcode::kBroadcast);
   Shape shape = ShapeUtil::FilterDimensions(
       [&](int64_t dim) {
@@ -1114,7 +1110,7 @@ Status CheckBroadcastLayout(HloInstruction* broadcast) {
 
 }  // namespace
 
-Status LayoutAssignment::CheckCallLayout(
+absl::Status LayoutAssignment::CheckCallLayout(
     HloInstruction* call, const ComputationLayout& computation_layout) {
   HloComputation* computation = call->to_apply();
   TF_RET_CHECK(computation->num_parameters() == call->operand_count());
@@ -1188,7 +1184,7 @@ absl::StatusOr<HloInstruction*> LayoutAssignment::CreateCopyWithNewLayout(
 // Creates a copy of the given operand if the operand's layout does not match
 // the given layout. This copy replaces the use in the given instruction. Tuple
 // operands will be deep-copied.
-Status LayoutAssignment::CopyOperandIfLayoutsDiffer(
+absl::Status LayoutAssignment::CopyOperandIfLayoutsDiffer(
     const ShapeLayout& operand_layout, HloInstruction* instruction,
     int64_t operand_no) {
   HloInstruction* operand = instruction->mutable_operand(operand_no);
@@ -1269,7 +1265,7 @@ void LayoutAssignment::SetupCopiedInstruction(const HloInstruction& instruction,
   copy->set_frontend_attributes(instruction.frontend_attributes());
 }
 
-Status LayoutAssignment::CheckLayouts(
+absl::Status LayoutAssignment::CheckLayouts(
     HloModule* module,
     const absl::flat_hash_set<absl::string_view>& execution_threads) {
   TF_ASSIGN_OR_RETURN(auto points_to_analysis,
@@ -1288,8 +1284,9 @@ Status LayoutAssignment::CheckLayouts(
       const PointsToSet& points_to_set =
           points_to_analysis->GetPointsToSet(instruction);
       TF_RETURN_IF_ERROR(points_to_set.ForEachElementWithStatus(
-          [&instruction](ShapeIndex index,
-                         const PointsToSet::BufferList& buffers) -> Status {
+          [&instruction](
+              ShapeIndex index,
+              const PointsToSet::BufferList& buffers) -> absl::Status {
             if (ShapeUtil::IsLeafIndex(instruction->shape(), index)) {
               const Shape& instruction_subshape =
                   ShapeUtil::GetSubshape(instruction->shape(), index);
@@ -1602,7 +1599,8 @@ std::unique_ptr<Layout> LayoutAssignment::ChooseOutputLayoutFromOperandLayout(
   return nullptr;
 }
 
-Status LayoutAssignment::PropagateConstraints(LayoutConstraints* constraints) {
+absl::Status LayoutAssignment::PropagateConstraints(
+    LayoutConstraints* constraints) {
   // Gathers all initial constraints in a worklist and propagates them in
   // depth-first order. DFS order seems to be better than BFS because a
   // constraint is propagated as far as possible before propagating unrelated
@@ -1683,7 +1681,7 @@ std::vector<std::pair<const HloInstruction*, int64_t>> GetArrayUsesOfBuffer(
 
 }  // namespace
 
-Status LayoutAssignment::PropagateUseConstraintToDefs(
+absl::Status LayoutAssignment::PropagateUseConstraintToDefs(
     const ShapeLayout& shape_layout, const HloInstruction* instruction,
     LayoutConstraints* constraints, int64_t priority,
     const HloInstruction* user) {
@@ -1694,7 +1692,7 @@ Status LayoutAssignment::PropagateUseConstraintToDefs(
   return points_to_set.ForEachElementWithStatus(
       [&shape_layout, this, priority, user](
           const ShapeIndex& index,
-          const PointsToSet::BufferList& buffers) -> Status {
+          const PointsToSet::BufferList& buffers) -> absl::Status {
         const auto& subshape =
             ShapeUtil::GetSubshape(shape_layout.shape(), index);
         if (ShapeUtil::IsLeafIndex(shape_layout.shape(), index) &&
@@ -1736,7 +1734,7 @@ bool InstructionShouldPropagateDepthFirst(const HloInstruction& hlo) {
 
 }  // namespace
 
-Status LayoutAssignment::PropagateOperandConstraint(
+absl::Status LayoutAssignment::PropagateOperandConstraint(
     const OperandLayoutConstraint& operand_constraint,
     LayoutConstraints* constraints) {
   if (IsAtMostRank1(operand_constraint.operand()->shape())) {
@@ -1900,7 +1898,7 @@ Status LayoutAssignment::PropagateOperandConstraint(
   return OkStatus();
 }
 
-Status LayoutAssignment::PropagateBufferConstraintToOperands(
+absl::Status LayoutAssignment::PropagateBufferConstraintToOperands(
     const BufferLayoutConstraint& buffer_constraint,
     LayoutConstraints* constraints) {
   const LogicalBuffer& buffer = buffer_constraint.buffer();
@@ -1965,7 +1963,7 @@ Status LayoutAssignment::PropagateBufferConstraintToOperands(
   return OkStatus();
 }
 
-Status LayoutAssignment::PropagateBufferConstraint(
+absl::Status LayoutAssignment::PropagateBufferConstraint(
     const BufferLayoutConstraint& buffer_constraint,
     LayoutConstraints* constraints) {
   // Only propagate array layouts.
@@ -1978,7 +1976,7 @@ Status LayoutAssignment::PropagateBufferConstraint(
   return PropagateBufferConstraintToUses(buffer_constraint, constraints);
 }
 
-Status LayoutAssignment::PropagateBufferConstraintToUses(
+absl::Status LayoutAssignment::PropagateBufferConstraintToUses(
     const BufferLayoutConstraint& buffer_constraint,
     LayoutConstraints* constraints) {
   VLOG(5) << "PropagateBufferConstraintToUses: "
@@ -2036,7 +2034,7 @@ Status LayoutAssignment::PropagateBufferConstraintToUses(
   return OkStatus();
 }
 
-Status LayoutAssignment::PropagateResultConstraint(
+absl::Status LayoutAssignment::PropagateResultConstraint(
     const ComputationLayoutConstraint& layout_constraint,
     LayoutConstraints* constraints) {
   ShapeLayout result_layout =
@@ -2112,7 +2110,7 @@ namespace {
 // to match the layout of its corresponding fusion instruction operand. Also,
 // set the layout of the fused root to match the layout of the fusion
 // instruction itself.
-Status SetFusionLayouts(HloInstruction* fusion) {
+absl::Status SetFusionLayouts(HloInstruction* fusion) {
   TF_RET_CHECK(fusion->opcode() == HloOpcode::kFusion);
   for (auto* fused_instruction :
        fusion->fused_instructions_computation()->MakeInstructionPostOrder()) {
@@ -2156,7 +2154,7 @@ Status SetFusionLayouts(HloInstruction* fusion) {
 
 }  // namespace
 
-Status LayoutAssignment::AssignLayouts(LayoutConstraints& constraints) {
+absl::Status LayoutAssignment::AssignLayouts(LayoutConstraints& constraints) {
   HloComputation* computation = constraints.computation();
   VLOG(2) << "Assigning layouts to computation: " << computation->name();
 
@@ -2284,7 +2282,7 @@ Status LayoutAssignment::AssignLayouts(LayoutConstraints& constraints) {
   return OkStatus();
 }
 
-Status LayoutAssignment::CalculateComputationLayout(
+absl::Status LayoutAssignment::CalculateComputationLayout(
     LayoutConstraints* constraints) {
   // Process instructions that contain nested computations and may require
   // additional layouts to be assigned on the instructions nested inside.
@@ -2308,10 +2306,11 @@ Status LayoutAssignment::CalculateComputationLayout(
     return change;
   };
 
-  auto SetCalleeLayout =
-      [this, UpdateLayout](const HloInstruction* result,
-                           absl::Span<const HloInstruction* const> operands,
-                           LayoutConstraints* callee, int priority) -> Status {
+  auto SetCalleeLayout = [this, UpdateLayout](
+                             const HloInstruction* result,
+                             absl::Span<const HloInstruction* const> operands,
+                             LayoutConstraints* callee,
+                             int priority) -> absl::Status {
     CHECK_NE(result, nullptr);
     ComputationLayoutConstraint* callee_constraint =
         callee->mutable_computation_constraint();
@@ -2409,7 +2408,8 @@ Status LayoutAssignment::CalculateComputationLayout(
   return OkStatus();
 }
 
-Status LayoutAssignment::ClearComputationLayouts(HloComputation* computation) {
+absl::Status LayoutAssignment::ClearComputationLayouts(
+    HloComputation* computation) {
   // Clear existing layouts of the instructions.  All layouts must be assigned
   // by the LayoutAssignment pass, except for those on parameters, the
   // computation result, and a couple special cases. The former two are
@@ -2435,7 +2435,7 @@ Status LayoutAssignment::ClearComputationLayouts(HloComputation* computation) {
   return OkStatus();
 }
 
-Status LayoutAssignment::RunOnComputation(
+absl::Status LayoutAssignment::RunOnComputation(
     LayoutConstraints* constraints,
     ChannelLayoutConstraints* channel_constraints) {
   HloComputation* computation = constraints->computation();
@@ -2517,7 +2517,7 @@ Status LayoutAssignment::RunOnComputation(
   return OkStatus();
 }
 
-Status LayoutAssignment::ConstrainChannelLayouts(
+absl::Status LayoutAssignment::ConstrainChannelLayouts(
     HloComputation* computation,
     ChannelLayoutConstraints* channel_constraints) {
   for (HloInstruction* instruction : computation->MakeInstructionPostOrder()) {
@@ -2532,7 +2532,7 @@ Status LayoutAssignment::ConstrainChannelLayouts(
   return OkStatus();
 }
 
-Status LayoutAssignment::PropagateComputationLayouts(
+absl::Status LayoutAssignment::PropagateComputationLayouts(
     HloComputation* computation, ComputationLayout* computation_layout) {
   ComputationLayout computed_computation_layout(
       computation->ComputeProgramShape(),
@@ -2947,7 +2947,7 @@ bool LayoutAssignment::IsAtMostRank1(const Shape& shape) {
   });
 }
 
-Status LayoutAssignment::Init(HloModule* module) {
+absl::Status LayoutAssignment::Init(HloModule* module) {
   computation_layouts_.clear();
   conditional_mismatch_.clear();
   current_priority_ = LayoutConstraint::kBeginningPriority;
@@ -2978,7 +2978,7 @@ Status LayoutAssignment::Init(HloModule* module) {
   return OkStatus();
 }
 
-Status LayoutAssignment::ClearPreviousPassSideEffects(
+absl::Status LayoutAssignment::ClearPreviousPassSideEffects(
     HloModule* module,
     const absl::flat_hash_set<absl::string_view>& execution_threads) {
   VLOG(5) << "Clearing previous side effects";
@@ -2993,8 +2993,8 @@ Status LayoutAssignment::ClearPreviousPassSideEffects(
   buffer_sets_cache_.clear();
   return OkStatus();
 }
-Status LayoutAssignment::AddCopyForOperand(HloInstruction* instruction,
-                                           int64_t operand_number) {
+absl::Status LayoutAssignment::AddCopyForOperand(HloInstruction* instruction,
+                                                 int64_t operand_number) {
   HloInstruction* operand = instruction->mutable_operand(operand_number);
   if (operand->opcode() != HloOpcode::kCopy || operand->user_count() > 1) {
     HloInstruction* copy =
diff --git a/third_party/xla/xla/service/layout_assignment.h b/third_party/xla/xla/service/layout_assignment.h
index 22586493917105..2d774be6aee485 100644
--- a/third_party/xla/xla/service/layout_assignment.h
+++ b/third_party/xla/xla/service/layout_assignment.h
@@ -299,8 +299,9 @@ class LayoutAssignment : public HloModulePass {
     OperandLayoutConstraint* InsertOperandLayoutConstraint(
         const HloInstruction* instruction, int64_t operand_no,
         const OperandLayoutConstraint& constraint);
-    Status SetResultLayout(LayoutAssignment* assignment,
-                           const Shape& shape_with_layout, int64_t priority);
+    absl::Status SetResultLayout(LayoutAssignment* assignment,
+                                 const Shape& shape_with_layout,
+                                 int64_t priority);
 
     const ComputationLayout& computation_layout() const {
       return computation_constraint_.computation_layout();
@@ -357,57 +358,63 @@ class LayoutAssignment : public HloModulePass {
   static bool IsAtMostRank1(const Shape& shape);
   // Convenience wrapper around SetOperandLayout for setting the layout of a
   // operand using a Layout object. The operand must be array-shaped.
-  Status SetArrayOperandLayout(const Layout& layout,
-                               const HloInstruction* instruction,
-                               int64_t operand_no, bool mandatory = true,
-                               bool dfs = true) {
+  absl::Status SetArrayOperandLayout(const Layout& layout,
+                                     const HloInstruction* instruction,
+                                     int64_t operand_no, bool mandatory = true,
+                                     bool dfs = true) {
     return SetArrayOperandLayout(layout, instruction, operand_no, mandatory,
                                  dfs, current_priority_);
   }
-  Status SetArrayOperandLayout(const Layout& layout,
-                               const HloInstruction* instruction,
-                               int64_t operand_no, bool mandatory, bool dfs,
-                               int64_t priority);
+  absl::Status SetArrayOperandLayout(const Layout& layout,
+                                     const HloInstruction* instruction,
+                                     int64_t operand_no, bool mandatory,
+                                     bool dfs, int64_t priority);
   // Convenience wrapper around SetBufferLayout. Sets the layouts of all buffers
   // created by the instruction to the layouts in the given shape. The
   // instruction must define every logical buffer in its output.
   // If `allow_alias` is false, the function will check that all output buffers
   // are defined by `instruction`, not aliased to an instruction elsewhere.
-  Status SetInstructionLayout(const Shape& shape_with_layout,
-                              const HloInstruction* instruction,
-                              bool mandatory = true, bool dfs = true,
-                              bool allow_alias = false) {
+  absl::Status SetInstructionLayout(const Shape& shape_with_layout,
+                                    const HloInstruction* instruction,
+                                    bool mandatory = true, bool dfs = true,
+                                    bool allow_alias = false) {
     return SetInstructionLayout(shape_with_layout, instruction, mandatory, dfs,
                                 allow_alias, current_priority_);
   }
-  Status SetInstructionLayout(const Shape& shape_with_layout,
-                              const HloInstruction* instruction, bool mandatory,
-                              bool dfs, bool allow_alias, int64_t priority);
+  absl::Status SetInstructionLayout(const Shape& shape_with_layout,
+                                    const HloInstruction* instruction,
+                                    bool mandatory, bool dfs, bool allow_alias,
+                                    int64_t priority);
   // Set the same given layout across all components of the instruction output.
   // It works the same as the API above if the output is a single array.
-  Status SetInstructionLayout(const Layout& layout,
-                              const HloInstruction* instruction,
-                              bool mandatory = true, bool dfs = true,
-                              bool allow_alias = false, int64_t priority = -1);
+  absl::Status SetInstructionLayout(const Layout& layout,
+                                    const HloInstruction* instruction,
+                                    bool mandatory = true, bool dfs = true,
+                                    bool allow_alias = false,
+                                    int64_t priority = -1);
   // Add a constraint on the layout of a LogicalBuffer, the layout of the
   // operand of the instruction, or the layout of the result of the computation,
   // respectively.
-  Status SetBufferLayout(const Layout& layout, const LogicalBuffer& buffer,
-                         bool mandatory = true, bool dfs = true) {
+  absl::Status SetBufferLayout(const Layout& layout,
+                               const LogicalBuffer& buffer,
+                               bool mandatory = true, bool dfs = true) {
     return SetBufferLayout(layout, buffer, mandatory, dfs, current_priority_);
   }
-  Status SetBufferLayout(const Layout& layout, const LogicalBuffer& buffer,
-                         bool mandatory, bool dfs, int64_t priority,
-                         const HloInstruction* from_user = nullptr);
-  Status SetOperandLayout(const Shape& shape_with_layout,
-                          const HloInstruction* instruction, int64_t operand_no,
-                          bool mandatory = true, bool dfs = true) {
+  absl::Status SetBufferLayout(const Layout& layout,
+                               const LogicalBuffer& buffer, bool mandatory,
+                               bool dfs, int64_t priority,
+                               const HloInstruction* from_user = nullptr);
+  absl::Status SetOperandLayout(const Shape& shape_with_layout,
+                                const HloInstruction* instruction,
+                                int64_t operand_no, bool mandatory = true,
+                                bool dfs = true) {
     return SetOperandLayout(shape_with_layout, instruction, operand_no,
                             mandatory, dfs, current_priority_);
   }
-  Status SetOperandLayout(const Shape& shape_with_layout,
-                          const HloInstruction* instruction, int64_t operand_no,
-                          bool mandatory, bool dfs, int64_t priority);
+  absl::Status SetOperandLayout(const Shape& shape_with_layout,
+                                const HloInstruction* instruction,
+                                int64_t operand_no, bool mandatory, bool dfs,
+                                int64_t priority);
   bool reverse_computation_order() const { return reverse_computation_order_; }
 
   ComputationLayout& saved_entry_computation_layout() {
@@ -450,13 +457,13 @@ class LayoutAssignment : public HloModulePass {
   //
   // Backends can override these methods with backend-specific propagation
   // rules.
-  virtual Status PropagateBufferConstraint(
+  virtual absl::Status PropagateBufferConstraint(
       const BufferLayoutConstraint& buffer_constraint,
       LayoutConstraints* constraints);
-  virtual Status PropagateOperandConstraint(
+  virtual absl::Status PropagateOperandConstraint(
       const OperandLayoutConstraint& operand_constraint,
       LayoutConstraints* constraints);
-  virtual Status PropagateResultConstraint(
+  virtual absl::Status PropagateResultConstraint(
       const ComputationLayoutConstraint& layout_constraint,
       LayoutConstraints* constraints);
 
@@ -465,11 +472,11 @@ class LayoutAssignment : public HloModulePass {
   }
   // Called after layouts of an instruction have been finalized to allow
   // subclasses to check for platform specific assumptions.
-  virtual Status Verify(const HloInstruction* instruction) {
+  virtual absl::Status Verify(const HloInstruction* instruction) {
     return OkStatus();
   }
 
-  Status PropagateUnconstraintedBuffers(LayoutConstraints* constraints);
+  absl::Status PropagateUnconstraintedBuffers(LayoutConstraints* constraints);
   const BufferLayoutConstraint* GetBufferLayoutConstraint(
       const LogicalBuffer& buffer) const;
   absl::StatusOr<const BufferLayoutConstraint*>
@@ -491,18 +498,17 @@ class LayoutAssignment : public HloModulePass {
                                           const ShapeIndex& index);
 
   // Propagates a buffer layout constraint into the operands that use it.
-  Status PropagateBufferConstraintToUses(
+  absl::Status PropagateBufferConstraintToUses(
       const BufferLayoutConstraint& buffer_constraint,
       LayoutConstraints* constraints);
 
   // Propagates a layout constraint on the use of the result of the given
   // instruction to the definitions of the LogicalBuffers which make up the
   // result.
-  Status PropagateUseConstraintToDefs(const ShapeLayout& shape_layout,
-                                      const HloInstruction* instruction,
-                                      LayoutConstraints* constraints,
-                                      int64_t priority,
-                                      const HloInstruction* user = nullptr);
+  absl::Status PropagateUseConstraintToDefs(
+      const ShapeLayout& shape_layout, const HloInstruction* instruction,
+      LayoutConstraints* constraints, int64_t priority,
+      const HloInstruction* user = nullptr);
 
   // Chooses a layout of operand `operand_no` of `instruction` that minimizes
   // the cost of `instruction`. `output_layout` is the layout of `instruction`.
@@ -546,17 +552,18 @@ class LayoutAssignment : public HloModulePass {
   // The operands of a call must match the layouts of parameters in the
   // ComputationLayout, and the call instruction itself must match the result
   // layout in the ComputationLayout.
-  Status CheckCallLayout(HloInstruction* call,
-                         const ComputationLayout& computation_layout);
+  absl::Status CheckCallLayout(HloInstruction* call,
+                               const ComputationLayout& computation_layout);
 
  private:
   // Initializes the layout assignment object for a new Run() call.
-  Status Init(HloModule* module);
+  absl::Status Init(HloModule* module);
 
   // Adds constraints which must be satisfied for correctness on all
   // backends. Called once prior to propagating constraints.
-  Status AddMandatoryConstraints(ChannelLayoutConstraints* channel_constraints,
-                                 LayoutConstraints* constraints);
+  absl::Status AddMandatoryConstraints(
+      ChannelLayoutConstraints* channel_constraints,
+      LayoutConstraints* constraints);
 
   // Return a vector containing the constraints which have been added to the
   // LayoutConstraints object since the construction of the object or since the
@@ -573,7 +580,7 @@ class LayoutAssignment : public HloModulePass {
   // layout of the instructions of a computation. This method is called after
   // all mandatory constraints have been added via AddMandatoryConstraints
   // and before propagating constraints.
-  virtual Status AddBackendConstraints(LayoutConstraints* constraints) {
+  virtual absl::Status AddBackendConstraints(LayoutConstraints* constraints) {
     return OkStatus();
   }
 
@@ -583,40 +590,40 @@ class LayoutAssignment : public HloModulePass {
   // computation instruction constraints.
   // Layouts constraints are added, then propagated until all LogicalBuffers in
   // the computation are constrained.
-  Status RunOnComputation(LayoutConstraints* constraints,
-                          ChannelLayoutConstraints* channel_constraints);
+  absl::Status RunOnComputation(LayoutConstraints* constraints,
+                                ChannelLayoutConstraints* channel_constraints);
 
   // Assign layouts to the instructions of a computation which satisfy the given
   // layout constraints. Copies may be added to satisfy the constraints. The
   // given LayoutConstraints must have layout constraints every logical buffer
   // in the computation.
-  Status AssignLayouts(LayoutConstraints& constraints);
+  absl::Status AssignLayouts(LayoutConstraints& constraints);
 
   // Propagates layout constraints from a set of initial constraints in order to
   // minimize the local cost of the computation. This propagation is *not*
   // required for correctness.
-  Status PropagateConstraints(LayoutConstraints* constraints);
+  absl::Status PropagateConstraints(LayoutConstraints* constraints);
 
-  Status PropagateBufferConstraintToOperands(
+  absl::Status PropagateBufferConstraintToOperands(
       const BufferLayoutConstraint& buffer_constraint,
       LayoutConstraints* constraints);
 
   // Check that all layouts in the module have been set and satisfy all
   // necessary conditions.
-  Status CheckLayouts(
+  absl::Status CheckLayouts(
       HloModule* module,
       const absl::flat_hash_set<absl::string_view>& execution_threads);
 
   // Computes the ComputationLayout of the given constraints based of the
   // layouts assigned to parameters and root instruction. Also propagate
   // constraints to computation nested inside.
-  Status CalculateComputationLayout(LayoutConstraints* constraints);
+  absl::Status CalculateComputationLayout(LayoutConstraints* constraints);
 
   // Clears all the layouts which can be cleared within a computation.
-  Status ClearComputationLayouts(HloComputation* computation);
+  absl::Status ClearComputationLayouts(HloComputation* computation);
 
   // Clears the side effects of a previous pass, like added copy instructions.
-  Status ClearPreviousPassSideEffects(
+  absl::Status ClearPreviousPassSideEffects(
       HloModule* module,
       const absl::flat_hash_set<absl::string_view>& execution_threads);
 
@@ -625,8 +632,8 @@ class LayoutAssignment : public HloModulePass {
   // This API propagates missing layout, and also checks that the caller
   // specified have been respected, by comparing those with the parameters and
   // root computation instruction.
-  Status PropagateComputationLayouts(HloComputation* computation,
-                                     ComputationLayout* computation_layout);
+  absl::Status PropagateComputationLayouts(
+      HloComputation* computation, ComputationLayout* computation_layout);
 
   // The pointer to the ComputationLayout passed as constructor parameter.
   ComputationLayout* entry_computation_layout_;
@@ -658,9 +665,9 @@ class LayoutAssignment : public HloModulePass {
   // Creates a copy of the given operand if the operand's layout does not match
   // the given layout. This copy replaces the use in the given instruction.
   // Tuple operands will be deep-copied.
-  virtual Status CopyOperandIfLayoutsDiffer(const ShapeLayout& operand_layout,
-                                            HloInstruction* instruction,
-                                            int64_t operand_no);
+  virtual absl::Status CopyOperandIfLayoutsDiffer(
+      const ShapeLayout& operand_layout, HloInstruction* instruction,
+      int64_t operand_no);
 
   // Registers a copy instruction added by the layout assignment pass.
   void RegisterAddedCopy(HloInstruction* copy) {
@@ -671,13 +678,15 @@ class LayoutAssignment : public HloModulePass {
   // Adds a copy for the operand of an instruction, unless such operand is
   // already a copy, and has a single user (which is forcibly the instruction
   // itself).
-  Status AddCopyForOperand(HloInstruction* instruction, int64_t operand_number);
+  absl::Status AddCopyForOperand(HloInstruction* instruction,
+                                 int64_t operand_number);
 
   // Apply the channel layout constraints by populating the channel_constraints
   // data structure passed in at constructor time. Eventually adds copies in
   // case two ends of a channel ended up with a different leyout.
-  Status ConstrainChannelLayouts(HloComputation* computation,
-                                 ChannelLayoutConstraints* channel_constraints);
+  absl::Status ConstrainChannelLayouts(
+      HloComputation* computation,
+      ChannelLayoutConstraints* channel_constraints);
 
   // Resets the input ChannelLayoutConstraints to the original copy received
   // from the constructor input.
@@ -688,7 +697,7 @@ class LayoutAssignment : public HloModulePass {
   }
 
   // Adds constraints related to host Send/Recv instructions.
-  Status BuildHostChannelConstraints(HloComputation* computation);
+  absl::Status BuildHostChannelConstraints(HloComputation* computation);
 
   // Module points to analysis that can be updated for cloned computations.
   std::unique_ptr<TuplePointsToAnalysis> points_to_analysis_;
diff --git a/third_party/xla/xla/service/layout_assignment_test.cc b/third_party/xla/xla/service/layout_assignment_test.cc
index d4a3982d2e1a88..4e9e9f5a7d145c 100644
--- a/third_party/xla/xla/service/layout_assignment_test.cc
+++ b/third_party/xla/xla/service/layout_assignment_test.cc
@@ -481,7 +481,7 @@ class OperandsMustBeTheSameLayoutAssignment : public LayoutAssignment {
       : LayoutAssignment(entry_computation_layout) {}
 
  protected:
-  Status PropagateBufferConstraint(
+  absl::Status PropagateBufferConstraint(
       const BufferLayoutConstraint& buffer_constraint,
       LayoutConstraints* constraints) override {
     const LogicalBuffer& buffer = buffer_constraint.buffer();
@@ -803,7 +803,7 @@ TEST_F(LayoutAssignmentTest, LayoutAssignmentToTupleSiblingOperand) {
   ComputationLayout computation_layout(
       m->entry_computation()->ComputeProgramShape());
   LayoutAssignment layout_assignment(&computation_layout);
-  Status error_status = layout_assignment.Run(m.get()).status();
+  absl::Status error_status = layout_assignment.Run(m.get()).status();
   EXPECT_TRUE(error_status.ok());
 }
 
@@ -820,7 +820,7 @@ TEST_F(LayoutAssignmentTest, InternalErrorOnBitcast) {
   ComputationLayout computation_layout(
       m->entry_computation()->ComputeProgramShape());
   LayoutAssignment layout_assignment(&computation_layout);
-  Status error_status = layout_assignment.Run(m.get()).status();
+  absl::Status error_status = layout_assignment.Run(m.get()).status();
   EXPECT_FALSE(error_status.ok());
   EXPECT_THAT(
       error_status.message(),
@@ -1328,7 +1328,7 @@ ENTRY %CustomCallLayoutConstrainedTupleResult (p0: f32[4,4]) -> (f32[4,4]{1,0},
   ExpectTupleLayoutIs(custom_call->shape(), {{1, 0}, {0, 1}});
 }
 
-Status AssignLayoutsToComputation(
+absl::Status AssignLayoutsToComputation(
     HloModule* m, ChannelLayoutConstraints* channel_constraints = nullptr) {
   if (!m->entry_computation_layout().result_layout().LayoutIsSet()) {
     m->mutable_entry_computation_layout()
diff --git a/third_party/xla/xla/service/layout_normalization.cc b/third_party/xla/xla/service/layout_normalization.cc
index 9974feb9541724..d966b5442b7829 100644
--- a/third_party/xla/xla/service/layout_normalization.cc
+++ b/third_party/xla/xla/service/layout_normalization.cc
@@ -58,7 +58,7 @@ class LayoutNormalizationVisitor : public DfsHloRewriteVisitor {
       : custom_call_transformer_(custom_call_transformer) {}
 
   // To handle a constant, just give the literal data a new layout.
-  Status HandleConstant(HloInstruction* hlo) override {
+  absl::Status HandleConstant(HloInstruction* hlo) override {
     Literal& literal = *Cast<HloConstantInstruction>(hlo)->mutable_literal();
     if (literal.shape().IsTuple()) {
       // TODO(cheshire): Tuple constants.
@@ -83,7 +83,7 @@ class LayoutNormalizationVisitor : public DfsHloRewriteVisitor {
 
   // Slice is layout-preserving, so handling is analoguous to elementwise unary,
   // and transposing the elements inside the metadata.
-  Status HandleSlice(HloInstruction* hlo) override {
+  absl::Status HandleSlice(HloInstruction* hlo) override {
     HloInstruction* operand = hlo->mutable_operand(0);
     const Shape& s = hlo->shape();
     const Shape& operand_shape = operand->shape();
@@ -118,7 +118,7 @@ class LayoutNormalizationVisitor : public DfsHloRewriteVisitor {
   //
   // Bitcast to descending layout and then bitcast back to make sure that shapes
   // match.
-  Status DefaultAction(HloInstruction* hlo) override {
+  absl::Status DefaultAction(HloInstruction* hlo) override {
     if (!hlo->user_count()) {
       // The local postcondition does not have to apply to the case when there
       // are no users.
@@ -144,7 +144,7 @@ class LayoutNormalizationVisitor : public DfsHloRewriteVisitor {
   //
   // With respect to layouts, concatenations are simple, as they are
   // layout-preserving.
-  Status HandleConcatenate(HloInstruction* hlo) override {
+  absl::Status HandleConcatenate(HloInstruction* hlo) override {
     const Shape& s = hlo->shape();
     int64_t orig_concat_dim = hlo->dimensions(0);
 
@@ -166,7 +166,7 @@ class LayoutNormalizationVisitor : public DfsHloRewriteVisitor {
     return OkStatus();
   }
 
-  Status HandleReduceWindow(HloInstruction* hlo) override {
+  absl::Status HandleReduceWindow(HloInstruction* hlo) override {
     if (hlo->shape().IsTuple()) {
       // TODO(cheshire): Handle variadic reductions.
       return OkStatus();
@@ -212,7 +212,7 @@ class LayoutNormalizationVisitor : public DfsHloRewriteVisitor {
   // Into:
   //
   //  A{I} -> broadcast[S']{I} -> bitcast[S]{L'}
-  Status HandleBroadcast(HloInstruction* hlo) override {
+  absl::Status HandleBroadcast(HloInstruction* hlo) override {
     VLOG(3) << "Input broadcast: " << hlo->ToString();
     auto s = hlo->shape();
     auto operand = hlo->mutable_operand(0);
@@ -240,7 +240,7 @@ class LayoutNormalizationVisitor : public DfsHloRewriteVisitor {
     return OkStatus();
   }
 
-  Status HandleIota(HloInstruction* hlo) override {
+  absl::Status HandleIota(HloInstruction* hlo) override {
     VLOG(3) << "Input iota: " << hlo->ToString();
     auto s = hlo->shape();
     auto normalized_shape = Normalize(s);
@@ -258,7 +258,7 @@ class LayoutNormalizationVisitor : public DfsHloRewriteVisitor {
   }
 
   // BitcastConvert is only layout-preserving if it doesn't change the rank.
-  Status HandleBitcastConvert(HloInstruction* hlo) override {
+  absl::Status HandleBitcastConvert(HloInstruction* hlo) override {
     // If the rank isn't changing this is just an unary op.
     if (hlo->shape().rank() == hlo->operand(0)->shape().rank()) {
       return HandleElementwiseUnary(hlo);
@@ -277,7 +277,7 @@ class LayoutNormalizationVisitor : public DfsHloRewriteVisitor {
   //    H_0{I} -> U{I} -> B{L}
   //
   // where {I} denotes default layout.
-  Status HandleElementwiseUnary(HloInstruction* hlo) override {
+  absl::Status HandleElementwiseUnary(HloInstruction* hlo) override {
     auto s = hlo->shape();
     auto operand = hlo->mutable_operand(0);
     auto operand_shape = operand->shape();
@@ -331,7 +331,7 @@ class LayoutNormalizationVisitor : public DfsHloRewriteVisitor {
   //         B{I} - bitcast{L}
   //        /
   //  A2{I}
-  Status HandleElementwiseBinary(HloInstruction* hlo) override {
+  absl::Status HandleElementwiseBinary(HloInstruction* hlo) override {
     auto s = hlo->shape();
     auto a = hlo->mutable_operand(0);
     auto b = hlo->mutable_operand(1);
@@ -363,7 +363,7 @@ class LayoutNormalizationVisitor : public DfsHloRewriteVisitor {
   //
   // A{I} -> R [S']{I} -> bitcast[S]{L2}
   //
-  Status HandleReshape(HloInstruction* hlo) override {
+  absl::Status HandleReshape(HloInstruction* hlo) override {
     auto s = hlo->shape();
     auto operand = hlo->mutable_operand(0);
     TF_RET_CHECK(ShapeUtil::ReshapeIsBitcast(s, operand->shape()));
@@ -401,7 +401,7 @@ class LayoutNormalizationVisitor : public DfsHloRewriteVisitor {
   //
   // where dim_0 is dimensions of the original transposition, and `o` denotes
   // permutation composition.
-  Status HandleTranspose(HloInstruction* hlo) override {
+  absl::Status HandleTranspose(HloInstruction* hlo) override {
     auto s = hlo->shape();
     auto operand = hlo->mutable_operand(0);
     auto operand_s = operand->shape();
@@ -442,7 +442,7 @@ class LayoutNormalizationVisitor : public DfsHloRewriteVisitor {
   //
   // Where S' is normalization of [S]{L'}, and transposition dimensions are
   // given by L'.
-  Status HandleCopy(HloInstruction* hlo) override {
+  absl::Status HandleCopy(HloInstruction* hlo) override {
     VLOG(3) << "Processing copy: " << hlo->ToString();
     auto s = hlo->shape();
     auto operand = hlo->mutable_operand(0);
@@ -461,7 +461,7 @@ class LayoutNormalizationVisitor : public DfsHloRewriteVisitor {
   }
 
   // The reverse HLO has a list of dimensions it reverses.
-  Status HandleReverse(HloInstruction* hlo) override {
+  absl::Status HandleReverse(HloInstruction* hlo) override {
     auto s = hlo->shape();
     auto operand = hlo->mutable_operand(0);
     TF_ASSIGN_OR_RETURN(auto a0, GetNormalizedInput(operand));
@@ -484,7 +484,7 @@ class LayoutNormalizationVisitor : public DfsHloRewriteVisitor {
 
   // Padding is layout-preserving, so we only have to permute values inside the
   // padding config.
-  Status HandlePad(HloInstruction* hlo) override {
+  absl::Status HandlePad(HloInstruction* hlo) override {
     auto s = hlo->shape();
     auto operand = hlo->mutable_operand(0);
     auto padded_by = hlo->mutable_operand(1);
@@ -515,7 +515,7 @@ class LayoutNormalizationVisitor : public DfsHloRewriteVisitor {
     return OkStatus();
   }
 
-  Status HandleCustomCall(HloInstruction* hlo) override {
+  absl::Status HandleCustomCall(HloInstruction* hlo) override {
     if (custom_call_transformer_) {
       TF_ASSIGN_OR_RETURN(
           std::optional<HloInstruction*> transformed_custom_call,
@@ -531,14 +531,14 @@ class LayoutNormalizationVisitor : public DfsHloRewriteVisitor {
 
   // Pushes down bitcast across the ternary select operation: same logic as
   // HandleElementwiseBinary.
-  Status HandleSelect(HloInstruction* hlo) override {
+  absl::Status HandleSelect(HloInstruction* hlo) override {
     return HandleTernary(hlo);
   }
 
   // DyanmicSlice is layout-preserving, so handling is analoguous to elementwise
   // unary, and transposing the elements inside the metadata, as well as the
   // operands specifying dimension sizes.
-  Status HandleDynamicSlice(HloInstruction* hlo) override {
+  absl::Status HandleDynamicSlice(HloInstruction* hlo) override {
     const Shape& s = hlo->shape();
     HloInstruction* operand = hlo->mutable_operand(0);
     const Shape& operand_shape = operand->shape();
@@ -570,7 +570,7 @@ class LayoutNormalizationVisitor : public DfsHloRewriteVisitor {
     return OkStatus();
   }
 
-  Status HandleDynamicUpdateSlice(HloInstruction* hlo) override {
+  absl::Status HandleDynamicUpdateSlice(HloInstruction* hlo) override {
     const Shape& s = hlo->shape();
     HloInstruction* operand = hlo->mutable_operand(0);
     HloInstruction* update = hlo->mutable_operand(1);
@@ -599,13 +599,13 @@ class LayoutNormalizationVisitor : public DfsHloRewriteVisitor {
     return OkStatus();
   }
 
-  Status HandleClamp(HloInstruction* hlo) override {
+  absl::Status HandleClamp(HloInstruction* hlo) override {
     return HandleTernary(hlo);
   }
 
  private:
   // Replace clamp/select ternary operation with a normalized one.
-  Status HandleTernary(HloInstruction* hlo) {
+  absl::Status HandleTernary(HloInstruction* hlo) {
     Shape s = hlo->shape();
     HloOpcode opcode = hlo->opcode();
     TF_RET_CHECK(opcode == HloOpcode::kClamp || opcode == HloOpcode::kSelect);
diff --git a/third_party/xla/xla/service/llvm_compiler.h b/third_party/xla/xla/service/llvm_compiler.h
index cd2699c2074fd6..4bd0e8d7c9d24f 100644
--- a/third_party/xla/xla/service/llvm_compiler.h
+++ b/third_party/xla/xla/service/llvm_compiler.h
@@ -58,11 +58,11 @@ class LLVMCompiler : public Compiler {
   void RemovePostOptimizationHook() { user_post_optimization_hook_ = nullptr; }
 
   // Bring in
-  //   StatusOr<std::unique_ptr<Executable>> RunBackend(
+  //   absl::StatusOr<std::unique_ptr<Executable>> RunBackend(
   //       std::unique_ptr<HloModule> module,
   //       se::StreamExecutor* stream_exec,
   //       se::DeviceMemoryAllocator* device_allocator)
-  //   StatusOr<std::unique_ptr<HloModule>> RunHloPasses(
+  //   absl::StatusOr<std::unique_ptr<HloModule>> RunHloPasses(
   //       std::unique_ptr<HloModule> module,
   //       se::StreamExecutor* stream_exec,
   //       se::DeviceMemoryAllocator* device_allocator)
diff --git a/third_party/xla/xla/service/logical_buffer_analysis.cc b/third_party/xla/xla/service/logical_buffer_analysis.cc
index 38bda9003bbc1a..1ec1f488b05c7d 100644
--- a/third_party/xla/xla/service/logical_buffer_analysis.cc
+++ b/third_party/xla/xla/service/logical_buffer_analysis.cc
@@ -51,7 +51,7 @@ LogicalBufferAnalysis::Run(const HloModule* module) {
   return std::move(analysis);
 }
 
-Status LogicalBufferAnalysis::Analyze() {
+absl::Status LogicalBufferAnalysis::Analyze() {
   // Empirically we usually have a few more logical buffers than instructions,
   // so reserve 10% more than the number of instructions to avoid frequent
   // resizes.
@@ -95,7 +95,8 @@ void LogicalBufferAnalysis::NewLogicalBuffer(HloInstruction* instruction,
   logical_buffers_.push_back(std::move(buffer));
 }
 
-Status LogicalBufferAnalysis::DefaultAction(HloInstruction* hlo_instruction) {
+absl::Status LogicalBufferAnalysis::DefaultAction(
+    HloInstruction* hlo_instruction) {
   // Create a logical buffer for each output of the instruction.
   ShapeUtil::ForEachSubshape(
       hlo_instruction->shape(),
@@ -106,38 +107,38 @@ Status LogicalBufferAnalysis::DefaultAction(HloInstruction* hlo_instruction) {
   return OkStatus();
 }
 
-Status LogicalBufferAnalysis::HandleGetTupleElement(HloInstruction*) {
+absl::Status LogicalBufferAnalysis::HandleGetTupleElement(HloInstruction*) {
   // GetTupleElement does not create buffers.
   return OkStatus();
 }
 
-Status LogicalBufferAnalysis::HandleAddDependency(
+absl::Status LogicalBufferAnalysis::HandleAddDependency(
     HloInstruction* add_dependency) {
   // AddDependency just forwards the value of its zero-th operand and does not
   // create buffers.
   return OkStatus();
 }
 
-Status LogicalBufferAnalysis::HandleCopy(HloInstruction* copy) {
+absl::Status LogicalBufferAnalysis::HandleCopy(HloInstruction* copy) {
   // The top-level buffer (index={}) for kCopy is newly created, but all other
   // buffers (in the case of a tuple shape) come from the operand
   NewLogicalBuffer(copy, /*index=*/{});
   return OkStatus();
 }
 
-Status LogicalBufferAnalysis::HandleBitcast(HloInstruction*) {
+absl::Status LogicalBufferAnalysis::HandleBitcast(HloInstruction*) {
   // A kBitcast instruction aliases its operand. That is, the buffer of its
   // result *is* the buffer of its operand.
   return OkStatus();
 }
 
-Status LogicalBufferAnalysis::HandleDomain(HloInstruction*) {
+absl::Status LogicalBufferAnalysis::HandleDomain(HloInstruction*) {
   // A kDomain instruction aliases its operand. That is, the buffer of its
   // result *is* the buffer of its operand.
   return OkStatus();
 }
 
-Status LogicalBufferAnalysis::HandleRecvDone(HloInstruction* recv_done) {
+absl::Status LogicalBufferAnalysis::HandleRecvDone(HloInstruction* recv_done) {
   // RecvDone produces a two-element tuple containing the data value (which
   // aliases part of its operand) and a token. Only the tuple index table and
   // the token are defined by the RecvDone.
@@ -146,7 +147,7 @@ Status LogicalBufferAnalysis::HandleRecvDone(HloInstruction* recv_done) {
   return OkStatus();
 }
 
-Status LogicalBufferAnalysis::HandleSend(HloInstruction* send) {
+absl::Status LogicalBufferAnalysis::HandleSend(HloInstruction* send) {
   // Send creates new buffers for the top-level tuple, the context (tuple
   // element at {1}), and the token (tuple element at {2}). Tuple element at {0}
   // is an alias of the Send operand, so we don't need to create a new Logical
@@ -157,7 +158,8 @@ Status LogicalBufferAnalysis::HandleSend(HloInstruction* send) {
   return OkStatus();
 }
 
-Status LogicalBufferAnalysis::HandleCopyStart(HloInstruction* copy_start) {
+absl::Status LogicalBufferAnalysis::HandleCopyStart(
+    HloInstruction* copy_start) {
   // CopyStart defines the tuple, target buffer at index {0}, and context at
   // index {2}.
   NewLogicalBuffer(copy_start, /*index=*/{});
@@ -166,19 +168,20 @@ Status LogicalBufferAnalysis::HandleCopyStart(HloInstruction* copy_start) {
   return OkStatus();
 }
 
-Status LogicalBufferAnalysis::HandleCopyDone(HloInstruction* copy_done) {
+absl::Status LogicalBufferAnalysis::HandleCopyDone(HloInstruction* copy_done) {
   // The output of CopyDone aliases with operand {0}. CopyDone doesn't create
   // any buffers.
   return OkStatus();
 }
 
-Status LogicalBufferAnalysis::HandleTuple(HloInstruction* tuple) {
+absl::Status LogicalBufferAnalysis::HandleTuple(HloInstruction* tuple) {
   // A Tuple instruction only creates the top-level buffer.
   NewLogicalBuffer(tuple, /*index=*/{});
   return OkStatus();
 }
 
-Status LogicalBufferAnalysis::HandleCustomCall(HloInstruction* custom_call) {
+absl::Status LogicalBufferAnalysis::HandleCustomCall(
+    HloInstruction* custom_call) {
   auto ccall = Cast<HloCustomCallInstruction>(custom_call);
   absl::flat_hash_set<ShapeIndex> aliased_outputs;
   for (const auto& pair : ccall->output_to_operand_aliasing()) {
@@ -200,7 +203,7 @@ Status LogicalBufferAnalysis::HandleCustomCall(HloInstruction* custom_call) {
 // accommodate the output-operand aliases (similar to HandleCustomCall).
 // TODO (sacer): We might want to consider the pairs discovered by
 // GetFusionInstructionInPlaceInputOutputPairs() here as well.
-Status LogicalBufferAnalysis::HandleFusion(HloInstruction* fusion) {
+absl::Status LogicalBufferAnalysis::HandleFusion(HloInstruction* fusion) {
   auto cfusion = Cast<HloFusionInstruction>(fusion);
   absl::flat_hash_set<ShapeIndex> aliased_outputs;
   for (const auto& pair : cfusion->output_to_operand_aliasing()) {
diff --git a/third_party/xla/xla/service/logical_buffer_analysis.h b/third_party/xla/xla/service/logical_buffer_analysis.h
index b7ca006ea6b023..cb2bde21028167 100644
--- a/third_party/xla/xla/service/logical_buffer_analysis.h
+++ b/third_party/xla/xla/service/logical_buffer_analysis.h
@@ -46,7 +46,7 @@ class LogicalBufferAnalysis : public DfsHloVisitorWithDefault {
 
  private:
   explicit LogicalBufferAnalysis(const HloModule* module) : module_(module) {}
-  Status Analyze();
+  absl::Status Analyze();
 
   // The module this analysis is performed on.
   const HloModule* module_;
@@ -56,19 +56,20 @@ class LogicalBufferAnalysis : public DfsHloVisitorWithDefault {
   // accessed with GetBuffer.
   void NewLogicalBuffer(HloInstruction* instruction, const ShapeIndex& index);
 
-  Status DefaultAction(HloInstruction* hlo_instruction) override;
-  Status HandleTuple(HloInstruction* tuple) override;
-  Status HandleGetTupleElement(HloInstruction* get_tuple_element) override;
-  Status HandleBitcast(HloInstruction* bitcast) override;
-  Status HandleDomain(HloInstruction* domain) override;
-  Status HandleCopy(HloInstruction* copy) override;
-  Status HandleCopyStart(HloInstruction* copy_start) override;
-  Status HandleCopyDone(HloInstruction* copy_done) override;
-  Status HandleRecvDone(HloInstruction* recv_done) override;
-  Status HandleSend(HloInstruction* send) override;
-  Status HandleAddDependency(HloInstruction* add_dependency) override;
-  Status HandleCustomCall(HloInstruction* custom_call) override;
-  Status HandleFusion(HloInstruction* fusion) override;
+  absl::Status DefaultAction(HloInstruction* hlo_instruction) override;
+  absl::Status HandleTuple(HloInstruction* tuple) override;
+  absl::Status HandleGetTupleElement(
+      HloInstruction* get_tuple_element) override;
+  absl::Status HandleBitcast(HloInstruction* bitcast) override;
+  absl::Status HandleDomain(HloInstruction* domain) override;
+  absl::Status HandleCopy(HloInstruction* copy) override;
+  absl::Status HandleCopyStart(HloInstruction* copy_start) override;
+  absl::Status HandleCopyDone(HloInstruction* copy_done) override;
+  absl::Status HandleRecvDone(HloInstruction* recv_done) override;
+  absl::Status HandleSend(HloInstruction* send) override;
+  absl::Status HandleAddDependency(HloInstruction* add_dependency) override;
+  absl::Status HandleCustomCall(HloInstruction* custom_call) override;
+  absl::Status HandleFusion(HloInstruction* fusion) override;
 
   // A map from the buffer ID to the logical buffer
   std::vector<std::unique_ptr<LogicalBuffer>> logical_buffers_;
diff --git a/third_party/xla/xla/service/map_inliner.cc b/third_party/xla/xla/service/map_inliner.cc
index eea543fe2d6e21..8d5915a681870a 100644
--- a/third_party/xla/xla/service/map_inliner.cc
+++ b/third_party/xla/xla/service/map_inliner.cc
@@ -39,11 +39,11 @@ class MapInlinerVisitor : public DfsHloVisitorWithDefault {
       : computation_(computation) {}
 
   // Default visitor action is to do nothing and return OK.
-  Status DefaultAction(HloInstruction* /*hlo_instruction*/) override {
+  absl::Status DefaultAction(HloInstruction* /*hlo_instruction*/) override {
     return OkStatus();
   }
 
-  Status HandleMap(HloInstruction* map) override;
+  absl::Status HandleMap(HloInstruction* map) override;
 
   // Runs the visitor on a computation.
   absl::StatusOr<bool> Run(HloComputation* computation);
@@ -63,7 +63,7 @@ absl::StatusOr<bool> MapInlinerVisitor::Run(HloComputation* computation) {
   return changed_;
 }
 
-Status MapInlinerVisitor::HandleMap(HloInstruction* map) {
+absl::Status MapInlinerVisitor::HandleMap(HloInstruction* map) {
   HloComputation* function = map->to_apply();
   HloInstruction& root = *function->root_instruction();
   // Only inlining functions that are simply a single operation until a better
diff --git a/third_party/xla/xla/service/mapped_ptr_container_sorter.h b/third_party/xla/xla/service/mapped_ptr_container_sorter.h
index 1b1dc5df19d5cc..e55f49179a3316 100644
--- a/third_party/xla/xla/service/mapped_ptr_container_sorter.h
+++ b/third_party/xla/xla/service/mapped_ptr_container_sorter.h
@@ -94,9 +94,9 @@ class MappedPtrContainerSorter {
   // - unmapped_index() returns an invalid index
   // - An internal error occurs. (This should theoretically not happen.)
   template <typename OrderedTy, typename UnorderedTy>
-  static Status Sort(MapPtrFn map_ptr, UnmappedPtrIndexFn unmapped_index,
-                     const OrderedTy& ordered_container,
-                     UnorderedTy& unordered_container);
+  static absl::Status Sort(MapPtrFn map_ptr, UnmappedPtrIndexFn unmapped_index,
+                           const OrderedTy& ordered_container,
+                           UnorderedTy& unordered_container);
 
  private:
   // A class for sorting the indices of the unordered_container.
@@ -114,8 +114,8 @@ class MappedPtrContainerSorter {
     // Specify the partial ordering value of a mapped element from the
     // unordered container. The partial ordering is amongst other mapped
     // elements.
-    Status AddMappedElement(size_t unordered_container_index,
-                            size_t partial_order);
+    absl::Status AddMappedElement(size_t unordered_container_index,
+                                  size_t partial_order);
 
     // Specify the index (amongst mapped elements), where an unmapped element
     // should be inserted. The unmapped element is inserted just after the
@@ -227,7 +227,8 @@ MappedPtrContainerSorter<PointedToTy>::InvalidIndexFn() {
 }
 
 template <typename PointedToTy>
-Status MappedPtrContainerSorter<PointedToTy>::SortedIndices::AddMappedElement(
+absl::Status
+MappedPtrContainerSorter<PointedToTy>::SortedIndices::AddMappedElement(
     size_t unordered_container_index, size_t partial_order) {
   if (partial_order >= mapped_element_indices_by_partial_order_.size()) {
     return InternalStrCat("invalid partial order: ", partial_order, " v max(",
@@ -440,7 +441,7 @@ void MappedPtrContainerSorter<PointedToTy>::Reorder(
 
 template <typename PointedToTy>
 template <typename OrderedTy, typename UnorderedTy>
-Status MappedPtrContainerSorter<PointedToTy>::Sort(
+absl::Status MappedPtrContainerSorter<PointedToTy>::Sort(
     MapPtrFn map_ptr, UnmappedPtrIndexFn unmapped_index,
     const OrderedTy& ordered_container, UnorderedTy& unordered_container) {
   std::vector<size_t> indices;
diff --git a/third_party/xla/xla/service/memory_space_propagation_test.cc b/third_party/xla/xla/service/memory_space_propagation_test.cc
index da2059817cd4d1..940a4ebbcc400e 100644
--- a/third_party/xla/xla/service/memory_space_propagation_test.cc
+++ b/third_party/xla/xla/service/memory_space_propagation_test.cc
@@ -29,7 +29,9 @@ class MemorySpacePropagationTest : public HloTestBase {
         verifier_(/*layout_sensitive=*/false, /*allow_mixed_precision*/ false) {
   }
 
-  Status Verify(HloModule* module) { return verifier_.Run(module).status(); }
+  absl::Status Verify(HloModule* module) {
+    return verifier_.Run(module).status();
+  }
 
  private:
   HloVerifier verifier_;
diff --git a/third_party/xla/xla/service/p2p_schedule_preparation.cc b/third_party/xla/xla/service/p2p_schedule_preparation.cc
index 7508e73810f2fb..a41aebc99e7fd9 100644
--- a/third_party/xla/xla/service/p2p_schedule_preparation.cc
+++ b/third_party/xla/xla/service/p2p_schedule_preparation.cc
@@ -236,7 +236,7 @@ static constexpr int kPipelinedParentNodeIdx = 1;
 // If a group forms a cycle with another group, records the other group as a
 // complement group.
 struct P2PGroup {
-  Status RecordP2POpForUnpipelinedGroup(HloSendRecvInstruction* p2p) {
+  absl::Status RecordP2POpForUnpipelinedGroup(HloSendRecvInstruction* p2p) {
     if (kind == kUnrecognized) {
       // Leave unrecognized P2P groups alone.
       return OkStatus();
@@ -251,7 +251,7 @@ struct P2PGroup {
     return OkStatus();
   }
 
-  Status RecordP2POpForPipelinedGroup(HloSendRecvInstruction* p2p) {
+  absl::Status RecordP2POpForPipelinedGroup(HloSendRecvInstruction* p2p) {
     if (kind == kUnrecognized) {
       // Leave unrecognized P2P groups alone.
       return OkStatus();
@@ -269,7 +269,7 @@ struct P2PGroup {
     return OkStatus();
   }
 
-  Status RecordWhileOpToPipelinedGroup(HloInstruction* while_op) {
+  absl::Status RecordWhileOpToPipelinedGroup(HloInstruction* while_op) {
     if (kind == kUnrecognized) {
       // Leave unrecognized P2P groups alone.
       return OkStatus();
@@ -304,7 +304,7 @@ struct P2PGroup {
 
   // Records the other group that forms a cycle with this group, assuming that
   // we handle only two groups that form a cycle.
-  Status RecordComplementGroup(P2PGroupMap& p2p_group_map) {
+  absl::Status RecordComplementGroup(P2PGroupMap& p2p_group_map) {
     CHECK(!complement_group_channel.has_value() && runtime_stream == kStream1);
     for (auto& [channel, p2p_group] : p2p_group_map) {
       if (&p2p_group == this ||
@@ -417,9 +417,9 @@ bool MayInvokeCollectiveOp(
 // Send operand of the while-op, we assume these two P2P groups belong to the
 // same pipelined P2P sequence. Adds the WhileOp to the pipelined group
 // representation in this case.
-Status MayAddWhileOpToPipelinedGroup(HloInstruction* while_op,
-                                     P2PInComputation& p2p_in_computation,
-                                     P2PGroupMap& p2p_group_map) {
+absl::Status MayAddWhileOpToPipelinedGroup(HloInstruction* while_op,
+                                           P2PInComputation& p2p_in_computation,
+                                           P2PGroupMap& p2p_group_map) {
   if (while_op->while_init()->opcode() != HloOpcode::kTuple) {
     // A while-init should contain the loop index variable. So if a while-init
     // is not a tuple, it only contains the loop index variable and shouldn't
@@ -455,7 +455,7 @@ Status MayAddWhileOpToPipelinedGroup(HloInstruction* while_op,
   return OkStatus();
 }
 
-Status OrderBefore(HloInstruction* i1, HloInstruction* i2) {
+absl::Status OrderBefore(HloInstruction* i1, HloInstruction* i2) {
   TF_RETURN_IF_ERROR(i1->AddControlDependencyTo(i2));
   VLOG(10) << "Add control predecessor " << i2->ToString();
   return OkStatus();
@@ -463,7 +463,7 @@ Status OrderBefore(HloInstruction* i1, HloInstruction* i2) {
 
 // Adds control dependence to enforce this ordering:
 //   recv => send => recv-done => send-done.
-Status ConnectP2P1NodeChain(const P2PGroupNode& node) {
+absl::Status ConnectP2P1NodeChain(const P2PGroupNode& node) {
   HloRecvDoneInstruction* recv_done = node.recv_done;
   HloRecvInstruction* recv = node.recv;
   HloSendDoneInstruction* send_done = node.send_done;
@@ -477,14 +477,14 @@ Status ConnectP2P1NodeChain(const P2PGroupNode& node) {
 // For an unpipelined Send-Recv chain, adds control dependence to enforce this
 // ordering:
 //   recv => send => recv-done => send-done.
-Status ConnectUnpipelinedP2P(const P2PGroup& p2p_group) {
+absl::Status ConnectUnpipelinedP2P(const P2PGroup& p2p_group) {
   return ConnectP2P1NodeChain(p2p_group.GetChild());
 }
 
 // For a single pipelined Send-Recv chain in a while-body, adds control
 // dependence toenforce this ordering:
 //   recv => send => recv-done => send-done
-Status ConnectPipelined1P2PChild(const P2PGroup& p2p_group) {
+absl::Status ConnectPipelined1P2PChild(const P2PGroup& p2p_group) {
   return ConnectP2P1NodeChain(p2p_group.GetChild());
 }
 
@@ -492,8 +492,8 @@ Status ConnectPipelined1P2PChild(const P2PGroup& p2p_group) {
 // enforce this ordering:
 //   recv.0 => send.0 => recv.1 => send.1 =>
 //   recv-done.0 => recv-done.1 => send-done.0 => send-done.1
-Status ConnectP2P2NodeChain(const P2PGroupNode& node0,
-                            const P2PGroupNode& node1) {
+absl::Status ConnectP2P2NodeChain(const P2PGroupNode& node0,
+                                  const P2PGroupNode& node1) {
   HloSendRecvInstruction* recv_done0 = node0.recv_done;
   HloRecvInstruction* recv0 = node0.recv;
   HloSendRecvInstruction* send_done0 = node0.send_done;
@@ -520,8 +520,8 @@ Status ConnectP2P2NodeChain(const P2PGroupNode& node0,
 // while-body computation, we enforce this ordering:
 //   recv.0 => send.0 => recv.1 => send.1 =>
 //   recv-done.0 => recv-done.1 => send-done.0 => send-done.1
-Status ConnectPipelined2P2PChild(const P2PGroup& p2p_group,
-                                 const P2PGroupMap& p2p_group_map) {
+absl::Status ConnectPipelined2P2PChild(const P2PGroup& p2p_group,
+                                       const P2PGroupMap& p2p_group_map) {
   return ConnectP2P2NodeChain(
       p2p_group.GetComplementGroup(p2p_group_map)->GetChild(),
       p2p_group.GetChild());
@@ -530,7 +530,7 @@ Status ConnectPipelined2P2PChild(const P2PGroup& p2p_group,
 // For a pipelined Send-Recv chain with one group in the while-body calling
 // computation, we enforce this ordering:
 //   recv => send => recv-done => send-done
-Status ConnectPipelined1P2PParent(const P2PGroup& p2p_group) {
+absl::Status ConnectPipelined1P2PParent(const P2PGroup& p2p_group) {
   return ConnectP2P1NodeChain(p2p_group.GetParent());
 }
 
@@ -538,8 +538,8 @@ Status ConnectPipelined1P2PParent(const P2PGroup& p2p_group) {
 // in the while-body calling computation, we enforce this ordering:
 //   recv.0 => send.0 => recv.1 => send.1 => =>
 //   recv-done.0 => recv-done.1 => send-done.0 => send-done.1
-Status ConnectPipelined2P2PParent(const P2PGroup& p2p_group,
-                                  const P2PGroupMap& p2p_group_map) {
+absl::Status ConnectPipelined2P2PParent(const P2PGroup& p2p_group,
+                                        const P2PGroupMap& p2p_group_map) {
   return ConnectP2P2NodeChain(
       p2p_group.GetComplementGroup(p2p_group_map)->GetParent(),
       p2p_group.GetParent());
@@ -550,8 +550,8 @@ Status ConnectPipelined2P2PParent(const P2PGroup& p2p_group,
 // enforece this ordering:
 //   recv.0 => send.0 => recv.1 => send.1 =>
 //   recv-done.0 => recv-done.1 => send-done.0 => send-done.1
-Status ConnectUnpipelined2P2P(const P2PGroup& p2p_group,
-                              const P2PGroupMap& p2p_group_map) {
+absl::Status ConnectUnpipelined2P2P(const P2PGroup& p2p_group,
+                                    const P2PGroupMap& p2p_group_map) {
   CHECK(p2p_group.runtime_stream == kStream1);
   return ConnectP2P2NodeChain(
       p2p_group.GetComplementGroup(p2p_group_map)->GetChild(),
@@ -562,7 +562,7 @@ Status ConnectUnpipelined2P2P(const P2PGroup& p2p_group,
 // groups them by channel IDs, records pipeline decision for groups and connects
 // groups that form a cycle for pipelining. Also records whether the computation
 // invokes collective operation directly or indirectly.
-Status GatherP2PGroupsAndCollectiveInfo(
+absl::Status GatherP2PGroupsAndCollectiveInfo(
     const HloComputation* computation, P2PInComputation& p2p_in_computation,
     P2PGroupMap& p2p_group_map,
     CollectiveInComputation& collective_in_computation) {
@@ -728,8 +728,8 @@ absl::StatusOr<std::pair<int, const P2PGroup*>> ConnectP2PChain(
   return std::make_pair(num_p2p_chains, pipelined_group);
 }
 
-Status OrderBefore(HloReachabilityMap* reachability, HloInstruction* a,
-                   HloInstruction* b) {
+absl::Status OrderBefore(HloReachabilityMap* reachability, HloInstruction* a,
+                         HloInstruction* b) {
   VLOG(10) << "OrderBefore " << a->ToString() << " " << b->ToString();
   if (!reachability->IsReachable(a, b)) {
     TF_RETURN_IF_ERROR(a->AddControlDependencyTo(b));
@@ -748,7 +748,7 @@ Status OrderBefore(HloReachabilityMap* reachability, HloInstruction* a,
 //      Recv-Done.1 => Send-Done.1 (pipelined 2)
 // We intend to schedule collective ops ordered before the beginning of such a
 // chain or after the ending of such a chain.
-Status LinearizeCollectivesWithOtherP2P(
+absl::Status LinearizeCollectivesWithOtherP2P(
     const P2PGroupMap& p2p_group_map, const P2PGroup& group,
     const CollectiveInComputation& collective_in_computation,
     const std::vector<HloInstruction*>::iterator& chain_start_iter,
@@ -835,7 +835,7 @@ Status LinearizeCollectivesWithOtherP2P(
 // Adds control dependence to linearize other collective ops with respect to
 // the given pipelined P2P chain in the computation for the pipelined
 // while-loop. All Collective ops should be scheduled before the chain.
-Status LinearizeCollectivesWithPipelinedP2PChild(
+absl::Status LinearizeCollectivesWithPipelinedP2PChild(
     const P2PGroupMap& p2p_group_map, const P2PGroup& group,
     const CollectiveInComputation& collective_in_computation,
     HloComputation* computation, HloReachabilityMap* reachability) {
diff --git a/third_party/xla/xla/service/qr_expander.cc b/third_party/xla/xla/service/qr_expander.cc
index e817b66b61d2c8..8eb3abfa353200 100644
--- a/third_party/xla/xla/service/qr_expander.cc
+++ b/third_party/xla/xla/service/qr_expander.cc
@@ -106,8 +106,8 @@ XlaOp Norm(std::vector<XlaOp> xs) {
 //   return (v, tau, beta)
 // TODO(phawkins): LAPACK's xLARFG implementation has code for handling
 // overflows in the norm/beta calculations. Perhaps do the same here.
-Status House(XlaOp x, XlaOp k, absl::Span<const int64_t> batch_dims,
-             const int64_t m, XlaOp* v, XlaOp* tau, XlaOp* beta) {
+absl::Status House(XlaOp x, XlaOp k, absl::Span<const int64_t> batch_dims,
+                   const int64_t m, XlaOp* v, XlaOp* tau, XlaOp* beta) {
   XlaBuilder* const builder = x.builder();
   TF_ASSIGN_OR_RETURN(Shape x_shape, builder->GetShape(x));
   const PrimitiveType type = x_shape.element_type();
diff --git a/third_party/xla/xla/service/reduce_decomposer.cc b/third_party/xla/xla/service/reduce_decomposer.cc
index 0e43381063f1b3..507be9637120a7 100644
--- a/third_party/xla/xla/service/reduce_decomposer.cc
+++ b/third_party/xla/xla/service/reduce_decomposer.cc
@@ -33,7 +33,7 @@ namespace {
 // Enforces property that all inputs to variadic reduction have same layout.
 class VariadicReductionLayoutEqualizer : public DfsHloRewriteVisitor {
  public:
-  Status HandleReduce(HloInstruction* hlo) override {
+  absl::Status HandleReduce(HloInstruction* hlo) override {
     auto reduce = Cast<HloReduceInstruction>(hlo);
     std::vector<HloInstruction*> new_inputs;
     bool changed = false;
@@ -70,7 +70,7 @@ class ReduceDecomposerVisitor : public DfsHloRewriteVisitor {
   explicit ReduceDecomposerVisitor(HloPredicate custom_layout_allowed)
       : custom_layout_allowed_(std::move(custom_layout_allowed)) {}
 
-  Status HandleReduce(HloInstruction* hlo) override {
+  absl::Status HandleReduce(HloInstruction* hlo) override {
     auto reduce = Cast<HloReduceInstruction>(hlo);
     auto shape = reduce->shape();
     if (custom_layout_allowed_ && custom_layout_allowed_(reduce)) {
diff --git a/third_party/xla/xla/service/reduce_scatter_combiner.cc b/third_party/xla/xla/service/reduce_scatter_combiner.cc
index 40563bec226022..468371fb4751d1 100644
--- a/third_party/xla/xla/service/reduce_scatter_combiner.cc
+++ b/third_party/xla/xla/service/reduce_scatter_combiner.cc
@@ -77,7 +77,8 @@ using ReduceScatterKey =
 // Combines the elements of to_combine into a single ReduceScatter op. All
 // entries in to_combine must be ReduceScatter ops with exactly one operand
 // and the same reduction operation.
-Status CombineReduceScatters(absl::Span<HloInstruction* const> to_combine) {
+absl::Status CombineReduceScatters(
+    absl::Span<HloInstruction* const> to_combine) {
   if (to_combine.size() < 2) {
     return OkStatus();
   }
diff --git a/third_party/xla/xla/service/reduce_window_rewriter.cc b/third_party/xla/xla/service/reduce_window_rewriter.cc
index 39b201e108cc29..934167beeed39d 100644
--- a/third_party/xla/xla/service/reduce_window_rewriter.cc
+++ b/third_party/xla/xla/service/reduce_window_rewriter.cc
@@ -76,7 +76,7 @@ static HloInstruction* GetAtIndex(HloInstruction* hlo,
 //   if rank(x) == 1:
 //   then: reshape_r2_r1(reduce-win(reshape_r1_r2(x)))
 //   else: no change
-Status ReduceWindowRewriter::ReplaceReduceWindowWithReshape(
+absl::Status ReduceWindowRewriter::ReplaceReduceWindowWithReshape(
     HloReduceWindowInstruction* reduce_window) {
   VLOG(2) << "Converting R1 reduce window: " << reduce_window->ToString();
 
@@ -434,7 +434,8 @@ absl::StatusOr<bool> ReduceWindowRewriter::TryOptimizeCumSumOrProd(
   std::vector<HloInstruction*> scans;
   auto status = ShapeUtil::ForEachSubshapeWithStatus(
       outer_reduce_window->shape(),
-      [&](const Shape& subshape, const ShapeIndex& shape_index) -> Status {
+      [&](const Shape& subshape,
+          const ShapeIndex& shape_index) -> absl::Status {
         if (!ShapeUtil::IsLeafIndex(outer_reduce_window->shape(),
                                     shape_index)) {
           return OkStatus();
diff --git a/third_party/xla/xla/service/reduce_window_rewriter.h b/third_party/xla/xla/service/reduce_window_rewriter.h
index 6fed774b27aadb..591056a3f094ba 100644
--- a/third_party/xla/xla/service/reduce_window_rewriter.h
+++ b/third_party/xla/xla/service/reduce_window_rewriter.h
@@ -58,10 +58,10 @@ class ReduceWindowRewriter : public HloModulePass {
       const absl::flat_hash_set<absl::string_view>& execution_threads) override;
 
  private:
-  Status ReplaceReduceWindowWithReshape(
+  absl::Status ReplaceReduceWindowWithReshape(
       HloReduceWindowInstruction* reduce_window);
 
-  StatusOr<bool> TryOptimizeCumSumOrProd(
+  absl::StatusOr<bool> TryOptimizeCumSumOrProd(
       HloReduceWindowInstruction* reduce_window);
 
   int64_t base_length_;
diff --git a/third_party/xla/xla/service/reshape_decomposer.cc b/third_party/xla/xla/service/reshape_decomposer.cc
index 2e3ab87569abc4..64fa392adf0851 100644
--- a/third_party/xla/xla/service/reshape_decomposer.cc
+++ b/third_party/xla/xla/service/reshape_decomposer.cc
@@ -25,7 +25,7 @@ namespace {
 
 class ReshapeDecomposerVisitor : public DfsHloRewriteVisitor {
  public:
-  Status HandleReshape(HloInstruction* reshape) override {
+  absl::Status HandleReshape(HloInstruction* reshape) override {
     HloInstruction* operand = reshape->mutable_operand(0);
     auto s = reshape->shape();
     auto s0 = operand->shape();
diff --git a/third_party/xla/xla/service/reshape_mover_test.cc b/third_party/xla/xla/service/reshape_mover_test.cc
index 02867a54a3b90f..d853d42bac89af 100644
--- a/third_party/xla/xla/service/reshape_mover_test.cc
+++ b/third_party/xla/xla/service/reshape_mover_test.cc
@@ -37,8 +37,8 @@ namespace m = xla::match;
 class ReshapeMoverTest : public HloTestBase {
  protected:
   // ReshapeMover relies on algsimp for cleanup.
-  Status RunPass(HloModule* module, bool change_expected,
-                 ReshapeMoverOptions options = ReshapeMoverOptions{}) {
+  absl::Status RunPass(HloModule* module, bool change_expected,
+                       ReshapeMoverOptions options = ReshapeMoverOptions{}) {
     TF_ASSIGN_OR_RETURN(bool changed,
                         RunHloPass(ReshapeMover(options), module));
     SCOPED_TRACE(module->ToString());
diff --git a/third_party/xla/xla/service/service.cc b/third_party/xla/xla/service/service.cc
index b22c081390eafe..42eeb8c8b85e4f 100644
--- a/third_party/xla/xla/service/service.cc
+++ b/third_party/xla/xla/service/service.cc
@@ -71,9 +71,9 @@ using absl::StrCat;
 using absl::StrFormat;
 
 // Records the arguments used to invoke a computation in an HloSnapshot proto.
-Status RecordArguments(const absl::Span<const ShapedBuffer* const> arguments,
-                       se::Stream* stream, TransferManager* transfer_manager,
-                       HloSnapshot* module) {
+absl::Status RecordArguments(
+    const absl::Span<const ShapedBuffer* const> arguments, se::Stream* stream,
+    TransferManager* transfer_manager, HloSnapshot* module) {
   module->clear_arguments();
   for (const ShapedBuffer* argument : arguments) {
     TF_ASSIGN_OR_RETURN(
@@ -85,8 +85,9 @@ Status RecordArguments(const absl::Span<const ShapedBuffer* const> arguments,
 }
 
 // Records the result of a computation in a HloSnapshot proto.
-Status RecordResult(const ShapedBuffer& result, se::Stream* stream,
-                    TransferManager* transfer_manager, HloSnapshot* module) {
+absl::Status RecordResult(const ShapedBuffer& result, se::Stream* stream,
+                          TransferManager* transfer_manager,
+                          HloSnapshot* module) {
   module->clear_result();
   TF_ASSIGN_OR_RETURN(
       Literal literal,
@@ -159,18 +160,18 @@ Service::Service(const ServiceOptions& options,
   }
 }
 
-Status Service::CreateChannelHandle(const CreateChannelHandleRequest* arg,
-                                    CreateChannelHandleResponse* result) {
+absl::Status Service::CreateChannelHandle(const CreateChannelHandleRequest* arg,
+                                          CreateChannelHandleResponse* result) {
   TF_ASSIGN_OR_RETURN(*result->mutable_channel(),
                       channel_tracker_.NewChannel(arg->channel_type()));
   return OkStatus();
 }
 
-Status Service::Unregister(const UnregisterRequest* arg,
-                           UnregisterResponse* result) {
-  Status status;
+absl::Status Service::Unregister(const UnregisterRequest* arg,
+                                 UnregisterResponse* result) {
+  absl::Status status;
   for (auto& data : arg->data()) {
-    Status unregister_status = allocation_tracker_.Unregister(data);
+    absl::Status unregister_status = allocation_tracker_.Unregister(data);
     if (!unregister_status.ok() && status.ok()) {
       status = unregister_status;
     }
@@ -179,8 +180,8 @@ Status Service::Unregister(const UnregisterRequest* arg,
 }
 
 // Deconstructs a previously-allocated global handle.
-Status Service::DeconstructTuple(const DeconstructTupleRequest* arg,
-                                 DeconstructTupleResponse* result) {
+absl::Status Service::DeconstructTuple(const DeconstructTupleRequest* arg,
+                                       DeconstructTupleResponse* result) {
   TF_ASSIGN_OR_RETURN(
       std::vector<GlobalDataHandle> elements,
       allocation_tracker_.DeconstructTuple(arg->tuple_handle()));
@@ -191,8 +192,8 @@ Status Service::DeconstructTuple(const DeconstructTupleRequest* arg,
   return OkStatus();
 }
 
-Status Service::ValidateResultShape(const Shape& client_shape,
-                                    const Shape& result_shape) {
+absl::Status Service::ValidateResultShape(const Shape& client_shape,
+                                          const Shape& result_shape) {
   TF_RETURN_IF_ERROR(ShapeUtil::ValidateShapeWithOptionalLayout(client_shape));
   if (!ShapeUtil::Compatible(client_shape, result_shape)) {
     return InvalidArgument(
@@ -421,7 +422,7 @@ Service::ExecuteParallelAndRegisterResult(
 
   // Wait for all executions to complete.
   for (int64_t i = 0, end = streams.size(); i < end; ++i) {
-    Status block_status = streams[i]->BlockHostUntilDone();
+    absl::Status block_status = streams[i]->BlockHostUntilDone();
     if (!block_status.ok()) {
       return Internal("failed to complete execution for stream %d: %s", i,
                       block_status.message());
@@ -530,8 +531,8 @@ Service::GetArguments(
   return replicated_arguments;
 }
 
-Status Service::ExecuteGraphParallel(const ExecuteGraphParallelRequest* arg,
-                                     ExecuteParallelResponse* result) {
+absl::Status Service::ExecuteGraphParallel(
+    const ExecuteGraphParallelRequest* arg, ExecuteParallelResponse* result) {
   VLOG(1) << "running execute-graph-parallel request";
 
   std::vector<std::vector<std::vector<const ShapedBuffer*>>> all_arguments;
@@ -655,7 +656,7 @@ Status Service::ExecuteGraphParallel(const ExecuteGraphParallelRequest* arg,
   // basically the same thing.
   ExecutionProfile profile;
   std::vector<GlobalDataHandle> outputs;
-  Status execution_status = OkStatus();
+  absl::Status execution_status = OkStatus();
 
   if (executable_ptrs.size() == 1) {
     absl::StatusOr<GlobalDataHandle> output_or_status =
@@ -714,8 +715,8 @@ Status Service::ExecuteGraphParallel(const ExecuteGraphParallelRequest* arg,
   return OkStatus();
 }
 
-Status Service::GetDeviceHandles(const GetDeviceHandlesRequest* arg,
-                                 GetDeviceHandlesResponse* result) {
+absl::Status Service::GetDeviceHandles(const GetDeviceHandlesRequest* arg,
+                                       GetDeviceHandlesResponse* result) {
   const int64_t available_device_count = execute_backend_->device_count();
   const int64_t replica_count = options_.number_of_replicas();
   if (replica_count <= 0) {
@@ -794,7 +795,8 @@ absl::StatusOr<std::unique_ptr<Executable>> Service::BuildExecutable(
   return executable;
 }
 
-Status Service::Compile(const CompileRequest* arg, CompileResponse* result) {
+absl::Status Service::Compile(const CompileRequest* arg,
+                              CompileResponse* result) {
   VLOG(1) << "running compile request";
   if (!arg->has_computation()) {
     return InvalidArgument("computations may not be empty");
@@ -835,7 +837,8 @@ Status Service::Compile(const CompileRequest* arg, CompileResponse* result) {
   return OkStatus();
 }
 
-Status Service::Execute(const ExecuteRequest* arg, ExecuteResponse* result) {
+absl::Status Service::Execute(const ExecuteRequest* arg,
+                              ExecuteResponse* result) {
   VLOG(1) << "running execute request";
   if (!arg->has_handle()) {
     return InvalidArgument("execution handle should not be empty");
@@ -905,8 +908,8 @@ Status Service::Execute(const ExecuteRequest* arg, ExecuteResponse* result) {
   return OkStatus();
 }
 
-Status Service::TransferToClient(const TransferToClientRequest* arg,
-                                 TransferToClientResponse* result) {
+absl::Status Service::TransferToClient(const TransferToClientRequest* arg,
+                                       TransferToClientResponse* result) {
   TF_ASSIGN_OR_RETURN(const ShapedBuffer* shaped_buffer,
                       allocation_tracker_.ResolveForReplica(arg->data(), 0));
 
@@ -948,8 +951,8 @@ Status Service::TransferToClient(const TransferToClientRequest* arg,
   return OkStatus();
 }
 
-Status Service::TransferToServer(const TransferToServerRequest* arg,
-                                 TransferToServerResponse* result) {
+absl::Status Service::TransferToServer(const TransferToServerRequest* arg,
+                                       TransferToServerResponse* result) {
   TF_ASSIGN_OR_RETURN(Literal literal,
                       Literal::CreateFromProto(arg->literal()));
   const Shape& shape = literal.shape();
@@ -991,8 +994,8 @@ Status Service::TransferToServer(const TransferToServerRequest* arg,
   return OkStatus();
 }
 
-Status Service::TransferToInfeed(const TransferToInfeedRequest* arg,
-                                 TransferToInfeedResponse* result) {
+absl::Status Service::TransferToInfeed(const TransferToInfeedRequest* arg,
+                                       TransferToInfeedResponse* result) {
   const int64_t replica_count = options_.number_of_replicas();
   if (arg->replica_id() < 0 || arg->replica_id() >= replica_count) {
     return FailedPrecondition(
@@ -1020,8 +1023,8 @@ Status Service::TransferToInfeed(const TransferToInfeedRequest* arg,
                                                                        literal);
 }
 
-Status Service::TransferFromOutfeed(const TransferFromOutfeedRequest* arg,
-                                    TransferFromOutfeedResponse* result) {
+absl::Status Service::TransferFromOutfeed(const TransferFromOutfeedRequest* arg,
+                                          TransferFromOutfeedResponse* result) {
   const int64_t replica_count = options_.number_of_replicas();
   if (arg->replica_id() < 0 || arg->replica_id() >= replica_count) {
     return FailedPrecondition(
@@ -1050,13 +1053,13 @@ Status Service::TransferFromOutfeed(const TransferFromOutfeedRequest* arg,
   return OkStatus();
 }
 
-Status Service::ResetDevice(const ResetDeviceRequest* arg,
-                            ResetDeviceResponse* result) {
+absl::Status Service::ResetDevice(const ResetDeviceRequest* arg,
+                                  ResetDeviceResponse* result) {
   return execute_backend_->ResetDevices();
 }
 
-Status Service::ComputeConstantGraph(const ComputeConstantGraphRequest* arg,
-                                     ComputeConstantResponse* result) {
+absl::Status Service::ComputeConstantGraph(
+    const ComputeConstantGraphRequest* arg, ComputeConstantResponse* result) {
   if (!arg->has_computation()) {
     return InvalidArgument("computations may not be empty");
   }
@@ -1117,14 +1120,15 @@ Status Service::ComputeConstantGraph(const ComputeConstantGraphRequest* arg,
   return OkStatus();
 }
 
-Status Service::GetShape(const GetShapeRequest* arg, GetShapeResponse* result) {
+absl::Status Service::GetShape(const GetShapeRequest* arg,
+                               GetShapeResponse* result) {
   TF_ASSIGN_OR_RETURN(const ShapedBuffer* buffer,
                       allocation_tracker_.ResolveForReplica(arg->data(), 0));
   *result->mutable_shape() = buffer->on_device_shape().ToProto();
   return OkStatus();
 }
 
-Status Service::GetComputationGraphStats(
+absl::Status Service::GetComputationGraphStats(
     const ComputationGraphStatsRequest* arg, ComputationStatsResponse* result) {
   if (!arg->has_computation()) {
     return InvalidArgument("Computations may not be empty.");
diff --git a/third_party/xla/xla/service/service.h b/third_party/xla/xla/service/service.h
index f606698b1a8582..897b12cd83e91e 100644
--- a/third_party/xla/xla/service/service.h
+++ b/third_party/xla/xla/service/service.h
@@ -84,28 +84,30 @@ class Service : public ServiceInterface {
   //
   // If the handle given is not currently allocated, a NOT_FOUND status is
   // returned.
-  Status Unregister(const UnregisterRequest* arg,
-                    UnregisterResponse* result) override;
+  absl::Status Unregister(const UnregisterRequest* arg,
+                          UnregisterResponse* result) override;
 
   // Deconstructs a tuple. Returns a newly created GlobalDataHandle for each
   // element in the tuple.
-  Status DeconstructTuple(const DeconstructTupleRequest* arg,
-                          DeconstructTupleResponse* result) override;
+  absl::Status DeconstructTuple(const DeconstructTupleRequest* arg,
+                                DeconstructTupleResponse* result) override;
 
   // Compiles a computation into an executable. The request contains the whole
   // computation graph. Returns the handle to the executable.
-  Status Compile(const CompileRequest* arg, CompileResponse* result) override;
+  absl::Status Compile(const CompileRequest* arg,
+                       CompileResponse* result) override;
 
   // Executes an executable with the provided global data passes as immutable
   // arguments. The request contains the handle to the executable. Returns
   // global data output and execution timing.
-  Status Execute(const ExecuteRequest* arg, ExecuteResponse* result) override;
+  absl::Status Execute(const ExecuteRequest* arg,
+                       ExecuteResponse* result) override;
 
   // Executes one or more computations in parallel with the provided global data
   // passed as immutable arguments. Returns global data output for each
   // computation.
-  Status ExecuteGraphParallel(const ExecuteGraphParallelRequest* arg,
-                              ExecuteParallelResponse* result) override;
+  absl::Status ExecuteGraphParallel(const ExecuteGraphParallelRequest* arg,
+                                    ExecuteParallelResponse* result) override;
 
   // Requests one or more device handles from the target.
   //
@@ -115,26 +117,27 @@ class Service : public ServiceInterface {
   // the first set of replicas, and the next R devices to the second set of
   // replicas, etc. Each returned device handle represents the device with the
   // replica id 0.
-  Status GetDeviceHandles(const GetDeviceHandlesRequest* arg,
-                          GetDeviceHandlesResponse* result) override;
+  absl::Status GetDeviceHandles(const GetDeviceHandlesRequest* arg,
+                                GetDeviceHandlesResponse* result) override;
 
   // Requests that global data be transferred to the client in literal form.
-  Status TransferToClient(const TransferToClientRequest* arg,
-                          TransferToClientResponse* result) override;
+  absl::Status TransferToClient(const TransferToClientRequest* arg,
+                                TransferToClientResponse* result) override;
 
   // Transfers data from a literal provided by the client, into device memory.
-  Status TransferToServer(const TransferToServerRequest* arg,
-                          TransferToServerResponse* result) override;
+  absl::Status TransferToServer(const TransferToServerRequest* arg,
+                                TransferToServerResponse* result) override;
 
   // Transfers data from a literal provided by the client, into the Infeed
   // buffer of the device.
-  Status TransferToInfeed(const TransferToInfeedRequest* arg,
-                          TransferToInfeedResponse* result) override;
+  absl::Status TransferToInfeed(const TransferToInfeedRequest* arg,
+                                TransferToInfeedResponse* result) override;
 
   // Transfers data from the Outfeed othe device to the literal provided by the
   // client.
-  Status TransferFromOutfeed(const TransferFromOutfeedRequest* arg,
-                             TransferFromOutfeedResponse* result) override;
+  absl::Status TransferFromOutfeed(
+      const TransferFromOutfeedRequest* arg,
+      TransferFromOutfeedResponse* result) override;
 
   // Resets devices, clearing all existing state on all the devices associated
   // with this service (including memory allocated on the devices).
@@ -145,25 +148,27 @@ class Service : public ServiceInterface {
   // ResetDevice should be called before an Execution that expect the device to
   // be in the reset state. For example, if the prior Execution modifies device
   // state (e.g., architectural state) that the next Execution depends on.
-  Status ResetDevice(const ResetDeviceRequest* arg,
-                     ResetDeviceResponse* result) override;
+  absl::Status ResetDevice(const ResetDeviceRequest* arg,
+                           ResetDeviceResponse* result) override;
 
-  Status ComputeConstantGraph(const ComputeConstantGraphRequest* arg,
-                              ComputeConstantResponse* result) override;
+  absl::Status ComputeConstantGraph(const ComputeConstantGraphRequest* arg,
+                                    ComputeConstantResponse* result) override;
 
   // Returns the shape (with layout) of an array associated with a given data
   // handle.
-  Status GetShape(const GetShapeRequest* arg,
-                  GetShapeResponse* result) override;
+  absl::Status GetShape(const GetShapeRequest* arg,
+                        GetShapeResponse* result) override;
 
   // Retrieves the statistics of a computation.
-  Status GetComputationGraphStats(const ComputationGraphStatsRequest* arg,
-                                  ComputationStatsResponse* result) override;
+  absl::Status GetComputationGraphStats(
+      const ComputationGraphStatsRequest* arg,
+      ComputationStatsResponse* result) override;
 
   // Creates a unique channel handle that can be used for Send/Recv
   // instructions.
-  Status CreateChannelHandle(const CreateChannelHandleRequest* arg,
-                             CreateChannelHandleResponse* result) override;
+  absl::Status CreateChannelHandle(
+      const CreateChannelHandleRequest* arg,
+      CreateChannelHandleResponse* result) override;
 
   // Returns the backend used to execute computations.
   const Backend& backend() const { return *execute_backend_; }
@@ -180,8 +185,8 @@ class Service : public ServiceInterface {
   // Convenience function which checks whether the given client_shape
   // (presumably passed by the client to set the result layout) is valid for the
   // given computation result shape.
-  static Status ValidateResultShape(const Shape& client_shape,
-                                    const Shape& result_shape);
+  static absl::Status ValidateResultShape(const Shape& client_shape,
+                                          const Shape& result_shape);
 
  private:
   // A private overload for Service itself, used by other methods within this
diff --git a/third_party/xla/xla/service/service_executable_run_options.h b/third_party/xla/xla/service/service_executable_run_options.h
index a2f57e6bbae635..49631dae39cdfd 100644
--- a/third_party/xla/xla/service/service_executable_run_options.h
+++ b/third_party/xla/xla/service/service_executable_run_options.h
@@ -64,7 +64,8 @@ class ServiceExecutableRunOptions {
       int device_ordinal,
       se::StreamPriority priority = se::StreamPriority::Default) const {
     if (!stream_borrower_) {
-      return Status(absl::StatusCode::kUnimplemented, "No stream borrower");
+      return absl::Status(absl::StatusCode::kUnimplemented,
+                          "No stream borrower");
     }
 
     TF_ASSIGN_OR_RETURN(
@@ -79,7 +80,8 @@ class ServiceExecutableRunOptions {
       se::StreamPriority priority = se::StreamPriority::Default) const {
     return stream_borrower_
                ? stream_borrower_(device_ordinal, num_streams, priority)
-               : Status(absl::StatusCode::kUnimplemented, "No stream borrower");
+               : absl::Status(absl::StatusCode::kUnimplemented,
+                              "No stream borrower");
   }
 
  private:
diff --git a/third_party/xla/xla/service/shape_inference.cc b/third_party/xla/xla/service/shape_inference.cc
index 10d2f5d3f75b44..9268479cba2005 100644
--- a/third_party/xla/xla/service/shape_inference.cc
+++ b/third_party/xla/xla/service/shape_inference.cc
@@ -78,7 +78,7 @@ bool CompatibleDimensionSizes(int64_t size_a, int64_t size_b) {
          size_a == size_b;
 }
 
-Status ExpectArray(const Shape& shape, absl::string_view op_type) {
+absl::Status ExpectArray(const Shape& shape, absl::string_view op_type) {
   if (!shape.IsArray()) {
     return InvalidArgument("Expected array argument for %s, but got %s.",
                            std::string(op_type), ShapeUtil::HumanString(shape));
@@ -86,10 +86,10 @@ Status ExpectArray(const Shape& shape, absl::string_view op_type) {
   return OkStatus();
 }
 
-Status VerifyReducerShape(const ProgramShape& reducer_shape,
-                          absl::Span<const Shape* const> init_value_shapes,
-                          absl::Span<const PrimitiveType> input_element_types,
-                          int64_t inputs) {
+absl::Status VerifyReducerShape(
+    const ProgramShape& reducer_shape,
+    absl::Span<const Shape* const> init_value_shapes,
+    absl::Span<const PrimitiveType> input_element_types, int64_t inputs) {
   if (reducer_shape.parameters_size() != inputs * 2) {
     return InvalidArgument(
         "Reduction function must take %d parameters, but "
@@ -791,7 +791,7 @@ absl::StatusOr<DimAndBound> InferMostSpecificDimAndBound(int64_t dim,
 
 namespace {
 
-Status ValidateDotDimensionNumbers(
+absl::Status ValidateDotDimensionNumbers(
     const Shape& lhs, const Shape& rhs,
     const DotDimensionNumbers& dimension_numbers) {
   // Check that dimension numbers are in range.
@@ -850,7 +850,7 @@ Status ValidateDotDimensionNumbers(
   TF_RETURN_IF_ERROR(ExpectArray(lhs, "lhs of dot"));
   TF_RETURN_IF_ERROR(ExpectArray(rhs, "rhs of dot"));
 
-  auto fail = [lhs, rhs](const std::string& addendum) -> Status {
+  auto fail = [lhs, rhs](const std::string& addendum) -> absl::Status {
     std::string message =
         StrFormat("Cannot infer shape for dot operation: %s <dot> %s.",
                   ShapeUtil::HumanString(lhs), ShapeUtil::HumanString(rhs));
@@ -3784,7 +3784,7 @@ ShapeInference::InferCollectivePermuteDoneShape(const Shape& operand_shape) {
   return to_apply.result();
 }
 
-static Status ValidateGatherDimensionNumbers(
+static absl::Status ValidateGatherDimensionNumbers(
     const Shape& input_shape, absl::Span<const int64_t> start_indices_shape,
     const GatherDimensionNumbers& dim_numbers) {
   if (!absl::c_is_sorted(dim_numbers.offset_dims())) {
@@ -4020,7 +4020,7 @@ static Status ValidateGatherDimensionNumbers(
 
 namespace {
 
-Status ValidateScatterDimensionNumbers(
+absl::Status ValidateScatterDimensionNumbers(
     const Shape& operand_shape, absl::Span<const int64_t> scatter_indices_shape,
     const Shape& updates_shape, const ScatterDimensionNumbers& dim_numbers) {
   // Validate update_window_dims in ScatterDimensionNumbers.
diff --git a/third_party/xla/xla/service/shaped_buffer_test.cc b/third_party/xla/xla/service/shaped_buffer_test.cc
index 28869830f984e9..f204f4e19e83cb 100644
--- a/third_party/xla/xla/service/shaped_buffer_test.cc
+++ b/third_party/xla/xla/service/shaped_buffer_test.cc
@@ -72,7 +72,8 @@ class TestAllocator : public se::DeviceMemoryAllocator {
                                   device_ordinal, this);
   }
 
-  Status Deallocate(int device_ordinal, se::DeviceMemoryBase mem) override {
+  absl::Status Deallocate(int device_ordinal,
+                          se::DeviceMemoryBase mem) override {
     if (mem.is_null()) {
       return OkStatus();
     }
diff --git a/third_party/xla/xla/service/sharding_propagation.cc b/third_party/xla/xla/service/sharding_propagation.cc
index dba8a717342596..183910e4e856f7 100644
--- a/third_party/xla/xla/service/sharding_propagation.cc
+++ b/third_party/xla/xla/service/sharding_propagation.cc
@@ -891,7 +891,7 @@ bool RemoveShardingMetadata(
 // instructions with a device assignment are on D. Further, annotate the root
 // instruction of the while body to ensure that HLO partitioning will keep the
 // entire while instruction on D.
-Status CheckAndUpdateDeviceAssignmentsInWhileBody(
+absl::Status CheckAndUpdateDeviceAssignmentsInWhileBody(
     HloInstruction* while_instruction) {
   auto bad_status = [](HloInstruction* instruction, int64_t device,
                        HloInstruction* channel_instruction,
@@ -1770,7 +1770,7 @@ int64_t ComputeNonRootUsers(const HloInstruction* instr) {
   return non_root_users;
 }
 
-/*static*/ Status ShardingPropagation::NormalizeDomain(
+/*static*/ absl::Status ShardingPropagation::NormalizeDomain(
     const DomainMetadata::Domain& domain, const DomainMetadata* metadata) {
   if (metadata != nullptr) {
     TF_ASSIGN_OR_RETURN(const auto& sharding_metadata,
diff --git a/third_party/xla/xla/service/sharding_propagation.h b/third_party/xla/xla/service/sharding_propagation.h
index 1d2fa03a932ad8..28505bbd58940f 100644
--- a/third_party/xla/xla/service/sharding_propagation.h
+++ b/third_party/xla/xla/service/sharding_propagation.h
@@ -130,8 +130,8 @@ class ShardingPropagation : public HloModulePass {
   // given domain. It will apply the sharding into the exit edges of the domain
   // and then rely on the rest of sharding propagation to ensure that the
   // intermediate nodes get the correct sharding.
-  static Status NormalizeDomain(const DomainMetadata::Domain& domain,
-                                const DomainMetadata* metadata);
+  static absl::Status NormalizeDomain(const DomainMetadata::Domain& domain,
+                                      const DomainMetadata* metadata);
 
   static std::optional<HloSharding> GetShardingFromUser(
       const HloInstruction& instruction, const HloInstruction& user,
diff --git a/third_party/xla/xla/service/slice_sinker.cc b/third_party/xla/xla/service/slice_sinker.cc
index dc7559444436c6..baf18e087ef774 100644
--- a/third_party/xla/xla/service/slice_sinker.cc
+++ b/third_party/xla/xla/service/slice_sinker.cc
@@ -164,8 +164,9 @@ std::optional<std::vector<HloInstruction*>> FindElementwiseOperationGroup(
 // Generates a new elementwise operation using the slice_sources as operands,
 // and replaces the uses of elementwise operation_on_slices with slices of the
 // new elementwise operations.
-Status SinkSlices(const std::vector<HloInstruction*>& slice_sources,
-                  const std::vector<HloInstruction*>& operation_on_slices) {
+absl::Status SinkSlices(
+    const std::vector<HloInstruction*>& slice_sources,
+    const std::vector<HloInstruction*>& operation_on_slices) {
   const Shape shape = slice_sources[0]->shape();
   PrimitiveType element_type = operation_on_slices[0]->shape().element_type();
   Shape new_shape = ShapeUtil::ChangeElementType(shape, element_type);
diff --git a/third_party/xla/xla/service/source_map_util.h b/third_party/xla/xla/service/source_map_util.h
index 5fd9db9dca751b..b5efa41186dded 100644
--- a/third_party/xla/xla/service/source_map_util.h
+++ b/third_party/xla/xla/service/source_map_util.h
@@ -25,9 +25,9 @@ namespace source_map_util {
 
 // Creates an INVALID_ARGUMENT status with the given format string.
 template <typename... Args>
-Status InvalidParameterArgument(const OpMetadata& op_metadata,
-                                const absl::FormatSpec<Args...>& format,
-                                const Args&... args) {
+absl::Status InvalidParameterArgument(const OpMetadata& op_metadata,
+                                      const absl::FormatSpec<Args...>& format,
+                                      const Args&... args) {
   std::string message = absl::StrFormat(format, args...);
   if (!op_metadata.source_file().empty()) {
     absl::StrAppendFormat(&message, " (%s:%d)", op_metadata.source_file(),
@@ -44,9 +44,10 @@ Status InvalidParameterArgument(const OpMetadata& op_metadata,
 // executable may be nullptr, but parameter_number should not be out of bounds
 // or a CHECK-failure may occur.
 template <typename... Args>
-Status InvalidParameterArgument(Executable* executable, int parameter_number,
-                                const absl::FormatSpec<Args...>& format,
-                                const Args&... args) {
+absl::Status InvalidParameterArgument(Executable* executable,
+                                      int parameter_number,
+                                      const absl::FormatSpec<Args...>& format,
+                                      const Args&... args) {
   if (executable != nullptr && executable->has_module()) {
     const HloModule& module = executable->module();
     const HloComputation& computation = *module.entry_computation();
diff --git a/third_party/xla/xla/service/space_to_batch_converter.cc b/third_party/xla/xla/service/space_to_batch_converter.cc
index a6149b197805ec..adaf2eb14daaae 100644
--- a/third_party/xla/xla/service/space_to_batch_converter.cc
+++ b/third_party/xla/xla/service/space_to_batch_converter.cc
@@ -65,7 +65,7 @@ namespace m = match;
 class ConvolutionVisitor {
  public:
   // Top-level function to begin space-to-batch conversion.
-  Status PerformSpaceToBatchOnConvolution(HloInstruction* convolution);
+  absl::Status PerformSpaceToBatchOnConvolution(HloInstruction* convolution);
 
   // Struct containing details about a convolution.
   struct ConvDetails {
@@ -164,23 +164,23 @@ class ConvolutionVisitor {
 
   // Perform space-to-batch propagation on the convolution. Assumes the
   // activations were already space-to-batched.
-  Status PropagateOnConv(HloInstruction* convolution);
+  absl::Status PropagateOnConv(HloInstruction* convolution);
 
   // Perform space-to-batch propagation on concatenate.
-  Status PropagateOnConcat(HloInstruction* concat);
+  absl::Status PropagateOnConcat(HloInstruction* concat);
 
   // Perform space-to-batch propagation on reverse.
-  Status PropagateOnReverse(HloInstruction* reverse);
+  absl::Status PropagateOnReverse(HloInstruction* reverse);
 
   // Perform space-to-batch propagation on pad.
-  Status PropagateOnPad(HloInstruction* pad);
+  absl::Status PropagateOnPad(HloInstruction* pad);
 
   // Perform space-to-batch propagation on slice.
-  Status PropagateOnSlice(HloInstruction* slice);
+  absl::Status PropagateOnSlice(HloInstruction* slice);
 
   // Perform space-to-batch propagation on the backprop filter convolution.
   // Assumes the activations and kernel were already space-to-batched.
-  Status PropagateOnBackpropFilterConv(HloInstruction* convolution);
+  absl::Status PropagateOnBackpropFilterConv(HloInstruction* convolution);
 
   // Method that checks validity of space-to-batch on a given convolution.
   bool IsConvSuitableForSpaceToBatch(HloInstruction* convolution);
@@ -190,7 +190,7 @@ class ConvolutionVisitor {
 
   // Once a convolution has been space-to-batch'ed, this function will
   // transitively propagate the space-to-batch-ness on rest of the graph.
-  Status PropagateOnUsers(HloInstruction* old_conv);
+  absl::Status PropagateOnUsers(HloInstruction* old_conv);
 
   // Generates masked output with valid data. This is useful when larger shapes
   // are generated due to space-to-batch.
@@ -2634,7 +2634,7 @@ absl::StatusOr<HloInstruction*> ConvolutionVisitor::BatchToSpace(
   return output_transpose;
 }
 
-Status ConvolutionVisitor::PropagateOnUsers(HloInstruction* old_conv) {
+absl::Status ConvolutionVisitor::PropagateOnUsers(HloInstruction* old_conv) {
   std::queue<std::pair<HloInstruction*, HloInstruction*>> propagation_worklist;
 
   if (old_conv->user_count() == 0) {
@@ -2727,7 +2727,7 @@ Status ConvolutionVisitor::PropagateOnUsers(HloInstruction* old_conv) {
   return OkStatus();
 }
 
-Status ConvolutionVisitor::PropagateOnConv(HloInstruction* convolution) {
+absl::Status ConvolutionVisitor::PropagateOnConv(HloInstruction* convolution) {
   auto activations_old = convolution->mutable_operand(0);
 
   CHECK(old_to_new_instrs_.contains(activations_old));
@@ -2934,7 +2934,7 @@ Status ConvolutionVisitor::PropagateOnConv(HloInstruction* convolution) {
   return OkStatus();
 }
 
-Status ConvolutionVisitor::PropagateOnConcat(HloInstruction* concat) {
+absl::Status ConvolutionVisitor::PropagateOnConcat(HloInstruction* concat) {
   auto first_operand = old_to_new_instrs_[concat->mutable_operand(0)];
   auto permute_dims = instr_to_dim_permute_map_[first_operand];
   const int64_t new_concat_dim =
@@ -2957,7 +2957,7 @@ Status ConvolutionVisitor::PropagateOnConcat(HloInstruction* concat) {
   return OkStatus();
 }
 
-Status ConvolutionVisitor::PropagateOnReverse(HloInstruction* reverse) {
+absl::Status ConvolutionVisitor::PropagateOnReverse(HloInstruction* reverse) {
   auto first_operand = old_to_new_instrs_[reverse->mutable_operand(0)];
   auto permute_dims = instr_to_dim_permute_map_[first_operand];
 
@@ -2978,7 +2978,7 @@ Status ConvolutionVisitor::PropagateOnReverse(HloInstruction* reverse) {
   return OkStatus();
 }
 
-Status ConvolutionVisitor::PropagateOnPad(HloInstruction* pad) {
+absl::Status ConvolutionVisitor::PropagateOnPad(HloInstruction* pad) {
   auto first_operand = old_to_new_instrs_[pad->mutable_operand(0)];
   auto permute_dims = instr_to_dim_permute_map_[first_operand];
 
@@ -3009,7 +3009,7 @@ Status ConvolutionVisitor::PropagateOnPad(HloInstruction* pad) {
   return OkStatus();
 }
 
-Status ConvolutionVisitor::PropagateOnSlice(HloInstruction* slice) {
+absl::Status ConvolutionVisitor::PropagateOnSlice(HloInstruction* slice) {
   auto operand = old_to_new_instrs_[slice->mutable_operand(0)];
   auto permute_dims = instr_to_dim_permute_map_[operand];
 
@@ -3230,7 +3230,7 @@ absl::StatusOr<HloInstruction*> ConvolutionVisitor::PropagateOnConstant(
   return new_consumer;
 }
 
-Status ConvolutionVisitor::PropagateOnBackpropFilterConv(
+absl::Status ConvolutionVisitor::PropagateOnBackpropFilterConv(
     HloInstruction* convolution) {
   auto activations_old = convolution->mutable_operand(0);
 
@@ -3865,7 +3865,7 @@ ConvolutionVisitor::ConvDetails ConvolutionVisitor::GetConvolutionDetails(
                      input_dim_size};
 }
 
-Status ConvolutionVisitor::PerformSpaceToBatchOnConvolution(
+absl::Status ConvolutionVisitor::PerformSpaceToBatchOnConvolution(
     HloInstruction* convolution) {
   if (!ConsumeFuel("space-to-batch-converter", [&] {
         return "Skipping space-to-batch propagation because fuel over\n";
diff --git a/third_party/xla/xla/service/stochastic_convert_decomposer.cc b/third_party/xla/xla/service/stochastic_convert_decomposer.cc
index 8829b42ed86456..e4a9adb0c1ef6d 100644
--- a/third_party/xla/xla/service/stochastic_convert_decomposer.cc
+++ b/third_party/xla/xla/service/stochastic_convert_decomposer.cc
@@ -32,8 +32,8 @@ limitations under the License.
 
 namespace xla {
 
-Status DecomposeStochasticConvert(HloComputation* comp,
-                                  HloInstruction* instruction) {
+absl::Status DecomposeStochasticConvert(HloComputation* comp,
+                                        HloInstruction* instruction) {
   CHECK(instruction->opcode() == HloOpcode::kStochasticConvert)
       << "requires a stochastic_convert instruction to decompose, but got: "
       << instruction->opcode();
diff --git a/third_party/xla/xla/service/sub_byte_normalization.cc b/third_party/xla/xla/service/sub_byte_normalization.cc
index d9cf2c8a322c49..1806bcbdedc7d8 100644
--- a/third_party/xla/xla/service/sub_byte_normalization.cc
+++ b/third_party/xla/xla/service/sub_byte_normalization.cc
@@ -88,7 +88,7 @@ absl::StatusOr<bool> SubByteNormalization::Run(
     HloModule* module,
     const absl::flat_hash_set<absl::string_view>& execution_threads) {
   bool changed = false;
-  FunctionVisitor visitor([&](HloInstruction* hlo) -> Status {
+  FunctionVisitor visitor([&](HloInstruction* hlo) -> absl::Status {
     auto* shape = hlo->mutable_shape();
     changed |= UpdateShape(shape, mode_);
     return OkStatus();
diff --git a/third_party/xla/xla/service/topk_rewriter.cc b/third_party/xla/xla/service/topk_rewriter.cc
index 257163c0faf7a2..aee82d6b655a92 100644
--- a/third_party/xla/xla/service/topk_rewriter.cc
+++ b/third_party/xla/xla/service/topk_rewriter.cc
@@ -460,7 +460,7 @@ class TopkDecomposerVisitor : public DfsHloRewriteVisitor {
   explicit TopkDecomposerVisitor(HloPredicate should_decompose)
       : should_decompose_(should_decompose) {}
 
-  Status HandleCustomCall(HloInstruction* inst) override {
+  absl::Status HandleCustomCall(HloInstruction* inst) override {
     if (should_decompose_ && !should_decompose_(inst)) {
       return OkStatus();
     }
@@ -472,7 +472,7 @@ class TopkDecomposerVisitor : public DfsHloRewriteVisitor {
     return DecomposeTopK(call, comparator);
   }
 
-  Status HandleTopK(HloInstruction* topk) override {
+  absl::Status HandleTopK(HloInstruction* topk) override {
     if (should_decompose_ && !should_decompose_(topk)) {
       return OkStatus();
     }
@@ -506,8 +506,8 @@ class TopkDecomposerVisitor : public DfsHloRewriteVisitor {
     return comparator;
   }
 
-  Status DecomposeTopK(HloInstruction* call,
-                       HloComputation* variadic_comparator) {
+  absl::Status DecomposeTopK(HloInstruction* call,
+                             HloComputation* variadic_comparator) {
     HloComputation* comp = call->parent();
     HloInstruction* input = call->mutable_operand(0);
     Shape iota_shape = input->shape();
diff --git a/third_party/xla/xla/service/tpu_computation_placer.cc b/third_party/xla/xla/service/tpu_computation_placer.cc
index ffccfebb3e6e80..e541ac0b970456 100644
--- a/third_party/xla/xla/service/tpu_computation_placer.cc
+++ b/third_party/xla/xla/service/tpu_computation_placer.cc
@@ -23,9 +23,6 @@ limitations under the License.
 namespace tensorflow {
 namespace tpu {
 
-template <typename T>
-using StatusOr = TpuComputationPlacer::StatusOr<T>;
-
 TpuComputationPlacer::TpuComputationPlacer() {
   placer_ = stream_executor::tpu::ExecutorApiFn()->TpuComputationPlacer_NewFn();
 }
@@ -34,13 +31,13 @@ TpuComputationPlacer::~TpuComputationPlacer() {
   stream_executor::tpu::ExecutorApiFn()->TpuComputationPlacer_FreeFn(placer_);
 }
 
-StatusOr<int> TpuComputationPlacer::DeviceId(int replica, int computation,
-                                             int replica_count,
-                                             int computation_count) {
+absl::StatusOr<int> TpuComputationPlacer::DeviceId(int replica, int computation,
+                                                   int replica_count,
+                                                   int computation_count) {
   LOG(FATAL) << "Unimplemented.";
 }
 
-StatusOr<xla::DeviceAssignment> TpuComputationPlacer::AssignDevices(
+absl::StatusOr<xla::DeviceAssignment> TpuComputationPlacer::AssignDevices(
     int replica_count, int computation_count) {
   StatusHelper status;
   xla::DeviceAssignment result(replica_count, computation_count);
@@ -59,7 +56,7 @@ StatusOr<xla::DeviceAssignment> TpuComputationPlacer::AssignDevices(
   return result;
 }
 
-/*static*/ StatusOr<xla::DeviceAssignment>
+/*static*/ absl::StatusOr<xla::DeviceAssignment>
 TpuComputationPlacer::AssignLocalDevices(TpuHostLocationExternal host_location,
                                          int replica_count,
                                          int computation_count) {
diff --git a/third_party/xla/xla/service/tpu_computation_placer.h b/third_party/xla/xla/service/tpu_computation_placer.h
index 60345010683a26..b1780f22b77672 100644
--- a/third_party/xla/xla/service/tpu_computation_placer.h
+++ b/third_party/xla/xla/service/tpu_computation_placer.h
@@ -26,19 +26,16 @@ namespace tpu {
 
 class TpuComputationPlacer : public xla::ComputationPlacer {
  public:
-  template <typename T>
-  using StatusOr = absl::StatusOr<T>;
-
   TpuComputationPlacer();
   ~TpuComputationPlacer() override;
 
-  StatusOr<int> DeviceId(int replica, int computation, int replica_count,
-                         int computation_count) override;
+  absl::StatusOr<int> DeviceId(int replica, int computation, int replica_count,
+                               int computation_count) override;
 
-  StatusOr<xla::DeviceAssignment> AssignDevices(int replica_count,
-                                                int computation_count) override;
+  absl::StatusOr<xla::DeviceAssignment> AssignDevices(
+      int replica_count, int computation_count) override;
 
-  static StatusOr<xla::DeviceAssignment> AssignLocalDevices(
+  static absl::StatusOr<xla::DeviceAssignment> AssignLocalDevices(
       TpuHostLocationExternal host_location, int replica_count,
       int computation_count);
 
diff --git a/third_party/xla/xla/service/transfer_manager.cc b/third_party/xla/xla/service/transfer_manager.cc
index 71ffbe72b68260..a01b4f04060a40 100644
--- a/third_party/xla/xla/service/transfer_manager.cc
+++ b/third_party/xla/xla/service/transfer_manager.cc
@@ -62,7 +62,7 @@ absl::StatusOr<Literal> TransferManager::TransferLiteralFromDevice(
   return std::move(literal);
 }
 
-Status TransferManager::TransferLiteralFromDevice(
+absl::Status TransferManager::TransferLiteralFromDevice(
     se::Stream* stream, const ShapedBuffer& device_buffer,
     const MutableBorrowingLiteral& literal,
     const TransferMetadata* transfer_metadata) {
@@ -70,11 +70,11 @@ Status TransferManager::TransferLiteralFromDevice(
   TF_RETURN_IF_ERROR(substream->WaitFor(stream));
   absl::Cleanup cleanup = [&]() { stream->ReturnSubStream(substream); };
 
-  Status ret;
+  absl::Status ret;
   tsl::Notification n;
   TransferLiteralFromDevice(
       substream, device_buffer, literal,
-      [&](Status status) {
+      [&](absl::Status status) {
         ret = status;
         n.Notify();
       },
@@ -83,7 +83,7 @@ Status TransferManager::TransferLiteralFromDevice(
   return ret;
 }
 
-Status TransferManager::TransferLiteralToDevice(
+absl::Status TransferManager::TransferLiteralToDevice(
     se::Stream* stream, const LiteralSlice& literal,
     const ShapedBuffer& device_buffer,
     const TransferMetadata* transfer_metadata) {
@@ -112,7 +112,7 @@ absl::StatusOr<Literal> TransferManager::TransferArrayFromDevice(
   return std::move(literal);
 }
 
-Status TransferManager::TransferArrayToDevice(
+absl::Status TransferManager::TransferArrayToDevice(
     se::Stream* stream, const LiteralSlice& literal,
     const se::DeviceMemoryBase& dest,
     const TransferMetadata* transfer_metadata) {
@@ -127,7 +127,7 @@ Status TransferManager::TransferArrayToDevice(
   return substream->BlockHostUntilDone();
 }
 
-Status TransferManager::TransferArrayToDeviceAsync(
+absl::Status TransferManager::TransferArrayToDeviceAsync(
     se::Stream* stream, const LiteralSlice& literal,
     const se::DeviceMemoryBase& dest,
     const TransferMetadata* transfer_metadata) {
@@ -139,9 +139,9 @@ Status TransferManager::TransferArrayToDeviceAsync(
                                       transfer_metadata);
 }
 
-Status TransferManager::ReadDynamicShapes(se::Stream* stream,
-                                          const ShapedBuffer* device_buffer,
-                                          Shape* device_shape) {
+absl::Status TransferManager::ReadDynamicShapes(
+    se::Stream* stream, const ShapedBuffer* device_buffer,
+    Shape* device_shape) {
   DCHECK(device_shape->is_dynamic());
   Shape original_device_shape = *device_shape;
   TF_RETURN_IF_ERROR(stream->BlockHostUntilDone());
@@ -223,19 +223,20 @@ Status TransferManager::ReadDynamicShapes(se::Stream* stream,
   return it->second.manager.get();
 }
 
-Status TransferManager::WriteTupleIndexTables(
+absl::Status TransferManager::WriteTupleIndexTables(
     se::Stream* stream, const ShapedBuffer& device_buffer) {
   TF_RETURN_IF_ERROR(WriteTupleIndexTablesAsync(stream, device_buffer));
   return stream->BlockHostUntilDone();
 }
 
-Status TransferManager::WriteTupleIndexTablesAsync(
+absl::Status TransferManager::WriteTupleIndexTablesAsync(
     se::Stream* stream, const ShapedBuffer& device_buffer) {
   VLOG(2) << "Writing tuple index tables for " << device_buffer;
 
   return ShapeUtil::ForEachSubshapeWithStatus(
       device_buffer.on_device_shape(),
-      [&](const Shape& device_subshape, const ShapeIndex& index) -> Status {
+      [&](const Shape& device_subshape,
+          const ShapeIndex& index) -> absl::Status {
         if (device_subshape.IsTuple() &&
             ShapeUtil::TupleElementCount(device_subshape) > 0) {
           se::DeviceMemoryBase device_memory = device_buffer.buffer(index);
@@ -258,7 +259,7 @@ Status TransferManager::WriteTupleIndexTablesAsync(
       });
 }
 
-Status TransferManager::WriteRootTupleIndexTable(
+absl::Status TransferManager::WriteRootTupleIndexTable(
     se::Stream* stream, const ShapedBuffer& device_buffer) {
   TF_RET_CHECK(device_buffer.on_device_shape().IsTuple());
   if (ShapeUtil::TupleElementCount(device_buffer.on_device_shape()) == 0) {
@@ -277,7 +278,7 @@ Status TransferManager::WriteRootTupleIndexTable(
       stream, elements, device_buffer.on_device_shape(), &device_memory);
 }
 
-Status TransferManager::WriteRootTupleIndexTable(
+absl::Status TransferManager::WriteRootTupleIndexTable(
     se::Stream* stream, const ShapeTree<MaybeOwningDeviceMemory>& buffer_tree) {
   TF_RET_CHECK(buffer_tree.shape().IsTuple());
   if (ShapeUtil::TupleElementCount(buffer_tree.shape()) == 0) {
diff --git a/third_party/xla/xla/service/transfer_manager.h b/third_party/xla/xla/service/transfer_manager.h
index b29a93ee6eb551..4ead3279bee2ac 100644
--- a/third_party/xla/xla/service/transfer_manager.h
+++ b/third_party/xla/xla/service/transfer_manager.h
@@ -81,7 +81,7 @@ class TransferManager {
       se::Stream* stream, const ShapedBuffer& device_buffer,
       const TransferMetadata* transfer_metadata = nullptr);
 
-  Status TransferLiteralFromDevice(
+  absl::Status TransferLiteralFromDevice(
       se::Stream* stream, const ShapedBuffer& device_buffer,
       const MutableBorrowingLiteral& literal,
       const TransferMetadata* transfer_metadata = nullptr);
@@ -122,7 +122,7 @@ class TransferManager {
   //
   // Optionally caller can specify platform-specific transfer metadata that
   // tells the actual implementation to do something special.
-  Status TransferLiteralToDevice(
+  absl::Status TransferLiteralToDevice(
       se::Stream* stream, const LiteralSlice& literal,
       const ShapedBuffer& device_buffer,
       const TransferMetadata* transfer_metadata = nullptr);
@@ -143,14 +143,14 @@ class TransferManager {
   //
   // Optionally caller can specify platform-specific transfer metadata that
   // tells the actual implementation to do something special.
-  virtual Status TransferLiteralToDeviceAsync(
+  virtual absl::Status TransferLiteralToDeviceAsync(
       se::Stream* stream, const LiteralSlice& literal,
       const ShapedBuffer& device_buffer,
       const TransferMetadata* transfer_metadata) = 0;
 
-  Status TransferLiteralToDeviceAsync(se::Stream* stream,
-                                      const LiteralSlice& literal,
-                                      const ShapedBuffer& device_buffer) {
+  absl::Status TransferLiteralToDeviceAsync(se::Stream* stream,
+                                            const LiteralSlice& literal,
+                                            const ShapedBuffer& device_buffer) {
     return TransferLiteralToDeviceAsync(stream, literal, device_buffer,
                                         nullptr);
   }
@@ -161,12 +161,12 @@ class TransferManager {
   //
   // Optionally caller can specify platform-specific transfer metadata that
   // tells the actual implementation to do something special.
-  Status TransferArrayToDevice(
+  absl::Status TransferArrayToDevice(
       se::Stream* stream, const LiteralSlice& literal,
       const se::DeviceMemoryBase& dest,
       const TransferMetadata* transfer_metadata = nullptr);
 
-  Status TransferArrayToDeviceAsync(
+  absl::Status TransferArrayToDeviceAsync(
       se::Stream* stream, const LiteralSlice& literal,
       const se::DeviceMemoryBase& dest,
       const TransferMetadata* transfer_metadata = nullptr);
@@ -181,39 +181,39 @@ class TransferManager {
   // shapes, and returns static shapes with dynamic shapes updated.
   // The shape of the buffer also have to be compatible with the host shape and
   // device shape.
-  virtual Status ReadDynamicShapes(se::Stream* stream,
-                                   const ShapedBuffer* device_buffer,
-                                   Shape* device_shape);
+  virtual absl::Status ReadDynamicShapes(se::Stream* stream,
+                                         const ShapedBuffer* device_buffer,
+                                         Shape* device_shape);
 
   // Transfers the given literal into the Infeed interface of the device,
   // using the given executor.
-  virtual Status TransferLiteralToInfeed(se::StreamExecutor* executor,
-                                         const LiteralSlice& literal) = 0;
+  virtual absl::Status TransferLiteralToInfeed(se::StreamExecutor* executor,
+                                               const LiteralSlice& literal) = 0;
 
   // Transfers the given literal from the Outfeed interface of the device,
   // using the given executor. The shape and layout are determined by the
   // shape and layout of `literal`.
-  virtual Status TransferLiteralFromOutfeed(
+  virtual absl::Status TransferLiteralFromOutfeed(
       se::StreamExecutor* executor, MutableBorrowingLiteral literal) = 0;
 
   // Resets the devices associated with this transfer manager.
-  virtual Status ResetDevices(
+  virtual absl::Status ResetDevices(
       absl::Span<se::StreamExecutor* const> executor) = 0;
 
   // Given an allocated ShapedBuffer, constructs the tuple index table(s) in
   // each buffer of the given ShapedBuffer corresponding to tuple shapes. If the
   // ShapedBuffer is array-shaped this method does nothing.
-  Status WriteTupleIndexTables(se::Stream* stream,
-                               const ShapedBuffer& device_buffer);
-  Status WriteTupleIndexTablesAsync(se::Stream* stream,
-                                    const ShapedBuffer& device_buffer);
+  absl::Status WriteTupleIndexTables(se::Stream* stream,
+                                     const ShapedBuffer& device_buffer);
+  absl::Status WriteTupleIndexTablesAsync(se::Stream* stream,
+                                          const ShapedBuffer& device_buffer);
 
   // Writes a tuple index buffer for the root of 'device_buffer', which must
   // be a tuple. Unlike WriteTupleIndexTables, only writes the root buffer,
   // rather than writing all subbuffers. This method is always asynchronous.
-  Status WriteRootTupleIndexTable(se::Stream* stream,
-                                  const ShapedBuffer& device_buffer);
-  Status WriteRootTupleIndexTable(
+  absl::Status WriteRootTupleIndexTable(se::Stream* stream,
+                                        const ShapedBuffer& device_buffer);
+  absl::Status WriteRootTupleIndexTable(
       se::Stream* stream,
       const ShapeTree<MaybeOwningDeviceMemory>& buffer_tree);
 
@@ -294,7 +294,7 @@ class TransferManager {
   // Writes the given device-memory pointers in 'elements' to the given region
   // to construct a tuple index table in the platform-specific tuple
   // representation.
-  virtual Status WriteSingleTupleIndexTable(
+  virtual absl::Status WriteSingleTupleIndexTable(
       se::Stream* stream, absl::Span<const se::DeviceMemoryBase> elements,
       const Shape& shape, se::DeviceMemoryBase* region) = 0;
 
diff --git a/third_party/xla/xla/service/transpose_folding.cc b/third_party/xla/xla/service/transpose_folding.cc
index b954f307f82d4f..2d6dbb2ab3e5a1 100644
--- a/third_party/xla/xla/service/transpose_folding.cc
+++ b/third_party/xla/xla/service/transpose_folding.cc
@@ -78,7 +78,7 @@ using InstructionOperandsPair =
     std::pair<HloInstruction*, TransposeFolding::OperandIndices>;
 
 // Folds the operands of `dot` that are foldable transposes.
-Status FoldTransposeIntoDot(InstructionOperandsPair& pair) {
+absl::Status FoldTransposeIntoDot(InstructionOperandsPair& pair) {
   HloInstruction* dot = pair.first;
 
   DotDimensionNumbers new_dot_dims = dot->dot_dimension_numbers();
diff --git a/third_party/xla/xla/service/tree_reduction_rewriter.cc b/third_party/xla/xla/service/tree_reduction_rewriter.cc
index 1c505358648bd8..633e3925ed3ebe 100644
--- a/third_party/xla/xla/service/tree_reduction_rewriter.cc
+++ b/third_party/xla/xla/service/tree_reduction_rewriter.cc
@@ -41,7 +41,7 @@ class ReductionRewriterVisitor : public DfsHloRewriteVisitor {
   explicit ReductionRewriterVisitor(int64_t reduce_window_size)
       : reduce_window_size_(reduce_window_size) {}
 
-  Status HandleReduce(HloInstruction *hlo) override {
+  absl::Status HandleReduce(HloInstruction *hlo) override {
     HloInstruction *reduced_op = hlo->mutable_operand(0);
     HloInstruction *initial_value = hlo->mutable_operand(1);
     const Shape &input_shape = reduced_op->shape();
diff --git a/third_party/xla/xla/service/tuple_points_to_analysis.cc b/third_party/xla/xla/service/tuple_points_to_analysis.cc
index c90e8836a12ff5..7e91cd787d2d7f 100644
--- a/third_party/xla/xla/service/tuple_points_to_analysis.cc
+++ b/third_party/xla/xla/service/tuple_points_to_analysis.cc
@@ -146,7 +146,7 @@ TuplePointsToAnalysis::Run(const HloModule* module) {
   return std::move(analysis);
 }
 
-Status TuplePointsToAnalysis::Analyze() {
+absl::Status TuplePointsToAnalysis::Analyze() {
   per_instruction_.clear();
   per_instruction_.reserve(module_->instruction_count());
 
@@ -177,7 +177,7 @@ Status TuplePointsToAnalysis::Analyze() {
   return OkStatus();
 }
 
-Status TuplePointsToAnalysis::PopulateDefinedBuffersAndAliases(
+absl::Status TuplePointsToAnalysis::PopulateDefinedBuffersAndAliases(
     const decltype(std::declval<HloComputation>()
                        .instructions())& instructions) {
   for (auto* instruction : instructions) {
@@ -199,7 +199,8 @@ Status TuplePointsToAnalysis::PopulateDefinedBuffersAndAliases(
   return OkStatus();
 }
 
-Status TuplePointsToAnalysis::DefaultAction(HloInstruction* hlo_instruction) {
+absl::Status TuplePointsToAnalysis::DefaultAction(
+    HloInstruction* hlo_instruction) {
   // Create trivial points-to set for instruction. Each points-to set at index i
   // contains a single element LogicalBuffer(hlo_instruction, i). This indicates
   // that this instruction is the source of all buffers in its own output.
@@ -220,7 +221,7 @@ Status TuplePointsToAnalysis::DefaultAction(HloInstruction* hlo_instruction) {
   return OkStatus();
 }
 
-Status TuplePointsToAnalysis::HandleGetTupleElement(
+absl::Status TuplePointsToAnalysis::HandleGetTupleElement(
     HloInstruction* get_tuple_element) {
   // GetTupleElement forwards a pointer to a particular element of the tuple
   // operand.
@@ -252,7 +253,7 @@ Status TuplePointsToAnalysis::HandleGetTupleElement(
   return OkStatus();
 }
 
-Status TuplePointsToAnalysis::HandleCopy(HloInstruction* copy) {
+absl::Status TuplePointsToAnalysis::HandleCopy(HloInstruction* copy) {
   // A kCopy instruction performs a shallow copy of the operand. The top-level
   // buffer (index={}) is newly created, but all other buffers (in the case of a
   // tuple shape) come from the operand
@@ -265,7 +266,7 @@ Status TuplePointsToAnalysis::HandleCopy(HloInstruction* copy) {
   return OkStatus();
 }
 
-Status TuplePointsToAnalysis::HandleBitcast(HloInstruction* bitcast) {
+absl::Status TuplePointsToAnalysis::HandleBitcast(HloInstruction* bitcast) {
   // A kBitcast instruction aliases its operand. That is, the buffer of its
   // result *is* the buffer of its operand, so just copy the operands points-to
   // set.
@@ -273,7 +274,7 @@ Status TuplePointsToAnalysis::HandleBitcast(HloInstruction* bitcast) {
   return OkStatus();
 }
 
-Status TuplePointsToAnalysis::HandleDomain(HloInstruction* domain) {
+absl::Status TuplePointsToAnalysis::HandleDomain(HloInstruction* domain) {
   // A kDomain instruction aliases its operand. That is, the buffer of its
   // result *is* the buffer of its operand, so just copy the operands points-to
   // set.
@@ -281,14 +282,14 @@ Status TuplePointsToAnalysis::HandleDomain(HloInstruction* domain) {
   return OkStatus();
 }
 
-Status TuplePointsToAnalysis::HandleAddDependency(
+absl::Status TuplePointsToAnalysis::HandleAddDependency(
     HloInstruction* add_dependency) {
   // AddDependency just forwards the value of its zero-th operand.
   CreateCopiedPointsToSet(add_dependency, add_dependency->operand(0));
   return OkStatus();
 }
 
-Status TuplePointsToAnalysis::HandleRecvDone(HloInstruction* recv_done) {
+absl::Status TuplePointsToAnalysis::HandleRecvDone(HloInstruction* recv_done) {
   // RecvDone aliases its input (Recv) tuple element {0} to element {0} of its
   // output. The other indices ({} and {1}) define their own buffers.
   PointsToSet& points_to_set = CreateEmptyPointsToSet(recv_done);
@@ -318,7 +319,8 @@ Status TuplePointsToAnalysis::HandleRecvDone(HloInstruction* recv_done) {
   return OkStatus();
 }
 
-Status TuplePointsToAnalysis::HandleAsyncStart(HloInstruction* async_start) {
+absl::Status TuplePointsToAnalysis::HandleAsyncStart(
+    HloInstruction* async_start) {
   // AsyncStart forwards its aliased operands to {0}.
   PointsToSet& points_to_set = CreateEmptyPointsToSet(async_start);
 
@@ -342,7 +344,8 @@ Status TuplePointsToAnalysis::HandleAsyncStart(HloInstruction* async_start) {
   return OkStatus();
 }
 
-Status TuplePointsToAnalysis::HandleAsyncUpdate(HloInstruction* async_update) {
+absl::Status TuplePointsToAnalysis::HandleAsyncUpdate(
+    HloInstruction* async_update) {
   // AsyncUpdate forwards its aliased operand to {}.
   PointsToSet& points_to_set = CreateEmptyPointsToSet(async_update);
   const PointsToSet& operand_points_to_set =
@@ -360,7 +363,8 @@ Status TuplePointsToAnalysis::HandleAsyncUpdate(HloInstruction* async_update) {
   return OkStatus();
 }
 
-Status TuplePointsToAnalysis::HandleAsyncDone(HloInstruction* async_done) {
+absl::Status TuplePointsToAnalysis::HandleAsyncDone(
+    HloInstruction* async_done) {
   // AsyncDone forwards its aliased operand.
   PointsToSet& points_to_set = CreateEmptyPointsToSet(async_done);
   const PointsToSet& operand_points_to_set =
@@ -383,7 +387,8 @@ Status TuplePointsToAnalysis::HandleAsyncDone(HloInstruction* async_done) {
   return OkStatus();
 }
 
-Status TuplePointsToAnalysis::HandleCopyStart(HloInstruction* copy_start) {
+absl::Status TuplePointsToAnalysis::HandleCopyStart(
+    HloInstruction* copy_start) {
   // CopyStart forwards its aliased operand to {1}.
   PointsToSet& points_to_set = CreateEmptyPointsToSet(copy_start);
   const PointsToSet& operand_points_to_set =
@@ -407,7 +412,7 @@ Status TuplePointsToAnalysis::HandleCopyStart(HloInstruction* copy_start) {
   return OkStatus();
 }
 
-Status TuplePointsToAnalysis::HandleCopyDone(HloInstruction* copy_done) {
+absl::Status TuplePointsToAnalysis::HandleCopyDone(HloInstruction* copy_done) {
   // CopyDone forwards its aliased operand.
   PointsToSet& points_to_set = CreateEmptyPointsToSet(copy_done);
   const PointsToSet& operand_points_to_set =
@@ -430,7 +435,7 @@ Status TuplePointsToAnalysis::HandleCopyDone(HloInstruction* copy_done) {
   return OkStatus();
 }
 
-Status TuplePointsToAnalysis::HandleSend(HloInstruction* send) {
+absl::Status TuplePointsToAnalysis::HandleSend(HloInstruction* send) {
   // Send creates a tuple of {aliased operand, U32 context, token}.
   PointsToSet& points_to_set = CreateEmptyPointsToSet(send);
 
@@ -469,7 +474,7 @@ Status TuplePointsToAnalysis::HandleSend(HloInstruction* send) {
   return OkStatus();
 }
 
-Status TuplePointsToAnalysis::HandleTuple(HloInstruction* tuple) {
+absl::Status TuplePointsToAnalysis::HandleTuple(HloInstruction* tuple) {
   absl::Span<HloInstruction* const> operands(tuple->operands());
   PointsToSet& points_to_set = CreateEmptyPointsToSet(tuple);
   points_to_set.AddPointedToBuffer(
@@ -508,7 +513,8 @@ Status TuplePointsToAnalysis::HandleTuple(HloInstruction* tuple) {
   return OkStatus();
 }
 
-Status TuplePointsToAnalysis::HandleCustomCall(HloInstruction* custom_call) {
+absl::Status TuplePointsToAnalysis::HandleCustomCall(
+    HloInstruction* custom_call) {
   auto ccall = Cast<HloCustomCallInstruction>(custom_call);
   PointsToSet& points_to_set = CreateEmptyPointsToSet(custom_call);
   absl::flat_hash_map<ShapeIndex, std::pair<int64_t, ShapeIndex>>
@@ -542,7 +548,7 @@ Status TuplePointsToAnalysis::HandleCustomCall(HloInstruction* custom_call) {
 // WARNING:
 // Adding this, which essentially does the same thing as HandleCustomCall
 // Not sure if it is really needed or it will break anything
-Status TuplePointsToAnalysis::HandleFusion(HloInstruction* fusion) {
+absl::Status TuplePointsToAnalysis::HandleFusion(HloInstruction* fusion) {
   auto cfusion = Cast<HloFusionInstruction>(fusion);
   PointsToSet& points_to_set = CreateEmptyPointsToSet(fusion);
   absl::flat_hash_map<ShapeIndex, std::pair<int64_t, ShapeIndex>>
@@ -573,7 +579,7 @@ Status TuplePointsToAnalysis::HandleFusion(HloInstruction* fusion) {
   return OkStatus();
 }
 
-Status TuplePointsToAnalysis::HandleOptimizationBarrier(
+absl::Status TuplePointsToAnalysis::HandleOptimizationBarrier(
     HloInstruction* barrier) {
   // A kOptimizationBarrier instruction is a no-op.
   CreateCopiedPointsToSet(barrier, barrier->operand(0));
@@ -602,7 +608,8 @@ bool TuplePointsToAnalysis::InstructionDefinesBufferAtIndex(
   return (buffers.size() == 1 && buffers[0]->instruction() == instruction);
 }
 
-Status TuplePointsToAnalysis::VerifyBuffer(const LogicalBuffer& buffer) const {
+absl::Status TuplePointsToAnalysis::VerifyBuffer(
+    const LogicalBuffer& buffer) const {
   if (!InstructionDefinesBufferAtIndex(buffer.instruction(), buffer.index())) {
     return FailedPrecondition(
         "LogicalBuffer %s is ill-defined: instruction %s does not define a "
@@ -654,7 +661,7 @@ TuplePointsToAnalysis::GetBuffersDefinedByInstruction(
   return PerInst(instruction)->instruction_defined_buffers;
 }
 
-Status TuplePointsToAnalysis::GatherBuffersDefinedByInstruction(
+absl::Status TuplePointsToAnalysis::GatherBuffersDefinedByInstruction(
     const HloInstruction* instruction,
     TuplePointsToAnalysis::BufferDefinitionVector* buffers) {
   GetPointsToSet(instruction)
diff --git a/third_party/xla/xla/service/tuple_points_to_analysis.h b/third_party/xla/xla/service/tuple_points_to_analysis.h
index 2b58b1c2e97ace..a80aeaa149ea00 100644
--- a/third_party/xla/xla/service/tuple_points_to_analysis.h
+++ b/third_party/xla/xla/service/tuple_points_to_analysis.h
@@ -132,7 +132,7 @@ class PointsToSet {
     });
   }
   template <typename Fn>
-  Status ForEachElementWithStatus(const Fn& fn) const {
+  absl::Status ForEachElementWithStatus(const Fn& fn) const {
     return tree_.ForEachElementWithStatus(
         [&fn](const ShapeIndex& index, const Elem& elem) {
           return fn(index, elem.buffers);
@@ -241,25 +241,26 @@ class TuplePointsToAnalysis : public DfsHloVisitorWithDefault {
   // which is not defined is a tuple element in a Tuple instruction. In this
   // case, the Tuple instruction does not define the LogicalBuffer, rather that
   // index aliases one of its operands.
-  Status VerifyBuffer(const LogicalBuffer& buffer) const;
-
-  Status DefaultAction(HloInstruction* hlo_instruction) override;
-  Status HandleTuple(HloInstruction* tuple) override;
-  Status HandleGetTupleElement(HloInstruction* get_tuple_element) override;
-  Status HandleAsyncStart(HloInstruction* async_start) override;
-  Status HandleAsyncUpdate(HloInstruction* async_update) override;
-  Status HandleAsyncDone(HloInstruction* async_done) override;
-  Status HandleBitcast(HloInstruction* bitcast) override;
-  Status HandleDomain(HloInstruction* domain) override;
-  Status HandleCopy(HloInstruction* copy) override;
-  Status HandleCopyStart(HloInstruction* copy_start) override;
-  Status HandleCopyDone(HloInstruction* copy_done) override;
-  Status HandleRecvDone(HloInstruction* recv_done) override;
-  Status HandleSend(HloInstruction* send) override;
-  Status HandleAddDependency(HloInstruction* add_dependency) override;
-  Status HandleCustomCall(HloInstruction* custom_call) override;
-  Status HandleFusion(HloInstruction* fusion) override;
-  Status HandleOptimizationBarrier(HloInstruction* barrier) override;
+  absl::Status VerifyBuffer(const LogicalBuffer& buffer) const;
+
+  absl::Status DefaultAction(HloInstruction* hlo_instruction) override;
+  absl::Status HandleTuple(HloInstruction* tuple) override;
+  absl::Status HandleGetTupleElement(
+      HloInstruction* get_tuple_element) override;
+  absl::Status HandleAsyncStart(HloInstruction* async_start) override;
+  absl::Status HandleAsyncUpdate(HloInstruction* async_update) override;
+  absl::Status HandleAsyncDone(HloInstruction* async_done) override;
+  absl::Status HandleBitcast(HloInstruction* bitcast) override;
+  absl::Status HandleDomain(HloInstruction* domain) override;
+  absl::Status HandleCopy(HloInstruction* copy) override;
+  absl::Status HandleCopyStart(HloInstruction* copy_start) override;
+  absl::Status HandleCopyDone(HloInstruction* copy_done) override;
+  absl::Status HandleRecvDone(HloInstruction* recv_done) override;
+  absl::Status HandleSend(HloInstruction* send) override;
+  absl::Status HandleAddDependency(HloInstruction* add_dependency) override;
+  absl::Status HandleCustomCall(HloInstruction* custom_call) override;
+  absl::Status HandleFusion(HloInstruction* fusion) override;
+  absl::Status HandleOptimizationBarrier(HloInstruction* barrier) override;
 
   std::string ToString() const;
 
@@ -280,11 +281,11 @@ class TuplePointsToAnalysis : public DfsHloVisitorWithDefault {
 
   // Perform the analysis. Should be called immediately after constructing the
   // object and before calling GetPointsToSet.
-  Status Analyze();
+  absl::Status Analyze();
 
   // Populates instruction-defined buffers and aliases for each instruction
   // in 'instructions'.
-  Status PopulateDefinedBuffersAndAliases(
+  absl::Status PopulateDefinedBuffersAndAliases(
       const decltype(std::declval<HloComputation>()
                          .instructions())& instructions);
 
@@ -298,8 +299,8 @@ class TuplePointsToAnalysis : public DfsHloVisitorWithDefault {
                                        const HloInstruction* src);
 
   // Adds the buffers defined by the given instruction to the given vector.
-  Status GatherBuffersDefinedByInstruction(const HloInstruction* instruction,
-                                           BufferDefinitionVector* buffers);
+  absl::Status GatherBuffersDefinedByInstruction(
+      const HloInstruction* instruction, BufferDefinitionVector* buffers);
 
   // Print points-to set for 'instruction' to 'output'.
   void InstructionToString(const HloInstruction* instruction,
diff --git a/third_party/xla/xla/service/while_loop_all_reduce_code_motion.cc b/third_party/xla/xla/service/while_loop_all_reduce_code_motion.cc
index eb0fe9bee4c4b4..0e2bf23d303a13 100644
--- a/third_party/xla/xla/service/while_loop_all_reduce_code_motion.cc
+++ b/third_party/xla/xla/service/while_loop_all_reduce_code_motion.cc
@@ -673,7 +673,7 @@ WhileInitContext CreateNewWhileInit(
 // When moving reduce-scatter outside the while body, change the associated
 // accumulation buffers to use the shape of the operand of the reduce-scatter
 // (i.e., the pre-scatter shape).
-Status ChangeAccumulatorShapesInLoopBodies(
+absl::Status ChangeAccumulatorShapesInLoopBodies(
     HloInstruction* old_while_instruction,
     const HloInstructionMap<std::vector<AccumulationContext>>&
         all_reduce_to_accumulations) {
@@ -893,7 +893,7 @@ HloInstruction* CreateNewWhileResult(
 // The all-reduce outputs are then added to the original accumulation buffers.
 // Creates a tuple that groups the while loop output and the accumulated
 // buffers and replaces all uses of the old while with this new tuple.
-Status AddSinkedAllReducesAndReplaceWhile(
+absl::Status AddSinkedAllReducesAndReplaceWhile(
     HloInstruction* while_instruction,
     const HloInstructionMap<std::vector<AccumulationContext>>&
         all_reduce_to_accumulations) {
diff --git a/third_party/xla/xla/service/while_loop_concat_code_motion.cc b/third_party/xla/xla/service/while_loop_concat_code_motion.cc
index d7cf929ae7c541..015d97aa82682a 100644
--- a/third_party/xla/xla/service/while_loop_concat_code_motion.cc
+++ b/third_party/xla/xla/service/while_loop_concat_code_motion.cc
@@ -652,9 +652,9 @@ std::vector<bool> TupleElementsUsedInCond(HloInstruction* loop) {
 
 // Adds copies to returned values to keep RewriteLoopWithConcatGroups simple:
 // the copies do not have other users and only appear once in the root tuple.
-Status AddCopiesToRoot(HloComputation* body,
-                       absl::Span<HloInstruction* const> param_gtes,
-                       ConcatGroups* groups) {
+absl::Status AddCopiesToRoot(HloComputation* body,
+                             absl::Span<HloInstruction* const> param_gtes,
+                             ConcatGroups* groups) {
   auto root = body->root_instruction();
   CHECK_EQ(root->opcode(), HloOpcode::kTuple);
   std::vector<HloInstruction*> copies(root->operand_count(), nullptr);
@@ -691,7 +691,7 @@ Status AddCopiesToRoot(HloComputation* body,
   return OkStatus();
 }
 
-Status RemoveCopiesFromRoot(HloComputation* body) {
+absl::Status RemoveCopiesFromRoot(HloComputation* body) {
   auto root = body->root_instruction();
   CHECK_EQ(root->opcode(), HloOpcode::kTuple);
   for (int64_t i = 0; i < root->operand_count(); ++i) {
@@ -703,9 +703,9 @@ Status RemoveCopiesFromRoot(HloComputation* body) {
   return OkStatus();
 }
 
-Status RewriteLoopWithConcatGroups(HloInstruction* loop,
-                                   absl::Span<HloInstruction* const> param_gtes,
-                                   ConcatGroups& groups) {
+absl::Status RewriteLoopWithConcatGroups(
+    HloInstruction* loop, absl::Span<HloInstruction* const> param_gtes,
+    ConcatGroups& groups) {
   VLOG(1) << "RewriteLoopWithConcatGroups with " << groups.Groups().size()
           << " groups.";
   // For simplicity, for each group, we rewrite the first element into full
diff --git a/third_party/xla/xla/service/while_loop_constant_sinking.cc b/third_party/xla/xla/service/while_loop_constant_sinking.cc
index d714cdb25a0216..339ea76b4d8be8 100644
--- a/third_party/xla/xla/service/while_loop_constant_sinking.cc
+++ b/third_party/xla/xla/service/while_loop_constant_sinking.cc
@@ -28,10 +28,9 @@ namespace {
 // `while_body_root` (which must be a tuple instruction) at index `tuple_index`.
 // This utility helps us replace an instruction in the while body with a
 // constant while still keeping it trivially loop invariant.
-Status ReplaceUsesWhileKeepingLoopInvariance(HloInstruction* old_instr,
-                                             HloInstruction* new_instr,
-                                             HloInstruction* while_body_root,
-                                             int64_t tuple_index) {
+absl::Status ReplaceUsesWhileKeepingLoopInvariance(
+    HloInstruction* old_instr, HloInstruction* new_instr,
+    HloInstruction* while_body_root, int64_t tuple_index) {
   CHECK_EQ(while_body_root->opcode(), HloOpcode::kTuple);
 
   std::vector<HloInstruction*> users;

From 1392171d25d4e341cfd7c702f3af591992dccef5 Mon Sep 17 00:00:00 2001
From: Kyle Lucke <klucke@google.com>
Date: Thu, 16 May 2024 15:04:34 -0700
Subject: [PATCH 308/478] Use absl::Status instead of xla::Status now that
 they're identical.

PiperOrigin-RevId: 634541223
---
 third_party/xla/xla/array.h                | 12 ++--
 third_party/xla/xla/layout_util.cc         | 10 +--
 third_party/xla/xla/layout_util.h          | 10 +--
 third_party/xla/xla/literal.cc             | 53 +++++++-------
 third_party/xla/xla/literal.h              | 84 +++++++++++-----------
 third_party/xla/xla/literal_comparison.cc  | 82 ++++++++++-----------
 third_party/xla/xla/literal_comparison.h   | 14 ++--
 third_party/xla/xla/literal_test.cc        | 24 +++----
 third_party/xla/xla/literal_util.h         |  6 +-
 third_party/xla/xla/map_util.h             |  2 +-
 third_party/xla/xla/protobuf_util.cc       |  8 +--
 third_party/xla/xla/protobuf_util.h        |  8 +--
 third_party/xla/xla/service_interface.h    | 64 +++++++++--------
 third_party/xla/xla/shape_layout.cc        |  4 +-
 third_party/xla/xla/shape_layout.h         |  4 +-
 third_party/xla/xla/shape_tree.h           | 23 +++---
 third_party/xla/xla/shape_util.cc          | 25 +++----
 third_party/xla/xla/shape_util.h           | 70 +++++++++---------
 third_party/xla/xla/shape_util_test.cc     |  2 +-
 third_party/xla/xla/sharding_op_util.cc    |  4 +-
 third_party/xla/xla/sharding_op_util.h     |  4 +-
 third_party/xla/xla/status_macros.cc       | 33 ++++-----
 third_party/xla/xla/status_macros.h        | 34 ++++-----
 third_party/xla/xla/status_macros_test.cc  | 22 +++---
 third_party/xla/xla/statusor.h             |  2 +-
 third_party/xla/xla/test_helpers.h         | 11 +--
 third_party/xla/xla/text_literal_reader.cc |  5 +-
 third_party/xla/xla/text_literal_writer.cc |  6 +-
 third_party/xla/xla/text_literal_writer.h  |  3 +-
 third_party/xla/xla/util.cc                | 12 ++--
 third_party/xla/xla/util.h                 | 34 ++++-----
 31 files changed, 349 insertions(+), 326 deletions(-)

diff --git a/third_party/xla/xla/array.h b/third_party/xla/xla/array.h
index 276b3b556c4a76..98ea2858ab4869 100644
--- a/third_party/xla/xla/array.h
+++ b/third_party/xla/xla/array.h
@@ -315,11 +315,11 @@ class Array {
   // Invokes a callback with the (indices, value_ptr) for each cell in the
   // array. If a callback returns a non-OK status, returns that else returns
   // OkStatus().
-  Status EachStatus(
-      absl::FunctionRef<Status(absl::Span<const int64_t>, T*)> f) {
+  absl::Status EachStatus(
+      absl::FunctionRef<absl::Status(absl::Span<const int64_t>, T*)> f) {
     OwnedBuffer<int64_t> index(sizes_.size, default_init_t{});
     for (int64_t i = 0; i < num_elements(); ++i, next_index(&index)) {
-      Status s = f(index.span(), &values_[i]);
+      absl::Status s = f(index.span(), &values_[i]);
       if (!s.ok()) {
         return s;
       }
@@ -330,11 +330,11 @@ class Array {
   // Invokes a callback with the (indices, value) for each cell in the array.
   // If a callback returns a non-OK status, returns that else returns
   // OkStatus().
-  Status EachStatus(
-      absl::FunctionRef<Status(absl::Span<const int64_t>, T)> f) const {
+  absl::Status EachStatus(
+      absl::FunctionRef<absl::Status(absl::Span<const int64_t>, T)> f) const {
     OwnedBuffer<int64_t> index(sizes_.size, default_init_t{});
     for (int64_t i = 0; i < num_elements(); ++i, next_index(&index)) {
-      Status s = f(index.span(), values_[i]);
+      absl::Status s = f(index.span(), values_[i]);
       if (!s.ok()) {
         return s;
       }
diff --git a/third_party/xla/xla/layout_util.cc b/third_party/xla/xla/layout_util.cc
index 03b83e22c20436..ed57d57cf4f286 100644
--- a/third_party/xla/xla/layout_util.cc
+++ b/third_party/xla/xla/layout_util.cc
@@ -204,7 +204,7 @@ Layout CreateDefaultLayoutForRank(int64_t rank) {
   LayoutUtil::SetToDefaultLayout(program_shape->mutable_result());
 }
 
-/* static */ Status LayoutUtil::ValidateLayoutInShape(
+/* static */ absl::Status LayoutUtil::ValidateLayoutInShape(
     const Shape& shape, bool allow_missing_layouts) {
   if (shape.IsTuple()) {
     // Tuple shape.
@@ -236,8 +236,8 @@ Layout CreateDefaultLayoutForRank(int64_t rank) {
   }
 }
 
-/* static */ Status LayoutUtil::ValidateLayoutForShape(const Layout& layout,
-                                                       const Shape& shape) {
+/* static */ absl::Status LayoutUtil::ValidateLayoutForShape(
+    const Layout& layout, const Shape& shape) {
   if (shape.IsTuple()) {
     return InvalidArgument("a single Layout is not valid for tuple shapes");
   }
@@ -586,7 +586,7 @@ Layout CreateDefaultLayoutForRank(int64_t rank) {
 namespace {
 
 // Internal helper for recursively copying layouts.
-Status CopyLayoutInternal(const Shape& src, Shape* dst) {
+absl::Status CopyLayoutInternal(const Shape& src, Shape* dst) {
   if (src.IsTuple() != dst->IsTuple()) {
     return InvalidArgument(
         "cannot copy layout from shape: shape structure differs");
@@ -619,7 +619,7 @@ Status CopyLayoutInternal(const Shape& src, Shape* dst) {
 }  // namespace
 
 /* static */
-Status LayoutUtil::CopyLayoutBetweenShapes(const Shape& src, Shape* dst) {
+absl::Status LayoutUtil::CopyLayoutBetweenShapes(const Shape& src, Shape* dst) {
   return CopyLayoutInternal(src, dst);
 }
 
diff --git a/third_party/xla/xla/layout_util.h b/third_party/xla/xla/layout_util.h
index 7cff0b3ce80ad1..ca48164965ec50 100644
--- a/third_party/xla/xla/layout_util.h
+++ b/third_party/xla/xla/layout_util.h
@@ -90,13 +90,13 @@ class LayoutUtil {
   // Validates that the layout within the given shape is correct. The check
   // is performed for all subshapes as well. If missing layouts are allowed
   // the check does not fail on array shapes without layouts.
-  static Status ValidateLayoutInShape(const Shape& shape,
-                                      bool allow_missing_layouts = false);
+  static absl::Status ValidateLayoutInShape(const Shape& shape,
+                                            bool allow_missing_layouts = false);
 
   // Validates that the provided layout satisfies invariants for the given
   // shape.
-  static Status ValidateLayoutForShape(const Layout& layout,
-                                       const Shape& shape);
+  static absl::Status ValidateLayoutForShape(const Layout& layout,
+                                             const Shape& shape);
 
   // Clears the layout in the given Shape. After this function is called,
   // HasLayout will return false for the shape.
@@ -242,7 +242,7 @@ class LayoutUtil {
   // tuples.  'src' and 'dst' need not be compatible but the two shapes must
   // have the same tuple structure (if any) and arrays must have the same
   // rank. within the shapes must have the same number of dimensions.
-  static Status CopyLayoutBetweenShapes(const Shape& src, Shape* dst);
+  static absl::Status CopyLayoutBetweenShapes(const Shape& src, Shape* dst);
 
   // Returns true if the layouts of lhs and rhs are equal, false
   // otherwise. Recursively compares layouts of tuples.
diff --git a/third_party/xla/xla/literal.cc b/third_party/xla/xla/literal.cc
index d7329d76b37120..fe35a577ebb37b 100644
--- a/third_party/xla/xla/literal.cc
+++ b/third_party/xla/xla/literal.cc
@@ -405,7 +405,7 @@ absl::StatusOr<std::string> LiteralBase::SerializeAsString() const {
 }
 
 template <typename NativeT>
-Status MutableLiteralBase::CopySliceFromInternal(
+absl::Status MutableLiteralBase::CopySliceFromInternal(
     const LiteralBase& src_literal, absl::Span<const int64_t> src_base,
     absl::Span<const int64_t> dest_base, absl::Span<const int64_t> copy_size) {
   auto linear_index = [](const Shape& shape,
@@ -674,8 +674,8 @@ void LiteralBase::Piece::CopyElementsWithDynamicBound(
   } while (IndexUtil::BumpIndices(bound_shape, absl::MakeSpan(index)));
 }
 
-Status LiteralBase::Piece::CopyFrom(const LiteralBase::Piece& src,
-                                    bool only_dynamic_bound) {
+absl::Status LiteralBase::Piece::CopyFrom(const LiteralBase::Piece& src,
+                                          bool only_dynamic_bound) {
   CHECK(subshape_ != nullptr);
   CHECK(src.subshape_ != nullptr);
   CHECK(LayoutUtil::IsDenseArray(subshape()))
@@ -745,10 +745,10 @@ void MutableLiteralBase::SetDynamicSize(int64_t dim_index,
   piece(shape_index).SetDynamicSize(dim_index, size);
 }
 
-Status MutableLiteralBase::CopyFrom(const LiteralSlice& src_literal,
-                                    const ShapeIndex& dest_shape_index,
-                                    const ShapeIndex& src_shape_index,
-                                    bool only_dynamic_bound) {
+absl::Status MutableLiteralBase::CopyFrom(const LiteralSlice& src_literal,
+                                          const ShapeIndex& dest_shape_index,
+                                          const ShapeIndex& src_shape_index,
+                                          bool only_dynamic_bound) {
   const Shape& dest_subshape =
       ShapeUtil::GetSubshape(shape(), dest_shape_index);
   const Shape& src_subshape =
@@ -799,8 +799,8 @@ Status MutableLiteralBase::CopyFrom(const LiteralSlice& src_literal,
       });
 }
 
-Status Literal::MoveFrom(Literal&& src_literal,
-                         const ShapeIndex& dest_shape_index) {
+absl::Status Literal::MoveFrom(Literal&& src_literal,
+                               const ShapeIndex& dest_shape_index) {
   const Shape& dest_subshape =
       ShapeUtil::GetSubshape(shape(), dest_shape_index);
   if (!ShapeUtil::Equal(dest_subshape, src_literal.shape())) {
@@ -832,10 +832,9 @@ Status Literal::MoveFrom(Literal&& src_literal,
   return OkStatus();
 }
 
-Status MutableLiteralBase::CopySliceFrom(const LiteralSlice& src_literal,
-                                         absl::Span<const int64_t> src_base,
-                                         absl::Span<const int64_t> dest_base,
-                                         absl::Span<const int64_t> copy_size) {
+absl::Status MutableLiteralBase::CopySliceFrom(
+    const LiteralSlice& src_literal, absl::Span<const int64_t> src_base,
+    absl::Span<const int64_t> dest_base, absl::Span<const int64_t> copy_size) {
   TF_RET_CHECK(LayoutUtil::IsDenseArray(shape())) << shape();
   TF_RET_CHECK(LayoutUtil::IsDenseArray(src_literal.shape()))
       << src_literal.shape();
@@ -844,7 +843,7 @@ Status MutableLiteralBase::CopySliceFrom(const LiteralSlice& src_literal,
   TF_RET_CHECK(shape().rank() == dest_base.size());
 
   return primitive_util::ArrayTypeSwitch<Status>(
-      [&](auto primitive_type_constant) -> Status {
+      [&](auto primitive_type_constant) -> absl::Status {
         using NativeT = NativeTypeOf<primitive_type_constant>;
         return CopySliceFromInternal<NativeT>(src_literal, src_base, dest_base,
                                               copy_size);
@@ -927,7 +926,7 @@ void MutableLiteralBase::PopulateInplaceInternal(
   }
 }
 
-Status MutableLiteralBase::PopulateInplace(
+absl::Status MutableLiteralBase::PopulateInplace(
     absl::FunctionRef<void(void*, absl::Span<const int64_t>)> populator) {
   TF_RET_CHECK(LayoutUtil::IsDenseArray(shape()))
       << __func__ << " is only supported for dense arrays: " << shape();
@@ -939,7 +938,7 @@ Status MutableLiteralBase::PopulateInplace(
   return OkStatus();
 }
 
-Status MutableLiteralBase::PopulateInplaceParallel(
+absl::Status MutableLiteralBase::PopulateInplaceParallel(
     absl::FunctionRef<void(void*, absl::Span<const int64_t>, int)> populator) {
   TF_RET_CHECK(LayoutUtil::IsDenseArray(shape()))
       << __func__ << " is only supported for dense arrays: " << shape();
@@ -1382,11 +1381,11 @@ std::optional<complex128> LiteralBase::GetAsComplex128(
       shape().element_type());
 }
 
-Status MutableLiteralBase::SetIntegralAsS64(
+absl::Status MutableLiteralBase::SetIntegralAsS64(
     absl::Span<const int64_t> multi_index, int64_t value) {
   CHECK(LayoutUtil::IsDenseArray(shape()));
   return primitive_util::PrimitiveTypeSwitch<Status>(
-      [&](auto primitive_type_constant) -> Status {
+      [&](auto primitive_type_constant) -> absl::Status {
         if constexpr (primitive_util::IsIntegralType(primitive_type_constant) ||
                       primitive_type_constant == PRED) {
           using NativeT = NativeTypeOf<primitive_type_constant>;
@@ -1399,8 +1398,8 @@ Status MutableLiteralBase::SetIntegralAsS64(
       shape().element_type());
 }
 
-Status MutableLiteralBase::SetFromDouble(absl::Span<const int64_t> multi_index,
-                                         double value) {
+absl::Status MutableLiteralBase::SetFromDouble(
+    absl::Span<const int64_t> multi_index, double value) {
   CHECK(LayoutUtil::IsDenseArray(shape()));
   if (!primitive_util::IsFloatingPointType(shape().element_type())) {
     return FailedPrecondition("Array element type is not integral: %s",
@@ -1695,8 +1694,8 @@ void ConvertBetweenNativeTypes(absl::Span<const NativeSrcT> src_data,
 }
 
 template <PrimitiveType kSrcType>
-Status ConvertIfDestTypeMatches(const LiteralBase& src_literal,
-                                MutableLiteralBase& dst_literal) {
+absl::Status ConvertIfDestTypeMatches(const LiteralBase& src_literal,
+                                      MutableLiteralBase& dst_literal) {
   DCHECK(dst_literal.shape().IsArray());
   using NativeSrcT = NativeTypeOf<kSrcType>;
   // Pass raw data Span/pointers to called template methods to avoid duplicating
@@ -1705,7 +1704,7 @@ Status ConvertIfDestTypeMatches(const LiteralBase& src_literal,
   void* dst_base = dst_literal.untyped_data();
   DCHECK_EQ(src_data.size(), dst_literal.element_count());
   return primitive_util::ArrayTypeSwitch<Status>(
-      [&](auto primitive_type_constant) -> Status {
+      [&](auto primitive_type_constant) -> absl::Status {
         if constexpr (primitive_util::IsComplexType(kSrcType) &&
                       !primitive_util::IsComplexType(primitive_type_constant)) {
           return Unimplemented("%s from type %s to type %s is not implemented.",
@@ -1741,7 +1740,7 @@ absl::StatusOr<Literal> ConvertSwitch(const LiteralBase& literal,
   Literal result(
       ShapeUtil::ChangeElementType(literal.shape(), primitive_dest_type));
   TF_RETURN_IF_ERROR(primitive_util::ArrayTypeSwitch<Status>(
-      [&](auto primitive_type_constant) -> Status {
+      [&](auto primitive_type_constant) -> absl::Status {
         return ConvertIfDestTypeMatches<primitive_type_constant>(literal,
                                                                  result);
       },
@@ -2336,8 +2335,8 @@ void* LiteralBase::Piece::untyped_data() {
 namespace {
 
 template <typename RepeatedFieldT, typename NativeT>
-Status CopyFromRepeatedField(absl::Span<NativeT> dest,
-                             const RepeatedFieldT& src) {
+absl::Status CopyFromRepeatedField(absl::Span<NativeT> dest,
+                                   const RepeatedFieldT& src) {
   if (dest.size() != src.size()) {
     return InvalidArgument(
         "Expected %lu elements in LiteralProto repeated field, has %d",
@@ -2349,7 +2348,7 @@ Status CopyFromRepeatedField(absl::Span<NativeT> dest,
 
 }  // namespace
 
-Status LiteralBase::Piece::CopyFromProto(const LiteralProto& proto) {
+absl::Status LiteralBase::Piece::CopyFromProto(const LiteralProto& proto) {
   // These conditions should have been checked in
   // MutableLiteralBase::CreateFromProto.
   TF_RET_CHECK(proto.has_shape());
diff --git a/third_party/xla/xla/literal.h b/third_party/xla/xla/literal.h
index 4f8ce8e57fa3b0..eb88978687fbc5 100644
--- a/third_party/xla/xla/literal.h
+++ b/third_party/xla/xla/literal.h
@@ -115,7 +115,7 @@ class LiteralBase {
   // compatibility.  If compatibility is required, you should use protobuf
   // serialization instead.
   template <typename OutputIterator>
-  Status Serialize(OutputIterator output) const {
+  absl::Status Serialize(OutputIterator output) const {
     return SerializeWithShapeProto(shape().ToProto(), output);
   }
 
@@ -525,8 +525,8 @@ class LiteralBase {
   void BuildPieceSubtree(const Shape& shape, Piece* piece);
 
   template <typename OutputIterator>
-  Status SerializeWithShapeProto(const ShapeProto& proto,
-                                 OutputIterator output) const;
+  absl::Status SerializeWithShapeProto(const ShapeProto& proto,
+                                       OutputIterator output) const;
 
   template <typename OutputIterator>
   class SerializeState {
@@ -919,10 +919,10 @@ class LiteralBase {
           .IgnoreError();
     }
     // Same as above, but the function has the type:
-    //    Status (const ShapeIndex& index, const Piece& piece)
+    //    absl::Status (const ShapeIndex& index, const Piece& piece)
     // The first non-OK return value is returned by the function.
     template <typename Fn>
-    Status ForEachSubpieceWithStatus(const Fn& func) const {
+    absl::Status ForEachSubpieceWithStatus(const Fn& func) const {
       ShapeIndex index;
       return ForEachHelper(func, *this, &index);
     }
@@ -948,10 +948,10 @@ class LiteralBase {
           .IgnoreError();
     }
     // Same as above, but the function has the type:
-    //    Status (const ShapeIndex& index, Piece& piece)
+    //    absl::Status (const ShapeIndex& index, Piece& piece)
     // The first non-OK return value is returned by the function.
     template <typename Fn>
-    Status ForEachMutableSubpieceWithStatus(const Fn& func) {
+    absl::Status ForEachMutableSubpieceWithStatus(const Fn& func) {
       ShapeIndex index;
       return ForEachMutableHelper(
           func, const_cast<xla::LiteralBase::Piece*>(this), &index);
@@ -985,11 +985,11 @@ class LiteralBase {
     // Copy the data from 'src' into this piece's buffer. Shapes of this piece
     // and src must be compatible. If only_dynamic_bound is true, only elements
     // within dynamic bounds will be copied.
-    Status CopyFrom(const Piece& src, bool only_dynamic_bound);
+    absl::Status CopyFrom(const Piece& src, bool only_dynamic_bound);
 
     // Copies the data from the given proto into this piece. The shape of this
     // piece must be equal (not just compatible) to the shape of the proto.
-    Status CopyFromProto(const LiteralProto& proto);
+    absl::Status CopyFromProto(const LiteralProto& proto);
 
     // See comments on ArrayValueState for detailed explanation.
     bool IsDetermined() const;
@@ -1048,8 +1048,8 @@ class LiteralBase {
     // The callable 'func' has the same signature as described above in
     // ForEachSubpiece*.
     template <typename Fn>
-    Status ForEachHelper(const Fn& func, const Piece& piece,
-                         ShapeIndex* index) const {
+    absl::Status ForEachHelper(const Fn& func, const Piece& piece,
+                               ShapeIndex* index) const {
       TF_RETURN_IF_ERROR(func(*index, piece));
       if (auto* tuple_rep = piece.GetTupleRep()) {
         for (int64_t i = 0; i < tuple_rep->children.size(); ++i) {
@@ -1079,8 +1079,8 @@ class LiteralBase {
       return true;
     }
     template <typename Fn>
-    Status ForEachMutableHelper(const Fn& func, Piece* piece,
-                                ShapeIndex* index) {
+    absl::Status ForEachMutableHelper(const Fn& func, Piece* piece,
+                                      ShapeIndex* index) {
       TF_RETURN_IF_ERROR(func(*index, piece));
       if (auto* tuple_rep = piece->GetTupleRep()) {
         for (int64_t i = 0; i < tuple_rep->children.size(); ++i) {
@@ -1170,10 +1170,10 @@ class MutableLiteralBase : public LiteralBase {
   // at 'dest_shape_index' must be compatible with the subshape of 'src_literal'
   // rooted at 'src_shape_index', but need not be arrays. If only_dynamic_bound
   // is true, only elements within dynamic bounds will be copied.
-  Status CopyFrom(const LiteralSlice& src_literal,
-                  const ShapeIndex& dest_shape_index = {},
-                  const ShapeIndex& src_shape_index = {},
-                  bool only_dynamic_bound = false);
+  absl::Status CopyFrom(const LiteralSlice& src_literal,
+                        const ShapeIndex& dest_shape_index = {},
+                        const ShapeIndex& src_shape_index = {},
+                        bool only_dynamic_bound = false);
 
   // Copies the values from src_literal, starting at src_base shape indexes,
   // to this literal, starting at dest_base, where the copy size in each
@@ -1185,10 +1185,10 @@ class MutableLiteralBase : public LiteralBase {
   // element, then copy_size must be 0 in these dimensions while the
   // corresponding base indices being 0.
   // This literal and 'src_literal' must be arrays.
-  Status CopySliceFrom(const LiteralSlice& src_literal,
-                       absl::Span<const int64_t> src_base,
-                       absl::Span<const int64_t> dest_base,
-                       absl::Span<const int64_t> copy_size);
+  absl::Status CopySliceFrom(const LiteralSlice& src_literal,
+                             absl::Span<const int64_t> src_base,
+                             absl::Span<const int64_t> dest_base,
+                             absl::Span<const int64_t> copy_size);
 
   // Copies one element from src_literal[src_index] to (*this)[dest_index].
   void CopyElementFrom(const LiteralSlice& src_literal,
@@ -1207,11 +1207,13 @@ class MutableLiteralBase : public LiteralBase {
 
   // As Set(), but truncates `value` to the literal element type before storing.
   // This literal must be an array.
-  Status SetIntegralAsS64(absl::Span<const int64_t> multi_index, int64_t value);
+  absl::Status SetIntegralAsS64(absl::Span<const int64_t> multi_index,
+                                int64_t value);
 
   // As Set(), but truncates `value` to the literal element type before storing.
   // This literal must be an array.
-  Status SetFromDouble(absl::Span<const int64_t> multi_index, double value);
+  absl::Status SetFromDouble(absl::Span<const int64_t> multi_index,
+                             double value);
 
   // Populate this literal with the given values. Examples:
   //
@@ -1247,14 +1249,14 @@ class MutableLiteralBase : public LiteralBase {
   //
   // This literal must have a dense layout.
   template <typename NativeT>
-  Status Populate(
+  absl::Status Populate(
       absl::FunctionRef<NativeT(absl::Span<const int64_t>)> generator);
 
   // A parallel version of Populate(). This can be used if the generator is
   // thread-safe and the values for the shape's different elements are
   // independent.
   template <typename NativeT>
-  Status PopulateParallel(
+  absl::Status PopulateParallel(
       absl::FunctionRef<NativeT(absl::Span<const int64_t>, int)> generator);
 
   // Similar to Populate() but takes a populator function that allows caller
@@ -1264,13 +1266,13 @@ class MutableLiteralBase : public LiteralBase {
   // that we can avoid templatizing the method for better code size.
   //
   // This literal must have a dense layout.
-  Status PopulateInplace(
+  absl::Status PopulateInplace(
       absl::FunctionRef<void(void*, absl::Span<const int64_t>)> populator);
 
   // A parallel version of PopulateInplace(). This can be used if the generator
   // is thread-safe and the values for the shape's different elements are
   // independent.
-  Status PopulateInplaceParallel(
+  absl::Status PopulateInplaceParallel(
       absl::FunctionRef<void(void*, absl::Span<const int64_t>, int)> populator);
 
   // Fills this literal with the given value.
@@ -1298,10 +1300,10 @@ class MutableLiteralBase : public LiteralBase {
   // Internal template helper for the Literal::CopySliceFrom(), matching its
   // arguments one by one.
   template <typename NativeT>
-  Status CopySliceFromInternal(const LiteralBase& src_literal,
-                               absl::Span<const int64_t> src_base,
-                               absl::Span<const int64_t> dest_base,
-                               absl::Span<const int64_t> copy_size);
+  absl::Status CopySliceFromInternal(const LiteralBase& src_literal,
+                                     absl::Span<const int64_t> src_base,
+                                     absl::Span<const int64_t> dest_base,
+                                     absl::Span<const int64_t> copy_size);
 
   // The literal may or may not own the storage of the shape. Creating/copying a
   // shape can incur significant overhead which in many case we'd like to avoid,
@@ -1313,9 +1315,9 @@ class MutableLiteralBase : public LiteralBase {
 
   // Implementation details shared between Populate() and PopulateParallel()
   //  template <typename NativeT, typename FnType>
-  //  Status PopulateInternal(const FnType& generator, bool parallel);
+  //  absl::Status PopulateInternal(const FnType& generator, bool parallel);
   template <typename NativeT>
-  Status PopulateInternal(
+  absl::Status PopulateInternal(
       absl::FunctionRef<NativeT(absl::Span<const int64_t>, int)> generator,
       bool parallel);
   void PopulateInplaceInternal(
@@ -1356,9 +1358,9 @@ class Literal : public MutableLiteralBase {
   // allocated in this literal for the subshape at dest_shape_index is
   // deallocated, and the respective buffers are replaced with those in
   // src_literal. Upon return, src_literal is set to a nil shape (empty tuple).
-  virtual Status MoveFrom(Literal&& src_literal,
-                          const ShapeIndex& dest_shape_index);
-  Status MoveFrom(Literal&& src_literal) {
+  virtual absl::Status MoveFrom(Literal&& src_literal,
+                                const ShapeIndex& dest_shape_index);
+  absl::Status MoveFrom(Literal&& src_literal) {
     return MoveFrom(std::move(src_literal), /*dest_shape_index=*/{});
   }
 
@@ -1546,8 +1548,8 @@ bool LiteralBase::Piece::DeserializeData(
 //   byte (PRED, S4, U4) are packed into bytes.  Otherwise, they are written in
 //   little-endian byte order.
 template <typename OutputIterator>
-Status LiteralBase::SerializeWithShapeProto(const ShapeProto& shape_proto,
-                                            OutputIterator output) const {
+absl::Status LiteralBase::SerializeWithShapeProto(const ShapeProto& shape_proto,
+                                                  OutputIterator output) const {
   SerializeState<OutputIterator> state(shape_proto, output);
   TF_RETURN_IF_ERROR(root_piece().ForEachSubpieceWithStatus(
       [&](const ShapeIndex& shape_index, const Piece& piece) -> absl::Status {
@@ -1866,7 +1868,7 @@ void MutableLiteralBase::PopulateR4FromArray4D(const Array4D<NativeT>& values) {
 }
 
 template <typename NativeT>
-TF_ATTRIBUTE_NOINLINE Status MutableLiteralBase::PopulateInternal(
+TF_ATTRIBUTE_NOINLINE absl::Status MutableLiteralBase::PopulateInternal(
     absl::FunctionRef<NativeT(absl::Span<const int64_t>, int)> generator,
     bool parallel) {
   const Shape& this_shape = shape();
@@ -1887,7 +1889,7 @@ TF_ATTRIBUTE_NOINLINE Status MutableLiteralBase::PopulateInternal(
 }
 
 template <typename NativeT>
-TF_ATTRIBUTE_NOINLINE Status MutableLiteralBase::Populate(
+TF_ATTRIBUTE_NOINLINE absl::Status MutableLiteralBase::Populate(
     absl::FunctionRef<NativeT(absl::Span<const int64_t>)> generator) {
   TF_RET_CHECK(LayoutUtil::IsDenseArray(shape()))
       << __func__ << " is only supported for dense arrays: " << shape();
@@ -1898,7 +1900,7 @@ TF_ATTRIBUTE_NOINLINE Status MutableLiteralBase::Populate(
       /*parallel=*/false);
 }
 template <typename NativeT>
-TF_ATTRIBUTE_NOINLINE Status MutableLiteralBase::PopulateParallel(
+TF_ATTRIBUTE_NOINLINE absl::Status MutableLiteralBase::PopulateParallel(
     absl::FunctionRef<NativeT(absl::Span<const int64_t>, int)> generator) {
   TF_RET_CHECK(LayoutUtil::IsDenseArray(shape()))
       << __func__ << " is only supported for dense arrays: " << shape();
diff --git a/third_party/xla/xla/literal_comparison.cc b/third_party/xla/xla/literal_comparison.cc
index 4da0d71d1ff772..6360b0f033405e 100644
--- a/third_party/xla/xla/literal_comparison.cc
+++ b/third_party/xla/xla/literal_comparison.cc
@@ -92,8 +92,8 @@ bool CompareEqual(NativeT lhs, NativeT rhs,
 }
 
 template <typename NativeT, typename UnsignedT>
-Status MakeBitwiseErrorStatus(NativeT lhs, NativeT rhs,
-                              absl::Span<const int64_t> multi_index) {
+absl::Status MakeBitwiseErrorStatus(NativeT lhs, NativeT rhs,
+                                    absl::Span<const int64_t> multi_index) {
   auto ulhs = Eigen::numext::bit_cast<UnsignedT>(lhs);
   auto urhs = Eigen::numext::bit_cast<UnsignedT>(rhs);
   auto lhs_double = static_cast<double>(lhs);
@@ -107,8 +107,8 @@ Status MakeBitwiseErrorStatus(NativeT lhs, NativeT rhs,
 }
 
 template <typename NativeT>
-Status MakeErrorStatus(NativeT lhs, NativeT rhs,
-                       absl::Span<const int64_t> multi_index) {
+absl::Status MakeErrorStatus(NativeT lhs, NativeT rhs,
+                             absl::Span<const int64_t> multi_index) {
   if constexpr (is_specialized_integral_v<NativeT>) {
     return InvalidArgument(
         "first mismatch at array index %s:\n  expected value: %s\n  actual "
@@ -137,9 +137,9 @@ Status MakeErrorStatus(NativeT lhs, NativeT rhs,
 //    - nullptr. In this case, the function will return once any mismatch is
 //      found between expected and actual.
 template <typename NativeT>
-Status Equal(LiteralSlice expected, LiteralSlice actual,
-             absl::Span<int64_t> multi_index, int64_t dimension,
-             Literal* mismatched = nullptr) {
+absl::Status Equal(LiteralSlice expected, LiteralSlice actual,
+                   absl::Span<int64_t> multi_index, int64_t dimension,
+                   Literal* mismatched = nullptr) {
   if (dimension == expected.shape().dimensions_size()) {
     NativeT expected_value = expected.Get<NativeT>(multi_index);
     NativeT actual_value = actual.Get<NativeT>(multi_index);
@@ -153,7 +153,7 @@ Status Equal(LiteralSlice expected, LiteralSlice actual,
                                              multi_index);
   }
 
-  Status result;
+  absl::Status result;
   int64_t upper_bound = expected.shape().dimensions(dimension);
   if (expected.shape().is_dynamic_dimension(dimension)) {
     // If the dimension is dynamic, we only want to check up until the actual
@@ -233,11 +233,11 @@ class NearComparator {
   // result. The comparison is ok() if all actual and expected elements are
   // within the given error bound. In case of error, the status contains a
   // detailed message about the discrepancy.
-  static Status Compare(const LiteralSlice& expected,
-                        const LiteralSlice& actual,
-                        const ShapeIndex& shape_index, ErrorSpec error,
-                        bool detailed_message,
-                        const MiscompareCallback& miscompare_callback) {
+  static absl::Status Compare(const LiteralSlice& expected,
+                              const LiteralSlice& actual,
+                              const ShapeIndex& shape_index, ErrorSpec error,
+                              bool detailed_message,
+                              const MiscompareCallback& miscompare_callback) {
     NearComparator<NativeT> comparator(expected, actual, shape_index, error,
                                        detailed_message, miscompare_callback);
     return comparator.Run();
@@ -285,7 +285,7 @@ class NearComparator {
         rel_error_buckets_(kErrorBucketBounds.size(), 0) {}
 
   // Runs the comparison between expected and actual literals.
-  Status Run() {
+  absl::Status Run() {
     // If the shapes mismatch, we simply fail the expectation instead of
     // printing out data, as it's a type error rather than a value error.
     TF_RETURN_IF_ERROR(EqualShapes(expected_.shape(), actual_.shape()));
@@ -627,21 +627,22 @@ class NearComparator {
   std::vector<int64_t> rel_error_buckets_;
 };
 
-Status EqualHelper(const LiteralSlice& expected, const LiteralSlice& actual,
-                   const ShapeIndex& shape_index,
-                   const MiscompareCallback& miscompare_callback) {
+absl::Status EqualHelper(const LiteralSlice& expected,
+                         const LiteralSlice& actual,
+                         const ShapeIndex& shape_index,
+                         const MiscompareCallback& miscompare_callback) {
   if (expected.shape().is_static() && actual.shape().is_static()) {
     TF_RETURN_IF_ERROR(EqualShapes(expected.shape(), actual.shape()));
   } else {
     TF_RETURN_IF_ERROR(EqualDynamicShapesAndDimensions(expected, actual));
   }
 
-  Status result;
+  absl::Status result;
   if (expected.shape().IsTuple()) {
     ShapeIndex next_index = shape_index;
     for (int i = 0; i < ShapeUtil::TupleElementCount(expected.shape()); ++i) {
       next_index.push_back(i);
-      Status tuple_result =
+      absl::Status tuple_result =
           EqualHelper(LiteralSlice(expected, {i}), LiteralSlice(actual, {i}),
                       next_index, miscompare_callback);
       if (miscompare_callback) {
@@ -691,10 +692,11 @@ Status EqualHelper(const LiteralSlice& expected, const LiteralSlice& actual,
 // Helper function for comparing two literals for nearness. Handles tuple-shapes
 // via recursion. shape_index is the ShapeIndex of expected (or actual)
 // currently being compared.
-Status NearHelper(const LiteralSlice& expected, const LiteralSlice& actual,
-                  const ShapeIndex& shape_index, const ErrorSpec& error,
-                  std::optional<bool> detailed_message,
-                  const MiscompareCallback& miscompare_callback) {
+absl::Status NearHelper(const LiteralSlice& expected,
+                        const LiteralSlice& actual,
+                        const ShapeIndex& shape_index, const ErrorSpec& error,
+                        std::optional<bool> detailed_message,
+                        const MiscompareCallback& miscompare_callback) {
   if (expected.shape().is_static() && actual.shape().is_static()) {
     TF_RETURN_IF_ERROR(EqualShapes(expected.shape(), actual.shape()));
   } else {
@@ -702,14 +704,14 @@ Status NearHelper(const LiteralSlice& expected, const LiteralSlice& actual,
   }
 
   if (expected.shape().IsTuple()) {
-    Status return_status;
+    absl::Status return_status;
     for (int64_t i = 0; i < ShapeUtil::TupleElementCount(expected.shape());
          ++i) {
       const auto expected_element = LiteralSlice(expected, {i});
       const auto actual_element = LiteralSlice(actual, {i});
       ShapeIndex element_index = shape_index;
       element_index.push_back(i);
-      Status element_result =
+      absl::Status element_result =
           NearHelper(expected_element, actual_element, element_index, error,
                      detailed_message, miscompare_callback);
       if (!element_result.ok()) {
@@ -740,7 +742,7 @@ Status NearHelper(const LiteralSlice& expected, const LiteralSlice& actual,
     bool use_detailed_message = detailed_message.value_or(
         ShapeUtil::ElementsIn(expected.shape()) >= 64);
     return primitive_util::PrimitiveTypeSwitch<Status>(
-        [&](auto primitive_type) -> Status {
+        [&](auto primitive_type) -> absl::Status {
           if constexpr (primitive_util::IsFloatingPointType(primitive_type) ||
                         primitive_util::IsComplexType(primitive_type)) {
             using NativeT = primitive_util::NativeTypeOf<primitive_type>;
@@ -761,7 +763,7 @@ Status NearHelper(const LiteralSlice& expected, const LiteralSlice& actual,
 
 }  // namespace
 
-Status EqualShapes(const Shape& expected, const Shape& actual) {
+absl::Status EqualShapes(const Shape& expected, const Shape& actual) {
   if (expected.element_type() != actual.element_type()) {
     return InvalidArgument("element type mismatch, want: %s got %s",
                            ShapeUtil::HumanString(expected),
@@ -776,7 +778,7 @@ Status EqualShapes(const Shape& expected, const Shape& actual) {
           ShapeUtil::TupleElementCount(actual));
     }
     for (int i = 0; i < expected.tuple_shapes_size(); ++i) {
-      Status result =
+      absl::Status result =
           EqualShapes(expected.tuple_shapes(i), actual.tuple_shapes(i));
       if (!result.ok()) {
         return AppendStatus(result, StrCat("mismatch in tuple index", i));
@@ -810,8 +812,8 @@ Status EqualShapes(const Shape& expected, const Shape& actual) {
   return OkStatus();
 }
 
-Status EqualDynamicShapesAndDimensions(const LiteralSlice& expected,
-                                       const LiteralSlice& actual) {
+absl::Status EqualDynamicShapesAndDimensions(const LiteralSlice& expected,
+                                             const LiteralSlice& actual) {
   TF_RETURN_IF_ERROR(EqualShapes(expected.shape(), actual.shape()));
   return ShapeUtil::ForEachSubshapeWithStatus(
       expected.shape(),
@@ -862,9 +864,9 @@ namespace {
 
 // If result is an error, extend the error message with the expected and actual
 // literals.
-Status EmitLiteralsInErrorMessage(const Status& result,
-                                  const LiteralSlice& expected,
-                                  const LiteralSlice& actual) {
+absl::Status EmitLiteralsInErrorMessage(const absl::Status& result,
+                                        const LiteralSlice& expected,
+                                        const LiteralSlice& actual) {
   if (result.ok()) {
     return result;
   }
@@ -875,24 +877,24 @@ Status EmitLiteralsInErrorMessage(const Status& result,
 
 }  // namespace
 
-Status Equal(const LiteralSlice& expected, const LiteralSlice& actual) {
+absl::Status Equal(const LiteralSlice& expected, const LiteralSlice& actual) {
   VLOG(1) << "expected:";
   XLA_VLOG_LINES(1, expected.ToString());
   VLOG(1) << "actual:";
   XLA_VLOG_LINES(1, actual.ToString());
-  Status result = EqualHelper(expected, actual, {}, nullptr);
+  absl::Status result = EqualHelper(expected, actual, {}, nullptr);
   return EmitLiteralsInErrorMessage(result, expected, actual);
 }
 
-Status Near(const LiteralSlice& expected, const LiteralSlice& actual,
-            const ErrorSpec& error, std::optional<bool> detailed_message,
-            const MiscompareCallback& miscompare_callback) {
+absl::Status Near(const LiteralSlice& expected, const LiteralSlice& actual,
+                  const ErrorSpec& error, std::optional<bool> detailed_message,
+                  const MiscompareCallback& miscompare_callback) {
   VLOG(1) << "Expected literal:";
   XLA_VLOG_LINES(1, expected.ToString());
   VLOG(1) << "Actual literal:";
   XLA_VLOG_LINES(1, actual.ToString());
-  Status result = NearHelper(expected, actual, /*shape_index=*/{}, error,
-                             detailed_message, miscompare_callback);
+  absl::Status result = NearHelper(expected, actual, /*shape_index=*/{}, error,
+                                   detailed_message, miscompare_callback);
   return EmitLiteralsInErrorMessage(result, expected, actual);
 }
 
diff --git a/third_party/xla/xla/literal_comparison.h b/third_party/xla/xla/literal_comparison.h
index 06d416f89dc4c8..61f0695c981e0c 100644
--- a/third_party/xla/xla/literal_comparison.h
+++ b/third_party/xla/xla/literal_comparison.h
@@ -36,17 +36,17 @@ namespace literal_comparison {
 
 // Returns ok if the given shapes have the same rank, dimension sizes, and
 // primitive types.
-Status EqualShapes(const Shape& expected, const Shape& actual);
+absl::Status EqualShapes(const Shape& expected, const Shape& actual);
 
 // Returns ok if the given literals share identical dynamic shapes and
 // dimension sizes.
-Status EqualDynamicShapesAndDimensions(const LiteralSlice& expected,
-                                       const LiteralSlice& actual);
+absl::Status EqualDynamicShapesAndDimensions(const LiteralSlice& expected,
+                                             const LiteralSlice& actual);
 
 // Returns ok if the expected and actual literals are (bitwise) equal for all
 // elements in the literal. Also, asserts that the rank, dimensions sizes, and
 // primitive type are equal.
-Status Equal(const LiteralSlice& expected, const LiteralSlice& actual);
+absl::Status Equal(const LiteralSlice& expected, const LiteralSlice& actual);
 
 // Structure that contains the distribution of absolute and relative errors,
 // bucketized into five buckets: [0.0001, 0.001, 0.01, 0.1, 1].
@@ -88,9 +88,9 @@ using MiscompareCallback = std::function<void(
 //
 // If miscompare_callback is nullptr, Near will return an error on the first
 // detected mismatch.
-Status Near(const LiteralSlice& expected, const LiteralSlice& actual,
-            const ErrorSpec& error, std::optional<bool> detailed_message,
-            const MiscompareCallback& miscompare_callback);
+absl::Status Near(const LiteralSlice& expected, const LiteralSlice& actual,
+                  const ErrorSpec& error, std::optional<bool> detailed_message,
+                  const MiscompareCallback& miscompare_callback);
 
 // Calling ToString on a literal with over 100 million elements takes around
 // 3 minutes.  The utility of printing a literal with >1000 elements is
diff --git a/third_party/xla/xla/literal_test.cc b/third_party/xla/xla/literal_test.cc
index 3701bf8b627320..37d0868334f210 100644
--- a/third_party/xla/xla/literal_test.cc
+++ b/third_party/xla/xla/literal_test.cc
@@ -1410,7 +1410,7 @@ TEST_F(LiteralUtilTest, CopyBetweenSameTuple) {
 TEST_F(LiteralUtilTest, CopyFromDifferentShapes) {
   auto matrix = LiteralUtil::CreateR2<float>({{1.0, 2.0}, {3.0, 4.0}});
   auto vector = LiteralUtil::CreateR1<float>({5.0, 7.0});
-  Status status = matrix.CopyFrom(vector);
+  absl::Status status = matrix.CopyFrom(vector);
   ASSERT_FALSE(status.ok());
   EXPECT_THAT(status.message(), HasSubstr("Destination subshape incompatible"));
 }
@@ -1819,7 +1819,7 @@ TEST_F(LiteralUtilTest, BitcastConvert) {
 
 TEST_F(LiteralUtilTest, BitcastConvertBetweenInvalidTypes) {
   Literal literal = LiteralUtil::CreateR0<uint32_t>(1234);
-  Status status =
+  absl::Status status =
       literal.BitcastConvert(ShapeUtil::ChangeElementType(literal.shape(), F64))
           .status();
   EXPECT_NE(OkStatus(), status);
@@ -2287,7 +2287,7 @@ TEST_F(LiteralUtilTest, InvalidProtoNoValues) {
   // Proto contains a shape, but no values.
   LiteralProto proto;
   *proto.mutable_shape() = ShapeUtil::MakeShape(F32, {3}).ToProto();
-  Status status = Literal::CreateFromProto(proto).status();
+  absl::Status status = Literal::CreateFromProto(proto).status();
   ASSERT_FALSE(status.ok());
   EXPECT_THAT(status.message(),
               HasSubstr("Expected 3 elements in LiteralProto"));
@@ -2297,7 +2297,7 @@ TEST_F(LiteralUtilTest, ValidProtoNoValues) {
   // Proto contains a shape, but no values.
   LiteralProto proto;
   *proto.mutable_shape() = ShapeUtil::MakeShape(F32, {3}).ToProto();
-  Status status =
+  absl::Status status =
       Literal::CreateFromProto(proto, /*prohibit_empty_literal=*/false)
           .status();
   EXPECT_TRUE(status.ok());
@@ -2311,7 +2311,7 @@ TEST_F(LiteralUtilTest, ValidProtoWithClearedValues) {
   // Clear values.
   proto.clear_preds();
   EXPECT_EQ(proto.preds_size(), 0);
-  Status status =
+  absl::Status status =
       Literal::CreateFromProto(proto, /*prohibit_empty_literal=*/false)
           .status();
   EXPECT_TRUE(status.ok());
@@ -2323,7 +2323,7 @@ TEST_F(LiteralUtilTest, InvalidProtoNoShape) {
   proto.add_preds(false);
   proto.add_preds(true);
   proto.add_preds(false);
-  Status status = Literal::CreateFromProto(proto).status();
+  absl::Status status = Literal::CreateFromProto(proto).status();
   ASSERT_FALSE(status.ok());
   EXPECT_THAT(status.message(), HasSubstr("LiteralProto has no shape"));
 }
@@ -2335,7 +2335,7 @@ TEST_F(LiteralUtilTest, InvalidProtoWrongContainer) {
   proto.add_preds(false);
   proto.add_preds(true);
   proto.add_preds(false);
-  Status status = Literal::CreateFromProto(proto).status();
+  absl::Status status = Literal::CreateFromProto(proto).status();
   ASSERT_FALSE(status.ok());
   EXPECT_THAT(status.message(),
               HasSubstr("Expected 3 elements in LiteralProto"));
@@ -2348,7 +2348,7 @@ TEST_F(LiteralUtilTest, InvalidProtoTooFewValues) {
   proto.add_f32s(1.0);
   proto.add_f32s(2.0);
   proto.add_f32s(3.0);
-  Status status = Literal::CreateFromProto(proto).status();
+  absl::Status status = Literal::CreateFromProto(proto).status();
   ASSERT_FALSE(status.ok());
   EXPECT_THAT(status.message(),
               HasSubstr("Expected 84 elements in LiteralProto"));
@@ -2361,7 +2361,7 @@ TEST_F(LiteralUtilTest, InvalidProtoTooManyValues) {
   proto.add_s32s(42);
   proto.add_s32s(-10);
   proto.add_s32s(100);
-  Status status = Literal::CreateFromProto(proto).status();
+  absl::Status status = Literal::CreateFromProto(proto).status();
   ASSERT_FALSE(status.ok());
   EXPECT_THAT(status.message(),
               HasSubstr("Expected 2 elements in LiteralProto"));
@@ -2376,7 +2376,7 @@ TEST_F(LiteralUtilTest, InvalidProtoMissingLayout) {
   proto.add_preds(false);
   proto.add_preds(true);
   proto.add_preds(false);
-  Status status = Literal::CreateFromProto(proto).status();
+  absl::Status status = Literal::CreateFromProto(proto).status();
   ASSERT_FALSE(status.ok());
   EXPECT_THAT(status.message(), HasSubstr("LiteralProto has no layout"));
 }
@@ -2394,7 +2394,7 @@ TEST_F(LiteralUtilTest, InvalidProtoTooFewTupleElements) {
   element0->add_preds(false);
   element0->add_preds(true);
 
-  Status status = Literal::CreateFromProto(proto).status();
+  absl::Status status = Literal::CreateFromProto(proto).status();
   ASSERT_FALSE(status.ok());
   EXPECT_THAT(status.message(), HasSubstr("Expected 2 tuple elements"));
 }
@@ -2419,7 +2419,7 @@ TEST_F(LiteralUtilTest, InvalidProtoTooManyTupleElements) {
   *element2->mutable_shape() = ShapeUtil::MakeShape(F32, {}).ToProto();
   element2->add_f32s(123.0);
 
-  Status status = Literal::CreateFromProto(proto).status();
+  absl::Status status = Literal::CreateFromProto(proto).status();
   ASSERT_FALSE(status.ok());
   EXPECT_THAT(status.message(), HasSubstr("Expected 2 tuple elements"));
 }
diff --git a/third_party/xla/xla/literal_util.h b/third_party/xla/xla/literal_util.h
index 5a27f1b510088a..fad4b5ed9f2826 100644
--- a/third_party/xla/xla/literal_util.h
+++ b/third_party/xla/xla/literal_util.h
@@ -251,7 +251,7 @@ class LiteralUtil {
 
   // Creates a literal with the supplied shape, and uses the provided value
   // generator to populate the literal's values.
-  // Returns the new literal object, or an error Status if failed.
+  // Returns the new literal object, or an error absl::Status if failed.
   template <PrimitiveType type, typename T = primitive_util::NativeTypeOf<type>>
   static absl::StatusOr<Literal> CreateLiteralWithGenerator(
       const Shape& shape,
@@ -260,7 +260,7 @@ class LiteralUtil {
   // Creates a literal with the supplied shape, and initializes the literal
   // values using a normal distribution with given mean and stddev standard
   // deviation, and using the engine as entropy generator.
-  // Returns the new literal object, or an error Status if failed.
+  // Returns the new literal object, or an error absl::Status if failed.
   template <PrimitiveType type, typename E,
             typename T = primitive_util::NativeTypeOf<type>>
   static absl::StatusOr<Literal> CreateRandomLiteral(const Shape& shape,
@@ -270,7 +270,7 @@ class LiteralUtil {
   // Creates a literal with the supplied shape, and initializes the literal
   // values using a normal distribution with given mean and stddev standard
   // deviation.
-  // Returns the new literal object, or an error Status if failed.
+  // Returns the new literal object, or an error absl::Status if failed.
   template <PrimitiveType type, typename T = primitive_util::NativeTypeOf<type>>
   static absl::StatusOr<Literal> CreateRandomLiteral(const Shape& shape, T mean,
                                                      T stddev);
diff --git a/third_party/xla/xla/map_util.h b/third_party/xla/xla/map_util.h
index a47c661f8f6798..f25caf0874d408 100644
--- a/third_party/xla/xla/map_util.h
+++ b/third_party/xla/xla/map_util.h
@@ -54,7 +54,7 @@ typename Collection::value_type::second_type& FindOrDie(
 // Like FindOrDie but returns an error instead of dying if `key` is not in
 // `container`.
 template <class Collection>
-StatusOr<
+absl::StatusOr<
     std::reference_wrapper<const typename Collection::value_type::second_type>>
 MaybeFind(const Collection& collection,
           const typename Collection::value_type::first_type& key) {
diff --git a/third_party/xla/xla/protobuf_util.cc b/third_party/xla/xla/protobuf_util.cc
index f9939110dd81a6..d465270316569c 100644
--- a/third_party/xla/xla/protobuf_util.cc
+++ b/third_party/xla/xla/protobuf_util.cc
@@ -49,10 +49,10 @@ size_t ProtobufHash(const tsl::protobuf::Message& m) {
   return absl::HashOf(serialized);
 }
 
-Status DumpProtoToDirectory(const tsl::protobuf::Message& message,
-                            const std::string& directory,
-                            const std::string& file_name,
-                            std::string* full_path) {
+absl::Status DumpProtoToDirectory(const tsl::protobuf::Message& message,
+                                  const std::string& directory,
+                                  const std::string& file_name,
+                                  std::string* full_path) {
   tsl::Env* env = tsl::Env::Default();
   TF_RETURN_IF_ERROR(env->RecursivelyCreateDir(directory));
   std::string safe_file_name = SanitizeFileName(file_name) + ".pb";
diff --git a/third_party/xla/xla/protobuf_util.h b/third_party/xla/xla/protobuf_util.h
index e739353cf826be..031596083df361 100644
--- a/third_party/xla/xla/protobuf_util.h
+++ b/third_party/xla/xla/protobuf_util.h
@@ -62,10 +62,10 @@ class ProtobufHashWrapper {
 //
 // If 'full_name' is not null then it is set to the name of the file the
 // protobuf was written to.
-Status DumpProtoToDirectory(const tsl::protobuf::Message& message,
-                            const std::string& directory,
-                            const std::string& file_name,
-                            std::string* full_path = nullptr);
+absl::Status DumpProtoToDirectory(const tsl::protobuf::Message& message,
+                                  const std::string& directory,
+                                  const std::string& file_name,
+                                  std::string* full_path = nullptr);
 
 // Registers a function that may either expand a dirpath or forward the original
 // dirpath along as-is.
diff --git a/third_party/xla/xla/service_interface.h b/third_party/xla/xla/service_interface.h
index c1bf58dc5c2a23..a3c974600d0afa 100644
--- a/third_party/xla/xla/service_interface.h
+++ b/third_party/xla/xla/service_interface.h
@@ -31,53 +31,57 @@ class ServiceInterface {
   ServiceInterface() {}
   virtual ~ServiceInterface() = default;
 
-  // TODO(b/31824348): Convert to use StatusOr.
-  virtual Status TransferToClient(const TransferToClientRequest* arg,
-                                  TransferToClientResponse* result) = 0;
+  // TODO(b/31824348): Convert to use absl::StatusOr.
+  virtual absl::Status TransferToClient(const TransferToClientRequest* arg,
+                                        TransferToClientResponse* result) = 0;
 
-  virtual Status TransferToServer(const TransferToServerRequest* arg,
-                                  TransferToServerResponse* result) = 0;
+  virtual absl::Status TransferToServer(const TransferToServerRequest* arg,
+                                        TransferToServerResponse* result) = 0;
 
-  virtual Status TransferToInfeed(const TransferToInfeedRequest* arg,
-                                  TransferToInfeedResponse* result) = 0;
+  virtual absl::Status TransferToInfeed(const TransferToInfeedRequest* arg,
+                                        TransferToInfeedResponse* result) = 0;
 
-  virtual Status TransferFromOutfeed(const TransferFromOutfeedRequest* arg,
-                                     TransferFromOutfeedResponse* result) = 0;
+  virtual absl::Status TransferFromOutfeed(
+      const TransferFromOutfeedRequest* arg,
+      TransferFromOutfeedResponse* result) = 0;
 
-  virtual Status ResetDevice(const ResetDeviceRequest* arg,
-                             ResetDeviceResponse* result) = 0;
+  virtual absl::Status ResetDevice(const ResetDeviceRequest* arg,
+                                   ResetDeviceResponse* result) = 0;
 
-  virtual Status Compile(const CompileRequest* arg,
-                         CompileResponse* result) = 0;
+  virtual absl::Status Compile(const CompileRequest* arg,
+                               CompileResponse* result) = 0;
 
-  virtual Status Execute(const ExecuteRequest* arg,
-                         ExecuteResponse* result) = 0;
+  virtual absl::Status Execute(const ExecuteRequest* arg,
+                               ExecuteResponse* result) = 0;
 
-  virtual Status ExecuteGraphParallel(const ExecuteGraphParallelRequest* arg,
-                                      ExecuteParallelResponse* result) = 0;
+  virtual absl::Status ExecuteGraphParallel(
+      const ExecuteGraphParallelRequest* arg,
+      ExecuteParallelResponse* result) = 0;
 
-  virtual Status DeconstructTuple(const DeconstructTupleRequest* arg,
-                                  DeconstructTupleResponse* result) = 0;
+  virtual absl::Status DeconstructTuple(const DeconstructTupleRequest* arg,
+                                        DeconstructTupleResponse* result) = 0;
 
-  virtual Status GetComputationGraphStats(
+  virtual absl::Status GetComputationGraphStats(
       const ComputationGraphStatsRequest* arg,
       ComputationStatsResponse* result) = 0;
 
-  virtual Status GetShape(const GetShapeRequest* arg,
-                          GetShapeResponse* result) = 0;
+  virtual absl::Status GetShape(const GetShapeRequest* arg,
+                                GetShapeResponse* result) = 0;
 
-  virtual Status CreateChannelHandle(const CreateChannelHandleRequest* arg,
-                                     CreateChannelHandleResponse* result) = 0;
+  virtual absl::Status CreateChannelHandle(
+      const CreateChannelHandleRequest* arg,
+      CreateChannelHandleResponse* result) = 0;
 
-  virtual Status GetDeviceHandles(const GetDeviceHandlesRequest* arg,
-                                  GetDeviceHandlesResponse* result) = 0;
+  virtual absl::Status GetDeviceHandles(const GetDeviceHandlesRequest* arg,
+                                        GetDeviceHandlesResponse* result) = 0;
 
-  virtual Status ComputeConstantGraph(const ComputeConstantGraphRequest* arg,
-                                      ComputeConstantResponse* result) = 0;
+  virtual absl::Status ComputeConstantGraph(
+      const ComputeConstantGraphRequest* arg,
+      ComputeConstantResponse* result) = 0;
 
   // Methods used by GlobalData.
-  virtual Status Unregister(const UnregisterRequest* arg,
-                            UnregisterResponse* result) = 0;
+  virtual absl::Status Unregister(const UnregisterRequest* arg,
+                                  UnregisterResponse* result) = 0;
 };
 
 }  // namespace xla
diff --git a/third_party/xla/xla/shape_layout.cc b/third_party/xla/xla/shape_layout.cc
index 1712bf789f370f..4ea3e43c9762f3 100644
--- a/third_party/xla/xla/shape_layout.cc
+++ b/third_party/xla/xla/shape_layout.cc
@@ -26,7 +26,7 @@ limitations under the License.
 
 namespace xla {
 
-Status ShapeLayout::CopyLayoutFromShape(const Shape& other_shape) {
+absl::Status ShapeLayout::CopyLayoutFromShape(const Shape& other_shape) {
   if (!ShapeUtil::Compatible(other_shape, shape_)) {
     return InvalidArgument("Shape %s is not compatible with shape %s",
                            ShapeUtil::HumanString(other_shape),
@@ -36,7 +36,7 @@ Status ShapeLayout::CopyLayoutFromShape(const Shape& other_shape) {
   return OkStatus();
 }
 
-Status ShapeLayout::AssignLayoutToShape(Shape* to_shape) const {
+absl::Status ShapeLayout::AssignLayoutToShape(Shape* to_shape) const {
   if (!ShapeUtil::Compatible(*to_shape, shape_)) {
     return InvalidArgument("Shape %s is not compatible with shape %s",
                            ShapeUtil::HumanString(*to_shape),
diff --git a/third_party/xla/xla/shape_layout.h b/third_party/xla/xla/shape_layout.h
index 9d6353e27f3a61..6262acb2851a1c 100644
--- a/third_party/xla/xla/shape_layout.h
+++ b/third_party/xla/xla/shape_layout.h
@@ -41,7 +41,7 @@ class ShapeLayout {
   // Assigns the layouts in this ShapeLayout to the Layout fields of the given
   // shape. 'to_shape' and the shape of the ShapeLayout object must be
   // compatible.
-  Status AssignLayoutToShape(Shape* to_shape) const;
+  absl::Status AssignLayoutToShape(Shape* to_shape) const;
 
   // Returns true if the Layouts in this ShapeLayout match the layouts in the
   // given shape. Returns false otherwise. If the given shape is not compatible
@@ -54,7 +54,7 @@ class ShapeLayout {
 
   // Copies the layout from the given shape into this ShapeLayout. 'other_shape'
   // must be compatible with the ShapeLayout's shape.
-  Status CopyLayoutFromShape(const Shape& other_shape);
+  absl::Status CopyLayoutFromShape(const Shape& other_shape);
 
   // Clears (Layout::Clear) all the Layouts stored in this object.
   void Clear();
diff --git a/third_party/xla/xla/shape_tree.h b/third_party/xla/xla/shape_tree.h
index 3daacaedbd875d..8be8e906318110 100644
--- a/third_party/xla/xla/shape_tree.h
+++ b/third_party/xla/xla/shape_tree.h
@@ -229,18 +229,19 @@ class ShapeTree {
     }
   }
 
-  // Like ForEach(Mutable)Element, but the callable returns a Status instead of
-  // void.  The first non-OK return value is returned by the ForEach* function.
-  Status ForEachElementWithStatus(
-      absl::FunctionRef<Status(const ShapeIndex&, const T&)> func) const {
+  // Like ForEach(Mutable)Element, but the callable returns a absl::Status
+  // instead of void.  The first non-OK return value is returned by the ForEach*
+  // function.
+  absl::Status ForEachElementWithStatus(
+      absl::FunctionRef<absl::Status(const ShapeIndex&, const T&)> func) const {
     for (const Node& node : nodes_) {
       TF_RETURN_IF_ERROR(func(node.first, node.second));
     }
     return OkStatus();
   }
 
-  Status ForEachMutableElementWithStatus(
-      absl::FunctionRef<Status(const ShapeIndex&, T*)> func) {
+  absl::Status ForEachMutableElementWithStatus(
+      absl::FunctionRef<absl::Status(const ShapeIndex&, T*)> func) {
     for (Node& node : nodes_) {
       TF_RETURN_IF_ERROR(func(node.first, &node.second));
     }
@@ -263,16 +264,16 @@ class ShapeTree {
     }
   }
 
-  Status ForEachElementPostOrderWithStatus(
-      absl::FunctionRef<Status(const ShapeIndex&, const T&)> func) const {
+  absl::Status ForEachElementPostOrderWithStatus(
+      absl::FunctionRef<absl::Status(const ShapeIndex&, const T&)> func) const {
     for (auto node = nodes_.rbegin(); node != nodes_.rend(); ++node) {
       TF_RETURN_IF_ERROR(func(node->first, node->second));
     }
     return OkStatus();
   }
 
-  Status ForEachMutableElementPostOrderWithStatus(
-      absl::FunctionRef<Status(const ShapeIndex&, T*)> func) {
+  absl::Status ForEachMutableElementPostOrderWithStatus(
+      absl::FunctionRef<absl::Status(const ShapeIndex&, T*)> func) {
     for (auto node = nodes_.rbegin(); node != nodes_.rend(); ++node) {
       TF_RETURN_IF_ERROR(func(node->first, &node->second));
     }
@@ -340,7 +341,7 @@ class ShapeTree {
     });
   }
 
-  StatusOr<ShapeTree<T>> SubShapeTree(const ShapeIndex& index) const {
+  absl::StatusOr<ShapeTree<T>> SubShapeTree(const ShapeIndex& index) const {
     TF_ASSIGN_OR_RETURN(const Shape* sub_shape,
                         ShapeUtil::TryGetSubshape(shape(), index));
     size_t count = ShapeUtil::SubshapeCount(*sub_shape);
diff --git a/third_party/xla/xla/shape_util.cc b/third_party/xla/xla/shape_util.cc
index b89e84e2d17f7e..976c06ab3bde08 100644
--- a/third_party/xla/xla/shape_util.cc
+++ b/third_party/xla/xla/shape_util.cc
@@ -437,7 +437,7 @@ ShapeUtil::MakeShapeWithDescendingLayoutAndSamePhysicalLayout(
   return new_shape;
 }
 
-/* static */ Status ShapeUtil::PopulateShape(
+/* static */ absl::Status ShapeUtil::PopulateShape(
     PrimitiveType element_type, absl::Span<const int64_t> dimensions,
     Shape* shape) {
   shape->Clear();
@@ -956,7 +956,7 @@ Shape ShapeUtil::PrependMajorDimension(int64_t bound, Shape shape) {
   return size;
 }
 
-/* static */ Status ShapeUtil::ValidateShapeWithOptionalLayoutInternal(
+/* static */ absl::Status ShapeUtil::ValidateShapeWithOptionalLayoutInternal(
     const Shape& shape) {
   if (shape.element_type() == PRIMITIVE_TYPE_INVALID ||
       !PrimitiveType_IsValid(shape.element_type())) {
@@ -1021,7 +1021,7 @@ Shape ShapeUtil::PrependMajorDimension(int64_t bound, Shape shape) {
   return OkStatus();
 }
 
-/* static */ Status ShapeUtil::ValidateShapeSize(const Shape& shape) {
+/* static */ absl::Status ShapeUtil::ValidateShapeSize(const Shape& shape) {
   VLOG(3) << "Validating shape size: " << ShapeUtil::HumanString(shape);
 
   if (!shape.IsArray()) {
@@ -1042,7 +1042,7 @@ Shape ShapeUtil::PrependMajorDimension(int64_t bound, Shape shape) {
   return OkStatus();
 }
 
-/* static */ Status ShapeUtil::ValidateShapeWithOptionalLayout(
+/* static */ absl::Status ShapeUtil::ValidateShapeWithOptionalLayout(
     const Shape& shape) {
   TF_RETURN_IF_ERROR(ValidateShapeWithOptionalLayoutInternal(shape));
 
@@ -1050,7 +1050,7 @@ Shape ShapeUtil::PrependMajorDimension(int64_t bound, Shape shape) {
                                            /*allow_missing_layouts=*/true);
 }
 
-/* static */ Status ShapeUtil::ValidateShape(const Shape& shape) {
+/* static */ absl::Status ShapeUtil::ValidateShape(const Shape& shape) {
   TF_RETURN_IF_ERROR(ValidateShapeWithOptionalLayoutInternal(shape));
 
   return LayoutUtil::ValidateLayoutInShape(shape);
@@ -1796,7 +1796,7 @@ ShapeUtil::DecomposeBitcastToTrt(const Shape& input_shape,
   return compatible;
 }
 
-/* static */ Status ShapeUtil::ForEachIndexWithStatus(
+/* static */ absl::Status ShapeUtil::ForEachIndexWithStatus(
     const Shape& shape, absl::Span<const int64_t> base,
     absl::Span<const int64_t> count, absl::Span<const int64_t> incr,
     const ForEachVisitorFunction& visitor_function) {
@@ -1827,7 +1827,7 @@ ShapeUtil::DecomposeBitcastToTrt(const Shape& input_shape,
                                              visitor_function));
 }
 
-/* static */ Status ShapeUtil::ForEachIndexParallelWithStatus(
+/* static */ absl::Status ShapeUtil::ForEachIndexParallelWithStatus(
     const Shape& shape, absl::Span<const int64_t> base,
     absl::Span<const int64_t> count, absl::Span<const int64_t> incr,
     const ForEachParallelVisitorFunction& visitor_function) {
@@ -1842,7 +1842,7 @@ ShapeUtil::DecomposeBitcastToTrt(const Shape& input_shape,
   TF_CHECK_OK(ForEachIndexParallelWithStatus(shape, visitor_function));
 }
 
-/* static */ Status ShapeUtil::ForEachIndexParallelWithStatus(
+/* static */ absl::Status ShapeUtil::ForEachIndexParallelWithStatus(
     const Shape& shape,
     const ForEachParallelVisitorFunction& visitor_function) {
   std::vector<int64_t> base(shape.dimensions_size());
@@ -1852,7 +1852,7 @@ ShapeUtil::DecomposeBitcastToTrt(const Shape& input_shape,
                                         visitor_function);
 }
 
-/* static */ Status ShapeUtil::ForEachIndexInternal(
+/* static */ absl::Status ShapeUtil::ForEachIndexInternal(
     const Shape& shape, absl::Span<const int64_t> base,
     absl::Span<const int64_t> count, absl::Span<const int64_t> incr,
     const ForEachVisitorFunction& visitor_function) {
@@ -1913,13 +1913,13 @@ struct ParallelState {
 
   absl::Mutex mu;
   tsl::thread::ThreadPool* pool;
-  Status status;  // Guarded by mu
+  absl::Status status;  // Guarded by mu
   absl::BlockingCounter counter;
 };
 
 }  // anonymous namespace
 
-/* static */ Status ShapeUtil::ForEachIndexInternalParallel(
+/* static */ absl::Status ShapeUtil::ForEachIndexInternalParallel(
     const Shape& shape, absl::Span<const int64_t> base,
     absl::Span<const int64_t> count, absl::Span<const int64_t> incr,
     const ForEachParallelVisitorFunction& visitor_function) {
@@ -2135,7 +2135,8 @@ Shape ShapeUtil::DeviceShapeToHostShape(Shape s) {
 }
 
 /*static*/
-Status ShapeUtil::ByteStrides(const Shape& shape, absl::Span<int64_t> strides) {
+absl::Status ShapeUtil::ByteStrides(const Shape& shape,
+                                    absl::Span<int64_t> strides) {
   TF_RET_CHECK(shape.IsArray());
   TF_RET_CHECK(shape.has_layout());
   TF_RET_CHECK(shape.dimensions_size() == strides.size());
diff --git a/third_party/xla/xla/shape_util.h b/third_party/xla/xla/shape_util.h
index 548d39961b9ce4..d6688c573cdbc2 100644
--- a/third_party/xla/xla/shape_util.h
+++ b/third_party/xla/xla/shape_util.h
@@ -478,19 +478,19 @@ class ShapeUtil {
       const Shape& shape);
 
   // As MakeShape, but the object to write to is passed in.
-  static Status PopulateShape(PrimitiveType element_type,
-                              absl::Span<const int64_t> dimensions,
-                              Shape* shape);
+  static absl::Status PopulateShape(PrimitiveType element_type,
+                                    absl::Span<const int64_t> dimensions,
+                                    Shape* shape);
 
   // Validates that the provided shape satisfies invariants.
-  static Status ValidateShape(const Shape& shape);
+  static absl::Status ValidateShape(const Shape& shape);
 
   // Validates the provided shape satisfies invariants, except those that
   // pertain to layout.
   //
   // Layout is optional for client-provided shapes, so that the compiler may
   // determine and assign an optimized layout.
-  static Status ValidateShapeWithOptionalLayout(const Shape& shape);
+  static absl::Status ValidateShapeWithOptionalLayout(const Shape& shape);
 
   // Returns whether the element type of the shape is integral (signed or
   // unsigned). Note that predicates are not considered integral here, since
@@ -612,24 +612,25 @@ class ShapeUtil {
                     });
   }
 
-  // Variants of ForEach(Mutable)Subshape which propagate Status from the
+  // Variants of ForEach(Mutable)Subshape which propagate absl::Status from the
   // visitor function.
   //
   // Visitor function must have the signature
   //
-  //   Status fn(const Shape& subshape, const ShapeIndex& index), or
-  //   Status fn(Shape* subshape, const ShapeIndex& index) (mutable version)
+  //   absl::Status fn(const Shape& subshape, const ShapeIndex& index), or
+  //   absl::Status fn(Shape* subshape, const ShapeIndex& index) (mutable
+  //   version)
   //
   template <typename Fn>
-  static Status ForEachSubshapeWithStatus(const Shape& shape, Fn&& fn) {
+  static absl::Status ForEachSubshapeWithStatus(const Shape& shape, Fn&& fn) {
     return ForEachMutableSubshapeWithStatus(
         const_cast<Shape*>(&shape),
-        [&](Shape* subshape, const ShapeIndex& index) -> Status {
+        [&](Shape* subshape, const ShapeIndex& index) -> absl::Status {
           return fn(*const_cast<const Shape*>(subshape), index);
         });
   }
   template <typename Fn>
-  static Status ForEachMutableSubshapeWithStatus(Shape* shape, Fn&& fn) {
+  static absl::Status ForEachMutableSubshapeWithStatus(Shape* shape, Fn&& fn) {
     ShapeIndex index;
     return ForEachMutableSubshapeWithStatusHelper(shape, fn, &index);
   }
@@ -661,26 +662,27 @@ class ShapeUtil {
         .IgnoreError();
   }
 
-  // Variants of ForEach(Mutable)SubshapePostOrder which propagate Status from
-  // the visitor function.
+  // Variants of ForEach(Mutable)SubshapePostOrder which propagate absl::Status
+  // from the visitor function.
   //
   // Visitor function must have the signature
   //
-  //   Status fn(const Shape& subshape, const ShapeIndex& index), or
-  //   Status fn(Shape* subshape, const ShapeIndex& index) (mutable version)
+  //   absl::Status fn(const Shape& subshape, const ShapeIndex& index), or
+  //   absl::Status fn(Shape* subshape, const ShapeIndex& index) (mutable
+  //   version)
   //
   template <typename Fn>
-  static Status ForEachSubshapePostOrderWithStatus(const Shape& shape,
-                                                   Fn&& fn) {
+  static absl::Status ForEachSubshapePostOrderWithStatus(const Shape& shape,
+                                                         Fn&& fn) {
     return ForEachMutableSubshapePostOrderWithStatus(
         const_cast<Shape*>(&shape),
-        [&](Shape* subshape, const ShapeIndex& index) -> Status {
+        [&](Shape* subshape, const ShapeIndex& index) -> absl::Status {
           return fn(*const_cast<const Shape*>(subshape), index);
         });
   }
   template <typename Fn>
-  static Status ForEachMutableSubshapePostOrderWithStatus(Shape* shape,
-                                                          Fn&& fn) {
+  static absl::Status ForEachMutableSubshapePostOrderWithStatus(Shape* shape,
+                                                                Fn&& fn) {
     ShapeIndex index;
     return ForEachMutableSubshapePostOrderWithStatusHelper(shape, fn, &index);
   }
@@ -881,7 +883,7 @@ class ShapeUtil {
   // count (index[i] < base[i] + count[i]), and calls the visitor_function
   // with the current index. The visitor_function visitor function should
   // return true if it wants to continue, or false otherwise.
-  static Status ForEachIndexWithStatus(
+  static absl::Status ForEachIndexWithStatus(
       const Shape& shape, absl::Span<const int64_t> base,
       absl::Span<const int64_t> count, absl::Span<const int64_t> incr,
       const ForEachVisitorFunction& visitor_function);
@@ -894,7 +896,7 @@ class ShapeUtil {
   };
 
   template <typename FnTy>
-  static Status ForEachIndexWithStatus(
+  static absl::Status ForEachIndexWithStatus(
       const Shape& shape, const IndexIterationSpace& iteration_space,
       FnTy&& function) {
     return ShapeUtil::ForEachIndexWithStatus(
@@ -915,7 +917,7 @@ class ShapeUtil {
   // These convenience wrappers don't take `base`, `count` and `incr`
   // explicitly, but iterate over every element in `shape` instead.
 
-  static Status ForEachIndexWithStatus(
+  static absl::Status ForEachIndexWithStatus(
       const Shape& shape, const ForEachVisitorFunction& visitor_function) {
     std::vector<int64_t> base(shape.dimensions_size());
     std::vector<int64_t> incr(shape.dimensions_size(), 1);
@@ -959,7 +961,7 @@ class ShapeUtil {
   // Returns the number of threads in the threadpool of ForEachIndexParallel*.
   static int GetForEachIndexParallelThreadCount();
 
-  static Status ForEachIndexParallelWithStatus(
+  static absl::Status ForEachIndexParallelWithStatus(
       const Shape& shape, absl::Span<const int64_t> base,
       absl::Span<const int64_t> count, absl::Span<const int64_t> incr,
       const ForEachParallelVisitorFunction& visitor_function);
@@ -970,7 +972,7 @@ class ShapeUtil {
       const Shape& shape,
       const ForEachParallelVisitorFunction& visitor_function);
 
-  static Status ForEachIndexParallelWithStatus(
+  static absl::Status ForEachIndexParallelWithStatus(
       const Shape& shape,
       const ForEachParallelVisitorFunction& visitor_function);
 
@@ -1012,7 +1014,8 @@ class ShapeUtil {
   // Computes byte strides of an array shape `shape`. `shape` must have a
   // layout. Ignores tiling. `strides` must have size equal to the number of
   // dimensions of `shape`.
-  static Status ByteStrides(const Shape& shape, absl::Span<int64_t> strides);
+  static absl::Status ByteStrides(const Shape& shape,
+                                  absl::Span<int64_t> strides);
   // Same as above but returns the stride array, or std::nullopt if error.
   static std::optional<absl::InlinedVector<int64_t, 4>> ByteStrides(
       const Shape& shape);
@@ -1033,17 +1036,18 @@ class ShapeUtil {
 
   // Validates the shape size is sane. This makes sure it's safe to do
   // calculations in int64_t without overflowing.
-  static Status ValidateShapeSize(const Shape& shape);
+  static absl::Status ValidateShapeSize(const Shape& shape);
 
   // Validates all of the non-layout properties of the shape -- this is a helper
   // used by both the layout-optional and layout-required public method.
-  static Status ValidateShapeWithOptionalLayoutInternal(const Shape& shape);
+  static absl::Status ValidateShapeWithOptionalLayoutInternal(
+      const Shape& shape);
 
   // Helper for ForEachSubshape which visits the subshapes of the given shape in
   // DFS pre-order starting with the index.
   template <typename Fn>
-  static Status ForEachMutableSubshapeWithStatusHelper(Shape* shape, Fn&& fn,
-                                                       ShapeIndex* index) {
+  static absl::Status ForEachMutableSubshapeWithStatusHelper(
+      Shape* shape, Fn&& fn, ShapeIndex* index) {
     TF_RETURN_IF_ERROR(fn(shape, *index));
     if (shape->IsTuple()) {
       for (int64_t i = 0; i < ShapeUtil::TupleElementCount(*shape); ++i) {
@@ -1059,7 +1063,7 @@ class ShapeUtil {
   // Helper for ForEachSubshapePost which visits the subshapes of the given
   // shape in DFS post-order.
   template <typename Fn>
-  static Status ForEachMutableSubshapePostOrderWithStatusHelper(
+  static absl::Status ForEachMutableSubshapePostOrderWithStatusHelper(
       Shape* shape, Fn&& fn, ShapeIndex* index) {
     if (shape->IsTuple()) {
       for (int64_t i = 0; i < ShapeUtil::TupleElementCount(*shape); ++i) {
@@ -1100,7 +1104,7 @@ class ShapeUtil {
     int64_t CalculateNumSteps() const;
   };
 
-  static Status ForEachIndexInternal(
+  static absl::Status ForEachIndexInternal(
       const Shape& shape, absl::Span<const int64_t> base,
       absl::Span<const int64_t> count, absl::Span<const int64_t> incr,
       const ForEachVisitorFunction& visitor_function);
@@ -1110,7 +1114,7 @@ class ShapeUtil {
       absl::Span<const int64_t> count, absl::Span<const int64_t> incr,
       const ForEachVisitorFunctionNoStatus& visitor_function);
 
-  static Status ForEachIndexInternalParallel(
+  static absl::Status ForEachIndexInternalParallel(
       const Shape& shape, absl::Span<const int64_t> base,
       absl::Span<const int64_t> count, absl::Span<const int64_t> incr,
       const ForEachParallelVisitorFunction& visitor_function);
diff --git a/third_party/xla/xla/shape_util_test.cc b/third_party/xla/xla/shape_util_test.cc
index 7314a0c464f48d..6ebda8e53ed79a 100644
--- a/third_party/xla/xla/shape_util_test.cc
+++ b/third_party/xla/xla/shape_util_test.cc
@@ -641,7 +641,7 @@ TEST(ShapeUtilTest, ForEachIndexWithStatus) {
     return true;
   };
 
-  Status error_status = ShapeUtil::ForEachIndexWithStatus(
+  absl::Status error_status = ShapeUtil::ForEachIndexWithStatus(
       shape, /*base=*/{0, 0}, /*count=*/{10, 10}, /*incr=*/{0, 1},
       increment_func);
 
diff --git a/third_party/xla/xla/sharding_op_util.cc b/third_party/xla/xla/sharding_op_util.cc
index 3b78a18b01f415..5ef8d7d9b06d51 100644
--- a/third_party/xla/xla/sharding_op_util.cc
+++ b/third_party/xla/xla/sharding_op_util.cc
@@ -38,8 +38,8 @@ std::string EncodeAttributes(absl::Span<const int64_t> unspecified_dims) {
                       absl::StrJoin(unspecified_dims, ","), "]");
 }
 
-Status ParseAttributes(absl::string_view opaque,
-                       std::vector<int64_t>* unspecified_dims) {
+absl::Status ParseAttributes(absl::string_view opaque,
+                             std::vector<int64_t>* unspecified_dims) {
   HloLexer lexer(opaque);
   while (lexer.Lex() != TokKind::kEof) {
     if (lexer.GetKind() != TokKind::kAttributeName) {
diff --git a/third_party/xla/xla/sharding_op_util.h b/third_party/xla/xla/sharding_op_util.h
index be71cab8464ff3..8b59f69a270777 100644
--- a/third_party/xla/xla/sharding_op_util.h
+++ b/third_party/xla/xla/sharding_op_util.h
@@ -31,8 +31,8 @@ namespace sharding_op_util {
 std::string EncodeAttributes(absl::Span<const int64_t> unspecified_dims);
 
 // Parses the opaque string of Sharding and auto/manual conversion custom ops.
-Status ParseAttributes(absl::string_view opaque,
-                       std::vector<int64_t>* unspecified_dims);
+absl::Status ParseAttributes(absl::string_view opaque,
+                             std::vector<int64_t>* unspecified_dims);
 
 }  // namespace sharding_op_util
 }  // namespace xla
diff --git a/third_party/xla/xla/status_macros.cc b/third_party/xla/xla/status_macros.cc
index 97254a2df5d5fa..3c4b8a8e608aa5 100644
--- a/third_party/xla/xla/status_macros.cc
+++ b/third_party/xla/xla/status_macros.cc
@@ -39,7 +39,7 @@ ABSL_CONST_INIT const char kPossibleAutoJitAlternative[] =
 
 // Log the error at the given severity, optionally with a stack trace.
 // If log_severity is NUM_SEVERITIES, nothing is logged.
-static void LogError(const Status& status, const char* filename, int line,
+static void LogError(const absl::Status& status, const char* filename, int line,
                      int log_severity, bool should_log_stack_trace) {
   if (ABSL_PREDICT_TRUE(log_severity != tsl::NUM_SEVERITIES)) {
     std::string stack_trace;
@@ -67,20 +67,21 @@ static void LogError(const Status& status, const char* filename, int line,
   }
 }
 
-// Make a Status with a code, error message and payload,
+// Make a absl::Status with a code, error message and payload,
 // and also send it to LOG(<log_severity>) using the given filename
 // and line (unless should_log is false, or log_severity is
 // NUM_SEVERITIES).  If should_log_stack_trace is true, the stack
 // trace is included in the log message (ignored if should_log is
 // false).
-static Status MakeError(const char* filename, int line, absl::StatusCode code,
-                        const std::string& message, bool should_log,
-                        int log_severity, bool should_log_stack_trace) {
+static absl::Status MakeError(const char* filename, int line,
+                              absl::StatusCode code, const std::string& message,
+                              bool should_log, int log_severity,
+                              bool should_log_stack_trace) {
   if (ABSL_PREDICT_FALSE(code == absl::StatusCode::kOk)) {
     LOG(ERROR) << "Cannot create error with status OK";
     code = absl::StatusCode::kUnknown;
   }
-  const Status status = Status(code, message);
+  const absl::Status status = absl::Status(code, message);
   if (ABSL_PREDICT_TRUE(should_log)) {
     LogError(status, filename, line, log_severity, should_log_stack_trace);
   }
@@ -109,7 +110,7 @@ MakeErrorStream::Impl::Impl(const char* file, int line, tsl::error::Code code,
       should_log_stack_trace_(false),
       make_error_stream_with_output_wrapper_(error_stream) {}
 
-MakeErrorStream::Impl::Impl(const Status& status,
+MakeErrorStream::Impl::Impl(const absl::Status& status,
                             PriorMessageHandling prior_message_handling,
                             const char* file, int line,
                             MakeErrorStream* error_stream)
@@ -134,20 +135,20 @@ MakeErrorStream::Impl::~Impl() {
   // Note: error messages refer to the public MakeErrorStream class.
 
   if (!is_done_) {
-    LOG(ERROR) << "MakeErrorStream destructed without getting Status: " << file_
-               << ":" << line_ << " " << stream_.str();
+    LOG(ERROR) << "MakeErrorStream destructed without getting absl::Status: "
+               << file_ << ":" << line_ << " " << stream_.str();
   }
 }
 
-Status MakeErrorStream::Impl::GetStatus() {
+absl::Status MakeErrorStream::Impl::GetStatus() {
   // Note: error messages refer to the public MakeErrorStream class.
 
-  // Getting a Status object out more than once is not harmful, but
+  // Getting a absl::Status object out more than once is not harmful, but
   // it doesn't match the expected pattern, where the stream is constructed
-  // as a temporary, loaded with a message, and then casted to Status.
+  // as a temporary, loaded with a message, and then casted to absl::Status.
   if (is_done_) {
-    LOG(ERROR) << "MakeErrorStream got Status more than once: " << file_ << ":"
-               << line_ << " " << stream_.str();
+    LOG(ERROR) << "MakeErrorStream got absl::Status more than once: " << file_
+               << ":" << line_ << " " << stream_.str();
   }
 
   is_done_ = true;
@@ -170,8 +171,8 @@ Status MakeErrorStream::Impl::GetStatus() {
 
 void MakeErrorStream::Impl::CheckNotDone() const {
   if (is_done_) {
-    LOG(ERROR) << "MakeErrorStream shift called after getting Status: " << file_
-               << ":" << line_ << " " << stream_.str();
+    LOG(ERROR) << "MakeErrorStream shift called after getting absl::Status: "
+               << file_ << ":" << line_ << " " << stream_.str();
   }
 }
 
diff --git a/third_party/xla/xla/status_macros.h b/third_party/xla/xla/status_macros.h
index 71b4815bb3b576..02621a8f2c1e01 100644
--- a/third_party/xla/xla/status_macros.h
+++ b/third_party/xla/xla/status_macros.h
@@ -40,13 +40,13 @@ extern const char kPossibleAutoJitAlternative[];
 // Stream object used to collect error messages in MAKE_ERROR macros
 // or append error messages with APPEND_ERROR.  It accepts any
 // arguments with operator<< to build an error string, and then has an
-// implicit cast operator to Status, which converts the
-// logged string to a Status object and returns it, after logging the
+// implicit cast operator to absl::Status, which converts the
+// logged string to a absl::Status object and returns it, after logging the
 // error.  At least one call to operator<< is required; a compile time
 // error will be generated if none are given. Errors will only be
 // logged by default for certain status codes, as defined in
 // IsLoggedByDefault. This class will give ERROR errors if you don't
-// retrieve a Status exactly once before destruction.
+// retrieve a absl::Status exactly once before destruction.
 //
 // The class converts into an intermediate wrapper object
 // MakeErrorStreamWithOutput to check that the error stream gets at least one
@@ -56,7 +56,7 @@ class MakeErrorStream {
   // Wrapper around MakeErrorStream that only allows for output. This
   // is created as output of the first operator<< call on
   // MakeErrorStream. The bare MakeErrorStream does not have a
-  // Status operator. The net effect of that is that you
+  // absl::Status operator. The net effect of that is that you
   // have to call operator<< at least once or else you'll get a
   // compile time error.
   class MakeErrorStreamWithOutput {
@@ -70,10 +70,10 @@ class MakeErrorStream {
       return *this;
     }
 
-    // Implicit cast operators to Status and StatusOr.
+    // Implicit cast operators to absl::Status and absl::StatusOr.
     // Exactly one of these must be called exactly once before destruction.
     // NOLINTNEXTLINE(google-explicit-constructor)
-    operator Status() { return wrapped_error_stream_->GetStatus(); }
+    operator absl::Status() { return wrapped_error_stream_->GetStatus(); }
     template <typename T>
     // NOLINTNEXTLINE(google-explicit-constructor)
     operator absl::StatusOr<T>() {
@@ -117,13 +117,14 @@ class MakeErrorStream {
    public:
     Impl(const char* file, int line, tsl::error::Code code,
          MakeErrorStream* error_stream, bool is_logged_by_default = true);
-    Impl(const Status& status, PriorMessageHandling prior_message_handling,
-         const char* file, int line, MakeErrorStream* error_stream);
+    Impl(const absl::Status& status,
+         PriorMessageHandling prior_message_handling, const char* file,
+         int line, MakeErrorStream* error_stream);
 
     ~Impl();
 
     // This must be called exactly once before destruction.
-    Status GetStatus();
+    absl::Status GetStatus();
 
     void CheckNotDone() const;
 
@@ -134,16 +135,16 @@ class MakeErrorStream {
 
     PriorMessageHandling prior_message_handling_ = kAppendToPriorMessage;
     std::string prior_message_;
-    bool is_done_;  // true after Status object has been returned
+    bool is_done_;  // true after absl::Status object has been returned
     std::ostringstream stream_;
     bool should_log_;
     int log_severity_;
     bool should_log_stack_trace_;
 
     // Wrapper around the MakeErrorStream object that has a
-    // Status conversion. The first << operator called on
+    // absl::Status conversion. The first << operator called on
     // MakeErrorStream will return this object, and only this object
-    // can implicitly convert to Status. The net effect of
+    // can implicitly convert to absl::Status. The net effect of
     // this is that you'll get a compile time error if you call
     // MAKE_ERROR etc. without adding any output.
     MakeErrorStreamWithOutput make_error_stream_with_output_wrapper_;
@@ -156,7 +157,7 @@ class MakeErrorStream {
   void CheckNotDone() const;
 
   // Returns the status. Used by MakeErrorStreamWithOutput.
-  Status GetStatus() const { return impl_->GetStatus(); }
+  absl::Status GetStatus() const { return impl_->GetStatus(); }
 
   // Store the actual data on the heap to reduce stack frame sizes.
   std::unique_ptr<Impl> impl_;
@@ -175,17 +176,18 @@ TF_ATTRIBUTE_NOINLINE MakeErrorStream::MakeErrorStream(const char* file,
 // that declares a variable.
 class StatusAdaptorForMacros {
  public:
-  explicit StatusAdaptorForMacros(Status status) : status_(std::move(status)) {}
+  explicit StatusAdaptorForMacros(absl::Status status)
+      : status_(std::move(status)) {}
 
   StatusAdaptorForMacros(const StatusAdaptorForMacros&) = delete;
   StatusAdaptorForMacros& operator=(const StatusAdaptorForMacros&) = delete;
 
   explicit operator bool() const { return ABSL_PREDICT_TRUE(status_.ok()); }
 
-  Status&& Consume() { return std::move(status_); }
+  absl::Status&& Consume() { return std::move(status_); }
 
  private:
-  Status status_;
+  absl::Status status_;
 };
 
 }  // namespace status_macros
diff --git a/third_party/xla/xla/status_macros_test.cc b/third_party/xla/xla/status_macros_test.cc
index 7573ccc638f577..817eb0a9f77c1f 100644
--- a/third_party/xla/xla/status_macros_test.cc
+++ b/third_party/xla/xla/status_macros_test.cc
@@ -27,37 +27,37 @@ limitations under the License.
 
 namespace xla {
 
-Status RetCheckFail() {
+absl::Status RetCheckFail() {
   TF_RET_CHECK(2 > 3);
   return OkStatus();
 }
 
-Status RetCheckFailWithExtraMessage() {
+absl::Status RetCheckFailWithExtraMessage() {
   TF_RET_CHECK(2 > 3) << "extra message";
   return OkStatus();
 }
 
-Status RetCheckSuccess() {
+absl::Status RetCheckSuccess() {
   TF_RET_CHECK(3 > 2);
   return OkStatus();
 }
 
 TEST(StatusMacros, RetCheckFailing) {
-  Status status = RetCheckFail();
+  absl::Status status = RetCheckFail();
   EXPECT_EQ(status.code(), tsl::error::INTERNAL);
   EXPECT_THAT(status.message(),
               ::testing::ContainsRegex("RET_CHECK failure.*2 > 3"));
 }
 
 TEST(StatusMacros, RetCheckFailingWithExtraMessage) {
-  Status status = RetCheckFailWithExtraMessage();
+  absl::Status status = RetCheckFailWithExtraMessage();
   EXPECT_EQ(status.code(), tsl::error::INTERNAL);
   EXPECT_THAT(status.message(),
               ::testing::ContainsRegex("RET_CHECK.*2 > 3 extra message"));
 }
 
 TEST(StatusMacros, RetCheckSucceeding) {
-  Status status = RetCheckSuccess();
+  absl::Status status = RetCheckSuccess();
   EXPECT_IS_OK(status);
 }
 
@@ -72,11 +72,11 @@ TEST(StatusMacros, AssignOrAssertOnOK) {
   EXPECT_EQ(42, result);
 }
 
-Status ReturnStatusOK() { return OkStatus(); }
+absl::Status ReturnStatusOK() { return OkStatus(); }
 
-Status ReturnStatusError() { return (tsl::errors::Internal("foobar")); }
+absl::Status ReturnStatusError() { return (tsl::errors::Internal("foobar")); }
 
-using StatusReturningFunction = std::function<Status()>;
+using StatusReturningFunction = std::function<absl::Status()>;
 
 absl::StatusOr<int> CallStatusReturningFunction(
     const StatusReturningFunction& func) {
@@ -97,7 +97,7 @@ TEST(StatusMacros, ReturnIfErrorOnError) {
 }
 
 TEST(StatusMacros, AssignOrReturnSuccessfully) {
-  Status status = []() {
+  absl::Status status = []() {
     TF_ASSIGN_OR_RETURN(int value, CreateIntSuccessfully());
     EXPECT_EQ(value, 42);
     return OkStatus();
@@ -106,7 +106,7 @@ TEST(StatusMacros, AssignOrReturnSuccessfully) {
 }
 
 TEST(StatusMacros, AssignOrReturnUnsuccessfully) {
-  Status status = []() {
+  absl::Status status = []() {
     TF_ASSIGN_OR_RETURN(int value, CreateIntUnsuccessfully());
     (void)value;
     return OkStatus();
diff --git a/third_party/xla/xla/statusor.h b/third_party/xla/xla/statusor.h
index a8704bfd45d065..66c02ab4e74489 100644
--- a/third_party/xla/xla/statusor.h
+++ b/third_party/xla/xla/statusor.h
@@ -20,7 +20,7 @@ limitations under the License.
 
 namespace xla {
 
-// Use steam_executor's StatusOr so we don't duplicate code.
+// Use steam_executor's absl::StatusOr so we don't duplicate code.
 using tsl::StatusOr;  // TENSORFLOW_STATUS_OK
 
 }  // namespace xla
diff --git a/third_party/xla/xla/test_helpers.h b/third_party/xla/xla/test_helpers.h
index 11425fd51fa35e..7a80fd2f1874d5 100644
--- a/third_party/xla/xla/test_helpers.h
+++ b/third_party/xla/xla/test_helpers.h
@@ -31,10 +31,13 @@ class Literal;
 namespace testing {
 
 namespace internal_status {
-inline const Status& GetStatus(const Status& status) { return status; }
+// TODO(b/340953531) Eliminate this function.
+inline const absl::Status& GetStatus(const absl::Status& status) {
+  return status;
+}
 
 template <typename T>
-inline const Status& GetStatus(const StatusOr<T>& status) {
+inline const absl::Status& GetStatus(const absl::StatusOr<T>& status) {
   return status.status();
 }
 }  // namespace internal_status
@@ -45,8 +48,8 @@ inline const Status& GetStatus(const StatusOr<T>& status) {
 // The following macros are similar to macros in gmock, but deliberately named
 // differently in order to avoid conflicts in files which include both.
 
-// Macros for testing the results of functions that return Status or
-// StatusOr<T> (for any type T).
+// Macros for testing the results of functions that return absl::Status or
+// absl::StatusOr<T> (for any type T).
 #define EXPECT_IS_OK(expression) \
   EXPECT_EQ(::absl::OkStatus(),  \
             xla::testing::internal_status::GetStatus(expression))
diff --git a/third_party/xla/xla/text_literal_reader.cc b/third_party/xla/xla/text_literal_reader.cc
index 06773cc229f03f..209790bc33eafa 100644
--- a/third_party/xla/xla/text_literal_reader.cc
+++ b/third_party/xla/xla/text_literal_reader.cc
@@ -43,7 +43,8 @@ absl::StatusOr<Literal> TextLiteralReader::ReadPath(absl::string_view path) {
   CHECK(!absl::EndsWith(path, ".gz"))
       << "TextLiteralReader no longer supports reading .gz files";
   std::unique_ptr<tsl::RandomAccessFile> file;
-  Status s = tsl::Env::Default()->NewRandomAccessFile(std::string(path), &file);
+  absl::Status s =
+      tsl::Env::Default()->NewRandomAccessFile(std::string(path), &file);
   if (!s.ok()) {
     return s;
   }
@@ -59,7 +60,7 @@ absl::StatusOr<Literal> TextLiteralReader::ReadAllLines() {
   tsl::io::RandomAccessInputStream stream(file_.get());
   tsl::io::BufferedInputStream buf(&stream, 65536);
   std::string shape_string;
-  Status s = buf.ReadLine(&shape_string);
+  absl::Status s = buf.ReadLine(&shape_string);
   if (!s.ok()) {
     return s;
   }
diff --git a/third_party/xla/xla/text_literal_writer.cc b/third_party/xla/xla/text_literal_writer.cc
index c07f66e6b5b73f..050eb5fe835adc 100644
--- a/third_party/xla/xla/text_literal_writer.cc
+++ b/third_party/xla/xla/text_literal_writer.cc
@@ -29,8 +29,8 @@ limitations under the License.
 
 namespace xla {
 
-/* static */ Status TextLiteralWriter::WriteToPath(const Literal& literal,
-                                                   absl::string_view path) {
+/* static */ absl::Status TextLiteralWriter::WriteToPath(
+    const Literal& literal, absl::string_view path) {
   std::unique_ptr<tsl::WritableFile> f;
   auto s = tsl::Env::Default()->NewWritableFile(std::string(path), &f);
   if (!s.ok()) {
@@ -42,7 +42,7 @@ namespace xla {
     return s;
   }
 
-  Status status;
+  absl::Status status;
   tsl::WritableFile* f_ptr = f.get();
   literal.EachCellAsString([f_ptr, &status](absl::Span<const int64_t> indices,
                                             const std::string& value) {
diff --git a/third_party/xla/xla/text_literal_writer.h b/third_party/xla/xla/text_literal_writer.h
index 6ca94296cf73d3..2ce5b368773d34 100644
--- a/third_party/xla/xla/text_literal_writer.h
+++ b/third_party/xla/xla/text_literal_writer.h
@@ -36,7 +36,8 @@ namespace xla {
 // This should be readable by xla::TextLiteralReader.
 class TextLiteralWriter {
  public:
-  static Status WriteToPath(const Literal& literal, absl::string_view path);
+  static absl::Status WriteToPath(const Literal& literal,
+                                  absl::string_view path);
 
  private:
   TextLiteralWriter(const TextLiteralWriter&) = delete;
diff --git a/third_party/xla/xla/util.cc b/third_party/xla/xla/util.cc
index 52fb915e8dfee9..585e9615c98d4f 100644
--- a/third_party/xla/xla/util.cc
+++ b/third_party/xla/xla/util.cc
@@ -68,7 +68,7 @@ std::vector<int64_t> ToMixedRadix(const int64_t n,
   return digits;
 }
 
-Status WithLogBacktrace(const Status& status) {
+absl::Status WithLogBacktrace(const absl::Status& status) {
   CHECK(!status.ok());
   VLOG(1) << status.ToString();
   VLOG(2) << tsl::CurrentStackTrace();
@@ -113,14 +113,16 @@ void ScopedLoggingTimer::StopAndLog() {
 
 ScopedLoggingTimer::~ScopedLoggingTimer() { StopAndLog(); }
 
-Status AddStatus(Status prior, absl::string_view context) {
+absl::Status AddStatus(absl::Status prior, absl::string_view context) {
   CHECK(!prior.ok());
-  return Status{prior.code(), absl::StrCat(context, ": ", prior.message())};
+  return absl::Status{prior.code(),
+                      absl::StrCat(context, ": ", prior.message())};
 }
 
-Status AppendStatus(Status prior, absl::string_view context) {
+absl::Status AppendStatus(absl::Status prior, absl::string_view context) {
   CHECK(!prior.ok());
-  return Status{prior.code(), absl::StrCat(prior.message(), ": ", context)};
+  return absl::Status{prior.code(),
+                      absl::StrCat(prior.message(), ": ", context)};
 }
 
 std::string Reindent(absl::string_view original,
diff --git a/third_party/xla/xla/util.h b/third_party/xla/xla/util.h
index 76653654085658..e9c132cf010dce 100644
--- a/third_party/xla/xla/util.h
+++ b/third_party/xla/xla/util.h
@@ -71,10 +71,10 @@ std::vector<int64_t> ToMixedRadix(int64_t n, absl::Span<const int64_t> bounds);
 
 // Logs the provided status message with a backtrace.
 //
-// For use by Status-factories, logs a backtrace at the point where the status
-// is created, such that we can use --vmodule=util=1 to see all status
+// For use by absl::Status-factories, logs a backtrace at the point where the
+// status is created, such that we can use --vmodule=util=1 to see all status
 // creation backtraces.
-Status WithLogBacktrace(const Status& status);
+absl::Status WithLogBacktrace(const absl::Status& status);
 
 // Ranks greater than 6 are very rare, so use InlinedVector<int64_t, 6> to store
 // the bounds and indices. And for the rare cases of ranks greater than 6,
@@ -207,10 +207,10 @@ void StridedCopy(D* dest, int64_t dest_stride, const S* src, int64_t src_stride,
 }
 
 // Adds some context information to the error message in a
-// Status.  This is useful as Statuses are
+// absl::Status.  This is useful as absl::Statuses are
 // propagated upwards.
-Status AddStatus(Status prior, absl::string_view context);
-Status AppendStatus(Status prior, absl::string_view context);
+absl::Status AddStatus(absl::Status prior, absl::string_view context);
+absl::Status AppendStatus(absl::Status prior, absl::string_view context);
 
 // The following three macros define a common set of code for creating
 // absl::Status errors with the given error_type, with the addition of adding
@@ -221,7 +221,7 @@ Status AppendStatus(Status prior, absl::string_view context);
 //
 // template <typename... Args>
 // struct ResourceExhausted {
-//   Status status;
+//   absl::Status status;
 // #if defined(PLATFORM_GOOGLE)
 //   // NOLINTNEXTLINE(google-explicit-constructor)
 //   ResourceExhausted(const absl::FormatSpec<Args...>& format, Args&&... args,
@@ -238,16 +238,16 @@ Status AppendStatus(Status prior, absl::string_view context);
 // #endif
 //
 //   // NOLINTNEXTLINE(google-explicit-constructor)
-//   operator Status() const { return status; }
+//   operator absl::Status() const { return status; }
 // };
 //
 #define XLA_ERROR_WITH_STRFORMAT_AND_BACKTRACE_PREFIX(error_type) \
   template <typename... Args>                                     \
   struct error_type {                                             \
-    Status status;
+    absl::Status status;
 #define XLA_ERROR_WITH_STRFORMAT_AND_BACKTRACE_SUFFIX(error_type)        \
   /* NOLINTNEXTLINE(google-explicit-constructor) */                      \
-  operator Status() const { return status; }                             \
+  operator absl::Status() const { return status; }                       \
   }                                                                      \
   ;                                                                      \
   /*Deduction guide to make variadic arguments play nice with default */ \
@@ -269,8 +269,8 @@ Status AppendStatus(Status prior, absl::string_view context);
 #else
 #define XLA_ERROR_WITH_STRFORMAT_AND_BACKTRACE(error_type)          \
   template <typename... Args>                                       \
-  Status error_type(const absl::FormatSpec<Args...>& format,        \
-                    const Args&... args) {                          \
+  absl::Status error_type(const absl::FormatSpec<Args...>& format,  \
+                          const Args&... args) {                    \
     return WithLogBacktrace(                                        \
         absl::error_type##Error(absl::StrFormat(format, args...))); \
   }
@@ -299,7 +299,7 @@ XLA_ERROR_WITH_STRFORMAT_AND_BACKTRACE(Unknown);
 //
 // template <typename... Args>
 // struct ResourceExhaustedStrCat {
-//   Status status;
+//   absl::Status status;
 // #if defined(PLATFORM_GOOGLE)
 //   // NOLINTNEXTLINE(google-explicit-constructor)
 //   ResourceExhaustedStrCat(Args&&... concat, absl::SourceLocation loc =
@@ -317,17 +317,17 @@ XLA_ERROR_WITH_STRFORMAT_AND_BACKTRACE(Unknown);
 // #endif
 //
 //   // NOLINTNEXTLINE(google-explicit-constructor)
-//   operator Status() const { return status; }
+//   operator absl::Status() const { return status; }
 // };
 //
 #define XLA_ERROR_WITH_STRCAT_AND_BACKTRACE_PREFIX(error_type) \
   template <typename... Args>                                  \
   struct error_type##StrCat {                                  \
-    Status status;                                             \
+    absl::Status status;                                       \
     /* NOLINTNEXTLINE(google-explicit-constructor) */
 #define XLA_ERROR_WITH_STRCAT_AND_BACKTRACE_SUFFIX(error_type)           \
   /* NOLINTNEXTLINE(google-explicit-constructor) */                      \
-  operator Status() const { return status; }                             \
+  operator absl::Status() const { return status; }                       \
   }                                                                      \
   ;                                                                      \
   /*Deduction guide to make variadic arguments play nice with default */ \
@@ -788,7 +788,7 @@ bool IsInt32(T x) {
 }
 
 template <typename T>
-Status EraseElementFromVector(std::vector<T>* container, const T& value) {
+absl::Status EraseElementFromVector(std::vector<T>* container, const T& value) {
   // absl::c_find returns a const_iterator which does not seem to work on
   // gcc 4.8.4, and this breaks the ubuntu/xla_gpu build bot.
   auto it = std::find(container->begin(), container->end(), value);

From 6ec902447b28138b7fc45720cd9cba6c341a1bd5 Mon Sep 17 00:00:00 2001
From: Daniel Kenji Toyama <kenjitoyama@google.com>
Date: Thu, 16 May 2024 15:11:32 -0700
Subject: [PATCH 309/478] Fix typo `s/Recieved/Received/`.

PiperOrigin-RevId: 634543297
---
 tensorflow/python/checkpoint/functional_saver.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorflow/python/checkpoint/functional_saver.py b/tensorflow/python/checkpoint/functional_saver.py
index 48d8f043458ca4..995e850b82d53b 100644
--- a/tensorflow/python/checkpoint/functional_saver.py
+++ b/tensorflow/python/checkpoint/functional_saver.py
@@ -284,7 +284,7 @@ def __init__(
 
           if (checkpoint_key, slice_spec) in self._keys_to_restore_fn:
             raise ValueError(
-                "Recieved multiple tensors with the same checkpoint key and "
+                "Received multiple tensors with the same checkpoint key and "
                 "slice spec. This is invalid because one will overwrite the "
                 "other in the checkpoint. This indicates a bug in the "
                 "Checkpoint key-generation.")

From 9c5e56cd91bb3206f9f9751bbfcc4426abca9cca Mon Sep 17 00:00:00 2001
From: Kyle Lucke <klucke@google.com>
Date: Thu, 16 May 2024 15:20:27 -0700
Subject: [PATCH 310/478] Use absl::Status instead of xla::Status now that
 they're identical.

PiperOrigin-RevId: 634545999
---
 .../xla/xla/service/cpu/cpu_compiler.cc       |  34 +++--
 .../xla/xla/service/cpu/cpu_compiler.h        |  12 +-
 .../xla/xla/service/cpu/cpu_executable.cc     |   4 +-
 .../xla/xla/service/cpu/cpu_executable.h      |   2 +-
 .../xla/service/cpu/cpu_layout_assignment.cc  |   2 +-
 .../xla/service/cpu/cpu_layout_assignment.h   |   2 +-
 .../xla/service/cpu/cpu_transfer_manager.cc   |  10 +-
 .../xla/service/cpu/cpu_transfer_manager.h    |  14 +-
 third_party/xla/xla/service/cpu/cpu_xfeed.cc  |  14 +-
 third_party/xla/xla/service/cpu/cpu_xfeed.h   |  14 +-
 .../xla/xla/service/cpu/dot_op_emitter.cc     |  47 +++---
 .../xla/xla/service/cpu/dot_op_emitter.h      |  16 +-
 third_party/xla/xla/service/cpu/ir_emitter.cc | 140 +++++++++---------
 third_party/xla/xla/service/cpu/ir_emitter.h  | 134 ++++++++---------
 .../xla/xla/service/cpu/ir_function.cc        |   2 +-
 third_party/xla/xla/service/cpu/ir_function.h |   2 +-
 .../xla/xla/service/cpu/mlir_emitter.cc       |   2 +-
 .../xla/xla/service/cpu/mlir_emitter.h        |   2 +-
 .../xla/service/cpu/onednn_matmul_rewriter.cc |  41 ++---
 .../xla/service/cpu/onednn_matmul_rewriter.h  |   2 +-
 .../xla/xla/service/cpu/onednn_memory_util.cc |   2 +-
 .../xla/xla/service/cpu/onednn_memory_util.h  |   3 +-
 .../xla/service/cpu/onednn_ops_rewriter.cc    |   8 +-
 .../xla/xla/service/cpu/onednn_ops_rewriter.h |   2 +-
 .../xla/xla/service/cpu/onednn_rewriter.h     |   2 +-
 .../service/cpu/parallel_task_assignment.cc   |   3 +-
 26 files changed, 262 insertions(+), 254 deletions(-)

diff --git a/third_party/xla/xla/service/cpu/cpu_compiler.cc b/third_party/xla/xla/service/cpu/cpu_compiler.cc
index ac9fa032c88c0a..6bed3bd4263cb4 100644
--- a/third_party/xla/xla/service/cpu/cpu_compiler.cc
+++ b/third_party/xla/xla/service/cpu/cpu_compiler.cc
@@ -351,13 +351,13 @@ class CollectProfileCandidates : public DfsHloVisitorWithDefault {
       : hlo_to_profile_idx_(hlo_to_profile_idx),
         assigned_indices_(assigned_indices) {}
 
-  Status DefaultAction(HloInstruction* hlo_instruction) override {
+  absl::Status DefaultAction(HloInstruction* hlo_instruction) override {
     hlo_to_profile_idx_->insert(
         {hlo_instruction, FindOrDie(assigned_indices_, hlo_instruction)});
     return absl::OkStatus();
   }
 
-  Status HandleCall(HloInstruction* call) override {
+  absl::Status HandleCall(HloInstruction* call) override {
     TF_RETURN_IF_ERROR(DefaultAction(call));
     CollectProfileCandidates candidates_for_call(hlo_to_profile_idx_,
                                                  assigned_indices_);
@@ -365,7 +365,7 @@ class CollectProfileCandidates : public DfsHloVisitorWithDefault {
     return absl::OkStatus();
   }
   // Recurse into "conditional" so we can profile inside of it.
-  Status HandleConditional(HloInstruction* conditional) override {
+  absl::Status HandleConditional(HloInstruction* conditional) override {
     TF_RETURN_IF_ERROR(DefaultAction(conditional));
 
     CollectProfileCandidates candidates_for_true(hlo_to_profile_idx_,
@@ -382,12 +382,16 @@ class CollectProfileCandidates : public DfsHloVisitorWithDefault {
   }
 
   // Skip constants, there is nothing to profile.
-  Status HandleConstant(HloInstruction*) override { return absl::OkStatus(); }
+  absl::Status HandleConstant(HloInstruction*) override {
+    return absl::OkStatus();
+  }
   // Skip parameters, they are a simple load.
-  Status HandleParameter(HloInstruction*) override { return absl::OkStatus(); }
+  absl::Status HandleParameter(HloInstruction*) override {
+    return absl::OkStatus();
+  }
   // It is important to recurse for "while" or else we risk overly coarse
   // profiling information.
-  Status HandleWhile(HloInstruction* xla_while) override {
+  absl::Status HandleWhile(HloInstruction* xla_while) override {
     TF_RETURN_IF_ERROR(DefaultAction(xla_while));
 
     CollectProfileCandidates candidates_for_condition(hlo_to_profile_idx_,
@@ -423,7 +427,7 @@ void AddHloVerifier(HloPassPipeline* pipeline, HloVerifierOpts&& opts = {},
 
 }  // namespace
 
-Status CpuCompiler::RunHloPassesThroughLayoutAssn(
+absl::Status CpuCompiler::RunHloPassesThroughLayoutAssn(
     HloModule* module, bool is_aot_compile,
     LLVMTargetMachineFeatures* target_machine_features, bool is_mlir_compile) {
   const DebugOptions& debug_options = module->config().debug_options();
@@ -696,7 +700,7 @@ Status CpuCompiler::RunHloPassesThroughLayoutAssn(
   return pipeline.Run(module).status();
 }
 
-Status CpuCompiler::RunHloPassesAfterLayoutAssn(
+absl::Status CpuCompiler::RunHloPassesAfterLayoutAssn(
     HloModule* module, bool is_aot_compile,
     LLVMTargetMachineFeatures* target_machine_features,
     const CompileOptions& compile_options, bool is_mlir_compile) {
@@ -799,10 +803,10 @@ Status CpuCompiler::RunHloPassesAfterLayoutAssn(
   return pipeline.Run(module).status();
 }
 
-Status CpuCompiler::RunHloPasses(HloModule* module, bool is_aot_compile,
-                                 llvm::TargetMachine* target_machine,
-                                 const CompileOptions& compile_options,
-                                 bool is_mlir_compile) {
+absl::Status CpuCompiler::RunHloPasses(HloModule* module, bool is_aot_compile,
+                                       llvm::TargetMachine* target_machine,
+                                       const CompileOptions& compile_options,
+                                       bool is_mlir_compile) {
   LLVMTargetMachineFeatures target_machine_features(target_machine);
   TF_RETURN_IF_ERROR(RunHloPassesThroughLayoutAssn(
       module, is_aot_compile, &target_machine_features, is_mlir_compile));
@@ -870,7 +874,7 @@ std::pair<LLVMCompiler::ModuleHook, LLVMCompiler::ModuleHook> GetIRModuleHooks(
           }};
 }
 
-Status VerifyLlvmModule(const llvm::Module& llvm_module) {
+absl::Status VerifyLlvmModule(const llvm::Module& llvm_module) {
   XLA_SCOPED_LOGGING_TIMER("CpuCompiler - Running LLVM verifier");
 
   std::string err;
@@ -885,7 +889,7 @@ Status VerifyLlvmModule(const llvm::Module& llvm_module) {
   return absl::OkStatus();
 }
 
-Status CreateHloProfilingArtifacts(
+absl::Status CreateHloProfilingArtifacts(
     const HloModule& module,
     absl::flat_hash_map<const HloInstruction*, int64_t>*
         instruction_to_profile_idx,
@@ -1404,7 +1408,7 @@ CpuCompiler::CompileAheadOfTime(std::unique_ptr<HloModuleGroup> module_group,
       // Run the LLVM verifier over the unoptimized LLVM IR.  If it fails, run
       // the pre-optimization IR dump hook before returning.
       {
-        Status verify_status = VerifyLlvmModule(*llvm_module);
+        absl::Status verify_status = VerifyLlvmModule(*llvm_module);
         if (!verify_status.ok() && pre_optimization_ir_hook) {
           pre_optimization_ir_hook(*llvm_module);
         }
diff --git a/third_party/xla/xla/service/cpu/cpu_compiler.h b/third_party/xla/xla/service/cpu/cpu_compiler.h
index 69ba7f498f3db7..3cc54c4a616e3f 100644
--- a/third_party/xla/xla/service/cpu/cpu_compiler.h
+++ b/third_party/xla/xla/service/cpu/cpu_compiler.h
@@ -195,19 +195,19 @@ class CpuCompiler : public LLVMCompiler {
 
   // Runs the HLO passes which are necessary for both optimizations and
   // correctness.
-  Status RunHloPasses(HloModule* module, bool is_aot_compile,
-                      llvm::TargetMachine* target_machine,
-                      const CompileOptions& compile_options,
-                      bool is_mlir_compile = false);
+  absl::Status RunHloPasses(HloModule* module, bool is_aot_compile,
+                            llvm::TargetMachine* target_machine,
+                            const CompileOptions& compile_options,
+                            bool is_mlir_compile = false);
 
   // Runs HLO passes up to and including layout assignment.
-  Status RunHloPassesThroughLayoutAssn(
+  absl::Status RunHloPassesThroughLayoutAssn(
       HloModule* module, bool /*is_aot_compile*/,
       LLVMTargetMachineFeatures* target_machine_features,
       bool is_mlir_compile = false);
 
   // Runs HLO passes after layout assignment.
-  Status RunHloPassesAfterLayoutAssn(
+  absl::Status RunHloPassesAfterLayoutAssn(
       HloModule* module, bool is_aot_compile,
       LLVMTargetMachineFeatures* target_machine_features,
       const CompileOptions& compile_options, bool is_mlir_compile);
diff --git a/third_party/xla/xla/service/cpu/cpu_executable.cc b/third_party/xla/xla/service/cpu/cpu_executable.cc
index 01a5007a3deb74..02fb083554ce1c 100644
--- a/third_party/xla/xla/service/cpu/cpu_executable.cc
+++ b/third_party/xla/xla/service/cpu/cpu_executable.cc
@@ -180,7 +180,7 @@ CpuExecutable::CreateBufferTable(se::DeviceMemoryAllocator* memory_allocator,
   return std::move(buffers);
 }
 
-Status CpuExecutable::ExecuteComputeFunction(
+absl::Status CpuExecutable::ExecuteComputeFunction(
     const ExecutableRunOptions* run_options,
     absl::Span<MaybeOwningDeviceMemory const> buffers,
     HloExecutionProfile* hlo_execution_profile) {
@@ -399,7 +399,7 @@ absl::StatusOr<ExecutionOutput> CpuExecutable::ExecuteAsyncOnStream(
     std::shared_ptr<std::vector<MaybeOwningDeviceMemory>> task_buffers;
     HloExecutionProfile* hlo_execution_profile;
 
-    Status operator()() {
+    absl::Status operator()() {
       return executable->ExecuteComputeFunction(
           &run_options.run_options(), *task_buffers, hlo_execution_profile);
     }
diff --git a/third_party/xla/xla/service/cpu/cpu_executable.h b/third_party/xla/xla/service/cpu/cpu_executable.h
index 2dc7e2d97009e3..5cf8aa83357a3d 100644
--- a/third_party/xla/xla/service/cpu/cpu_executable.h
+++ b/third_party/xla/xla/service/cpu/cpu_executable.h
@@ -73,7 +73,7 @@ class CpuExecutable : public Executable {
 
   // Calls the generated function performing the computation with the given
   // arguments using the supplied buffers.
-  Status ExecuteComputeFunction(
+  absl::Status ExecuteComputeFunction(
       const ExecutableRunOptions* run_options,
       absl::Span<MaybeOwningDeviceMemory const> buffers,
       HloExecutionProfile* hlo_execution_profile);
diff --git a/third_party/xla/xla/service/cpu/cpu_layout_assignment.cc b/third_party/xla/xla/service/cpu/cpu_layout_assignment.cc
index 35371ea946369f..48f5ae59bfe6c3 100644
--- a/third_party/xla/xla/service/cpu/cpu_layout_assignment.cc
+++ b/third_party/xla/xla/service/cpu/cpu_layout_assignment.cc
@@ -120,7 +120,7 @@ static bool OperandsAndResultMustHaveRowMajorLayout(
   return false;
 }
 
-Status CpuLayoutAssignment::AddBackendConstraints(
+absl::Status CpuLayoutAssignment::AddBackendConstraints(
     LayoutConstraints* constraints) {
   ShouldMakeOperandColMajorCache cache;
 
diff --git a/third_party/xla/xla/service/cpu/cpu_layout_assignment.h b/third_party/xla/xla/service/cpu/cpu_layout_assignment.h
index 35ecde418bc60f..26e155458f569f 100644
--- a/third_party/xla/xla/service/cpu/cpu_layout_assignment.h
+++ b/third_party/xla/xla/service/cpu/cpu_layout_assignment.h
@@ -37,7 +37,7 @@ class CpuLayoutAssignment : public LayoutAssignment {
   ~CpuLayoutAssignment() override {}
 
  protected:
-  Status AddBackendConstraints(LayoutConstraints* constraints) override;
+  absl::Status AddBackendConstraints(LayoutConstraints* constraints) override;
 
   const TargetMachineFeatures& target_machine_features_;
 };
diff --git a/third_party/xla/xla/service/cpu/cpu_transfer_manager.cc b/third_party/xla/xla/service/cpu/cpu_transfer_manager.cc
index 27e62fc1af723d..fd68f2996357f6 100644
--- a/third_party/xla/xla/service/cpu/cpu_transfer_manager.cc
+++ b/third_party/xla/xla/service/cpu/cpu_transfer_manager.cc
@@ -43,19 +43,19 @@ CpuTransferManager::CpuTransferManager()
     : GenericTransferManager(se::host::kHostPlatformId,
                              /*pointer_size=*/sizeof(void*)) {}
 
-Status CpuTransferManager::TransferLiteralToInfeed(
+absl::Status CpuTransferManager::TransferLiteralToInfeed(
     se::StreamExecutor* executor, const LiteralSlice& literal) {
   return TransferLiteralToInfeedOnCpu(executor->device_ordinal(), literal);
 }
 
-Status CpuTransferManager::TransferLiteralFromOutfeed(
+absl::Status CpuTransferManager::TransferLiteralFromOutfeed(
     se::StreamExecutor* executor, MutableBorrowingLiteral literal) {
   return TransferLiteralFromOutfeedOnCpu(executor->device_ordinal(), literal);
 }
 
-Status CpuTransferManager::ReadDynamicShapes(se::Stream* stream,
-                                             const ShapedBuffer* device_buffer,
-                                             Shape* device_shape) {
+absl::Status CpuTransferManager::ReadDynamicShapes(
+    se::Stream* stream, const ShapedBuffer* device_buffer,
+    Shape* device_shape) {
   if (stream != nullptr) {
     // When a stream is presented, respect the stream dependency.
     return TransferManager::ReadDynamicShapes(stream, device_buffer,
diff --git a/third_party/xla/xla/service/cpu/cpu_transfer_manager.h b/third_party/xla/xla/service/cpu/cpu_transfer_manager.h
index 9cdf0478a47cb2..ed7f0fe3f7b383 100644
--- a/third_party/xla/xla/service/cpu/cpu_transfer_manager.h
+++ b/third_party/xla/xla/service/cpu/cpu_transfer_manager.h
@@ -37,10 +37,10 @@ class CpuTransferManager : public GenericTransferManager {
   CpuTransferManager();
   ~CpuTransferManager() override {}
 
-  Status TransferLiteralToInfeed(se::StreamExecutor* executor,
-                                 const LiteralSlice& literal) override;
-  Status TransferLiteralFromOutfeed(se::StreamExecutor* executor,
-                                    MutableBorrowingLiteral literal) override;
+  absl::Status TransferLiteralToInfeed(se::StreamExecutor* executor,
+                                       const LiteralSlice& literal) override;
+  absl::Status TransferLiteralFromOutfeed(
+      se::StreamExecutor* executor, MutableBorrowingLiteral literal) override;
 
   bool CanShapedBufferBeAccessedNow(
       se::StreamExecutor* executor,
@@ -54,9 +54,9 @@ class CpuTransferManager : public GenericTransferManager {
     return true;
   }
 
-  Status ReadDynamicShapes(se::Stream* stream,
-                           const ShapedBuffer* device_buffer,
-                           Shape* device_shape) override;
+  absl::Status ReadDynamicShapes(se::Stream* stream,
+                                 const ShapedBuffer* device_buffer,
+                                 Shape* device_shape) override;
 
  private:
   bool PackSubbyteTypes() const override { return true; }
diff --git a/third_party/xla/xla/service/cpu/cpu_xfeed.cc b/third_party/xla/xla/service/cpu/cpu_xfeed.cc
index b6fea6cb0d3851..8b4ea238f518bc 100644
--- a/third_party/xla/xla/service/cpu/cpu_xfeed.cc
+++ b/third_party/xla/xla/service/cpu/cpu_xfeed.cc
@@ -102,8 +102,8 @@ absl::StatusOr<cpu::runtime::XfeedBuffer*> TransferBufferToInfeedInternal(
   return queued_buffer;
 }
 
-Status TransferBufferToInfeed(int device_ordinal, int64_t size,
-                              const void* source) {
+absl::Status TransferBufferToInfeed(int device_ordinal, int64_t size,
+                                    const void* source) {
   TF_ASSIGN_OR_RETURN(cpu::runtime::XfeedBuffer * buffer,
                       TransferBufferToInfeedInternal(size, source));
 
@@ -175,8 +175,8 @@ absl::StatusOr<Shape> TransferTupleBuffersFromOutfeed(
 }
 }  // namespace
 
-Status TransferLiteralToInfeedOnCpu(int device_ordinal,
-                                    const LiteralSlice& literal) {
+absl::Status TransferLiteralToInfeedOnCpu(int device_ordinal,
+                                          const LiteralSlice& literal) {
   const Shape& shape = literal.shape();
   VLOG(2) << "Transferring literal to infeed with shape: "
           << ShapeUtil::HumanString(shape);
@@ -221,8 +221,8 @@ Status TransferLiteralToInfeedOnCpu(int device_ordinal,
   return OkStatus();
 }
 
-Status TransferLiteralFromOutfeedOnCpu(int device_ordinal,
-                                       MutableBorrowingLiteral literal) {
+absl::Status TransferLiteralFromOutfeedOnCpu(int device_ordinal,
+                                             MutableBorrowingLiteral literal) {
   if (!literal.shape().IsTuple()) {
     int64_t size =
         cpu::runtime::GetByteSizeRequirement(literal.shape(), sizeof(void*));
@@ -275,7 +275,7 @@ Status TransferLiteralFromOutfeedOnCpu(int device_ordinal,
   return OkStatus();
 }
 
-Status ReadDynamicShapesOnCpu(
+absl::Status ReadDynamicShapesOnCpu(
     const ShapedBuffer* device_buffer, Shape* device_shape,
     HloCostAnalysis::ShapeSizeFunction shape_size_fn) {
   TF_RET_CHECK(device_shape->is_dynamic());
diff --git a/third_party/xla/xla/service/cpu/cpu_xfeed.h b/third_party/xla/xla/service/cpu/cpu_xfeed.h
index 26512839d50174..7d09ec862d84d1 100644
--- a/third_party/xla/xla/service/cpu/cpu_xfeed.h
+++ b/third_party/xla/xla/service/cpu/cpu_xfeed.h
@@ -30,17 +30,17 @@ limitations under the License.
 namespace xla {
 
 // Helper function to transfers to infeed on CPU.
-Status TransferLiteralToInfeedOnCpu(int device_ordinal,
-                                    const LiteralSlice& literal);
+absl::Status TransferLiteralToInfeedOnCpu(int device_ordinal,
+                                          const LiteralSlice& literal);
 
 // Helper function to transfers from outfeed on CPU.
-Status TransferLiteralFromOutfeedOnCpu(int device_ordinal,
-                                       MutableBorrowingLiteral literal);
+absl::Status TransferLiteralFromOutfeedOnCpu(int device_ordinal,
+                                             MutableBorrowingLiteral literal);
 
 // Helper function to retrieve dynamic shape on CPU.
-Status ReadDynamicShapesOnCpu(const ShapedBuffer* device_buffer,
-                              Shape* device_shape,
-                              HloCostAnalysis::ShapeSizeFunction shape_size_fn);
+absl::Status ReadDynamicShapesOnCpu(
+    const ShapedBuffer* device_buffer, Shape* device_shape,
+    HloCostAnalysis::ShapeSizeFunction shape_size_fn);
 }  // namespace xla
 
 #endif  // XLA_SERVICE_CPU_CPU_XFEED_H_
diff --git a/third_party/xla/xla/service/cpu/dot_op_emitter.cc b/third_party/xla/xla/service/cpu/dot_op_emitter.cc
index c6bad1e54ec045..9548b9919ac722 100644
--- a/third_party/xla/xla/service/cpu/dot_op_emitter.cc
+++ b/third_party/xla/xla/service/cpu/dot_op_emitter.cc
@@ -134,21 +134,21 @@ class DotOpEmitter {
                         const TargetMachineFeatures& target_machine_features);
 
   // Emits the IR to perform the dot operation.
-  Status Emit();
+  absl::Status Emit();
 
   // Emits the IR to perform the batch dot operation.
-  Status EmitBatch();
+  absl::Status EmitBatch();
 
  private:
   // Emits instructions to perform a scalar dot product (a multiply of the
   // LHS and RHS) and store the results in the target.
-  Status EmitScalarDot();
+  absl::Status EmitScalarDot();
 
   // Emits a call to the CPU runtime to perform the matrix multiply.
-  Status EmitCallToRuntime();
+  absl::Status EmitCallToRuntime();
 
   // Emits a call to the CPU runtime to perform the batch matrix multiply.
-  Status EmitCallToBatchRuntime();
+  absl::Status EmitCallToBatchRuntime();
 
   // Represents the dimensions of a matrix-matrix multiply operation.
   struct MatMultDims {
@@ -192,7 +192,7 @@ class DotOpEmitter {
   void EmitTiledLlvmIrGemm();
 
   // Lowers the dot operation through MLIR's linalg.matmul.
-  Status EmitLinalgMatmul();
+  absl::Status EmitLinalgMatmul();
 
   // Lowers the dot operation as a naive nested loop that computes the result
   // one element at a time.
@@ -264,7 +264,7 @@ DotOpEmitter::DotOpEmitter(
       hlo_module_config_(hlo_module_config),
       target_machine_features_(target_machine_features) {}
 
-Status DotOpEmitter::EmitLinalgMatmul() {
+absl::Status DotOpEmitter::EmitLinalgMatmul() {
   Shape operand_shapes[] = {dot_info_.lhs_shape, dot_info_.rhs_shape};
   llvm::Value* operand_ptrs[] = {lhs_array_.GetBasePointer(),
                                  rhs_array_.GetBasePointer()};
@@ -529,7 +529,7 @@ void DotOpEmitter::EmitTiledLlvmIrGemv() {
   }
 }
 
-Status DotOpEmitter::Emit() {
+absl::Status DotOpEmitter::Emit() {
   // The dot operation performs a sum of products over dimension 0 of the left
   // hand side operand and dimension 1 of the right hand side operand.
   //
@@ -584,7 +584,7 @@ Status DotOpEmitter::Emit() {
   }
 }
 
-Status DotOpEmitter::EmitBatch() {
+absl::Status DotOpEmitter::EmitBatch() {
   // The dot operation performs a sum of products over dimension 0 of the left
   // hand side operand and dimension 1 of the right hand side operand.
   //
@@ -756,7 +756,7 @@ void DotOpEmitter::EmitNaiveLlvmIrGemm() {
   b_->SetInsertPoint(loop_nest.GetOuterLoopExitBasicBlock());
 }
 
-Status DotOpEmitter::EmitScalarDot() {
+absl::Status DotOpEmitter::EmitScalarDot() {
   // A scalar dot is just a scalar multiply.
   llvm::Value* result;
   // Use the same index_type for all tensor accesses in the same kernel.
@@ -791,7 +791,7 @@ Status DotOpEmitter::EmitScalarDot() {
   return OkStatus();
 }
 
-Status DotOpEmitter::EmitCallToRuntime() {
+absl::Status DotOpEmitter::EmitCallToRuntime() {
   // The signature of the Eigen runtime matmul function is:
   //
   //   (void)(void* run_options, float* out, float* lhs, float* rhs,
@@ -902,7 +902,7 @@ Status DotOpEmitter::EmitCallToRuntime() {
   return OkStatus();
 }
 
-Status DotOpEmitter::EmitCallToBatchRuntime() {
+absl::Status DotOpEmitter::EmitCallToBatchRuntime() {
   // The signature of the runtime batch matmul function is:
   //
   //   (void)(void* run_options, float* out, float* lhs, float* rhs,
@@ -1218,7 +1218,7 @@ DotImplementationStrategy GetDotImplementationStrategy(
   return DotImplementationStrategy::kNaiveLlvmIr;
 }
 
-Status EmitNonBatchDotOperation(
+absl::Status EmitNonBatchDotOperation(
     DotInfo dot_info, std::string hlo_name,
     const llvm_ir::IrArray& target_array, const llvm_ir::IrArray& lhs_array,
     const llvm_ir::IrArray& rhs_array, const llvm_ir::IrArray* addend_array,
@@ -1270,7 +1270,8 @@ llvm_ir::IrArray CollapseFirstNDims(llvm::IRBuilder<>* b,
                           std::move(new_shape));
 }
 
-Status ValidateDotDimensionNumbers(const DotDimensionNumbers& dim_numbers) {
+absl::Status ValidateDotDimensionNumbers(
+    const DotDimensionNumbers& dim_numbers) {
   // Checks some invariants that do not hold in general, but DotDecomposer
   // should have established for us.  This is just a debugging aid.
   TF_RET_CHECK(dim_numbers.lhs_contracting_dimensions_size() == 1);
@@ -1357,7 +1358,7 @@ bool PotentiallyImplementedAsEigenMatmul(
   return impl_strategy == DotImplementationStrategy::kEigen;
 }
 
-Status EmitBatchDotOperation(
+absl::Status EmitBatchDotOperation(
     const HloInstruction& dot, const llvm_ir::IrArray& target_array,
     const llvm_ir::IrArray& lhs_array, const llvm_ir::IrArray& rhs_array,
     llvm::Value* executable_run_options_value, llvm::IRBuilder<>* b,
@@ -1481,15 +1482,13 @@ bool DotOperandsAndResultMustHaveRowMajorLayout(
          impl_strategy == DotImplementationStrategy::kEigen;
 }
 
-Status EmitDotOperation(const HloInstruction& dot,
-                        const llvm_ir::IrArray& target_array,
-                        const llvm_ir::IrArray& lhs_array,
-                        const llvm_ir::IrArray& rhs_array,
-                        const llvm_ir::IrArray* addend_array,
-                        llvm::Value* executable_run_options_value,
-                        llvm::IRBuilder<>* b, mlir::MLIRContext* mlir_context,
-                        const HloModuleConfig& hlo_module_config,
-                        const TargetMachineFeatures& target_machine_features) {
+absl::Status EmitDotOperation(
+    const HloInstruction& dot, const llvm_ir::IrArray& target_array,
+    const llvm_ir::IrArray& lhs_array, const llvm_ir::IrArray& rhs_array,
+    const llvm_ir::IrArray* addend_array,
+    llvm::Value* executable_run_options_value, llvm::IRBuilder<>* b,
+    mlir::MLIRContext* mlir_context, const HloModuleConfig& hlo_module_config,
+    const TargetMachineFeatures& target_machine_features) {
   // This routine assumes that the dot operation is not in a parallelized
   // enclosing computation.
   CHECK(dot.parent()
diff --git a/third_party/xla/xla/service/cpu/dot_op_emitter.h b/third_party/xla/xla/service/cpu/dot_op_emitter.h
index 9d4f56179601f5..58e3afee737f7b 100644
--- a/third_party/xla/xla/service/cpu/dot_op_emitter.h
+++ b/third_party/xla/xla/service/cpu/dot_op_emitter.h
@@ -57,15 +57,13 @@ std::optional<int64_t> ProfitableToMakeDotOperandColumnMajor(
 // dimensions as the result, and the result is computed as `addend_array` +
 // dot(`lhs_array`, `rhs_array`).  A non-null `addend_array` is only supported
 // for Matrix-vector products.
-Status EmitDotOperation(const HloInstruction& dot,
-                        const llvm_ir::IrArray& target_array,
-                        const llvm_ir::IrArray& lhs_array,
-                        const llvm_ir::IrArray& rhs_array,
-                        const llvm_ir::IrArray* addend_array,
-                        llvm::Value* executable_run_options_value,
-                        llvm::IRBuilder<>* b, mlir::MLIRContext* mlir_context,
-                        const HloModuleConfig& hlo_module_config,
-                        const TargetMachineFeatures& target_machine_features);
+absl::Status EmitDotOperation(
+    const HloInstruction& dot, const llvm_ir::IrArray& target_array,
+    const llvm_ir::IrArray& lhs_array, const llvm_ir::IrArray& rhs_array,
+    const llvm_ir::IrArray* addend_array,
+    llvm::Value* executable_run_options_value, llvm::IRBuilder<>* b,
+    mlir::MLIRContext* mlir_context, const HloModuleConfig& hlo_module_config,
+    const TargetMachineFeatures& target_machine_features);
 }  // namespace cpu
 }  // namespace xla
 
diff --git a/third_party/xla/xla/service/cpu/ir_emitter.cc b/third_party/xla/xla/service/cpu/ir_emitter.cc
index dff2ee5e7cde59..d6a5a832d03ba8 100644
--- a/third_party/xla/xla/service/cpu/ir_emitter.cc
+++ b/third_party/xla/xla/service/cpu/ir_emitter.cc
@@ -126,7 +126,7 @@ IrEmitter::IrEmitter(mlir::MLIRContext* mlir_context,
       target_machine_features_(*target_machine_features),
       emit_code_for_msan_(emit_code_for_msan) {
   b_.setFastMathFlags(llvm_ir::GetCpuFastMathFlags(hlo_module_config_));
-  Status s = GatherComputationsByAllocationType(
+  absl::Status s = GatherComputationsByAllocationType(
       &hlo_module, &thread_local_computations_, &global_computations_);
   absl::c_sort(thread_local_computations_);
   absl::c_sort(global_computations_);
@@ -255,7 +255,7 @@ void IrEmitter::InitializeIrFunction(const std::string& function_name) {
                                    module_, &b_, num_dynamic_loop_bounds_);
 }
 
-Status IrEmitter::HandleBitcast(HloInstruction* bitcast) {
+absl::Status IrEmitter::HandleBitcast(HloInstruction* bitcast) {
   VLOG(2) << "HandleBitcast: " << bitcast->ToString();
   emitted_value_[bitcast] = GetEmittedValueFor(bitcast->operand(0));
   return OkStatus();
@@ -277,7 +277,7 @@ llvm::Constant* IrEmitter::EmitGlobalForLiteral(const Literal& literal) {
   return result_global;
 }
 
-Status IrEmitter::EmitConstantGlobals() {
+absl::Status IrEmitter::EmitConstantGlobals() {
   for (const BufferAllocation& allocation : assignment_.Allocations()) {
     if (!allocation.is_constant()) {
       continue;
@@ -300,14 +300,14 @@ Status IrEmitter::EmitConstantGlobals() {
   return OkStatus();
 }
 
-Status IrEmitter::HandleConstant(HloInstruction* constant) {
+absl::Status IrEmitter::HandleConstant(HloInstruction* constant) {
   VLOG(2) << "HandleConstant: " << constant->ToString();
   // IrEmitter::EmitConstantGlobals has already taken care of emitting the body
   // of the constant.
   return EmitTargetAddressForOp(constant);
 }
 
-Status IrEmitter::HandleCopy(HloInstruction* copy) {
+absl::Status IrEmitter::HandleCopy(HloInstruction* copy) {
   if (copy->shape().IsTuple() ||
       (copy->shape().IsArray() &&
        LayoutUtil::Equal(copy->operand(0)->shape().layout(),
@@ -382,7 +382,8 @@ void IrEmitter::AttachDereferenceableMetadataForLoad(llvm::LoadInst* load,
   }
 }
 
-Status IrEmitter::HandleGetTupleElement(HloInstruction* get_tuple_element) {
+absl::Status IrEmitter::HandleGetTupleElement(
+    HloInstruction* get_tuple_element) {
   // A tuple is an array of pointers, one for each operand. Each pointer points
   // to the output buffer of its corresponding operand. A GetTupleElement
   // instruction forwards a pointer to the tuple element buffer at the given
@@ -395,13 +396,13 @@ Status IrEmitter::HandleGetTupleElement(HloInstruction* get_tuple_element) {
   return OkStatus();
 }
 
-Status IrEmitter::HandleSelect(HloInstruction* select) {
+absl::Status IrEmitter::HandleSelect(HloInstruction* select) {
   auto pred = select->operand(0);
   TF_RET_CHECK(pred->shape().element_type() == PRED);
   return DefaultAction(select);
 }
 
-Status IrEmitter::HandleInfeed(HloInstruction* instruction) {
+absl::Status IrEmitter::HandleInfeed(HloInstruction* instruction) {
   HloInfeedInstruction* infeed = Cast<HloInfeedInstruction>(instruction);
   VLOG(2) << "HandleInfeed: " << infeed->ToString();
 
@@ -461,8 +462,8 @@ Status IrEmitter::HandleInfeed(HloInstruction* instruction) {
   return OkStatus();
 }
 
-Status IrEmitter::EmitXfeedTransfer(XfeedKind kind, const Shape& shape,
-                                    llvm::Value* program_buffer_address) {
+absl::Status IrEmitter::EmitXfeedTransfer(XfeedKind kind, const Shape& shape,
+                                          llvm::Value* program_buffer_address) {
   int64_t length = ByteSizeOf(shape);
   if (length < 0 || length > std::numeric_limits<int32_t>::max()) {
     return InvalidArgument(
@@ -526,7 +527,7 @@ Status IrEmitter::EmitXfeedTransfer(XfeedKind kind, const Shape& shape,
   return OkStatus();
 }
 
-Status IrEmitter::HandleOutfeed(HloInstruction* outfeed) {
+absl::Status IrEmitter::HandleOutfeed(HloInstruction* outfeed) {
   // Outfeed produces no useful result, but it does return a token[] that can be
   // threaded through to other side effecting operations to ensure ordering.  In
   // the IR emitter we treat this token as a normal u8[] and thus need to insert
@@ -556,7 +557,7 @@ Status IrEmitter::HandleOutfeed(HloInstruction* outfeed) {
   return OkStatus();
 }
 
-Status IrEmitter::HandleSort(HloInstruction* hlo) {
+absl::Status IrEmitter::HandleSort(HloInstruction* hlo) {
   const HloSortInstruction* sort = Cast<HloSortInstruction>(hlo);
   TF_RETURN_IF_ERROR(EmitTargetAddressForOp(sort));
   Shape keys_shape = sort->keys()->shape();
@@ -651,7 +652,7 @@ Status IrEmitter::HandleSort(HloInstruction* hlo) {
   return OkStatus();
 }
 
-Status IrEmitter::HandleTuple(HloInstruction* tuple) {
+absl::Status IrEmitter::HandleTuple(HloInstruction* tuple) {
   TF_RETURN_IF_ERROR(EmitTargetAddressForOp(tuple));
   llvm::SmallVector<llvm::Value*> base_ptrs;
   for (auto operand : tuple->operands()) {
@@ -661,7 +662,7 @@ Status IrEmitter::HandleTuple(HloInstruction* tuple) {
   return OkStatus();
 }
 
-Status IrEmitter::HandleReduceWindow(HloInstruction* reduce_window) {
+absl::Status IrEmitter::HandleReduceWindow(HloInstruction* reduce_window) {
   // Pseudo code for reduce window:
   //
   //   for (coordinates O in the output)
@@ -677,12 +678,13 @@ Status IrEmitter::HandleReduceWindow(HloInstruction* reduce_window) {
   // that works.
   bool saved_allow_reassociation = allow_reassociation_;
   allow_reassociation_ = true;
-  Status status = DefaultAction(reduce_window);
+  absl::Status status = DefaultAction(reduce_window);
   allow_reassociation_ = saved_allow_reassociation;
   return status;
 }
 
-Status IrEmitter::HandleSelectAndScatter(HloInstruction* select_and_scatter) {
+absl::Status IrEmitter::HandleSelectAndScatter(
+    HloInstruction* select_and_scatter) {
   CHECK_EQ(select_and_scatter->operand_count(), 3);
   const auto operand = select_and_scatter->operand(0);
   const auto source = select_and_scatter->operand(1);
@@ -864,7 +866,7 @@ Status IrEmitter::HandleSelectAndScatter(HloInstruction* select_and_scatter) {
   return OkStatus();
 }
 
-Status IrEmitter::HandleDot(HloInstruction* dot) {
+absl::Status IrEmitter::HandleDot(HloInstruction* dot) {
   auto lhs = dot->operand(0);
   auto rhs = dot->operand(1);
   TF_RETURN_IF_ERROR(ElementTypesSameAndSupported(
@@ -900,7 +902,7 @@ Status IrEmitter::HandleDot(HloInstruction* dot) {
                           hlo_module_config_, target_machine_features_);
 }
 
-Status IrEmitter::HandleConvolution(HloInstruction* convolution) {
+absl::Status IrEmitter::HandleConvolution(HloInstruction* convolution) {
   auto lhs = convolution->operand(0);
   auto rhs = convolution->operand(1);
   TF_RETURN_IF_ERROR(ElementTypesSameAndSupported(
@@ -1085,7 +1087,7 @@ Status IrEmitter::HandleConvolution(HloInstruction* convolution) {
   return DefaultAction(convolution);
 }
 
-Status IrEmitter::HandleFft(HloInstruction* fft) {
+absl::Status IrEmitter::HandleFft(HloInstruction* fft) {
   auto operand = fft->operand(0);
   TF_RETURN_IF_ERROR(ElementTypesSameAndSupported(
       /*instruction=*/*fft, /*operands=*/{operand},
@@ -1137,7 +1139,7 @@ Status IrEmitter::HandleFft(HloInstruction* fft) {
   return OkStatus();
 }
 
-Status IrEmitter::HandleAllReduceSingleReplica(HloInstruction* crs) {
+absl::Status IrEmitter::HandleAllReduceSingleReplica(HloInstruction* crs) {
   // When there is a single replica, a cross replica sum is the identity
   // function, and the buffer assignment expects a copy.
   //
@@ -1196,7 +1198,7 @@ static bool DataTypeIsSupportedByReduceScatter(PrimitiveType datatype) {
   }
 }
 
-Status IrEmitter::HandleAllReduceMultipleReplica(HloInstruction* crs) {
+absl::Status IrEmitter::HandleAllReduceMultipleReplica(HloInstruction* crs) {
   CHECK_GE(crs->operand_count(), 1);
   PrimitiveType datatype = crs->operand(0)->shape().element_type();
   TF_RETURN_IF_ERROR(EmitTargetAddressForOp(crs));
@@ -1280,7 +1282,7 @@ Status IrEmitter::HandleAllReduceMultipleReplica(HloInstruction* crs) {
   return OkStatus();
 }
 
-Status IrEmitter::HandleAllReduce(HloInstruction* crs) {
+absl::Status IrEmitter::HandleAllReduce(HloInstruction* crs) {
   if (hlo_module_config_.replica_count() == 1 &&
       hlo_module_config_.num_partitions() == 1) {
     return HandleAllReduceSingleReplica(crs);
@@ -1288,7 +1290,7 @@ Status IrEmitter::HandleAllReduce(HloInstruction* crs) {
   return HandleAllReduceMultipleReplica(crs);
 }
 
-Status IrEmitter::HandleReduceScatter(HloInstruction* rs) {
+absl::Status IrEmitter::HandleReduceScatter(HloInstruction* rs) {
   CHECK_EQ(rs->operand_count(), 1);
   PrimitiveType datatype = rs->operand(0)->shape().element_type();
   TF_RETURN_IF_ERROR(EmitTargetAddressForOp(rs));
@@ -1344,7 +1346,7 @@ Status IrEmitter::HandleReduceScatter(HloInstruction* rs) {
   return OkStatus();
 }
 
-Status IrEmitter::HandleAllToAll(HloInstruction* instruction) {
+absl::Status IrEmitter::HandleAllToAll(HloInstruction* instruction) {
   auto* instr = Cast<HloAllToAllInstruction>(instruction);
   TF_RETURN_IF_ERROR(EmitTargetAddressForOp(instruction));
   CHECK(!instr->split_dimension() && instr->shape().IsTuple())
@@ -1400,7 +1402,7 @@ Status IrEmitter::HandleAllToAll(HloInstruction* instruction) {
   return OkStatus();
 }
 
-Status IrEmitter::HandleAllGather(HloInstruction* instruction) {
+absl::Status IrEmitter::HandleAllGather(HloInstruction* instruction) {
   TF_RETURN_IF_ERROR(EmitTargetAddressForOp(instruction));
 
   std::string replica_groups =
@@ -1451,7 +1453,7 @@ Status IrEmitter::HandleAllGather(HloInstruction* instruction) {
   return OkStatus();
 }
 
-Status IrEmitter::HandleCollectivePermute(HloInstruction* crs) {
+absl::Status IrEmitter::HandleCollectivePermute(HloInstruction* crs) {
   auto* instr = Cast<HloCollectivePermuteInstruction>(crs);
   TF_RETURN_IF_ERROR(EmitTargetAddressForOp(instr));
   std::string source_target_pairs = absl::StrJoin(
@@ -1488,7 +1490,7 @@ Status IrEmitter::HandleCollectivePermute(HloInstruction* crs) {
   return OkStatus();
 }
 
-Status IrEmitter::HandlePartitionId(HloInstruction* hlo) {
+absl::Status IrEmitter::HandlePartitionId(HloInstruction* hlo) {
   TF_RETURN_IF_ERROR(EmitTargetAddressForOp(hlo));
   TF_ASSIGN_OR_RETURN(BufferAllocation::Slice output_slice,
                       assignment_.GetUniqueSlice(hlo, {}));
@@ -1500,7 +1502,7 @@ Status IrEmitter::HandlePartitionId(HloInstruction* hlo) {
   return OkStatus();
 }
 
-Status IrEmitter::HandleReplicaId(HloInstruction* hlo) {
+absl::Status IrEmitter::HandleReplicaId(HloInstruction* hlo) {
   TF_RETURN_IF_ERROR(EmitTargetAddressForOp(hlo));
   TF_ASSIGN_OR_RETURN(BufferAllocation::Slice output_slice,
                       assignment_.GetUniqueSlice(hlo, {}));
@@ -1512,7 +1514,7 @@ Status IrEmitter::HandleReplicaId(HloInstruction* hlo) {
   return OkStatus();
 }
 
-Status IrEmitter::HandleParameter(HloInstruction* parameter) {
+absl::Status IrEmitter::HandleParameter(HloInstruction* parameter) {
   VLOG(2) << "HandleParameter: " << parameter->ToString();
   return EmitTargetAddressForOp(parameter);
 }
@@ -1998,7 +2000,7 @@ absl::StatusOr<bool> IrEmitter::EmitVectorizedReduce(
   return true;
 }
 
-Status IrEmitter::HandleReduce(HloInstruction* reduce) {
+absl::Status IrEmitter::HandleReduce(HloInstruction* reduce) {
   auto arg = reduce->mutable_operand(0);
   auto init_value = reduce->mutable_operand(1);
   absl::Span<const int64_t> dimensions(reduce->dimensions());
@@ -2027,21 +2029,21 @@ Status IrEmitter::HandleReduce(HloInstruction* reduce) {
   return DefaultAction(reduce);
 }
 
-Status IrEmitter::HandleSend(HloInstruction* send) {
+absl::Status IrEmitter::HandleSend(HloInstruction* send) {
   // TODO(b/33942983): Support Send/Recv on CPU.
   return Unimplemented("Send is not implemented on CPU.");
 }
 
-Status IrEmitter::HandleSendDone(HloInstruction* send_done) {
+absl::Status IrEmitter::HandleSendDone(HloInstruction* send_done) {
   // TODO(b/33942983): Support Send/Recv on CPU.
   return Unimplemented("Send-done is not implemented on CPU.");
 }
 
-Status IrEmitter::HandleScatter(HloInstruction*) {
+absl::Status IrEmitter::HandleScatter(HloInstruction*) {
   return Unimplemented("Scatter is not implemented on CPUs.");
 }
 
-Status IrEmitter::HandleSlice(HloInstruction* slice) {
+absl::Status IrEmitter::HandleSlice(HloInstruction* slice) {
   VLOG(2) << "HandleSlice: " << slice->ToString();
   auto operand = slice->operand(0);
   // The code below emits a sequential loop nest. For the parallel backend, use
@@ -2178,7 +2180,7 @@ Status IrEmitter::HandleSlice(HloInstruction* slice) {
   return OkStatus();
 }
 
-Status IrEmitter::HandleDynamicSlice(HloInstruction* dynamic_slice) {
+absl::Status IrEmitter::HandleDynamicSlice(HloInstruction* dynamic_slice) {
   if (ShapeUtil::IsScalar(dynamic_slice->shape())) {
     TF_RETURN_IF_ERROR(EmitTargetAddressForOp(dynamic_slice));
     return EmitMemcpy(*dynamic_slice->operand(0), *dynamic_slice);
@@ -2186,7 +2188,7 @@ Status IrEmitter::HandleDynamicSlice(HloInstruction* dynamic_slice) {
   return DefaultAction(dynamic_slice);
 }
 
-Status IrEmitter::HandleDynamicUpdateSlice(
+absl::Status IrEmitter::HandleDynamicUpdateSlice(
     HloInstruction* dynamic_update_slice) {
   auto update = dynamic_update_slice->operand(1);
   if (ShapeUtil::IsScalar(dynamic_update_slice->shape())) {
@@ -2203,17 +2205,17 @@ Status IrEmitter::HandleDynamicUpdateSlice(
   return DefaultAction(dynamic_update_slice);
 }
 
-Status IrEmitter::HandleRecv(HloInstruction* recv) {
+absl::Status IrEmitter::HandleRecv(HloInstruction* recv) {
   // TODO(b/33942983): Support Send/Recv on CPU.
   return Unimplemented("Recv is not implemented on CPU.");
 }
 
-Status IrEmitter::HandleRecvDone(HloInstruction* recv_done) {
+absl::Status IrEmitter::HandleRecvDone(HloInstruction* recv_done) {
   // TODO(b/33942983): Support Send/Recv on CPU.
   return Unimplemented("Recv-done is not implemented on CPU.");
 }
 
-Status IrEmitter::HandlePad(HloInstruction* pad) {
+absl::Status IrEmitter::HandlePad(HloInstruction* pad) {
   // CPU backend does not properly handle negative padding but this is ok
   // because negative padding should be removed by the algebraic simplifier.
   for (auto& padding_dimension : pad->padding_config().dimensions()) {
@@ -2272,7 +2274,7 @@ Status IrEmitter::HandlePad(HloInstruction* pad) {
   return OkStatus();
 }
 
-Status IrEmitter::HandleFusion(HloInstruction* fusion) {
+absl::Status IrEmitter::HandleFusion(HloInstruction* fusion) {
   auto* root = fusion->fused_expression_root();
   if (llvm_ir::CanEmitFusedDynamicUpdateSliceInPlace(fusion, assignment_)) {
     VLOG(3) << "HandleFusion FusedDynamicUpdateSliceInPlace";
@@ -2327,7 +2329,7 @@ Status IrEmitter::HandleFusion(HloInstruction* fusion) {
   }
 }
 
-Status IrEmitter::HandleCall(HloInstruction* call) {
+absl::Status IrEmitter::HandleCall(HloInstruction* call) {
   HloComputation* computation = call->to_apply();
   llvm::Function* call_ir_function = FindOrDie(
       emitted_functions_, ComputationToEmit{computation, allow_reassociation_});
@@ -2372,7 +2374,7 @@ Status IrEmitter::HandleCall(HloInstruction* call) {
   return OkStatus();
 }
 
-Status IrEmitter::HandleSliceToDynamic(HloInstruction* hlo) {
+absl::Status IrEmitter::HandleSliceToDynamic(HloInstruction* hlo) {
   TF_RETURN_IF_ERROR(EmitTargetAddressForOp(hlo));
   std::vector<llvm::Value*> dynamic_dims;
   int32_t raw_data_size =
@@ -2400,7 +2402,7 @@ Status IrEmitter::HandleSliceToDynamic(HloInstruction* hlo) {
   //     dest_index = delinearize(linearize(i, dynamic_dim), static_dim)
   //     dest[dest_index] = source[i]
   auto loop_body_emitter =
-      [&](const llvm_ir::IrArray::Index& array_index) -> Status {
+      [&](const llvm_ir::IrArray::Index& array_index) -> absl::Status {
     llvm::Value* source_element =
         GetIrArrayFor(hlo->operand(0)).EmitReadArrayElement(array_index, &b_);
     llvm::Value* linear_index = array_index.Linearize(dynamic_dims, &b_);
@@ -2415,7 +2417,7 @@ Status IrEmitter::HandleSliceToDynamic(HloInstruction* hlo) {
       .EmitLoop(IrName(hlo));
 }
 
-Status IrEmitter::HandlePadToStatic(HloInstruction* hlo) {
+absl::Status IrEmitter::HandlePadToStatic(HloInstruction* hlo) {
   TF_RETURN_IF_ERROR(EmitTargetAddressForOp(hlo));
 
   TF_ASSIGN_OR_RETURN(BufferAllocation::Slice data_slice,
@@ -2463,7 +2465,7 @@ Status IrEmitter::HandlePadToStatic(HloInstruction* hlo) {
   //     source_index = delinearize(inearize(i, dynamic_dim), static_dim)
   //     dest[i] = source[source_index]
   auto loop_body_emitter =
-      [&](const llvm_ir::IrArray::Index& array_index) -> Status {
+      [&](const llvm_ir::IrArray::Index& array_index) -> absl::Status {
     llvm::Value* linear_index = array_index.Linearize(dynamic_dims, &b_);
     llvm_ir::IrArray::Index source_index(linear_index, input_shape, &b_);
     llvm::Value* source_element =
@@ -2480,7 +2482,7 @@ Status IrEmitter::HandlePadToStatic(HloInstruction* hlo) {
   return OkStatus();
 }
 
-Status IrEmitter::HandleTopK(HloInstruction* hlo) {
+absl::Status IrEmitter::HandleTopK(HloInstruction* hlo) {
   TF_RETURN_IF_ERROR(EmitTargetAddressForOp(hlo));
   const HloInstruction* input = hlo->operand(0);
   const int64_t k = hlo->shape().tuple_shapes(0).dimensions().back();
@@ -2520,8 +2522,8 @@ Status IrEmitter::HandleTopK(HloInstruction* hlo) {
 }
 
 #if defined(INTEL_MKL) && defined(ENABLE_ONEDNN_V3)
-Status IrEmitter::HandleOneDnnMatMulCalls(HloInstruction* custom_call,
-                                          std::string runtime_symbol_name) {
+absl::Status IrEmitter::HandleOneDnnMatMulCalls(
+    HloInstruction* custom_call, std::string runtime_symbol_name) {
   // We would like to emit LLVM IR for the following function call
   //      custom_call_target(void* result, void** args)
   // args can be thought of an array of pointers allocated on the stack,
@@ -2650,7 +2652,7 @@ Status IrEmitter::HandleOneDnnMatMulCalls(HloInstruction* custom_call,
   return OkStatus();
 }
 
-Status IrEmitter::HandleOneDnnLayerNorm(HloInstruction* custom_call) {
+absl::Status IrEmitter::HandleOneDnnLayerNorm(HloInstruction* custom_call) {
   //      args[0]: ptr to nargs
   //      args[1]: ptr to ExecutableRunOptions
   //      args[2]: ptr to OneDnnLayerNormConfig
@@ -2727,7 +2729,7 @@ Status IrEmitter::HandleOneDnnLayerNorm(HloInstruction* custom_call) {
   return OkStatus();
 }
 
-Status IrEmitter::HandleOneDnnSoftmax(HloInstruction* custom_call) {
+absl::Status IrEmitter::HandleOneDnnSoftmax(HloInstruction* custom_call) {
   auto input = custom_call->operand(0);
   llvm_ir::IrArray input_array(GetIrArrayFor(input));
   auto input_stack_alloca = GetAllocaAndEmitMemrefInfo(b_, input_array);
@@ -2751,7 +2753,7 @@ Status IrEmitter::HandleOneDnnSoftmax(HloInstruction* custom_call) {
 }
 #endif  // INTEL_MKL && ENABLE_ONEDNN_V3
 
-Status IrEmitter::HandleCustomCall(HloInstruction* custom_call) {
+absl::Status IrEmitter::HandleCustomCall(HloInstruction* custom_call) {
   if (custom_call->custom_call_target() == "PadToStatic") {
     return HandlePadToStatic(custom_call);
   }
@@ -2856,7 +2858,7 @@ Status IrEmitter::HandleCustomCall(HloInstruction* custom_call) {
   return OkStatus();
 }
 
-Status IrEmitter::HandleWhile(HloInstruction* xla_while) {
+absl::Status IrEmitter::HandleWhile(HloInstruction* xla_while) {
   // Precondition: Condition computation must return a scalar bool.
   HloComputation* condition = xla_while->while_condition();
   TF_RET_CHECK(ShapeUtil::IsScalar(condition->root_instruction()->shape()) &&
@@ -2867,7 +2869,7 @@ Status IrEmitter::HandleWhile(HloInstruction* xla_while) {
   TF_RETURN_IF_ERROR(ShapeUtil::ForEachSubshapeWithStatus(
       xla_while->shape(),
       [this, &xla_while](const Shape& /*subshape*/,
-                         const ShapeIndex& index) -> Status {
+                         const ShapeIndex& index) -> absl::Status {
         auto check = [this](const HloInstruction* a, const HloInstruction* b,
                             const ShapeIndex& index) -> absl::Status {
           const BufferAllocation::Slice slice_a =
@@ -3236,7 +3238,7 @@ void IrEmitter::EmitTransferElements(llvm::Value* target, llvm::Value* source,
   }
 }
 
-Status IrEmitter::HandleConcatenate(HloInstruction* concatenate) {
+absl::Status IrEmitter::HandleConcatenate(HloInstruction* concatenate) {
   absl::Span<HloInstruction* const> operands(concatenate->operands());
   std::string failure_reason;
   TF_ASSIGN_OR_RETURN(
@@ -3253,7 +3255,7 @@ Status IrEmitter::HandleConcatenate(HloInstruction* concatenate) {
   return DefaultAction(concatenate);
 }
 
-Status IrEmitter::HandleConditional(HloInstruction* conditional) {
+absl::Status IrEmitter::HandleConditional(HloInstruction* conditional) {
   auto branch_index = conditional->operand(0);
   int num_branches = conditional->branch_count();
   TF_RET_CHECK(ShapeUtil::IsScalar(branch_index->shape()) &&
@@ -3362,25 +3364,25 @@ Status IrEmitter::HandleConditional(HloInstruction* conditional) {
   return OkStatus();
 }
 
-Status IrEmitter::HandleAfterAll(HloInstruction* after_all) {
+absl::Status IrEmitter::HandleAfterAll(HloInstruction* after_all) {
   TF_RET_CHECK(ByteSizeOf(after_all->shape()) == 0);
   // No code to generate, but we need to emit an address for book-keeping.
   TF_RETURN_IF_ERROR(EmitTargetAddressForOp(after_all));
   return OkStatus();
 }
 
-Status IrEmitter::HandleAddDependency(HloInstruction* add_dependency) {
+absl::Status IrEmitter::HandleAddDependency(HloInstruction* add_dependency) {
   // AddDedendency just forwards its zero-th operand.
   emitted_value_[add_dependency] =
       GetEmittedValueFor(add_dependency->operand(0));
   return OkStatus();
 }
 
-Status IrEmitter::HandleRng(HloInstruction* rng) {
+absl::Status IrEmitter::HandleRng(HloInstruction* rng) {
   return Unimplemented("Rng should be expanded for CPU.");
 }
 
-Status IrEmitter::HandleRngGetAndUpdateState(HloInstruction* rng_state) {
+absl::Status IrEmitter::HandleRngGetAndUpdateState(HloInstruction* rng_state) {
   VLOG(2) << "RngGetAndUpdateState: " << rng_state->ToString();
   llvm::Value* old_state = llvm_ir::RngGetAndUpdateState(
       Cast<HloRngGetAndUpdateStateInstruction>(rng_state)->delta(), module_,
@@ -3398,7 +3400,7 @@ Status IrEmitter::HandleRngGetAndUpdateState(HloInstruction* rng_state) {
   return OkStatus();
 }
 
-Status IrEmitter::FinishVisit(HloInstruction* root) {
+absl::Status IrEmitter::FinishVisit(HloInstruction* root) {
   // When this method is called, we should have already emitted an IR value for
   // the root (return) op. The IR value holds the address of the buffer holding
   // the value. If the root is a constant or parameter, we perform a memcpy from
@@ -3591,7 +3593,7 @@ bool IsHloVeryCheap(const HloInstruction* hlo) {
 }
 }  // namespace
 
-Status IrEmitter::Preprocess(HloInstruction* hlo) {
+absl::Status IrEmitter::Preprocess(HloInstruction* hlo) {
   VLOG(3) << "Visiting: " << hlo->ToString();
   // When profiling is enabled, trace the same HLOs that the profiler does.
   if (instruction_to_profile_idx_.count(hlo) ||
@@ -3604,7 +3606,7 @@ Status IrEmitter::Preprocess(HloInstruction* hlo) {
   return OkStatus();
 }
 
-Status IrEmitter::Postprocess(HloInstruction* hlo) {
+absl::Status IrEmitter::Postprocess(HloInstruction* hlo) {
   if (auto* prof_counter = GetProfileCounterFor(*hlo)) {
     profiling_state_.RecordCycleDelta(&b_, hlo, prof_counter);
   }
@@ -3764,7 +3766,7 @@ llvm::Value* IrEmitter::EmitBufferPointer(const BufferAllocation::Slice& slice,
   }
 }
 
-Status IrEmitter::EmitTargetAddressForOp(const HloInstruction* op) {
+absl::Status IrEmitter::EmitTargetAddressForOp(const HloInstruction* op) {
   const Shape& target_shape = op->shape();
   TF_ASSIGN_OR_RETURN(const BufferAllocation::Slice slice,
                       assignment_.GetUniqueTopLevelSlice(op));
@@ -3774,13 +3776,13 @@ Status IrEmitter::EmitTargetAddressForOp(const HloInstruction* op) {
   return OkStatus();
 }
 
-Status IrEmitter::EmitTargetElementLoop(
+absl::Status IrEmitter::EmitTargetElementLoop(
     HloInstruction* target_op,
     const llvm_ir::ElementGenerator& element_generator) {
   return EmitTargetElementLoop(target_op, /*desc=*/"", element_generator);
 }
 
-Status IrEmitter::EmitTargetElementLoop(
+absl::Status IrEmitter::EmitTargetElementLoop(
     HloInstruction* target_op, absl::string_view desc,
     const llvm_ir::ElementGenerator& element_generator) {
   VLOG(2) << "EmitTargetElementLoop: " << target_op->ToString();
@@ -3833,8 +3835,8 @@ Status IrEmitter::EmitTargetElementLoop(
   return OkStatus();
 }
 
-Status IrEmitter::EmitMemcpy(const HloInstruction& source,
-                             const HloInstruction& destination) {
+absl::Status IrEmitter::EmitMemcpy(const HloInstruction& source,
+                                   const HloInstruction& destination) {
   llvm::Value* source_value = GetEmittedValueFor(&source);
   llvm::Value* destination_value = GetEmittedValueFor(&destination);
   int64_t source_size = ByteSizeOf(source.shape());
@@ -3844,7 +3846,7 @@ Status IrEmitter::EmitMemcpy(const HloInstruction& source,
   return OkStatus();
 }
 
-Status IrEmitter::ElementTypesSameAndSupported(
+absl::Status IrEmitter::ElementTypesSameAndSupported(
     const HloInstruction& instruction,
     absl::Span<const HloInstruction* const> operands,
     absl::Span<const PrimitiveType> supported_types) {
@@ -3863,7 +3865,7 @@ Status IrEmitter::ElementTypesSameAndSupported(
   return OkStatus();
 }
 
-Status IrEmitter::DefaultAction(HloInstruction* hlo) {
+absl::Status IrEmitter::DefaultAction(HloInstruction* hlo) {
   ElementalIrEmitter::HloToElementGeneratorMap operand_to_generator;
   for (const HloInstruction* operand : hlo->operands()) {
     operand_to_generator[operand] = [=](const llvm_ir::IrArray::Index& index) {
diff --git a/third_party/xla/xla/service/cpu/ir_emitter.h b/third_party/xla/xla/service/cpu/ir_emitter.h
index bc99acd4f2b2f9..21a371810116ac 100644
--- a/third_party/xla/xla/service/cpu/ir_emitter.h
+++ b/third_party/xla/xla/service/cpu/ir_emitter.h
@@ -124,7 +124,7 @@ class IrEmitter : public DfsHloVisitorWithDefault,
   llvm::IRBuilder<>* builder() { return &b_; }
 
   // Emit an LLVM global variable for every constant buffer allocation.
-  Status EmitConstantGlobals();
+  absl::Status EmitConstantGlobals();
 
  protected:
   //
@@ -132,55 +132,57 @@ class IrEmitter : public DfsHloVisitorWithDefault,
   //
   // Default action which emits code for most operations. Operations which are
   // special in some way are handled explicitly in HandleFoo methods.
-  Status DefaultAction(HloInstruction* hlo) override;
-
-  Status HandleAllGather(HloInstruction* instruction) override;
-  Status HandleAllToAll(HloInstruction* instruction) override;
-  Status HandleBitcast(HloInstruction* bitcast) override;
-  Status HandleConstant(HloInstruction* constant) override;
-  Status HandleCopy(HloInstruction* copy) override;
-  Status HandleGetTupleElement(HloInstruction* get_tuple_element) override;
-  Status HandleSelect(HloInstruction* select) override;
-  Status HandleDot(HloInstruction* dot) override;
-  Status HandleConvolution(HloInstruction* convolution) override;
-  Status HandleFft(HloInstruction* fft) override;
-  Status HandleAllReduce(HloInstruction* crs) override;
-  Status HandleReduceScatter(HloInstruction* crs) override;
-  Status HandleCollectivePermute(HloInstruction* crs) override;
-  Status HandleInfeed(HloInstruction* instruction) override;
-  Status HandleOutfeed(HloInstruction* outfeed) override;
-  Status HandleSort(HloInstruction* hlo) override;
-  Status HandleParameter(HloInstruction* parameter) override;
-  Status HandleReduce(HloInstruction* reduce) override;
-  Status HandleReduceWindow(HloInstruction* reduce_window) override;
-  Status HandleSelectAndScatter(HloInstruction* select_and_scatter) override;
-  Status HandleSend(HloInstruction* send) override;
-  Status HandleSendDone(HloInstruction* send_done) override;
-  Status HandleSlice(HloInstruction* slice) override;
-  Status HandleDynamicSlice(HloInstruction* dynamic_slice) override;
-  Status HandleDynamicUpdateSlice(
+  absl::Status DefaultAction(HloInstruction* hlo) override;
+
+  absl::Status HandleAllGather(HloInstruction* instruction) override;
+  absl::Status HandleAllToAll(HloInstruction* instruction) override;
+  absl::Status HandleBitcast(HloInstruction* bitcast) override;
+  absl::Status HandleConstant(HloInstruction* constant) override;
+  absl::Status HandleCopy(HloInstruction* copy) override;
+  absl::Status HandleGetTupleElement(
+      HloInstruction* get_tuple_element) override;
+  absl::Status HandleSelect(HloInstruction* select) override;
+  absl::Status HandleDot(HloInstruction* dot) override;
+  absl::Status HandleConvolution(HloInstruction* convolution) override;
+  absl::Status HandleFft(HloInstruction* fft) override;
+  absl::Status HandleAllReduce(HloInstruction* crs) override;
+  absl::Status HandleReduceScatter(HloInstruction* crs) override;
+  absl::Status HandleCollectivePermute(HloInstruction* crs) override;
+  absl::Status HandleInfeed(HloInstruction* instruction) override;
+  absl::Status HandleOutfeed(HloInstruction* outfeed) override;
+  absl::Status HandleSort(HloInstruction* hlo) override;
+  absl::Status HandleParameter(HloInstruction* parameter) override;
+  absl::Status HandleReduce(HloInstruction* reduce) override;
+  absl::Status HandleReduceWindow(HloInstruction* reduce_window) override;
+  absl::Status HandleSelectAndScatter(
+      HloInstruction* select_and_scatter) override;
+  absl::Status HandleSend(HloInstruction* send) override;
+  absl::Status HandleSendDone(HloInstruction* send_done) override;
+  absl::Status HandleSlice(HloInstruction* slice) override;
+  absl::Status HandleDynamicSlice(HloInstruction* dynamic_slice) override;
+  absl::Status HandleDynamicUpdateSlice(
       HloInstruction* dynamic_update_slice) override;
-  Status HandleRecv(HloInstruction* recv) override;
-  Status HandleRecvDone(HloInstruction* recv_done) override;
-  Status HandlePad(HloInstruction* pad) override;
-  Status HandleTuple(HloInstruction* tuple) override;
-  Status HandleFusion(HloInstruction* fusion) override;
-  Status HandleCall(HloInstruction* call) override;
-  Status HandleCustomCall(HloInstruction* custom_call) override;
-  Status HandleWhile(HloInstruction* xla_while) override;
-  Status HandleConcatenate(HloInstruction* concatenate) override;
-  Status HandleConditional(HloInstruction* conditional) override;
-  Status HandleScatter(HloInstruction* scatter) override;
-  Status HandleAfterAll(HloInstruction* after_all) override;
-  Status HandleAddDependency(HloInstruction* add_dependency) override;
-  Status HandlePartitionId(HloInstruction* hlo) override;
-  Status HandleReplicaId(HloInstruction* hlo) override;
-  Status HandleRng(HloInstruction* rng) override;
-  Status HandleRngGetAndUpdateState(HloInstruction* rng_state) override;
-  Status FinishVisit(HloInstruction* root) override;
-
-  Status Preprocess(HloInstruction* hlo) override;
-  Status Postprocess(HloInstruction* hlo) override;
+  absl::Status HandleRecv(HloInstruction* recv) override;
+  absl::Status HandleRecvDone(HloInstruction* recv_done) override;
+  absl::Status HandlePad(HloInstruction* pad) override;
+  absl::Status HandleTuple(HloInstruction* tuple) override;
+  absl::Status HandleFusion(HloInstruction* fusion) override;
+  absl::Status HandleCall(HloInstruction* call) override;
+  absl::Status HandleCustomCall(HloInstruction* custom_call) override;
+  absl::Status HandleWhile(HloInstruction* xla_while) override;
+  absl::Status HandleConcatenate(HloInstruction* concatenate) override;
+  absl::Status HandleConditional(HloInstruction* conditional) override;
+  absl::Status HandleScatter(HloInstruction* scatter) override;
+  absl::Status HandleAfterAll(HloInstruction* after_all) override;
+  absl::Status HandleAddDependency(HloInstruction* add_dependency) override;
+  absl::Status HandlePartitionId(HloInstruction* hlo) override;
+  absl::Status HandleReplicaId(HloInstruction* hlo) override;
+  absl::Status HandleRng(HloInstruction* rng) override;
+  absl::Status HandleRngGetAndUpdateState(HloInstruction* rng_state) override;
+  absl::Status FinishVisit(HloInstruction* root) override;
+
+  absl::Status Preprocess(HloInstruction* hlo) override;
+  absl::Status Postprocess(HloInstruction* hlo) override;
 
   // A convenient helper for calling BufferAssignment::GetUniqueSlice.
   BufferAllocation::Slice GetAllocationSlice(
@@ -189,16 +191,16 @@ class IrEmitter : public DfsHloVisitorWithDefault,
   }
 
  private:
-  Status HandleSliceToDynamic(HloInstruction* hlo);
-  Status HandlePadToStatic(HloInstruction* hlo);
-  Status HandleTopK(HloInstruction* hlo);
-  Status HandleAllReduceSingleReplica(HloInstruction* crs);
-  Status HandleAllReduceMultipleReplica(HloInstruction* crs);
+  absl::Status HandleSliceToDynamic(HloInstruction* hlo);
+  absl::Status HandlePadToStatic(HloInstruction* hlo);
+  absl::Status HandleTopK(HloInstruction* hlo);
+  absl::Status HandleAllReduceSingleReplica(HloInstruction* crs);
+  absl::Status HandleAllReduceMultipleReplica(HloInstruction* crs);
 #if defined(INTEL_MKL) && defined(ENABLE_ONEDNN_V3)
-  Status HandleOneDnnMatMulCalls(HloInstruction* hlo,
-                                 std::string runtime_symbol_name);
-  Status HandleOneDnnSoftmax(HloInstruction* hlo);
-  Status HandleOneDnnLayerNorm(HloInstruction* hlo);
+  absl::Status HandleOneDnnMatMulCalls(HloInstruction* hlo,
+                                       std::string runtime_symbol_name);
+  absl::Status HandleOneDnnSoftmax(HloInstruction* hlo);
+  absl::Status HandleOneDnnLayerNorm(HloInstruction* hlo);
 #endif  // INTEL_MKL && ENABLE_ONEDNN_V3
   // Private helper to initialize an IR function for the computation.
   void InitializeIrFunction(const std::string& function_name);
@@ -316,7 +318,7 @@ class IrEmitter : public DfsHloVisitorWithDefault,
 
   // Verifies that the element types of all of the given operand instructions
   // match and are of one of the given supported types.
-  Status ElementTypesSameAndSupported(
+  absl::Status ElementTypesSameAndSupported(
       const HloInstruction& instruction,
       absl::Span<const HloInstruction* const> operands,
       absl::Span<const PrimitiveType> supported_types);
@@ -331,23 +333,23 @@ class IrEmitter : public DfsHloVisitorWithDefault,
   // in the loop name.
   //
   // TODO(jingyue): target_op should be a `const HloInstruction*`.
-  Status EmitTargetElementLoop(
+  absl::Status EmitTargetElementLoop(
       HloInstruction* target_op,
       const llvm_ir::ElementGenerator& element_generator);
-  Status EmitTargetElementLoop(
+  absl::Status EmitTargetElementLoop(
       HloInstruction* target_op, absl::string_view desc,
       const llvm_ir::ElementGenerator& element_generator);
 
   // Emits a memcpy from the source instruction's result value to the
   // destination's.  Both source and destination must have an entry in the
   // emitted_value_ table.
-  Status EmitMemcpy(const HloInstruction& source,
-                    const HloInstruction& destination);
+  absl::Status EmitMemcpy(const HloInstruction& source,
+                          const HloInstruction& destination);
 
   // Emits IR to compute the target address of the buffer for the given op.
   // After calling this function, you can get a pointer to this buffer by
   // calling GetIrArrayForOp or GetEmittedValueFor.
-  Status EmitTargetAddressForOp(const HloInstruction* op);
+  absl::Status EmitTargetAddressForOp(const HloInstruction* op);
 
   // Structurizes "array_elements" into an MD array that represents "shape".
   // This is a recursive function, and "dimension_index" indicates the index of
@@ -652,8 +654,8 @@ class IrEmitter : public DfsHloVisitorWithDefault,
 
   // Emit IR to transfer between a {infeed,outfeed} buffer and an in-program
   // address.
-  Status EmitXfeedTransfer(XfeedKind kind, const Shape& shape,
-                           llvm::Value* program_buffer_address);
+  absl::Status EmitXfeedTransfer(XfeedKind kind, const Shape& shape,
+                                 llvm::Value* program_buffer_address);
 
   // Returns a ConstExpr bitcast.
   llvm::Constant* EmitGlobalForLiteral(const Literal& literal);
diff --git a/third_party/xla/xla/service/cpu/ir_function.cc b/third_party/xla/xla/service/cpu/ir_function.cc
index 69961501a37b8b..cf66fe382e75be 100644
--- a/third_party/xla/xla/service/cpu/ir_function.cc
+++ b/third_party/xla/xla/service/cpu/ir_function.cc
@@ -234,7 +234,7 @@ std::vector<llvm::Value*> GetArrayFunctionCallArguments(
 
 // Emits a call to a runtime fork/join function which dispatches parallel
 // calls to 'parallel_function' (and joins threads before returning).
-Status EmitCallToParallelForkJoin(
+absl::Status EmitCallToParallelForkJoin(
     const std::vector<llvm::Value*>& arguments, const Shape& shape,
     absl::Span<const int64_t> dimension_partition_counts, llvm::IRBuilder<>* b,
     llvm::Function* parallel_function, absl::string_view name) {
diff --git a/third_party/xla/xla/service/cpu/ir_function.h b/third_party/xla/xla/service/cpu/ir_function.h
index 47034675bdf8e6..73a85867438a9d 100644
--- a/third_party/xla/xla/service/cpu/ir_function.h
+++ b/third_party/xla/xla/service/cpu/ir_function.h
@@ -145,7 +145,7 @@ std::vector<llvm::Value*> GetArrayFunctionCallArguments(
 
 // Emits a call to a runtime fork/join function which dispatches parallel
 // calls to 'parallel_function' (and joins threads before returning).
-Status EmitCallToParallelForkJoin(
+absl::Status EmitCallToParallelForkJoin(
     const std::vector<llvm::Value*>& arguments, const Shape& shape,
     absl::Span<const int64_t> dimension_partition_counts, llvm::IRBuilder<>* b,
     llvm::Function* parallel_function, absl::string_view name);
diff --git a/third_party/xla/xla/service/cpu/mlir_emitter.cc b/third_party/xla/xla/service/cpu/mlir_emitter.cc
index 8d3a28815ee01b..396f567cf2ddd4 100644
--- a/third_party/xla/xla/service/cpu/mlir_emitter.cc
+++ b/third_party/xla/xla/service/cpu/mlir_emitter.cc
@@ -82,7 +82,7 @@ void BuildViewForBuffer(llvm::SmallVectorImpl<llvm::Value *> *args,
 }
 }  // namespace
 
-Status EmitMlirFuncAndCall(
+absl::Status EmitMlirFuncAndCall(
     mlir::MLIRContext *context, llvm::IRBuilder<> *b, const Shape &result_shape,
     llvm::ArrayRef<Shape> operand_shapes, llvm::Value *result_ptr,
     llvm::ArrayRef<llvm::Value *> operand_ptrs, llvm::StringRef func_name,
diff --git a/third_party/xla/xla/service/cpu/mlir_emitter.h b/third_party/xla/xla/service/cpu/mlir_emitter.h
index c7a8480e4e9d73..af1b1626af2a50 100644
--- a/third_party/xla/xla/service/cpu/mlir_emitter.h
+++ b/third_party/xla/xla/service/cpu/mlir_emitter.h
@@ -32,7 +32,7 @@ namespace cpu {
 // `emitter` and create a call, passing it the buffers defined by
 // resultShape/resultPtr and operandShapes/operandPtrs. The function is added to
 // the LLVM module at `b`s insertion point.
-Status EmitMlirFuncAndCall(
+absl::Status EmitMlirFuncAndCall(
     mlir::MLIRContext *context, llvm::IRBuilder<> *b, const Shape &result_shape,
     llvm::ArrayRef<Shape> operand_shapes, llvm::Value *result_ptr,
     llvm::ArrayRef<llvm::Value *> operand_ptrs, llvm::StringRef func_name,
diff --git a/third_party/xla/xla/service/cpu/onednn_matmul_rewriter.cc b/third_party/xla/xla/service/cpu/onednn_matmul_rewriter.cc
index a792a5ba52130d..1fc542088d1531 100644
--- a/third_party/xla/xla/service/cpu/onednn_matmul_rewriter.cc
+++ b/third_party/xla/xla/service/cpu/onednn_matmul_rewriter.cc
@@ -43,7 +43,7 @@ namespace {
 namespace m = match;
 namespace pu = ::xla::cpu::onednn_pattern_utils_internal;
 
-inline Status ValidateDotDimensionNumbers(
+inline absl::Status ValidateDotDimensionNumbers(
     const DotDimensionNumbers& dim_numbers) {
   // Checks some invariants that do not hold in general, but DotDecomposer
   // should have established for us.
@@ -276,8 +276,8 @@ auto GELUActivation(HloInstruction* instr, HloInstruction** src) {
 // OneDNN matmul can fuse add operation with automatic broadcasting along the
 // addend's dimensions that are 1s. When compatible, Broadcast can be replaced
 // by Bitcast, which is much cheaper. Compute new shape for the Bitcast.
-StatusOr<Shape> AdjustBiasShape(const HloInstruction* broadcast_instr,
-                                const Shape& dot_shape) {
+absl::StatusOr<Shape> AdjustBiasShape(const HloInstruction* broadcast_instr,
+                                      const Shape& dot_shape) {
   if (broadcast_instr->opcode() != HloOpcode::kBroadcast) {
     return absl::InvalidArgumentError(
         "Hlo instruction is not a Broadcast insruction.");
@@ -428,7 +428,7 @@ class OneDnnMatMulRewriteVisitor : public DfsHloRewriteVisitor {
  public:
   // Matches patterns for possible MatMul fusions that are supported by oneDNN
   // library. Matched HLO instruction(s) are replaced by custom call.
-  Status HandleDot(HloInstruction* instr) override {
+  absl::Status HandleDot(HloInstruction* instr) override {
     HloInstruction* dot_instr;
     auto pattern = m::Op(&dot_instr).WithOpcode(HloOpcode::kDot);
     if (!Match(instr, pattern)) return OkStatus();
@@ -463,7 +463,7 @@ class OneDnnMatMulRewriteVisitor : public DfsHloRewriteVisitor {
     return OkStatus();
   }
 
-  Status HandleAdd(HloInstruction* instr) override {
+  absl::Status HandleAdd(HloInstruction* instr) override {
     // Try to do a fusion for Dot(onednn-matmul) + Add. However,
     // HLO Add instruction might receive the addends after additional
     // processing like Broadcast, Bitcast, Convert, etc. is applied to the raw
@@ -616,7 +616,7 @@ class OneDnnMatMulRewriteVisitor : public DfsHloRewriteVisitor {
     return OkStatus();
   }
 
-  Status HandleMaximum(HloInstruction* instr) override {
+  absl::Status HandleMaximum(HloInstruction* instr) override {
     HloInstruction* matmul_call;
     HloInstruction* intermediate_instr = nullptr;
     HloInstruction* optional_bitcast = nullptr;
@@ -699,7 +699,7 @@ class OneDnnMatMulRewriteVisitor : public DfsHloRewriteVisitor {
     return OkStatus();
   }
 
-  Status HandleMultiply(HloInstruction* instr) override {
+  absl::Status HandleMultiply(HloInstruction* instr) override {
     HloInstruction* matmul_call;
     HloInstruction* intermediate_instr = nullptr;
     HloInstruction* src;
@@ -760,10 +760,11 @@ class OneDnnMatMulRewriteVisitor : public DfsHloRewriteVisitor {
     return OkStatus();
   }
 
-  Status FuseActivation(OneDnnMatMulConfig_FusionKind kind,
-                        HloInstruction* activation, HloInstruction* matmul,
-                        HloInstruction* intermediate_instr = nullptr,
-                        HloInstruction* optional_bitcast = nullptr) {
+  absl::Status FuseActivation(OneDnnMatMulConfig_FusionKind kind,
+                              HloInstruction* activation,
+                              HloInstruction* matmul,
+                              HloInstruction* intermediate_instr = nullptr,
+                              HloInstruction* optional_bitcast = nullptr) {
     TF_ASSIGN_OR_RETURN(auto backend_config,
                         matmul->backend_config<BackendConfig>());
     auto* matmul_config = backend_config.mutable_onednn_matmul_config();
@@ -811,7 +812,7 @@ class OneDnnMatMulRewriteVisitor : public DfsHloRewriteVisitor {
   //      lhs:    [batch_dims,contracting_dim] to [batch_dims,1,contracting_dim]
   //      rhs:    [batch_dims,contracting_dim] to [batch_dims,contracting_dim,1]
   //      result: [batch_dims] to [batch_dims,1,1]
-  StatusOr<HloInstruction*> ReconfigureDotDimensions(
+  absl::StatusOr<HloInstruction*> ReconfigureDotDimensions(
       HloInstruction* dot_instr) {
     HloInstruction* lhs = dot_instr->mutable_operand(0);
     HloInstruction* rhs = dot_instr->mutable_operand(1);
@@ -900,7 +901,7 @@ class OneDnnPostRewriteVisitor : public DfsHloRewriteVisitor {
 #endif
   }
 
-  Status HandleCustomCall(HloInstruction* custom_call) override {
+  absl::Status HandleCustomCall(HloInstruction* custom_call) override {
     HloInstruction* matmul;
     if (Match(custom_call, OneDnnMatmulInstr(&matmul))) {
       return HandleCustomCallInternal<dnnl::matmul::primitive_desc>(
@@ -911,7 +912,7 @@ class OneDnnPostRewriteVisitor : public DfsHloRewriteVisitor {
   }
 
   template <typename PrimDesc>
-  Status HandleCustomCallInternal(HloInstruction* custom_call) {
+  absl::Status HandleCustomCallInternal(HloInstruction* custom_call) {
     auto scratch_add = AddScratch<PrimDesc>(custom_call);
     if (scratch_add.ok()) {
       custom_call = *scratch_add;
@@ -926,10 +927,10 @@ class OneDnnPostRewriteVisitor : public DfsHloRewriteVisitor {
   }
 
   template <typename>
-  Status SetWeightsPrepack(HloInstruction*, bool);
+  absl::Status SetWeightsPrepack(HloInstruction*, bool);
 
   template <typename>
-  Status SetUserScratch(HloInstruction*, bool);
+  absl::Status SetUserScratch(HloInstruction*, bool);
 
   template <typename>
   bool GetWeightsPrepack(HloInstruction*);
@@ -940,7 +941,7 @@ class OneDnnPostRewriteVisitor : public DfsHloRewriteVisitor {
   // Add scratch for matmul by changing the result of custom-call to
   // tuple(result, scratch)
   template <typename PrimDesc>
-  StatusOr<HloInstruction*> AddScratch(HloInstruction* custom_call) {
+  absl::StatusOr<HloInstruction*> AddScratch(HloInstruction* custom_call) {
     if (GetUserScratch<PrimDesc>(custom_call)) {
       return custom_call;
     }
@@ -964,7 +965,7 @@ class OneDnnPostRewriteVisitor : public DfsHloRewriteVisitor {
   }
 
   template <typename PrimDesc>
-  StatusOr<HloInstruction*> PrepackWeights(HloInstruction* custom_call) {
+  absl::StatusOr<HloInstruction*> PrepackWeights(HloInstruction* custom_call) {
     if (GetWeightsPrepack<PrimDesc>(custom_call)) {
       return custom_call;
     }
@@ -1042,7 +1043,7 @@ EMIT_GET_BACKEND_CONFIG_SPECIALIZATION(GetWeightsPrepack,
 #define EMIT_SET_BACKEND_CONFIG_SPECIALIZATION(SETTER, PRIM_DESC, CONFIG_TYPE, \
                                                CONFIG, FIELD)                  \
   template <>                                                                  \
-  inline Status OneDnnPostRewriteVisitor::SETTER<PRIM_DESC>(                   \
+  inline absl::Status OneDnnPostRewriteVisitor::SETTER<PRIM_DESC>(             \
       HloInstruction * custom_call, bool value) {                              \
     TF_ASSIGN_OR_RETURN(auto backend_config,                                   \
                         custom_call->backend_config<BackendConfig>());         \
@@ -1060,7 +1061,7 @@ EMIT_SET_BACKEND_CONFIG_SPECIALIZATION(SetUserScratch,
                                        OneDnnMatMulConfig, onednn_matmul_config,
                                        user_scratchpad);
 
-StatusOr<bool> OneDnnMatMulRewriter::Run(
+absl::StatusOr<bool> OneDnnMatMulRewriter::Run(
     HloModule* module,
     const absl::flat_hash_set<absl::string_view>& execution_threads) {
   OneDnnMatMulRewriteVisitor visitor;
diff --git a/third_party/xla/xla/service/cpu/onednn_matmul_rewriter.h b/third_party/xla/xla/service/cpu/onednn_matmul_rewriter.h
index 36cab7ee949c37..eaeea210021666 100644
--- a/third_party/xla/xla/service/cpu/onednn_matmul_rewriter.h
+++ b/third_party/xla/xla/service/cpu/onednn_matmul_rewriter.h
@@ -41,7 +41,7 @@ class OneDnnMatMulRewriter : public HloModulePass {
   absl::string_view name() const override { return "onednn-matmul-rewriter"; }
 
   using HloPassInterface::Run;
-  StatusOr<bool> Run(
+  absl::StatusOr<bool> Run(
       HloModule* module,
       const absl::flat_hash_set<absl::string_view>& execution_threads) override;
 
diff --git a/third_party/xla/xla/service/cpu/onednn_memory_util.cc b/third_party/xla/xla/service/cpu/onednn_memory_util.cc
index fd0b6f927957a0..6ab913161a7e4e 100644
--- a/third_party/xla/xla/service/cpu/onednn_memory_util.cc
+++ b/third_party/xla/xla/service/cpu/onednn_memory_util.cc
@@ -169,7 +169,7 @@ int64_t MemrefInfo::GetChannels() const { return pod_->dims[pod_->rank - 1]; }
 
 int64_t MemrefInfo::GetRank() const { return pod_->rank; }
 
-StatusOr<dnnl::memory::desc> TransposeLastTwoDims(
+absl::StatusOr<dnnl::memory::desc> TransposeLastTwoDims(
     const dnnl::memory::desc& md) {
   int64_t ndims = md.get_ndims();
   if (ndims < 2) {
diff --git a/third_party/xla/xla/service/cpu/onednn_memory_util.h b/third_party/xla/xla/service/cpu/onednn_memory_util.h
index 5793c1bdf4f1ed..c0c956a32dc0b1 100644
--- a/third_party/xla/xla/service/cpu/onednn_memory_util.h
+++ b/third_party/xla/xla/service/cpu/onednn_memory_util.h
@@ -119,7 +119,8 @@ class MemrefInfo {
   MemrefInfoPOD* pod_;
 };
 
-StatusOr<dnnl::memory::desc> TransposeLastTwoDims(const dnnl::memory::desc& md);
+absl::StatusOr<dnnl::memory::desc> TransposeLastTwoDims(
+    const dnnl::memory::desc& md);
 #define TRANSPOSE_LAST_TWO_DIMS_IF(pred, mem_desc)        \
   if (pred) {                                             \
     auto trans_mem_desc = TransposeLastTwoDims(mem_desc); \
diff --git a/third_party/xla/xla/service/cpu/onednn_ops_rewriter.cc b/third_party/xla/xla/service/cpu/onednn_ops_rewriter.cc
index 06c2137a782d50..25530aa59c8d7b 100644
--- a/third_party/xla/xla/service/cpu/onednn_ops_rewriter.cc
+++ b/third_party/xla/xla/service/cpu/onednn_ops_rewriter.cc
@@ -385,7 +385,7 @@ bool MatchFlaxLayerNorm(HloInstruction* instr, HloInstruction** src,
 
 class OneDnnOpsRewriterVisitor : public DfsHloRewriteVisitor {
  public:
-  Status HandleAdd(HloInstruction* instr) override {
+  absl::Status HandleAdd(HloInstruction* instr) override {
     HloInstruction *src, *scale, *bias;
     float eps;
     bool is_bf16orfp16_convert = false;
@@ -445,7 +445,7 @@ class OneDnnOpsRewriterVisitor : public DfsHloRewriteVisitor {
     return OkStatus();
   }
 
-  Status HandleConvert(HloInstruction* instr) override {
+  absl::Status HandleConvert(HloInstruction* instr) override {
     HloInstruction* custom_call;
     HloInstruction* convert_instr;
     auto pattern =
@@ -478,7 +478,7 @@ class OneDnnOpsRewriterVisitor : public DfsHloRewriteVisitor {
     return OkStatus();
   }
 
-  Status HandleDivide(HloInstruction* divide_instr) override {
+  absl::Status HandleDivide(HloInstruction* divide_instr) override {
     if (divide_instr->HasControlDependencies()) return OkStatus();
     if (!IsSupportedType(divide_instr->shape().element_type()))
       return OkStatus();
@@ -495,7 +495,7 @@ class OneDnnOpsRewriterVisitor : public DfsHloRewriteVisitor {
   }
 };
 
-StatusOr<bool> OneDnnOpsRewriter::Run(
+absl::StatusOr<bool> OneDnnOpsRewriter::Run(
     HloModule* module,
     const absl::flat_hash_set<absl::string_view>& execution_threads) {
   OneDnnOpsRewriterVisitor visitor;
diff --git a/third_party/xla/xla/service/cpu/onednn_ops_rewriter.h b/third_party/xla/xla/service/cpu/onednn_ops_rewriter.h
index ea62f33ebcfb99..8e777d8889ad86 100644
--- a/third_party/xla/xla/service/cpu/onednn_ops_rewriter.h
+++ b/third_party/xla/xla/service/cpu/onednn_ops_rewriter.h
@@ -32,7 +32,7 @@ class OneDnnOpsRewriter : public HloModulePass {
   absl::string_view name() const override { return "onednn-ops-rewriter"; }
 
   using HloPassInterface::Run;
-  StatusOr<bool> Run(
+  absl::StatusOr<bool> Run(
       HloModule* module,
       const absl::flat_hash_set<absl::string_view>& execution_threads) override;
 };
diff --git a/third_party/xla/xla/service/cpu/onednn_rewriter.h b/third_party/xla/xla/service/cpu/onednn_rewriter.h
index a1ba3205c96805..53fd5c0f977c48 100644
--- a/third_party/xla/xla/service/cpu/onednn_rewriter.h
+++ b/third_party/xla/xla/service/cpu/onednn_rewriter.h
@@ -33,7 +33,7 @@ class OneDnnRewriter : public HloModulePass {
   absl::string_view name() const override { return "onednn-rewriter"; }
 
   using HloPassInterface::Run;
-  StatusOr<bool> Run(
+  absl::StatusOr<bool> Run(
       HloModule* module,
       const absl::flat_hash_set<absl::string_view>& execution_threads) override;
 };
diff --git a/third_party/xla/xla/service/cpu/parallel_task_assignment.cc b/third_party/xla/xla/service/cpu/parallel_task_assignment.cc
index 6d103759585896..b82707a367e246 100644
--- a/third_party/xla/xla/service/cpu/parallel_task_assignment.cc
+++ b/third_party/xla/xla/service/cpu/parallel_task_assignment.cc
@@ -133,7 +133,8 @@ ParallelTaskAssignment::ParallelTaskAssignment(
   // Run cost analysis on 'module'.
   auto cost_analysis = std::make_unique<HloCostAnalysis>(shape_size);
   HloComputation* computation = module->entry_computation();
-  Status status = computation->root_instruction()->Accept(cost_analysis.get());
+  absl::Status status =
+      computation->root_instruction()->Accept(cost_analysis.get());
   if (status.ok()) {
     // Set default cost model based on 'cost_analysis'.
     cost_model_ = std::make_unique<DefaultCostModel>(

From db6be9037d4fa657413f73694b715581cf84c387 Mon Sep 17 00:00:00 2001
From: David Dunleavy <ddunleavy@google.com>
Date: Thu, 16 May 2024 15:56:08 -0700
Subject: [PATCH 311/478] Remove unnecessary flags from XLA ARM build now that
 hermetic Python has landed

In a followup I will check the disabled ARM tests to see which can be reenabled after this change.

PiperOrigin-RevId: 634555688
---
 third_party/xla/.kokoro/linux/build.sh | 1 -
 1 file changed, 1 deletion(-)

diff --git a/third_party/xla/.kokoro/linux/build.sh b/third_party/xla/.kokoro/linux/build.sh
index dc495910bfc4e5..f9a5e399fd8424 100644
--- a/third_party/xla/.kokoro/linux/build.sh
+++ b/third_party/xla/.kokoro/linux/build.sh
@@ -72,7 +72,6 @@ else
 
     if is_linux_cpu_arm64_job ; then
         TAGS_FILTER="$TAGS_FILTER,-no_aarch64"
-        ADDITIONAL_FLAGS="$ADDITIONAL_FLAGS --action_env PYTHON_BIN_PATH=/usr/bin/python3.11 --python_path=/usr/bin/python3.11"
         # Some cross-compile tests are not working for XLA Linux Aarch64.
         # TODO(ddunleavy): Revisit these when hermetic python is available.
         TARGET_FILTERS="$TARGET_FILTERS -//xla/python_api:xla_shape_test -//xla/python_api:xla_literal_test -//xla/service:xla_aot_compile_stablehlo_cpu_test -//xla/tests:local_client_aot_test"

From ba66234ce939170e07f0656d2f8a5ed9fca9bf0f Mon Sep 17 00:00:00 2001
From: Fangrui Song <maskray@google.com>
Date: Thu, 16 May 2024 16:54:03 -0700
Subject: [PATCH 312/478] Integrate LLVM at llvm/llvm-project@c86a53d75995

Updates LLVM usage to match
[c86a53d75995](https://github.com/llvm/llvm-project/commit/c86a53d75995)

PiperOrigin-RevId: 634569842
---
 third_party/llvm/workspace.bzl | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/third_party/llvm/workspace.bzl b/third_party/llvm/workspace.bzl
index 4aca2a545e9fbf..4dc1a665db1e39 100644
--- a/third_party/llvm/workspace.bzl
+++ b/third_party/llvm/workspace.bzl
@@ -4,8 +4,8 @@ load("//third_party:repo.bzl", "tf_http_archive")
 
 def repo(name):
     """Imports LLVM."""
-    LLVM_COMMIT = "a383b3cca3383128123b33bc73ffc0199ec573e4"
-    LLVM_SHA256 = "f70bda6dd60f40263a377850a2c43556e4e36ed2a30b16c6ccb211b29eab2230"
+    LLVM_COMMIT = "c86a53d75995d65601d5745b2fceb2d50a970645"
+    LLVM_SHA256 = "a62373ace2a6f442b212ae16eced5e0d45d00fb04ae8f94e05dd7712661301ea"
 
     tf_http_archive(
         name = name,

From b1c954c84045289346aa888c4f87e9b506631189 Mon Sep 17 00:00:00 2001
From: "Jae H. Yoo" <jaeyoo@google.com>
Date: Thu, 16 May 2024 18:25:13 -0700
Subject: [PATCH 313/478] Fix the unused return value error in XNNPACK
 weightcache test.

PiperOrigin-RevId: 634592513
---
 tensorflow/lite/delegates/xnnpack/weight_cache_test.cc | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/tensorflow/lite/delegates/xnnpack/weight_cache_test.cc b/tensorflow/lite/delegates/xnnpack/weight_cache_test.cc
index 4cf49e2144cbe4..3d88e1e2dd0f03 100644
--- a/tensorflow/lite/delegates/xnnpack/weight_cache_test.cc
+++ b/tensorflow/lite/delegates/xnnpack/weight_cache_test.cc
@@ -144,7 +144,7 @@ TEST(MMapHandleTest, MapExistingFileWorks) {
 
   TempFileDesc tmp_file;
   ASSERT_TRUE(tmp_file.IsOpen());
-  write(tmp_file.GetFd(), payload.c_str(), size(payload));
+  static_cast<void>(write(tmp_file.GetFd(), payload.c_str(), size(payload)));
   tmp_file.Close();
 
   MMapHandle handle;
@@ -165,7 +165,7 @@ TEST(MMapHandleTest, MoveConstructs) {
 
   TempFileDesc tmp_file;
   ASSERT_TRUE(tmp_file.IsOpen());
-  write(tmp_file.GetFd(), payload.c_str(), size(payload));
+  static_cast<void>(write(tmp_file.GetFd(), payload.c_str(), size(payload)));
   tmp_file.Close();
 
   MMapHandle handle;

From 9fedb4a3ceb10ffcea4b4055c98ac2406a45b2fb Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Thu, 16 May 2024 18:47:39 -0700
Subject: [PATCH 314/478] Update TFRT dependency to use revision
 http://github.com/tensorflow/runtime/commit/edb0d2c6f5e343c83ea121817dc2599ad5453d5c.

PiperOrigin-RevId: 634597659
---
 third_party/tf_runtime/workspace.bzl                          | 4 ++--
 .../xla/third_party/tsl/third_party/tf_runtime/workspace.bzl  | 4 ++--
 2 files changed, 4 insertions(+), 4 deletions(-)

diff --git a/third_party/tf_runtime/workspace.bzl b/third_party/tf_runtime/workspace.bzl
index 2b6fd2f3c8a0d5..04d0e390c8dfe3 100644
--- a/third_party/tf_runtime/workspace.bzl
+++ b/third_party/tf_runtime/workspace.bzl
@@ -6,8 +6,8 @@ def repo():
     """Imports TFRT."""
 
     # Attention: tools parse and update these lines.
-    TFRT_COMMIT = "76dfd9f8757c12739429c8178d1b0c5167862ae7"
-    TFRT_SHA256 = "e607dd6d42ed911021615a7ac5e41ad26190724888e2742544e33bc28a97ff71"
+    TFRT_COMMIT = "edb0d2c6f5e343c83ea121817dc2599ad5453d5c"
+    TFRT_SHA256 = "97f7bfcbff025da3005e59b9ffe1bcb06b439874e3e2cd28a17d9287193d6901"
 
     tf_http_archive(
         name = "tf_runtime",
diff --git a/third_party/xla/third_party/tsl/third_party/tf_runtime/workspace.bzl b/third_party/xla/third_party/tsl/third_party/tf_runtime/workspace.bzl
index 2b6fd2f3c8a0d5..04d0e390c8dfe3 100644
--- a/third_party/xla/third_party/tsl/third_party/tf_runtime/workspace.bzl
+++ b/third_party/xla/third_party/tsl/third_party/tf_runtime/workspace.bzl
@@ -6,8 +6,8 @@ def repo():
     """Imports TFRT."""
 
     # Attention: tools parse and update these lines.
-    TFRT_COMMIT = "76dfd9f8757c12739429c8178d1b0c5167862ae7"
-    TFRT_SHA256 = "e607dd6d42ed911021615a7ac5e41ad26190724888e2742544e33bc28a97ff71"
+    TFRT_COMMIT = "edb0d2c6f5e343c83ea121817dc2599ad5453d5c"
+    TFRT_SHA256 = "97f7bfcbff025da3005e59b9ffe1bcb06b439874e3e2cd28a17d9287193d6901"
 
     tf_http_archive(
         name = "tf_runtime",

From 7414b5ac8f29018ae52814ef520a297fc572d210 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Thu, 16 May 2024 19:28:59 -0700
Subject: [PATCH 315/478] Clean up the build tags for the gpu runner test
 target. Removes this test from *san due to this:
 go/cuda#sanitizers-for-cuda-host-code.

PiperOrigin-RevId: 634604218
---
 tensorflow/core/tfrt/gpu/kernel/BUILD | 7 +++++--
 1 file changed, 5 insertions(+), 2 deletions(-)

diff --git a/tensorflow/core/tfrt/gpu/kernel/BUILD b/tensorflow/core/tfrt/gpu/kernel/BUILD
index 2b79330c050d82..16599c77cbd4e0 100644
--- a/tensorflow/core/tfrt/gpu/kernel/BUILD
+++ b/tensorflow/core/tfrt/gpu/kernel/BUILD
@@ -79,9 +79,12 @@ tf_cuda_cc_test(
     name = "gpu_runner_test",
     srcs = ["gpu_runner_test.cc"],
     tags = [
-        "config-cuda-only",
+        "gpu",  # Only enables test on GPU.
         "no_oss",  # This test only runs with GPU.
-        "requires-gpu-nvidia",
+        "noasan",
+        "nomsan",
+        "noopt",
+        "notsan",
     ],
     deps = [
         ":gpu_runner",

From 68a8100dddd3c1938916f0651c98ce519e59ced9 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Thu, 16 May 2024 19:37:16 -0700
Subject: [PATCH 316/478] Minor change to add logic for finding all lines with
 same id.

PiperOrigin-RevId: 634605755
---
 tensorflow/core/profiler/utils/xplane_utils.h          |  1 +
 .../third_party/tsl/tsl/profiler/utils/xplane_utils.cc | 10 ++++++++++
 .../third_party/tsl/tsl/profiler/utils/xplane_utils.h  |  1 +
 3 files changed, 12 insertions(+)

diff --git a/tensorflow/core/profiler/utils/xplane_utils.h b/tensorflow/core/profiler/utils/xplane_utils.h
index f78350760106b2..75ed0d1b3ed330 100644
--- a/tensorflow/core/profiler/utils/xplane_utils.h
+++ b/tensorflow/core/profiler/utils/xplane_utils.h
@@ -27,6 +27,7 @@ namespace profiler {
 
 using tsl::profiler::AddFlowsToXplane;               // NOLINT
 using tsl::profiler::AggregateXPlane;                // NOLINT
+using tsl::profiler::FindLinesWithId;                // NOLINT
 using tsl::profiler::FindLineWithId;                 // NOLINT
 using tsl::profiler::FindLineWithName;               // NOLINT
 using tsl::profiler::FindMutablePlanes;              // NOLINT
diff --git a/third_party/xla/third_party/tsl/tsl/profiler/utils/xplane_utils.cc b/third_party/xla/third_party/tsl/tsl/profiler/utils/xplane_utils.cc
index 88c7e30b76eee5..3e50fee7f7ce4a 100644
--- a/third_party/xla/third_party/tsl/tsl/profiler/utils/xplane_utils.cc
+++ b/third_party/xla/third_party/tsl/tsl/profiler/utils/xplane_utils.cc
@@ -224,6 +224,16 @@ const XLine* FindLineWithId(const XPlane& plane, int64_t id) {
       Find(plane.lines(), [id](const XLine* line) { return line->id() == id; });
   return (i != -1) ? &plane.lines(i) : nullptr;
 }
+std::vector<const XLine*> FindLinesWithId(const XPlane& plane, int64_t id) {
+  std::vector<int> indices = FindAll(
+      plane.lines(), [id](const XLine* line) { return line->id() == id; });
+  std::vector<const XLine*> lines;
+  lines.reserve(indices.size());
+  for (int index : indices) {
+    lines.push_back(&plane.lines(index));
+  }
+  return lines;
+}
 
 const XLine* FindLineWithName(const XPlane& plane, absl::string_view name) {
   int i = Find(plane.lines(),
diff --git a/third_party/xla/third_party/tsl/tsl/profiler/utils/xplane_utils.h b/third_party/xla/third_party/tsl/tsl/profiler/utils/xplane_utils.h
index d0177632263455..8ea1429c1d90d2 100644
--- a/third_party/xla/third_party/tsl/tsl/profiler/utils/xplane_utils.h
+++ b/third_party/xla/third_party/tsl/tsl/profiler/utils/xplane_utils.h
@@ -83,6 +83,7 @@ std::vector<XPlane*> FindMutablePlanesWithPrefix(XSpace* space,
 
 // Returns the plane with the given id/name or nullptr if not found.
 const XLine* FindLineWithId(const XPlane& plane, int64_t id);
+std::vector<const XLine*> FindLinesWithId(const XPlane& plane, int64_t id);
 const XLine* FindLineWithName(const XPlane& plane, absl::string_view name);
 
 XStat* FindOrAddMutableStat(const XStatMetadata& stat_metadata, XEvent* event);

From ce5b30c9bf49db4f8fffc07352e945e6546ad276 Mon Sep 17 00:00:00 2001
From: Doyeon Kim <doyeonkim@google.com>
Date: Thu, 16 May 2024 20:24:39 -0700
Subject: [PATCH 317/478] Enable weight-only quantization with StableHLO opset
 in TF Quantizer

This CL integrates StableHLO weight-only quantization preset to TF Quantizer API. The weight-only quantization is enabled when the quantization method is set to `METHOD_STATIC_RANGE_WEIGHT_ONLY_INT8` and opset is set to `quant_opts_pb2.STABLEHLO`.

PiperOrigin-RevId: 634614727
---
 .../mlir/quantization/stablehlo/cc/config.cc  |  1 -
 .../stablehlo/quantization_config.proto       |  4 +-
 .../mlir/quantization/tensorflow/python/BUILD |  1 +
 .../integration_test/quantize_model_test.py   | 18 +++++++-
 .../tensorflow/python/quantize_model.cc       | 46 +++++++++++++++----
 .../tensorflow/python/quantize_model.py       |  4 +-
 6 files changed, 59 insertions(+), 15 deletions(-)

diff --git a/tensorflow/compiler/mlir/quantization/stablehlo/cc/config.cc b/tensorflow/compiler/mlir/quantization/stablehlo/cc/config.cc
index 145533ddb07cc9..1522c68f300cba 100644
--- a/tensorflow/compiler/mlir/quantization/stablehlo/cc/config.cc
+++ b/tensorflow/compiler/mlir/quantization/stablehlo/cc/config.cc
@@ -176,7 +176,6 @@ void ExpandWeightOnlyPtqPreset(QuantizationConfig& config) {
   // user-provided explicit `QuantizationSpec`s will be appended.
   QuantizationSpecs new_specs{};
   *new_specs.add_specs() = GetDefaultWeightOnlyPtqSpec();
-  // TODO: b/307625297 - Add per-channel weight only support.
 
   // Append user-provided specs to override existing specs.
   const QuantizationSpecs& previous_specs = config.specs();
diff --git a/tensorflow/compiler/mlir/quantization/stablehlo/quantization_config.proto b/tensorflow/compiler/mlir/quantization/stablehlo/quantization_config.proto
index a307e1d927eea3..49e8161df3a749 100644
--- a/tensorflow/compiler/mlir/quantization/stablehlo/quantization_config.proto
+++ b/tensorflow/compiler/mlir/quantization/stablehlo/quantization_config.proto
@@ -77,8 +77,8 @@ message StaticRangePtqPreset {
   bool enable_full_int_quantization = 3;
 }
 
-// Applies int8 per-tensor weight-only post-training quantization for all
-// dot_general op.
+// Applies int8 per-channel weight-only post-training quantization for all
+// dot_general and convolution ops.
 message WeightOnlyPtqPreset {}
 
 // Metadata specific to the input TensorFlow SavedModel, which may be required
diff --git a/tensorflow/compiler/mlir/quantization/tensorflow/python/BUILD b/tensorflow/compiler/mlir/quantization/tensorflow/python/BUILD
index 6630119fbaaac1..c0a472ca8f2e26 100644
--- a/tensorflow/compiler/mlir/quantization/tensorflow/python/BUILD
+++ b/tensorflow/compiler/mlir/quantization/tensorflow/python/BUILD
@@ -45,6 +45,7 @@ cc_library(
         "//tensorflow/compiler/mlir/quantization/stablehlo/cc:saved_model_export",
         "//tensorflow/compiler/mlir/quantization/stablehlo/cc:saved_model_import",
         "//tensorflow/compiler/mlir/quantization/stablehlo/cc:types",
+        "//tensorflow/compiler/mlir/quantization/stablehlo/cc:weight_only_ptq",
         "//tensorflow/compiler/mlir/quantization/stablehlo/cc/calibration:component",
         "//tensorflow/compiler/mlir/quantization/stablehlo/cc/calibration:statistics",
         "//tensorflow/compiler/mlir/quantization/tensorflow:exported_model_proto_cc",
diff --git a/tensorflow/compiler/mlir/quantization/tensorflow/python/integration_test/quantize_model_test.py b/tensorflow/compiler/mlir/quantization/tensorflow/python/integration_test/quantize_model_test.py
index 53228d30bd53fc..b44c788bc10f7b 100644
--- a/tensorflow/compiler/mlir/quantization/tensorflow/python/integration_test/quantize_model_test.py
+++ b/tensorflow/compiler/mlir/quantization/tensorflow/python/integration_test/quantize_model_test.py
@@ -5501,6 +5501,7 @@ def test_einsum_model(
 
   @parameterized.named_parameters(
       ('to_xla_per_tensor', quant_opts_pb2.XLA, False),
+      ('stablehlo_per_channel', quant_opts_pb2.STABLEHLO, True),
   )
   @test_util.run_in_graph_and_eager_modes
   def test_matmul_model(
@@ -5542,8 +5543,14 @@ def test_matmul_model(
     )
     output_graphdef = output_loader.get_meta_graph_def_from_tags(tags).graph_def
 
+    if target_opset == quant_opts_pb2.XLA:
+      self.assertTrue(self._contains_op(output_graphdef, 'XlaDotV2'))
+    elif target_opset == quant_opts_pb2.STABLEHLO:
+      # This is to verify the invocation of StableHLO quantizer works. More
+      # thorough functional tests are in StableHLO quantizer directory.
+      self.assertTrue(self._contains_op(output_graphdef, 'XlaCallModule'))
+
     # Due to other meta data, the compression is not exactly 1/4.
-    self.assertTrue(self._contains_op(output_graphdef, 'XlaDotV2'))
     self.assertLess(
         testing.get_size_ratio(
             self._output_saved_model_path, self._input_saved_model_path
@@ -5553,6 +5560,7 @@ def test_matmul_model(
 
   @parameterized.named_parameters(
       ('to_xla_per_tensor', quant_opts_pb2.XLA, False),
+      ('stablehlo_per_channel', quant_opts_pb2.STABLEHLO, True),
       # TODO: b/289761265 - [Converter Component][TF-Quantizer] Improve Weight-
       # only Quantization
       # Enable this back once new weight-only quantizer is supported for per-
@@ -5612,7 +5620,7 @@ def test_conv_model(
         0.3,
     )
 
-    if enable_per_channel_quantization:
+    if enable_per_channel_quantization and target_opset == quant_opts_pb2.XLA:
       per_channel_size_attr = attr_value_pb2.AttrValue(
           list=attr_value_pb2.AttrValue.ListValue(
               shape=[
@@ -5631,6 +5639,12 @@ def test_conv_model(
               output_graphdef, 'Const', '_output_shapes', per_channel_size_attr
           )
       )
+    if target_opset == quant_opts_pb2.XLA:
+      self.assertTrue(self._contains_op(output_graphdef, 'XlaConvV2'))
+    elif target_opset == quant_opts_pb2.STABLEHLO:
+      # This is to verify the invocation of StableHLO quantizer works. More
+      # thorough functional tests are in StableHLO quantizer directory.
+      self.assertTrue(self._contains_op(output_graphdef, 'XlaCallModule'))
 
     input_tensor = array_ops.constant(
         np.random.uniform(low=0, high=0.1, size=input_shape),
diff --git a/tensorflow/compiler/mlir/quantization/tensorflow/python/quantize_model.cc b/tensorflow/compiler/mlir/quantization/tensorflow/python/quantize_model.cc
index aef2732549ede2..e38310879184ef 100644
--- a/tensorflow/compiler/mlir/quantization/tensorflow/python/quantize_model.cc
+++ b/tensorflow/compiler/mlir/quantization/tensorflow/python/quantize_model.cc
@@ -47,6 +47,7 @@ limitations under the License.
 #include "tensorflow/compiler/mlir/quantization/stablehlo/cc/saved_model_export.h"
 #include "tensorflow/compiler/mlir/quantization/stablehlo/cc/saved_model_import.h"
 #include "tensorflow/compiler/mlir/quantization/stablehlo/cc/types.h"
+#include "tensorflow/compiler/mlir/quantization/stablehlo/cc/weight_only_ptq.h"
 #include "tensorflow/compiler/mlir/quantization/stablehlo/quantization_config.pb.h"
 #include "tensorflow/compiler/mlir/quantization/tensorflow/cc/convert_asset_args.h"
 #include "tensorflow/compiler/mlir/quantization/tensorflow/cc/run_passes.h"
@@ -80,6 +81,7 @@ using ::mlir::quant::stablehlo::PostCalibrationComponent;
 using ::mlir::quant::stablehlo::PreCalibrationComponent;
 using ::mlir::quant::stablehlo::RunCalibrationPasses;
 using ::mlir::quant::stablehlo::UpdateFunctionAliases;
+using ::mlir::quant::stablehlo::WeightOnlyPtqComponent;
 using ::stablehlo::quantization::AddCalibrationStatistics;
 using ::stablehlo::quantization::ChangeToQuantizedFilename;
 using ::stablehlo::quantization::DebuggerConfig;
@@ -223,6 +225,20 @@ QuantizationConfig GetQuantizationConfigForStaticRangePtq(
   return ExpandPresets(PopulateDefaults(quantization_config));
 }
 
+QuantizationConfig GetQuantizationConfigForWeightOnlyPtq(
+    const QuantizationOptions &quantization_options) {
+  QuantizationConfig quantization_config{};
+  quantization_config.mutable_weight_only_ptq_preset();
+  // When targeting server TPUs quantized types should be unpacked into
+  // integer ops.
+  quantization_config.mutable_pipeline_config()->set_unpack_quantized_types(
+      true);
+  *quantization_config.mutable_debugger_config() =
+      quantization_options.debugger_config();
+
+  return ExpandPresets(PopulateDefaults(quantization_config));
+}
+
 absl::StatusOr<ExportedModel> QuantizePtqModelPreCalibrationImpl(
     mlir::ModuleOp module_op, mlir::MLIRContext *context,
     const QuantizationOptions &quantization_options,
@@ -385,6 +401,7 @@ absl::StatusOr<ExportedModel> QuantizeWeightOnly(
         "Failed to get function alias: ", function_aliases.status().message()));
   }
 
+  const bool is_stablehlo = quantization_options.op_set() == OpSet::STABLEHLO;
   absl::StatusOr<mlir::OwningOpRef<mlir::ModuleOp>> module =
       ImportAndPreprocessSavedModel(
           saved_model_path,
@@ -392,7 +409,8 @@ absl::StatusOr<ExportedModel> QuantizeWeightOnly(
            quantization_options.signature_keys().end()},
           {quantization_options.tags().begin(),
            quantization_options.tags().end()},
-          context.get(), /*is_inliner_run=*/true, /*run_tf_to_stablehlo=*/false,
+          context.get(), /*is_inliner_run=*/true,
+          /*run_tf_to_stablehlo=*/is_stablehlo,
           /*deserialize_xla_call_module=*/false, *function_aliases);
   if (!module.status().ok()) {
     return absl::InternalError(
@@ -401,14 +419,24 @@ absl::StatusOr<ExportedModel> QuantizeWeightOnly(
   }
   mlir::OwningOpRef<mlir::ModuleOp> module_ref = std::move(module).value();
 
-  TF_RETURN_IF_ERROR(RunPasses(
-      kTfQuantWeightOnlyStepName,
-      /*add_passes_func=*/
-      [&quantization_options](mlir::PassManager &pm) {
-        AddQuantizeWeightOnlyPasses(pm, quantization_options,
-                                    kTfQuantWeightOnlyStepName);
-      },
-      *context, *module_ref));
+  // Use StableHLO Quantizer option if opset is specified.
+  if (is_stablehlo) {
+    const QuantizationConfig quantization_config =
+        GetQuantizationConfigForWeightOnlyPtq(quantization_options);
+
+    WeightOnlyPtqComponent weight_only_ptq_component(context.get());
+    TF_ASSIGN_OR_RETURN(*module_ref, weight_only_ptq_component.Run(
+                                         *module_ref, quantization_config));
+  } else {
+    TF_RETURN_IF_ERROR(RunPasses(
+        kTfQuantWeightOnlyStepName,
+        /*add_passes_func=*/
+        [&quantization_options](mlir::PassManager &pm) {
+          AddQuantizeWeightOnlyPasses(pm, quantization_options,
+                                      kTfQuantWeightOnlyStepName);
+        },
+        *context, *module_ref));
+  }
 
   return ModuleOpToExportedModel(
       *module_ref, context.get(), kTfQuantWeightOnlyStepName,
diff --git a/tensorflow/compiler/mlir/quantization/tensorflow/python/quantize_model.py b/tensorflow/compiler/mlir/quantization/tensorflow/python/quantize_model.py
index 91c82710c489a9..f7dec2d2a5dee7 100644
--- a/tensorflow/compiler/mlir/quantization/tensorflow/python/quantize_model.py
+++ b/tensorflow/compiler/mlir/quantization/tensorflow/python/quantize_model.py
@@ -741,10 +741,12 @@ def _populate_quantization_options_default_values(
   if (quantization_options.op_set == quant_opts_pb2.OpSet.STABLEHLO) and (
       quantization_options.quantization_method.preset_method
       != _PresetMethod.METHOD_STATIC_RANGE_INT8
+      and quantization_options.quantization_method.preset_method
+      != _PresetMethod.METHOD_STATIC_RANGE_WEIGHT_ONLY_INT8
   ):
     raise ValueError(
         'StableHLO quantized opset currently only supports static range'
-        ' quantization via TF Quantizer.'
+        ' quantization and weight-only quantizationvia TF Quantizer.'
     )
 
   # Set `force_graph_mode_calibration` to True to avoid skipping op execution,

From 5005bae4baa0fa27502f902a85312f7b5ee7d829 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Thu, 16 May 2024 22:08:30 -0700
Subject: [PATCH 318/478] Open source pybind lib for populating sparse core
 layouts.

PiperOrigin-RevId: 634632061
---
 tensorflow/BUILD                              |  1 +
 tensorflow/core/tpu/kernels/BUILD             | 15 +++++++
 tensorflow/python/BUILD                       |  3 ++
 tensorflow/python/tpu/BUILD                   | 21 +++++++++
 .../python/tpu/_pywrap_sparse_core_layout.pyi | 24 +++++++++++
 .../python/tpu/pywrap_sparse_core_layout.cc   | 43 +++++++++++++++++++
 .../tools/def_file_filter/symbols_pybind.txt  |  6 +++
 .../tools/def_file_filter/symbols_pybind.txt  |  6 +++
 8 files changed, 119 insertions(+)
 create mode 100644 tensorflow/python/tpu/_pywrap_sparse_core_layout.pyi
 create mode 100644 tensorflow/python/tpu/pywrap_sparse_core_layout.cc

diff --git a/tensorflow/BUILD b/tensorflow/BUILD
index 71487e2aec0bee..0b2455baa5f7f3 100644
--- a/tensorflow/BUILD
+++ b/tensorflow/BUILD
@@ -1416,6 +1416,7 @@ tf_cc_shared_library(
         "//tensorflow/core/grappler:grappler_item_builder",
         "//tensorflow/core/kernels:data_service_ops",
         "//tensorflow/core/kernels:dataset_ops",
+        "//tensorflow/core/tpu/kernels:sparse_core_layout",
         "//tensorflow/core/platform:logging",
         "//tensorflow/core/platform:path",
         "//tensorflow/core/platform:stacktrace_handler",
diff --git a/tensorflow/core/tpu/kernels/BUILD b/tensorflow/core/tpu/kernels/BUILD
index 50334814c993b3..7d55de7bbc1de3 100644
--- a/tensorflow/core/tpu/kernels/BUILD
+++ b/tensorflow/core/tpu/kernels/BUILD
@@ -1481,6 +1481,21 @@ cc_library(
     visibility = ["//visibility:public"],
 )
 
+cc_library(
+    name = "_pywrap_sparse_core_layout_header_only",
+    srcs = [],
+    hdrs = ["sparse_core_layout.h"],
+    visibility = ["//tensorflow/python/tpu:__pkg__"],  # ONLY for `_pywrap_sparse_core_layout`.
+    deps = [
+        ":sparse_core_layout_proto_cc",
+        "//tensorflow/core/platform:stringpiece",
+        "@com_google_absl//absl/log",
+        "@com_google_absl//absl/log:check",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/status:statusor",
+    ],
+)
+
 cc_library(
     name = "sparse_core_layout",
     srcs = ["sparse_core_layout.cc"],
diff --git a/tensorflow/python/BUILD b/tensorflow/python/BUILD
index 571881111137af..b9a762c607b4eb 100644
--- a/tensorflow/python/BUILD
+++ b/tensorflow/python/BUILD
@@ -285,6 +285,7 @@ py_strict_library(
         "//tensorflow/python/summary/writer",
         "//tensorflow/python/tools:module_util",
         "//tensorflow/python/tools/api/generator:create_python_api",
+        "//tensorflow/python/tpu:_pywrap_sparse_core_layout",
         "//tensorflow/python/tpu:datasets",
         "//tensorflow/python/tpu:functional",
         "//tensorflow/python/tpu:preempted_hook_py",
@@ -799,6 +800,7 @@ pywrap_tensorflow_macro(
         "//tensorflow/core/platform:stacktrace_handler",
         "//tensorflow/core/profiler:profiler_impl",
         "//tensorflow/core/profiler/internal:print_model_analysis",
+        "//tensorflow/core/tpu/kernels:sparse_core_layout",
         "//tensorflow/core/util:determinism",
         "//tensorflow/distribute/experimental/rpc/kernels:rpc_ops",
         "//tensorflow/dtensor/cc:dtensor_device_cc",
@@ -920,6 +922,7 @@ filegroup(
         "//tensorflow/core/platform:statusor",  # tfe
         "//tensorflow/core/profiler/internal:print_model_analysis",  # tfprof
         "//tensorflow/core/profiler/rpc/client:profiler_client_impl",  # profiler
+        "//tensorflow/core/tpu/kernels:sparse_core_layout",  # sparse_core_layouts
         "//tensorflow/core/util:determinism",  # determinism
         "//tensorflow/core/util:port",  # util_port
         "//tensorflow/core/util/tensor_bundle",  # checkpoint_reader
diff --git a/tensorflow/python/tpu/BUILD b/tensorflow/python/tpu/BUILD
index 108e643b1e4ea9..c9c2b7f75a0710 100644
--- a/tensorflow/python/tpu/BUILD
+++ b/tensorflow/python/tpu/BUILD
@@ -1015,6 +1015,26 @@ tf_proto_library(
 # )
 # copybara:uncomment_end
 
+tf_python_pybind_extension(
+    name = "_pywrap_sparse_core_layout",
+    srcs = ["pywrap_sparse_core_layout.cc"],
+    enable_stub_generation = True,
+    features = ["-layering_check"],
+    pytype_srcs = [
+        "_pywrap_sparse_core_layout.pyi",
+    ],
+    deps = [
+        "//tensorflow/core/tpu/kernels:_pywrap_sparse_core_layout_header_only",
+        "//tensorflow/python/lib/core:pybind11_lib",
+        "//tensorflow/python/lib/core:pybind11_status",
+        "//tensorflow/python/lib/core:pybind11_status_headers",
+        "//third_party/python_runtime:headers",  # buildcleaner: keep
+        "@pybind11",
+        "@pybind11_abseil//pybind11_abseil:status_casters",
+        "@pybind11_protobuf//pybind11_protobuf:native_proto_caster",
+    ],
+)
+
 tf_python_pybind_extension(
     name = "_pywrap_tpu_embedding",
     srcs = ["pywrap_tpu_embedding.cc"],
@@ -1039,6 +1059,7 @@ py_strict_library(
     name = "tpu_embedding_v3",
     srcs = ["tpu_embedding_v3.py"],
     deps = [
+        ":_pywrap_sparse_core_layout",
         ":_pywrap_tpu_embedding",
         ":tpu_embedding_base",
         ":tpu_embedding_v2_utils",
diff --git a/tensorflow/python/tpu/_pywrap_sparse_core_layout.pyi b/tensorflow/python/tpu/_pywrap_sparse_core_layout.pyi
new file mode 100644
index 00000000000000..778093406836bc
--- /dev/null
+++ b/tensorflow/python/tpu/_pywrap_sparse_core_layout.pyi
@@ -0,0 +1,24 @@
+# Copyright 2023 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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.
+# ==============================================================================
+
+from typing import Any
+
+class SparseCoreLayoutStacker:
+    def __init__(self, num_partitions: int, sparse_cores_per_partition: int) -> None: ...
+    def AddTable(self, table_name: str, table_height: int, table_width: int, group: str, output_samples: int) -> None: ...
+    def GetLayouts(self, *args, **kwargs) -> Any: ...
+    def SetActivationMemoryBytesLimit(self, arg0: int) -> None: ...
+    def SetStackingEnabled(self, arg0: bool) -> None: ...
+    def SetVariableShardBytesLimit(self, arg0: int) -> None: ...
diff --git a/tensorflow/python/tpu/pywrap_sparse_core_layout.cc b/tensorflow/python/tpu/pywrap_sparse_core_layout.cc
new file mode 100644
index 00000000000000..785efe5d682a25
--- /dev/null
+++ b/tensorflow/python/tpu/pywrap_sparse_core_layout.cc
@@ -0,0 +1,43 @@
+/* Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+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.
+==============================================================================*/
+#include <pybind11/pybind11.h>
+
+#include "pybind11/cast.h"  // from @pybind11
+#include "pybind11/detail/common.h"  // from @pybind11
+#include "pybind11_abseil/status_casters.h"  // from @pybind11_abseil
+#include "pybind11_protobuf/native_proto_caster.h"  // from @pybind11_protobuf
+#include "tensorflow/core/tpu/kernels/sparse_core_layout.h"
+#include "tensorflow/core/tpu/kernels/sparse_core_layout.pb.h"
+
+namespace tensorflow::tpu {
+
+namespace py = pybind11;
+
+PYBIND11_MODULE(_pywrap_sparse_core_layout, m) {
+  py::class_<SparseCoreLayoutStacker>(m, "SparseCoreLayoutStacker")
+      .def(py::init<int, int>(), py::arg("num_partitions"),
+           py::arg("sparse_cores_per_partition"))
+      .def("SetActivationMemoryBytesLimit",
+           &SparseCoreLayoutStacker::SetActivationMemoryBytesLimit)
+      .def("SetVariableShardBytesLimit",
+           &SparseCoreLayoutStacker::SetVariableShardBytesLimit)
+      .def("SetStackingEnabled", &SparseCoreLayoutStacker::SetStackingEnabled)
+      .def("AddTable", &SparseCoreLayoutStacker::AddTable,
+           py::arg("table_name"), py::arg("table_height"),
+           py::arg("table_width"), py::arg("group"), py::arg("output_samples"))
+      .def("GetLayouts", &SparseCoreLayoutStacker::GetLayouts);
+}
+
+}  // namespace tensorflow::tpu
diff --git a/tensorflow/tools/def_file_filter/symbols_pybind.txt b/tensorflow/tools/def_file_filter/symbols_pybind.txt
index 78c42c6f454c5c..c9faf5ad6b8375 100644
--- a/tensorflow/tools/def_file_filter/symbols_pybind.txt
+++ b/tensorflow/tools/def_file_filter/symbols_pybind.txt
@@ -590,3 +590,9 @@ tensorflow::dtensor::SetIteratorElementLayouts
 tensorflow::dtensor::Mesh
 tensorflow::dtensor::Layout
 tensorflow::dtensor::Layout::LayoutType
+
+[//tensorflow/core/tpu/kernels:sparse_core_layout] #SparseCoreLayoutStacker
+tensorflow::tpu::SparseCoreLayoutStacker::AddTable
+tensorflow::tpu::SparseCoreLayoutStacker::SparseCoreLayoutStacker
+tensorflow::tpu::SparseCoreLayoutStacker::AddTable
+tensorflow::tpu::SparseCoreLayoutStacker::GetLayouts
\ No newline at end of file
diff --git a/third_party/xla/third_party/tsl/tools/def_file_filter/symbols_pybind.txt b/third_party/xla/third_party/tsl/tools/def_file_filter/symbols_pybind.txt
index 78c42c6f454c5c..c9faf5ad6b8375 100644
--- a/third_party/xla/third_party/tsl/tools/def_file_filter/symbols_pybind.txt
+++ b/third_party/xla/third_party/tsl/tools/def_file_filter/symbols_pybind.txt
@@ -590,3 +590,9 @@ tensorflow::dtensor::SetIteratorElementLayouts
 tensorflow::dtensor::Mesh
 tensorflow::dtensor::Layout
 tensorflow::dtensor::Layout::LayoutType
+
+[//tensorflow/core/tpu/kernels:sparse_core_layout] #SparseCoreLayoutStacker
+tensorflow::tpu::SparseCoreLayoutStacker::AddTable
+tensorflow::tpu::SparseCoreLayoutStacker::SparseCoreLayoutStacker
+tensorflow::tpu::SparseCoreLayoutStacker::AddTable
+tensorflow::tpu::SparseCoreLayoutStacker::GetLayouts
\ No newline at end of file

From 41361300a4e865613908025d6bdc392df0c0760e Mon Sep 17 00:00:00 2001
From: Ionel Gog <icgog@google.com>
Date: Thu, 16 May 2024 22:18:36 -0700
Subject: [PATCH 319/478] [IFRT] Use llvm::DenseSet instead of llvm::SmallSet
 when verifying devices.

PiperOrigin-RevId: 634633826
---
 third_party/xla/xla/python/ifrt/ir/ifrt_dialect.cc | 3 +--
 third_party/xla/xla/python/ifrt/ir/ifrt_ops.cc     | 7 ++++---
 2 files changed, 5 insertions(+), 5 deletions(-)

diff --git a/third_party/xla/xla/python/ifrt/ir/ifrt_dialect.cc b/third_party/xla/xla/python/ifrt/ir/ifrt_dialect.cc
index 3c0c044d5ee8e1..265c4deea601fc 100644
--- a/third_party/xla/xla/python/ifrt/ir/ifrt_dialect.cc
+++ b/third_party/xla/xla/python/ifrt/ir/ifrt_dialect.cc
@@ -20,7 +20,6 @@ limitations under the License.
 #include "absl/status/statusor.h"
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/STLFunctionalExtras.h"
-#include "llvm/ADT/SmallSet.h"
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/ADT/TypeSwitch.h"
 #include "llvm/Support/Casting.h"
@@ -225,7 +224,7 @@ IfrtDevicesAttr::operator llvm::ArrayRef<int>() const { return getIds(); }
 mlir::LogicalResult IfrtDevicesAttr::verify(
     llvm::function_ref<::mlir::InFlightDiagnostic()> emitError,
     llvm::ArrayRef<int> ids) {
-  llvm::SmallSet<int, 4> device_set;
+  llvm::DenseSet<int> device_set;
   for (int id : ids) {
     if (id < 0) {
       return emitError() << "Device list has negative logical id " << id;
diff --git a/third_party/xla/xla/python/ifrt/ir/ifrt_ops.cc b/third_party/xla/xla/python/ifrt/ir/ifrt_ops.cc
index b268243fde98f5..ab6468d518c005 100644
--- a/third_party/xla/xla/python/ifrt/ir/ifrt_ops.cc
+++ b/third_party/xla/xla/python/ifrt/ir/ifrt_ops.cc
@@ -23,6 +23,7 @@ limitations under the License.
 
 #include "absl/status/statusor.h"
 #include "llvm/ADT/ArrayRef.h"
+#include "llvm/ADT/DenseSet.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/SmallSet.h"
 #include "llvm/ADT/SmallVector.h"
@@ -146,7 +147,7 @@ mlir::LogicalResult VerifyDevicePlacement(
     mlir::Operation* op, llvm::ArrayRef<int> devices,
     llvm::ArrayRef<IfrtArrayType> inputs,
     llvm::ArrayRef<IfrtArrayType> outputs) {
-  llvm::SmallSet<int, 4> device_set;
+  llvm::DenseSet<int> device_set;
   device_set.insert(devices.begin(), devices.end());
 
   for (const IfrtArrayType input : inputs) {
@@ -239,7 +240,7 @@ mlir::LogicalResult ReshardOp::verify() {
 }
 
 mlir::LogicalResult AssembleOp::verify() {
-  llvm::SmallVector<int, 4> input_devices;
+  std::vector<int> input_devices;
   for (const mlir::Value input : getInputs()) {
     const auto array = llvm::cast<IfrtArrayType>(input.getType());
     if (array.getDevices().size() != 1) {
@@ -259,7 +260,7 @@ mlir::LogicalResult AssembleOp::verify() {
 }
 
 mlir::LogicalResult DisassembleOp::verify() {
-  llvm::SmallVector<int, 4> output_devices;
+  std::vector<int> output_devices;
   for (const mlir::Value output : getOutputs()) {
     const auto array = llvm::cast<IfrtArrayType>(output.getType());
     if (array.getDevices().size() != 1) {

From d8237f6576a783724dfe85e6678227047e754e11 Mon Sep 17 00:00:00 2001
From: Henning Becker <hebecker@google.com>
Date: Thu, 16 May 2024 22:39:00 -0700
Subject: [PATCH 320/478] Move variable declaration into preprocessor branch

This avoid warnings about unused variables when the ROCm branch is taken.

PiperOrigin-RevId: 634637222
---
 third_party/xla/xla/service/gpu/make_batch_pointers.cc | 6 ++----
 1 file changed, 2 insertions(+), 4 deletions(-)

diff --git a/third_party/xla/xla/service/gpu/make_batch_pointers.cc b/third_party/xla/xla/service/gpu/make_batch_pointers.cc
index f0de161a35e3ce..46e06e7ddfde07 100644
--- a/third_party/xla/xla/service/gpu/make_batch_pointers.cc
+++ b/third_party/xla/xla/service/gpu/make_batch_pointers.cc
@@ -48,16 +48,14 @@ absl::Status MakeBatchPointers(se::Stream* stream,
                                se::DeviceMemoryBase base_ptr,
                                size_t stride_bytes, size_t n,
                                se::DeviceMemoryBase ptrs_out) {
-  static constexpr size_t kThreads = 128;
-
-  se::StreamExecutor* executor = stream->parent();
-
 #if TENSORFLOW_USE_ROCM
   stream_executor::gpu::rocm_MakeBatchPointers(
       se::gpu::AsGpuStreamValue(stream),
       reinterpret_cast<char*>(base_ptr.opaque()), stride_bytes, n,
       reinterpret_cast<void**>(ptrs_out.opaque()));
 #else
+  se::StreamExecutor* executor = stream->parent();
+  static constexpr size_t kThreads = 128;
 
   TF_ASSIGN_OR_RETURN(
       auto kernel,

From 55d75e94e0cd4e1d7d14a5893b006bbe180968cc Mon Sep 17 00:00:00 2001
From: Adrian Kuegel <akuegel@google.com>
Date: Thu, 16 May 2024 23:03:21 -0700
Subject: [PATCH 321/478] Normalize layouts for Scatter.

Unfortunately this requires to temporarily use unsorted update_window_dims
attribute. ScatterSimplifier can turn this into valid HLO again.

PiperOrigin-RevId: 634641295
---
 third_party/xla/xla/service/BUILD             |   5 +
 .../xla/xla/service/gather_scatter_utils.cc   |   5 +
 .../xla/xla/service/gather_simplifier_test.cc |   4 +-
 .../xla/xla/service/gpu/gpu_compiler.cc       |   8 ++
 .../xla/xla/service/layout_normalization.cc   | 107 +++++++++++++++++
 .../xla/service/layout_normalization_test.cc  | 109 +++++++++++++++++-
 .../xla/service/scatter_simplifier_test.cc    |   8 +-
 7 files changed, 237 insertions(+), 9 deletions(-)

diff --git a/third_party/xla/xla/service/BUILD b/third_party/xla/xla/service/BUILD
index e9c6d033688b6d..7df3a319664561 100644
--- a/third_party/xla/xla/service/BUILD
+++ b/third_party/xla/xla/service/BUILD
@@ -7468,7 +7468,11 @@ xla_cc_test(
     ],
     deps = [
         ":layout_normalization",
+        ":scatter_simplifier",
+        "//xla/hlo/ir:hlo",
         "//xla/tests:hlo_test_base",
+        "@com_google_absl//absl/strings:string_view",
+        "@local_tsl//tsl/platform:status",
         "@local_tsl//tsl/platform:test",
         "@local_tsl//tsl/platform:test_main",
     ],
@@ -7493,6 +7497,7 @@ cc_library(
     deps = [
         ":hlo_creation_utils",
         "//xla:permutation_util",
+        "//xla:shape_util",
         "//xla/hlo/ir:hlo",
         "@com_google_absl//absl/types:span",
     ],
diff --git a/third_party/xla/xla/service/gather_scatter_utils.cc b/third_party/xla/xla/service/gather_scatter_utils.cc
index 223879d4936fe2..897f1e7429b50e 100644
--- a/third_party/xla/xla/service/gather_scatter_utils.cc
+++ b/third_party/xla/xla/service/gather_scatter_utils.cc
@@ -19,6 +19,7 @@ limitations under the License.
 #include <utility>
 #include <vector>
 
+#include "xla/layout_util.h"
 #include "xla/permutation_util.h"
 #include "xla/service/hlo_creation_utils.h"
 
@@ -68,6 +69,10 @@ absl::StatusOr<HloInstruction*> MaybeTranspose(
     return operand;
   }
   TF_ASSIGN_OR_RETURN(auto* result, MakeTransposeHlo(operand, permutation));
+  // Assign the default layout to the transpose. This method is also used after
+  // layout normalization, and before, we don't care about the layout.
+  *result->mutable_shape()->mutable_layout() =
+      LayoutUtil::GetDefaultLayoutForShape(result->shape());
   return result;
 }
 
diff --git a/third_party/xla/xla/service/gather_simplifier_test.cc b/third_party/xla/xla/service/gather_simplifier_test.cc
index 5f00fe83a01e28..61b8bc716e120b 100644
--- a/third_party/xla/xla/service/gather_simplifier_test.cc
+++ b/third_party/xla/xla/service/gather_simplifier_test.cc
@@ -100,7 +100,7 @@ TEST_F(GatherSimplifierTest, MakesStartIndexMapIdentity) {
 
   RunAndFilecheckHloRewrite(kModuleStr, GatherSimplifier(), R"(
   %operand = f32[33,34,35]{2,1,0} parameter(0)
-           CHECK: %[[OPERAND:.*]] = f32[35,33,34]{0,2,1} transpose(%operand)
+           CHECK: %[[OPERAND:.*]] = f32[35,33,34]{2,1,0} transpose(%operand)
            CHECK: %[[GATHER:.*]] = f32[42,3,1,2]{{.*}} gather(%[[OPERAND]],
       CHECK-SAME:    start_index_map={0,1,2},
            CHECK: ROOT {{.*}} = f32[42,1,2,3]{{.*}} transpose(%[[GATHER]])
@@ -126,7 +126,7 @@ TEST_F(GatherSimplifierTest, CollapsesSomeDims) {
   RunAndFilecheckHloRewrite(kModuleStr, GatherSimplifier(), R"(
            CHECK: %[[GATHER:.*]] = f32[42,1,7,1]{3,2,1,0} gather(
            CHECK: %[[COLLAPSED:.*]] = f32[42,7]{1,0} reshape(%[[GATHER]])
-           CHECK: ROOT {{.*}} = f32[7,42]{0,1} transpose(%[[COLLAPSED]]),
+           CHECK: ROOT {{.*}} = f32[7,42]{1,0} transpose(%[[COLLAPSED]]),
       CHECK-SAME: dimensions={1,0}
   )");
 }
diff --git a/third_party/xla/xla/service/gpu/gpu_compiler.cc b/third_party/xla/xla/service/gpu/gpu_compiler.cc
index 609c92872f05e1..c9bd0c8812d8c5 100644
--- a/third_party/xla/xla/service/gpu/gpu_compiler.cc
+++ b/third_party/xla/xla/service/gpu/gpu_compiler.cc
@@ -1224,6 +1224,11 @@ absl::Status GpuCompiler::OptimizeHloModule(
     // duplicate or NOPs, so remove them with algebraic simplification and CSE.
     layout_normalization_pipeline.AddPass<HloPassFix<AlgebraicSimplifier>>(
         simplifier_options);
+    // Layout normalization will create broadcasts that are not canonical.
+    layout_normalization_pipeline.AddPass<BroadcastCanonicalizer>();
+    // Layout normalization will create scatters that are not simplified and
+    // also have unsorted update_window_dims.
+    layout_normalization_pipeline.AddPass<ScatterSimplifier>();
   }
   TF_RETURN_IF_ERROR(layout_normalization_pipeline.Run(hlo_module).status());
   // Run target-specific HLO optimization passes after layout assignment.
@@ -1376,6 +1381,9 @@ absl::Status GpuCompiler::OptimizeHloPostLayoutAssignment(
       // Remove any redundant operations (such as bitcasts) introduced by layout
       // normalization.
       pipeline.AddPass<HloPassFix<AlgebraicSimplifier>>(simplifier_options);
+      // Layout normalization will create scatters that are not simplified and
+      // also have unsorted update_window_dims.
+      pipeline.AddPass<ScatterSimplifier>();
     }
     pipeline.AddPass<BroadcastCanonicalizer>();
 
diff --git a/third_party/xla/xla/service/layout_normalization.cc b/third_party/xla/xla/service/layout_normalization.cc
index d966b5442b7829..6ff486dd0aedc9 100644
--- a/third_party/xla/xla/service/layout_normalization.cc
+++ b/third_party/xla/xla/service/layout_normalization.cc
@@ -25,6 +25,7 @@ limitations under the License.
 #include "xla/hlo/ir/dfs_hlo_visitor_with_default.h"
 #include "xla/hlo/ir/hlo_casting_utils.h"
 #include "xla/hlo/ir/hlo_instruction.h"
+#include "xla/hlo/ir/hlo_instructions.h"
 #include "xla/hlo/ir/hlo_module.h"
 #include "xla/layout_util.h"
 #include "xla/permutation_util.h"
@@ -377,6 +378,112 @@ class LayoutNormalizationVisitor : public DfsHloRewriteVisitor {
     return OkStatus();
   }
 
+  // Scatter is layout-preserving regarding the scatter operands, so we only
+  // have to permute values inside the ScatterDimensionNumbers.
+  Status HandleScatter(HloInstruction* hlo) override {
+    auto* scatter = Cast<HloScatterInstruction>(hlo);
+    std::vector<HloInstruction*> normalized_operands;
+    normalized_operands.reserve(scatter->scatter_operand_count());
+    Shape operand_shape = scatter->scatter_operands().front()->shape();
+    for (HloInstruction* operand : scatter->scatter_operands()) {
+      if (operand->shape().layout() != operand_shape.layout()) {
+        return FailedPrecondition(
+            "All scatter operands must have the same layout");
+      }
+      TF_ASSIGN_OR_RETURN(auto normalized_operand, GetNormalizedInput(operand));
+      normalized_operands.push_back(normalized_operand);
+    }
+    std::vector<HloInstruction*> normalized_updates;
+    normalized_updates.reserve(scatter->scatter_operand_count());
+    Shape update_shape = scatter->scatter_updates().front()->shape();
+    for (HloInstruction* operand : scatter->scatter_updates()) {
+      if (operand->shape().layout() != update_shape.layout()) {
+        return FailedPrecondition(
+            "All scatter updates must have the same layout");
+      }
+      TF_ASSIGN_OR_RETURN(auto normalized_update, GetNormalizedInput(operand));
+      normalized_updates.push_back(normalized_update);
+    }
+    TF_ASSIGN_OR_RETURN(auto normalized_indices,
+                        GetNormalizedInput(scatter->scatter_indices()));
+
+    // The scatter operands are normalized by applying a permutation such that
+    // perm(layout) = standard layout -> inverse layout permutation is applied.
+    auto indices_permutation = InversePermutation(
+        ToTransposeDimensions(scatter->scatter_indices()->shape().layout()));
+
+    auto layout_permutation =
+        ToTransposeDimensions(scatter->scatter_operands()[0]->shape().layout());
+    auto operand_permutation = InversePermutation(layout_permutation);
+
+    auto update_permutation = InversePermutation(
+        ToTransposeDimensions(scatter->scatter_updates()[0]->shape().layout()));
+
+    // scatter_dims_to_operand_dims -> mapping from scatter dimensions to
+    // operand dimensions. scatter dimension i corresponds to
+    // scatter_dims_to_operand_dims[i] operand dimension.
+
+    const auto& dims = scatter->scatter_dimension_numbers();
+    ScatterDimensionNumbers normalized_dims;
+    normalized_dims.set_index_vector_dim(
+        indices_permutation[dims.index_vector_dim()]);
+    for (int64_t dim : dims.scatter_dims_to_operand_dims()) {
+      normalized_dims.add_scatter_dims_to_operand_dims(
+          operand_permutation[dim]);
+    }
+    std::vector<int64_t> normalized_update_window_dims;
+    normalized_update_window_dims.reserve(dims.update_window_dims_size());
+    for (int64_t dim : dims.update_window_dims()) {
+      normalized_update_window_dims.push_back(update_permutation[dim]);
+    }
+
+    // Now reorder 'normalized_update_window_dims' and 'inserted_window_dims'
+    // according to the output permutation, so that the window dimensions again
+    // appear in the same order as in the output. First we need to build a
+    // combined array of window dimensions. Note: 'inserted_window_dims' and
+    // 'update_window_dims' must be sorted according to shape inference/hlo
+    // verifier. We will temporarily create an unsorted update_window_dims
+    // attribute and rely on ScatterSimplifier to clean this up.
+    std::vector<int64_t> window_dimensions(operand_permutation.size());
+    for (int64_t i = 0, j = 0, k = 0; i < window_dimensions.size(); ++i) {
+      if (j < dims.inserted_window_dims_size() &&
+          dims.inserted_window_dims(j) == i) {
+        window_dimensions[i] = -1;
+        ++j;
+      } else {
+        window_dimensions[i] = normalized_update_window_dims[k];
+        ++k;
+      }
+    }
+    std::vector<int64_t> permuted_window_dimensions =
+        ComposePermutations(window_dimensions, layout_permutation);
+    for (int64_t i = 0; i < permuted_window_dimensions.size(); ++i) {
+      if (permuted_window_dimensions[i] == -1) {
+        normalized_dims.add_inserted_window_dims(i);
+      } else {
+        normalized_dims.add_update_window_dims(permuted_window_dimensions[i]);
+      }
+    }
+
+    auto normalized_shape = normalized_operands.front()->shape();
+    if (scatter->shape().IsTuple()) {
+      std::vector<Shape> tuple_shapes;
+      tuple_shapes.reserve(normalized_operands.size());
+      for (HloInstruction* operand : normalized_operands) {
+        tuple_shapes.push_back(operand->shape());
+      }
+      normalized_shape = ShapeUtil::MakeTupleShape(tuple_shapes);
+    }
+    auto normalized_scatter = hlo->AddInstruction(HloInstruction::CreateScatter(
+        normalized_shape, normalized_operands, normalized_indices,
+        normalized_updates, scatter->to_apply(), normalized_dims,
+        scatter->indices_are_sorted(), scatter->unique_indices()));
+    SetVisited(*normalized_scatter);
+    auto bc_to_orig = MakeBitcastHlo(normalized_scatter, scatter->shape());
+    TF_RETURN_IF_ERROR(ReplaceInstruction(scatter, bc_to_orig));
+    return OkStatus();
+  }
+
   // For bitcasting transposes, converts:
   //
   // A{I} -> bitcast[S]{L} -> transpose{L2}
diff --git a/third_party/xla/xla/service/layout_normalization_test.cc b/third_party/xla/xla/service/layout_normalization_test.cc
index 6c5da347cb1d58..1b74f8bb44c325 100644
--- a/third_party/xla/xla/service/layout_normalization_test.cc
+++ b/third_party/xla/xla/service/layout_normalization_test.cc
@@ -15,10 +15,15 @@ limitations under the License.
 
 #include "xla/service/layout_normalization.h"
 
+#include <functional>
 #include <optional>
 #include <utility>
 
+#include "absl/strings/string_view.h"
+#include "xla/hlo/ir/hlo_module.h"
+#include "xla/service/scatter_simplifier.h"
 #include "xla/tests/hlo_test_base.h"
+#include "tsl/platform/status.h"
 #include "tsl/platform/test.h"
 
 namespace xla {
@@ -26,9 +31,11 @@ namespace {
 
 class LayoutNormalizationTest : public HloTestBase {
  public:
-  void CheckLayoutNormalization(absl::string_view hlo,
-                                std::optional<absl::string_view> expected) {
-    RunAndFilecheckHloRewrite(hlo, LayoutNormalization{}, expected);
+  void CheckLayoutNormalization(
+      absl::string_view hlo, std::optional<absl::string_view> expected,
+      std::function<void(HloModule*)> after_pass_checks = nullptr) {
+    RunAndFilecheckHloRewrite(hlo, LayoutNormalization{}, expected,
+                              after_pass_checks);
   }
 };
 
@@ -769,5 +776,101 @@ ENTRY main {
 )");
 }
 
+TEST_F(LayoutNormalizationTest, Scatter) {
+  const char* hlo = R"(
+HloModule simplified_scatter
+
+region_0.10 {
+  Arg_0.11 = s16[] parameter(0)
+  Arg_1.12 = s16[] parameter(1)
+  ROOT maximum.13 = s16[] maximum(Arg_0.11, Arg_1.12)
+}
+
+ENTRY main.17 {
+  p0 = s16[3,2,2,14,16]{0,1,4,3,2} parameter(0)
+  p1 = s32[6,11,8,2]{0,2,1,3} parameter(1)
+  p2 = s16[6,11,8,3,5]{0,1,4,3,2} parameter(2)
+  ROOT scatter = s16[3,2,2,14,16]{0,1,4,3,2} scatter(p0, p1, p2), update_window_dims={3,4}, inserted_window_dims={1,2,3}, scatter_dims_to_operand_dims={4,0}, index_vector_dim=3, to_apply=region_0.10
+}
+)";
+
+  CheckLayoutNormalization(
+      hlo, R"(
+// CHECK: scatter({{.*}}),
+// CHECK-SAME: update_window_dims={2,1}, inserted_window_dims={0,1,3}, scatter_dims_to_operand_dims={2,4}, index_vector_dim=0, to_apply=%region_0.10
+)",
+      // Run the ScatterSimplifier afterwards, otherwise the verifier will
+      // complain!
+      [](HloModule* module) {
+        TF_CHECK_OK(ScatterSimplifier().Run(module).status());
+      });
+}
+
+TEST_F(LayoutNormalizationTest, SimplifiedScatter) {
+  const char* hlo = R"(
+HloModule simplified_scatter
+
+region_0.10 {
+  Arg_0.11 = s16[] parameter(0)
+  Arg_1.12 = s16[] parameter(1)
+  ROOT maximum.13 = s16[] maximum(Arg_0.11, Arg_1.12)
+}
+
+ENTRY main.17 {
+  p0 = s16[16,3,2,2,14]{0,4,3,2,1} parameter(0)
+  p1 = s32[528,2]{1,0} parameter(1)
+  p2 = s16[528,5,3,1,1,1]{1,2,0,5,4,3} parameter(2)
+  ROOT scatter = s16[16,3,2,2,14]{0,4,3,2,1} scatter(p0, p1, p2), update_window_dims={1,2,3,4,5}, inserted_window_dims={}, scatter_dims_to_operand_dims={0,1}, index_vector_dim=1, to_apply=region_0.10
+}
+)";
+
+  CheckLayoutNormalization(
+      hlo, R"(
+// CHECK: scatter({{.*}}),
+// CHECK-SAME: update_window_dims={4,0,1,2,5}, inserted_window_dims={}, scatter_dims_to_operand_dims={4,0}, index_vector_dim=1, to_apply=%region_0.10
+)",
+      // Run the ScatterSimplifier afterwards, otherwise the verifier will
+      // complain!
+      [](HloModule* module) {
+        TF_CHECK_OK(ScatterSimplifier().Run(module).status());
+      });
+}
+
+TEST_F(LayoutNormalizationTest, VariadicScatter) {
+  const char* hlo = R"(
+HloModule simplified_scatter
+
+region_0.10 {
+  Arg_0.11 = s16[] parameter(0)
+  Arg_1.12 = s16[] parameter(1)
+  Arg_2.13 = s16[] parameter(2)
+  Arg_3.14 = s16[] parameter(3)
+  maximum.15 = s16[] maximum(Arg_0.11, Arg_1.12)
+  maximum.16 = s16[] maximum(Arg_2.13, Arg_3.14)
+  ROOT res = (s16[], s16[]) tuple(maximum.15, maximum.16)
+}
+
+ENTRY main.17 {
+  p0 = s16[16,3,2,2,14]{0,4,3,2,1} parameter(0)
+  p1 = s16[16,3,2,2,14]{0,4,3,2,1} parameter(1)
+  p2 = s32[528,2]{1,0} parameter(2)
+  p3 = s16[528,5,3,1,1,1]{1,2,0,5,4,3} parameter(3)
+  p4 = s16[528,5,3,1,1,1]{1,2,0,5,4,3} parameter(4)
+  ROOT scatter = (s16[16,3,2,2,14]{0,4,3,2,1}, s16[16,3,2,2,14]{0,4,3,2,1}) scatter(p0, p1, p2, p3, p4), update_window_dims={1,2,3,4,5}, inserted_window_dims={}, scatter_dims_to_operand_dims={0,1}, index_vector_dim=1, to_apply=region_0.10
+}
+)";
+
+  CheckLayoutNormalization(
+      hlo, R"(
+// CHECK: scatter({{.*}}),
+// CHECK-SAME: update_window_dims={4,0,1,2,5}, inserted_window_dims={}, scatter_dims_to_operand_dims={4,0}, index_vector_dim=1, to_apply=%region_0.10
+)",
+      // Run the ScatterSimplifier afterwards, otherwise the verifier will
+      // complain!
+      [](HloModule* module) {
+        TF_CHECK_OK(ScatterSimplifier().Run(module).status());
+      });
+}
+
 }  // namespace
 }  // namespace xla
diff --git a/third_party/xla/xla/service/scatter_simplifier_test.cc b/third_party/xla/xla/service/scatter_simplifier_test.cc
index 36d82b4b487eee..dc813e506dedac 100644
--- a/third_party/xla/xla/service/scatter_simplifier_test.cc
+++ b/third_party/xla/xla/service/scatter_simplifier_test.cc
@@ -152,7 +152,7 @@ TEST_F(ScatterSimplifierTest, MovesIndexVectorDim) {
     })";
 
   RunAndFilecheckHloRewrite(kModuleStr, ScatterSimplifier(), R"(
-           CHECK: %[[TRANSPOSED_INDICES:.*]] = s32[1,2]{0,1}
+           CHECK: %[[TRANSPOSED_INDICES:.*]] = s32[1,2]{1,0}
       CHECK-SAME:     transpose(%indices), dimensions={1,0}
            CHECK: scatter(%operand, %[[TRANSPOSED_INDICES]], %update),
       CHECK-SAME:     index_vector_dim=1
@@ -183,9 +183,9 @@ TEST_F(ScatterSimplifierTest, TransformsUpdatesAndOperandUsingScatterDims) {
     })";
 
   RunAndFilecheckHloRewrite(kModuleStr, ScatterSimplifier(), R"(
-           CHECK: %[[T_OPERAND:.*]] = f32[5,3,4]{0,2,1} transpose(%operand),
+           CHECK: %[[T_OPERAND:.*]] = f32[5,3,4]{2,1,0} transpose(%operand),
       CHECK-SAME:     dimensions={2,0,1}
-           CHECK: %[[T_UPDATES:.*]] = f32[2,3,1,1]{1,3,2,0} transpose(%update),
+           CHECK: %[[T_UPDATES:.*]] = f32[2,3,1,1]{3,2,1,0} transpose(%update),
       CHECK-SAME:     dimensions={0,3,1,2}
            CHECK: %[[SCATTER:.*]] = {{.*}} scatter(
       CHECK-SAME:     %[[T_OPERAND]], %indices, %[[T_UPDATES]])
@@ -218,7 +218,7 @@ TEST_F(ScatterSimplifierTest, MakesScatterDimensionsLeadingInUpdates) {
     })";
 
   RunAndFilecheckHloRewrite(kModuleStr, ScatterSimplifier(), R"(
-           CHECK: %[[TRANSPOSED_UPDATES:.*]] = f32[1,2]{0,1}
+           CHECK: %[[TRANSPOSED_UPDATES:.*]] = f32[1,2]{1,0}
       CHECK-SAME:     transpose(%update), dimensions={1,0}
            CHECK: scatter(
       CHECK-SAME:     %[[TRANSPOSED_UPDATES]]

From 88f115a6efa6e191eba3b6a10af02c29687b24b2 Mon Sep 17 00:00:00 2001
From: Henning Becker <hebecker@google.com>
Date: Thu, 16 May 2024 23:05:43 -0700
Subject: [PATCH 322/478] Fix string-replace mistake in ROCm code

This was introduced by the `xla::Status`/`absl::Status` cleanup by accident.

PiperOrigin-RevId: 634641828
---
 third_party/xla/xla/stream_executor/rocm/rocm_dnn.cc | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/third_party/xla/xla/stream_executor/rocm/rocm_dnn.cc b/third_party/xla/xla/stream_executor/rocm/rocm_dnn.cc
index 046c86df225408..cf5f9a9c80b72f 100644
--- a/third_party/xla/xla/stream_executor/rocm/rocm_dnn.cc
+++ b/third_party/xla/xla/stream_executor/rocm/rocm_dnn.cc
@@ -1975,7 +1975,7 @@ template <typename Base>
 class MIOpenDescriptorCommon : public MixinBase<Base> {
  public:
   bool ok() const { return status_.ok(); }
-  absl::Status absl::Status() const { return status_; }
+  absl::Status Status() const { return status_; }
 
  protected:
   void SetFailure(const absl::Status& status) { status_.Update(status); }

From 8a74312d921a28d28f8b90effe03d6320958939b Mon Sep 17 00:00:00 2001
From: Henning Becker <hebecker@google.com>
Date: Thu, 16 May 2024 23:14:09 -0700
Subject: [PATCH 323/478] Add missing ROCm dependencies in xla/tests

All the changed targets include `rocm_config.h`,
therefore they need to depend on `:rocm_headers`
directly to pass the layering check.

PiperOrigin-RevId: 634643380
---
 third_party/xla/xla/tests/BUILD | 30 ++++++++++++++++++++++++------
 1 file changed, 24 insertions(+), 6 deletions(-)

diff --git a/third_party/xla/xla/tests/BUILD b/third_party/xla/xla/tests/BUILD
index ea9aa2f35d7ef4..4bd2e01e65a6dc 100644
--- a/third_party/xla/xla/tests/BUILD
+++ b/third_party/xla/xla/tests/BUILD
@@ -814,7 +814,10 @@ xla_test(
         "@com_google_absl//absl/types:span",
         "@local_tsl//tsl/platform:ml_dtypes",
         "@ml_dtypes//:float8",
-    ],
+    ] + if_rocm_is_configured([
+        # keep sorted
+        "@local_config_rocm//rocm:rocm_headers",
+    ]),
 )
 
 cc_library(
@@ -916,7 +919,10 @@ xla_test(
         "@com_google_absl//absl/strings",
         "@local_tsl//tsl/platform:test",
         "@local_tsl//tsl/platform:test_benchmark",
-    ],
+    ] + if_rocm_is_configured([
+        # keep sorted
+        "@local_config_rocm//rocm:rocm_headers",
+    ]),
 )
 
 # Run dot tests with auto-tuning disabled.  This just does a basic sanity check
@@ -950,7 +956,10 @@ xla_test(
         "@com_google_absl//absl/strings",
         "@local_tsl//tsl/platform:test",
         "@local_tsl//tsl/platform:test_benchmark",
-    ],
+    ] + if_rocm_is_configured([
+        # keep sorted
+        "@local_config_rocm//rocm:rocm_headers",
+    ]),
 )
 
 # Run dot tests with dot canonicalization after the layout assignment pass.
@@ -989,7 +998,10 @@ xla_test(
         "@com_google_absl//absl/strings",
         "@local_tsl//tsl/platform:test",
         "@local_tsl//tsl/platform:test_benchmark",
-    ],
+    ] + if_rocm_is_configured([
+        # keep sorted
+        "@local_config_rocm//rocm:rocm_headers",
+    ]),
 )
 
 xla_test(
@@ -1062,7 +1074,10 @@ xla_test(
         "@com_google_absl//absl/strings",
         "@local_tsl//tsl/platform:test",
         "@local_tsl//tsl/platform:test_benchmark",
-    ],
+    ] + if_rocm_is_configured([
+        # keep sorted
+        "@local_config_rocm//rocm:rocm_headers",
+    ]),
 )
 
 xla_test(
@@ -1909,7 +1924,10 @@ xla_test(
         "@com_google_absl//absl/strings",
         "@com_google_absl//absl/strings:str_format",
         "@local_tsl//tsl/platform:test",
-    ],
+    ] + if_rocm_is_configured([
+        # keep sorted
+        "@local_config_rocm//rocm:rocm_headers",
+    ]),
 )
 
 xla_test(

From 24a30f65bfa0ebcc596f33830552671d96b4aa4c Mon Sep 17 00:00:00 2001
From: Henning Becker <hebecker@google.com>
Date: Thu, 16 May 2024 23:53:06 -0700
Subject: [PATCH 324/478] Revive the PLATFORM_GOOGLE code in the ROCm wrappers

The ROCm wrappers is basically a hand-written implementation layer library.

We bypass it when `PLATFORM_GOOGLE` is set and just statically link the ROCm
and HIP libraries instead. The bypass code was broken since we didn't use it
for a long time.

So this CL is fixing a bunch of issues with the code and makes it work with
the latest ROCm version.

PiperOrigin-RevId: 634650804
---
 .../xla/xla/stream_executor/rocm/hipsparse_wrapper.h  |  5 ++++-
 .../xla/xla/stream_executor/rocm/rocblas_wrapper.h    |  6 +++---
 .../xla/xla/stream_executor/rocm/roctracer_wrapper.h  | 11 ++++++-----
 3 files changed, 13 insertions(+), 9 deletions(-)

diff --git a/third_party/xla/xla/stream_executor/rocm/hipsparse_wrapper.h b/third_party/xla/xla/stream_executor/rocm/hipsparse_wrapper.h
index b4bcc7d8f3944e..d66b2426d50d84 100644
--- a/third_party/xla/xla/stream_executor/rocm/hipsparse_wrapper.h
+++ b/third_party/xla/xla/stream_executor/rocm/hipsparse_wrapper.h
@@ -20,12 +20,15 @@ limitations under the License.
 #ifndef XLA_STREAM_EXECUTOR_ROCM_HIPSPARSE_WRAPPER_H_
 #define XLA_STREAM_EXECUTOR_ROCM_HIPSPARSE_WRAPPER_H_
 
+#include "rocm/rocm_config.h"
+
 #if (TF_ROCM_VERSION >= 50200)
 #include "rocm/include/hipsparse/hipsparse.h"
 #else
 #include "rocm/include/hipsparse.h"
 #endif
 #include "xla/stream_executor/platform/dso_loader.h"
+#include "xla/stream_executor/platform/platform.h"
 #include "xla/stream_executor/platform/port.h"
 #include "tsl/platform/env.h"
 
@@ -38,7 +41,7 @@ namespace wrap {
   struct WrapperShim__##__name {                    \
     template <typename... Args>                     \
     hipsparseStatus_t operator()(Args... args) {    \
-      hipSparseStatus_t retval = ::__name(args...); \
+      hipsparseStatus_t retval = ::__name(args...); \
       return retval;                                \
     }                                               \
   } __name;
diff --git a/third_party/xla/xla/stream_executor/rocm/rocblas_wrapper.h b/third_party/xla/xla/stream_executor/rocm/rocblas_wrapper.h
index 3d444ab83a0ee6..436b3d3692822b 100644
--- a/third_party/xla/xla/stream_executor/rocm/rocblas_wrapper.h
+++ b/third_party/xla/xla/stream_executor/rocm/rocblas_wrapper.h
@@ -27,23 +27,23 @@ limitations under the License.
 #include "xla/stream_executor/platform/dso_loader.h"
 #include "xla/stream_executor/platform/port.h"
 #include "tsl/platform/env.h"
+#include "tsl/platform/platform.h"
 
 namespace stream_executor {
 namespace wrap {
 
-using stream_executor::internal::CachedDsoLoader::GetRocblasDsoHandle;
-
 #ifdef PLATFORM_GOOGLE
 #define ROCBLAS_API_WRAPPER(__name)               \
   struct WrapperShim__##__name {                  \
     constexpr static const char* kName = #__name; \
     template <typename... Args>                   \
     rocblas_status operator()(Args... args) {     \
-      return ::__name(args...);                   \
+      return (::__name)(args...);                 \
     }                                             \
   } __name;
 
 #else
+using stream_executor::internal::CachedDsoLoader::GetRocblasDsoHandle;
 
 #define ROCBLAS_API_WRAPPER(__name)                                      \
   static struct DynLoadShim__##__name {                                  \
diff --git a/third_party/xla/xla/stream_executor/rocm/roctracer_wrapper.h b/third_party/xla/xla/stream_executor/rocm/roctracer_wrapper.h
index cbfe9e2e2cfe67..b42751bb53e0cf 100644
--- a/third_party/xla/xla/stream_executor/rocm/roctracer_wrapper.h
+++ b/third_party/xla/xla/stream_executor/rocm/roctracer_wrapper.h
@@ -31,16 +31,17 @@ limitations under the License.
 #include "xla/stream_executor/platform/dso_loader.h"
 #include "xla/stream_executor/platform/port.h"
 #include "tsl/platform/env.h"
+#include "tsl/platform/platform.h"
 
 namespace stream_executor {
 namespace wrap {
 
 #ifdef PLATFORM_GOOGLE
 
-#define ROCTRACER_API_WRAPPER(API_NAME)                          \
-  template <typename... Args>                                    \
-  auto API_NAME()(Args... args)->decltype(::API_NAME(args...)) { \
-    return ::API_NAME(args...);                                  \
+#define ROCTRACER_API_WRAPPER(API_NAME)                            \
+  template <typename... Args>                                      \
+  auto API_NAME(Args... args) -> decltype((::API_NAME)(args...)) { \
+    return (::API_NAME)(args...);                                  \
   }
 
 #else
@@ -52,7 +53,7 @@ namespace wrap {
     static FuncPtrT loaded = []() -> FuncPtrT {                               \
       static const char* kName = #API_NAME;                                   \
       void* f;                                                                \
-      auto s = tsl::Env::Default() -> GetSymbolFromLibrary(                   \
+      auto s = tsl::Env::Default()->GetSymbolFromLibrary(                     \
           stream_executor::internal::CachedDsoLoader::GetRoctracerDsoHandle() \
               .value(),                                                       \
           kName, &f);                                                         \

From 07e48f1c97963b9918e9ae6aaba79143226162de Mon Sep 17 00:00:00 2001
From: Johannes Reifferscheid <jreiffers@google.com>
Date: Thu, 16 May 2024 23:57:48 -0700
Subject: [PATCH 325/478] Fix scatter bounds checks for unsigned indices.

I accidentally used sle where we need ule, resulting in incorrect
results for values >= 2**31 / 2**63.

PiperOrigin-RevId: 634651640
---
 .../xla/xla/service/gpu/fusions/scatter_mlir.cc       |  8 +++-----
 .../xla/xla/service/gpu/fusions/scatter_mlir_test.cc  | 11 ++++++++---
 2 files changed, 11 insertions(+), 8 deletions(-)

diff --git a/third_party/xla/xla/service/gpu/fusions/scatter_mlir.cc b/third_party/xla/xla/service/gpu/fusions/scatter_mlir.cc
index e90b948275889c..372b6926953e27 100644
--- a/third_party/xla/xla/service/gpu/fusions/scatter_mlir.cc
+++ b/third_party/xla/xla/service/gpu/fusions/scatter_mlir.cc
@@ -237,17 +237,15 @@ absl::Status MlirScatterFusion::EmitEntryFunction(
             index = b.create<ma::IndexCastUIOp>(b.getIndexType(), index);
           } else {
             index = b.create<ma::IndexCastOp>(b.getIndexType(), index);
-            auto c0 = b.create<ma::ConstantIndexOp>(0);
-            in_bounds = b.create<ma::AndIOp>(
-                in_bounds,
-                b.create<ma::CmpIOp>(ma::CmpIPredicate::sge, index, c0));
           }
           Value ub = b.create<ma::ConstantIndexOp>(
               scatter_operand->shape().dimensions(i) -
               scatter_update->shape().dimensions(i + 1));
+          // One bounds check is enough even for signed indices: `sge 0` is
+          // implied by `ule ub`, because `ub >= 0`.
           in_bounds = b.create<ma::AndIOp>(
               in_bounds,
-              b.create<ma::CmpIOp>(ma::CmpIPredicate::sle, index, ub));
+              b.create<ma::CmpIOp>(ma::CmpIPredicate::ule, index, ub));
           indices[i] = b.create<ma::AddIOp>(index, indices[i]);
         }
         // Call scatter's computation if is_in_bounds.
diff --git a/third_party/xla/xla/service/gpu/fusions/scatter_mlir_test.cc b/third_party/xla/xla/service/gpu/fusions/scatter_mlir_test.cc
index 86ce79992c819c..a961a560288d1f 100644
--- a/third_party/xla/xla/service/gpu/fusions/scatter_mlir_test.cc
+++ b/third_party/xla/xla/service/gpu/fusions/scatter_mlir_test.cc
@@ -207,7 +207,7 @@ TEST_F(MlirScatterFusionTest, Scatter_UniqueIndices) {
     // CHECK:      xla_gpu.pure_call @scatter_indices(%[[OPERAND]], %[[INDICES]]
     // CHECK-SAME:  %[[UPDATES]], %[[SLICE_ID]], %[[C0]])
 
-    // CHECK:      %[[IN_BOUNDS:.*]] = arith.andi
+    // CHECK:      %[[IN_BOUNDS:.*]] = arith.cmpi ule
     // CHECK:      scf.if %[[IN_BOUNDS]] -> (tensor<10x5xf32>) {
     // CHECK:        %[[CURRENT:.*]] = xla_gpu.pure_call @scatter_operand(
     // CHECK-SAME:  %[[OPERAND]], %[[INDICES]], %[[UPDATES]],
@@ -255,6 +255,10 @@ TEST_F(MlirScatterFusionTest, Scatter_Unsigned) {
   )";
   TF_ASSERT_OK(EmitAndCheckIR(kHloString, R"(
     // CHECK: func.func @fused_computation(
+    // CHECK: %[[PARAM:.*]] = xla_gpu.pure_call @scatter_indices
+    // CHECK: %[[CVT:.*]] = builtin.unrealized_conversion_cast %[[PARAM]]
+    // CHECK: %[[INDEX:.*]] = arith.index_castui %[[CVT]]
+    // CHECK: arith.cmpi ule, %[[INDEX]]
   )"));
   EXPECT_TRUE(RunAndCompareNoHloPasses(kHloString, ErrorSpec{1e-3}));
 }
@@ -302,7 +306,7 @@ TEST_F(MlirScatterFusionTest, Scatter_Add) {
     // CHECK-SAME:    %[[OUT:[a-zA-Z0-9]*]]: tensor<10x5xf32>
 
     // CHECK: %[[UPD_ELEM:.*]] = xla_gpu.pure_call @scatter_update
-    // CHECK: %[[IN_BOUNDS:.*]] = arith.andi
+    // CHECK: %[[IN_BOUNDS:.*]] = arith.cmpi ule
     // CHECK: scf.if %[[IN_BOUNDS]] -> (tensor<10x5xf32>) {
     // CHECK:   %[[RMW:.*]] = xla_gpu.atomic_rmw %[[OUT]]
     // CHECK:   ^bb0(%[[CUR_VALUE:.*]]: f32):
@@ -359,7 +363,7 @@ TEST_F(MlirScatterFusionTest, Scatter_Overwrite) {
     // CHECK-SAME:    %[[OUT:[a-zA-Z0-9]*]]: tensor<10x5xf32>
 
     // CHECK: %[[UPD_ELEM:.*]] = xla_gpu.pure_call @scatter_update
-    // CHECK: %[[IN_BOUNDS:.*]] = arith.andi
+    // CHECK: %[[IN_BOUNDS:.*]] = arith.cmpi ule
     // CHECK: scf.if %[[IN_BOUNDS]] -> (tensor<10x5xf32>) {
     // CHECK:   %[[RMW:.*]] = xla_gpu.atomic_rmw %[[OUT]]
     // CHECK:   ^bb0(%[[CUR_VALUE:.*]]: f32):
@@ -372,6 +376,7 @@ TEST_F(MlirScatterFusionTest, Scatter_Overwrite) {
   )"));
   EXPECT_TRUE(RunAndCompareNoHloPasses(kHloString, ErrorSpec{1e-3}));
 }
+
 }  // namespace
 }  // namespace gpu
 }  // namespace xla

From 60c746c1e22ccdcf3100d17f89ea3327fdd90227 Mon Sep 17 00:00:00 2001
From: Henning Becker <hebecker@google.com>
Date: Fri, 17 May 2024 00:15:33 -0700
Subject: [PATCH 326/478] Fix TF Grappler ROCm build issue

The `gcnArch` property had been deprecated for a while (at least for 4 years).
`gcnArchName` is the replacement and since it's already available for so
long it shouldn't break anyone.

PiperOrigin-RevId: 634656093
---
 tensorflow/core/grappler/clusters/utils.cc | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorflow/core/grappler/clusters/utils.cc b/tensorflow/core/grappler/clusters/utils.cc
index d428b6632e6be9..e38807af1c797e 100644
--- a/tensorflow/core/grappler/clusters/utils.cc
+++ b/tensorflow/core/grappler/clusters/utils.cc
@@ -135,7 +135,7 @@ DeviceProperties GetLocalGPUInfo(PlatformDeviceId platform_device_id) {
                        properties.memoryClockRate * 2);
 
   (*device.mutable_environment())["architecture"] =
-      strings::StrCat("gfx", properties.gcnArch);
+      strings::StrCat("gfx", properties.gcnArchName);
 #endif
 
   return device;

From 98e84d2d3fb1f6910baf8a1e3f3d96f1df6a9e58 Mon Sep 17 00:00:00 2001
From: Henning Becker <hebecker@google.com>
Date: Fri, 17 May 2024 00:30:53 -0700
Subject: [PATCH 327/478] Disable NCCL persistent plan allocator on ROCm

It had already been disabled for OSS builds, this is only making
sure it's also disabled when `PLATFORM_GOOGLE` is set to true.

It's a non-functional change - just some ugly macro hackery.

PiperOrigin-RevId: 634659230
---
 .../xla/xla/service/gpu/runtime/nccl_api.cc   | 29 ++++++++++++-------
 1 file changed, 19 insertions(+), 10 deletions(-)

diff --git a/third_party/xla/xla/service/gpu/runtime/nccl_api.cc b/third_party/xla/xla/service/gpu/runtime/nccl_api.cc
index 0224fdd5cf85d5..49abfa4636b44e 100644
--- a/third_party/xla/xla/service/gpu/runtime/nccl_api.cc
+++ b/third_party/xla/xla/service/gpu/runtime/nccl_api.cc
@@ -55,6 +55,12 @@ limitations under the License.
 #include "third_party/nccl/nccl.h"
 #endif  // TENSORFLOW_USE_ROCM
 
+#if (defined(PLATFORM_GOOGLE) && !defined(TENSORFLOW_USE_ROCM))
+#define WITH_PERSISTENT_PLAN_ALLOCATOR_SUPPORT true
+#else
+#define WITH_PERSISTENT_PLAN_ALLOCATOR_SUPPORT false
+#endif
+
 namespace xla::gpu {
 
 //==-----------------------------------------------------------------------===//
@@ -184,7 +190,7 @@ static ncclComm_t Cast(NcclApi::NcclCommHandle comm) {
   return reinterpret_cast<ncclComm_t>(comm);
 }
 
-#ifdef PLATFORM_GOOGLE
+#if WITH_PERSISTENT_PLAN_ALLOCATOR_SUPPORT
 static ncclPersistentPlanAllocator* Cast(
     NcclApi::NcclPersistentPlanAllocatorHandle handle) {
   return reinterpret_cast<ncclPersistentPlanAllocator*>(handle);
@@ -199,7 +205,7 @@ static NcclApi::NcclPersistentPlanAllocatorHandle Cast(
     ncclPersistentPlanAllocator* ptr) {
   return reinterpret_cast<NcclApi::NcclPersistentPlanAllocatorHandle>(ptr);
 }
-#endif  // PLATFORM_GOOGLE
+#endif  // WITH_PERSISTENT_PLAN_ALLOCATOR_SUPPORT
 
 //==-----------------------------------------------------------------------===//
 // NcclApi::PersistentPlanAllocator
@@ -217,7 +223,7 @@ PersistentPlanAllocator::PersistentPlanAllocator(
       stream_(stream) {
   // NCCL persistent plan allocator is implemented as NCCL patch that is not yet
   // open sourced and can't be used from OSS XLA.
-#ifdef PLATFORM_GOOGLE
+#if WITH_PERSISTENT_PLAN_ALLOCATOR_SUPPORT
   auto* nccl_allocator = new ncclPersistentPlanAllocator;
   nccl_allocator->ctl = this;
 
@@ -238,13 +244,13 @@ PersistentPlanAllocator::PersistentPlanAllocator(
   };
 
   handle_ = Cast(nccl_allocator);
-#endif  // PLATFORM_GOOGLE
+#endif  // WITH_PERSISTENT_PLAN_ALLOCATOR_SUPPORT
 }
 
 PersistentPlanAllocator::~PersistentPlanAllocator() {
-#ifdef PLATFORM_GOOGLE
+#if WITH_PERSISTENT_PLAN_ALLOCATOR_SUPPORT
   delete Cast(handle_);
-#endif  // PLATFORM_GOOGLE
+#endif  // WITH_PERSISTENT_PLAN_ALLOCATOR_SUPPORT
 }
 
 absl::StatusOr<se::DeviceMemoryBase>
@@ -266,22 +272,22 @@ absl::Status PersistentPlanAllocator::Deallocate(se::DeviceMemoryBase mem) {
 ScopedPersistentPlanAllocator::ScopedPersistentPlanAllocator(
     NcclCommHandle comm, tsl::RCReference<PersistentPlanAllocator> allocator)
     : comm_(comm), allocator_(std::move(allocator)) {
-#ifdef PLATFORM_GOOGLE
+#if WITH_PERSISTENT_PLAN_ALLOCATOR_SUPPORT
   XLA_NCCL_CHECK(
       ncclCommGetPersistentPlanAllocator(Cast(comm_), Cast(&recover_)))
       << "Failed to get NCCL persistent plan allocator";
   XLA_NCCL_CHECK(ncclCommSetPersistentPlanAllocator(Cast(comm_),
                                                     Cast(allocator_->handle())))
       << "Failed to set NCCL persistent plan allocator";
-#endif  // PLATFORM_GOOGLE
+#endif  // WITH_PERSISTENT_PLAN_ALLOCATOR_SUPPORT
 }
 
 ScopedPersistentPlanAllocator::~ScopedPersistentPlanAllocator() {
-#ifdef PLATFORM_GOOGLE
+#if WITH_PERSISTENT_PLAN_ALLOCATOR_SUPPORT
   XLA_NCCL_CHECK(
       ncclCommSetPersistentPlanAllocator(Cast(comm_), Cast(recover_)))
       << "Failed to set NCCL persistent plan allocator";
-#endif  // PLATFORM_GOOGLE
+#endif  // WITH_PERSISTENT_PLAN_ALLOCATOR_SUPPORT
 }
 
 //==-----------------------------------------------------------------------===//
@@ -655,6 +661,9 @@ DefaultNcclApi::DeregisterBuffer(NcclCommHandle comm,
 #if (NCCL_VERSION_CODE >= 21901)
   return XLA_NCCL_STATUS(
       ncclCommDeregister(Cast(comm), reinterpret_cast<void*>(handle)));
+#else
+  return absl::UnimplementedError(
+      "ncclCommDeregister is unavailable in this build.");
 #endif  // NCCL_VERSION_CODE >= 21901
 }
 }  // namespace xla::gpu

From f699a428c74421ed1ae98ab9e957415779695b9d Mon Sep 17 00:00:00 2001
From: Henning Becker <hebecker@google.com>
Date: Fri, 17 May 2024 00:38:41 -0700
Subject: [PATCH 328/478] Add missing ROCm dependency in gpu/runtime

`:convolution_thunk` includes `rocm_config.h`, hence we need to depend
on `rocm_headers` to make the layering check happy.

PiperOrigin-RevId: 634661039
---
 third_party/xla/xla/service/gpu/runtime/BUILD | 6 +++++-
 1 file changed, 5 insertions(+), 1 deletion(-)

diff --git a/third_party/xla/xla/service/gpu/runtime/BUILD b/third_party/xla/xla/service/gpu/runtime/BUILD
index 3bc13cac2f9665..ef8a8134628a05 100644
--- a/third_party/xla/xla/service/gpu/runtime/BUILD
+++ b/third_party/xla/xla/service/gpu/runtime/BUILD
@@ -217,6 +217,7 @@ cc_library(
         "//xla/stream_executor/cuda:cuda_driver",
         "//xla/stream_executor/cuda:cuda_executor",
     ]) + if_rocm_is_configured([
+        "//third_party/gpus/rocm:rocm_headers",
         "@local_config_rocm//rocm:rccl",
         "//xla/stream_executor/rocm:rocm_driver",
         "//xla/stream_executor/rocm:rocm_executor",
@@ -527,7 +528,10 @@ cc_library(
         "@com_google_absl//absl/synchronization",
         "@com_google_absl//absl/types:span",
         "@local_tsl//tsl/platform:errors",
-    ] + if_rocm_is_configured(["//xla/service/gpu:stream_executor_util"]),
+    ] + if_rocm_is_configured([
+        # keep sorted
+        "//xla/service/gpu:stream_executor_util",
+    ]),
 )
 
 cc_library(

From 7b3ae219f2e9a286927b5f2d4a257d49fda57795 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Fri, 17 May 2024 01:08:45 -0700
Subject: [PATCH 329/478] Automated Code Change

PiperOrigin-RevId: 634667656
---
 third_party/xla/xla/python/py_array.cc | 30 +++++++++++++-------------
 1 file changed, 15 insertions(+), 15 deletions(-)

diff --git a/third_party/xla/xla/python/py_array.cc b/third_party/xla/xla/python/py_array.cc
index 320f2bfdd20849..897872c9a68bcb 100644
--- a/third_party/xla/xla/python/py_array.cc
+++ b/third_party/xla/xla/python/py_array.cc
@@ -112,8 +112,8 @@ PjRtBuffer* GetPjrtBuffer(ifrt::Array* ifrt_array) {
   return arr->pjrt_buffers().front().get();
 }
 
-StatusOr<const Shape*> XlaDynamicShape(ifrt::Array* ifrt_array,
-                                       std::optional<Shape>& scratch) {
+absl::StatusOr<const Shape*> XlaDynamicShape(ifrt::Array* ifrt_array,
+                                             std::optional<Shape>& scratch) {
   auto* pjrt_buffer = GetPjrtBuffer(ifrt_array);
 
   if (!scratch) {
@@ -632,7 +632,7 @@ Status PyArray::set_arrays(nb::object obj) {
   return OkStatus();
 }
 
-StatusOr<PyArray> PyArray::FullyReplicatedShard() {
+absl::StatusOr<PyArray> PyArray::FullyReplicatedShard() {
   auto& cached = GetStorage().fully_replicated_array;
   if (!cached.is_none()) {
     return nb::cast<PyArray>(cached);
@@ -663,7 +663,7 @@ Status PyArray::BlockUntilReady() const {
   return AwaitBuffersReady(absl::MakeConstSpan(&ifrt_array, 1));
 }
 
-StatusOr<size_t> PyArray::GetOnDeviceSizeInBytes() {
+absl::StatusOr<size_t> PyArray::GetOnDeviceSizeInBytes() {
   if (ifrt_array() == nullptr) {
     return InvalidArgument(
         "GetOnDeviceSizeInBytes() called on deleted or donated buffer");
@@ -689,7 +689,7 @@ absl::Status PyArray::BlockUntilResultStatusIsReady() {
   return result_status.Await();
 }
 
-StatusOr<nb::object> PyArray::SingleDeviceArrayToNumpyArray() {
+absl::StatusOr<nb::object> PyArray::SingleDeviceArrayToNumpyArray() {
   TF_ASSIGN_OR_RETURN(auto arr, FullyReplicatedShard());
   auto result = arr.GetStorage().host_value.AsNumPyArray(
       arr.GetStorage().dynamic_shape, arr.ifrt_array());
@@ -703,7 +703,7 @@ Status PyArray::CopySingleDeviceArrayToHostAsync() {
       arr.GetStorage().dynamic_shape, arr.ifrt_array());
 }
 
-StatusOr<PyArray> PyArray::AssertUnsharded(std::string_view api) {
+absl::StatusOr<PyArray> PyArray::AssertUnsharded(std::string_view api) {
   if (ifrt_array() == nullptr) {
     return InvalidArgument("%s( called on deleted or donated buffer", api);
   }
@@ -719,7 +719,7 @@ StatusOr<PyArray> PyArray::AssertUnsharded(std::string_view api) {
   return py_arrays[0];
 }
 
-StatusOr<std::uintptr_t> PyArray::UnsafeBufferPointer() {
+absl::StatusOr<std::uintptr_t> PyArray::UnsafeBufferPointer() {
   TF_ASSIGN_OR_RETURN(auto arr, AssertUnsharded("UnsafeBufferPointer"));
 
   return py_client()->pjrt_client()->UnsafeBufferPointer(
@@ -800,9 +800,9 @@ nb::dict PyArray::CudaArrayInterface() {
   return result;
 }
 
-StatusOr<nb::object> CudaArrayInterfaceToBuffer(const nb::dict& cai,
-                                                nb_class_ptr<PyClient> client,
-                                                std::optional<int> device_id) {
+absl::StatusOr<nb::object> CudaArrayInterfaceToBuffer(
+    const nb::dict& cai, nb_class_ptr<PyClient> client,
+    std::optional<int> device_id) {
   if (!cai.contains("data")) {
     return absl::InvalidArgumentError(
         "CUDA Array Interface does not define `data`");
@@ -862,7 +862,7 @@ StatusOr<nb::object> CudaArrayInterfaceToBuffer(const nb::dict& cai,
       (version == 3 && (!cai.contains("stream") || cai["stream"].is_none()));
   TF_ASSIGN_OR_RETURN(
       std::intptr_t stream,
-      ([is_default_stream, cai, device]() -> StatusOr<std::intptr_t> {
+      ([is_default_stream, cai, device]() -> absl::StatusOr<std::intptr_t> {
         if (is_default_stream) {
           return device->GetStreamForExternalReadyEvents();
         } else {
@@ -987,8 +987,8 @@ PyArray::Storage::~PyArray_Storage() {
   ifrt_array.reset();
 }
 
-StatusOr<PyArray> PyArray::CopyToDeviceWithSharding(ifrt::DeviceList devices,
-                                                    nb::object dst_sharding) {
+absl::StatusOr<PyArray> PyArray::CopyToDeviceWithSharding(
+    ifrt::DeviceList devices, nb::object dst_sharding) {
   auto* ifrt_array_ptr = ifrt_array();
   ifrt::MemoryKind dst_memory_kind =
       CreateIfRtMemoryKindFromSharding(dst_sharding);
@@ -1069,7 +1069,7 @@ StatusOr<PyArray> PyArray::CopyToDeviceWithSharding(ifrt::DeviceList devices,
                  /*skip_checks=*/true, result_status());
 }
 
-StatusOr<PyArray> PyArray::BatchedDevicePut(
+absl::StatusOr<PyArray> PyArray::BatchedDevicePut(
     nb::object aval, nb::object sharding, std::vector<nb::object> xs,
     absl::Span<const PyDevice* const> dst_devices, bool committed,
     bool force_copy, PjRtClient::HostBufferSemantics host_buffer_semantics,
@@ -1407,7 +1407,7 @@ bool IsZeroCopyableCpuBuffer(const PjRtBuffer* buf) {
 PyHostValue::PyHostValue() = default;
 PyHostValue::~PyHostValue() = default;
 
-StatusOr<nb::object> PyHostValue::AsNumPyArray(
+absl::StatusOr<nb::object> PyHostValue::AsNumPyArray(
     std::optional<Shape>& dynamic_shape_holder, ifrt::Array* ifrt_array) {
   if (ifrt_array->IsDeleted()) {
     return InvalidArgument("DeviceArray has been deleted.");

From 3fd41a5abb6a2fae523ca9d5cf9847086ea79b5c Mon Sep 17 00:00:00 2001
From: Henning Becker <hebecker@google.com>
Date: Fri, 17 May 2024 01:09:53 -0700
Subject: [PATCH 330/478] ROCm dependency fixes in service/gpu

Adding some more direct dependencies to make the layering check happy.

PiperOrigin-RevId: 634667868
---
 third_party/xla/xla/service/gpu/BUILD | 19 +++++++++++++------
 1 file changed, 13 insertions(+), 6 deletions(-)

diff --git a/third_party/xla/xla/service/gpu/BUILD b/third_party/xla/xla/service/gpu/BUILD
index 2af7bf63b6dfe7..a28d711aa25f39 100644
--- a/third_party/xla/xla/service/gpu/BUILD
+++ b/third_party/xla/xla/service/gpu/BUILD
@@ -1731,16 +1731,20 @@ cc_library(
         "@local_tsl//tsl/platform:errors",
         "@local_tsl//tsl/platform:status",
         "@local_tsl//tsl/platform:statusor",
-    ] + if_cuda_is_configured([
+    ] + if_gpu_is_configured([
+        #keep sorted
+        "//xla/stream_executor:host_or_device_scalar",
+    ]) + if_cuda_is_configured([
+        #keep sorted
         "//xla/stream_executor/cuda:cublas_lt_header",
         "//xla/stream_executor/cuda:cublas_plugin",
         "@local_tsl//tsl/platform:tensor_float_32_hdr_lib",
-        "//xla/stream_executor:host_or_device_scalar",
     ]) + if_rocm_is_configured([
-        "//xla/stream_executor/rocm:hipblas_lt_header",
-        "//xla/stream_executor/rocm:amdhipblaslt_plugin",
-        "//xla/stream_executor:host_or_device_scalar",
+        #keep sorted
+        "//third_party/gpus/rocm:rocm_headers",
         "//xla/stream_executor/platform:dso_loader",
+        "//xla/stream_executor/rocm:amdhipblaslt_plugin",
+        "//xla/stream_executor/rocm:hipblas_lt_header",
     ]) + if_static([
         "@local_tsl//tsl/platform:tensor_float_32_utils",
     ]),
@@ -4010,7 +4014,10 @@ cc_library(
         "@local_tsl//tsl/platform:env",
         "@local_tsl//tsl/platform:errors",
         "@local_tsl//tsl/platform:statusor",
-    ],
+    ] + if_rocm_is_configured([
+        # keep sorted
+        "@local_config_rocm//rocm:rocm_headers",
+    ]),
 )
 
 cc_library(

From 2a66297f4fa1c5411e6f139a930803d9ba35710d Mon Sep 17 00:00:00 2001
From: Johannes Reifferscheid <jreiffers@google.com>
Date: Fri, 17 May 2024 01:11:09 -0700
Subject: [PATCH 331/478] Disable MultiOutputLoopFeedingMap with MLIR emitters.

We don't support nested fusions. This creates one because the
map gets wrapped in a fusion, and the computation already
contains a fusion.

PiperOrigin-RevId: 634668045
---
 third_party/xla/xla/tests/multioutput_fusion_test.cc | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/third_party/xla/xla/tests/multioutput_fusion_test.cc b/third_party/xla/xla/tests/multioutput_fusion_test.cc
index c2220217b12de2..845a0e2f6628ee 100644
--- a/third_party/xla/xla/tests/multioutput_fusion_test.cc
+++ b/third_party/xla/xla/tests/multioutput_fusion_test.cc
@@ -244,6 +244,11 @@ XLA_TEST_F(MultiOutputFusionTest,
 }
 
 XLA_TEST_F(MultiOutputFusionTest, MultiOutputLoopFeedingMap) {
+#ifdef XLA_TEST_BACKEND_GPU
+  if (GetDebugOptionsForTest().xla_gpu_enable_mlir_emitters()) {
+    GTEST_SKIP() << "Nested fusions not supported on GPU with MLIR emitters.";
+  }
+#endif
   const char* testcase = R"(
     HloModule m, is_scheduled=true
 

From e22b292005d9452516917ae4068cc4da491692e7 Mon Sep 17 00:00:00 2001
From: Thomas Joerg <tjoerg@google.com>
Date: Fri, 17 May 2024 01:43:01 -0700
Subject: [PATCH 332/478] Exclude //xla/service/gpu/... targets for CPU CI
 runs.

PiperOrigin-RevId: 634674988
---
 third_party/xla/.kokoro/linux/build.sh | 1 +
 1 file changed, 1 insertion(+)

diff --git a/third_party/xla/.kokoro/linux/build.sh b/third_party/xla/.kokoro/linux/build.sh
index f9a5e399fd8424..b1ac5b53a8c0aa 100644
--- a/third_party/xla/.kokoro/linux/build.sh
+++ b/third_party/xla/.kokoro/linux/build.sh
@@ -69,6 +69,7 @@ if is_linux_gpu_job ; then
 else
     TAGS_FILTER="$TAGS_FILTER,-gpu,-requires-gpu-nvidia"
     ADDITIONAL_FLAGS="$ADDITIONAL_FLAGS --config=nonccl"
+    TARGET_FILTERS="$TARGET_FILTERS -//xla/service/gpu/..."
 
     if is_linux_cpu_arm64_job ; then
         TAGS_FILTER="$TAGS_FILTER,-no_aarch64"

From 04386ec4f9bac821ce48cfbe21b1de9045f3ebe1 Mon Sep 17 00:00:00 2001
From: Adrian Kuegel <akuegel@google.com>
Date: Fri, 17 May 2024 01:44:40 -0700
Subject: [PATCH 333/478] Move unique_indices from Gather to Scatter where it
 belongs.

Somehow the attribute unique_indices ended up in the documentation for Gather,
but it is an attribute for Scatter.

PiperOrigin-RevId: 634675288
---
 third_party/xla/docs/operation_semantics.md | 12 ++++++------
 1 file changed, 6 insertions(+), 6 deletions(-)

diff --git a/third_party/xla/docs/operation_semantics.md b/third_party/xla/docs/operation_semantics.md
index d2828b642ad846..5a3626b2931e8b 100644
--- a/third_party/xla/docs/operation_semantics.md
+++ b/third_party/xla/docs/operation_semantics.md
@@ -1416,7 +1416,6 @@ For a more intuitive description, see the "Informal Description" section below.
 | `collapsed_slice_dims` | `ArraySlice<int64>` | The set of dimensions in each slice that are collapsed away. These dimensions must have size 1. |
 | `start_index_map`      | `ArraySlice<int64>` | A map that describes how to map indices in `start_indices` to legal indices into operand. |
 | `indices_are_sorted`   | `bool`              | Whether the indices are guaranteed to be sorted by the caller. |
-| `unique_indices`       | `bool`              | Whether the indices are guaranteed to be unique by the caller. |
 
 For convenience, we label dimensions in the output array not in `offset_dims`
 as `batch_dims`.
@@ -1486,11 +1485,6 @@ If `indices_are_sorted` is set to true then XLA can assume that `start_indices`
 are sorted (in ascending `start_index_map` order) by the user. If they are not
 then the semantics is implementation defined.
 
-If `unique_indices` is set to true then XLA can assume that all elements
-scattered to are unique. So XLA could use non-atomic operations. If
-`unique_indices` is set to true and the indices being scattered to are not
-unique then the semantics is implementation defined.
-
 ### Informal Description and Examples
 
 Informally, every index `Out` in the output array corresponds to an element `E`
@@ -2382,6 +2376,7 @@ Arguments                      | Type                  | Semantics
 `inserted_window_dims`         | `ArraySlice<int64>`   | The set of *window dimensions* that must be inserted into `updates` shape.
 `scatter_dims_to_operand_dims` | `ArraySlice<int64>`   | A dimensions map from the scatter indices to the operand index space. This array is interpreted as mapping `i` to `scatter_dims_to_operand_dims[i]` . It has to be one-to-one and total.
 `indices_are_sorted`           | `bool`                | Whether the indices are guaranteed to be sorted by the caller.
+`unique_indices`               | `bool`                | Whether the indices are guaranteed to be unique by the caller.
 
 Where:
 
@@ -2484,6 +2479,11 @@ If `indices_are_sorted` is set to true then XLA can assume that `start_indices`
 are sorted (in ascending `start_index_map` order) by the user. If they are not
 then the semantics is implementation defined.
 
+If `unique_indices` is set to true then XLA can assume that all elements
+scattered to are unique. So XLA could use non-atomic operations. If
+`unique_indices` is set to true and the indices being scattered to are not
+unique then the semantics is implementation defined.
+
 Informally, the scatter op can be viewed as an _inverse_ of the gather op, i.e.
 the scatter op updates the elements in the input that are extracted by the
 corresponding gather op.

From 2fdd8e9f181a6e21d4e0cb8ebc443e8fb65d14a6 Mon Sep 17 00:00:00 2001
From: Dmitri Gribenko <dmitrig@google.com>
Date: Fri, 17 May 2024 01:47:11 -0700
Subject: [PATCH 334/478] Integrate LLVM at llvm/llvm-project@5b7088c3619e

Updates LLVM usage to match
[5b7088c3619e](https://github.com/llvm/llvm-project/commit/5b7088c3619e)

PiperOrigin-RevId: 634675819
---
 third_party/llvm/workspace.bzl | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/third_party/llvm/workspace.bzl b/third_party/llvm/workspace.bzl
index 4dc1a665db1e39..7d72ce6255348f 100644
--- a/third_party/llvm/workspace.bzl
+++ b/third_party/llvm/workspace.bzl
@@ -4,8 +4,8 @@ load("//third_party:repo.bzl", "tf_http_archive")
 
 def repo(name):
     """Imports LLVM."""
-    LLVM_COMMIT = "c86a53d75995d65601d5745b2fceb2d50a970645"
-    LLVM_SHA256 = "a62373ace2a6f442b212ae16eced5e0d45d00fb04ae8f94e05dd7712661301ea"
+    LLVM_COMMIT = "5b7088c3619e95eaa8e154765470ddb4d3859fa4"
+    LLVM_SHA256 = "21c667d6ea0c4141ec02bc96e788d1cda2f55063626661cdc21665fb61025abc"
 
     tf_http_archive(
         name = name,

From 245ed12bc41aa740cbc1c93261292c06623dcec1 Mon Sep 17 00:00:00 2001
From: Henning Becker <hebecker@google.com>
Date: Fri, 17 May 2024 01:50:51 -0700
Subject: [PATCH 335/478] Fix ROCm XLA profiling code

- Add missing dependencies to make the layering check happy
- Add missing cast to function call to make latest clang happy
- Add missing mutex_lock to fix thread safety analysis error
- Replace flat_hash_map by a node_hash_map since the ValueType is not moveable or copyable.

PiperOrigin-RevId: 634676567
---
 .../xla/xla/backends/profiler/gpu/BUILD       | 22 ++++++++++++-------
 .../backends/profiler/gpu/rocm_collector.cc   | 11 +++++-----
 2 files changed, 20 insertions(+), 13 deletions(-)

diff --git a/third_party/xla/xla/backends/profiler/gpu/BUILD b/third_party/xla/xla/backends/profiler/gpu/BUILD
index ecf4b7647670e9..8f7429713fcd0c 100644
--- a/third_party/xla/xla/backends/profiler/gpu/BUILD
+++ b/third_party/xla/xla/backends/profiler/gpu/BUILD
@@ -222,23 +222,27 @@ tsl_gpu_library(
     visibility = ["//visibility:public"],
     deps = [
         "//xla/stream_executor/rocm:roctracer_wrapper",
+        "//xla/tsl/util:env_var",
         "@com_google_absl//absl/container:fixed_array",
         "@com_google_absl//absl/container:flat_hash_map",
+        "@com_google_absl//absl/container:flat_hash_set",
         "@com_google_absl//absl/container:node_hash_map",
         "@com_google_absl//absl/container:node_hash_set",
+        "@com_google_absl//absl/strings",
+        "@com_google_absl//absl/strings:str_format",
         "@com_google_absl//absl/types:optional",
-        "@local_tsl//tsl/platform:env",
+        "@local_tsl//tsl/platform:abi",
+        "@local_tsl//tsl/platform:env_time",
         "@local_tsl//tsl/platform:errors",
-        "@local_tsl//tsl/platform:logging",
         "@local_tsl//tsl/platform:macros",
-        "@local_tsl//tsl/platform:platform_port",
+        "@local_tsl//tsl/platform:mutex",
+        "@local_tsl//tsl/platform:status",
+        "@local_tsl//tsl/platform:thread_annotations",
+        "@local_tsl//tsl/platform:types",
         "@local_tsl//tsl/profiler/backends/cpu:annotation_stack",
         "@local_tsl//tsl/profiler/lib:profiler_factory",
         "@local_tsl//tsl/profiler/lib:profiler_interface",
-        "@local_tsl//tsl/profiler/protobuf:xplane_proto_cc",
         "@local_tsl//tsl/profiler/utils:parse_annotation",
-        "@local_tsl//tsl/profiler/utils:time_utils",
-        "@local_tsl//tsl/profiler/utils:trace_utils",
         "@local_tsl//tsl/profiler/utils:xplane_builder",
         "@local_tsl//tsl/profiler/utils:xplane_schema",
         "@local_tsl//tsl/profiler/utils:xplane_utils",
@@ -256,16 +260,18 @@ tsl_gpu_library(
         "//xla/stream_executor/rocm:roctracer_wrapper",
         "@com_google_absl//absl/container:fixed_array",
         "@com_google_absl//absl/container:flat_hash_map",
+        "@com_google_absl//absl/container:flat_hash_set",
         "@com_google_absl//absl/container:node_hash_map",
         "@com_google_absl//absl/container:node_hash_set",
-        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/types:optional",
         "@local_tsl//tsl/platform:env",
         "@local_tsl//tsl/platform:errors",
         "@local_tsl//tsl/platform:logging",
         "@local_tsl//tsl/platform:macros",
         "@local_tsl//tsl/platform:platform_port",
+        "@local_tsl//tsl/platform:status",
+        "@local_tsl//tsl/platform:types",
         "@local_tsl//tsl/profiler/backends/cpu:annotation_stack",
-        "@local_tsl//tsl/profiler/protobuf:xplane_proto_cc",
         "@local_tsl//tsl/profiler/utils:time_utils",
     ],
 )
diff --git a/third_party/xla/xla/backends/profiler/gpu/rocm_collector.cc b/third_party/xla/xla/backends/profiler/gpu/rocm_collector.cc
index 41b21c486eb340..f817a5e424d953 100644
--- a/third_party/xla/xla/backends/profiler/gpu/rocm_collector.cc
+++ b/third_party/xla/xla/backends/profiler/gpu/rocm_collector.cc
@@ -18,6 +18,7 @@ limitations under the License.
 
 #include "absl/container/fixed_array.h"
 #include "absl/container/flat_hash_set.h"
+#include "absl/container/node_hash_map.h"
 #include "absl/strings/str_cat.h"
 #include "absl/strings/str_format.h"
 #include "absl/strings/str_join.h"
@@ -223,8 +224,8 @@ class PerDeviceCollector {
     stats.occupancy_pct /= device_properties_.maxThreadsPerMultiProcessor;
 
     err = hipOccupancyMaxPotentialBlockSize(
-        &stats.min_grid_size, &stats.suggested_block_size, params.func_ptr,
-        params.dynamic_smem_size, 0);
+        &stats.min_grid_size, &stats.suggested_block_size,
+        static_cast<const void*>(params.func_ptr), params.dynamic_smem_size, 0);
 
     if (err != hipError_t::hipSuccess) {
       return {};
@@ -578,8 +579,7 @@ class RocmTraceCollectorImpl : public profiler::RocmTraceCollector {
         num_activity_events_(0),
         start_walltime_ns_(start_walltime_ns),
         start_gputime_ns_(start_gputime_ns),
-        num_gpus_(options.num_gpus),
-        per_device_collector_(options.num_gpus) {}
+        num_gpus_(options.num_gpus) {}
 
   void AddEvent(RocmTracerEvent&& event, bool is_auxiliary) override;
   void Flush() override;
@@ -615,7 +615,7 @@ class RocmTraceCollectorImpl : public profiler::RocmTraceCollector {
 
   const std::vector<RocmTracerEvent> ApiActivityInfoExchange();
 
-  absl::flat_hash_map<const uint32_t, PerDeviceCollector> per_device_collector_;
+  absl::node_hash_map<uint32_t, PerDeviceCollector> per_device_collector_;
 };
 //==========
 
@@ -730,6 +730,7 @@ RocmTraceCollectorImpl::ApiActivityInfoExchange() {
   std::vector<RocmTracerEvent> aggregated_events;
 
   // Copy info from activity events to API callback events
+  mutex_lock lock{event_maps_mutex_};
   for (auto& api_iter : api_events_map_) {
     RocmTracerEvent& api_event = api_iter.second;
     auto activity_event =

From 69c461ad23290062d7830c2659f8394044803f92 Mon Sep 17 00:00:00 2001
From: Sheng Yang <yang.sheng@intel.com>
Date: Fri, 17 May 2024 01:51:10 -0700
Subject: [PATCH 336/478] PR #11767: Add SYCL build script

Imported from GitHub PR https://github.com/openxla/xla/pull/11767

This is a sub PR of https://github.com/openxla/xla/pull/9042, it adds a script to build xla with `--config=sycl`.
Here is how to use:
```bash
workspace=$1
xla_path=$workspace/xla
cd $workspace
git clone -b yang/ci https://github.com/Intel-tensorflow/xla xla
bash $xla_path/build_tools/sycl/build.sh $workspace
```
Copybara import of the project:

--
d811bba1548f0c983da473a0940b1ea3d8b2c436 by Sheng, Yang <yang.sheng@intel.com>:

Add SYCL build script

--
0c2500b703b73bc340ccc9e3817c7ce54cc8d718 by Sheng, Yang <yang.sheng@intel.com>:

change path

Merging this change closes #11767

PiperOrigin-RevId: 634676627
---
 .../xla/build_tools/configure/configure.py    |  3 ++
 third_party/xla/build_tools/sycl/build.sh     | 46 +++++++++++++++++++
 third_party/xla/build_tools/sycl/build_xla.sh | 25 ++++++++++
 third_party/xla/build_tools/sycl/clean.sh     | 19 ++++++++
 .../xla/build_tools/sycl/install_bazel.sh     | 26 +++++++++++
 .../xla/build_tools/sycl/install_oneapi.sh    | 30 ++++++++++++
 6 files changed, 149 insertions(+)
 create mode 100644 third_party/xla/build_tools/sycl/build.sh
 create mode 100644 third_party/xla/build_tools/sycl/build_xla.sh
 create mode 100644 third_party/xla/build_tools/sycl/clean.sh
 create mode 100644 third_party/xla/build_tools/sycl/install_bazel.sh
 create mode 100644 third_party/xla/build_tools/sycl/install_oneapi.sh

diff --git a/third_party/xla/build_tools/configure/configure.py b/third_party/xla/build_tools/configure/configure.py
index 663e4b8724280d..2178a8b71dbc72 100755
--- a/third_party/xla/build_tools/configure/configure.py
+++ b/third_party/xla/build_tools/configure/configure.py
@@ -183,6 +183,7 @@ class Backend(ArgparseableEnum):
   CPU = enum.auto()
   CUDA = enum.auto()
   ROCM = enum.auto()
+  SYCL = enum.auto()
 
 
 class HostCompiler(ArgparseableEnum):
@@ -402,6 +403,8 @@ def to_bazelrc_lines(
         rc.append("build --config nonccl")
     elif self.backend == Backend.ROCM:
       pass
+    elif self.backend == Backend.SYCL:
+      rc.append("build --config sycl")
 
     # Lines that are added for every backend
     if dpav.ld_library_path:
diff --git a/third_party/xla/build_tools/sycl/build.sh b/third_party/xla/build_tools/sycl/build.sh
new file mode 100644
index 00000000000000..5ed6beb86c8a3d
--- /dev/null
+++ b/third_party/xla/build_tools/sycl/build.sh
@@ -0,0 +1,46 @@
+#!/usr/bin/env bash
+# Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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.
+# ==============================================================================
+#
+#
+# A script to build XLA sycl target.
+#
+# Required input:
+#     workspace: the local to do this buidl
+
+if [ $# -lt 1 ];then
+  echo "Error: workspace not set."
+  exit 1
+fi
+
+workspace=$1
+
+if [ -e ${workspace} ];then
+  time_stamp=$(date +%s%N)
+  echo "Warning: ${workspace} exist."
+  workspace=$workspace/$time_stamp
+  echo "Will use $workspace as new workspace"
+fi
+
+mkdir -p $workspace
+
+xla_path=$workspace/xla
+cd $workspace
+git clone -b yang/ci https://github.com/Intel-tensorflow/xla xla
+bash $xla_path/build_tools/sycl/install_bazel.sh $workspace
+bash $xla_path/build_tools/sycl/install_oneapi.sh $workspace install
+bash $xla_path/build_tools/sycl/build_xla.sh $workspace
+bash $xla_path/build_tools/sycl/clean.sh $workspace
+
diff --git a/third_party/xla/build_tools/sycl/build_xla.sh b/third_party/xla/build_tools/sycl/build_xla.sh
new file mode 100644
index 00000000000000..38e1c00e3dca21
--- /dev/null
+++ b/third_party/xla/build_tools/sycl/build_xla.sh
@@ -0,0 +1,25 @@
+#!/usr/bin/env bash
+# Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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.
+# ==============================================================================
+workspace=$1
+cd $workspace/xla
+
+if [ -z ${SYCL_TOOLKIT_PATH+x} ];
+then
+  export SYCL_TOOLKIT_PATH=$workspace/oneapi/compiler/2024.1/
+fi
+bazel_bin=$(ls $workspace/bazel/)
+./configure.py --backend=SYCL --host_compiler=GCC
+$workspace/bazel/$bazel_bin build --config=verbose_logs -s --verbose_failures --nocheck_visibility //xla/tools:run_hlo_module
diff --git a/third_party/xla/build_tools/sycl/clean.sh b/third_party/xla/build_tools/sycl/clean.sh
new file mode 100644
index 00000000000000..b72aa75a84a622
--- /dev/null
+++ b/third_party/xla/build_tools/sycl/clean.sh
@@ -0,0 +1,19 @@
+#!/usr/bin/env bash
+# Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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.
+# ==============================================================================
+workspace=$1
+SCRIPT_PATH="`dirname \"$0\"`"
+bash $SCRIPT_PATH/install_oneapi.sh $workspace remove
+rm -rf $workspace
diff --git a/third_party/xla/build_tools/sycl/install_bazel.sh b/third_party/xla/build_tools/sycl/install_bazel.sh
new file mode 100644
index 00000000000000..773dcee27cf2fa
--- /dev/null
+++ b/third_party/xla/build_tools/sycl/install_bazel.sh
@@ -0,0 +1,26 @@
+#!/usr/bin/env bash
+# Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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.
+# ==============================================================================
+workspace=$1
+cd $workspace/xla
+bazel_version=$(head -1 .bazelversion)
+bazel_base_url="https://github.com/bazelbuild/bazel/releases/download"
+bazel_url=$bazel_base_url/$bazel_version/"bazel-$bazel_version-linux-x86_64"
+echo $bazel_url
+mkdir -p $workspace/bazel
+cd $workspace/bazel
+wget $bazel_url
+bazel_bin=$(ls)
+chmod +x $bazel_bin
diff --git a/third_party/xla/build_tools/sycl/install_oneapi.sh b/third_party/xla/build_tools/sycl/install_oneapi.sh
new file mode 100644
index 00000000000000..67ffe93f107cbb
--- /dev/null
+++ b/third_party/xla/build_tools/sycl/install_oneapi.sh
@@ -0,0 +1,30 @@
+#!/usr/bin/env bash
+# Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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.
+# ==============================================================================
+if [ -z ${SYCL_TOOLKIT_PATH+x} ];
+then
+workspace=$1
+action=$2
+echo "Install Intel OneAPI in $workspace/oneapi"
+cd $workspace
+mkdir -p oneapi
+if ! [ -f $workspace/l_BaseKit_p_2024.1.0.596.sh ]; then
+  echo "Download oneAPI package"
+  wget https://registrationcenter-download.intel.com/akdlm/IRC_NAS/fdc7a2bc-b7a8-47eb-8876-de6201297144/l_BaseKit_p_2024.1.0.596.sh
+fi
+bash l_BaseKit_p_2024.1.0.596.sh -a -s --eula accept --action $action --install-dir $workspace/oneapi --log-dir $workspace/oneapi/log --download-cache $workspace/oneapi/cache --components=intel.oneapi.lin.dpcpp-cpp-compiler:intel.oneapi.lin.mkl.devel
+else
+  echo "SYCL_TOOLKIT_PATH set to $SYCL_TOOLKIT_PATH", skip install/remove oneAPI;
+fi

From d3ead792df5afb39028579b0aba60b76f29b834f Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Fri, 17 May 2024 02:08:25 -0700
Subject: [PATCH 337/478] compat: Update forward compatibility horizon to
 2024-05-17

PiperOrigin-RevId: 634680645
---
 tensorflow/python/compat/compat.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorflow/python/compat/compat.py b/tensorflow/python/compat/compat.py
index a236aaf853f105..4fc96bfdaa9f54 100644
--- a/tensorflow/python/compat/compat.py
+++ b/tensorflow/python/compat/compat.py
@@ -29,7 +29,7 @@
 # This value changes every day with an automatic CL. It can be modified in code
 # via `forward_compatibility_horizon()` or with the environment variable
 # TF_FORWARD_COMPATIBILITY_DELTA_DAYS, which is added to the compatibility date.
-_FORWARD_COMPATIBILITY_HORIZON = datetime.date(2024, 5, 16)
+_FORWARD_COMPATIBILITY_HORIZON = datetime.date(2024, 5, 17)
 _FORWARD_COMPATIBILITY_DELTA_DAYS_VAR_NAME = "TF_FORWARD_COMPATIBILITY_DELTA_DAYS"
 _FORWARD_COMPATIBILITY_DATE_NUMBER = None
 

From 4cdcd36fa7fc2c7624ee791f820c767388f93ecb Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Fri, 17 May 2024 02:09:34 -0700
Subject: [PATCH 338/478] Update GraphDef version to 1865.

PiperOrigin-RevId: 634680889
---
 tensorflow/core/public/version.h | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorflow/core/public/version.h b/tensorflow/core/public/version.h
index ea887161e0298c..79913c71835d8b 100644
--- a/tensorflow/core/public/version.h
+++ b/tensorflow/core/public/version.h
@@ -108,7 +108,7 @@ limitations under the License.
 
 #define TF_GRAPH_DEF_VERSION_MIN_PRODUCER 0
 #define TF_GRAPH_DEF_VERSION_MIN_CONSUMER 0
-#define TF_GRAPH_DEF_VERSION 1864  // Updated: 2024/5/16
+#define TF_GRAPH_DEF_VERSION 1865  // Updated: 2024/5/17
 
 // Checkpoint compatibility versions (the versions field in SavedSliceMeta).
 //

From fed68cd9555c732396c90ca064755c564c52987c Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Fri, 17 May 2024 02:16:07 -0700
Subject: [PATCH 339/478] Automated Code Change

PiperOrigin-RevId: 634682199
---
 .../kernels/batching_util/adaptive_shared_batch_scheduler.h | 4 ++--
 .../core/kernels/batching_util/shared_batch_scheduler.h     | 6 +++---
 2 files changed, 5 insertions(+), 5 deletions(-)

diff --git a/tensorflow/core/kernels/batching_util/adaptive_shared_batch_scheduler.h b/tensorflow/core/kernels/batching_util/adaptive_shared_batch_scheduler.h
index 28e8f4396fd668..9e6a4b2829aa5a 100644
--- a/tensorflow/core/kernels/batching_util/adaptive_shared_batch_scheduler.h
+++ b/tensorflow/core/kernels/batching_util/adaptive_shared_batch_scheduler.h
@@ -639,7 +639,7 @@ void AdaptiveSharedBatchScheduler<TaskType>::CallbackWrapper(
             "ProcessBatch", {{"batch_size_before_padding", batch->size()},
                              {"_r", 2} /*root_event*/});
       },
-      profiler::ContextType::kAdaptiveSharedBatchScheduler,
+      tsl::profiler::ContextType::kAdaptiveSharedBatchScheduler,
       batch->traceme_context_id());
   const int64_t start_time = batch->creation_time_micros();
   callback(std::unique_ptr<Batch<TaskType>>(
@@ -798,7 +798,7 @@ Status ASBSQueue<TaskType>::Schedule(std::unique_ptr<TaskType>* task) {
                 "ASBSQueue::Schedule",
                 {{"batching_input_task_size", task_size}});
           },
-          profiler::ContextType::kAdaptiveSharedBatchScheduler,
+          tsl::profiler::ContextType::kAdaptiveSharedBatchScheduler,
           this->current_batch_->traceme_context_id());
       current_batch_->AddTask(std::move(task));
       num_enqueued_tasks_++;
diff --git a/tensorflow/core/kernels/batching_util/shared_batch_scheduler.h b/tensorflow/core/kernels/batching_util/shared_batch_scheduler.h
index 7e638224369e82..500dfc662ee359 100644
--- a/tensorflow/core/kernels/batching_util/shared_batch_scheduler.h
+++ b/tensorflow/core/kernels/batching_util/shared_batch_scheduler.h
@@ -959,7 +959,7 @@ Status Queue<TaskType>::ScheduleWithLazySplit(std::unique_ptr<TaskType>* task) {
             return profiler::TraceMeEncode("ScheduleOutputTask",
                                            {{"size", task_handles[i]->size()}});
           },
-          profiler::ContextType::kSharedBatchScheduler,
+          tsl::profiler::ContextType::kSharedBatchScheduler,
           task_handle_batches_.back()->traceme_context_id());
 
       task_handle_batches_.back()->AddTask(std::move(task_handles[i]));
@@ -1040,7 +1040,7 @@ Status Queue<TaskType>::ScheduleWithoutOrEagerSplitImpl(
           return profiler::TraceMeEncode("ScheduleOutputTask",
                                          {{"size", output_tasks[i]->size()}});
         },
-        profiler::ContextType::kSharedBatchScheduler,
+        tsl::profiler::ContextType::kSharedBatchScheduler,
         batches.back()->traceme_context_id());
     batches.back()->AddTask(std::move(output_tasks[i]));
   }
@@ -1332,7 +1332,7 @@ void Queue<TaskType>::ProcessBatch(
             "ProcessBatch", {{"batch_size_before_padding", batch->size()},
                              {"_r", 2} /*root_event*/});
       },
-      profiler::ContextType::kSharedBatchScheduler,
+      tsl::profiler::ContextType::kSharedBatchScheduler,
       batch->traceme_context_id());
 
   if (std::holds_alternative<ProcessBatchCallbackWithoutPaddingTasks>(

From d065e49e407f22d1edf98fb231e982f8fdb260ad Mon Sep 17 00:00:00 2001
From: Henning Becker <hebecker@google.com>
Date: Fri, 17 May 2024 02:26:53 -0700
Subject: [PATCH 340/478] Fix various code style and build issues in
 stream_executor/rocm

- Adding namespace to `string` -> `std::string` (Not sure why this compiled in the first place)
- Remove unused local variables
- Avoid unintended fallthrough in switch statement
- Fix member variable initialization order in a constructor
- Avoid read of unintialized variable

PiperOrigin-RevId: 634684537
---
 .../xla/xla/stream_executor/rocm/BUILD        | 188 +++++++++++++-----
 .../xla/xla/stream_executor/rocm/rocm_blas.cc |   4 +-
 .../xla/xla/stream_executor/rocm/rocm_dnn.cc  |  43 ++--
 .../xla/stream_executor/rocm/rocm_driver.cc   |  20 +-
 .../xla/stream_executor/rocm/rocm_executor.cc |  17 +-
 .../xla/stream_executor/rocm/rocm_platform.cc |   2 +-
 .../xla/stream_executor/rocm/rocm_platform.h  |   4 +-
 7 files changed, 182 insertions(+), 96 deletions(-)

diff --git a/third_party/xla/xla/stream_executor/rocm/BUILD b/third_party/xla/xla/stream_executor/rocm/BUILD
index a629c4efbbbd68..d8ce19414a63fe 100644
--- a/third_party/xla/xla/stream_executor/rocm/BUILD
+++ b/third_party/xla/xla/stream_executor/rocm/BUILD
@@ -45,6 +45,7 @@ cc_library(
         "//xla/stream_executor/gpu:gpu_diagnostics_header",
         "//xla/stream_executor/platform",
         "@local_tsl//tsl/platform:platform_port",
+        "@local_tsl//tsl/platform:logging",
     ]),
 )
 
@@ -56,17 +57,25 @@ cc_library(
         "rocm_driver.h",
     ]),
     deps = if_rocm_is_configured([
+        # keep sorted
         ":rocm_diagnostics",
-        "@com_google_absl//absl/base",
-        "@com_google_absl//absl/container:inlined_vector",
-        "@com_google_absl//absl/container:node_hash_map",
-        "@com_google_absl//absl/strings",
         "//xla/stream_executor",
+        "//xla/stream_executor/gpu:gpu_diagnostics_header",
         "//xla/stream_executor/gpu:gpu_driver_header",
         "//xla/stream_executor/platform",
         "//xla/stream_executor/platform:dso_loader",
+        "@com_google_absl//absl/base",
+        "@com_google_absl//absl/container:inlined_vector",
+        "@com_google_absl//absl/container:node_hash_map",
+        "@com_google_absl//absl/memory",
+        "@com_google_absl//absl/strings",
+        "@com_google_absl//absl/strings:str_format",
+        "@com_google_absl//absl/synchronization",
+        "@local_config_rocm//rocm:hip",
         "@local_config_rocm//rocm:rocm_headers",
         "@local_tsl//tsl/platform:env",
+        "@local_tsl//tsl/platform:errors",
+        "@local_tsl//tsl/platform:logging",
         "@local_tsl//tsl/platform:numbers",
         "@local_tsl//tsl/platform:stacktrace",
     ]),
@@ -80,17 +89,24 @@ cc_library(
         "rocm_driver.h",
     ]),
     deps = if_rocm_is_configured([
-        "@com_google_absl//absl/base:core_headers",
-        "@com_google_absl//absl/log",
-        "@com_google_absl//absl/status",
-        "@com_google_absl//absl/strings",
+        # keep sorted
         "//xla/stream_executor",
+        "//xla/stream_executor/gpu:gpu_driver_header",
         "//xla/stream_executor/gpu:gpu_runtime_header",
         "//xla/stream_executor/gpu:gpu_types_header",
-        "//xla/stream_executor/gpu:gpu_driver_header",
         "//xla/stream_executor/platform",
         "//xla/stream_executor/platform:dso_loader",
+        "@com_google_absl//absl/base:core_headers",
+        "@com_google_absl//absl/container:node_hash_map",
+        "@com_google_absl//absl/log",
+        "@com_google_absl//absl/memory",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/status:statusor",
+        "@com_google_absl//absl/strings",
+        "@com_google_absl//absl/synchronization",
         "@local_config_rocm//rocm:rocm_headers",
+        "@local_tsl//tsl/platform:env",
+        "@local_tsl//tsl/platform:logging",
         "@local_tsl//tsl/platform:statusor",
     ]),
 )
@@ -99,10 +115,12 @@ cc_library(
     name = "rocm_collectives",
     srcs = if_rocm_is_configured(["rocm_collectives.cc"]),
     deps = if_rocm_is_configured([
-        "@com_google_absl//absl/status",
-        "@com_google_absl//absl/strings",
+        # keep sorted
         "//xla/stream_executor/gpu:gpu_collectives_header",
         "//xla/stream_executor/gpu:gpu_driver_header",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/status:statusor",
+        "@com_google_absl//absl/strings",
     ]),
 )
 
@@ -111,12 +129,13 @@ cc_library(
     srcs = [],
     hdrs = if_rocm_is_configured(["rocm_activation.h"]),
     deps = if_rocm_is_configured([
+        # keep sorted
         ":rocm_driver",
-        "@local_config_rocm//rocm:rocm_headers",
         "//xla/stream_executor",
         "//xla/stream_executor:stream_executor_interface",
         "//xla/stream_executor/gpu:gpu_activation",
         "//xla/stream_executor/platform",
+        "@local_config_rocm//rocm:rocm_headers",
     ]),
 )
 
@@ -124,6 +143,7 @@ cc_library(
     name = "rocm_event",
     srcs = if_rocm_is_configured(["rocm_event.cc"]),
     deps = if_rocm_is_configured([
+        # keep sorted
         ":rocm_driver",
         "//xla/stream_executor",
         "//xla/stream_executor/gpu:gpu_event_header",
@@ -136,30 +156,40 @@ cc_library(
     name = "rocm_executor",
     srcs = if_rocm_is_configured(["rocm_executor.cc"]),
     deps = if_rocm_is_configured([
+        # keep sorted
+        ":rocm_collectives",
         ":rocm_diagnostics",
         ":rocm_driver",
         ":rocm_event",
         ":rocm_kernel",
         ":rocm_platform_id",
         ":rocm_runtime",
-        ":rocm_collectives",
-        "@com_google_absl//absl/functional:any_invocable",
-        "@com_google_absl//absl/strings",
         "//xla/stream_executor",
         "//xla/stream_executor:kernel",
         "//xla/stream_executor:plugin_registry",
+        "//xla/stream_executor:stream_executor_headers",
         "//xla/stream_executor/gpu:gpu_activation_header",
+        "//xla/stream_executor/gpu:gpu_collectives_header",
+        "//xla/stream_executor/gpu:gpu_command_buffer",
+        "//xla/stream_executor/gpu:gpu_driver_header",
         "//xla/stream_executor/gpu:gpu_event",
+        "//xla/stream_executor/gpu:gpu_executor_header",
         "//xla/stream_executor/gpu:gpu_kernel_header",
-        "//xla/stream_executor/gpu:gpu_command_buffer",
         "//xla/stream_executor/gpu:gpu_runtime_header",
         "//xla/stream_executor/gpu:gpu_stream",
         "//xla/stream_executor/gpu:gpu_timer",
+        "//xla/stream_executor/integrations:device_mem_allocator",
         "//xla/stream_executor/platform",
         "//xla/stream_executor/platform:dso_loader",
-        "//xla/stream_executor/integrations:device_mem_allocator",
+        "@com_google_absl//absl/base",
+        "@com_google_absl//absl/functional:any_invocable",
+        "@com_google_absl//absl/strings",
+        "@com_google_absl//absl/strings:str_format",
         "@local_tsl//tsl/platform:env",
+        "@local_tsl//tsl/platform:errors",
         "@local_tsl//tsl/platform:fingerprint",
+        "@local_tsl//tsl/platform:logging",
+        "@local_tsl//tsl/platform:statusor",
     ]),
     alwayslink = True,
 )
@@ -186,16 +216,23 @@ cc_library(
     hdrs = if_rocm_is_configured(["rocm_platform.h"]),
     visibility = ["//visibility:public"],
     deps = if_rocm_is_configured([
+        # keep sorted
+        ":rocm_collectives",
         ":rocm_driver",
         ":rocm_executor",
         ":rocm_platform_id",
         ":rocm_runtime",
-        ":rocm_collectives",
-        "@com_google_absl//absl/base",
-        "@com_google_absl//absl/memory",
         "//xla/stream_executor",  # buildcleaner: keep
         "//xla/stream_executor:executor_cache",
+        "//xla/stream_executor:stream_executor_headers",
+        "//xla/stream_executor/gpu:gpu_driver_header",
+        "//xla/stream_executor/gpu:gpu_executor_header",
         "//xla/stream_executor/platform",
+        "@com_google_absl//absl/base",
+        "@com_google_absl//absl/memory",
+        "@com_google_absl//absl/strings:str_format",
+        "@com_google_absl//absl/synchronization",
+        "@local_tsl//tsl/platform:errors",
     ]),
     alwayslink = True,  # Registers itself with the PlatformManager.
 )
@@ -223,17 +260,19 @@ cc_library(
 
 cc_library(
     name = "rocblas_wrapper",
-    srcs = if_rocm_is_configured(["rocblas_wrapper.h"]),
     hdrs = if_rocm_is_configured(["rocblas_wrapper.h"]),
     deps = if_rocm_is_configured([
+        # keep sorted
         ":rocblas_if_static",
         ":rocm_executor",
         ":rocm_platform_id",
-        "@local_config_rocm//rocm:rocm_headers",
+        "//xla/stream_executor/gpu:gpu_activation",
         "//xla/stream_executor/platform",
         "//xla/stream_executor/platform:dso_loader",
-        "@local_tsl//tsl/platform:env",
         "//xla/tsl/util:determinism_for_kernels",
+        "@local_config_rocm//rocm:rocm_headers",
+        "@local_tsl//tsl/platform",
+        "@local_tsl//tsl/platform:env",
     ]),
     alwayslink = True,
 )
@@ -244,26 +283,35 @@ cc_library(
     hdrs = if_rocm_is_configured(["rocm_blas.h"]),
     visibility = ["//visibility:public"],
     deps = if_rocm_is_configured([
+        # keep sorted
+        ":hipblas_lt_header",
         ":rocblas_if_static",
         ":rocblas_wrapper",
-        ":hipblas_lt_header",
         ":rocm_executor",
+        ":rocm_helpers",
         ":rocm_platform_id",
-        "@eigen_archive//:eigen3",
         "//xla/stream_executor",
+        "//xla/stream_executor:blas",
         "//xla/stream_executor:host_or_device_scalar",
         "//xla/stream_executor:plugin_registry",
+        "//xla/stream_executor:stream_executor_interface",
         "//xla/stream_executor/gpu:gpu_activation",
+        "//xla/stream_executor/gpu:gpu_blas_lt",
+        "//xla/stream_executor/gpu:gpu_executor_header",
         "//xla/stream_executor/gpu:gpu_helpers_header",
         "//xla/stream_executor/gpu:gpu_stream_header",
         "//xla/stream_executor/gpu:gpu_timer",
-        "//xla/stream_executor/gpu:gpu_blas_lt",
         "//xla/stream_executor/platform",
-        "//xla/stream_executor:blas",
         "//xla/stream_executor/platform:dso_loader",
+        "//xla/tsl/util:determinism_hdr_lib",
+        "@com_google_absl//absl/base:core_headers",
         "@com_google_absl//absl/strings",
+        "@com_google_absl//absl/strings:str_format",
+        "@com_google_absl//absl/synchronization",
         "@com_google_absl//absl/types:span",
+        "@eigen_archive//:eigen3",
         "@local_config_rocm//rocm:rocm_headers",
+        "@local_tsl//tsl/platform:logging",
     ]),
     alwayslink = True,
 )
@@ -288,20 +336,23 @@ cc_library(
     hdrs = if_rocm_is_configured(["rocm_fft.h"]),
     visibility = ["//visibility:public"],
     deps = if_rocm_is_configured([
+        # keep sorted
         ":hipfft_if_static",
         ":rocm_platform_id",
         "//xla/stream_executor",
         "//xla/stream_executor:fft",
         "//xla/stream_executor:plugin_registry",
+        "//xla/stream_executor:stream_executor_headers",
         "//xla/stream_executor/gpu:gpu_activation",
-        "//xla/stream_executor/gpu:gpu_helpers_header",
         "//xla/stream_executor/gpu:gpu_executor_header",
-        "//xla/stream_executor/gpu:gpu_stream_header",
+        "//xla/stream_executor/gpu:gpu_helpers_header",
         "//xla/stream_executor/gpu:gpu_kernel_header",
+        "//xla/stream_executor/gpu:gpu_stream_header",
         "//xla/stream_executor/platform",
         "//xla/stream_executor/platform:dso_loader",
         "@local_config_rocm//rocm:rocm_headers",
         "@local_tsl//tsl/platform:env",
+        "@local_tsl//tsl/platform:logging",
     ]),
     alwayslink = True,
 )
@@ -331,29 +382,41 @@ cc_library(
     ],
     visibility = ["//visibility:public"],
     deps = if_rocm_is_configured([
+        # keep sorted
         ":miopen_if_static",
         ":rocm_diagnostics",
         ":rocm_driver",
         ":rocm_executor",
+        ":rocm_helpers",
         ":rocm_platform_id",
-        "@eigen_archive//:eigen3",
         "//xla/stream_executor",
+        "//xla/stream_executor:device_memory_allocator",
         "//xla/stream_executor:dnn",
         "//xla/stream_executor:plugin_registry",
-        "//xla/stream_executor/gpu:gpu_activation_header",
+        "//xla/stream_executor:stream_interface",
+        "//xla/stream_executor/gpu:gpu_activation",
+        "//xla/stream_executor/gpu:gpu_driver_header",
+        "//xla/stream_executor/gpu:gpu_executor_header",
         "//xla/stream_executor/gpu:gpu_stream_header",
         "//xla/stream_executor/gpu:gpu_timer",
         "//xla/stream_executor/gpu:gpu_types_header",
-        "//xla/stream_executor:device_memory_allocator",
         "//xla/stream_executor/platform",
         "//xla/stream_executor/platform:dso_loader",
+        "//xla/tsl/util:determinism_for_kernels",
+        "//xla/tsl/util:env_var",
         "@com_google_absl//absl/algorithm:container",
+        "@com_google_absl//absl/base:core_headers",
         "@com_google_absl//absl/strings",
+        "@com_google_absl//absl/strings:str_format",
+        "@com_google_absl//absl/synchronization",
         "@com_google_absl//absl/types:span",
+        "@eigen_archive//:eigen3",
         "@local_config_rocm//rocm:rocm_headers",
+        "@local_tsl//tsl/platform:env",
         "@local_tsl//tsl/platform:env_impl",
-        "//xla/tsl/util:env_var",
-        "//xla/tsl/util:determinism_for_kernels",
+        "@local_tsl//tsl/platform:errors",
+        "@local_tsl//tsl/platform:hash",
+        "@local_tsl//tsl/platform:logging",
     ]),
     alwayslink = True,
 )
@@ -448,7 +511,6 @@ cc_library(
 
 cc_library(
     name = "hipsolver_wrapper",
-    srcs = if_rocm_is_configured(["hipsolver_wrapper.h"]),
     hdrs = if_rocm_is_configured(["hipsolver_wrapper.h"]),
     deps = if_rocm_is_configured([
         ":rocm_executor",
@@ -478,21 +540,31 @@ cc_library(
         "hip_blas_utils.h",
     ]),
     deps = if_rocm_is_configured([
-        ":rocm_executor",
-        ":rocm_platform_id",
-        ":rocblas_plugin",
+        # keep sorted
         ":hip_blas_utils",
         ":hipblas_lt_header",
-        "@local_config_rocm//rocm:rocm_headers",
-        "@local_tsl//tsl/platform:status",
-        "@local_tsl//tsl/platform:errors",
+        ":rocblas_plugin",
+        ":rocm_executor",
+        ":rocm_platform_id",
+        "//xla:shape_util",
         "//xla:status",
-        "//xla/stream_executor/platform",
-        "//xla/stream_executor/platform:dso_loader",
+        "//xla:status_macros",
+        "//xla:types",
+        "//xla:util",
         "//xla/stream_executor",
-        "//xla/stream_executor/gpu:gpu_helpers_header",
+        "//xla/stream_executor:host_or_device_scalar",
+        "//xla/stream_executor/gpu:gpu_activation",
         "//xla/stream_executor/gpu:gpu_blas_lt",
-        "//xla:util",
+        "//xla/stream_executor/gpu:gpu_helpers_header",
+        "//xla/stream_executor/gpu:gpu_stream_header",
+        "//xla/stream_executor/gpu:gpu_timer",
+        "//xla/stream_executor/platform",
+        "//xla/stream_executor/platform:dso_loader",
+        "@com_google_absl//absl/status",
+        "@local_config_rocm//rocm:rocm_headers",
+        "@local_tsl//tsl/platform:env",
+        "@local_tsl//tsl/platform:errors",
+        "@local_tsl//tsl/platform:status",
     ]) + if_static([
         ":hipblaslt_if_static",
     ]),
@@ -508,13 +580,19 @@ cc_library(
     ]),
     visibility = ["//visibility:public"],
     deps = if_rocm_is_configured([
-        "@local_config_rocm//rocm:rocm_headers",
-        "@local_tsl//tsl/platform:status",
-        "@local_tsl//tsl/platform:errors",
+        # keep sorted
         "//xla:status",
-        "//xla/stream_executor:host_or_device_scalar",
+        "//xla:types",
         "//xla/stream_executor",
+        "//xla/stream_executor:host_or_device_scalar",
+        "//xla/stream_executor/gpu:gpu_blas_lt",
         "//xla/stream_executor/platform",
+        "//xla/stream_executor/platform:dso_loader",
+        "@com_google_absl//absl/status",
+        "@local_config_rocm//rocm:rocm_headers",
+        "@local_tsl//tsl/platform:env",
+        "@local_tsl//tsl/platform:errors",
+        "@local_tsl//tsl/platform:status",
     ]),
 )
 
@@ -523,13 +601,15 @@ cc_library(
     srcs = if_rocm_is_configured(["hip_blas_utils.cc"]),
     hdrs = if_rocm_is_configured(["hip_blas_utils.h"]),
     deps = if_rocm_is_configured([
-        ":rocblas_plugin",
+        # keep sorted
         ":hipblas_lt_header",
+        ":rocblas_plugin",
+        "//xla/stream_executor",
+        "@com_google_absl//absl/status",
         "@com_google_absl//absl/strings",
         "@local_config_rocm//rocm:rocm_headers",
-        "@local_tsl//tsl/platform:status",
         "@local_tsl//tsl/platform:errors",
-        "//xla/stream_executor",
+        "@local_tsl//tsl/platform:status",
     ]),
 )
 
@@ -552,12 +632,14 @@ cc_library(
     srcs = if_rocm_is_configured(["roctracer_wrapper.h"]),
     hdrs = if_rocm_is_configured(["roctracer_wrapper.h"]),
     deps = if_rocm_is_configured([
+        # keep sorted
         ":rocm_executor",
         ":rocm_platform_id",
         ":roctracer_if_static",
-        "@local_config_rocm//rocm:rocm_headers",
         "//xla/stream_executor/platform",
         "//xla/stream_executor/platform:dso_loader",
+        "@local_config_rocm//rocm:rocm_headers",
+        "@local_tsl//tsl/platform",
         "@local_tsl//tsl/platform:env",
     ]),
     alwayslink = True,
diff --git a/third_party/xla/xla/stream_executor/rocm/rocm_blas.cc b/third_party/xla/xla/stream_executor/rocm/rocm_blas.cc
index 8ca9f9ca9fc653..ec20077bde6e6f 100644
--- a/third_party/xla/xla/stream_executor/rocm/rocm_blas.cc
+++ b/third_party/xla/xla/stream_executor/rocm/rocm_blas.cc
@@ -75,7 +75,7 @@ RocBlasType_t<T> *complex_cast(DeviceMemory<T> *a) {
   return reinterpret_cast<RocBlasType_t<T> *>(GpuMemoryMutable(a));
 }
 
-static string ToString(rocblas_status status) {
+static std::string ToString(rocblas_status status) {
 #define XVAL(x) \
   case x:       \
     return #x
@@ -1251,7 +1251,7 @@ IMPL_DoBlasGemmBatched(float, wrap::rocblas_sgemm_strided_batched)
   }
 }
 
-absl::Status ROCMBlas::GetVersion(string *version) {
+absl::Status ROCMBlas::GetVersion(std::string *version) {
 #if TF_ROCM_VERSION >= 60300  // Not yet available in ROCM-6.1
   absl::MutexLock lock{&mu_};
   size_t len = 0;
diff --git a/third_party/xla/xla/stream_executor/rocm/rocm_dnn.cc b/third_party/xla/xla/stream_executor/rocm/rocm_dnn.cc
index cf5f9a9c80b72f..f4f5910771c03b 100644
--- a/third_party/xla/xla/stream_executor/rocm/rocm_dnn.cc
+++ b/third_party/xla/xla/stream_executor/rocm/rocm_dnn.cc
@@ -95,7 +95,7 @@ static absl::Status PopulateProfileFromTimer(
   return absl::OkStatus();
 }
 
-string ToString(miopenStatus_t status) {
+std::string ToString(miopenStatus_t status) {
   switch (status) {
     case miopenStatusSuccess:
       return "miopenStatusSuccess";
@@ -121,8 +121,8 @@ string ToString(miopenStatus_t status) {
   }
 }
 
-string ToString(miopenConvFwdAlgorithm_t algorithm) {
-  string s;
+std::string ToString(miopenConvFwdAlgorithm_t algorithm) {
+  std::string s;
   switch (algorithm) {
     case miopenConvolutionFwdAlgoGEMM:
       s = "GEMM";
@@ -143,8 +143,8 @@ string ToString(miopenConvFwdAlgorithm_t algorithm) {
   return s;
 }
 
-string ToString(miopenConvBwdWeightsAlgorithm_t algorithm) {
-  string s;
+std::string ToString(miopenConvBwdWeightsAlgorithm_t algorithm) {
+  std::string s;
   switch (algorithm) {
     case miopenConvolutionBwdWeightsAlgoGEMM:
       s = "GEMM";
@@ -162,8 +162,8 @@ string ToString(miopenConvBwdWeightsAlgorithm_t algorithm) {
   return s;
 }
 
-string ToString(miopenConvBwdDataAlgorithm_t algorithm) {
-  string s;
+std::string ToString(miopenConvBwdDataAlgorithm_t algorithm) {
+  std::string s;
   switch (algorithm) {
     case miopenConvolutionBwdDataAlgoGEMM:
       s = "GEMM";
@@ -187,8 +187,8 @@ string ToString(miopenConvBwdDataAlgorithm_t algorithm) {
   return s;
 }
 
-string ToString(miopenConvAlgorithm_t algorithm) {
-  string s;
+std::string ToString(miopenConvAlgorithm_t algorithm) {
+  std::string s;
   switch (algorithm) {
     case miopenConvolutionAlgoGEMM:
       s = "GEMM";
@@ -1908,6 +1908,9 @@ miopenDataType_t ToMIOpenDataType(
       return miopenHalf;
     case dnn::DataType::kInt8:
       if (data_layout == dnn::DataLayout::kBatchDepthYX) return miopenInt8;
+      LOG(FATAL)
+          << "The kInt8 data type only supports the kBatchDepthYX data layout!";
+      break;
     case dnn::DataType::kDouble:
       LOG(FATAL)
           << "Unsupported DNN data type: tf.float64 (dnn::DataType::kDouble)";
@@ -1963,12 +1966,12 @@ class MixinBase<void> {};
 
 }  // namespace
 
-#define RETURN_IF_MIOPEN_ERROR(STATUS, ...)                              \
-  if (!SE_PREDICT_TRUE((STATUS) == miopenStatusSuccess)) {               \
-    string error_msg = absl::StrCat(ToString(STATUS), " ", __VA_ARGS__); \
-    SetFailure(::absl::UnknownError(error_msg));                         \
-    LOG(ERROR) << error_msg;                                             \
-    return;                                                              \
+#define RETURN_IF_MIOPEN_ERROR(STATUS, ...)                                   \
+  if (!SE_PREDICT_TRUE((STATUS) == miopenStatusSuccess)) {                    \
+    std::string error_msg = absl::StrCat(ToString(STATUS), " ", __VA_ARGS__); \
+    SetFailure(::absl::UnknownError(error_msg));                              \
+    LOG(ERROR) << error_msg;                                                  \
+    return;                                                                   \
   }
 
 template <typename Base>
@@ -2202,7 +2205,7 @@ class MIOpenRnnSequenceTensorDescriptor
         data_type_(data_type) {
     miopenTensorDescriptor_t handle = nullptr;
     if (seq_length <= 0) {
-      string error_msg =
+      std::string error_msg =
           absl::StrCat("sequence length must be positive: ", seq_length);
       LOG(ERROR) << error_msg;
       SetFailure(absl::UnknownError(error_msg));
@@ -3819,8 +3822,8 @@ absl::Status MIOpenSupport::GetMIOpenConvolveAlgorithmsFindMode(
       break;
     }
     default: {
-      return absl::InternalError("Unexpected convolution kind " +
-                                 static_cast<int>(kind));
+      return absl::InternalError(absl::StrFormat(
+          "Unexpected convolution kind %d", static_cast<int>(kind)));
       break;
     }
   }
@@ -4985,14 +4988,13 @@ class RocmFusedConvRunner : public dnn::FusedConvRunner {
         leakyrelu_alpha_(leakyrelu_alpha),
         side_input_scale_f32_(static_cast<float>(side_input_scale)),
 
+        activation_mode_(activation),
         dnn_input_nd_(dnn_input_nd),
         dnn_output_nd_(dnn_output_nd),
         dnn_filter_(dnn_filter),
         dnn_bias_nd_(dnn_bias_nd),
         dnn_conv_(dnn_conv),
 
-        activation_mode_(activation),
-
         input_nd_(std::move(input_nd)),
         output_nd_(std::move(output_nd)),
         filter_(std::move(filter)),
@@ -5029,7 +5031,6 @@ class RocmFusedConvRunner : public dnn::FusedConvRunner {
     int d1 = 1, d2 = 1;
     bool bNCHW = (dnn_output_nd_.layout() != dnn::DataLayout::kBatchYXDepth);
     batch = dims_output[0];
-    int w, h = 1;
     if (bNCHW) {
       d1 = dims_output[1];
       for (int i = 2; i < rank; i++) d2 *= dims_output[i];
diff --git a/third_party/xla/xla/stream_executor/rocm/rocm_driver.cc b/third_party/xla/xla/stream_executor/rocm/rocm_driver.cc
index e2e1db3355d56d..ba4a7913ecfcb0 100644
--- a/third_party/xla/xla/stream_executor/rocm/rocm_driver.cc
+++ b/third_party/xla/xla/stream_executor/rocm/rocm_driver.cc
@@ -78,7 +78,7 @@ namespace gpu {
 
 // Formats hipError_t to output prettified values into a log stream.
 // Error summaries taken from:
-string ToString(hipError_t result) {
+std::string ToString(hipError_t result) {
 #define OSTREAM_ROCM_ERROR(__name) \
   case hipError##__name:           \
     return "HIP_ERROR_" #__name;
@@ -156,7 +156,7 @@ tsl::thread::ThreadPool* GetDriverExecutor() {
 
 }  // namespace
 
-string MemorySpaceString(MemorySpace memory_space) {
+std::string MemorySpaceString(MemorySpace memory_space) {
   switch (memory_space) {
     case MemorySpace::kHost:
       return "host";
@@ -251,7 +251,7 @@ namespace {
 // Returns a stringified device number associated with pointer, primarily for
 // logging purposes. Returns "?" if the device could not be successfully
 // queried.
-string ROCMPointerToDeviceString(hipDeviceptr_t pointer) {
+std::string ROCMPointerToDeviceString(hipDeviceptr_t pointer) {
   auto value = GpuDriver::GetPointerDevice(pointer);
   if (value.ok()) {
     return absl::StrCat(value.value());
@@ -263,7 +263,7 @@ string ROCMPointerToDeviceString(hipDeviceptr_t pointer) {
 // Returns a stringified memory space associated with pointer, primarily for
 // logging purposes. Returns "?" if the memory space could not be successfully
 // queried.
-string ROCMPointerToMemorySpaceString(hipDeviceptr_t pointer) {
+std::string ROCMPointerToMemorySpaceString(hipDeviceptr_t pointer) {
   auto value = GpuDriver::GetPointerMemorySpace(pointer);
   if (value.ok()) {
     return MemorySpaceString(value.value());
@@ -276,7 +276,8 @@ string ROCMPointerToMemorySpaceString(hipDeviceptr_t pointer) {
 // permitted between the "from" and "to" pointers' associated contexts,
 // primarily for logging purposes. Returns "error" if an error is encountered
 // in the process of querying.
-string ROCMPointersToCanAccessString(hipDeviceptr_t from, hipDeviceptr_t to) {
+std::string ROCMPointersToCanAccessString(hipDeviceptr_t from,
+                                          hipDeviceptr_t to) {
   auto from_context = GpuDriver::GetPointerContext(from);
   if (!from_context.ok()) {
     LOG(ERROR) << "could not retrieve source pointer's context: "
@@ -339,7 +340,7 @@ static absl::Status InternalInit() {
 }
 
 /* static */ absl::Status GpuDriver::GetDeviceName(hipDevice_t device,
-                                                   string* device_name) {
+                                                   std::string* device_name) {
   static const size_t kCharLimit = 64;
   absl::InlinedVector<char, 4> chars(kCharLimit);
   RETURN_IF_ROCM_ERROR(
@@ -1763,7 +1764,8 @@ struct BitPatternToValue {
 /* static */ absl::StatusOr<MemorySpace> GpuDriver::GetPointerMemorySpace(
     hipDeviceptr_t pointer) {
   unsigned int value;
-  hipError_t result = hipSuccess;
+  hipError_t result = wrap::hipPointerGetAttribute(
+      &value, HIP_POINTER_ATTRIBUTE_MEMORY_TYPE, pointer);
   if (result == hipSuccess) {
     switch (value) {
       case hipMemoryTypeDevice:
@@ -2057,8 +2059,8 @@ static absl::StatusOr<T> GetSimpleAttribute(hipDevice_t device,
   return true;
 }
 
-/* static */ string GpuDriver::GetPCIBusID(hipDevice_t device) {
-  string pci_bus_id;
+/* static */ std::string GpuDriver::GetPCIBusID(hipDevice_t device) {
+  std::string pci_bus_id;
   static const int kBufferSize = 64;
   absl::InlinedVector<char, 4> chars(kBufferSize);
   chars[kBufferSize - 1] = '\0';
diff --git a/third_party/xla/xla/stream_executor/rocm/rocm_executor.cc b/third_party/xla/xla/stream_executor/rocm/rocm_executor.cc
index 6096c2b1e85183..c9e81b77a94ea4 100644
--- a/third_party/xla/xla/stream_executor/rocm/rocm_executor.cc
+++ b/third_party/xla/xla/stream_executor/rocm/rocm_executor.cc
@@ -245,7 +245,7 @@ absl::Status GpuExecutor::Init() {
 // Arg: strip_exe: if true, remove the name of the executable itself from the
 //                 returned string. Example: calling this from /usr/bin/foo
 //                 would return /usr/bin.
-static string GetBinaryDir(bool strip_exe) {
+static std::string GetBinaryDir(bool strip_exe) {
   char exe_path[PATH_MAX] = {0};
   CHECK_NE(readlink("/proc/self/exe", exe_path, sizeof(exe_path) - 1), -1);
   // Make sure it's null-terminated:
@@ -253,8 +253,8 @@ static string GetBinaryDir(bool strip_exe) {
 
   if (strip_exe) {
     // The exe is the last component of the path, so remove one component.
-    string ret = exe_path;
-    std::vector<string> components = absl::StrSplit(exe_path, '/');
+    std::string ret = exe_path;
+    std::vector<std::string> components = absl::StrSplit(exe_path, '/');
     components.pop_back();
     return absl::StrJoin(components, "/");
   }
@@ -265,7 +265,7 @@ absl::Status GpuExecutor::GetKernel(const MultiKernelLoaderSpec& spec,
                                     Kernel* kernel) {
   GpuKernel* rocm_kernel = AsGpuKernel(kernel);
   hipModule_t module = nullptr;
-  const string* kernel_name;
+  const std::string* kernel_name;
 
   if (spec.has_cuda_cubin_in_memory()) {
     kernel_name = &spec.cuda_cubin_in_memory().kernel_name();
@@ -909,7 +909,8 @@ GpuContext* GpuExecutor::gpu_context() { return context_; }
 //
 // For anything more complicated/prod-focused than this, you'll likely want to
 // turn to gsys' topology modeling.
-static int TryToReadNumaNode(const string& pci_bus_id, int device_ordinal) {
+static int TryToReadNumaNode(const std::string& pci_bus_id,
+                             int device_ordinal) {
   VLOG(2) << "trying to read NUMA node for device ordinal: " << device_ordinal;
   static const int kUnknownNumaNode = -1;
 
@@ -983,7 +984,7 @@ GpuExecutor::CreateDeviceDescription(int device_ordinal) {
 
   {
     int version = GpuDriver::GetDriverVersion().value_or(-1);
-    string augmented_driver_version = absl::StrFormat(
+    std::string augmented_driver_version = absl::StrFormat(
         "%d (%s)", version,
         rocm::DriverVersionStatusToString(Diagnostician::FindDsoVersion())
             .c_str());
@@ -991,7 +992,7 @@ GpuExecutor::CreateDeviceDescription(int device_ordinal) {
   }
 
   {
-    string pci_bus_id = GpuDriver::GetPCIBusID(device);
+    std::string pci_bus_id = GpuDriver::GetPCIBusID(device);
 
     // Lower the hex characters to match sysfs.
     pci_bus_id = absl::AsciiStrToLower(pci_bus_id);
@@ -1040,7 +1041,7 @@ GpuExecutor::CreateDeviceDescription(int device_ordinal) {
   }
 
   {
-    string device_name;
+    std::string device_name;
     TF_RETURN_IF_ERROR(GpuDriver::GetDeviceName(device, &device_name));
     builder.set_name(device_name);
   }
diff --git a/third_party/xla/xla/stream_executor/rocm/rocm_platform.cc b/third_party/xla/xla/stream_executor/rocm/rocm_platform.cc
index 5ef0a43ce8bcbd..0ac3540c4e627d 100644
--- a/third_party/xla/xla/stream_executor/rocm/rocm_platform.cc
+++ b/third_party/xla/xla/stream_executor/rocm/rocm_platform.cc
@@ -102,7 +102,7 @@ int ROCmPlatform::VisibleDeviceCount() const {
   return GpuDriver::GetDeviceCount();
 }
 
-const string& ROCmPlatform::Name() const { return name_; }
+const std::string& ROCmPlatform::Name() const { return name_; }
 
 absl::StatusOr<std::unique_ptr<DeviceDescription>>
 ROCmPlatform::DescriptionForDevice(int ordinal) const {
diff --git a/third_party/xla/xla/stream_executor/rocm/rocm_platform.h b/third_party/xla/xla/stream_executor/rocm/rocm_platform.h
index 05934b1245ab6a..e25035da503d7c 100644
--- a/third_party/xla/xla/stream_executor/rocm/rocm_platform.h
+++ b/third_party/xla/xla/stream_executor/rocm/rocm_platform.h
@@ -59,7 +59,7 @@ class ROCmPlatform : public Platform {
   // Returns -1 as a sentinel on internal failure (and logs the error).
   int VisibleDeviceCount() const override;
 
-  const string& Name() const override;
+  const std::string& Name() const override;
 
   absl::StatusOr<std::unique_ptr<DeviceDescription>> DescriptionForDevice(
       int ordinal) const override;
@@ -77,7 +77,7 @@ class ROCmPlatform : public Platform {
   void InspectNumaNodes();
 
   // This platform's name.
-  string name_;
+  std::string name_;
 
   // mutex that guards internal state.
   mutable absl::Mutex mu_;

From eb2fec4c8012131ff68d2604a27d7dfb5217a865 Mon Sep 17 00:00:00 2001
From: Quentin Khan <qkhan@google.com>
Date: Fri, 17 May 2024 02:49:37 -0700
Subject: [PATCH 341/478] Fix compiler errors in TFLite XNNPack delegate
 weight_cache_test.

PiperOrigin-RevId: 634689677
---
 tensorflow/lite/delegates/xnnpack/weight_cache_test.cc | 8 +++++---
 1 file changed, 5 insertions(+), 3 deletions(-)

diff --git a/tensorflow/lite/delegates/xnnpack/weight_cache_test.cc b/tensorflow/lite/delegates/xnnpack/weight_cache_test.cc
index 3d88e1e2dd0f03..6ecbba3b17b8b2 100644
--- a/tensorflow/lite/delegates/xnnpack/weight_cache_test.cc
+++ b/tensorflow/lite/delegates/xnnpack/weight_cache_test.cc
@@ -60,7 +60,7 @@ constexpr const char kTempFileTemplate[] =
 class TempFileDesc {
  public:
   static constexpr struct AutoClose {
-  } kAutoCLose;
+  } kAutoCLose{};
 
 #if defined(_MSC_VER)
   TempFileDesc() : fd_() {
@@ -144,7 +144,8 @@ TEST(MMapHandleTest, MapExistingFileWorks) {
 
   TempFileDesc tmp_file;
   ASSERT_TRUE(tmp_file.IsOpen());
-  static_cast<void>(write(tmp_file.GetFd(), payload.c_str(), size(payload)));
+  ASSERT_EQ(write(tmp_file.GetFd(), payload.c_str(), size(payload)),
+            size(payload));
   tmp_file.Close();
 
   MMapHandle handle;
@@ -165,7 +166,8 @@ TEST(MMapHandleTest, MoveConstructs) {
 
   TempFileDesc tmp_file;
   ASSERT_TRUE(tmp_file.IsOpen());
-  static_cast<void>(write(tmp_file.GetFd(), payload.c_str(), size(payload)));
+  ASSERT_EQ(write(tmp_file.GetFd(), payload.c_str(), size(payload)),
+            size(payload));
   tmp_file.Close();
 
   MMapHandle handle;

From 8874baeb65f873cf4b63a65656279aff788b1ad9 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Tam=C3=A1s=20Danyluk?= <tdanyluk@google.com>
Date: Fri, 17 May 2024 02:51:08 -0700
Subject: [PATCH 342/478] [XLA:GPU] Support fusion of dynamic-slice into triton
 gemm.

This change fuses dynamic-slice into triton gemm kernels provided the slice is taken along the major-most dimension, leaving all other dimensions the same. No further fusion occurs in operands of the dynamic slice, meaning the resulting triton gemm must take in all operands of the dynamic slice as parameters.

Autotuning can handle dynamic-slice just fine, because of the instruction's semantics, which ensures that it never reads out of bounds (the offsets are clamped to a valid region).

The original author of this CL is jvstokes and I did a comprehensive refactor.

PiperOrigin-RevId: 634690002
---
 .../xla/xla/service/gpu/gemm_fusion_test.cc   |  96 ++++++
 .../xla/xla/service/gpu/ir_emitter_triton.cc  | 114 ++++++-
 .../xla/service/gpu/ir_emitter_triton_test.cc | 297 ++++++++++++++++++
 .../gpu/triton_fusion_analysis_test.cc        |  56 ++++
 .../service/gpu/triton_tiling_propagation.cc  |  61 +++-
 5 files changed, 609 insertions(+), 15 deletions(-)

diff --git a/third_party/xla/xla/service/gpu/gemm_fusion_test.cc b/third_party/xla/xla/service/gpu/gemm_fusion_test.cc
index d83bbb59541e5d..4a718df503a01c 100644
--- a/third_party/xla/xla/service/gpu/gemm_fusion_test.cc
+++ b/third_party/xla/xla/service/gpu/gemm_fusion_test.cc
@@ -258,6 +258,102 @@ ENTRY e {
   EXPECT_FALSE(GemmFusion(cc).Run(module.get()).value());
 }
 
+TEST_F(GemmFusionTest, DynamicSliceIsFused) {
+  TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr<VerifiedHloModule> module,
+                          ParseAndReturnVerifiedModule(R"(
+ENTRY e {
+  dot_lhs = f32[2,18] parameter(0)
+  dynamic_slice_input = f32[2,64,2] parameter(1)
+  start_index0 = s32[] parameter(2)
+  start_index1_2 = s32[] constant(0)
+  dynamic_slice = f32[1,64,2] dynamic-slice(dynamic_slice_input, start_index0, start_index1_2, start_index1_2),
+                  dynamic_slice_sizes={1,64,2}
+  reshape = f32[64,2] reshape(dynamic_slice)
+  ROOT dot = f16[18,64] dot(dot_lhs, reshape),
+             lhs_contracting_dims={0}, rhs_contracting_dims={1}
+})"));
+
+  EXPECT_TRUE(GemmFusion(se::CudaComputeCapability{
+                             se::CudaComputeCapability::AMPERE, 0})
+                  .Run(module.get())
+                  .value());
+  EXPECT_THAT(module->entry_computation()->root_instruction(),
+              GmockMatch((m::Fusion(m::Parameter(), m::Parameter(),
+                                    m::Parameter(), m::Constant()))));
+}
+
+TEST_F(GemmFusionTest, DynamicSlicesAreFusedEvenIfTheyShareIndices) {
+  TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr<VerifiedHloModule> module,
+                          ParseAndReturnVerifiedModule(R"(
+ENTRY e {
+  p0 = f32[2,64,2] parameter(0)
+  p1 = s32[] parameter(1)
+  p2 = s32[] parameter(2)
+  p3 = s32[] parameter(3)
+  ds0 = f32[1,64,2] dynamic-slice(p0, p1, p2, p3), dynamic_slice_sizes={1,64,2}
+  a = f32[64,2] reshape(ds0)
+  ds1 = f32[1,64,2] dynamic-slice(p0, p3, p2, p1), dynamic_slice_sizes={1,64,2}
+  b = f32[64,2] reshape(ds1)
+  ROOT d = f16[64,64] dot(a, b),
+           lhs_contracting_dims={1}, rhs_contracting_dims={1}
+})"));
+
+  EXPECT_TRUE(GemmFusion(se::CudaComputeCapability{
+                             se::CudaComputeCapability::AMPERE, 0})
+                  .Run(module.get())
+                  .value());
+  // TODO(b/339810582): Don't duplicate scalar parameters to dot fusions,
+  // because they are never tiled differently.
+  // TODO(b/339814210): Don't count scalar parameters towards dot fusion
+  // parameter limit.
+  EXPECT_THAT(
+      module->entry_computation()->root_instruction(),
+      GmockMatch((m::Fusion(m::Parameter(), m::Parameter(), m::Parameter(),
+                            m::Parameter(), m::Parameter(), m::Parameter(),
+                            m::Parameter(), m::Parameter()))));
+}
+
+TEST_F(GemmFusionTest, DoNotFuseDynamicSliceOfNonMajorFragments) {
+  TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr<VerifiedHloModule> module,
+                          ParseAndReturnVerifiedModule(R"(
+ENTRY e {
+  dot_lhs = f32[2,4]{1,0} parameter(0)
+  dynamic_slice_input = f32[4,5,2]{2,1,0} parameter(1)
+  c0 = s32[] constant(0)
+  c2 = s32[] constant(2)
+  dynamic_slice = f32[4,1,2]{2,1,0} dynamic-slice(dynamic_slice_input, c0, c2, c0),
+                  dynamic_slice_sizes={4,1,2}
+  reshape = f32[4,2]{1,0} reshape(dynamic_slice)
+  ROOT dot = f32[4,4]{1,0} dot(dot_lhs, reshape),
+             lhs_contracting_dims={0}, rhs_contracting_dims={1}
+})"));
+  const se::CudaComputeCapability cc{se::CudaComputeCapability::AMPERE, 0};
+  // FusionDecision "Unsupported dynamic slice on non-major-most dimension."
+  EXPECT_FALSE(GemmFusion(cc).Run(module.get()).value());
+}
+
+TEST_F(GemmFusionTest, CanFuseDynamicSliceOfContractingDimIfItIsMajor) {
+  TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr<VerifiedHloModule> module,
+                          ParseAndReturnVerifiedModule(R"(
+ENTRY e {
+  dot_lhs = f32[2,4]{1,0} parameter(0)
+  dynamic_slice_input = f32[5,5]{1,0} parameter(1)
+  start_index0 = s32[] constant(2)
+  start_index1 = s32[] constant(0)
+  dynamic_slice = f32[2,5]{1,0} dynamic-slice(dynamic_slice_input, start_index0, start_index1),
+                  dynamic_slice_sizes={2,5}
+  ROOT d = f32[4,5]{1,0} dot(dot_lhs, dynamic_slice),
+           lhs_contracting_dims={0}, rhs_contracting_dims={0}
+})"));
+  EXPECT_TRUE(GemmFusion(se::CudaComputeCapability{
+                             se::CudaComputeCapability::AMPERE, 0})
+                  .Run(module.get())
+                  .value());
+  EXPECT_THAT(module->entry_computation()->root_instruction(),
+              GmockMatch((m::Fusion(m::Parameter(), m::Parameter(),
+                                    m::Constant(), m::Constant()))));
+}
+
 TEST_F(GemmFusionTest, SliceToDegenerateIsSkipped) {
   TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr<VerifiedHloModule> module,
                           ParseAndReturnVerifiedModule(R"(
diff --git a/third_party/xla/xla/service/gpu/ir_emitter_triton.cc b/third_party/xla/xla/service/gpu/ir_emitter_triton.cc
index 3692c5e51d4207..bd4412efcee26b 100644
--- a/third_party/xla/xla/service/gpu/ir_emitter_triton.cc
+++ b/third_party/xla/xla/service/gpu/ir_emitter_triton.cc
@@ -886,12 +886,14 @@ absl::StatusOr<Value> EmitScope(
     absl::flat_hash_map<const HloInstruction*, Value>& values) {
   for (const HloInstruction* hlo : instructions) {
     Value result;
-    if (hlo->opcode() == HloOpcode::kConcatenate) {
+    if (hlo->opcode() == HloOpcode::kConcatenate ||
+        hlo->opcode() == HloOpcode::kDynamicSlice) {
       // Parameter loads and their concatenations are handled outside EmitScope.
       TF_RET_CHECK(values.contains(hlo)) << hlo->ToString();
       continue;
     } else if (hlo->opcode() == HloOpcode::kParameter) {
-      if (hlo->users()[0]->opcode() == HloOpcode::kConcatenate) {
+      if (hlo->users()[0]->opcode() == HloOpcode::kConcatenate ||
+          hlo->users()[0]->opcode() == HloOpcode::kDynamicSlice) {
         continue;
       }
       TF_RET_CHECK(values.contains(hlo)) << hlo->ToString();
@@ -1376,6 +1378,18 @@ class MatMulEmitterHelper {
         values);
   }
 
+  int64_t GetNonContractingDimIdxForOperandScope(
+      TritonFusionAnalysis::Scope scope) {
+    if (scope == TritonFusionAnalysis::Scope::LHS) {
+      return dims_.lhs_noncontracting_dim_idx;
+    } else if (scope == TritonFusionAnalysis::Scope::RHS) {
+      return dims_.rhs_noncontracting_dim_idx;
+    } else {
+      CHECK(false) << "This shouldn't be called for the output scope.";
+    }
+  }
+
+  // bases: The base pointers of each argument.
   absl::StatusOr<Value> EmitTensorPointer(
       const HloInstruction* hlo, const Side& side, ValueRange bases,
       Value pid_k, std::vector<int32_t>& boundary_checks) {
@@ -1491,7 +1505,74 @@ class MatMulEmitterHelper {
             select_value,
             EmitMultiSelect(b_, pid_offset, concat_boundaries, input_bounds));
         bounds.push_back(select_value);
+        tensor_offsets.push_back(Cst32(specs.front()->at(0).slice_start));
+      } else if (hlo->opcode() == HloOpcode::kDynamicSlice &&
+                 (side.scope == TritonFusionAnalysis::Scope::LHS ||
+                  side.scope == TritonFusionAnalysis::Scope::RHS) &&
+                 properties.index ==
+                     GetNonContractingDimIdxForOperandScope(side.scope)) {
+        // Here we compute the offset of where we should read the slice from.
+        // TODO(b/323255699): Add support for slices of the contracting dim.
+        // Dynamic slices are guaranteed to only be offset along the majormost
+        // dimension.
+
+        // The only fragment of the non-contracting dim of the dot's input in
+        // the current scope:
+        TF_RET_CHECK(specs.back()->size() == 1);
+        const TensorIterationSpec::IterationSpecFragment
+            only_fragment_of_nc_dim = specs.back()->at(0);
+        // The majormost dim index in the dynamic slice's output.
+        const int majormost_dim = hlo->shape().layout().minor_to_major().back();
+
+        // dynamic slice operands are (input, start_index0, start_index1, ...)
+        // so the start index corresponding to the ith dimension is bases[i+1].
+        Value majormost_dim_start_index_ptr_val = bases[majormost_dim + 1];
+        Value majormost_dim_start_index_val = b_.create<mt::LoadOp>(
+            majormost_dim_start_index_ptr_val, mt::CacheModifier::NONE,
+            mt::EvictionPolicy::NORMAL,
+            /*isVolatile=*/false);
+        Value majormost_dim_start_index_lower_limit_val =
+            CreateConst(b_, majormost_dim_start_index_val.getType(), 0);
+        int64_t majormost_dim_start_index_upper_limit =
+            hlo->operand(0)->shape().dimensions(majormost_dim) -
+            hlo->dynamic_slice_sizes().at(majormost_dim);
+        Value majormost_dim_start_index_upper_limit_val =
+            CreateConst(b_, majormost_dim_start_index_val.getType(),
+                        majormost_dim_start_index_upper_limit);
+        // Our Triton codegen only supports signed integers so far.
+        majormost_dim_start_index_val =
+            b_.create<ma::MaxSIOp>(majormost_dim_start_index_val,
+                                   majormost_dim_start_index_lower_limit_val);
+        majormost_dim_start_index_val =
+            b_.create<ma::MinSIOp>(majormost_dim_start_index_val,
+                                   majormost_dim_start_index_upper_limit_val);
+
+        // How many "rows" (non-contracting dim values) are there in a slice of
+        // size 1?
+        int64_t rows_per_majormost_dim = 1;
+        for (int i = 0; i < hlo->shape().dimensions().size() - 1; ++i) {
+          rows_per_majormost_dim *= hlo->shape().dimensions_minor(i);
+        }
+        rows_per_majormost_dim =
+            rows_per_majormost_dim / only_fragment_of_nc_dim.stride;
+        Value rows_per_majormost_dim_val = Cst32(rows_per_majormost_dim);
+
+        Value tensor_offset_val_i32 = b_.create<ma::MulIOp>(
+            majormost_dim_start_index_val, rows_per_majormost_dim_val);
+        tensor_offsets.push_back(tensor_offset_val_i32);
+
+        // tt.make_tensor_ptr expects an i64 for shape and size, but expects
+        // i32 for offsets. We extend the offset to calculate the upper bound.
+        Value tensor_offset_val_i64 =
+            b_.create<ma::ExtSIOp>(i64_ty_, tensor_offset_val_i32);
+        Value sliced_count_val = Cst64(only_fragment_of_nc_dim.sliced_count);
+        Value upper_bound_val =
+            b_.create<ma::AddIOp>(tensor_offset_val_i64, sliced_count_val);
+        bounds.push_back(upper_bound_val);
+
+        block_offsets.push_back(pid_offset);
       } else {
+        tensor_offsets.push_back(Cst32(specs.front()->at(0).slice_start));
         block_offsets.push_back(pid_offset);
         int64_t count = specs.front()->at(0).count;
         if (side.scope == TritonFusionAnalysis::Scope::OUTPUT &&
@@ -1507,7 +1588,6 @@ class MatMulEmitterHelper {
           boundary_checks.push_back(bounds.size() - 1);
         }
       }
-      tensor_offsets.push_back(Cst32(specs.front()->at(0).slice_start));
       block_dims.push_back(properties.block_size);
       dim_order.emplace(dim_order.begin(), dim_order.size());
       return absl::OkStatus();
@@ -1651,13 +1731,17 @@ absl::StatusOr<LaunchDimensions> GetMatMulLaunchDimensions(
   return launch_config.launch_dims;
 }
 
-SmallVector<Value> GetArguments(mlir::triton::FuncOp fn,
-                                const HloInstruction& input) {
+absl::StatusOr<SmallVector<Value>> GetArguments(mlir::triton::FuncOp fn,
+                                                const HloInstruction& input) {
   if (input.opcode() == HloOpcode::kParameter) {
-    return {fn.getArgument(input.parameter_number())};
-  } else if (input.opcode() == HloOpcode::kConcatenate) {
+    return {{fn.getArgument(input.parameter_number())}};
+  } else if (input.opcode() == HloOpcode::kConcatenate ||
+             input.opcode() == HloOpcode::kDynamicSlice) {
+    // As defined in GemmFusion, all inputs of concatenate and dynamic slice are
+    // parameters.
     SmallVector<Value> result;
     for (const HloInstruction* operand : input.operands()) {
+      TF_RET_CHECK(operand->opcode() == HloOpcode::kParameter);
       result.push_back(fn.getArgument(operand->parameter_number()));
     }
     return result;
@@ -1674,7 +1758,8 @@ ConstHloInstructionSet ScopeInputs(const TritonFusionAnalysis& analysis,
   ConstHloInstructionSet result;
   for (const HloInstruction* parameter : analysis.ScopeParameters(scope)) {
     if (absl::c_any_of(parameter->users(), [](const HloInstruction* user) {
-          return user->opcode() == HloOpcode::kConcatenate;
+          return user->opcode() == HloOpcode::kConcatenate ||
+                 user->opcode() == HloOpcode::kDynamicSlice;
         })) {
       // Concatenation is always the only user of its parameters by
       // construction.
@@ -2188,9 +2273,11 @@ absl::Status EmitMatMul(mlir::OpBuilder builder,
     for (const HloInstruction* input : ScopeInputs(analysis, side.scope)) {
       TF_RET_CHECK(
           iter_args_to_inputs.insert({iter_args.size(), input}).second);
+      TF_ASSIGN_OR_RETURN(SmallVector<Value> arguments,
+                          GetArguments(fn, *input));
       TF_ASSIGN_OR_RETURN(Value tensor_ptr,
                           emitter.EmitTensorPointer(
-                              input, side, GetArguments(fn, *input), pid_k,
+                              input, side, arguments, pid_k,
                               iter_args_to_boundary_checks[iter_args.size()]));
       iter_args.push_back(tensor_ptr);
     }
@@ -2214,10 +2301,11 @@ absl::Status EmitMatMul(mlir::OpBuilder builder,
     for (const HloInstruction* input :
          ScopeInputs(analysis, TritonFusionAnalysis::Scope::OUTPUT)) {
       std::vector<int32_t> boundary_checks;
-      TF_ASSIGN_OR_RETURN(
-          Value tensor_pointer,
-          emitter.EmitTensorPointer(input, out, GetArguments(fn, *input), pid_k,
-                                    boundary_checks));
+      TF_ASSIGN_OR_RETURN(SmallVector<Value> arguments,
+                          GetArguments(fn, *input));
+      TF_ASSIGN_OR_RETURN(Value tensor_pointer,
+                          emitter.EmitTensorPointer(input, out, arguments,
+                                                    pid_k, boundary_checks));
       TF_RET_CHECK(values_out
                        .insert({input, EmitParameterLoad(b, tensor_pointer,
                                                          boundary_checks)})
diff --git a/third_party/xla/xla/service/gpu/ir_emitter_triton_test.cc b/third_party/xla/xla/service/gpu/ir_emitter_triton_test.cc
index d98bda3e05fb70..966716b6201b35 100644
--- a/third_party/xla/xla/service/gpu/ir_emitter_triton_test.cc
+++ b/third_party/xla/xla/service/gpu/ir_emitter_triton_test.cc
@@ -15,6 +15,7 @@ limitations under the License.
 
 #include "xla/service/gpu/ir_emitter_triton.h"
 
+#include <cstdlib>
 #include <iterator>
 #include <limits>
 #include <memory>
@@ -27,7 +28,9 @@ limitations under the License.
 #include <gtest/gtest.h>
 #include "absl/algorithm/container.h"
 #include "absl/status/status.h"
+#include "absl/strings/str_cat.h"
 #include "absl/strings/string_view.h"
+#include "absl/strings/substitute.h"
 #include "absl/types/span.h"
 #include "llvm/IR/LLVMContext.h"
 #include "llvm/Support/raw_ostream.h"
@@ -117,6 +120,12 @@ class TritonGemmTest : public TritonTest {
     debug_options.set_xla_gpu_gemm_rewrite_size_threshold(0);
     return debug_options;
   }
+
+  void MatchHloModule(HloModule& module, absl::string_view pattern) {
+    TF_ASSERT_OK_AND_ASSIGN(bool filecheck_result,
+                            RunFileCheck(module.ToString(), pattern));
+    EXPECT_TRUE(filecheck_result);
+  }
 };
 
 class TritonGemmTestWithSplitK : public TritonGemmTest {
@@ -1208,6 +1217,59 @@ CHECK:     %[[BLOCK_BASE_PTR:.*]] = tt.addptr %[[ARG_PTR]], %[[OFFSET]]
 )"));
 }
 
+TEST_F(TritonFilecheckTest, CodegenDynamicSliceWithCorrectOffsets) {
+  // The start index(es) for the non-majormost dimension(s) are constant zero(s)
+  // because we don't support dynamic slice on those dimensions.
+  constexpr absl::string_view kHloText = R"(
+HloModule t
+
+triton_gemm {
+  dot_lhs = f32[2,4] parameter(0)
+  dynamic_slice_input = f32[4,5,2] parameter(1)
+  start_index0 = s32[] parameter(2)
+  start_index1 = s32[] parameter(3)
+  start_index2 = s32[] parameter(4)
+  dynamic_slice = f32[1,5,2] dynamic-slice(dynamic_slice_input, start_index0, start_index1, start_index2), dynamic_slice_sizes={1,5,2}
+  bitcast = f32[5,2] bitcast(dynamic_slice)
+  ROOT dot = f32[4,5] dot(dot_lhs, bitcast), lhs_contracting_dims={0}, rhs_contracting_dims={1}
+}
+
+ENTRY e {
+  dot_lhs = f32[2,4] parameter(0)
+  dynamic_slice_input = f32[4,5,2] parameter(1)
+  start_index0 = s32[] parameter(2)
+  start_index1 = s32[] constant(0)
+  start_index2 = s32[] constant(0)
+  ROOT fusion = f32[4,5] fusion(dot_lhs, dynamic_slice_input, start_index0, start_index1, start_index2),
+       kind=kCustom, calls=triton_gemm,
+       backend_config={
+         "fusion_backend_config":{
+           "kind":"__triton_gemm","triton_gemm_config":{
+             "block_m":"32","block_n":"32","block_k":"32","split_k":"1",
+             "num_stages":"1","num_warps":"4","num_ctas":"1"}}}
+})";
+
+  TritonGemmConfig config(16, 64, 32, 1, 1, 2);
+  ASSERT_THAT(CreateTritonIrAndFileCheck(kHloText, config, EmitMatMul,
+                                         "triton_gemm", R"(
+CHECK:     tt.func @triton_fn({{[^,]*}}, %[[DYNAMIC_SLICE_INPUT:[^:]*]]: !tt.ptr<f32> {{[^,]*}}, %[[START_INDEX0_PTR:[^:]*]]: !tt.ptr<i32>
+CHECK-DAG:   %[[C0_i32:.*]] = arith.constant 0 : i32
+CHECK-DAG:   %[[C1_i64:.*]] = arith.constant 1 : i64
+CHECK-DAG:   %[[C2_i64:.*]] = arith.constant 2 : i64
+CHECK-DAG:   %[[C3_i32:.*]] = arith.constant 3 : i32
+CHECK-DAG:   %[[C5_i32:.*]] = arith.constant 5 : i32
+CHECK-DAG:   %[[C5_i64:.*]] = arith.constant 5 : i64
+CHECK-DAG:   %[[START_INDEX0:.*]] = tt.load %[[START_INDEX0_PTR]] : !tt.ptr<i32>
+CHECK-DAG:   %[[SEMI_CLAMPED_START_INDEX0:.*]] = arith.maxsi %[[START_INDEX0]], %[[C0_i32]] : i32
+CHECK-DAG:   %[[CLAMPED_START_INDEX0:.*]] = arith.minsi %[[SEMI_CLAMPED_START_INDEX0]], %[[C3_i32]] : i32
+CHECK-DAG:   %[[ROW_OFFSET:.*]] = arith.muli %[[CLAMPED_START_INDEX0]], %[[C5_i32]] : i32
+CHECK-DAG:   %[[ROW_OFFSET_i64:.*]] = arith.extsi %[[ROW_OFFSET]] : i32 to i64
+CHECK-DAG:   %[[ROW_LIMIT:.*]] = arith.addi %[[ROW_OFFSET_i64]], %[[C5_i64]] : i64
+CHECK-DAG:   tt.make_tensor_ptr %[[DYNAMIC_SLICE_INPUT]], [%[[C2_i64]], %[[ROW_LIMIT]]], [%[[C1_i64]], %[[C2_i64]]], [%[[C0_i32]], %[[ROW_OFFSET]]]
+)"),
+              tsl::testing::IsOk());
+}
+
 TEST_F(TritonFilecheckTest, SparseDot) {
   const char* kHloText = R"(
 HloModule t
@@ -2173,6 +2235,241 @@ ENTRY e {
   EXPECT_TRUE(RunAndCompare(hlo_text, ErrorSpec{/*aabs=*/1e-6, /*arel=*/1e-6}));
 }
 
+TEST_F(TritonGemmTest, DynamicSliceIsSupportedInLhsEndToEnd) {
+  // The select is used to restrict the start index to values that make sense.
+  // If it was constant, then the dynamic-slice would be optimized to slice. It
+  // is not strictly needed, because we also support clamping the indices.
+  // The start index(es) for the non-majormost dimension(s) are constant zero(s)
+  // because we don't support dynamic slice on those dimensions.
+  constexpr absl::string_view kHloText = R"(
+HloModule m
+
+ENTRY e {
+  dot_lhs = f32[2,4] parameter(0)
+  dynamic_slice_input = f32[7,2] parameter(1)
+  pred0 = pred[] parameter(2)
+  c1 = s32[] constant(1)
+  c2 = s32[] constant(2)
+  start_index0 = s32[] select(pred0, c1, c2)
+  start_index1 = s32[] constant(0)
+  dynamic_slice = f32[5,2] dynamic-slice(dynamic_slice_input, start_index0, start_index1),
+                  dynamic_slice_sizes={5,2}
+  ROOT dot = f32[4,5] dot(dot_lhs, dynamic_slice),
+          lhs_contracting_dims={0}, rhs_contracting_dims={1}
+})";
+  TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
+                          GetOptimizedModule(kHloText));
+  EXPECT_THAT(
+      module->entry_computation()->root_instruction(),
+      GmockMatch(m::Fusion(m::Parameter(), m::Parameter(),
+                           m::Fusion(m::Parameter()), m::Constant())
+                     .WithFusionKind(HloInstruction::FusionKind::kCustom)));
+  // Check that it's not optimized away.
+  MatchHloModule(*module, "; CHECK: dynamic-slice(");
+  EXPECT_TRUE(RunAndCompare(kHloText, ErrorSpec{/*aabs=*/1e-4, /*arel=*/1e-6}));
+}
+
+TEST_F(TritonGemmTest, DynamicSliceIsSupportedInRhs) {
+  // The start index(es) for the non-majormost dimension(s) are constant zero(s)
+  // because we don't support dynamic slice on those dimensions.
+  constexpr absl::string_view kHloText = R"(
+HloModule m
+
+triton_gemm {
+  dynamic_slice_input = f32[7,2] parameter(0)
+  dot_rhs = f32[2,4] parameter(1)
+  start_index0 = s32[] parameter(2)
+  start_index1 = s32[] parameter(3)
+  dynamic_slice = f32[5,2] dynamic-slice(dynamic_slice_input, start_index0, start_index1),
+                  dynamic_slice_sizes={5,2}
+  ROOT dot = f32[5, 4] dot(dynamic_slice, dot_rhs),
+          lhs_contracting_dims={1}, rhs_contracting_dims={0}
+}
+
+ENTRY e {
+  dynamic_slice_input = f32[7,2] parameter(0)
+  dot_rhs = f32[2,4] parameter(1)
+  start_index0 = s32[] constant(1)
+  start_index1 = s32[] constant(0)
+  ROOT fusion = f32[5,4] fusion(dynamic_slice_input, dot_rhs, start_index0, start_index1),
+       kind=kCustom, calls=triton_gemm,
+       backend_config={
+         "fusion_backend_config":{
+           "kind":"__triton_gemm","triton_gemm_config":{
+             "block_m":"32","block_n":"32","block_k":"32","split_k":"1",
+             "num_stages":"1","num_warps":"4","num_ctas":"1"}}}
+})";
+
+  EXPECT_TRUE(RunAndCompareNoHloPasses(
+      kHloText, ErrorSpec{/*aabs=*/1e-4, /*arel=*/1e-6}));
+}
+
+class TritonGemmDynamicSliceClampingTest
+    : public TritonTest,
+      public ::testing::WithParamInterface<int> {};
+
+TEST_P(TritonGemmDynamicSliceClampingTest,
+       DynamicSliceIsSupportedWhenTheStartIndexNeedsClamping) {
+  // The start index(es) for the non-majormost dimension(s) are constant zero(s)
+  // because we don't support dynamic slice on those dimensions.
+
+  const std::string hlo_text = absl::Substitute(R"(
+HloModule m
+
+triton_gemm {
+  dynamic_slice_input = f32[7,2] parameter(0)
+  dot_rhs = f32[2,4] parameter(1)
+  start_index0 = s32[] parameter(2)
+  start_index1 = s32[] parameter(3)
+  dynamic_slice = f32[5,2] dynamic-slice(dynamic_slice_input, start_index0, start_index1),
+                  dynamic_slice_sizes={5,2}
+  ROOT dot = f32[5, 4] dot(dynamic_slice, dot_rhs),
+          lhs_contracting_dims={1}, rhs_contracting_dims={0}
+}
+
+ENTRY e {
+  dynamic_slice_input = f32[7,2] parameter(0)
+  dot_rhs = f32[2,4] parameter(1)
+  start_index0 = s32[] constant($0)
+  start_index1 = s32[] constant(0)
+  ROOT fusion = f32[5,4] fusion(dynamic_slice_input, dot_rhs, start_index0, start_index1),
+       kind=kCustom, calls=triton_gemm,
+       backend_config={
+         "fusion_backend_config":{
+           "kind":"__triton_gemm","triton_gemm_config":{
+             "block_m":"32","block_n":"32","block_k":"32","split_k":"1",
+             "num_stages":"1","num_warps":"4","num_ctas":"1"}}}
+})",
+                                                GetParam());
+
+  EXPECT_TRUE(RunAndCompareNoHloPasses(
+      hlo_text, ErrorSpec{/*aabs=*/1e-4, /*arel=*/1e-6}));
+}
+
+std::string OffsetParamToString(const ::testing::TestParamInfo<int>& data) {
+  return absl::StrCat("WithOffsetEq", data.param < 0 ? "Negative" : "",
+                      std::abs(data.param));
+}
+
+INSTANTIATE_TEST_SUITE_P(All, TritonGemmDynamicSliceClampingTest,
+                         ::testing::Values(-100, 3, 999), OffsetParamToString);
+
+TEST_F(TritonGemmTest, DynamicSliceOfMajormostContractingDimIsSupported) {
+  // Tests that dynamic-slice works on the majormost dimension even if that
+  // dimension is contracted.
+  // The start index(es) for the non-majormost dimension(s) are constant zero(s)
+  // because we don't support dynamic slice on those dimensions.
+  constexpr absl::string_view kHloText = R"(
+HloModule m
+
+triton_gemm {
+  dot_lhs = f32[2,4] parameter(0)
+  dynamic_slice_input = f32[5,4] parameter(1)
+  start_index0 = s32[] parameter(2)
+  start_index1 = s32[] parameter(3)
+  dynamic_slice = f32[2,4] dynamic-slice(dynamic_slice_input, start_index0, start_index1),
+                  dynamic_slice_sizes={2,4}
+  ROOT dot = f32[4,4] dot(dot_lhs, dynamic_slice),
+             lhs_contracting_dims={0}, rhs_contracting_dims={0}
+}
+
+ENTRY e {
+  dot_lhs = f32[2,4] parameter(0)
+  dynamic_slice_input = f32[5,4] parameter(1)
+  start_index0 = s32[] constant(2)
+  start_index1 = s32[] constant(0)
+  ROOT fusion = f32[4,4] fusion(dot_lhs, dynamic_slice_input, start_index0, start_index1),
+       kind=kCustom, calls=triton_gemm,
+       backend_config={
+         "fusion_backend_config":{
+           "kind":"__triton_gemm","triton_gemm_config":{
+             "block_m":"32","block_n":"32","block_k":"32","split_k":"1",
+             "num_stages":"1","num_warps":"4","num_ctas":"1"}}}
+})";
+
+  EXPECT_TRUE(RunAndCompareNoHloPasses(
+      kHloText, ErrorSpec{/*aabs=*/1e-4, /*arel=*/1e-6}));
+}
+
+TEST_F(TritonGemmTest, DynamicSliceOfMajormostBatchDimIsSupported) {
+  // Tests that dynamic-slice works on the majormost dimension even if that
+  // dimension is a batch.
+  // The start index(es) for the non-majormost dimension(s) are constant zero(s)
+  // because we don't support dynamic slice on those dimensions.
+  constexpr absl::string_view kHloText = R"(
+HloModule m
+
+triton_gemm {
+  dot_lhs = f32[2,2,4] parameter(0)
+  dynamic_slice_input = f32[7,2,4] parameter(1)
+  start_index0 = s32[] parameter(2)
+  start_index1 = s32[] parameter(3)
+  start_index2 = s32[] parameter(4)
+  dynamic_slice = f32[2,2,4] dynamic-slice(dynamic_slice_input, start_index0, start_index1, start_index2),
+                  dynamic_slice_sizes={2,2,4}
+  ROOT dot = f32[2,4,4] dot(dot_lhs, dynamic_slice),
+             lhs_batch_dims={0}, rhs_batch_dims={0}, lhs_contracting_dims={1}, rhs_contracting_dims={1}
+}
+
+ENTRY e {
+  dot_lhs = f32[2,2,4] parameter(0)
+  dynamic_slice_input = f32[7,2,4] parameter(1)
+  start_index0 = s32[] constant(2)
+  start_index1 = s32[] constant(0)
+  start_index2 = s32[] constant(0)
+  ROOT fusion = f32[2,4,4] fusion(dot_lhs, dynamic_slice_input, start_index0, start_index1, start_index2),
+       kind=kCustom, calls=triton_gemm,
+       backend_config={
+         "fusion_backend_config":{
+           "kind":"__triton_gemm","triton_gemm_config":{
+             "block_m":"32","block_n":"32","block_k":"32","split_k":"1",
+             "num_stages":"1","num_warps":"4","num_ctas":"1"}}}
+})";
+
+  EXPECT_TRUE(RunAndCompareNoHloPasses(
+      kHloText, ErrorSpec{/*aabs=*/1e-4, /*arel=*/1e-6}));
+}
+
+TEST_F(TritonGemmTest, DynamicSliceSingleDimensionIntoReshapeIsSupported) {
+  // This directly tests the targeted use case (b/307922364) of iterating over
+  // layer weights and extracting them with dynamic slice.
+  // The start index(es) for the non-majormost dimension(s) are constant zero(s)
+  // because we don't support dynamic slice on those dimensions.
+  constexpr absl::string_view kHloText = R"(
+HloModule m
+
+triton_gemm {
+  dot_lhs = f32[2,4] parameter(0)
+  dynamic_slice_input = f32[4,5,2] parameter(1)
+  start_index0 = s32[] parameter(2)
+  start_index1 = s32[] parameter(3)
+  start_index2 = s32[] parameter(4)
+  dynamic_slice = f32[1,5,2] dynamic-slice(dynamic_slice_input, start_index0, start_index1, start_index2),
+                             dynamic_slice_sizes={1,5,2}
+  reshape = f32[5,2] reshape(dynamic_slice)
+  ROOT d = f32[4,5] dot(dot_lhs, reshape),
+           lhs_contracting_dims={0}, rhs_contracting_dims={1}
+}
+
+ENTRY e {
+  dot_lhs = f32[2,4] parameter(0)
+  dynamic_slice_input = f32[4,5,2] parameter(1)
+  start_index0 = s32[] constant(3)
+  start_index1 = s32[] constant(0)
+  start_index2 = s32[] constant(0)
+  ROOT fusion = f32[4,5] fusion(dot_lhs, dynamic_slice_input, start_index0, start_index1, start_index2),
+       kind=kCustom, calls=triton_gemm,
+       backend_config={
+         "fusion_backend_config":{
+           "kind":"__triton_gemm","triton_gemm_config":{
+             "block_m":"32","block_n":"32","block_k":"32","split_k":"1",
+             "num_stages":"1","num_warps":"4","num_ctas":"1"}}}
+})";
+
+  EXPECT_TRUE(RunAndCompareNoHloPasses(
+      kHloText, ErrorSpec{/*aabs=*/1e-4, /*arel=*/1e-6}));
+}
+
 TEST_F(TritonGemmTestAny,
        DoNotFuseConcatenationOfSplitNonContractingDimension) {
   if (SkipBF16Tests()) {
diff --git a/third_party/xla/xla/service/gpu/triton_fusion_analysis_test.cc b/third_party/xla/xla/service/gpu/triton_fusion_analysis_test.cc
index cfa200d5511cd8..0ff9e93f6f9bb6 100644
--- a/third_party/xla/xla/service/gpu/triton_fusion_analysis_test.cc
+++ b/third_party/xla/xla/service/gpu/triton_fusion_analysis_test.cc
@@ -674,6 +674,62 @@ ENTRY e {
   (void)analysis;
 }
 
+TEST_F(TritonDotAnalysisTest, DynamicSliceIsSupported) {
+  TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr<VerifiedHloModule> module,
+                          ParseAndReturnVerifiedModule(R"(
+triton_gemm {
+  dot_lhs = f32[2,18]{1,0} parameter(0)
+  dynamic_slice_input = f32[96,2]{1,0} parameter(1)
+  start_index0 = s32[] parameter(2)
+  start_index1 = s32[] parameter(3)
+  dynamic_slice = f32[64,2]{1,0} dynamic-slice(dynamic_slice_input,
+                                               start_index0, start_index1),
+                  dynamic_slice_sizes={64,2}
+  ROOT dot = f32[18,64]{1,0} dot(dot_lhs, dynamic_slice),
+             lhs_contracting_dims={0}, rhs_contracting_dims={1}
+}
+
+ENTRY e {
+  dot_lhs = f32[2,18]{1,0} parameter(0)
+  dynamic_slice_input = f32[96,2]{1,0} parameter(1)
+  start_index0 = s32[] parameter(2)
+  start_index1 = s32[] parameter(3)
+  ROOT triton_gemm_d = f32[18,64]{1,0} fusion(dot_lhs, dynamic_slice_input,
+                                              start_index0, start_index1),
+      kind=kCustom,
+      calls=triton_gemm,
+      backend_config={"kind":"__triton_gemm"}
+}
+)"));
+
+  const HloComputation* dot_computation =
+      module->entry_computation()->root_instruction()->called_computations()[0];
+  TF_ASSERT_OK_AND_ASSIGN(const auto analysis,
+                          TritonFusionAnalysis::Execute(*dot_computation));
+  const HloInstruction* p0 = dot_computation->parameter_instruction(0);
+  const HloInstruction* p1 = dot_computation->parameter_instruction(1);
+  EXPECT_EQ(*analysis.ScopeParameters(TritonFusionAnalysis::Scope::LHS).begin(),
+            p0);
+  EXPECT_EQ(*analysis.ScopeParameters(TritonFusionAnalysis::Scope::RHS).begin(),
+            p1);
+  EXPECT_THAT(*analysis.IterSpec(TritonFusionAnalysis::Scope::LHS, p0, 0),
+              ElementsAre(FieldsAre(/*stride=*/18, /*count=*/2,
+                                    /*slice_start=*/0, /*sliced_count=*/2,
+                                    /*subfragments=*/ElementsAre(2))));
+  EXPECT_THAT(*analysis.IterSpec(TritonFusionAnalysis::Scope::LHS, p0, 1),
+              ElementsAre(FieldsAre(/*stride=*/1, /*count=*/18,
+                                    /*slice_start=*/0, /*sliced_count=*/18,
+                                    /*subfragments=*/ElementsAre(18))));
+  EXPECT_THAT(*analysis.IterSpec(TritonFusionAnalysis::Scope::RHS, p1, 0),
+              ElementsAre(FieldsAre(/*stride=*/2, /*count=*/96,
+                                    /*slice_start=*/0, /*sliced_count=*/96,
+                                    /*subfragments=*/ElementsAre(96))));
+  EXPECT_THAT(*analysis.IterSpec(TritonFusionAnalysis::Scope::RHS, p1, 1),
+              ElementsAre(FieldsAre(/*stride=*/1, /*count=*/2,
+                                    /*slice_start=*/0, /*sliced_count=*/2,
+                                    /*subfragments=*/ElementsAre(2))));
+}
+
 TEST_F(TritonDotAnalysisTest, SparseDot) {
   TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr<VerifiedHloModule> module,
                           ParseAndReturnVerifiedModule(R"(
diff --git a/third_party/xla/xla/service/gpu/triton_tiling_propagation.cc b/third_party/xla/xla/service/gpu/triton_tiling_propagation.cc
index aa3bf579377151..58b9492fda2bd3 100644
--- a/third_party/xla/xla/service/gpu/triton_tiling_propagation.cc
+++ b/third_party/xla/xla/service/gpu/triton_tiling_propagation.cc
@@ -39,6 +39,7 @@ limitations under the License.
 #include "xla/hlo/ir/hlo_instructions.h"
 #include "xla/hlo/ir/hlo_opcode.h"
 #include "xla/hlo/utils/hlo_query.h"
+#include "xla/layout.h"
 #include "xla/permutation_util.h"
 #include "xla/service/gpu/triton_support.h"
 #include "xla/service/instruction_fusion.h"
@@ -840,6 +841,29 @@ DimOrderMapOrError GetPropagatedDimOrdersForDimAlteringOp(
               fragment->sliced_count());
         }
       }
+    } else if (hlo.opcode() == HloOpcode::kDynamicSlice) {
+      // All operands after idx 0 are scalar indices. As such, we do not want
+      // to explicitly define dim orders.
+      if (dst != &hlo && hlo.operand_index(dst) >= 1) {
+        continue;
+      }
+      const auto dynamic_slice = Cast<HloDynamicSliceInstruction>(&hlo);
+      dst_logical.resize(src_logical.size());
+      for (int dim = 0; dim < src_logical.size(); ++dim) {
+        dst_logical[dim] = src_logical[dim];
+        if (dynamic_slice->slice_sizes(dim) != dst->shape().dimensions(dim)) {
+          if (dst_logical[dim].size() > 1) {
+            return FusionDecision("Slicing of fragmented dimension.");
+          }
+          auto fragment = dst_logical[dim].front();
+          fragment->set_count(dst->shape().dimensions(dim));
+
+          // As we do not know which section of the tensor we keep, we retain
+          // the whole part.
+          fragment->set_slice(fragment->slice_start(),
+                              dst->shape().dimensions(dim));
+        }
+      }
     } else {
       return FusionDecision("Function called on a wrong instruction.");
     }
@@ -889,9 +913,10 @@ DimOrderMapOrError GetPropagatedDimOrders(const HloInstruction& hlo,
   if (hlo.opcode() != HloOpcode::kParameter &&
       direction == TransformDirection::kOutputToInput &&
       absl::c_any_of(hlo.users(), [](const HloInstruction* user) {
-        return user->opcode() == HloOpcode::kConcatenate;
+        return (user->opcode() == HloOpcode::kConcatenate ||
+                user->opcode() == HloOpcode::kDynamicSlice);
       })) {
-    return "No fusion into concatenations";
+    return "No fusion into concatenations or dynamic slice.";
   }
   if (hlo.opcode() == HloOpcode::kParameter ||
       hlo_query::IsScalarConstant(&hlo)) {
@@ -941,6 +966,38 @@ DimOrderMapOrError GetPropagatedDimOrders(const HloInstruction& hlo,
     if (direction != TransformDirection::kOutputToInput) {
       return "Unsupported slice direction.";
     }
+
+    return GetPropagatedDimOrdersForDimAlteringOp(hlo, direction, src_dim_order,
+                                                  properties);
+  } else if (hlo.opcode() == HloOpcode::kDynamicSlice &&
+             direction == TransformDirection::kOutputToInput) {
+    // We handle the dynamic slice within EmitTensorPointer, which is only
+    // used for GEMM fusions.
+    if (!std::holds_alternative<DotProperties>(properties)) {
+      return "Dynamic slices for now are only supported in GEMM fusions.";
+    }
+
+    // Similar to normal slice, we cannot slice a non-major-most dimension as
+    // that would introduce non-contiguous strides under tiling. The existing
+    // check against this in GetRequirementsIfSupportedOrder is not suitable for
+    // dynamic slices, so we instead check for this here.
+    const HloInstruction* input = hlo.operand(0);
+    Layout in_layout = input->shape().layout();
+    int64_t majormost =
+        in_layout.minor_to_major(in_layout.minor_to_major_size() - 1);
+    const HloDynamicSliceInstruction* dynamic_slice =
+        Cast<HloDynamicSliceInstruction>(&hlo);
+
+    for (int i = 0; i < input->shape().dimensions_size(); ++i) {
+      if (i == majormost) {
+        continue;
+      } else if (input->shape().dimensions(i) !=
+                 dynamic_slice->slice_sizes(i)) {
+        return FusionDecision(
+            "Unsupported dynamic slice on non-major-most dimension.");
+      }
+    }
+
     return GetPropagatedDimOrdersForDimAlteringOp(hlo, direction, src_dim_order,
                                                   properties);
   } else if (hlo.opcode() == HloOpcode::kReshape) {

From 6c1a8e4dcc8b6a05cfbb3c47bb85ef5f9e8d8800 Mon Sep 17 00:00:00 2001
From: Adrian Kuegel <akuegel@google.com>
Date: Fri, 17 May 2024 03:12:45 -0700
Subject: [PATCH 343/478] Followup to Scatter Layout normalization

The code already implicitly assumes that ScatterSimplifier has run before. We
cannot normalize if there are more than 1 "scatter" (batch) dimensions. Add a
check for that and adjust a test case that would in fact be incorrectly
normalized (but hidden from the verifier by running ScatterSimplifier
afterwards).

PiperOrigin-RevId: 634695331
---
 .../xla/xla/service/layout_normalization.cc     | 17 ++++++++++++++++-
 .../xla/service/layout_normalization_test.cc    |  8 ++++----
 2 files changed, 20 insertions(+), 5 deletions(-)

diff --git a/third_party/xla/xla/service/layout_normalization.cc b/third_party/xla/xla/service/layout_normalization.cc
index 6ff486dd0aedc9..e5cc2ffdf8ebab 100644
--- a/third_party/xla/xla/service/layout_normalization.cc
+++ b/third_party/xla/xla/service/layout_normalization.cc
@@ -404,6 +404,22 @@ class LayoutNormalizationVisitor : public DfsHloRewriteVisitor {
       TF_ASSIGN_OR_RETURN(auto normalized_update, GetNormalizedInput(operand));
       normalized_updates.push_back(normalized_update);
     }
+
+    // Since normalization might reorder the 'scatter_updates' operands
+    // differently than the 'scatter_indices' update, we have no way to specify
+    // the order of 'scatter' (batch) dimensions, as that is not an attribute in
+    // ScatterDimensionNumbers. Scatter implicitly assumes that the 'scatter'
+    // dimensions appear in the same order in 'scatter_updates' and
+    // 'scatter_indices'. So we require that there is just a single
+    // 'scatter' dimension. This is ensured by the ScatterSimplifier pass.
+    const auto& dims = scatter->scatter_dimension_numbers();
+    if (scatter->scatter_updates().front()->shape().rank() -
+            dims.update_window_dims_size() >
+        1) {
+      return FailedPrecondition(
+          "There should be just a single scatter dimension. Make sure to run "
+          "ScatterSimplifier before LayoutNormalization");
+    }
     TF_ASSIGN_OR_RETURN(auto normalized_indices,
                         GetNormalizedInput(scatter->scatter_indices()));
 
@@ -423,7 +439,6 @@ class LayoutNormalizationVisitor : public DfsHloRewriteVisitor {
     // operand dimensions. scatter dimension i corresponds to
     // scatter_dims_to_operand_dims[i] operand dimension.
 
-    const auto& dims = scatter->scatter_dimension_numbers();
     ScatterDimensionNumbers normalized_dims;
     normalized_dims.set_index_vector_dim(
         indices_permutation[dims.index_vector_dim()]);
diff --git a/third_party/xla/xla/service/layout_normalization_test.cc b/third_party/xla/xla/service/layout_normalization_test.cc
index 1b74f8bb44c325..d2b9d92d2fb934 100644
--- a/third_party/xla/xla/service/layout_normalization_test.cc
+++ b/third_party/xla/xla/service/layout_normalization_test.cc
@@ -788,16 +788,16 @@ region_0.10 {
 
 ENTRY main.17 {
   p0 = s16[3,2,2,14,16]{0,1,4,3,2} parameter(0)
-  p1 = s32[6,11,8,2]{0,2,1,3} parameter(1)
-  p2 = s16[6,11,8,3,5]{0,1,4,3,2} parameter(2)
-  ROOT scatter = s16[3,2,2,14,16]{0,1,4,3,2} scatter(p0, p1, p2), update_window_dims={3,4}, inserted_window_dims={1,2,3}, scatter_dims_to_operand_dims={4,0}, index_vector_dim=3, to_apply=region_0.10
+  p1 = s32[2,11]{0,1} parameter(1)
+  p2 = s16[11,3,5]{2,0,1} parameter(2)
+  ROOT scatter = s16[3,2,2,14,16]{0,1,4,3,2} scatter(p0, p1, p2), update_window_dims={1,2}, inserted_window_dims={1,2,3}, scatter_dims_to_operand_dims={4,0}, index_vector_dim=0, to_apply=region_0.10
 }
 )";
 
   CheckLayoutNormalization(
       hlo, R"(
 // CHECK: scatter({{.*}}),
-// CHECK-SAME: update_window_dims={2,1}, inserted_window_dims={0,1,3}, scatter_dims_to_operand_dims={2,4}, index_vector_dim=0, to_apply=%region_0.10
+// CHECK-SAME: update_window_dims={2,0}, inserted_window_dims={0,1,3}, scatter_dims_to_operand_dims={2,4}, index_vector_dim=1, to_apply=%region_0.10
 )",
       // Run the ScatterSimplifier afterwards, otherwise the verifier will
       // complain!

From abb166cd38ad303b6fd6c5521440a654824837fa Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Fri, 17 May 2024 04:04:13 -0700
Subject: [PATCH 344/478] Automated Code Change

PiperOrigin-RevId: 634707918
---
 .../lite/kernels/parse_example/example_proto_fast_parsing.h | 2 +-
 tensorflow/lite/kernels/parse_example/parse_example.cc      | 6 +++---
 2 files changed, 4 insertions(+), 4 deletions(-)

diff --git a/tensorflow/lite/kernels/parse_example/example_proto_fast_parsing.h b/tensorflow/lite/kernels/parse_example/example_proto_fast_parsing.h
index 971ddad7f95feb..f25e10f235d0e7 100644
--- a/tensorflow/lite/kernels/parse_example/example_proto_fast_parsing.h
+++ b/tensorflow/lite/kernels/parse_example/example_proto_fast_parsing.h
@@ -450,7 +450,7 @@ struct FeatureProtos {
 // Map from feature name to FeatureProtos for that feature.
 using FeatureProtosMap = absl::flat_hash_map<StringPiece, FeatureProtos>;
 
-string ExampleName(const gtl::ArraySlice<tstring> example_names, int n);
+string ExampleName(const absl::Span<const tstring> example_names, int n);
 
 // Return the number of bytes elements parsed, or -1 on error. If out is null,
 // this method simply counts the number of elements without any copying.
diff --git a/tensorflow/lite/kernels/parse_example/parse_example.cc b/tensorflow/lite/kernels/parse_example/parse_example.cc
index 86ac20625c7cde..0d3e8657c2a7fe 100644
--- a/tensorflow/lite/kernels/parse_example/parse_example.cc
+++ b/tensorflow/lite/kernels/parse_example/parse_example.cc
@@ -430,7 +430,7 @@ void CopySparseBufferToTensor(tf::DataType dtype, size_t offset,
   }
 }
 
-inline void CopyToBuffer(tf::gtl::ArraySlice<tstring> vec, char* tensor_buffer,
+inline void CopyToBuffer(absl::Span<const tstring> vec, char* tensor_buffer,
                          int num_examples, int batch_size,
                          int elements_per_stride) {
   int i = 0, k = 0;
@@ -454,7 +454,7 @@ inline void CopyToBuffer(tf::gtl::ArraySlice<tstring> vec, char* tensor_buffer,
 
 Status FastParseExampleLite(
     const FastParseExampleConfig& config, const TfLiteTensor* serialized,
-    tf::gtl::ArraySlice<tstring> example_names, bool* quick_filter,
+    absl::Span<const tstring> example_names, bool* quick_filter,
     int quick_filter_size, const std::unique_ptr<ConfigIndex>& config_index,
     int config_index_size, SeededHasher* hasher, TfLiteResult* result,
     std::map<absl::string_view, int>& stats, TfLiteContext* context) {
@@ -633,7 +633,7 @@ Status FastParseExampleLite(
         memcpy(tensor_buffer + sizeof(int32_t) * (i + 1), &offset_i,
                sizeof(int32_t));
       }
-      tf::gtl::ArraySlice<tstring> slice(vec.data(), vec.size());
+      absl::Span<const tstring> slice(vec.data(), vec.size());
       CopyToBuffer(slice, tensor_buffer + start, count, batch_size,
                    elements_per_stride);
     }

From cd5779ddd0a3a07b580b99952bc1f4bfc643b46d Mon Sep 17 00:00:00 2001
From: Harsha H S <hsharsha@users.noreply.github.com>
Date: Fri, 17 May 2024 04:30:35 -0700
Subject: [PATCH 345/478] PR #12515: [ROCm] Fix build break introduced in
 ffa7bb5 and df736d7

Imported from GitHub PR https://github.com/openxla/xla/pull/12515

Fix build break introduced in ffa7bb5 and df736d7
Copybara import of the project:

--
499687ea87e38be681cb978d15ebd314377623c4 by Harsha HS <harsha.havanurshamsundara@amd.com>:

Fix build break introduced in ffa7bb5 and df736d7

Merging this change closes #12515

PiperOrigin-RevId: 634713247
---
 .../xla/xla/service/gpu/ir_emitter_triton_rocm.cc        | 9 ++++++---
 1 file changed, 6 insertions(+), 3 deletions(-)

diff --git a/third_party/xla/xla/service/gpu/ir_emitter_triton_rocm.cc b/third_party/xla/xla/service/gpu/ir_emitter_triton_rocm.cc
index 1a7aa92c62a9eb..e4331885baa08c 100644
--- a/third_party/xla/xla/service/gpu/ir_emitter_triton_rocm.cc
+++ b/third_party/xla/xla/service/gpu/ir_emitter_triton_rocm.cc
@@ -69,19 +69,22 @@ absl::Status CreateTritonPipeline(
   // Based on make_ttgir() in
   // @triton//:third_party/nvidia/backend/compiler.py
   pm.addPass(mt::createConvertTritonToTritonGPUPass(
-      config.num_warps, threadsPerWarp, config.num_ctas, ccAsInt));
+      absl::StrFormat("cuda:%u", ccAsInt), config.num_warps, threadsPerWarp,
+      config.num_ctas));
   pm.addPass(mt::gpu::createCoalescePass());
   pm.addPass(mt::gpu::createRemoveLayoutConversionsPass());
   pm.addPass(mt::gpu::createOptimizeThreadLocalityPass());
   pm.addPass(mt::gpu::createAccelerateMatmulPass(ccAsInt));
   pm.addPass(mt::gpu::createRemoveLayoutConversionsPass());
-  pm.addPass(mt::gpu::createOptimizeDotOperandsPass());
+  // TODO ROCm Check if we want to compare MI100 and greater
+  pm.addPass(mt::gpu::createOptimizeDotOperandsPass(true));
   pm.addPass(mlir::createCSEPass());
   pm.addPass(mt::gpu::createPipelinePass(config.num_stages, config.num_warps,
                                          config.num_ctas, ccAsInt));
   pm.addPass(mt::gpu::createPrefetchPass());
 
-  pm.addPass(mt::gpu::createOptimizeDotOperandsPass());
+  // TODO ROCm Check if we want to compare MI100 and greater
+  pm.addPass(mt::gpu::createOptimizeDotOperandsPass(true));
   pm.addPass(mt::gpu::createRemoveLayoutConversionsPass());
   pm.addPass(mt::gpu::createReduceDataDuplicationPass());
   pm.addPass(mt::gpu::createReorderInstructionsPass());

From ce8a3bd8926738d7d8160e8ee3019baa04e9c722 Mon Sep 17 00:00:00 2001
From: Henning Becker <hebecker@google.com>
Date: Fri, 17 May 2024 05:54:46 -0700
Subject: [PATCH 346/478] Fix dependencies in ROCM Triton emitter

PiperOrigin-RevId: 634730178
---
 third_party/xla/xla/service/gpu/BUILD | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/third_party/xla/xla/service/gpu/BUILD b/third_party/xla/xla/service/gpu/BUILD
index a28d711aa25f39..275499abcb330e 100644
--- a/third_party/xla/xla/service/gpu/BUILD
+++ b/third_party/xla/xla/service/gpu/BUILD
@@ -593,11 +593,13 @@ cc_library(
         "@triton//:TritonNvidiaGPUTransforms",
         "@triton//:TritonGPUToLLVM",
         "@triton//:TritonToTritonGPU",
+        "@triton//:TritonGPUTransforms",
     ]) + if_cuda_is_configured([
         "@triton//third_party/nvidia:NVGPUToLLVM",
         "@triton//third_party/nvidia:TritonNVIDIAGPUToLLVM",
-        "@triton//:TritonGPUTransforms",
         "@triton//:TritonLLVMIR",
+    ]) + if_rocm_is_configured([
+        "@local_tsl//tsl/platform:rocm_rocdl_path",
     ]),
 )
 

From a316efcb8e2eb54030f04f8c8aea8922fa9a50f2 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Fri, 17 May 2024 05:58:27 -0700
Subject: [PATCH 347/478] Automated Code Change

PiperOrigin-RevId: 634731291
---
 .../xla/third_party/tsl/tsl/platform/env.cc   | 169 ++++-----
 .../xla/third_party/tsl/tsl/platform/env.h    | 276 ++++++++-------
 .../third_party/tsl/tsl/platform/errors.cc    |  37 +-
 .../xla/third_party/tsl/tsl/platform/errors.h | 230 ++++++------
 .../tsl/tsl/platform/errors_test.cc           |  32 +-
 .../tsl/tsl/platform/file_system.cc           |  52 +--
 .../tsl/tsl/platform/file_system.h            | 330 +++++++++---------
 .../tsl/tsl/platform/file_system_helper.cc    |  27 +-
 .../tsl/tsl/platform/file_system_helper.h     |   6 +-
 .../tsl/tsl/platform/null_file_system.h       |  45 +--
 .../tsl/tsl/platform/ram_file_system.h        | 107 +++---
 .../tsl/tsl/platform/retrying_file_system.h   |  89 ++---
 .../tsl/platform/retrying_file_system_test.cc | 144 ++++----
 .../tsl/tsl/platform/retrying_utils.cc        |  19 +-
 .../tsl/tsl/platform/retrying_utils.h         |  16 +-
 .../tsl/tsl/platform/retrying_utils_test.cc   |  31 +-
 .../third_party/tsl/tsl/platform/status.cc    |  35 +-
 .../xla/third_party/tsl/tsl/platform/status.h |  29 +-
 .../tsl/tsl/platform/status_matchers.cc       |   2 +-
 .../tsl/tsl/platform/status_matchers.h        |   8 +-
 .../tsl/tsl/platform/status_matchers_test.cc  |  75 ++--
 .../tsl/tsl/platform/status_test.cc           |  40 +--
 .../tsl/tsl/platform/status_to_from_proto.cc  |  12 +-
 .../tsl/tsl/platform/status_to_from_proto.h   |   4 +-
 .../tsl/tsl/platform/statusor_test.cc         | 202 ++++++-----
 25 files changed, 1044 insertions(+), 973 deletions(-)

diff --git a/third_party/xla/third_party/tsl/tsl/platform/env.cc b/third_party/xla/third_party/tsl/tsl/platform/env.cc
index 2649aa00a5d33c..45199b50ebd94b 100644
--- a/third_party/xla/third_party/tsl/tsl/platform/env.cc
+++ b/third_party/xla/third_party/tsl/tsl/platform/env.cc
@@ -55,11 +55,11 @@ constexpr size_t kCopyFileBufferSize = 128 * 1024;
 
 class FileSystemRegistryImpl : public FileSystemRegistry {
  public:
-  Status Register(const std::string& scheme, Factory factory) override;
-  Status Register(const std::string& scheme,
-                  std::unique_ptr<FileSystem> filesystem) override;
+  absl::Status Register(const std::string& scheme, Factory factory) override;
+  absl::Status Register(const std::string& scheme,
+                        std::unique_ptr<FileSystem> filesystem) override;
   FileSystem* Lookup(const std::string& scheme) override;
-  Status GetRegisteredFileSystemSchemes(
+  absl::Status GetRegisteredFileSystemSchemes(
       std::vector<std::string>* schemes) override;
 
  private:
@@ -68,25 +68,25 @@ class FileSystemRegistryImpl : public FileSystemRegistry {
       TF_GUARDED_BY(mu_);
 };
 
-Status FileSystemRegistryImpl::Register(const std::string& scheme,
-                                        FileSystemRegistry::Factory factory) {
+absl::Status FileSystemRegistryImpl::Register(
+    const std::string& scheme, FileSystemRegistry::Factory factory) {
   mutex_lock lock(mu_);
   if (!registry_.emplace(scheme, std::unique_ptr<FileSystem>(factory()))
            .second) {
     return errors::AlreadyExists("File factory for ", scheme,
                                  " already registered");
   }
-  return OkStatus();
+  return absl::OkStatus();
 }
 
-Status FileSystemRegistryImpl::Register(
+absl::Status FileSystemRegistryImpl::Register(
     const std::string& scheme, std::unique_ptr<FileSystem> filesystem) {
   mutex_lock lock(mu_);
   if (!registry_.emplace(scheme, std::move(filesystem)).second) {
     return errors::AlreadyExists("File system for ", scheme,
                                  " already registered");
   }
-  return OkStatus();
+  return absl::OkStatus();
 }
 
 FileSystem* FileSystemRegistryImpl::Lookup(const std::string& scheme) {
@@ -98,19 +98,19 @@ FileSystem* FileSystemRegistryImpl::Lookup(const std::string& scheme) {
   return found->second.get();
 }
 
-Status FileSystemRegistryImpl::GetRegisteredFileSystemSchemes(
+absl::Status FileSystemRegistryImpl::GetRegisteredFileSystemSchemes(
     std::vector<std::string>* schemes) {
   mutex_lock lock(mu_);
   for (const auto& e : registry_) {
     schemes->push_back(e.first);
   }
-  return OkStatus();
+  return absl::OkStatus();
 }
 
 Env::Env() : file_system_registry_(new FileSystemRegistryImpl) {}
 
-Status Env::GetFileSystemForFile(const std::string& fname,
-                                 FileSystem** result) {
+absl::Status Env::GetFileSystemForFile(const std::string& fname,
+                                       FileSystem** result) {
   StringPiece scheme, host, path;
   io::ParseURI(fname, &scheme, &host, &path);
   FileSystem* file_system = file_system_registry_->Lookup(std::string(scheme));
@@ -123,25 +123,26 @@ Status Env::GetFileSystemForFile(const std::string& fname,
                                  "' not implemented (file: '", fname, "')");
   }
   *result = file_system;
-  return OkStatus();
+  return absl::OkStatus();
 }
 
-Status Env::GetRegisteredFileSystemSchemes(std::vector<std::string>* schemes) {
+absl::Status Env::GetRegisteredFileSystemSchemes(
+    std::vector<std::string>* schemes) {
   return file_system_registry_->GetRegisteredFileSystemSchemes(schemes);
 }
 
-Status Env::RegisterFileSystem(const std::string& scheme,
-                               FileSystemRegistry::Factory factory) {
+absl::Status Env::RegisterFileSystem(const std::string& scheme,
+                                     FileSystemRegistry::Factory factory) {
   return file_system_registry_->Register(scheme, std::move(factory));
 }
 
-Status Env::RegisterFileSystem(const std::string& scheme,
-                               std::unique_ptr<FileSystem> filesystem) {
+absl::Status Env::RegisterFileSystem(const std::string& scheme,
+                                     std::unique_ptr<FileSystem> filesystem) {
   return file_system_registry_->Register(scheme, std::move(filesystem));
 }
 
-Status Env::SetOption(const std::string& scheme, const std::string& key,
-                      const std::string& value) {
+absl::Status Env::SetOption(const std::string& scheme, const std::string& key,
+                            const std::string& value) {
   FileSystem* file_system = file_system_registry_->Lookup(scheme);
   if (!file_system) {
     return errors::Unimplemented("File system scheme '", scheme,
@@ -150,8 +151,8 @@ Status Env::SetOption(const std::string& scheme, const std::string& key,
   return file_system->SetOption(key, value);
 }
 
-Status Env::SetOption(const std::string& scheme, const std::string& key,
-                      const std::vector<string>& values) {
+absl::Status Env::SetOption(const std::string& scheme, const std::string& key,
+                            const std::vector<string>& values) {
   FileSystem* file_system = file_system_registry_->Lookup(scheme);
   if (!file_system) {
     return errors::Unimplemented("File system scheme '", scheme,
@@ -160,8 +161,8 @@ Status Env::SetOption(const std::string& scheme, const std::string& key,
   return file_system->SetOption(key, values);
 }
 
-Status Env::SetOption(const std::string& scheme, const std::string& key,
-                      const std::vector<int64_t>& values) {
+absl::Status Env::SetOption(const std::string& scheme, const std::string& key,
+                            const std::vector<int64_t>& values) {
   FileSystem* file_system = file_system_registry_->Lookup(scheme);
   if (!file_system) {
     return errors::Unimplemented("File system scheme '", scheme,
@@ -170,8 +171,8 @@ Status Env::SetOption(const std::string& scheme, const std::string& key,
   return file_system->SetOption(key, values);
 }
 
-Status Env::SetOption(const std::string& scheme, const std::string& key,
-                      const std::vector<double>& values) {
+absl::Status Env::SetOption(const std::string& scheme, const std::string& key,
+                            const std::vector<double>& values) {
   FileSystem* file_system = file_system_registry_->Lookup(scheme);
   if (!file_system) {
     return errors::Unimplemented("File system scheme '", scheme,
@@ -180,7 +181,7 @@ Status Env::SetOption(const std::string& scheme, const std::string& key,
   return file_system->SetOption(key, values);
 }
 
-Status Env::FlushFileSystemCaches() {
+absl::Status Env::FlushFileSystemCaches() {
   std::vector<string> schemes;
   TF_RETURN_IF_ERROR(GetRegisteredFileSystemSchemes(&schemes));
   for (const string& scheme : schemes) {
@@ -189,45 +190,45 @@ Status Env::FlushFileSystemCaches() {
         GetFileSystemForFile(io::CreateURI(scheme, "", ""), &fs));
     fs->FlushCaches();
   }
-  return OkStatus();
+  return absl::OkStatus();
 }
 
-Status Env::NewRandomAccessFile(const string& fname,
-                                std::unique_ptr<RandomAccessFile>* result) {
+absl::Status Env::NewRandomAccessFile(
+    const string& fname, std::unique_ptr<RandomAccessFile>* result) {
   FileSystem* fs;
   TF_RETURN_IF_ERROR(GetFileSystemForFile(fname, &fs));
   return fs->NewRandomAccessFile(fname, result);
 }
 
-Status Env::NewReadOnlyMemoryRegionFromFile(
+absl::Status Env::NewReadOnlyMemoryRegionFromFile(
     const string& fname, std::unique_ptr<ReadOnlyMemoryRegion>* result) {
   FileSystem* fs;
   TF_RETURN_IF_ERROR(GetFileSystemForFile(fname, &fs));
   return fs->NewReadOnlyMemoryRegionFromFile(fname, result);
 }
 
-Status Env::NewWritableFile(const string& fname,
-                            std::unique_ptr<WritableFile>* result) {
+absl::Status Env::NewWritableFile(const string& fname,
+                                  std::unique_ptr<WritableFile>* result) {
   FileSystem* fs;
   TF_RETURN_IF_ERROR(GetFileSystemForFile(fname, &fs));
   return fs->NewWritableFile(fname, result);
 }
 
-Status Env::NewAppendableFile(const string& fname,
-                              std::unique_ptr<WritableFile>* result) {
+absl::Status Env::NewAppendableFile(const string& fname,
+                                    std::unique_ptr<WritableFile>* result) {
   FileSystem* fs;
   TF_RETURN_IF_ERROR(GetFileSystemForFile(fname, &fs));
   return fs->NewAppendableFile(fname, result);
 }
 
-Status Env::FileExists(const string& fname) {
+absl::Status Env::FileExists(const string& fname) {
   FileSystem* fs;
   TF_RETURN_IF_ERROR(GetFileSystemForFile(fname, &fs));
   return fs->FileExists(fname);
 }
 
 bool Env::FilesExist(const std::vector<string>& files,
-                     std::vector<Status>* status) {
+                     std::vector<absl::Status>* status) {
   std::unordered_map<string, std::vector<string>> files_per_fs;
   for (const auto& file : files) {
     StringPiece scheme, host, path;
@@ -235,18 +236,18 @@ bool Env::FilesExist(const std::vector<string>& files,
     files_per_fs[string(scheme)].push_back(file);
   }
 
-  std::unordered_map<string, Status> per_file_status;
+  std::unordered_map<string, absl::Status> per_file_status;
   bool result = true;
   for (auto itr : files_per_fs) {
     FileSystem* file_system = file_system_registry_->Lookup(itr.first);
     bool fs_result;
-    std::vector<Status> local_status;
-    std::vector<Status>* fs_status = status ? &local_status : nullptr;
+    std::vector<absl::Status> local_status;
+    std::vector<absl::Status>* fs_status = status ? &local_status : nullptr;
     if (!file_system) {
       fs_result = false;
       if (fs_status) {
-        Status s = errors::Unimplemented("File system scheme '", itr.first,
-                                         "' not implemented");
+        absl::Status s = errors::Unimplemented("File system scheme '",
+                                               itr.first, "' not implemented");
         local_status.resize(itr.second.size(), s);
       }
     } else {
@@ -272,75 +273,76 @@ bool Env::FilesExist(const std::vector<string>& files,
   return result;
 }
 
-Status Env::GetChildren(const string& dir, std::vector<string>* result) {
+absl::Status Env::GetChildren(const string& dir, std::vector<string>* result) {
   FileSystem* fs;
   TF_RETURN_IF_ERROR(GetFileSystemForFile(dir, &fs));
   return fs->GetChildren(dir, result);
 }
 
-Status Env::GetMatchingPaths(const string& pattern,
-                             std::vector<string>* results) {
+absl::Status Env::GetMatchingPaths(const string& pattern,
+                                   std::vector<string>* results) {
   FileSystem* fs;
   TF_RETURN_IF_ERROR(GetFileSystemForFile(pattern, &fs));
   return fs->GetMatchingPaths(pattern, results);
 }
 
-Status Env::DeleteFile(const string& fname) {
+absl::Status Env::DeleteFile(const string& fname) {
   FileSystem* fs;
   TF_RETURN_IF_ERROR(GetFileSystemForFile(fname, &fs));
   return fs->DeleteFile(fname);
 }
 
-Status Env::RecursivelyCreateDir(const string& dirname) {
+absl::Status Env::RecursivelyCreateDir(const string& dirname) {
   FileSystem* fs;
   TF_RETURN_IF_ERROR(GetFileSystemForFile(dirname, &fs));
   return fs->RecursivelyCreateDir(dirname);
 }
 
-Status Env::CreateDir(const string& dirname) {
+absl::Status Env::CreateDir(const string& dirname) {
   FileSystem* fs;
   TF_RETURN_IF_ERROR(GetFileSystemForFile(dirname, &fs));
   return fs->CreateDir(dirname);
 }
 
-Status Env::DeleteDir(const string& dirname) {
+absl::Status Env::DeleteDir(const string& dirname) {
   FileSystem* fs;
   TF_RETURN_IF_ERROR(GetFileSystemForFile(dirname, &fs));
   return fs->DeleteDir(dirname);
 }
 
-Status Env::Stat(const string& fname, FileStatistics* stat) {
+absl::Status Env::Stat(const string& fname, FileStatistics* stat) {
   FileSystem* fs;
   TF_RETURN_IF_ERROR(GetFileSystemForFile(fname, &fs));
   return fs->Stat(fname, stat);
 }
 
-Status Env::IsDirectory(const string& fname) {
+absl::Status Env::IsDirectory(const string& fname) {
   FileSystem* fs;
   TF_RETURN_IF_ERROR(GetFileSystemForFile(fname, &fs));
   return fs->IsDirectory(fname);
 }
 
-Status Env::HasAtomicMove(const string& path, bool* has_atomic_move) {
+absl::Status Env::HasAtomicMove(const string& path, bool* has_atomic_move) {
   FileSystem* fs;
   TF_RETURN_IF_ERROR(GetFileSystemForFile(path, &fs));
   return fs->HasAtomicMove(path, has_atomic_move);
 }
 
-Status Env::DeleteRecursively(const string& dirname, int64_t* undeleted_files,
-                              int64_t* undeleted_dirs) {
+absl::Status Env::DeleteRecursively(const string& dirname,
+                                    int64_t* undeleted_files,
+                                    int64_t* undeleted_dirs) {
   FileSystem* fs;
   TF_RETURN_IF_ERROR(GetFileSystemForFile(dirname, &fs));
   return fs->DeleteRecursively(dirname, undeleted_files, undeleted_dirs);
 }
 
-Status Env::GetFileSize(const string& fname, uint64* file_size) {
+absl::Status Env::GetFileSize(const string& fname, uint64* file_size) {
   FileSystem* fs;
   TF_RETURN_IF_ERROR(GetFileSystemForFile(fname, &fs));
   return fs->GetFileSize(fname, file_size);
 }
 
-Status Env::RenameFile(const string& src, const string& target) {
+absl::Status Env::RenameFile(const string& src, const string& target) {
   FileSystem* src_fs;
   FileSystem* target_fs;
   TF_RETURN_IF_ERROR(GetFileSystemForFile(src, &src_fs));
@@ -352,7 +354,7 @@ Status Env::RenameFile(const string& src, const string& target) {
   return src_fs->RenameFile(src, target);
 }
 
-Status Env::CopyFile(const string& src, const string& target) {
+absl::Status Env::CopyFile(const string& src, const string& target) {
   FileSystem* src_fs;
   FileSystem* target_fs;
   TF_RETURN_IF_ERROR(GetFileSystemForFile(src, &src_fs));
@@ -464,9 +466,9 @@ Thread::~Thread() {}
 
 EnvWrapper::~EnvWrapper() {}
 
-Status ReadFileToString(Env* env, const string& fname, string* data) {
+absl::Status ReadFileToString(Env* env, const string& fname, string* data) {
   uint64 file_size;
-  Status s = env->GetFileSize(fname, &file_size);
+  absl::Status s = env->GetFileSize(fname, &file_size);
   if (!s.ok()) {
     return s;
   }
@@ -493,10 +495,10 @@ Status ReadFileToString(Env* env, const string& fname, string* data) {
   return s;
 }
 
-Status WriteStringToFile(Env* env, const string& fname,
-                         const StringPiece& data) {
+absl::Status WriteStringToFile(Env* env, const string& fname,
+                               const StringPiece& data) {
   std::unique_ptr<WritableFile> file;
-  Status s = env->NewWritableFile(fname, &file);
+  absl::Status s = env->NewWritableFile(fname, &file);
   if (!s.ok()) {
     return s;
   }
@@ -507,8 +509,8 @@ Status WriteStringToFile(Env* env, const string& fname,
   return s;
 }
 
-Status FileSystemCopyFile(FileSystem* src_fs, const string& src,
-                          FileSystem* target_fs, const string& target) {
+absl::Status FileSystemCopyFile(FileSystem* src_fs, const string& src,
+                                FileSystem* target_fs, const string& target) {
   std::unique_ptr<RandomAccessFile> src_file;
   TF_RETURN_IF_ERROR(src_fs->NewRandomAccessFile(src, &src_file));
 
@@ -525,7 +527,7 @@ Status FileSystemCopyFile(FileSystem* src_fs, const string& src,
 
   uint64 offset = 0;
   std::unique_ptr<char[]> scratch(new char[kCopyFileBufferSize]);
-  Status s = OkStatus();
+  absl::Status s = absl::OkStatus();
   while (s.ok()) {
     StringPiece result;
     s = src_file->Read(offset, kCopyFileBufferSize, &result, scratch.get());
@@ -550,11 +552,11 @@ class FileStream : public protobuf::io::ZeroCopyInputStream {
     return true;
   }
   int64_t ByteCount() const override { return pos_; }
-  Status status() const { return status_; }
+  absl::Status status() const { return status_; }
 
   bool Next(const void** data, int* size) override {
     StringPiece result;
-    Status s = file_->Read(pos_, kBufSize, &result, scratch_);
+    absl::Status s = file_->Read(pos_, kBufSize, &result, scratch_);
     if (result.empty()) {
       status_ = s;
       return false;
@@ -570,21 +572,21 @@ class FileStream : public protobuf::io::ZeroCopyInputStream {
 
   RandomAccessFile* file_;
   int64_t pos_;
-  Status status_;
+  absl::Status status_;
   char scratch_[kBufSize];
 };
 
 }  // namespace
 
-Status WriteBinaryProto(Env* env, const string& fname,
-                        const protobuf::MessageLite& proto) {
+absl::Status WriteBinaryProto(Env* env, const string& fname,
+                              const protobuf::MessageLite& proto) {
   string serialized;
   proto.AppendToString(&serialized);
   return WriteStringToFile(env, fname, serialized);
 }
 
-Status ReadBinaryProto(Env* env, const string& fname,
-                       protobuf::MessageLite* proto) {
+absl::Status ReadBinaryProto(Env* env, const string& fname,
+                             protobuf::MessageLite* proto) {
   std::unique_ptr<RandomAccessFile> file;
   TF_RETURN_IF_ERROR(env->NewRandomAccessFile(fname, &file));
   std::unique_ptr<FileStream> stream(new FileStream(file.get()));
@@ -595,11 +597,11 @@ Status ReadBinaryProto(Env* env, const string& fname,
     TF_RETURN_IF_ERROR(stream->status());
     return errors::DataLoss("Can't parse ", fname, " as binary proto");
   }
-  return OkStatus();
+  return absl::OkStatus();
 }
 
-Status WriteTextProto(Env* env, const string& fname,
-                      const protobuf::Message& proto) {
+absl::Status WriteTextProto(Env* env, const string& fname,
+                            const protobuf::Message& proto) {
   string serialized;
   if (!protobuf::TextFormat::PrintToString(proto, &serialized)) {
     return errors::FailedPrecondition("Unable to convert proto to text.");
@@ -607,7 +609,8 @@ Status WriteTextProto(Env* env, const string& fname,
   return WriteStringToFile(env, fname, serialized);
 }
 
-Status ReadTextProto(Env* env, const string& fname, protobuf::Message* proto) {
+absl::Status ReadTextProto(Env* env, const string& fname,
+                           protobuf::Message* proto) {
   std::unique_ptr<RandomAccessFile> file;
   TF_RETURN_IF_ERROR(env->NewRandomAccessFile(fname, &file));
   std::unique_ptr<FileStream> stream(new FileStream(file.get()));
@@ -616,19 +619,19 @@ Status ReadTextProto(Env* env, const string& fname, protobuf::Message* proto) {
     TF_RETURN_IF_ERROR(stream->status());
     return errors::DataLoss("Can't parse ", fname, " as text proto");
   }
-  return OkStatus();
+  return absl::OkStatus();
 }
 
-Status ReadTextOrBinaryProto(Env* env, const string& fname,
-                             protobuf::Message* proto) {
+absl::Status ReadTextOrBinaryProto(Env* env, const string& fname,
+                                   protobuf::Message* proto) {
   if (ReadTextProto(env, fname, proto).ok()) {
-    return OkStatus();
+    return absl::OkStatus();
   }
   return ReadBinaryProto(env, fname, proto);
 }
 
-Status ReadTextOrBinaryProto(Env* env, const string& fname,
-                             protobuf::MessageLite* proto) {
+absl::Status ReadTextOrBinaryProto(Env* env, const string& fname,
+                                   protobuf::MessageLite* proto) {
   return ReadBinaryProto(env, fname, proto);
 }
 
diff --git a/third_party/xla/third_party/tsl/tsl/platform/env.h b/third_party/xla/third_party/tsl/tsl/platform/env.h
index 35b446a99445a5..6953ac8722976a 100644
--- a/third_party/xla/third_party/tsl/tsl/platform/env.h
+++ b/third_party/xla/third_party/tsl/tsl/platform/env.h
@@ -79,16 +79,16 @@ class Env {
   /// for the file system related (non-virtual) functions that follow.
   /// Returned FileSystem object is still owned by the Env object and will
   // (might) be destroyed when the environment is destroyed.
-  virtual Status GetFileSystemForFile(const std::string& fname,
-                                      FileSystem** result);
+  virtual absl::Status GetFileSystemForFile(const std::string& fname,
+                                            FileSystem** result);
 
   /// \brief Returns the file system schemes registered for this Env.
-  virtual Status GetRegisteredFileSystemSchemes(
+  virtual absl::Status GetRegisteredFileSystemSchemes(
       std::vector<std::string>* schemes);
 
   /// \brief Register a file system for a scheme.
-  virtual Status RegisterFileSystem(const std::string& scheme,
-                                    FileSystemRegistry::Factory factory);
+  virtual absl::Status RegisterFileSystem(const std::string& scheme,
+                                          FileSystemRegistry::Factory factory);
 
   /// \brief Register a modular file system for a scheme.
   ///
@@ -96,23 +96,23 @@ class Env {
   ///
   /// TODO(b/139060984): After all filesystems are converted, make this be the
   /// canonical registration function.
-  virtual Status RegisterFileSystem(const std::string& scheme,
-                                    std::unique_ptr<FileSystem> filesystem);
+  virtual absl::Status RegisterFileSystem(
+      const std::string& scheme, std::unique_ptr<FileSystem> filesystem);
 
-  Status SetOption(const std::string& scheme, const std::string& key,
-                   const std::string& value);
+  absl::Status SetOption(const std::string& scheme, const std::string& key,
+                         const std::string& value);
 
-  Status SetOption(const std::string& scheme, const std::string& key,
-                   const std::vector<string>& values);
+  absl::Status SetOption(const std::string& scheme, const std::string& key,
+                         const std::vector<string>& values);
 
-  Status SetOption(const std::string& scheme, const std::string& key,
-                   const std::vector<int64_t>& values);
+  absl::Status SetOption(const std::string& scheme, const std::string& key,
+                         const std::vector<int64_t>& values);
 
-  Status SetOption(const std::string& scheme, const std::string& key,
-                   const std::vector<double>& values);
+  absl::Status SetOption(const std::string& scheme, const std::string& key,
+                         const std::vector<double>& values);
 
   /// \brief Flush filesystem caches for all registered filesystems.
-  Status FlushFileSystemCaches();
+  absl::Status FlushFileSystemCaches();
 
   /// \brief Creates a brand new random access read-only file with the
   /// specified name.
@@ -127,14 +127,15 @@ class Env {
   /// The ownership of the returned RandomAccessFile is passed to the caller
   /// and the object should be deleted when is not used. The file object
   /// shouldn't live longer than the Env object.
-  Status NewRandomAccessFile(const std::string& fname,
-                             std::unique_ptr<RandomAccessFile>* result);
+  absl::Status NewRandomAccessFile(const std::string& fname,
+                                   std::unique_ptr<RandomAccessFile>* result);
 
-  Status NewRandomAccessFile(const std::string& fname, TransactionToken* token,
-                             std::unique_ptr<RandomAccessFile>* result) {
+  absl::Status NewRandomAccessFile(const std::string& fname,
+                                   TransactionToken* token,
+                                   std::unique_ptr<RandomAccessFile>* result) {
     // We duplicate these methods due to Google internal coding style prevents
     // virtual functions with default arguments. See PR #41615.
-    return OkStatus();
+    return absl::OkStatus();
   }
 
   /// \brief Creates an object that writes to a new file with the specified
@@ -150,12 +151,13 @@ class Env {
   /// The ownership of the returned WritableFile is passed to the caller
   /// and the object should be deleted when is not used. The file object
   /// shouldn't live longer than the Env object.
-  Status NewWritableFile(const std::string& fname,
-                         std::unique_ptr<WritableFile>* result);
+  absl::Status NewWritableFile(const std::string& fname,
+                               std::unique_ptr<WritableFile>* result);
 
-  Status NewWritableFile(const std::string& fname, TransactionToken* token,
-                         std::unique_ptr<WritableFile>* result) {
-    return OkStatus();
+  absl::Status NewWritableFile(const std::string& fname,
+                               TransactionToken* token,
+                               std::unique_ptr<WritableFile>* result) {
+    return absl::OkStatus();
   }
 
   /// \brief Creates an object that either appends to an existing file, or
@@ -170,12 +172,13 @@ class Env {
   /// The ownership of the returned WritableFile is passed to the caller
   /// and the object should be deleted when is not used. The file object
   /// shouldn't live longer than the Env object.
-  Status NewAppendableFile(const std::string& fname,
-                           std::unique_ptr<WritableFile>* result);
+  absl::Status NewAppendableFile(const std::string& fname,
+                                 std::unique_ptr<WritableFile>* result);
 
-  Status NewAppendableFile(const std::string& fname, TransactionToken* token,
-                           std::unique_ptr<WritableFile>* result) {
-    return OkStatus();
+  absl::Status NewAppendableFile(const std::string& fname,
+                                 TransactionToken* token,
+                                 std::unique_ptr<WritableFile>* result) {
+    return absl::OkStatus();
   }
   /// \brief Creates a readonly region of memory with the file context.
   ///
@@ -188,30 +191,30 @@ class Env {
   /// The ownership of the returned ReadOnlyMemoryRegion is passed to the caller
   /// and the object should be deleted when is not used. The memory region
   /// object shouldn't live longer than the Env object.
-  Status NewReadOnlyMemoryRegionFromFile(
+  absl::Status NewReadOnlyMemoryRegionFromFile(
       const std::string& fname, std::unique_ptr<ReadOnlyMemoryRegion>* result);
 
-  Status NewReadOnlyMemoryRegionFromFile(
+  absl::Status NewReadOnlyMemoryRegionFromFile(
       const std::string& fname, TransactionToken* token,
       std::unique_ptr<ReadOnlyMemoryRegion>* result) {
-    return OkStatus();
+    return absl::OkStatus();
   }
 
   /// Returns OK if the named path exists and NOT_FOUND otherwise.
-  Status FileExists(const std::string& fname);
+  absl::Status FileExists(const std::string& fname);
 
-  Status FileExists(const std::string& fname, TransactionToken* token) {
-    return OkStatus();
+  absl::Status FileExists(const std::string& fname, TransactionToken* token) {
+    return absl::OkStatus();
   }
 
   /// Returns true if all the listed files exist, false otherwise.
   /// if status is not null, populate the vector with a detailed status
   /// for each file.
   bool FilesExist(const std::vector<string>& files,
-                  std::vector<Status>* status);
+                  std::vector<absl::Status>* status);
 
   bool FilesExist(const std::vector<string>& files, TransactionToken* token,
-                  std::vector<Status>* status) {
+                  std::vector<absl::Status>* status) {
     return true;
   }
 
@@ -219,11 +222,11 @@ class Env {
   /// directory. The names are relative to "dir".
   ///
   /// Original contents of *results are dropped.
-  Status GetChildren(const std::string& dir, std::vector<string>* result);
+  absl::Status GetChildren(const std::string& dir, std::vector<string>* result);
 
-  Status GetChildren(const std::string& dir, TransactionToken* token,
-                     std::vector<string>* result) {
-    return OkStatus();
+  absl::Status GetChildren(const std::string& dir, TransactionToken* token,
+                           std::vector<string>* result) {
+    return absl::OkStatus();
   }
 
   /// \brief Returns true if the path matches the given pattern. The wildcards
@@ -235,19 +238,20 @@ class Env {
   /// that pattern. *results is cleared.
   ///
   /// More details about `pattern` in FileSystem::GetMatchingPaths.
-  virtual Status GetMatchingPaths(const std::string& pattern,
-                                  std::vector<string>* results);
+  virtual absl::Status GetMatchingPaths(const std::string& pattern,
+                                        std::vector<string>* results);
 
-  Status GetMatchingPaths(const std::string& pattern, TransactionToken* token,
-                          std::vector<string>* results) {
-    return OkStatus();
+  absl::Status GetMatchingPaths(const std::string& pattern,
+                                TransactionToken* token,
+                                std::vector<string>* results) {
+    return absl::OkStatus();
   }
 
   /// Deletes the named file.
-  Status DeleteFile(const std::string& fname);
+  absl::Status DeleteFile(const std::string& fname);
 
-  Status DeleteFile(const std::string& fname, TransactionToken* token) {
-    return OkStatus();
+  absl::Status DeleteFile(const std::string& fname, TransactionToken* token) {
+    return absl::OkStatus();
   }
 
   /// \brief Deletes the specified directory and all subdirectories and files
@@ -274,12 +278,15 @@ class Env {
   ///  * PERMISSION_DENIED - dirname or some descendant is not writable
   ///  * UNIMPLEMENTED - Some underlying functions (like Delete) are not
   ///                    implemented
-  Status DeleteRecursively(const std::string& dirname, int64_t* undeleted_files,
-                           int64_t* undeleted_dirs);
+  absl::Status DeleteRecursively(const std::string& dirname,
+                                 int64_t* undeleted_files,
+                                 int64_t* undeleted_dirs);
 
-  Status DeleteRecursively(const std::string& dirname, TransactionToken* token,
-                           int64_t* undeleted_files, int64_t* undeleted_dirs) {
-    return OkStatus();
+  absl::Status DeleteRecursively(const std::string& dirname,
+                                 TransactionToken* token,
+                                 int64_t* undeleted_files,
+                                 int64_t* undeleted_dirs) {
+    return absl::OkStatus();
   }
 
   /// \brief Creates the specified directory and all the necessary
@@ -287,35 +294,35 @@ class Env {
   ///  * OK - successfully created the directory and sub directories, even if
   ///         they were already created.
   ///  * PERMISSION_DENIED - dirname or some subdirectory is not writable.
-  Status RecursivelyCreateDir(const std::string& dirname);
+  absl::Status RecursivelyCreateDir(const std::string& dirname);
 
-  Status RecursivelyCreateDir(const std::string& dirname,
-                              TransactionToken* token) {
-    return OkStatus();
+  absl::Status RecursivelyCreateDir(const std::string& dirname,
+                                    TransactionToken* token) {
+    return absl::OkStatus();
   }
   /// \brief Creates the specified directory. Typical return codes
   ///  * OK - successfully created the directory.
   ///  * ALREADY_EXISTS - directory already exists.
   ///  * PERMISSION_DENIED - dirname is not writable.
-  Status CreateDir(const std::string& dirname);
+  absl::Status CreateDir(const std::string& dirname);
 
-  Status CreateDir(const std::string& dirname, TransactionToken* token) {
-    return OkStatus();
+  absl::Status CreateDir(const std::string& dirname, TransactionToken* token) {
+    return absl::OkStatus();
   }
 
   /// Deletes the specified directory.
-  Status DeleteDir(const std::string& dirname);
+  absl::Status DeleteDir(const std::string& dirname);
 
-  Status DeleteDir(const std::string& dirname, TransactionToken* token) {
-    return OkStatus();
+  absl::Status DeleteDir(const std::string& dirname, TransactionToken* token) {
+    return absl::OkStatus();
   }
 
   /// Obtains statistics for the given path.
-  Status Stat(const std::string& fname, FileStatistics* stat);
+  absl::Status Stat(const std::string& fname, FileStatistics* stat);
 
-  Status Stat(const std::string& fname, TransactionToken* token,
-              FileStatistics* stat) {
-    return OkStatus();
+  absl::Status Stat(const std::string& fname, TransactionToken* token,
+                    FileStatistics* stat) {
+    return absl::OkStatus();
   }
 
   /// \brief Returns whether the given path is a directory or not.
@@ -325,7 +332,7 @@ class Env {
   ///  * NOT_FOUND - The path entry does not exist.
   ///  * PERMISSION_DENIED - Insufficient permissions.
   ///  * UNIMPLEMENTED - The file factory doesn't support directories.
-  Status IsDirectory(const std::string& fname);
+  absl::Status IsDirectory(const std::string& fname);
 
   /// \brief Returns whether the given path is on a file system
   /// that has atomic move capabilities. This can be used
@@ -337,63 +344,66 @@ class Env {
   ///         so has_atomic_move holds the above information.
   ///  * UNIMPLEMENTED - The file system of the path hasn't been implemented in
   ///  TF
-  Status HasAtomicMove(const std::string& path, bool* has_atomic_move);
+  absl::Status HasAtomicMove(const std::string& path, bool* has_atomic_move);
 
   /// Stores the size of `fname` in `*file_size`.
-  Status GetFileSize(const std::string& fname, uint64* file_size);
+  absl::Status GetFileSize(const std::string& fname, uint64* file_size);
 
-  Status GetFileSize(const std::string& fname, TransactionToken* token,
-                     uint64* file_size) {
-    return OkStatus();
+  absl::Status GetFileSize(const std::string& fname, TransactionToken* token,
+                           uint64* file_size) {
+    return absl::OkStatus();
   }
 
   /// \brief Renames file src to target. If target already exists, it will be
   /// replaced.
-  Status RenameFile(const std::string& src, const std::string& target);
+  absl::Status RenameFile(const std::string& src, const std::string& target);
 
-  Status RenameFile(const std::string& src, const std::string& target,
-                    TransactionToken* token) {
-    return OkStatus();
+  absl::Status RenameFile(const std::string& src, const std::string& target,
+                          TransactionToken* token) {
+    return absl::OkStatus();
   }
 
   /// \brief Copy the src to target.
-  Status CopyFile(const std::string& src, const std::string& target);
+  absl::Status CopyFile(const std::string& src, const std::string& target);
 
-  Status CopyFile(const std::string& src, const std::string& target,
-                  TransactionToken* token) {
-    return OkStatus();
+  absl::Status CopyFile(const std::string& src, const std::string& target,
+                        TransactionToken* token) {
+    return absl::OkStatus();
   }
 
   /// \brief starts a new transaction on the filesystem that handles filename
-  Status StartTransaction(const std::string& filename,
-                          TransactionToken** token) {
+  absl::Status StartTransaction(const std::string& filename,
+                                TransactionToken** token) {
     *token = nullptr;
-    return OkStatus();
+    return absl::OkStatus();
   }
 
   /// \brief Adds `path` to transaction in `token` if token belongs to
   /// filesystem that handles the path.
-  Status AddToTransaction(const std::string& path, TransactionToken* token) {
-    return OkStatus();
+  absl::Status AddToTransaction(const std::string& path,
+                                TransactionToken* token) {
+    return absl::OkStatus();
   }
 
   /// \brief Get token for `path` or start a new transaction and add `path` to
   /// it.
-  Status GetTokenOrStartTransaction(const std::string& path,
-                                    TransactionToken** token) {
+  absl::Status GetTokenOrStartTransaction(const std::string& path,
+                                          TransactionToken** token) {
     *token = nullptr;
-    return OkStatus();
+    return absl::OkStatus();
   }
 
   /// \brief Returns the transaction for `path` or nullptr in `token`
-  Status GetTransactionForPath(const std::string& path,
-                               TransactionToken** token) {
+  absl::Status GetTransactionForPath(const std::string& path,
+                                     TransactionToken** token) {
     *token = nullptr;
-    return OkStatus();
+    return absl::OkStatus();
   }
 
   /// \brief Finalizes the transaction
-  Status EndTransaction(TransactionToken* token) { return OkStatus(); }
+  absl::Status EndTransaction(TransactionToken* token) {
+    return absl::OkStatus();
+  }
 
   /// \brief Returns the absolute path of the current executable. It resolves
   /// symlinks if there is any.
@@ -469,8 +479,8 @@ class Env {
   // OK from the function.
   // Otherwise returns nullptr in "*handle" and an error status from the
   // function.
-  virtual Status LoadDynamicLibrary(const char* library_filename,
-                                    void** handle) = 0;
+  virtual absl::Status LoadDynamicLibrary(const char* library_filename,
+                                          void** handle) = 0;
 
   // \brief Get a pointer to a symbol from a dynamic library.
   //
@@ -478,8 +488,9 @@ class Env {
   // On success, store a pointer to the located symbol in "*symbol" and return
   // OK from the function. Otherwise, returns nullptr in "*symbol" and an error
   // status from the function.
-  virtual Status GetSymbolFromLibrary(void* handle, const char* symbol_name,
-                                      void** symbol) = 0;
+  virtual absl::Status GetSymbolFromLibrary(void* handle,
+                                            const char* symbol_name,
+                                            void** symbol) = 0;
 
   // \brief build the name of dynamic library.
   //
@@ -511,17 +522,18 @@ class EnvWrapper : public Env {
   /// Returns the target to which this Env forwards all calls
   Env* target() const { return target_; }
 
-  Status GetFileSystemForFile(const std::string& fname,
-                              FileSystem** result) override {
+  absl::Status GetFileSystemForFile(const std::string& fname,
+                                    FileSystem** result) override {
     return target_->GetFileSystemForFile(fname, result);
   }
 
-  Status GetRegisteredFileSystemSchemes(std::vector<string>* schemes) override {
+  absl::Status GetRegisteredFileSystemSchemes(
+      std::vector<string>* schemes) override {
     return target_->GetRegisteredFileSystemSchemes(schemes);
   }
 
-  Status RegisterFileSystem(const std::string& scheme,
-                            FileSystemRegistry::Factory factory) override {
+  absl::Status RegisterFileSystem(
+      const std::string& scheme, FileSystemRegistry::Factory factory) override {
     return target_->RegisterFileSystem(scheme, factory);
   }
 
@@ -549,12 +561,12 @@ class EnvWrapper : public Env {
                          absl::AnyInvocable<void()> closure) override {
     target_->SchedClosureAfter(micros, std::move(closure));
   }
-  Status LoadDynamicLibrary(const char* library_filename,
-                            void** handle) override {
+  absl::Status LoadDynamicLibrary(const char* library_filename,
+                                  void** handle) override {
     return target_->LoadDynamicLibrary(library_filename, handle);
   }
-  Status GetSymbolFromLibrary(void* handle, const char* symbol_name,
-                              void** symbol) override {
+  absl::Status GetSymbolFromLibrary(void* handle, const char* symbol_name,
+                                    void** symbol) override {
     return target_->GetSymbolFromLibrary(handle, symbol_name, symbol);
   }
   std::string FormatLibraryFileName(const std::string& name,
@@ -608,49 +620,53 @@ struct ThreadOptions {
 
 /// A utility routine: copy contents of `src` in file system `src_fs`
 /// to `target` in file system `target_fs`.
-Status FileSystemCopyFile(FileSystem* src_fs, const std::string& src,
-                          FileSystem* target_fs, const std::string& target);
+absl::Status FileSystemCopyFile(FileSystem* src_fs, const std::string& src,
+                                FileSystem* target_fs,
+                                const std::string& target);
 
 /// A utility routine: reads contents of named file into `*data`
-Status ReadFileToString(Env* env, const std::string& fname, std::string* data);
+absl::Status ReadFileToString(Env* env, const std::string& fname,
+                              std::string* data);
 
 /// A utility routine: write contents of `data` to file named `fname`
 /// (overwriting existing contents, if any).
-Status WriteStringToFile(Env* env, const std::string& fname,
-                         const StringPiece& data);
+absl::Status WriteStringToFile(Env* env, const std::string& fname,
+                               const StringPiece& data);
 
 /// Write binary representation of "proto" to the named file.
-Status WriteBinaryProto(Env* env, const std::string& fname,
-                        const protobuf::MessageLite& proto);
+absl::Status WriteBinaryProto(Env* env, const std::string& fname,
+                              const protobuf::MessageLite& proto);
 
 /// Reads contents of named file and parse as binary encoded proto data
 /// and store into `*proto`.
-Status ReadBinaryProto(Env* env, const std::string& fname,
-                       protobuf::MessageLite* proto);
+absl::Status ReadBinaryProto(Env* env, const std::string& fname,
+                             protobuf::MessageLite* proto);
 
 /// Write the text representation of "proto" to the named file.
-inline Status WriteTextProto(Env* /* env */, const std::string& /* fname */,
-                             const protobuf::MessageLite& /* proto */) {
+inline absl::Status WriteTextProto(Env* /* env */,
+                                   const std::string& /* fname */,
+                                   const protobuf::MessageLite& /* proto */) {
   return errors::Unimplemented("Can't write text protos with protolite.");
 }
-Status WriteTextProto(Env* env, const std::string& fname,
-                      const protobuf::Message& proto);
+absl::Status WriteTextProto(Env* env, const std::string& fname,
+                            const protobuf::Message& proto);
 
 /// Read contents of named file and parse as text encoded proto data
 /// and store into `*proto`.
-inline Status ReadTextProto(Env* /* env */, const std::string& /* fname */,
-                            protobuf::MessageLite* /* proto */) {
+inline absl::Status ReadTextProto(Env* /* env */,
+                                  const std::string& /* fname */,
+                                  protobuf::MessageLite* /* proto */) {
   return errors::Unimplemented("Can't parse text protos with protolite.");
 }
-Status ReadTextProto(Env* env, const std::string& fname,
-                     protobuf::Message* proto);
+absl::Status ReadTextProto(Env* env, const std::string& fname,
+                           protobuf::Message* proto);
 
 /// Read contents of named file and parse as either text or binary encoded proto
 /// data and store into `*proto`.
-Status ReadTextOrBinaryProto(Env* env, const std::string& fname,
-                             protobuf::Message* proto);
-Status ReadTextOrBinaryProto(Env* env, const std::string& fname,
-                             protobuf::MessageLite* proto);
+absl::Status ReadTextOrBinaryProto(Env* env, const std::string& fname,
+                                   protobuf::Message* proto);
+absl::Status ReadTextOrBinaryProto(Env* env, const std::string& fname,
+                                   protobuf::MessageLite* proto);
 
 // START_SKIP_DOXYGEN
 
diff --git a/third_party/xla/third_party/tsl/tsl/platform/errors.cc b/third_party/xla/third_party/tsl/tsl/platform/errors.cc
index 7f61fb678954ed..6c732a47849113 100644
--- a/third_party/xla/third_party/tsl/tsl/platform/errors.cc
+++ b/third_party/xla/third_party/tsl/tsl/platform/errors.cc
@@ -175,72 +175,73 @@ absl::StatusCode ErrnoToCode(int err_number) {
 
 }  // namespace
 
-Status IOError(const string& context, int err_number) {
+absl::Status IOError(const string& context, int err_number) {
   auto code = ErrnoToCode(err_number);
-  return Status(code, strings::StrCat(context, "; ", strerror(err_number)));
+  return absl::Status(code,
+                      strings::StrCat(context, "; ", strerror(err_number)));
 }
 
-bool IsAborted(const Status& status) {
+bool IsAborted(const absl::Status& status) {
   return status.code() == tsl::error::Code::ABORTED;
 }
 
-bool IsAlreadyExists(const Status& status) {
+bool IsAlreadyExists(const absl::Status& status) {
   return status.code() == tsl::error::Code::ALREADY_EXISTS;
 }
 
-bool IsCancelled(const Status& status) {
+bool IsCancelled(const absl::Status& status) {
   return status.code() == tsl::error::Code::CANCELLED;
 }
 
-bool IsDataLoss(const Status& status) {
+bool IsDataLoss(const absl::Status& status) {
   return status.code() == tsl::error::Code::DATA_LOSS;
 }
 
-bool IsDeadlineExceeded(const Status& status) {
+bool IsDeadlineExceeded(const absl::Status& status) {
   return status.code() == tsl::error::Code::DEADLINE_EXCEEDED;
 }
 
-bool IsFailedPrecondition(const Status& status) {
+bool IsFailedPrecondition(const absl::Status& status) {
   return status.code() == tsl::error::Code::FAILED_PRECONDITION;
 }
 
-bool IsInternal(const Status& status) {
+bool IsInternal(const absl::Status& status) {
   return status.code() == tsl::error::Code::INTERNAL;
 }
 
-bool IsInvalidArgument(const Status& status) {
+bool IsInvalidArgument(const absl::Status& status) {
   return status.code() == tsl::error::Code::INVALID_ARGUMENT;
 }
 
-bool IsNotFound(const Status& status) {
+bool IsNotFound(const absl::Status& status) {
   return status.code() == tsl::error::Code::NOT_FOUND;
 }
 
-bool IsOutOfRange(const Status& status) {
+bool IsOutOfRange(const absl::Status& status) {
   return status.code() == tsl::error::Code::OUT_OF_RANGE;
 }
 
-bool IsPermissionDenied(const Status& status) {
+bool IsPermissionDenied(const absl::Status& status) {
   return status.code() == tsl::error::Code::PERMISSION_DENIED;
 }
 
-bool IsResourceExhausted(const Status& status) {
+bool IsResourceExhausted(const absl::Status& status) {
   return status.code() == tsl::error::Code::RESOURCE_EXHAUSTED;
 }
 
-bool IsUnauthenticated(const Status& status) {
+bool IsUnauthenticated(const absl::Status& status) {
   return status.code() == tsl::error::Code::UNAUTHENTICATED;
 }
 
-bool IsUnavailable(const Status& status) {
+bool IsUnavailable(const absl::Status& status) {
   return status.code() == tsl::error::Code::UNAVAILABLE;
 }
 
-bool IsUnimplemented(const Status& status) {
+bool IsUnimplemented(const absl::Status& status) {
   return status.code() == tsl::error::Code::UNIMPLEMENTED;
 }
 
-bool IsUnknown(const Status& status) {
+bool IsUnknown(const absl::Status& status) {
   return status.code() == tsl::error::Code::UNKNOWN;
 }
 
diff --git a/third_party/xla/third_party/tsl/tsl/platform/errors.h b/third_party/xla/third_party/tsl/tsl/platform/errors.h
index 9008dedad8270c..1d86af35cd3efd 100644
--- a/third_party/xla/third_party/tsl/tsl/platform/errors.h
+++ b/third_party/xla/third_party/tsl/tsl/platform/errors.h
@@ -88,11 +88,11 @@ inline const strings::AlphaNum& PrepareForStrCat(const strings::AlphaNum& a) {
 }  // namespace internal
 
 // Maps UNIX errors into a Status.
-Status IOError(const string& context, int err_number);
+absl::Status IOError(const string& context, int err_number);
 
 // Returns all payloads from a Status as a key-value map.
 inline std::unordered_map<std::string, std::string> GetPayloads(
-    const ::tsl::Status& status) {
+    const absl::Status& status) {
   std::unordered_map<std::string, std::string> payloads;
   status.ForEachPayload(
       [&payloads](::tsl::StringPiece key, const absl::Cord& value) {
@@ -104,7 +104,7 @@ inline std::unordered_map<std::string, std::string> GetPayloads(
 // Inserts all given payloads into the given status. Will overwrite existing
 // payloads if they exist with the same key.
 inline void InsertPayloads(
-    ::tsl::Status& status,
+    absl::Status& status,
     const std::unordered_map<std::string, std::string>& payloads) {
   for (const auto& payload : payloads) {
     status.SetPayload(payload.first, absl::Cord(payload.second));
@@ -113,7 +113,7 @@ inline void InsertPayloads(
 
 // Copies all payloads from one Status to another. Will overwrite existing
 // payloads in the destination if they exist with the same key.
-inline void CopyPayloads(const ::tsl::Status& from, ::tsl::Status& to) {
+inline void CopyPayloads(const absl::Status& from, absl::Status& to) {
   from.ForEachPayload([&to](::tsl::StringPiece key, const absl::Cord& value) {
     to.SetPayload(key, value);
   });
@@ -121,22 +121,22 @@ inline void CopyPayloads(const ::tsl::Status& from, ::tsl::Status& to) {
 
 #if defined(PLATFORM_GOOGLE)
 // Creates a new status with the given code, message and payloads.
-inline ::tsl::Status Create(
+inline absl::Status Create(
     absl::StatusCode code, ::tsl::StringPiece message,
     const std::unordered_map<std::string, std::string>& payloads,
     absl::SourceLocation loc = absl::SourceLocation::current()) {
-  Status status(code, message, loc);
+  absl::Status status(code, message, loc);
   InsertPayloads(status, payloads);
   return status;
 }
 // Returns a new Status, replacing its message with the given.
-inline ::tsl::Status CreateWithUpdatedMessage(const ::tsl::Status& status,
-                                              ::tsl::StringPiece message) {
+inline absl::Status CreateWithUpdatedMessage(const absl::Status& status,
+                                             ::tsl::StringPiece message) {
   auto locations = status.GetSourceLocations();
   auto initial_loc =
       locations.empty() ? absl::SourceLocation::current() : locations[0];
-  Status new_status = Create(static_cast<absl::StatusCode>(status.code()),
-                             message, GetPayloads(status), initial_loc);
+  absl::Status new_status = Create(static_cast<absl::StatusCode>(status.code()),
+                                   message, GetPayloads(status), initial_loc);
   if (locations.size() > 1) {
     for (auto loc : locations.subspan(1)) {
       new_status.AddSourceLocation(loc);
@@ -165,7 +165,7 @@ inline ::tsl::Status CreateWithUpdatedMessage(const ::tsl::Status& status,
 // context put it on a new line, since it is possible for there
 // to be several layers of additional context.
 template <typename... Args>
-void AppendToMessage(::tsl::Status* status, Args... args) {
+void AppendToMessage(absl::Status* status, Args... args) {
   auto new_status = CreateWithUpdatedMessage(
       *status, ::tsl::strings::StrCat(status->message(), "\n\t", args...));
   CopyPayloads(*status, new_status);
@@ -199,13 +199,13 @@ void AppendToMessage(::tsl::Status* status, Args... args) {
 
 // CANCELLED
 template <typename... Args>
-::tsl::Status Cancelled(Args... args) {
-  return ::tsl::Status(absl::StatusCode::kCancelled,
-                       ::tsl::strings::StrCat(
-                           ::tsl::errors::internal::PrepareForStrCat(args)...));
+absl::Status Cancelled(Args... args) {
+  return absl::Status(absl::StatusCode::kCancelled,
+                      ::tsl::strings::StrCat(
+                          ::tsl::errors::internal::PrepareForStrCat(args)...));
 }
 template <typename... Args>
-::tsl::Status CancelledWithPayloads(
+absl::Status CancelledWithPayloads(
     const ::tsl::StringPiece& message,
     const std::unordered_map<std::string, std::string>& payloads) {
   return errors::Create(absl::StatusCode::kCancelled, message, payloads);
@@ -213,10 +213,10 @@ ::tsl::Status CancelledWithPayloads(
 
 // InvalidArgument
 template <typename... Args>
-::tsl::Status InvalidArgument(Args... args) {
-  return ::tsl::Status(absl::StatusCode::kInvalidArgument,
-                       ::tsl::strings::StrCat(
-                           ::tsl::errors::internal::PrepareForStrCat(args)...));
+absl::Status InvalidArgument(Args... args) {
+  return absl::Status(absl::StatusCode::kInvalidArgument,
+                      ::tsl::strings::StrCat(
+                          ::tsl::errors::internal::PrepareForStrCat(args)...));
 }
 
 #if defined(PLATFORM_GOOGLE)
@@ -225,7 +225,7 @@ template <typename Arg1, typename Arg2, typename Arg3, typename Arg4>
 ::absl::Status InvalidArgument(
     Arg1 arg1, Arg2 arg2, Arg3 arg3, Arg4 arg4,
     absl::SourceLocation loc = absl::SourceLocation::current()) {
-  return ::tsl::Status(
+  return absl::Status(
       absl::StatusCode::kInvalidArgument,
       ::tsl::strings::StrCat(::tsl::errors::internal::PrepareForStrCat(arg1),
                              ::tsl::errors::internal::PrepareForStrCat(arg2),
@@ -237,7 +237,7 @@ template <typename Arg1, typename Arg2, typename Arg3>
 ::absl::Status InvalidArgument(
     Arg1 arg1, Arg2 arg2, Arg3 arg3,
     absl::SourceLocation loc = absl::SourceLocation::current()) {
-  return ::tsl::Status(
+  return absl::Status(
       absl::StatusCode::kInvalidArgument,
       ::tsl::strings::StrCat(::tsl::errors::internal::PrepareForStrCat(arg1),
                              ::tsl::errors::internal::PrepareForStrCat(arg2),
@@ -248,7 +248,7 @@ template <typename Arg1, typename Arg2>
 ::absl::Status InvalidArgument(
     Arg1 arg1, Arg2 arg2,
     absl::SourceLocation loc = absl::SourceLocation::current()) {
-  return ::tsl::Status(
+  return absl::Status(
       absl::StatusCode::kInvalidArgument,
       ::tsl::strings::StrCat(::tsl::errors::internal::PrepareForStrCat(arg1),
                              ::tsl::errors::internal::PrepareForStrCat(arg2)),
@@ -257,7 +257,7 @@ ::absl::Status InvalidArgument(
 template <typename Arg1>
 ::absl::Status InvalidArgument(
     Arg1 arg1, absl::SourceLocation loc = absl::SourceLocation::current()) {
-  return ::tsl::Status(
+  return absl::Status(
       absl::StatusCode::kInvalidArgument,
       ::tsl::strings::StrCat(::tsl::errors::internal::PrepareForStrCat(arg1)),
       loc);
@@ -302,10 +302,10 @@ ::absl::Status InvalidArgumentWithPayloads(
 
 // NotFound
 template <typename... Args>
-::tsl::Status NotFound(Args... args) {
-  return ::tsl::Status(absl::StatusCode::kNotFound,
-                       ::tsl::strings::StrCat(
-                           ::tsl::errors::internal::PrepareForStrCat(args)...));
+absl::Status NotFound(Args... args) {
+  return absl::Status(absl::StatusCode::kNotFound,
+                      ::tsl::strings::StrCat(
+                          ::tsl::errors::internal::PrepareForStrCat(args)...));
 }
 #if defined(PLATFORM_GOOGLE)
 // Specialized overloads to capture source location for up to three arguments.
@@ -313,7 +313,7 @@ template <typename Arg1, typename Arg2, typename Arg3>
 ::absl::Status NotFound(
     Arg1 arg1, Arg2 arg2, Arg3 arg3,
     absl::SourceLocation loc = absl::SourceLocation::current()) {
-  return ::tsl::Status(
+  return absl::Status(
       absl::StatusCode::kNotFound,
       ::tsl::strings::StrCat(::tsl::errors::internal::PrepareForStrCat(arg1),
                              ::tsl::errors::internal::PrepareForStrCat(arg2),
@@ -324,7 +324,7 @@ template <typename Arg1, typename Arg2>
 ::absl::Status NotFound(
     Arg1 arg1, Arg2 arg2,
     absl::SourceLocation loc = absl::SourceLocation::current()) {
-  return ::tsl::Status(
+  return absl::Status(
       absl::StatusCode::kNotFound,
       ::tsl::strings::StrCat(::tsl::errors::internal::PrepareForStrCat(arg1),
                              ::tsl::errors::internal::PrepareForStrCat(arg2)),
@@ -333,7 +333,7 @@ ::absl::Status NotFound(
 template <typename Arg1>
 ::absl::Status NotFound(
     Arg1 arg1, absl::SourceLocation loc = absl::SourceLocation::current()) {
-  return ::tsl::Status(
+  return absl::Status(
       absl::StatusCode::kNotFound,
       ::tsl::strings::StrCat(::tsl::errors::internal::PrepareForStrCat(arg1)),
       loc);
@@ -377,13 +377,13 @@ ::absl::Status NotFoundWithPayloads(
 
 // AlreadyExists
 template <typename... Args>
-::tsl::Status AlreadyExists(Args... args) {
-  return ::tsl::Status(absl::StatusCode::kAlreadyExists,
-                       ::tsl::strings::StrCat(
-                           ::tsl::errors::internal::PrepareForStrCat(args)...));
+absl::Status AlreadyExists(Args... args) {
+  return absl::Status(absl::StatusCode::kAlreadyExists,
+                      ::tsl::strings::StrCat(
+                          ::tsl::errors::internal::PrepareForStrCat(args)...));
 }
 template <typename... Args>
-::tsl::Status AlreadyExistsWithPayloads(
+absl::Status AlreadyExistsWithPayloads(
     const ::tsl::StringPiece& message,
     const std::unordered_map<std::string, std::string>& payloads) {
   return errors::Create(absl::StatusCode::kAlreadyExists, message, payloads);
@@ -391,13 +391,13 @@ ::tsl::Status AlreadyExistsWithPayloads(
 
 // ResourceExhausted
 template <typename... Args>
-::tsl::Status ResourceExhausted(Args... args) {
-  return ::tsl::Status(absl::StatusCode::kResourceExhausted,
-                       ::tsl::strings::StrCat(
-                           ::tsl::errors::internal::PrepareForStrCat(args)...));
+absl::Status ResourceExhausted(Args... args) {
+  return absl::Status(absl::StatusCode::kResourceExhausted,
+                      ::tsl::strings::StrCat(
+                          ::tsl::errors::internal::PrepareForStrCat(args)...));
 }
 template <typename... Args>
-::tsl::Status ResourceExhaustedWithPayloads(
+absl::Status ResourceExhaustedWithPayloads(
     const ::tsl::StringPiece& message,
     const std::unordered_map<std::string, std::string>& payloads) {
   return errors::Create(absl::StatusCode::kResourceExhausted, message,
@@ -406,13 +406,13 @@ ::tsl::Status ResourceExhaustedWithPayloads(
 
 // Unavailable
 template <typename... Args>
-::tsl::Status Unavailable(Args... args) {
-  return ::tsl::Status(absl::StatusCode::kUnavailable,
-                       ::tsl::strings::StrCat(
-                           ::tsl::errors::internal::PrepareForStrCat(args)...));
+absl::Status Unavailable(Args... args) {
+  return absl::Status(absl::StatusCode::kUnavailable,
+                      ::tsl::strings::StrCat(
+                          ::tsl::errors::internal::PrepareForStrCat(args)...));
 }
 template <typename... Args>
-::tsl::Status UnavailableWithPayloads(
+absl::Status UnavailableWithPayloads(
     const ::tsl::StringPiece& message,
     const std::unordered_map<std::string, std::string>& payloads) {
   return errors::Create(absl::StatusCode::kUnavailable, message, payloads);
@@ -420,13 +420,13 @@ ::tsl::Status UnavailableWithPayloads(
 
 // FailedPrecondition
 template <typename... Args>
-::tsl::Status FailedPrecondition(Args... args) {
-  return ::tsl::Status(absl::StatusCode::kFailedPrecondition,
-                       ::tsl::strings::StrCat(
-                           ::tsl::errors::internal::PrepareForStrCat(args)...));
+absl::Status FailedPrecondition(Args... args) {
+  return absl::Status(absl::StatusCode::kFailedPrecondition,
+                      ::tsl::strings::StrCat(
+                          ::tsl::errors::internal::PrepareForStrCat(args)...));
 }
 template <typename... Args>
-::tsl::Status FailedPreconditionWithPayloads(
+absl::Status FailedPreconditionWithPayloads(
     const ::tsl::StringPiece& message,
     const std::unordered_map<std::string, std::string>& payloads) {
   return errors::Create(absl::StatusCode::kFailedPrecondition, message,
@@ -435,13 +435,13 @@ ::tsl::Status FailedPreconditionWithPayloads(
 
 // OutOfRange
 template <typename... Args>
-::tsl::Status OutOfRange(Args... args) {
-  return ::tsl::Status(absl::StatusCode::kOutOfRange,
-                       ::tsl::strings::StrCat(
-                           ::tsl::errors::internal::PrepareForStrCat(args)...));
+absl::Status OutOfRange(Args... args) {
+  return absl::Status(absl::StatusCode::kOutOfRange,
+                      ::tsl::strings::StrCat(
+                          ::tsl::errors::internal::PrepareForStrCat(args)...));
 }
 template <typename... Args>
-::tsl::Status OutOfRangeWithPayloads(
+absl::Status OutOfRangeWithPayloads(
     const ::tsl::StringPiece& message,
     const std::unordered_map<std::string, std::string>& payloads) {
   return errors::Create(absl::StatusCode::kOutOfRange, message, payloads);
@@ -449,13 +449,13 @@ ::tsl::Status OutOfRangeWithPayloads(
 
 // Unimplemented
 template <typename... Args>
-::tsl::Status Unimplemented(Args... args) {
-  return ::tsl::Status(absl::StatusCode::kUnimplemented,
-                       ::tsl::strings::StrCat(
-                           ::tsl::errors::internal::PrepareForStrCat(args)...));
+absl::Status Unimplemented(Args... args) {
+  return absl::Status(absl::StatusCode::kUnimplemented,
+                      ::tsl::strings::StrCat(
+                          ::tsl::errors::internal::PrepareForStrCat(args)...));
 }
 template <typename... Args>
-::tsl::Status UnimplementedWithPayloads(
+absl::Status UnimplementedWithPayloads(
     const ::tsl::StringPiece& message,
     const std::unordered_map<std::string, std::string>& payloads) {
   return errors::Create(absl::StatusCode::kUnimplemented, message, payloads);
@@ -463,13 +463,13 @@ ::tsl::Status UnimplementedWithPayloads(
 
 // Internal
 template <typename... Args>
-::tsl::Status Internal(Args... args) {
-  return ::tsl::Status(absl::StatusCode::kInternal,
-                       ::tsl::strings::StrCat(
-                           ::tsl::errors::internal::PrepareForStrCat(args)...));
+absl::Status Internal(Args... args) {
+  return absl::Status(absl::StatusCode::kInternal,
+                      ::tsl::strings::StrCat(
+                          ::tsl::errors::internal::PrepareForStrCat(args)...));
 }
 template <typename... Args>
-::tsl::Status InternalWithPayloads(
+absl::Status InternalWithPayloads(
     const ::tsl::StringPiece& message,
     const std::unordered_map<std::string, std::string>& payloads) {
   return errors::Create(absl::StatusCode::kInternal, message, payloads);
@@ -477,13 +477,13 @@ ::tsl::Status InternalWithPayloads(
 
 // Aborted
 template <typename... Args>
-::tsl::Status Aborted(Args... args) {
-  return ::tsl::Status(absl::StatusCode::kAborted,
-                       ::tsl::strings::StrCat(
-                           ::tsl::errors::internal::PrepareForStrCat(args)...));
+absl::Status Aborted(Args... args) {
+  return absl::Status(absl::StatusCode::kAborted,
+                      ::tsl::strings::StrCat(
+                          ::tsl::errors::internal::PrepareForStrCat(args)...));
 }
 template <typename... Args>
-::tsl::Status AbortedWithPayloads(
+absl::Status AbortedWithPayloads(
     const ::tsl::StringPiece& message,
     const std::unordered_map<std::string, std::string>& payloads) {
   return errors::Create(absl::StatusCode::kAborted, message, payloads);
@@ -491,13 +491,13 @@ ::tsl::Status AbortedWithPayloads(
 
 // DeadlineExceeded
 template <typename... Args>
-::tsl::Status DeadlineExceeded(Args... args) {
-  return ::tsl::Status(absl::StatusCode::kDeadlineExceeded,
-                       ::tsl::strings::StrCat(
-                           ::tsl::errors::internal::PrepareForStrCat(args)...));
+absl::Status DeadlineExceeded(Args... args) {
+  return absl::Status(absl::StatusCode::kDeadlineExceeded,
+                      ::tsl::strings::StrCat(
+                          ::tsl::errors::internal::PrepareForStrCat(args)...));
 }
 template <typename... Args>
-::tsl::Status DeadlineExceededWithPayloads(
+absl::Status DeadlineExceededWithPayloads(
     const ::tsl::StringPiece& message,
     const std::unordered_map<std::string, std::string>& payloads) {
   return errors::Create(absl::StatusCode::kDeadlineExceeded, message, payloads);
@@ -505,13 +505,13 @@ ::tsl::Status DeadlineExceededWithPayloads(
 
 // DataLoss
 template <typename... Args>
-::tsl::Status DataLoss(Args... args) {
-  return ::tsl::Status(absl::StatusCode::kDataLoss,
-                       ::tsl::strings::StrCat(
-                           ::tsl::errors::internal::PrepareForStrCat(args)...));
+absl::Status DataLoss(Args... args) {
+  return absl::Status(absl::StatusCode::kDataLoss,
+                      ::tsl::strings::StrCat(
+                          ::tsl::errors::internal::PrepareForStrCat(args)...));
 }
 template <typename... Args>
-::tsl::Status DataLossWithPayloads(
+absl::Status DataLossWithPayloads(
     const ::tsl::StringPiece& message,
     const std::unordered_map<std::string, std::string>& payloads) {
   return errors::Create(absl::StatusCode::kDataLoss, message, payloads);
@@ -519,26 +519,26 @@ ::tsl::Status DataLossWithPayloads(
 
 // Unknown
 template <typename... Args>
-::tsl::Status Unknown(Args... args) {
-  return ::tsl::Status(absl::StatusCode::kUnknown,
-                       ::tsl::strings::StrCat(
-                           ::tsl::errors::internal::PrepareForStrCat(args)...));
+absl::Status Unknown(Args... args) {
+  return absl::Status(absl::StatusCode::kUnknown,
+                      ::tsl::strings::StrCat(
+                          ::tsl::errors::internal::PrepareForStrCat(args)...));
 }
 template <typename... Args>
-::tsl::Status UnknownPayloads(
+absl::Status UnknownPayloads(
     const ::tsl::StringPiece& message,
     const std::unordered_map<std::string, std::string>& payloads) {
   return errors::Create(absl::StatusCode::kUnknown, message, payloads);
 }
 // PermissionDenied
 template <typename... Args>
-::tsl::Status PermissionDenied(Args... args) {
-  return ::tsl::Status(absl::StatusCode::kPermissionDenied,
-                       ::tsl::strings::StrCat(
-                           ::tsl::errors::internal::PrepareForStrCat(args)...));
+absl::Status PermissionDenied(Args... args) {
+  return absl::Status(absl::StatusCode::kPermissionDenied,
+                      ::tsl::strings::StrCat(
+                          ::tsl::errors::internal::PrepareForStrCat(args)...));
 }
 template <typename... Args>
-::tsl::Status PermissionDeniedWithPayloads(
+absl::Status PermissionDeniedWithPayloads(
     const ::tsl::StringPiece& message,
     const std::unordered_map<std::string, std::string>& payloads) {
   return errors::Create(absl::StatusCode::kPermissionDenied, message, payloads);
@@ -546,34 +546,34 @@ ::tsl::Status PermissionDeniedWithPayloads(
 
 // Unauthenticated
 template <typename... Args>
-::tsl::Status Unauthenticated(Args... args) {
-  return ::tsl::Status(absl::StatusCode::kUnauthenticated,
-                       ::tsl::strings::StrCat(
-                           ::tsl::errors::internal::PrepareForStrCat(args)...));
+absl::Status Unauthenticated(Args... args) {
+  return absl::Status(absl::StatusCode::kUnauthenticated,
+                      ::tsl::strings::StrCat(
+                          ::tsl::errors::internal::PrepareForStrCat(args)...));
 }
 template <typename... Args>
-::tsl::Status UnauthenticatedWithPayloads(
+absl::Status UnauthenticatedWithPayloads(
     const ::tsl::StringPiece& message,
     const std::unordered_map<std::string, std::string>& payloads) {
   return errors::Create(absl::StatusCode::kUnauthenticated, message, payloads);
 }
 
-bool IsAborted(const Status& status);
-bool IsAlreadyExists(const Status& status);
-bool IsCancelled(const Status& status);
-bool IsDataLoss(const Status& status);
-bool IsDeadlineExceeded(const Status& status);
-bool IsFailedPrecondition(const Status& status);
-bool IsInternal(const Status& status);
-bool IsInvalidArgument(const Status& status);
-bool IsNotFound(const Status& status);
-bool IsOutOfRange(const Status& status);
-bool IsPermissionDenied(const Status& status);
-bool IsResourceExhausted(const Status& status);
-bool IsUnauthenticated(const Status& status);
-bool IsUnavailable(const Status& status);
-bool IsUnimplemented(const Status& status);
-bool IsUnknown(const Status& status);
+bool IsAborted(const absl::Status& status);
+bool IsAlreadyExists(const absl::Status& status);
+bool IsCancelled(const absl::Status& status);
+bool IsDataLoss(const absl::Status& status);
+bool IsDeadlineExceeded(const absl::Status& status);
+bool IsFailedPrecondition(const absl::Status& status);
+bool IsInternal(const absl::Status& status);
+bool IsInvalidArgument(const absl::Status& status);
+bool IsNotFound(const absl::Status& status);
+bool IsOutOfRange(const absl::Status& status);
+bool IsPermissionDenied(const absl::Status& status);
+bool IsResourceExhausted(const absl::Status& status);
+bool IsUnauthenticated(const absl::Status& status);
+bool IsUnavailable(const absl::Status& status);
+bool IsUnimplemented(const absl::Status& status);
+bool IsUnknown(const absl::Status& status);
 
 // Produces a formatted string pattern from the name which can uniquely identify
 // this node upstream to produce an informative error message. The pattern
@@ -610,10 +610,10 @@ inline std::string FormatFunctionForError(absl::string_view name) {
   return strings::StrCat("{{function_node ", name, "}}");
 }
 
-inline Status ReplaceErrorFromNonCommunicationOps(const Status s,
-                                                  absl::string_view op_name) {
+inline absl::Status ReplaceErrorFromNonCommunicationOps(
+    const absl::Status s, absl::string_view op_name) {
   assert(::tsl::errors::IsUnavailable(s));
-  return Status(
+  return absl::Status(
       absl::StatusCode::kInternal,
       strings::StrCat(
           s.message(), "\nExecuting non-communication op <", op_name,
diff --git a/third_party/xla/third_party/tsl/tsl/platform/errors_test.cc b/third_party/xla/third_party/tsl/tsl/platform/errors_test.cc
index 3607b90c2d7b3e..88a3a5a78f72a5 100644
--- a/third_party/xla/third_party/tsl/tsl/platform/errors_test.cc
+++ b/third_party/xla/third_party/tsl/tsl/platform/errors_test.cc
@@ -21,7 +21,7 @@ limitations under the License.
 namespace tsl {
 
 TEST(AppendToMessageTest, PayloadsAreCopied) {
-  Status status = errors::Aborted("Aborted Error Message");
+  absl::Status status = errors::Aborted("Aborted Error Message");
   status.SetPayload("payload_key", absl::Cord("payload_value"));
   errors::AppendToMessage(&status, "Appended Message");
 
@@ -30,22 +30,22 @@ TEST(AppendToMessageTest, PayloadsAreCopied) {
 }
 
 TEST(Status, GetAllPayloads) {
-  Status s_error(absl::StatusCode::kInternal, "Error message");
+  absl::Status s_error(absl::StatusCode::kInternal, "Error message");
   s_error.SetPayload("Error key", absl::Cord("foo"));
   auto payloads_error_status = errors::GetPayloads(s_error);
   ASSERT_EQ(payloads_error_status.size(), 1);
   ASSERT_EQ(payloads_error_status["Error key"], "foo");
 
-  Status s_ok = Status();
+  absl::Status s_ok = absl::Status();
   auto payloads_ok_status = errors::GetPayloads(s_ok);
   ASSERT_TRUE(payloads_ok_status.empty());
 }
 
 TEST(Status, OKStatusInsertPayloadsFromErrorStatus) {
   // An OK status will should not change after InsertPayloads() calls.
-  Status s_error(absl::StatusCode::kInternal, "Error message");
+  absl::Status s_error(absl::StatusCode::kInternal, "Error message");
   s_error.SetPayload("Error key", absl::Cord("foo"));
-  Status s_ok = Status();
+  absl::Status s_ok = absl::Status();
 
   errors::InsertPayloads(s_ok, errors::GetPayloads(s_error));
   auto payloads_ok_status = errors::GetPayloads(s_ok);
@@ -54,19 +54,19 @@ TEST(Status, OKStatusInsertPayloadsFromErrorStatus) {
 
 TEST(Status, ErrorStatusInsertPayloadsFromOKStatus) {
   // An InsertPayloads() call should not take effect from empty inputs.
-  Status s_error(absl::StatusCode::kInternal, "Error message");
+  absl::Status s_error(absl::StatusCode::kInternal, "Error message");
   s_error.SetPayload("Error key", absl::Cord("foo"));
-  Status s_ok = Status();
+  absl::Status s_ok = absl::Status();
 
   errors::InsertPayloads(s_error, errors::GetPayloads(s_ok));
   ASSERT_EQ(s_error.GetPayload("Error key"), "foo");
 }
 
 TEST(Status, ErrorStatusInsertPayloadsFromErrorStatus) {
-  Status s_error1(absl::StatusCode::kInternal, "Error message");
+  absl::Status s_error1(absl::StatusCode::kInternal, "Error message");
   s_error1.SetPayload("Error key 1", absl::Cord("foo"));
   s_error1.SetPayload("Error key 2", absl::Cord("bar"));
-  Status s_error2(absl::StatusCode::kInternal, "Error message");
+  absl::Status s_error2(absl::StatusCode::kInternal, "Error message");
   s_error2.SetPayload("Error key", absl::Cord("bar"));
   ASSERT_EQ(s_error2.GetPayload("Error key"), "bar");
 
@@ -79,22 +79,22 @@ TEST(Status, ErrorStatusInsertPayloadsFromErrorStatus) {
 
 #if defined(PLATFORM_GOOGLE)
 
-Status GetError() {
+absl::Status GetError() {
   return absl::InvalidArgumentError("An invalid argument error");
 }
 
-Status PropagateError() {
+absl::Status PropagateError() {
   TF_RETURN_IF_ERROR(GetError());
   return absl::OkStatus();
 }
 
-Status PropagateError2() {
+absl::Status PropagateError2() {
   TF_RETURN_IF_ERROR(PropagateError());
   return absl::OkStatus();
 }
 
 TEST(Status, StackTracePropagation) {
-  Status s = PropagateError2();
+  absl::Status s = PropagateError2();
   auto sources = s.GetSourceLocations();
   ASSERT_EQ(sources.size(), 3);
 
@@ -105,16 +105,16 @@ TEST(Status, StackTracePropagation) {
 }
 
 TEST(Status, SourceLocationsPreservedByAppend) {
-  Status s = PropagateError2();
+  absl::Status s = PropagateError2();
   ASSERT_EQ(s.GetSourceLocations().size(), 3);
   errors::AppendToMessage(&s, "A new message.");
   ASSERT_EQ(s.GetSourceLocations().size(), 3);
 }
 
 TEST(Status, SourceLocationsPreservedByUpdate) {
-  Status s = PropagateError2();
+  absl::Status s = PropagateError2();
   ASSERT_EQ(s.GetSourceLocations().size(), 3);
-  Status s2 = errors::CreateWithUpdatedMessage(s, "New message.");
+  absl::Status s2 = errors::CreateWithUpdatedMessage(s, "New message.");
   ASSERT_EQ(s2.GetSourceLocations().size(), 3);
 }
 
diff --git a/third_party/xla/third_party/tsl/tsl/platform/file_system.cc b/third_party/xla/third_party/tsl/tsl/platform/file_system.cc
index 022782ee47c5e7..cbca921d50b545 100644
--- a/third_party/xla/third_party/tsl/tsl/platform/file_system.cc
+++ b/third_party/xla/third_party/tsl/tsl/platform/file_system.cc
@@ -74,31 +74,33 @@ string FileSystem::TranslateName(const string& name) const {
   return this->CleanPath(path);
 }
 
-Status FileSystem::IsDirectory(const string& name, TransactionToken* token) {
+absl::Status FileSystem::IsDirectory(const string& name,
+                                     TransactionToken* token) {
   // Check if path exists.
   // TODO(sami):Forward token to other methods once migration is complete.
   TF_RETURN_IF_ERROR(FileExists(name));
   FileStatistics stat;
   TF_RETURN_IF_ERROR(Stat(name, &stat));
   if (stat.is_directory) {
-    return OkStatus();
+    return absl::OkStatus();
   }
-  return Status(absl::StatusCode::kFailedPrecondition, "Not a directory");
+  return absl::Status(absl::StatusCode::kFailedPrecondition, "Not a directory");
 }
 
-Status FileSystem::HasAtomicMove(const string& path, bool* has_atomic_move) {
+absl::Status FileSystem::HasAtomicMove(const string& path,
+                                       bool* has_atomic_move) {
   *has_atomic_move = true;
-  return OkStatus();
+  return absl::OkStatus();
 }
 
 void FileSystem::FlushCaches(TransactionToken* token) {}
 
 bool FileSystem::FilesExist(const std::vector<string>& files,
                             TransactionToken* token,
-                            std::vector<Status>* status) {
+                            std::vector<absl::Status>* status) {
   bool result = true;
   for (const auto& file : files) {
-    Status s = FileExists(file);
+    absl::Status s = FileExists(file);
     result &= s.ok();
     if (status != nullptr) {
       status->push_back(s);
@@ -110,17 +112,17 @@ bool FileSystem::FilesExist(const std::vector<string>& files,
   return result;
 }
 
-Status FileSystem::DeleteRecursively(const string& dirname,
-                                     TransactionToken* token,
-                                     int64_t* undeleted_files,
-                                     int64_t* undeleted_dirs) {
+absl::Status FileSystem::DeleteRecursively(const string& dirname,
+                                           TransactionToken* token,
+                                           int64_t* undeleted_files,
+                                           int64_t* undeleted_dirs) {
   CHECK_NOTNULL(undeleted_files);
   CHECK_NOTNULL(undeleted_dirs);
 
   *undeleted_files = 0;
   *undeleted_dirs = 0;
   // Make sure that dirname exists;
-  Status exists_status = FileExists(dirname);
+  absl::Status exists_status = FileExists(dirname);
   if (!exists_status.ok()) {
     (*undeleted_dirs)++;
     return exists_status;
@@ -128,7 +130,7 @@ Status FileSystem::DeleteRecursively(const string& dirname,
 
   // If given path to a single file, we should just delete it.
   if (!IsDirectory(dirname).ok()) {
-    Status delete_root_status = DeleteFile(dirname);
+    absl::Status delete_root_status = DeleteFile(dirname);
     if (!delete_root_status.ok()) (*undeleted_files)++;
     return delete_root_status;
   }
@@ -136,7 +138,7 @@ Status FileSystem::DeleteRecursively(const string& dirname,
   std::deque<string> dir_q;      // Queue for the BFS
   std::vector<string> dir_list;  // List of all dirs discovered
   dir_q.push_back(dirname);
-  Status ret;  // Status to be returned.
+  absl::Status ret;  // Status to be returned.
   // Do a BFS on the directory to discover all the sub-directories. Remove all
   // children that are files along the way. Then cleanup and remove the
   // directories in reverse order.;
@@ -146,7 +148,7 @@ Status FileSystem::DeleteRecursively(const string& dirname,
     dir_list.push_back(dir);
     std::vector<string> children;
     // GetChildren might fail if we don't have appropriate permissions.
-    Status s = GetChildren(dir, &children);
+    absl::Status s = GetChildren(dir, &children);
     ret.Update(s);
     if (!s.ok()) {
       (*undeleted_dirs)++;
@@ -160,7 +162,7 @@ Status FileSystem::DeleteRecursively(const string& dirname,
       } else {
         // Delete file might fail because of permissions issues or might be
         // unimplemented.
-        Status del_status = DeleteFile(child_path);
+        absl::Status del_status = DeleteFile(child_path);
         ret.Update(del_status);
         if (!del_status.ok()) {
           (*undeleted_files)++;
@@ -174,7 +176,7 @@ Status FileSystem::DeleteRecursively(const string& dirname,
   for (const string& dir : dir_list) {
     // Delete dir might fail because of permissions issues or might be
     // unimplemented.
-    Status s = DeleteDir(dir);
+    absl::Status s = DeleteDir(dir);
     ret.Update(s);
     if (!s.ok()) {
       (*undeleted_dirs)++;
@@ -183,19 +185,19 @@ Status FileSystem::DeleteRecursively(const string& dirname,
   return ret;
 }
 
-Status FileSystem::RecursivelyCreateDir(const string& dirname,
-                                        TransactionToken* token) {
+absl::Status FileSystem::RecursivelyCreateDir(const string& dirname,
+                                              TransactionToken* token) {
   StringPiece scheme, host, remaining_dir;
   this->ParseURI(dirname, &scheme, &host, &remaining_dir);
   std::vector<StringPiece> sub_dirs;
   while (!remaining_dir.empty()) {
     std::string current_entry = this->CreateURI(scheme, host, remaining_dir);
-    Status exists_status = FileExists(current_entry);
+    absl::Status exists_status = FileExists(current_entry);
     if (exists_status.ok()) {
       // FileExists cannot differentiate between existence of a file or a
       // directory, hence we need an additional test as we must not assume that
       // a path to a file is a path to a parent directory.
-      Status directory_status = IsDirectory(current_entry);
+      absl::Status directory_status = IsDirectory(current_entry);
       if (directory_status.ok()) {
         break;  // We need to start creating directories from here.
       } else if (directory_status.code() == absl::StatusCode::kUnimplemented) {
@@ -221,16 +223,16 @@ Status FileSystem::RecursivelyCreateDir(const string& dirname,
   string built_path(remaining_dir);
   for (const StringPiece sub_dir : sub_dirs) {
     built_path = this->JoinPath(built_path, sub_dir);
-    Status status = CreateDir(this->CreateURI(scheme, host, built_path));
+    absl::Status status = CreateDir(this->CreateURI(scheme, host, built_path));
     if (!status.ok() && status.code() != absl::StatusCode::kAlreadyExists) {
       return status;
     }
   }
-  return OkStatus();
+  return absl::OkStatus();
 }
 
-Status FileSystem::CopyFile(const string& src, const string& target,
-                            TransactionToken* token) {
+absl::Status FileSystem::CopyFile(const string& src, const string& target,
+                                  TransactionToken* token) {
   return FileSystemCopyFile(this, src, this, target);
 }
 
diff --git a/third_party/xla/third_party/tsl/tsl/platform/file_system.h b/third_party/xla/third_party/tsl/tsl/platform/file_system.h
index 8f7bd875e35bc3..a25a4760357ca6 100644
--- a/third_party/xla/third_party/tsl/tsl/platform/file_system.h
+++ b/third_party/xla/third_party/tsl/tsl/platform/file_system.h
@@ -69,17 +69,17 @@ class FileSystem {
   ///
   /// The ownership of the returned RandomAccessFile is passed to the caller
   /// and the object should be deleted when is not used.
-  virtual tsl::Status NewRandomAccessFile(
+  virtual absl::Status NewRandomAccessFile(
       const std::string& fname, std::unique_ptr<RandomAccessFile>* result) {
     return NewRandomAccessFile(fname, nullptr, result);
   }
 
-  virtual tsl::Status NewRandomAccessFile(
+  virtual absl::Status NewRandomAccessFile(
       const std::string& fname, TransactionToken* token,
       std::unique_ptr<RandomAccessFile>* result) {
     // We duplicate these methods due to Google internal coding style prevents
     // virtual functions with default arguments. See PR #41615.
-    return OkStatus();
+    return absl::OkStatus();
   }
 
   /// \brief Creates an object that writes to a new file with the specified
@@ -94,15 +94,15 @@ class FileSystem {
   ///
   /// The ownership of the returned WritableFile is passed to the caller
   /// and the object should be deleted when is not used.
-  virtual tsl::Status NewWritableFile(const std::string& fname,
-                                      std::unique_ptr<WritableFile>* result) {
+  virtual absl::Status NewWritableFile(const std::string& fname,
+                                       std::unique_ptr<WritableFile>* result) {
     return NewWritableFile(fname, nullptr, result);
   }
 
-  virtual tsl::Status NewWritableFile(const std::string& fname,
-                                      TransactionToken* token,
-                                      std::unique_ptr<WritableFile>* result) {
-    return OkStatus();
+  virtual absl::Status NewWritableFile(const std::string& fname,
+                                       TransactionToken* token,
+                                       std::unique_ptr<WritableFile>* result) {
+    return absl::OkStatus();
   }
 
   /// \brief Creates an object that either appends to an existing file, or
@@ -116,15 +116,15 @@ class FileSystem {
   ///
   /// The ownership of the returned WritableFile is passed to the caller
   /// and the object should be deleted when is not used.
-  virtual tsl::Status NewAppendableFile(const std::string& fname,
-                                        std::unique_ptr<WritableFile>* result) {
+  virtual absl::Status NewAppendableFile(
+      const std::string& fname, std::unique_ptr<WritableFile>* result) {
     return NewAppendableFile(fname, nullptr, result);
   }
 
-  virtual tsl::Status NewAppendableFile(const std::string& fname,
-                                        TransactionToken* token,
-                                        std::unique_ptr<WritableFile>* result) {
-    return OkStatus();
+  virtual absl::Status NewAppendableFile(
+      const std::string& fname, TransactionToken* token,
+      std::unique_ptr<WritableFile>* result) {
+    return absl::OkStatus();
   }
 
   /// \brief Creates a readonly region of memory with the file context.
@@ -137,50 +137,51 @@ class FileSystem {
   ///
   /// The ownership of the returned ReadOnlyMemoryRegion is passed to the caller
   /// and the object should be deleted when is not used.
-  virtual tsl::Status NewReadOnlyMemoryRegionFromFile(
+  virtual absl::Status NewReadOnlyMemoryRegionFromFile(
       const std::string& fname, std::unique_ptr<ReadOnlyMemoryRegion>* result) {
     return NewReadOnlyMemoryRegionFromFile(fname, nullptr, result);
   }
 
-  virtual tsl::Status NewReadOnlyMemoryRegionFromFile(
+  virtual absl::Status NewReadOnlyMemoryRegionFromFile(
       const std::string& fname, TransactionToken* token,
       std::unique_ptr<ReadOnlyMemoryRegion>* result) {
-    return OkStatus();
+    return absl::OkStatus();
   }
 
   /// Returns OK if the named path exists and NOT_FOUND otherwise.
-  virtual tsl::Status FileExists(const std::string& fname) {
+  virtual absl::Status FileExists(const std::string& fname) {
     return FileExists(fname, nullptr);
   }
 
-  virtual tsl::Status FileExists(const std::string& fname,
-                                 TransactionToken* token) {
-    return OkStatus();
+  virtual absl::Status FileExists(const std::string& fname,
+                                  TransactionToken* token) {
+    return absl::OkStatus();
   }
 
   /// Returns true if all the listed files exist, false otherwise.
   /// if status is not null, populate the vector with a detailed status
   /// for each file.
   virtual bool FilesExist(const std::vector<string>& files,
-                          std::vector<Status>* status) {
+                          std::vector<absl::Status>* status) {
     return FilesExist(files, nullptr, status);
   }
 
   virtual bool FilesExist(const std::vector<string>& files,
-                          TransactionToken* token, std::vector<Status>* status);
+                          TransactionToken* token,
+                          std::vector<absl::Status>* status);
 
   /// \brief Returns the immediate children in the given directory.
   ///
   /// The returned paths are relative to 'dir'.
-  virtual tsl::Status GetChildren(const std::string& dir,
-                                  std::vector<string>* result) {
+  virtual absl::Status GetChildren(const std::string& dir,
+                                   std::vector<string>* result) {
     return GetChildren(dir, nullptr, result);
   }
 
-  virtual tsl::Status GetChildren(const std::string& dir,
-                                  TransactionToken* token,
-                                  std::vector<string>* result) {
-    return OkStatus();
+  virtual absl::Status GetChildren(const std::string& dir,
+                                   TransactionToken* token,
+                                   std::vector<string>* result) {
+    return absl::OkStatus();
   }
 
   /// \brief Given a pattern, stores in *results the set of paths that matches
@@ -205,15 +206,15 @@ class FileSystem {
   ///  * OK - no errors
   ///  * UNIMPLEMENTED - Some underlying functions (like GetChildren) are not
   ///                    implemented
-  virtual tsl::Status GetMatchingPaths(const std::string& pattern,
-                                       std::vector<string>* results) {
+  virtual absl::Status GetMatchingPaths(const std::string& pattern,
+                                        std::vector<string>* results) {
     return GetMatchingPaths(pattern, nullptr, results);
   }
 
-  virtual tsl::Status GetMatchingPaths(const std::string& pattern,
-                                       TransactionToken* token,
-                                       std::vector<string>* results) {
-    return OkStatus();
+  virtual absl::Status GetMatchingPaths(const std::string& pattern,
+                                        TransactionToken* token,
+                                        std::vector<string>* results) {
+    return absl::OkStatus();
   }
 
   /// \brief Checks if the given filename matches the pattern.
@@ -224,23 +225,23 @@ class FileSystem {
   virtual bool Match(const std::string& filename, const std::string& pattern);
 
   /// \brief Obtains statistics for the given path.
-  virtual tsl::Status Stat(const std::string& fname, FileStatistics* stat) {
+  virtual absl::Status Stat(const std::string& fname, FileStatistics* stat) {
     return Stat(fname, nullptr, stat);
   }
 
-  virtual tsl::Status Stat(const std::string& fname, TransactionToken* token,
-                           FileStatistics* stat) {
-    return OkStatus();
+  virtual absl::Status Stat(const std::string& fname, TransactionToken* token,
+                            FileStatistics* stat) {
+    return absl::OkStatus();
   }
 
   /// \brief Deletes the named file.
-  virtual tsl::Status DeleteFile(const std::string& fname) {
+  virtual absl::Status DeleteFile(const std::string& fname) {
     return DeleteFile(fname, nullptr);
   }
 
-  virtual tsl::Status DeleteFile(const std::string& fname,
-                                 TransactionToken* token) {
-    return OkStatus();
+  virtual absl::Status DeleteFile(const std::string& fname,
+                                  TransactionToken* token) {
+    return absl::OkStatus();
   }
 
   /// \brief Creates the specified directory.
@@ -248,13 +249,13 @@ class FileSystem {
   ///  * OK - successfully created the directory.
   ///  * ALREADY_EXISTS - directory with name dirname already exists.
   ///  * PERMISSION_DENIED - dirname is not writable.
-  virtual tsl::Status CreateDir(const std::string& dirname) {
+  virtual absl::Status CreateDir(const std::string& dirname) {
     return CreateDir(dirname, nullptr);
   }
 
-  virtual tsl::Status CreateDir(const std::string& dirname,
-                                TransactionToken* token) {
-    return OkStatus();
+  virtual absl::Status CreateDir(const std::string& dirname,
+                                 TransactionToken* token) {
+    return absl::OkStatus();
   }
 
   /// \brief Creates the specified directory and all the necessary
@@ -263,21 +264,21 @@ class FileSystem {
   ///  * OK - successfully created the directory and sub directories, even if
   ///         they were already created.
   ///  * PERMISSION_DENIED - dirname or some subdirectory is not writable.
-  virtual tsl::Status RecursivelyCreateDir(const std::string& dirname) {
+  virtual absl::Status RecursivelyCreateDir(const std::string& dirname) {
     return RecursivelyCreateDir(dirname, nullptr);
   }
 
-  virtual tsl::Status RecursivelyCreateDir(const std::string& dirname,
-                                           TransactionToken* token);
+  virtual absl::Status RecursivelyCreateDir(const std::string& dirname,
+                                            TransactionToken* token);
 
   /// \brief Deletes the specified directory.
-  virtual tsl::Status DeleteDir(const std::string& dirname) {
+  virtual absl::Status DeleteDir(const std::string& dirname) {
     return DeleteDir(dirname, nullptr);
   }
 
-  virtual tsl::Status DeleteDir(const std::string& dirname,
-                                TransactionToken* token) {
-    return OkStatus();
+  virtual absl::Status DeleteDir(const std::string& dirname,
+                                 TransactionToken* token) {
+    return absl::OkStatus();
   }
 
   /// \brief Deletes the specified directory and all subdirectories and files
@@ -304,48 +305,49 @@ class FileSystem {
   ///  * PERMISSION_DENIED - dirname or some descendant is not writable
   ///  * UNIMPLEMENTED - Some underlying functions (like Delete) are not
   ///                    implemented
-  virtual tsl::Status DeleteRecursively(const std::string& dirname,
-                                        int64_t* undeleted_files,
-                                        int64_t* undeleted_dirs) {
+  virtual absl::Status DeleteRecursively(const std::string& dirname,
+                                         int64_t* undeleted_files,
+                                         int64_t* undeleted_dirs) {
     return DeleteRecursively(dirname, nullptr, undeleted_files, undeleted_dirs);
   }
 
-  virtual tsl::Status DeleteRecursively(const std::string& dirname,
-                                        TransactionToken* token,
-                                        int64_t* undeleted_files,
-                                        int64_t* undeleted_dirs);
+  virtual absl::Status DeleteRecursively(const std::string& dirname,
+                                         TransactionToken* token,
+                                         int64_t* undeleted_files,
+                                         int64_t* undeleted_dirs);
 
   /// \brief Stores the size of `fname` in `*file_size`.
-  virtual tsl::Status GetFileSize(const std::string& fname, uint64* file_size) {
+  virtual absl::Status GetFileSize(const std::string& fname,
+                                   uint64* file_size) {
     return GetFileSize(fname, nullptr, file_size);
   }
 
-  virtual tsl::Status GetFileSize(const std::string& fname,
-                                  TransactionToken* token, uint64* file_size) {
-    return OkStatus();
+  virtual absl::Status GetFileSize(const std::string& fname,
+                                   TransactionToken* token, uint64* file_size) {
+    return absl::OkStatus();
   }
 
   /// \brief Overwrites the target if it exists.
-  virtual tsl::Status RenameFile(const std::string& src,
-                                 const std::string& target) {
+  virtual absl::Status RenameFile(const std::string& src,
+                                  const std::string& target) {
     return RenameFile(src, target, nullptr);
   }
 
-  virtual tsl::Status RenameFile(const std::string& src,
-                                 const std::string& target,
-                                 TransactionToken* token) {
-    return OkStatus();
+  virtual absl::Status RenameFile(const std::string& src,
+                                  const std::string& target,
+                                  TransactionToken* token) {
+    return absl::OkStatus();
   }
 
   /// \brief Copy the src to target.
-  virtual tsl::Status CopyFile(const std::string& src,
-                               const std::string& target) {
+  virtual absl::Status CopyFile(const std::string& src,
+                                const std::string& target) {
     return CopyFile(src, target, nullptr);
   }
 
-  virtual tsl::Status CopyFile(const std::string& src,
-                               const std::string& target,
-                               TransactionToken* token);
+  virtual absl::Status CopyFile(const std::string& src,
+                                const std::string& target,
+                                TransactionToken* token);
 
   /// \brief Translate an URI to a filename for the FileSystem implementation.
   ///
@@ -365,12 +367,12 @@ class FileSystem {
   ///  * NOT_FOUND - The path entry does not exist.
   ///  * PERMISSION_DENIED - Insufficient permissions.
   ///  * UNIMPLEMENTED - The file factory doesn't support directories.
-  virtual tsl::Status IsDirectory(const std::string& fname) {
+  virtual absl::Status IsDirectory(const std::string& fname) {
     return IsDirectory(fname, nullptr);
   }
 
-  virtual tsl::Status IsDirectory(const std::string& fname,
-                                  TransactionToken* token);
+  virtual absl::Status IsDirectory(const std::string& fname,
+                                   TransactionToken* token);
 
   /// \brief Returns whether the given path is on a file system
   /// that has atomic move capabilities. This can be used
@@ -382,7 +384,8 @@ class FileSystem {
   ///         so has_atomic_move holds the above information.
   ///  * UNIMPLEMENTED - The file system of the path hasn't been implemented in
   ///  TF
-  virtual Status HasAtomicMove(const std::string& path, bool* has_atomic_move);
+  virtual absl::Status HasAtomicMove(const std::string& path,
+                                     bool* has_atomic_move);
 
   /// \brief Flushes any cached filesystem objects from memory.
   virtual void FlushCaches() { FlushCaches(nullptr); }
@@ -476,67 +479,67 @@ class FileSystem {
   // Transaction related API
 
   /// \brief Starts a new transaction
-  virtual tsl::Status StartTransaction(TransactionToken** token) {
+  virtual absl::Status StartTransaction(TransactionToken** token) {
     *token = nullptr;
-    return OkStatus();
+    return absl::OkStatus();
   }
 
   /// \brief Adds `path` to transaction in `token`
-  virtual tsl::Status AddToTransaction(const std::string& path,
-                                       TransactionToken* token) {
-    return OkStatus();
+  virtual absl::Status AddToTransaction(const std::string& path,
+                                        TransactionToken* token) {
+    return absl::OkStatus();
   }
 
   /// \brief Ends transaction
-  virtual tsl::Status EndTransaction(TransactionToken* token) {
-    return OkStatus();
+  virtual absl::Status EndTransaction(TransactionToken* token) {
+    return absl::OkStatus();
   }
 
   /// \brief Get token for `path` or start a new transaction and add `path` to
   /// it.
-  virtual tsl::Status GetTokenOrStartTransaction(const std::string& path,
-                                                 TransactionToken** token) {
+  virtual absl::Status GetTokenOrStartTransaction(const std::string& path,
+                                                  TransactionToken** token) {
     *token = nullptr;
-    return OkStatus();
+    return absl::OkStatus();
   }
 
   /// \brief Return transaction for `path` or nullptr in `token`
-  virtual tsl::Status GetTransactionForPath(const std::string& path,
-                                            TransactionToken** token) {
+  virtual absl::Status GetTransactionForPath(const std::string& path,
+                                             TransactionToken** token) {
     *token = nullptr;
-    return OkStatus();
+    return absl::OkStatus();
   }
 
   /// \brief Decode transaction to human readable string.
   virtual std::string DecodeTransaction(const TransactionToken* token);
 
   /// \brief Set File System Configuration Options
-  virtual Status SetOption(const string& key, const string& value) {
+  virtual absl::Status SetOption(const string& key, const string& value) {
     return errors::Unimplemented("SetOption");
   }
 
   /// \brief Set File System Configuration Option
-  virtual tsl::Status SetOption(const std::string& name,
-                                const std::vector<string>& values) {
+  virtual absl::Status SetOption(const std::string& name,
+                                 const std::vector<string>& values) {
     return errors::Unimplemented("SetOption");
   }
 
   /// \brief Set File System Configuration Option
-  virtual tsl::Status SetOption(const std::string& name,
-                                const std::vector<int64_t>& values) {
+  virtual absl::Status SetOption(const std::string& name,
+                                 const std::vector<int64_t>& values) {
     return errors::Unimplemented("SetOption");
   }
 
   /// \brief Set File System Configuration Option
-  virtual tsl::Status SetOption(const std::string& name,
-                                const std::vector<double>& values) {
+  virtual absl::Status SetOption(const std::string& name,
+                                 const std::vector<double>& values) {
     return errors::Unimplemented("SetOption");
   }
 
   /// \brief Set File System ACL checker.
   ///
   /// No checks are enforced if a FileAcl is never set.
-  virtual tsl::Status SetFileAcl(std::shared_ptr<FileAcl> file_acl) {
+  virtual absl::Status SetFileAcl(std::shared_ptr<FileAcl> file_acl) {
     return errors::Unimplemented("SetFileAcl");
   }
 
@@ -581,48 +584,49 @@ class WrappedFileSystem : public FileSystem {
  public:
   TF_USE_FILESYSTEM_METHODS_WITH_NO_TRANSACTION_SUPPORT;
 
-  tsl::Status NewRandomAccessFile(
+  absl::Status NewRandomAccessFile(
       const std::string& fname, TransactionToken* token,
       std::unique_ptr<RandomAccessFile>* result) override {
     return fs_->NewRandomAccessFile(fname, (token ? token : token_), result);
   }
 
-  tsl::Status NewWritableFile(const std::string& fname, TransactionToken* token,
-                              std::unique_ptr<WritableFile>* result) override {
+  absl::Status NewWritableFile(const std::string& fname,
+                               TransactionToken* token,
+                               std::unique_ptr<WritableFile>* result) override {
     return fs_->NewWritableFile(fname, (token ? token : token_), result);
   }
 
-  tsl::Status NewAppendableFile(
+  absl::Status NewAppendableFile(
       const std::string& fname, TransactionToken* token,
       std::unique_ptr<WritableFile>* result) override {
     return fs_->NewAppendableFile(fname, (token ? token : token_), result);
   }
 
-  tsl::Status NewReadOnlyMemoryRegionFromFile(
+  absl::Status NewReadOnlyMemoryRegionFromFile(
       const std::string& fname, TransactionToken* token,
       std::unique_ptr<ReadOnlyMemoryRegion>* result) override {
     return fs_->NewReadOnlyMemoryRegionFromFile(fname, (token ? token : token_),
                                                 result);
   }
 
-  tsl::Status FileExists(const std::string& fname,
-                         TransactionToken* token) override {
+  absl::Status FileExists(const std::string& fname,
+                          TransactionToken* token) override {
     return fs_->FileExists(fname, (token ? token : token_));
   }
 
   bool FilesExist(const std::vector<string>& files, TransactionToken* token,
-                  std::vector<Status>* status) override {
+                  std::vector<absl::Status>* status) override {
     return fs_->FilesExist(files, (token ? token : token_), status);
   }
 
-  tsl::Status GetChildren(const std::string& dir, TransactionToken* token,
-                          std::vector<string>* result) override {
+  absl::Status GetChildren(const std::string& dir, TransactionToken* token,
+                           std::vector<string>* result) override {
     return fs_->GetChildren(dir, (token ? token : token_), result);
   }
 
-  tsl::Status GetMatchingPaths(const std::string& pattern,
-                               TransactionToken* token,
-                               std::vector<string>* results) override {
+  absl::Status GetMatchingPaths(const std::string& pattern,
+                                TransactionToken* token,
+                                std::vector<string>* results) override {
     return fs_->GetMatchingPaths(pattern, (token ? token : token_), results);
   }
 
@@ -630,51 +634,51 @@ class WrappedFileSystem : public FileSystem {
     return fs_->Match(filename, pattern);
   }
 
-  tsl::Status Stat(const std::string& fname, TransactionToken* token,
-                   FileStatistics* stat) override {
+  absl::Status Stat(const std::string& fname, TransactionToken* token,
+                    FileStatistics* stat) override {
     return fs_->Stat(fname, (token ? token : token_), stat);
   }
 
-  tsl::Status DeleteFile(const std::string& fname,
-                         TransactionToken* token) override {
+  absl::Status DeleteFile(const std::string& fname,
+                          TransactionToken* token) override {
     return fs_->DeleteFile(fname, (token ? token : token_));
   }
 
-  tsl::Status CreateDir(const std::string& dirname,
-                        TransactionToken* token) override {
+  absl::Status CreateDir(const std::string& dirname,
+                         TransactionToken* token) override {
     return fs_->CreateDir(dirname, (token ? token : token_));
   }
 
-  tsl::Status RecursivelyCreateDir(const std::string& dirname,
-                                   TransactionToken* token) override {
+  absl::Status RecursivelyCreateDir(const std::string& dirname,
+                                    TransactionToken* token) override {
     return fs_->RecursivelyCreateDir(dirname, (token ? token : token_));
   }
 
-  tsl::Status DeleteDir(const std::string& dirname,
-                        TransactionToken* token) override {
+  absl::Status DeleteDir(const std::string& dirname,
+                         TransactionToken* token) override {
     return fs_->DeleteDir(dirname, (token ? token : token_));
   }
 
-  tsl::Status DeleteRecursively(const std::string& dirname,
-                                TransactionToken* token,
-                                int64_t* undeleted_files,
-                                int64_t* undeleted_dirs) override {
+  absl::Status DeleteRecursively(const std::string& dirname,
+                                 TransactionToken* token,
+                                 int64_t* undeleted_files,
+                                 int64_t* undeleted_dirs) override {
     return fs_->DeleteRecursively(dirname, (token ? token : token_),
                                   undeleted_files, undeleted_dirs);
   }
 
-  tsl::Status GetFileSize(const std::string& fname, TransactionToken* token,
-                          uint64* file_size) override {
+  absl::Status GetFileSize(const std::string& fname, TransactionToken* token,
+                           uint64* file_size) override {
     return fs_->GetFileSize(fname, (token ? token : token_), file_size);
   }
 
-  tsl::Status RenameFile(const std::string& src, const std::string& target,
-                         TransactionToken* token) override {
+  absl::Status RenameFile(const std::string& src, const std::string& target,
+                          TransactionToken* token) override {
     return fs_->RenameFile(src, target, (token ? token : token_));
   }
 
-  tsl::Status CopyFile(const std::string& src, const std::string& target,
-                       TransactionToken* token) override {
+  absl::Status CopyFile(const std::string& src, const std::string& target,
+                        TransactionToken* token) override {
     return fs_->CopyFile(src, target, (token ? token : token_));
   }
 
@@ -682,13 +686,13 @@ class WrappedFileSystem : public FileSystem {
     return fs_->TranslateName(name);
   }
 
-  tsl::Status IsDirectory(const std::string& fname,
-                          TransactionToken* token) override {
+  absl::Status IsDirectory(const std::string& fname,
+                           TransactionToken* token) override {
     return fs_->IsDirectory(fname, (token ? token : token_));
   }
 
-  Status HasAtomicMove(const std::string& path,
-                       bool* has_atomic_move) override {
+  absl::Status HasAtomicMove(const std::string& path,
+                             bool* has_atomic_move) override {
     return fs_->HasAtomicMove(path, has_atomic_move);
   }
 
@@ -702,26 +706,26 @@ class WrappedFileSystem : public FileSystem {
     return fs_->Basename(path);
   }
 
-  tsl::Status StartTransaction(TransactionToken** token) override {
+  absl::Status StartTransaction(TransactionToken** token) override {
     return fs_->StartTransaction(token);
   }
 
-  tsl::Status AddToTransaction(const std::string& path,
-                               TransactionToken* token) override {
+  absl::Status AddToTransaction(const std::string& path,
+                                TransactionToken* token) override {
     return fs_->AddToTransaction(path, (token ? token : token_));
   }
 
-  tsl::Status EndTransaction(TransactionToken* token) override {
+  absl::Status EndTransaction(TransactionToken* token) override {
     return fs_->EndTransaction(token);
   }
 
-  tsl::Status GetTransactionForPath(const std::string& path,
-                                    TransactionToken** token) override {
+  absl::Status GetTransactionForPath(const std::string& path,
+                                     TransactionToken** token) override {
     return fs_->GetTransactionForPath(path, token);
   }
 
-  tsl::Status GetTokenOrStartTransaction(const std::string& path,
-                                         TransactionToken** token) override {
+  absl::Status GetTokenOrStartTransaction(const std::string& path,
+                                          TransactionToken** token) override {
     return fs_->GetTokenOrStartTransaction(path, token);
   }
 
@@ -749,7 +753,7 @@ class RandomAccessFile {
   ///
   /// This is an optional operation that may not be implemented by every
   /// filesystem.
-  virtual tsl::Status Name(StringPiece* result) const {
+  virtual absl::Status Name(StringPiece* result) const {
     return errors::Unimplemented("This filesystem does not support Name()");
   }
 
@@ -768,12 +772,12 @@ class RandomAccessFile {
   /// because of EOF.
   ///
   /// Safe for concurrent use by multiple threads.
-  virtual tsl::Status Read(uint64 offset, size_t n, StringPiece* result,
-                           char* scratch) const = 0;
+  virtual absl::Status Read(uint64 offset, size_t n, StringPiece* result,
+                            char* scratch) const = 0;
 
 #if defined(TF_CORD_SUPPORT)
   /// \brief Read up to `n` bytes from the file starting at `offset`.
-  virtual tsl::Status Read(uint64 offset, size_t n, absl::Cord* cord) const {
+  virtual absl::Status Read(uint64 offset, size_t n, absl::Cord* cord) const {
     return errors::Unimplemented(
         "Read(uint64, size_t, absl::Cord*) is not "
         "implemented");
@@ -795,15 +799,15 @@ class WritableFile {
   virtual ~WritableFile() = default;
 
   /// \brief Append 'data' to the file.
-  virtual tsl::Status Append(StringPiece data) = 0;
+  virtual absl::Status Append(StringPiece data) = 0;
 
 #if defined(TF_CORD_SUPPORT)
   // \brief Append 'data' to the file.
-  virtual tsl::Status Append(const absl::Cord& cord) {
+  virtual absl::Status Append(const absl::Cord& cord) {
     for (StringPiece chunk : cord.Chunks()) {
       TF_RETURN_IF_ERROR(Append(chunk));
     }
-    return OkStatus();
+    return absl::OkStatus();
   }
 #endif
 
@@ -814,7 +818,7 @@ class WritableFile {
   /// Typical return codes (not guaranteed to be exhaustive):
   ///  * OK
   ///  * Other codes, as returned from Flush()
-  virtual tsl::Status Close() = 0;
+  virtual absl::Status Close() = 0;
 
   /// \brief Flushes the file and optionally syncs contents to filesystem.
   ///
@@ -826,13 +830,13 @@ class WritableFile {
   /// eventually flush the contents.  If the OS or machine crashes
   /// after a successful flush, the contents may or may not be
   /// persisted, depending on the implementation.
-  virtual tsl::Status Flush() = 0;
+  virtual absl::Status Flush() = 0;
 
   // \brief Returns the name of the file.
   ///
   /// This is an optional operation that may not be implemented by every
   /// filesystem.
-  virtual tsl::Status Name(StringPiece* result) const {
+  virtual absl::Status Name(StringPiece* result) const {
     return errors::Unimplemented("This filesystem does not support Name()");
   }
 
@@ -842,14 +846,14 @@ class WritableFile {
   /// of the file have been persisted to the filesystem; if the OS
   /// or machine crashes after a successful Sync, the contents should
   /// be properly saved.
-  virtual tsl::Status Sync() = 0;
+  virtual absl::Status Sync() = 0;
 
   /// \brief Retrieves the current write position in the file, or -1 on
   /// error.
   ///
   /// This is an optional operation, subclasses may choose to return
   /// errors::Unimplemented.
-  virtual tsl::Status Tell(int64_t* position) {
+  virtual absl::Status Tell(int64_t* position) {
     *position = -1;
     return errors::Unimplemented("This filesystem does not support Tell()");
   }
@@ -903,11 +907,11 @@ class FileSystemRegistry {
   typedef std::function<FileSystem*()> Factory;
 
   virtual ~FileSystemRegistry() = default;
-  virtual tsl::Status Register(const std::string& scheme, Factory factory) = 0;
-  virtual tsl::Status Register(const std::string& scheme,
-                               std::unique_ptr<FileSystem> filesystem) = 0;
+  virtual absl::Status Register(const std::string& scheme, Factory factory) = 0;
+  virtual absl::Status Register(const std::string& scheme,
+                                std::unique_ptr<FileSystem> filesystem) = 0;
   virtual FileSystem* Lookup(const std::string& scheme) = 0;
-  virtual tsl::Status GetRegisteredFileSystemSchemes(
+  virtual absl::Status GetRegisteredFileSystemSchemes(
       std::vector<std::string>* schemes) = 0;
 };
 
diff --git a/third_party/xla/third_party/tsl/tsl/platform/file_system_helper.cc b/third_party/xla/third_party/tsl/tsl/platform/file_system_helper.cc
index bc132ba02a56b5..04dc6a0420516c 100644
--- a/third_party/xla/third_party/tsl/tsl/platform/file_system_helper.cc
+++ b/third_party/xla/third_party/tsl/tsl/platform/file_system_helper.cc
@@ -121,18 +121,19 @@ static inline int GetFirstGlobbingEntry(const std::vector<std::string>& dirs) {
 
 }  // namespace
 
-Status GetMatchingPaths(FileSystem* fs, Env* env, const string& pattern,
-                        std::vector<string>* results) {
+absl::Status GetMatchingPaths(FileSystem* fs, Env* env, const string& pattern,
+                              std::vector<string>* results) {
   // Check that `fs`, `env` and `results` are non-null.
   if (fs == nullptr || env == nullptr || results == nullptr) {
-    return Status(absl::StatusCode::kInvalidArgument,
-                  "Filesystem calls GetMatchingPaths with nullptr arguments");
+    return absl::Status(
+        absl::StatusCode::kInvalidArgument,
+        "Filesystem calls GetMatchingPaths with nullptr arguments");
   }
 
   // By design, we don't match anything on empty pattern
   results->clear();
   if (pattern.empty()) {
-    return OkStatus();
+    return absl::OkStatus();
   }
 
   // The pattern can contain globbing characters at multiple levels, e.g.:
@@ -155,7 +156,7 @@ Status GetMatchingPaths(FileSystem* fs, Env* env, const string& pattern,
     if (fs->FileExists(pattern).ok()) {
       results->emplace_back(pattern);
     }
-    return OkStatus();
+    return absl::OkStatus();
   }
 
   // To expand the globbing, we do a BFS from `dirs[matching_index-1]`.
@@ -205,7 +206,7 @@ Status GetMatchingPaths(FileSystem* fs, Env* env, const string& pattern,
 
       // Get all children of `parent`. If this fails, return early.
       std::vector<std::string> children;
-      Status s = fs->GetChildren(parent, &children);
+      absl::Status s = fs->GetChildren(parent, &children);
       if (s.code() == absl::StatusCode::kPermissionDenied) {
         return;
       }
@@ -220,13 +221,13 @@ Status GetMatchingPaths(FileSystem* fs, Env* env, const string& pattern,
       // We also check that children match the pattern in parallel, for speedup.
       // We store the status of the match and `IsDirectory` in
       // `children_status` array, one element for each children.
-      std::vector<Status> children_status(children.size());
+      std::vector<absl::Status> children_status(children.size());
       auto handle_children = [&fs, &match_pattern, &parent, &children,
                               &children_status](int j) {
         const std::string path = io::JoinPath(parent, children[j]);
         if (!fs->Match(path, match_pattern)) {
-          children_status[j] =
-              Status(absl::StatusCode::kCancelled, "Operation not needed");
+          children_status[j] = absl::Status(absl::StatusCode::kCancelled,
+                                            "Operation not needed");
         } else {
           children_status[j] = fs->IsDirectory(path);
         }
@@ -263,11 +264,11 @@ Status GetMatchingPaths(FileSystem* fs, Env* env, const string& pattern,
     std::swap(expand_queue, next_expand_queue);
   }
 
-  return OkStatus();
+  return absl::OkStatus();
 }
 
-StatusOr<bool> FileExists(Env* env, const string& fname) {
-  Status status = env->FileExists(fname);
+absl::StatusOr<bool> FileExists(Env* env, const string& fname) {
+  absl::Status status = env->FileExists(fname);
   if (errors::IsNotFound(status)) {
     return false;
   }
diff --git a/third_party/xla/third_party/tsl/tsl/platform/file_system_helper.h b/third_party/xla/third_party/tsl/tsl/platform/file_system_helper.h
index 5cc2cdc5a1e898..e9e7df6aa68907 100644
--- a/third_party/xla/third_party/tsl/tsl/platform/file_system_helper.h
+++ b/third_party/xla/third_party/tsl/tsl/platform/file_system_helper.h
@@ -44,8 +44,8 @@ namespace internal {
 //   results: will be cleared and may not be null.
 //
 // Returns an error status if any call to 'fs' failed.
-Status GetMatchingPaths(FileSystem* fs, Env* env, const string& pattern,
-                        std::vector<string>* results);
+absl::Status GetMatchingPaths(FileSystem* fs, Env* env, const string& pattern,
+                              std::vector<string>* results);
 
 // Given a file path, determines whether the file exists. This helper simplifies
 // the use of Env::FileExists.
@@ -56,7 +56,7 @@ Status GetMatchingPaths(FileSystem* fs, Env* env, const string& pattern,
 //
 // Returns true if the file exists, false if it does not exist, or an error
 // Status.
-StatusOr<bool> FileExists(Env* env, const string& fname);
+absl::StatusOr<bool> FileExists(Env* env, const string& fname);
 
 }  // namespace internal
 }  // namespace tsl
diff --git a/third_party/xla/third_party/tsl/tsl/platform/null_file_system.h b/third_party/xla/third_party/tsl/tsl/platform/null_file_system.h
index 77b8142ee357b7..c04d2c1f0d6056 100644
--- a/third_party/xla/third_party/tsl/tsl/platform/null_file_system.h
+++ b/third_party/xla/third_party/tsl/tsl/platform/null_file_system.h
@@ -38,67 +38,72 @@ class NullFileSystem : public FileSystem {
 
   TF_USE_FILESYSTEM_METHODS_WITH_NO_TRANSACTION_SUPPORT;
 
-  Status NewRandomAccessFile(
+  absl::Status NewRandomAccessFile(
       const string& fname, TransactionToken* token,
       std::unique_ptr<RandomAccessFile>* result) override {
     return errors::Unimplemented("NewRandomAccessFile unimplemented");
   }
 
-  Status NewWritableFile(const string& fname, TransactionToken* token,
-                         std::unique_ptr<WritableFile>* result) override {
+  absl::Status NewWritableFile(const string& fname, TransactionToken* token,
+                               std::unique_ptr<WritableFile>* result) override {
     return errors::Unimplemented("NewWritableFile unimplemented");
   }
 
-  Status NewAppendableFile(const string& fname, TransactionToken* token,
-                           std::unique_ptr<WritableFile>* result) override {
+  absl::Status NewAppendableFile(
+      const string& fname, TransactionToken* token,
+      std::unique_ptr<WritableFile>* result) override {
     return errors::Unimplemented("NewAppendableFile unimplemented");
   }
 
-  Status NewReadOnlyMemoryRegionFromFile(
+  absl::Status NewReadOnlyMemoryRegionFromFile(
       const string& fname, TransactionToken* token,
       std::unique_ptr<ReadOnlyMemoryRegion>* result) override {
     return errors::Unimplemented(
         "NewReadOnlyMemoryRegionFromFile unimplemented");
   }
 
-  Status FileExists(const string& fname, TransactionToken* token) override {
+  absl::Status FileExists(const string& fname,
+                          TransactionToken* token) override {
     return errors::Unimplemented("FileExists unimplemented");
   }
 
-  Status GetChildren(const string& dir, TransactionToken* token,
-                     std::vector<string>* result) override {
+  absl::Status GetChildren(const string& dir, TransactionToken* token,
+                           std::vector<string>* result) override {
     return errors::Unimplemented("GetChildren unimplemented");
   }
 
-  Status GetMatchingPaths(const string& pattern, TransactionToken* token,
-                          std::vector<string>* results) override {
+  absl::Status GetMatchingPaths(const string& pattern, TransactionToken* token,
+                                std::vector<string>* results) override {
     return internal::GetMatchingPaths(this, Env::Default(), pattern, results);
   }
 
-  Status DeleteFile(const string& fname, TransactionToken* token) override {
+  absl::Status DeleteFile(const string& fname,
+                          TransactionToken* token) override {
     return errors::Unimplemented("DeleteFile unimplemented");
   }
 
-  Status CreateDir(const string& dirname, TransactionToken* token) override {
+  absl::Status CreateDir(const string& dirname,
+                         TransactionToken* token) override {
     return errors::Unimplemented("CreateDir unimplemented");
   }
 
-  Status DeleteDir(const string& dirname, TransactionToken* token) override {
+  absl::Status DeleteDir(const string& dirname,
+                         TransactionToken* token) override {
     return errors::Unimplemented("DeleteDir unimplemented");
   }
 
-  Status GetFileSize(const string& fname, TransactionToken* token,
-                     uint64* file_size) override {
+  absl::Status GetFileSize(const string& fname, TransactionToken* token,
+                           uint64* file_size) override {
     return errors::Unimplemented("GetFileSize unimplemented");
   }
 
-  Status RenameFile(const string& src, const string& target,
-                    TransactionToken* token) override {
+  absl::Status RenameFile(const string& src, const string& target,
+                          TransactionToken* token) override {
     return errors::Unimplemented("RenameFile unimplemented");
   }
 
-  Status Stat(const string& fname, TransactionToken* token,
-              FileStatistics* stat) override {
+  absl::Status Stat(const string& fname, TransactionToken* token,
+                    FileStatistics* stat) override {
     return errors::Unimplemented("Stat unimplemented");
   }
 };
diff --git a/third_party/xla/third_party/tsl/tsl/platform/ram_file_system.h b/third_party/xla/third_party/tsl/tsl/platform/ram_file_system.h
index 1b51653b716c3e..245eacfe465daa 100644
--- a/third_party/xla/third_party/tsl/tsl/platform/ram_file_system.h
+++ b/third_party/xla/third_party/tsl/tsl/platform/ram_file_system.h
@@ -49,13 +49,13 @@ class RamRandomAccessFile : public RandomAccessFile, public WritableFile {
       : name_(name), data_(cord) {}
   ~RamRandomAccessFile() override {}
 
-  Status Name(StringPiece* result) const override {
+  absl::Status Name(StringPiece* result) const override {
     *result = name_;
-    return OkStatus();
+    return absl::OkStatus();
   }
 
-  Status Read(uint64 offset, size_t n, StringPiece* result,
-              char* scratch) const override {
+  absl::Status Read(uint64 offset, size_t n, StringPiece* result,
+                    char* scratch) const override {
     if (offset >= data_->size()) {
       return errors::OutOfRange("");
     }
@@ -72,26 +72,26 @@ class RamRandomAccessFile : public RandomAccessFile, public WritableFile {
     if (left < n) {
       return errors::OutOfRange("");
     }
-    return OkStatus();
+    return absl::OkStatus();
   }
 
-  Status Append(StringPiece data) override {
+  absl::Status Append(StringPiece data) override {
     data_->append(data.data(), data.size());
-    return OkStatus();
+    return absl::OkStatus();
   }
 
 #if defined(TF_CORD_SUPPORT)
-  Status Append(const absl::Cord& cord) override {
+  absl::Status Append(const absl::Cord& cord) override {
     data_->append(cord.char_begin(), cord.char_end());
-    return OkStatus();
+    return absl::OkStatus();
   }
 #endif
 
-  Status Close() override { return OkStatus(); }
-  Status Flush() override { return OkStatus(); }
-  Status Sync() override { return OkStatus(); }
+  absl::Status Close() override { return absl::OkStatus(); }
+  absl::Status Flush() override { return absl::OkStatus(); }
+  absl::Status Sync() override { return absl::OkStatus(); }
 
-  Status Tell(int64_t* position) override {
+  absl::Status Tell(int64_t* position) override {
     *position = -1;
     return errors::Unimplemented("This filesystem does not support Tell()");
   }
@@ -107,7 +107,7 @@ class RamFileSystem : public FileSystem {
  public:
   TF_USE_FILESYSTEM_METHODS_WITH_NO_TRANSACTION_SUPPORT;
 
-  Status NewRandomAccessFile(
+  absl::Status NewRandomAccessFile(
       const std::string& fname_, TransactionToken* token,
       std::unique_ptr<RandomAccessFile>* result) override {
     mutex_lock m(mu_);
@@ -121,11 +121,12 @@ class RamFileSystem : public FileSystem {
     }
     *result = std::unique_ptr<RandomAccessFile>(
         new RamRandomAccessFile(fname, fs_[fname]));
-    return OkStatus();
+    return absl::OkStatus();
   }
 
-  Status NewWritableFile(const std::string& fname_, TransactionToken* token,
-                         std::unique_ptr<WritableFile>* result) override {
+  absl::Status NewWritableFile(const std::string& fname_,
+                               TransactionToken* token,
+                               std::unique_ptr<WritableFile>* result) override {
     mutex_lock m(mu_);
     auto fname = StripRamFsPrefix(fname_);
 
@@ -137,11 +138,12 @@ class RamFileSystem : public FileSystem {
     }
     *result = std::unique_ptr<WritableFile>(
         new RamRandomAccessFile(fname, fs_[fname]));
-    return OkStatus();
+    return absl::OkStatus();
   }
 
-  Status NewAppendableFile(const std::string& fname_, TransactionToken* token,
-                           std::unique_ptr<WritableFile>* result) override {
+  absl::Status NewAppendableFile(
+      const std::string& fname_, TransactionToken* token,
+      std::unique_ptr<WritableFile>* result) override {
     mutex_lock m(mu_);
     auto fname = StripRamFsPrefix(fname_);
 
@@ -153,25 +155,25 @@ class RamFileSystem : public FileSystem {
     }
     *result = std::unique_ptr<WritableFile>(
         new RamRandomAccessFile(fname, fs_[fname]));
-    return OkStatus();
+    return absl::OkStatus();
   }
 
-  Status NewReadOnlyMemoryRegionFromFile(
+  absl::Status NewReadOnlyMemoryRegionFromFile(
       const std::string& fname, TransactionToken* token,
       std::unique_ptr<ReadOnlyMemoryRegion>* result) override {
     return errors::Unimplemented("");
   }
 
-  Status FileExists(const std::string& fname_,
-                    TransactionToken* token) override {
+  absl::Status FileExists(const std::string& fname_,
+                          TransactionToken* token) override {
     FileStatistics stat;
     auto fname = StripRamFsPrefix(fname_);
 
     return Stat(fname, token, &stat);
   }
 
-  Status GetChildren(const std::string& dir_, TransactionToken* token,
-                     std::vector<std::string>* result) override {
+  absl::Status GetChildren(const std::string& dir_, TransactionToken* token,
+                           std::vector<std::string>* result) override {
     mutex_lock m(mu_);
     auto dir = StripRamFsPrefix(dir_);
 
@@ -185,11 +187,12 @@ class RamFileSystem : public FileSystem {
       ++it;
     }
 
-    return OkStatus();
+    return absl::OkStatus();
   }
 
-  Status GetMatchingPaths(const std::string& pattern_, TransactionToken* token,
-                          std::vector<std::string>* results) override {
+  absl::Status GetMatchingPaths(const std::string& pattern_,
+                                TransactionToken* token,
+                                std::vector<std::string>* results) override {
     mutex_lock m(mu_);
     auto pattern = StripRamFsPrefix(pattern_);
 
@@ -199,11 +202,11 @@ class RamFileSystem : public FileSystem {
         results->push_back("ram://" + it->first);
       }
     }
-    return OkStatus();
+    return absl::OkStatus();
   }
 
-  Status Stat(const std::string& fname_, TransactionToken* token,
-              FileStatistics* stat) override {
+  absl::Status Stat(const std::string& fname_, TransactionToken* token,
+                    FileStatistics* stat) override {
     mutex_lock m(mu_);
     auto fname = StripRamFsPrefix(fname_);
 
@@ -216,30 +219,30 @@ class RamFileSystem : public FileSystem {
       stat->is_directory = false;
       stat->length = fs_[fname]->size();
       stat->mtime_nsec = 0;
-      return OkStatus();
+      return absl::OkStatus();
     }
 
     stat->is_directory = true;
     stat->length = 0;
     stat->mtime_nsec = 0;
-    return OkStatus();
+    return absl::OkStatus();
   }
 
-  Status DeleteFile(const std::string& fname_,
-                    TransactionToken* token) override {
+  absl::Status DeleteFile(const std::string& fname_,
+                          TransactionToken* token) override {
     mutex_lock m(mu_);
     auto fname = StripRamFsPrefix(fname_);
 
     if (fs_.find(fname) != fs_.end()) {
       fs_.erase(fname);
-      return OkStatus();
+      return absl::OkStatus();
     }
 
     return errors::NotFound("");
   }
 
-  Status CreateDir(const std::string& dirname_,
-                   TransactionToken* token) override {
+  absl::Status CreateDir(const std::string& dirname_,
+                         TransactionToken* token) override {
     mutex_lock m(mu_);
     auto dirname = StripRamFsPrefix(dirname_);
 
@@ -250,15 +253,15 @@ class RamFileSystem : public FileSystem {
     }
 
     fs_[dirname] = nullptr;
-    return OkStatus();
+    return absl::OkStatus();
   }
 
-  Status RecursivelyCreateDir(const std::string& dirname_,
-                              TransactionToken* token) override {
+  absl::Status RecursivelyCreateDir(const std::string& dirname_,
+                                    TransactionToken* token) override {
     auto dirname = StripRamFsPrefix(dirname_);
 
     std::vector<std::string> dirs = StrSplit(dirname, "/");
-    Status last_status;
+    absl::Status last_status;
     std::string dir = dirs[0];
     last_status = CreateDir(dir, token);
 
@@ -269,8 +272,8 @@ class RamFileSystem : public FileSystem {
     return last_status;
   }
 
-  Status DeleteDir(const std::string& dirname_,
-                   TransactionToken* token) override {
+  absl::Status DeleteDir(const std::string& dirname_,
+                         TransactionToken* token) override {
     mutex_lock m(mu_);
     auto dirname = StripRamFsPrefix(dirname_);
 
@@ -283,11 +286,11 @@ class RamFileSystem : public FileSystem {
     }
     fs_.erase(dirname);
 
-    return OkStatus();
+    return absl::OkStatus();
   }
 
-  Status GetFileSize(const std::string& fname_, TransactionToken* token,
-                     uint64* file_size) override {
+  absl::Status GetFileSize(const std::string& fname_, TransactionToken* token,
+                           uint64* file_size) override {
     mutex_lock m(mu_);
     auto fname = StripRamFsPrefix(fname_);
 
@@ -296,13 +299,13 @@ class RamFileSystem : public FileSystem {
         return errors::InvalidArgument("Not a file");
       }
       *file_size = fs_[fname]->size();
-      return OkStatus();
+      return absl::OkStatus();
     }
     return errors::NotFound("");
   }
 
-  Status RenameFile(const std::string& src_, const std::string& target_,
-                    TransactionToken* token) override {
+  absl::Status RenameFile(const std::string& src_, const std::string& target_,
+                          TransactionToken* token) override {
     mutex_lock m(mu_);
     auto src = StripRamFsPrefix(src_);
     auto target = StripRamFsPrefix(target_);
@@ -310,7 +313,7 @@ class RamFileSystem : public FileSystem {
     if (fs_.find(src) != fs_.end()) {
       fs_[target] = fs_[src];
       fs_.erase(fs_.find(src));
-      return OkStatus();
+      return absl::OkStatus();
     }
     return errors::NotFound("");
   }
diff --git a/third_party/xla/third_party/tsl/tsl/platform/retrying_file_system.h b/third_party/xla/third_party/tsl/tsl/platform/retrying_file_system.h
index 591423b4fe3ec7..3db1665a4f1867 100644
--- a/third_party/xla/third_party/tsl/tsl/platform/retrying_file_system.h
+++ b/third_party/xla/third_party/tsl/tsl/platform/retrying_file_system.h
@@ -40,21 +40,23 @@ class RetryingFileSystem : public FileSystem {
 
   TF_USE_FILESYSTEM_METHODS_WITH_NO_TRANSACTION_SUPPORT;
 
-  Status NewRandomAccessFile(
+  absl::Status NewRandomAccessFile(
       const string& filename, TransactionToken* token,
       std::unique_ptr<RandomAccessFile>* result) override;
 
-  Status NewWritableFile(const string& filename, TransactionToken* token,
-                         std::unique_ptr<WritableFile>* result) override;
+  absl::Status NewWritableFile(const string& filename, TransactionToken* token,
+                               std::unique_ptr<WritableFile>* result) override;
 
-  Status NewAppendableFile(const string& filename, TransactionToken* token,
-                           std::unique_ptr<WritableFile>* result) override;
+  absl::Status NewAppendableFile(
+      const string& filename, TransactionToken* token,
+      std::unique_ptr<WritableFile>* result) override;
 
-  Status NewReadOnlyMemoryRegionFromFile(
+  absl::Status NewReadOnlyMemoryRegionFromFile(
       const string& filename, TransactionToken* token,
       std::unique_ptr<ReadOnlyMemoryRegion>* result) override;
 
-  Status FileExists(const string& fname, TransactionToken* token) override {
+  absl::Status FileExists(const string& fname,
+                          TransactionToken* token) override {
     return RetryingUtils::CallWithRetries(
         [this, &fname, token]() {
           return base_file_system_->FileExists(fname, token);
@@ -62,8 +64,8 @@ class RetryingFileSystem : public FileSystem {
         retry_config_);
   }
 
-  Status GetChildren(const string& dir, TransactionToken* token,
-                     std::vector<string>* result) override {
+  absl::Status GetChildren(const string& dir, TransactionToken* token,
+                           std::vector<string>* result) override {
     return RetryingUtils::CallWithRetries(
         [this, &dir, result, token]() {
           return base_file_system_->GetChildren(dir, token, result);
@@ -71,8 +73,8 @@ class RetryingFileSystem : public FileSystem {
         retry_config_);
   }
 
-  Status GetMatchingPaths(const string& pattern, TransactionToken* token,
-                          std::vector<string>* result) override {
+  absl::Status GetMatchingPaths(const string& pattern, TransactionToken* token,
+                                std::vector<string>* result) override {
     return RetryingUtils::CallWithRetries(
         [this, &pattern, result, token]() {
           return base_file_system_->GetMatchingPaths(pattern, token, result);
@@ -80,8 +82,8 @@ class RetryingFileSystem : public FileSystem {
         retry_config_);
   }
 
-  Status Stat(const string& fname, TransactionToken* token,
-              FileStatistics* stat) override {
+  absl::Status Stat(const string& fname, TransactionToken* token,
+                    FileStatistics* stat) override {
     return RetryingUtils::CallWithRetries(
         [this, &fname, stat, token]() {
           return base_file_system_->Stat(fname, token, stat);
@@ -89,7 +91,8 @@ class RetryingFileSystem : public FileSystem {
         retry_config_);
   }
 
-  Status DeleteFile(const string& fname, TransactionToken* token) override {
+  absl::Status DeleteFile(const string& fname,
+                          TransactionToken* token) override {
     return RetryingUtils::DeleteWithRetries(
         [this, &fname, token]() {
           return base_file_system_->DeleteFile(fname, token);
@@ -97,7 +100,8 @@ class RetryingFileSystem : public FileSystem {
         retry_config_);
   }
 
-  Status CreateDir(const string& dirname, TransactionToken* token) override {
+  absl::Status CreateDir(const string& dirname,
+                         TransactionToken* token) override {
     return RetryingUtils::CallWithRetries(
         [this, &dirname, token]() {
           return base_file_system_->CreateDir(dirname, token);
@@ -105,7 +109,8 @@ class RetryingFileSystem : public FileSystem {
         retry_config_);
   }
 
-  Status DeleteDir(const string& dirname, TransactionToken* token) override {
+  absl::Status DeleteDir(const string& dirname,
+                         TransactionToken* token) override {
     return RetryingUtils::DeleteWithRetries(
         [this, &dirname, token]() {
           return base_file_system_->DeleteDir(dirname, token);
@@ -113,8 +118,8 @@ class RetryingFileSystem : public FileSystem {
         retry_config_);
   }
 
-  Status GetFileSize(const string& fname, TransactionToken* token,
-                     uint64* file_size) override {
+  absl::Status GetFileSize(const string& fname, TransactionToken* token,
+                           uint64* file_size) override {
     return RetryingUtils::CallWithRetries(
         [this, &fname, file_size, token]() {
           return base_file_system_->GetFileSize(fname, token, file_size);
@@ -122,8 +127,8 @@ class RetryingFileSystem : public FileSystem {
         retry_config_);
   }
 
-  Status RenameFile(const string& src, const string& target,
-                    TransactionToken* token) override {
+  absl::Status RenameFile(const string& src, const string& target,
+                          TransactionToken* token) override {
     return RetryingUtils::CallWithRetries(
         [this, &src, &target, token]() {
           return base_file_system_->RenameFile(src, target, token);
@@ -131,7 +136,8 @@ class RetryingFileSystem : public FileSystem {
         retry_config_);
   }
 
-  Status IsDirectory(const string& dirname, TransactionToken* token) override {
+  absl::Status IsDirectory(const string& dirname,
+                           TransactionToken* token) override {
     return RetryingUtils::CallWithRetries(
         [this, &dirname, token]() {
           return base_file_system_->IsDirectory(dirname, token);
@@ -139,14 +145,15 @@ class RetryingFileSystem : public FileSystem {
         retry_config_);
   }
 
-  Status HasAtomicMove(const string& path, bool* has_atomic_move) override {
+  absl::Status HasAtomicMove(const string& path,
+                             bool* has_atomic_move) override {
     // this method does not need to be retried
     return base_file_system_->HasAtomicMove(path, has_atomic_move);
   }
 
-  Status DeleteRecursively(const string& dirname, TransactionToken* token,
-                           int64_t* undeleted_files,
-                           int64_t* undeleted_dirs) override {
+  absl::Status DeleteRecursively(const string& dirname, TransactionToken* token,
+                                 int64_t* undeleted_files,
+                                 int64_t* undeleted_dirs) override {
     return RetryingUtils::DeleteWithRetries(
         [this, &dirname, token, undeleted_files, undeleted_dirs]() {
           return base_file_system_->DeleteRecursively(
@@ -177,12 +184,12 @@ class RetryingRandomAccessFile : public RandomAccessFile {
                            const RetryConfig& retry_config)
       : base_file_(std::move(base_file)), retry_config_(retry_config) {}
 
-  Status Name(StringPiece* result) const override {
+  absl::Status Name(StringPiece* result) const override {
     return base_file_->Name(result);
   }
 
-  Status Read(uint64 offset, size_t n, StringPiece* result,
-              char* scratch) const override {
+  absl::Status Read(uint64 offset, size_t n, StringPiece* result,
+                    char* scratch) const override {
     return RetryingUtils::CallWithRetries(
         [this, offset, n, result, scratch]() {
           return base_file_->Read(offset, n, result, scratch);
@@ -206,26 +213,26 @@ class RetryingWritableFile : public WritableFile {
     Close().IgnoreError();
   }
 
-  Status Append(StringPiece data) override {
+  absl::Status Append(StringPiece data) override {
     return RetryingUtils::CallWithRetries(
         [this, &data]() { return base_file_->Append(data); }, retry_config_);
   }
-  Status Close() override {
+  absl::Status Close() override {
     return RetryingUtils::CallWithRetries(
         [this]() { return base_file_->Close(); }, retry_config_);
   }
-  Status Flush() override {
+  absl::Status Flush() override {
     return RetryingUtils::CallWithRetries(
         [this]() { return base_file_->Flush(); }, retry_config_);
   }
-  Status Name(StringPiece* result) const override {
+  absl::Status Name(StringPiece* result) const override {
     return base_file_->Name(result);
   }
-  Status Sync() override {
+  absl::Status Sync() override {
     return RetryingUtils::CallWithRetries(
         [this]() { return base_file_->Sync(); }, retry_config_);
   }
-  Status Tell(int64_t* position) override {
+  absl::Status Tell(int64_t* position) override {
     return RetryingUtils::CallWithRetries(
         [this, &position]() { return base_file_->Tell(position); },
         retry_config_);
@@ -239,7 +246,7 @@ class RetryingWritableFile : public WritableFile {
 }  // namespace retrying_internals
 
 template <typename Underlying>
-Status RetryingFileSystem<Underlying>::NewRandomAccessFile(
+absl::Status RetryingFileSystem<Underlying>::NewRandomAccessFile(
     const string& filename, TransactionToken* token,
     std::unique_ptr<RandomAccessFile>* result) {
   std::unique_ptr<RandomAccessFile> base_file;
@@ -251,11 +258,11 @@ Status RetryingFileSystem<Underlying>::NewRandomAccessFile(
       retry_config_));
   result->reset(new retrying_internals::RetryingRandomAccessFile(
       std::move(base_file), retry_config_));
-  return OkStatus();
+  return absl::OkStatus();
 }
 
 template <typename Underlying>
-Status RetryingFileSystem<Underlying>::NewWritableFile(
+absl::Status RetryingFileSystem<Underlying>::NewWritableFile(
     const string& filename, TransactionToken* token,
     std::unique_ptr<WritableFile>* result) {
   std::unique_ptr<WritableFile> base_file;
@@ -266,11 +273,11 @@ Status RetryingFileSystem<Underlying>::NewWritableFile(
       retry_config_));
   result->reset(new retrying_internals::RetryingWritableFile(
       std::move(base_file), retry_config_));
-  return OkStatus();
+  return absl::OkStatus();
 }
 
 template <typename Underlying>
-Status RetryingFileSystem<Underlying>::NewAppendableFile(
+absl::Status RetryingFileSystem<Underlying>::NewAppendableFile(
     const string& filename, TransactionToken* token,
     std::unique_ptr<WritableFile>* result) {
   std::unique_ptr<WritableFile> base_file;
@@ -282,11 +289,11 @@ Status RetryingFileSystem<Underlying>::NewAppendableFile(
       retry_config_));
   result->reset(new retrying_internals::RetryingWritableFile(
       std::move(base_file), retry_config_));
-  return OkStatus();
+  return absl::OkStatus();
 }
 
 template <typename Underlying>
-Status RetryingFileSystem<Underlying>::NewReadOnlyMemoryRegionFromFile(
+absl::Status RetryingFileSystem<Underlying>::NewReadOnlyMemoryRegionFromFile(
     const string& filename, TransactionToken* token,
     std::unique_ptr<ReadOnlyMemoryRegion>* result) {
   return RetryingUtils::CallWithRetries(
diff --git a/third_party/xla/third_party/tsl/tsl/platform/retrying_file_system_test.cc b/third_party/xla/third_party/tsl/tsl/platform/retrying_file_system_test.cc
index 605ce2eb7ea3c3..522c59f565e0b1 100644
--- a/third_party/xla/third_party/tsl/tsl/platform/retrying_file_system_test.cc
+++ b/third_party/xla/third_party/tsl/tsl/platform/retrying_file_system_test.cc
@@ -24,7 +24,7 @@ limitations under the License.
 namespace tsl {
 namespace {
 
-typedef std::vector<std::tuple<string, Status>> ExpectedCalls;
+typedef std::vector<std::tuple<string, absl::Status>> ExpectedCalls;
 
 ExpectedCalls CreateRetriableErrors(const string& method, int n) {
   ExpectedCalls expected_calls;
@@ -47,7 +47,7 @@ class MockCallSequence {
         << "the next expected call: " << std::get<0>(calls_.front());
   }
 
-  Status ConsumeNextCall(const string& method) {
+  absl::Status ConsumeNextCall(const string& method) {
     EXPECT_FALSE(calls_.empty()) << "No more calls were expected.";
     auto call = calls_.front();
     calls_.erase(calls_.begin());
@@ -62,11 +62,11 @@ class MockCallSequence {
 class MockRandomAccessFile : public RandomAccessFile {
  public:
   explicit MockRandomAccessFile(const ExpectedCalls& calls) : calls_(calls) {}
-  Status Name(StringPiece* result) const override {
+  absl::Status Name(StringPiece* result) const override {
     return calls_.ConsumeNextCall("Name");
   }
-  Status Read(uint64 offset, size_t n, StringPiece* result,
-              char* scratch) const override {
+  absl::Status Read(uint64 offset, size_t n, StringPiece* result,
+                    char* scratch) const override {
     return calls_.ConsumeNextCall("Read");
   }
 
@@ -77,16 +77,16 @@ class MockRandomAccessFile : public RandomAccessFile {
 class MockWritableFile : public WritableFile {
  public:
   explicit MockWritableFile(const ExpectedCalls& calls) : calls_(calls) {}
-  Status Append(StringPiece data) override {
+  absl::Status Append(StringPiece data) override {
     return calls_.ConsumeNextCall("Append");
   }
-  Status Close() override { return calls_.ConsumeNextCall("Close"); }
-  Status Flush() override { return calls_.ConsumeNextCall("Flush"); }
-  Status Name(StringPiece* result) const override {
+  absl::Status Close() override { return calls_.ConsumeNextCall("Close"); }
+  absl::Status Flush() override { return calls_.ConsumeNextCall("Flush"); }
+  absl::Status Name(StringPiece* result) const override {
     return calls_.ConsumeNextCall("Name");
   }
-  Status Sync() override { return calls_.ConsumeNextCall("Sync"); }
-  Status Tell(int64_t* position) override {
+  absl::Status Sync() override { return calls_.ConsumeNextCall("Sync"); }
+  absl::Status Tell(int64_t* position) override {
     return calls_.ConsumeNextCall("Tell");
   }
 
@@ -101,79 +101,85 @@ class MockFileSystem : public FileSystem {
 
   TF_USE_FILESYSTEM_METHODS_WITH_NO_TRANSACTION_SUPPORT;
 
-  Status NewRandomAccessFile(
+  absl::Status NewRandomAccessFile(
       const string& fname, TransactionToken* token,
       std::unique_ptr<RandomAccessFile>* result) override {
     *result = std::move(random_access_file_to_return);
     return calls_.ConsumeNextCall("NewRandomAccessFile");
   }
 
-  Status NewWritableFile(const string& fname, TransactionToken* token,
-                         std::unique_ptr<WritableFile>* result) override {
+  absl::Status NewWritableFile(const string& fname, TransactionToken* token,
+                               std::unique_ptr<WritableFile>* result) override {
     *result = std::move(writable_file_to_return);
     return calls_.ConsumeNextCall("NewWritableFile");
   }
 
-  Status NewAppendableFile(const string& fname, TransactionToken* token,
-                           std::unique_ptr<WritableFile>* result) override {
+  absl::Status NewAppendableFile(
+      const string& fname, TransactionToken* token,
+      std::unique_ptr<WritableFile>* result) override {
     *result = std::move(writable_file_to_return);
     return calls_.ConsumeNextCall("NewAppendableFile");
   }
 
-  Status NewReadOnlyMemoryRegionFromFile(
+  absl::Status NewReadOnlyMemoryRegionFromFile(
       const string& fname, TransactionToken* token,
       std::unique_ptr<ReadOnlyMemoryRegion>* result) override {
     return calls_.ConsumeNextCall("NewReadOnlyMemoryRegionFromFile");
   }
 
-  Status FileExists(const string& fname, TransactionToken* token) override {
+  absl::Status FileExists(const string& fname,
+                          TransactionToken* token) override {
     return calls_.ConsumeNextCall("FileExists");
   }
 
-  Status GetChildren(const string& dir, TransactionToken* token,
-                     std::vector<string>* result) override {
+  absl::Status GetChildren(const string& dir, TransactionToken* token,
+                           std::vector<string>* result) override {
     return calls_.ConsumeNextCall("GetChildren");
   }
 
-  Status GetMatchingPaths(const string& dir, TransactionToken* token,
-                          std::vector<string>* result) override {
+  absl::Status GetMatchingPaths(const string& dir, TransactionToken* token,
+                                std::vector<string>* result) override {
     return calls_.ConsumeNextCall("GetMatchingPaths");
   }
 
-  Status Stat(const string& fname, TransactionToken* token,
-              FileStatistics* stat) override {
+  absl::Status Stat(const string& fname, TransactionToken* token,
+                    FileStatistics* stat) override {
     return calls_.ConsumeNextCall("Stat");
   }
 
-  Status DeleteFile(const string& fname, TransactionToken* token) override {
+  absl::Status DeleteFile(const string& fname,
+                          TransactionToken* token) override {
     return calls_.ConsumeNextCall("DeleteFile");
   }
 
-  Status CreateDir(const string& dirname, TransactionToken* token) override {
+  absl::Status CreateDir(const string& dirname,
+                         TransactionToken* token) override {
     return calls_.ConsumeNextCall("CreateDir");
   }
 
-  Status DeleteDir(const string& dirname, TransactionToken* token) override {
+  absl::Status DeleteDir(const string& dirname,
+                         TransactionToken* token) override {
     return calls_.ConsumeNextCall("DeleteDir");
   }
 
-  Status GetFileSize(const string& fname, TransactionToken* token,
-                     uint64* file_size) override {
+  absl::Status GetFileSize(const string& fname, TransactionToken* token,
+                           uint64* file_size) override {
     return calls_.ConsumeNextCall("GetFileSize");
   }
 
-  Status RenameFile(const string& src, const string& target,
-                    TransactionToken* token) override {
+  absl::Status RenameFile(const string& src, const string& target,
+                          TransactionToken* token) override {
     return calls_.ConsumeNextCall("RenameFile");
   }
 
-  Status IsDirectory(const string& dirname, TransactionToken* token) override {
+  absl::Status IsDirectory(const string& dirname,
+                           TransactionToken* token) override {
     return calls_.ConsumeNextCall("IsDirectory");
   }
 
-  Status DeleteRecursively(const string& dirname, TransactionToken* token,
-                           int64_t* undeleted_files,
-                           int64_t* undeleted_dirs) override {
+  absl::Status DeleteRecursively(const string& dirname, TransactionToken* token,
+                                 int64_t* undeleted_files,
+                                 int64_t* undeleted_dirs) override {
     return calls_.ConsumeNextCall("DeleteRecursively");
   }
 
@@ -193,14 +199,15 @@ class MockFileSystem : public FileSystem {
 
 TEST(RetryingFileSystemTest, NewRandomAccessFile_ImmediateSuccess) {
   // Configure the mock base random access file.
-  ExpectedCalls expected_file_calls({std::make_tuple("Name", OkStatus()),
-                                     std::make_tuple("Read", OkStatus())});
+  ExpectedCalls expected_file_calls(
+      {std::make_tuple("Name", absl::OkStatus()),
+       std::make_tuple("Read", absl::OkStatus())});
   std::unique_ptr<RandomAccessFile> base_file(
       new MockRandomAccessFile(expected_file_calls));
 
   // Configure the mock base file system.
   ExpectedCalls expected_fs_calls(
-      {std::make_tuple("NewRandomAccessFile", OkStatus())});
+      {std::make_tuple("NewRandomAccessFile", absl::OkStatus())});
   std::unique_ptr<MockFileSystem> base_fs(
       new MockFileSystem(expected_fs_calls));
   base_fs->random_access_file_to_return = std::move(base_file);
@@ -226,13 +233,13 @@ TEST(RetryingFileSystemTest, NewRandomAccessFile_SuccessWith3rdTry) {
   ExpectedCalls expected_file_calls(
       {std::make_tuple("Read", errors::Unavailable("Something is wrong")),
        std::make_tuple("Read", errors::Unavailable("Wrong again")),
-       std::make_tuple("Read", OkStatus())});
+       std::make_tuple("Read", absl::OkStatus())});
   std::unique_ptr<RandomAccessFile> base_file(
       new MockRandomAccessFile(expected_file_calls));
 
   // Configure the mock base file system.
   ExpectedCalls expected_fs_calls(
-      {std::make_tuple("NewRandomAccessFile", OkStatus())});
+      {std::make_tuple("NewRandomAccessFile", absl::OkStatus())});
   std::unique_ptr<MockFileSystem> base_fs(
       new MockFileSystem(expected_fs_calls));
   base_fs->random_access_file_to_return = std::move(base_file);
@@ -258,7 +265,7 @@ TEST(RetryingFileSystemTest, NewRandomAccessFile_AllRetriesFailed) {
 
   // Configure the mock base file system.
   ExpectedCalls expected_fs_calls(
-      {std::make_tuple("NewRandomAccessFile", OkStatus())});
+      {std::make_tuple("NewRandomAccessFile", absl::OkStatus())});
   std::unique_ptr<MockFileSystem> base_fs(
       new MockFileSystem(expected_fs_calls));
   base_fs->random_access_file_to_return = std::move(base_file);
@@ -289,7 +296,7 @@ TEST(RetryingFileSystemTest, NewRandomAccessFile_NoRetriesForSomeErrors) {
 
   // Configure the mock base file system.
   ExpectedCalls expected_fs_calls(
-      {std::make_tuple("NewRandomAccessFile", OkStatus())});
+      {std::make_tuple("NewRandomAccessFile", absl::OkStatus())});
   std::unique_ptr<MockFileSystem> base_fs(
       new MockFileSystem(expected_fs_calls));
   base_fs->random_access_file_to_return = std::move(base_file);
@@ -310,15 +317,16 @@ TEST(RetryingFileSystemTest, NewRandomAccessFile_NoRetriesForSomeErrors) {
 
 TEST(RetryingFileSystemTest, NewWritableFile_ImmediateSuccess) {
   // Configure the mock base random access file.
-  ExpectedCalls expected_file_calls({std::make_tuple("Name", OkStatus()),
-                                     std::make_tuple("Sync", OkStatus()),
-                                     std::make_tuple("Close", OkStatus())});
+  ExpectedCalls expected_file_calls(
+      {std::make_tuple("Name", absl::OkStatus()),
+       std::make_tuple("Sync", absl::OkStatus()),
+       std::make_tuple("Close", absl::OkStatus())});
   std::unique_ptr<WritableFile> base_file(
       new MockWritableFile(expected_file_calls));
 
   // Configure the mock base file system.
   ExpectedCalls expected_fs_calls(
-      {std::make_tuple("NewWritableFile", OkStatus())});
+      {std::make_tuple("NewWritableFile", absl::OkStatus())});
   std::unique_ptr<MockFileSystem> base_fs(
       new MockFileSystem(expected_fs_calls));
   base_fs->writable_file_to_return = std::move(base_file);
@@ -342,14 +350,14 @@ TEST(RetryingFileSystemTest, NewWritableFile_SuccessWith3rdTry) {
   ExpectedCalls expected_file_calls(
       {std::make_tuple("Sync", errors::Unavailable("Something is wrong")),
        std::make_tuple("Sync", errors::Unavailable("Something is wrong again")),
-       std::make_tuple("Sync", OkStatus()),
-       std::make_tuple("Close", OkStatus())});
+       std::make_tuple("Sync", absl::OkStatus()),
+       std::make_tuple("Close", absl::OkStatus())});
   std::unique_ptr<WritableFile> base_file(
       new MockWritableFile(expected_file_calls));
 
   // Configure the mock base file system.
   ExpectedCalls expected_fs_calls(
-      {std::make_tuple("NewWritableFile", OkStatus())});
+      {std::make_tuple("NewWritableFile", absl::OkStatus())});
   std::unique_ptr<MockFileSystem> base_fs(
       new MockFileSystem(expected_fs_calls));
   base_fs->writable_file_to_return = std::move(base_file);
@@ -370,13 +378,13 @@ TEST(RetryingFileSystemTest, NewWritableFile_SuccessWith3rdTry_ViaDestructor) {
       {std::make_tuple("Close", errors::Unavailable("Something is wrong")),
        std::make_tuple("Close",
                        errors::Unavailable("Something is wrong again")),
-       std::make_tuple("Close", OkStatus())});
+       std::make_tuple("Close", absl::OkStatus())});
   std::unique_ptr<WritableFile> base_file(
       new MockWritableFile(expected_file_calls));
 
   // Configure the mock base file system.
   ExpectedCalls expected_fs_calls(
-      {std::make_tuple("NewWritableFile", OkStatus())});
+      {std::make_tuple("NewWritableFile", absl::OkStatus())});
   std::unique_ptr<MockFileSystem> base_fs(
       new MockFileSystem(expected_fs_calls));
   base_fs->writable_file_to_return = std::move(base_file);
@@ -395,14 +403,14 @@ TEST(RetryingFileSystemTest, NewAppendableFile_SuccessWith3rdTry) {
   ExpectedCalls expected_file_calls(
       {std::make_tuple("Sync", errors::Unavailable("Something is wrong")),
        std::make_tuple("Sync", errors::Unavailable("Something is wrong again")),
-       std::make_tuple("Sync", OkStatus()),
-       std::make_tuple("Close", OkStatus())});
+       std::make_tuple("Sync", absl::OkStatus()),
+       std::make_tuple("Close", absl::OkStatus())});
   std::unique_ptr<WritableFile> base_file(
       new MockWritableFile(expected_file_calls));
 
   // Configure the mock base file system.
   ExpectedCalls expected_fs_calls(
-      {std::make_tuple("NewAppendableFile", OkStatus())});
+      {std::make_tuple("NewAppendableFile", absl::OkStatus())});
   std::unique_ptr<MockFileSystem> base_fs(
       new MockFileSystem(expected_fs_calls));
   base_fs->writable_file_to_return = std::move(base_file);
@@ -420,13 +428,13 @@ TEST(RetryingFileSystemTest, NewAppendableFile_SuccessWith3rdTry) {
 TEST(RetryingFileSystemTest, NewWritableFile_AllRetriesFailed) {
   // Configure the mock base random access file.
   ExpectedCalls expected_file_calls = CreateRetriableErrors("Sync", 11);
-  expected_file_calls.emplace_back(std::make_tuple("Close", OkStatus()));
+  expected_file_calls.emplace_back(std::make_tuple("Close", absl::OkStatus()));
   std::unique_ptr<WritableFile> base_file(
       new MockWritableFile(expected_file_calls));
 
   // Configure the mock base file system.
   ExpectedCalls expected_fs_calls(
-      {std::make_tuple("NewWritableFile", OkStatus())});
+      {std::make_tuple("NewWritableFile", absl::OkStatus())});
   std::unique_ptr<MockFileSystem> base_fs(
       new MockFileSystem(expected_fs_calls));
   base_fs->writable_file_to_return = std::move(base_file);
@@ -448,7 +456,7 @@ TEST(RetryingFileSystemTest,
   ExpectedCalls expected_fs_calls(
       {std::make_tuple("NewReadOnlyMemoryRegionFromFile",
                        errors::Unavailable("Something is wrong")),
-       std::make_tuple("NewReadOnlyMemoryRegionFromFile", OkStatus())});
+       std::make_tuple("NewReadOnlyMemoryRegionFromFile", absl::OkStatus())});
   std::unique_ptr<MockFileSystem> base_fs(
       new MockFileSystem(expected_fs_calls));
   RetryingFileSystem<MockFileSystem> fs(
@@ -478,7 +486,7 @@ TEST(RetryingFileSystemTest, GetChildren_SuccessWith2ndTry) {
   ExpectedCalls expected_fs_calls(
       {std::make_tuple("GetChildren",
                        errors::Unavailable("Something is wrong")),
-       std::make_tuple("GetChildren", OkStatus())});
+       std::make_tuple("GetChildren", absl::OkStatus())});
   std::unique_ptr<MockFileSystem> base_fs(
       new MockFileSystem(expected_fs_calls));
   RetryingFileSystem<MockFileSystem> fs(
@@ -505,7 +513,7 @@ TEST(RetryingFileSystemTest, GetMatchingPaths_SuccessWith2ndTry) {
   ExpectedCalls expected_fs_calls(
       {std::make_tuple("GetMatchingPaths",
                        errors::Unavailable("Something is wrong")),
-       std::make_tuple("GetMatchingPaths", OkStatus())});
+       std::make_tuple("GetMatchingPaths", absl::OkStatus())});
   std::unique_ptr<MockFileSystem> base_fs(
       new MockFileSystem(expected_fs_calls));
   RetryingFileSystem<MockFileSystem> fs(
@@ -532,7 +540,7 @@ TEST(RetryingFileSystemTest, GetMatchingPaths_AllRetriesFailed) {
 TEST(RetryingFileSystemTest, DeleteFile_SuccessWith2ndTry) {
   ExpectedCalls expected_fs_calls(
       {std::make_tuple("DeleteFile", errors::Unavailable("Something is wrong")),
-       std::make_tuple("DeleteFile", OkStatus())});
+       std::make_tuple("DeleteFile", absl::OkStatus())});
   std::unique_ptr<MockFileSystem> base_fs(
       new MockFileSystem(expected_fs_calls));
   RetryingFileSystem<MockFileSystem> fs(
@@ -556,7 +564,7 @@ TEST(RetryingFileSystemTest, DeleteFile_AllRetriesFailed) {
 TEST(RetryingFileSystemTest, CreateDir_SuccessWith2ndTry) {
   ExpectedCalls expected_fs_calls(
       {std::make_tuple("CreateDir", errors::Unavailable("Something is wrong")),
-       std::make_tuple("CreateDir", OkStatus())});
+       std::make_tuple("CreateDir", absl::OkStatus())});
   std::unique_ptr<MockFileSystem> base_fs(
       new MockFileSystem(expected_fs_calls));
   RetryingFileSystem<MockFileSystem> fs(
@@ -580,7 +588,7 @@ TEST(RetryingFileSystemTest, CreateDir_AllRetriesFailed) {
 TEST(RetryingFileSystemTest, DeleteDir_SuccessWith2ndTry) {
   ExpectedCalls expected_fs_calls(
       {std::make_tuple("DeleteDir", errors::Unavailable("Something is wrong")),
-       std::make_tuple("DeleteDir", OkStatus())});
+       std::make_tuple("DeleteDir", absl::OkStatus())});
   std::unique_ptr<MockFileSystem> base_fs(
       new MockFileSystem(expected_fs_calls));
   RetryingFileSystem<MockFileSystem> fs(
@@ -605,7 +613,7 @@ TEST(RetryingFileSystemTest, GetFileSize_SuccessWith2ndTry) {
   ExpectedCalls expected_fs_calls(
       {std::make_tuple("GetFileSize",
                        errors::Unavailable("Something is wrong")),
-       std::make_tuple("GetFileSize", OkStatus())});
+       std::make_tuple("GetFileSize", absl::OkStatus())});
   std::unique_ptr<MockFileSystem> base_fs(
       new MockFileSystem(expected_fs_calls));
   RetryingFileSystem<MockFileSystem> fs(
@@ -631,7 +639,7 @@ TEST(RetryingFileSystemTest, GetFileSize_AllRetriesFailed) {
 TEST(RetryingFileSystemTest, RenameFile_SuccessWith2ndTry) {
   ExpectedCalls expected_fs_calls(
       {std::make_tuple("RenameFile", errors::Unavailable("Something is wrong")),
-       std::make_tuple("RenameFile", OkStatus())});
+       std::make_tuple("RenameFile", absl::OkStatus())});
   std::unique_ptr<MockFileSystem> base_fs(
       new MockFileSystem(expected_fs_calls));
   RetryingFileSystem<MockFileSystem> fs(
@@ -655,7 +663,7 @@ TEST(RetryingFileSystemTest, RenameFile_AllRetriesFailed) {
 TEST(RetryingFileSystemTest, Stat_SuccessWith2ndTry) {
   ExpectedCalls expected_fs_calls(
       {std::make_tuple("Stat", errors::Unavailable("Something is wrong")),
-       std::make_tuple("Stat", OkStatus())});
+       std::make_tuple("Stat", absl::OkStatus())});
   std::unique_ptr<MockFileSystem> base_fs(
       new MockFileSystem(expected_fs_calls));
   RetryingFileSystem<MockFileSystem> fs(
@@ -693,7 +701,7 @@ TEST(RetryingFileSystemTest, FileExists_AllRetriesFailed) {
 TEST(RetryingFileSystemTest, FileExists_SuccessWith2ndTry) {
   ExpectedCalls expected_fs_calls(
       {std::make_tuple("FileExists", errors::Unavailable("Something is wrong")),
-       std::make_tuple("FileExists", OkStatus())});
+       std::make_tuple("FileExists", absl::OkStatus())});
   std::unique_ptr<MockFileSystem> base_fs(
       new MockFileSystem(expected_fs_calls));
   RetryingFileSystem<MockFileSystem> fs(
@@ -706,7 +714,7 @@ TEST(RetryingFileSystemTest, IsDirectory_SuccessWith2ndTry) {
   ExpectedCalls expected_fs_calls(
       {std::make_tuple("IsDirectory",
                        errors::Unavailable("Something is wrong")),
-       std::make_tuple("IsDirectory", OkStatus())});
+       std::make_tuple("IsDirectory", absl::OkStatus())});
   std::unique_ptr<MockFileSystem> base_fs(
       new MockFileSystem(expected_fs_calls));
   RetryingFileSystem<MockFileSystem> fs(
@@ -731,7 +739,7 @@ TEST(RetryingFileSystemTest, DeleteRecursively_SuccessWith2ndTry) {
   ExpectedCalls expected_fs_calls(
       {std::make_tuple("DeleteRecursively",
                        errors::Unavailable("Something is wrong")),
-       std::make_tuple("DeleteRecursively", OkStatus())});
+       std::make_tuple("DeleteRecursively", absl::OkStatus())});
   std::unique_ptr<MockFileSystem> base_fs(
       new MockFileSystem(expected_fs_calls));
   RetryingFileSystem<MockFileSystem> fs(
diff --git a/third_party/xla/third_party/tsl/tsl/platform/retrying_utils.cc b/third_party/xla/third_party/tsl/tsl/platform/retrying_utils.cc
index 3beb3f46110a24..14459e93b61ef3 100644
--- a/third_party/xla/third_party/tsl/tsl/platform/retrying_utils.cc
+++ b/third_party/xla/third_party/tsl/tsl/platform/retrying_utils.cc
@@ -54,8 +54,8 @@ double GenerateUniformRandomNumberBetween(double a, double b) {
 
 }  // namespace
 
-Status RetryingUtils::CallWithRetries(const std::function<Status()>& f,
-                                      const RetryConfig& config) {
+absl::Status RetryingUtils::CallWithRetries(
+    const std::function<absl::Status()>& f, const RetryConfig& config) {
   return CallWithRetries(
       f,
       [](int64_t micros) {
@@ -64,8 +64,8 @@ Status RetryingUtils::CallWithRetries(const std::function<Status()>& f,
       config);
 }
 
-Status RetryingUtils::CallWithRetries(
-    const std::function<Status()>& f,
+absl::Status RetryingUtils::CallWithRetries(
+    const std::function<absl::Status()>& f,
     const std::function<void(int64_t)>& sleep_usec, const RetryConfig& config) {
   int retries = 0;
   while (true) {
@@ -76,7 +76,7 @@ Status RetryingUtils::CallWithRetries(
     if (retries >= config.max_retries) {
       // Return AbortedError, so that it doesn't get retried again somewhere
       // at a higher level.
-      return Status(
+      return absl::Status(
           absl::StatusCode::kAborted,
           strings::StrCat(
               "All ", config.max_retries,
@@ -98,14 +98,15 @@ Status RetryingUtils::CallWithRetries(
   }
 }
 
-Status RetryingUtils::DeleteWithRetries(
-    const std::function<Status()>& delete_func, const RetryConfig& config) {
+absl::Status RetryingUtils::DeleteWithRetries(
+    const std::function<absl::Status()>& delete_func,
+    const RetryConfig& config) {
   bool is_retried = false;
   return RetryingUtils::CallWithRetries(
       [delete_func, &is_retried]() {
-        const Status status = delete_func();
+        const absl::Status status = delete_func();
         if (is_retried && status.code() == error::NOT_FOUND) {
-          return OkStatus();
+          return absl::OkStatus();
         }
         is_retried = true;
         return status;
diff --git a/third_party/xla/third_party/tsl/tsl/platform/retrying_utils.h b/third_party/xla/third_party/tsl/tsl/platform/retrying_utils.h
index 3252da2637c4d2..470b6a8f183412 100644
--- a/third_party/xla/third_party/tsl/tsl/platform/retrying_utils.h
+++ b/third_party/xla/third_party/tsl/tsl/platform/retrying_utils.h
@@ -51,20 +51,22 @@ class RetryingUtils {
   /// If initial_delay_microseconds is zero, no delays will be made between
   /// retries.
   /// If all retries failed, returns the last error status.
-  static Status CallWithRetries(const std::function<Status()>& f,
-                                const RetryConfig& config);
+  static absl::Status CallWithRetries(const std::function<absl::Status()>& f,
+                                      const RetryConfig& config);
 
   /// sleep_usec is a function that sleeps for the given number of microseconds.
-  static Status CallWithRetries(const std::function<Status()>& f,
-                                const std::function<void(int64_t)>& sleep_usec,
-                                const RetryConfig& config);
+  static absl::Status CallWithRetries(
+      const std::function<absl::Status()>& f,
+      const std::function<void(int64_t)>& sleep_usec,
+      const RetryConfig& config);
   /// \brief A retrying wrapper for a function that deletes a resource.
   ///
   /// The function takes care of the scenario when a delete operation
   /// returns a failure but succeeds under the hood: if a retry returns
   /// NOT_FOUND, the whole operation is considered a success.
-  static Status DeleteWithRetries(const std::function<Status()>& delete_func,
-                                  const RetryConfig& config);
+  static absl::Status DeleteWithRetries(
+      const std::function<absl::Status()>& delete_func,
+      const RetryConfig& config);
 };
 
 // Given the total number of retries attempted, returns a randomized duration of
diff --git a/third_party/xla/third_party/tsl/tsl/platform/retrying_utils_test.cc b/third_party/xla/third_party/tsl/tsl/platform/retrying_utils_test.cc
index c0b3ad7b651e16..5d55ec31cc2f20 100644
--- a/third_party/xla/third_party/tsl/tsl/platform/retrying_utils_test.cc
+++ b/third_party/xla/third_party/tsl/tsl/platform/retrying_utils_test.cc
@@ -32,7 +32,9 @@ TEST(RetryingUtilsTest, CallWithRetries_RetryDelays) {
   std::function<void(int64_t)> sleep = [&requested_delays](int64_t delay) {
     requested_delays.emplace_back(delay / 1000000.0);
   };
-  std::function<Status()> f = []() { return errors::Unavailable("Failed."); };
+  std::function<absl::Status()> f = []() {
+    return errors::Unavailable("Failed.");
+  };
 
   const auto& status = RetryingUtils::CallWithRetries(
       f, sleep, RetryConfig(500000 /* init_delay_time_us */));
@@ -58,9 +60,9 @@ TEST(RetryingUtilsTest, CallWithRetries_RetryDelays) {
 }
 
 TEST(RetryingUtilsTest, CallWithRetries_NotFoundIsNotRetried) {
-  std::vector<Status> results(
+  std::vector<absl::Status> results(
       {errors::Unavailable("Failed."), errors::NotFound("Not found.")});
-  std::function<Status()> f = [&results]() {
+  std::function<absl::Status()> f = [&results]() {
     auto result = results[0];
     results.erase(results.begin());
     return result;
@@ -70,11 +72,11 @@ TEST(RetryingUtilsTest, CallWithRetries_NotFoundIsNotRetried) {
 }
 
 TEST(RetryingUtilsTest, CallWithRetries_ImmediateSuccess) {
-  std::vector<Status> results({OkStatus()});
+  std::vector<absl::Status> results({absl::OkStatus()});
   std::function<void(int64_t)> sleep = [](int64_t delay) {
     ADD_FAILURE() << "Unexpected call to sleep.";
   };
-  std::function<Status()> f = [&results]() {
+  std::function<absl::Status()> f = [&results]() {
     auto result = results[0];
     results.erase(results.begin());
     return result;
@@ -84,10 +86,10 @@ TEST(RetryingUtilsTest, CallWithRetries_ImmediateSuccess) {
 }
 
 TEST(RetryingUtilsTest, CallWithRetries_EventualSuccess) {
-  std::vector<Status> results({errors::Unavailable("Failed."),
-                               errors::Unavailable("Failed again."),
-                               OkStatus()});
-  std::function<Status()> f = [&results]() {
+  std::vector<absl::Status> results({errors::Unavailable("Failed."),
+                                     errors::Unavailable("Failed again."),
+                                     absl::OkStatus()});
+  std::function<absl::Status()> f = [&results]() {
     auto result = results[0];
     results.erase(results.begin());
     return result;
@@ -97,7 +99,7 @@ TEST(RetryingUtilsTest, CallWithRetries_EventualSuccess) {
 }
 
 TEST(RetryingUtilsTest, DeleteWithRetries_ImmediateSuccess) {
-  std::vector<Status> delete_results({OkStatus()});
+  std::vector<absl::Status> delete_results({absl::OkStatus()});
   const auto delete_func = [&delete_results]() {
     auto result = delete_results[0];
     delete_results.erase(delete_results.begin());
@@ -108,7 +110,8 @@ TEST(RetryingUtilsTest, DeleteWithRetries_ImmediateSuccess) {
 }
 
 TEST(RetryingUtilsTest, DeleteWithRetries_EventualSuccess) {
-  std::vector<Status> delete_results({errors::Unavailable(""), OkStatus()});
+  std::vector<absl::Status> delete_results(
+      {errors::Unavailable(""), absl::OkStatus()});
   const auto delete_func = [&delete_results]() {
     auto result = delete_results[0];
     delete_results.erase(delete_results.begin());
@@ -119,7 +122,7 @@ TEST(RetryingUtilsTest, DeleteWithRetries_EventualSuccess) {
 }
 
 TEST(RetryingUtilsTest, DeleteWithRetries_PermissionDeniedNotRetried) {
-  std::vector<Status> delete_results(
+  std::vector<absl::Status> delete_results(
       {errors::Unavailable(""), errors::PermissionDenied("")});
   const auto delete_func = [&delete_results]() {
     auto result = delete_results[0];
@@ -131,7 +134,7 @@ TEST(RetryingUtilsTest, DeleteWithRetries_PermissionDeniedNotRetried) {
 }
 
 TEST(RetryingUtilsTest, DeleteWithRetries_SuccessThroughFileNotFound) {
-  std::vector<Status> delete_results(
+  std::vector<absl::Status> delete_results(
       {errors::Unavailable(""), errors::NotFound("")});
   const auto delete_func = [&delete_results]() {
     auto result = delete_results[0];
@@ -143,7 +146,7 @@ TEST(RetryingUtilsTest, DeleteWithRetries_SuccessThroughFileNotFound) {
 }
 
 TEST(RetryingUtilsTest, DeleteWithRetries_FirstNotFoundReturnedAsIs) {
-  std::vector<Status> delete_results({errors::NotFound("")});
+  std::vector<absl::Status> delete_results({errors::NotFound("")});
   const auto delete_func = [&delete_results]() {
     auto result = delete_results[0];
     delete_results.erase(delete_results.begin());
diff --git a/third_party/xla/third_party/tsl/tsl/platform/status.cc b/third_party/xla/third_party/tsl/tsl/platform/status.cc
index 3f9dd1ecb58427..4256bd2f2a44e9 100644
--- a/third_party/xla/third_party/tsl/tsl/platform/status.cc
+++ b/third_party/xla/third_party/tsl/tsl/platform/status.cc
@@ -112,7 +112,7 @@ namespace errors {
 static constexpr const char kStackTraceProtoUrl[] =
     "type.googleapis.com/tensorflow.StackTracePayload";
 
-void SetStackTrace(::tsl::Status& status, std::vector<StackFrame> stack_trace) {
+void SetStackTrace(absl::Status& status, std::vector<StackFrame> stack_trace) {
   // Given the StackFrame fields are (a) line number (b) filename (c) function
   // name, we can safely assume that there is no `\n` in there.
   // Thus, we can serialize as strings using a simple new line delimiter.
@@ -134,7 +134,7 @@ void SetStackTrace(::tsl::Status& status, std::vector<StackFrame> stack_trace) {
                     absl::Cord(absl::StrJoin(items, "\n")));
 }
 
-std::vector<StackFrame> GetStackTrace(const ::tsl::Status& status) {
+std::vector<StackFrame> GetStackTrace(const absl::Status& status) {
   std::vector<StackFrame> stack_trace;
   absl::optional<absl::Cord> maybe_serialized_payload =
       status.GetPayload(kStackTraceProtoUrl);
@@ -163,7 +163,7 @@ const char* NullTerminatedMessage(const absl::Status& status) {
 }
 #endif
 
-std::string* TfCheckOpHelperOutOfLine(const ::tsl::Status& v, const char* msg) {
+std::string* TfCheckOpHelperOutOfLine(const absl::Status& v, const char* msg) {
   std::string r("Non-OK-status: ");
   r += msg;
   r += " status: ";
@@ -174,8 +174,8 @@ std::string* TfCheckOpHelperOutOfLine(const ::tsl::Status& v, const char* msg) {
 
 StatusGroup::StatusGroup() {}
 
-StatusGroup::StatusGroup(std::initializer_list<Status> statuses) {
-  for (const Status& s : statuses) {
+StatusGroup::StatusGroup(std::initializer_list<absl::Status> statuses) {
+  for (const absl::Status& s : statuses) {
     Update(s);
   }
 }
@@ -183,11 +183,11 @@ StatusGroup::StatusGroup(std::initializer_list<Status> statuses) {
 static constexpr const char kDerivedStatusProtoUrl[] =
     "type.googleapis.com/tensorflow.DerivedStatus";
 
-Status StatusGroup::MakeDerived(const Status& s) {
+absl::Status StatusGroup::MakeDerived(const absl::Status& s) {
   if (IsDerived(s)) {
     return s;
   } else {
-    Status derived(s);
+    absl::Status derived(s);
     // TODO(b/200167936): Serialize an instance of DerivedStatus proto instead
     // of using the string directly. The string is never used so it is not
     // causing any issues at the moment.
@@ -196,7 +196,7 @@ Status StatusGroup::MakeDerived(const Status& s) {
   }
 }
 
-bool StatusGroup::IsDerived(const Status& s) {
+bool StatusGroup::IsDerived(const absl::Status& s) {
   return s.GetPayload(kDerivedStatusProtoUrl).has_value();
 }
 
@@ -204,7 +204,7 @@ void StatusGroup::ConfigureLogHistory() {
   StatusLogSink::GetInstance()->enable();
 }
 
-void StatusGroup::Update(const Status& s) {
+void StatusGroup::Update(const absl::Status& s) {
   if (s.ok()) {
     ++num_ok_;
   } else {
@@ -241,25 +241,26 @@ std::unordered_map<std::string, absl::Cord> StatusGroup::GetPayloads() const {
   return payloads;
 }
 
-Status MakeStatus(absl::StatusCode code, absl::string_view message,
-                  const std::unordered_map<std::string, absl::Cord>& payloads) {
-  Status status(code, message);
+absl::Status MakeStatus(
+    absl::StatusCode code, absl::string_view message,
+    const std::unordered_map<std::string, absl::Cord>& payloads) {
+  absl::Status status(code, message);
   for (const auto& payload : payloads) {
     status.SetPayload(payload.first, payload.second);
   }
   return status;
 }
 
-std::string MakeString(const Status& status) {
+std::string MakeString(const absl::Status& status) {
   return absl::StrCat(absl::StatusCodeToString(status.code()), ": ",
                       status.message());
 }
 
 // Summarize all the status objects in the StatusGroup. This is used when
 // individual Status objects in the StatusGroup are not already summarized.
-Status StatusGroup::as_summary_status() const {
+absl::Status StatusGroup::as_summary_status() const {
   if (ok_) {
-    return OkStatus();
+    return absl::OkStatus();
   }
 
   // Gather recent logs as a string
@@ -322,9 +323,9 @@ Status StatusGroup::as_summary_status() const {
 
 // Concatenate all the status objects in the StatusGroup. This is used when
 // individual Status objects in the StatusGroup are already summarized Status.
-Status StatusGroup::as_concatenated_status() const {
+absl::Status StatusGroup::as_concatenated_status() const {
   if (ok_) {
-    return OkStatus();
+    return absl::OkStatus();
   }
 
   // If only one root status is found, return it directly.
diff --git a/third_party/xla/third_party/tsl/tsl/platform/status.h b/third_party/xla/third_party/tsl/tsl/platform/status.h
index b9aca1e9285a11..84954ff485a48b 100644
--- a/third_party/xla/third_party/tsl/tsl/platform/status.h
+++ b/third_party/xla/third_party/tsl/tsl/platform/status.h
@@ -133,9 +133,9 @@ inline const char* NullTerminatedMessage(const absl::Status& status) {
 // TODO(b/197552541) Move this namespace to errors.h.
 namespace errors {
 
-void SetStackTrace(::tsl::Status& status, std::vector<StackFrame> stack_trace);
+void SetStackTrace(absl::Status& status, std::vector<StackFrame> stack_trace);
 
-std::vector<StackFrame> GetStackTrace(const ::tsl::Status& status);
+std::vector<StackFrame> GetStackTrace(const absl::Status& status);
 }  // namespace errors
 
 // Helper class to manage multiple child status values.
@@ -144,12 +144,12 @@ class StatusGroup {
   StatusGroup();
   // Constructor to form a StatusGroup from any N set of Status arguments.
   // Usage: StatusGroup({status_a, status_b, status_c});
-  StatusGroup(std::initializer_list<Status> statuses);
+  StatusGroup(std::initializer_list<absl::Status> statuses);
 
   // Utility function to mark a Status as derived. By marking derived status,
   // Derived status messages are ignored when reporting errors to end users.
-  static Status MakeDerived(const Status& s);
-  static bool IsDerived(const Status& s);
+  static absl::Status MakeDerived(const absl::Status& s);
+  static bool IsDerived(const absl::Status& s);
 
   // Enable warning and error log collection for appending to the aggregated
   // status. This function may be called more than once.
@@ -163,15 +163,15 @@ class StatusGroup {
   std::unordered_map<std::string, absl::Cord> GetPayloads() const;
 
   // Return a merged status with combined child status messages with a summary.
-  Status as_summary_status() const;
+  absl::Status as_summary_status() const;
   // Return a merged status with combined child status messages with
   // concatenation.
-  Status as_concatenated_status() const;
+  absl::Status as_concatenated_status() const;
 
   bool ok() const { return ok_; }
 
   // Augment this group with the child status `status`.
-  void Update(const Status& status);
+  void Update(const absl::Status& status);
 
   // Attach recent warning and error log messages
   void AttachLogMessages();
@@ -183,25 +183,24 @@ class StatusGroup {
 
   // Maintain a sorted collection of statuses.
   struct CompareStatus {
-    bool operator()(const Status& a, const Status& b) const {
+    bool operator()(const absl::Status& a, const absl::Status& b) const {
       return a.ToString() > b.ToString();
     }
   };
   // Using std::set instead of absl::btree_set to keep size for certain
   // dependent libraries under the limit.
-  std::set<Status, CompareStatus> derived_;
-  std::set<Status, CompareStatus> non_derived_;
+  std::set<absl::Status, CompareStatus> derived_;
+  std::set<absl::Status, CompareStatus> non_derived_;
 
   std::vector<std::string> recent_logs_;  // recent warning and error logs
 };
 
+typedef std::function<void(const absl::Status&)> StatusCallback;
 
-typedef std::function<void(const Status&)> StatusCallback;
-
-extern ::tsl::string* TfCheckOpHelperOutOfLine(const ::tsl::Status& v,
+extern ::tsl::string* TfCheckOpHelperOutOfLine(const absl::Status& v,
                                                const char* msg);
 
-inline ::tsl::string* TfCheckOpHelper(::tsl::Status v, const char* msg) {
+inline ::tsl::string* TfCheckOpHelper(absl::Status v, const char* msg) {
   if (v.ok()) return nullptr;
   return TfCheckOpHelperOutOfLine(v, msg);
 }
diff --git a/third_party/xla/third_party/tsl/tsl/platform/status_matchers.cc b/third_party/xla/third_party/tsl/tsl/platform/status_matchers.cc
index 7546a7b20770ed..77422d564f8cda 100644
--- a/third_party/xla/third_party/tsl/tsl/platform/status_matchers.cc
+++ b/third_party/xla/third_party/tsl/tsl/platform/status_matchers.cc
@@ -40,7 +40,7 @@ void StatusIsMatcherCommonImpl::DescribeNegationTo(std::ostream* os) const {
 }
 
 bool StatusIsMatcherCommonImpl::MatchAndExplain(
-    const Status& status,
+    const absl::Status& status,
     ::testing::MatchResultListener* result_listener) const {
   ::testing::StringMatchResultListener inner_listener;
 
diff --git a/third_party/xla/third_party/tsl/tsl/platform/status_matchers.h b/third_party/xla/third_party/tsl/tsl/platform/status_matchers.h
index ee2144dca8a698..cb8eba40783093 100644
--- a/third_party/xla/third_party/tsl/tsl/platform/status_matchers.h
+++ b/third_party/xla/third_party/tsl/tsl/platform/status_matchers.h
@@ -115,10 +115,12 @@ void PrintTo(const StatusOr<T>& status_or, std::ostream* os) {
 namespace testing {
 namespace internal_status {
 
-inline const Status& GetStatus(const Status& status) { return status; }
+inline const absl::Status& GetStatus(const absl::Status& status) {
+  return status;
+}
 
 template <typename T>
-inline const Status& GetStatus(const StatusOr<T>& status) {
+inline const absl::Status& GetStatus(const StatusOr<T>& status) {
   return status.status();
 }
 
@@ -211,7 +213,7 @@ class StatusIsMatcherCommonImpl {
 
   void DescribeNegationTo(std::ostream* os) const;
 
-  bool MatchAndExplain(const Status& status,
+  bool MatchAndExplain(const absl::Status& status,
                        ::testing::MatchResultListener* result_listener) const;
 
  private:
diff --git a/third_party/xla/third_party/tsl/tsl/platform/status_matchers_test.cc b/third_party/xla/third_party/tsl/tsl/platform/status_matchers_test.cc
index 70fc6191167c3b..ea0e0d489c24bf 100644
--- a/third_party/xla/third_party/tsl/tsl/platform/status_matchers_test.cc
+++ b/third_party/xla/third_party/tsl/tsl/platform/status_matchers_test.cc
@@ -76,13 +76,13 @@ std::string ExplainMatch(const Matcher<T>& matcher, const V& value) {
 }
 
 TEST(IsOkAndHoldsTest, MatchesValue) {
-  StatusOr<std::string> status_or_message("Hello, world");
+  absl::StatusOr<std::string> status_or_message("Hello, world");
   EXPECT_THAT(status_or_message, IsOkAndHolds("Hello, world"));
   EXPECT_THAT(status_or_message, IsOkAndHolds(HasSubstr("Hello,")));
 }
 
 TEST(IsOkAndHoldsTest, MatchesContainer) {
-  StatusOr<std::vector<std::string>> status_or_messages =
+  absl::StatusOr<std::vector<std::string>> status_or_messages =
       std::vector<std::string>{"Hello, world", "Hello, tf"};
   EXPECT_THAT(status_or_messages,
               IsOkAndHolds(ElementsAre("Hello, world", "Hello, tf")));
@@ -91,23 +91,23 @@ TEST(IsOkAndHoldsTest, MatchesContainer) {
 }
 
 TEST(IsOkAndHoldsTest, DoesNotMatchStatus) {
-  StatusOr<std::string> status_or_message =
+  absl::StatusOr<std::string> status_or_message =
       errors::InvalidArgument("Invalid argument");
   EXPECT_THAT(status_or_message, Not(IsOkAndHolds("Hello, world")));
 }
 
 TEST(IsOkAndHoldsTest, DoesNotMatchValue) {
-  StatusOr<std::string> status_or_message("Hello, tf");
+  absl::StatusOr<std::string> status_or_message("Hello, tf");
   EXPECT_THAT(status_or_message, Not(IsOkAndHolds("Hello, world")));
 }
 
 TEST(IsOkAndHoldsTest, DoesNotMatchContainer) {
-  StatusOr<std::vector<int>> status_or_container({1, 2, 3});
+  absl::StatusOr<std::vector<int>> status_or_container({1, 2, 3});
   EXPECT_THAT(status_or_container, Not(IsOkAndHolds(ElementsAre(4, 5, 6))));
 }
 
 TEST(IsOkAndHoldsTest, DescribeExpectedValue) {
-  Matcher<StatusOr<std::string>> is_ok_and_has_substr =
+  Matcher<absl::StatusOr<std::string>> is_ok_and_has_substr =
       IsOkAndHolds(HasSubstr("Hello"));
   EXPECT_EQ(Describe(is_ok_and_has_substr),
             "is OK and has a value that has substring \"Hello\"");
@@ -116,20 +116,22 @@ TEST(IsOkAndHoldsTest, DescribeExpectedValue) {
 }
 
 TEST(IsOkAndHoldsTest, ExplainNotMatchingStatus) {
-  Matcher<StatusOr<int>> is_ok_and_less_than = IsOkAndHolds(LessThan(100));
-  StatusOr<int> status = errors::Unknown("Unknown");
+  Matcher<absl::StatusOr<int>> is_ok_and_less_than =
+      IsOkAndHolds(LessThan(100));
+  absl::StatusOr<int> status = errors::Unknown("Unknown");
   EXPECT_THAT(ExplainMatch(is_ok_and_less_than, status),
               HasSubstr("which has status UNKNOWN: Unknown"));
 }
 
 TEST(IsOkAndHoldsTest, ExplainNotMatchingValue) {
-  Matcher<StatusOr<int>> is_ok_and_less_than = IsOkAndHolds(LessThan(100));
+  Matcher<absl::StatusOr<int>> is_ok_and_less_than =
+      IsOkAndHolds(LessThan(100));
   EXPECT_EQ(ExplainMatch(is_ok_and_less_than, 120),
             "which contains value 120, which is 20 more than 100");
 }
 
 TEST(IsOkAndHoldsTest, ExplainNotMatchingContainer) {
-  Matcher<StatusOr<std::vector<int>>> is_ok_and_less_than =
+  Matcher<absl::StatusOr<std::vector<int>>> is_ok_and_less_than =
       IsOkAndHolds(ElementsAre(1, 2, 3));
   std::vector<int> actual{4, 5, 6};
   EXPECT_THAT(ExplainMatch(is_ok_and_less_than, actual),
@@ -137,20 +139,20 @@ TEST(IsOkAndHoldsTest, ExplainNotMatchingContainer) {
 }
 
 TEST(StatusIsTest, MatchesOK) {
-  EXPECT_THAT(OkStatus(), StatusIs(error::OK));
-  StatusOr<std::string> message("Hello, world");
+  EXPECT_THAT(absl::OkStatus(), StatusIs(error::OK));
+  absl::StatusOr<std::string> message("Hello, world");
   EXPECT_THAT(message, StatusIs(error::OK));
 }
 
 TEST(StatusIsTest, DoesNotMatchOk) {
   EXPECT_THAT(errors::DeadlineExceeded("Deadline exceeded"),
               Not(StatusIs(error::OK)));
-  StatusOr<std::string> status = errors::NotFound("Not found");
+  absl::StatusOr<std::string> status = errors::NotFound("Not found");
   EXPECT_THAT(status, Not(StatusIs(error::OK)));
 }
 
 TEST(StatusIsTest, MatchesStatus) {
-  Status s = errors::Cancelled("Cancelled");
+  absl::Status s = errors::Cancelled("Cancelled");
   EXPECT_THAT(s, StatusIs(error::CANCELLED));
   EXPECT_THAT(s, StatusIs(error::CANCELLED, "Cancelled"));
   EXPECT_THAT(s, StatusIs(_, "Cancelled"));
@@ -161,7 +163,7 @@ TEST(StatusIsTest, MatchesStatus) {
 }
 
 TEST(StatusIsTest, StatusOrMatchesStatus) {
-  StatusOr<int> s = errors::InvalidArgument("Invalid Argument");
+  absl::StatusOr<int> s = errors::InvalidArgument("Invalid Argument");
   EXPECT_THAT(s, StatusIs(error::INVALID_ARGUMENT));
   EXPECT_THAT(s, StatusIs(error::INVALID_ARGUMENT, "Invalid Argument"));
   EXPECT_THAT(s, StatusIs(_, "Invalid Argument"));
@@ -172,7 +174,7 @@ TEST(StatusIsTest, StatusOrMatchesStatus) {
 }
 
 TEST(StatusIsTest, DoesNotMatchStatus) {
-  Status s = errors::Internal("Internal");
+  absl::Status s = errors::Internal("Internal");
   EXPECT_THAT(s, Not(StatusIs(error::FAILED_PRECONDITION)));
   EXPECT_THAT(s, Not(StatusIs(error::INTERNAL, "Failed Precondition")));
   EXPECT_THAT(s, Not(StatusIs(_, "Failed Precondition")));
@@ -180,7 +182,7 @@ TEST(StatusIsTest, DoesNotMatchStatus) {
 }
 
 TEST(StatusIsTest, StatusOrDoesNotMatchStatus) {
-  StatusOr<int> s = errors::FailedPrecondition("Failed Precondition");
+  absl::StatusOr<int> s = errors::FailedPrecondition("Failed Precondition");
   EXPECT_THAT(s, Not(StatusIs(error::INTERNAL)));
   EXPECT_THAT(s, Not(StatusIs(error::FAILED_PRECONDITION, "Internal")));
   EXPECT_THAT(s, Not(StatusIs(_, "Internal")));
@@ -188,7 +190,7 @@ TEST(StatusIsTest, StatusOrDoesNotMatchStatus) {
 }
 
 TEST(StatusIsTest, DescribeExpectedValue) {
-  Matcher<Status> status_is =
+  Matcher<absl::Status> status_is =
       StatusIs(error::UNAVAILABLE, std::string("Unavailable"));
   EXPECT_EQ(Describe(status_is),
             "has a status code that is equal to UNAVAILABLE, "
@@ -196,7 +198,7 @@ TEST(StatusIsTest, DescribeExpectedValue) {
 }
 
 TEST(StatusIsTest, DescribeNegatedExpectedValue) {
-  Matcher<StatusOr<std::string>> status_is =
+  Matcher<absl::StatusOr<std::string>> status_is =
       StatusIs(error::ABORTED, std::string("Aborted"));
   EXPECT_EQ(DescribeNegation(status_is),
             "has a status code that isn't equal to ABORTED, "
@@ -204,60 +206,61 @@ TEST(StatusIsTest, DescribeNegatedExpectedValue) {
 }
 
 TEST(StatusIsTest, ExplainNotMatchingErrorCode) {
-  Matcher<Status> status_is = StatusIs(error::NOT_FOUND, _);
-  const Status status = errors::AlreadyExists("Already exists");
+  Matcher<absl::Status> status_is = StatusIs(error::NOT_FOUND, _);
+  const absl::Status status = errors::AlreadyExists("Already exists");
   EXPECT_EQ(ExplainMatch(status_is, status), "whose status code is wrong");
 }
 
 TEST(StatusIsTest, ExplainNotMatchingErrorMessage) {
-  Matcher<Status> status_is = StatusIs(error::NOT_FOUND, "Not found");
-  const Status status = errors::NotFound("Already exists");
+  Matcher<absl::Status> status_is = StatusIs(error::NOT_FOUND, "Not found");
+  const absl::Status status = errors::NotFound("Already exists");
   EXPECT_EQ(ExplainMatch(status_is, status), "whose error message is wrong");
 }
 
 TEST(StatusIsTest, ExplainStatusOrNotMatchingErrorCode) {
-  Matcher<StatusOr<int>> status_is = StatusIs(error::ALREADY_EXISTS, _);
-  const StatusOr<int> status_or = errors::NotFound("Not found");
+  Matcher<absl::StatusOr<int>> status_is = StatusIs(error::ALREADY_EXISTS, _);
+  const absl::StatusOr<int> status_or = errors::NotFound("Not found");
   EXPECT_EQ(ExplainMatch(status_is, status_or), "whose status code is wrong");
 }
 
 TEST(StatusIsTest, ExplainStatusOrNotMatchingErrorMessage) {
-  Matcher<StatusOr<int>> status_is =
+  Matcher<absl::StatusOr<int>> status_is =
       StatusIs(error::ALREADY_EXISTS, "Already exists");
-  const StatusOr<int> status_or = errors::AlreadyExists("Not found");
+  const absl::StatusOr<int> status_or = errors::AlreadyExists("Not found");
   EXPECT_EQ(ExplainMatch(status_is, status_or), "whose error message is wrong");
 }
 
 TEST(StatusIsTest, ExplainStatusOrHasValue) {
-  Matcher<StatusOr<int>> status_is =
+  Matcher<absl::StatusOr<int>> status_is =
       StatusIs(error::RESOURCE_EXHAUSTED, "Resource exhausted");
-  const StatusOr<int> value = -1;
+  const absl::StatusOr<int> value = -1;
   EXPECT_EQ(ExplainMatch(status_is, value), "whose status code is wrong");
 }
 
 TEST(IsOkTest, MatchesOK) {
-  EXPECT_THAT(OkStatus(), IsOk());
-  StatusOr<std::string> message = std::string("Hello, world");
+  EXPECT_THAT(absl::OkStatus(), IsOk());
+  absl::StatusOr<std::string> message = std::string("Hello, world");
   EXPECT_THAT(message, IsOk());
 }
 
 TEST(IsOkTest, DoesNotMatchOK) {
   EXPECT_THAT(errors::PermissionDenied("Permission denied"), Not(IsOk()));
-  StatusOr<std::string> status = errors::Unauthenticated("Unauthenticated");
+  absl::StatusOr<std::string> status =
+      errors::Unauthenticated("Unauthenticated");
   EXPECT_THAT(status, Not(IsOk()));
 }
 
 TEST(IsOkTest, DescribeExpectedValue) {
-  Matcher<Status> status_is_ok = IsOk();
+  Matcher<absl::Status> status_is_ok = IsOk();
   EXPECT_EQ(Describe(status_is_ok), "is OK");
-  Matcher<StatusOr<std::string>> status_or_is_ok = IsOk();
+  Matcher<absl::StatusOr<std::string>> status_or_is_ok = IsOk();
   EXPECT_EQ(Describe(status_or_is_ok), "is OK");
 }
 
 TEST(IsOkTest, DescribeNegatedExpectedValue) {
-  Matcher<Status> status_is_ok = IsOk();
+  Matcher<absl::Status> status_is_ok = IsOk();
   EXPECT_EQ(DescribeNegation(status_is_ok), "is not OK");
-  Matcher<StatusOr<std::string>> status_or_is_ok = IsOk();
+  Matcher<absl::StatusOr<std::string>> status_or_is_ok = IsOk();
   EXPECT_EQ(DescribeNegation(status_or_is_ok), "is not OK");
 }
 
diff --git a/third_party/xla/third_party/tsl/tsl/platform/status_test.cc b/third_party/xla/third_party/tsl/tsl/platform/status_test.cc
index bccf969e63b662..b95de35e181be1 100644
--- a/third_party/xla/third_party/tsl/tsl/platform/status_test.cc
+++ b/third_party/xla/third_party/tsl/tsl/platform/status_test.cc
@@ -34,7 +34,7 @@ using ::tsl::testing::IsOk;
 using ::tsl::testing::StatusIs;
 
 TEST(ToStringTest, PayloadsArePrinted) {
-  Status status = errors::Aborted("Aborted Error Message");
+  absl::Status status = errors::Aborted("Aborted Error Message");
   status.SetPayload("payload_key", absl::Cord(absl::StrFormat(
                                        "payload_value %c%c%c", 1, 2, 3)));
 
@@ -44,7 +44,7 @@ TEST(ToStringTest, PayloadsArePrinted) {
 }
 
 TEST(ToStringTest, MatchesAbslStatus) {
-  Status status = errors::Aborted("Aborted Error Message");
+  absl::Status status = errors::Aborted("Aborted Error Message");
   status.SetPayload("payload_key", absl::Cord(absl::StrFormat(
                                        "payload_value %c%c%c", 1, 2, 3)));
 
@@ -57,7 +57,7 @@ TEST(ToStringTest, MatchesAbslStatus) {
 }
 
 TEST(StackTrace, SerializeAndDeserializeCorrectly) {
-  Status status = errors::Aborted("Aborted Error Message");
+  absl::Status status = errors::Aborted("Aborted Error Message");
   std::vector<StackFrame> stack_trace;
   stack_trace.push_back(StackFrame("filename_1", 33, "func_name_1"));
   stack_trace.push_back(StackFrame("filename_2", 66, "func_name_2"));
@@ -72,11 +72,11 @@ TEST(StackTrace, SerializeAndDeserializeCorrectly) {
 }
 
 TEST(StatusGroupTest, DeterministicOrderWithoutPayloads) {
-  Status status_a = errors::Aborted("Status A");
-  Status status_b = errors::Aborted("Status B");
-  Status status_c = errors::Aborted("Status C");
+  absl::Status status_a = errors::Aborted("Status A");
+  absl::Status status_b = errors::Aborted("Status B");
+  absl::Status status_c = errors::Aborted("Status C");
 
-  Status combined =
+  absl::Status combined =
       StatusGroup({status_a, status_b, status_c}).as_summary_status();
 
   EXPECT_EQ(combined,
@@ -94,14 +94,14 @@ TEST(StatusGroupTest, DeterministicOrderWithoutPayloads) {
 }
 
 TEST(StatusGroupTest, DeterministicOrderWithPayloads) {
-  Status status_a = errors::Aborted("Status A");
+  absl::Status status_a = errors::Aborted("Status A");
   status_a.SetPayload("payload_key", absl::Cord("payload_value_a"));
-  Status status_b = errors::Aborted("Status B");
+  absl::Status status_b = errors::Aborted("Status B");
   status_b.SetPayload("payload_key", absl::Cord("payload_value_b"));
-  Status status_c = errors::Aborted("Status C");
+  absl::Status status_c = errors::Aborted("Status C");
   status_c.SetPayload("payload_key", absl::Cord("payload_value_c"));
 
-  Status combined =
+  absl::Status combined =
       StatusGroup({status_a, status_b, status_c}).as_summary_status();
   ASSERT_TRUE(combined.GetPayload("payload_key").has_value());
   std::string payload(combined.GetPayload("payload_key").value());
@@ -127,16 +127,16 @@ TEST(StatusGroupTest, DeterministicOrderWithPayloads) {
 }
 
 TEST(StatusGroupTest, PayloadsMergedProperly) {
-  Status status_a = errors::Aborted("Status A");
+  absl::Status status_a = errors::Aborted("Status A");
   status_a.SetPayload("payload_key_a",
                       absl::Cord(std::string("payload_value_a")));
-  Status status_b = errors::Aborted("Status B");
+  absl::Status status_b = errors::Aborted("Status B");
   status_b.SetPayload("payload_key_b",
                       absl::Cord(std::string("payload_value_b")));
-  Status status_c = errors::Aborted("Status C");
+  absl::Status status_c = errors::Aborted("Status C");
   status_c.SetPayload("payload_key_c",
                       absl::Cord(std::string("payload_value_c")));
-  Status derived_status_c =
+  absl::Status derived_status_c =
       StatusGroup::MakeDerived(errors::Aborted("Status C"));
   derived_status_c.SetPayload(
       "payload_key_c", absl::Cord(std::string("derived_payload_value_c")));
@@ -144,14 +144,14 @@ TEST(StatusGroupTest, PayloadsMergedProperly) {
   StatusGroup status_group({status_a, status_b, status_c, derived_status_c});
   EXPECT_THAT(status_group.GetPayloads(), ::testing::SizeIs(3));
 
-  Status combined = status_group.as_summary_status();
+  absl::Status combined = status_group.as_summary_status();
   EXPECT_EQ(combined.GetPayload("payload_key_a"), "payload_value_a");
   EXPECT_EQ(combined.GetPayload("payload_key_b"), "payload_value_b");
   EXPECT_EQ(combined.GetPayload("payload_key_c"), "payload_value_c");
 }
 
 TEST(Status, ErrorStatusForEachPayloadIteratesOverAll) {
-  Status s(absl::StatusCode::kInternal, "Error message");
+  absl::Status s(absl::StatusCode::kInternal, "Error message");
   s.SetPayload("key1", absl::Cord("value1"));
   s.SetPayload("key2", absl::Cord("value2"));
   s.SetPayload("key3", absl::Cord("value3"));
@@ -168,7 +168,7 @@ TEST(Status, ErrorStatusForEachPayloadIteratesOverAll) {
 }
 
 TEST(Status, OkStatusForEachPayloadNoIteration) {
-  Status s = OkStatus();
+  absl::Status s = absl::OkStatus();
   s.SetPayload("key1", absl::Cord("value1"));
   s.SetPayload("key2", absl::Cord("value2"));
   s.SetPayload("key3", absl::Cord("value3"));
@@ -182,7 +182,7 @@ TEST(Status, OkStatusForEachPayloadNoIteration) {
 }
 
 TEST(Status, SaveOKStatusToProto) {
-  tensorflow::StatusProto status_proto = StatusToProto(OkStatus());
+  tensorflow::StatusProto status_proto = StatusToProto(absl::OkStatus());
   EXPECT_EQ(status_proto.code(), error::OK);
   EXPECT_THAT(status_proto.message(), IsEmpty());
 }
@@ -195,7 +195,7 @@ TEST(Status, SaveErrorStatusToProto) {
 }
 
 TEST(Status, SaveEmptyStatusToProto) {
-  tensorflow::StatusProto status_proto = StatusToProto(Status());
+  tensorflow::StatusProto status_proto = StatusToProto(absl::Status());
   EXPECT_EQ(status_proto.code(), error::OK);
   EXPECT_THAT(status_proto.message(), IsEmpty());
 }
diff --git a/third_party/xla/third_party/tsl/tsl/platform/status_to_from_proto.cc b/third_party/xla/third_party/tsl/tsl/platform/status_to_from_proto.cc
index 250e2e1ed4f670..96ad290f92c71a 100644
--- a/third_party/xla/third_party/tsl/tsl/platform/status_to_from_proto.cc
+++ b/third_party/xla/third_party/tsl/tsl/platform/status_to_from_proto.cc
@@ -22,7 +22,7 @@ limitations under the License.
 
 namespace tsl {
 
-tensorflow::StatusProto StatusToProto(const Status& s) {
+tensorflow::StatusProto StatusToProto(const absl::Status& s) {
   tensorflow::StatusProto status_proto;
   if (s.ok()) {
     return status_proto;
@@ -36,13 +36,13 @@ tensorflow::StatusProto StatusToProto(const Status& s) {
 }
 
 #if defined(PLATFORM_GOOGLE)
-Status StatusFromProto(const tensorflow::StatusProto& proto,
-                       absl::SourceLocation loc) {
+absl::Status StatusFromProto(const tensorflow::StatusProto& proto,
+                             absl::SourceLocation loc) {
   if (proto.code() == tensorflow::error::OK) {
-    return OkStatus();
+    return absl::OkStatus();
   }
-  return Status(static_cast<absl::StatusCode>(proto.code()), proto.message(),
-                loc);
+  return absl::Status(static_cast<absl::StatusCode>(proto.code()),
+                      proto.message(), loc);
 }
 #else
 Status StatusFromProto(const tensorflow::StatusProto& proto) {
diff --git a/third_party/xla/third_party/tsl/tsl/platform/status_to_from_proto.h b/third_party/xla/third_party/tsl/tsl/platform/status_to_from_proto.h
index 6abbe78dc0ef69..9891737f08159c 100644
--- a/third_party/xla/third_party/tsl/tsl/platform/status_to_from_proto.h
+++ b/third_party/xla/third_party/tsl/tsl/platform/status_to_from_proto.h
@@ -28,11 +28,11 @@ namespace tsl {
 // Symbol not found: tensorflow11StatusProtoC1EPN6protobuf5ArenaEb
 
 // Converts a `Status` to a `StatusProto`.
-tensorflow::StatusProto StatusToProto(const Status& s);
+tensorflow::StatusProto StatusToProto(const absl::Status& s);
 
 #if defined(PLATFORM_GOOGLE)
 // Constructs a `Status` from a `StatusProto`.
-Status StatusFromProto(
+absl::Status StatusFromProto(
     const tensorflow::StatusProto& proto,
     absl::SourceLocation loc = absl::SourceLocation::current());
 #else
diff --git a/third_party/xla/third_party/tsl/tsl/platform/statusor_test.cc b/third_party/xla/third_party/tsl/tsl/platform/statusor_test.cc
index 9cfc44de7167ed..fd0ee7886073b4 100644
--- a/third_party/xla/third_party/tsl/tsl/platform/statusor_test.cc
+++ b/third_party/xla/third_party/tsl/tsl/platform/statusor_test.cc
@@ -67,7 +67,7 @@ class NoDefaultConstructor {
 static_assert(!std::is_default_constructible<NoDefaultConstructor>(),
               "Should not be default-constructible.");
 
-StatusOr<std::unique_ptr<int>> ReturnUniquePtr() {
+absl::StatusOr<std::unique_ptr<int>> ReturnUniquePtr() {
   // Uses implicit constructor from T&&
   return std::unique_ptr<int>(new int(0));
 }
@@ -75,25 +75,25 @@ StatusOr<std::unique_ptr<int>> ReturnUniquePtr() {
 TEST(StatusOr, NullPointerStatusOr) {
   // As a very special case, null-plain-pointer StatusOr used to be an
   // error. Test that it no longer is.
-  StatusOr<int*> null_status(nullptr);
+  absl::StatusOr<int*> null_status(nullptr);
   EXPECT_TRUE(null_status.ok());
   EXPECT_EQ(null_status.value(), nullptr);
 }
 
 TEST(StatusOr, TestNoDefaultConstructorInitialization) {
   // Explicitly initialize it with an error code.
-  StatusOr<NoDefaultConstructor> statusor(errors::Cancelled(""));
+  absl::StatusOr<NoDefaultConstructor> statusor(errors::Cancelled(""));
   EXPECT_FALSE(statusor.ok());
   EXPECT_EQ(statusor.status().code(), absl::StatusCode::kCancelled);
 
   // Default construction of StatusOr initializes it with an UNKNOWN error code.
-  StatusOr<NoDefaultConstructor> statusor2;
+  absl::StatusOr<NoDefaultConstructor> statusor2;
   EXPECT_FALSE(statusor2.ok());
   EXPECT_EQ(statusor2.status().code(), absl::StatusCode::kUnknown);
 }
 
 TEST(StatusOr, TestMoveOnlyInitialization) {
-  StatusOr<std::unique_ptr<int>> thing(ReturnUniquePtr());
+  absl::StatusOr<std::unique_ptr<int>> thing(ReturnUniquePtr());
   ASSERT_TRUE(thing.ok());
   EXPECT_EQ(0, *thing.value());
   int* previous = thing.value().get();
@@ -105,12 +105,12 @@ TEST(StatusOr, TestMoveOnlyInitialization) {
 }
 
 TEST(StatusOr, TestMoveOnlyStatusCtr) {
-  StatusOr<std::unique_ptr<int>> thing(errors::Cancelled(""));
+  absl::StatusOr<std::unique_ptr<int>> thing(errors::Cancelled(""));
   ASSERT_FALSE(thing.ok());
 }
 
 TEST(StatusOr, TestMoveOnlyValueExtraction) {
-  StatusOr<std::unique_ptr<int>> thing(ReturnUniquePtr());
+  absl::StatusOr<std::unique_ptr<int>> thing(ReturnUniquePtr());
   ASSERT_TRUE(thing.ok());
   std::unique_ptr<int> ptr = std::move(thing).value();
   EXPECT_EQ(0, *ptr);
@@ -121,7 +121,7 @@ TEST(StatusOr, TestMoveOnlyValueExtraction) {
 }
 
 TEST(StatusOr, TestMoveOnlyConversion) {
-  StatusOr<std::unique_ptr<const int>> const_thing(ReturnUniquePtr());
+  absl::StatusOr<std::unique_ptr<const int>> const_thing(ReturnUniquePtr());
   EXPECT_TRUE(const_thing.ok());
   EXPECT_EQ(0, *const_thing.value());
 
@@ -135,7 +135,7 @@ TEST(StatusOr, TestMoveOnlyConversion) {
 
 TEST(StatusOr, TestMoveOnlyVector) {
   // Sanity check that StatusOr<MoveOnly> works in vector.
-  std::vector<StatusOr<std::unique_ptr<int>>> vec;
+  std::vector<absl::StatusOr<std::unique_ptr<int>>> vec;
   vec.push_back(ReturnUniquePtr());
   vec.resize(2);
   auto another_vec = std::move(vec);
@@ -144,11 +144,13 @@ TEST(StatusOr, TestMoveOnlyVector) {
 }
 
 TEST(StatusOr, TestMoveWithValuesAndErrors) {
-  StatusOr<std::string> status_or(std::string(1000, '0'));
-  StatusOr<std::string> value1(std::string(1000, '1'));
-  StatusOr<std::string> value2(std::string(1000, '2'));
-  StatusOr<std::string> error1(Status(absl::StatusCode::kUnknown, "error1"));
-  StatusOr<std::string> error2(Status(absl::StatusCode::kUnknown, "error2"));
+  absl::StatusOr<std::string> status_or(std::string(1000, '0'));
+  absl::StatusOr<std::string> value1(std::string(1000, '1'));
+  absl::StatusOr<std::string> value2(std::string(1000, '2'));
+  absl::StatusOr<std::string> error1(
+      absl::Status(absl::StatusCode::kUnknown, "error1"));
+  absl::StatusOr<std::string> error2(
+      absl::Status(absl::StatusCode::kUnknown, "error2"));
 
   ASSERT_TRUE(status_or.ok());
   EXPECT_EQ(std::string(1000, '0'), status_or.value());
@@ -175,11 +177,13 @@ TEST(StatusOr, TestMoveWithValuesAndErrors) {
 }
 
 TEST(StatusOr, TestCopyWithValuesAndErrors) {
-  StatusOr<std::string> status_or(std::string(1000, '0'));
-  StatusOr<std::string> value1(std::string(1000, '1'));
-  StatusOr<std::string> value2(std::string(1000, '2'));
-  StatusOr<std::string> error1(Status(absl::StatusCode::kUnknown, "error1"));
-  StatusOr<std::string> error2(Status(absl::StatusCode::kUnknown, "error2"));
+  absl::StatusOr<std::string> status_or(std::string(1000, '0'));
+  absl::StatusOr<std::string> value1(std::string(1000, '1'));
+  absl::StatusOr<std::string> value2(std::string(1000, '2'));
+  absl::StatusOr<std::string> error1(
+      absl::Status(absl::StatusCode::kUnknown, "error1"));
+  absl::StatusOr<std::string> error2(
+      absl::Status(absl::StatusCode::kUnknown, "error2"));
 
   ASSERT_TRUE(status_or.ok());
   EXPECT_EQ(std::string(1000, '0'), status_or.value());
@@ -212,13 +216,13 @@ TEST(StatusOr, TestCopyWithValuesAndErrors) {
 }
 
 TEST(StatusOr, TestDefaultCtor) {
-  StatusOr<int> thing;
+  absl::StatusOr<int> thing;
   EXPECT_FALSE(thing.ok());
   EXPECT_EQ(thing.status().code(), absl::StatusCode::kUnknown);
 }
 
 TEST(StatusOrDeathTest, TestDefaultCtorValue) {
-  StatusOr<int> thing;
+  absl::StatusOr<int> thing;
 #ifdef ABSL_HAVE_EXCEPTIONS
   try {
     thing.value();
@@ -230,7 +234,7 @@ TEST(StatusOrDeathTest, TestDefaultCtorValue) {
   EXPECT_DEATH(thing.value(), "");
 #endif
 
-  const StatusOr<int> thing2;
+  const absl::StatusOr<int> thing2;
 #ifdef ABSL_HAVE_EXCEPTIONS
   try {
     thing.value();
@@ -244,93 +248,94 @@ TEST(StatusOrDeathTest, TestDefaultCtorValue) {
 }
 
 TEST(StatusOr, TestStatusCtor) {
-  StatusOr<int> thing(Status(absl::StatusCode::kCancelled, ""));
+  absl::StatusOr<int> thing(absl::Status(absl::StatusCode::kCancelled, ""));
   EXPECT_FALSE(thing.ok());
   EXPECT_EQ(thing.status().code(), absl::StatusCode::kCancelled);
 }
 
 TEST(StatusOr, TestValueCtor) {
   const int kI = 4;
-  const StatusOr<int> thing(kI);
+  const absl::StatusOr<int> thing(kI);
   EXPECT_TRUE(thing.ok());
   EXPECT_EQ(kI, thing.value());
 }
 
 TEST(StatusOr, TestCopyCtorStatusOk) {
   const int kI = 4;
-  const StatusOr<int> original(kI);
-  const StatusOr<int> copy(original);
+  const absl::StatusOr<int> original(kI);
+  const absl::StatusOr<int> copy(original);
   EXPECT_EQ(copy.status(), original.status());
   EXPECT_EQ(original.value(), copy.value());
 }
 
 TEST(StatusOr, TestCopyCtorStatusNotOk) {
-  StatusOr<int> original(Status(absl::StatusCode::kCancelled, ""));
-  StatusOr<int> copy(original);
+  absl::StatusOr<int> original(absl::Status(absl::StatusCode::kCancelled, ""));
+  absl::StatusOr<int> copy(original);
   EXPECT_EQ(copy.status(), original.status());
 }
 
 TEST(StatusOr, TestCopyCtorNonAssignable) {
   const int kI = 4;
   CopyNoAssign value(kI);
-  StatusOr<CopyNoAssign> original(value);
-  StatusOr<CopyNoAssign> copy(original);
+  absl::StatusOr<CopyNoAssign> original(value);
+  absl::StatusOr<CopyNoAssign> copy(original);
   EXPECT_EQ(copy.status(), original.status());
   EXPECT_EQ(original.value().foo_, copy.value().foo_);
 }
 
 TEST(StatusOr, TestCopyCtorStatusOKConverting) {
   const int kI = 4;
-  StatusOr<int> original(kI);
-  StatusOr<double> copy(original);
+  absl::StatusOr<int> original(kI);
+  absl::StatusOr<double> copy(original);
   EXPECT_EQ(copy.status(), original.status());
   EXPECT_DOUBLE_EQ(original.value(), copy.value());
 }
 
 TEST(StatusOr, TestCopyCtorStatusNotOkConverting) {
-  StatusOr<int> original(Status(absl::StatusCode::kCancelled, ""));
-  StatusOr<double> copy(original);
+  absl::StatusOr<int> original(absl::Status(absl::StatusCode::kCancelled, ""));
+  absl::StatusOr<double> copy(original);
   EXPECT_EQ(copy.status(), original.status());
 }
 
 TEST(StatusOr, TestAssignmentStatusOk) {
   const int kI = 4;
-  StatusOr<int> source(kI);
-  StatusOr<int> target;
+  absl::StatusOr<int> source(kI);
+  absl::StatusOr<int> target;
   target = source;
   EXPECT_EQ(target.status(), source.status());
   EXPECT_EQ(source.value(), target.value());
 }
 
 TEST(StatusOr, TestAssignmentStatusNotOk) {
-  StatusOr<int> source(Status(absl::StatusCode::kCancelled, ""));
-  StatusOr<int> target;
+  absl::StatusOr<int> source(absl::Status(absl::StatusCode::kCancelled, ""));
+  absl::StatusOr<int> target;
   target = source;
   EXPECT_EQ(target.status(), source.status());
 }
 
 TEST(StatusOr, TestStatus) {
-  StatusOr<int> good(4);
+  absl::StatusOr<int> good(4);
   EXPECT_TRUE(good.ok());
-  StatusOr<int> bad(Status(absl::StatusCode::kCancelled, ""));
+  absl::StatusOr<int> bad(absl::Status(absl::StatusCode::kCancelled, ""));
   EXPECT_FALSE(bad.ok());
-  EXPECT_EQ(bad.status(), Status(absl::StatusCode::kCancelled, ""));
+  EXPECT_EQ(bad.status(), absl::Status(absl::StatusCode::kCancelled, ""));
 }
 
 TEST(StatusOr, TestValue) {
   const int kI = 4;
-  StatusOr<int> thing(kI);
+  absl::StatusOr<int> thing(kI);
   EXPECT_EQ(kI, thing.value());
 }
 
 TEST(StatusOr, TestValueConst) {
   const int kI = 4;
-  const StatusOr<int> thing(kI);
+  const absl::StatusOr<int> thing(kI);
   EXPECT_EQ(kI, thing.value());
 }
 
 TEST(StatusOrDeathTest, TestValueNotOk) {
-  StatusOr<int> thing(Status(absl::StatusCode::kCancelled, "cancelled"));
+  absl::StatusOr<int> thing(
+      absl::Status(absl::StatusCode::kCancelled, "cancelled"));
 #ifdef ABSL_HAVE_EXCEPTIONS
   try {
     thing.value();
@@ -344,7 +349,7 @@ TEST(StatusOrDeathTest, TestValueNotOk) {
 }
 
 TEST(StatusOrDeathTest, TestValueNotOkConst) {
-  const StatusOr<int> thing(Status(absl::StatusCode::kUnknown, ""));
+  const absl::StatusOr<int> thing(absl::Status(absl::StatusCode::kUnknown, ""));
 #ifdef ABSL_HAVE_EXCEPTIONS
   try {
     thing.value();
@@ -358,13 +363,13 @@ TEST(StatusOrDeathTest, TestValueNotOkConst) {
 }
 
 TEST(StatusOr, TestPointerDefaultCtor) {
-  StatusOr<int*> thing;
+  absl::StatusOr<int*> thing;
   EXPECT_FALSE(thing.ok());
   EXPECT_EQ(thing.status().code(), absl::StatusCode::kUnknown);
 }
 
 TEST(StatusOrDeathTest, TestPointerDefaultCtorValue) {
-  StatusOr<int*> thing;
+  absl::StatusOr<int*> thing;
 #ifdef ABSL_HAVE_EXCEPTIONS
   try {
     thing.value();
@@ -378,94 +383,97 @@ TEST(StatusOrDeathTest, TestPointerDefaultCtorValue) {
 }
 
 TEST(StatusOr, TestPointerStatusCtor) {
-  StatusOr<int*> thing(Status(absl::StatusCode::kCancelled, ""));
+  absl::StatusOr<int*> thing(absl::Status(absl::StatusCode::kCancelled, ""));
   EXPECT_FALSE(thing.ok());
-  EXPECT_EQ(thing.status(), Status(absl::StatusCode::kCancelled, ""));
+  EXPECT_EQ(thing.status(), absl::Status(absl::StatusCode::kCancelled, ""));
 }
 
 TEST(StatusOr, TestPointerValueCtor) {
   const int kI = 4;
-  StatusOr<const int*> thing(&kI);
+  absl::StatusOr<const int*> thing(&kI);
   EXPECT_TRUE(thing.ok());
   EXPECT_EQ(&kI, thing.value());
 }
 
 TEST(StatusOr, TestPointerCopyCtorStatusOk) {
   const int kI = 0;
-  StatusOr<const int*> original(&kI);
-  StatusOr<const int*> copy(original);
+  absl::StatusOr<const int*> original(&kI);
+  absl::StatusOr<const int*> copy(original);
   EXPECT_EQ(copy.status(), original.status());
   EXPECT_EQ(original.value(), copy.value());
 }
 
 TEST(StatusOr, TestPointerCopyCtorStatusNotOk) {
-  StatusOr<int*> original(Status(absl::StatusCode::kCancelled, ""));
-  StatusOr<int*> copy(original);
+  absl::StatusOr<int*> original(absl::Status(absl::StatusCode::kCancelled, ""));
+  absl::StatusOr<int*> copy(original);
   EXPECT_EQ(copy.status(), original.status());
 }
 
 TEST(StatusOr, TestPointerCopyCtorStatusOKConverting) {
   Derived derived;
-  StatusOr<Derived*> original(&derived);
-  StatusOr<Base2*> copy(original);
+  absl::StatusOr<Derived*> original(&derived);
+  absl::StatusOr<Base2*> copy(original);
   EXPECT_EQ(copy.status(), original.status());
   EXPECT_EQ(static_cast<const Base2*>(original.value()), copy.value());
 }
 
 TEST(StatusOr, TestPointerCopyCtorStatusNotOkConverting) {
-  StatusOr<Derived*> original(Status(absl::StatusCode::kCancelled, ""));
-  StatusOr<Base2*> copy(original);
+  absl::StatusOr<Derived*> original(
+      absl::Status(absl::StatusCode::kCancelled, ""));
+  absl::StatusOr<Base2*> copy(original);
   EXPECT_EQ(copy.status(), original.status());
 }
 
 TEST(StatusOr, TestPointerAssignmentStatusOk) {
   const int kI = 0;
-  StatusOr<const int*> source(&kI);
-  StatusOr<const int*> target;
+  absl::StatusOr<const int*> source(&kI);
+  absl::StatusOr<const int*> target;
   target = source;
   EXPECT_EQ(target.status(), source.status());
   EXPECT_EQ(source.value(), target.value());
 }
 
 TEST(StatusOr, TestPointerAssignmentStatusNotOk) {
-  StatusOr<int*> source(Status(absl::StatusCode::kCancelled, ""));
-  StatusOr<int*> target;
+  absl::StatusOr<int*> source(absl::Status(absl::StatusCode::kCancelled, ""));
+  absl::StatusOr<int*> target;
   target = source;
   EXPECT_EQ(target.status(), source.status());
 }
 
 TEST(StatusOr, TestPointerStatus) {
   const int kI = 0;
-  StatusOr<const int*> good(&kI);
+  absl::StatusOr<const int*> good(&kI);
   EXPECT_TRUE(good.ok());
-  StatusOr<const int*> bad(Status(absl::StatusCode::kCancelled, ""));
-  EXPECT_EQ(bad.status(), Status(absl::StatusCode::kCancelled, ""));
+  absl::StatusOr<const int*> bad(
+      absl::Status(absl::StatusCode::kCancelled, ""));
+  EXPECT_EQ(bad.status(), absl::Status(absl::StatusCode::kCancelled, ""));
 }
 
 TEST(StatusOr, TestPointerValue) {
   const int kI = 0;
-  StatusOr<const int*> thing(&kI);
+  absl::StatusOr<const int*> thing(&kI);
   EXPECT_EQ(&kI, thing.value());
 }
 
 TEST(StatusOr, TestPointerValueConst) {
   const int kI = 0;
-  const StatusOr<const int*> thing(&kI);
+  const absl::StatusOr<const int*> thing(&kI);
   EXPECT_EQ(&kI, thing.value());
 }
 
 TEST(StatusOr, TestArrowOperator) {
-  StatusOr<std::unique_ptr<int>> uptr = ReturnUniquePtr();
+  absl::StatusOr<std::unique_ptr<int>> uptr = ReturnUniquePtr();
   EXPECT_EQ(*uptr->get(), 0);
 }
 
 TEST(StatusOr, TestStarOperator) {
-  StatusOr<std::unique_ptr<int>> uptr = ReturnUniquePtr();
+  absl::StatusOr<std::unique_ptr<int>> uptr = ReturnUniquePtr();
   EXPECT_EQ(**uptr, 0);
 }
 
 TEST(StatusOr, TestStarOperatorDeath) {
-  StatusOr<Base1> error(Status(absl::StatusCode::kCancelled, "cancelled"));
+  absl::StatusOr<Base1> error(
+      absl::Status(absl::StatusCode::kCancelled, "cancelled"));
   EXPECT_DEATH(*error, "cancelled");
 }
 
@@ -478,7 +486,7 @@ TEST(StatusOr, TestStarOperatorDeath) {
 //   v.reserve(v.capacity() + 10);
 // }
 
-static StatusOr<int> MakeStatus() { return 100; }
+static absl::StatusOr<int> MakeStatus() { return 100; }
 // A factory to help us benchmark the various factory styles. All of
 // the factory methods are marked as non-inlineable so as to more
 // accurately simulate calling a factory for which you do not have
@@ -502,25 +510,25 @@ class BenchmarkFactory {
   // A more sophisticated factory, which returns a status to indicate
   // the result of the operation. The factory result is populated into
   // the user provided pointer result.
-  Status ArgumentFactory(T** result) TF_ATTRIBUTE_NOINLINE {
+  absl::Status ArgumentFactory(T** result) TF_ATTRIBUTE_NOINLINE {
     *result = value_;
-    return OkStatus();
+    return absl::OkStatus();
   }
 
-  Status ArgumentFactoryFail(T** result) TF_ATTRIBUTE_NOINLINE {
+  absl::Status ArgumentFactoryFail(T** result) TF_ATTRIBUTE_NOINLINE {
     *result = nullptr;
-    return Status(absl::StatusCode::kCancelled, "");
+    return absl::Status(absl::StatusCode::kCancelled, "");
   }
 
-  Status ArgumentFactoryFailShortMsg(T** result) TF_ATTRIBUTE_NOINLINE {
+  absl::Status ArgumentFactoryFailShortMsg(T** result) TF_ATTRIBUTE_NOINLINE {
     *result = nullptr;
-    return Status(absl::StatusCode::kInternal, "");
+    return absl::Status(absl::StatusCode::kInternal, "");
   }
 
-  Status ArgumentFactoryFailLongMsg(T** result) TF_ATTRIBUTE_NOINLINE {
+  absl::Status ArgumentFactoryFailLongMsg(T** result) TF_ATTRIBUTE_NOINLINE {
     *result = nullptr;
-    return Status(absl::StatusCode::kInternal,
-                  "a big string of message junk that will never be read");
+    return absl::Status(absl::StatusCode::kInternal,
+                        "a big string of message junk that will never be read");
   }
 
   // A factory that returns a StatusOr<T*>. If the factory operation
@@ -531,16 +539,16 @@ class BenchmarkFactory {
   }
 
   StatusOr<T*> StatusOrFactoryFail() TF_ATTRIBUTE_NOINLINE {
-    return Status(absl::StatusCode::kCancelled, "");
+    return absl::Status(absl::StatusCode::kCancelled, "");
   }
 
   StatusOr<T*> StatusOrFactoryFailShortMsg() TF_ATTRIBUTE_NOINLINE {
-    return Status(absl::StatusCode::kInternal, "");
+    return absl::Status(absl::StatusCode::kInternal, "");
   }
 
   StatusOr<T*> StatusOrFactoryFailLongMsg() TF_ATTRIBUTE_NOINLINE {
-    return Status(absl::StatusCode::kInternal,
-                  "a big string of message junk that will never be read");
+    return absl::Status(absl::StatusCode::kInternal,
+                        "a big string of message junk that will never be read");
   }
 
  private:
@@ -594,7 +602,7 @@ void BM_ArgumentFactory(::testing::benchmark::State& state) {
   BenchmarkFactory<BenchmarkType> factory;
   for (auto s : state) {
     BenchmarkType* result = nullptr;
-    Status status = factory.ArgumentFactory(&result);
+    absl::Status status = factory.ArgumentFactory(&result);
     if (status.ok() && result != nullptr) {
       result->DoWork();
     }
@@ -607,7 +615,7 @@ BENCHMARK(BM_ArgumentFactory);
 void BM_StatusOrFactory(::testing::benchmark::State& state) {
   BenchmarkFactory<BenchmarkType> factory;
   for (auto s : state) {
-    StatusOr<BenchmarkType*> result = factory.StatusOrFactory();
+    absl::StatusOr<BenchmarkType*> result = factory.StatusOrFactory();
     if (result.ok()) {
       result.value()->DoWork();
     }
@@ -622,7 +630,7 @@ void BM_ArgumentFactoryFail(::testing::benchmark::State& state) {
   BenchmarkFactory<BenchmarkType> factory;
   for (auto s : state) {
     BenchmarkType* result = nullptr;
-    Status status = factory.ArgumentFactoryFail(&result);
+    absl::Status status = factory.ArgumentFactoryFail(&result);
     if (status.ok() && result != nullptr) {
       result->DoWork();
     }
@@ -635,7 +643,7 @@ BENCHMARK(BM_ArgumentFactoryFail);
 void BM_StatusOrFactoryFail(::testing::benchmark::State& state) {
   BenchmarkFactory<BenchmarkType> factory;
   for (auto s : state) {
-    StatusOr<BenchmarkType*> result = factory.StatusOrFactoryFail();
+    absl::StatusOr<BenchmarkType*> result = factory.StatusOrFactoryFail();
     if (result.ok()) {
       result.value()->DoWork();
     }
@@ -650,7 +658,7 @@ void BM_ArgumentFactoryFailShortMsg(::testing::benchmark::State& state) {
   BenchmarkFactory<BenchmarkType> factory;
   for (auto s : state) {
     BenchmarkType* result = nullptr;
-    Status status = factory.ArgumentFactoryFailShortMsg(&result);
+    absl::Status status = factory.ArgumentFactoryFailShortMsg(&result);
     if (status.ok() && result != nullptr) {
       result->DoWork();
     }
@@ -663,7 +671,8 @@ BENCHMARK(BM_ArgumentFactoryFailShortMsg);
 void BM_StatusOrFactoryFailShortMsg(::testing::benchmark::State& state) {
   BenchmarkFactory<BenchmarkType> factory;
   for (auto s : state) {
-    StatusOr<BenchmarkType*> result = factory.StatusOrFactoryFailShortMsg();
+    absl::StatusOr<BenchmarkType*> result =
+        factory.StatusOrFactoryFailShortMsg();
     if (result.ok()) {
       result.value()->DoWork();
     }
@@ -678,7 +687,7 @@ void BM_ArgumentFactoryFailLongMsg(::testing::benchmark::State& state) {
   BenchmarkFactory<BenchmarkType> factory;
   for (auto s : state) {
     BenchmarkType* result = nullptr;
-    Status status = factory.ArgumentFactoryFailLongMsg(&result);
+    absl::Status status = factory.ArgumentFactoryFailLongMsg(&result);
     if (status.ok() && result != nullptr) {
       result->DoWork();
     }
@@ -691,7 +700,8 @@ BENCHMARK(BM_ArgumentFactoryFailLongMsg);
 void BM_StatusOrFactoryFailLongMsg(::testing::benchmark::State& state) {
   BenchmarkFactory<BenchmarkType> factory;
   for (auto s : state) {
-    StatusOr<BenchmarkType*> result = factory.StatusOrFactoryFailLongMsg();
+    absl::StatusOr<BenchmarkType*> result =
+        factory.StatusOrFactoryFailLongMsg();
     if (result.ok()) {
       result.value()->DoWork();
     }
@@ -701,22 +711,22 @@ BENCHMARK(BM_StatusOrFactoryFailLongMsg);
 
 #if defined(PLATFORM_GOOGLE)
 
-StatusOr<int> GetError() {
+absl::StatusOr<int> GetError() {
   return absl::InvalidArgumentError("An invalid argument error");
 }
 
-StatusOr<int> PropagateError() {
+absl::StatusOr<int> PropagateError() {
   TF_ASSIGN_OR_RETURN(int a, GetError());
   return a;
 }
 
-StatusOr<int> PropagateError2() {
+absl::StatusOr<int> PropagateError2() {
   TF_ASSIGN_OR_RETURN(int a, PropagateError());
   return a;
 }
 
 TEST(Status, StackTracePropagation) {
-  StatusOr<int> s = PropagateError2();
+  absl::StatusOr<int> s = PropagateError2();
   auto sources = s.status().GetSourceLocations();
   ASSERT_EQ(sources.size(), 3);
 

From a81a393f72b9d25a36b978646a205b050109a8fb Mon Sep 17 00:00:00 2001
From: Dmitri Gribenko <dmitrig@google.com>
Date: Fri, 17 May 2024 06:11:38 -0700
Subject: [PATCH 348/478] Integrate LLVM at llvm/llvm-project@5a20a07fce88

Updates LLVM usage to match
[5a20a07fce88](https://github.com/llvm/llvm-project/commit/5a20a07fce88)

PiperOrigin-RevId: 634736134
---
 third_party/llvm/workspace.bzl | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/third_party/llvm/workspace.bzl b/third_party/llvm/workspace.bzl
index 7d72ce6255348f..f017f7ecff0479 100644
--- a/third_party/llvm/workspace.bzl
+++ b/third_party/llvm/workspace.bzl
@@ -4,8 +4,8 @@ load("//third_party:repo.bzl", "tf_http_archive")
 
 def repo(name):
     """Imports LLVM."""
-    LLVM_COMMIT = "5b7088c3619e95eaa8e154765470ddb4d3859fa4"
-    LLVM_SHA256 = "21c667d6ea0c4141ec02bc96e788d1cda2f55063626661cdc21665fb61025abc"
+    LLVM_COMMIT = "5a20a07fce88d54cf01cafde489bfc2fc447acc0"
+    LLVM_SHA256 = "fb33549dd5fef2b3e8a214cbbbd5f23962a273c7c0871a70d1fe58fc41bdaded"
 
     tf_http_archive(
         name = name,

From c287cd48573c39c58dce7443d1526ac99aefca82 Mon Sep 17 00:00:00 2001
From: Adrian Kuegel <akuegel@google.com>
Date: Fri, 17 May 2024 06:40:31 -0700
Subject: [PATCH 349/478] IsSimplifiedScatter should also return false if
 update_window_dims is unsorted.

We started relying on that with LayoutNormalization creating Scatters that
don't have sorted update_window_dims, and want ScatterSimplifier to clean this
up.

PiperOrigin-RevId: 634748724
---
 third_party/xla/xla/service/BUILD             |  2 ++
 .../xla/xla/service/scatter_simplifier.cc     |  3 +-
 .../xla/service/scatter_simplifier_test.cc    | 29 +++++++++++++++++++
 3 files changed, 33 insertions(+), 1 deletion(-)

diff --git a/third_party/xla/xla/service/BUILD b/third_party/xla/xla/service/BUILD
index 7df3a319664561..39336da669e985 100644
--- a/third_party/xla/xla/service/BUILD
+++ b/third_party/xla/xla/service/BUILD
@@ -7430,9 +7430,11 @@ xla_cc_test(
     name = "scatter_simplifier_test",
     srcs = ["scatter_simplifier_test.cc"],
     deps = [
+        ":hlo_parser",
         ":hlo_pass",
         ":hlo_pass_pipeline",
         ":scatter_simplifier",
+        "//xla/hlo/ir:hlo",
         "//xla/tests:hlo_test_base",
         "//xla/tests:xla_internal_test_main",
     ],
diff --git a/third_party/xla/xla/service/scatter_simplifier.cc b/third_party/xla/xla/service/scatter_simplifier.cc
index b6b12a3dc3b9ad..c5265cd61eb585 100644
--- a/third_party/xla/xla/service/scatter_simplifier.cc
+++ b/third_party/xla/xla/service/scatter_simplifier.cc
@@ -215,10 +215,11 @@ bool ScatterSimplifier::IsSimplifiedScatter(
       !IsIdentityPermutation(dims.scatter_dims_to_operand_dims());
   bool scatter_dim_not_first =
       absl::c_linear_search(dims.update_window_dims(), 0);
+  bool update_window_dims_sorted = absl::c_is_sorted(dims.update_window_dims());
 
   return !(nonstandard_index_vector_dim || num_scatter_dims > 1 ||
            scatter_indices_reordered || scatter_dim_not_first ||
-           !dims.inserted_window_dims().empty());
+           !update_window_dims_sorted || !dims.inserted_window_dims().empty());
 }
 
 bool ScatterSimplifier::InstructionMatchesPattern(HloInstruction* inst) {
diff --git a/third_party/xla/xla/service/scatter_simplifier_test.cc b/third_party/xla/xla/service/scatter_simplifier_test.cc
index dc813e506dedac..b10054347a4fc7 100644
--- a/third_party/xla/xla/service/scatter_simplifier_test.cc
+++ b/third_party/xla/xla/service/scatter_simplifier_test.cc
@@ -17,6 +17,8 @@ limitations under the License.
 
 #include <optional>
 
+#include "xla/hlo/ir/hlo_casting_utils.h"
+#include "xla/service/hlo_parser.h"
 #include "xla/service/hlo_pass_fix.h"
 #include "xla/service/hlo_pass_pipeline.h"
 #include "xla/tests/hlo_test_base.h"
@@ -252,5 +254,32 @@ TEST_F(ScatterSimplifierTest, ZeroDimScatterIndices) {
     )");
 }
 
+TEST_F(ScatterSimplifierTest,
+       IsSimplifiedScatterReturnsFalseForUnsortedWindowDims) {
+  constexpr absl::string_view kModuleStr = R"(
+    HloModule scatter_simplifier
+
+    scatter_computation {
+      %p0 = f32[] parameter(0)
+      ROOT result = f32[] parameter(1)
+    }
+
+    ENTRY kernel_entry {
+      operand = f32[3,2] parameter(0)
+      indices = s32[1,1] parameter(1)
+      update = f32[1,2,2] parameter(2)
+      ROOT scatter = f32[3,2] scatter(operand, indices, update),
+          to_apply=scatter_computation,
+          update_window_dims={2,1},
+          inserted_window_dims={},
+          scatter_dims_to_operand_dims={0},
+          index_vector_dim=1
+    })";
+  auto module = ParseAndReturnUnverifiedModule(kModuleStr).value();
+  auto scatter = module->entry_computation()->root_instruction();
+  EXPECT_FALSE(ScatterSimplifier::IsSimplifiedScatter(
+      Cast<HloScatterInstruction>(scatter)));
+}
+
 }  // namespace
 }  // namespace xla

From ce41ee4f396bde91aac3360197708ae5ec7fdee8 Mon Sep 17 00:00:00 2001
From: Ilia Sergachev <isergachev@nvidia.com>
Date: Fri, 17 May 2024 06:59:40 -0700
Subject: [PATCH 350/478] PR #12520: [GPU] Add new flag
 xla_gpu_exclude_nondeterministic_ops.

Imported from GitHub PR https://github.com/openxla/xla/pull/12520

It's more granular than the existing --xla_gpu_deterministic_ops because it allows doing an autotuning compilation with non-deterministic ops disabled.

--xla_gpu_deterministic_ops is a superset of --xla_gpu_exclude_nondeterministic_ops, so --xla_gpu_deterministic_ops=true will be setting --xla_gpu_exclude_nondeterministic_ops=true too.
Copybara import of the project:

--
4e2837457dc426154bf80f321c001c916a7d3677 by Ilia Sergachev <isergachev@nvidia.com>:

[GPU] Add new flag xla_gpu_exclude_nondeterministic_ops.

Merging this change closes #12520

PiperOrigin-RevId: 634756524
---
 third_party/xla/xla/debug_options_flags.cc    |  5 +++
 .../xla/service/gpu/conv_algorithm_picker.cc  |  8 +++-
 .../xla/xla/service/gpu/determinism_test.cc   | 44 +++++++++++++++++--
 .../xla/xla/service/gpu/fusions/custom.cc     |  3 +-
 .../xla/service/gpu/gemm_algorithm_picker.cc  |  3 +-
 .../xla/xla/service/gpu/gpu_compiler.cc       |  3 +-
 .../xla/service/gpu/ir_emitter_unnested.cc    | 12 +++--
 .../service/gpu/tests/gemm_rewrite_test.cc    |  2 +-
 third_party/xla/xla/xla.proto                 | 18 ++++++--
 9 files changed, 80 insertions(+), 18 deletions(-)

diff --git a/third_party/xla/xla/debug_options_flags.cc b/third_party/xla/xla/debug_options_flags.cc
index 7232a0a8b157bc..ba4561802cf8fb 100644
--- a/third_party/xla/xla/debug_options_flags.cc
+++ b/third_party/xla/xla/debug_options_flags.cc
@@ -1012,6 +1012,11 @@ void MakeDebugOptionsFlags(std::vector<tsl::Flag>* flag_list,
                 bool_setter_for(&DebugOptions::set_xla_gpu_deterministic_ops),
                 debug_options->xla_gpu_deterministic_ops(),
                 "Guarantees run-to-run determinism on GPU."));
+  flag_list->push_back(tsl::Flag(
+      "xla_gpu_exclude_nondeterministic_ops",
+      bool_setter_for(&DebugOptions::set_xla_gpu_exclude_nondeterministic_ops),
+      debug_options->xla_gpu_exclude_nondeterministic_ops(),
+      "Excludes non-deterministic ops from compiled executables."));
   flag_list->push_back(tsl::Flag(
       "xla_gpu_disable_async_collectives",
       setter_for_xla_gpu_disable_async_collectives,
diff --git a/third_party/xla/xla/service/gpu/conv_algorithm_picker.cc b/third_party/xla/xla/service/gpu/conv_algorithm_picker.cc
index 2a6dcae8417b9b..cda0d5a700d238 100644
--- a/third_party/xla/xla/service/gpu/conv_algorithm_picker.cc
+++ b/third_party/xla/xla/service/gpu/conv_algorithm_picker.cc
@@ -769,7 +769,9 @@ GpuConvAlgorithmPicker::PickBestAlgorithmNoCacheCuda(
 
   const bool cudnn_frontend_enabled =
       debug_options.xla_gpu_enable_cudnn_frontend();
-  const bool deterministic_ops = debug_options.xla_gpu_deterministic_ops();
+  const bool deterministic_ops =
+      debug_options.xla_gpu_deterministic_ops() ||
+      debug_options.xla_gpu_exclude_nondeterministic_ops();
   bool allow_tf32 = true;
   // TODO(b/284371623): Properly set allow_tf32 even if instr==nullptr, which is
   // the case when running an AOT compiled executable with runtime autotuning.
@@ -879,7 +881,9 @@ GpuConvAlgorithmPicker::PickBestAlgorithmNoCacheRocm(
 
   const DebugOptions& debug_options =
       instr->GetModule()->config().debug_options();
-  const bool deterministic_ops = debug_options.xla_gpu_deterministic_ops();
+  const bool deterministic_ops =
+      debug_options.xla_gpu_deterministic_ops() ||
+      debug_options.xla_gpu_exclude_nondeterministic_ops();
   const bool allow_tf32 = absl::c_all_of(
       instr->precision_config().operand_precision(),
       [](int precision) { return precision <= PrecisionConfig::HIGH; });
diff --git a/third_party/xla/xla/service/gpu/determinism_test.cc b/third_party/xla/xla/service/gpu/determinism_test.cc
index 43799901f2a0eb..02f5ba86a64254 100644
--- a/third_party/xla/xla/service/gpu/determinism_test.cc
+++ b/third_party/xla/xla/service/gpu/determinism_test.cc
@@ -38,7 +38,7 @@ namespace gpu {
 class DeterminismTest : public GpuCodegenTest {
  public:
   DeterminismTest() : debug_options_(HloTestBase::GetDebugOptionsForTest()) {
-    debug_options_.set_xla_gpu_deterministic_ops(true);
+    debug_options_.set_xla_gpu_exclude_nondeterministic_ops(true);
     // Randomize timer durations to better test autotuning does not introduce
     // nondeterminism.
     se::gpu::GpuTimer::ReturnRandomDurationsForTesting();
@@ -109,7 +109,7 @@ ENTRY e {
   AssertDeterminism(kHloText);
 }
 
-TEST_F(DeterminismTest, TritonDot) {
+TEST_F(DeterminismTest, DeterministicTritonGemmUsesDefaultConfig) {
 #if GOOGLE_CUDA
   auto comp = backend()
                   .default_stream_executor()
@@ -130,11 +130,47 @@ ENTRY e {
   ROOT d = f32[128,128] dot(p0_convert, p1), lhs_contracting_dims={1}, rhs_contracting_dims={0}
 })";
 
-  debug_options_.set_xla_gpu_triton_gemm_any(true);
-  MatchOptimizedHlo(kHloText, R"(; CHECK: __triton_gemm)");
+  // Disable autotuning.
+  debug_options_.set_xla_gpu_deterministic_ops(true);
+  // Check that triton is used but without autotuning (default config).
+  AutotunerUtil::ClearAutotuneResults();
+  MatchOptimizedHlo(kHloText, R"(
+    CHECK: __triton_gemm
+    CHECK: {"block_m":"32","block_n":"32","block_k":"32","split_k":"1","num_stages":"1","num_warps":"4","num_ctas":"1"}
+  )");
   AssertDeterminism(kHloText, /*num_runs=*/3);
 }
 
+TEST_F(DeterminismTest, ExcludingNonDeterministicOpsDoesNotDisableAutotuning) {
+#if GOOGLE_CUDA
+  auto comp = backend()
+                  .default_stream_executor()
+                  ->GetDeviceDescription()
+                  .cuda_compute_capability();
+  if (!comp.IsAtLeast(se::CudaComputeCapability::VOLTA)) {
+    GTEST_SKIP() << "Triton not used on pre-Volta GPUs";
+  }
+#elif TENSORFLOW_USE_ROCM
+  GTEST_SKIP() << "Triton Gemm rewriter is not yet supported on ROCM";
+#endif  // TENSORFLOW_USE_ROCM
+
+  debug_options_.set_xla_gpu_cublas_fallback(false);
+  ASSERT_FALSE(debug_options_.xla_gpu_deterministic_ops());
+  AutotunerUtil::ClearAutotuneResults();
+  // The default config is not used when autotuning is on.
+  MatchOptimizedHlo(R"(
+ENTRY e {
+  p0 = bf16[128,128] parameter(0)
+  p0_convert = f32[128,128] convert(p0)
+  p1 = f32[128,128] parameter(1)
+  ROOT d = f32[128,128] dot(p0_convert, p1), lhs_contracting_dims={1}, rhs_contracting_dims={0}
+})",
+                    R"(
+    CHECK: __triton_gemm
+    CHECK-NOT: {"block_m":"32","block_n":"32","block_k":"32","split_k":"1","num_stages":"1","num_warps":"4","num_ctas":"1"}
+  )");
+}
+
 TEST_F(DeterminismTest, Conv) {
   constexpr absl::string_view kHloText = R"(
 ENTRY e {
diff --git a/third_party/xla/xla/service/gpu/fusions/custom.cc b/third_party/xla/xla/service/gpu/fusions/custom.cc
index e6dc11a46cb0be..3efdae3ca831dd 100644
--- a/third_party/xla/xla/service/gpu/fusions/custom.cc
+++ b/third_party/xla/xla/service/gpu/fusions/custom.cc
@@ -335,7 +335,8 @@ absl::StatusOr<FusionEmissionResult> EmitGemm(
   }
 
   bool deterministic_ops =
-      ir_emitter_context.debug_options().xla_gpu_deterministic_ops();
+      ir_emitter_context.debug_options().xla_gpu_deterministic_ops() ||
+      ir_emitter_context.debug_options().xla_gpu_exclude_nondeterministic_ops();
 
   TF_ASSIGN_OR_RETURN(
       GemmConfig config,
diff --git a/third_party/xla/xla/service/gpu/gemm_algorithm_picker.cc b/third_party/xla/xla/service/gpu/gemm_algorithm_picker.cc
index e43703f362b525..2207d9eac9999d 100644
--- a/third_party/xla/xla/service/gpu/gemm_algorithm_picker.cc
+++ b/third_party/xla/xla/service/gpu/gemm_algorithm_picker.cc
@@ -108,7 +108,8 @@ class GemmAutotuner {
     TF_ASSIGN_OR_RETURN(stream_, autotune_config_.GetStream());
     const DebugOptions& debug_options =
         gemm->GetModule()->config().debug_options();
-    deterministic_ops_ = debug_options.xla_gpu_deterministic_ops();
+    deterministic_ops_ = debug_options.xla_gpu_deterministic_ops() ||
+                         debug_options.xla_gpu_exclude_nondeterministic_ops();
     solutions_limit_ = debug_options.xla_gpu_autotune_max_solutions();
 
     TF_ASSIGN_OR_RETURN(auto gemm_config, GemmConfig::For(gemm));
diff --git a/third_party/xla/xla/service/gpu/gpu_compiler.cc b/third_party/xla/xla/service/gpu/gpu_compiler.cc
index c9bd0c8812d8c5..87126e13307473 100644
--- a/third_party/xla/xla/service/gpu/gpu_compiler.cc
+++ b/third_party/xla/xla/service/gpu/gpu_compiler.cc
@@ -691,7 +691,8 @@ absl::Status RunOptimizationPasses(
   // handle it.
   pipeline.AddPass<ZeroSizedHloElimination>();
 
-  if (debug_options.xla_gpu_deterministic_ops()) {
+  if (debug_options.xla_gpu_deterministic_ops() ||
+      debug_options.xla_gpu_exclude_nondeterministic_ops()) {
     // Scatter can be indeterministic if indices are not unique or a non
     // associative combiner function is used. Eliminate these Scatter ops.
     pipeline.AddPass<ScatterExpander>(
diff --git a/third_party/xla/xla/service/gpu/ir_emitter_unnested.cc b/third_party/xla/xla/service/gpu/ir_emitter_unnested.cc
index 5961fe8d68bce1..99f469752da5f8 100644
--- a/third_party/xla/xla/service/gpu/ir_emitter_unnested.cc
+++ b/third_party/xla/xla/service/gpu/ir_emitter_unnested.cc
@@ -656,7 +656,9 @@ absl::Status IrEmitterUnnested::EmitGemmThunk(
   }
 
   bool deterministic_ops =
-      ir_emitter_context_->debug_options().xla_gpu_deterministic_ops();
+      ir_emitter_context_->debug_options().xla_gpu_deterministic_ops() ||
+      ir_emitter_context_->debug_options()
+          .xla_gpu_exclude_nondeterministic_ops();
 
   TF_ASSIGN_OR_RETURN(
       GemmConfig config,
@@ -1690,11 +1692,13 @@ absl::Status IrEmitterUnnested::EmitFusion(const HloFusionInstruction* instr,
 
 absl::Status IrEmitterUnnested::AssertNonDeterminismIsOkay(
     const std::string& op_name) {
-  if (ir_emitter_context_->debug_options().xla_gpu_deterministic_ops()) {
+  if (ir_emitter_context_->debug_options().xla_gpu_deterministic_ops() ||
+      ir_emitter_context_->debug_options()
+          .xla_gpu_exclude_nondeterministic_ops()) {
     return Unimplemented(
         "HLO instruction %s does not have a deterministic implementation, "
-        "but run-to-run determinism is required by "
-        "--xla_gpu_deterministic_ops.",
+        "but run-to-run determinism is required by --xla_gpu_deterministic_ops "
+        "or --xla_gpu_exclude_nondeterministic_ops.",
         op_name);
   }
   return absl::OkStatus();
diff --git a/third_party/xla/xla/service/gpu/tests/gemm_rewrite_test.cc b/third_party/xla/xla/service/gpu/tests/gemm_rewrite_test.cc
index cca4fc6aa957c3..4dd99c541a1477 100644
--- a/third_party/xla/xla/service/gpu/tests/gemm_rewrite_test.cc
+++ b/third_party/xla/xla/service/gpu/tests/gemm_rewrite_test.cc
@@ -202,7 +202,7 @@ ENTRY AddDotsFunc {
   auto get_module = [&]() {
     HloModuleConfig config;
     DebugOptions debug_options = GetDebugOptionsForTest();
-    debug_options.set_xla_gpu_deterministic_ops(true);
+    debug_options.set_xla_gpu_exclude_nondeterministic_ops(true);
     config.set_debug_options(debug_options);
     return ParseAndReturnVerifiedModule(hlo_text, config);
   };
diff --git a/third_party/xla/xla/xla.proto b/third_party/xla/xla/xla.proto
index ea742ebc06cd61..2960c21d183601 100644
--- a/third_party/xla/xla/xla.proto
+++ b/third_party/xla/xla/xla.proto
@@ -363,9 +363,9 @@ message DebugOptions {
   int32 xla_gpu_force_compilation_parallelism = 147;
   bool xla_gpu_enable_llvm_module_compilation_parallelism = 268;
 
-  // Guarantees run-to-run determinism. At present, the HLO ops Scatter and
-  // SelectAndScatter do not have deterministic XLA:GPU implementations.
-  // Compilation errors out if these ops are encountered.
+  // Guarantees run-to-run determinism.
+  // This flag implies --xla_gpu_exclude_nondeterministic_ops and in addition
+  // disables autotuning.
   bool xla_gpu_deterministic_ops = 148;
 
   // Paths to files with LLVM code.
@@ -771,7 +771,17 @@ message DebugOptions {
   // If true, will enable host memory offloading on a device.
   bool xla_gpu_enable_host_memory_offloading = 296;
 
-  // Next id: 297
+  // Excludes non-deterministic ops from compiled executables.
+  // Unlike --xla_gpu_deterministic_ops does not disable autotuning - the
+  // compilation itself can be non-deterministic.
+  // At present, the HLO op SelectAndScatter does not have a
+  // deterministic XLA:GPU implementation.
+  // Compilation errors out if SelectAndScatter is encountered.
+  // Scatter ops can non-deterministic by default; these get converted to
+  // a deterministic implementation.
+  bool xla_gpu_exclude_nondeterministic_ops = 297;
+
+  // Next id: 298
 
   // Extra options to pass to the compilation backend (e.g. LLVM); specific
   // interpretation of these values is left to the backend.

From 2a6c0d66749adf848115ab0964efe9235adef33f Mon Sep 17 00:00:00 2001
From: Benjamin Chetioui <bchetioui@google.com>
Date: Fri, 17 May 2024 08:23:50 -0700
Subject: [PATCH 351/478] [XLA:GPU] Add support for destructuring collapsing
 reshapes in symbolic tiles.

Take a reshape `[6,8] reshape([48])`. The indexing map going through this\
reshape from output to input would be\
&nbsp;&nbsp;&nbsp;&nbsp;`(d0, d1) -> (d0 * 8 + d1)`.\
Previously to this change, there was no support for extracting sizes and\
strides from a multivariate expression.

Deriving a `size_map` for the indexing map above is as simple as creating a map\
&nbsp;&nbsp;&nbsp;&nbsp;`(s0, s1) -> (s0 * s1)`,\
which is easy enough. However, deriving a `stride_map` for the expression is\
complicated. Conceptually, the stride of the composite expression should\
correspond to the stride of the minormost dimension involved in the reshape\
along which we capture more than a single element.

This holds because there are restrictions on what a valid tiling can be when\
going through a collapse. Let `s` be an `n`-dimensional shape that is fully\
collapsed. In order to be propagated successfully through the collapse, the\
pattern of the tiling of `s` has to look like\
&nbsp;&nbsp;&nbsp;&nbsp;`(1*, partial_dim?, full_dims*, 1*)`\
where `full_dims` are dimensions that are captured completely, and\
`partial_dim` is a dimension that can be captured with an arbitrary tile.\
This restriction is necessary to ensure that the gap between two elements\
captured in the expression is always the same (i.e., the set of elements can\
be described using a single `stride` and is thus a tile).

Based on the above, algorithm for extracting the `stride` could therefore be\
represented as a series of nested `if` statements\
&nbsp;&nbsp;&nbsp;&nbsp;`if size0 != 1 then stride0 else (if size1 != 1 then stride1 else ...)`\
where `{size,stride}i` corresponds to the `i-th` major {size,stride}.

We implement a utility function that allows us to generate `if` statements as\
affine expressions.

PiperOrigin-RevId: 634784163
---
 third_party/xla/xla/service/gpu/model/BUILD   |   5 +-
 .../xla/service/gpu/model/symbolic_tile.cc    | 298 +++++++++++++++++-
 .../service/gpu/model/symbolic_tile_test.cc   |  88 ++++++
 3 files changed, 372 insertions(+), 19 deletions(-)

diff --git a/third_party/xla/xla/service/gpu/model/BUILD b/third_party/xla/xla/service/gpu/model/BUILD
index ae876e97aa6832..759476d142d306 100644
--- a/third_party/xla/xla/service/gpu/model/BUILD
+++ b/third_party/xla/xla/service/gpu/model/BUILD
@@ -546,12 +546,13 @@ xla_cc_test(
         ":indexing_analysis",
         ":indexing_test_utils",
         ":symbolic_tile",
-        "//xla/hlo/ir:hlo",
-        "//xla/tests:hlo_test_base",
         "//xla/tests:xla_internal_test_main",
+        "@com_google_absl//absl/log:check",
         "@com_google_absl//absl/strings",
         "@com_google_absl//absl/strings:string_view",
+        "@com_google_absl//absl/types:span",
         "@com_google_googletest//:gtest",
+        "@llvm-project//llvm:Support",
         "@llvm-project//mlir:IR",
         "@local_tsl//tsl/platform:test",
     ],
diff --git a/third_party/xla/xla/service/gpu/model/symbolic_tile.cc b/third_party/xla/xla/service/gpu/model/symbolic_tile.cc
index 07c6881f13a5b1..8fca5c039347e6 100644
--- a/third_party/xla/xla/service/gpu/model/symbolic_tile.cc
+++ b/third_party/xla/xla/service/gpu/model/symbolic_tile.cc
@@ -43,9 +43,12 @@ namespace xla {
 namespace gpu {
 namespace {
 
+using ::mlir::AffineConstantExpr;
+using ::mlir::AffineDimExpr;
 using ::mlir::AffineExpr;
 using ::mlir::AffineExprKind;
 using ::mlir::AffineMap;
+using ::mlir::AffineSymbolExpr;
 using ::mlir::getAffineConstantExpr;
 using ::mlir::getAffineDimExpr;
 using ::mlir::MLIRContext;
@@ -170,13 +173,268 @@ std::optional<SizeAndStrideExpression> ExtractSizeAndStrideFromFloorDiv(
   return std::nullopt;
 }
 
+// Given an n-ary summation of expressions `e0 + e1 + ... + e{n-1}` with
+// arbitrary order of association, returns the vector `(e0, e1, ..., e{n-1})`.
+// The order of the returned subexpressions is not guaranteed to match the order
+// in which they appear in the original expression.
+//
+// AffineExprKind::Add should be the operation that binds the least tightly,
+// allowing us to simply recursively destructure expressions until we reach an
+// AffineExprKind that is not an AffineExprKind::Add.
+//
+// Note that this will only work correctly for expressions that do no
+// factoring/grouping of summands such as `(d0 + d1) * c` or `(d0 + d1) mod c`.
+// It's unclear at this point whether this restriction will prove problematic,
+// but it isn't really worth thinking about until we are sure this actually
+// has practical implications.
+std::vector<AffineExpr> DestructureSummation(AffineExpr expr) {
+  std::vector<AffineExpr> summands;
+
+  switch (expr.getKind()) {
+    case AffineExprKind::Add: {
+      const auto add = llvm::cast<mlir::AffineBinaryOpExpr>(expr);
+
+      absl::c_move(DestructureSummation(add.getLHS()),
+                   std::back_inserter(summands));
+      absl::c_move(DestructureSummation(add.getRHS()),
+                   std::back_inserter(summands));
+      break;
+    }
+    default:
+      // The expression is not a sum.
+      summands.push_back(expr);
+      break;
+  }
+
+  return summands;
+}
+
 std::optional<SizeAndStrideExpression> ExtractSizeAndStride(
-    AffineExpr strided_indexing, absl::Span<Interval const> symbol_intervals) {
+    AffineExpr strided_indexing, absl::Span<Interval const> dimension_intervals,
+    absl::Span<Interval const> symbol_intervals);
+
+// Given a multivariate summation of strided indexing expression, extracts a
+// size and a stride for each summand. Returns std::nullopt if extraction fails
+// for any of the summands.
+std::optional<std::vector<SizeAndStrideExpression>>
+ExtractSizesAndStridesFromMultivariateSummation(
+    AffineExpr summation, absl::Span<Interval const> dimension_intervals,
+    absl::Span<Interval const> symbol_intervals) {
+  std::vector<AffineExpr> summands = DestructureSummation(summation);
+
+  std::vector<SizeAndStrideExpression> sizes_and_strides;
+  sizes_and_strides.reserve(summands.size());
+  for (AffineExpr summand : summands) {
+    std::optional<SizeAndStrideExpression> maybe_size_and_stride =
+        ExtractSizeAndStride(summand, dimension_intervals, symbol_intervals);
+    if (!maybe_size_and_stride.has_value()) {
+      VLOG(1) << "Couldn't extract size and stride from "
+              << AffineMapPrinter().ToString(summand);
+      return std::nullopt;
+    }
+    sizes_and_strides.push_back(*maybe_size_and_stride);
+  }
+  return sizes_and_strides;
+}
+
+// Given a list of sizes and strides, returns the product of all sizes.
+AffineExpr CombineSizes(
+    absl::Span<SizeAndStrideExpression const> sizes_and_strides) {
+  CHECK(!sizes_and_strides.empty());
+  AffineExpr product =
+      getAffineConstantExpr(1, sizes_and_strides[0].size.getContext());
+  for (const SizeAndStrideExpression& size_and_stride : sizes_and_strides) {
+    product = product * size_and_stride.size;
+  }
+  return product;
+}
+
+// Returns an affine expression logically equivalent to
+//   `eq_param != 1 ? true_expr : false_expr`.
+// `eq_param` is assumed to be able to be in the inclusive range
+//    {1, 2, ..., eq_param_inclusive_upper_bound}.
+AffineExpr IfNeqOne(AffineExpr eq_param, AffineExpr true_expr,
+                    AffineExpr false_expr,
+                    int64_t eq_param_inclusive_upper_bound) {
+  // Let e = eq_param, and b = eq_param_inclusive_bound, then we have:
+  //     1 <= e <= b
+  // <=> -b <= e - b - 1 <= -1              (subtract (b + 1))
+  // <=> 1 <= b + 1 - e <= b                (negate)
+  // <=> 0 <= (b + 1 - e) floordiv b <= 1   (divide by b)
+  //
+  // Since (b + 1 - e) floordiv b is an integer, it can only take values 0 or 1.
+  // Let's prove that
+  //   (b + 1 - e) floordiv b = 1 <=> e = 1.
+  //
+  // * If e = 1, then (b + 1 - e) floordiv b = (b + 1 - 1) floordiv b = 1.
+  // * If e != 1, then 1 < e since 1 is the lower bound for e.
+  //     1 < e <=> -e < -1                       (negate)
+  //           <=> b + 1 - e < b                 (add b + 1)
+  //           <=> (b - e + 1) floordiv b < 1.   (divide by b)
+  //   We also know that 0 <= (b + 1 - e) floordiv b. Therefore, we have that
+  //     (b - e + 1) floordiv b = 0.
+  //
+  // Thus,
+  //   (b + 1 - e) floordiv b = 1 <=> e = 1, and
+  //   (b + 1 - e) floordiv b = 0 <=> e != 1
+  // hold.
+  AffineExpr b = getAffineConstantExpr(eq_param_inclusive_upper_bound,
+                                       eq_param.getContext());
+  AffineExpr condition = mlir::getAffineBinaryOpExpr(AffineExprKind::FloorDiv,
+                                                     b + 1 - eq_param, b);
+
+  return condition * false_expr + (1 - condition) * true_expr;
+}
+
+// Sorts a list of `SizeAndStrideExpression`s by stride. There is a precondition
+// that all strides are constant.
+void SortByStride(std::vector<SizeAndStrideExpression>& sizes_and_strides) {
+  absl::c_sort(sizes_and_strides, [](const SizeAndStrideExpression& sas1,
+                                     const SizeAndStrideExpression& sas2) {
+    int64_t stride1 = llvm::cast<AffineConstantExpr>(sas1.stride).getValue();
+    int64_t stride2 = llvm::cast<AffineConstantExpr>(sas2.stride).getValue();
+    return stride1 < stride2;
+  });
+}
+
+// Given a list of sizes and strides, combines the strides into a single
+// expression if it is possible.
+//
+// The current implementation expects that each size captures a single dimension
+// parameter.
+//
+// Let s be an n-dimensional shape that we want to fully collapse. In order to
+// be propagated successfully through the collapse, the pattern of the tiling of
+// s will have to look like the following (in row-major order):
+//   (1*, partial_dim?, full_dims*, 1*)
+// where full_dims are dimensions along which we capture all the elements
+// we can based on the corresponding stride, and partial_dim is a dimension that
+// can be captured with an arbitrary tile.
+//
+// In that case, the stride will be the stride corresponding to the minormost
+// dimension in which we capture more than a single element. This corresponds
+// to the size expression `e` with the smallest stride such that `e` evaluates
+// to another value than 1. Algorithmically, this can be represented as a series
+// of nested if statements:
+//   if size0 != 1 then stride0 else (if size1 != 1 then stride1 else ...)
+// where {size,stride}i = size_and_strides[i].{size,stride} (sizes_and_strides
+// being sorted in ascending order of stride).
+//
+// We generate this nest.
+//
+// If all the sizes are 1, then return a zero stride. Note that this
+// value is arbitrarily chosen.
+std::optional<AffineExpr> CombineStrides(
+    std::vector<SizeAndStrideExpression> sizes_and_strides,
+    absl::Span<Interval const> dimension_intervals) {
+  CHECK(!sizes_and_strides.empty());
+  for (const SizeAndStrideExpression& size_and_stride : sizes_and_strides) {
+    if (size_and_stride.stride.getKind() != AffineExprKind::Constant) {
+      VLOG(1) << "Attempted to combine non-constant stride: "
+              << AffineMapPrinter().ToString(size_and_stride.stride);
+      return std::nullopt;
+    }
+
+    // We know the exact bounds of dimension parameters, since they correspond
+    // to parameters of the initial indexing map. It follows that if a size
+    // expression is exactly a dimension parameter, we know its exact bounds.
+    //
+    // If a size is not exactly a dimension parameter, then it is dubious
+    // whether we know the bounds---and may thus calculate wrong strides.
+    if (size_and_stride.size.getKind() != AffineExprKind::DimId) {
+      VLOG(1) << "Attempted to combine strides but got non-dimension size "
+              << AffineMapPrinter().ToString(size_and_stride.size);
+      return std::nullopt;
+    }
+  }
+
+  SortByStride(sizes_and_strides);
+
+  for (auto [dim_id, size_and_stride] : llvm::enumerate(sizes_and_strides)) {
+    int64_t stride =
+        llvm::cast<AffineConstantExpr>(size_and_stride.stride).getValue();
+
+    // The minormost stride can be anything, but we expect every subsequent
+    // stride to be exactly `p_stride * p_size` where `p_size` is the upper
+    // bound of the size expression of the previous dimension and `p_stride` is
+    // its stride expression.
+    //
+    // For simplicity, we assume that each size expression captures a single
+    // dimension parameter.
+    if (dim_id > 0) {
+      const SizeAndStrideExpression& previous_size_and_stride =
+          sizes_and_strides[dim_id - 1];
+      const auto previous_dimension =
+          llvm::cast<AffineDimExpr>(previous_size_and_stride.size);
+      const Interval& previous_size_interval =
+          dimension_intervals[previous_dimension.getPosition()];
+      if (previous_size_interval.lower != 0) {
+        // TODO(bchetioui): I think we may need to handle this to have reshapes
+        // working well with concatenations. Nevertheless, we can take a look
+        // later.
+        VLOG(1) << "Attempted to combine strides but got dimension "
+                << AffineMapPrinter().ToString(previous_dimension)
+                << " with lower bound " << previous_size_interval.lower
+                << " != 0";
+        return std::nullopt;
+      }
+
+      int64_t previous_stride =
+          llvm::cast<AffineConstantExpr>(previous_size_and_stride.stride)
+              .getValue();
+      // We need to add 1 to the upper bound of the interval to describe the
+      // number of elements being captured, since the interval bounds are
+      // inclusive.
+      if ((previous_size_interval.upper + 1) * previous_stride != stride) {
+        VLOG(1) << "Attempted to combine strides but stride did not grow "
+                << "exactly as expected: got "
+                << (previous_size_interval.upper + 1) << " * "
+                << previous_stride << " != " << stride;
+        return std::nullopt;
+      }
+    }
+  }
+
+  // Produce a nested if statement as described in the function's documentation.
+  MLIRContext* ctx = sizes_and_strides[0].stride.getContext();
+  AffineExpr nested_if = getAffineConstantExpr(0, ctx);
+  for (auto size_and_stride_it = sizes_and_strides.rbegin();
+       size_and_stride_it != sizes_and_strides.rend(); ++size_and_stride_it) {
+    AffineExpr size = size_and_stride_it->size;
+    AffineExpr stride = size_and_stride_it->stride;
+    const Interval& size_interval =
+        dimension_intervals[llvm::cast<AffineDimExpr>(size).getPosition()];
+    nested_if = IfNeqOne(size, stride, nested_if, size_interval.upper + 1);
+  }
+
+  return nested_if;
+}
+
+// See documentation of `CombineSizes` and `CombineStrides` for an explanation
+// of how sizes and strides are combined.
+std::optional<SizeAndStrideExpression> CombineSizesAndStrides(
+    std::vector<SizeAndStrideExpression> sizes_and_strides,
+    absl::Span<Interval const> dimension_intervals) {
+  CHECK(!sizes_and_strides.empty());
+  AffineExpr size = CombineSizes(sizes_and_strides);
+  std::optional<AffineExpr> stride =
+      CombineStrides(std::move(sizes_and_strides), dimension_intervals);
+  if (!stride.has_value()) {
+    return std::nullopt;
+  }
+
+  // TODO(b/326998704): handle reshape constraints here.
+  return SizeAndStrideExpression{size, *stride};
+}
+
+std::optional<SizeAndStrideExpression> ExtractSizeAndStride(
+    AffineExpr strided_indexing, absl::Span<Interval const> dimension_intervals,
+    absl::Span<Interval const> symbol_intervals) {
   MLIRContext* ctx = strided_indexing.getContext();
   // Deal with the symbol case (capturing a whole untiled dimension).
   // TODO(b/330906085): concatenating across a reduction dimension needs to be
   // handled by this code.
-  if (auto symbol = llvm::dyn_cast<mlir::AffineSymbolExpr>(strided_indexing)) {
+  if (auto symbol = llvm::dyn_cast<AffineSymbolExpr>(strided_indexing)) {
     const Interval& symbol_interval = symbol_intervals[symbol.getPosition()];
     if (symbol_interval.lower != 0) {
       return std::nullopt;
@@ -199,11 +457,10 @@ std::optional<SizeAndStrideExpression> ExtractSizeAndStride(
       AffineExpr lhs = mul.getLHS();
       // The stride may not be fully collapsed if it is negative; in that case,
       // we need to extract the negative multiplier first.
-      if (const auto rhs =
-              llvm::dyn_cast<mlir::AffineConstantExpr>(mul.getRHS());
+      if (const auto rhs = llvm::dyn_cast<AffineConstantExpr>(mul.getRHS());
           rhs && rhs.getValue() == -1) {
         std::optional<SizeAndStrideExpression> maybe_size_and_stride =
-            ExtractSizeAndStride(lhs, symbol_intervals);
+            ExtractSizeAndStride(lhs, dimension_intervals, symbol_intervals);
         if (!maybe_size_and_stride.has_value()) {
           return std::nullopt;
         }
@@ -224,7 +481,7 @@ std::optional<SizeAndStrideExpression> ExtractSizeAndStride(
       auto floor_div = llvm::cast<mlir::AffineBinaryOpExpr>(strided_indexing);
       return ExtractSizeAndStrideFromFloorDiv(floor_div.getLHS(),
                                               floor_div.getRHS());
-    };
+    }
     case mlir::AffineExprKind::Constant:
       return SizeAndStrideExpression{/*size=*/getAffineConstantExpr(1, ctx),
                                      /*stride=*/getAffineConstantExpr(0, ctx)};
@@ -234,12 +491,19 @@ std::optional<SizeAndStrideExpression> ExtractSizeAndStride(
       // It's currently not checked separately, but RTVars shouldn't appear in
       // the strided indexing expressions.
       return std::nullopt;
-    case mlir::AffineExprKind::Add:
+    case mlir::AffineExprKind::Add: {
       // TODO(b/328427138): this should only be necessary in the multivariate
       // case, and will be implemented later.
-      VLOG(1) << "Encountered complex strided indexing expression "
-              << printer.ToString(strided_indexing);
-      return std::nullopt;
+      std::optional<std::vector<SizeAndStrideExpression>>
+          maybe_sizes_and_strides =
+              ExtractSizesAndStridesFromMultivariateSummation(
+                  strided_indexing, dimension_intervals, symbol_intervals);
+      if (!maybe_sizes_and_strides.has_value()) {
+        return std::nullopt;
+      }
+      return CombineSizesAndStrides(std::move(*maybe_sizes_and_strides),
+                                    dimension_intervals);
+    }
     case mlir::AffineExprKind::CeilDiv:
       break;
   };
@@ -323,6 +587,7 @@ AffineExpr SimplifyAffineExpr(const AffineExpr& expr,
     std::optional<SizeAndStrideExpression> maybe_size_and_stride =
         ExtractSizeAndStride(SimplifyAffineExpr(composite_indexing - offset,
                                                 /*reference=*/indexing_map),
+                             indexing_map.GetDimensionBounds(),
                              indexing_map.GetSymbolBounds());
     if (!maybe_size_and_stride.has_value()) {
       VLOG(1) << "No size and stride extracted";
@@ -333,19 +598,18 @@ AffineExpr SimplifyAffineExpr(const AffineExpr& expr,
   }
 
   // Eliminate negative strides and recalculate offsets.
+  // TODO(b/340555497): handle normalization of more complex expressions.
   std::vector<AffineExpr> dim_replacements, sym_replacements;
   for (auto [offset, size, stride] :
        llvm::zip(offset_expressions, size_expressions, stride_expressions)) {
-    auto constant = llvm::dyn_cast<mlir::AffineConstantExpr>(stride);
-    if (!constant) {
+    auto constant = llvm::dyn_cast<AffineConstantExpr>(stride);
+    if (constant && constant.getValue() < 0) {
+      offset = offset + size * stride - stride;
+      stride = -stride;
+    } else if (!constant) {
       AffineMapPrinter printer;
       VLOG(1) << "Unexpected non-constant stride expression: "
               << printer.ToString(stride);
-      return std::nullopt;
-    }
-    if (constant.getValue() < 0) {
-      offset = offset + size * stride - stride;
-      stride = -stride;
     }
   }
 
diff --git a/third_party/xla/xla/service/gpu/model/symbolic_tile_test.cc b/third_party/xla/xla/service/gpu/model/symbolic_tile_test.cc
index d3611b33d676db..97128095df6f4b 100644
--- a/third_party/xla/xla/service/gpu/model/symbolic_tile_test.cc
+++ b/third_party/xla/xla/service/gpu/model/symbolic_tile_test.cc
@@ -15,13 +15,22 @@ limitations under the License.
 
 #include "xla/service/gpu/model/symbolic_tile.h"
 
+#include <cstdint>
 #include <optional>
+#include <vector>
 
 #include <gmock/gmock.h>
 #include <gtest/gtest.h>
+#include "absl/log/check.h"
 #include "absl/strings/ascii.h"
 #include "absl/strings/str_cat.h"
 #include "absl/strings/string_view.h"
+#include "absl/types/span.h"
+#include "llvm/ADT/STLExtras.h"
+#include "llvm/ADT/SmallVector.h"
+#include "llvm/Support/Casting.h"
+#include "mlir/IR/AffineExpr.h"  // from @llvm-project
+#include "mlir/IR/AffineMap.h"  // from @llvm-project
 #include "xla/service/gpu/model/affine_map_printer.h"
 #include "xla/service/gpu/model/indexing_analysis.h"
 #include "xla/service/gpu/model/indexing_test_utils.h"
@@ -31,6 +40,10 @@ namespace xla {
 namespace gpu {
 namespace {
 
+using ::llvm::SmallVector;
+using ::mlir::AffineExpr;
+using ::mlir::AffineMap;
+using ::testing::ElementsAre;
 using ::testing::ExplainMatchResult;
 using ::testing::Optional;
 using ::testing::StrEq;
@@ -72,6 +85,29 @@ MATCHER_P3(MatchSymbolicTile, offset_map_string, size_map_string,
       arg, result_listener);
 }
 
+std::vector<int64_t> EvaluateMapAt(AffineMap affine_map,
+                                   absl::Span<int64_t const> parameters) {
+  CHECK_EQ(affine_map.getNumSymbols(), parameters.size());
+  CHECK_EQ(affine_map.getNumDims(), 0);
+
+  SmallVector<AffineExpr> symbol_replacements = llvm::to_vector(
+      llvm::map_range(parameters, [affine_map](const int64_t v) -> AffineExpr {
+        return mlir::getAffineConstantExpr(v, affine_map.getContext());
+      }));
+
+  AffineMap simplified_affine_map =
+      mlir::simplifyAffineMap(affine_map.replaceDimsAndSymbols(
+          /*dimReplacements=*/{}, symbol_replacements, /*numResultDims=*/0,
+          /*numResultSyms=*/0));
+
+  SmallVector<int64_t> results = llvm::to_vector(llvm::map_range(
+      simplified_affine_map.getResults(), [](AffineExpr result) -> int64_t {
+        return llvm::cast<mlir::AffineConstantExpr>(result).getValue();
+      }));
+
+  return std::vector<int64_t>(results.begin(), results.end());
+}
+
 using SymbolicTileTest = IndexingTestBase;
 
 TEST_F(SymbolicTileTest, CanPropagateTileFromDotOutputToInputs) {
@@ -134,6 +170,58 @@ TEST_F(SymbolicTileTest,
           "()[s0, s1] -> (0, 1, 1, 1)")));
 }
 
+TEST_F(SymbolicTileTest,
+       CanPropagateTileThroughNonTrivialSplitReshapeFromOutputToInput) {
+  auto input_indexing = GetOutputToInputIndexing(ParseAndGetRoot(R"(
+    HloModule m
+    ENTRY e {
+      p0 = f32[192,4]{1,0} parameter(0)
+      ROOT bitcast = s8[4,8,6,4]{3,2,1,0} bitcast(p0)
+    }
+  )"));
+
+  std::optional<SymbolicTile> symbolic_tile =
+      SymbolicTile::FromIndexingMap(*input_indexing.indexing_maps[0].begin());
+
+  EXPECT_THAT(
+      symbolic_tile,
+      Optional(MatchSymbolicTile(
+          "()[s0, s1, s2, s3] -> (0, 0)",
+          "()[s0, s1, s2, s3] -> ((s0 * s1) * s2, s3)",
+          // Collapsed dimensions force us to create nested conditionals, since
+          // the stride of the output corresponds to the stride of the minormost
+          // expression along which elements are captured in the composite
+          // expression. Hence, the resulting expression is very ugly.
+          "()[s0, s1, s2, s3] -> "
+          "(((-s2 + 7) floordiv 6) * (((-s1 + 9) floordiv 8) * "
+          "((-((-s0 + 5) floordiv 4) + 1) * 48) + "
+          "(-((-s1 + 9) floordiv 8) + 1) * 6) + -((-s2 + 7) floordiv 6) + 1, "
+          "1)")));
+
+  // Capturing elements along dimensions 0, 1, and 2 makes the stride equal to
+  // 1.
+  EXPECT_THAT(EvaluateMapAt(symbolic_tile->stride_map(), {4, 8, 6, 4}),
+              ElementsAre(1, 1));
+  // Capturing elements along dimension 2 makes the stride equal to 1.
+  EXPECT_THAT(EvaluateMapAt(symbolic_tile->stride_map(), {1, 1, 6, 4}),
+              ElementsAre(1, 1));
+  // Capturing elements only along dimension 1 makes the stride equal to
+  // the length of dimension 2 (6).
+  EXPECT_THAT(EvaluateMapAt(symbolic_tile->stride_map(), {1, 8, 1, 4}),
+              ElementsAre(6, 1));
+  // Capturing elements only along dimension 0 makes the stride equal to the
+  // product of the lengths of dimensions 1 and 2 (8 * 6).
+  EXPECT_THAT(EvaluateMapAt(symbolic_tile->stride_map(), {2, 1, 1, 4}),
+              ElementsAre(48, 1));
+  // Capturing elements along dimension 0 and dimension 1 makes the stride
+  // equal to the length of dimension 2 (6).
+  EXPECT_THAT(EvaluateMapAt(symbolic_tile->stride_map(), {2, 8, 1, 4}),
+              ElementsAre(6, 1));
+  // Capturing a single element in the collapsed dimensions makes the stride 0.
+  EXPECT_THAT(EvaluateMapAt(symbolic_tile->stride_map(), {1, 1, 1, 4}),
+              ElementsAre(0, 1));
+}
+
 TEST_F(SymbolicTileTest, FailsToPropagateTileThroughNonTrivialReshape) {
   auto input_indexing = GetOutputToInputIndexing(ParseAndGetRoot(R"(
     HloModule m

From d12d0e010cacd8d96ecce3deeda0fcca1874e01f Mon Sep 17 00:00:00 2001
From: Dmitri Gribenko <dmitrig@google.com>
Date: Fri, 17 May 2024 08:35:02 -0700
Subject: [PATCH 352/478] Integrate LLVM at llvm/llvm-project@a68d20e98605

Updates LLVM usage to match
[a68d20e98605](https://github.com/llvm/llvm-project/commit/a68d20e98605)

PiperOrigin-RevId: 634786916
---
 third_party/llvm/workspace.bzl | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/third_party/llvm/workspace.bzl b/third_party/llvm/workspace.bzl
index f017f7ecff0479..29dfa535cd78f4 100644
--- a/third_party/llvm/workspace.bzl
+++ b/third_party/llvm/workspace.bzl
@@ -4,8 +4,8 @@ load("//third_party:repo.bzl", "tf_http_archive")
 
 def repo(name):
     """Imports LLVM."""
-    LLVM_COMMIT = "5a20a07fce88d54cf01cafde489bfc2fc447acc0"
-    LLVM_SHA256 = "fb33549dd5fef2b3e8a214cbbbd5f23962a273c7c0871a70d1fe58fc41bdaded"
+    LLVM_COMMIT = "a68d20e986053ec571223a9f3ead3e146a27dc82"
+    LLVM_SHA256 = "4c53512522cfd625a75aa7f201e6932f18ebf3fbe8ee89e05883c3d30ca935ac"
 
     tf_http_archive(
         name = name,

From 42ac0dde8f7855c87dfea2217b03b8ee7e829ae9 Mon Sep 17 00:00:00 2001
From: Quentin Khan <qkhan@google.com>
Date: Fri, 17 May 2024 09:07:00 -0700
Subject: [PATCH 353/478] Clean-up TFLite benchmark tools TFLite resolver.

PiperOrigin-RevId: 634795383
---
 .../lite/tools/benchmark/benchmark_tflite_model.cc | 14 ++------------
 1 file changed, 2 insertions(+), 12 deletions(-)

diff --git a/tensorflow/lite/tools/benchmark/benchmark_tflite_model.cc b/tensorflow/lite/tools/benchmark/benchmark_tflite_model.cc
index d775122fe9c1fc..7728723940a481 100644
--- a/tensorflow/lite/tools/benchmark/benchmark_tflite_model.cc
+++ b/tensorflow/lite/tools/benchmark/benchmark_tflite_model.cc
@@ -1100,18 +1100,8 @@ TfLiteStatus BenchmarkTfLiteModel::LoadModel() {
 
 std::unique_ptr<tflite::OpResolver> BenchmarkTfLiteModel::GetOpResolver()
     const {
-  tflite::ops::builtin::BuiltinOpResolver* resolver = nullptr;
-  // When --use_xnnpack is explicitly set to false, skip applying the default
-  // XNNPACK delegate in TfLite runtime so that the original execution path
-  // based on the unmodified model graph is still exercised.
-  if (params_.HasParam("use_xnnpack") &&
-      params_.HasValueSet<bool>("use_xnnpack") &&
-      !params_.Get<bool>("use_xnnpack")) {
-    resolver =
-        new tflite::ops::builtin::BuiltinOpResolverWithoutDefaultDelegates();
-  } else {
-    resolver = new tflite::ops::builtin::BuiltinOpResolver();
-  }
+  tflite::ops::builtin::BuiltinOpResolver* resolver =
+      new tflite::ops::builtin::BuiltinOpResolverWithoutDefaultDelegates();
   RegisterSelectedOps(resolver);
   return std::unique_ptr<tflite::OpResolver>(resolver);
 }

From 4d73219fed7d243d0d43217f7343bbbc77b7702b Mon Sep 17 00:00:00 2001
From: Goran Flegar <gflegar@google.com>
Date: Fri, 17 May 2024 09:15:25 -0700
Subject: [PATCH 354/478] Integrate Triton up to 25b4212a9

PiperOrigin-RevId: 634797641
---
 .../temporary/fix_register_constraints.patch  | 46 -------------------
 third_party/triton/temporary/fp8_splat.patch  | 41 -----------------
 third_party/triton/temporary/series.bzl       |  5 +-
 third_party/triton/workspace.bzl              |  4 +-
 .../temporary/fix_register_constraints.patch  | 46 -------------------
 .../triton/temporary/fp8_splat.patch          | 41 -----------------
 .../third_party/triton/temporary/series.bzl   |  5 +-
 .../xla/third_party/triton/workspace.bzl      |  4 +-
 8 files changed, 6 insertions(+), 186 deletions(-)
 delete mode 100644 third_party/triton/temporary/fix_register_constraints.patch
 delete mode 100644 third_party/triton/temporary/fp8_splat.patch
 delete mode 100644 third_party/xla/third_party/triton/temporary/fix_register_constraints.patch
 delete mode 100644 third_party/xla/third_party/triton/temporary/fp8_splat.patch

diff --git a/third_party/triton/temporary/fix_register_constraints.patch b/third_party/triton/temporary/fix_register_constraints.patch
deleted file mode 100644
index 2caeac268e39db..00000000000000
--- a/third_party/triton/temporary/fix_register_constraints.patch
+++ /dev/null
@@ -1,46 +0,0 @@
-This should be fixed properly and upstreamed to OpenAI, tracked in b/336865530.
-
-diff --git a/third_party/nvidia/lib/NVGPUToLLVM/NVGPUToLLVMPass.cpp b/third_party/nvidia/lib/NVGPUToLLVM/NVGPUToLLVMPass.cpp
---- a/third_party/nvidia/lib/NVGPUToLLVM/NVGPUToLLVMPass.cpp
-+++ b/third_party/nvidia/lib/NVGPUToLLVM/NVGPUToLLVMPass.cpp
-@@ -480,10 +480,36 @@ public:
-   std::vector<std::string>
-   getOutputConstraints(ttn::WGMMAWaitGroupOp op) const {
-     auto outputStructType = cast<LLVM::LLVMStructType>(op.getType());
--    uint32_t numOutputRegs = outputStructType.getBody().size();
--    std::string output =
--        outputStructType.getBody().front().isF32() ? "=f" : "=r";
--    return std::vector<std::string>(numOutputRegs, output);
-+    std::vector<std::string> outputConstraints;
-+    outputConstraints.reserve(outputStructType.getBody().size());
-+    for (mlir::Type type : outputStructType.getBody()) {
-+      if (type.isF32()) {
-+        outputConstraints.push_back("=f");
-+        continue;
-+      } else if (type.isF64()) {
-+        outputConstraints.push_back("=d");
-+        continue;
-+      }
-+      unsigned bitwidth = isa<LLVM::LLVMPointerType>(type) ?
-+          64 : type.getIntOrFloatBitWidth();
-+      switch (bitwidth) {
-+        case 1:
-+          outputConstraints.push_back("=b");
-+          break;
-+        case 16:
-+          outputConstraints.push_back("=h");
-+          break;
-+        case 32:
-+          outputConstraints.push_back("=r");
-+          break;
-+        case 64:
-+          outputConstraints.push_back("=l");
-+          break;
-+        default:
-+          assert(false && "unsupported bitwidth");
-+      }
-+    }
-+    return outputConstraints;
-   }
- 
-   OperandsAndConstraints
diff --git a/third_party/triton/temporary/fp8_splat.patch b/third_party/triton/temporary/fp8_splat.patch
deleted file mode 100644
index 7a709e85419821..00000000000000
--- a/third_party/triton/temporary/fp8_splat.patch
+++ /dev/null
@@ -1,41 +0,0 @@
-diff --git a/include/triton/Conversion/MLIRTypes.h b/include/triton/Conversion/MLIRTypes.h
---- a/include/triton/Conversion/MLIRTypes.h
-+++ b/include/triton/Conversion/MLIRTypes.h
-@@ -26,12 +26,6 @@ inline Type f32Ty(MLIRContext *ctx) { re
- inline Type f64Ty(MLIRContext *ctx) { return FloatType::getF64(ctx); }
- inline Type bf16Ty(MLIRContext *ctx) { return FloatType::getBF16(ctx); }
- 
--inline bool isFloat(Type type) {
--  return type.isF32() || type.isF64() || type.isF16() || type.isF128();
--}
--
--inline bool isInt(Type type) { return type.isIntOrFloat() && !isFloat(type); }
--
- } // namespace type
- } // namespace triton
- } // namespace mlir
-diff --git a/lib/Conversion/TritonGPUToLLVM/ViewOpToLLVM.cpp b/lib/Conversion/TritonGPUToLLVM/ViewOpToLLVM.cpp
---- a/lib/Conversion/TritonGPUToLLVM/ViewOpToLLVM.cpp
-+++ b/lib/Conversion/TritonGPUToLLVM/ViewOpToLLVM.cpp
-@@ -74,17 +74,18 @@ struct ArithConstantSplatOpConversion
-     auto values = mlir::dyn_cast<SplatElementsAttr>(op.getValue());
-     auto elemType = values.getElementType();
-     Attribute val;
--    if (elemType.isBF16() || type::isFloat(elemType)) {
-+    if (isa<FloatType>(elemType)) {
-       val = values.getValues<FloatAttr>()[0];
--    } else if (type::isInt(elemType)) {
-+    } else if (isa<IntegerType>(elemType)) {
-       val = values.getValues<IntegerAttr>()[0];
-     } else {
-       llvm::errs() << "ArithConstantSplatOpConversion get unsupported type: "
-                    << value.getType() << "\n";
-       return failure();
-     }
--    auto constOp = rewriter.create<LLVM::ConstantOp>(loc, elemType, val);
-     auto typeConverter = getTypeConverter();
-+    auto constOp = rewriter.create<LLVM::ConstantOp>(
-+        loc, typeConverter->convertType(elemType), val);
-     auto llStruct = SplatOpConversion::convertSplatLikeOp(
-         elemType, op.getType(), constOp, typeConverter, rewriter, loc);
-     rewriter.replaceOp(op, llStruct);
diff --git a/third_party/triton/temporary/series.bzl b/third_party/triton/temporary/series.bzl
index a5dfa7d5eae3f6..53a5059fe432d7 100644
--- a/third_party/triton/temporary/series.bzl
+++ b/third_party/triton/temporary/series.bzl
@@ -5,7 +5,4 @@ These are created temporarily and should be moved to the first copybara workflow
 internal patch during the next triton integration process.
 """
 
-temporary_patch_list = [
-    "//third_party/triton/temporary:fp8_splat.patch",
-    "//third_party/triton/temporary:fix_register_constraints.patch",
-]
+temporary_patch_list = []
diff --git a/third_party/triton/workspace.bzl b/third_party/triton/workspace.bzl
index 4c21ab7d7c9099..a257f1f3e44645 100644
--- a/third_party/triton/workspace.bzl
+++ b/third_party/triton/workspace.bzl
@@ -8,8 +8,8 @@ load("//third_party/triton/xla_extensions:series.bzl", "extensions_files_patch_l
 def repo():
     """Imports Triton."""
 
-    TRITON_COMMIT = "cl633653394"
-    TRITON_SHA256 = "4bf1b25ce6f7ce6ab6fb4e2403b24c8c0bf71f6a9f5c08f0125fdb44bcb371a8"
+    TRITON_COMMIT = "cl634675237"
+    TRITON_SHA256 = "7151d057ee8443c2f45cbe18a7435a42f37e18f562e5d238b844b6e09fc560e6"
     tf_http_archive(
         name = "triton",
         sha256 = TRITON_SHA256,
diff --git a/third_party/xla/third_party/triton/temporary/fix_register_constraints.patch b/third_party/xla/third_party/triton/temporary/fix_register_constraints.patch
deleted file mode 100644
index 2caeac268e39db..00000000000000
--- a/third_party/xla/third_party/triton/temporary/fix_register_constraints.patch
+++ /dev/null
@@ -1,46 +0,0 @@
-This should be fixed properly and upstreamed to OpenAI, tracked in b/336865530.
-
-diff --git a/third_party/nvidia/lib/NVGPUToLLVM/NVGPUToLLVMPass.cpp b/third_party/nvidia/lib/NVGPUToLLVM/NVGPUToLLVMPass.cpp
---- a/third_party/nvidia/lib/NVGPUToLLVM/NVGPUToLLVMPass.cpp
-+++ b/third_party/nvidia/lib/NVGPUToLLVM/NVGPUToLLVMPass.cpp
-@@ -480,10 +480,36 @@ public:
-   std::vector<std::string>
-   getOutputConstraints(ttn::WGMMAWaitGroupOp op) const {
-     auto outputStructType = cast<LLVM::LLVMStructType>(op.getType());
--    uint32_t numOutputRegs = outputStructType.getBody().size();
--    std::string output =
--        outputStructType.getBody().front().isF32() ? "=f" : "=r";
--    return std::vector<std::string>(numOutputRegs, output);
-+    std::vector<std::string> outputConstraints;
-+    outputConstraints.reserve(outputStructType.getBody().size());
-+    for (mlir::Type type : outputStructType.getBody()) {
-+      if (type.isF32()) {
-+        outputConstraints.push_back("=f");
-+        continue;
-+      } else if (type.isF64()) {
-+        outputConstraints.push_back("=d");
-+        continue;
-+      }
-+      unsigned bitwidth = isa<LLVM::LLVMPointerType>(type) ?
-+          64 : type.getIntOrFloatBitWidth();
-+      switch (bitwidth) {
-+        case 1:
-+          outputConstraints.push_back("=b");
-+          break;
-+        case 16:
-+          outputConstraints.push_back("=h");
-+          break;
-+        case 32:
-+          outputConstraints.push_back("=r");
-+          break;
-+        case 64:
-+          outputConstraints.push_back("=l");
-+          break;
-+        default:
-+          assert(false && "unsupported bitwidth");
-+      }
-+    }
-+    return outputConstraints;
-   }
- 
-   OperandsAndConstraints
diff --git a/third_party/xla/third_party/triton/temporary/fp8_splat.patch b/third_party/xla/third_party/triton/temporary/fp8_splat.patch
deleted file mode 100644
index 7a709e85419821..00000000000000
--- a/third_party/xla/third_party/triton/temporary/fp8_splat.patch
+++ /dev/null
@@ -1,41 +0,0 @@
-diff --git a/include/triton/Conversion/MLIRTypes.h b/include/triton/Conversion/MLIRTypes.h
---- a/include/triton/Conversion/MLIRTypes.h
-+++ b/include/triton/Conversion/MLIRTypes.h
-@@ -26,12 +26,6 @@ inline Type f32Ty(MLIRContext *ctx) { re
- inline Type f64Ty(MLIRContext *ctx) { return FloatType::getF64(ctx); }
- inline Type bf16Ty(MLIRContext *ctx) { return FloatType::getBF16(ctx); }
- 
--inline bool isFloat(Type type) {
--  return type.isF32() || type.isF64() || type.isF16() || type.isF128();
--}
--
--inline bool isInt(Type type) { return type.isIntOrFloat() && !isFloat(type); }
--
- } // namespace type
- } // namespace triton
- } // namespace mlir
-diff --git a/lib/Conversion/TritonGPUToLLVM/ViewOpToLLVM.cpp b/lib/Conversion/TritonGPUToLLVM/ViewOpToLLVM.cpp
---- a/lib/Conversion/TritonGPUToLLVM/ViewOpToLLVM.cpp
-+++ b/lib/Conversion/TritonGPUToLLVM/ViewOpToLLVM.cpp
-@@ -74,17 +74,18 @@ struct ArithConstantSplatOpConversion
-     auto values = mlir::dyn_cast<SplatElementsAttr>(op.getValue());
-     auto elemType = values.getElementType();
-     Attribute val;
--    if (elemType.isBF16() || type::isFloat(elemType)) {
-+    if (isa<FloatType>(elemType)) {
-       val = values.getValues<FloatAttr>()[0];
--    } else if (type::isInt(elemType)) {
-+    } else if (isa<IntegerType>(elemType)) {
-       val = values.getValues<IntegerAttr>()[0];
-     } else {
-       llvm::errs() << "ArithConstantSplatOpConversion get unsupported type: "
-                    << value.getType() << "\n";
-       return failure();
-     }
--    auto constOp = rewriter.create<LLVM::ConstantOp>(loc, elemType, val);
-     auto typeConverter = getTypeConverter();
-+    auto constOp = rewriter.create<LLVM::ConstantOp>(
-+        loc, typeConverter->convertType(elemType), val);
-     auto llStruct = SplatOpConversion::convertSplatLikeOp(
-         elemType, op.getType(), constOp, typeConverter, rewriter, loc);
-     rewriter.replaceOp(op, llStruct);
diff --git a/third_party/xla/third_party/triton/temporary/series.bzl b/third_party/xla/third_party/triton/temporary/series.bzl
index a5dfa7d5eae3f6..53a5059fe432d7 100644
--- a/third_party/xla/third_party/triton/temporary/series.bzl
+++ b/third_party/xla/third_party/triton/temporary/series.bzl
@@ -5,7 +5,4 @@ These are created temporarily and should be moved to the first copybara workflow
 internal patch during the next triton integration process.
 """
 
-temporary_patch_list = [
-    "//third_party/triton/temporary:fp8_splat.patch",
-    "//third_party/triton/temporary:fix_register_constraints.patch",
-]
+temporary_patch_list = []
diff --git a/third_party/xla/third_party/triton/workspace.bzl b/third_party/xla/third_party/triton/workspace.bzl
index 4c21ab7d7c9099..a257f1f3e44645 100644
--- a/third_party/xla/third_party/triton/workspace.bzl
+++ b/third_party/xla/third_party/triton/workspace.bzl
@@ -8,8 +8,8 @@ load("//third_party/triton/xla_extensions:series.bzl", "extensions_files_patch_l
 def repo():
     """Imports Triton."""
 
-    TRITON_COMMIT = "cl633653394"
-    TRITON_SHA256 = "4bf1b25ce6f7ce6ab6fb4e2403b24c8c0bf71f6a9f5c08f0125fdb44bcb371a8"
+    TRITON_COMMIT = "cl634675237"
+    TRITON_SHA256 = "7151d057ee8443c2f45cbe18a7435a42f37e18f562e5d238b844b6e09fc560e6"
     tf_http_archive(
         name = "triton",
         sha256 = TRITON_SHA256,

From ee44de4ea0e4d6b19e3452843b0e8e38e20f250d Mon Sep 17 00:00:00 2001
From: Benjamin Chetioui <bchetioui@google.com>
Date: Fri, 17 May 2024 09:45:52 -0700
Subject: [PATCH 355/478] [XLA:GPU][NFC] Make SymbolicTile's
 `DestructureSummation` util more idiomatic.

PiperOrigin-RevId: 634805719
---
 .../xla/service/gpu/model/symbolic_tile.cc    | 36 ++++++++++---------
 1 file changed, 19 insertions(+), 17 deletions(-)

diff --git a/third_party/xla/xla/service/gpu/model/symbolic_tile.cc b/third_party/xla/xla/service/gpu/model/symbolic_tile.cc
index 8fca5c039347e6..836bc243954812 100644
--- a/third_party/xla/xla/service/gpu/model/symbolic_tile.cc
+++ b/third_party/xla/xla/service/gpu/model/symbolic_tile.cc
@@ -173,6 +173,24 @@ std::optional<SizeAndStrideExpression> ExtractSizeAndStrideFromFloorDiv(
   return std::nullopt;
 }
 
+// See documentation of `DestructureSummation` for an explanation of the
+// algorithm.
+void DestructureSummationImpl(AffineExpr expr,
+                              std::vector<AffineExpr>& summands) {
+  switch (expr.getKind()) {
+    case AffineExprKind::Add: {
+      const auto add = llvm::cast<mlir::AffineBinaryOpExpr>(expr);
+      DestructureSummationImpl(add.getLHS(), summands);
+      DestructureSummationImpl(add.getRHS(), summands);
+      break;
+    }
+    default:
+      // The expression is not a sum.
+      summands.push_back(expr);
+      break;
+  }
+}
+
 // Given an n-ary summation of expressions `e0 + e1 + ... + e{n-1}` with
 // arbitrary order of association, returns the vector `(e0, e1, ..., e{n-1})`.
 // The order of the returned subexpressions is not guaranteed to match the order
@@ -189,23 +207,7 @@ std::optional<SizeAndStrideExpression> ExtractSizeAndStrideFromFloorDiv(
 // has practical implications.
 std::vector<AffineExpr> DestructureSummation(AffineExpr expr) {
   std::vector<AffineExpr> summands;
-
-  switch (expr.getKind()) {
-    case AffineExprKind::Add: {
-      const auto add = llvm::cast<mlir::AffineBinaryOpExpr>(expr);
-
-      absl::c_move(DestructureSummation(add.getLHS()),
-                   std::back_inserter(summands));
-      absl::c_move(DestructureSummation(add.getRHS()),
-                   std::back_inserter(summands));
-      break;
-    }
-    default:
-      // The expression is not a sum.
-      summands.push_back(expr);
-      break;
-  }
-
+  DestructureSummationImpl(expr, summands);
   return summands;
 }
 

From 45fa10c3fc23062fa32331b5e6239ec74560badf Mon Sep 17 00:00:00 2001
From: Deqiang Chen <deqiangc@google.com>
Date: Fri, 17 May 2024 10:07:52 -0700
Subject: [PATCH 356/478] Make returned future immdiately ready if it is TPU
 used only. In such case, the variable tensor is obtained at runtime through
 the returned name instead and there is no need to wait for the restored
 tensor to be ready.

PiperOrigin-RevId: 634812708
---
 .../core/tfrt/mlrt/kernel/ifrt_ops_kernel.cc  | 36 +++++++++++--------
 1 file changed, 21 insertions(+), 15 deletions(-)

diff --git a/tensorflow/core/tfrt/mlrt/kernel/ifrt_ops_kernel.cc b/tensorflow/core/tfrt/mlrt/kernel/ifrt_ops_kernel.cc
index e0acdcce10df8c..899b189046ebfb 100644
--- a/tensorflow/core/tfrt/mlrt/kernel/ifrt_ops_kernel.cc
+++ b/tensorflow/core/tfrt/mlrt/kernel/ifrt_ops_kernel.cc
@@ -288,22 +288,28 @@ absl::Status MlrtIfrtLoadVariableKernel::InvokeHelper() {
   if (used_by_host()) {
     TF_RETURN_IF_ERROR(
         ifrt_restore_tensor_registry.SetUsedByHost(runtime_name));
-  }
 
-  xla::ifrt::Future<tensorflow::Tensor> restored_tensor_future =
-      ifrt_restore_tensor_registry.GetRestoredTensor(runtime_name);
-
-  restored_tensor_future.OnReady(
-      [tensor_promise = std::move(tensor_promise)](
-          absl::StatusOr<tensorflow::Tensor> restored_tensor) mutable {
-        if (!restored_tensor.ok()) {
-          std::move(tensor_promise).SetError(restored_tensor.status());
-          return;
-        }
-        std::move(tensor_promise)
-            .Set<tensorflow::tfrt_stub::FallbackTensor>(
-                tensorflow::tfrt_stub::FallbackTensor(*restored_tensor));
-      });
+    xla::ifrt::Future<tensorflow::Tensor> restored_tensor_future =
+        ifrt_restore_tensor_registry.GetRestoredTensor(runtime_name);
+
+    restored_tensor_future.OnReady(
+        [tensor_promise = std::move(tensor_promise)](
+            absl::StatusOr<tensorflow::Tensor> restored_tensor) mutable {
+          if (!restored_tensor.ok()) {
+            std::move(tensor_promise).SetError(restored_tensor.status());
+            return;
+          }
+          std::move(tensor_promise)
+              .Set<tensorflow::tfrt_stub::FallbackTensor>(
+                  tensorflow::tfrt_stub::FallbackTensor(*restored_tensor));
+        });
+  } else {
+    // If not used by host, set the future to be ready immediately with an empty
+    // tensor so that it does not block the graph execution.
+    std::move(tensor_promise)
+        .Set<tensorflow::tfrt_stub::FallbackTensor>(
+            tensorflow::tfrt_stub::FallbackTensor());
+  }
   // Return the name as the key
   tensorflow::Tensor key_tensor(tensorflow::DT_STRING, {});
   key_tensor.scalar<tsl::tstring>()() = runtime_name;

From d1d37b52a4a30e6a848cf2d83a6551e5f9dd2ef3 Mon Sep 17 00:00:00 2001
From: Eugene Zhulenev <ezhulenev@google.com>
Date: Fri, 17 May 2024 10:16:56 -0700
Subject: [PATCH 357/478] [stream_executor:host] Rework HostExecutor to avoid
 depending on LLVM by default

PiperOrigin-RevId: 634815409
---
 third_party/xla/xla/stream_executor/BUILD     |   2 +-
 .../xla/xla/stream_executor/host/BUILD        |  98 +++----
 .../host/host_execution_engine.h              | 157 -----------
 .../xla/stream_executor/host/host_executor.cc |  36 +--
 .../xla/stream_executor/host/host_executor.h  |  17 +-
 .../xla/stream_executor/host/host_kernel.cc   |   7 +-
 .../xla/stream_executor/host/host_kernel.h    |  33 ++-
 .../stream_executor/host/host_kernel_test.cc  |   2 +-
 ..._engine.cc => jit_host_kernel_function.cc} | 243 +++++++++++-------
 .../host/jit_host_kernel_function.h           |  53 ++++
 10 files changed, 326 insertions(+), 322 deletions(-)
 delete mode 100644 third_party/xla/xla/stream_executor/host/host_execution_engine.h
 rename third_party/xla/xla/stream_executor/host/{host_execution_engine.cc => jit_host_kernel_function.cc} (62%)
 create mode 100644 third_party/xla/xla/stream_executor/host/jit_host_kernel_function.h

diff --git a/third_party/xla/xla/stream_executor/BUILD b/third_party/xla/xla/stream_executor/BUILD
index 248876a1104577..d3cc88ebc1314f 100644
--- a/third_party/xla/xla/stream_executor/BUILD
+++ b/third_party/xla/xla/stream_executor/BUILD
@@ -574,7 +574,7 @@ cc_library(
     name = "kernel_spec",
     srcs = ["kernel_spec.cc"],
     hdrs = ["kernel_spec.h"],
-    visibility = ["//visibility:private"],
+    visibility = [":internal"],
     deps = [
         "@com_google_absl//absl/status:statusor",
         "@com_google_absl//absl/strings",
diff --git a/third_party/xla/xla/stream_executor/host/BUILD b/third_party/xla/xla/stream_executor/host/BUILD
index 6ad6b619df78c7..cd52f851ca9ec8 100644
--- a/third_party/xla/xla/stream_executor/host/BUILD
+++ b/third_party/xla/xla/stream_executor/host/BUILD
@@ -89,57 +89,11 @@ cc_library(
     hdrs = ["host_kernel_c_api.h"],
 )
 
-cc_library(
-    name = "host_execution_engine",
-    srcs = ["host_execution_engine.cc"],
-    hdrs = ["host_execution_engine.h"],
-    deps = [
-        ":host_kernel_c_api",
-        "@com_google_absl//absl/log:check",
-        "@com_google_absl//absl/status",
-        "@com_google_absl//absl/status:statusor",
-        "@com_google_absl//absl/strings",
-        "@com_google_absl//absl/strings:str_format",
-        "@com_google_absl//absl/synchronization",
-        "@com_google_absl//absl/types:span",
-        "@llvm-project//llvm:Analysis",
-        "@llvm-project//llvm:AsmParser",
-        "@llvm-project//llvm:Core",
-        "@llvm-project//llvm:ExecutionEngine",
-        "@llvm-project//llvm:JITLink",
-        "@llvm-project//llvm:OrcJIT",
-        "@llvm-project//llvm:OrcShared",
-        "@llvm-project//llvm:Passes",
-        "@llvm-project//llvm:Support",
-        "@llvm-project//llvm:Target",
-        "@llvm-project//llvm:TargetParser",
-        "@llvm-project//llvm:TransformUtils",
-        "@llvm-project//llvm:ir_headers",
-        "@local_tsl//tsl/platform:statusor",
-    ] + if_llvm_aarch32_available([
-        "@llvm-project//llvm:ARMAsmParser",
-        "@llvm-project//llvm:ARMCodeGen",
-    ]) + if_llvm_aarch64_available([
-        "@llvm-project//llvm:AArch64AsmParser",
-        "@llvm-project//llvm:AArch64CodeGen",
-    ]) + if_llvm_powerpc_available([
-        "@llvm-project//llvm:PowerPCAsmParser",
-        "@llvm-project//llvm:PowerPCCodeGen",
-    ]) + if_llvm_system_z_available([
-        "@llvm-project//llvm:SystemZAsmParser",
-        "@llvm-project//llvm:SystemZCodeGen",
-    ]) + if_llvm_x86_available([
-        "@llvm-project//llvm:X86AsmParser",
-        "@llvm-project//llvm:X86CodeGen",
-    ]),
-)
-
 cc_library(
     name = "host_kernel",
     srcs = ["host_kernel.cc"],
     hdrs = ["host_kernel.h"],
     deps = [
-        ":host_execution_engine",
         ":host_kernel_c_api",
         "//xla/stream_executor",
         "//xla/stream_executor:device_memory",
@@ -156,6 +110,7 @@ xla_cc_test(
         ":host_kernel",
         ":host_kernel_c_api",
         ":host_platform",
+        ":jit_host_kernel_function",
         "//xla/stream_executor",
         "//xla/stream_executor:device_memory",
         "//xla/stream_executor:kernel_factory",
@@ -176,7 +131,6 @@ cc_library(
         "host_executor.h",
     ],
     deps = [
-        ":host_execution_engine",
         ":host_kernel",
         ":host_stream",
         "//xla/stream_executor",
@@ -213,3 +167,53 @@ xla_cc_test(
         "@local_tsl//tsl/platform:test_main",
     ],
 )
+
+cc_library(
+    name = "jit_host_kernel_function",
+    srcs = ["jit_host_kernel_function.cc"],
+    hdrs = ["jit_host_kernel_function.h"],
+    deps = [
+        ":host_executor",
+        ":host_kernel",
+        ":host_kernel_c_api",
+        "//xla/stream_executor:kernel_spec",
+        "//xla/stream_executor/platform",
+        "@com_google_absl//absl/log:check",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/status:statusor",
+        "@com_google_absl//absl/strings",
+        "@com_google_absl//absl/strings:str_format",
+        "@com_google_absl//absl/synchronization",
+        "@com_google_absl//absl/types:span",
+        "@llvm-project//llvm:Analysis",
+        "@llvm-project//llvm:AsmParser",
+        "@llvm-project//llvm:Core",
+        "@llvm-project//llvm:ExecutionEngine",
+        "@llvm-project//llvm:JITLink",
+        "@llvm-project//llvm:OrcJIT",
+        "@llvm-project//llvm:OrcShared",
+        "@llvm-project//llvm:Passes",
+        "@llvm-project//llvm:Support",
+        "@llvm-project//llvm:Target",
+        "@llvm-project//llvm:TargetParser",
+        "@llvm-project//llvm:TransformUtils",
+        "@llvm-project//llvm:ir_headers",
+        "@local_tsl//tsl/platform:statusor",
+    ] + if_llvm_aarch32_available([
+        "@llvm-project//llvm:ARMAsmParser",
+        "@llvm-project//llvm:ARMCodeGen",
+    ]) + if_llvm_aarch64_available([
+        "@llvm-project//llvm:AArch64AsmParser",
+        "@llvm-project//llvm:AArch64CodeGen",
+    ]) + if_llvm_powerpc_available([
+        "@llvm-project//llvm:PowerPCAsmParser",
+        "@llvm-project//llvm:PowerPCCodeGen",
+    ]) + if_llvm_system_z_available([
+        "@llvm-project//llvm:SystemZAsmParser",
+        "@llvm-project//llvm:SystemZCodeGen",
+    ]) + if_llvm_x86_available([
+        "@llvm-project//llvm:X86AsmParser",
+        "@llvm-project//llvm:X86CodeGen",
+    ]),
+    alwayslink = 1,  # static kernel function loader registration
+)
diff --git a/third_party/xla/xla/stream_executor/host/host_execution_engine.h b/third_party/xla/xla/stream_executor/host/host_execution_engine.h
deleted file mode 100644
index 9215ad5ed5e3af..00000000000000
--- a/third_party/xla/xla/stream_executor/host/host_execution_engine.h
+++ /dev/null
@@ -1,157 +0,0 @@
-/* Copyright 2024 The OpenXLA Authors.
-
-Licensed under the Apache License, Version 2.0 (the "License");
-you may not use this file except in compliance with the License.
-You may obtain a copy of the License at
-
-    http://www.apache.org/licenses/LICENSE-2.0
-
-Unless required by applicable law or agreed to in writing, software
-distributed under the License is distributed on an "AS IS" BASIS,
-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.
-==============================================================================*/
-
-#ifndef XLA_STREAM_EXECUTOR_HOST_HOST_EXECUTION_ENGINE_H_
-#define XLA_STREAM_EXECUTOR_HOST_HOST_EXECUTION_ENGINE_H_
-
-#include <functional>
-#include <memory>
-#include <string>
-#include <utility>
-#include <vector>
-
-#include "absl/status/statusor.h"
-#include "absl/strings/string_view.h"
-#include "absl/types/span.h"
-#include "llvm/ExecutionEngine/JITEventListener.h"
-#include "llvm/ExecutionEngine/Orc/LLJIT.h"
-#include "llvm/ExecutionEngine/SectionMemoryManager.h"
-#include "llvm/IR/LLVMContext.h"
-#include "llvm/Support/CodeGen.h"
-#include "llvm/Support/Error.h"
-#include "llvm/Support/MemoryBuffer.h"
-#include "llvm/Target/TargetMachine.h"
-#include "xla/stream_executor/host/host_kernel_c_api.h"
-
-namespace stream_executor::host {
-
-class RuntimeExecutionEngine {
- public:
-  using ExportedFunctionPtr = const SE_HOST_Kernel *;
-
-  // Callback to run optimization passes on the compiled LLVM module.
-  using OptimizingTransformer = std::function<llvm::Error(llvm::Module *)>;
-
-  // Callback to construct an optimizing transformer for the given options.
-  using MakeOptimizingTransformer =
-      std::function<OptimizingTransformer(llvm::TargetMachine *targetMachine)>;
-
-  //------------------------------------------------------------------------- //
-  // Options for creating execution engine from an LLVM module.
-  //------------------------------------------------------------------------- //
-
-  struct JitOptions {
-    // User-provided codegen optimization level.
-    llvm::CodeGenOptLevel opt_level = llvm::CodeGenOptLevel::Default;
-
-    // User-provided target machine specification.
-    std::shared_ptr<llvm::TargetMachine> target_machine = nullptr;
-
-    // User-provided builder for the optimizing transformer.
-    MakeOptimizingTransformer make_optimizing_transformer;
-
-    // User-provided memory mapper for allocating memory for executables.
-    llvm::SectionMemoryManager::MemoryMapper *section_memory_mapper = nullptr;
-
-    // Notify the llvm's global GDB notifications listener.
-    bool enable_gdb_listener = false;
-
-    // Notify the llvm's global Perf notifications listener.
-    bool enable_perf_listener = false;
-  };
-
-  // Creates a new execution engine by compiling the provided LLVM module to
-  // a native executable using LLVM ORC stack.
-  static absl::StatusOr<std::unique_ptr<RuntimeExecutionEngine>>
-  CreateFromModule(std::unique_ptr<llvm::LLVMContext> ctx,
-                   std::unique_ptr<llvm::Module> module, JitOptions options,
-                   absl::Span<const std::string_view> exported);
-
-  //------------------------------------------------------------------------- //
-
-  // Returns a pointer to the exported function.
-  absl::Span<const ExportedFunctionPtr> exported() const { return exported_; }
-
-  ExportedFunctionPtr exported(unsigned ordinal) const {
-    return exported_[ordinal];
-  }
-
-  // Return a memory buffer with a object file behind this execution engine. Can
-  // be null if execution engine didn't save the compiled object file.
-  std::unique_ptr<llvm::MemoryBuffer> obj_file() const;
-
- private:
-  RuntimeExecutionEngine(bool enable_gdb_listener, bool enable_perf_listener);
-
-  // We build execution engine on top of the ORC LLJIT API, which owns all
-  // compiled/loaded object files and does the linking at run time.
-  //
-  // TODO(ezhulenev): Instead of keeping LLJIT alive we should be able to keep
-  // only llvm::orc::JITDylibSP owning main dylib and the object layer owning
-  // memory-mapped regions holding object files. Once we are done with
-  // executable compilation this jit is defunct because it holds an expired
-  // weak_ptr to an llvm::orc::TargetMachine instance.
-  std::unique_ptr<llvm::orc::LLJIT> jit_;
-
-  // Pointers to resolved exported functions. Indexed by function ordinal.
-  std::vector<ExportedFunctionPtr> exported_;
-
-  // Object file behind the compiled executable. Can be null.
-  std::unique_ptr<llvm::MemoryBuffer> obj_file_;
-
-  llvm::JITEventListener *gdb_listener_ = nullptr;
-  llvm::JITEventListener *perf_listener_ = nullptr;
-};
-
-// Virtual base class that owns jit-compiled function.
-class HostExecutionEngine {
- public:
-  virtual ~HostExecutionEngine() = default;
-  virtual SE_HOST_Kernel *kernel() const = 0;
-};
-
-class LlvmExecutionEngine : public HostExecutionEngine {
- public:
-  SE_HOST_Kernel *kernel() const override { return kernel_; }
-  // TODO(tsilytskyi): clean up kernel_
-  ~LlvmExecutionEngine() override = default;
-
-  static absl::StatusOr<std::unique_ptr<LlvmExecutionEngine>> CreateFromLlvmIr(
-      absl::string_view name, absl::string_view entry, absl::string_view ir,
-      absl::Span<const std::string> options);
-
- private:
-  explicit LlvmExecutionEngine(
-      std::unique_ptr<RuntimeExecutionEngine> exec_engine)
-      : engine_(std::move(exec_engine)) {
-    kernel_ = reinterpret_cast<SE_HOST_Kernel *>(engine_->exported(0));
-  };
-  std::unique_ptr<RuntimeExecutionEngine> engine_;
-  SE_HOST_Kernel *kernel_;
-};
-
-class CppExecutionEngine : public HostExecutionEngine {
- public:
-  ~CppExecutionEngine() override = default;
-  SE_HOST_Kernel *kernel() const override { return kernel_; }
-
- private:
-  CppExecutionEngine() = default;
-  SE_HOST_Kernel *kernel_ = nullptr;
-};
-
-}  // namespace stream_executor::host
-
-#endif  // XLA_STREAM_EXECUTOR_HOST_HOST_EXECUTION_ENGINE_H_
diff --git a/third_party/xla/xla/stream_executor/host/host_executor.cc b/third_party/xla/xla/stream_executor/host/host_executor.cc
index 8a9c812bb7ecdc..28b69aa213a0e1 100644
--- a/third_party/xla/xla/stream_executor/host/host_executor.cc
+++ b/third_party/xla/xla/stream_executor/host/host_executor.cc
@@ -24,6 +24,7 @@ limitations under the License.
 #include <memory>
 #include <string>
 #include <utility>
+#include <vector>
 
 #include "absl/functional/any_invocable.h"
 #include "absl/log/check.h"
@@ -32,12 +33,10 @@ limitations under the License.
 #include "absl/status/statusor.h"
 #include "absl/strings/string_view.h"
 #include "absl/synchronization/notification.h"
-#include "absl/types/span.h"
 #include "xla/stream_executor/device_description.h"
 #include "xla/stream_executor/device_memory.h"
 #include "xla/stream_executor/event.h"
 #include "xla/stream_executor/event_interface.h"
-#include "xla/stream_executor/host/host_execution_engine.h"
 #include "xla/stream_executor/host/host_kernel.h"
 #include "xla/stream_executor/host/host_stream.h"
 #include "xla/stream_executor/kernel_spec.h"
@@ -56,6 +55,16 @@ HostStream* AsHostStream(Stream* stream) {
   return dynamic_cast<HostStream*>(stream->implementation());
 }
 
+static std::vector<HostExecutor::KernelFunctionLoader>&
+KernelFunctionLoaderRegistry() {
+  static auto* registry = new std::vector<HostExecutor::KernelFunctionLoader>();
+  return *registry;
+}
+
+void HostExecutor::RegisterKernelFunctionLoader(KernelFunctionLoader loader) {
+  KernelFunctionLoaderRegistry().push_back(std::move(loader));
+}
+
 absl::Status HostExecutor::Init() { return absl::OkStatus(); }
 
 absl::StatusOr<std::unique_ptr<Kernel>> HostExecutor::CreateKernel() {
@@ -69,25 +78,16 @@ absl::Status HostExecutor::GetKernel(const MultiKernelLoaderSpec& spec,
 
   VLOG(3) << "GetKernel on kernel " << kernel << " : " << kernel->name();
 
-  if (spec.has_llvm_host_kernel()) {
-    const LlvmHostKernel& llvm_host_kernel = spec.llvm_host_kernel();
-    const absl::string_view name = llvm_host_kernel.kernel_name();
-    const absl::string_view entry = llvm_host_kernel.entrypoint();
-    const absl::string_view ir = llvm_host_kernel.ir();
-    const absl::Span<const std::string> options = llvm_host_kernel.options();
-
-    TF_ASSIGN_OR_RETURN(
-        auto execution_engine,
-        LlvmExecutionEngine::CreateFromLlvmIr(name, entry, ir, options));
-    host_kernel->SetExecutionEngine(std::move(execution_engine));
+  for (auto& loader : KernelFunctionLoaderRegistry()) {
+    auto loaded = loader(spec);
+    if (!loaded.has_value()) continue;
+
+    TF_ASSIGN_OR_RETURN(auto kernel_function, *std::move(loaded));
+    host_kernel->SetExecutionEngine(std::move(kernel_function));
     return absl::OkStatus();
-  } else if (false /* TODO(tsilytskyi): Implement CppHostKernel */) {
-    // host_kernel->SetExecutionEngine(std::make_unique<CppExecutionEngine>());
-  } else {
-    return absl::InternalError("No method of loading host kernel provided");
   }
 
-  return absl::UnimplementedError("Not Implemented");
+  return absl::InternalError("No method of loading host kernel provided");
 }
 
 absl::Status HostExecutor::Launch(Stream* stream, const ThreadDim& thread_dims,
diff --git a/third_party/xla/xla/stream_executor/host/host_executor.h b/third_party/xla/xla/stream_executor/host/host_executor.h
index c57a078e6f6b51..b1ce6f9e42db2a 100644
--- a/third_party/xla/xla/stream_executor/host/host_executor.h
+++ b/third_party/xla/xla/stream_executor/host/host_executor.h
@@ -21,7 +21,10 @@ limitations under the License.
 
 #include <cstddef>
 #include <cstdint>
+#include <functional>
 #include <memory>
+#include <optional>
+#include <variant>
 
 #include "absl/functional/any_invocable.h"
 #include "absl/status/status.h"
@@ -29,11 +32,13 @@ limitations under the License.
 #include "xla/stream_executor/device_description.h"
 #include "xla/stream_executor/device_memory.h"
 #include "xla/stream_executor/event.h"
+#include "xla/stream_executor/host/host_kernel.h"
 #include "xla/stream_executor/host_memory_allocation.h"
 #include "xla/stream_executor/kernel.h"
 #include "xla/stream_executor/kernel_spec.h"
 #include "xla/stream_executor/launch_dim.h"
 #include "xla/stream_executor/memory_allocation.h"
+#include "xla/stream_executor/platform.h"
 #include "xla/stream_executor/stream_executor.h"
 #include "xla/stream_executor/stream_executor_interface.h"
 
@@ -51,6 +56,15 @@ namespace host {
 // See stream_executor.h for description of the below operations.
 class HostExecutor : public StreamExecutor {
  public:
+  // A function that loads a kernel function from a given spec. If spec is not
+  // supported it returns an empty optional.
+  using KernelFunctionLoader = std::function<std::optional<
+      absl::StatusOr<std::unique_ptr<HostKernel::KernelFunction>>>(
+      const MultiKernelLoaderSpec& spec)>;
+
+  // Registers a kernel function loader in a static registry.
+  static void RegisterKernelFunctionLoader(KernelFunctionLoader loader);
+
   HostExecutor(Platform* platform, int device_ordinal)
       : StreamExecutor(platform), device_ordinal_(device_ordinal) {}
 
@@ -138,8 +152,7 @@ class HostExecutor : public StreamExecutor {
   std::unique_ptr<EventInterface> CreateEventImplementation() override;
 
   absl::StatusOr<std::unique_ptr<Stream>> CreateStream(
-      std::optional<std::variant<StreamPriority, int>> priority =
-          std::nullopt) override;
+      std::optional<std::variant<StreamPriority, int>> priority) override;
 
  private:
   int device_ordinal_;
diff --git a/third_party/xla/xla/stream_executor/host/host_kernel.cc b/third_party/xla/xla/stream_executor/host/host_kernel.cc
index b208ec453b9e0e..faee158cf8c03b 100644
--- a/third_party/xla/xla/stream_executor/host/host_kernel.cc
+++ b/third_party/xla/xla/stream_executor/host/host_kernel.cc
@@ -16,6 +16,7 @@ limitations under the License.
 #include "xla/stream_executor/host/host_kernel.h"
 
 #include <cstdint>
+#include <memory>
 #include <vector>
 
 #include "absl/status/status.h"
@@ -27,7 +28,7 @@ limitations under the License.
 namespace stream_executor::host {
 
 HostKernel::HostKernel(unsigned arity, SE_HOST_Kernel* kernel)
-    : arity_(arity), kernel_(kernel) {}
+    : function_(std::make_unique<KernelFunctionPtr>(kernel)), arity_(arity) {}
 
 absl::Status HostKernel::Launch(
     const ThreadDim& thread_dims,
@@ -46,6 +47,8 @@ absl::Status HostKernel::Launch(
   // for different threads (blocks) concurrently. For now it's the most trivial
   // implementation that runs tasks sequentially.
 
+  SE_HOST_Kernel* kernel = function_->kernel();
+
   for (uint64_t z = 0; z < thread_dims.z; ++z) {
     for (uint64_t y = 0; y < thread_dims.y; ++y) {
       for (uint64_t x = 0; x < thread_dims.x; ++x) {
@@ -54,7 +57,7 @@ absl::Status HostKernel::Launch(
         SE_HOST_KernelCallFrame call_frame = {
             &kernel_thread_dims, &kernel_thread, args.size(), args.data()};
 
-        SE_HOST_KernelError* error = (*kernel_)(&call_frame);
+        SE_HOST_KernelError* error = (*kernel)(&call_frame);
 
         if (error != nullptr) {
           return absl::InternalError("Failed to call host kernel");
diff --git a/third_party/xla/xla/stream_executor/host/host_kernel.h b/third_party/xla/xla/stream_executor/host/host_kernel.h
index b06fbe76b244f2..99d132e1e77e68 100644
--- a/third_party/xla/xla/stream_executor/host/host_kernel.h
+++ b/third_party/xla/xla/stream_executor/host/host_kernel.h
@@ -26,9 +26,9 @@ limitations under the License.
 #include "absl/status/statusor.h"
 #include "absl/types/span.h"
 #include "xla/stream_executor/device_memory.h"
-#include "xla/stream_executor/host/host_execution_engine.h"
 #include "xla/stream_executor/host/host_kernel_c_api.h"
 #include "xla/stream_executor/kernel.h"
+#include "xla/stream_executor/kernel_spec.h"
 #include "xla/stream_executor/launch_dim.h"
 
 namespace stream_executor::host {
@@ -37,6 +37,27 @@ class HostExecutor;
 
 class HostKernel : public Kernel {
  public:
+  // Virtual base class that owns the function behind the host kernel. It can be
+  // a function in a jit-compiled LLVM module or simply a pointer to the
+  // in-process function written in C++. HostKernel is responsible for launching
+  // the kernel function owned by the KernelFunction with given user-provided
+  // arguments potentially on a thread pool.
+  class KernelFunction {
+   public:
+    virtual ~KernelFunction() = default;
+    virtual SE_HOST_Kernel* kernel() const = 0;
+  };
+
+  // A wrapper around function pointer that implements SE_HOST_Kernel API.
+  class KernelFunctionPtr final : public KernelFunction {
+   public:
+    explicit KernelFunctionPtr(SE_HOST_Kernel* ptr) : ptr_(ptr) {}
+    SE_HOST_Kernel* kernel() const override { return ptr_; }
+
+   private:
+    SE_HOST_Kernel* ptr_;  // not owned
+  };
+
   explicit HostKernel() = default;
 
   // TODO(tsilytskyi): make this implementation detail private
@@ -58,18 +79,16 @@ class HostKernel : public Kernel {
   void SetArity(unsigned arity) { arity_ = arity; };
   unsigned Arity() const override { return arity_; };
 
-  template <typename T>
+  template <typename T,
+            std::enable_if_t<std::is_base_of_v<KernelFunction, T>>* = nullptr>
   void SetExecutionEngine(std::unique_ptr<T> execution_engine) {
-    static_assert(std::is_base_of<HostExecutionEngine, T>::value,
-                  "T is not derived from HostExecutionEngine");
-    execution_engine_ = std::move(execution_engine);
+    function_ = std::move(execution_engine);
   }
 
  private:
-  std::unique_ptr<HostExecutionEngine> execution_engine_;
+  std::unique_ptr<KernelFunction> function_;
 
   unsigned arity_;
-  SE_HOST_Kernel* kernel_ = nullptr;
 };
 
 inline const HostKernel* AsHostKernel(const Kernel* kernel) {
diff --git a/third_party/xla/xla/stream_executor/host/host_kernel_test.cc b/third_party/xla/xla/stream_executor/host/host_kernel_test.cc
index cebe2b3220c7fd..aa1bed20c361db 100644
--- a/third_party/xla/xla/stream_executor/host/host_kernel_test.cc
+++ b/third_party/xla/xla/stream_executor/host/host_kernel_test.cc
@@ -100,7 +100,7 @@ TEST(HostKernelTest, Addition) {
   EXPECT_EQ(out, expected);
 }
 
-TEST(HostKernelTest, LlvmAddition) {
+TEST(HostKernelTest, JitAddition) {
   std::vector<int32_t> lhs = {1, 2, 3, 4};
   std::vector<int32_t> rhs = {5, 6, 7, 8};
   std::vector<int32_t> out = {0, 0, 0, 0};
diff --git a/third_party/xla/xla/stream_executor/host/host_execution_engine.cc b/third_party/xla/xla/stream_executor/host/jit_host_kernel_function.cc
similarity index 62%
rename from third_party/xla/xla/stream_executor/host/host_execution_engine.cc
rename to third_party/xla/xla/stream_executor/host/jit_host_kernel_function.cc
index 0ed748e633bd05..c46ad5cd5771ed 100644
--- a/third_party/xla/xla/stream_executor/host/host_execution_engine.cc
+++ b/third_party/xla/xla/stream_executor/host/jit_host_kernel_function.cc
@@ -13,12 +13,13 @@ See the License for the specific language governing permissions and
 limitations under the License.
 ==============================================================================*/
 
-#include "xla/stream_executor/host/host_execution_engine.h"
+#include "xla/stream_executor/host/jit_host_kernel_function.h"
 
 #include <functional>
 #include <memory>
 #include <optional>
 #include <string>
+#include <string_view>
 #include <utility>
 #include <vector>
 
@@ -50,9 +51,11 @@ limitations under the License.
 #include "llvm/IR/DerivedTypes.h"
 #include "llvm/IR/Instructions.h"
 #include "llvm/IR/LLVMContext.h"
+#include "llvm/IR/Module.h"
 #include "llvm/IR/PassManager.h"
 #include "llvm/Passes/OptimizationLevel.h"
 #include "llvm/Passes/PassBuilder.h"
+#include "llvm/Support/CodeGen.h"
 #include "llvm/Support/Error.h"
 #include "llvm/Support/MemoryBuffer.h"
 #include "llvm/Support/SourceMgr.h"
@@ -61,13 +64,14 @@ limitations under the License.
 #include "llvm/Target/TargetMachine.h"
 #include "llvm/TargetParser/Triple.h"
 #include "llvm/Transforms/Utils/Cloning.h"
+#include "xla/stream_executor/host/host_executor.h"
+#include "xla/stream_executor/host/host_kernel.h"
+#include "xla/stream_executor/host/host_kernel_c_api.h"
+#include "xla/stream_executor/kernel_spec.h"
+#include "xla/stream_executor/platform/initialize.h"
 #include "tsl/platform/statusor.h"
 
-namespace stream_executor {
-namespace host {
-
-using absl::InternalError;
-using absl::StrFormat;
+namespace stream_executor::host {
 
 using llvm::Expected;
 using llvm::MemoryBuffer;
@@ -132,86 +136,116 @@ class WeakCompiler : public IRCompileLayer::IRCompiler {
 
 }  // namespace
 
-RuntimeExecutionEngine::RuntimeExecutionEngine(bool enable_gdb_listener,
-                                               bool enable_perf_listener) {
-  if (enable_gdb_listener)
-    gdb_listener_ = llvm::JITEventListener::createGDBRegistrationListener();
-  if (enable_perf_listener)
-    perf_listener_ = llvm::JITEventListener::createPerfJITEventListener();
-}
+namespace internal {
+// A minimal LLVM ORC JIT compilation stack to jit-compile LLVM modules to
+// executable functions.
+class ExecutionEngine {
+ public:
+  using ExportedFunctionPtr = const SE_HOST_Kernel *;
 
-std::unique_ptr<MemoryBuffer> RuntimeExecutionEngine::obj_file() const {
-  return obj_file_ ? MemoryBuffer::getMemBuffer(obj_file_->getMemBufferRef())
-                   : nullptr;
-}
+  // Callback to run optimization passes on the compiled LLVM module.
+  using OptimizingTransformer = std::function<llvm::Error(llvm::Module *)>;
 
-// -------------------------------------------------------------------------- //
+  // Callback to construct an optimizing transformer for the given options.
+  using MakeOptimizingTransformer =
+      std::function<OptimizingTransformer(llvm::TargetMachine *targetMachine)>;
 
-namespace {
-using llvm::DenseMap;
+  // Options for creating execution engine from an LLVM module.
+  struct Options {
+    // User-provided codegen optimization level.
+    llvm::CodeGenOptLevel opt_level = llvm::CodeGenOptLevel::Default;
 
-// Intercept object compilation to save the object file corresponding to the
-// XLA executable in the execution engine.
-class ExecutionEngineObjectCache : public llvm::ObjectCache {
- public:
-  void notifyObjectCompiled(const llvm::Module *m,
-                            llvm::MemoryBufferRef objBuffer) override;
+    // User-provided target machine specification.
+    std::shared_ptr<llvm::TargetMachine> target_machine = nullptr;
 
-  std::unique_ptr<llvm::MemoryBuffer> getObject(const llvm::Module *m) override;
+    // User-provided builder for the optimizing transformer.
+    MakeOptimizingTransformer make_optimizing_transformer;
 
-  // Transfer memory buffer from the cache to the caller.
-  std::unique_ptr<llvm::MemoryBuffer> stealObject(const llvm::Module *m);
+    // User-provided memory mapper for allocating memory for executables.
+    llvm::SectionMemoryManager::MemoryMapper *section_memory_mapper = nullptr;
+
+    // Notify the llvm's global GDB notifications listener.
+    bool enable_gdb_listener = false;
+
+    // Notify the llvm's global Perf notifications listener.
+    bool enable_perf_listener = false;
+  };
+
+  // Creates a new execution engine by compiling the provided LLVM module to
+  // a native executable using LLVM ORC stack.
+  static absl::StatusOr<std::unique_ptr<ExecutionEngine>> CreateFromModule(
+      std::unique_ptr<llvm::LLVMContext> ctx,
+      std::unique_ptr<llvm::Module> module, Options options,
+      absl::Span<const std::string_view> exported);
+
+  // Returns a pointer to the exported function.
+  absl::Span<const ExportedFunctionPtr> exported() const { return exported_; }
+
+  ExportedFunctionPtr exported(unsigned ordinal) const {
+    return exported_[ordinal];
+  }
+
+  // Return a memory buffer with a object file behind this execution engine. Can
+  // be null if execution engine didn't save the compiled object file.
+  std::unique_ptr<llvm::MemoryBuffer> obj_file() const;
 
  private:
-  DenseMap<const llvm::Module *, std::unique_ptr<llvm::MemoryBuffer>> objs_;
-};
-}  // namespace
+  ExecutionEngine(bool enable_gdb_listener, bool enable_perf_listener);
 
-void ExecutionEngineObjectCache::notifyObjectCompiled(
-    const llvm::Module *m, llvm::MemoryBufferRef objBuffer) {
-  objs_[m] = llvm::MemoryBuffer::getMemBufferCopy(
-      objBuffer.getBuffer(), objBuffer.getBufferIdentifier());
-}
+  // We build execution engine on top of the ORC LLJIT API, which owns all
+  // compiled/loaded object files and does the linking at run time.
+  //
+  // TODO(ezhulenev): Instead of keeping LLJIT alive we should be able to keep
+  // only llvm::orc::JITDylibSP owning main dylib and the object layer owning
+  // memory-mapped regions holding object files. Once we are done with
+  // executable compilation this jit is defunct because it holds an expired
+  // weak_ptr to an llvm::orc::TargetMachine instance.
+  std::unique_ptr<llvm::orc::LLJIT> jit_;
 
-std::unique_ptr<llvm::MemoryBuffer> ExecutionEngineObjectCache::getObject(
-    const llvm::Module *m) {
-  auto it = objs_.find(m);
-  if (it == objs_.end()) return nullptr;
-  return llvm::MemoryBuffer::getMemBuffer(it->second->getMemBufferRef());
-}
+  // Pointers to resolved exported functions. Indexed by function ordinal.
+  std::vector<ExportedFunctionPtr> exported_;
+
+  // Object file behind the compiled executable. Can be null.
+  std::unique_ptr<llvm::MemoryBuffer> obj_file_;
+
+  llvm::JITEventListener *gdb_listener_ = nullptr;
+  llvm::JITEventListener *perf_listener_ = nullptr;
+};
 
-std::unique_ptr<llvm::MemoryBuffer> ExecutionEngineObjectCache::stealObject(
-    const llvm::Module *m) {
-  auto it = objs_.find(m);
-  if (it == objs_.end()) return nullptr;
-  return std::move(it->second);
+ExecutionEngine::ExecutionEngine(bool enable_gdb_listener,
+                                 bool enable_perf_listener) {
+  if (enable_gdb_listener)
+    gdb_listener_ = llvm::JITEventListener::createGDBRegistrationListener();
+  if (enable_perf_listener)
+    perf_listener_ = llvm::JITEventListener::createPerfJITEventListener();
 }
 
-// -------------------------------------------------------------------------- //
+std::unique_ptr<MemoryBuffer> ExecutionEngine::obj_file() const {
+  return obj_file_ ? MemoryBuffer::getMemBuffer(obj_file_->getMemBufferRef())
+                   : nullptr;
+}
 
-// llvm_ir::DumpToString() is not used here, because we don't want to add too
-// many dependencies to the runtime.
 static std::string ToString(const llvm::Error &err) {
   std::string str;
   llvm::raw_string_ostream(str) << err;
   return str;
 }
 
-/*static*/ absl::StatusOr<std::unique_ptr<RuntimeExecutionEngine>>
-RuntimeExecutionEngine::CreateFromModule(
-    std::unique_ptr<llvm::LLVMContext> ctx,
-    std::unique_ptr<llvm::Module> module, JitOptions options,
-    absl::Span<const std::string_view> exported) {
-  auto engine =
-      std::unique_ptr<RuntimeExecutionEngine>(new RuntimeExecutionEngine(
-          options.enable_gdb_listener, options.enable_perf_listener));
+absl::StatusOr<std::unique_ptr<ExecutionEngine>>
+ExecutionEngine::CreateFromModule(std::unique_ptr<llvm::LLVMContext> ctx,
+                                  std::unique_ptr<llvm::Module> module,
+                                  Options options,
+                                  absl::Span<const std::string_view> exported) {
+  auto engine = std::unique_ptr<ExecutionEngine>(new ExecutionEngine(
+      options.enable_gdb_listener, options.enable_perf_listener));
 
   // We'll need module pointer later to lookup object file in the cache.
   llvm::Module *module_ptr = module.get();
 
   // Set up the target machine details.
-  if (!options.target_machine)
-    return InternalError("target machine was not provided");
+  if (!options.target_machine) {
+    return absl::InternalError("target machine was not provided");
+  }
   module->setDataLayout(options.target_machine->createDataLayout());
   module->setTargetTriple(options.target_machine->getTargetTriple().str());
 
@@ -224,8 +258,8 @@ RuntimeExecutionEngine::CreateFromModule(
   auto transformer =
       options.make_optimizing_transformer(options.target_machine.get());
   if (auto err = transformer(module_ptr))
-    return InternalError(
-        StrFormat("failed to run optimization pipeline: %s", ToString(err)));
+    return absl::InternalError(absl::StrFormat(
+        "failed to run optimization pipeline: %s", ToString(err)));
 
   // Callback to create the object layer with a user-provided section memory
   // mapper and JIT event listeners.
@@ -256,9 +290,10 @@ RuntimeExecutionEngine::CreateFromModule(
   auto executorProcessControl = SelfExecutorProcessControl::Create(
       nullptr, std::make_unique<InPlaceTaskDispatcher>());
 
-  if (auto err = executorProcessControl.takeError())
-    return InternalError(StrFormat(
+  if (auto err = executorProcessControl.takeError()) {
+    return absl::InternalError(absl::StrFormat(
         "failed to create executor process control: %s", ToString(err)));
+  }
 
   // TODO(b/286475799): Concurrent compilation leads to spurious memory
   // corruptions and segfaults at run time, however nothing shows up in tsan
@@ -274,17 +309,19 @@ RuntimeExecutionEngine::CreateFromModule(
                  .setNumCompileThreads(0)  // disable multi-threading
                  .create();
 
-  if (auto err = jit.takeError())
-    return InternalError(
-        StrFormat("failed to construct LLJIT: %s", ToString(err)));
+  if (auto err = jit.takeError()) {
+    return absl::InternalError(
+        absl::StrFormat("failed to construct LLJIT: %s", ToString(err)));
+  }
 
   lljit_lock.reset();
 
   // Register input module with the LLJIT.
   ThreadSafeModule tsm(std::move(module), std::move(ctx));
-  if (auto err = (*jit)->addIRModule(std::move(tsm)))
-    return InternalError(
-        StrFormat("failed to add source module: %s", ToString(err)));
+  if (auto err = (*jit)->addIRModule(std::move(tsm))) {
+    return absl::InternalError(
+        absl::StrFormat("failed to add source module: %s", ToString(err)));
+  }
 
   llvm::DataLayout data_layout = (*jit)->getDataLayout();
 
@@ -295,15 +332,17 @@ RuntimeExecutionEngine::CreateFromModule(
     //   - Do we need to mangle function name?
     //   - Do we need to verify/adapt function proto to expected API?
     Expected<ExecutorAddr> addr = (*jit)->lookup(name);
-    if (auto err = addr.takeError())
-      return InternalError(StrFormat(
+    if (auto err = addr.takeError()) {
+      return absl::InternalError(absl::StrFormat(
           "failed to compile exported function %s: %s", name, ToString(err)));
+    }
 
     // Check that we found an address of an exported function.
     auto ptr = addr->toPtr<ExportedFunctionPtr>();
-    if (!ptr)
-      return InternalError(
-          StrFormat("exported function %s resolved to null", name));
+    if (!ptr) {
+      return absl::InternalError(
+          absl::StrFormat("exported function %s resolved to null", name));
+    }
 
     engine->exported_.push_back(ptr);
   }
@@ -313,6 +352,14 @@ RuntimeExecutionEngine::CreateFromModule(
   return std::move(engine);
 }
 
+}  // namespace internal
+
+JitHostKernelFunction::JitHostKernelFunction(
+    std::unique_ptr<internal::ExecutionEngine> exec_engine)
+    : engine_(std::move(exec_engine)) {
+  kernel_ = reinterpret_cast<SE_HOST_Kernel *>(engine_->exported(0));
+};
+
 static std::function<llvm::Error(llvm::Module *)>
 MakeOptimizingTransformerForJit(llvm::TargetMachine *targetMachine) {
   return [targetMachine](llvm::Module *m) -> llvm::Error {
@@ -342,11 +389,11 @@ MakeOptimizingTransformerForJit(llvm::TargetMachine *targetMachine) {
   };
 }
 
-absl::StatusOr<std::unique_ptr<LlvmExecutionEngine>>
-LlvmExecutionEngine::CreateFromLlvmIr(absl::string_view name,
-                                      absl::string_view entry,
-                                      absl::string_view ir,
-                                      absl::Span<const std::string> options) {
+absl::StatusOr<std::unique_ptr<HostKernel::KernelFunction>>
+JitHostKernelFunction::CreateFromLlvmIr(absl::string_view name,
+                                        absl::string_view entry,
+                                        absl::string_view ir,
+                                        absl::Span<const std::string> options) {
   llvm::InitializeNativeTarget();
   llvm::InitializeNativeTargetAsmPrinter();
   auto llvm_ctx = std::make_unique<llvm::LLVMContext>();
@@ -369,7 +416,7 @@ LlvmExecutionEngine::CreateFromLlvmIr(absl::string_view name,
       llvm::StringRef(target_machine.get()->getTargetTriple().getTriple()));
 
   // Construct options for the XLA runtime execution engine.
-  RuntimeExecutionEngine::JitOptions engine_options;
+  internal::ExecutionEngine::Options engine_options;
   engine_options.target_machine = std::move(target_machine.get());
   engine_options.make_optimizing_transformer = MakeOptimizingTransformerForJit;
 
@@ -377,13 +424,35 @@ LlvmExecutionEngine::CreateFromLlvmIr(absl::string_view name,
 
   // Compile input module to the native function.
   TF_ASSIGN_OR_RETURN(auto engine,
-                      RuntimeExecutionEngine::CreateFromModule(
+                      internal::ExecutionEngine::CreateFromModule(
                           std::move(llvm_ctx), std::move(llvm_module),
                           std::move(engine_options), exported));
 
-  return std::unique_ptr<LlvmExecutionEngine>(
-      new LlvmExecutionEngine(std::move(engine)));
+  return std::unique_ptr<HostKernel::KernelFunction>(
+      new JitHostKernelFunction(std::move(engine)));
 }
 
-}  // namespace host
-}  // namespace stream_executor
+static void RegisterJitKernelFunctionLoader() {
+  using CompiledFunction = std::optional<
+      absl::StatusOr<std::unique_ptr<HostKernel::KernelFunction>>>;
+
+  HostExecutor::RegisterKernelFunctionLoader(
+      [](const MultiKernelLoaderSpec &spec) -> CompiledFunction {
+        if (!spec.has_llvm_host_kernel()) return std::nullopt;
+
+        const LlvmHostKernel &llvm_host_kernel = spec.llvm_host_kernel();
+        std::string_view name = llvm_host_kernel.kernel_name();
+        std::string_view entry = llvm_host_kernel.entrypoint();
+        std::string_view ir = llvm_host_kernel.ir();
+        absl::Span<const std::string> options = llvm_host_kernel.options();
+
+        return JitHostKernelFunction::CreateFromLlvmIr(name, entry, ir,
+                                                       options);
+      });
+}
+
+}  // namespace stream_executor::host
+
+STREAM_EXECUTOR_REGISTER_MODULE_INITIALIZER(
+    jot_kernel_function_loader,
+    stream_executor::host::RegisterJitKernelFunctionLoader());
diff --git a/third_party/xla/xla/stream_executor/host/jit_host_kernel_function.h b/third_party/xla/xla/stream_executor/host/jit_host_kernel_function.h
new file mode 100644
index 00000000000000..73991761db0109
--- /dev/null
+++ b/third_party/xla/xla/stream_executor/host/jit_host_kernel_function.h
@@ -0,0 +1,53 @@
+/* Copyright 2024 The OpenXLA Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+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.
+==============================================================================*/
+
+#ifndef XLA_STREAM_EXECUTOR_HOST_JIT_HOST_KERNEL_FUNCTION_H_
+#define XLA_STREAM_EXECUTOR_HOST_JIT_HOST_KERNEL_FUNCTION_H_
+
+#include <memory>
+#include <string>
+
+#include "absl/status/statusor.h"
+#include "absl/strings/string_view.h"
+#include "absl/types/span.h"
+#include "xla/stream_executor/host/host_kernel.h"
+#include "xla/stream_executor/host/host_kernel_c_api.h"
+
+namespace stream_executor::host {
+
+namespace internal {
+class ExecutionEngine;
+}
+
+// A host kernel function compiled from LLVM IR at run time
+class JitHostKernelFunction : public HostKernel::KernelFunction {
+ public:
+  SE_HOST_Kernel *kernel() const override { return kernel_; }
+
+  static absl::StatusOr<std::unique_ptr<HostKernel::KernelFunction>>
+  CreateFromLlvmIr(absl::string_view name, absl::string_view entry,
+                   absl::string_view ir, absl::Span<const std::string> options);
+
+ private:
+  explicit JitHostKernelFunction(
+      std::unique_ptr<internal::ExecutionEngine> exec_engine);
+
+  std::unique_ptr<internal::ExecutionEngine> engine_;
+  SE_HOST_Kernel *kernel_;
+};
+
+}  // namespace stream_executor::host
+
+#endif  // XLA_STREAM_EXECUTOR_HOST_JIT_HOST_KERNEL_FUNCTION_H_

From 317bee106b3fd214b1053b9b2d04bc7c4b9bd04c Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Fri, 17 May 2024 10:19:21 -0700
Subject: [PATCH 358/478] Fix up some unreachable code warnings on mobile
 platforms

PiperOrigin-RevId: 634816157
---
 tensorflow/lite/kernels/internal/optimized/cpu_check.cc | 4 ++--
 tensorflow/lite/kernels/kernel_util.cc                  | 7 +++----
 2 files changed, 5 insertions(+), 6 deletions(-)

diff --git a/tensorflow/lite/kernels/internal/optimized/cpu_check.cc b/tensorflow/lite/kernels/internal/optimized/cpu_check.cc
index 8fd17a7e33a03e..565299a8004acf 100644
--- a/tensorflow/lite/kernels/internal/optimized/cpu_check.cc
+++ b/tensorflow/lite/kernels/internal/optimized/cpu_check.cc
@@ -42,9 +42,9 @@ bool DetectDotprodByLinuxAuxvMethod() {
 bool DetectArmNeonDotprod() {
 #if defined __linux__ && defined __aarch64__
   return DetectDotprodByLinuxAuxvMethod();
-#endif
-
+#else
   return false;
+#endif
 }
 
 }  // namespace tflite
diff --git a/tensorflow/lite/kernels/kernel_util.cc b/tensorflow/lite/kernels/kernel_util.cc
index 58fd99f8fb15be..39f7bc7da53a49 100644
--- a/tensorflow/lite/kernels/kernel_util.cc
+++ b/tensorflow/lite/kernels/kernel_util.cc
@@ -572,12 +572,11 @@ int TfLiteTypeGetSize(TfLiteType type) {
 bool IsMobilePlatform() {
 #if defined(ANDROID) || defined(__ANDROID__)
   return true;
-#elif defined(__APPLE__)
-#if TARGET_IPHONE_SIMULATOR || TARGET_OS_IPHONE
+#elif defined(__APPLE__) && (TARGET_IPHONE_SIMULATOR || TARGET_OS_IPHONE)
   return true;
-#endif
-#endif
+#else
   return false;
+#endif
 }
 
 bool HasUnspecifiedDimension(const TfLiteTensor* tensor) {

From 37dba8b98931e371d983b2d6e6db1b9f99509bba Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Fri, 17 May 2024 10:24:20 -0700
Subject: [PATCH 359/478] disables `detect_odr_violation` for
 `validator_runner_test`

Runtimes-on-demand is causing AddressSanitizer to be unhappy with `validator_runner_test` due to some ODR violation. Given that we don't use what's causing this error, we're turning it off.

PiperOrigin-RevId: 634817533
---
 .../lite/experimental/acceleration/mini_benchmark/BUILD    | 7 +++++++
 1 file changed, 7 insertions(+)

diff --git a/tensorflow/lite/experimental/acceleration/mini_benchmark/BUILD b/tensorflow/lite/experimental/acceleration/mini_benchmark/BUILD
index 945eab9ffe521f..a9990f90a1b7f3 100644
--- a/tensorflow/lite/experimental/acceleration/mini_benchmark/BUILD
+++ b/tensorflow/lite/experimental/acceleration/mini_benchmark/BUILD
@@ -1103,6 +1103,13 @@ cc_test(
 cc_test(
     name = "validator_runner_test",
     srcs = ["validator_runner_test.cc"],
+    # Shared libraries that are built and loaded into internal Google applications are generally
+    # unsupported. One consequence of this lack of support is that AddressSanitizer's ODR checking
+    # will produce spurious failures with an upcoming Crosstool change. The workaround for this is
+    # to disable the ODR checking on affected tests.
+    env = {
+        "ASAN_OPTIONS": "detect_odr_violation=0",
+    },
     tags = [
         "no_mac",
         "no_windows",

From db3425ca7c2eb95c4000ba502fc406cc0435d32d Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Fri, 17 May 2024 10:24:47 -0700
Subject: [PATCH 360/478] Migrate away from llvm::StringRef::equals

Note that llvm::StringRef::equals has been deprecated upstream.

PiperOrigin-RevId: 634817679
---
 .../odml_converter/transforms/outline_composites.cc      | 2 +-
 .../compiler/mlir/lite/utils/perception_ops_utils.cc     | 7 +++----
 .../mlir/quantization/common/lift_as_function_call.cc    | 2 +-
 .../passes/bridge/convert_tf_quant_ops_to_mhlo.cc        | 4 ++--
 .../mlir/quantization/tensorflow/cc/constant_fold.cc     | 2 +-
 .../tensorflow/passes/insert_custom_aggregation_ops.cc   | 6 +++---
 .../tensorflow/passes/lift_hashtable_ops_as_args.cc      | 9 +++++----
 .../quantization/tensorflow/passes/prepare_lifting.cc    | 4 ++--
 .../tensorflow/utils/tf_to_xla_attribute_utils.cc        | 6 +++---
 .../transforms/executor_tpuv1_island_coarsening.cc       | 2 +-
 .../compiler/mlir/tensorflow/transforms/tfg-to-tfe.cc    | 2 +-
 .../compiler/mlir/tensorflow/utils/export_utils.cc       | 6 +++---
 12 files changed, 26 insertions(+), 26 deletions(-)

diff --git a/tensorflow/compiler/mlir/lite/stablehlo/odml_converter/transforms/outline_composites.cc b/tensorflow/compiler/mlir/lite/stablehlo/odml_converter/transforms/outline_composites.cc
index b66b8f7e51ff60..821ba4fa7e4d2f 100644
--- a/tensorflow/compiler/mlir/lite/stablehlo/odml_converter/transforms/outline_composites.cc
+++ b/tensorflow/compiler/mlir/lite/stablehlo/odml_converter/transforms/outline_composites.cc
@@ -94,7 +94,7 @@ bool HasSplatArg(Operation* op, float val, int opr_num) {
 // Determines if the given op is semantically that of the gauss error function.
 bool MatchERF(Operation* op) {
   if (auto custom_call = llvm::dyn_cast_or_null<stablehlo::CustomCallOp>(op)) {
-    return custom_call.getCallTargetName().equals("mhlo.erf");
+    return custom_call.getCallTargetName() == "mhlo.erf";
   }
   return llvm::isa<chlo::ErfOp>(op);
 }
diff --git a/tensorflow/compiler/mlir/lite/utils/perception_ops_utils.cc b/tensorflow/compiler/mlir/lite/utils/perception_ops_utils.cc
index 98cc1048bced90..f6595331c02415 100644
--- a/tensorflow/compiler/mlir/lite/utils/perception_ops_utils.cc
+++ b/tensorflow/compiler/mlir/lite/utils/perception_ops_utils.cc
@@ -140,8 +140,7 @@ LogicalResult ConvertMaxUnpoolingFunc::VerifySignature() {
     return func_.emitWarning() << "'padding' attribute for " << kMaxUnpooling
                                << " is not set or not a string";
   }
-  if (!padding.getValue().equals("VALID") &&
-      !padding.getValue().equals("SAME")) {
+  if (padding.getValue() != "VALID" && padding.getValue() != "SAME") {
     return func_.emitWarning()
            << "Padding for " << kMaxUnpooling << " must be 'SAME' or 'VALID'";
   }
@@ -174,9 +173,9 @@ LogicalResult ConvertMaxUnpoolingFunc::CreateCustomOptions(
     return func_.emitError() << "'padding' attribute for " << kMaxUnpooling
                              << " is not set or not a string";
   }
-  if (padding.getValue().equals("VALID")) {
+  if (padding.getValue() == "VALID") {
     pool_params.padding = kTfLitePaddingValid;
-  } else if (padding.getValue().equals("SAME")) {
+  } else if (padding.getValue() == "SAME") {
     pool_params.padding = kTfLitePaddingSame;
   } else {
     return func_.emitError()
diff --git a/tensorflow/compiler/mlir/quantization/common/lift_as_function_call.cc b/tensorflow/compiler/mlir/quantization/common/lift_as_function_call.cc
index a507b816426e8b..bf894948d4cec8 100644
--- a/tensorflow/compiler/mlir/quantization/common/lift_as_function_call.cc
+++ b/tensorflow/compiler/mlir/quantization/common/lift_as_function_call.cc
@@ -270,7 +270,7 @@ LogicalResult SetAttributeMap(MLIRContext& context,
     if (const auto string_attr =
             mlir::dyn_cast_or_null<StringAttr>(attribute.getValue());
         string_attr != nullptr &&
-        string_attr.getValue().equals(kNullAttributeValue)) {
+        string_attr.getValue() == kNullAttributeValue) {
       continue;
     }
 
diff --git a/tensorflow/compiler/mlir/quantization/stablehlo/passes/bridge/convert_tf_quant_ops_to_mhlo.cc b/tensorflow/compiler/mlir/quantization/stablehlo/passes/bridge/convert_tf_quant_ops_to_mhlo.cc
index 965da4ff998635..5575a7516fccc9 100644
--- a/tensorflow/compiler/mlir/quantization/stablehlo/passes/bridge/convert_tf_quant_ops_to_mhlo.cc
+++ b/tensorflow/compiler/mlir/quantization/stablehlo/passes/bridge/convert_tf_quant_ops_to_mhlo.cc
@@ -192,12 +192,12 @@ FailureOr<ElementsAttr> ConvertPaddingAttr(
 
   const int64_t padding_nums_size = 2 * (rhs_shape.getRank() - 2);
   padding_nums.reserve(padding_nums_size);
-  if (conv_padding.strref().equals("EXPLICIT")) {
+  if (conv_padding.strref() == "EXPLICIT") {
     for (auto padding_elem :
          op.getExplicitPaddingAttr().template getAsRange<IntegerAttr>()) {
       padding_nums.push_back(padding_elem.getInt());
     }
-  } else if (conv_padding.strref().equals("VALID")) {
+  } else if (conv_padding.strref() == "VALID") {
     padding_nums.resize(padding_nums_size, 0);
   } else {
     padding_nums.resize(padding_nums_size);
diff --git a/tensorflow/compiler/mlir/quantization/tensorflow/cc/constant_fold.cc b/tensorflow/compiler/mlir/quantization/tensorflow/cc/constant_fold.cc
index 565adebfe52300..64695d6719885d 100644
--- a/tensorflow/compiler/mlir/quantization/tensorflow/cc/constant_fold.cc
+++ b/tensorflow/compiler/mlir/quantization/tensorflow/cc/constant_fold.cc
@@ -71,7 +71,7 @@ LogicalResult FoldOperation(OpBuilder& builder, Operation* op,
 bool IsOperationFoldable(Operation* op) {
   if (isa<TF::ConstOp>(op)) return true;
 
-  if (!op->getDialect()->getNamespace().equals("tf") || !TF::CanBeFolded(op)) {
+  if (op->getDialect()->getNamespace() != "tf" || !TF::CanBeFolded(op)) {
     return false;
   }
 
diff --git a/tensorflow/compiler/mlir/quantization/tensorflow/passes/insert_custom_aggregation_ops.cc b/tensorflow/compiler/mlir/quantization/tensorflow/passes/insert_custom_aggregation_ops.cc
index 470d134d527a0f..c5d7ca8e47f6f9 100644
--- a/tensorflow/compiler/mlir/quantization/tensorflow/passes/insert_custom_aggregation_ops.cc
+++ b/tensorflow/compiler/mlir/quantization/tensorflow/passes/insert_custom_aggregation_ops.cc
@@ -69,9 +69,9 @@ bool IsCallToQuantizableLiftedFunction(Operation *op) {
 
   TF::PartitionedCallOp call_op = dyn_cast_or_null<TF::PartitionedCallOp>(op);
   return call_op && call_op->hasAttrOfType<StringAttr>(kQuantTraitAttrName) &&
-         call_op->getAttrOfType<StringAttr>(kQuantTraitAttrName)
-             .getValue()
-             .equals(QuantTraitValues[QuantizationTrait::FullyQuantizable]);
+         call_op->getAttrOfType<StringAttr>(kQuantTraitAttrName).getValue() ==
+             llvm::StringRef(
+                 QuantTraitValues[QuantizationTrait::FullyQuantizable]);
 }
 
 // Returns the composite function name.
diff --git a/tensorflow/compiler/mlir/quantization/tensorflow/passes/lift_hashtable_ops_as_args.cc b/tensorflow/compiler/mlir/quantization/tensorflow/passes/lift_hashtable_ops_as_args.cc
index f48d15dd81cbdb..18ee96bfe9422e 100644
--- a/tensorflow/compiler/mlir/quantization/tensorflow/passes/lift_hashtable_ops_as_args.cc
+++ b/tensorflow/compiler/mlir/quantization/tensorflow/passes/lift_hashtable_ops_as_args.cc
@@ -70,14 +70,15 @@ bool IsHashTableOp(Operation* op) {
 
 // Checks if the function is the main or initializer function.
 bool IsMainOrInitializerFunction(ModuleOp module, func::FuncOp func) {
-  if (func.getSymName().equals(tensorflow::kImportModelDefaultGraphFuncName) ||
-      func.getSymName().equals(kTfQuantSaveFuncName)) {
+  if (func.getSymName() ==
+          llvm::StringRef(tensorflow::kImportModelDefaultGraphFuncName) ||
+      func.getSymName() == kTfQuantSaveFuncName) {
     return true;
   }
 
   for (func::FuncOp init_func :
        tf_saved_model::GetInitializerFunctions(module)) {
-    if (func.getSymName().equals(init_func.getSymName())) {
+    if (func.getSymName() == init_func.getSymName()) {
       return true;
     }
   }
@@ -118,7 +119,7 @@ bool IsResourceInitialized(ModuleOp module_op, Operation* hash_table) {
        tf_saved_model::GetInitializerFunctions(module_op)) {
     for (Operation& op : init_func_op.getBody().getOps()) {
       StringRef other_shared_name = GetSharedName(&op);
-      if (IsHashTableOp(&op) && other_shared_name.equals(shared_name)) {
+      if (IsHashTableOp(&op) && other_shared_name == shared_name) {
         return true;
       }
     }
diff --git a/tensorflow/compiler/mlir/quantization/tensorflow/passes/prepare_lifting.cc b/tensorflow/compiler/mlir/quantization/tensorflow/passes/prepare_lifting.cc
index b0a84d71c84182..a87245345f6987 100644
--- a/tensorflow/compiler/mlir/quantization/tensorflow/passes/prepare_lifting.cc
+++ b/tensorflow/compiler/mlir/quantization/tensorflow/passes/prepare_lifting.cc
@@ -157,8 +157,8 @@ LogicalResult MatchSupportedAffineOp(Operation* op, Value& binding_output,
   bool is_supported_affine_op = false;
   if (llvm::isa<TF::Conv2DOp, TF::Conv3DOp, TF::DepthwiseConv2dNativeOp>(op)) {
     if (const auto data_format = op->getAttrOfType<StringAttr>("data_format")) {
-      is_supported_affine_op = data_format.getValue().equals("NHWC") ||
-                               data_format.getValue().equals("NDHWC");
+      is_supported_affine_op =
+          data_format.getValue() == "NHWC" || data_format.getValue() == "NDHWC";
     }
   } else if (llvm::isa<TF::BatchMatMulV2Op>(op)) {
     if (const auto adj_y = op->getAttrOfType<BoolAttr>("adj_y")) {
diff --git a/tensorflow/compiler/mlir/quantization/tensorflow/utils/tf_to_xla_attribute_utils.cc b/tensorflow/compiler/mlir/quantization/tensorflow/utils/tf_to_xla_attribute_utils.cc
index 7054f049d1369c..b22726de30aeaa 100644
--- a/tensorflow/compiler/mlir/quantization/tensorflow/utils/tf_to_xla_attribute_utils.cc
+++ b/tensorflow/compiler/mlir/quantization/tensorflow/utils/tf_to_xla_attribute_utils.cc
@@ -160,7 +160,7 @@ Value CalculatePaddingAndPadIfNeeded(OpBuilder &builder, Location loc,
   bool has_dynamic_spatial_dim = absl::c_any_of(
       spatial_dims,
       [&input_shape](int64_t dim) { return input_shape.isDynamicDim(dim); });
-  if (conv_padding.strref().equals("SAME") && has_dynamic_spatial_dim) {
+  if (conv_padding.strref() == "SAME" && has_dynamic_spatial_dim) {
     return PadForDynamicShapedInputSamePadding(
         builder, loc, input, filter, input_zp_value, strides, dilations,
         conv_padding, padding, num_dims);
@@ -168,7 +168,7 @@ Value CalculatePaddingAndPadIfNeeded(OpBuilder &builder, Location loc,
 
   ShapedType filter_shape = mlir::cast<ShapedType>(filter.getType());
   SmallVector<int32_t> padding_values(2 * num_dims, 0);
-  if (conv_padding.strref().equals("EXPLICIT")) {
+  if (conv_padding.strref() == "EXPLICIT") {
     if (explicit_paddings.size() != 2 * num_dims) {
       emitError(loc,
                 absl::StrFormat(
@@ -182,7 +182,7 @@ Value CalculatePaddingAndPadIfNeeded(OpBuilder &builder, Location loc,
       padding_values[2 * i + 1] =
           mlir::cast<IntegerAttr>(explicit_paddings[2 * i + 1]).getInt();
     }
-  } else if (conv_padding.strref().equals("SAME")) {
+  } else if (conv_padding.strref() == "SAME") {
     for (int i : spatial_dims) {
       int input_size = input_shape.getDimSize(i);
       int filter_size = filter_shape.getDimSize(i - 1);
diff --git a/tensorflow/compiler/mlir/tensorflow/transforms/executor_tpuv1_island_coarsening.cc b/tensorflow/compiler/mlir/tensorflow/transforms/executor_tpuv1_island_coarsening.cc
index 2f456248c381af..b75f081d1a0064 100644
--- a/tensorflow/compiler/mlir/tensorflow/transforms/executor_tpuv1_island_coarsening.cc
+++ b/tensorflow/compiler/mlir/tensorflow/transforms/executor_tpuv1_island_coarsening.cc
@@ -395,7 +395,7 @@ bool is_valid_special_tpu_op(
 
     bool op_has_inconsistent_cluster_name =
         wrapped_op_cluster_name.has_value() &&
-        !wrapped_op_cluster_name.value().equals(cluster_name);
+        wrapped_op_cluster_name.value() != cluster_name;
 
     if (op_has_inconsistent_cluster_name) {
       return false;
diff --git a/tensorflow/compiler/mlir/tensorflow/transforms/tfg-to-tfe.cc b/tensorflow/compiler/mlir/tensorflow/transforms/tfg-to-tfe.cc
index 8180b4116ef21b..68d50e54a1bce0 100644
--- a/tensorflow/compiler/mlir/tensorflow/transforms/tfg-to-tfe.cc
+++ b/tensorflow/compiler/mlir/tensorflow/transforms/tfg-to-tfe.cc
@@ -115,7 +115,7 @@ static void FilterOutBlockArgControlDep(
 static void SplitNextIteration(Block &block) {
   // TODO(b/207144333): Supports callback for unregistered ops
   block.walk([&](Operation *op) {
-    if (!op->getName().getStringRef().equals("tfg.NextIteration")) return;
+    if (op->getName().getStringRef() != "tfg.NextIteration") return;
     mlir::OpBuilder builder(op);
 
     llvm::SmallVector<Value, 2> new_operands;
diff --git a/tensorflow/compiler/mlir/tensorflow/utils/export_utils.cc b/tensorflow/compiler/mlir/tensorflow/utils/export_utils.cc
index 7bb35d158f00cd..96ba0afd096a16 100644
--- a/tensorflow/compiler/mlir/tensorflow/utils/export_utils.cc
+++ b/tensorflow/compiler/mlir/tensorflow/utils/export_utils.cc
@@ -282,12 +282,12 @@ static bool IsRefTypeControlOp(mlir::Operation* op) {
   if (!op_name_or_status.ok()) return false;
 
   auto op_name = std::move(op_name_or_status).value();
-  if (op_name.equals("NextIteration"))
+  if (op_name == "NextIteration")
     return mlir::isa<mlir::TF::TensorFlowRefType>(
         mlir::getElementTypeOrSelf(op->getOperand(0).getType()));
 
-  if (op_name.equals("Enter") || op_name.equals("Exit") ||
-      op_name.equals("Switch") || op_name.equals("Merge")) {
+  if (op_name == "Enter" || op_name == "Exit" || op_name == "Switch" ||
+      op_name == "Merge") {
     return mlir::isa<mlir::TF::TensorFlowRefType>(
         getElementTypeOrSelf(op->getResult(0).getType()));
   }

From e116891194c9e89b962d4a7e436d20c190600c35 Mon Sep 17 00:00:00 2001
From: Hyeontaek Lim <hyeontaek@google.com>
Date: Fri, 17 May 2024 10:25:43 -0700
Subject: [PATCH 361/478] [IFRT] Allow kReuseInput in RemapArrays if the number
 of input arrays is 1

`xla::ifrt::Client::RemapArrays()` allows `kReuseInput` for projection (view).
Previously, it expected the number of input arrays and output arrays is both 1.
This restriction has been relaxed to expecting only the number of input arrays
to be 1 because each of output arrays can be considered a projection of the
input arrays. In theory, the user could run N `RemapArrays()` for 1:1 remapping
to get 1:N remapping anyway to bypass the previous restriction, making the
restriction not effective.

PiperOrigin-RevId: 634818012
---
 third_party/xla/xla/python/ifrt/client.h      | 11 +++-
 .../xla/python/ifrt/remap_impl_test_lib.cc    | 62 ++++++++++++-------
 .../xla/xla/python/pjrt_ifrt/pjrt_remap.cc    |  4 +-
 3 files changed, 49 insertions(+), 28 deletions(-)

diff --git a/third_party/xla/xla/python/ifrt/client.h b/third_party/xla/xla/python/ifrt/client.h
index f6481e803f0a27..ef42e3b3b42f88 100644
--- a/third_party/xla/xla/python/ifrt/client.h
+++ b/third_party/xla/xla/python/ifrt/client.h
@@ -120,8 +120,15 @@ class Client : public llvm::RTTIExtends<Client, llvm::RTTIRoot> {
   // Remaps shards across input `Array`s to create new `Array`s based on `plan`.
   // This array remapping is a metadata-only operation that can shuffle or
   // extract shards without changing their per-shard interpretation and causing
-  // data copy/transfer. Using `ArrayCopySemantics::kAlwaysCopy` has an
-  // undefined behavior.
+  // data copy/transfer.
+  //
+  // There are constraints on `semantics`:
+  //
+  // * `ArrayCopySemantics::kAlwaysCopy` has an undefined behavior because
+  // `RemapArrays` does not copy data.
+  // * `ArrayCopySemantics::kReuseInput` is allowed only if the number of inputs
+  // is 1. This is safe because each input shard can be used only once.
+  // * `ArrayCopySemantics::kDonateInput` is always allowed.
   virtual absl::StatusOr<std::vector<tsl::RCReference<xla::ifrt::Array>>>
   RemapArrays(const RemapPlan& plan,
               absl::Span<tsl::RCReference<xla::ifrt::Array>> arrays,
diff --git a/third_party/xla/xla/python/ifrt/remap_impl_test_lib.cc b/third_party/xla/xla/python/ifrt/remap_impl_test_lib.cc
index 9486e0409412b7..a55d97d13998e4 100644
--- a/third_party/xla/xla/python/ifrt/remap_impl_test_lib.cc
+++ b/third_party/xla/xla/python/ifrt/remap_impl_test_lib.cc
@@ -243,11 +243,12 @@ TEST(RemapImplTest, InterleaveArrays) {
                           CreateArray(client.get(), /*base_values=*/{100, 106},
                                       /*device_indices=*/{2, 3}));
 
-  EXPECT_THAT(client->RemapArrays(plan, absl::MakeSpan(arrays),
-                                  ArrayCopySemantics::kReuseInput),
-              StatusIs(absl::StatusCode::kInvalidArgument,
-                       HasSubstr("kDonateInput is required if multiple inputs "
-                                 "or outputs are used")));
+  EXPECT_THAT(
+      client->RemapArrays(plan, absl::MakeSpan(arrays),
+                          ArrayCopySemantics::kReuseInput),
+      StatusIs(
+          absl::StatusCode::kInvalidArgument,
+          HasSubstr("kDonateInput is required if multiple inputs are used")));
 
   TF_ASSERT_OK_AND_ASSIGN(
       auto out_arrays, client->RemapArrays(plan, absl::MakeSpan(arrays),
@@ -311,27 +312,40 @@ TEST(RemapImplTest, DeinterleaveArrays) {
       CreateArray(client.get(), /*base_values=*/{0, 100, 6, 106},
                   /*device_indices=*/{0, 2, 1, 3}));
 
-  EXPECT_THAT(client->RemapArrays(plan, absl::MakeSpan(arrays),
-                                  ArrayCopySemantics::kReuseInput),
-              StatusIs(absl::StatusCode::kInvalidArgument,
-                       HasSubstr("kDonateInput is required if multiple inputs "
-                                 "or outputs are used")));
+  {
+    TF_ASSERT_OK_AND_ASSIGN(
+        auto out_arrays, client->RemapArrays(plan, absl::MakeSpan(arrays),
+                                             ArrayCopySemantics::kReuseInput));
 
-  TF_ASSERT_OK_AND_ASSIGN(
-      auto out_arrays, client->RemapArrays(plan, absl::MakeSpan(arrays),
-                                           ArrayCopySemantics::kDonateInput));
+    ASSERT_THAT(out_arrays, SizeIs(2));
+    // `out_arrays[0].shards[0] == arrays[0].shards[0]`
+    // `out_arrays[0].shards[1] == arrays[0].shards[2]`
+    // `out_arrays[1].shards[0] == arrays[0].shards[1]`
+    // `out_arrays[1].shards[1] == arrays[0].shards[3]`
+    AssertArrayContent(client.get(), out_arrays[0].get(),
+                       /*base_values=*/{0, 6},
+                       /*device_indices=*/{0, 1});
+    AssertArrayContent(client.get(), out_arrays[1].get(),
+                       /*base_values=*/{100, 106},
+                       /*device_indices=*/{2, 3});
+  }
+  {
+    TF_ASSERT_OK_AND_ASSIGN(
+        auto out_arrays, client->RemapArrays(plan, absl::MakeSpan(arrays),
+                                             ArrayCopySemantics::kDonateInput));
 
-  ASSERT_THAT(out_arrays, SizeIs(2));
-  // `out_arrays[0].shards[0] == arrays[0].shards[0]`
-  // `out_arrays[0].shards[1] == arrays[0].shards[2]`
-  // `out_arrays[1].shards[0] == arrays[0].shards[1]`
-  // `out_arrays[1].shards[1] == arrays[0].shards[3]`
-  AssertArrayContent(client.get(), out_arrays[0].get(),
-                     /*base_values=*/{0, 6},
-                     /*device_indices=*/{0, 1});
-  AssertArrayContent(client.get(), out_arrays[1].get(),
-                     /*base_values=*/{100, 106},
-                     /*device_indices=*/{2, 3});
+    ASSERT_THAT(out_arrays, SizeIs(2));
+    // `out_arrays[0].shards[0] == arrays[0].shards[0]`
+    // `out_arrays[0].shards[1] == arrays[0].shards[2]`
+    // `out_arrays[1].shards[0] == arrays[0].shards[1]`
+    // `out_arrays[1].shards[1] == arrays[0].shards[3]`
+    AssertArrayContent(client.get(), out_arrays[0].get(),
+                       /*base_values=*/{0, 6},
+                       /*device_indices=*/{0, 1});
+    AssertArrayContent(client.get(), out_arrays[1].get(),
+                       /*base_values=*/{100, 106},
+                       /*device_indices=*/{2, 3});
+  }
 }
 
 }  // namespace
diff --git a/third_party/xla/xla/python/pjrt_ifrt/pjrt_remap.cc b/third_party/xla/xla/python/pjrt_ifrt/pjrt_remap.cc
index f72a1761b6e75b..87e2046a268347 100644
--- a/third_party/xla/xla/python/pjrt_ifrt/pjrt_remap.cc
+++ b/third_party/xla/xla/python/pjrt_ifrt/pjrt_remap.cc
@@ -52,10 +52,10 @@ PjRtCompatibleClientRemapArrays(
           arrays[i]->DebugString());
     }
   }
-  if (plan.input_specs.size() > 1 || plan.output_specs.size() > 1) {
+  if (plan.input_specs.size() > 1) {
     if (semantics != ArrayCopySemantics::kDonateInput) {
       return InvalidArgument(
-          "kDonateInput is required if multiple inputs or outputs are used");
+          "kDonateInput is required if multiple inputs are used");
     }
   }
 

From b960aaf3c647e8b118c49c959a1eda90386b8825 Mon Sep 17 00:00:00 2001
From: Harsha H S <hsharsha@users.noreply.github.com>
Date: Fri, 17 May 2024 10:31:41 -0700
Subject: [PATCH 362/478] PR #12637: [ROCm] Fix build break due to 63c33b and
 f5ab2a4

Imported from GitHub PR https://github.com/openxla/xla/pull/12637

Copybara import of the project:

--
9a090b699df98d7c2a27311fbf353d4a0390e87f by Harsha HS <harsha.havanurshamsundara@amd.com>:

[ROCm] Fix build break due to 63c33b and f5ab2a4

Merging this change closes #12637

PiperOrigin-RevId: 634819780
---
 third_party/xla/xla/service/gpu/BUILD         | 2 +-
 third_party/xla/xla/service/gpu/runtime/BUILD | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/third_party/xla/xla/service/gpu/BUILD b/third_party/xla/xla/service/gpu/BUILD
index 275499abcb330e..92d6d3d5596b2b 100644
--- a/third_party/xla/xla/service/gpu/BUILD
+++ b/third_party/xla/xla/service/gpu/BUILD
@@ -1743,10 +1743,10 @@ cc_library(
         "@local_tsl//tsl/platform:tensor_float_32_hdr_lib",
     ]) + if_rocm_is_configured([
         #keep sorted
-        "//third_party/gpus/rocm:rocm_headers",
         "//xla/stream_executor/platform:dso_loader",
         "//xla/stream_executor/rocm:amdhipblaslt_plugin",
         "//xla/stream_executor/rocm:hipblas_lt_header",
+        "@local_config_rocm//rocm:rocm_headers",
     ]) + if_static([
         "@local_tsl//tsl/platform:tensor_float_32_utils",
     ]),
diff --git a/third_party/xla/xla/service/gpu/runtime/BUILD b/third_party/xla/xla/service/gpu/runtime/BUILD
index ef8a8134628a05..ae2c70bfb0981a 100644
--- a/third_party/xla/xla/service/gpu/runtime/BUILD
+++ b/third_party/xla/xla/service/gpu/runtime/BUILD
@@ -217,7 +217,7 @@ cc_library(
         "//xla/stream_executor/cuda:cuda_driver",
         "//xla/stream_executor/cuda:cuda_executor",
     ]) + if_rocm_is_configured([
-        "//third_party/gpus/rocm:rocm_headers",
+        "@local_config_rocm//rocm:rocm_headers",
         "@local_config_rocm//rocm:rccl",
         "//xla/stream_executor/rocm:rocm_driver",
         "//xla/stream_executor/rocm:rocm_executor",

From d4dc6cd8faceff89bab579a31ebd2bf1cbd9d2d1 Mon Sep 17 00:00:00 2001
From: Greg Olechwierowicz <olechwierowicz@google.com>
Date: Fri, 17 May 2024 10:49:46 -0700
Subject: [PATCH 363/478] [XLA:GPU][NFC] Print reachability between two
 instructions on VLOG.

PiperOrigin-RevId: 634825414
---
 third_party/xla/xla/service/BUILD             |  5 ++++-
 .../xla/service/collective_combiner_utils.h   | 21 +++++++++++++------
 2 files changed, 19 insertions(+), 7 deletions(-)

diff --git a/third_party/xla/xla/service/BUILD b/third_party/xla/xla/service/BUILD
index 39336da669e985..51e4a64911d110 100644
--- a/third_party/xla/xla/service/BUILD
+++ b/third_party/xla/xla/service/BUILD
@@ -2650,7 +2650,6 @@ cc_library(
     name = "collective_combiner_utils",
     hdrs = ["collective_combiner_utils.h"],
     deps = [
-        ":hlo_domain_map",
         "//xla:shape_util",
         "//xla:status_macros",
         "//xla:xla_data_proto_cc",
@@ -2659,6 +2658,10 @@ cc_library(
         "@com_google_absl//absl/container:flat_hash_map",
         "@com_google_absl//absl/container:flat_hash_set",
         "@com_google_absl//absl/functional:function_ref",
+        "@com_google_absl//absl/log",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/status:statusor",
+        "@com_google_absl//absl/types:span",
         "@local_tsl//tsl/platform:errors",
     ],
 )
diff --git a/third_party/xla/xla/service/collective_combiner_utils.h b/third_party/xla/xla/service/collective_combiner_utils.h
index e7b2ecc573670a..29ef053dc6d60b 100644
--- a/third_party/xla/xla/service/collective_combiner_utils.h
+++ b/third_party/xla/xla/service/collective_combiner_utils.h
@@ -16,18 +16,20 @@ limitations under the License.
 #ifndef XLA_SERVICE_COLLECTIVE_COMBINER_UTILS_H_
 #define XLA_SERVICE_COLLECTIVE_COMBINER_UTILS_H_
 
-#include <functional>
-#include <utility>
+#include <cstdint>
+#include <memory>
+#include <optional>
 #include <vector>
 
 #include "absl/container/flat_hash_map.h"
 #include "absl/container/flat_hash_set.h"
 #include "absl/functional/function_ref.h"
+#include "absl/log/log.h"
+#include "absl/status/status.h"
+#include "absl/status/statusor.h"
+#include "absl/types/span.h"
 #include "xla/hlo/ir/hlo_instruction.h"
-#include "xla/hlo/ir/hlo_instructions.h"
-#include "xla/hlo/ir/hlo_opcode.h"
 #include "xla/hlo/ir/hlo_reachability.h"
-#include "xla/service/hlo_domain_map.h"
 #include "xla/shape_util.h"
 #include "xla/status_macros.h"
 #include "xla/xla_data.pb.h"
@@ -123,7 +125,14 @@ absl::StatusOr<bool> CombineInstructionsByKey(
       // We can't combine dependent instructions.
       bool is_reachable =
           absl::c_any_of(to_combine, [&](HloInstruction* to_combine_inst) {
-            return reachability->IsReachable(to_combine_inst, instruction);
+            bool reachable =
+                reachability->IsReachable(to_combine_inst, instruction);
+            if (reachable) {
+              VLOG(2) << "<< Instruction {" << instruction->ToShortString()
+                      << "} is reachable from {"
+                      << to_combine_inst->ToShortString() << "}";
+            }
+            return reachable;
           });
       if (is_reachable) {
         VLOG(1) << "Instruction is reachable.";

From dcb77c5847dc17fc545c84627c77e51bb5455512 Mon Sep 17 00:00:00 2001
From: Derek Murray <mrry@google.com>
Date: Fri, 17 May 2024 11:45:53 -0700
Subject: [PATCH 364/478] [Function inlining] Add flag to enable pruning before
 function calls are inlined.

Currently, the behavior of lowered functional ("V2") control flow is subtly different from V1 control flow. In V1 control flow, the fetches for a `Session::Run()` call will determine whether or not  operations within the loop will be pruned. However, in V2 control flow, pruning may not always occur for multi-device functions, because the graph is partitioned (inserting stateful Send/Recv nodes) into multiple sub-functions, before the individual sub-functions are pruned.
While this should not affect the correctness of the execution, it may cause unnecessary computations and data transfers.

To address this problem, we introduce a new flag to enable pruning when inlining function calls, including as part of control flow V2 lowering, which inlines the condition and body functions of while loops, and true and false branches of conditionals.

Set `TF_FLAG_ENABLE_FUNCTION_PRUNING_BEFORE_INLINING=1` to enable this option.

PiperOrigin-RevId: 634842835
---
 tensorflow/core/common_runtime/BUILD          | 10 ++-
 tensorflow/core/common_runtime/function.cc    | 44 -----------
 .../core/common_runtime/function_def_utils.cc | 49 ++++++++++++
 .../core/common_runtime/function_def_utils.h  |  9 +++
 .../common_runtime/lower_function_call_op.cc  | 16 ++++
 .../lower_function_call_op_test.cc            | 78 +++++++++++++++++++
 .../common_runtime/lower_while_op_test.cc     | 66 ++++++++++++++++
 tensorflow/core/config/flag_defs.h            |  3 +
 tensorflow/core/config/flags_api_wrapper.cc   |  1 +
 tensorflow/core/framework/function_testlib.cc | 23 ++++++
 tensorflow/core/framework/function_testlib.h  |  2 +
 tensorflow/python/flags_pybind.pyi            |  1 +
 12 files changed, 256 insertions(+), 46 deletions(-)

diff --git a/tensorflow/core/common_runtime/BUILD b/tensorflow/core/common_runtime/BUILD
index fd46e483f39906..62617578abb5b2 100644
--- a/tensorflow/core/common_runtime/BUILD
+++ b/tensorflow/core/common_runtime/BUILD
@@ -936,8 +936,9 @@ cc_library(
         ":function_body",
         ":graph_constructor",
         "//tensorflow/core:framework",
-        "//tensorflow/core:graph",
-        "//tensorflow/core:protos_all_cc",
+        "//tensorflow/core/lib/core:status",
+        "//tensorflow/core/platform:refcount",
+        "@local_tsl//tsl/platform:errors",
     ],
 )
 
@@ -1267,8 +1268,10 @@ cc_library(
         "//tensorflow/core:framework",
         "//tensorflow/core:graph",
         "//tensorflow/core:lib",
+        "//tensorflow/core/config:flag_defs",
         "//tensorflow/core/platform:refcount",
         "@com_google_absl//absl/algorithm:container",
+        "@com_google_absl//absl/types:span",
     ],
 )
 
@@ -3128,6 +3131,7 @@ tf_cc_test(
         "//tensorflow/core:test",
         "//tensorflow/core:test_main",
         "//tensorflow/core:testlib",
+        "//tensorflow/core/config:flag_defs",
     ],
 )
 
@@ -3193,6 +3197,7 @@ tf_cc_test(
         "//tensorflow/cc:client_session",
         "//tensorflow/cc:function_ops",
         "//tensorflow/cc:ops",
+        "//tensorflow/cc:scope",
         "//tensorflow/core:all_kernels",
         "//tensorflow/core:framework",
         "//tensorflow/core:framework_internal",
@@ -3201,6 +3206,7 @@ tf_cc_test(
         "//tensorflow/core:test",
         "//tensorflow/core:test_main",
         "//tensorflow/core:testlib",
+        "//tensorflow/core/config:flag_defs",
         "@com_google_absl//absl/strings",
         "@com_google_googletest//:gtest",
     ],
diff --git a/tensorflow/core/common_runtime/function.cc b/tensorflow/core/common_runtime/function.cc
index d4891e6f906138..8a55581ed19262 100644
--- a/tensorflow/core/common_runtime/function.cc
+++ b/tensorflow/core/common_runtime/function.cc
@@ -897,50 +897,6 @@ Status FunctionLibraryRuntimeImpl::ReleaseHandle(Handle handle) {
   return parent_status;
 }
 
-namespace {
-
-// Removes all stateless nodes that do not contribute to a return
-// value from the function body. Unlike `RemoveDeadNodes()`, which is
-// triggered by `OptimizerOptions.do_function_inlining`, this pass
-// ignores the SINK node, from which (by definition) all nodes are
-// reverse reachable, and preserves all nodes that are reachable from
-// control output nodes.
-//
-// TODO(ezhulenev, skyewm): Function body should not have special treatment of
-// stateful ops, graph should encode nodes that must execute with `control_ret`
-// and `control_output`.
-void PruneFunctionBody(const FunctionDef& fdef, Graph* g) {
-  VLOG(2) << "Pruning function body: function_name=" << fdef.signature().name();
-
-  // `control_ret` nodes must be always executed.
-  std::unordered_set<StringPiece, StringPieceHasher> control_ret_nodes;
-  for (const auto& control_ret : fdef.control_ret()) {
-    control_ret_nodes.insert(control_ret.second);
-  }
-
-  std::unordered_set<const Node*> nodes;
-  for (auto n : g->nodes()) {
-    // NOTE(mrry): "_Retval" nodes are stateful, and so will be added
-    // to the seed set of `nodes`. "_Arg" nodes are also stateful, but we
-    // specifically exclude them as seeds, to avoid unconditionally executing
-    // unused argument nodes (e.g. in a function like `lambda x, y: y`).
-    // TODO(mrry): Investigate whether the `n->IsControlFlow()` test is
-    // still needed. It would be preferable to prune entire loops and/or
-    // conditionals if they are not used in the graph.
-    if (n->IsControlFlow() ||
-        (n->op_def().is_stateful() && n->type_string() != kArgOp) ||
-        (control_ret_nodes.find(n->name()) != control_ret_nodes.end())) {
-      nodes.insert(n);
-    }
-  }
-  bool changed = PruneForReverseReachability(g, std::move(nodes));
-  if (changed) {
-    FixupSourceAndSinkEdges(g);
-  }
-}
-
-}  // namespace
-
 Status FunctionLibraryRuntimeImpl::CreateItem(Item** item) {
   const FunctionBody* fbody;
   FunctionLibraryRuntime* flr;
diff --git a/tensorflow/core/common_runtime/function_def_utils.cc b/tensorflow/core/common_runtime/function_def_utils.cc
index c4d87114291ac5..ee7d1bb867c3a8 100644
--- a/tensorflow/core/common_runtime/function_def_utils.cc
+++ b/tensorflow/core/common_runtime/function_def_utils.cc
@@ -15,18 +15,25 @@ limitations under the License.
 
 #include "tensorflow/core/common_runtime/function_def_utils.h"
 
+#include <unordered_set>
 #include <utility>
 #include <vector>
 
+#include "absl/container/flat_hash_set.h"
+#include "absl/strings/string_view.h"
+#include "absl/types/span.h"
 #include "tensorflow/core/common_runtime/function_body.h"
 #include "tensorflow/core/common_runtime/graph_constructor.h"
 #include "tensorflow/core/framework/function.h"
 #include "tensorflow/core/framework/node_def_util.h"
+#include "tensorflow/core/graph/algorithm.h"
 #include "tensorflow/core/graph/control_flow.h"
 #include "tensorflow/core/graph/graph.h"
 #include "tensorflow/core/graph/graph_debug_info_builder.h"
 #include "tensorflow/core/platform/refcount.h"
 #include "tsl/platform/errors.h"
+#include "tsl/platform/hash.h"
+#include "tsl/platform/logging.h"
 
 namespace tensorflow {
 
@@ -106,4 +113,46 @@ Status FunctionDefToBodyHelper(const FunctionDef& fdef, const AttrSlice& attrs,
                                  get_func_sig, fbody);
 }
 
+// TODO(ezhulenev, skyewm): Function body should not have special treatment of
+// stateful ops, graph should encode nodes that must execute with `control_ret`
+// and `control_output`.
+void PruneFunctionBody(const FunctionDef& fdef, Graph* g,
+                       absl::Span<Node*> additional_root_nodes) {
+  VLOG(2) << "Pruning function body: function_name=" << fdef.signature().name()
+          << " #nodes = " << g->num_nodes();
+
+  // `control_ret` nodes must be always executed.
+  absl::flat_hash_set<absl::string_view, tsl::StringPieceHasher>
+      control_ret_nodes;
+  for (const auto& control_ret : fdef.control_ret()) {
+    control_ret_nodes.insert(control_ret.second);
+  }
+
+  std::unordered_set<const Node*> nodes;
+  for (auto n : additional_root_nodes) {
+    nodes.insert(n);
+  }
+  for (auto n : g->nodes()) {
+    // NOTE(mrry): "_Retval" nodes are stateful, and so will be added
+    // to the seed set of `nodes`. "_Arg" nodes are also stateful, but we
+    // specifically exclude them as seeds, to avoid unconditionally executing
+    // unused argument nodes (e.g. in a function like `lambda x, y: y`).
+    // TODO(mrry): Investigate whether the `n->IsControlFlow()` test is
+    // still needed. It would be preferable to prune entire loops and/or
+    // conditionals if they are not used in the graph.
+    if (n->IsControlFlow() ||
+        (n->op_def().is_stateful() &&
+         n->type_string() != FunctionLibraryDefinition::kArgOp) ||
+        (control_ret_nodes.find(n->name()) != control_ret_nodes.end())) {
+      nodes.insert(n);
+    }
+  }
+  bool changed = PruneForReverseReachability(g, std::move(nodes));
+  if (changed) {
+    VLOG(2) << "Pruned function body and changed: function_name="
+            << fdef.signature().name() << " #nodes = " << g->num_nodes();
+    FixupSourceAndSinkEdges(g);
+  }
+}
+
 }  // end namespace tensorflow
diff --git a/tensorflow/core/common_runtime/function_def_utils.h b/tensorflow/core/common_runtime/function_def_utils.h
index 1d60ce3b38c43c..b5f92660dc7f22 100644
--- a/tensorflow/core/common_runtime/function_def_utils.h
+++ b/tensorflow/core/common_runtime/function_def_utils.h
@@ -57,6 +57,15 @@ Status FunctionDefToBodyHelper(
     const std::function<Status(const string&, const OpDef**)>& get_func_sig,
     std::unique_ptr<FunctionBody>* fbody);
 
+// Removes all stateless nodes that do not contribute to a return
+// value from the function body. Unlike `RemoveDeadNodes()`, which is
+// triggered by `OptimizerOptions.do_function_inlining`, this pass
+// ignores the SINK node, from which (by definition) all nodes are
+// reverse reachable, and preserves all nodes that are reachable from
+// control output nodes.
+void PruneFunctionBody(const FunctionDef& fdef, Graph* g,
+                       absl::Span<Node*> additional_root_nodes = {});
+
 }  // end namespace tensorflow
 
 #endif  // TENSORFLOW_CORE_COMMON_RUNTIME_FUNCTION_DEF_UTILS_H_
diff --git a/tensorflow/core/common_runtime/lower_function_call_op.cc b/tensorflow/core/common_runtime/lower_function_call_op.cc
index b05509e5246e9e..c5226b4eefcc85 100644
--- a/tensorflow/core/common_runtime/lower_function_call_op.cc
+++ b/tensorflow/core/common_runtime/lower_function_call_op.cc
@@ -18,9 +18,11 @@ limitations under the License.
 #include <utility>
 
 #include "absl/algorithm/container.h"
+#include "absl/types/span.h"
 #include "tensorflow/core/common_runtime/function_def_utils.h"
 #include "tensorflow/core/common_runtime/inline_function_utils.h"
 #include "tensorflow/core/common_runtime/lower_function_call_inline_policy.h"
+#include "tensorflow/core/config/flag_defs.h"
 #include "tensorflow/core/framework/node_def_util.h"
 #include "tensorflow/core/graph/graph.h"
 #include "tensorflow/core/graph/graph_node_util.h"
@@ -87,6 +89,20 @@ Status RewriteFunctionCallNode(Node* n, Graph* g,
   TF_RETURN_IF_ERROR(
       FunctionDefToBodyHelper(std::move(fdef), n->attrs(), &flib_def, &fbody));
 
+  if (flags::Global().enable_function_pruning_before_inlining.value()) {
+    // TODO(b/341325107): Enable this path by default and remove the flag.
+    VLOG(2) << "Pruning enabled before inlining";
+    // NOTE(mrry): We pass `fbody->arg_nodes` as an additional set of roots,
+    // because otherwise the `FunctionBody` state will become inconsistent.
+    // The unused `Identity` nodes will be colocated with the arguments, and
+    // pruned in a subsequent pass.
+    PruneFunctionBody(
+        fbody->record->fdef(), fbody->graph,
+        absl::Span<Node*>(fbody->arg_nodes.data(), fbody->arg_nodes.size()));
+  } else {
+    VLOG(2) << "Pruning disabled before inlining";
+  }
+
   Status can_inline_function_call =
       ValidateInlining(n, fbody.get(), inline_options);
   if (can_inline_function_call.ok()) {
diff --git a/tensorflow/core/common_runtime/lower_function_call_op_test.cc b/tensorflow/core/common_runtime/lower_function_call_op_test.cc
index 600c422c8f546c..cc131835a38811 100644
--- a/tensorflow/core/common_runtime/lower_function_call_op_test.cc
+++ b/tensorflow/core/common_runtime/lower_function_call_op_test.cc
@@ -23,6 +23,7 @@ limitations under the License.
 #include "tensorflow/core/common_runtime/graph_constructor.h"
 #include "tensorflow/core/common_runtime/graph_runner.h"
 #include "tensorflow/core/common_runtime/lower_functional_ops.h"
+#include "tensorflow/core/config/flag_defs.h"
 #include "tensorflow/core/framework/function_testlib.h"
 #include "tensorflow/core/framework/node_def_util.h"
 #include "tensorflow/core/framework/op.h"
@@ -133,6 +134,83 @@ TEST(LowerFunctionCallTest, InlineFunctionCall) {
   }
 }
 
+TEST(LowerFunctionCallTest, InlineFunctionCallAfterPruning) {
+  flags::Global().enable_function_pruning_before_inlining.reset(true);
+  using FDH = FunctionDefHelper;
+
+  std::unique_ptr<Graph> graph(new Graph(OpRegistry::Global()));
+
+  FunctionDefLibrary f_lib_proto;
+
+  // `add` node is not required to compute regular output `o`, but it must
+  // execute because it is in `control_ret`.
+  // The `div` node and the unused arguments `j` and `k` should be pruned.
+  *(f_lib_proto.add_function()) = FDH::Create(
+      "AddAndMul", {"i: int32", "j: int32", "k: int32"}, {"o: int32"}, {},
+      {{{"add"}, "Add", {"i", "i"}, {{"T", DT_INT32}}},
+       {{"div"}, "FloorDiv", {"i", "i"}, {{"T", DT_INT32}}},
+       {{"ret"}, "Mul", {"i", "i"}, {{"T", DT_INT32}}}},
+      /*ret_def=*/{{"o", "ret:z:0"}},
+      /*control_ret_def=*/{{"must_execute", "add"}});
+
+  // Construct a graph:
+  //   A = Placeholder[dtype=int32]
+  //   F = PartitionedCall[f=AddAndMul](a)
+  //   B = Identity(func, ^func)
+  Scope root = Scope::NewRootScope().ExitOnError();
+  TF_ASSERT_OK(root.graph()->AddFunctionLibrary(f_lib_proto));
+  auto x = ops::Placeholder(root.WithOpName("X"), DT_INT32);
+  auto y = ops::Placeholder(root.WithOpName("Y"), DT_INT32);
+  auto z = ops::Placeholder(root.WithOpName("Z"), DT_INT32);
+  Node* function_call;
+  std::vector<NodeBuilder::NodeOut> inputs({NodeBuilder::NodeOut(x.node()),
+                                            NodeBuilder::NodeOut(y.node()),
+                                            NodeBuilder::NodeOut(z.node())});
+  TF_ASSERT_OK(NodeBuilder("F", "PartitionedCall", &root.graph()->flib_def())
+                   .Input(inputs)
+                   .Attr("Tin", {DT_INT32, DT_INT32, DT_INT32})
+                   .Attr("Tout", {DT_INT32})
+                   .Attr("f", FuncAttr("AddAndMul"))
+                   .Finalize(root.graph(), &function_call));
+  TF_ASSERT_OK(root.DoShapeInference(function_call));
+
+  auto b = ops::Identity(root.WithOpName("B"), Output(function_call, 0));
+  root.graph()->AddControlEdge(function_call, b.node());
+
+  TF_ASSERT_OK(root.ToGraph(graph.get()));
+  TF_ASSERT_OK(Rewrite(&graph));
+
+  // Verify the resultant graph has no PartitionedCall ops and function body was
+  // inlined into the main graph.
+  int partitioned_call_count = 0;
+  int add_count = 0;
+  int mul_count = 0;
+  int floor_div_count = 0;
+  for (const auto* op : graph->op_nodes()) {
+    if (op->IsPartitionedCall()) partitioned_call_count++;
+    if (op->type_string() == "Add") add_count++;
+    if (op->type_string() == "Mul") mul_count++;
+    if (op->type_string() == "FloorDiv") floor_div_count++;
+  }
+
+  ASSERT_EQ(partitioned_call_count, 0);
+  ASSERT_EQ(add_count, 1);
+  ASSERT_EQ(mul_count, 1);
+  ASSERT_EQ(floor_div_count, 0);
+
+  // Verify execution.
+  ClientSession session(root, SessionOptionsWithInlining());
+  {
+    ClientSession::FeedType feeds;
+    feeds.emplace(Output(x.node()), Input::Initializer(10));
+    std::vector<Tensor> out_tensors;
+    TF_ASSERT_OK(session.Run(feeds, {Output(b)}, &out_tensors));
+    EXPECT_EQ(out_tensors.size(), 1);
+    EXPECT_EQ(out_tensors[0].scalar<int>()(), 100);
+  }
+  flags::Global().enable_function_pruning_before_inlining.reset(false);
+}
+
 TEST(LowerFunctionCallTest, DoNotInlineTpuOrXlaFunctions) {
   std::unique_ptr<Graph> graph(new Graph(OpRegistry::Global()));
 
diff --git a/tensorflow/core/common_runtime/lower_while_op_test.cc b/tensorflow/core/common_runtime/lower_while_op_test.cc
index e1237cd5555b5f..b57145c73167ff 100644
--- a/tensorflow/core/common_runtime/lower_while_op_test.cc
+++ b/tensorflow/core/common_runtime/lower_while_op_test.cc
@@ -13,10 +13,14 @@ See the License for the specific language governing permissions and
 limitations under the License.
 ==============================================================================*/
 
+#include <memory>
+#include <vector>
+
 #include <gtest/gtest.h>
 #include "absl/strings/match.h"
 #include "tensorflow/cc/client/client_session.h"
 #include "tensorflow/cc/framework/ops.h"
+#include "tensorflow/cc/framework/scope.h"
 #include "tensorflow/cc/ops/array_ops.h"
 #include "tensorflow/cc/ops/control_flow_ops_internal.h"
 #include "tensorflow/cc/ops/function_ops.h"
@@ -24,15 +28,18 @@ limitations under the License.
 #include "tensorflow/core/common_runtime/graph_constructor.h"
 #include "tensorflow/core/common_runtime/graph_runner.h"
 #include "tensorflow/core/common_runtime/lower_functional_ops.h"
+#include "tensorflow/core/config/flag_defs.h"
 #include "tensorflow/core/framework/attr_value.pb.h"
 #include "tensorflow/core/framework/function_testlib.h"
 #include "tensorflow/core/framework/node_def_util.h"
 #include "tensorflow/core/framework/op.h"
 #include "tensorflow/core/framework/types.pb.h"
+#include "tensorflow/core/graph/graph.h"
 #include "tensorflow/core/graph/graph_def_builder.h"
 #include "tensorflow/core/lib/core/status_test_util.h"
 #include "tensorflow/core/lib/strings/str_util.h"
 #include "tensorflow/core/platform/test.h"
+#include "tsl/lib/core/status_test_util.h"
 
 namespace tensorflow {
 namespace {
@@ -173,6 +180,65 @@ TEST(LowerWhileOpTest, Simple) {
   }
 }
 
+static void DanglingNodeTestHelper(int expected_count) {
+  std::unique_ptr<Graph> graph(new Graph(OpRegistry::Global()));
+
+  // Add test functions for cond and body.
+  FunctionDefLibrary f_lib_proto;
+  *f_lib_proto.add_function() =
+      test::function::XTimesTwoWithDanglingFloorDivNode();
+  *f_lib_proto.add_function() = test::function::LessThanOrEqualToN(8);
+
+  Scope root = Scope::NewRootScope().ExitOnError();
+  TF_ASSERT_OK(root.graph()->AddFunctionLibrary(f_lib_proto));
+  auto a = ops::Placeholder(root.WithOpName("A"), DT_INT32);
+  Node* while_node;
+  std::vector<NodeBuilder::NodeOut> inputs({NodeBuilder::NodeOut(a.node())});
+  AttrValue cond_func;
+  cond_func.mutable_func()->set_name("LessThanOrEqualToN");
+  AttrValue body_func;
+  body_func.mutable_func()->set_name("XTimesTwoWithDanglingFloorDivNode");
+  TF_ASSERT_OK(
+      NodeBuilder("while", "While", &root.graph()->flib_def())
+          .Input(inputs)
+          .Attr("T", {DT_INT32})
+          .Attr("cond", cond_func)
+          .Attr("body", body_func)
+          .Attr("parallel_iterations", 100)
+          .Attr(LowerFunctionalOpsPass::kLowerUsingSwitchMergeAttr, true)
+          .Finalize(root.graph(), &while_node));
+  auto c = ops::Identity(
+      root.WithOpName("C").WithControlDependencies(Output(while_node)),
+      Output(while_node));
+  TF_ASSERT_OK(root.DoShapeInference(while_node));
+  TF_ASSERT_OK(root.ToGraph(graph.get()));
+
+  TF_ASSERT_OK(Rewrite(&graph));
+
+  int mul_count = 0;
+  int floor_div_count = 0;
+
+  for (const auto* op : graph->op_nodes()) {
+    if (op->type_string() == "Mul") {
+      mul_count++;
+    }
+    if (op->type_string() == "FloorDiv") {
+      floor_div_count++;
+    }
+  }
+
+  ASSERT_EQ(mul_count, 1);
+  ASSERT_EQ(floor_div_count, expected_count);
+}
+
+TEST(LowerWhileOpTest, DanglingNode) { DanglingNodeTestHelper(1); }
+
+TEST(LowerWhileOpTest, DanglingNodeWithPruning) {
+  flags::Global().enable_function_pruning_before_inlining.reset(true);
+  DanglingNodeTestHelper(0);
+  flags::Global().enable_function_pruning_before_inlining.reset(false);
+}
+
 TEST(LowerWhileOpTest, ForwardAssignedInputDevice) {
   std::unique_ptr<Graph> graph(new Graph(OpRegistry::Global()));
 
diff --git a/tensorflow/core/config/flag_defs.h b/tensorflow/core/config/flag_defs.h
index bc91a9bcc4247b..a773fbb1b20c1c 100644
--- a/tensorflow/core/config/flag_defs.h
+++ b/tensorflow/core/config/flag_defs.h
@@ -61,6 +61,9 @@ class Flags {
                   "propagated during while op lowering to switch/merge ops.")
   TF_DECLARE_FLAG(enable_tf2min_ici_weight, false,
                   "If true, ici weight optimization will be used in tf2/min.")
+  // TODO(b/341325107): Make this behavior the default and remove the flag.
+  TF_DECLARE_FLAG(enable_function_pruning_before_inlining, false,
+                  "If true, functions will be pruned before inlining.")
   // LINT.ThenChange(//tensorflow/core/config/flags_api_wrapper.cc)
 };
 
diff --git a/tensorflow/core/config/flags_api_wrapper.cc b/tensorflow/core/config/flags_api_wrapper.cc
index e6c2192523cca9..a4925180c1ef53 100644
--- a/tensorflow/core/config/flags_api_wrapper.cc
+++ b/tensorflow/core/config/flags_api_wrapper.cc
@@ -55,5 +55,6 @@ PYBIND11_MODULE(flags_pybind, m) {
   TF_PY_DECLARE_FLAG(enable_aggressive_constant_replication);
   TF_PY_DECLARE_FLAG(enable_colocation_key_propagation_in_while_op_lowering);
   TF_PY_DECLARE_FLAG(enable_tf2min_ici_weight)
+  TF_PY_DECLARE_FLAG(enable_function_pruning_before_inlining)
   // LINT.ThenChange(//tensorflow/core/config/flag_defs.h)
 };
diff --git a/tensorflow/core/framework/function_testlib.cc b/tensorflow/core/framework/function_testlib.cc
index 7303228b935bbf..ae06188b8bc83a 100644
--- a/tensorflow/core/framework/function_testlib.cc
+++ b/tensorflow/core/framework/function_testlib.cc
@@ -15,8 +15,11 @@ limitations under the License.
 
 #include "tensorflow/core/framework/function_testlib.h"
 
+#include <cstdint>
+
 #include "tensorflow/core/framework/function.h"
 #include "tensorflow/core/framework/node_def.pb.h"
+#include "tensorflow/core/framework/tensor.h"
 #include "tensorflow/core/framework/tensor_testutil.h"
 #include "tensorflow/core/framework/versions.pb.h"
 #include "tensorflow/core/lib/core/threadpool.h"
@@ -180,6 +183,26 @@ FunctionDef XTimesTwoWithControlOutput() {
       {{"dummy", "dummy"}});
 }
 
+FunctionDef XTimesTwoWithDanglingFloorDivNode() {
+  const Tensor kTwo = test::AsScalar<int64_t>(2);
+  return FDH::Define(
+      // Name
+      "XTimesTwoWithDanglingFloorDivNode",
+      // Args
+      {"x: T"},
+      // Return values
+      {"y: T"},
+      // Attr def
+      {"T: {float, double, int32, int64}"},
+      // Nodes
+      {
+          {{"two"}, "Const", {}, {{"value", kTwo}, {"dtype", DT_INT64}}},
+          {{"scale"}, "Cast", {"two"}, {{"SrcT", DT_INT64}, {"DstT", "$T"}}},
+          {{"z"}, "FloorDiv", {"x", "scale"}, {{"T", "$T"}}},
+          {{"y"}, "Mul", {"x", "scale"}, {{"T", "$T"}}},
+      });
+}
+
 FunctionDef TwoDeviceMult() {
   const Tensor kTwo = test::AsScalar<int64_t>(2);
   const Tensor kThree = test::AsScalar<int64_t>(3);
diff --git a/tensorflow/core/framework/function_testlib.h b/tensorflow/core/framework/function_testlib.h
index 0aba30650d5f84..06e0c3a6d36ca9 100644
--- a/tensorflow/core/framework/function_testlib.h
+++ b/tensorflow/core/framework/function_testlib.h
@@ -74,6 +74,8 @@ FunctionDef XTimesTwo();
 FunctionDef XTimesTwoWithControlInput();
 // Same as `XTimesTwo` above, but with a `dummy` control output node.
 FunctionDef XTimesTwoWithControlOutput();
+// Same as `XTimesTwo` above, but with a dangling `FloorDiv` node.
+FunctionDef XTimesTwoWithDanglingFloorDivNode();
 
 // x: T -> cpu(x * 2) + cpu(x * 3).
 FunctionDef TwoDeviceTimesFive();
diff --git a/tensorflow/python/flags_pybind.pyi b/tensorflow/python/flags_pybind.pyi
index c08b02ddcb9b5d..b34ed2f4b68c19 100644
--- a/tensorflow/python/flags_pybind.pyi
+++ b/tensorflow/python/flags_pybind.pyi
@@ -21,6 +21,7 @@ class Flag:
 class Flags:
     enable_aggressive_constant_replication: Flag
     enable_colocation_key_propagation_in_while_op_lowering: Flag
+    enable_function_pruning_before_inlining: Flag
     enable_nested_function_shape_inference: Flag
     enable_quantized_dtypes_training: Flag
     enable_tf2min_ici_weight: Flag

From a8b7969a726dcbb5c3fd628a6b0d89827ab454f8 Mon Sep 17 00:00:00 2001
From: Henning Becker <hebecker@google.com>
Date: Fri, 17 May 2024 11:50:03 -0700
Subject: [PATCH 365/478] Remove duplicate test macros from TSL

`tsl/platform/test.h` used to define a macro `DISABLED_ON_GPU_ROCM` which also exists in
`xla/tests/test_macros.h`. This led to build failures if both headers were included.

Since `test_macros.h` contains a huge set of similar macros let's remove the one
from `platform/test.h` and use `test_macros.h` everywhere.

I'm also fixing a CUDA include issue along the way.

PiperOrigin-RevId: 634844088
---
 tensorflow/core/common_runtime/gpu/BUILD                    | 2 ++
 tensorflow/core/common_runtime/gpu/gpu_device_test.cc       | 6 ++++++
 third_party/xla/third_party/tsl/tsl/platform/test.h         | 6 ------
 third_party/xla/xla/service/gpu/BUILD                       | 1 +
 .../xla/xla/service/gpu/conv_layout_normalization_test.cc   | 1 +
 5 files changed, 10 insertions(+), 6 deletions(-)

diff --git a/tensorflow/core/common_runtime/gpu/BUILD b/tensorflow/core/common_runtime/gpu/BUILD
index 84ec94ba673ff8..8fdd8521c83cc6 100644
--- a/tensorflow/core/common_runtime/gpu/BUILD
+++ b/tensorflow/core/common_runtime/gpu/BUILD
@@ -328,6 +328,7 @@ tf_cuda_cc_test(
         "//tensorflow/core/kernels:ops_util",
         "@local_tsl//tsl/framework:device_id",
         "@local_xla//xla/stream_executor/gpu:gpu_cudamallocasync_allocator_header",
+        "@local_xla//xla/tests:test_macros_header",
     ],
 )
 
@@ -388,6 +389,7 @@ tf_cuda_cc_test(
         "//tensorflow/core/common_runtime:direct_session_internal",
         "//tensorflow/core/kernels:ops_util",
         "@local_xla//xla/stream_executor/gpu:gpu_cudamallocasync_allocator_header",
+        "@local_xla//xla/tests:test_macros_header",
     ],
 )
 
diff --git a/tensorflow/core/common_runtime/gpu/gpu_device_test.cc b/tensorflow/core/common_runtime/gpu/gpu_device_test.cc
index 8f453f442e2950..7a9a0c2661f571 100644
--- a/tensorflow/core/common_runtime/gpu/gpu_device_test.cc
+++ b/tensorflow/core/common_runtime/gpu/gpu_device_test.cc
@@ -23,6 +23,7 @@ limitations under the License.
 
 #include "xla/stream_executor/gpu/gpu_cudamallocasync_allocator.h"
 #include "xla/stream_executor/gpu/gpu_init.h"
+#include "xla/tests/test_macros.h"
 #include "tensorflow/core/common_runtime/gpu/gpu_process_state.h"
 #include "tensorflow/core/platform/env.h"
 #include "tensorflow/core/platform/errors.h"
@@ -37,6 +38,11 @@ limitations under the License.
 #include "tensorflow/core/tfrt/common/pjrt_util.h"
 #endif  // TF_GPU_USE_PJRT
 
+#if GOOGLE_CUDA
+// Needed for CUDA_VERSION preprocessor directive
+#include "third_party/gpus/cuda/include/cuda.h"
+#endif
+
 namespace tensorflow {
 namespace {
 
diff --git a/third_party/xla/third_party/tsl/tsl/platform/test.h b/third_party/xla/third_party/tsl/tsl/platform/test.h
index 313bfe5f0ea3dd..c8950099cad936 100644
--- a/third_party/xla/third_party/tsl/tsl/platform/test.h
+++ b/third_party/xla/third_party/tsl/tsl/platform/test.h
@@ -47,12 +47,6 @@ limitations under the License.
 #endif
 #include <gmock/gmock.h>  // IWYU pragma: export
 
-#define DISABLED_ON_GPU_ROCM(X) X
-#if TENSORFLOW_USE_ROCM
-#undef DISABLED_ON_GPU_ROCM
-#define DISABLED_ON_GPU_ROCM(X) DISABLED_##X
-#endif  // TENSORFLOW_USE_ROCM
-
 namespace tsl {
 namespace testing {
 
diff --git a/third_party/xla/xla/service/gpu/BUILD b/third_party/xla/xla/service/gpu/BUILD
index 92d6d3d5596b2b..30b13e0f98b3b0 100644
--- a/third_party/xla/xla/service/gpu/BUILD
+++ b/third_party/xla/xla/service/gpu/BUILD
@@ -4873,6 +4873,7 @@ xla_cc_test(
         "//xla/hlo/ir:hlo",
         "//xla/service/gpu/tests:gpu_codegen_test",  # fixdeps: keep
         "//xla/tests:hlo_test_base",
+        "//xla/tests:test_macros_header",
         "@local_tsl//tsl/platform:test",
         "@local_tsl//tsl/platform:test_main",
     ],
diff --git a/third_party/xla/xla/service/gpu/conv_layout_normalization_test.cc b/third_party/xla/xla/service/gpu/conv_layout_normalization_test.cc
index 56629fe39bee3f..206416b332f8c3 100644
--- a/third_party/xla/xla/service/gpu/conv_layout_normalization_test.cc
+++ b/third_party/xla/xla/service/gpu/conv_layout_normalization_test.cc
@@ -16,6 +16,7 @@ limitations under the License.
 #include "xla/error_spec.h"
 #include "xla/hlo/ir/hlo_module.h"
 #include "xla/tests/hlo_test_base.h"
+#include "xla/tests/test_macros.h"
 #include "tsl/platform/test.h"
 
 namespace xla {

From 5e1bc5b467949e740aacd8e5231fa2dcd0b9d1dc Mon Sep 17 00:00:00 2001
From: Dmitri Gribenko <dmitrig@google.com>
Date: Fri, 17 May 2024 11:53:22 -0700
Subject: [PATCH 366/478] Integrate LLVM at llvm/llvm-project@371eccd5dfed

Updates LLVM usage to match
[371eccd5dfed](https://github.com/llvm/llvm-project/commit/371eccd5dfed)

PiperOrigin-RevId: 634844947
---
 third_party/llvm/workspace.bzl | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/third_party/llvm/workspace.bzl b/third_party/llvm/workspace.bzl
index 29dfa535cd78f4..1893eaa66f2c53 100644
--- a/third_party/llvm/workspace.bzl
+++ b/third_party/llvm/workspace.bzl
@@ -4,8 +4,8 @@ load("//third_party:repo.bzl", "tf_http_archive")
 
 def repo(name):
     """Imports LLVM."""
-    LLVM_COMMIT = "a68d20e986053ec571223a9f3ead3e146a27dc82"
-    LLVM_SHA256 = "4c53512522cfd625a75aa7f201e6932f18ebf3fbe8ee89e05883c3d30ca935ac"
+    LLVM_COMMIT = "371eccd5dfed88c8e76449233d8388c12be3464b"
+    LLVM_SHA256 = "31e51b40d2890d6389dc818ab45bf1395b239b9e1d5b280b58f63e260b6deae8"
 
     tf_http_archive(
         name = name,

From bc8d197ecd56272d6f4f8633f5b3d274dcec09c0 Mon Sep 17 00:00:00 2001
From: Eugene Zhulenev <ezhulenev@google.com>
Date: Fri, 17 May 2024 12:25:05 -0700
Subject: [PATCH 367/478] [xla:python] NFC: Fix all warnings in xla_compiler.cc

PiperOrigin-RevId: 634854042
---
 third_party/xla/xla/python/BUILD           |  5 +++++
 third_party/xla/xla/python/xla_compiler.cc | 17 +++++++++++------
 2 files changed, 16 insertions(+), 6 deletions(-)

diff --git a/third_party/xla/xla/python/BUILD b/third_party/xla/xla/python/BUILD
index 48501d4600b401..ada2f76f96b0ff 100644
--- a/third_party/xla/xla/python/BUILD
+++ b/third_party/xla/xla/python/BUILD
@@ -1087,6 +1087,8 @@ cc_library(
         ":py_client",
         ":types",
         # placeholder for index annotation deps
+        "@com_google_absl//absl/base:core_headers",
+        "@com_google_absl//absl/container:inlined_vector",
         "@com_google_absl//absl/hash",
         "@com_google_absl//absl/status",
         "@com_google_absl//absl/strings",
@@ -1126,7 +1128,10 @@ cc_library(
         "//xla/service:name_uniquer",
         "//xla/service:tuple_simplifier",
         "@local_tsl//tsl/lib/strings:proto_serialization",
+        "@local_tsl//tsl/platform:env",
+        "@local_tsl//tsl/platform:errors",
         "@local_tsl//tsl/platform:logging",
+        "@local_tsl//tsl/platform:statusor",
     ],
 )
 
diff --git a/third_party/xla/xla/python/xla_compiler.cc b/third_party/xla/xla/python/xla_compiler.cc
index 1e4cbf1aebe7c5..27022340dbc323 100644
--- a/third_party/xla/xla/python/xla_compiler.cc
+++ b/third_party/xla/xla/python/xla_compiler.cc
@@ -23,11 +23,14 @@ limitations under the License.
 #include <utility>
 #include <vector>
 
+#include "absl/base/thread_annotations.h"
+#include "absl/container/inlined_vector.h"
 #include "absl/hash/hash.h"
 #include "absl/status/status.h"
 #include "absl/strings/str_cat.h"
 #include "absl/strings/str_format.h"
 #include "absl/strings/str_join.h"
+#include "absl/strings/str_split.h"
 #include "absl/strings/string_view.h"
 #include "absl/synchronization/mutex.h"
 #include "absl/types/span.h"
@@ -77,12 +80,14 @@ limitations under the License.
 #include "xla/service/tuple_simplifier.h"
 #include "xla/shape.h"
 #include "xla/shape_util.h"
-#include "xla/statusor.h"
 #include "xla/util.h"
 #include "xla/xla.pb.h"
 #include "xla/xla_data.pb.h"
 #include "tsl/lib/strings/proto_serialization.h"
+#include "tsl/platform/env.h"
+#include "tsl/platform/errors.h"
 #include "tsl/platform/logging.h"
+#include "tsl/platform/statusor.h"
 
 namespace nanobind {
 namespace detail {
@@ -617,24 +622,24 @@ void BuildXlaCompilerSubmodule(nb::module_& m) {
    public:
     ComputationWrapper(const HloComputation* comp,
                        const std::shared_ptr<HloModule> module)
-        : comp{comp}, module{module} {}
-    absl::string_view name() const { return comp->name(); }
+        : comp_(comp), module_(module) {}
+    absl::string_view name() const { return comp_->name(); }
     void render_html(const std::string& filename) {
       std::string html = xla::ValueOrThrow(RenderGraph(
-          *comp, /*label=*/"", comp->parent()->config().debug_options(),
+          *comp_, /*label=*/"", comp_->parent()->config().debug_options(),
           RenderedGraphFormat::kHtml, HloRenderOptions()));
       xla::ThrowIfError(tsl::WriteStringToFile(
           tsl::Env::Default(), absl::StrCat(filename, ".html"), html));
     }
 
    private:
-    const HloComputation* comp;
+    const HloComputation* comp_;
     // The module owns the computations: if its destructor is called, the
     // computations are freed. To prevent that from happening in cases where the
     // module Python object goes out of scope and gets garbage collected before
     // the computations, we keep a shared_ptr to the module that originated the
     // computation.
-    const std::shared_ptr<HloModule> module;
+    const std::shared_ptr<HloModule> module_;
   };
 
   nb::class_<ComputationWrapper> hlo_computation_class(m, "HloComputation");

From 59683afae964dc56fe3d9f1fce4ab7a1dfe68a1e Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Fri, 17 May 2024 12:45:50 -0700
Subject: [PATCH 368/478] Use default relative error tolerance for Cbrt.

PiperOrigin-RevId: 634859424
---
 .../tests/exhaustive/exhaustive_unary_test_f32_or_smaller.cc  | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/third_party/xla/xla/tests/exhaustive/exhaustive_unary_test_f32_or_smaller.cc b/third_party/xla/xla/tests/exhaustive/exhaustive_unary_test_f32_or_smaller.cc
index db3fc8d4d8aff4..56240f06343359 100644
--- a/third_party/xla/xla/tests/exhaustive/exhaustive_unary_test_f32_or_smaller.cc
+++ b/third_party/xla/xla/tests/exhaustive/exhaustive_unary_test_f32_or_smaller.cc
@@ -388,10 +388,10 @@ UNARY_TEST_FLOAT_32_BITS_OR_LESS(Cbrt, {
   auto error_spec_gen = +[](NativeT x) {
     NativeT eps = std::numeric_limits<NativeT>::epsilon();
     NativeT min = std::numeric_limits<NativeT>::min();
-    // Allow a small absollute error (e.g. 9e-16 for F32).
+    // Allow a small absolute error (e.g. 9e-16 for F32).
     // This corresponds to a 0.5% relative error for the
     // smallest normalized floating point values.
-    return ErrorSpec{.abs_err = std::cbrt(min) / 200, .rel_err = 10 * eps};
+    return ErrorSpec{.abs_err = std::cbrt(min) / 200, .rel_err = 50 * eps};
   };
   Run(Cbrt, std::cbrt, error_spec_gen);
 })

From 43e54f62fb15fc03ac437bd490b955170e79894e Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Fri, 17 May 2024 12:46:46 -0700
Subject: [PATCH 369/478] Improve errors upon service shutdown.

PiperOrigin-RevId: 634859665
---
 .../distributed_runtime/coordination/BUILD    |  1 +
 .../coordination/coordination_service.cc      | 84 ++++++++++++++-----
 .../coordination/coordination_service.h       |  4 +
 .../coordination/coordination_service_agent.h | 13 ++-
 .../coordination/coordination_service_test.cc |  6 +-
 5 files changed, 79 insertions(+), 29 deletions(-)

diff --git a/third_party/xla/xla/tsl/distributed_runtime/coordination/BUILD b/third_party/xla/xla/tsl/distributed_runtime/coordination/BUILD
index bf67649bb5474f..137197420cfab2 100644
--- a/third_party/xla/xla/tsl/distributed_runtime/coordination/BUILD
+++ b/third_party/xla/xla/tsl/distributed_runtime/coordination/BUILD
@@ -72,6 +72,7 @@ tsl_gpu_library(
         ":coordination_service_error_util",
         "//xla/tsl/distributed_runtime:call_options",
         "//xla/tsl/util:device_name_utils",
+        "@com_google_absl//absl/base:core_headers",
         "@com_google_absl//absl/container:flat_hash_map",
         "@com_google_absl//absl/container:flat_hash_set",
         "@com_google_absl//absl/hash",
diff --git a/third_party/xla/xla/tsl/distributed_runtime/coordination/coordination_service.cc b/third_party/xla/xla/tsl/distributed_runtime/coordination/coordination_service.cc
index 8c3a0fb1a57edf..9b0dc57084d020 100644
--- a/third_party/xla/xla/tsl/distributed_runtime/coordination/coordination_service.cc
+++ b/third_party/xla/xla/tsl/distributed_runtime/coordination/coordination_service.cc
@@ -26,6 +26,7 @@ limitations under the License.
 #include <utility>
 #include <vector>
 
+#include "absl/base/thread_annotations.h"
 #include "absl/container/flat_hash_map.h"
 #include "absl/container/flat_hash_set.h"
 #include "absl/hash/hash.h"
@@ -225,6 +226,7 @@ class CoordinationServiceStandaloneImpl : public CoordinationServiceInterface {
     CoordinatedTaskState state_ = CoordinatedTaskState::TASKSTATE_DISCONNECTED;
     absl::Status status_;
     absl::Mutex last_heartbeat_mu_;
+    // TODO(b/339231167): Migrate to use absl annotations.
     uint64_t last_heartbeat_us_ TF_GUARDED_BY(last_heartbeat_mu_);
     // This denotes the deadline after which we stop accepting heartbeats from a
     // disconnected task. This grace period accounts for the lag time between
@@ -264,10 +266,8 @@ class CoordinationServiceStandaloneImpl : public CoordinationServiceInterface {
   absl::flat_hash_map<std::string, std::vector<StatusOrValueCallback>> get_cb_
       TF_GUARDED_BY(kv_mu_);
 
-  absl::Mutex check_staleness_thread_shutdown_mu_;
   absl::CondVar check_staleness_thread_cv_;
-  bool shutting_down_ TF_GUARDED_BY(check_staleness_thread_shutdown_mu_) =
-      false;
+  bool shutting_down_ ABSL_GUARDED_BY(state_mu_) = false;
   std::unique_ptr<Thread> check_staleness_thread_;
 
   absl::flat_hash_map<std::string, BarrierState> barriers_
@@ -311,6 +311,7 @@ absl::Status CoordinationServiceStandaloneImpl::TaskState::RecordHeartbeat(
     uint64_t task_incarnation) {
   if (!status_.ok()) return status_;
   if (task_incarnation != task_incarnation_) {
+    // TODO(b/339231167): Migrate to use absl error API.
     return MakeCoordinationError(errors::Aborted(
         "Incarnation ID mismatch: expecting ", task_incarnation_, " but got ",
         task_incarnation, ". This means the remote task has restarted."));
@@ -389,9 +390,9 @@ void CoordinationServiceStandaloneImpl::StartCheckStaleness() {
         absl::flat_hash_map<std::string, BarrierState*> expired_barriers;
         while (true) {
           {
-            absl::MutexLock l(&check_staleness_thread_shutdown_mu_);
-            check_staleness_thread_cv_.WaitWithTimeout(
-                &check_staleness_thread_shutdown_mu_, absl::Seconds(1));
+            absl::MutexLock l(&state_mu_);
+            check_staleness_thread_cv_.WaitWithTimeout(&state_mu_,
+                                                       absl::Seconds(1));
             if (shutting_down_) {
               return;
             }
@@ -507,6 +508,11 @@ void CoordinationServiceStandaloneImpl::Stop(bool shut_staleness_thread) {
   }
   {
     absl::MutexLock l(&state_mu_);
+    // Indicate that the service is shutting down and stop accepting new RPCs.
+    shutting_down_ = true;
+    // Stop the heartbeat thread.
+    check_staleness_thread_cv_.SignalAll();
+    // Fail all ongoing barriers.
     for (auto& [barrier_id, barrier] : barriers_) {
       if (!barrier.passed) {
         absl::Status error = MakeCoordinationError(errors::Aborted(absl::StrCat(
@@ -516,22 +522,19 @@ void CoordinationServiceStandaloneImpl::Stop(bool shut_staleness_thread) {
       }
     }
     barriers_.clear();
-    // Cluster state is used in `PassBarrier` and it needs to be cleared after
-    // it.
+    // Erase cluster state.
+    // Note: sequence matters here, this must happen after barrier clean-up as
+    // the state is used in `PassBarrier`.
     cluster_state_.clear();
   }
-  {
-    absl::MutexLock l(&check_staleness_thread_shutdown_mu_);
-    shutting_down_ = true;
-    check_staleness_thread_cv_.SignalAll();
-  }
+  // Destroy thread outside of the mutex.
   if (shut_staleness_thread) {
     check_staleness_thread_.reset();
   }
 }
 
 bool CoordinationServiceStandaloneImpl::ServiceHasStopped() const {
-  return cluster_state_.empty();
+  return shutting_down_;
 }
 
 // Helper to log progress to having waited for all tasks.
@@ -562,6 +565,15 @@ absl::Status CoordinationServiceStandaloneImpl::RegisterTask(
   std::string error_message;
   {
     absl::MutexLock l(&state_mu_);
+    if (ServiceHasStopped()) {
+      return MakeCoordinationError(absl::InternalError(absl::StrCat(
+          "Coordination service has stopped. RegisterTask() from task: ",
+          task_name,
+          " failed. This usually implies an earlier error that caused "
+          "coordination service to shut down before the workers disconnect "
+          "gracefully. Check the task leader's logs for an earlier error to "
+          "debug the root cause.")));
+    }
     if (!cluster_state_.contains(task_name)) {
       // Note: return early here as unexpected task register errors should not
       // be propagated to other tasks.
@@ -632,6 +644,11 @@ void CoordinationServiceStandaloneImpl::WaitForAllTasks(
     StatusCallback done) {
   {
     absl::MutexLock l(&state_mu_);
+    if (ServiceHasStopped()) {
+      done(MakeCoordinationError(absl::InternalError(
+          "Coordination service has stopped. WaitForAllTasks() failed.")));
+      return;
+    }
     const auto& task_state = cluster_state_.find(GetTaskName(task));
     // Collect task device info for the first time that task
     // has called WaitForAllTasks(). This will be aggregated when the barrier
@@ -655,8 +672,13 @@ void CoordinationServiceStandaloneImpl::ShutdownTaskAsync(
     absl::Status status;
     {
       absl::MutexLock l(&state_mu_);
-      // Disconnect task from service individually.
-      status = DisconnectTask(task);
+      if (ServiceHasStopped()) {
+        status = MakeCoordinationError(absl::InternalError(
+            "Coordination service has stopped. ShutdownTaskAsync() failed."));
+      } else {
+        // Disconnect task from service individually.
+        status = DisconnectTask(task);
+      }
     }
     done(status);
   }
@@ -672,7 +694,12 @@ absl::Status CoordinationServiceStandaloneImpl::DisconnectTask(
     const CoordinatedTask& task) {
   const std::string task_name = GetTaskName(task);
   // Check if task is valid and not already disconnected.
-  if (!cluster_state_.contains(task_name)) {
+  if (ServiceHasStopped()) {
+    return MakeCoordinationError(absl::InternalError(
+        absl::StrCat("Coordination service has stopped. DisconnectTask() "
+                     "failed for task_name=",
+                     task_name)));
+  } else if (!cluster_state_.contains(task_name)) {
     return MakeCoordinationError(errors::InvalidArgument(
         "Unexpected disconnect request with task_name=", task_name));
   } else if (cluster_state_[task_name]->GetState() ==
@@ -709,7 +736,10 @@ absl::Status CoordinationServiceStandaloneImpl::ReportTaskError(
   const std::string task_name = GetTaskName(task);
   {
     absl::MutexLock l(&state_mu_);
-    if (!cluster_state_.contains(task_name)) {
+    if (ServiceHasStopped()) {
+      return MakeCoordinationError(absl::InternalError(
+          "Coordination service has stopped. ReportTaskError() failed."));
+    } else if (!cluster_state_.contains(task_name)) {
       return MakeCoordinationError(
           errors::InvalidArgument("Unexpected request from task ", task_name));
     } else if (cluster_state_[task_name]->GetState() !=
@@ -754,12 +784,18 @@ absl::Status CoordinationServiceStandaloneImpl::RecordHeartbeat(
   absl::Status s = absl::OkStatus();
   {
     absl::MutexLock l(&state_mu_);
-    if (!cluster_state_.contains(task_name)) {
+    if (ServiceHasStopped()) {
+      return MakeCoordinationError(absl::InternalError(absl::StrCat(
+          "Coordination service has stopped. RecordHeartbeat() from task: ",
+          task_name,
+          " failed. This usually implies an earlier error that caused "
+          "coordination service to shut down before the workers disconnect "
+          "gracefully. Check the task leader's logs for an earlier error to "
+          "debug the root cause.")));
+    } else if (!cluster_state_.contains(task_name)) {
       return MakeCoordinationError(errors::InvalidArgument(
           "Unexpected heartbeat request from task: ", task_name,
-          ". This usually implies an earlier error that caused coordination "
-          "service to shut down before the workers disconnect. Check the task "
-          "leader's logs for an earlier error to debug the root cause."));
+          ". This usually implies a configuration error."));
     }
     if (!cluster_state_[task_name]->GetStatus().ok()) {
       return cluster_state_[task_name]->GetStatus();
@@ -1194,6 +1230,10 @@ void CoordinationServiceStandaloneImpl::BarrierAsync(
 absl::Status CoordinationServiceStandaloneImpl::CancelBarrier(
     const std::string& barrier_id, const CoordinatedTask& task) {
   absl::MutexLock l(&state_mu_);
+  if (ServiceHasStopped()) {
+    return MakeCoordinationError(absl::InternalError(
+        "Coordination service has stopped. CancelBarrier() failed."));
+  }
   auto [it, inserted] = barriers_.try_emplace(barrier_id);
   auto* barrier = &it->second;
   if (inserted) {
diff --git a/third_party/xla/xla/tsl/distributed_runtime/coordination/coordination_service.h b/third_party/xla/xla/tsl/distributed_runtime/coordination/coordination_service.h
index fbf3206c5e95d3..627cce09186f9b 100644
--- a/third_party/xla/xla/tsl/distributed_runtime/coordination/coordination_service.h
+++ b/third_party/xla/xla/tsl/distributed_runtime/coordination/coordination_service.h
@@ -117,6 +117,7 @@ class CoordinationServiceInterface {
 
   // Register a task to the service.
   // Possible service errors:
+  //   - Internal: Service has shut down.
   //   - InvalidArgument: Unexpected task request.
   //   - Aborted: (1) task is in error state, or (2) task is in connected state
   //       with a different incarnation, indicating that it restarted.
@@ -136,6 +137,7 @@ class CoordinationServiceInterface {
   // specified in the config, blocks until all tasks reach the barrier before
   // disconnecting together.
   // Possible service errors:
+  //   - Internal: Service has shut down.
   //   - InvalidArgument: Unexpected task request.
   //   - FailedPrecondition: task has already disconnected.
   virtual void ShutdownTaskAsync(const tensorflow::CoordinatedTask& task,
@@ -143,12 +145,14 @@ class CoordinationServiceInterface {
 
   // Disconnects task from the service and cleans up its internal error state.
   // Possible service errors:
+  //   - Internal: Service has shut down.
   //   - InvalidArgument: Unexpected task request.
   //   - FailedPrecondition: task has already disconnected.
   virtual absl::Status ResetTask(const tensorflow::CoordinatedTask& task) = 0;
 
   // Update the heartbeat timestamp of a task. This should only be invoked on
   // the leader of the cluster.
+  //   - Internal: Service has shut down.
   virtual absl::Status RecordHeartbeat(const tensorflow::CoordinatedTask& task,
                                        uint64_t incarnation) = 0;
 
diff --git a/third_party/xla/xla/tsl/distributed_runtime/coordination/coordination_service_agent.h b/third_party/xla/xla/tsl/distributed_runtime/coordination/coordination_service_agent.h
index 418ca50b0db122..d9cddbe7d9642b 100644
--- a/third_party/xla/xla/tsl/distributed_runtime/coordination/coordination_service_agent.h
+++ b/third_party/xla/xla/tsl/distributed_runtime/coordination/coordination_service_agent.h
@@ -52,9 +52,7 @@ class Env;
 // reported by the user via `agent->ReportError()`.
 //
 // Possible service errors:
-//    - Internal: Coordination service is not enabled.
-//                        If it was previously accessible, coordination service
-//                        has been shut down.
+//    - Internal: Coordination service has shut down or has not been enabled.
 //    - Aborted: Incarnation mismatch during heartbeat (either remote
 //                       task or coordination service has restarted).
 //    - Unavailable: Heartbeat timeout from remote task (failed,
@@ -100,6 +98,7 @@ class CoordinationServiceAgent {
   //   - register itself as a task to the service
   //   - start a thread to periodically send heartbeat message with the service
   // Possible service errors:
+  //   - Internal: Coordination service has shut down.
   //   - FailedPrecondition: Agent is not in DISCONNECTED state.
   //   - InvalidArgument: Unexpected task registration
   //   - Aborted: Duplicate task registration (agent will retry connecting until
@@ -109,6 +108,7 @@ class CoordinationServiceAgent {
   // Wait for all tasks to be up and registered. The call blocks until all tasks
   // in the cluster are up, or some error occurs.
   // Possible service errors:
+  //   - Internal: Coordination service has shut down.
   //   - FailedPrecondition: Agent is not in CONNECTED state.
   //   - InvalidArgument: Unexpected task request
   virtual absl::Status WaitForAllTasks(
@@ -136,6 +136,7 @@ class CoordinationServiceAgent {
   // Note that the error payload will set `is_reported_error` to true, to
   // distinguish user-specified errors from internal service or RPC failures.
   // Possible service errors:
+  //   - Internal: Coordination service has shut down.
   //   - FailedPrecondition: Uninitialized/disconnected/already in error state.
   //   - InvalidArgument: Unexpected task request
   virtual absl::Status ReportError(const absl::Status& error) = 0;
@@ -147,6 +148,7 @@ class CoordinationServiceAgent {
   // the barrier times out, this agent will still disconnect, while an error is
   // reported to other agents that did not reach the barrier on time.
   // Possible service errors:
+  //   - Internal: Coordination service has shut down.
   //   - InvalidArgument: Unexpected task request.
   //   - FailedPrecondition: Task was in error state (note: agent is still
   //                         shut down forcefully).
@@ -154,6 +156,7 @@ class CoordinationServiceAgent {
 
   // Disconnect from the service, and clean up the internal error status.
   // Possible service errors:
+  //   - Internal: Coordination service has shut down.
   //   - InvalidArgument: Unexpected task request.
   //   - FailedPrecondition: task is not in error state/has already
   //       disconnected.
@@ -233,7 +236,8 @@ class CoordinationServiceAgent {
   //      first WaitAtBarrier() + timeout duration.
   //   - Cancelled: One of the tasks called CancelBarrier().
   //   - Aborted: Service is shutting down.
-  //   - Internal: Any participating task is in ERROR state.
+  //   - Internal: Any participating task is in ERROR state, or service has shut
+  //     down.
   //   - InvalidArgument: (1) Conflicting tasks specified by different agents
   //       for the same barrier, (2) one of the participating tasks is not in
   //       the cluster, or (3) task making the request is not included in the
@@ -253,6 +257,7 @@ class CoordinationServiceAgent {
   // Current and future WaitAtBarrier() calls with the same id will return a
   // CANCELLED error status.
   // Possible service errors:
+  //   - Internal: Coordination service has shut down.
   //   - FailedPrecondition: Barrier has already been passed.
   virtual absl::Status CancelBarrier(std::string_view barrier_id) = 0;
   virtual void CancelBarrierAsync(std::string_view barrier_id,
diff --git a/third_party/xla/xla/tsl/distributed_runtime/coordination/coordination_service_test.cc b/third_party/xla/xla/tsl/distributed_runtime/coordination/coordination_service_test.cc
index d5d0dea523396b..8e222353e078e0 100644
--- a/third_party/xla/xla/tsl/distributed_runtime/coordination/coordination_service_test.cc
+++ b/third_party/xla/xla/tsl/distributed_runtime/coordination/coordination_service_test.cc
@@ -473,9 +473,9 @@ TEST_F(CoordinateTwoTasksTest,
       absl::ToInt64Microseconds(2 * kHeartbeatTimeout));
   // Unexpected heartbeat from unregistered tasks since service state has been
   // reset.
-  EXPECT_TRUE(absl::IsInvalidArgument(
+  EXPECT_TRUE(absl::IsInternal(
       coord_service_->RecordHeartbeat(task_0_, incarnation_0_)));
-  EXPECT_TRUE(absl::IsInvalidArgument(
+  EXPECT_TRUE(absl::IsInternal(
       coord_service_->RecordHeartbeat(task_1_, incarnation_1_)));
 }
 
@@ -1448,7 +1448,7 @@ TEST_F(CoordinateTwoTasksTest,
   // service has stopped yet, which should fail.
   absl::Status s = coord_service_->RecordHeartbeat(task_1_, incarnation_1_);
 
-  EXPECT_TRUE(absl::IsInvalidArgument(s)) << s;
+  EXPECT_TRUE(absl::IsInternal(s)) << s;
 }
 
 TEST_F(CoordinateTwoTasksTest, BarrierFailsIfServiceHasStopped) {

From d928140cd73aae8f01d2f62b60140db995b17c65 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Fri, 17 May 2024 14:01:03 -0700
Subject: [PATCH 370/478] Remove reference to `mlir::Operations` in `Thunks`.

The `mlir::Operation` in `Thunk` (`Thunk::op_`) is never set to a value different from `nullptr`.

Conveniently, this also removes the need for `Thunk::ClearCompileTimeInfo`, which in turn leaves `ForAllThunks` (in `compile_module_to_llvm_ir.cc`) unused.

I'm doing this first so that I can then (a) move `ForAllThunks` to its own file, and (b) change its behavior to actually iterate over all `Thunks`: for example, `ForAllThunks` currently does not iterate over the `embedded_thunk` in `AddressComputationThunk`. This behavior will be safer to change after there's no remaining users of the function.

PiperOrigin-RevId: 634879646
---
 .../service/gpu/compile_module_to_llvm_ir.cc  |  5 +-
 .../xla/xla/service/gpu/gpu_compiler.cc       |  2 -
 .../gpu/runtime/address_computation_thunk.cc  |  2 +-
 .../runtime/address_computation_thunk_test.cc | 63 +++++++++----------
 .../gpu/runtime/command_buffer_thunk_test.cc  | 24 +++----
 .../xla/xla/service/gpu/runtime/copy_thunk.h  |  2 -
 .../xla/service/gpu/runtime/kernel_thunk.h    |  2 -
 .../xla/service/gpu/runtime/memset_thunk.h    |  4 --
 .../xla/xla/service/gpu/runtime/thunk.cc      | 10 +--
 .../xla/xla/service/gpu/runtime/thunk.h       | 16 +----
 .../xla/service/gpu/runtime/while_thunk.cc    |  4 +-
 11 files changed, 46 insertions(+), 88 deletions(-)

diff --git a/third_party/xla/xla/service/gpu/compile_module_to_llvm_ir.cc b/third_party/xla/xla/service/gpu/compile_module_to_llvm_ir.cc
index 40b3e20b683fe8..371d0c5e8c35c4 100644
--- a/third_party/xla/xla/service/gpu/compile_module_to_llvm_ir.cc
+++ b/third_party/xla/xla/service/gpu/compile_module_to_llvm_ir.cc
@@ -234,10 +234,7 @@ absl::StatusOr<CompileModuleResults> CompileModuleToLlvmIr(
     RecordHloToLlvmDuration(end_usecs - start_usecs);
   }
 
-  auto thunk_sequence = ir_emitter->ConsumeThunkSequence();
-  ForAllThunks([](Thunk* thunk) { thunk->ClearCompileTimeInfo(); },
-               thunk_sequence.get());
-  results.executable = std::move(thunk_sequence);
+  results.executable = ir_emitter->ConsumeThunkSequence();
 
   return results;
 }
diff --git a/third_party/xla/xla/service/gpu/gpu_compiler.cc b/third_party/xla/xla/service/gpu/gpu_compiler.cc
index 87126e13307473..5d789ff717081f 100644
--- a/third_party/xla/xla/service/gpu/gpu_compiler.cc
+++ b/third_party/xla/xla/service/gpu/gpu_compiler.cc
@@ -404,8 +404,6 @@ GpuThunkAotCompilationResult::LoadExecutable(
       ir_emitter->EmitHloComputation(hlo_module->entry_computation()));
   std::unique_ptr<ThunkSequence> thunk_sequence =
       ir_emitter->ConsumeThunkSequence();
-  ForAllThunks([](Thunk* thunk) { thunk->ClearCompileTimeInfo(); },
-               thunk_sequence.get());
 
   // Get all other fields required by GpuExecutable.
   std::vector<GpuExecutable::ConstantInfo> constants =
diff --git a/third_party/xla/xla/service/gpu/runtime/address_computation_thunk.cc b/third_party/xla/xla/service/gpu/runtime/address_computation_thunk.cc
index 26d7d3aacddfe2..2e49de4220e2df 100644
--- a/third_party/xla/xla/service/gpu/runtime/address_computation_thunk.cc
+++ b/third_party/xla/xla/service/gpu/runtime/address_computation_thunk.cc
@@ -54,7 +54,7 @@ AddressComputationThunk::AddressComputationThunk(
     std::vector<std::optional<uint64_t>> offset_byte_sizes)
     : Thunk(Kind::kAddressComputation, thunk_info),
       embedded_thunk_(std::make_unique<SequentialThunk>(
-          ThunkInfo(thunk_info.op), std::move(*embedded_thunk))),
+          ThunkInfo(), std::move(*embedded_thunk))),
       embedded_thunk_arguments_(std::move(arguments)),
       fake_allocations_(std::move(fake_allocations)),
       offset_buffer_indices_(std::move(offset_buffer_indices)),
diff --git a/third_party/xla/xla/service/gpu/runtime/address_computation_thunk_test.cc b/third_party/xla/xla/service/gpu/runtime/address_computation_thunk_test.cc
index 3cbf691d9b2f76..bd51204ab2f480 100644
--- a/third_party/xla/xla/service/gpu/runtime/address_computation_thunk_test.cc
+++ b/third_party/xla/xla/service/gpu/runtime/address_computation_thunk_test.cc
@@ -130,15 +130,14 @@ TEST(AddressComputationThunkTest, SlicedGemm) {
   // Creating embedded GEMM thunk.
   ThunkSequence seq;
   seq.emplace_back(std::make_unique<GemmThunk>(
-      Thunk::ThunkInfo(nullptr), config.value(), slice_lhs_fake, slice_rhs,
-      slice_out, slice_workspace, /*deterministic=*/true));
+      Thunk::ThunkInfo(), config.value(), slice_lhs_fake, slice_rhs, slice_out,
+      slice_workspace, /*deterministic=*/true));
 
   // Wrapping address computation thunk around the GEMM thunk.
   std::vector<BufferAllocation::Slice> lhs_offsets{slice_lhs_offset_0,
                                                    slice_lhs_offset_1};
   AddressComputationThunk thunk(
-      Thunk::ThunkInfo(nullptr),
-      std::make_unique<ThunkSequence>(std::move(seq)),
+      Thunk::ThunkInfo(), std::make_unique<ThunkSequence>(std::move(seq)),
       {slice_lhs, slice_rhs, slice_out, slice_workspace},
       std::move(fake_allocations),
       {lhs_offsets, std::nullopt, std::nullopt, std::nullopt},
@@ -284,7 +283,7 @@ TEST(AddressComputationThunkTest, SlicedNonContiguousGemm) {
   // Creating embedded GEMM thunk.
   ThunkSequence seq;
   seq.emplace_back(std::make_unique<GemmThunk>(
-      Thunk::ThunkInfo(nullptr), config.value(), slice_lhs_fake, slice_rhs_fake,
+      Thunk::ThunkInfo(), config.value(), slice_lhs_fake, slice_rhs_fake,
       slice_out, slice_workspace, /*deterministic=*/true));
 
   // Wrapping address computation thunk around the GEMM thunk.
@@ -293,8 +292,7 @@ TEST(AddressComputationThunkTest, SlicedNonContiguousGemm) {
   std::vector<BufferAllocation::Slice> rhs_offsets{slice_rhs_offset_0,
                                                    slice_rhs_offset_1};
   AddressComputationThunk thunk(
-      Thunk::ThunkInfo(nullptr),
-      std::make_unique<ThunkSequence>(std::move(seq)),
+      Thunk::ThunkInfo(), std::make_unique<ThunkSequence>(std::move(seq)),
       {slice_lhs, slice_rhs, slice_out, slice_workspace},
       std::move(fake_allocations),
       {lhs_offsets, rhs_offsets, std::nullopt, std::nullopt},
@@ -449,7 +447,7 @@ TEST(AddressComputationThunkTest, MulipleSlicedOperandsGemm) {
   // Creating embedded GEMM thunk.
   ThunkSequence seq;
   seq.emplace_back(std::make_unique<GemmThunk>(
-      Thunk::ThunkInfo(nullptr), config.value(), slice_lhs_fake, slice_rhs_fake,
+      Thunk::ThunkInfo(), config.value(), slice_lhs_fake, slice_rhs_fake,
       slice_out, slice_workspace, /*deterministic=*/true));
 
   // Wrapping address computation thunk around the GEMM thunk.
@@ -458,8 +456,7 @@ TEST(AddressComputationThunkTest, MulipleSlicedOperandsGemm) {
   std::vector<BufferAllocation::Slice> rhs_offsets{slice_rhs_offset_0,
                                                    slice_rhs_offset_1};
   AddressComputationThunk thunk(
-      Thunk::ThunkInfo(nullptr),
-      std::make_unique<ThunkSequence>(std::move(seq)),
+      Thunk::ThunkInfo(), std::make_unique<ThunkSequence>(std::move(seq)),
       {slice_lhs, slice_rhs, slice_out, slice_workspace},
       std::move(fake_allocations),
       {lhs_offsets, rhs_offsets, std::nullopt, std::nullopt},
@@ -627,7 +624,7 @@ TEST(AddressComputationThunkTest, SlicedMemcpy) {
   // Creating embedded custom call thunk.
   ThunkSequence seq;
   seq.emplace_back(std::make_unique<CustomCallThunk>(
-      Thunk::ThunkInfo(nullptr), registration->handler, operands, results,
+      Thunk::ThunkInfo(), registration->handler, operands, results,
       /*attributes=*/CustomCallThunk::AttributesMap(),
       /*called_computation=*/nullptr));
 
@@ -635,9 +632,9 @@ TEST(AddressComputationThunkTest, SlicedMemcpy) {
   std::vector<BufferAllocation::Slice> slice_offsets{
       slice_offset_0, slice_offset_1, slice_offset_2, slice_offset_3};
   AddressComputationThunk thunk(
-      Thunk::ThunkInfo(nullptr),
-      std::make_unique<ThunkSequence>(std::move(seq)), {slice_src, slice_dst},
-      std::move(fake_allocations), {slice_offsets, std::nullopt},
+      Thunk::ThunkInfo(), std::make_unique<ThunkSequence>(std::move(seq)),
+      {slice_src, slice_dst}, std::move(fake_allocations),
+      {slice_offsets, std::nullopt},
       {ShapeUtil::MakeShape(PrimitiveType::S32, {8, 8, 10, 8}), std::nullopt},
       // Make sure to pass a dst shape with the same rank as src shape (i.e.
       // original slice result and not bitcasted one)
@@ -785,7 +782,7 @@ TEST(AddressComputationThunkTest, SlicedOutputMemcpy) {
   // Creating embedded custom call thunk.
   ThunkSequence seq;
   seq.emplace_back(std::make_unique<CustomCallThunk>(
-      Thunk::ThunkInfo(nullptr), registration->handler, operands, results,
+      Thunk::ThunkInfo(), registration->handler, operands, results,
       /*attributes=*/CustomCallThunk::AttributesMap(),
       /*called_computation=*/nullptr));
 
@@ -797,9 +794,9 @@ TEST(AddressComputationThunkTest, SlicedOutputMemcpy) {
       slice_dst_offset_0, slice_dst_offset_1, slice_dst_offset_2,
       slice_dst_offset_3};
   AddressComputationThunk thunk(
-      Thunk::ThunkInfo(nullptr),
-      std::make_unique<ThunkSequence>(std::move(seq)), {slice_src, slice_dst},
-      std::move(fake_allocations), {slice_src_offsets, slice_dst_offsets},
+      Thunk::ThunkInfo(), std::make_unique<ThunkSequence>(std::move(seq)),
+      {slice_src, slice_dst}, std::move(fake_allocations),
+      {slice_src_offsets, slice_dst_offsets},
       {ShapeUtil::MakeShape(PrimitiveType::S32, {8, 8, 10, 2}),
        ShapeUtil::MakeShape(PrimitiveType::S32, {2, 2, 2, 2})},
       // Make sure to pass a dst shape with the same rank as src shape (i.e.
@@ -966,15 +963,14 @@ TEST(AddressComputationThunkTest, SlicedGemmArbitraryArgumentOrder) {
   // Creating embedded GEMM thunk.
   ThunkSequence seq;
   seq.emplace_back(std::make_unique<GemmThunk>(
-      Thunk::ThunkInfo(nullptr), config.value(), slice_lhs_fake, slice_rhs_fake,
+      Thunk::ThunkInfo(), config.value(), slice_lhs_fake, slice_rhs_fake,
       slice_out_fake, slice_workspace_fake, /*deterministic=*/true));
 
   // Wrapping address computation thunk around the GEMM thunk.
   std::vector<BufferAllocation::Slice> lhs_offsets{slice_lhs_offset_0,
                                                    slice_lhs_offset_1};
   AddressComputationThunk thunk(
-      Thunk::ThunkInfo(nullptr),
-      std::make_unique<ThunkSequence>(std::move(seq)),
+      Thunk::ThunkInfo(), std::make_unique<ThunkSequence>(std::move(seq)),
       {slice_lhs, slice_rhs, slice_out, slice_workspace},
       std::move(fake_allocations),
       {lhs_offsets, std::nullopt, std::nullopt, std::nullopt},
@@ -1115,15 +1111,14 @@ TEST(AddressComputationThunkTest, SlicedGemmArbitraryNumberOfArguments) {
   // Creating embedded GEMM thunk.
   ThunkSequence seq;
   seq.emplace_back(std::make_unique<GemmThunk>(
-      Thunk::ThunkInfo(nullptr), config.value(), slice_lhs_fake, slice_rhs_fake,
+      Thunk::ThunkInfo(), config.value(), slice_lhs_fake, slice_rhs_fake,
       slice_out_fake, slice_workspace_fake, /*deterministic=*/true));
 
   // Wrapping address computation thunk around the GEMM thunk.
   std::vector<BufferAllocation::Slice> lhs_offsets{slice_lhs_offset_0,
                                                    slice_lhs_offset_1};
   AddressComputationThunk thunk(
-      Thunk::ThunkInfo(nullptr),
-      std::make_unique<ThunkSequence>(std::move(seq)),
+      Thunk::ThunkInfo(), std::make_unique<ThunkSequence>(std::move(seq)),
       {slice_lhs, slice_rhs, slice_out, slice_workspace},
       std::move(fake_allocations),
       {lhs_offsets, std::nullopt, std::nullopt, std::nullopt},
@@ -1257,15 +1252,14 @@ TEST(AddressComputationThunkTest, SlicedTupledOperandGemm) {
   // Creating embedded GEMM thunk.
   ThunkSequence seq;
   seq.emplace_back(std::make_unique<GemmThunk>(
-      Thunk::ThunkInfo(nullptr), config.value(), slice_lhs_fake, slice_rhs,
-      slice_out, slice_workspace, /*deterministic=*/true));
+      Thunk::ThunkInfo(), config.value(), slice_lhs_fake, slice_rhs, slice_out,
+      slice_workspace, /*deterministic=*/true));
 
   // Wrapping address computation thunk around the GEMM thunk.
   std::vector<BufferAllocation::Slice> lhs_offsets{slice_lhs_offset_0,
                                                    slice_lhs_offset_1};
   AddressComputationThunk thunk(
-      Thunk::ThunkInfo(nullptr),
-      std::make_unique<ThunkSequence>(std::move(seq)),
+      Thunk::ThunkInfo(), std::make_unique<ThunkSequence>(std::move(seq)),
       {slice_lhs, slice_rhs, slice_out, slice_workspace},
       std::move(fake_allocations),
       {lhs_offsets, std::nullopt, std::nullopt, std::nullopt},
@@ -1427,7 +1421,7 @@ TEST(AddressComputationThunkTest, SlicedMemcpyOOB) {
   // Creating embedded custom call thunk.
   ThunkSequence seq;
   seq.emplace_back(std::make_unique<CustomCallThunk>(
-      Thunk::ThunkInfo(nullptr), registration->handler, operands, results,
+      Thunk::ThunkInfo(), registration->handler, operands, results,
       /*attributes=*/CustomCallThunk::AttributesMap(),
       /*called_computation=*/nullptr));
 
@@ -1439,9 +1433,9 @@ TEST(AddressComputationThunkTest, SlicedMemcpyOOB) {
       slice_dst_offset_0, slice_dst_offset_1, slice_dst_offset_2,
       slice_dst_offset_3};
   AddressComputationThunk thunk(
-      Thunk::ThunkInfo(nullptr),
-      std::make_unique<ThunkSequence>(std::move(seq)), {slice_src, slice_dst},
-      std::move(fake_allocations), {slice_src_offsets, slice_dst_offsets},
+      Thunk::ThunkInfo(), std::make_unique<ThunkSequence>(std::move(seq)),
+      {slice_src, slice_dst}, std::move(fake_allocations),
+      {slice_src_offsets, slice_dst_offsets},
       {ShapeUtil::MakeShape(PrimitiveType::S32, {8, 8, 10, 2}),
        ShapeUtil::MakeShape(PrimitiveType::S32, {2, 2, 2, 2})},
       // Make sure to pass a dst shape with the same rank as src shape (i.e.
@@ -1609,15 +1603,14 @@ TEST(AddressComputationThunkTest, SlicedOperandsSameBufferGemm) {
   // Creating embedded GEMM thunk.
   ThunkSequence seq;
   seq.emplace_back(std::make_unique<GemmThunk>(
-      Thunk::ThunkInfo(nullptr), config.value(), slice_lhs_fake, slice_rhs_fake,
+      Thunk::ThunkInfo(), config.value(), slice_lhs_fake, slice_rhs_fake,
       slice_out_fake, slice_workspace_fake, /*deterministic=*/true));
 
   // Wrapping address computation thunk around the GEMM thunk.
   std::vector<BufferAllocation::Slice> lhs_offsets{slice_lhs_offset_0,
                                                    slice_lhs_offset_1};
   AddressComputationThunk thunk(
-      Thunk::ThunkInfo(nullptr),
-      std::make_unique<ThunkSequence>(std::move(seq)),
+      Thunk::ThunkInfo(), std::make_unique<ThunkSequence>(std::move(seq)),
       {slice_lhs, slice_rhs, slice_out, slice_workspace},
       std::move(fake_allocations),
       {lhs_offsets, std::nullopt, std::nullopt, std::nullopt},
diff --git a/third_party/xla/xla/service/gpu/runtime/command_buffer_thunk_test.cc b/third_party/xla/xla/service/gpu/runtime/command_buffer_thunk_test.cc
index b1aa3a7b3f3aaa..e8ab612924474b 100644
--- a/third_party/xla/xla/service/gpu/runtime/command_buffer_thunk_test.cc
+++ b/third_party/xla/xla/service/gpu/runtime/command_buffer_thunk_test.cc
@@ -132,7 +132,7 @@ TEST(CommandBufferThunkTest, MemcpyCmd) {
   commands.Emplace<MemcpyDeviceToDeviceCmd>(s0, slice_b, slice_a, byte_length);
 
   // Construct a thunk with command sequence.
-  CommandBufferThunk thunk(std::move(commands), Thunk::ThunkInfo(nullptr));
+  CommandBufferThunk thunk(std::move(commands), Thunk::ThunkInfo());
 
   se::StreamExecutorMemoryAllocator allocator(executor);
   ServiceExecutableRunOptions run_options;
@@ -186,7 +186,7 @@ TEST(CommandBufferThunkTest, MemzeroCmd) {
   commands.Emplace<MemzeroCmd>(s0, slice_a);
 
   // Construct a thunk with command sequence.
-  CommandBufferThunk thunk(std::move(commands), Thunk::ThunkInfo(nullptr));
+  CommandBufferThunk thunk(std::move(commands), Thunk::ThunkInfo());
 
   ServiceExecutableRunOptions run_options;
   se::StreamExecutorMemoryAllocator allocator(executor);
@@ -228,7 +228,7 @@ TEST(CommandBufferThunkTest, Memset32Cmd) {
   commands.Emplace<Memset32Cmd>(s0, slice_a, int32_t{84});
 
   // Construct a thunk with command sequence.
-  CommandBufferThunk thunk(std::move(commands), Thunk::ThunkInfo(nullptr));
+  CommandBufferThunk thunk(std::move(commands), Thunk::ThunkInfo());
 
   ServiceExecutableRunOptions run_options;
   se::StreamExecutorMemoryAllocator allocator(executor);
@@ -267,7 +267,7 @@ TEST(CommandBufferThunkTest, Memset32CmdOnDifferentStreams) {
   commands.Emplace<Memset32Cmd>(s1, slice1, int32_t{34});
 
   // Construct a thunk with command sequence.
-  CommandBufferThunk thunk(std::move(commands), Thunk::ThunkInfo(nullptr));
+  CommandBufferThunk thunk(std::move(commands), Thunk::ThunkInfo());
 
   ServiceExecutableRunOptions run_options;
   se::StreamExecutorMemoryAllocator allocator(executor);
@@ -320,7 +320,7 @@ TEST(CommandBufferThunkTest, LaunchCmd) {
                               /*shmem_bytes=*/0);
 
   // Construct a thunk with command sequence.
-  CommandBufferThunk thunk(std::move(commands), Thunk::ThunkInfo(nullptr));
+  CommandBufferThunk thunk(std::move(commands), Thunk::ThunkInfo());
 
   ServiceExecutableRunOptions run_options;
   se::StreamExecutorMemoryAllocator allocator(executor);
@@ -416,7 +416,7 @@ TEST(CommandBufferThunkTest, CustomAddKernelLaunchCmd) {
                               /*shmem_bytes=*/0);
 
   // Construct a thunk with command sequence.
-  CommandBufferThunk thunk(std::move(commands), Thunk::ThunkInfo(nullptr));
+  CommandBufferThunk thunk(std::move(commands), Thunk::ThunkInfo());
 
   ServiceExecutableRunOptions run_options;
   se::StreamExecutorMemoryAllocator allocator(executor);
@@ -531,7 +531,7 @@ TEST(CommandBufferThunkTest, GemmCmd) {
                             /*deterministic=*/true);
 
   // Construct a thunk with command sequence.
-  CommandBufferThunk thunk(std::move(commands), Thunk::ThunkInfo(nullptr));
+  CommandBufferThunk thunk(std::move(commands), Thunk::ThunkInfo());
 
   ServiceExecutableRunOptions run_options;
   se::StreamExecutorMemoryAllocator allocator(executor);
@@ -631,7 +631,7 @@ TEST(CommandBufferThunkTest, MultipleLaunchCmd) {
                               /*shmem_bytes=*/0);
 
   // Construct a thunk with command sequence.
-  CommandBufferThunk thunk(std::move(commands), Thunk::ThunkInfo(nullptr));
+  CommandBufferThunk thunk(std::move(commands), Thunk::ThunkInfo());
 
   ServiceExecutableRunOptions run_options;
   se::StreamExecutorMemoryAllocator allocator(executor);
@@ -746,7 +746,7 @@ TEST(CommandBufferThunkTest, IfCmd) {
   commands.Emplace<IfCmd>(s0, slice_p, std::move(then_commands));
 
   // Construct a thunk with command sequence.
-  CommandBufferThunk thunk(std::move(commands), Thunk::ThunkInfo(nullptr));
+  CommandBufferThunk thunk(std::move(commands), Thunk::ThunkInfo());
 
   ServiceExecutableRunOptions run_options;
   se::StreamExecutorMemoryAllocator allocator(executor);
@@ -845,7 +845,7 @@ TEST(CommandBufferThunkTest, IfElseCmd) {
                               std::move(else_commands));
 
   // Construct a thunk with command sequence.
-  CommandBufferThunk thunk(std::move(commands), Thunk::ThunkInfo(nullptr));
+  CommandBufferThunk thunk(std::move(commands), Thunk::ThunkInfo());
 
   ServiceExecutableRunOptions run_options;
   se::StreamExecutorMemoryAllocator allocator(executor);
@@ -934,7 +934,7 @@ TEST(CommandBufferThunkTest, CaseCmd) {
   commands.Emplace<CaseCmd>(s0, slice_i, std::move(branches));
 
   // Construct a thunk with command sequence.
-  CommandBufferThunk thunk(std::move(commands), Thunk::ThunkInfo(nullptr));
+  CommandBufferThunk thunk(std::move(commands), Thunk::ThunkInfo());
 
   ServiceExecutableRunOptions run_options;
   se::StreamExecutorMemoryAllocator allocator(executor);
@@ -1013,7 +1013,7 @@ TEST(CommandBufferThunkTest, ForCmd) {
                            std::move(body_commands));
 
   // Construct a thunk with command sequence.
-  CommandBufferThunk thunk(std::move(commands), Thunk::ThunkInfo(nullptr));
+  CommandBufferThunk thunk(std::move(commands), Thunk::ThunkInfo());
 
   ServiceExecutableRunOptions run_options;
   se::StreamExecutorMemoryAllocator allocator(executor);
diff --git a/third_party/xla/xla/service/gpu/runtime/copy_thunk.h b/third_party/xla/xla/service/gpu/runtime/copy_thunk.h
index 4bbd962ed6e91a..0f57cef2ec3e23 100644
--- a/third_party/xla/xla/service/gpu/runtime/copy_thunk.h
+++ b/third_party/xla/xla/service/gpu/runtime/copy_thunk.h
@@ -45,8 +45,6 @@ class DeviceToDeviceCopyThunk : public Thunk {
 
   absl::Status ExecuteOnStream(const ExecuteParams& params) override;
 
-  void ClearCompileTimeInfo() override { Thunk::ClearCompileTimeInfo(); }
-
   const BufferAllocation::Slice& source() const { return source_buffer_; }
   const BufferAllocation::Slice& destination() const {
     return destination_buffer_;
diff --git a/third_party/xla/xla/service/gpu/runtime/kernel_thunk.h b/third_party/xla/xla/service/gpu/runtime/kernel_thunk.h
index 1b50ce6be28694..0f8c36b80a3065 100644
--- a/third_party/xla/xla/service/gpu/runtime/kernel_thunk.h
+++ b/third_party/xla/xla/service/gpu/runtime/kernel_thunk.h
@@ -82,8 +82,6 @@ class KernelThunk : public Thunk {
   absl::Status Initialize(const InitializeParams& params) override;
   absl::Status ExecuteOnStream(const ExecuteParams& params) override;
 
-  void ClearCompileTimeInfo() override { Thunk::ClearCompileTimeInfo(); }
-
   const std::vector<BufferAllocation::Slice>& arguments() const {
     return args_;
   }
diff --git a/third_party/xla/xla/service/gpu/runtime/memset_thunk.h b/third_party/xla/xla/service/gpu/runtime/memset_thunk.h
index e1eef4c39c5d3a..6d7cf63ebbbafb 100644
--- a/third_party/xla/xla/service/gpu/runtime/memset_thunk.h
+++ b/third_party/xla/xla/service/gpu/runtime/memset_thunk.h
@@ -37,8 +37,6 @@ class MemzeroThunk : public Thunk {
 
   absl::Status ExecuteOnStream(const ExecuteParams& params) override;
 
-  void ClearCompileTimeInfo() override { Thunk::ClearCompileTimeInfo(); }
-
   const BufferAllocation::Slice& destination() const { return dest_; }
 
  private:
@@ -57,8 +55,6 @@ class Memset32BitValueThunk : public Thunk {
 
   absl::Status ExecuteOnStream(const ExecuteParams& params) override;
 
-  void ClearCompileTimeInfo() override { Thunk::ClearCompileTimeInfo(); }
-
   const BufferAllocation::Slice& destination() const { return dest_; }
   uint32_t value() const { return value_; }
 
diff --git a/third_party/xla/xla/service/gpu/runtime/thunk.cc b/third_party/xla/xla/service/gpu/runtime/thunk.cc
index f066aed07aaa94..0c95623db733be 100644
--- a/third_party/xla/xla/service/gpu/runtime/thunk.cc
+++ b/third_party/xla/xla/service/gpu/runtime/thunk.cc
@@ -31,7 +31,6 @@ limitations under the License.
 #include "absl/strings/str_cat.h"
 #include "absl/strings/string_view.h"
 #include "absl/types/span.h"
-#include "mlir/IR/Operation.h"  // from @llvm-project
 #include "xla/executable_run_options.h"
 #include "xla/ffi/execution_context.h"
 #include "xla/hlo/ir/hlo_instruction.h"
@@ -340,16 +339,9 @@ bool IsReductionCollective(Thunk::Kind kind) {
          kind == Thunk::kNcclReduceScatterStart;
 }
 
-Thunk::ThunkInfo Thunk::ThunkInfo::WithProfileAnnotation(mlir::Operation* op) {
-  ThunkInfo thunk_info(op);
-  thunk_info.profile_annotation =
-      mlir::mhlo::GetDebugNameFromLocation(op->getLoc());
-  return thunk_info;
-}
-
 Thunk::ThunkInfo Thunk::ThunkInfo::WithProfileAnnotation(
     const HloInstruction* instr) {
-  ThunkInfo thunk_info(nullptr);
+  ThunkInfo thunk_info;
   thunk_info.profile_annotation = instr->name();
   auto gpu_backend_config = instr->backend_config<GpuBackendConfig>();
   if (gpu_backend_config.ok()) {
diff --git a/third_party/xla/xla/service/gpu/runtime/thunk.h b/third_party/xla/xla/service/gpu/runtime/thunk.h
index 7acd446896fe91..a0c3601f2db0f4 100644
--- a/third_party/xla/xla/service/gpu/runtime/thunk.h
+++ b/third_party/xla/xla/service/gpu/runtime/thunk.h
@@ -183,14 +183,10 @@ class Thunk {
   };
 
   struct ThunkInfo {
-    explicit ThunkInfo(mlir::Operation* op) : op(op) {}
-    static ThunkInfo WithProfileAnnotation(mlir::Operation* op);
+    ThunkInfo() = default;  // Disable implicit constructors.
     static ThunkInfo WithProfileAnnotation(const HloInstruction* instr);
 
     std::string profile_annotation;
-    // TODO(b/304613751): This is only needed by the LMHLO. Remove this when
-    // LMHLO is removed from the runtime pipeline.
-    mlir::Operation* op;
 
     ExecutionStreamId execution_stream_id = kDefaultExecutionStreamId;
   };
@@ -409,7 +405,6 @@ class Thunk {
   Thunk(Kind kind, ThunkInfo thunk_info)
       : kind_(kind),
         profile_annotation_(thunk_info.profile_annotation),
-        op_(thunk_info.op),
         execution_stream_id_(thunk_info.execution_stream_id) {}
   virtual ~Thunk() = default;
   Thunk(const Thunk&) = delete;
@@ -419,11 +414,6 @@ class Thunk {
   Kind kind() const { return kind_; }
   std::string_view profile_annotation() const { return profile_annotation_; }
 
-  // Only valid during compilation, i.e., lowering thunks to kernel-launch
-  // related XLA runtime custom calls). nullptr at runtime. MLIR codegen will
-  // cease the practice of lowering thunks to XLA runtime custom calls.
-  mlir::Operation* op() { return op_; }
-
   // Prepares thunk for execution.
   //
   // This may be called multiple times. Its main purpose is to pass resource
@@ -451,9 +441,6 @@ class Thunk {
   // Precondition: Initialize(initialize_params) has been called.
   virtual absl::Status ExecuteOnStream(const ExecuteParams& params) = 0;
 
-  // Clears metadata that is only valid during compile time.
-  virtual void ClearCompileTimeInfo() { op_ = nullptr; }
-
   static absl::string_view KindToString(Thunk::Kind kind);
 
   ExecutionStreamId execution_stream_id() const { return execution_stream_id_; }
@@ -464,7 +451,6 @@ class Thunk {
  private:
   Kind kind_;
   std::string profile_annotation_;
-  mlir::Operation* op_;
   ExecutionStreamId execution_stream_id_;
 };
 
diff --git a/third_party/xla/xla/service/gpu/runtime/while_thunk.cc b/third_party/xla/xla/service/gpu/runtime/while_thunk.cc
index cbd9ce3134133e..56c312252da203 100644
--- a/third_party/xla/xla/service/gpu/runtime/while_thunk.cc
+++ b/third_party/xla/xla/service/gpu/runtime/while_thunk.cc
@@ -44,9 +44,9 @@ WhileThunk::WhileThunk(
     : Thunk(Kind::kWhile, thunk_info),
       condition_result_buffer_index_(condition_result_buffer_index),
       condition_thunk_sequence_(std::make_unique<SequentialThunk>(
-          ThunkInfo(thunk_info.op), std::move(*condition_thunk_sequence))),
+          ThunkInfo(), std::move(*condition_thunk_sequence))),
       body_thunk_sequence_(std::make_unique<SequentialThunk>(
-          ThunkInfo(thunk_info.op), std::move(*body_thunk_sequence))),
+          ThunkInfo(), std::move(*body_thunk_sequence))),
       trip_count_(trip_count) {}
 
 absl::Status WhileThunk::Prepare(const PrepareParams& params,

From da07e34ae4f3adbd3df88b956a9b14ab3141c350 Mon Sep 17 00:00:00 2001
From: Sara Smoot <sarasmoot@google.com>
Date: Fri, 17 May 2024 14:12:34 -0700
Subject: [PATCH 371/478] [XLA:CPU] Remove unnecessary include of
 absl/log/check.h.

This resolves naming conflicts with tsl and absl.

PiperOrigin-RevId: 634883049
---
 third_party/xla/xla/service/cpu/BUILD             |  2 --
 third_party/xla/xla/service/cpu/onednn_softmax.cc | 10 ----------
 2 files changed, 12 deletions(-)

diff --git a/third_party/xla/xla/service/cpu/BUILD b/third_party/xla/xla/service/cpu/BUILD
index 24dc3d2d833696..0918e2d6c03e49 100644
--- a/third_party/xla/xla/service/cpu/BUILD
+++ b/third_party/xla/xla/service/cpu/BUILD
@@ -1193,7 +1193,6 @@ cc_library(
         "@com_google_absl//absl/algorithm:container",
         "@com_google_absl//absl/base:core_headers",
         "@com_google_absl//absl/base:dynamic_annotations",
-        "@com_google_absl//absl/log",
         "@com_google_absl//absl/status",
         "@com_google_absl//absl/strings",
         "@com_google_absl//absl/types:span",
@@ -1652,7 +1651,6 @@ cc_library(
         "//xla:executable_run_options",
         "@com_google_absl//absl/base:core_headers",
         "@com_google_absl//absl/base:dynamic_annotations",
-        "@com_google_absl//absl/log:check",
         "@eigen_archive//:eigen3",
         "@local_tsl//tsl/platform:blocking_counter",
         "@local_tsl//tsl/platform:env",
diff --git a/third_party/xla/xla/service/cpu/onednn_softmax.cc b/third_party/xla/xla/service/cpu/onednn_softmax.cc
index 5af6de54078596..d2a55adb2d10ac 100644
--- a/third_party/xla/xla/service/cpu/onednn_softmax.cc
+++ b/third_party/xla/xla/service/cpu/onednn_softmax.cc
@@ -20,16 +20,6 @@ limitations under the License.
 #include <initializer_list>
 #include <vector>
 
-// Both "absl/log/check.h" and "third_party/tsl/platform/logging.h"
-// are transitively included in bazel. Both of them define similar CHECK macros.
-// Explicitly including the Abseil header first because the TSL version has
-// undefs.
-
-// Otherwise, we would get redefinition error.
-// clang-format off
-#include "absl/log/check.h"
-// clang-format on
-
 #include "dnnl.hpp"
 #include "absl/base/dynamic_annotations.h"
 #include "xla/executable_run_options.h"

From df8d64ea75d10861fba0a80d04de7aef4e682e4f Mon Sep 17 00:00:00 2001
From: "Jae H. Yoo" <jaeyoo@google.com>
Date: Fri, 17 May 2024 14:23:09 -0700
Subject: [PATCH 372/478] Fix Flatbuffer's upstream error on
 `GetTemporaryPointer()`

PiperOrigin-RevId: 634885829
---
 RELEASE.md                                            | 2 ++
 tensorflow/lite/tools/cmake/modules/flatbuffers.cmake | 5 +++--
 third_party/flatbuffers/workspace.bzl                 | 4 ++--
 3 files changed, 7 insertions(+), 4 deletions(-)

diff --git a/RELEASE.md b/RELEASE.md
index 2e59555b7610b0..0b33b6b9afef4e 100644
--- a/RELEASE.md
+++ b/RELEASE.md
@@ -87,6 +87,8 @@
     * The Python TF Lite Interpreter bindings now have an option
       `experimental_default_delegate_latest_features` to enable all default
       delegate features.
+    * Flatbuffer version update:
+        * `GetTemporaryPointer()` bug fixed.
 
 * `tf.data`
     * Add `wait` to `tf.data.Dataset.load`. If `True`, for snapshots written
diff --git a/tensorflow/lite/tools/cmake/modules/flatbuffers.cmake b/tensorflow/lite/tools/cmake/modules/flatbuffers.cmake
index a768574923932a..8febb1fbd1c706 100644
--- a/tensorflow/lite/tools/cmake/modules/flatbuffers.cmake
+++ b/tensorflow/lite/tools/cmake/modules/flatbuffers.cmake
@@ -23,8 +23,9 @@ OverridableFetchContent_Declare(
   flatbuffers
   GIT_REPOSITORY https://github.com/google/flatbuffers
   # Sync with tensorflow/third_party/flatbuffers/workspace.bzl
-  GIT_TAG v24.3.25
-  GIT_SHALLOW TRUE
+  GIT_TAG e6463926479bd6b330cbcf673f7e917803fd5831
+  # NOTE: b/340264458 - `GIT_SHALLOW TRUE` works for tag name only.
+  GIT_SHALLOW FALSE
   GIT_PROGRESS TRUE
   SOURCE_DIR "${CMAKE_BINARY_DIR}/flatbuffers"
 )
diff --git a/third_party/flatbuffers/workspace.bzl b/third_party/flatbuffers/workspace.bzl
index 129880229e7dcb..b564d0331abe15 100644
--- a/third_party/flatbuffers/workspace.bzl
+++ b/third_party/flatbuffers/workspace.bzl
@@ -5,10 +5,10 @@ load("//third_party:repo.bzl", "tf_http_archive", "tf_mirror_urls")
 # _FLATBUFFERS_GIT_COMMIT / _FLATBUFFERS_SHA256 were added due to an urgent change being made to
 # Flatbuffers that needed to be updated in order for Flatbuffers/TfLite be compatible with Android
 # API level >= 23. They can be removed next flatbuffers offical release / update.
-_FLATBUFFERS_GIT_COMMIT = "595bf0007ab1929570c7671f091313c8fc20644e"
+_FLATBUFFERS_GIT_COMMIT = "e6463926479bd6b330cbcf673f7e917803fd5831"
 
 # curl -L https://github.com/google/flatbuffers/archive/<_FLATBUFFERS_GIT_COMMIT>.tar.gz | shasum -a 256
-_FLATBUFFERS_SHA256 = "987300083ec1f1b095d5596ef8fb657ba46c45d786bc866a5e9029d7590a5e48"
+_FLATBUFFERS_SHA256 = "c9c6b8653597ed7ee5c62243979010bd0f09b29a46be414505bc5b58a874bb17"
 
 def repo():
     tf_http_archive(

From 672697d6121301c61afa12a8cae6adfa1735fe5b Mon Sep 17 00:00:00 2001
From: Fangrui Song <maskray@google.com>
Date: Fri, 17 May 2024 15:19:57 -0700
Subject: [PATCH 373/478] Integrate LLVM at llvm/llvm-project@1e5f29af81a5

Updates LLVM usage to match
[1e5f29af81a5](https://github.com/llvm/llvm-project/commit/1e5f29af81a5)

PiperOrigin-RevId: 634900249
---
 third_party/llvm/workspace.bzl | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/third_party/llvm/workspace.bzl b/third_party/llvm/workspace.bzl
index 1893eaa66f2c53..e20b531ed2fe3b 100644
--- a/third_party/llvm/workspace.bzl
+++ b/third_party/llvm/workspace.bzl
@@ -4,8 +4,8 @@ load("//third_party:repo.bzl", "tf_http_archive")
 
 def repo(name):
     """Imports LLVM."""
-    LLVM_COMMIT = "371eccd5dfed88c8e76449233d8388c12be3464b"
-    LLVM_SHA256 = "31e51b40d2890d6389dc818ab45bf1395b239b9e1d5b280b58f63e260b6deae8"
+    LLVM_COMMIT = "1e5f29af81a5f6fda308074f6345b9fba4faa71c"
+    LLVM_SHA256 = "f31b9a08cb74cffd89a0bc87b6a7ef327b54bc67af25e1503683d40817b22f2a"
 
     tf_http_archive(
         name = name,

From 7325fefecdc5873a509198f0fc8462b10e5f2028 Mon Sep 17 00:00:00 2001
From: Kyle Lucke <klucke@google.com>
Date: Fri, 17 May 2024 15:26:16 -0700
Subject: [PATCH 374/478] Use StreamExecutorInterface::CreateEvent to create
 events in send & recv thunks.

PiperOrigin-RevId: 634901606
---
 tensorflow/compiler/jit/kernels/BUILD         |  2 ++
 tensorflow/compiler/jit/kernels/xla_ops.cc    | 28 ++++++++---------
 .../jit/xla_host_recv_device_context.cc       |  2 +-
 .../jit/xla_host_recv_device_context.h        | 14 ++++-----
 .../jit/xla_host_send_device_context.cc       |  2 +-
 .../jit/xla_host_send_device_context.h        | 14 ++++-----
 .../xla_host_send_recv_device_context_test.cc | 28 +++++++++--------
 third_party/xla/xla/executable_run_options.h  |  8 ++---
 .../xla/pjrt/pjrt_stream_executor_client.cc   | 30 +++++++++----------
 .../service/gpu/runtime/send_recv_thunk.cc    | 18 +++++------
 .../xla/service/gpu/runtime/send_recv_thunk.h |  8 ++---
 11 files changed, 77 insertions(+), 77 deletions(-)

diff --git a/tensorflow/compiler/jit/kernels/BUILD b/tensorflow/compiler/jit/kernels/BUILD
index 0ac326c61fb3ec..d173564b7fd10d 100644
--- a/tensorflow/compiler/jit/kernels/BUILD
+++ b/tensorflow/compiler/jit/kernels/BUILD
@@ -59,8 +59,10 @@ cc_library(
         "//tensorflow/compiler/jit:xla_compile_util",
         "//tensorflow/core/platform:refcount",
         "@com_google_absl//absl/status",
+        "@com_google_absl//absl/status:statusor",
         "@com_google_absl//absl/strings",
         "@local_xla//xla/pjrt:pjrt_client",
+        "@local_xla//xla/tsl/concurrency:async_value",
     ],
     alwayslink = 1,
 )
diff --git a/tensorflow/compiler/jit/kernels/xla_ops.cc b/tensorflow/compiler/jit/kernels/xla_ops.cc
index 9d75388cfbbe80..5a29e8ef36e9b3 100644
--- a/tensorflow/compiler/jit/kernels/xla_ops.cc
+++ b/tensorflow/compiler/jit/kernels/xla_ops.cc
@@ -24,10 +24,10 @@ limitations under the License.
 #include <string_view>
 #include <tuple>
 #include <utility>
-#include <variant>
 
 #include "absl/container/flat_hash_map.h"
 #include "absl/status/status.h"
+#include "absl/status/statusor.h"
 #include "absl/strings/str_cat.h"
 #include "absl/types/optional.h"
 #include "tensorflow/compiler/jit/device_compilation_profiler.h"
@@ -52,7 +52,7 @@ limitations under the License.
 #include "xla/executable_run_options.h"
 #include "xla/pjrt/pjrt_client.h"
 #include "xla/service/gpu/gpu_executable_run_options.h"
-#include "xla/statusor.h"
+#include "xla/tsl/concurrency/async_value_ref.h"
 #include "tensorflow/core/framework/allocator.h"
 #include "tensorflow/core/framework/node_def_util.h"
 #include "tensorflow/core/framework/op_kernel.h"
@@ -224,7 +224,7 @@ xla::SendDeviceMemoryFunction GetSendDeviceMemoryFunction(
           int64_t channel_id, se::Stream* stream, const xla::Shape& shape,
           const se::DeviceMemoryBase& device_memory_base,
           const absl::flat_hash_map<std::string, std::string>& frontend_attrs)
-          -> absl::StatusOr<tsl::AsyncValueRef<se::Event>> {
+          -> absl::StatusOr<tsl::AsyncValueRef<std::unique_ptr<se::Event>>> {
         auto iter = frontend_attrs.find("_xla_host_transfer_rendezvous");
 
         // Generate the Rendezvous key.
@@ -244,12 +244,10 @@ xla::SendDeviceMemoryFunction GetSendDeviceMemoryFunction(
         RendezvousInterface::ParsedKey parsed_key;
         TF_RETURN_IF_ERROR(Rendezvous::ParseKey(rendezvous_key, &parsed_key));
 
-        tsl::AsyncValueRef<se::Event> done_event =
-            tsl::MakeConstructedAsyncValueRef<se::Event>(stream->parent());
-        if (!done_event->Init()) {
-          return errors::Internal(
-              "Failed to initialize done event (channel_id=%d)", channel_id);
-        }
+        TF_ASSIGN_OR_RETURN(auto event, stream->parent()->CreateEvent());
+        tsl::AsyncValueRef<std::unique_ptr<se::Event>> done_event =
+            tsl::MakeConstructedAsyncValueRef<std::unique_ptr<se::Event>>(
+                std::move(event));
 
         Rendezvous::Args args;
         // Rendezvous::Args owns the device context pointer.
@@ -273,7 +271,7 @@ xla::RecvDeviceMemoryFunction GetRecvDeviceMemoryFunction(
           int64_t channel_id, se::Stream* stream, const xla::Shape& shape,
           se::DeviceMemoryBase* device_memory_base,
           const absl::flat_hash_map<std::string, std::string>& frontend_attrs)
-          -> absl::StatusOr<tsl::AsyncValueRef<se::Event>> {
+          -> absl::StatusOr<tsl::AsyncValueRef<std::unique_ptr<se::Event>>> {
         auto iter = frontend_attrs.find("_xla_host_transfer_rendezvous");
 
         // Generate the Rendezvous key.
@@ -293,12 +291,10 @@ xla::RecvDeviceMemoryFunction GetRecvDeviceMemoryFunction(
         RendezvousInterface::ParsedKey parsed_key;
         TF_RETURN_IF_ERROR(Rendezvous::ParseKey(rendezvous_key, &parsed_key));
 
-        tsl::AsyncValueRef<se::Event> done_event =
-            tsl::MakeConstructedAsyncValueRef<se::Event>(stream->parent());
-        if (!done_event->Init()) {
-          return errors::Internal(
-              "Failed to initialize done event (channel_id=%d)", channel_id);
-        }
+        TF_ASSIGN_OR_RETURN(auto event, stream->parent()->CreateEvent());
+        tsl::AsyncValueRef<std::unique_ptr<se::Event>> done_event =
+            tsl::MakeConstructedAsyncValueRef<std::unique_ptr<se::Event>>(
+                std::move(event));
 
         Rendezvous::Args args;
         // Rendezvous::Args owns the device context pointer.
diff --git a/tensorflow/compiler/jit/xla_host_recv_device_context.cc b/tensorflow/compiler/jit/xla_host_recv_device_context.cc
index 54f22fe59fa0bf..ae3c149d5d1387 100644
--- a/tensorflow/compiler/jit/xla_host_recv_device_context.cc
+++ b/tensorflow/compiler/jit/xla_host_recv_device_context.cc
@@ -38,7 +38,7 @@ void XlaHostRecvDeviceContext::CopyDeviceTensorToCPU(
     done(status);
     return;
   }
-  status = stream_->RecordEvent(&done_event_.get());
+  status = stream_->RecordEvent(done_event_.get().get());
   if (!status.ok()) {
     done(status);
     return;
diff --git a/tensorflow/compiler/jit/xla_host_recv_device_context.h b/tensorflow/compiler/jit/xla_host_recv_device_context.h
index 8938fd9c9e0c17..028fd4efd68091 100644
--- a/tensorflow/compiler/jit/xla_host_recv_device_context.h
+++ b/tensorflow/compiler/jit/xla_host_recv_device_context.h
@@ -36,8 +36,8 @@ namespace tensorflow {
 //  Tensor device_tensor(device_allocator, DT_FLOAT, TensorShape({2, 2}));
 //  se::DeviceMemoryBase gpu_dst{device_tensor.data(), 4 * sizeof(float)};
 //  xla::Shape shape(xla::F32, {2, 2}, {}, {})
-//  tsl::AsyncValueRef<se::Event> done_event =
-//      tsl::MakeConstructedAsyncValueRef<se::Event>(stream.parent());
+//  tsl::AsyncValueRef<std::unique_ptr<se::Event>> done_event =
+//      tsl::MakeConstructedAsyncValueRef<std::unique_ptr<se::Event>>(stream.parent());
 //  done_event->Init();
 //  Tensor dest_cpu_tensor;
 //
@@ -48,10 +48,10 @@ namespace tensorflow {
 
 class XlaHostRecvDeviceContext : public DeviceContext {
  public:
-  XlaHostRecvDeviceContext(se::Stream* stream,
-                           const se::DeviceMemoryBase& device_memory_base,
-                           const xla::Shape& shape,
-                           tsl::AsyncValueRef<se::Event>& done_event)
+  XlaHostRecvDeviceContext(
+      se::Stream* stream, const se::DeviceMemoryBase& device_memory_base,
+      const xla::Shape& shape,
+      tsl::AsyncValueRef<std::unique_ptr<se::Event>>& done_event)
       : stream_(stream),
         device_memory_base_(device_memory_base),
         shape_(shape),
@@ -82,7 +82,7 @@ class XlaHostRecvDeviceContext : public DeviceContext {
   // not an issue here since only DeviceMemoryBase methods/members are used.
   const se::DeviceMemoryBase device_memory_base_;
   const xla::Shape shape_;
-  tsl::AsyncValueRef<se::Event> done_event_;
+  tsl::AsyncValueRef<std::unique_ptr<se::Event>> done_event_;
 
   XlaHostRecvDeviceContext(const XlaHostRecvDeviceContext&) = delete;
   void operator=(const XlaHostRecvDeviceContext&) = delete;
diff --git a/tensorflow/compiler/jit/xla_host_send_device_context.cc b/tensorflow/compiler/jit/xla_host_send_device_context.cc
index 5d106c8dc3e073..3d1a9a9f5228c6 100644
--- a/tensorflow/compiler/jit/xla_host_send_device_context.cc
+++ b/tensorflow/compiler/jit/xla_host_send_device_context.cc
@@ -28,7 +28,7 @@ void XlaHostSendDeviceContext::CopyCPUTensorToDevice(
     done(status);
     return;
   }
-  status = stream_->RecordEvent(&done_event_.get());
+  status = stream_->RecordEvent(done_event_.get().get());
   if (!status.ok()) {
     done(status);
     return;
diff --git a/tensorflow/compiler/jit/xla_host_send_device_context.h b/tensorflow/compiler/jit/xla_host_send_device_context.h
index d7a254770c969e..f4e4e9a2535341 100644
--- a/tensorflow/compiler/jit/xla_host_send_device_context.h
+++ b/tensorflow/compiler/jit/xla_host_send_device_context.h
@@ -37,8 +37,8 @@ namespace tensorflow {
 //  Tensor device_tensor(device_allocator, DT_FLOAT, TensorShape({2, 2}));
 //  se::DeviceMemoryBase gpu_dst{device_tensor.data(), 4 * sizeof(float)};
 //  xla::Shape shape(xla::F32, {2, 2}, {}, {})
-//  tsl::AsyncValueRef<se::Event> done_event =
-//      tsl::MakeConstructedAsyncValueRef<se::Event>(stream.parent());
+//  tsl::AsyncValueRef<std::unique_ptr<se::Event>> done_event =
+//      tsl::MakeConstructedAsyncValueRef<std::unique_ptr<se::Event>>(stream.parent());
 //  done_event->Init();
 //
 //  XlaHostSendDeviceContext device_context(&stream, &gpu_dst,
@@ -48,10 +48,10 @@ namespace tensorflow {
 
 class XlaHostSendDeviceContext : public DeviceContext {
  public:
-  XlaHostSendDeviceContext(se::Stream* stream,
-                           se::DeviceMemoryBase* device_memory_base,
-                           const xla::Shape& shape,
-                           tsl::AsyncValueRef<se::Event>& done_event)
+  XlaHostSendDeviceContext(
+      se::Stream* stream, se::DeviceMemoryBase* device_memory_base,
+      const xla::Shape& shape,
+      tsl::AsyncValueRef<std::unique_ptr<se::Event>>& done_event)
       : stream_(stream),
         device_memory_base_(device_memory_base),
         shape_(shape),
@@ -79,7 +79,7 @@ class XlaHostSendDeviceContext : public DeviceContext {
   se::Stream* stream_;                        // Not owned.
   se::DeviceMemoryBase* device_memory_base_;  // Not owned.
   const xla::Shape shape_;
-  tsl::AsyncValueRef<se::Event> done_event_;
+  tsl::AsyncValueRef<std::unique_ptr<se::Event>> done_event_;
 
   XlaHostSendDeviceContext(const XlaHostSendDeviceContext&) = delete;
   void operator=(const XlaHostSendDeviceContext&) = delete;
diff --git a/tensorflow/compiler/jit/xla_host_send_recv_device_context_test.cc b/tensorflow/compiler/jit/xla_host_send_recv_device_context_test.cc
index 16f42d1dbe1a0d..62da04c3e7510f 100644
--- a/tensorflow/compiler/jit/xla_host_send_recv_device_context_test.cc
+++ b/tensorflow/compiler/jit/xla_host_send_recv_device_context_test.cc
@@ -79,9 +79,10 @@ TEST_F(XlaHostSendRecvDeviceContextTest, CopyDeviceTensorToCPU) {
       stream->Memcpy(&gpu_dst, origin_cpu_tensor.data(), gpu_dst.size()));
   TF_ASSERT_OK(stream->BlockHostUntilDone());
 
-  tsl::AsyncValueRef<se::Event> done_event =
-      tsl::MakeConstructedAsyncValueRef<se::Event>(executor);
-  done_event->Init();
+  TF_ASSERT_OK_AND_ASSIGN(auto se_event, executor->CreateEvent());
+  tsl::AsyncValueRef<std::unique_ptr<se::Event>> done_event =
+      tsl::MakeConstructedAsyncValueRef<std::unique_ptr<se::Event>>(
+          std::move(se_event));
   XlaHostRecvDeviceContext* device_context =
       new XlaHostRecvDeviceContext(stream.get(), gpu_dst, shape, done_event);
   TF_ASSERT_OK(device_context->CopyDeviceTensorToCPUSync(
@@ -108,9 +109,10 @@ TEST_F(XlaHostSendRecvDeviceContextTest, CopyCPUTensorToDevice) {
   xla::Shape shape;
   TF_ASSERT_OK(TensorShapeToXLAShape(DT_FLOAT, TensorShape({2, 2}), &shape));
 
-  tsl::AsyncValueRef<se::Event> done_event =
-      tsl::MakeConstructedAsyncValueRef<se::Event>(executor);
-  done_event->Init();
+  TF_ASSERT_OK_AND_ASSIGN(auto se_event, executor->CreateEvent());
+  tsl::AsyncValueRef<std::unique_ptr<se::Event>> done_event =
+      tsl::MakeConstructedAsyncValueRef<std::unique_ptr<se::Event>>(
+          std::move(se_event));
   XlaHostSendDeviceContext* device_context =
       new XlaHostSendDeviceContext(stream.get(), &gpu_dst, shape, done_event);
   TF_ASSERT_OK(device_context->CopyCPUTensorToDeviceSync(
@@ -141,17 +143,19 @@ TEST_F(XlaHostSendRecvDeviceContextTest, RoundTrip) {
   xla::Shape shape;
   TF_ASSERT_OK(TensorShapeToXLAShape(DT_FLOAT, TensorShape({2, 2}), &shape));
 
-  tsl::AsyncValueRef<se::Event> send_done_event =
-      tsl::MakeConstructedAsyncValueRef<se::Event>(executor);
-  send_done_event->Init();
+  TF_ASSERT_OK_AND_ASSIGN(auto se_event, executor->CreateEvent());
+  tsl::AsyncValueRef<std::unique_ptr<se::Event>> send_done_event =
+      tsl::MakeConstructedAsyncValueRef<std::unique_ptr<se::Event>>(
+          std::move(se_event));
   XlaHostSendDeviceContext* send_device_context = new XlaHostSendDeviceContext(
       stream.get(), &gpu_dst, shape, send_done_event);
   TF_ASSERT_OK(send_device_context->CopyCPUTensorToDeviceSync(
       &origin_cpu_tensor, device_.get(), &device_tensor));
 
-  tsl::AsyncValueRef<se::Event> recv_done_event =
-      tsl::MakeConstructedAsyncValueRef<se::Event>(executor);
-  recv_done_event->Init();
+  TF_ASSERT_OK_AND_ASSIGN(auto recv_se_event, executor->CreateEvent());
+  tsl::AsyncValueRef<std::unique_ptr<se::Event>> recv_done_event =
+      tsl::MakeConstructedAsyncValueRef<std::unique_ptr<se::Event>>(
+          std::move(recv_se_event));
   XlaHostRecvDeviceContext* recv_device_context = new XlaHostRecvDeviceContext(
       stream.get(), gpu_dst, shape, recv_done_event);
   TF_ASSERT_OK(recv_device_context->CopyDeviceTensorToCPUSync(
diff --git a/third_party/xla/xla/executable_run_options.h b/third_party/xla/xla/executable_run_options.h
index f02e5e73d68c3a..c6a4897c2067ec 100644
--- a/third_party/xla/xla/executable_run_options.h
+++ b/third_party/xla/xla/executable_run_options.h
@@ -103,8 +103,8 @@ using ThenExecuteFunction =
 // recorded on a `stream` once the send operation is completed and data was
 // copied from the `src` memory. `frontend_attrs` contains frontend specific
 // attributes for the send.
-using SendDeviceMemoryFunction =
-    std::function<absl::StatusOr<tsl::AsyncValueRef<stream_executor::Event>>(
+using SendDeviceMemoryFunction = std::function<
+    absl::StatusOr<tsl::AsyncValueRef<std::unique_ptr<stream_executor::Event>>>(
         int64_t channel_id, stream_executor::Stream* stream, const Shape& shape,
         const stream_executor::DeviceMemoryBase& src,
         const absl::flat_hash_map<std::string, std::string>& frontend_attrs)>;
@@ -113,8 +113,8 @@ using SendDeviceMemoryFunction =
 // recorded on a `stream` once the recv operation is completed and data was
 // copied into the `dst` memory. `frontend_attrs` contains frontend specific
 // attributes for the receive.
-using RecvDeviceMemoryFunction =
-    std::function<absl::StatusOr<tsl::AsyncValueRef<stream_executor::Event>>(
+using RecvDeviceMemoryFunction = std::function<
+    absl::StatusOr<tsl::AsyncValueRef<std::unique_ptr<stream_executor::Event>>>(
         int64_t channel_id, stream_executor::Stream* stream, const Shape& shape,
         stream_executor::DeviceMemoryBase* dst,
         const absl::flat_hash_map<std::string, std::string>& frontend_attrs)>;
diff --git a/third_party/xla/xla/pjrt/pjrt_stream_executor_client.cc b/third_party/xla/xla/pjrt/pjrt_stream_executor_client.cc
index b0e11d18f260a7..3f2ab8a28e0a54 100644
--- a/third_party/xla/xla/pjrt/pjrt_stream_executor_client.cc
+++ b/third_party/xla/xla/pjrt/pjrt_stream_executor_client.cc
@@ -2427,7 +2427,7 @@ static SendDeviceMemoryFunction ConvertSendCallbacksToSendFunction(
              int64_t channel_id, se::Stream* stream, const Shape& shape,
              const se::DeviceMemoryBase& src,
              const absl::flat_hash_map<std::string, std::string>&)
-             -> absl::StatusOr<AsyncValueRef<se::Event>> {
+             -> absl::StatusOr<AsyncValueRef<std::unique_ptr<se::Event>>> {
     VLOG(3) << "Send " << src.size() << " bytes to channel #" << channel_id
             << " (shape=" << shape.ToString() << ")";
 
@@ -2441,10 +2441,9 @@ static SendDeviceMemoryFunction ConvertSendCallbacksToSendFunction(
     // Allocate event that will signal completion of send operation. We do not
     // actually track the completion of the send callback, we only have to keep
     // the device memory long enough to complete the memcpy command.
-    auto done_event = MakeConstructedAsyncValueRef<se::Event>(stream->parent());
-    if (!done_event->Init())
-      return Internal("Failed to initialize done event (channel_id=%d)",
-                      channel_id);
+    TF_ASSIGN_OR_RETURN(auto se_event, stream->parent()->CreateEvent());
+    auto done_event = MakeConstructedAsyncValueRef<std::unique_ptr<se::Event>>(
+        std::move(se_event));
 
     thread_pool->Schedule([done_event, stream, src, channel_id, shape, send] {
       tsl::profiler::TraceMe trace([&] {
@@ -2461,7 +2460,7 @@ static SendDeviceMemoryFunction ConvertSendCallbacksToSendFunction(
         done_event.SetError(status);
         return;
       }
-      status = stream->RecordEvent(&done_event.get());
+      status = stream->RecordEvent(done_event.get().get());
       if (!status.ok()) {
         done_event.SetError(status);
         return;
@@ -2494,9 +2493,9 @@ static SendDeviceMemoryFunction ConvertSendCallbacksToSendFunction(
 namespace {
 class StreamExecutorCopyToDeviceStream : public CopyToDeviceStream {
  public:
-  StreamExecutorCopyToDeviceStream(int64_t channel_id, se::Stream* stream,
-                                   se::DeviceMemoryBase dst,
-                                   AsyncValueRef<se::Event> done)
+  StreamExecutorCopyToDeviceStream(
+      int64_t channel_id, se::Stream* stream, se::DeviceMemoryBase dst,
+      AsyncValueRef<std::unique_ptr<se::Event>> done)
       : CopyToDeviceStream(dst.size(), /*granule_bytes=*/1),
         channel_id_(channel_id),
         stream_(stream),
@@ -2558,7 +2557,7 @@ class StreamExecutorCopyToDeviceStream : public CopyToDeviceStream {
     // responsibility to synchronize with this event before submitting any new
     // computations to the stream.
     if (complete) {
-      auto recorded = stream_->RecordEvent(&done_.get());
+      auto recorded = stream_->RecordEvent(done_.get().get());
       if (!recorded.ok()) {
         done_.SetError(recorded);
         return PjRtFuture<>(done_.GetError());
@@ -2576,7 +2575,7 @@ class StreamExecutorCopyToDeviceStream : public CopyToDeviceStream {
 
   // Async value will become available after we'll submit the last memcpy
   // operation, and the event will be recorded on the stream.
-  AsyncValueRef<se::Event> done_;
+  AsyncValueRef<std::unique_ptr<se::Event>> done_;
 };
 }  // namespace
 
@@ -2600,7 +2599,7 @@ static RecvDeviceMemoryFunction ConvertRecvCallbacksToRecvFunction(
   return [callbacks](int64_t channel_id, se::Stream* stream, const Shape& shape,
                      se::DeviceMemoryBase* dst,
                      const absl::flat_hash_map<std::string, std::string>&)
-             -> absl::StatusOr<AsyncValueRef<se::Event>> {
+             -> absl::StatusOr<AsyncValueRef<std::unique_ptr<se::Event>>> {
     VLOG(3) << "Recv from channel #" << channel_id
             << " (shape=" << shape.ToString() << ")";
 
@@ -2620,10 +2619,9 @@ static RecvDeviceMemoryFunction ConvertRecvCallbacksToRecvFunction(
     // Allocate event that will signal completion of recv operation. We record
     // it on a stream after submitting the memcpy for the last chunk (see
     // `StreamExecutorCopyToDeviceStream` implementation above).
-    auto done_event = MakeConstructedAsyncValueRef<se::Event>(stream->parent());
-    if (!done_event->Init())
-      return Internal("Failed to initialize done event (channel_id=%d)",
-                      channel_id);
+    TF_ASSIGN_OR_RETURN(auto event, stream->parent()->CreateEvent());
+    auto done_event = MakeConstructedAsyncValueRef<std::unique_ptr<se::Event>>(
+        std::move(event));
 
     recv->callback({shape}, std::make_unique<StreamExecutorCopyToDeviceStream>(
                                 channel_id, stream, *dst, done_event));
diff --git a/third_party/xla/xla/service/gpu/runtime/send_recv_thunk.cc b/third_party/xla/xla/service/gpu/runtime/send_recv_thunk.cc
index a7ff02de4beb70..00a22bc207c59b 100644
--- a/third_party/xla/xla/service/gpu/runtime/send_recv_thunk.cc
+++ b/third_party/xla/xla/service/gpu/runtime/send_recv_thunk.cc
@@ -69,9 +69,9 @@ static absl::StatusOr<bool> ShouldSkip(
 // SendRecvAsyncEvents
 //===----------------------------------------------------------------------===//
 
-absl::Status SendRecvAsyncEvents::Emplace(se::StreamExecutor* executor,
-                                          int32_t channel_id,
-                                          tsl::AsyncValueRef<se::Event> event) {
+absl::Status SendRecvAsyncEvents::Emplace(
+    se::StreamExecutor* executor, int32_t channel_id,
+    tsl::AsyncValueRef<std::unique_ptr<se::Event>> event) {
   Key key = {executor, channel_id};
 
   absl::MutexLock lock(&mutex_);
@@ -82,8 +82,8 @@ absl::Status SendRecvAsyncEvents::Emplace(se::StreamExecutor* executor,
       "Async send/recv event already exists (channel_id=%d)", channel_id));
 }
 
-absl::StatusOr<AsyncValueRef<se::Event>> SendRecvAsyncEvents::Extract(
-    se::StreamExecutor* executor, int32_t channel_id) {
+absl::StatusOr<AsyncValueRef<std::unique_ptr<se::Event>>>
+SendRecvAsyncEvents::Extract(se::StreamExecutor* executor, int32_t channel_id) {
   Key key = {executor, channel_id};
 
   absl::MutexLock lock(&mutex_);
@@ -135,7 +135,7 @@ absl::Status SendThunk::ExecuteOnStream(const ExecuteParams& params) {
   // Send buffer to a handler registered with the executable.
   if (auto* send = params.send_device_memory_function) {
     TF_ASSIGN_OR_RETURN(
-        AsyncValueRef<se::Event> done,
+        AsyncValueRef<std::unique_ptr<se::Event>> done,
         (*send)(channel_id_, stream, shape_, src, frontend_attrs_));
     return events_->Emplace(stream->parent(), channel_id_, std::move(done));
   }
@@ -176,7 +176,7 @@ absl::Status SendDoneThunk::ExecuteOnStream(const ExecuteParams& params) {
   VLOG(5) << "Completed Send operation: channel_id=" << channel_id_;
 
   // Once event is recorded we can add a stream dependency.
-  return params.stream->WaitFor(&done_event.get());
+  return params.stream->WaitFor(done_event.get().get());
 }
 
 //===----------------------------------------------------------------------===//
@@ -221,7 +221,7 @@ absl::Status RecvThunk::ExecuteOnStream(const ExecuteParams& params) {
   // Recv buffer from a handler registered with the run options.
   if (auto* recv = params.recv_device_memory_function) {
     TF_ASSIGN_OR_RETURN(
-        AsyncValueRef<se::Event> done,
+        AsyncValueRef<std::unique_ptr<se::Event>> done,
         (*recv)(channel_id_, stream, shape_, &dst, frontend_attrs_));
     return events_->Emplace(stream->parent(), channel_id_, std::move(done));
   }
@@ -261,7 +261,7 @@ absl::Status RecvDoneThunk::ExecuteOnStream(const ExecuteParams& params) {
   VLOG(5) << "Completed Recv operation: channel=" << channel_id_;
 
   // Once event is recorded we can add a stream dependency.
-  return params.stream->WaitFor(&done_event.get());
+  return params.stream->WaitFor(done_event.get().get());
 }
 
 }  // namespace xla::gpu
diff --git a/third_party/xla/xla/service/gpu/runtime/send_recv_thunk.h b/third_party/xla/xla/service/gpu/runtime/send_recv_thunk.h
index 6494e1dcc79e64..b146efc6799151 100644
--- a/third_party/xla/xla/service/gpu/runtime/send_recv_thunk.h
+++ b/third_party/xla/xla/service/gpu/runtime/send_recv_thunk.h
@@ -63,18 +63,18 @@ class SendRecvAsyncEvents {
  public:
   // Emplace a new send/recv completion event.
   absl::Status Emplace(se::StreamExecutor* executor, int32_t channel_id,
-                       tsl::AsyncValueRef<se::Event> event);
+                       tsl::AsyncValueRef<std::unique_ptr<se::Event>> event);
 
   // Extract a send/recv completion event.
-  absl::StatusOr<tsl::AsyncValueRef<se::Event>> Extract(
+  absl::StatusOr<tsl::AsyncValueRef<std::unique_ptr<se::Event>>> Extract(
       se::StreamExecutor* executor, int32_t channel_id);
 
  private:
   using Key = std::pair<se::StreamExecutor*, /*channel_id=*/int64_t>;
 
   absl::Mutex mutex_;
-  absl::flat_hash_map<Key, tsl::AsyncValueRef<se::Event>> events_
-      ABSL_GUARDED_BY(mutex_);
+  absl::flat_hash_map<Key, tsl::AsyncValueRef<std::unique_ptr<se::Event>>>
+      events_ ABSL_GUARDED_BY(mutex_);
 };
 
 //===----------------------------------------------------------------------===//

From 74236756f03cf8bc8af636bf294ba6636dcc96a4 Mon Sep 17 00:00:00 2001
From: David Dunleavy <ddunleavy@google.com>
Date: Fri, 17 May 2024 16:59:09 -0700
Subject: [PATCH 375/478] Tag `//xla/tools/hlo_opt:hlo_opt_gpu` with the `gpu`
 tag to avoid building on CPU builds of XLA

PiperOrigin-RevId: 634922729
---
 third_party/xla/xla/tools/hlo_opt/BUILD | 1 +
 1 file changed, 1 insertion(+)

diff --git a/third_party/xla/xla/tools/hlo_opt/BUILD b/third_party/xla/xla/tools/hlo_opt/BUILD
index 1e042b61849b5d..7657f4624f7bda 100644
--- a/third_party/xla/xla/tools/hlo_opt/BUILD
+++ b/third_party/xla/xla/tools/hlo_opt/BUILD
@@ -57,6 +57,7 @@ cc_library(
     name = "gpu_opt",
     testonly = True,
     srcs = if_gpu_is_configured(["gpu_opt.cc"]),
+    tags = ["gpu"],
     deps = [
         ":opt_lib",
         "//xla:debug_options_flags",

From a1e0516b8308138f52ea36508f9b0d8cad353764 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Fri, 17 May 2024 18:06:43 -0700
Subject: [PATCH 376/478] Migrate coord service and tests to use absl libraries
 (mutex, thread annotation, status) directly.

PiperOrigin-RevId: 634935787
---
 .../distributed_runtime/coordination/BUILD    |  22 +--
 .../coordination/coordination_service.cc      | 125 +++++++++---------
 .../coordination_service_agent_test.cc        |   7 +-
 .../coordination_service_error_util_test.cc   |  11 +-
 ...ordination_service_recoverable_job_test.cc |  45 ++++---
 .../coordination/coordination_service_test.cc |  45 ++++---
 6 files changed, 130 insertions(+), 125 deletions(-)

diff --git a/third_party/xla/xla/tsl/distributed_runtime/coordination/BUILD b/third_party/xla/xla/tsl/distributed_runtime/coordination/BUILD
index 137197420cfab2..5810a6ebcbc602 100644
--- a/third_party/xla/xla/tsl/distributed_runtime/coordination/BUILD
+++ b/third_party/xla/xla/tsl/distributed_runtime/coordination/BUILD
@@ -1,6 +1,7 @@
 load("@local_tsl//tsl/platform:build_config.bzl", "tf_proto_library", "tsl_cc_test")
 load("@local_tsl//tsl/platform:rules_cc.bzl", "cc_library")
 load("//xla/tsl:tsl.bzl", "if_oss", "internal_visibility", "tsl_gpu_library")
+load("//xla/tsl:tsl.default.bzl", "tsl_grpc_cc_dependencies")
 
 package(
     # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
@@ -26,7 +27,7 @@ tsl_cc_test(
     srcs = ["coordination_service_error_util_test.cc"],
     deps = [
         ":coordination_service_error_util",
-        "@local_tsl//tsl/platform:errors",
+        "@com_google_absl//absl/status",
         "@local_tsl//tsl/platform:status",
         "@local_tsl//tsl/platform:test",
         "@local_tsl//tsl/platform:test_main",
@@ -83,12 +84,9 @@ tsl_gpu_library(
         "@com_google_absl//absl/synchronization",
         "@com_google_absl//absl/time",
         "@local_tsl//tsl/platform:env",
-        "@local_tsl//tsl/platform:errors",
         "@local_tsl//tsl/platform:macros",
         "@local_tsl//tsl/platform:random",
         "@local_tsl//tsl/platform:status",
-        "@local_tsl//tsl/platform:strcat",
-        "@local_tsl//tsl/platform:thread_annotations",
         "@local_tsl//tsl/protobuf:coordination_config_proto_cc",
         "@local_tsl//tsl/protobuf:coordination_service_proto_cc",
     ],
@@ -116,19 +114,20 @@ tsl_cc_test(
         ":coordination_service_impl",
         ":test_device_proto_cc",
         "//xla/tsl/distributed_runtime:call_options",
+        "@com_google_absl//absl/base:core_headers",
+        "@com_google_absl//absl/log",
         "@com_google_absl//absl/status",
+        "@com_google_absl//absl/status:statusor",
         "@com_google_absl//absl/strings",
         "@com_google_absl//absl/synchronization",
         "@com_google_absl//absl/time",
         "@local_tsl//tsl/lib/core:status_test_util",
         "@local_tsl//tsl/platform:env",
         "@local_tsl//tsl/platform:errors",
-        "@local_tsl//tsl/platform:mutex",
         "@local_tsl//tsl/platform:random",
         "@local_tsl//tsl/platform:status",
         "@local_tsl//tsl/platform:test",
         "@local_tsl//tsl/platform:test_main",
-        "@local_tsl//tsl/platform:thread_annotations",
         "@local_tsl//tsl/platform:types",
         "@local_tsl//tsl/protobuf:coordination_config_proto_cc",
         "@local_tsl//tsl/protobuf:coordination_service_proto_cc",
@@ -172,14 +171,13 @@ tsl_cc_test(
         "@com_google_absl//absl/log:check",
         "@com_google_absl//absl/memory",
         "@com_google_absl//absl/status",
+        "@com_google_absl//absl/status:statusor",
         "@com_google_absl//absl/time",
         "@local_tsl//tsl/lib/core:status_test_util",
         "@local_tsl//tsl/platform:env",
         "@local_tsl//tsl/platform:env_impl",
-        "@local_tsl//tsl/platform:errors",
         "@local_tsl//tsl/platform:protobuf",
         "@local_tsl//tsl/platform:status",
-        "@local_tsl//tsl/platform:statusor",
         "@local_tsl//tsl/platform:test",
         "@local_tsl//tsl/platform:test_main",
         "@local_tsl//tsl/protobuf:coordination_config_proto_cc_impl",
@@ -222,20 +220,22 @@ tsl_cc_test(
         "//xla/tsl/distributed_runtime/rpc:async_service_interface",
         "//xla/tsl/distributed_runtime/rpc/coordination:grpc_coordination_client",
         "//xla/tsl/distributed_runtime/rpc/coordination:grpc_coordination_service_impl",
+        "@com_google_absl//absl/base:core_headers",
         "@com_google_absl//absl/container:flat_hash_map",
+        "@com_google_absl//absl/log",
         "@com_google_absl//absl/memory",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/synchronization",
         "@local_tsl//tsl/lib/core:status_test_util",
         "@local_tsl//tsl/platform:env",
         "@local_tsl//tsl/platform:env_impl",
-        "@local_tsl//tsl/platform:errors",
-        "@local_tsl//tsl/platform:mutex",
         "@local_tsl//tsl/platform:status",
         "@local_tsl//tsl/platform:test",
         "@local_tsl//tsl/platform:test_main",
         "@local_tsl//tsl/protobuf:coordination_config_proto_cc_impl",
         "@local_tsl//tsl/protobuf:coordination_service_proto_cc_impl",
         "@local_tsl//tsl/protobuf:distributed_runtime_payloads_proto_cc_impl",
-    ],
+    ] + tsl_grpc_cc_dependencies(),
 )
 
 filegroup(
diff --git a/third_party/xla/xla/tsl/distributed_runtime/coordination/coordination_service.cc b/third_party/xla/xla/tsl/distributed_runtime/coordination/coordination_service.cc
index 9b0dc57084d020..e18086dbadfb65 100644
--- a/third_party/xla/xla/tsl/distributed_runtime/coordination/coordination_service.cc
+++ b/third_party/xla/xla/tsl/distributed_runtime/coordination/coordination_service.cc
@@ -43,10 +43,8 @@ limitations under the License.
 #include "xla/tsl/distributed_runtime/coordination/coordination_service_error_util.h"
 #include "xla/tsl/util/device_name_utils.h"
 #include "tsl/platform/env.h"
-#include "tsl/platform/errors.h"
 #include "tsl/platform/random.h"
 #include "tsl/platform/status.h"
-#include "tsl/platform/thread_annotations.h"
 #include "tsl/protobuf/coordination_config.pb.h"
 #include "tsl/protobuf/coordination_service.pb.h"
 
@@ -110,7 +108,7 @@ class CoordinationServiceStandaloneImpl : public CoordinationServiceInterface {
       std::function<DeviceInfo(const DeviceInfo& devices)>
           post_aggregate_device_fn) override;
 
-  void LogConnectStatusLocked() const TF_EXCLUSIVE_LOCKS_REQUIRED(state_mu_);
+  void LogConnectStatusLocked() const ABSL_EXCLUSIVE_LOCKS_REQUIRED(state_mu_);
 
   absl::Status RegisterTask(const CoordinatedTask& task,
                             uint64_t incarnation) override;
@@ -142,11 +140,11 @@ class CoordinationServiceStandaloneImpl : public CoordinationServiceInterface {
 
  private:
   const DeviceInfo& ListClusterDevices() override
-      TF_EXCLUSIVE_LOCKS_REQUIRED(state_mu_);
+      ABSL_EXCLUSIVE_LOCKS_REQUIRED(state_mu_);
   uint64_t GetServiceIncarnation() override;
   void StartCheckStaleness();  // Checks both heartbeat and barrier timeouts.
   void Stop(bool shut_staleness_thread = true);
-  bool ServiceHasStopped() const TF_EXCLUSIVE_LOCKS_REQUIRED(state_mu_);
+  bool ServiceHasStopped() const ABSL_EXCLUSIVE_LOCKS_REQUIRED(state_mu_);
   // Report service error to a specified task.
   void ReportServiceErrorToTaskAsync(const CoordinatedTask& destination_task,
                                      absl::Status error);
@@ -155,16 +153,16 @@ class CoordinationServiceStandaloneImpl : public CoordinationServiceInterface {
   // Note: SetTaskError() must be called before propagating its error.
   void PropagateError(const CoordinatedTask& source_task,
                       bool is_reported_by_task = false)
-      TF_LOCKS_EXCLUDED(state_mu_);
+      ABSL_LOCKS_EXCLUDED(state_mu_);
   void SetTaskError(absl::string_view task_name, absl::Status error)
-      TF_EXCLUSIVE_LOCKS_REQUIRED(state_mu_);
-  void AggregateClusterDevices() TF_EXCLUSIVE_LOCKS_REQUIRED(state_mu_);
+      ABSL_EXCLUSIVE_LOCKS_REQUIRED(state_mu_);
+  void AggregateClusterDevices() ABSL_EXCLUSIVE_LOCKS_REQUIRED(state_mu_);
   absl::Status DisconnectTask(const CoordinatedTask& task)
-      TF_EXCLUSIVE_LOCKS_REQUIRED(state_mu_);
+      ABSL_EXCLUSIVE_LOCKS_REQUIRED(state_mu_);
 
   struct BarrierState {
     bool passed = false;
-    absl::Status result = errors::Unknown(
+    absl::Status result = absl::UnknownError(
         "Invalid barrier result.");  // Only valid if `passed` is true.
     uint64_t deadline_in_micros = 0;
     int num_pending_tasks = 0;
@@ -176,7 +174,7 @@ class CoordinationServiceStandaloneImpl : public CoordinationServiceInterface {
   };
   void PassBarrier(absl::string_view barrier_id, absl::Status result,
                    BarrierState* barrier)
-      TF_EXCLUSIVE_LOCKS_REQUIRED(state_mu_);
+      ABSL_EXCLUSIVE_LOCKS_REQUIRED(state_mu_);
   // Check if participating tasks are specified correctly across barrier calls.
   bool ValidateTaskArgs(
       const std::vector<CoordinatedTask>& tasks_args,
@@ -226,8 +224,7 @@ class CoordinationServiceStandaloneImpl : public CoordinationServiceInterface {
     CoordinatedTaskState state_ = CoordinatedTaskState::TASKSTATE_DISCONNECTED;
     absl::Status status_;
     absl::Mutex last_heartbeat_mu_;
-    // TODO(b/339231167): Migrate to use absl annotations.
-    uint64_t last_heartbeat_us_ TF_GUARDED_BY(last_heartbeat_mu_);
+    uint64_t last_heartbeat_us_ ABSL_GUARDED_BY(last_heartbeat_mu_);
     // This denotes the deadline after which we stop accepting heartbeats from a
     // disconnected task. This grace period accounts for the lag time between
     // the service recording the state change and the agent stopping heartbeats.
@@ -257,24 +254,24 @@ class CoordinationServiceStandaloneImpl : public CoordinationServiceInterface {
 
   absl::Mutex state_mu_;
   absl::flat_hash_map<std::string, std::unique_ptr<TaskState>> cluster_state_
-      TF_GUARDED_BY(state_mu_);
-  DeviceInfo cluster_devices_ TF_GUARDED_BY(state_mu_);
+      ABSL_GUARDED_BY(state_mu_);
+  DeviceInfo cluster_devices_ ABSL_GUARDED_BY(state_mu_);
 
   absl::Mutex kv_mu_;
   // Ordered map to store config key-values
-  std::map<std::string, std::string> kv_store_ TF_GUARDED_BY(kv_mu_);
+  std::map<std::string, std::string> kv_store_ ABSL_GUARDED_BY(kv_mu_);
   absl::flat_hash_map<std::string, std::vector<StatusOrValueCallback>> get_cb_
-      TF_GUARDED_BY(kv_mu_);
+      ABSL_GUARDED_BY(kv_mu_);
 
   absl::CondVar check_staleness_thread_cv_;
   bool shutting_down_ ABSL_GUARDED_BY(state_mu_) = false;
   std::unique_ptr<Thread> check_staleness_thread_;
 
   absl::flat_hash_map<std::string, BarrierState> barriers_
-      TF_GUARDED_BY(state_mu_);
+      ABSL_GUARDED_BY(state_mu_);
   // For now, we assume there won't be many simultaneous barriers so we simply
   // use a set.
-  absl::flat_hash_set<std::string> ongoing_barriers_ TF_GUARDED_BY(state_mu_);
+  absl::flat_hash_set<std::string> ongoing_barriers_ ABSL_GUARDED_BY(state_mu_);
 
   absl::flat_hash_set<std::string> recoverable_jobs_;
 
@@ -311,10 +308,9 @@ absl::Status CoordinationServiceStandaloneImpl::TaskState::RecordHeartbeat(
     uint64_t task_incarnation) {
   if (!status_.ok()) return status_;
   if (task_incarnation != task_incarnation_) {
-    // TODO(b/339231167): Migrate to use absl error API.
-    return MakeCoordinationError(errors::Aborted(
+    return MakeCoordinationError(absl::AbortedError(absl::StrCat(
         "Incarnation ID mismatch: expecting ", task_incarnation_, " but got ",
-        task_incarnation, ". This means the remote task has restarted."));
+        task_incarnation, ". This means the remote task has restarted.")));
   }
   absl::MutexLock l(&last_heartbeat_mu_);
   last_heartbeat_us_ = Env::Default()->NowMicros();
@@ -413,11 +409,14 @@ void CoordinationServiceStandaloneImpl::StartCheckStaleness() {
                        << " stale?=" << is_stale;
               if (is_stale) {
                 stale_task_names.push_back(task_name);
-                status = MakeCoordinationError(errors::Unavailable(
-                    "Task ", task_name,
-                    " heartbeat timeout. This indicates that the remote task "
-                    "has failed, got preempted, or crashed unexpectedly. Check "
-                    "the task logs for an earlier error to debug further."));
+                status = MakeCoordinationError(absl::UnavailableError(
+                    absl::StrCat("Task ", task_name,
+                                 " heartbeat timeout. This indicates that the "
+                                 "remote task "
+                                 "has failed, got preempted, or crashed "
+                                 "unexpectedly. Check "
+                                 "the task logs for an earlier error to debug "
+                                 "further.")));
                 SetTaskError(task_name, status);
               }
             }
@@ -469,8 +468,8 @@ void CoordinationServiceStandaloneImpl::StartCheckStaleness() {
                   }
                 }
               }
-              const absl::Status error =
-                  MakeCoordinationError(errors::DeadlineExceeded(absl::StrCat(
+              const absl::Status error = MakeCoordinationError(
+                  absl::DeadlineExceededError(absl::StrCat(
                       "Barrier timed out. Barrier_id: ", barrier_id,
                       ". Timed out task names:\n", pending_tasks)));
               PassBarrier(barrier_id, error, barrier);
@@ -498,7 +497,7 @@ void CoordinationServiceStandaloneImpl::Stop(bool shut_staleness_thread) {
     absl::MutexLock l(&kv_mu_);
     for (const auto& [key, get_kv_callbacks] : get_cb_) {
       for (const auto& get_kv_callback : get_kv_callbacks) {
-        get_kv_callback(errors::Cancelled(
+        get_kv_callback(absl::CancelledError(
             absl::StrCat("Coordination service is shutting down. Cancelling "
                          "GetKeyValue() for key: ",
                          key)));
@@ -515,9 +514,10 @@ void CoordinationServiceStandaloneImpl::Stop(bool shut_staleness_thread) {
     // Fail all ongoing barriers.
     for (auto& [barrier_id, barrier] : barriers_) {
       if (!barrier.passed) {
-        absl::Status error = MakeCoordinationError(errors::Aborted(absl::StrCat(
-            "Barrier failed because service is shutting down. Barrier_id: ",
-            barrier_id)));
+        absl::Status error =
+            MakeCoordinationError(absl::AbortedError(absl::StrCat(
+                "Barrier failed because service is shutting down. Barrier_id: ",
+                barrier_id)));
         PassBarrier(barrier_id, error, &barrier);
       }
     }
@@ -577,8 +577,8 @@ absl::Status CoordinationServiceStandaloneImpl::RegisterTask(
     if (!cluster_state_.contains(task_name)) {
       // Note: return early here as unexpected task register errors should not
       // be propagated to other tasks.
-      return MakeCoordinationError(errors::InvalidArgument(
-          "Unexpected task registered with task_name=", task_name));
+      return MakeCoordinationError(absl::InvalidArgumentError(absl::StrCat(
+          "Unexpected task registered with task_name=", task_name)));
     }
 
     auto* task_cluster_state = cluster_state_[task_name].get();
@@ -587,7 +587,7 @@ absl::Status CoordinationServiceStandaloneImpl::RegisterTask(
 
     if (task_state == CoordinatedTaskState::TASKSTATE_DISCONNECTED ||
         (allow_new_incarnation_to_reconnect_ &&
-         (errors::IsUnavailable(task_status) &&
+         (absl::IsUnavailable(task_status) &&
           task_status.GetPayload(CoordinationErrorPayloadKey())))) {
       // The task is allowed to register itself if:
       // - this task is currently disconnected (registering for the first time
@@ -631,7 +631,7 @@ absl::Status CoordinationServiceStandaloneImpl::RegisterTask(
                        "subsequent connect attempt.");
     }
     LOG(ERROR) << error_message;
-    error = MakeCoordinationError(errors::Aborted(error_message), task);
+    error = MakeCoordinationError(absl::AbortedError(error_message), task);
     SetTaskError(task_name, error);
   }
   assert(!error.ok());
@@ -700,12 +700,12 @@ absl::Status CoordinationServiceStandaloneImpl::DisconnectTask(
                      "failed for task_name=",
                      task_name)));
   } else if (!cluster_state_.contains(task_name)) {
-    return MakeCoordinationError(errors::InvalidArgument(
-        "Unexpected disconnect request with task_name=", task_name));
+    return MakeCoordinationError(absl::InvalidArgumentError(absl::StrCat(
+        "Unexpected disconnect request with task_name=", task_name)));
   } else if (cluster_state_[task_name]->GetState() ==
              CoordinatedTaskState::TASKSTATE_DISCONNECTED) {
-    return MakeCoordinationError(errors::FailedPrecondition(
-        "The task is already disconnected: ", task_name));
+    return MakeCoordinationError(absl::FailedPreconditionError(
+        absl::StrCat("The task is already disconnected: ", task_name)));
   }
 
   // Disconnect task and fail any ongoing barriers.
@@ -713,7 +713,7 @@ absl::Status CoordinationServiceStandaloneImpl::DisconnectTask(
       /*grace_period_duration_us=*/heartbeat_timeout_ms_ * 1000);
   for (const auto& barrier_id :
        cluster_state_[task_name]->GetOngoingBarriers()) {
-    absl::Status error = MakeCoordinationError(errors::Internal(absl::StrCat(
+    absl::Status error = MakeCoordinationError(absl::InternalError(absl::StrCat(
         "Barrier failed from a disconnected task. Barrier Id: ", barrier_id,
         ", Task: ", task_name)));
     PassBarrier(barrier_id, error, &barriers_[barrier_id]);
@@ -740,11 +740,11 @@ absl::Status CoordinationServiceStandaloneImpl::ReportTaskError(
       return MakeCoordinationError(absl::InternalError(
           "Coordination service has stopped. ReportTaskError() failed."));
     } else if (!cluster_state_.contains(task_name)) {
-      return MakeCoordinationError(
-          errors::InvalidArgument("Unexpected request from task ", task_name));
+      return MakeCoordinationError(absl::InvalidArgumentError(
+          absl::StrCat("Unexpected request from task ", task_name)));
     } else if (cluster_state_[task_name]->GetState() !=
                CoordinatedTaskState::TASKSTATE_CONNECTED) {
-      return MakeCoordinationError(errors::FailedPrecondition(
+      return MakeCoordinationError(absl::FailedPreconditionError(
           "The task is not connected or already has an error."));
     } else {
       SetTaskError(task_name, error);
@@ -793,9 +793,9 @@ absl::Status CoordinationServiceStandaloneImpl::RecordHeartbeat(
           "gracefully. Check the task leader's logs for an earlier error to "
           "debug the root cause.")));
     } else if (!cluster_state_.contains(task_name)) {
-      return MakeCoordinationError(errors::InvalidArgument(
-          "Unexpected heartbeat request from task: ", task_name,
-          ". This usually implies a configuration error."));
+      return MakeCoordinationError(absl::InvalidArgumentError(
+          absl::StrCat("Unexpected heartbeat request from task: ", task_name,
+                       ". This usually implies a configuration error.")));
     }
     if (!cluster_state_[task_name]->GetStatus().ok()) {
       return cluster_state_[task_name]->GetStatus();
@@ -807,9 +807,9 @@ absl::Status CoordinationServiceStandaloneImpl::RecordHeartbeat(
                Env::Default()->NowMicros() >
                    cluster_state_[task_name]
                        ->GetDisconnectedGracePeriodMicros()) {
-      return MakeCoordinationError(errors::InvalidArgument(
+      return MakeCoordinationError(absl::InvalidArgumentError(absl::StrCat(
           "Task with task_name=", task_name,
-          " must be registered before sending heartbeat messages"));
+          " must be registered before sending heartbeat messages")));
     }
     s = cluster_state_[task_name]->RecordHeartbeat(incarnation);
   }
@@ -958,8 +958,8 @@ absl::Status CoordinationServiceStandaloneImpl::InsertKeyValue(
   const std::string norm_key = NormalizeKey(key);
   absl::MutexLock l(&kv_mu_);
   if (kv_store_.find(norm_key) != kv_store_.end()) {
-    return MakeCoordinationError(
-        errors::AlreadyExists("Config key ", key, " already exists."));
+    return MakeCoordinationError(absl::AlreadyExistsError(
+        absl::StrCat("Config key ", key, " already exists.")));
   }
   kv_store_.emplace(norm_key, value);
   auto iter = get_cb_.find(norm_key);
@@ -997,7 +997,7 @@ absl::StatusOr<std::string> CoordinationServiceStandaloneImpl::TryGetKeyValue(
   absl::MutexLock l(&kv_mu_);
   const auto& iter = kv_store_.find(norm_key);
   if (iter == kv_store_.end()) {
-    return errors::NotFound("Config key ", key, " not found.");
+    return absl::NotFoundError(absl::StrCat("Config key ", key, " not found."));
   }
   return iter->second;
 }
@@ -1057,7 +1057,7 @@ void CoordinationServiceStandaloneImpl::SetTaskError(
   cluster_state_[task_name]->SetError(error);
   for (const auto& barrier_id :
        cluster_state_[task_name]->GetOngoingBarriers()) {
-    absl::Status error = MakeCoordinationError(errors::Internal(absl::StrCat(
+    absl::Status error = MakeCoordinationError(absl::InternalError(absl::StrCat(
         "Barrier failed from a task error. Barrier Id: ", barrier_id,
         ", Task: ", task_name)));
     PassBarrier(barrier_id, error, &barriers_[barrier_id]);
@@ -1088,7 +1088,7 @@ void CoordinationServiceStandaloneImpl::BarrierAsync(
 
   if (!participating_tasks.empty() && !among_participating_tasks) {
     const std::string task_name = GetTaskName(task);
-    absl::Status error = MakeCoordinationError(errors::InvalidArgument(
+    absl::Status error = MakeCoordinationError(absl::InvalidArgumentError(
         absl::StrCat("A non-participating task (", GetTaskName(task),
                      ") called the barrier: ", barrier_id)));
     {
@@ -1138,7 +1138,7 @@ void CoordinationServiceStandaloneImpl::BarrierAsync(
         // barrier.
         const std::string task_name = GetTaskName(task);
         if (!cluster_state_.contains(task_name)) {
-          absl::Status error = MakeCoordinationError(errors::InvalidArgument(
+          absl::Status error = MakeCoordinationError(absl::InvalidArgumentError(
               absl::StrCat("Unexpected task (", task_name,
                            ") that is not in the cluster called the barrier. "
                            "Barrier Id: ",
@@ -1157,7 +1157,7 @@ void CoordinationServiceStandaloneImpl::BarrierAsync(
       const std::string task_name = GetTaskName(pending_task.first);
       if (cluster_state_[task_name]->GetState() ==
           CoordinatedTaskState::TASKSTATE_ERROR) {
-        absl::Status error = MakeCoordinationError(errors::Internal(
+        absl::Status error = MakeCoordinationError(absl::InternalError(
             absl::StrCat("Task (", task_name,
                          ") is already in error before the barrier "
                          "was called. Barrier Id: ",
@@ -1208,7 +1208,7 @@ void CoordinationServiceStandaloneImpl::BarrierAsync(
   if (!ValidateTaskArgs(participating_tasks, barrier->tasks_at_barrier,
                         cluster_state_.size())) {
     absl::Status error =
-        MakeCoordinationError(errors::InvalidArgument(absl::StrCat(
+        MakeCoordinationError(absl::InvalidArgumentError(absl::StrCat(
             "Conflicting tasks specified for the same barrier: ", barrier_id)));
     PassBarrier(barrier_id, error, barrier);
     return;
@@ -1243,14 +1243,15 @@ absl::Status CoordinationServiceStandaloneImpl::CancelBarrier(
   }
   // Barrier has already been passed.
   if (barrier->passed) {
-    return MakeCoordinationError(errors::FailedPrecondition(absl::StrCat(
+    return MakeCoordinationError(absl::FailedPreconditionError(absl::StrCat(
         "Barrier (", barrier_id, ") has already been passed with status code: ",
         barrier->result.code())));
   }
 
   // Cancel barrier.
-  absl::Status cancelled = MakeCoordinationError(errors::Cancelled(absl::StrCat(
-      "Barrier (", barrier_id, ") is cancelled by task: ", GetTaskName(task))));
+  absl::Status cancelled = MakeCoordinationError(absl::CancelledError(
+      absl::StrCat("Barrier (", barrier_id,
+                   ") is cancelled by task: ", GetTaskName(task))));
   PassBarrier(barrier_id, cancelled, barrier);
 
   VLOG(3) << "Barrier (" << barrier_id << ") is cancelled.";
@@ -1285,7 +1286,7 @@ void CoordinationServiceStandaloneImpl::PassBarrier(
                     "crashed early or too slow / hanging. Check the logs for "
                     "an earlier error to identify the root cause.";
     }
-    absl::Status shutdown_error = MakeCoordinationError(errors::Internal(
+    absl::Status shutdown_error = MakeCoordinationError(absl::InternalError(
         absl::StrCat("Shutdown barrier has been passed with status: '",
                      barrier->result.ToString(),
                      "', but this task is not at the barrier yet.")));
diff --git a/third_party/xla/xla/tsl/distributed_runtime/coordination/coordination_service_agent_test.cc b/third_party/xla/xla/tsl/distributed_runtime/coordination/coordination_service_agent_test.cc
index 419471938b8981..db0cadf874bacc 100644
--- a/third_party/xla/xla/tsl/distributed_runtime/coordination/coordination_service_agent_test.cc
+++ b/third_party/xla/xla/tsl/distributed_runtime/coordination/coordination_service_agent_test.cc
@@ -31,9 +31,7 @@ limitations under the License.
 #include "xla/tsl/distributed_runtime/coordination/coordination_client.h"
 #include "tsl/lib/core/status_test_util.h"
 #include "tsl/platform/env.h"
-#include "tsl/platform/errors.h"
 #include "tsl/platform/status.h"
-#include "tsl/platform/statusor.h"
 #include "tsl/platform/test.h"
 #include "tsl/protobuf/coordination_config.pb.h"
 #include "tsl/protobuf/coordination_service.pb.h"
@@ -250,7 +248,7 @@ TEST_F(CoordinationServiceAgentTest, GetKeyValue_Timeout_ReturnError) {
 
   auto result = agent_->GetKeyValue(test_key, /*timeout=*/absl::Seconds(1));
 
-  EXPECT_EQ(result.status().code(), error::DEADLINE_EXCEEDED);
+  EXPECT_TRUE(absl::IsDeadlineExceeded(result.status()));
   // Needed to tear down test safely since agent dtor would cancel pending
   // calls, which would reference deallocated call_opts.
   owned_done(absl::CancelledError("error"));
@@ -275,8 +273,7 @@ TEST_F(CoordinationServiceAgentTest,
   InitializeAgent();
 
   auto result = agent_->GetKeyValue(test_key, /*timeout=*/absl::Seconds(1));
-  EXPECT_EQ(result.status().code(), error::DEADLINE_EXCEEDED);
-
+  EXPECT_TRUE(absl::IsDeadlineExceeded(result.status()));
   // Delayed server response: set key-value response, and invoke done callback.
   auto kv = owned_response->mutable_kv();
   kv->set_key(test_key);
diff --git a/third_party/xla/xla/tsl/distributed_runtime/coordination/coordination_service_error_util_test.cc b/third_party/xla/xla/tsl/distributed_runtime/coordination/coordination_service_error_util_test.cc
index 9f1baada2ce84c..3c19fa5759a207 100644
--- a/third_party/xla/xla/tsl/distributed_runtime/coordination/coordination_service_error_util_test.cc
+++ b/third_party/xla/xla/tsl/distributed_runtime/coordination/coordination_service_error_util_test.cc
@@ -16,8 +16,7 @@ limitations under the License.
 
 #include <string>
 
-#include "tsl/platform/errors.h"
-#include "tsl/platform/status.h"
+#include "absl/status/status.h"
 #include "tsl/platform/test.h"
 #include "tsl/protobuf/coordination_service.pb.h"
 namespace tsl {
@@ -26,7 +25,7 @@ using ::tensorflow::CoordinatedTask;
 using ::tensorflow::CoordinationServiceError;
 
 TEST(CoordinationServiceErrorUtil, MakeCoordinationErrorWithEmptyPayload) {
-  absl::Status error = errors::Internal("Test Error");
+  absl::Status error = absl::InternalError("Test Error");
 
   absl::Status coordination_error = MakeCoordinationError(error);
 
@@ -38,7 +37,7 @@ TEST(CoordinationServiceErrorUtil, MakeCoordinationErrorWithEmptyPayload) {
 }
 
 TEST(CoordinationServiceErrorUtil, MakeCoordinationErrorWithErrorOrigin) {
-  absl::Status error = errors::Internal("Test Error");
+  absl::Status error = absl::InternalError("Test Error");
   CoordinatedTask source_task;
   source_task.set_job_name("test_worker");
   source_task.set_task_id(7);
@@ -57,7 +56,7 @@ TEST(CoordinationServiceErrorUtil, MakeCoordinationErrorWithErrorOrigin) {
 }
 
 TEST(CoordinationServiceErrorUtil, MakeCoordinationErrorWithUserReportedError) {
-  absl::Status error = errors::Internal("Test Error");
+  absl::Status error = absl::InternalError("Test Error");
   CoordinatedTask source_task;
   source_task.set_job_name("test_worker");
   source_task.set_task_id(7);
@@ -78,7 +77,7 @@ TEST(CoordinationServiceErrorUtil, MakeCoordinationErrorWithUserReportedError) {
 }
 
 TEST(CoordinationServiceErrorUtil, MakeCoordinationErrorWithPayload) {
-  absl::Status error = errors::Internal("Test Error");
+  absl::Status error = absl::InternalError("Test Error");
   CoordinationServiceError payload;
   CoordinatedTask* source_task = payload.mutable_source_task();
   source_task->set_job_name("test_worker");
diff --git a/third_party/xla/xla/tsl/distributed_runtime/coordination/coordination_service_recoverable_job_test.cc b/third_party/xla/xla/tsl/distributed_runtime/coordination/coordination_service_recoverable_job_test.cc
index 37f128b53ff415..836fb8803e02f7 100644
--- a/third_party/xla/xla/tsl/distributed_runtime/coordination/coordination_service_recoverable_job_test.cc
+++ b/third_party/xla/xla/tsl/distributed_runtime/coordination/coordination_service_recoverable_job_test.cc
@@ -15,10 +15,16 @@ limitations under the License.
 #include <memory>
 #include <string>
 #include <utility>
-#include <vector>
 
+#include "grpcpp/server.h"
+#include "grpcpp/server_builder.h"
+#include "grpcpp/support/channel_arguments.h"
+#include "absl/base/thread_annotations.h"
 #include "absl/container/flat_hash_map.h"
+#include "absl/log/log.h"
 #include "absl/memory/memory.h"
+#include "absl/status/status.h"
+#include "absl/synchronization/mutex.h"
 #include "xla/tsl/distributed_runtime/coordination/coordination_client.h"
 #include "xla/tsl/distributed_runtime/coordination/coordination_service.h"
 #include "xla/tsl/distributed_runtime/coordination/coordination_service_agent.h"
@@ -27,7 +33,6 @@ limitations under the License.
 #include "xla/tsl/distributed_runtime/rpc/coordination/grpc_coordination_service_impl.h"
 #include "tsl/lib/core/status_test_util.h"
 #include "tsl/platform/env.h"
-#include "tsl/platform/mutex.h"
 #include "tsl/platform/status.h"
 #include "tsl/platform/test.h"
 #include "tsl/platform/threadpool.h"
@@ -46,12 +51,12 @@ constexpr char kServiceLeader[] = "/job:parameter_server/replica:0/task:0";
 class TestCoordinationClientCache : public CoordinationClientCache {
  public:
   void AddTask(const std::string& target, CoordinationClient* client) {
-    mutex_lock l(clients_mu_);
+    absl::MutexLock l(&clients_mu_);
     clients_.emplace(target, client);
   }
 
-  CoordinationClient* GetClient(const string& target) override {
-    mutex_lock l(clients_mu_);
+  CoordinationClient* GetClient(const std::string& target) override {
+    absl::MutexLock l(&clients_mu_);
     if (auto it = clients_.find(target); it != clients_.end()) {
       return it->second;
     }
@@ -59,15 +64,15 @@ class TestCoordinationClientCache : public CoordinationClientCache {
   }
 
   std::unique_ptr<CoordinationClient> GetOwnedClient(
-      const string& target) override {
+      const std::string& target) override {
     LOG(ERROR) << "GetOwnedClient is not supported.";
     return nullptr;
   }
 
  private:
-  mutex clients_mu_;
+  absl::Mutex clients_mu_;
   absl::flat_hash_map<std::string, CoordinationClient*> clients_
-      TF_GUARDED_BY(clients_mu_);
+      ABSL_GUARDED_BY(clients_mu_);
 };
 
 class TestCoordinationServiceTaskState {
@@ -229,37 +234,37 @@ class CoordinationServiceRecoverableJobTest : public ::testing::Test {
 TEST_F(CoordinationServiceRecoverableJobTest,
        UnrecoverableWorkerFailurePropagated) {
   Initialize();
-  TF_ASSERT_OK(state_worker_0_.ReportError(errors::Internal("Test Error.")));
+  TF_ASSERT_OK(state_worker_0_.ReportError(absl::InternalError("Test Error.")));
 
   // For unrecoverable task, error propagates to all connected tasks.
-  EXPECT_TRUE(errors::IsInternal(state_ps_0_.GetStatus()));
-  EXPECT_TRUE(errors::IsInternal(state_ps_1_.GetStatus()));
-  EXPECT_TRUE(errors::IsInternal(state_worker_0_.GetStatus()));
-  EXPECT_TRUE(errors::IsInternal(state_worker_1_.GetStatus()));
+  EXPECT_TRUE(absl::IsInternal(state_ps_0_.GetStatus()));
+  EXPECT_TRUE(absl::IsInternal(state_ps_1_.GetStatus()));
+  EXPECT_TRUE(absl::IsInternal(state_worker_0_.GetStatus()));
+  EXPECT_TRUE(absl::IsInternal(state_worker_1_.GetStatus()));
 }
 
 TEST_F(CoordinationServiceRecoverableJobTest,
        UnrecoverablePSFailurePropagated) {
   Initialize();
-  TF_ASSERT_OK(state_ps_0_.ReportError(errors::Internal("Test Error.")));
+  TF_ASSERT_OK(state_ps_0_.ReportError(absl::InternalError("Test Error.")));
 
   // For unrecoverable task, error propagates to all connected tasks.
-  EXPECT_TRUE(errors::IsInternal(state_ps_0_.GetStatus()));
-  EXPECT_TRUE(errors::IsInternal(state_ps_1_.GetStatus()));
-  EXPECT_TRUE(errors::IsInternal(state_worker_0_.GetStatus()));
-  EXPECT_TRUE(errors::IsInternal(state_worker_1_.GetStatus()));
+  EXPECT_TRUE(absl::IsInternal(state_ps_0_.GetStatus()));
+  EXPECT_TRUE(absl::IsInternal(state_ps_1_.GetStatus()));
+  EXPECT_TRUE(absl::IsInternal(state_worker_0_.GetStatus()));
+  EXPECT_TRUE(absl::IsInternal(state_worker_1_.GetStatus()));
 }
 
 TEST_F(CoordinationServiceRecoverableJobTest,
        RecoverableWorkerFailureNotPropagated) {
   AddJobToRecoverableJobs(kWorkerJobName);
   Initialize();
-  TF_ASSERT_OK(state_worker_0_.ReportError(errors::Internal("Test Error.")));
+  TF_ASSERT_OK(state_worker_0_.ReportError(absl::InternalError("Test Error.")));
 
   // For recoverable task, error does not propagate.
   EXPECT_TRUE(state_ps_0_.GetStatus().ok());
   EXPECT_TRUE(state_ps_1_.GetStatus().ok());
-  EXPECT_TRUE(errors::IsInternal(state_worker_0_.GetStatus()));
+  EXPECT_TRUE(absl::IsInternal(state_worker_0_.GetStatus()));
   EXPECT_TRUE(state_worker_1_.GetStatus().ok());
 }
 
diff --git a/third_party/xla/xla/tsl/distributed_runtime/coordination/coordination_service_test.cc b/third_party/xla/xla/tsl/distributed_runtime/coordination/coordination_service_test.cc
index 8e222353e078e0..8eba77c43211f7 100644
--- a/third_party/xla/xla/tsl/distributed_runtime/coordination/coordination_service_test.cc
+++ b/third_party/xla/xla/tsl/distributed_runtime/coordination/coordination_service_test.cc
@@ -15,15 +15,21 @@ limitations under the License.
 
 #include "xla/tsl/distributed_runtime/coordination/coordination_service.h"
 
+#include <cstdint>
 #include <memory>
 #include <string>
 #include <unordered_map>
 #include <utility>
 #include <vector>
 
+#include "absl/base/thread_annotations.h"
+#include "absl/log/log.h"
 #include "absl/status/status.h"
+#include "absl/status/statusor.h"
 #include "absl/strings/match.h"
 #include "absl/strings/str_cat.h"
+#include "absl/strings/string_view.h"
+#include "absl/synchronization/mutex.h"
 #include "absl/synchronization/notification.h"
 #include "absl/time/time.h"
 #include "xla/tsl/distributed_runtime/call_options.h"
@@ -32,12 +38,9 @@ limitations under the License.
 #include "xla/tsl/distributed_runtime/coordination/test_device.pb.h"
 #include "tsl/lib/core/status_test_util.h"
 #include "tsl/platform/env.h"
-#include "tsl/platform/errors.h"
-#include "tsl/platform/mutex.h"
 #include "tsl/platform/random.h"
 #include "tsl/platform/status.h"
 #include "tsl/platform/test.h"
-#include "tsl/platform/thread_annotations.h"
 #include "tsl/platform/types.h"
 #include "tsl/protobuf/coordination_config.pb.h"
 #include "tsl/protobuf/coordination_service.pb.h"
@@ -81,7 +84,7 @@ class TestCoordinationClient : public CoordinationClient {
   TestCoordinationClient() = default;
 
   absl::Status GetStatus() {
-    mutex_lock l(mu_);
+    absl::MutexLock l(&mu_);
     return status_;
   }
 
@@ -95,7 +98,7 @@ class TestCoordinationClient : public CoordinationClient {
                               const ReportErrorToTaskRequest* request,
                               ReportErrorToTaskResponse* response,
                               StatusCallback done) override {
-    mutex_lock l(mu_);
+    absl::MutexLock l(&mu_);
     status_ = absl::Status(static_cast<absl::StatusCode>(request->error_code()),
                            request->error_message());
     done(absl::OkStatus());
@@ -104,7 +107,7 @@ class TestCoordinationClient : public CoordinationClient {
 #define UNIMPLEMENTED(method)                                         \
   void method##Async(const method##Request* request,                  \
                      method##Response* response, StatusCallback done) \
-      override{done(errors::Unimplemented(#method "Async"));          \
+      override{done(absl::UnimplementedError(#method "Async"));       \
   }
 
   UNIMPLEMENTED(WaitForAllTasks);
@@ -122,7 +125,7 @@ class TestCoordinationClient : public CoordinationClient {
 #define UNIMPLEMENTED_WITH_CALL_OPTS(method)                                 \
   void method##Async(CallOptions* call_opts, const method##Request* request, \
                      method##Response* response, StatusCallback done)        \
-      override{done(errors::Unimplemented(#method "Async"));                 \
+      override{done(absl::UnimplementedError(#method "Async"));              \
   }
 
   UNIMPLEMENTED_WITH_CALL_OPTS(GetKeyValue);
@@ -131,8 +134,8 @@ class TestCoordinationClient : public CoordinationClient {
 #undef UNIMPLEMENTED_WITH_CALL_OPTS
 
  private:
-  mutex mu_;
-  absl::Status status_ TF_GUARDED_BY(mu_);
+  absl::Mutex mu_;
+  absl::Status status_ ABSL_GUARDED_BY(mu_);
 };
 
 class TestCoordinationClientCache : public CoordinationClientCache {
@@ -436,8 +439,8 @@ TEST(CoordinationServiceTest, RegisterTask_AlreadyInError_Fails) {
   // Task connects to coordination service.
   TF_ASSERT_OK(coord_service->RegisterTask(task_0, /*incarnation=*/0));
   // Arbitrarily set task to be in error.
-  TF_ASSERT_OK(
-      coord_service->ReportTaskError(task_0, errors::Internal("test_error")));
+  TF_ASSERT_OK(coord_service->ReportTaskError(
+      task_0, absl::InternalError("test_error")));
 
   // Registration should fail since task already registered previously.
   const absl::Status status =
@@ -1073,7 +1076,7 @@ TEST_F(CoordinationBarrierTest, BarrierReturnsPreviousError) {
         n_0.Notify();
       });
   TF_ASSERT_OK(GetCoordinationService()->ReportTaskError(
-      GetTask(0), errors::Internal("test_error")));
+      GetTask(0), absl::InternalError("test_error")));
   // Block until barrier has failed due to task error.
   n_0.WaitForNotification();
   // Same response should be returned immediately.
@@ -1201,7 +1204,7 @@ TEST_F(CoordinationBarrierTest, BarrierFailsIfTaskIsAlreadyInError) {
   absl::Duration timeout = absl::Seconds(5);
   // Set task 0 to error state.
   TF_ASSERT_OK(GetCoordinationService()->ReportTaskError(
-      GetTask(0), errors::Internal("test_error")));
+      GetTask(0), absl::InternalError("test_error")));
   absl::Status barrier_status;
 
   GetCoordinationService()->BarrierAsync(
@@ -1225,7 +1228,7 @@ TEST_F(CoordinationBarrierTest, BarrierFailsUponTaskError) {
         n0.Notify();
       });
   TF_ASSERT_OK(GetCoordinationService()->ReportTaskError(
-      GetTask(0), errors::Internal("test_error")));
+      GetTask(0), absl::InternalError("test_error")));
   n0.WaitForNotification();
 
   EXPECT_TRUE(absl::IsInternal(barrier_status));
@@ -1534,8 +1537,8 @@ TEST_F(CoordinateTwoTasksTest, UnrecoverableTaskPropagatesError) {
   TF_EXPECT_OK(coord_service_->RegisterTask(task_0_, incarnation_0_));
   TF_EXPECT_OK(coord_service_->RegisterTask(task_1_, incarnation_1_));
 
-  TF_ASSERT_OK(
-      coord_service_->ReportTaskError(task_0_, errors::Internal("test_error")));
+  TF_ASSERT_OK(coord_service_->ReportTaskError(
+      task_0_, absl::InternalError("test_error")));
 
   EXPECT_TRUE(absl::IsInternal(
       coord_service_->RecordHeartbeat(task_0_, incarnation_0_)));
@@ -1551,8 +1554,8 @@ TEST_F(CoordinateTwoTasksTest, RecoverableTaskWillNotPropagateError) {
   TF_EXPECT_OK(coord_service_->RegisterTask(task_0_, incarnation_0_));
   TF_EXPECT_OK(coord_service_->RegisterTask(task_1_, incarnation_1_));
 
-  TF_ASSERT_OK(
-      coord_service_->ReportTaskError(task_0_, errors::Internal("test_error")));
+  TF_ASSERT_OK(coord_service_->ReportTaskError(
+      task_0_, absl::InternalError("test_error")));
 
   EXPECT_TRUE(absl::IsInternal(
       coord_service_->RecordHeartbeat(task_0_, incarnation_0_)));
@@ -1570,8 +1573,8 @@ TEST_F(CoordinateTwoTasksTest,
   TF_EXPECT_OK(coord_service_->RegisterTask(task_0_, incarnation_0_));
   TF_EXPECT_OK(coord_service_->RegisterTask(task_1_, incarnation_1_));
 
-  TF_ASSERT_OK(
-      coord_service_->ReportTaskError(task_0_, errors::Internal("test_error")));
+  TF_ASSERT_OK(coord_service_->ReportTaskError(
+      task_0_, absl::InternalError("test_error")));
 
   EXPECT_TRUE(absl::IsInternal(
       coord_service_->RecordHeartbeat(task_0_, incarnation_0_)));
@@ -1595,7 +1598,7 @@ TEST_F(CoordinateTwoTasksTest, UnavailableTaskCanReconnect) {
   TF_EXPECT_OK(coord_service_->RegisterTask(task_0_, incarnation_0_));
 
   TF_ASSERT_OK(coord_service_->ReportTaskError(
-      task_0_, MakeCoordinationError(errors::Unavailable("test_error"))));
+      task_0_, MakeCoordinationError(absl::UnavailableError("test_error"))));
 
   TF_EXPECT_OK(coord_service_->RegisterTask(task_0_, incarnation_0_new_));
 }

From e8372d5129847773fe2f019106b35fbb89d35306 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Fri, 17 May 2024 18:51:22 -0700
Subject: [PATCH 377/478] Adopts the syntax
 'xla_tpu_auto_spmd_partitioning_memory_budget_ratio=-1.2' where the new '-'
 sign indicates that we are disabling the soft memory constraint.

PiperOrigin-RevId: 634942516
---
 .../experimental/auto_sharding/auto_sharding.cc    | 14 +++++++++++---
 .../auto_sharding/auto_sharding_option.h           |  5 +++++
 2 files changed, 16 insertions(+), 3 deletions(-)

diff --git a/third_party/xla/xla/hlo/experimental/auto_sharding/auto_sharding.cc b/third_party/xla/xla/hlo/experimental/auto_sharding/auto_sharding.cc
index 1564d1834f3cc1..5cba98f86637db 100644
--- a/third_party/xla/xla/hlo/experimental/auto_sharding/auto_sharding.cc
+++ b/third_party/xla/xla/hlo/experimental/auto_sharding/auto_sharding.cc
@@ -19,6 +19,7 @@ limitations under the License.
 #include <climits>
 #include <cstddef>
 #include <cstdint>
+#include <cstdlib>
 #include <functional>
 #include <iterator>
 #include <limits>
@@ -94,7 +95,6 @@ namespace xla {
 namespace spmd {
 
 namespace {
-constexpr double kOverbudgetCoeff = 1e6;
 constexpr double kSaltiplier = 0.0;  // This value (0.0) disables salting.
 }  // namespace
 
@@ -1972,7 +1972,11 @@ AutoShardingSolverResult CallSolver(
   request.mutable_s_hint()->Add(s_hint.begin(), s_hint.end());
   request.mutable_solver_timeout()->set_solver_timeout_in_seconds(
       solver_timeout_in_seconds);
-  request.mutable_overbudget_coeff()->set_coeff(kOverbudgetCoeff);
+  // Only apply soft memory constraints if the overbudget coeff is nonnegative.
+  if (option.memory_overbudget_coeff >= 0.0) {
+    request.mutable_overbudget_coeff()->set_coeff(
+        option.memory_overbudget_coeff);
+  }
   request.set_crash_at_infinity_costs_check(!option.try_multiple_mesh_shapes);
   request.set_compute_iis(compute_iis);
   request.set_saltiplier(kSaltiplier);
@@ -4038,7 +4042,11 @@ absl::StatusOr<AutoShardingResult> AutoShardingImplementation::RunAutoSharding(
           << " = " << memory_lower_bound_gb * option_.memory_budget_ratio
           << " GB";
       option_.memory_budget_per_device =
-          memory_lower_bound * option_.memory_budget_ratio;
+          memory_lower_bound * std::abs(option_.memory_budget_ratio);
+      // TODO(b/341299984): Document this flag syntax, or automate the behavior.
+      if (option_.memory_budget_ratio < 0) {
+        option_.memory_overbudget_coeff = -1.0;  // Disables the soft constraint
+      }
     } else if (option_.memory_budget_per_device > 0) {
       option_.memory_budget_per_device = original_memory_budget *
                                          original_device_mesh.num_elements() /
diff --git a/third_party/xla/xla/hlo/experimental/auto_sharding/auto_sharding_option.h b/third_party/xla/xla/hlo/experimental/auto_sharding/auto_sharding_option.h
index 2983d5ac9d7999..47b6985ca8d4da 100644
--- a/third_party/xla/xla/hlo/experimental/auto_sharding/auto_sharding_option.h
+++ b/third_party/xla/xla/hlo/experimental/auto_sharding/auto_sharding_option.h
@@ -26,6 +26,7 @@ namespace xla {
 
 static constexpr double kDeviceMeshAlpha = 1.0;
 static constexpr double kDeviceMeshBeta = 1.0;
+static constexpr double kOverbudgetCoeff = 1e6;
 
 // Options for the autosharding pass
 struct AutoShardingOption {
@@ -63,6 +64,10 @@ struct AutoShardingOption {
   // Enabled when memory_budget_per_device == 0;
   float memory_budget_ratio = 1.1;
 
+  // Controls the penalty associated with violating memory constraints; if
+  // negative, the memory budget is instead imposed as a hard constraint.
+  float memory_overbudget_coeff = kOverbudgetCoeff;
+
   // Overwrite the all gather cost with the input all reduce cost.
   bool force_override_all_gather_cost = false;
   double all_gather_cost = 0;

From 17c7b86dda238621de095e639bb7b74546eab38c Mon Sep 17 00:00:00 2001
From: "Jae H. Yoo" <jaeyoo@google.com>
Date: Fri, 17 May 2024 18:56:48 -0700
Subject: [PATCH 378/478] Fix tf.lite's lite_v2_test by using the new Keras 3
 API.

Wrapping TF quantization functions inside Keras 3 layers is required.

PiperOrigin-RevId: 634943217
---
 tensorflow/lite/python/lite_v2_test.py | 72 ++++++++++++++++++++------
 1 file changed, 57 insertions(+), 15 deletions(-)

diff --git a/tensorflow/lite/python/lite_v2_test.py b/tensorflow/lite/python/lite_v2_test.py
index 9a9a502b829e35..687c87d03108a8 100644
--- a/tensorflow/lite/python/lite_v2_test.py
+++ b/tensorflow/lite/python/lite_v2_test.py
@@ -391,9 +391,19 @@ def _createV2QATSavedModelWithFloatOpsAtEnd(self):
     """Create a simple QAT SavedModel that includes float ops at the end."""
     saved_model_dir = os.path.join(self.get_temp_dir(), 'qat_float_ops_at_end')
     input_tensor = tf.keras.layers.Input((32, 32, 128))
-    x = tf.quantization.fake_quant_with_min_max_args(input_tensor, -3.0, 3.0)
+
+    class _FakeQuantArgsLayer(tf.keras.layers.Layer):
+      """A fake quantization layer with fake_quant_with_min_max_args.
+
+      Keras 3 requires wrapping the tf function inside Keras layer.
+      """
+
+      def call(self, x):
+        return tf.quantization.fake_quant_with_min_max_args(x, -3.0, 3.0)
+
+    x = _FakeQuantArgsLayer()(input_tensor)
     x = tf.keras.layers.Conv2D(1, (3, 3))(x)
-    x = tf.quantization.fake_quant_with_min_max_args(x, -3.0, 3.0)
+    x = _FakeQuantArgsLayer()(x)
     # Exclude the quantization of the following Dense layer by not putting
     # fake quant layer after the dense layer.
     output_tensor = tf.keras.layers.Dense(1, activation='sigmoid')(x)
@@ -1607,10 +1617,19 @@ def _createV2QATSavedModel(self, shape):
     input_name = 'input'
     output_name = 'scores'
 
+    class _FakeQuantArgsLayer(tf.keras.layers.Layer):
+      """A fake quantization layer with fake_quant_with_min_max_args.
+
+      Keras 3 requires wrapping the tf function inside Keras layer.
+      """
+
+      def call(self, x):
+        return tf.quantization.fake_quant_with_min_max_args(x, -3.0, 3.0)
+
     input_tensor = tf.keras.layers.Input((32, 32, 128), name=input_name)
-    x = tf.quantization.fake_quant_with_min_max_args(input_tensor, -3.0, 3.0)
+    x = _FakeQuantArgsLayer()(input_tensor)
     x = tf.keras.layers.Conv2D(1, (3, 3))(x)
-    x = tf.quantization.fake_quant_with_min_max_args(x, -3.0, 3.0)
+    x = _FakeQuantArgsLayer()(x)
     scores = tf.keras.layers.Reshape((-1,), name=output_name)(x)
     model = tf.keras.Model(input_tensor, scores)
     model.save(saved_model_dir)
@@ -2679,9 +2698,18 @@ def testKerasFullyConnectedOutputShape3D(self):
         batch_size=1, shape=[3, 3], name='input_tensor', dtype=tf.float32
     )
 
-    x = tf.quantization.fake_quant_with_min_max_args(input_tensor, -3.0, 3.0)
+    class _FakeQuantArgsLayer(tf.keras.layers.Layer):
+      """A fake quantization layer with fake_quant_with_min_max_args.
+
+      Keras 3 requires wrapping the tf function inside Keras layer.
+      """
+
+      def call(self, x):
+        return tf.quantization.fake_quant_with_min_max_args(x, -3.0, 3.0)
+
+    x = _FakeQuantArgsLayer()(input_tensor)
     x = tf.keras.layers.Dense(3)(x)
-    x = tf.quantization.fake_quant_with_min_max_args(x, -3.0, 3.0)
+    x = _FakeQuantArgsLayer()(x)
     model = tf.keras.Model(input_tensor, x)
 
     model.compile(
@@ -2736,14 +2764,20 @@ def call(self, inputs):
             inputs, filters, [*inputs.shape[:-1], 24], 1
         )
 
+    class _FakeQuantVarsLayer(tf.keras.layers.Layer):
+      """A fake quantization layer with fake_quant_with_min_max_vars.
+
+      Keras 3 requires wrapping the tf function inside Keras layer.
+      """
+
+      def call(self, x):
+        return tf.quantization.fake_quant_with_min_max_vars(
+            x, -3.0, 3.0, narrow_range=True)
+
     inp = tf.keras.Input(shape=(6, 8, 48), batch_size=1)
-    x = tf.quantization.fake_quant_with_min_max_vars(
-        inp, -3.0, 3.0, narrow_range=True
-    )
+    x = _FakeQuantVarsLayer()(inp)
     x = QuantConv2DTransposed()(x)
-    x = tf.quantization.fake_quant_with_min_max_vars(
-        x, -3.0, 3.0, narrow_range=True
-    )
+    x = _FakeQuantVarsLayer()(x)
 
     model = tf.keras.Model(inp, x)
 
@@ -3512,10 +3546,18 @@ def call(self, inputs):
             result, -3.0, 3.0, narrow_range=True
         )
 
+    class _FakeQuantVarsLayer(tf.keras.layers.Layer):
+      """A fake quantization layer with fake_quant_with_min_max_vars.
+
+      Keras 3 requires wrapping the tf function inside Keras layer.
+      """
+
+      def call(self, x):
+        return tf.quantization.fake_quant_with_min_max_vars(
+            x, -3.0, 3.0, narrow_range=True)
+
     inp = tf.keras.Input(shape=(6, 8, 6), batch_size=1)
-    x = tf.quantization.fake_quant_with_min_max_vars(
-        inp, -3.0, 3.0, narrow_range=True
-    )
+    x = _FakeQuantVarsLayer()(inp)
     x = QuantConv2DTransposedWithBiasAndActivation()(x)
 
     model = tf.keras.Model(inp, x)

From b409682867f59e7b29703395f3ea7f68b6a5b320 Mon Sep 17 00:00:00 2001
From: Tongfei Guo <tongfei@google.com>
Date: Fri, 17 May 2024 19:07:04 -0700
Subject: [PATCH 379/478] [XLA:SPMD] Support partitioning kCall.

PiperOrigin-RevId: 634944833
---
 .../xla/xla/service/spmd/spmd_partitioner.cc  | 27 ++++++++++++++++-
 .../xla/xla/service/spmd/spmd_partitioner.h   |  1 +
 .../xla/service/spmd/spmd_partitioner_test.cc | 30 +++++++++++++++++++
 3 files changed, 57 insertions(+), 1 deletion(-)

diff --git a/third_party/xla/xla/service/spmd/spmd_partitioner.cc b/third_party/xla/xla/service/spmd/spmd_partitioner.cc
index 2d99310671f39b..1579d73344efd7 100644
--- a/third_party/xla/xla/service/spmd/spmd_partitioner.cc
+++ b/third_party/xla/xla/service/spmd/spmd_partitioner.cc
@@ -2266,6 +2266,30 @@ std::vector<ReplicaGroup> SpmdPartitioningVisitor::CreateReplicaGroups(
   return device_groups;
 }
 
+Status SpmdPartitioningVisitor::HandleCall(HloInstruction* hlo) {
+  std::vector<HloInstruction*> call_args;
+  HloComputation* computation = hlo->called_computations()[0];
+  for (int64_t i = 0; i < hlo->operand_count(); ++i) {
+    // Shardings of the computation parameter and its argument must be
+    // the same.
+    computation->parameter_instruction(i)->set_sharding(
+        hlo->operand(i)->sharding());
+    call_args.push_back(GetPartitionedHlo(hlo->operand(i)).hlo());
+  }
+
+  TF_RETURN_IF_ERROR(partitioner_
+                         ->PartitionComputation(computation, hlo->sharding(),
+                                                next_channel_id_, logger_,
+                                                call_graph_)
+                         .status());
+  SetPartitionedHlo(hlo, [&] {
+    return b_.AddInstruction(HloInstruction::CreateCall(
+        MakePartitionedShape(hlo->shape(), hlo->sharding()), call_args,
+        hlo->called_computations()[0]));
+  });
+  return OkStatus();
+}
+
 Status SpmdPartitioningVisitor::DefaultAction(HloInstruction* hlo) {
   if (hlo->HasSideEffect() && !hlo->sharding().HasUniqueDevice()) {
     return Unimplemented("Side-effect ops cannot be replicated: %s",
@@ -2343,7 +2367,8 @@ Status SpmdPartitioningVisitor::Preprocess(HloInstruction* hlo) {
       hlo->opcode() != HloOpcode::kParameter &&
       hlo->opcode() != HloOpcode::kWhile && hlo->opcode() != HloOpcode::kRng &&
       hlo->opcode() != HloOpcode::kOutfeed &&
-      hlo->opcode() != HloOpcode::kAllReduce) {
+      hlo->opcode() != HloOpcode::kAllReduce &&
+      hlo->opcode() != HloOpcode::kCall) {
     const bool has_manual_sharding =
         hlo->sharding().IsManual() ||
         (hlo->sharding().IsTuple() &&
diff --git a/third_party/xla/xla/service/spmd/spmd_partitioner.h b/third_party/xla/xla/service/spmd/spmd_partitioner.h
index c137a821c333a8..836250bffb6ba8 100644
--- a/third_party/xla/xla/service/spmd/spmd_partitioner.h
+++ b/third_party/xla/xla/service/spmd/spmd_partitioner.h
@@ -555,6 +555,7 @@ class SpmdPartitioningVisitor : public DfsHloVisitorWithDefault {
   Status DefaultAction(HloInstruction* hlo) override;
   Status HandleAllReduce(HloInstruction* hlo) override;
   Status HandleBroadcast(HloInstruction* hlo) override;
+  Status HandleCall(HloInstruction* hlo) override;
   Status HandleConstant(HloInstruction* hlo) override;
   Status HandleCustomCall(HloInstruction* hlo) override;
   Status HandleDot(HloInstruction* hlo) override;
diff --git a/third_party/xla/xla/service/spmd/spmd_partitioner_test.cc b/third_party/xla/xla/service/spmd/spmd_partitioner_test.cc
index e542984ae13152..85352c6451247d 100644
--- a/third_party/xla/xla/service/spmd/spmd_partitioner_test.cc
+++ b/third_party/xla/xla/service/spmd/spmd_partitioner_test.cc
@@ -233,6 +233,36 @@ ENTRY entry {
           op::Shape("s32[1,3]")));
 }
 
+TEST_P(SpmdPartitioningTest, PartitionCall) {
+  absl::string_view hlo_string = R"(
+HloModule jit_f
+
+g {
+  Arg_0.6 = s32[8,2]{1,0} parameter(0), sharding={devices=[2,2]<=[4]}
+  constant.0 = s32[] constant(2), sharding={replicated}
+  broadcast.0 = s32[8,2]{1,0} broadcast(constant.0), dimensions={}, sharding={devices=[2,2]<=[4]}
+  ROOT multiply.9 = s32[8,2]{1,0} multiply(Arg_0.6, broadcast.0), sharding={devices=[2,2]<=[4]}
+}
+
+ENTRY main {
+  Arg_0.1 = s32[8,2]{1,0} parameter(0), sharding={devices=[2,2]<=[4]}
+  constant.1 = s32[] constant(3), sharding={replicated}
+  broadcast.1 = s32[8,2]{1,0} broadcast(constant.1), dimensions={}, sharding={devices=[2,2]<=[4]}
+  multiply.4 = s32[8,2]{1,0} multiply(Arg_0.1, broadcast.1), sharding={devices=[2,2]<=[4]}
+  ROOT call = s32[8,2]{1,0} call(multiply.4), to_apply=g, sharding={devices=[2,2]<=[4]}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_DEFAULT","device_type":"DEVICE_TYPE_HOST","used_scoped_memory_configs":[]}
+})";
+  TF_ASSERT_OK_AND_ASSIGN(auto module,
+                          PartitionComputation(hlo_string, /*num_devices=*/4));
+  VLOG(1) << module->ToString();
+  HloInstruction* root = module->entry_computation()->root_instruction();
+  EXPECT_THAT(root, AllOf(op::Call(), op::Shape("s32[4,1]")));
+  HloInstruction* call_comp_root =
+      root->called_computations()[0]->root_instruction();
+  EXPECT_THAT(call_comp_root, AllOf(op::Multiply(op::Parameter(0),
+                                                 op::Broadcast(op::Constant())),
+                                    op::Shape("s32[4,1]")));
+}
+
 TEST_P(SpmdPartitioningTest, TiledToReplicated) {
   absl::string_view hlo_string = R"(
 HloModule module

From d3f2cdbcc66683191fee58f06235d4539d3d1209 Mon Sep 17 00:00:00 2001
From: Yifan Jiang <yifanjiang@google.com>
Date: Fri, 17 May 2024 19:08:01 -0700
Subject: [PATCH 380/478] In PjRtStreamExecutorBuffer::Delete, fix a bug
 causing memory corruption.

The bug is that we don't wait for events using the buffer on the Compute stream, and also the definition events. This can cause a race condition where the buffer is deleted before it is read/written, which can lead to memory corruptions.

For events on the compute stream, the fix is to schedule the deallocation on the compute stream.
For the definition events, the fix is to wait for them on a stream in a stream pool if they are not defined on the compute stream.

Added a fixed-size stream pool to avoid the overhead and deadlock that could be introduced by BorrowStreamFromPool, since we previously saw deadlocks between the cuStreamCreate and other cuda calls like cuMemHostAlloc.

This change probably wouldn't be needed if all writes are scheduled on the compute stream, but that is not the case today. We have a lot of usages of the h2d and d2d stream.

PiperOrigin-RevId: 634944963
---
 .../xla/xla/pjrt/local_device_state.cc        | 13 +++
 third_party/xla/xla/pjrt/local_device_state.h | 10 ++-
 .../xla/pjrt/pjrt_stream_executor_client.cc   | 80 ++++++++++++++++---
 .../xla/pjrt/pjrt_stream_executor_client.h    | 13 +--
 4 files changed, 100 insertions(+), 16 deletions(-)

diff --git a/third_party/xla/xla/pjrt/local_device_state.cc b/third_party/xla/xla/pjrt/local_device_state.cc
index fca99bcd391f58..bf52adbe6d1173 100644
--- a/third_party/xla/xla/pjrt/local_device_state.cc
+++ b/third_party/xla/xla/pjrt/local_device_state.cc
@@ -80,6 +80,10 @@ LocalDeviceState::LocalDeviceState(se::StreamExecutor* executor,
   for (int i = 0; i < num_device_to_device_streams; ++i) {
     device_to_device_streams_.emplace_back(create_stream());
   }
+  fixed_size_pool_usage_streams_.reserve(kNumFixedSizePoolUsageStreams);
+  for (int i = 0; i < kNumFixedSizePoolUsageStreams; ++i) {
+    fixed_size_pool_usage_streams_.emplace_back(create_stream());
+  }
   external_ready_event_streams_.reserve(kNumExternalReadyEventStreams);
   for (int i = 0; i < kNumExternalReadyEventStreams; ++i) {
     external_ready_event_streams_.emplace_back(create_stream());
@@ -167,6 +171,15 @@ se::Stream* LocalDeviceState::GetDeviceToDeviceStream() {
   return device_to_device_streams_.at(i).get();
 }
 
+se::Stream* LocalDeviceState::GetFixedSizePoolUsageStream() {
+  absl::MutexLock lock(&mu_);
+  int i = next_fixed_size_pool_usage_stream_;
+  next_fixed_size_pool_usage_stream_ =
+      (next_fixed_size_pool_usage_stream_ + 1) %
+      fixed_size_pool_usage_streams_.size();
+  return fixed_size_pool_usage_streams_.at(i).get();
+}
+
 se::Stream* LocalDeviceState::GetExternalReadyEventStream() {
   absl::MutexLock lock(&mu_);
   int i = next_external_ready_event_stream_;
diff --git a/third_party/xla/xla/pjrt/local_device_state.h b/third_party/xla/xla/pjrt/local_device_state.h
index deece8131acecb..9cc0db39776e66 100644
--- a/third_party/xla/xla/pjrt/local_device_state.h
+++ b/third_party/xla/xla/pjrt/local_device_state.h
@@ -136,6 +136,11 @@ class LocalDeviceState {
   // fashion amongst the available streams.
   se::Stream* GetDeviceToDeviceStream();
 
+  // Returns a usage stream. Allocates streams in a round-robin fashion amongst
+  // the available streams. When the overhead from BorrowStreamFromPool is too
+  // large for a use case, consider using this API instead.
+  se::Stream* GetFixedSizePoolUsageStream();
+
   // Return a stream that should be used to track when an externally-managed
   // buffer is ready. This is intended to support dlpack on GPU. Allocates
   // streams in a round-robin fashion amongst the available streams.
@@ -151,7 +156,7 @@ class LocalDeviceState {
   // Returns a vector of device to device streams.
   std::vector<se::Stream*> GetDeviceToDeviceStreams();
 
-  // Returns a stream from a pool. The stream is guaranteed not to have any
+  // Borrows a stream from a pool. The stream is guaranteed not to have any
   // currently outstanding work at its tail.
   std::unique_ptr<se::Stream> BorrowStreamFromPool();
   // Returns a stream to the pool. The caller must ensure the stream does not
@@ -213,15 +218,18 @@ class LocalDeviceState {
   std::unique_ptr<se::Stream> host_to_device_stream_;
   std::vector<std::unique_ptr<se::Stream>> device_to_host_streams_;
   std::vector<std::unique_ptr<se::Stream>> device_to_device_streams_;
+  std::vector<std::unique_ptr<se::Stream>> fixed_size_pool_usage_streams_;
   std::vector<std::unique_ptr<se::Stream>> external_ready_event_streams_;
 
   static constexpr int kNumDeviceToHostStreams = 4;
   static constexpr int kNumDeviceToDeviceStreams = 4;
+  static constexpr int kNumFixedSizePoolUsageStreams = 4;
   static constexpr int kNumExternalReadyEventStreams = 4;
 
   absl::Mutex mu_;
   int next_device_to_host_stream_ ABSL_GUARDED_BY(mu_) = 0;
   int next_device_to_device_stream_ ABSL_GUARDED_BY(mu_) = 0;
+  int next_fixed_size_pool_usage_stream_ ABSL_GUARDED_BY(mu_) = 0;
   int next_external_ready_event_stream_ ABSL_GUARDED_BY(mu_) = 0;
 
   std::random_device prng_seed_device_ ABSL_GUARDED_BY(mu_);
diff --git a/third_party/xla/xla/pjrt/pjrt_stream_executor_client.cc b/third_party/xla/xla/pjrt/pjrt_stream_executor_client.cc
index 3f2ab8a28e0a54..efc10c37ce5d5b 100644
--- a/third_party/xla/xla/pjrt/pjrt_stream_executor_client.cc
+++ b/third_party/xla/xla/pjrt/pjrt_stream_executor_client.cc
@@ -1486,32 +1486,82 @@ PjRtStreamExecutorBuffer::Release(bool wait_for_operations_to_complete) {
   } else {
     if (local_device_state->allocation_model() ==
         LocalDeviceState::kComputeSynchronized) {
-      std::unique_ptr<se::Stream> block_stream;
+      se::Stream* block_stream = nullptr;
       for (const auto& stream_and_event : events) {
+        VLOG(2)
+            << "Checking whether need to wait for stream_and_event: stream: "
+            << stream_and_event.stream
+            << "; event: " << stream_and_event.event.get()
+            << "; reference_held: " << stream_and_event.reference_held
+            << "; is_predetermined_error: "
+            << stream_and_event.event->IsPredeterminedError();
         // We only need to do something for events that didn't already acquire a
         // reference to the buffer, and also which the compute stream didn't
         // already wait for. Based on our heuristics this rare case should only
         // occur when a buffer was copied to a device and then never used there.
         // In that case we get a new stream and use it to hold onto a reference
         // to the buffer until the events are complete.
+        //
+        // It is also important that we check IsPredeterminedError before
+        // checking DefinedOn(compute_stream) because otherwise DefinedOn would
+        // indefinitely wait since the event is never recorded when the buffer
+        // is predetermined error.
         if (!stream_and_event.event->IsPredeterminedError() &&
             !stream_and_event.reference_held &&
             !stream_and_event.event->DefinedOn(
                 local_device_state->compute_stream()) &&
             !stream_and_event.event->IsComplete()) {
           if (block_stream == nullptr) {
-            block_stream = local_device_state->BorrowStreamFromPool();
+            block_stream = local_device_state->GetFixedSizePoolUsageStream();
           }
-          stream_and_event.event->WaitForEventOnStream(block_stream.get());
+          VLOG(2) << "Waiting for stream_and_event: stream: "
+                  << stream_and_event.stream
+                  << "; event: " << stream_and_event.event.get()
+                  << "; reference_held: " << stream_and_event.reference_held
+                  << "; is_predetermined_error: "
+                  << stream_and_event.event->IsPredeterminedError();
+          stream_and_event.event->WaitForEventOnStream(block_stream);
+        }
+      }
+      for (const auto& definition_event : device_buffer->definition_events()) {
+        VLOG(2) << "Checking whether need to wait for definition_event: "
+                << definition_event.get() << "; is_predetermined_error: "
+                << definition_event->IsPredeterminedError();
+        // Here we wait for the definition events to complete on block_stream as
+        // well, if they are not on the compute stream and not also recorded as
+        // usage events.
+        //
+        // It is also important that we check IsPredeterminedError before
+        // checking DefinedOn(compute_stream) because otherwise DefinedOn would
+        // indefinitely wait since the event is never recorded when the buffer
+        // is predetermined error.
+        //
+        // Since it's possible that definition_event.SetSequencingEvent()
+        // is called on a different host thread than this host thread, when in
+        // future more conditions are added to this check, we should be careful
+        // about whether we put them before the DefinedOn check or after it.
+        // For example, we shouldn't add an IsDefined() check before the
+        // DefinedOn() check here because that could potentially cause a
+        // shortcut where we don't wait for
+        // definition_event.SetSequencingEvent() on the other thread and
+        // eventually cause memory corruption.
+        if (!definition_event->IsPredeterminedError() &&
+            !definition_event->DefinedOn(
+                local_device_state->compute_stream()) &&
+            !definition_event->IsComplete()) {
+          if (block_stream == nullptr) {
+            block_stream = local_device_state->GetFixedSizePoolUsageStream();
+          }
+          VLOG(2) << "Waiting for definition_event: " << definition_event.get()
+                  << "; is_predetermined_error: "
+                  << definition_event->IsPredeterminedError();
+          definition_event->WaitForEventOnStream(block_stream);
         }
       }
       if (block_stream != nullptr) {
-        se::Stream* block_stream_ptr = block_stream.release();
         TF_RETURN_IF_ERROR(local_device_state->ThenExecuteCallback(
-            block_stream_ptr,
-            [device_buffer, block_stream_ptr, local_device_state]() {
-              local_device_state->ReturnStreamToPool(
-                  std::unique_ptr<se::Stream>(block_stream_ptr));
+            block_stream, [device_buffer]() {
+              // Drops device_buffer shared pointer.
             }));
       }
     }
@@ -1521,8 +1571,20 @@ PjRtStreamExecutorBuffer::Release(bool wait_for_operations_to_complete) {
 
 void PjRtStreamExecutorBuffer::Delete() {
   VLOG(1) << "PjRtStreamExecutorBuffer::Delete";
+
   // When wait_for_reads_to_complete is false, Release should never fail.
-  TF_CHECK_OK(Release(/*wait_for_operations_to_complete=*/false).status());
+  absl::StatusOr<std::shared_ptr<TrackedDeviceBuffer>> tracked_device_buffer =
+      Release(/*wait_for_operations_to_complete=*/false);
+
+  // The only usage events that
+  // Release(/*wait_for_operations_to_complete=*/false) doesn't wait for are
+  // events defined on the compute stream. So we schedule the deallocation on
+  // the compute stream so that there wouldn't be use-after-free usages on the
+  // compute stream.
+  TF_CHECK_OK(tracked_device_buffer.status());
+  TF_CHECK_OK(device_->local_device_state()->ThenRelease(
+      device_->local_device_state()->compute_stream(),
+      std::move(tracked_device_buffer.value())));
 }
 
 bool PjRtStreamExecutorBuffer::IsDeleted() {
diff --git a/third_party/xla/xla/pjrt/pjrt_stream_executor_client.h b/third_party/xla/xla/pjrt/pjrt_stream_executor_client.h
index cc879e4f773b1e..c2f55e743cf910 100644
--- a/third_party/xla/xla/pjrt/pjrt_stream_executor_client.h
+++ b/third_party/xla/xla/pjrt/pjrt_stream_executor_client.h
@@ -798,12 +798,13 @@ class PjRtStreamExecutorBuffer : public PjRtBuffer {
   // Similar to Delete, drops the buffer's reference to its associated device
   // memory, leaving the buffer in an invalid state, but returns the
   // TrackedDeviceBuffer rather than freeing the device memory, so that another
-  // framework can take ownership of it. The buffer returned from Release may
-  // be safely dropped at any time even if it still has pending async
-  // operations. The client should call GetReadyFuture()->Await() before calling
-  // Release with wait_for_operations_to_complete=false, to ensure that the host
-  // has synchronized past any outstanding write operations to the buffer. If
-  // wait_for_operations_to_complete=true the host will block until any
+  // framework can take ownership of it.
+  //
+  // When called with wait_for_operations_to_complete=false, the buffer returned
+  // from Release should be dropped on the compute stream, since the only events
+  // that Release doesn't wait for are events defined on the compute stream.
+  //
+  // If wait_for_operations_to_complete=true, the host will block until any
   // potentially outstanding asynchronous operations have completed before
   // returning, in which case it is safe to read or mutate the returned buffer.
   // If the buffer was shared via an external reference it is the client's

From 4a6548c44420d2792d1e33d10587e1f89ee8e132 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Fri, 17 May 2024 19:24:31 -0700
Subject: [PATCH 381/478] [xla:ffi] Fix return types in `xla::ffi::Expected<T,
 E>`.

PiperOrigin-RevId: 634947036
---
 third_party/xla/xla/ffi/api/api.h | 35 ++++++++++++++++++++++---------
 1 file changed, 25 insertions(+), 10 deletions(-)

diff --git a/third_party/xla/xla/ffi/api/api.h b/third_party/xla/xla/ffi/api/api.h
index 655d0ba1b656d5..d0299022a630b2 100644
--- a/third_party/xla/xla/ffi/api/api.h
+++ b/third_party/xla/xla/ffi/api/api.h
@@ -820,18 +820,32 @@ class Unexpected;
 template <typename T, typename E>
 class Expected {
  public:
-  Expected(T value) : data_(std::move(value)) {}  // NOLINT
-  Expected(Unexpected<E> u);                      // NOLINT
+  constexpr Expected(T value) : data_(std::move(value)) {}  // NOLINT
+  constexpr Expected(Unexpected<E> u);                      // NOLINT
 
-  operator bool() const {  // NOLINT
+  constexpr operator bool() const {  // NOLINT
     return has_value();
   }
-  T operator*() const { return value(); }
-  T* operator->() const { return &value(); }
 
-  bool has_value() const { return std::holds_alternative<T>(data_); }
-  T value() const { return std::get<T>(data_); }
-  E error() const { return std::get<E>(data_); }
+  constexpr T& operator*() & { return value(); }
+  constexpr const T& operator*() const& { return value(); }
+  constexpr T&& operator*() && { return std::move(value()); }
+  constexpr const T& operator*() const&& { return std::move(value()); }
+
+  constexpr T* operator->() { return &value(); }
+  constexpr const T* operator->() const { return &value(); }
+
+  constexpr bool has_value() const { return std::holds_alternative<T>(data_); }
+
+  constexpr T& value() & { return std::get<T>(data_); }
+  constexpr const T& value() const& { return std::get<T>(data_); }
+  constexpr T&& value() && { return std::get<T>(std::move(data_)); }
+  constexpr const T& value() const&& { return std::get<T>(std::move(data_)); }
+
+  constexpr E& error() & { return std::get<E>(data_); }
+  constexpr const E& error() const& { return std::get<E>(data_); }
+  constexpr E&& error() && { return std::get<E>(std::move(data_)); }
+  constexpr const E&& error() const&& { return std::get<E>(std::move(data_)); }
 
  private:
   std::variant<T, E> data_;
@@ -840,7 +854,7 @@ class Expected {
 template <typename E>
 class Unexpected {
  public:
-  explicit Unexpected(E error) : error_(std::move(error)) {}
+  explicit constexpr Unexpected(E error) : error_(std::move(error)) {}
 
  private:
   template <typename, typename>
@@ -852,7 +866,8 @@ class Unexpected {
 Unexpected(const char*) -> Unexpected<std::string>;
 
 template <typename T, typename E>
-Expected<T, E>::Expected(Unexpected<E> u) : data_(std::move(u.error_)) {}
+constexpr Expected<T, E>::Expected(Unexpected<E> u)
+    : data_(std::move(u.error_)) {}
 
 //===----------------------------------------------------------------------===//
 // Type-safe wrapper for accessing a variable number of arguments.

From 3c5c0edbe25d51e7652542067691a1d8a22bb6be Mon Sep 17 00:00:00 2001
From: Ce Zheng <zce@google.com>
Date: Fri, 17 May 2024 20:55:53 -0700
Subject: [PATCH 382/478] [XLA] Support more cases in
 IotaTileAssignment::Transpose.

PiperOrigin-RevId: 634960053
---
 third_party/xla/xla/hlo/ir/BUILD              |  2 +
 third_party/xla/xla/hlo/ir/tile_assignment.cc | 66 ++++++++++++++++++-
 .../xla/xla/tests/tile_assignment_test.cc     | 41 ++++++++++++
 3 files changed, 108 insertions(+), 1 deletion(-)

diff --git a/third_party/xla/xla/hlo/ir/BUILD b/third_party/xla/xla/hlo/ir/BUILD
index d3ca0ed792e56c..8f15554a314061 100644
--- a/third_party/xla/xla/hlo/ir/BUILD
+++ b/third_party/xla/xla/hlo/ir/BUILD
@@ -212,6 +212,8 @@ cc_library(
         "//xla:printer",
         "//xla:util",
         "@com_google_absl//absl/algorithm:container",
+        "@com_google_absl//absl/container:inlined_vector",
+        "@com_google_absl//absl/log:check",
         "@com_google_absl//absl/types:span",
     ],
 )
diff --git a/third_party/xla/xla/hlo/ir/tile_assignment.cc b/third_party/xla/xla/hlo/ir/tile_assignment.cc
index 2db424833bc85f..e3bd49952c198b 100644
--- a/third_party/xla/xla/hlo/ir/tile_assignment.cc
+++ b/third_party/xla/xla/hlo/ir/tile_assignment.cc
@@ -22,6 +22,8 @@ limitations under the License.
 #include <string>
 #include <utility>
 
+#include "absl/container/inlined_vector.h"
+#include "absl/log/check.h"
 #include "absl/types/span.h"
 #include "xla/array.h"
 #include "xla/util.h"
@@ -129,6 +131,52 @@ TransposeKind GetTransposeKind(absl::Span<const int64_t> dims,
   return kind;
 }
 
+// Try to split canonicalized reshape_dims and transpose_perm so that
+// reshape_dims transposed with split_transpose_perm will result in
+// non_one_dims (must not contain 1s, as the name suggests), returns true iff
+// such a split is found, and fills the results in split_transpose_perm.
+bool TryDecanonicalize(absl::Span<const int64_t> non_one_dims,
+                       absl::Span<const int64_t> reshape_dims,
+                       absl::Span<const int> transpose_perm,
+                       absl::InlinedVector<int, 6>& split_transpose_perm) {
+  DCHECK_LT(reshape_dims.size(), non_one_dims.size());
+  DCHECK_EQ(transpose_perm.size(), reshape_dims.size());
+  split_transpose_perm.clear();
+  absl::InlinedVector<int, 6> split_counts(reshape_dims.size());
+  int non_one_idx = 0;
+  for (int i = 0; i < reshape_dims.size() && non_one_idx < non_one_dims.size();
+       ++i) {
+    int split_dim = transpose_perm[i];
+    int64_t target = reshape_dims[transpose_perm[i]];
+    int64_t cand = non_one_dims[non_one_idx];
+    int start_idx = non_one_idx;
+    while (target % cand == 0) {
+      target /= cand;
+      if (++non_one_idx >= non_one_dims.size()) {
+        break;
+      }
+      cand = non_one_dims[non_one_idx];
+    }
+    if (target != 1) {
+      return false;
+    }
+    split_counts[split_dim] = non_one_idx - start_idx;
+  }
+  absl::InlinedVector<int, 6> old_to_new_dims(split_counts.size());
+  for (int i = 1; i < old_to_new_dims.size(); ++i) {
+    old_to_new_dims[i] = old_to_new_dims[i - 1] + split_counts[i - 1];
+  }
+  split_transpose_perm.reserve(non_one_dims.size());
+  for (int i = 0; i < old_to_new_dims.size(); ++i) {
+    const int old_dim = transpose_perm[i];
+    for (int j = 0; j < split_counts[old_dim]; ++j) {
+      split_transpose_perm.push_back(old_to_new_dims[old_dim] + j);
+    }
+  }
+  CHECK_EQ(split_transpose_perm.size(), non_one_dims.size());
+  return true;
+}
+
 }  // namespace
 
 /*static*/ IotaTileAssignment IotaTileAssignment::Create(
@@ -245,7 +293,23 @@ std::optional<IotaTileAssignment> IotaTileAssignment::Transpose(
     CHECK_EQ(reshape_ndims_, new_perm.size());
     return IotaTileAssignment::Create(new_dims, reshape_dims, new_perm);
   }
-  // TODO(b/281892190): Handle remaining patterns and remove nullopt path.
+  absl::InlinedVector<int, 6> split_transpose_perm;
+  if (TryDecanonicalize(non_one_dims, reshape_dims, transpose_perm,
+                        split_transpose_perm)) {
+    absl::InlinedVector<int, 6> new_perm;
+    new_perm.reserve(non_one_dims.size());
+    for (int i = 0; i < ndims_; ++i) {
+      if (dims[perm[i]] == 1) continue;
+      new_perm.push_back(split_transpose_perm[one_to_non_one[perm[i]]]);
+    }
+    absl::InlinedVector<int64_t, 6> new_reshape_dims(
+        split_transpose_perm.size());
+    for (int i = 0; i < non_one_dims.size(); ++i) {
+      new_reshape_dims[split_transpose_perm[i]] = non_one_dims[i];
+    }
+    return IotaTileAssignment::Create(new_dims, new_reshape_dims, new_perm);
+  }
+  // TODO(b/341371396): Handle remaining patterns and remove nullopt path.
   return std::nullopt;
 }
 
diff --git a/third_party/xla/xla/tests/tile_assignment_test.cc b/third_party/xla/xla/tests/tile_assignment_test.cc
index d28fb6e754964c..a0dc21b1f38599 100644
--- a/third_party/xla/xla/tests/tile_assignment_test.cc
+++ b/third_party/xla/xla/tests/tile_assignment_test.cc
@@ -281,6 +281,47 @@ TEST_P(FormattedTileAssignmentTest, TransposeIotaTileWithDegernateDims) {
       ElementsAre(0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15));
 }
 
+TEST_P(FormattedTileAssignmentTest,
+       TransposeIotaTileSplittingCanonicalizedReshapeDims) {
+  TileAssignment tile({8, 2, 16}, {16, 16}, {1, 0});
+  if (ShouldConvertToV1()) {
+    tile = TileAssignment(tile.shared_array());
+  }
+  TileAssignment xposed = tile.Transpose({0, 2, 1});
+  EXPECT_NE(xposed, tile);
+  EXPECT_EQ(xposed, TileAssignment({8, 16, 2}, {16, 8, 2}, {1, 0, 2}));
+  EXPECT_EQ(xposed.num_dimensions(), 3);
+  EXPECT_EQ(xposed.dim(0), 8);
+  EXPECT_EQ(xposed.dim(1), 16);
+  EXPECT_EQ(xposed.dim(2), 2);
+  EXPECT_EQ(xposed(0, 0, 0), 0);
+  EXPECT_EQ(xposed({2, 7, 1}), 117);
+  EXPECT_EQ(xposed.iota().has_value(), !ShouldConvertToV1());
+  EXPECT_TRUE(xposed.UsesDevice(0));
+  EXPECT_TRUE(xposed.UsesDevice(255));
+  EXPECT_FALSE(xposed.UsesDevice(256));
+  EXPECT_THAT(
+      ToVectorUsingEach(xposed),
+      ElementsAre(
+          0, 1, 16, 17, 32, 33, 48, 49, 64, 65, 80, 81, 96, 97, 112, 113, 128,
+          129, 144, 145, 160, 161, 176, 177, 192, 193, 208, 209, 224, 225, 240,
+          241, 2, 3, 18, 19, 34, 35, 50, 51, 66, 67, 82, 83, 98, 99, 114, 115,
+          130, 131, 146, 147, 162, 163, 178, 179, 194, 195, 210, 211, 226, 227,
+          242, 243, 4, 5, 20, 21, 36, 37, 52, 53, 68, 69, 84, 85, 100, 101, 116,
+          117, 132, 133, 148, 149, 164, 165, 180, 181, 196, 197, 212, 213, 228,
+          229, 244, 245, 6, 7, 22, 23, 38, 39, 54, 55, 70, 71, 86, 87, 102, 103,
+          118, 119, 134, 135, 150, 151, 166, 167, 182, 183, 198, 199, 214, 215,
+          230, 231, 246, 247, 8, 9, 24, 25, 40, 41, 56, 57, 72, 73, 88, 89, 104,
+          105, 120, 121, 136, 137, 152, 153, 168, 169, 184, 185, 200, 201, 216,
+          217, 232, 233, 248, 249, 10, 11, 26, 27, 42, 43, 58, 59, 74, 75, 90,
+          91, 106, 107, 122, 123, 138, 139, 154, 155, 170, 171, 186, 187, 202,
+          203, 218, 219, 234, 235, 250, 251, 12, 13, 28, 29, 44, 45, 60, 61, 76,
+          77, 92, 93, 108, 109, 124, 125, 140, 141, 156, 157, 172, 173, 188,
+          189, 204, 205, 220, 221, 236, 237, 252, 253, 14, 15, 30, 31, 46, 47,
+          62, 63, 78, 79, 94, 95, 110, 111, 126, 127, 142, 143, 158, 159, 174,
+          175, 190, 191, 206, 207, 222, 223, 238, 239, 254, 255));
+}
+
 TEST_P(FormattedTileAssignmentTest, TransposeNoopIotaTile) {
   TileAssignment tile({4, 4}, {4, 4}, {1, 0});
   if (ShouldConvertToV1()) {

From 2d26ae911d88fd5e36bd53577cbd17adbd16925a Mon Sep 17 00:00:00 2001
From: Yin Zhang <yinzz@google.com>
Date: Fri, 17 May 2024 23:03:51 -0700
Subject: [PATCH 383/478] Add ReadFileTraceMetadata function to support trace
 viewer processes list reading

PiperOrigin-RevId: 634977692
---
 .../core/profiler/convert/trace_viewer/BUILD  |  1 +
 .../convert/trace_viewer/trace_events.cc      | 35 ++++++++++++++++++-
 .../convert/trace_viewer/trace_events.h       |  4 +++
 .../xla/third_party/tsl/tsl/lib/io/BUILD      |  1 +
 4 files changed, 40 insertions(+), 1 deletion(-)

diff --git a/tensorflow/core/profiler/convert/trace_viewer/BUILD b/tensorflow/core/profiler/convert/trace_viewer/BUILD
index 8b2ff03a343174..bdb65463505483 100644
--- a/tensorflow/core/profiler/convert/trace_viewer/BUILD
+++ b/tensorflow/core/profiler/convert/trace_viewer/BUILD
@@ -118,6 +118,7 @@ cc_library(
         "@com_google_absl//absl/status",
         "@com_google_absl//absl/strings",
         "@com_google_absl//absl/types:optional",
+        "@local_tsl//tsl/lib/io:iterator",
         "@local_tsl//tsl/platform:status",
         "@local_tsl//tsl/profiler/utils:timespan",
     ],
diff --git a/tensorflow/core/profiler/convert/trace_viewer/trace_events.cc b/tensorflow/core/profiler/convert/trace_viewer/trace_events.cc
index 35ac11941f4d0e..c514988f8f2b12 100644
--- a/tensorflow/core/profiler/convert/trace_viewer/trace_events.cc
+++ b/tensorflow/core/profiler/convert/trace_viewer/trace_events.cc
@@ -28,6 +28,7 @@ limitations under the License.
 #include "absl/base/internal/endian.h"
 #include "absl/log/check.h"
 #include "absl/log/log.h"
+#include "absl/status/status.h"
 #include "absl/strings/string_view.h"
 #include "tensorflow/core/platform/file_system.h"
 #include "tensorflow/core/platform/macros.h"
@@ -36,6 +37,7 @@ limitations under the License.
 #include "tensorflow/core/profiler/convert/trace_viewer/trace_viewer_visibility.h"
 #include "tensorflow/core/profiler/protobuf/trace_events.pb.h"
 #include "tensorflow/core/profiler/protobuf/trace_events_raw.pb.h"
+#include "tsl/lib/io/iterator.h"
 #include "tsl/lib/io/table.h"
 #include "tsl/lib/io/table_builder.h"
 #include "tsl/lib/io/table_options.h"
@@ -170,6 +172,36 @@ std::vector<std::vector<const TraceEvent*>> GetEventsByLevel(
   return events_by_level;
 }
 
+absl::Status ReadFileTraceMetadata(std::string& filepath, Trace* trace) {
+  // 1. Open the file.
+  uint64_t file_size;
+  TF_RETURN_IF_ERROR(tsl::Env::Default()->GetFileSize(filepath, &file_size));
+
+  tsl::FileSystem* file_system;
+  TF_RETURN_IF_ERROR(
+      tsl::Env::Default()->GetFileSystemForFile(filepath, &file_system));
+
+  std::unique_ptr<tsl::RandomAccessFile> file;
+  TF_RETURN_IF_ERROR(file_system->NewRandomAccessFile(filepath, &file));
+
+  tsl::table::Options options;
+  options.block_size = 20 * 1024 * 1024;
+  tsl::table::Table* table = nullptr;
+  TF_RETURN_IF_ERROR(
+      tsl::table::Table::Open(options, file.get(), file_size, &table));
+  std::unique_ptr<tsl::table::Table> table_deleter(table);
+
+  std::unique_ptr<tsl::table::Iterator> iterator(table->NewIterator());
+  if (iterator == nullptr) return absl::UnknownError("Could not open table");
+
+  // 2. Read the metadata.
+  iterator->SeekToFirst();
+  if (!ReadTraceMetadata(iterator.get(), kTraceMetadataKey, trace)) {
+    return absl::UnknownError("Could not parse Trace proto");
+  }
+  return absl::OkStatus();
+}
+
 // Store the contents of this container in an sstable file. The format is as
 // follows:
 //
@@ -265,7 +297,8 @@ tsl::Status DoLoadFromLevelDbTable(
   // Read the metadata.
   iterator->SeekToFirst();
   if (!ReadTraceMetadata(iterator.get(), kTraceMetadataKey, &trace)) {
-    return tsl::errors::Unknown("Could not parse Trace proto");
+    return absl::UnknownError(
+        "Could not parse Trace proto to read trace metadata");
   }
 
   if (filter) filter->SetUp(trace);
diff --git a/tensorflow/core/profiler/convert/trace_viewer/trace_events.h b/tensorflow/core/profiler/convert/trace_viewer/trace_events.h
index 68a27c22f86b67..fb2ae7766a6c58 100644
--- a/tensorflow/core/profiler/convert/trace_viewer/trace_events.h
+++ b/tensorflow/core/profiler/convert/trace_viewer/trace_events.h
@@ -37,6 +37,7 @@ limitations under the License.
 #include "tensorflow/core/profiler/lib/context_types.h"
 #include "tensorflow/core/profiler/protobuf/task.pb.h"
 #include "tensorflow/core/profiler/protobuf/trace_events.pb.h"
+#include "tsl/lib/io/table.h"
 #include "tsl/platform/errors.h"
 #include "tsl/platform/file_system.h"
 #include "tsl/platform/status.h"
@@ -65,6 +66,9 @@ tsl::Status DoLoadFromLevelDbTable(
     const std::function<TraceEvent*(const TraceEvent&)>& copy_event_to_arena,
     const std::function<void(TraceEvent*)>& add_arena_event);
 
+// Reads the trace metadata from a file with given path
+absl::Status ReadFileTraceMetadata(std::string& filepath, Trace* trace);
+
 std::vector<std::vector<const TraceEvent*>> GetEventsByLevel(
     const Trace& trace, std::vector<TraceEvent*>& events);
 
diff --git a/third_party/xla/third_party/tsl/tsl/lib/io/BUILD b/third_party/xla/third_party/tsl/tsl/lib/io/BUILD
index e92e2896dc9241..c103dcfdc5a417 100644
--- a/third_party/xla/third_party/tsl/tsl/lib/io/BUILD
+++ b/third_party/xla/third_party/tsl/tsl/lib/io/BUILD
@@ -18,6 +18,7 @@ package(
         "//tensorflow/core:__pkg__",
         "//tensorflow/core/lib/io:__subpackages__",
         "//tsl/profiler:__subpackages__",
+        "//tensorflow/core/profiler:__subpackages__",
     ]),
     licenses = ["notice"],
 )

From 52da3cb2cc44de19a54e0ae79c41608cafacdc0e Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Fri, 17 May 2024 23:25:10 -0700
Subject: [PATCH 384/478] Automated Code Change

PiperOrigin-RevId: 634980543
---
 third_party/xla/xla/pjrt/distributed/topology_util.cc | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/third_party/xla/xla/pjrt/distributed/topology_util.cc b/third_party/xla/xla/pjrt/distributed/topology_util.cc
index e2fbd874756223..40a9d63b8cadea 100644
--- a/third_party/xla/xla/pjrt/distributed/topology_util.cc
+++ b/third_party/xla/xla/pjrt/distributed/topology_util.cc
@@ -78,7 +78,7 @@ static std::string GetGlobalTopologyKey(std::string_view platform) {
 static absl::StatusOr<std::vector<LocalTopologyProto>> GetAllLocalTopologies(
     std::string_view platform, int num_nodes, KeyValueStoreInterface* kv_store,
     absl::Duration timeout) {
-  std::vector<StatusOr<std::string>> local_topology_strs(num_nodes);
+  std::vector<absl::StatusOr<std::string>> local_topology_strs(num_nodes);
 
   // TODO(ezhulenev): Should a thread pool become a function argument?
   tsl::thread::ThreadPool thread_pool(

From 12dfa7e3d5b112df4b3dd81a5d779181996aca76 Mon Sep 17 00:00:00 2001
From: "Jae H. Yoo" <jaeyoo@google.com>
Date: Fri, 17 May 2024 23:53:18 -0700
Subject: [PATCH 385/478] Fix keras model saving error into TF SavedModel
 format.

Now that keras 3 doesn't support model saving format TF SavedModel directly via `model.save()`, we need to use `tf.saved_model.save()` to save the model.

PiperOrigin-RevId: 634985205
---
 tensorflow/lite/python/util_test.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorflow/lite/python/util_test.py b/tensorflow/lite/python/util_test.py
index 3a8b212115c376..5c09fd47171bd8 100644
--- a/tensorflow/lite/python/util_test.py
+++ b/tensorflow/lite/python/util_test.py
@@ -274,7 +274,7 @@ def _generate_integer_tflite_model(quantization_type=dtypes.int8,
     # Convert TF Model to an Integer Quantized TFLite Model
     converter = tf.lite.TFLiteConverter.from_keras_model(model)
   else:
-    model.save(saved_model_dir)
+    tf.saved_model.save(model, saved_model_dir)
     converter = tf.lite.TFLiteConverter.from_saved_model(saved_model_dir)
   converter.optimizations = {tf.lite.Optimize.DEFAULT}
 

From a254edd08a499dc8a5a8f572b08baeb13c2308a6 Mon Sep 17 00:00:00 2001
From: "Jae H. Yoo" <jaeyoo@google.com>
Date: Sat, 18 May 2024 00:01:40 -0700
Subject: [PATCH 386/478] Fix nightly build / cmake build error due to missing
 argument

PiperOrigin-RevId: 634986604
---
 tensorflow/lite/tools/evaluation/utils.cc | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/tensorflow/lite/tools/evaluation/utils.cc b/tensorflow/lite/tools/evaluation/utils.cc
index 98f46247ab1c5b..6628911f79441e 100644
--- a/tensorflow/lite/tools/evaluation/utils.cc
+++ b/tensorflow/lite/tools/evaluation/utils.cc
@@ -185,7 +185,9 @@ TfLiteDelegatePtr CreateHexagonDelegate(
 #endif  // TFLITE_ENABLE_HEXAGON
 
 #ifdef TFLITE_WITHOUT_XNNPACK
-TfLiteDelegatePtr CreateXNNPACKDelegate(int num_threads, bool force_fp16) {
+TfLiteDelegatePtr CreateXNNPACKDelegate(
+    int num_threads, bool force_fp16,
+    const char* experimental_weight_cache_file_path) {
   return tools::CreateNullDelegate();
 }
 #else  // !defined(TFLITE_WITHOUT_XNNPACK)

From 62593a118b1a1a402ee58d6e390a7d1165987b37 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Sat, 18 May 2024 02:02:15 -0700
Subject: [PATCH 387/478] Update GraphDef version to 1866.

PiperOrigin-RevId: 635005462
---
 tensorflow/core/public/version.h | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorflow/core/public/version.h b/tensorflow/core/public/version.h
index 79913c71835d8b..51c8ef32c1956a 100644
--- a/tensorflow/core/public/version.h
+++ b/tensorflow/core/public/version.h
@@ -108,7 +108,7 @@ limitations under the License.
 
 #define TF_GRAPH_DEF_VERSION_MIN_PRODUCER 0
 #define TF_GRAPH_DEF_VERSION_MIN_CONSUMER 0
-#define TF_GRAPH_DEF_VERSION 1865  // Updated: 2024/5/17
+#define TF_GRAPH_DEF_VERSION 1866  // Updated: 2024/5/18
 
 // Checkpoint compatibility versions (the versions field in SavedSliceMeta).
 //

From 80e063704b9e88e5fb3ec41216315528af15faa3 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Sat, 18 May 2024 02:02:22 -0700
Subject: [PATCH 388/478] compat: Update forward compatibility horizon to
 2024-05-18

PiperOrigin-RevId: 635005488
---
 tensorflow/python/compat/compat.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorflow/python/compat/compat.py b/tensorflow/python/compat/compat.py
index 4fc96bfdaa9f54..5dd1128fc1715a 100644
--- a/tensorflow/python/compat/compat.py
+++ b/tensorflow/python/compat/compat.py
@@ -29,7 +29,7 @@
 # This value changes every day with an automatic CL. It can be modified in code
 # via `forward_compatibility_horizon()` or with the environment variable
 # TF_FORWARD_COMPATIBILITY_DELTA_DAYS, which is added to the compatibility date.
-_FORWARD_COMPATIBILITY_HORIZON = datetime.date(2024, 5, 17)
+_FORWARD_COMPATIBILITY_HORIZON = datetime.date(2024, 5, 18)
 _FORWARD_COMPATIBILITY_DELTA_DAYS_VAR_NAME = "TF_FORWARD_COMPATIBILITY_DELTA_DAYS"
 _FORWARD_COMPATIBILITY_DATE_NUMBER = None
 

From 03e7962cc5e70e1dd78a0e4b3dc97f0af57bfc5b Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Sat, 18 May 2024 03:00:30 -0700
Subject: [PATCH 389/478] Add unit test coverage for IFRT call op kernel impl

PiperOrigin-RevId: 635013418
---
 tensorflow/core/tfrt/ifrt/testdata/BUILD      |   5 +-
 tensorflow/core/tfrt/kernels/BUILD            |  51 +++++
 .../tfrt/kernels/ifrt_program_ops_test.cc     | 177 ++++++++++++++++++
 3 files changed, 232 insertions(+), 1 deletion(-)
 create mode 100644 tensorflow/core/tfrt/kernels/ifrt_program_ops_test.cc

diff --git a/tensorflow/core/tfrt/ifrt/testdata/BUILD b/tensorflow/core/tfrt/ifrt/testdata/BUILD
index 948ce54ab983a7..0fcbd0de55f217 100644
--- a/tensorflow/core/tfrt/ifrt/testdata/BUILD
+++ b/tensorflow/core/tfrt/ifrt/testdata/BUILD
@@ -1,6 +1,9 @@
 package(
     # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
-    default_visibility = ["//tensorflow/core/tfrt/ifrt:__subpackages__"],
+    default_visibility = [
+        "//tensorflow/core/tfrt/ifrt:__subpackages__",
+        "//tensorflow/core/tfrt/kernels:__subpackages__",
+    ],
     licenses = ["notice"],
 )
 
diff --git a/tensorflow/core/tfrt/kernels/BUILD b/tensorflow/core/tfrt/kernels/BUILD
index fa342c7a6bcb17..39cfca154b506a 100644
--- a/tensorflow/core/tfrt/kernels/BUILD
+++ b/tensorflow/core/tfrt/kernels/BUILD
@@ -34,6 +34,57 @@ cc_library(
     alwayslink = 1,
 )
 
+tf_cc_test(
+    name = "ifrt_program_ops_test",
+    srcs = ["ifrt_program_ops_test.cc"],
+    data = [
+        "//tensorflow/core/tfrt/ifrt/testdata",
+    ],
+    tags = ["no_oss"],
+    deps = [
+        ":ifrt_program_ops",
+        "//tensorflow/compiler/mlir/tensorflow",
+        "//tensorflow/compiler/tf2xla:xla_helpers",
+        "//tensorflow/core:core_cpu_base",
+        "//tensorflow/core:framework",
+        "//tensorflow/core:test",
+        "//tensorflow/core:test_main",
+        "//tensorflow/core:testlib",
+        "//tensorflow/core/framework:tensor_matcher",
+        "//tensorflow/core/framework:types_proto_cc",
+        "//tensorflow/core/lib/gtl:cleanup",
+        "//tensorflow/core/platform:resource_loader",
+        "//tensorflow/core/platform:status",
+        "//tensorflow/core/tfrt/ifrt:ifrt_executable_registry",
+        "//tensorflow/core/tfrt/ifrt:ifrt_loaded_variable_registry",
+        "//tensorflow/core/tfrt/ifrt:ifrt_restore_tensor_registry",
+        "//tensorflow/core/tfrt/ifrt:ifrt_serving_core_selector",
+        "//tensorflow/core/tfrt/ifrt:ifrt_serving_executable",
+        "//tensorflow/core/tfrt/ifrt:tf_host_callback",
+        "//tensorflow/core/tfrt/ops:ifrt_program_ops_op_lib",
+        "@com_google_absl//absl/log:check",
+        "@com_google_absl//absl/status:statusor",
+        "@com_google_absl//absl/strings",
+        "@com_google_absl//absl/strings:string_view",
+        "@com_google_absl//absl/types:span",
+        "@com_google_googletest//:gtest_main",
+        "@llvm-project//mlir:AllPassesAndDialects",
+        "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:Parser",
+        "@local_tsl//tsl/framework:serving_device_selector",
+        "@local_tsl//tsl/framework/test_util:mock_serving_device_selector",
+        "@local_tsl//tsl/lib/core:status_test_util",
+        "@local_tsl//tsl/platform:env",
+        "@local_tsl//tsl/platform:status",
+        "@local_tsl//tsl/platform:statusor",
+        "@local_xla//xla/pjrt/cpu:cpu_client",
+        "@local_xla//xla/python/ifrt",
+        "@local_xla//xla/python/ifrt:test_util",
+        "@local_xla//xla/python/pjrt_ifrt",
+        "@tf_runtime//:hostcontext",
+    ],
+)
+
 cc_library(
     name = "stream_ops",
     srcs = ["stream_ops.cc"],
diff --git a/tensorflow/core/tfrt/kernels/ifrt_program_ops_test.cc b/tensorflow/core/tfrt/kernels/ifrt_program_ops_test.cc
new file mode 100644
index 00000000000000..c7a2783bccb7f0
--- /dev/null
+++ b/tensorflow/core/tfrt/kernels/ifrt_program_ops_test.cc
@@ -0,0 +1,177 @@
+/* Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+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.
+==============================================================================*/
+
+#include <cstdint>
+#include <memory>
+#include <string>
+#include <utility>
+#include <vector>
+
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+#include "absl/log/check.h"
+#include "absl/status/statusor.h"
+#include "absl/strings/str_cat.h"
+#include "absl/strings/string_view.h"
+#include "absl/types/span.h"
+#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
+#include "mlir/IR/DialectRegistry.h"  // from @llvm-project
+#include "mlir/IR/MLIRContext.h"  // from @llvm-project
+#include "mlir/IR/OwningOpRef.h"  // from @llvm-project
+#include "mlir/InitAllDialects.h"  // from @llvm-project
+#include "mlir/Parser/Parser.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/tensorflow/dialect_registration.h"
+#include "tensorflow/compiler/tf2xla/xla_helpers.h"
+#include "xla/pjrt/cpu/cpu_client.h"
+#include "xla/python/ifrt/array.h"
+#include "xla/python/ifrt/client.h"
+#include "xla/python/ifrt/test_util.h"
+#include "xla/python/pjrt_ifrt/pjrt_client.h"
+#include "tensorflow/core/framework/fake_input.h"
+#include "tensorflow/core/framework/node_def_builder.h"
+#include "tensorflow/core/framework/tensor.h"
+#include "tensorflow/core/framework/tensor_matcher.h"
+#include "tensorflow/core/framework/tensor_shape.h"
+#include "tensorflow/core/framework/tensor_testutil.h"
+#include "tensorflow/core/framework/types.pb.h"
+#include "tensorflow/core/kernels/ops_testutil.h"
+#include "tensorflow/core/lib/gtl/cleanup.h"
+#include "tensorflow/core/platform/resource_loader.h"
+#include "tensorflow/core/platform/status.h"
+#include "tensorflow/core/platform/test.h"
+#include "tensorflow/core/tfrt/ifrt/ifrt_executable_registry.h"
+#include "tensorflow/core/tfrt/ifrt/ifrt_loaded_variable_registry.h"
+#include "tensorflow/core/tfrt/ifrt/ifrt_restore_tensor_registry.h"
+#include "tensorflow/core/tfrt/ifrt/ifrt_serving_core_selector.h"
+#include "tensorflow/core/tfrt/ifrt/ifrt_serving_executable.h"
+#include "tensorflow/core/tfrt/ifrt/tf_host_callback.h"
+#include "tsl/framework/serving_device_selector.h"
+#include "tsl/framework/test_util/mock_serving_device_selector.h"
+#include "tsl/lib/core/status_test_util.h"
+#include "tsl/platform/env.h"
+#include "tsl/platform/status.h"
+#include "tsl/platform/statusor.h"
+#include "tsl/platform/threadpool.h"
+#include "tfrt/host_context/concurrent_work_queue.h"  // from @tf_runtime
+
+namespace tensorflow {
+namespace tfrt_stub {
+namespace {
+
+using tensorflow::ifrt_serving::IfrtLoadedVariableRegistry;
+using tensorflow::ifrt_serving::IfrtRestoreTensorRegistry;
+using tensorflow::ifrt_serving::IfrtServingCoreSelector;
+using tensorflow::ifrt_serving::IfrtServingExecutable;
+using tensorflow::ifrt_serving::ServingExecutableRegistry;
+using tensorflow::test::AsTensor;
+using tensorflow::test::TensorEq;
+using ::testing::Return;
+
+const bool kUnused =
+    (xla::ifrt::test_util::RegisterClientFactory(
+         []() -> absl::StatusOr<std::shared_ptr<xla::ifrt::Client>> {
+           xla::CpuClientOptions options;
+           options.cpu_device_count = 4;
+           TF_ASSIGN_OR_RETURN(auto pjrt_client,
+                               xla::GetTfrtCpuClient(options));
+           return std::shared_ptr<xla::ifrt::Client>(
+               xla::ifrt::PjRtClient::Create(std::move(pjrt_client)));
+         }),
+     true);
+
+class IfrtCallOpTest : public OpsTestBase {
+ protected:
+  Status Init(int64_t program_id, int num_inputs, DataType input_type,
+              const std::vector<int>& variable_arg_indices,
+              const std::vector<DataType>& output_type_list) {
+    TF_CHECK_OK(NodeDefBuilder("op", "IfrtCall")
+                    .Input(FakeInput(num_inputs, input_type))
+                    .Attr("program_id", program_id)
+                    .Attr("variable_arg_indices", variable_arg_indices)
+                    .Attr("Tout", output_type_list)
+                    .Finalize(node_def()));
+    return InitOp();
+  }
+};
+
+TEST_F(IfrtCallOpTest, Basic) {
+  int64_t program_id = 123;
+  TF_ASSERT_OK(Init(
+      /*program_id=*/program_id,
+      /*num_inputs=*/2,
+      /*input_type=*/DT_INT32,
+      /*variable_arg_indices=*/{},
+      /*output_type_list=*/{DT_INT32}));
+
+  // TODO(b/324451111): Extract the executable creation to a common function
+  // to be shared by all tests (e.g.
+  // tensorflow/core/tfrt/ifrt/ifrt_serving_executable_test.cc)
+  tsl::test_util::MockServingDeviceSelector selector;
+  EXPECT_CALL(selector, ReserveDevice(absl::StrCat(program_id)))
+      .Times(1)
+      .WillOnce(Return(tsl::DeviceReservation(0, /*selector=*/nullptr)));
+
+  auto core_selector = std::make_unique<IfrtServingCoreSelector>(&selector);
+
+  constexpr absl::string_view kDataDirectory =
+      "tensorflow/core/tfrt/ifrt/testdata";
+  std::string mlir_module_path = tensorflow::GetDataDependencyFilepath(
+      absl::StrCat(kDataDirectory, "/executable.mlir"));
+  mlir::DialectRegistry registry;
+  mlir::registerAllDialects(registry);
+  mlir::RegisterAllTensorFlowDialects(registry);
+  mlir::MLIRContext context(registry);
+  auto mlir_module =
+      mlir::parseSourceFile<mlir::ModuleOp>(mlir_module_path, &context);
+
+  TF_ASSERT_OK_AND_ASSIGN(auto client, xla::ifrt::test_util::GetClient());
+
+  constexpr int kMaxParallelism = 16;
+  auto thread_pool = std::make_unique<tsl::thread::ThreadPool>(
+      tsl::Env::Default(), tsl::ThreadOptions(), "IfrtSharding",
+      kMaxParallelism);
+
+  auto work_queue = tfrt::CreateMultiThreadedWorkQueue(
+      /*num_threads=*/4, /*num_blocking_threads=*/4);
+
+  TF_ASSERT_OK_AND_ASSIGN(auto device_mgr,
+                          ifrt_serving::CreateTfStaticDeviceMgr());
+
+  IfrtLoadedVariableRegistry ifrt_loaded_variable_registry;
+  IfrtRestoreTensorRegistry ifrt_restore_tensor_registry;
+
+  TF_ASSERT_OK_AND_ASSIGN(
+      auto executable,
+      IfrtServingExecutable::Create(
+          program_id, "test", "main", std::move(mlir_module), client,
+          thread_pool.get(), &ifrt_loaded_variable_registry,
+          &ifrt_restore_tensor_registry, work_queue.get(), device_mgr_.get(),
+          tensorflow::IdentityShapeRepresentationFn(), core_selector.get()));
+
+  TF_ASSERT_OK_AND_ASSIGN(
+      ServingExecutableRegistry::Handle handle,
+      ServingExecutableRegistry::Register(program_id, std::move(executable)));
+  auto handle_cleaner = gtl::MakeCleanup([&handle] { handle.Release(); });
+
+  AddInputFromArray<int32_t>(TensorShape({1, 3}), {1, 2, 3});
+  AddInputFromArray<int32_t>(TensorShape({3, 1}), {1, 2, 3});
+  TF_ASSERT_OK(RunOpKernel());
+  Tensor expected_out = AsTensor<int32_t>({14}, TensorShape({1, 1}));
+  EXPECT_THAT(*GetOutput(0), TensorEq(expected_out));
+}
+
+}  // namespace
+}  // namespace tfrt_stub
+}  // namespace tensorflow

From 4c655fd456988cd90717bdd3d66501fb5f728ed1 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Sat, 18 May 2024 03:10:50 -0700
Subject: [PATCH 390/478] Automated Code Change

PiperOrigin-RevId: 635015194
---
 tensorflow/core/common_runtime/device_mgr.h   |  4 ++--
 .../core/common_runtime/direct_session.h      |  4 ++--
 .../common_runtime/direct_session_test.cc     | 10 ++++----
 .../core/common_runtime/dynamic_device_mgr.cc |  2 +-
 tensorflow/core/common_runtime/function.cc    | 24 +++++++++----------
 5 files changed, 22 insertions(+), 22 deletions(-)

diff --git a/tensorflow/core/common_runtime/device_mgr.h b/tensorflow/core/common_runtime/device_mgr.h
index 6e8420de1733cb..82688a2311a345 100644
--- a/tensorflow/core/common_runtime/device_mgr.h
+++ b/tensorflow/core/common_runtime/device_mgr.h
@@ -65,7 +65,7 @@ class DeviceMgr {
 
   // Clears given containers of all devices if 'container' is
   // non-empty. Otherwise, clears default containers of all devices.
-  virtual void ClearContainers(gtl::ArraySlice<string> containers) const = 0;
+  virtual void ClearContainers(absl::Span<const string> containers) const = 0;
 
   virtual int NumDeviceType(const string& type) const = 0;
 
@@ -102,7 +102,7 @@ class DynamicDeviceMgr : public DeviceMgr {
   string DeviceMappingString() const override;
   Status LookupDevice(StringPiece name, Device** device) const override;
   bool ContainsDevice(int64_t device_incarnation) const override;
-  void ClearContainers(gtl::ArraySlice<string> containers) const override;
+  void ClearContainers(absl::Span<const string> containers) const override;
   int NumDeviceType(const string& type) const override;
   int NumDevices() const override;
   Device* HostCPU() const override;
diff --git a/tensorflow/core/common_runtime/direct_session.h b/tensorflow/core/common_runtime/direct_session.h
index 6421862af65af0..a755befa5bf0d8 100644
--- a/tensorflow/core/common_runtime/direct_session.h
+++ b/tensorflow/core/common_runtime/direct_session.h
@@ -241,8 +241,8 @@ class DirectSession : public Session {
   // Retrieves an already existing set of executors to run 'inputs' and
   // 'outputs', or creates and caches them for future use.
   ::tensorflow::Status GetOrCreateExecutors(
-      gtl::ArraySlice<string> inputs, gtl::ArraySlice<string> outputs,
-      gtl::ArraySlice<string> target_nodes,
+      absl::Span<const string> inputs, absl::Span<const string> outputs,
+      absl::Span<const string> target_nodes,
       ExecutorsAndKeys** executors_and_keys, RunStateArgs* run_state_args);
 
   // Creates a set of executors to run the subgraph defined by
diff --git a/tensorflow/core/common_runtime/direct_session_test.cc b/tensorflow/core/common_runtime/direct_session_test.cc
index f278dfc233bbf9..355f16d57aac6b 100644
--- a/tensorflow/core/common_runtime/direct_session_test.cc
+++ b/tensorflow/core/common_runtime/direct_session_test.cc
@@ -68,9 +68,9 @@ limitations under the License.
 namespace tensorflow {
 namespace {
 
-CallableOptions MakeCallableOptions(gtl::ArraySlice<string> feeds,
-                                    gtl::ArraySlice<string> fetches,
-                                    gtl::ArraySlice<string> targets) {
+CallableOptions MakeCallableOptions(absl::Span<const string> feeds,
+                                    absl::Span<const string> fetches,
+                                    absl::Span<const string> targets) {
   CallableOptions ret;
   for (const string& feed : feeds) {
     ret.add_feed(feed);
@@ -1444,7 +1444,7 @@ TEST(DirectSessionTest, SessionSyncRun_DeepGraph) {
   std::vector<Node*> nodes;
   nodes.reserve(1024);
 
-  auto make_expensive_noop = [&g](gtl::ArraySlice<Node*> control_deps) {
+  auto make_expensive_noop = [&g](absl::Span<Node* const> control_deps) {
     Node* ret;
     auto builder = NodeBuilder(g.NewName("N"), "ExpensiveNoop");
     for (Node* control_dep : control_deps) {
@@ -2850,7 +2850,7 @@ class DirectSessionCollectiveTest : public ::testing::Test {
             {{"group_size", 2},
              {"group_key", 1},
              {"instance_key", instance_key},
-             {"subdiv_offsets", gtl::ArraySlice<int32>({0})},
+             {"subdiv_offsets", absl::Span<const int32>({0})},
              {"merge_op", "Add"},
              {"final_op", "Div"},
              {"T", DT_FLOAT}},
diff --git a/tensorflow/core/common_runtime/dynamic_device_mgr.cc b/tensorflow/core/common_runtime/dynamic_device_mgr.cc
index 325bbfd97b9849..06ccc121440394 100644
--- a/tensorflow/core/common_runtime/dynamic_device_mgr.cc
+++ b/tensorflow/core/common_runtime/dynamic_device_mgr.cc
@@ -127,7 +127,7 @@ bool DynamicDeviceMgr::ContainsDevice(int64_t device_incarnation) const {
 }
 
 void DynamicDeviceMgr::ClearContainers(
-    gtl::ArraySlice<string> containers) const {
+    absl::Span<const string> containers) const {
   Status s;
   tf_shared_lock l(devices_mu_);
   for (const auto& it : dynamic_devices_) {
diff --git a/tensorflow/core/common_runtime/function.cc b/tensorflow/core/common_runtime/function.cc
index 8a55581ed19262..3cb7aa98a95ec9 100644
--- a/tensorflow/core/common_runtime/function.cc
+++ b/tensorflow/core/common_runtime/function.cc
@@ -170,13 +170,13 @@ class FunctionLibraryRuntimeOverlay : public FunctionLibraryRuntime {
 
   Status GetRetTypes(Handle h, DataTypeVector* ret_types) override;
 
-  void Run(const Options& opts, Handle handle, gtl::ArraySlice<Tensor> args,
+  void Run(const Options& opts, Handle handle, absl::Span<const Tensor> args,
            std::vector<Tensor>* rets, DoneCallback done) override;
 
   void Run(const Options& opts, Handle handle, CallFrameInterface* call_frame,
            DoneCallback done) override;
 
-  Status RunSync(Options opts, Handle handle, gtl::ArraySlice<Tensor> args,
+  Status RunSync(Options opts, Handle handle, absl::Span<const Tensor> args,
                  std::vector<Tensor>* rets) override;
 
   Status RunSync(Options opts, Handle handle,
@@ -240,7 +240,7 @@ Status FunctionLibraryRuntimeOverlay::GetRetTypes(Handle h,
 }
 
 void FunctionLibraryRuntimeOverlay::Run(const Options& opts, Handle handle,
-                                        gtl::ArraySlice<Tensor> args,
+                                        absl::Span<const Tensor> args,
                                         std::vector<Tensor>* rets,
                                         DoneCallback done) {
   base_flr_->Run(opts, handle, args, rets, std::move(done));
@@ -253,7 +253,7 @@ void FunctionLibraryRuntimeOverlay::Run(const Options& opts, Handle handle,
 }
 
 Status FunctionLibraryRuntimeOverlay::RunSync(Options opts, Handle handle,
-                                              gtl::ArraySlice<Tensor> args,
+                                              absl::Span<const Tensor> args,
                                               std::vector<Tensor>* rets) {
   return base_flr_->RunSync(std::move(opts), handle, args, rets);
 }
@@ -354,11 +354,11 @@ class FunctionLibraryRuntimeImpl : public FunctionLibraryRuntime {
   Status CreateKernel(const std::shared_ptr<const NodeProperties>& props,
                       OpKernel** kernel) override;
 
-  void Run(const Options& opts, Handle handle, gtl::ArraySlice<Tensor> args,
+  void Run(const Options& opts, Handle handle, absl::Span<const Tensor> args,
            std::vector<Tensor>* rets, DoneCallback done) override;
   void Run(const Options& opts, Handle handle, CallFrameInterface* frame,
            DoneCallback done) override;
-  Status RunSync(Options opts, Handle handle, gtl::ArraySlice<Tensor> args,
+  Status RunSync(Options opts, Handle handle, absl::Span<const Tensor> args,
                  std::vector<Tensor>* rets) override;
   Status RunSync(Options opts, Handle handle,
                  CallFrameInterface* call_frame) override;
@@ -453,7 +453,7 @@ class FunctionLibraryRuntimeImpl : public FunctionLibraryRuntime {
   bool IsLocalTarget(const InstantiateOptions& options) const;
   AttrValueMap FixAttrs(const AttrSlice& attrs);
   void RunRemote(const Options& opts, Handle handle,
-                 gtl::ArraySlice<Tensor> args, std::vector<Tensor>* rets,
+                 absl::Span<const Tensor> args, std::vector<Tensor>* rets,
                  Item* item, DoneCallback done);
 
   // TODO(fishx): Avoid using std::unique_ptr for PrivateIntraProcessRendezvous,
@@ -1008,7 +1008,7 @@ void FunctionLibraryRuntimeImpl::ExecutorArgsFromOptions(
 }
 
 void FunctionLibraryRuntimeImpl::RunRemote(const Options& opts, Handle handle,
-                                           gtl::ArraySlice<Tensor> args,
+                                           absl::Span<const Tensor> args,
                                            std::vector<Tensor>* rets,
                                            Item* item, DoneCallback done) {
   string target_device = parent_->GetDeviceName(handle);
@@ -1098,7 +1098,7 @@ void FunctionLibraryRuntimeImpl::RunRemote(const Options& opts, Handle handle,
 }
 
 void FunctionLibraryRuntimeImpl::Run(const Options& opts, Handle handle,
-                                     gtl::ArraySlice<Tensor> args,
+                                     absl::Span<const Tensor> args,
                                      std::vector<Tensor>* rets,
                                      DoneCallback done) {
   if (opts.cancellation_manager && opts.cancellation_manager->IsCancelled()) {
@@ -1159,7 +1159,7 @@ void FunctionLibraryRuntimeImpl::Run(const Options& opts, Handle handle,
         return tsl::profiler::TraceMeEncode(
             "FunctionRun", {{"id", run_opts.step_id}, {"_r", 1}});
       },
-      profiler::ContextType::kTfExecutor, *exec_args.function_trace_id,
+      tsl::profiler::ContextType::kTfExecutor, *exec_args.function_trace_id,
       tsl::profiler::TraceMeLevel::kInfo);
 
   bool allow_dead_tensors = run_opts.allow_dead_tensors;
@@ -1231,7 +1231,7 @@ void FunctionLibraryRuntimeImpl::Run(const Options& opts, Handle handle,
         return tsl::profiler::TraceMeEncode("FunctionRun",
                                             {{"id", opts.step_id}, {"_r", 1}});
       },
-      profiler::ContextType::kTfExecutor, *exec_args.function_trace_id,
+      tsl::profiler::ContextType::kTfExecutor, *exec_args.function_trace_id,
       tsl::profiler::TraceMeLevel::kInfo);
 
   item->exec->RunAsync(exec_args, std::move(done));
@@ -1278,7 +1278,7 @@ Status FunctionLibraryRuntimeImpl::PrepareRunSync(
 }
 
 Status FunctionLibraryRuntimeImpl::RunSync(Options opts, Handle handle,
-                                           gtl::ArraySlice<Tensor> args,
+                                           absl::Span<const Tensor> args,
                                            std::vector<Tensor>* rets) {
   Item* item = nullptr;
   std::unique_ptr<PrivateIntraProcessRendezvous> rendezvous;

From 530da27481c5819f8ee7ed9fa1175a86bc72e5a1 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Sat, 18 May 2024 05:40:53 -0700
Subject: [PATCH 391/478] Move `ForAllThunks` to its own file.

`ForAllThunks` is currently unused, and the new version has a few improvements:
* For `Thunks` that have nested `Thunks`, e.g. `ConditionalThunk`, it invokes `fn` on both the parent `Thunk` and all nested children. The previous version would only invoke `fn` on the nested children.
* It discovers nested children in all `Thunk` types, including `AddressComputationThunks` and `CommandBufferThunks`.
* It uses more modern types -- e.g. `absl::FunctionRef`.

PiperOrigin-RevId: 635036109
---
 .../service/gpu/compile_module_to_llvm_ir.cc  |  23 ---
 third_party/xla/xla/service/gpu/runtime/BUILD |  34 +++++
 .../gpu/runtime/address_computation_thunk.h   |   2 +
 .../gpu/runtime/command_buffer_thunk.h        |   2 +
 .../xla/service/gpu/runtime/for_all_thunks.cc | 128 ++++++++++++++++
 .../xla/service/gpu/runtime/for_all_thunks.h  |  34 +++++
 .../gpu/runtime/for_all_thunks_test.cc        | 139 ++++++++++++++++++
 7 files changed, 339 insertions(+), 23 deletions(-)
 create mode 100644 third_party/xla/xla/service/gpu/runtime/for_all_thunks.cc
 create mode 100644 third_party/xla/xla/service/gpu/runtime/for_all_thunks.h
 create mode 100644 third_party/xla/xla/service/gpu/runtime/for_all_thunks_test.cc

diff --git a/third_party/xla/xla/service/gpu/compile_module_to_llvm_ir.cc b/third_party/xla/xla/service/gpu/compile_module_to_llvm_ir.cc
index 371d0c5e8c35c4..e3197f495d471d 100644
--- a/third_party/xla/xla/service/gpu/compile_module_to_llvm_ir.cc
+++ b/third_party/xla/xla/service/gpu/compile_module_to_llvm_ir.cc
@@ -102,29 +102,6 @@ void RemoveUnusedAndUninitializedGlobals(
 
 }  // namespace
 
-void ForAllThunks(const std::function<void(Thunk*)>& fn,
-                  ThunkSequence* thunk_sequence) {
-  for (std::unique_ptr<Thunk>& thunk : *thunk_sequence) {
-    if (thunk->kind() == Thunk::kConditional) {
-      auto* cond_thunk = tensorflow::down_cast<ConditionalThunk*>(thunk.get());
-      for (const std::unique_ptr<SequentialThunk>& branch_thunks :
-           cond_thunk->branch_thunks()) {
-        ForAllThunks(fn, &branch_thunks->thunks());
-      }
-    } else if (thunk->kind() == Thunk::kSequential) {
-      auto* sequential_thunk =
-          tensorflow::down_cast<SequentialThunk*>(thunk.get());
-      ForAllThunks(fn, &sequential_thunk->thunks());
-    } else if (thunk->kind() == Thunk::kWhile) {
-      auto* while_thunk = tensorflow::down_cast<WhileThunk*>(thunk.get());
-      ForAllThunks(fn, &while_thunk->condition_thunk_sequence()->thunks());
-      ForAllThunks(fn, &while_thunk->body_thunk_sequence()->thunks());
-    } else {
-      fn(thunk.get());
-    }
-  }
-}
-
 absl::StatusOr<CompileModuleResults> CompileModuleToLlvmIr(
     HloModule* hlo_module, llvm::LLVMContext* llvm_context,
     const std::string& target_triple, const std::string& data_layout,
diff --git a/third_party/xla/xla/service/gpu/runtime/BUILD b/third_party/xla/xla/service/gpu/runtime/BUILD
index ae2c70bfb0981a..f13bfe7b3f474e 100644
--- a/third_party/xla/xla/service/gpu/runtime/BUILD
+++ b/third_party/xla/xla/service/gpu/runtime/BUILD
@@ -1123,6 +1123,40 @@ cc_library(
     ],
 )
 
+cc_library(
+    name = "for_all_thunks",
+    srcs = ["for_all_thunks.cc"],
+    hdrs = ["for_all_thunks.h"],
+    deps = [
+        ":address_computation_thunk",
+        ":command_buffer_thunk",
+        ":conditional_thunk",
+        ":sequential_thunk",
+        ":thunk",
+        ":while_thunk",
+        "@com_google_absl//absl/functional:function_ref",
+        "@local_tsl//tsl/platform:casts",
+    ],
+)
+
+cc_test(
+    name = "for_all_thunks_test",
+    srcs = ["for_all_thunks_test.cc"],
+    deps = [
+        ":address_computation_thunk",
+        ":command_buffer_cmd",
+        ":command_buffer_thunk",
+        ":conditional_thunk",
+        ":for_all_thunks",
+        ":sequential_thunk",
+        ":thunk",
+        ":while_thunk",
+        "//xla/service:buffer_assignment",
+        "@com_google_absl//absl/status",
+        "@com_google_googletest//:gtest_main",
+    ],
+)
+
 cc_library(
     name = "triangular_solve_thunk",
     srcs = if_gpu_is_configured(["triangular_solve_thunk.cc"]),
diff --git a/third_party/xla/xla/service/gpu/runtime/address_computation_thunk.h b/third_party/xla/xla/service/gpu/runtime/address_computation_thunk.h
index dffe9c5da1d8f6..bfc70574d975d5 100644
--- a/third_party/xla/xla/service/gpu/runtime/address_computation_thunk.h
+++ b/third_party/xla/xla/service/gpu/runtime/address_computation_thunk.h
@@ -56,6 +56,8 @@ class AddressComputationThunk : public Thunk {
   AddressComputationThunk(const AddressComputationThunk&) = delete;
   AddressComputationThunk& operator=(const AddressComputationThunk&) = delete;
 
+  const Thunk* embedded_thunk() const { return embedded_thunk_.get(); }
+
   absl::Status Prepare(const PrepareParams& params,
                        ResourceRequests& resource_requests) override;
   absl::Status Initialize(const InitializeParams& params) override;
diff --git a/third_party/xla/xla/service/gpu/runtime/command_buffer_thunk.h b/third_party/xla/xla/service/gpu/runtime/command_buffer_thunk.h
index c315f412bb3cce..3c05aada9607e7 100644
--- a/third_party/xla/xla/service/gpu/runtime/command_buffer_thunk.h
+++ b/third_party/xla/xla/service/gpu/runtime/command_buffer_thunk.h
@@ -39,6 +39,8 @@ class CommandBufferThunk : public Thunk {
   CommandBufferThunk(CommandBufferCmdSequence commands, ThunkInfo thunk_info,
                      std::optional<ThunkSequence> thunks = std::nullopt);
 
+  const std::optional<ThunkSequence>& thunks() const { return thunks_; }
+
   absl::Status Prepare(const PrepareParams& params,
                        ResourceRequests& resource_requests) override;
   absl::Status Initialize(const InitializeParams& params) override;
diff --git a/third_party/xla/xla/service/gpu/runtime/for_all_thunks.cc b/third_party/xla/xla/service/gpu/runtime/for_all_thunks.cc
new file mode 100644
index 00000000000000..77932caad72fa0
--- /dev/null
+++ b/third_party/xla/xla/service/gpu/runtime/for_all_thunks.cc
@@ -0,0 +1,128 @@
+/* Copyright 2024 The OpenXLA Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+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.
+==============================================================================*/
+
+#include "xla/service/gpu/runtime/for_all_thunks.h"
+
+#include <memory>
+#include <optional>
+
+#include "absl/functional/function_ref.h"
+#include "xla/service/gpu/runtime/address_computation_thunk.h"
+#include "xla/service/gpu/runtime/command_buffer_thunk.h"
+#include "xla/service/gpu/runtime/conditional_thunk.h"
+#include "xla/service/gpu/runtime/sequential_thunk.h"
+#include "xla/service/gpu/runtime/thunk.h"
+#include "xla/service/gpu/runtime/while_thunk.h"
+#include "tsl/platform/casts.h"
+
+namespace xla::gpu {
+
+void ForAllThunks(absl::FunctionRef<void(const Thunk*)> fn,
+                  const Thunk* thunk) {
+  // Invoke `fn` with the `Thunk` itself first...
+  fn(thunk);
+  // ... and then handle all nested `Thunks` recursively.
+  switch (thunk->kind()) {
+    case Thunk::kAddressComputation:
+      ForAllThunks(fn,
+                   tensorflow::down_cast<const AddressComputationThunk*>(thunk)
+                       ->embedded_thunk());
+      break;
+    case Thunk::kCommandBuffer:
+      if (const std::optional<ThunkSequence>& sequence =
+              tensorflow::down_cast<const CommandBufferThunk*>(thunk)->thunks();
+          sequence.has_value()) {
+        ForAllThunks(fn, &sequence.value());
+      }
+      break;
+    case Thunk::kConditional:
+      for (const std::unique_ptr<SequentialThunk>& branch :
+           tensorflow::down_cast<const ConditionalThunk*>(thunk)
+               ->branch_thunks()) {
+        ForAllThunks(fn, branch.get());
+      }
+      break;
+    case Thunk::kSequential:
+      ForAllThunks(
+          fn, &tensorflow::down_cast<const SequentialThunk*>(thunk)->thunks());
+      break;
+    case Thunk::kWhile:
+      ForAllThunks(fn, tensorflow::down_cast<const WhileThunk*>(thunk)
+                           ->condition_thunk_sequence());
+      ForAllThunks(fn, tensorflow::down_cast<const WhileThunk*>(thunk)
+                           ->body_thunk_sequence());
+      break;
+    case Thunk::kCholesky:
+    case Thunk::kConvolution:
+    case Thunk::kConvolutionReorder:
+    case Thunk::kCopy:
+    case Thunk::kCopyDone:
+    case Thunk::kCubSort:
+    case Thunk::kCublasLtMatmul:
+    case Thunk::kCustomCall:
+    case Thunk::kCustomKernel:
+    case Thunk::kCuDnn:
+    case Thunk::kFft:
+    case Thunk::kFusedMHA:
+    case Thunk::kGemm:
+    case Thunk::kInfeed:
+    case Thunk::kKernel:
+    case Thunk::kMemset32BitValue:
+    case Thunk::kMemzero:
+    case Thunk::kNcclAllGather:
+    case Thunk::kNcclAllGatherStart:
+    case Thunk::kNcclAllGatherDone:
+    case Thunk::kNcclAllReduce:
+    case Thunk::kNcclAllReduceStart:
+    case Thunk::kNcclAllReduceDone:
+    case Thunk::kNcclCollectiveBroadcast:
+    case Thunk::kNcclCollectiveBroadcastStart:
+    case Thunk::kNcclCollectiveBroadcastDone:
+    case Thunk::kNcclCollectivePermute:
+    case Thunk::kNcclCollectivePermuteStart:
+    case Thunk::kNcclCollectivePermuteDone:
+    case Thunk::kNcclReduceScatter:
+    case Thunk::kNcclReduceScatterStart:
+    case Thunk::kNcclReduceScatterDone:
+    case Thunk::kNcclAllToAll:
+    case Thunk::kNcclAllToAllStart:
+    case Thunk::kNcclAllToAllDone:
+    case Thunk::kNcclSend:
+    case Thunk::kNcclSendDone:
+    case Thunk::kNcclRecv:
+    case Thunk::kNcclRecvDone:
+    case Thunk::kNorm:
+    case Thunk::kOutfeed:
+    case Thunk::kPartitionId:
+    case Thunk::kRecv:
+    case Thunk::kRecvDone:
+    case Thunk::kReplicaId:
+    case Thunk::kSend:
+    case Thunk::kSendDone:
+    case Thunk::kTriangularSolve:
+    case Thunk::kWaitForStreams:
+      // No default. All `Thunk::Kinds` must be handled.
+      break;
+  }
+}
+
+void ForAllThunks(absl::FunctionRef<void(const Thunk*)> fn,
+                  const ThunkSequence* thunks) {
+  for (const std::unique_ptr<Thunk>& thunk : *thunks) {
+    ForAllThunks(fn, thunk.get());
+  }
+}
+
+}  // namespace xla::gpu
diff --git a/third_party/xla/xla/service/gpu/runtime/for_all_thunks.h b/third_party/xla/xla/service/gpu/runtime/for_all_thunks.h
new file mode 100644
index 00000000000000..6f6fc61c34c427
--- /dev/null
+++ b/third_party/xla/xla/service/gpu/runtime/for_all_thunks.h
@@ -0,0 +1,34 @@
+/* Copyright 2024 The OpenXLA Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+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.
+==============================================================================*/
+
+#ifndef XLA_SERVICE_GPU_RUNTIME_FOR_ALL_THUNKS_H_
+#define XLA_SERVICE_GPU_RUNTIME_FOR_ALL_THUNKS_H_
+
+#include "absl/functional/function_ref.h"
+#include "xla/service/gpu/runtime/thunk.h"
+
+namespace xla::gpu {
+
+// Recursively invokes `fn` for all `Thunks` in `root`, including those nested
+// within other `Thunks` (e.g. the condition `Thunk` within a `WhileThunk`).
+void ForAllThunks(absl::FunctionRef<void(const Thunk*)> fn, const Thunk* thunk);
+
+// Same as above but for a `ThunkSequence` root.
+void ForAllThunks(absl::FunctionRef<void(const Thunk*)> fn,
+                  const ThunkSequence* thunks);
+
+}  // namespace xla::gpu
+
+#endif  // XLA_SERVICE_GPU_RUNTIME_FOR_ALL_THUNKS_H_
diff --git a/third_party/xla/xla/service/gpu/runtime/for_all_thunks_test.cc b/third_party/xla/xla/service/gpu/runtime/for_all_thunks_test.cc
new file mode 100644
index 00000000000000..32f7bc6d144af6
--- /dev/null
+++ b/third_party/xla/xla/service/gpu/runtime/for_all_thunks_test.cc
@@ -0,0 +1,139 @@
+/* Copyright 2024 The OpenXLA Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+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.
+==============================================================================*/
+
+#include "xla/service/gpu/runtime/for_all_thunks.h"
+
+#include <memory>
+#include <utility>
+#include <vector>
+
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+#include "absl/status/status.h"
+#include "xla/service/buffer_assignment.h"
+#include "xla/service/gpu/runtime/address_computation_thunk.h"
+#include "xla/service/gpu/runtime/command_buffer_cmd.h"
+#include "xla/service/gpu/runtime/command_buffer_thunk.h"
+#include "xla/service/gpu/runtime/conditional_thunk.h"
+#include "xla/service/gpu/runtime/sequential_thunk.h"
+#include "xla/service/gpu/runtime/thunk.h"
+#include "xla/service/gpu/runtime/while_thunk.h"
+
+namespace xla::gpu {
+namespace {
+
+using ::testing::IsSupersetOf;
+using ::testing::UnorderedElementsAre;
+
+// Invokes `ForAllThunks` on the `root` and returns a `vector` containing all
+// iterated `Thunks`.
+std::vector<const Thunk*> GetAllThunks(Thunk* root) {
+  std::vector<const Thunk*> thunks;
+  ForAllThunks([&](const Thunk* thunk) { thunks.push_back(thunk); }, root);
+  return thunks;
+}
+
+// A dummy `Thunk` that does nothing.
+struct DummyThunk : public Thunk {
+  DummyThunk() : Thunk(Thunk::Kind::kGemm, Thunk::ThunkInfo()) {}
+  absl::Status ExecuteOnStream(const ExecuteParams& params) override {
+    return absl::OkStatus();
+  }
+};
+
+TEST(ForAllThunksTest, SingleThunk) {
+  DummyThunk thunk;
+  EXPECT_THAT(GetAllThunks(&thunk), UnorderedElementsAre(&thunk));
+}
+
+TEST(ForAllThunksTest, AddressComputationThunk) {
+  auto thunk = std::make_unique<DummyThunk>();
+  Thunk* thunk_ptr = thunk.get();
+
+  auto thunk_sequence = std::make_unique<ThunkSequence>();
+  thunk_sequence->push_back(std::move(thunk));
+
+  AddressComputationThunk address_computation_thunk(
+      Thunk::ThunkInfo(), std::move(thunk_sequence), {}, {}, {}, {}, {}, {});
+  EXPECT_THAT(
+      GetAllThunks(&address_computation_thunk),
+      // `AddressComputationThunk` wraps the `embedded_thunk` in a
+      // `SequentialThunk`, which is why iterate over more than the
+      // two expected `Thunks`.
+      IsSupersetOf<const Thunk*>({thunk_ptr, &address_computation_thunk}));
+}
+
+TEST(ForAllThunksTest, CommandBufferThunk) {
+  auto thunk = std::make_unique<DummyThunk>();
+  Thunk* thunk_ptr = thunk.get();
+
+  ThunkSequence thunk_sequence;
+  thunk_sequence.push_back(std::move(thunk));
+
+  CommandBufferThunk command_buffer_thunk(CommandBufferCmdSequence(),
+                                          Thunk::ThunkInfo(),
+                                          std::move(thunk_sequence));
+  EXPECT_THAT(GetAllThunks(&command_buffer_thunk),
+              UnorderedElementsAre(thunk_ptr, &command_buffer_thunk));
+}
+
+TEST(ForAllThunksTest, ConditionalThunk) {
+  auto thunk = std::make_unique<DummyThunk>();
+  Thunk* thunk_ptr = thunk.get();
+
+  ThunkSequence thunk_sequence;
+  thunk_sequence.push_back(std::move(thunk));
+
+  auto sequential_thunk = std::make_unique<SequentialThunk>(
+      Thunk::ThunkInfo(), std::move(thunk_sequence));
+  SequentialThunk* sequential_thunk_ptr = sequential_thunk.get();
+
+  ConditionalThunkConfig config;
+  config.branch_thunks.push_back(std::move(sequential_thunk));
+  ConditionalThunk conditional_thunk(Thunk::ThunkInfo(), std::move(config),
+                                     BufferAllocation::Slice());
+
+  EXPECT_THAT(GetAllThunks(&conditional_thunk),
+              UnorderedElementsAre(thunk_ptr, sequential_thunk_ptr,
+                                   &conditional_thunk));
+}
+
+TEST(ForAllThunksTest, WhileThunk) {
+  auto condition_thunk = std::make_unique<DummyThunk>();
+  Thunk* condition_thunk_ptr = condition_thunk.get();
+
+  auto condition_thunk_sequence = std::make_unique<ThunkSequence>();
+  condition_thunk_sequence->push_back(std::move(condition_thunk));
+
+  auto body_thunk = std::make_unique<DummyThunk>();
+  Thunk* body_thunk_ptr = body_thunk.get();
+
+  auto body_thunk_sequence = std::make_unique<ThunkSequence>();
+  body_thunk_sequence->push_back(std::move(body_thunk));
+
+  WhileThunk while_thunk(Thunk::ThunkInfo(), BufferAllocation::Slice(),
+                         std::move(condition_thunk_sequence),
+                         std::move(body_thunk_sequence));
+
+  EXPECT_THAT(GetAllThunks(&while_thunk),
+              // `WhileThunk` wraps the `condition_thunk_sequence` and
+              // `body_thunk_sequence` in `SequentialThunks`, which is why
+              // iterate over more than the three expected `Thunks`.
+              IsSupersetOf<const Thunk*>(
+                  {condition_thunk_ptr, body_thunk_ptr, &while_thunk}));
+}
+
+}  // namespace
+}  // namespace xla::gpu

From 70205668d8eaa0164c2f1c7a01a90e0b8922d46d Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Sat, 18 May 2024 06:00:07 -0700
Subject: [PATCH 392/478] Extract `ExecutionStreamIds` from `Thunks` instead of
 the now-deprecated `GpuBackendConfig::operation_queue_id`.

PiperOrigin-RevId: 635038291
---
 third_party/xla/xla/service/gpu/BUILD         |  1 +
 .../xla/xla/service/gpu/gpu_executable.cc     | 34 +++++++------------
 2 files changed, 13 insertions(+), 22 deletions(-)

diff --git a/third_party/xla/xla/service/gpu/BUILD b/third_party/xla/xla/service/gpu/BUILD
index 30b13e0f98b3b0..afe18d35d71af7 100644
--- a/third_party/xla/xla/service/gpu/BUILD
+++ b/third_party/xla/xla/service/gpu/BUILD
@@ -922,6 +922,7 @@ cc_library(
         "//xla/service:stream_pool",
         "//xla/service:xla_debug_info_manager",
         "//xla/service/gpu/runtime:annotation",
+        "//xla/service/gpu/runtime:for_all_thunks",
         "//xla/service/gpu/runtime:nccl_clique",
         "//xla/service/gpu/runtime:nccl_clique_key",
         "//xla/service/gpu/runtime:thunk",
diff --git a/third_party/xla/xla/service/gpu/gpu_executable.cc b/third_party/xla/xla/service/gpu/gpu_executable.cc
index 25343eb9c93512..cd0121dcdb77e2 100644
--- a/third_party/xla/xla/service/gpu/gpu_executable.cc
+++ b/third_party/xla/xla/service/gpu/gpu_executable.cc
@@ -53,6 +53,7 @@ limitations under the License.
 #include "xla/service/gpu/gpu_constants.h"
 #include "xla/service/gpu/gpu_executable_run_options.h"
 #include "xla/service/gpu/runtime/annotation.h"
+#include "xla/service/gpu/runtime/for_all_thunks.h"
 #include "xla/service/gpu/runtime/nccl_clique.h"
 #include "xla/service/gpu/runtime/nccl_clique_key.h"
 #include "xla/service/gpu/runtime/thunk.h"
@@ -119,28 +120,19 @@ static bool NeedsAsyncCommsStream(Thunk& thunk) {
   }
 }
 
-// Traverses operations in HLO module and collects execution stream ids
-// requested by HLO operations. At run time thunks may use additional streams to
-// launch compute operations in addition to a main one.
-//
-// TODO(ezhulenev): Execution stream requirements should be queried from thunks
-// directly and not from HLO module that might be missing.
+// Returns the set of `ExecutionStreamIds` requested by all `Thunks` in the
+// `GpuExecutable`. At run time `Thunks` may use additional streams to launch
+// compute operations in parallel.
 static absl::flat_hash_set<ExecutionStreamId> GetExecutionStreamIds(
-    const HloModule& module) {
+    const ThunkSequence& thunks) {
   absl::flat_hash_set<ExecutionStreamId> stream_ids;
-  for (const HloComputation* comp : module.computations()) {
-    for (const HloInstruction* hlo : comp->instructions()) {
-      if (hlo->has_backend_config() &&
-          hlo->backend_config<GpuBackendConfig>().ok()) {
-        int64_t op_queue_id = hlo->backend_config<GpuBackendConfig>()
-                                  .value()
-                                  .operation_queue_id();
-        if (op_queue_id > 0) {
-          stream_ids.insert(ExecutionStreamId(op_queue_id));
+  ForAllThunks(
+      [&](const Thunk* thunk) {
+        if (thunk->execution_stream_id() > 0) {
+          stream_ids.insert(thunk->execution_stream_id());
         }
-      }
-    }
-  }
+      },
+      &thunks);
   return stream_ids;
 }
 
@@ -158,9 +150,7 @@ GpuExecutable::GpuExecutable(GpuExecutable::Params params)
       dnn_compiled_graphs_(std::move(params.dnn_compiled_graphs)),
       gpu_version_(params.gpu_version),
       thunks_(std::move(params.executable)),
-      execution_stream_ids_(has_module()
-                                ? GetExecutionStreamIds(module())
-                                : absl::flat_hash_set<ExecutionStreamId>()),
+      execution_stream_ids_(GetExecutionStreamIds(*thunks_)),
       module_name_(params.module_name),
       output_shape_(params.output_shape),
       allocations_(std::move(params.mlir_allocations)),

From 4f584aeb5e1a143f84f9522e1bb309012b74da5f Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Sat, 18 May 2024 06:27:11 -0700
Subject: [PATCH 393/478] Automated Code Change

PiperOrigin-RevId: 635041988
---
 tensorflow/tools/api/tests/convert_from_multiline.cc | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorflow/tools/api/tests/convert_from_multiline.cc b/tensorflow/tools/api/tests/convert_from_multiline.cc
index afde8f9f122f77..9ad5a12be7d378 100644
--- a/tensorflow/tools/api/tests/convert_from_multiline.cc
+++ b/tensorflow/tools/api/tests/convert_from_multiline.cc
@@ -48,7 +48,7 @@ Status ConvertFilesFromMultiline(const string& input_dir,
     TF_RETURN_IF_ERROR(
         tensorflow::WriteStringToFile(env, output_path, contents));
   }
-  return OkStatus();
+  return absl::OkStatus();
 }
 }  // namespace
 }  // namespace tensorflow

From 113315ce187e34feb94110d94994833dbe48f6da Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Sat, 18 May 2024 08:11:37 -0700
Subject: [PATCH 394/478] Automated Code Change

PiperOrigin-RevId: 635053948
---
 tensorflow/dtensor/cc/dtensor_device.cc | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/tensorflow/dtensor/cc/dtensor_device.cc b/tensorflow/dtensor/cc/dtensor_device.cc
index 6632f22be4f6d7..ca1460e6260990 100644
--- a/tensorflow/dtensor/cc/dtensor_device.cc
+++ b/tensorflow/dtensor/cc/dtensor_device.cc
@@ -2826,7 +2826,7 @@ void ExperimentalSetDefaultLayout(const std::string& serialized_layout,
   StatusOr<Layout> layout = Layout::FromString(serialized_layout);
   if (!layout.ok()) {
     RETURN_STATUS(status, TF_INTERNAL,
-                  tsl::NullTerminatedMessage(layout.status()));
+                  absl::StatusMessageAsCStr(layout.status()));
   }
   DTensorDevice* device = reinterpret_cast<DTensorDevice*>(device_info);
   device->SetDefaultLayout(layout.value());
@@ -2842,7 +2842,7 @@ void ExperimentalSetDefaultMesh(const std::string& serialized_mesh,
   StatusOr<Mesh> mesh = Mesh::FromString(serialized_mesh);
   if (!mesh.ok()) {
     RETURN_STATUS(status, TF_INTERNAL,
-                  tsl::NullTerminatedMessage(mesh.status()));
+                  absl::StatusMessageAsCStr(mesh.status()));
   }
   DTensorDevice* device = reinterpret_cast<DTensorDevice*>(device_info);
   device->SetDefaultMesh(mesh.value());

From 81e1e204055cb1e4c564d7c81ee83784df9c64dd Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Sat, 18 May 2024 08:16:33 -0700
Subject: [PATCH 395/478] Automated Code Change

PiperOrigin-RevId: 635054554
---
 third_party/xla/third_party/tsl/tsl/profiler/lib/BUILD   | 9 +++++++--
 .../tsl/tsl/profiler/lib/profiler_collection.cc          | 2 +-
 .../tsl/tsl/profiler/lib/profiler_collection.h           | 1 +
 .../tsl/tsl/profiler/lib/profiler_controller.cc          | 1 -
 .../tsl/tsl/profiler/lib/profiler_controller.h           | 1 +
 .../tsl/tsl/profiler/lib/profiler_factory_test.cc        | 2 +-
 .../third_party/tsl/tsl/profiler/lib/profiler_lock.cc    | 2 +-
 .../xla/third_party/tsl/tsl/profiler/lib/profiler_lock.h | 1 +
 .../tsl/tsl/profiler/lib/profiler_lock_test.cc           | 2 +-
 .../third_party/tsl/tsl/profiler/lib/profiler_session.cc | 2 +-
 .../third_party/tsl/tsl/profiler/lib/profiler_session.h  | 1 +
 11 files changed, 16 insertions(+), 8 deletions(-)

diff --git a/third_party/xla/third_party/tsl/tsl/profiler/lib/BUILD b/third_party/xla/third_party/tsl/tsl/profiler/lib/BUILD
index b6b43562847e0e..85ff01b6a987ce 100644
--- a/third_party/xla/third_party/tsl/tsl/profiler/lib/BUILD
+++ b/third_party/xla/third_party/tsl/tsl/profiler/lib/BUILD
@@ -64,6 +64,7 @@ cc_library(
         "//tsl/platform:logging",
         "//tsl/platform:status",
         "//tsl/profiler/protobuf:xplane_proto_cc",
+        "@com_google_absl//absl/status",
     ],
 )
 
@@ -112,12 +113,12 @@ tsl_cc_test(
         ":profiler_factory_impl",
         ":profiler_interface",
         "//tsl/platform:macros",
-        "//tsl/platform:status",
         "//tsl/platform:test",
         "//tsl/platform:test_main",
         "//tsl/profiler/protobuf:profiler_options_proto_cc",
         "//tsl/profiler/protobuf:xplane_proto_cc",
         "@com_google_absl//absl/memory",
+        "@com_google_absl//absl/status",
     ],
 )
 
@@ -149,6 +150,7 @@ cc_library(
         "//tsl/platform:errors",
         "//tsl/platform:macros",
         "//tsl/platform:statusor",
+        "@com_google_absl//absl/status:statusor",
         "@com_google_absl//absl/strings:string_view",
         "@local_xla//xla/tsl/util:env_var",
     ],
@@ -159,9 +161,9 @@ tsl_cc_test(
     srcs = ["profiler_lock_test.cc"],
     deps = [
         ":profiler_lock",
-        "//tsl/platform:statusor",
         "//tsl/platform:test",
         "//tsl/platform:test_main",
+        "@com_google_absl//absl/status:statusor",
     ],
 )
 
@@ -178,6 +180,7 @@ cc_library(
         "//tsl/platform:types",
         "//tsl/profiler/protobuf:profiler_options_proto_cc",
         "//tsl/profiler/protobuf:xplane_proto_cc",
+        "@com_google_absl//absl/status",
     ] + if_not_android([
         ":profiler_interface",
         ":profiler_lock",
@@ -206,6 +209,7 @@ cc_library(
         "//tsl/profiler/protobuf:profiler_options_proto_cc",
         "//tsl/profiler/protobuf:xplane_proto_cc",
         "@com_google_absl//absl/memory",
+        "@com_google_absl//absl/status",
     ] + if_not_android([
         ":profiler_collection",
         ":profiler_factory",
@@ -354,6 +358,7 @@ cc_library(
         ":profiler_interface",
         "//tsl/platform:status",
         "//tsl/profiler/protobuf:xplane_proto_cc",
+        "@com_google_absl//absl/status",
     ],
 )
 
diff --git a/third_party/xla/third_party/tsl/tsl/profiler/lib/profiler_collection.cc b/third_party/xla/third_party/tsl/tsl/profiler/lib/profiler_collection.cc
index 762a23cfe1b3e0..f3ffec62bff7f7 100644
--- a/third_party/xla/third_party/tsl/tsl/profiler/lib/profiler_collection.cc
+++ b/third_party/xla/third_party/tsl/tsl/profiler/lib/profiler_collection.cc
@@ -18,7 +18,7 @@ limitations under the License.
 #include <utility>
 #include <vector>
 
-#include "tsl/platform/status.h"
+#include "absl/status/status.h"
 #include "tsl/profiler/lib/profiler_interface.h"
 #include "tsl/profiler/protobuf/xplane.pb.h"
 
diff --git a/third_party/xla/third_party/tsl/tsl/profiler/lib/profiler_collection.h b/third_party/xla/third_party/tsl/tsl/profiler/lib/profiler_collection.h
index a73e1725d292f5..c3bede9af47c8d 100644
--- a/third_party/xla/third_party/tsl/tsl/profiler/lib/profiler_collection.h
+++ b/third_party/xla/third_party/tsl/tsl/profiler/lib/profiler_collection.h
@@ -18,6 +18,7 @@ limitations under the License.
 #include <memory>
 #include <vector>
 
+#include "absl/status/status.h"
 #include "tsl/platform/status.h"
 #include "tsl/profiler/lib/profiler_interface.h"
 #include "tsl/profiler/protobuf/xplane.pb.h"
diff --git a/third_party/xla/third_party/tsl/tsl/profiler/lib/profiler_controller.cc b/third_party/xla/third_party/tsl/tsl/profiler/lib/profiler_controller.cc
index 4820b97f03b84b..55fc42706dfea5 100644
--- a/third_party/xla/third_party/tsl/tsl/profiler/lib/profiler_controller.cc
+++ b/third_party/xla/third_party/tsl/tsl/profiler/lib/profiler_controller.cc
@@ -19,7 +19,6 @@ limitations under the License.
 
 #include "tsl/platform/errors.h"
 #include "tsl/platform/logging.h"
-#include "tsl/platform/status.h"
 #include "tsl/profiler/lib/profiler_interface.h"
 #include "tsl/profiler/protobuf/xplane.pb.h"
 
diff --git a/third_party/xla/third_party/tsl/tsl/profiler/lib/profiler_controller.h b/third_party/xla/third_party/tsl/tsl/profiler/lib/profiler_controller.h
index d797960c6d3b26..ed88f8ec26b561 100644
--- a/third_party/xla/third_party/tsl/tsl/profiler/lib/profiler_controller.h
+++ b/third_party/xla/third_party/tsl/tsl/profiler/lib/profiler_controller.h
@@ -17,6 +17,7 @@ limitations under the License.
 
 #include <memory>
 
+#include "absl/status/status.h"
 #include "tsl/platform/status.h"
 #include "tsl/profiler/lib/profiler_interface.h"
 #include "tsl/profiler/protobuf/xplane.pb.h"
diff --git a/third_party/xla/third_party/tsl/tsl/profiler/lib/profiler_factory_test.cc b/third_party/xla/third_party/tsl/tsl/profiler/lib/profiler_factory_test.cc
index 94fa5c6868f4b2..a1188b9fa5563d 100644
--- a/third_party/xla/third_party/tsl/tsl/profiler/lib/profiler_factory_test.cc
+++ b/third_party/xla/third_party/tsl/tsl/profiler/lib/profiler_factory_test.cc
@@ -18,8 +18,8 @@ limitations under the License.
 #include <utility>
 
 #include "absl/memory/memory.h"
+#include "absl/status/status.h"
 #include "tsl/platform/macros.h"
-#include "tsl/platform/status.h"
 #include "tsl/platform/test.h"
 #include "tsl/profiler/lib/profiler_interface.h"
 #include "tsl/profiler/protobuf/profiler_options.pb.h"
diff --git a/third_party/xla/third_party/tsl/tsl/profiler/lib/profiler_lock.cc b/third_party/xla/third_party/tsl/tsl/profiler/lib/profiler_lock.cc
index 08ce17a8bebd03..d32ea96fd2bf69 100644
--- a/third_party/xla/third_party/tsl/tsl/profiler/lib/profiler_lock.cc
+++ b/third_party/xla/third_party/tsl/tsl/profiler/lib/profiler_lock.cc
@@ -16,10 +16,10 @@ limitations under the License.
 
 #include <atomic>
 
+#include "absl/status/statusor.h"
 #include "xla/tsl/util/env_var.h"
 #include "tsl/platform/errors.h"
 #include "tsl/platform/macros.h"
-#include "tsl/platform/statusor.h"
 
 namespace tsl {
 namespace profiler {
diff --git a/third_party/xla/third_party/tsl/tsl/profiler/lib/profiler_lock.h b/third_party/xla/third_party/tsl/tsl/profiler/lib/profiler_lock.h
index c4e8b40303f4e1..478dae87b8a399 100644
--- a/third_party/xla/third_party/tsl/tsl/profiler/lib/profiler_lock.h
+++ b/third_party/xla/third_party/tsl/tsl/profiler/lib/profiler_lock.h
@@ -17,6 +17,7 @@ limitations under the License.
 
 #include <utility>
 
+#include "absl/status/statusor.h"
 #include "absl/strings/string_view.h"
 #include "tsl/platform/statusor.h"
 
diff --git a/third_party/xla/third_party/tsl/tsl/profiler/lib/profiler_lock_test.cc b/third_party/xla/third_party/tsl/tsl/profiler/lib/profiler_lock_test.cc
index 4a1d75eb6b1b89..2ddc56fb0b9a8d 100644
--- a/third_party/xla/third_party/tsl/tsl/profiler/lib/profiler_lock_test.cc
+++ b/third_party/xla/third_party/tsl/tsl/profiler/lib/profiler_lock_test.cc
@@ -16,7 +16,7 @@ limitations under the License.
 
 #include <utility>
 
-#include "tsl/platform/statusor.h"
+#include "absl/status/statusor.h"
 #include "tsl/platform/test.h"
 
 namespace tsl {
diff --git a/third_party/xla/third_party/tsl/tsl/profiler/lib/profiler_session.cc b/third_party/xla/third_party/tsl/tsl/profiler/lib/profiler_session.cc
index c54f7969b4154f..f48b5ecb6f3f2c 100644
--- a/third_party/xla/third_party/tsl/tsl/profiler/lib/profiler_session.cc
+++ b/third_party/xla/third_party/tsl/tsl/profiler/lib/profiler_session.cc
@@ -19,10 +19,10 @@ limitations under the License.
 #include <utility>
 
 #include "absl/memory/memory.h"
+#include "absl/status/status.h"
 #include "tsl/platform/errors.h"
 #include "tsl/platform/logging.h"
 #include "tsl/platform/mutex.h"
-#include "tsl/platform/status.h"
 #include "tsl/profiler/protobuf/profiler_options.pb.h"
 #include "tsl/profiler/protobuf/xplane.pb.h"
 
diff --git a/third_party/xla/third_party/tsl/tsl/profiler/lib/profiler_session.h b/third_party/xla/third_party/tsl/tsl/profiler/lib/profiler_session.h
index 7a04f42c188a4e..b503f428ff30d5 100644
--- a/third_party/xla/third_party/tsl/tsl/profiler/lib/profiler_session.h
+++ b/third_party/xla/third_party/tsl/tsl/profiler/lib/profiler_session.h
@@ -19,6 +19,7 @@ limitations under the License.
 #include <memory>
 #include <vector>
 
+#include "absl/status/status.h"
 #include "tsl/platform/mutex.h"
 #include "tsl/platform/platform.h"
 #include "tsl/platform/status.h"

From ed39df35f4eef252adb61a5e24f9a28a7d0281ac Mon Sep 17 00:00:00 2001
From: Edward Schwartz <schwartzedward@google.com>
Date: Sat, 18 May 2024 09:40:22 -0700
Subject: [PATCH 396/478] disable flaky mwms_pjrt_gpu_test_xla_2gpu test

PiperOrigin-RevId: 635064164
---
 tensorflow/python/distribute/BUILD | 1 +
 1 file changed, 1 insertion(+)

diff --git a/tensorflow/python/distribute/BUILD b/tensorflow/python/distribute/BUILD
index 8bcd53d37efb49..5a17d482598288 100644
--- a/tensorflow/python/distribute/BUILD
+++ b/tensorflow/python/distribute/BUILD
@@ -658,6 +658,7 @@ cuda_py_strict_test(
     srcs_version = "PY3",
     tags = [
         "multi_and_single_gpu",
+        "notap",  # TODO(b/341375925): Re-enable this test when flakiness is fixed.
     ],
     xla_enabled = True,
     deps = [

From c57ab0a3a34316bf986ccfb010125b40625c987a Mon Sep 17 00:00:00 2001
From: George Karpenkov <cheshire@google.com>
Date: Sat, 18 May 2024 14:14:58 -0700
Subject: [PATCH 397/478] [XLA] Error out when encountering errors during flag
 parsing

Currently we crash on unknown flags left after XLA_FLAGS parsing, but all the
errors are just silently ignored. This leaves a large footgun which ignores
failed flags.

PiperOrigin-RevId: 635098018
---
 third_party/xla/xla/parse_flags_from_env.cc   | 23 +++----
 third_party/xla/xla/parse_flags_from_env.h    | 17 ++---
 .../xla/xla/parse_flags_from_env_test.cc      | 64 +++++++++++++++----
 .../service/gpu_compilation_environment.cc    |  5 +-
 .../gpu_compilation_environment_test.cc       | 14 ----
 5 files changed, 70 insertions(+), 53 deletions(-)

diff --git a/third_party/xla/xla/parse_flags_from_env.cc b/third_party/xla/xla/parse_flags_from_env.cc
index 0f58671ebff7df..4da2c9a0de2613 100644
--- a/third_party/xla/xla/parse_flags_from_env.cc
+++ b/third_party/xla/xla/parse_flags_from_env.cc
@@ -186,18 +186,15 @@ static absl::flat_hash_map<std::string, EnvArgv>& EnvArgvs() {
 // Used to protect accesses to env_argvs.
 static absl::Mutex env_argv_mu(absl::kConstInit);
 
-bool ParseFlagsFromEnvAndDieIfUnknown(absl::string_view envvar,
+static void DieIfEnvHasUnknownFlagsLeft(absl::string_view envvar);
+
+void ParseFlagsFromEnvAndDieIfUnknown(absl::string_view envvar,
                                       const std::vector<tsl::Flag>& flag_list) {
-  bool parsed_recognized_flags =
-      ParseFlagsFromEnvAndIgnoreUnknown(envvar, flag_list);
-  if (!parsed_recognized_flags) {
-    return false;
-  } else {
-    return !DieIfEnvHasUnknownFlagsLeft(envvar);
-  }
+  ParseFlagsFromEnvAndIgnoreUnknown(envvar, flag_list);
+  DieIfEnvHasUnknownFlagsLeft(envvar);
 }
 
-bool ParseFlagsFromEnvAndIgnoreUnknown(
+void ParseFlagsFromEnvAndIgnoreUnknown(
     absl::string_view envvar, const std::vector<tsl::Flag>& flag_list) {
   absl::MutexLock lock(&env_argv_mu);
   auto* env_argv = &EnvArgvs()[envvar];
@@ -210,10 +207,12 @@ bool ParseFlagsFromEnvAndIgnoreUnknown(
     }
   }
 
-  return tsl::Flags::Parse(&env_argv->argc, env_argv->argv.data(), flag_list);
+  QCHECK(tsl::Flags::Parse(&env_argv->argc, env_argv->argv.data(), flag_list))
+      << "Flag parsing failed.\n"
+      << tsl::Flags::Usage(getenv(std::string(envvar).c_str()), flag_list);
 }
 
-bool DieIfEnvHasUnknownFlagsLeft(absl::string_view envvar) {
+static void DieIfEnvHasUnknownFlagsLeft(absl::string_view envvar) {
   absl::MutexLock lock(&env_argv_mu);
   auto* env_argv = &EnvArgvs()[envvar];
   SetArgvFromEnv(envvar, env_argv);
@@ -225,9 +224,7 @@ bool DieIfEnvHasUnknownFlagsLeft(absl::string_view envvar) {
     LOG(QFATAL) << "Unknown flag" << (unknown_flags.size() > 1 ? "s" : "")
                 << " in " << envvar << ": "
                 << absl::StrJoin(unknown_flags, " ");
-    return true;
   }
-  return false;
 }
 
 // Testing only.
diff --git a/third_party/xla/xla/parse_flags_from_env.h b/third_party/xla/xla/parse_flags_from_env.h
index 01d476f22fa3dc..e135a664d93b47 100644
--- a/third_party/xla/xla/parse_flags_from_env.h
+++ b/third_party/xla/xla/parse_flags_from_env.h
@@ -57,26 +57,21 @@ limitations under the License.
 namespace xla {
 
 // Calls tsl::Flags::Parse(argc, argv, flag_list) against any as yet
-// unrecognized flags passed in the environment variable `envvar`, and returns
-// its return value.
+// unrecognized flags passed in the environment variable `envvar`.
 //
-// Raises a fatal error if any flags in `envvar` were not recognized.
-bool ParseFlagsFromEnvAndDieIfUnknown(absl::string_view envvar,
+// Raises a fatal error if any flags in `envvar` were not recognized, or if flag
+// parsing failed.
+void ParseFlagsFromEnvAndDieIfUnknown(absl::string_view envvar,
                                       const std::vector<tsl::Flag>& flag_list);
 
 // Calls tsl::Flags::Parse(argc, argv, flag_list) against any as yet
 // unrecognized flags passed in the environment variable `envvar`, and returns
 // its return value.
 //
-// Best effort and ignores unknown flags
-bool ParseFlagsFromEnvAndIgnoreUnknown(absl::string_view envvar,
+// Ignores unknown flags, raises a fatal if flag parsing failed.
+void ParseFlagsFromEnvAndIgnoreUnknown(absl::string_view envvar,
                                        const std::vector<tsl::Flag>& flag_list);
 
-// Raises a fatal error if there are unrecognized flags left in the specified
-// environment variable `envvar`, besides the fake argv[0]. Returns false if
-// there are no unrecognized flags.
-bool DieIfEnvHasUnknownFlagsLeft(absl::string_view envvar);
-
 // Used only for testing.  Not to be used by clients.
 void ResetFlagsFromEnvForTesting(absl::string_view envvar, int** pargc,
                                  std::vector<char*>** pargv);
diff --git a/third_party/xla/xla/parse_flags_from_env_test.cc b/third_party/xla/xla/parse_flags_from_env_test.cc
index f00cb309c12a96..0f68b41b7510be 100644
--- a/third_party/xla/xla/parse_flags_from_env_test.cc
+++ b/third_party/xla/xla/parse_flags_from_env_test.cc
@@ -53,12 +53,11 @@ static void TestParseFlagsFromEnv(const char* msg) {
       tsl::Flag("single_quoted", &single_quoted, ""),
       tsl::Flag("double_quoted", &double_quoted, ""),
   };
-  bool parsed_ok = ParseFlagsFromEnvAndDieIfUnknown("TF_XLA_FLAGS", flag_list);
+  ParseFlagsFromEnvAndDieIfUnknown("TF_XLA_FLAGS", flag_list);
   CHECK_EQ(*pargc, 1) << msg;
   const std::vector<char*>& argv_second = *pargv;
   CHECK_NE(argv_second[0], nullptr) << msg;
   CHECK_EQ(argv_second[1], nullptr) << msg;
-  CHECK(parsed_ok) << msg;
   CHECK(simple) << msg;
   CHECK_EQ(with_value, "a_value") << msg;
   CHECK_EQ(embedded_quotes, "single'double\"") << msg;
@@ -156,6 +155,56 @@ TEST(ParseFlagsFromEnv, EnvAndFlag) {
   }
 }
 
+TEST(ParseFlagsFromEnv, ErrorOutOnFlagFailure) {
+  const char* env = "--int_flag=3parsefailure";
+
+  if (env == nullptr) {
+    // Might be set from previous tests.
+    tsl::unsetenv("TF_XLA_FLAGS");
+  } else {
+    tsl::setenv("TF_XLA_FLAGS", env, /*overwrite=*/true);
+  }
+  tsl::SubProcess child;
+  std::vector<std::string> argv;
+  argv.push_back(binary_name);
+  argv.push_back("--recursing");
+  child.SetProgram(binary_name, argv);
+  child.SetChannelAction(tsl::CHAN_STDOUT, tsl::ACTION_PIPE);
+  child.SetChannelAction(tsl::CHAN_STDERR, tsl::ACTION_PIPE);
+  EXPECT_TRUE(child.Start());
+  std::string stdout_str;
+  std::string stderr_str;
+
+  // Expecting failure.
+  int child_status = child.Communicate(nullptr, &stdout_str, &stderr_str);
+  EXPECT_NE(child_status, 0);
+}
+
+TEST(ParseFlagsFromEnv, ErrorOutOnUnknownFlag) {
+  const char* env = "--int_flag=3 --unknown_flag=value";
+
+  if (env == nullptr) {
+    // Might be set from previous tests.
+    tsl::unsetenv("TF_XLA_FLAGS");
+  } else {
+    tsl::setenv("TF_XLA_FLAGS", env, /*overwrite=*/true);
+  }
+  tsl::SubProcess child;
+  std::vector<std::string> argv;
+  argv.push_back(binary_name);
+  argv.push_back("--recursing");
+  child.SetProgram(binary_name, argv);
+  child.SetChannelAction(tsl::CHAN_STDOUT, tsl::ACTION_PIPE);
+  child.SetChannelAction(tsl::CHAN_STDERR, tsl::ACTION_PIPE);
+  EXPECT_TRUE(child.Start());
+  std::string stdout_str;
+  std::string stderr_str;
+
+  // Expecting failure.
+  int child_status = child.Communicate(nullptr, &stdout_str, &stderr_str);
+  EXPECT_NE(child_status, 0);
+}
+
 }  // namespace xla
 
 int main(int argc, char* argv[]) {
@@ -169,15 +218,8 @@ int main(int argc, char* argv[]) {
       tsl::Flag("int_flag", &int_flag, "An integer flag to test with"),
   };
   std::string usage = tsl::Flags::Usage(argv[0], flag_list);
-  bool parse_ok =
-      xla::ParseFlagsFromEnvAndDieIfUnknown("TF_XLA_FLAGS", flag_list);
-  if (!parse_ok) {
-    LOG(QFATAL) << "can't parse from environment\n" << usage;
-  }
-  parse_ok = tsl::Flags::Parse(&argc, argv, flag_list);
-  if (!parse_ok) {
-    LOG(QFATAL) << usage;
-  }
+  xla::ParseFlagsFromEnvAndDieIfUnknown("TF_XLA_FLAGS", flag_list);
+  tsl::Flags::Parse(&argc, argv, flag_list);
   if (recursing) {
     printf("%d\n", int_flag);
     exit(0);
diff --git a/third_party/xla/xla/service/gpu_compilation_environment.cc b/third_party/xla/xla/service/gpu_compilation_environment.cc
index 4376c975544c0c..c525726fa3f033 100644
--- a/third_party/xla/xla/service/gpu_compilation_environment.cc
+++ b/third_party/xla/xla/service/gpu_compilation_environment.cc
@@ -66,10 +66,7 @@ absl::StatusOr<GpuCompilationEnvironment> CreateGpuCompEnvFromEnvVar() {
   GpuCompilationEnvironment env;
   std::vector<tsl::Flag> flag_objects;
   InitializeFlagsForGpuCompEnv(&flag_objects, &env);
-  bool result = ParseFlagsFromEnvAndIgnoreUnknown("XLA_FLAGS", flag_objects);
-  if (!result) {
-    return InvalidArgument("Could not parse XLA_FLAGS.");
-  }
+  ParseFlagsFromEnvAndIgnoreUnknown("XLA_FLAGS", flag_objects);
   return env;
 }
 
diff --git a/third_party/xla/xla/service/gpu_compilation_environment_test.cc b/third_party/xla/xla/service/gpu_compilation_environment_test.cc
index 85dde778f8aee3..170d928b2665d1 100644
--- a/third_party/xla/xla/service/gpu_compilation_environment_test.cc
+++ b/third_party/xla/xla/service/gpu_compilation_environment_test.cc
@@ -71,13 +71,6 @@ TEST(CreateGpuCompEnvFromFlagStringsTest, InvalidFlagName) {
   ASSERT_EQ(flags.size(), 1);
 }
 
-TEST(CreateGpuCompEnvFromFlagStringsTest, InvalidFlagValue) {
-  std::vector<std::string> flags = {"--dummy_flag=foo"};
-
-  EXPECT_THAT(CreateGpuCompEnvFromFlagStrings(flags, /*strict=*/false),
-              StatusIs(tsl::error::INVALID_ARGUMENT));
-}
-
 TEST(CreateGpuCompEnvFromEnvVarTest, ValidFlags) {
   set_xla_flags_env_var("--dummy_flag=4");
 
@@ -87,13 +80,6 @@ TEST(CreateGpuCompEnvFromEnvVarTest, ValidFlags) {
   ASSERT_EQ(gpu_comp_env.dummy_flag(), 4);
 }
 
-TEST(CreateGpuCompEnvFromEnvVarTest, InvalidFlagValue) {
-  set_xla_flags_env_var("--dummy_flag=foo");
-
-  EXPECT_THAT(CreateGpuCompEnvFromEnvVar(),
-              StatusIs(tsl::error::INVALID_ARGUMENT));
-}
-
 TEST(InitializeMissingFieldsFromXLAFlagsTest, BothProtoAndEnvVarUnset) {
   set_xla_flags_env_var("");
   GpuCompilationEnvironment env;

From 626f974e888b18fb14ce2d7bf9d047c8d799c004 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Sun, 19 May 2024 02:02:23 -0700
Subject: [PATCH 398/478] compat: Update forward compatibility horizon to
 2024-05-19

PiperOrigin-RevId: 635189458
---
 tensorflow/python/compat/compat.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorflow/python/compat/compat.py b/tensorflow/python/compat/compat.py
index 5dd1128fc1715a..4ab2aa5a2a9cd1 100644
--- a/tensorflow/python/compat/compat.py
+++ b/tensorflow/python/compat/compat.py
@@ -29,7 +29,7 @@
 # This value changes every day with an automatic CL. It can be modified in code
 # via `forward_compatibility_horizon()` or with the environment variable
 # TF_FORWARD_COMPATIBILITY_DELTA_DAYS, which is added to the compatibility date.
-_FORWARD_COMPATIBILITY_HORIZON = datetime.date(2024, 5, 18)
+_FORWARD_COMPATIBILITY_HORIZON = datetime.date(2024, 5, 19)
 _FORWARD_COMPATIBILITY_DELTA_DAYS_VAR_NAME = "TF_FORWARD_COMPATIBILITY_DELTA_DAYS"
 _FORWARD_COMPATIBILITY_DATE_NUMBER = None
 

From e667913c92c33045a28f4cb4259a0e8d7c36647b Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Sun, 19 May 2024 02:03:48 -0700
Subject: [PATCH 399/478] Update GraphDef version to 1867.

PiperOrigin-RevId: 635189739
---
 tensorflow/core/public/version.h | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorflow/core/public/version.h b/tensorflow/core/public/version.h
index 51c8ef32c1956a..bee0841d5d51bf 100644
--- a/tensorflow/core/public/version.h
+++ b/tensorflow/core/public/version.h
@@ -108,7 +108,7 @@ limitations under the License.
 
 #define TF_GRAPH_DEF_VERSION_MIN_PRODUCER 0
 #define TF_GRAPH_DEF_VERSION_MIN_CONSUMER 0
-#define TF_GRAPH_DEF_VERSION 1866  // Updated: 2024/5/18
+#define TF_GRAPH_DEF_VERSION 1867  // Updated: 2024/5/19
 
 // Checkpoint compatibility versions (the versions field in SavedSliceMeta).
 //

From 08d2478fb123fc6cd1744ada4e3b709b6fd1a557 Mon Sep 17 00:00:00 2001
From: Clive Verghese <cliveverghese@google.com>
Date: Sun, 19 May 2024 10:52:52 -0700
Subject: [PATCH 400/478] Cleanup dependency on tracing.h

PiperOrigin-RevId: 635253325
---
 tensorflow/core/common_runtime/BUILD          |  2 ++
 .../base_collective_executor.cc               |  1 -
 .../core/common_runtime/direct_session.cc     |  1 -
 tensorflow/core/common_runtime/executor.cc    | 22 +++++++++----------
 .../core/common_runtime/executor_test.cc      |  1 -
 .../core/common_runtime/process_util.cc       | 11 +++++-----
 .../single_threaded_executor_test.cc          |  1 -
 .../core/common_runtime/threadpool_device.cc  |  1 -
 8 files changed, 19 insertions(+), 21 deletions(-)

diff --git a/tensorflow/core/common_runtime/BUILD b/tensorflow/core/common_runtime/BUILD
index 62617578abb5b2..7fea998611bb50 100644
--- a/tensorflow/core/common_runtime/BUILD
+++ b/tensorflow/core/common_runtime/BUILD
@@ -683,6 +683,7 @@ cc_library(
         "@com_google_absl//absl/strings",
         "@com_google_absl//absl/time",
         "@com_google_absl//absl/types:optional",
+        "@local_tsl//tsl/platform:env",
     ],
     alwayslink = 1,
 )
@@ -1572,6 +1573,7 @@ cc_library(
         "//tensorflow/core:framework",
         "//tensorflow/core:lib",
         "//tensorflow/core:lib_internal",
+        "@local_tsl//tsl/platform:env",
     ],
 )
 
diff --git a/tensorflow/core/common_runtime/base_collective_executor.cc b/tensorflow/core/common_runtime/base_collective_executor.cc
index 3136f6bab840b1..4c7ec29953cbbc 100644
--- a/tensorflow/core/common_runtime/base_collective_executor.cc
+++ b/tensorflow/core/common_runtime/base_collective_executor.cc
@@ -35,7 +35,6 @@ limitations under the License.
 #include "tensorflow/core/lib/strings/strcat.h"
 #include "tensorflow/core/platform/macros.h"
 #include "tensorflow/core/platform/refcount.h"
-#include "tensorflow/core/platform/tracing.h"
 #include "tensorflow/core/platform/types.h"
 #include "tensorflow/core/profiler/lib/connected_traceme.h"
 #include "tensorflow/core/profiler/lib/scoped_memory_debug_annotation.h"
diff --git a/tensorflow/core/common_runtime/direct_session.cc b/tensorflow/core/common_runtime/direct_session.cc
index 50285f87b2283c..46a9fd5943fb40 100644
--- a/tensorflow/core/common_runtime/direct_session.cc
+++ b/tensorflow/core/common_runtime/direct_session.cc
@@ -73,7 +73,6 @@ limitations under the License.
 #include "tensorflow/core/platform/cpu_info.h"
 #include "tensorflow/core/platform/logging.h"
 #include "tensorflow/core/platform/mutex.h"
-#include "tensorflow/core/platform/tracing.h"
 #include "tensorflow/core/platform/types.h"
 #include "tensorflow/core/profiler/lib/connected_traceme.h"
 #include "tensorflow/core/profiler/lib/device_profiler_session.h"
diff --git a/tensorflow/core/common_runtime/executor.cc b/tensorflow/core/common_runtime/executor.cc
index 026719175535f5..cb6186e6ff9b84 100644
--- a/tensorflow/core/common_runtime/executor.cc
+++ b/tensorflow/core/common_runtime/executor.cc
@@ -70,7 +70,6 @@ limitations under the License.
 #include "tensorflow/core/platform/profile_utils/cpu_utils.h"
 #include "tensorflow/core/platform/status.h"
 #include "tensorflow/core/platform/thread_annotations.h"
-#include "tensorflow/core/platform/tracing.h"
 #include "tensorflow/core/platform/types.h"
 #include "tensorflow/core/profiler/lib/annotated_traceme.h"
 #include "tensorflow/core/profiler/lib/connected_traceme.h"
@@ -82,6 +81,7 @@ limitations under the License.
 #include "tensorflow/core/util/determinism.h"
 #include "tensorflow/core/util/managed_stack_trace.h"
 #include "tensorflow/core/util/tensor_slice_reader_cache.h"
+#include "tsl/platform/tracing.h"
 
 namespace tensorflow {
 
@@ -381,7 +381,7 @@ class ExecutorState {
   // Step-local container.
   ScopedStepContainer* step_container_;
   StepStatsCollectorInterface* const stats_collector_;
-  const tracing::EventCollector* const event_collector_;
+  const tsl::tracing::EventCollector* const event_collector_;
   Context context_;
 
   // QUESTION: Make it a checkpoint::TensorSliceReaderCacheWrapper
@@ -435,8 +435,8 @@ ExecutorState<PropagatorStateType>::ExecutorState(
       tensor_store_(args.tensor_store),
       step_container_(args.step_container),
       stats_collector_(args.stats_collector),
-      event_collector_(
-          tracing::GetEventCollector(tracing::EventCategory::kCompute)),
+      event_collector_(tsl::tracing::GetEventCollector(
+          tsl::tracing::EventCategory::kCompute)),
       context_(ContextKind::kThread),
       slice_reader_cache_(new checkpoint::TensorSliceReaderCacheWrapper),
       call_frame_(args.call_frame),
@@ -563,15 +563,15 @@ struct ExecutorState<PropagatorStateType>::AsyncState {
 
 // Returns true if `item` might be traced by the given trace and event
 // collectors. Returns false only if `item` definitely will not be traced.
-bool MightTrace(const tracing::EventCollector* event_collector,
+bool MightTrace(const tsl::tracing::EventCollector* event_collector,
                 bool is_expensive) {
   // Tracing will only be enabled if either `event_collector` is non null,
   // or `trace_collector` is non-null and enabled for this particular kernel.
   // Although `profiler::TraceMe`, `profiler::ScopedAnnotation`, and
-  // `tracing::ScopedRegion` check subsets of these properties internally in
-  // their constructors, the cost of passing the necessary arguments to them can
-  // be significant, so we avoid constructing them in the common case (when we
-  // know they will not be used).
+  // `tsl::tracing::ScopedRegion` check subsets of these properties internally
+  // in their constructors, the cost of passing the necessary arguments to them
+  // can be significant, so we avoid constructing them in the common case (when
+  // we know they will not be used).
   if (event_collector != nullptr) {
     return true;
   }
@@ -594,8 +594,8 @@ Status ExecutorState<PropagatorStateType>::ProcessSync(
   const bool is_expensive = kernel_stats_->IsExpensive(item);
 
   if (TF_PREDICT_FALSE(MightTrace(event_collector_, is_expensive))) {
-    tracing::ScopedRegion region(tracing::EventCategory::kCompute,
-                                 op_kernel->name_view());
+    tsl::tracing::ScopedRegion region(tsl::tracing::EventCategory::kCompute,
+                                      op_kernel->name_view());
     profiler::AnnotatedTraceMe activity(
         [op_kernel, &ctx] {
           return op_kernel->TraceString(
diff --git a/tensorflow/core/common_runtime/executor_test.cc b/tensorflow/core/common_runtime/executor_test.cc
index 6a847b3f1b2f9f..cb683e21a707f9 100644
--- a/tensorflow/core/common_runtime/executor_test.cc
+++ b/tensorflow/core/common_runtime/executor_test.cc
@@ -46,7 +46,6 @@ limitations under the License.
 #include "tensorflow/core/platform/strcat.h"
 #include "tensorflow/core/platform/test.h"
 #include "tensorflow/core/platform/test_benchmark.h"
-#include "tensorflow/core/platform/tracing.h"
 #include "tensorflow/core/public/session_options.h"
 
 namespace tensorflow {
diff --git a/tensorflow/core/common_runtime/process_util.cc b/tensorflow/core/common_runtime/process_util.cc
index 53d73952219ecb..e0fa771c4b8280 100644
--- a/tensorflow/core/common_runtime/process_util.cc
+++ b/tensorflow/core/common_runtime/process_util.cc
@@ -26,9 +26,9 @@ limitations under the License.
 #include "tensorflow/core/platform/byte_order.h"
 #include "tensorflow/core/platform/cpu_info.h"
 #include "tensorflow/core/platform/logging.h"
-#include "tensorflow/core/platform/tracing.h"
 #include "tensorflow/core/platform/types.h"
 #include "tensorflow/core/util/util.h"
+#include "tsl/platform/tracing.h"
 
 namespace tensorflow {
 
@@ -166,14 +166,15 @@ thread::ThreadPool* NewThreadPoolFromSessionOptions(
 }
 
 void SchedClosure(absl::AnyInvocable<void()> closure) {
-  if (!tracing::EventCollector::IsEnabled()) {
+  if (!tsl::tracing::EventCollector::IsEnabled()) {
     return Env::Default()->SchedClosure(std::move(closure));
   }
-  uint64 id = tracing::GetUniqueArg();
-  tracing::RecordEvent(tracing::EventCategory::kScheduleClosure, id);
+  uint64 id = tsl::tracing::GetUniqueArg();
+  tsl::tracing::RecordEvent(tsl::tracing::EventCategory::kScheduleClosure, id);
 
   Env::Default()->SchedClosure([id, closure = std::move(closure)]() mutable {
-    tracing::ScopedRegion region(tracing::EventCategory::kRunClosure, id);
+    tsl::tracing::ScopedRegion region(tsl::tracing::EventCategory::kRunClosure,
+                                      id);
     closure();
   });
 }
diff --git a/tensorflow/core/common_runtime/single_threaded_executor_test.cc b/tensorflow/core/common_runtime/single_threaded_executor_test.cc
index 04fcd51647efc8..a53e65d7a7a513 100644
--- a/tensorflow/core/common_runtime/single_threaded_executor_test.cc
+++ b/tensorflow/core/common_runtime/single_threaded_executor_test.cc
@@ -41,7 +41,6 @@ limitations under the License.
 #include "tensorflow/core/platform/logging.h"
 #include "tensorflow/core/platform/test.h"
 #include "tensorflow/core/platform/test_benchmark.h"
-#include "tensorflow/core/platform/tracing.h"
 #include "tensorflow/core/public/session_options.h"
 
 namespace tensorflow {
diff --git a/tensorflow/core/common_runtime/threadpool_device.cc b/tensorflow/core/common_runtime/threadpool_device.cc
index cd77afcf53dbd3..a06e9e90b7ba16 100644
--- a/tensorflow/core/common_runtime/threadpool_device.cc
+++ b/tensorflow/core/common_runtime/threadpool_device.cc
@@ -41,7 +41,6 @@ info. It does not have any negative impact on performance. */
 #include "tensorflow/core/framework/types.h"
 #include "tensorflow/core/graph/types.h"
 #include "tensorflow/core/lib/hash/hash.h"
-#include "tensorflow/core/platform/tracing.h"
 #include "tensorflow/core/platform/types.h"
 #include "tensorflow/core/public/session_options.h"
 #include "tensorflow/core/util/port.h"

From 8776f75dd11c05a2f422eecf659844232abb967c Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Sun, 19 May 2024 11:08:14 -0700
Subject: [PATCH 401/478] Automated Code Change

PiperOrigin-RevId: 635255967
---
 .../core/common_runtime/optimized_function_graph_info.h      | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/tensorflow/core/common_runtime/optimized_function_graph_info.h b/tensorflow/core/common_runtime/optimized_function_graph_info.h
index b2bd9af5bb1c5a..dd05b026ebfed2 100644
--- a/tensorflow/core/common_runtime/optimized_function_graph_info.h
+++ b/tensorflow/core/common_runtime/optimized_function_graph_info.h
@@ -71,9 +71,10 @@ struct OptimizedFunctionGraphInfo {
   OptimizedFunctionGraphInfo(OptimizedFunctionGraphInfo& info) = delete;
   OptimizedFunctionGraphInfo& operator=(OptimizedFunctionGraphInfo& info) =
       delete;
-  OptimizedFunctionGraphInfo(OptimizedFunctionGraphInfo&& info) = default;
-  OptimizedFunctionGraphInfo& operator=(OptimizedFunctionGraphInfo&& info) =
+  OptimizedFunctionGraphInfo(OptimizedFunctionGraphInfo&& info) noexcept =
       default;
+  OptimizedFunctionGraphInfo& operator=(
+      OptimizedFunctionGraphInfo&& info) noexcept = default;
 
   // Converts from the struct to OptimizedFunctionGraph proto.
   static OptimizedFunctionGraph ToProto(const OptimizedFunctionGraphInfo& info);

From e33a5ea470d1ef02107b595737d23670f7a1eda4 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Sun, 19 May 2024 14:12:09 -0700
Subject: [PATCH 402/478] Migrate coord grpc client and service to use absl
 libraries (mutex, thread annotation, status) directly.

PiperOrigin-RevId: 635277739
---
 .../rpc/coordination/BUILD                    | 15 ++++++++----
 .../coordination/grpc_coordination_client.cc  | 24 ++++++++++++-------
 .../grpc_coordination_service_impl.cc         |  6 ++---
 .../grpc_coordination_service_impl.h          | 14 ++++++-----
 4 files changed, 36 insertions(+), 23 deletions(-)

diff --git a/third_party/xla/xla/tsl/distributed_runtime/rpc/coordination/BUILD b/third_party/xla/xla/tsl/distributed_runtime/rpc/coordination/BUILD
index e21bcbd7370868..838c2cbdf5ab5c 100644
--- a/third_party/xla/xla/tsl/distributed_runtime/rpc/coordination/BUILD
+++ b/third_party/xla/xla/tsl/distributed_runtime/rpc/coordination/BUILD
@@ -21,12 +21,14 @@ cc_library(
         "//xla/tsl/distributed_runtime/rpc:grpc_client_cq_tag",
         "//xla/tsl/distributed_runtime/rpc:grpc_state",
         "//xla/tsl/distributed_runtime/rpc:grpc_util",
-        "@local_tsl//tsl/platform:mutex",
+        "@com_google_absl//absl/base:core_headers",
+        "@com_google_absl//absl/log",
+        "@com_google_absl//absl/synchronization",
+        "@local_tsl//tsl/platform:env",
         "@local_tsl//tsl/platform:protobuf",
         "@local_tsl//tsl/platform:status",
-        "@local_tsl//tsl/platform:thread_annotations",
         "@local_tsl//tsl/protobuf:coordination_service_proto_cc",
-    ],
+    ] + tsl_grpc_cc_dependencies(),
 )
 
 cc_library(
@@ -34,14 +36,17 @@ cc_library(
     srcs = ["grpc_coordination_service_impl.cc"],
     hdrs = ["grpc_coordination_service_impl.h"],
     deps = [
+        "//xla/tsl/distributed_runtime/coordination:coordination_service",
         "//xla/tsl/distributed_runtime/coordination:coordination_service_agent",
         "//xla/tsl/distributed_runtime/coordination:coordination_service_rpc_handler",
         "//xla/tsl/distributed_runtime/rpc:async_service_interface",
         "//xla/tsl/distributed_runtime/rpc:grpc_call",
         "//xla/tsl/distributed_runtime/rpc:grpc_util",
+        "@com_google_absl//absl/base:core_headers",
+        "@com_google_absl//absl/log",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/synchronization",
         "@local_tsl//tsl/platform:env",
-        "@local_tsl//tsl/platform:mutex",
-        "@local_tsl//tsl/platform:thread_annotations",
         "@local_tsl//tsl/protobuf:coordination_service_cc_grpc_proto",
         "@local_tsl//tsl/protobuf:coordination_service_proto_cc",
     ] + tsl_grpc_cc_dependencies(),
diff --git a/third_party/xla/xla/tsl/distributed_runtime/rpc/coordination/grpc_coordination_client.cc b/third_party/xla/xla/tsl/distributed_runtime/rpc/coordination/grpc_coordination_client.cc
index 7c2c520362bc16..db74a713a7f132 100644
--- a/third_party/xla/xla/tsl/distributed_runtime/rpc/coordination/grpc_coordination_client.cc
+++ b/third_party/xla/xla/tsl/distributed_runtime/rpc/coordination/grpc_coordination_client.cc
@@ -15,22 +15,28 @@ limitations under the License.
 
 #include "xla/tsl/distributed_runtime/rpc/coordination/grpc_coordination_client.h"
 
+#include <cstddef>
 #include <memory>
 #include <string>
 #include <unordered_map>
 #include <utility>
 #include <vector>
 
+#include "grpcpp/channel.h"
+#include "grpcpp/completion_queue.h"
+#include "grpcpp/generic/generic_stub.h"
+#include "absl/base/thread_annotations.h"
+#include "absl/log/log.h"
+#include "absl/synchronization/mutex.h"
 #include "xla/tsl/distributed_runtime/call_options.h"
 #include "xla/tsl/distributed_runtime/coordination/coordination_client.h"
 #include "xla/tsl/distributed_runtime/rpc/grpc_channel.h"
 #include "xla/tsl/distributed_runtime/rpc/grpc_client_cq_tag.h"
 #include "xla/tsl/distributed_runtime/rpc/grpc_state.h"
 #include "xla/tsl/distributed_runtime/rpc/grpc_util.h"
-#include "tsl/platform/mutex.h"
+#include "tsl/platform/env.h"
 #include "tsl/platform/protobuf.h"
 #include "tsl/platform/status.h"
-#include "tsl/platform/thread_annotations.h"
 #include "tsl/protobuf/coordination_service.pb.h"
 
 namespace tsl {
@@ -281,7 +287,7 @@ class GrpcCoordinationClientCache : public CoordinationClientCache {
   ~GrpcCoordinationClientCache() override = default;
 
   CoordinationClient* GetClient(const std::string& target) override {
-    mutex_lock l(clients_mu_);
+    absl::MutexLock l(&clients_mu_);
     auto it = clients_.find(target);
     if (it == clients_.end()) {
       SharedGrpcChannelPtr channel = channel_cache_->FindWorkerChannel(target);
@@ -306,15 +312,15 @@ class GrpcCoordinationClientCache : public CoordinationClientCache {
   }
 
  private:
-  mutex assignment_mu_;
+  absl::Mutex assignment_mu_;
   std::unordered_map<std::string, size_t> target_assignments_
-      TF_GUARDED_BY(assignment_mu_);
-  size_t next_round_robin_assignment_ TF_GUARDED_BY(assignment_mu_);
+      ABSL_GUARDED_BY(assignment_mu_);
+  size_t next_round_robin_assignment_ ABSL_GUARDED_BY(assignment_mu_);
 
   size_t AssignClientToThread(const std::string& target) {
     // Round-robin target assignment, but keeps the same target on the same
     // polling thread always, as this is important for gRPC performance
-    mutex_lock lock(assignment_mu_);
+    absl::MutexLock l(&assignment_mu_);
     auto it = target_assignments_.find(target);
     if (it == target_assignments_.end()) {
       it = target_assignments_
@@ -326,9 +332,9 @@ class GrpcCoordinationClientCache : public CoordinationClientCache {
   }
 
   std::shared_ptr<GrpcChannelCache> channel_cache_;
-  mutable mutex clients_mu_;
+  mutable absl::Mutex clients_mu_;
   std::unordered_map<std::string, std::unique_ptr<CoordinationClient>> clients_
-      TF_GUARDED_BY(clients_mu_);
+      ABSL_GUARDED_BY(clients_mu_);
   std::vector<GrpcCoordinationClientThread> threads_;
 };
 
diff --git a/third_party/xla/xla/tsl/distributed_runtime/rpc/coordination/grpc_coordination_service_impl.cc b/third_party/xla/xla/tsl/distributed_runtime/rpc/coordination/grpc_coordination_service_impl.cc
index 86ef14c89f1d2b..72160cf8d1b0a3 100644
--- a/third_party/xla/xla/tsl/distributed_runtime/rpc/coordination/grpc_coordination_service_impl.cc
+++ b/third_party/xla/xla/tsl/distributed_runtime/rpc/coordination/grpc_coordination_service_impl.cc
@@ -15,7 +15,7 @@ limitations under the License.
 
 #include "xla/tsl/distributed_runtime/rpc/coordination/grpc_coordination_service_impl.h"
 
-#include "tsl/platform/mutex.h"
+#include "absl/synchronization/mutex.h"
 #include "tsl/platform/threadpool.h"
 
 namespace tsl {
@@ -30,7 +30,7 @@ GrpcCoordinationServiceImpl::GrpcCoordinationServiceImpl(
 void GrpcCoordinationServiceImpl::HandleRPCsLoop() {
 #define ENQUEUE_REQUEST(method)                                               \
   do {                                                                        \
-    tf_shared_lock l(shutdown_mu_);                                           \
+    absl::ReaderMutexLock l(&shutdown_mu_);                                   \
     if (shutdown_) {                                                          \
       continue;                                                               \
     }                                                                         \
@@ -80,7 +80,7 @@ void GrpcCoordinationServiceImpl::HandleRPCsLoop() {
 }
 
 void GrpcCoordinationServiceImpl::Shutdown() {
-  mutex_lock l(shutdown_mu_);
+  absl::MutexLock l(&shutdown_mu_);
   shutdown_ = true;
   // This enqueues a special event (with a null tag) that causes the completion
   // queue to be shut down on the polling thread.
diff --git a/third_party/xla/xla/tsl/distributed_runtime/rpc/coordination/grpc_coordination_service_impl.h b/third_party/xla/xla/tsl/distributed_runtime/rpc/coordination/grpc_coordination_service_impl.h
index 21441ad290a26c..89574550e48616 100644
--- a/third_party/xla/xla/tsl/distributed_runtime/rpc/coordination/grpc_coordination_service_impl.h
+++ b/third_party/xla/xla/tsl/distributed_runtime/rpc/coordination/grpc_coordination_service_impl.h
@@ -21,13 +21,15 @@ limitations under the License.
 #include "grpcpp/alarm.h"
 #include "grpcpp/completion_queue.h"
 #include "grpcpp/server_builder.h"
+#include "absl/base/thread_annotations.h"
+#include "absl/status/status.h"
+#include "absl/synchronization/mutex.h"
+#include "xla/tsl/distributed_runtime/coordination/coordination_service.h"
 #include "xla/tsl/distributed_runtime/coordination/coordination_service_agent.h"
 #include "xla/tsl/distributed_runtime/coordination/coordination_service_rpc_handler.h"
 #include "xla/tsl/distributed_runtime/rpc/async_service_interface.h"
 #include "xla/tsl/distributed_runtime/rpc/grpc_call.h"
 #include "xla/tsl/distributed_runtime/rpc/grpc_util.h"
-#include "tsl/platform/mutex.h"
-#include "tsl/platform/thread_annotations.h"
 #include "tsl/platform/threadpool.h"
 #include "tsl/protobuf/coordination_service.grpc.pb.h"
 #include "tsl/protobuf/coordination_service.pb.h"
@@ -59,10 +61,10 @@ class GrpcCoordinationServiceImpl : public AsyncServiceInterface {
 #define HANDLER(method)                                                       \
   void method##Handler(CoordCall<tensorflow::method##Request,                 \
                                  tensorflow::method##Response>* call) {       \
-    tf_shared_lock l(shutdown_mu_);                                           \
+    absl::ReaderMutexLock l(&shutdown_mu_);                                   \
     if (shutdown_) {                                                          \
       call->SendResponse(ToGrpcStatus(                                        \
-          errors::Internal("Coordination service has been shut down.")));     \
+          absl::InternalError("Coordination service has been shut down.")));  \
       return;                                                                 \
     }                                                                         \
     compute_pool_.Schedule([this, call]() {                                   \
@@ -101,8 +103,8 @@ class GrpcCoordinationServiceImpl : public AsyncServiceInterface {
   thread::ThreadPool& compute_pool_;
   CoordinationServiceRpcHandler rpc_handler_;
 
-  mutex shutdown_mu_;
-  bool shutdown_ TF_GUARDED_BY(shutdown_mu_);
+  absl::Mutex shutdown_mu_;
+  bool shutdown_ ABSL_GUARDED_BY(shutdown_mu_);
   std::unique_ptr<::grpc::Alarm> shutdown_alarm_;
 
   std::unique_ptr<::grpc::ServerCompletionQueue> cq_;

From ac1b250ac2bb8ccf8cf8bbe3081afa9e4cb36319 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Sun, 19 May 2024 18:27:52 -0700
Subject: [PATCH 403/478] Refactor common test code to a util

PiperOrigin-RevId: 635307400
---
 tensorflow/core/tfrt/ifrt/BUILD               |  47 ++-
 .../tfrt/ifrt/ifrt_serving_executable_test.cc | 329 ++----------------
 .../ifrt/ifrt_serving_executable_test_util.cc |  95 +++++
 .../ifrt/ifrt_serving_executable_test_util.h  |  82 +++++
 tensorflow/core/tfrt/kernels/BUILD            |  13 +-
 .../tfrt/kernels/ifrt_program_ops_test.cc     |  68 +---
 6 files changed, 253 insertions(+), 381 deletions(-)
 create mode 100644 tensorflow/core/tfrt/ifrt/ifrt_serving_executable_test_util.cc
 create mode 100644 tensorflow/core/tfrt/ifrt/ifrt_serving_executable_test_util.h

diff --git a/tensorflow/core/tfrt/ifrt/BUILD b/tensorflow/core/tfrt/ifrt/BUILD
index 0bfb4aa0cb5752..48b9e6303e1c27 100644
--- a/tensorflow/core/tfrt/ifrt/BUILD
+++ b/tensorflow/core/tfrt/ifrt/BUILD
@@ -299,6 +299,41 @@ cc_library(
     ],
 )
 
+cc_library(
+    name = "ifrt_serving_executable_test_util",
+    testonly = True,
+    srcs = ["ifrt_serving_executable_test_util.cc"],
+    hdrs = ["ifrt_serving_executable_test_util.h"],
+    data = [
+        "//tensorflow/core/tfrt/ifrt/testdata",
+    ],
+    deps = [
+        ":ifrt_loaded_variable_registry",
+        ":ifrt_restore_tensor_registry",
+        ":ifrt_serving_core_selector",
+        ":ifrt_serving_executable",
+        ":tf_host_callback",
+        "//tensorflow/compiler/mlir/tensorflow",
+        "//tensorflow/compiler/tf2xla:xla_helpers",
+        "//tensorflow/core:core_cpu_base",
+        "//tensorflow/core/framework:tensor",
+        "//tensorflow/core/framework:types_proto_cc",
+        "//tensorflow/core/platform:resource_loader",
+        "@com_google_absl//absl/status:statusor",
+        "@com_google_absl//absl/strings",
+        "@com_google_absl//absl/strings:string_view",
+        "@llvm-project//mlir:AllPassesAndDialects",
+        "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:Parser",
+        "@local_tsl//tsl/framework/test_util:mock_serving_device_selector",
+        "@local_tsl//tsl/platform:env",
+        "@local_tsl//tsl/platform:status",
+        "@local_xla//xla/python/ifrt",
+        "@local_xla//xla/python/ifrt:test_util",
+        "@tf_runtime//:hostcontext",
+    ],
+)
+
 tf_cc_test(
     name = "ifrt_restore_tensor_registry_test",
     srcs = ["ifrt_restore_tensor_registry_test.cc"],
@@ -418,13 +453,10 @@ tf_cc_test(
     ],
     tags = ["no_oss"],
     deps = [
-        ":ifrt_loaded_variable_registry",
         ":ifrt_restore_tensor_registry",
         ":ifrt_serving_core_selector",
         ":ifrt_serving_executable",
-        ":sharding_utils",
-        ":tf_host_callback",
-        "//tensorflow/compiler/mlir/tensorflow",
+        ":ifrt_serving_executable_test_util",
         "//tensorflow/compiler/tf2xla:xla_helpers",
         "//tensorflow/core:core_cpu_base",
         "//tensorflow/core:framework",
@@ -436,29 +468,22 @@ tf_cc_test(
         "//tensorflow/core/framework:tensor_matcher",
         "//tensorflow/core/framework:tensor_testutil",
         "//tensorflow/core/framework:types_proto_cc",
-        "//tensorflow/core/platform:resource_loader",
         "@com_google_absl//absl/log:check",
         "@com_google_absl//absl/status",
         "@com_google_absl//absl/status:statusor",
         "@com_google_absl//absl/strings",
         "@com_google_absl//absl/types:span",
         "@com_google_googletest//:gtest_main",
-        "@llvm-project//mlir:AllPassesAndDialects",
-        "@llvm-project//mlir:IR",
-        "@llvm-project//mlir:Parser",
         "@local_tsl//tsl/framework:serving_device_selector",
         "@local_tsl//tsl/framework/test_util:mock_serving_device_selector",
         "@local_tsl//tsl/platform:env",
         "@local_tsl//tsl/platform:statusor",
         "@local_tsl//tsl/platform:tstring",
-        "@local_xla//xla/hlo/ir:hlo",
         "@local_xla//xla/python/ifrt",
         "@local_xla//xla/python/ifrt:test_util",
         "@local_xla//xla/python/pjrt_ifrt:tfrt_cpu_client_test_lib",
-        "@local_xla//xla/tsl/concurrency:ref_count",
         "@tf_runtime//:basic_kernels_alwayslink",
         "@tf_runtime//:core_runtime_alwayslink",
-        "@tf_runtime//:hostcontext",
         "@tf_runtime//:test_kernels_alwayslink",
         "@tf_runtime//backends/cpu:core_runtime_alwayslink",
         "@tf_runtime//backends/cpu:tf_ops_alwayslink",
diff --git a/tensorflow/core/tfrt/ifrt/ifrt_serving_executable_test.cc b/tensorflow/core/tfrt/ifrt/ifrt_serving_executable_test.cc
index 5de01ff67139ae..aca38747cff03d 100644
--- a/tensorflow/core/tfrt/ifrt/ifrt_serving_executable_test.cc
+++ b/tensorflow/core/tfrt/ifrt/ifrt_serving_executable_test.cc
@@ -29,44 +29,36 @@ limitations under the License.
 #include "absl/strings/str_cat.h"
 #include "absl/strings/string_view.h"
 #include "absl/types/span.h"
-#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
-#include "mlir/IR/DialectRegistry.h"  // from @llvm-project
-#include "mlir/IR/MLIRContext.h"  // from @llvm-project
-#include "mlir/IR/OwningOpRef.h"  // from @llvm-project
-#include "mlir/InitAllDialects.h"  // from @llvm-project
-#include "mlir/Parser/Parser.h"  // from @llvm-project
-#include "tensorflow/compiler/mlir/tensorflow/dialect_registration.h"
-#include "tensorflow/compiler/tf2xla/xla_helpers.h"
-#include "xla/hlo/ir/hlo_sharding.h"
 #include "xla/python/ifrt/array.h"
 #include "xla/python/ifrt/client.h"
 #include "xla/python/ifrt/future.h"
 #include "xla/python/ifrt/test_util.h"
-#include "xla/tsl/concurrency/ref_count.h"
-#include "tensorflow/core/common_runtime/device_mgr.h"
 #include "tensorflow/core/framework/tensor.h"
 #include "tensorflow/core/framework/tensor_matcher.h"
 #include "tensorflow/core/framework/tensor_shape.h"
 #include "tensorflow/core/framework/tensor_testutil.h"
 #include "tensorflow/core/framework/types.pb.h"
-#include "tensorflow/core/platform/resource_loader.h"
 #include "tensorflow/core/platform/test.h"
-#include "tensorflow/core/tfrt/ifrt/ifrt_loaded_variable_registry.h"
 #include "tensorflow/core/tfrt/ifrt/ifrt_restore_tensor_registry.h"
 #include "tensorflow/core/tfrt/ifrt/ifrt_serving_core_selector.h"
-#include "tensorflow/core/tfrt/ifrt/sharding_utils.h"
-#include "tensorflow/core/tfrt/ifrt/tf_host_callback.h"
+#include "tensorflow/core/tfrt/ifrt/ifrt_serving_executable_test_util.h"
 #include "tsl/framework/serving_device_selector.h"
 #include "tsl/framework/test_util/mock_serving_device_selector.h"
 #include "tsl/platform/env.h"
 #include "tsl/platform/statusor.h"
 #include "tsl/platform/threadpool.h"
 #include "tsl/platform/tstring.h"
-#include "tfrt/host_context/concurrent_work_queue.h"  // from @tf_runtime
 
 namespace tensorflow {
 namespace ifrt_serving {
 namespace {
+using tensorflow::ifrt_serving::test_utils::GetMlirModulePath;
+using ::tensorflow::test::AsTensor;
+using ::tensorflow::test::TensorEq;
+using ::testing::ElementsAre;
+using ::testing::Return;
+using ::tsl::testing::StatusIs;
+
 struct VariableInputTestParam {
   std::vector<tensorflow::Tensor> in_tensors;
   std::vector<bool>
@@ -76,12 +68,6 @@ struct VariableInputTestParam {
 };
 using VariableInputTest = ::testing::TestWithParam<VariableInputTestParam>;
 
-using ::tensorflow::test::AsTensor;
-using ::tensorflow::test::TensorEq;
-using ::testing::ElementsAre;
-using ::testing::Return;
-using ::tsl::testing::StatusIs;
-
 const tsl::thread::ThreadPool& GetThreadPool() {
   constexpr int kMaxParallelism = 16;
   static auto* const thread_pool =
@@ -89,7 +75,6 @@ const tsl::thread::ThreadPool& GetThreadPool() {
                                   "IfrtSharding", kMaxParallelism);
   return *thread_pool;
 }
-
 class IfrtServingExecutableTest : public ::testing::Test {
  protected:
   explicit IfrtServingExecutableTest() {
@@ -106,48 +91,13 @@ class IfrtServingExecutableTest : public ::testing::Test {
 };
 
 TEST_F(IfrtServingExecutableTest, Basic) {
-  // Create test input module
-  constexpr absl::string_view kDataDirectory =
-      "tensorflow/core/tfrt/ifrt/testdata";
-  std::string mlir_module_path = tensorflow::GetDataDependencyFilepath(
-      absl::StrCat(kDataDirectory, "/executable.mlir"));
-
-  mlir::DialectRegistry registry;
-  mlir::registerAllDialects(registry);
-  mlir::RegisterAllTensorFlowDialects(registry);
-
-  mlir::MLIRContext context(registry);
-
-  mlir::OwningOpRef<mlir::ModuleOp> mlir_module =
-      mlir::parseSourceFile<mlir::ModuleOp>(mlir_module_path, &context);
-
-  ASSERT_TRUE(mlir_module);
-
+  test_utils::IfrtServingExecutableTestHelper helper(&selector_);
   int64_t program_id = 123456;
   EXPECT_CALL(selector_, ReserveDevice(absl::StrCat(program_id)))
       .Times(1)
       .WillOnce(Return(tsl::DeviceReservation(0, /*selector=*/nullptr)));
-
-  // Create contexts required for the compiler execution.
-  TF_ASSERT_OK_AND_ASSIGN(std::shared_ptr<xla::ifrt::Client> client,
-                          xla::ifrt::test_util::GetClient());
-
-  IfrtLoadedVariableRegistry ifrt_loaded_variable_registry;
-  IfrtRestoreTensorRegistry ifrt_restore_tensor_registry;
-  std::unique_ptr<tfrt::ConcurrentWorkQueue> work_queue =
-      tfrt::CreateMultiThreadedWorkQueue(
-          /*num_threads=*/4, /*num_blocking_threads=*/4);
-  TF_ASSERT_OK_AND_ASSIGN(
-      std::unique_ptr<tensorflow::StaticDeviceMgr> device_mgr,
-      CreateTfStaticDeviceMgr());
-
-  TF_ASSERT_OK_AND_ASSIGN(
-      auto executable,
-      IfrtServingExecutable ::Create(
-          program_id, "test", "main", std::move(mlir_module), client,
-          &GetThreadPool(), &ifrt_loaded_variable_registry,
-          &ifrt_restore_tensor_registry, work_queue.get(), device_mgr.get(),
-          tensorflow::IdentityShapeRepresentationFn(), core_selector_.get()));
+  auto executable =
+      helper.MakeExecutable(program_id, GetMlirModulePath("executable.mlir"));
 
   auto x = AsTensor<int32_t>({1, 2, 3}, tensorflow::TensorShape({1, 3}));
   auto y = AsTensor<int32_t>({1, 2, 3}, tensorflow::TensorShape({3, 1}));
@@ -163,49 +113,14 @@ TEST_F(IfrtServingExecutableTest, Basic) {
 }
 
 TEST_F(IfrtServingExecutableTest, MultipleShapes) {
-  // Create test input module
-  constexpr absl::string_view kDataDirectory =
-      "tensorflow/core/tfrt/ifrt/testdata";
-  std::string mlir_module_path = tensorflow::GetDataDependencyFilepath(
-      absl::StrCat(kDataDirectory, "/executable.mlir"));
-
-  mlir::DialectRegistry registry;
-  mlir::registerAllDialects(registry);
-  mlir::RegisterAllTensorFlowDialects(registry);
-
-  mlir::MLIRContext context(registry);
-
-  mlir::OwningOpRef<mlir::ModuleOp> mlir_module =
-      mlir::parseSourceFile<mlir::ModuleOp>(mlir_module_path, &context);
-
-  ASSERT_TRUE(mlir_module);
-
+  test_utils::IfrtServingExecutableTestHelper helper(&selector_);
   int64_t program_id = 123456;
   EXPECT_CALL(selector_, ReserveDevice(absl::StrCat(program_id)))
       .Times(6)
       .WillRepeatedly(
           [](::testing::Unused) { return tsl::DeviceReservation(0, nullptr); });
-
-  // Create contexts required for the compiler execution.
-  TF_ASSERT_OK_AND_ASSIGN(std::shared_ptr<xla::ifrt::Client> client,
-                          xla::ifrt::test_util::GetClient());
-
-  IfrtLoadedVariableRegistry ifrt_loaded_variable_registry;
-  IfrtRestoreTensorRegistry ifrt_restore_tensor_registry;
-  std::unique_ptr<tfrt::ConcurrentWorkQueue> work_queue =
-      tfrt::CreateMultiThreadedWorkQueue(
-          /*num_threads=*/4, /*num_blocking_threads=*/4);
-  TF_ASSERT_OK_AND_ASSIGN(
-      std::unique_ptr<tensorflow::StaticDeviceMgr> device_mgr,
-      CreateTfStaticDeviceMgr());
-
-  TF_ASSERT_OK_AND_ASSIGN(
-      auto executable,
-      IfrtServingExecutable ::Create(
-          program_id, "test", "main", std::move(mlir_module), client,
-          &GetThreadPool(), &ifrt_loaded_variable_registry,
-          &ifrt_restore_tensor_registry, work_queue.get(), device_mgr.get(),
-          tensorflow::IdentityShapeRepresentationFn(), core_selector_.get()));
+  auto executable =
+      helper.MakeExecutable(program_id, GetMlirModulePath("executable.mlir"));
 
   auto x1 = AsTensor<int32_t>({1, 2, 3}, tensorflow::TensorShape({1, 3}));
   auto y1 = AsTensor<int32_t>({1, 2, 3}, tensorflow::TensorShape({3, 1}));
@@ -236,49 +151,14 @@ TEST_F(IfrtServingExecutableTest, MultipleShapes) {
 }
 
 TEST_F(IfrtServingExecutableTest, ReturnFailOnUncompiledShapeAfterFrozen) {
-  // Create test input module
-  constexpr absl::string_view kDataDirectory =
-      "tensorflow/core/tfrt/ifrt/testdata";
-  std::string mlir_module_path = tensorflow::GetDataDependencyFilepath(
-      absl::StrCat(kDataDirectory, "/executable.mlir"));
-
-  mlir::DialectRegistry registry;
-  mlir::registerAllDialects(registry);
-  mlir::RegisterAllTensorFlowDialects(registry);
-
-  mlir::MLIRContext context(registry);
-
-  mlir::OwningOpRef<mlir::ModuleOp> mlir_module =
-      mlir::parseSourceFile<mlir::ModuleOp>(mlir_module_path, &context);
-
-  ASSERT_TRUE(mlir_module);
-
+  test_utils::IfrtServingExecutableTestHelper helper(&selector_);
   int64_t program_id = 123456;
   EXPECT_CALL(selector_, ReserveDevice(absl::StrCat(program_id)))
       .Times(3)
       .WillRepeatedly(
           [](::testing::Unused) { return tsl::DeviceReservation(0, nullptr); });
-
-  // Create contexts required for the compiler execution.
-  TF_ASSERT_OK_AND_ASSIGN(std::shared_ptr<xla::ifrt::Client> client,
-                          xla::ifrt::test_util::GetClient());
-
-  IfrtLoadedVariableRegistry ifrt_loaded_variable_registry;
-  IfrtRestoreTensorRegistry ifrt_restore_tensor_registry;
-  std::unique_ptr<tfrt::ConcurrentWorkQueue> work_queue =
-      tfrt::CreateMultiThreadedWorkQueue(
-          /*num_threads=*/4, /*num_blocking_threads=*/4);
-  TF_ASSERT_OK_AND_ASSIGN(
-      std::unique_ptr<tensorflow::StaticDeviceMgr> device_mgr,
-      CreateTfStaticDeviceMgr());
-
-  TF_ASSERT_OK_AND_ASSIGN(
-      auto executable,
-      IfrtServingExecutable ::Create(
-          program_id, "test", "main", std::move(mlir_module), client,
-          &GetThreadPool(), &ifrt_loaded_variable_registry,
-          &ifrt_restore_tensor_registry, work_queue.get(), device_mgr.get(),
-          tensorflow::IdentityShapeRepresentationFn(), core_selector_.get()));
+  auto executable =
+      helper.MakeExecutable(program_id, GetMlirModulePath("executable.mlir"));
 
   auto x1 = AsTensor<int32_t>({1, 2, 3}, tensorflow::TensorShape({1, 3}));
   auto y1 = AsTensor<int32_t>({1, 2, 3}, tensorflow::TensorShape({3, 1}));
@@ -310,47 +190,11 @@ TEST_F(IfrtServingExecutableTest, ReturnFailOnUncompiledShapeAfterFrozen) {
 }
 
 TEST_F(IfrtServingExecutableTest, Spmd) {
-  // Create test input module
-  constexpr absl::string_view kDataDirectory =
-      "tensorflow/core/tfrt/ifrt/testdata";
-  std::string mlir_module_path = tensorflow::GetDataDependencyFilepath(
-      absl::StrCat(kDataDirectory, "/spmd_executable.mlir"));
-
-  mlir::DialectRegistry registry;
-  mlir::registerAllDialects(registry);
-  mlir::RegisterAllTensorFlowDialects(registry);
-
-  mlir::MLIRContext context(registry);
-
-  mlir::OwningOpRef<mlir::ModuleOp> mlir_module =
-      mlir::parseSourceFile<mlir::ModuleOp>(mlir_module_path, &context);
-
-  ASSERT_TRUE(mlir_module);
-
+  test_utils::IfrtServingExecutableTestHelper helper(&selector_);
   int64_t program_id = 111111;
-
   EXPECT_CALL(selector_, ReserveDevice(absl::StrCat(program_id))).Times(0);
-
-  // Create contexts required for the compiler execution.
-  TF_ASSERT_OK_AND_ASSIGN(std::shared_ptr<xla::ifrt::Client> client,
-                          xla::ifrt::test_util::GetClient());
-
-  IfrtLoadedVariableRegistry ifrt_loaded_variable_registry;
-  IfrtRestoreTensorRegistry ifrt_restore_tensor_registry;
-  std::unique_ptr<tfrt::ConcurrentWorkQueue> work_queue =
-      tfrt::CreateMultiThreadedWorkQueue(
-          /*num_threads=*/4, /*num_blocking_threads=*/4);
-  TF_ASSERT_OK_AND_ASSIGN(
-      std::unique_ptr<tensorflow::StaticDeviceMgr> device_mgr,
-      CreateTfStaticDeviceMgr());
-
-  TF_ASSERT_OK_AND_ASSIGN(
-      auto executable,
-      IfrtServingExecutable ::Create(
-          program_id, "test", "main", std::move(mlir_module), client,
-          &GetThreadPool(), &ifrt_loaded_variable_registry,
-          &ifrt_restore_tensor_registry, work_queue.get(), device_mgr.get(),
-          tensorflow::IdentityShapeRepresentationFn(), core_selector_.get()));
+  auto executable = helper.MakeExecutable(
+      program_id, GetMlirModulePath("spmd_executable.mlir"));
 
   auto x = AsTensor<int32_t>({1, 2, 3, 4, 5, 6, 7, 8},
                              tensorflow::TensorShape({4, 2}));
@@ -371,47 +215,11 @@ TEST_F(IfrtServingExecutableTest, Spmd) {
 }
 
 TEST_F(IfrtServingExecutableTest, SpmdTwoReturns) {
-  // Create test input module
-  constexpr absl::string_view kDataDirectory =
-      "tensorflow/core/tfrt/ifrt/testdata";
-  std::string mlir_module_path = tensorflow::GetDataDependencyFilepath(
-      absl::StrCat(kDataDirectory, "/spmd_executable_two_returns.mlir"));
-
-  mlir::DialectRegistry registry;
-  mlir::registerAllDialects(registry);
-  mlir::RegisterAllTensorFlowDialects(registry);
-
-  mlir::MLIRContext context(registry);
-
-  mlir::OwningOpRef<mlir::ModuleOp> mlir_module =
-      mlir::parseSourceFile<mlir::ModuleOp>(mlir_module_path, &context);
-
-  ASSERT_TRUE(mlir_module);
-
+  test_utils::IfrtServingExecutableTestHelper helper(&selector_);
   int64_t program_id = 111111;
-
   EXPECT_CALL(selector_, ReserveDevice(absl::StrCat(program_id))).Times(0);
-
-  // Create contexts required for the compiler execution.
-  TF_ASSERT_OK_AND_ASSIGN(std::shared_ptr<xla::ifrt::Client> client,
-                          xla::ifrt::test_util::GetClient());
-
-  IfrtLoadedVariableRegistry ifrt_loaded_variable_registry;
-  IfrtRestoreTensorRegistry ifrt_restore_tensor_registry;
-  std::unique_ptr<tfrt::ConcurrentWorkQueue> work_queue =
-      tfrt::CreateMultiThreadedWorkQueue(
-          /*num_threads=*/4, /*num_blocking_threads=*/4);
-  TF_ASSERT_OK_AND_ASSIGN(
-      std::unique_ptr<tensorflow::StaticDeviceMgr> device_mgr,
-      CreateTfStaticDeviceMgr());
-
-  TF_ASSERT_OK_AND_ASSIGN(
-      auto executable,
-      IfrtServingExecutable ::Create(
-          program_id, "test", "main", std::move(mlir_module), client,
-          &GetThreadPool(), &ifrt_loaded_variable_registry,
-          &ifrt_restore_tensor_registry, work_queue.get(), device_mgr.get(),
-          tensorflow::IdentityShapeRepresentationFn(), core_selector_.get()));
+  auto executable = helper.MakeExecutable(
+      program_id, GetMlirModulePath("spmd_executable_two_returns.mlir"));
 
   auto x = AsTensor<int32_t>({1, 2, 3, 4, 5, 6, 7, 8},
                              tensorflow::TensorShape({4, 2}));
@@ -436,50 +244,14 @@ TEST_F(IfrtServingExecutableTest, SpmdTwoReturns) {
 }
 
 TEST_F(IfrtServingExecutableTest, NoReturn) {
-  // Create test input module
-  constexpr absl::string_view kDataDirectory =
-      "tensorflow/core/tfrt/ifrt/testdata";
-  std::string mlir_module_path = tensorflow::GetDataDependencyFilepath(
-      absl::StrCat(kDataDirectory, "/executable_no_return.mlir"));
-
-  mlir::DialectRegistry registry;
-  mlir::registerAllDialects(registry);
-  mlir::RegisterAllTensorFlowDialects(registry);
-
-  mlir::MLIRContext context(registry);
-
-  mlir::OwningOpRef<mlir::ModuleOp> mlir_module =
-      mlir::parseSourceFile<mlir::ModuleOp>(mlir_module_path, &context);
-
-  ASSERT_TRUE(mlir_module);
-
+  test_utils::IfrtServingExecutableTestHelper helper(&selector_);
   int64_t program_id = 111111;
-
   EXPECT_CALL(selector_, ReserveDevice(absl::StrCat(program_id)))
       .Times(1)
       .WillRepeatedly(
           [](::testing::Unused) { return tsl::DeviceReservation(0, nullptr); });
-
-  // Create contexts required for the compiler execution.
-  TF_ASSERT_OK_AND_ASSIGN(std::shared_ptr<xla::ifrt::Client> client,
-                          xla::ifrt::test_util::GetClient());
-
-  IfrtLoadedVariableRegistry ifrt_loaded_variable_registry;
-  IfrtRestoreTensorRegistry ifrt_restore_tensor_registry;
-  std::unique_ptr<tfrt::ConcurrentWorkQueue> work_queue =
-      tfrt::CreateMultiThreadedWorkQueue(
-          /*num_threads=*/4, /*num_blocking_threads=*/4);
-  TF_ASSERT_OK_AND_ASSIGN(
-      std::unique_ptr<tensorflow::StaticDeviceMgr> device_mgr,
-      CreateTfStaticDeviceMgr());
-
-  TF_ASSERT_OK_AND_ASSIGN(
-      auto executable,
-      IfrtServingExecutable ::Create(
-          program_id, "test", "main", std::move(mlir_module), client,
-          &GetThreadPool(), &ifrt_loaded_variable_registry,
-          &ifrt_restore_tensor_registry, work_queue.get(), device_mgr.get(),
-          tensorflow::IdentityShapeRepresentationFn(), core_selector_.get()));
+  auto executable = helper.MakeExecutable(
+      program_id, GetMlirModulePath("executable_no_return.mlir"));
 
   auto x = AsTensor<int32_t>({1, 2, 3}, tensorflow::TensorShape({1, 3}));
   auto y = AsTensor<int32_t>({1, 2, 3}, tensorflow::TensorShape({3, 1}));
@@ -492,52 +264,17 @@ TEST_F(IfrtServingExecutableTest, NoReturn) {
 }
 
 TEST_P(VariableInputTest, InterleaveVariable) {
-  // Create test input module
-  constexpr absl::string_view kDataDirectory =
-      "tensorflow/core/tfrt/ifrt/testdata";
-  std::string mlir_module_path = tensorflow::GetDataDependencyFilepath(
-      absl::StrCat(kDataDirectory, "/executable_long_inputs.mlir"));
-
-  mlir::DialectRegistry registry;
-  mlir::registerAllDialects(registry);
-  mlir::RegisterAllTensorFlowDialects(registry);
-
-  mlir::MLIRContext context(registry);
-
-  mlir::OwningOpRef<mlir::ModuleOp> mlir_module =
-      mlir::parseSourceFile<mlir::ModuleOp>(mlir_module_path, &context);
-
-  ASSERT_TRUE(mlir_module);
-
   tsl::test_util::MockServingDeviceSelector device_selector;
-  IfrtServingCoreSelector core_selector(&device_selector);
+  test_utils::IfrtServingExecutableTestHelper helper(&device_selector);
   int64_t program_id = 111111;
-
   EXPECT_CALL(device_selector, ReserveDevice(absl::StrCat(program_id)))
       .Times(1)
       .WillRepeatedly(
           [](::testing::Unused) { return tsl::DeviceReservation(0, nullptr); });
-
-  // Create contexts required for the compiler execution.
-  TF_ASSERT_OK_AND_ASSIGN(std::shared_ptr<xla::ifrt::Client> client,
-                          xla::ifrt::test_util::GetClient());
-
-  IfrtLoadedVariableRegistry ifrt_loaded_variable_registry;
-  IfrtRestoreTensorRegistry ifrt_restore_tensor_registry;
-  std::unique_ptr<tfrt::ConcurrentWorkQueue> work_queue =
-      tfrt::CreateMultiThreadedWorkQueue(
-          /*num_threads=*/4, /*num_blocking_threads=*/4);
-  TF_ASSERT_OK_AND_ASSIGN(
-      std::unique_ptr<tensorflow::StaticDeviceMgr> device_mgr,
-      CreateTfStaticDeviceMgr());
-
-  TF_ASSERT_OK_AND_ASSIGN(
-      auto executable,
-      IfrtServingExecutable ::Create(
-          program_id, "test", "main", std::move(mlir_module), client,
-          &GetThreadPool(), &ifrt_loaded_variable_registry,
-          &ifrt_restore_tensor_registry, work_queue.get(), device_mgr.get(),
-          tensorflow::IdentityShapeRepresentationFn(), &core_selector));
+  auto executable = helper.MakeExecutable(
+      program_id, GetMlirModulePath("executable_long_inputs.mlir"));
+  IfrtRestoreTensorRegistry* ifrt_restore_tensor_registry =
+      helper.ifrt_restore_tensor_registry();
 
   std::vector<tensorflow::Tensor> inputs;
   std::vector<int> loaded_variable_indices;
@@ -552,8 +289,8 @@ TEST_P(VariableInputTest, InterleaveVariable) {
                            .shape = GetParam().in_tensors[i].shape()},
           .tensor_future = input_tensor_future};
       std::string variable_name = absl::StrCat("variable_", i);
-      ASSERT_OK(ifrt_restore_tensor_registry.TryRegister(variable_name,
-                                                         restore_tensor_info));
+      ASSERT_OK(ifrt_restore_tensor_registry->TryRegister(variable_name,
+                                                          restore_tensor_info));
       loaded_variable_indices.push_back(i);
       input_tensor_promise.Set(GetParam().in_tensors[i]);
       // Use string tensor containing the key (name) in place of variable
diff --git a/tensorflow/core/tfrt/ifrt/ifrt_serving_executable_test_util.cc b/tensorflow/core/tfrt/ifrt/ifrt_serving_executable_test_util.cc
new file mode 100644
index 00000000000000..49ca7b3a9bbd98
--- /dev/null
+++ b/tensorflow/core/tfrt/ifrt/ifrt_serving_executable_test_util.cc
@@ -0,0 +1,95 @@
+/* Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+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.
+==============================================================================*/
+
+#include "tensorflow/core/tfrt/ifrt/ifrt_serving_executable_test_util.h"
+
+#include <cstdint>
+#include <memory>
+#include <string>
+#include <utility>
+
+#include "absl/status/statusor.h"
+#include "absl/strings/str_cat.h"
+#include "absl/strings/string_view.h"
+#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
+#include "mlir/IR/MLIRContext.h"  // from @llvm-project
+#include "mlir/IR/OwningOpRef.h"  // from @llvm-project
+#include "mlir/InitAllDialects.h"  // from @llvm-project
+#include "mlir/Parser/Parser.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/tensorflow/dialect_registration.h"
+#include "tensorflow/compiler/tf2xla/xla_helpers.h"
+#include "xla/python/ifrt/test_util.h"
+#include "tensorflow/core/framework/types.pb.h"
+#include "tensorflow/core/platform/resource_loader.h"
+#include "tensorflow/core/tfrt/ifrt/ifrt_serving_core_selector.h"
+#include "tensorflow/core/tfrt/ifrt/ifrt_serving_executable.h"
+#include "tensorflow/core/tfrt/ifrt/tf_host_callback.h"
+#include "tsl/framework/test_util/mock_serving_device_selector.h"
+#include "tsl/platform/env.h"
+#include "tsl/platform/status.h"
+#include "tsl/platform/threadpool.h"
+#include "tfrt/host_context/concurrent_work_queue.h"  // from @tf_runtime
+namespace tensorflow {
+namespace ifrt_serving {
+namespace test_utils {
+
+inline constexpr absl::string_view kMlirModulePath =
+    "tensorflow/core/tfrt/ifrt/testdata/";
+
+std::string GetMlirModulePath(absl::string_view module_name) {
+  return tensorflow::GetDataDependencyFilepath(
+      absl::StrCat(kMlirModulePath, module_name));
+}
+
+IfrtServingExecutableTestHelper::IfrtServingExecutableTestHelper(
+    tsl::test_util::MockServingDeviceSelector* device_selector)
+    : device_selector_(device_selector) {
+  core_selector_ = std::make_unique<IfrtServingCoreSelector>(device_selector_);
+  auto client_or = xla::ifrt::test_util::GetClient();
+  TF_CHECK_OK(client_or.status());
+  client_ = std::move(client_or.value());
+
+  thread_pool_ = std::make_unique<tsl::thread::ThreadPool>(
+      tsl::Env::Default(), tsl::ThreadOptions(), "IfrtSharding",
+      kThreadPoolNumThreads);
+  work_queue_ = tfrt::CreateMultiThreadedWorkQueue(
+      /*num_threads=*/4, /*num_blocking_threads=*/4);
+
+  auto device_mgr_or = ifrt_serving::CreateTfStaticDeviceMgr();
+  TF_CHECK_OK(device_mgr_or.status());
+  device_mgr_ = std::move(device_mgr_or.value());
+
+  mlir::registerAllDialects(registry_);
+  mlir::RegisterAllTensorFlowDialects(registry_);
+  context_ = std::make_unique<mlir::MLIRContext>(registry_);
+}
+
+std::unique_ptr<IfrtServingExecutable>
+IfrtServingExecutableTestHelper::MakeExecutable(int64_t program_id,
+                                                std::string mlir_module_path) {
+  auto mlir_module =
+      mlir::parseSourceFile<mlir::ModuleOp>(mlir_module_path, context_.get());
+  auto executable_or = IfrtServingExecutable::Create(
+      program_id, "test", "main", std::move(mlir_module), client_,
+      thread_pool_.get(), &ifrt_loaded_variable_registry_,
+      &ifrt_restore_tensor_registry_, work_queue_.get(), device_mgr_.get(),
+      tensorflow::IdentityShapeRepresentationFn(), core_selector_.get());
+  TF_CHECK_OK(executable_or.status());
+  return std::move(executable_or.value());
+}
+
+}  // namespace test_utils
+}  // namespace ifrt_serving
+}  // namespace tensorflow
diff --git a/tensorflow/core/tfrt/ifrt/ifrt_serving_executable_test_util.h b/tensorflow/core/tfrt/ifrt/ifrt_serving_executable_test_util.h
new file mode 100644
index 00000000000000..2fc669c08a5062
--- /dev/null
+++ b/tensorflow/core/tfrt/ifrt/ifrt_serving_executable_test_util.h
@@ -0,0 +1,82 @@
+/* Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+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.
+==============================================================================*/
+
+#ifndef TENSORFLOW_CORE_TFRT_IFRT_IFRT_SERVING_EXECUTABLE_TEST_UTIL_H_
+#define TENSORFLOW_CORE_TFRT_IFRT_IFRT_SERVING_EXECUTABLE_TEST_UTIL_H_
+
+#include <cstdint>
+#include <memory>
+#include <string>
+
+#include "absl/strings/string_view.h"
+#include "mlir/IR/DialectRegistry.h"  // from @llvm-project
+#include "mlir/IR/MLIRContext.h"  // from @llvm-project
+#include "xla/python/ifrt/array.h"
+#include "xla/python/ifrt/client.h"
+#include "xla/python/ifrt/test_util.h"
+#include "tensorflow/core/common_runtime/device_mgr.h"
+#include "tensorflow/core/framework/types.pb.h"
+#include "tensorflow/core/tfrt/ifrt/ifrt_loaded_variable_registry.h"
+#include "tensorflow/core/tfrt/ifrt/ifrt_restore_tensor_registry.h"
+#include "tensorflow/core/tfrt/ifrt/ifrt_serving_core_selector.h"
+#include "tensorflow/core/tfrt/ifrt/ifrt_serving_executable.h"
+#include "tsl/framework/test_util/mock_serving_device_selector.h"
+#include "tsl/platform/threadpool.h"
+#include "tfrt/host_context/concurrent_work_queue.h"  // from @tf_runtime
+
+namespace tensorflow {
+namespace ifrt_serving {
+namespace test_utils {
+
+// A test helper class to create and IfrtServingExecutable.
+class IfrtServingExecutableTestHelper {
+ public:
+  explicit IfrtServingExecutableTestHelper(
+      tsl::test_util::MockServingDeviceSelector* device_selector);
+
+  // Creates an IfrtServingExecutable with the given program id.
+  // Note the instance of this class must outlive the returned
+  // IfrtServingExecutable.
+  std::unique_ptr<IfrtServingExecutable> MakeExecutable(
+      int64_t program_id, std::string mlir_module_path);
+
+  IfrtRestoreTensorRegistry* ifrt_restore_tensor_registry() {
+    return &ifrt_restore_tensor_registry_;
+  }
+
+ private:
+  static constexpr int kThreadPoolNumThreads = 16;
+
+  tsl::test_util::MockServingDeviceSelector* device_selector_;  // Not owned.
+  std::unique_ptr<IfrtServingCoreSelector> core_selector_;
+  std::shared_ptr<xla::ifrt::Client> client_;
+  std::unique_ptr<tsl::thread::ThreadPool> thread_pool_;
+  IfrtLoadedVariableRegistry ifrt_loaded_variable_registry_;
+  IfrtRestoreTensorRegistry ifrt_restore_tensor_registry_;
+  std::unique_ptr<tfrt::ConcurrentWorkQueue> work_queue_;
+  std::unique_ptr<tensorflow::StaticDeviceMgr> device_mgr_;
+
+  mlir::DialectRegistry registry_;
+  std::unique_ptr<mlir::MLIRContext> context_;
+};
+
+// Returns the path to the MLIR module for the given module name.
+std::string GetMlirModulePath(absl::string_view module_name);
+
+}  // namespace test_utils
+}  // namespace ifrt_serving
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_CORE_TFRT_IFRT_IFRT_SERVING_EXECUTABLE_TEST_UTIL_H_
diff --git a/tensorflow/core/tfrt/kernels/BUILD b/tensorflow/core/tfrt/kernels/BUILD
index 39cfca154b506a..a6c6ba58c6659a 100644
--- a/tensorflow/core/tfrt/kernels/BUILD
+++ b/tensorflow/core/tfrt/kernels/BUILD
@@ -43,7 +43,6 @@ tf_cc_test(
     tags = ["no_oss"],
     deps = [
         ":ifrt_program_ops",
-        "//tensorflow/compiler/mlir/tensorflow",
         "//tensorflow/compiler/tf2xla:xla_helpers",
         "//tensorflow/core:core_cpu_base",
         "//tensorflow/core:framework",
@@ -53,14 +52,9 @@ tf_cc_test(
         "//tensorflow/core/framework:tensor_matcher",
         "//tensorflow/core/framework:types_proto_cc",
         "//tensorflow/core/lib/gtl:cleanup",
-        "//tensorflow/core/platform:resource_loader",
         "//tensorflow/core/platform:status",
         "//tensorflow/core/tfrt/ifrt:ifrt_executable_registry",
-        "//tensorflow/core/tfrt/ifrt:ifrt_loaded_variable_registry",
-        "//tensorflow/core/tfrt/ifrt:ifrt_restore_tensor_registry",
-        "//tensorflow/core/tfrt/ifrt:ifrt_serving_core_selector",
-        "//tensorflow/core/tfrt/ifrt:ifrt_serving_executable",
-        "//tensorflow/core/tfrt/ifrt:tf_host_callback",
+        "//tensorflow/core/tfrt/ifrt:ifrt_serving_executable_test_util",
         "//tensorflow/core/tfrt/ops:ifrt_program_ops_op_lib",
         "@com_google_absl//absl/log:check",
         "@com_google_absl//absl/status:statusor",
@@ -68,20 +62,15 @@ tf_cc_test(
         "@com_google_absl//absl/strings:string_view",
         "@com_google_absl//absl/types:span",
         "@com_google_googletest//:gtest_main",
-        "@llvm-project//mlir:AllPassesAndDialects",
-        "@llvm-project//mlir:IR",
-        "@llvm-project//mlir:Parser",
         "@local_tsl//tsl/framework:serving_device_selector",
         "@local_tsl//tsl/framework/test_util:mock_serving_device_selector",
         "@local_tsl//tsl/lib/core:status_test_util",
-        "@local_tsl//tsl/platform:env",
         "@local_tsl//tsl/platform:status",
         "@local_tsl//tsl/platform:statusor",
         "@local_xla//xla/pjrt/cpu:cpu_client",
         "@local_xla//xla/python/ifrt",
         "@local_xla//xla/python/ifrt:test_util",
         "@local_xla//xla/python/pjrt_ifrt",
-        "@tf_runtime//:hostcontext",
     ],
 )
 
diff --git a/tensorflow/core/tfrt/kernels/ifrt_program_ops_test.cc b/tensorflow/core/tfrt/kernels/ifrt_program_ops_test.cc
index c7a2783bccb7f0..292bae8da04685 100644
--- a/tensorflow/core/tfrt/kernels/ifrt_program_ops_test.cc
+++ b/tensorflow/core/tfrt/kernels/ifrt_program_ops_test.cc
@@ -15,7 +15,6 @@ limitations under the License.
 
 #include <cstdint>
 #include <memory>
-#include <string>
 #include <utility>
 #include <vector>
 
@@ -26,14 +25,6 @@ limitations under the License.
 #include "absl/strings/str_cat.h"
 #include "absl/strings/string_view.h"
 #include "absl/types/span.h"
-#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
-#include "mlir/IR/DialectRegistry.h"  // from @llvm-project
-#include "mlir/IR/MLIRContext.h"  // from @llvm-project
-#include "mlir/IR/OwningOpRef.h"  // from @llvm-project
-#include "mlir/InitAllDialects.h"  // from @llvm-project
-#include "mlir/Parser/Parser.h"  // from @llvm-project
-#include "tensorflow/compiler/mlir/tensorflow/dialect_registration.h"
-#include "tensorflow/compiler/tf2xla/xla_helpers.h"
 #include "xla/pjrt/cpu/cpu_client.h"
 #include "xla/python/ifrt/array.h"
 #include "xla/python/ifrt/client.h"
@@ -48,33 +39,23 @@ limitations under the License.
 #include "tensorflow/core/framework/types.pb.h"
 #include "tensorflow/core/kernels/ops_testutil.h"
 #include "tensorflow/core/lib/gtl/cleanup.h"
-#include "tensorflow/core/platform/resource_loader.h"
 #include "tensorflow/core/platform/status.h"
 #include "tensorflow/core/platform/test.h"
 #include "tensorflow/core/tfrt/ifrt/ifrt_executable_registry.h"
-#include "tensorflow/core/tfrt/ifrt/ifrt_loaded_variable_registry.h"
-#include "tensorflow/core/tfrt/ifrt/ifrt_restore_tensor_registry.h"
-#include "tensorflow/core/tfrt/ifrt/ifrt_serving_core_selector.h"
-#include "tensorflow/core/tfrt/ifrt/ifrt_serving_executable.h"
-#include "tensorflow/core/tfrt/ifrt/tf_host_callback.h"
+#include "tensorflow/core/tfrt/ifrt/ifrt_serving_executable_test_util.h"
 #include "tsl/framework/serving_device_selector.h"
 #include "tsl/framework/test_util/mock_serving_device_selector.h"
 #include "tsl/lib/core/status_test_util.h"
-#include "tsl/platform/env.h"
 #include "tsl/platform/status.h"
 #include "tsl/platform/statusor.h"
-#include "tsl/platform/threadpool.h"
-#include "tfrt/host_context/concurrent_work_queue.h"  // from @tf_runtime
 
 namespace tensorflow {
 namespace tfrt_stub {
 namespace {
 
-using tensorflow::ifrt_serving::IfrtLoadedVariableRegistry;
-using tensorflow::ifrt_serving::IfrtRestoreTensorRegistry;
-using tensorflow::ifrt_serving::IfrtServingCoreSelector;
-using tensorflow::ifrt_serving::IfrtServingExecutable;
 using tensorflow::ifrt_serving::ServingExecutableRegistry;
+using tensorflow::ifrt_serving::test_utils::GetMlirModulePath;
+using tensorflow::ifrt_serving::test_utils::IfrtServingExecutableTestHelper;
 using tensorflow::test::AsTensor;
 using tensorflow::test::TensorEq;
 using ::testing::Return;
@@ -115,50 +96,13 @@ TEST_F(IfrtCallOpTest, Basic) {
       /*variable_arg_indices=*/{},
       /*output_type_list=*/{DT_INT32}));
 
-  // TODO(b/324451111): Extract the executable creation to a common function
-  // to be shared by all tests (e.g.
-  // tensorflow/core/tfrt/ifrt/ifrt_serving_executable_test.cc)
   tsl::test_util::MockServingDeviceSelector selector;
+  IfrtServingExecutableTestHelper helper(&selector);
   EXPECT_CALL(selector, ReserveDevice(absl::StrCat(program_id)))
       .Times(1)
       .WillOnce(Return(tsl::DeviceReservation(0, /*selector=*/nullptr)));
-
-  auto core_selector = std::make_unique<IfrtServingCoreSelector>(&selector);
-
-  constexpr absl::string_view kDataDirectory =
-      "tensorflow/core/tfrt/ifrt/testdata";
-  std::string mlir_module_path = tensorflow::GetDataDependencyFilepath(
-      absl::StrCat(kDataDirectory, "/executable.mlir"));
-  mlir::DialectRegistry registry;
-  mlir::registerAllDialects(registry);
-  mlir::RegisterAllTensorFlowDialects(registry);
-  mlir::MLIRContext context(registry);
-  auto mlir_module =
-      mlir::parseSourceFile<mlir::ModuleOp>(mlir_module_path, &context);
-
-  TF_ASSERT_OK_AND_ASSIGN(auto client, xla::ifrt::test_util::GetClient());
-
-  constexpr int kMaxParallelism = 16;
-  auto thread_pool = std::make_unique<tsl::thread::ThreadPool>(
-      tsl::Env::Default(), tsl::ThreadOptions(), "IfrtSharding",
-      kMaxParallelism);
-
-  auto work_queue = tfrt::CreateMultiThreadedWorkQueue(
-      /*num_threads=*/4, /*num_blocking_threads=*/4);
-
-  TF_ASSERT_OK_AND_ASSIGN(auto device_mgr,
-                          ifrt_serving::CreateTfStaticDeviceMgr());
-
-  IfrtLoadedVariableRegistry ifrt_loaded_variable_registry;
-  IfrtRestoreTensorRegistry ifrt_restore_tensor_registry;
-
-  TF_ASSERT_OK_AND_ASSIGN(
-      auto executable,
-      IfrtServingExecutable::Create(
-          program_id, "test", "main", std::move(mlir_module), client,
-          thread_pool.get(), &ifrt_loaded_variable_registry,
-          &ifrt_restore_tensor_registry, work_queue.get(), device_mgr_.get(),
-          tensorflow::IdentityShapeRepresentationFn(), core_selector.get()));
+  auto executable =
+      helper.MakeExecutable(program_id, GetMlirModulePath("executable.mlir"));
 
   TF_ASSERT_OK_AND_ASSIGN(
       ServingExecutableRegistry::Handle handle,

From ff7a81e4775c6876ac3e384f6c69d181c9a93f69 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Sun, 19 May 2024 19:44:28 -0700
Subject: [PATCH 404/478] Automated Code Change

PiperOrigin-RevId: 635317185
---
 .../quantization/stablehlo/python/pywrap_quantization_lib.cc     | 1 +
 1 file changed, 1 insertion(+)

diff --git a/tensorflow/compiler/mlir/quantization/stablehlo/python/pywrap_quantization_lib.cc b/tensorflow/compiler/mlir/quantization/stablehlo/python/pywrap_quantization_lib.cc
index 3b5ece120bdeb0..517bd117348072 100644
--- a/tensorflow/compiler/mlir/quantization/stablehlo/python/pywrap_quantization_lib.cc
+++ b/tensorflow/compiler/mlir/quantization/stablehlo/python/pywrap_quantization_lib.cc
@@ -25,6 +25,7 @@ limitations under the License.
 #include "tensorflow/compiler/mlir/quantization/stablehlo/cc/weight_only_ptq.h"
 #include "tensorflow/compiler/mlir/quantization/stablehlo/quantization_config.pb.h"
 #include "tensorflow/compiler/mlir/quantization/tensorflow/python/py_function_lib.h"
+#include "tensorflow/core/protobuf/meta_graph.pb.h"
 
 namespace stablehlo::quantization::pywrap {
 

From f7d9db78dc6226758806d77dd0fc3c0ba598e5cd Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Sun, 19 May 2024 23:04:13 -0700
Subject: [PATCH 405/478] Automated Code Change

PiperOrigin-RevId: 635350042
---
 tensorflow/dtensor/mlir/BUILD                     |  1 +
 tensorflow/dtensor/mlir/function_renaming.cc      |  3 ++-
 .../mlir/handle_cross_cluster_dependencies.cc     | 10 ++++++++--
 tensorflow/dtensor/mlir/handle_sparsetensors.cc   | 15 ++++++++-------
 4 files changed, 19 insertions(+), 10 deletions(-)

diff --git a/tensorflow/dtensor/mlir/BUILD b/tensorflow/dtensor/mlir/BUILD
index e92270975a2f94..5cd39d3619908d 100644
--- a/tensorflow/dtensor/mlir/BUILD
+++ b/tensorflow/dtensor/mlir/BUILD
@@ -273,6 +273,7 @@ cc_library(
         "@com_google_absl//absl/types:span",
         "@llvm-project//llvm:Support",
         "@llvm-project//mlir:FuncDialect",
+        "@llvm-project//mlir:FunctionInterfaces",
         "@llvm-project//mlir:IR",
         "@llvm-project//mlir:Pass",
         "@llvm-project//mlir:SideEffectInterfaces",
diff --git a/tensorflow/dtensor/mlir/function_renaming.cc b/tensorflow/dtensor/mlir/function_renaming.cc
index 0c485e91878d15..3453986221730c 100644
--- a/tensorflow/dtensor/mlir/function_renaming.cc
+++ b/tensorflow/dtensor/mlir/function_renaming.cc
@@ -17,13 +17,14 @@ limitations under the License.
 #include <string>
 
 #include "absl/strings/str_cat.h"
+#include "llvm/ADT/STLExtras.h"
 #include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
 #include "mlir/IR/Attributes.h"  // from @llvm-project
+#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
 #include "mlir/IR/Operation.h"  // from @llvm-project
 #include "mlir/IR/SymbolTable.h"  // from @llvm-project
 #include "mlir/Pass/Pass.h"  // from @llvm-project
 #include "tensorflow/dtensor/cc/constants.h"
-#include "tensorflow/dtensor/mlir/dtensor_mlir_passes.h"
 
 namespace tensorflow {
 namespace dtensor {
diff --git a/tensorflow/dtensor/mlir/handle_cross_cluster_dependencies.cc b/tensorflow/dtensor/mlir/handle_cross_cluster_dependencies.cc
index 6877f77dd25f37..5244e18e13c724 100644
--- a/tensorflow/dtensor/mlir/handle_cross_cluster_dependencies.cc
+++ b/tensorflow/dtensor/mlir/handle_cross_cluster_dependencies.cc
@@ -22,17 +22,23 @@ limitations under the License.
 #include "mlir/IR/Builders.h"  // from @llvm-project
 #include "mlir/IR/BuiltinOps.h"  // from @llvm-project
 #include "mlir/IR/BuiltinTypes.h"  // from @llvm-project
+#include "mlir/IR/DialectRegistry.h"  // from @llvm-project
+#include "mlir/IR/MLIRContext.h"  // from @llvm-project
 #include "mlir/IR/Operation.h"  // from @llvm-project
 #include "mlir/IR/UseDefLists.h"  // from @llvm-project
 #include "mlir/IR/Value.h"  // from @llvm-project
-#include "mlir/Support/DebugStringHelper.h"  // from @llvm-project
+#include "mlir/IR/Visitors.h"  // from @llvm-project
+#include "mlir/Pass/Pass.h"  // from @llvm-project
 #include "mlir/Support/LLVM.h"  // from @llvm-project
 #include "mlir/Support/LogicalResult.h"  // from @llvm-project
-#include "mlir/Transforms/Passes.h"  // from @llvm-project
 #include "mlir/Transforms/RegionUtils.h"  // from @llvm-project
 #include "tensorflow/compiler/mlir/tensorflow/ir/tf_attributes.h"
 #include "tensorflow/compiler/mlir/tensorflow/ir/tf_device.h"
+#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
+#include "tensorflow/dtensor/cc/dstatus.h"
+#include "tensorflow/dtensor/cc/tensor_layout.h"
 #include "tensorflow/dtensor/mlir/dtensor_dialect/ir/dialect.h"
+#include "tensorflow/dtensor/mlir/dtensor_dialect/ir/dtensor_attributes.h"
 #include "tensorflow/dtensor/mlir/ir/tf_dtensor.h"
 #include "tensorflow/dtensor/mlir/layout_parsing.h"
 #include "tensorflow/dtensor/mlir/spmd_expander_common.h"
diff --git a/tensorflow/dtensor/mlir/handle_sparsetensors.cc b/tensorflow/dtensor/mlir/handle_sparsetensors.cc
index 678d8a41841114..1464c47beaf4a4 100644
--- a/tensorflow/dtensor/mlir/handle_sparsetensors.cc
+++ b/tensorflow/dtensor/mlir/handle_sparsetensors.cc
@@ -20,6 +20,8 @@ limitations under the License.
 #include <vector>
 
 #include "absl/container/flat_hash_set.h"
+#include "absl/strings/str_cat.h"
+#include "absl/strings/str_join.h"
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/ADT/StringRef.h"
 #include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
@@ -27,23 +29,22 @@ limitations under the License.
 #include "mlir/IR/Builders.h"  // from @llvm-project
 #include "mlir/IR/BuiltinAttributes.h"  // from @llvm-project
 #include "mlir/IR/BuiltinOps.h"  // from @llvm-project
+#include "mlir/IR/BuiltinTypeInterfaces.h"  // from @llvm-project
 #include "mlir/IR/BuiltinTypes.h"  // from @llvm-project
 #include "mlir/IR/Operation.h"  // from @llvm-project
-#include "mlir/IR/SymbolTable.h"  // from @llvm-project
+#include "mlir/IR/OperationSupport.h"  // from @llvm-project
+#include "mlir/IR/Value.h"  // from @llvm-project
+#include "mlir/IR/ValueRange.h"  // from @llvm-project
 #include "mlir/IR/Visitors.h"  // from @llvm-project
+#include "mlir/Interfaces/FunctionInterfaces.h"  // from @llvm-project
 #include "mlir/Pass/Pass.h"  // from @llvm-project
 #include "mlir/Pass/PassManager.h"  // from @llvm-project
 #include "mlir/Support/LLVM.h"  // from @llvm-project
 #include "mlir/Support/LogicalResult.h"  // from @llvm-project
-#include "mlir/Transforms/Passes.h"  // from @llvm-project
 #include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
 #include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops_n_z.h"
-#include "tensorflow/compiler/mlir/tensorflow/utils/attribute_utils.h"
 #include "tensorflow/dtensor/cc/constants.h"
-#include "tensorflow/dtensor/mlir/device_utils.h"
-#include "tensorflow/dtensor/mlir/dtensor_mlir_passes.h"
-#include "tensorflow/dtensor/mlir/op_utils.h"
-#include "tensorflow/dtensor/mlir/spmd_expander_common.h"
+#include "tensorflow/dtensor/cc/dstatus.h"
 #include "tensorflow/dtensor/mlir/value_utils.h"
 
 namespace tensorflow {

From 887a686efd8859f90a3f4d04a9e79f13e9acb6af Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Mon, 20 May 2024 00:48:27 -0700
Subject: [PATCH 406/478] Automated Code Change

PiperOrigin-RevId: 635370016
---
 tensorflow/core/runtime_fallback/conversion/BUILD         | 6 +++++-
 tensorflow/core/runtime_fallback/conversion/conversion.cc | 6 +++++-
 2 files changed, 10 insertions(+), 2 deletions(-)

diff --git a/tensorflow/core/runtime_fallback/conversion/BUILD b/tensorflow/core/runtime_fallback/conversion/BUILD
index 9b90ab2aa3d062..bf533491c77136 100644
--- a/tensorflow/core/runtime_fallback/conversion/BUILD
+++ b/tensorflow/core/runtime_fallback/conversion/BUILD
@@ -29,11 +29,15 @@ cc_library(
     srcs = ["conversion.cc"],
     hdrs = ["conversion.h"],
     deps = [
+        "//tensorflow/core:framework",
         "//tensorflow/core:protos_all_cc",
+        "//tensorflow/core/common_runtime/eager:tensor_handle",
+        "//tensorflow/core/framework:tensor",
+        "//tensorflow/core/platform:status",
         "//tensorflow/core/runtime_fallback/kernel:kernel_fallback_tensor",
         "//tensorflow/core/runtime_fallback/kernel:tensor_util",
+        "//tensorflow/core/runtime_fallback/runtime:kernel_utils",
         "//tensorflow/core/runtime_fallback/runtime:runtime_fallback_tensor",
-        "//tensorflow/core/tfrt/utils:error_util",
         "@tf_runtime//:hostcontext",
         "@tf_runtime//:tensor",
     ],
diff --git a/tensorflow/core/runtime_fallback/conversion/conversion.cc b/tensorflow/core/runtime_fallback/conversion/conversion.cc
index 71d97860aaa0af..f7a8fe018d18e2 100644
--- a/tensorflow/core/runtime_fallback/conversion/conversion.cc
+++ b/tensorflow/core/runtime_fallback/conversion/conversion.cc
@@ -20,11 +20,15 @@ limitations under the License.
 
 #include <utility>
 
+#include "tensorflow/core/common_runtime/eager/tensor_handle.h"
+#include "tensorflow/core/framework/device.h"
+#include "tensorflow/core/framework/tensor.h"
 #include "tensorflow/core/framework/types.pb.h"
+#include "tensorflow/core/platform/status.h"
 #include "tensorflow/core/runtime_fallback/kernel/kernel_fallback_tensor.h"
 #include "tensorflow/core/runtime_fallback/kernel/tensor_util.h"
+#include "tensorflow/core/runtime_fallback/runtime/kernel_utils.h"
 #include "tensorflow/core/runtime_fallback/runtime/runtime_fallback_tensor.h"
-#include "tensorflow/core/tfrt/utils/error_util.h"
 #include "tfrt/host_context/async_value_ref.h"  // from @tf_runtime
 #include "tfrt/host_context/device.h"  // from @tf_runtime
 #include "tfrt/tensor/conversion_registry.h"  // from @tf_runtime

From 4cf0de11b67d1684a5961744d1dba616999b4c4a Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Mon, 20 May 2024 02:02:16 -0700
Subject: [PATCH 407/478] Update GraphDef version to 1868.

PiperOrigin-RevId: 635385236
---
 tensorflow/core/public/version.h | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorflow/core/public/version.h b/tensorflow/core/public/version.h
index bee0841d5d51bf..d5df864e20d0ca 100644
--- a/tensorflow/core/public/version.h
+++ b/tensorflow/core/public/version.h
@@ -108,7 +108,7 @@ limitations under the License.
 
 #define TF_GRAPH_DEF_VERSION_MIN_PRODUCER 0
 #define TF_GRAPH_DEF_VERSION_MIN_CONSUMER 0
-#define TF_GRAPH_DEF_VERSION 1867  // Updated: 2024/5/19
+#define TF_GRAPH_DEF_VERSION 1868  // Updated: 2024/5/20
 
 // Checkpoint compatibility versions (the versions field in SavedSliceMeta).
 //

From 52f378fde18b0446f225f20f96dd741e93695b82 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Mon, 20 May 2024 02:02:19 -0700
Subject: [PATCH 408/478] compat: Update forward compatibility horizon to
 2024-05-20

PiperOrigin-RevId: 635385257
---
 tensorflow/python/compat/compat.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorflow/python/compat/compat.py b/tensorflow/python/compat/compat.py
index 4ab2aa5a2a9cd1..5beb9afea1416a 100644
--- a/tensorflow/python/compat/compat.py
+++ b/tensorflow/python/compat/compat.py
@@ -29,7 +29,7 @@
 # This value changes every day with an automatic CL. It can be modified in code
 # via `forward_compatibility_horizon()` or with the environment variable
 # TF_FORWARD_COMPATIBILITY_DELTA_DAYS, which is added to the compatibility date.
-_FORWARD_COMPATIBILITY_HORIZON = datetime.date(2024, 5, 19)
+_FORWARD_COMPATIBILITY_HORIZON = datetime.date(2024, 5, 20)
 _FORWARD_COMPATIBILITY_DELTA_DAYS_VAR_NAME = "TF_FORWARD_COMPATIBILITY_DELTA_DAYS"
 _FORWARD_COMPATIBILITY_DATE_NUMBER = None
 

From 718f8209623a2f2b232c5395b295777d4db462c0 Mon Sep 17 00:00:00 2001
From: Penporn Koanantakool <penporn@google.com>
Date: Mon, 20 May 2024 04:23:22 -0700
Subject: [PATCH 409/478] [xla:ffi] Fix msan error for
 //xla/tests:custom_call_test_cpu.

PiperOrigin-RevId: 635414868
---
 third_party/xla/xla/ffi/BUILD         | 1 +
 third_party/xla/xla/ffi/call_frame.cc | 5 +++++
 2 files changed, 6 insertions(+)

diff --git a/third_party/xla/xla/ffi/BUILD b/third_party/xla/xla/ffi/BUILD
index 2f597ceb904bee..3fe3e55c5d8e97 100644
--- a/third_party/xla/xla/ffi/BUILD
+++ b/third_party/xla/xla/ffi/BUILD
@@ -24,6 +24,7 @@ cc_library(
         "//xla/ffi/api:c_api_internal",
         "//xla/stream_executor:device_memory",
         "@com_google_absl//absl/algorithm:container",
+        "@com_google_absl//absl/base:dynamic_annotations",
         "@com_google_absl//absl/container:flat_hash_map",
         "@com_google_absl//absl/log:check",
         "@com_google_absl//absl/types:span",
diff --git a/third_party/xla/xla/ffi/call_frame.cc b/third_party/xla/xla/ffi/call_frame.cc
index ee17f61c7e29fb..cb5ae84822bf4f 100644
--- a/third_party/xla/xla/ffi/call_frame.cc
+++ b/third_party/xla/xla/ffi/call_frame.cc
@@ -26,6 +26,7 @@ limitations under the License.
 #include <vector>
 
 #include "absl/algorithm/container.h"
+#include "absl/base/dynamic_annotations.h"
 #include "absl/log/check.h"
 #include "absl/types/span.h"
 #include "xla/ffi/api/c_api.h"
@@ -83,12 +84,16 @@ void CallFrameBuilder::AddBufferArg(se::DeviceMemoryBase memory,
                                     PrimitiveType type,
                                     absl::Span<const int64_t> dims) {
   args_.push_back(Buffer{memory, type, {dims.begin(), dims.end()}});
+  ABSL_ANNOTATE_MEMORY_IS_INITIALIZED(
+      args_.back().dims.data(), sizeof(int64_t) * args_.back().dims.size());
 }
 
 void CallFrameBuilder::AddBufferRet(se::DeviceMemoryBase memory,
                                     PrimitiveType type,
                                     absl::Span<const int64_t> dims) {
   rets_.push_back(Buffer{memory, type, {dims.begin(), dims.end()}});
+  ABSL_ANNOTATE_MEMORY_IS_INITIALIZED(
+      rets_.back().dims.data(), sizeof(int64_t) * rets_.back().dims.size());
 }
 
 void CallFrameBuilder::AddAttributes(AttributesMap attrs) {

From 312661bcf1f7e40cf93e395f4070b511caf22c58 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Mon, 20 May 2024 07:35:58 -0700
Subject: [PATCH 410/478] Print mismatches for UnorderedElements() of different
 sizes.

Changes the behavior of UnorderedElements()/UnorderedElementsAreArray() to print items-without-matchers and matchers-without-items in the case where the actual and expected are different sizes.

PiperOrigin-RevId: 635451316
---
 tensorflow/compiler/jit/node_matchers_test.cc | 28 +++++++++++++++----
 1 file changed, 22 insertions(+), 6 deletions(-)

diff --git a/tensorflow/compiler/jit/node_matchers_test.cc b/tensorflow/compiler/jit/node_matchers_test.cc
index 8edb3e456c4c00..6f37d5617b6ce6 100644
--- a/tensorflow/compiler/jit/node_matchers_test.cc
+++ b/tensorflow/compiler/jit/node_matchers_test.cc
@@ -15,6 +15,8 @@ limitations under the License.
 
 #include "tensorflow/compiler/jit/node_matchers.h"
 
+#include <string>
+
 #include "tensorflow/cc/framework/ops.h"
 #include "tensorflow/cc/ops/array_ops.h"
 #include "tensorflow/cc/ops/const_op.h"
@@ -117,12 +119,26 @@ TEST(NodeMatchers, CheckControlDependence) {
   EXPECT_THAT(placeholder_d.node(),
               NodeWith(Name("placeholder_d"), CtrlDeps()));
 
-  EXPECT_EQ(
-      Explain(placeholder_c.node(), NodeWith(CtrlDeps())),
-      "ctrl_deps, which has 2 elements, does not match expected: is empty");
-  EXPECT_EQ(Explain(placeholder_d.node(), NodeWith(CtrlDeps(NodeWith()))),
-            "ctrl_deps does not match expected: has 1 element and that element "
-            "is any node");
+  // TODO(griffithjames): Exactly match these explanations.
+  //
+  // When the OSS build has been updated to include the new error messages, the
+  // Explain() expectations can be exact strings again.
+  {
+    const std::string explanation =
+        Explain(placeholder_c.node(), NodeWith(CtrlDeps()));
+    EXPECT_NE(explanation.find("ctrl_deps, which has 2 elements"),
+              std::string::npos);
+    EXPECT_NE(explanation.find("does not match expected: is empty"),
+              std::string::npos);
+  }
+  {
+    const std::string explanation =
+        Explain(placeholder_d.node(), NodeWith(CtrlDeps(NodeWith())));
+    EXPECT_NE(explanation.find("ctrl_deps"), std::string::npos);
+    EXPECT_NE(explanation.find("does not match expected: has 1 element and "
+                               "that element is any node"),
+              std::string::npos);
+  }
 }
 
 TEST(NodeMatchers, ConstValue) {

From d0e577447a11533887249e87e0f9b701d00d527c Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Mon, 20 May 2024 07:54:11 -0700
Subject: [PATCH 411/478] Add float/double template specializations for
 `TensorEq` matcher.

Adds template specializations that improve floating-point equality comparison for float and double tensors.

PiperOrigin-RevId: 635455030
---
 tensorflow/core/framework/tensor_matcher.cc   | 33 ++++++++++++++-----
 .../core/framework/tensor_matcher_test.cc     | 31 +++++++++++++++++
 2 files changed, 55 insertions(+), 9 deletions(-)

diff --git a/tensorflow/core/framework/tensor_matcher.cc b/tensorflow/core/framework/tensor_matcher.cc
index 5e159eac21d6d7..c16f700d467e31 100644
--- a/tensorflow/core/framework/tensor_matcher.cc
+++ b/tensorflow/core/framework/tensor_matcher.cc
@@ -41,13 +41,32 @@ namespace {
 
 using tensorflow::Tensor;
 
+template <typename T>
+::testing::Matcher<absl::Span<const T>> MakePointwiseMatcher(
+    absl::Span<const T> target) {
+  return ::testing::MatcherCast<absl::Span<const T>>(
+      ::testing::Pointwise(::testing::Eq(), target));
+}
+
+template <>
+::testing::Matcher<absl::Span<const float>> MakePointwiseMatcher(
+    absl::Span<const float> target) {
+  return ::testing::MatcherCast<absl::Span<const float>>(
+      ::testing::Pointwise(::testing::FloatEq(), target));
+}
+
+template <>
+::testing::Matcher<absl::Span<const double>> MakePointwiseMatcher(
+    absl::Span<const double> target) {
+  return ::testing::MatcherCast<absl::Span<const double>>(
+      ::testing::Pointwise(::testing::DoubleEq(), target));
+}
+
 template <typename T>
 bool MatchAndExplainPointwise(absl::Span<const T> value,
                               absl::Span<const T> target,
                               ::testing::MatchResultListener* listener) {
-  auto matcher = ::testing::MatcherCast<absl::Span<const T>>(
-      ::testing::Pointwise(::testing::Eq(), target));
-  return matcher.MatchAndExplain(value, listener);
+  return MakePointwiseMatcher<T>(target).MatchAndExplain(value, listener);
 }
 
 class TensorEqMatcherImpl : public ::testing::MatcherInterface<const Tensor&> {
@@ -62,9 +81,7 @@ class TensorEqMatcherImpl : public ::testing::MatcherInterface<const Tensor&> {
   case tensorflow::DataTypeToEnum<T>::value: {             \
     *os << ", and tensor data ";                           \
     absl::Span<const T> data(target_.unaligned_flat<T>()); \
-    ::testing::MatcherCast<absl::Span<const T>>(           \
-        ::testing::Pointwise(::testing::Eq(), data))       \
-        .DescribeTo(os);                                   \
+    MakePointwiseMatcher<T>(data).DescribeTo(os);          \
     break;                                                 \
   }
       TF_CALL_POD_STRING_TYPES(CASE_TYPE);
@@ -84,9 +101,7 @@ class TensorEqMatcherImpl : public ::testing::MatcherInterface<const Tensor&> {
   case tensorflow::DataTypeToEnum<T>::value: {             \
     *os << ", or tensor data ";                            \
     absl::Span<const T> data(target_.unaligned_flat<T>()); \
-    ::testing::MatcherCast<absl::Span<const T>>(           \
-        ::testing::Pointwise(::testing::Eq(), data))       \
-        .DescribeNegationTo(os);                           \
+    MakePointwiseMatcher<T>(data).DescribeNegationTo(os);  \
     break;                                                 \
   }
       TF_CALL_POD_STRING_TYPES(CASE_TYPE);
diff --git a/tensorflow/core/framework/tensor_matcher_test.cc b/tensorflow/core/framework/tensor_matcher_test.cc
index 7e93bde4c964e8..614ea6e7ae353a 100644
--- a/tensorflow/core/framework/tensor_matcher_test.cc
+++ b/tensorflow/core/framework/tensor_matcher_test.cc
@@ -14,6 +14,7 @@ limitations under the License.
 ==============================================================================*/
 #include "tensorflow/core/framework/tensor_matcher.h"
 
+#include <cmath>
 #include <cstdint>
 #include <string>
 #include <vector>
@@ -26,7 +27,9 @@ namespace tensorflow {
 namespace test {
 namespace {
 
+using ::testing::DoubleEq;
 using ::testing::ElementsAre;
+using ::testing::FloatEq;
 
 TEST(TensorMatcherTest, BasicPod) {
   std::vector<Tensor> expected;
@@ -50,6 +53,34 @@ TEST(TensorMatcherTest, BasicString) {
               ElementsAre(TensorEq(Tensor(s1)), TensorEq(Tensor(s2))));
 }
 
+TEST(TensorMatcherTest, FloatComparisonUsesTolerance) {
+  // Two floats that are *nearly* equal.
+  float f1(1);
+  float f2 = std::nextafter(f1, f1 + 1);
+
+  // Direct equality checks should fail, but use of the specialized `FloatEq`
+  // should succeed since this matcher applies ULP-based comparison.
+  // go/matchers#FpMatchers
+  ASSERT_NE(f1, f2);
+  ASSERT_THAT(f1, FloatEq(f2));
+
+  EXPECT_THAT(Tensor(f1), TensorEq(Tensor(f2)));
+}
+
+TEST(TensorMatcherTest, DoubleComparisonUsesTolerance) {
+  // Two doubles that are *nearly* equal.
+  double d1(1);
+  double d2 = std::nextafter(d1, d1 + 1);
+
+  // Direct equality checks should fail, but use of the specialized `DoubleEq`
+  // should succeed since this matcher applies ULP-based comparison.
+  // go/matchers#FpMatchers
+  ASSERT_NE(d1, d2);
+  ASSERT_THAT(d1, DoubleEq(d2));
+
+  EXPECT_THAT(Tensor(d1), TensorEq(Tensor(d2)));
+}
+
 }  // namespace
 }  // namespace test
 }  // namespace tensorflow

From 62c13e80217aad59e0510a5476dd18138878977d Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Mon, 20 May 2024 09:32:15 -0700
Subject: [PATCH 412/478] Compute stats should clone the feature configs before
 updating them.

PiperOrigin-RevId: 635480121
---
 tensorflow/python/tpu/tpu_embedding_v3.py     | 35 ++++++++++---------
 .../python/tpu/tpu_embedding_v3_test.py       |  7 +++-
 2 files changed, 24 insertions(+), 18 deletions(-)

diff --git a/tensorflow/python/tpu/tpu_embedding_v3.py b/tensorflow/python/tpu/tpu_embedding_v3.py
index c32823017f1685..9fcfa060542070 100644
--- a/tensorflow/python/tpu/tpu_embedding_v3.py
+++ b/tensorflow/python/tpu/tpu_embedding_v3.py
@@ -307,6 +307,20 @@ def read_from_device(self, device):
 )
 
 
+def _clone_feature_config(feature_config):
+  old_to_new_table = {}
+  new_features = []
+
+  for old_feature in nest.flatten(feature_config):
+    feature = copy.copy(old_feature)
+    if feature.table not in old_to_new_table:
+      old_to_new_table[feature.table] = copy.copy(feature.table)
+    feature.table = old_to_new_table[feature.table]
+    new_features.append(feature)
+
+  return nest.pack_sequence_as(feature_config, new_features)
+
+
 def _stack_tables_with_same_table_dim_and_optimizer(
     table_config: Sequence[TableConfig],
     flat_features: Sequence[Tuple[Any, FeatureConfig]],
@@ -496,7 +510,7 @@ def __init__(
     # We do a clone on the feature_config here as we will alter settings in it
     # and we don't want the user to see these. We can't just use clone here
     # as we need to maintain some object relationships.
-    super().__init__(self._clone_feature_config(feature_config), optimizer)
+    super().__init__(_clone_feature_config(feature_config), optimizer)
     self._strategy = distribute_lib.get_strategy()
     if not isinstance(
         self._strategy, (tpu_strategy.TPUStrategy, tpu_strategy.TPUStrategyV2)
@@ -621,19 +635,6 @@ def _update_sparse_core_buffer_size_after_table_stacking(self):
             ]
         )
 
-  def _clone_feature_config(self, feature_config):
-    old_to_new_table = {}
-    new_features = []
-
-    for old_feature in nest.flatten(feature_config):
-      feature = copy.copy(old_feature)
-      if feature.table not in old_to_new_table:
-        old_to_new_table[feature.table] = copy.copy(feature.table)
-      feature.table = old_to_new_table[feature.table]
-      new_features.append(feature)
-
-    return nest.pack_sequence_as(feature_config, new_features)
-
   @property
   def embedding_tables(
       self,
@@ -1094,16 +1095,16 @@ def compute_sparse_core_stats(
       sparse_core_embedding_config: Optional[SparseCoreEmbeddingConfig] = None,
   ) -> Tuple[Any, Any]:
     """Computes the max_ids/unique ids settings from the input features."""
-
+    copy_feature_config = _clone_feature_config(feature_config)
     table_config = []
-    for feature in nest.flatten(feature_config):
+    for feature in nest.flatten(copy_feature_config):
       table_config.append(feature.table)
 
     for table in table_config:
       if table.optimizer is None:
         table.optimizer = optimizer
 
-    flat_features = nest.flatten_with_joined_string_paths(feature_config)
+    flat_features = nest.flatten_with_joined_string_paths(copy_feature_config)
 
     s = _stack_tables_with_same_table_dim_and_optimizer(
         table_config,
diff --git a/tensorflow/python/tpu/tpu_embedding_v3_test.py b/tensorflow/python/tpu/tpu_embedding_v3_test.py
index 75d9edff9bbdb7..433eac09c06759 100644
--- a/tensorflow/python/tpu/tpu_embedding_v3_test.py
+++ b/tensorflow/python/tpu/tpu_embedding_v3_test.py
@@ -684,6 +684,9 @@ def test_compute_sparse_core_stats_and_pass_it_to_api(self):
             optimizer=tpu_embedding_v2_utils.SGD(learning_rate=1.0),
         )
     )
+    self.assertEqual(
+        feature_config.table.dim, 127, 'Unexpected update to FeatureConfig'
+    )
 
     sparse_core_embedding_config = tpu_embedding_v3.SparseCoreEmbeddingConfig(
         disable_table_stacking=False,
@@ -708,7 +711,9 @@ def step(data):
       return strategy.run(step, args=(data,))
 
     result = test_fn()
-
+    self.assertEqual(
+        feature_config.table.dim, 127, 'Unexpected update to FeatureConfig'
+    )
     mid_level_api_cpu = tpu_embedding_for_serving.TPUEmbeddingForServing(
         feature_config=feature_config,
         optimizer=tpu_embedding_v2_utils.SGD(learning_rate=1.0),

From 943e1935d4ca1b3721b6ca8dbcd9f9c64778535c Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Mon, 20 May 2024 10:20:48 -0700
Subject: [PATCH 413/478] Clean up include and build file

PiperOrigin-RevId: 635495307
---
 tensorflow/core/tfrt/ifrt/BUILD                            | 7 +------
 tensorflow/core/tfrt/ifrt/ifrt_executable_registry.h       | 1 +
 tensorflow/core/tfrt/ifrt/ifrt_loaded_variable_registry.cc | 1 -
 tensorflow/core/tfrt/ifrt/ifrt_loaded_variable_registry.h  | 1 +
 tensorflow/core/tfrt/ifrt/ifrt_loaded_variable_utils.h     | 3 +++
 tensorflow/core/tfrt/ifrt/ifrt_model_context.cc            | 4 ----
 tensorflow/core/tfrt/ifrt/ifrt_model_context.h             | 5 -----
 tensorflow/core/tfrt/ifrt/ifrt_serving_executable.h        | 2 --
 tensorflow/core/tfrt/ifrt/sharding_utils.cc                | 1 -
 tensorflow/core/tfrt/ifrt/sharding_utils.h                 | 4 +---
 tensorflow/core/tfrt/ifrt/sharding_utils_test.cc           | 1 +
 11 files changed, 8 insertions(+), 22 deletions(-)

diff --git a/tensorflow/core/tfrt/ifrt/BUILD b/tensorflow/core/tfrt/ifrt/BUILD
index 48b9e6303e1c27..8b3d61c8bbbc12 100644
--- a/tensorflow/core/tfrt/ifrt/BUILD
+++ b/tensorflow/core/tfrt/ifrt/BUILD
@@ -146,6 +146,7 @@ cc_library(
     srcs = ["ifrt_loaded_variable_registry.cc"],
     hdrs = ["ifrt_loaded_variable_registry.h"],
     deps = [
+        "@com_google_absl//absl/base:core_headers",
         "@com_google_absl//absl/container:flat_hash_map",
         "@com_google_absl//absl/container:flat_hash_set",
         "@com_google_absl//absl/functional:any_invocable",
@@ -171,15 +172,11 @@ cc_library(
         ":ifrt_serving_core_selector",
         "//tensorflow/compiler/tf2xla:xla_helpers",
         "//tensorflow/core:core_cpu_base",
-        "@com_google_absl//absl/container:flat_hash_map",
         "@com_google_absl//absl/status",
-        "@com_google_absl//absl/status:statusor",
         "@com_google_absl//absl/strings",
-        "@com_google_absl//absl/synchronization",
         "@local_tsl//tsl/platform:env",
         "@local_tsl//tsl/platform:errors",
         "@local_xla//xla/python/ifrt",
-        "@local_xla//xla/tsl/concurrency:ref_count",
         "@tf_runtime//:hostcontext",
     ],
 )
@@ -220,7 +217,6 @@ cc_library(
         "//tensorflow/core:framework",
         "//tensorflow/core:protos_all_cc",
         "//tensorflow/core/platform:status",
-        "//tensorflow/core/platform:statusor",
         "//tensorflow/core/tpu/kernels:sharding_utils",
         "@com_google_absl//absl/container:btree",
         "@com_google_absl//absl/log",
@@ -233,7 +229,6 @@ cc_library(
         "@local_tsl//tsl/platform:env",
         "@local_tsl//tsl/platform:errors",
         "@local_tsl//tsl/platform:statusor",
-        "@local_xla//xla:executable_run_options",
         "@local_xla//xla/hlo/ir:hlo",
         "@local_xla//xla/python/ifrt",
         "@local_xla//xla/python/pjrt_ifrt:xla_ifrt",
diff --git a/tensorflow/core/tfrt/ifrt/ifrt_executable_registry.h b/tensorflow/core/tfrt/ifrt/ifrt_executable_registry.h
index c57f0c0f353387..5275b683027602 100644
--- a/tensorflow/core/tfrt/ifrt/ifrt_executable_registry.h
+++ b/tensorflow/core/tfrt/ifrt/ifrt_executable_registry.h
@@ -22,6 +22,7 @@ limitations under the License.
 
 #include "absl/base/thread_annotations.h"
 #include "absl/container/flat_hash_map.h"
+#include "absl/status/status.h"
 #include "absl/status/statusor.h"
 #include "absl/synchronization/mutex.h"
 #include "tensorflow/core/tfrt/ifrt/ifrt_serving_executable.h"
diff --git a/tensorflow/core/tfrt/ifrt/ifrt_loaded_variable_registry.cc b/tensorflow/core/tfrt/ifrt/ifrt_loaded_variable_registry.cc
index 5ef8f3cbf77e03..73651cf7ed5cf1 100644
--- a/tensorflow/core/tfrt/ifrt/ifrt_loaded_variable_registry.cc
+++ b/tensorflow/core/tfrt/ifrt/ifrt_loaded_variable_registry.cc
@@ -22,7 +22,6 @@ limitations under the License.
 #include "absl/status/status.h"
 #include "absl/status/statusor.h"
 #include "absl/strings/str_cat.h"
-#include "absl/strings/string_view.h"
 #include "absl/synchronization/mutex.h"
 #include "tsl/platform/statusor.h"
 
diff --git a/tensorflow/core/tfrt/ifrt/ifrt_loaded_variable_registry.h b/tensorflow/core/tfrt/ifrt/ifrt_loaded_variable_registry.h
index 0c66af49a50fc4..e799c571a246b7 100644
--- a/tensorflow/core/tfrt/ifrt/ifrt_loaded_variable_registry.h
+++ b/tensorflow/core/tfrt/ifrt/ifrt_loaded_variable_registry.h
@@ -18,6 +18,7 @@ limitations under the License.
 
 #include <string>
 
+#include "absl/base/thread_annotations.h"
 #include "absl/container/flat_hash_map.h"
 #include "absl/container/flat_hash_set.h"
 #include "absl/functional/any_invocable.h"
diff --git a/tensorflow/core/tfrt/ifrt/ifrt_loaded_variable_utils.h b/tensorflow/core/tfrt/ifrt/ifrt_loaded_variable_utils.h
index dbf464ce4540d9..4d07d1a3771a8a 100644
--- a/tensorflow/core/tfrt/ifrt/ifrt_loaded_variable_utils.h
+++ b/tensorflow/core/tfrt/ifrt/ifrt_loaded_variable_utils.h
@@ -20,6 +20,9 @@ limitations under the License.
 #include <string>
 
 #include "absl/status/status.h"
+#include "absl/status/statusor.h"
+#include "absl/strings/string_view.h"
+#include "tensorflow/compiler/mlir/tfrt/transforms/ifrt/ifrt_types.h"
 #include "xla/python/ifrt/client.h"
 #include "tensorflow/core/framework/resource_handle.h"
 #include "tensorflow/core/tfrt/ifrt/ifrt_config.pb.h"
diff --git a/tensorflow/core/tfrt/ifrt/ifrt_model_context.cc b/tensorflow/core/tfrt/ifrt/ifrt_model_context.cc
index 2d5b6b79002089..191fc765715ba7 100644
--- a/tensorflow/core/tfrt/ifrt/ifrt_model_context.cc
+++ b/tensorflow/core/tfrt/ifrt/ifrt_model_context.cc
@@ -16,12 +16,8 @@ limitations under the License.
 
 #include "tensorflow/core/tfrt/ifrt/ifrt_model_context.h"
 
-#include <utility>
 
 #include "absl/status/status.h"
-#include "absl/strings/string_view.h"
-#include "xla/python/ifrt/array.h"
-#include "xla/tsl/concurrency/ref_count.h"
 #include "tsl/platform/errors.h"
 #include "tsl/platform/threadpool.h"
 
diff --git a/tensorflow/core/tfrt/ifrt/ifrt_model_context.h b/tensorflow/core/tfrt/ifrt/ifrt_model_context.h
index 76813a3a80937a..6b4a600b85a3cb 100644
--- a/tensorflow/core/tfrt/ifrt/ifrt_model_context.h
+++ b/tensorflow/core/tfrt/ifrt/ifrt_model_context.h
@@ -17,19 +17,14 @@ limitations under the License.
 #define TENSORFLOW_CORE_TFRT_IFRT_IFRT_MODEL_CONTEXT_H_
 
 #include <memory>
-#include <string>
 #include <utility>
 #include <vector>
 
-#include "absl/container/flat_hash_map.h"
 #include "absl/status/status.h"
-#include "absl/status/statusor.h"
 #include "absl/strings/string_view.h"
-#include "absl/synchronization/mutex.h"
 #include "tensorflow/compiler/tf2xla/xla_helpers.h"
 #include "xla/python/ifrt/array.h"
 #include "xla/python/ifrt/client.h"
-#include "xla/tsl/concurrency/ref_count.h"
 #include "tensorflow/core/common_runtime/device_mgr.h"
 #include "tensorflow/core/tfrt/ifrt/ifrt_executable_registry.h"
 #include "tensorflow/core/tfrt/ifrt/ifrt_loaded_variable_registry.h"
diff --git a/tensorflow/core/tfrt/ifrt/ifrt_serving_executable.h b/tensorflow/core/tfrt/ifrt/ifrt_serving_executable.h
index 1d224d9345c9dd..f7a67e680a07e0 100644
--- a/tensorflow/core/tfrt/ifrt/ifrt_serving_executable.h
+++ b/tensorflow/core/tfrt/ifrt/ifrt_serving_executable.h
@@ -25,14 +25,12 @@ limitations under the License.
 
 #include "absl/base/thread_annotations.h"
 #include "absl/container/flat_hash_map.h"
-#include "absl/container/flat_hash_set.h"
 #include "absl/log/log.h"
 #include "absl/status/statusor.h"
 #include "absl/strings/string_view.h"
 #include "absl/synchronization/mutex.h"
 #include "absl/types/span.h"
 #include "mlir/IR/BuiltinOps.h"  // from @llvm-project
-#include "mlir/IR/MLIRContext.h"  // from @llvm-project
 #include "mlir/IR/OwningOpRef.h"  // from @llvm-project
 #include "tensorflow/compiler/mlir/tfrt/transforms/ifrt/ifrt_types.h"
 #include "tensorflow/compiler/tf2xla/xla_helpers.h"
diff --git a/tensorflow/core/tfrt/ifrt/sharding_utils.cc b/tensorflow/core/tfrt/ifrt/sharding_utils.cc
index 13e0c2941f4201..73ddec940649e2 100644
--- a/tensorflow/core/tfrt/ifrt/sharding_utils.cc
+++ b/tensorflow/core/tfrt/ifrt/sharding_utils.cc
@@ -52,7 +52,6 @@ limitations under the License.
 #include "tensorflow/core/framework/types.h"
 #include "tensorflow/core/framework/types.pb.h"
 #include "tensorflow/core/platform/status.h"
-#include "tensorflow/core/platform/statusor.h"
 #include "tensorflow/core/tfrt/ifrt/ifrt_tensor_utils.h"
 #include "tensorflow/core/tpu/kernels/sharding_utils.h"
 #include "tsl/platform/errors.h"
diff --git a/tensorflow/core/tfrt/ifrt/sharding_utils.h b/tensorflow/core/tfrt/ifrt/sharding_utils.h
index 96c99a061d2855..12f04bdc990692 100644
--- a/tensorflow/core/tfrt/ifrt/sharding_utils.h
+++ b/tensorflow/core/tfrt/ifrt/sharding_utils.h
@@ -17,15 +17,13 @@ limitations under the License.
 #define TENSORFLOW_CORE_TFRT_IFRT_SHARDING_UTILS_H_
 
 #include "absl/status/statusor.h"
-#include "xla/executable_run_options.h"
+#include "absl/types/span.h"
 #include "xla/hlo/ir/hlo_sharding.h"
 #include "xla/python/ifrt/array.h"
 #include "xla/python/ifrt/client.h"
 #include "xla/python/ifrt/device.h"
-#include "xla/python/ifrt/future.h"
 #include "xla/tsl/concurrency/ref_count.h"
 #include "tensorflow/core/framework/tensor.h"
-#include "tensorflow/core/platform/statusor.h"
 #include "tsl/platform/threadpool.h"
 
 namespace tensorflow {
diff --git a/tensorflow/core/tfrt/ifrt/sharding_utils_test.cc b/tensorflow/core/tfrt/ifrt/sharding_utils_test.cc
index c0fc64851868e9..ded49b2c0f14c0 100644
--- a/tensorflow/core/tfrt/ifrt/sharding_utils_test.cc
+++ b/tensorflow/core/tfrt/ifrt/sharding_utils_test.cc
@@ -28,6 +28,7 @@ limitations under the License.
 #include "xla/python/ifrt/array.h"
 #include "xla/python/ifrt/client.h"
 #include "xla/python/ifrt/device.h"
+#include "xla/python/ifrt/memory.h"
 #include "xla/python/ifrt/shape.h"
 #include "xla/python/ifrt/sharding.h"
 #include "xla/python/ifrt/test_util.h"

From c0da8716b09c04f075d5219066e2c29cb971db92 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Mon, 20 May 2024 10:33:39 -0700
Subject: [PATCH 414/478] Migrate away from llvm::StringRef::equals

Note that llvm::StringRef::equals has been deprecated upstream.

PiperOrigin-RevId: 635499769
---
 tensorflow/dtensor/mlir/expansions/reduce_spmd_expander.cc | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/tensorflow/dtensor/mlir/expansions/reduce_spmd_expander.cc b/tensorflow/dtensor/mlir/expansions/reduce_spmd_expander.cc
index 1ee4a60cd4d3f6..4905d49c8faf1a 100644
--- a/tensorflow/dtensor/mlir/expansions/reduce_spmd_expander.cc
+++ b/tensorflow/dtensor/mlir/expansions/reduce_spmd_expander.cc
@@ -132,11 +132,11 @@ Status ExtractDims<mlir::TF::BiasAddGradOp>(
   auto data_format = bias_add_grad_op.getDataFormat();
   // rank is at least 2 (required by BiasAddGrad).
   int rank = ValueRank(bias_add_grad_op->getOperand(0));
-  if (data_format.equals("NHWC")) {
+  if (data_format == "NHWC") {
     for (int dim = 0; dim < rank - 1; ++dim) {
       reduced_dims->push_back(dim);
     }
-  } else if (data_format.equals("NCHW")) {
+  } else if (data_format == "NCHW") {
     for (int dim = 0; dim < rank; ++dim) {
       if (dim == 1) continue;
       reduced_dims->push_back(dim);

From 0f3594eda11236795339002136ae1f1e4009fe66 Mon Sep 17 00:00:00 2001
From: Kanglan Tang <kanglan@google.com>
Date: Mon, 20 May 2024 10:43:00 -0700
Subject: [PATCH 415/478] Update np.float_ and np.complex_ to match Numpy 2.0.

Numpy 2.0 removed the aliases np.float_ and np.complex_. Direct usages of np.float_ and np.complex_ are replaced with np.float64 and np.complex128, respectively.
See https://github.com/numpy/numpy/issues/24743 and https://numpy.org/devdocs/numpy_2_0_migration_guide.html#changes-to-namespaces.

Consistent with JAX, TF retains dtype wrappers for float_ and complex_. This is considered useful for determining the default float/complex type, providing compatibility with older code that might still rely on these aliases.

PiperOrigin-RevId: 635502855
---
 tensorflow/python/ops/numpy_ops/np_dtypes.py  | 22 +++++++++------
 .../python/ops/numpy_ops/tests/np_test.py     | 28 +++++++++----------
 .../python/ops/numpy_ops/tests/test_util.py   |  4 +--
 3 files changed, 30 insertions(+), 24 deletions(-)

diff --git a/tensorflow/python/ops/numpy_ops/np_dtypes.py b/tensorflow/python/ops/numpy_ops/np_dtypes.py
index 00e2cdfc22964b..ec3608df67cb7b 100644
--- a/tensorflow/python/ops/numpy_ops/np_dtypes.py
+++ b/tensorflow/python/ops/numpy_ops/np_dtypes.py
@@ -27,10 +27,6 @@
 tf_export.tf_export('experimental.numpy.bool_', v1=[]).export_constant(
     __name__, 'bool_'
 )
-complex_ = np.complex_
-tf_export.tf_export('experimental.numpy.complex_', v1=[]).export_constant(
-    __name__, 'complex_'
-)
 complex128 = np.complex128
 tf_export.tf_export('experimental.numpy.complex128', v1=[]).export_constant(
     __name__, 'complex128'
@@ -39,10 +35,6 @@
 tf_export.tf_export('experimental.numpy.complex64', v1=[]).export_constant(
     __name__, 'complex64'
 )
-float_ = np.float_
-tf_export.tf_export('experimental.numpy.float_', v1=[]).export_constant(
-    __name__, 'float_'
-)
 float16 = np.float16
 tf_export.tf_export('experimental.numpy.float16', v1=[]).export_constant(
     __name__, 'float16'
@@ -107,6 +99,20 @@
 tf_export.tf_export('experimental.numpy.unicode_', v1=[]).export_constant(
     __name__, 'unicode_'
 )
+if int(np.__version__.split('.')[0]) < 2:
+  complex_ = np.complex_
+  float_ = np.float_
+else:
+  # Aliases np.complex_ and np.float_ have been removed in Numpy 2.0. Use
+  # np.complex128 and np.float64 instead.
+  complex_ = np.complex128
+  float_ = np.float64
+tf_export.tf_export('experimental.numpy.complex_', v1=[]).export_constant(
+    __name__, 'complex_'
+)
+tf_export.tf_export('experimental.numpy.float_', v1=[]).export_constant(
+    __name__, 'float_'
+)
 
 
 iinfo = np.iinfo
diff --git a/tensorflow/python/ops/numpy_ops/tests/np_test.py b/tensorflow/python/ops/numpy_ops/tests/np_test.py
index d27c39a896dea7..fde60098841096 100644
--- a/tensorflow/python/ops/numpy_ops/tests/np_test.py
+++ b/tensorflow/python/ops/numpy_ops/tests/np_test.py
@@ -62,7 +62,7 @@
 all_dtypes = number_dtypes + bool_dtypes
 
 
-python_scalar_dtypes = [tnp.bool_, tnp.int_, tnp.float_, tnp.complex_]
+python_scalar_dtypes = [tnp.bool_, tnp.int_, tnp.float64, tnp.complex128]
 # pylint: disable=unnecessary-lambda,g-long-lambda,expression-not-assigned
 
 def _valid_dtypes_for_shape(shape, dtypes):
@@ -434,7 +434,7 @@ def wrapper(*args, **kw):
     if inexact and not any(
         tnp.issubdtype(tnp.result_type(x).as_numpy_dtype, tnp.inexact)
         for x in flat_args):
-      dtype = tnp.result_type(tnp.float_, *flat_args)
+      dtype = tnp.result_type(tnp.float64, *flat_args)
     else:
       dtype = tnp.result_type(*flat_args)
     dtype = dtype.as_numpy_dtype
@@ -479,7 +479,7 @@ class LaxBackedNumpyTests(jtu.TestCase):
 
   def _GetArgsMaker(self, rng, shapes, dtypes, onp_arrays=True):
     def f():
-      out = [rng(shape, dtype or tnp.float_)
+      out = [rng(shape, dtype or tnp.float64)
              for shape, dtype in zip(shapes, dtypes)]
       return out if onp_arrays else [tnp.asarray(a) for a in out]
     return f
@@ -1775,14 +1775,14 @@ def testAverage(self, shape, dtype, axis, weights_shape, returned, rng_factory):
        "arg": arg, "ndmin": ndmin, "dtype": dtype}
       for i, (arg, dtype) in enumerate([
           ([True, False, True], tnp.bool_),
-          (3., tnp.float_),
+          (3., tnp.float64),
           ([1, 2, 3], tnp.int_),
-          ([1., 2., 3.], tnp.float_),
+          ([1., 2., 3.], tnp.float64),
           ([[1, 2], [3, 4], [5, 6]], tnp.int_),
-          ([[1, 2.], [3, 4], [5, 6]], tnp.float_),
-          ([[1., 2j], [3., 4.], [5., 6.]], tnp.complex_),
-          ([[3, onp.array(2, dtype=tnp.float_), 1],
-           onp.arange(3., dtype=tnp.float_)], tnp.float_),  # pylint: disable=bad-continuation
+          ([[1, 2.], [3, 4], [5, 6]], tnp.float64),
+          ([[1., 2j], [3., 4.], [5., 6.]], tnp.complex128),
+          ([[3, onp.array(2, dtype=tnp.float64), 1],
+           onp.arange(3., dtype=tnp.float64)], tnp.float64),  # pylint: disable=bad-continuation
       ])
       for ndmin in [None, onp.ndim(arg), onp.ndim(arg) + 1, onp.ndim(arg) + 2]))
   def testArray(self, arg, ndmin, dtype):
@@ -2055,14 +2055,14 @@ def testAstype(self):
 
   def testOnpMean(self):
     # from https://github.com/google/jax/issues/125
-    x = tnp.add(tnp.eye(3, dtype=tnp.float_), 0.)
+    x = tnp.add(tnp.eye(3, dtype=tnp.float64), 0.)
     ans = onp.mean(x)
     self.assertAllClose(ans, onp.array(1./3), check_dtypes=False)
 
   @jtu.disable
   def testArangeOnFloats(self):
     # from https://github.com/google/jax/issues/145
-    expected = onp.arange(0.0, 1.0, 0.1, dtype=tnp.float_)
+    expected = onp.arange(0.0, 1.0, 0.1, dtype=tnp.float64)
     ans = tnp.arange(0.0, 1.0, 0.1)
     self.assertAllClose(expected, ans, check_dtypes=True)
 
@@ -2431,8 +2431,8 @@ def testIssue453(self):
   @named_parameters(jtu.cases_from_list(
       {"testcase_name": "_op={}_dtype={}".format(op, pytype.__name__),
        "pytype": pytype, "dtype": dtype, "op": op}
-      for pytype, dtype in [(int, tnp.int_), (float, tnp.float_),
-                            (bool, tnp.bool_), (complex, tnp.complex_)]
+      for pytype, dtype in [(int, tnp.int_), (float, tnp.float64),
+                            (bool, tnp.bool_), (complex, tnp.complex128)]
       for op in ["atleast_1d", "atleast_2d", "atleast_3d"]))
   def testAtLeastNdLiterals(self, pytype, dtype, op):
     # Fixes: https://github.com/google/jax/issues/634
@@ -2470,7 +2470,7 @@ def testArange(self):
         onp.arange(2, 13, dtype=int),
         check_dtypes=True)
     self.assertAllClose(tnp.arange(0, 1, -0.5),
-                        onp.arange(0, 1, -0.5, dtype=tnp.float_),
+                        onp.arange(0, 1, -0.5, dtype=tnp.float64),
                         check_dtypes=True)
 
     self.assertRaises(TypeError, lambda: tnp.arange())
diff --git a/tensorflow/python/ops/numpy_ops/tests/test_util.py b/tensorflow/python/ops/numpy_ops/tests/test_util.py
index cf178a3f5dfbcc..086b182bf2a628 100644
--- a/tensorflow/python/ops/numpy_ops/tests/test_util.py
+++ b/tensorflow/python/ops/numpy_ops/tests/test_util.py
@@ -76,8 +76,8 @@
 python_scalar_dtypes = {
   bool: onp.dtype(onp.bool_),
   int: onp.dtype(onp.int_),
-  float: onp.dtype(onp.float_),
-  complex: onp.dtype(onp.complex_),
+  float: onp.dtype(onp.float64),
+  complex: onp.dtype(onp.complex128),
 }
 
 

From 9631d5f31bb25c0d775f41adbd07120c12121feb Mon Sep 17 00:00:00 2001
From: David Dunleavy <ddunleavy@google.com>
Date: Mon, 20 May 2024 10:47:16 -0700
Subject: [PATCH 416/478] Reenable/tag previously broken tests on ARM builds of
 XLA

Some tests work now because of hermetic python in XLA, while cross-compilation still breaks other tests.

PiperOrigin-RevId: 635504270
---
 third_party/xla/.kokoro/linux/build.sh | 3 ---
 third_party/xla/xla/service/BUILD      | 1 +
 third_party/xla/xla/tests/BUILD        | 1 +
 3 files changed, 2 insertions(+), 3 deletions(-)

diff --git a/third_party/xla/.kokoro/linux/build.sh b/third_party/xla/.kokoro/linux/build.sh
index b1ac5b53a8c0aa..c32a54520a5a05 100644
--- a/third_party/xla/.kokoro/linux/build.sh
+++ b/third_party/xla/.kokoro/linux/build.sh
@@ -73,9 +73,6 @@ else
 
     if is_linux_cpu_arm64_job ; then
         TAGS_FILTER="$TAGS_FILTER,-no_aarch64"
-        # Some cross-compile tests are not working for XLA Linux Aarch64.
-        # TODO(ddunleavy): Revisit these when hermetic python is available.
-        TARGET_FILTERS="$TARGET_FILTERS -//xla/python_api:xla_shape_test -//xla/python_api:xla_literal_test -//xla/service:xla_aot_compile_stablehlo_cpu_test -//xla/tests:local_client_aot_test"
         RBE_FLAGS="--config=rbe_cross_compile_linux_arm64_xla --jobs=150"
     else
         RBE_FLAGS="--config=rbe_linux_cpu --jobs=150"
diff --git a/third_party/xla/xla/service/BUILD b/third_party/xla/xla/service/BUILD
index 51e4a64911d110..9b91d9ec8a80fc 100644
--- a/third_party/xla/xla/service/BUILD
+++ b/third_party/xla/xla/service/BUILD
@@ -7784,6 +7784,7 @@ xla_cc_test(
     name = "xla_aot_compile_stablehlo_cpu_test",
     srcs = ["xla_aot_compile_stablehlo_cpu_test.cc"],
     data = [":xla_aot_compile_stablehlo_test_cpu_executable"],
+    tags = ["no_aarch64"],  # b/341355246
     deps = [
         ":cpu_plugin",
         ":platform_util",
diff --git a/third_party/xla/xla/tests/BUILD b/third_party/xla/xla/tests/BUILD
index 4bd2e01e65a6dc..0b8f0cf6889870 100644
--- a/third_party/xla/xla/tests/BUILD
+++ b/third_party/xla/xla/tests/BUILD
@@ -2509,6 +2509,7 @@ xla_cc_test(
         ":local_client_aot_test_computation.o",
     ],
     linkstatic = 1,
+    tags = ["no_aarch64"],  # b/341355246
     deps = [
         "//xla:executable_run_options",
         "@com_google_absl//absl/base:dynamic_annotations",

From ac0724c3b087c6a86aad133237f42071b2d90b81 Mon Sep 17 00:00:00 2001
From: George Karpenkov <cheshire@google.com>
Date: Mon, 20 May 2024 10:51:27 -0700
Subject: [PATCH 417/478] Simplify JAX lowering rules for cumulative sum

Upstream fix has landed => removing CPU workaround.

PiperOrigin-RevId: 635505632
---
 third_party/xla/xla/python/xla_client.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/third_party/xla/xla/python/xla_client.py b/third_party/xla/xla/python/xla_client.py
index af0590c12c0f8a..b649011958ff32 100644
--- a/third_party/xla/xla/python/xla_client.py
+++ b/third_party/xla/xla/python/xla_client.py
@@ -49,7 +49,7 @@
 
 # Just an internal arbitrary increasing number to help with backward-compatible
 # changes. In JAX, reference this via jax._src.lib.xla_extension_version.
-_version = 265
+_version = 266
 
 # Version number for MLIR:Python components.
 mlir_api_version = 56

From 647277e7a137fe50d44c4ccb730315227b60b6fc Mon Sep 17 00:00:00 2001
From: David Dunleavy <ddunleavy@google.com>
Date: Mon, 20 May 2024 10:51:34 -0700
Subject: [PATCH 418/478] Remove unnecessary uses of `gpu_any` backend

If a test is intended for the v100 backend, that will run in OSS, so no need for the extra `gpu_any` backend.

The usage in hlo_opt should've been removed already in a previous change.

PiperOrigin-RevId: 635505667
---
 third_party/xla/xla/service/gpu/BUILD   | 4 ++--
 third_party/xla/xla/tools/hlo_opt/BUILD | 2 +-
 2 files changed, 3 insertions(+), 3 deletions(-)

diff --git a/third_party/xla/xla/service/gpu/BUILD b/third_party/xla/xla/service/gpu/BUILD
index afe18d35d71af7..f3e0c8f11adc03 100644
--- a/third_party/xla/xla/service/gpu/BUILD
+++ b/third_party/xla/xla/service/gpu/BUILD
@@ -1677,7 +1677,7 @@ xla_test(
     srcs = if_gpu_is_configured(["gemm_algorithm_picker_test.cc"]),
     backends = [
         "gpu_v100",
-    ] + if_oss(["gpu_any"]),
+    ],
     deps = [
         ":autotuner_util",
         ":backend_configs_cc",
@@ -1974,7 +1974,7 @@ xla_test(
     srcs = if_gpu_is_configured(["conv_algorithm_picker_test.cc"]),
     backends = [
         "gpu_v100",
-    ] + if_oss(["gpu_any"]),
+    ],
     tags = [
         "noasan",
         "nomsan",
diff --git a/third_party/xla/xla/tools/hlo_opt/BUILD b/third_party/xla/xla/tools/hlo_opt/BUILD
index 7657f4624f7bda..4d69ab619d70d0 100644
--- a/third_party/xla/xla/tools/hlo_opt/BUILD
+++ b/third_party/xla/xla/tools/hlo_opt/BUILD
@@ -177,7 +177,7 @@ lit_test_suite(
     ]),
     cfg = "//xla:lit.cfg.py",
     data = [":test_utilities"],
-    default_tags = tf_cuda_tests_tags() + ["gpu_any"],
+    default_tags = tf_cuda_tests_tags(),
     tags_override = {
         "gpu_hlo_ptx.hlo": ["no_rocm"],
     },

From 138d910519c12b83bc01a4529bf0e0945a273cf1 Mon Sep 17 00:00:00 2001
From: Yifan Jiang <yifanjiang@google.com>
Date: Mon, 20 May 2024 11:00:01 -0700
Subject: [PATCH 419/478] Use AllocateDestinationBuffer instead of
 CreateUninitializedBuffer for AsyncHostToDeviceTransferManager::Create. This
 is to synchronize the h2d stream with the compute stream.

With this change, PjRtStreamExecutorBuffer::Delete doesn't need to schedule deallocation on compute stream if there are no other use cases that don't wait for the compute steam before writing to the HBM.

PiperOrigin-RevId: 635508746
---
 .../xla/xla/pjrt/gpu/se_gpu_pjrt_client.cc    | 13 ++-
 .../xla/pjrt/pjrt_stream_executor_client.cc   | 88 ++++++++-----------
 .../xla/pjrt/pjrt_stream_executor_client.h    | 19 ++++
 3 files changed, 65 insertions(+), 55 deletions(-)

diff --git a/third_party/xla/xla/pjrt/gpu/se_gpu_pjrt_client.cc b/third_party/xla/xla/pjrt/gpu/se_gpu_pjrt_client.cc
index a1be7fe67058db..a976e9cd9e6a05 100644
--- a/third_party/xla/xla/pjrt/gpu/se_gpu_pjrt_client.cc
+++ b/third_party/xla/xla/pjrt/gpu/se_gpu_pjrt_client.cc
@@ -135,9 +135,18 @@ class AsyncHostToDeviceTransferManager
       // event will block the buffer usage until the transfer is done.
       definition_events.push_back(
           std::make_shared<BufferSequencingEvent>(client->thread_pool()));
+      TF_ASSIGN_OR_RETURN(Shape compact_shape,
+                          client->client()
+                              ->backend()
+                              .transfer_manager()
+                              ->ChooseCompactLayoutForShape(shape));
+      LocalDeviceState* local_device = device->local_device_state();
+      se::Stream* h2d_stream = local_device->host_to_device_stream();
       TF_ASSIGN_OR_RETURN(auto buffer,
-                          client->CreateUninitializedBuffer(
-                              shape, device, definition_events.back()));
+                          AllocateDestinationBuffer(
+                              compact_shape, device, local_device, h2d_stream,
+                              /*is_uninitialized_create=*/true, client,
+                              definition_events.back()));
       // Get a temporary hold just so we can fish out a shared_ptr to the
       // TrackedDeviceBuffer. It's ok to drop the hold before return the
       // buffers, because the invariants of this class ensure that the buffer
diff --git a/third_party/xla/xla/pjrt/pjrt_stream_executor_client.cc b/third_party/xla/xla/pjrt/pjrt_stream_executor_client.cc
index efc10c37ce5d5b..59bb40d776a86e 100644
--- a/third_party/xla/xla/pjrt/pjrt_stream_executor_client.cc
+++ b/third_party/xla/xla/pjrt/pjrt_stream_executor_client.cc
@@ -414,24 +414,42 @@ void RecordUsage(PjRtStreamExecutorBuffer::ScopedHold device_buffer,
                                  retain_buffer_until_completion);
 }
 
-// Allocates the device buffers for a buffer that will be used as the
-// destination of a copy, either from the host or another device. copy_stream
-// may be nullptr, e.g., when allocating a buffer for a cross-host copy. If the
-// buffer is a tuple then the tuple tables are allocated, and all necessary
-// synchronization for them is dealt with, before the buffer is returned.
-//
-// It is safe to delete the returned PjRtBuffer without further
-// synchronization if an error occurs before the buffer is used.
-//
-// The caller may optionally provide a definition event to be recorded in
-// the buffer.
-// TODO(phawkins): replace on_host_shape here with on_device_shape.
+// Adds necessary synchronization after a copy has been enqueued to a buffer.
+// definition_event was added when the buffer was allocated, but has not yet
+// had an event recorded.
+absl::Status AddDestinationBufferSynchronization(
+    LocalDeviceState* local_device,
+    PjRtStreamExecutorBuffer::ScopedHold device_buffer,
+    std::shared_ptr<BufferSequencingEvent> definition_event,
+    se::Stream* copy_stream) {
+  absl::StatusOr<EventPool::Handle> event_or =
+      local_device->event_pool().ThenAllocateAndRecordEvent(copy_stream);
+  if (!event_or.ok()) {
+    StallStreamOnError(local_device, copy_stream);
+    return event_or.status();
+  }
+  definition_event->SetSequencingEvent(std::move(event_or).value(),
+                                       copy_stream);
+  // prefer_to_retain_reference=false means don't retain a memory reference
+  // until the transfer is complete when using the ComputeSynchronized
+  // allocation model. This is a heuristic because in the common case
+  // destination buffers will be used on the compute stream and therefore don't
+  // require any synchronization before being freed. If the buffer is allocated
+  // and never used, the free will take longer and this is assumed to be ok.
+  RecordUsage(std::move(device_buffer), local_device, local_device,
+              definition_event, copy_stream,
+              /*prefer_to_retain_reference=*/false);
+  return OkStatus();
+}
+
+}  // namespace
+
 absl::StatusOr<std::unique_ptr<PjRtStreamExecutorBuffer>>
 AllocateDestinationBuffer(
     const Shape& on_host_shape, PjRtDevice* device,
     LocalDeviceState* local_device, se::Stream* copy_stream,
     bool is_uninitialized_create, PjRtStreamExecutorClient* client,
-    std::shared_ptr<BufferSequencingEvent> definition_event = nullptr) {
+    std::shared_ptr<BufferSequencingEvent> definition_event) {
   if (on_host_shape.IsTuple() && on_host_shape.tuple_shapes_size() == 0) {
     return InvalidArgument("Can't make a buffer from an empty tuple");
   }
@@ -547,36 +565,6 @@ AllocateDestinationBuffer(
   return py_buffer;
 }
 
-// Adds necessary synchronization after a copy has been enqueued to a buffer.
-// definition_event was added when the buffer was allocated, but has not yet
-// had an event recorded.
-absl::Status AddDestinationBufferSynchronization(
-    LocalDeviceState* local_device,
-    PjRtStreamExecutorBuffer::ScopedHold device_buffer,
-    std::shared_ptr<BufferSequencingEvent> definition_event,
-    se::Stream* copy_stream) {
-  absl::StatusOr<EventPool::Handle> event_or =
-      local_device->event_pool().ThenAllocateAndRecordEvent(copy_stream);
-  if (!event_or.ok()) {
-    StallStreamOnError(local_device, copy_stream);
-    return event_or.status();
-  }
-  definition_event->SetSequencingEvent(std::move(event_or).value(),
-                                       copy_stream);
-  // prefer_to_retain_reference=false means don't retain a memory reference
-  // until the transfer is complete when using the ComputeSynchronized
-  // allocation model. This is a heuristic because in the common case
-  // destination buffers will be used on the compute stream and therefore don't
-  // require any synchronization before being freed. If the buffer is allocated
-  // and never used, the free will take longer and this is assumed to be ok.
-  RecordUsage(std::move(device_buffer), local_device, local_device,
-              definition_event, copy_stream,
-              /*prefer_to_retain_reference=*/false);
-  return OkStatus();
-}
-
-}  // namespace
-
 PjRtStreamExecutorBuffer::ScopedHold::~ScopedHold() {
   if (ok()) {
     parent_->DropHold(type_, buffer().get());
@@ -1573,18 +1561,12 @@ void PjRtStreamExecutorBuffer::Delete() {
   VLOG(1) << "PjRtStreamExecutorBuffer::Delete";
 
   // When wait_for_reads_to_complete is false, Release should never fail.
-  absl::StatusOr<std::shared_ptr<TrackedDeviceBuffer>> tracked_device_buffer =
-      Release(/*wait_for_operations_to_complete=*/false);
-
+  //
   // The only usage events that
   // Release(/*wait_for_operations_to_complete=*/false) doesn't wait for are
-  // events defined on the compute stream. So we schedule the deallocation on
-  // the compute stream so that there wouldn't be use-after-free usages on the
-  // compute stream.
-  TF_CHECK_OK(tracked_device_buffer.status());
-  TF_CHECK_OK(device_->local_device_state()->ThenRelease(
-      device_->local_device_state()->compute_stream(),
-      std::move(tracked_device_buffer.value())));
+  // events defined on the compute stream. All streams other than the compute
+  // stream are expected to WaitFor compute stream before any write operations.
+  TF_CHECK_OK(Release(/*wait_for_operations_to_complete=*/false).status());
 }
 
 bool PjRtStreamExecutorBuffer::IsDeleted() {
diff --git a/third_party/xla/xla/pjrt/pjrt_stream_executor_client.h b/third_party/xla/xla/pjrt/pjrt_stream_executor_client.h
index c2f55e743cf910..cac6e2691f3b6c 100644
--- a/third_party/xla/xla/pjrt/pjrt_stream_executor_client.h
+++ b/third_party/xla/xla/pjrt/pjrt_stream_executor_client.h
@@ -876,6 +876,25 @@ class PjRtStreamExecutorBuffer : public PjRtBuffer {
   PjRtFuture<>::Promise definition_promise_ ABSL_GUARDED_BY(mu_);
 };
 
+// Allocates the device buffers for a buffer that will be used as the
+// destination of a copy, either from the host or another device. copy_stream
+// may be nullptr, e.g., when allocating a buffer for a cross-host copy. If the
+// buffer is a tuple then the tuple tables are allocated, and all necessary
+// synchronization for them is dealt with, before the buffer is returned.
+//
+// It is safe to delete the returned PjRtBuffer without further
+// synchronization if an error occurs before the buffer is used.
+//
+// The caller may optionally provide a definition event to be recorded in
+// the buffer.
+// TODO(phawkins): replace on_host_shape here with on_device_shape.
+absl::StatusOr<std::unique_ptr<PjRtStreamExecutorBuffer>>
+AllocateDestinationBuffer(
+    const Shape& on_host_shape, PjRtDevice* device,
+    LocalDeviceState* local_device, se::Stream* copy_stream,
+    bool is_uninitialized_create, PjRtStreamExecutorClient* client,
+    std::shared_ptr<BufferSequencingEvent> definition_event = nullptr);
+
 // Wraps one or more XLA LocalExecutables (one per partition, as specified by
 // the build options).
 class PjRtStreamExecutorLoadedExecutable : public PjRtLoadedExecutable {

From e817f66475dfc8bd1e38446a1695d0b1eb680e38 Mon Sep 17 00:00:00 2001
From: Kyle Lucke <klucke@google.com>
Date: Mon, 20 May 2024 11:01:22 -0700
Subject: [PATCH 420/478] Remove the use of xla::OkStatus now that it's just an
 alias to absl::OkStatus.

PiperOrigin-RevId: 635509418
---
 third_party/xla/xla/python/ifrt/shape.cc | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/third_party/xla/xla/python/ifrt/shape.cc b/third_party/xla/xla/python/ifrt/shape.cc
index 0d6fb7c1963cbc..fb66cbe464cf87 100644
--- a/third_party/xla/xla/python/ifrt/shape.cc
+++ b/third_party/xla/xla/python/ifrt/shape.cc
@@ -22,11 +22,11 @@ limitations under the License.
 #include <variant>
 
 #include "absl/container/inlined_vector.h"
+#include "absl/status/status.h"
 #include "absl/strings/str_cat.h"
 #include "absl/strings/str_join.h"
 #include "absl/strings/string_view.h"
 #include "xla/python/ifrt/shape.pb.h"
-#include "xla/status.h"
 #include "xla/statusor.h"
 #include "xla/util.h"
 #include "tsl/platform/errors.h"
@@ -111,7 +111,7 @@ absl::StatusOr<DynamicShape> DynamicShape::Create(Shape shape,
               return InvalidArgument(
                   "Shape and tag must have the same number of dimensions.");
             }
-            return xla::OkStatus();
+            return absl::OkStatus();
           },
       },
       tag));

From 777c11afce7d6df16a0dfb00b258251d870d900b Mon Sep 17 00:00:00 2001
From: Kyle Lucke <klucke@google.com>
Date: Mon, 20 May 2024 11:04:16 -0700
Subject: [PATCH 421/478] Use absl::Status instead of xla::Status now that
 they're identical.

PiperOrigin-RevId: 635510656
---
 .../canonicalize_all_gather_for_cse_test.cc   |  3 +-
 .../xla/service/spmd/convolution_handler.cc   |  2 +-
 .../xla/service/spmd/custom_call_handler.cc   |  7 +-
 .../xla/xla/service/spmd/dot_handler.cc       | 10 +-
 .../xla/xla/service/spmd/fft_handler.cc       |  2 +-
 .../service/spmd/gather_scatter_handler.cc    |  4 +-
 .../xla/xla/service/spmd/spmd_partitioner.cc  | 73 +++++++-------
 .../xla/xla/service/spmd/spmd_partitioner.h   | 98 +++++++++----------
 .../spmd/stateful_rng_spmd_partitioner.cc     |  4 +-
 .../spmd/stateful_rng_spmd_partitioner.h      |  4 +-
 .../spmd/whole_graph_manual_pass_test.cc      |  3 +-
 .../xla/stream_executor/tpu/tpu_executable.h  |  2 +-
 .../tpu/tpu_executable_interface.cc           |  6 +-
 .../tpu/tpu_executable_interface.h            |  2 +-
 .../xla/stream_executor/tpu/tpu_executor.cc   | 43 ++++----
 .../xla/stream_executor/tpu/tpu_executor.h    |  6 +-
 .../tpu/tpu_executor_interface.h              |  4 +-
 .../stream_executor/tpu/tpu_node_context.cc   |  4 +-
 .../stream_executor/tpu/tpu_node_context.h    |  6 +-
 .../xla/stream_executor/tpu/tpu_platform.cc   |  7 +-
 .../xla/stream_executor/tpu/tpu_platform.h    | 11 +--
 .../stream_executor/tpu/tpu_profiler_c_api.h  |  2 +-
 .../tpu/tpu_transfer_manager.cc               |  5 +-
 .../tpu/tpu_transfer_manager.h                |  5 +-
 .../tpu/tpu_transfer_manager_interface.h      |  4 +-
 25 files changed, 151 insertions(+), 166 deletions(-)

diff --git a/third_party/xla/xla/service/spmd/canonicalize_all_gather_for_cse_test.cc b/third_party/xla/xla/service/spmd/canonicalize_all_gather_for_cse_test.cc
index b7e3b2a8c7a8b1..bcba61a080748d 100644
--- a/third_party/xla/xla/service/spmd/canonicalize_all_gather_for_cse_test.cc
+++ b/third_party/xla/xla/service/spmd/canonicalize_all_gather_for_cse_test.cc
@@ -43,7 +43,8 @@ class AllGatherCanonicalizeTest : public HloTestBase {
     TF_RETURN_IF_ERROR(pipeline.Run(module.get()).status());
     return absl::StatusOr<std::unique_ptr<HloModule>>(std::move(module));
   }
-  Status RunPassOnModule(HloModule* module, int64_t distance_threshold = 100) {
+  absl::Status RunPassOnModule(HloModule* module,
+                               int64_t distance_threshold = 100) {
     HloPassPipeline pipeline("all-gather-cse");
     pipeline.AddPass<CanonicalizeAllGatherForCSE>();
     TF_RETURN_IF_ERROR(pipeline.Run(module).status());
diff --git a/third_party/xla/xla/service/spmd/convolution_handler.cc b/third_party/xla/xla/service/spmd/convolution_handler.cc
index 2a88430331db97..84b9d65ff11177 100644
--- a/third_party/xla/xla/service/spmd/convolution_handler.cc
+++ b/third_party/xla/xla/service/spmd/convolution_handler.cc
@@ -1017,7 +1017,7 @@ absl::StatusOr<HloInstruction*> PartitionConvolution(
   return nullptr;
 }
 
-Status SpmdPartitioningVisitor::HandleConvolution(HloInstruction* hlo) {
+absl::Status SpmdPartitioningVisitor::HandleConvolution(HloInstruction* hlo) {
   if (hlo->sharding().HasUniqueDevice()) {
     return DefaultAction(hlo);
   }
diff --git a/third_party/xla/xla/service/spmd/custom_call_handler.cc b/third_party/xla/xla/service/spmd/custom_call_handler.cc
index 88052735952263..919ced6bc29837 100644
--- a/third_party/xla/xla/service/spmd/custom_call_handler.cc
+++ b/third_party/xla/xla/service/spmd/custom_call_handler.cc
@@ -84,7 +84,8 @@ constexpr char kSPMDOpRotateRight[] = "_SPMDInternalOp_RotateRight";
 
 }  // namespace
 
-Status SpmdPartitioningVisitor::HandleCustomCallTopK(HloInstruction* hlo) {
+absl::Status SpmdPartitioningVisitor::HandleCustomCallTopK(
+    HloInstruction* hlo) {
   if (!hlo->operand(0)->has_sharding()) {
     return DefaultAction(hlo);
   }
@@ -254,7 +255,7 @@ Status SpmdPartitioningVisitor::HandleCustomCallTopK(HloInstruction* hlo) {
   return OkStatus();
 }
 
-Status SpmdPartitioningVisitor::HandleCustomCallSPMDInternal_RotateRight(
+absl::Status SpmdPartitioningVisitor::HandleCustomCallSPMDInternal_RotateRight(
     HloInstruction* hlo) {
   TF_ASSIGN_OR_RETURN(auto attrs, ParseOpaqueAsAttributes(hlo));
   auto dim_it = attrs.find("dimension");
@@ -390,7 +391,7 @@ std::unique_ptr<HloInstruction> CreateCustomCallSPMDInternal_RotateRight(
                                           kSPMDOpRotateRight, opaque);
 }
 
-Status SpmdPartitioningVisitor::HandleCustomCall(HloInstruction* hlo) {
+absl::Status SpmdPartitioningVisitor::HandleCustomCall(HloInstruction* hlo) {
   if (auto* partitioner = GetCustomCallPartitioner(hlo->custom_call_target())) {
     return partitioner->Partition(this, hlo);
   }
diff --git a/third_party/xla/xla/service/spmd/dot_handler.cc b/third_party/xla/xla/service/spmd/dot_handler.cc
index 47909ff72a5801..d19f074b6a93ef 100644
--- a/third_party/xla/xla/service/spmd/dot_handler.cc
+++ b/third_party/xla/xla/service/spmd/dot_handler.cc
@@ -64,7 +64,7 @@ namespace {
 using hlo_sharding_util::GroupedSharding;
 }  // namespace
 
-Status SpmdPartitioningVisitor::HandleDot(HloInstruction* hlo) {
+absl::Status SpmdPartitioningVisitor::HandleDot(HloInstruction* hlo) {
   DotConvDimsMapping mapping;
   const auto& dnums = hlo->dot_dimension_numbers();
   int64_t next_output_dim = 0;
@@ -4264,7 +4264,7 @@ absl::StatusOr<HloInstruction*> PartitionDot(
 
 }  // namespace
 
-Status SpmdPartitioningVisitor::HandleDotHelper(
+absl::Status SpmdPartitioningVisitor::HandleDotHelper(
     HloInstruction* hlo, const DotConvDimsMapping& dims_mapping,
     absl::FunctionRef<absl::StatusOr<HloInstruction*>(
         HloInstruction*, HloInstruction*, SpmdBuilder*,
@@ -4362,7 +4362,7 @@ FindInputNodesIfOnlyDependOnSmallOperands(HloInstruction* hlo) {
 //
 // Later optimization passes (TpuPadSliceMover) will merge the dynamic slice
 // with the input nodes.
-Status SinkInputNodesIntoWindowedDotGeneralLoopOnContractingDimensions(
+absl::Status SinkInputNodesIntoWindowedDotGeneralLoopOnContractingDimensions(
     HloInstruction* loop, int64_t non_windowed_operand_index) {
   auto input_tuple = loop->mutable_operand(0);
   auto old_operand = input_tuple->mutable_operand(non_windowed_operand_index);
@@ -4496,7 +4496,7 @@ bool CheckOperandsRecursive(
 //
 // Later optimization passes (TpuPadSliceMover) will merge the dynamic slice
 // with the input nodes (broadcast).
-Status MoveUsersIntoWindowedDotGeneralLoopOnNonContractingDimensions(
+absl::Status MoveUsersIntoWindowedDotGeneralLoopOnNonContractingDimensions(
     HloInstruction* loop, const SpmdPartitionerOptions& options) {
   CHECK_EQ(loop->user_count(), 1);
   // There should be a single direct user of the while loop, which is the
@@ -4954,7 +4954,7 @@ Status MoveUsersIntoWindowedDotGeneralLoopOnNonContractingDimensions(
 
 }  // namespace
 
-Status SpmdPartitioningVisitor::DoCodeMotionForWindowedDotGeneralLoops(
+absl::Status SpmdPartitioningVisitor::DoCodeMotionForWindowedDotGeneralLoops(
     HloComputation* computation, const SpmdPartitionerOptions& options) {
   for (auto& loop : windowed_dot_general_loops_) {
     if (loop.windowed_in_contracting_dims || loop.windowed_in_batch_dims ||
diff --git a/third_party/xla/xla/service/spmd/fft_handler.cc b/third_party/xla/xla/service/spmd/fft_handler.cc
index 76cf9abedd8ea6..25984d96c981ea 100644
--- a/third_party/xla/xla/service/spmd/fft_handler.cc
+++ b/third_party/xla/xla/service/spmd/fft_handler.cc
@@ -346,7 +346,7 @@ HloInstruction* SliceValidData(HloInstruction* hlo, const Shape& target_shape,
 }  // namespace
 
 // Distributed FFT using the algorithm described in go/tpu-spmd-fft.
-Status SpmdPartitioningVisitor::HandleFft(HloInstruction* hlo) {
+absl::Status SpmdPartitioningVisitor::HandleFft(HloInstruction* hlo) {
   if (hlo->operand(0)->shape().rank() < 3 || hlo->fft_type() != FftType::FFT) {
     return DefaultAction(hlo);
   }
diff --git a/third_party/xla/xla/service/spmd/gather_scatter_handler.cc b/third_party/xla/xla/service/spmd/gather_scatter_handler.cc
index e8b97b25d4d175..c088e041cec891 100644
--- a/third_party/xla/xla/service/spmd/gather_scatter_handler.cc
+++ b/third_party/xla/xla/service/spmd/gather_scatter_handler.cc
@@ -825,7 +825,7 @@ absl::StatusOr<HloInstruction*> PartitionGather(
 
 }  // namespace
 
-Status SpmdPartitioningVisitor::HandleGather(HloInstruction* hlo) {
+absl::Status SpmdPartitioningVisitor::HandleGather(HloInstruction* hlo) {
   if (hlo->sharding().HasUniqueDevice()) {
     return DefaultAction(hlo);
   }
@@ -1613,7 +1613,7 @@ absl::StatusOr<HloInstruction*> PartitionScatter(
 
 }  // namespace
 
-Status SpmdPartitioningVisitor::HandleScatter(HloInstruction* hlo) {
+absl::Status SpmdPartitioningVisitor::HandleScatter(HloInstruction* hlo) {
   if (hlo->sharding().HasUniqueDevice()) {
     return DefaultAction(hlo);
   }
diff --git a/third_party/xla/xla/service/spmd/spmd_partitioner.cc b/third_party/xla/xla/service/spmd/spmd_partitioner.cc
index 1579d73344efd7..9fe9eb43b3293a 100644
--- a/third_party/xla/xla/service/spmd/spmd_partitioner.cc
+++ b/third_party/xla/xla/service/spmd/spmd_partitioner.cc
@@ -209,7 +209,7 @@ bool ShouldKeepSharding(const HloInstruction* hlo) {
 
 // Clears all sharding attributes from instructions in the module. This must be
 // called only after all SPMD transformation is complete.
-Status ClearShardingAttributes(
+absl::Status ClearShardingAttributes(
     HloModule* module,
     const absl::flat_hash_set<absl::string_view>& execution_threads) {
   for (HloComputation* computation : module->computations(execution_threads)) {
@@ -2290,7 +2290,7 @@ Status SpmdPartitioningVisitor::HandleCall(HloInstruction* hlo) {
   return OkStatus();
 }
 
-Status SpmdPartitioningVisitor::DefaultAction(HloInstruction* hlo) {
+absl::Status SpmdPartitioningVisitor::DefaultAction(HloInstruction* hlo) {
   if (hlo->HasSideEffect() && !hlo->sharding().HasUniqueDevice()) {
     return Unimplemented("Side-effect ops cannot be replicated: %s",
                          hlo->ToString());
@@ -2335,7 +2335,7 @@ Status SpmdPartitioningVisitor::DefaultAction(HloInstruction* hlo) {
   return OkStatus();
 }
 
-Status SpmdPartitioningVisitor::Preprocess(HloInstruction* hlo) {
+absl::Status SpmdPartitioningVisitor::Preprocess(HloInstruction* hlo) {
   visiting_hlo_ = hlo;
   b_.set_visiting_hlo(hlo);
   // Temporarily replace manual sharding to one-device sharding so that the
@@ -2501,7 +2501,7 @@ Status SpmdPartitioningVisitor::Preprocess(HloInstruction* hlo) {
   return OkStatus();
 }
 
-Status SpmdPartitioningVisitor::Postprocess(HloInstruction* hlo) {
+absl::Status SpmdPartitioningVisitor::Postprocess(HloInstruction* hlo) {
   logger_->RegisterLogEntry(hlo, b_.derived_instructions(hlo));
   visiting_hlo_ = nullptr;
   b_.set_visiting_hlo(nullptr);
@@ -2540,7 +2540,7 @@ Status SpmdPartitioningVisitor::Postprocess(HloInstruction* hlo) {
   return OkStatus();
 }
 
-Status SpmdPartitioningVisitor::HandleElementwise(HloInstruction* hlo) {
+absl::Status SpmdPartitioningVisitor::HandleElementwise(HloInstruction* hlo) {
   std::vector<HloInstruction*> new_operands;
   for (HloInstruction* operand : hlo->operands()) {
     new_operands.push_back(
@@ -2553,7 +2553,7 @@ Status SpmdPartitioningVisitor::HandleElementwise(HloInstruction* hlo) {
   return OkStatus();
 }
 
-Status SpmdPartitioningVisitor::HandleConcatenate(HloInstruction* hlo) {
+absl::Status SpmdPartitioningVisitor::HandleConcatenate(HloInstruction* hlo) {
   const HloSharding& sharding = hlo->sharding();
   if (sharding.IsTileMaximal()) {
     return DefaultAction(hlo);
@@ -2650,7 +2650,7 @@ Status SpmdPartitioningVisitor::HandleConcatenate(HloInstruction* hlo) {
   return OkStatus();
 }
 
-Status SpmdPartitioningVisitor::HandleSlice(HloInstruction* hlo) {
+absl::Status SpmdPartitioningVisitor::HandleSlice(HloInstruction* hlo) {
   const HloSharding& sharding = hlo->sharding();
   if (sharding.IsTileMaximal()) {
     return DefaultAction(hlo);
@@ -2679,7 +2679,7 @@ Status SpmdPartitioningVisitor::HandleSlice(HloInstruction* hlo) {
   return OkStatus();
 }
 
-Status SpmdPartitioningVisitor::HandleSort(HloInstruction* hlo) {
+absl::Status SpmdPartitioningVisitor::HandleSort(HloInstruction* hlo) {
   HloSharding sharding = hlo->sharding();
   int64_t input_count = 1;
   if (hlo->shape().IsTuple()) {
@@ -2938,7 +2938,7 @@ Status SpmdPartitioningVisitor::HandleSort(HloInstruction* hlo) {
   return OkStatus();
 }
 
-Status SpmdPartitioningVisitor::HandleTranspose(HloInstruction* hlo) {
+absl::Status SpmdPartitioningVisitor::HandleTranspose(HloInstruction* hlo) {
   const HloSharding& sharding = hlo->sharding();
   if (sharding.IsTileMaximal()) {
     return DefaultAction(hlo);
@@ -2961,7 +2961,7 @@ Status SpmdPartitioningVisitor::HandleTranspose(HloInstruction* hlo) {
   return OkStatus();
 }
 
-Status SpmdPartitioningVisitor::HandleReshape(HloInstruction* hlo) {
+absl::Status SpmdPartitioningVisitor::HandleReshape(HloInstruction* hlo) {
   const HloSharding& sharding = hlo->sharding();
   if (sharding.IsTileMaximal()) {
     return DefaultAction(hlo);
@@ -3248,7 +3248,7 @@ Status SpmdPartitioningVisitor::HandleReshape(HloInstruction* hlo) {
   return OkStatus();
 }
 
-Status SpmdPartitioningVisitor::HandleIota(HloInstruction* hlo) {
+absl::Status SpmdPartitioningVisitor::HandleIota(HloInstruction* hlo) {
   const HloSharding& sharding = hlo->sharding();
   if (sharding.IsTileMaximal()) {
     return DefaultAction(hlo);
@@ -3283,7 +3283,8 @@ Status SpmdPartitioningVisitor::HandleIota(HloInstruction* hlo) {
   return OkStatus();
 }
 
-Status SpmdPartitioningVisitor::HandleSingleDevice(const HloInstruction* hlo) {
+absl::Status SpmdPartitioningVisitor::HandleSingleDevice(
+    const HloInstruction* hlo) {
   TF_RET_CHECK(hlo->sharding().HasUniqueDevice());
   int64_t device = hlo->sharding().GetUniqueDevice();
   const HloSharding sharding = HloSharding::AssignDevice(device);
@@ -3339,7 +3340,7 @@ Status SpmdPartitioningVisitor::HandleSingleDevice(const HloInstruction* hlo) {
   return OkStatus();
 }
 
-Status SpmdPartitioningVisitor::HandleAllReduce(HloInstruction* hlo) {
+absl::Status SpmdPartitioningVisitor::HandleAllReduce(HloInstruction* hlo) {
   if (hlo->IsCrossReplicaAllReduce() && hlo->operand_count() == 1) {
     return HandleElementwise(hlo);
   }
@@ -3389,7 +3390,7 @@ Status SpmdPartitioningVisitor::HandleAllReduce(HloInstruction* hlo) {
   return DefaultAction(hlo);
 }
 
-Status SpmdPartitioningVisitor::HandleBroadcast(HloInstruction* hlo) {
+absl::Status SpmdPartitioningVisitor::HandleBroadcast(HloInstruction* hlo) {
   if (hlo->sharding().IsTileMaximal()) {
     return DefaultAction(hlo);
   }
@@ -3416,7 +3417,7 @@ Status SpmdPartitioningVisitor::HandleBroadcast(HloInstruction* hlo) {
   return OkStatus();
 }
 
-Status SpmdPartitioningVisitor::HandleConstant(HloInstruction* hlo) {
+absl::Status SpmdPartitioningVisitor::HandleConstant(HloInstruction* hlo) {
   const Literal& literal = hlo->literal();
   if (literal.shape().IsTuple() ||
       (!hlo->sharding().IsTileMaximal() &&
@@ -3436,7 +3437,7 @@ Status SpmdPartitioningVisitor::HandleConstant(HloInstruction* hlo) {
   return OkStatus();
 }
 
-Status SpmdPartitioningVisitor::HandleDynamicSlice(HloInstruction* hlo) {
+absl::Status SpmdPartitioningVisitor::HandleDynamicSlice(HloInstruction* hlo) {
   if (hlo->sharding().IsTileMaximal()) {
     return DefaultAction(hlo);
   }
@@ -3473,7 +3474,8 @@ Status SpmdPartitioningVisitor::HandleDynamicSlice(HloInstruction* hlo) {
   return OkStatus();
 }
 
-Status SpmdPartitioningVisitor::HandleDynamicUpdateSlice(HloInstruction* hlo) {
+absl::Status SpmdPartitioningVisitor::HandleDynamicUpdateSlice(
+    HloInstruction* hlo) {
   if (hlo->sharding().IsTileMaximal()) {
     return DefaultAction(hlo);
   }
@@ -3679,7 +3681,8 @@ Status SpmdPartitioningVisitor::HandleDynamicUpdateSlice(HloInstruction* hlo) {
   return OkStatus();
 }
 
-Status SpmdPartitioningVisitor::HandleGetTupleElement(HloInstruction* hlo) {
+absl::Status SpmdPartitioningVisitor::HandleGetTupleElement(
+    HloInstruction* hlo) {
   if (hlo->sharding().IsManual()) {
     return DefaultAction(hlo);
   }
@@ -3698,7 +3701,7 @@ Status SpmdPartitioningVisitor::HandleGetTupleElement(HloInstruction* hlo) {
   return OkStatus();
 }
 
-Status SpmdPartitioningVisitor::HandleInfeed(HloInstruction* hlo) {
+absl::Status SpmdPartitioningVisitor::HandleInfeed(HloInstruction* hlo) {
   const Shape& shape = ShapeUtil::GetTupleElementShape(hlo->shape(), 0);
   auto token = GetPartitionedHlo(hlo->operand(0)).hlo();
   if (ShapeUtil::GetLeafCount(shape) == 0) {
@@ -3830,7 +3833,7 @@ Status SpmdPartitioningVisitor::HandleInfeed(HloInstruction* hlo) {
   return OkStatus();
 }
 
-Status SpmdPartitioningVisitor::HandlePad(HloInstruction* hlo) {
+absl::Status SpmdPartitioningVisitor::HandlePad(HloInstruction* hlo) {
   if (hlo->sharding().IsTileMaximal()) {
     return DefaultAction(hlo);
   }
@@ -3859,7 +3862,7 @@ Status SpmdPartitioningVisitor::HandlePad(HloInstruction* hlo) {
   return OkStatus();
 }
 
-Status SpmdPartitioningVisitor::HandleParameter(HloInstruction* hlo) {
+absl::Status SpmdPartitioningVisitor::HandleParameter(HloInstruction* hlo) {
   SetPartitionedHlo(hlo, [&]() {
     auto shard_shape = MakePartitionedShape(hlo->shape(), hlo->sharding());
     auto new_param = b_.AddInstruction(HloInstruction::CreateParameter(
@@ -3873,7 +3876,7 @@ Status SpmdPartitioningVisitor::HandleParameter(HloInstruction* hlo) {
   return OkStatus();
 }
 
-Status SpmdPartitioningVisitor::HandleReduce(HloInstruction* hlo) {
+absl::Status SpmdPartitioningVisitor::HandleReduce(HloInstruction* hlo) {
   int64_t input_count = 1;
   if (hlo->shape().IsTuple()) {
     input_count = hlo->shape().tuple_shapes_size();
@@ -4011,7 +4014,7 @@ Status SpmdPartitioningVisitor::HandleReduce(HloInstruction* hlo) {
   return OkStatus();
 }
 
-Status SpmdPartitioningVisitor::HandleReverse(HloInstruction* hlo) {
+absl::Status SpmdPartitioningVisitor::HandleReverse(HloInstruction* hlo) {
   auto reverse = Cast<HloReverseInstruction>(hlo);
   if (reverse->sharding().IsTileMaximal()) {
     return DefaultAction(hlo);
@@ -4031,7 +4034,7 @@ Status SpmdPartitioningVisitor::HandleReverse(HloInstruction* hlo) {
   return OkStatus();
 }
 
-Status SpmdPartitioningVisitor::HandleWhile(HloInstruction* hlo) {
+absl::Status SpmdPartitioningVisitor::HandleWhile(HloInstruction* hlo) {
   const HloSharding& sharding = hlo->sharding();
 
   // Shardings for the body parameter, body root, and cond parameter must be
@@ -4063,7 +4066,7 @@ Status SpmdPartitioningVisitor::HandleWhile(HloInstruction* hlo) {
   return OkStatus();
 }
 
-Status SpmdPartitioningVisitor::HandleConditional(HloInstruction* hlo) {
+absl::Status SpmdPartitioningVisitor::HandleConditional(HloInstruction* hlo) {
   std::vector<HloInstruction*> branch_args;
   for (int64_t i = 0; i < hlo->branch_count(); ++i) {
     HloComputation* computation = hlo->branch_computation(i);
@@ -4101,11 +4104,12 @@ Status SpmdPartitioningVisitor::HandleConditional(HloInstruction* hlo) {
   return OkStatus();
 }
 
-Status SpmdPartitioningVisitor::HandleOptimizationBarrier(HloInstruction* hlo) {
+absl::Status SpmdPartitioningVisitor::HandleOptimizationBarrier(
+    HloInstruction* hlo) {
   return HandleElementwise(hlo);
 }
 
-Status SpmdPartitioningVisitor::HandleOutfeed(HloInstruction* hlo) {
+absl::Status SpmdPartitioningVisitor::HandleOutfeed(HloInstruction* hlo) {
   if (hlo->sharding().HasUniqueDevice()) {
     return HandleSingleDevice(hlo);
   }
@@ -4287,7 +4291,7 @@ Status SpmdPartitioningVisitor::HandleOutfeed(HloInstruction* hlo) {
   return OkStatus();
 }
 
-Status SpmdPartitioningVisitor::HandleRng(HloInstruction* hlo) {
+absl::Status SpmdPartitioningVisitor::HandleRng(HloInstruction* hlo) {
   if (hlo->sharding().HasUniqueDevice()) {
     return HandleSingleDevice(hlo);
   }
@@ -4356,7 +4360,7 @@ Status SpmdPartitioningVisitor::HandleRng(HloInstruction* hlo) {
   return OkStatus();
 }
 
-Status SpmdPartitioningVisitor::HandleReduceWindow(HloInstruction* hlo) {
+absl::Status SpmdPartitioningVisitor::HandleReduceWindow(HloInstruction* hlo) {
   if (hlo->sharding().IsTileMaximal()) {
     return DefaultAction(hlo);
   }
@@ -4418,7 +4422,8 @@ Status SpmdPartitioningVisitor::HandleReduceWindow(HloInstruction* hlo) {
   return OkStatus();
 }
 
-Status SpmdPartitioningVisitor::HandleSelectAndScatter(HloInstruction* hlo) {
+absl::Status SpmdPartitioningVisitor::HandleSelectAndScatter(
+    HloInstruction* hlo) {
   if (hlo->sharding().IsTileMaximal()) {
     return DefaultAction(hlo);
   }
@@ -4639,7 +4644,7 @@ Status SpmdPartitioningVisitor::HandleSelectAndScatter(HloInstruction* hlo) {
   return OkStatus();
 }
 
-Status SpmdPartitioningVisitor::HandleTuple(HloInstruction* hlo) {
+absl::Status SpmdPartitioningVisitor::HandleTuple(HloInstruction* hlo) {
   std::vector<HloInstruction*> new_operands;
   for (int64_t i = 0; i < hlo->operand_count(); ++i) {
     new_operands.push_back(
@@ -4676,7 +4681,7 @@ absl::StatusOr<bool> SpmdPartitioningVisitor::DoPartition(
   return changed_;
 }
 
-Status SpmdPartitioningVisitor::HandlePartitionId(HloInstruction* hlo) {
+absl::Status SpmdPartitioningVisitor::HandlePartitionId(HloInstruction* hlo) {
   if (hlo->has_sharding() && hlo->sharding().IsManual()) {
     hlo->set_sharding(HloSharding::AssignDevice(0));
     return DefaultAction(hlo);
@@ -5204,7 +5209,7 @@ absl::StatusOr<bool> SpmdPartitioner::Run(
   return changed;
 }
 
-Status SpmdPartitioner::PreprocessSharding(
+absl::Status SpmdPartitioner::PreprocessSharding(
     HloModule* module,
     const absl::flat_hash_set<absl::string_view>& execution_threads) {
   for (HloComputation* computation : module->computations(execution_threads)) {
@@ -5272,7 +5277,7 @@ Status SpmdPartitioner::PreprocessSharding(
   return OkStatus();
 }
 
-Status SpmdPartitioner::PreprocessHlos(
+absl::Status SpmdPartitioner::PreprocessHlos(
     HloModule* module,
     const absl::flat_hash_set<absl::string_view>& execution_threads) {
   auto skip_copy_operands = [](HloInstruction* operand,
diff --git a/third_party/xla/xla/service/spmd/spmd_partitioner.h b/third_party/xla/xla/service/spmd/spmd_partitioner.h
index 836250bffb6ba8..700377f27ba1ba 100644
--- a/third_party/xla/xla/service/spmd/spmd_partitioner.h
+++ b/third_party/xla/xla/service/spmd/spmd_partitioner.h
@@ -326,7 +326,7 @@ class SpmdPartitioner : public HloModulePass {
 
   // Verifies that the sharding of instructions in the module are valid, and
   // also fill in missing sharding information.
-  virtual Status PreprocessSharding(
+  virtual absl::Status PreprocessSharding(
       HloModule* module,
       const absl::flat_hash_set<absl::string_view>& execution_threads);
 
@@ -347,7 +347,7 @@ class SpmdPartitioner : public HloModulePass {
   // Preprocesses the graph to simplify some communication patterns. E.g., merge
   // pad->slice into a single pad with potentially negative padding to avoid
   // multiple halo exchanges.
-  Status PreprocessHlos(
+  absl::Status PreprocessHlos(
       HloModule* module,
       const absl::flat_hash_set<absl::string_view>& execution_threads);
 
@@ -552,59 +552,59 @@ class SpmdPartitioningVisitor : public DfsHloVisitorWithDefault {
 
   SpmdPartitioningVisitor(const SpmdPartitioningVisitor& src);
 
-  Status DefaultAction(HloInstruction* hlo) override;
-  Status HandleAllReduce(HloInstruction* hlo) override;
-  Status HandleBroadcast(HloInstruction* hlo) override;
-  Status HandleCall(HloInstruction* hlo) override;
-  Status HandleConstant(HloInstruction* hlo) override;
-  Status HandleCustomCall(HloInstruction* hlo) override;
-  Status HandleDot(HloInstruction* hlo) override;
-  Status HandleDynamicSlice(HloInstruction* hlo) override;
-  Status HandleDynamicUpdateSlice(HloInstruction* hlo) override;
-  Status HandleFft(HloInstruction* hlo) override;
-  Status HandleGather(HloInstruction* hlo) override;
-  Status HandleGetTupleElement(HloInstruction* hlo) override;
-  Status HandleInfeed(HloInstruction* hlo) override;
-  Status HandleOptimizationBarrier(HloInstruction* hlo) override;
-  Status HandleOutfeed(HloInstruction* hlo) override;
-  Status HandlePad(HloInstruction* hlo) override;
-  Status HandleParameter(HloInstruction* hlo) override;
-  Status HandleReduce(HloInstruction* hlo) override;
-  Status HandleReverse(HloInstruction* hlo) override;
-  Status HandleWhile(HloInstruction* hlo) override;
-  Status HandleConditional(HloInstruction* hlo) override;
-  Status HandleReduceWindow(HloInstruction* hlo) override;
-  Status HandleSelectAndScatter(HloInstruction* hlo) override;
-  Status HandleTuple(HloInstruction* hlo) override;
-  Status HandleRng(HloInstruction* hlo) override;
-  Status HandleConvolution(HloInstruction* hlo) override;
-  Status HandleConcatenate(HloInstruction* hlo) override;
-  Status HandleScatter(HloInstruction* hlo) override;
-  Status HandleSlice(HloInstruction* hlo) override;
-  Status HandleSort(HloInstruction* hlo) override;
-  Status HandleTranspose(HloInstruction* hlo) override;
-  Status HandleReshape(HloInstruction* hlo) override;
-  Status HandleIota(HloInstruction* hlo) override;
-  Status HandlePartitionId(HloInstruction* hlo) override;
+  absl::Status DefaultAction(HloInstruction* hlo) override;
+  absl::Status HandleAllReduce(HloInstruction* hlo) override;
+  absl::Status HandleBroadcast(HloInstruction* hlo) override;
+  absl::Status HandleCall(HloInstruction* hlo) override;
+  absl::Status HandleConstant(HloInstruction* hlo) override;
+  absl::Status HandleCustomCall(HloInstruction* hlo) override;
+  absl::Status HandleDot(HloInstruction* hlo) override;
+  absl::Status HandleDynamicSlice(HloInstruction* hlo) override;
+  absl::Status HandleDynamicUpdateSlice(HloInstruction* hlo) override;
+  absl::Status HandleFft(HloInstruction* hlo) override;
+  absl::Status HandleGather(HloInstruction* hlo) override;
+  absl::Status HandleGetTupleElement(HloInstruction* hlo) override;
+  absl::Status HandleInfeed(HloInstruction* hlo) override;
+  absl::Status HandleOptimizationBarrier(HloInstruction* hlo) override;
+  absl::Status HandleOutfeed(HloInstruction* hlo) override;
+  absl::Status HandlePad(HloInstruction* hlo) override;
+  absl::Status HandleParameter(HloInstruction* hlo) override;
+  absl::Status HandleReduce(HloInstruction* hlo) override;
+  absl::Status HandleReverse(HloInstruction* hlo) override;
+  absl::Status HandleWhile(HloInstruction* hlo) override;
+  absl::Status HandleConditional(HloInstruction* hlo) override;
+  absl::Status HandleReduceWindow(HloInstruction* hlo) override;
+  absl::Status HandleSelectAndScatter(HloInstruction* hlo) override;
+  absl::Status HandleTuple(HloInstruction* hlo) override;
+  absl::Status HandleRng(HloInstruction* hlo) override;
+  absl::Status HandleConvolution(HloInstruction* hlo) override;
+  absl::Status HandleConcatenate(HloInstruction* hlo) override;
+  absl::Status HandleScatter(HloInstruction* hlo) override;
+  absl::Status HandleSlice(HloInstruction* hlo) override;
+  absl::Status HandleSort(HloInstruction* hlo) override;
+  absl::Status HandleTranspose(HloInstruction* hlo) override;
+  absl::Status HandleReshape(HloInstruction* hlo) override;
+  absl::Status HandleIota(HloInstruction* hlo) override;
+  absl::Status HandlePartitionId(HloInstruction* hlo) override;
 
   // Implementation of dot partitioning given DotGeneralDimsMapping.
-  Status HandleDotHelper(HloInstruction* hlo,
-                         const DotConvDimsMapping& dims_mapping,
-                         absl::FunctionRef<absl::StatusOr<HloInstruction*>(
-                             HloInstruction*, HloInstruction*, SpmdBuilder*,
-                             const Window& conv_window)>
-                             create_sharded_dot);
+  absl::Status HandleDotHelper(
+      HloInstruction* hlo, const DotConvDimsMapping& dims_mapping,
+      absl::FunctionRef<absl::StatusOr<HloInstruction*>(
+          HloInstruction*, HloInstruction*, SpmdBuilder*,
+          const Window& conv_window)>
+          create_sharded_dot);
 
   // Common handle for elementwise HLOs.
-  Status HandleElementwise(HloInstruction* hlo);
+  absl::Status HandleElementwise(HloInstruction* hlo);
 
   // Common handle for HLOs that runs on a single device.
-  Status HandleSingleDevice(const HloInstruction* hlo);
+  absl::Status HandleSingleDevice(const HloInstruction* hlo);
 
   // CustomCall handlers per call target.
-  Status HandleCustomCallTopK(HloInstruction* hlo);
+  absl::Status HandleCustomCallTopK(HloInstruction* hlo);
   // Convenient custom ops defined by the partitioner itself.
-  Status HandleCustomCallSPMDInternal_RotateRight(HloInstruction* hlo);
+  absl::Status HandleCustomCallSPMDInternal_RotateRight(HloInstruction* hlo);
 
   virtual std::unique_ptr<SpmdPartitioningVisitor> Clone() const;
 
@@ -692,13 +692,13 @@ class SpmdPartitioningVisitor : public DfsHloVisitorWithDefault {
   };
 
  protected:
-  Status Preprocess(HloInstruction* hlo) override;
-  Status Postprocess(HloInstruction* hlo) override;
+  absl::Status Preprocess(HloInstruction* hlo) override;
+  absl::Status Postprocess(HloInstruction* hlo) override;
 
   // Performs code motion for windowed dot-general loops in
   // windowed_dot_general_loops_. Invoked after the visitor finishes traversing
   // the graph.
-  Status DoCodeMotionForWindowedDotGeneralLoops(
+  absl::Status DoCodeMotionForWindowedDotGeneralLoops(
       HloComputation* computation, const SpmdPartitionerOptions& options);
 
   bool changed_;
diff --git a/third_party/xla/xla/service/spmd/stateful_rng_spmd_partitioner.cc b/third_party/xla/xla/service/spmd/stateful_rng_spmd_partitioner.cc
index 2d575fa6a627b6..f267d737d8b087 100644
--- a/third_party/xla/xla/service/spmd/stateful_rng_spmd_partitioner.cc
+++ b/third_party/xla/xla/service/spmd/stateful_rng_spmd_partitioner.cc
@@ -24,7 +24,7 @@ limitations under the License.
 namespace xla {
 namespace spmd {
 
-Status StatefulRngSpmdPartitioningVisitor::HandleRngGetAndUpdateState(
+absl::Status StatefulRngSpmdPartitioningVisitor::HandleRngGetAndUpdateState(
     HloInstruction* hlo) {
   if (hlo->sharding().HasUniqueDevice()) {
     return HandleSingleDevice(hlo);
@@ -57,7 +57,7 @@ StatefulRngSpmdPartitioner::CreateVisitor(
       next_channel_id, logger, std::move(options), this, call_graph);
 }
 
-Status StatefulRngSpmdPartitioner::PreprocessSharding(
+absl::Status StatefulRngSpmdPartitioner::PreprocessSharding(
     HloModule* module,
     const absl::flat_hash_set<absl::string_view>& execution_threads) {
   // For rng-get-and-update-status with no sharding, set sharding to be
diff --git a/third_party/xla/xla/service/spmd/stateful_rng_spmd_partitioner.h b/third_party/xla/xla/service/spmd/stateful_rng_spmd_partitioner.h
index dd3e2635076a8a..1fad5d1621ffef 100644
--- a/third_party/xla/xla/service/spmd/stateful_rng_spmd_partitioner.h
+++ b/third_party/xla/xla/service/spmd/stateful_rng_spmd_partitioner.h
@@ -40,7 +40,7 @@ class StatefulRngSpmdPartitioningVisitor
                                       collective_ops_creator, next_channel_id,
                                       logger, std::move(options), partitioner,
                                       call_graph) {}
-  Status HandleRngGetAndUpdateState(HloInstruction* hlo) override;
+  absl::Status HandleRngGetAndUpdateState(HloInstruction* hlo) override;
 };
 
 class StatefulRngSpmdPartitioner : public spmd::SpmdPartitioner {
@@ -66,7 +66,7 @@ class StatefulRngSpmdPartitioner : public spmd::SpmdPartitioner {
       spmd::SpmdPartitionerOptions options,
       const CallGraph& call_graph) override;
 
-  Status PreprocessSharding(
+  absl::Status PreprocessSharding(
       HloModule* module,
       const absl::flat_hash_set<absl::string_view>& execution_threads) override;
   bool CanSideEffectingHaveReplicatedSharding(
diff --git a/third_party/xla/xla/service/spmd/whole_graph_manual_pass_test.cc b/third_party/xla/xla/service/spmd/whole_graph_manual_pass_test.cc
index 9270589e8d9cf3..0720e5b7da647b 100644
--- a/third_party/xla/xla/service/spmd/whole_graph_manual_pass_test.cc
+++ b/third_party/xla/xla/service/spmd/whole_graph_manual_pass_test.cc
@@ -45,7 +45,8 @@ class WholeGraphManualPassTest : public HloTestBase {
     TF_RETURN_IF_ERROR(pipeline.Run(module.get()).status());
     return absl::StatusOr<std::unique_ptr<HloModule>>(std::move(module));
   }
-  Status RunPassOnModule(HloModule* module, int64_t distance_threshold = 100) {
+  absl::Status RunPassOnModule(HloModule* module,
+                               int64_t distance_threshold = 100) {
     HloPassPipeline pipeline("all-gather-cse");
     pipeline.AddPass<WholeGraphManualPass>();
     TF_RETURN_IF_ERROR(pipeline.Run(module).status());
diff --git a/third_party/xla/xla/stream_executor/tpu/tpu_executable.h b/third_party/xla/xla/stream_executor/tpu/tpu_executable.h
index 2c5db60ce71c9b..bee80cb8489a59 100644
--- a/third_party/xla/xla/stream_executor/tpu/tpu_executable.h
+++ b/third_party/xla/xla/stream_executor/tpu/tpu_executable.h
@@ -62,7 +62,7 @@ class TpuExecutable : public xla::TpuExecutableInterface {
       absl::string_view serialized);
 
  private:
-  Status LoadProgramAndEnqueueToStream(
+  absl::Status LoadProgramAndEnqueueToStream(
       const ServiceExecutableRunOptions& run_options,
       absl::Span<const stream_executor::DeviceMemoryBase> arguments,
       stream_executor::DeviceMemoryBase result,
diff --git a/third_party/xla/xla/stream_executor/tpu/tpu_executable_interface.cc b/third_party/xla/xla/stream_executor/tpu/tpu_executable_interface.cc
index 07f6de17cc7d09..6a5435d17fbffb 100644
--- a/third_party/xla/xla/stream_executor/tpu/tpu_executable_interface.cc
+++ b/third_party/xla/xla/stream_executor/tpu/tpu_executable_interface.cc
@@ -51,9 +51,9 @@ namespace xla {
 namespace {
 
 // Write the tuple index buffers (arrays of pointers).
-static Status PopulateResultTupleBuffers(const ShapedBuffer& result,
-                                         se::Stream* stream,
-                                         se::Stream* transfer_stream) {
+static absl::Status PopulateResultTupleBuffers(const ShapedBuffer& result,
+                                               se::Stream* stream,
+                                               se::Stream* transfer_stream) {
   TF_ASSIGN_OR_RETURN(
       auto transfer_manager,
       TransferManager::GetForPlatform(stream->parent()->GetPlatform()));
diff --git a/third_party/xla/xla/stream_executor/tpu/tpu_executable_interface.h b/third_party/xla/xla/stream_executor/tpu/tpu_executable_interface.h
index ff037c07a532f5..b4e24551a80493 100644
--- a/third_party/xla/xla/stream_executor/tpu/tpu_executable_interface.h
+++ b/third_party/xla/xla/stream_executor/tpu/tpu_executable_interface.h
@@ -69,7 +69,7 @@ class TpuExecutableInterface : public Executable {
       std::vector<ExecutionInput>* arguments, se::Stream* stream,
       se::Stream* transfer_stream = nullptr);
 
-  virtual Status LoadProgramAndEnqueueToStream(
+  virtual absl::Status LoadProgramAndEnqueueToStream(
       const ServiceExecutableRunOptions& run_options,
       absl::Span<const stream_executor::DeviceMemoryBase> arguments,
       stream_executor::DeviceMemoryBase result,
diff --git a/third_party/xla/xla/stream_executor/tpu/tpu_executor.cc b/third_party/xla/xla/stream_executor/tpu/tpu_executor.cc
index ffab1983d879a0..5b493b104d432d 100644
--- a/third_party/xla/xla/stream_executor/tpu/tpu_executor.cc
+++ b/third_party/xla/xla/stream_executor/tpu/tpu_executor.cc
@@ -45,13 +45,9 @@ limitations under the License.
 namespace stream_executor {
 namespace tpu {
 
-namespace {
-using absl::Status;
-}  // namespace
-
 TpuExecutor::~TpuExecutor() { ExecutorApiFn()->TpuExecutor_FreeFn(executor_); }
 
-Status TpuExecutor::Init() {
+absl::Status TpuExecutor::Init() {
   StatusHelper status;
   ExecutorApiFn()->TpuExecutor_InitFn(executor_, status.c_status);
   return status.status();
@@ -61,14 +57,14 @@ bool TpuExecutor::SynchronizeAllActivity() {
   return ExecutorApiFn()->TpuExecutor_SynchronizeAllActivityFn(executor_);
 }
 
-Status TpuExecutor::BlockHostUntilDone(Stream* stream) {
+absl::Status TpuExecutor::BlockHostUntilDone(Stream* stream) {
   StatusHelper status;
   ExecutorApiFn()->TpuExecutor_BlockHostUntilDoneFn(
       executor_, get_stream(stream->implementation()), status.c_status);
   return status.status();
 }
 
-Status TpuExecutor::GetStatus(Stream* stream) {
+absl::Status TpuExecutor::GetStatus(Stream* stream) {
   StatusHelper status;
   ExecutorApiFn()->TpuExecutor_GetStatusFn(
       executor_, get_stream(stream->implementation()), status.c_status);
@@ -95,7 +91,7 @@ bool TpuExecutor::CreateStreamDependency(Stream* dependent, Stream* other) {
       get_stream(other->implementation()));
 }
 
-Status TpuExecutor::AllocateEvent(Event* event) {
+absl::Status TpuExecutor::AllocateEvent(Event* event) {
   StatusHelper status;
   auto se_event = tpu_platform().LookupEvent(event->implementation());
   ExecutorApiFn()->TpuExecutor_AllocateEventFn(executor_, se_event,
@@ -103,7 +99,7 @@ Status TpuExecutor::AllocateEvent(Event* event) {
   return status.status();
 }
 
-Status TpuExecutor::DeallocateEvent(Event* event) {
+absl::Status TpuExecutor::DeallocateEvent(Event* event) {
   tpu_platform().EraseEvent(event->implementation());
   return absl::OkStatus();
 }
@@ -115,8 +111,8 @@ stream_executor::Event::Status TpuExecutor::PollForEventStatus(
       ExecutorApiFn()->TpuExecutor_PollForEventStatusFn(executor_, se_event));
 }
 
-Status TpuExecutor::RecordEvent(Stream* stream,
-                                ::stream_executor::Event* event) {
+absl::Status TpuExecutor::RecordEvent(Stream* stream,
+                                      ::stream_executor::Event* event) {
   StatusHelper status;
   auto se_event = tpu_platform().LookupEvent(event->implementation());
   ExecutorApiFn()->TpuExecutor_RecordEventFn(
@@ -125,8 +121,8 @@ Status TpuExecutor::RecordEvent(Stream* stream,
   return status.status();
 }
 
-Status TpuExecutor::WaitForEvent(Stream* stream,
-                                 ::stream_executor::Event* event) {
+absl::Status TpuExecutor::WaitForEvent(Stream* stream,
+                                       ::stream_executor::Event* event) {
   StatusHelper status;
   auto se_event = tpu_platform().LookupEvent(event->implementation());
   ExecutorApiFn()->TpuExecutor_WaitForEventFn(
@@ -215,8 +211,8 @@ void TpuExecutor::DequeueOutfeed(int32_t outfeed_queue_index,
   done(status.status());
 }
 
-Status TpuExecutor::EnqueueInfeed(int32_t infeed_queue_index,
-                                  absl::Span<const uint8_t> bytes) {
+absl::Status TpuExecutor::EnqueueInfeed(int32_t infeed_queue_index,
+                                        absl::Span<const uint8_t> bytes) {
   StatusHelper status;
   ExecutorApiFn()->TpuExecutor_EnqueueInfeedFn(executor_, infeed_queue_index,
                                                bytes.data(), bytes.size(),
@@ -246,7 +242,7 @@ absl::Status TpuExecutor::Memcpy(
   return status.status();
 }
 
-Status TpuExecutor::SynchronousMemcpy(
+absl::Status TpuExecutor::SynchronousMemcpy(
     ::stream_executor::DeviceMemoryBase* device_dst, const void* host_src,
     uint64_t size) {
   StatusHelper status;
@@ -256,7 +252,7 @@ Status TpuExecutor::SynchronousMemcpy(
   return status.status();
 }
 
-Status TpuExecutor::SynchronousMemcpy(
+absl::Status TpuExecutor::SynchronousMemcpy(
     void* host_dst, const ::stream_executor::DeviceMemoryBase& device_src,
     uint64_t size) {
   StatusHelper status;
@@ -272,13 +268,14 @@ bool TpuExecutor::MemcpyDeviceToDevice(
   LOG(FATAL) << __func__ << " not supported on TpuExecutor";
 }
 
-Status TpuExecutor::UnloadAllPrograms() {
+absl::Status TpuExecutor::UnloadAllPrograms() {
   StatusHelper status;
   ExecutorApiFn()->TpuExecutor_UnloadAllProgramsFn(executor_, status.c_status);
   return status.status();
 }
 
-Status TpuExecutor::EnqueueCompactionOnStreamForHbm(Stream* compaction_stream) {
+absl::Status TpuExecutor::EnqueueCompactionOnStreamForHbm(
+    Stream* compaction_stream) {
   StatusHelper status;
   ExecutorApiFn()->TpuExecutor_EnqueueCompactionOnStreamForHbmFn(
       executor_, get_stream(compaction_stream->implementation()),
@@ -287,12 +284,12 @@ Status TpuExecutor::EnqueueCompactionOnStreamForHbm(Stream* compaction_stream) {
 }
 
 struct HostCallbackContext {
-  absl::AnyInvocable<Status() &&> callback;
+  absl::AnyInvocable<absl::Status() &&> callback;
 };
 
 TSL_Status* HostCallbackTrampoline(void* ctx) {
   HostCallbackContext* host_ctx = reinterpret_cast<HostCallbackContext*>(ctx);
-  Status status = std::move(host_ctx->callback)();
+  absl::Status status = std::move(host_ctx->callback)();
   TSL_Status* c_status = ExecutorApiFn()->TpuStatus_CreateFn(
       status.raw_code(), absl::StatusMessageAsCStr(status));
   delete host_ctx;
@@ -300,14 +297,14 @@ TSL_Status* HostCallbackTrampoline(void* ctx) {
 }
 
 bool TpuExecutor::HostCallback(Stream* stream,
-                               absl::AnyInvocable<Status() &&> callback) {
+                               absl::AnyInvocable<absl::Status() &&> callback) {
   HostCallbackContext* ctx = new HostCallbackContext{std::move(callback)};
   return ExecutorApiFn()->TpuExecutor_HostCallbackFn(
       executor_, get_stream(stream->implementation()), &HostCallbackTrampoline,
       ctx);
 }
 
-TpuExecutor::StatusOr<std::unique_ptr<::stream_executor::DeviceDescription>>
+absl::StatusOr<std::unique_ptr<::stream_executor::DeviceDescription>>
 TpuExecutor::CreateDeviceDescription() const {
   StatusHelper status;
   SE_DeviceDescription* description =
diff --git a/third_party/xla/xla/stream_executor/tpu/tpu_executor.h b/third_party/xla/xla/stream_executor/tpu/tpu_executor.h
index 74b17765ef5b92..db41a12f552dad 100644
--- a/third_party/xla/xla/stream_executor/tpu/tpu_executor.h
+++ b/third_party/xla/xla/stream_executor/tpu/tpu_executor.h
@@ -51,8 +51,6 @@ namespace tpu {
 
 class TpuExecutor : public tensorflow::tpu::TpuExecutorInterface {
  public:
-  template <typename T>
-  using StatusOr = ::absl::StatusOr<T>;
   using StatusCallback = std::function<void(const absl::Status&)>;
 
   TpuExecutor(::tensorflow::tpu::TpuPlatformInterface* platform,
@@ -72,7 +70,7 @@ class TpuExecutor : public tensorflow::tpu::TpuExecutorInterface {
 
   absl::Status BlockHostUntilDone(Stream* stream) override;
 
-  StatusOr<std::unique_ptr<DeviceDescription>> CreateDeviceDescription()
+  absl::StatusOr<std::unique_ptr<DeviceDescription>> CreateDeviceDescription()
       const override;
 
   bool CreateStreamDependency(Stream* dependent, Stream* other) override;
@@ -147,7 +145,7 @@ class TpuExecutor : public tensorflow::tpu::TpuExecutorInterface {
   int device_ordinal() const override { return device_ordinal_; }
   // TODO(henrytan): convert this to override once the base interface is changed
   // to TpuExecutorInterface.
-  StatusOr<std::unique_ptr<
+  absl::StatusOr<std::unique_ptr<
       tensorflow::tpu::TpuExecutorInterface::TemporaryDeviceMemory>>
   CreateTemporaryDeviceMemory(int64_t memory_space, int64_t byte_offset,
                               int64_t size) override {
diff --git a/third_party/xla/xla/stream_executor/tpu/tpu_executor_interface.h b/third_party/xla/xla/stream_executor/tpu/tpu_executor_interface.h
index f11218fc6b60eb..7816efbd5c2637 100644
--- a/third_party/xla/xla/stream_executor/tpu/tpu_executor_interface.h
+++ b/third_party/xla/xla/stream_executor/tpu/tpu_executor_interface.h
@@ -38,8 +38,6 @@ class TpuExecutorInterface : public stream_executor::StreamExecutor {
  public:
   explicit TpuExecutorInterface(stream_executor::Platform* platform)
       : StreamExecutor(platform) {}
-  template <typename T>
-  using StatusOr = absl::StatusOr<T>;
 
   class TemporaryDeviceMemory {
    public:
@@ -47,7 +45,7 @@ class TpuExecutorInterface : public stream_executor::StreamExecutor {
     virtual stream_executor::DeviceMemoryBase AsDeviceMemoryBase() const = 0;
   };
 
-  virtual StatusOr<std::unique_ptr<TemporaryDeviceMemory>>
+  virtual absl::StatusOr<std::unique_ptr<TemporaryDeviceMemory>>
   CreateTemporaryDeviceMemory(int64_t memory_space, int64_t byte_offset,
                               int64_t size) {
     LOG(FATAL) << "Unimplemented.";
diff --git a/third_party/xla/xla/stream_executor/tpu/tpu_node_context.cc b/third_party/xla/xla/stream_executor/tpu/tpu_node_context.cc
index 3734ce92663064..1a026c9c340242 100644
--- a/third_party/xla/xla/stream_executor/tpu/tpu_node_context.cc
+++ b/third_party/xla/xla/stream_executor/tpu/tpu_node_context.cc
@@ -28,10 +28,8 @@ limitations under the License.
 namespace tensorflow {
 namespace tpu {
 
-using absl::StatusOr;
-
 /*static*/
-StatusOr<std::unique_ptr<TpuNodeContext>> TpuNodeContext::Create(
+absl::StatusOr<std::unique_ptr<TpuNodeContext>> TpuNodeContext::Create(
     int device_ordinal) {
   StatusHelper status;
   XLA_TpuNodeContext* node_context =
diff --git a/third_party/xla/xla/stream_executor/tpu/tpu_node_context.h b/third_party/xla/xla/stream_executor/tpu/tpu_node_context.h
index 47e1d4d216682e..95fde9a67f9cbf 100644
--- a/third_party/xla/xla/stream_executor/tpu/tpu_node_context.h
+++ b/third_party/xla/xla/stream_executor/tpu/tpu_node_context.h
@@ -39,10 +39,8 @@ namespace tpu {
 // individual nodes.
 class TpuNodeContext final {
  public:
-  template <typename T>
-  using StatusOr = absl::StatusOr<T>;
-
-  static StatusOr<std::unique_ptr<TpuNodeContext>> Create(int device_ordinal);
+  static absl::StatusOr<std::unique_ptr<TpuNodeContext>> Create(
+      int device_ordinal);
 
   explicit TpuNodeContext(int device_ordinal, XLA_TpuNodeContext* node_context)
       : device_ordinal_(device_ordinal), node_context_(node_context) {
diff --git a/third_party/xla/xla/stream_executor/tpu/tpu_platform.cc b/third_party/xla/xla/stream_executor/tpu/tpu_platform.cc
index 2649985da24024..73bfc99f635fdf 100644
--- a/third_party/xla/xla/stream_executor/tpu/tpu_platform.cc
+++ b/third_party/xla/xla/stream_executor/tpu/tpu_platform.cc
@@ -47,9 +47,6 @@ namespace tpu {
 const ::stream_executor::Platform::Id TpuPlatform::kId = GetTpuPlatformId();
 TpuPlatform* tpu_registered_platform = nullptr;
 
-template <typename T>
-using StatusOr = ::absl::StatusOr<T>;
-
 TpuPlatform::TpuPlatform() : name_("TPU") {
   platform_ = stream_executor::tpu::ExecutorApiFn()->TpuPlatform_NewFn();
   CHECK(platform_ != nullptr);
@@ -99,13 +96,13 @@ int TpuPlatform::VisibleDeviceCount() const {
       ->TpuPlatform_VisibleDeviceCountFn(platform_);
 }
 
-StatusOr<::stream_executor::StreamExecutor*> TpuPlatform::GetExecutor(
+absl::StatusOr<::stream_executor::StreamExecutor*> TpuPlatform::GetExecutor(
     const ::stream_executor::StreamExecutorConfig& config) {
   return executor_cache_.GetOrCreate(
       config, [&]() { return GetUncachedExecutor(config); });
 }
 
-StatusOr<std::unique_ptr<::stream_executor::StreamExecutor>>
+absl::StatusOr<std::unique_ptr<::stream_executor::StreamExecutor>>
 TpuPlatform::GetUncachedExecutor(
     const ::stream_executor::StreamExecutorConfig& config) {
   SE_StreamExecutorConfig* c_config = stream_executor::tpu::ExecutorApiFn()
diff --git a/third_party/xla/xla/stream_executor/tpu/tpu_platform.h b/third_party/xla/xla/stream_executor/tpu/tpu_platform.h
index a55470cbef3b7a..cb0e6c99acdc87 100644
--- a/third_party/xla/xla/stream_executor/tpu/tpu_platform.h
+++ b/third_party/xla/xla/stream_executor/tpu/tpu_platform.h
@@ -49,9 +49,6 @@ class TpuPlatform : public ::tensorflow::tpu::TpuPlatformInterface {
 
   static const ::stream_executor::Platform::Id kId;
 
-  template <typename T>
-  using StatusOr = ::absl::StatusOr<T>;
-
   TpuPlatform();
 
   ~TpuPlatform() override;
@@ -82,22 +79,22 @@ class TpuPlatform : public ::tensorflow::tpu::TpuPlatformInterface {
     LOG(FATAL) << "Not yet implemented";
   }
 
-  StatusOr<std::unique_ptr<::stream_executor::DeviceDescription>>
+  absl::StatusOr<std::unique_ptr<::stream_executor::DeviceDescription>>
   DescriptionForDevice(int ordinal) const override {
     LOG(FATAL) << "Not yet implemented";
   }
 
-  StatusOr<::stream_executor::StreamExecutor*> ExecutorForDevice(
+  absl::StatusOr<::stream_executor::StreamExecutor*> ExecutorForDevice(
       int ordinal) override {
     stream_executor::StreamExecutorConfig config;
     config.ordinal = ordinal;
     return GetExecutor(config);
   }
 
-  StatusOr<::stream_executor::StreamExecutor*> GetExecutor(
+  absl::StatusOr<::stream_executor::StreamExecutor*> GetExecutor(
       const ::stream_executor::StreamExecutorConfig& config) override;
 
-  StatusOr<std::unique_ptr<::stream_executor::StreamExecutor>>
+  absl::StatusOr<std::unique_ptr<::stream_executor::StreamExecutor>>
   GetUncachedExecutor(
       const ::stream_executor::StreamExecutorConfig& config) override;
 
diff --git a/third_party/xla/xla/stream_executor/tpu/tpu_profiler_c_api.h b/third_party/xla/xla/stream_executor/tpu/tpu_profiler_c_api.h
index 3f0572a66132b9..448dcba8efc1fa 100644
--- a/third_party/xla/xla/stream_executor/tpu/tpu_profiler_c_api.h
+++ b/third_party/xla/xla/stream_executor/tpu/tpu_profiler_c_api.h
@@ -61,7 +61,7 @@ TFTPU_CAPI_EXPORT void TpuProfiler_CollectData(TpuProfiler* tpu_profiler,
                                                uint8_t* buffer,
                                                size_t* size_in_bytes);
 
-// Status helpers to create TFStatus for Profiler.
+// absl::Status helpers to create TFStatus for Profiler.
 TF_Status* TpuStatus_New();
 void TpuStatus_Free(TF_Status* status);
 const char* TpuStatus_Message(TF_Status* status);
diff --git a/third_party/xla/xla/stream_executor/tpu/tpu_transfer_manager.cc b/third_party/xla/xla/stream_executor/tpu/tpu_transfer_manager.cc
index 8576c8bfcc0fb9..1273551804c408 100644
--- a/third_party/xla/xla/stream_executor/tpu/tpu_transfer_manager.cc
+++ b/third_party/xla/xla/stream_executor/tpu/tpu_transfer_manager.cc
@@ -48,9 +48,6 @@ limitations under the License.
 namespace tensorflow {
 namespace tpu {
 
-template <typename T>
-using StatusOr = absl::StatusOr<T>;
-
 TpuTransferManager::TpuTransferManager() {
   manager_ = stream_executor::tpu::ExecutorApiFn()->TpuTransferManager_NewFn();
 }
@@ -246,7 +243,7 @@ int64_t TpuTransferManager::GetByteSizeRequirement(
   return size_in_bytes;
 }
 
-StatusOr<xla::Shape> TpuTransferManager::ChooseCompactLayoutForShape(
+absl::StatusOr<xla::Shape> TpuTransferManager::ChooseCompactLayoutForShape(
     const xla::Shape& host_shape) const {
   XLA_Shape c_host_shape;
   ApiConverter::ToC(host_shape, &c_host_shape);
diff --git a/third_party/xla/xla/stream_executor/tpu/tpu_transfer_manager.h b/third_party/xla/xla/stream_executor/tpu/tpu_transfer_manager.h
index 4f441dcf239f0f..0efd2fa558a647 100644
--- a/third_party/xla/xla/stream_executor/tpu/tpu_transfer_manager.h
+++ b/third_party/xla/xla/stream_executor/tpu/tpu_transfer_manager.h
@@ -43,9 +43,6 @@ class TpuTransferManager : public xla::TpuTransferManagerInterface {
   TpuTransferManager();
   ~TpuTransferManager() override;
 
-  template <typename T>
-  using StatusOr = absl::StatusOr<T>;
-
   stream_executor::Platform::Id PlatformId() const override;
 
   xla::Shape HostShapeToDeviceShape(
@@ -79,7 +76,7 @@ class TpuTransferManager : public xla::TpuTransferManagerInterface {
 
   int64_t GetByteSizeRequirement(const xla::Shape& shape) const override;
 
-  StatusOr<xla::Shape> ChooseCompactLayoutForShape(
+  absl::StatusOr<xla::Shape> ChooseCompactLayoutForShape(
       const xla::Shape& host_shape) const override;
 
   bool CanShapedBufferBeAccessedNow(
diff --git a/third_party/xla/xla/stream_executor/tpu/tpu_transfer_manager_interface.h b/third_party/xla/xla/stream_executor/tpu/tpu_transfer_manager_interface.h
index a2bd574914db0d..62a9647e7650b8 100644
--- a/third_party/xla/xla/stream_executor/tpu/tpu_transfer_manager_interface.h
+++ b/third_party/xla/xla/stream_executor/tpu/tpu_transfer_manager_interface.h
@@ -29,11 +29,11 @@ namespace xla {
 
 class TpuTransferManagerInterface : public xla::TransferManager {
  public:
-  virtual Status TransferBuffersToInfeed(
+  virtual absl::Status TransferBuffersToInfeed(
       se::StreamExecutor* executor,
       const std::deque<tensorflow::tpu::NoncopyableBuffer>& buffers) = 0;
 
-  virtual Status LinearizeToBuffers(
+  virtual absl::Status LinearizeToBuffers(
       const LiteralSlice& literal, const Shape& device_shape,
       std::deque<tensorflow::tpu::NoncopyableBuffer>* buffers) = 0;
 

From dbb27477c35200ebce4c6da1cd5546d32363bb7c Mon Sep 17 00:00:00 2001
From: Kyle Lucke <klucke@google.com>
Date: Mon, 20 May 2024 11:12:00 -0700
Subject: [PATCH 422/478] Use absl::Status instead of xla::Status now that
 they're identical.

PiperOrigin-RevId: 635513502
---
 third_party/xla/xla/client/client.cc          | 39 +++++++++---------
 third_party/xla/xla/client/client.h           |  9 +++--
 third_party/xla/xla/client/global_data.cc     |  2 +-
 third_party/xla/xla/client/local_client.cc    | 14 +++----
 third_party/xla/xla/client/local_client.h     | 17 ++++----
 third_party/xla/xla/client/padding.cc         |  6 +--
 third_party/xla/xla/client/padding.h          |  6 +--
 third_party/xla/xla/client/xla_builder.cc     | 24 +++++------
 third_party/xla/xla/client/xla_builder.h      | 40 +++++++++----------
 .../xla/xla/client/xla_builder_test.cc        |  6 +--
 third_party/xla/xla/python/jax_jit.cc         |  2 +-
 .../xla/xla/python/outfeed_receiver_test.cc   |  7 ++--
 third_party/xla/xla/python/pmap_lib.cc        |  4 +-
 third_party/xla/xla/python/py_array.cc        | 24 +++++------
 third_party/xla/xla/python/py_array.h         | 14 +++----
 third_party/xla/xla/python/py_values.h        |  2 +-
 third_party/xla/xla/python/util.cc            |  4 +-
 third_party/xla/xla/python/util.h             |  2 +-
 .../python/xplane_to_profile_instructions.cc  |  4 +-
 .../python/xplane_to_profile_instructions.h   |  4 +-
 20 files changed, 117 insertions(+), 113 deletions(-)

diff --git a/third_party/xla/xla/client/client.cc b/third_party/xla/xla/client/client.cc
index 1d2fdab9029afa..d79d978dfeae54 100644
--- a/third_party/xla/xla/client/client.cc
+++ b/third_party/xla/xla/client/client.cc
@@ -49,7 +49,7 @@ absl::StatusOr<Literal> Client::Transfer(const GlobalData& data,
 
   VLOG(1) << "making transfer request";
   VLOG(3) << "TransferToClientRequest: {" << request.DebugString() << "}";
-  Status s = stub_->TransferToClient(&request, &response);
+  absl::Status s = stub_->TransferToClient(&request, &response);
   VLOG(1) << "done with request";
 
   if (!s.ok()) {
@@ -76,7 +76,7 @@ absl::StatusOr<std::unique_ptr<GlobalData>> Client::TransferToServer(
 
   VLOG(1) << "making transfer to server request";
   VLOG(3) << "TransferToServerRequest: {" << request.DebugString() << "}";
-  Status s = stub_->TransferToServer(&request, &response);
+  absl::Status s = stub_->TransferToServer(&request, &response);
   VLOG(1) << "done with request";
 
   if (!s.ok()) {
@@ -93,8 +93,9 @@ absl::StatusOr<std::unique_ptr<GlobalData>> Client::TransferToServer(
   return std::make_unique<GlobalData>(stub_, response.data());
 }
 
-Status Client::TransferToInfeed(const LiteralSlice& literal, int64_t replica_id,
-                                const DeviceHandle* device_handle) {
+absl::Status Client::TransferToInfeed(const LiteralSlice& literal,
+                                      int64_t replica_id,
+                                      const DeviceHandle* device_handle) {
   TransferToInfeedRequest request;
   *request.mutable_literal() = literal.ToProto();
   if (device_handle) {
@@ -105,7 +106,7 @@ Status Client::TransferToInfeed(const LiteralSlice& literal, int64_t replica_id,
 
   VLOG(1) << "making transfer to infeed request";
   VLOG(3) << "TransferToInfeedRequest: {" << request.DebugString() << "}";
-  Status s = stub_->TransferToInfeed(&request, &response);
+  absl::Status s = stub_->TransferToInfeed(&request, &response);
   VLOG(1) << "done with request";
 
   if (!s.ok()) {
@@ -130,7 +131,7 @@ absl::StatusOr<Literal> Client::TransferFromOutfeed(
 
   VLOG(1) << "making transfer from outfeed request";
   VLOG(3) << "TransferFromOutfeedRequest: {" << request.DebugString() << "}";
-  Status s = stub_->TransferFromOutfeed(&request, &response);
+  absl::Status s = stub_->TransferFromOutfeed(&request, &response);
   VLOG(1) << "done with request";
 
   if (!s.ok()) {
@@ -147,13 +148,13 @@ absl::StatusOr<Literal> Client::TransferFromOutfeed(
   return Literal::CreateFromProto(response.literal());
 }
 
-Status Client::ResetDevice() {
+absl::Status Client::ResetDevice() {
   ResetDeviceRequest request;
   ResetDeviceResponse response;
 
   VLOG(1) << "making reset device request";
   VLOG(3) << "ResetDeviceRequest: {" << request.DebugString() << "}";
-  Status s = stub_->ResetDevice(&request, &response);
+  absl::Status s = stub_->ResetDevice(&request, &response);
   VLOG(1) << "done with request";
 
   if (!s.ok()) {
@@ -192,7 +193,7 @@ absl::StatusOr<Literal> Client::ComputeConstant(
   ComputeConstantResponse response;
 
   VLOG(2) << "making compute-constant-graph request";
-  Status s = stub_->ComputeConstantGraph(&request, &response);
+  absl::Status s = stub_->ComputeConstantGraph(&request, &response);
   VLOG(2) << "done with request";
 
   if (!s.ok()) {
@@ -238,7 +239,7 @@ absl::StatusOr<ExecutionHandle> Client::Compile(
 
   CompileResponse response;
   VLOG(1) << "making compile request: " << request.ShortDebugString();
-  Status s = stub_->Compile(&request, &response);
+  absl::Status s = stub_->Compile(&request, &response);
   VLOG(1) << "done with request";
 
   if (!s.ok()) {
@@ -260,7 +261,7 @@ absl::StatusOr<std::unique_ptr<GlobalData>> Client::Execute(
 
   ExecuteResponse response;
   VLOG(1) << "making execute request: " << request.ShortDebugString();
-  Status s = stub_->Execute(&request, &response);
+  absl::Status s = stub_->Execute(&request, &response);
   VLOG(1) << "done with request";
 
   if (!s.ok()) {
@@ -342,7 +343,7 @@ Client::ExecuteParallel(absl::Span<const XlaComputationInstance> computations) {
   ExecuteParallelResponse response;
   VLOG(1) << "making execute-graph-parallel request: "
           << request.ShortDebugString();
-  Status s = stub_->ExecuteGraphParallel(&request, &response);
+  absl::Status s = stub_->ExecuteGraphParallel(&request, &response);
   VLOG(1) << "done with request";
 
   if (!s.ok()) {
@@ -372,7 +373,7 @@ absl::StatusOr<std::vector<DeviceHandle>> Client::GetDeviceHandles(
 
   GetDeviceHandlesResponse response;
   VLOG(1) << "making get device request: " << request.ShortDebugString();
-  Status s = stub_->GetDeviceHandles(&request, &response);
+  absl::Status s = stub_->GetDeviceHandles(&request, &response);
   VLOG(1) << "done with request";
 
   if (!s.ok()) {
@@ -389,13 +390,13 @@ absl::StatusOr<std::vector<DeviceHandle>> Client::GetDeviceHandles(
   return device_handles;
 }
 
-Status Client::Unregister(const GlobalData& data) {
+absl::Status Client::Unregister(const GlobalData& data) {
   UnregisterRequest request;
   *request.add_data() = data.handle();
   UnregisterResponse response;
 
   VLOG(1) << "making unregister request";
-  Status s = stub_->Unregister(&request, &response);
+  absl::Status s = stub_->Unregister(&request, &response);
   VLOG(1) << "done with request";
 
   return s;
@@ -408,7 +409,7 @@ Client::DeconstructTuple(const GlobalData& data) {
   DeconstructTupleResponse response;
 
   VLOG(1) << "making DestructTuple request";
-  Status s = stub_->DeconstructTuple(&request, &response);
+  absl::Status s = stub_->DeconstructTuple(&request, &response);
   VLOG(1) << "done with request";
 
   if (!s.ok()) {
@@ -433,7 +434,7 @@ absl::StatusOr<ComputationStats> Client::GetComputationStats(
   ComputationStatsResponse response;
 
   VLOG(1) << "making computation graph stats request";
-  Status s = stub_->GetComputationGraphStats(&request, &response);
+  absl::Status s = stub_->GetComputationGraphStats(&request, &response);
   VLOG(1) << "done with request";
 
   if (!s.ok()) {
@@ -455,7 +456,7 @@ absl::StatusOr<Shape> Client::GetShape(const GlobalData& data) {
   GetShapeResponse response;
 
   VLOG(1) << "making get shape request";
-  Status s = stub_->GetShape(&request, &response);
+  absl::Status s = stub_->GetShape(&request, &response);
   VLOG(1) << "done with request";
 
   if (!s.ok()) {
@@ -493,7 +494,7 @@ absl::StatusOr<ChannelHandle> Client::CreateChannelHandleByType(
   CreateChannelHandleResponse response;
 
   VLOG(1) << "making create channel handle request";
-  Status s = stub_->CreateChannelHandle(&request, &response);
+  absl::Status s = stub_->CreateChannelHandle(&request, &response);
   VLOG(1) << "done with request";
 
   if (!s.ok()) {
diff --git a/third_party/xla/xla/client/client.h b/third_party/xla/xla/client/client.h
index e4c48d58359958..9831182c9c40e0 100644
--- a/third_party/xla/xla/client/client.h
+++ b/third_party/xla/xla/client/client.h
@@ -151,8 +151,9 @@ class Client {
   // device_handle and replica_id together specify a particular device; a device
   // assigned for the given replica_id among the replicas that the given device
   // handle belongs to.
-  Status TransferToInfeed(const LiteralSlice& literal, int64_t replica_id = 0,
-                          const DeviceHandle* device_handle = nullptr);
+  absl::Status TransferToInfeed(const LiteralSlice& literal,
+                                int64_t replica_id = 0,
+                                const DeviceHandle* device_handle = nullptr);
 
   // Transfers from the Outfeed of the device.
   //
@@ -164,7 +165,7 @@ class Client {
       const DeviceHandle* device_handle = nullptr);
 
   // Resets the device, clearing all existing state on the device.
-  Status ResetDevice();
+  absl::Status ResetDevice();
 
   // Executes the computation with the given arguments and transfers the result
   // to the client as a literal. Parameters are defined the same as for
@@ -195,7 +196,7 @@ class Client {
       const Layout* output_layout = nullptr) const;
 
   // Unregister the memory for the given GlobalData on the device.
-  Status Unregister(const GlobalData& data);
+  absl::Status Unregister(const GlobalData& data);
 
   // Returns a vector of global data handles that point to the tuple elements.
   absl::StatusOr<std::vector<std::unique_ptr<GlobalData>>> DeconstructTuple(
diff --git a/third_party/xla/xla/client/global_data.cc b/third_party/xla/xla/client/global_data.cc
index 1235d1e3c50d6b..5fb1ca9df693ac 100644
--- a/third_party/xla/xla/client/global_data.cc
+++ b/third_party/xla/xla/client/global_data.cc
@@ -37,7 +37,7 @@ void ReleaseHandles(ServiceInterface* parent,
     *request.add_data() = handle;
   }
   UnregisterResponse response;
-  Status status = parent->Unregister(&request, &response);
+  absl::Status status = parent->Unregister(&request, &response);
   VLOG(1) << "Done with request";
   if (!status.ok()) {
     LOG(WARNING) << "Failed to unregister handles: " << status
diff --git a/third_party/xla/xla/client/local_client.cc b/third_party/xla/xla/client/local_client.cc
index 59d4ecf6a6e82d..f8a29b126dacd3 100644
--- a/third_party/xla/xla/client/local_client.cc
+++ b/third_party/xla/xla/client/local_client.cc
@@ -51,7 +51,7 @@ LocalExecutable::LocalExecutable(std::unique_ptr<Executable> executable,
       << "Must have a valid device ordinal that the executable was built for.";
 }
 
-Status LocalExecutable::ValidateExecutionOptions(
+absl::Status LocalExecutable::ValidateExecutionOptions(
     const ExecutableRunOptions& run_options, const Backend& backend) {
   if (run_options.stream() != nullptr) {
     if (!run_options.stream()->ok()) {
@@ -207,7 +207,7 @@ static std::shared_ptr<HloSnapshot> DumpArguments(
   for (const ShapedBuffer* arg : arguments) {
     auto literal = std::make_shared<Literal>(arg->on_host_shape());
     backend->transfer_manager()->TransferLiteralFromDevice(
-        stream, *arg, literal.get(), [snapshot, literal](Status status) {
+        stream, *arg, literal.get(), [snapshot, literal](absl::Status status) {
           if (!status.ok()) {
             LOG(ERROR) << "TransferLiteralFromDevice for HLO snapshot inputs "
                           "failed: "
@@ -227,7 +227,7 @@ static void DumpOutputsAndSaveSnapshot(const Backend* backend,
   auto literal = std::make_shared<Literal>(outputs.on_host_shape());
   backend->transfer_manager()->TransferLiteralFromDevice(
       stream, outputs, literal.get(),
-      [snapshot{std::move(snapshot)}, literal](Status status) {
+      [snapshot{std::move(snapshot)}, literal](absl::Status status) {
         if (status.ok()) {
           *snapshot->mutable_result() = literal->ToProto();
         } else {
@@ -471,16 +471,16 @@ absl::StatusOr<const ShapedBuffer*> LocalClient::GlobalDataToShapedBuffer(
   return local_service_->GlobalDataToShapedBuffer(data, replica_number);
 }
 
-Status LocalClient::TransferToInfeedLocal(const LiteralSlice& literal,
-                                          int device_ordinal) {
+absl::Status LocalClient::TransferToInfeedLocal(const LiteralSlice& literal,
+                                                int device_ordinal) {
   TF_ASSIGN_OR_RETURN(se::StreamExecutor * executor,
                       backend().stream_executor(device_ordinal));
   return backend().transfer_manager()->TransferLiteralToInfeed(executor,
                                                                literal);
 }
 
-Status LocalClient::TransferFromOutfeedLocal(int device_ordinal,
-                                             MutableBorrowingLiteral literal) {
+absl::Status LocalClient::TransferFromOutfeedLocal(
+    int device_ordinal, MutableBorrowingLiteral literal) {
   TF_ASSIGN_OR_RETURN(se::StreamExecutor * executor,
                       backend().stream_executor(device_ordinal));
   return backend().transfer_manager()->TransferLiteralFromOutfeed(executor,
diff --git a/third_party/xla/xla/client/local_client.h b/third_party/xla/xla/client/local_client.h
index d6e382192cab86..a73210b0af7f63 100644
--- a/third_party/xla/xla/client/local_client.h
+++ b/third_party/xla/xla/client/local_client.h
@@ -85,8 +85,8 @@ class LocalExecutable {
   //
   // The given ExecutableRunOptions override any values from TF_XLA_FLAGS
   // environment variable.
-  Status ValidateExecutionOptions(const ExecutableRunOptions& run_options,
-                                  const Backend& backend);
+  absl::Status ValidateExecutionOptions(const ExecutableRunOptions& run_options,
+                                        const Backend& backend);
 
   // Returns a literal containing the contents of the given ShapedBuffer.
   absl::StatusOr<Literal> LiteralFromShapedBuffer(
@@ -102,7 +102,7 @@ class LocalExecutable {
   int build_device_ordinal() const { return build_options_.device_ordinal(); }
 
   template <typename T>
-  StatusOr<T> AsyncCallAndBlockHostUntilDone(
+  absl::StatusOr<T> AsyncCallAndBlockHostUntilDone(
       absl::Span<Shape const* const> argument_shapes,
       const ExecutableRunOptions& run_options,
       std::function<StatusOr<T>(const ExecutableRunOptions&)> async_callback) {
@@ -110,8 +110,8 @@ class LocalExecutable {
                         RunHelper(argument_shapes, run_options));
     ExecutableRunOptions options = options_and_stream.first.run_options();
     options.set_device_ordinal(-1);
-    StatusOr<T> result = async_callback(options);
-    Status block_status = options.stream()->BlockHostUntilDone();
+    absl::StatusOr<T> result = async_callback(options);
+    absl::Status block_status = options.stream()->BlockHostUntilDone();
     TF_RETURN_IF_ERROR(result.status());
     TF_RETURN_IF_ERROR(block_status);
     return result;
@@ -188,15 +188,16 @@ class LocalClient : public Client {
   // TODO(b/69670845): Remove the 'Local' from the name when LocalClient does
   // not inherit from Client and there is no possibility of confusion with
   // Client::TransferToInfeed.
-  Status TransferToInfeedLocal(const LiteralSlice& literal, int device_ordinal);
+  absl::Status TransferToInfeedLocal(const LiteralSlice& literal,
+                                     int device_ordinal);
 
   // Transfer and return a value from the outfeed of the given device. The
   // shape of the object to transfer is determined by `literal`'s shape.
   // TODO(b/69670845): Remove the 'Local' from the name when LocalClient does
   // not inherit from Client and there is no possibility of confusion with
   // Client::TransferFromOutfeed.
-  Status TransferFromOutfeedLocal(int device_ordinal,
-                                  MutableBorrowingLiteral literal);
+  absl::Status TransferFromOutfeedLocal(int device_ordinal,
+                                        MutableBorrowingLiteral literal);
 
   // Returns the device ordinal that corresponds to the given replica number.
   //
diff --git a/third_party/xla/xla/client/padding.cc b/third_party/xla/xla/client/padding.cc
index f36c1cd18a505b..22c4901a512362 100644
--- a/third_party/xla/xla/client/padding.cc
+++ b/third_party/xla/xla/client/padding.cc
@@ -25,9 +25,9 @@ limitations under the License.
 
 namespace xla {
 
-Status ValidatePaddingValues(absl::Span<const int64_t> input_dimensions,
-                             absl::Span<const int64_t> window_dimensions,
-                             absl::Span<const int64_t> window_strides) {
+absl::Status ValidatePaddingValues(absl::Span<const int64_t> input_dimensions,
+                                   absl::Span<const int64_t> window_dimensions,
+                                   absl::Span<const int64_t> window_strides) {
   bool ok = input_dimensions.size() == window_dimensions.size() &&
             input_dimensions.size() == window_strides.size();
   if (!ok) {
diff --git a/third_party/xla/xla/client/padding.h b/third_party/xla/xla/client/padding.h
index 5c4a34b0d09220..50ed6e58057ad8 100644
--- a/third_party/xla/xla/client/padding.h
+++ b/third_party/xla/xla/client/padding.h
@@ -41,9 +41,9 @@ enum class Padding {
 // Validates that the slices are acceptable for determining padding -- this can
 // be used to check the preconditions of MakePadding below to produce an error
 // message that can be returned to the user.
-Status ValidatePaddingValues(absl::Span<const int64_t> input_dimensions,
-                             absl::Span<const int64_t> window_dimensions,
-                             absl::Span<const int64_t> window_strides);
+absl::Status ValidatePaddingValues(absl::Span<const int64_t> input_dimensions,
+                                   absl::Span<const int64_t> window_dimensions,
+                                   absl::Span<const int64_t> window_strides);
 
 // Returns the padding needed for the base area, given the base area dimensions,
 // window dimensions, strides, and the type of padding.
diff --git a/third_party/xla/xla/client/xla_builder.cc b/third_party/xla/xla/client/xla_builder.cc
index 9d32c48ea9afad..e6b4e955aa69dd 100644
--- a/third_party/xla/xla/client/xla_builder.cc
+++ b/third_party/xla/xla/client/xla_builder.cc
@@ -112,8 +112,8 @@ bool InstrIsSetBound(const HloInstructionProto* instr_proto) {
   return false;
 }
 
-Status NormalizeAndAssignSharing(HloInstructionProto* instr,
-                                 const OpSharding& op_sharding) {
+absl::Status NormalizeAndAssignSharing(HloInstructionProto* instr,
+                                       const OpSharding& op_sharding) {
   // Normalize tuple sharding and fail the call if the sharding is invalid.
   Shape shape(instr->shape());
   TF_ASSIGN_OR_RETURN(HloSharding sharding,
@@ -532,7 +532,7 @@ XlaBuilder::XlaBuilder(const std::string& computation_name)
 
 XlaBuilder::~XlaBuilder() = default;
 
-XlaOp XlaBuilder::ReportError(const Status& error) {
+XlaOp XlaBuilder::ReportError(const absl::Status& error) {
   CHECK(!error.ok());
   if (die_immediately_on_error_) {
     LOG(FATAL) << "error building computation: " << error;
@@ -687,16 +687,16 @@ void XlaBuilder::IsConstantVisitor(const int64_t op_handle, int depth,
   visited->insert(op_handle);
 }
 
-Status XlaBuilder::SetInstructionFrontendAttribute(const XlaOp op,
-                                                   std::string attribute,
-                                                   std::string value) {
+absl::Status XlaBuilder::SetInstructionFrontendAttribute(const XlaOp op,
+                                                         std::string attribute,
+                                                         std::string value) {
   TF_ASSIGN_OR_RETURN(auto instr_proto, LookUpMutableInstruction(op));
   auto* frontend_attributes = instr_proto->mutable_frontend_attributes();
   (*frontend_attributes->mutable_map())[attribute] = std::move(value);
   return OkStatus();
 }
 
-Status XlaBuilder::SetInstructionSharding(
+absl::Status XlaBuilder::SetInstructionSharding(
     XlaOp op, const std::optional<OpSharding>& sharding) {
   TF_ASSIGN_OR_RETURN(auto instr_proto, LookUpMutableInstruction(op));
   if (!sharding.has_value()) {
@@ -717,7 +717,7 @@ XlaComputation XlaBuilder::BuildAndNoteError() {
   return std::move(build_status).value();
 }
 
-Status XlaBuilder::GetCurrentStatus() const {
+absl::Status XlaBuilder::GetCurrentStatus() const {
   if (!first_error_.ok()) {
     std::string backtrace;
     first_error_backtrace_.Dump(tsl::DebugWriteToString, &backtrace);
@@ -804,7 +804,7 @@ absl::StatusOr<XlaComputation> XlaBuilder::Build(
   return std::move(computation);
 }
 
-/* static */ Status XlaBuilder::PopulateInputOutputAliasAndBufferDonor(
+/* static */ absl::Status XlaBuilder::PopulateInputOutputAliasAndBufferDonor(
     HloModuleProto* module, const ProgramShape& program_shape,
     const std::vector<InputOutputAlias>& input_output_aliases,
     const absl::flat_hash_set<HloBufferDonorConfig::BufferDonor>&
@@ -1987,7 +1987,7 @@ XlaOp XlaBuilder::SparseDot(
   });
 }
 
-Status XlaBuilder::VerifyConvolution(
+absl::Status XlaBuilder::VerifyConvolution(
     const Shape& lhs_shape, const Shape& rhs_shape,
     const ConvolutionDimensionNumbers& dimension_numbers) const {
   if (lhs_shape.rank() != rhs_shape.rank()) {
@@ -3333,7 +3333,7 @@ XlaOp XlaBuilder::ConditionalImpl(
   });
 }
 
-Status XlaBuilder::CheckOpBuilder(XlaOp op) const {
+absl::Status XlaBuilder::CheckOpBuilder(XlaOp op) const {
   if (this != op.builder()) {
     return InvalidArgument(
         "XlaOp with handle %d is built by builder '%s', but is trying to use "
@@ -4605,7 +4605,7 @@ XlaBuilder::CreateDefaultConvDimensionNumbers(int num_spatial_dims) {
   return dimension_numbers;
 }
 
-/* static */ Status XlaBuilder::Validate(
+/* static */ absl::Status XlaBuilder::Validate(
     const ConvolutionDimensionNumbers& dnum) {
   if (dnum.input_spatial_dimensions_size() < 2) {
     return FailedPrecondition("input spacial dimension < 2: %d",
diff --git a/third_party/xla/xla/client/xla_builder.h b/third_party/xla/xla/client/xla_builder.h
index 84fa7588c69016..d88c90b2cf8896 100644
--- a/third_party/xla/xla/client/xla_builder.h
+++ b/third_party/xla/xla/client/xla_builder.h
@@ -336,7 +336,7 @@ class XlaBuilder {
       int num_spatial_dims = 2);
 
   // Returns an error if the convolution dimension numbers have conflicts.
-  static Status Validate(const ConvolutionDimensionNumbers& dnum);
+  static absl::Status Validate(const ConvolutionDimensionNumbers& dnum);
 
   // Returns a new XlaBuilder whose resultant Computation is used only by this
   // XlaBuilder. The sub-XlaBuilder has the same die_immediately_on_error
@@ -385,11 +385,11 @@ class XlaBuilder {
   // building the computation when they make a final call to Build().
   //
   // See also set_die_immediately_on_error().
-  Status first_error() const { return first_error_; }
+  absl::Status first_error() const { return first_error_; }
 
   // Returns the current status of the builder, complete with the stack trace
   // information.
-  Status GetCurrentStatus() const;
+  absl::Status GetCurrentStatus() const;
 
   // Returns the shape of the given op.
   absl::StatusOr<Shape> GetShape(XlaOp op) const;
@@ -416,15 +416,15 @@ class XlaBuilder {
   // * dying if die_immediately_on_error_ is true.
   // Returns an XlaOp with an invalid handle but a valid builder. This value can
   // be returned in place of a value in APIs that return an XlaOp.
-  XlaOp ReportError(const Status& error);
+  XlaOp ReportError(const absl::Status& error);
 
-  // A helper function that converts a StatusOr<XlaOp> into an XlaOp.
-  // If the Status was an error, reports the error to builder and returns an
-  // invalid XlaOp handle.
+  // A helper function that converts a absl::StatusOr<XlaOp> into an XlaOp.
+  // If the absl::Status was an error, reports the error to builder and returns
+  // an invalid XlaOp handle.
   XlaOp ReportErrorOrReturn(const absl::StatusOr<XlaOp>& op);
 
-  // A helper function that runs a function that returns a StatusOr<XlaOp> and
-  // returns an XlaOp.
+  // A helper function that runs a function that returns a absl::StatusOr<XlaOp>
+  // and returns an XlaOp.
   XlaOp ReportErrorOrReturn(
       absl::FunctionRef<absl::StatusOr<XlaOp>()> op_creator);
 
@@ -475,15 +475,15 @@ class XlaBuilder {
   // "value". If the attribute already existed, then its value is updated.
   //
   // The attribute is only added to the HloInstruction, not to the builder.
-  Status SetInstructionFrontendAttribute(XlaOp op, std::string attribute,
-                                         std::string value);
+  absl::Status SetInstructionFrontendAttribute(XlaOp op, std::string attribute,
+                                               std::string value);
 
   // Looks up the HloInstruction and sets the sharding. If the sharding already
   // existed, then its value is updated.
   //
   // The sharding is only added to the HloInstruction, not to the builder.
-  Status SetInstructionSharding(XlaOp op,
-                                const std::optional<OpSharding>& sharding);
+  absl::Status SetInstructionSharding(
+      XlaOp op, const std::optional<OpSharding>& sharding);
 
   // Returns shapes for the operands.
   absl::StatusOr<std::vector<Shape>> GetOperandShapes(
@@ -1117,7 +1117,7 @@ class XlaBuilder {
                          bool* is_constant) const;
 
   // Checks bounds for convolution parameters.
-  Status VerifyConvolution(
+  absl::Status VerifyConvolution(
       const Shape& lhs_shape, const Shape& rhs_shape,
       const ConvolutionDimensionNumbers& dimension_numbers) const;
 
@@ -1125,7 +1125,7 @@ class XlaBuilder {
 
   // Populates the module with the input/output alias information stored within
   // the input_output_aliases vector.
-  static Status PopulateInputOutputAliasAndBufferDonor(
+  static absl::Status PopulateInputOutputAliasAndBufferDonor(
       HloModuleProto* module, const ProgramShape& program_shape,
       const std::vector<InputOutputAlias>& input_output_aliases,
       const absl::flat_hash_set<HloBufferDonorConfig::BufferDonor>&
@@ -1139,7 +1139,7 @@ class XlaBuilder {
 
   // The first error encountered while building the computation.
   // This is OK until the first error is encountered.
-  Status first_error_;
+  absl::Status first_error_;
 
   // The saved stack trace from the point at which the first error occurred.
   tsl::SavedStackTrace first_error_backtrace_;
@@ -1661,7 +1661,7 @@ class XlaBuilder {
  protected:
   // Returns OK status if the given op was built using this builder. Otherwise,
   // returns an error.
-  Status CheckOpBuilder(XlaOp op) const;
+  absl::Status CheckOpBuilder(XlaOp op) const;
 
  private:
   XlaOp AllGatherImpl(XlaOp operand, int64_t all_gather_dimension,
@@ -1703,7 +1703,7 @@ class XlaBuilder {
   // Here, InstructionType is either const HloInstructionProto* or non-const
   // HloInstructionProto*.
   template <typename InstructionType>
-  StatusOr<InstructionType> LookUpInstructionByHandleInternal(
+  absl::StatusOr<InstructionType> LookUpInstructionByHandleInternal(
       int64_t handle) const {
     auto it = handle_to_index_.find(handle);
     if (it == handle_to_index_.end()) {
@@ -1723,11 +1723,11 @@ class XlaBuilder {
   // Here, InstructionType is either const HloInstructionProto* or non-const
   // HloInstructionProto*.
   //
-  // TODO(hinsu): Return const pointer within StatusOr and use
+  // TODO(hinsu): Return const pointer within absl::StatusOr and use
   // absl::implicit_cast at callsites. This requires implicit_cast support in
   // absl::StatusOr similar to absl::StatusOr.
   template <typename InstructionType>
-  StatusOr<InstructionType> LookUpInstructionInternal(XlaOp op) const {
+  absl::StatusOr<InstructionType> LookUpInstructionInternal(XlaOp op) const {
     TF_RETURN_IF_ERROR(CheckOpBuilder(op));
     return LookUpInstructionByHandleInternal<InstructionType>(op.handle());
   }
diff --git a/third_party/xla/xla/client/xla_builder_test.cc b/third_party/xla/xla/client/xla_builder_test.cc
index c1c9640aa841f1..9d3ae1850a9092 100644
--- a/third_party/xla/xla/client/xla_builder_test.cc
+++ b/third_party/xla/xla/client/xla_builder_test.cc
@@ -754,7 +754,7 @@ TEST(XlaBuilderTest, BuildWithSpecificRootWithWrongBuilder) {
   Parameter(&b, 0, shape, "param");
   const XlaOp other_param = Parameter(&other_b, 0, shape, "other_param");
 
-  Status status = b.Build(other_param).status();
+  absl::Status status = b.Build(other_param).status();
   ASSERT_IS_NOT_OK(status);
   EXPECT_THAT(
       status.message(),
@@ -1202,7 +1202,7 @@ TEST(XlaBuilderTest, DynamicSelectNotCompatible) {
   auto gte0 = GetTupleElement(p0, 0);  // f32[4,<=5,6]
   auto gte1 = GetTupleElement(p0, 1);  // f32[4,5,<=6]
   Select(pred, gte0, gte1);
-  Status status = BuildHloModule(b).status();
+  absl::Status status = BuildHloModule(b).status();
   ASSERT_IS_OK(status);
 }
 
@@ -1299,7 +1299,7 @@ TEST(XlaBuilderTest, ConvolutionWithPreferredElementType) {
 TEST(XlaBuilderTest, AfterAllWithNonTokenOperands) {
   XlaBuilder b(TestName());
   AfterAll(&b, {CreateToken(&b), ConstantR0<float>(&b, 1.0)});
-  Status status = b.Build().status();
+  absl::Status status = b.Build().status();
   ASSERT_IS_NOT_OK(status);
   EXPECT_THAT(status.message(),
               ::testing::HasSubstr("All operands to AfterAll must be tokens"));
diff --git a/third_party/xla/xla/python/jax_jit.cc b/third_party/xla/xla/python/jax_jit.cc
index 20ff684cd38ca4..fa6125ded04f6c 100644
--- a/third_party/xla/xla/python/jax_jit.cc
+++ b/third_party/xla/xla/python/jax_jit.cc
@@ -62,7 +62,7 @@ namespace jax {
 namespace nb = nanobind;
 
 // TODO(phawkins): Add support for Tracers.
-// TODO(jblespiau): Use absl Status.
+// TODO(jblespiau): Use absl absl::Status.
 
 namespace {
 
diff --git a/third_party/xla/xla/python/outfeed_receiver_test.cc b/third_party/xla/xla/python/outfeed_receiver_test.cc
index bdd794e8b6eff3..447e67cda2786c 100644
--- a/third_party/xla/xla/python/outfeed_receiver_test.cc
+++ b/third_party/xla/xla/python/outfeed_receiver_test.cc
@@ -33,8 +33,8 @@ namespace xla {
 
 namespace {
 
-Status CompileAndExecute(XlaBuilder* builder, XlaOp root, int device_id,
-                         PjRtClient* client) {
+absl::Status CompileAndExecute(XlaBuilder* builder, XlaOp root, int device_id,
+                               PjRtClient* client) {
   XlaComputation computation = builder->Build(root).value();
 
   CompileOptions compile_options;
@@ -79,7 +79,8 @@ class Accumulator {
 
 // TODO(necula): update this test for the TFRT CPU client, which current does
 // not support non-local devices.
-// StatusOr<std::unique_ptr<PjRtClient>> GetCpuClientWithNonLocalDevice() {
+// absl::StatusOr<std::unique_ptr<PjRtClient>> GetCpuClientWithNonLocalDevice()
+// {
 //   TF_ASSIGN_OR_RETURN(se::Platform * platform,
 //                       PlatformUtil::GetPlatform("Host"));
 //   if (platform->VisibleDeviceCount() <= 0) {
diff --git a/third_party/xla/xla/python/pmap_lib.cc b/third_party/xla/xla/python/pmap_lib.cc
index d8706a9db9c90f..97d96c8b64a398 100644
--- a/third_party/xla/xla/python/pmap_lib.cc
+++ b/third_party/xla/xla/python/pmap_lib.cc
@@ -125,8 +125,8 @@ struct ShardArgResult {
 // `arg`: The object to shard across `devices`. If a `Array`,
 //   a fast-path will be executed if it's already correctly sharded.
 //
-// Returns a failure Status when an unrecoverable error occurred, so we don't
-// need to fallback to Python.
+// Returns a failure absl::Status when an unrecoverable error occurred, so we
+// don't need to fallback to Python.
 //
 // Both `devices` and `sharding_spec` has the same length.
 absl::StatusOr<ShardArgResult> ShardArg(
diff --git a/third_party/xla/xla/python/py_array.cc b/third_party/xla/xla/python/py_array.cc
index 897872c9a68bcb..49d6d1c66e173f 100644
--- a/third_party/xla/xla/python/py_array.cc
+++ b/third_party/xla/xla/python/py_array.cc
@@ -142,7 +142,7 @@ tsl::RCReference<ifrt::Array> CreateIfRtArrayFromSingleDeviceShardedPyArrays(
     nb_dtype dtype, absl::Span<const int64_t> shape,
     absl::Span<const PyArray> py_arrays) {
   if (py_arrays.empty()) {
-    // TODO(hyeontaek): Return a Status.
+    // TODO(hyeontaek): Return a absl::Status.
     throw nb::value_error("At least one array must be provided.");
   }
   std::vector<tsl::RCReference<ifrt::Array>> ifrt_arrays;
@@ -181,7 +181,7 @@ tsl::RCReference<ifrt::Array> CreateIfRtArrayFromSingleDeviceShardedPyArrays(
 
   auto ifrt_dtype = DtypeToIfRtDType(dtype);
   if (!ifrt_dtype.ok()) {
-    // TODO(hyeontaek): Return a Status.
+    // TODO(hyeontaek): Return a absl::Status.
     throw nb::value_error(ifrt_dtype.status().ToString().c_str());
   }
   auto ifrt_array = client->AssembleArrayFromSingleDeviceArrays(
@@ -192,7 +192,7 @@ tsl::RCReference<ifrt::Array> CreateIfRtArrayFromSingleDeviceShardedPyArrays(
                                      /*shard_shapes=*/std::move(shapes)),
       absl::MakeSpan(ifrt_arrays), ifrt::ArrayCopySemantics::kReuseInput);
   if (!ifrt_array.ok()) {
-    // TODO(hyeontaek): Return a Status.
+    // TODO(hyeontaek): Return a absl::Status.
     throw nb::value_error(ifrt_array.status().ToString().c_str());
   }
   return *std::move(ifrt_array);
@@ -554,7 +554,7 @@ nb::object PyArray::arrays() {
   return nb::cast(py_arrays_cached());
 }
 
-Status PyArray::set_arrays(nb::object obj) {
+absl::Status PyArray::set_arrays(nb::object obj) {
   if (obj.is_none()) {
     SetIfrtArray(tsl::RCReference<ifrt::Array>());
     py_arrays().clear();
@@ -653,7 +653,7 @@ absl::StatusOr<PyArray> PyArray::FullyReplicatedShard() {
   return nb::cast<PyArray>(cached);
 }
 
-Status PyArray::BlockUntilReady() const {
+absl::Status PyArray::BlockUntilReady() const {
   nb::gil_scoped_release gil_release;
   if (ifrt_array() == nullptr) {
     return InvalidArgument(
@@ -697,7 +697,7 @@ absl::StatusOr<nb::object> PyArray::SingleDeviceArrayToNumpyArray() {
   return result;
 }
 
-Status PyArray::CopySingleDeviceArrayToHostAsync() {
+absl::Status PyArray::CopySingleDeviceArrayToHostAsync() {
   TF_ASSIGN_OR_RETURN(auto arr, FullyReplicatedShard());
   return arr.GetStorage().host_value.CopyToHostAsync(
       arr.GetStorage().dynamic_shape, arr.ifrt_array());
@@ -932,7 +932,7 @@ absl::StatusOr<nb::object> CudaArrayInterfaceToBuffer(
                                             std::move(ifrt_array), false, true);
 }
 
-Status PyArray::Delete() {
+absl::Status PyArray::Delete() {
   for (auto& arr : py_arrays()) {
     TF_RETURN_IF_ERROR(arr.Delete());
   }
@@ -1236,7 +1236,7 @@ bool HasDefaultLayout(const Layout& layout) {
 }
 
 int PyArray_bf_getbuffer(PyObject* exporter, Py_buffer* view, int flags) {
-  Status status = [&]() -> absl::Status {
+  absl::Status status = [&]() -> absl::Status {
     PyArray py_array = nb::borrow<PyArray>(exporter);
     if (py_array.ifrt_array() == nullptr) {
       // TODO(phawkins): why is this happening?
@@ -1466,8 +1466,8 @@ absl::StatusOr<nb::object> PyHostValue::AsNumPyArray(
   return value_;
 }
 
-Status PyHostValue::CopyToHostAsync(std::optional<Shape>& dynamic_shape_holder,
-                                    ifrt::Array* ifrt_array) {
+absl::Status PyHostValue::CopyToHostAsync(
+    std::optional<Shape>& dynamic_shape_holder, ifrt::Array* ifrt_array) {
   if (ready_.IsValid()) {
     // The array value has been populated, so CopyToHostAsync has been called.
     return OkStatus();
@@ -1520,7 +1520,7 @@ Status PyHostValue::CopyToHostAsync(std::optional<Shape>& dynamic_shape_holder,
                                         ifrt::ArrayCopySemantics::kReuseInput);
   // Make sure the destination of the copy remains alive until the copy is done.
   value_.inc_ref();
-  ready_.OnReady([array{value_.ptr()}](Status status) {
+  ready_.OnReady([array{value_.ptr()}](absl::Status status) {
     GlobalPyRefManager()->AddGarbage(nb::steal(array));
   });
   value_.attr("flags").attr("writeable") = nb::bool_(false);
@@ -1559,7 +1559,7 @@ PyType_Slot PyArray_slots[] = {
 
 }  // namespace
 
-Status PyArray::RegisterTypes(nb::module_& m) {
+absl::Status PyArray::RegisterTypes(nb::module_& m) {
   std::string name =
       absl::StrCat(nb::cast<std::string>(m.attr("__name__")), ".ArrayImpl");
 
diff --git a/third_party/xla/xla/python/py_array.h b/third_party/xla/xla/python/py_array.h
index 2824fd858fb24d..75df6229ee59a1 100644
--- a/third_party/xla/xla/python/py_array.h
+++ b/third_party/xla/xla/python/py_array.h
@@ -63,8 +63,8 @@ class PyHostValue {
   PyHostValue& operator=(const PyHostValue&) = delete;
   PyHostValue& operator=(PyHostValue&&) = delete;
 
-  Status CopyToHostAsync(std::optional<Shape>& dynamic_shape_holder,
-                         ifrt::Array* ifrt_array);
+  absl::Status CopyToHostAsync(std::optional<Shape>& dynamic_shape_holder,
+                               ifrt::Array* ifrt_array);
 
   absl::StatusOr<nanobind::object> AsNumPyArray(
       std::optional<Shape>& dynamic_shape_holder, ifrt::Array* ifrt_array);
@@ -152,7 +152,7 @@ class PyArray : public nanobind::object {
       tsl::RCReference<ifrt::Array> ifrt_array, nanobind::object sharding,
       bool weak_type, bool committed, bool skip_checks);
 
-  static Status RegisterTypes(nanobind::module_& m);
+  static absl::Status RegisterTypes(nanobind::module_& m);
 
   using Storage = PyArray_Storage;
 
@@ -238,7 +238,7 @@ class PyArray : public nanobind::object {
   const std::vector<PyArray>& py_arrays_cached();
 
   nanobind::object arrays();
-  Status set_arrays(nanobind::object obj);
+  absl::Status set_arrays(nanobind::object obj);
   absl::StatusOr<PyArray> FullyReplicatedShard();
 
   int num_shards() const {
@@ -258,17 +258,17 @@ class PyArray : public nanobind::object {
     return arg.type().is(PyArray::type());
   }
 
-  Status BlockUntilReady() const;
+  absl::Status BlockUntilReady() const;
 
   absl::Status BlockUntilResultStatusIsReady();
 
   absl::StatusOr<size_t> GetOnDeviceSizeInBytes();
   absl::StatusOr<nanobind::object> SingleDeviceArrayToNumpyArray();
-  Status CopySingleDeviceArrayToHostAsync();
+  absl::Status CopySingleDeviceArrayToHostAsync();
   nanobind::dict CudaArrayInterface();
   absl::StatusOr<std::uintptr_t> UnsafeBufferPointer();
 
-  Status Delete();
+  absl::Status Delete();
 
   bool IsDeleted() const;
 
diff --git a/third_party/xla/xla/python/py_values.h b/third_party/xla/xla/python/py_values.h
index ef889372e0ff7c..7f19ba1fbb71c0 100644
--- a/third_party/xla/xla/python/py_values.h
+++ b/third_party/xla/xla/python/py_values.h
@@ -72,7 +72,7 @@ struct DevicePutResult {
 // to execute.
 //
 // May throw exceptions from nanobind in addition to failing via an error
-// Status. (We could catch these if needed, but there seems little point.)
+// absl::Status. (We could catch these if needed, but there seems little point.)
 struct DevicePutOptions {
   bool squash_64bit_types = false;
   bool allow_zero_copy = true;
diff --git a/third_party/xla/xla/python/util.cc b/third_party/xla/xla/python/util.cc
index d03de01d88f341..724f71772ed213 100644
--- a/third_party/xla/xla/python/util.cc
+++ b/third_party/xla/xla/python/util.cc
@@ -29,7 +29,7 @@ limitations under the License.
 
 namespace xla {
 
-Status AwaitBuffersReady(absl::Span<ifrt::Array* const> ifrt_arrays) {
+absl::Status AwaitBuffersReady(absl::Span<ifrt::Array* const> ifrt_arrays) {
   ifrt::Future<> future;
   if (ifrt_arrays.size() == 1) {
     future = ifrt_arrays[0]->GetReadyFuture();
@@ -42,7 +42,7 @@ Status AwaitBuffersReady(absl::Span<ifrt::Array* const> ifrt_arrays) {
     future = ifrt::JoinFutures(absl::MakeSpan(futures));
   }
 
-  Status s = future.Await();
+  absl::Status s = future.Await();
   if (!s.ok()) {
     // Fix up error string because some clients rely on it.
     if (s.message() == "GetReadyFuture() called on deleted or donated buffer") {
diff --git a/third_party/xla/xla/python/util.h b/third_party/xla/xla/python/util.h
index 3e9f26d2bbd086..20977ffa0be36c 100644
--- a/third_party/xla/xla/python/util.h
+++ b/third_party/xla/xla/python/util.h
@@ -27,7 +27,7 @@ namespace xla {
 
 // Requests if given buffers are ready, awaits for results and returns OK if
 // all of the buffers are ready or the last non-ok status.
-Status AwaitBuffersReady(absl::Span<ifrt::Array* const> ifrt_arrays);
+absl::Status AwaitBuffersReady(absl::Span<ifrt::Array* const> ifrt_arrays);
 
 }  // namespace xla
 
diff --git a/third_party/xla/xla/python/xplane_to_profile_instructions.cc b/third_party/xla/xla/python/xplane_to_profile_instructions.cc
index 772e1821a82a6b..07137cd1d923b2 100644
--- a/third_party/xla/xla/python/xplane_to_profile_instructions.cc
+++ b/third_party/xla/xla/python/xplane_to_profile_instructions.cc
@@ -177,7 +177,7 @@ void GetXPlaneHloModuleInfo(
 
 }  // namespace
 
-Status ConvertXplaneUnderLogdirToProfiledInstructionsProto(
+absl::Status ConvertXplaneUnderLogdirToProfiledInstructionsProto(
     const std::string& logdir, tensorflow::profiler::ProfiledInstructionsProto*
                                    profiled_instructions_proto) {
   // Find the xplane files for each host under logdir.
@@ -202,7 +202,7 @@ Status ConvertXplaneUnderLogdirToProfiledInstructionsProto(
                                                   profiled_instructions_proto);
 }
 
-Status ConvertXplaneToProfiledInstructionsProto(
+absl::Status ConvertXplaneToProfiledInstructionsProto(
     std::vector<tensorflow::profiler::XSpace> xspaces,
     tensorflow::profiler::ProfiledInstructionsProto*
         profiled_instructions_proto) {
diff --git a/third_party/xla/xla/python/xplane_to_profile_instructions.h b/third_party/xla/xla/python/xplane_to_profile_instructions.h
index 8971dc651030e8..b32bfd2fbb9220 100644
--- a/third_party/xla/xla/python/xplane_to_profile_instructions.h
+++ b/third_party/xla/xla/python/xplane_to_profile_instructions.h
@@ -36,14 +36,14 @@ struct HloLatencyInfo {
 
 // Convert XSpace to ProfiledInstructionsProto. This function will aggregate
 // all the xplane.pb info into ProfiledInstructionsProto.
-Status ConvertXplaneToProfiledInstructionsProto(
+absl::Status ConvertXplaneToProfiledInstructionsProto(
     std::vector<tensorflow::profiler::XSpace> xspaces,
     tensorflow::profiler::ProfiledInstructionsProto*
         profiled_instructions_proto);
 
 // Convert XSpace to ProfiledInstructionsProto. This function will aggregate
 // all the xplane.pb info under logdir into ProfiledInstructionsProto.
-Status ConvertXplaneUnderLogdirToProfiledInstructionsProto(
+absl::Status ConvertXplaneUnderLogdirToProfiledInstructionsProto(
     const std::string& logdir, tensorflow::profiler::ProfiledInstructionsProto*
                                    profiled_instructions_proto);
 

From 190d56b4c1a21829bf69e23ac7c3ab6875137f16 Mon Sep 17 00:00:00 2001
From: Weiyi Wang <weiyiw@google.com>
Date: Mon, 20 May 2024 11:12:10 -0700
Subject: [PATCH 423/478] Temporarily disable xnnpack cache test on Android.

PiperOrigin-RevId: 635513547
---
 tensorflow/lite/delegates/xnnpack/BUILD | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/tensorflow/lite/delegates/xnnpack/BUILD b/tensorflow/lite/delegates/xnnpack/BUILD
index a2841aecca92cd..46c327528076e8 100644
--- a/tensorflow/lite/delegates/xnnpack/BUILD
+++ b/tensorflow/lite/delegates/xnnpack/BUILD
@@ -2850,6 +2850,9 @@ cc_test(
 cc_test(
     name = "weight_cache_test",
     srcs = ["weight_cache_test.cc"],
+    tags = [
+        "tflite_disable_mobile_test",  # TODO - b/341104412 - enable on Android
+    ],
     deps = [
         ":test_main",
         ":weight_cache",

From 8f3b195bf51ee99bcab42f8eaff44b7bb268e220 Mon Sep 17 00:00:00 2001
From: Ce Zheng <zce@google.com>
Date: Mon, 20 May 2024 11:15:48 -0700
Subject: [PATCH 424/478] [XLA] Support transposing iota tile assignment with
 tile dimensions combining multiple reshape dimensions.

PiperOrigin-RevId: 635514612
---
 third_party/xla/xla/hlo/ir/tile_assignment.cc | 67 +++++++++++++++----
 .../xla/xla/tests/tile_assignment_test.cc     | 27 ++++++++
 2 files changed, 81 insertions(+), 13 deletions(-)

diff --git a/third_party/xla/xla/hlo/ir/tile_assignment.cc b/third_party/xla/xla/hlo/ir/tile_assignment.cc
index e3bd49952c198b..fe82af0a35c416 100644
--- a/third_party/xla/xla/hlo/ir/tile_assignment.cc
+++ b/third_party/xla/xla/hlo/ir/tile_assignment.cc
@@ -293,21 +293,62 @@ std::optional<IotaTileAssignment> IotaTileAssignment::Transpose(
     CHECK_EQ(reshape_ndims_, new_perm.size());
     return IotaTileAssignment::Create(new_dims, reshape_dims, new_perm);
   }
-  absl::InlinedVector<int, 6> split_transpose_perm;
-  if (TryDecanonicalize(non_one_dims, reshape_dims, transpose_perm,
-                        split_transpose_perm)) {
-    absl::InlinedVector<int, 6> new_perm;
-    new_perm.reserve(non_one_dims.size());
-    for (int i = 0; i < ndims_; ++i) {
-      if (dims[perm[i]] == 1) continue;
-      new_perm.push_back(split_transpose_perm[one_to_non_one[perm[i]]]);
+  if (reshape_ndims_ > non_one_dims.size()) {
+    // If there are fewer non-one tile dimensions than reshape dimensions,
+    // try grouping reshape dimensions together to see if they form the
+    // identical tile dimensions, then transpose them in groups.
+    absl::InlinedVector<absl::InlinedVector<int, 2>, 6> grouped_reshape_dims(
+        non_one_dims.size());
+    int transpose_perm_idx = 0;
+    for (int i = 0, n = non_one_dims.size();
+         i < n && transpose_perm_idx < reshape_ndims_; ++i) {
+      int reshape_dim_idx = transpose_perm[transpose_perm_idx];
+      int64_t cand = reshape_dims[reshape_dim_idx];
+      int64_t target = non_one_dims[i];
+      while (target % cand == 0) {
+        target /= cand;
+        grouped_reshape_dims[i].push_back(reshape_dim_idx);
+        if (++transpose_perm_idx >= reshape_ndims_) {
+          break;
+        }
+        reshape_dim_idx = transpose_perm[transpose_perm_idx];
+        cand = reshape_dims[reshape_dim_idx];
+      }
+      if (target != 1) {
+        return std::nullopt;
+      }
     }
-    absl::InlinedVector<int64_t, 6> new_reshape_dims(
-        split_transpose_perm.size());
-    for (int i = 0; i < non_one_dims.size(); ++i) {
-      new_reshape_dims[split_transpose_perm[i]] = non_one_dims[i];
+    absl::InlinedVector<int, 6> flattened_transpose_perm;
+    flattened_transpose_perm.reserve(reshape_ndims_);
+    for (int i = 0; i < perm.size(); ++i) {
+      const int dim = perm[i];
+      if (one_to_non_one[dim] < 0) {
+        continue;
+      }
+      auto& group = grouped_reshape_dims[one_to_non_one[dim]];
+      flattened_transpose_perm.insert(flattened_transpose_perm.end(),
+                                      group.begin(), group.end());
+    }
+    CHECK_EQ(flattened_transpose_perm.size(), reshape_ndims_);
+    return IotaTileAssignment::Create(new_dims, reshape_dims,
+                                      flattened_transpose_perm);
+  } else {
+    absl::InlinedVector<int, 6> split_transpose_perm;
+    if (TryDecanonicalize(non_one_dims, reshape_dims, transpose_perm,
+                          split_transpose_perm)) {
+      absl::InlinedVector<int, 6> new_perm;
+      new_perm.reserve(non_one_dims.size());
+      for (int i = 0; i < ndims_; ++i) {
+        if (dims[perm[i]] == 1) continue;
+        new_perm.push_back(split_transpose_perm[one_to_non_one[perm[i]]]);
+      }
+      absl::InlinedVector<int64_t, 6> new_reshape_dims(
+          split_transpose_perm.size());
+      for (int i = 0; i < non_one_dims.size(); ++i) {
+        new_reshape_dims[split_transpose_perm[i]] = non_one_dims[i];
+      }
+      return IotaTileAssignment::Create(new_dims, new_reshape_dims, new_perm);
     }
-    return IotaTileAssignment::Create(new_dims, new_reshape_dims, new_perm);
   }
   // TODO(b/341371396): Handle remaining patterns and remove nullopt path.
   return std::nullopt;
diff --git a/third_party/xla/xla/tests/tile_assignment_test.cc b/third_party/xla/xla/tests/tile_assignment_test.cc
index a0dc21b1f38599..728d48d87c6f04 100644
--- a/third_party/xla/xla/tests/tile_assignment_test.cc
+++ b/third_party/xla/xla/tests/tile_assignment_test.cc
@@ -322,6 +322,33 @@ TEST_P(FormattedTileAssignmentTest,
           175, 190, 191, 206, 207, 222, 223, 238, 239, 254, 255));
 }
 
+TEST_P(FormattedTileAssignmentTest,
+       TransposeIotaTileGroupingCanonicalizedReshapeDims) {
+  TileAssignment tile({1, 4, 16}, {4, 4, 4}, {1, 0, 2});
+  if (ShouldConvertToV1()) {
+    tile = TileAssignment(tile.shared_array());
+  }
+  TileAssignment xposed = tile.Transpose({2, 0, 1});
+  EXPECT_NE(xposed, tile);
+  EXPECT_EQ(xposed, TileAssignment({16, 1, 4}, {4, 4, 4}, {0, 2, 1}));
+  EXPECT_EQ(xposed.num_dimensions(), 3);
+  EXPECT_EQ(xposed.dim(0), 16);
+  EXPECT_EQ(xposed.dim(1), 1);
+  EXPECT_EQ(xposed.dim(2), 4);
+  EXPECT_EQ(xposed(0, 0, 0), 0);
+  EXPECT_EQ(xposed({7, 0, 3}), 31);
+  EXPECT_EQ(xposed.iota().has_value(), !ShouldConvertToV1());
+  EXPECT_TRUE(xposed.UsesDevice(0));
+  EXPECT_TRUE(xposed.UsesDevice(63));
+  EXPECT_FALSE(xposed.UsesDevice(64));
+  EXPECT_THAT(ToVectorUsingEach(xposed),
+              ElementsAre(0, 4, 8, 12, 1, 5, 9, 13, 2, 6, 10, 14, 3, 7, 11, 15,
+                          16, 20, 24, 28, 17, 21, 25, 29, 18, 22, 26, 30, 19,
+                          23, 27, 31, 32, 36, 40, 44, 33, 37, 41, 45, 34, 38,
+                          42, 46, 35, 39, 43, 47, 48, 52, 56, 60, 49, 53, 57,
+                          61, 50, 54, 58, 62, 51, 55, 59, 63));
+}
+
 TEST_P(FormattedTileAssignmentTest, TransposeNoopIotaTile) {
   TileAssignment tile({4, 4}, {4, 4}, {1, 0});
   if (ShouldConvertToV1()) {

From 6fda9f2f3cf305ff2fc3f1f3b07cbb11890fd395 Mon Sep 17 00:00:00 2001
From: Dan Foreman-Mackey <danfm@google.com>
Date: Mon, 20 May 2024 11:39:46 -0700
Subject: [PATCH 425/478] Add missing parentheses in XLA custom calls doc.

When reading through custom_call.md, I noticed these two missing parentheses.

PiperOrigin-RevId: 635522222
---
 third_party/xla/docs/custom_call.md | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/third_party/xla/docs/custom_call.md b/third_party/xla/docs/custom_call.md
index 7f4bf765b88b67..ef56830e08c764 100644
--- a/third_party/xla/docs/custom_call.md
+++ b/third_party/xla/docs/custom_call.md
@@ -76,7 +76,7 @@ XLA_FFI_DEFINE_HANDLER(handler, do_custom_call,
                        ffi::Ffi::Bind()
                            .Arg<Buffer>()
                            .Arg<Buffer>()
-                           .Ret<Buffer>();
+                           .Ret<Buffer>());
 
 // Registers `handler` with and XLA FFI on a "Host" platform.
 XLA_FFI_REGISTER_HANDLER(xla::ffi::GetXlaFfiApi(), "do_custom_call",
@@ -118,7 +118,7 @@ XLA_FFI_DEFINE_HANDLER(handler, do_custom_call,
                            .Ctx<xla::ffi::PlatformStream<CUstream>>()
                            .Arg<Buffer>()
                            .Arg<Buffer>()
-                           .Ret<Buffer>();
+                           .Ret<Buffer>());
 
 XLA_FFI_REGISTER_HANDLER(xla::ffi::GetXlaFfiApi(), "do_custom_call",
                          "CUDA", handler);

From c67475d9bb225d0472a60f4d4bca25e5013f3b03 Mon Sep 17 00:00:00 2001
From: Kyle Lucke <klucke@google.com>
Date: Mon, 20 May 2024 11:49:30 -0700
Subject: [PATCH 426/478] Replace the use of xla::OkStatus with absl::OkStatus
 now that they're the same.

PiperOrigin-RevId: 635524889
---
 .../xla/xla/service/cpu/cpu_executable.cc     |  2 +-
 .../xla/service/cpu/cpu_layout_assignment.cc  |  2 +-
 third_party/xla/xla/service/cpu/cpu_xfeed.cc  | 16 +--
 .../xla/xla/service/cpu/dot_op_emitter.cc     | 14 +--
 third_party/xla/xla/service/cpu/ir_emitter.cc | 98 +++++++++----------
 .../xla/xla/service/cpu/ir_function.cc        |  2 +-
 .../xla/xla/service/cpu/mlir_emitter.cc       |  2 +-
 .../xla/service/cpu/onednn_matmul_rewriter.cc | 34 ++++---
 .../xla/service/cpu/onednn_ops_rewriter.cc    | 19 ++--
 .../memory_space_assignment/algorithm.cc      | 10 +-
 .../memory_space_assignment/allocation.cc     | 14 +--
 .../memory_space_assignment/allocation.h      |  8 +-
 .../cost_analysis_test.cc                     |  2 +-
 .../memory_bound_loop_optimizer.cc            |  2 +-
 .../memory_bound_loop_optimizer_test.cc       |  4 +-
 .../memory_space_assignment.cc                | 20 ++--
 .../memory_space_assignment_test.cc           | 20 ++--
 17 files changed, 136 insertions(+), 133 deletions(-)

diff --git a/third_party/xla/xla/service/cpu/cpu_executable.cc b/third_party/xla/xla/service/cpu/cpu_executable.cc
index 02fb083554ce1c..c6c9337a8b94b6 100644
--- a/third_party/xla/xla/service/cpu/cpu_executable.cc
+++ b/third_party/xla/xla/service/cpu/cpu_executable.cc
@@ -243,7 +243,7 @@ absl::Status CpuExecutable::ExecuteComputeFunction(
     return Internal("CustomCall failed: %s", *error_message);
   }
 
-  return OkStatus();
+  return absl::OkStatus();
 }
 
 absl::StatusOr<ExecutionOutput> CpuExecutable::CreateResultShapedBuffer(
diff --git a/third_party/xla/xla/service/cpu/cpu_layout_assignment.cc b/third_party/xla/xla/service/cpu/cpu_layout_assignment.cc
index 48f5ae59bfe6c3..d4907a88f364fc 100644
--- a/third_party/xla/xla/service/cpu/cpu_layout_assignment.cc
+++ b/third_party/xla/xla/service/cpu/cpu_layout_assignment.cc
@@ -185,7 +185,7 @@ absl::Status CpuLayoutAssignment::AddBackendConstraints(
       }
     }
   }
-  return OkStatus();
+  return absl::OkStatus();
 }
 }  // namespace cpu
 }  // namespace xla
diff --git a/third_party/xla/xla/service/cpu/cpu_xfeed.cc b/third_party/xla/xla/service/cpu/cpu_xfeed.cc
index 8b4ea238f518bc..332ce3b1c126fb 100644
--- a/third_party/xla/xla/service/cpu/cpu_xfeed.cc
+++ b/third_party/xla/xla/service/cpu/cpu_xfeed.cc
@@ -111,7 +111,7 @@ absl::Status TransferBufferToInfeed(int device_ordinal, int64_t size,
       cpu::runtime::GetXfeedManager(device_ordinal);
   xfeed_manager->infeed()->EnqueueBuffersAtomically({buffer});
 
-  return OkStatus();
+  return absl::OkStatus();
 }
 
 absl::StatusOr<Shape> TransferBuffersFromOutfeedInternal(
@@ -218,7 +218,7 @@ absl::Status TransferLiteralToInfeedOnCpu(int device_ordinal,
   xfeed_manager->infeed()->EnqueueBuffersAtomically(buffers);
 
   std::move(cleanup).Cancel();
-  return OkStatus();
+  return absl::OkStatus();
 }
 
 absl::Status TransferLiteralFromOutfeedOnCpu(int device_ordinal,
@@ -242,7 +242,7 @@ absl::Status TransferLiteralFromOutfeedOnCpu(int device_ordinal,
     TF_RET_CHECK(size == cpu::runtime::GetByteSizeRequirement(received_shape,
                                                               sizeof(void*)));
     *literal.mutable_shape_do_not_use() = received_shape;
-    return OkStatus();
+    return absl::OkStatus();
   }
 
   if (ShapeUtil::IsNestedTuple(literal.shape())) {
@@ -272,7 +272,7 @@ absl::Status TransferLiteralFromOutfeedOnCpu(int device_ordinal,
       cpu::runtime::GetByteSizeRequirement(received_shape, sizeof(void*)));
 
   TF_RET_CHECK(ShapeUtil::Equal(literal.shape(), literal.shape()));
-  return OkStatus();
+  return absl::OkStatus();
 }
 
 absl::Status ReadDynamicShapesOnCpu(
@@ -286,12 +286,12 @@ absl::Status ReadDynamicShapesOnCpu(
         const Shape& buffer_shape =
             ShapeUtil::GetSubshape(*device_shape, index);
         if (buffer_shape.IsTuple()) {
-          return OkStatus();
+          return absl::OkStatus();
         }
         Shape& device_sub_shape =
             *ShapeUtil::GetMutableSubshape(device_shape, index);
         if (device_sub_shape.is_static()) {
-          return OkStatus();
+          return absl::OkStatus();
         }
         const void* memory = buffer.opaque();
 
@@ -310,12 +310,12 @@ absl::Status ReadDynamicShapesOnCpu(
         for (int64_t i = 0; i < device_sub_shape.rank(); ++i) {
           device_sub_shape.mutable_dimensions()[i] = metadata_buffer[i];
         }
-        return OkStatus();
+        return absl::OkStatus();
       }));
   device_shape->clear_dynamic_dimensions();
 
   TF_RET_CHECK(ShapeUtil::DynamicShapeIsCompatible(*device_shape,
                                                    original_device_shape));
-  return OkStatus();
+  return absl::OkStatus();
 }
 }  // namespace xla
diff --git a/third_party/xla/xla/service/cpu/dot_op_emitter.cc b/third_party/xla/xla/service/cpu/dot_op_emitter.cc
index 9548b9919ac722..ef06f3bb08e338 100644
--- a/third_party/xla/xla/service/cpu/dot_op_emitter.cc
+++ b/third_party/xla/xla/service/cpu/dot_op_emitter.cc
@@ -566,15 +566,15 @@ absl::Status DotOpEmitter::Emit() {
                                        target_machine_features_)) {
     case DotImplementationStrategy::kNaiveLlvmIr:
       EmitNaiveLlvmIrGemm();
-      return OkStatus();
+      return absl::OkStatus();
 
     case DotImplementationStrategy::kTiledLlvmIrGemv:
       EmitTiledLlvmIrGemv();
-      return OkStatus();
+      return absl::OkStatus();
 
     case DotImplementationStrategy::kTiledLlvmIrGemm:
       EmitTiledLlvmIrGemm();
-      return OkStatus();
+      return absl::OkStatus();
 
     case DotImplementationStrategy::kLinalgMatmul:
       return EmitLinalgMatmul();
@@ -788,7 +788,7 @@ absl::Status DotOpEmitter::EmitScalarDot() {
     result = b_->CreateFMul(lhs_value, rhs_value);
   }
   target_array_.EmitWriteArrayElement(/*index=*/element_index, result, b_);
-  return OkStatus();
+  return absl::OkStatus();
 }
 
 absl::Status DotOpEmitter::EmitCallToRuntime() {
@@ -899,7 +899,7 @@ absl::Status DotOpEmitter::EmitCallToRuntime() {
                   b_->getInt64(mat_mult_dims.m), b_->getInt64(mat_mult_dims.n),
                   b_->getInt64(mat_mult_dims.k), b_->getInt32(transpose_lhs),
                   b_->getInt32(transpose_rhs)});
-  return OkStatus();
+  return absl::OkStatus();
 }
 
 absl::Status DotOpEmitter::EmitCallToBatchRuntime() {
@@ -982,7 +982,7 @@ absl::Status DotOpEmitter::EmitCallToBatchRuntime() {
        b_->getInt64(mat_mult_dims.k), b_->getInt64(lhs_shape.dimensions(0)),
        b_->getInt32(static_cast<uint32_t>(transpose_lhs)),
        b_->getInt32(static_cast<uint32_t>(transpose_rhs))});
-  return OkStatus();
+  return absl::OkStatus();
 }
 
 DotOpEmitter::MatMultDims DotOpEmitter::GetMatMultDims() const {
@@ -1282,7 +1282,7 @@ absl::Status ValidateDotDimensionNumbers(
       absl::c_equal(batch_dim_numbers, dim_numbers.lhs_batch_dimensions()));
   TF_RET_CHECK(
       absl::c_equal(batch_dim_numbers, dim_numbers.rhs_batch_dimensions()));
-  return OkStatus();
+  return absl::OkStatus();
 }
 
 // Slice out the inner array at batch index `batch_index` from `outer_array`.
diff --git a/third_party/xla/xla/service/cpu/ir_emitter.cc b/third_party/xla/xla/service/cpu/ir_emitter.cc
index d6a5a832d03ba8..7a0689790f1498 100644
--- a/third_party/xla/xla/service/cpu/ir_emitter.cc
+++ b/third_party/xla/xla/service/cpu/ir_emitter.cc
@@ -258,7 +258,7 @@ void IrEmitter::InitializeIrFunction(const std::string& function_name) {
 absl::Status IrEmitter::HandleBitcast(HloInstruction* bitcast) {
   VLOG(2) << "HandleBitcast: " << bitcast->ToString();
   emitted_value_[bitcast] = GetEmittedValueFor(bitcast->operand(0));
-  return OkStatus();
+  return absl::OkStatus();
 }
 
 llvm::Constant* IrEmitter::EmitGlobalForLiteral(const Literal& literal) {
@@ -297,7 +297,7 @@ absl::Status IrEmitter::EmitConstantGlobals() {
                 global_for_const);
   }
 
-  return OkStatus();
+  return absl::OkStatus();
 }
 
 absl::Status IrEmitter::HandleConstant(HloInstruction* constant) {
@@ -393,7 +393,7 @@ absl::Status IrEmitter::HandleGetTupleElement(
   emitted_value_[get_tuple_element] = llvm_ir::EmitGetTupleElement(
       shape, get_tuple_element->tuple_index(), MinimumAlignmentForShape(shape),
       GetEmittedValueFor(operand), IrShapeType(operand->shape()), &b_);
-  return OkStatus();
+  return absl::OkStatus();
 }
 
 absl::Status IrEmitter::HandleSelect(HloInstruction* select) {
@@ -459,7 +459,7 @@ absl::Status IrEmitter::HandleInfeed(HloInstruction* instruction) {
         EmitXfeedTransfer(XfeedKind::kInfeed, data_shape, data_address));
   }
 
-  return OkStatus();
+  return absl::OkStatus();
 }
 
 absl::Status IrEmitter::EmitXfeedTransfer(XfeedKind kind, const Shape& shape,
@@ -524,7 +524,7 @@ absl::Status IrEmitter::EmitXfeedTransfer(XfeedKind kind, const Shape& shape,
                   acquired_pointer, shape_ptr, b_.getInt32(shape_length)},
                  b_.getVoidTy());
 
-  return OkStatus();
+  return absl::OkStatus();
 }
 
 absl::Status IrEmitter::HandleOutfeed(HloInstruction* outfeed) {
@@ -554,7 +554,7 @@ absl::Status IrEmitter::HandleOutfeed(HloInstruction* outfeed) {
                                          tuple_element_shape, tuple_element));
   }
 
-  return OkStatus();
+  return absl::OkStatus();
 }
 
 absl::Status IrEmitter::HandleSort(HloInstruction* hlo) {
@@ -649,7 +649,7 @@ absl::Status IrEmitter::HandleSort(HloInstruction* hlo) {
   if (sort->values_count() > 0) {
     llvm_ir::EmitTuple(GetIrArrayFor(sort), destination_addresses, &b_);
   }
-  return OkStatus();
+  return absl::OkStatus();
 }
 
 absl::Status IrEmitter::HandleTuple(HloInstruction* tuple) {
@@ -659,7 +659,7 @@ absl::Status IrEmitter::HandleTuple(HloInstruction* tuple) {
     base_ptrs.push_back(GetEmittedValueFor(operand));
   }
   llvm_ir::EmitTuple(GetIrArrayFor(tuple), base_ptrs, &b_);
-  return OkStatus();
+  return absl::OkStatus();
 }
 
 absl::Status IrEmitter::HandleReduceWindow(HloInstruction* reduce_window) {
@@ -863,7 +863,7 @@ absl::Status IrEmitter::HandleSelectAndScatter(
   output_array.EmitWriteArrayElement(selected_index, scatter_value, &b_);
 
   SetToFirstInsertPoint(source_loops.GetOuterLoopExitBasicBlock(), &b_);
-  return OkStatus();
+  return absl::OkStatus();
 }
 
 absl::Status IrEmitter::HandleDot(HloInstruction* dot) {
@@ -1075,7 +1075,7 @@ absl::Status IrEmitter::HandleConvolution(HloInstruction* convolution) {
       EmitCallToFunc(fn_name, args, b_.getVoidTy(), /*does_not_throw=*/true,
                      /*only_accesses_arg_memory=*/true);
 
-      return OkStatus();
+      return absl::OkStatus();
     }
   }
   // This is a completely un-optimized version of convolution just to
@@ -1136,7 +1136,7 @@ absl::Status IrEmitter::HandleFft(HloInstruction* fft) {
                  /*only_accesses_arg_memory=*/false,
                  /*only_accesses_inaccessible_mem_or_arg_mem=*/true);
 
-  return OkStatus();
+  return absl::OkStatus();
 }
 
 absl::Status IrEmitter::HandleAllReduceSingleReplica(HloInstruction* crs) {
@@ -1171,7 +1171,7 @@ absl::Status IrEmitter::HandleAllReduceSingleReplica(HloInstruction* crs) {
            /*SrcAlign=*/llvm::Align(1), ShapeUtil::ByteSizeOf(operand_shape));
   }
   llvm_ir::EmitTuple(GetIrArrayFor(crs), operand_ptrs, &b_);
-  return OkStatus();
+  return absl::OkStatus();
 }
 
 // Data types supported by ReduceScatter and AllReduce.
@@ -1279,7 +1279,7 @@ absl::Status IrEmitter::HandleAllReduceMultipleReplica(HloInstruction* crs) {
        /*output_buffers=*/output_buffers},
       b_.getVoidTy());
 
-  return OkStatus();
+  return absl::OkStatus();
 }
 
 absl::Status IrEmitter::HandleAllReduce(HloInstruction* crs) {
@@ -1343,7 +1343,7 @@ absl::Status IrEmitter::HandleReduceScatter(HloInstruction* rs) {
        /*output_buffer=*/output_buffer},
       b_.getVoidTy());
 
-  return OkStatus();
+  return absl::OkStatus();
 }
 
 absl::Status IrEmitter::HandleAllToAll(HloInstruction* instruction) {
@@ -1399,7 +1399,7 @@ absl::Status IrEmitter::HandleAllToAll(HloInstruction* instruction) {
                  b_.getVoidTy());
 
   llvm_ir::EmitTuple(GetIrArrayFor(instruction), output_buffer_ptrs, &b_);
-  return OkStatus();
+  return absl::OkStatus();
 }
 
 absl::Status IrEmitter::HandleAllGather(HloInstruction* instruction) {
@@ -1450,7 +1450,7 @@ absl::Status IrEmitter::HandleAllGather(HloInstruction* instruction) {
       b_.getVoidTy());
 
   llvm_ir::EmitTuple(GetIrArrayFor(instruction), output_buffer_ptrs, &b_);
-  return OkStatus();
+  return absl::OkStatus();
 }
 
 absl::Status IrEmitter::HandleCollectivePermute(HloInstruction* crs) {
@@ -1487,7 +1487,7 @@ absl::Status IrEmitter::HandleCollectivePermute(HloInstruction* crs) {
        /*source_target_pairs_size=*/b_.getInt32(source_target_pairs.size())},
       b_.getVoidTy());
 
-  return OkStatus();
+  return absl::OkStatus();
 }
 
 absl::Status IrEmitter::HandlePartitionId(HloInstruction* hlo) {
@@ -1499,7 +1499,7 @@ absl::Status IrEmitter::HandlePartitionId(HloInstruction* hlo) {
                  {/*run_options=*/GetExecutableRunOptionsArgument(),
                   /*output_buffer=*/output_buffer},
                  b_.getVoidTy());
-  return OkStatus();
+  return absl::OkStatus();
 }
 
 absl::Status IrEmitter::HandleReplicaId(HloInstruction* hlo) {
@@ -1511,7 +1511,7 @@ absl::Status IrEmitter::HandleReplicaId(HloInstruction* hlo) {
                  {/*run_options=*/GetExecutableRunOptionsArgument(),
                   /*output_buffer=*/output_buffer},
                  b_.getVoidTy());
-  return OkStatus();
+  return absl::OkStatus();
 }
 
 absl::Status IrEmitter::HandleParameter(HloInstruction* parameter) {
@@ -2019,7 +2019,7 @@ absl::Status IrEmitter::HandleReduce(HloInstruction* reduce) {
     if (vectorization_successful) {
       VLOG(1) << "Successfully vectorized reduction " << reduce->ToString()
               << "\n";
-      return OkStatus();
+      return absl::OkStatus();
     } else {
       VLOG(1) << "Could not vectorize reduction " << reduce->ToString() << ": "
               << vectorization_failure_reason;
@@ -2060,7 +2060,7 @@ absl::Status IrEmitter::HandleSlice(HloInstruction* slice) {
   TF_RETURN_IF_ERROR(EmitTargetAddressForOp(slice));
 
   if (ShapeUtil::IsZeroElementArray(slice->shape())) {
-    return OkStatus();
+    return absl::OkStatus();
   }
 
   const Layout& layout = operand->shape().layout();
@@ -2177,7 +2177,7 @@ absl::Status IrEmitter::HandleSlice(HloInstruction* slice) {
     SetToFirstInsertPoint(loops.GetOuterLoopExitBasicBlock(), &b_);
   }
 
-  return OkStatus();
+  return absl::OkStatus();
 }
 
 absl::Status IrEmitter::HandleDynamicSlice(HloInstruction* dynamic_slice) {
@@ -2271,7 +2271,7 @@ absl::Status IrEmitter::HandlePad(HloInstruction* pad) {
   output_array.EmitWriteArrayElement(output_index, operand_data, &b_);
 
   SetToFirstInsertPoint(loops.GetOuterLoopExitBasicBlock(), &b_);
-  return OkStatus();
+  return absl::OkStatus();
 }
 
 absl::Status IrEmitter::HandleFusion(HloInstruction* fusion) {
@@ -2323,7 +2323,7 @@ absl::Status IrEmitter::HandleFusion(HloInstruction* fusion) {
         *dot, target_array, lhs_array, rhs_array, &addend_array,
         GetExecutableRunOptionsArgument(), &b_, mlir_context_,
         hlo_module_config_, target_machine_features_));
-    return OkStatus();
+    return absl::OkStatus();
   } else {
     return Unimplemented("Fusion kind not implemented on CPU");
   }
@@ -2371,7 +2371,7 @@ absl::Status IrEmitter::HandleCall(HloInstruction* call) {
     EmitGlobalCall(*computation, computation->name());
   }
 
-  return OkStatus();
+  return absl::OkStatus();
 }
 
 absl::Status IrEmitter::HandleSliceToDynamic(HloInstruction* hlo) {
@@ -2410,7 +2410,7 @@ absl::Status IrEmitter::HandleSliceToDynamic(HloInstruction* hlo) {
     llvm_ir::IrArray::Index dest_index(linear_index, data_array.GetShape(),
                                        &b_);
     data_array.EmitWriteArrayElement(dest_index, source_element, &b_);
-    return OkStatus();
+    return absl::OkStatus();
   };
   return llvm_ir::LoopEmitter(loop_body_emitter, data_array.GetShape(),
                               dynamic_dims, &b_)
@@ -2471,7 +2471,7 @@ absl::Status IrEmitter::HandlePadToStatic(HloInstruction* hlo) {
     llvm::Value* source_element =
         GetIrArrayFor(hlo->operand(0)).EmitReadArrayElement(source_index, &b_);
     data_array.EmitWriteArrayElement(array_index, source_element, &b_);
-    return OkStatus();
+    return absl::OkStatus();
   };
   TF_RETURN_IF_ERROR(
       llvm_ir::LoopEmitter(loop_body_emitter, input_shape, dynamic_dims, &b_)
@@ -2479,7 +2479,7 @@ absl::Status IrEmitter::HandlePadToStatic(HloInstruction* hlo) {
 
   // Emit static tensor and dynamic sizes as one tuple.
   llvm_ir::EmitTuple(GetIrArrayFor(hlo), tuple_operand_ptrs, &b_);
-  return OkStatus();
+  return absl::OkStatus();
 }
 
 absl::Status IrEmitter::HandleTopK(HloInstruction* hlo) {
@@ -2518,7 +2518,7 @@ absl::Status IrEmitter::HandleTopK(HloInstruction* hlo) {
 
   llvm_ir::EmitTuple(GetIrArrayFor(hlo), {out_values_ptr, out_indices_ptr},
                      &b_);
-  return OkStatus();
+  return absl::OkStatus();
 }
 
 #if defined(INTEL_MKL) && defined(ENABLE_ONEDNN_V3)
@@ -2649,7 +2649,7 @@ absl::Status IrEmitter::HandleOneDnnMatMulCalls(
     scratch_stack_alloca.EmitLifetimeEnd();
   }
 
-  return OkStatus();
+  return absl::OkStatus();
 }
 
 absl::Status IrEmitter::HandleOneDnnLayerNorm(HloInstruction* custom_call) {
@@ -2726,7 +2726,7 @@ absl::Status IrEmitter::HandleOneDnnLayerNorm(HloInstruction* custom_call) {
   b_.CreateLifetimeEnd(args_ptr, b_.getInt64(-1));
   result_stack_alloca.EmitLifetimeEnd();
 
-  return OkStatus();
+  return absl::OkStatus();
 }
 
 absl::Status IrEmitter::HandleOneDnnSoftmax(HloInstruction* custom_call) {
@@ -2749,7 +2749,7 @@ absl::Status IrEmitter::HandleOneDnnSoftmax(HloInstruction* custom_call) {
   input_stack_alloca.EmitLifetimeEnd();
   result_stack_alloca.EmitLifetimeEnd();
 
-  return OkStatus();
+  return absl::OkStatus();
 }
 #endif  // INTEL_MKL && ENABLE_ONEDNN_V3
 
@@ -2855,7 +2855,7 @@ absl::Status IrEmitter::HandleCustomCall(HloInstruction* custom_call) {
           CustomCallApiVersion_Name(typed_custom_call->api_version()));
   }
 
-  return OkStatus();
+  return absl::OkStatus();
 }
 
 absl::Status IrEmitter::HandleWhile(HloInstruction* xla_while) {
@@ -2883,7 +2883,7 @@ absl::Status IrEmitter::HandleWhile(HloInstruction* xla_while) {
                 a->ToString(), slice_a.ToString(), b->ToString(),
                 slice_b.ToString());
           }
-          return OkStatus();
+          return absl::OkStatus();
         };
         TF_RETURN_IF_ERROR(check(xla_while, xla_while->operand(0), index));
         TF_RETURN_IF_ERROR(check(
@@ -2894,7 +2894,7 @@ absl::Status IrEmitter::HandleWhile(HloInstruction* xla_while) {
                   index));
         TF_RETURN_IF_ERROR(check(
             xla_while, xla_while->while_body()->root_instruction(), index));
-        return OkStatus();
+        return absl::OkStatus();
       }));
 
   // Set emitted value to that of 'init' with which it shares an allocation.
@@ -2946,7 +2946,7 @@ absl::Status IrEmitter::HandleWhile(HloInstruction* xla_while) {
   llvm_fn->insert(llvm_fn->end(), exit_bb);
   b_.SetInsertPoint(exit_bb);
 
-  return OkStatus();
+  return absl::OkStatus();
 }
 
 absl::StatusOr<bool> IrEmitter::EmitFastConcatenate(
@@ -3246,7 +3246,7 @@ absl::Status IrEmitter::HandleConcatenate(HloInstruction* concatenate) {
       EmitFastConcatenate(concatenate, operands, &failure_reason));
   if (successful) {
     VLOG(1) << "Emitted fast concatenate for " << concatenate->ToString();
-    return OkStatus();
+    return absl::OkStatus();
   }
 
   VLOG(1) << "Could not emit fast concatenate for " << concatenate->ToString()
@@ -3302,7 +3302,7 @@ absl::Status IrEmitter::HandleConditional(HloInstruction* conditional) {
                    IrName(conditional, "_false"));
 
     SetToFirstInsertPoint(if_data.after_block, &b_);
-    return OkStatus();
+    return absl::OkStatus();
   }
   // We emit a switch statement to LLVM:
   // switch (branch_index) {
@@ -3361,21 +3361,21 @@ absl::Status IrEmitter::HandleConditional(HloInstruction* conditional) {
   }
 
   SetToFirstInsertPoint(after_block, &b_);
-  return OkStatus();
+  return absl::OkStatus();
 }
 
 absl::Status IrEmitter::HandleAfterAll(HloInstruction* after_all) {
   TF_RET_CHECK(ByteSizeOf(after_all->shape()) == 0);
   // No code to generate, but we need to emit an address for book-keeping.
   TF_RETURN_IF_ERROR(EmitTargetAddressForOp(after_all));
-  return OkStatus();
+  return absl::OkStatus();
 }
 
 absl::Status IrEmitter::HandleAddDependency(HloInstruction* add_dependency) {
   // AddDedendency just forwards its zero-th operand.
   emitted_value_[add_dependency] =
       GetEmittedValueFor(add_dependency->operand(0));
-  return OkStatus();
+  return absl::OkStatus();
 }
 
 absl::Status IrEmitter::HandleRng(HloInstruction* rng) {
@@ -3397,7 +3397,7 @@ absl::Status IrEmitter::HandleRngGetAndUpdateState(HloInstruction* rng_state) {
   store->setAlignment(llvm::Align(IrEmitter::MinimumAlignmentForPrimitiveType(
       rng_state->shape().element_type())));
 
-  return OkStatus();
+  return absl::OkStatus();
 }
 
 absl::Status IrEmitter::FinishVisit(HloInstruction* root) {
@@ -3425,7 +3425,7 @@ absl::Status IrEmitter::FinishVisit(HloInstruction* root) {
   // computations since it includes cycles spent in computations invoked by
   // While, Call etc.
   record_complete_computation(GetProfileCounterFor(*root->parent()));
-  return OkStatus();
+  return absl::OkStatus();
 }
 
 template <typename T>
@@ -3603,7 +3603,7 @@ absl::Status IrEmitter::Preprocess(HloInstruction* hlo) {
                                     GetExecutableRunOptionsArgument());
     profiling_state_.RecordCycleStart(&b_, hlo);
   }
-  return OkStatus();
+  return absl::OkStatus();
 }
 
 absl::Status IrEmitter::Postprocess(HloInstruction* hlo) {
@@ -3616,7 +3616,7 @@ absl::Status IrEmitter::Postprocess(HloInstruction* hlo) {
        hlo->parent()->IsEntryComputation())) {
     tracing_state_.EmitTracingEnd(&b_, hlo, GetExecutableRunOptionsArgument());
   }
-  return OkStatus();
+  return absl::OkStatus();
 }
 
 llvm_ir::IrArray IrEmitter::GetIrArrayFor(const HloInstruction* hlo) {
@@ -3773,7 +3773,7 @@ absl::Status IrEmitter::EmitTargetAddressForOp(const HloInstruction* op) {
   llvm::Value* addr = EmitBufferPointer(slice, target_shape);
   addr->setName(IrName(op));
   emitted_value_[op] = addr;
-  return OkStatus();
+  return absl::OkStatus();
 }
 
 absl::Status IrEmitter::EmitTargetElementLoop(
@@ -3832,7 +3832,7 @@ absl::Status IrEmitter::EmitTargetElementLoop(
               .EmitLoop(IrName(target_op)));
     }
   }
-  return OkStatus();
+  return absl::OkStatus();
 }
 
 absl::Status IrEmitter::EmitMemcpy(const HloInstruction& source,
@@ -3843,7 +3843,7 @@ absl::Status IrEmitter::EmitMemcpy(const HloInstruction& source,
   // TODO(b/63762267): Be more aggressive about specifying alignment.
   MemCpy(destination_value, /*DstAlign=*/llvm::Align(1), source_value,
          /*SrcAlign=*/llvm::Align(1), source_size);
-  return OkStatus();
+  return absl::OkStatus();
 }
 
 absl::Status IrEmitter::ElementTypesSameAndSupported(
@@ -3862,7 +3862,7 @@ absl::Status IrEmitter::ElementTypesSameAndSupported(
                          PrimitiveType_Name(primitive_type),
                          HloOpcodeString(instruction.opcode()));
   }
-  return OkStatus();
+  return absl::OkStatus();
 }
 
 absl::Status IrEmitter::DefaultAction(HloInstruction* hlo) {
diff --git a/third_party/xla/xla/service/cpu/ir_function.cc b/third_party/xla/xla/service/cpu/ir_function.cc
index cf66fe382e75be..1c6180604b6d88 100644
--- a/third_party/xla/xla/service/cpu/ir_function.cc
+++ b/third_party/xla/xla/service/cpu/ir_function.cc
@@ -329,7 +329,7 @@ absl::Status EmitCallToParallelForkJoin(
   // Emit call to parallel fork/join.
   b->CreateCall(fork_join_func, fork_join_arguments);
 
-  return OkStatus();
+  return absl::OkStatus();
 }
 
 }  // namespace cpu
diff --git a/third_party/xla/xla/service/cpu/mlir_emitter.cc b/third_party/xla/xla/service/cpu/mlir_emitter.cc
index 396f567cf2ddd4..62939d896b6eff 100644
--- a/third_party/xla/xla/service/cpu/mlir_emitter.cc
+++ b/third_party/xla/xla/service/cpu/mlir_emitter.cc
@@ -133,7 +133,7 @@ absl::Status EmitMlirFuncAndCall(
   }
   b->CreateCall(func, op_vals);
 
-  return OkStatus();
+  return absl::OkStatus();
 }
 
 }  // namespace cpu
diff --git a/third_party/xla/xla/service/cpu/onednn_matmul_rewriter.cc b/third_party/xla/xla/service/cpu/onednn_matmul_rewriter.cc
index 1fc542088d1531..57418d4959e145 100644
--- a/third_party/xla/xla/service/cpu/onednn_matmul_rewriter.cc
+++ b/third_party/xla/xla/service/cpu/onednn_matmul_rewriter.cc
@@ -55,7 +55,7 @@ inline absl::Status ValidateDotDimensionNumbers(
       absl::c_equal(batch_dim_numbers, dim_numbers.lhs_batch_dimensions()));
   TF_RET_CHECK(
       absl::c_equal(batch_dim_numbers, dim_numbers.rhs_batch_dimensions()));
-  return OkStatus();
+  return absl::OkStatus();
 }
 
 // Whether the element type of instr is compatible with oneDNN kernels.
@@ -431,11 +431,12 @@ class OneDnnMatMulRewriteVisitor : public DfsHloRewriteVisitor {
   absl::Status HandleDot(HloInstruction* instr) override {
     HloInstruction* dot_instr;
     auto pattern = m::Op(&dot_instr).WithOpcode(HloOpcode::kDot);
-    if (!Match(instr, pattern)) return OkStatus();
+    if (!Match(instr, pattern)) return absl::OkStatus();
 
     TF_RETURN_IF_ERROR(
         ValidateDotDimensionNumbers(dot_instr->dot_dimension_numbers()));
-    if (!OneDnnMatMulRewriter::ShouldRewrite(dot_instr)) return OkStatus();
+    if (!OneDnnMatMulRewriter::ShouldRewrite(dot_instr))
+      return absl::OkStatus();
     TF_ASSIGN_OR_RETURN(dot_instr, ReconfigureDotDimensions(dot_instr));
     auto dot_dim_numbers = dot_instr->dot_dimension_numbers();
     const Shape& lhs_shape = dot_instr->operand(0)->shape();
@@ -460,7 +461,7 @@ class OneDnnMatMulRewriteVisitor : public DfsHloRewriteVisitor {
     matmul_config->set_transpose_b(transpose_b);
     TF_RETURN_IF_ERROR(matmul_call->set_backend_config(backend_config));
     TF_RETURN_IF_ERROR(ReplaceInstruction(dot_instr, matmul_call));
-    return OkStatus();
+    return absl::OkStatus();
   }
 
   absl::Status HandleAdd(HloInstruction* instr) override {
@@ -496,7 +497,8 @@ class OneDnnMatMulRewriteVisitor : public DfsHloRewriteVisitor {
         m::Op(&addend_intermediate));
 
     if (Match(instr, pattern)) {
-      if (!IsSupportedType(dot->shape().element_type())) return OkStatus();
+      if (!IsSupportedType(dot->shape().element_type()))
+        return absl::OkStatus();
       // TODO(intel-tf): Remove the condition below when the fusion Dot +
       // Add(bias) + Add(e.g., residual) is enabled.
       if (!dot->backend_config<BackendConfig>()
@@ -506,7 +508,7 @@ class OneDnnMatMulRewriteVisitor : public DfsHloRewriteVisitor {
           dot->backend_config<BackendConfig>()
                   ->mutable_onednn_matmul_config()
                   ->fused_ops(0) == OneDnnMatMulConfig::BIAS) {
-        return OkStatus();
+        return absl::OkStatus();
       }
       std::vector<HloInstruction*> new_operands;
       for (auto operand : dot->operands()) {
@@ -533,7 +535,7 @@ class OneDnnMatMulRewriteVisitor : public DfsHloRewriteVisitor {
           m::Convert(&addend, m::Op()),
           m::Broadcast(&optional_addend_broadcast, m::Op(&addend)),
           m::Op(&addend));
-      if (!Match(addend_intermediate, addend_pattern)) return OkStatus();
+      if (!Match(addend_intermediate, addend_pattern)) return absl::OkStatus();
 
       if (optional_addend_broadcast && addend->shape().rank() != 1) {
         auto new_shape =
@@ -543,7 +545,7 @@ class OneDnnMatMulRewriteVisitor : public DfsHloRewriteVisitor {
               HloInstruction::CreateBitcast(new_shape.value(), addend));
         } else {
           VLOG(2) << new_shape.status();
-          return OkStatus();
+          return absl::OkStatus();
         }
       }
 
@@ -552,7 +554,7 @@ class OneDnnMatMulRewriteVisitor : public DfsHloRewriteVisitor {
           IsOperandFusible(addend, dot)) {
         new_operands.push_back(addend);
       } else {
-        return OkStatus();
+        return absl::OkStatus();
       }
 
       auto matmul_call = Cast<HloCustomCallInstruction>(instr->AddInstruction(
@@ -613,7 +615,7 @@ class OneDnnMatMulRewriteVisitor : public DfsHloRewriteVisitor {
       TF_RETURN_IF_ERROR(ReplaceInstruction(instr, new_instr));
     }
 
-    return OkStatus();
+    return absl::OkStatus();
   }
 
   absl::Status HandleMaximum(HloInstruction* instr) override {
@@ -631,7 +633,7 @@ class OneDnnMatMulRewriteVisitor : public DfsHloRewriteVisitor {
       return FuseActivation(OneDnnMatMulConfig::RELU, instr, matmul_call,
                             intermediate_instr, optional_bitcast);
     }
-    return OkStatus();
+    return absl::OkStatus();
   }
 
   auto ELUActivation(HloInstruction* instr, HloInstruction** src) {
@@ -662,7 +664,7 @@ class OneDnnMatMulRewriteVisitor : public DfsHloRewriteVisitor {
                               intermediate_instr);
       }
     }
-    return OkStatus();
+    return absl::OkStatus();
   }
 
   Status HandleTanh(HloInstruction* instr) override {
@@ -678,7 +680,7 @@ class OneDnnMatMulRewriteVisitor : public DfsHloRewriteVisitor {
       return FuseActivation(OneDnnMatMulConfig::TANH, instr, matmul_call,
                             intermediate_instr);
     }
-    return OkStatus();
+    return absl::OkStatus();
   }
 
   Status HandleClamp(HloInstruction* instr) override {
@@ -696,7 +698,7 @@ class OneDnnMatMulRewriteVisitor : public DfsHloRewriteVisitor {
       return FuseActivation(OneDnnMatMulConfig::RELU6, instr, matmul_call,
                             intermediate_instr);
     }
-    return OkStatus();
+    return absl::OkStatus();
   }
 
   absl::Status HandleMultiply(HloInstruction* instr) override {
@@ -757,7 +759,7 @@ class OneDnnMatMulRewriteVisitor : public DfsHloRewriteVisitor {
 
       TF_RETURN_IF_ERROR(ReplaceInstruction(instr, new_instr));
     }
-    return OkStatus();
+    return absl::OkStatus();
   }
 
   absl::Status FuseActivation(OneDnnMatMulConfig_FusionKind kind,
@@ -923,7 +925,7 @@ class OneDnnPostRewriteVisitor : public DfsHloRewriteVisitor {
     if (!weights_prepack.ok()) {
       VLOG(2) << weights_prepack.status();
     }
-    return OkStatus();
+    return absl::OkStatus();
   }
 
   template <typename>
diff --git a/third_party/xla/xla/service/cpu/onednn_ops_rewriter.cc b/third_party/xla/xla/service/cpu/onednn_ops_rewriter.cc
index 25530aa59c8d7b..95242415afbdde 100644
--- a/third_party/xla/xla/service/cpu/onednn_ops_rewriter.cc
+++ b/third_party/xla/xla/service/cpu/onednn_ops_rewriter.cc
@@ -401,7 +401,7 @@ class OneDnnOpsRewriterVisitor : public DfsHloRewriteVisitor {
                                     &is_producer_bf16orfp16, &convert_instr);
     }
 
-    if (!found_ln) return OkStatus();
+    if (!found_ln) return absl::OkStatus();
 
     const Shape& src_shape = src->shape();
     auto scale_type = scale->shape().element_type();
@@ -442,7 +442,7 @@ class OneDnnOpsRewriterVisitor : public DfsHloRewriteVisitor {
       TF_RETURN_IF_ERROR(ReplaceInstruction(instr, ln_call));
     }
 
-    return OkStatus();
+    return absl::OkStatus();
   }
 
   absl::Status HandleConvert(HloInstruction* instr) override {
@@ -455,12 +455,13 @@ class OneDnnOpsRewriterVisitor : public DfsHloRewriteVisitor {
                                 .WithOneUser()
                                 .WithElementType(PrimitiveType::F32));
 
-    if (!IsSupportedType(instr->shape().element_type())) return OkStatus();
+    if (!IsSupportedType(instr->shape().element_type()))
+      return absl::OkStatus();
     if (Match(instr, pattern)) {
       bool is_bf16orfp16_convert =
           (convert_instr->shape().element_type() == PrimitiveType::BF16) ||
           (convert_instr->shape().element_type() == PrimitiveType::F16);
-      if (!is_bf16orfp16_convert) return OkStatus();
+      if (!is_bf16orfp16_convert) return absl::OkStatus();
       HloInstruction* producer = instr->mutable_operand(0)->mutable_operand(0);
       HloInstruction* newinp =
           producer->AddInstruction(HloInstruction::CreateConvert(
@@ -475,15 +476,15 @@ class OneDnnOpsRewriterVisitor : public DfsHloRewriteVisitor {
       TF_RETURN_IF_ERROR(ReplaceInstruction(instr, updated_call));
     }
 
-    return OkStatus();
+    return absl::OkStatus();
   }
 
   absl::Status HandleDivide(HloInstruction* divide_instr) override {
-    if (divide_instr->HasControlDependencies()) return OkStatus();
+    if (divide_instr->HasControlDependencies()) return absl::OkStatus();
     if (!IsSupportedType(divide_instr->shape().element_type()))
-      return OkStatus();
+      return absl::OkStatus();
     std::optional<HloInstruction*> producer = MatchSoftmax(divide_instr);
-    if (producer == std::nullopt) return OkStatus();
+    if (producer == std::nullopt) return absl::OkStatus();
 
     const Shape& output_shape = divide_instr->shape();
     HloInstruction* softmax_call =
@@ -491,7 +492,7 @@ class OneDnnOpsRewriterVisitor : public DfsHloRewriteVisitor {
             output_shape, {producer.value()}, "__onednn$softmax"));
     TF_RETURN_IF_ERROR(ReplaceInstruction(divide_instr, softmax_call));
 
-    return OkStatus();
+    return absl::OkStatus();
   }
 };
 
diff --git a/third_party/xla/xla/service/memory_space_assignment/algorithm.cc b/third_party/xla/xla/service/memory_space_assignment/algorithm.cc
index c8cafe0bf8410e..b2b9ef640b5545 100644
--- a/third_party/xla/xla/service/memory_space_assignment/algorithm.cc
+++ b/third_party/xla/xla/service/memory_space_assignment/algorithm.cc
@@ -1033,7 +1033,7 @@ absl::Status MsaAlgorithm::OptimizeMemoryBoundLoop(int loop_start_idx,
       }
     }
   }
-  return OkStatus();
+  return absl::OkStatus();
 }
 
 namespace {
@@ -1477,7 +1477,7 @@ absl::StatusOr<HeapSimulator::Result<HloValue>> MsaAlgorithm::Finish() {
         VLOG(2) << "Repacking.";
         auto repack_status =
             options_.repacker->Repack(absl::MakeSpan(repack_allocation_blocks));
-        CHECK_EQ(repack_status.status(), OkStatus());
+        CHECK_EQ(repack_status.status(), absl::OkStatus());
         VLOG(2) << "Repack complete. Modified = " << *repack_status;
         // For debug and testing purpose, also update allocations if
         // repack_after_every_allocation is on.
@@ -1525,7 +1525,7 @@ absl::StatusOr<HeapSimulator::Result<HloValue>> MsaAlgorithm::Finish() {
     VLOG(2) << "Final Repacking.";
     auto repack_status =
         options_.repacker->Repack(absl::MakeSpan(repack_allocation_blocks));
-    CHECK_EQ(repack_status.status(), OkStatus());
+    CHECK_EQ(repack_status.status(), absl::OkStatus());
     VLOG(2) << "Final Repack complete. Modified = " << *repack_status;
   }
 
@@ -3377,7 +3377,7 @@ void MsaAlgorithm::ImportRepackedSlicedAllocation(
 absl::Status MsaAlgorithm::AreRepackedSlicesValid(
     const RepackAllocationBlock& block) {
   if (!block.repacked_slice_data.has_value()) {
-    return OkStatus();
+    return absl::OkStatus();
   }
   if (!block.original_slice_data.has_value()) {
     return InvalidArgumentStrCat(
@@ -3418,7 +3418,7 @@ absl::Status MsaAlgorithm::AreRepackedSlicesValid(
         "mappings.");
   }
 
-  return OkStatus();
+  return absl::OkStatus();
 }
 
 void MsaAlgorithm::UncommitPendingChunks(
diff --git a/third_party/xla/xla/service/memory_space_assignment/allocation.cc b/third_party/xla/xla/service/memory_space_assignment/allocation.cc
index 9f01ddbd9ab791..00e22355fc6239 100644
--- a/third_party/xla/xla/service/memory_space_assignment/allocation.cc
+++ b/third_party/xla/xla/service/memory_space_assignment/allocation.cc
@@ -177,7 +177,7 @@ absl::Status Allocation::UpdateUses(HloComputation* computation,
     TF_RETURN_IF_ERROR(use.instruction->ReplaceOperandWith(
         use.operand_number, replacement_instruction));
   }
-  return OkStatus();
+  return absl::OkStatus();
 }
 
 bool Allocation::is_copy_like_allocation() const {
@@ -262,7 +262,7 @@ bool PinnedAllocation::operator==(const Allocation& other) const {
 absl::Status PinnedAllocation::Process() {
   if (is_scoped_allocation()) {
     // Nothing to do here for scoped allocations.
-    return OkStatus();
+    return absl::OkStatus();
   }
   HloInstruction* producing_instruction = AddGetTupleElements();
   HloComputation* computation = producing_instruction->parent();
@@ -626,7 +626,7 @@ absl::Status SlicedCopyAllocation::CreateBitcastConcat(
       slices.front()->parent()->AddInstruction(HloInstruction::CreateCustomCall(
           shape, slices,
           xla::memory_space_assignment::kConcatBitcastCustomCall));
-  return OkStatus();
+  return absl::OkStatus();
 }
 
 std::string SlicedCopyAllocation::SliceDetail::ToString() const {
@@ -695,7 +695,7 @@ absl::Status SlicedCopyAllocation::SliceDetail::CreateAsyncSlice(
                                      slice, {ShapeUtil::MakeShape(S32, {})}));
   copy_start = copy_done->mutable_operand(0);
 
-  return OkStatus();
+  return absl::OkStatus();
 }
 
 bool SlicedCopyAllocation::operator==(const Allocation& other) const {
@@ -732,7 +732,7 @@ absl::Status MirroredAllocation::Process() {
   set_original_defining_position(original_allocation_.defining_position());
   if (is_scoped_allocation()) {
     // Nothing to do here for scoped allocations.
-    return OkStatus();
+    return absl::OkStatus();
   }
   HloInstruction* producing_instruction = AddGetTupleElements();
   HloComputation* computation = producing_instruction->parent();
@@ -789,7 +789,7 @@ absl::Status ParentAllocation::Process() {
 
   if (is_scoped_allocation()) {
     // Nothing to do here for scoped allocations.
-    return OkStatus();
+    return absl::OkStatus();
   }
   HloInstruction* final_instruction = AddGetTupleElements();
   HloComputation* computation = final_instruction->parent();
@@ -809,7 +809,7 @@ absl::Status ParentAllocation::PostProcess() {
                           original_defining_position().index));
   while_body->set_root_instruction(new_while_body_root,
                                    /*accept_different_shape=*/true);
-  return OkStatus();
+  return absl::OkStatus();
 }
 
 void ParentAllocation::MarkIfNeeded(
diff --git a/third_party/xla/xla/service/memory_space_assignment/allocation.h b/third_party/xla/xla/service/memory_space_assignment/allocation.h
index 3d01439aee7623..f523adc5df420f 100644
--- a/third_party/xla/xla/service/memory_space_assignment/allocation.h
+++ b/third_party/xla/xla/service/memory_space_assignment/allocation.h
@@ -210,7 +210,7 @@ class PinnedAllocation final : public Allocation {
   bool is_copy_allocation() const override { return false; }
   bool is_sliced_copy_allocation() const override { return false; }
   absl::Status Process() override;
-  absl::Status PostProcess() override { return OkStatus(); }
+  absl::Status PostProcess() override { return absl::OkStatus(); }
   void MarkIfNeeded(absl::flat_hash_set<const Allocation*>& needed_allocations)
       const override;
   void MarkNeeded(absl::flat_hash_set<const Allocation*>& needed_allocations)
@@ -247,7 +247,7 @@ class CopyAllocation final : public Allocation {
   bool is_copy_allocation() const override { return true; }
   bool is_sliced_copy_allocation() const override { return false; }
   absl::Status Process() override;
-  absl::Status PostProcess() override { return OkStatus(); }
+  absl::Status PostProcess() override { return absl::OkStatus(); }
   void MarkIfNeeded(absl::flat_hash_set<const Allocation*>& needed_allocations)
       const override;
   void MarkNeeded(absl::flat_hash_set<const Allocation*>& needed_allocations)
@@ -349,7 +349,7 @@ class SlicedCopyAllocation final : public Allocation {
   // MemorySpaceAssignment::Process() calls Process() to create asynchronous
   // slice copies, and a bitcast-concat call to glue the slices back together.
   absl::Status Process() override;
-  absl::Status PostProcess() override { return OkStatus(); }
+  absl::Status PostProcess() override { return absl::OkStatus(); }
   // Marks the allocation as needed.
   void MarkIfNeeded(absl::flat_hash_set<const Allocation*>& needed_allocations)
       const override;
@@ -405,7 +405,7 @@ class MirroredAllocation final : public Allocation {
   bool is_copy_allocation() const override { return false; }
   bool is_sliced_copy_allocation() const override { return false; }
   absl::Status Process() override;
-  absl::Status PostProcess() override { return OkStatus(); }
+  absl::Status PostProcess() override { return absl::OkStatus(); }
   void MarkIfNeeded(absl::flat_hash_set<const Allocation*>& needed_allocations)
       const override;
   void MarkNeeded(absl::flat_hash_set<const Allocation*>& needed_allocations)
diff --git a/third_party/xla/xla/service/memory_space_assignment/cost_analysis_test.cc b/third_party/xla/xla/service/memory_space_assignment/cost_analysis_test.cc
index 5340551b0c1374..47df6481df60f9 100644
--- a/third_party/xla/xla/service/memory_space_assignment/cost_analysis_test.cc
+++ b/third_party/xla/xla/service/memory_space_assignment/cost_analysis_test.cc
@@ -65,7 +65,7 @@ class MemorySpaceAssignmentCostAnalysisTest : public HloTestBase {
     TF_ASSIGN_OR_RETURN(
         cost_analysis_,
         CostAnalysis::Create(*hlo_cost_analysis_costs_, options_, *module));
-    return OkStatus();
+    return absl::OkStatus();
   }
 
   CostAnalysisOptions options_;
diff --git a/third_party/xla/xla/service/memory_space_assignment/memory_bound_loop_optimizer.cc b/third_party/xla/xla/service/memory_space_assignment/memory_bound_loop_optimizer.cc
index 55035f996def94..3e3cf4a10a0cba 100644
--- a/third_party/xla/xla/service/memory_space_assignment/memory_bound_loop_optimizer.cc
+++ b/third_party/xla/xla/service/memory_space_assignment/memory_bound_loop_optimizer.cc
@@ -158,7 +158,7 @@ absl::Status MemoryBoundLoopOptimizer::Initialize() {
   for (const HloBuffer* buffer : buffers_to_process) {
     MaybeCreateLoopValue(*buffer, loop_computation);
   }
-  return OkStatus();
+  return absl::OkStatus();
 }
 
 void MemoryBoundLoopOptimizer::MaybeCreateLoopValue(
diff --git a/third_party/xla/xla/service/memory_space_assignment/memory_bound_loop_optimizer_test.cc b/third_party/xla/xla/service/memory_space_assignment/memory_bound_loop_optimizer_test.cc
index 15b50eac85f4d0..0d4fdcdbe4c7b2 100644
--- a/third_party/xla/xla/service/memory_space_assignment/memory_bound_loop_optimizer_test.cc
+++ b/third_party/xla/xla/service/memory_space_assignment/memory_bound_loop_optimizer_test.cc
@@ -119,7 +119,7 @@ class MemoryBoundLoopOptimizerTest : public HloTestBase {
     TF_ASSIGN_OR_RETURN(live_range_,
                         HloLiveRange::Run(module->schedule(), *alias_analysis_,
                                           module->entry_computation()));
-    return OkStatus();
+    return absl::OkStatus();
   }
 
   absl::StatusOr<MemoryBoundLoopOptimizer*> CreateOptimizer(
@@ -496,7 +496,7 @@ ENTRY Entry {
         }
       }
     }
-    return OkStatus();
+    return absl::OkStatus();
   }
 
  private:
diff --git a/third_party/xla/xla/service/memory_space_assignment/memory_space_assignment.cc b/third_party/xla/xla/service/memory_space_assignment/memory_space_assignment.cc
index b2ede602ceb643..ae22ed48ffb0ed 100644
--- a/third_party/xla/xla/service/memory_space_assignment/memory_space_assignment.cc
+++ b/third_party/xla/xla/service/memory_space_assignment/memory_space_assignment.cc
@@ -81,7 +81,7 @@ absl::Status EnsureInstructionAndOperandsInserted(
     HloInstruction* new_instruction, HloInstructionSequence* new_sequence,
     absl::flat_hash_set<HloInstruction*>* inserted_instructions) {
   if (inserted_instructions->contains(new_instruction)) {
-    return OkStatus();
+    return absl::OkStatus();
   }
   return InsertInstructionAndEnsureOperandsInserted(
       new_instruction, new_sequence, inserted_instructions);
@@ -100,7 +100,7 @@ absl::Status InsertInstructionAndEnsureOperandsInserted(
   VLOG(4) << "inserting: " << new_instruction->ToShortString();
   new_sequence->push_back(new_instruction);
   TF_RET_CHECK(inserted_instructions->insert(new_instruction).second);
-  return OkStatus();
+  return absl::OkStatus();
 }
 
 std::string InstructionScheduleToString(const HloLiveRange& hlo_live_range) {
@@ -387,7 +387,7 @@ absl::Status MemorySpaceAssignment::FindAllocationSequence(
                                         options_.size_fn,
                                         heap_simulator_options)
                          .status());
-  return OkStatus();
+  return absl::OkStatus();
 }
 
 float MemorySpaceAssignment::ComputeEstimatedElapsedTime(
@@ -510,7 +510,7 @@ absl::Status MemorySpaceAssignment::Process(
       TF_RETURN_IF_ERROR(allocation->PostProcess());
     }
   }
-  return OkStatus();
+  return absl::OkStatus();
 }
 
 absl::Status MemorySpaceAssignment::ExportAndColorBuffers() {
@@ -578,7 +578,7 @@ absl::Status MemorySpaceAssignment::ExportAndColorBuffers() {
       }
     }
   }
-  return OkStatus();
+  return absl::OkStatus();
 }
 
 void MemorySpaceAssignment::RemoveAssignmentForInstruction(
@@ -701,7 +701,7 @@ absl::Status MemorySpaceAssignment::SimplifyGraph() {
     }
   }
 
-  return OkStatus();
+  return absl::OkStatus();
 }
 
 namespace {
@@ -1019,7 +1019,7 @@ absl::Status MemorySpaceAssignment::FixSchedule() {
 
   TF_RETURN_IF_ERROR(schedule.Update());
 
-  return OkStatus();
+  return absl::OkStatus();
 }
 
 absl::Status MemorySpaceAssignment::VerifyAndExportHeapSimulatorTrace() {
@@ -1065,7 +1065,7 @@ absl::Status MemorySpaceAssignment::VerifyAndExportHeapSimulatorTrace() {
       }
     }
     interval_tree.Add(start_time, end_time - 1, chunk);
-    return OkStatus();
+    return absl::OkStatus();
   };
 
   // Go through all instructions in the module to ensure CopyStart/CopyDone
@@ -1171,7 +1171,7 @@ absl::Status MemorySpaceAssignment::VerifyAndExportHeapSimulatorTrace() {
                 << ")";
         TF_RETURN_IF_ERROR(add_allocation_and_verify(
             start_time, earliest_computation_start_time - 1, chunk, value));
-        return OkStatus();
+        return absl::OkStatus();
       };
 
       if (last_use_instruction &&
@@ -1238,7 +1238,7 @@ absl::Status MemorySpaceAssignment::VerifyAndExportHeapSimulatorTrace() {
   }
   VLOG(1) << "Max memory usage ignoring fragmentation: " << max_memory_usage;
 
-  return OkStatus();
+  return absl::OkStatus();
 }
 
 }  // namespace memory_space_assignment
diff --git a/third_party/xla/xla/service/memory_space_assignment/memory_space_assignment_test.cc b/third_party/xla/xla/service/memory_space_assignment/memory_space_assignment_test.cc
index b694fb852c5d50..7517fa78ca21ef 100644
--- a/third_party/xla/xla/service/memory_space_assignment/memory_space_assignment_test.cc
+++ b/third_party/xla/xla/service/memory_space_assignment/memory_space_assignment_test.cc
@@ -10364,7 +10364,7 @@ class SlicedPrefetchTest : public MemorySpaceAssignmentTestBase {
       }
     }
 
-    return OkStatus();
+    return absl::OkStatus();
   }
 
   // REQUIRES:
@@ -10394,7 +10394,7 @@ class SlicedPrefetchTest : public MemorySpaceAssignmentTestBase {
       }
     }
 
-    return OkStatus();
+    return absl::OkStatus();
   }
 
   // REQUIRES:
@@ -10447,7 +10447,7 @@ class SlicedPrefetchTest : public MemorySpaceAssignmentTestBase {
                     first_slice_start_after_schedule_after, "."));
     }
 
-    return OkStatus();
+    return absl::OkStatus();
   }
 
   // REQUIRES:
@@ -10474,7 +10474,7 @@ class SlicedPrefetchTest : public MemorySpaceAssignmentTestBase {
       }
     }
 
-    return OkStatus();
+    return absl::OkStatus();
   }
 
   // REQUIRES:
@@ -10507,7 +10507,7 @@ class SlicedPrefetchTest : public MemorySpaceAssignmentTestBase {
       }
     }
 
-    return OkStatus();
+    return absl::OkStatus();
   }
 
   // REQUIRES:
@@ -10537,7 +10537,7 @@ class SlicedPrefetchTest : public MemorySpaceAssignmentTestBase {
       }
     }
 
-    return OkStatus();
+    return absl::OkStatus();
   }
 
   // Return an OK status iff:
@@ -10629,10 +10629,10 @@ class SlicedPrefetchTest : public MemorySpaceAssignmentTestBase {
     TF_RETURN_IF_ERROR(
         ConcatBitcastAfterSliceDones(entry_schedule, schedule_to_class));
 
-    return OkStatus();
+    return absl::OkStatus();
   }
 
-  // Returns OkStatus iff:
+  // Returns absl::OkStatus iff:
   // - Each slice is assigned a chunk that is the same size as the slice
   //   instruction's shape.
   // - When the slices of sliced_copy_result are sorted in expected spatial
@@ -10699,7 +10699,7 @@ class SlicedPrefetchTest : public MemorySpaceAssignmentTestBase {
             << (result_chunk.has_value() ? result_chunk->ToString()
                                          : "no chunk assigned");
     if (sorted_slices.empty()) {
-      return OkStatus();
+      return absl::OkStatus();
     }
 
     // Check that slices are assigned contiguous chunks that are spatially
@@ -10763,7 +10763,7 @@ class SlicedPrefetchTest : public MemorySpaceAssignmentTestBase {
                              ", to match its shape."));
     }
 
-    return OkStatus();
+    return absl::OkStatus();
   }
 
   SlicedPrefetchTest() {

From f2fd583e3988a12677158ebd072d7411e13329fd Mon Sep 17 00:00:00 2001
From: Kyle Lucke <klucke@google.com>
Date: Mon, 20 May 2024 12:06:36 -0700
Subject: [PATCH 427/478] Replace the use of xla::OkStatus with absl::OkStatus
 now that they're the same.

PiperOrigin-RevId: 635530537
---
 third_party/xla/xla/array.h               |  8 +++---
 third_party/xla/xla/layout_util.cc        | 14 +++++-----
 third_party/xla/xla/literal.cc            | 34 +++++++++++------------
 third_party/xla/xla/literal.h             | 20 ++++++-------
 third_party/xla/xla/literal_comparison.cc |  8 +++---
 third_party/xla/xla/shape_layout.cc       |  4 +--
 third_party/xla/xla/shape_tree.h          |  8 +++---
 third_party/xla/xla/shape_util.cc         | 20 ++++++-------
 third_party/xla/xla/shape_util.h          | 12 ++++----
 third_party/xla/xla/sharding_op_util.cc   |  2 +-
 third_party/xla/xla/status_macros_test.cc | 12 ++++----
 third_party/xla/xla/util.h                |  2 +-
 12 files changed, 72 insertions(+), 72 deletions(-)

diff --git a/third_party/xla/xla/array.h b/third_party/xla/xla/array.h
index 98ea2858ab4869..fe33a49335ac26 100644
--- a/third_party/xla/xla/array.h
+++ b/third_party/xla/xla/array.h
@@ -314,7 +314,7 @@ class Array {
 
   // Invokes a callback with the (indices, value_ptr) for each cell in the
   // array. If a callback returns a non-OK status, returns that else returns
-  // OkStatus().
+  // absl::OkStatus().
   absl::Status EachStatus(
       absl::FunctionRef<absl::Status(absl::Span<const int64_t>, T*)> f) {
     OwnedBuffer<int64_t> index(sizes_.size, default_init_t{});
@@ -324,12 +324,12 @@ class Array {
         return s;
       }
     }
-    return OkStatus();
+    return absl::OkStatus();
   }
 
   // Invokes a callback with the (indices, value) for each cell in the array.
   // If a callback returns a non-OK status, returns that else returns
-  // OkStatus().
+  // absl::OkStatus().
   absl::Status EachStatus(
       absl::FunctionRef<absl::Status(absl::Span<const int64_t>, T)> f) const {
     OwnedBuffer<int64_t> index(sizes_.size, default_init_t{});
@@ -339,7 +339,7 @@ class Array {
         return s;
       }
     }
-    return OkStatus();
+    return absl::OkStatus();
   }
 
   // Returns the value at the cell specified by the indexes. The number of
diff --git a/third_party/xla/xla/layout_util.cc b/third_party/xla/xla/layout_util.cc
index ed57d57cf4f286..59381c738b4818 100644
--- a/third_party/xla/xla/layout_util.cc
+++ b/third_party/xla/xla/layout_util.cc
@@ -215,11 +215,11 @@ Layout CreateDefaultLayoutForRank(int64_t rank) {
       TF_RETURN_IF_ERROR(
           ValidateLayoutInShape(element_shape, allow_missing_layouts));
     }
-    return OkStatus();
+    return absl::OkStatus();
   } else if (shape.IsArray()) {
     if (!shape.has_layout()) {
       if (allow_missing_layouts) {
-        return OkStatus();
+        return absl::OkStatus();
       }
       return InvalidArgument("shape %s does not have a layout",
                              ShapeUtil::HumanString(shape));
@@ -232,7 +232,7 @@ Layout CreateDefaultLayoutForRank(int64_t rank) {
           "shape of primitive type %s should not have a layout",
           PrimitiveType_Name(shape.element_type()));
     }
-    return OkStatus();
+    return absl::OkStatus();
   }
 }
 
@@ -248,7 +248,7 @@ Layout CreateDefaultLayoutForRank(int64_t rank) {
           "shape of primitive type %s should not have a non-trivial layout",
           PrimitiveType_Name(shape.element_type()));
     }
-    return OkStatus();
+    return absl::OkStatus();
   }
 
   if (layout.minor_to_major_size() != shape.rank()) {
@@ -358,7 +358,7 @@ Layout CreateDefaultLayoutForRank(int64_t rank) {
                   "physical shape: %s",
                   shape.ShortDebugString());
             }
-            return OkStatus();
+            return absl::OkStatus();
           }));
       if (layout.index_primitive_type() != PRIMITIVE_TYPE_INVALID &&
           !primitive_util::IsUnsignedIntegralType(
@@ -419,7 +419,7 @@ Layout CreateDefaultLayoutForRank(int64_t rank) {
                            layout.element_size_in_bits());
   }
 
-  return OkStatus();
+  return absl::OkStatus();
 }
 
 /* static */ void LayoutUtil::ClearLayout(Shape* shape) {
@@ -613,7 +613,7 @@ absl::Status CopyLayoutInternal(const Shape& src, Shape* dst) {
       dst->clear_layout();
     }
   }
-  return OkStatus();
+  return absl::OkStatus();
 }
 
 }  // namespace
diff --git a/third_party/xla/xla/literal.cc b/third_party/xla/xla/literal.cc
index fe35a577ebb37b..a11ddc71c549e5 100644
--- a/third_party/xla/xla/literal.cc
+++ b/third_party/xla/xla/literal.cc
@@ -458,7 +458,7 @@ absl::Status MutableLiteralBase::CopySliceFromInternal(
                             stride_config.dimensions, stride_config.step,
                             copy_proc);
   }
-  return OkStatus();
+  return absl::OkStatus();
 }
 
 void MutableLiteralBase::CopyElementFrom(const LiteralSlice& src_literal,
@@ -521,10 +521,10 @@ void MutableLiteralBase::CopyElementFrom(const LiteralSlice& src_literal,
                 ShapeUtil::TupleElementCount(piece->subshape()),
                 proto_element->tuple_literals_size());
           }
-          return OkStatus();
+          return absl::OkStatus();
         }
         if (piece->subshape().element_type() == TOKEN) {
-          return OkStatus();
+          return absl::OkStatus();
         }
 
         CHECK(piece->subshape().IsArray());
@@ -536,7 +536,7 @@ void MutableLiteralBase::CopyElementFrom(const LiteralSlice& src_literal,
           TF_RETURN_IF_ERROR(piece->CopyFromProto(*proto_element));
         }
 
-        return OkStatus();
+        return absl::OkStatus();
       }));
 
   return std::move(literal);
@@ -691,7 +691,7 @@ absl::Status LiteralBase::Piece::CopyFrom(const LiteralBase::Piece& src,
       DeallocateBuffers();
     }
     array_value_state_ = src.array_value_state_;
-    return OkStatus();
+    return absl::OkStatus();
   } else {
     CHECK(src.array_value_state_ == ArrayValueState::kKnown);
     if (array_value_state_ == ArrayValueState::kUndetermined ||
@@ -724,7 +724,7 @@ absl::Status LiteralBase::Piece::CopyFrom(const LiteralBase::Piece& src,
     memcpy(dynamic_size_buffer(), src.dynamic_size_buffer(),
            src.dynamic_size_buffer_bytes());
   }
-  return OkStatus();
+  return absl::OkStatus();
 }
 
 void MutableLiteralBase::SetDynamicSize(int64_t dim_index, int32_t size) {
@@ -771,7 +771,7 @@ absl::Status MutableLiteralBase::CopyFrom(const LiteralSlice& src_literal,
   return mutable_root_piece().ForEachMutableSubpieceWithStatus(
       [&](const ShapeIndex& index, Piece* piece) {
         if (!piece->subshape().IsArray()) {
-          return OkStatus();
+          return absl::OkStatus();
         }
 
         // Determine if this index is in the part of this literal that we want
@@ -784,7 +784,7 @@ absl::Status MutableLiteralBase::CopyFrom(const LiteralSlice& src_literal,
           }
         }
         if (!in_subtree_to_copy) {
-          return OkStatus();
+          return absl::OkStatus();
         }
         // Construct the index of the corresponding piece in the source literal.
         ShapeIndex src_piece_index = src_shape_index;
@@ -795,7 +795,7 @@ absl::Status MutableLiteralBase::CopyFrom(const LiteralSlice& src_literal,
         TF_RETURN_IF_ERROR(
             piece->CopyFrom(src_literal.piece(src_piece_index),
                             /*only_dynamic_bound=*/only_dynamic_bound));
-        return OkStatus();
+        return absl::OkStatus();
       });
 }
 
@@ -829,7 +829,7 @@ absl::Status Literal::MoveFrom(Literal&& src_literal,
   src_literal.root_piece_ = Piece();
   src_literal.root_piece_.set_subshape(src_literal.shape_.get());
 
-  return OkStatus();
+  return absl::OkStatus();
 }
 
 absl::Status MutableLiteralBase::CopySliceFrom(
@@ -935,7 +935,7 @@ absl::Status MutableLiteralBase::PopulateInplace(
         return populator(dest, indexes);
       },
       /*parallel=*/false);
-  return OkStatus();
+  return absl::OkStatus();
 }
 
 absl::Status MutableLiteralBase::PopulateInplaceParallel(
@@ -944,7 +944,7 @@ absl::Status MutableLiteralBase::PopulateInplaceParallel(
       << __func__ << " is only supported for dense arrays: " << shape();
   PopulateInplaceInternal(populator,
                           /*parallel=*/element_count() > 32);
-  return OkStatus();
+  return absl::OkStatus();
 }
 
 Literal LiteralBase::Relayout(const Layout& new_layout,
@@ -1390,7 +1390,7 @@ absl::Status MutableLiteralBase::SetIntegralAsS64(
                       primitive_type_constant == PRED) {
           using NativeT = NativeTypeOf<primitive_type_constant>;
           Set<NativeT>(multi_index, static_cast<NativeT>(value));
-          return OkStatus();
+          return absl::OkStatus();
         }
         return FailedPrecondition("Array element type is not integral: %s",
                                   PrimitiveType_Name(shape().element_type()));
@@ -1411,7 +1411,7 @@ absl::Status MutableLiteralBase::SetFromDouble(
         Set<NativeT>(multi_index, static_cast<NativeT>(value));
       },
       shape().element_type());
-  return OkStatus();
+  return absl::OkStatus();
 }
 
 namespace {
@@ -1715,7 +1715,7 @@ absl::Status ConvertIfDestTypeMatches(const LiteralBase& src_literal,
           ConvertBetweenNativeTypes<NativeSrcT, NativeDestT>(src_data,
                                                              dst_base);
         }
-        return OkStatus();
+        return absl::OkStatus();
       },
       dst_literal.shape().element_type());
 }
@@ -2343,7 +2343,7 @@ absl::Status CopyFromRepeatedField(absl::Span<NativeT> dest,
         dest.size(), src.size());
   }
   std::copy(src.begin(), src.end(), dest.begin());
-  return OkStatus();
+  return absl::OkStatus();
 }
 
 }  // namespace
@@ -2514,7 +2514,7 @@ absl::Status LiteralBase::Piece::CopyFromProto(const LiteralProto& proto) {
       return InvalidArgument("Is called on unsupported shape: %s",
                              ShapeUtil::HumanString(subshape()));
   }
-  return OkStatus();
+  return absl::OkStatus();
 }
 
 bool LiteralBase::Piece::IsKnown() const {
diff --git a/third_party/xla/xla/literal.h b/third_party/xla/xla/literal.h
index eb88978687fbc5..0ea05a201d8188 100644
--- a/third_party/xla/xla/literal.h
+++ b/third_party/xla/xla/literal.h
@@ -913,7 +913,7 @@ class LiteralBase {
       return ForEachHelper(
                  [&func](const ShapeIndex& index, const Piece& piece) {
                    func(index, piece);
-                   return OkStatus();
+                   return absl::OkStatus();
                  },
                  *this, &index)
           .IgnoreError();
@@ -942,7 +942,7 @@ class LiteralBase {
       return ForEachMutableHelper(
                  [&func](const ShapeIndex& index, Piece* piece) {
                    func(index, piece);
-                   return OkStatus();
+                   return absl::OkStatus();
                  },
                  const_cast<xla::LiteralBase::Piece*>(this), &index)
           .IgnoreError();
@@ -1059,7 +1059,7 @@ class LiteralBase {
           index->pop_back();
         }
       }
-      return OkStatus();
+      return absl::OkStatus();
     }
     template <typename Fn>
     bool ForEachHelperBool(const Fn& func, const Piece& piece,
@@ -1090,7 +1090,7 @@ class LiteralBase {
           index->pop_back();
         }
       }
-      return OkStatus();
+      return absl::OkStatus();
     }
 
     // Recursive helper for EqualElements.
@@ -1555,7 +1555,7 @@ absl::Status LiteralBase::SerializeWithShapeProto(const ShapeProto& shape_proto,
       [&](const ShapeIndex& shape_index, const Piece& piece) -> absl::Status {
         const Shape& subshape = piece.subshape();
         if (subshape.IsTuple()) {
-          return OkStatus();
+          return absl::OkStatus();
         }
         if (!subshape.IsArray()) {
           return InvalidArgument("Shape cannot be serialized: %s",
@@ -1567,11 +1567,11 @@ absl::Status LiteralBase::SerializeWithShapeProto(const ShapeProto& shape_proto,
               piece.SerializeData<NativeT>(state);
             },
             subshape.element_type());
-        return OkStatus();
+        return absl::OkStatus();
       }));
   DCHECK_EQ(state.num_written(), SerializedSize().value())
       << shape().ToString();
-  return OkStatus();
+  return absl::OkStatus();
 }
 
 template <typename InputIterator>
@@ -1589,7 +1589,7 @@ absl::StatusOr<Literal> Literal::Deserialize(InputIterator begin,
           [&](const ShapeIndex& shape_index, Piece* piece) -> absl::Status {
             const Shape& subshape = piece->subshape();
             if (subshape.IsTuple()) {
-              return OkStatus();
+              return absl::OkStatus();
             }
             if (!subshape.IsArray()) {
               return InvalidArgument("Shape cannot be deserialized: %s",
@@ -1606,7 +1606,7 @@ absl::StatusOr<Literal> Literal::Deserialize(InputIterator begin,
                   "Failed to deserialize all data for shape: %s",
                   shape.ToString());
             }
-            return OkStatus();
+            return absl::OkStatus();
           }));
   DCHECK_EQ(state.num_read(), ShapeUtil::SerializedSize(shape).value())
       << shape.ToString();
@@ -1885,7 +1885,7 @@ TF_ATTRIBUTE_NOINLINE absl::Status MutableLiteralBase::PopulateInternal(
         *static_cast<NativeT*>(dest) = generator(indices, thread_id);
       },
       parallel);
-  return OkStatus();
+  return absl::OkStatus();
 }
 
 template <typename NativeT>
diff --git a/third_party/xla/xla/literal_comparison.cc b/third_party/xla/xla/literal_comparison.cc
index 6360b0f033405e..cb2e2cf64539df 100644
--- a/third_party/xla/xla/literal_comparison.cc
+++ b/third_party/xla/xla/literal_comparison.cc
@@ -148,7 +148,7 @@ absl::Status Equal(LiteralSlice expected, LiteralSlice actual,
     if (mismatched) {
       mismatched->Set<bool>(multi_index, !result);
     }
-    return result ? OkStatus()
+    return result ? absl::OkStatus()
                   : MakeErrorStatus<NativeT>(expected_value, actual_value,
                                              multi_index);
   }
@@ -300,7 +300,7 @@ class NearComparator {
     CompareLiterals();
 
     if (num_mismatches_ == 0) {
-      return OkStatus();
+      return absl::OkStatus();
     } else if (!VLOG_IS_ON(1) && miscompare_callback_ != nullptr) {
       miscompare_callback_(
           expected_, actual_, mismatches_, shape_index_,
@@ -809,7 +809,7 @@ absl::Status EqualShapes(const Shape& expected, const Shape& actual) {
     }
   }
   // Non-array, non-tuple shapes are trivially equivalent.
-  return OkStatus();
+  return absl::OkStatus();
 }
 
 absl::Status EqualDynamicShapesAndDimensions(const LiteralSlice& expected,
@@ -856,7 +856,7 @@ absl::Status EqualDynamicShapesAndDimensions(const LiteralSlice& expected,
           }
         }
 
-        return OkStatus();
+        return absl::OkStatus();
       });
 }
 
diff --git a/third_party/xla/xla/shape_layout.cc b/third_party/xla/xla/shape_layout.cc
index 4ea3e43c9762f3..a810e0bda9a896 100644
--- a/third_party/xla/xla/shape_layout.cc
+++ b/third_party/xla/xla/shape_layout.cc
@@ -33,7 +33,7 @@ absl::Status ShapeLayout::CopyLayoutFromShape(const Shape& other_shape) {
                            ShapeUtil::HumanString(shape()));
   }
   shape_ = other_shape;
-  return OkStatus();
+  return absl::OkStatus();
 }
 
 absl::Status ShapeLayout::AssignLayoutToShape(Shape* to_shape) const {
@@ -43,7 +43,7 @@ absl::Status ShapeLayout::AssignLayoutToShape(Shape* to_shape) const {
                            ShapeUtil::HumanString(shape()));
   }
   *to_shape = shape_;
-  return OkStatus();
+  return absl::OkStatus();
 }
 
 void ShapeLayout::SetToDefaultLayout() {
diff --git a/third_party/xla/xla/shape_tree.h b/third_party/xla/xla/shape_tree.h
index 8be8e906318110..3d514d23e53158 100644
--- a/third_party/xla/xla/shape_tree.h
+++ b/third_party/xla/xla/shape_tree.h
@@ -237,7 +237,7 @@ class ShapeTree {
     for (const Node& node : nodes_) {
       TF_RETURN_IF_ERROR(func(node.first, node.second));
     }
-    return OkStatus();
+    return absl::OkStatus();
   }
 
   absl::Status ForEachMutableElementWithStatus(
@@ -245,7 +245,7 @@ class ShapeTree {
     for (Node& node : nodes_) {
       TF_RETURN_IF_ERROR(func(node.first, &node.second));
     }
-    return OkStatus();
+    return absl::OkStatus();
   }
 
   // Like the above, but traverses in post-order.  Note children are visited in
@@ -269,7 +269,7 @@ class ShapeTree {
     for (auto node = nodes_.rbegin(); node != nodes_.rend(); ++node) {
       TF_RETURN_IF_ERROR(func(node->first, node->second));
     }
-    return OkStatus();
+    return absl::OkStatus();
   }
 
   absl::Status ForEachMutableElementPostOrderWithStatus(
@@ -277,7 +277,7 @@ class ShapeTree {
     for (auto node = nodes_.rbegin(); node != nodes_.rend(); ++node) {
       TF_RETURN_IF_ERROR(func(node->first, &node->second));
     }
-    return OkStatus();
+    return absl::OkStatus();
   }
 
   // Maps each element to generate a new tree with the same shape.
diff --git a/third_party/xla/xla/shape_util.cc b/third_party/xla/xla/shape_util.cc
index 976c06ab3bde08..739dacdc26266c 100644
--- a/third_party/xla/xla/shape_util.cc
+++ b/third_party/xla/xla/shape_util.cc
@@ -929,7 +929,7 @@ Shape ShapeUtil::PrependMajorDimension(int64_t bound, Shape shape) {
       shape,
       [&](const Shape& subshape, const ShapeIndex& index) -> absl::Status {
         if (subshape.IsTuple()) {
-          return OkStatus();
+          return absl::OkStatus();
         }
         if (!subshape.IsArray()) {
           return InvalidArgument("Shape cannot be serialiized: %s",
@@ -950,7 +950,7 @@ Shape ShapeUtil::PrependMajorDimension(int64_t bound, Shape shape) {
         } else {
           size += ByteSizeOfElements(subshape);
         }
-        return OkStatus();
+        return absl::OkStatus();
       }));
 
   return size;
@@ -971,7 +971,7 @@ Shape ShapeUtil::PrependMajorDimension(int64_t bound, Shape shape) {
       TF_RETURN_IF_ERROR(
           ValidateShapeWithOptionalLayoutInternal(element_shape));
     }
-    return OkStatus();
+    return absl::OkStatus();
   }
 
   // Non-tuple shape.
@@ -993,7 +993,7 @@ Shape ShapeUtil::PrependMajorDimension(int64_t bound, Shape shape) {
           primitive_util::LowercasePrimitiveTypeName(shape.element_type()),
           shape.ShortDebugString());
     }
-    return OkStatus();
+    return absl::OkStatus();
   }
 
   bool any_overflows = false;
@@ -1018,14 +1018,14 @@ Shape ShapeUtil::PrependMajorDimension(int64_t bound, Shape shape) {
   }
 
   TF_RETURN_IF_ERROR(ValidateShapeSize(shape));
-  return OkStatus();
+  return absl::OkStatus();
 }
 
 /* static */ absl::Status ShapeUtil::ValidateShapeSize(const Shape& shape) {
   VLOG(3) << "Validating shape size: " << ShapeUtil::HumanString(shape);
 
   if (!shape.IsArray()) {
-    return OkStatus();
+    return absl::OkStatus();
   }
 
   auto [extent_product, extent_overflow] =
@@ -1039,7 +1039,7 @@ Shape ShapeUtil::PrependMajorDimension(int64_t bound, Shape shape) {
   }
 
   VLOG(3) << "Shape size is valid: " << dense_shape_size;
-  return OkStatus();
+  return absl::OkStatus();
 }
 
 /* static */ absl::Status ShapeUtil::ValidateShapeWithOptionalLayout(
@@ -1858,7 +1858,7 @@ ShapeUtil::DecomposeBitcastToTrt(const Shape& input_shape,
     const ForEachVisitorFunction& visitor_function) {
   ForEachState s(shape, base, count, incr);
   if (s.IsZeroElementArray()) {
-    return OkStatus();
+    return absl::OkStatus();
   }
   // Allows handling R0 arrays, such that the visitor function will be called
   // once with the proper empty indexes.
@@ -1872,7 +1872,7 @@ ShapeUtil::DecomposeBitcastToTrt(const Shape& input_shape,
     // Increments dimensions in minor to major order.
     n = s.IncrementDim();
   }
-  return OkStatus();
+  return absl::OkStatus();
 }
 
 /* static */ void ShapeUtil::ForEachIndexInternalNoStatus(
@@ -2146,7 +2146,7 @@ absl::Status ShapeUtil::ByteStrides(const Shape& shape,
     strides.at(i) = stride;
     stride *= shape.dimensions(i);
   }
-  return OkStatus();
+  return absl::OkStatus();
 }
 
 /*static*/
diff --git a/third_party/xla/xla/shape_util.h b/third_party/xla/xla/shape_util.h
index d6688c573cdbc2..0339df04b20755 100644
--- a/third_party/xla/xla/shape_util.h
+++ b/third_party/xla/xla/shape_util.h
@@ -583,7 +583,7 @@ class ShapeUtil {
     ForEachSubshapeWithStatus(shape, [&](const Shape& subshape,
                                          const ShapeIndex& index) {
       fn(subshape, index);
-      return OkStatus();
+      return absl::OkStatus();
     }).IgnoreError();
   }
   template <typename Fn>
@@ -591,7 +591,7 @@ class ShapeUtil {
     ForEachMutableSubshapeWithStatus(shape, [&](Shape* subshape,
                                                 const ShapeIndex& index) {
       fn(subshape, index);
-      return OkStatus();
+      return absl::OkStatus();
     }).IgnoreError();
   }
 
@@ -648,7 +648,7 @@ class ShapeUtil {
     ForEachSubshapePostOrderWithStatus(shape, [&](const Shape& subshape,
                                                   const ShapeIndex& index) {
       fn(subshape, index);
-      return OkStatus();
+      return absl::OkStatus();
     }).IgnoreError();
   }
   template <typename Fn>
@@ -657,7 +657,7 @@ class ShapeUtil {
         shape,
         [&](Shape* subshape, const ShapeIndex& index) {
           fn(subshape, index);
-          return OkStatus();
+          return absl::OkStatus();
         })
         .IgnoreError();
   }
@@ -1057,7 +1057,7 @@ class ShapeUtil {
         index->pop_back();
       }
     }
-    return OkStatus();
+    return absl::OkStatus();
   }
 
   // Helper for ForEachSubshapePost which visits the subshapes of the given
@@ -1074,7 +1074,7 @@ class ShapeUtil {
       }
     }
     TF_RETURN_IF_ERROR(fn(shape, *index));
-    return OkStatus();
+    return absl::OkStatus();
   }
 
   // Keeps track of the iteration state for the ForEach...Internal routines
diff --git a/third_party/xla/xla/sharding_op_util.cc b/third_party/xla/xla/sharding_op_util.cc
index 5ef8d7d9b06d51..d2dc84ffc3d405 100644
--- a/third_party/xla/xla/sharding_op_util.cc
+++ b/third_party/xla/xla/sharding_op_util.cc
@@ -61,7 +61,7 @@ absl::Status ParseAttributes(absl::string_view opaque,
                                    attr_name);
     }
   }
-  return OkStatus();
+  return absl::OkStatus();
 }
 
 }  // namespace sharding_op_util
diff --git a/third_party/xla/xla/status_macros_test.cc b/third_party/xla/xla/status_macros_test.cc
index 817eb0a9f77c1f..99eaebdf96b697 100644
--- a/third_party/xla/xla/status_macros_test.cc
+++ b/third_party/xla/xla/status_macros_test.cc
@@ -29,17 +29,17 @@ namespace xla {
 
 absl::Status RetCheckFail() {
   TF_RET_CHECK(2 > 3);
-  return OkStatus();
+  return absl::OkStatus();
 }
 
 absl::Status RetCheckFailWithExtraMessage() {
   TF_RET_CHECK(2 > 3) << "extra message";
-  return OkStatus();
+  return absl::OkStatus();
 }
 
 absl::Status RetCheckSuccess() {
   TF_RET_CHECK(3 > 2);
-  return OkStatus();
+  return absl::OkStatus();
 }
 
 TEST(StatusMacros, RetCheckFailing) {
@@ -72,7 +72,7 @@ TEST(StatusMacros, AssignOrAssertOnOK) {
   EXPECT_EQ(42, result);
 }
 
-absl::Status ReturnStatusOK() { return OkStatus(); }
+absl::Status ReturnStatusOK() { return absl::OkStatus(); }
 
 absl::Status ReturnStatusError() { return (tsl::errors::Internal("foobar")); }
 
@@ -100,7 +100,7 @@ TEST(StatusMacros, AssignOrReturnSuccessfully) {
   absl::Status status = []() {
     TF_ASSIGN_OR_RETURN(int value, CreateIntSuccessfully());
     EXPECT_EQ(value, 42);
-    return OkStatus();
+    return absl::OkStatus();
   }();
   EXPECT_IS_OK(status);
 }
@@ -109,7 +109,7 @@ TEST(StatusMacros, AssignOrReturnUnsuccessfully) {
   absl::Status status = []() {
     TF_ASSIGN_OR_RETURN(int value, CreateIntUnsuccessfully());
     (void)value;
-    return OkStatus();
+    return absl::OkStatus();
   }();
   EXPECT_FALSE(status.ok());
   EXPECT_EQ(status.code(), tsl::error::INTERNAL);
diff --git a/third_party/xla/xla/util.h b/third_party/xla/xla/util.h
index e9c132cf010dce..91a2b50b38a26c 100644
--- a/third_party/xla/xla/util.h
+++ b/third_party/xla/xla/util.h
@@ -794,7 +794,7 @@ absl::Status EraseElementFromVector(std::vector<T>* container, const T& value) {
   auto it = std::find(container->begin(), container->end(), value);
   TF_RET_CHECK(it != container->end());
   container->erase(it);
-  return OkStatus();
+  return absl::OkStatus();
 }
 
 // Takes a sequence of unpacked n-bit values, such that every byte stores one

From 11fd24abfd3ca49807c04c844c7f7f69dcb6e570 Mon Sep 17 00:00:00 2001
From: Kyle Lucke <klucke@google.com>
Date: Mon, 20 May 2024 12:10:17 -0700
Subject: [PATCH 428/478] Move Event creation into each StreamExecutorInterface
 derived class.

PiperOrigin-RevId: 635531604
---
 .../stream_executor/stream_executor.cc        | 14 ++++------
 .../stream_executor/stream_executor_test.cc   | 19 ++++++--------
 .../xla/xla/backends/interpreter/executor.h   |  7 ++---
 .../xla/stream_executor/cuda/cuda_executor.cc | 10 +++----
 third_party/xla/xla/stream_executor/event.cc  | 18 ++++---------
 third_party/xla/xla/stream_executor/event.h   |  7 ++---
 .../xla/stream_executor/gpu/gpu_executor.h    |  4 +--
 .../xla/stream_executor/host/host_executor.cc |  8 ++----
 .../xla/stream_executor/host/host_executor.h  |  3 +--
 .../stream_executor/mock_stream_executor.h    |  3 ---
 .../xla/stream_executor/rocm/rocm_executor.cc | 10 +++----
 .../stream_executor_interface.h               |  7 -----
 .../stream_executor/stream_executor_pimpl.cc  |  9 -------
 .../stream_executor/stream_executor_pimpl.h   |  1 -
 third_party/xla/xla/stream_executor/tpu/BUILD |  1 +
 .../xla/stream_executor/tpu/tpu_executor.cc   | 26 +++++++++----------
 .../xla/stream_executor/tpu/tpu_executor.h    |  4 +--
 17 files changed, 48 insertions(+), 103 deletions(-)

diff --git a/tensorflow/c/experimental/stream_executor/stream_executor.cc b/tensorflow/c/experimental/stream_executor/stream_executor.cc
index 6e706366d47127..65b31f8cfb8f1c 100644
--- a/tensorflow/c/experimental/stream_executor/stream_executor.cc
+++ b/tensorflow/c/experimental/stream_executor/stream_executor.cc
@@ -407,10 +407,6 @@ class CStreamExecutor : public StreamExecutor {
     return stream_executor_->host_callback(&device_, stream_handle,
                                            &HostCallbackTrampoline, ctx);
   }
-  absl::Status AllocateEvent(Event* event) override {
-    DCHECK(event != nullptr);
-    return static_cast<CEvent*>(event->implementation())->Create();
-  }
   absl::Status DeallocateEvent(Event* event) override {
     static_cast<CEvent*>(event->implementation())->Destroy();
     return absl::OkStatus();
@@ -551,12 +547,12 @@ class CStreamExecutor : public StreamExecutor {
     return builder.Build();
   }
 
-  // Each call creates a new instance of the platform-specific implementation of
-  // the corresponding interface type.
-  std::unique_ptr<EventInterface> CreateEventImplementation() override {
-    return std::unique_ptr<EventInterface>(
-        new CEvent(&device_, stream_executor_));
+  absl::StatusOr<std::unique_ptr<Event>> CreateEvent() override {
+    auto c_event = std::make_unique<CEvent>(&device_, stream_executor_);
+    TF_RETURN_IF_ERROR(c_event->Create());
+    return std::make_unique<Event>(this, std::move(c_event));
   }
+
   absl::StatusOr<std::unique_ptr<Stream>> CreateStream(
       std::optional<std::variant<StreamPriority, int>> priority =
           std::nullopt) override {
diff --git a/tensorflow/c/experimental/stream_executor/stream_executor_test.cc b/tensorflow/c/experimental/stream_executor/stream_executor_test.cc
index 56f25a5811293e..680a1d9d1db1f5 100644
--- a/tensorflow/c/experimental/stream_executor/stream_executor_test.cc
+++ b/tensorflow/c/experimental/stream_executor/stream_executor_test.cc
@@ -342,11 +342,10 @@ TEST_F(StreamExecutorTest, CreateEvent) {
 
   StreamExecutor* executor = GetExecutor(0);
   ASSERT_FALSE(event_created);
-  Event* event = new Event(executor);
-  event->Init();
+  TF_ASSERT_OK_AND_ASSIGN(auto event, executor->CreateEvent());
   ASSERT_TRUE(event_created);
   ASSERT_FALSE(event_deleted);
-  delete event;
+  event.reset();
   ASSERT_TRUE(event_deleted);
 }
 
@@ -365,11 +364,10 @@ TEST_F(StreamExecutorTest, PollForEventStatus) {
   };
 
   StreamExecutor* executor = GetExecutor(0);
-  Event event(executor);
-  event.Init();
-  ASSERT_EQ(event.PollForStatus(), Event::Status::kComplete);
+  TF_ASSERT_OK_AND_ASSIGN(auto event, executor->CreateEvent());
+  ASSERT_EQ(event->PollForStatus(), Event::Status::kComplete);
   event_status = SE_EVENT_ERROR;
-  ASSERT_EQ(event.PollForStatus(), Event::Status::kError);
+  ASSERT_EQ(event->PollForStatus(), Event::Status::kError);
 }
 
 TEST_F(StreamExecutorTest, RecordAndWaitForEvent) {
@@ -403,14 +401,13 @@ TEST_F(StreamExecutorTest, RecordAndWaitForEvent) {
   };
 
   StreamExecutor* executor = GetExecutor(0);
-  Event event(executor);
-  event.Init();
+  TF_ASSERT_OK_AND_ASSIGN(auto event, executor->CreateEvent());
   TF_ASSERT_OK_AND_ASSIGN(auto stream, executor->CreateStream());
   ASSERT_FALSE(record_called);
-  TF_ASSERT_OK(stream->RecordEvent(&event));
+  TF_ASSERT_OK(stream->RecordEvent(event.get()));
   ASSERT_TRUE(record_called);
   ASSERT_FALSE(wait_called);
-  TF_ASSERT_OK(stream->WaitFor(&event));
+  TF_ASSERT_OK(stream->WaitFor(event.get()));
   ASSERT_TRUE(wait_called);
 }
 
diff --git a/third_party/xla/xla/backends/interpreter/executor.h b/third_party/xla/xla/backends/interpreter/executor.h
index 7ca85d7c22048f..b77ca8d120ec7c 100644
--- a/third_party/xla/xla/backends/interpreter/executor.h
+++ b/third_party/xla/xla/backends/interpreter/executor.h
@@ -114,8 +114,6 @@ class XlaInterpreterExecutor : public StreamExecutor {
   bool HostCallback(Stream *stream,
                     absl::AnyInvocable<absl::Status() &&> callback) override;
 
-  absl::Status AllocateEvent(Event *event) override { return absl::OkStatus(); }
-
   absl::Status DeallocateEvent(Event *event) override {
     return absl::OkStatus();
   }
@@ -156,9 +154,8 @@ class XlaInterpreterExecutor : public StreamExecutor {
   bool CanEnablePeerAccessTo(StreamExecutorInterface *other) override {
     return true;
   }
-
-  std::unique_ptr<EventInterface> CreateEventImplementation() override {
-    return nullptr;
+  absl::StatusOr<std::unique_ptr<Event>> CreateEvent() override {
+    return std::make_unique<Event>(this, nullptr);
   }
 
   absl::StatusOr<std::unique_ptr<Stream>> CreateStream(
diff --git a/third_party/xla/xla/stream_executor/cuda/cuda_executor.cc b/third_party/xla/xla/stream_executor/cuda/cuda_executor.cc
index cdf5e85bb2e2df..bf2918195cb359 100644
--- a/third_party/xla/xla/stream_executor/cuda/cuda_executor.cc
+++ b/third_party/xla/xla/stream_executor/cuda/cuda_executor.cc
@@ -758,10 +758,6 @@ bool GpuExecutor::HostCallback(Stream* stream,
   delete callback;
 }
 
-absl::Status GpuExecutor::AllocateEvent(Event* event) {
-  return AsGpuEvent(event)->Init();
-}
-
 absl::Status GpuExecutor::DeallocateEvent(Event* event) {
   return AsGpuEvent(event)->Destroy();
 }
@@ -950,8 +946,10 @@ absl::Status FillBlockDimLimit(GpuDeviceHandle device,
   return absl::OkStatus();
 }
 
-std::unique_ptr<EventInterface> GpuExecutor::CreateEventImplementation() {
-  return std::unique_ptr<EventInterface>(new GpuEvent(this));
+absl::StatusOr<std::unique_ptr<Event>> GpuExecutor::CreateEvent() {
+  auto gpu_event = std::make_unique<GpuEvent>(this);
+  TF_RETURN_IF_ERROR(gpu_event->Init());
+  return std::make_unique<Event>(this, std::move(gpu_event));
 }
 
 absl::StatusOr<std::unique_ptr<Stream>> GpuExecutor::CreateStream(
diff --git a/third_party/xla/xla/stream_executor/event.cc b/third_party/xla/xla/stream_executor/event.cc
index fc11f3c4797b35..5cca5e4c1e78f5 100644
--- a/third_party/xla/xla/stream_executor/event.cc
+++ b/third_party/xla/xla/stream_executor/event.cc
@@ -16,6 +16,8 @@ limitations under the License.
 #include "xla/stream_executor/event.h"
 
 #include <cstdint>
+#include <memory>
+#include <utility>
 
 #include "absl/log/log.h"
 #include "absl/status/status.h"
@@ -24,9 +26,9 @@ limitations under the License.
 
 namespace stream_executor {
 
-Event::Event(StreamExecutorInterface* stream_exec)
-    : stream_exec_(stream_exec),
-      implementation_(stream_exec_->CreateEventImplementation()) {}
+Event::Event(StreamExecutorInterface* stream_exec,
+             std::unique_ptr<EventInterface> implementation)
+    : stream_exec_(stream_exec), implementation_(std::move(implementation)) {}
 
 Event::~Event() {
   // Deal with nullptr implementation_, as this event may have been std::moved.
@@ -41,16 +43,6 @@ Event::~Event() {
 Event::Event(Event&&) = default;
 Event& Event::operator=(Event&&) = default;
 
-bool Event::Init() {
-  auto status = stream_exec_->AllocateEvent(this);
-  if (!status.ok()) {
-    LOG(ERROR) << status.message();
-    return false;
-  }
-
-  return true;
-}
-
 Event::Status Event::PollForStatus() {
   return stream_exec_->PollForEventStatus(this);
 }
diff --git a/third_party/xla/xla/stream_executor/event.h b/third_party/xla/xla/stream_executor/event.h
index c73bfa470e2166..64c871471178d3 100644
--- a/third_party/xla/xla/stream_executor/event.h
+++ b/third_party/xla/xla/stream_executor/event.h
@@ -43,15 +43,12 @@ class Event {
     kComplete,
   };
 
-  explicit Event(StreamExecutorInterface* stream_exec);  // NOLINT
+  Event(StreamExecutorInterface* stream_exec,
+        std::unique_ptr<EventInterface> implementation);
 
   // Releases any resources held by the Event object.
   ~Event();
 
-  // Performs any platform-specific or potentially error-generating
-  // initialization.
-  bool Init();
-
   // Returns the current Status for the event.
   Status PollForStatus();
 
diff --git a/third_party/xla/xla/stream_executor/gpu/gpu_executor.h b/third_party/xla/xla/stream_executor/gpu/gpu_executor.h
index c7bd541f2d0c8a..e72d60a0f582f0 100644
--- a/third_party/xla/xla/stream_executor/gpu/gpu_executor.h
+++ b/third_party/xla/xla/stream_executor/gpu/gpu_executor.h
@@ -233,8 +233,6 @@ class GpuExecutor : public StreamExecutor {
 
   bool CreateStreamDependency(Stream* dependent, Stream* other) override;
 
-  absl::Status AllocateEvent(Event* event) override;
-
   absl::Status DeallocateEvent(Event* event) override;
 
   absl::Status RecordEvent(Stream* stream, Event* event) override;
@@ -271,7 +269,7 @@ class GpuExecutor : public StreamExecutor {
 
   dnn::DnnSupport* AsDnn() override;
 
-  std::unique_ptr<EventInterface> CreateEventImplementation() override;
+  absl::StatusOr<std::unique_ptr<Event>> CreateEvent() override;
 
   absl::StatusOr<std::unique_ptr<Stream>> CreateStream(
       std::optional<std::variant<StreamPriority, int>> priority =
diff --git a/third_party/xla/xla/stream_executor/host/host_executor.cc b/third_party/xla/xla/stream_executor/host/host_executor.cc
index 28b69aa213a0e1..159d60d3357630 100644
--- a/third_party/xla/xla/stream_executor/host/host_executor.cc
+++ b/third_party/xla/xla/stream_executor/host/host_executor.cc
@@ -235,8 +235,8 @@ class HostEvent : public EventInterface {
   std::shared_ptr<absl::Notification> notification_;
 };
 
-std::unique_ptr<EventInterface> HostExecutor::CreateEventImplementation() {
-  return std::unique_ptr<EventInterface>(new HostEvent());
+absl::StatusOr<std::unique_ptr<Event>> HostExecutor::CreateEvent() {
+  return std::make_unique<Event>(this, std::make_unique<HostEvent>());
 }
 
 static HostEvent* AsHostEvent(Event* event) {
@@ -244,10 +244,6 @@ static HostEvent* AsHostEvent(Event* event) {
   return static_cast<HostEvent*>(event->implementation());
 }
 
-absl::Status HostExecutor::AllocateEvent(Event* /*event*/) {
-  return absl::OkStatus();
-}
-
 absl::Status HostExecutor::DeallocateEvent(Event* /*event*/) {
   return absl::OkStatus();
 }
diff --git a/third_party/xla/xla/stream_executor/host/host_executor.h b/third_party/xla/xla/stream_executor/host/host_executor.h
index b1ce6f9e42db2a..2ee91d60a466ee 100644
--- a/third_party/xla/xla/stream_executor/host/host_executor.h
+++ b/third_party/xla/xla/stream_executor/host/host_executor.h
@@ -119,7 +119,6 @@ class HostExecutor : public StreamExecutor {
   bool HostCallback(Stream* stream,
                     absl::AnyInvocable<absl::Status() &&> callback) override;
 
-  absl::Status AllocateEvent(Event* event) override;
   absl::Status DeallocateEvent(Event* event) override;
   absl::Status RecordEvent(Stream* stream, Event* event) override;
   absl::Status WaitForEvent(Stream* stream, Event* event) override;
@@ -149,7 +148,7 @@ class HostExecutor : public StreamExecutor {
     return true;
   }
 
-  std::unique_ptr<EventInterface> CreateEventImplementation() override;
+  absl::StatusOr<std::unique_ptr<Event>> CreateEvent() override;
 
   absl::StatusOr<std::unique_ptr<Stream>> CreateStream(
       std::optional<std::variant<StreamPriority, int>> priority) override;
diff --git a/third_party/xla/xla/stream_executor/mock_stream_executor.h b/third_party/xla/xla/stream_executor/mock_stream_executor.h
index 21f9147b9f6440..8f6fc02f863aad 100644
--- a/third_party/xla/xla/stream_executor/mock_stream_executor.h
+++ b/third_party/xla/xla/stream_executor/mock_stream_executor.h
@@ -133,7 +133,6 @@ class MockStreamExecutor : public StreamExecutorInterface {
   MOCK_METHOD(bool, HostCallback,
               (Stream * stream, absl::AnyInvocable<absl::Status() &&> callback),
               (override));
-  MOCK_METHOD(absl::Status, AllocateEvent, (Event * event), (override));
   MOCK_METHOD(absl::Status, DeallocateEvent, (Event * event), (override));
   MOCK_METHOD(absl::Status, RecordEvent, (Stream * stream, Event* event),
               (override));
@@ -161,8 +160,6 @@ class MockStreamExecutor : public StreamExecutorInterface {
   MOCK_METHOD(blas::BlasSupport*, AsBlas, (), (override));
   MOCK_METHOD(fft::FftSupport*, AsFft, (), (override));
   MOCK_METHOD(dnn::DnnSupport*, AsDnn, (), (override));
-  MOCK_METHOD(std::unique_ptr<EventInterface>, CreateEventImplementation, (),
-              (override));
   MOCK_METHOD(absl::StatusOr<std::unique_ptr<Kernel>>, CreateKernel, (),
               (override));
   MOCK_METHOD(absl::StatusOr<std::unique_ptr<CommandBuffer>>,
diff --git a/third_party/xla/xla/stream_executor/rocm/rocm_executor.cc b/third_party/xla/xla/stream_executor/rocm/rocm_executor.cc
index c9e81b77a94ea4..d8b29d23208ff2 100644
--- a/third_party/xla/xla/stream_executor/rocm/rocm_executor.cc
+++ b/third_party/xla/xla/stream_executor/rocm/rocm_executor.cc
@@ -660,10 +660,6 @@ bool GpuExecutor::HostCallback(Stream* stream,
   delete callback;
 }
 
-absl::Status GpuExecutor::AllocateEvent(Event* event) {
-  return AsGpuEvent(event)->Init();
-}
-
 absl::Status GpuExecutor::DeallocateEvent(Event* event) {
   return AsGpuEvent(event)->Destroy();
 }
@@ -856,8 +852,10 @@ absl::Status FillBlockDimLimit(GpuDeviceHandle device,
   return absl::OkStatus();
 }
 
-std::unique_ptr<EventInterface> GpuExecutor::CreateEventImplementation() {
-  return std::unique_ptr<EventInterface>(new GpuEvent(this));
+absl::StatusOr<std::unique_ptr<Event>> GpuExecutor::CreateEvent() {
+  auto gpu_event = std::make_unique<GpuEvent>(this);
+  TF_RETURN_IF_ERROR(gpu_event->Init());
+  return std::make_unique<Event>(this, std::move(gpu_event));
 }
 
 absl::StatusOr<std::unique_ptr<Stream>> GpuExecutor::CreateStream(
diff --git a/third_party/xla/xla/stream_executor/stream_executor_interface.h b/third_party/xla/xla/stream_executor/stream_executor_interface.h
index 628577725014d5..95b23ef1231594 100644
--- a/third_party/xla/xla/stream_executor/stream_executor_interface.h
+++ b/third_party/xla/xla/stream_executor/stream_executor_interface.h
@@ -263,9 +263,6 @@ class StreamExecutorInterface {
   virtual bool HostCallback(Stream* stream,
                             absl::AnyInvocable<absl::Status() &&> callback) = 0;
 
-  // Performs platform-specific allocation and initialization of an event.
-  virtual absl::Status AllocateEvent(Event* event) = 0;
-
   // Performs platform-specific deallocation and cleanup of an event.
   virtual absl::Status DeallocateEvent(Event* event) = 0;
 
@@ -357,10 +354,6 @@ class StreamExecutorInterface {
   // underlying platform.
   virtual dnn::DnnSupport* AsDnn() { return nullptr; }
 
-  // Each call creates a new instance of the platform-specific implementation of
-  // the corresponding interface type.
-  virtual std::unique_ptr<EventInterface> CreateEventImplementation() = 0;
-
   // Creates a new Kernel object.
   // TODO(klucke) Combine with GetKernel.
   virtual absl::StatusOr<std::unique_ptr<Kernel>> CreateKernel() {
diff --git a/third_party/xla/xla/stream_executor/stream_executor_pimpl.cc b/third_party/xla/xla/stream_executor/stream_executor_pimpl.cc
index 2866a7d7d50aeb..efd27f4ec49f48 100644
--- a/third_party/xla/xla/stream_executor/stream_executor_pimpl.cc
+++ b/third_party/xla/xla/stream_executor/stream_executor_pimpl.cc
@@ -78,13 +78,4 @@ const DeviceDescription& StreamExecutor::GetDeviceDescription() const {
   return *device_description_;
 }
 
-absl::StatusOr<std::unique_ptr<Event>> StreamExecutor::CreateEvent() {
-  auto event = std::make_unique<Event>(this);
-  if (event->Init()) {
-    return std::move(event);
-  }
-
-  return absl::InternalError("Failed to create event.");
-}
-
 }  // namespace stream_executor
diff --git a/third_party/xla/xla/stream_executor/stream_executor_pimpl.h b/third_party/xla/xla/stream_executor/stream_executor_pimpl.h
index 41546f59ea5321..f57ace700cc9a0 100644
--- a/third_party/xla/xla/stream_executor/stream_executor_pimpl.h
+++ b/third_party/xla/xla/stream_executor/stream_executor_pimpl.h
@@ -71,7 +71,6 @@ class StreamExecutor : public StreamExecutorInterface {
   const Platform* GetPlatform() const override { return platform_; }
   const DeviceDescription& GetDeviceDescription() const override;
   int64_t GetMemoryLimitBytes() const override { return memory_limit_bytes_; }
-  absl::StatusOr<std::unique_ptr<Event>> CreateEvent() override;
 
  private:
   // Reader/writer lock for mutable data structures on this StreamExecutor.
diff --git a/third_party/xla/xla/stream_executor/tpu/BUILD b/third_party/xla/xla/stream_executor/tpu/BUILD
index 2f205cf03a97e9..5de3a5589a76bb 100644
--- a/third_party/xla/xla/stream_executor/tpu/BUILD
+++ b/third_party/xla/xla/stream_executor/tpu/BUILD
@@ -323,6 +323,7 @@ cc_library(
         "@com_google_absl//absl/synchronization",
         "@com_google_absl//absl/types:span",
         "@local_tsl//tsl/platform:casts",
+        "@local_tsl//tsl/platform:errors",
         "@local_tsl//tsl/platform:logging",
         "@local_tsl//tsl/platform:status",
         "@local_tsl//tsl/platform:types",
diff --git a/third_party/xla/xla/stream_executor/tpu/tpu_executor.cc b/third_party/xla/xla/stream_executor/tpu/tpu_executor.cc
index 5b493b104d432d..4b515f330eb119 100644
--- a/third_party/xla/xla/stream_executor/tpu/tpu_executor.cc
+++ b/third_party/xla/xla/stream_executor/tpu/tpu_executor.cc
@@ -39,6 +39,7 @@ limitations under the License.
 #include "xla/stream_executor/tpu/tpu_stream.h"
 #include "xla/stream_executor/tpu/tpu_topology.h"
 #include "xla/tsl/c/tsl_status.h"
+#include "tsl/platform/errors.h"
 #include "tsl/platform/logging.h"  // IWYU pragma: keep
 #include "tsl/platform/status.h"
 
@@ -91,14 +92,6 @@ bool TpuExecutor::CreateStreamDependency(Stream* dependent, Stream* other) {
       get_stream(other->implementation()));
 }
 
-absl::Status TpuExecutor::AllocateEvent(Event* event) {
-  StatusHelper status;
-  auto se_event = tpu_platform().LookupEvent(event->implementation());
-  ExecutorApiFn()->TpuExecutor_AllocateEventFn(executor_, se_event,
-                                               status.c_status);
-  return status.status();
-}
-
 absl::Status TpuExecutor::DeallocateEvent(Event* event) {
   tpu_platform().EraseEvent(event->implementation());
   return absl::OkStatus();
@@ -142,12 +135,17 @@ absl::StatusOr<std::unique_ptr<Stream>> TpuExecutor::CreateStream(
   return std::move(stream);
 }
 
-// Called by Event::Event
-std::unique_ptr<EventInterface> TpuExecutor::CreateEventImplementation() {
-  SE_Event* tpu_event = ExecutorApiFn()->TpuEvent_NewFn(executor_);
-  auto ptr = std::make_unique<TpuEvent>(tpu_event);
-  tpu_platform().InsertEvent(ptr.get(), tpu_event);
-  return ptr;
+absl::StatusOr<std::unique_ptr<Event>> TpuExecutor::CreateEvent() {
+  SE_Event* se_event = ExecutorApiFn()->TpuEvent_NewFn(executor_);
+  auto tpu_event = std::make_unique<TpuEvent>(se_event);
+  tpu_platform().InsertEvent(tpu_event.get(), se_event);
+
+  StatusHelper status;
+  ExecutorApiFn()->TpuExecutor_AllocateEventFn(executor_, se_event,
+                                               status.c_status);
+  TF_RETURN_IF_ERROR(status.status());
+
+  return std::make_unique<Event>(this, std::move(tpu_event));
 }
 
 DeviceMemoryBase TpuExecutor::Allocate(uint64_t size, int64_t memory_space) {
diff --git a/third_party/xla/xla/stream_executor/tpu/tpu_executor.h b/third_party/xla/xla/stream_executor/tpu/tpu_executor.h
index db41a12f552dad..066db0c1dc319b 100644
--- a/third_party/xla/xla/stream_executor/tpu/tpu_executor.h
+++ b/third_party/xla/xla/stream_executor/tpu/tpu_executor.h
@@ -66,8 +66,6 @@ class TpuExecutor : public tensorflow::tpu::TpuExecutorInterface {
 
   DeviceMemoryBase Allocate(uint64_t size, int64_t memory_space) override;
 
-  absl::Status AllocateEvent(Event* event) override;
-
   absl::Status BlockHostUntilDone(Stream* stream) override;
 
   absl::StatusOr<std::unique_ptr<DeviceDescription>> CreateDeviceDescription()
@@ -102,7 +100,7 @@ class TpuExecutor : public tensorflow::tpu::TpuExecutorInterface {
       std::optional<std::variant<StreamPriority, int>> priority =
           std::nullopt) override;
 
-  std::unique_ptr<EventInterface> CreateEventImplementation() override;
+  absl::StatusOr<std::unique_ptr<Event>> CreateEvent() override;
 
   bool HostCallback(Stream* stream,
                     absl::AnyInvocable<absl::Status() &&> callback) override;

From 1dff6b31e9e3b3776f4b5504f5f460f9c666b985 Mon Sep 17 00:00:00 2001
From: Junwhan Ahn <junwhan@google.com>
Date: Mon, 20 May 2024 12:11:22 -0700
Subject: [PATCH 429/478] Remove `PjRtStreamExecutorMemorySpace` specialization
 from `PjRtStreamExecutorClient`

This allows us to use existing host memory spaces for GPUs without having to reimplement them for stream executor clients.

PiperOrigin-RevId: 635531851
---
 .../eager/context_distributed_manager.cc      |  2 -
 .../core/common_runtime/gpu/gpu_device.cc     |  2 -
 .../examples/axpy/stablehlo_compile_test.cc   | 13 ++--
 third_party/xla/xla/pjrt/gpu/BUILD            |  1 +
 .../xla/xla/pjrt/gpu/se_gpu_pjrt_client.cc    | 61 ++++++++++++-------
 .../xla/xla/pjrt/gpu/se_gpu_pjrt_client.h     | 14 +----
 .../xla/xla/pjrt/interpreter_device.cc        |  1 -
 .../xla/pjrt/pjrt_stream_executor_client.cc   | 24 +-------
 .../xla/pjrt/pjrt_stream_executor_client.h    |  9 +--
 .../pjrt/pjrt_stream_executor_client_test.cc  |  6 +-
 10 files changed, 54 insertions(+), 79 deletions(-)

diff --git a/tensorflow/core/common_runtime/eager/context_distributed_manager.cc b/tensorflow/core/common_runtime/eager/context_distributed_manager.cc
index 535427558d36f5..35c583927433de 100644
--- a/tensorflow/core/common_runtime/eager/context_distributed_manager.cc
+++ b/tensorflow/core/common_runtime/eager/context_distributed_manager.cc
@@ -351,11 +351,9 @@ absl::Status CreateClientOnce(
   }
 
   if (use_creation_info) {
-    auto memory_spaces = xla::BuildMemorySpaces(pjrt_devices);
     std::unique_ptr<xla::PjRtClient> pjrt_client =
         std::make_unique<xla::StreamExecutorGpuClient>(
             platform_name, info->local_client, std::move(pjrt_devices),
-            std::move(memory_spaces),
             /*process_index=*/node_id,
             /*allocator=*/std::move(info->allocator),
             /*host_memory_allocator=*/std::move(info->host_memory_allocator),
diff --git a/tensorflow/core/common_runtime/gpu/gpu_device.cc b/tensorflow/core/common_runtime/gpu/gpu_device.cc
index 16f071c16c37c8..6f8761b4985b23 100644
--- a/tensorflow/core/common_runtime/gpu/gpu_device.cc
+++ b/tensorflow/core/common_runtime/gpu/gpu_device.cc
@@ -1923,11 +1923,9 @@ Status BaseGPUDeviceFactory::CreateDevices(
 #else   // TENSORFLOW_USE_ROCM
       auto platform_name = xla::CudaName();
 #endif  // TENSORFLOW_USE_ROCM
-      auto memory_spaces = xla::BuildMemorySpaces(pjrt_devices);
       std::unique_ptr<xla::PjRtClient> pjrt_client =
           std::make_unique<xla::StreamExecutorGpuClient>(
               platform_name, xla_client, std::move(pjrt_devices),
-              std::move(memory_spaces),
               /*process_index=*/numa_node,
               /*allocator=*/std::move(allocator_adapter),
               /*host_memory_allocator=*/std::move(pjrt_gpu_host_allocator),
diff --git a/third_party/xla/xla/examples/axpy/stablehlo_compile_test.cc b/third_party/xla/xla/examples/axpy/stablehlo_compile_test.cc
index f8c54b33cb4d2f..7072f49a6544ca 100644
--- a/third_party/xla/xla/examples/axpy/stablehlo_compile_test.cc
+++ b/third_party/xla/xla/examples/axpy/stablehlo_compile_test.cc
@@ -72,13 +72,12 @@ TEST(StableHloAxpyTest, LoadAndRunCpuExecutable) {
 
   // The PjRtStreamExecutorClient will allow us to compile and execute
   // computations on the device we just configured.
-  auto pjrt_se_client = PjRtStreamExecutorClient(
-      "cpu", local_client, std::move(devices),
-      std::vector<std::unique_ptr<PjRtStreamExecutorMemorySpace>>(),
-      /*process_index=*/0, /*allocator=*/nullptr,
-      /*host_memory_allocator=*/nullptr,
-      /*should_stage_host_to_device_transfers=*/false,
-      /*gpu_run_options=*/nullptr);
+  auto pjrt_se_client =
+      PjRtStreamExecutorClient("cpu", local_client, std::move(devices),
+                               /*process_index=*/0, /*allocator=*/nullptr,
+                               /*host_memory_allocator=*/nullptr,
+                               /*should_stage_host_to_device_transfers=*/false,
+                               /*gpu_run_options=*/nullptr);
 
   // Read StableHLO program to string.
   std::string program_path = tsl::io::JoinPath(
diff --git a/third_party/xla/xla/pjrt/gpu/BUILD b/third_party/xla/xla/pjrt/gpu/BUILD
index 69eb9f9a635d64..f3223cb52380f6 100644
--- a/third_party/xla/xla/pjrt/gpu/BUILD
+++ b/third_party/xla/xla/pjrt/gpu/BUILD
@@ -87,6 +87,7 @@ cc_library(
         "//xla/stream_executor/integrations:device_mem_allocator",
         "//xla/stream_executor/integrations:tf_allocator_adapter",
         "//xla/tsl/util:env_var",
+        "@com_google_absl//absl/algorithm:container",
         "@com_google_absl//absl/base:core_headers",
         "@com_google_absl//absl/container:flat_hash_map",
         "@com_google_absl//absl/container:inlined_vector",
diff --git a/third_party/xla/xla/pjrt/gpu/se_gpu_pjrt_client.cc b/third_party/xla/xla/pjrt/gpu/se_gpu_pjrt_client.cc
index a976e9cd9e6a05..92682b59739146 100644
--- a/third_party/xla/xla/pjrt/gpu/se_gpu_pjrt_client.cc
+++ b/third_party/xla/xla/pjrt/gpu/se_gpu_pjrt_client.cc
@@ -28,6 +28,7 @@ limitations under the License.
 #include <variant>
 #include <vector>
 
+#include "absl/algorithm/container.h"
 #include "absl/base/thread_annotations.h"
 #include "absl/container/flat_hash_map.h"
 #include "absl/container/inlined_vector.h"
@@ -478,6 +479,43 @@ class AsyncHostToDeviceTransferManager
   }
 };
 
+StreamExecutorGpuClient::StreamExecutorGpuClient(
+    std::string platform_name, LocalClient* client,
+    std::vector<std::unique_ptr<PjRtStreamExecutorDevice>> devices,
+    int process_index, std::unique_ptr<se::DeviceMemoryAllocator> allocator,
+    std::unique_ptr<tsl::Allocator> host_memory_allocator,
+    bool should_stage_host_to_device_transfers,
+    std::unique_ptr<gpu::GpuExecutableRunOptions> gpu_run_options)
+    : xla::PjRtStreamExecutorClient(
+          platform_name, client, std::move(devices), process_index,
+          std::move(allocator), std::move(host_memory_allocator),
+          should_stage_host_to_device_transfers, std::move(gpu_run_options)),
+      topology_(xla::StreamExecutorGpuTopologyDescription::Create(
+          tsl::Fingerprint64(platform_name), platform_name,
+          devices_.back()->device_kind(), devices_)) {
+  for (auto* device : addressable_devices()) {
+    // Use the device id to construct a globally unique memory space id. We do
+    // not promise that memory space ids and device ids are the same.
+    const int id = device->id();
+    auto memory_space =
+        std::make_unique<StreamExecutorGpuHbmMemorySpace>(id, device);
+    tensorflow::down_cast<PjRtStreamExecutorDevice*>(device)->AttachMemorySpace(
+        memory_space.get());
+    owned_memory_spaces_.push_back(std::move(memory_space));
+  }
+  for (const std::unique_ptr<PjRtMemorySpace>& memory_space :
+       owned_memory_spaces_) {
+    memory_spaces_.push_back(memory_space.get());
+  }
+
+  // We don't promise anything about the order of memory spaces, but this
+  // sorting is done for consistency with the device list that's sorted above.
+  absl::c_sort(memory_spaces_,
+               [](const PjRtMemorySpace* a, const PjRtMemorySpace* b) {
+                 return a->id() < b->id();
+               });
+}
+
 absl::string_view StreamExecutorGpuClient::platform_version() const {
 #define STRINGIFY2(X) #X
 #define STRINGIFY(X) STRINGIFY2(X)
@@ -992,23 +1030,6 @@ absl::Status BuildDistributedDevices(
   return OkStatus();
 }
 
-std::vector<std::unique_ptr<PjRtStreamExecutorMemorySpace>> BuildMemorySpaces(
-    absl::Span<const std::unique_ptr<PjRtStreamExecutorDevice>> devices) {
-  std::vector<std::unique_ptr<PjRtStreamExecutorMemorySpace>> memory_spaces;
-  for (const auto& device : devices) {
-    if (device->IsAddressable()) {
-      // Use the device id to construct a globally unique memory space id. We
-      // do not promise that memory space ids and device ids are the same.
-      const int id = device->id();
-      auto memory_space =
-          std::make_unique<StreamExecutorGpuHbmMemorySpace>(id, device.get());
-      device->AttachMemorySpace(memory_space.get());
-      memory_spaces.push_back(std::move(memory_space));
-    }
-  }
-  return memory_spaces;
-}
-
 std::string MakeComputeCapabilityString(const se::DeviceDescription* desc) {
   se::GpuComputeCapability cc = desc->gpu_compute_capability();
   if (std::holds_alternative<se::CudaComputeCapability>(cc)) {
@@ -1142,12 +1163,10 @@ absl::StatusOr<std::unique_ptr<PjRtClient>> GetStreamExecutorGpuClient(
       pjrt_platform_name, std::move(local_device_states), options.node_id,
       options.num_nodes, &devices, gpu_run_options.get(), kv_store,
       options.enable_mock_nccl));
-  auto memory_spaces = BuildMemorySpaces(devices);
 
   return std::unique_ptr<PjRtClient>(std::make_unique<StreamExecutorGpuClient>(
-      pjrt_platform_name, xla_client, std::move(devices),
-      std::move(memory_spaces), options.node_id, std::move(allocator),
-      std::move(host_memory_allocator),
+      pjrt_platform_name, xla_client, std::move(devices), options.node_id,
+      std::move(allocator), std::move(host_memory_allocator),
       options.should_stage_host_to_device_transfers,
       std::move(gpu_run_options)));
 }
diff --git a/third_party/xla/xla/pjrt/gpu/se_gpu_pjrt_client.h b/third_party/xla/xla/pjrt/gpu/se_gpu_pjrt_client.h
index 79af364f3bea33..59efaf3226fd40 100644
--- a/third_party/xla/xla/pjrt/gpu/se_gpu_pjrt_client.h
+++ b/third_party/xla/xla/pjrt/gpu/se_gpu_pjrt_client.h
@@ -205,19 +205,10 @@ class StreamExecutorGpuClient : public xla::PjRtStreamExecutorClient {
   StreamExecutorGpuClient(
       std::string platform_name, LocalClient* client,
       std::vector<std::unique_ptr<PjRtStreamExecutorDevice>> devices,
-      std::vector<std::unique_ptr<PjRtStreamExecutorMemorySpace>> memory_spaces,
       int process_index, std::unique_ptr<se::DeviceMemoryAllocator> allocator,
       std::unique_ptr<tsl::Allocator> host_memory_allocator,
       bool should_stage_host_to_device_transfers,
-      std::unique_ptr<gpu::GpuExecutableRunOptions> gpu_run_options)
-      : xla::PjRtStreamExecutorClient(
-            platform_name, client, std::move(devices), std::move(memory_spaces),
-            process_index, std::move(allocator),
-            std::move(host_memory_allocator),
-            should_stage_host_to_device_transfers, std::move(gpu_run_options)),
-        topology_(xla::StreamExecutorGpuTopologyDescription::Create(
-            tsl::Fingerprint64(platform_name), platform_name,
-            devices_.back()->device_kind(), devices_)) {}
+      std::unique_ptr<gpu::GpuExecutableRunOptions> gpu_run_options);
 
   absl::StatusOr<xla::DeviceAssignment> GetDefaultDeviceAssignment(
       int num_replicas, int num_partitions) const override;
@@ -283,9 +274,6 @@ absl::Status BuildDistributedDevices(
     absl::Duration get_local_topology_timeout = absl::Minutes(2),
     absl::Duration get_global_topology_timeout = absl::Minutes(5));
 
-std::vector<std::unique_ptr<PjRtStreamExecutorMemorySpace>> BuildMemorySpaces(
-    absl::Span<const std::unique_ptr<PjRtStreamExecutorDevice>> devices);
-
 struct GpuClientOptions {
   GpuAllocatorConfig allocator_config;
 
diff --git a/third_party/xla/xla/pjrt/interpreter_device.cc b/third_party/xla/xla/pjrt/interpreter_device.cc
index cfa12b96392645..52b52d15a22033 100644
--- a/third_party/xla/xla/pjrt/interpreter_device.cc
+++ b/third_party/xla/xla/pjrt/interpreter_device.cc
@@ -56,7 +56,6 @@ absl::StatusOr<std::unique_ptr<PjRtClient>> GetInterpreterClient() {
 
   return std::unique_ptr<PjRtClient>(std::make_unique<PjRtStreamExecutorClient>(
       "interpreter", client, std::move(devices),
-      std::vector<std::unique_ptr<PjRtStreamExecutorMemorySpace>>(),
       /*process_index=*/0, /*allocator=*/nullptr,
       /*host_memory_allocator=*/nullptr,
       /*should_stage_host_to_device_transfers=*/false,
diff --git a/third_party/xla/xla/pjrt/pjrt_stream_executor_client.cc b/third_party/xla/xla/pjrt/pjrt_stream_executor_client.cc
index 59bb40d776a86e..4b9a0ac56d0a96 100644
--- a/third_party/xla/xla/pjrt/pjrt_stream_executor_client.cc
+++ b/third_party/xla/xla/pjrt/pjrt_stream_executor_client.cc
@@ -154,12 +154,9 @@ namespace xla {
 PjRtStreamExecutorMemorySpace::PjRtStreamExecutorMemorySpace(
     int id, PjRtDevice* device, absl::string_view kind, int kind_id)
     : id_(id), device_(device), kind_(kind), kind_id_(kind_id) {
+  DCHECK(device_ != nullptr && device_->client() != nullptr);
+  auto* client = device_->client();
   to_string_ = absl::StrFormat("MEMORY_SPACE_%i", id_);
-}
-
-void PjRtStreamExecutorMemorySpace::SetClient(PjRtClient* client) {
-  // We have to define debug_string_ here because process_index() and
-  // platform_name() requires client to be initialized.
   debug_string_ = absl::StrFormat(
       "PjRtStreamExecutorMemory(id=%i, process_index=%i, client=%s)", id_,
       client->process_index(), client->platform_name());
@@ -232,7 +229,6 @@ class CpuAllocator : public tsl::Allocator {
 PjRtStreamExecutorClient::PjRtStreamExecutorClient(
     std::string platform_name, LocalClient* client,
     std::vector<std::unique_ptr<PjRtStreamExecutorDevice>> devices,
-    std::vector<std::unique_ptr<PjRtStreamExecutorMemorySpace>> memory_spaces,
     int process_index, std::unique_ptr<se::DeviceMemoryAllocator> allocator,
     std::unique_ptr<tsl::Allocator> host_memory_allocator,
     bool should_stage_host_to_device_transfers,
@@ -244,7 +240,6 @@ PjRtStreamExecutorClient::PjRtStreamExecutorClient(
       owned_allocator_(std::move(allocator)),
       owned_devices_(std::move(devices)),
       process_index_(process_index),
-      owned_memory_spaces_(std::move(memory_spaces)),
       should_stage_host_to_device_transfers_(
           should_stage_host_to_device_transfers),
       gpu_run_options_(std::move(gpu_run_options)),
@@ -281,18 +276,6 @@ PjRtStreamExecutorClient::PjRtStreamExecutorClient(
                [](const PjRtDevice* a, const PjRtDevice* b) {
                  return a->local_device_id() < b->local_device_id();
                });
-
-  for (const std::unique_ptr<PjRtStreamExecutorMemorySpace>& memory_space :
-       owned_memory_spaces_) {
-    memory_spaces_.push_back(memory_space.get());
-    memory_space->SetClient(this);
-  }
-  // We don't promise anything about the order of memory spaces, but this
-  // sorting is done for consistency with the device list that's sorted above.
-  absl::c_sort(memory_spaces_,
-               [](const PjRtMemorySpace* a, const PjRtMemorySpace* b) {
-                 return a->id() < b->id();
-               });
 }
 
 absl::StatusOr<DeviceAssignment>
@@ -1401,8 +1384,7 @@ PjRtStreamExecutorBuffer::PjRtStreamExecutorBuffer(
     : client_(tensorflow::down_cast<PjRtStreamExecutorClient*>(client)),
       on_device_shape_(std::move(on_device_shape)),
       device_(tensorflow::down_cast<PjRtStreamExecutorDevice*>(device)),
-      memory_space_(
-          tensorflow::down_cast<PjRtStreamExecutorMemorySpace*>(memory_space)),
+      memory_space_(memory_space),
       device_buffer_(std::move(device_buffer)) {
   for (int i = 0; i < ScopedHold::Type::kMaxValue; ++i) {
     holds_[i] = 0;
diff --git a/third_party/xla/xla/pjrt/pjrt_stream_executor_client.h b/third_party/xla/xla/pjrt/pjrt_stream_executor_client.h
index cac6e2691f3b6c..ad33f57e6cd5c8 100644
--- a/third_party/xla/xla/pjrt/pjrt_stream_executor_client.h
+++ b/third_party/xla/xla/pjrt/pjrt_stream_executor_client.h
@@ -228,9 +228,6 @@ class PjRtStreamExecutorMemorySpace : public PjRtMemorySpace {
   PjRtStreamExecutorMemorySpace(int id, PjRtDevice* device,
                                 absl::string_view kind, int kind_id);
 
-  // Must set client exactly once.
-  void SetClient(PjRtClient* client);
-
   PjRtClient* client() const override { return device_->client(); }
 
   absl::Span<PjRtDevice* const> devices() const override {
@@ -262,7 +259,6 @@ class PjRtStreamExecutorClient : public PjRtClient {
   explicit PjRtStreamExecutorClient(
       std::string platform_name, LocalClient* client,
       std::vector<std::unique_ptr<PjRtStreamExecutorDevice>> devices,
-      std::vector<std::unique_ptr<PjRtStreamExecutorMemorySpace>> memory_spaces,
       int process_index, std::unique_ptr<se::DeviceMemoryAllocator> allocator,
       std::unique_ptr<tsl::Allocator> host_memory_allocator,
       bool should_stage_host_to_device_transfers,
@@ -506,8 +502,7 @@ class PjRtStreamExecutorClient : public PjRtClient {
   std::vector<PjRtDevice*> addressable_devices_;
   int process_index_;
 
-  std::vector<std::unique_ptr<PjRtStreamExecutorMemorySpace>>
-      owned_memory_spaces_;
+  std::vector<std::unique_ptr<PjRtMemorySpace>> owned_memory_spaces_;
   // Pointers to `owned_memory_spaces_`.
   std::vector<PjRtMemorySpace*> memory_spaces_;
 
@@ -867,7 +862,7 @@ class PjRtStreamExecutorBuffer : public PjRtBuffer {
   PjRtStreamExecutorClient* const client_;
   const Shape on_device_shape_;
   PjRtStreamExecutorDevice* const device_;
-  PjRtStreamExecutorMemorySpace* const memory_space_;
+  PjRtMemorySpace* const memory_space_;
 
   mutable absl::Mutex mu_;
   std::shared_ptr<TrackedDeviceBuffer> device_buffer_ ABSL_GUARDED_BY(mu_);
diff --git a/third_party/xla/xla/pjrt/pjrt_stream_executor_client_test.cc b/third_party/xla/xla/pjrt/pjrt_stream_executor_client_test.cc
index 26df8698348323..d34d5c3c54740f 100644
--- a/third_party/xla/xla/pjrt/pjrt_stream_executor_client_test.cc
+++ b/third_party/xla/xla/pjrt/pjrt_stream_executor_client_test.cc
@@ -56,14 +56,10 @@ absl::StatusOr<std::unique_ptr<PjRtStreamExecutorClient>> GetClient() {
       /*allow_event_reuse=*/false, /*use_callback_stream=*/false);
   auto device = std::make_unique<PjRtStreamExecutorDevice>(
       0, std::move(device_state), "cpu");
-  auto memory_space = std::make_unique<PjRtStreamExecutorMemorySpace>(
-      0, device.get(), "device", 1234);
   std::vector<std::unique_ptr<PjRtStreamExecutorDevice>> devices;
   devices.emplace_back(std::move(device));
-  std::vector<std::unique_ptr<PjRtStreamExecutorMemorySpace>> memory_spaces;
-  memory_spaces.emplace_back(std::move(memory_space));
   return std::make_unique<PjRtStreamExecutorClient>(
-      "cpu", local_client, std::move(devices), std::move(memory_spaces),
+      "cpu", local_client, std::move(devices),
       /*process_index=*/0, /*allocator=*/nullptr,
       /*host_memory_allocator=*/nullptr,
       /*should_stage_host_to_device_transfers=*/false,

From c00fe0a1dbbbadc58fd2229dbb80b402e0af83c6 Mon Sep 17 00:00:00 2001
From: Kyle Lucke <klucke@google.com>
Date: Mon, 20 May 2024 12:17:07 -0700
Subject: [PATCH 430/478] Automated Code Change

PiperOrigin-RevId: 635533415
---
 tensorflow/compiler/mlir/lite/transforms/legalize_tf.cc | 1 -
 tensorflow/core/tpu/kernels/tpu_execute_op.cc           | 1 -
 tensorflow/core/tpu/tpu_execute.cc                      | 1 -
 3 files changed, 3 deletions(-)

diff --git a/tensorflow/compiler/mlir/lite/transforms/legalize_tf.cc b/tensorflow/compiler/mlir/lite/transforms/legalize_tf.cc
index b18227bdddc1ee..2011b6d33ccd45 100644
--- a/tensorflow/compiler/mlir/lite/transforms/legalize_tf.cc
+++ b/tensorflow/compiler/mlir/lite/transforms/legalize_tf.cc
@@ -58,7 +58,6 @@ limitations under the License.
 #include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops_a_m.h"
 #include "tensorflow/compiler/mlir/tensorflow/transforms/lower_tf.h"
 #include "tensorflow/compiler/mlir/tensorflow/utils/mangling_util.h"
-#include "xla/status.h"
 #include "xla/statusor.h"
 #include "tensorflow/core/framework/tensor.pb.h"
 #include "tensorflow/core/framework/tensor_shape.pb.h"
diff --git a/tensorflow/core/tpu/kernels/tpu_execute_op.cc b/tensorflow/core/tpu/kernels/tpu_execute_op.cc
index 7c23a7558eea05..7a9812be9ae43b 100644
--- a/tensorflow/core/tpu/kernels/tpu_execute_op.cc
+++ b/tensorflow/core/tpu/kernels/tpu_execute_op.cc
@@ -42,7 +42,6 @@ limitations under the License.
 #include "xla/shape.h"
 #include "xla/shape_tree.h"
 #include "xla/shape_util.h"
-#include "xla/status.h"
 #include "xla/status_macros.h"
 #include "xla/statusor.h"
 #include "xla/stream_executor/device_memory.h"
diff --git a/tensorflow/core/tpu/tpu_execute.cc b/tensorflow/core/tpu/tpu_execute.cc
index 66b756213c46a6..8daf680e09314f 100644
--- a/tensorflow/core/tpu/tpu_execute.cc
+++ b/tensorflow/core/tpu/tpu_execute.cc
@@ -46,7 +46,6 @@ limitations under the License.
 #include "xla/shape_layout.h"
 #include "xla/shape_tree.h"
 #include "xla/shape_util.h"
-#include "xla/status.h"
 #include "xla/status_macros.h"
 #include "xla/statusor.h"
 #include "xla/stream_executor/device_memory.h"

From 42f89ec974dfbb18a1166e8464296a7858e93724 Mon Sep 17 00:00:00 2001
From: Clive Verghese <cliveverghese@google.com>
Date: Mon, 20 May 2024 12:39:41 -0700
Subject: [PATCH 431/478] Cleanup dependency on tracing.h

PiperOrigin-RevId: 635539097
---
 tensorflow/compiler/jit/xla_device.cc              |  1 -
 tensorflow/core/debug/debug_grpc_io_utils_test.cc  |  1 -
 tensorflow/core/debug/debug_grpc_testlib.cc        |  1 -
 .../core/distributed_runtime/master_session.cc     |  2 +-
 tensorflow/core/distributed_runtime/rpc/BUILD      |  1 +
 .../distributed_runtime/rpc/grpc_master_service.cc |  1 -
 .../distributed_runtime/rpc/grpc_remote_master.cc  |  4 ++--
 .../distributed_runtime/rpc/grpc_remote_worker.cc  |  1 -
 .../distributed_runtime/rpc/grpc_worker_service.cc |  2 +-
 tensorflow/core/distributed_runtime/worker.cc      |  2 +-
 tensorflow/core/framework/dataset.h                |  1 -
 tensorflow/core/framework/run_handler.cc           | 13 +++++++------
 tensorflow/core/kernels/collective_nccl.cc         |  1 -
 .../core/kernels/collective_nccl_all_to_all.cc     |  1 -
 .../core/kernels/collective_nccl_broadcaster.cc    |  1 -
 .../core/kernels/collective_nccl_gatherer.cc       |  1 -
 tensorflow/core/kernels/collective_nccl_reducer.cc |  1 -
 .../data/experimental/map_and_batch_dataset_op.cc  |  1 -
 tensorflow/core/kernels/function_ops.cc            |  1 -
 tensorflow/core/nccl/collective_communicator.cc    |  1 -
 tensorflow/core/tfrt/run_handler_thread_pool/BUILD |  1 +
 .../tfrt/run_handler_thread_pool/run_handler.cc    | 14 +++++++-------
 22 files changed, 21 insertions(+), 32 deletions(-)

diff --git a/tensorflow/compiler/jit/xla_device.cc b/tensorflow/compiler/jit/xla_device.cc
index b5b0c16422ccab..471f54571d2b53 100644
--- a/tensorflow/compiler/jit/xla_device.cc
+++ b/tensorflow/compiler/jit/xla_device.cc
@@ -52,7 +52,6 @@ limitations under the License.
 #include "tensorflow/core/lib/core/status.h"
 #include "tensorflow/core/platform/logging.h"
 #include "tensorflow/core/platform/stream_executor_no_cuda.h"
-#include "tensorflow/core/platform/tracing.h"
 #include "tensorflow/core/profiler/lib/traceme.h"
 #include "tensorflow/core/public/session_options.h"
 #include "tensorflow/core/public/version.h"
diff --git a/tensorflow/core/debug/debug_grpc_io_utils_test.cc b/tensorflow/core/debug/debug_grpc_io_utils_test.cc
index 235bdaed604270..3eaf3651126528 100644
--- a/tensorflow/core/debug/debug_grpc_io_utils_test.cc
+++ b/tensorflow/core/debug/debug_grpc_io_utils_test.cc
@@ -24,7 +24,6 @@ limitations under the License.
 #include "tensorflow/core/lib/strings/stringprintf.h"
 #include "tensorflow/core/platform/env.h"
 #include "tensorflow/core/platform/test.h"
-#include "tensorflow/core/platform/tracing.h"
 
 namespace tensorflow {
 
diff --git a/tensorflow/core/debug/debug_grpc_testlib.cc b/tensorflow/core/debug/debug_grpc_testlib.cc
index 4927caf5a3285a..2bc06061c459f7 100644
--- a/tensorflow/core/debug/debug_grpc_testlib.cc
+++ b/tensorflow/core/debug/debug_grpc_testlib.cc
@@ -23,7 +23,6 @@ limitations under the License.
 #include "tensorflow/core/lib/strings/str_util.h"
 #include "tensorflow/core/platform/env.h"
 #include "tensorflow/core/platform/protobuf.h"
-#include "tensorflow/core/platform/tracing.h"
 
 namespace tensorflow {
 
diff --git a/tensorflow/core/distributed_runtime/master_session.cc b/tensorflow/core/distributed_runtime/master_session.cc
index 5593963988d9e5..41ce2df923cad3 100644
--- a/tensorflow/core/distributed_runtime/master_session.cc
+++ b/tensorflow/core/distributed_runtime/master_session.cc
@@ -60,10 +60,10 @@ limitations under the License.
 #include "tensorflow/core/platform/logging.h"
 #include "tensorflow/core/platform/macros.h"
 #include "tensorflow/core/platform/mutex.h"
-#include "tensorflow/core/platform/tracing.h"
 #include "tensorflow/core/protobuf/config.pb.h"
 #include "tensorflow/core/public/session_options.h"
 #include "tensorflow/core/util/device_name_utils.h"
+#include "tsl/platform/tracing.h"
 #include "tsl/protobuf/coordination_config.pb.h"
 
 namespace tensorflow {
diff --git a/tensorflow/core/distributed_runtime/rpc/BUILD b/tensorflow/core/distributed_runtime/rpc/BUILD
index a9f2370ab81fa1..959a8abae6518a 100644
--- a/tensorflow/core/distributed_runtime/rpc/BUILD
+++ b/tensorflow/core/distributed_runtime/rpc/BUILD
@@ -241,6 +241,7 @@ cc_library(
         "//tensorflow/core/profiler/lib:traceme",
         "//tensorflow/core/protobuf:master_proto_cc",
         "@com_google_absl//absl/time",
+        "@local_tsl//tsl/platform:env",
         "@local_tsl//tsl/platform:retrying_utils",
     ],
     alwayslink = 1,
diff --git a/tensorflow/core/distributed_runtime/rpc/grpc_master_service.cc b/tensorflow/core/distributed_runtime/rpc/grpc_master_service.cc
index abc9d12969147e..de9bd049d0d468 100644
--- a/tensorflow/core/distributed_runtime/rpc/grpc_master_service.cc
+++ b/tensorflow/core/distributed_runtime/rpc/grpc_master_service.cc
@@ -41,7 +41,6 @@ limitations under the License.
 #include "tensorflow/core/distributed_runtime/rpc/grpc_util.h"
 #include "tensorflow/core/platform/logging.h"
 #include "tensorflow/core/platform/macros.h"
-#include "tensorflow/core/platform/tracing.h"
 #include "tensorflow/core/profiler/lib/traceme.h"
 #include "tensorflow/core/protobuf/master.pb.h"
 
diff --git a/tensorflow/core/distributed_runtime/rpc/grpc_remote_master.cc b/tensorflow/core/distributed_runtime/rpc/grpc_remote_master.cc
index c47dc31fc3ef4e..d180b4dc236451 100644
--- a/tensorflow/core/distributed_runtime/rpc/grpc_remote_master.cc
+++ b/tensorflow/core/distributed_runtime/rpc/grpc_remote_master.cc
@@ -27,10 +27,10 @@ limitations under the License.
 #include "tensorflow/core/lib/core/status.h"
 #include "tensorflow/core/lib/strings/strcat.h"
 #include "tensorflow/core/platform/env.h"
-#include "tensorflow/core/platform/tracing.h"
 #include "tensorflow/core/profiler/lib/traceme.h"
 #include "tensorflow/core/protobuf/master.pb.h"
 #include "tsl/platform/retrying_utils.h"
+#include "tsl/platform/tracing.h"
 
 namespace tensorflow {
 
@@ -116,7 +116,7 @@ class GrpcRemoteMaster : public MasterInterface {
   // Start tracing, attaching a unique ID to both the trace and the RPC.
   tsl::profiler::TraceMe* NewTraceRpc(StringPiece name,
                                       ::grpc::ClientContext* ctx) {
-    string trace_id = strings::StrCat(tracing::GetUniqueArg());
+    string trace_id = strings::StrCat(tsl::tracing::GetUniqueArg());
     ctx->AddMetadata(GrpcIdKey(), trace_id);
     return new tsl::profiler::TraceMe(
         [&] { return strings::StrCat(name, ":", trace_id); },
diff --git a/tensorflow/core/distributed_runtime/rpc/grpc_remote_worker.cc b/tensorflow/core/distributed_runtime/rpc/grpc_remote_worker.cc
index 05bf31918ef543..d1ad012a28f222 100644
--- a/tensorflow/core/distributed_runtime/rpc/grpc_remote_worker.cc
+++ b/tensorflow/core/distributed_runtime/rpc/grpc_remote_worker.cc
@@ -33,7 +33,6 @@ limitations under the License.
 #include "tensorflow/core/lib/core/threadpool.h"
 #include "tensorflow/core/lib/strings/str_util.h"
 #include "tensorflow/core/platform/logging.h"
-#include "tensorflow/core/platform/tracing.h"
 #include "tensorflow/core/protobuf/transport_options.pb.h"
 #include "tensorflow/core/protobuf/worker.pb.h"
 #include "tensorflow/core/util/env_var.h"
diff --git a/tensorflow/core/distributed_runtime/rpc/grpc_worker_service.cc b/tensorflow/core/distributed_runtime/rpc/grpc_worker_service.cc
index a990f41af261a2..3fac884a73bd11 100644
--- a/tensorflow/core/distributed_runtime/rpc/grpc_worker_service.cc
+++ b/tensorflow/core/distributed_runtime/rpc/grpc_worker_service.cc
@@ -52,10 +52,10 @@ limitations under the License.
 #include "tensorflow/core/lib/strings/stringprintf.h"
 #include "tensorflow/core/platform/logging.h"
 #include "tensorflow/core/platform/mutex.h"
-#include "tensorflow/core/platform/tracing.h"
 #include "tensorflow/core/profiler/lib/scoped_memory_debug_annotation.h"
 #include "tensorflow/core/protobuf/transport_options.pb.h"
 #include "tensorflow/core/protobuf/worker.pb.h"
+#include "tsl/platform/tracing.h"
 #include "tsl/protobuf/rpc_options.pb.h"
 
 namespace tensorflow {
diff --git a/tensorflow/core/distributed_runtime/worker.cc b/tensorflow/core/distributed_runtime/worker.cc
index 42708852b36f36..0922b04de0b0f8 100644
--- a/tensorflow/core/distributed_runtime/worker.cc
+++ b/tensorflow/core/distributed_runtime/worker.cc
@@ -27,8 +27,8 @@ limitations under the License.
 #include "tensorflow/core/distributed_runtime/tensor_coding.h"
 #include "tensorflow/core/distributed_runtime/worker_session.h"
 #include "tensorflow/core/framework/collective.h"
-#include "tensorflow/core/platform/tracing.h"
 #include "tensorflow/core/profiler/lib/device_profiler_session.h"
+#include "tsl/platform/tracing.h"
 #include "tsl/protobuf/distributed_runtime_payloads.pb.h"
 
 namespace tensorflow {
diff --git a/tensorflow/core/framework/dataset.h b/tensorflow/core/framework/dataset.h
index effa997eca60f9..394aa9a2d6e86e 100644
--- a/tensorflow/core/framework/dataset.h
+++ b/tensorflow/core/framework/dataset.h
@@ -60,7 +60,6 @@ limitations under the License.
 #include "tensorflow/core/platform/mutex.h"
 #include "tensorflow/core/platform/refcount.h"
 #include "tensorflow/core/platform/status.h"
-#include "tensorflow/core/platform/tracing.h"
 #include "tsl/framework/allocator.h"
 #include "tsl/platform/errors.h"
 #include "tsl/platform/thread_annotations.h"
diff --git a/tensorflow/core/framework/run_handler.cc b/tensorflow/core/framework/run_handler.cc
index 5ff63469ae4aa9..e68025643817d8 100644
--- a/tensorflow/core/framework/run_handler.cc
+++ b/tensorflow/core/framework/run_handler.cc
@@ -31,8 +31,8 @@ limitations under the License.
 #include "tensorflow/core/platform/mutex.h"
 #include "tensorflow/core/platform/numa.h"
 #include "tensorflow/core/platform/setround.h"
-#include "tensorflow/core/platform/tracing.h"
 #include "tensorflow/core/profiler/lib/traceme.h"
+#include "tsl/platform/tracing.h"
 
 namespace tensorflow {
 namespace {
@@ -67,9 +67,10 @@ RunHandlerEnvironment::EnvThread* RunHandlerEnvironment::CreateThread(
 RunHandlerEnvironment::Task RunHandlerEnvironment::CreateTask(
     std::function<void()> f) {
   uint64 id = 0;
-  if (tracing::EventCollector::IsEnabled()) {
-    id = tracing::GetUniqueArg();
-    tracing::RecordEvent(tracing::EventCategory::kScheduleClosure, id);
+  if (tsl::tracing::EventCollector::IsEnabled()) {
+    id = tsl::tracing::GetUniqueArg();
+    tsl::tracing::RecordEvent(tsl::tracing::EventCategory::kScheduleClosure,
+                              id);
   }
   return Task{
       std::unique_ptr<TaskImpl>(new TaskImpl{
@@ -82,8 +83,8 @@ RunHandlerEnvironment::Task RunHandlerEnvironment::CreateTask(
 
 void RunHandlerEnvironment::ExecuteTask(const Task& t) {
   WithContext wc(t.f->context);
-  tracing::ScopedRegion region(tracing::EventCategory::kRunClosure,
-                               t.f->trace_id);
+  tsl::tracing::ScopedRegion region(tsl::tracing::EventCategory::kRunClosure,
+                                    t.f->trace_id);
   t.f->f();
 }
 
diff --git a/tensorflow/core/kernels/collective_nccl.cc b/tensorflow/core/kernels/collective_nccl.cc
index 9a1cb900f3bf11..c44680b27124aa 100644
--- a/tensorflow/core/kernels/collective_nccl.cc
+++ b/tensorflow/core/kernels/collective_nccl.cc
@@ -18,7 +18,6 @@ limitations under the License.
 
 #include "tensorflow/core/common_runtime/collective_util.h"
 #include "tensorflow/core/nccl/nccl_manager.h"
-#include "tensorflow/core/platform/tracing.h"
 #include "tensorflow/core/profiler/lib/traceme.h"
 
 namespace tensorflow {
diff --git a/tensorflow/core/kernels/collective_nccl_all_to_all.cc b/tensorflow/core/kernels/collective_nccl_all_to_all.cc
index a24183ee71e40b..a531f19b976cec 100644
--- a/tensorflow/core/kernels/collective_nccl_all_to_all.cc
+++ b/tensorflow/core/kernels/collective_nccl_all_to_all.cc
@@ -18,7 +18,6 @@ limitations under the License.
 
 #include "tensorflow/core/common_runtime/collective_util.h"
 #include "tensorflow/core/nccl/nccl_manager.h"
-#include "tensorflow/core/platform/tracing.h"
 #include "tensorflow/core/profiler/lib/traceme.h"
 
 namespace tensorflow {
diff --git a/tensorflow/core/kernels/collective_nccl_broadcaster.cc b/tensorflow/core/kernels/collective_nccl_broadcaster.cc
index 12b7f5222a9dcf..d3a2ec9218a002 100644
--- a/tensorflow/core/kernels/collective_nccl_broadcaster.cc
+++ b/tensorflow/core/kernels/collective_nccl_broadcaster.cc
@@ -18,7 +18,6 @@ limitations under the License.
 
 #include "tensorflow/core/common_runtime/collective_util.h"
 #include "tensorflow/core/nccl/nccl_manager.h"
-#include "tensorflow/core/platform/tracing.h"
 #include "tensorflow/core/profiler/lib/traceme.h"
 
 namespace tensorflow {
diff --git a/tensorflow/core/kernels/collective_nccl_gatherer.cc b/tensorflow/core/kernels/collective_nccl_gatherer.cc
index a029e5bb6604f2..f6d8e63a748877 100644
--- a/tensorflow/core/kernels/collective_nccl_gatherer.cc
+++ b/tensorflow/core/kernels/collective_nccl_gatherer.cc
@@ -18,7 +18,6 @@ limitations under the License.
 
 #include "tensorflow/core/common_runtime/collective_util.h"
 #include "tensorflow/core/nccl/nccl_manager.h"
-#include "tensorflow/core/platform/tracing.h"
 #include "tensorflow/core/profiler/lib/traceme.h"
 
 namespace tensorflow {
diff --git a/tensorflow/core/kernels/collective_nccl_reducer.cc b/tensorflow/core/kernels/collective_nccl_reducer.cc
index 46fbc2dc12d087..ebc2129f4c4e03 100644
--- a/tensorflow/core/kernels/collective_nccl_reducer.cc
+++ b/tensorflow/core/kernels/collective_nccl_reducer.cc
@@ -18,7 +18,6 @@ limitations under the License.
 
 #include "tensorflow/core/common_runtime/collective_util.h"
 #include "tensorflow/core/nccl/nccl_manager.h"
-#include "tensorflow/core/platform/tracing.h"
 #include "tensorflow/core/profiler/lib/traceme.h"
 
 namespace tensorflow {
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 13173f183a9d8c..d88506b4176a29 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
@@ -38,7 +38,6 @@ limitations under the License.
 #include "tensorflow/core/platform/env_time.h"
 #include "tensorflow/core/platform/status.h"
 #include "tensorflow/core/platform/stringprintf.h"
-#include "tensorflow/core/platform/tracing.h"
 #include "tensorflow/core/profiler/lib/traceme.h"
 #include "tensorflow/core/profiler/lib/traceme_encode.h"
 
diff --git a/tensorflow/core/kernels/function_ops.cc b/tensorflow/core/kernels/function_ops.cc
index 3f4bb1bb96ed4d..7c865b62d0452b 100644
--- a/tensorflow/core/kernels/function_ops.cc
+++ b/tensorflow/core/kernels/function_ops.cc
@@ -31,7 +31,6 @@ limitations under the License.
 #include "tensorflow/core/framework/register_types.h"
 #include "tensorflow/core/graph/algorithm.h"
 #include "tensorflow/core/platform/macros.h"
-#include "tensorflow/core/platform/tracing.h"
 #include "tensorflow/core/profiler/lib/traceme.h"
 #include "tensorflow/core/util/device_name_utils.h"
 
diff --git a/tensorflow/core/nccl/collective_communicator.cc b/tensorflow/core/nccl/collective_communicator.cc
index 1555ed9aac46a0..ac62f471087a21 100644
--- a/tensorflow/core/nccl/collective_communicator.cc
+++ b/tensorflow/core/nccl/collective_communicator.cc
@@ -21,7 +21,6 @@ limitations under the License.
 
 #include "absl/memory/memory.h"
 #include "tensorflow/core/nccl/nccl_manager.h"
-#include "tensorflow/core/platform/tracing.h"
 #include "tensorflow/core/profiler/lib/traceme.h"
 
 namespace tensorflow {
diff --git a/tensorflow/core/tfrt/run_handler_thread_pool/BUILD b/tensorflow/core/tfrt/run_handler_thread_pool/BUILD
index a7c5fb24092b17..a70727213d4e0c 100644
--- a/tensorflow/core/tfrt/run_handler_thread_pool/BUILD
+++ b/tensorflow/core/tfrt/run_handler_thread_pool/BUILD
@@ -57,6 +57,7 @@ cc_library(
         "//tensorflow/core/protobuf:for_core_protos_cc",
         "//tensorflow/core/tfrt/runtime:work_queue_interface",
         "@eigen_archive//:eigen3",
+        "@local_tsl//tsl/platform:env",
         "@tf_runtime//:hostcontext",
     ],
 )
diff --git a/tensorflow/core/tfrt/run_handler_thread_pool/run_handler.cc b/tensorflow/core/tfrt/run_handler_thread_pool/run_handler.cc
index c0147eb6392fb5..aafa0155f14426 100644
--- a/tensorflow/core/tfrt/run_handler_thread_pool/run_handler.cc
+++ b/tensorflow/core/tfrt/run_handler_thread_pool/run_handler.cc
@@ -32,13 +32,13 @@ limitations under the License.
 #include "tensorflow/core/platform/denormal.h"
 #include "tensorflow/core/platform/mutex.h"
 #include "tensorflow/core/platform/setround.h"
-#include "tensorflow/core/platform/tracing.h"
 #include "tensorflow/core/profiler/lib/connected_traceme.h"
 #include "tensorflow/core/profiler/lib/traceme.h"
 #include "tensorflow/core/profiler/lib/traceme_encode.h"
 #include "tensorflow/core/tfrt/run_handler_thread_pool/run_handler.h"
 #include "tensorflow/core/tfrt/run_handler_thread_pool/run_handler_util.h"
 #include "tensorflow/core/tfrt/runtime/work_queue_interface.h"
+#include "tsl/platform/tracing.h"
 #include "tfrt/host_context/async_dispatch.h"  // from @tf_runtime
 
 namespace tfrt {
@@ -72,10 +72,10 @@ RunHandlerEnvironment::EnvThread* RunHandlerEnvironment::CreateThread(
 
 RunHandlerEnvironment::Task RunHandlerEnvironment::CreateTask(TaskFunction f) {
   uint64_t id = 0;
-  if (tensorflow::tracing::EventCollector::IsEnabled()) {
-    id = tensorflow::tracing::GetUniqueArg();
-    tensorflow::tracing::RecordEvent(
-        tensorflow::tracing::EventCategory::kScheduleClosure, id);
+  if (tsl::tracing::EventCollector::IsEnabled()) {
+    id = tsl::tracing::GetUniqueArg();
+    tsl::tracing::RecordEvent(tsl::tracing::EventCategory::kScheduleClosure,
+                              id);
   }
   return Task{
       std::unique_ptr<TaskImpl>(new TaskImpl{
@@ -88,8 +88,8 @@ RunHandlerEnvironment::Task RunHandlerEnvironment::CreateTask(TaskFunction f) {
 
 void RunHandlerEnvironment::ExecuteTask(const Task& t) {
   tensorflow::WithContext wc(t.f->context);
-  tensorflow::tracing::ScopedRegion region(
-      tensorflow::tracing::EventCategory::kRunClosure, t.f->trace_id);
+  tsl::tracing::ScopedRegion region(tsl::tracing::EventCategory::kRunClosure,
+                                    t.f->trace_id);
   t.f->f();
 }
 

From f4b0af12faf167c7967d3925f3c2f20250453470 Mon Sep 17 00:00:00 2001
From: Michael Levesque-Dion <mlevesquedion@google.com>
Date: Mon, 20 May 2024 13:03:35 -0700
Subject: [PATCH 432/478] Use == instead of equals method (deprecated)

PiperOrigin-RevId: 635545562
---
 third_party/stablehlo/temporary.patch         | 55 +++++++++++++++++++
 .../xla/third_party/stablehlo/temporary.patch | 55 +++++++++++++++++++
 2 files changed, 110 insertions(+)

diff --git a/third_party/stablehlo/temporary.patch b/third_party/stablehlo/temporary.patch
index abb329aec579e4..85a49110b59f95 100755
--- a/third_party/stablehlo/temporary.patch
+++ b/third_party/stablehlo/temporary.patch
@@ -200,6 +200,28 @@ diff --ruN a/stablehlo/stablehlo/conversions/tosa/tests/nullary.mlir b/stablehlo
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<10xf64>
    return %0 : tensor<10xf64>
  }
+diff --ruN a/stablehlo/stablehlo/dialect/AssemblyFormat.cpp b/stablehlo/stablehlo/dialect/AssemblyFormat.cpp
+--- stablehlo/stablehlo/dialect/AssemblyFormat.cpp
++++ stablehlo/stablehlo/dialect/AssemblyFormat.cpp
+@@ -305,8 +305,7 @@
+ bool isCommutativeNoRegionMatchingDialect(OperationName innerOp,
+                                           StringRef reduceOpDialect) {
+   auto innerOpDialect = innerOp.getDialect();
+-  return innerOpDialect &&
+-         innerOpDialect->getNamespace().equals(reduceOpDialect) &&
++  return innerOpDialect && innerOpDialect->getNamespace() == reduceOpDialect &&
+          innerOp.hasTrait<mlir::OpTrait::NOperands<2>::Impl>() &&
+          innerOp.hasTrait<mlir::OpTrait::OneResult>() &&
+          (innerOp.hasTrait<mlir::hlo::OpTrait::IsCommutative>() ||
+@@ -359,7 +358,7 @@
+   // Check E5.
+   LLVM_DEBUG(llvm::dbgs() << "Checking ReduceOp compact print E5\n");
+   auto retOp = block.getTerminator();
+-  if (!retOp->getName().stripDialect().equals("return")) return false;
++  if (retOp->getName().stripDialect() != "return") return false;
+ 
+   return llvm::equal(innerOp.getResults(), retOp->getOperands());
+ }
 diff --ruN a/stablehlo/stablehlo/experimental/BUILD.bazel b/stablehlo/stablehlo/experimental/BUILD.bazel
 --- stablehlo/stablehlo/experimental/BUILD.bazel
 +++ stablehlo/stablehlo/experimental/BUILD.bazel
@@ -2589,4 +2611,37 @@ diff --ruN a/stablehlo/stablehlo/experimental/transforms/StablehloRefineShapes.c
 +}  // namespace experimental
 +}  // namespace stablehlo
 +}  // namespace mlir
+diff --ruN a/stablehlo/stablehlo/reference/Api.cpp b/stablehlo/stablehlo/reference/Api.cpp
+--- stablehlo/stablehlo/reference/Api.cpp
++++ stablehlo/stablehlo/reference/Api.cpp
+@@ -51,7 +51,7 @@
+   auto functions = module.getOps<func::FuncOp>();
+ 
+   for (auto funcOp : functions)
+-    if (funcOp.getSymName().equals(mainName)) return funcOp;
++    if (funcOp.getSymName() == mainName) return funcOp;
+ 
+   bool isSingleFunction =
+       std::distance(functions.begin(), functions.end()) == 1;
+@@ -68,7 +68,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/transforms/StablehloRefineShapes.cpp b/stablehlo/stablehlo/transforms/StablehloRefineShapes.cpp
+--- stablehlo/stablehlo/transforms/StablehloRefineShapes.cpp
++++ stablehlo/stablehlo/transforms/StablehloRefineShapes.cpp
+@@ -764,7 +764,7 @@
+ 
+     // Clean up operand buffers after refinement
+     // Must do in this pattern to avoid needing multiple refinement iterations
+-    if (op.getCallTargetName().equals(kCustomCallOperandBarrierTarget)) {
++    if (op.getCallTargetName() == kCustomCallOperandBarrierTarget) {
+       Value operand = op.getOperand(0);
+       if (operand.getType() == op.getResult(0).getType()) {
+         op.replaceAllUsesWith(ValueRange(operand));
 
diff --git a/third_party/xla/third_party/stablehlo/temporary.patch b/third_party/xla/third_party/stablehlo/temporary.patch
index abb329aec579e4..85a49110b59f95 100755
--- a/third_party/xla/third_party/stablehlo/temporary.patch
+++ b/third_party/xla/third_party/stablehlo/temporary.patch
@@ -200,6 +200,28 @@ diff --ruN a/stablehlo/stablehlo/conversions/tosa/tests/nullary.mlir b/stablehlo
    %0 = stablehlo.constant dense<0.000000e+00> : tensor<10xf64>
    return %0 : tensor<10xf64>
  }
+diff --ruN a/stablehlo/stablehlo/dialect/AssemblyFormat.cpp b/stablehlo/stablehlo/dialect/AssemblyFormat.cpp
+--- stablehlo/stablehlo/dialect/AssemblyFormat.cpp
++++ stablehlo/stablehlo/dialect/AssemblyFormat.cpp
+@@ -305,8 +305,7 @@
+ bool isCommutativeNoRegionMatchingDialect(OperationName innerOp,
+                                           StringRef reduceOpDialect) {
+   auto innerOpDialect = innerOp.getDialect();
+-  return innerOpDialect &&
+-         innerOpDialect->getNamespace().equals(reduceOpDialect) &&
++  return innerOpDialect && innerOpDialect->getNamespace() == reduceOpDialect &&
+          innerOp.hasTrait<mlir::OpTrait::NOperands<2>::Impl>() &&
+          innerOp.hasTrait<mlir::OpTrait::OneResult>() &&
+          (innerOp.hasTrait<mlir::hlo::OpTrait::IsCommutative>() ||
+@@ -359,7 +358,7 @@
+   // Check E5.
+   LLVM_DEBUG(llvm::dbgs() << "Checking ReduceOp compact print E5\n");
+   auto retOp = block.getTerminator();
+-  if (!retOp->getName().stripDialect().equals("return")) return false;
++  if (retOp->getName().stripDialect() != "return") return false;
+ 
+   return llvm::equal(innerOp.getResults(), retOp->getOperands());
+ }
 diff --ruN a/stablehlo/stablehlo/experimental/BUILD.bazel b/stablehlo/stablehlo/experimental/BUILD.bazel
 --- stablehlo/stablehlo/experimental/BUILD.bazel
 +++ stablehlo/stablehlo/experimental/BUILD.bazel
@@ -2589,4 +2611,37 @@ diff --ruN a/stablehlo/stablehlo/experimental/transforms/StablehloRefineShapes.c
 +}  // namespace experimental
 +}  // namespace stablehlo
 +}  // namespace mlir
+diff --ruN a/stablehlo/stablehlo/reference/Api.cpp b/stablehlo/stablehlo/reference/Api.cpp
+--- stablehlo/stablehlo/reference/Api.cpp
++++ stablehlo/stablehlo/reference/Api.cpp
+@@ -51,7 +51,7 @@
+   auto functions = module.getOps<func::FuncOp>();
+ 
+   for (auto funcOp : functions)
+-    if (funcOp.getSymName().equals(mainName)) return funcOp;
++    if (funcOp.getSymName() == mainName) return funcOp;
+ 
+   bool isSingleFunction =
+       std::distance(functions.begin(), functions.end()) == 1;
+@@ -68,7 +68,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/transforms/StablehloRefineShapes.cpp b/stablehlo/stablehlo/transforms/StablehloRefineShapes.cpp
+--- stablehlo/stablehlo/transforms/StablehloRefineShapes.cpp
++++ stablehlo/stablehlo/transforms/StablehloRefineShapes.cpp
+@@ -764,7 +764,7 @@
+ 
+     // Clean up operand buffers after refinement
+     // Must do in this pattern to avoid needing multiple refinement iterations
+-    if (op.getCallTargetName().equals(kCustomCallOperandBarrierTarget)) {
++    if (op.getCallTargetName() == kCustomCallOperandBarrierTarget) {
+       Value operand = op.getOperand(0);
+       if (operand.getType() == op.getResult(0).getType()) {
+         op.replaceAllUsesWith(ValueRange(operand));
 

From ba8c22bacab54feab12b93385d6262ff9bc807ca Mon Sep 17 00:00:00 2001
From: Kyle Lucke <klucke@google.com>
Date: Mon, 20 May 2024 13:17:19 -0700
Subject: [PATCH 433/478] Replace the use of xla::OkStatus with absl::OkStatus
 now that they're the same.

PiperOrigin-RevId: 635549245
---
 .../backends/interpreter/executable_base.cc   |   2 +-
 .../profiler/cpu/metadata_collector.cc        |   6 +-
 .../backends/profiler/plugin/plugin_tracer.cc |   6 +-
 .../xla/backends/profiler/tpu/tpu_tracer.cc   |   6 +-
 third_party/xla/xla/ffi/ffi_api.cc            |   2 +-
 .../xla/xla/hlo/evaluator/hlo_evaluator.cc    | 128 +++++++++---------
 .../evaluator/hlo_evaluator_typed_visitor.h   | 104 +++++++-------
 .../xla/xla/hlo/utils/hlo_sharding_util.cc    |   8 +-
 .../xla/pjrt/cpu/abstract_tfrt_cpu_buffer.cc  |   4 +-
 third_party/xla/xla/pjrt/cpu/cpu_client.cc    |   8 +-
 .../pjrt/distributed/client_server_test.cc    |  24 ++--
 third_party/xla/xla/pjrt/distributed/util.h   |   2 +-
 .../xla/xla/pjrt/gpu/se_gpu_pjrt_client.cc    |   4 +-
 .../xla/pjrt/gpu/se_gpu_pjrt_client_test.cc   |  15 +-
 .../xla/service/cpu/tests/cpu_noalias_test.cc |   2 +-
 .../service/heap_simulator/heap_simulator.cc  |   4 +-
 .../stream_executor_interface.h               |   4 +-
 .../xla/tools/hlo_bisect/hlo_bisect_state.cc  |   2 +-
 .../xla/tools/hlo_bisect/hlo_bisect_utils.cc  |   2 +-
 third_party/xla/xla/tools/hlo_opt/opt_main.cc |   2 +-
 .../functional_hlo_runner.cc                  |  10 +-
 .../hlo_to_mhlo/hlo_function_importer.cc      |   4 +-
 .../hlo_to_mhlo/hlo_module_importer.cc        |   2 +-
 .../xla/translate/mhlo_to_hlo/translate.cc    |   2 +-
 .../xla/xla/tsl/util/byte_swap_array.h        |   2 +-
 25 files changed, 178 insertions(+), 177 deletions(-)

diff --git a/third_party/xla/xla/backends/interpreter/executable_base.cc b/third_party/xla/xla/backends/interpreter/executable_base.cc
index 0bb1639daa58a0..a58c811c87e257 100644
--- a/third_party/xla/xla/backends/interpreter/executable_base.cc
+++ b/third_party/xla/xla/backends/interpreter/executable_base.cc
@@ -171,7 +171,7 @@ InterpreterExecutableBase::AllocateOutputMemoryWithInputReuse(
                 alias->ToString());
           }
         }
-        return OkStatus();
+        return absl::OkStatus();
       }));
 
   se::StreamExecutor* executor = stream->parent();
diff --git a/third_party/xla/xla/backends/profiler/cpu/metadata_collector.cc b/third_party/xla/xla/backends/profiler/cpu/metadata_collector.cc
index 0fc767d5d69e7f..ebb73c2ca5fc7b 100644
--- a/third_party/xla/xla/backends/profiler/cpu/metadata_collector.cc
+++ b/third_party/xla/xla/backends/profiler/cpu/metadata_collector.cc
@@ -47,7 +47,7 @@ class MetadataCollector : public tsl::profiler::ProfilerInterface {
       xla::XlaDebugInfoManager::Get()->StartTracing();
       trace_active_ = true;
     }
-    return OkStatus();
+    return absl::OkStatus();
   }
 
   absl::Status Stop() override {
@@ -55,7 +55,7 @@ class MetadataCollector : public tsl::profiler::ProfilerInterface {
       xla::XlaDebugInfoManager::Get()->StopTracing(&debug_info_);
       trace_active_ = false;
     }
-    return OkStatus();
+    return absl::OkStatus();
   }
 
   absl::Status CollectData(tsl::profiler::XSpace* space) override {
@@ -70,7 +70,7 @@ class MetadataCollector : public tsl::profiler::ProfilerInterface {
       }
       debug_info_.clear();
     }
-    return OkStatus();
+    return absl::OkStatus();
   }
 
  private:
diff --git a/third_party/xla/xla/backends/profiler/plugin/plugin_tracer.cc b/third_party/xla/xla/backends/profiler/plugin/plugin_tracer.cc
index 5ce73a669a785f..57fb7980cf06ed 100644
--- a/third_party/xla/xla/backends/profiler/plugin/plugin_tracer.cc
+++ b/third_party/xla/xla/backends/profiler/plugin/plugin_tracer.cc
@@ -146,7 +146,7 @@ absl::Status PluginTracer::Start() {
   args.profiler = profiler_;
   RETURN_STATUS_IF_PLUGIN_PROFILER_ERROR(profiler_api_->start(&args),
                                          profiler_api_);
-  return OkStatus();
+  return absl::OkStatus();
 }
 
 absl::Status PluginTracer::Stop() {
@@ -154,7 +154,7 @@ absl::Status PluginTracer::Stop() {
   args.profiler = profiler_;
   RETURN_STATUS_IF_PLUGIN_PROFILER_ERROR(profiler_api_->stop(&args),
                                          profiler_api_);
-  return OkStatus();
+  return absl::OkStatus();
 }
 
 absl::Status PluginTracer::CollectData(XSpace* space) {
@@ -172,7 +172,7 @@ absl::Status PluginTracer::CollectData(XSpace* space) {
       plane->Swap(&tpu_plane);
     }
   }
-  return OkStatus();
+  return absl::OkStatus();
 }
 
 }  // namespace profiler
diff --git a/third_party/xla/xla/backends/profiler/tpu/tpu_tracer.cc b/third_party/xla/xla/backends/profiler/tpu/tpu_tracer.cc
index 031cf651bf408d..7488645cf40f2a 100644
--- a/third_party/xla/xla/backends/profiler/tpu/tpu_tracer.cc
+++ b/third_party/xla/xla/backends/profiler/tpu/tpu_tracer.cc
@@ -126,7 +126,7 @@ absl::Status TpuTracer::Start() {
     LOG(ERROR) << "TPU tracer failed to start.";
     return status.status();
   }
-  return OkStatus();
+  return absl::OkStatus();
 }
 
 absl::Status TpuTracer::Stop() {
@@ -137,7 +137,7 @@ absl::Status TpuTracer::Stop() {
     LOG(ERROR) << "TPU tracer failed to stop.";
     return status.status();
   }
-  return OkStatus();
+  return absl::OkStatus();
 }
 
 absl::Status TpuTracer::CollectData(XSpace* space) {
@@ -164,7 +164,7 @@ absl::Status TpuTracer::CollectData(XSpace* space) {
     LOG(ERROR) << "TPU tracer failed to collect data.";
     return status.status();
   }
-  return OkStatus();
+  return absl::OkStatus();
 }
 
 // Initializes TpuProfilerApiFns. The initialization may not be successful if
diff --git a/third_party/xla/xla/ffi/ffi_api.cc b/third_party/xla/xla/ffi/ffi_api.cc
index 0bc59e99f51248..02e6b69e4f9064 100644
--- a/third_party/xla/xla/ffi/ffi_api.cc
+++ b/third_party/xla/xla/ffi/ffi_api.cc
@@ -138,7 +138,7 @@ static absl::Status RegisterHandler(std::string_view name,
     return absl::InvalidArgumentError(
         absl::StrCat("Duplicate FFI handler registration for ", name,
                      " on a platform ", platform));
-  return OkStatus();
+  return absl::OkStatus();
 }
 
 absl::StatusOr<HandlerRegistration> FindHandler(std::string_view name,
diff --git a/third_party/xla/xla/hlo/evaluator/hlo_evaluator.cc b/third_party/xla/xla/hlo/evaluator/hlo_evaluator.cc
index 7e0a9b1627e5a4..9df1e652d24dd5 100644
--- a/third_party/xla/xla/hlo/evaluator/hlo_evaluator.cc
+++ b/third_party/xla/xla/hlo/evaluator/hlo_evaluator.cc
@@ -1121,7 +1121,7 @@ absl::Status HloEvaluator::EvaluateParameterFromCallerArgument(
   TF_RETURN_IF_ERROR(evaluated_[parameter].CopyFrom(
       caller_operand_literal, /*dest_shape_index=*/shape_index,
       /*src_shape_index=*/shape_index));
-  return OkStatus();
+  return absl::OkStatus();
 }
 
 std::vector<int64_t> HloEvaluator::GetS64Indices(
@@ -1163,7 +1163,7 @@ absl::Status HloEvaluator::EvaluateInternal(
   // Don't need to evaluate this instruction again if it has already been
   // evaluated.
   if (IsAlreadyEvaluated(instruction, shape_index)) {
-    return OkStatus();
+    return absl::OkStatus();
   }
 
   if (!recursively_evaluate_nonconstant_operands) {
@@ -1228,7 +1228,7 @@ absl::Status HloEvaluator::EvaluateInternal(
           evaluated_[instruction] =
               Literal::CreateFromShapeWithUnknownLeafArrays(
                   instruction->shape());
-          return OkStatus();
+          return absl::OkStatus();
         }
       }
     }
@@ -1237,7 +1237,7 @@ absl::Status HloEvaluator::EvaluateInternal(
   TF_RETURN_IF_ERROR(Preprocess(instruction));
   TF_RETURN_IF_ERROR(instruction->Visit(this));
   TF_RETURN_IF_ERROR(Postprocess(instruction));
-  return OkStatus();
+  return absl::OkStatus();
 }
 
 absl::Status HloEvaluator::HandleBitcast(const HloInstruction* bitcast) {
@@ -1251,7 +1251,7 @@ absl::Status HloEvaluator::HandleBitcast(const HloInstruction* bitcast) {
   memcpy(result.untyped_data(), operand_literal.untyped_data(),
          result.size_bytes());
   evaluated_[bitcast] = std::move(result);
-  return OkStatus();
+  return absl::OkStatus();
 }
 
 absl::Status HloEvaluator::HandleBitcastConvert(const HloInstruction* convert) {
@@ -1261,7 +1261,7 @@ absl::Status HloEvaluator::HandleBitcastConvert(const HloInstruction* convert) {
       GetEvaluatedLiteralFor(operand).BitcastConvert(convert->shape()));
 
   evaluated_[convert] = std::move(result);
-  return OkStatus();
+  return absl::OkStatus();
 }
 
 absl::Status HloEvaluator::HandleGetDimensionSize(
@@ -1278,7 +1278,7 @@ absl::Status HloEvaluator::HandleGetDimensionSize(
   if (dynamic_size != nullptr) {
     evaluated_[get_dimension_size] =
         GetEvaluatedLiteralFor(dynamic_size).Clone();
-    return OkStatus();
+    return absl::OkStatus();
   }
 
   const Shape& shape = get_dimension_size->operand(0)->shape();
@@ -1286,7 +1286,7 @@ absl::Status HloEvaluator::HandleGetDimensionSize(
   output.PopulateWithValue(
       static_cast<int32_t>(shape.dimensions(get_dimension_size->dimension())));
   evaluated_[get_dimension_size] = std::move(output);
-  return OkStatus();
+  return absl::OkStatus();
 }
 
 absl::Status HloEvaluator::HandleSetDimensionSize(
@@ -1301,7 +1301,7 @@ absl::Status HloEvaluator::HandleSetDimensionSize(
   result.SetDynamicSize(set_dimension_size->dimension(),
                         size_literal.Get<int32_t>({}));
   evaluated_[set_dimension_size] = std::move(result);
-  return OkStatus();
+  return absl::OkStatus();
 }
 
 absl::Status HloEvaluator::HandleParameter(const HloInstruction* parameter) {
@@ -1313,7 +1313,7 @@ absl::Status HloEvaluator::HandleParameter(const HloInstruction* parameter) {
     }
     evaluated_[parameter] =
         Literal::CreateFromShapeWithUnknownLeafArrays(parameter->shape());
-    return OkStatus();
+    return absl::OkStatus();
   }
 
   if (!arg_literals_.empty()) {
@@ -1332,7 +1332,7 @@ absl::Status HloEvaluator::HandleParameter(const HloInstruction* parameter) {
 #endif
   }
 
-  return OkStatus();
+  return absl::OkStatus();
 }
 
 absl::Status HloEvaluator::HandleInfeed(const HloInstruction* infeed) {
@@ -1343,24 +1343,24 @@ absl::Status HloEvaluator::HandleInfeed(const HloInstruction* infeed) {
   }
   evaluated_[infeed] =
       Literal::CreateFromShapeWithUnknownLeafArrays(infeed->shape());
-  return OkStatus();
+  return absl::OkStatus();
 }
 
 absl::Status HloEvaluator::HandleConstant(const HloInstruction*) {
-  return OkStatus();
+  return absl::OkStatus();
 }
 
 absl::Status HloEvaluator::HandleReshape(const HloInstruction* reshape) {
   TF_ASSIGN_OR_RETURN(evaluated_[reshape],
                       GetEvaluatedLiteralFor(reshape->operand(0))
                           .Reshape(reshape->shape().dimensions()));
-  return OkStatus();
+  return absl::OkStatus();
 }
 
 absl::Status HloEvaluator::HandleTranspose(const HloInstruction* transpose) {
   evaluated_[transpose] = GetEvaluatedLiteralFor(transpose->operand(0))
                               .Transpose(transpose->dimensions());
-  return OkStatus();
+  return absl::OkStatus();
 }
 
 absl::Status HloEvaluator::HandleConcatenate(
@@ -1402,7 +1402,7 @@ absl::Status HloEvaluator::HandleConcatenate(
   }
 
   evaluated_[concatenate] = std::move(result_literal);
-  return OkStatus();
+  return absl::OkStatus();
 }
 
 absl::Status HloEvaluator::HandleIsFinite(const HloInstruction* is_finite) {
@@ -1420,7 +1420,7 @@ absl::Status HloEvaluator::HandleIsFinite(const HloInstruction* is_finite) {
               },
               GetEvaluatedLiteralFor(operand));
           TF_ASSIGN_OR_RETURN(evaluated_[is_finite], std::move(result_or));
-          return OkStatus();
+          return absl::OkStatus();
         }
         return InvalidArgument(
             "expected element type in shape to be floating point, but got: %s",
@@ -1440,7 +1440,7 @@ absl::Status HloEvaluator::HandleReal(const HloInstruction* real) {
               real, [](NativeT elem_operand) { return elem_operand; },
               GetEvaluatedLiteralFor(operand));
           TF_ASSIGN_OR_RETURN(evaluated_[real], std::move(result_or));
-          return OkStatus();
+          return absl::OkStatus();
         }
         if constexpr (primitive_util::IsComplexType(primitive_type_constant)) {
           using NativeT = primitive_util::NativeTypeOf<primitive_type_constant>;
@@ -1450,7 +1450,7 @@ absl::Status HloEvaluator::HandleReal(const HloInstruction* real) {
                   [](NativeT elem_operand) { return std::real(elem_operand); },
                   GetEvaluatedLiteralFor(operand));
           TF_ASSIGN_OR_RETURN(evaluated_[real], std::move(result_or));
-          return OkStatus();
+          return absl::OkStatus();
         }
         LOG(FATAL) << "HandleReal: unknown/unhandled primitive type: "
                    << PrimitiveType_Name(operand->shape().element_type());
@@ -1469,7 +1469,7 @@ absl::Status HloEvaluator::HandleImag(const HloInstruction* imag) {
               imag, [](NativeT elem_operand) { return NativeT(0); },
               GetEvaluatedLiteralFor(operand));
           TF_ASSIGN_OR_RETURN(evaluated_[imag], std::move(result_or));
-          return OkStatus();
+          return absl::OkStatus();
         }
         if constexpr (primitive_util::IsComplexType(primitive_type_constant)) {
           using NativeT = primitive_util::NativeTypeOf<primitive_type_constant>;
@@ -1479,7 +1479,7 @@ absl::Status HloEvaluator::HandleImag(const HloInstruction* imag) {
                   [](NativeT elem_operand) { return std::imag(elem_operand); },
                   GetEvaluatedLiteralFor(operand));
           TF_ASSIGN_OR_RETURN(evaluated_[imag], std::move(result_or));
-          return OkStatus();
+          return absl::OkStatus();
         }
         LOG(FATAL) << "HandleImag: unknown/unhandled primitive type: "
                    << PrimitiveType_Name(operand->shape().element_type());
@@ -1504,7 +1504,7 @@ absl::Status HloEvaluator::HandleComplex(const HloInstruction* complex) {
                     imag.Get<typename NativeT::value_type>(multi_index));
               }));
           evaluated_[complex] = std::move(result);
-          return OkStatus();
+          return absl::OkStatus();
         }
         LOG(FATAL) << "HandleComplex: unknown/unhandled primitive type: "
                    << PrimitiveType_Name(complex->shape().element_type());
@@ -1534,7 +1534,7 @@ absl::Status HloEvaluator::HandleCompare(const HloInstruction* compare) {
           TF_ASSIGN_OR_RETURN(evaluated_[compare],
                               Compare<NativeT>(compare->shape(), comparison,
                                                lhs_literal, rhs_literal));
-          return OkStatus();
+          return absl::OkStatus();
         }
         LOG(FATAL) << "HandleCompare: unknown primitive type: "
                    << PrimitiveType_Name(element_type);
@@ -1592,7 +1592,7 @@ absl::Status HloEvaluator::HandleTuple(const HloInstruction* tuple) {
   } else {
     evaluated_[tuple] = std::move(new_result);
   }
-  return OkStatus();
+  return absl::OkStatus();
 }
 
 namespace {
@@ -1739,7 +1739,7 @@ class FftTransform {
                       input_strides, input_shape.rank(), 0, 0, base_case);
     }
 
-    return OkStatus();
+    return absl::OkStatus();
   }
 
  private:
@@ -2237,7 +2237,7 @@ class FftTransform {
       }
     }
 
-    return OkStatus();
+    return absl::OkStatus();
   }
 
  private:
@@ -2256,7 +2256,7 @@ absl::Status HloEvaluator::HandleFft(const HloInstruction* fft) {
   TF_RETURN_IF_ERROR(transform.ComputeFft(fft, input_literal, &output_literal));
   evaluated_[fft] = std::move(output_literal);
 
-  return OkStatus();
+  return absl::OkStatus();
 }
 
 // Returns an ShapeUtil::IndexIterationSpace that iterates over the output batch
@@ -2392,7 +2392,7 @@ class OutputBatchIndexToInputIndex {
       TF_RET_CHECK(start_index.has_value());
       index_vector_[i] = *start_index;
     }
-    return OkStatus();
+    return absl::OkStatus();
   }
 
   // Populates input_index_.
@@ -2574,7 +2574,7 @@ absl::Status HloEvaluator::HandleGather(const HloInstruction* gather) {
   const Shape& operand_shape = operand.shape();
   if (ShapeUtil::IsZeroElementArray(operand_shape)) {
     evaluated_[gather] = std::move(result);
-    return OkStatus();
+    return absl::OkStatus();
   }
 
   auto gather_inner_loop_body =
@@ -2629,7 +2629,7 @@ absl::Status HloEvaluator::HandleGather(const HloInstruction* gather) {
   TF_RETURN_IF_ERROR(ShapeUtil::ForEachIndexWithStatus(
       shape, start_indices_iteration_space, gather_outer_loop_body));
   evaluated_[gather] = std::move(result);
-  return OkStatus();
+  return absl::OkStatus();
 }
 
 namespace {
@@ -2796,7 +2796,7 @@ class UpdateScatterIndexToInputIndex {
       index_vector_[i] =
           *scatter_indices_.GetIntegralAsS64(index_vector_index_);
     }
-    return OkStatus();
+    return absl::OkStatus();
   }
 
   // Populates input_index_.
@@ -3045,7 +3045,7 @@ absl::Status HloEvaluator::HandleScatter(const HloInstruction* hlo) {
       updates[0]->shape(), scatter_indices_iteration_space,
       scatter_outer_loop_body));
   evaluated_[scatter] = std::move(result);
-  return OkStatus();
+  return absl::OkStatus();
 }
 
 absl::Status HloEvaluator::HandleBroadcast(const HloInstruction* broadcast) {
@@ -3073,12 +3073,12 @@ absl::Status HloEvaluator::HandleBroadcast(const HloInstruction* broadcast) {
       evaluated_[broadcast],
       operand.Broadcast(broadcast->shape(), broadcast->dimensions()));
 
-  return OkStatus();
+  return absl::OkStatus();
 }
 
 absl::Status HloEvaluator::HandleAfterAll(const HloInstruction* after_all) {
   evaluated_[after_all] = LiteralUtil::CreateToken();
-  return OkStatus();
+  return absl::OkStatus();
 }
 
 absl::Status HloEvaluator::HandleAddDependency(
@@ -3086,7 +3086,7 @@ absl::Status HloEvaluator::HandleAddDependency(
   // AddDedendency just forwards its zero-th operand.
   evaluated_[add_dependency] =
       GetEvaluatedLiteralFor(add_dependency->operand(0)).Clone();
-  return OkStatus();
+  return absl::OkStatus();
 }
 
 absl::Status HloEvaluator::HandleGetTupleElement(
@@ -3115,7 +3115,7 @@ absl::Status HloEvaluator::HandleGetTupleElement(
 absl::Status HloEvaluator::HandleCopy(const HloInstruction* copy) {
   TF_RET_CHECK(ShapeUtil::Compatible(copy->shape(), copy->operand(0)->shape()));
   evaluated_[copy] = GetEvaluatedLiteralFor(copy->operand(0)).Clone();
-  return OkStatus();
+  return absl::OkStatus();
 }
 
 absl::Status HloEvaluator::HandleAsyncStart(const HloInstruction* async_start) {
@@ -3146,7 +3146,7 @@ absl::Status HloEvaluator::HandleAsyncStart(const HloInstruction* async_start) {
   TF_RETURN_IF_ERROR(evaluated_[async_start].MoveFrom(
       std::move(result), /*dest_shape_index=*/{1}));
 
-  return OkStatus();
+  return absl::OkStatus();
 }
 
 absl::Status HloEvaluator::HandleAsyncUpdate(
@@ -3157,7 +3157,7 @@ absl::Status HloEvaluator::HandleAsyncUpdate(
   TF_RETURN_IF_ERROR(evaluated_[async_update].CopyFrom(operand_tuple_literal,
                                                        /*dest_shape_index=*/{},
                                                        /*src_shape_index=*/{}));
-  return OkStatus();
+  return absl::OkStatus();
 }
 
 absl::Status HloEvaluator::HandleAsyncDone(const HloInstruction* async_done) {
@@ -3167,7 +3167,7 @@ absl::Status HloEvaluator::HandleAsyncDone(const HloInstruction* async_done) {
   TF_RETURN_IF_ERROR(evaluated_[async_done].CopyFrom(operand_tuple_literal,
                                                      /*dest_shape_index=*/{},
                                                      /*src_shape_index=*/{1}));
-  return OkStatus();
+  return absl::OkStatus();
 }
 
 absl::Status HloEvaluator::HandleCopyStart(const HloInstruction* copy_start) {
@@ -3187,7 +3187,7 @@ absl::Status HloEvaluator::HandleCopyStart(const HloInstruction* copy_start) {
   evaluated_[copy_start] = LiteralUtil::MakeTuple(
       {&GetEvaluatedLiteralFor(copy_start->operand(0)),
        &GetEvaluatedLiteralFor(copy_start->operand(0)), &context_literal});
-  return OkStatus();
+  return absl::OkStatus();
 }
 
 absl::Status HloEvaluator::HandleCopyDone(const HloInstruction* copy_done) {
@@ -3204,7 +3204,7 @@ absl::Status HloEvaluator::HandleCopyDone(const HloInstruction* copy_done) {
   TF_RETURN_IF_ERROR(evaluated_[copy_done].CopyFrom(operand_tuple_literal,
                                                     /*dest_shape_index=*/{},
                                                     /*src_shape_index=*/{0}));
-  return OkStatus();
+  return absl::OkStatus();
 }
 
 absl::Status HloEvaluator::HandleCall(const HloInstruction* call) {
@@ -3226,7 +3226,7 @@ absl::Status HloEvaluator::HandleCall(const HloInstruction* call) {
                       embedded_evaluator->Evaluate(*computation, arg_literals));
 
   evaluated_[call] = std::move(result);
-  return OkStatus();
+  return absl::OkStatus();
 }
 
 absl::Status HloEvaluator::HandleFusion(const HloInstruction* fusion) {
@@ -3264,7 +3264,7 @@ absl::Status HloEvaluator::HandleFusion(const HloInstruction* fusion) {
                                           *readded_computation, arg_literals));
 
   evaluated_[fusion] = std::move(result);
-  return OkStatus();
+  return absl::OkStatus();
 }
 
 absl::Status HloEvaluator::HandleConditional(
@@ -3293,7 +3293,7 @@ absl::Status HloEvaluator::HandleConditional(
                           {&branch_computation_arg}));
 
   evaluated_[conditional] = std::move(result);
-  return OkStatus();
+  return absl::OkStatus();
 }
 
 absl::Status HloEvaluator::HandleConvert(const HloInstruction* convert) {
@@ -3302,7 +3302,7 @@ absl::Status HloEvaluator::HandleConvert(const HloInstruction* convert) {
   TF_ASSIGN_OR_RETURN(Literal result, GetEvaluatedLiteralFor(operand).Convert(
                                           convert->shape().element_type()));
   evaluated_[convert] = std::move(result);
-  return OkStatus();
+  return absl::OkStatus();
 }
 
 absl::Status HloEvaluator::HandleDynamicSlice(
@@ -3355,7 +3355,7 @@ absl::Status HloEvaluator::HandleDynamicSlice(
   };
   TF_RETURN_IF_ERROR(result.PopulateInplace(func));
   evaluated_[dynamic_slice] = std::move(result);
-  return OkStatus();
+  return absl::OkStatus();
 }
 
 absl::Status HloEvaluator::HandleDynamicUpdateSlice(const HloInstruction* dus) {
@@ -3407,7 +3407,7 @@ absl::Status HloEvaluator::HandleDynamicUpdateSlice(const HloInstruction* dus) {
                                   func);
   evaluated_[dus] = std::move(result);
 
-  return OkStatus();
+  return absl::OkStatus();
 }
 
 absl::Status HloEvaluator::HandleSelect(const HloInstruction* select) {
@@ -3422,7 +3422,7 @@ absl::Status HloEvaluator::HandleSelect(const HloInstruction* select) {
     } else {
       evaluated_[select] = on_false.Clone();
     }
-    return OkStatus();
+    return absl::OkStatus();
   }
 
   return DefaultAction(select);
@@ -3497,7 +3497,7 @@ absl::Status HloEvaluator::HandleWhile(const HloInstruction* while_hlo) {
         visitor_shape_index_.size() != 1 ||
         parsed_while_loop->static_while_loop->induction_var_index !=
             visitor_shape_index_[0]) {
-      return OkStatus();
+      return absl::OkStatus();
     }
     Shape induction_var_shape =
         ShapeUtil::GetSubshape(while_hlo->shape(), visitor_shape_index_);
@@ -3508,7 +3508,7 @@ absl::Status HloEvaluator::HandleWhile(const HloInstruction* while_hlo) {
     TF_RETURN_IF_ERROR(evaluated_[while_hlo].CopyFrom(
         induction_var_val, /*dest_shape_index=*/visitor_shape_index_,
         /*src_shape_index=*/{}));
-    return OkStatus();
+    return absl::OkStatus();
   }
   bool keep_going = true;
   int64_t iteration_count = 0;
@@ -3544,7 +3544,7 @@ absl::Status HloEvaluator::HandleWhile(const HloInstruction* while_hlo) {
     }
   }
   evaluated_[while_hlo] = std::move(lcv);
-  return OkStatus();
+  return absl::OkStatus();
 }
 
 namespace {
@@ -3776,7 +3776,7 @@ absl::Status HloEvaluator::HandleReverse(const HloInstruction* reverse) {
       }));
 
   evaluated_[reverse] = std::move(result);
-  return OkStatus();
+  return absl::OkStatus();
 }
 
 absl::Status HloEvaluator::HandleSelectAndScatter(
@@ -3870,7 +3870,7 @@ absl::Status HloEvaluator::HandleSelectAndScatter(
       IndexUtil::BumpIndices(source->shape(), absl::MakeSpan(source_index)));
 
   evaluated_[select_and_scatter] = std::move(result);
-  return OkStatus();
+  return absl::OkStatus();
 }
 
 absl::Status HloEvaluator::HandleSlice(const HloInstruction* slice) {
@@ -3905,7 +3905,7 @@ absl::Status HloEvaluator::HandleSlice(const HloInstruction* slice) {
   Literal result(shape);
   TF_RETURN_IF_ERROR(result.PopulateInplaceParallel(func));
   evaluated_[slice] = std::move(result);
-  return OkStatus();
+  return absl::OkStatus();
 }
 
 absl::Status HloEvaluator::HandleSort(const HloInstruction* sort) {
@@ -4012,7 +4012,7 @@ absl::Status HloEvaluator::HandleSort(const HloInstruction* sort) {
     absl::c_copy(lhs, std::back_inserter(tmp));
     absl::c_copy(rhs, std::back_inserter(tmp));
     absl::c_copy(tmp, output.begin());
-    return OkStatus();
+    return absl::OkStatus();
   };
   auto* env = tsl::Env::Default();
   const int max_parallelism = tsl::port::MaxParallelism();
@@ -4031,7 +4031,7 @@ absl::Status HloEvaluator::HandleSort(const HloInstruction* sort) {
                       HloEvaluator* embedded_evaluator) -> absl::Status {
     // Base case: inputs with 0 or 1 elements are already sorted.
     if (to_sort.size() < 2) {
-      return OkStatus();
+      return absl::OkStatus();
     }
     size_t halfway = to_sort.size() / 2;
     auto lhs = to_sort.subspan(/*pos=*/0, halfway);
@@ -4109,7 +4109,7 @@ absl::Status HloEvaluator::HandleSort(const HloInstruction* sort) {
         std::rotate(ub, i, i + 1);
       }
     }
-    return OkStatus();
+    return absl::OkStatus();
   };
 
   // Iterate through each dimension except 'sort_dim'.
@@ -4163,7 +4163,7 @@ absl::Status HloEvaluator::HandleSort(const HloInstruction* sort) {
 
     evaluated_[sort] = std::move(result_tuple);
   }
-  return OkStatus();
+  return absl::OkStatus();
 }
 
 absl::Status HloEvaluator::HandleStochasticConvert(
@@ -4179,7 +4179,7 @@ absl::Status HloEvaluator::HandleStochasticConvert(
   TF_ASSIGN_OR_RETURN(
       evaluated_[stochastic_convert],
       StochasticConvertOp(operand_literal, random_literal, result_shape));
-  return OkStatus();
+  return absl::OkStatus();
 }
 
 static bool IsScalarAdd(HloComputation* computation) {
@@ -4419,7 +4419,7 @@ absl::Status HloEvaluator::HandleReduce(const HloInstruction* hlo) {
     TF_ASSIGN_OR_RETURN(evaluated_[reduce],
                         evaluated_[reduce].ConvertToShape(reduce->shape()));
   }
-  return OkStatus();
+  return absl::OkStatus();
 }
 
 absl::Status HloEvaluator::HandleReduceWindow(const HloInstruction* hlo) {
@@ -4560,7 +4560,7 @@ absl::Status HloEvaluator::HandleReduceWindow(const HloInstruction* hlo) {
   }
   VLOG(2) << "Final result is:" << result.ToString() << "\n";
   evaluated_[reduce_window] = std::move(result);
-  return OkStatus();
+  return absl::OkStatus();
 }
 
 absl::Status HloEvaluator::HandleMap(const HloInstruction* map) {
@@ -4592,7 +4592,7 @@ absl::Status HloEvaluator::HandleMap(const HloInstruction* map) {
         return computed_result;
       }));
   evaluated_[map] = std::move(result);
-  return OkStatus();
+  return absl::OkStatus();
 }
 
 absl::Status HloEvaluator::HandleCustomCall(const HloInstruction* custom_call) {
@@ -4613,7 +4613,7 @@ absl::Status HloEvaluator::HandleCustomCall(const HloInstruction* custom_call) {
       auto output, custom_call_handler_(custom_call, absl::MakeSpan(operands)));
 
   evaluated_[custom_call] = std::move(output);
-  return OkStatus();
+  return absl::OkStatus();
 }
 
 absl::Status HloEvaluator::Preprocess(const HloInstruction* hlo) {
@@ -4647,7 +4647,7 @@ absl::Status HloEvaluator::Postprocess(const HloInstruction* hlo) {
                                           hlo_shape.layout())) {
     evaluated_.at(hlo) = evaluated_.at(hlo).Relayout(hlo_shape);
   }
-  return OkStatus();
+  return absl::OkStatus();
 }
 
 namespace {
diff --git a/third_party/xla/xla/hlo/evaluator/hlo_evaluator_typed_visitor.h b/third_party/xla/xla/hlo/evaluator/hlo_evaluator_typed_visitor.h
index 9d11bbc5d0f906..8bc1e2c889e785 100644
--- a/third_party/xla/xla/hlo/evaluator/hlo_evaluator_typed_visitor.h
+++ b/third_party/xla/xla/hlo/evaluator/hlo_evaluator_typed_visitor.h
@@ -167,7 +167,7 @@ class HloEvaluatorTypedVisitor : public ConstDfsHloVisitorWithDefault {
                         ElementWiseUnaryOp(abs, [](NativeT elem_operand) {
                           return elem_operand;
                         }));
-    return OkStatus();
+    return absl::OkStatus();
   }
 
   template <typename NativeT,
@@ -177,7 +177,7 @@ class HloEvaluatorTypedVisitor : public ConstDfsHloVisitorWithDefault {
                         ElementWiseUnaryOp(abs, [](NativeT elem_operand) {
                           return std::abs(elem_operand);
                         }));
-    return OkStatus();
+    return absl::OkStatus();
   }
 
   template <typename NativeT,
@@ -192,7 +192,7 @@ class HloEvaluatorTypedVisitor : public ConstDfsHloVisitorWithDefault {
             abs, [](NativeT elem_operand) { return std::abs(elem_operand); },
             operand_literal)));
 
-    return OkStatus();
+    return absl::OkStatus();
   }
 
   absl::Status HandleAbs(const HloInstruction* abs) override {
@@ -214,7 +214,7 @@ class HloEvaluatorTypedVisitor : public ConstDfsHloVisitorWithDefault {
           ElementWiseUnaryOp(round, [](ElementwiseT elem_operand) {
             return std::round(elem_operand);
           }));
-      return OkStatus();
+      return absl::OkStatus();
     }
     return UnsupportedTypeError(round);
   }
@@ -228,7 +228,7 @@ class HloEvaluatorTypedVisitor : public ConstDfsHloVisitorWithDefault {
           ElementWiseUnaryOp(round, [](ElementwiseT elem_operand) {
             return std::nearbyint(elem_operand);
           }));
-      return OkStatus();
+      return absl::OkStatus();
     }
     return UnsupportedTypeError(round);
   }
@@ -240,7 +240,7 @@ class HloEvaluatorTypedVisitor : public ConstDfsHloVisitorWithDefault {
           ElementWiseUnaryOp(ceil, [](ElementwiseT elem_operand) {
             return std::ceil(elem_operand);
           }));
-      return OkStatus();
+      return absl::OkStatus();
     }
     return UnsupportedTypeError(ceil);
   }
@@ -252,7 +252,7 @@ class HloEvaluatorTypedVisitor : public ConstDfsHloVisitorWithDefault {
           ElementWiseUnaryOp(erf, [](ElementwiseT elem_operand) {
             return std::erf(elem_operand);
           }));
-      return OkStatus();
+      return absl::OkStatus();
     }
     return UnsupportedTypeError(erf);
   }
@@ -262,7 +262,7 @@ class HloEvaluatorTypedVisitor : public ConstDfsHloVisitorWithDefault {
                         ElementWiseUnaryOp(exp, [](ElementwiseT elem_operand) {
                           return std::exp(elem_operand);
                         }));
-    return OkStatus();
+    return absl::OkStatus();
   }
 
   absl::Status HandleExpm1(const HloInstruction* expm1) override {
@@ -272,7 +272,7 @@ class HloEvaluatorTypedVisitor : public ConstDfsHloVisitorWithDefault {
           ElementWiseUnaryOp(expm1, [](ElementwiseT elem_operand) {
             return std::expm1(elem_operand);
           }));
-      return OkStatus();
+      return absl::OkStatus();
     }
     return UnsupportedTypeError(expm1);
   }
@@ -284,7 +284,7 @@ class HloEvaluatorTypedVisitor : public ConstDfsHloVisitorWithDefault {
           ElementWiseUnaryOp(floor, [](ElementwiseT elem_operand) {
             return std::floor(elem_operand);
           }));
-      return OkStatus();
+      return absl::OkStatus();
     }
     return UnsupportedTypeError(floor);
   }
@@ -294,7 +294,7 @@ class HloEvaluatorTypedVisitor : public ConstDfsHloVisitorWithDefault {
                         ElementWiseUnaryOp(log, [](ElementwiseT elem_operand) {
                           return std::log(elem_operand);
                         }));
-    return OkStatus();
+    return absl::OkStatus();
   }
 
   absl::Status HandleLog1p(const HloInstruction* log1p) override {
@@ -304,7 +304,7 @@ class HloEvaluatorTypedVisitor : public ConstDfsHloVisitorWithDefault {
           ElementWiseUnaryOp(log1p, [](ElementwiseT elem_operand) {
             return std::log1p(elem_operand);
           }));
-      return OkStatus();
+      return absl::OkStatus();
     }
     return UnsupportedTypeError(log1p);
   }
@@ -322,7 +322,7 @@ class HloEvaluatorTypedVisitor : public ConstDfsHloVisitorWithDefault {
               return ~elem_operand;
             }
           }));
-      return OkStatus();
+      return absl::OkStatus();
     }
     return UnsupportedTypeError(not_);
   }
@@ -338,7 +338,7 @@ class HloEvaluatorTypedVisitor : public ConstDfsHloVisitorWithDefault {
         ElementWiseUnaryOp(negate, [](ElementwiseT elem_operand) {
           return NativeT(-type(elem_operand));
         }));
-    return OkStatus();
+    return absl::OkStatus();
   }
 
   template <typename NativeT, typename std::enable_if_t<
@@ -349,7 +349,7 @@ class HloEvaluatorTypedVisitor : public ConstDfsHloVisitorWithDefault {
         parent_->evaluated_[negate],
         ElementWiseUnaryOp(
             negate, [](ElementwiseT elem_operand) { return -elem_operand; }));
-    return OkStatus();
+    return absl::OkStatus();
   }
 
   absl::Status HandleNegate(const HloInstruction* negate) override {
@@ -363,7 +363,7 @@ class HloEvaluatorTypedVisitor : public ConstDfsHloVisitorWithDefault {
           return static_cast<ElementwiseT>(1) /
                  (static_cast<ElementwiseT>(1) + std::exp(-elem_operand));
         }));
-    return OkStatus();
+    return absl::OkStatus();
   }
 
   absl::Status HandleSign(const HloInstruction* sign) override {
@@ -386,7 +386,7 @@ class HloEvaluatorTypedVisitor : public ConstDfsHloVisitorWithDefault {
             return 0 == abs_val ? ElementwiseT(0) : elem_operand / abs_val;
           }
         }));
-    return OkStatus();
+    return absl::OkStatus();
   }
 
   absl::Status HandleAtan2(const HloInstruction* atan2) override {
@@ -396,7 +396,7 @@ class HloEvaluatorTypedVisitor : public ConstDfsHloVisitorWithDefault {
                                                         ElementwiseT rhs_elem) {
                             return std::atan2(lhs_elem, rhs_elem);
                           }));
-      return OkStatus();
+      return absl::OkStatus();
     }
     if constexpr (is_complex_v<ElementwiseT>) {
       TF_ASSIGN_OR_RETURN(
@@ -406,7 +406,7 @@ class HloEvaluatorTypedVisitor : public ConstDfsHloVisitorWithDefault {
             auto i = ElementwiseT(0.0, 1.0);
             return (-i) * (std::log((x + i * y) / std::sqrt(x * x + y * y)));
           }));
-      return OkStatus();
+      return absl::OkStatus();
     }
     return UnsupportedTypeError(atan2);
   }
@@ -416,7 +416,7 @@ class HloEvaluatorTypedVisitor : public ConstDfsHloVisitorWithDefault {
                         ElementWiseUnaryOp(tanh, [](ElementwiseT elem_operand) {
                           return std::tanh(elem_operand);
                         }));
-    return OkStatus();
+    return absl::OkStatus();
   }
 
   absl::Status HandleMultiply(const HloInstruction* multiply) override {
@@ -427,7 +427,7 @@ class HloEvaluatorTypedVisitor : public ConstDfsHloVisitorWithDefault {
               return ElementwiseT(ToArithmeticSafeType(lhs_elem) *
                                   ToArithmeticSafeType(rhs_elem));
             }));
-    return OkStatus();
+    return absl::OkStatus();
   }
 
   absl::Status HandleSubtract(const HloInstruction* subtract) override {
@@ -438,7 +438,7 @@ class HloEvaluatorTypedVisitor : public ConstDfsHloVisitorWithDefault {
               return ElementwiseT(ToArithmeticSafeType(lhs_elem) -
                                   ToArithmeticSafeType(rhs_elem));
             }));
-    return OkStatus();
+    return absl::OkStatus();
   }
 
   absl::Status HandleAdd(const HloInstruction* add) override {
@@ -448,7 +448,7 @@ class HloEvaluatorTypedVisitor : public ConstDfsHloVisitorWithDefault {
                           return ElementwiseT(ToArithmeticSafeType(lhs_elem) +
                                               ToArithmeticSafeType(rhs_elem));
                         }));
-    return OkStatus();
+    return absl::OkStatus();
   }
 
   absl::Status HandleDivide(const HloInstruction* divide) override {
@@ -475,7 +475,7 @@ class HloEvaluatorTypedVisitor : public ConstDfsHloVisitorWithDefault {
               }
               return lhs_elem / rhs_elem;
             }));
-    return OkStatus();
+    return absl::OkStatus();
   }
 
   absl::Status HandleMaximum(const HloInstruction* maximum) override {
@@ -493,7 +493,7 @@ class HloEvaluatorTypedVisitor : public ConstDfsHloVisitorWithDefault {
             }
             return std::max(lhs, rhs);
           }));
-      return OkStatus();
+      return absl::OkStatus();
     }
     return UnsupportedTypeError(maximum);
   }
@@ -513,7 +513,7 @@ class HloEvaluatorTypedVisitor : public ConstDfsHloVisitorWithDefault {
             }
             return std::min(lhs, rhs);
           }));
-      return OkStatus();
+      return absl::OkStatus();
     }
     return UnsupportedTypeError(minimum);
   }
@@ -566,7 +566,7 @@ class HloEvaluatorTypedVisitor : public ConstDfsHloVisitorWithDefault {
             return static_cast<ElementwiseT>(std::pow(lhs_el, rhs_el));
           }
         }));
-    return OkStatus();
+    return absl::OkStatus();
   }
 
   absl::Status HandleSqrt(const HloInstruction* sqrt) override {
@@ -574,7 +574,7 @@ class HloEvaluatorTypedVisitor : public ConstDfsHloVisitorWithDefault {
                         ElementWiseUnaryOp(sqrt, [](ElementwiseT elem_operand) {
                           return std::sqrt(elem_operand);
                         }));
-    return OkStatus();
+    return absl::OkStatus();
   }
 
   absl::Status HandleCbrt(const HloInstruction* cbrt) override {
@@ -584,7 +584,7 @@ class HloEvaluatorTypedVisitor : public ConstDfsHloVisitorWithDefault {
           ElementWiseUnaryOp(cbrt, [](ElementwiseT elem_operand) {
             return std::cbrt(elem_operand);
           }));
-      return OkStatus();
+      return absl::OkStatus();
     }
     return UnsupportedTypeError(cbrt);
   }
@@ -595,7 +595,7 @@ class HloEvaluatorTypedVisitor : public ConstDfsHloVisitorWithDefault {
         ElementWiseUnaryOp(rsqrt, [](ElementwiseT elem_operand) {
           return static_cast<ElementwiseT>(1) / std::sqrt(elem_operand);
         }));
-    return OkStatus();
+    return absl::OkStatus();
   }
 
   absl::Status HandleRemainder(const HloInstruction* remainder) override {
@@ -621,7 +621,7 @@ class HloEvaluatorTypedVisitor : public ConstDfsHloVisitorWithDefault {
                   return std::fmod(lhs_el, rhs_el);
                 }
               }));
-      return OkStatus();
+      return absl::OkStatus();
     }
     return UnsupportedTypeError(remainder);
   }
@@ -634,7 +634,7 @@ class HloEvaluatorTypedVisitor : public ConstDfsHloVisitorWithDefault {
                               [](ElementwiseT lhs_el, ElementwiseT rhs_el) {
                                 return lhs_el & rhs_el;
                               }));
-      return OkStatus();
+      return absl::OkStatus();
     }
     return UnsupportedTypeError(and_inst);
   }
@@ -646,7 +646,7 @@ class HloEvaluatorTypedVisitor : public ConstDfsHloVisitorWithDefault {
                                                           ElementwiseT rhs_el) {
                             return lhs_el | rhs_el;
                           }));
-      return OkStatus();
+      return absl::OkStatus();
     }
     return UnsupportedTypeError(or_inst);
   }
@@ -659,7 +659,7 @@ class HloEvaluatorTypedVisitor : public ConstDfsHloVisitorWithDefault {
                               [](ElementwiseT lhs_el, ElementwiseT rhs_el) {
                                 return lhs_el ^ rhs_el;
                               }));
-      return OkStatus();
+      return absl::OkStatus();
     }
     return UnsupportedTypeError(xor_inst);
   }
@@ -674,7 +674,7 @@ class HloEvaluatorTypedVisitor : public ConstDfsHloVisitorWithDefault {
                                        ? 0
                                        : (lhs_elem << rhs_elem);
                           }));
-      return OkStatus();
+      return absl::OkStatus();
     }
     return UnsupportedTypeError(shl);
   }
@@ -694,7 +694,7 @@ class HloEvaluatorTypedVisitor : public ConstDfsHloVisitorWithDefault {
                   return static_cast<ElementwiseT>(lhs_signed >> rhs_elem);
                 }
               }));
-      return OkStatus();
+      return absl::OkStatus();
     }
     return UnsupportedTypeError(shr);
   }
@@ -714,7 +714,7 @@ class HloEvaluatorTypedVisitor : public ConstDfsHloVisitorWithDefault {
                             return static_cast<ElementwiseT>(
                                 static_cast<UnsignedT>(lhs_elem) >> rhs_elem);
                           }));
-      return OkStatus();
+      return absl::OkStatus();
     }
     return UnsupportedTypeError(shr);
   }
@@ -740,7 +740,7 @@ class HloEvaluatorTypedVisitor : public ConstDfsHloVisitorWithDefault {
           parent_->evaluated_[clamp],
           ElementwiseTernaryOp(clamp,
                                std::move(ConvertTernaryFunction(clamp_op))));
-      return OkStatus();
+      return absl::OkStatus();
     }
     return UnsupportedTypeError(clamp);
   }
@@ -757,7 +757,7 @@ class HloEvaluatorTypedVisitor : public ConstDfsHloVisitorWithDefault {
         };
     TF_ASSIGN_OR_RETURN(parent_->evaluated_[select],
                         ElementwiseTernaryOp(select, std::move(select_op)));
-    return OkStatus();
+    return absl::OkStatus();
   }
 
   absl::Status HandleConvolutionWithLiterals(const HloInstruction* conv,
@@ -958,7 +958,7 @@ class HloEvaluatorTypedVisitor : public ConstDfsHloVisitorWithDefault {
     TF_RETURN_IF_ERROR(result.PopulateParallel<ReturnT>(func));
 
     parent_->evaluated_[conv] = std::move(result);
-    return OkStatus();
+    return absl::OkStatus();
   }
 
   absl::Status HandleConvolution(const HloInstruction* conv) override {
@@ -1092,7 +1092,7 @@ class HloEvaluatorTypedVisitor : public ConstDfsHloVisitorWithDefault {
     result.PopulateR2FromArray2D(*result_array);
     parent_->evaluated_[dot] =
         std::move(result).Convert(dot->shape().element_type()).value();
-    return OkStatus();
+    return absl::OkStatus();
   }
 
   template <typename NativeT, typename std::enable_if_t<
@@ -1209,7 +1209,7 @@ class HloEvaluatorTypedVisitor : public ConstDfsHloVisitorWithDefault {
         }));
 
     parent_->evaluated_[dot] = std::move(result);
-    return OkStatus();
+    return absl::OkStatus();
   }
 
   absl::Status HandleDotSlowPath(const HloInstruction* dot) {
@@ -1307,7 +1307,7 @@ class HloEvaluatorTypedVisitor : public ConstDfsHloVisitorWithDefault {
                                     step, func);
 
     parent_->evaluated_[pad] = std::move(result);
-    return OkStatus();
+    return absl::OkStatus();
   }
 
   absl::Status HandleClz(const HloInstruction* clz) override {
@@ -1321,7 +1321,7 @@ class HloEvaluatorTypedVisitor : public ConstDfsHloVisitorWithDefault {
                                       std::numeric_limits<ReturnT>::is_signed;
             return (unsigned_digits - 1) - Log2Floor<uint64_t>(elem_operand);
           }));
-      return OkStatus();
+      return absl::OkStatus();
     }
     return UnsupportedTypeError(clz);
   }
@@ -1335,7 +1335,7 @@ class HloEvaluatorTypedVisitor : public ConstDfsHloVisitorWithDefault {
             return std::bitset<CHAR_BIT * sizeof(ReturnT)>(elem_operand)
                 .count();
           }));
-      return OkStatus();
+      return absl::OkStatus();
     }
     return UnsupportedTypeError(popcnt);
   }
@@ -1348,7 +1348,7 @@ class HloEvaluatorTypedVisitor : public ConstDfsHloVisitorWithDefault {
           ElementWiseUnaryOp(sin, [](ElementwiseT elem_operand) {
             return std::sin(elem_operand);
           }));
-      return OkStatus();
+      return absl::OkStatus();
     }
     return UnsupportedTypeError(sin);
   }
@@ -1361,7 +1361,7 @@ class HloEvaluatorTypedVisitor : public ConstDfsHloVisitorWithDefault {
           ElementWiseUnaryOp(cos, [](ElementwiseT elem_operand) {
             return std::cos(elem_operand);
           }));
-      return OkStatus();
+      return absl::OkStatus();
     }
     return UnsupportedTypeError(cos);
   }
@@ -1373,7 +1373,7 @@ class HloEvaluatorTypedVisitor : public ConstDfsHloVisitorWithDefault {
           ElementWiseUnaryOp(tan, [](ElementwiseT elem_operand) {
             return std::tan(elem_operand);
           }));
-      return OkStatus();
+      return absl::OkStatus();
     }
     return UnsupportedTypeError(tan);
   }
@@ -1474,7 +1474,7 @@ class HloEvaluatorTypedVisitor : public ConstDfsHloVisitorWithDefault {
           }
           return reduced_result;
         }));
-    return OkStatus();
+    return absl::OkStatus();
   }
 
   template <typename NativeT,
@@ -1501,7 +1501,7 @@ class HloEvaluatorTypedVisitor : public ConstDfsHloVisitorWithDefault {
             return true;
           });
       parent_->evaluated_[iota] = std::move(result);
-      return OkStatus();
+      return absl::OkStatus();
     }
     return UnsupportedTypeError(iota);
   }
@@ -1566,7 +1566,7 @@ class HloEvaluatorTypedVisitor : public ConstDfsHloVisitorWithDefault {
                                      " is not implemented.");
       }
       parent_->evaluated_[random] = std::move(result);
-      return OkStatus();
+      return absl::OkStatus();
     }
     if constexpr (std::is_integral_v<ElementwiseT>) {
       switch (distribution) {
@@ -1599,7 +1599,7 @@ class HloEvaluatorTypedVisitor : public ConstDfsHloVisitorWithDefault {
                                      " is not implemented.");
       }
       parent_->evaluated_[random] = std::move(result);
-      return OkStatus();
+      return absl::OkStatus();
     }
     return UnsupportedTypeError(random);
   }
diff --git a/third_party/xla/xla/hlo/utils/hlo_sharding_util.cc b/third_party/xla/xla/hlo/utils/hlo_sharding_util.cc
index 40cd1ccb77e603..8e28d90f5c9e21 100644
--- a/third_party/xla/xla/hlo/utils/hlo_sharding_util.cc
+++ b/third_party/xla/xla/hlo/utils/hlo_sharding_util.cc
@@ -538,7 +538,7 @@ bool MergeShardingIfCompatible(const HloSharding& to_merge,
               *device = *it1;
               gm1.erase(it1);
               gm2.erase(it2);
-              return OkStatus();
+              return absl::OkStatus();
             } else if (*it1 < *it2) {
               it1++;
             } else {
@@ -3160,11 +3160,11 @@ absl::Status CanonicalizeLayoutAfterShardingPropagation(
     HloModule* module, bool update_output_layout,
     bool update_parameters_layout) {
   if (!update_output_layout && !update_parameters_layout) {
-    return OkStatus();
+    return absl::OkStatus();
   }
   if (!module->layout_canonicalization_callback()) {
     LOG(INFO) << "There is no registered layout_canonicalization_callback.";
-    return OkStatus();
+    return absl::OkStatus();
   }
   TF_ASSIGN_OR_RETURN(auto shapes_with_layout,
                       module->layout_canonicalization_callback()(*module));
@@ -3190,7 +3190,7 @@ absl::Status CanonicalizeLayoutAfterShardingPropagation(
     }
   }
 
-  return OkStatus();
+  return absl::OkStatus();
 }
 
 }  // namespace hlo_sharding_util
diff --git a/third_party/xla/xla/pjrt/cpu/abstract_tfrt_cpu_buffer.cc b/third_party/xla/xla/pjrt/cpu/abstract_tfrt_cpu_buffer.cc
index d6f851f8c6febd..f0247de19f532b 100644
--- a/third_party/xla/xla/pjrt/cpu/abstract_tfrt_cpu_buffer.cc
+++ b/third_party/xla/xla/pjrt/cpu/abstract_tfrt_cpu_buffer.cc
@@ -946,7 +946,7 @@ AbstractAsyncHostToHostMemoryTransferManager::TransferRawDataToSubBuffer(
       event->SetStateConcrete();
     }
   });
-  return OkStatus();
+  return absl::OkStatus();
 }
 
 void AbstractAsyncHostToHostMemoryTransferManager::SetBufferError(
@@ -981,7 +981,7 @@ AbstractAsyncHostToHostMemoryTransferManager::PopulateAsyncTransferManagerData(
     buffer_sizes.push_back(buffer_size);
   }
 
-  return OkStatus();
+  return absl::OkStatus();
 }
 
 }  // namespace xla
diff --git a/third_party/xla/xla/pjrt/cpu/cpu_client.cc b/third_party/xla/xla/pjrt/cpu/cpu_client.cc
index a0972f795f3e03..73ea144c7e35c9 100644
--- a/third_party/xla/xla/pjrt/cpu/cpu_client.cc
+++ b/third_party/xla/xla/pjrt/cpu/cpu_client.cc
@@ -1176,7 +1176,7 @@ absl::Status TfrtCpuExecutable::SetUpDonation(bool tuple_inputs) {
   TF_ASSIGN_OR_RETURN(parameters_that_must_be_donated_,
                       ComputeParametersThatMustBeDonated(
                           *cpu_executable_->shared_module(), tuple_inputs));
-  return OkStatus();
+  return absl::OkStatus();
 }
 
 namespace {
@@ -1330,7 +1330,7 @@ absl::Status TfrtCpuExecutable::CheckBufferCompatibilities(
           i, input_buffer_sizes_in_bytes_[i], buffer->BufferSizes()[0]);
     }
   }
-  return OkStatus();
+  return absl::OkStatus();
 }
 
 absl::StatusOr<PjRtLoadedExecutable::Result> TfrtCpuExecutable::ExecuteHelper(
@@ -1437,7 +1437,7 @@ absl::StatusOr<PjRtLoadedExecutable::Result> TfrtCpuExecutable::ExecuteHelper(
           tracked_buffer = donation_transaction->device_buffer();
           tracked_buffers.emplace_back(/*can_donate=*/true, tracked_buffer);
           donation_transactions.push_back(std::move(*donation_transaction));
-          return OkStatus();
+          return absl::OkStatus();
         }
       }
       tracked_buffer = tfrt_buffer->AcquireUsage(execute_event);
@@ -1445,7 +1445,7 @@ absl::StatusOr<PjRtLoadedExecutable::Result> TfrtCpuExecutable::ExecuteHelper(
         return InvalidArgument(
             "Invalid buffer passed: buffer has been deleted or donated.");
       tracked_buffers.emplace_back(/*can_donate=*/false, tracked_buffer);
-      return OkStatus();
+      return absl::OkStatus();
     };
     TF_RETURN_IF_ERROR(get_buffer(i));
 
diff --git a/third_party/xla/xla/pjrt/distributed/client_server_test.cc b/third_party/xla/xla/pjrt/distributed/client_server_test.cc
index f6a4b1fae9828c..ef0f29181400f2 100644
--- a/third_party/xla/xla/pjrt/distributed/client_server_test.cc
+++ b/third_party/xla/xla/pjrt/distributed/client_server_test.cc
@@ -154,7 +154,7 @@ TEST_F(ClientServerTest, ConnectAndShutdownAreBarriers) {
       TF_RET_CHECK(shutdown_count == num_nodes);
     }
 
-    return OkStatus();
+    return absl::OkStatus();
   };
 
   std::vector<absl::Status> statuses(num_nodes);
@@ -234,7 +234,7 @@ TEST_F(ClientServerTest, ConnectAndEnumerateDevices) {
         std::string value,
         client->BlockingKeyValueGet("key2", absl::InfiniteDuration()));
     TF_RET_CHECK(value == "value2");
-    return OkStatus();
+    return absl::OkStatus();
   };
   auto thread1_fn = [&]() -> absl::Status {
     auto client = GetClient(/*node_id=*/1);
@@ -260,7 +260,7 @@ TEST_F(ClientServerTest, ConnectAndEnumerateDevices) {
         client->BlockingKeyValueGet("key1", absl::InfiniteDuration()));
     TF_RET_CHECK(value == "value1");
     TF_RETURN_IF_ERROR(client->KeyValueSet("key2", "value2"));
-    return OkStatus();
+    return absl::OkStatus();
   };
 
   std::vector<std::function<absl::Status()>> functions = {thread0_fn,
@@ -314,7 +314,7 @@ TEST_F(ClientServerTest, EnumerateElevenDevices) {
     TF_RET_CHECK(
         xla::protobuf_util::ProtobufEquals(topology, expected_topology))
         << topology.DebugString();
-    return OkStatus();
+    return absl::OkStatus();
   };
 
   std::vector<absl::Status> statuses(num_nodes);
@@ -350,7 +350,7 @@ TEST_F(ClientServerTest, ZeroInitTimeoutShouldStillWaitForOtherTasks) {
     }
     TF_RETURN_IF_ERROR(client->Connect());
 
-    return OkStatus();
+    return absl::OkStatus();
   };
 
   std::vector<absl::Status> statuses(num_nodes);
@@ -380,13 +380,13 @@ TEST_F(ClientServerTest, ClientsTerminateShutdownIfAnyClientGoesAway) {
     TF_RETURN_IF_ERROR(client->Connect());
 
     if (node_id == 0) {
-      return OkStatus();
+      return absl::OkStatus();
     }
 
     // The call to Shutdown() should be interrupted if a worker stops issuing
     // heartbeats.
     TF_RETURN_IF_ERROR(client->Shutdown());
-    return OkStatus();
+    return absl::OkStatus();
   };
 
   std::vector<absl::Status> statuses(num_nodes);
@@ -428,10 +428,10 @@ TEST_F(ClientServerTest, ClientsReceiveMissedHeartbeatIfAnyClientGoesAway) {
     TF_RETURN_IF_ERROR(client->Connect());
 
     if (node_id == 0) {
-      return OkStatus();
+      return absl::OkStatus();
     }
     shutdown.WaitForNotification();
-    return OkStatus();
+    return absl::OkStatus();
   };
 
   std::vector<absl::Status> statuses(num_nodes);
@@ -479,7 +479,7 @@ TEST_F(ClientServerTest, ClientsTerminateIfServiceGoesAway) {
     shutdown.WaitForNotification();
 
     TF_RETURN_IF_ERROR(client->Shutdown());
-    return OkStatus();
+    return absl::OkStatus();
   };
 
   std::vector<absl::Status> statuses(num_nodes);
@@ -514,7 +514,7 @@ TEST_F(ClientServerTest, LateClientsAreOk) {
     absl::SleepFor(absl::Milliseconds(200) * node_id);
     TF_RETURN_IF_ERROR(client->Connect());
     TF_RETURN_IF_ERROR(client->Shutdown());
-    return OkStatus();
+    return absl::OkStatus();
   };
 
   std::vector<absl::Status> statuses(num_nodes);
@@ -552,7 +552,7 @@ TEST_F(ClientServerTest, ConnectEventuallyTimesOutIfAClientDoesNotShowUp) {
 
     TF_RETURN_IF_ERROR(client->Connect());
     TF_RETURN_IF_ERROR(client->Shutdown());
-    return OkStatus();
+    return absl::OkStatus();
   };
 
   // Note: one fewer thread than 'num_nodes'.
diff --git a/third_party/xla/xla/pjrt/distributed/util.h b/third_party/xla/xla/pjrt/distributed/util.h
index be07f318ab7246..b0fc713b191138 100644
--- a/third_party/xla/xla/pjrt/distributed/util.h
+++ b/third_party/xla/xla/pjrt/distributed/util.h
@@ -23,7 +23,7 @@ namespace xla {
 
 inline absl::Status FromGrpcStatus(const ::grpc::Status& s) {
   if (s.ok()) {
-    return OkStatus();
+    return absl::OkStatus();
   } else {
     return absl::Status(static_cast<absl::StatusCode>(s.error_code()),
                         s.error_message());
diff --git a/third_party/xla/xla/pjrt/gpu/se_gpu_pjrt_client.cc b/third_party/xla/xla/pjrt/gpu/se_gpu_pjrt_client.cc
index 92682b59739146..dc0768aa1473d6 100644
--- a/third_party/xla/xla/pjrt/gpu/se_gpu_pjrt_client.cc
+++ b/third_party/xla/xla/pjrt/gpu/se_gpu_pjrt_client.cc
@@ -302,7 +302,7 @@ class AsyncHostToDeviceTransferManager
           (*ptr)();
           delete ptr;
         }));
-    return OkStatus();
+    return absl::OkStatus();
   }
 
   absl::Status TransferRawDataToBuffer(
@@ -1027,7 +1027,7 @@ absl::Status BuildDistributedDevices(
         });
   }
 #endif  // GOOGLE_CUDA
-  return OkStatus();
+  return absl::OkStatus();
 }
 
 std::string MakeComputeCapabilityString(const se::DeviceDescription* desc) {
diff --git a/third_party/xla/xla/pjrt/gpu/se_gpu_pjrt_client_test.cc b/third_party/xla/xla/pjrt/gpu/se_gpu_pjrt_client_test.cc
index 42a5170dcc2541..c178cf38a2930d 100644
--- a/third_party/xla/xla/pjrt/gpu/se_gpu_pjrt_client_test.cc
+++ b/third_party/xla/xla/pjrt/gpu/se_gpu_pjrt_client_test.cc
@@ -191,7 +191,7 @@ TEST(StreamExecutorGpuClientTest, SendRecvChunked) {
         float* data = reinterpret_cast<float*>(chunk.data());
         sent_value[0] = data[0];
         sent_value[1] = data[1];
-        return OkStatus();
+        return absl::OkStatus();
       }};
 
   // Recv buffer from host.
@@ -206,7 +206,7 @@ TEST(StreamExecutorGpuClientTest, SendRecvChunked) {
         *reinterpret_cast<float*>(chunk1.data()) = 6.0f;
         TF_CHECK_OK(stream->AddChunk(std::move(chunk1)).Await());
 
-        return OkStatus();
+        return absl::OkStatus();
       }};
 
   // Callbacks for point-to-point communication ops.
@@ -243,9 +243,10 @@ TEST(StreamExecutorGpuClientTest, SendErrorNoDeadLock) {
 
   // No-op Recv handler.
   RecvCallback recv_callback = {
-      /*channel_id=*/2,
-      [&](const PjRtTransferMetadata& m,
-          std::unique_ptr<CopyToDeviceStream> stream) { return OkStatus(); }};
+      /*channel_id=*/2, [&](const PjRtTransferMetadata& m,
+                            std::unique_ptr<CopyToDeviceStream> stream) {
+        return absl::OkStatus();
+      }};
 
   // Callbacks for point-to-point communication ops.
   std::vector<std::vector<SendCallback>> send_callbacks = {{send_callback}};
@@ -271,7 +272,7 @@ TEST(StreamExecutorGpuClientTest, RecvErrorNoDeadLock) {
   // No-op Send handler.
   SendCallback send_callback = {
       /*channel_id=*/1, [&](const PjRtTransferMetadata&, PjRtChunk, int64_t,
-                            bool) { return OkStatus(); }};
+                            bool) { return absl::OkStatus(); }};
 
   // Invalid Recv handler that tries to add invalid chunk.
   RecvCallback recv_callback = {
@@ -280,7 +281,7 @@ TEST(StreamExecutorGpuClientTest, RecvErrorNoDeadLock) {
         auto chunk = PjRtChunk::AllocateDefault(10 * sizeof(float));
         stream->AddChunk(std::move(chunk)).Await().IgnoreError();
         // Return ok status to proceed to corresponding recv-done call.
-        return OkStatus();
+        return absl::OkStatus();
       }};
 
   // Callbacks for point-to-point communication ops.
diff --git a/third_party/xla/xla/service/cpu/tests/cpu_noalias_test.cc b/third_party/xla/xla/service/cpu/tests/cpu_noalias_test.cc
index ca64a1263c3e58..8426238fbeb288 100644
--- a/third_party/xla/xla/service/cpu/tests/cpu_noalias_test.cc
+++ b/third_party/xla/xla/service/cpu/tests/cpu_noalias_test.cc
@@ -80,7 +80,7 @@ TEST_F(CpuNoAliasTest, Concat) {
       std::make_unique<DependencyHloOrdering>(hlo_module.get()),
       backend().compiler()->BufferSizeBytesFunction(),
       [](LogicalBuffer::Color) { return /*alignment=*/1; });
-  ASSERT_EQ(status_or_buffer_assn.status(), OkStatus());
+  ASSERT_EQ(status_or_buffer_assn.status(), absl::OkStatus());
 
   llvm::LLVMContext context;
   llvm_ir::AliasAnalysis aa(*hlo_module, *status_or_buffer_assn.value(),
diff --git a/third_party/xla/xla/service/heap_simulator/heap_simulator.cc b/third_party/xla/xla/service/heap_simulator/heap_simulator.cc
index 737b9e9351b894..7ca0490f6db776 100644
--- a/third_party/xla/xla/service/heap_simulator/heap_simulator.cc
+++ b/third_party/xla/xla/service/heap_simulator/heap_simulator.cc
@@ -383,7 +383,7 @@ absl::Status HeapSimulator::RunComputation(
       Free(value, value->instruction());
     }
   }
-  return OkStatus();
+  return absl::OkStatus();
 }
 
 HeapSimulator::HeapSimulator(
@@ -2038,7 +2038,7 @@ absl::Status GlobalDecreasingSizeBestFitHeap<BufferType>::
                           "have caught such a condition earlier.");
   }
 
-  return OkStatus();
+  return absl::OkStatus();
 }
 
 // Future opportunities:
diff --git a/third_party/xla/xla/stream_executor/stream_executor_interface.h b/third_party/xla/xla/stream_executor/stream_executor_interface.h
index 95b23ef1231594..aa5f04ced0af43 100644
--- a/third_party/xla/xla/stream_executor/stream_executor_interface.h
+++ b/third_party/xla/xla/stream_executor/stream_executor_interface.h
@@ -104,8 +104,8 @@ class StreamExecutorInterface {
   // Loads a module for the platform this StreamExecutor is acting upon.
   //
   // `spec` describes the module to be loaded.  On success writes the handle for
-  // the loaded module to `module_handle` and returns OkStatus().  Otherwise,
-  // returns the error which has occurred.
+  // the loaded module to `module_handle` and returns absl::OkStatus().
+  // Otherwise, returns the error which has occurred.
   virtual absl::Status LoadModule(const MultiModuleLoaderSpec& spec,
                                   ModuleHandle* module_handle) {
     return absl::UnimplementedError("Not Implemented");
diff --git a/third_party/xla/xla/tools/hlo_bisect/hlo_bisect_state.cc b/third_party/xla/xla/tools/hlo_bisect/hlo_bisect_state.cc
index a8dfde80578563..aeabb3193bd439 100644
--- a/third_party/xla/xla/tools/hlo_bisect/hlo_bisect_state.cc
+++ b/third_party/xla/xla/tools/hlo_bisect/hlo_bisect_state.cc
@@ -325,7 +325,7 @@ absl::Status HloBisectState::ExpectModuleIsBuggy() {
   // Verify that the current module has a bug.
   TF_ASSIGN_OR_RETURN(bool has_bug, RunModule(*module_));
   if (has_bug) {
-    return OkStatus();
+    return absl::OkStatus();
   }
 
   // Check for the bug checker stability.
diff --git a/third_party/xla/xla/tools/hlo_bisect/hlo_bisect_utils.cc b/third_party/xla/xla/tools/hlo_bisect/hlo_bisect_utils.cc
index c86f4cc0435286..0e381514275a7b 100644
--- a/third_party/xla/xla/tools/hlo_bisect/hlo_bisect_utils.cc
+++ b/third_party/xla/xla/tools/hlo_bisect/hlo_bisect_utils.cc
@@ -156,7 +156,7 @@ absl::Status DumpHloModule(HloModule* module, const std::string& file_name,
     LOG(FATAL) << "Unexpected output format: " << output_format;
   }
 
-  return OkStatus();
+  return absl::OkStatus();
 }
 
 }  // namespace
diff --git a/third_party/xla/xla/tools/hlo_opt/opt_main.cc b/third_party/xla/xla/tools/hlo_opt/opt_main.cc
index 57ff4480f3a085..c7c1904dceb66a 100644
--- a/third_party/xla/xla/tools/hlo_opt/opt_main.cc
+++ b/third_party/xla/xla/tools/hlo_opt/opt_main.cc
@@ -176,7 +176,7 @@ absl::Status RunOpt(int argc, char** argv, const HloOptConfig& opts) {
     TF_RETURN_IF_ERROR(
         tsl::WriteStringToFile(tsl::Env::Default(), opts.output_file, output));
   }
-  return OkStatus();
+  return absl::OkStatus();
 }
 
 }  // namespace
diff --git a/third_party/xla/xla/tools/multihost_hlo_runner/functional_hlo_runner.cc b/third_party/xla/xla/tools/multihost_hlo_runner/functional_hlo_runner.cc
index a44fc62f300217..e8a9966fd2abab 100644
--- a/third_party/xla/xla/tools/multihost_hlo_runner/functional_hlo_runner.cc
+++ b/third_party/xla/xla/tools/multihost_hlo_runner/functional_hlo_runner.cc
@@ -451,7 +451,7 @@ absl::Status FunctionalHloRunner::DumpOutput(
       }
     }
   }
-  return OkStatus();
+  return absl::OkStatus();
 }
 
 absl::Span<PjRtDevice* const> FunctionalHloRunner::GetLocalDevices(
@@ -504,7 +504,7 @@ absl::Status FunctionalHloRunner::LoadAndRunAndDump(
                                       compile_options, running_options,
                                       hlo_files, input_format));
   return dump_output_to.empty()
-             ? OkStatus()
+             ? absl::OkStatus()
              : FunctionalHloRunner::DumpOutput(output, dump_output_to, task_id);
 }
 
@@ -573,7 +573,7 @@ absl::Status FunctionalHloRunner::LoadAndCompile(
                          debug_options, preproc_options, compile_options)
                          .status());
 
-  return OkStatus();
+  return absl::OkStatus();
 }
 
 absl::StatusOr<std::unique_ptr<HloModule>>
@@ -749,7 +749,7 @@ absl::Status FunctionalHloRunner::PrepareHloModuleForCompilation(
             /*remove_comm=*/false, /*remove_host_transfer=*/true});
     TF_RETURN_IF_ERROR(pipeline.Run(hlo_module).status());
   }
-  return OkStatus();
+  return absl::OkStatus();
 }
 
 CompileOptions FunctionalHloRunner::CompleteCompileOptions(
@@ -915,7 +915,7 @@ absl::Status EnsureSingleTupleForFlattening(const HloModule& module) {
             ->shape()
             .ToString());
   }
-  return OkStatus();
+  return absl::OkStatus();
 }
 
 }  // namespace
diff --git a/third_party/xla/xla/translate/hlo_to_mhlo/hlo_function_importer.cc b/third_party/xla/xla/translate/hlo_to_mhlo/hlo_function_importer.cc
index 8407deef789292..a2e74df3d6052b 100644
--- a/third_party/xla/xla/translate/hlo_to_mhlo/hlo_function_importer.cc
+++ b/third_party/xla/xla/translate/hlo_to_mhlo/hlo_function_importer.cc
@@ -1408,7 +1408,7 @@ absl::StatusOr<mlir::Operation*> HloFunctionImporter::ImportInstructionImpl(
       }
       return ImportOldStyleAsyncStart<mlir::mhlo::SendOp>(
           attributes, operands, loc, async_bundled_tuple, func_builder, "send_",
-          [](auto) { return OkStatus(); });
+          [](auto) { return absl::OkStatus(); });
     }
     case HloOpcode::kSendDone: {
       return ImportOldStyleAsyncDone(attributes, operands, loc, result_type,
@@ -1442,7 +1442,7 @@ absl::StatusOr<mlir::Operation*> HloFunctionImporter::ImportInstructionImpl(
       }
       return ImportOldStyleAsyncStart<mlir::mhlo::RecvOp>(
           attributes, operands, loc, async_bundled_tuple, func_builder, "recv_",
-          [](auto) { return OkStatus(); });
+          [](auto) { return absl::OkStatus(); });
     }
     case HloOpcode::kRecvDone: {
       return ImportOldStyleAsyncDone(attributes, operands, loc, result_type,
diff --git a/third_party/xla/xla/translate/hlo_to_mhlo/hlo_module_importer.cc b/third_party/xla/xla/translate/hlo_to_mhlo/hlo_module_importer.cc
index c0fb1ffbcf5c9d..299c7b3333500b 100644
--- a/third_party/xla/xla/translate/hlo_to_mhlo/hlo_module_importer.cc
+++ b/third_party/xla/xla/translate/hlo_to_mhlo/hlo_module_importer.cc
@@ -154,7 +154,7 @@ absl::Status HloModuleImporter::Import(const HloModule& hlo_module) {
                            flatten_computation_args_result_)
                            .status());
 
-  return OkStatus();
+  return absl::OkStatus();
 }
 
 absl::Status HloModuleImporter::Import(const HloModuleProto& module_proto) {
diff --git a/third_party/xla/xla/translate/mhlo_to_hlo/translate.cc b/third_party/xla/xla/translate/mhlo_to_hlo/translate.cc
index 31a0e1b789b4cc..c7d30c709a6c06 100644
--- a/third_party/xla/xla/translate/mhlo_to_hlo/translate.cc
+++ b/third_party/xla/xla/translate/mhlo_to_hlo/translate.cc
@@ -105,7 +105,7 @@ absl::Status ConvertMlirHloToHloViaBuilder(
   auto hlo_module = computation.proto();
   hlo_proto->mutable_hlo_module()->Swap(&hlo_module);
 
-  return OkStatus();
+  return absl::OkStatus();
 }
 
 mlir::LogicalResult MlirHloToHloTextTranslateFunction(
diff --git a/third_party/xla/xla/tsl/util/byte_swap_array.h b/third_party/xla/xla/tsl/util/byte_swap_array.h
index 43ee9f77a38bac..a2ff2a864ee2dd 100644
--- a/third_party/xla/xla/tsl/util/byte_swap_array.h
+++ b/third_party/xla/xla/tsl/util/byte_swap_array.h
@@ -95,7 +95,7 @@ namespace tsl {
 //  bytes_per_elem: Number of bytes in each element of the array
 //  array_len: Number of elements in the array
 //
-// Returns: OkStatus() on success, -1 otherwise
+// Returns: absl::OkStatus() on success, -1 otherwise
 //
 absl::Status ByteSwapArray(char *array, size_t bytes_per_elem, int array_len);
 

From b8f00f09d1a8ba201f6163005a0a1e1d10bbbc68 Mon Sep 17 00:00:00 2001
From: Kyle Lucke <klucke@google.com>
Date: Mon, 20 May 2024 13:49:14 -0700
Subject: [PATCH 434/478] Remove never called TpuExecutor_DeallocateEvent.

PiperOrigin-RevId: 635558747
---
 third_party/xla/xla/stream_executor/tpu/tpu_executor_c_api.h   | 3 ---
 .../xla/xla/stream_executor/tpu/tpu_executor_init_fns.inc      | 1 -
 2 files changed, 4 deletions(-)

diff --git a/third_party/xla/xla/stream_executor/tpu/tpu_executor_c_api.h b/third_party/xla/xla/stream_executor/tpu/tpu_executor_c_api.h
index 68ddb6c03da435..065a6b18daee39 100644
--- a/third_party/xla/xla/stream_executor/tpu/tpu_executor_c_api.h
+++ b/third_party/xla/xla/stream_executor/tpu/tpu_executor_c_api.h
@@ -64,8 +64,6 @@ SE_TpuTopology_Core* TpuExecutor_GetCoreLocation(SE_StreamExecutor* executor);
 
 void TpuExecutor_AllocateEvent(SE_StreamExecutor* executor, SE_Event* event,
                                TF_Status* status);
-void TpuExecutor_DeallocateEvent(SE_StreamExecutor* executor, SE_Event* event,
-                                 TF_Status* status);
 int TpuExecutor_PollForEventStatus(SE_StreamExecutor* executor,
                                    SE_Event* event);
 void TpuExecutor_RecordEvent(SE_StreamExecutor* executor, SE_Stream* stream,
@@ -399,7 +397,6 @@ struct TfTpu_ExecutorApiFn {
   TFTPU_ADD_FN_IN_STRUCT(TpuExecutor_GetStatus);
   TFTPU_ADD_FN_IN_STRUCT(TpuExecutor_GetCoreLocation);
   TFTPU_ADD_FN_IN_STRUCT(TpuExecutor_AllocateEvent);
-  TFTPU_ADD_FN_IN_STRUCT(TpuExecutor_DeallocateEvent);
   TFTPU_ADD_FN_IN_STRUCT(TpuExecutor_PollForEventStatus);
   TFTPU_ADD_FN_IN_STRUCT(TpuExecutor_RecordEvent);
   TFTPU_ADD_FN_IN_STRUCT(TpuExecutor_WaitForEvent);
diff --git a/third_party/xla/xla/stream_executor/tpu/tpu_executor_init_fns.inc b/third_party/xla/xla/stream_executor/tpu/tpu_executor_init_fns.inc
index 4b5f8b64890efd..b4260442b5d2ae 100644
--- a/third_party/xla/xla/stream_executor/tpu/tpu_executor_init_fns.inc
+++ b/third_party/xla/xla/stream_executor/tpu/tpu_executor_init_fns.inc
@@ -27,7 +27,6 @@ absl::Status SetExecutorStructFn(
   TFTPU_SET_FN(executor_fn, TpuExecutor_GetStatus);
   TFTPU_SET_FN(executor_fn, TpuExecutor_GetCoreLocation);
   TFTPU_SET_FN(executor_fn, TpuExecutor_AllocateEvent);
-  TFTPU_SET_FN(executor_fn, TpuExecutor_DeallocateEvent);
   TFTPU_SET_FN(executor_fn, TpuExecutor_PollForEventStatus);
   TFTPU_SET_FN(executor_fn, TpuExecutor_RecordEvent);
   TFTPU_SET_FN(executor_fn, TpuExecutor_WaitForEvent);

From 1d1b6a77f9ede8936283f03394b5b9d2bac08887 Mon Sep 17 00:00:00 2001
From: Derek Murray <mrry@google.com>
Date: Mon, 20 May 2024 13:52:10 -0700
Subject: [PATCH 435/478] When pruning function bodies, enable pruning
 side-effect-free "stateful" ops.

Previously, function pruning would leave in place ops marked as "stateful" that are also side-effect free, even when they have no downstream dependencies. In the case of "ResourceGather" this op can trigger a relatively large amount of CPU work and memory bandwidth consumption, so it is valuable to be able to prune it.

PiperOrigin-RevId: 635559636
---
 .../core/common_runtime/function_def_utils.cc | 23 ++++++++++++++++--
 .../lower_function_call_op_test.cc            | 24 ++++++++++++++-----
 2 files changed, 39 insertions(+), 8 deletions(-)

diff --git a/tensorflow/core/common_runtime/function_def_utils.cc b/tensorflow/core/common_runtime/function_def_utils.cc
index ee7d1bb867c3a8..6cffeabc9205d4 100644
--- a/tensorflow/core/common_runtime/function_def_utils.cc
+++ b/tensorflow/core/common_runtime/function_def_utils.cc
@@ -31,6 +31,7 @@ limitations under the License.
 #include "tensorflow/core/graph/graph.h"
 #include "tensorflow/core/graph/graph_debug_info_builder.h"
 #include "tensorflow/core/platform/refcount.h"
+#include "tensorflow/core/platform/types.h"
 #include "tsl/platform/errors.h"
 #include "tsl/platform/hash.h"
 #include "tsl/platform/logging.h"
@@ -113,6 +114,25 @@ Status FunctionDefToBodyHelper(const FunctionDef& fdef, const AttrSlice& attrs,
                                  get_func_sig, fbody);
 }
 
+namespace {
+bool PrunableStatefulNode(const Node* n) {
+  // This set contains ops that are marked as "stateful" in their op
+  // registration, but can be pruned from a function graph if nothing depends
+  // on them. Typically, these are operations that are "impure" but have no
+  // side effects. For example, "ResourceGather" reads from a resource variable
+  // and can produce different results on each invocation (due to variable
+  // updates) but it does not itself modify the variable.
+  // TODO(b/341721055): Consolidate this set with other side effect modeling.
+  static const absl::flat_hash_set<string>* prunable_stateful_ops =
+      new absl::flat_hash_set<string>{
+          FunctionLibraryDefinition::kArgOp,
+          "ResourceGather",
+          "ResourceGatherNd",
+      };
+  return prunable_stateful_ops->contains(n->type_string());
+}
+}  // namespace
+
 // TODO(ezhulenev, skyewm): Function body should not have special treatment of
 // stateful ops, graph should encode nodes that must execute with `control_ret`
 // and `control_output`.
@@ -141,8 +161,7 @@ void PruneFunctionBody(const FunctionDef& fdef, Graph* g,
     // still needed. It would be preferable to prune entire loops and/or
     // conditionals if they are not used in the graph.
     if (n->IsControlFlow() ||
-        (n->op_def().is_stateful() &&
-         n->type_string() != FunctionLibraryDefinition::kArgOp) ||
+        (n->op_def().is_stateful() && !PrunableStatefulNode(n)) ||
         (control_ret_nodes.find(n->name()) != control_ret_nodes.end())) {
       nodes.insert(n);
     }
diff --git a/tensorflow/core/common_runtime/lower_function_call_op_test.cc b/tensorflow/core/common_runtime/lower_function_call_op_test.cc
index cc131835a38811..ea3de9500b9d3c 100644
--- a/tensorflow/core/common_runtime/lower_function_call_op_test.cc
+++ b/tensorflow/core/common_runtime/lower_function_call_op_test.cc
@@ -146,15 +146,23 @@ TEST(LowerFunctionCallTest, InlineFunctionCallAfterPruning) {
   // execute because it is in `control_ret`.
   // The `div` node and the unused arguments `j` and `k` should be pruned.
   *(f_lib_proto.add_function()) = FDH::Create(
-      "AddAndMul", {"i: int32", "j: int32", "k: int32"}, {"o: int32"}, {},
+      "AddAndMul", {"i: int32", "j: int32", "k: int32", "r: resource"},
+      {"o: int32"}, {},
       {{{"add"}, "Add", {"i", "i"}, {{"T", DT_INT32}}},
        {{"div"}, "FloorDiv", {"i", "i"}, {{"T", DT_INT32}}},
+       {{"gather"},
+        "ResourceGather",
+        {"r", "i"},
+        {{"Tindices", DT_INT32}, {"dtype", DT_FLOAT}}},
        {{"ret"}, "Mul", {"i", "i"}, {{"T", DT_INT32}}}},
       /*ret_def=*/{{"o", "ret:z:0"}},
       /*control_ret_def=*/{{"must_execute", "add"}});
 
   // Construct a graph:
-  //   A = Placeholder[dtype=int32]
+  //   X = Placeholder[dtype=int32]
+  //   Y = Placeholder[dtype=int32]
+  //   Z = Placeholder[dtype=int32]
+  //   R = Placeholder[dtype=resource]
   //   F = PartitionedCall[f=AddAndMul](a)
   //   B = Identity(func, ^func)
   Scope root = Scope::NewRootScope().ExitOnError();
@@ -162,13 +170,14 @@ TEST(LowerFunctionCallTest, InlineFunctionCallAfterPruning) {
   auto x = ops::Placeholder(root.WithOpName("X"), DT_INT32);
   auto y = ops::Placeholder(root.WithOpName("Y"), DT_INT32);
   auto z = ops::Placeholder(root.WithOpName("Z"), DT_INT32);
+  auto r = ops::Placeholder(root.WithOpName("R"), DT_RESOURCE);
   Node* function_call;
-  std::vector<NodeBuilder::NodeOut> inputs({NodeBuilder::NodeOut(x.node()),
-                                            NodeBuilder::NodeOut(y.node()),
-                                            NodeBuilder::NodeOut(z.node())});
+  std::vector<NodeBuilder::NodeOut> inputs(
+      {NodeBuilder::NodeOut(x.node()), NodeBuilder::NodeOut(y.node()),
+       NodeBuilder::NodeOut(z.node()), NodeBuilder::NodeOut(r.node())});
   TF_ASSERT_OK(NodeBuilder("F", "PartitionedCall", &root.graph()->flib_def())
                    .Input(inputs)
-                   .Attr("Tin", {DT_INT32, DT_INT32, DT_INT32})
+                   .Attr("Tin", {DT_INT32, DT_INT32, DT_INT32, DT_RESOURCE})
                    .Attr("Tout", {DT_INT32})
                    .Attr("f", FuncAttr("AddAndMul"))
                    .Finalize(root.graph(), &function_call));
@@ -186,17 +195,20 @@ TEST(LowerFunctionCallTest, InlineFunctionCallAfterPruning) {
   int add_count = 0;
   int mul_count = 0;
   int floor_div_count = 0;
+  int resource_gather_count = 0;
   for (const auto* op : graph->op_nodes()) {
     if (op->IsPartitionedCall()) partitioned_call_count++;
     if (op->type_string() == "Add") add_count++;
     if (op->type_string() == "Mul") mul_count++;
     if (op->type_string() == "FloorDiv") floor_div_count++;
+    if (op->type_string() == "ResourceGather") resource_gather_count++;
   }
 
   ASSERT_EQ(partitioned_call_count, 0);
   ASSERT_EQ(add_count, 1);
   ASSERT_EQ(mul_count, 1);
   ASSERT_EQ(floor_div_count, 0);
+  ASSERT_EQ(resource_gather_count, 0);
 
   // Verify execution.
   ClientSession session(root, SessionOptionsWithInlining());

From 98a7f0b73ceaf05539677ab43ebe3ae655080ab7 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Mon, 20 May 2024 14:19:23 -0700
Subject: [PATCH 436/478] Automated Code Change

PiperOrigin-RevId: 635568697
---
 tensorflow/core/common_runtime/next_pluggable_device/BUILD       | 1 +
 .../plugin_coordination_service_agent_helper.h                   | 1 +
 2 files changed, 2 insertions(+)

diff --git a/tensorflow/core/common_runtime/next_pluggable_device/BUILD b/tensorflow/core/common_runtime/next_pluggable_device/BUILD
index 41fd6207bd188d..58ccde7524408f 100644
--- a/tensorflow/core/common_runtime/next_pluggable_device/BUILD
+++ b/tensorflow/core/common_runtime/next_pluggable_device/BUILD
@@ -304,6 +304,7 @@ cc_library(
         "//tensorflow/c:kernels_hdrs",
         "//tensorflow/c:tf_status_helper",
         "//tensorflow/core/common_runtime/next_pluggable_device:flags",
+        "@com_google_absl//absl/flags:flag",
     ],
 )
 
diff --git a/tensorflow/core/common_runtime/next_pluggable_device/plugin_coordination_service_agent_helper.h b/tensorflow/core/common_runtime/next_pluggable_device/plugin_coordination_service_agent_helper.h
index 5d2bc60d74fc4f..a5adfa50656039 100644
--- a/tensorflow/core/common_runtime/next_pluggable_device/plugin_coordination_service_agent_helper.h
+++ b/tensorflow/core/common_runtime/next_pluggable_device/plugin_coordination_service_agent_helper.h
@@ -18,6 +18,7 @@ limitations under the License.
 
 #include <memory>
 
+#include "absl/flags/flag.h"
 #include "tensorflow/c/kernels.h"
 #include "tensorflow/c/tf_status_helper.h"
 #include "tensorflow/core/common_runtime/next_pluggable_device/c_plugin_coordination_service_agent.h"

From 651e9b8c8a372aebbedd9ba00a390802dfd294c1 Mon Sep 17 00:00:00 2001
From: Sandeep Dasgupta <sdasgup@google.com>
Date: Mon, 20 May 2024 14:19:45 -0700
Subject: [PATCH 437/478] Add SavedModel to StableHLO Converter to TensorFlow
 pip package

PiperOrigin-RevId: 635568823
---
 RELEASE.md                                    |   1 +
 tensorflow/BUILD                              |   1 +
 .../tensorflow_to_stablehlo/BUILD             |   3 -
 .../tensorflow_to_stablehlo/README.md         |  75 +++++++-
 .../tensorflow_to_stablehlo/python/BUILD      |  56 +++++-
 .../python/pywrap_tensorflow_to_stablehlo.cc  | 182 +++++-------------
 .../pywrap_tensorflow_to_stablehlo_lib.cc     | 141 ++++++++++++++
 .../pywrap_tensorflow_to_stablehlo_lib.h      |  67 +++++++
 tensorflow/python/BUILD                       |   3 +
 .../tools/def_file_filter/symbols_pybind.txt  |   4 +
 .../tools/def_file_filter/symbols_pybind.txt  |   4 +
 11 files changed, 390 insertions(+), 147 deletions(-)
 create mode 100644 tensorflow/compiler/mlir/quantization/tensorflow_to_stablehlo/python/pywrap_tensorflow_to_stablehlo_lib.cc
 create mode 100644 tensorflow/compiler/mlir/quantization/tensorflow_to_stablehlo/python/pywrap_tensorflow_to_stablehlo_lib.h

diff --git a/RELEASE.md b/RELEASE.md
index 0b33b6b9afef4e..8287988507e571 100644
--- a/RELEASE.md
+++ b/RELEASE.md
@@ -31,6 +31,7 @@
       been added to TF binary distributions (Python wheels).
 * Replace `DebuggerOptions` of TensorFlow Quantizer, and migrate to
   `DebuggerConfig` of StableHLO Quantizer.
+* Add TensorFlow to StableHLO converter to TensorFlow pip package.
 
 ## Keras
 
diff --git a/tensorflow/BUILD b/tensorflow/BUILD
index 0b2455baa5f7f3..a4cd4af8975bc2 100644
--- a/tensorflow/BUILD
+++ b/tensorflow/BUILD
@@ -1382,6 +1382,7 @@ tf_cc_shared_library(
         "//tensorflow/compiler/mlir/quantization/common/quantization_lib:quantization_config",
         "//tensorflow/compiler/mlir/lite/sparsity:sparsify_model",
         "//tensorflow/compiler/mlir/quantization/stablehlo/python:pywrap_quantization_lib_impl",
+        "//tensorflow/compiler/mlir/quantization/tensorflow_to_stablehlo/python:pywrap_tensorflow_to_stablehlo_lib_impl",
         "//tensorflow/compiler/mlir/quantization/tensorflow/calibrator:custom_aggregator_op",
         "//tensorflow/compiler/mlir/quantization/tensorflow/python:quantize_model_cc_impl",
         "//tensorflow/compiler/mlir/quantization/tensorflow:passes",
diff --git a/tensorflow/compiler/mlir/quantization/tensorflow_to_stablehlo/BUILD b/tensorflow/compiler/mlir/quantization/tensorflow_to_stablehlo/BUILD
index 922fc87b9d8ddf..ecc3c9e7ca6ebe 100644
--- a/tensorflow/compiler/mlir/quantization/tensorflow_to_stablehlo/BUILD
+++ b/tensorflow/compiler/mlir/quantization/tensorflow_to_stablehlo/BUILD
@@ -6,12 +6,10 @@ load("//tensorflow/compiler/mlir/quantization/stablehlo:internal_visibility_allo
 package_group(
     name = "internal_visibility_allowlist_package",
     packages = [
-        "//learning/brain/mlir/quantization/stablehlo/python/integration_test/...",
         "//tensorflow/compiler/mlir/lite/...",
         "//tensorflow/compiler/mlir/quantization/...",
         "//tensorflow/compiler/mlir/tf2xla/transforms/...",
         "//tensorflow/lite/...",
-        "//third_party/cloud_tpu/inference_converter/...",  # TPU Inference Converter V1
     ] + internal_visibility_allowlist(),
 )
 
@@ -80,7 +78,6 @@ tf_cc_binary(
 glob_lit_tests(
     name = "all_tests",
     data = [":test_utilities"],
-    # TODO: b/288344501 - Enable OSS tests again when stable-quant-opt works well.
     default_tags = [
         "no_oss",
         "no_pip",
diff --git a/tensorflow/compiler/mlir/quantization/tensorflow_to_stablehlo/README.md b/tensorflow/compiler/mlir/quantization/tensorflow_to_stablehlo/README.md
index 8a6fd0f5105c88..a65de3c38df001 100644
--- a/tensorflow/compiler/mlir/quantization/tensorflow_to_stablehlo/README.md
+++ b/tensorflow/compiler/mlir/quantization/tensorflow_to_stablehlo/README.md
@@ -1,13 +1,11 @@
-## Tensorflow SavedModel to StableHLO (tf-to-stablehlo-translate)
+# Tensorflow SavedModel to StableHLO (tf-to-stablehlo-translate)
 
-### Description
-
-This tool converts TensorFlow models (SavedModel or MLIR module) to StableHLO
-MLIR modules, preserving model structure and signatures. It enables seamless
+Converts TensorFlow models (SavedModel or MLIR module) to StableHLO MLIR
+modules, preserving model structure and signatures. It enables seamless
 integration of TensorFlow models into MLIR-based compiler frameworks for further
 optimization and deployment.
 
-### Usage
+## C++ APIs
 
 ```bash
 tf-to-stablehlo-translate \
@@ -58,3 +56,68 @@ tf-to-stablehlo-translate <saved-model-path> --input-arg-shapes=1,12:1,12:1,12
 * TensorFlow
 * MLIR
 * Abseil (absl)
+
+## Python APIs
+
+
+### `savedmodel_to_stablehlo`
+
+Converts a TensorFlow SavedModel into StableHLO bytecode.
+
+```Python
+from tensorflow.compiler.mlir.quantization.tensorflow_to_stablehlo.python import pywrap_tensorflow_to_stablehlo as tf2shlo
+
+stablehlo_bytes = tf2shlo.savedmodel_to_stablehlo(
+        input_path="/path/to/your/savedmodel",
+        exported_model_signatures=["serving_default"],
+        tag_names=["serve"],
+        input_arg_shapes_str="1,28,28,3::32"
+)
+
+```
+
+#### Arguments:
+
+* `input_path` (required): Path to your SavedModel directory.
+* `exported_model_signatures` (optional): List of signature names to convert.
+                                          Defaults to ["serving_default"].
+* `tag_names` (optional): List of tags associated with the SavedModel. Defaults
+                          to ["serve"].
+* `input_arg_shapes_str` (optional): A string representation of input argument
+                                     shapes for 'main' entry-point, separating
+                                     tensors with ':', dimension with ',', and
+                                     using '?' for unknown sizes. For example,
+                                     `input-arg-shapes=1,2::1,?` expresses
+                                     argument shapes `[1,2], [] and [1,?]`.
+
+#### Error Handling
+
+An exception will be raised with details about the error.
+
+### `tensorflow_module_to_stablehlo`
+
+Converts a TensorFlow MLIR module string into StableHLO bytecode.
+
+```Python
+from tensorflow.compiler.mlir.quantization.tensorflow_to_stablehlo.python import pywrap_tensorflow_to_stablehlo as tf2shlo
+
+stablehlo_bytes = tf2shlo.tensorflow_module_to_stablehlo(
+    module_op_str="your_tensorflow_mlir_module_string",
+    input_arg_shapes_str="1,28,28,3::32"
+)
+```
+
+#### Arguments:
+
+* `module_op_str` (required): String containing the TensorFlow MLIR module.
+* `input_arg_shapes_str` (optional): A string representation of input argument
+                                     shapes for 'main' entry-point, separating
+                                     tensors with ':', dimension with ',', and
+                                     using '?' for unknown sizes. For example,
+                                     `input-arg-shapes=1,2::1,?` expresses
+                                     argument shapes `[1,2], [] and [1,?]`.
+
+#### Error Handling
+
+Return `py::none()` (equivalent to Python's `None`) if there's an error. An
+exception will be raised with details about the error.
diff --git a/tensorflow/compiler/mlir/quantization/tensorflow_to_stablehlo/python/BUILD b/tensorflow/compiler/mlir/quantization/tensorflow_to_stablehlo/python/BUILD
index 86066ca5e6b6c9..f7a1c77026d215 100644
--- a/tensorflow/compiler/mlir/quantization/tensorflow_to_stablehlo/python/BUILD
+++ b/tensorflow/compiler/mlir/quantization/tensorflow_to_stablehlo/python/BUILD
@@ -1,5 +1,6 @@
 load(
     "//tensorflow:tensorflow.default.bzl",
+    "get_compatible_with_portable",
     "tf_py_strict_test",
     "tf_python_pybind_extension",
 )
@@ -20,6 +21,7 @@ package(
     default_visibility = [
         ":internal_visibility_allowlist_package",
         "//tensorflow:__pkg__",
+        "//tensorflow/python:__pkg__",
     ],
     licenses = ["notice"],
 )
@@ -45,20 +47,60 @@ package(
 # )
 # copybara:uncomment_end
 
-tf_python_pybind_extension(
-    name = "pywrap_tensorflow_to_stablehlo",
-    srcs = ["pywrap_tensorflow_to_stablehlo.cc"],
-    pytype_srcs = ["pywrap_tensorflow_to_stablehlo.pyi"],
+# This is a header-only target. The purpose of `pywrap_tensorflow_to_stablehlo_lib_*` targets is to expose only
+# the symbols that are required by `pywrap_tensorflow_to_stablehlo` that translates them to python functions.
+# The only intended use case of this library is by `pywrap_tensorflow_to_stablehlo`. Not letting
+# `pywrap_tensorflow_to_stablehlo` directly depend on sub-libraries like `static_range_srq` and instead haiving
+# a consolidated impl library `pywrap_tensorflow_to_stablehlo_lib_impl` allows the maintainers to avoid
+# declaring multiple impl libraries to `libtensorflow_cc` and `lib_pywrap_tensorflow_internal`,
+# which is required to avoid ODR violations.
+cc_library(
+    name = "pywrap_tensorflow_to_stablehlo_lib_header_only",
+    srcs = [],
+    hdrs = ["pywrap_tensorflow_to_stablehlo_lib.h"],
+    compatible_with = get_compatible_with_portable(),
+    visibility = ["//visibility:private"],  # ONLY for `pywrap_tensorflow_to_stablehlo`.
+    deps = [
+        "@com_google_absl//absl/status:statusor",
+        "@com_google_absl//absl/strings:string_view",
+    ],
+)
+
+# See the comments for `pywrap_tensorflow_to_stablehlo_lib_header_only`.
+cc_library(
+    name = "pywrap_tensorflow_to_stablehlo_lib_impl",
+    srcs = ["pywrap_tensorflow_to_stablehlo_lib.cc"],
+    hdrs = ["pywrap_tensorflow_to_stablehlo_lib.h"],
+    compatible_with = get_compatible_with_portable(),
+    visibility = [
+        "//tensorflow:__pkg__",  # For libtensorflow_cc.so.
+        "//tensorflow/python:__pkg__",  # For lib_pywrap_tensorflow_internal.so.
+    ],
     deps = [
-        "//tensorflow/compiler/mlir/quantization/tensorflow/python:type_casters",
         "//tensorflow/compiler/mlir/quantization/tensorflow_to_stablehlo:tf_to_stablehlo",
         "//tensorflow/compiler/mlir/tensorflow",
-        "@com_google_absl//absl/strings:str_format",
+        "//tensorflow/core:lib",
+        "//third_party/python_runtime:headers",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/status:statusor",
+        "@com_google_absl//absl/strings",
+        "@com_google_absl//absl/strings:string_view",
         "@llvm-project//llvm:Support",
-        "@llvm-project//mlir:CAPIIR",
+        "@llvm-project//mlir:BytecodeWriter",
         "@llvm-project//mlir:IR",
         "@llvm-project//mlir:Parser",
         "@llvm-project//mlir:Support",
+    ],
+)
+
+tf_python_pybind_extension(
+    name = "pywrap_tensorflow_to_stablehlo",
+    srcs = ["pywrap_tensorflow_to_stablehlo.cc"],
+    pytype_srcs = ["pywrap_tensorflow_to_stablehlo.pyi"],
+    # Each dependency MUST be either header-only or exclusive.
+    deps = [
+        ":pywrap_tensorflow_to_stablehlo_lib_header_only",
+        "//third_party/python_runtime:headers",
         "@pybind11",
         "@pybind11_abseil//pybind11_abseil:absl_casters",
         "@pybind11_abseil//pybind11_abseil:status_casters",
diff --git a/tensorflow/compiler/mlir/quantization/tensorflow_to_stablehlo/python/pywrap_tensorflow_to_stablehlo.cc b/tensorflow/compiler/mlir/quantization/tensorflow_to_stablehlo/python/pywrap_tensorflow_to_stablehlo.cc
index 54792f908026bb..1d1f775f5dfda1 100644
--- a/tensorflow/compiler/mlir/quantization/tensorflow_to_stablehlo/python/pywrap_tensorflow_to_stablehlo.cc
+++ b/tensorflow/compiler/mlir/quantization/tensorflow_to_stablehlo/python/pywrap_tensorflow_to_stablehlo.cc
@@ -12,159 +12,79 @@ 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.
 ==============================================================================*/
+#include <Python.h>
 
-#include "absl/strings/str_format.h"
-#include "llvm/Support/ToolOutputFile.h"
-#include "mlir/IR/OwningOpRef.h"  // from @llvm-project
-#include "mlir/Parser/Parser.h"  // from @llvm-project
-#include "mlir/Support/FileUtilities.h"  // from @llvm-project
 #include "pybind11/pybind11.h"  // from @pybind11
 #include "pybind11/pytypes.h"  // from @pybind11
-#include "pybind11/stl.h"  // from @pybind11   // IWYU pragma: keep
 #include "pybind11_abseil/absl_casters.h"  // from @pybind11_abseil    // IWYU pragma: keep
 #include "pybind11_abseil/status_casters.h"  // from @pybind11_abseil  // IWYU pragma: keep
-#include "tensorflow/compiler/mlir/quantization/tensorflow/python/type_casters.h"  // IWYU pragma: keep
-#include "tensorflow/compiler/mlir/quantization/tensorflow_to_stablehlo/tf_to_stablehlo.h"
-#include "tensorflow/compiler/mlir/tensorflow/dialect_registration.h"
+#include "tensorflow/compiler/mlir/quantization/tensorflow_to_stablehlo/python/pywrap_tensorflow_to_stablehlo_lib.h"
 
 namespace py = pybind11;
 
-namespace mlir::pywrap {
+namespace {
 
-absl::StatusOr<std::string> ModuleToBytecode(ModuleOp module) {
-  std::string bytecode;
-  llvm::raw_string_ostream os(bytecode);
-  mlir::BytecodeWriterConfig config;
-  if (mlir::failed(mlir::writeBytecodeToFile(module, os, config))) {
-    return absl::InvalidArgumentError("mlir::writeBytecodeToFile failed");
-  }
-  return bytecode;
-}
-
-absl::StatusOr<std::string> ExportModule(ModuleOp module) {
-  auto output_filename = absl::StrFormat(
-      "%s/tf_module.mlir", std::filesystem::temp_directory_path());
-
-  std::string error_msg;
-  auto output = openOutputFile(output_filename, &error_msg);
-  if (output == nullptr) {
-    return absl::AbortedError(
-        absl::StrCat("Unable to open output path: ", error_msg));
-  }
-
-  std::string result;
-  llvm::raw_string_ostream os(result);
-  OpPrintingFlags printing_flags;
-  module.print(os, printing_flags);
-
-  output->os() << result;
-  output->keep();
-
-  return output_filename;
-}
-
-py::bytes PywrapSavedModelToStablehlo(
-    absl::string_view input_path,
-    const std::vector<std::string>& exported_model_signatures =
-        {"serving_default"},
-    const std::vector<std::string>& tag_names = {"serve"},
-    absl::string_view input_arg_shapes_str = "") {
-  mlir::DialectRegistry registry;
-  RegisterAllTensorFlowDialects(registry);
-  mlir::MLIRContext context(registry);
-  context.loadAllAvailableDialects();
-
-  auto module =
-      TfToStablehlo(input_path, &context, exported_model_signatures, tag_names,
-                    input_arg_shapes_str, /*is_input_mlir_module=*/false);
-
-  if (!module.ok()) {
-    PyErr_SetString(PyExc_ValueError,
-                    "failed to converted TensorFlow to StableHLO");
-    return {};
-  }
-
-  auto bytecode = ModuleToBytecode(module.value().get());
-  if (!bytecode.ok()) {
-    PyErr_SetString(PyExc_ValueError, "failed to write module to bytecode");
-    return {};
-  }
-
-  return bytecode.value();
-}
-
-py::bytes PywrapTfModuleToStablehlo(
-    absl::string_view module_op_str,
-    absl::string_view input_arg_shapes_str = "") {
-  mlir::DialectRegistry registry;
-  RegisterAllTensorFlowDialects(registry);
-  mlir::MLIRContext context(registry);
-  context.loadAllAvailableDialects();
-
-  auto tf_module = mlir::parseSourceString<ModuleOp>(module_op_str, &context);
-  if (!tf_module) {
-    PyErr_SetString(PyExc_ValueError, "failed to parse TF module string");
-    return {};
-  }
-
-  auto mlir_file_path = ExportModule(*tf_module);
-  if (!mlir_file_path.ok()) {
-    PyErr_SetString(PyExc_ValueError,
-                    "failed to write TF module to a temporary file");
-    return {};
-  }
-
-  auto module = TfToStablehlo(
-      *mlir_file_path, &context, /*exported_model_signatures=*/{},
-      /*tag_names=*/{}, input_arg_shapes_str, /*is_input_mlir_module=*/true);
-
-  if (!module.ok()) {
-    PyErr_SetString(PyExc_ValueError,
-                    "failed to converted TensorFlow to StableHLO");
-    return {};
-  }
-
-  auto bytecode = ModuleToBytecode(module.value().get());
-  if (!bytecode.ok()) {
-    PyErr_SetString(PyExc_ValueError, "failed to write module to bytecode");
-    return {};
-  }
-
-  return bytecode.value();
-}
+using mlir::tensorflow_to_stablehlo::pywrap::PywrapSavedModelToStablehlo;
+using mlir::tensorflow_to_stablehlo::pywrap::PywrapTfModuleToStablehlo;
 
-}  // namespace mlir::pywrap
+}  // namespace
 
 PYBIND11_MODULE(pywrap_tensorflow_to_stablehlo, m) {
   m.doc() = "TensorFlow to StableHLO APIs.";
 
   // LINT.IfChange(savedmodel_to_stablehlo)
-  m.def("savedmodel_to_stablehlo", &mlir::pywrap::PywrapSavedModelToStablehlo,
-        R"pbdoc(
-        This tool converts TensorFlow SavedModel to StableHLO.
+  m.def(
+      "savedmodel_to_stablehlo",
+      [](absl::string_view input_path,
+         const std::vector<std::string>& exported_model_signatures =
+             {"serving_default"},
+         const std::vector<std::string>& tag_names = {"serve"},
+         absl::string_view input_arg_shapes_str = "") -> py::bytes {
+        auto module_bytecode =
+            PywrapSavedModelToStablehlo(input_path, exported_model_signatures,
+                                        tag_names, input_arg_shapes_str);
+        if (!module_bytecode.ok()) {
+          PyErr_SetString(PyExc_ValueError,
+                          module_bytecode.status().ToString().c_str());
+          throw py::error_already_set();
+        }
+        return py::bytes(module_bytecode.value());
+      },
+      R"pbdoc(
+        Converts a TensorFlow SavedModel into StableHLO bytecode.
 
         * input-path: The path to the input TensorFlow SavedModel.
         * exported-model-signatures: Comma-separated list of exported model
-          signatures to convert. Ignored for MLIR input.
-        * tag_names: Comma-separated list of tags for loading SavedModel. Ignored for MLIR
-          input.
+          signatures to convert.
+        * tag_names: Comma-separated list of tags for loading SavedModel.
         * input-arg-shapes: A string representation of input argument shapes for
           'main' entry-point, separating tensors with ':', dimension with ',', and
           using '?' for unknown sizes. For example, 'input-arg-shapes=1,2::1,?'
           expresses argument shapes [1,2], [] and [1,?].
         )pbdoc",
-        py::arg("input_path"),
-        py::arg("exported_model_signatures") =
-            std::vector<std::string>{"serving_default"},
-        py::arg("tag_names") = std::vector<std::string>{"serve"},
-        py::arg("input_arg_shapes_str") = "");
-  // LINT.ThenChange(pywrap_tensorflow_to_stablehlo.pyi:tensorflow_to_stablehlo)
+      py::arg("input_path"),
+      py::arg("exported_model_signatures") =
+          std::vector<std::string>{"serving_default"},
+      py::arg("tag_names") = std::vector<std::string>{"serve"},
+      py::arg("input_arg_shapes_str") = "");
+  // LINT.ThenChange(pywrap_tensorflow_to_stablehlo.pyi:savedmodel_to_stablehlo)
   //
-  // LINT.IfChange(tensorflow_mlir_to_stablehlo)
-  m.def("tensorflow_module_to_stablehlo",
-        &mlir::pywrap::PywrapTfModuleToStablehlo,
-        R"pbdoc(
-        This tool converts TensorFlow mlir module string to StableHLO.
+  // LINT.IfChange(tensorflow_module_to_stablehlo)
+  m.def(
+      "tensorflow_module_to_stablehlo",
+      [](absl::string_view module_op_str,
+         absl::string_view input_arg_shapes_str) -> py::bytes {
+        auto module_bytecode =
+            PywrapTfModuleToStablehlo(module_op_str, input_arg_shapes_str);
+        if (!module_bytecode.ok()) {
+          PyErr_SetString(PyExc_ValueError,
+                          module_bytecode.status().ToString().c_str());
+          throw py::error_already_set();
+        }
+        return py::bytes(module_bytecode.value());
+      },
+      R"pbdoc(
+        Converts a TensorFlow MLIR module string into StableHLO bytecode.
 
         * module: TensorFlow MLIR module string.
         * input-arg-shapes: A string representation of input argument shapes for
@@ -172,6 +92,6 @@ PYBIND11_MODULE(pywrap_tensorflow_to_stablehlo, m) {
           using '?' for unknown sizes. For example, 'input-arg-shapes=1,2::1,?'
           expresses argument shapes [1,2], [] and [1,?].
         )pbdoc",
-        py::arg("module"), py::arg("input_arg_shapes_str") = "");
-  // LINT.ThenChange(pywrap_tensorflow_to_stablehlo.pyi:tensorflow_mlir_to_stablehlo)
+      py::arg("module"), py::arg("input_arg_shapes_str") = "");
+  // LINT.ThenChange(pywrap_tensorflow_to_stablehlo.pyi:tensorflow_module_to_stablehlo)
 }
diff --git a/tensorflow/compiler/mlir/quantization/tensorflow_to_stablehlo/python/pywrap_tensorflow_to_stablehlo_lib.cc b/tensorflow/compiler/mlir/quantization/tensorflow_to_stablehlo/python/pywrap_tensorflow_to_stablehlo_lib.cc
new file mode 100644
index 00000000000000..cbd535a861482f
--- /dev/null
+++ b/tensorflow/compiler/mlir/quantization/tensorflow_to_stablehlo/python/pywrap_tensorflow_to_stablehlo_lib.cc
@@ -0,0 +1,141 @@
+/* Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+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.
+==============================================================================*/
+
+#include "tensorflow/compiler/mlir/quantization/tensorflow_to_stablehlo/python/pywrap_tensorflow_to_stablehlo_lib.h"
+
+#include <string>
+#include <vector>
+
+#include "absl/status/status.h"
+#include "absl/status/statusor.h"
+#include "absl/strings/str_cat.h"
+#include "absl/strings/string_view.h"
+#include "llvm/Support/ToolOutputFile.h"
+#include "llvm/Support/raw_ostream.h"
+#include "mlir/Bytecode/BytecodeWriter.h"  // from @llvm-project
+#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
+#include "mlir/IR/DialectRegistry.h"  // from @llvm-project
+#include "mlir/IR/MLIRContext.h"  // from @llvm-project
+#include "mlir/IR/OperationSupport.h"  // from @llvm-project
+#include "mlir/Parser/Parser.h"  // from @llvm-project
+#include "mlir/Support/FileUtilities.h"  // from @llvm-project
+#include "mlir/Support/LogicalResult.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/quantization/tensorflow_to_stablehlo/tf_to_stablehlo.h"
+#include "tensorflow/compiler/mlir/tensorflow/dialect_registration.h"
+#include "tensorflow/core/platform/path.h"
+
+namespace mlir::tensorflow_to_stablehlo::pywrap {
+
+absl::StatusOr<std::string> ModuleToBytecode(ModuleOp module) {
+  std::string bytecode;
+  llvm::raw_string_ostream os(bytecode);
+  mlir::BytecodeWriterConfig config;
+  if (mlir::failed(mlir::writeBytecodeToFile(module, os, config))) {
+    return absl::InvalidArgumentError("mlir::writeBytecodeToFile failed");
+  }
+  return bytecode;
+}
+
+absl::StatusOr<std::string> ExportModule(ModuleOp module) {
+  const std::string output_filename = tensorflow::io::GetTempFilename(".mlir");
+  std::string error_msg;
+  auto output = openOutputFile(output_filename, &error_msg);
+  if (output == nullptr) {
+    return absl::UnknownError(
+        absl::StrCat("Unable to open output path: ", error_msg));
+  }
+
+  std::string result;
+  llvm::raw_string_ostream os(result);
+  OpPrintingFlags printing_flags;
+  module.print(os, printing_flags);
+
+  output->os() << result;
+  output->keep();
+
+  return output_filename;
+}
+
+absl::StatusOr<std::string> PywrapSavedModelToStablehlo(
+    absl::string_view input_path,
+    const std::vector<std::string>& exported_model_signatures,
+    const std::vector<std::string>& tag_names,
+    absl::string_view input_arg_shapes_str) {
+  mlir::DialectRegistry registry;
+  RegisterAllTensorFlowDialects(registry);
+  mlir::MLIRContext context(registry);
+  context.loadAllAvailableDialects();
+
+  auto module =
+      TfToStablehlo(input_path, &context, exported_model_signatures, tag_names,
+                    input_arg_shapes_str, /*is_input_mlir_module=*/false);
+
+  if (!module.ok()) {
+    return absl::UnknownError(
+        absl::StrCat("Failed to convert SavedModel to StableHLO: ",
+                     module.status().message()));
+  }
+
+  auto bytecode = ModuleToBytecode(module.value().get());
+  if (!bytecode.ok()) {
+    return absl::UnknownError(
+        absl::StrCat("Failed to serialize MLIR module to bytecode: ",
+                     bytecode.status().message()));
+  }
+
+  return bytecode.value();
+}
+
+absl::StatusOr<std::string> PywrapTfModuleToStablehlo(
+    absl::string_view module_op_str, absl::string_view input_arg_shapes_str) {
+  mlir::DialectRegistry registry;
+  RegisterAllTensorFlowDialects(registry);
+  mlir::MLIRContext context(registry);
+  context.loadAllAvailableDialects();
+
+  auto tf_module = mlir::parseSourceString<ModuleOp>(module_op_str, &context);
+  if (!tf_module) {
+    return absl::UnknownError("Failed to parse MLIR module");
+  }
+
+  auto mlir_file_path = ExportModule(*tf_module);
+  if (!mlir_file_path.ok()) {
+    return absl::UnknownError(
+        absl::StrCat("Failed to write MLIR module to file.",
+                     mlir_file_path.status().message()));
+  }
+
+  auto module = TfToStablehlo(*mlir_file_path, &context,
+                              /*exported_model_signatures=*/{},
+                              /*tag_names=*/{}, input_arg_shapes_str,
+                              /*is_input_mlir_module=*/true);
+
+  if (!module.ok()) {
+    return absl::UnknownError(
+        absl::StrCat(" Failed to convert SavedModel to StableHLO: ",
+                     module.status().message()));
+  }
+
+  auto bytecode = ModuleToBytecode(module.value().get());
+  if (!bytecode.ok()) {
+    return absl::UnknownError(
+        absl::StrCat("Failed to serialize MLIR module to bytecode: ",
+                     bytecode.status().message()));
+  }
+
+  return bytecode.value();
+}
+
+}  // namespace mlir::tensorflow_to_stablehlo::pywrap
diff --git a/tensorflow/compiler/mlir/quantization/tensorflow_to_stablehlo/python/pywrap_tensorflow_to_stablehlo_lib.h b/tensorflow/compiler/mlir/quantization/tensorflow_to_stablehlo/python/pywrap_tensorflow_to_stablehlo_lib.h
new file mode 100644
index 00000000000000..c79ed32b990dd6
--- /dev/null
+++ b/tensorflow/compiler/mlir/quantization/tensorflow_to_stablehlo/python/pywrap_tensorflow_to_stablehlo_lib.h
@@ -0,0 +1,67 @@
+/* Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+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.
+==============================================================================*/
+
+#ifndef TENSORFLOW_COMPILER_MLIR_QUANTIZATION_TENSORFLOW_TO_STABLEHLO_PYTHON_PYWRAP_TENSORFLOW_TO_STABLEHLO_LIB_H_
+#define TENSORFLOW_COMPILER_MLIR_QUANTIZATION_TENSORFLOW_TO_STABLEHLO_PYTHON_PYWRAP_TENSORFLOW_TO_STABLEHLO_LIB_H_
+
+#include <string>
+#include <vector>
+
+#include "absl/status/statusor.h"
+#include "absl/strings/string_view.h"
+
+namespace mlir::tensorflow_to_stablehlo::pywrap {
+
+// Converts a TensorFlow SavedModel to a StableHLO MLIR module and serializes it
+// to bytecode.
+//
+// Args:
+//   input_path: The path to the SavedModel directory.
+//    exported_model_signatures: Comma-separated list of exported model
+//   signatures to convert. tag_names: Comma-separated list of tags for loading
+//    SavedModel.
+//   input_arg_shapes_str: A string representation of input argument
+//    shapes for 'main' entry-point, separating tensors with ':', dimension
+//    with ',', and using '?' for unknown sizes. For example,
+//    'input-arg-shapes=1,2::1,?' expresses argument shapes [1,2], [] and [1,?].
+//
+// Returns:
+//   An absl::StatusOr containing the serialized bytecode of the StableHLO
+//   module on success, or an error status on failure.
+absl::StatusOr<std::string> PywrapSavedModelToStablehlo(
+    absl::string_view input_path,
+    const std::vector<std::string>& exported_model_signatures,
+    const std::vector<std::string>& tag_names,
+    absl::string_view input_arg_shapes_str);
+
+// Converts a TensorFlow MLIR module string to a StableHLO MLIR module and
+// serializes it to bytecode.
+//
+// Args:
+//   module_op_str: TensorFlow MLIR module string.
+//   input_arg_shapes_str: A string representation of input argument
+//    shapes for 'main' entry-point, separating tensors with ':', dimension
+//    with ',', and using '?' for unknown sizes. For example,
+//    'input-arg-shapes=1,2::1,?' expresses argument shapes [1,2], [] and [1,?].
+//
+// Returns:
+//   An absl::StatusOr containing the serialized bytecode of the StableHLO
+//   module on success, or an error status on failure.
+absl::StatusOr<std::string> PywrapTfModuleToStablehlo(
+    absl::string_view module_op_str, absl::string_view input_arg_shapes_str);
+
+}  // namespace mlir::tensorflow_to_stablehlo::pywrap
+
+#endif  // TENSORFLOW_COMPILER_MLIR_QUANTIZATION_TENSORFLOW_TO_STABLEHLO_PYTHON_PYWRAP_TENSORFLOW_TO_STABLEHLO_LIB_H_
diff --git a/tensorflow/python/BUILD b/tensorflow/python/BUILD
index b9a762c607b4eb..199cf34c4f27ab 100644
--- a/tensorflow/python/BUILD
+++ b/tensorflow/python/BUILD
@@ -134,6 +134,7 @@ py_strict_library(
         ":pywrap_tensorflow",
         ":pywrap_tfe",
         "//tensorflow/compiler/mlir/quantization/tensorflow/python:quantize_model",
+        "//tensorflow/compiler/mlir/quantization/tensorflow_to_stablehlo/python:pywrap_tensorflow_to_stablehlo",
         "//tensorflow/core:protos_all_py",
         "//tensorflow/dtensor/python:dtensor",
         "//tensorflow/python/autograph",
@@ -768,6 +769,7 @@ pywrap_tensorflow_macro(
         "//tensorflow/cc/saved_model:metrics_impl",
         "//tensorflow/compiler/mlir/quantization/stablehlo/python:pywrap_quantization_lib_impl",
         "//tensorflow/compiler/mlir/quantization/tensorflow/python:quantize_model_cc_impl",
+        "//tensorflow/compiler/mlir/quantization/tensorflow_to_stablehlo/python:pywrap_tensorflow_to_stablehlo_lib_impl",
         "//tensorflow/compiler/mlir/tensorflow/c:mlir_c_api_registration",
         "//tensorflow/compiler/tf2tensorrt:op_converter_registry_impl",
         "//tensorflow/compiler/tf2xla:tf2xla_opset",
@@ -883,6 +885,7 @@ filegroup(
         "//tensorflow/compiler/jit:flags",  # tfe
         "//tensorflow/compiler/jit:get_compiler_ir",  # tfe
         "//tensorflow/compiler/mlir/quantization/tensorflow/python:quantize_model_cc_impl",  # quantization
+        "//tensorflow/compiler/mlir/quantization/tensorflow_to_stablehlo/python:pywrap_tensorflow_to_stablehlo_lib_impl",  # tensorflow_to_stablehlo
         "//tensorflow/compiler/tf2xla:tf2xla_opset",  # pywrap_xla_ops
         "//tensorflow/core:framework_internal_impl",  # op_def_registry
         "//tensorflow/core:lib_internal_impl",  # device_lib
diff --git a/tensorflow/tools/def_file_filter/symbols_pybind.txt b/tensorflow/tools/def_file_filter/symbols_pybind.txt
index c9faf5ad6b8375..69aea896b57102 100644
--- a/tensorflow/tools/def_file_filter/symbols_pybind.txt
+++ b/tensorflow/tools/def_file_filter/symbols_pybind.txt
@@ -566,6 +566,10 @@ tensorflow::quantization::QuantizeStaticRangePtq
 tensorflow::quantization::QuantizeDynamicRangePtq
 tensorflow::quantization::QuantizeWeightOnly
 
+[//tensorflow/compiler/mlir/quantization/tensorflow_to_stablehlo/python:pywrap_tensorflow_to_stablehlo_lib_impl] # tensorflow_to_stablehlo
+mlir::tensorflow_to_stablehlo::pywrap::PywrapSavedModelToStablehlo
+mlir::tensorflow_to_stablehlo::pywrap::PywrapTfModuleToStablehlo
+
 [//tensorflow/dtensor/cc:dtensor_device_cc] # DTensor
 tensorflow::dtensor::AllocateDTensorDevice
 tensorflow::dtensor::AddMesh
diff --git a/third_party/xla/third_party/tsl/tools/def_file_filter/symbols_pybind.txt b/third_party/xla/third_party/tsl/tools/def_file_filter/symbols_pybind.txt
index c9faf5ad6b8375..69aea896b57102 100644
--- a/third_party/xla/third_party/tsl/tools/def_file_filter/symbols_pybind.txt
+++ b/third_party/xla/third_party/tsl/tools/def_file_filter/symbols_pybind.txt
@@ -566,6 +566,10 @@ tensorflow::quantization::QuantizeStaticRangePtq
 tensorflow::quantization::QuantizeDynamicRangePtq
 tensorflow::quantization::QuantizeWeightOnly
 
+[//tensorflow/compiler/mlir/quantization/tensorflow_to_stablehlo/python:pywrap_tensorflow_to_stablehlo_lib_impl] # tensorflow_to_stablehlo
+mlir::tensorflow_to_stablehlo::pywrap::PywrapSavedModelToStablehlo
+mlir::tensorflow_to_stablehlo::pywrap::PywrapTfModuleToStablehlo
+
 [//tensorflow/dtensor/cc:dtensor_device_cc] # DTensor
 tensorflow::dtensor::AllocateDTensorDevice
 tensorflow::dtensor::AddMesh

From 43d498cfd19de12282d416ebd6b7993f592695cd Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Mon, 20 May 2024 15:03:27 -0700
Subject: [PATCH 438/478] Automated Code Change

PiperOrigin-RevId: 635581142
---
 tensorflow/core/common_runtime/graph_execution_state.cc | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/tensorflow/core/common_runtime/graph_execution_state.cc b/tensorflow/core/common_runtime/graph_execution_state.cc
index c2e115b80c3cb9..744e72fcf75d57 100644
--- a/tensorflow/core/common_runtime/graph_execution_state.cc
+++ b/tensorflow/core/common_runtime/graph_execution_state.cc
@@ -778,11 +778,11 @@ Status GraphExecutionState::OptimizeGraph(
         node_names.insert(node->name());
       }
     }
-    for (const auto& feed : item.feed) {
+    for (auto& feed : item.feed) {
       SafeTensorId tensor_id = ParseTensorName(feed.first);
       if (node_names.find(tensor_id.node()) == node_names.end()) {
         return errors::InvalidArgument("Invalid feed, no such node in graph: ",
-                                       feed.first);
+                                       std::move(feed.first));
       }
     }
     for (const auto& fetch : item.fetch) {

From 482560633fd8f87e957626e8bfcda55b405fe85f Mon Sep 17 00:00:00 2001
From: Quoc Truong <quoct@google.com>
Date: Mon, 20 May 2024 15:08:23 -0700
Subject: [PATCH 439/478] Add support for building and pushing to Artifact
 Registry for Linux ARM64 images.

PiperOrigin-RevId: 635582651
---
 ci/official/containers/linux_arm64/build.sh | 8 +++++++-
 1 file changed, 7 insertions(+), 1 deletion(-)

diff --git a/ci/official/containers/linux_arm64/build.sh b/ci/official/containers/linux_arm64/build.sh
index 5d6a40658bd782..611d5f48ac0084 100755
--- a/ci/official/containers/linux_arm64/build.sh
+++ b/ci/official/containers/linux_arm64/build.sh
@@ -40,11 +40,15 @@ else
   fi
 fi
 
+# TODO(b/341050361): When these steps are verified, removed the GCR image code.
+AR_IMAGE_PATH="us-central1-docker.pkg.dev/tensorflow-sigs/tensorflow/build-arm64"
+
 # Build for both JAX and TF usage.  We do these in one place because they share
 # almost all of the same cache layers
 export DOCKER_BUILDKIT=1
 for target in jax tf; do
   IMAGE="gcr.io/tensorflow-sigs/build-arm64:$target-$TAG"
+  AR_IMAGE="$AR_IMAGE_PATH:$target-$TAG"
   docker pull "$IMAGE" || true
   # Due to some flakiness of resources pulled in the build, allow the docker
   # command to reattempt build a few times in the case of failure (b/302558736)
@@ -55,7 +59,7 @@ for target in jax tf; do
     --build-arg REQUIREMENTS_FILE=jax.requirements.txt \
     --target=$target \
     --cache-from "$IMAGE" \
-    -t "$IMAGE"  . && break
+    -t "$IMAGE" -t "$AR_IMAGE" . && break
   done
   final=$?
   if [ $final -ne 0 ]; then
@@ -66,5 +70,7 @@ for target in jax tf; do
   if [[ -n "$KOKORO_BUILD_ID" ]]; then
     gcloud auth configure-docker
     docker push "$IMAGE"
+    gcloud auth configure-docker us-central1-docker.pkg.dev
+    docker push "$AR_IMAGE"
   fi
 done

From dec844f37e94d7fca0f802fb87b583836aed5f13 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Mon, 20 May 2024 16:17:06 -0700
Subject: [PATCH 440/478] Add restore_uid to TPUEmbeddingShardedVariable.

PiperOrigin-RevId: 635600618
---
 tensorflow/python/tpu/tpu_embedding_v3.py | 20 +++++++++++++++++---
 1 file changed, 17 insertions(+), 3 deletions(-)

diff --git a/tensorflow/python/tpu/tpu_embedding_v3.py b/tensorflow/python/tpu/tpu_embedding_v3.py
index 9fcfa060542070..c1eb333915a0d8 100644
--- a/tensorflow/python/tpu/tpu_embedding_v3.py
+++ b/tensorflow/python/tpu/tpu_embedding_v3.py
@@ -1864,7 +1864,9 @@ def _stack_gradients(self, gradients):
 # this file is OSSed.
 def extract_variable_info(
     kwargs: Any,
-) -> Tuple[str, Tuple[int, ...], dtypes.DType, Callable[[], Any]]:
+) -> Tuple[
+    str, Tuple[int, ...], dtypes.DType, Callable[[], Any], Optional[int]
+]:
   """Extracts the variable creation attributes from the kwargs.
 
   Args:
@@ -1872,8 +1874,13 @@ def extract_variable_info(
       scope.
 
   Returns:
-    A tuple of variable name, shape, dtype, initialization function.
+    A tuple of variable name, shape, dtype, initialization function,
+    restore_uid.
   """
+
+  def get_restore_uid(initial_value: Callable[..., Any]) -> int | None:
+    return getattr(initial_value, "restore_uid", None)
+
   if isinstance(kwargs["initial_value"], functools.partial) and (
       "shape" in kwargs["initial_value"].keywords
       or kwargs["initial_value"].args
@@ -1888,6 +1895,7 @@ def extract_variable_info(
         shape,
         kwargs["initial_value"].keywords.get("dtype", kwargs["dtype"]),
         kwargs["initial_value"].func,
+        get_restore_uid(kwargs["initial_value"].func),
     )
   elif (
       "shape" not in kwargs
@@ -1909,6 +1917,7 @@ def extract_variable_info(
         kwargs["shape"],
         kwargs["dtype"],
         kwargs["initial_value"],
+        get_restore_uid(kwargs["initial_value"]),
     )
 
 
@@ -1962,7 +1971,9 @@ def _create_sharded_variable(next_creator, *args, **kwargs):
           "shard_info must be in arguments of the init function."
       )
 
-    name, shape, dtype, unwrapped_initial_value = extract_variable_info(kwargs)
+    name, shape, dtype, unwrapped_initial_value, restore_uid = (
+        extract_variable_info(kwargs)
+    )
 
     shape = ops.tensor_shape.TensorShape(shape)
     num_devices = num_replicas * num_cores_per_replica
@@ -2006,6 +2017,9 @@ def _create_sharded_variable(next_creator, *args, **kwargs):
     result = TPUEmbeddingShardedVariable(
         strategy, variables, tf_variables.VariableAggregation.NONE, None
     )
+    if restore_uid is not None:
+      result._maybe_initialize_trackable()  # pylint: disable=protected-access
+      result._update_uid = restore_uid  # pylint: disable=protected-access
     return result
 
   return _create_sharded_variable

From e03a6f2385f7787d829cded83339173ce1ca85db Mon Sep 17 00:00:00 2001
From: Kuy Mainwaring <kuym@google.com>
Date: Mon, 20 May 2024 16:43:02 -0700
Subject: [PATCH 441/478] [XLA:GPU] Clang-tidy cleanup for
 xla/service/gpu/amdgpu_compiler.cc

PiperOrigin-RevId: 635607046
---
 third_party/xla/xla/service/gpu/BUILD              | 1 +
 third_party/xla/xla/service/gpu/amdgpu_compiler.cc | 5 ++++-
 2 files changed, 5 insertions(+), 1 deletion(-)

diff --git a/third_party/xla/xla/service/gpu/BUILD b/third_party/xla/xla/service/gpu/BUILD
index f3e0c8f11adc03..60169c2a9a9907 100644
--- a/third_party/xla/xla/service/gpu/BUILD
+++ b/third_party/xla/xla/service/gpu/BUILD
@@ -3998,6 +3998,7 @@ cc_library(
         "//xla/service:convert_mover",
         "//xla/service:dot_dimension_merger",
         "//xla/service:float_normalization",
+        "//xla/service:float_support",
         "//xla/service:hlo_constant_folding",
         "//xla/service:hlo_module_config",
         "//xla/service:hlo_pass",
diff --git a/third_party/xla/xla/service/gpu/amdgpu_compiler.cc b/third_party/xla/xla/service/gpu/amdgpu_compiler.cc
index 4224fe11d372d3..d5936c62a91244 100644
--- a/third_party/xla/xla/service/gpu/amdgpu_compiler.cc
+++ b/third_party/xla/xla/service/gpu/amdgpu_compiler.cc
@@ -25,11 +25,13 @@ limitations under the License.
 #include "xla/hlo/ir/hlo_computation.h"
 #include "xla/hlo/ir/hlo_instruction.h"
 #include "xla/hlo/ir/hlo_module.h"
+#include "xla/hlo/ir/hlo_opcode.h"
 #include "xla/service/algebraic_simplifier.h"
 #include "xla/service/call_inliner.h"
 #include "xla/service/convert_mover.h"
 #include "xla/service/dot_dimension_merger.h"
 #include "xla/service/float_normalization.h"
+#include "xla/service/float_support.h"
 #include "xla/service/gpu/autotuner_util.h"
 #include "xla/service/gpu/conv_algorithm_picker.h"
 #include "xla/service/gpu/cublas_pad_for_gemms.h"
@@ -70,7 +72,8 @@ namespace gpu {
 
 namespace {
 
-struct ConvBfloat16Support : public FloatSupport {
+class ConvBfloat16Support : public FloatSupport {
+ public:
   explicit ConvBfloat16Support(const se::RocmComputeCapability& rocm)
       : FloatSupport(BF16),
         // TODO: MIOpen does not support bf16 convolutions yet

From e922bfc0d2f65da969f9a36fa7d888bb1f763cb2 Mon Sep 17 00:00:00 2001
From: Kuy Mainwaring <kuym@google.com>
Date: Mon, 20 May 2024 17:03:36 -0700
Subject: [PATCH 442/478] [XLA:GPU] Clang-tidy cleanup for
 xla/service/gpu/gpu_conv_rewriter_test.cc

PiperOrigin-RevId: 635612157
---
 third_party/xla/xla/service/gpu/gpu_conv_rewriter_test.cc | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/third_party/xla/xla/service/gpu/gpu_conv_rewriter_test.cc b/third_party/xla/xla/service/gpu/gpu_conv_rewriter_test.cc
index add6e0f53ea506..ee5d2a5dbcce4e 100644
--- a/third_party/xla/xla/service/gpu/gpu_conv_rewriter_test.cc
+++ b/third_party/xla/xla/service/gpu/gpu_conv_rewriter_test.cc
@@ -46,8 +46,8 @@ namespace m = ::xla::match;
 class GpuConvRewriterTest : public HloTestBase {
  public:
   GpuConvRewriterTest()
-      : HloTestBase(/*layout_sensitive=*/true,
-                    /*allow_mixed_precision=*/false) {
+      : HloTestBase(/*verifier_layout_sensitive=*/true,
+                    /*allow_mixed_precision_in_hlo_verifier=*/false) {
     for (int i = 0; i < 2; ++i) {
       WindowDimension* window_dim = default_conv_window_.add_dimensions();
       window_dim->set_size(1);

From bdfcc5062b12d7592614d22ddb0a5cd85e11b176 Mon Sep 17 00:00:00 2001
From: Shawn Wang <shawnw@nvidia.com>
Date: Mon, 20 May 2024 17:12:09 -0700
Subject: [PATCH 443/478] PR #12848: [XLA:GPU] Add force_update function for
 conditional commands

Imported from GitHub PR https://github.com/openxla/xla/pull/12848

Copybara import of the project:

--
9f8daced1e91da4abca328406b8c42cfcc250d20 by Shawn Wang <shawnw@nvidia.com>:

update force_update for conditional commands

--
1ffd8e0f59158bb40ad35f0ed205c48c1287cc18 by Shawn Wang <shawnw@nvidia.com>:

fix

Merging this change closes #12848

PiperOrigin-RevId: 635614435
---
 .../service/gpu/runtime/command_buffer_cmd.cc   | 17 +++++++++++++++++
 .../service/gpu/runtime/command_buffer_cmd.h    | 10 ++++++++++
 2 files changed, 27 insertions(+)

diff --git a/third_party/xla/xla/service/gpu/runtime/command_buffer_cmd.cc b/third_party/xla/xla/service/gpu/runtime/command_buffer_cmd.cc
index 75a5d704c33bff..0d4305db9cc8a5 100644
--- a/third_party/xla/xla/service/gpu/runtime/command_buffer_cmd.cc
+++ b/third_party/xla/xla/service/gpu/runtime/command_buffer_cmd.cc
@@ -856,6 +856,8 @@ absl::Status IfCmd::Record(const Thunk::ExecuteParams& execute_params,
       CreateBuilder(&then_commands_, &execute_params, &record_params));
 }
 
+bool IfCmd::force_update() { return then_commands_.force_update(); }
+
 CommandBufferCmd::BufferUsageVector IfCmd::buffers() {
   absl::flat_hash_set<CommandBufferCmd::BufferUsage> buffers;
   buffers.emplace(pred_, MemoryAccess::kRead);
@@ -900,6 +902,10 @@ absl::Status IfElseCmd::Record(const Thunk::ExecuteParams& execute_params,
       CreateBuilder(&else_commands_, &execute_params, &record_params));
 }
 
+bool IfElseCmd::force_update() {
+  return (then_commands_.force_update() || else_commands_.force_update());
+}
+
 CommandBufferCmd::BufferUsageVector IfElseCmd::buffers() {
   absl::flat_hash_set<CommandBufferCmd::BufferUsage> buffers;
   buffers.emplace(pred_, MemoryAccess::kRead);
@@ -945,6 +951,11 @@ absl::Status CaseCmd::Record(const Thunk::ExecuteParams& execute_params,
                                              &execute_params, &record_params));
 }
 
+bool CaseCmd::force_update() {
+  return absl::c_any_of(branches_commands_,
+                        [](const auto& seq) { return seq.force_update(); });
+}
+
 CommandBufferCmd::BufferUsageVector CaseCmd::buffers() {
   absl::flat_hash_set<CommandBufferCmd::BufferUsage> buffers;
   buffers.emplace(index_, MemoryAccess::kRead);
@@ -990,6 +1001,8 @@ absl::Status ForCmd::Record(const Thunk::ExecuteParams& execute_params,
       CreateBuilder(&body_commands_, &execute_params, &record_params));
 }
 
+bool ForCmd::force_update() { return body_commands_.force_update(); }
+
 CommandBufferCmd::BufferUsageVector ForCmd::buffers() {
   absl::flat_hash_set<CommandBufferCmd::BufferUsage> buffers;
   buffers.emplace(loop_counter_, MemoryAccess::kWrite);
@@ -1036,6 +1049,10 @@ absl::Status WhileCmd::Record(const Thunk::ExecuteParams& execute_params,
       CreateBuilder(&body_commands_, &execute_params, &record_params));
 }
 
+bool WhileCmd::force_update() {
+  return (cond_commands_.force_update() || body_commands_.force_update());
+}
+
 CommandBufferCmd::BufferUsageVector WhileCmd::buffers() {
   absl::flat_hash_set<CommandBufferCmd::BufferUsage> buffers;
   buffers.emplace(pred_, MemoryAccess::kWrite);
diff --git a/third_party/xla/xla/service/gpu/runtime/command_buffer_cmd.h b/third_party/xla/xla/service/gpu/runtime/command_buffer_cmd.h
index 9cdecc75531a90..0000eb0c4a221f 100644
--- a/third_party/xla/xla/service/gpu/runtime/command_buffer_cmd.h
+++ b/third_party/xla/xla/service/gpu/runtime/command_buffer_cmd.h
@@ -586,6 +586,8 @@ class IfCmd : public CommandBufferCmd {
                       const RecordParams& record_params,
                       se::CommandBuffer* command_buffer) override;
 
+  bool force_update() override;
+
   BufferUsageVector buffers() override;
 
  private:
@@ -610,6 +612,8 @@ class IfElseCmd : public CommandBufferCmd {
                       const RecordParams& record_params,
                       se::CommandBuffer* command_buffer) override;
 
+  bool force_update() override;
+
   BufferUsageVector buffers() override;
 
  private:
@@ -634,6 +638,8 @@ class CaseCmd : public CommandBufferCmd {
                       const RecordParams& record_params,
                       se::CommandBuffer* command_buffer) override;
 
+  bool force_update() override;
+
   BufferUsageVector buffers() override;
 
  private:
@@ -658,6 +664,8 @@ class ForCmd : public CommandBufferCmd {
                       const RecordParams& record_params,
                       se::CommandBuffer* command_buffer) override;
 
+  bool force_update() override;
+
   BufferUsageVector buffers() override;
 
  private:
@@ -683,6 +691,8 @@ class WhileCmd : public CommandBufferCmd {
                       const RecordParams& record_params,
                       se::CommandBuffer* command_buffer) override;
 
+  bool force_update() override;
+
   BufferUsageVector buffers() override;
 
  private:

From 32e5b6ff793d5cbff41a99b85887b6058cb434fb Mon Sep 17 00:00:00 2001
From: Kuy Mainwaring <kuym@google.com>
Date: Mon, 20 May 2024 17:27:23 -0700
Subject: [PATCH 444/478] [XLA:GPU] Clang-tidy cleanup for
 xla/service/gpu/cudnn_workspace_rewriter.cc

PiperOrigin-RevId: 635617583
---
 third_party/xla/xla/service/gpu/BUILD           |  4 +++-
 .../xla/service/gpu/cudnn_workspace_rewriter.cc | 17 +++++++----------
 2 files changed, 10 insertions(+), 11 deletions(-)

diff --git a/third_party/xla/xla/service/gpu/BUILD b/third_party/xla/xla/service/gpu/BUILD
index 60169c2a9a9907..bf2d24bc389c47 100644
--- a/third_party/xla/xla/service/gpu/BUILD
+++ b/third_party/xla/xla/service/gpu/BUILD
@@ -2995,6 +2995,8 @@ cc_library(
         ":stream_executor_util",
         "@com_google_absl//absl/log",
         "@com_google_absl//absl/log:check",
+        "@com_google_absl//absl/container:flat_hash_set",
+        "@com_google_absl//absl/container:inlined_vector",
         "@com_google_absl//absl/status",
         "@com_google_absl//absl/status:statusor",
         "@com_google_absl//absl/strings:string_view",
@@ -3008,7 +3010,7 @@ cc_library(
         "//xla/stream_executor/cuda:cudnn_plugin",
         "@local_tsl//tsl/platform:errors",
         "@local_tsl//tsl/platform:statusor",
-        "@com_google_absl//absl/container:flat_hash_set",
+        "//xla:status_macros",
     ]),
 )
 
diff --git a/third_party/xla/xla/service/gpu/cudnn_workspace_rewriter.cc b/third_party/xla/xla/service/gpu/cudnn_workspace_rewriter.cc
index 8e37ea1be3416b..55c2c4f913cf8f 100644
--- a/third_party/xla/xla/service/gpu/cudnn_workspace_rewriter.cc
+++ b/third_party/xla/xla/service/gpu/cudnn_workspace_rewriter.cc
@@ -16,15 +16,16 @@ limitations under the License.
 #include "xla/service/gpu/cudnn_workspace_rewriter.h"
 
 #include <optional>
-#include <string>
+#include <utility>
 #include <vector>
 
+#include "absl/container/flat_hash_set.h"
+#include "absl/container/inlined_vector.h"
 #include "absl/log/check.h"
 #include "absl/log/log.h"
 #include "absl/status/status.h"
 #include "absl/status/statusor.h"
 #include "absl/strings/string_view.h"
-#include "third_party/gpus/cudnn/cudnn_version.h"
 #include "xla/hlo/ir/dfs_hlo_visitor_with_default.h"
 #include "xla/hlo/ir/hlo_casting_utils.h"
 #include "xla/hlo/ir/hlo_clone_context.h"
@@ -32,17 +33,16 @@ limitations under the License.
 #include "xla/hlo/ir/hlo_instruction.h"
 #include "xla/hlo/ir/hlo_instructions.h"
 #include "xla/hlo/ir/hlo_module.h"
-#include "xla/hlo/ir/hlo_opcode.h"
-#include "xla/primitive_util.h"
 #include "xla/service/gpu/backend_configs.pb.h"
 #include "xla/service/gpu/cublas_cudnn.h"
 #include "xla/service/gpu/gpu_fused_mha_runner.h"
-#include "xla/service/gpu/ir_emission_utils.h"
 #include "xla/service/gpu/stream_executor_util.h"
+#include "xla/shape.h"
+#include "xla/shape_util.h"
+#include "xla/status_macros.h"
 #include "xla/stream_executor/cuda/cuda_dnn.h"
-#include "xla/stream_executor/cuda/cudnn_frontend_helpers.h"
+#include "xla/stream_executor/dnn.h"
 #include "xla/util.h"
-#include "tsl/platform/errors.h"
 #include "tsl/platform/statusor.h"
 
 namespace xla {
@@ -50,9 +50,6 @@ namespace gpu {
 
 namespace {
 
-namespace fe = cudnn_frontend;
-namespace graph = fe::graph;
-
 // create cuDNN graphs from HloCustomCall
 absl::StatusOr<se::gpu::CudnnGraph> HloCustomCallToCuDnnGraph(
     se::dnn::DnnSupport& dnn_support,

From 97a794b0199b6e507e3f577861856ff09357a3cd Mon Sep 17 00:00:00 2001
From: Farzin Houshmand <farzinh@google.com>
Date: Mon, 20 May 2024 19:22:14 -0700
Subject: [PATCH 445/478] Introduce MakeConstantWithShape in hlo_creation_util
 file.

PiperOrigin-RevId: 635641775
---
 third_party/xla/xla/service/BUILD             |  2 ++
 .../xla/xla/service/hlo_creation_utils.cc     | 17 +++++++++++++++
 .../xla/xla/service/hlo_creation_utils.h      |  4 ++++
 .../xla/xla/service/while_loop_unroller.cc    | 21 +++----------------
 4 files changed, 26 insertions(+), 18 deletions(-)

diff --git a/third_party/xla/xla/service/BUILD b/third_party/xla/xla/service/BUILD
index 9b91d9ec8a80fc..fa49f73567299d 100644
--- a/third_party/xla/xla/service/BUILD
+++ b/third_party/xla/xla/service/BUILD
@@ -2033,6 +2033,7 @@ cc_library(
         "//xla/client/lib:comparators",
         "//xla/hlo/ir:hlo",
         "@com_google_absl//absl/algorithm:container",
+        "@com_google_absl//absl/log",
         "@com_google_absl//absl/log:check",
         "@com_google_absl//absl/strings",
         "@com_google_absl//absl/types:span",
@@ -3247,6 +3248,7 @@ cc_library(
         ":call_inliner",
         ":collective_ops_utils",
         ":flatten_call_graph",
+        ":hlo_creation_utils",
         ":hlo_cse",
         ":hlo_pass",
         ":pattern_matcher",
diff --git a/third_party/xla/xla/service/hlo_creation_utils.cc b/third_party/xla/xla/service/hlo_creation_utils.cc
index b6d735b0f1fb76..3ed1087be0c2f8 100644
--- a/third_party/xla/xla/service/hlo_creation_utils.cc
+++ b/third_party/xla/xla/service/hlo_creation_utils.cc
@@ -26,6 +26,7 @@ limitations under the License.
 
 #include "absl/algorithm/container.h"
 #include "absl/log/check.h"
+#include "absl/log/log.h"
 #include "absl/strings/str_cat.h"
 #include "absl/strings/string_view.h"
 #include "absl/types/span.h"
@@ -869,4 +870,20 @@ HloInstruction* ExpandDegenerateReshape(HloInstruction* inst) {
   return nullptr;
 }
 
+std::unique_ptr<HloInstruction> MakeConstantWithShape(const Shape& shape,
+                                                      int64_t value) {
+  return primitive_util::PrimitiveTypeSwitch<std::unique_ptr<HloInstruction>>(
+      [&](auto literal_constant) -> std::unique_ptr<HloInstruction> {
+        if constexpr (primitive_util::IsIntegralType(literal_constant)) {
+          using NativeT = primitive_util::NativeTypeOf<literal_constant>;
+          auto constant = HloInstruction::CreateConstant(
+              LiteralUtil::CreateR0(static_cast<NativeT>(value)));
+          *constant->mutable_shape() = shape;
+          return std::move(constant);
+        }
+        LOG(FATAL) << "Literal is of non-integral type";
+      },
+      shape.element_type());
+}
+
 }  // namespace xla
diff --git a/third_party/xla/xla/service/hlo_creation_utils.h b/third_party/xla/xla/service/hlo_creation_utils.h
index 95313b9396e2f2..8ca3ac5ed716e0 100644
--- a/third_party/xla/xla/service/hlo_creation_utils.h
+++ b/third_party/xla/xla/service/hlo_creation_utils.h
@@ -390,6 +390,10 @@ absl::StatusOr<std::unique_ptr<HloComputation>> CreateComputationWithSignature(
 // adding and removing reshapes that changes only a single dimension.
 HloInstruction* ExpandDegenerateReshape(HloInstruction* inst);
 
+// Creates an integral constant with the given shape and integer value.
+std::unique_ptr<HloInstruction> MakeConstantWithShape(const Shape& shape,
+                                                      int64_t value);
+
 }  // namespace xla
 
 #endif  // XLA_SERVICE_HLO_CREATION_UTILS_H_
diff --git a/third_party/xla/xla/service/while_loop_unroller.cc b/third_party/xla/xla/service/while_loop_unroller.cc
index 841cac903c9b9e..179ea488b47d22 100644
--- a/third_party/xla/xla/service/while_loop_unroller.cc
+++ b/third_party/xla/xla/service/while_loop_unroller.cc
@@ -45,6 +45,7 @@ limitations under the License.
 #include "xla/service/call_inliner.h"
 #include "xla/service/collective_ops_utils.h"
 #include "xla/service/flatten_call_graph.h"
+#include "xla/service/hlo_creation_utils.h"
 #include "xla/service/hlo_cse.h"
 #include "xla/service/hlo_pass_fix.h"
 #include "xla/service/pattern_matcher.h"
@@ -67,22 +68,6 @@ const int kUnrollTripCountThreshold = 64;
 const int kUnrollInstructionCountThreshold = 800;
 const int kUnrollExpandFactorThreshold = 10000;
 
-std::unique_ptr<HloInstruction> GetConstantWithShape(const Shape& shape,
-                                                     int64_t value) {
-  return primitive_util::PrimitiveTypeSwitch<std::unique_ptr<HloInstruction>>(
-      [&](auto literal_constant) -> std::unique_ptr<HloInstruction> {
-        if constexpr (primitive_util::IsIntegralType(literal_constant)) {
-          using NativeT = primitive_util::NativeTypeOf<literal_constant>;
-          auto constant = HloInstruction::CreateConstant(
-              LiteralUtil::CreateR0(static_cast<NativeT>(value)));
-          *constant->mutable_shape() = shape;
-          return std::move(constant);
-        }
-        LOG(FATAL) << "literal is of non-integral type";
-      },
-      shape.element_type());
-}
-
 // Helper function to create a condition for a single iteration while loop in
 // the form of 'i <= init_value' where i is the induction variable.
 std::unique_ptr<HloComputation> MakeTrivialLoopCondition(
@@ -99,7 +84,7 @@ std::unique_ptr<HloComputation> MakeTrivialLoopCondition(
           param_instruction.value(), induction_idx));
 
   HloInstruction* init_value_constant = condition_builder.AddInstruction(
-      GetConstantWithShape(indvar_instruction->shape(), init_value));
+      MakeConstantWithShape(indvar_instruction->shape(), init_value));
 
   return condition_builder.Build(
       condition_builder.AddInstruction(HloInstruction::CreateCompare(
@@ -155,7 +140,7 @@ UnrollSingleIterationOfTrivialLoop(HloInstruction* while_op,
     }
 
     HloInstruction* induction_value_constant = while_body_clone->AddInstruction(
-        GetConstantWithShape(induction_var_hlo->shape(), induction_value));
+        MakeConstantWithShape(induction_var_hlo->shape(), induction_value));
 
     // Finds all the uses of induction var within the while body and replace it
     // with the constant.

From 0cefd08f5b91e0d0678a4d26b6a688b0e973b022 Mon Sep 17 00:00:00 2001
From: Shawn Wang <shawnw@nvidia.com>
Date: Mon, 20 May 2024 19:42:32 -0700
Subject: [PATCH 446/478] PR #12845: [XLA:GPU] add workspace rewrite for FP8
 gemm

Imported from GitHub PR https://github.com/openxla/xla/pull/12845

Add  workspace rewrite for FP8 with amax fusion.
Copybara import of the project:

--
a2b7f41caff2924798021d130566379e65984750 by Shawn Wang <shawnw@nvidia.com>:

add workspace rewrite for FP8 gemm

Merging this change closes #12845

PiperOrigin-RevId: 635646533
---
 .../xla/xla/service/gpu/backend_configs.proto |  1 +
 .../xla/xla/service/gpu/gemm_rewriter.cc      | 15 ++++++++++----
 .../xla/service/gpu/ir_emitter_unnested.cc    | 15 ++++++--------
 .../service/gpu/tests/gemm_rewrite_test.cc    | 20 +++++++++----------
 4 files changed, 28 insertions(+), 23 deletions(-)

diff --git a/third_party/xla/xla/service/gpu/backend_configs.proto b/third_party/xla/xla/service/gpu/backend_configs.proto
index 7f225008337be8..18264a76795cb6 100644
--- a/third_party/xla/xla/service/gpu/backend_configs.proto
+++ b/third_party/xla/xla/service/gpu/backend_configs.proto
@@ -104,6 +104,7 @@ message GemmBackendConfig {
 
   optional bool grad_x = 16;
   optional bool grad_y = 17;
+  bool damax_output = 18;
 }
 
 // Backend config for bitcast operation generated from MLIR MHLO dialect.
diff --git a/third_party/xla/xla/service/gpu/gemm_rewriter.cc b/third_party/xla/xla/service/gpu/gemm_rewriter.cc
index dfb72f2dc0cba9..df4b9696c54870 100644
--- a/third_party/xla/xla/service/gpu/gemm_rewriter.cc
+++ b/third_party/xla/xla/service/gpu/gemm_rewriter.cc
@@ -1391,6 +1391,12 @@ class GemmRewriterVisitor : public DfsHloRewriteVisitor {
     HloInstruction *gemm_and_damax =
         instr->AddInstruction(existing_gemm->CloneWithNewShape(tuple_shape));
 
+    TF_ASSIGN_OR_RETURN(auto gpu_config,
+                        gemm_and_damax->backend_config<GpuBackendConfig>());
+    GemmBackendConfig &config = *gpu_config.mutable_gemm_backend_config();
+    config.set_damax_output(true);
+    TF_RETURN_IF_ERROR(gemm_and_damax->set_backend_config(gpu_config));
+
     // Obtain D and DAmax separately from the output tuple.
     HloInstruction *d =
         instr->AddInstruction(HloInstruction::CreateGetTupleElement(
@@ -2262,9 +2268,10 @@ class GemmWorkspaceRewriteVisitor : public DfsHloRewriteVisitor {
           has_aux_output,
           xla::gpu::gpublas_lt::EpilogueHasAuxiliaryOutput(epilogue));
 
-      if (!((has_aux_output && instr->shape().IsTuple() &&
-             instr->shape().tuple_shapes_size() == 2) ||
-            (!has_aux_output && instr->shape().IsArray()))) {
+      if (!((instr->shape().IsTuple() &&
+             instr->shape().tuple_shapes_size() ==
+                 has_aux_output + config.damax_output() + 1) ||
+            instr->shape().IsArray())) {
         return absl::OkStatus();
       }
     } else if (instr->custom_call_target() != kGemmCallTarget ||
@@ -2314,7 +2321,7 @@ class GemmWorkspaceRewriteVisitor : public DfsHloRewriteVisitor {
       custom_call->set_output_to_operand_aliasing({{{0}, {2, {}}}});
     }
 
-    if (has_aux_output) {
+    if (instr->shape().IsTuple()) {
       for (auto user : instr->users()) {
         auto user_get_tuple =
             dynamic_cast<HloGetTupleElementInstruction *>(user);
diff --git a/third_party/xla/xla/service/gpu/ir_emitter_unnested.cc b/third_party/xla/xla/service/gpu/ir_emitter_unnested.cc
index 99f469752da5f8..c6ac156b1a4cc2 100644
--- a/third_party/xla/xla/service/gpu/ir_emitter_unnested.cc
+++ b/third_party/xla/xla/service/gpu/ir_emitter_unnested.cc
@@ -29,6 +29,7 @@ limitations under the License.
 #include <variant>
 #include <vector>
 
+#include "mlir/AsmParser/AsmParser.h"  // from @llvm-project
 #include "absl/container/flat_hash_map.h"
 #include "absl/container/flat_hash_set.h"
 #include "absl/container/inlined_vector.h"
@@ -55,7 +56,6 @@ limitations under the License.
 #include "llvm/IR/LLVMContext.h"
 #include "llvm/IR/Module.h"
 #include "llvm/Linker/Linker.h"
-#include "mlir/AsmParser/AsmParser.h"  // from @llvm-project
 #include "mlir/Dialect/LLVMIR/LLVMDialect.h"  // from @llvm-project
 #include "mlir/Dialect/MemRef/Transforms/Passes.h"  // from @llvm-project
 #include "mlir/IR/Attributes.h"  // from @llvm-project
@@ -758,9 +758,9 @@ absl::Status IrEmitterUnnested::EmitCublasLtMatmulThunkF8(
 
   TF_ASSIGN_OR_RETURN(bool has_vector_bias,
                       xla::gpu::gpublas_lt::EpilogueAddsVectorBias(epilogue));
-  bool has_damax = instr->shape().IsTuple();
-  xla::ShapeIndex output_index =
-      has_damax ? xla::ShapeIndex{0} : xla::ShapeIndex{};
+
+  TF_RET_CHECK(instr->shape().IsTuple());
+  xla::ShapeIndex output_index = xla::ShapeIndex{0};
 
   TF_ASSIGN_OR_RETURN(
       bool has_aux_output,
@@ -805,7 +805,7 @@ absl::Status IrEmitterUnnested::EmitCublasLtMatmulThunkF8(
   }
 
   BufferAllocation::Slice d_amax;
-  if (has_damax) {
+  if (config.damax_output()) {
     TF_ASSIGN_OR_RETURN(d_amax, GetAllocationSliceForHlo(instr, {1}));
   }
 
@@ -822,10 +822,7 @@ absl::Status IrEmitterUnnested::EmitCublasLtMatmulThunkF8(
   BufferAllocation::Slice aux;  // Not used.
   TF_RET_CHECK(!has_aux_output);
   std::optional<BufferAllocation::Slice> workspace_buffer;
-  if (instr->shape().IsTuple() &&
-      (instr->shape().tuple_shapes_size() - has_aux_output - 1)) {
-    TF_RET_CHECK((has_aux_output && instr->shape().tuple_shapes_size() == 3) ||
-                 (!has_aux_output && instr->shape().tuple_shapes_size() == 2));
+  if (instr->shape().tuple_shapes_size() - config.damax_output() == 2) {
     TF_ASSIGN_OR_RETURN(workspace_buffer,
                         GetAllocationSliceForHlo(
                             instr, {instr->shape().tuple_shapes_size() - 1}));
diff --git a/third_party/xla/xla/service/gpu/tests/gemm_rewrite_test.cc b/third_party/xla/xla/service/gpu/tests/gemm_rewrite_test.cc
index 4dd99c541a1477..24f95bd68af283 100644
--- a/third_party/xla/xla/service/gpu/tests/gemm_rewrite_test.cc
+++ b/third_party/xla/xla/service/gpu/tests/gemm_rewrite_test.cc
@@ -5654,8 +5654,8 @@ TEST_P(ParameterizedFp8GemmRewriteTest,
 ; CHECK-NEXT:    [[XS1:%[^ ]+]] = f32[] convert([[P3]])
 ; CHECK-NEXT:    [[C1:%[^ ]+]] = f32[] constant(1)
 ; CHECK-PTX-NEXT:    [[B:%[^ ]+]] = bf16[16]{0} parameter(4)
-; CHECK-PTX-NEXT:    ROOT [[OUT:%[^ ]+]] = bf16[16,16]{1,0} custom-call([[P0]], [[P1_TRANSPOSE]], [[XS]], [[XS1]], [[C1]], /*index=5*/[[C1]], [[B]]),
-; CHECK-GCN-NEXT:    [[OUT:%[^ ]+]] = f32[16,16]{1,0} custom-call([[P0]], [[P1_TRANSPOSE]], [[XS]], [[XS1]], [[C1]], /*index=5*/[[C1]]),
+; CHECK-PTX-NEXT:    [[OUT:%[^ ]+]] = (bf16[16,16]{1,0}, s8[{{[0-9]+}}]{0}) custom-call([[P0]], [[P1_TRANSPOSE]], [[XS]], [[XS1]], [[C1]], /*index=5*/[[C1]], [[B]]),
+; CHECK-GCN-NEXT:    [[OUT:%[^ ]+]] = (f32[16,16]{1,0}, s8[{{[0-9]+}}]{0}) custom-call([[P0]], [[P1_TRANSPOSE]], [[XS]], [[XS1]], [[C1]], /*index=5*/[[C1]]),
 ; CHECK:           custom_call_target="__cublas$lt$matmul$f8",
 ; CHECK:           backend_config={
 ; CHECK-DAG:         "alpha_real":1
@@ -5742,8 +5742,8 @@ TEST_P(ParameterizedFp8GemmRewriteTest,
 ; CHECK-NEXT:    [[P3:%[^ ]+]] = bf16[] parameter(3)
 ; CHECK-NEXT:    [[XS1:%[^ ]+]] = f32[] convert([[P3]])
 ; CHECK-NEXT:    [[C1:%[^ ]+]] = f32[] constant(1)
-; CHECK-PTX-NEXT:    ROOT [[OUT:%[^ ]+]] = bf16[16,16]{1,0} custom-call([[P0]], [[P1_TRANSPOSE]], [[XS]], [[XS1]], [[C1]], /*index=5*/[[C1]]),
-; CHECK-GCN-NEXT:    [[OUT:%[^ ]+]] = f32[16,16]{1,0} custom-call([[P0]], [[P1_TRANSPOSE]], [[XS]], [[XS1]], [[C1]], /*index=5*/[[C1]]),
+; CHECK-PTX-NEXT:    [[OUT:%[^ ]+]] = (bf16[16,16]{1,0}, s8[{{[0-9]+}}]{0}) custom-call([[P0]], [[P1_TRANSPOSE]], [[XS]], [[XS1]], [[C1]], /*index=5*/[[C1]]),
+; CHECK-GCN-NEXT:    [[OUT:%[^ ]+]] = (f32[16,16]{1,0}, s8[{{[0-9]+}}]{0}) custom-call([[P0]], [[P1_TRANSPOSE]], [[XS]], [[XS1]], [[C1]], /*index=5*/[[C1]]),
 ; CHECK:           custom_call_target="__cublas$lt$matmul$f8",
 ; CHECK:           backend_config={
 ; CHECK-DAG:         "alpha_real":1
@@ -7443,8 +7443,8 @@ TEST_P(ParameterizedFp8GemmRewriteTest, ScaledABScaledDWithDAmaxF8) {
 ; CHECK-PTX-NEXT:    [[C2:%[^ ]+]] = f32[] constant(1)
 ; CHECK-PTX-NEXT:    [[P4:%[^ ]+]] = f32[] parameter(4)
 ; CHECK-PTX-NEXT:    [[P4_INV:%[^ ]+]] = f32[] divide([[C2]], [[P4]])
-; CHECK-PTX-NEXT:    [[OUT:%[^ ]+]] = (<<F8E4M3>>[16,16]{1,0}, f32[]) custom-call([[P0]], [[P1_TRANSPOSE]], [[P2]], [[P3]], [[C1]], /*index=5*/[[P4_INV]]),
-; CHECK-GCN-NEXT:    [[OUT:%[^ ]+]] = f32[16,16]{1,0} custom-call([[P0]], [[P1_TRANSPOSE]], [[P2]], [[P3]], [[C1]], /*index=5*/[[C1]]),
+; CHECK-PTX-NEXT:    [[OUT:%[^ ]+]] = (<<F8E4M3>>[16,16]{1,0}, f32[], s8[{{[0-9]+}}]{0}) custom-call([[P0]], [[P1_TRANSPOSE]], [[P2]], [[P3]], [[C1]], /*index=5*/[[P4_INV]]),
+; CHECK-GCN-NEXT:    [[OUT:%[^ ]+]] = (f32[16,16]{1,0}, s8[{{[0-9]+}}]{0}) custom-call([[P0]], [[P1_TRANSPOSE]], [[P2]], [[P3]], [[C1]], /*index=5*/[[C1]]),
 ; CHECK:           custom_call_target="__cublas$lt$matmul$f8",
 ; CHECK:           backend_config={
 ; CHECK-DAG:         "alpha_real":1
@@ -7533,8 +7533,8 @@ TEST_P(ParameterizedFp8GemmRewriteTest,
 ; CHECK-PTX-NEXT:    [[P4:%[^ ]+]] = f16[] parameter(4)
 ; CHECK-PTX-NEXT:    [[P4_INV:%[^ ]+]] = f16[] divide([[C2]], [[P4]])
 ; CHECK-PTX-NEXT:    [[P4_INV_CONVERT:%[^ ]+]] = f32[] convert([[P4_INV]])
-; CHECK-PTX-NEXT:    [[OUT:%[^ ]+]] = (<<F8E4M3>>[16,16]{1,0}, f32[]) custom-call([[P0]], [[P1_TRANSPOSE]], [[P2_CONVERT]], [[P3_CONVERT]], [[C1]], /*index=5*/[[P4_INV_CONVERT]]),
-; CHECK-GCN-NEXT:    [[OUT:%[^ ]+]] = f16[16,16]{1,0} custom-call([[P0]], [[P1_TRANSPOSE]], [[P2_CONVERT]], [[P3_CONVERT]], [[C1]], /*index=5*/[[C1]]),
+; CHECK-PTX-NEXT:    [[OUT:%[^ ]+]] = (<<F8E4M3>>[16,16]{1,0}, f32[], s8[{{[0-9]+}}]{0}) custom-call([[P0]], [[P1_TRANSPOSE]], [[P2_CONVERT]], [[P3_CONVERT]], [[C1]], /*index=5*/[[P4_INV_CONVERT]]),
+; CHECK-GCN-NEXT:    [[OUT:%[^ ]+]] = (f16[16,16]{1,0}, s8[{{[0-9]+}}]{0}) custom-call([[P0]], [[P1_TRANSPOSE]], [[P2_CONVERT]], [[P3_CONVERT]], [[C1]], /*index=5*/[[C1]]),
 ; CHECK:           custom_call_target="__cublas$lt$matmul$f8",
 ; CHECK:           backend_config={
 ; CHECK-DAG:         "alpha_real":1
@@ -7620,8 +7620,8 @@ TEST_P(ParameterizedFp8GemmRewriteTest,
 ; CHECK-PTX-NEXT:    [[C2:%[^ ]+]] = f32[] constant(1)
 ; CHECK-PTX-NEXT:    [[P4:%[^ ]+]] = f32[] parameter(4)
 ; CHECK-PTX-NEXT:    [[P4_INV:%[^ ]+]] = f32[] divide([[C2]], [[P4]])
-; CHECK-PTX-NEXT:    [[OUT:%[^ ]+]] = (<<F8E4M3>>[16,16]{1,0}, f32[]) custom-call([[P0]], [[P1_TRANSPOSE]], [[P2]], [[P3]], [[C1]], /*index=5*/[[P4_INV]]),
-; CHECK-GCN-NEXT:    [[OUT:%[^ ]+]] = f32[16,16]{1,0} custom-call([[P0]], [[P1_TRANSPOSE]], [[P2]], [[P3]], [[C1]], /*index=5*/[[C1]]),
+; CHECK-PTX-NEXT:    [[OUT:%[^ ]+]] = (<<F8E4M3>>[16,16]{1,0}, f32[], s8[{{[0-9]+}}]{0}) custom-call([[P0]], [[P1_TRANSPOSE]], [[P2]], [[P3]], [[C1]], /*index=5*/[[P4_INV]]),
+; CHECK-GCN-NEXT:    [[OUT:%[^ ]+]] = (f32[16,16]{1,0}, s8[{{[0-9]+}}]{0}) custom-call([[P0]], [[P1_TRANSPOSE]], [[P2]], [[P3]], [[C1]], /*index=5*/[[C1]]),
 ; CHECK:           custom_call_target="__cublas$lt$matmul$f8",
 ; CHECK:           backend_config={
 ; CHECK-DAG:         "alpha_real":1

From a51ee7da2ff54d595bf568ec5976f86899a68eca Mon Sep 17 00:00:00 2001
From: Johannes Reifferscheid <jreiffers@google.com>
Date: Tue, 21 May 2024 00:28:15 -0700
Subject: [PATCH 447/478] Internal change only.

PiperOrigin-RevId: 635704825
---
 third_party/xla/xla/service/gpu/ir_emitter_unnested.cc | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/third_party/xla/xla/service/gpu/ir_emitter_unnested.cc b/third_party/xla/xla/service/gpu/ir_emitter_unnested.cc
index c6ac156b1a4cc2..8d3015fed82aef 100644
--- a/third_party/xla/xla/service/gpu/ir_emitter_unnested.cc
+++ b/third_party/xla/xla/service/gpu/ir_emitter_unnested.cc
@@ -29,7 +29,6 @@ limitations under the License.
 #include <variant>
 #include <vector>
 
-#include "mlir/AsmParser/AsmParser.h"  // from @llvm-project
 #include "absl/container/flat_hash_map.h"
 #include "absl/container/flat_hash_set.h"
 #include "absl/container/inlined_vector.h"
@@ -56,6 +55,7 @@ limitations under the License.
 #include "llvm/IR/LLVMContext.h"
 #include "llvm/IR/Module.h"
 #include "llvm/Linker/Linker.h"
+#include "mlir/AsmParser/AsmParser.h"  // from @llvm-project
 #include "mlir/Dialect/LLVMIR/LLVMDialect.h"  // from @llvm-project
 #include "mlir/Dialect/MemRef/Transforms/Passes.h"  // from @llvm-project
 #include "mlir/IR/Attributes.h"  // from @llvm-project

From fde16c701672a992c75c9ee19e6cd2da4bc420f7 Mon Sep 17 00:00:00 2001
From: Johannes Reifferscheid <jreiffers@google.com>
Date: Tue, 21 May 2024 00:55:38 -0700
Subject: [PATCH 448/478] Rename Buffer to GetBuffer.

InterpreterValue::Buffer having the same name as Buffer is considered
invalid by GCC.

sed "s/\.Buffer/.GetBuffer/g" -i `find . -name '*.cc'`

PiperOrigin-RevId: 635710161
---
 .../mlir/tools/mlir_interpreter/dialects/arith.cc  |  4 ++--
 .../mlir_interpreter/dialects/bufferization.cc     |  2 +-
 .../tools/mlir_interpreter/dialects/builtin.cc     |  2 +-
 .../mlir/tools/mlir_interpreter/dialects/memref.cc | 14 +++++++-------
 .../mlir/tools/mlir_interpreter/dialects/mhlo.cc   |  2 +-
 .../mlir/tools/mlir_interpreter/dialects/vector.cc |  6 +++---
 .../mlir_interpreter/framework/interpreter.cc      |  6 +++---
 .../framework/interpreter_value.cc                 |  2 +-
 .../mlir_interpreter/framework/interpreter_value.h |  2 +-
 .../mlir_replay/public/execution_trace_utils.cc    |  2 +-
 10 files changed, 21 insertions(+), 21 deletions(-)

diff --git a/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/arith.cc b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/arith.cc
index d48edff6f37634..18b208dc71014b 100644
--- a/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/arith.cc
+++ b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/arith.cc
@@ -44,9 +44,9 @@ InterpreterValue Bitcast(InterpreterState&, arith::BitcastOp op,
     TensorOrMemref<decltype(dummy)> result;
     result.view = {};
     if (shaped_ty) {
-      result.buffer = in.Clone().Buffer();
+      result.buffer = in.Clone().GetBuffer();
     } else {
-      result.buffer = in.AsUnitTensor().Buffer();
+      result.buffer = in.AsUnitTensor().GetBuffer();
     }
     return {result};
   });
diff --git a/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/bufferization.cc b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/bufferization.cc
index 6bd9239876e22c..950b8eff988608 100644
--- a/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/bufferization.cc
+++ b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/bufferization.cc
@@ -57,7 +57,7 @@ InterpreterValue AllocTensor(
 InterpreterValue Clone(InterpreterState& state, bufferization::CloneOp,
                        const InterpreterValue& in) {
   if (auto* stats = state.GetOptions().stats) {
-    stats->heap_size += in.Buffer()->GetByteSize();
+    stats->heap_size += in.GetBuffer()->GetByteSize();
     stats->peak_heap_size = std::max(stats->peak_heap_size, stats->heap_size);
     ++stats->num_allocations;
   }
diff --git a/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/builtin.cc b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/builtin.cc
index 48b8008dc07cfc..ccd1546e59d11c 100644
--- a/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/builtin.cc
+++ b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/builtin.cc
@@ -45,7 +45,7 @@ llvm::SmallVector<InterpreterValue> UnrealizedConversionCast(
         return {DispatchScalarType(r, [&](auto dummy) -> InterpreterValue {
           TensorOrMemref<decltype(dummy)> result;
           result.view = args[0].View();
-          result.buffer = args[0].Buffer();
+          result.buffer = args[0].GetBuffer();
           return {result};
         })};
       }
diff --git a/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/memref.cc b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/memref.cc
index 64b28342f88707..040e67864049da 100644
--- a/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/memref.cc
+++ b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/memref.cc
@@ -40,7 +40,7 @@ namespace {
 InterpreterValue Load(InterpreterState& state, memref::LoadOp,
                       const InterpreterValue& memref,
                       ArrayRef<int64_t> indices) {
-  if (!memref.Buffer()) {
+  if (!memref.GetBuffer()) {
     state.AddFailure("null pointer dereference.");
     return {};
   }
@@ -69,11 +69,11 @@ InterpreterValue Alloc(InterpreterState& state, memref::AllocOp alloc,
   auto result =
       InterpreterValue::MakeTensor(ty.getElementType(), std::move(shape));
   if (auto* stats = state.GetOptions().stats) {
-    stats->heap_size += result.Buffer()->GetByteSize();
+    stats->heap_size += result.GetBuffer()->GetByteSize();
     stats->peak_heap_size = std::max(stats->peak_heap_size, stats->heap_size);
     ++stats->num_allocations;
   }
-  result.Buffer()->SetAllocatedBy(alloc);
+  result.GetBuffer()->SetAllocatedBy(alloc);
   return result;
 }
 
@@ -83,18 +83,18 @@ InterpreterValue AllocA(InterpreterState&, memref::AllocaOp alloc,
   auto shape = ReplaceDynamicVals(ty.getShape(), dynamic_sizes);
   auto result =
       InterpreterValue::MakeTensor(ty.getElementType(), std::move(shape));
-  result.Buffer()->SetIsAlloca();
-  result.Buffer()->SetAllocatedBy(alloc);
+  result.GetBuffer()->SetIsAlloca();
+  result.GetBuffer()->SetAllocatedBy(alloc);
   return result;
 }
 
 void Dealloc(InterpreterState& state, memref::DeallocOp op,
              InterpreterValue memref) {
-  if (!memref.Buffer()) {
+  if (!memref.GetBuffer()) {
     state.AddFailure("attempting to deallocate null pointer.");
     return;
   }
-  auto buffer = memref.Buffer();
+  auto buffer = memref.GetBuffer();
   const auto& view = memref.View();
   if (auto* stats = state.GetOptions().stats) {
     stats->heap_size -= buffer->GetByteSize();
diff --git a/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/mhlo.cc b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/mhlo.cc
index bf77f451c84e11..18a6c1f0ce34c0 100644
--- a/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/mhlo.cc
+++ b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/mhlo.cc
@@ -70,7 +70,7 @@ InterpreterValue BitcastConvert(InterpreterState&, mhlo::BitcastConvertOp op,
   auto result = DispatchScalarType(ty, [&](auto dummy) -> InterpreterValue {
     TensorOrMemref<decltype(dummy)> result;
     result.view = {};
-    result.buffer = in.Clone().Buffer();
+    result.buffer = in.Clone().GetBuffer();
     return {result};
   });
   auto& out_view = result.View();
diff --git a/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/vector.cc b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/vector.cc
index 7305b7e1b4e649..f01a4f6ece711a 100644
--- a/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/vector.cc
+++ b/third_party/xla/xla/mlir/tools/mlir_interpreter/dialects/vector.cc
@@ -155,7 +155,7 @@ InterpreterValue Bitcast(InterpreterState&, vector::BitCastOp op,
                          const InterpreterValue& vector) {
   ShapedType ty = cast<ShapedType>(op->getResultTypes()[0]);
   auto flattened = vector.CoerceLayout({});
-  auto buffer = flattened.Buffer();
+  auto buffer = flattened.GetBuffer();
   auto view = flattened.View();
   view.sizes = llvm::to_vector(ty.getShape());
   view.strides = BufferView::GetDefaultStrides(view.sizes);
@@ -190,7 +190,7 @@ InterpreterValue Broadcast(InterpreterState&, vector::BroadcastOp broadcast,
 void CompressStore(InterpreterState& state, vector::CompressStoreOp,
                    InterpreterValue dst, ArrayRef<int64_t> indices,
                    TensorOrMemref<bool> mask, InterpreterValue value) {
-  auto dst_buffer = dst.Buffer();
+  auto dst_buffer = dst.GetBuffer();
   const auto& dst_view = dst.View();
   if (dst_view.strides.back() != 1) {
     state.AddFailure("trailing dimension must be continguous");
@@ -202,7 +202,7 @@ void CompressStore(InterpreterState& state, vector::CompressStoreOp,
     return;
   }
 
-  auto src_buffer = value.Buffer();
+  auto src_buffer = value.GetBuffer();
   const auto& src_view = value.View();
 
   // TODO(jreiffers): Bounds checks.
diff --git a/third_party/xla/xla/mlir/tools/mlir_interpreter/framework/interpreter.cc b/third_party/xla/xla/mlir/tools/mlir_interpreter/framework/interpreter.cc
index 21de4ea6ffc7cf..f5476bf74296ba 100644
--- a/third_party/xla/xla/mlir/tools/mlir_interpreter/framework/interpreter.cc
+++ b/third_party/xla/xla/mlir/tools/mlir_interpreter/framework/interpreter.cc
@@ -124,9 +124,9 @@ void InterpreterState::AddFailure(llvm::StringRef failure) {
 
 void InterpreterScope::Verify() const {
   for (auto& [_, value] : values_) {
-    if (value.IsTensor() && value.Buffer() &&
-        !value.Buffer()->GetFailure().empty()) {
-      state_.AddFailure(value.Buffer()->GetFailure());
+    if (value.IsTensor() && value.GetBuffer() &&
+        !value.GetBuffer()->GetFailure().empty()) {
+      state_.AddFailure(value.GetBuffer()->GetFailure());
       break;
     }
   }
diff --git a/third_party/xla/xla/mlir/tools/mlir_interpreter/framework/interpreter_value.cc b/third_party/xla/xla/mlir/tools/mlir_interpreter/framework/interpreter_value.cc
index cd5c2ac4190308..7c0a3e756c764f 100644
--- a/third_party/xla/xla/mlir/tools/mlir_interpreter/framework/interpreter_value.cc
+++ b/third_party/xla/xla/mlir/tools/mlir_interpreter/framework/interpreter_value.cc
@@ -322,7 +322,7 @@ bool Tuple::operator==(const Tuple& other) const {
   return true;
 }
 
-std::shared_ptr<Buffer> InterpreterValue::Buffer() const {
+std::shared_ptr<Buffer> InterpreterValue::GetBuffer() const {
   return std::visit(
       [](const auto& it) -> std::shared_ptr<interpreter::Buffer> {
         if constexpr (is_tensor_or_memref_v<decltype(it)>) {
diff --git a/third_party/xla/xla/mlir/tools/mlir_interpreter/framework/interpreter_value.h b/third_party/xla/xla/mlir/tools/mlir_interpreter/framework/interpreter_value.h
index 90fb8e74eceb33..69228c4edc7c13 100644
--- a/third_party/xla/xla/mlir/tools/mlir_interpreter/framework/interpreter_value.h
+++ b/third_party/xla/xla/mlir/tools/mlir_interpreter/framework/interpreter_value.h
@@ -99,7 +99,7 @@ struct InterpreterValue {
   BufferView& View();
   const BufferView& View() const;
   // Returns the underlying tensor's buffer. Must be a tensor.
-  std::shared_ptr<Buffer> Buffer() const;
+  std::shared_ptr<Buffer> GetBuffer() const;
 
   bool IsTensor() const;
 
diff --git a/third_party/xla/xla/mlir/tools/mlir_replay/public/execution_trace_utils.cc b/third_party/xla/xla/mlir/tools/mlir_replay/public/execution_trace_utils.cc
index 90357dde29d622..e90b51ed6e0733 100644
--- a/third_party/xla/xla/mlir/tools/mlir_replay/public/execution_trace_utils.cc
+++ b/third_party/xla/xla/mlir/tools/mlir_replay/public/execution_trace_utils.cc
@@ -318,7 +318,7 @@ absl::StatusOr<InterpreterValue> LiteralToValue(
   return {DispatchScalarType(type, [&](auto dummy) -> InterpreterValue {
     TensorOrMemref<decltype(dummy)> cast;
     cast.view = result.View();
-    cast.buffer = result.Buffer();
+    cast.buffer = result.GetBuffer();
     return {cast};
   })};
 }

From f53aae362873d6df132d08b63b495adf644d67d9 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Tue, 21 May 2024 01:06:27 -0700
Subject: [PATCH 449/478] Automated Code Change

PiperOrigin-RevId: 635712904
---
 third_party/xla/xla/service/gpu/BUILD                         | 4 +++-
 third_party/xla/xla/service/gpu/dot_operand_converter_test.cc | 1 +
 third_party/xla/xla/service/gpu/dot_sparsity_rewriter_test.cc | 1 +
 third_party/xla/xla/service/gpu/fusion_merger_test.cc         | 1 +
 4 files changed, 6 insertions(+), 1 deletion(-)

diff --git a/third_party/xla/xla/service/gpu/BUILD b/third_party/xla/xla/service/gpu/BUILD
index bf2d24bc389c47..a739805c376c56 100644
--- a/third_party/xla/xla/service/gpu/BUILD
+++ b/third_party/xla/xla/service/gpu/BUILD
@@ -1832,6 +1832,7 @@ xla_cc_test(
     use_gpu = True,
     deps = [
         ":dot_sparsity_rewriter",
+        "//xla:xla_data_proto_cc",
         "//xla/hlo/ir:hlo",
         "//xla/tests:hlo_test_base",
         "//xla/tests:xla_internal_test_main",  # fixdeps: keep
@@ -2634,6 +2635,7 @@ xla_cc_test(
         ":gpu_device_info_for_tests",
         ":gpu_fusible",
         "//xla:shape_util",
+        "//xla:xla_data_proto_cc",
         "//xla/hlo/ir:hlo",
         "//xla/service:hlo_cost_analysis",
         "//xla/service:pattern_matcher",
@@ -5318,7 +5320,7 @@ xla_test(
             "@local_tsl//tsl/platform:statusor",
         ],
         ["@local_tsl//tsl/platform:test_main"],  # b/317293391
-    ),
+    ) + ["//xla:xla_data_proto_cc"],
 )
 
 cc_library(
diff --git a/third_party/xla/xla/service/gpu/dot_operand_converter_test.cc b/third_party/xla/xla/service/gpu/dot_operand_converter_test.cc
index 9a36a288e01b6b..63b0017012f419 100644
--- a/third_party/xla/xla/service/gpu/dot_operand_converter_test.cc
+++ b/third_party/xla/xla/service/gpu/dot_operand_converter_test.cc
@@ -25,6 +25,7 @@ limitations under the License.
 #include "xla/hlo/utils/hlo_matchers.h"
 #include "xla/primitive_util.h"
 #include "xla/tests/hlo_test_base.h"
+#include "xla/xla_data.pb.h"
 #include "tsl/platform/statusor.h"
 
 namespace xla::gpu {
diff --git a/third_party/xla/xla/service/gpu/dot_sparsity_rewriter_test.cc b/third_party/xla/xla/service/gpu/dot_sparsity_rewriter_test.cc
index 315600f9495e09..c608f8d814410a 100644
--- a/third_party/xla/xla/service/gpu/dot_sparsity_rewriter_test.cc
+++ b/third_party/xla/xla/service/gpu/dot_sparsity_rewriter_test.cc
@@ -23,6 +23,7 @@ limitations under the License.
 #include "xla/hlo/ir/hlo_instruction.h"
 #include "xla/hlo/ir/hlo_instructions.h"
 #include "xla/tests/hlo_test_base.h"
+#include "xla/xla_data.pb.h"
 #include "tsl/platform/statusor.h"
 
 namespace xla {
diff --git a/third_party/xla/xla/service/gpu/fusion_merger_test.cc b/third_party/xla/xla/service/gpu/fusion_merger_test.cc
index 447ffce64998cd..de45a4b9d3273e 100644
--- a/third_party/xla/xla/service/gpu/fusion_merger_test.cc
+++ b/third_party/xla/xla/service/gpu/fusion_merger_test.cc
@@ -32,6 +32,7 @@ limitations under the License.
 #include "xla/shape.h"
 #include "xla/shape_util.h"
 #include "xla/tests/hlo_test_base.h"
+#include "xla/xla_data.pb.h"
 
 namespace xla {
 namespace gpu {

From ea5adb0f08787c01697a4142d1cd2f2ec6361db2 Mon Sep 17 00:00:00 2001
From: Johannes Reifferscheid <jreiffers@google.com>
Date: Tue, 21 May 2024 01:19:06 -0700
Subject: [PATCH 450/478] MLIR emitters: Vectorize transposes with small
 element types.

Special thanks to github user lingzhi98, who proposed and benchmarked
this optimization and implemented it in https://github.com/openxla/xla/pull/12552.
It turned out that the Tiling abstraction made this harder than it needs to be.
This is an independent reimplementation that removes the Tiling abstraction, but
reuses lingzhi98's tests and most of their heuristic for enabling vectorization.

The differences in the heuristic are:
- respect the shared memory budget
- enable vectorization if the dimensions are divisible by 2, not just when
  they are divisible by 64.

PiperOrigin-RevId: 635715722
---
 third_party/xla/xla/service/gpu/fusions/BUILD |   8 +-
 .../gpu/fusions/mlir/elemental_hlo_to_mlir.cc |   7 +
 .../gpu/fusions/mlir/elemental_hlo_to_mlir.h  |   4 +
 .../xla/service/gpu/fusions/reduction_base.cc |   3 +-
 .../xla/service/gpu/fusions/reduction_mlir.cc |   5 +-
 .../xla/xla/service/gpu/fusions/tiling_util.h |  11 +-
 .../xla/service/gpu/fusions/transpose_mlir.cc | 248 ++++++++----------
 .../xla/service/gpu/fusions/transpose_mlir.h  |  21 +-
 .../gpu/fusions/transpose_mlir_test.cc        | 197 ++++++++++++--
 .../service/gpu/model/indexing_analysis.cc    |  10 +-
 .../xla/service/gpu/model/indexing_analysis.h |   3 +-
 11 files changed, 328 insertions(+), 189 deletions(-)

diff --git a/third_party/xla/xla/service/gpu/fusions/BUILD b/third_party/xla/xla/service/gpu/fusions/BUILD
index 6ba5e41a4e9f2e..3796e9a40f3bab 100644
--- a/third_party/xla/xla/service/gpu/fusions/BUILD
+++ b/third_party/xla/xla/service/gpu/fusions/BUILD
@@ -457,7 +457,7 @@ cc_library(
     srcs = ["transpose_mlir.cc"],
     hdrs = ["transpose_mlir.h"],
     deps = [
-        ":tiling_util",
+        ":fusion_emitter",
         "//xla:permutation_util",
         "//xla:shape_util",
         "//xla:status",
@@ -475,16 +475,14 @@ cc_library(
         "//xla/service/gpu/model:indexing_analysis",
         "//xla/service/gpu/model:indexing_map",
         "@com_google_absl//absl/algorithm:container",
-        "@com_google_absl//absl/container:inlined_vector",
-        "@com_google_absl//absl/log",
+        "@com_google_absl//absl/container:flat_hash_map",
         "@com_google_absl//absl/log:check",
-        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/types:span",
         "@llvm-project//llvm:Support",
         "@llvm-project//mlir:FuncDialect",
         "@llvm-project//mlir:IR",
         "@llvm-project//mlir:Support",
         "@llvm-project//mlir:TensorDialect",
-        "@local_tsl//tsl/platform:statusor",
     ],
 )
 
diff --git a/third_party/xla/xla/service/gpu/fusions/mlir/elemental_hlo_to_mlir.cc b/third_party/xla/xla/service/gpu/fusions/mlir/elemental_hlo_to_mlir.cc
index 5433fa6d31f9ef..192997dcf706bf 100644
--- a/third_party/xla/xla/service/gpu/fusions/mlir/elemental_hlo_to_mlir.cc
+++ b/third_party/xla/xla/service/gpu/fusions/mlir/elemental_hlo_to_mlir.cc
@@ -699,6 +699,13 @@ SmallVector<Value> ApplyAffineMap(mlir::AffineMap map, ValueRange dims,
   return result;
 }
 
+SmallVector<Value> ApplyIndexing(const IndexingMap& map, ValueRange dims,
+                                 ValueRange symbols, ImplicitLocOpBuilder& b) {
+  SmallVector<Value> results;
+  b.createOrFold<ApplyIndexingOp>(results, dims, symbols, map);
+  return results;
+}
+
 Value CheckConstraint(mlir::Value constrained_value, Interval range,
                       ImplicitLocOpBuilder& b) {
   auto lb = b.create<ConstantOp>(b.getIndexAttr(range.lower));
diff --git a/third_party/xla/xla/service/gpu/fusions/mlir/elemental_hlo_to_mlir.h b/third_party/xla/xla/service/gpu/fusions/mlir/elemental_hlo_to_mlir.h
index 4add69c840b0d8..bb20a88951a38d 100644
--- a/third_party/xla/xla/service/gpu/fusions/mlir/elemental_hlo_to_mlir.h
+++ b/third_party/xla/xla/service/gpu/fusions/mlir/elemental_hlo_to_mlir.h
@@ -91,6 +91,10 @@ llvm::SmallVector<mlir::Value> ApplyAffineMap(mlir::AffineMap map,
                                               mlir::ValueRange dims,
                                               mlir::ValueRange symbols,
                                               mlir::ImplicitLocOpBuilder& b);
+llvm::SmallVector<mlir::Value> ApplyIndexing(const IndexingMap& map,
+                                             mlir::ValueRange dims,
+                                             mlir::ValueRange symbols,
+                                             mlir::ImplicitLocOpBuilder& b);
 
 // Checks all the constraints and dimension ranges in the map.
 mlir::Value CheckConstraints(const IndexingMap& map, mlir::ValueRange dims,
diff --git a/third_party/xla/xla/service/gpu/fusions/reduction_base.cc b/third_party/xla/xla/service/gpu/fusions/reduction_base.cc
index aa3dccfa8a6d3d..93e4c65852089a 100644
--- a/third_party/xla/xla/service/gpu/fusions/reduction_base.cc
+++ b/third_party/xla/xla/service/gpu/fusions/reduction_base.cc
@@ -345,8 +345,7 @@ std::optional<IndexingMap> ReductionInfo::ComputeThreadIdToOutputIndexing(
 
   auto block_offsets = GetBlockOffsetsForTiling(tiling_, ctx);
   auto thread_ids = DelinearizeInBoundsIndex(mlir::getAffineDimExpr(0, ctx),
-                                             tiling_.GetThreadsPerBlock(),
-                                             tiling_.GetThreadStrides());
+                                             tiling_.GetThreadsPerBlock());
 
   auto physical_shape =
       ShapeUtil::DeleteDimensions(hero.dimensions(), hero.operand(0)->shape());
diff --git a/third_party/xla/xla/service/gpu/fusions/reduction_mlir.cc b/third_party/xla/xla/service/gpu/fusions/reduction_mlir.cc
index 8c1dcdaf3b7bf9..7a90f21c8d787b 100644
--- a/third_party/xla/xla/service/gpu/fusions/reduction_mlir.cc
+++ b/third_party/xla/xla/service/gpu/fusions/reduction_mlir.cc
@@ -204,9 +204,8 @@ llvm::SmallVector<Value> MlirReductionFusion::EmitReduction(
   auto thread_id = state.thread_and_block_ids[0];
   Value cst_true = b.create<ma::ConstantOp>(b.getOneAttr(b.getI1Type()));
 
-  auto delinearized =
-      DelinearizeInBoundsIndex(mlir::getAffineDimExpr(0, ctx),
-                               threads_per_block, tiling.GetThreadStrides());
+  auto delinearized = DelinearizeInBoundsIndex(mlir::getAffineDimExpr(0, ctx),
+                                               threads_per_block);
   auto thread_ids = mlir_converter::ApplyAffineMap(
       mlir::AffineMap::get(1, 0, delinearized, ctx), {thread_id}, {}, b);
 
diff --git a/third_party/xla/xla/service/gpu/fusions/tiling_util.h b/third_party/xla/xla/service/gpu/fusions/tiling_util.h
index 977dbdaf986f72..bc65a733419307 100644
--- a/third_party/xla/xla/service/gpu/fusions/tiling_util.h
+++ b/third_party/xla/xla/service/gpu/fusions/tiling_util.h
@@ -40,15 +40,14 @@ namespace gpu {
 // Used by reduction and transpose emitters.
 class Tiling {
  public:
-  Tiling(absl::InlinedVector<int64_t, 4> shape,
-         absl::InlinedVector<int64_t, 4> tile_sizes,
-         absl::InlinedVector<int64_t, 4> num_threads,
+  Tiling(absl::Span<const int64_t> shape, absl::Span<const int64_t> tile_sizes,
+         absl::Span<const int64_t> num_threads,
          // By default, don't unroll anything.
          absl::InlinedVector<bool, 4> loops_to_unroll = {})
-      : shape_(shape),
-        tile_sizes_per_thread_(tile_sizes),
+      : shape_{shape.begin(), shape.end()},
+        tile_sizes_per_thread_{tile_sizes.begin(), tile_sizes.end()},
         tile_sizes_per_block_(shape.size()),
-        num_threads_(num_threads),
+        num_threads_{num_threads.begin(), num_threads.end()},
         num_blocks_(shape.size()),
         loops_to_unroll_(loops_to_unroll) {
     for (int64_t i = 0; i < shape.size(); ++i) {
diff --git a/third_party/xla/xla/service/gpu/fusions/transpose_mlir.cc b/third_party/xla/xla/service/gpu/fusions/transpose_mlir.cc
index 79111485562508..0927dbbcaffc49 100644
--- a/third_party/xla/xla/service/gpu/fusions/transpose_mlir.cc
+++ b/third_party/xla/xla/service/gpu/fusions/transpose_mlir.cc
@@ -14,17 +14,16 @@ limitations under the License.
 ==============================================================================*/
 #include "xla/service/gpu/fusions/transpose_mlir.h"
 
+#include <algorithm>
 #include <cstdint>
 #include <iterator>
 #include <optional>
 #include <vector>
 
 #include "absl/algorithm/container.h"
-#include "absl/container/inlined_vector.h"
+#include "absl/container/flat_hash_map.h"
 #include "absl/log/check.h"
-#include "absl/log/log.h"
-#include "absl/status/status.h"
-#include "llvm/ADT/ArrayRef.h"
+#include "absl/types/span.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/SmallVector.h"
 #include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
@@ -33,6 +32,7 @@ limitations under the License.
 #include "mlir/IR/AffineMap.h"  // from @llvm-project
 #include "mlir/IR/BuiltinTypes.h"  // from @llvm-project
 #include "mlir/IR/ImplicitLocOpBuilder.h"  // from @llvm-project
+#include "mlir/IR/MLIRContext.h"  // from @llvm-project
 #include "mlir/IR/TypeRange.h"  // from @llvm-project
 #include "mlir/IR/Value.h"  // from @llvm-project
 #include "mlir/IR/ValueRange.h"  // from @llvm-project
@@ -42,10 +42,11 @@ limitations under the License.
 #include "xla/hlo/ir/hlo_instructions.h"
 #include "xla/mlir/utils/type_util.h"
 #include "xla/permutation_util.h"
+#include "xla/primitive_util.h"
+#include "xla/service/gpu/fusions/fusion_emitter.h"
 #include "xla/service/gpu/fusions/mlir/computation_partitioner.h"
 #include "xla/service/gpu/fusions/mlir/elemental_hlo_to_mlir.h"
 #include "xla/service/gpu/fusions/mlir/ir/xla_gpu_ops.h"
-#include "xla/service/gpu/fusions/tiling_util.h"
 #include "xla/service/gpu/hlo_fusion_analysis.h"
 #include "xla/service/gpu/ir_emission_utils.h"
 #include "xla/service/gpu/launch_dimensions.h"
@@ -62,7 +63,6 @@ namespace {
 
 using llvm::SmallVector;
 using mlir::AffineExpr;
-using mlir::AffineMap;
 using mlir::MLIRContext;
 using mlir::RankedTensorType;
 using mlir::Value;
@@ -72,48 +72,32 @@ using mlir::func::ReturnOp;
 using mlir::tensor::ExtractOp;
 using mlir::tensor::InsertOp;
 using mlir_converter::ApplyAffineMap;
-using mlir_converter::CallTargetProvider;
-using mlir_converter::PartitionedComputation;
+using mlir_converter::ApplyIndexing;
 
-Tiling ComputeTransposeTiling(const TransposeDescription& tiled_transpose) {
-  constexpr int kNumRows = 4;
-  static_assert(WarpSize() % kNumRows == 0);
-
-  // 3D view over the output shape.
-  Vector3 transposed_dims = tiled_transpose.dimensions;
-  Vector3 permutation = tiled_transpose.permutation;
-
-  // Note: the supported permutations are their own inverses. Therefore we
-  // always use the permutation, even when we want the inverse.
-  CHECK((permutation == Vector3{0, 2, 1}) || (permutation == Vector3{2, 1, 0}));
-
-  absl::InlinedVector<int64_t, 4> input_dims{transposed_dims[permutation[0]],
-                                             transposed_dims[permutation[1]],
-                                             transposed_dims[permutation[2]]};
-
-  // We tile along the minor dimensions pre- and post-transpose.
-  absl::InlinedVector<int64_t, 4> tile_sizes{1, 1, 1};
-  tile_sizes[permutation[2]] = WarpSize() / kNumRows;
-  absl::InlinedVector<int64_t, 4> num_threads{1, 1, WarpSize()};
-  num_threads[permutation[2]] = kNumRows;
-
-  return Tiling(input_dims, tile_sizes, num_threads);
-}
+constexpr int kNumRows = 4;
+constexpr int kBaseBlockSize = WarpSize();
+constexpr int kNumThreadsPerBlock = 128;
 
 }  // namespace
 
 MlirTransposeFusion::MlirTransposeFusion(const HloFusionAnalysis& analysis)
     : analysis_(analysis),
-      tiling_(ComputeTransposeTiling(analysis.tiled_transpose())) {
+      transpose_(analysis.tiled_transpose()),
+      permutation_(transpose_.permutation),
+      input_shape_(Permute(transpose_.dimensions, permutation_)) {
   ConstHloInstructionSet transposes_to_tile;
   int index = 0;
+  int64_t shmem_usage = 0;
+  int max_element_bytes = 0;
   for (auto [root, hero] :
        llvm::zip(analysis_.fusion_roots(), analysis_.fusion_heroes())) {
     if (auto transpose = GetDescriptionForTiledTransposeEmitter(*root, *hero)) {
       transposes_to_tile.insert(hero);
       shmem_transpose_roots_.push_back(root);
+      int size = primitive_util::ByteWidth(hero->shape().element_type());
+      max_element_bytes = std::max(max_element_bytes, size);
+      shmem_usage += kBaseBlockSize * (kBaseBlockSize + 1) * size;
       shmem_transpose_root_indices_.push_back(index);
-      permutation_ = transpose->permutation;
     } else {
       side_output_roots_.push_back(root);
       side_output_root_indices_.push_back(index);
@@ -121,112 +105,73 @@ MlirTransposeFusion::MlirTransposeFusion(const HloFusionAnalysis& analysis)
     ++index;
   }
   shmem_transposes_ = {transposes_to_tile.begin(), transposes_to_tile.end()};
+
+  auto compute_block_sizes = [this](int vector_size) {
+    vector_size_ = vector_size;
+    block_size_ = kBaseBlockSize * vector_size_;
+    block_sizes_ = {1, 1, block_size_};
+    block_sizes_[permutation_[2]] = block_size_;
+    block_counts_ = {CeilOfRatio(input_shape_[0], block_sizes_[0]),
+                     CeilOfRatio(input_shape_[1], block_sizes_[1]),
+                     CeilOfRatio(input_shape_[2], block_sizes_[2])};
+  };
+  // Compute initial block sizes without vectorization. We use the result to
+  // determine whether we can vectorize.
+  compute_block_sizes(1);
+
+  // Enable vectorization if we have enough work, enough shared memory and
+  // the input dimensions are divisible by the vector size. Vectorizing loads
+  // for large data types does not help (there's already enough parallelism).
+  const auto& device = analysis_.device_info();
+  bool enough_work = Product(block_counts_) * kNumThreadsPerBlock >=
+                     4 * device.core_count() * device.threads_per_core_limit();
+  bool enough_shmem = shmem_usage * 4 <= device.shared_memory_per_block();
+  bool aligned_dims =
+      (input_shape_[2] % 2 == 0) && (input_shape_[permutation_[2]] % 2 == 0);
+  if (max_element_bytes < 4 && enough_work && enough_shmem && aligned_dims) {
+    compute_block_sizes(2);
+  }
 }
 
 std::optional<IndexingMap> MlirTransposeFusion::ComputeThreadIdToOutputIndexing(
     int64_t root_index, MLIRContext* mlir_context) const {
   const auto& hero = analysis_.fusion_hero(root_index).instruction();
-  // The block offsets are permuted, but the thread offsets remain the same.
-  auto block_offset = GetBlockOffsetsForTiling(tiling_, mlir_context)
-                          .getSubMap(std::vector<unsigned>{permutation_.begin(),
-                                                           permutation_.end()});
-  auto thread_offset = GetThreadOffsetsForTiling(tiling_, mlir_context);
-  auto permuted_tiled_shape =
-      ShapeUtil::MakeShape(U8, Permute(tiling_.GetShape(), permutation_));
-
-  auto map = ComposeIndexingMaps(
-      GetIndexingMapForTiling(
-          block_offset, thread_offset, tiling_.GetNumThreadsPerBlock(),
-          tiling_.GetNumBlocks(), tiling_.GetThreadTileSize(),
-          permuted_tiled_shape.dimensions()),
-      GetBitcastMap(permuted_tiled_shape, hero.shape(), mlir_context));
-  map.Simplify(GetIndexingMapForInstruction);
-  return map;
-}
-
-IndexingMap MlirTransposeFusion::ComputeThreadIdToInputIndexing(
-    const HloInstruction& hero, MLIRContext* mlir_context) const {
-  auto map = ComposeIndexingMaps(
-      GetIndexingMapForTiling(tiling_, mlir_context),
-      GetBitcastMap(tiling_.GetXlaShape(), hero.operand(0)->shape(),
-                    mlir_context));
-  map.Simplify(GetIndexingMapForInstruction);
-  return map;
+  return GetIndexing(/*input=*/false, hero.shape(), mlir_context);
 }
 
 std::optional<IndexingMap> MlirTransposeFusion::ComputeThreadIdToInputIndexing(
     int64_t root_index, int64_t hero_operand_index,
     MLIRContext* mlir_context) const {
   const auto& hero = analysis_.fusion_hero(root_index).instruction();
-  const auto& root = analysis_.fusion_root(root_index).instruction();
-  if (!GetDescriptionForTiledTransposeEmitter(root, hero)) {
-    // Non-transpose roots are elementwise by definition.
-    return ComputeThreadIdToOutputIndexing(root_index, mlir_context);
-  }
-  return ComputeThreadIdToInputIndexing(hero, mlir_context);
+  return GetIndexing(/*input=*/true, hero.operand(hero_operand_index)->shape(),
+                     mlir_context);
 }
 
 LaunchDimensions MlirTransposeFusion::launch_dimensions() const {
-  return LaunchDimensions(tiling_.GetNumBlocks(),
-                          tiling_.GetNumThreadsPerBlock());
+  return LaunchDimensions(Product(block_counts_), kNumThreadsPerBlock);
 }
 
-// Returns an indexing map with block_x, block_y, block_z set to 0.
-IndexingMap GetSharedMemoryWriteIndexingMap(
-    const IndexingMap& thread_id_indexing, int loop_dim) {
-  auto* mlir_context = thread_id_indexing.GetMLIRContext();
-
-  AffineExpr c0 = mlir::getAffineConstantExpr(0, mlir_context);
-  AffineExpr th_x = mlir::getAffineDimExpr(0, mlir_context);
-  SmallVector<AffineExpr, 3> tile_sizes(3);
-  mlir::bindSymbolsList(mlir_context, llvm::MutableArrayRef(tile_sizes));
-  SmallVector<AffineExpr, 3> shared_memory_indices = {
-      th_x.floorDiv(32) + 4 * tile_sizes[loop_dim], th_x % 32};
-  for (auto [index, range_val] :
-       llvm::enumerate(thread_id_indexing.GetRangeVars())) {
-    if (range_val.range.NumElements() == 1) {
-      shared_memory_indices.insert(shared_memory_indices.begin() + index, c0);
-      break;
-    }
-  }
-
-  IndexingMap shmem_write_indexing{
-      AffineMap::get(thread_id_indexing.GetDimensionCount(),
-                     thread_id_indexing.GetSymbolCount(), shared_memory_indices,
-                     mlir_context),
-      thread_id_indexing.GetDimVars(), thread_id_indexing.GetRangeVars(),
-      thread_id_indexing.GetRTVars(), thread_id_indexing.GetConstraints()};
-  return shmem_write_indexing;
-}
-
-// Returns an indexing map with block_x, block_y, block_z set to 0 and swapped
-// 2nd and 3rd results.
-IndexingMap GetSharedMemoryReadIndexingMap(
-    const IndexingMap& thread_id_indexing, Vector3 permutation) {
-  IndexingMap write_indexing = GetSharedMemoryWriteIndexingMap(
-      thread_id_indexing, /*loop_dim=*/permutation[2]);
-  llvm::SmallVector<unsigned, 3> positions;
-  absl::c_copy(permutation, std::back_inserter(positions));
-  return IndexingMap{write_indexing.GetAffineMap().getSubMap(positions),
-                     write_indexing.GetDimVars(), write_indexing.GetRangeVars(),
-                     write_indexing.GetRTVars(),
-                     write_indexing.GetConstraints()};
+IndexingMap MlirTransposeFusion::GetSharedMemoryIndexing(
+    bool read, mlir::MLIRContext* ctx) const {
+  auto thread_offsets =
+      Permute(GetThreadOffsets(ctx), read ? Vector3{0, 1, 2} : permutation_);
+  return {mlir::AffineMap::get(6, 2, thread_offsets, ctx),
+          DimVarsFromTensorSizes({kNumThreadsPerBlock, 1, 1, 1, 1, 1}),
+          RangeVarsFromTensorSizes({block_size_ / kNumRows, vector_size_}),
+          {}};
 }
 
 MlirTransposeFusion::WriteResult MlirTransposeFusion::EmitWriteToShMemMlir(
     mlir::ImplicitLocOpBuilder& builder, FuncOp entry_function,
     const HloFusionInstruction& fusion,
-    const PartitionedComputation& root_computation,
-    const CallTargetProvider& call_target_provider,
+    const mlir_converter::PartitionedComputation& root_computation,
+    const mlir_converter::CallTargetProvider& call_target_provider,
     ValueRange output_args) const {
-  std::vector<int64_t> shmem_tensor_size(tiling_.GetBlockTileSize().begin(),
-                                         tiling_.GetBlockTileSize().end());
-  // Avoid bank conflict.
-  ++shmem_tensor_size.back();
-
   MLIRContext* ctx = builder.getContext();
+  auto shmem_tensor_size = block_sizes_;
+  // Avoid bank conflicts.
+  ++shmem_tensor_size.back();
 
-  WriteResult write_result;
   // Allocate shared memory.
   SmallVector<Value> inits;
   for (auto* transpose : shmem_transposes_) {
@@ -242,21 +187,16 @@ MlirTransposeFusion::WriteResult MlirTransposeFusion::EmitWriteToShMemMlir(
     inits.push_back(entry_function.getArgument(num_inputs + index));
   }
 
-  auto tile_indexing = GetIndexingMapForTiling(tiling_, ctx);
-  IndexingMap shmem_write_indexing =
-      GetSharedMemoryWriteIndexingMap(tile_indexing, permutation_[2]);
+  IndexingMap write_indexing = GetSharedMemoryIndexing(/*read=*/false, ctx);
   auto body_builder = [&](ValueRange output_tensors, ValueRange dim_values,
                           ValueRange symbol_values) -> SmallVector<Value> {
     auto input_indices = [&](const HloInstruction* instr) {
-      auto map = ComposeIndexingMaps(
-          tile_indexing,
-          GetBitcastMap(tiling_.GetXlaShape(), instr->shape(), ctx));
-      return ApplyAffineMap(map.GetAffineMap(), dim_values, symbol_values,
-                            builder);
+      return ApplyIndexing(GetIndexing(/*input=*/true, instr->shape(), ctx),
+                           dim_values, symbol_values, builder);
     };
     SmallVector<Value> result_tensors;
-    auto shmem_indices = ApplyAffineMap(shmem_write_indexing.GetAffineMap(),
-                                        dim_values, symbol_values, builder);
+    auto shmem_indices =
+        ApplyIndexing(write_indexing, dim_values, symbol_values, builder);
     for (auto [transpose, output] :
          llvm::zip(shmem_transposes_, output_tensors)) {
       // Emit loop that writes subgraphs of transpose operands to shmem.
@@ -290,8 +230,10 @@ MlirTransposeFusion::WriteResult MlirTransposeFusion::EmitWriteToShMemMlir(
     return result_tensors;
   };
 
+  auto indexing = GetIndexing(
+      /*input=*/true, shmem_transposes_.front()->operand(0)->shape(), ctx);
   auto written_vector =
-      EmitThreadLoopNest(builder, inits, tile_indexing, body_builder);
+      EmitThreadLoopNest(builder, inits, indexing, body_builder);
   ValueRange written = written_vector;
   auto shmem_tensors = written.take_front(shmem_transposes_.size());
 
@@ -316,15 +258,14 @@ void MlirTransposeFusion::EmitReadFromShMemMlir(
     const WriteResult& written) const {
   auto* mlir_context = builder.getContext();
   auto output_indexing = *ComputeThreadIdToOutputIndexing(0, mlir_context);
-  auto shmem_output_indexing =
-      GetSharedMemoryReadIndexingMap(output_indexing, permutation_);
+  auto shmem_read_indexing =
+      GetSharedMemoryIndexing(/*read=*/true, mlir_context);
   auto result_tensors = EmitThreadLoopNest(
       builder, written.updated_outputs, output_indexing,
       [&](ValueRange output_tensors, ValueRange dim_values,
           ValueRange symbol_values) -> SmallVector<Value> {
-        auto shmem_indices =
-            ApplyAffineMap(shmem_output_indexing.GetAffineMap(), dim_values,
-                           symbol_values, builder);
+        auto shmem_indices = ApplyIndexing(shmem_read_indexing, dim_values,
+                                           symbol_values, builder);
         absl::flat_hash_map<const HloInstruction*, llvm::SmallVector<Value>>
             transpose_values;
         for (auto [transpose, shmem] :
@@ -379,5 +320,48 @@ absl::Status MlirTransposeFusion::EmitEntryFunction(
   return absl::OkStatus();
 }
 
+llvm::SmallVector<mlir::AffineExpr, 4> MlirTransposeFusion::GetThreadOffsets(
+    mlir::MLIRContext* ctx) const {
+  auto thread = mlir::getAffineDimExpr(
+      KernelFusionInterface::kIndexingMapThreadIdxDims[0], ctx);
+  auto loop = mlir::getAffineSymbolExpr(0, ctx);
+  auto vector = mlir::getAffineSymbolExpr(1, ctx);
+  int loop_stride = block_size_ * kNumRows;
+  auto linear_index = loop * loop_stride + thread * vector_size_ + vector;
+  return DelinearizeInBoundsIndex(linear_index, block_sizes_);
+}
+
+IndexingMap MlirTransposeFusion::GetIndexing(bool input,
+                                             const xla::Shape& shape,
+                                             mlir::MLIRContext* ctx) const {
+  auto raw_id = mlir::getAffineDimExpr(
+      KernelFusionInterface::kIndexingMapBlockIdxDims[0], ctx);
+  auto block_ids = Permute(DelinearizeInBoundsIndex(raw_id, block_counts_),
+                           input ? Vector3{0, 1, 2} : permutation_);
+  auto thread_offsets = GetThreadOffsets(ctx);
+  llvm::SmallVector<AffineExpr, 3> offsets;
+  for (auto [block_id, block_size, thread] :
+       llvm::zip(block_ids, block_sizes_, thread_offsets)) {
+    offsets.push_back(block_id * block_size + thread);
+  }
+  IndexingMap result{
+      mlir::AffineMap::get(6, 2, offsets, ctx),
+      DimVarsFromTensorSizes(
+          {kNumThreadsPerBlock, 1, 1, Product(block_counts_), 1, 1}),
+      RangeVarsFromTensorSizes({block_size_ / kNumRows, vector_size_}),
+      {}};
+  auto normalized_shape =
+      input ? ShapeUtil::MakeShape(shape.element_type(), input_shape_)
+            : ShapeUtil::MakeShape(shape.element_type(), transpose_.dimensions);
+  for (auto [size, dim] : llvm::zip(normalized_shape.dimensions(),
+                                    result.GetAffineMap().getResults())) {
+    result.AddConstraint(dim, {0, size - 1});
+  }
+  result =
+      ComposeIndexingMaps(result, GetBitcastMap(normalized_shape, shape, ctx));
+  result.Simplify(GetIndexingMapForInstruction);
+  return result;
+}
+
 }  // namespace gpu
 }  // namespace xla
diff --git a/third_party/xla/xla/service/gpu/fusions/transpose_mlir.h b/third_party/xla/xla/service/gpu/fusions/transpose_mlir.h
index 5607b48e7d9d7a..e4613ff8e3a6d5 100644
--- a/third_party/xla/xla/service/gpu/fusions/transpose_mlir.h
+++ b/third_party/xla/xla/service/gpu/fusions/transpose_mlir.h
@@ -21,6 +21,7 @@ limitations under the License.
 
 #include "llvm/ADT/SmallVector.h"
 #include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
+#include "mlir/IR/AffineMap.h"  // from @llvm-project
 #include "mlir/IR/ImplicitLocOpBuilder.h"  // from @llvm-project
 #include "mlir/IR/MLIRContext.h"  // from @llvm-project
 #include "mlir/IR/Value.h"  // from @llvm-project
@@ -28,8 +29,8 @@ limitations under the License.
 #include "xla/hlo/ir/hlo_instructions.h"
 #include "xla/service/gpu/fusions/mlir/computation_partitioner.h"
 #include "xla/service/gpu/fusions/mlir/mlir_fusion_emitter.h"
-#include "xla/service/gpu/fusions/tiling_util.h"
 #include "xla/service/gpu/hlo_fusion_analysis.h"
+#include "xla/service/gpu/ir_emission_utils.h"
 #include "xla/service/gpu/launch_dimensions.h"
 #include "xla/service/gpu/model/indexing_map.h"
 #include "xla/status.h"
@@ -59,9 +60,6 @@ class MlirTransposeFusion : public MlirFusionEmitterBase {
       mlir::MLIRContext* mlir_context) const override;
 
  protected:
-  IndexingMap ComputeThreadIdToInputIndexing(
-      const HloInstruction& hero, mlir::MLIRContext* mlir_context) const;
-
   absl::Status EmitEntryFunction(
       const mlir_converter::PartitionedComputations& computations,
       const mlir_converter::CallTargetProvider& call_targets,
@@ -93,8 +91,21 @@ class MlirTransposeFusion : public MlirFusionEmitterBase {
 
  private:
   const HloFusionAnalysis& analysis_;
-  Tiling tiling_;
+
+  IndexingMap GetIndexing(bool input, const xla::Shape& shape,
+                          mlir::MLIRContext* ctx) const;
+  IndexingMap GetSharedMemoryIndexing(bool read, mlir::MLIRContext* ctx) const;
+  llvm::SmallVector<mlir::AffineExpr, 4> GetThreadOffsets(
+      mlir::MLIRContext* ctx) const;
+
+  TransposeDescription transpose_;
   Vector3 permutation_;
+  std::vector<int64_t> input_shape_;
+  std::vector<int64_t> block_sizes_;  // In input elements.
+  std::vector<int64_t> block_counts_;
+  int vector_size_;
+  int block_size_;
+
   std::vector<const HloInstruction*> shmem_transposes_;
   std::vector<const HloInstruction*> shmem_transpose_roots_;
   std::vector<int> shmem_transpose_root_indices_;
diff --git a/third_party/xla/xla/service/gpu/fusions/transpose_mlir_test.cc b/third_party/xla/xla/service/gpu/fusions/transpose_mlir_test.cc
index b2343d916537cf..f2d3a299d3162b 100644
--- a/third_party/xla/xla/service/gpu/fusions/transpose_mlir_test.cc
+++ b/third_party/xla/xla/service/gpu/fusions/transpose_mlir_test.cc
@@ -49,9 +49,9 @@ TEST_F(MlirTransposeFusionTest, ThreadIndexing021) {
   EXPECT_THAT(
       fusion.ComputeThreadIdToInputIndexing(0, 0, &mlir_context_)->ToString(),
       MatchIndexingString(R"(
-        (d0, d1, d2, d3, d4, d5)[s0, s1, s2] -> (
+        (d0, d1, d2, d3, d4, d5)[s0, s1] -> (
           d3 floordiv 2,
-          d0 floordiv 32 + s1 * 4,
+          d0 floordiv 32 + s0 * 4,
           (d3 mod 2) * 32 + d0 mod 32
         )
         domain:
@@ -62,16 +62,15 @@ TEST_F(MlirTransposeFusionTest, ThreadIndexing021) {
         d4 in [0, 0]
         d5 in [0, 0]
 
-        s0 in [0, 0]
-        s1 in [0, 7]
-        s2 in [0, 0]
+        s0 in [0, 7]
+        s1 in [0, 0]
       )"));
   EXPECT_THAT(
       fusion.ComputeThreadIdToOutputIndexing(0, &mlir_context_)->ToString(),
       MatchIndexingString(R"(
-        (d0, d1, d2, d3, d4, d5)[s0, s1, s2] -> (
+        (d0, d1, d2, d3, d4, d5)[s0, s1] -> (
           d3 floordiv 2,
-          (d3 mod 2) * 32 + s1 * 4 + d0 floordiv 32,
+          (d3 mod 2) * 32 + s0 * 4 + d0 floordiv 32,
           d0 mod 32
         )
         domain:
@@ -82,9 +81,8 @@ TEST_F(MlirTransposeFusionTest, ThreadIndexing021) {
         d4 in [0, 0]
         d5 in [0, 0]
 
-        s0 in [0, 0]
-        s1 in [0, 7]
-        s2 in [0, 0]
+        s0 in [0, 7]
+        s1 in [0, 0]
       )"));
 }
 
@@ -108,9 +106,9 @@ TEST_F(MlirTransposeFusionTest, ThreadIndexing201) {
   EXPECT_THAT(
       fusion.ComputeThreadIdToInputIndexing(0, 0, &mlir_context_)->ToString(),
       MatchIndexingString(R"(
-        (d0, d1, d2, d3, d4, d5)[s0, s1, s2] -> (
+        (d0, d1, d2, d3, d4, d5)[s0, s1] -> (
           d3 floordiv 2,
-          (d3 * 32 + s1 * 4) mod 64 + d0 floordiv 32,
+          (d3 * 32 + s0 * 4) mod 64 + d0 floordiv 32,
           d0 mod 32
         )
         domain:
@@ -121,15 +119,14 @@ TEST_F(MlirTransposeFusionTest, ThreadIndexing201) {
         d4 in [0, 0]
         d5 in [0, 0]
 
-        s0 in [0, 0]
-        s1 in [0, 7]
-        s2 in [0, 0]
+        s0 in [0, 7]
+        s1 in [0, 0]
       )"));
   EXPECT_THAT(
       fusion.ComputeThreadIdToOutputIndexing(0, &mlir_context_)->ToString(),
       MatchIndexingString(R"(
-        (d0, d1, d2, d3, d4, d5)[s0, s1, s2] -> (
-          d0 floordiv 32 + s1 * 4,
+        (d0, d1, d2, d3, d4, d5)[s0, s1] -> (
+          d0 floordiv 32 + s0 * 4,
           d3 floordiv 2,
           (d3 mod 2) * 32 + d0 mod 32
         )
@@ -141,9 +138,117 @@ TEST_F(MlirTransposeFusionTest, ThreadIndexing201) {
         d4 in [0, 0]
         d5 in [0, 0]
 
-        s0 in [0, 0]
-        s1 in [0, 7]
-        s2 in [0, 0]
+        s0 in [0, 7]
+        s1 in [0, 0]
+      )"));
+}
+
+TEST_F(MlirTransposeFusionTest, ThreadIndexingVectorized021) {
+  TF_ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(R"(
+    HloModule module
+    fusion {
+      %input = f16[8192,64,64] parameter(0)
+      ROOT transpose = f16[8192,64,64] transpose(%input), dimensions={0,2,1}
+    }
+    ENTRY entry {
+      %input = f16[8192,64,64] parameter(0)
+      ROOT %fusion = f16[8192,64,64] fusion(%input), kind=kInput, calls=fusion
+    }
+  )"));
+
+  auto* root = module->entry_computation()->root_instruction();
+  auto analysis = AnalyzeFusion(*root, device_info_);
+
+  MlirTransposeFusion fusion(analysis);
+  EXPECT_THAT(
+      fusion.ComputeThreadIdToInputIndexing(0, 0, &mlir_context_)->ToString(),
+      MatchIndexingString(R"(
+        (d0, d1, d2, d3, d4, d5)[s0, s1] -> (
+          d3,
+          d0 floordiv 32 + s0 * 4,
+          (d0 mod 32) * 2 + s1
+        )
+        domain:
+        d0 in [0, 127]
+        d1 in [0, 0]
+        d2 in [0, 0]
+        d3 in [0, 8191]
+        d4 in [0, 0]
+        d5 in [0, 0]
+        s0 in [0, 15]
+        s1 in [0, 1]
+      )"));
+  EXPECT_THAT(
+      fusion.ComputeThreadIdToOutputIndexing(0, &mlir_context_)->ToString(),
+      MatchIndexingString(R"(
+        (d0, d1, d2, d3, d4, d5)[s0, s1] -> (
+          d3,
+          d0 floordiv 32 + s0 * 4,
+          (d0 mod 32) * 2 + s1
+        )
+        domain:
+        d0 in [0, 127]
+        d1 in [0, 0]
+        d2 in [0, 0]
+        d3 in [0, 8191]
+        d4 in [0, 0]
+        d5 in [0, 0]
+        s0 in [0, 15]
+        s1 in [0, 1]
+      )"));
+}
+
+TEST_F(MlirTransposeFusionTest, ThreadIndexingVectorized210) {
+  TF_ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(R"(
+    HloModule module
+    fusion {
+      %input = f16[64,64,8192] parameter(0)
+      ROOT transpose = f16[8192,64,64] transpose(%input), dimensions={2,1,0}
+    }
+    ENTRY entry {
+      %input = f16[64,64,8192] parameter(0)
+      ROOT %fusion = f16[8192,64,64] fusion(%input), kind=kInput, calls=fusion
+    })"));
+
+  auto* root = module->entry_computation()->root_instruction();
+  auto analysis = AnalyzeFusion(*root, device_info_);
+  MlirTransposeFusion fusion(analysis);
+
+  EXPECT_THAT(
+      fusion.ComputeThreadIdToInputIndexing(0, 0, &mlir_context_)->ToString(),
+      MatchIndexingString(R"(
+        (d0, d1, d2, d3, d4, d5)[s0, s1] -> (
+          d0 floordiv 32 + s0 * 4,
+          d3 floordiv 128,
+          (d3 mod 128) * 64 + s1 + (d0 mod 32) * 2
+        )
+        domain:
+        d0 in [0, 127]
+        d1 in [0, 0]
+        d2 in [0, 0]
+        d3 in [0, 8191]
+        d4 in [0, 0]
+        d5 in [0, 0]
+        s0 in [0, 15]
+        s1 in [0, 1]
+      )"));
+  EXPECT_THAT(
+      fusion.ComputeThreadIdToOutputIndexing(0, &mlir_context_)->ToString(),
+      MatchIndexingString(R"(
+        (d0, d1, d2, d3, d4, d5)[s0, s1] -> (
+          (d3 mod 128) * 64 + s0 * 4 + d0 floordiv 32,
+          d3 floordiv 128,
+          (d0 mod 32) * 2 + s1
+        )
+        domain:
+        d0 in [0, 127]
+        d1 in [0, 0]
+        d2 in [0, 0]
+        d3 in [0, 8191]
+        d4 in [0, 0]
+        d5 in [0, 0]
+        s0 in [0, 15]
+        s1 in [0, 1]
       )"));
 }
 
@@ -163,7 +268,7 @@ TEST_F(MlirTransposeFusionTest, FusedTranspose021) {
         calls=%fused_computation
     }
   )";
-  TF_ASSERT_OK(EmitAndCheckIR(kHloString, R"(
+  TF_EXPECT_OK(EmitAndCheckIR(kHloString, R"(
     // CHECK-LABEL: func.func @fused_computation(
     // CHECK-SAME:   }, %[[OUT:.*]]: tensor<20x170x160xf32>
     //
@@ -205,7 +310,7 @@ TEST_F(MlirTransposeFusionTest, FusedTranspose210) {
         calls=%fused_computation
     }
   )";
-  TF_ASSERT_OK(EmitAndCheckIR(kHloString, R"(
+  TF_EXPECT_OK(EmitAndCheckIR(kHloString, R"(
     // CHECK-LABEL: func.func @fused_computation(
     // CHECK-SAME:   }, %[[OUT:.*]]: tensor<170x160x20xf32>
     //
@@ -246,7 +351,7 @@ TEST_F(MlirTransposeFusionTest, Transpose021_Parameter) {
         calls=%fused_computation
     }
   )";
-  TF_ASSERT_OK(EmitAndCheckIR(kHloString, R"(
+  TF_EXPECT_OK(EmitAndCheckIR(kHloString, R"(
     // CHECK-LABEL: func.func @fused_computation(
     // CHECK-SAME:   }, %[[OUT:.*]]: tensor<20x170x160xf32>
     //
@@ -286,7 +391,7 @@ TEST_F(MlirTransposeFusionTest, Transpose021_NoEpilogue) {
         calls=%fused_computation
     }
   )";
-  TF_ASSERT_OK(EmitAndCheckIR(kHloString, R"(
+  TF_EXPECT_OK(EmitAndCheckIR(kHloString, R"(
     // CHECK:       func.func private @fused_computation__epilogue__
     // CHECK-NEXT:  return %
   )"));
@@ -317,13 +422,13 @@ TEST_F(MlirTransposeFusionTest, Transpose_2D) {
     HloModule Transpose
 
     %fused_computation {
-      %param_0 = f64[100, 200] parameter(0)
-      ROOT %transpose= f64[200,100] transpose(f64[100, 200] %param_0),
+      %param_0 = f64[64, 64] parameter(0)
+      ROOT %transpose= f64[64,64] transpose( %param_0),
         dimensions={1,0}
     }
     ENTRY main {
-      %param = f64[100, 200] parameter(0)
-      ROOT %fusion = f64[200,100] fusion(%param), kind=kInput,
+      %param = f64[64,64] parameter(0)
+      ROOT %fusion = f64[64,64] fusion(%param), kind=kInput,
         calls=%fused_computation
     }
   )";
@@ -453,6 +558,42 @@ TEST_F(MlirTransposeFusionTest, SideOutputs) {
   EXPECT_TRUE(RunAndCompareNoHloPasses(kHloString, ErrorSpec{1e-3}));
 }
 
+TEST_F(MlirTransposeFusionTest, VectorizedTranspose021) {
+  auto kHloString = R"(
+    HloModule Transpose
+    %fused_computation {
+      %p0 = bf16[256,128,128] parameter(0)
+      %transpose = bf16[256,128,128] transpose(%p0), dimensions={0,2,1}
+    }
+    ENTRY main {
+      %param = bf16[256,128,128] parameter(0)
+      ROOT %fusion = bf16[256,128,128] fusion(%param), kind=kInput,
+        calls=%fused_computation
+    }
+  )";
+  TF_EXPECT_OK(EmitAndCheckIR(
+      kHloString, "// CHECK: xla_gpu.allocate_shared : tensor<1x64x65xbf16>"));
+  EXPECT_TRUE(RunAndCompareNoHloPasses(kHloString, ErrorSpec{1e-3}));
+}
+
+TEST_F(MlirTransposeFusionTest, VectorizedTranspose210) {
+  auto kHloString = R"(
+    HloModule Transpose
+    %fused_computation {
+      %p0 = bf16[256,128,128] parameter(0)
+      %transpose = bf16[128,128,256] transpose(%p0), dimensions={2,1,0}
+    }
+    ENTRY main {
+      %param = bf16[256,128,128] parameter(0)
+      ROOT %fusion = bf16[128,128,256] fusion(%param), kind=kInput,
+        calls=%fused_computation
+    }
+  )";
+  TF_EXPECT_OK(EmitAndCheckIR(
+      kHloString, "// CHECK: xla_gpu.allocate_shared : tensor<64x1x65xbf16>"));
+  EXPECT_TRUE(RunAndCompareNoHloPasses(kHloString, ErrorSpec{1e-3}));
+}
+
 }  // namespace
 }  // namespace gpu
 }  // namespace xla
diff --git a/third_party/xla/xla/service/gpu/model/indexing_analysis.cc b/third_party/xla/xla/service/gpu/model/indexing_analysis.cc
index fd971af23bb3c6..e7a56b62d45c75 100644
--- a/third_party/xla/xla/service/gpu/model/indexing_analysis.cc
+++ b/third_party/xla/xla/service/gpu/model/indexing_analysis.cc
@@ -1141,8 +1141,7 @@ IndexingMap CreateIdentityMap(const Shape& shape, MLIRContext* mlir_context) {
 }
 
 llvm::SmallVector<AffineExpr, 4> DelinearizeInBoundsIndex(
-    AffineExpr linear, absl::Span<const int64_t> sizes,
-    absl::Span<const int64_t> strides) {
+    AffineExpr linear, absl::Span<const int64_t> sizes) {
   llvm::SmallVector<AffineExpr, 4> result;
   result.reserve(sizes.size());
   if (absl::c_linear_search(sizes, 0)) {
@@ -1152,6 +1151,7 @@ llvm::SmallVector<AffineExpr, 4> DelinearizeInBoundsIndex(
     return result;
   }
 
+  auto strides = ComputeStrides(sizes);
   for (auto [size, stride] : llvm::zip(sizes, strides)) {
     result.push_back(linear.floorDiv(stride) % size);
   }
@@ -1197,8 +1197,7 @@ IndexingMap GetIndexingMapFromLogicalToPhysicalLayout(
 AffineMap GetBlockOffsetsForTiling(const Tiling& tiling,
                                    MLIRContext* mlir_context) {
   auto offsets = DelinearizeInBoundsIndex(getAffineDimExpr(3, mlir_context),
-                                          tiling.GetBlockCounts(),
-                                          tiling.GetBlockStrides());
+                                          tiling.GetBlockCounts());
   for (auto&& [offset, tile_size] :
        llvm::zip(offsets, tiling.GetBlockTileSize())) {
     offset = offset * tile_size;
@@ -1209,8 +1208,7 @@ AffineMap GetBlockOffsetsForTiling(const Tiling& tiling,
 AffineMap GetThreadOffsetsForTiling(const Tiling& tiling,
                                     MLIRContext* mlir_context) {
   auto offsets = DelinearizeInBoundsIndex(getAffineDimExpr(0, mlir_context),
-                                          tiling.GetThreadsPerBlock(),
-                                          tiling.GetThreadStrides());
+                                          tiling.GetThreadsPerBlock());
   for (int dim = 0; dim < tiling.GetShape().size(); ++dim) {
     if (tiling.GetThreadTileSize()[dim] > 1) {
       offsets[dim] = offsets[dim] + getAffineSymbolExpr(dim, mlir_context) *
diff --git a/third_party/xla/xla/service/gpu/model/indexing_analysis.h b/third_party/xla/xla/service/gpu/model/indexing_analysis.h
index 76fc8ab39e476c..8d6125e71645dc 100644
--- a/third_party/xla/xla/service/gpu/model/indexing_analysis.h
+++ b/third_party/xla/xla/service/gpu/model/indexing_analysis.h
@@ -184,8 +184,7 @@ IndexingMap CreateIdentityMap(const Shape& shape,
                               mlir::MLIRContext* mlir_context);
 
 llvm::SmallVector<mlir::AffineExpr, 4> DelinearizeInBoundsIndex(
-    mlir::AffineExpr linear, absl::Span<const int64_t> sizes,
-    absl::Span<const int64_t> strides);
+    mlir::AffineExpr linear, absl::Span<const int64_t> sizes);
 
 // Returns the output-to-input indexing map of the first output of `instr`
 IndexingMap GetIndexingMapForInstruction(const HloInstruction* instr,

From aa8bb6a0db6bf51e44a60936c5de4eccd603e287 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Tue, 21 May 2024 01:21:04 -0700
Subject: [PATCH 451/478] Automated Code Change

PiperOrigin-RevId: 635716090
---
 tensorflow/core/common_runtime/executor.cc | 43 +++++++++++-----------
 1 file changed, 22 insertions(+), 21 deletions(-)

diff --git a/tensorflow/core/common_runtime/executor.cc b/tensorflow/core/common_runtime/executor.cc
index cb6186e6ff9b84..46b44f7dc86935 100644
--- a/tensorflow/core/common_runtime/executor.cc
+++ b/tensorflow/core/common_runtime/executor.cc
@@ -578,7 +578,8 @@ bool MightTrace(const tsl::tracing::EventCollector* event_collector,
 
   if (tsl::profiler::ScopedAnnotation::IsEnabled()) return true;
 
-  return profiler::TraceMe::Active(profiler::GetTFTraceMeLevel(is_expensive));
+  return tsl::profiler::TraceMe::Active(
+      tsl::profiler::GetTFTraceMeLevel(is_expensive));
 }
 
 template <class PropagatorStateType>
@@ -599,9 +600,9 @@ Status ExecutorState<PropagatorStateType>::ProcessSync(
     profiler::AnnotatedTraceMe activity(
         [op_kernel, &ctx] {
           return op_kernel->TraceString(
-              ctx, /*verbose=*/profiler::TfOpDetailsEnabled());
+              ctx, /*verbose=*/tsl::profiler::TfOpDetailsEnabled());
         },
-        profiler::GetTFTraceMeLevel(is_expensive));
+        tsl::profiler::GetTFTraceMeLevel(is_expensive));
     device->Compute(op_kernel, &ctx);
   } else if (kernel_stats_->HasExpensiveMarker(item)) {
     KernelTimer timer;
@@ -641,9 +642,9 @@ void ExecutorState<PropagatorStateType>::ProcessAsync(
     profiler::AnnotatedTraceMe activity(
         [async_kernel, state] {
           return async_kernel->TraceString(
-              state->ctx, /*verbose=*/profiler::TfOpDetailsEnabled());
+              state->ctx, /*verbose=*/tsl::profiler::TfOpDetailsEnabled());
         },
-        profiler::GetTFTraceMeLevel(/*is_expensive=*/false));
+        tsl::profiler::GetTFTraceMeLevel(/*is_expensive=*/false));
 
     // Trace async op start.
     profiler::TraceMeProducer producer(
@@ -654,7 +655,7 @@ void ExecutorState<PropagatorStateType>::ProcessAsync(
                {"kernel_type", async_kernel->type_string()},
                {"step_id", step_id_}});
         },
-        profiler::ContextType::kTfExecutor);
+        tsl::profiler::ContextType::kTfExecutor);
 
     auto done = [this, state, activity_id, ctx_id = producer.GetContextId()]() {
       // Trace async op done.
@@ -666,7 +667,7 @@ void ExecutorState<PropagatorStateType>::ProcessAsync(
                  {"kernel_type", state->item->kernel->type_string()},
                  {"step_id", step_id_}});
           },
-          profiler::ContextType::kTfExecutor, ctx_id);
+          tsl::profiler::ContextType::kTfExecutor, ctx_id);
 
       Device* device = immutable_state_.params().device;
       NodeExecStatsInterface* stats = state->stats;  // Shorthand
@@ -728,8 +729,8 @@ void ExecutorState<PropagatorStateType>::ProcessConstTensor(
 template <class PropagatorStateType>
 void ExecutorState<PropagatorStateType>::Process(const TaggedNode& tagged_node,
                                                  int64_t scheduled_nsec) {
-  profiler::TraceMe traceme("ExecutorState::Process Scheduled",
-                            profiler::TraceMeLevel::kVerbose);
+  tsl::profiler::TraceMe traceme("ExecutorState::Process Scheduled",
+                                 tsl::profiler::TraceMeLevel::kVerbose);
   TaggedNodeReadyQueue inline_ready;
   inline_ready.push_back(tagged_node);
   return ProcessInline(&inline_ready, scheduled_nsec);
@@ -824,8 +825,8 @@ void ExecutorState<PropagatorStateType>::ProcessInline(
                 "ExecutorState::Process",
                 {{"id", step_id_}, {"iter_num", tagged_node.get_iter_num()}});
           },
-          profiler::ContextType::kTfExecutor, trace_id_,
-          profiler::TraceMeLevel::kInfo);
+          tsl::profiler::ContextType::kTfExecutor, trace_id_,
+          tsl::profiler::TraceMeLevel::kInfo);
       last_iter_num = current_iter_num;
     }
     inline_ready->pop_front();
@@ -1265,7 +1266,7 @@ bool ExecutorState<PropagatorStateType>::NodeDone(
 template <class PropagatorStateType>
 void ExecutorState<PropagatorStateType>::ScheduleReady(
     TaggedNodeSeq* ready, TaggedNodeReadyQueue* inline_ready) {
-  profiler::TraceMe activity(
+  tsl::profiler::TraceMe activity(
       [&]() {
         return strings::StrCat(
             "ExecutorState::ScheduleReady#",
@@ -1273,7 +1274,7 @@ void ExecutorState<PropagatorStateType>::ScheduleReady(
             ",inline_ready_size=",
             (inline_ready == nullptr ? -1 : inline_ready->size()), "#");
       },
-      profiler::GetTFTraceMeLevel(/*is_expensive=*/false));
+      tsl::profiler::GetTFTraceMeLevel(/*is_expensive=*/false));
   DCHECK(!ready->empty());
 
   int64_t scheduled_nsec = 0;
@@ -1352,14 +1353,14 @@ void ExecutorState<PropagatorStateType>::ScheduleReady(
           TaggedNodeSeq ready_chunk{it, end};
           RunTask(
               [this, ready_chunk = std::move(ready_chunk), scheduled_nsec]() {
-                profiler::TraceMe activity(
+                tsl::profiler::TraceMe activity(
                     [&]() {
                       return strings::StrCat(
                           "ExecutorState::ScheduleReady::"
                           "ChildThreadExpensiveNodes#",
                           "ready_chunk_size=", ready_chunk.size(), "#");
                     },
-                    profiler::GetTFTraceMeLevel(/*is_expensive=*/false));
+                    tsl::profiler::GetTFTraceMeLevel(/*is_expensive=*/false));
                 for (auto& tagged_node : ready_chunk) {
                   RunTask(std::bind(&ExecutorState::Process, this, tagged_node,
                                     scheduled_nsec),
@@ -1465,8 +1466,8 @@ void ExecutorState<PropagatorStateType>::Finish() {
             return profiler::TraceMeEncode("ExecutorDoneCallback",
                                            {{"id", step_id}});
           },
-          profiler::ContextType::kTfExecutor, trace_id,
-          profiler::TraceMeLevel::kInfo);
+          tsl::profiler::ContextType::kTfExecutor, trace_id,
+          tsl::profiler::TraceMeLevel::kInfo);
       done_cb(status);
     });
     return;
@@ -1488,8 +1489,8 @@ void ExecutorState<PropagatorStateType>::Finish() {
               return profiler::TraceMeEncode("ExecutorDoneCallback",
                                              {{"id", step_id}});
             },
-            profiler::ContextType::kTfExecutor, trace_id,
-            profiler::TraceMeLevel::kInfo);
+            tsl::profiler::ContextType::kTfExecutor, trace_id,
+            tsl::profiler::TraceMeLevel::kInfo);
         done_cb(status);
       });
     });
@@ -1503,8 +1504,8 @@ void ExecutorState<PropagatorStateType>::Finish() {
             return profiler::TraceMeEncode("ExecutorDoneCallback",
                                            {{"id", step_id}});
           },
-          profiler::ContextType::kTfExecutor, trace_id,
-          profiler::TraceMeLevel::kInfo);
+          tsl::profiler::ContextType::kTfExecutor, trace_id,
+          tsl::profiler::TraceMeLevel::kInfo);
       done_cb(status);
     });
   }

From c0ac2d38c481eb1c48f0950d3711bceb4b311003 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Tue, 21 May 2024 02:02:26 -0700
Subject: [PATCH 452/478] Update GraphDef version to 1869.

PiperOrigin-RevId: 635725513
---
 tensorflow/core/public/version.h | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorflow/core/public/version.h b/tensorflow/core/public/version.h
index d5df864e20d0ca..c436412c0e88c6 100644
--- a/tensorflow/core/public/version.h
+++ b/tensorflow/core/public/version.h
@@ -108,7 +108,7 @@ limitations under the License.
 
 #define TF_GRAPH_DEF_VERSION_MIN_PRODUCER 0
 #define TF_GRAPH_DEF_VERSION_MIN_CONSUMER 0
-#define TF_GRAPH_DEF_VERSION 1868  // Updated: 2024/5/20
+#define TF_GRAPH_DEF_VERSION 1869  // Updated: 2024/5/21
 
 // Checkpoint compatibility versions (the versions field in SavedSliceMeta).
 //

From 1eb72a6a6646ee509573317278ac89a3d5ccffea Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Tue, 21 May 2024 02:02:27 -0700
Subject: [PATCH 453/478] compat: Update forward compatibility horizon to
 2024-05-21

PiperOrigin-RevId: 635725518
---
 tensorflow/python/compat/compat.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorflow/python/compat/compat.py b/tensorflow/python/compat/compat.py
index 5beb9afea1416a..0a9176e41d5a19 100644
--- a/tensorflow/python/compat/compat.py
+++ b/tensorflow/python/compat/compat.py
@@ -29,7 +29,7 @@
 # This value changes every day with an automatic CL. It can be modified in code
 # via `forward_compatibility_horizon()` or with the environment variable
 # TF_FORWARD_COMPATIBILITY_DELTA_DAYS, which is added to the compatibility date.
-_FORWARD_COMPATIBILITY_HORIZON = datetime.date(2024, 5, 20)
+_FORWARD_COMPATIBILITY_HORIZON = datetime.date(2024, 5, 21)
 _FORWARD_COMPATIBILITY_DELTA_DAYS_VAR_NAME = "TF_FORWARD_COMPATIBILITY_DELTA_DAYS"
 _FORWARD_COMPATIBILITY_DATE_NUMBER = None
 

From 688998aabec710e5ae27c63abd7868b83f9c69e1 Mon Sep 17 00:00:00 2001
From: Gunhyun Park <gunhyun@google.com>
Date: Tue, 21 May 2024 02:26:22 -0700
Subject: [PATCH 454/478] Add unbounded dynamism test for FftOp.

PiperOrigin-RevId: 635731691
---
 .../xla/xla/client/xla_builder_test.cc        | 96 +++++++++++++++++++
 .../xla/xla/service/shape_inference.cc        | 22 +++--
 .../xla/xla/service/shape_inference_test.cc   | 81 ++++++++++++++++
 3 files changed, 192 insertions(+), 7 deletions(-)

diff --git a/third_party/xla/xla/client/xla_builder_test.cc b/third_party/xla/xla/client/xla_builder_test.cc
index 9d3ae1850a9092..fad65f35384248 100644
--- a/third_party/xla/xla/client/xla_builder_test.cc
+++ b/third_party/xla/xla/client/xla_builder_test.cc
@@ -1240,6 +1240,54 @@ TEST(XlaBuilderTest, DotWithPreferredElementType) {
       ShapeUtil::Equal(ShapeUtil::MakeShape(U32, {2, 2}), result_shape));
 }
 
+TEST(XlaBuilderTest, FftWithFFT) {
+  XlaBuilder b(TestName());
+  TF_ASSERT_OK_AND_ASSIGN(const Shape operand, ParseShape("c64[5, <=10]"));
+  const std::vector<int64_t> fft_length = {5, 10};
+  TF_ASSERT_OK_AND_ASSIGN(const Shape expected, ParseShape("c64[5, <=10]"));
+  Fft(Parameter(&b, 0, operand, "operand"), /*fft_type=*/FftType::FFT,
+      fft_length);
+  TF_ASSERT_OK_AND_ASSIGN(const auto module, BuildHloModule(b));
+  EXPECT_THAT(GetRoot(*module),
+              GmockMatch(m::Op().WithShapeEqualTo(&expected)));
+}
+
+TEST(XlaBuilderTest, FftWithIFFT) {
+  XlaBuilder b(TestName());
+  TF_ASSERT_OK_AND_ASSIGN(const Shape operand, ParseShape("c64[5, <=10]"));
+  const std::vector<int64_t> fft_length = {5, 10};
+  TF_ASSERT_OK_AND_ASSIGN(const Shape expected, ParseShape("c64[5, <=10]"));
+  Fft(Parameter(&b, 0, operand, "operand"), /*fft_type=*/FftType::IFFT,
+      fft_length);
+  TF_ASSERT_OK_AND_ASSIGN(const auto module, BuildHloModule(b));
+  EXPECT_THAT(GetRoot(*module),
+              GmockMatch(m::Op().WithShapeEqualTo(&expected)));
+}
+
+TEST(XlaBuilderTest, FftWithRFFT) {
+  XlaBuilder b(TestName());
+  TF_ASSERT_OK_AND_ASSIGN(const Shape operand, ParseShape("f64[10, <=5]"));
+  const std::vector<int64_t> fft_length = {5};
+  TF_ASSERT_OK_AND_ASSIGN(const Shape expected, ParseShape("c128[10, <=3]"));
+  Fft(Parameter(&b, 0, operand, "operand"), /*fft_type=*/FftType::RFFT,
+      fft_length);
+  TF_ASSERT_OK_AND_ASSIGN(const auto module, BuildHloModule(b));
+  EXPECT_THAT(GetRoot(*module),
+              GmockMatch(m::Op().WithShapeEqualTo(&expected)));
+}
+
+TEST(XlaBuilderTest, FftWithIRFFT) {
+  XlaBuilder b(TestName());
+  TF_ASSERT_OK_AND_ASSIGN(const Shape operand, ParseShape("c128[10, <=3]"));
+  const std::vector<int64_t> fft_length = {5};
+  TF_ASSERT_OK_AND_ASSIGN(const Shape expected, ParseShape("f64[10, <=5]"));
+  Fft(Parameter(&b, 0, operand, "operand"), /*fft_type=*/FftType::IRFFT,
+      fft_length);
+  TF_ASSERT_OK_AND_ASSIGN(const auto module, BuildHloModule(b));
+  EXPECT_THAT(GetRoot(*module),
+              GmockMatch(m::Op().WithShapeEqualTo(&expected)));
+}
+
 TEST(XlaBuilderTest, SparseDot) {
   XlaBuilder b(TestName());
   auto lhs = Parameter(&b, 0, ShapeUtil::MakeShape(F32, {10, 16}), "lhs");
@@ -2513,6 +2561,54 @@ TEST(XlaBuilderTest, UnboundedDynamicUpdateSlice) {
               GmockMatch(m::Op().WithShapeEqualTo(&expected)));
 }
 
+TEST(XlaBuilderTest, UnboundedFftWithFFT) {
+  XlaBuilder b(TestName());
+  TF_ASSERT_OK_AND_ASSIGN(const Shape operand, ParseShape("c64[2, <=5, ?]"));
+  const std::vector<int64_t> fft_length = {5, 10};
+  TF_ASSERT_OK_AND_ASSIGN(const Shape expected, ParseShape("c64[2, <=5, ?]"));
+  Fft(Parameter(&b, 0, operand, "operand"), /*fft_type=*/FftType::FFT,
+      fft_length);
+  TF_ASSERT_OK_AND_ASSIGN(const auto module, BuildHloModule(b));
+  EXPECT_THAT(GetRoot(*module),
+              GmockMatch(m::Op().WithShapeEqualTo(&expected)));
+}
+
+TEST(XlaBuilderTest, UnboundedFftWithIFFT) {
+  XlaBuilder b(TestName());
+  TF_ASSERT_OK_AND_ASSIGN(const Shape operand, ParseShape("c64[2, <=5, ?]"));
+  const std::vector<int64_t> fft_length = {5, 10};
+  TF_ASSERT_OK_AND_ASSIGN(const Shape expected, ParseShape("c64[2, <=5, ?]"));
+  Fft(Parameter(&b, 0, operand, "operand"), /*fft_type=*/FftType::IFFT,
+      fft_length);
+  TF_ASSERT_OK_AND_ASSIGN(const auto module, BuildHloModule(b));
+  EXPECT_THAT(GetRoot(*module),
+              GmockMatch(m::Op().WithShapeEqualTo(&expected)));
+}
+
+TEST(XlaBuilderTest, UnboundedFftWithRFFT) {
+  XlaBuilder b(TestName());
+  TF_ASSERT_OK_AND_ASSIGN(const Shape operand, ParseShape("f64[2, <=5, ?]"));
+  const std::vector<int64_t> fft_length = {5, 10};
+  TF_ASSERT_OK_AND_ASSIGN(const Shape expected, ParseShape("c128[2, <=5, 6]"));
+  Fft(Parameter(&b, 0, operand, "operand"), /*fft_type=*/FftType::RFFT,
+      fft_length);
+  TF_ASSERT_OK_AND_ASSIGN(const auto module, BuildHloModule(b));
+  EXPECT_THAT(GetRoot(*module),
+              GmockMatch(m::Op().WithShapeEqualTo(&expected)));
+}
+
+TEST(XlaBuilderTest, UnboundedFftWithIRFFT) {
+  XlaBuilder b(TestName());
+  TF_ASSERT_OK_AND_ASSIGN(const Shape operand, ParseShape("c128[2, <=5, ?]"));
+  const std::vector<int64_t> fft_length = {5, 10};
+  TF_ASSERT_OK_AND_ASSIGN(const Shape expected, ParseShape("f64[2, <=5, 10]"));
+  Fft(Parameter(&b, 0, operand, "operand"), /*fft_type=*/FftType::IRFFT,
+      fft_length);
+  TF_ASSERT_OK_AND_ASSIGN(const auto module, BuildHloModule(b));
+  EXPECT_THAT(GetRoot(*module),
+              GmockMatch(m::Op().WithShapeEqualTo(&expected)));
+}
+
 TEST(XlaBuilderTest, UnboundedGather) {
   XlaBuilder b(TestName());
   TF_ASSERT_OK_AND_ASSIGN(const Shape operand, ParseShape("f32[3, 4, 2]"));
diff --git a/third_party/xla/xla/service/shape_inference.cc b/third_party/xla/xla/service/shape_inference.cc
index 9268479cba2005..fd7c82825680c4 100644
--- a/third_party/xla/xla/service/shape_inference.cc
+++ b/third_party/xla/xla/service/shape_inference.cc
@@ -2248,8 +2248,10 @@ ShapeInference::InferScalarBroadcastShape(absl::Span<const Shape> shapes) {
       }
       RET_CHECK_RANK(in);
       for (int i = 0; i < fft_rank; i++) {
-        if (in.dimensions(in.dimensions_size() - fft_rank + i) !=
-            fft_length[i]) {
+        if (!IsUnboundedDynamicSize(
+                in.dimensions(in.dimensions_size() - fft_rank + i)) &&
+            in.dimensions(in.dimensions_size() - fft_rank + i) !=
+                fft_length[i]) {
           return InvalidArgument(
               "RFFT requires innermost dimensions match fft_length but "
               "dimension %d is %d and should be %d.",
@@ -2264,6 +2266,9 @@ ShapeInference::InferScalarBroadcastShape(absl::Span<const Shape> shapes) {
       if (fft_length[fft_rank - 1] != 0) {
         result.set_dimensions(result.dimensions_size() - 1,
                               fft_length[fft_rank - 1] / 2 + 1);
+        if (in.is_unbounded_dynamic_dimension(result.dimensions_size() - 1)) {
+          result.set_dynamic_dimension(result.dimensions_size() - 1, false);
+        }
       }
       return result;
     }
@@ -2286,18 +2291,21 @@ ShapeInference::InferScalarBroadcastShape(absl::Span<const Shape> shapes) {
         }
       }
       // The size of zero-sized dimensions is preserved.
-      if ((in.dimensions(in.dimensions_size() - 1) != 0 ||
-           fft_length[fft_rank - 1] != 0) &&
-          in.dimensions(in.dimensions_size() - 1) !=
-              fft_length[fft_rank - 1] / 2 + 1) {
+      int64_t last_in_dimension_size = in.dimensions(in.dimensions_size() - 1);
+      if ((last_in_dimension_size != 0 || fft_length[fft_rank - 1] != 0) &&
+          !IsUnboundedDynamicSize(last_in_dimension_size) &&
+          last_in_dimension_size != fft_length[fft_rank - 1] / 2 + 1) {
         return InvalidArgument(
             "IRFFT requires innermost dimension matches fft_length/2+1, but "
             "dimension %d is %d and should be %d.",
-            in.dimensions_size() - 1, in.dimensions(in.dimensions_size() - 1),
+            in.dimensions_size() - 1, last_in_dimension_size,
             fft_length[fft_rank - 1] / 2 + 1);
       }
       result.set_dimensions(result.dimensions_size() - 1,
                             fft_length[fft_rank - 1]);
+      if (in.is_unbounded_dynamic_dimension(result.dimensions_size() - 1)) {
+        result.set_dynamic_dimension(result.dimensions_size() - 1, false);
+      }
       return result;
     }
     default:
diff --git a/third_party/xla/xla/service/shape_inference_test.cc b/third_party/xla/xla/service/shape_inference_test.cc
index 2ec1968a94b7c2..14c3e804563815 100644
--- a/third_party/xla/xla/service/shape_inference_test.cc
+++ b/third_party/xla/xla/service/shape_inference_test.cc
@@ -987,6 +987,16 @@ TEST_F(ShapeInferenceTest, InferFftShapeTestFftRanks) {
   fft::Fail(shape, type, {64, 32, 16, 8}, fft::unsupported_rank);
 }
 
+TEST_F(ShapeInferenceTest, InferFftShapeTestFftRanksBounded) {
+  FftType type = FftType::FFT;
+  const Shape shape = ShapeUtil::MakeShape(C64, {16, 8}, {false, true});
+  fft::Fail(shape, type, {}, fft::unsupported_rank);
+  fft::Pass(shape, type, {8}, shape);
+  fft::Pass(shape, type, {16, 8}, shape);
+  fft::Fail(shape, type, {32, 16, 8}, fft::invalid_rank);
+  fft::Fail(shape, type, {64, 32, 16, 8}, fft::unsupported_rank);
+}
+
 TEST_F(ShapeInferenceTest, InferFftShapeTestFftTypes) {
   FftType type = FftType::FFT;
   const Shape shape_f32 = ShapeUtil::MakeShape(F32, {16, 8});
@@ -1005,6 +1015,16 @@ TEST_F(ShapeInferenceTest, InferFftShapeTestIfftRanks) {
   fft::Fail(shape, type, {64, 32, 16, 8}, fft::unsupported_rank);
 }
 
+TEST_F(ShapeInferenceTest, InferFftShapeTestIfftRanksBounded) {
+  FftType type = FftType::IFFT;
+  const Shape shape = ShapeUtil::MakeShape(C64, {16, 8}, {false, true});
+  fft::Fail(shape, type, {}, fft::unsupported_rank);
+  fft::Pass(shape, type, {8}, shape);
+  fft::Pass(shape, type, {16, 8}, shape);
+  fft::Fail(shape, type, {32, 16, 8}, fft::invalid_rank);
+  fft::Fail(shape, type, {64, 32, 16, 8}, fft::unsupported_rank);
+}
+
 TEST_F(ShapeInferenceTest, InferFftShapeTestIfftTypes) {
   FftType type = FftType::IFFT;
   const Shape shape_f32 = ShapeUtil::MakeShape(F32, {16, 8});
@@ -1042,6 +1062,12 @@ TEST_F(ShapeInferenceTest, InferFftShapeTestRfftDimensions) {
   const Shape shape_out = ShapeUtil::MakeShape(C64, {16, 5});
   fft::Pass(even_shape_in, type, {16, 8}, shape_out);
   fft::Pass(odd_shape_in, type, {16, 9}, shape_out);
+
+  const Shape bounded_shape_in =
+      ShapeUtil::MakeShape(F32, {16, 8}, {false, true});
+  const Shape bounded_shape_out =
+      ShapeUtil::MakeShape(C64, {16, 5}, {false, true});
+  fft::Pass(bounded_shape_in, type, {16, 8}, bounded_shape_out);
 }
 
 TEST_F(ShapeInferenceTest, InferFftShapeTestRfftTypes) {
@@ -1080,6 +1106,12 @@ TEST_F(ShapeInferenceTest, InferFftShapeTestIrfftDimensions) {
   const Shape odd_shape_out = ShapeUtil::MakeShape(F32, {16, 9});
   fft::Pass(shape, type, {16, 8}, even_shape_out);
   fft::Pass(shape, type, {16, 9}, odd_shape_out);
+
+  const Shape bounded_shape_in =
+      ShapeUtil::MakeShape(C64, {16, 5}, {false, true});
+  const Shape bounded_shape_out =
+      ShapeUtil::MakeShape(F32, {16, 9}, {false, true});
+  fft::Pass(bounded_shape_in, type, {16, 9}, bounded_shape_out);
 }
 
 TEST_F(ShapeInferenceTest, InferFftShapeTestIrfftTypes) {
@@ -4528,6 +4560,55 @@ TEST_F(ShapeInferenceTest, UnboundedDynamicUpdateSlice) {
       << " expected: " << ShapeUtil::HumanString(expected);
 }
 
+TEST_F(ShapeInferenceTest, UnboundedFftWithFFT) {
+  TF_ASSERT_OK_AND_ASSIGN(const Shape operand, ParseShape("c64[2, <=5, ?]"));
+  const std::vector<int64_t> fft_length = {5, 10};
+  TF_ASSERT_OK_AND_ASSIGN(const Shape expected, ParseShape("c64[2, <=5, ?]"));
+  TF_ASSERT_OK_AND_ASSIGN(const Shape inferred_shape,
+                          ShapeInference::InferFftShape(
+                              operand, /*fft_type=*/FftType::FFT, fft_length));
+  EXPECT_TRUE(ShapeUtil::Equal(inferred_shape, expected))
+      << "inferred: " << ShapeUtil::HumanString(inferred_shape)
+      << " expected: " << ShapeUtil::HumanString(expected);
+}
+
+TEST_F(ShapeInferenceTest, UnboundedFftWithIFFT) {
+  TF_ASSERT_OK_AND_ASSIGN(const Shape operand, ParseShape("c64[2, <=5, ?]"));
+  const std::vector<int64_t> fft_length = {5, 10};
+  TF_ASSERT_OK_AND_ASSIGN(const Shape expected, ParseShape("c64[2, <=5, ?]"));
+  TF_ASSERT_OK_AND_ASSIGN(const Shape inferred_shape,
+                          ShapeInference::InferFftShape(
+                              operand, /*fft_type=*/FftType::IFFT, fft_length));
+  EXPECT_TRUE(ShapeUtil::Equal(inferred_shape, expected))
+      << "inferred: " << ShapeUtil::HumanString(inferred_shape)
+      << " expected: " << ShapeUtil::HumanString(expected);
+}
+
+TEST_F(ShapeInferenceTest, UnboundedFftWithRFFT) {
+  TF_ASSERT_OK_AND_ASSIGN(const Shape operand, ParseShape("f64[2, <=5, ?]"));
+  const std::vector<int64_t> fft_length = {5, 10};
+  TF_ASSERT_OK_AND_ASSIGN(const Shape expected, ParseShape("c128[2, <=5, 6]"));
+  TF_ASSERT_OK_AND_ASSIGN(const Shape inferred_shape,
+                          ShapeInference::InferFftShape(
+                              operand, /*fft_type=*/FftType::RFFT, fft_length));
+  EXPECT_TRUE(ShapeUtil::Equal(inferred_shape, expected))
+      << "inferred: " << ShapeUtil::HumanString(inferred_shape)
+      << " expected: " << ShapeUtil::HumanString(expected);
+}
+
+TEST_F(ShapeInferenceTest, UnboundedFftWithIRFFT) {
+  TF_ASSERT_OK_AND_ASSIGN(const Shape operand, ParseShape("c128[2, <=5, ?]"));
+  const std::vector<int64_t> fft_length = {5, 10};
+  TF_ASSERT_OK_AND_ASSIGN(const Shape expected, ParseShape("f64[2, <=5, 10]"));
+  TF_ASSERT_OK_AND_ASSIGN(
+      const Shape inferred_shape,
+      ShapeInference::InferFftShape(operand, /*fft_type=*/FftType::IRFFT,
+                                    fft_length));
+  EXPECT_TRUE(ShapeUtil::Equal(inferred_shape, expected))
+      << "inferred: " << ShapeUtil::HumanString(inferred_shape)
+      << " expected: " << ShapeUtil::HumanString(expected);
+}
+
 TEST_F(ShapeInferenceTest, UnboundedGather) {
   TF_ASSERT_OK_AND_ASSIGN(const Shape operand, ParseShape("f32[3, 4, 2]"));
   TF_ASSERT_OK_AND_ASSIGN(const Shape start_indices,

From f30bdc86c168313211c03d44dbf2288c3afd3e19 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Tue, 21 May 2024 02:35:06 -0700
Subject: [PATCH 455/478] Automated Code Change

PiperOrigin-RevId: 635733471
---
 .../c/experimental/ops/gen/common/case_format.cc       |  3 +++
 .../c/experimental/ops/gen/common/case_format_test.cc  |  1 +
 tensorflow/c/experimental/ops/gen/common/controller.cc | 10 ++++++++--
 .../c/experimental/ops/gen/common/path_config.cc       |  2 ++
 .../c/experimental/ops/gen/common/source_code.cc       |  3 +++
 tensorflow/c/experimental/ops/gen/common/view_util.cc  |  2 ++
 6 files changed, 19 insertions(+), 2 deletions(-)

diff --git a/tensorflow/c/experimental/ops/gen/common/case_format.cc b/tensorflow/c/experimental/ops/gen/common/case_format.cc
index 9b8e955356db07..1e9d123005e8a4 100644
--- a/tensorflow/c/experimental/ops/gen/common/case_format.cc
+++ b/tensorflow/c/experimental/ops/gen/common/case_format.cc
@@ -14,6 +14,9 @@ limitations under the License.
 ==============================================================================*/
 #include "tensorflow/c/experimental/ops/gen/common/case_format.h"
 
+#include "tensorflow/core/platform/str_util.h"
+#include "tensorflow/core/platform/types.h"
+
 namespace tensorflow {
 namespace generator {
 
diff --git a/tensorflow/c/experimental/ops/gen/common/case_format_test.cc b/tensorflow/c/experimental/ops/gen/common/case_format_test.cc
index 37bc5be753fd64..302bcc42453169 100644
--- a/tensorflow/c/experimental/ops/gen/common/case_format_test.cc
+++ b/tensorflow/c/experimental/ops/gen/common/case_format_test.cc
@@ -15,6 +15,7 @@ limitations under the License.
 #include "tensorflow/c/experimental/ops/gen/common/case_format.h"
 
 #include "tensorflow/core/platform/test.h"
+#include "tensorflow/core/platform/types.h"
 
 namespace tensorflow {
 namespace generator {
diff --git a/tensorflow/c/experimental/ops/gen/common/controller.cc b/tensorflow/c/experimental/ops/gen/common/controller.cc
index a8e02f41011d32..cafb57c0919403 100644
--- a/tensorflow/c/experimental/ops/gen/common/controller.cc
+++ b/tensorflow/c/experimental/ops/gen/common/controller.cc
@@ -15,11 +15,17 @@ limitations under the License.
 #include "tensorflow/c/experimental/ops/gen/common/controller.h"
 
 #include "absl/strings/substitute.h"
+#include "tensorflow/c/experimental/ops/gen/common/path_config.h"
+#include "tensorflow/c/experimental/ops/gen/common/source_code.h"
+#include "tensorflow/c/experimental/ops/gen/model/op_spec.h"
+#include "tensorflow/core/framework/api_def.pb.h"
 #include "tensorflow/core/framework/op.h"
-#include "tensorflow/core/lib/io/path.h"
-#include "tensorflow/core/lib/strings/str_util.h"
+#include "tensorflow/core/framework/op_def.pb.h"
+#include "tensorflow/core/framework/op_gen_lib.h"
 #include "tensorflow/core/platform/env.h"
 #include "tensorflow/core/platform/logging.h"
+#include "tensorflow/core/platform/path.h"
+#include "tsl/platform/status.h"
 
 namespace tensorflow {
 namespace generator {
diff --git a/tensorflow/c/experimental/ops/gen/common/path_config.cc b/tensorflow/c/experimental/ops/gen/common/path_config.cc
index d9e3881bf15580..b8f84d5f31f4d3 100644
--- a/tensorflow/c/experimental/ops/gen/common/path_config.cc
+++ b/tensorflow/c/experimental/ops/gen/common/path_config.cc
@@ -16,7 +16,9 @@ limitations under the License.
 
 #include <iostream>
 
+#include "absl/strings/str_join.h"
 #include "tensorflow/core/lib/strings/str_util.h"
+#include "tensorflow/core/platform/types.h"
 
 namespace tensorflow {
 namespace generator {
diff --git a/tensorflow/c/experimental/ops/gen/common/source_code.cc b/tensorflow/c/experimental/ops/gen/common/source_code.cc
index ea4db53d167109..ea2b66fac7cd27 100644
--- a/tensorflow/c/experimental/ops/gen/common/source_code.cc
+++ b/tensorflow/c/experimental/ops/gen/common/source_code.cc
@@ -14,9 +14,12 @@ limitations under the License.
 ==============================================================================*/
 #include "tensorflow/c/experimental/ops/gen/common/source_code.h"
 
+#include "absl/strings/ascii.h"
 #include "absl/strings/match.h"
+#include "absl/strings/str_cat.h"
 #include "tensorflow/core/lib/strings/str_util.h"
 #include "tensorflow/core/platform/logging.h"
+#include "tensorflow/core/platform/stringpiece.h"
 
 namespace tensorflow {
 namespace generator {
diff --git a/tensorflow/c/experimental/ops/gen/common/view_util.cc b/tensorflow/c/experimental/ops/gen/common/view_util.cc
index a14c7e38b63b46..7c8717067b08fe 100644
--- a/tensorflow/c/experimental/ops/gen/common/view_util.cc
+++ b/tensorflow/c/experimental/ops/gen/common/view_util.cc
@@ -14,7 +14,9 @@ limitations under the License.
 ==============================================================================*/
 #include "tensorflow/c/experimental/ops/gen/common/view_util.h"
 
+#include "absl/strings/str_join.h"
 #include "absl/strings/substitute.h"
+#include "tensorflow/core/platform/types.h"
 
 namespace tensorflow {
 namespace generator {

From 104615d728a0eb136c7a72189702320f98f8fd23 Mon Sep 17 00:00:00 2001
From: Oleg Shyshkov <shyshkov@google.com>
Date: Tue, 21 May 2024 02:50:22 -0700
Subject: [PATCH 456/478] Re-land: [XLA:GPU] Store fusion_roots and
 fusion_heroes as HloInstructionAdaptor in HloFusionAnalysis.

Reverts 23ac61b9f8a787925a089772a70e288fe42a336c

PiperOrigin-RevId: 635737157
---
 third_party/xla/xla/service/gpu/BUILD         |  1 +
 third_party/xla/xla/service/gpu/fusions/BUILD |  6 ++
 .../xla/service/gpu/fusions/concatenate.cc    |  5 +-
 .../xla/xla/service/gpu/fusions/fusions.cc    |  7 +-
 .../fusions/in_place_dynamic_update_slice.h   |  4 +-
 .../in_place_dynamic_update_slice_mlir.cc     |  4 +-
 .../in_place_dynamic_update_slice_mlir.h      |  4 +-
 .../service/gpu/fusions/input_slices_mlir.cc  | 22 +++--
 .../xla/xla/service/gpu/fusions/loop_mlir.cc  | 14 +--
 .../fusions/mlir/computation_partitioner.cc   |  4 +-
 .../xla/xla/service/gpu/fusions/reduction.cc  | 17 ++--
 .../xla/service/gpu/fusions/reduction_base.cc |  8 +-
 .../gpu/fusions/reduction_base_test.cc        |  5 +-
 .../xla/service/gpu/fusions/reduction_mlir.cc | 12 ++-
 .../xla/xla/service/gpu/fusions/transpose.cc  | 14 ++-
 .../xla/service/gpu/fusions/transpose_mlir.cc | 11 ++-
 .../xla/service/gpu/hlo_fusion_analysis.cc    | 91 ++++++++++---------
 .../xla/xla/service/gpu/hlo_fusion_analysis.h | 35 ++++---
 .../service/gpu/model/coalescing_analysis.cc  | 10 +-
 .../gpu/model/fusion_analysis_cache_test.cc   |  9 +-
 20 files changed, 155 insertions(+), 128 deletions(-)

diff --git a/third_party/xla/xla/service/gpu/BUILD b/third_party/xla/xla/service/gpu/BUILD
index a739805c376c56..a925edb9986b2c 100644
--- a/third_party/xla/xla/service/gpu/BUILD
+++ b/third_party/xla/xla/service/gpu/BUILD
@@ -4436,6 +4436,7 @@ cc_library(
         "@com_google_absl//absl/log",
         "@com_google_absl//absl/log:check",
         "@com_google_absl//absl/strings:string_view",
+        "@com_google_absl//absl/types:span",
         "@llvm-project//llvm:Support",
     ],
 )
diff --git a/third_party/xla/xla/service/gpu/fusions/BUILD b/third_party/xla/xla/service/gpu/fusions/BUILD
index 3796e9a40f3bab..94ccd305743800 100644
--- a/third_party/xla/xla/service/gpu/fusions/BUILD
+++ b/third_party/xla/xla/service/gpu/fusions/BUILD
@@ -732,6 +732,7 @@ cc_library(
         "//xla/hlo/ir:hlo",
         "//xla/service:buffer_assignment",
         "//xla/service/gpu:hlo_fusion_analysis",
+        "//xla/service/gpu:hlo_traversal",
         "//xla/service/gpu:ir_emission_utils",
         "//xla/service/gpu:ir_emitter",
         "//xla/service/gpu:ir_emitter_context",
@@ -812,6 +813,7 @@ xla_cc_test(
         "//xla/service/gpu:gpu_device_info_for_tests",
         "//xla/service/gpu:hlo_fusion_analysis",
         "//xla/service/gpu:ir_emitter_context",
+        "//xla/service/gpu/model:indexing_analysis",
         "//xla/service/gpu/model:indexing_test_utils",
         "//xla/stream_executor:device_description",
         "//xla/tests:hlo_test_base",
@@ -833,6 +835,7 @@ cc_library(
         "//xla:shape_util",
         "//xla/hlo/ir:hlo",
         "//xla/service/gpu:hlo_fusion_analysis",
+        "//xla/service/gpu:hlo_traversal",
         "//xla/service/gpu:ir_emission_utils",
         "//xla/service/gpu:reduction_utils",
         "//xla/service/gpu/fusions/mlir:computation_partitioner",
@@ -845,6 +848,7 @@ cc_library(
         "@com_google_absl//absl/algorithm:container",
         "@com_google_absl//absl/container:flat_hash_map",
         "@com_google_absl//absl/container:flat_hash_set",
+        "@com_google_absl//absl/log:check",
         "@com_google_absl//absl/status",
         "@com_google_absl//absl/types:span",
         "@llvm-project//llvm:Support",
@@ -969,6 +973,7 @@ cc_library(
         ":fusion_emitter",
         ":tiling_util",
         "//xla:permutation_util",
+        "//xla:shape_util",
         "//xla:status",
         "//xla:util",
         "//xla/hlo/ir:hlo",
@@ -1058,6 +1063,7 @@ cc_library(
         "//xla:xla_data_proto_cc",
         "//xla/hlo/ir:hlo",
         "//xla/service/gpu:hlo_fusion_analysis",
+        "//xla/service/gpu:hlo_traversal",
         "//xla/service/gpu:launch_dimensions",
         "//xla/service/gpu/fusions/mlir:computation_partitioner",
         "//xla/service/gpu/fusions/mlir:elemental_hlo_to_mlir",
diff --git a/third_party/xla/xla/service/gpu/fusions/concatenate.cc b/third_party/xla/xla/service/gpu/fusions/concatenate.cc
index 865970112ff958..5963005a248d44 100644
--- a/third_party/xla/xla/service/gpu/fusions/concatenate.cc
+++ b/third_party/xla/xla/service/gpu/fusions/concatenate.cc
@@ -111,8 +111,9 @@ absl::Status ConcatenateFusion::EmitKernel(
       for (const auto& [output, root] :
            llvm::zip_equal(outputs, analysis_.fusion_roots())) {
         llvm_ir::IrArray::Index root_index = result_index.SourceIndexOfBitcast(
-            concat.shape(), root->shape(), builder);
-        TF_ASSIGN_OR_RETURN(auto generator, fused_emitter.GetGenerator(*root));
+            concat.shape(), root.shape(), builder);
+        TF_ASSIGN_OR_RETURN(auto generator,
+                            fused_emitter.GetGenerator(root.instruction()));
         TF_ASSIGN_OR_RETURN(llvm::Value * value, generator(root_index));
         output.EmitWriteArrayElement(root_index, value, builder);
       }
diff --git a/third_party/xla/xla/service/gpu/fusions/fusions.cc b/third_party/xla/xla/service/gpu/fusions/fusions.cc
index 3f524403bd70ce..c7b622ca87567c 100644
--- a/third_party/xla/xla/service/gpu/fusions/fusions.cc
+++ b/third_party/xla/xla/service/gpu/fusions/fusions.cc
@@ -77,7 +77,7 @@ bool IsParameterOrGteOfParameter(const HloInstruction* instr) {
 
 bool IsDynamicUpdateSliceFusion(const HloFusionAnalysis& analysis) {
   return absl::c_all_of(
-      analysis.fusion_root_adaptors(), [](const HloInstructionAdaptor& root) {
+      analysis.fusion_roots(), [](const HloInstructionAdaptor& root) {
         return root.opcode() == HloOpcode::kDynamicUpdateSlice ||
                (root.opcode() == HloOpcode::kBitcast &&
                 root.GetOperand(0).opcode() == HloOpcode::kDynamicUpdateSlice);
@@ -89,7 +89,8 @@ bool IsDynamicUpdateSliceFusion(const HloFusionAnalysis& analysis) {
 std::optional<absl::StatusOr<std::unique_ptr<FusionInterface>>>
 HloFusionInfo::GetCopyFusion() const {
   std::vector<BufferAllocation::Slice> src_buffers;
-  for (auto* root : analysis().fusion_roots()) {
+  for (const HloInstructionAdaptor& root_adaptor : analysis().fusion_roots()) {
+    const HloInstruction* root = &root_adaptor.instruction();
     if (root->opcode() != HloOpcode::kCopy ||
         root->operand(0)->opcode() != HloOpcode::kParameter ||
         !LayoutUtil::Equal(root->operand(0)->shape().layout(),
@@ -127,7 +128,7 @@ HloFusionInfo::GetCopyFusion() const {
 
 bool HloFusionInfo::CanEmitDynamicUpdateSliceInPlace() const {
   auto ret = CanEmitFusedDynamicUpdateSliceInPlaceForGpu(
-      instr_, buffer_assignment_, analysis().fusion_root_adaptors());
+      instr_, buffer_assignment_, analysis().fusion_roots());
   return ret.ok() && *ret;
 }
 
diff --git a/third_party/xla/xla/service/gpu/fusions/in_place_dynamic_update_slice.h b/third_party/xla/xla/service/gpu/fusions/in_place_dynamic_update_slice.h
index 2f1c12ca642ea2..741d6b339a6478 100644
--- a/third_party/xla/xla/service/gpu/fusions/in_place_dynamic_update_slice.h
+++ b/third_party/xla/xla/service/gpu/fusions/in_place_dynamic_update_slice.h
@@ -65,8 +65,8 @@ class InPlaceDynamicUpdateSliceFusion : public KernelFusionEmitterBase {
  public:
   explicit InPlaceDynamicUpdateSliceFusion(const HloFusionAnalysis& analysis)
       : analysis_(analysis),
-        dus_ops_(GetOutputDefiningDynamicUpdateSlices(
-            analysis.fusion_root_adaptors())) {}
+        dus_ops_(
+            GetOutputDefiningDynamicUpdateSlices(analysis.fusion_roots())) {}
   LaunchDimensions launch_dimensions() const override;
 
   std::optional<IndexingMap> ComputeThreadIdToOutputIndexing(
diff --git a/third_party/xla/xla/service/gpu/fusions/in_place_dynamic_update_slice_mlir.cc b/third_party/xla/xla/service/gpu/fusions/in_place_dynamic_update_slice_mlir.cc
index 02d4e327358481..b300bcf1a23d36 100644
--- a/third_party/xla/xla/service/gpu/fusions/in_place_dynamic_update_slice_mlir.cc
+++ b/third_party/xla/xla/service/gpu/fusions/in_place_dynamic_update_slice_mlir.cc
@@ -153,9 +153,9 @@ absl::Status MlirInPlaceDynamicUpdateSliceFusion::EmitEntryFunction(
               ProvideParameter(root_computation, dus_instr, kDUSUpdateIndex,
                                input_indices, call_targets, entry_function, b);
           // Handle bitcasts under the DUS.
-          if (dus_instr->shape() != root->shape()) {
+          if (dus_instr->shape() != root.shape()) {
             update_indices = ApplyAffineMap(
-                GetBitcastMap(dus_instr->shape(), root->shape(), b.getContext())
+                GetBitcastMap(dus_instr->shape(), root.shape(), b.getContext())
                     .GetAffineMap(),
                 update_indices, {}, b);
           }
diff --git a/third_party/xla/xla/service/gpu/fusions/in_place_dynamic_update_slice_mlir.h b/third_party/xla/xla/service/gpu/fusions/in_place_dynamic_update_slice_mlir.h
index efcfce8b2c6ede..b3086ea29af6a6 100644
--- a/third_party/xla/xla/service/gpu/fusions/in_place_dynamic_update_slice_mlir.h
+++ b/third_party/xla/xla/service/gpu/fusions/in_place_dynamic_update_slice_mlir.h
@@ -47,8 +47,8 @@ class MlirInPlaceDynamicUpdateSliceFusion : public MlirFusionEmitterBase {
   explicit MlirInPlaceDynamicUpdateSliceFusion(
       const HloFusionAnalysis& analysis)
       : analysis_(analysis),
-        dus_ops_(GetOutputDefiningDynamicUpdateSlices(
-            analysis.fusion_root_adaptors())) {}
+        dus_ops_(
+            GetOutputDefiningDynamicUpdateSlices(analysis.fusion_roots())) {}
 
   LaunchDimensions launch_dimensions() const override;
 
diff --git a/third_party/xla/xla/service/gpu/fusions/input_slices_mlir.cc b/third_party/xla/xla/service/gpu/fusions/input_slices_mlir.cc
index 4ac9b0427f3b71..806bb1d4328907 100644
--- a/third_party/xla/xla/service/gpu/fusions/input_slices_mlir.cc
+++ b/third_party/xla/xla/service/gpu/fusions/input_slices_mlir.cc
@@ -41,6 +41,7 @@ limitations under the License.
 #include "xla/service/gpu/fusions/mlir/computation_partitioner.h"
 #include "xla/service/gpu/fusions/mlir/elemental_hlo_to_mlir.h"
 #include "xla/service/gpu/fusions/mlir/ir/xla_gpu_ops.h"
+#include "xla/service/gpu/hlo_traversal.h"
 #include "xla/service/gpu/launch_dimensions.h"
 #include "xla/service/gpu/model/indexing_analysis.h"
 #include "xla/service/gpu/model/indexing_map.h"
@@ -69,11 +70,16 @@ MlirInputSlicesFusion::ComputeThreadIdToOutputIndexing(
 std::vector<mlir_converter::EpilogueSpecification>
 MlirInputSlicesFusion::GetEpilogues(const HloFusionInstruction& fusion,
                                     mlir::MLIRContext* mlir_context) const {
+  std::vector<const HloInstruction*> roots;
+  roots.reserve(analysis_.fusion_root_count());
+  for (const auto& root : analysis_.fusion_roots()) {
+    roots.push_back(&root.instruction());
+  }
+
   // We don't actually use epilogues here, but this is how we tell the base
   // class not to emit code for the slices.
   return {mlir_converter::EpilogueSpecification::FromOutputIndexing(
-      analysis_, analysis_.fusion_roots(), analysis_.fusion_roots(), *this,
-      mlir_context)};
+      analysis_, roots, roots, *this, mlir_context)};
 }
 
 LaunchDimensions MlirInputSlicesFusion::launch_dimensions() const {
@@ -94,7 +100,8 @@ absl::Status MlirInputSlicesFusion::EmitEntryFunction(
   builder.setInsertionPointToStart(entry_function.addEntryBlock());
 
   auto launch_dims = launch_dimensions();
-  const auto& shape = analysis_.fusion_roots().front()->operand(0)->shape();
+  const auto& shape =
+      analysis_.fusion_root(0).instruction().operand(0)->shape();
   auto input_indexing = GetDefaultThreadIdIndexingMap(
       launch_dims, unroll_factor_, shape, builder.getContext());
 
@@ -115,8 +122,8 @@ absl::Status MlirInputSlicesFusion::EmitEntryFunction(
         result_tensors.reserve(output_tensor_args.size());
 
         absl::flat_hash_map<const HloInstruction*, mlir::Value> input_values;
-        for (const auto* root : analysis_.fusion_roots()) {
-          const auto* arg = root->operand(0);
+        for (const HloInstructionAdaptor& root : analysis_.fusion_roots()) {
+          const auto* arg = root.instruction().operand(0);
           if (auto& value = input_values[arg]; !value) {
             value =
                 builder.create<PureCallOp>(call_targets(arg), input_operands)
@@ -137,8 +144,9 @@ absl::Status MlirInputSlicesFusion::EmitEntryFunction(
                 auto output_indices = mlir_converter::ApplyAffineMap(
                     output_indexing->GetAffineMap(), dim_values, symbol_values,
                     then_builder);
-                const auto* arg =
-                    analysis_.fusion_roots()[output_index]->operand(0);
+                const auto* arg = analysis_.fusion_root(output_index)
+                                      .instruction()
+                                      .operand(0);
                 auto inserted = then_builder.create<mlir::tensor::InsertOp>(
                     input_values[arg], output, output_indices);
                 then_builder.create<mlir::scf::YieldOp>(inserted.getResult());
diff --git a/third_party/xla/xla/service/gpu/fusions/loop_mlir.cc b/third_party/xla/xla/service/gpu/fusions/loop_mlir.cc
index d52043f117efec..7bf0df17db3aec 100644
--- a/third_party/xla/xla/service/gpu/fusions/loop_mlir.cc
+++ b/third_party/xla/xla/service/gpu/fusions/loop_mlir.cc
@@ -65,7 +65,7 @@ std::optional<IndexingMap> MlirLoopFusion::ComputeThreadIdToOutputIndexing(
   auto launch_dims = launch_dimensions();
   return GetDefaultThreadIdIndexingMap(
       launch_dims, config_.unroll_factor,
-      GetIndexShape(analysis_.fusion_roots()[root_index]->shape()), ctx);
+      GetIndexShape(analysis_.fusion_root(root_index).shape()), ctx);
 }
 
 std::optional<IndexingMap> MlirLoopFusion::ComputeThreadIdToInputIndexing(
@@ -93,8 +93,8 @@ std::optional<IndexingMap> MlirLoopFusion::ComputeThreadIdToInputIndexing(
 
 LaunchDimensions MlirLoopFusion::launch_dimensions() const {
   return CalculateLaunchDimensions(
-      GetIndexShape(analysis_.fusion_roots()[0]->shape()),
-      analysis_.device_info(), config_);
+      GetIndexShape(analysis_.fusion_root(0).shape()), analysis_.device_info(),
+      config_);
 }
 
 absl::Status MlirLoopFusion::EmitEntryFunction(
@@ -113,13 +113,13 @@ absl::Status MlirLoopFusion::EmitEntryFunction(
   auto output_tensor_args =
       entry_function.getArguments().drop_front(num_inputs);
   llvm::SmallVector<const Shape*> result_shapes;
-  for (const auto* root : analysis_.fusion_roots()) {
-    if (root->shape().IsTuple()) {
-      for (const auto& shape : root->shape().tuple_shapes()) {
+  for (const HloInstructionAdaptor& root : analysis_.fusion_roots()) {
+    if (root.shape().IsTuple()) {
+      for (const auto& shape : root.shape().tuple_shapes()) {
         result_shapes.push_back(&shape);
       }
     } else {
-      result_shapes.push_back(&root->shape());
+      result_shapes.push_back(&root.shape());
     }
   }
 
diff --git a/third_party/xla/xla/service/gpu/fusions/mlir/computation_partitioner.cc b/third_party/xla/xla/service/gpu/fusions/mlir/computation_partitioner.cc
index 97c3cba8b3892f..20bb07a321019d 100644
--- a/third_party/xla/xla/service/gpu/fusions/mlir/computation_partitioner.cc
+++ b/third_party/xla/xla/service/gpu/fusions/mlir/computation_partitioner.cc
@@ -131,11 +131,11 @@ EpilogueSpecification EpilogueSpecification::FromOutputIndexing(
       root_to_hero;
   for (auto [root, hero] :
        llvm::zip(analysis.fusion_roots(), analysis.fusion_heroes())) {
-    root_to_hero[root] = hero;
+    root_to_hero[&root.instruction()] = &hero.instruction();
   }
   absl::flat_hash_map<const HloInstruction*, int> root_to_index;
   for (auto [index, root] : llvm::enumerate(analysis.fusion_roots())) {
-    root_to_index[root] = root_to_index.size();
+    root_to_index[&root.instruction()] = root_to_index.size();
   }
 
   result.root_indexing.reserve(roots.size());
diff --git a/third_party/xla/xla/service/gpu/fusions/reduction.cc b/third_party/xla/xla/service/gpu/fusions/reduction.cc
index 9856b45a516d8a..83093677aabf26 100644
--- a/third_party/xla/xla/service/gpu/fusions/reduction.cc
+++ b/third_party/xla/xla/service/gpu/fusions/reduction.cc
@@ -53,6 +53,7 @@ limitations under the License.
 #include "xla/service/gpu/fusions/thunk_util.h"
 #include "xla/service/gpu/fusions/tiling_util.h"
 #include "xla/service/gpu/hlo_fusion_analysis.h"
+#include "xla/service/gpu/hlo_traversal.h"
 #include "xla/service/gpu/ir_emission_utils.h"
 #include "xla/service/gpu/ir_emitter_context.h"
 #include "xla/service/gpu/ir_emitter_nested.h"
@@ -147,10 +148,10 @@ class ReductionEmitter {
         index_ty_(GetIndexType(fusion, reduction_codegen_info.GetTiling(),
                                elemental_emitter_.builder())) {
     for (auto hero : analysis.fusion_heroes()) {
-      if (hero->opcode() == HloOpcode::kReduce) {
-        for (int i = 0; i < hero->operand_count() / 2; ++i) {
+      if (hero.opcode() == HloOpcode::kReduce) {
+        for (int i = 0; i < hero.instruction().operand_count() / 2; ++i) {
           CHECK(LayoutUtil::IsMonotonicWithDim0Major(
-              hero->operand(i)->shape().layout()))
+              hero.instruction().operand(i)->shape().layout()))
               << "reduction-layout-normalizer must run before code generation";
         }
       }
@@ -1012,10 +1013,10 @@ absl::StatusOr<FusionEmissionResult> ReductionEmitter::EmitInitializers() {
         return absl::OkStatus();
       }));
 
-  absl::Span<const HloInstruction* const> fusion_roots =
+  absl::Span<HloInstructionAdaptor const> fusion_roots =
       analysis_.fusion_roots();
   for (int i = 0; i < fusion_roots.size(); ++i) {
-    const HloInstruction* fusion_root = fusion_roots[i];
+    const HloInstruction* fusion_root = &fusion_roots[i].instruction();
 
     if (IsReductionFromOrToContiguousDimensions(*fusion_root)) {
       TF_ASSIGN_OR_RETURN(
@@ -1046,9 +1047,9 @@ absl::Status ReductionEmitter::EmitKernel(
   ReductionOutputMap result_ir_arrays;
 
   int ir_arrays_idx = 0;
-  for (const HloInstruction* root : analysis_.fusion_roots()) {
-    int get_num_results = GetNumOutputs(root->shape());
-    result_ir_arrays[root] =
+  for (const HloInstructionAdaptor& root : analysis_.fusion_roots()) {
+    int get_num_results = GetNumOutputs(root.shape());
+    result_ir_arrays[&root.instruction()] =
         absl::MakeSpan(outputs).subspan(ir_arrays_idx, get_num_results);
     ir_arrays_idx += get_num_results;
   }
diff --git a/third_party/xla/xla/service/gpu/fusions/reduction_base.cc b/third_party/xla/xla/service/gpu/fusions/reduction_base.cc
index 93e4c65852089a..00057e7528bc2f 100644
--- a/third_party/xla/xla/service/gpu/fusions/reduction_base.cc
+++ b/third_party/xla/xla/service/gpu/fusions/reduction_base.cc
@@ -134,7 +134,7 @@ ReductionGroups GroupDisjointReductions(const HloFusionAnalysis& analysis,
     disjoint_sets[root].Get() = root;
     reachable_outputs[root].insert(root);
     result.is_reduction_root.push_back(
-        IsRealReductionHero(root.instruction(), *hero));
+        IsRealReductionHero(root.instruction(), hero.instruction()));
     if (result.is_reduction_root.back()) {
       roots_with_reduction.insert(root);
     } else if (first_non_reduction_root) {
@@ -337,7 +337,7 @@ std::optional<IndexingMap> ReductionInfo::ComputeThreadIdToOutputIndexing(
     auto map = ComposeIndexingMaps(
         GetIndexingMapForTiling(tiling_, ctx),
         GetBitcastMap(tiling_.GetXlaShape(),
-                      analysis_.fusion_roots()[root_index]->shape(), ctx));
+                      analysis_.fusion_root(root_index).shape(), ctx));
     AddGroupIdConstraint(map, root_index, ctx);
     return map;
   }
@@ -431,8 +431,8 @@ std::optional<IndexingMap> ReductionInfo::ComputeThreadIdToInputIndexing(
   if (!groups_.is_reduction_root[root_index]) {
     return ComposeIndexingMaps(
         *ComputeThreadIdToOutputIndexing(root_index, ctx),
-        *ComputeOutputToInputIndexing(analysis_.fusion_roots()[root_index], 0,
-                                      ctx)
+        *ComputeOutputToInputIndexing(
+             &analysis_.fusion_root(root_index).instruction(), 0, ctx)
              .indexing_maps[hero_operand_index]
              .begin());
   }
diff --git a/third_party/xla/xla/service/gpu/fusions/reduction_base_test.cc b/third_party/xla/xla/service/gpu/fusions/reduction_base_test.cc
index 6a43851f555bc7..674765142cf4fe 100644
--- a/third_party/xla/xla/service/gpu/fusions/reduction_base_test.cc
+++ b/third_party/xla/xla/service/gpu/fusions/reduction_base_test.cc
@@ -26,6 +26,7 @@ limitations under the License.
 #include "xla/service/gpu/gpu_device_info_for_tests.h"
 #include "xla/service/gpu/hlo_fusion_analysis.h"
 #include "xla/service/gpu/ir_emitter_context.h"
+#include "xla/service/gpu/model/indexing_analysis.h"
 #include "xla/service/gpu/model/indexing_test_utils.h"
 #include "xla/stream_executor/device_description.h"
 #include "xla/tests/hlo_test_base.h"
@@ -492,8 +493,8 @@ TEST_F(ReductionTest, TwoGroups) {
   FakeReductionFusion fusion(analysis);
 
   EXPECT_THAT(fusion.reduction_info().GetGroups().grouped_roots,
-              ElementsAre(ElementsAre(analysis.fusion_roots()[0]),
-                          ElementsAre(analysis.fusion_roots()[1])));
+              ElementsAre(ElementsAre(&analysis.fusion_root(0).instruction()),
+                          ElementsAre(&analysis.fusion_root(1).instruction())));
 }
 
 TEST_F(ReductionTest, OneGroup) {
diff --git a/third_party/xla/xla/service/gpu/fusions/reduction_mlir.cc b/third_party/xla/xla/service/gpu/fusions/reduction_mlir.cc
index 7a90f21c8d787b..c36bc3b61c0cae 100644
--- a/third_party/xla/xla/service/gpu/fusions/reduction_mlir.cc
+++ b/third_party/xla/xla/service/gpu/fusions/reduction_mlir.cc
@@ -47,10 +47,12 @@ limitations under the License.
 #include "xla/service/gpu/fusions/mlir/type_util.h"
 #include "xla/service/gpu/fusions/reduction_base.h"
 #include "xla/service/gpu/hlo_fusion_analysis.h"
+#include "xla/service/gpu/hlo_traversal.h"
 #include "xla/service/gpu/ir_emission_utils.h"
 #include "xla/service/gpu/model/indexing_analysis.h"
 #include "xla/service/gpu/model/indexing_map.h"
 #include "xla/service/gpu/reduction_utils.h"
+#include "xla/shape.h"
 #include "xla/shape_util.h"
 
 namespace xla {
@@ -80,9 +82,9 @@ struct MlirReductionFusion::EmitterState {
         computation(computations.FindPartitionedComputation(
             fusion.fused_instructions_computation())) {
     int index = 0;
-    for (auto root : owner.analysis().fusion_roots()) {
-      fusion_result_index_starts[root] = index;
-      index += root->shape().IsTuple() ? root->shape().tuple_shapes_size() : 1;
+    for (const auto& root : owner.analysis().fusion_roots()) {
+      fusion_result_index_starts[&root.instruction()] = index;
+      index += root.shape().IsTuple() ? root.shape().tuple_shapes_size() : 1;
     }
   }
 
@@ -132,9 +134,11 @@ MlirReductionFusion::MlirReductionFusion(const HloFusionAnalysis& analysis)
   reduction_roots_.resize(num_groups);
 
   absl::flat_hash_set<const HloInstruction*> seen_heroes;
-  for (auto [root, hero, is_reduction, group_id] :
+  for (auto [root_adaptor, hero_adaptor, is_reduction, group_id] :
        llvm::zip(analysis.fusion_roots(), analysis.fusion_heroes(),
                  groups.is_reduction_root, groups.group_id_per_root)) {
+    const HloInstruction* root = &root_adaptor.instruction();
+    const HloInstruction* hero = &hero_adaptor.instruction();
     if (is_reduction) {
       if (seen_heroes.insert(hero).second) {
         reduction_heroes_[group_id].push_back(hero);
diff --git a/third_party/xla/xla/service/gpu/fusions/transpose.cc b/third_party/xla/xla/service/gpu/fusions/transpose.cc
index 319941d8e5d59f..ee66e531a9cf7d 100644
--- a/third_party/xla/xla/service/gpu/fusions/transpose.cc
+++ b/third_party/xla/xla/service/gpu/fusions/transpose.cc
@@ -48,6 +48,7 @@ limitations under the License.
 #include "xla/service/llvm_ir/fused_ir_emitter.h"
 #include "xla/service/llvm_ir/ir_array.h"
 #include "xla/service/llvm_ir/llvm_util.h"
+#include "xla/shape_util.h"
 #include "xla/status.h"
 #include "xla/util.h"
 
@@ -128,7 +129,8 @@ TransposeFusion::TransposeFusion(const se::DeviceDescription& gpu_device_info,
           ComputeTransposeTiling(gpu_device_info, analysis.tiled_transpose())) {
   for (auto [root, hero] :
        llvm::zip(analysis_.fusion_roots(), analysis_.fusion_heroes())) {
-    if (auto transpose = GetDescriptionForTiledTransposeEmitter(*root, *hero)) {
+    if (auto transpose = GetDescriptionForTiledTransposeEmitter(
+            root.instruction(), hero.instruction())) {
       permutation_ = transpose->permutation;
       break;
     }
@@ -165,18 +167,20 @@ absl::Status TransposeFusion::EmitKernel(IrEmitterContext& ir_emitter_context,
 
   for (const auto& [output_idx, root] : llvm::enumerate(hlo_roots)) {
     const auto& hero = analysis_.fusion_hero(output_idx).instruction();
-    auto transpose_descr = GetDescriptionForTiledTransposeEmitter(*root, hero);
+    auto transpose_descr =
+        GetDescriptionForTiledTransposeEmitter(root.instruction(), hero);
     if (transpose_descr.has_value()) {
       auto iterator_inserted = transposes_to_roots.insert(std::make_pair(
           &hero, std::vector<std::pair<int64_t, const HloInstruction*>>{
-                     {output_idx, root}}));
+                     {output_idx, &root.instruction()}}));
       if (iterator_inserted.second) {
         transposes.push_back(*transpose_descr);
       } else {
-        iterator_inserted.first->second.push_back({output_idx, root});
+        iterator_inserted.first->second.push_back(
+            {output_idx, &root.instruction()});
       }
     } else {
-      extra_outputs.push_back({output_idx, root});
+      extra_outputs.push_back({output_idx, &root.instruction()});
     }
   }
 
diff --git a/third_party/xla/xla/service/gpu/fusions/transpose_mlir.cc b/third_party/xla/xla/service/gpu/fusions/transpose_mlir.cc
index 0927dbbcaffc49..23cf134412f0ff 100644
--- a/third_party/xla/xla/service/gpu/fusions/transpose_mlir.cc
+++ b/third_party/xla/xla/service/gpu/fusions/transpose_mlir.cc
@@ -91,15 +91,16 @@ MlirTransposeFusion::MlirTransposeFusion(const HloFusionAnalysis& analysis)
   int max_element_bytes = 0;
   for (auto [root, hero] :
        llvm::zip(analysis_.fusion_roots(), analysis_.fusion_heroes())) {
-    if (auto transpose = GetDescriptionForTiledTransposeEmitter(*root, *hero)) {
-      transposes_to_tile.insert(hero);
-      shmem_transpose_roots_.push_back(root);
-      int size = primitive_util::ByteWidth(hero->shape().element_type());
+    if (auto transpose = GetDescriptionForTiledTransposeEmitter(
+            root.instruction(), hero.instruction())) {
+      transposes_to_tile.insert(&hero.instruction());
+      shmem_transpose_roots_.push_back(&root.instruction());
+      int size = primitive_util::ByteWidth(hero.shape().element_type());
       max_element_bytes = std::max(max_element_bytes, size);
       shmem_usage += kBaseBlockSize * (kBaseBlockSize + 1) * size;
       shmem_transpose_root_indices_.push_back(index);
     } else {
-      side_output_roots_.push_back(root);
+      side_output_roots_.push_back(&root.instruction());
       side_output_root_indices_.push_back(index);
     }
     ++index;
diff --git a/third_party/xla/xla/service/gpu/hlo_fusion_analysis.cc b/third_party/xla/xla/service/gpu/hlo_fusion_analysis.cc
index d6460c00861fa6..765fd6b80354d8 100644
--- a/third_party/xla/xla/service/gpu/hlo_fusion_analysis.cc
+++ b/third_party/xla/xla/service/gpu/hlo_fusion_analysis.cc
@@ -23,9 +23,11 @@ limitations under the License.
 #include <vector>
 
 #include "absl/algorithm/container.h"
+#include "absl/container/inlined_vector.h"
 #include "absl/log/check.h"
 #include "absl/log/log.h"
 #include "absl/strings/string_view.h"
+#include "absl/types/span.h"
 #include "llvm/ADT/STLExtras.h"
 #include "xla/hlo/ir/hlo_instruction.h"
 #include "xla/hlo/ir/hlo_instructions.h"
@@ -46,16 +48,19 @@ namespace {
 
 // Returns true if the fusion output contains non-strided slices only.
 bool IsInputFusibleNonStridedSlices(
-    const std::vector<const HloInstruction*>& fusion_roots) {
-  return absl::c_all_of(fusion_roots, IsSliceWithUnitStrides);
+    const absl::Span<const HloInstructionAdaptor> fusion_roots) {
+  return absl::c_all_of(fusion_roots, [&](const HloInstructionAdaptor& root) {
+    return IsSliceWithUnitStrides(&root.instruction());
+  });
 }
 
 // Returns true if all slice inputs in a tuple are equal (ignoring type).
 bool AllSliceInputsAreCompatible(
-    const std::vector<const HloInstruction*>& fusion_roots) {
-  const Shape& first_slice_operand_shape = fusion_roots[0]->operand(0)->shape();
-  return absl::c_all_of(fusion_roots, [&](const HloInstruction* slice) {
-    return ShapeUtil::EqualIgnoringElementType(slice->operand(0)->shape(),
+    const absl::Span<const HloInstructionAdaptor> fusion_roots) {
+  const Shape& first_slice_operand_shape =
+      fusion_roots[0].GetOperand(0).shape();
+  return absl::c_all_of(fusion_roots, [&](const HloInstructionAdaptor& slice) {
+    return ShapeUtil::EqualIgnoringElementType(slice.GetOperand(0).shape(),
                                                first_slice_operand_shape);
   });
 }
@@ -66,13 +71,14 @@ bool AllSliceInputsAreCompatible(
 //   * Either the root has a traspose hero with the same normalized dimensions
 //   * Or the root output shape is equal to the the transpose input shape
 std::optional<TransposeDescription> FindConsistentTransposeHero(
-    const std::vector<const HloInstruction*>& hlo_roots,
-    const std::vector<const HloInstruction*>& heroes) {
+    const absl::InlinedVector<HloInstructionAdaptor, 2>& hlo_roots,
+    const absl::InlinedVector<HloInstructionAdaptor, 2>& heroes) {
   std::optional<TransposeDescription> tiled_transpose_hero;
   std::vector<const HloInstruction*> non_transpose_roots;
 
   for (auto [root, hero] : llvm::zip(hlo_roots, heroes)) {
-    if (auto tr = GetDescriptionForTiledTransposeEmitter(*root, *hero)) {
+    if (auto tr = GetDescriptionForTiledTransposeEmitter(root.instruction(),
+                                                         hero.instruction())) {
       if (!tiled_transpose_hero) {
         // First transpose hero found.
         tiled_transpose_hero = tr;
@@ -81,7 +87,7 @@ std::optional<TransposeDescription> FindConsistentTransposeHero(
         return std::nullopt;
       }
     } else {
-      non_transpose_roots.push_back(root);
+      non_transpose_roots.push_back(&root.instruction());
     }
   }
 
@@ -100,14 +106,14 @@ std::optional<TransposeDescription> FindConsistentTransposeHero(
   return tiled_transpose_hero;
 }
 
-int SmallestBitWidth(absl::Span<const HloInstruction* const> args) {
+int SmallestBitWidth(absl::Span<HloInstructionAdaptor const> args) {
   int bits = std::numeric_limits<int>::max();
-  for (const HloInstruction* operand : args) {
-    if (!operand->shape().IsArray()) continue;
+  for (const HloInstructionAdaptor& operand : args) {
+    if (!operand.shape().IsArray()) continue;
     bits = std::min(
-        bits, operand->shape().element_type() == PRED
+        bits, operand.shape().element_type() == PRED
                   ? 8
-                  : primitive_util::BitWidth(operand->shape().element_type()));
+                  : primitive_util::BitWidth(operand.shape().element_type()));
   }
   return bits;
 }
@@ -117,8 +123,8 @@ int SmallestBitWidth(absl::Span<const HloInstruction* const> args) {
 HloFusionAnalysis::HloFusionAnalysis(
     FusionBackendConfig fusion_backend_config,
     std::unique_ptr<HloFusionAdaptor> fusion,
-    std::vector<const HloInstruction*> fusion_roots,
-    std::vector<const HloInstruction*> fusion_heroes,
+    absl::InlinedVector<HloInstructionAdaptor, 2> fusion_roots,
+    absl::InlinedVector<HloInstructionAdaptor, 2> fusion_heroes,
     const se::DeviceDescription* device_info,
     std::optional<TransposeDescription> tiled_transpose,
     HloFusionAnalysis::InputOutputInfo input_output_info)
@@ -135,16 +141,15 @@ HloFusionAnalysis HloFusionAnalysis::Create(
     FusionBackendConfig backend_config,
     std::unique_ptr<HloFusionAdaptor> fusion,
     const se::DeviceDescription* device_info) {
-  std::vector<const HloInstruction*> roots;
-  std::vector<const HloInstruction*> heroes;
-  for (auto root : fusion->GetRoots()) {
-    roots.push_back(&root.instruction());
-    heroes.push_back(&FindNonTrivialHero(root).instruction());
+  absl::InlinedVector<HloInstructionAdaptor, 2> roots = fusion->GetRoots();
+  absl::InlinedVector<HloInstructionAdaptor, 2> heroes;
+  for (auto root : roots) {
+    heroes.push_back(FindNonTrivialHero(root));
   }
 
-  std::vector<const HloInstruction*> fusion_arguments;
-  FindFusionArguments(*fusion, [&](auto argument) {
-    fusion_arguments.push_back(&argument.instruction());
+  absl::InlinedVector<HloInstructionAdaptor, 2> fusion_arguments;
+  FindFusionArguments(*fusion, [&](const HloInstructionAdaptor& argument) {
+    fusion_arguments.push_back(argument);
   });
 
   InputOutputInfo input_output_info{
@@ -179,15 +184,15 @@ bool HloFusionAnalysis::HasConsistentTransposeHeros() const {
 }
 
 static bool UseConcatenateFusion(
-    const std::vector<const HloInstruction*>& roots,
-    const std::vector<const HloInstruction*>& heroes) {
+    absl::Span<const HloInstructionAdaptor> roots,
+    absl::Span<const HloInstructionAdaptor> heroes) {
   if (heroes.size() != 1) return false;
-  if (heroes.front()->opcode() != HloOpcode::kConcatenate) return false;
+  if (heroes.front().opcode() != HloOpcode::kConcatenate) return false;
   // The concat emitter does not support multiple outputs yet. TODO(csigg): fix.
-  if (roots.front()->shape().IsTuple()) return false;
+  if (roots.front().shape().IsTuple()) return false;
   // Limit the number of operands because the concat emitter produces code for
   // each operand, hurting occupancy.
-  if (heroes.front()->operand_count() > 4) return false;
+  if (heroes.front().instruction().operand_count() > 4) return false;
   // The loop emitter is faster when warp divergence and occupancy are both low.
   // TODO(csigg): exclude this case.
   return true;
@@ -221,30 +226,30 @@ HloFusionAnalysis::EmitterFusionKind HloFusionAnalysis::GetEmitterFusionKind()
     return EmitterFusionKind::kLoop;
   }
 
-  const HloInstruction* first_reduce_hero = nullptr;
+  std::optional<HloInstructionAdaptor> first_reduce_hero;
   for (auto [root, hero] : llvm::zip(fusion_roots_, fusion_heroes_)) {
-    if (IsRealReductionHero(*root, *hero)) {
+    if (IsRealReductionHero(root.instruction(), hero.instruction())) {
       first_reduce_hero = hero;
       break;
     }
   }
-  if (first_reduce_hero != nullptr) {
+  if (first_reduce_hero.has_value()) {
     bool valid_shapes = true;
-    Shape hero_operand_shape = first_reduce_hero->operand(0)->shape();
+    Shape hero_operand_shape = first_reduce_hero->GetOperand(0).shape();
     for (auto [root, hero] : llvm::zip(fusion_roots_, fusion_heroes_)) {
-      if (root == first_reduce_hero) {
+      if (root == *first_reduce_hero) {
         continue;
       }
-      if (!IsRealReductionHero(*root, *hero)) {
+      if (!IsRealReductionHero(root.instruction(), hero.instruction())) {
         // Needs to have a compatible shape to the reduce operand (compatible
         // meaning same number of elements).
-        if (ShapeUtil::ElementsIn(root->shape()) !=
+        if (ShapeUtil::ElementsIn(root.shape()) !=
             ShapeUtil::ElementsIn(hero_operand_shape)) {
           valid_shapes = false;
           break;
         }
-      } else if (!AreReductionsMultiOutputFusionCompatible(hero,
-                                                           first_reduce_hero)) {
+      } else if (!AreReductionsMultiOutputFusionCompatible(
+                     &hero.instruction(), &first_reduce_hero->instruction())) {
         valid_shapes = false;
         break;
       }
@@ -267,7 +272,7 @@ HloFusionAnalysis::EmitterFusionKind HloFusionAnalysis::GetEmitterFusionKind()
     return EmitterFusionKind::kLoop;
   }
 
-  if (fusion_roots_[0]->opcode() == HloOpcode::kScatter) {
+  if (fusion_roots_[0].opcode() == HloOpcode::kScatter) {
     return EmitterFusionKind::kScatter;
   }
 
@@ -282,15 +287,15 @@ const HloInstruction* HloFusionAnalysis::FindHeroReduction() const {
   if (GetEmitterFusionKind() != EmitterFusionKind::kReduction) {
     return nullptr;
   }
-  auto roots = fusion_roots();
+  const auto& roots = fusion_roots();
   CHECK(!roots.empty());
   // We always use the first reduce root that triggers unnested reduction
   // emitter as the hero reduction, since all the reductions are required to
   // have the same shape and layout as verified by
   // `IsFusedReductionOutputConsistent()`.
   for (auto [root, hero] : llvm::zip(roots, fusion_heroes_)) {
-    if (IsRealReductionHero(*root, *hero)) {
-      return hero;
+    if (IsRealReductionHero(root.instruction(), hero.instruction())) {
+      return &hero.instruction();
     }
   }
   LOG(FATAL) << "Did not find a hero reduction";
diff --git a/third_party/xla/xla/service/gpu/hlo_fusion_analysis.h b/third_party/xla/xla/service/gpu/hlo_fusion_analysis.h
index c384a2e9eaf52c..146224b394579f 100644
--- a/third_party/xla/xla/service/gpu/hlo_fusion_analysis.h
+++ b/third_party/xla/xla/service/gpu/hlo_fusion_analysis.h
@@ -63,27 +63,19 @@ class HloFusionAnalysis {
 
   const HloFusionAdaptor& fusion() const { return *fusion_; }
 
-  const std::vector<const HloInstruction*>& fusion_roots() const {
+  const absl::InlinedVector<HloInstructionAdaptor, 2>& fusion_roots() const {
     return fusion_roots_;
   }
-
-  // TODO(b/336597139): Merge with fusion_roots() and only return
-  // HloInstructionAdaptor. This function is added temporarily to make
-  // transition easier and avoid breaking the existing code.
-  absl::InlinedVector<HloInstructionAdaptor, 2> fusion_root_adaptors() const {
-    return fusion_->GetRoots();
-  }
-
   HloInstructionAdaptor fusion_root(int64_t i) const {
-    return HloInstructionAdaptor(*fusion_roots_[i], fusion_.get());
+    return fusion_roots_[i];
   }
   int64_t fusion_root_count() const { return fusion_roots_.size(); }
 
-  const std::vector<const HloInstruction*>& fusion_heroes() const {
+  const absl::InlinedVector<HloInstructionAdaptor, 2>& fusion_heroes() const {
     return fusion_heroes_;
   }
   HloInstructionAdaptor fusion_hero(int64_t i) const {
-    return HloInstructionAdaptor(*fusion_heroes_[i], fusion_.get());
+    return fusion_heroes_[i];
   }
   int64_t fusion_hero_count() const { return fusion_heroes_.size(); }
 
@@ -113,8 +105,8 @@ class HloFusionAnalysis {
  private:
   HloFusionAnalysis(FusionBackendConfig fusion_backend_config,
                     std::unique_ptr<HloFusionAdaptor> fusion,
-                    std::vector<const HloInstruction*> fusion_roots,
-                    std::vector<const HloInstruction*> fusion_heroes,
+                    absl::InlinedVector<HloInstructionAdaptor, 2> fusion_roots,
+                    absl::InlinedVector<HloInstructionAdaptor, 2> fusion_heroes,
                     const se::DeviceDescription* device_info,
                     std::optional<TransposeDescription> tiled_transpose,
                     InputOutputInfo input_output_info);
@@ -122,11 +114,18 @@ class HloFusionAnalysis {
   bool HasConsistentTransposeHeros() const;
 
   FusionBackendConfig fusion_backend_config_;
+
+  // Owning pointer to the fusion adaptor object.
   std::unique_ptr<HloFusionAdaptor> fusion_;
-  // TODO(shyshkov): Change fusion_roots_ and fusion_heroes_ to store
-  // HloInstructionAdaptor.
-  std::vector<const HloInstruction*> fusion_roots_;
-  std::vector<const HloInstruction*> fusion_heroes_;
+
+  // A list of all roots of the fusion. The instruction adaptors have `fusion_`
+  // as their parent and should not outlive `fusion_`.
+  absl::InlinedVector<HloInstructionAdaptor, 2> fusion_roots_;
+
+  // A list of all heroes of the fusion. The instruction adaptors have `fusion_`
+  // as their parent and should not outlive `fusion_`.
+  absl::InlinedVector<HloInstructionAdaptor, 2> fusion_heroes_;
+
   const se::DeviceDescription* device_info_;
   std::optional<TransposeDescription> tiled_transpose_;
   InputOutputInfo input_output_info_;
diff --git a/third_party/xla/xla/service/gpu/model/coalescing_analysis.cc b/third_party/xla/xla/service/gpu/model/coalescing_analysis.cc
index 61afe6bba95be5..c391ade9379383 100644
--- a/third_party/xla/xla/service/gpu/model/coalescing_analysis.cc
+++ b/third_party/xla/xla/service/gpu/model/coalescing_analysis.cc
@@ -158,8 +158,8 @@ std::optional<GroupedByOpIndexingMap> GetThreadIdToInputMemoryLayoutsMaps(
   for (const auto& [root_index, hero] :
        llvm::enumerate(fusion_analysis.fusion_heroes())) {
     for (const auto& [hero_operand_index, hero_operand] :
-         llvm::enumerate(hero->operands())) {
-      if (hero_operand->shape().rank() == 0) {
+         llvm::enumerate(hero.GetOperands())) {
+      if (hero_operand.shape().rank() == 0) {
         continue;
       }
       // Compute thread ID -> hero operand indexing map.
@@ -170,11 +170,9 @@ std::optional<GroupedByOpIndexingMap> GetThreadIdToInputMemoryLayoutsMaps(
         return std::nullopt;
       }
       // Compute indexing from output to inputs for logical layout.
-      HloInstructionAdaptor hero_operand_adaptor(*hero_operand,
-                                                 &fusion_adaptor);
       GroupedByOpIndexingMap instr_indexing_keyed_by_operands =
-          ComputeGroupedOutputToInputIndexing(
-              fusion_adaptor, hero_operand_adaptor, mlir_context);
+          ComputeGroupedOutputToInputIndexing(fusion_adaptor, hero_operand,
+                                              mlir_context);
       // For every operand compute thread ID -> physical layout of operand
       // indexing map.
       for (const HloInstruction* operand : operands) {
diff --git a/third_party/xla/xla/service/gpu/model/fusion_analysis_cache_test.cc b/third_party/xla/xla/service/gpu/model/fusion_analysis_cache_test.cc
index d507a88a4d803b..1750c38c3c41ae 100644
--- a/third_party/xla/xla/service/gpu/model/fusion_analysis_cache_test.cc
+++ b/third_party/xla/xla/service/gpu/model/fusion_analysis_cache_test.cc
@@ -56,18 +56,15 @@ TEST_F(FusionAnalysisCacheTest, CachesAndInvalidates) {
   auto* negate = computation->GetInstructionWithName("n0");
   auto* fusion = module->entry_computation()->root_instruction();
 
-  EXPECT_THAT(cache_.Get(*fusion).fusion_roots(),
-              ::testing::ElementsAre(negate));
+  EXPECT_EQ(&cache_.Get(*fusion).fusion_root(0).instruction(), negate);
 
   computation->set_root_instruction(broadcast);
 
-  EXPECT_THAT(cache_.Get(*fusion).fusion_roots(),
-              ::testing::ElementsAre(negate))
+  EXPECT_EQ(&cache_.Get(*fusion).fusion_root(0).instruction(), negate)
       << "Analysis should be cached.";
 
   cache_.Invalidate(*fusion);
-  EXPECT_THAT(cache_.Get(*fusion).fusion_roots(),
-              ::testing::ElementsAre(broadcast))
+  EXPECT_EQ(&cache_.Get(*fusion).fusion_root(0).instruction(), broadcast)
       << "Analysis should have been recomputed";
 }
 

From b2040515fa7f087ee7a945f342ed50d1fe0d51b3 Mon Sep 17 00:00:00 2001
From: Mohammed Anany <manany@google.com>
Date: Tue, 21 May 2024 03:00:18 -0700
Subject: [PATCH 457/478] [Triton] Re-enable createRemoveLayoutConversionPass.

PiperOrigin-RevId: 635739578
---
 third_party/xla/xla/service/gpu/ir_emitter_triton.cc      | 2 +-
 third_party/xla/xla/service/gpu/ir_emitter_triton_cuda.cc | 7 +------
 2 files changed, 2 insertions(+), 7 deletions(-)

diff --git a/third_party/xla/xla/service/gpu/ir_emitter_triton.cc b/third_party/xla/xla/service/gpu/ir_emitter_triton.cc
index bd4412efcee26b..a0b882de6fe4c4 100644
--- a/third_party/xla/xla/service/gpu/ir_emitter_triton.cc
+++ b/third_party/xla/xla/service/gpu/ir_emitter_triton.cc
@@ -2034,7 +2034,7 @@ absl::Status EmitMatMul(mlir::OpBuilder builder,
         if (node->opcode() != HloOpcode::kConvert) {
           return false;
         }
-        auto in_type =
+        Type in_type =
             TritonType(builder, node->operand(0)->shape().element_type());
         Type out_type = TritonType(builder, node->shape().element_type());
         return in_type.getIntOrFloatBitWidth() <= 8 && out_type.isF32();
diff --git a/third_party/xla/xla/service/gpu/ir_emitter_triton_cuda.cc b/third_party/xla/xla/service/gpu/ir_emitter_triton_cuda.cc
index 0606d115ced97f..ec2c6e8136f6b3 100644
--- a/third_party/xla/xla/service/gpu/ir_emitter_triton_cuda.cc
+++ b/third_party/xla/xla/service/gpu/ir_emitter_triton_cuda.cc
@@ -83,12 +83,7 @@ absl::Status CreateTritonPipeline(
   }
 
   pm.addPass(mt::gpu::createOptimizeDotOperandsPass(ccCuda.IsAtLeastAmpere()));
-  // We need to disable this pass because it undoes the hoisting of dot_operand
-  // layout conversion done in
-  // triton/lib/Dialect/TritonGPU/Transforms/OptimizeDotOperands.cpp in
-  // HoistLayoutConversion pattern.
-  // Bug: b/331360119
-  // pm.addPass(mt::gpu::createRemoveLayoutConversionsPass());
+  pm.addPass(mt::gpu::createRemoveLayoutConversionsPass());
   pm.addPass(mt::gpu::createReduceDataDuplicationPass());
   pm.addPass(mt::gpu::createReorderInstructionsPass());
   pm.addPass(mlir::createCSEPass());

From aef9a9e199ef24a37b2a44a7ce8b330a7a68fa8b Mon Sep 17 00:00:00 2001
From: Chris Jones <cjfj@google.com>
Date: Tue, 21 May 2024 03:13:13 -0700
Subject: [PATCH 458/478] [xla:gpu] Fail gracefully (i.e. don't segfault) if
 Triton MLIR isn't parsable.

PiperOrigin-RevId: 635742591
---
 third_party/xla/xla/service/gpu/ir_emitter_unnested.cc | 1 +
 1 file changed, 1 insertion(+)

diff --git a/third_party/xla/xla/service/gpu/ir_emitter_unnested.cc b/third_party/xla/xla/service/gpu/ir_emitter_unnested.cc
index 8d3015fed82aef..d4461a2ad75582 100644
--- a/third_party/xla/xla/service/gpu/ir_emitter_unnested.cc
+++ b/third_party/xla/xla/service/gpu/ir_emitter_unnested.cc
@@ -1598,6 +1598,7 @@ absl::Status IrEmitterUnnested::EmitTritonCustomCall(
 
     auto triton_module =
         mlir::parseSourceString<mlir::ModuleOp>(call.ir, &mlir_context);
+    TF_RET_CHECK(triton_module);
     auto triton_fn =
         triton_module->lookupSymbol<mlir::triton::FuncOp>(call.name);
     triton_fn.setName(kernel_name);

From 1f5ff66a4bfe3040644ae17b36b8cb38cd837f70 Mon Sep 17 00:00:00 2001
From: Adrian Kuegel <akuegel@google.com>
Date: Tue, 21 May 2024 03:37:22 -0700
Subject: [PATCH 459/478] Automated Code Change

Fix header includes

PiperOrigin-RevId: 635748303
---
 third_party/xla/xla/service/gpu/fusions/BUILD | 28 ++++++++++++++++---
 .../gpu/fusions/mlir_emitter_test_base.h      |  2 ++
 .../xla/xla/service/gpu/fusions/reduction.cc  |  4 +++
 .../xla/service/gpu/fusions/reduction_base.cc |  2 ++
 .../xla/service/gpu/fusions/reduction_mlir.h  |  1 +
 .../xla/xla/service/gpu/fusions/scatter.cc    |  8 ++++++
 .../xla/xla/service/gpu/fusions/scatter.h     |  2 ++
 .../xla/service/gpu/fusions/scatter_mlir.cc   |  1 +
 .../service/gpu/fusions/scatter_mlir_test.cc  |  1 +
 .../xla/service/gpu/fusions/scatter_test.cc   |  1 +
 .../xla/service/gpu/fusions/tiling_util.cc    |  2 ++
 .../xla/xla/service/gpu/fusions/tiling_util.h |  1 +
 .../xla/xla/service/gpu/fusions/transpose.cc  |  3 ++
 .../xla/service/gpu/fusions/transpose_mlir.cc |  1 +
 .../xla/service/gpu/fusions/transpose_mlir.h  |  1 +
 .../gpu/fusions/transpose_mlir_test.cc        |  1 +
 .../xla/service/gpu/fusions/transpose_test.cc |  2 +-
 .../xla/xla/service/gpu/fusions/triton.cc     | 11 ++++++--
 .../xla/xla/service/gpu/fusions/triton.h      |  2 ++
 .../xla/service/gpu/fusions/triton_test.cc    |  2 +-
 20 files changed, 67 insertions(+), 9 deletions(-)

diff --git a/third_party/xla/xla/service/gpu/fusions/BUILD b/third_party/xla/xla/service/gpu/fusions/BUILD
index 94ccd305743800..c4bc16601c40d5 100644
--- a/third_party/xla/xla/service/gpu/fusions/BUILD
+++ b/third_party/xla/xla/service/gpu/fusions/BUILD
@@ -445,7 +445,6 @@ xla_cc_test(
         "//xla/service/gpu:hlo_fusion_analysis",
         "//xla/service/gpu/model:indexing_test_utils",
         "//xla/tests:xla_internal_test_main",
-        "@com_google_absl//absl/strings",
         "@com_google_googletest//:gtest",
         "@local_tsl//tsl/lib/core:status_test_util",
         "@local_tsl//tsl/platform:statusor",
@@ -476,6 +475,8 @@ cc_library(
         "//xla/service/gpu/model:indexing_map",
         "@com_google_absl//absl/algorithm:container",
         "@com_google_absl//absl/container:flat_hash_map",
+        "@com_google_absl//absl/container:inlined_vector",
+        "@com_google_absl//absl/log",
         "@com_google_absl//absl/log:check",
         "@com_google_absl//absl/types:span",
         "@llvm-project//llvm:Support",
@@ -533,6 +534,7 @@ cc_library(
         ":loop",
         "//xla:shape_util",
         "//xla:status",
+        "//xla:xla_data_proto_cc",
         "//xla/hlo/ir:hlo",
         "//xla/service/gpu:hlo_fusion_analysis",
         "//xla/service/gpu:ir_emission_utils",
@@ -541,6 +543,7 @@ cc_library(
         "//xla/service/gpu:launch_dimensions",
         "//xla/service/gpu:parallel_loop_emitter",
         "//xla/service/gpu/model:indexing_analysis",
+        "//xla/service/gpu/model:indexing_map",
         "//xla/service/llvm_ir:fused_ir_emitter",
         "//xla/service/llvm_ir:ir_array",
         "//xla/service/llvm_ir:llvm_util",
@@ -549,6 +552,8 @@ cc_library(
         "@com_google_absl//absl/types:span",
         "@llvm-project//llvm:Support",
         "@llvm-project//llvm:ir_headers",
+        "@llvm-project//mlir:IR",
+        "@local_tsl//tsl/platform:statusor",
     ],
 )
 
@@ -566,6 +571,7 @@ xla_cc_test(
         "//xla/tests:hlo_test_base",
         "//xla/tests:xla_internal_test_main",
         "@com_google_googletest//:gtest",
+        "@llvm-project//mlir:IR",
         "@local_tsl//tsl/platform:statusor",
     ],
 )
@@ -578,6 +584,7 @@ cc_library(
     deps = [
         "//xla:shape_util",
         "//xla:util",
+        "//xla:xla_data_proto_cc",
         "//xla/service/gpu:ir_emission_utils",
         "//xla/service/gpu:target_util",
         "//xla/service/llvm_ir:ir_array",
@@ -602,9 +609,12 @@ cc_library(
     local_defines = if_cuda_is_configured(["GOOGLE_CUDA=1"]),
     deps = [
         ":fusion_emitter",
+        "//xla:shape_util",
+        "//xla:status_macros",
         "//xla:statusor",
         "//xla/hlo/ir:hlo",
         "//xla/service/gpu:hlo_fusion_analysis",
+        "//xla/service/gpu:hlo_traversal",
         "//xla/service/gpu:ir_emission_utils",
         "//xla/service/gpu:ir_emitter_context",
         "//xla/service/gpu:ir_emitter_triton",
@@ -614,14 +624,17 @@ cc_library(
         "//xla/service/gpu:matmul_utils",
         "//xla/service/gpu:triton_fusion_analysis",
         "//xla/service/gpu/runtime:kernel_thunk",
+        "//xla/service/llvm_ir:ir_array",
         "//xla/service/llvm_ir:llvm_util",
         "@com_google_absl//absl/log",
         "@com_google_absl//absl/log:check",
         "@com_google_absl//absl/status",
+        "@com_google_absl//absl/status:statusor",
+        "@com_google_absl//absl/strings",
+        "@com_google_absl//absl/strings:string_view",
+        "@llvm-project//llvm:Support",
         "@llvm-project//llvm:ir_headers",
-        "@llvm-project//mlir:IR",
         "@llvm-project//mlir:Support",
-        "@local_tsl//tsl/platform:errors",
         "@local_tsl//tsl/platform:statusor",
     ],
 )
@@ -635,6 +648,7 @@ xla_cc_test(
         ":triton",
         "//xla/service/gpu:hlo_fusion_analysis",
         "//xla/stream_executor:device_description",
+        "//xla/stream_executor:device_description_proto_cc",
         "//xla/tests:hlo_test_base",
         "//xla/tests:xla_internal_test_main",
         "@com_google_googletest//:gtest",
@@ -729,6 +743,7 @@ cc_library(
         "//xla:status",
         "//xla:status_macros",
         "//xla:util",
+        "//xla:xla_data_proto_cc",
         "//xla/hlo/ir:hlo",
         "//xla/service:buffer_assignment",
         "//xla/service/gpu:hlo_fusion_analysis",
@@ -779,6 +794,7 @@ cc_library(
         "//xla:shape_util",
         "//xla:union_find",
         "//xla:util",
+        "//xla:xla_data_proto_cc",
         "//xla/hlo/ir:hlo",
         "//xla/hlo/utils:hlo_query",
         "//xla/service/gpu:gpu_fusible",
@@ -798,6 +814,7 @@ cc_library(
         "@com_google_absl//absl/container:node_hash_map",
         "@com_google_absl//absl/log",
         "@com_google_absl//absl/log:check",
+        "@com_google_absl//absl/types:span",
         "@llvm-project//llvm:Support",
         "@llvm-project//mlir:IR",
     ],
@@ -976,6 +993,7 @@ cc_library(
         "//xla:shape_util",
         "//xla:status",
         "//xla:util",
+        "//xla:xla_data_proto_cc",
         "//xla/hlo/ir:hlo",
         "//xla/service/gpu:hlo_fusion_analysis",
         "//xla/service/gpu:ir_emission_utils",
@@ -988,6 +1006,7 @@ cc_library(
         "//xla/service/llvm_ir:fused_ir_emitter",
         "//xla/service/llvm_ir:ir_array",
         "//xla/service/llvm_ir:llvm_util",
+        "//xla/service/llvm_ir:loop_emitter",
         "@com_google_absl//absl/container:flat_hash_map",
         "@com_google_absl//absl/container:inlined_vector",
         "@com_google_absl//absl/log:check",
@@ -997,6 +1016,7 @@ cc_library(
         "@llvm-project//llvm:Support",
         "@llvm-project//llvm:ir_headers",
         "@llvm-project//mlir:IR",
+        "@local_tsl//tsl/platform:statusor",
     ],
 )
 
@@ -1007,13 +1027,13 @@ xla_cc_test(
         ":fusions",
         ":transpose",
         "//xla:status_macros",
-        "//xla:statusor",
         "//xla/service/gpu:gpu_device_info_for_tests",
         "//xla/service/gpu:hlo_fusion_analysis",
         "//xla/service/gpu/model:indexing_test_utils",
         "//xla/stream_executor:device_description",
         "//xla/tests:hlo_test_base",
         "//xla/tests:xla_internal_test_main",
+        "@com_google_absl//absl/status:statusor",
         "@com_google_googletest//:gtest",
         "@llvm-project//mlir:IR",
         "@local_tsl//tsl/platform:statusor",
diff --git a/third_party/xla/xla/service/gpu/fusions/mlir_emitter_test_base.h b/third_party/xla/xla/service/gpu/fusions/mlir_emitter_test_base.h
index a299c2ea4007ba..61dbeffddfb240 100644
--- a/third_party/xla/xla/service/gpu/fusions/mlir_emitter_test_base.h
+++ b/third_party/xla/xla/service/gpu/fusions/mlir_emitter_test_base.h
@@ -21,6 +21,7 @@ limitations under the License.
 #include <string>
 #include <string_view>
 
+#include "absl/status/status.h"
 #include "absl/status/statusor.h"
 #include "mlir/IR/MLIRContext.h"  // from @llvm-project
 #include "mlir/Pass/PassManager.h"  // from @llvm-project
@@ -30,6 +31,7 @@ limitations under the License.
 #include "xla/service/gpu/model/affine_map_printer.h"
 #include "xla/stream_executor/device_description.h"
 #include "xla/tests/hlo_test_base.h"
+#include "xla/xla.pb.h"
 
 namespace xla {
 namespace gpu {
diff --git a/third_party/xla/xla/service/gpu/fusions/reduction.cc b/third_party/xla/xla/service/gpu/fusions/reduction.cc
index 83093677aabf26..1eba5e2024c4d8 100644
--- a/third_party/xla/xla/service/gpu/fusions/reduction.cc
+++ b/third_party/xla/xla/service/gpu/fusions/reduction.cc
@@ -28,9 +28,11 @@ limitations under the License.
 #include "absl/container/inlined_vector.h"
 #include "absl/container/node_hash_map.h"
 #include "absl/log/check.h"
+#include "absl/status/status.h"
 #include "absl/strings/str_cat.h"
 #include "absl/strings/string_view.h"
 #include "absl/types/span.h"
+#include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/Twine.h"
 #include "llvm/IR/Constants.h"
 #include "llvm/IR/DerivedTypes.h"
@@ -50,6 +52,7 @@ limitations under the License.
 #include "xla/service/buffer_assignment.h"
 #include "xla/service/gpu/elemental_ir_emitter.h"
 #include "xla/service/gpu/fusions/fusion_emitter.h"
+#include "xla/service/gpu/fusions/reduction_base.h"
 #include "xla/service/gpu/fusions/thunk_util.h"
 #include "xla/service/gpu/fusions/tiling_util.h"
 #include "xla/service/gpu/hlo_fusion_analysis.h"
@@ -77,6 +80,7 @@ limitations under the License.
 #include "xla/status_macros.h"
 #include "xla/stream_executor/device_description.h"
 #include "xla/util.h"
+#include "xla/xla_data.pb.h"
 #include "tsl/platform/errors.h"
 #include "tsl/platform/logging.h"
 #include "tsl/platform/status.h"
diff --git a/third_party/xla/xla/service/gpu/fusions/reduction_base.cc b/third_party/xla/xla/service/gpu/fusions/reduction_base.cc
index 00057e7528bc2f..0d63d73b3339af 100644
--- a/third_party/xla/xla/service/gpu/fusions/reduction_base.cc
+++ b/third_party/xla/xla/service/gpu/fusions/reduction_base.cc
@@ -28,6 +28,7 @@ limitations under the License.
 #include "absl/container/node_hash_map.h"
 #include "absl/log/check.h"
 #include "absl/log/log.h"
+#include "absl/types/span.h"
 #include "llvm/ADT/STLExtras.h"
 #include "mlir/IR/AffineExpr.h"  // from @llvm-project
 #include "mlir/IR/MLIRContext.h"  // from @llvm-project
@@ -49,6 +50,7 @@ limitations under the License.
 #include "xla/stream_executor/launch_dim.h"
 #include "xla/union_find.h"
 #include "xla/util.h"
+#include "xla/xla_data.pb.h"
 
 namespace xla {
 namespace gpu {
diff --git a/third_party/xla/xla/service/gpu/fusions/reduction_mlir.h b/third_party/xla/xla/service/gpu/fusions/reduction_mlir.h
index 12bc8124064245..9b2243afd52ac7 100644
--- a/third_party/xla/xla/service/gpu/fusions/reduction_mlir.h
+++ b/third_party/xla/xla/service/gpu/fusions/reduction_mlir.h
@@ -19,6 +19,7 @@ limitations under the License.
 
 #include "absl/status/status.h"
 #include "mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h"  // from @llvm-project
+#include "mlir/IR/MLIRContext.h"  // from @llvm-project
 #include "xla/hlo/ir/hlo_instruction.h"
 #include "xla/hlo/ir/hlo_instructions.h"
 #include "xla/service/gpu/fusions/mlir/computation_partitioner.h"
diff --git a/third_party/xla/xla/service/gpu/fusions/scatter.cc b/third_party/xla/xla/service/gpu/fusions/scatter.cc
index 002d8d3dfb67a9..eb7bfd3086b484 100644
--- a/third_party/xla/xla/service/gpu/fusions/scatter.cc
+++ b/third_party/xla/xla/service/gpu/fusions/scatter.cc
@@ -28,16 +28,22 @@ limitations under the License.
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/IR/IRBuilder.h"
 #include "llvm/IR/Value.h"
+#include "mlir/IR/AffineExpr.h"  // from @llvm-project
+#include "mlir/IR/AffineMap.h"  // from @llvm-project
+#include "mlir/IR/MLIRContext.h"  // from @llvm-project
 #include "xla/hlo/ir/hlo_casting_utils.h"
 #include "xla/hlo/ir/hlo_computation.h"
 #include "xla/hlo/ir/hlo_instructions.h"
+#include "xla/hlo/ir/hlo_opcode.h"
 #include "xla/service/gpu/elemental_ir_emitter.h"
 #include "xla/service/gpu/fusions/loop.h"
+#include "xla/service/gpu/hlo_fusion_analysis.h"
 #include "xla/service/gpu/ir_emission_utils.h"
 #include "xla/service/gpu/ir_emitter_context.h"
 #include "xla/service/gpu/ir_emitter_nested.h"
 #include "xla/service/gpu/launch_dimensions.h"
 #include "xla/service/gpu/model/indexing_analysis.h"
+#include "xla/service/gpu/model/indexing_map.h"
 #include "xla/service/gpu/parallel_loop_emitter.h"
 #include "xla/service/llvm_ir/fused_ir_emitter.h"
 #include "xla/service/llvm_ir/ir_array.h"
@@ -45,6 +51,8 @@ limitations under the License.
 #include "xla/shape.h"
 #include "xla/shape_util.h"
 #include "xla/status.h"
+#include "xla/xla_data.pb.h"
+#include "tsl/platform/statusor.h"
 
 namespace xla {
 namespace gpu {
diff --git a/third_party/xla/xla/service/gpu/fusions/scatter.h b/third_party/xla/xla/service/gpu/fusions/scatter.h
index 289328fd7a7c14..f054a3ecf81f9d 100644
--- a/third_party/xla/xla/service/gpu/fusions/scatter.h
+++ b/third_party/xla/xla/service/gpu/fusions/scatter.h
@@ -20,12 +20,14 @@ limitations under the License.
 
 #include "absl/log/check.h"
 #include "llvm/IR/IRBuilder.h"
+#include "mlir/IR/MLIRContext.h"  // from @llvm-project
 #include "xla/hlo/ir/hlo_instructions.h"
 #include "xla/hlo/ir/hlo_opcode.h"
 #include "xla/service/gpu/fusions/fusion_emitter.h"
 #include "xla/service/gpu/hlo_fusion_analysis.h"
 #include "xla/service/gpu/ir_emitter_context.h"
 #include "xla/service/gpu/launch_dimensions.h"
+#include "xla/service/gpu/model/indexing_map.h"
 #include "xla/service/llvm_ir/ir_array.h"
 #include "xla/status.h"
 
diff --git a/third_party/xla/xla/service/gpu/fusions/scatter_mlir.cc b/third_party/xla/xla/service/gpu/fusions/scatter_mlir.cc
index 372b6926953e27..18d357a4316f5e 100644
--- a/third_party/xla/xla/service/gpu/fusions/scatter_mlir.cc
+++ b/third_party/xla/xla/service/gpu/fusions/scatter_mlir.cc
@@ -41,6 +41,7 @@ limitations under the License.
 #include "xla/service/gpu/fusions/mlir/computation_partitioner.h"
 #include "xla/service/gpu/fusions/mlir/elemental_hlo_to_mlir.h"
 #include "xla/service/gpu/fusions/mlir/ir/xla_gpu_ops.h"
+#include "xla/service/gpu/hlo_fusion_analysis.h"
 #include "xla/service/gpu/launch_dimensions.h"
 #include "xla/service/gpu/model/indexing_analysis.h"
 #include "xla/service/gpu/model/indexing_map.h"
diff --git a/third_party/xla/xla/service/gpu/fusions/scatter_mlir_test.cc b/third_party/xla/xla/service/gpu/fusions/scatter_mlir_test.cc
index a961a560288d1f..489e588cec4248 100644
--- a/third_party/xla/xla/service/gpu/fusions/scatter_mlir_test.cc
+++ b/third_party/xla/xla/service/gpu/fusions/scatter_mlir_test.cc
@@ -14,6 +14,7 @@ limitations under the License.
 ==============================================================================*/
 #include "xla/service/gpu/fusions/scatter_mlir.h"
 
+#include <gmock/gmock.h>
 #include <gtest/gtest.h>
 #include "xla/error_spec.h"
 #include "xla/service/gpu/fusions/mlir_emitter_test_base.h"
diff --git a/third_party/xla/xla/service/gpu/fusions/scatter_test.cc b/third_party/xla/xla/service/gpu/fusions/scatter_test.cc
index 07e2da47a6a6b0..ae55b19ac4c3a0 100644
--- a/third_party/xla/xla/service/gpu/fusions/scatter_test.cc
+++ b/third_party/xla/xla/service/gpu/fusions/scatter_test.cc
@@ -18,6 +18,7 @@ limitations under the License.
 
 #include <gmock/gmock.h>
 #include <gtest/gtest.h>
+#include "mlir/IR/MLIRContext.h"  // from @llvm-project
 #include "xla/service/gpu/fusions/fusions.h"
 #include "xla/service/gpu/gpu_device_info_for_tests.h"
 #include "xla/service/gpu/hlo_fusion_analysis.h"
diff --git a/third_party/xla/xla/service/gpu/fusions/tiling_util.cc b/third_party/xla/xla/service/gpu/fusions/tiling_util.cc
index 83f9e5fcfb1da7..9ad085faff3a99 100644
--- a/third_party/xla/xla/service/gpu/fusions/tiling_util.cc
+++ b/third_party/xla/xla/service/gpu/fusions/tiling_util.cc
@@ -20,6 +20,7 @@ limitations under the License.
 #include <string>
 #include <vector>
 
+#include "absl/container/inlined_vector.h"
 #include "absl/log/check.h"
 #include "absl/strings/str_cat.h"
 #include "absl/types/span.h"
@@ -38,6 +39,7 @@ limitations under the License.
 #include "xla/service/llvm_ir/llvm_util.h"
 #include "xla/shape_util.h"
 #include "xla/util.h"
+#include "xla/xla_data.pb.h"
 #include "tsl/platform/statusor.h"
 
 namespace xla {
diff --git a/third_party/xla/xla/service/gpu/fusions/tiling_util.h b/third_party/xla/xla/service/gpu/fusions/tiling_util.h
index bc65a733419307..732f39a5daa4df 100644
--- a/third_party/xla/xla/service/gpu/fusions/tiling_util.h
+++ b/third_party/xla/xla/service/gpu/fusions/tiling_util.h
@@ -31,6 +31,7 @@ limitations under the License.
 #include "xla/shape.h"
 #include "xla/shape_util.h"
 #include "xla/util.h"
+#include "xla/xla_data.pb.h"
 
 namespace xla {
 namespace gpu {
diff --git a/third_party/xla/xla/service/gpu/fusions/transpose.cc b/third_party/xla/xla/service/gpu/fusions/transpose.cc
index ee66e531a9cf7d..7fb3fab849fa8a 100644
--- a/third_party/xla/xla/service/gpu/fusions/transpose.cc
+++ b/third_party/xla/xla/service/gpu/fusions/transpose.cc
@@ -48,9 +48,12 @@ limitations under the License.
 #include "xla/service/llvm_ir/fused_ir_emitter.h"
 #include "xla/service/llvm_ir/ir_array.h"
 #include "xla/service/llvm_ir/llvm_util.h"
+#include "xla/service/llvm_ir/loop_emitter.h"
 #include "xla/shape_util.h"
 #include "xla/status.h"
 #include "xla/util.h"
+#include "xla/xla_data.pb.h"
+#include "tsl/platform/statusor.h"
 
 namespace xla {
 namespace gpu {
diff --git a/third_party/xla/xla/service/gpu/fusions/transpose_mlir.cc b/third_party/xla/xla/service/gpu/fusions/transpose_mlir.cc
index 23cf134412f0ff..df65dbb55ea5e1 100644
--- a/third_party/xla/xla/service/gpu/fusions/transpose_mlir.cc
+++ b/third_party/xla/xla/service/gpu/fusions/transpose_mlir.cc
@@ -22,6 +22,7 @@ limitations under the License.
 
 #include "absl/algorithm/container.h"
 #include "absl/container/flat_hash_map.h"
+#include "absl/container/inlined_vector.h"
 #include "absl/log/check.h"
 #include "absl/types/span.h"
 #include "llvm/ADT/STLExtras.h"
diff --git a/third_party/xla/xla/service/gpu/fusions/transpose_mlir.h b/third_party/xla/xla/service/gpu/fusions/transpose_mlir.h
index e4613ff8e3a6d5..af2869ecd32ec6 100644
--- a/third_party/xla/xla/service/gpu/fusions/transpose_mlir.h
+++ b/third_party/xla/xla/service/gpu/fusions/transpose_mlir.h
@@ -26,6 +26,7 @@ limitations under the License.
 #include "mlir/IR/MLIRContext.h"  // from @llvm-project
 #include "mlir/IR/Value.h"  // from @llvm-project
 #include "mlir/IR/ValueRange.h"  // from @llvm-project
+#include "mlir/Support/LLVM.h"  // from @llvm-project
 #include "xla/hlo/ir/hlo_instructions.h"
 #include "xla/service/gpu/fusions/mlir/computation_partitioner.h"
 #include "xla/service/gpu/fusions/mlir/mlir_fusion_emitter.h"
diff --git a/third_party/xla/xla/service/gpu/fusions/transpose_mlir_test.cc b/third_party/xla/xla/service/gpu/fusions/transpose_mlir_test.cc
index f2d3a299d3162b..3236aec62889c6 100644
--- a/third_party/xla/xla/service/gpu/fusions/transpose_mlir_test.cc
+++ b/third_party/xla/xla/service/gpu/fusions/transpose_mlir_test.cc
@@ -14,6 +14,7 @@ limitations under the License.
 ==============================================================================*/
 #include "xla/service/gpu/fusions/transpose_mlir.h"
 
+#include <gmock/gmock.h>
 #include <gtest/gtest.h>
 #include "xla/error_spec.h"
 #include "xla/service/gpu/fusions/mlir_emitter_test_base.h"
diff --git a/third_party/xla/xla/service/gpu/fusions/transpose_test.cc b/third_party/xla/xla/service/gpu/fusions/transpose_test.cc
index a9b40704ff7842..6bb842bf9dcb5c 100644
--- a/third_party/xla/xla/service/gpu/fusions/transpose_test.cc
+++ b/third_party/xla/xla/service/gpu/fusions/transpose_test.cc
@@ -19,13 +19,13 @@ limitations under the License.
 
 #include <gmock/gmock.h>
 #include <gtest/gtest.h>
+#include "absl/status/statusor.h"
 #include "mlir/IR/MLIRContext.h"  // from @llvm-project
 #include "xla/service/gpu/fusions/fusions.h"
 #include "xla/service/gpu/gpu_device_info_for_tests.h"
 #include "xla/service/gpu/hlo_fusion_analysis.h"
 #include "xla/service/gpu/model/indexing_test_utils.h"
 #include "xla/status_macros.h"
-#include "xla/statusor.h"
 #include "xla/stream_executor/device_description.h"
 #include "xla/tests/hlo_test_base.h"
 #include "tsl/platform/statusor.h"
diff --git a/third_party/xla/xla/service/gpu/fusions/triton.cc b/third_party/xla/xla/service/gpu/fusions/triton.cc
index 30a98194e04200..37c2bb69dc77cb 100644
--- a/third_party/xla/xla/service/gpu/fusions/triton.cc
+++ b/third_party/xla/xla/service/gpu/fusions/triton.cc
@@ -22,14 +22,18 @@ limitations under the License.
 
 #include "absl/log/check.h"
 #include "absl/log/log.h"
+#include "absl/strings/str_cat.h"
+#include "absl/strings/string_view.h"
+#include "llvm/ADT/STLExtras.h"
 #include "llvm/IR/Function.h"
 #include "llvm/IR/IRBuilder.h"
-#include "mlir/IR/Operation.h"  // from @llvm-project
 #include "mlir/Support/LLVM.h"  // from @llvm-project
 #include "xla/hlo/ir/hlo_computation.h"
 #include "xla/hlo/ir/hlo_instruction.h"
 #include "xla/hlo/ir/hlo_instructions.h"
+#include "xla/hlo/ir/hlo_opcode.h"
 #include "xla/service/gpu/fusions/fusion_emitter.h"
+#include "xla/service/gpu/hlo_traversal.h"
 #include "xla/service/gpu/ir_emission_utils.h"
 #include "xla/service/gpu/ir_emitter_context.h"
 #include "xla/service/gpu/kernel_arguments.h"
@@ -38,9 +42,10 @@ limitations under the License.
 #include "xla/service/gpu/matmul_utils.h"
 #include "xla/service/gpu/runtime/kernel_thunk.h"
 #include "xla/service/gpu/triton_fusion_analysis.h"
+#include "xla/service/llvm_ir/ir_array.h"
 #include "xla/service/llvm_ir/llvm_util.h"
-#include "xla/statusor.h"
-#include "tsl/platform/errors.h"
+#include "xla/shape.h"
+#include "xla/status_macros.h"
 #include "tsl/platform/statusor.h"
 
 #if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
diff --git a/third_party/xla/xla/service/gpu/fusions/triton.h b/third_party/xla/xla/service/gpu/fusions/triton.h
index c584ca8c56fb7e..7f6ceb6dd70462 100644
--- a/third_party/xla/xla/service/gpu/fusions/triton.h
+++ b/third_party/xla/xla/service/gpu/fusions/triton.h
@@ -17,10 +17,12 @@ limitations under the License.
 
 #include <optional>
 
+#include "absl/status/statusor.h"
 #include "xla/hlo/ir/hlo_instructions.h"
 #include "xla/service/gpu/fusions/fusion_emitter.h"
 #include "xla/service/gpu/hlo_fusion_analysis.h"
 #include "xla/service/gpu/ir_emitter_context.h"
+#include "xla/service/gpu/launch_dimensions.h"
 #include "xla/statusor.h"
 
 namespace xla {
diff --git a/third_party/xla/xla/service/gpu/fusions/triton_test.cc b/third_party/xla/xla/service/gpu/fusions/triton_test.cc
index 49b8f6d966444b..4d85c47cd9c284 100644
--- a/third_party/xla/xla/service/gpu/fusions/triton_test.cc
+++ b/third_party/xla/xla/service/gpu/fusions/triton_test.cc
@@ -16,11 +16,11 @@ limitations under the License.
 
 #include <optional>
 
-#include <gmock/gmock.h>
 #include <gtest/gtest.h>
 #include "xla/service/gpu/fusions/fusions.h"
 #include "xla/service/gpu/hlo_fusion_analysis.h"
 #include "xla/stream_executor/device_description.h"
+#include "xla/stream_executor/device_description.pb.h"
 #include "xla/tests/hlo_test_base.h"
 #include "tsl/platform/statusor.h"
 

From 39a2d6995e445826243cef18307187a25d2c1c2f Mon Sep 17 00:00:00 2001
From: Johannes Reifferscheid <jreiffers@google.com>
Date: Tue, 21 May 2024 05:01:17 -0700
Subject: [PATCH 460/478] Replace all remaining uses of affine.apply with
 apply_indexing.

Also split indexing maps into one map per result, to not block
LICM.

PiperOrigin-RevId: 635767137
---
 .../service/gpu/fusions/concatenate_mlir.cc   |  10 +-
 .../gpu/fusions/concatenate_mlir_test.cc      |  17 +-
 .../in_place_dynamic_update_slice_mlir.cc     |  11 +-
 ...in_place_dynamic_update_slice_mlir_test.cc |  16 +-
 .../service/gpu/fusions/input_slices_mlir.cc  |   9 +-
 .../xla/xla/service/gpu/fusions/mlir/BUILD    |   4 +-
 .../fusions/mlir/computation_partitioner.cc   |  21 ++-
 .../fusions/mlir/computation_partitioner.h    |   5 +-
 .../mlir/computation_partitioner_test.cc      |   5 +-
 .../gpu/fusions/mlir/elemental_hlo_to_mlir.cc |  90 +++++------
 .../gpu/fusions/mlir/elemental_hlo_to_mlir.h  |   6 +-
 .../mlir/elemental_hlo_to_mlir_test.cc        | 146 +++++++++++-------
 .../xla/service/gpu/fusions/reduction_mlir.cc |  24 +--
 .../xla/service/gpu/fusions/scatter_mlir.cc   |  11 +-
 .../service/gpu/fusions/scatter_mlir_test.cc  |   8 +-
 .../xla/service/gpu/fusions/transpose_mlir.cc |   3 +-
 .../xla/xla/service/gpu/model/indexing_map.h  |   5 +
 17 files changed, 208 insertions(+), 183 deletions(-)

diff --git a/third_party/xla/xla/service/gpu/fusions/concatenate_mlir.cc b/third_party/xla/xla/service/gpu/fusions/concatenate_mlir.cc
index f4e2f18a5fdd6d..d53df96e35589e 100644
--- a/third_party/xla/xla/service/gpu/fusions/concatenate_mlir.cc
+++ b/third_party/xla/xla/service/gpu/fusions/concatenate_mlir.cc
@@ -118,18 +118,16 @@ absl::Status MlirConcatenateFusion::EmitEntryFunction(
         [&, operand_index = operand_index](
             ValueRange output_tensors, ValueRange dim_values,
             ValueRange symbol_values) -> SmallVector<Value> {
-      auto input_indices =
-          mlir_converter::ApplyAffineMap(thread_id_to_input_map.GetAffineMap(),
-                                         dim_values, symbol_values, builder);
+      auto input_indices = mlir_converter::ApplyIndexing(
+          thread_id_to_input_map, dim_values, symbol_values, builder);
 
       auto result_scalar = mlir_converter::ProvideParameter(
           root_computation, concat, operand_index, input_indices, call_targets,
           entry_function, builder);
       absl::flat_hash_map<const HloInstruction*, llvm::SmallVector<Value>>
           hero_value{{concat, result_scalar}};
-      auto output_indices =
-          mlir_converter::ApplyAffineMap(thread_id_to_output_map.GetAffineMap(),
-                                         dim_values, symbol_values, builder);
+      auto output_indices = mlir_converter::ApplyIndexing(
+          thread_id_to_output_map, dim_values, symbol_values, builder);
       auto result_scalars = EmitEpilogue(
           /*epilogue_index=*/0, computations, entry_function, hero_value,
           output_indices, builder)[&analysis_.fusion_root(0).instruction()];
diff --git a/third_party/xla/xla/service/gpu/fusions/concatenate_mlir_test.cc b/third_party/xla/xla/service/gpu/fusions/concatenate_mlir_test.cc
index 781ebf563f66ea..dce6eb3e7a0f94 100644
--- a/third_party/xla/xla/service/gpu/fusions/concatenate_mlir_test.cc
+++ b/third_party/xla/xla/service/gpu/fusions/concatenate_mlir_test.cc
@@ -102,10 +102,9 @@ TEST_F(MlirConcatenateFusionTest, StandAloneConcatenate) {
     }
   )";
   TF_ASSERT_OK(EmitAndCheckIR(kHloString, R"(
-    // CHECK-DAG: #[[MAP_1:.*]] = affine_map<()[s0, s1] -> (s0 + s1 * 128)>
-    // CHECK-DAG: #[[MAP_2:.*]] = affine_map<()[s0, s1] -> ((s0 + s1 * 128) mod 400)>
-    // CHECK-DAG: #[[MAP_3:.*]] = affine_map<()[s0, s1] -> ((s0 + s1 * 128) mod 400 + 200)>
-    // CHECK-DAG: #[[MAP_4:.*]] = affine_map<()[s0, s1] -> ((s0 + s1 * 128) mod 400 + 600)>
+    // CHECK-DAG: #[[MAP_1:.*]] = affine_map<(d0, d1) -> ((d1 * 128 + d0) mod 400)>
+    // CHECK-DAG: #[[MAP_2:.*]] = affine_map<(d0, d1) -> ((d1 * 128 + d0) mod 400 + 200)>
+    // CHECK-DAG: #[[MAP_3:.*]] = affine_map<(d0, d1) -> ((d1 * 128 + d0) mod 400 + 600)>
 
     // CHECK-LABEL: fused_computation
     // CHECK-SAME:    %[[ARG_0:[a-zA-Z0-9]*]]: {{[^,]*}},
@@ -116,19 +115,19 @@ TEST_F(MlirConcatenateFusionTest, StandAloneConcatenate) {
     // CHECK: %[[THREAD_ID:.*]] = gpu.thread_id x
     // CHECK: %[[BLOCK_ID:.*]] = gpu.block_id x
 
-    // CHECK: %[[INPUT_INDEX_1:.*]] = affine.apply #[[MAP_2]]()[%[[THREAD_ID]], %[[BLOCK_ID]]]
+    // CHECK: %[[INPUT_INDEX_1:.*]] = xla_gpu.apply_indexing #[[MAP_1]]
     // CHECK: %[[IF_1:.*]] = scf.if
     // CHECK:   %[[VAL_1:.*]] = xla_gpu.pure_call @fused_computation_param0
     // CHECK:   %[[INSERTED_1:.*]] = tensor.insert %[[VAL_1:.*]] into %[[OUTPUT]][%[[INPUT_INDEX_1]]]
 
     // CHECK: %[[IF_2:.*]] = scf.if
     // CHECK:   %[[VAL_2:.*]] = xla_gpu.pure_call @fused_computation_param1
-    // CHECK:   %[[OUTPUT_INDEX_2:.*]] = affine.apply #[[MAP_3]]()[%[[THREAD_ID]], %[[BLOCK_ID]]]
+    // CHECK:   %[[OUTPUT_INDEX_2:.*]] = xla_gpu.apply_indexing #[[MAP_2]]
     // CHECK:   %[[INSERTED_2:.*]] = tensor.insert %[[VAL_2:.*]] into {{.*}}[%[[OUTPUT_INDEX_2]]]
 
     // CHECK: %[[IF_3:.*]] = scf.if
     // CHECK:   %[[VAL_3:.*]] = xla_gpu.pure_call @fused_computation_param2
-    // CHECK:   %[[OUTPUT_INDEX_3:.*]] = affine.apply #[[MAP_4]]()[%[[THREAD_ID]], %[[BLOCK_ID]]]
+    // CHECK:   %[[OUTPUT_INDEX_3:.*]] = xla_gpu.apply_indexing #[[MAP_3]]
     // CHECK:   %[[INSERTED_3:.*]] = tensor.insert %[[VAL_3:.*]] into {{.*}}[%[[OUTPUT_INDEX_3]]]
   )"));
   EXPECT_TRUE(RunAndCompareNoHloPasses(kHloString, ErrorSpec{1e-3}));
@@ -153,7 +152,7 @@ TEST_F(MlirConcatenateFusionTest, PrologueEpilogue) {
     }
   )";
   TF_ASSERT_OK(EmitAndCheckIR(kHloString, R"(
-    // CHECK: #[[MAP:.*]] = affine_map<()[s0] -> (s0 + 64)>
+    // CHECK: #[[MAP:.*]] = affine_map<(d0) -> (d0 + 64)>
 
     // CHECK-LABEL: fused_computation
     // CHECK-DAG: %[[C_63:.*]] = arith.constant 63
@@ -167,7 +166,7 @@ TEST_F(MlirConcatenateFusionTest, PrologueEpilogue) {
     // CHECK:   scf.yield %[[INSERTED_1]]
 
     // CHECK: %[[VAL_2_1:.*]] = xla_gpu.pure_call @fused_computation_exp({{.*}}, %[[THREAD_ID]])
-    // CHECK: %[[INDEX_2:.*]] = affine.apply #[[MAP]]()[%[[THREAD_ID]]]
+    // CHECK: %[[INDEX_2:.*]] = xla_gpu.apply_indexing #[[MAP]](%[[THREAD_ID]]
     // CHECK: %[[VAL_2_2:.*]] = xla_gpu.pure_call @fused_computation__epilogue__neg({{.*}}, %[[INDEX_2]], %[[VAL_2_1]])
     // CHECK: %[[INSERTED_2:.*]] = tensor.insert %[[VAL_2_2:.*]] into {{.*}}[%[[INDEX_2]]]
 
diff --git a/third_party/xla/xla/service/gpu/fusions/in_place_dynamic_update_slice_mlir.cc b/third_party/xla/xla/service/gpu/fusions/in_place_dynamic_update_slice_mlir.cc
index b300bcf1a23d36..a6f019727865e9 100644
--- a/third_party/xla/xla/service/gpu/fusions/in_place_dynamic_update_slice_mlir.cc
+++ b/third_party/xla/xla/service/gpu/fusions/in_place_dynamic_update_slice_mlir.cc
@@ -55,7 +55,7 @@ using mlir::ValueRange;
 using mlir::arith::AddIOp;
 using mlir::func::ReturnOp;
 using mlir::tensor::InsertOp;
-using mlir_converter::ApplyAffineMap;
+using mlir_converter::ApplyIndexing;
 using mlir_converter::CallTargetProvider;
 using mlir_converter::ClampIndex;
 using mlir_converter::PartitionedComputations;
@@ -122,8 +122,8 @@ absl::Status MlirInPlaceDynamicUpdateSliceFusion::EmitEntryFunction(
       b, output_tensor_args, indexing,
       [&](ValueRange output_tensors, ValueRange dim_values,
           ValueRange symbol_values) -> llvm::SmallVector<Value> {
-        auto input_indices = ApplyAffineMap(indexing.GetAffineMap(), dim_values,
-                                            symbol_values, b);
+        auto input_indices =
+            ApplyIndexing(indexing, dim_values, symbol_values, b);
         llvm::SmallVector<Value> results;
         for (auto [instr, root, output] :
              llvm::zip(dus_ops_, analysis_.fusion_roots(), output_tensors)) {
@@ -154,9 +154,8 @@ absl::Status MlirInPlaceDynamicUpdateSliceFusion::EmitEntryFunction(
                                input_indices, call_targets, entry_function, b);
           // Handle bitcasts under the DUS.
           if (dus_instr->shape() != root.shape()) {
-            update_indices = ApplyAffineMap(
-                GetBitcastMap(dus_instr->shape(), root.shape(), b.getContext())
-                    .GetAffineMap(),
+            update_indices = ApplyIndexing(
+                GetBitcastMap(dus_instr->shape(), root.shape(), b.getContext()),
                 update_indices, {}, b);
           }
           results.push_back(
diff --git a/third_party/xla/xla/service/gpu/fusions/in_place_dynamic_update_slice_mlir_test.cc b/third_party/xla/xla/service/gpu/fusions/in_place_dynamic_update_slice_mlir_test.cc
index 53d34c4aa02ee4..e22ea62a6f5d36 100644
--- a/third_party/xla/xla/service/gpu/fusions/in_place_dynamic_update_slice_mlir_test.cc
+++ b/third_party/xla/xla/service/gpu/fusions/in_place_dynamic_update_slice_mlir_test.cc
@@ -100,8 +100,8 @@ TEST_F(MlirInPlaceDynamicUpdateSliceFusionTest, SimpleDUS) {
     }
   )";
   TF_ASSERT_OK(EmitAndCheckIR(kHloString, R"(
-    // CHECK-DAG: #[[MAP_1:.*]] = affine_map<()[s0] -> (s0 floordiv 6)>
-    // CHECK-DAG: #[[MAP_2:.*]] = affine_map<()[s0] -> (s0 mod 6)>
+    // CHECK-DAG: #[[MAP_1:.*]] = affine_map<(d0) -> (d0 floordiv 6)>
+    // CHECK-DAG: #[[MAP_2:.*]] = affine_map<(d0) -> (d0 mod 6)>
     // CHECK:     func.func @fused_computation
     // CHECK-SAME:  %arg0: tensor<20x30xf32>
     // CHECK-SAME:  %arg1: tensor<5x6xf32>
@@ -112,8 +112,8 @@ TEST_F(MlirInPlaceDynamicUpdateSliceFusionTest, SimpleDUS) {
     // CHECK-DAG:   %[[C_15:.*]] = arith.constant 15
     // CHECK-DAG:   %[[C_0:.*]] = arith.constant 0
     // CHECK:       %[[THREAD_ID:.*]] = gpu.thread_id  x
-    // CHECK:       %[[INPUT_INDEX_0:.*]] = affine.apply #[[MAP_1]]()[%[[THREAD_ID]]]
-    // CHECK:       %[[INPUT_INDEX_1:.*]] = affine.apply #[[MAP_2]]()[%[[THREAD_ID]]]
+    // CHECK:       %[[INPUT_INDEX_0:.*]] = xla_gpu.apply_indexing #[[MAP_1]](%[[THREAD_ID]] in [0, 29])
+    // CHECK:       %[[INPUT_INDEX_1:.*]] = xla_gpu.apply_indexing #[[MAP_2]](%[[THREAD_ID]] in [0, 29])
     // CHECK:       %[[I0:.*]] = xla_gpu.pure_call @fused_computation_i0
     // CHECK:       %[[I1:.*]] = xla_gpu.pure_call @fused_computation_i1
     // CHECK:       %[[IDX0:.*]] = arith.index_cast %[[I0]]
@@ -151,8 +151,8 @@ TEST_F(MlirInPlaceDynamicUpdateSliceFusionTest, OutOfBoundDUS) {
     }
   )";
   TF_ASSERT_OK(EmitAndCheckIR(kHloString, R"(
-    // CHECK-DAG: #[[MAP_1:.*]] = affine_map<()[s0] -> (s0 floordiv 3)>
-    // CHECK-DAG: #[[MAP_2:.*]] = affine_map<()[s0] -> (s0 mod 3)>
+    // CHECK-DAG: #[[MAP_1:.*]] = affine_map<(d0) -> (d0 floordiv 3)>
+    // CHECK-DAG: #[[MAP_2:.*]] = affine_map<(d0) -> (d0 mod 3)>
     // CHECK:     func.func @fused_computation
     // CHECK-SAME:  %arg0: tensor<7x8xf32>
     // CHECK-SAME:  %arg1: tensor<2x3xf32>
@@ -162,8 +162,8 @@ TEST_F(MlirInPlaceDynamicUpdateSliceFusionTest, OutOfBoundDUS) {
     // CHECK-DAG:   %[[C_5:.*]] = arith.constant 5
     // CHECK-DAG:   %[[C_0:.*]] = arith.constant 0
     // CHECK:       %[[THREAD_ID:.*]] = gpu.thread_id  x
-    // CHECK:       %[[INPUT_INDEX_0:.*]] = affine.apply #[[MAP_1]]()[%[[THREAD_ID]]]
-    // CHECK:       %[[INPUT_INDEX_1:.*]] = affine.apply #[[MAP_2]]()[%[[THREAD_ID]]]
+    // CHECK:       %[[INPUT_INDEX_0:.*]] = xla_gpu.apply_indexing #[[MAP_1]](%[[THREAD_ID]] in [0, 5])
+    // CHECK:       %[[INPUT_INDEX_1:.*]] = xla_gpu.apply_indexing #[[MAP_2]](%[[THREAD_ID]] in [0, 5])
     // CHECK:       %[[I0:.*]] = xla_gpu.pure_call @fused_computation_i0
     // CHECK:       %[[I1:.*]] = xla_gpu.pure_call @fused_computation_i1
     // CHECK:       %[[IDX0:.*]] = arith.index_cast %[[I0]]
diff --git a/third_party/xla/xla/service/gpu/fusions/input_slices_mlir.cc b/third_party/xla/xla/service/gpu/fusions/input_slices_mlir.cc
index 806bb1d4328907..6b9f280745c322 100644
--- a/third_party/xla/xla/service/gpu/fusions/input_slices_mlir.cc
+++ b/third_party/xla/xla/service/gpu/fusions/input_slices_mlir.cc
@@ -113,8 +113,8 @@ absl::Status MlirInputSlicesFusion::EmitEntryFunction(
       builder, output_tensor_args, input_indexing,
       [&](ValueRange output_tensors, ValueRange dim_values,
           ValueRange symbol_values) -> SmallVector<Value> {
-        auto input_indices = mlir_converter::ApplyAffineMap(
-            input_indexing.GetAffineMap(), dim_values, symbol_values, builder);
+        auto input_indices = mlir_converter::ApplyIndexing(
+            input_indexing, dim_values, symbol_values, builder);
         SmallVector<Value> input_operands(
             entry_function.getArguments().take_front(num_inputs));
         absl::c_copy(input_indices, std::back_inserter(input_operands));
@@ -141,9 +141,8 @@ absl::Status MlirInputSlicesFusion::EmitEntryFunction(
               [&, output_index = output_index, output = output](
                   mlir::OpBuilder b, mlir::Location loc) {
                 mlir::ImplicitLocOpBuilder then_builder(loc, b);
-                auto output_indices = mlir_converter::ApplyAffineMap(
-                    output_indexing->GetAffineMap(), dim_values, symbol_values,
-                    then_builder);
+                auto output_indices = mlir_converter::ApplyIndexing(
+                    *output_indexing, dim_values, symbol_values, then_builder);
                 const auto* arg = analysis_.fusion_root(output_index)
                                       .instruction()
                                       .operand(0);
diff --git a/third_party/xla/xla/service/gpu/fusions/mlir/BUILD b/third_party/xla/xla/service/gpu/fusions/mlir/BUILD
index 645bafc935eac7..35d2cf8bcd8a9d 100644
--- a/third_party/xla/xla/service/gpu/fusions/mlir/BUILD
+++ b/third_party/xla/xla/service/gpu/fusions/mlir/BUILD
@@ -26,6 +26,7 @@ cc_library(
         "//xla/service/gpu:hlo_fusion_analysis",
         "//xla/service/gpu/fusions:fusion_emitter",
         "//xla/service/gpu/model:indexing_analysis",
+        "//xla/service/gpu/model:indexing_map",
         "//xla/service/llvm_ir:llvm_util",
         "//xla/translate/hlo_to_mhlo:hlo_utils",
         "@com_google_absl//absl/algorithm:container",
@@ -49,7 +50,7 @@ xla_cc_test(
     deps = [
         ":computation_partitioner",
         "//xla/hlo/ir:hlo",
-        "//xla/service/gpu:launch_dimensions",
+        "//xla/service/gpu/model:indexing_analysis",
         "//xla/tests:hlo_test_base",
         "//xla/tests:xla_internal_test_main",
         "@com_google_googletest//:gtest",
@@ -119,6 +120,7 @@ xla_cc_test(
         "//xla/service:hlo_parser",
         "//xla/service/gpu:launch_dimensions",
         "//xla/service/gpu/fusions/mlir/ir:xla_gpu",
+        "//xla/service/gpu/model:indexing_analysis",
         "//xla/service/gpu/model:indexing_map",
         "//xla/service/llvm_ir:llvm_util",
         "//xla/stream_executor:launch_dim",
diff --git a/third_party/xla/xla/service/gpu/fusions/mlir/computation_partitioner.cc b/third_party/xla/xla/service/gpu/fusions/mlir/computation_partitioner.cc
index 20bb07a321019d..3edf17babc28b4 100644
--- a/third_party/xla/xla/service/gpu/fusions/mlir/computation_partitioner.cc
+++ b/third_party/xla/xla/service/gpu/fusions/mlir/computation_partitioner.cc
@@ -48,6 +48,7 @@ limitations under the License.
 #include "xla/service/gpu/fusions/mlir/type_util.h"
 #include "xla/service/gpu/hlo_fusion_analysis.h"
 #include "xla/service/gpu/model/indexing_analysis.h"
+#include "xla/service/gpu/model/indexing_map.h"
 #include "xla/service/llvm_ir/llvm_util.h"
 #include "xla/shape.h"
 #include "xla/translate/hlo_to_mhlo/hlo_utils.h"
@@ -114,8 +115,8 @@ EpilogueSpecification EpilogueSpecification::FromIdentityIndexing(
   absl::c_copy(root->shape().dimensions(),
                std::back_inserter(result.index_ranges));
   result.roots.push_back(root);
-  result.root_indexing.push_back(mlir::AffineMap::getMultiDimIdentityMap(
-      root->shape().rank(), mlir_context));
+  result.root_indexing.push_back(
+      CreateIdentityMap(root->shape(), mlir_context));
   result.heroes.push_back(hero);
   return result;
 }
@@ -155,9 +156,8 @@ EpilogueSpecification EpilogueSpecification::FromOutputIndexing(
     auto* hero = root_to_hero[root];
     auto epilogue_indexing = ComputeEpilogueInputToOutputIndexing(
         {*hero, &analysis.fusion()}, {*root, &analysis.fusion()}, mlir_context);
-    auto root_indexing = ComposeIndexingMaps(*indexing, epilogue_indexing);
-
-    result.root_indexing.push_back(root_indexing.GetAffineMap());
+    result.root_indexing.push_back(
+        ComposeIndexingMaps(*indexing, epilogue_indexing));
   }
   result.heroes = heroes;
   result.roots = roots;
@@ -266,7 +266,7 @@ PartitionedComputation::PartitionedComputation(
     };
 
     std::vector<const HloInstruction*> roots;
-    std::vector<mlir::AffineMap> root_indexing;
+    std::vector<IndexingMap> root_indexing;
     const xla::Shape* first_root_shape = nullptr;
     for (auto* instruction : instructions) {
       if (instruction->user_count() == 0 ||
@@ -280,17 +280,16 @@ PartitionedComputation::PartitionedComputation(
             root_indexing.push_back(root_indexing.front());
           } else {
             // Bitcast from the first root to the target shape.
-            auto bitcast = GetBitcastMap(*first_root_shape,
-                                         instruction->shape(), mlir_context);
-            root_indexing.push_back(bitcast.GetAffineMap());
+            root_indexing.push_back(GetBitcastMap(
+                *first_root_shape, instruction->shape(), mlir_context));
           }
         } else {
           first_root_shape = &instruction->shape();
           while (first_root_shape->IsTuple()) {
             first_root_shape = &first_root_shape->tuple_shapes()[0];
           }
-          root_indexing.push_back(mlir::AffineMap::getMultiDimIdentityMap(
-              first_root_shape->rank(), mlir_context));
+          root_indexing.push_back(
+              CreateIdentityMap(*first_root_shape, mlir_context));
         }
       }
     }
diff --git a/third_party/xla/xla/service/gpu/fusions/mlir/computation_partitioner.h b/third_party/xla/xla/service/gpu/fusions/mlir/computation_partitioner.h
index 4a5d25361c828f..ecded59f580d25 100644
--- a/third_party/xla/xla/service/gpu/fusions/mlir/computation_partitioner.h
+++ b/third_party/xla/xla/service/gpu/fusions/mlir/computation_partitioner.h
@@ -30,6 +30,7 @@ limitations under the License.
 #include "xla/hlo/ir/hlo_instruction.h"
 #include "xla/service/gpu/fusions/fusion_emitter.h"
 #include "xla/service/gpu/hlo_fusion_analysis.h"
+#include "xla/service/gpu/model/indexing_map.h"
 
 namespace xla {
 namespace gpu {
@@ -56,7 +57,7 @@ struct EpilogueSpecification {
 
   // Indexing maps for each root output. All maps must have the same number of
   // input dimensions.
-  std::vector<mlir::AffineMap> root_indexing;
+  std::vector<IndexingMap> root_indexing;
 };
 
 // Partitions an HLO computation into subgraphs so that all users of a node have
@@ -107,7 +108,7 @@ class PartitionedComputation {
     std::vector<int64_t> index_ranges;
 
     // Maps from raw indices to root indices.
-    std::vector<mlir::AffineMap> root_indexing;
+    std::vector<IndexingMap> root_indexing;
 
     // For values that are function arguments (not function calls), stores
     // the mapping from value to the starting argument index. The arguments
diff --git a/third_party/xla/xla/service/gpu/fusions/mlir/computation_partitioner_test.cc b/third_party/xla/xla/service/gpu/fusions/mlir/computation_partitioner_test.cc
index 044d6f068ee1e1..7fb007ad3c9bd5 100644
--- a/third_party/xla/xla/service/gpu/fusions/mlir/computation_partitioner_test.cc
+++ b/third_party/xla/xla/service/gpu/fusions/mlir/computation_partitioner_test.cc
@@ -28,6 +28,7 @@ limitations under the License.
 #include "mlir/IR/MLIRContext.h"  // from @llvm-project
 #include "xla/hlo/ir/hlo_instruction.h"
 #include "xla/hlo/ir/hlo_opcode.h"
+#include "xla/service/gpu/model/indexing_analysis.h"
 #include "xla/tests/hlo_test_base.h"
 
 namespace xla {
@@ -163,7 +164,9 @@ TEST_F(ComputationPartitionerTest, Epilogue) {
       {fused_computation->GetInstructionWithName("log"),
        fused_computation->GetInstructionWithName("sign")},
       /*index_ranges=*/{1, 42},
-      {mlir::AffineMap::get(1, 0, mlir::getAffineDimExpr(0, &mlir_context_))}};
+      {CreateIdentityMap(
+          fused_computation->root_instruction()->shape().tuple_shapes(0),
+          &mlir_context_)}};
   PartitionedComputations fusion(fused_computation, &mlir_context_, {epilogue});
 
   mlir::ImplicitLocOpBuilder builder(mlir::UnknownLoc::get(&mlir_context_),
diff --git a/third_party/xla/xla/service/gpu/fusions/mlir/elemental_hlo_to_mlir.cc b/third_party/xla/xla/service/gpu/fusions/mlir/elemental_hlo_to_mlir.cc
index 192997dcf706bf..81d7ac0431523e 100644
--- a/third_party/xla/xla/service/gpu/fusions/mlir/elemental_hlo_to_mlir.cc
+++ b/third_party/xla/xla/service/gpu/fusions/mlir/elemental_hlo_to_mlir.cc
@@ -233,10 +233,7 @@ absl::StatusOr<SmallVector<Value>> EmitReduce(
   auto body =
       [&](ValueRange iter_args, ValueRange dim_values,
           ValueRange symbol_values) -> absl::StatusOr<SmallVector<Value>> {
-    auto indices =
-        b.create<ApplyIndexingOp>(dim_values, symbol_values, indexing_map)
-            .getResults();
-
+    auto indices = ApplyIndexing(indexing_map, dim_values, symbol_values, b);
     SmallVector<Value> args{iter_args};
     for (int i = 0; i < instr->operand_count() / 2; ++i) {
       TF_ASSIGN_OR_RETURN(
@@ -294,10 +291,7 @@ absl::StatusOr<SmallVector<Value>> EmitReduceWindow(
   auto body =
       [&](ValueRange iter_args, ValueRange dim_values,
           ValueRange symbol_values) -> absl::StatusOr<SmallVector<Value>> {
-    auto indices =
-        b.create<ApplyIndexingOp>(dim_values, symbol_values, indexing_map)
-            .getResults();
-
+    auto indices = ApplyIndexing(indexing_map, dim_values, symbol_values, b);
     SmallVector<Value> args{iter_args};
     for (auto [index, input] : llvm::enumerate(reduce_window->inputs())) {
       TF_ASSIGN_OR_RETURN(
@@ -484,8 +478,7 @@ SmallVector<SmallVector<Value>> GetInputIndices(
   for (auto& maps : indexing.indexing_maps) {
     CHECK_EQ(maps.size(), 1);
     CHECK(!maps.begin()->IsUndefined());
-    b.createOrFold<ApplyIndexingOp>(indices.emplace_back(), output_indices,
-                                    ValueRange{}, *maps.begin());
+    indices.push_back(ApplyIndexing(*maps.begin(), output_indices, {}, b));
   }
   return indices;
 }
@@ -574,19 +567,17 @@ absl::StatusOr<SmallVector<Value>> EmitDotLoop(
       [&](ValueRange iter_args, ValueRange dim_values,
           ValueRange symbol_values) -> absl::StatusOr<SmallVector<Value>> {
     auto lhs_indices =
-        b.create<ApplyIndexingOp>(dim_values, symbol_values, lhs_indexing_map);
-    auto rhs_indices = b.create<ApplyIndexingOp>(
-        dim_values, symbol_values.take_front(rhs_symbol_count),
-        rhs_indexing_map);
-
-    TF_ASSIGN_OR_RETURN(
-        Value lhs_value,
-        GetSingleOperandValue(operand_provider, instr,
-                              /*operand_index=*/0, lhs_indices.getResults()));
-    TF_ASSIGN_OR_RETURN(
-        Value rhs_value,
-        GetSingleOperandValue(operand_provider, instr,
-                              /*operand_index=*/1, rhs_indices.getResults()));
+        ApplyIndexing(lhs_indexing_map, dim_values, symbol_values, b);
+    auto rhs_indices =
+        ApplyIndexing(rhs_indexing_map, dim_values,
+                      symbol_values.take_front(rhs_symbol_count), b);
+
+    TF_ASSIGN_OR_RETURN(Value lhs_value, GetSingleOperandValue(
+                                             operand_provider, instr,
+                                             /*operand_index=*/0, lhs_indices));
+    TF_ASSIGN_OR_RETURN(Value rhs_value, GetSingleOperandValue(
+                                             operand_provider, instr,
+                                             /*operand_index=*/1, rhs_indices));
     Value accum = iter_args[0];
 
     TF_ASSIGN_OR_RETURN(
@@ -687,22 +678,14 @@ Value ApplyAffineExpr(mlir::AffineExpr expr, ValueRange dims,
   return b.createOrFold<mlir::affine::AffineApplyOp>(expr, args);
 }
 
-SmallVector<Value> ApplyAffineMap(mlir::AffineMap map, ValueRange dims,
-                                  ValueRange symbols, ImplicitLocOpBuilder& b) {
-  CHECK_EQ(map.getNumDims(), dims.size());
-  CHECK_EQ(map.getNumSymbols(), symbols.size());
-  SmallVector<Value> result;
-  result.reserve(map.getNumResults());
-  for (auto expr : map.getResults()) {
-    result.push_back(ApplyAffineExpr(expr, dims, symbols, b));
-  }
-  return result;
-}
-
 SmallVector<Value> ApplyIndexing(const IndexingMap& map, ValueRange dims,
                                  ValueRange symbols, ImplicitLocOpBuilder& b) {
   SmallVector<Value> results;
-  b.createOrFold<ApplyIndexingOp>(results, dims, symbols, map);
+  for (unsigned int i = 0; i < map.GetAffineMap().getNumResults(); ++i) {
+    SmallVector<Value> result;
+    b.createOrFold<ApplyIndexingOp>(result, dims, symbols, map.GetSubMap(i));
+    results.append(result);
+  }
   return results;
 }
 
@@ -720,11 +703,25 @@ Value CheckConstraint(mlir::Value constrained_value, Interval range,
 
 Value CheckConstraints(const IndexingMap& map, ValueRange dims,
                        ValueRange symbols, ImplicitLocOpBuilder& b) {
+  llvm::SmallVector<mlir::AffineExpr> expressions;
+  for (auto&& [expression, _] : map.GetConstraints()) {
+    expressions.push_back(expression);
+  }
+
+  // Construct an indexing for the constraints, so we can use `apply_indexing`.
+  auto input_map = map.GetAffineMap();
+  IndexingMap constraints_map{
+      mlir::AffineMap::get(input_map.getNumDims(), input_map.getNumSymbols(),
+                           expressions, input_map.getContext()),
+      map.GetDimVars(), map.GetRangeVars(), map.GetRTVars()};
+  llvm::SmallVector<Value> constraints_values =
+      ApplyIndexing(constraints_map, dims, symbols, b);
+
   Value ret = b.create<ConstantOp>(b.getIntegerAttr(b.getI1Type(), 1));
-  for (auto&& [expression, range] : map.GetConstraints()) {
+  for (auto&& [value, expression_and_range] :
+       llvm::zip(constraints_values, map.GetConstraints())) {
     ret = b.create<AndIOp>(
-        ret, CheckConstraint(ApplyAffineExpr(expression, dims, symbols, b),
-                             range, b));
+        ret, CheckConstraint(value, expression_and_range.second, b));
   }
   for (auto&& [index, bound] : llvm::enumerate(map.GetDimensionBounds())) {
     ret = b.create<AndIOp>(ret, CheckConstraint(dims[index], bound, b));
@@ -844,9 +841,7 @@ absl::StatusOr<SmallVector<Value>> HloToMlir(
                          *first_shape, *operand_index_shape)) {
           auto operand_map =
               GetBitcastMap(*first_shape, *operand_index_shape, mlir_context);
-          operand_indices =
-              builder.create<ApplyIndexingOp>(indices, operand_map)
-                  .getResults();
+          operand_indices = ApplyIndexing(operand_map, indices, {}, builder);
         } else {
           operand_indices = indices;
         }
@@ -1237,15 +1232,10 @@ absl::StatusOr<SmallVector<Value>> SubgraphToMlir(
                    std::back_inserter(results));
       continue;
     }
-    TF_RET_CHECK(indexing.getNumDims() + indexing.getNumSymbols() ==
-                 indices.size())
-        << "Incorrect number of indices (got " << indices.size()
-        << ", expected " << indexing.getNumDims() << " dims and "
-        << indexing.getNumSymbols() << " symbols) in " << subgraph.ToString();
-    int num_dims = indexing.getNumDims();
+    int num_dims = indexing.GetAffineMap().getNumDims();
     auto root_indices =
-        ApplyAffineMap(indexing, /*dims=*/indices.take_front(num_dims),
-                       /*symbols=*/indices.drop_front(num_dims), builder);
+        ApplyIndexing(indexing, /*dims=*/indices.take_front(num_dims),
+                      /*symbols=*/indices.drop_front(num_dims), builder);
     TF_ASSIGN_OR_RETURN(auto root_results, emit_instr(root, root_indices));
     results.append(root_results.begin(), root_results.end());
   }
diff --git a/third_party/xla/xla/service/gpu/fusions/mlir/elemental_hlo_to_mlir.h b/third_party/xla/xla/service/gpu/fusions/mlir/elemental_hlo_to_mlir.h
index bb20a88951a38d..01a705873b0673 100644
--- a/third_party/xla/xla/service/gpu/fusions/mlir/elemental_hlo_to_mlir.h
+++ b/third_party/xla/xla/service/gpu/fusions/mlir/elemental_hlo_to_mlir.h
@@ -86,11 +86,7 @@ mlir::Value ApplyAffineExpr(mlir::AffineExpr expr, mlir::ValueRange dims,
                             mlir::ValueRange symbols,
                             mlir::ImplicitLocOpBuilder& b);
 
-// Creates affine.apply ops for each result of the given map.
-llvm::SmallVector<mlir::Value> ApplyAffineMap(mlir::AffineMap map,
-                                              mlir::ValueRange dims,
-                                              mlir::ValueRange symbols,
-                                              mlir::ImplicitLocOpBuilder& b);
+// Creates an `apply_indexing` op for the given map.
 llvm::SmallVector<mlir::Value> ApplyIndexing(const IndexingMap& map,
                                              mlir::ValueRange dims,
                                              mlir::ValueRange symbols,
diff --git a/third_party/xla/xla/service/gpu/fusions/mlir/elemental_hlo_to_mlir_test.cc b/third_party/xla/xla/service/gpu/fusions/mlir/elemental_hlo_to_mlir_test.cc
index 4ffe6e5613a436..195db04a7b06f7 100644
--- a/third_party/xla/xla/service/gpu/fusions/mlir/elemental_hlo_to_mlir_test.cc
+++ b/third_party/xla/xla/service/gpu/fusions/mlir/elemental_hlo_to_mlir_test.cc
@@ -39,6 +39,8 @@ limitations under the License.
 #include "xla/service/gpu/fusions/mlir/computation_partitioner.h"
 #include "xla/service/gpu/fusions/mlir/ir/xla_gpu_ops.h"
 #include "xla/service/gpu/launch_dimensions.h"
+#include "xla/service/gpu/model/indexing_analysis.h"
+#include "xla/service/gpu/model/indexing_map.h"
 #include "xla/service/hlo_parser.h"
 #include "xla/service/llvm_ir/llvm_util.h"
 #include "xla/status_macros.h"
@@ -233,11 +235,12 @@ TEST_F(ElementalHloToMlirTest, ReduceWindow) {
     // CHECK:      %[[INIT:.*]] = tensor.extract %[[ARG1]][]
     // CHECK:      %[[RET:.*]] = scf.for %[[I:.*]] = %[[C0]] to %[[C7]]
     // CHECK-SAME:   step %[[C1]] iter_args(%[[ACC:.*]] = %[[INIT]])
-    // CHECK:      %[[J:[a-z0-9]+]]:2 = xla_gpu.apply_indexing
-    // CHECK-SAME:     affine_map<(d0, d1)[s0] -> (d0 * 4, d1 + s0 - 3)>
-    // CHECK-SAME:     (%[[Y]] in [0, 2], %[[Z]] in [0, 7])[%[[I]] in [0, 6]]
+    // CHECK:      %[[J0:.*]] = xla_gpu.apply_indexing affine_map<(d0) -> (d0 * 4)>
+    // CHECK-SAME:     (%[[Y]] in [0, 2])
+    // CHECK:      %[[J1:.*]] = xla_gpu.apply_indexing affine_map<(d0)[s0] -> (d0 + s0 - 3)>
+    // CHECK-SAME:     (%[[Z]] in [0, 7])[%[[I]] in [0, 6]]
     // CHECK:          %[[VAL:.*]] = tensor.extract %[[ARG0]]
-    // CHECK-SAME:        [%[[X]], %[[J]]#0, %[[J]]#1]
+    // CHECK-SAME:        [%[[X]], %[[J0]], %[[J1]]]
     // CHECK:          %[[UPD:.*]] = func.call @add_sum(%[[ACC]],
     // CHECK-SAME:                                      %[[VAL]])
     // CHECK:          scf.yield %[[UPD]]
@@ -405,9 +408,8 @@ TEST_F(ElementalHloToMlirTest, Pad) {
     // CHECK-DAG:    %[[C1:.*]] = arith.constant 1
     // CHECK-DAG:    %[[C4:.*]] = arith.constant 4
     // CHECK-DAG:    %[[C7:.*]] = arith.constant 7
-    // CHECK:        %[[CONSTRAINT_VAL:.*]] = affine.apply
-    // CHECK-SAME:     <()[s0] -> (s0 - ((s0 - 1) floordiv 2) * 2 - 1)>
-    // CHECK-SAME:     ()[%[[X]]]
+    // CHECK:        %[[CONSTRAINT_VAL:.*]] = xla_gpu.apply_indexing
+    // CHECK-SAME:     <(d0) -> ((d0 - 1) mod 2)>(%[[X]] in [1, 7])
     // CHECK:        %[[CONSTRAINT:.*]] = arith.cmpi eq, %[[CONSTRAINT_VAL]], %[[C0]]
     // CHECK-DAG:        %[[X_L:.*]] = arith.cmpi sge, %[[X]], %[[C1]]
     // CHECK-DAG:        %[[X_H:.*]] = arith.cmpi sle, %[[X]], %[[C7]]
@@ -418,10 +420,13 @@ TEST_F(ElementalHloToMlirTest, Pad) {
     // CHECK:        %[[Y_BOUNDS:.*]] = arith.andi %[[Y_L]], %[[Y_H]]
     // CHECK:        %[[FROM_INPUT:.*]] = arith.andi %[[X_AND_CONSTRAINT]], %[[Y_BOUNDS]]
     // CHECK:        %[[RET:.*]] = scf.if %[[FROM_INPUT]]
-    // CHECK:          %[[IN:.*]]:2 = xla_gpu.apply_indexing
-    // CHECK-SAME:         <(d0, d1) -> ((d0 - 1) floordiv 2, d1 - 4)>
-    // CHECK-SAME:         (%[[X]] in [1, 7], %[[Y]] in [4, 7])
-    // CHECK:          %[[VAL:.*]] = tensor.extract %[[ARG0]][%[[IN]]#0, %[[IN]]#1]
+    // CHECK:          %[[IN0:.*]] = xla_gpu.apply_indexing
+    // CHECK-SAME:         <(d0) -> ((d0 - 1) floordiv 2)>
+    // CHECK-SAME:         (%[[X]] in [1, 7])
+    // CHECK:          %[[IN1:.*]] = xla_gpu.apply_indexing
+    // CHECK-SAME:         <(d0) -> (d0 - 4)>
+    // CHECK-SAME:         (%[[Y]] in [4, 7])
+    // CHECK:          %[[VAL:.*]] = tensor.extract %[[ARG0]][%[[IN0]], %[[IN1]]]
     // CHECK:          scf.yield %[[VAL]]
     // CHECK:        } else {
     // CHECK:          %[[PAD_VAL:.*]] = tensor.extract %[[ARG1]][]
@@ -447,9 +452,8 @@ TEST_F(ElementalHloToMlirTest, PadUnsigned) {
     // CHECK-DAG:    %[[C1:.*]] = arith.constant 1
     // CHECK-DAG:    %[[C4:.*]] = arith.constant 4
     // CHECK-DAG:    %[[C7:.*]] = arith.constant 7
-    // CHECK:        %[[CONSTRAINT_VAL:.*]] = affine.apply
-    // CHECK-SAME:     <()[s0] -> (s0 - ((s0 - 1) floordiv 2) * 2 - 1)>
-    // CHECK-SAME:     ()[%[[X]]]
+    // CHECK:        %[[CONSTRAINT_VAL:.*]] = xla_gpu.apply_indexing
+    // CHECK-SAME:     <(d0) -> ((d0 - 1) mod 2)>(%[[X]] in [1, 7])
     // CHECK:        %[[CONSTRAINT:.*]] = arith.cmpi eq, %[[CONSTRAINT_VAL]], %[[C0]]
     // CHECK-DAG:        %[[X_L:.*]] = arith.cmpi sge, %[[X]], %[[C1]]
     // CHECK-DAG:        %[[X_H:.*]] = arith.cmpi sle, %[[X]], %[[C7]]
@@ -460,10 +464,13 @@ TEST_F(ElementalHloToMlirTest, PadUnsigned) {
     // CHECK:        %[[Y_BOUNDS:.*]] = arith.andi %[[Y_L]], %[[Y_H]]
     // CHECK:        %[[FROM_INPUT:.*]] = arith.andi %[[X_AND_CONSTRAINT]], %[[Y_BOUNDS]]
     // CHECK:        %[[RET:.*]] = scf.if %[[FROM_INPUT]]
-    // CHECK:          %[[IN:.*]]:2 = xla_gpu.apply_indexing
-    // CHECK-SAME:         <(d0, d1) -> ((d0 - 1) floordiv 2, d1 - 4)>
-    // CHECK-SAME:         (%[[X]] in [1, 7], %[[Y]] in [4, 7])
-    // CHECK:          %[[VAL:.*]] = tensor.extract %[[ARG0]][%[[IN]]#0, %[[IN]]#1]
+    // CHECK:          %[[IN0:.*]] = xla_gpu.apply_indexing
+    // CHECK-SAME:         <(d0) -> ((d0 - 1) floordiv 2)>
+    // CHECK-SAME:         (%[[X]] in [1, 7])
+    // CHECK:          %[[IN1:.*]] = xla_gpu.apply_indexing
+    // CHECK-SAME:         <(d0) -> (d0 - 4)>
+    // CHECK-SAME:         (%[[Y]] in [4, 7])
+    // CHECK:          %[[VAL:.*]] = tensor.extract %[[ARG0]][%[[IN0]], %[[IN1]]]
     // CHECK:          %[[CAST0:.*]] = builtin.unrealized_conversion_cast %[[VAL]]
     // CHECK:          scf.yield %[[CAST0]]
     // CHECK:        } else {
@@ -784,10 +791,13 @@ TEST_F(ElementalHloToMlirTest, ConvolutionSimple) {
     // CHECK-NEXT: %[[R1:.+]] = scf.for %[[Y:.+]] = %[[C0]] to %[[C5]] step %[[C1]] iter_args(%[[A1:.+]] = %[[A0]]) -> (f32) {
     // CHECK-NEXT: %[[R2:.+]] = scf.for %[[I:.+]] = %[[C0]] to %[[C4]] step %[[C1]] iter_args(%[[ACC:.+]] = %[[A1]]) -> (f32) {
     // CHECK:      %[[R3:.+]] = scf.if {{.+}} -> (f32) {
-    // CHECK:      %[[XX:.+]]:2 = xla_gpu.apply_indexing
-    // CHECK-SAME:   affine_map<(d0, d1)[s0, s1] -> (d0 + s0, d1 + s1)>
-    // CHECK-SAME:   (%[[W]] in [0, 5], %[[H]] in [0, 7])[%[[X]] in [0, 2], %[[Y]] in [0, 4]]
-    // CHECK-DAG:    %[[VL:.+]] = tensor.extract %[[LHS]][%[[B]], %[[XX]]#0, %[[XX]]#1, %[[I]]] : tensor<2x8x12x4xf32>
+    // CHECK:      %[[XX0:.+]] = xla_gpu.apply_indexing
+    // CHECK-SAME:   affine_map<(d0)[s0] -> (d0 + s0)>
+    // CHECK-SAME:   (%[[W]] in [0, 5])[%[[X]] in [0, 2]]
+    // CHECK:      %[[XX1:.+]] = xla_gpu.apply_indexing
+    // CHECK-SAME:   affine_map<(d0)[s0] -> (d0 + s0)>
+    // CHECK-SAME:   (%[[H]] in [0, 7])[%[[Y]] in [0, 4]]
+    // CHECK-DAG:    %[[VL:.+]] = tensor.extract %[[LHS]][%[[B]], %[[XX0]], %[[XX1]], %[[I]]] : tensor<2x8x12x4xf32>
     // CHECK-DAG:    %[[VR:.+]] = tensor.extract %[[RHS]][%[[I]], %[[X]], %[[Y]], %[[O]]] : tensor<4x3x5x16xf32>
     // CHECK:        %[[MUL:.+]] = arith.mulf %[[VL]], %[[VR]] : f32
     // CHECK-NEXT:   %[[ADD:.+]] = arith.addf %[[ACC]], %[[MUL]] : f32
@@ -827,11 +837,13 @@ TEST_F(ElementalHloToMlirTest, ConvolutionWithWindowStrides) {
     // CHECK-NEXT: %[[R1:.+]] = scf.for %[[Y:.+]] = %[[C0]] to %[[C5]] step %[[C1]] iter_args(%[[A1:.+]] = %[[A0]]) -> (f32) {
     // CHECK-NEXT: %[[R2:.+]] = scf.for %[[I:.+]] = %[[C0]] to %[[C4]] step %[[C1]] iter_args(%[[ACC:.+]] = %[[A1]]) -> (f32) {
     // CHECK:      %[[R3:.+]] = scf.if {{.+}} -> (f32) {
-
-    // CHECK:        %[[XX:.+]]:2 = xla_gpu.apply_indexing
-    // CHECK-SAME:     affine_map<(d0, d1)[s0, s1] -> (d0 * 2 + s0, d1 * 2 + s1)>
-    // CHECK-SAME:     (%[[W]] in [0, 2], %[[H]] in [0, 3])[%[[X]] in [0, 2], %[[Y]] in [0, 4]]
-    // CHECK-DAG:    %[[VL:.+]] = tensor.extract %[[LHS]][%[[B]], %[[XX]]#0, %[[XX]]#1, %[[I]]] : tensor<2x8x12x4xf32>
+    // CHECK:        %[[XX0:.+]] = xla_gpu.apply_indexing
+    // CHECK-SAME:     affine_map<(d0)[s0] -> (d0 * 2 + s0)>
+    // CHECK-SAME:     (%[[W]] in [0, 2])[%[[X]] in [0, 2]]
+    // CHECK:        %[[XX1:.+]] = xla_gpu.apply_indexing
+    // CHECK-SAME:     affine_map<(d0)[s0] -> (d0 * 2 + s0)>
+    // CHECK-SAME:     (%[[H]] in [0, 3])[%[[Y]] in [0, 4]]
+    // CHECK-DAG:    %[[VL:.+]] = tensor.extract %[[LHS]][%[[B]], %[[XX0]], %[[XX1]], %[[I]]] : tensor<2x8x12x4xf32>
     // CHECK-DAG:    %[[VR:.+]] = tensor.extract %[[RHS]][%[[I]], %[[X]], %[[Y]], %[[O]]] : tensor<4x3x5x16xf32>
     // CHECK:        %[[MUL:.+]] = arith.mulf %[[VL]], %[[VR]] : f32
     // CHECK-NEXT:   %[[ADD:.+]] = arith.addf %[[ACC]], %[[MUL]] : f32
@@ -873,19 +885,22 @@ TEST_F(ElementalHloToMlirTest, ConvolutionWithPadding) {
     // CHECK:      %[[R0:.+]] = scf.for %[[X:.+]] = %[[C0]] to %[[C3]] step %[[C1]] iter_args(%[[A0:.+]] = %[[INIT]]) -> (f32) {
     // CHECK-NEXT: %[[R1:.+]] = scf.for %[[Y:.+]] = %[[C0]] to %[[C5]] step %[[C1]] iter_args(%[[A1:.+]] = %[[A0]]) -> (f32) {
     // CHECK-NEXT: %[[R2:.+]] = scf.for %[[I:.+]] = %[[C0]] to %[[C4]] step %[[C1]] iter_args(%[[ACC:.+]] = %[[A1]]) -> (f32) {
-    // CHECK-DAG:  %[[TESTX:.+]] = affine.apply affine_map<()[s0, s1] -> (s0 + s1)>()[%[[X]], %[[W]]]
+    // CHECK-DAG:  %[[TESTX:.+]] = xla_gpu.apply_indexing affine_map<(d0)[s0] -> (d0 + s0)>(%[[W]] in [0, 7])[%[[X]] in [0, 2]]
     // CHECK-DAG:  %[[TXGE:.+]] = arith.cmpi sge, %[[TESTX]], %[[C1]] : index
     // CHECK-DAG:  %[[TXLE:.+]] = arith.cmpi sle, %[[TESTX]], %[[C8]] : index
     // CHECK-DAG:  %[[TX:.+]] = arith.andi %[[TXGE]], %[[TXLE]] : i1
-    // CHECK-DAG:  %[[TESTY:.+]] = affine.apply affine_map<()[s0, s1] -> (s0 + s1)>()[%[[Y]], %[[H]]]
+    // CHECK-DAG:  %[[TESTY:.+]] = xla_gpu.apply_indexing affine_map<(d0)[s0] -> (d0 + s0)>(%[[H]] in [0, 11])[%[[Y]] in [0, 4]]
     // CHECK-DAG:  %[[TYGE:.+]] = arith.cmpi sge, %[[TESTY]], %[[C2]] : index
     // CHECK-DAG:  %[[TYLE:.+]] = arith.cmpi sle, %[[TESTY]], %[[C13]] : index
     // CHECK-DAG:  %[[TY:.+]] = arith.andi %[[TYGE]], %[[TYLE]] : i1
     // CHECK:      %[[R3:.+]] = scf.if {{.+}} -> (f32) {
-    // CHECK:        %[[XX:.+]]:2 = xla_gpu.apply_indexing
-    // CHECK-SAME:     affine_map<(d0, d1)[s0, s1] -> (d0 + s0 - 1, d1 + s1 - 2)>
-    // CHECK-SAME:     (%[[W]] in [0, 7], %[[H]] in [0, 11])[%[[X]] in [0, 2], %[[Y]] in [0, 4]]
-    // CHECK-DAG:    %[[VL:.+]] = tensor.extract %[[LHS]][%[[B]], %[[XX]]#0, %[[XX]]#1, %[[I]]] : tensor<2x8x12x4xf32>
+    // CHECK:        %[[XX0:.+]] = xla_gpu.apply_indexing
+    // CHECK-SAME:     affine_map<(d0)[s0] -> (d0 + s0 - 1)>
+    // CHECK-SAME:     (%[[W]] in [0, 7])[%[[X]] in [0, 2]]
+    // CHECK:        %[[XX1:.+]] = xla_gpu.apply_indexing
+    // CHECK-SAME:     affine_map<(d0)[s0] -> (d0 + s0 - 2)>
+    // CHECK-SAME:     (%[[H]] in [0, 11])[%[[Y]] in [0, 4]]
+    // CHECK-DAG:    %[[VL:.+]] = tensor.extract %[[LHS]][%[[B]], %[[XX0]], %[[XX1]], %[[I]]] : tensor<2x8x12x4xf32>
     // CHECK-DAG:    %[[VR:.+]] = tensor.extract %[[RHS]][%[[I]], %[[X]], %[[Y]], %[[O]]] : tensor<4x3x5x16xf32>
     // CHECK:        %[[MUL:.+]] = arith.mulf %[[VL]], %[[VR]] : f32
     // CHECK-NEXT:   %[[ADD:.+]] = arith.addf %[[ACC]], %[[MUL]] : f32
@@ -924,15 +939,18 @@ TEST_F(ElementalHloToMlirTest, ConvolutionWithLhsDilation) {
     // CHECK:      %[[R0:.+]] = scf.for %[[X:.+]] = %[[C0]] to %[[C3]] step %[[C1]] iter_args(%[[A0:.+]] = %[[INIT]]) -> (f32) {
     // CHECK-NEXT: %[[R1:.+]] = scf.for %[[Y:.+]] = %[[C0]] to %[[C5]] step %[[C1]] iter_args(%[[A1:.+]] = %[[A0]]) -> (f32) {
     // CHECK-NEXT: %[[R2:.+]] = scf.for %[[I:.+]] = %[[C0]] to %[[C4]] step %[[C1]] iter_args(%[[ACC:.+]] = %[[A1]]) -> (f32) {
-    // CHECK-DAG:  %[[TESTX:.+]] = affine.apply affine_map<()[s0, s1] -> ((s0 + s1) mod 2)>()[%[[X]], %[[W]]]
+    // CHECK-DAG:  %[[TESTX:.+]] = xla_gpu.apply_indexing affine_map<(d0)[s0] -> ((d0 + s0) mod 2)>(%[[W]] in [0, 12])[%[[X]] in [0, 2]]
     // CHECK-DAG:  %[[TX:.+]] = arith.cmpi eq, %[[TESTX]], %[[C0]] : index
-    // CHECK-DAG:  %[[TESTY:.+]] = affine.apply affine_map<()[s0, s1] -> ((s0 + s1) mod 2)>()[%[[Y]], %[[H]]]
+    // CHECK-DAG:  %[[TESTY:.+]] = xla_gpu.apply_indexing affine_map<(d0)[s0] -> ((d0 + s0) mod 2)>(%[[H]] in [0, 18])[%[[Y]] in [0, 4]]
     // CHECK-DAG:  %[[TY:.+]] = arith.cmpi eq, %[[TESTY]], %[[C0]] : index
     // CHECK:      %[[R3:.+]] = scf.if {{.+}} -> (f32) {
-    // CHECK:        %[[XX:.+]]:2 = xla_gpu.apply_indexing
-    // CHECK-SAME:     affine_map<(d0, d1)[s0, s1] -> ((d0 + s0) floordiv 2, (d1 + s1) floordiv 2)>
-    // CHECK-SAME:     (%[[W]] in [0, 12], %[[H]] in [0, 18])[%[[X]] in [0, 2], %[[Y]] in [0, 4]]
-    // CHECK-DAG:    %[[VL:.+]] = tensor.extract %[[LHS]][%[[B]], %[[XX]]#0, %[[XX]]#1, %[[I]]] : tensor<2x8x12x4xf32>
+    // CHECK:        %[[XX0:.+]] = xla_gpu.apply_indexing
+    // CHECK-SAME:     affine_map<(d0)[s0] -> ((d0 + s0) floordiv 2)>
+    // CHECK-SAME:     (%[[W]] in [0, 12])[%[[X]] in [0, 2]]
+    // CHECK:        %[[XX1:.+]] = xla_gpu.apply_indexing
+    // CHECK-SAME:     affine_map<(d0)[s0] -> ((d0 + s0) floordiv 2)>
+    // CHECK-SAME:     (%[[H]] in [0, 18])[%[[Y]] in [0, 4]]
+    // CHECK-DAG:    %[[VL:.+]] = tensor.extract %[[LHS]][%[[B]], %[[XX0]], %[[XX1]], %[[I]]] : tensor<2x8x12x4xf32>
     // CHECK-DAG:    %[[VR:.+]] = tensor.extract %[[RHS]][%[[I]], %[[X]], %[[Y]], %[[O]]] : tensor<4x3x5x16xf32>
     // CHECK:        %[[MUL:.+]] = arith.mulf %[[VL]], %[[VR]] : f32
     // CHECK-NEXT:   %[[ADD:.+]] = arith.addf %[[ACC]], %[[MUL]] : f32
@@ -972,10 +990,13 @@ TEST_F(ElementalHloToMlirTest, ConvolutionWithRhsDilation) {
     // CHECK-NEXT: %[[R1:.+]] = scf.for %[[Y:.+]] = %[[C0]] to %[[C5]] step %[[C1]] iter_args(%[[A1:.+]] = %[[A0]]) -> (f32) {
     // CHECK-NEXT: %[[R2:.+]] = scf.for %[[I:.+]] = %[[C0]] to %[[C4]] step %[[C1]] iter_args(%[[ACC:.+]] = %[[A1]]) -> (f32) {
     // CHECK:      %[[R3:.+]] = scf.if {{.+}} -> (f32) {
-    // CHECK:        %[[XX:.+]]:2 = xla_gpu.apply_indexing
-    // CHECK-SAME:     affine_map<(d0, d1)[s0, s1] -> (d0 + s0 * 2, d1 + s1 * 2)>
-    // CHECK-SAME:     (%[[W]] in [0, 3], %[[H]] in [0, 3])[%[[X]] in [0, 2], %[[Y]] in [0, 4]]
-    // CHECK-DAG:    %[[VL:.+]] = tensor.extract %[[LHS]][%[[B]], %[[XX]]#0, %[[XX]]#1, %[[I]]] : tensor<2x8x12x4xf32>
+    // CHECK:        %[[XX0:.+]] = xla_gpu.apply_indexing
+    // CHECK-SAME:     affine_map<(d0)[s0] -> (d0 + s0 * 2)>
+    // CHECK-SAME:     (%[[W]] in [0, 3])[%[[X]] in [0, 2]]
+    // CHECK:        %[[XX1:.+]] = xla_gpu.apply_indexing
+    // CHECK-SAME:     affine_map<(d0)[s0] -> (d0 + s0 * 2)>
+    // CHECK-SAME:     (%[[H]] in [0, 3])[%[[Y]] in [0, 4]]
+    // CHECK-DAG:    %[[VL:.+]] = tensor.extract %[[LHS]][%[[B]], %[[XX0]], %[[XX1]], %[[I]]] : tensor<2x8x12x4xf32>
     // CHECK-DAG:    %[[VR:.+]] = tensor.extract %[[RHS]][%[[I]], %[[X]], %[[Y]], %[[O]]] : tensor<4x3x5x16xf32>
     // CHECK:        %[[MUL:.+]] = arith.mulf %[[VL]], %[[VR]] : f32
     // CHECK-NEXT:   %[[ADD:.+]] = arith.addf %[[ACC]], %[[MUL]] : f32
@@ -1015,11 +1036,19 @@ TEST_F(ElementalHloToMlirTest, ConvolutionWithFeatureGroupCount) {
     // CHECK-NEXT: %[[R1:.+]] = scf.for %[[Y:.+]] = %[[C0]] to %[[C5]] step %[[C1]] iter_args(%[[A1:.+]] = %[[A0]]) -> (f32) {
     // CHECK-NEXT: %[[R2:.+]] = scf.for %[[I:.+]] = %[[C0]] to %[[C2]] step %[[C1]] iter_args(%[[ACC:.+]] = %[[A1]]) -> (f32) {
     // CHECK:      %[[R3:.+]] = scf.if {{.+}} -> (f32) {
-    // CHECK:        %[[XX:.+]]:3 = xla_gpu.apply_indexing
-    // CHECK-SAME:     affine_map<(d0, d1, d2)[s0, s1, s2] -> (d0 + s0, d1 + s1, (d2 floordiv 8) * 2 + s2)>
-    // CHECK-SAME:     (%[[W]] in [0, 5], %[[H]] in [0, 7], %[[O]] in [0, 15])
-    // CHECK-SAME:     [%[[X]] in [0, 2], %[[Y]] in [0, 4], %[[I]] in [0, 1]]
-    // CHECK-DAG:    %[[VL:.+]] = tensor.extract %[[LHS]][%[[B]], %[[XX]]#0, %[[XX]]#1, %[[XX]]#2] : tensor<2x8x12x4xf32>
+    // CHECK:        %[[XX0:.+]] = xla_gpu.apply_indexing
+    // CHECK-SAME:     affine_map<(d0)[s0] -> (d0 + s0)>
+    // CHECK-SAME:     (%[[W]] in [0, 5])
+    // CHECK-SAME:     [%[[X]] in [0, 2]]
+    // CHECK:        %[[XX1:.+]] = xla_gpu.apply_indexing
+    // CHECK-SAME:     affine_map<(d0)[s0] -> (d0 + s0)>
+    // CHECK-SAME:     (%[[H]] in [0, 7])
+    // CHECK-SAME:     [%[[Y]] in [0, 4]]
+    // CHECK:        %[[XX2:.+]] = xla_gpu.apply_indexing
+    // CHECK-SAME:     affine_map<(d0)[s0] -> ((d0 floordiv 8) * 2 + s0)>
+    // CHECK-SAME:     (%[[O]] in [0, 15])
+    // CHECK-SAME:     [%[[I]] in [0, 1]]
+    // CHECK-DAG:    %[[VL:.+]] = tensor.extract %[[LHS]][%[[B]], %[[XX0]], %[[XX1]], %[[XX2]]] : tensor<2x8x12x4xf32>
     // CHECK-DAG:    %[[VR:.+]] = tensor.extract %[[RHS]][%[[I]], %[[X]], %[[Y]], %[[O]]] : tensor<2x3x5x16xf32>
     // CHECK:        %[[MUL:.+]] = arith.mulf %[[VL]], %[[VR]] : f32
     // CHECK-NEXT:   %[[ADD:.+]] = arith.addf %[[ACC]], %[[MUL]] : f32
@@ -1061,11 +1090,15 @@ TEST_F(ElementalHloToMlirTest, ConvolutionWithBatchGroupCount) {
     // CHECK-NEXT: %[[R2:.+]] = scf.for %[[I:.+]] = %[[C0]] to %[[C4]] step %[[C1]] iter_args(%[[A2:.+]] = %[[A1]]) -> (f32) {
     // CHECK-NEXT: %[[R3:.+]] = scf.for %[[G:.+]] = %[[C0]] to %[[C2]] step %[[C1]] iter_args(%[[ACC:.+]] = %[[A2]]) -> (f32) {
     // CHECK:      %[[R4:.+]] = scf.if {{.+}} -> (f32) {
-    // CHECK:        %[[XX:.+]]:2 = xla_gpu.apply_indexing
-    // CHECK-SAME:     affine_map<(d0, d1)[s0, s1] -> (d0 + s0, d1 + s1)>
-    // CHECK-SAME:     (%[[W]] in [0, 5], %[[H]] in [0, 7])
-    // CHECK-SAME:     [%[[X]] in [0, 2], %[[Y]] in [0, 4]]
-    // CHECK-DAG:    %[[VL:.+]] = tensor.extract %[[LHS]][%[[G]], %[[XX]]#0, %[[XX]]#1, %[[I]]] : tensor<2x8x12x4xf32>
+    // CHECK:        %[[XX0:.+]] = xla_gpu.apply_indexing
+    // CHECK-SAME:     affine_map<(d0)[s0] -> (d0 + s0)>
+    // CHECK-SAME:     (%[[W]] in [0, 5])
+    // CHECK-SAME:     [%[[X]] in [0, 2]]
+    // CHECK:        %[[XX1:.+]] = xla_gpu.apply_indexing
+    // CHECK-SAME:     affine_map<(d0)[s0] -> (d0 + s0)>
+    // CHECK-SAME:     (%[[H]] in [0, 7])
+    // CHECK-SAME:     [%[[Y]] in [0, 4]]
+    // CHECK-DAG:    %[[VL:.+]] = tensor.extract %[[LHS]][%[[G]], %[[XX0]], %[[XX1]], %[[I]]] : tensor<2x8x12x4xf32>
     // CHECK-DAG:    %[[VR:.+]] = tensor.extract %[[RHS]][%[[I]], %[[X]], %[[Y]], %[[O]]] : tensor<4x3x5x16xf32>
     // CHECK:        %[[MUL:.+]] = arith.mulf %[[VL]], %[[VR]] : f32
     // CHECK-NEXT:   %[[ADD:.+]] = arith.addf %[[ACC]], %[[MUL]] : f32
@@ -1236,8 +1269,11 @@ TEST_F(ElementalHloToMlirTest, Epilogue) {
         epilogue.roots.push_back(entry->GetInstructionWithName("add"));
         epilogue.index_ranges = {2, 16, 17};
         epilogue.root_indexing.push_back(
-            mlir::AffineMap::getMultiDimIdentityMap(3, &context_)
-                .getSubMap({0, 2, 1}));
+            IndexingMap{mlir::AffineMap::getMultiDimIdentityMap(3, &context_)
+                            .getSubMap({0, 2, 1}),
+                        DimVarsFromTensorSizes({2, 17, 17}),
+                        {},
+                        {}});
         return epilogue;
       }));
 }
diff --git a/third_party/xla/xla/service/gpu/fusions/reduction_mlir.cc b/third_party/xla/xla/service/gpu/fusions/reduction_mlir.cc
index c36bc3b61c0cae..4e9effcabb2e8a 100644
--- a/third_party/xla/xla/service/gpu/fusions/reduction_mlir.cc
+++ b/third_party/xla/xla/service/gpu/fusions/reduction_mlir.cc
@@ -208,10 +208,12 @@ llvm::SmallVector<Value> MlirReductionFusion::EmitReduction(
   auto thread_id = state.thread_and_block_ids[0];
   Value cst_true = b.create<ma::ConstantOp>(b.getOneAttr(b.getI1Type()));
 
-  auto delinearized = DelinearizeInBoundsIndex(mlir::getAffineDimExpr(0, ctx),
-                                               threads_per_block);
-  auto thread_ids = mlir_converter::ApplyAffineMap(
-      mlir::AffineMap::get(1, 0, delinearized, ctx), {thread_id}, {}, b);
+  auto thread_indexing = GetBitcastMap(
+      ShapeUtil::MakeShapeWithDescendingLayout(
+          U8, {tiling.GetNumThreadsPerBlock()}),
+      ShapeUtil::MakeShapeWithDescendingLayout(U8, threads_per_block), ctx);
+  auto thread_ids =
+      mlir_converter::ApplyIndexing(thread_indexing, {thread_id}, {}, b);
 
   auto warp_id = b.create<ma::DivUIOp>(
       reduction_info().IsRowReduction()
@@ -246,7 +248,7 @@ llvm::SmallVector<Value> MlirReductionFusion::EmitReduction(
     if (epilogue.roots.empty()) return outputs;
 
     llvm::SmallVector<Value> epilogue_input_symbols(
-        epilogue.root_indexing.front().getNumSymbols(), zero);
+        epilogue.root_indexing.front().GetAffineMap().getNumSymbols(), zero);
     auto epilogue_input_indices = state.thread_and_block_ids;
     epilogue_input_indices.append(epilogue_input_symbols);
     auto values =
@@ -257,7 +259,7 @@ llvm::SmallVector<Value> MlirReductionFusion::EmitReduction(
         *ComputeThreadIdToOutputIndexing(first_root_index, ctx),
         state.thread_and_block_ids, {}, b);
     for (auto [index, root] : llvm::enumerate(epilogue.roots)) {
-      auto output_indices = mlir_converter::ApplyAffineMap(
+      auto output_indices = mlir_converter::ApplyIndexing(
           epilogue.root_indexing[index], state.thread_and_block_ids,
           epilogue_input_symbols, b);
       for (auto [result_index, result] : llvm::enumerate(values.at(root))) {
@@ -373,16 +375,16 @@ HloValueMap MlirReductionFusion::EmitterState::EmitPerThreadReducedElements(
 
   auto body_builder = [&](ValueRange iter_args, ValueRange dim_values,
                           ValueRange symbol_values) -> SmallVector<Value> {
-    auto tile_indices = mlir_converter::ApplyAffineMap(
-        tile_indexing.GetAffineMap(), dim_values, symbol_values, builder);
+    auto tile_indices = mlir_converter::ApplyIndexing(tile_indexing, dim_values,
+                                                      symbol_values, builder);
 
     llvm::SmallVector<Value> results(iter_args.size(), nullptr);
     auto get_input_indices = [&](auto* hero, bool is_reduction) {
       const auto& input_shape =
           is_reduction ? hero->operand(0)->shape() : hero->shape();
-      return mlir_converter::ApplyAffineMap(
-          GetBitcastMap(tiling.GetXlaShape(), input_shape, builder.getContext())
-              .GetAffineMap(),
+      return mlir_converter::ApplyIndexing(
+          GetBitcastMap(tiling.GetXlaShape(), input_shape,
+                        builder.getContext()),
           tile_indices, {}, builder);
     };
     for (auto* reduction : reductions) {
diff --git a/third_party/xla/xla/service/gpu/fusions/scatter_mlir.cc b/third_party/xla/xla/service/gpu/fusions/scatter_mlir.cc
index 18d357a4316f5e..2bf6c52a76d095 100644
--- a/third_party/xla/xla/service/gpu/fusions/scatter_mlir.cc
+++ b/third_party/xla/xla/service/gpu/fusions/scatter_mlir.cc
@@ -53,7 +53,8 @@ namespace xla {
 namespace gpu {
 namespace {
 
-namespace ma = mlir::arith;
+namespace ma = ::mlir::arith;
+namespace scf = ::mlir::scf;
 
 using llvm::SmallVector;
 using mlir::Location;
@@ -62,13 +63,10 @@ using mlir::Value;
 using mlir::ValueRange;
 using mlir::func::ReturnOp;
 using mlir::tensor::InsertOp;
-using mlir_converter::ApplyAffineMap;
 using mlir_converter::CallTargetProvider;
 using mlir_converter::PartitionedComputations;
 using mlir_converter::ProvideParameter;
 
-namespace scf = ::mlir::scf;
-
 }  // namespace
 
 bool MlirScatterFusion::IsSupported(const HloFusionAnalysis& analysis) {
@@ -212,9 +210,8 @@ absl::Status MlirScatterFusion::EmitEntryFunction(
       [&](ValueRange output_tensors, ValueRange dim_values,
           ValueRange symbol_values) -> SmallVector<Value> {
         // Extract input element.
-        auto update_tensor_indices =
-            ApplyAffineMap(thread_id_to_update_map.GetAffineMap(), dim_values,
-                           symbol_values, b);
+        auto update_tensor_indices = mlir_converter::ApplyIndexing(
+            thread_id_to_update_map, dim_values, symbol_values, b);
         auto update_elem = ProvideParameter(
             root_computation, scatter, kScatterUpdateIndex,
             update_tensor_indices, call_targets, entry_function, b)[0];
diff --git a/third_party/xla/xla/service/gpu/fusions/scatter_mlir_test.cc b/third_party/xla/xla/service/gpu/fusions/scatter_mlir_test.cc
index 489e588cec4248..174145e8ff9448 100644
--- a/third_party/xla/xla/service/gpu/fusions/scatter_mlir_test.cc
+++ b/third_party/xla/xla/service/gpu/fusions/scatter_mlir_test.cc
@@ -185,8 +185,8 @@ TEST_F(MlirScatterFusionTest, Scatter_UniqueIndices) {
     }
   )";
   TF_ASSERT_OK(EmitAndCheckIR(kHloString, R"(
-    // CHECK: #[[$MAP0:.*]] = affine_map<()[s0] -> (s0 floordiv 2)>
-    // CHECK: #[[$MAP1:.*]] = affine_map<()[s0] -> (s0 mod 2)>
+    // CHECK: #[[$MAP0:.*]] = affine_map<(d0) -> (d0 floordiv 2)>
+    // CHECK: #[[$MAP1:.*]] = affine_map<(d0) -> (d0 mod 2)>
 
     // CHECK-LABEL: func.func @fused_computation(
     // CHECK-SAME:    %[[OPERAND:[a-zA-Z0-9]*]]: tensor<10x5xf32>
@@ -198,8 +198,8 @@ TEST_F(MlirScatterFusionTest, Scatter_UniqueIndices) {
     // CHECK-DAG:       %[[C9:.*]] = arith.constant 9 : index
 
     // CHECK:      %[[TH_X:.*]] = gpu.thread_id  x
-    // CHECK:      %[[SLICE_ID:.*]] = affine.apply #[[$MAP0]]()[%[[TH_X]]]
-    // CHECK:      %[[SLICE_X:.*]] = affine.apply #[[$MAP1]]()[%[[TH_X]]]
+    // CHECK:      %[[SLICE_ID:.*]] = xla_gpu.apply_indexing #[[$MAP0]](%[[TH_X]]
+    // CHECK:      %[[SLICE_X:.*]] = xla_gpu.apply_indexing #[[$MAP1]](%[[TH_X]]
 
     // CHECK:      %[[UPD_ELEM:.*]] = xla_gpu.pure_call @scatter_update(
     // CHECK-SAME:  %[[OPERAND]], %[[INDICES]], %[[UPDATES]],
diff --git a/third_party/xla/xla/service/gpu/fusions/transpose_mlir.cc b/third_party/xla/xla/service/gpu/fusions/transpose_mlir.cc
index df65dbb55ea5e1..9d76a1d09237ae 100644
--- a/third_party/xla/xla/service/gpu/fusions/transpose_mlir.cc
+++ b/third_party/xla/xla/service/gpu/fusions/transpose_mlir.cc
@@ -72,7 +72,6 @@ using mlir::func::FuncOp;
 using mlir::func::ReturnOp;
 using mlir::tensor::ExtractOp;
 using mlir::tensor::InsertOp;
-using mlir_converter::ApplyAffineMap;
 using mlir_converter::ApplyIndexing;
 
 constexpr int kNumRows = 4;
@@ -286,7 +285,7 @@ void MlirTransposeFusion::EmitReadFromShMemMlir(
                        computations.epilogues().front().root_indexing,
                        shmem_transpose_root_indices_)) {
           llvm::SmallVector<Value> indices =
-              ApplyAffineMap(indexing, dim_values, symbol_values, builder);
+              ApplyIndexing(indexing, dim_values, symbol_values, builder);
           results[root_index] = builder.create<InsertOp>(
               result_scalars.at(root).front(), results[root_index], indices);
         }
diff --git a/third_party/xla/xla/service/gpu/model/indexing_map.h b/third_party/xla/xla/service/gpu/model/indexing_map.h
index 69021749385506..ee0a18ae1ef9cf 100644
--- a/third_party/xla/xla/service/gpu/model/indexing_map.h
+++ b/third_party/xla/xla/service/gpu/model/indexing_map.h
@@ -351,6 +351,11 @@ class IndexingMap {
   // Does `symbol` correspond to an RTVar?
   bool IsRTVarSymbol(mlir::AffineSymbolExpr symbol) const;
 
+  IndexingMap GetSubMap(unsigned int result_index) const {
+    return {affine_map_.getSubMap({result_index}), dim_vars_, range_vars_,
+            rt_vars_, constraints_};
+  }
+
  private:
   IndexingMap() = default;
 

From 35bc8b9bbc6babe1210fc7cb6304b82afc8651dd Mon Sep 17 00:00:00 2001
From: Kyle Lucke <klucke@google.com>
Date: Tue, 21 May 2024 06:14:19 -0700
Subject: [PATCH 461/478] Remove the use of xla::OkStatus now that it's just an
 alias to absl::OkStatus.

PiperOrigin-RevId: 635783707
---
 .../xla/xla/pjrt/c/pjrt_c_api_wrapper_impl.cc       | 13 ++++++-------
 third_party/xla/xla/pjrt/distributed/BUILD          |  1 +
 .../xla/xla/pjrt/distributed/client_server_test.cc  | 11 ++++++-----
 third_party/xla/xla/python/BUILD                    |  1 +
 third_party/xla/xla/python/inspect_sharding.cc      |  3 ++-
 third_party/xla/xla/translate/mhlo_to_hlo/BUILD     |  1 +
 .../xla/xla/translate/mhlo_to_hlo/layout_util.cc    |  4 +++-
 7 files changed, 20 insertions(+), 14 deletions(-)

diff --git a/third_party/xla/xla/pjrt/c/pjrt_c_api_wrapper_impl.cc b/third_party/xla/xla/pjrt/c/pjrt_c_api_wrapper_impl.cc
index ab1a5a1a3b754c..68fd8f7b8edaaf 100644
--- a/third_party/xla/xla/pjrt/c/pjrt_c_api_wrapper_impl.cc
+++ b/third_party/xla/xla/pjrt/c/pjrt_c_api_wrapper_impl.cc
@@ -62,7 +62,6 @@ limitations under the License.
 #include "xla/service/hlo.pb.h"
 #include "xla/shape.h"
 #include "xla/shape_util.h"
-#include "xla/status.h"
 #include "xla/util.h"
 #include "xla/xla.pb.h"
 #include "xla/xla_data.pb.h"
@@ -109,7 +108,7 @@ static absl::Status PopulateExecutableCostAnalysis(
                       executable->get()->GetCostAnalysis());
   // If no output, return empty result
   if (properties.empty()) {
-    return xla::OkStatus();
+    return absl::OkStatus();
   }
 
   // Copy each returned property to cost analysis vectors in PJRT_Executable
@@ -141,7 +140,7 @@ static absl::Status PopulateExecutableCostAnalysis(
     ++i;
   }
 
-  return xla::OkStatus();
+  return absl::OkStatus();
 }
 
 static absl::Status PopulateExecutableOutputElementTypes(
@@ -167,7 +166,7 @@ static absl::Status PopulateExecutableOutputElementTypes(
     out_types.push_back(ConvertToPjRtBufferType(element_type));
   }
 
-  return xla::OkStatus();
+  return absl::OkStatus();
 }
 
 static absl::Status PopulateExecutableOutputDimensions(
@@ -201,7 +200,7 @@ static absl::Status PopulateExecutableOutputDimensions(
     }
   }
 
-  return xla::OkStatus();
+  return absl::OkStatus();
 }
 
 static absl::Status PopulateExecutableOutputMemoryKinds(
@@ -230,7 +229,7 @@ static absl::Status PopulateExecutableOutputMemoryKinds(
     memory_kind_sizes.push_back(memory.size());
   }
 
-  return xla::OkStatus();
+  return absl::OkStatus();
 }
 
 class CApiKeyValueStore : public xla::KeyValueStoreInterface {
@@ -1101,7 +1100,7 @@ static absl::Status VerifyOptimizedProgramArgs(
       PJRT_Executable_OptimizedProgram_Args_STRUCT_SIZE, args->struct_size));
   TF_RETURN_IF_ERROR(ActualStructSizeIsGreaterOrEqual(
       "PJRT_Program", PJRT_Program_STRUCT_SIZE, args->program->struct_size));
-  return xla::OkStatus();
+  return absl::OkStatus();
 }
 
 static absl::StatusOr<std::shared_ptr<xla::HloModule>>
diff --git a/third_party/xla/xla/pjrt/distributed/BUILD b/third_party/xla/xla/pjrt/distributed/BUILD
index c0ae37c0b3e3b7..20533c2cbb1073 100644
--- a/third_party/xla/xla/pjrt/distributed/BUILD
+++ b/third_party/xla/xla/pjrt/distributed/BUILD
@@ -152,6 +152,7 @@ xla_cc_test(
         "//xla:protobuf_util",
         "//xla:status_macros",
         "//xla/tsl/distributed_runtime/coordination:coordination_service_agent",
+        "@com_google_absl//absl/status",
         "@com_google_absl//absl/strings",
         "@com_google_absl//absl/synchronization",
         "@com_google_absl//absl/time",
diff --git a/third_party/xla/xla/pjrt/distributed/client_server_test.cc b/third_party/xla/xla/pjrt/distributed/client_server_test.cc
index ef0f29181400f2..9fe0741a83ae85 100644
--- a/third_party/xla/xla/pjrt/distributed/client_server_test.cc
+++ b/third_party/xla/xla/pjrt/distributed/client_server_test.cc
@@ -20,6 +20,7 @@ limitations under the License.
 #include <string_view>
 #include <vector>
 
+#include "absl/status/status.h"
 #include "absl/strings/str_cat.h"
 #include "absl/strings/string_view.h"
 #include "absl/synchronization/barrier.h"
@@ -583,7 +584,7 @@ TEST_F(ClientServerTest, WaitAtBarrier_Succeed) {
         client->WaitAtBarrier("barrier_2", kBarrierTimeout, std::nullopt));
 
     TF_RETURN_IF_ERROR(client->Shutdown());
-    return xla::OkStatus();
+    return absl::OkStatus();
   };
 
   std::vector<absl::Status> statuses(num_nodes);
@@ -621,7 +622,7 @@ TEST_F(ClientServerTest, WaitAtBarrier_Timeout) {
     TF_RETURN_IF_ERROR(barrier_status);
 
     TF_RETURN_IF_ERROR(client->Shutdown());
-    return xla::OkStatus();
+    return absl::OkStatus();
   };
 
   std::vector<absl::Status> statuses(num_nodes);
@@ -658,7 +659,7 @@ TEST_F(ClientServerTest, WaitAtBarrier_TimeoutWithDifferentBarrierId) {
         client->WaitAtBarrier(barrier_id, kBarrierTimeout, std::nullopt));
 
     TF_RETURN_IF_ERROR(client->Shutdown());
-    return xla::OkStatus();
+    return absl::OkStatus();
   };
 
   std::vector<absl::Status> statuses(num_nodes);
@@ -689,7 +690,7 @@ TEST_F(ClientServerTest, WaitAtBarrier_FailWithSameBarrierId) {
         client->WaitAtBarrier("barrier_1", kBarrierTimeout, std::nullopt));
 
     TF_RETURN_IF_ERROR(client->Shutdown());
-    return xla::OkStatus();
+    return absl::OkStatus();
   };
 
   std::vector<absl::Status> statuses(num_nodes);
@@ -721,7 +722,7 @@ TEST_F(ClientServerTest, WaitAtBarrierSubset_Succeeds) {
     }
 
     TF_RETURN_IF_ERROR(client->Shutdown());
-    return xla::OkStatus();
+    return absl::OkStatus();
   };
 
   std::vector<absl::Status> statuses(num_nodes);
diff --git a/third_party/xla/xla/python/BUILD b/third_party/xla/xla/python/BUILD
index ada2f76f96b0ff..cc3c219dbdfebb 100644
--- a/third_party/xla/xla/python/BUILD
+++ b/third_party/xla/xla/python/BUILD
@@ -567,6 +567,7 @@ cc_library(
         "//xla/hlo/ir:hlo",
         "//xla/service:custom_call_sharding_helper",
         "//xla/service/spmd:spmd_partitioner",
+        "@com_google_absl//absl/status",
     ],
     # Always register 'InspectSharding' custom partitioning handler.
     alwayslink = 1,
diff --git a/third_party/xla/xla/python/inspect_sharding.cc b/third_party/xla/xla/python/inspect_sharding.cc
index 2ab386983b0a51..dfa03f37f01e01 100644
--- a/third_party/xla/xla/python/inspect_sharding.cc
+++ b/third_party/xla/xla/python/inspect_sharding.cc
@@ -21,6 +21,7 @@ limitations under the License.
 #include <string>
 #include <utility>
 
+#include "absl/status/status.h"
 #include "xla/service/custom_call_sharding_helper.h"
 #include "xla/service/spmd/spmd_partitioner_util.h"
 
@@ -85,7 +86,7 @@ class InspectShardingCallPartitioner : public xla::CustomCallPartitioner {
     partitioner->SetPartitionedHlo(
         instruction,
         partitioner->GetPartitionedHlo(instruction->mutable_operand(0)));
-    return xla::OkStatus();
+    return absl::OkStatus();
   }
   HloSharding PropagateUserSharding(
       const HloInstruction* instruction, const HloInstruction* user,
diff --git a/third_party/xla/xla/translate/mhlo_to_hlo/BUILD b/third_party/xla/xla/translate/mhlo_to_hlo/BUILD
index 77654530a5c80b..1386aa586a555c 100644
--- a/third_party/xla/xla/translate/mhlo_to_hlo/BUILD
+++ b/third_party/xla/xla/translate/mhlo_to_hlo/BUILD
@@ -43,6 +43,7 @@ cc_library(
         "//xla:xla_data_proto_cc",
         "//xla/client:xla_builder",
         "//xla/hlo/ir:hlo",
+        "@com_google_absl//absl/status",
     ],
 )
 
diff --git a/third_party/xla/xla/translate/mhlo_to_hlo/layout_util.cc b/third_party/xla/xla/translate/mhlo_to_hlo/layout_util.cc
index 88261a847874f5..cf219a5a13bdfc 100644
--- a/third_party/xla/xla/translate/mhlo_to_hlo/layout_util.cc
+++ b/third_party/xla/xla/translate/mhlo_to_hlo/layout_util.cc
@@ -15,6 +15,8 @@ limitations under the License.
 
 #include "xla/translate/mhlo_to_hlo/layout_util.h"
 
+#include "absl/status/status.h"
+
 namespace mlir {
 
 // Rewrites the layout of xla_shape if there is tiled sharding.
@@ -58,7 +60,7 @@ absl::Status RewriteLayoutWithShardedShape(
             : per_device_xla_shape);
     *xla_shape->mutable_layout() = per_device_xla_shape.layout();
   }
-  return xla::OkStatus();
+  return absl::OkStatus();
 }
 
 // There is a shape_representation_fn or sharding for an output, this function

From 6a90cce5e743dd4a936c3b2600783fdeb9eaf56a Mon Sep 17 00:00:00 2001
From: Quentin Khan <qkhan@google.com>
Date: Tue, 21 May 2024 06:47:26 -0700
Subject: [PATCH 462/478] Reverts 42ac0dde8f7855c87dfea2217b03b8ee7e829ae9

PiperOrigin-RevId: 635790293
---
 .../lite/tools/benchmark/benchmark_tflite_model.cc | 14 ++++++++++++--
 1 file changed, 12 insertions(+), 2 deletions(-)

diff --git a/tensorflow/lite/tools/benchmark/benchmark_tflite_model.cc b/tensorflow/lite/tools/benchmark/benchmark_tflite_model.cc
index 7728723940a481..d775122fe9c1fc 100644
--- a/tensorflow/lite/tools/benchmark/benchmark_tflite_model.cc
+++ b/tensorflow/lite/tools/benchmark/benchmark_tflite_model.cc
@@ -1100,8 +1100,18 @@ TfLiteStatus BenchmarkTfLiteModel::LoadModel() {
 
 std::unique_ptr<tflite::OpResolver> BenchmarkTfLiteModel::GetOpResolver()
     const {
-  tflite::ops::builtin::BuiltinOpResolver* resolver =
-      new tflite::ops::builtin::BuiltinOpResolverWithoutDefaultDelegates();
+  tflite::ops::builtin::BuiltinOpResolver* resolver = nullptr;
+  // When --use_xnnpack is explicitly set to false, skip applying the default
+  // XNNPACK delegate in TfLite runtime so that the original execution path
+  // based on the unmodified model graph is still exercised.
+  if (params_.HasParam("use_xnnpack") &&
+      params_.HasValueSet<bool>("use_xnnpack") &&
+      !params_.Get<bool>("use_xnnpack")) {
+    resolver =
+        new tflite::ops::builtin::BuiltinOpResolverWithoutDefaultDelegates();
+  } else {
+    resolver = new tflite::ops::builtin::BuiltinOpResolver();
+  }
   RegisterSelectedOps(resolver);
   return std::unique_ptr<tflite::OpResolver>(resolver);
 }

From a68d77d0434df8c7d2cff790f2bb1c08c4c88742 Mon Sep 17 00:00:00 2001
From: Ilia Sergachev <isergachev@nvidia.com>
Date: Tue, 21 May 2024 06:57:26 -0700
Subject: [PATCH 463/478] PR #12897: [GPU] Fix cuDNN GEMM scalar constants test
 condition.

Imported from GitHub PR https://github.com/openxla/xla/pull/12897

Copybara import of the project:

--
71b57fd71d603bc8c346283796f6859761eef014 by Ilia Sergachev <isergachev@nvidia.com>:

[GPU] Fix cuDNN GEMM scalar constants test condition.

Merging this change closes #12897

PiperOrigin-RevId: 635792570
---
 third_party/xla/xla/service/gpu/fusions/cudnn_test.cc | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/third_party/xla/xla/service/gpu/fusions/cudnn_test.cc b/third_party/xla/xla/service/gpu/fusions/cudnn_test.cc
index afae7eccee16c2..ba664e1029854e 100644
--- a/third_party/xla/xla/service/gpu/fusions/cudnn_test.cc
+++ b/third_party/xla/xla/service/gpu/fusions/cudnn_test.cc
@@ -522,6 +522,9 @@ ENTRY e {
 }
 
 TEST_F(CuDnnFusionLevel2Test, ConstantExecutesCorrectly) {
+  if (!IsAtLeastCuDnn91()) {
+    GTEST_SKIP() << "Fused scalar constants require cuDNN 9.1+.";
+  }
   EXPECT_TRUE(RunAndCompare(R"(
 fusion1 {
   x = bf16[16,32] parameter(0)

From 7990149b38ecb9bf87cd37ed15f929a2ccd94514 Mon Sep 17 00:00:00 2001
From: Johannes Reifferscheid <jreiffers@google.com>
Date: Tue, 21 May 2024 06:58:01 -0700
Subject: [PATCH 464/478] Rename `RewriteSumIf` to `RemoveSummands`.

It was pointed out to me that this code is terrible. Let this
be the first part of my formal apology.

PiperOrigin-RevId: 635792761
---
 .../xla/xla/service/gpu/model/indexing_map.cc | 61 ++++++++-----------
 1 file changed, 25 insertions(+), 36 deletions(-)

diff --git a/third_party/xla/xla/service/gpu/model/indexing_map.cc b/third_party/xla/xla/service/gpu/model/indexing_map.cc
index 80390bb6c119e8..e96b4491add5b1 100644
--- a/third_party/xla/xla/service/gpu/model/indexing_map.cc
+++ b/third_party/xla/xla/service/gpu/model/indexing_map.cc
@@ -87,11 +87,10 @@ class AffineExprSimplifier {
   // - Rewrites a / 100 to 0 when a is known to be less than 100.
   mlir::AffineExpr RewriteFloorDiv(mlir::AffineBinaryOpExpr div);
 
-  mlir::AffineExpr RewriteSum(
-      mlir::AffineExpr expr,
-      const std::function<mlir::AffineExpr(mlir::AffineExpr)>& map);
-
-  mlir::AffineExpr RewriteSumIf(
+  // Removes summands from arbitrarily nested sums (e.g, ((a+b)+c)) if `pred`
+  // returns true. In this example, `pred` is evaluated on `a`, `b` and `c`, not
+  // on `a+b`.
+  mlir::AffineExpr RemoveSummands(
       mlir::AffineExpr expr, const std::function<bool(mlir::AffineExpr)>& pred);
 
   // Attempts to simplify the expression, but doesn't attempt to simplify the
@@ -152,15 +151,15 @@ AffineExpr AffineExprSimplifier::RewriteMod(AffineBinaryOpExpr mod) {
   int64_t multiplier_gcd = -1;
 
   int64_t extracted_constant = 0;
-  auto new_lhs = RewriteSumIf(lhs_simplified, [&](AffineExpr expr) {
+  auto new_lhs = RemoveSummands(lhs_simplified, [&](AffineExpr expr) {
     if (auto cst = mlir::dyn_cast<AffineConstantExpr>(expr);
         cst && cst.getValue() >= m) {
       extracted_constant += cst.getValue();
-      return false;
+      return true;
     }
     if (auto multiplier = GetConstantRhs(expr, AffineExprKind::Mul)) {
       if (*multiplier % m == 0) {
-        return false;
+        return true;
       }
 
       if (multiplier_gcd == -1) {
@@ -168,12 +167,12 @@ AffineExpr AffineExprSimplifier::RewriteMod(AffineBinaryOpExpr mod) {
       } else {
         multiplier_gcd = std::gcd(multiplier_gcd, *multiplier);
       }
-      return true;
+      return false;
     }
     auto range = range_evaluator_->ComputeExpressionRange(expr);
     no_multiplier_range.lower += range.lower;
     no_multiplier_range.upper += range.upper;
-    return true;
+    return false;
   });
   new_lhs = new_lhs + (extracted_constant % m);
 
@@ -181,12 +180,12 @@ AffineExpr AffineExprSimplifier::RewriteMod(AffineBinaryOpExpr mod) {
   if (m % multiplier_gcd == 0 && no_multiplier_range.lower >= 0 &&
       no_multiplier_range.upper < multiplier_gcd) {
     // Remove everything that doesn't have a multiplier.
-    new_lhs = RewriteSumIf(new_lhs, [&](AffineExpr expr) {
+    new_lhs = RemoveSummands(new_lhs, [&](AffineExpr expr) {
       if (GetConstantRhs(expr, AffineExprKind::Mul)) {
-        return true;
+        return false;
       }
       extracted = extracted + expr;
-      return false;
+      return true;
     });
   }
   return new_lhs % mod.getRHS() + extracted;
@@ -246,17 +245,16 @@ AffineExpr AffineExprSimplifier::RewriteFloorDiv(AffineBinaryOpExpr div) {
   int64_t max_remaining_multiplier_gcd = -1;
   AffineExpr zero = getAffineConstantExpr(0, mlir_context);
   AffineExpr extracted = zero;
-  auto new_dividend = RewriteSumIf(lhs_simplified, [&](AffineExpr expr) {
+  auto new_dividend = RemoveSummands(lhs_simplified, [&](AffineExpr expr) {
     if (auto multiplier = GetConstantRhs(expr, AffineExprKind::Mul)) {
       // (x * 7 + ...) / 3 -> can't extract. We could extract x * 2 and keep
       // one x, but we currently have no reason to do that.
-
       if (*multiplier % d == 0) {
         int64_t factor = *multiplier / d;
         extracted =
             extracted + mlir::cast<AffineBinaryOpExpr>(expr).getLHS() * factor;
         // Remove from dividend.
-        return false;
+        return true;
       }
 
       if (*multiplier > 0) {
@@ -267,14 +265,14 @@ AffineExpr AffineExprSimplifier::RewriteFloorDiv(AffineBinaryOpExpr div) {
         }
         max_remaining_multiplier_gcd =
             std::max(max_remaining_multiplier_gcd, std::gcd(*multiplier, d));
-        return true;
+        return false;
       }
     }
     auto range = range_evaluator_->ComputeExpressionRange(expr);
     no_multiplier_range.lower += range.lower;
     no_multiplier_range.upper += range.upper;
     // Not a constant multiplier, keep in dividend.
-    return true;
+    return false;
   });
 
   // If we removed everything, skip the div.
@@ -286,9 +284,9 @@ AffineExpr AffineExprSimplifier::RewriteFloorDiv(AffineBinaryOpExpr div) {
     if (no_multiplier_range.lower >= 0 &&
         no_multiplier_range.upper < multiplier_gcd) {
       // Remove everything that doesn't have a multiplier.
-      new_dividend = RewriteSumIf(new_dividend, [&](AffineExpr expr) {
+      new_dividend = RemoveSummands(new_dividend, [&](AffineExpr expr) {
         auto mult = GetConstantRhs(expr, AffineExprKind::Mul);
-        return mult.has_value();
+        return !mult.has_value();
       });
     }
   }
@@ -297,7 +295,7 @@ AffineExpr AffineExprSimplifier::RewriteFloorDiv(AffineBinaryOpExpr div) {
   // (x * 128 + y) // 192 -> (x * 2 + y // 64) // 3
   if (max_remaining_multiplier_gcd > 1) {
     AffineExpr partially_extracted = getAffineConstantExpr(0, mlir_context);
-    new_dividend = RewriteSumIf(new_dividend, [&](AffineExpr expr) {
+    new_dividend = RemoveSummands(new_dividend, [&](AffineExpr expr) {
       if (auto multiplier = GetConstantRhs(expr, AffineExprKind::Mul);
           multiplier && (*multiplier > 0) &&
           ((*multiplier % max_remaining_multiplier_gcd) == 0)) {
@@ -306,9 +304,9 @@ AffineExpr AffineExprSimplifier::RewriteFloorDiv(AffineBinaryOpExpr div) {
             partially_extracted +
             expr_lhs * (*multiplier / max_remaining_multiplier_gcd);
         // Remove from dividend.
-        return false;
+        return true;
       }
-      return true;
+      return false;
     });
     return extracted + (partially_extracted +
                         new_dividend.floorDiv(max_remaining_multiplier_gcd))
@@ -337,27 +335,18 @@ std::optional<int64_t> AffineExprSimplifier::GetConstantRhs(
   return bound.lower;
 }
 
-AffineExpr AffineExprSimplifier::RewriteSum(
-    AffineExpr expr, const std::function<AffineExpr(AffineExpr)>& map) {
-  if (expr.getKind() == AffineExprKind::Add) {
-    auto add = mlir::dyn_cast<AffineBinaryOpExpr>(expr);
-    return RewriteSum(add.getLHS(), map) + RewriteSum(add.getRHS(), map);
-  }
-  return map(expr);
-}
-
-AffineExpr AffineExprSimplifier::RewriteSumIf(
+AffineExpr AffineExprSimplifier::RemoveSummands(
     AffineExpr expr, const std::function<bool(AffineExpr)>& pred) {
   if (expr.getKind() == AffineExprKind::Add) {
     auto add = mlir::dyn_cast<AffineBinaryOpExpr>(expr);
-    auto lhs = RewriteSumIf(add.getLHS(), pred);
-    auto rhs = RewriteSumIf(add.getRHS(), pred);
+    auto lhs = RemoveSummands(add.getLHS(), pred);
+    auto rhs = RemoveSummands(add.getRHS(), pred);
     if (lhs == add.getLHS() && rhs == add.getRHS()) {
       return add;
     }
     return lhs + rhs;
   }
-  return pred(expr) ? expr : mlir::getAffineConstantExpr(0, expr.getContext());
+  return pred(expr) ? mlir::getAffineConstantExpr(0, expr.getContext()) : expr;
 }
 
 // Compares the two expression by their AST. The ordering is arbitrary but

From da437e36ad449b1c4b1ed22b7a879f496e5016a6 Mon Sep 17 00:00:00 2001
From: Oleg Shyshkov <shyshkov@google.com>
Date: Tue, 21 May 2024 07:09:47 -0700
Subject: [PATCH 465/478] [Triton] Fix a use-after-free bug in
 LinearLayout::compose.

`output.apply(bases)` returns a temporary object and `llvm::make_second_range` doesn't copy or take ownership. As a result, `newBases` points to deallocated memory.

Also rename `newBases` to `outerBases`, because there is `newBases` variable in the outer scope.

PiperOrigin-RevId: 635795849
---
 .../temporary/linear_layout_compose_asan.patch | 18 ++++++++++++++++++
 third_party/triton/temporary/series.bzl        |  4 +++-
 .../temporary/linear_layout_compose_asan.patch | 18 ++++++++++++++++++
 .../third_party/triton/temporary/series.bzl    |  4 +++-
 4 files changed, 42 insertions(+), 2 deletions(-)
 create mode 100644 third_party/triton/temporary/linear_layout_compose_asan.patch
 create mode 100644 third_party/xla/third_party/triton/temporary/linear_layout_compose_asan.patch

diff --git a/third_party/triton/temporary/linear_layout_compose_asan.patch b/third_party/triton/temporary/linear_layout_compose_asan.patch
new file mode 100644
index 00000000000000..eff83a166ac4a3
--- /dev/null
+++ b/third_party/triton/temporary/linear_layout_compose_asan.patch
@@ -0,0 +1,18 @@
+==== triton/lib/Tools/LinearLayout.cpp#2 - /google/src/cloud/shyshkov/triton_asan/triton/lib/Tools/LinearLayout.cpp ====
+# action=edit type=text
+--- triton/lib/Tools/LinearLayout.cpp	2024-05-17 09:15:25.000000000 -0700
++++ triton/lib/Tools/LinearLayout.cpp	2024-05-21 06:27:58.000000000 -0700
+@@ -397,9 +397,11 @@
+       for (auto [outDim, b] : llvm::zip(getOutDimNames(), basis)) {
+         bases.push_back({outDim, b});
+       }
+-      auto newBases = llvm::make_second_range(outer.apply(bases));
++
++      auto outerBases =
++          llvm::to_vector(llvm::make_second_range(outer.apply(bases)));
+       newInDimBases.push_back(
+-          std::vector<int32_t>(newBases.begin(), newBases.end()));
++          std::vector<int32_t>(outerBases.begin(), outerBases.end()));
+     }
+   }
+   return LinearLayout(std::move(newBases), outer.getOutDimNames());
diff --git a/third_party/triton/temporary/series.bzl b/third_party/triton/temporary/series.bzl
index 53a5059fe432d7..b3d935c048fadd 100644
--- a/third_party/triton/temporary/series.bzl
+++ b/third_party/triton/temporary/series.bzl
@@ -5,4 +5,6 @@ These are created temporarily and should be moved to the first copybara workflow
 internal patch during the next triton integration process.
 """
 
-temporary_patch_list = []
+temporary_patch_list = [
+    "//third_party/triton/temporary:linear_layout_compose_asan.patch",
+]
diff --git a/third_party/xla/third_party/triton/temporary/linear_layout_compose_asan.patch b/third_party/xla/third_party/triton/temporary/linear_layout_compose_asan.patch
new file mode 100644
index 00000000000000..eff83a166ac4a3
--- /dev/null
+++ b/third_party/xla/third_party/triton/temporary/linear_layout_compose_asan.patch
@@ -0,0 +1,18 @@
+==== triton/lib/Tools/LinearLayout.cpp#2 - /google/src/cloud/shyshkov/triton_asan/triton/lib/Tools/LinearLayout.cpp ====
+# action=edit type=text
+--- triton/lib/Tools/LinearLayout.cpp	2024-05-17 09:15:25.000000000 -0700
++++ triton/lib/Tools/LinearLayout.cpp	2024-05-21 06:27:58.000000000 -0700
+@@ -397,9 +397,11 @@
+       for (auto [outDim, b] : llvm::zip(getOutDimNames(), basis)) {
+         bases.push_back({outDim, b});
+       }
+-      auto newBases = llvm::make_second_range(outer.apply(bases));
++
++      auto outerBases =
++          llvm::to_vector(llvm::make_second_range(outer.apply(bases)));
+       newInDimBases.push_back(
+-          std::vector<int32_t>(newBases.begin(), newBases.end()));
++          std::vector<int32_t>(outerBases.begin(), outerBases.end()));
+     }
+   }
+   return LinearLayout(std::move(newBases), outer.getOutDimNames());
diff --git a/third_party/xla/third_party/triton/temporary/series.bzl b/third_party/xla/third_party/triton/temporary/series.bzl
index 53a5059fe432d7..b3d935c048fadd 100644
--- a/third_party/xla/third_party/triton/temporary/series.bzl
+++ b/third_party/xla/third_party/triton/temporary/series.bzl
@@ -5,4 +5,6 @@ These are created temporarily and should be moved to the first copybara workflow
 internal patch during the next triton integration process.
 """
 
-temporary_patch_list = []
+temporary_patch_list = [
+    "//third_party/triton/temporary:linear_layout_compose_asan.patch",
+]

From 4001ea62313a4ed03f1e6fbc1cdaa7b138f2ddb0 Mon Sep 17 00:00:00 2001
From: Johannes Reifferscheid <jreiffers@google.com>
Date: Tue, 21 May 2024 07:14:34 -0700
Subject: [PATCH 466/478] Tighten reduction loop bounds for MLIR emitter.

See the added comment for an example of a problematic bound
that is currently created.

PiperOrigin-RevId: 635796929
---
 .../xla/service/gpu/fusions/reduction_base.cc | 15 ++++++
 .../gpu/fusions/reduction_base_test.cc        | 46 +++++++++++++++++++
 2 files changed, 61 insertions(+)

diff --git a/third_party/xla/xla/service/gpu/fusions/reduction_base.cc b/third_party/xla/xla/service/gpu/fusions/reduction_base.cc
index 0d63d73b3339af..145b4107e30fba 100644
--- a/third_party/xla/xla/service/gpu/fusions/reduction_base.cc
+++ b/third_party/xla/xla/service/gpu/fusions/reduction_base.cc
@@ -309,6 +309,21 @@ ReductionInfo ReductionInfo::Create(const HloFusionAnalysis& analysis,
   absl::InlinedVector<int64_t, 4> tile_per_thread{
       reduction_tiling[0], reduction_tiling[1],
       reduction_tiling[2] / vector_size};
+  if (for_mlir) {
+    // The indexing map simplifier does not currently handle this correctly,
+    // leading to loop bounds that are too large.
+    // TODO(jreiffers): Implement tightening of ranges based on constraints
+    // instead. For example, based on:
+    //
+    //   s1 in [0, 127]
+    //   d0 floordiv 32 + s1 * 32 in [0, 63]
+    //
+    // Tighten the bound of s1 to [0, 1].
+    for (int i = 0; i < num_threads.size(); ++i) {
+      tile_per_thread[i] = std::min(
+          tile_per_thread[i], CeilOfRatio(tiled_shape[i], num_threads[i]));
+    }
+  }
   if (rows_per_warp > 1) {
     // If we produce more than one element per thread, that means the reduced
     // dimension is small and it can't be tiled - we already have more threads
diff --git a/third_party/xla/xla/service/gpu/fusions/reduction_base_test.cc b/third_party/xla/xla/service/gpu/fusions/reduction_base_test.cc
index 674765142cf4fe..7cd9eec245b823 100644
--- a/third_party/xla/xla/service/gpu/fusions/reduction_base_test.cc
+++ b/third_party/xla/xla/service/gpu/fusions/reduction_base_test.cc
@@ -530,6 +530,52 @@ TEST_F(ReductionTest, OneGroup) {
               SizeIs(1));
 }
 
+TEST_F(ReductionTest, MlirColumnReduction) {
+  auto module = ParseAndReturnVerifiedModule(R"(
+    add {
+      b = f32[] parameter(1)
+      a = f32[] parameter(0)
+      ROOT out = f32[] add(a, b)
+    }
+    fusion {
+      %p0 = f32[192,64,1536] parameter(0)
+      %c0 = f32[] constant(0)
+      ROOT reduce = f32[192,1536] reduce(p0, c0), dimensions={1}, to_apply=add
+    }
+    ENTRY entry {
+      %p0 = f32[192,64,1536] parameter(0)
+      ROOT %fusion = f32[192,1536] fusion(%p0), kind=kInput, calls=fusion
+    })")
+                    .value();
+
+  auto* root = module->entry_computation()->root_instruction();
+  auto analysis = AnalyzeFusion(*root, device_info_);
+
+  FakeMlirReductionFusion fusion(analysis);
+
+  EXPECT_THAT(
+      fusion.ComputeThreadIdToInputIndexing(0, 0, &mlir_context_)->ToString(),
+      MatchIndexingString(R"(
+        (d0, d1, d2, d3, d4, d5)[s0, s1, s2] -> (
+          d3 floordiv 48,
+          d0 floordiv 32 + s1 * 32,
+          (d3 mod 48) * 32 + d0 mod 32
+        )
+        domain:
+        d0 in [0, 1023]
+        d1 in [0, 0]
+        d2 in [0, 0]
+        d3 in [0, 9215]
+        d4 in [0, 0]
+        d5 in [0, 0]
+        s0 in [0, 0]
+        s1 in [0, 1]
+        s2 in [0, 0]
+        (d3 mod 48) * 32 + d0 mod 32 in [0, 1535]
+        d0 floordiv 32 + s1 * 32 in [0, 63]
+      )"));
+}
+
 }  // namespace
 }  // namespace gpu
 }  // namespace xla

From 0f7658d6ec8a0bb8521e2c15a30c22199993f802 Mon Sep 17 00:00:00 2001
From: Johannes Reifferscheid <jreiffers@google.com>
Date: Tue, 21 May 2024 08:28:19 -0700
Subject: [PATCH 467/478] Split sum-modification from sum-visitation.

Currently, we're both computing values and modifying sums in one step, which is
not very nice. This splits the modification parts from the traversal parts
where feasible. There's some almost-duplication in the mod/div simplifiers
which could probably be cleaned up, but so far I failed to find a good name for
it.

PiperOrigin-RevId: 635815457
---
 .../xla/xla/service/gpu/model/indexing_map.cc | 70 ++++++++++++-------
 1 file changed, 44 insertions(+), 26 deletions(-)

diff --git a/third_party/xla/xla/service/gpu/model/indexing_map.cc b/third_party/xla/xla/service/gpu/model/indexing_map.cc
index e96b4491add5b1..168d77f849e9d9 100644
--- a/third_party/xla/xla/service/gpu/model/indexing_map.cc
+++ b/third_party/xla/xla/service/gpu/model/indexing_map.cc
@@ -92,6 +92,8 @@ class AffineExprSimplifier {
   // on `a+b`.
   mlir::AffineExpr RemoveSummands(
       mlir::AffineExpr expr, const std::function<bool(mlir::AffineExpr)>& pred);
+  void VisitSummands(mlir::AffineExpr expr,
+                     const std::function<void(mlir::AffineExpr)>& visit);
 
   // Attempts to simplify the expression, but doesn't attempt to simplify the
   // result further.
@@ -147,9 +149,6 @@ AffineExpr AffineExprSimplifier::RewriteMod(AffineBinaryOpExpr mod) {
            *mul;
   }
 
-  Interval no_multiplier_range{0, 0};
-  int64_t multiplier_gcd = -1;
-
   int64_t extracted_constant = 0;
   auto new_lhs = RemoveSummands(lhs_simplified, [&](AffineExpr expr) {
     if (auto cst = mlir::dyn_cast<AffineConstantExpr>(expr);
@@ -161,20 +160,26 @@ AffineExpr AffineExprSimplifier::RewriteMod(AffineBinaryOpExpr mod) {
       if (*multiplier % m == 0) {
         return true;
       }
+    }
+    return false;
+  });
+  new_lhs = new_lhs + (extracted_constant % m);
 
+  Interval no_multiplier_range{0, 0};
+  int64_t multiplier_gcd = -1;
+  VisitSummands(new_lhs, [&](AffineExpr expr) {
+    if (auto multiplier = GetConstantRhs(expr, AffineExprKind::Mul)) {
       if (multiplier_gcd == -1) {
         multiplier_gcd = *multiplier;
       } else {
         multiplier_gcd = std::gcd(multiplier_gcd, *multiplier);
       }
-      return false;
+    } else {
+      auto range = range_evaluator_->ComputeExpressionRange(expr);
+      no_multiplier_range.lower += range.lower;
+      no_multiplier_range.upper += range.upper;
     }
-    auto range = range_evaluator_->ComputeExpressionRange(expr);
-    no_multiplier_range.lower += range.lower;
-    no_multiplier_range.upper += range.upper;
-    return false;
   });
-  new_lhs = new_lhs + (extracted_constant % m);
 
   mlir::AffineExpr extracted = getAffineConstantExpr(0, mod.getContext());
   if (m % multiplier_gcd == 0 && no_multiplier_range.lower >= 0 &&
@@ -240,9 +245,6 @@ AffineExpr AffineExprSimplifier::RewriteFloorDiv(AffineBinaryOpExpr div) {
   }
 
   Interval no_multiplier_range{0, 0};
-  int64_t multiplier_gcd = -1;
-  // The maximum GCD of any remaining multiplier inside the div and the divisor.
-  int64_t max_remaining_multiplier_gcd = -1;
   AffineExpr zero = getAffineConstantExpr(0, mlir_context);
   AffineExpr extracted = zero;
   auto new_dividend = RemoveSummands(lhs_simplified, [&](AffineExpr expr) {
@@ -256,21 +258,7 @@ AffineExpr AffineExprSimplifier::RewriteFloorDiv(AffineBinaryOpExpr div) {
         // Remove from dividend.
         return true;
       }
-
-      if (*multiplier > 0) {
-        if (multiplier_gcd == -1) {
-          multiplier_gcd = *multiplier;
-        } else {
-          multiplier_gcd = std::gcd(multiplier_gcd, *multiplier);
-        }
-        max_remaining_multiplier_gcd =
-            std::max(max_remaining_multiplier_gcd, std::gcd(*multiplier, d));
-        return false;
-      }
     }
-    auto range = range_evaluator_->ComputeExpressionRange(expr);
-    no_multiplier_range.lower += range.lower;
-    no_multiplier_range.upper += range.upper;
     // Not a constant multiplier, keep in dividend.
     return false;
   });
@@ -280,6 +268,25 @@ AffineExpr AffineExprSimplifier::RewriteFloorDiv(AffineBinaryOpExpr div) {
     return extracted;
   }
 
+  int64_t multiplier_gcd = -1;
+  // The maximum GCD of any remaining multiplier inside the div and the divisor.
+  int64_t max_remaining_multiplier_gcd = -1;
+  VisitSummands(new_dividend, [&](AffineExpr summand) {
+    if (auto multiplier = GetConstantRhs(summand, AffineExprKind::Mul)) {
+      if (multiplier_gcd == -1) {
+        multiplier_gcd = *multiplier;
+      } else {
+        multiplier_gcd = std::gcd(multiplier_gcd, *multiplier);
+      }
+      max_remaining_multiplier_gcd =
+          std::max(max_remaining_multiplier_gcd, std::gcd(*multiplier, d));
+    } else {
+      auto range = range_evaluator_->ComputeExpressionRange(summand);
+      no_multiplier_range.lower += range.lower;
+      no_multiplier_range.upper += range.upper;
+    }
+  });
+
   if ((d % multiplier_gcd) == 0) {
     if (no_multiplier_range.lower >= 0 &&
         no_multiplier_range.upper < multiplier_gcd) {
@@ -349,6 +356,17 @@ AffineExpr AffineExprSimplifier::RemoveSummands(
   return pred(expr) ? mlir::getAffineConstantExpr(0, expr.getContext()) : expr;
 }
 
+void AffineExprSimplifier::VisitSummands(
+    mlir::AffineExpr expr, const std::function<void(mlir::AffineExpr)>& visit) {
+  if (expr.getKind() == AffineExprKind::Add) {
+    auto add = mlir::dyn_cast<AffineBinaryOpExpr>(expr);
+    VisitSummands(add.getLHS(), visit);
+    VisitSummands(add.getRHS(), visit);
+  } else {
+    visit(expr);
+  }
+}
+
 // Compares the two expression by their AST. The ordering is arbitrary but
 // similar to what MLIR's simplifier does.
 int CompareExprs(AffineExpr a, AffineExpr b) {

From aa9f630ee546a93aa90c581a70651a0b55019ae1 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Tue, 21 May 2024 08:32:58 -0700
Subject: [PATCH 468/478] Move schema_fbs_with_reflection to mlir/lite/schema.

PiperOrigin-RevId: 635816622
---
 tensorflow/compiler/mlir/lite/BUILD                |  2 +-
 .../compiler/mlir/lite/flatbuffer_to_string.cc     |  2 +-
 tensorflow/compiler/mlir/lite/schema/BUILD         | 14 ++++++++++++++
 tensorflow/lite/schema/BUILD                       |  2 +-
 4 files changed, 17 insertions(+), 3 deletions(-)

diff --git a/tensorflow/compiler/mlir/lite/BUILD b/tensorflow/compiler/mlir/lite/BUILD
index 0114f28b476694..7e49b1d028ce69 100644
--- a/tensorflow/compiler/mlir/lite/BUILD
+++ b/tensorflow/compiler/mlir/lite/BUILD
@@ -1055,8 +1055,8 @@ tf_native_cc_binary(
     name = "flatbuffer_to_string",
     srcs = ["flatbuffer_to_string.cc"],
     deps = [
+        "//tensorflow/compiler/mlir/lite/schema:schema_fbs_with_reflection",
         "//tensorflow/lite/core:model_builder",
-        "//tensorflow/lite/schema:schema_fbs_with_reflection",
         "@flatbuffers",
     ],
 )
diff --git a/tensorflow/compiler/mlir/lite/flatbuffer_to_string.cc b/tensorflow/compiler/mlir/lite/flatbuffer_to_string.cc
index b3e7e8e633e0da..df28f501ef7656 100644
--- a/tensorflow/compiler/mlir/lite/flatbuffer_to_string.cc
+++ b/tensorflow/compiler/mlir/lite/flatbuffer_to_string.cc
@@ -25,7 +25,7 @@ limitations under the License.
 
 #include "flatbuffers/flatbuffers.h"  // from @flatbuffers
 #include "flatbuffers/minireflect.h"  // from @flatbuffers
-#include "tensorflow/lite/schema/reflection/schema_generated.h"
+#include "tensorflow/compiler/mlir/lite/schema/reflection/schema_generated.h"
 #if FLATBUFFERS_LITTLEENDIAN == 0
 #include "tensorflow/lite/core/model_builder.h"
 #endif
diff --git a/tensorflow/compiler/mlir/lite/schema/BUILD b/tensorflow/compiler/mlir/lite/schema/BUILD
index bb5b4434f6f5ad..34b799a9738741 100644
--- a/tensorflow/compiler/mlir/lite/schema/BUILD
+++ b/tensorflow/compiler/mlir/lite/schema/BUILD
@@ -30,3 +30,17 @@ flatbuffer_cc_library(
     ],
     out_prefix = "mutable/",
 )
+
+# Generic schema for inference on device (but with reflections makes bigger).
+flatbuffer_cc_library(
+    name = "schema_fbs_with_reflection",
+    srcs = ["schema.fbs"],
+    compatible_with = get_compatible_with_portable(),
+    flatc_args = [
+        "--reflect-types",
+        "--reflect-names",
+        "--no-union-value-namespacing",
+        "--gen-object-api",
+    ],
+    out_prefix = "reflection/",
+)
diff --git a/tensorflow/lite/schema/BUILD b/tensorflow/lite/schema/BUILD
index 6a005359ccba7a..29d09105e59950 100644
--- a/tensorflow/lite/schema/BUILD
+++ b/tensorflow/lite/schema/BUILD
@@ -134,7 +134,7 @@ flatbuffer_cc_library(
 # Generic schema for inference on device (but with reflections makes bigger).
 flatbuffer_cc_library(
     name = "schema_fbs_with_reflection",
-    srcs = ["schema.fbs"],
+    srcs = ["//tensorflow/compiler/mlir/lite/schema:schema.fbs"],
     compatible_with = get_compatible_with_portable(),
     flatc_args = [
         "--reflect-types",

From d83a73e56299684f1f863e9aa7c923ebc38a0c54 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Tue, 21 May 2024 08:55:32 -0700
Subject: [PATCH 469/478] Migrate away from llvm::StringRef::equals

Note that llvm::StringRef::equals has been deprecated upstream.

PiperOrigin-RevId: 635822562
---
 .../runtime_fallback/kernel/kernel_fallback_execute_compat.cc   | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorflow/core/runtime_fallback/kernel/kernel_fallback_execute_compat.cc b/tensorflow/core/runtime_fallback/kernel/kernel_fallback_execute_compat.cc
index d056f90d45638d..aa48bcf6be10f0 100644
--- a/tensorflow/core/runtime_fallback/kernel/kernel_fallback_execute_compat.cc
+++ b/tensorflow/core/runtime_fallback/kernel/kernel_fallback_execute_compat.cc
@@ -286,7 +286,7 @@ tfrt::AsyncValueRef<tfrt::Chain> KernelFallbackExecuteCompatCoreRuntimeDispatch(
   // TODO(b/176997538): Skip checking dtypes for tf._BatchFunctionFallback op
   // due to b/176997538. Remove the skipping once the SavedModel lowering
   // problem is fixed.
-  if (!status.ok() && !op_name.equals("_BatchFunctionFallback")) {
+  if (!status.ok() && op_name != "_BatchFunctionFallback") {
     KernelFallbackEmitError(exec_ctx, &fallback_request_state, op_name,
                             &op_chain, results, status);
     return op_chain;

From 54cfe0f6ffc5bc427b9eb2e8a1df456a7f070636 Mon Sep 17 00:00:00 2001
From: Artsiom Ablavatski <artsiom@google.com>
Date: Tue, 21 May 2024 09:09:37 -0700
Subject: [PATCH 470/478] Align DRQ TransposeConv XNNPack test with reference
 TFLite implementation

PiperOrigin-RevId: 635826787
---
 ...namically_quantized_transpose_conv_test.cc |  44 +----
 ...mically_quantized_transpose_conv_tester.cc | 173 ++----------------
 ...amically_quantized_transpose_conv_tester.h |   6 +-
 3 files changed, 22 insertions(+), 201 deletions(-)

diff --git a/tensorflow/lite/delegates/xnnpack/dynamically_quantized_transpose_conv_test.cc b/tensorflow/lite/delegates/xnnpack/dynamically_quantized_transpose_conv_test.cc
index 3b83541b81b215..de863e4f1e2125 100644
--- a/tensorflow/lite/delegates/xnnpack/dynamically_quantized_transpose_conv_test.cc
+++ b/tensorflow/lite/delegates/xnnpack/dynamically_quantized_transpose_conv_test.cc
@@ -27,11 +27,8 @@ namespace tflite {
 namespace xnnpack {
 
 TEST(DynamicallyQuantizedTransposeConvTest, 2x2Stride2) {
-  auto delegate_options = TfLiteXNNPackDelegateOptionsDefault();
-  delegate_options.flags |=
-      TFLITE_XNNPACK_DELEGATE_FLAG_ENABLE_SUBGRAPH_RESHAPING;
   std::unique_ptr<TfLiteDelegate, decltype(&TfLiteXNNPackDelegateDelete)>
-      xnnpack_delegate(TfLiteXNNPackDelegateCreate(&delegate_options),
+      xnnpack_delegate(TfLiteXNNPackDelegateCreate(nullptr),
                        TfLiteXNNPackDelegateDelete);
 
   std::random_device random_device;
@@ -55,11 +52,8 @@ TEST(DynamicallyQuantizedTransposeConvTest, 2x2Stride2) {
 }
 
 TEST(DynamicallyQuantizedTransposeConvTest, 3x3Stride2) {
-  auto delegate_options = TfLiteXNNPackDelegateOptionsDefault();
-  delegate_options.flags |=
-      TFLITE_XNNPACK_DELEGATE_FLAG_ENABLE_SUBGRAPH_RESHAPING;
   std::unique_ptr<TfLiteDelegate, decltype(&TfLiteXNNPackDelegateDelete)>
-      xnnpack_delegate(TfLiteXNNPackDelegateCreate(&delegate_options),
+      xnnpack_delegate(TfLiteXNNPackDelegateCreate(nullptr),
                        TfLiteXNNPackDelegateDelete);
   std::random_device random_device;
   auto rng = std::mt19937(random_device());
@@ -82,11 +76,8 @@ TEST(DynamicallyQuantizedTransposeConvTest, 3x3Stride2) {
 }
 
 TEST(DynamicallyQuantizedTransposeConvTest, 4x4Stride2) {
-  auto delegate_options = TfLiteXNNPackDelegateOptionsDefault();
-  delegate_options.flags |=
-      TFLITE_XNNPACK_DELEGATE_FLAG_ENABLE_SUBGRAPH_RESHAPING;
   std::unique_ptr<TfLiteDelegate, decltype(&TfLiteXNNPackDelegateDelete)>
-      xnnpack_delegate(TfLiteXNNPackDelegateCreate(&delegate_options),
+      xnnpack_delegate(TfLiteXNNPackDelegateCreate(nullptr),
                        TfLiteXNNPackDelegateDelete);
 
   std::random_device random_device;
@@ -110,11 +101,8 @@ TEST(DynamicallyQuantizedTransposeConvTest, 4x4Stride2) {
 }
 
 TEST(DynamicallyQuantizedTransposeConvTest, 4x4Stride4) {
-  auto delegate_options = TfLiteXNNPackDelegateOptionsDefault();
-  delegate_options.flags |=
-      TFLITE_XNNPACK_DELEGATE_FLAG_ENABLE_SUBGRAPH_RESHAPING;
   std::unique_ptr<TfLiteDelegate, decltype(&TfLiteXNNPackDelegateDelete)>
-      xnnpack_delegate(TfLiteXNNPackDelegateCreate(&delegate_options),
+      xnnpack_delegate(TfLiteXNNPackDelegateCreate(nullptr),
                        TfLiteXNNPackDelegateDelete);
 
   std::random_device random_device;
@@ -138,11 +126,8 @@ TEST(DynamicallyQuantizedTransposeConvTest, 4x4Stride4) {
 }
 
 TEST(DynamicallyQuantizedTransposeConvTest, SmallKernelWithSamePadding) {
-  auto delegate_options = TfLiteXNNPackDelegateOptionsDefault();
-  delegate_options.flags |=
-      TFLITE_XNNPACK_DELEGATE_FLAG_ENABLE_SUBGRAPH_RESHAPING;
   std::unique_ptr<TfLiteDelegate, decltype(&TfLiteXNNPackDelegateDelete)>
-      xnnpack_delegate(TfLiteXNNPackDelegateCreate(&delegate_options),
+      xnnpack_delegate(TfLiteXNNPackDelegateCreate(nullptr),
                        TfLiteXNNPackDelegateDelete);
 
   std::random_device random_device;
@@ -169,11 +154,8 @@ TEST(DynamicallyQuantizedTransposeConvTest, SmallKernelWithSamePadding) {
 }
 
 TEST(DynamicallyQuantizedTransposeConvTest, SmallKernelWithValidPadding) {
-  auto delegate_options = TfLiteXNNPackDelegateOptionsDefault();
-  delegate_options.flags |=
-      TFLITE_XNNPACK_DELEGATE_FLAG_ENABLE_SUBGRAPH_RESHAPING;
   std::unique_ptr<TfLiteDelegate, decltype(&TfLiteXNNPackDelegateDelete)>
-      xnnpack_delegate(TfLiteXNNPackDelegateCreate(&delegate_options),
+      xnnpack_delegate(TfLiteXNNPackDelegateCreate(nullptr),
                        TfLiteXNNPackDelegateDelete);
   std::random_device random_device;
   auto rng = std::mt19937(random_device());
@@ -199,11 +181,8 @@ TEST(DynamicallyQuantizedTransposeConvTest, SmallKernelWithValidPadding) {
 }
 
 TEST(DynamicallyQuantizedTransposeConvTest, StrideWithSamePadding) {
-  auto delegate_options = TfLiteXNNPackDelegateOptionsDefault();
-  delegate_options.flags |=
-      TFLITE_XNNPACK_DELEGATE_FLAG_ENABLE_SUBGRAPH_RESHAPING;
   std::unique_ptr<TfLiteDelegate, decltype(&TfLiteXNNPackDelegateDelete)>
-      xnnpack_delegate(TfLiteXNNPackDelegateCreate(&delegate_options),
+      xnnpack_delegate(TfLiteXNNPackDelegateCreate(nullptr),
                        TfLiteXNNPackDelegateDelete);
 
   std::random_device random_device;
@@ -234,11 +213,8 @@ TEST(DynamicallyQuantizedTransposeConvTest, StrideWithSamePadding) {
 }
 
 TEST(DynamicallyQuantizedTransposeConvTest, StrideWithValidPadding) {
-  auto delegate_options = TfLiteXNNPackDelegateOptionsDefault();
-  delegate_options.flags |=
-      TFLITE_XNNPACK_DELEGATE_FLAG_ENABLE_SUBGRAPH_RESHAPING;
   std::unique_ptr<TfLiteDelegate, decltype(&TfLiteXNNPackDelegateDelete)>
-      xnnpack_delegate(TfLiteXNNPackDelegateCreate(&delegate_options),
+      xnnpack_delegate(TfLiteXNNPackDelegateCreate(nullptr),
                        TfLiteXNNPackDelegateDelete);
 
   std::random_device random_device;
@@ -272,8 +248,6 @@ TEST(DynamicallyQuantizedTransposeConvTest, MultiThreading) {
   TfLiteXNNPackDelegateOptions delegate_options =
       TfLiteXNNPackDelegateOptionsDefault();
   delegate_options.num_threads = 2;
-  delegate_options.flags |=
-      TFLITE_XNNPACK_DELEGATE_FLAG_ENABLE_SUBGRAPH_RESHAPING;
   std::unique_ptr<TfLiteDelegate, decltype(&TfLiteXNNPackDelegateDelete)>
       xnnpack_delegate(TfLiteXNNPackDelegateCreate(&delegate_options),
                        TfLiteXNNPackDelegateDelete);
@@ -308,8 +282,6 @@ TEST(DynamicallyQuantizedTransposeConvTest, MultiThreading) {
 TEST(DynamicallyQuantizedTransposeConvTest, WeightsCache) {
   TfLiteXNNPackDelegateOptions delegate_options =
       TfLiteXNNPackDelegateOptionsDefault();
-  delegate_options.flags |=
-      TFLITE_XNNPACK_DELEGATE_FLAG_ENABLE_SUBGRAPH_RESHAPING;
   std::unique_ptr<TfLiteXNNPackDelegateWeightsCache,
                   decltype(&TfLiteXNNPackDelegateWeightsCacheDelete)>
       weights_cache(TfLiteXNNPackDelegateWeightsCacheCreate(),
diff --git a/tensorflow/lite/delegates/xnnpack/dynamically_quantized_transpose_conv_tester.cc b/tensorflow/lite/delegates/xnnpack/dynamically_quantized_transpose_conv_tester.cc
index bbc7110920ae08..35c1b590c7396a 100644
--- a/tensorflow/lite/delegates/xnnpack/dynamically_quantized_transpose_conv_tester.cc
+++ b/tensorflow/lite/delegates/xnnpack/dynamically_quantized_transpose_conv_tester.cc
@@ -49,23 +49,19 @@ void DynamicallyQuantizedTransposeConvTester::Test(
   const std::vector<int8_t> kernel_data = GenerateKernelData();
   const std::vector<float> bias_data = GenerateBiasData();
   const std::vector<float> kernel_scale_data = GenerateKernelScaleData();
-  std::vector<char> drq_buffer =
-      CreateDRQTfLiteModel(kernel_data, bias_data, kernel_scale_data);
-  std::vector<char> dequantize_buffer =
-      CreateDequantizeTfLiteModel(kernel_data, bias_data, kernel_scale_data);
-  const Model* drq_model = GetModel(drq_buffer.data());
-  const Model* dequantize_model = GetModel(dequantize_buffer.data());
+  std::vector<char> buffer =
+      CreateTfLiteModel(kernel_data, bias_data, kernel_scale_data);
+  const Model* model = GetModel(buffer.data());
 
   std::unique_ptr<Interpreter> delegate_interpreter;
-  ASSERT_EQ(
-      InterpreterBuilder(
-          drq_model, ::tflite::ops::builtin::BuiltinOpResolverWithXNNPACK())(
-          &delegate_interpreter),
-      kTfLiteOk);
+  ASSERT_EQ(InterpreterBuilder(
+                model, ::tflite::ops::builtin::BuiltinOpResolverWithXNNPACK())(
+                &delegate_interpreter),
+            kTfLiteOk);
   std::unique_ptr<Interpreter> default_interpreter;
   ASSERT_EQ(
       InterpreterBuilder(
-          dequantize_model,
+          model,
           ::tflite::ops::builtin::BuiltinOpResolverWithoutDefaultDelegates())(
           &default_interpreter),
       kTfLiteOk);
@@ -79,11 +75,11 @@ void DynamicallyQuantizedTransposeConvTester::Test(
   ASSERT_EQ(delegate_interpreter->outputs().size(), 1);
   ASSERT_EQ(default_interpreter->outputs().size(), 1);
 
-  ASSERT_EQ(delegate_interpreter->ModifyGraphWithDelegate(delegate), kTfLiteOk);
-
   ASSERT_EQ(delegate_interpreter->AllocateTensors(), kTfLiteOk);
   ASSERT_EQ(default_interpreter->AllocateTensors(), kTfLiteOk);
 
+  ASSERT_EQ(delegate_interpreter->ModifyGraphWithDelegate(delegate), kTfLiteOk);
+
   if (weights_cache_ != nullptr) {
     TfLiteXNNPackDelegateWeightsCacheFinalizeHard(weights_cache_);
   }
@@ -110,12 +106,11 @@ void DynamicallyQuantizedTransposeConvTester::Test(
   int different_output_values = 0;
   for (size_t i = 0; i < num_output_values; i++) {
     if (std::abs(default_output_data[i] - xnnpack_output_data[i]) >
-        0.1 * std::abs(default_output_data[i])) {
+        0.005 * std::abs(default_output_data[i])) {
       ++different_output_values;
     }
   }
-
-  if (different_output_values > 0.055 * num_output_values) {
+  if (different_output_values > 0.05 * num_output_values) {
     GTEST_FAIL() << (float)different_output_values / num_output_values * 100.f
                  << "% of output values differ";
   }
@@ -182,7 +177,7 @@ DynamicallyQuantizedTransposeConvTester::GenerateKernelScaleData() const {
   return kernel_scale;
 }
 
-std::vector<char> DynamicallyQuantizedTransposeConvTester::CreateDRQTfLiteModel(
+std::vector<char> DynamicallyQuantizedTransposeConvTester::CreateTfLiteModel(
     const std::vector<int8_t>& filter_data, const std::vector<float>& bias_data,
     const std::vector<float>& kernel_scale) const {
   /*************************** Define operator codes **************************/
@@ -297,147 +292,5 @@ std::vector<char> DynamicallyQuantizedTransposeConvTester::CreateDRQTfLiteModel(
                            builder.GetBufferPointer() + builder.GetSize());
 }
 
-std::vector<char>
-DynamicallyQuantizedTransposeConvTester::CreateDequantizeTfLiteModel(
-    const std::vector<int8_t>& filter_data, const std::vector<float>& bias_data,
-    const std::vector<float>& kernel_scale) const {
-  /*************************** Define operator codes **************************/
-  flatbuffers::FlatBufferBuilder builder;
-  std::vector<flatbuffers::Offset<OperatorCode>> operator_codes{
-      {CreateOperatorCode(builder, BuiltinOperator_TRANSPOSE_CONV)}};
-  const int dequantize_operator_code = operator_codes.size();
-  operator_codes.emplace_back(
-      CreateOperatorCode(builder, BuiltinOperator_DEQUANTIZE));
-
-  /****************************** Define buffers ******************************/
-  std::vector<flatbuffers::Offset<tflite::Buffer>> buffers{
-      {CreateBuffer(builder, builder.CreateVector({}))}};
-
-  int filter_buffer_id = 0;
-  const int quantized_filter_buffer_id = buffers.size();
-  buffers.emplace_back(CreateBuffer(
-      builder,
-      builder.CreateVector(reinterpret_cast<const uint8_t*>(filter_data.data()),
-                           sizeof(int8_t) * filter_data.size())));
-
-  int bias_buffer_id = buffers.size();
-  buffers.emplace_back(CreateBuffer(
-      builder,
-      builder.CreateVector(reinterpret_cast<const uint8_t*>(bias_data.data()),
-                           sizeof(float) * bias_data.size())));
-  const std::array<int32_t, 4> output_shape{
-      {BatchSize(), OutputHeight(), OutputWidth(), OutputChannels()}};
-  const int output_shape_buffer_id = buffers.size();
-  buffers.emplace_back(CreateBuffer(
-      builder, builder.CreateVector(
-                   reinterpret_cast<const uint8_t*>(output_shape.data()),
-                   sizeof(int32_t) * output_shape.size())));
-
-  /****************************** Define tensors ******************************/
-  const std::vector<int32_t> filter_shape = {OutputChannels(), KernelHeight(),
-                                             KernelWidth(), InputChannels()};
-  const std::vector<int32_t> bias_shape = {OutputChannels()};
-  const std::array<int32_t, 4> input_shape{
-      {BatchSize(), InputHeight(), InputWidth(), InputChannels()}};
-
-  std::vector<flatbuffers::Offset<tflite::Tensor>> tensors;
-  const int quantized_filter_tensor_id = tensors.size();
-  tensors.emplace_back(CreateTensor(
-      builder,
-      builder.CreateVector<int32_t>(filter_shape.data(), filter_shape.size()),
-      /*type=*/TensorType_INT8,
-      /*buffer=*/quantized_filter_buffer_id,
-      /*name=*/0,
-      CreateQuantizationParameters(
-          builder, /*min=*/0, /*max=*/0,
-          builder.CreateVector<float>(kernel_scale),
-          builder.CreateVector<int64_t>(
-              std::vector<int64_t>(OutputChannels(), 0)))));
-
-  const int input_tensor_id = tensors.size();
-  tensors.emplace_back(CreateTensor(
-      builder,
-      builder.CreateVector<int32_t>(input_shape.data(), input_shape.size()),
-      TensorType_FLOAT32));
-
-  const int filter_tensor_id = tensors.size();
-  tensors.emplace_back(CreateTensor(
-      builder,
-      builder.CreateVector<int32_t>(filter_shape.data(), filter_shape.size()),
-      TensorType_FLOAT32,
-      /*buffer=*/filter_buffer_id));
-
-  const int bias_tensor_id = tensors.size();
-  tensors.emplace_back(CreateTensor(
-      builder,
-      builder.CreateVector<int32_t>(bias_shape.data(), bias_shape.size()),
-      TensorType_FLOAT32, bias_buffer_id));
-
-  const int output_tensor_id = tensors.size();
-  tensors.emplace_back(CreateTensor(
-      builder,
-      builder.CreateVector<int32_t>(output_shape.data(), output_shape.size()),
-      TensorType_FLOAT32));
-
-  const int output_shape_tensor_id = tensors.size();
-  const std::array<int32_t, 1> output_shape_shape{{4}};
-  tensors.emplace_back(
-      CreateTensor(builder,
-                   builder.CreateVector<int32_t>(output_shape_shape.data(),
-                                                 output_shape_shape.size()),
-                   TensorType_INT32, output_shape_buffer_id));
-
-  /***************************** Define operators *****************************/
-  std::vector<flatbuffers::Offset<tflite::Operator>> operators;
-
-  const std::array<int32_t, 1> dequantize_filter_inputs{
-      {quantized_filter_tensor_id}};
-  const std::array<int32_t, 1> dequantize_filter_outputs{{filter_tensor_id}};
-  operators.emplace_back(CreateOperator(
-      builder, /*opcode_index=*/dequantize_operator_code,
-      builder.CreateVector<int32_t>(dequantize_filter_inputs.data(),
-                                    dequantize_filter_inputs.size()),
-      builder.CreateVector<int32_t>(dequantize_filter_outputs.data(),
-                                    dequantize_filter_outputs.size())));
-
-  std::vector<int32_t> op_inputs{
-      {output_shape_tensor_id, filter_tensor_id, input_tensor_id}};
-  op_inputs.push_back(bias_tensor_id);
-  const std::array<int32_t, 1> op_outputs{{output_tensor_id}};
-  const flatbuffers::Offset<TransposeConvOptions> transpose_conv_options =
-      CreateTransposeConvOptions(builder, Padding(), StrideWidth(),
-                                 StrideHeight());
-  operators.emplace_back(CreateOperator(
-      builder, /*opcode_index=*/0,
-      builder.CreateVector<int32_t>(op_inputs.data(), op_inputs.size()),
-      builder.CreateVector<int32_t>(op_outputs.data(), op_outputs.size()),
-      BuiltinOptions_TransposeConvOptions, transpose_conv_options.Union()));
-
-  /****************************** Define subgraph *****************************/
-  const std::array<int32_t, 1> subgraph_inputs{{input_tensor_id}};
-  const std::array<int32_t, 1> subgraph_outputs{{output_tensor_id}};
-  const flatbuffers::Offset<SubGraph> subgraph = CreateSubGraph(
-      builder, builder.CreateVector(tensors.data(), tensors.size()),
-      builder.CreateVector<int32_t>(subgraph_inputs.data(),
-                                    subgraph_inputs.size()),
-      builder.CreateVector<int32_t>(subgraph_outputs.data(),
-                                    subgraph_outputs.size()),
-      builder.CreateVector(operators.data(), operators.size()));
-
-  const flatbuffers::Offset<flatbuffers::String> description =
-      builder.CreateString("TransposeConv model");
-
-  const flatbuffers::Offset<Model> model_buffer = CreateModel(
-      builder, TFLITE_SCHEMA_VERSION,
-      builder.CreateVector(operator_codes.data(), operator_codes.size()),
-      builder.CreateVector(&subgraph, 1), description,
-      builder.CreateVector(buffers.data(), buffers.size()));
-
-  builder.Finish(model_buffer);
-
-  return std::vector<char>(builder.GetBufferPointer(),
-                           builder.GetBufferPointer() + builder.GetSize());
-}
-
 }  // namespace xnnpack
 }  // namespace tflite
diff --git a/tensorflow/lite/delegates/xnnpack/dynamically_quantized_transpose_conv_tester.h b/tensorflow/lite/delegates/xnnpack/dynamically_quantized_transpose_conv_tester.h
index e7dc7cd9536608..5c69808c836b33 100644
--- a/tensorflow/lite/delegates/xnnpack/dynamically_quantized_transpose_conv_tester.h
+++ b/tensorflow/lite/delegates/xnnpack/dynamically_quantized_transpose_conv_tester.h
@@ -183,11 +183,7 @@ class DynamicallyQuantizedTransposeConvTester {
   std::vector<int8_t> GenerateKernelData() const;
   std::vector<float> GenerateBiasData() const;
   std::vector<float> GenerateKernelScaleData() const;
-  std::vector<char> CreateDRQTfLiteModel(
-      const std::vector<int8_t>& filter_data,
-      const std::vector<float>& bias_data,
-      const std::vector<float>& kernel_scale) const;
-  std::vector<char> CreateDequantizeTfLiteModel(
+  std::vector<char> CreateTfLiteModel(
       const std::vector<int8_t>& filter_data,
       const std::vector<float>& bias_data,
       const std::vector<float>& kernel_scale) const;

From 84d52706e1a7a0eb0ef0703e09016a124af7f451 Mon Sep 17 00:00:00 2001
From: Dragan Mladjenovic <Dragan.Mladjenovic@amd.com>
Date: Tue, 21 May 2024 17:04:07 +0000
Subject: [PATCH 471/478] Merge fixes after 240521

---
 tensorflow/core/grappler/clusters/utils.cc                 | 5 -----
 .../xla/xla/backends/profiler/gpu/rocm_collector.cc        | 5 -----
 third_party/xla/xla/mlir_hlo/tests/BUILD                   | 7 -------
 third_party/xla/xla/service/gpu/model/BUILD                | 4 ----
 4 files changed, 21 deletions(-)

diff --git a/tensorflow/core/grappler/clusters/utils.cc b/tensorflow/core/grappler/clusters/utils.cc
index a69a102cf361ea..708fdb1a2baacc 100644
--- a/tensorflow/core/grappler/clusters/utils.cc
+++ b/tensorflow/core/grappler/clusters/utils.cc
@@ -135,7 +135,6 @@ DeviceProperties GetLocalGPUInfo(PlatformDeviceId platform_device_id) {
   device.set_bandwidth(properties.memoryBusWidth / 8 *
                        properties.memoryClockRate * 2);
 
-<<<<<<< HEAD
   std::string gcnFullName = properties.gcnArchName;
   std::vector<std::string> tokens = absl::StrSplit(gcnFullName, ':');
   std::string gcnName = gcnFullName;
@@ -143,10 +142,6 @@ DeviceProperties GetLocalGPUInfo(PlatformDeviceId platform_device_id) {
       gcnName = tokens[0];
   }
   (*device.mutable_environment())["architecture"] = gcnName;
-=======
-  (*device.mutable_environment())["architecture"] =
-      strings::StrCat("gfx", properties.gcnArchName);
->>>>>>> upstream/master
 #endif
 
   return device;
diff --git a/third_party/xla/xla/backends/profiler/gpu/rocm_collector.cc b/third_party/xla/xla/backends/profiler/gpu/rocm_collector.cc
index 8405e6079cf14f..f817a5e424d953 100644
--- a/third_party/xla/xla/backends/profiler/gpu/rocm_collector.cc
+++ b/third_party/xla/xla/backends/profiler/gpu/rocm_collector.cc
@@ -224,13 +224,8 @@ class PerDeviceCollector {
     stats.occupancy_pct /= device_properties_.maxThreadsPerMultiProcessor;
 
     err = hipOccupancyMaxPotentialBlockSize(
-<<<<<<< HEAD
-        &stats.min_grid_size, &stats.suggested_block_size, static_cast<const void*>(params.func_ptr),
-        params.dynamic_smem_size, 0);
-=======
         &stats.min_grid_size, &stats.suggested_block_size,
         static_cast<const void*>(params.func_ptr), params.dynamic_smem_size, 0);
->>>>>>> upstream/master
 
     if (err != hipError_t::hipSuccess) {
       return {};
diff --git a/third_party/xla/xla/mlir_hlo/tests/BUILD b/third_party/xla/xla/mlir_hlo/tests/BUILD
index 239d706014233a..4db81efeb2dd52 100644
--- a/third_party/xla/xla/mlir_hlo/tests/BUILD
+++ b/third_party/xla/xla/mlir_hlo/tests/BUILD
@@ -26,14 +26,7 @@ package(
         tags = [
             "nomsan",  # The execution engine doesn't work with msan, see b/248097619.
         ],
-<<<<<<< HEAD
-        #Build failing due to the double deps definition. Should be resolved by next weekly thanks to 
-        #https://github.com/llvm/llvm-project/commit/dd5797505ebc2dbfdd58927c4f0a11a1256696eb
-        #deps = ["@pypi_lit//:pkg"],
-=======
         deps = ["@pypi_lit//:pkg"],
-        deps = ["@pypi_lit//:pkg"],
->>>>>>> upstream/master
     )
     for src in glob(["**/*.mlir"])
 ]
diff --git a/third_party/xla/xla/service/gpu/model/BUILD b/third_party/xla/xla/service/gpu/model/BUILD
index cfeb4206cc88f2..759476d142d306 100644
--- a/third_party/xla/xla/service/gpu/model/BUILD
+++ b/third_party/xla/xla/service/gpu/model/BUILD
@@ -549,12 +549,8 @@ xla_cc_test(
         "//xla/tests:xla_internal_test_main",
         "@com_google_absl//absl/log:check",
         "@com_google_absl//absl/strings",
-<<<<<<< HEAD
-        #"@com_google_absl//absl/strings:string_view",
-=======
         "@com_google_absl//absl/strings:string_view",
         "@com_google_absl//absl/types:span",
->>>>>>> upstream/master
         "@com_google_googletest//:gtest",
         "@llvm-project//llvm:Support",
         "@llvm-project//mlir:IR",

From 6af30d6bfc2fe52a9e5dbc88ca9afd8f905a97a3 Mon Sep 17 00:00:00 2001
From: Dragan Mladjenovic <Dragan.Mladjenovic@amd.com>
Date: Tue, 21 May 2024 21:00:23 +0000
Subject: [PATCH 472/478] Fix tensorflow/core/kernels/conv_ops_gpu.cc build
 failure

---
 tensorflow/core/kernels/conv_ops_gpu.cc | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/tensorflow/core/kernels/conv_ops_gpu.cc b/tensorflow/core/kernels/conv_ops_gpu.cc
index d1526b44aba08c..60238d5cd39f47 100644
--- a/tensorflow/core/kernels/conv_ops_gpu.cc
+++ b/tensorflow/core/kernels/conv_ops_gpu.cc
@@ -363,9 +363,9 @@ StatusOr<AutotuneEntry<se::dnn::ConvOp>> AutotuneUnfusedConv(
       return absl::InvalidArgumentError("No DNN in stream executor.");
     }
     if (!dnn->GetMIOpenConvolveAlgorithms(
-            kind, se::dnn::ToDataType<T>::value, stream, input_desc, input_ptr,
-            filter_desc, filter_ptr, output_desc, output_ptr, conv_desc,
-            &scratch_allocator, &algorithms)) {
+            kind, se::dnn::ToDataType<T>::value, se::dnn::ToDataType<T>::value,
+            stream, input_desc, input_ptr, filter_desc, filter_ptr, output_desc,
+            output_ptr, conv_desc, &scratch_allocator, &algorithms)) {
       return errors::Unknown(
           "Failed to get convolution algorithm. This is probably "
           "because MIOpen failed to initialize, so try looking to "

From cf09f6adea48688d1775cf3a272377239a4c5f40 Mon Sep 17 00:00:00 2001
From: Dragan Mladjenovic <Dragan.Mladjenovic@amd.com>
Date: Sat, 25 May 2024 08:16:07 +0000
Subject: [PATCH 473/478] [ROCM][DO NOT UPSTREAM] Pin libamdhip64.so version

To avoid loading both system and locald_rocm_config version of the
library. WORKAROUND!
---
 .../third_party/tsl/tsl/platform/default/dso_loader.cc    | 8 +++++++-
 1 file changed, 7 insertions(+), 1 deletion(-)

diff --git a/third_party/xla/third_party/tsl/tsl/platform/default/dso_loader.cc b/third_party/xla/third_party/tsl/tsl/platform/default/dso_loader.cc
index a835a81489367a..39c09808530185 100644
--- a/third_party/xla/third_party/tsl/tsl/platform/default/dso_loader.cc
+++ b/third_party/xla/third_party/tsl/tsl/platform/default/dso_loader.cc
@@ -181,7 +181,13 @@ absl::StatusOr<void*> GetHipblasltDsoHandle() {
   return GetDsoHandle("hipblaslt", "");
 }
 
-absl::StatusOr<void*> GetHipDsoHandle() { return GetDsoHandle("amdhip64", ""); }
+#if (TF_ROCM_VERSION >= 60000)
+#define HIP_SO_VERSION "6"
+#else
+#define HIP_SO_VERSION "5"
+#endif
+
+absl::StatusOr<void*> GetHipDsoHandle() { return GetDsoHandle("amdhip64", HIP_SO_VERSION); }
 
 }  // namespace DsoLoader
 

From 4da4f80dcb5db94904404741f730f6b902ef4b62 Mon Sep 17 00:00:00 2001
From: Dragan Mladjenovic <Dragan.Mladjenovic@amd.com>
Date: Tue, 28 May 2024 17:04:04 +0000
Subject: [PATCH 474/478] Fix
 //tensorflow/core/common_runtime/gpu:gpu_device_test_gpu failure

---
 tensorflow/core/common_runtime/gpu/BUILD | 1 +
 1 file changed, 1 insertion(+)

diff --git a/tensorflow/core/common_runtime/gpu/BUILD b/tensorflow/core/common_runtime/gpu/BUILD
index 3c61cb528b6e38..d5e984dc44b8bd 100644
--- a/tensorflow/core/common_runtime/gpu/BUILD
+++ b/tensorflow/core/common_runtime/gpu/BUILD
@@ -308,6 +308,7 @@ tf_cuda_cc_test(
     srcs = [
         "gpu_device_test.cc",
     ],
+    extra_copts = ["-DXLA_TEST_BACKEND_GPU=1"], # This is a mess
     features = ["-layering_check"],
     tags = tf_cuda_tests_tags(),
     deps = [

From 473684840d7f029d5ca20bb0dbf84f978bd45cb7 Mon Sep 17 00:00:00 2001
From: Dragan Mladjenovic <Dragan.Mladjenovic@amd.com>
Date: Wed, 29 May 2024 17:22:18 +0000
Subject: [PATCH 475/478] [ROCm] Fix profiler deadlock due to
 3d7a4720b02e643faccd13f9a450213b5dfe17b7

---
 tensorflow/compiler/tests/BUILD                             | 2 +-
 third_party/xla/xla/backends/profiler/gpu/rocm_collector.cc | 6 +++---
 2 files changed, 4 insertions(+), 4 deletions(-)

diff --git a/tensorflow/compiler/tests/BUILD b/tensorflow/compiler/tests/BUILD
index d255b67ccff83f..956296e561f205 100644
--- a/tensorflow/compiler/tests/BUILD
+++ b/tensorflow/compiler/tests/BUILD
@@ -2082,7 +2082,7 @@ cuda_py_strict_test(
 
 cuda_py_strict_test(
     name = "dense_layer_test",
-    size = "medium",
+    size = "large",
     srcs = ["dense_layer_test.py"],
     tags = [
         "no_pip",  # TODO(b/149738646): fix pip install so these tests run on kokoro pip
diff --git a/third_party/xla/xla/backends/profiler/gpu/rocm_collector.cc b/third_party/xla/xla/backends/profiler/gpu/rocm_collector.cc
index f817a5e424d953..d18b0a5d513842 100644
--- a/third_party/xla/xla/backends/profiler/gpu/rocm_collector.cc
+++ b/third_party/xla/xla/backends/profiler/gpu/rocm_collector.cc
@@ -613,9 +613,10 @@ class RocmTraceCollectorImpl : public profiler::RocmTraceCollector {
   absl::flat_hash_map<uint32_t, RocmTracerEvent> auxiliary_api_events_map_
       TF_GUARDED_BY(event_maps_mutex_);
 
-  const std::vector<RocmTracerEvent> ApiActivityInfoExchange();
+  const std::vector<RocmTracerEvent> ApiActivityInfoExchange()
+      TF_EXCLUSIVE_LOCKS_REQUIRED(event_maps_mutex_);
 
-  absl::node_hash_map<uint32_t, PerDeviceCollector> per_device_collector_;
+  absl::flat_hash_map<uint32_t, PerDeviceCollector> per_device_collector_;
 };
 //==========
 
@@ -730,7 +731,6 @@ RocmTraceCollectorImpl::ApiActivityInfoExchange() {
   std::vector<RocmTracerEvent> aggregated_events;
 
   // Copy info from activity events to API callback events
-  mutex_lock lock{event_maps_mutex_};
   for (auto& api_iter : api_events_map_) {
     RocmTracerEvent& api_event = api_iter.second;
     auto activity_event =

From cc4a257ce03254e6868817a821f99dcf0b50116c Mon Sep 17 00:00:00 2001
From: Dragan Mladjenovic <Dragan.Mladjenovic@amd.com>
Date: Mon, 3 Jun 2024 17:46:35 +0000
Subject: [PATCH 476/478] Disable some failing XLA tests

---
 third_party/xla/xla/service/gpu/BUILD         | 4 +---
 third_party/xla/xla/service/gpu/fusions/BUILD | 6 ++++++
 third_party/xla/xla/service/gpu/tests/BUILD   | 1 +
 third_party/xla/xla/tests/BUILD               | 3 +++
 4 files changed, 11 insertions(+), 3 deletions(-)

diff --git a/third_party/xla/xla/service/gpu/BUILD b/third_party/xla/xla/service/gpu/BUILD
index ad80d33e8a2697..24af9775bee0d8 100644
--- a/third_party/xla/xla/service/gpu/BUILD
+++ b/third_party/xla/xla/service/gpu/BUILD
@@ -4637,9 +4637,6 @@ cc_library(
 xla_test(
     name = "cudnn_fused_conv_rewriter_test",
     srcs = ["cudnn_fused_conv_rewriter_test.cc"],
-    tags = [
-        "no_rocm"
-    ],
     backend_tags = {
         "gpu_a100": [
             "noasan",
@@ -5870,6 +5867,7 @@ xla_cc_test(
 xla_test(
     name = "determinism_test",
     srcs = if_gpu_is_configured(["determinism_test.cc"]),
+    tags = ["no_rocm"],
     backends = [
         "gpu_a100",
     ],
diff --git a/third_party/xla/xla/service/gpu/fusions/BUILD b/third_party/xla/xla/service/gpu/fusions/BUILD
index 60ce65ac397d68..14fd6dda63eb45 100644
--- a/third_party/xla/xla/service/gpu/fusions/BUILD
+++ b/third_party/xla/xla/service/gpu/fusions/BUILD
@@ -84,6 +84,7 @@ cc_library(
 xla_cc_test(
     name = "in_place_dynamic_update_slice_mlir_test",
     srcs = ["in_place_dynamic_update_slice_mlir_test.cc"],
+    tags = ["no_rocm"], #TODO(rocm) TEMP
     use_gpu = True,
     deps = [
         ":in_place_dynamic_update_slice_mlir",
@@ -438,6 +439,7 @@ cc_library(
 xla_cc_test(
     name = "scatter_mlir_test",
     srcs = ["scatter_mlir_test.cc"],
+    tags = ["no_rocm"],
     use_gpu = True,
     deps = [
         ":mlir_emitter_test_base",
@@ -491,6 +493,7 @@ cc_library(
 xla_cc_test(
     name = "transpose_mlir_test",
     srcs = ["transpose_mlir_test.cc"],
+    tags = ["no_rocm"], #TODO(rocm) TEMP
     use_gpu = True,
     deps = [
         ":mlir_emitter_test_base",
@@ -883,6 +886,7 @@ cc_library(
 xla_cc_test(
     name = "reduction_mlir_test",
     srcs = ["reduction_mlir_test.cc"],
+    tags = ["no_rocm"], #TODO(rocm) TEMP
     use_gpu = True,
     deps = [
         ":mlir_emitter_test_base",
@@ -969,6 +973,7 @@ cc_library(
 xla_cc_test(
     name = "concatenate_mlir_test",
     srcs = ["concatenate_mlir_test.cc"],
+    tags = ["no_rocm"], #TODO(rocm) TEMP
     use_gpu = True,
     deps = [
         ":concatenate_mlir",
@@ -1107,6 +1112,7 @@ cc_library(
 xla_cc_test(
     name = "input_slices_mlir_test",
     srcs = ["input_slices_mlir_test.cc"],
+    tags = ["no_rocm"], #TODO(rocm) TEMP
     use_gpu = True,
     deps = [
         ":input_slices_mlir",
diff --git a/third_party/xla/xla/service/gpu/tests/BUILD b/third_party/xla/xla/service/gpu/tests/BUILD
index 5740336c265a8c..0a80cf497fa3ad 100644
--- a/third_party/xla/xla/service/gpu/tests/BUILD
+++ b/third_party/xla/xla/service/gpu/tests/BUILD
@@ -447,6 +447,7 @@ xla_cc_test(
 xla_test(
     name = "gpu_triton_custom_call_test",
     srcs = ["gpu_triton_custom_call_test.cc"],
+    tags = ["no_rocm"], #TODO(rocm) TEMP
     backends = [
         "gpu_a100",
         "gpu_v100",
diff --git a/third_party/xla/xla/tests/BUILD b/third_party/xla/xla/tests/BUILD
index 06e9148e37a79f..f7451a7dc54f7f 100644
--- a/third_party/xla/xla/tests/BUILD
+++ b/third_party/xla/xla/tests/BUILD
@@ -900,6 +900,7 @@ xla_test(
     shard_count = 20,
     tags = [
         "optonly",
+        "no_rocm",
     ],
     deps = [
         ":client_library_test_base",
@@ -937,6 +938,7 @@ xla_test(
         "optonly",
         # TODO(b/151340488): Timed out on 2020-03-12.
         "nozapfhahn",
+        "no_rocm",
     ],
     deps = [
         ":client_library_test_base",
@@ -1055,6 +1057,7 @@ xla_test(
     shard_count = 20,
     tags = [
         "optonly",
+        "no_rocm",
     ],
     deps = [
         ":client_library_test_base",

From b3fcb02114adb34329aa642319d72b9c46a8306b Mon Sep 17 00:00:00 2001
From: Dragan Mladjenovic <Dragan.Mladjenovic@amd.com>
Date: Tue, 4 Jun 2024 04:01:02 +0000
Subject: [PATCH 477/478] Fix build error in
 third_party/xla/xla/service/memory_space_assignment/cost_analysis.h

---
 .../xla/xla/service/memory_space_assignment/cost_analysis.h      | 1 -
 1 file changed, 1 deletion(-)

diff --git a/third_party/xla/xla/service/memory_space_assignment/cost_analysis.h b/third_party/xla/xla/service/memory_space_assignment/cost_analysis.h
index 4f1e685d150b0c..2955e8d1985adb 100644
--- a/third_party/xla/xla/service/memory_space_assignment/cost_analysis.h
+++ b/third_party/xla/xla/service/memory_space_assignment/cost_analysis.h
@@ -115,7 +115,6 @@ class HloCostAnalysisCosts : public BaseCosts {
   float ComputeSeconds(const HloInstruction& instruction) override;
 
  private:
-  HloCostAnalysisCosts() = default;
 
   const HloCostAnalysis& hlo_cost_analysis_;
 };

From a344018ef114ebbca1abf15ae3fc9c4f4871b9d7 Mon Sep 17 00:00:00 2001
From: Dragan Mladjenovic <Dragan.Mladjenovic@amd.com>
Date: Tue, 4 Jun 2024 15:31:35 +0000
Subject: [PATCH 478/478] Fix xla/service/gpu/gpu_compiler_test.cc build

---
 third_party/xla/xla/service/gpu/BUILD                | 4 ++++
 third_party/xla/xla/service/gpu/gpu_compiler_test.cc | 1 +
 2 files changed, 5 insertions(+)

diff --git a/third_party/xla/xla/service/gpu/BUILD b/third_party/xla/xla/service/gpu/BUILD
index 24af9775bee0d8..dd52523ee63f03 100644
--- a/third_party/xla/xla/service/gpu/BUILD
+++ b/third_party/xla/xla/service/gpu/BUILD
@@ -3625,6 +3625,10 @@ cc_library(
 xla_test(
     name = "gpu_compiler_test",
     srcs = if_gpu_is_configured(["gpu_compiler_test.cc"]),
+    backend_tags = {
+        "gpu_a100": ["no_rocm"],
+        "gpu_v100": ["no_rocm"],
+    },
     backends = ["gpu"],
     data = ["gpu_compiler_test_autotune_db.textproto"],
     deps = [
diff --git a/third_party/xla/xla/service/gpu/gpu_compiler_test.cc b/third_party/xla/xla/service/gpu/gpu_compiler_test.cc
index c5a77f2a007d3c..416a2745bf1f25 100644
--- a/third_party/xla/xla/service/gpu/gpu_compiler_test.cc
+++ b/third_party/xla/xla/service/gpu/gpu_compiler_test.cc
@@ -44,6 +44,7 @@ limitations under the License.
 #include "xla/stream_executor/device_description.h"
 #include "xla/tests/filecheck.h"
 #include "xla/tests/hlo_test_base.h"
+#include "xla/tests/test_macros.h"
 #include "tsl/lib/core/status_test_util.h"
 #include "tsl/platform/casts.h"
 #include "tsl/platform/env.h"