void DnnPoolingOp<T>::Compute(OpKernelContext* context,
                              se::dnn::PoolingMode pooling_mode,
                              const std::vector<int32>& size,
                              const std::vector<int32>& stride, Padding padding,
                              std::vector<int64> explicit_paddings,
                              TensorFormat data_format, const Tensor& tensor_in,
                              const TensorShape& tensor_out_shape,
                              bool propagate_nans) {
  Tensor* tensor_out = nullptr;
  OP_REQUIRES_OK(context,
                 context->allocate_output(0, tensor_out_shape, &tensor_out));
  if (tensor_in.shape().num_elements() == 0) {
    return;
  }

  PoolParameters params{
      context,           size,        stride,           padding,
      explicit_paddings, data_format, tensor_in.shape()};
  if (!context->status().ok()) {
    return;
  }

  int batch_size = params.tensor_in_batch;
  int depth = params.depth;
  int tensor_in_cols = params.tensor_in_cols;
  int tensor_in_rows = params.tensor_in_rows;

#if CUDNN_VERSION < 7300
  /// Earlier versions do not support NHWC format, so we need to convert it
  /// to NCHW before calling cudnn. We need to get rid of this once it is done
  Tensor transformed_input;
  if (data_format == FORMAT_NHWC) {
    OP_REQUIRES_OK(context, context->allocate_temp(
                                DataTypeToEnum<T>::value,
                                ShapeFromFormat(FORMAT_NCHW, tensor_in.shape(),
                                                data_format),
                                &transformed_input));
    functor::NHWCToNCHW<GPUDevice, T, 4>()(context->eigen_device<Device>(),
                                           tensor_in.tensor<T, 4>(),
                                           transformed_input.tensor<T, 4>());
  } else {
    transformed_input = tensor_in;
  }
  Tensor transformed_output;
  if (data_format == FORMAT_NHWC) {
    OP_REQUIRES_OK(context, context->allocate_temp(
                                DataTypeToEnum<T>::value,
                                ShapeFromFormat(FORMAT_NCHW, tensor_out_shape,
                                                data_format),
                                &transformed_output));
  } else {
    transformed_output = *tensor_out;
  }
  se::dnn::DataLayout data_layout = se::dnn::DataLayout::kBatchDepthYX;
#else
  Tensor transformed_input = tensor_in;
  auto& transformed_output = *tensor_out;
  se::dnn::DataLayout data_layout;
  switch (data_format) {
    case FORMAT_NHWC:
      data_layout = se::dnn::DataLayout::kBatchYXDepth;
      break;
    case FORMAT_NCHW:
      data_layout = se::dnn::DataLayout::kBatchDepthYX;
      break;
    case FORMAT_NCHW_VECT_C:
      // NCHW_VECT_C is not supported by cudnnPoolingForward(), but can be
      // emulated via NHWC.
      data_layout = se::dnn::DataLayout::kBatchYXDepth;
      batch_size *= depth / 4;
      depth = 4;
      break;
    default:
      OP_REQUIRES(context, false,
                  errors::InvalidArgument("Unsupported format: ",
                                          ToString(data_format)));
  }
#endif

  int64_t vertical_padding = params.pad_top;
  int64_t horizontal_padding = params.pad_left;

  if (padding == EXPLICIT && (params.pad_top != params.pad_bottom ||
                              params.pad_left != params.pad_right)) {
    // cuDNN only supports padding the same amount on the left and right sides,
    // and on the top and bottom sides. So we manually create a new padded
    // input tensor such that we can pass it to cuDNN.
    const int64_t common_padding_rows =
        std::min(params.pad_top, params.pad_bottom);
    const int64_t common_padding_cols =
        std::min(params.pad_left, params.pad_right);

    Tensor padded_input;
    const int64_t padding_rows_diff =
        std::abs(params.pad_top - params.pad_bottom);
    const int64_t padding_cols_diff =
        std::abs(params.pad_left - params.pad_right);

    const int64_t new_in_rows = tensor_in_rows + padding_rows_diff;
    const int64_t new_in_cols = tensor_in_cols + padding_cols_diff;

    OP_REQUIRES_OK(
        context,
        context->allocate_temp(DataTypeToEnum<T>::value,
                               ShapeFromFormat(data_format, batch_size,
                                               new_in_rows, new_in_cols, depth),
                               &padded_input));
    const int64_t input_pad_top = params.pad_top - common_padding_rows;
    const int64_t input_pad_bottom = params.pad_bottom - common_padding_rows;
    const int64_t input_pad_left = params.pad_left - common_padding_cols;
    const int64_t input_pad_right = params.pad_right - common_padding_cols;

    bool in_bounds =
        FastBoundsCheck(input_pad_top, std::numeric_limits<int>::max()) &&
        FastBoundsCheck(input_pad_bottom, std::numeric_limits<int>::max()) &&
        FastBoundsCheck(input_pad_left, std::numeric_limits<int>::max()) &&
        FastBoundsCheck(input_pad_right, std::numeric_limits<int>::max());
    if (!in_bounds) {
      context->SetStatus(errors::InvalidArgument("Padding is too large."));
      return;
    }

    // We need to call the const version of transformed_input.tensor()
    const Tensor& const_transformed_input = transformed_input;
    OP_REQUIRES_OK(
        context,
        PadInputWithNegativeInf<T>()(
            context->eigen_device<GPUDevice>(),
            To32Bit(const_transformed_input.tensor<T, 4>()),
            static_cast<int>(input_pad_top), static_cast<int>(input_pad_bottom),
            static_cast<int>(input_pad_left), static_cast<int>(input_pad_right),
            To32Bit(padded_input.tensor<T, 4>()), data_format));
    transformed_input = padded_input;
    vertical_padding = common_padding_rows;
    horizontal_padding = common_padding_cols;
    tensor_in_rows = new_in_rows;
    tensor_in_cols = new_in_cols;
  }

  se::dnn::PoolingDescriptor pooling_desc;
  pooling_desc.set_pooling_mode(pooling_mode)
      .set_window_height(params.window_rows)
      .set_window_width(params.window_cols)
      .set_vertical_stride(params.row_stride)
      .set_horizontal_stride(params.col_stride)
      .set_vertical_padding(vertical_padding)
      .set_horizontal_padding(horizontal_padding)
      .set_propagate_nans(propagate_nans);

  se::dnn::BatchDescriptor input_desc;
  input_desc.set_count(batch_size)
      .set_height(tensor_in_rows)
      .set_width(tensor_in_cols)
      .set_feature_map_count(depth)
      .set_layout(data_layout);

  se::dnn::BatchDescriptor output_desc;
  output_desc.set_count(batch_size)
      .set_height(params.out_height)
      .set_width(params.out_width)
      .set_feature_map_count(depth)
      .set_layout(data_layout);

  auto input_data =
      AsDeviceMemory(reinterpret_cast<const typename RawType<T>::type*>(
                         transformed_input.template flat<T>().data()),
                     transformed_input.template flat<T>().size());

  auto output_data =
      AsDeviceMemory(reinterpret_cast<const typename RawType<T>::type*>(
                         transformed_output.template flat<T>().data()),
                     transformed_output.template flat<T>().size());

  auto* stream = context->op_device_context()->stream();
  OP_REQUIRES(context, stream, errors::Internal("No GPU stream available."));

#if TENSORFLOW_USE_ROCM
  static int64 PoolingScratchSize = GetDnnWorkspaceLimit(
      // default value is in bytes despite the name of the environment variable
      "TF_CUDNN_WORKSPACE_LIMIT_IN_MB", 1LL << 32  // 4GB
  );

  DnnScratchAllocator scratch_allocator(PoolingScratchSize, context);
  bool status =
      stream
          ->ThenPoolForward(pooling_desc, input_desc, input_data, output_desc,
                            &output_data, &scratch_allocator)
          .ok();
#else
  bool status = stream
                    ->ThenPoolForward(pooling_desc, input_desc, input_data,
                                      output_desc, &output_data)
                    .ok();
#endif
  OP_REQUIRES(context, status,
              errors::Internal("dnn PoolForward launch failed"));
#if CUDNN_VERSION < 7300
  if (data_format == FORMAT_NHWC) {
    /// Transform the output data from NCHW back to NHWC
    auto toConstTensor = [](const Tensor& x) -> const Tensor { return x; };
    using RT = typename RawType<T>::type;
    functor::NCHWToNHWC<GPUDevice, RT, 4>()(
        context->eigen_device<Device>(),
        toConstTensor(transformed_output).template tensor<RT, 4>(),
        tensor_out->tensor<RT, 4>());
  }
#endif
}