[XLA:Original Value] Propagate OriginalValue information when hoisting loop-invariant instructions.

When `HoistLoopInvariantInstructions` creates a new while loop, the `OriginalValue` for the new while instruction is constructed. Elements from the original while tuple retain their original values. For hoisted instructions, their original values are copied, and the original instruction name is prefixed to indicate they are loop-invariant values hoisted out of the while loop.

Also, introduced `GetOriginalCallInstructionName` and use it in `CallInliner`.

PiperOrigin-RevId: 814772914

Author: Google-ML-Automation

xla/hlo/ir/hlo_original_value.cc

@@ -309,4 +309,24 @@ bool OriginalValue::IsCompatibleWith(const Shape& shape) const {
   return tree().IsStructurallyCompatible(shape);
 }
 
+std::optional<std::string> OriginalValue::GetOriginalCallLikeInstructions()
+    const {
+  if (is_synthetic_call()) {
+    // Synthetic call are transparent and hence resulting in empty call
+    // instructions.
+    return "";
+  }
+  if (IsEmpty()) {
+    // Currently we don't track original call information separately and rely
+    // on the first leaf to find the original call information. So if there are
+    // no leaves we return std::nullopt.
+    return std::nullopt;
+  }
+  auto original_array = original_arrays().begin()->second;
+  if (!original_array.has_value()) {
+    return std::nullopt;
+  }
+  return original_array->instruction_name;
+}
+
 }  // namespace xla

xla/hlo/ir/hlo_original_value.h

@@ -135,6 +135,12 @@ class OriginalValue {
     return !(*this == other);
   }
 
+  // Gets the (partial) call hierarchy string of the original call instructions
+  // that this OriginalValue is associated with. Returns std::nullopt if this
+  // OriginalValue is not associated with a call instruction or the call
+  // hierarchy is lost (e.g., after complicated optimizations).
+  std::optional<std::string> GetOriginalCallLikeInstructions() const;
+
   template <typename H>
   friend H AbslHashValue(H h, const OriginalValue& value) {
     h = H::combine(std::move(h), value.is_synthetic_call());

xla/service/BUILD

@@ -4707,10 +4707,10 @@ xla_cc_test(
         "//xla/hlo/utils:hlo_matchers",
         "//xla/tests:xla_internal_test_main",
         "//xla/tsl/lib/core:status_test_util",
+        "//xla/tsl/platform:statusor",
         "@com_google_absl//absl/algorithm:container",
         "@com_google_absl//absl/status:statusor",
         "@com_google_absl//absl/strings:string_view",
-        "@tsl//tsl/platform:statusor",
     ],
 )
 
@@ -4829,6 +4829,8 @@ cc_library(
         "//xla/hlo/pass:hlo_pass",
         "//xla/hlo/transforms/simplifiers:hlo_dce",
         "//xla/hlo/transforms/simplifiers:tuple_simplifier",
+        "//xla/tsl/platform:errors",
+        "//xla/tsl/platform:statusor",
         "@com_google_absl//absl/algorithm:container",
         "@com_google_absl//absl/container:flat_hash_map",
         "@com_google_absl//absl/container:flat_hash_set",
@@ -4837,8 +4839,6 @@ cc_library(
         "@com_google_absl//absl/log:check",
         "@com_google_absl//absl/status:statusor",
         "@com_google_absl//absl/strings",
-        "@tsl//tsl/platform:errors",
-        "@tsl//tsl/platform:statusor",
     ],
 )
 
@@ -4857,8 +4857,8 @@ xla_cc_test(
         "//xla/hlo/utils:hlo_matchers",
         "//xla/tests:xla_internal_test_main",
         "//xla/tsl/lib/core:status_test_util",
+        "//xla/tsl/platform:statusor",
         "@com_google_absl//absl/log",
-        "@tsl//tsl/platform:statusor",
     ],
 )
 
