Add the ability to verify assumptions about layout assignment after p…

…ropagation is done. PiperOrigin-RevId: 158696212
TeeeJay · Jun 12, 2017 · 9456bcf · 9456bcf
1 parent b085231
commit 9456bcf
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 0 deletions.
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,