cudnn_bn_stats_finalize.cu.h

// Copyright (c) 2024 PaddlePaddle 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.

#pragma once

#include "paddle/phi/backends/gpu/gpu_dnn.h"
#include "paddle/phi/kernels/fusion/gpu/cudnn_fusion_helper.h"

namespace phi {
namespace fusion {

template <typename T>
using BatchNormParamType =
    typename backends::gpu::CudnnDataType<T>::BatchNormParamType;

#if CUDNN_VERSION >= 8000

template <typename T>
struct BNStatsFinalizeArgs {
  BNStatsFinalizeArgs() {
    dtype = backends::gpu::CudnnDataType<T>::type;
    param_dtype = backends::gpu::CudnnDataType<BatchNormParamType<T>>::type;
    format = CUDNN_TENSOR_NHWC;
  }

  void Set(const std::vector<int> &param_shape) {
    PADDLE_ENFORCE_EQ(
        param_shape.size(),
        4U,
        common::errors::InvalidArgument(
            "The size of param_shape is expected to 4. But received "
            "param_shape's size is %d, param_shape is [%s].",
            param_shape.size(),
            make_ddim(param_shape)));

    in_desc.set(param_shape, format, param_dtype);
    out_desc.set(param_shape, format, dtype);
  }

  cudnnDataType_t dtype;
  cudnnDataType_t param_dtype;
  cudnnTensorFormat_t format;

  backends::gpu::TensorDescriptor in_desc;
  backends::gpu::TensorDescriptor out_desc;
};

template <typename T>
class CudnnBNStatsFinalize {
 public:
  CudnnBNStatsFinalize(const GPUContext &dev_ctx,
                       const std::vector<int> &param_shape)
      : train_op_(CUDNN_FUSED_BN_FINALIZE_STATISTICS_TRAINING),
        inference_op_(CUDNN_FUSED_BN_FINALIZE_STATISTICS_INFERENCE) {
    args_.Set(param_shape);
  }
  ~CudnnBNStatsFinalize() {}

  void Forward(const GPUContext &dev_ctx,
               const DenseTensor &sum,
               const DenseTensor &sum_of_squares,
               const DenseTensor &scale,
               const DenseTensor &bias,
               DenseTensor *saved_mean,
               DenseTensor *saved_invstd,
               DenseTensor *running_mean,
               DenseTensor *running_var,
               DenseTensor *equiv_scale,
               DenseTensor *equiv_bias,
               double eps,
               float momentum,
               int64_t ele_count,
               bool is_train) {
    if (is_train) {
      TrainInit(dev_ctx);
    } else {
      InferenceInit(dev_ctx);
    }
    auto &op = is_train ? train_op_ : inference_op_;

    // Set variant_param for both inference_op_ and train_op_
    float *sum_ptr = const_cast<float *>(sum.data<float>());
    float *sum_of_squares_ptr =
        const_cast<float *>(sum_of_squares.data<float>());
    float *scale_ptr = const_cast<float *>(scale.data<float>());
    float *bias_ptr = const_cast<float *>(bias.data<float>());
    float *saved_mean_ptr = dev_ctx.template Alloc<float>(
        saved_mean, saved_mean->numel() * sizeof(float));
    float *saved_invstd_ptr = dev_ctx.template Alloc<float>(
        saved_invstd, saved_invstd->numel() * sizeof(float));
    float *running_mean_ptr = dev_ctx.template Alloc<float>(
        running_mean, running_mean->numel() * sizeof(float));
    float *running_var_ptr = dev_ctx.template Alloc<float>(
        running_var, running_var->numel() * sizeof(float));
    T *equiv_scale_ptr = dev_ctx.template Alloc<T>(
        equiv_scale, equiv_scale->numel() * sizeof(T));
    T *equiv_bias_ptr =
        dev_ctx.template Alloc<T>(equiv_bias, equiv_bias->numel() * sizeof(T));
    op.SetOpVariantParamAttrPtr(CUDNN_PTR_BN_SCALE, scale_ptr);
    op.SetOpVariantParamAttrPtr(CUDNN_PTR_BN_BIAS, bias_ptr);
    op.SetOpVariantParamAttrPtr(CUDNN_PTR_BN_RUNNING_MEAN, running_mean_ptr);
    op.SetOpVariantParamAttrPtr(CUDNN_PTR_BN_RUNNING_VAR, running_var_ptr);
    op.SetOpVariantParamAttrPtr(CUDNN_PTR_BN_EQSCALE, equiv_scale_ptr);
    op.SetOpVariantParamAttrPtr(CUDNN_PTR_BN_EQBIAS, equiv_bias_ptr);
    op.SetOpVariantParamAttrPtr<double>(CUDNN_SCALAR_DOUBLE_BN_EPSILON, &eps);

    // Set extra variant_param only for train_op_:
    if (is_train) {
      op.SetOpVariantParamAttrPtr(CUDNN_PTR_YSUM, sum_ptr);
      op.SetOpVariantParamAttrPtr(CUDNN_PTR_YSQSUM, sum_of_squares_ptr);
      op.SetOpVariantParamAttrPtr(CUDNN_PTR_BN_SAVED_MEAN, saved_mean_ptr);
      op.SetOpVariantParamAttrPtr(CUDNN_PTR_BN_SAVED_INVSTD, saved_invstd_ptr);
      double avg_factor = 1.0 - momentum;
      op.SetOpVariantParamAttrPtr(CUDNN_SCALAR_INT64_T_BN_ACCUMULATION_COUNT,
                                  &ele_count);
      op.SetOpVariantParamAttrPtr(CUDNN_SCALAR_DOUBLE_BN_EXP_AVG_FACTOR,
                                  &avg_factor);
    }
    // fused op execute
    auto handle = dev_ctx.cudnn_handle();
    op.Execute(handle);
  }

