[XLA:GPU] Implement BatchNormThunk as a call into cudnn.

Using cudnn for these calls is disabled by default, because it's not a
performance win on our benchmarks.

PiperOrigin-RevId: 179882911

1 files changed, 17 insertions, 0 deletions

diff --git a/tensorflow/compiler/xla/service/layout_assignment.h b/tensorflow/compiler/xla/service/layout_assignment.h
index 0b97fba744..77cd454e52 100644
--- a/tensorflow/compiler/xla/service/layout_assignment.h
+++ b/tensorflow/compiler/xla/service/layout_assignment.h
@@ -247,6 +247,19 @@ class LayoutAssignment : public HloPassInterface {
       const ResultLayoutConstraint& layout_constraint,
       LayoutConstraints* constraints);
 
+  // By default LayoutAssignment ensures that inputs and ouptuts of CustomCalls
+  // have the "major-first" layout (i.e.  {n, n-1, ..., 0}).
+  //
+  // If this function returns true, LayoutAssignment does not set a layout for
+  // the given CustomCall.  It's up to the backend to set one in
+  // AddBackendConstraints, if necessary.
+  //
+  // Precondition: instruction->opcode() == HloOpcode::kCustomCall.
+  virtual bool CustomCallRequiresMajorFirstLayout(
+      const HloInstruction* /*instruction*/) {
+    return true;
+  }
+
   // Called after layouts of an instruction have been finalized to allow
   // subclasses to check for platform specific assumptions.
   virtual Status Verify(const HloInstruction* instruction) {
@@ -315,6 +328,10 @@ class LayoutAssignment : public HloPassInterface {
   // required for correctness.
   Status PropagateConstraints(LayoutConstraints* constraints);
 
+  // Check that all layouts in the module have been set and satisfy all
+  // necessary conditions.
+  Status CheckLayouts(HloModule* module);
+
   ComputationLayout* entry_computation_layout_;
 
  protected: