conv_common.hpp 4.64 KB
Newer Older
Chao Liu's avatar
Chao Liu committed
1
2
#ifndef CONV_COMMON_HPP
#define CONV_COMMON_HPP
3

Chao Liu's avatar
Chao Liu committed
4
#include "ConstantTensorDescriptor.hpp"
Chao Liu's avatar
Chao Liu committed
5

Jehandad Khan's avatar
Jehandad Khan committed
6
7

typedef enum ConvolutionDir{ Forward=0, BackwardData=1, BackwardWeights=2};
Chao Liu's avatar
Chao Liu committed
8
9
// this is ugly, only for 4d
template <class InDesc, class WeiDesc>
10
constexpr auto get_convolution_output_default_4d_tensor_descriptor(InDesc, WeiDesc)
Chao Liu's avatar
Chao Liu committed
11
{
Chao Liu's avatar
Chao Liu committed
12
13
    using namespace ck;

Chao Liu's avatar
Chao Liu committed
14
15
16
17
18
19
20
21
    constexpr auto in_desc  = InDesc{};
    constexpr auto wei_desc = WeiDesc{};

    constexpr auto I0 = Number<0>{};
    constexpr auto I1 = Number<1>{};
    constexpr auto I2 = Number<2>{};
    constexpr auto I3 = Number<3>{};

Chao Liu's avatar
Chao Liu committed
22
23
    static_assert(in_desc.GetNumOfDimension() == 4, "input nDim is not 4");
    static_assert(wei_desc.GetNumOfDimension() == 4, "weight nDim is not 4");
Chao Liu's avatar
Chao Liu committed
24
25
26
27
28
29
30
31
    static_assert(in_desc.GetLength(I1) == wei_desc.GetLength(I1),
                  "input & weight dimension not consistent");

    constexpr auto N  = in_desc.GetLength(I0);
    constexpr auto HI = in_desc.GetLength(I2);
    constexpr auto WI = in_desc.GetLength(I3);

    constexpr auto K = wei_desc.GetLength(I0);
Chao Liu's avatar
Chao Liu committed
32
33
    constexpr auto Y = wei_desc.GetLength(I2);
    constexpr auto X = wei_desc.GetLength(I3);
Chao Liu's avatar
Chao Liu committed
34

Chao Liu's avatar
Chao Liu committed
35
36
    constexpr auto HO = HI + 1 - Y;
    constexpr auto WO = WI + 1 - X;
37

Chao Liu's avatar
Chao Liu committed
38
    return make_ConstantTensorDescriptor_packed(Sequence<N, K, HO, WO>{});
39
40
}

41
42
43
44
45
46
47
48
template <class InDesc,
          class WeiDesc,
          class ConvStrides,
          class ConvDilations,
          class LowerPads,
          class UpperPads>
constexpr auto get_convolution_with_padding_output_default_4d_tensor_descriptor(
    InDesc, WeiDesc, ConvStrides, ConvDilations, LowerPads, UpperPads)
