[XLA:GPU] move TransposeFolding after simplifier pipeline

Combination of DotDecompose, AlgebraicSimplifier, and TransposeFolding might never reach a fixed point and stuck in a rewrite loop.

Moving TransposeFolding in simplification-2 pipeline should result in a similar output as we keep the relative order.

Also:

- Added a few tests to detect / record cases when we run into rewrite loop

- Log warning when HloPassFix reaches iteraction limit as this is likely a similar bug

PiperOrigin-RevId: 715279302

Author: metaflow

xla/hlo/pass/hlo_pass_fix.h

@@ -19,6 +19,7 @@ limitations under the License.
 #include <algorithm>
 #include <type_traits>
 
+#include "absl/log/log.h"
 #include "absl/status/statusor.h"
 #include "xla/hlo/ir/hlo_module.h"
 #include "xla/hlo/ir/hlo_module_group.h"
@@ -76,8 +77,8 @@ class HloPassFix : public Pass {
       VLOG(3) << "changed_this_iteration: " << changed_this_iteration;
       ++iteration_count;
       if (iteration_count == kIterationLimit) {
-        VLOG(1) << "Unexpectedly high number of iterations in HLO passes, "
-                   "exiting fixed point loop.";
+        LOG(WARNING) << "Unexpectedly high number of iterations in HLO passes, "
+                        "exiting fixed point loop.";
         // Return false in case this is fixed point is nested.
         return false;
       }
@@ -98,9 +99,9 @@ class HloPassFix : public Pass {
               << !run_state->changed_last_iteration.empty();
       run_state->IncrementIteration();
       if (run_state->iteration == kIterationLimit) {
-        VLOG(1) << "Unexpectedly high number of iterations in HLO passes '"
-                << Pass::name() << "' for module '" << module->name()
-                << "'. Exiting fixed point loop.";
+        LOG(WARNING) << "Unexpectedly high number of iterations in HLO passes '"
+                     << Pass::name() << "' for module '" << module->name()
+                     << "'. Exiting fixed point loop.";
         // Clear changed and abort in case this is fixed point is nested.
         run_state->changed.clear();
         break;

xla/service/gpu/gpu_compiler.cc

@@ -821,7 +821,6 @@ absl::Status RunOptimizationPasses(
     pipeline.AddPass<HloConstantFolding>();
     pipeline.AddPass<ConditionalSimplifier>();
     pipeline.AddPass<RealImagExpander>();
-    pipeline.AddPass<TransposeFolding>(CanFoldTransposeOperandIntoDot);
     pipeline.AddPass<HloCSE>(/*is_layout_sensitive=*/false);
     pipeline.AddPass<HloDCE>();
   }();
@@ -835,6 +834,7 @@ absl::Status RunOptimizationPasses(
     pipeline.AddPass<ConvertMover>();
     pipeline.AddPass<GpuAlgebraicSimplifier>(layout_insensitive_algsimp_opts,
                                              gpu_version);
+    pipeline.AddPass<TransposeFolding>(CanFoldTransposeOperandIntoDot);
   }();
 
   pipeline.AddPass<HloComputationDeduplicator>(

xla/service/gpu/gpu_compiler_test.cc

@@ -1568,6 +1568,62 @@ TEST_F(PassOrderTest,
                   "remove-no-op-reduce-precision-algebraic-simplifier");
 }
 
+// Tests that passes are converging and pipelines reach a fix point.
+class FixPointTest : public HloTestBase {
+ public:
+  void ExpectPipelinesReachFixedPoint(absl::string_view module_text) {
+    std::unique_ptr<HloModule> optimized_module;
+    TF_ASSERT_OK_AND_ASSIGN(
+        std::unique_ptr<VerifiedHloModule> module,
+        ParseAndReturnVerifiedModule(module_text, GetModuleConfigForTest()));
+    TF_ASSERT_OK_AND_ASSIGN(optimized_module,
+                            GetOptimizedModule(std::move(module)));
+
+    std::string last_pipeline_name;
+    int count = 0;
+    for (const HloPassMetadata& pass_metadata :
+         optimized_module->metadata()->proto().pass_metadata()) {
+      if (pass_metadata.pass_name() != "pipeline-start") {
+        continue;
+      }
+      VLOG(2) << "pipeline: " << pass_metadata.pipeline_name();
+      if (pass_metadata.pipeline_name() != last_pipeline_name) {
+        count = 0;
+        last_pipeline_name = pass_metadata.pipeline_name();
+      }
+      count++;
+      // 25 is a default iteration limit of HloPassFix.
+      EXPECT_LT(count, 25) << "Pipeline '" << pass_metadata.pipeline_name()
+                           << "' ran " << count
+                           << " times. That is likely an indication that the "
+                              "pipeline is not reaching a fixed point.";
+    }
+  }
+};
+
+TEST_F(FixPointTest, Constant) {
+  ExpectPipelinesReachFixedPoint(R"(ENTRY main {
+  ROOT constant = f32[] constant(0)
+})");
+}
+
+TEST_F(FixPointTest, ReshapeTranspose) {
+  ExpectPipelinesReachFixedPoint(R"(ENTRY main {
+p0 = f32[1024,4096]{1,0} parameter(0)
+reshape = f32[1024,1024,4]{2,1,0} reshape(p0)
+ROOT transpose = f32[4,1024,1024]{2,1,0} transpose(reshape), dimensions={2,1,0}
+})");
+}
+
+TEST_F(FixPointTest, DotWithBatchDims) {
+  // Reduced test case for b/383729716.
+  ExpectPipelinesReachFixedPoint(R"(ENTRY main {
+p0 = f32[8,4,64]{2,1,0} parameter(0)
+p1 = f32[4,64,1024] parameter(1)
+ROOT dot = f32[4,8,1024]{2,1,0} dot(p0, p1), lhs_batch_dims={1}, lhs_contracting_dims={2}, rhs_batch_dims={0}, rhs_contracting_dims={1}
+})");
+}
+
 }  // namespace
 }  // namespace gpu
 }  // namespace xla