Skip to content

GitLab

Switch branch/tag
Find file Normal viewHistoryPermalink
dynamic_buffer.hpp 9.09 KB
Newer Older
Chao Liu's avatar
Use DynamicBuffer instead of raw pointer (#32)  
Chao Liu committed
1
2
3
4
5
6
7
8
9
10
11
12
13
14
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
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
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 
#ifndef CK_DYNAMIC_BUFFER_HPP
#define CK_DYNAMIC_BUFFER_HPP

namespace ck {

#include "amd_buffer_addressing_v2.hpp"

template <AddressSpace BufferAddressSpace, typename T, typename ElementSpaceSize>
struct DynamicBuffer
{
    using type = T;

    T* p_data_;
    ElementSpaceSize element_space_size_;

    __host__ __device__ constexpr DynamicBuffer(T* p_data, ElementSpaceSize element_space_size)
        : p_data_{p_data}, element_space_size_{element_space_size}
    {
    }

    __host__ __device__ static constexpr AddressSpace GetAddressSpace()
    {
        return BufferAddressSpace;
    }

    __host__ __device__ constexpr const T& operator[](index_t i) const { return p_data_[i]; }

    __host__ __device__ constexpr T& operator()(index_t i) { return p_data_[i]; }

    template <typename X,
              typename std::enable_if<
                  is_same<typename scalar_type<remove_cv_t<remove_reference_t<X>>>::type,
                          typename scalar_type<remove_cv_t<remove_reference_t<T>>>::type>::value,
                  bool>::type = false>
    __host__ __device__ constexpr const auto Get(index_t i, bool is_valid_offset) const
    {
        // X contains multiple T
        constexpr index_t scalar_per_t_vector =
            scalar_type<remove_cv_t<remove_reference_t<T>>>::vector_size;

        constexpr index_t scalar_per_x_vector =
            scalar_type<remove_cv_t<remove_reference_t<X>>>::vector_size;

        static_assert(scalar_per_x_vector % scalar_per_t_vector == 0,
                      "wrong! X need to be multiple T");

        constexpr index_t t_per_x = scalar_per_x_vector / scalar_per_t_vector;

        if constexpr(GetAddressSpace() == AddressSpace::Global)
        {
#if CK_USE_AMD_BUFFER_ADDRESSING
            return amd_buffer_load_v2<remove_cv_t<remove_reference_t<T>>, t_per_x>(
                p_data_, i, is_valid_offset, element_space_size_);
#else
            return is_valid_offset ? *reinterpret_cast<const X*>(&p_data_[i]) : X{0};
#endif
        }
        else
        {
            return is_valid_offset ? *reinterpret_cast<const X*>(&p_data_[i]) : X{0};
        }
    }

    template <typename X,
              typename std::enable_if<
                  is_same<typename scalar_type<remove_cv_t<remove_reference_t<X>>>::type,
                          typename scalar_type<remove_cv_t<remove_reference_t<T>>>::type>::value,
                  bool>::type = false>
    __host__ __device__ void Set(index_t i, bool is_valid_offset, const X& x)
    {
        // X contains multiple T
        constexpr index_t scalar_per_t_vector =
            scalar_type<remove_cv_t<remove_reference_t<T>>>::vector_size;

        constexpr index_t scalar_per_x_vector =
            scalar_type<remove_cv_t<remove_reference_t<X>>>::vector_size;

        static_assert(scalar_per_x_vector % scalar_per_t_vector == 0,
                      "wrong! X need to be multiple T");

        constexpr index_t t_per_x = scalar_per_x_vector / scalar_per_t_vector;

        if constexpr(GetAddressSpace() == AddressSpace::Global)
        {
#if CK_USE_AMD_BUFFER_ADDRESSING
            amd_buffer_store_v2<remove_cv_t<remove_reference_t<T>>, t_per_x>(
                x, p_data_, i, is_valid_offset, element_space_size_);
#else
            if(is_valid_offset)
            {
                *reinterpret_cast<X*>(&p_data_[i]) = x;
            }
#endif
        }
        else if constexpr(GetAddressSpace() == AddressSpace::Lds)
        {
            if(is_valid_offset)
            {
#if !CK_WORKAROUND_SWDEV_XXXXXX_INT8_DS_WRITE_ISSUE
                *reinterpret_cast<X*>(&p_data_[i]) = x;
#else
                // HACK: compiler would lower IR "store<i8, 16> address_space(3)" into inefficient
Chao Liu's avatar
DL GEMM fp32/fp16/int8 (#41)  
Chao Liu committed
103 
                // ISA, so I try to let compiler emit IR "store<i32, 4>" which would be lower to
Chao Liu's avatar
Use DynamicBuffer instead of raw pointer (#32)  
Chao Liu committed
104
105
106
107
108
109 
                // ds_write_b128
                // TODO: remove this after compiler fix
                if constexpr(is_same<typename scalar_type<remove_cv_t<remove_reference_t<T>>>::type,
                                     int8_t>::value)
                {
                    static_assert(
Chao Liu's avatar
DL GEMM fp32/fp16/int8 (#41)  
Chao Liu committed
110
111
112
113
114
115
116
117 
                        (is_same<remove_cv_t<remove_reference_t<T>>, int8_t>::value &&
                         is_same<remove_cv_t<remove_reference_t<X>>, int8_t>::value) ||
                            (is_same<remove_cv_t<remove_reference_t<T>>, int8_t>::value &&
                             is_same<remove_cv_t<remove_reference_t<X>>, int8x2_t>::value) ||
                            (is_same<remove_cv_t<remove_reference_t<T>>, int8_t>::value &&
                             is_same<remove_cv_t<remove_reference_t<X>>, int8x4_t>::value) ||
                            (is_same<remove_cv_t<remove_reference_t<T>>, int8x4_t>::value &&
                             is_same<remove_cv_t<remove_reference_t<X>>, int8x4_t>::value) ||
Chao Liu's avatar
Use DynamicBuffer instead of raw pointer (#32)  
Chao Liu committed
118
119
120
121
122
123 
                            (is_same<remove_cv_t<remove_reference_t<T>>, int8x8_t>::value &&
                             is_same<remove_cv_t<remove_reference_t<X>>, int8x8_t>::value) ||
                            (is_same<remove_cv_t<remove_reference_t<T>>, int8x16_t>::value &&
                             is_same<remove_cv_t<remove_reference_t<X>>, int8x16_t>::value),
                        "wrong! not implemented for this combination, please add implementation");
Chao Liu's avatar
DL GEMM fp32/fp16/int8 (#41)  
Chao Liu committed
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150 
                    if constexpr(is_same<remove_cv_t<remove_reference_t<T>>, int8_t>::value &&
                                 is_same<remove_cv_t<remove_reference_t<X>>, int8_t>::value)
                    {
                        // HACK: cast pointer of x is bad
                        // TODO: remove this after compiler fix
                        *reinterpret_cast<int8_t*>(&p_data_[i]) =
                            *reinterpret_cast<const int8_t*>(&x);
                    }
                    else if constexpr(is_same<remove_cv_t<remove_reference_t<T>>, int8_t>::value &&
                                      is_same<remove_cv_t<remove_reference_t<X>>, int8x2_t>::value)
                    {
                        // HACK: cast pointer of x is bad
                        // TODO: remove this after compiler fix
                        *reinterpret_cast<int16_t*>(&p_data_[i]) =
                            *reinterpret_cast<const int16_t*>(&x);
                    }
                    else if constexpr(is_same<remove_cv_t<remove_reference_t<T>>, int8_t>::value &&
                                      is_same<remove_cv_t<remove_reference_t<X>>, int8x4_t>::value)
                    {
                        // HACK: cast pointer of x is bad
                        // TODO: remove this after compiler fix
                        *reinterpret_cast<int32_t*>(&p_data_[i]) =
                            *reinterpret_cast<const int32_t*>(&x);
                    }
                    else if constexpr(is_same<remove_cv_t<remove_reference_t<T>>,
                                              int8x4_t>::value &&
                                      is_same<remove_cv_t<remove_reference_t<X>>, int8x4_t>::value)
Chao Liu's avatar
Use DynamicBuffer instead of raw pointer (#32)  
Chao Liu committed
151
152
153
154
155
156 
                    {
                        // HACK: cast pointer of x is bad
                        // TODO: remove this after compiler fix
                        *reinterpret_cast<int32_t*>(&p_data_[i]) =
                            *reinterpret_cast<const int32_t*>(&x);
                    }
Chao Liu's avatar
DL GEMM fp32/fp16/int8 (#41)  
Chao Liu committed
157
158
159 
                    else if constexpr(is_same<remove_cv_t<remove_reference_t<T>>,
                                              int8x8_t>::value &&
                                      is_same<remove_cv_t<remove_reference_t<X>>, int8x8_t>::value)
Chao Liu's avatar
Use DynamicBuffer instead of raw pointer (#32)  
Chao Liu committed
160
161
162
163
164
165 
                    {
                        // HACK: cast pointer of x is bad
                        // TODO: remove this after compiler fix
                        *reinterpret_cast<int32x2_t*>(&p_data_[i]) =
                            *reinterpret_cast<const int32x2_t*>(&x);
                    }
Chao Liu's avatar
DL GEMM fp32/fp16/int8 (#41)  
Chao Liu committed
166
167
168 
                    else if constexpr(is_same<remove_cv_t<remove_reference_t<T>>,
                                              int8x16_t>::value &&
                                      is_same<remove_cv_t<remove_reference_t<X>>, int8x16_t>::value)
Chao Liu's avatar
Use DynamicBuffer instead of raw pointer (#32)  
Chao Liu committed
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
194
195
196
197
198
199
200
201
202
203
204
205
206 
                    {
                        // HACK: cast pointer of x is bad
                        // TODO: remove this after compiler fix
                        *reinterpret_cast<int32x4_t*>(&p_data_[i]) =
                            *reinterpret_cast<const int32x4_t*>(&x);
                    }
                }
                else
                {
                    *reinterpret_cast<X*>(&p_data_[i]) = x;
                }
#endif
            }
        }
        else
        {
            if(is_valid_offset)
            {
                *reinterpret_cast<X*>(&p_data_[i]) = x;
            }
        }
    }

    __host__ __device__ static constexpr bool IsStaticBuffer() { return false; }

    __host__ __device__ static constexpr bool IsDynamicBuffer() { return true; }
};

template <AddressSpace BufferAddressSpace = AddressSpace::Generic,
          typename T,
          typename ElementSpaceSize>
__host__ __device__ constexpr auto make_dynamic_buffer(T* p, ElementSpaceSize element_space_size)
{
    return DynamicBuffer<BufferAddressSpace, T, ElementSpaceSize>{p, element_space_size};
}

} // namespace ck
#endif