From b678b55e332af661e5a3b58880399ea3b81679cd Mon Sep 17 00:00:00 2001
From: Trent Lo <trentl@nvidia.com>
Date: Fri, 2 Oct 2020 14:52:36 -0700
Subject: [PATCH] [XLA/GPU] clear operands for removed HLOs.

This avoids possible accesses to the Null operands.
---
 tensorflow/compiler/xla/service/hlo_computation.cc | 2 ++
 tensorflow/compiler/xla/service/hlo_instruction.h  | 4 ++++
 2 files changed, 6 insertions(+)

diff --git a/tensorflow/compiler/xla/service/hlo_computation.cc b/tensorflow/compiler/xla/service/hlo_computation.cc
index 38c4da1f182..5d695b9c20f 100644
--- a/tensorflow/compiler/xla/service/hlo_computation.cc
+++ b/tensorflow/compiler/xla/service/hlo_computation.cc
@@ -315,6 +315,8 @@ Status HloComputation::RemoveInstructionImpl(HloInstruction* instruction,
   (*inst_it->second)->set_parent(nullptr);
   to_be_deleted_.emplace_back(inst_it->second->release());
   to_be_deleted_.back()->DetachFromOperandsAndUsers();
+  // Clear all operands to avoid Null operands.
+  to_be_deleted_.back()->RemoveAllOperands();
   to_be_deleted_.back()->MarkAsDead();
   instructions_.erase(inst_it->second);
   instruction_iterators_.erase(inst_it);
diff --git a/tensorflow/compiler/xla/service/hlo_instruction.h b/tensorflow/compiler/xla/service/hlo_instruction.h
index e21ae719e4d..9f4f7f048ad 100644
--- a/tensorflow/compiler/xla/service/hlo_instruction.h
+++ b/tensorflow/compiler/xla/service/hlo_instruction.h
@@ -1911,6 +1911,10 @@ class HloInstruction {
   // by factory methods.
   HloInstruction(HloOpcode opcode, const Shape& shape);
 
+  void RemoveAllOperands() {
+    operands_.clear();
+  }
+
   void RemoveOperandAt(int index) {
     operands_.erase(operands_.begin() + index);
   }