Skip to content

GitLab

Switch branch/tag
Find file Normal viewHistoryPermalink
conv_common.hip.hpp 3.42 KB
Newer Older
Chao Liu's avatar
refactor  
Chao Liu committed
1 
#pragma once
Chao Liu's avatar
change file extension to hip.hpp and hip.cpp  
Chao Liu committed
2 
#include "ConstantTensorDescriptor.hip.hpp"
Chao Liu's avatar
refactor  
Chao Liu committed
3
4
5
6
7
8
9
10
11
12
13
14
15
16 

// this is ugly, only for 4d
template <class InDesc, class WeiDesc>
__host__ __device__ constexpr auto get_convolution_output_default_4d_tensor_descriptor(InDesc,
                                                                                       WeiDesc)
{
    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
17
18 
    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
19
20
21
22
23
24
25
26 
    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
27
28 
    constexpr auto Y = wei_desc.GetLength(I2);
    constexpr auto X = wei_desc.GetLength(I3);
Chao Liu's avatar
refactor  
Chao Liu committed
29
Chao Liu's avatar
clean up  
Chao Liu committed
30
31 
    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
32
Chao Liu's avatar
adding implicit gemm v3  
Chao Liu committed
33 
    return make_ConstantTensorDescriptor_default_rank_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
34
35
36
37
38
39
40
41
42
43
44
45
46
47 
}

template <class InDesc, class WeiDesc, class LowerPads, class UpperPads>
__host__ __device__ constexpr auto get_convolution_with_padding_output_default_4d_tensor_descriptor(
    InDesc, WeiDesc, LowerPads, UpperPads)
{
    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
48
49 
    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
50
51
52
53
54
55
56
57 
    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
58
59 
    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
60
61
62
63
64
65
66 

    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
67
68 
    constexpr auto HO = HI + HPadLow + HPadUp + 1 - Y;
    constexpr auto WO = WI + WPadLow + WPadUp + 1 - X;
Chao Liu's avatar
refactor  
Chao Liu committed
69
Chao Liu's avatar
adding implicit gemm v3  
Chao Liu committed
70 
    return make_ConstantTensorDescriptor_default_rank_packed(Sequence<N, K, HO, WO>{});
Chao Liu's avatar
refactor  
Chao Liu committed
71 
}
Chao Liu's avatar
unroll even-odd loop  
Chao Liu committed
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94 

template <class InDesc, class WeiDesc, class OutDesc>
__host__ __device__ constexpr std::size_t calculate_convolution_flops(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);

    return std::size_t(2) * N * K * Ho * Wo * C * Y * X;
}