From a610534ef1766f8fcd80a4bb6ddef4918f09bd30 Mon Sep 17 00:00:00 2001 From: Adrian Kuegel Date: Wed, 13 Nov 2024 00:39:23 -0800 Subject: [PATCH] Avoid unnecessary overhead in HloExtractor (NFC) Currently we iterate over all instructions in the old module to find the instructions which have been cloned. But CloneContext offers direct access to them, so we should be using that instead. PiperOrigin-RevId: 696017622 --- xla/tools/hlo_extractor.cc | 12 +++++------- 1 file changed, 5 insertions(+), 7 deletions(-) diff --git a/xla/tools/hlo_extractor.cc b/xla/tools/hlo_extractor.cc index f16dbfa3ed004..571549372d22b 100644 --- a/xla/tools/hlo_extractor.cc +++ b/xla/tools/hlo_extractor.cc @@ -168,13 +168,11 @@ class ExtractionVisitor : public ConstDfsHloVisitorWithDefault { // Rename HLOs so that their name matches the original. By default, // HLOs get new unique names when adding a new entry computation to // a module. - for (auto computation : old_module_->MakeComputationPostOrder()) { - for (auto old_instruction : computation->MakeInstructionPostOrder()) { - if (auto new_instruction = - clone_context_.FindInstruction(old_instruction)) { - new_instruction->SetAndSanitizeName(old_instruction->name()); - } - } + for (const auto& instruction_mapping : + clone_context_.cloned_instructions()) { + auto old_instruction = instruction_mapping.first; + auto new_instruction = instruction_mapping.second; + new_instruction->SetAndSanitizeName(old_instruction->name()); } // For the extra created instructions (e.g., the ones created when replacing // with broadcasted zeros), we make sure they have unique names without