Revert TrivialDce pass

The need for this was due to a bug in MLIR's GreedyPatternRewriteDriver (llvm/llvm-project#86765) which has been fixed (llvm/llvm-project#86990), so the TrivialDce pass is no longer needed.

Reverts 04b31c7

PiperOrigin-RevId: 621959286

Author: Michael Levesque-Dion

tensorflow/lite/kernels/lstm_test.cc

@@ -1304,7 +1304,7 @@ TEST_P(LstmOpTest, Cifg_Peephole_Projection_LayerNorm) {
 
   static const auto* tolerance_per_type =
       new std::map<TensorType, float>{{TensorType_FLOAT32, 0.00001f},
-                                      {TensorType_UINT8, 0.000971057f},
+                                      {TensorType_UINT8, 0.001f},
                                       {TensorType_INT8, 0.001f}};
   VerifyGoldens(&lstm, tolerance_per_type->at(weight_type));
 }

third_party/stablehlo/temporary.patch

@@ -177,7 +177,7 @@ diff --ruN a/stablehlo/stablehlo/CMakeLists.txt b/stablehlo/stablehlo/CMakeLists
 diff --ruN a/stablehlo/stablehlo/experimental/BUILD.bazel b/stablehlo/stablehlo/experimental/BUILD.bazel
 --- stablehlo/stablehlo/experimental/BUILD.bazel
 +++ stablehlo/stablehlo/experimental/BUILD.bazel
-@@ -0,0 +1,116 @@
+@@ -0,0 +1,115 @@
 +# Copyright 2023 The StableHLO Authors. All Rights Reserved.
 +#
 +# Licensed under the Apache License, Version 2.0 (the "License");
@@ -251,7 +251,6 @@ diff --ruN a/stablehlo/stablehlo/experimental/BUILD.bazel b/stablehlo/stablehlo/
 +        "transforms/ChloRecomposeOps.cpp",
 +        "transforms/StablehloCanonicalizeDynamism.cpp",
 +        "transforms/StablehloRefineShapes.cpp",
