Skip to content

GitLab

Switch branch/tag
Find file Normal viewHistoryPermalink
transpose_vectors.hpp 7.85 KB
Newer Older
Chao Liu's avatar
add license in file (#303)  
Chao Liu committed
1 
// SPDX-License-Identifier: MIT
Illia Silin's avatar
update copyright headers (#726)  
Illia Silin committed
2 
// Copyright (c) 2018-2023, Advanced Micro Devices, Inc. All rights reserved.
Chao Liu's avatar
add license in file (#303)  
Chao Liu committed
3
Chao Liu's avatar
Absolute include path (#281)  
Chao Liu committed
4 
#pragma once
Chao Liu's avatar
FP16 data in-register transpose (#41)  
Chao Liu committed
5
Chao Liu's avatar
Absolute include path (#281)  
Chao Liu committed
6 
#include "ck/ck.hpp"
Chao Liu's avatar
FP16 data in-register transpose (#41)  
Chao Liu committed
7
8
9
10
11 
#include "statically_indexed_array.hpp"
#include "data_type.hpp"

namespace ck {
Chao Liu's avatar
Shuffle in thread (#13)  
Chao Liu committed
12
13
14 
// S: scalar type
// NX: # of vector before transpose
// NY: # of vector after transpose
Chao Liu's avatar
FP16 data in-register transpose (#41)  
Chao Liu committed
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39 
template <typename S,
          index_t NX,
          index_t NY,
          typename enable_if<is_scalar_type<S>::value, bool>::type = false>
struct transpose_vectors;

// transpose fp16 2x2
__device__ void transpose_fp16_2x2(const half2_t& x0, const half2_t& x1, half2_t& y0, half2_t& y1)
{
#if 0
    static constexpr auto I0 = Number<0>{};
    static constexpr auto I1 = Number<1>{};

    const vector_type<half_t, 2> vx0{x0}, vx1{x1};
    vector_type<half_t, 2> vy0, vy1;

    vy0.template AsType<half_t>()(I0) = vx0.template AsType<half_t>()[I0];
    vy0.template AsType<half_t>()(I1) = vx1.template AsType<half_t>()[I0];

    vy1.template AsType<half_t>()(I0) = vx0.template AsType<half_t>()[I1];
    vy1.template AsType<half_t>()(I1) = vx1.template AsType<half_t>()[I1];

    y0 = vy0.template AsType<half2_t>()[I0];
    y1 = vy1.template AsType<half2_t>()[I0];
#else
Anthony Chang's avatar
work around inline asm potential hazard using intrinsic (#416)  
Anthony Chang committed
40
41
42
43
44
45
46
47
48 
    constexpr int32_t m0 = 0x05040100;
    constexpr int32_t m1 = 0x07060302;

    // ex: v_perm_b32(0x 11 22 33 44, 0x 55 66 77 88, 0x 05 01 04 00) -> 0x33774488
    //                   -- -- -- --     -- -- -- --      -  -  -  -
    //             index  7  6  5  4      3  2  1  0     33 77 44 88
    // index is reversed because of little endianness (least significant bits first)
    y0 = bit_cast<half2_t>(__builtin_amdgcn_perm(bit_cast<int32_t>(x1), bit_cast<int32_t>(x0), m0));
    y1 = bit_cast<half2_t>(__builtin_amdgcn_perm(bit_cast<int32_t>(x1), bit_cast<int32_t>(x0), m1));
Chao Liu's avatar
FP16 data in-register transpose (#41)  
Chao Liu committed
49
50
51
52
53
54 
#endif
}

template <index_t NX, index_t NY>
struct transpose_vectors<half_t, NX, NY>
{
Anthony Chang's avatar
use inline asm for 4x4 int8 transposition (#187)  
Anthony Chang committed
55 
    // we got [NY * NX] amount of S data to be transposed
Chao Liu's avatar
FP16 data in-register transpose (#41)  
Chao Liu committed
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76 
    static constexpr index_t s_per_x = NY;
    static constexpr index_t s_per_y = NX;

    using S  = half_t;
    using VX = vector_type<half_t, s_per_x>;
    using VY = vector_type<half_t, s_per_y>;

    __device__ void operator()(const StaticallyIndexedArray<const VX&, NX>& vx_tuple,
                               StaticallyIndexedArray<VY&, NY>& vy_tuple)
    {
        static constexpr auto I1 = Number<1>{};
        static constexpr auto I2 = Number<2>{};

        static_assert((NX % 2 == 0 && NY % 2 == 0), "wrong!");

        // loop over 2x2 tile and transpose data from vx_tuple into vy_tuple
        static_for<0, NY, 2>{}([&](auto iy) {
            static_for<0, NX, 2>{}([&](auto ix) {
                // reference to 2 half2_t data from vx_tuple
                const auto& x_s2_0 = vx_tuple[ix].template AsType<half2_t>()[iy / I2];
                const auto& x_s2_1 = vx_tuple[ix + I1].template AsType<half2_t>()[iy / I2];
Chao Liu's avatar
Shuffle in thread (#13)  
Chao Liu committed
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102 

                // reference to 2 half2_t data from vy_tuple
                auto& y_s2_0 = vy_tuple(iy).template AsType<half2_t>()(ix / I2);
                auto& y_s2_1 = vy_tuple(iy + I1).template AsType<half2_t>()(ix / I2);

                // transpose
                transpose_fp16_2x2(x_s2_0, x_s2_1, y_s2_0, y_s2_1);
            });
        });
    }

    // FIXME: duplicated code
    __device__ void operator()(const StaticallyIndexedArray<VX, NX>& vx_tuple,
                               StaticallyIndexedArray<VY, NY>& vy_tuple)
    {
        static constexpr auto I1 = Number<1>{};
        static constexpr auto I2 = Number<2>{};

        static_assert((NX % 2 == 0 && NY % 2 == 0), "wrong!");

        // loop over 2x2 tile and transpose data from vx_tuple into vy_tuple
        static_for<0, NY, 2>{}([&](auto iy) {
            static_for<0, NX, 2>{}([&](auto ix) {
                // reference to 2 half2_t data from vx_tuple
                const auto& x_s2_0 = vx_tuple[ix].template AsType<half2_t>()[iy / I2];
                const auto& x_s2_1 = vx_tuple[ix + I1].template AsType<half2_t>()[iy / I2];
Chao Liu's avatar
FP16 data in-register transpose (#41)  
Chao Liu committed
103
104
105
106
107
108
109
110
111
112
113
114 

                // reference to 2 half2_t data from vy_tuple
                auto& y_s2_0 = vy_tuple(iy).template AsType<half2_t>()(ix / I2);
                auto& y_s2_1 = vy_tuple(iy + I1).template AsType<half2_t>()(ix / I2);

                // transpose
                transpose_fp16_2x2(x_s2_0, x_s2_1, y_s2_0, y_s2_1);
            });
        });
    }
};
Anthony Chang's avatar
use inline asm for 4x4 int8 transposition (#187)  
Anthony Chang committed
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135 
// transpose int8 4x4
__device__ void transpose_int8_4x4(const int8x4_t& x0,
                                   const int8x4_t& x1,
                                   const int8x4_t& x2,
                                   const int8x4_t& x3,
                                   int8x4_t& y0,
                                   int8x4_t& y1,
                                   int8x4_t& y2,
                                   int8x4_t& y3)
{
    int32_t t0, t1;
    int32_t z0, z1, z2, z3;
    constexpr int32_t m0 = 0x05010400;
    constexpr int32_t m1 = 0x05040100;
    constexpr int32_t m2 = 0x07060302;
    constexpr int32_t m3 = 0x07030602;

    // ex: v_perm_b32(0x 11 22 33 44, 0x 55 66 77 88, 0x 05 01 04 00) -> 0x33774488
    //                   -- -- -- --     -- -- -- --      -  -  -  -
    //             index  7  6  5  4      3  2  1  0     33 77 44 88
    // index is reversed because of little endianness (least significant bits first)
Anthony Chang's avatar
work around inline asm potential hazard using intrinsic (#416)  
Anthony Chang committed
136
137
138
139
140
141
142
143 
    t0 = __builtin_amdgcn_perm(bit_cast<int32_t>(x1), bit_cast<int32_t>(x0), m0);
    t1 = __builtin_amdgcn_perm(bit_cast<int32_t>(x3), bit_cast<int32_t>(x2), m0);
    z0 = __builtin_amdgcn_perm(bit_cast<int32_t>(t1), bit_cast<int32_t>(t0), m1);
    z1 = __builtin_amdgcn_perm(bit_cast<int32_t>(t1), bit_cast<int32_t>(t0), m2);
    t0 = __builtin_amdgcn_perm(bit_cast<int32_t>(x1), bit_cast<int32_t>(x0), m3);
    t1 = __builtin_amdgcn_perm(bit_cast<int32_t>(x3), bit_cast<int32_t>(x2), m3);
    z2 = __builtin_amdgcn_perm(bit_cast<int32_t>(t1), bit_cast<int32_t>(t0), m1);
    z3 = __builtin_amdgcn_perm(bit_cast<int32_t>(t1), bit_cast<int32_t>(t0), m2);
Anthony Chang's avatar
use inline asm for 4x4 int8 transposition (#187)  
Anthony Chang committed
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193 

    y0 = bit_cast<int8x4_t>(z0);
    y1 = bit_cast<int8x4_t>(z1);
    y2 = bit_cast<int8x4_t>(z2);
    y3 = bit_cast<int8x4_t>(z3);
}

template <index_t NX, index_t NY>
struct transpose_vectors<int8_t, NX, NY>
{
    // we got [NY * NX] amount of S data to be transposed
    static constexpr index_t s_per_x = NY;
    static constexpr index_t s_per_y = NX;

    using S  = int8_t;
    using VX = vector_type<int8_t, s_per_x>;
    using VY = vector_type<int8_t, s_per_y>;

    __device__ void operator()(const StaticallyIndexedArray<const VX&, NX>& vx_tuple,
                               StaticallyIndexedArray<VY&, NY>& vy_tuple)
    {
        static constexpr auto I1 = Number<1>{};
        static constexpr auto I2 = Number<2>{};
        static constexpr auto I3 = Number<3>{};
        static constexpr auto I4 = Number<4>{};

        static_assert((NX % 4 == 0 && NY % 4 == 0), "wrong!");

        // loop over 4x4 tile and transpose data from vx_tuple into vy_tuple
        static_for<0, NY, 4>{}([&](auto iy) {
            static_for<0, NX, 4>{}([&](auto ix) {
                // reference to 4 int8 data from vx_tuple
                const auto& x_s4_0 = vx_tuple[ix].template AsType<int8x4_t>()[iy / I4];
                const auto& x_s4_1 = vx_tuple[ix + I1].template AsType<int8x4_t>()[iy / I4];
                const auto& x_s4_2 = vx_tuple[ix + I2].template AsType<int8x4_t>()[iy / I4];
                const auto& x_s4_3 = vx_tuple[ix + I3].template AsType<int8x4_t>()[iy / I4];

                // reference to 4 int8 data from vy_tuple
                auto& y_s4_0 = vy_tuple(iy).template AsType<int8x4_t>()(ix / I4);
                auto& y_s4_1 = vy_tuple(iy + I1).template AsType<int8x4_t>()(ix / I4);
                auto& y_s4_2 = vy_tuple(iy + I2).template AsType<int8x4_t>()(ix / I4);
                auto& y_s4_3 = vy_tuple(iy + I3).template AsType<int8x4_t>()(ix / I4);

                // transpose
                transpose_int8_4x4(x_s4_0, x_s4_1, x_s4_2, x_s4_3, y_s4_0, y_s4_1, y_s4_2, y_s4_3);
            });
        });
    }
};
Chao Liu's avatar
FP16 data in-register transpose (#41)  
Chao Liu committed
194 
} // namespace ck