Skip to content

GitLab

Switch branch/tag
Find file Normal viewHistoryPermalink
conv_common.hpp 4.33 KB
Newer Older
Chao Liu's avatar
refactor  
Chao Liu committed
1 
#pragma once
Chao Liu's avatar
remove .hip extension  
Chao Liu committed
2 
#include "ConstantTensorDescriptor.hpp"
Chao Liu's avatar
refactor  
Chao Liu committed
3
4
5 

// this is ugly, only for 4d
template <class InDesc, class WeiDesc>
Chao Liu's avatar
use vectorized read and write for threadwise generic tensor copy  
Chao Liu committed
6 
constexpr auto get_convolution_output_default_4d_tensor_descriptor(InDesc, WeiDesc)
Chao Liu's avatar
refactor  
Chao Liu committed
7
8
9
10
11
12
13
14
15 
{
    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
adding implicit gemm v3  
Chao Liu committed
16
17 
    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
refactor  
Chao Liu committed
18
19
20
21
22
23
24
25 
    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
clean up  
Chao Liu committed
26
27 
    constexpr auto Y = wei_desc.GetLength(I2);
    constexpr auto X = wei_desc.GetLength(I3);
Chao Liu's avatar
refactor  
Chao Liu committed
28
Chao Liu's avatar
clean up  
Chao Liu committed
29
30 
    constexpr auto HO = HI + 1 - Y;
    constexpr auto WO = WI + 1 - X;
Chao Liu's avatar
padding works (sort of), but code looks ugly. Tuned some resnet configs  
Chao Liu committed
31
Chao Liu's avatar
refactor  
Chao Liu committed
32 
    return make_ConstantTensorDescriptor_packed(Sequence<N, K, HO, WO>{});
Chao Liu's avatar
padding works (sort of), but code looks ugly. Tuned some resnet configs  
Chao Liu committed
33
34
35 
}

template <class InDesc, class WeiDesc, class LowerPads, class UpperPads>
Chao Liu's avatar
use vectorized read and write for threadwise generic tensor copy  
Chao Liu committed
36
37
38
39 
constexpr auto get_convolution_with_padding_output_default_4d_tensor_descriptor(InDesc,
                                                                                WeiDesc,
                                                                                LowerPads,
                                                                                UpperPads)
Chao Liu's avatar
padding works (sort of), but code looks ugly. Tuned some resnet configs  
Chao Liu committed
40
41
42
43
44
45
46
47
48 
{
    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
adding implicit gemm v3  
Chao Liu committed
49
50 
    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
padding works (sort of), but code looks ugly. Tuned some resnet configs  
Chao Liu committed
51
52
53
54
55
56
57
58 
    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
clean up  
Chao Liu committed
59
60 
    constexpr auto Y = wei_desc.GetLength(I2);
    constexpr auto X = wei_desc.GetLength(I3);
Chao Liu's avatar
padding works (sort of), but code looks ugly. Tuned some resnet configs  
Chao Liu committed
61
62
63
64
65
66
67 

    constexpr auto HPadLow = LowerPads{}.Get(I0);
    constexpr auto WPadLow = LowerPads{}.Get(I1);

    constexpr auto HPadUp = UpperPads{}.Get(I0);
    constexpr auto WPadUp = UpperPads{}.Get(I1);
Chao Liu's avatar
clean up  
Chao Liu committed
68
69 
    constexpr auto HO = HI + HPadLow + HPadUp + 1 - Y;
    constexpr auto WO = WI + WPadLow + WPadUp + 1 - X;
Chao Liu's avatar
refactor  
Chao Liu committed
70
Chao Liu's avatar
refactor  
Chao Liu committed
71 
    return make_ConstantTensorDescriptor_packed(Sequence<N, K, HO, WO>{});
Chao Liu's avatar
refactor  
Chao Liu committed
72 
}
Chao Liu's avatar
unroll even-odd loop  
Chao Liu committed
73
74 

template <class InDesc, class WeiDesc, class OutDesc>
Chao Liu's avatar
refactor  
Chao Liu committed
75 
constexpr std::size_t calculate_convolution_flops(InDesc, WeiDesc, OutDesc)
Chao Liu's avatar
unroll even-odd loop  
Chao Liu committed
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95 
{
    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
refactor  
Chao Liu committed
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116 

template <class Float, class InDesc, class WeiDesc, class OutDesc>
constexpr std::size_t calculate_convolution_memory_size(Float, InDesc, WeiDesc, OutDesc)
{
    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);
Chao Liu's avatar
use vectorized read and write for threadwise generic tensor copy  
Chao Liu committed
117
118 
    return sizeof(Float) *
           (InDesc::GetElementSpace() + WeiDesc::GetElementSpace() + OutDesc::GetElementSpace());
Chao Liu's avatar
refactor  
Chao Liu committed
119 
}