[XLA] Remove unused InstructionRequiresInputLayoutEqualToOutputLayout hook

PiperOrigin-RevId: 156843392

2 files changed, 16 insertions, 26 deletions

diff --git a/tensorflow/compiler/xla/service/layout_assignment.cc b/tensorflow/compiler/xla/service/layout_assignment.cc
index a8366ae794..82417253ab 100644
--- a/tensorflow/compiler/xla/service/layout_assignment.cc
+++ b/tensorflow/compiler/xla/service/layout_assignment.cc
@@ -711,8 +711,7 @@ std::unique_ptr<Layout> LayoutAssignment::ChooseOperandLayoutFromOutputLayout(
   CHECK(ShapeUtil::IsArray(instruction->shape()) &&
         ShapeUtil::IsArray(operand->shape()));
 
-  if ((instruction->IsElementwiseOnOperand(operand_no) ||
-       InstructionRequiresInputLayoutEqualToOutputLayout(instruction)) &&
+  if (instruction->IsElementwiseOnOperand(operand_no) &&
       !ShapeUtil::IsScalar(operand->shape()) &&
       ShapeUtil::Rank(operand->shape()) ==
           ShapeUtil::Rank(instruction->shape())) {
diff --git a/tensorflow/compiler/xla/service/layout_assignment.h b/tensorflow/compiler/xla/service/layout_assignment.h
index ae53702a28..cd62d6b007 100644
--- a/tensorflow/compiler/xla/service/layout_assignment.h
+++ b/tensorflow/compiler/xla/service/layout_assignment.h
@@ -258,6 +258,21 @@ class LayoutAssignment : public HloPassInterface {
                                       const HloInstruction* instruction,
                                       LayoutConstraints* constraints);
 
+  // Chooses a layout of operand `operand_no` of `instruction` that minimizes
+  // the cost of `instruction`. `output_layout` is the layout of `instruction`.
+  // Returns null if it can't decide the best layout.
+  // Precondition: `instruction` and the operand are array-shaped.
+  std::unique_ptr<Layout> ChooseOperandLayoutFromOutputLayout(
+      const Layout& output_layout, const HloInstruction* instruction,
+      int64 operand_no);
+  // Given the layout of `user`'s `operand_no`-th operand, chooses a layout of
+  // `user` that minimizes its cost on that operand.  Returns null if it can't
+  // decide the best layout.
+  // Precondition: `user` and the operand are array-shaped.
+  std::unique_ptr<Layout> ChooseOutputLayoutFromOperandLayout(
+      const Layout& operand_layout, const HloInstruction* user,
+      int64 operand_no);
+
  private:
   // Adds constraints which must be satisfied for correctness on all
   // backends. Called once prior to propagating constraints.
@@ -273,15 +288,6 @@ class LayoutAssignment : public HloPassInterface {
     return Status::OK();
   }
 
-  // This method can be overriden to mark instructions as requiring the operands
-  // to have the same layout as the result, for performance or correctness. This
-  // will propagate constraints through the instruction from the result into the
-  // operands.
-  virtual bool InstructionRequiresInputLayoutEqualToOutputLayout(
-      const HloInstruction* instruction) {
-    return false;
-  }
-
   // Construct contraints and assign layouts to all instructions in the
   // computation satisfying the given ComputationLayout. Layouts constraints are
   // added, then propagated until all LogicalBuffers in the computation are
@@ -301,21 +307,6 @@ class LayoutAssignment : public HloPassInterface {
   // required for correctness.
   Status PropagateConstraints(LayoutConstraints* constraints);
 
-  // Chooses a layout of operand `operand_no` of `instruction` that minimizes
-  // the cost of `instruction`. `output_layout` is the layout of `instruction`.
-  // Returns null if it can't decide the best layout.
-  // Precondition: `instruction` and the operand are array-shaped.
-  std::unique_ptr<Layout> ChooseOperandLayoutFromOutputLayout(
-      const Layout& output_layout, const HloInstruction* instruction,
-      int64 operand_no);
-  // Given the layout of `user`'s `operand_no`-th operand, chooses a layout of
-  // `user` that minimizes its cost on that operand.  Returns null if it can't
-  // decide the best layout.
-  // Precondition: `user` and the operand are array-shaped.
-  std::unique_ptr<Layout> ChooseOutputLayoutFromOperandLayout(
-      const Layout& operand_layout, const HloInstruction* user,
-      int64 operand_no);
-
   ComputationLayout* entry_computation_layout_;
 
   // Map containing the layouts of all computations assigned so