message: For cuDNN >= 7.4, undo the work-around for NHWC with >2 strides for conv backward data.
  In cuDNN 7.4 release notes, they fixed the performance issue that we
  work-arounded, and it improves resnet50 a lot.
diff: @@ -65,8 +65,8 @@ HeuristicLayoutAssignment(const HloInstruction* instr,
  VLOG(2) << "Using heuristic to figure out layouts for " << instr->ToString();
  - // Empirically we've found with Volta and cudnn 7 that backward-input convs
  - // with stride are significantly faster with NCHW layouts.
  + // Empirically we've found with Volta and cudnn <= 7.3 that backward-input
  + // convs with stride are significantly faster with NCHW layouts.
  //
  // We could have used a mixed layout combination, e.g. (NHWC, NCHW, NCHW),
  // which on paper gives good performance. However, there are two observations:
  @@ -75,12 +75,18 @@ HeuristicLayoutAssignment(const HloInstruction* instr,
  // * we've also observed that for mixed layouts, cuDNN transposes data back
  // and forth from a different layout combination. If we end up with
  // transposes anyway, we prefer to have them in XLA, as they can be fused.
  - // TODO(timshen): Figure out the exact condition. This may be achieved by
  - // auto-tuning layouts offline.
  - if (instr->custom_call_target() == kCudnnConvBackwardInputCallTarget &&
  + if (auto* dnn = stream_executor->AsDnn()) {
  + auto version_status = dnn->GetVersion();
  + if (version_status.ok()) {
  + auto version = version_status.ConsumeValueOrDie();
  + if (std::make_tuple(version.major_version(), version.minor_version()) <=
  + std::make_tuple(7, 3) &&
  + instr->custom_call_target() == kCudnnConvBackwardInputCallTarget &&
  window_util::HasStride(instr->window())) {
  return kAllNCHW;
  }
  + }
  + }
  // For other Volta f16 convolutions, use NHWC.
  return kAllNHWC;
  