 private:
  void TrainInit(const GPUContext &dev_ctx) {
    // Set constant_param for train op
    train_op_.SetOpConstParamAttr({CUDNN_PARAM_YSUM_PLACEHOLDER,
                                   CUDNN_PARAM_YSQSUM_PLACEHOLDER,
                                   CUDNN_PARAM_BN_SCALE_PLACEHOLDER,
                                   CUDNN_PARAM_BN_BIAS_PLACEHOLDER,
                                   CUDNN_PARAM_BN_SAVED_MEAN_PLACEHOLDER,
                                   CUDNN_PARAM_BN_SAVED_INVSTD_PLACEHOLDER,
                                   CUDNN_PARAM_BN_RUNNING_MEAN_PLACEHOLDER,
                                   CUDNN_PARAM_BN_RUNNING_VAR_PLACEHOLDER,
                                   CUDNN_PARAM_BN_EQSCALE_PLACEHOLDER,
                                   CUDNN_PARAM_BN_EQBIAS_PLACEHOLDER},
                                  CUDNN_PTR_16B_ALIGNED);
    // Set input and output desc for train op
    train_op_.SetOpConstParamDesc(
        {CUDNN_PARAM_YSTATS_DESC, CUDNN_PARAM_BN_SCALEBIAS_MEANVAR_DESC},
        args_.in_desc.desc());
    train_op_.SetOpConstParamDesc(CUDNN_PARAM_BN_EQSCALEBIAS_DESC,
                                  args_.out_desc.desc());

    // Get workspace
    auto handle = dev_ctx.cudnn_handle();
    train_op_.SetOpConstParamAttr(CUDNN_PARAM_BN_MODE,
                                  CUDNN_BATCHNORM_SPATIAL_PERSISTENT);
    // Check workspace size, also creates plan.
    size_t workspace_size_bytes = train_op_.GetWorkspaceSizeInBytes(handle);
    PADDLE_ENFORCE_EQ(workspace_size_bytes,
                      0U,
                      common::errors::InvalidArgument(
                          "Unexpected non-zero workspace size for "
                          "CudnnBNStatsFinalize."));
    train_op_.SetOpVariantParamAttrPtr(CUDNN_PTR_WORKSPACE,
                                       static_cast<void *>(nullptr));
    train_op_.SetOpVariantParamAttrPtr(CUDNN_PTR_WORKSPACE,
                                       &workspace_size_bytes);
  }

  void InferenceInit(const GPUContext &dev_ctx) {
    // Set constant_param for inference op
    inference_op_.SetOpConstParamAttr({CUDNN_PARAM_BN_SCALE_PLACEHOLDER,
                                       CUDNN_PARAM_BN_BIAS_PLACEHOLDER,
                                       CUDNN_PARAM_BN_RUNNING_MEAN_PLACEHOLDER,
                                       CUDNN_PARAM_BN_RUNNING_VAR_PLACEHOLDER,
                                       CUDNN_PARAM_BN_EQSCALE_PLACEHOLDER,
                                       CUDNN_PARAM_BN_EQBIAS_PLACEHOLDER},
                                      CUDNN_PTR_16B_ALIGNED);
    // Set input and output desc for inference op
    inference_op_.SetOpConstParamDesc(CUDNN_PARAM_BN_SCALEBIAS_MEANVAR_DESC,
                                      args_.in_desc.desc());
    inference_op_.SetOpConstParamDesc(CUDNN_PARAM_BN_EQSCALEBIAS_DESC,
                                      args_.out_desc.desc());

    // Get workspace
    auto handle = dev_ctx.cudnn_handle();
    inference_op_.SetOpConstParamAttr(CUDNN_PARAM_BN_MODE,
                                      CUDNN_BATCHNORM_SPATIAL_PERSISTENT);
    // Check workspace size, also creates plan.
    size_t workspace_size_bytes = inference_op_.GetWorkspaceSizeInBytes(handle);
    PADDLE_ENFORCE_EQ(workspace_size_bytes,
                      0U,
                      common::errors::InvalidArgument(
                          "Unexpected non-zero workspace size for "
                          "CudnnBNStatsFinalize."));
    inference_op_.SetOpVariantParamAttrPtr(CUDNN_PTR_WORKSPACE,
                                           static_cast<void *>(nullptr));
    inference_op_.SetOpVariantParamAttrPtr(CUDNN_PTR_WORKSPACE,
                                           &workspace_size_bytes);
  }

  BNStatsFinalizeArgs<T> args_;
  CudnnFusionOp train_op_;
  CudnnFusionOp inference_op_;
};
#endif
}  // namespace fusion
}  // namespace phi