[XLA:TPU] HostOffloader: Refactor HandleInputStreaming to occur befor…

…e tracing the graph for custom calls.

Currently, input streaming gets handled in arbitrary order, after we have already possibly searched the graph for the corresponding streamed custom calls. This refactors the input streaming to occur before any graph searching, which makes more sense and prepares for output streaming.

PiperOrigin-RevId: 629208199

third_party/xla/xla/service/host_offloader.cc

@@ -761,12 +761,12 @@ absl::StatusOr<bool> HostOffloader::Run(
   // Run HloAliasAnalysis on module.
   TF_ASSIGN_OR_RETURN(alias_analysis_, HloAliasAnalysis::Run(module));
 
+  // Handle streamed parameters first.
+  TF_RETURN_IF_ERROR(HandleInputStreaming(module->entry_computation()));
+
   // Iterate over all instructions and look for XLA host offload annotations.
   for (HloComputation* computation :
        module->MakeNonfusionComputations(execution_threads)) {
-    if (computation->IsEntryComputation()) {
-      TF_RETURN_IF_ERROR(HandleInputStreaming(computation));
-    }
     for (HloInstruction* instruction :
          computation->MakeInstructionPostOrder()) {
       if (instruction->opcode() != HloOpcode::kCustomCall) {