From f23efde5109d0f387ecdc535deaf9fb45de5d391 Mon Sep 17 00:00:00 2001 From: Yunxing Dai Date: Fri, 30 Nov 2018 11:13:46 -0800 Subject: [PATCH] Keep the old instruction alive when replacing instructions in GetDimensionSizeRewriter. Otherwise it could create use-after-free if we delete and iterate through instructions at the same time. PiperOrigin-RevId: 223544664 --- .../xla/service/hlo_get_dimension_size_rewriter.cc | 9 ++------- 1 file changed, 2 insertions(+), 7 deletions(-) diff --git a/tensorflow/compiler/xla/service/hlo_get_dimension_size_rewriter.cc b/tensorflow/compiler/xla/service/hlo_get_dimension_size_rewriter.cc index 631b3ad735f..c919dbd82d3 100644 --- a/tensorflow/compiler/xla/service/hlo_get_dimension_size_rewriter.cc +++ b/tensorflow/compiler/xla/service/hlo_get_dimension_size_rewriter.cc @@ -39,7 +39,7 @@ StatusOr ReplaceGetSize(HloInstruction* instr) { uint32 size = instr->operand(0)->shape().dimensions(instr->dimension()); HloInstruction* new_instr = computation->AddInstruction( HloInstruction::CreateConstant(LiteralUtil::CreateR0(size))); - TF_RETURN_IF_ERROR(computation->ReplaceInstruction(instr, new_instr)); + TF_RETURN_IF_ERROR(instr->ReplaceAllUsesWith(new_instr)); return true; } @@ -50,12 +50,7 @@ StatusOr HloGetDimensionSizeRewriter::Run(HloModule* module) { HloProto proto; *proto.mutable_hlo_module() = module->ToProto(); for (auto* computation : module->computations()) { - // Replacing instructions will change the instruction list in the - // computation. So instead of iterating computation->instructions() - // directly, we make a copy of the list to avoid use-after-free. - std::vector instrs(computation->instruction_count()); - absl::c_copy(computation->instructions(), instrs.begin()); - for (auto instruction : instrs) { + for (auto instruction : computation->instructions()) { TF_ASSIGN_OR_RETURN(bool replaced, ReplaceGetSize(instruction)); changed = changed || replaced; }