[XLA] Check buffer definition before use

PiperOrigin-RevId: 355338830
Change-Id: I6a0eef763abdadddbb5a53f7695473f1aa2a1a58
This commit is contained in:
Yuanzhong Xu 2021-02-03 00:20:06 -08:00 committed by TensorFlower Gardener
parent bd14bb7cf4
commit 238f47d671

View File

@ -1419,6 +1419,10 @@ Status LayoutAssignment::PropagateOperandConstraint(
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 @@ Status LayoutAssignment::PropagateOperandConstraint(
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,