Function data_gradient_cudnn

src/backend/cuda/convolveNN.cpp:288–348

#ifdef WITH_CUDNN
template<typename T>
Array<T> data_gradient_cudnn(const Array<T> &incoming_gradient,
                             const Array<T> &original_signal,
                             const Array<T> &original_filter,
                             const Array<T> &convolved_output, af::dim4 stride,
                             af::dim4 padding, af::dim4 dilation) {
    UNUSED(convolved_output);
    auto cudnn = nnHandle();

    dim4 sDims = original_signal.dims();
    dim4 fDims = original_filter.dims();

    cudnnDataType_t cudnn_dtype = getCudnnDataType<T>();

    // create x descriptor
    auto dx_descriptor = toCudnn<cudnnTensorDescriptor_t>(original_signal);
    auto dy_descriptor = toCudnn<cudnnTensorDescriptor_t>(incoming_gradient);

    // create output filter gradient descriptor
    auto w_descriptor = make_handle<cudnnFilterDescriptor_t>();

    CUDNN_CHECK(cuda::cudnnSetFilter4dDescriptor(w_descriptor, cudnn_dtype,
                                                 CUDNN_TENSOR_NCHW, fDims[3],
                                                 fDims[2], fDims[1], fDims[0]));

    // create convolution descriptor
    auto convolution_descriptor = make_handle<cudnnConvolutionDescriptor_t>();

    CUDNN_CHECK(cuda::cudnnSetConvolution2dDescriptor(
        convolution_descriptor, padding[1], padding[0], stride[1], stride[0],
        dilation[1], dilation[0], CUDNN_CONVOLUTION, cudnn_dtype));

    cudnnConvolutionBwdDataAlgo_t bwd_data_convolution_algorithm;
    if ((dilation[0] == 1 && dilation[1] == 1) || is_same<T, half>::value) {
        bwd_data_convolution_algorithm = CUDNN_CONVOLUTION_BWD_DATA_ALGO_1;
    } else {
        bwd_data_convolution_algorithm = CUDNN_CONVOLUTION_BWD_DATA_ALGO_0;
    }

    // figure out scratch space memory requirements
    size_t workspace_bytes;
    CUDNN_CHECK(cuda::cudnnGetConvolutionBackwardDataWorkspaceSize(
        cudnn, w_descriptor, dy_descriptor, convolution_descriptor,
        dx_descriptor, bwd_data_convolution_algorithm, &workspace_bytes));

    dim4 odims(sDims[0], sDims[1], sDims[2], sDims[3]);
    Array<T> out = createEmptyArray<T>(odims);

    auto workspace_buffer = memAlloc<char>(workspace_bytes);

    // perform convolution
    auto alpha = scalar<scale_type<T>>(1.0);
    auto beta  = scalar<scale_type<T>>(0.0);

    CUDNN_CHECK(cuda::cudnnConvolutionBackwardData(
        cudnn, &alpha, w_descriptor, original_filter.get(), dy_descriptor,
        incoming_gradient.get(), convolution_descriptor,
        bwd_data_convolution_algorithm, (void *)workspace_buffer.get(),
        workspace_bytes, &beta, dx_descriptor, out.device()));