void Compute(OpKernelContext* context) override {
    const Tensor& filter = context->input(1);
    const TensorShape& filter_shape = filter.shape();

    const Tensor& out_backprop = context->input(2);
    const TensorShape& out_backprop_shape = out_backprop.shape();

    TensorShape input_shape;
    if (takes_shape_) {
      const Tensor& input_sizes = context->input(0);
      OP_REQUIRES_OK(context, tensor::MakeShape(input_sizes, &input_shape));
    } else {
      input_shape = context->input(0).shape();
    }

    ConvBackpropDimensions dims;
    OP_REQUIRES_OK(context, ConvBackpropComputeDimensionsV2(
                                "Conv3DBackpropInputOp", /*num_spatial_dims=*/3,
                                input_shape, filter_shape, out_backprop_shape,
                                dilation_, stride_, padding_,
                                /*explicit_paddings=*/{}, data_format_, &dims));

    Tensor* in_backprop;
    OP_REQUIRES_OK(context,
                   context->allocate_output(0, input_shape, &in_backprop));

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

    bool is_grouped_convolution = filter_shape.dim_size(3) != dims.in_depth;
    if (!is_grouped_convolution && dims.filter_size(0) == 1 &&
        dims.filter_size(1) == 1 && dims.filter_size(2) == 1 &&
        dims.dilation(0) == 1 && dims.dilation(1) == 1 &&
        dims.dilation(2) == 1 && dims.stride(0) == 1 && dims.stride(1) == 1 &&
        dims.stride(2) == 1 && data_format_ == FORMAT_NHWC) {
      const uint64 m = dims.batch_size * dims.input_size(0) *
                       dims.input_size(1) * dims.input_size(2);
      const uint64 k = dims.out_depth;
      const uint64 n = dims.in_depth;

      auto a_ptr = AsDeviceMemory(out_backprop.template flat<T>().data(),
                                  out_backprop.template flat<T>().size());
      auto b_ptr = AsDeviceMemory(filter.template flat<T>().data(),
                                  filter.template flat<T>().size());
      auto c_ptr = AsDeviceMemory(in_backprop->template flat<T>().data(),
                                  in_backprop->template flat<T>().size());

      auto transpose = se::blas::Transpose::kTranspose;
      auto no_transpose = se::blas::Transpose::kNoTranspose;

      bool blas_launch_status =
          stream
              ->ThenBlasGemm(transpose, no_transpose, n, m, k, 1.0f, b_ptr, k,
                             a_ptr, k, 0.0f, &c_ptr, n)
              .ok();
      if (!blas_launch_status) {
        context->SetStatus(errors::Internal("Blas SGEMM launch failed : m=", m,
                                            ", n=", n, ", k=", k));
      }
      return;
    } else if (!is_grouped_convolution &&
               dims.filter_size(0) == dims.input_size(0) &&
               dims.filter_size(1) == dims.input_size(1) &&
               dims.filter_size(2) == dims.input_size(2) &&
               padding_ == Padding::VALID && data_format_ == FORMAT_NHWC) {
      const uint64 m = dims.batch_size;
      const uint64 k = dims.out_depth;
      const uint64 n = dims.input_size(0) * dims.input_size(1) *
                       dims.input_size(2) * dims.in_depth;

      auto a_ptr = AsDeviceMemory(out_backprop.template flat<T>().data(),
                                  out_backprop.template flat<T>().size());
      auto b_ptr = AsDeviceMemory(filter.template flat<T>().data(),
                                  filter.template flat<T>().size());
      auto c_ptr = AsDeviceMemory(in_backprop->template flat<T>().data(),
                                  in_backprop->template flat<T>().size());

      auto transpose = se::blas::Transpose::kTranspose;
      auto no_transpose = se::blas::Transpose::kNoTranspose;

      bool blas_launch_status =
          stream
              ->ThenBlasGemm(transpose, no_transpose, n, m, k, 1.0f, b_ptr, k,
                             a_ptr, k, 0.0f, &c_ptr, n)
              .ok();
      if (!blas_launch_status) {
        context->SetStatus(errors::Internal("Blas SGEMM launch failed : m=", m,
                                            ", n=", n, ", k=", k));
      }
      return;
    }

    int padding_planes = dims.SpatialPadding(padding_, 0);
    int padding_rows = dims.SpatialPadding(padding_, 1);
    int padding_cols = dims.SpatialPadding(padding_, 2);
    const bool planes_odd = (padding_planes % 2 != 0);
    const bool rows_odd = (padding_rows % 2 != 0);
    const bool cols_odd = (padding_cols % 2 != 0);

    TensorShape compatible_input_shape;
    if (rows_odd || cols_odd || planes_odd) {
      // cuDNN only supports the same amount of padding on both sides.
      compatible_input_shape = {
          dims.batch_size,
          dims.in_depth,
          dims.input_size(0) + planes_odd,
          dims.input_size(1) + rows_odd,
          dims.input_size(2) + cols_odd,
      };
    } else {
      compatible_input_shape = {dims.batch_size, dims.in_depth,
                                dims.input_size(0), dims.input_size(1),
                                dims.input_size(2)};
    }

    CHECK(padding_rows >= 0 && padding_cols >= 0 && padding_planes >= 0)
        << "Negative paddings: (" << padding_rows << ", " << padding_cols
        << ", " << padding_planes << ")";

#if GOOGLE_CUDA
    const bool compute_in_nhwc =
        CUDNN_VERSION >= 8000 && DataTypeToEnum<T>::value == DT_HALF;
#else
    // fast NDHWC implementation is a CUDA only feature
    const bool compute_in_nhwc = false;
#endif
    const TensorFormat compute_data_format =
        (compute_in_nhwc && data_format_ == FORMAT_NHWC) ? FORMAT_NHWC
                                                         : FORMAT_NCHW;

    VLOG(3) << "Compute Conv3DBackpropInput with cuDNN:"
            << " data_format=" << ToString(data_format_)
            << " compute_data_format=" << ToString(compute_data_format);

    constexpr auto kComputeInNHWC =
        std::make_tuple(se::dnn::DataLayout::kBatchYXDepth,
                        se::dnn::FilterLayout::kOutputYXInput);
    constexpr auto kComputeInNCHW =
        std::make_tuple(se::dnn::DataLayout::kBatchDepthYX,
                        se::dnn::FilterLayout::kOutputInputYX);

    se::dnn::DataLayout compute_data_layout;
    se::dnn::FilterLayout filter_layout;

    std::tie(compute_data_layout, filter_layout) =
        compute_data_format == FORMAT_NHWC ? kComputeInNHWC : kComputeInNCHW;

    se::dnn::BatchDescriptor input_desc(3);
    input_desc.set_count(dims.batch_size)
        .set_spatial_dim(DimIndex::X, compatible_input_shape.dim_size(4))
        .set_spatial_dim(DimIndex::Y, compatible_input_shape.dim_size(3))
        .set_spatial_dim(DimIndex::Z, compatible_input_shape.dim_size(2))
        .set_feature_map_count(dims.in_depth)
        .set_layout(compute_data_layout);
    se::dnn::BatchDescriptor output_desc(3);
    output_desc.set_count(dims.batch_size)
        .set_spatial_dim(DimIndex::X, dims.output_size(2))
        .set_spatial_dim(DimIndex::Y, dims.output_size(1))
        .set_spatial_dim(DimIndex::Z, dims.output_size(0))
        .set_feature_map_count(dims.out_depth)
        .set_layout(compute_data_layout);
    se::dnn::FilterDescriptor filter_desc(3);
    filter_desc.set_spatial_dim(DimIndex::X, dims.filter_size(2))
        .set_spatial_dim(DimIndex::Y, dims.filter_size(1))
        .set_spatial_dim(DimIndex::Z, dims.filter_size(0))
        .set_input_feature_map_count(filter_shape.dim_size(3))
        .set_output_feature_map_count(filter_shape.dim_size(4))
        .set_layout(filter_layout);
    se::dnn::ConvolutionDescriptor conv_desc(3);
    conv_desc.set_dilation_rate(DimIndex::X, dims.dilation(2))
        .set_dilation_rate(DimIndex::Y, dims.dilation(1))
        .set_dilation_rate(DimIndex::Z, dims.dilation(0))
        .set_filter_stride(DimIndex::X, dims.stride(2))
        .set_filter_stride(DimIndex::Y, dims.stride(1))
        .set_filter_stride(DimIndex::Z, dims.stride(0))
        .set_zero_padding(DimIndex::X, padding_cols / 2)
        .set_zero_padding(DimIndex::Y, padding_rows / 2)
        .set_zero_padding(DimIndex::Z, padding_planes / 2)
        .set_group_count(dims.in_depth / filter_shape.dim_size(3));

    // Shape: out, in, z, y, x.
    Tensor transformed_filter;
    auto dst_format =
        compute_data_format == FORMAT_NCHW ? FORMAT_OIHW : FORMAT_OHWI;
    TensorShape dst_shape =
        dst_format == FORMAT_OIHW
            ? TensorShape({filter_shape.dim_size(4), filter_shape.dim_size(3),
                           dims.filter_size(0), dims.filter_size(1),
                           dims.filter_size(2)})
            : TensorShape({filter_shape.dim_size(4), dims.filter_size(0),
                           dims.filter_size(1), dims.filter_size(2),
                           filter_shape.dim_size(3)});
    OP_REQUIRES_OK(context,
                   context->allocate_temp(DataTypeToEnum<T>::value, dst_shape,
                                          &transformed_filter));

    functor::TransformFilter<GPUDevice, T, int, 5>()(
        context->eigen_device<GPUDevice>(), dst_format,
        To32Bit(filter.tensor<T, 5>()),
        To32Bit(transformed_filter.tensor<T, 5>()));

    // Shape: batch, filters, z, y, x.
    Tensor transformed_out_backprop;
    if (data_format_ == FORMAT_NHWC && compute_data_format == FORMAT_NCHW) {
      TensorShape nchw_shape = {dims.batch_size, dims.out_depth,
                                dims.output_size(0), dims.output_size(1),
                                dims.output_size(2)};
      if (dims.out_depth > 1) {
        OP_REQUIRES_OK(context, context->allocate_temp(
                                    DataTypeToEnum<T>::value, nchw_shape,
                                    &transformed_out_backprop));
        functor::NHWCToNCHW<GPUDevice, T, 5>()(
            context->eigen_device<GPUDevice>(), out_backprop.tensor<T, 5>(),
            transformed_out_backprop.tensor<T, 5>());
      } else {
        CHECK(transformed_out_backprop.CopyFrom(out_backprop, nchw_shape));
      }
    } else {
      transformed_out_backprop = out_backprop;
    }
    // Shape: batch, filters, z, y, x.
    Tensor pre_transformed_in_backprop;
    OP_REQUIRES_OK(context,
                   context->allocate_temp(
                       DataTypeToEnum<T>::value,
                       ShapeFromFormat(compute_data_format,
                                       compatible_input_shape.dim_size(0),
                                       {{compatible_input_shape.dim_size(2),
                                         compatible_input_shape.dim_size(3),
                                         compatible_input_shape.dim_size(4)}},
                                       compatible_input_shape.dim_size(1)),
                       &pre_transformed_in_backprop));

    auto out_backprop_ptr =
        AsDeviceMemory(transformed_out_backprop.template flat<T>().data(),
                       transformed_out_backprop.template flat<T>().size());
    auto filter_ptr =
        AsDeviceMemory(transformed_filter.template flat<T>().data(),
                       transformed_filter.template flat<T>().size());
    auto in_backprop_ptr =
        AsDeviceMemory(pre_transformed_in_backprop.template flat<T>().data(),
                       pre_transformed_in_backprop.template flat<T>().size());

    static int64 ConvolveBackwardDataScratchSize = GetDnnWorkspaceLimit(
        "TF_CUDNN_WORKSPACE_LIMIT_IN_MB", 1LL << 32);  // 4GB by default

    const int device_id = stream->parent()->device_ordinal();
    // To make sure the Conv3DBackpropInputV2 get the correct dtype, we infer
    // the dtype from 2nd input, i.e., out_backprop.
    DataType dtype = context->input(2).dtype();
    const ConvParameters conv_parameters = {
        dims.batch_size,
        dims.in_depth,
        {{dims.input_size(0), dims.input_size(1), dims.input_size(2)}},
        compute_data_format,
        dims.out_depth,
        {{dims.filter_size(0), dims.filter_size(1), dims.filter_size(2)}},
        {{dims.dilation(0), dims.dilation(1), dims.dilation(2)}},
        {{dims.stride(0), dims.stride(1), dims.stride(2)}},
        {{padding_planes, padding_rows, padding_cols}},
        dtype,
        device_id,
        conv_desc.group_count()};

    using se::dnn::AlgorithmConfig;
    using se::dnn::AlgorithmDesc;
    using se::dnn::ProfileResult;
#if TENSORFLOW_USE_ROCM
    // cudnn_use_autotune is applicable only the CUDA flow
    // for ROCm/MIOpen, we need to call GetMIOpenConvolveAlgorithms explicitly
    // if we do not have a cached algorithm_config for this conv_parameters
    cudnn_use_autotune_ = true;
#endif
    AlgorithmConfig algorithm_config;

    if (cudnn_use_autotune_ && !AutoTuneConv3dBwdData::GetInstance()->Find(
                                   conv_parameters, &algorithm_config)) {
      std::vector<std::unique_ptr<se::dnn::ConvolveExecutionPlan>> plans;
#if GOOGLE_CUDA
      std::vector<AlgorithmDesc> algorithms;
      std::vector<AlgorithmConfig> configs;
      if (CudnnUseFrontend()) {
        OP_REQUIRES(context,
                    stream->parent()->GetConvolveExecutionPlans(
                        se::dnn::ConvolutionKind::BACKWARD_DATA,
                        se::dnn::ToDataType<T>::value, stream, input_desc,
                        filter_desc, output_desc, conv_desc, &plans),
                    errors::Unknown(
                        "Failed to get convolution execution plan. This is "
                        "probably because cuDNN failed to initialize, so try "
                        "looking to see if a warning log message was printed "
                        "above."));
        for (const auto& plan : plans) {
          configs.push_back(AlgorithmConfig(
              AlgorithmDesc{plan->getTag(), plan->get_raw_desc()},
              plan->getWorkspaceSize()));
        }
      } else {
        OP_REQUIRES(context,
                    stream->parent()->GetConvolveBackwardDataAlgorithms(
                        conv_parameters.ShouldIncludeWinogradNonfusedAlgo<T>(
                            stream->parent()),
                        &algorithms),
                    errors::Unknown(
                        "Failed to get convolution execution plan. This is "
                        "probably because cuDNN failed to initialize, so try "
                        "looking to see if a warning log message was printed "
                        "above."));
        for (const auto& algorithm : algorithms) {
          configs.push_back(AlgorithmConfig(algorithm));
        }
      }

      se::TfAllocatorAdapter tf_allocator_adapter(
          context->device()->GetAllocator({}), stream);
      se::RedzoneAllocator rz_allocator(stream, &tf_allocator_adapter,
                                        se::GpuAsmOpts());
      se::DeviceMemory<T> in_backprop_ptr_rz(
          WrapRedzoneBestEffort(&rz_allocator, in_backprop_ptr));

      std::vector<tensorflow::AutotuneResult> results;
      for (auto& profile_config : configs) {
        // TODO(zhengxq): profile each algorithm multiple times to better
        // accuracy.
        DnnScratchAllocator scratch_allocator(ConvolveBackwardDataScratchSize,
                                              context);
        se::RedzoneAllocator rz_scratch_allocator(
            stream, &tf_allocator_adapter, se::GpuAsmOpts(),
            /*memory_limit=*/ConvolveBackwardDataScratchSize);
        se::ScratchAllocator* allocator_used =
            !RedzoneCheckDisabled()
                ? static_cast<se::ScratchAllocator*>(&rz_scratch_allocator)
                : static_cast<se::ScratchAllocator*>(&scratch_allocator);
        ProfileResult profile_result;

        Status cudnn_launch_status;
        if (CudnnUseFrontend()) {
          cudnn_launch_status = stream->ConvolveBackwardDataWithExecutionPlan(
              filter_desc, filter_ptr, output_desc, out_backprop_ptr, conv_desc,
              input_desc, &in_backprop_ptr_rz, allocator_used, profile_config,
              &profile_result);
        } else {
          cudnn_launch_status = stream->ConvolveBackwardDataWithAlgorithm(
              filter_desc, filter_ptr, output_desc, out_backprop_ptr, conv_desc,
              input_desc, &in_backprop_ptr_rz, allocator_used, profile_config,
              &profile_result);
        }

        if (cudnn_launch_status.ok() && profile_result.is_valid()) {
          results.emplace_back();
          auto& result = results.back();
          if (CudnnUseFrontend()) {
            result.mutable_cuda_conv_plan()->set_exec_plan_id(
                profile_config.algorithm()->exec_plan_id());
          } else {
            result.mutable_conv()->set_algorithm(
                profile_config.algorithm()->algo_id());
            result.mutable_conv()->set_tensor_ops_enabled(
                profile_config.algorithm()->tensor_ops_enabled());
          }

          result.set_scratch_bytes(
              !RedzoneCheckDisabled()
                  ? rz_scratch_allocator.TotalAllocatedBytesExcludingRedzones()
                  : scratch_allocator.TotalByteSize());
          *result.mutable_run_time() = proto_utils::ToDurationProto(
              absl::Milliseconds(profile_result.elapsed_time_in_ms()));

          // TODO(george): they don't do results at all??
          CheckRedzones(rz_scratch_allocator, &result);
          CheckRedzones(rz_allocator, &result);
        } else {
          // When CuDNN frontend APIs are used, we need to make sure the
          // profiling results are one-to-one mapping of the "plans". So, we
          // insert dummy results when the excution fails.
          results.emplace_back();
          auto& result = results.back();
          result.mutable_failure()->set_kind(AutotuneResult::UNKNOWN);
          result.mutable_failure()->set_msg(
              absl::StrCat("Profiling failure on CUDNN engine: ",
                           profile_config.algorithm()->exec_plan_id()));
        }
      }
#elif TENSORFLOW_USE_ROCM
      DnnScratchAllocator scratch_allocator(ConvolveBackwardDataScratchSize,
                                            context);
      std::vector<ProfileResult> algorithms;
      CHECK(stream->parent()->GetMIOpenConvolveAlgorithms(
          se::dnn::ConvolutionKind::BACKWARD_DATA,
          se::dnn::ToDataType<T>::value, stream, input_desc, in_backprop_ptr,
          filter_desc, filter_ptr, output_desc, out_backprop_ptr, conv_desc,
          &scratch_allocator, &algorithms));
      std::vector<tensorflow::AutotuneResult> results;
      for (auto miopen_algorithm : algorithms) {
        auto profile_algorithm = miopen_algorithm.algorithm();
        ProfileResult profile_result;
        auto miopen_launch_status = stream->ConvolveBackwardDataWithAlgorithm(
            filter_desc, filter_ptr, output_desc, out_backprop_ptr, conv_desc,
            input_desc, &in_backprop_ptr, &scratch_allocator,
            AlgorithmConfig(profile_algorithm, miopen_algorithm.scratch_size()),
            &profile_result);
        if (miopen_launch_status.ok()) {
          if (profile_result.is_valid()) {
            results.emplace_back();
            auto& result = results.back();
            result.mutable_conv()->set_algorithm(profile_algorithm.algo_id());
            result.mutable_conv()->set_tensor_ops_enabled(
                profile_algorithm.tensor_ops_enabled());
            result.set_scratch_bytes(scratch_allocator.TotalByteSize());
            *result.mutable_run_time() = proto_utils::ToDurationProto(
                absl::Milliseconds(profile_result.elapsed_time_in_ms()));
          }
        }
      }
#endif
      LogConvAutotuneResults(se::dnn::ConvolutionKind::BACKWARD_DATA,
                             se::dnn::ToDataType<T>::value, in_backprop_ptr,
                             filter_ptr, out_backprop_ptr, input_desc,
                             filter_desc, output_desc, conv_desc,
                             stream->parent(), results);
      if (CudnnUseFrontend()) {
        OP_REQUIRES_OK(context, BestCudnnConvAlgorithm(results, &plans,
                                                       &algorithm_config));
      } else {
        OP_REQUIRES_OK(context, BestCudnnConvAlgorithm(results, nullptr,
                                                       &algorithm_config));
      }
      AutoTuneConv3dBwdData::GetInstance()->Insert(conv_parameters,
                                                   algorithm_config);
    }

    Status cudnn_launch_status;
    DnnScratchAllocator scratch_allocator(ConvolveBackwardDataScratchSize,
                                          context);
    if (CudnnUseFrontend()) {
      if (algorithm_config.algorithm().has_value()) {
        VLOG(4) << "Conv3DBackpropInput Execution Plan: "
                << algorithm_config.algorithm()->exec_plan_id();
      } else {
        VLOG(4) << "Convolution AutoTune has been turned off";
      }
      cudnn_launch_status = stream->ConvolveBackwardDataWithExecutionPlan(
          filter_desc, filter_ptr, output_desc, out_backprop_ptr, conv_desc,
          input_desc, &in_backprop_ptr, &scratch_allocator, algorithm_config,
          nullptr);
    } else {
      cudnn_launch_status = stream->ConvolveBackwardDataWithAlgorithm(
          filter_desc, filter_ptr, output_desc, out_backprop_ptr, conv_desc,
          input_desc, &in_backprop_ptr, &scratch_allocator, algorithm_config,
          nullptr);
    }

    if (!cudnn_launch_status.ok()) {
      context->SetStatus(cudnn_launch_status);
    }

    if (rows_odd || cols_odd || planes_odd) {
      Tensor in_backprop_remove_padding;
      OP_REQUIRES_OK(
          context, context->allocate_temp(
                       DataTypeToEnum<T>::value,
                       ShapeFromFormat(compute_data_format, dims.batch_size,
                                       {{dims.input_size(0), dims.input_size(1),
                                         dims.input_size(2)}},
                                       dims.in_depth),
                       &in_backprop_remove_padding));

      // Remove the padding for odd spatial dimensions.
      functor::PadInput<GPUDevice, T, int, 5>()(
          context->eigen_device<GPUDevice>(),
          To32Bit(const_cast<const Tensor&>(pre_transformed_in_backprop)
                      .tensor<T, 5>()),
          {{0, 0, 0}}, {{-planes_odd, -rows_odd, -cols_odd}},
          To32Bit(in_backprop_remove_padding.tensor<T, 5>()),
          compute_data_format, T{});

      pre_transformed_in_backprop = in_backprop_remove_padding;
    }

    if (data_format_ == FORMAT_NHWC && compute_data_format == FORMAT_NCHW) {
      auto toConstTensor = [](const Tensor& x) -> const Tensor { return x; };
      functor::NCHWToNHWC<GPUDevice, T, 5>()(
          context->eigen_device<GPUDevice>(),
          toConstTensor(pre_transformed_in_backprop).template tensor<T, 5>(),
          in_backprop->tensor<T, 5>());
    } else {
      *in_backprop = pre_transformed_in_backprop;
    }
  }