conv_fusion_op.cu 24.2 KB
Newer Older
Q
qingqing01 已提交
1 2 3 4 5 6 7 8 9 10 11 12 13 14
/* Copyright (c) 2016 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. */

15
#include <array>
16

17
#include "paddle/fluid/framework/conv_search_cache.h"
Q
qingqing01 已提交
18
#include "paddle/fluid/framework/op_registry.h"
19
#include "paddle/fluid/operators/conv_op.h"
20
#include "paddle/fluid/platform/device/gpu/gpu_dnn.h"
21
#include "paddle/phi/kernels/funcs/padding.h"
22
#include "paddle/phi/kernels/gpudnn/conv_gpudnn_info.h"
Q
qingqing01 已提交
23

24
DECLARE_int64(cudnn_exhaustive_search_times);
Q
qingqing01 已提交
25 26 27 28

namespace paddle {
namespace operators {

R
ronnywang 已提交
29
#if PADDLE_WITH_HIP || CUDNN_VERSION >= 7100
Q
qingqing01 已提交
30 31 32 33 34
using ScopedTensorDescriptor = platform::ScopedTensorDescriptor;
using ScopedFilterDescriptor = platform::ScopedFilterDescriptor;
using ScopedConvolutionDescriptor = platform::ScopedConvolutionDescriptor;
using ScopedActivationDescriptor = platform::ScopedActivationDescriptor;
using DataLayout = platform::DataLayout;
35
using framework::AlgorithmsCache;
36
using framework::ConvSearchCache;
X
xiaoxiaohehe001 已提交
37
using framework::SearchFuseResult;
38

Q
qingqing01 已提交
39 40 41 42 43 44 45
template <typename T>
using ScalingParamType = typename platform::CudnnDataType<T>::ScalingParamType;

template <typename T>
class CUDNNConvFusionOpKernel : public framework::OpKernel<T> {
 public:
  void Compute(const framework::ExecutionContext& ctx) const override {
L
Leo Chen 已提交
46
    auto& dev_ctx = ctx.template device_context<phi::GPUContext>();
47 48 49 50 51
    auto* input = ctx.Input<phi::DenseTensor>("Input");
    auto* filter = ctx.Input<phi::DenseTensor>("Filter");
    auto* bias = ctx.Input<phi::DenseTensor>("Bias");
    auto* residual = ctx.Input<phi::DenseTensor>("ResidualData");
    auto* output = ctx.Output<phi::DenseTensor>("Output");
52
    dev_ctx.template Alloc<T>(output, output->numel() * sizeof(T));
Q
qingqing01 已提交
53 54 55 56 57

    std::vector<int> strides = ctx.Attr<std::vector<int>>("strides");
    std::vector<int> paddings = ctx.Attr<std::vector<int>>("paddings");
    std::vector<int> dilations = ctx.Attr<std::vector<int>>("dilations");
    const std::string activation = ctx.Attr<std::string>("activation");
58 59 60 61 62 63 64 65 66
    std::string data_format = ctx.Attr<std::string>("data_format");
    PADDLE_ENFORCE_NE(
        data_format,
        "NHWC",
        platform::errors::PermissionDenied(
            "Operator(Conv2DFusion) in cuDNN only supports data format of "
            "channel first (NCHW) now. But received: data_format = '%s'.",
            data_format));

Q
qingqing01 已提交
67 68 69 70 71 72 73 74
    int groups = ctx.Attr<int>("groups");
    int64_t user_workspace_size =
        static_cast<size_t>(ctx.Attr<int>("workspace_size_MB"));
    bool exhaustive_search =
        FLAGS_cudnn_exhaustive_search || ctx.Attr<bool>("exhaustive_search");

    const T* filter_data = filter->data<T>();
    const T* bias_data = bias->data<T>();
75 76 77 78

    const std::string padding_algorithm =
        ctx.Attr<std::string>("padding_algorithm");

79 80
    phi::DenseTensor transformed_input_channel(input->dtype());
    phi::DenseTensor transformed_output(output->dtype());
81 82
    transformed_input_channel = *input;
    transformed_output = *output;
83 84
    T* output_data = transformed_output.data<T>();

Q
qingqing01 已提交
85
    const T* residual_data = residual ? residual->data<T>() : output_data;
86

87 88 89
    // update padding and dilation
    auto in_dims = transformed_input_channel.dims();
    auto filter_dims = filter->dims();
90
    framework::DDim in_data_dims = phi::slice_ddim(in_dims, 2, in_dims.size());
91 92

    framework::DDim filter_data_dims =
93 94
        phi::slice_ddim(filter_dims, 2, filter_dims.size());
    std::vector<int> ksize = phi::vectorize<int>(filter_data_dims);
95 96
    UpdatePaddingAndDilation(
        &paddings, &dilations, padding_algorithm, in_data_dims, strides, ksize);
97 98

    int data_dim = strides.size();  // 2d or 3d
99
    bool is_sys_pad = phi::funcs::IsSymmetricPadding(paddings, data_dim);
100

101
    phi::DenseTensor transformed_input;
102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118
    std::vector<int> padding_common(data_dim, 0);
    if (!is_sys_pad) {
      std::vector<int> padding_diff(data_dim);
      std::vector<int> new_input_shape_vec(data_dim + 2);
      new_input_shape_vec[0] = transformed_input_channel.dims()[0];
      new_input_shape_vec[1] = transformed_input_channel.dims()[1];

      std::vector<int> input_pad(transformed_input_channel.dims().size() * 2,
                                 0);
      for (size_t i = 0; i < data_dim; ++i) {
        padding_diff[i] = std::abs(paddings[2 * i] - paddings[2 * i + 1]);
        padding_common[i] = std::min(paddings[2 * i], paddings[2 * i + 1]);
        new_input_shape_vec[i + 2] =
            transformed_input_channel.dims()[i + 2] + padding_diff[i];
        input_pad[2 * i + 4] = paddings[2 * i] - padding_common[i];
        input_pad[2 * i + 4 + 1] = paddings[2 * i + 1] - padding_common[i];
      }
119
      framework::DDim new_input_shape(phi::make_ddim(new_input_shape_vec));
120
      transformed_input.Resize(new_input_shape);
L
Leo Chen 已提交
121
      auto& dev_ctx = ctx.template device_context<phi::GPUContext>();
122 123

      transformed_input =
L
Leo Chen 已提交
124
          ctx.AllocateTmpTensor<T, phi::GPUContext>(new_input_shape, dev_ctx);
125 126 127 128
      const int rank = transformed_input_channel.dims().size();
      T pad_value(0.0);
      switch (rank) {
        case 4: {
L
Leo Chen 已提交
129
          phi::funcs::PadFunction<phi::GPUContext, T, 4>(
130 131 132 133
              dev_ctx,
              input_pad,
              transformed_input_channel,
              pad_value,
134 135 136
              &transformed_input);
        } break;
        case 5: {
L
Leo Chen 已提交
137
          phi::funcs::PadFunction<phi::GPUContext, T, 5>(
138 139 140 141
              dev_ctx,
              input_pad,
              transformed_input_channel,
              pad_value,
142 143 144
              &transformed_input);
        } break;
        default:
145
          PADDLE_THROW(platform::errors::PermissionDenied(
146 147
              "Operator Conv2DFusion expects Input to be a 4-D or 5-D "
              "phi::DenseTensor. "
148
              "But received the actual dimension = %d, shape = [%s].",
149 150
              rank,
              transformed_input_channel.dims()));
151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166
      }

    } else {
      transformed_input = transformed_input_channel;
      if (paddings.size() == data_dim) {
        for (size_t i = 0; i < data_dim; ++i) {
          padding_common[i] = paddings[i];
        }
      } else {
        for (size_t i = 0; i < data_dim; ++i) {
          padding_common[i] = paddings[2 * i];
        }
      }
    }

    const T* input_data = transformed_input.data<T>();
Q
qingqing01 已提交
167 168 169 170 171 172 173 174 175 176 177 178

    // ------------------- cudnn descriptors ---------------------
    ScopedTensorDescriptor input_desc;
    ScopedTensorDescriptor output_desc;
    ScopedFilterDescriptor filter_desc;
    ScopedTensorDescriptor bias_desc;
    ScopedConvolutionDescriptor conv_desc;
    ScopedActivationDescriptor act_desc;
    DataLayout layout = DataLayout::kNCHW;
    if (input->dims().size() == 5) {
      layout = DataLayout::kNCDHW;
    }
R
ronnywang 已提交
179 180 181
#ifdef PADDLE_WITH_HIP
    miopenConvolutionDescriptor_t cudnn_conv_desc =
        conv_desc.descriptor<T>(padding_common, strides, dilations);
182
    PADDLE_ENFORCE_GPU_SUCCESS(
R
ronnywang 已提交
183 184 185 186 187 188
        platform::dynload::miopenSetConvolutionGroupCount(cudnn_conv_desc,
                                                          groups));
    // Now only support NCHW
    std::vector<int> bias_dim = {
        1, static_cast<int>(transformed_output.dims()[1]), 1, 1};
    miopenTensorDescriptor_t cudnn_input_desc = input_desc.descriptor<T>(
189
        layout, phi::vectorize<int>(transformed_input.dims()));
R
ronnywang 已提交
190
    miopenTensorDescriptor_t cudnn_output_desc = output_desc.descriptor<T>(
191
        layout, phi::vectorize<int>(transformed_output.dims()));
192
    miopenTensorDescriptor_t cudnn_filter_desc =
193
        filter_desc.descriptor<T>(layout, phi::vectorize<int>(filter->dims()));
R
ronnywang 已提交
194 195 196 197
    miopenTensorDescriptor_t cudnn_bias_desc =
        bias_desc.descriptor<T>(layout, bias_dim);
    miopenActivationDescriptor_t cudnn_act_desc =
        act_desc.descriptor<T>(activation);
Q
qingqing01 已提交
198

R
ronnywang 已提交
199 200 201 202
    miopenConvFwdAlgorithm_t algo;
    auto handle = dev_ctx.cudnn_handle();
    auto workspace_handle = dev_ctx.cudnn_workspace_handle();

203 204
    auto x_dims = phi::vectorize(transformed_input.dims());
    auto f_dims = phi::vectorize(filter->dims());
R
ronnywang 已提交
205 206

    size_t workspace_size = 0;
207
    PADDLE_ENFORCE_GPU_SUCCESS(
R
ronnywang 已提交
208
        platform::dynload::miopenConvolutionForwardGetWorkSpaceSize(
209 210 211 212 213 214
            handle,
            cudnn_filter_desc,
            cudnn_input_desc,
            cudnn_conv_desc,
            cudnn_output_desc,
            &workspace_size));
R
ronnywang 已提交
215 216 217
    int find_count;
    miopenConvAlgoPerf_t find_result;
    auto cudnn_find_func = [&](void* cudnn_workspace_ptr) {
218
      PADDLE_ENFORCE_GPU_SUCCESS(
R
ronnywang 已提交
219
          platform::dynload::miopenFindConvolutionForwardAlgorithm(
220 221 222 223 224 225 226 227
              handle,
              cudnn_input_desc,
              input_data,
              cudnn_filter_desc,
              filter_data,
              cudnn_conv_desc,
              cudnn_output_desc,
              output_data,
228
              phi::kNUM_CUDNN_FWD_ALGS,
229 230 231 232 233
              &find_count,
              &find_result,
              cudnn_workspace_ptr,
              workspace_size,
              false));
R
ronnywang 已提交
234 235 236 237 238 239 240 241
    };
    workspace_handle.RunFuncSync(cudnn_find_func, workspace_size);
    algo = find_result.fwd_algo;
    VLOG(3) << "cuDNN forward algo " << algo;

    {
      ScalingParamType<T> alpha = 1.0f, beta = 0.0f;
      auto cudnn_func = [&](void* cudnn_workspace) {
242 243 244 245 246 247 248 249 250 251 252 253 254 255
        PADDLE_ENFORCE_GPU_SUCCESS(
            platform::dynload::miopenConvolutionForward(handle,
                                                        &alpha,
                                                        cudnn_input_desc,
                                                        input_data,
                                                        cudnn_filter_desc,
                                                        filter_data,
                                                        cudnn_conv_desc,
                                                        algo,
                                                        &beta,
                                                        cudnn_output_desc,
                                                        output_data,
                                                        cudnn_workspace,
                                                        workspace_size));
R
ronnywang 已提交
256 257
      };
      workspace_handle.RunFunc(cudnn_func, workspace_size);
258
      PADDLE_ENFORCE_GPU_SUCCESS(
259 260 261 262 263 264 265
          platform::dynload::miopenConvolutionForwardBias(handle,
                                                          &alpha,
                                                          cudnn_bias_desc,
                                                          bias_data,
                                                          &beta,
                                                          cudnn_output_desc,
                                                          output_data));
R
ronnywang 已提交
266
      if (activation != "identity") {
267 268 269 270 271 272 273 274 275
        PADDLE_ENFORCE_GPU_SUCCESS(
            platform::dynload::miopenActivationForward(handle,
                                                       cudnn_act_desc,
                                                       &alpha,
                                                       cudnn_output_desc,
                                                       output_data,
                                                       &beta,
                                                       cudnn_output_desc,
                                                       output_data));
R
ronnywang 已提交
276 277
      }
      if (residual) {
278 279 280 281 282 283 284 285 286 287 288 289
        PADDLE_ENFORCE_GPU_SUCCESS(
            platform::dynload::miopenOpTensor(handle,
                                              miopenTensorOpAdd,
                                              &alpha,
                                              cudnn_output_desc,
                                              output_data,
                                              &alpha,
                                              cudnn_output_desc,
                                              residual_data,
                                              &beta,
                                              cudnn_output_desc,
                                              output_data));
R
ronnywang 已提交
290 291 292
      }
    }
#else  // PADDLE_WITH_HIP
Q
qingqing01 已提交
293
    cudnnConvolutionDescriptor_t cudnn_conv_desc =
294
        conv_desc.descriptor<T>(padding_common, strides, dilations);
295 296
    PADDLE_ENFORCE_GPU_SUCCESS(platform::dynload::cudnnSetConvolutionGroupCount(
        cudnn_conv_desc, groups));
Q
qingqing01 已提交
297 298

    cudnnTensorDescriptor_t cudnn_input_desc = input_desc.descriptor<T>(
299
        layout, phi::vectorize<int>(transformed_input.dims()));
Q
qingqing01 已提交
300
    cudnnTensorDescriptor_t cudnn_output_desc = output_desc.descriptor<T>(
301
        layout, phi::vectorize<int>(transformed_output.dims()));
302
    cudnnFilterDescriptor_t cudnn_filter_desc =
303
        filter_desc.descriptor<T>(layout, phi::vectorize<int>(filter->dims()));
Q
qingqing01 已提交
304
    // Now only support NCHW
305 306
    std::vector<int> bias_dim = {
        1, static_cast<int>(transformed_output.dims()[1]), 1, 1};
Q
qingqing01 已提交
307 308 309 310 311 312 313
    cudnnTensorDescriptor_t cudnn_bias_desc =
        bias_desc.descriptor<T>(layout, bias_dim);
    cudnnActivationDescriptor_t cudnn_act_desc =
        act_desc.descriptor<T>(activation);

    // ------------------- cudnn conv workspace ---------------------
    size_t workspace_size_in_bytes;  // final workspace to allocate.
314
    size_t workspace_size_limit = 0;
Q
qingqing01 已提交
315 316
    if (FLAGS_conv_workspace_size_limit > 0 || user_workspace_size > 0) {
      int64_t max_user_size =
317
          std::min(static_cast<int64_t>(FLAGS_conv_workspace_size_limit),
Q
qingqing01 已提交
318 319 320 321 322 323 324
                   user_workspace_size);
      workspace_size_limit = max_user_size * 1024 * 1024;
    }

    // ------------------- cudnn conv algorithm ---------------------
    cudnnConvolutionFwdAlgo_t algo;
    auto handle = dev_ctx.cudnn_handle();
C
chengduo 已提交
325
    auto workspace_handle = dev_ctx.cudnn_workspace_handle();
X
xiaoxiaohehe001 已提交
326
    auto dtype = platform::CudnnDataType<T>::type;
Q
qingqing01 已提交
327

328
    PADDLE_ENFORCE_GPU_SUCCESS(platform::dynload::cudnnSetConvolutionMathType(
329
        cudnn_conv_desc, CUDNN_DEFAULT_MATH));
X
xiaoxiaohehe001 已提交
330 331 332 333
    if (dtype == CUDNN_DATA_HALF) {
      PADDLE_ENFORCE_GPU_SUCCESS(platform::dynload::cudnnSetConvolutionMathType(
          cudnn_conv_desc, CUDNN_TENSOR_OP_MATH));
    }
A
AshburnLee 已提交
334
#if CUDA_VERSION >= 11000 && CUDNN_VERSION >= 8000
A
AshburnLee 已提交
335
    if (!platform::allow_tf32_cudnn) {
336 337
      PADDLE_ENFORCE_GPU_SUCCESS(platform::dynload::cudnnSetConvolutionMathType(
          cudnn_conv_desc, CUDNN_FMA_MATH));
A
AshburnLee 已提交
338
    }
A
AshburnLee 已提交
339
#endif  // CUDA_VERSION >= 11000 && CUDNN_VERSION >= 8000
Q
qingqing01 已提交
340

341 342
    auto x_dims = phi::vectorize(transformed_input.dims());
    auto f_dims = phi::vectorize(filter->dims());
343
    if (!exhaustive_search) {
344
#if CUDNN_VERSION >= 8000
345 346 347 348
      int perf_count;
      int best_algo_idx = 0;
      size_t tmp_size = 0;
      std::unique_ptr<cudnnConvolutionFwdAlgoPerf_t[]> perf_results(
349
          new cudnnConvolutionFwdAlgoPerf_t[phi::kNUM_CUDNN_FWD_ALGS]);
350
      PADDLE_ENFORCE_GPU_SUCCESS(
351
          platform::dynload::cudnnGetConvolutionForwardAlgorithm_v7(
352 353 354 355 356
              handle,
              cudnn_input_desc,
              cudnn_filter_desc,
              cudnn_conv_desc,
              cudnn_output_desc,
357
              phi::kNUM_CUDNN_FWD_ALGS,
358
              &perf_count,
359 360
              perf_results.get()));
      algo = (perf_results.get())[best_algo_idx].algo;
X
xiaoxiaohehe001 已提交
361
#else
362
      PADDLE_ENFORCE_GPU_SUCCESS(
X
xiaoxiaohehe001 已提交
363
          platform::dynload::cudnnGetConvolutionForwardAlgorithm(
364 365 366 367 368
              handle,
              cudnn_input_desc,
              cudnn_filter_desc,
              cudnn_conv_desc,
              cudnn_output_desc,
X
xiaoxiaohehe001 已提交
369 370 371 372
              CUDNN_CONVOLUTION_FWD_SPECIFY_WORKSPACE_LIMIT,
              workspace_size_limit,
              &algo));
#endif
373
      PADDLE_ENFORCE_GPU_SUCCESS(
X
xiaoxiaohehe001 已提交
374
          platform::dynload::cudnnGetConvolutionForwardWorkspaceSize(
375 376 377 378 379
              handle,
              cudnn_input_desc,
              cudnn_filter_desc,
              cudnn_conv_desc,
              cudnn_output_desc,
X
xiaoxiaohehe001 已提交
380 381 382 383
              algo,
              &workspace_size_in_bytes));
      if (workspace_size_in_bytes > workspace_size_limit)
        workspace_size_limit = workspace_size_in_bytes;
384
      VLOG(3) << "cuDNN forward algo " << algo;
Q
qingqing01 已提交
385
    } else {
X
xiaoxiaohehe001 已提交
386 387
      std::function<SearchFuseResult<cudnnConvolutionFwdAlgo_t>()> search_func =
          [&]() -> SearchFuseResult<cudnnConvolutionFwdAlgo_t> {
Q
qingqing01 已提交
388
        int returned_algo_count;
X
xiaoxiaohehe001 已提交
389
        SearchFuseResult<cudnnConvolutionFwdAlgo_t> fwd_result;
390
        std::array<cudnnConvolutionFwdAlgoPerf_t, phi::kNUM_CUDNN_FWD_ALGS>
Q
qingqing01 已提交
391
            fwd_perf_stat;
C
chengduo 已提交
392
        auto cudnn_find_func = [&](void* cudnn_workspace) {
393
          PADDLE_ENFORCE_GPU_SUCCESS(
C
chengduo 已提交
394
              platform::dynload::cudnnFindConvolutionForwardAlgorithmEx(
395 396 397 398 399 400 401 402
                  handle,
                  cudnn_input_desc,
                  input_data,
                  cudnn_filter_desc,
                  filter_data,
                  cudnn_conv_desc,
                  cudnn_output_desc,
                  output_data,
403
                  phi::kNUM_CUDNN_FWD_ALGS,
404 405 406 407
                  &returned_algo_count,
                  fwd_perf_stat.data(),
                  cudnn_workspace,
                  workspace_size_limit));
C
chengduo 已提交
408
        };
409
        workspace_handle.RunFuncSync(cudnn_find_func, workspace_size_limit);
Q
qingqing01 已提交
410 411 412 413 414 415
        VLOG(3) << "Perf result: (algo: stat, time, memory)";
        for (int i = 0; i < returned_algo_count; ++i) {
          const auto& stat = fwd_perf_stat[i];
          VLOG(3) << stat.algo << ": " << stat.status << " " << stat.time << " "
                  << stat.memory;
        }
X
xiaoxiaohehe001 已提交
416 417 418 419 420 421 422 423 424 425 426 427 428 429 430 431 432 433 434 435 436 437 438

        PADDLE_ENFORCE_GPU_SUCCESS(
            platform::dynload::cudnnGetConvolutionForwardWorkspaceSize(
                handle,
                cudnn_input_desc,
                cudnn_filter_desc,
                cudnn_conv_desc,
                cudnn_output_desc,
                fwd_perf_stat[0].algo,
                &workspace_size_in_bytes));
        // PADDLE_ENFORCE_LE(
        //     workspace_size_in_bytes,
        //     workspace_size_limit,
        //     platform::errors::InvalidArgument(
        //         "The actual workspace size to be allocated for cuDNN is
        //         expected " "to be less than the limit. But received: the
        //         actual workspace " "size = %d, limit = %d.",
        //         workspace_size_in_bytes,
        //         workspace_size_limit));

        fwd_result.algo = fwd_perf_stat[0].algo;
        fwd_result.workspace_size = workspace_size_in_bytes;
        return fwd_result;
Q
qingqing01 已提交
439
      };
X
xiaoxiaohehe001 已提交
440
      AlgorithmsCache<SearchFuseResult<cudnnConvolutionFwdAlgo_t>>& algo_cache =
441
          *(framework::ConvSearchCache::Instance().GetConvFusion());
Q
qingqing01 已提交
442
      int search_times = ctx.Attr<int>("search_times");
X
xiaoxiaohehe001 已提交
443
      SearchFuseResult<cudnnConvolutionFwdAlgo_t> algo_result;
Q
qingqing01 已提交
444 445
      search_times = std::max(
          static_cast<int>(FLAGS_cudnn_exhaustive_search_times), search_times);
446
      // TODO(dangqingqing): Unify this if-else.
Q
qingqing01 已提交
447 448 449 450
      if (search_times > 0) {
        // The searched algo will be cached by `search_times` times for
        // different input dimension. For other dimensions, select the algo
        // of closest area.
X
xiaoxiaohehe001 已提交
451
        algo_result = algo_cache.GetAlgorithm(
452
            x_dims[2] * x_dims[3], search_times, 0, search_func);
X
xiaoxiaohehe001 已提交
453 454
        algo = algo_result.algo;
        workspace_size_in_bytes = algo_result.workspace_size;
Q
qingqing01 已提交
455
      } else {
X
xiaoxiaohehe001 已提交
456 457 458 459 460 461 462 463 464 465
        algo_result = algo_cache.GetAlgorithm(x_dims,
                                              f_dims,
                                              strides,
                                              paddings,
                                              dilations,
                                              0,
                                              dtype,
                                              search_func);
        algo = algo_result.algo;
        workspace_size_in_bytes = algo_result.workspace_size;
Q
qingqing01 已提交
466 467 468
      }
      VLOG(3) << "choose algo " << algo;
    }
N
nhzlx 已提交
469
    if ((activation == "identity") && (!residual)) {
470 471 472 473 474 475
      // Only the CUDNN_CONVOLUTION_FWD_ALGO_IMPLICIT_PRECOMP_GEMM algo is
      // enabled with CUDNN_ACTIVATION_IDENTITY in cuDNN lib.
      // But test in some case, the speed is slower, change to use
      // cudnnConvolutionForward and cudnnAddTensor
      // ------------- cudnn conv forward and bias add ---------------------
      ScalingParamType<T> alpha = 1.0f, beta = 0.0f;
C
chengduo 已提交
476
      auto cudnn_func = [&](void* cudnn_workspace) {
477 478 479 480 481 482 483 484 485 486 487 488 489 490
        PADDLE_ENFORCE_GPU_SUCCESS(
            platform::dynload::cudnnConvolutionForward(handle,
                                                       &alpha,
                                                       cudnn_input_desc,
                                                       input_data,
                                                       cudnn_filter_desc,
                                                       filter_data,
                                                       cudnn_conv_desc,
                                                       algo,
                                                       cudnn_workspace,
                                                       workspace_size_in_bytes,
                                                       &beta,
                                                       cudnn_output_desc,
                                                       output_data));
C
chengduo 已提交
491 492
      };
      workspace_handle.RunFunc(cudnn_func, workspace_size_in_bytes);
493 494 495 496 497 498 499 500
      PADDLE_ENFORCE_GPU_SUCCESS(
          platform::dynload::cudnnAddTensor(handle,
                                            &alpha,
                                            cudnn_bias_desc,
                                            bias_data,
                                            &alpha,
                                            cudnn_output_desc,
                                            output_data));
501 502 503 504 505 506 507
    } else {
      if (activation == "identity") {
        algo = CUDNN_CONVOLUTION_FWD_ALGO_IMPLICIT_PRECOMP_GEMM;
      }
      // ------------------- cudnn conv+bias+act forward --------------------
      ScalingParamType<T> alpha1 = 1.0f;
      ScalingParamType<T> alpha2 = residual ? 1.0f : 0.0f;
C
chengduo 已提交
508
      auto cudnn_func = [&](void* cudnn_workspace) {
509
        PADDLE_ENFORCE_GPU_SUCCESS(
510
            platform::dynload::cudnnConvolutionBiasActivationForward(
511 512 513 514 515 516 517 518 519 520 521 522 523 524 525 526 527 528
                handle,
                &alpha1,
                cudnn_input_desc,
                input_data,
                cudnn_filter_desc,
                filter_data,
                cudnn_conv_desc,
                algo,
                cudnn_workspace,
                workspace_size_in_bytes,
                &alpha2,
                cudnn_output_desc,
                residual_data,
                cudnn_bias_desc,
                bias_data,
                cudnn_act_desc,
                cudnn_output_desc,
                output_data));
C
chengduo 已提交
529 530
      };
      workspace_handle.RunFunc(cudnn_func, workspace_size_in_bytes);
531
    }
R
ronnywang 已提交
532
#endif
Q
qingqing01 已提交
533 534
    std::vector<int> channels = ctx.Attr<std::vector<int>>("split_channels");
    if (channels.size()) {
535
      auto outs = ctx.MultiOutput<phi::DenseTensor>("Outputs");
Q
qingqing01 已提交
536 537
      if (x_dims[0] == 1) {
        // share data with Output
538
        phi::DenseTensor t;
Q
qingqing01 已提交
539 540 541 542 543 544 545 546 547 548 549 550
        t.ShareDataWith(*output);
        auto y_dims = output->dims();
        t.Resize({y_dims[1], y_dims[2], y_dims[3]});
        int s = 0;
        for (size_t i = 0; i < channels.size(); ++i) {
          int e = s + channels[i];
          outs[i]->ShareDataWith(t.Slice(s, e));
          outs[i]->Resize({x_dims[0], channels[i], y_dims[2], y_dims[3]});
          s = e;
        }
      } else {
        // TODO(qingiqng): do copy when batch size large than 1
551
        PADDLE_THROW(platform::errors::Unimplemented(
552
            "Input with batch size greater than 1 is unsupported. The received "
553
            "batch size is %d, Input's shape is [%s].",
554 555
            x_dims[0],
            phi::make_ddim(x_dims)));
Q
qingqing01 已提交
556 557
      }
    }
Q
qingqing01 已提交
558 559
  }
};
D
Dang Qingqing 已提交
560
#endif
Q
qingqing01 已提交
561 562 563 564 565

}  // namespace operators
}  // namespace paddle

namespace ops = paddle::operators;
R
ronnywang 已提交
566
#if CUDNN_VERSION >= 7100
567 568 569 570 571
REGISTER_OP_CUDA_KERNEL(
    conv2d_fusion,
    ops::CUDNNConvFusionOpKernel<float>,
    ops::CUDNNConvFusionOpKernel<double>,
    ops::CUDNNConvFusionOpKernel<paddle::platform::float16>);
D
Dang Qingqing 已提交
572
#endif
R
ronnywang 已提交
573 574 575
#ifdef PADDLE_WITH_HIP
REGISTER_OP_CUDA_KERNEL(conv2d_fusion, ops::CUDNNConvFusionOpKernel<float>);
#endif