@@ -4873,13 +4873,16 @@ cc_library(
         "//xla/hlo/analysis:while_loop_analysis",
         "//xla/hlo/ir:hlo",
         "//xla/hlo/pass:hlo_pass",
+        "//xla/tsl/platform:errors",
+        "//xla/tsl/platform:statusor",
         "@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/container:inlined_vector",
         "@com_google_absl//absl/log",
         "@com_google_absl//absl/log:check",
         "@com_google_absl//absl/status:statusor",
+        "@com_google_absl//absl/strings",
         "@com_google_absl//absl/strings:string_view",
     ],
 )
@@ -4895,6 +4898,7 @@ xla_cc_test(
         "//xla/hlo/testlib:hlo_hardware_independent_test_base",
         "//xla/hlo/utils:hlo_matchers",
         "//xla/tests:xla_internal_test_main",
+        "//xla/tsl/platform:statusor",
         "@com_google_googletest//:gtest",
         "@tsl//tsl/platform:statusor",
     ],

xla/service/call_inliner.cc

@@ -214,9 +214,21 @@ class SubcomputationInsertionVisitor : public DfsHloVisitorWithDefault {
       new_hlo_pointer->set_original_value(nullptr);
       return;
     }
+    std::optional<std::string> call_instructions =
+        call_original_value->GetOriginalCallLikeInstructions();
+    if (!call_instructions.has_value()) {
+      // If the call instruction is lost, we must drop the original values
+      // on the inlined instructions because the call hierarchy is lost.
+      new_hlo_pointer->set_original_value(nullptr);
+      return;
+    }
     new_hlo_pointer->CopyOriginalValue(hlo, /*clone=*/true,
                                        /*issue_warning=*/true);
-    if (call_original_value->is_synthetic_call()) {
+    if (call_instructions->empty()) {
+      // Empty call instructions means the call is synthetic and hence the
+      // inlined instruction do not need to be prefixed with the call
+      // instructions. Hence we can just return here to have the copied original
+      // value to be used.
       return;
     }
     std::shared_ptr<OriginalValue> original_value =
@@ -227,12 +239,8 @@ class SubcomputationInsertionVisitor : public DfsHloVisitorWithDefault {
     for (auto& pair : original_value->mutable_original_arrays()) {
       std::optional<OriginalArray>& original_array = pair.second;
       if (original_array.has_value()) {
-        std::string call_instruction_name =
-            call_original_value->original_arrays()
-                .begin()
-                ->second->instruction_name;
         original_array->instruction_name = absl::StrCat(
-            call_instruction_name, "/", original_array->instruction_name);
+            *call_instructions, "/", original_array->instruction_name);
       }
     }
   }

xla/service/while_loop_expensive_invariant_code_motion.cc

@@ -17,7 +17,10 @@ limitations under the License.
 
 #include <cstdint>
 #include <iterator>
+#include <memory>
+#include <optional>
 #include <string>
+#include <utility>
 #include <vector>
 
 #include "absl/algorithm/container.h"
@@ -27,10 +30,20 @@ limitations under the License.
 #include "absl/log/check.h"
 #include "absl/log/log.h"
 #include "absl/status/statusor.h"
+#include "absl/strings/str_cat.h"
 #include "absl/strings/string_view.h"
 #include "xla/hlo/analysis/while_loop_analysis.h"
+#include "xla/hlo/ir/hlo_computation.h"
+#include "xla/hlo/ir/hlo_instruction.h"
+#include "xla/hlo/ir/hlo_opcode.h"
+#include "xla/hlo/ir/hlo_original_value.h"
+#include "xla/hlo/ir/hlo_print_options.h"
+#include "xla/map_util.h"
 #include "xla/service/while_util.h"
+#include "xla/shape.h"
 #include "xla/shape_util.h"
