quantize_ins.cpp 3 KB
Newer Older
1
2
3
4
#include <migraphx/quantize_ins.hpp>
#include <migraphx/program.hpp>
#include <migraphx/instruction.hpp>
#include <migraphx/iterator_for.hpp>
5
#include <migraphx/op/fp_conversion.hpp>
6
7
8
9
10
11
#include <migraphx/stringutils.hpp>
#include <utility>

namespace migraphx {
inline namespace MIGRAPHX_INLINE_NS {

Shucai Xiao's avatar
Shucai Xiao committed
12
13
14
instruction_ref
insert_fp16(program& prog,
            instruction_ref& ins,
15
16
            shape::type_t type, 
            std::unordered_map<instruction_ref, instruction_ref>& map_fp16)
17
{
Shucai Xiao's avatar
Shucai Xiao committed
18
    if(map_fp16.count(ins) > 0)
19
20
21
22
    {
        return map_fp16[ins];
    }

Shucai Xiao's avatar
Shucai Xiao committed
23
24
    assert(ins->get_shape().type() == shape::float_type ||
           ins->get_shape().type() == shape::double_type);
25
26
27
    instruction_ref ins_fp16{};
    if(ins == std::prev(prog.end()))
    {
28
        ins_fp16 = prog.add_instruction(op::fp_conversion{type}, ins);
29
30
31
32
33
34
35
36
37
38
    }
    else
    {
        ins_fp16 = prog.insert_instruction(std::next(ins), op::fp_conversion{}, ins);
    }
    map_fp16[ins] = ins_fp16;

    return ins_fp16;
}

39
40
void quantize_ins(program& prog, const std::vector<std::string>& ins_names)
{
41
    std::unordered_map<instruction_ref, instruction_ref> map_fp16;
Shucai Xiao's avatar
Shucai Xiao committed
42
    for(auto ins : iterator_for(prog))
43
    {
44
45
        auto name_it = std::find(ins_names.begin(), ins_names.end(), ins->name());
        if(name_it == ins_names.end())
46
47
48
        {
            continue;
        }
49

50
51
52
        shape::type_t orig_type = ins->get_shape().type();
        // process all inputs, if input is a fp32 or fp64, convert it 
        // to a fp16 by adding a fp_conversion operator. 
53
        auto inputs = ins->inputs();
Shucai Xiao's avatar
Shucai Xiao committed
54
        for(auto input : inputs)
55
56
        {
            auto s = input->get_shape();
Shucai Xiao's avatar
Shucai Xiao committed
57
            if(s.type() == shape::float_type || s.type() == shape::double_type)
58
            {
59
60
61
62
63
64
65
66
67
68
69
                // if the input is a fp_conversion operator, uses its input
                // as its current input
                instruction_ref input_fp16{};
                if (input->name() == "fp_conversion")
                {
                    input_fp16 = input->inputs().front();
                }
                else
                {
                    input_fp16 = insert_fp16(prog, input, shape::half_type, map_fp16);
                }
70
71
72
                instruction::replace_argument(ins, input, input_fp16, false);
            }
        }
73
        // recompute the output shape
74
75
        ins->recompute_ins_shape();

76
77
78
        // If output is not the original type, add another instruction
        // to convert it back to the original type
        if(ins->get_shape().type() != orig_type)
79
        {
80
81
82
83
84
85
86
87
88
89
90
            instruction_ref ins_orig_type{};
            if (ins == std::prev(prog.end()))
            {
                ins_orig_type = prog.add_instruction(op::fp_conversion{orig_type}, ins);
            }
            else
            {
                ins_orig_type = prog.insert_instruction(std::next(ins), op::fp_conversion{orig_type}, ins);
            }
            
            prog.replace_instruction(ins, ins_orig_type);
91
        }
92
93
94
95
96
    }
}

} // namespace MIGRAPHX_INLINE_NS
} // namespace migraphx