kern.cuh 2.4 KB
Newer Older
1 2 3 4
/**
 * \file dnn/src/cuda/convolution/chanwise/kern.cuh
 * MegEngine is Licensed under the Apache License, Version 2.0 (the "License")
 *
5
 * Copyright (c) 2014-2021 Megvii Inc. All rights reserved.
6 7 8 9 10 11 12 13 14 15
 *
 * Unless required by applicable law or agreed to in writing,
 * software distributed under the License is distributed on an
 * "AS IS" BASIS, WITHOUT ARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 */
#pragma once

#include "src/cuda/utils.cuh"

#include <cuda_runtime.h>
M
Megvii Engine Team 已提交
16
#include <stdint.h>
17 18 19 20 21 22 23 24 25 26

#if MEGDNN_CC_HOST
#include "src/cuda/convolution/helper.h"
#endif

namespace megdnn {
namespace cuda {
namespace convolution {
namespace chanwise {

M
Megvii Engine Team 已提交
27 28 29
struct Param {
    uint32_t batch, src_chl, src_h, src_w, chl_mul, flt_h, flt_w, out_h, out_w, pad_h,
            pad_w, stride_h, stride_w, dilation_h, dilation_w;
30
#if MEGDNN_CC_HOST
M
Megvii Engine Team 已提交
31
    static Param from_fwd_args(const ForwardSizeArgs& args) {
32
#define U(v) static_cast<uint32_t>(v)
M
Megvii Engine Team 已提交
33 34 35 36 37 38 39 40 41 42
        auto&& src = args.src_layout->shape;
        auto&& dst = args.dst_layout->shape;
        auto&& fm = args.filter_meta;
        size_t c_pos, hw_pos;
        if (fm.format == param::Convolution::Format::NCHW) {
            c_pos = 1;
            hw_pos = 2;
        } else {
            c_pos = 3;
            hw_pos = 1;
43
        }
M
Megvii Engine Team 已提交
44 45 46 47 48 49 50 51 52
        return {
                U(src[0]),          U(src[c_pos]),     U(src[hw_pos]),
                U(src[hw_pos + 1]), U(fm.ocpg),        U(fm.spatial[0]),
                U(fm.spatial[1]),   U(dst[hw_pos]),    U(dst[hw_pos + 1]),
                U(fm.padding[0]),   U(fm.padding[1]),  U(fm.stride[0]),
                U(fm.stride[1]),    U(fm.dilation[0]), U(fm.dilation[1]),
        };
#undef U
    }
53
#endif
M
Megvii Engine Team 已提交
54
};
55

M
Megvii Engine Team 已提交
56 57 58 59
template <typename T>
void run_bwd_data_small(
        T* src_grad, const T* dst_grad, const T* flt, const Param& param,
        cudaStream_t stream);
60

M
Megvii Engine Team 已提交
61 62 63 64
template <typename T>
void run_bwd_data(
        T* src_grad, const T* dst_grad, const T* flt, const Param& param,
        cudaStream_t stream);
65

66 67 68 69
template <typename T>
void run_bwd_depthwise_large_filter(
        T* dst, const T* src, const T* flt, const Param& param, cudaStream_t stream);

M
Megvii Engine Team 已提交
70 71 72 73
template <typename T>
void run_bwd_filter(
        T* filter_grad, const T* src, const T* dst_grad, const Param& param,
        cudaStream_t stream);
74

M
Megvii Engine Team 已提交
75 76 77 78
}  // namespace chanwise
}  // namespace convolution
}  // namespace cuda
}  // namespace megdnn
79 80

// vim: ft=cpp syntax=cpp.doxygen