text: /* Copyright 2016 The TensorFlow Authors. All Rights Reserved.
  
  Licensed under the Apache License, Version 2.0 (the "License");
  you may not use this file except in compliance with the License.
  You may obtain a copy of the License at
  
      http://www.apache.org/licenses/LICENSE-2.0
  
  Unless required by applicable law or agreed to in writing, software
  distributed under the License is distributed on an "AS IS" BASIS,
  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  See the License for the specific language governing permissions and
  limitations under the License.
  ==============================================================================*/
  
  // Implements convolution operations with other kernels baked into the
  // processing, to optimize latency and memory usage:
  //  - Conv2D + BiasAdd + <Activation>
  //  - Conv2D + FusedBatchNorm + <Activation>
  //
  // Activation: Relu, Relu6, Elu, etc...
  //
  // Kernels for convolutions fused with image transformations (resize and mirror
  // padding) defined in `conv_ops_fused_image_transform.cc`.
  //
  // For the CPU device we implement fusion with an Eigen tensor contraction
  // output kernel. For the GPU device we rely on CuDNN primitives.
  //
  // NOTE: GPU only supports fusion of Conv2D + BiasAdd + <optional Relu>.
  
  #ifndef TENSORFLOW_CORE_KERNELS_CONV_OPS_FUSED_IMPL_H_
  #define TENSORFLOW_CORE_KERNELS_CONV_OPS_FUSED_IMPL_H_
  
  #define USE_EIGEN_TENSOR
  #define EIGEN_USE_THREADS
  
  #if GOOGLE_CUDA
  #define EIGEN_USE_GPU
  #endif  // GOOGLE_CUDA
  
  #include <string>
  #include <vector>
  
  #include "absl/strings/str_cat.h"
  #include "absl/strings/str_join.h"
  #include "absl/strings/substitute.h"
  #include "tensorflow/core/framework/bounds_check.h"
  #include "tensorflow/core/framework/op_kernel.h"
  #include "tensorflow/core/framework/register_types.h"
  #include "tensorflow/core/framework/tensor.h"
  #include "tensorflow/core/framework/tensor_shape.h"
  #include "tensorflow/core/kernels/conv_2d.h"
  #include "tensorflow/core/kernels/conv_ops.h"
  #include "tensorflow/core/kernels/ops_util.h"
  #include "tensorflow/core/util/tensor_format.h"
  #include "tensorflow/core/util/use_cudnn.h"
  
  #if GOOGLE_CUDA
  #include "cuda/include/cudnn.h"
  #include "tensorflow/core/kernels/conv_ops_gpu.h"
  #include "tensorflow/core/platform/stream_executor.h"
  #include "tensorflow/core/util/proto/proto_utils.h"
  #endif  // GOOGLE_CUDA
  
  namespace tensorflow {
  
  class AutotuneResult;
  
  typedef Eigen::ThreadPoolDevice CPUDevice;
  typedef Eigen::GpuDevice GPUDevice;
  
  // Supported Conv2D fusions. Not all of them supported on all type of devices.
  enum class FusedComputationType {
    kUndefined,
    // NOTE(ezhulenev): CuDNN `cudnnConvolutionBiasActivationForward` supports
    // identity activation function, it in theory should allow to fuse convolution
    // with BiasAdd, but in practice it doesn't work, cuDNN ignores this parameter
    // and always does Relu activation.
    kBiasAdd,                  // CPU
    kBiasAddWithRelu,          // CPU and GPU
    kBiasAddWithRelu6,         // CPU
    kBiasAddWithElu,           // CPU
    kFusedBatchNorm,           // CPU
    kFusedBatchNormWithRelu,   // CPU
    kFusedBatchNormWithRelu6,  // CPU
    kFusedBatchNormWithElu     // CPU
  };
  
  // We have to pass around additional arguments for all possible fusion types.
  struct FusedComputationArgs {
    float epsilon = 0.0;  // Used by `FusedBatchNorm` fusion only
  };
  
  template <typename Device, typename T>
  struct LaunchFusedConv2DOp {
    void operator()(OpKernelContext* context, bool use_cudnn,
                    bool cudnn_use_autotune, const Tensor& input,
                    const Tensor& filter, FusedComputationType fusion,
                    const FusedComputationArgs& fusion_args,
                    const Conv2DParameters& params,
                    const Conv2DDimensions& dimensions, Tensor* output);
  };
  
  // Type alias for the tensor contraction output mapper.
  template <typename Scalar, typename Index>
  using ContractionOutputMapper =
      Eigen::internal::blas_data_mapper<Scalar, Index, Eigen::ColMajor>;
  
  // Returns input expression without any transformations.
  struct Identity {
    template <typename XprType>
    static auto apply(XprType expr) -> XprType {
      return expr;
    };
  };
  
  // Applies `Relu` to the passed input expression.
  struct Relu {
    template <typename XprType>
    static auto apply(XprType expr)
        -> decltype(expr.cwiseMax(std::declval<typename XprType::Scalar>())) {
      return expr.cwiseMax(static_cast<typename XprType::Scalar>(0));
    };
  };
  
  // Applies `Relu6` to the passed input expression.
  struct Relu6 {
    template <typename XprType>
    static auto apply(XprType expr)
        -> decltype(expr.cwiseMax(std::declval<typename XprType::Scalar>())
                        .cwiseMin(std::declval<typename XprType::Scalar>())) {
      return expr.cwiseMax(static_cast<typename XprType::Scalar>(0))
          .cwiseMin(static_cast<typename XprType::Scalar>(6));
    };
  };
  
  // Applies `Elu` to the passed input expression.
  struct Elu {
    template <typename XprType>
    static auto apply(XprType expr) -> decltype(
        (expr < std::declval<typename XprType::Scalar>())
            .select(expr.exp() -
                        expr.constant(std::declval<typename XprType::Scalar>()),
                    expr)) {
      return (expr < static_cast<typename XprType::Scalar>(0))
          .select(expr.exp() -
                      expr.constant(static_cast<typename XprType::Scalar>(1)),
                  expr);
    };
  };
  
  // TensorContraction swaps lhs with rhs, and changes layout from RowMajor
  // (default in Tensorflow) to ColMajor (preferred in Eigen), and computes matmul
  // using these tensors.
  //
  // TensorContraction output matrix (before reshape) has a ColMajor layout, and
  // has dimensions:
  //  - rows: output_channels
  //  - cols: all other dimensions
  //
  // First element in every column is:
  //   [batch ??, height ??, width ??, out_channel = i]
  //
  // We do not know what are the values of the 'batch', 'height', and 'width' here
  // (if we know original dimensions, they can be computed from 'j').
  //
  // Each column of an output block is a continuous slice along the output channel
  // dimension, so we can use it to efficiently compute any transformation that
  // depends only on a channel value (e.g. add channel bias).
  
  // Output kernel that fuses BiasAdd operation into the output of tensor
  // contraction + activation function defined by Activation.
  template <typename T, typename Activation = Identity>
  struct BiasAddOutputKernel {
    explicit BiasAddOutputKernel(const T* bias_data) : bias_data(bias_data) {}
  
    template <typename Index, typename Scalar>
    EIGEN_ALWAYS_INLINE void operator()(
        const ContractionOutputMapper<Scalar, Index>& output_mapper,
        const Eigen::TensorContractionParams& params, Index i, Index j,
        Index num_rows, Index num_cols) const {
      DCHECK(params.swapped_arguments);
  
      const T* bias_base = bias_data + i;
      typename TTypes<T>::UnalignedConstTensor bias(bias_base, num_rows);
  
      for (int col = 0; col < num_cols; ++col) {
        T* output_base = &output_mapper(0, col);
        typename TTypes<T>::UnalignedTensor output(output_base, num_rows);
        const auto expr = output + bias;
        output = Activation::template apply<decltype(expr)>(expr);
      }
    }
  
   private:
    const T* bias_data;
  };
  
  // Output kernel that fuses FusedBatchNorm operation into the output of tensor
  // contraction + activation function defined by Activation.
  template <typename T, typename Activation = Identity>
  struct FusedBatchNormOutputKernel {
    FusedBatchNormOutputKernel(T epsilon, const T* scaling_factor_data,
                               const T* offset_data, const T* estimated_mean_data)
        : epsilon(epsilon),
          scaling_factor_data(scaling_factor_data),
          offset_data(offset_data),
          estimated_mean_data(estimated_mean_data) {}
  
    template <typename Index, typename Scalar>
    EIGEN_ALWAYS_INLINE void operator()(
        const ContractionOutputMapper<Scalar, Index>& output_mapper,
        const Eigen::TensorContractionParams& params, Index i, Index j,
        Index num_rows, Index num_cols) const {
      DCHECK(params.swapped_arguments);
  
      const T* scaling_factor_base = scaling_factor_data + i;
      const T* offset_base = offset_data + i;
      const T* mean_base = estimated_mean_data + i;
  
      typename TTypes<T>::UnalignedConstTensor scaling_factor(scaling_factor_base,
                                                              num_rows);
      typename TTypes<T>::UnalignedConstTensor offset(offset_base, num_rows);
      typename TTypes<T>::UnalignedConstTensor mean(mean_base, num_rows);
  
      for (int col = 0; col < num_cols; ++col) {
        T* output_base = &output_mapper(0, col);
        typename TTypes<T>::UnalignedTensor output(output_base, num_rows);
  
        auto scaled = (output - mean) * scaling_factor;
        auto shifted = scaled + offset;
  
        output = Activation::template apply<decltype(shifted)>(shifted);
      }
    }
  
   private:
    T epsilon;
    const T* scaling_factor_data;
    const T* offset_data;
    const T* estimated_mean_data;
  };
  
  // Type aliases for the output kernels, purely for the sake of better launch
  // dispatching code readability.
  template <typename T>
  using WithBiasAdd = BiasAddOutputKernel<T>;
  template <typename T>
  using WithBiasAddAndRelu = BiasAddOutputKernel<T, Relu>;
  template <typename T>
  using WithBiasAddAndRelu6 = BiasAddOutputKernel<T, Relu6>;
  template <typename T>
  using WithBiasAddAndElu = BiasAddOutputKernel<T, Elu>;
  template <typename T>
  using WithFusedBatchNorm = FusedBatchNormOutputKernel<T>;
  template <typename T>
  using WithFusedBatchNormAndRelu = FusedBatchNormOutputKernel<T, Relu>;
  template <typename T>
  using WithFusedBatchNormAndRelu6 = FusedBatchNormOutputKernel<T, Relu6>;
  template <typename T>
  using WithFusedBatchNormAndElu = FusedBatchNormOutputKernel<T, Elu>;
  
  // This is CPU-only implementation that uses Eigen contraction output kernels.
  //
  // Dispatch 2D convolution to the appropriate primitive operation:
  //   (1) MatMul for the case of 1x1 convolution.
  //   (2) MatMul for the case when filter size equals to the input size.
  //   (3) General spatial 2D convolution for all other cases.
  template <typename T>
  class LaunchFusedConv2DWithOutputKernel {
   public:
    LaunchFusedConv2DWithOutputKernel(int row_stride, int col_stride,      //
                                      int row_dilation, int col_dilation,  //
                                      Padding padding)
        : row_stride_(row_stride),
          col_stride_(col_stride),
          row_dilation_(row_dilation),
          col_dilation_(col_dilation),
          padding_(padding) {}
  
    template <typename OutputKernel>
    void operator()(const OutputKernel& output_kernel, OpKernelContext* ctx,
                    const Tensor& input, const Tensor& filter, Tensor* output) {
      if (filter.dim_size(0) == 1 && filter.dim_size(1) == 1 &&
          row_stride_ == 1 && col_stride_ == 1) {
        int conv_width = 1;  // Width for the convolution step.
        for (int i = 0; i < 3; ++i) {
          conv_width *= output->dim_size(i);
        }
  
        Eigen::array<Eigen::IndexPair<Eigen::DenseIndex>, 1> dim_pair;
        dim_pair[0] = Eigen::IndexPair<Eigen::DenseIndex>(1, 0);
        functor::MatMulConvFunctor<CPUDevice, T, OutputKernel>()(
            ctx->eigen_device<CPUDevice>(),
            output->shaped<T, 2>({conv_width, filter.dim_size(3)}),
            input.shaped<T, 2>({conv_width, filter.dim_size(2)}),
            filter.shaped<T, 2>({filter.dim_size(2), filter.dim_size(3)}),
            dim_pair, output_kernel);
  
      } else if (filter.dim_size(0) == input.dim_size(1) &&
                 filter.dim_size(1) == input.dim_size(2) && row_dilation_ == 1 &&
                 col_dilation_ == 1 && padding_ == VALID) {
        // If the input data and filter have the same height/width,
        // reduce the 2D convolution to matrix multiplication.
        const auto k =  // Length of reduction dimension.
            filter.dim_size(0) * filter.dim_size(1) * filter.dim_size(2);
  
        Eigen::array<Eigen::IndexPair<Eigen::DenseIndex>, 1> dim_pair;
        dim_pair[0] = Eigen::IndexPair<Eigen::DenseIndex>(1, 0);
        functor::MatMulConvFunctor<CPUDevice, T, OutputKernel>()(
            ctx->eigen_device<CPUDevice>(),
            output->shaped<T, 2>({input.dim_size(0), filter.dim_size(3)}),
            input.shaped<T, 2>({input.dim_size(0), k}),
            filter.shaped<T, 2>({k, filter.dim_size(3)}), dim_pair,
            output_kernel);
  
      } else {
        functor::SpatialConvolution<CPUDevice, T, OutputKernel>()(
            ctx->eigen_device<CPUDevice>(), output->tensor<T, 4>(),
            input.tensor<T, 4>(), filter.tensor<T, 4>(), row_stride_, col_stride_,
            row_dilation_, col_dilation_, BrainPadding2EigenPadding(padding_),
            output_kernel);
      }
    }
  
   private:
    int row_stride_;
    int col_stride_;
    int row_dilation_;
    int col_dilation_;
    const Padding padding_;
  };
  
  template <typename T>
  struct LaunchFusedConv2DOp<CPUDevice, T> {
    void operator()(OpKernelContext* context, bool use_cudnn,
                    bool cudnn_use_autotune, const Tensor& input,
                    const Tensor& filter, const FusedComputationType fusion,
                    const FusedComputationArgs& fusion_args,
                    const Conv2DParameters& params,
                    const Conv2DDimensions& dimensions, Tensor* output) {
      OP_REQUIRES(context, dimensions.in_depth == filter.dim_size(2),
                  errors::Unimplemented("Fused conv implementation does not "
                                        "support grouped convolutions for now."));
      OP_REQUIRES(context, params.data_format == FORMAT_NHWC,
                  errors::Unimplemented("Fused conv implementation only supports "
                                        "NHWC tensor format for now."));
  
      BiasAddArgs bias_add;
      FusedBatchNormArgs fused_batch_norm;
  
      LaunchFusedConv2DWithOutputKernel<T> conv2d(
          dimensions.stride_rows, dimensions.stride_cols,
          dimensions.dilation_rows, dimensions.dilation_cols, params.padding);
  
      switch (fusion) {
        case FusedComputationType::kUndefined:
          OP_REQUIRES_OK(context, errors::Internal("Fusion type is undefined"));
  
        case FusedComputationType::kBiasAdd:
          OP_REQUIRES_OK(context, InitBiasAddArgs(context, &bias_add));
          conv2d(WithBiasAdd<T>(bias_add.bias_add_data), context, input, filter,
                 output);
          break;
  
        case FusedComputationType::kBiasAddWithRelu:
          OP_REQUIRES_OK(context, InitBiasAddArgs(context, &bias_add));
          conv2d(WithBiasAddAndRelu<T>(bias_add.bias_add_data), context, input,
                 filter, output);
          break;
  
        case FusedComputationType::kBiasAddWithRelu6:
          OP_REQUIRES_OK(context, InitBiasAddArgs(context, &bias_add));
          conv2d(WithBiasAddAndRelu6<T>(bias_add.bias_add_data), context, input,
                 filter, output);
          break;
  
        case FusedComputationType::kBiasAddWithElu:
          OP_REQUIRES_OK(context, InitBiasAddArgs(context, &bias_add));
          conv2d(WithBiasAddAndElu<T>(bias_add.bias_add_data), context, input,
                 filter, output);
          break;
  
        case FusedComputationType::kFusedBatchNorm:
          OP_REQUIRES_OK(context,
                         InitFusedBatchNormArgs(context, fusion_args.epsilon,
                                                &fused_batch_norm));
          conv2d(WithFusedBatchNorm<T>(fusion_args.epsilon,
                                       fused_batch_norm.scaling_factor.data(),
                                       fused_batch_norm.offset_data,
                                       fused_batch_norm.estimated_mean_data),
                 context, input, filter, output);
          break;
  
        case FusedComputationType::kFusedBatchNormWithRelu:
          OP_REQUIRES_OK(context,
                         InitFusedBatchNormArgs(context, fusion_args.epsilon,
                                                &fused_batch_norm));
          conv2d(WithFusedBatchNormAndRelu<T>(
                     fusion_args.epsilon, fused_batch_norm.scaling_factor.data(),
                     fused_batch_norm.offset_data,
                     fused_batch_norm.estimated_mean_data),
                 context, input, filter, output);
          break;
  
        case FusedComputationType::kFusedBatchNormWithRelu6:
          OP_REQUIRES_OK(context,
                         InitFusedBatchNormArgs(context, fusion_args.epsilon,
                                                &fused_batch_norm));
          conv2d(WithFusedBatchNormAndRelu6<T>(
                     fusion_args.epsilon, fused_batch_norm.scaling_factor.data(),
                     fused_batch_norm.offset_data,
                     fused_batch_norm.estimated_mean_data),
                 context, input, filter, output);
          break;
  
        case FusedComputationType::kFusedBatchNormWithElu:
          OP_REQUIRES_OK(context,
                         InitFusedBatchNormArgs(context, fusion_args.epsilon,
                                                &fused_batch_norm));
          conv2d(WithFusedBatchNormAndElu<T>(
                     fusion_args.epsilon, fused_batch_norm.scaling_factor.data(),
                     fused_batch_norm.offset_data,
                     fused_batch_norm.estimated_mean_data),
                 context, input, filter, output);
          break;
      }
    }
  
   private:
    struct BiasAddArgs {
      const T* bias_add_data = nullptr;
    };
  
    struct FusedBatchNormArgs {
      const T* scale_data = nullptr;
      const T* offset_data = nullptr;
      const T* estimated_mean_data = nullptr;
      const T* estimated_variance_data = nullptr;
  
      // Precomputed expression:
      //   scaling_factor = (estimated_variance + epsilon).rsqrt() * scale
      Eigen::Tensor<T, 1, Eigen::RowMajor> scaling_factor;
    };
  
  #define TF_REQUIRES(EXP, STATUS) \
    if (!TF_PREDICT_TRUE(EXP)) return (STATUS)
  
    void InitDataPtr(const Tensor& tensor, const T** ptr) const {
      *ptr = reinterpret_cast<const T*>(tensor.tensor_data().data());
    }
  
    Status InitBiasAddArgs(OpKernelContext* context, BiasAddArgs* args) const {
      // Bias of the following dimensions: [ output_depth ]
      const Tensor& bias = context->input(2);
  
      TF_REQUIRES(bias.dims() == 1,
                  errors::InvalidArgument("bias must be 1-dimensional",
                                          bias.shape().DebugString()));
  
      InitDataPtr(bias, &args->bias_add_data);
  
      return Status::OK();
    }
  
    Status InitFusedBatchNormArgs(OpKernelContext* context, float epsilon,
                                  FusedBatchNormArgs* args) const {
      const Tensor& scale = context->input(2);
      const Tensor& offset = context->input(3);
      const Tensor& estimated_mean = context->input(4);
      const Tensor& estimated_variance = context->input(5);
  
      TF_REQUIRES(scale.dims() == 1,
                  errors::InvalidArgument("scale must be 1-dimensional",
                                          scale.shape().DebugString()));
      TF_REQUIRES(offset.dims() == 1,
                  errors::InvalidArgument("offset must be 1-dimensional",
                                          offset.shape().DebugString()));
      TF_REQUIRES(estimated_mean.dims() == 1,
                  errors::InvalidArgument("estimated_mean must be 1-dimensional",
                                          estimated_mean.shape().DebugString()));
      TF_REQUIRES(
          estimated_variance.dims() == 1,
          errors::InvalidArgument("estimated_variance must be 1-dimensional",
                                  estimated_variance.shape().DebugString()));
  
      InitDataPtr(scale, &args->scale_data);
      InitDataPtr(offset, &args->offset_data);
      InitDataPtr(estimated_mean, &args->estimated_mean_data);
      InitDataPtr(estimated_variance, &args->estimated_variance_data);
  
      // Precompute scaling factor once for all output blocks (kernels).
      args->scaling_factor =
          (estimated_variance.flat<T>() + static_cast<T>(epsilon)).rsqrt() *
          scale.flat<T>();
  
      return Status::OK();
    }
  
  #undef TF_REQUIRES
  };
  
  #if GOOGLE_CUDA
  
  // Encapsulate the default shape information that is used by the convolution
  // operation, and add an activation mode for the fusion.
  class FusedConvParameters : public ConvParameters {
   public:
    FusedConvParameters(const ConvParameters& base,
                        const se::dnn::ActivationMode activation_mode)
        : ConvParameters(base), activation_mode_(activation_mode) {}
  
    string ToString() const {
      return absl::StrCat(ConvParameters::ToString(), ", ", activation_mode_);
    }
  
   private:
    friend bool operator==(const FusedConvParameters& lhs,
                           const FusedConvParameters& rhs);
  
    using ParameterDataType =
        std::tuple<ConvParameters::ParameterDataType, se::dnn::ActivationMode>;
  
    ParameterDataType get_data_as_tuple() const {
      return std::make_tuple(ConvParameters::get_data_as_tuple(),
                             activation_mode_);
    }
  
    se::dnn::ActivationMode activation_mode_;
  };
  
  inline bool operator==(const FusedConvParameters& lhs,
                         const FusedConvParameters& rhs) {
    return lhs.get_data_as_tuple() == rhs.get_data_as_tuple();
  }
  
  inline bool operator!=(const FusedConvParameters& lhs,
                         const FusedConvParameters& rhs) {
    return !(lhs == rhs);
  }
  
  // A dummy type to group forward convolution autotune results together.
  struct FusedConvAutoTuneGroup {
    static string name() { return "FusedConv"; }
  };
  
  using AutoTuneFusedConv =
      AutoTuneSingleton<FusedConvAutoTuneGroup, FusedConvParameters,
                        se::dnn::AlgorithmConfig>;
  
  inline int64 ConvolveScratchSize() {
    static int64 convolve_scratch_size = GetDnnWorkspaceLimit(
        // default value is in bytes despite the name of the environment variable
        "TF_CUDNN_WORKSPACE_LIMIT_IN_MB", 1LL << 32  // 4GB
    );
    return convolve_scratch_size;
  }
  
  // Finds the best convolutiun algorithm for the given ConvLaunch (cuda
  // convolution on the stream) and parameters, by running all possible
  // algorithms and measuring execution time.
  // TODO(ezhulenev): Move it to conv_ops_gpu.h and share with conv_ops.cc.
  template <typename T, typename ConvLaunch, typename LogFunc>
  Status FindBestConvolveAlgorithm(const FusedConvParameters& params,
                                   const ConvLaunch launch,
                                   OpKernelContext* context, se::Stream* stream,
                                   const LogFunc& log,
                                   se::dnn::AlgorithmConfig* algorithm_config) {
    // Check if we already have an algorithm selected for the given parameters.
    if (AutoTuneFusedConv::GetInstance()->Find(params, algorithm_config)) {
      return Status::OK();
    }
  
    // Find all candidate algorithms.
    std::vector<se::dnn::AlgorithmDesc> algorithms;
    if (!stream->parent()->GetConvolveAlgorithms(
            params.ShouldIncludeWinogradNonfusedAlgo<T>(stream->parent()),
            &algorithms)) {
      return errors::Unknown(
          "Failed to get convolution algorithm. This is probably "
          "because cuDNN failed to initialize, so try looking to "
          "see if a warning log message was printed above.");
    }
  
    std::vector<tensorflow::AutotuneResult> results;
    for (auto profile_algorithm : algorithms) {
      DnnScratchAllocator scratch_allocator(ConvolveScratchSize(), context);
      se::dnn::ProfileResult profile_result;
  
      bool cudnn_launch_status =
          launch(se::dnn::AlgorithmConfig(profile_algorithm), &scratch_allocator,
                 &profile_result);
  
      if (cudnn_launch_status && 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()));
      }
    }
    // Only log on an AutoTuneFusedConv cache miss.
    log(results);
    TF_RETURN_IF_ERROR(BestCudnnConvAlgorithm(results, algorithm_config));
    AutoTuneFusedConv::GetInstance()->Insert(params, *algorithm_config);
    return Status::OK();
  }
  
  template <typename T>
  struct LaunchFusedConv2DOp<GPUDevice, T> {
    void operator()(OpKernelContext* context, bool use_cudnn,
                    bool cudnn_use_autotune, const Tensor& input_param,
                    const Tensor& filter, FusedComputationType fusion,
                    const FusedComputationArgs& fusion_args,
                    const Conv2DParameters& params,
                    const Conv2DDimensions& dimensions, Tensor* output) {
      OP_REQUIRES(
          context,
          params.data_format == FORMAT_NHWC || params.data_format == FORMAT_NCHW,
          errors::Unimplemented("Fused conv implementation only supports "
                                "NHWC and HCHW tensor formats for now."));
  
      auto* stream = context->op_device_context()->stream();
      OP_REQUIRES(context, stream, errors::Internal("No GPU stream available."));
      OP_REQUIRES(
          context, use_cudnn,
          errors::Unimplemented("FusedConv2D for GPU is not currently supported "
                                "without cudnn"));
  
      OP_REQUIRES(
          context, fusion == FusedComputationType::kBiasAddWithRelu,
          errors::Unimplemented("FusedConv2D implementation only supports "
                                "fusing with `BiasAdd + Relu` for now."));
  
      Tensor input = input_param;
  
      const int64 in_batch = GetTensorDim(input, params.data_format, 'N');
      int64 in_rows = GetTensorDim(input, params.data_format, 'H');
      int64 in_cols = GetTensorDim(input, params.data_format, 'W');
      const int64 in_depths = GetTensorDim(input, params.data_format, 'C');
  
      const int64 patch_rows = filter.dim_size(0);
      const int64 patch_cols = filter.dim_size(1);
      const int64 patch_depths = filter.dim_size(2);
  
      int64 padding_rows = 0;
      int64 padding_cols = 0;
      const int64 out_batch = GetTensorDim(*output, params.data_format, 'N');
      const int64 out_rows = GetTensorDim(*output, params.data_format, 'H');
      const int64 out_cols = GetTensorDim(*output, params.data_format, 'W');
      const int64 out_depths = GetTensorDim(*output, params.data_format, 'C');
  
      // Bias of the following dimensions: [ output_depth ]
      const Tensor& bias = context->input(2);
      OP_REQUIRES(context, bias.dims() == 1,
                  errors::InvalidArgument("bias must be 1-dimensional",
                                          bias.shape().DebugString()));
      OP_REQUIRES(context, bias.dim_size(0) == out_depths,
                  errors::InvalidArgument("bias depth must be equal to out depth",
                                          bias.shape().DebugString()));
  
      if (params.padding == SAME) {
        // Total padding on rows and cols is
        // Pr = (R' - 1) * S + (Kr - 1) * Dr + 1 - R
        // Pc = (C' - 1) * S + (Kc - 1) * Dc + 1 - C
        // where (R', C') are output dimensions, (R, C) are input dimensions, S
        // is stride, (Dr, Dc) are dilations, (Kr, Kc) are filter dimensions.
        // We pad Pr/2 on the left and Pr - Pr/2 on the right, Pc/2 on the top
        // and Pc - Pc/2 on the bottom.  When Pr or Pc is odd, this means
        // we pad more on the right and bottom than on the top and left.
        padding_rows = std::max<int>(
            0, (out_rows - 1) * dimensions.stride_rows +
                   (patch_rows - 1) * dimensions.dilation_rows + 1 - in_rows);
        padding_cols = std::max<int>(
            0, (out_cols - 1) * dimensions.stride_cols +
                   (patch_cols - 1) * dimensions.dilation_cols + 1 - in_cols);
        const bool rows_odd = (padding_rows % 2 != 0);
        const bool cols_odd = (padding_cols % 2 != 0);
        if (rows_odd || cols_odd) {
          Tensor transformed_input;
          int64 new_in_rows = in_rows + rows_odd;
          int64 new_in_cols = in_cols + cols_odd;
          OP_REQUIRES_OK(context,
                         context->allocate_temp(
                             DataTypeToEnum<T>::value,
                             ShapeFromFormat(params.data_format, in_batch,
                                             new_in_rows, new_in_cols, in_depths),
                             &transformed_input));
  
          functor::PadInput<GPUDevice, T, int, 4>()(
              context->eigen_device<GPUDevice>(),
              To32Bit(input_param.tensor<T, 4>()), {{0, 0}},
              {{rows_odd, cols_odd}}, To32Bit(transformed_input.tensor<T, 4>()),
              params.data_format);
  
          input = transformed_input;
          in_rows = new_in_rows;
          in_cols = new_in_cols;
        }
      }
  
      if (params.data_format == FORMAT_NHWC) {
        // Convert the input tensor from NHWC to NCHW.
        TensorShape nchw_shape =
            ShapeFromFormat(FORMAT_NCHW, in_batch, in_rows, in_cols, in_depths);
        if (in_depths > 1) {
          Tensor transformed_input;
          OP_REQUIRES_OK(context,
                         context->allocate_temp(DataTypeToEnum<T>::value,
                                                nchw_shape, &transformed_input));
          functor::NHWCToNCHW<GPUDevice, T, 4>()(
              context->eigen_device<GPUDevice>(),
              const_cast<const Tensor&>(input).tensor<T, 4>(),
              transformed_input.tensor<T, 4>());
          input = transformed_input;
        } else {
          // If depth <= 1, then just reshape.
          CHECK(input.CopyFrom(input, nchw_shape));  // Crash OK
        }
      }
  
      CHECK(padding_rows >= 0) << "Negative padding rows";  // Crash OK
      CHECK(padding_cols >= 0) << "Negative padding cols";  // Crash OK
  
      se::dnn::ActivationMode dnn_activation_mode;
      switch (fusion) {
        case FusedComputationType::kBiasAddWithRelu:
          dnn_activation_mode = se::dnn::ActivationMode::kRelu;
          break;
        default:
          LOG(FATAL) << "Unsupported fusion type";  // Crash OK
      }
  
      se::dnn::BatchDescriptor input_desc;
      input_desc.set_count(in_batch)
          .set_feature_map_count(in_depths)
          .set_height(in_rows)
          .set_width(in_cols)
          .set_layout(se::dnn::DataLayout::kBatchDepthYX);
      se::dnn::FilterDescriptor filter_desc;
      filter_desc.set_input_filter_height(patch_rows)
          .set_input_filter_width(patch_cols)
          .set_input_feature_map_count(patch_depths)
          .set_output_feature_map_count(filter.dim_size(3));
      se::dnn::BatchDescriptor bias_desc;
      bias_desc.set_count(1)
          .set_height(1)
          .set_width(1)
          .set_feature_map_count(out_depths)
          .set_layout(se::dnn::DataLayout::kBatchDepthYX);
      se::dnn::ConvolutionDescriptor conv_desc;
      conv_desc.set_vertical_dilation_rate(dimensions.dilation_rows)
          .set_horizontal_dilation_rate(dimensions.dilation_cols)
          .set_vertical_filter_stride(dimensions.stride_rows)
          .set_horizontal_filter_stride(dimensions.stride_cols)
          .set_zero_padding_height(padding_rows / 2)
          .set_zero_padding_width(padding_cols / 2)
          .set_group_count(in_depths / patch_depths);
      se::dnn::BatchDescriptor output_desc;
      output_desc.set_count(out_batch)
          .set_height(out_rows)
          .set_width(out_cols)
          .set_feature_map_count(out_depths)
          .set_layout(se::dnn::DataLayout::kBatchDepthYX);
  
      Tensor transformed_filter;
      OP_REQUIRES_OK(context,
                     context->allocate_temp(
                         DataTypeToEnum<T>::value,
                         TensorShape({filter.dim_size(3), filter.dim_size(2),
                                      filter.dim_size(0), filter.dim_size(1)}),
                         &transformed_filter));
      functor::TransformFilter<GPUDevice, T, int, 4>()(
          context->eigen_device<GPUDevice>(), FORMAT_OIHW,
          To32Bit(filter.tensor<T, 4>()),
          To32Bit(transformed_filter.tensor<T, 4>()));
  
      Tensor transformed_output;
      if (params.data_format == FORMAT_NHWC) {
        // Only allocate temporary memory when a layout transformation is needed.
        OP_REQUIRES_OK(context,
                       context->allocate_temp(
                           DataTypeToEnum<T>::value,
                           ShapeFromFormat(FORMAT_NCHW, out_batch, out_rows,
                                           out_cols, out_depths),
                           &transformed_output));
      } else {
        transformed_output = *output;
      }
  
      const auto tensor_on_device = [](const Tensor& t) -> se::DeviceMemory<T> {
        return AsDeviceMemory(t.template flat<T>().data(),
                              t.template flat<T>().size());
      };
  
      se::DeviceMemory<T> input_ptr = tensor_on_device(input);
      se::DeviceMemory<T> filter_ptr = tensor_on_device(transformed_filter);
      se::DeviceMemory<T> bias_ptr = tensor_on_device(bias);
      se::DeviceMemory<T> output_ptr = tensor_on_device(transformed_output);
  
      // We do not use side inputs, so we can safely pass nullptr.
      se::DeviceMemory<T> side_input_ptr =
          AsDeviceMemory(static_cast<T*>(nullptr), 0);
  
      int device_id = stream->parent()->device_ordinal();
      DataType dtype = input.dtype();
      FusedConvParameters conv_parameters = {
          {
              in_batch,                      // batch
              in_depths,                     // in_depths
              {{in_rows,                     // in_rows
                in_cols}},                   // in_cols
              FORMAT_NCHW,                   // compute_data_format
              out_depths,                    // out_depths
              {{patch_rows,                  // filter_rows
                patch_cols,                  // filter_cols
                patch_depths}},              // filter_depths
              {{dimensions.dilation_rows,    // dilation_rows
                dimensions.dilation_cols}},  // dilation_cols
              {{dimensions.stride_rows,      // stride_rows
                dimensions.stride_cols}},    // stride_cols
              {{padding_rows,                // padding_rows
                padding_cols}},              // padding_cols
              dtype,                         // tensor datatype
              device_id,                     // device_id
          },
          dnn_activation_mode  // activation_mode
      };
  
      // Launch fused convolution with given parameters and scratch allocator.
      // Record profile result into `profile_result` if it's not nullptr.
      const auto launch = [&](se::dnn::AlgorithmConfig algorithm_config,
                              DnnScratchAllocator* scratch_allocator,
                              se::dnn::ProfileResult* profile_result) -> bool {
        return stream
            ->ThenFusedConvolveWithAlgorithm(
                input_desc, input_ptr,                     // input
                /*conv_input_scale=*/1.0,                  // input_scale
                filter_desc, filter_ptr,                   // filter
                conv_desc,                                 // conv
                side_input_ptr, /*side_input_scale=*/0.0,  // side_input
                bias_desc, bias_ptr,                       // bias
                dnn_activation_mode,                       // activation
                output_desc, &output_ptr,                  // output
                scratch_allocator, algorithm_config, profile_result)
            .ok();
      };
  
      se::dnn::AlgorithmConfig algorithm_config;
      if (cudnn_use_autotune) {
        auto status = FindBestConvolveAlgorithm<T>(
            conv_parameters, launch, context, stream,
            [&](absl::Span<const tensorflow::AutotuneResult> results) {
              LogFusedConvAutotuneResults(
                  context->op_kernel().def(), input, transformed_filter,
                  transformed_output, bias, nullptr, stream->parent(), results);
            },
            &algorithm_config);
        OP_REQUIRES_OK(context, status);
      }
  
      DnnScratchAllocator scratch_allocator(ConvolveScratchSize(), context);
      bool cudnn_launch_status = launch(algorithm_config, &scratch_allocator,
                                        /*profile_result=*/nullptr);
      OP_REQUIRES(
          context, cudnn_launch_status,
          errors::Internal(absl::Substitute(
              "cuDNN launch failure: input shape($0) filter shape($1)",
              input.shape().DebugString(), filter.shape().DebugString())));
  
      // Convert the output tensor back from NCHW to NHWC.
      if (params.data_format == FORMAT_NHWC) {
        functor::NCHWToNHWC<GPUDevice, T, 4>()(
            context->eigen_device<GPUDevice>(),
            const_cast<const Tensor&>(transformed_output).tensor<T, 4>(),
            output->tensor<T, 4>());
      }
    }
  };
  
  #endif  // GOOGLE_CUDA
  
  template <typename Device, typename T>
  class FusedConv2DOp : public OpKernel {
   public:
    explicit FusedConv2DOp(OpKernelConstruction* context) : OpKernel(context) {
      OP_REQUIRES_OK(context, InitConv2DParameters(context, &params_));
  
      OP_REQUIRES_OK(context, context->GetAttr("use_cudnn_on_gpu", &use_cudnn_));
      use_cudnn_ &= CanUseCudnn();
      cudnn_use_autotune_ = CudnnUseAutotune();
  
      // 'fused_ops' and 'num_args' attributes are specified by the Grappler
      // Remapper optimizer (see grappler/optimizers/remapper.cc).
  
      std::vector<string> fused_ops;
      OP_REQUIRES_OK(context, context->GetAttr("fused_ops", &fused_ops));
      OP_REQUIRES(context, !fused_ops.empty(),
                  errors::InvalidArgument(
                      "Fused Conv2D must have at least one fused op."));
  
      int num_args;
      OP_REQUIRES_OK(context, context->GetAttr("num_args", &num_args));
  
      // TODO(ezhulenev): Add support for fusion element-wise op chains defined
      // at runtime, e.g. Relu+Sqrt+Tanh+etc.
  
      using FCT = FusedComputationType;
      std::vector<std::pair<FusedConv2DPattern, FusedComputationType>> mappings =
          {{{{"BiasAdd"}, true}, FCT::kBiasAdd},
           {{{"BiasAdd", "Relu"}, false}, FCT::kBiasAddWithRelu},
           {{{"BiasAdd", "Relu6"}, true}, FCT::kBiasAddWithRelu6},
           {{{"BiasAdd", "Elu"}, true}, FCT::kBiasAddWithElu},
           {{{"FusedBatchNorm"}, true}, FCT::kFusedBatchNorm},
           {{{"FusedBatchNorm", "Relu"}, true}, FCT::kFusedBatchNormWithRelu},
           {{{"FusedBatchNorm", "Relu6"}, true}, FCT::kFusedBatchNormWithRelu6},
           {{{"FusedBatchNorm", "Elu"}, true}, FCT::kFusedBatchNormWithElu}};
  
      // Match op fusion to one of hte supported patterns.
      for (const auto& mapping : mappings) {
        const FusedConv2DPattern& pattern = mapping.first;
        if (FusedOpsMatchAndSupportedOnDevice(fused_ops, pattern.fused_ops,
                                              pattern.cpu_only)) {
          fused_computation_ = mapping.second;
        }
      }
      if (fused_computation_ == FusedComputationType::kUndefined) {
        OP_REQUIRES(context, false,
                    errors::Unimplemented("Fusion is not implemented: [",
                                          absl::StrJoin(fused_ops, ","), "]"));
      }
  
      // Depending on a picked fusion type validate fusion-specific arguments.
      if (fused_computation_ == FusedComputationType::kBiasAdd ||
          fused_computation_ == FusedComputationType::kBiasAddWithRelu ||
          fused_computation_ == FusedComputationType::kBiasAddWithRelu6 ||
          fused_computation_ == FusedComputationType::kBiasAddWithElu) {
        OP_REQUIRES(context, num_args == 1,
                    errors::InvalidArgument(
                        "Fused Conv2D must have one extra argument: bias."));
      }
  
      if (fused_computation_ == FusedComputationType::kFusedBatchNorm ||
          fused_computation_ == FusedComputationType::kFusedBatchNormWithRelu ||
          fused_computation_ == FusedComputationType::kFusedBatchNormWithRelu6 ||
          fused_computation_ == FusedComputationType::kFusedBatchNormWithElu) {
        OP_REQUIRES(
            context, num_args == 4,
            errors::InvalidArgument("Fused FusedBatchNorm must have four extra "
                                    "arguments: scale, offset, mean, variance."));
        OP_REQUIRES_OK(context, context->GetAttr("epsilon", &epsilon_));
      }
    }
  
    void Compute(OpKernelContext* context) override {
      // Input tensor is of the following dimensions:
      // [ batch, in_rows, in_cols, in_depth ]
      const Tensor& input = context->input(0);
  
      // Input filter is of the following dimensions:
      // [ filter_rows, filter_cols, in_depth, out_depth]
      const Tensor& filter = context->input(1);
  
      Conv2DDimensions dimensions;
      OP_REQUIRES_OK(context,
                     ComputeConv2DDimension(params_, input, filter, &dimensions));
  
      TensorShape out_shape = ShapeFromFormat(
          params_.data_format, dimensions.batch, dimensions.out_rows,
          dimensions.out_cols, dimensions.out_depth);
  
      // Output tensor is of the following dimensions:
      // [ in_batch, out_rows, out_cols, out_depth ]
      Tensor* output = nullptr;
      OP_REQUIRES_OK(context, context->allocate_output(0, out_shape, &output));
  
      VLOG(2) << "FusedConv2D: in_depth = " << dimensions.in_depth
              << ", patch_depth = " << dimensions.patch_depth
              << ", input_cols = " << dimensions.input_cols
              << ", filter_cols = " << dimensions.filter_cols
              << ", input_rows = " << dimensions.input_rows
              << ", filter_rows = " << dimensions.filter_rows
              << ", stride_rows = " << dimensions.stride_rows
              << ", stride_cols = " << dimensions.stride_cols
              << ", dilation_rows = " << dimensions.dilation_rows
              << ", dilation_cols = " << dimensions.dilation_cols
              << ", out_depth = " << dimensions.out_depth;
  
      // If there is nothing to compute, return.
      if (out_shape.num_elements() == 0) {
        return;
      }
  
      FusedComputationArgs args;
      args.epsilon = epsilon_;
  
      LaunchFusedConv2DOp<Device, T>()(context, use_cudnn_, cudnn_use_autotune_,
                                       input, filter, fused_computation_, args,
                                       params_, dimensions, output);
    }
  
   private:
    struct FusedConv2DPattern {
      std::vector<string> fused_ops;
      bool cpu_only;
    };
  
    bool FusedOpsMatchAndSupportedOnDevice(const std::vector<string>& fused_ops,
                                           const std::vector<string>& expected,
                                           bool cpu_only) const {
      if (std::is_same<Device, GPUDevice>::value && cpu_only) {
        return false;
      }
      return fused_ops == expected;
    }
  
    Conv2DParameters params_;
    bool use_cudnn_;
    bool cudnn_use_autotune_;
  
    FusedComputationType fused_computation_ = FusedComputationType::kUndefined;
  
    float epsilon_;  // Used only in FusedBatchNorm fusion
  
    TF_DISALLOW_COPY_AND_ASSIGN(FusedConv2DOp);
  };
  
  // Registration of the CPU implementations.
  #define REGISTER_FUSED_CPU_CONV2D(T)                                  \
    REGISTER_KERNEL_BUILDER(                                            \
        Name("_FusedConv2D").Device(DEVICE_CPU).TypeConstraint<T>("T"), \
        FusedConv2DOp<CPUDevice, T>);
  
  #if GOOGLE_CUDA
  
  #define DECLARE_FUNCTOR_GPU_SPEC(T)                                      \
    template <>                                                            \
    void TransformFilter<GPUDevice, T, int, 4>::operator()(                \
        const GPUDevice& d, FilterTensorFormat dst_filter_format,          \
        typename TTypes<T, 4, int>::ConstTensor in,                        \
        typename TTypes<T, 4, int>::Tensor out);                           \
    extern template struct TransformFilter<GPUDevice, T, int, 4>;          \
    template <>                                                            \
    void PadInput<GPUDevice, T, int, 4>::operator()(                       \
        const GPUDevice& d, typename TTypes<T, 4, int>::ConstTensor in,    \
        const std::array<int, 2>& padding_left,                            \
        const std::array<int, 2>& padding_right,                           \
        typename TTypes<T, 4, int>::Tensor out, TensorFormat data_format); \
    extern template struct PadInput<GPUDevice, T, int, 4>
  
  // Registration of the GPU implementations.
  #define REGISTER_FUSED_GPU_CONV2D(T)                                  \
    REGISTER_KERNEL_BUILDER(                                            \
        Name("_FusedConv2D").Device(DEVICE_GPU).TypeConstraint<T>("T"), \
        FusedConv2DOp<GPUDevice, T>);
  
  #endif  // GOOGLE_CUDA
  
  }  // namespace tensorflow
  
  #endif  // TENSORFLOW_CORE_KERNELS_CONV_OPS_FUSED_IMPL_H_
  