Add the ability to verify assumptions about layout assignment after propagation

tensorflower-gardener · tensorflower-gardener · commit 9456bcf1765a · 2017-06-12T03:47:09.000-07:00
is done.

PiperOrigin-RevId: 158696212
diff --git a/tensorflow/compiler/xla/service/layout_assignment.cc b/tensorflow/compiler/xla/service/layout_assignment.cc
@@ -1265,6 +1265,9 @@ Status LayoutAssignment::AssignLayouts(const LayoutConstraints& constraints,
       TF_RETURN_IF_ERROR(SetFusionLayouts(instruction));
     }
 
+    // Execute extra verification step once the layout has been finalized.
+    TF_RETURN_IF_ERROR(Verify(instruction));
+
     // Verify all layouts in the shape have been set.
     TF_RET_CHECK(LayoutUtil::HasLayout(instruction->shape()));
   }
diff --git a/tensorflow/compiler/xla/service/layout_assignment.h b/tensorflow/compiler/xla/service/layout_assignment.h
@@ -246,6 +246,12 @@ class LayoutAssignment : public HloPassInterface {
       const ResultLayoutConstraint& layout_constraint,
       LayoutConstraints* constraints);
 
+  // Called after layouts of an instruction have been finalized to allow
+  // subclasses to check for platform specific assumptions.
+  virtual Status Verify(const HloInstruction* instruction) {
+    return Status::OK();
+  }
+
   // Propagates a buffer layout constraint into the operands that use it.
   Status PropagateBufferConstraintToUses(
       const BufferLayoutConstraint& layout_constraint,

Original file line number	Diff line number	Diff line change
`@@ -1265,6 +1265,9 @@ Status LayoutAssignment::AssignLayouts(const LayoutConstraints& constraints,`
`1265`	`1265`	`TF_RETURN_IF_ERROR(SetFusionLayouts(instruction));`
`1266`	`1266`	`}`
`1267`	`1267`
	`1268`	`+ // Execute extra verification step once the layout has been finalized.`
	`1269`	`+ TF_RETURN_IF_ERROR(Verify(instruction));`
	`1270`	`+`
`1268`	`1271`	`// Verify all layouts in the shape have been set.`
`1269`	`1272`	`TF_RET_CHECK(LayoutUtil::HasLayout(instruction->shape()));`
`1270`	`1273`	`}`