-+        "transforms/StablehloTrivialDce.cpp",
 +    ],
 +    hdrs = [
 +        "transforms/Passes.h",
@@ -1904,7 +1903,7 @@ diff --ruN a/stablehlo/stablehlo/experimental/tools/StablehloOptMain.cpp b/stabl
 diff --ruN a/stablehlo/stablehlo/experimental/transforms/CMakeLists.txt b/stablehlo/stablehlo/experimental/transforms/CMakeLists.txt
 --- stablehlo/stablehlo/experimental/transforms/CMakeLists.txt
 +++ stablehlo/stablehlo/experimental/transforms/CMakeLists.txt
-@@ -0,0 +1,41 @@
+@@ -0,0 +1,40 @@
 +# Copyright 2023 The StableHLO Authors.
 +#
 +# Licensed under the Apache License, Version 2.0 (the "License");
@@ -1928,7 +1927,6 @@ diff --ruN a/stablehlo/stablehlo/experimental/transforms/CMakeLists.txt b/stable
 +  ChloRecomposeOps.cpp
 +  StablehloCanonicalizeDynamism.cpp
 +  StablehloRefineShapes.cpp
-+  StablehloTrivialDce.cpp
 +
 +  DEPENDS
 +  ExperimentalPassesIncGen
@@ -2173,7 +2171,7 @@ diff --ruN a/stablehlo/stablehlo/experimental/transforms/Passes.h b/stablehlo/st
 diff --ruN a/stablehlo/stablehlo/experimental/transforms/Passes.td b/stablehlo/stablehlo/experimental/transforms/Passes.td
 --- stablehlo/stablehlo/experimental/transforms/Passes.td
 +++ stablehlo/stablehlo/experimental/transforms/Passes.td
-@@ -0,0 +1,55 @@
+@@ -0,0 +1,39 @@
 +/* Copyright 2023 The StableHLO Authors.
 +
 +Licensed under the Apache License, Version 2.0 (the "License");
@@ -2213,22 +2211,6 @@ diff --ruN a/stablehlo/stablehlo/experimental/transforms/Passes.td b/stablehlo/s
 +  }];
 +  let dependentDialects = ["chlo::ChloDialect"];
 +}
-+
-+def StablehloTrivialDcePass : Pass<"experimental-stablehlo-trivial-dce", "ModuleOp"> {
-+  let summary = "(Experimental) Performs a single bottom up pass to remove values that are trivially dead.";
-+  let description = [{
-+    An experimental pass to remove dead values prior to running other passes
-+    that may fail to converge otherwise. For example, running shape refinement
-+    on a program that has a lot of dead values can fail because shape refinement
-+    is top down and removing values causes a new iteration to be triggered, and
-+    removing all the dead values with a top down traversal can take a lot of
-+    iterations (10+), which is slow.
-+
-+    Performing a single pass should be fast, and doing it bottom up means that
-+    values that are transitively dead can be removed since leaf values will be
-+    processed first.
-+  }];
-+}
 diff --ruN a/stablehlo/stablehlo/experimental/transforms/StablehloCanonicalizeDynamism.cpp b/stablehlo/stablehlo/experimental/transforms/StablehloCanonicalizeDynamism.cpp
 --- stablehlo/stablehlo/experimental/transforms/StablehloCanonicalizeDynamism.cpp
 +++ stablehlo/stablehlo/experimental/transforms/StablehloCanonicalizeDynamism.cpp
@@ -2578,73 +2560,6 @@ diff --ruN a/stablehlo/stablehlo/experimental/transforms/StablehloRefineShapes.c
 +}  // namespace experimental
 +}  // namespace stablehlo
 +}  // namespace mlir
-diff --ruN a/stablehlo/stablehlo/experimental/transforms/StablehloTrivialDce.cpp b/stablehlo/stablehlo/experimental/transforms/StablehloTrivialDce.cpp
---- stablehlo/stablehlo/experimental/transforms/StablehloTrivialDce.cpp
-+++ stablehlo/stablehlo/experimental/transforms/StablehloTrivialDce.cpp
-@@ -0,0 +1,63 @@
-+/* Copyright 2022 The StableHLO 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 <cstdint>
-+
-+#include "llvm/ADT/SmallVector.h"
-+#include "mlir/Dialect/Func/IR/FuncOps.h"
-+#include "mlir/IR/PatternMatch.h"
-+#include "mlir/Support/LogicalResult.h"
-+#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
-+#include "stablehlo/experimental/transforms/Passes.h"
-+
-+namespace mlir {
-+namespace stablehlo {
-+namespace experimental {
-+
-+#define GEN_PASS_DEF_STABLEHLOTRIVIALDCEPASS
-+#include "stablehlo/experimental/transforms/Passes.h.inc"
-+
-+namespace {
-+
-+struct StablehloTrivialDcePass
-+    : public impl::StablehloTrivialDcePassBase<StablehloTrivialDcePass> {
-+  using StablehloTrivialDcePassBase::StablehloTrivialDcePassBase;
-+
-+  void runOnOperation() override {
-+    GreedyRewriteConfig config;
-+
-+    // Hardcode defaults for stability.
-+    config.enableRegionSimplification = true;
-+    config.maxNumRewrites = GreedyRewriteConfig::kNoLimit;
-+    config.strictMode = GreedyRewriteStrictness::AnyOp;
-+
-+    // Run a single bottom up pass.
-+    config.useTopDownTraversal = false;
-+    config.maxIterations = 1;
-+
-+    // Running a greedy rewrite will cause trivially dead values to be removed.
-+    // Doing it without patterns ensures that no other changes are made to the
-+    // IR. Doing it bottom-up ensures that values that are transitively dead are
-+    // also removed. Although 1 pass should be enough,
-+    // applyPatternsAndFoldGreedily will want to run at least 1 more iteration
-+    // to confirm convergence, but we don't need to check for convergence, so we
-+    // ignore the return value.
-+    (void)applyPatternsAndFoldGreedily(getOperation(), RewritePatternSet(&getContext()), config);
-+  }
-+};
-+
-+}  // namespace
-+}  // namespace experimental
-+}  // namespace stablehlo
-+}  // namespace mlir
 diff --ruN a/stablehlo/stablehlo/tests/chlo/chlo_legalize_to_stablehlo.mlir b/stablehlo/stablehlo/tests/chlo/chlo_legalize_to_stablehlo.mlir
 --- stablehlo/stablehlo/tests/chlo/chlo_legalize_to_stablehlo.mlir
 +++ stablehlo/stablehlo/tests/chlo/chlo_legalize_to_stablehlo.mlir

third_party/xla/third_party/stablehlo/temporary.patch

@@ -177,7 +177,7 @@ diff --ruN a/stablehlo/stablehlo/CMakeLists.txt b/stablehlo/stablehlo/CMakeLists
 diff --ruN a/stablehlo/stablehlo/experimental/BUILD.bazel b/stablehlo/stablehlo/experimental/BUILD.bazel
 --- stablehlo/stablehlo/experimental/BUILD.bazel
 +++ stablehlo/stablehlo/experimental/BUILD.bazel
-@@ -0,0 +1,116 @@
+@@ -0,0 +1,115 @@
 +# Copyright 2023 The StableHLO Authors. All Rights Reserved.
 +#
 +# Licensed under the Apache License, Version 2.0 (the "License");
@@ -251,7 +251,6 @@ diff --ruN a/stablehlo/stablehlo/experimental/BUILD.bazel b/stablehlo/stablehlo/
 +        "transforms/ChloRecomposeOps.cpp",
 +        "transforms/StablehloCanonicalizeDynamism.cpp",
 +        "transforms/StablehloRefineShapes.cpp",
-+        "transforms/StablehloTrivialDce.cpp",
 +    ],
 +    hdrs = [
 +        "transforms/Passes.h",
@@ -1904,7 +1903,7 @@ diff --ruN a/stablehlo/stablehlo/experimental/tools/StablehloOptMain.cpp b/stabl
 diff --ruN a/stablehlo/stablehlo/experimental/transforms/CMakeLists.txt b/stablehlo/stablehlo/experimental/transforms/CMakeLists.txt
 --- stablehlo/stablehlo/experimental/transforms/CMakeLists.txt
 +++ stablehlo/stablehlo/experimental/transforms/CMakeLists.txt
-@@ -0,0 +1,41 @@
+@@ -0,0 +1,40 @@
 +# Copyright 2023 The StableHLO Authors.
 +#
 +# Licensed under the Apache License, Version 2.0 (the "License");
@@ -1928,7 +1927,6 @@ diff --ruN a/stablehlo/stablehlo/experimental/transforms/CMakeLists.txt b/stable
 +  ChloRecomposeOps.cpp
 +  StablehloCanonicalizeDynamism.cpp
 +  StablehloRefineShapes.cpp
-+  StablehloTrivialDce.cpp
 +
 +  DEPENDS
 +  ExperimentalPassesIncGen
@@ -2173,7 +2171,7 @@ diff --ruN a/stablehlo/stablehlo/experimental/transforms/Passes.h b/stablehlo/st
 diff --ruN a/stablehlo/stablehlo/experimental/transforms/Passes.td b/stablehlo/stablehlo/experimental/transforms/Passes.td
 --- stablehlo/stablehlo/experimental/transforms/Passes.td
 +++ stablehlo/stablehlo/experimental/transforms/Passes.td
-@@ -0,0 +1,55 @@
+@@ -0,0 +1,39 @@
 +/* Copyright 2023 The StableHLO Authors.
 +
 +Licensed under the Apache License, Version 2.0 (the "License");
@@ -2213,22 +2211,6 @@ diff --ruN a/stablehlo/stablehlo/experimental/transforms/Passes.td b/stablehlo/s
 +  }];
 +  let dependentDialects = ["chlo::ChloDialect"];
 +}
-+
-+def StablehloTrivialDcePass : Pass<"experimental-stablehlo-trivial-dce", "ModuleOp"> {
-+  let summary = "(Experimental) Performs a single bottom up pass to remove values that are trivially dead.";
-+  let description = [{
-+    An experimental pass to remove dead values prior to running other passes
-+    that may fail to converge otherwise. For example, running shape refinement
-+    on a program that has a lot of dead values can fail because shape refinement
-+    is top down and removing values causes a new iteration to be triggered, and
-+    removing all the dead values with a top down traversal can take a lot of
-+    iterations (10+), which is slow.
-+
-+    Performing a single pass should be fast, and doing it bottom up means that
-+    values that are transitively dead can be removed since leaf values will be
-+    processed first.
-+  }];
-+}
 diff --ruN a/stablehlo/stablehlo/experimental/transforms/StablehloCanonicalizeDynamism.cpp b/stablehlo/stablehlo/experimental/transforms/StablehloCanonicalizeDynamism.cpp
 --- stablehlo/stablehlo/experimental/transforms/StablehloCanonicalizeDynamism.cpp
 +++ stablehlo/stablehlo/experimental/transforms/StablehloCanonicalizeDynamism.cpp
@@ -2578,73 +2560,6 @@ diff --ruN a/stablehlo/stablehlo/experimental/transforms/StablehloRefineShapes.c
 +}  // namespace experimental
 +}  // namespace stablehlo
 +}  // namespace mlir
-diff --ruN a/stablehlo/stablehlo/experimental/transforms/StablehloTrivialDce.cpp b/stablehlo/stablehlo/experimental/transforms/StablehloTrivialDce.cpp
---- stablehlo/stablehlo/experimental/transforms/StablehloTrivialDce.cpp
-+++ stablehlo/stablehlo/experimental/transforms/StablehloTrivialDce.cpp
-@@ -0,0 +1,63 @@
-+/* Copyright 2022 The StableHLO 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 <cstdint>
-+
-+#include "llvm/ADT/SmallVector.h"
-+#include "mlir/Dialect/Func/IR/FuncOps.h"
-+#include "mlir/IR/PatternMatch.h"
-+#include "mlir/Support/LogicalResult.h"
-+#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
-+#include "stablehlo/experimental/transforms/Passes.h"
-+
-+namespace mlir {
-+namespace stablehlo {
-+namespace experimental {
-+
-+#define GEN_PASS_DEF_STABLEHLOTRIVIALDCEPASS
-+#include "stablehlo/experimental/transforms/Passes.h.inc"
-+
-+namespace {
-+
-+struct StablehloTrivialDcePass
-+    : public impl::StablehloTrivialDcePassBase<StablehloTrivialDcePass> {
-+  using StablehloTrivialDcePassBase::StablehloTrivialDcePassBase;
-+
-+  void runOnOperation() override {
-+    GreedyRewriteConfig config;
-+
-+    // Hardcode defaults for stability.
-+    config.enableRegionSimplification = true;
-+    config.maxNumRewrites = GreedyRewriteConfig::kNoLimit;
-+    config.strictMode = GreedyRewriteStrictness::AnyOp;
-+
-+    // Run a single bottom up pass.
-+    config.useTopDownTraversal = false;
-+    config.maxIterations = 1;
-+
-+    // Running a greedy rewrite will cause trivially dead values to be removed.
-+    // Doing it without patterns ensures that no other changes are made to the
-+    // IR. Doing it bottom-up ensures that values that are transitively dead are
-+    // also removed. Although 1 pass should be enough,
-+    // applyPatternsAndFoldGreedily will want to run at least 1 more iteration
-+    // to confirm convergence, but we don't need to check for convergence, so we
-+    // ignore the return value.
-+    (void)applyPatternsAndFoldGreedily(getOperation(), RewritePatternSet(&getContext()), config);
-+  }
-+};
-+
-+}  // namespace
-+}  // namespace experimental
-+}  // namespace stablehlo
-+}  // namespace mlir
 diff --ruN a/stablehlo/stablehlo/tests/chlo/chlo_legalize_to_stablehlo.mlir b/stablehlo/stablehlo/tests/chlo/chlo_legalize_to_stablehlo.mlir
 --- stablehlo/stablehlo/tests/chlo/chlo_legalize_to_stablehlo.mlir
 +++ stablehlo/stablehlo/tests/chlo/chlo_legalize_to_stablehlo.mlir

third_party/xla/xla/pjrt/distributed/client.cc

@@ -48,17 +48,17 @@ class DistributedRuntimeCoordinationServiceClient
       : DistributedRuntimeCoordinationServiceClient(channel, Options()) {}
   ~DistributedRuntimeCoordinationServiceClient() override;
 
-  xla::Status Connect() override;
-  xla::Status Shutdown() override;
+  absl::Status Connect() override;
+  absl::Status Shutdown() override;
   absl::StatusOr<std::string> BlockingKeyValueGet(
       std::string_view key, absl::Duration timeout) override;
   absl::StatusOr<std::vector<std::pair<std::string, std::string>>>
   KeyValueDirGet(std::string_view key) override;
-  xla::Status KeyValueSet(std::string_view key,
-                          std::string_view value) override;
-  xla::Status KeyValueDelete(std::string_view key) override;
-  xla::Status WaitAtBarrier(std::string barrier_id,
-                            absl::Duration timeout) override;
+  absl::Status KeyValueSet(std::string_view key,
+                           std::string_view value) override;
+  absl::Status KeyValueDelete(std::string_view key) override;
+  absl::Status WaitAtBarrier(std::string barrier_id,
+                             absl::Duration timeout) override;
   absl::StatusOr<tsl::CoordinationServiceAgent*> GetCoordinationServiceAgent()
       override;
 
@@ -107,7 +107,7 @@ DistributedRuntimeCoordinationServiceClient::
 DistributedRuntimeCoordinationServiceClient::
     ~DistributedRuntimeCoordinationServiceClient() = default;
 
-xla::Status DistributedRuntimeCoordinationServiceClient::Connect() {
+absl::Status DistributedRuntimeCoordinationServiceClient::Connect() {
   const absl::Time deadline =
       absl::Now() +
       absl::Milliseconds(config_.cluster_register_timeout_in_ms());
@@ -130,7 +130,7 @@ xla::Status DistributedRuntimeCoordinationServiceClient::Connect() {
   return s;
 }
 
-xla::Status DistributedRuntimeCoordinationServiceClient::Shutdown() {
+absl::Status DistributedRuntimeCoordinationServiceClient::Shutdown() {
   LOG(INFO) << "Distributed task shutdown initiated.";
   Status s = coord_agent_->Shutdown();
   LOG(INFO) << "Distributed task shutdown result: " << s;
@@ -162,17 +162,17 @@ DistributedRuntimeCoordinationServiceClient::KeyValueDirGet(
   return kvs;
 }
 
-xla::Status DistributedRuntimeCoordinationServiceClient::KeyValueDelete(
+absl::Status DistributedRuntimeCoordinationServiceClient::KeyValueDelete(
     std::string_view key) {
   return coord_agent_->DeleteKeyValue(key);
 }
 
-xla::Status DistributedRuntimeCoordinationServiceClient::KeyValueSet(
+absl::Status DistributedRuntimeCoordinationServiceClient::KeyValueSet(
     std::string_view key, std::string_view value) {
   return coord_agent_->InsertKeyValue(key, value);
 }
 
-xla::Status DistributedRuntimeCoordinationServiceClient::WaitAtBarrier(
+absl::Status DistributedRuntimeCoordinationServiceClient::WaitAtBarrier(
     std::string barrier_id, absl::Duration timeout) {
   return coord_agent_->WaitAtBarrier(barrier_id, timeout, /*tasks=*/{});
 }