+#include "xla/tsl/platform/errors.h"
+#include "xla/tsl/platform/statusor.h"
 #include "xla/util.h"
 
 namespace xla {
@@ -101,6 +114,32 @@ static void CreateLoopInvariantCopy(
             old_instruction->CloneWithNewOperands(old_instruction->shape(),
                                                   new_operands));
 
+        std::optional<std::string> original_call_instructions;
+        if (while_instr->original_value() != nullptr) {
+          original_call_instructions =
+              while_instr->original_value()->GetOriginalCallLikeInstructions();
+        }
+        if (original_call_instructions.has_value() &&
+            old_instruction->original_value() != nullptr) {
+          std::string original_call_prefix;
+          if (!original_call_instructions->empty()) {
+            // We only add the wildcard iteration count if the call-like
+            // instruction is available.
+            original_call_prefix =
+                absl::StrCat(*original_call_instructions, "#*/");
+          }
+
+          auto new_original_value = std::make_shared<OriginalValue>(
+              *old_instruction->original_value());
+          for (auto& [shape_index, original_array] :
+               new_original_value->mutable_original_arrays()) {
+            if (original_array) {
+              original_array->instruction_name = absl::StrCat(
+                  original_call_prefix, original_array->instruction_name);
+            }
+          }
+          new_instruction->set_original_value(std::move(new_original_value));
+        }
         info.hoisted_copy = new_instruction;
       }
 

xla/service/while_loop_expensive_invariant_code_motion_test.cc

@@ -23,8 +23,8 @@ limitations under the License.
 #include "xla/hlo/parser/hlo_parser.h"
 #include "xla/hlo/testlib/hlo_hardware_independent_test_base.h"
 #include "xla/hlo/utils/hlo_matchers.h"
+#include "xla/tsl/platform/statusor.h"
 #include "xla/util.h"
-#include "tsl/platform/statusor.h"
 
 namespace xla {
 namespace {
@@ -268,5 +268,75 @@ ENTRY entry {
   EXPECT_FALSE(simplified_loop);
 }
 
+TEST_F(WhileLoopExpensiveInvariantCodeMotionTest, HoistWithOriginalValue) {
+  const char* const hlo_string = R"(
+HloModule licm_ov_test
+
+body {
+  p_body = (f32[8,8], f32[16, 8]) parameter(0)
+  b = f32[16, 8] get-tuple-element(p_body), index=1
+  const = f32[] constant(1.0)
+  lhs = f32[8, 16] broadcast(const), dimensions={}, origin={{"lhs.1"}}
+  dot = f32[8,8] dot(lhs, b), lhs_contracting_dims={1}, rhs_contracting_dims={0}, origin={{"dot.1"}}
+  a = f32[8,8] get-tuple-element(p_body), index=0
+  add = f32[8,8] add(a, dot)
+  ROOT root = (f32[8,8], f32[16,8]) tuple(add, b)
+}
+
+condition {
+  p_cond = (f32[8,8], f32[16, 8]) parameter(0)
+  ROOT result = pred[] constant(true)
+}
+
+ENTRY entry {
+  param0 = f32[8,8] parameter(0)
+  param1 = f32[16, 8] parameter(1)
+  while_init = (f32[8,8], f32[16,8]) tuple(param0, param1)
+  ROOT while0 = (f32[8,8], f32[16, 8]) while(while_init), condition=condition, body=body, origin={({"while.5" {0}},{"while.5" {1}})}
+}
+)";
+  TF_ASSERT_OK_AND_ASSIGN(auto m, ParseAndReturnVerifiedModule(hlo_string));
+  HloComputation* body = m->GetComputationWithName("body");
+  HloInstruction* dot = body->GetInstructionWithName("dot");
+  HloInstruction* lhs = body->GetInstructionWithName("lhs");
+
+  TF_ASSERT_OK_AND_ASSIGN(
+      bool simplified_loop,
+      WhileLoopExpensiveInvariantCodeMotion(
+          /*worth_hoisting_individually=*/HloPredicateIsOp<HloOpcode::kDot>)
+          .Run(m.get()));
+  EXPECT_TRUE(simplified_loop);
+
+  HloInstruction* transformed_while = nullptr;
+  for (auto* instr : m->entry_computation()->instructions()) {
+    if (instr->opcode() == HloOpcode::kWhile) {
+      transformed_while = instr;
+      break;
+    }
+  }
+  ASSERT_NE(transformed_while, nullptr);
+
+  HloInstruction* hoisted_dot = nullptr;
+  HloInstruction* hoisted_lhs = nullptr;
+  for (auto* instr : m->entry_computation()->instructions()) {
+    if (instr->opcode() == HloOpcode::kDot && instr->shape() == dot->shape()) {
+      hoisted_dot = instr;
+    }
+    if (instr->opcode() == HloOpcode::kBroadcast &&
+        instr->shape() == lhs->shape()) {
+      hoisted_lhs = instr;
+    }
+  }
+  ASSERT_NE(hoisted_dot, nullptr);
+  ASSERT_NE(hoisted_lhs, nullptr);
+  ASSERT_NE(hoisted_dot->original_value(), nullptr);
+  EXPECT_EQ(hoisted_dot->original_value()->ToString(), "{\"while.5#*/dot.1\"}");
+  ASSERT_NE(hoisted_lhs->original_value(), nullptr);
+  EXPECT_EQ(hoisted_lhs->original_value()->ToString(), "{\"while.5#*/lhs.1\"}");
+  ASSERT_NE(transformed_while->original_value(), nullptr);
+  EXPECT_EQ(transformed_while->original_value()->ToString(),
+            "({\"while.5\" {0}}, {\"while.5\" {1}}, {\"while.5#*/dot.1\"})");
+}
+
 }  // namespace
 }  // namespace xla