49
{
Chao Liu's avatar
Chao Liu committed
50
51
    using namespace ck;

52
53
54
55
56
57
58
59
    constexpr auto in_desc  = InDesc{};
    constexpr auto wei_desc = WeiDesc{};

    constexpr auto I0 = Number<0>{};
    constexpr auto I1 = Number<1>{};
    constexpr auto I2 = Number<2>{};
    constexpr auto I3 = Number<3>{};

Chao Liu's avatar
Chao Liu committed
60
61
    static_assert(in_desc.GetNumOfDimension() == 4, "input nDim is not 4");
    static_assert(wei_desc.GetNumOfDimension() == 4, "weight nDim is not 4");
62
63
64
    static_assert(in_desc.GetLength(I1) == wei_desc.GetLength(I1),
                  "input & weight dimension not consistent");

65
66
67
    constexpr index_t N  = in_desc.GetLength(I0);
    constexpr index_t Hi = in_desc.GetLength(I2);
    constexpr index_t Wi = in_desc.GetLength(I3);
68

69
70
71
    constexpr index_t K = wei_desc.GetLength(I0);
    constexpr index_t Y = wei_desc.GetLength(I2);
    constexpr index_t X = wei_desc.GetLength(I3);
72

73
74
    constexpr index_t HPadLow = LowerPads{}.Get(I0);
    constexpr index_t WPadLow = LowerPads{}.Get(I1);
75

76
77
    constexpr index_t HPadUp = UpperPads{}.Get(I0);
    constexpr index_t WPadUp = UpperPads{}.Get(I1);
78

79
80
    constexpr index_t YEff = (Y - 1) * ConvDilations{}[0] + 1;
    constexpr index_t XEff = (X - 1) * ConvDilations{}[1] + 1;
Chao Liu's avatar
Chao Liu committed
81

82
83
84
85
    constexpr index_t Ho = (Hi + HPadLow + HPadUp - YEff) / ConvStrides{}[0] + 1;
    constexpr index_t Wo = (Wi + WPadLow + WPadUp - XEff) / ConvStrides{}[1] + 1;

    return make_ConstantTensorDescriptor_packed(Sequence<N, K, Ho, Wo>{});
Chao Liu's avatar
Chao Liu committed
86
}
Chao Liu's avatar
Chao Liu committed
87
88

template <class InDesc, class WeiDesc, class OutDesc>
Chao Liu's avatar
Chao Liu committed
89
constexpr std::size_t calculate_convolution_flops(InDesc, WeiDesc, OutDesc)
Chao Liu's avatar
Chao Liu committed
90
{
Chao Liu's avatar
Chao Liu committed
91
92
    using namespace ck;

Chao Liu's avatar
Chao Liu committed
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
    constexpr auto wei_desc = WeiDesc{};
    constexpr auto out_desc = OutDesc{};

    constexpr auto I0 = Number<0>{};
    constexpr auto I1 = Number<1>{};
    constexpr auto I2 = Number<2>{};
    constexpr auto I3 = Number<3>{};

    constexpr index_t N  = out_desc.GetLength(I0);
    constexpr index_t K  = out_desc.GetLength(I1);
    constexpr index_t Ho = out_desc.GetLength(I2);
    constexpr index_t Wo = out_desc.GetLength(I3);

    constexpr index_t C = wei_desc.GetLength(I1);
    constexpr index_t Y = wei_desc.GetLength(I2);
    constexpr index_t X = wei_desc.GetLength(I3);

    return std::size_t(2) * N * K * Ho * Wo * C * Y * X;
}
Chao Liu's avatar
Chao Liu committed
112
113
114
115

template <class Float, class InDesc, class WeiDesc, class OutDesc>
constexpr std::size_t calculate_convolution_memory_size(Float, InDesc, WeiDesc, OutDesc)
{
Chao Liu's avatar
Chao Liu committed
116
117
    using namespace ck;

Chao Liu's avatar
Chao Liu committed
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
    constexpr auto wei_desc = WeiDesc{};
    constexpr auto out_desc = OutDesc{};

    constexpr auto I0 = Number<0>{};
    constexpr auto I1 = Number<1>{};
    constexpr auto I2 = Number<2>{};
    constexpr auto I3 = Number<3>{};

    constexpr index_t N  = out_desc.GetLength(I0);
    constexpr index_t K  = out_desc.GetLength(I1);
    constexpr index_t Ho = out_desc.GetLength(I2);
    constexpr index_t Wo = out_desc.GetLength(I3);

    constexpr index_t C = wei_desc.GetLength(I1);
    constexpr index_t Y = wei_desc.GetLength(I2);
    constexpr index_t X = wei_desc.GetLength(I3);

135
136
    return sizeof(Float) *
           (InDesc::GetElementSpace() + WeiDesc::GetElementSpace() + OutDesc::GetElementSpace());
Chao Liu's avatar
Chao Liu committed
137
}
138
139

#endif