[XLA][Numerics][HLO Original Value] Try to infer original value of GetTupleElement operations

If the tuple operand of an GetTupleElement instruction is a Tuple operation and the instruction does not have an original value yet, returns the original value of the corresponding tuple element instead.

PiperOrigin-RevId: 827373266

Author: jcai19

xla/hlo/ir/hlo_instruction.cc

@@ -6043,7 +6043,13 @@ void HloInstruction::set_output_to_operand_aliasing(
 }
 
 std::shared_ptr<OriginalValue> HloInstruction::original_value() const {
-  return original_value_;
+  if (original_value_ != nullptr || opcode_ != HloOpcode::kGetTupleElement) {
+    return original_value_;
+  }
+  const HloInstruction* tuple = operand(0);
+  return tuple->opcode() == HloOpcode::kTuple
+             ? tuple->operand(tuple_index())->original_value()
+             : nullptr;
 }
 
 void HloInstruction::set_original_value(

xla/hlo/ir/hlo_original_value_test.cc

@@ -435,5 +435,24 @@ ENTRY main {
   EXPECT_EQ(p0->original_value(), p1->original_value());
 }
 
+TEST_F(OriginalValueHloTest, InferGetTupleElementOriginalValue) {
+  const char* hlo_string = R"(
+HloModule test
+
+ENTRY main {
+  p0 = f32[] parameter(0), origin={{"p0"}}
+  p1 = f32[] parameter(1)
+  tuple = (f32[], f32[]) tuple(p0, p1)
+  ROOT gte = f32[] get-tuple-element(tuple), index=0
+}
+)";
+  TF_ASSERT_OK_AND_ASSIGN(auto module,
+                          ParseAndReturnVerifiedModule(hlo_string));
+  const HloInstruction* gte = module->entry_computation()->root_instruction();
+
+  EXPECT_NE(gte->original_value(), nullptr);
+  EXPECT_EQ(gte->original_value()->ToString(), R"({"p0"})");
+}
+
 }  // namespace
 }  // namespace xla