xla/service/while_loop_invariant_code_motion.cc

@@ -17,7 +17,10 @@ limitations under the License.
 
 #include <cstdint>
 #include <iterator>
+#include <memory>
+#include <optional>
 #include <string>
+#include <utility>
 #include <vector>
 
 #include "absl/algorithm/container.h"
@@ -26,21 +29,24 @@ limitations under the License.
 #include "absl/container/inlined_vector.h"
 #include "absl/log/check.h"
 #include "absl/log/log.h"
+#include "absl/strings/str_cat.h"
 #include "absl/strings/string_view.h"
 #include "xla/hlo/analysis/while_loop_analysis.h"
 #include "xla/hlo/ir/hlo_computation.h"
 #include "xla/hlo/ir/hlo_instruction.h"
 #include "xla/hlo/ir/hlo_opcode.h"
+#include "xla/hlo/ir/hlo_original_value.h"
+#include "xla/hlo/ir/hlo_print_options.h"
 #include "xla/hlo/transforms/simplifiers/hlo_dce.h"
 #include "xla/hlo/transforms/simplifiers/tuple_simplifier.h"
 #include "xla/map_util.h"
 #include "xla/service/compile_time_cap.h"
 #include "xla/service/while_util.h"
 #include "xla/shape.h"
 #include "xla/shape_util.h"
+#include "xla/tsl/platform/errors.h"
+#include "xla/tsl/platform/statusor.h"
 #include "xla/util.h"
-#include "tsl/platform/errors.h"
-#include "tsl/platform/statusor.h"
 
 namespace xla {
 
@@ -93,6 +99,33 @@ static void CreateLoopInvariantCopy(
           parent_of_while->AddInstruction(old_instruction->CloneWithNewOperands(
               old_instruction->shape(), new_operands));
 
+      std::optional<std::string> original_call_instructions;
+      if (while_instr->original_value() != nullptr) {
+        original_call_instructions =
+            while_instr->original_value()->GetOriginalCallLikeInstructions();
+      }
+      if (original_call_instructions.has_value() &&
+          old_instruction->original_value() != nullptr) {
+        std::string original_call_prefix;
+        if (!original_call_instructions->empty()) {
+          // We only add the wildcard iteration count if the call-like
+          // instruction is available.
+          original_call_prefix =
+              absl::StrCat(*original_call_instructions, "#*/");
+        }
+
+        auto new_original_value =
+            std::make_shared<OriginalValue>(*old_instruction->original_value());
+        for (auto& [shape_index, original_array] :
+             new_original_value->mutable_original_arrays()) {
+          if (original_array) {
+            original_array->instruction_name = absl::StrCat(
+                original_call_prefix, original_array->instruction_name);
+          }
+        }
+        new_instruction->set_original_value(std::move(new_original_value));
+      }
+
       InsertOrDie(hoisted_instructions, old_instruction, new_instruction);
 
       // Approximately half of the instructions that would normally be present