[XLA] Check buffer definition before use
PiperOrigin-RevId: 355338830
Change-Id: I6a0eef763abdadddbb5a53f7695473f1aa2a1a58
diff --git a/tensorflow/compiler/xla/service/layout_assignment.cc b/tensorflow/compiler/xla/service/layout_assignment.cc
index 55569cf..2885005 100644
--- a/tensorflow/compiler/xla/service/layout_assignment.cc
+++ b/tensorflow/compiler/xla/service/layout_assignment.cc
@@ -1419,6 +1419,10 @@
if (subshape.rank() != operand->shape().rank()) {
return Status::OK();
}
+ if (!constraints->points_to_analysis()
+ .InstructionDefinesBufferAtIndex(user, shape_index)) {
+ return Status::OK();
+ }
// TODO(b/67641796): Are there cases except fusion that use this code
// path?
TF_ASSIGN_OR_RETURN(
@@ -1450,6 +1454,10 @@
if (subshape.rank() <= 1) {
return Status::OK();
}
+ if (!constraints->points_to_analysis().InstructionDefinesBufferAtIndex(
+ user, shape_index)) {
+ return Status::OK();
+ }
TF_ASSIGN_OR_RETURN(
const LogicalBuffer* buffer,
constraints->points_to_analysis().GetBufferDefinedAt(user,
