quantization.cpp 9.79 KB
Newer Older
Shucai Xiao's avatar
Shucai Xiao committed
1
#include <migraphx/quantization.hpp>
2
3
4
#include <migraphx/program.hpp>
#include <migraphx/instruction.hpp>
#include <migraphx/iterator_for.hpp>
5
#include <migraphx/op/convert.hpp>
Shucai Xiao's avatar
Shucai Xiao committed
6
7
8
9
10
#include <migraphx/op/dot.hpp>
#include <migraphx/op/quant_dot.hpp>
#include <migraphx/op/convolution.hpp>
#include <migraphx/op/quant_convolution.hpp>
#include <migraphx/op/multibroadcast.hpp>
11
#include <migraphx/stringutils.hpp>
12
#include <migraphx/ranges.hpp>
13
14
15
16
17
#include <utility>

namespace migraphx {
inline namespace MIGRAPHX_INLINE_NS {

Shucai Xiao's avatar
Shucai Xiao committed
18
instruction_ref insert_quant_ins(program& prog,
Shucai Xiao's avatar
Shucai Xiao committed
19
20
21
22
23
                                 instruction_ref& ins,
                                 shape::type_t type,
                                 std::unordered_map<instruction_ref, instruction_ref>& map_ins,
                                 float scale = 1.0f,
                                 float shift = 0.0f)
24
{
Shucai Xiao's avatar
Shucai Xiao committed
25
    if(map_ins.count(ins) > 0)
26
    {
Shucai Xiao's avatar
Shucai Xiao committed
27
        return map_ins[ins];
28
29
    }

Shucai Xiao's avatar
Shucai Xiao committed
30
    assert(ins->get_shape().type() == shape::float_type ||
Shucai Xiao's avatar
Shucai Xiao committed
31
32
33
           ins->get_shape().type() == shape::double_type ||
           ins->get_shape().type() == shape::int32_type);
    instruction_ref quant_ins{};
Shucai Xiao's avatar
Shucai Xiao committed
34
    quant_ins    = prog.insert_instruction(std::next(ins), op::convert{type, scale, shift}, ins);
Shucai Xiao's avatar
Shucai Xiao committed
35
    map_ins[ins] = quant_ins;
36

Shucai Xiao's avatar
Shucai Xiao committed
37
    return quant_ins;
38
39
}

Shucai Xiao's avatar
Shucai Xiao committed
40
41
42
// This function is to convert any instructions specified in the input
// from double or float to float16 by inserting a convert operator.
// For the conversion, there could be cases of overflowing, but it
Shucai Xiao's avatar
Shucai Xiao committed
43
// is very rare in the area of deeping learning, so we just do a
Shucai Xiao's avatar
Shucai Xiao committed
44
// truncate of the input to get the fp16.
45
void quantize(program& prog, const std::vector<std::string>& ins_names)
46
{
47
    std::unordered_map<instruction_ref, instruction_ref> map_fp16;
Shucai Xiao's avatar
Shucai Xiao committed
48
    for(auto ins : iterator_for(prog))
49
    {
50
        // all indicates every instruction is converted
Shucai Xiao's avatar
Shucai Xiao committed
51
        if((not contains(ins_names, "all")) and (not contains(ins_names, ins->name())))
52
53
54
        {
            continue;
        }
55

56
        shape::type_t orig_type = ins->get_shape().type();
Shucai Xiao's avatar
Shucai Xiao committed
57
        // process all inputs, if input is a fp32 or fp64, convert it
58
        // to a fp16 by adding a convert operator.
59
        auto inputs = ins->inputs();
60
        std::vector<instruction_ref> converted_inputs;
Shucai Xiao's avatar
Shucai Xiao committed
61
        for(auto input : inputs)
62
63
        {
            auto s = input->get_shape();
Shucai Xiao's avatar
Shucai Xiao committed
64
            if(s.type() == shape::float_type || s.type() == shape::double_type)
65
            {
66
                // if the input is a convert operator, uses its input
67
68
                // as its current input
                instruction_ref input_fp16{};
69
                if(input->name() == "convert")
70
71
72
73
74
                {
                    input_fp16 = input->inputs().front();
                }
                else
                {
Shucai Xiao's avatar
Shucai Xiao committed
75
                    input_fp16 = insert_quant_ins(prog, input, shape::half_type, map_fp16);
76
                }
77
                converted_inputs.push_back(input_fp16);
78
            }
79
80
81
82
83
84
            else
            {
                converted_inputs.push_back(input);
            }
        }

85
        // no change for the input, go to the next instruction
Shucai Xiao's avatar
Shucai Xiao committed
86
        if(inputs == converted_inputs)
87
        {
88
            continue;
Shucai Xiao's avatar
Shucai Xiao committed
89
90
91
92
93
94
        }

        auto op        = ins->get_operator();
        auto ins_shape = compute_shape(op, converted_inputs);
        if(ins_shape.type() != orig_type)
        {
Shucai Xiao's avatar
Shucai Xiao committed
95
96
97
98
99
            // check the dead code case to avoid assert
            bool output_empty = ins->outputs().empty();
            auto ins_orig_type =
                prog.insert_instruction(std::next(ins), op::convert{orig_type}, ins);
            if(!output_empty)
100
            {
Shucai Xiao's avatar
Shucai Xiao committed
101
                prog.replace_instruction(ins, ins_orig_type);
Shucai Xiao's avatar
Shucai Xiao committed
102
            }
Shucai Xiao's avatar
Shucai Xiao committed
103
104
105
106
107
108
109
110
111
        }

        prog.replace_instruction(ins, op, converted_inputs);
    }
}

void quantize(program& prog) { quantize(prog, {"all"}); }

// int8 quantization is different from fp16 since int8 can only handle value
Shucai Xiao's avatar
Shucai Xiao committed
112
// -128 ~ 127. To convert the float or double to int8, we need a scale and
Shucai Xiao's avatar
Shucai Xiao committed
113
// a shift, then the convert can be done as v_int8 = fp * scale + shift.
Shucai Xiao's avatar
Shucai Xiao committed
114
// To simplify the changes, we consider shift as 0.0f for now.
Shucai Xiao's avatar
Shucai Xiao committed
115
116
117
118
void quantize_int8(program& prog, const std::vector<std::string>& ins_names)
{
    // For now, we only support the int8 quantization of gemm and convolution
    std::vector<std::string> op_names = {"dot", "convolution"};
Shucai Xiao's avatar
Shucai Xiao committed
119
    if(!std::all_of(ins_names.begin(), ins_names.end(), [&](auto name) {
Shucai Xiao's avatar
Shucai Xiao committed
120
           return (std::find(op_names.begin(), op_names.end(), name) != op_names.end());
Shucai Xiao's avatar
Shucai Xiao committed
121
       }))
Shucai Xiao's avatar
Shucai Xiao committed
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
    {
        MIGRAPHX_THROW("QUANTIZE_INT8: only support DOT and CONVOLUTION operation");
    }

    // tmp value used just testing
    std::vector<std::pair<float, float>> int8_param{{1.0f, 0.0f}, {1.0f, 0.0f}, {1.0f, 0.0f}};

    std::unordered_map<instruction_ref, instruction_ref> map_quant_ins;
    for(auto ins : iterator_for(prog))
    {
        if(not contains(ins_names, ins->name()))
        {
            continue;
        }

        shape::type_t orig_type = ins->get_shape().type();

        // for the dot operator, there could be 2 or 3 input arguments
        // if the 3rd argument is available, convert it to an int32.
        std::vector<instruction_ref> converted_inputs;

        // process all inputs, if input is a fp32 or fp64, convert it
Shucai Xiao's avatar
Shucai Xiao committed
144
        // to a int8 type by adding a convert operator and replace
Shucai Xiao's avatar
Shucai Xiao committed
145
        // the operator with the corresponding int8 version
Shucai Xiao's avatar
Shucai Xiao committed
146
        auto inputs             = ins->inputs();
Shucai Xiao's avatar
Shucai Xiao committed
147
148
149
        std::size_t param_index = 0;
        for(auto input : inputs)
        {
Shucai Xiao's avatar
Shucai Xiao committed
150
151
            // In general, the target_type is int8, but for the dot
            // operation, if it has 3 inputs, then the last one should
Shucai Xiao's avatar
Shucai Xiao committed
152
153
            // be converted to int32_type
            shape::type_t quant_type = shape::int8_type;
Shucai Xiao's avatar
Shucai Xiao committed
154
            if(ins->name() == "dot" and inputs.size() == 3 and input == inputs.back())
Shucai Xiao's avatar
Shucai Xiao committed
155
            {
Shucai Xiao's avatar
Shucai Xiao committed
156
157
                quant_type = shape::int32_type;
            }
Shucai Xiao's avatar
Shucai Xiao committed
158

Shucai Xiao's avatar
Shucai Xiao committed
159
            auto param = int8_param[param_index++];
Shucai Xiao's avatar
Shucai Xiao committed
160
161
162
            auto s     = input->get_shape();
            if(s.type() == shape::float_type || s.type() == shape::double_type ||
               s.type() == shape::int32_type)
Shucai Xiao's avatar
Shucai Xiao committed
163
164
165
166
167
168
169
            {
                // if the input is a convert operator, uses its input
                // as its current input
                instruction_ref quant_input{};
                if(input->name() == "convert")
                {
                    auto tmp_ins = input->inputs().front();
Shucai Xiao's avatar
Shucai Xiao committed
170
                    if(tmp_ins->get_shape().type() == quant_type)
Shucai Xiao's avatar
Shucai Xiao committed
171
172
173
174
175
                    {
                        quant_input = input->inputs().front();
                    }
                    else
                    {
Shucai Xiao's avatar
Shucai Xiao committed
176
177
                        quant_input = insert_quant_ins(
                            prog, input, quant_type, map_quant_ins, param.first, param.second);
Shucai Xiao's avatar
Shucai Xiao committed
178
179
180
                    }
                }
                else
181
                {
Shucai Xiao's avatar
Shucai Xiao committed
182
183
                    quant_input = insert_quant_ins(
                        prog, input, quant_type, map_quant_ins, param.first, param.second);
184
                }
Shucai Xiao's avatar
Shucai Xiao committed
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
                converted_inputs.push_back(quant_input);
            }
            else
            {
                converted_inputs.push_back(input);
            }
        }

        // no change for the input, go to the next instruction
        if(inputs == converted_inputs)
        {
            continue;
        }

        auto op        = ins->get_operator();
        auto ins_shape = compute_shape(op, converted_inputs);
        if(ins_shape.type() != orig_type)
        {
            // check the dead code case to avoid assert
            bool output_empty = ins->outputs().empty();
            // this conversion can be only from int32 to float or double
            auto ins_orig_type =
                prog.insert_instruction(std::next(ins), op::convert{orig_type}, ins);
            if(!output_empty)
            {
                prog.replace_instruction(ins, ins_orig_type);
211
            }
212
        }
Shucai Xiao's avatar
Shucai Xiao committed
213

Shucai Xiao's avatar
Shucai Xiao committed
214
215
216
        // When converting from other types to int8_type, there are parameters
        // used as scale and shift(.0f), which will generate results diffrent from
        // the original results. To adjust the output to be "correct(approximatly
Shucai Xiao's avatar
Shucai Xiao committed
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
        // equal)", we need additional calculation for the adjustment
        if (ins->name() == "dot")
        {
            auto dot_op = any_cast<op::dot>(ins->get_operator());
            int32_t quant_alpha = static_cast<int32_t>(dot_op.alpha / (int8_param[0].first * int8_param[1].first) + 0.5f);
            int32_t quant_beta = static_cast<int32_t>(dot_op.beta + 0.5f);
            prog.replace_instruction(ins, op::quant_dot{quant_alpha, quant_beta}, converted_inputs);
        }
        else if (ins->name() == "convolution")
        {
            // Current MIOpen convolution does not support alpha and beta, 
            // so we need a separate multiply to adjust the output
            auto conv_op = any_cast<op::convolution>(ins->get_operator());
            auto padding  = conv_op.padding;
            auto stride   = conv_op.stride;
            auto dilation = conv_op.dilation;
            auto padding_mode = conv_op.padding_mode;
            auto group = conv_op.group;
            auto adjust_factor = 1.0 / int8_param[0].first * int8_param[1].first;

            auto conv_res = prog.insert_instruction(ins, op::quant_convolution{padding, stride, dilation, padding_mode, group}, converted_inputs);
            auto conv_lens = conv_res->get_shape().lens();
            auto fl = prog.add_literal(literal(adjust_factor));
            auto adj_fact = prog.insert_instruction(ins, op::multibroadcast{conv_lens}, fl);
            prog.replace_instruction(ins, adj_fact);
        }
        else
        {
            MIGRAPHX_THROW("INT8_QUANTIZE: does not support operator" + ins->name());
        }
Shucai Xiao's avatar
Shucai Xiao committed
247

Shucai Xiao's avatar
Shucai Xiao committed
248
        prog.replace_instruction(ins, op, converted_inputs);
249
250
251
252
253
    }
}

} // namespace MIGRAPHX_INLINE_NS
} // namespace migraphx