f8_utils.hpp 7.77 KB
Newer Older
1
2
3
4
5
6
7
8
9
// SPDX-License-Identifier: MIT
// Copyright (c) 2018-2023, Advanced Micro Devices, Inc. All rights reserved.

#pragma once

#include "ck/utility/statically_indexed_array.hpp"

namespace ck {

Rostyslav Geyyer's avatar
Format  
Rostyslav Geyyer committed
10
11
using f8_t   = uint8_t;
using half_t = _Float16;
12
13

// fp8 rounding modes
14
15
// use standard for rounding to nearest, the faster one
// use stochastic for stochastic rounding, helps to avoid error accumulation
16
17
18
19
20
21
enum class f8_rounding_mode
{
    standard,
    stochastic
};

Rostyslav Geyyer's avatar
Rostyslav Geyyer committed
22
23
template <typename T, bool negative_zero_nan, bool clip, bool stoch>
__host__ __device__ f8_t run_cast_to_f8(T x, uint32_t rng)
24
{
Rostyslav Geyyer's avatar
Rostyslav Geyyer committed
25
26
27
    // check data type
    constexpr bool is_half  = std::is_same<T, half_t>::value;
    constexpr bool is_float = std::is_same<T, float>::value;
28

Rostyslav Geyyer's avatar
Rostyslav Geyyer committed
29
    // fp8 exponent/mantissa layout
Rostyslav Geyyer's avatar
Format  
Rostyslav Geyyer committed
30
    constexpr int f8_exp  = 4;
Rostyslav Geyyer's avatar
Rostyslav Geyyer committed
31
    constexpr int f8_mant = 3;
32

Rostyslav Geyyer's avatar
Rostyslav Geyyer committed
33
    // resulting type exponent/mantissa layout
Rostyslav Geyyer's avatar
Format  
Rostyslav Geyyer committed
34
    constexpr int type_exp  = is_half ? 5 : 8;
Rostyslav Geyyer's avatar
Rostyslav Geyyer committed
35
    constexpr int type_mant = is_half ? 10 : 23;
36
37

    int exponent;
Rostyslav Geyyer's avatar
Rostyslav Geyyer committed
38
39
    uint32_t head, mantissa, sign;
    // nan code is same for float and half
Rostyslav Geyyer's avatar
Format  
Rostyslav Geyyer committed
40
    constexpr uint8_t nan_code  = 0x80;
Rostyslav Geyyer's avatar
Rostyslav Geyyer committed
41
    constexpr uint32_t nan_mask = is_half ? 0x7C00 : 0x7F800000;
42

Rostyslav Geyyer's avatar
Rostyslav Geyyer committed
43
    // convert to bitwise
Rostyslav Geyyer's avatar
Format  
Rostyslav Geyyer committed
44
45
    typedef typename std::conditional<std::is_same<T, half_t>::value, uint16_t, uint32_t>::type
        T_bitwise;
Rostyslav Geyyer's avatar
Rostyslav Geyyer committed
46
    T_bitwise x_bitwise = *(reinterpret_cast<T_bitwise*>(&x));
47

Rostyslav Geyyer's avatar
Rostyslav Geyyer committed
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
    // unpack the input, depends on datatype
    if constexpr(is_float)
    {
        head     = x_bitwise & 0xFF800000;
        mantissa = x_bitwise & 0x7FFFFF;
        exponent = (head >> type_mant) & 0xFF;
        sign     = head >> (type_exp + type_mant);
    }
    else if constexpr(is_half)
    {
        head     = x_bitwise & 0xFC00;
        mantissa = x_bitwise & 0x3FF;
        exponent = (head >> type_mant) & 0x1F;
        sign     = head >> (type_exp + type_mant);
    }

Rostyslav Geyyer's avatar
Format  
Rostyslav Geyyer committed
64
65
66
67
68
    uint32_t signed_inf   = (sign << (type_exp + type_mant)) + (((1 << type_exp) - 1) << type_mant);
    uint32_t drop_mask    = (1 << (type_mant - f8_mant)) - 1;
    constexpr int max_exp = (1 << f8_exp) - (negative_zero_nan ? 1 : 2);
    constexpr int exp_low_cutoff =
        (1 << (type_exp - 1)) - (1 << (f8_exp - 1)) + 1 - (negative_zero_nan ? 1 : 0);
69

Rostyslav Geyyer's avatar
Rostyslav Geyyer committed
70
    if constexpr(negative_zero_nan)
71
    {
Rostyslav Geyyer's avatar
Rostyslav Geyyer committed
72
73
        if((x_bitwise & nan_mask) == nan_mask)
            return nan_code;
74
75
76
    }
    else
    {
Rostyslav Geyyer's avatar
Rostyslav Geyyer committed
77
        if((x_bitwise & nan_mask) == nan_mask)
78
79
80
            return signed_inf + (mantissa != 0 ? 1 : 0);
    }

Rostyslav Geyyer's avatar
Rostyslav Geyyer committed
81
    // check if x is 0.0
82
83
84
85
86
    if(x_bitwise == 0)
        return 0;

    exponent -= exp_low_cutoff - 1;
    if(exponent <= 0)
Rostyslav Geyyer's avatar
Rostyslav Geyyer committed
87
88
89
        drop_mask = (1 << (type_mant - f8_mant + 1 - exponent)) - 1;
    mantissa += 1 << type_mant;
    // apply random number if needed
90
    mantissa += (stoch ? rng : mantissa) & drop_mask;
Rostyslav Geyyer's avatar
Rostyslav Geyyer committed
91
    if(mantissa >= (2 << type_mant))
92
93
94
95
    {
        mantissa >>= 1;
        exponent++;
    }
Rostyslav Geyyer's avatar
Rostyslav Geyyer committed
96
    mantissa >>= (type_mant - f8_mant);
97

Rostyslav Geyyer's avatar
Rostyslav Geyyer committed
98
    // check negative exponent
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
    if(exponent <= 0)
    {
        if(x_bitwise == 0)
            return 0;
        else
        {
            // subnormal range; represented by a subnormal float8 (exponent 0)
            // and involves loss of accuracy
            mantissa >>= 1 - exponent;
            exponent = 0;
        }
    }
    // above range: quantize to maximum possible float of the same sign
    else if(exponent > max_exp)
    {
        if(clip)
        {
Rostyslav Geyyer's avatar
Rostyslav Geyyer committed
116
            mantissa = (1 << f8_mant) - 1;
117
118
119
120
121
122
123
            exponent = max_exp;
        }
        else
        {
            return signed_inf;
        }
    }
Rostyslav Geyyer's avatar
Rostyslav Geyyer committed
124
125

    // check if x is 0.0 or -0.0
126
    if(exponent == 0 && mantissa == 0)
Rostyslav Geyyer's avatar
Rostyslav Geyyer committed
127
128
129
130
131
132
133
134
135
136
137
138
139
140
        return negative_zero_nan ? 0 : (sign << (f8_exp + f8_mant));
    mantissa &= (1 << f8_mant) - 1;
    return (sign << (f8_exp + f8_mant)) | (exponent << f8_mant) | mantissa;
}

template <typename T, bool negative_zero_nan, bool clip, bool stoch>
__host__ __device__ f8_t cast_to_f8(T x, uint32_t rng)
{
    // check datatype
    constexpr bool is_half  = std::is_same<T, half_t>::value;
    constexpr bool is_float = std::is_same<T, float>::value;
    static_assert(is_half || is_float, "Only half and float can be casted to f8.");

    return run_cast_to_f8<T, negative_zero_nan, clip, stoch>(x, rng);
141
142
}

Rostyslav Geyyer's avatar
Rostyslav Geyyer committed
143
144
template <typename T, bool negative_zero_nan>
__host__ __device__ T run_cast_from_f8(f8_t x)
145
{
Rostyslav Geyyer's avatar
Rostyslav Geyyer committed
146
147
148
149
    // check data type
    constexpr bool is_half  = std::is_same<T, half_t>::value;
    constexpr bool is_float = std::is_same<T, float>::value;

150
    // fp8 exponent/mantissa layout
Rostyslav Geyyer's avatar
Format  
Rostyslav Geyyer committed
151
    constexpr int f8_exp  = 4;
Rostyslav Geyyer's avatar
Rostyslav Geyyer committed
152
    constexpr int f8_mant = 3;
153

Rostyslav Geyyer's avatar
Rostyslav Geyyer committed
154
    // resulting type exponent/mantissa layout
Rostyslav Geyyer's avatar
Format  
Rostyslav Geyyer committed
155
    constexpr int type_exp  = is_half ? 5 : 8;
Rostyslav Geyyer's avatar
Rostyslav Geyyer committed
156
157
158
159
160
161
162
163
164
165
166
    constexpr int type_mant = is_half ? 10 : 23;

    // prepare the codes
    constexpr uint8_t nan_code = 0x80;
    T fInf, fNegInf, fNaN, fNeg0;
    if constexpr(is_half)
    {
        constexpr uint16_t ihInf    = 0x7C00;
        constexpr uint16_t ihNegInf = 0xFC00;
        constexpr uint16_t ihNaN    = 0x7C01;
        constexpr uint16_t ihNeg0   = 0x8000;
Rostyslav Geyyer's avatar
Format  
Rostyslav Geyyer committed
167
168
169
170
        fInf                        = *(reinterpret_cast<const half_t*>(&ihInf));
        fNegInf                     = *(reinterpret_cast<const half_t*>(&ihNegInf));
        fNaN                        = *(reinterpret_cast<const half_t*>(&ihNaN));
        fNeg0                       = *(reinterpret_cast<const half_t*>(&ihNeg0));
Rostyslav Geyyer's avatar
Rostyslav Geyyer committed
171
172
173
174
175
176
177
    }
    else if constexpr(is_float)
    {
        constexpr uint32_t ifInf    = 0x7F800000;
        constexpr uint32_t ifNegInf = 0xFF800000;
        constexpr uint32_t ifNaN    = 0x7F800001;
        constexpr uint32_t ifNeg0   = 0x80000000;
Rostyslav Geyyer's avatar
Format  
Rostyslav Geyyer committed
178
179
180
181
        fInf                        = *(reinterpret_cast<const float*>(&ifInf));
        fNegInf                     = *(reinterpret_cast<const float*>(&ifNegInf));
        fNaN                        = *(reinterpret_cast<const float*>(&ifNaN));
        fNeg0                       = *(reinterpret_cast<const float*>(&ifNeg0));
Rostyslav Geyyer's avatar
Rostyslav Geyyer committed
182
    }
183
184

    // unpack the input
Rostyslav Geyyer's avatar
Rostyslav Geyyer committed
185
186
187
    uint32_t sign     = x >> (f8_exp + f8_mant);
    uint32_t mantissa = x & ((1 << f8_mant) - 1);
    int exponent      = (x & 0x7F) >> f8_mant;
188

Rostyslav Geyyer's avatar
Format  
Rostyslav Geyyer committed
189
190
    constexpr int exp_low_cutoff =
        (1 << (type_exp - 1)) - (1 << (f8_exp - 1)) + 1 - (negative_zero_nan ? 1 : 0);
Rostyslav Geyyer's avatar
Rostyslav Geyyer committed
191
    typename std::conditional<std::is_same<T, half_t>::value, uint16_t, uint32_t>::type retval;
Rostyslav Geyyer's avatar
Rostyslav Geyyer committed
192
193

    if constexpr(negative_zero_nan)
194
    {
Rostyslav Geyyer's avatar
Rostyslav Geyyer committed
195
        if(x == nan_code)
196
197
198
199
            return fNaN;
    }
    else
    {
Rostyslav Geyyer's avatar
Rostyslav Geyyer committed
200
        if(x == nan_code)
201
            return fNeg0;
Rostyslav Geyyer's avatar
Rostyslav Geyyer committed
202
        if(exponent == ((1 << f8_exp) - 1))
203
204
205
206
207
208
209
            return (mantissa == 0) ? (sign ? fNegInf : fInf) : fNaN;
    }

    // subnormal input
    if(exponent == 0)
    {
        // guaranteed mantissa!=0 since cases 0x0 and 0x80 are handled above
Rostyslav Geyyer's avatar
Rostyslav Geyyer committed
210
        int sh = 1 + __builtin_clz(mantissa) - ((1 + type_exp + type_mant) - f8_mant);
211
212
213
214
215
216
217
218
219
        mantissa <<= sh;
        exponent += 1 - sh;
        /*
        exponent++;
        while(mantissa<(1<<wm)) {
        mantissa <<= 1;
        exponent--;
        }
        */
Rostyslav Geyyer's avatar
Rostyslav Geyyer committed
220
        mantissa &= ((1 << f8_mant) - 1);
221
222
    }
    exponent += exp_low_cutoff - 1;
Rostyslav Geyyer's avatar
Rostyslav Geyyer committed
223
    mantissa <<= type_mant - f8_mant;
224
225
226
227

    // subnormal output (occurs when T=half, we=5, negative_zero_nan=true)
    if(exponent <= 0)
    {
Rostyslav Geyyer's avatar
Rostyslav Geyyer committed
228
        mantissa |= 1 << type_mant;
229
230
231
232
        mantissa >>= 1 - exponent;
        exponent = 0;
    }

Rostyslav Geyyer's avatar
Rostyslav Geyyer committed
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
    retval = (sign << (type_exp + type_mant)) | (exponent << type_mant) | mantissa;
    return *(reinterpret_cast<const T*>(&retval));
}

template <typename T, bool negative_zero_nan>
__host__ __device__ T cast_from_f8(f8_t x)
{
    // check datatype
    constexpr bool is_half  = std::is_same<T, half_t>::value;
    constexpr bool is_float = std::is_same<T, float>::value;
    static_assert(is_half || is_float, "only half and float are supported.");

    // check if x is 0.0
    if(x == 0)
        return static_cast<T>(0);

    return run_cast_from_f8<T, negative_zero_nan>(x);
250
251
252
}

} // namespace ck